from __future__ import annotations

from pathlib import Path

import gradio as gr
import numpy as np
import plotly.graph_objects as go
import torch
from model import JEncoder
from plotly.subplots import make_subplots
from safetensors.torch import load_file

ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "pocket-jepa"
MODEL = JEncoder()
MODEL.load_state_dict(load_file(ARTIFACT_DIR / "model.safetensors"))
MODEL.eval()
ATLAS = np.load(ARTIFACT_DIR / "j_space.npz")
EMBEDDINGS = ATLAS["embeddings"]
IMAGES = ATLAS["images"]
LABELS = ATLAS["labels"]


def explore(index: int, mask_size: int) -> tuple[go.Figure, dict]:
    index = int(index) % len(IMAGES)
    mask_size = int(mask_size)
    image = IMAGES[index].copy()
    masked = image.copy()
    start = (8 - mask_size) // 2
    masked[start : start + mask_size, start : start + mask_size] = 0
    with torch.inference_mode():
        embedding = MODEL(
            torch.from_numpy(masked[None].astype(np.float32))
        ).numpy()[0]
    embedding /= max(np.linalg.norm(embedding), 1e-8)
    distances = 1 - EMBEDDINGS @ embedding
    nearest = np.argsort(distances)[:5]
    panels = [image, masked, *IMAGES[nearest[:3]]]
    titles = [
        f"Original: {LABELS[index]}",
        "Masked context",
        *[f"J-neighbor: {LABELS[item]}" for item in nearest[:3]],
    ]
    figure = make_subplots(rows=1, cols=5, subplot_titles=titles)
    for panel_index, (panel, title) in enumerate(zip(panels, titles, strict=True)):
        figure.add_trace(
            go.Heatmap(
                z=np.flipud(panel),
                colorscale="Viridis",
                showscale=False,
                name=title,
            ),
            row=1,
            col=panel_index + 1,
        )
    figure.update_layout(
        title="Masked prediction in learned J-space",
        template="plotly_dark",
        height=310,
        margin=dict(t=70, b=20),
    )
    return figure, {
        "query_label": int(LABELS[index]),
        "nearest_labels": LABELS[nearest].astype(int).tolist(),
        "cosine_distances": np.round(distances[nearest], 4).tolist(),
    }


with gr.Blocks(title="Pocket JEPA") as demo:
    gr.Markdown(
        "# Pocket JEPA\n"
        "Hide the center of a digit and inspect the nearest complete images in "
        "the learned joint-embedding prediction space."
    )
    with gr.Row():
        index = gr.Slider(0, len(IMAGES) - 1, 0, step=1, label="Image")
        mask = gr.Slider(2, 5, 3, step=1, label="Center mask")
    run = gr.Button("Explore J-space", variant="primary")
    gallery = gr.Plot()
    metrics = gr.JSON()
    run.click(explore, [index, mask], [gallery, metrics])
    demo.load(explore, [index, mask], [gallery, metrics])


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