#!/usr/bin/env python3
"""Generate DISCA pattern videos using an existing Extended-251 FFFF reference."""

from __future__ import annotations

import argparse
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 torchvision.io import read_video

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


MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json"
WEIGHTS_DEFAULT = (
    REPO_ROOT
    / "training_runs/disca_layer17_stage1_1000p_4gpu_b16_2000steps"
    / "checkpoint_step_2000/predictor.safetensors"
)
OUTPUT_DEFAULT = (
    REPO_ROOT
    / "evaluation_runs/vbench8_extended_disca_chunk0_fppf_stage1_step2000"
)
REFERENCE_DEFAULT = (
    REPO_ROOT / "evaluation_runs/vbench8_extended_stage1_step2000_20260901"
)
PATTERN_TO_CANDIDATE_STEPS = {
    "FPFF": [1],
    "FPPF": [1, 2],
    "FPPP": [1, 2, 3],
}


def strategy_configs(patterns: list[str] | tuple[str, ...]) -> list[dict[str, Any]]:
    configs: list[dict[str, Any]] = []
    for pattern in patterns:
        candidate_steps = PATTERN_TO_CANDIDATE_STEPS[pattern]
        configs.append(
            {
                "name": f"disca_{pattern.lower()}",
                "policy": "fppf",
                "pattern": pattern,
                "candidate_steps": candidate_steps,
                "allow_chunk0_predictor": True,
                "beta": None,
                "threshold": None,
                "target_accepts": 7 * len(candidate_steps),
                "head": None,
                "input_variant": "disca",
            }
        )
    return configs


def load_models(
    weights: Path, device: torch.device
) -> tuple[Any, Any, Any, Any, None, None, Any]:
    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("[setup] loading frozen Teacher and DISCA Predictor", flush=True)
    pipeline = base.build_pipeline(
        config, REPO_ROOT / "checkpoints/self_forcing_dmd.pt", vae, device
    )
    predictor = base.load_predictor(
        pipeline.generator.model,
        {
            "source_layer": 17,
            "weights": weights,
            "input_variant": "disca",
            "gate_mode": "baseline",
        },
        device,
    )
    if predictor.input_variant != "disca":
        raise RuntimeError("Loaded Predictor is not the DISCA input variant")
    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 read_u8_video(path: Path) -> torch.Tensor:
    if not path.is_file():
        raise FileNotFoundError(path)
    frames, _, _ = read_video(str(path), pts_unit="sec", output_format="TCHW")
    frames = frames.to(device="cpu", dtype=torch.uint8).contiguous()
    if frames.ndim != 4 or frames.shape[0] != 81 or frames.shape[1] != 3:
        raise RuntimeError(f"Unexpected decoded video shape {tuple(frames.shape)}: {path}")
    return frames


def artifact_paths(
    output_root: Path, mapping_row: dict[str, Any], strategy_name: str
) -> tuple[Path, Path]:
    suite = str(mapping_row["prompt_suite"])
    suite_index = int(mapping_row["suite_index"])
    global_index = int(mapping_row["global_index"])
    return (
        output_root
        / "generated_videos"
        / strategy_name
        / suite
        / f"{suite_index:03d}.mp4",
        output_root
        / "generation_metrics/per_prompt"
        / strategy_name
        / f"global_{global_index:04d}.json",
    )


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


