# Authors: Robert Luke <mail@robertluke.net>
#
# License: BSD (3-clause)


import os

import mne
import numpy as np
import pytest

import mne_nirs
from mne_nirs.experimental_design import make_first_level_design_matrix
from mne_nirs.statistics import run_glm
from mne_nirs.statistics._glm_level_first import _compute_contrast
from mne_nirs.utils._io import _tidy_long_to_wide, glm_to_tidy


@pytest.mark.filterwarnings("ignore:.*more comprehensive.*:")
def test_io():
    num_chans = 6
    fnirs_data_folder = mne.datasets.fnirs_motor.data_path()
    fnirs_raw_dir = os.path.join(fnirs_data_folder, "Participant-1")
    raw_intensity = mne.io.read_raw_nirx(fnirs_raw_dir).load_data()
    raw_intensity.resample(0.2)
    raw_intensity.annotations.description[:] = [
        "e" + d.replace(".", "p") for d in raw_intensity.annotations.description
    ]
    raw_od = mne.preprocessing.nirs.optical_density(raw_intensity)
    raw_haemo = mne.preprocessing.nirs.beer_lambert_law(raw_od, ppf=0.1)
    raw_haemo = mne_nirs.channels.get_long_channels(raw_haemo)
    raw_haemo.pick(picks=range(num_chans))
    design_matrix = make_first_level_design_matrix(
        raw_intensity,
        hrf_model="spm",
        stim_dur=5.0,
        drift_order=3,
        drift_model="polynomial",
    )
    glm_est = run_glm(raw_haemo, design_matrix)
    df = glm_to_tidy(raw_haemo, glm_est.data, design_matrix)
    assert df.shape == (48, 12)
    assert set(df.columns) == {
        "ch_name",
        "Condition",
        "df",
        "mse",
        "p_value",
        "t",
        "theta",
        "Source",
        "Detector",
        "Chroma",
        "Significant",
        "se",
    }
    num_conds = 8  # triggers (1, 2, 3, 15) + 3 drifts + constant
    assert df.shape[0] == num_chans * num_conds
    assert len(df["se"]) == 48
    assert sum(df["se"]) > 0  # Check isn't nan
    assert len(df["df"]) == 48
    assert sum(df["df"]) > 0  # Check isn't nan
    assert len(df["p_value"]) == 48
    assert sum(df["p_value"]) > 0  # Check isn't nan
    assert len(df["theta"]) == 48
    assert sum(df["theta"]) > 0  # Check isn't nan
    assert len(df["t"]) == 48
    assert sum(df["t"]) > -99999  # Check isn't nan

    contrast_matrix = np.eye(design_matrix.shape[1])
    basic_conts = dict(
        [(column, contrast_matrix[i]) for i, column in enumerate(design_matrix.columns)]
    )
    contrast_LvR = basic_conts["e2p0"] - basic_conts["e3p0"]

    contrast = _compute_contrast(glm_est.data, contrast_LvR)
    df = glm_to_tidy(raw_haemo, contrast, design_matrix)
    assert df.shape == (6, 10)
    assert set(df.columns) == {
        "ch_name",
        "ContrastType",
        "z_score",
        "stat",
        "p_value",
        "effect",
        "Source",
        "Detector",
        "Chroma",
        "Significant",
    }

    contrast = _compute_contrast(glm_est.data, contrast_LvR, contrast_type="F")
    df = glm_to_tidy(raw_haemo, contrast, design_matrix, wide=False)
    df = _tidy_long_to_wide(df)
    assert df.shape == (6, 10)
    assert set(df.columns) == {
        "ch_name",
        "ContrastType",
        "z_score",
        "stat",
        "p_value",
        "effect",
        "Source",
        "Detector",
        "Chroma",
        "Significant",
    }

    with pytest.raises(TypeError, match="Unknown statistic type"):
        glm_to_tidy(raw_haemo, [1, 2, 3], design_matrix, wide=False)
