#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import statistics
from functools import partial
from pathlib import Path
from typing import Any

import torch
from safetensors.torch import load_file

from alpamayo_lite.action_space import RelativeDeltaActionSpace
from alpamayo_lite.config import ExperimentConfig, load_config
from alpamayo_lite.data import OverfitSample, load_manifest
from alpamayo_lite.diffusion import FlowMatching
from alpamayo_lite.evaluation import (
    save_prediction_arrays,
    save_trajectory_plot,
    trajectory_metrics,
)
from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV
from alpamayo_lite.prompts import build_expert_prefix, build_vlm_prefix
from alpamayo_lite.runtime import (
    empty_device_cache,
    parameter_counts,
    resolve_dtype,
    seed_everything,
    select_device,
    write_json,
)
from alpamayo_lite.vlm import (
    extract_prefix_kv,
    generate_reasoning,
    load_vlm,
    vlm_text_config,
)


def _resolve(cfg: ExperimentConfig, value: str) -> Path:
    path = Path(value).expanduser()
    return path.resolve() if path.is_absolute() else (cfg.root / path).resolve()


def _read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _summary(values: list[float]) -> dict[str, float]:
    return {
        "mean": statistics.fmean(values),
        "std": statistics.pstdev(values),
        "min": min(values),
        "max": max(values),
    }


def _evaluate_expert(
    stage2_dir: Path,
    vlm: Any,
    prefix: PrefixKV,
    sample: OverfitSample,
    device: torch.device,
    dtype: torch.dtype,
    num_rollouts: int,
    base_seed: int,
) -> dict[str, Any]:
    expert_config = _read_json(stage2_dir / "expert_config.json")
    metrics_path = stage2_dir / "stage2_metrics.json"
    saved_metrics = _read_json(metrics_path) if metrics_path.exists() else {}
    architecture = ExpertArchitecture(**expert_config["architecture"])
    expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
        device=device, dtype=dtype
    )
    expert.load_state_dict(load_file(stage2_dir / "expert.safetensors"), strict=True)
    expert.requires_grad_(False)
    expert.eval()

    action_shape = tuple(int(value) for value in expert_config["action_shape"])
    if len(action_shape) != 2:
        raise ValueError(f"Invalid action shape in {stage2_dir}: {action_shape}")
    flow = FlowMatching(
        action_shape=action_shape,
        inference_steps=int(expert_config["inference_steps"]),
    )
    action_space = RelativeDeltaActionSpace(
        torch.tensor(expert_config["action_mean"], device=device),
        torch.tensor(expert_config["action_scale"], device=device),
    ).to(device)
    history, future = sample.load_trajectory()
    history_xy = history[None, :, :2].to(device)
    target_xy = future[None, :, :2].to(device)
    target_action = action_space.trajectory_to_action(history_xy, target_xy).float()

    rollouts: list[dict[str, float | int]] = []
    first_prediction: tuple[torch.Tensor, torch.Tensor] | None = None
    for index in range(num_rollouts):
        seed = base_seed + index
        seed_everything(seed)
        predicted_action = flow.sample(
            batch_size=1,
            step_fn=partial(expert, prefix=prefix),
            device=device,
            dtype=dtype,
        ).float()
        predicted_xy = action_space.action_to_trajectory(predicted_action, history_xy)
        record: dict[str, float | int] = {
            "seed": seed,
            "action_mse": float(
                torch.nn.functional.mse_loss(predicted_action, target_action).cpu()
            ),
            **trajectory_metrics(predicted_xy, target_xy),
        }
        if not torch.isfinite(predicted_action).all() or not torch.isfinite(predicted_xy).all():
            raise RuntimeError(f"Non-finite reload prediction for seed {seed}")
        rollouts.append(record)
        if first_prediction is None:
            first_prediction = (predicted_action, predicted_xy)

    assert first_prediction is not None
    first_action, first_xy = first_prediction
    save_prediction_arrays(
        stage2_dir / "reload_prediction.npz",
        history[:, :2],
        target_action[0],
        first_action[0],
        target_xy[0],
        first_xy[0],
    )
    save_trajectory_plot(
        stage2_dir / "reload_trajectory.png",
        history[:, :2],
        target_xy[0],
        first_xy[0],
    )

    metric_names = ("action_mse", "ade_m", "fde_m", "max_error_m")
    result = {
        "stage2_dir": str(stage2_dir),
        "checkpoint_steps": int(
            expert_config.get("completed_steps", saved_metrics.get("steps", 0))
        ),
        "expert_parameters": parameter_counts(expert),
        "architecture": expert_config["architecture"],
        "action_mean": expert_config["action_mean"],
        "action_scale": expert_config["action_scale"],
        "action_shape": list(action_shape),
        "inference_steps": flow.inference_steps,
        "num_rollouts": num_rollouts,
        "rollouts": rollouts,
        "summary": {
            name: _summary([float(record[name]) for record in rollouts])
            for name in metric_names
        },
    }
    del expert
    empty_device_cache(device)
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/local_overfit.yaml")
    parser.add_argument(
        "--stage2-dir",
        action="append",
        dest="stage2_dirs",
        help="Stage-2 directory to reload; repeat to compare checkpoints.",
    )
    parser.add_argument("--num-rollouts", type=int, default=16)
    parser.add_argument("--base-seed", type=int, default=None)
    parser.add_argument("--output", default="outputs/reload_evaluation.json")
    args = parser.parse_args()
    if args.num_rollouts <= 0:
        raise ValueError("--num-rollouts must be positive")

    cfg = load_config(args.config)
    seed_everything(cfg.seed)
    device = select_device()
    dtype = resolve_dtype(cfg.vlm.dtype, device)
    stage2_values = args.stage2_dirs or [cfg.paths.stage2_dir]
    stage2_dirs = [_resolve(cfg, value) for value in stage2_values]
    sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
    history, _ = sample.load_trajectory()

    stage1_dir = cfg.resolve(cfg.paths.stage1_dir)
    processor, vlm = load_vlm(stage1_dir, cfg.vlm, device, dtype)
    vlm.requires_grad_(False)
    vlm.eval()
    generated_reasoning = generate_reasoning(
        vlm,
        processor,
        sample.load_image(),
        build_vlm_prefix(history),
        device,
        dtype,
        cfg.vlm.max_reasoning_tokens,
    )
    prefix = extract_prefix_kv(
        vlm,
        processor,
        sample.load_image(),
        build_expert_prefix(history, generated_reasoning),
        device,
        dtype,
    )
    base_seed = args.base_seed if args.base_seed is not None else cfg.seed + 1
    models = [
        _evaluate_expert(
            stage2_dir,
            vlm,
            prefix,
            sample,
            device,
            dtype,
            args.num_rollouts,
            base_seed,
        )
        for stage2_dir in stage2_dirs
    ]
    result = {
        "stage": "reloaded_pipeline_evaluation",
        "device": str(device),
        "dtype": str(dtype),
        "stage1_dir": str(stage1_dir),
        "vlm_parameters": parameter_counts(vlm),
        "generated_reasoning": generated_reasoning,
        "target_reasoning": sample.reasoning,
        "reasoning_exact_match": generated_reasoning == sample.reasoning,
        "prefix_layers": prefix.num_layers,
        "prefix_sequence_length": prefix.sequence_length,
        "prefix_key_shape": list(prefix.keys[0].shape),
        "models": models,
    }
    output = _resolve(cfg, args.output)
    write_json(output, result)
    print(json.dumps(result, indent=2, ensure_ascii=False))
    del vlm
    empty_device_cache(device)


if __name__ == "__main__":
    main()