"""Tests for :mod:`neuropose.analyzer.features`."""

from __future__ import annotations

import math

import numpy as np
import pytest

from neuropose.analyzer.features import (
    AlignmentDiagnostics,
    FeatureStatistics,
    extract_feature_statistics,
    extract_joint_angles,
    find_peaks,
    normalize_pose_sequence,
    pad_sequences,
    predictions_to_numpy,
    procrustes_align,
)
from neuropose.io import VideoPredictions

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


def _make_predictions(num_frames: int, num_persons: int = 1) -> VideoPredictions:
    """Build a minimal VideoPredictions object for tests."""
    frames = {}
    for i in range(num_frames):
        frames[f"frame_{i:06d}"] = {
            "boxes": [[0.0, 0.0, 1.0, 1.0, 0.9]] * num_persons,
            "poses3d": [[[float(i), float(i) * 2, float(i) * 3], [0.0, 0.0, 0.0]]] * num_persons,
            "poses2d": [[[0.0, 0.0], [1.0, 1.0]]] * num_persons,
        }
    return VideoPredictions.model_validate(
        {
            "metadata": {
                "frame_count": num_frames,
                "fps": 30.0,
                "width": 640,
                "height": 480,
            },
            "frames": frames,
        }
    )


# ---------------------------------------------------------------------------
# predictions_to_numpy
# ---------------------------------------------------------------------------


class TestPredictionsToNumpy:
    def test_single_person_shape(self) -> None:
        predictions = _make_predictions(num_frames=4)
        arr = predictions_to_numpy(predictions)
        assert arr.shape == (4, 2, 3)
        assert arr.dtype == np.float64

    def test_values_preserved(self) -> None:
        predictions = _make_predictions(num_frames=3)
        arr = predictions_to_numpy(predictions)
        # Frame i has joint 0 at (i, 2i, 3i) per _make_predictions.
        for i in range(3):
            np.testing.assert_allclose(arr[i, 0], [i, 2 * i, 3 * i])
            np.testing.assert_allclose(arr[i, 1], [0, 0, 0])

    def test_person_index_out_of_range(self) -> None:
        predictions = _make_predictions(num_frames=2, num_persons=1)
        with pytest.raises(ValueError, match="out of range"):
            predictions_to_numpy(predictions, person_index=1)

    def test_multi_person_with_explicit_index(self) -> None:
        predictions = _make_predictions(num_frames=2, num_persons=2)
        arr = predictions_to_numpy(predictions, person_index=1)
        assert arr.shape == (2, 2, 3)

    def test_empty_predictions_raises(self) -> None:
        predictions = _make_predictions(num_frames=0)
        with pytest.raises(ValueError, match="zero frames"):
            predictions_to_numpy(predictions)


# ---------------------------------------------------------------------------
# normalize_pose_sequence
# ---------------------------------------------------------------------------


