#!/usr/bin/env python3
"""Standalone inference for the Universal Upscale model.

This repository intentionally ships no training code. `model_scripted.pt` already
contains the network graph, the trained weights and the tiled inference logic,
so this script only needs torch + Pillow + numpy.
"""
import argparse
from pathlib import Path
import warnings

import numpy as np
import torch
from PIL import Image

warnings.filterwarnings("ignore", category=FutureWarning, message=".*torch.jit.load.*")

MAX_OUTPUT_PIXELS = 64_000_000


def load_model(path, repo_id=None, revision="main", device=None):
    local = Path(path)
    if not local.exists():
        script_local = Path(__file__).resolve().parent / path
        if script_local.exists():
            local = script_local
    if not local.exists():
        from huggingface_hub import hf_hub_download

        if not repo_id:
            raise SystemExit("model_scripted.pt not found locally: pass --repo-id <user>/<repo> to download it.")
        local = Path(hf_hub_download(repo_id=repo_id, filename="model_scripted.pt", revision=revision))
    model = torch.jit.load(str(local), map_location=device if device is not None else "cpu")
    model.eval()
    return model


def pil_to_tensor(img):
    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
    return torch.from_numpy(arr).permute(2, 0, 1).contiguous().unsqueeze(0)


def tensor_to_pil(t):
    arr = t.detach().float().clamp(0.0, 1.0).permute(1, 2, 0).cpu().numpy()
    return Image.fromarray((arr * 255.0).round().astype(np.uint8), mode="RGB")


def compute_target(img, args):
    w, h = img.size
    if args.scale is not None:
        scale = float(args.scale)
    elif args.target_long_side is not None:
        scale = float(args.target_long_side) / float(max(w, h))
    elif args.target_width or args.target_height:
        candidates = []
        if args.target_width:
            candidates.append(float(args.target_width) / float(w))
        if args.target_height:
            candidates.append(float(args.target_height) / float(h))
        scale = min(candidates)
    else:
        raise SystemExit("Provide --scale, --target-long-side, --target-width or --target-height.")
    if scale <= 0:
        raise SystemExit("Target scale must be positive.")

    out_w, out_h = max(1, int(round(w * scale))), max(1, int(round(h * scale)))
    canvas_w, canvas_h = out_w, out_h
    if args.pad_to_target and args.target_width and args.target_height:
        canvas_w, canvas_h = int(args.target_width), int(args.target_height)
        if canvas_w < out_w or canvas_h < out_h:
            raise SystemExit("--pad-to-target canvas must contain the aspect-preserving output.")
    return out_w, out_h, canvas_w, canvas_h, scale


def main():
    ap = argparse.ArgumentParser(description="Universal image upscaler (self-contained TorchScript model).")
    ap.add_argument("--input", required=True, help="Input image (.png/.jpg/.webp/.tif)")
    ap.add_argument("--output", required=True, help="Output image (.png/.jpg/.webp/.tif)")
    ap.add_argument("--model", default="model_scripted.pt", help="Path to the TorchScript model")
    ap.add_argument("--repo-id", default=None, help="Download the model from this Hub repo instead")
    ap.add_argument("--scale", type=float, default=None, help="Arbitrary scale factor, e.g. 4 or 10")
    ap.add_argument("--target-long-side", type=int, default=None, help="Target long side in pixels, e.g. 3000")
    ap.add_argument("--target-width", type=int, default=None)
    ap.add_argument("--target-height", type=int, default=None)
    ap.add_argument("--pad-to-target", action="store_true", help="Pad the canvas to the exact target size")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--amp", action="store_true", help="fp16 autocast (CUDA only)")
    ap.add_argument("--webp-lossless", action="store_true")
    args = ap.parse_args()

    device = torch.device(args.device)
    img = Image.open(args.input).convert("RGB")
    out_w, out_h, canvas_w, canvas_h, scale = compute_target(img, args)
    if canvas_w * canvas_h > MAX_OUTPUT_PIXELS:
        raise SystemExit(f"Output {canvas_w}x{canvas_h} is too large ({canvas_w * canvas_h} px). Lower --scale/--target.")

    print(f"Input {img.width}x{img.height} -> output {out_w}x{out_h} (scale {scale:.4f}, device {device})")

    if scale <= 1.0001:
        out_img = img.resize((out_w, out_h), Image.LANCZOS)
    else:
        model = load_model(args.model, args.repo_id, device=device)
        lr = pil_to_tensor(img).to(device)
        with torch.autocast(device_type=device.type, enabled=bool(args.amp and device.type == "cuda")):
            out = model(lr, out_h, out_w)
        out_img = tensor_to_pil(out.squeeze(0))

    if canvas_w != out_w or canvas_h != out_h:
        canvas = Image.new("RGB", (canvas_w, canvas_h), (0, 0, 0))
        canvas.paste(out_img, ((canvas_w - out_w) // 2, (canvas_h - out_h) // 2))
        out_img = canvas

    suffix = Path(args.output).suffix.lower()
    kwargs = {}
    dpi = img.info.get("dpi")
    if dpi:
        kwargs["dpi"] = dpi
    if suffix in (".jpg", ".jpeg"):
        kwargs.update(quality=95, subsampling=0)
    elif suffix == ".webp":
        kwargs.update(quality=100, lossless=bool(args.webp_lossless))
    elif suffix in (".tif", ".tiff"):
        kwargs.update(compression="tiff_lzw")

    out_img.save(args.output, **kwargs)
    print(f"Saved: {args.output}")


if __name__ == "__main__":
    main()