"""Upscale a saved direct-runtime H3 latent with the learned 3D model."""

import argparse
import json
import time
from pathlib import Path

import torch

from h3_blackwell_runtime.latent_upscaler import load_h3_latent_upscaler, upscale_h3_latent


parser = argparse.ArgumentParser()
parser.add_argument("--latent", type=Path, required=True)
parser.add_argument("--model", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--scale", type=float, default=2.0)
parser.add_argument("--precision", choices=("float16", "bfloat16", "float32"), default="float16")
args = parser.parse_args()

dtype = getattr(torch, args.precision)
state = torch.load(args.latent, map_location="cpu", weights_only=False)
if not isinstance(state, dict) or "latent" not in state:
    state = {"latent": state}
source = state["latent"]

torch.cuda.synchronize()
started = time.perf_counter()
model = load_h3_latent_upscaler(args.model, dtype=dtype)
torch.cuda.synchronize()
loaded = time.perf_counter()
upscaled = upscale_h3_latent(model, source, scale=args.scale).cpu()
torch.cuda.synchronize()
finished = time.perf_counter()

result = dict(state)
result["latent"] = upscaled
result["width"] = upscaled.shape[-1] * 16
result["height"] = upscaled.shape[-2] * 16
result["upscale"] = {
    "model": args.model.name,
    "scale": args.scale,
    "precision": args.precision,
    "source_shape": tuple(source.shape),
    "output_shape": tuple(upscaled.shape),
}
args.output.parent.mkdir(parents=True, exist_ok=True)
torch.save(result, args.output)
print(json.dumps({
    "output": str(args.output),
    "source_shape": tuple(source.shape),
    "output_shape": tuple(upscaled.shape),
    "model_load_seconds": loaded - started,
    "upscale_seconds": finished - loaded,
}, indent=2))