#!/usr/bin/env python3
"""Upscale di un'immagine con lo Student 9x distillato (ONNX) alla scala richiesta.

Lo Student esportato dal notebook Colab e` nativamente 9x (staged upsample 3x3).
Per una scala diversa il modello non viene riaddestrato: si lavora sull'output 9x

  * scala < 9  -> inferenza a 9x, poi downsampling d'area sulla scala richiesta;
  * scala > 9  -> inferenza a 9x, poi upsampling bicubico.

L'inferenza e` a tile in spazio LR (tile e pad si intendono in pixel LR), come nella
cella "Usare il modello esportato" del notebook. L'eventuale canale alfa dei PNG non
attraversa la rete, che e` a 3 canali: viene ricampionato a parte e riaccodato.

Esempio:
    python scripts/upscale_onnx.py a.png --scale 6
"""
from __future__ import annotations

import argparse
import time
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort


def build_session(onnx_path: Path, providers_mode: str) -> ort.InferenceSession:
    """Session ONNX con arena disattivata: senza, la RAM cresce a ogni tile.

    Con l'inferenza a tile le shape d'input cambiano da tile a tile; il pattern
    allocator di ONNX Runtime trattiene i buffer del tile piu` grande e la RSS sale
    run dopo run fino a esaurire la memoria (processo ucciso senza traceback).
    """
    options = ort.SessionOptions()
    options.enable_cpu_mem_arena = False
    options.enable_mem_pattern = False
    options.log_severity_level = 3
    providers = (["CUDAExecutionProvider", "CPUExecutionProvider"]
                 if providers_mode == "auto" else ["CPUExecutionProvider"])
    return ort.InferenceSession(str(onnx_path), sess_options=options, providers=providers)


def load_image(path: Path) -> tuple[np.ndarray, np.ndarray | None]:
    """Carica l'immagine in BGR separando l'alpha, se presente."""
    img = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
    if img is None:
        raise SystemExit(f"Impossibile leggere {path}")
    if img.dtype != np.uint8:
        img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    if img.ndim == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    alpha = None
    if img.shape[2] == 4:
        img, alpha = img[:, :, :3].copy(), img[:, :, 3].copy()
    return img, alpha


def native_scale(session: ort.InferenceSession) -> int:
    """Deduce la scala nativa del modello da una probe 8x9 (evita di fissarla a mano)."""
    name = session.get_inputs()[0].name
    out = session.run(None, {name: np.zeros((1, 3, 8, 9), dtype=np.float32)})[0]
    ratio = max(out.shape[-2] / 8.0, out.shape[-1] / 9.0)
    if abs(ratio - round(ratio)) > 1e-6 or ratio < 1:
        raise SystemExit(f"Scala nativa non intera: input 8x9 -> output {out.shape}")
    return int(round(ratio))


