#!/usr/bin/env python
"""Fit the TeaCache rescale polynomial for a base model.

TeaCache ships per-model 4th-degree coefficients that map the cheap input-side
relative L1 distance onto the output-side distance.  Those were fitted for
bidirectional Wan on a 50-step schedule with an ``e0``-only indicator, so they do
not transfer to a 4-step block-causal model with the modulated-input indicator.
This script reruns TeaCache's own fitting procedure on the target model.

    python calibrate.py --base self_forcing --num-prompts 8 \
        --out results/coeff_self_forcing.json
"""

import argparse
import json
import os
import sys

ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)

from harness import BASES, enter_base, load_pipeline  # noqa: E402


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--base", choices=sorted(BASES), required=True)
    p.add_argument("--indicator", default="modulated_input",
                   choices=["modulated_input", "e0"])
    p.add_argument("--num-prompts", type=int, default=8)
    p.add_argument("--num-output-frames", type=int, default=21)
    p.add_argument("--degree", type=int, default=1,
                   help="Degree 4 (TeaCache default) is ill-conditioned on the narrow "
                        "distance range a 4-step schedule produces; degree 1 stays "
                        "monotone and is safe to extrapolate per token.")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--out", required=True)
    args = p.parse_args()

    out_path = args.out if os.path.isabs(args.out) else os.path.join(ROOT, args.out)
    base = enter_base(args.base)

    import numpy as np
    import torch

    from utils.dataset import TextDataset
    from utils.misc import set_seed

    from cachelib import CacheController, build_method, cached_inference, install

    set_seed(args.seed)
    torch.set_grad_enabled(False)
    device = torch.device("cuda")
    pipeline = load_pipeline(base)

    method = build_method("calibrate", indicator=args.indicator)
    num_steps = len(pipeline.denoising_step_list)
    controller = CacheController(method, num_steps=num_steps,
                                 forced_steps=(0, num_steps - 1))
    install(pipeline.generator.model, controller)

    prompt_file = ("prompts/MovieGenVideoBench.txt"
                   if os.path.exists("prompts/MovieGenVideoBench.txt")
                   else "prompts/demos.txt")
    dataset = TextDataset(prompt_path=prompt_file)

    for i in range(min(args.num_prompts, len(dataset))):
        set_seed(args.seed)
        noise = torch.randn([1, args.num_output_frames, 16, 60, 104],
                            device=device, dtype=torch.bfloat16)
        cached_inference(pipeline, controller, noise, [dataset[i]["prompts"]],
                         decode=False)
        print(f"[{i}] samples={len(method.samples)}", flush=True)

    xs = np.array([s["x"] for s in method.samples], dtype=np.float64)
    ys = np.array([s["y"] for s in method.samples], dtype=np.float64)
    coeffs = np.polyfit(xs, ys, args.degree)
    pred = np.poly1d(coeffs)(xs)

    result = {
        "base": args.base,
        "indicator": args.indicator,
        "degree": args.degree,
        "num_samples": int(xs.size),
        "coefficients": [float(c) for c in coeffs],
        "x_range": [float(xs.min()), float(xs.max())],
        "y_range": [float(ys.min()), float(ys.max())],
        "fit_r2": float(1 - ((ys - pred) ** 2).sum() / ((ys - ys.mean()) ** 2).sum()),
        "per_step_mean_x": {
            str(k): float(xs[[s["step"] == k for s in method.samples]].mean())
            for k in sorted({s["step"] for s in method.samples})
        },
        "per_step_mean_y": {
            str(k): float(ys[[s["step"] == k for s in method.samples]].mean())
            for k in sorted({s["step"] for s in method.samples})
        },
        "samples": method.samples,
    }
    print(json.dumps({k: v for k, v in result.items() if k != "samples"}, indent=2))

    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    with open(out_path, "w") as f:
        json.dump(result, f, indent=2)
    print(f"wrote {out_path}")


if __name__ == "__main__":
    main()