from __future__ import annotations

import csv
from pathlib import Path

import numpy as np
import onnxruntime as ort
import pytest
import torch

from mouse_control.train import TrajectoryDataset, _load_csv, train


def _write_csv(path: Path, rows: list[list[tuple[int, int]]]) -> None:
    """Write rows in the project's quirky single-quoted-pair format."""
    with open(path, "w", newline="") as f:
        writer = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
        for row in rows:
            # First column is the target (last keypoint), then 10 keypoints.
            target_x, target_y = row[-1]
            cells = [f"{target_x},{target_y}"] + [f"{x},{y}" for x, y in row]
            writer.writerow(cells)


@pytest.fixture
def synthetic_csv(tmp_path: Path) -> Path:
    """Tiny linear trajectories from origin -> target."""
    path = tmp_path / "synth.csv"
    rows = []
    rng = np.random.default_rng(0)
    for _ in range(20):
        tx, ty = int(rng.integers(-150, 151)), int(rng.integers(-150, 151))
        keypoints = [
            (int(round(tx * i / 9)), int(round(ty * i / 9))) for i in range(10)
        ]
        rows.append(keypoints)
    _write_csv(path, rows)
    return path


def test_load_csv_shapes(synthetic_csv: Path) -> None:
    inputs, labels = _load_csv(synthetic_csv)
    assert inputs.shape == (20, 1, 2)
    assert labels.shape == (20, 10, 2)
    assert inputs.dtype == torch.float32
    assert labels.dtype == torch.float32


def test_load_csv_round_trip(synthetic_csv: Path) -> None:
    """The target column (input) must equal the last keypoint of the labels."""
    inputs, labels = _load_csv(synthetic_csv)
    last_keypoint = labels[:, -1, :]
    np.testing.assert_array_equal(inputs.squeeze(1).numpy(), last_keypoint.numpy())


def test_dataset_indexing(synthetic_csv: Path) -> None:
    inputs, labels = _load_csv(synthetic_csv)
    ds = TrajectoryDataset(inputs, labels)
    assert len(ds) == 20
    x, y = ds[0]
    assert x.shape == (1, 2)
    assert y.shape == (10, 2)


def test_load_csv_missing_file_raises(tmp_path: Path) -> None:
    with pytest.raises(FileNotFoundError):
        _load_csv(tmp_path / "missing.csv")


def test_train_end_to_end_produces_single_file_onnx(
    synthetic_csv: Path, tmp_path: Path,
) -> None:
    """A quick training run must export a valid, self-contained ONNX model."""
    output = tmp_path / "trained.onnx"
    train(
        train_csv=synthetic_csv,
        test_csv=synthetic_csv,
        output=output,
        epochs=3,
        batch_size=8,
    )
    assert output.exists()
    assert output.stat().st_size > 1000
    # external_data=False -> no .data sidecar
    assert not output.with_suffix(output.suffix + ".data").exists()

    session = ort.InferenceSession(str(output))
    inp = np.array([[[100.0, 50.0]]], dtype=np.float32)
    out = session.run(None, {"input": inp})[0]
    assert out.shape == (1, 10, 2)


def test_train_with_real_data(train_csv: Path, test_csv: Path, tmp_path: Path) -> None:
    """Smoke test against the project's actual data."""
    output = tmp_path / "real.onnx"
    train(
        train_csv=train_csv,
        test_csv=test_csv,
        output=output,
        epochs=2,
        batch_size=32,
    )
    assert output.exists()