def run_prompt(
    *,
    mapping_row: dict[str, Any],
    strategy: dict[str, Any],
    reference_root: Path,
    output_root: Path,
    seed: int,
    vae: Any,
    pipeline: Any,
    predictor: Any,
    text_encoder: Any,
    lpips_model: Any,
    device: torch.device,
) -> None:
    global_index = int(mapping_row["global_index"])
    suite = str(mapping_row["prompt_suite"])
    suite_index = int(mapping_row["suite_index"])
    prompt = str(mapping_row["extended_prompt"])
    strategy_name = str(strategy["name"])
    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 = text_encoder(text_prompts=[prompt])
    latent, diagnostic = extended.generate_rollout(
        pipeline=pipeline,
        conditional_dict=conditional,
        seed=seed,
        device=device,
        predictor=predictor,
        head=None,
        config=strategy,
    )
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        pixels = vae.decode_to_pixel(latent, use_cache=False)
    prediction_u8 = base.pixels_to_u8(pixels)
    video_path, record_path = artifact_paths(output_root, mapping_row, strategy_name)
    extended.atomic_video(prediction_u8, video_path)

    # The original run retained FFFF MP4s, not its pre-encoding uint8 frames.
    # Decode both sides so the reused reference and DISCA prediction receive
    # the same H.264/CRF-18 treatment before full-81-frame pixel metrics.
    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=lpips_model,
        batch_size=4,
        device=device,
    )
    record = {
        "status": "complete",
        "strategy": strategy_name,
        "policy": "fppf",
        "pattern": strategy["pattern"],
        "input_variant": strategy["input_variant"],
        "candidate_steps": strategy["candidate_steps"],
        "allow_chunk0_predictor": strategy["allow_chunk0_predictor"],
        "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)),
    }
    extended.atomic_json(record_path, record)
    print(
        f"[result] global={global_index} strategy={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(vae.model, "clear_cache"):
        vae.model.clear_cache()
    del (
        conditional,
        latent,
        pixels,
        prediction_u8,
        reference_decoded,
        prediction_decoded,
        metrics,
    )
    torch.cuda.empty_cache()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu", default=extended.PHYSICAL_GPU)
    parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT)
    parser.add_argument("--predictor-weights", type=Path, default=WEIGHTS_DEFAULT)
    parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT)
    parser.add_argument("--reference-root", type=Path, default=REFERENCE_DEFAULT)
    parser.add_argument(
        "--pattern",
        action="append",
        choices=tuple(PATTERN_TO_CANDIDATE_STEPS),
        default=None,
        help="DISCA denoise pattern; repeat to generate several. Default: FPPF.",
    )
    parser.add_argument("--shard-index", type=int, default=0)
    parser.add_argument("--num-shards", type=int, default=1)
    parser.add_argument(
        "--global-index",
        dest="global_indices",
        action="append",
        type=int,
        default=None,
        help="Process only this global mapping index; may be repeated.",
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--overwrite", action="store_true")
    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 of shards")
    weights = args.predictor_weights.resolve()
    if not weights.is_file():
        raise FileNotFoundError(weights)
    mapping = extended.read_mapping(args.mapping.resolve())
    patterns = args.pattern or ["FPPF"]
    if len(set(patterns)) != len(patterns):
        raise ValueError("--pattern values must be unique")
    strategies = strategy_configs(patterns)
    reference_root = args.reference_root.resolve()
    reference_records = reference_root / "generation_metrics/per_prompt/ffff"
    reference_videos = reference_root / "generated_videos/ffff"
    if len(list(reference_records.glob("global_*.json"))) != 251:
        raise ValueError(f"Incomplete FFFF reference records: {reference_records}")
    if len(list(reference_videos.rglob("*.mp4"))) != 251:
        raise ValueError(f"Incomplete FFFF reference videos: {reference_videos}")
    if args.global_indices is None:
        rows = [
            row
            for position, row in enumerate(mapping)
            if position % args.num_shards == args.shard_index
        ]
    else:
        requested = set(args.global_indices)
        mapping_by_global = {int(row["global_index"]): row for row in mapping}
        unknown = sorted(requested - set(mapping_by_global))
        if unknown:
            raise ValueError(f"Unknown global indices: {unknown}")
        rows = [row for row in mapping if int(row["global_index"]) in requested]

    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",
    )
    print(
        f"[setup] physical_gpu={args.gpu} shard={args.shard_index}/{args.num_shards} "
        f"prompts={len(rows)} strategies={','.join(strategy['name'] for strategy in strategies)} "
        "reference=existing_ffff",
        flush=True,
    )
    device = torch.device("cuda")
    torch.set_grad_enabled(False)
    models = load_models(weights, device)
    completed = 0
    try:
        for row in rows:
            if not args.overwrite and is_prompt_complete(
                output_root, row, strategies
            ):
                completed += 1
                print(
                    f"[cached] {completed}/{len(rows)} global={row['global_index']}",
                    flush=True,
                )
                continue
            for strategy in strategies:
                video, record = artifact_paths(
                    output_root, row, str(strategy["name"])
                )
                if not args.overwrite and video.is_file() and record.is_file():
                    continue
                run_prompt(
                    mapping_row=row,
                    strategy=strategy,
                    reference_root=reference_root,
                    output_root=output_root,
                    seed=args.seed,
                    vae=models[0],
                    pipeline=models[1],
                    predictor=models[2],
                    text_encoder=models[3],
                    lpips_model=models[6],
                    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} prompts={completed}/{len(rows)}", flush=True)


if __name__ == "__main__":
    main()