from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch

from .data import Pi05OverfitDataset
from .flow import FlowMatching
from .model import Pi05Config, Pi05Mini
from .runtime import choose_device, write_json


@torch.no_grad()
def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    cfg = checkpoint["config"]
    root = Path(__file__).resolve().parents[2]
    device = choose_device(cfg.get("device", "auto"))
    data_manifest = Path(manifest) if manifest else root / "data/overfit/overfit_manifest.jsonl"
    dataset = Pi05OverfitDataset(
        data_manifest, image_size=cfg["data"]["image_size"],
        action_horizon=cfg["data"]["action_horizon"],
        action_statistics=checkpoint["action_statistics"]
    )
    model = Pi05Mini(Pi05Config(**checkpoint["model_config"]))
    model.load_state_dict(checkpoint["model"])
    model.to(device).eval()
    flow = FlowMatching(cfg["inference"]["flow_steps"])
    rows = []
    figure, axes = plt.subplots(1, len(dataset), figsize=(4 * len(dataset), 4), squeeze=False)
    for index, sample in enumerate(dataset):
        image = sample["image"][None].to(device)
        tokens = sample["prefix_tokens"][None].to(device)
        # Autoregressively predict semantic subtask followed by the complete
        # discretized action stream (ACTION_START, 32 scalar tokens, ACTION_END).
        zero_action = torch.zeros(1, dataset.action_horizon, 2, device=device)
        generated: list[int] = []
        expected_tokens = 1 + 2 + dataset.action_horizon * 2
        current_tokens = tokens
        for generation_step in range(expected_tokens):
            semantic = model(
                image, current_tokens, zero_action,
                torch.zeros(1, device=device),
                generated_token_start=tokens.shape[1]
            )
            if generation_step == 0:
                candidates = torch.tensor(
                    dataset.tokenizer.subtasks.token_ids, device=device
                )
                candidate_logits = semantic.text_logits[0, -1, candidates]
                next_token = int(candidates[candidate_logits.argmax()].cpu())
            else:
                next_token = int(semantic.text_logits[0, -1].argmax().cpu())
            generated.append(next_token)
            current_tokens = torch.cat(
                (current_tokens, torch.tensor([[next_token]], device=device)), dim=1
            )
        predicted_token = generated[0]
        vlm_action_tokens = torch.tensor(generated[1:], dtype=torch.long)
        try:
            vlm_normalized = dataset.tokenizer.decode_actions(vlm_action_tokens)
            vlm_xy = dataset.action_normalizer.decode(
                vlm_normalized, sample["history_xy"]
            ).numpy()
            vlm_valid = True
        except ValueError:
            vlm_xy = np.full_like(sample["future_xy"].numpy(), np.nan)
            vlm_valid = False
        conditioned_tokens = dataset.tokenizer.build_low_level_prefix(
            sample["history_xy"], predicted_token
        )[None].to(device)
        def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
            return model(
                image, conditioned_tokens, x, t
            ).vector_field
        normalized = flow.sample(
            velocity, (1, dataset.action_horizon, 2), device,
            steps=cfg["inference"]["flow_steps"],
            generator=None,
        )[0].cpu()
        predicted_xy = dataset.action_normalizer.decode(
            normalized, sample["history_xy"]
        ).numpy()
        target_xy = sample["future_xy"].numpy()
        diffusion_mse = float(np.mean((predicted_xy - target_xy) ** 2))
        vlm_mse = float(np.mean((vlm_xy - target_xy) ** 2)) if vlm_valid else None
        target_label = dataset.tokenizer.subtasks.decode(sample["subtask_target"])
        try:
            predicted_label = dataset.tokenizer.subtasks.decode(predicted_token)
        except ValueError:
            predicted_label = f"token_{predicted_token}"
        rows.append({
            "sample_id": sample["sample_id"], "target_subtask": target_label,
            "task_instruction": sample["task_instruction"],
            "predicted_subtask": predicted_label,
            "vlm_action_tokens_valid": vlm_valid,
            "vlm_trajectory_mse": vlm_mse,
            "diffusion_trajectory_mse": diffusion_mse,
        })
        ax = axes[0, index]
        ax.plot(target_xy[:, 0], target_xy[:, 1], "o-", label="target")
        ax.plot(vlm_xy[:, 0], vlm_xy[:, 1], "s--", label="VLM tokens")
        ax.plot(predicted_xy[:, 0], predicted_xy[:, 1], "x--", label="diffusion")
        ax.scatter([0], [0], c="black", marker="s", label="ego")
        vlm_text = "invalid" if vlm_mse is None else f"{vlm_mse:.3f}"
        ax.set_title(
            f"{predicted_label}\nVLM={vlm_text}, diffusion={diffusion_mse:.3f}"
        )
        ax.set_xlabel("forward x (m)")
        ax.set_ylabel("left y (m)")
        ax.axis("equal")
        ax.grid(True, alpha=.3)
        ax.legend(fontsize=8)
    figure.tight_layout()
    output_dir = Path(checkpoint_path).parent
    plot_path = output_dir / "trajectory_comparison.png"
    figure.savefig(plot_path, dpi=160)
    plt.close(figure)
    write_json(output_dir / "evaluation.json", rows)
    print(rows)
    return plot_path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt")
    parser.add_argument("--manifest")
    args = parser.parse_args()
    evaluate(args.checkpoint, args.manifest)


if __name__ == "__main__":
    main()