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 MeshGraphGAT
from safetensors.torch import load_file

PROJECT_DIR = Path(__file__).resolve().parent
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "meshgraph-gat"
GRAPH = np.load(PROJECT_DIR / "data" / "meshgraph.npz")
PREPROCESSING = np.load(ARTIFACT_DIR / "preprocessing.npz")
FEATURES = torch.from_numpy(
    (
        GRAPH["features"].astype(np.float32) - PREPROCESSING["mean"]
    )
    / PREPROCESSING["scale"]
)
ADJACENCY = torch.from_numpy(GRAPH["adjacency"].astype(np.float32))
MODEL = MeshGraphGAT(FEATURES.shape[1])
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 inspect_node(node: int) -> tuple[go.Figure, dict]:
    node = int(node)
    logits, attention = MODEL(FEATURES, ADJACENCY, return_attention=True)
    risk = torch.softmax(logits, dim=1)[node, 1]
    neighbors = torch.nonzero(ADJACENCY[node] > 0).flatten()
    weights = attention[:, node, neighbors].mean(0)
    order = torch.argsort(weights, descending=True)
    neighbor_ids = neighbors[order].tolist()
    sorted_weights = weights[order].tolist()
    figure = go.Figure(go.Bar(x=[str(value) for value in neighbor_ids], y=sorted_weights))
    figure.update_layout(
        template="plotly_dark",
        title=f"Mean four-head attention from node {node}",
        xaxis_title="Neighbor node",
        yaxis_title="Attention",
    )
    metrics = {
        "node": node,
        "compromise_probability": float(risk),
        "predicted_compromised": bool(
            float(risk) >= float(PREPROCESSING["threshold"])
        ),
        "true_label": int(GRAPH["labels"][node]),
        "degree": len(neighbor_ids),
        "verified_test_roc_auc": REPORT["gat_test"]["roc_auc"],
    }
    return figure, metrics


with gr.Blocks(title="MeshGraph GAT") as demo:
    gr.Markdown(
        "# MeshGraph GAT\n"
        "Inspect which enterprise communication edges a four-head graph-attention "
        "network uses when estimating compromise risk."
    )
    node = gr.Slider(0, len(FEATURES) - 1, value=25, step=1, label="Node")
    initial = inspect_node(25)
    chart = gr.Plot(value=initial[0])
    metrics = gr.JSON(value=initial[1])
    button = gr.Button("Inspect attention", variant="primary")
    button.click(inspect_node, inputs=node, outputs=[chart, metrics])


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