from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from pi06mini.data import Pi06ExperienceDataset, 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_experience_dataset_contract_and_intervention_label(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 = Pi06ExperienceDataset(manifest, image_size=32)
    assert len(dataset) == 4
    assert {item.source_type for item in dataset.episodes} == {
        "demo_success",
        "autonomous_understeer",
        "autonomous_stop",
        "intervention_correction",
    }
    assert float(dataset.normalized_returns.min()) == -1.0
    assert float(dataset.normalized_returns.max()) == 0.0

    items = [dataset[index] for index in range(len(dataset))]
    intervention = next(item for item in items if item["is_intervention"])
    assert bool(intervention["positive_condition"])
    assert torch.all(intervention["positive_steps"][intervention["intervention_mask"]])
    assert intervention["image"].shape == (3, 32, 32)
    assert intervention["state"].shape == (16,)
    assert intervention["actions"].shape == (16, 2)
    assert intervention["returns"].shape == (16,)
    assert intervention["return_target"].shape == ()
    assert intervention["value_target"].shape == ()
    assert intervention["value_histories"].shape == (17, 8, 2)
    assert intervention["value_targets"].shape == (17,)
    assert intervention["advantages"].shape == (16,)
    assert intervention["advantage_label"] == "positive"
    assert intervention["reasoning"] == "Steer left."
    decoded = dataset.action_normalizer.decode(
        intervention["actions"], intervention["history_xy"]
    )
    torch.testing.assert_close(decoded, intervention["future_xy"], atol=1e-5, rtol=1e-5)


def test_learned_value_predictions_drive_advantage_relabeling(tmp_path: Path) -> None:
    manifest = _make_source_project(tmp_path / "source")
    dataset = Pi06ExperienceDataset(manifest, image_size=32)
    before = dataset.advantages.clone()
    predicted = torch.zeros(len(dataset.episodes), dataset.action_horizon + 1)
    predicted[:, :-1] = torch.linspace(-0.9, -0.1, dataset.action_horizon)

    dataset.relabel_from_value_predictions(predicted)

    assert not torch.allclose(dataset.advantages, before)
    intervention_index = next(
        index
        for index, episode in enumerate(dataset.episodes)
        if episode.source_type == "intervention_correction"
    )
    assert bool(dataset.positive_episodes[intervention_index])