from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from pi07mini.data import (
    Pi07SteerableDataset,
    build_policy_trajectories,
    load_manifest,
    materialize_subset,
)


def _make_source_project(root: Path) -> Path:
    processed = root / "data" / "processed"
    (processed / "images").mkdir(parents=True)
    (processed / "trajectories").mkdir(parents=True)
    pixels = np.zeros((40, 64, 3), dtype=np.uint8)
    pixels[..., 0] = 80
    pixels[..., 1] = np.arange(64, dtype=np.uint8)[None, :]
    Image.fromarray(pixels).save(processed / "images" / "sample.jpg")
    hx = np.linspace(-7.0, 0.0, 8, dtype=np.float32)
    history = np.stack((hx, 0.01 * hx, np.zeros(8)), axis=-1)
    fx = np.linspace(1.0, 24.0, 16, dtype=np.float32)
    future = np.stack((fx, 0.015 * fx**2, np.zeros(16)), axis=-1)
    np.savez_compressed(
        processed / "trajectories" / "sample.npz",
        history_xyz=history,
        future_xyz=future,
    )
    manifest = processed / "overfit_manifest.jsonl"
    manifest.write_text(
        json.dumps(
            {
                "clip_id": "clip-a",
                "t0_us": 123,
                "reasoning": "Steer left.",
                "image_path": "data/processed/images/sample.jpg",
                "trajectory_path": "data/processed/trajectories/sample.npz",
            }
        )
        + "\n",
        encoding="utf-8",
    )
    return manifest


def test_policy_variants_and_metadata() -> None:
    x = torch.linspace(-7.0, 0.0, 8)
    history = torch.stack((x, torch.zeros_like(x)), -1)
    fx = torch.linspace(1.0, 24.0, 16)
    future = torch.stack((fx, 0.02 * fx.square()), -1)
    policies = build_policy_trajectories(history, future)
    assert [policy.kind for policy in policies] == [
        "expert",
        "fast",
        "slow",
        "understeer",
        "stop",
    ]
    metadata = {policy.kind: policy.metadata.as_dict() for policy in policies}
    assert metadata["expert"] == {
        "speed": "normal",
        "quality": "expert",
        "mistake": "none",
        "control": "human",
    }
    assert metadata["understeer"]["mistake"] == "understeer"
    assert metadata["stop"]["quality"] == "poor"
    assert policies[0].metadata.as_token_metadata() == {
        "speed_steps": 1_000,
        "quality": 5,
        "mistake": False,
        "control_mode": "ee",
        "source_type": "demonstration",
    }
    assert torch.allclose(policies[-1].future_xy[-1], policies[-1].future_xy[4])


def test_dataset_contract_and_roundtrip(tmp_path: Path) -> None:
    source = _make_source_project(tmp_path / "source")
    manifest = materialize_subset(source, tmp_path / "target")
    assert load_manifest(manifest)[0].image_path.is_file()
    dataset = Pi07SteerableDataset(
        manifest,
        image_size=32,
        prompt_dropout=0.0,
        max_rtc_delay_steps=12,
    )
    assert len(dataset) == 5
    assert {policy.kind for policy in dataset.policies} == {
        "expert",
        "fast",
        "slow",
        "understeer",
        "stop",
    }
    item = dataset[0]
    assert item["pixel_values"].shape == (3, 32, 32)
    assert item["history_pixel_values"].shape == (6, 3, 32, 32)
    assert item["history_states"].shape == (6, 2)
    assert item["current_state"].shape == (2,)
    assert item["actions"].shape == (16, 2)
    assert item["subgoal_heatmap"].shape == (1, 32, 32)
    assert item["subgoal_pixel_values"].shape == (3, 32, 32)
    assert item["ambiguous_history_pair"].shape == (2, 6, 2)
    assert 0 <= item["rtc_delay_steps"] <= 12
    assert item["rtc_target_actions"].shape == (16, 2)
    assert item["rtc_valid_mask"].shape == (16,)
    assert item["reasoning"] == "Steer left."
    assert item["prompt_component_mask"].all()
    assert item["token_metadata"]["quality"] == 5
    assert item["control_mode"] == "ee"
    decoded = dataset.action_normalizer.decode(item["actions"], item["history_xy"])
    torch.testing.assert_close(decoded, item["future_xy"], atol=1e-5, rtol=1e-5)


def test_dataset_is_deterministic_per_index(tmp_path: Path) -> None:
    source = _make_source_project(tmp_path / "source")
    dataset = Pi07SteerableDataset(
        source,
        image_size=24,
        prompt_dropout=0.5,
        seed=11,
    )
    first = dataset[3]
    second = dataset[3]
    assert first["prompt_components"] == second["prompt_components"]
    assert first["rtc_delay_steps"] == second["rtc_delay_steps"]
    torch.testing.assert_close(first["rtc_target_actions"], second["rtc_target_actions"])