"""NanoVSR — real-time video super-resolution, interactive Gradio demo.

Paper: "NanoVSR: Towards Real-Time Video Super-Resolution on Edge Devices" (ECCV 2026)
Code:  https://github.com/filippawlicki/nanovsr
"""
import os
import tempfile
import time

import cv2
import gradio as gr
import imageio.v2 as imageio
import numpy as np
import torch

try:
    import spaces
except ImportError:
    class _NoOpGPU:
        def __call__(self, fn=None, **kwargs):
            if fn is not None:
                return fn
            return lambda f: f

    class spaces:
        GPU = _NoOpGPU()

from inference import (
    DEFAULT_MODEL,
    MODEL_INFO,
    SCALE,
    get_model,
    run_chunked_inference,
)

ASSETS_DIR = os.path.join(os.path.dirname(__file__), "assets")
EXAMPLE_VIDEO = os.path.join(ASSETS_DIR, "examples", "reds4_000_lr.mp4")
TEASER_VIDEO = os.path.join(ASSETS_DIR, "teaser.mp4")

gr.set_static_paths([ASSETS_DIR])

MAX_SIDE_CHOICES = [180, 240, 320, 480]
DIVIDER_PX = 8


def _even(n):
    return n - (n % 2)


def _read_video(path, max_frames):
    reader = imageio.get_reader(path)
    meta = reader.get_meta_data()
    fps = meta.get("fps", 25.0) or 25.0

    frames = []
    for i, frame in enumerate(reader):
        if max_frames is not None and i >= max_frames:
            break
        if frame.shape[-1] == 4:
            frame = frame[..., :3]
        frames.append(frame)
    reader.close()

    if not frames:
        raise gr.Error("Could not read any frames from the input video.")
    return frames, float(fps)


def _resize_to_cap(frames, max_side):
    h, w = frames[0].shape[:2]
    longest = max(h, w)
    if longest <= max_side:
        new_h, new_w = _even(h), _even(w)
    else:
        scale = max_side / longest
        new_h, new_w = _even(int(round(h * scale))), _even(int(round(w * scale)))
    new_h, new_w = max(new_h, 8), max(new_w, 8)

    if (new_h, new_w) == (h, w):
        return frames, (h, w)
    resized = [cv2.resize(f, (new_w, new_h), interpolation=cv2.INTER_AREA) for f in frames]
    return resized, (h, w)


def _label(img, text):
    img = np.ascontiguousarray(img)
    scale = max(0.5, img.shape[0] / 720.0)
    thickness = max(1, int(round(2 * scale)))
    origin = (int(16 * scale), int(34 * scale))
    cv2.putText(img, text, origin, cv2.FONT_HERSHEY_SIMPLEX, scale, (0, 0, 0), thickness * 3, cv2.LINE_AA)
    cv2.putText(img, text, origin, cv2.FONT_HERSHEY_SIMPLEX, scale, (255, 255, 255), thickness, cv2.LINE_AA)
    return img


def _make_comparison_frame(lq_rgb, sr_rgb):
    h, w = sr_rgb.shape[:2]
    base = cv2.resize(lq_rgb, (w, h), interpolation=cv2.INTER_NEAREST)
    base = _label(base, "Input (LR, nearest-upscaled)")
    sr = _label(np.ascontiguousarray(sr_rgb), "NanoVSR")
    divider = np.full((h, DIVIDER_PX, 3), 255, dtype=np.uint8)
    return np.hstack([base, divider, sr])


def _estimate_duration(model_name, frames, chunk_size, use_fp16, progress=None):
    del use_fp16, progress
    n = len(frames)
    factor = {"NanoVSR-226k (fastest)": 0.15, "NanoVSR-644k (baseline)": 0.25,
              "NanoVSR-1.7M": 0.4, "NanoVSR-5.4M (best quality)": 0.8}.get(model_name, 0.3)
    return int(min(180, max(25, 20 + n * factor)))


@spaces.GPU(duration=_estimate_duration)
def _gpu_infer(model_name, frames, chunk_size, use_fp16, progress=gr.Progress()):
    model = get_model(model_name)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    def cb(done, total):
        progress(done / total, desc=f"Upscaling frames ({done}/{total}) on {device.upper()}")

    progress(0, desc="Starting inference...")
    return run_chunked_inference(model, frames, chunk_size, device, use_fp16, progress_cb=cb)


def upscale(video_path, model_name, chunk_size, max_frames, max_side, use_fp16,
            progress=gr.Progress()):
    if not video_path:
        raise gr.Error("Please upload or select a low-resolution video first.")

    progress(0, desc="Reading video...")
    frames, fps = _read_video(video_path, max_frames)
    frames, orig_size = _resize_to_cap(frames, max_side)
    in_h, in_w = frames[0].shape[:2]

    t0 = time.time()
    sr_frames = _gpu_infer(model_name, frames, int(chunk_size), bool(use_fp16), progress)
    elapsed = time.time() - t0

    progress(0.95, desc="Encoding output video...")
    sr_h, sr_w = sr_frames.shape[1:3]

    out_frames = [_make_comparison_frame(frames[i], sr_frames[i]) for i in range(len(frames))]
    preview_before = cv2.resize(frames[0], (sr_w, sr_h), interpolation=cv2.INTER_NEAREST)

    out_fd, out_path = tempfile.mkstemp(suffix=f"_nanovsr_x{SCALE}.mp4")
    os.close(out_fd)
    out_fps = max(1.0, min(fps, 60.0))
    writer = imageio.get_writer(out_path, fps=out_fps, codec="libx264",
                                 ffmpeg_params=["-crf", "18"], macro_block_size=1)
    for f in out_frames:
        writer.append_data(f)
    writer.close()

    resize_note = (f" (downscaled from {orig_size[1]}x{orig_size[0]} to fit the "
                    f"{max_side}px cap)" if (orig_size[1], orig_size[0]) != (in_w, in_h) else "")
    status = (
        f"**Model:** {model_name} · **{MODEL_INFO[model_name]['psnr']}**\n\n"
        f"**Resolution:** {in_w}x{in_h}{resize_note} → {sr_w}x{sr_h} (x{SCALE})\n\n"
        f"**Frames processed:** {len(frames)} · **Inference time:** {elapsed:.2f}s "
        f"({len(frames) / max(elapsed, 1e-6):.1f} FPS)"
    )
    preview = (preview_before, np.ascontiguousarray(sr_frames[0]))
    return out_path, preview, status


