"""iPhoneBlur Restore on ZeroGPU.

Full resolution, no lossy step anywhere in the pipeline: the photo is decoded, restored
in overlapping 512 px tiles and handed back as PNG.

Tiles are kept at 512 px on purpose. The channel attention averages over whatever it is
given, so feeding a whole photo changes those statistics away from what the model was
trained on. Measured on the test split, 512 px tiling scored 30.25 dB against 29.59 dB
for whole-image inference, and on one image the gap was 8 dB. Tiling is a quality
decision here, not only a memory one.
"""
import os
import re
import tempfile
import threading
import time

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

from restore_net import RestoreNet

try:                      # iPhones shoot HEIC by default, Pillow cannot read it unaided
    import pillow_heif
    pillow_heif.register_heif_opener()
    HEIF = True
except Exception:
    HEIF = False

try:
    import spaces
except ImportError:                                   # runs locally without ZeroGPU too
    class _Stub:
        @staticmethod
        def GPU(*a, **kw):
            if a and callable(a[0]):
                return a[0]

            def wrap(fn):
                return fn
            return wrap
    spaces = _Stub()

REPO = os.environ.get("MODEL_REPO", "Shafi99/iPhoneBlur_Finetune_Models")
CKPT = os.environ.get("MODEL_FILE", "nafnet_final_weights.pth")
TILE, OVERLAP = 512, 48
BATCH = 8
MAX_PIXELS = 40_000_000       # refuse absurd inputs rather than hang
Image.MAX_IMAGE_PIXELS = None


def _infer_config(sd):
    width = sd["intro.weight"].shape[0]

    def count(prefix):
        idx = {}
        for k in sd:
            m = re.match(r"%s\.(\d+)\.(\d+)\." % prefix, k)
            if m:
                idx.setdefault(int(m.group(1)), set()).add(int(m.group(2)))
        return tuple(len(idx[i]) for i in sorted(idx))

    mid = {int(m.group(1)) for k in sd for m in [re.match(r"middle_blks\.(\d+)\.", k)] if m}
    return dict(img_channel=sd["intro.weight"].shape[1], width=width,
                enc_blk_nums=count("encoders"), dec_blk_nums=count("decoders"),
                middle_blk_num=len(mid))


def _load():
    path = hf_hub_download(REPO, CKPT)
    raw = torch.load(path, map_location="cpu", weights_only=False)
    if isinstance(raw, dict):
        for key in ("params", "params_ema", "model_state_dict", "state_dict", "model", "net_g"):
            if key in raw and isinstance(raw[key], dict):
                raw = raw[key]
                break
    sd = {}
    for k, v in raw.items():
        for prefix in ("module.", "model.", "net_g."):
            if k.startswith(prefix):
                k = k[len(prefix):]
        sd[k] = v
    net = RestoreNet(**_infer_config(sd))
    net.load_state_dict(sd)          # strict: the layout must match exactly
    return net.eval()


# ZeroGPU wants the model placed on cuda at module level, not inside the GPU function.
model = _load()
if torch.cuda.is_available() or os.environ.get("SPACES_ZERO_GPU"):
    model = model.to("cuda")

# Second copy, kept on the CPU, built only if the GPU allowance runs out. Measured cost
# is roughly 54 s per 512 px tile on one core and 24 s on two, so the fallback works at a
# reduced size: a full 1920x1080 photo would be 15 tiles and several minutes.
CPU_MAX_SIDE = int(os.environ.get("CPU_MAX_SIDE", "1024"))
_cpu_model = None
_cpu_lock = threading.Lock()


def _cpu_net():
    global _cpu_model
    with _cpu_lock:
        if _cpu_model is None:
            torch.set_num_threads(max(1, os.cpu_count() or 1))
            _cpu_model = _load()
    return _cpu_model


def _window():
    ramp = np.minimum(1.0, (np.arange(TILE) + 1) / OVERLAP)
    ramp = np.minimum(ramp, ramp[::-1])
    return np.outer(ramp, ramp).astype(np.float32)


def _grid(n):
    """Padded length and tile origins along one axis.

    Deliberately identical to the server pipeline in api/_core.py, including padding the
    canvas up to a whole number of strides. A simpler placement that skips that padding
    measured up to 0.28 dB worse on the test split, so the two must not drift apart.
    """
    stride = TILE - 2 * OVERLAP
    padded = max(TILE, n + (-(n - TILE) % stride if n > TILE else 0))
    pos = list(range(0, padded - TILE + 1, stride))
    if pos[-1] != padded - TILE:
        pos.append(padded - TILE)
    return padded, pos


