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 BitMLP
from packing import load_binary_model
from safetensors.torch import load_file
from sklearn.datasets import load_digits

ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "bitforge-1bit"
FP32 = BitMLP("fp32")
FP32.load_state_dict(load_file(ARTIFACT_DIR / "fp32.safetensors"))
BINARY = load_binary_model(ARTIFACT_DIR / "binary_weights.npz")
TERNARY = BitMLP("ternary")
TERNARY.load_state_dict(load_file(ARTIFACT_DIR / "ternary_qat.safetensors"))
for model in [FP32, BINARY, TERNARY]:
    model.eval()
DIGITS = load_digits()
IMAGES = (DIGITS.images / 16.0).astype(np.float32)
LABELS = DIGITS.target


def compare(index: int) -> tuple[go.Figure, dict]:
    index = int(index) % len(IMAGES)
    image = torch.from_numpy(IMAGES[index][None])
    predictions = {}
    with torch.inference_mode():
        for name, model in [
            ("FP32", FP32),
            ("Binary packed", BINARY),
            ("Ternary", TERNARY),
        ]:
            predictions[name] = torch.softmax(model(image), dim=1).numpy()[0]
    figure = go.Figure()
    for name, probabilities in predictions.items():
        figure.add_trace(
            go.Bar(
                name=name,
                x=list(range(10)),
                y=probabilities,
            )
        )
    figure.update_layout(
        title=f"Digit {int(LABELS[index])}: precision variants",
        xaxis_title="Predicted class",
        yaxis_title="Probability",
        barmode="group",
        template="plotly_dark",
    )
    return figure, {
        name: {
            "prediction": int(values.argmax()),
            "confidence": round(float(values.max()), 4),
        }
        for name, values in predictions.items()
    }


with gr.Blocks(title="BitForge 1-bit") as demo:
    gr.Markdown(
        "# BitForge 1-bit\n"
        "Compare full-precision, packed one-bit matrix weights, and ternary weights "
        "on the same digit."
    )
    index = gr.Slider(0, len(IMAGES) - 1, 0, step=1, label="Digit sample")
    run = gr.Button("Compare precision", variant="primary")
    chart = gr.Plot()
    results = gr.JSON()
    run.click(compare, index, [chart, results])
    demo.load(compare, index, [chart, results])


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