class TestNormalize:
    def test_uniform_preserves_ratio(self) -> None:
        # (frames, joints, 3) — one joint per frame, two frames.
        seq = np.array(
            [
                [[0.0, 0.0, 0.0]],
                [[3.0, 6.0, 9.0]],
            ]
        )
        # Ranges: x=3, y=6, z=9. Uniform scale = 9. All values / 9.
        result = normalize_pose_sequence(seq, axis_wise=False)
        np.testing.assert_allclose(result, seq / 9.0)

    def test_axis_wise_each_axis_to_unit_range(self) -> None:
        seq = np.array(
            [
                [[0.0, 0.0, 0.0]],
                [[3.0, 6.0, 9.0]],
            ]
        )
        result = normalize_pose_sequence(seq, axis_wise=True)
        # Per-axis normalization → each axis's max becomes 1.
        np.testing.assert_allclose(result[0, 0], [0.0, 0.0, 0.0])
        np.testing.assert_allclose(result[1, 0], [1.0, 1.0, 1.0])

    def test_does_not_mutate_input(self) -> None:
        seq = np.array([[[0.0, 0.0, 0.0]], [[1.0, 2.0, 3.0]]])
        before = seq.copy()
        normalize_pose_sequence(seq)
        np.testing.assert_array_equal(seq, before)

    def test_degenerate_sequence_rejected(self) -> None:
        seq = np.zeros((3, 2, 3))
        with pytest.raises(ValueError, match="degenerate"):
            normalize_pose_sequence(seq)

    def test_bad_shape_rejected(self) -> None:
        seq = np.zeros((3, 2))  # Missing the xyz axis.
        with pytest.raises(ValueError, match="expected"):
            normalize_pose_sequence(seq)

    def test_axis_wise_with_zero_axis_keeps_it_zero(self) -> None:
        # Sequence where the Z axis never moves — axis_wise should not
        # divide by zero; the Z column should remain at 0.
        seq = np.array(
            [
                [[0.0, 0.0, 5.0]],
                [[4.0, 8.0, 5.0]],
            ]
        )
        result = normalize_pose_sequence(seq, axis_wise=True)
        np.testing.assert_allclose(result[:, 0, 2], [0.0, 0.0])


# ---------------------------------------------------------------------------
# pad_sequences
# ---------------------------------------------------------------------------


class TestPadSequences:
    def test_pads_to_max_when_target_length_none(self) -> None:
        a = np.zeros((3, 2, 3))
        b = np.zeros((5, 2, 3))
        padded = pad_sequences([a, b])
        assert all(seq.shape[0] == 5 for seq in padded)

    def test_pads_to_explicit_target_length(self) -> None:
        a = np.zeros((3, 2, 3))
        padded = pad_sequences([a], target_length=10)
        assert padded[0].shape == (10, 2, 3)

    def test_edge_padding_repeats_last_frame(self) -> None:
        a = np.array([[[1.0, 2.0, 3.0]]])  # shape (1, 1, 3)
        padded = pad_sequences([a], target_length=4)
        # All 4 frames should equal the original single frame.
        for i in range(4):
            np.testing.assert_allclose(padded[0][i, 0], [1.0, 2.0, 3.0])

    def test_truncates_longer_than_target(self) -> None:
        a = np.zeros((10, 2, 3))
        padded = pad_sequences([a], target_length=4)
        assert padded[0].shape == (4, 2, 3)

    def test_does_not_mutate_input(self) -> None:
        a = np.zeros((3, 2, 3))
        pad_sequences([a], target_length=5)
        assert a.shape == (3, 2, 3)

    def test_mismatched_trailing_shape_rejected(self) -> None:
        a = np.zeros((3, 2, 3))
        b = np.zeros((3, 4, 3))  # Different joint count.
        with pytest.raises(ValueError, match="trailing shape"):
            pad_sequences([a, b])

    def test_empty_input_with_target(self) -> None:
        assert pad_sequences([], target_length=5) == []

    def test_empty_input_without_target_raises(self) -> None:
        with pytest.raises(ValueError, match="empty"):
            pad_sequences([])


# ---------------------------------------------------------------------------
# extract_joint_angles
# ---------------------------------------------------------------------------


