"""Gradio Space demo for the Universal Upscale model."""
from pathlib import Path

import gradio_client.utils as gutils

# Monkey-patch gradio_client schema parser to handle boolean additionalProperties in Gradio 5.
_orig_json_schema = gutils._json_schema_to_python_type


def _safe_json_schema_to_python_type(schema, defs=None):
    if isinstance(schema, bool):
        return "Any" if schema else "Never"
    return _orig_json_schema(schema, defs)


gutils._json_schema_to_python_type = _safe_json_schema_to_python_type

import gradio as gr
import numpy as np
import torch
from PIL import Image

try:
    import spaces
    gpu_decorator = spaces.GPU
except ImportError:
    def gpu_decorator(fn):
        return fn

MODEL_PATH = Path(__file__).resolve().parent / "model_scripted.pt"


def load_model(device: str):
    model = torch.jit.load(str(MODEL_PATH), map_location=device)
    model.eval()
    return model


MAX_SIDE = 4096
# Largest scale factor verified end-to-end for the published checkpoint. The slider is capped
# here; "Target Long Side" still accepts any size up to MAX_SIDE.
MAX_SCALE = 9.2


def resolve_factor(input_size, scale, target_long_side):
    """Scale factor to apply: an explicit target size wins, scale requests stay under MAX_SCALE."""
    target = float(target_long_side) if target_long_side is not None else 0.0
    if target > 0:
        return target / float(max(input_size)), ""
    factor = float(scale if scale is not None else 4.0)
    if factor > MAX_SCALE:
        return MAX_SCALE, f" (clamped to the tested maximum {MAX_SCALE:g}x)"
    return factor, ""


@gpu_decorator
def upscale(image, scale: float = 4.0, target_long_side: float = 0.0):
    if image is None:
        return None, "No image provided."
    img = image if isinstance(image, Image.Image) else Image.open(image)
    img = img.convert("RGB")

    factor, clamped = resolve_factor(img.size, scale, target_long_side)

    out_w, out_h = max(1, int(round(img.width * factor))), max(1, int(round(img.height * factor)))
    if max(out_w, out_h) > MAX_SIDE:
        shrink = MAX_SIDE / float(max(out_w, out_h))
        out_w, out_h = max(1, int(out_w * shrink)), max(1, int(out_h * shrink))

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_model(device)

    lr = torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0).permute(2, 0, 1).contiguous().unsqueeze(0).to(device)
    with torch.no_grad():
        out = model(lr, out_h, out_w)
    arr = (out.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0).round().astype(np.uint8)
    res_img = Image.fromarray(arr)
    status = f"Input: {img.width}x{img.height} | Output: {res_img.width}x{res_img.height} (Scale: {factor:.2f}x) | Device: {device}{clamped}"
    return res_img, status


example_files = [
    [str(Path(__file__).resolve().parent / name)]
    for name in ("examples/input.png", "examples/demo_input.png")
    if (Path(__file__).resolve().parent / name).exists()
]

demo = gr.Interface(
    fn=upscale,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Slider(minimum=1.0, maximum=9.2, value=4.0, step=0.2, label="Scale Factor"),
        gr.Number(value=0, label="Or Target Long Side (px) [0 = use scale factor]"),
    ],
    outputs=[
        gr.Image(type="pil", label="Upscaled Image"),
        gr.Textbox(label="Status / Image Info"),
    ],
    title="Universal Upscale",
    description=(
        "Arbitrary-scale super-resolution: upscale images to any scale factor from 1x up to the tested maximum "
        "9.2x, or to an exact target resolution. Built-in tiled inference enables high-resolution "
        "upscaling without memory issues (outputs are capped at 4096 px on the long side)."
    ),
    examples=example_files if example_files else None,
)

if __name__ == "__main__":
    demo.launch()