from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from pi05mini.data import Pi05OverfitDataset, load_manifest, materialize_subset
from pi05mini.tokenizer import (
    CoordinateTokenizer,
    MiniTokenizer,
    SubtaskVocabulary,
    TaskVocabulary,
)


def _make_source_project(root: Path) -> Path:
    processed = root / "data" / "processed"
    image_dir = processed / "images"
    trajectory_dir = processed / "trajectories"
    image_dir.mkdir(parents=True)
    trajectory_dir.mkdir(parents=True)

    image_path = image_dir / "sample.jpg"
    pixels = np.zeros((48, 80, 3), dtype=np.uint8)
    pixels[..., 0] = 64
    pixels[..., 1] = np.arange(80, dtype=np.uint8)[None, :]
    Image.fromarray(pixels).save(image_path)

    history_x = np.linspace(-7.0, 0.0, 8, dtype=np.float32)
    history = np.stack((history_x, 0.02 * history_x, np.zeros(8)), axis=-1)
    future_x = np.linspace(1.0, 32.0, 16, dtype=np.float32)
    future = np.stack((future_x, 0.002 * future_x**2, np.zeros(16)), axis=-1)
    trajectory_path = trajectory_dir / "sample.npz"
    np.savez_compressed(trajectory_path, history_xyz=history, future_xyz=future)

    manifest = processed / "overfit_manifest.jsonl"
    row = {
        "clip_id": "clip-a",
        "t0_us": 1234,
        "task_instruction": "Follow the road through the construction zone.",
        "reasoning": "Steer left through the construction lane.",
        "image_path": "data/processed/images/sample.jpg",
        "trajectory_path": "data/processed/trajectories/sample.npz",
    }
    manifest.write_text(json.dumps(row) + "\n", encoding="utf-8")
    return manifest


def test_coordinate_tokens_are_literal_bytes_and_round_trip() -> None:
    tokenizer = CoordinateTokenizer(x_range=(-10.0, 10.0), y_range=(-5.0, 5.0))
    coordinates = torch.tensor([[-10.0, -5.0], [0.0, 0.0], [10.0, 5.0]])
    tokens = tokenizer.encode(coordinates)

    assert tokens.dtype == torch.long
    assert tokens.tolist()[0] == [0, 0]
    assert tokens.tolist()[-1] == [255, 255]
    assert 0 <= int(tokens.min()) <= int(tokens.max()) <= 255
    decoded = tokenizer.decode(tokens)
    assert torch.allclose(decoded, coordinates, atol=10.1 / 255.0)


def test_task_subtask_vocabulary_and_hierarchical_prefix_contract() -> None:
    tokenizer = MiniTokenizer()
    history = torch.zeros(8, 2)
    prefix = tokenizer.build_prefix(history, "Follow the road.")
    other_task = tokenizer.build_prefix(history, "Clean the kitchen.")
    subtask = tokenizer.encode_subtask("Steer left around the cones")
    low_level = tokenizer.build_low_level_prefix(history, subtask)

    assert prefix.shape == (22,)
    assert prefix[:3].tolist() == [tokenizer.BOS, tokenizer.IMAGE, tokenizer.STATE]
    assert prefix[-3].item() == tokenizer.TASK
    assert tokenizer.tasks.decode(prefix[-2]) == "follow_road"
    assert prefix[-1].item() == tokenizer.SUBTASK_QUERY
    assert prefix[3:19].min() >= 0 and prefix[3:19].max() <= 255
    assert other_task[-2] != prefix[-2]
    assert low_level.shape == (21,)
    torch.testing.assert_close(low_level[:19], prefix[:19])
    assert low_level[-2].item() == tokenizer.SUBTASK_COMMAND
    assert low_level[-1].item() == subtask
    assert tokenizer.TASK not in low_level
    assert tokenizer.subtasks.decode(subtask) == "steer_left"
    assert TaskVocabulary().classify("Clean all dishes in the kitchen") == "clean_kitchen"
    assert SubtaskVocabulary().classify("Slow down for a pedestrian") == "slow_down"


def test_materialize_and_dataset_are_self_contained_and_deterministic(tmp_path: Path) -> None:
    source_manifest = _make_source_project(tmp_path / "source")
    copied_manifest = materialize_subset(source_manifest, tmp_path / "target", max_samples=1)
    records = load_manifest(copied_manifest)

    assert len(records) == 1
    assert records[0].image_path.is_file()
    assert records[0].trajectory_path.is_file()

    dataset = Pi05OverfitDataset(
        copied_manifest,
        image_size=(32, 48),
        augmentation_factor=3,
        seed=19,
    )
    same_dataset = Pi05OverfitDataset(
        copied_manifest,
        image_size=(32, 48),
        augmentation_factor=3,
        seed=19,
        action_statistics=dataset.action_normalizer.as_dict(),
    )

    assert len(dataset) == 3
    original = dataset[0]
    augmented = dataset[1]
    repeated = same_dataset[1]
    assert original["image"].shape == (3, 32, 48)
    assert original["prefix_tokens"].shape == (22,)
    assert original["low_level_prefix_tokens"].shape == (21,)
    assert original["task_instruction"].startswith("Follow the road")
    assert original["actions"].shape == (16, 2)
    assert original["subtask_target"].shape == ()
    assert original["subtask_label"] == "steer_left"
    assert original["reasoning"].startswith("Steer left")
    assert original["sample_id"].endswith(":aug0")
    assert augmented["sample_id"].endswith(":aug1")
    assert not torch.equal(original["future_xy"], augmented["future_xy"])
    assert torch.equal(augmented["future_xy"], repeated["future_xy"])
    assert torch.equal(augmented["image"], repeated["image"])

    reconstructed = dataset.action_normalizer.decode(
        original["actions"], original["history_xy"]
    )
    assert torch.allclose(reconstructed, original["future_xy"], atol=1e-5)

    statistics_path = dataset.save_action_statistics(tmp_path / "action_stats.json")
    statistics = json.loads(statistics_path.read_text(encoding="utf-8"))
    assert len(statistics["mean"]) == 2
    assert len(statistics["scale"]) == 2
    assert min(statistics["scale"]) >= 0.05