import os

import pytest
import numpy as np

import mdevaluate
from mdevaluate import free_energy_landscape as fel


@pytest.fixture
def trajectory(request):
    return mdevaluate.open(os.path.join(os.path.dirname(__file__), 'data/pore'))


def test_get_fel(trajectory):
    test_array = np.array(
        [
            0.0,
            13.162354034697204,
            5.327100985208421,
            9.558746399158396,
            4.116475238453127,
            6.305715728953043,
            3.231102391108276,
            5.896478799115712,
            8.381981206446293,
            5.1191684352849816,
            5.361112857237105,
            8.053932845998895,
            6.895396051256847,
            7.588888886900885,
            11.223429636542576,
            3.779149304024221,
            40.64319010769286,
            93.1120609754045,
            136.99287780099627,
            171.4403749377496,
        ]
    )

    oxygens_water = trajectory.subset(atom_name="OW", residue_name="SOL")
    r, energy_differences = fel.get_fel(
        oxygens_water,
        os.path.join(os.path.dirname(__file__), 'data/pore'),
        "cylindrical",
        225,
        edge=0.05,
        radiusmin=0.05,
        radiusmax=2.05,
        z=[-np.inf, np.inf],
        overwrite=True,
    )

    assert (energy_differences == test_array).all()