THEME = gr.themes.Soft(primary_hue="blue", secondary_hue="cyan")

CSS = """
.status-box { font-size: 0.95em; }
.badge-row { display: flex !important; gap: 8px; justify-content: center; flex-wrap: wrap; margin: 0; }
.badge-row a, .badge-row img { display: inline-block !important; margin: 0 !important; }
"""

with gr.Blocks(title="NanoVSR: Towards Real-Time Video Super-Resolution on Edge Devices") as demo:
    gr.Markdown(
        """
        <div align="center">

        # NanoVSR — Real-Time 4× Video Super-Resolution on Edge Devices

        <p class="badge-row" style="display:flex; gap:8px; justify-content:center; flex-wrap:wrap; margin:0;">
        <a href="https://arxiv.org/abs/2607.10495"><img src="https://img.shields.io/badge/arXiv-2607.10495-b31b1b" alt="arXiv"></a>
        <a href="https://github.com/filippawlicki/nanovsr"><img src="https://img.shields.io/badge/GitHub-nanovsr-181717?logo=github" alt="GitHub"></a>
        <a href="https://github.com/filippawlicki/nanovsr/blob/main/LICENSE"><img src="https://img.shields.io/badge/License-MIT-green.svg" alt="License"></a>
        </p>

        </div>

        A fully convolutional, bidirectional-recurrent video super-resolution model that reaches
        **27–44 FPS on a Jetson Orin NX** while upscaling 4×. Upload a low-resolution clip below and
        watch NanoVSR sharpen it in real time — or press one of the examples to try it instantly.
        """
    )
    gr.HTML(
        f'<video id="header-video" autoplay muted loop playsinline '
        f'style="width:100%; border-radius:12px; display:block;">'
        f'<source src="/gradio_api/file={TEASER_VIDEO}" type="video/mp4"></video>'
    )

    with gr.Row():
        with gr.Column(scale=1):
            video_in = gr.Video(label="Low-resolution input video", sources=["upload"])
            gr.Examples(
                examples=[[EXAMPLE_VIDEO]],
                inputs=[video_in],
                label="Example (REDS4 clip '000', the paper's own benchmark set)",
            )
            model_dd = gr.Dropdown(
                choices=list(MODEL_INFO.keys()), value=DEFAULT_MODEL, label="Model",
                info="Bigger models = higher quality, more compute.",
            )
            with gr.Accordion("Advanced options", open=False):
                chunk_size = gr.Slider(4, 30, value=15, step=1, label="Temporal chunk size (T)",
                                        info="Frames processed per forward pass. The paper's edge setting is 15.")
                max_frames = gr.Slider(8, 150, value=60, step=1, label="Max frames to process",
                                        info="Caps runtime; trim long videos to the first N frames.")
                max_side = gr.Dropdown(MAX_SIDE_CHOICES, value=320, label="Max input side (px)",
                                        info="Input is downscaled (never upscaled) to this cap before "
                                             "running NanoVSR, which expects genuinely low-res input.")
                fp16 = gr.Checkbox(value=True, label="FP16 inference (GPU only, faster)")
            run_btn = gr.Button("Upscale ▶", variant="primary")

        with gr.Column(scale=1):
            video_out = gr.Video(label="NanoVSR result", buttons=["download"])
            slider_out = gr.ImageSlider(label="Before / after (first frame)", buttons=[])
            status_out = gr.Markdown(elem_classes=["status-box"])

    run_btn.click(
        fn=upscale,
        inputs=[video_in, model_dd, chunk_size, max_frames, max_side, fp16],
        outputs=[video_out, slider_out, status_out],
    )

    with gr.Accordion("Model zoo & citation", open=False):
        model_list_md = "\n".join(f"- **{name}** — {info['psnr']}" for name, info in MODEL_INFO.items())
        citation_md = model_list_md + "\n\n" + (
"""All models perform 4× upscaling; larger models trade speed for quality. Full details,
training code and TensorRT deployment scripts are in the
[GitHub repo](https://github.com/filippawlicki/nanovsr).

```bibtex
@misc{pawlicki2026nanovsr,
      title={NanoVSR: Towards Real-Time Video Super-Resolution on Edge Devices},
      author={Filip Pawlicki and Marcel Kańduła and Marcin Pucek and Kamil Dobies},
      year={2026},
      eprint={2607.10495},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2607.10495},
}
```"""
        )
        gr.Markdown(citation_md)

if __name__ == "__main__":
    demo.queue().launch(theme=THEME, css=CSS)