class TestExtractJointAngles:
    def test_right_angle(self) -> None:
        # Three joints forming a right angle at joint 1.
        # joint 0 at (1, 0, 0), joint 1 at origin, joint 2 at (0, 1, 0).
        sequence = np.array(
            [
                [[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
            ]
        )
        angles = extract_joint_angles(sequence, triplets=[(0, 1, 2)])
        assert angles.shape == (1, 1)
        assert angles[0, 0] == pytest.approx(math.pi / 2)

    def test_collinear_gives_pi(self) -> None:
        sequence = np.array(
            [
                [[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [-1.0, 0.0, 0.0]],
            ]
        )
        angles = extract_joint_angles(sequence, triplets=[(0, 1, 2)])
        assert angles[0, 0] == pytest.approx(math.pi)

    def test_multiple_triplets(self) -> None:
        sequence = np.array(
            [
                [
                    [1.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0],
                    [0.0, 1.0, 0.0],
                    [0.0, 0.0, 1.0],
                ],
            ]
        )
        # Right angle at 1 (first triplet) and right angle at 1 again
        # using joint 3 as the other arm — still 90°.
        angles = extract_joint_angles(sequence, triplets=[(0, 1, 2), (0, 1, 3)])
        assert angles.shape == (1, 2)
        assert angles[0, 0] == pytest.approx(math.pi / 2)
        assert angles[0, 1] == pytest.approx(math.pi / 2)

    def test_zero_length_vector_yields_nan(self) -> None:
        # Joints 0 and 1 coincide → v1 is the zero vector → NaN angle.
        sequence = np.array(
            [
                [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
            ]
        )
        angles = extract_joint_angles(sequence, triplets=[(0, 1, 2)])
        assert math.isnan(angles[0, 0])

    def test_out_of_range_index_rejected(self) -> None:
        sequence = np.zeros((1, 3, 3))
        with pytest.raises(ValueError, match="out of range"):
            extract_joint_angles(sequence, triplets=[(0, 1, 10)])


# ---------------------------------------------------------------------------
# extract_feature_statistics
# ---------------------------------------------------------------------------


class TestExtractFeatureStatistics:
    def test_basic_stats(self) -> None:
        values = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        stats = extract_feature_statistics(values)
        assert isinstance(stats, FeatureStatistics)
        assert stats.mean == pytest.approx(3.0)
        assert stats.min == pytest.approx(1.0)
        assert stats.max == pytest.approx(5.0)
        assert stats.range == pytest.approx(4.0)
        assert stats.std == pytest.approx(np.std(values))

    def test_rejects_2d(self) -> None:
        values = np.zeros((3, 3))
        with pytest.raises(ValueError, match="1D"):
            extract_feature_statistics(values)

    def test_rejects_empty(self) -> None:
        with pytest.raises(ValueError, match="empty"):
            extract_feature_statistics(np.array([]))


# ---------------------------------------------------------------------------
# find_peaks
# ---------------------------------------------------------------------------


class TestFindPeaks:
    def test_sine_wave_peaks(self) -> None:
        # A sine wave over two full cycles has two peaks at quarter
        # cycles — roughly at t=pi/2 and t=5pi/2 given 4pi duration.
        t = np.linspace(0, 4 * np.pi, 401)
        values = np.sin(t)
        indices = find_peaks(values)
        assert indices.ndim == 1
        assert len(indices) == 2

    def test_flat_signal_has_no_peaks(self) -> None:
        indices = find_peaks(np.zeros(100))
        assert indices.size == 0

    def test_rejects_2d_input(self) -> None:
        with pytest.raises(ValueError, match="1D"):
            find_peaks(np.zeros((5, 5)))


# ---------------------------------------------------------------------------
# procrustes_align
# ---------------------------------------------------------------------------


def _rotation_matrix_z(angle_rad: float) -> np.ndarray:
    """Rotation matrix about the Z axis."""
    c, s = np.cos(angle_rad), np.sin(angle_rad)
    return np.array(
        [
            [c, -s, 0.0],
            [s, c, 0.0],
            [0.0, 0.0, 1.0],
        ]
    )


def _skeleton(num_joints: int = 8, seed: int = 0) -> np.ndarray:
    """A deterministic, non-degenerate single-frame skeleton."""
    rng = np.random.default_rng(seed)
    return rng.standard_normal((num_joints, 3))


class TestProcrustesAlignPerSequence:
    def test_identical_sequences_yield_identity_transform(self) -> None:
        sequence = _skeleton()[np.newaxis, :, :].repeat(3, axis=0)  # (3, 8, 3)
        aligned, target, diag = procrustes_align(sequence, sequence, mode="per_sequence")
        np.testing.assert_allclose(aligned, sequence, atol=1e-10)
        np.testing.assert_array_equal(target, sequence)
        assert diag.mode == "per_sequence"
        assert diag.rotation_deg == pytest.approx(0.0, abs=1e-6)
        assert diag.translation == pytest.approx(0.0, abs=1e-9)
        assert diag.scale == pytest.approx(1.0)

    def test_recovers_known_rotation(self) -> None:
        # Build a reference sequence; construct the source by rotating it
        # about Z, then verify alignment returns the reference up to
        # floating-point error.
        rotation = _rotation_matrix_z(np.deg2rad(37.0))
        reference = _skeleton(num_joints=10)[np.newaxis, :, :].repeat(4, axis=0)
        source = reference @ rotation.T
        aligned, _, diag = procrustes_align(source, reference, mode="per_sequence")
        np.testing.assert_allclose(aligned, reference, atol=1e-8)
        # The recovered rotation's magnitude should be the original 37°.
        assert diag.rotation_deg == pytest.approx(37.0, abs=1e-4)

    def test_recovers_known_translation(self) -> None:
        reference = _skeleton()[np.newaxis, :, :].repeat(5, axis=0)
        translation = np.array([10.0, -4.5, 2.25])
        source = reference + translation
        aligned, _, diag = procrustes_align(source, reference, mode="per_sequence")
        np.testing.assert_allclose(aligned, reference, atol=1e-9)
        # rotation_deg may be numerically tiny but not exactly 0.
        assert diag.rotation_deg == pytest.approx(0.0, abs=1e-4)
        assert diag.translation == pytest.approx(np.linalg.norm(translation), rel=1e-6)

    def test_recovers_combined_rotation_and_translation(self) -> None:
        rotation = _rotation_matrix_z(np.deg2rad(-12.0))
        translation = np.array([1.0, 2.0, 3.0])
        reference = _skeleton(num_joints=6)[np.newaxis, :, :].repeat(3, axis=0)
        source = reference @ rotation.T + translation
        aligned, _, diag = procrustes_align(source, reference, mode="per_sequence")
        np.testing.assert_allclose(aligned, reference, atol=1e-8)
        assert diag.rotation_deg == pytest.approx(12.0, abs=1e-4)
        assert diag.translation == pytest.approx(np.linalg.norm(translation), rel=1e-4)

    def test_scale_flag_recovers_known_scale(self) -> None:
        reference = _skeleton()[np.newaxis, :, :].repeat(2, axis=0)
        source = reference * 0.5
        aligned, _, diag = procrustes_align(source, reference, mode="per_sequence", scale=True)
        np.testing.assert_allclose(aligned, reference, atol=1e-8)
        assert diag.scale == pytest.approx(2.0, rel=1e-6)

    def test_scale_flag_off_leaves_scale_at_one(self) -> None:
        reference = _skeleton()[np.newaxis, :, :].repeat(2, axis=0)
        source = reference * 0.5
        _, _, diag = procrustes_align(source, reference, mode="per_sequence", scale=False)
        assert diag.scale == pytest.approx(1.0)

    def test_rejects_mismatched_shapes(self) -> None:
        a = np.zeros((4, 8, 3))
        b = np.zeros((4, 7, 3))
        with pytest.raises(ValueError, match="same shape"):
            procrustes_align(a, b)

    def test_rejects_wrong_trailing_axis(self) -> None:
        a = np.zeros((4, 8, 2))
        b = np.zeros((4, 8, 2))
        with pytest.raises(ValueError, match="joints, 3"):
            procrustes_align(a, b)

    def test_rejects_unknown_mode(self) -> None:
        a = np.zeros((2, 4, 3))
        with pytest.raises(ValueError, match="unknown mode"):
            procrustes_align(a, a, mode="nope")  # type: ignore[arg-type]

    def test_does_not_mutate_inputs(self) -> None:
        source = _skeleton()[np.newaxis, :, :].repeat(3, axis=0).copy()
        target = (source @ _rotation_matrix_z(np.deg2rad(10.0)).T).copy()
        source_before = source.copy()
        target_before = target.copy()
        procrustes_align(source, target, mode="per_sequence")
        np.testing.assert_array_equal(source, source_before)
        np.testing.assert_array_equal(target, target_before)

    def test_returns_alignment_diagnostics_dataclass(self) -> None:
        a = _skeleton()[np.newaxis, :, :].repeat(2, axis=0)
        _, _, diag = procrustes_align(a, a)
        assert isinstance(diag, AlignmentDiagnostics)


class TestProcrustesAlignPerFrame:
    def test_per_frame_recovers_varying_rotations(self) -> None:
        # Each frame is rotated by a different angle; per_frame alignment
        # should recover each frame independently.
        num_frames = 4
        reference_frame = _skeleton(num_joints=6)
        angles = np.deg2rad([5.0, -10.0, 20.0, 45.0])
        reference = np.stack([reference_frame for _ in range(num_frames)], axis=0)
        source = np.stack([reference_frame @ _rotation_matrix_z(a).T for a in angles], axis=0)
        aligned, _, diag = procrustes_align(source, reference, mode="per_frame")
        np.testing.assert_allclose(aligned, reference, atol=1e-8)
        assert diag.mode == "per_frame"
        # The max rotation across frames should be 45°.
        assert diag.rotation_deg_max == pytest.approx(45.0, abs=1e-4)
        # The mean rotation across frames should be 20°.
        assert diag.rotation_deg == pytest.approx(20.0, abs=1e-4)

    def test_per_frame_with_identical_sequences_yields_zero(self) -> None:
        sequence = _skeleton(num_joints=5)[np.newaxis, :, :].repeat(3, axis=0)
        aligned, _, diag = procrustes_align(sequence, sequence, mode="per_frame")
        np.testing.assert_allclose(aligned, sequence, atol=1e-10)
        # Per-frame SVD on a symmetric covariance is numerically ambiguous
        # in axis selection, so the fitted rotation can be a few micro-
        # degrees off zero; the residual positions are still exact.
        assert diag.rotation_deg == pytest.approx(0.0, abs=1e-3)
        assert diag.rotation_deg_max == pytest.approx(0.0, abs=1e-3)
        assert diag.translation == pytest.approx(0.0, abs=1e-9)


# ---------------------------------------------------------------------------
# DTW with align= (integration)
# ---------------------------------------------------------------------------


class TestDtwAlignIntegration:
    """Smoke tests: align= routes through procrustes_align correctly.

    Depth tests of the DTW path itself live in test_analyzer_dtw.
    """

    def test_dtw_all_with_alignment_cancels_rigid_offset(self) -> None:
        pytest.importorskip("fastdtw")
        from neuropose.analyzer.dtw import dtw_all

        rotation = _rotation_matrix_z(np.deg2rad(30.0))
        translation = np.array([5.0, -2.0, 1.0])
        reference = _skeleton(num_joints=6)[np.newaxis, :, :].repeat(4, axis=0)
        source = reference @ rotation.T + translation
        baseline = dtw_all(source, reference, align="none")
        aligned_result = dtw_all(source, reference, align="procrustes_per_sequence")
        assert baseline.distance > 0.0
        assert aligned_result.distance == pytest.approx(0.0, abs=1e-6)

    def test_dtw_align_rejects_mismatched_frame_counts(self) -> None:
        pytest.importorskip("fastdtw")
        from neuropose.analyzer.dtw import dtw_all

        a = np.zeros((5, 3, 3))
        b = np.zeros((6, 3, 3))
        with pytest.raises(ValueError, match="matching frame counts"):
            dtw_all(a, b, align="procrustes_per_sequence")