"""Testing stats reporting."""
# Authors: Alex Gramfort <alexandre.gramfort@inria.fr>
#
# License: BSD-3-Clause

import itertools

import pytest
import numpy as np

import mne
from mne.datasets import testing

from mne_bids import BIDSPath, write_raw_bids
from mne_bids.stats import count_events
from mne_bids.read import _from_tsv
from mne_bids.write import _write_tsv

data_path = testing.data_path(download=False)


def _make_dataset(root, subjects, tasks=(None,), runs=(None,), sessions=(None,)):
    raw_fname = data_path / "MEG" / "sample" / "sample_audvis_trunc_raw.fif"
    raw = mne.io.read_raw(raw_fname)
    raw.info["line_freq"] = 60.0
    events = mne.find_events(raw)
    event_id = {
        "auditory/left": 1,
        "auditory/right": 2,
        "visual/left": 3,
        "visual/right": 4,
        "face": 5,
        "button": 32,
    }

    for subject, session, task, run in itertools.product(
        subjects, sessions, tasks, runs
    ):
        bids_path = BIDSPath(
            subject=subject,
            session=session,
            run=run,
            task=task,
            root=root,
        )
        write_raw_bids(
            raw,
            bids_path,
            events=events,
            event_id=event_id,
            overwrite=True,
            verbose=False,
        )

    return root, events, event_id


def _check_counts(
    counts, events, event_id, subjects, tasks=(None,), runs=(None,), sessions=(None,)
):
    if (sessions[0] is None) and (runs[0] is None):
        assert np.all(counts.index == subjects)
    else:
        assert np.all(counts.index.levels[0] == subjects)

    if (sessions[0] is not None) and (runs[0] is not None):
        assert np.all(counts.index.levels[1] == sessions)
        assert np.all(counts.index.levels[2] == runs)
    elif sessions[0] is not None:
        assert np.all(counts.index.levels[1] == sessions)
    elif runs[0] is not None:
        assert np.all(counts.index.levels[1] == runs)

    assert np.all(counts.columns.levels[0] == tasks)
    assert sorted(counts.columns.levels[1]) == sorted(event_id.keys())
    for k, v in event_id.items():
        key = (subjects[0],)
        if sessions[0] is not None:
            key += (sessions[0],)
        if runs[0] is not None:
            key += (runs[0],)
        key = key if len(key) > 1 else key[0]

        assert counts.at[key, (tasks[0], k)] == (events[:, 2] == v).sum()


@pytest.mark.parametrize(
    ("subjects", "tasks", "runs", "sessions"),
    [
        (["01"], ["task1"], ["01"], ["01"]),
        (["01", "02"], ["task1"], ["01"], ["01"]),
        (["01", "02"], ["task1", "task2"], ["01"], ["01"]),
        (["01"], ["task1", "task2"], [None], ["01"]),
        (["01"], ["task1", "task2"], ["01"], [None]),
        (["01"], ["task1", "task2"], [None], [None]),
    ],
)
@testing.requires_testing_data
def test_count_events(tmp_path, subjects, tasks, runs, sessions):
    """Test the event counts."""
    pytest.importorskip("pandas")
    root, events, event_id = _make_dataset(tmp_path, subjects, tasks, runs, sessions)

    counts = count_events(root)

    _check_counts(counts, events, event_id, subjects, tasks, runs, sessions)


@testing.requires_testing_data
def test_count_events_bids_path(tmp_path):
    """Test the event counts passing a BIDSPath."""
    pytest.importorskip("pandas")
    root, events, event_id = _make_dataset(
        tmp_path, subjects=["01", "02"], tasks=["task1"]
    )

    with pytest.raises(ValueError, match="datatype .*anat.* is not supported"):
        bids_path = BIDSPath(root=root, subject="01", datatype="anat")
        count_events(bids_path)

    bids_path = BIDSPath(root=root, subject="01", datatype="meg")
    counts = count_events(bids_path)

    _check_counts(counts, events, event_id, subjects=["01"], tasks=["task1"])


@testing.requires_testing_data
def test_count_no_events_file(tmp_path):
    """Test count_events with no event present."""
    pytest.importorskip("pandas")
    raw_fname = data_path / "MEG" / "sample" / "sample_audvis_trunc_raw.fif"
    raw = mne.io.read_raw(raw_fname)
    raw.info["line_freq"] = 60.0
    root = str(tmp_path)

    bids_path = BIDSPath(
        subject="01",
        task="task1",
        root=root,
    )
    write_raw_bids(raw, bids_path, overwrite=True, verbose=False)

    with pytest.raises(ValueError, match="No events files found."):
        count_events(root)


@testing.requires_testing_data
def test_count_no_events_column(tmp_path):
    """Test case where events.tsv doesn't contain [stim,trial]_type column."""
    pytest.importorskip("pandas")
    subject, task, run, session, datatype = "01", "task1", "01", "01", "meg"
    root, events, event_id = _make_dataset(
        tmp_path, [subject], [task], [run], [session]
    )

    # Delete the `stim_type` column.
    events_tsv_fpath = BIDSPath(
        root=root,
        subject=subject,
        task=task,
        run=run,
        session=session,
        datatype=datatype,
        suffix="events",
        extension=".tsv",
    ).fpath
    events_tsv = _from_tsv(events_tsv_fpath)
    events_tsv["stim_type"] = events_tsv["trial_type"]
    del events_tsv["trial_type"]
    _write_tsv(fname=events_tsv_fpath, dictionary=events_tsv, overwrite=True)

    counts = count_events(root)
    _check_counts(counts, events, event_id, [subject], [task], [run], [session])
