#!/usr/bin/env python3
"""Reproduce the benchmark pipeline for a SINGLE video, end-to-end.

Runs the EXACT same DCVC-RT readiness patch-selection used to generate the
benchmark codec assets (deterministic on the pytorch fallback path), then feeds
the resulting canvases to the release model — giving the same result a benchmark
eval would produce for this video.

ALL DCVC selection params live in ``<model>/processor/preprocessor_config.json`` under
``codec.dcvc`` and are read by the pipeline via ``codec_dcvc_config`` — NOT env
vars. The file defaults to the ``b50`` benchmark config. ``--config`` optionally
patches that file to a named preset (cap12 / b50 / s95_b50) before generating;
``--config current`` (default) uses the file as-is.

Note: inference uses trust_remote_code; the Mage-VL-Ported remote code targets a
newer transformers. Under the eval env (transformers 4.57.3) pass the equivalent
``--model .../Mage-VL-Base`` (identical weights, 4.57.3-compatible code).

Usage:
    python neural_codec/reproduce_bench.py --video V.mp4 --question "..." [--config b50]
    python neural_codec/reproduce_bench.py --video V.mp4 --gen_only   # just write canvases
"""
from __future__ import annotations
import argparse, os, subprocess, sys
from pathlib import Path

_HERE = os.path.dirname(os.path.abspath(__file__))
_MODEL_DEFAULT = os.path.dirname(_HERE)  # Mage-VL-Ported/

# Named presets (written into codec.dcvc by --config). Values are config keys.
CONFIGS = {
    "cap12":   {"per_frame_cap_ratio": 1.2, "bottom_atten": 1.0, "bottom_band": 0.10, "threshold_scale": 1.0},
    "b50":     {"per_frame_cap_ratio": 1.2, "bottom_atten": 0.5, "bottom_band": 0.10, "threshold_scale": 1.0},
    "s95_b50": {"per_frame_cap_ratio": 1.2, "bottom_atten": 0.5, "bottom_band": 0.10, "threshold_scale": 0.95},
}


def _patch_config_file(config: str):
    """Optionally write a named preset's selection knobs into the model's
    preprocessor_config.json (codec.dcvc), so the readiness pipeline picks them
    up. ``config='current'`` leaves the file untouched (use whatever is set)."""
    if config == "current":
        return
    import json, collections
    p = os.path.join(_MODEL_DEFAULT, "processor", "preprocessor_config.json")
    with open(p, encoding="utf-8") as f:
        c = json.load(f, object_pairs_hook=collections.OrderedDict)
    c["codec"]["dcvc"].update(CONFIGS[config])
    tmp = p + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(c, f, indent=2)
    os.replace(tmp, p)  # atomic: never leave a truncated config on write error
    print(f"[reproduce] wrote preset '{config}' into {p} (codec.dcvc): {CONFIGS[config]}")


def gen_canvases(video: str, out_dir: str, config: str, gpu: int) -> str:
    # ALL DCVC selection params come from preprocessor_config.json's codec.dcvc
    # (read by the pipeline via codec_dcvc_config) — NOT env / CLI. --config may
    # patch that file first.
    _patch_config_file(config)
    env = dict(os.environ)
    env["CUDA_VISIBLE_DEVICES"] = str(gpu)
    env["http_proxy"] = ""; env["https_proxy"] = ""
    os.makedirs(out_dir, exist_ok=True)
    cmd = [sys.executable, os.path.join(_HERE, "dcvc_readiness_gen.py"),
           "--video", video, "--out_dir", out_dir]
    print(f"[reproduce] gen (config={config}): {' '.join(cmd)}", flush=True)
    subprocess.run(cmd, check=True, env=env)
    return out_dir


def run_inference(model_dir: str, asset_dir: str, question: str, max_new_tokens: int,
                  max_pixels, gpu: int) -> str:
    import torch
    sys.path.insert(0, _HERE)                                   # codec_loader
    sys.path.insert(0, os.path.join(model_dir, "processor"))    # MageVLProcessor + siblings
    from codec_loader import build_inputs_from_assets  # noqa: E402
    from processing_magevl import MageVLProcessor  # noqa: E402
    from transformers import AutoModelForCausalLM
    device = torch.device(f"cuda:{gpu}" if torch.cuda.is_available() else "cpu")
    processor = MageVLProcessor.from_pretrained(model_dir)      # subdir-aware
    model = AutoModelForCausalLM.from_pretrained(
        os.path.join(model_dir, "transformer"), trust_remote_code=True,
        torch_dtype="auto").to(device).eval()
    messages = [{"role": "user", "content": [{"type": "video"},
                                             {"type": "text", "text": question}]}]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = build_inputs_from_assets(processor, asset_dir, text,
                                      max_pixels=max_pixels, device=device)
    inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
    with torch.inference_mode():
        gen = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    new = gen[0, inputs["input_ids"].shape[1]:]
    return processor.tokenizer.decode(new, skip_special_tokens=True).strip()


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--video", required=True)
    ap.add_argument("--question", default="Describe this video in detail.")
    ap.add_argument("--config", default="current", choices=["current"]+list(CONFIGS),
                    help="current=use preprocessor_config.json as-is; a preset patches codec.dcvc first")
    ap.add_argument("--model", default=_MODEL_DEFAULT)
    ap.add_argument("--out_dir", default=None, help="asset dir (default /tmp/reproduce/<stem>)")
    ap.add_argument("--gpu", type=int, default=0)
    ap.add_argument("--max_new_tokens", type=int, default=256)
    ap.add_argument("--max_pixels", type=int, default=None,
                    help="canvas pixel budget; default None -> codec.dcvc.max_pixels "
                         "(150000) from preprocessor_config.json")
    ap.add_argument("--gen_only", action="store_true")
    args = ap.parse_args()

    out_dir = args.out_dir or f"/tmp/reproduce/{Path(args.video).stem}_{args.config}"
    gen_canvases(args.video, out_dir, args.config, args.gpu)
    nc = len([p for ext in ("jpg", "png") for p in Path(out_dir).glob(f"canvas_*.{ext}")])
    print(f"[reproduce] wrote {nc} canvases -> {out_dir}")
    if args.gen_only:
        return
    ans = run_inference(args.model, out_dir, args.question, args.max_new_tokens,
                        args.max_pixels, args.gpu)
    print("\n================ ANSWER ================\n" + ans + "\n")


if __name__ == "__main__":
    main()