#!/usr/bin/env python3
"""Generate original-concat FPPP videos against the existing FFFF reference."""

from __future__ import annotations

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 evaluate_single_block_fppf as base
from scripts import generate_vbench8_extended_disca as runner

import lpips
import torch
from omegaconf import OmegaConf

from utils.wan_wrapper import WanTextEncoder, WanVAEWrapper


WEIGHTS_DEFAULT = (
    REPO_ROOT
    / "training_runs/layer17_stage1_1000p_4gpu_b16_2000steps"
    / "checkpoint_step_2000/predictor.safetensors"
)
OUTPUT_DEFAULT = (
    REPO_ROOT
    / "evaluation_runs/vbench8_extended_concat_fppp_stage1_step2000_20260902"
)


def strategy_configs(patterns: list[str] | tuple[str, ...]) -> list[dict[str, Any]]:
    if patterns != ["FPPP"] and tuple(patterns) != ("FPPP",):
        raise ValueError("Original-concat evaluation supports exactly --pattern FPPP")
    return [
        {
            "name": "concat_fppp",
            "policy": "fppf",
            "pattern": "FPPP",
            "candidate_steps": [1, 2, 3],
            "allow_chunk0_predictor": False,
            "beta": None,
            "threshold": None,
            "target_accepts": 18,
            "head": None,
            "input_variant": "self_forcing",
        }
    ]


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 original concat 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": "self_forcing",
            "gate_mode": "baseline",
        },
        device,
    )
    if predictor.input_variant != "self_forcing":
        raise RuntimeError("Loaded Predictor is not the original concat 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 main() -> None:
    runner.WEIGHTS_DEFAULT = WEIGHTS_DEFAULT
    runner.OUTPUT_DEFAULT = OUTPUT_DEFAULT
    runner.strategy_configs = strategy_configs
    runner.load_models = load_models
    runner.main()


if __name__ == "__main__":
    main()