from __future__ import annotations

from pathlib import Path

import gradio as gr
import numpy as np
import torch
from model import PocketDenoiser
from PIL import Image
from safetensors.torch import load_file

STEPS = 50
ARTIFACT = Path(__file__).resolve().parent / "artifacts" / "pocket-diffusion"
MODEL = PocketDenoiser(diffusion_steps=STEPS)
MODEL.load_state_dict(load_file(ARTIFACT / "model.safetensors"))
MODEL.eval()


def generate_digit(label: int, seed: int, guidance: float) -> Image.Image:
    generator = torch.Generator().manual_seed(seed)
    betas = torch.linspace(1e-4, 0.025, STEPS)
    alphas = 1.0 - betas
    cumulative = torch.cumprod(alphas, dim=0)
    pixels = torch.randn(1, 64, generator=generator)
    labels = torch.tensor([label])
    null_labels = torch.tensor([10])
    with torch.no_grad():
        for step in reversed(range(STEPS)):
            timesteps = torch.tensor([step])
            conditional = MODEL(pixels, timesteps, labels)
            unconditional = MODEL(pixels, timesteps, null_labels)
            noise_prediction = unconditional + guidance * (conditional - unconditional)
            alpha = alphas[step]
            mean = (
                pixels - (1 - alpha) / torch.sqrt(1 - cumulative[step]) * noise_prediction
            ) / torch.sqrt(alpha)
            if step:
                pixels = mean + torch.sqrt(betas[step]) * torch.randn(
                    pixels.shape,
                    generator=generator,
                )
            else:
                pixels = mean
    image = torch.clamp((pixels[0] + 1) / 2, 0, 1).reshape(8, 8).numpy()
    array = np.clip(image * 255, 0, 255).astype(np.uint8)
    return Image.fromarray(array, mode="L").resize(
        (512, 512),
        Image.Resampling.NEAREST,
    )


with gr.Blocks(title="PocketDiffusion") as demo:
    gr.Markdown("# PocketDiffusion\nGenerate a digit through 50 reverse-denoising steps.")
    with gr.Row():
        label = gr.Slider(0, 9, value=8, step=1, label="Digit")
        seed = gr.Slider(0, 100_000, value=2032, step=1, label="Noise seed")
        guidance = gr.Slider(1.0, 4.0, value=3.0, step=0.1, label="Guidance")
    output = gr.Image(value=generate_digit(8, 2032, 3.0), label="Generated glyph")
    button = gr.Button("Denoise", variant="primary")
    button.click(generate_digit, inputs=[label, seed, guidance], outputs=output)


if __name__ == "__main__":
    demo.launch()