#!/usr/bin/env python3
"""Generate ATC FPPF/FPPP videos against the existing Extended-251 FFFF run."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from scripts import generate_vbench8_extended_strategies as extended

import lpips
import torch
from omegaconf import OmegaConf
from safetensors import safe_open

from scripts import evaluate_single_block_fppf as base
from scripts.generate_vbench8_extended_disca import artifact_paths, read_u8_video
from utils.wan_wrapper import WanTextEncoder, WanVAEWrapper


MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json"
REFERENCE_DEFAULT = (
    REPO_ROOT / "evaluation_runs/vbench8_extended_stage1_step2000_20260901"
)
OUTPUT_DEFAULT = (
    REPO_ROOT / "evaluation_runs/vbench8_extended_atc_stage1_step2000_20260901"
)
PATTERN_TO_STEPS = {"FPPF": [1, 2], "FPPP": [1, 2, 3]}


def read_predictor_config(weights: Path) -> dict[str, Any]:
    with safe_open(weights, framework="pt", device="cpu") as handle:
        metadata = handle.metadata() or {}
    raw = metadata.get("predictor_config")
    if raw is None:
        raise ValueError(f"ATC checkpoint has no predictor_config metadata: {weights}")
    config = json.loads(raw)
    if config.get("input_variant") != "atc":
        raise ValueError(f"Checkpoint is not ATC: {config}")
    return config


def strategy_configs(
    scope: str, patterns: list[str] | tuple[str, ...]
) -> list[dict[str, Any]]:
    return [
        {
            "name": f"atc_{scope}_{pattern.lower()}",
            "policy": "fppf",
            "pattern": pattern,
            "candidate_steps": PATTERN_TO_STEPS[pattern],
            "allow_chunk0_predictor": False,
            "beta": None,
            "threshold": None,
            "target_accepts": 6 * len(PATTERN_TO_STEPS[pattern]),
            "head": None,
            "input_variant": "atc",
            "atc_previous_scope": scope,
        }
        for pattern in patterns
    ]


def load_models(
    weights: Path, scope: str, device: torch.device
) -> tuple[Any, Any, Any, Any, None, None, Any]:
    predictor_config = read_predictor_config(weights)
    if predictor_config.get("atc_previous_scope") != scope:
        raise ValueError(
            f"Checkpoint scope={predictor_config.get('atc_previous_scope')} "
            f"does not match requested scope={scope}"
        )
    print("[setup] loading VAE", flush=True)
    vae = WanVAEWrapper().to(device=device, dtype=torch.bfloat16).eval()
    config = OmegaConf.merge(
        OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"),
        OmegaConf.load(REPO_ROOT / "configs/self_forcing_sid.yaml"),
    )
    print(f"[setup] loading Teacher and ATC scope={scope}", flush=True)
    pipeline = base.build_pipeline(
        config, REPO_ROOT / "checkpoints/self_forcing_dmd.pt", vae, device
    )
    predictor = base.load_predictor(
        pipeline.generator.model,
        {
            **predictor_config,
            "weights": weights,
            "atc_collect_diagnostics": False,
        },
        device,
    )
    text_encoder = WanTextEncoder().to(device=device, dtype=torch.bfloat16).eval()
    text_encoder.requires_grad_(False)
    print("[setup] loading LPIPS", flush=True)
    lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
    lpips_model.requires_grad_(False)
    return vae, pipeline, predictor, text_encoder, None, None, lpips_model


def is_complete(
    output_root: Path,
    row: dict[str, Any],
    strategies: list[dict[str, Any]],
) -> bool:
    return all(
        all(path.is_file() for path in artifact_paths(output_root, row, config["name"]))
        for config in strategies
    )


def run_prompt(
    *,
    row: dict[str, Any],
    strategies: list[dict[str, Any]],
    reference_root: Path,
    output_root: Path,
    seed: int,
    models: tuple[Any, ...],
    device: torch.device,
) -> None:
    global_index = int(row["global_index"])
    suite = str(row["prompt_suite"])
    suite_index = int(row["suite_index"])
    prompt = str(row["extended_prompt"])
    reference_path = (
        reference_root / "generated_videos/ffff" / suite / f"{suite_index:03d}.mp4"
    )
    print(f"[prompt] global={global_index} suite={suite}/{suite_index}", flush=True)
    conditional = models[3](text_prompts=[prompt])
    for strategy in strategies:
        name = str(strategy["name"])
        video_path, record_path = artifact_paths(output_root, row, name)
        if video_path.is_file() and record_path.is_file():
            print(f"[cached-strategy] global={global_index} strategy={name}", flush=True)
            continue
        latent, diagnostic = extended.generate_rollout(
            pipeline=models[1],
            conditional_dict=conditional,
            seed=seed,
            device=device,
            predictor=models[2],
            head=None,
            config=strategy,
        )
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            pixels = models[0].decode_to_pixel(latent, use_cache=False)
        prediction_u8 = base.pixels_to_u8(pixels)
        extended.atomic_video(prediction_u8, video_path)
        reference_decoded = read_u8_video(reference_path)
        prediction_decoded = read_u8_video(video_path)
        metrics = base.frame_metrics(
            reference_u8=reference_decoded,
            prediction_u8=prediction_decoded,
            lpips_model=models[6],
            batch_size=4,
            device=device,
        )
        extended.atomic_json(
            record_path,
            {
                "status": "complete",
                "strategy": name,
                "policy": "fppf",
                "pattern": strategy["pattern"],
                "input_variant": "atc",
                "atc_previous_scope": strategy["atc_previous_scope"],
                "candidate_steps": strategy["candidate_steps"],
                "allow_chunk0_predictor": False,
                "global_index": global_index,
                "prompt_suite": suite,
                "suite_index": suite_index,
                "prompt": prompt,
                "seed": seed,
                "generation": {
                    key: value
                    for key, value in diagnostic.items()
                    if key != "decisions"
                },
                "decisions": diagnostic["decisions"],
                "pixel_metrics_vs_ffff": extended.compact_metrics(metrics),
                "pixel_metric_input": (
                    "MP4-decoded uint8 RGB, all 81 frames on both sides"
                ),
                "reference_video": str(reference_path),
                "video": str(video_path.relative_to(output_root)),
            },
        )
        print(
            f"[result] global={global_index} strategy={name} "
            f"accept={diagnostic['accepted_predictor_calls']} "
            f"psnr={metrics['psnr']:.4f} ssim={metrics['ssim']:.6f} "
            f"lpips={metrics['lpips']:.6f}",
            flush=True,
        )
        if hasattr(models[0].model, "clear_cache"):
            models[0].model.clear_cache()
        del latent, pixels, prediction_u8, reference_decoded, prediction_decoded
        torch.cuda.empty_cache()
    del conditional


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu", default=extended.PHYSICAL_GPU)
    parser.add_argument("--scope", choices=("chunk", "last_frame"), required=True)
    parser.add_argument("--predictor-weights", type=Path, required=True)
    parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT)
    parser.add_argument("--reference-root", type=Path, default=REFERENCE_DEFAULT)
    parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT)
    parser.add_argument("--pattern", action="append", choices=tuple(PATTERN_TO_STEPS))
    parser.add_argument("--shard-index", type=int, default=0)
    parser.add_argument("--num-shards", type=int, default=1)
    parser.add_argument("--global-index", action="append", type=int, default=None)
    parser.add_argument("--seed", type=int, default=0)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.num_shards < 1 or not 0 <= args.shard_index < args.num_shards:
        raise ValueError("Invalid shard index/number")
    weights = args.predictor_weights.resolve()
    mapping = extended.read_mapping(args.mapping.resolve())
    if args.global_index is None:
        rows = [
            row for index, row in enumerate(mapping)
            if index % args.num_shards == args.shard_index
        ]
    else:
        requested = set(args.global_index)
        rows = [row for row in mapping if int(row["global_index"]) in requested]
        if len(rows) != len(requested):
            raise ValueError("Unknown global index")
    patterns = args.pattern or ["FPPF", "FPPP"]
    strategies = strategy_configs(args.scope, patterns)
    reference_root = args.reference_root.resolve()
    if len(list((reference_root / "generated_videos/ffff").rglob("*.mp4"))) != 251:
        raise ValueError("Incomplete FFFF reference")
    output_root = args.output_root.resolve()
    output_root.mkdir(parents=True, exist_ok=True)
    extended.write_shard_manifest(
        output_root,
        str(args.gpu),
        args.shard_index,
        args.num_shards,
        rows,
        strategies,
        "running",
    )
    device = torch.device("cuda")
    torch.set_grad_enabled(False)
    models = load_models(weights, args.scope, device)
    completed = 0
    try:
        for row in rows:
            if is_complete(output_root, row, strategies):
                completed += 1
                print(f"[cached] {completed}/{len(rows)} global={row['global_index']}", flush=True)
                continue
            run_prompt(
                row=row,
                strategies=strategies,
                reference_root=reference_root,
                output_root=output_root,
                seed=args.seed,
                models=models,
                device=device,
            )
            completed += 1
            print(f"[progress] {completed}/{len(rows)} prompts", flush=True)
    finally:
        extended.write_shard_manifest(
            output_root,
            str(args.gpu),
            args.shard_index,
            args.num_shards,
            rows,
            strategies,
            "complete" if completed == len(rows) else "failed",
        )
    print(f"[complete] gpu={args.gpu} scope={args.scope} prompts={completed}", flush=True)


if __name__ == "__main__":
    main()