"""Tests for scroll training pipeline."""
from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
import pytest

from tools.scroll.trainer import load_scroll_data, train_scroll, _augment_scroll


def _make_synthetic_scroll_trace(mode="target"):
    """Create a synthetic scroll trace."""
    distance = {"target": 1500, "fast": 5000, "precise": 400}[mode]
    direction = "down"
    start = 2000
    target = start + distance

    events = []
    n_events = 20
    for i in range(n_events):
        frac = (i + 1) / n_events
        delta = int(distance / n_events * (1 + 0.2 * np.random.randn()))
        delta = max(20, delta)
        t = int(frac * 800 + np.random.normal(0, 10))
        events.append({"deltaY": delta, "deltaMode": 0, "t": max(0, t)})

    events.sort(key=lambda e: e["t"])
    events[0]["t"] = 0

    return {
        "meta": {
            "mode": mode,
            "start_scrollY": start,
            "target_scrollY": target,
            "end_scrollY": target + 5,
            "distance": distance,
            "direction": direction,
            "duration_ms": events[-1]["t"],
            "viewport_height": 900,
        },
        "events": events,
    }


@pytest.fixture
def synthetic_scroll_file(tmp_path):
    traces_path = tmp_path / "scroll_traces.jsonl"
    lines = []
    for mode in ["target", "fast", "precise"]:
        for _ in range(10):
            lines.append(json.dumps(_make_synthetic_scroll_trace(mode)))
    traces_path.write_text("\n".join(lines), encoding="utf-8")
    return traces_path


class TestLoadScrollData:
    def test_returns_correct_shapes(self, synthetic_scroll_file):
        seq, cond = load_scroll_data(synthetic_scroll_file, seq_len=32)
        assert seq.shape[1] == 32
        assert seq.shape[2] == 2  # (delta_norm, log_dt)
        assert cond.shape[1] == 7
        assert len(seq) > 0


class TestAugment:
    def test_4x_augmentation(self, synthetic_scroll_file):
        seq, cond = load_scroll_data(synthetic_scroll_file, seq_len=32)
        n = len(seq)
        seq_aug, cond_aug = _augment_scroll(seq, cond)
        assert len(seq_aug) == n * 4


class TestTrainScroll:
    def test_produces_model_files(self, synthetic_scroll_file, tmp_path):
        output_dir = tmp_path / "scroll_models"
        train_scroll(
            data_path=synthetic_scroll_file,
            output_dir=output_dir,
            epochs=3,
            batch_size=8,
        )
        assert (output_dir / "scroll_model.pt").exists()
        assert (output_dir / "scroll_config.json").exists()