import spaces  # MUST come before torch
import os

import pathlib
import tempfile
import time

import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from torchvision.utils import save_image

from dit import MFDiT

MODEL_ID = "yscript/agent-genpath"
INPUT_SIZE = 256
# Backbone hyper-parameters exactly as in configs/sample.json from the paper repo.
BACKBONE = dict(
    input_size=256,
    patch_size=8,
    in_channels=3,
    dim=768,
    depth=12,
    num_heads=12,
    num_classes=None,
)

# --- Model load (module scope, eager .to("cuda"); ZeroGPU intercepts) -----------
model = MFDiT(**BACKBONE)
_ckpt = hf_hub_download(MODEL_ID, "dit.pth")
_state = torch.load(_ckpt, map_location="cpu", weights_only=True)
_state = {k.removeprefix("module."): v for k, v in _state.items()}
model.load_state_dict(_state, strict=True)
model.to("cuda").eval()
del _state


@torch.no_grad()
def euler_sample(batch, steps, seed):
    """Reference Euler sampler from generation/sample.py (Agent-GenPath repo)."""
    torch.manual_seed(seed)
    device = next(model.parameters()).device
    x = torch.randn(batch, 3, INPUT_SIZE, INPUT_SIZE, device=device)
    for i in range(steps):
        t = torch.full((batch,), 1 - i / steps, device=device)
        x = x - model(x, t, torch.zeros_like(t), y=None) / steps
    return (x.clamp(-1, 1) + 1) / 2


def _duration(seed, batch, steps, *args, **kwargs):
    # Positional args mirror generate()'s signature (Gradio also injects progress=,
    # which lands in *args). Measured on ZeroGPU: ~14 ms per image per Euler step
    # (showcase batch=8 x 50 steps = 5.5 s GPU time). ~2x headroom for dispatch
    # and first-call warmup; capped well below for the slider maximum.
    try:
        batch = max(1, min(int(batch), 8))
        steps = max(1, min(int(steps), 100))
    except (TypeError, ValueError):
        batch, steps = 8, 50
    return min(60, 4 + 0.016 * steps * batch)


@spaces.GPU(duration=_duration)
def generate(seed=6657, batch=8, steps=50, progress=gr.Progress(track_tqdm=True)):
    """Sample synthetic pathology image patches with the pretrained Agent-GenPath flow-matching DiT.

    Args:
        seed: RNG seed for the whole batch; same seed + settings reproduces the same patches.
        batch: number of 256x256 RGB patches to synthesize in one go (1-8).
        steps: number of Euler integration steps (paper showcase uses 50).
        progress: Gradio progress bar.

    Returns:
        List of file paths to the generated PNG patches.
    """
    seed = int(seed)
    batch = max(1, min(int(batch), 8))
    steps = max(1, min(int(steps), 100))
    t0 = time.perf_counter()
    imgs = euler_sample(batch, steps, seed)
    dt = time.perf_counter() - t0
    out = (
        pathlib.Path(tempfile.gettempdir())
        / f"agent-genpath-{seed}-{batch}-{steps}-{os.getpid()}"
    )
    out.mkdir(parents=True, exist_ok=True)
    paths = []
    for i, image in enumerate(imgs.cpu()):
        p = out / f"sample_{i + 1:05d}.png"
        save_image(image, p)
        paths.append(str(p))
    print(
        f"[agent-genpath] seed={seed} batch={batch} steps={steps} took {dt:.1f}s",
        flush=True,
    )
    return paths


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

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Agent-GenPath — Pathology Patch Synthesis
        **Agent-Driven Generative Priors for Pathology Image Synthesis and Segmentation** ([arXiv 2611.06701](https://arxiv.org/abs/2611.06701))

        Unconditional flow-matching Diffusion Transformer (DiT) trained on **TCGA-200K**, a
        217,030-patch pathology dataset curated by the paper's pathology agent system. The same
        backbone serves as a generative prior for downstream semantic segmentation. This demo
        samples 256×256 synthetic H&E patches from the released checkpoint
        ([`yscript/agent-genpath`](https://huggingface.co/yscript/agent-genpath)) with the
        authors' reference Euler sampler.

        *Code: [script-Yang/Agent-GenPath](https://github.com/script-Yang/Agent-GenPath) · Paper: [arXiv 2611.06701](https://arxiv.org/abs/2611.06701)*
        """
    )
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            seed = gr.Number(label="Seed", value=6657, precision=0)
            batch = gr.Slider(1, 8, value=8, step=1, label="Batch size")
            steps = gr.Slider(10, 100, value=50, step=5, label="Euler steps")
        run = gr.Button("Sample patches", variant="primary")
        gallery = gr.Gallery(
            label="Synthetic pathology patches",
            columns=4,
            height=520,
            format="png",
            type="filepath",
            buttons=["download", "fullscreen"],
        )
        gr.Examples(
            examples=[[6657, 8, 50]],
            inputs=[seed, batch, steps],
            outputs=gallery,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
            label="Paper showcase batch (seed 6657 · 8 patches · 50 Euler steps)",
        )
    run.click(generate, inputs=[seed, batch, steps], outputs=gallery, api_name="generate")

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