def upscale_tiled(session, bgr: np.ndarray, scale: float, native: int,
                  tile: int, pad: int) -> np.ndarray:
    """Inferenza a tile in spazio LR; ogni tile viene ricampionato native -> scale."""
    input_name = session.get_inputs()[0].name
    height, width = bgr.shape[:2]
    out_h, out_w = int(round(height * scale)), int(round(width * scale))
    canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8)
    interp = cv2.INTER_AREA if native > scale else cv2.INTER_CUBIC
    tiles = 0
    t0 = time.time()
    for top in range(0, height, tile):
        for left in range(0, width, tile):
            bottom, right = min(top + tile, height), min(left + tile, width)
            # Tile + bordo di padding: il padding viene usato dall'interpolazione e
            # poi tagliato, cosi` le giunzioni tra tile non restano visibili.
            y0, y1 = max(top - pad, 0), min(bottom + pad, height)
            x0, x1 = max(left - pad, 0), min(right + pad, width)
            patch = bgr[y0:y1, x0:x1]
            tensor = (cv2.cvtColor(patch, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0)
            tensor = np.ascontiguousarray(tensor.transpose(2, 0, 1)[None, :, :, :])
            sr = np.clip(session.run(None, {input_name: tensor})[0][0], 0.0, 1.0)
            sr = cv2.cvtColor((sr.transpose(1, 2, 0) * 255.0).round().astype(np.uint8),
                              cv2.COLOR_RGB2BGR)
            # Il tile paddato va ricampionato alla SUA dimensione, non a quella del
            # ritaglio interno: cosi` l'offset del ritaglio resta dentro i confini.
            full_h, full_w = int(round((y1 - y0) * scale)), int(round((x1 - x0) * scale))
            sr = cv2.resize(sr, (full_w, full_h), interpolation=interp)
            th, tw = int(round((bottom - top) * scale)), int(round((right - left) * scale))
            oy = min(int(round((top - y0) * scale)), max(0, full_h - th))
            ox = min(int(round((left - x0) * scale)), max(0, full_w - tw))
            th = min(th, full_h - oy)
            tw = min(tw, full_w - ox)
            oy_px, ox_px = int(round(top * scale)), int(round(left * scale))
            th = min(th, out_h - oy_px)
            tw = min(tw, out_w - ox_px)
            canvas[oy_px:oy_px + th, ox_px:ox_px + tw] = sr[oy:oy + th, ox:ox + tw]
            tiles += 1
            print(f"  tile {tiles}: LR {y1 - y0}x{x1 - x0} -> {th}x{tw} "
                  f"[{time.time() - t0:.1f}s]", flush=True)
    return canvas


def psnr(a: np.ndarray, b: np.ndarray) -> float:
    mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
    return 99.0 if mse == 0 else float(10.0 * np.log10(255.0 ** 2 / mse))


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("input", help="immagine da ingrandire")
    ap.add_argument("--onnx", default="student_sr9x.onnx", help="ONNX dello Student distillato")
    ap.add_argument("--scale", type=float, default=6.0, help="scala richiesta (il modello e` 9x)")
    ap.add_argument("--output", default=None, help="percorso output (default output_<nome>_<scala>x.png)")
    ap.add_argument("--tile", type=int, default=64, help="lato tile in pixel LR (0 = immagine intera)")
    ap.add_argument("--pad", type=int, default=6, help="overlap di padding in pixel LR")
    ap.add_argument("--alpha-interp", choices=["nearest", "linear"], default="nearest",
                    help="ricampionamento dell'alpha (nearest evita aloni sui bordi)")
    ap.add_argument("--no-baseline", action="store_true", help="non generare il confronto bicubico")
    ap.add_argument("--providers", choices=["auto", "cpu"], default="auto")
    args = ap.parse_args()

    src = Path(args.input)
    if not src.is_file():
        raise SystemExit(f"File non trovato: {src}")
    onnx_path = Path(args.onnx)
    if not onnx_path.is_file():
        raise SystemExit(f"ONNX non trovato: {onnx_path}")
    if args.scale <= 0:
        raise SystemExit("--scale deve essere > 0")

    session = build_session(onnx_path, args.providers)
    native = native_scale(session)
    bgr, alpha = load_image(src)
    height, width = bgr.shape[:2]
    out_h, out_w = int(round(height * args.scale)), int(round(width * args.scale))
    tile = height if args.tile <= 0 else min(int(args.tile), height)
    print(f"[modello] {onnx_path} | scala nativa {native}x | tile {tile}+pad {args.pad}")
    print(f"[input]   {src} {width}x{height}"
          + (" (con alpha)" if alpha is not None else ""))
    print(f"[output]  {out_w}x{out_h} ({args.scale}x)")

    result = upscale_tiled(session, bgr, args.scale, native, tile, int(args.pad))

    out_path = Path(args.output) if args.output else src.with_name(
        f"output_{src.stem}_{args.scale:g}x.png")
    if alpha is not None:
        interp = cv2.INTER_NEAREST if args.alpha_interp == "nearest" else cv2.INTER_LINEAR
        alpha_out = cv2.resize(alpha, (out_w, out_h), interpolation=interp)
        result = np.dstack([result, alpha_out])  # BGR + alpha = BGRA, da salvare cosi` com'e`
    if not cv2.imwrite(str(out_path), result):
        raise SystemExit(f"Scrittura fallita: {out_path}")
    print(f"[salvato] {out_path}")

    if not args.no_baseline:
        interp = cv2.INTER_AREA if native > args.scale else cv2.INTER_CUBIC
        bicubic = cv2.resize(bgr, (out_w, out_h), interpolation=interp)
        bic_path = out_path.with_name(f"{out_path.stem}_bicubic{out_path.suffix}")
        if alpha is not None:
            bicubic = np.dstack([bicubic, cv2.resize(alpha, (out_w, out_h),
                                                     interpolation=cv2.INTER_NEAREST)])
        cv2.imwrite(str(bic_path), bicubic)
        print(f"[baseline] {bic_path} | PSNR modello vs bicubica "
              f"{psnr(result if result.shape[2] == 3 else result[:, :, :3], bicubic[:, :, :3]):.2f} dB")


if __name__ == "__main__":
    main()