import os

import pytest
import numpy as np

import mdevaluate
import mdevaluate.extra.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(
        [
            174.46253634,
            174.60905476,
            178.57658092,
            182.43001192,
            180.57916378,
            176.49886217,
            178.96018547,
            181.13561782,
            178.31026314,
            176.08903996,
            180.71215345,
            181.59703135,
            180.34329368,
            187.02474488,
            197.99167477,
            214.05788031,
            245.58571282,
            287.52457507,
            331.53492965,
        ]
    )

    OW = trajectory.subset(atom_name="OW")

    box = trajectory[0].box
    box_voxels = (np.diag(box) // [0.05, 0.05, 0.05] + [1, 1, 1]) * [0.05, 0.05, 0.05]
    occupation_matrix = fel.occupation_matrix(OW, skip=0, segments=1000)
    radius_maxima = 0.05 * 3 ** (1 / 2) + 0.05 / 100
    maxima_matrix = fel.find_maxima(
        occupation_matrix,
        box=box_voxels,
        radius=radius_maxima,
        pore_geometry="cylindrical"
    )
    maxima_matrix = fel.add_distances(maxima_matrix, "cylindrical", np.diag(box) / 2)
    r_bins = np.arange(0, 1, 0.02)
    distance_bins = np.arange(0.05, 2.05, 0.1)
    energy_df = fel.distance_resolved_energies(
        maxima_matrix, distance_bins, r_bins, box, "cylindrical", 225
    )
    result = fel.find_energy_maxima(energy_df, r_min=0.05, r_max=0.15)

    assert (np.round(np.array(result["energy"])) == np.round(test_array)).all()