from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from pi0mini.data import ActionNormalizer, Pi0OverfitDataset, load_manifest, materialize_subset


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)

    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_dir / "sample.jpg")

    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)
    np.savez_compressed(
        trajectory_dir / "sample.npz", history_xyz=history, future_xyz=future
    )

    manifest = processed / "overfit_manifest.jsonl"
    manifest.write_text(
        json.dumps(
            {
                "clip_id": "clip-a",
                "t0_us": 1234,
                "image_path": "data/processed/images/sample.jpg",
                "trajectory_path": "data/processed/trajectories/sample.npz",
            }
        )
        + "\n",
        encoding="utf-8",
    )
    return manifest


def test_action_normalizer_is_reversible() -> None:
    history = torch.tensor([[-1.0, 0.0], [0.0, 0.1]])
    future = torch.tensor([[1.0, 0.2], [2.5, 0.4], [4.0, 0.7]])
    normalizer = ActionNormalizer.fit([history], [future])
    actions = normalizer.encode(history, future)

    assert actions.shape == (3, 2)
    assert torch.allclose(normalizer.decode(actions, history), future, atol=1e-6)
    restored = ActionNormalizer.from_dict(normalizer.as_dict())
    assert torch.allclose(restored.decode(actions, history), future, atol=1e-6)


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

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

    dataset = Pi0OverfitDataset(copied_manifest, seed=19)
    repeated = Pi0OverfitDataset(
        copied_manifest,
        seed=19,
        action_statistics=dataset.action_normalizer.as_dict(),
    )
    assert len(dataset) == 4
    assert dataset.state_dim == 16

    original = dataset[0]
    augmented = dataset[1]
    repeated_augmented = repeated[1]
    assert set(original) == {
        "image",
        "state",
        "actions",
        "history_xy",
        "future_xy",
        "sample_id",
    }
    assert original["image"].shape == (3, 64, 64)
    assert original["state"].shape == (16,)
    assert original["actions"].shape == (16, 2)
    assert original["history_xy"].shape == (8, 2)
    assert original["future_xy"].shape == (16, 2)
    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["image"], repeated_augmented["image"])
    assert torch.equal(augmented["state"], repeated_augmented["state"])
    assert torch.equal(augmented["actions"], repeated_augmented["actions"])

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


def test_statistics_can_be_saved(tmp_path: Path) -> None:
    source_manifest = _make_source_project(tmp_path / "source")
    manifest = materialize_subset(source_manifest, tmp_path / "target")
    dataset = Pi0OverfitDataset(manifest)
    path = dataset.save_action_statistics(tmp_path / "stats.json")
    values = json.loads(path.read_text(encoding="utf-8"))

    assert len(values["mean"]) == 2
    assert len(values["scale"]) == 2
    assert min(values["scale"]) >= 0.05