# Authors: The MNE-Python contributors.
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

import numpy as np
import pytest
from numpy.testing import assert_allclose, assert_array_equal

from mne import Annotations, events_from_annotations
from mne.chpi import read_head_pos
from mne.datasets import testing
from mne.io import read_raw_fif
from mne.preprocessing import (
    annotate_break,
    annotate_movement,
    annotate_muscle_zscore,
    compute_average_dev_head_t,
)
from mne.tests.test_annotations import _assert_annotations_equal
from mne.transforms import _angle_dist_between_rigid, quat_to_rot, rot_to_quat

data_path = testing.data_path(download=False)
sss_path = data_path / "SSS"
pos_fname = sss_path / "test_move_anon_raw.pos"
raw_fname = sss_path / "test_move_anon_raw.fif"


@testing.requires_testing_data
@pytest.mark.parametrize("meas_date", (None, "orig"))
def test_movement_annotation_head_correction(meas_date):
    """Test correct detection movement artifact and dev_head_t."""
    raw = read_raw_fif(raw_fname, allow_maxshield="yes").load_data()
    pos = read_head_pos(pos_fname)
    if meas_date is None:
        raw.set_meas_date(None)
    else:
        assert meas_date == "orig"
    raw_unannot = raw.copy()

    # Check 5 rotation segments are detected
    annot_rot, [] = annotate_movement(raw, pos, rotation_velocity_limit=5)
    assert annot_rot.orig_time == raw.info["meas_date"]
    assert annot_rot.duration.size == 5

    # Check 2 translation vel. segments are detected
    annot_tra, [] = annotate_movement(raw, pos, translation_velocity_limit=0.05)
    assert annot_tra.duration.size == 2

    # Check 1 movement distance segment is detected
    annot_dis, _ = annotate_movement(raw, pos, mean_distance_limit=0.02)
    assert annot_dis.duration.size == 1

    # Check correct trans mat
    annot_all_2 = annotate_movement(
        raw,
        pos,
        rotation_velocity_limit=5,
        translation_velocity_limit=0.05,
        mean_distance_limit=0.02,
    )[0]
    assert (
        annot_rot.orig_time
        == annot_tra.orig_time
        == annot_dis.orig_time
        == raw.info["meas_date"]
    )
    annot_all = annot_rot + annot_tra + annot_dis
    _assert_annotations_equal(annot_all_2, annot_all)
    assert annot_all.orig_time == raw.info["meas_date"]
    raw.set_annotations(annot_all)
    dev_head_t = compute_average_dev_head_t(raw, pos)["trans"]

    dev_head_t_ori = np.array(
        [
            [0.9957292, -0.08688804, 0.03120615, 0.00698271],
            [0.09020767, 0.9875856, -0.12859731, -0.0159098],
            [-0.01964518, 0.1308631, 0.99120578, 0.07258289],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )

    assert_allclose(dev_head_t_ori, dev_head_t, rtol=1e-5, atol=0)

    with pytest.raises(ValueError, match="Number of .* must match .*"):
        compute_average_dev_head_t([raw], [pos] * 2)
    # Using two identical ones should be identical ...
    dev_head_t_double = compute_average_dev_head_t([raw] * 2, [pos] * 2)["trans"]
    assert_allclose(dev_head_t, dev_head_t_double)
    # ... unannotated and annotated versions differ ...
    dev_head_t_unannot = compute_average_dev_head_t(raw_unannot, pos)["trans"]
    rot_tol = 1.5e-3
    mov_tol = 1e-3
    assert not np.allclose(
        dev_head_t_unannot[:3, :3],
        dev_head_t[:3, :3],
        atol=rot_tol,
        rtol=0,
    )
    assert not np.allclose(
        dev_head_t_unannot[:3, 3],
        dev_head_t[:3, 3],
        atol=mov_tol,
        rtol=0,
    )
    # ... and Averaging the two is close to (but not identical!) to operating on the two
    # files. Note they shouldn't be identical because there are more time points
    # included in the unannotated version!
    dev_head_t_naive = np.eye(4)
    dev_head_t_naive[:3, :3] = quat_to_rot(
        np.mean(
            rot_to_quat(np.array([dev_head_t[:3, :3], dev_head_t_unannot[:3, :3]])),
            axis=0,
        )
    )
    dev_head_t_naive[:3, 3] = np.mean(
        [dev_head_t[:3, 3], dev_head_t_unannot[:3, 3]], axis=0
    )
    dev_head_t_combo = compute_average_dev_head_t([raw, raw_unannot], [pos] * 2)[
        "trans"
    ]
    unit_kw = dict(distance_units="mm", angle_units="deg")
    deg_annot_combo, mm_annot_combo = _angle_dist_between_rigid(
        dev_head_t,
        dev_head_t_combo,
        **unit_kw,
    )
    deg_unannot_combo, mm_unannot_combo = _angle_dist_between_rigid(
        dev_head_t_unannot,
        dev_head_t_combo,
        **unit_kw,
    )
    deg_annot_unannot, mm_annot_unannot = _angle_dist_between_rigid(
        dev_head_t,
        dev_head_t_unannot,
        **unit_kw,
    )
    deg_combo_naive, mm_combo_naive = _angle_dist_between_rigid(
        dev_head_t_combo,
        dev_head_t_naive,
        **unit_kw,
    )
    # combo<->naive closer than combo<->annotated closer than annotated<->unannotated
    assert 0.05 < deg_combo_naive < deg_annot_combo < deg_annot_unannot < 1.5
    assert 0.1 < mm_combo_naive < mm_annot_combo < mm_annot_unannot < 2
    # combo<->naive closer than combo<->unannotated closer than annotated<->unannotated
    assert 0.05 < deg_combo_naive < deg_unannot_combo < deg_annot_unannot < 1.5
    assert 0.12 < mm_combo_naive < mm_unannot_combo < mm_annot_unannot < 2.0

    # Smoke test skipping time due to previous annotations.
    raw.set_annotations(Annotations([raw.times[0]], 0.1, "bad"))
    annot_dis, _ = annotate_movement(raw, pos, mean_distance_limit=0.02)
    assert annot_dis.duration.size == 1

    # really far should warn
    pos[:, 4] += 5
    with pytest.warns(RuntimeWarning, match="Implausible head position"):
        compute_average_dev_head_t(raw, pos)


@testing.requires_testing_data
@pytest.mark.parametrize("meas_date", (None, "orig"))
def test_muscle_annotation(meas_date, events):
    """Test correct detection muscle artifacts."""
    raw = read_raw_fif(raw_fname, allow_maxshield="yes").load_data()
    if meas_date is None:
        raw.set_meas_date(None)
    raw.notch_filter([50, 110, 150])
    # Check 2 muscle segments are detected
    annot_muscle, scores = annotate_muscle_zscore(raw, ch_type="mag", threshold=10)
    assert annot_muscle.orig_time == raw.info["meas_date"]
    onset = annot_muscle.onset * raw.info["sfreq"]
    if meas_date is not None:
        onset -= raw.first_samp
    onset = onset.astype(int)
    assert_array_equal(scores[onset].astype(int), np.array([23, 10]))
    assert annot_muscle.duration.size == 2
    raw.set_annotations(annot_muscle)


@testing.requires_testing_data
@pytest.mark.parametrize("meas_date", (None, "orig"))
def test_muscle_annotation_without_meeg_data(meas_date):
    """Call annotate_muscle_zscore with data without meg or eeg."""
    raw = read_raw_fif(raw_fname, allow_maxshield="yes")
    if meas_date is None:
        raw.set_meas_date(None)
    raw.crop(0, 0.1).load_data()
    raw.pick("stim")
    with pytest.raises(ValueError, match="No M/EEG channel types found"):
        annotate_muscle_zscore(raw, threshold=10)


@pytest.mark.parametrize("meas_date", (None, "orig"))
@testing.requires_testing_data
def test_annotate_breaks(meas_date):
    """Test annotate_breaks."""
    raw = read_raw_fif(raw_fname, allow_maxshield="yes")
    if meas_date is None:
        raw.set_meas_date(None)

    annots = Annotations(
        onset=[12, 15, 16, 20, 21],
        duration=[1, 1, 1, 2, 0.5],
        description=["test"],
        orig_time=raw.info["meas_date"],
    )

    if raw.info["meas_date"] is None:
        annots.onset -= raw.first_time

    raw.set_annotations(annots)

    min_break_duration = 0.5
    t_start_after_previous = 0.1
    t_stop_before_next = 0.1

    expected_onsets = np.array(
        [
            raw.first_time,
            13 + t_start_after_previous,
            17 + t_start_after_previous,
            22 + t_start_after_previous,
        ]
    )

    if raw.info["meas_date"] is None:
        expected_onsets -= raw.first_time

    expected_durations = np.array(
        [
            12 - raw.first_time - t_stop_before_next,
            15 - 13 - t_start_after_previous - t_stop_before_next,
            20 - 17 - t_start_after_previous - t_stop_before_next,
            raw._last_time - 22 - t_start_after_previous,
        ]
    )

    break_annots = annotate_break(
        raw=raw,
        min_break_duration=min_break_duration,
        t_start_after_previous=t_start_after_previous,
        t_stop_before_next=t_stop_before_next,
    )

    assert break_annots.orig_time == raw.info["meas_date"]
    assert_allclose(break_annots.onset, expected_onsets)
    assert_allclose(break_annots.duration, expected_durations)
    assert all(description == "BAD_break" for description in break_annots.description)

    # try setting the annotations, this should not omit anything
    raw.set_annotations(break_annots)
    current_annotations = raw.annotations
    if raw.info["meas_date"] is None:
        current_annotations.onset -= raw.first_time
    raw.set_annotations(current_annotations + break_annots)

    # reset before next test
    raw.set_annotations(annots)

    # `ignore` parameter should be respected
    raw.annotations.description[0] = "BAD_"
    break_annots = annotate_break(
        raw=raw,
        min_break_duration=min_break_duration,
        t_start_after_previous=t_start_after_previous,
        t_stop_before_next=t_stop_before_next,
    )

    assert_allclose(break_annots.onset, expected_onsets[[True, False, True, True]])
    assert_allclose(
        break_annots.duration,
        [15 - raw.first_time - t_stop_before_next] + list(expected_durations[2:]),
    )

    # try setting the annotations, this should not omit anything
    raw.set_annotations(break_annots)
    current_annotations = raw.annotations
    if raw.info["meas_date"] is None:
        current_annotations.onset -= raw.first_time
    raw.set_annotations(current_annotations + break_annots)

    # Restore annotations for next test
    raw.set_annotations(annots)
    raw.annotations.description[0] = "test"

    # Test with events
    events, _ = events_from_annotations(raw=raw)
    raw.set_annotations(None)

    expected_onsets = np.array(
        [
            raw.first_time,
            12 + t_start_after_previous,
            15 + t_start_after_previous,
            16 + t_start_after_previous,
            20 + t_start_after_previous,
            21 + t_start_after_previous,
        ]
    )

    expected_durations = np.array(
        [
            12 - raw.first_time - t_stop_before_next,
            15 - 12 - t_start_after_previous - t_stop_before_next,
            16 - 15 - t_start_after_previous - t_stop_before_next,
            20 - 16 - t_start_after_previous - t_stop_before_next,
            21 - 20 - t_start_after_previous - t_stop_before_next,
            raw._last_time - 21 - t_start_after_previous,
        ]
    )

    break_annots = annotate_break(
        raw=raw,
        events=events,
        min_break_duration=min_break_duration,
        t_start_after_previous=t_start_after_previous,
        t_stop_before_next=t_stop_before_next,
    )

    if raw.info["meas_date"] is None:
        expected_onsets -= raw.first_time

    assert_allclose(break_annots.onset, expected_onsets)
    assert_allclose(break_annots.duration, expected_durations)

    # try setting the annotations, this should not omit anything
    raw.set_annotations(break_annots)
    current_annotations = raw.annotations
    if raw.info["meas_date"] is None:
        current_annotations.onset -= raw.first_time
    raw.set_annotations(current_annotations + break_annots)

    # reset before next test
    raw.set_annotations(annots)

    # Not finding any break periods
    break_annots = annotate_break(
        raw=raw,
        events=events,
        min_break_duration=1000,
    )

    assert len(break_annots) == 0

    # Implausible parameters (would produce break annot of duration < 0)
    with pytest.raises(ValueError, match="must be greater than 0"):
        annotate_break(
            raw=raw,
            min_break_duration=5,
            t_start_after_previous=5,
            t_stop_before_next=5,
        )

    # Empty events array
    with pytest.raises(ValueError, match="events array must not be empty"):
        annotate_break(raw=raw, events=np.array([]))

    # Invalid `ignore` value
    with pytest.raises(TypeError, match="must be an instance of str"):
        annotate_break(raw=raw, ignore=("foo", 1))

    # No annotations to work with
    raw.set_annotations(None)
    with pytest.raises(ValueError, match="Could not find.*annotations"):
        annotate_break(raw=raw)