"""Turn on the two speed knobs that actually help Chronos-2 on a GPU.

Weight-only INT8 is smaller, not faster, on an RTX 3090. What *does* help:

  1. TF32  — use the GPU's fast float math (off by default in PyTorch).
  2. torch.compile — fuse kernels so a tiny 120M model is not launch-bound.

Usage after snapshot_download of this repo:

    from load import load
    from fast_infer import speedup
    pipe = speedup(load(".", device="cuda"))

Or from the command line:

    python fast_infer.py --repo oxfrug/chronos-2-int8-torchao
    python fast_infer.py --fp32 amazon/chronos-2   # original model, same knobs
"""
from __future__ import annotations

import argparse
import time

import numpy as np
import torch
from chronos import Chronos2Pipeline


def enable_tf32() -> None:
    if not torch.cuda.is_available():
        return
    torch.set_float32_matmul_precision("high")
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True


def speedup(pipe: Chronos2Pipeline, compile: bool = True, tf32: bool = True) -> Chronos2Pipeline:
    """Return the same pipeline with TF32 and (optionally) torch.compile."""
    if tf32:
        enable_tf32()
    if compile and torch.cuda.is_available():
        pipe.model = torch.compile(pipe.model, mode="reduce-overhead")
    return pipe


def _predict_once(pipe: Chronos2Pipeline, y: np.ndarray, h: int) -> None:
    pipe.predict_quantiles(inputs=[y], prediction_length=h, quantile_levels=[0.5])


def bench(pipe: Chronos2Pipeline, n_warm: int = 3, n_run: int = 8, h: int = 24) -> float:
    y = np.sin(np.linspace(0, 40, 512)).astype(np.float32)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    for _ in range(n_warm):
        _predict_once(pipe, y, h)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n_run):
        _predict_once(pipe, y, h)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    return 1000.0 * (time.perf_counter() - t0) / n_run


def main() -> None:
    ap = argparse.ArgumentParser(description="Load Chronos-2 with TF32 + torch.compile and time one call.")
    ap.add_argument("--repo", default="oxfrug/chronos-2-int8-torchao")
    ap.add_argument("--fp32", default=None, help="Load this HF id with Chronos2Pipeline.from_pretrained instead.")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--no-compile", action="store_true")
    ap.add_argument("--no-tf32", action="store_true")
    args = ap.parse_args()

    if args.fp32:
        pipe = Chronos2Pipeline.from_pretrained(
            args.fp32,
            device_map=args.device,
            dtype=torch.float32,
        )
    else:
        from pathlib import Path
        import importlib.util

        here = Path(__file__).resolve().parent
        load_py = here / "load.py"
        if load_py.exists():
            spec = importlib.util.spec_from_file_location("c2int8", load_py)
            mod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(mod)
            pipe = mod.load(str(here if (here / "model.safetensors").exists() else args.repo), device=args.device)
        else:
            spec = importlib.util.spec_from_file_location("c2int8", "load.py")
            raise SystemExit("need load.py next to this script, or pass --fp32 amazon/chronos-2")

    pipe = speedup(pipe, compile=not args.no_compile, tf32=not args.no_tf32)
    print("first call compiles; later calls are the ones that count", flush=True)
    ms = bench(pipe)
    print(f"{ms:.2f} ms / call  (512 context, horizon 24, one series)")
    print("device", args.device, "compile", not args.no_compile, "tf32", not args.no_tf32)


if __name__ == "__main__":
    main()