def _tile_count(w, h):
    return len(_grid(w)[1]) * len(_grid(h)[1])


def _restore(arr, device, net=None):
    """arr: float32 HxWx3 in [0,1]. Overlapping 512 tiles, linear ramp blend."""
    h, w, _ = arr.shape
    ph, ys = _grid(h)
    pw, xs = _grid(w)
    if ph > h or pw > w:
        arr = np.pad(arr, ((0, ph - h), (0, pw - w), (0, 0)), mode="reflect")
    out = np.zeros((ph, pw, 3), np.float32)
    acc = np.zeros((ph, pw, 1), np.float32)
    win = _window()[..., None]
    coords = [(y, x) for y in ys for x in xs]
    for i in range(0, len(coords), BATCH):
        chunk = coords[i:i + BATCH]
        batch = np.stack([arr[y:y + TILE, x:x + TILE].transpose(2, 0, 1) for y, x in chunk])
        with torch.no_grad():
            res = (net or model)(torch.from_numpy(batch).to(device)).float().cpu().numpy()
        for (y, x), r in zip(chunk, res):
            out[y:y + TILE, x:x + TILE] += r.transpose(1, 2, 0) * win
            acc[y:y + TILE, x:x + TILE] += win
    return np.clip(out[:h, :w] / np.maximum(acc[:h, :w], 1e-6), 0.0, 1.0)


def _duration(image):
    """Ask for as little GPU time as is safe.

    Quota is charged per visitor and a shorter request also earns better queue
    priority, so this tracks the measured cost rather than padding it: about
    0.15 s per 512 px tile once the model is resident, plus a fixed allowance for
    getting it there.
    """
    if image is None:
        return 25
    n = _tile_count(*getattr(image, "size", (512, 512)))
    return int(min(180, 20 + n * 0.4))


@spaces.GPU(duration=_duration)
def _run(image):
    """The GPU half. Everything else stays off the GPU clock."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    arr = np.asarray(image, np.float32) / 255.0
    return Image.fromarray((_restore(arr, device) * 255.0 + 0.5).astype(np.uint8))


def _run_on_cpu(image):
    """Slow path, used only when the GPU allowance is exhausted.

    Shrunk to CPU_MAX_SIDE first, because the cost is linear in tiles and a full
    resolution photo would take several minutes on these cores.
    """
    k = min(1.0, float(CPU_MAX_SIDE) / max(image.size))
    small = image if k >= 1.0 else image.resize(
        (max(1, round(image.width * k)), max(1, round(image.height * k))), Image.LANCZOS)
    arr = np.asarray(small, np.float32) / 255.0
    out = _restore(arr, "cpu", _cpu_net())
    return Image.fromarray((out * 255.0 + 0.5).astype(np.uint8)), small.size


def _out_of_quota(err):
    t = str(err).lower()
    return ("quota" in t or "exceeded" in t or "gpu task aborted" in t
            or "zerogpu" in t or "no gpu" in t)


def _restore_any(image):
    """GPU first, CPU if the allowance is gone. Returns (image, how, size)."""
    try:
        return _run(image), "gpu", image.size
    except Exception as e:
        if not _out_of_quota(e):
            raise
        result, size = _run_on_cpu(image)
        return result, "cpu", size


def _as_image(x):
    """Accept what either caller sends: a PIL image from the page, or the file
    descriptor dict for an uploaded file. Any decode failure becomes a sentence the
    person can act on rather than a stack trace."""
    if isinstance(x, Image.Image):
        return x
    if isinstance(x, dict):
        x = x.get("path") or x.get("url")
    if not isinstance(x, (str, bytes, os.PathLike)):
        raise gr.Error("Could not read that file as an image.")
    try:
        im = Image.open(x)
        im.load()
        return im
    except Exception:
        ext = str(os.path.splitext(str(x))[1]).lower()
        if ext in (".heic", ".heif") and not HEIF:
            raise gr.Error("That is a HEIC photo and this server cannot open it. "
                           "On iPhone: Settings, Camera, Formats, Most Compatible.")
        raise gr.Error("That file could not be opened as an image. It may be damaged "
                       "or in a format this server does not support.")


def _check(image):
    if image is None:
        raise gr.Error("Choose a photo first.")
    image = _as_image(image)
    if image.width * image.height > MAX_PIXELS:
        raise gr.Error("That photo is larger than 40 megapixels. Please resize it a little.")
    return image.convert("RGB")


def _outdir():
    """A fresh directory per request. A fixed filename would let two visitors
    overwrite each other's result."""
    return tempfile.mkdtemp(prefix="restore_")


