"""
İvme-Denoise-v1 (İvmetron-Guanine1) — 2.97 M-parameter blind image denoiser.

Demo Space for IvmeLabs/Ivme-Denoise-v1: a tiny Bayesian-shrinkage U-Net
that beats DnCNN on every standard benchmark at 65 % of its FLOPs.
"""
import os

import spaces  # MUST come before torch
import gradio as gr
import numpy as np
import torch  # noqa: F401  (imported after spaces on purpose)
from PIL import Image

from modeling_ivmetron import load_model, denoise_tiled, estimate_sigma_map

MODEL_ID = "IvmeLabs/Ivme-Denoise-v1"

# Model is tiny (12 MB) — load eagerly at module scope, on the GPU.
model = load_model(MODEL_ID, device="cuda")

MAX_SIDE = 3072  # cap gigantic uploads; up to this size the tiled path stays fast

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""


def _to_float_array(pil_img):
    """PIL image -> HxWx3 float32 in [0,1], downscaled if huge."""
    img = pil_img.convert("RGB")
    w, h = img.size
    if max(w, h) > MAX_SIDE:
        scale = MAX_SIDE / max(w, h)
        img = img.resize(
            (max(1, int(round(w * scale))), max(1, int(round(h * scale)))),
            Image.LANCZOS,
        )
    return np.asarray(img, dtype=np.float32) / 255.0


@spaces.GPU(duration=10)
def denoise_image(image, sigma_mode, sigma_value, progress=gr.Progress()):
    """Denoise an sRGB photograph with İvme-Denoise-v1.

    Args:
        image: the noisy input photograph (sRGB, any resolution).
        sigma_mode: "Auto (estimate per-pixel sigma map)" or "Manual sigma".
        sigma_value: manual noise level on the 0-255 scale (5-50), used only
            when sigma_mode is "Manual sigma".
    Returns:
        The denoised image.
    """
    if image is None:
        raise gr.Error("Please upload an image first.")
    progress(0.1, desc="Preparing image")
    img = _to_float_array(image)

    if str(sigma_mode).startswith("Auto"):
        progress(0.3, desc="Estimating noise map")
        sigma = estimate_sigma_map(img)
        mode_note = f"auto σ ≈ {float(np.asarray(sigma).mean()):.1f}"
    else:
        sigma = float(sigma_value)
        mode_note = f"manual σ = {sigma:.0f}"

    progress(0.5, desc="Denoising")
    out = denoise_tiled(model, img, sigma, device="cuda")

    progress(1.0, desc=f"Done ({mode_note})")
    return Image.fromarray((np.clip(out, 0, 1) * 255).round().astype(np.uint8))


with gr.Blocks() as demo:
    gr.Markdown(
        f"""
        # İvme-Denoise-v1 · *Margay 1*
        **2.97 M-parameter blind image denoiser** (İvmetron-Guanine1) — every
        non-linearity is the exact posterior mean of a Gaussian scale mixture.
        Beats DnCNN on every standard benchmark using **35 % fewer FLOPs**, within
        0.22–0.54 dB of DRUNet / Restormer at a fraction of their size.
        \n
        Model: [{MODEL_ID}](https://huggingface.co/{MODEL_ID}) · Apache-2.0 ·
        12 MB weights · 17.9 ms per 512² frame
        """
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(
                    type="pil", label="Noisy input", height=340,
                    sources=["upload", "clipboard"],
                )
                sigma_mode = gr.Radio(
                    ["Auto (estimate per-pixel sigma map)", "Manual sigma"],
                    value="Auto (estimate per-pixel sigma map)",
                    label="Noise level",
                )
                sigma_value = gr.Slider(
                    5, 50, value=15, step=1,
                    label="Manual σ (0–255 scale)",
                    info="Only used in Manual mode. The model is trained for σ ∈ [5, 50].",
                )
                run_btn = gr.Button("Denoise", variant="primary")
            with gr.Column():
                image_out = gr.Image(type="pil", label="Denoised output", height=340)

        gr.Examples(
            examples=[
                ["examples/park_original.png",
                 "Auto (estimate per-pixel sigma map)", 15],
                ["examples/cat_noisy_sigma15.png", "Manual sigma", 15],
                ["examples/owl_noisy_sigma25.png", "Manual sigma", 25],
                ["examples/night_city_noisy_sigma50.png", "Manual sigma", 50],
            ],
            inputs=[image_in, sigma_mode, sigma_value],
            outputs=[image_out],
            fn=denoise_image,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.Markdown(
            """
            **How to use** — upload any sRGB photo (real high-ISO camera output works
            best). Leave the noise level on *Auto* to let the built-in robust
            estimator produce a per-pixel σ map; or switch to *Manual* and set σ
            yourself (the benchmark convention: 15 / 25 / 50 on the 0–255 scale).

            **Notes** — sRGB only (no linear RAW/HDR). Huge uploads are
            downscaled to 3072 px on the long side. Images above ~5 MP are
            denoised with tiled inference (896-px tiles, 96-px overlap, Hann
            blending) exactly as in the reference implementation.

            *Example images: the park frame is from the model card
            (Apache-2.0); the others are CC0 images with synthetic Gaussian
            noise at the model's benchmark σ levels.*
            """
        )

    run_btn.click(
        denoise_image,
        inputs=[image_in, sigma_mode, sigma_value],
        outputs=[image_out],
        api_name="denoise",
        concurrency_limit=1,
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)