import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# `~/.cache/torch` is not writable in the Space image and RAFT weights are fetched
# through torch.hub, so point TORCH_HOME (and matplotlib's config dir) somewhere writable.
os.environ.setdefault("TORCH_HOME", "/tmp/torch_home")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/mplconfig")
os.makedirs(os.path.join(os.environ["TORCH_HOME"], "hub", "checkpoints"), exist_ok=True)
os.makedirs(os.environ["MPLCONFIGDIR"], exist_ok=True)

import spaces  # noqa: E402  (must precede any CUDA-touching import)

import tempfile  # noqa: E402
import time  # noqa: E402
from typing import List, Tuple  # noqa: E402

import gradio as gr  # noqa: E402
import imageio.v3 as iio  # noqa: E402
import imageio_ffmpeg  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402
from torchvision.models.optical_flow import Raft_Large_Weights, raft_large  # noqa: E402

from diffusers import DDIMScheduler  # noqa: E402
from pipeline.stream_diffvsr_pipeline import (  # noqa: E402
    ControlNetModel,
    StreamDiffVSRPipeline,
    UNet2DConditionModel,
)
from temporal_autoencoder.autoencoder_tiny import TemporalAutoencoderTiny  # noqa: E402

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

MODEL_ID = "Jamichsu/Stream-DiffVSR"
SCALE = 4  # the model is a x4 super-resolver
MULTIPLE_OF = 8  # UNet needs the low-res side divisible by 8

# ---------------------------------------------------------------- model load
controlnet = ControlNetModel.from_pretrained(MODEL_ID, subfolder="controlnet")
unet = UNet2DConditionModel.from_pretrained(MODEL_ID, subfolder="unet")
vae = TemporalAutoencoderTiny.from_pretrained(MODEL_ID, subfolder="vae")
scheduler = DDIMScheduler.from_pretrained(MODEL_ID, subfolder="scheduler")

pipe = StreamDiffVSRPipeline.from_pretrained(
    MODEL_ID,
    controlnet=controlnet,
    unet=unet,
    vae=vae,
    scheduler=scheduler,
    safety_checker=None,
    feature_extractor=None,
    requires_safety_checker=False,
    custom_pipeline=None,
)
pipe.set_progress_bar_config(desc="Stream-DiffVSR")
pipe = pipe.to("cuda")

of_model = raft_large(weights=Raft_Large_Weights.DEFAULT).eval().to("cuda")
of_model.requires_grad_(False)

FFMPEG = imageio_ffmpeg.get_ffmpeg_exe()