def restore(image):
    """The built-in page: both halves for the slider, plus a lossless download."""
    image = _check(image)
    t0 = time.time()
    try:
        result, how, size = _restore_any(image)
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(_friendly(e))
    dt = time.time() - t0
    path = os.path.join(_outdir(), "restored.png")
    result.save(path, "PNG", optimize=False)     # lossless, nothing added after the model
    n = _tile_count(*size)
    note = "%d by %d px, %d tiles, %.1f s on the %s" % (size[0], size[1], n, dt, how.upper())
    if how == "cpu":
        note += ". The shared GPU was busy, so this ran on the processor at reduced size."
    return (image.resize(size) if size != image.size else image, result), path, note


def _friendly(e):
    t = str(e).lower()
    if "memory" in t or "alloc" in t:
        return "That photo was too large to process. Try one under about 40 megapixels."
    if "cuda" in t or "gpu" in t:
        return "The shared GPU could not take this one. Try again in a moment."
    return "Something went wrong while restoring that photo. Try another one."


def restore_api(image):
    """Lean endpoint for a standalone front end.

    Returns two files rather than one: a small preview to show straight away, and the
    lossless PNG. Both come back as URLs, so the browser only pulls the big one if the
    visitor actually clicks download. Measured on a 0.9 Mbps line, the PNG alone took
    23.6 s to transfer while the GPU work took 4.2 s, so the preview is what makes this
    usable on a slow connection. Nothing is re-compressed in the file people keep.
    """
    try:
        result, how, _size = _restore_any(_check(image))
    except gr.Error:
        raise
    except Exception as e:                      # never let a raw traceback reach the page
        raise gr.Error(_friendly(e))
    d = _outdir()
    png = os.path.join(d, "restored.png")
    preview = os.path.join(d, "preview.webp")
    result.save(png, "PNG", optimize=False)
    result.save(preview, "WEBP", quality=92, method=4)
    return preview, png


CSS = ".contain{max-width:1100px!important}"

with gr.Blocks(title="iPhoneBlur Restore") as demo:
    gr.Markdown(
        "# iPhoneBlur Restore\n"
        "Sharpen photos spoiled by motion blur. Processed at full resolution and returned "
        "as PNG, so nothing is re-compressed after restoring. Drag the line to compare."
    )
    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.Image(type="pil", label="Blurry photo", sources=["upload", "clipboard"],
                           image_mode="RGB", format="png", height=320)
            go = gr.Button("Restore", variant="primary")
            info = gr.Markdown("")
            dl = gr.DownloadButton("Download full resolution PNG", visible=True)
        with gr.Column(scale=2):
            out = gr.ImageSlider(label="Before and after", type="pil", format="png",
                                 slider_position=50, height=520)
    gr.Markdown(
        "Built from the iPhoneBlur benchmark of real handheld phone blur. "
        "Trained on pairs from one device, so other cameras and other kinds of blur "
        "(out of focus, very dark, heavy noise) may not improve."
    )
    go.click(restore, inputs=inp, outputs=[out, dl, info], show_progress="full")

    # Same model, lean payload, for the standalone page hosted elsewhere. Real
    # components rather than gr.api, because those serialise the result as a PNG file
    # the caller can fetch; gr.api returned the repr of the image object instead.
    with gr.Row(visible=False):
        api_in = gr.Image(type="pil", image_mode="RGB", format="png")
        api_preview = gr.File()      # gr.File passes bytes through untouched,
        api_full = gr.File()         # gr.Image would re-encode and break "lossless"
        api_btn = gr.Button("api")
    api_btn.click(restore_api, inputs=api_in, outputs=[api_preview, api_full],
                  api_name="restore_image")

if __name__ == "__main__":
    demo.queue(max_size=20).launch(css=CSS, theme=gr.themes.Soft(), show_error=True)