from __future__ import annotations

import json
from pathlib import Path

import gradio as gr
import numpy as np
import plotly.graph_objects as go
import torch
from model import ConditionalNeuralProcess
from safetensors.torch import load_file

ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "neural-process-pocket"
MODEL = ConditionalNeuralProcess()
MODEL.load_state_dict(load_file(ARTIFACT_DIR / "model.safetensors"))
MODEL.eval()
REPORT = json.loads((ARTIFACT_DIR / "evaluation.json").read_text(encoding="utf-8"))


@torch.inference_mode()
def infer_function(
    amplitude: float,
    phase: float,
    frequency: float,
    context_points: int,
    seed: int,
) -> tuple[go.Figure, dict]:
    rng = np.random.default_rng(int(seed))
    context_x = np.sort(rng.uniform(-5, 5, int(context_points))).astype(np.float32)
    context_y = amplitude * np.sin(frequency * context_x + phase)
    target_x = np.linspace(-5, 5, 300, dtype=np.float32)
    true_y = amplitude * np.sin(frequency * target_x + phase)
    mean, std = MODEL(
        torch.from_numpy(context_x)[None, :, None],
        torch.from_numpy(context_y.astype(np.float32))[None, :, None],
        torch.from_numpy(target_x)[None, :, None],
    )
    mean = mean[0, :, 0].numpy()
    std = std[0, :, 0].numpy()
    figure = go.Figure()
    figure.add_trace(
        go.Scatter(
            x=np.concatenate([target_x, target_x[::-1]]),
            y=np.concatenate([mean - 1.645 * std, (mean + 1.645 * std)[::-1]]),
            fill="toself",
            name="90% predictive interval",
            line={"color": "rgba(0,0,0,0)"},
        )
    )
    figure.add_trace(go.Scatter(x=target_x, y=true_y, name="True function"))
    figure.add_trace(go.Scatter(x=target_x, y=mean, name="CNP mean"))
    figure.add_trace(
        go.Scatter(x=context_x, y=context_y, mode="markers", name="Context")
    )
    figure.update_layout(template="plotly_dark", xaxis_title="x", yaxis_title="y")
    metrics = {
        "live_rmse": float(np.sqrt(np.mean((mean - true_y) ** 2))),
        "context_points": int(context_points),
        "verified_500_task_rmse": REPORT["benchmark"][
            "conditional_neural_process"
        ]["rmse"],
        "verified_90_percent_coverage": REPORT["benchmark"][
            "conditional_neural_process"
        ]["coverage_90"],
    }
    return figure, metrics


with gr.Blocks(title="Neural Process Pocket") as demo:
    gr.Markdown(
        "# Neural Process Pocket\n"
        "Give the model a few observations and watch it infer a complete function "
        "distribution with a learned uncertainty band."
    )
    with gr.Row():
        amplitude = gr.Slider(0.1, 5.0, value=2.5, step=0.1, label="Amplitude")
        phase = gr.Slider(0, 6.28, value=1.0, step=0.05, label="Phase")
        frequency = gr.Slider(0.8, 1.2, value=1.0, step=0.02, label="Frequency")
        context = gr.Slider(3, 20, value=5, step=1, label="Context points")
        seed = gr.Slider(0, 100_000, value=2141, step=1, label="Seed")
    initial = infer_function(2.5, 1.0, 1.0, 5, 2141)
    chart = gr.Plot(value=initial[0])
    metrics = gr.JSON(value=initial[1])
    button = gr.Button("Infer function distribution", variant="primary")
    button.click(
        infer_function,
        inputs=[amplitude, phase, frequency, context, seed],
        outputs=[chart, metrics],
    )


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