"""Minimal Gradio frontend for raw burst reconstruction."""

from pathlib import Path
import shutil
import tempfile

import spaces
import gradio as gr
import numpy as np


@spaces.GPU(duration=120)
def reconstruct(files, scale):
    # Import CUDA modules only after a GPU has been allocated.
    from handheld_super_resolution.config import Config
    from handheld_super_resolution.monte_carlo import generate_noise_lut
    from handheld_super_resolution.noise_lut import save_noise_lut
    from handheld_super_resolution.utils_dng import read_dng_noise_profile
    from handheld_super_resolution.super_resolution import process

    with tempfile.TemporaryDirectory(prefix="handheld-") as directory:
        burst = Path(directory)
        for index, source in enumerate(sorted(files, key=lambda p: Path(p).name)):
            shutil.copyfile(source, burst / f"{index:04d}.dng")

        alpha, beta = read_dng_noise_profile(burst / "0000.dng")
        lut = generate_noise_lut(alpha, beta, show_progress=False)
        lut_path = burst / "noise.npz"
        save_noise_lut(lut_path, lut, trials=10_000_000, seed=0)
        config = Config(scale=float(scale))
        config.noise_model.lut_path = lut_path
        config.robustness.save_mask = False
        result, _ = process(burst, config)
        return np.round(np.clip(np.nan_to_num(result), 0, 1) * 255).astype(np.uint8)


def run(files, scale):
    if not files or len(files) < 2:
        raise gr.Error("Upload at least two DNG images from the same burst.")
    if any(Path(path).suffix.lower() != ".dng" for path in files):
        raise gr.Error("This pipeline accepts raw DNG images.")
    try:
        return reconstruct(files, scale)
    except Exception as error:
        raise gr.Error(f"Reconstruction failed: {error}") from error


with gr.Blocks(delete_cache=(3600, 3600)) as demo:
    gr.Markdown("# Handheld Super-Resolution\nUpload raw DNG images from one burst. "
                "The first image alphabetically is the reference. "
                "Noise calibration is automatic; the first run may take longer.")
    images = gr.File(label="DNG burst", file_count="multiple", file_types=[".dng"], type="filepath")
    scale = gr.Radio([1, 2], value=2, label="Output scale")
    button = gr.Button("Reconstruct", variant="primary")
    result = gr.Image(label="Result", format="png")
    button.click(run, inputs=[images, scale], outputs=result, concurrency_limit=1)


if __name__ == "__main__":
    demo.queue(max_size=8).launch()