# ------------------------------------------------------------------- helpers
def _round_down(value: int, multiple: int = MULTIPLE_OF) -> int:
    return max(multiple, (int(value) // multiple) * multiple)


def _read_frames(path: str, max_frames: int) -> Tuple[List[Image.Image], float]:
    try:
        fps = float(iio.immeta(path, plugin="pyav").get("fps") or 24.0)
    except Exception:
        fps = 24.0
    if not np.isfinite(fps) or fps <= 0:
        fps = 24.0

    frames: List[Image.Image] = []
    for frame in iio.imiter(path, plugin="pyav"):
        frames.append(Image.fromarray(np.asarray(frame)).convert("RGB"))
        if len(frames) >= max_frames:
            break
    if not frames:
        raise gr.Error("Could not read any frame from this video.")
    return frames, fps


def _lowres_size(w: int, h: int, max_input_size: int) -> Tuple[int, int]:
    """Downscale to `max_input_size` on the long side and to a 16:9 pixel budget.

    Bounding the pixel count (not just the long side) keeps the runtime — and hence
    the ZeroGPU reservation — predictable for any aspect ratio.
    """
    budget = max_input_size * max_input_size * 9.0 / 16.0
    scale = min(1.0, max_input_size / float(max(w, h)), (budget / float(w * h)) ** 0.5)
    return _round_down(w * scale), _round_down(h * scale)


def _prepare_lowres(frames: List[Image.Image], max_input_size: int) -> List[Image.Image]:
    w, h = _lowres_size(*frames[0].size, max_input_size=max_input_size)
    return [f if f.size == (w, h) else f.resize((w, h), Image.BICUBIC) for f in frames]


def _write_video(frames: List[Image.Image], fps: float, suffix: str) -> str:
    w, h = frames[0].size
    path = tempfile.NamedTemporaryFile(suffix=f"_{suffix}.mp4", delete=False).name
    writer = imageio_ffmpeg.write_frames(
        path,
        (w, h),
        fps=fps,
        codec="libx264",
        pix_fmt_in="rgb24",
        pix_fmt_out="yuv420p",
        macro_block_size=1,
        quality=None,
        output_params=["-crf", "17", "-preset", "veryfast", "-movflags", "+faststart"],
    )
    writer.send(None)
    for frame in frames:
        writer.send(np.asarray(frame, dtype=np.uint8).tobytes())
    writer.close()
    return path


def _estimate_duration(*args, **kwargs) -> int:
    """ZeroGPU duration estimator, fitted to runtime measured on this Space.

    Measured on ZeroGPU at 320x176 LR: 0.67 s/frame at 4 steps, 1.20 s/frame at 8 steps,
    i.e. ~0.15 s/frame fixed + ~0.125 s/frame/step, scaling with the LR pixel count.
    """
    values = list(args)

    def _pick(name, index, default):
        if name in kwargs and kwargs[name] is not None:
            return kwargs[name]
        if len(values) > index and values[index] is not None:
            return values[index]
        return default

    try:
        num_frames = int(_pick("num_frames", 1, 30))
        max_input_size = int(_pick("max_input_size", 2, 320))
        steps = int(_pick("num_inference_steps", 3, 4))
    except (TypeError, ValueError):
        num_frames, max_input_size, steps = 30, 320, 4

    # assume a 16:9 clip for the pixel-count scaling against the 312x176 reference
    # `_lowres_size` bounds the LR frame to a 16:9 pixel budget, so this is an upper bound
    px = (max_input_size * max_input_size * 9.0 / 16.0) / (312.0 * 176.0)
    seconds = 5.0 + num_frames * (0.15 + 0.125 * steps) * px * 1.25
    return int(min(115, max(20, round(seconds))))


# ----------------------------------------------------------------- inference
@spaces.GPU(duration=_estimate_duration)
def upscale_video(
    video: str,
    num_frames: int = 30,
    max_input_size: int = 320,
    num_inference_steps: int = 4,
    seed: int = 42,
    progress=gr.Progress(track_tqdm=True),
) -> Tuple[str, str, str]:
    """Upscale a low-resolution video 4x with Stream-DiffVSR.

    Args:
        video: path to the low-resolution input video.
        num_frames: how many frames from the start of the clip to process.
        max_input_size: the input is downscaled so its longest side is at most
            this many pixels before the 4x upscale.
        num_inference_steps: denoising steps per frame (the model is distilled for 4).
        seed: RNG seed for the initial latents.

    Returns:
        The Stream-DiffVSR 4x video, a bicubic 4x video for reference, and a status line.
    """
    if not video:
        raise gr.Error("Please upload a low-resolution video first.")

    num_frames = int(num_frames)
    num_inference_steps = int(num_inference_steps)
    max_input_size = int(max_input_size)
    seed = int(seed)

    frames, fps = _read_frames(video, num_frames)
    lowres = _prepare_lowres(frames, max_input_size)
    in_w, in_h = lowres[0].size
    out_w, out_h = in_w * SCALE, in_h * SCALE

    generator = torch.Generator(device="cuda").manual_seed(seed)

    start = time.perf_counter()
    with torch.inference_mode():
        result = pipe(
            "",
            lowres,
            num_inference_steps=num_inference_steps,
            guidance_scale=0,
            of_model=of_model,
            generator=generator,
        )
    elapsed = time.perf_counter() - start

    sr_frames = [f[0] if isinstance(f, (list, tuple)) else f for f in result.images]
    bicubic_frames = [f.resize((out_w, out_h), Image.BICUBIC) for f in lowres]

    sr_path = _write_video(sr_frames, fps, "streamdiffvsr")
    bicubic_path = _write_video(bicubic_frames, fps, "bicubic")

    status = (
        f"{len(sr_frames)} frames · {in_w}x{in_h} → {out_w}x{out_h} · "
        f"{num_inference_steps} steps · {elapsed:.1f}s "
        f"({elapsed / max(len(sr_frames), 1):.2f}s/frame)"
    )
    return sr_path, bicubic_path, status


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

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# Stream-DiffVSR · 4× Video Super-Resolution
Low-latency streamable video super-resolution via auto-regressive diffusion — a
causal, 4-step distilled denoiser with Auto-regressive Temporal Guidance (ARTG)
and a temporal-aware decoder, so every frame is restored from **past frames only**.

[Paper](https://huggingface.co/papers/2512.23709) ·
[Model](https://huggingface.co/Jamichsu/Stream-DiffVSR) ·
[Code](https://github.com/jamichss/Stream-DiffVSR) ·
[Project page](https://jamichss.github.io/stream-diffvsr-project-page/)
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                video_in = gr.Video(
                    label="Low-resolution video",
                    sources=["upload"],
                    height=300,
                )
                run_btn = gr.Button("Upscale 4×", variant="primary")
                status_out = gr.Markdown()
            with gr.Column(scale=2):
                with gr.Row():
                    sr_out = gr.Video(label="Stream-DiffVSR 4×", autoplay=True, loop=True)
                    bicubic_out = gr.Video(label="Bicubic 4× (reference)", autoplay=True, loop=True)

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                num_frames = gr.Slider(
                    label="Frames to process",
                    minimum=4,
                    maximum=40,
                    step=1,
                    value=30,
                )
                max_input_size = gr.Slider(
                    label="Max input size (long side, px)",
                    minimum=128,
                    maximum=384,
                    step=32,
                    value=320,
                    info="The clip is downscaled to this long side (and a matching 16:9 pixel budget) before the 4× upscale.",
                )
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Denoising steps",
                    minimum=1,
                    maximum=8,
                    step=1,
                    value=4,
                    info="The model is distilled for 4 steps.",
                )
                seed = gr.Number(label="Seed", value=42, precision=0)

        gr.Examples(
            examples=[
                ["examples/portrait_lr.mp4"],
                ["examples/city_night_lr.mp4"],
                ["examples/waterfall_lr.mp4"],
            ],
            inputs=[video_in],
            outputs=[sr_out, bicubic_out, status_out],
            fn=upscale_video,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples (low-resolution 320×176 clips)",
        )

    run_btn.click(
        fn=upscale_video,
        inputs=[video_in, num_frames, max_input_size, num_inference_steps, seed],
        outputs=[sr_out, bicubic_out, status_out],
        api_name="upscale_video",
    )

if __name__ == "__main__":
    # Gradio 6 moved `theme` / `css` from the Blocks constructor to launch()
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)