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,
            12.87438176,
            4.95868203,
            11.02055197,
            5.44195534,
            6.73933442,
            3.30971789,
            6.10424055,
            8.56153733,
            5.45777331,
            5.64545817,
            8.42100423,
            6.28132121,
            7.4777172,
            11.64839354,
            4.52566354,
            40.84730838,
            93.86241602,
            140.3039937,
            173.55970021,
        ]
    )

    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 (np.round(energy_differences) == np.round(test_array)).all()