"""Tests for Flow Matching trajectory generator."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest
import torch

from ai_mouse.generator import generate
from tools.models import TrajectoryFlowModel


@pytest.fixture
def model_dir(tmp_path):
    """Create temp dir with Flow model artifacts."""
    model = TrajectoryFlowModel(seq_len=64)
    torch.save(model.state_dict(), tmp_path / "flow_model.pt")

    click_dist = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
    (tmp_path / "click_dist.json").write_text(json.dumps(click_dist))

    duration_dist = {
        "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
        "params": [
            {"mu_log": 5.5, "sigma_log": 0.3},
            {"mu_log": 5.8, "sigma_log": 0.3},
            {"mu_log": 6.0, "sigma_log": 0.3},
            {"mu_log": 6.2, "sigma_log": 0.3},
            {"mu_log": 6.5, "sigma_log": 0.3},
            {"mu_log": 6.7, "sigma_log": 0.3},
            {"mu_log": 6.9, "sigma_log": 0.3},
            {"mu_log": 7.0, "sigma_log": 0.3},
        ],
    }
    (tmp_path / "duration_dist.json").write_text(json.dumps(duration_dist))

    train_config = {
        "seq_len": 64,
        "d_model": 128,
        "nhead": 4,
        "num_layers": 4,
        "dim_feedforward": 256,
        "cond_dim": 3,
    }
    (tmp_path / "train_config.json").write_text(json.dumps(train_config))

    return tmp_path


class TestGenerate:
    def test_returns_list_of_tuples(self, model_dir):
        result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
        assert isinstance(result, list)
        assert all(isinstance(p, tuple) and len(p) == 3 for p in result)
        # All elements are ints
        for p in result:
            assert all(isinstance(v, int) for v in p)

    def test_timestamps_monotonically_increasing(self, model_dir):
        result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
        times = [p[2] for p in result]
        for i in range(1, len(times)):
            assert times[i] >= times[i - 1]

    def test_starts_near_start(self, model_dir):
        start = (100, 200)
        result = generate(start=start, end=(500, 400), model_dir=str(model_dir))
        first = result[0]
        assert abs(first[0] - start[0]) < 30
        assert abs(first[1] - start[1]) < 30

    def test_ends_near_end(self, model_dir):
        end = (500, 400)
        result = generate(start=(100, 200), end=end, model_dir=str(model_dir))
        # Last two are click events; the one before is last movement point
        last_move = result[-3]
        assert abs(last_move[0] - end[0]) < 30
        assert abs(last_move[1] - end[1]) < 30

    def test_last_two_are_click_events(self, model_dir):
        result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
        down = result[-2]
        up = result[-1]
        # Same x, y for click down and up
        assert down[0] == up[0]
        assert down[1] == up[1]
        # Up timestamp > down timestamp
        assert up[2] > down[2]
        # Click duration within bounds
        assert 20 <= up[2] - down[2] <= 300

    def test_different_z_gives_different_paths(self, model_dir):
        r1 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
        r2 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
        points1 = [(p[0], p[1]) for p in r1[:-2]]
        points2 = [(p[0], p[1]) for p in r2[:-2]]
        assert points1 != points2

    def test_n_points_parameter(self, model_dir):
        result = generate(
            start=(100, 200), end=(500, 400), n_points=32, model_dir=str(model_dir)
        )
        # 32 move points + 2 click events = 34
        assert len(result) == 34


class TestPostProcessing:
    def test_dt_diversity_preserved(self, model_dir):
        """After removing speed_profile + median clip, multiple generations
        should differ in their Δt sequences (not all identical)."""
        results = [generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
                   for _ in range(5)]
        # Extract Δt sequences (only move events, not click events)
        dts = []
        for r in results:
            moves = r[:-2]
            dt_seq = [moves[i+1][2] - moves[i][2] for i in range(len(moves)-1)]
            dts.append(dt_seq)
        # At least 2 of the 5 sequences should differ at any given index
        for i in range(min(len(d) for d in dts)):
            values = {tuple([d[i]]) for d in dts}
            if len(values) > 1:
                return  # at least one position has variation — pass
        pytest.fail("All 5 Δt sequences are identical at every position — diversity collapsed")


class TestGaussianSmooth:
    def test_endpoints_preserved(self):
        from ai_mouse.generator import _gaussian_smooth
        x = np.array([1.0, 5.0, 3.0, 7.0, 2.0], dtype=np.float64)
        smoothed = _gaussian_smooth(x, sigma=1.0)
        assert smoothed[0] == 1.0
        assert smoothed[-1] == 2.0

    def test_smooths_high_frequency(self):
        """A high-frequency square wave should have reduced amplitude after smoothing."""
        from ai_mouse.generator import _gaussian_smooth
        x = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=np.float64)
        smoothed = _gaussian_smooth(x, sigma=1.0)
        # Interior amplitude should be reduced
        interior_orig = x[2:-2]
        interior_smooth = smoothed[2:-2]
        assert interior_smooth.std() < interior_orig.std()

    def test_constant_signal_unchanged(self):
        from ai_mouse.generator import _gaussian_smooth
        x = np.full(20, 0.5, dtype=np.float64)
        smoothed = _gaussian_smooth(x, sigma=1.0)
        np.testing.assert_allclose(smoothed, x, rtol=1e-6)

    def test_short_array_returns_unchanged(self):
        """Arrays shorter than the kernel are returned unchanged."""
        from ai_mouse.generator import _gaussian_smooth
        x = np.array([1.0, 2.0, 3.0], dtype=np.float64)
        smoothed = _gaussian_smooth(x, sigma=1.0)
        np.testing.assert_allclose(smoothed, x, rtol=1e-6)