import tempfile

import mne

from meggie.datatypes.evoked.evoked import Evoked

from meggie_statistics.utilities.stats import prepare_data_for_permutation
from meggie_statistics.utilities.stats import permutation_analysis
from meggie.utilities.testing import create_evoked_conditions_experiment


def test_permutation_analysis():
    with tempfile.TemporaryDirectory() as dirpath:
        # generate dataset with generate_datasets to tmp folder (n=10)
        experiment = create_evoked_conditions_experiment(
            dirpath, "permutation_test_experiment", n_subjects=10
        )

        for subject_name, subject in experiment.subjects.items():
            raw = subject.get_raw()

            raw.resample(100)
            subject.save()

            events = mne.find_events(raw)
            evoked_la = mne.Epochs(raw, events, event_id=1).average()
            evoked_ra = mne.Epochs(raw, events, event_id=2).average()
            content = {"la": evoked_la, "ra": evoked_ra}
            meggie_evoked = Evoked(
                "Evoked", subject.evoked_directory, {}, content=content
            )
            subject.add(meggie_evoked, "evoked")

        experiment.activate_subject(subject_name)

        design = "within-subjects"
        conditions = ["la", "ra"]
        groups = {"1": list(experiment.subjects.keys())}
        location_limits = ("ch_type", "grad")
        time_limits = None
        threshold = 0.01
        n_permutations = 1000
        info, data, adjacency = prepare_data_for_permutation(
            experiment=experiment,
            design=design,
            groups=groups,
            item_type="evoked",
            item_name="Evoked",
            location_limits=location_limits,
            time_limits=time_limits,
            frequency_limits=None,
            data_format=("locations", "times"),
            no_threading=True,
        )

        results = permutation_analysis(
            data=data,
            design=design,
            conditions=conditions,
            groups=groups,
            threshold=threshold,
            adjacency=adjacency,
            n_permutations=n_permutations,
            no_threading=True,
            random_state=10,
        )

        # Correct cluster shape
        assert results["1"][0].shape == (71, 202)
