import json
from pathlib import Path

import spaces
import gradio as gr
import pandas as pd
import torch
from huggingface_hub import snapshot_download

from model import classify, load_classifier


ROOT = Path(__file__).resolve().parent
deployment = json.loads((ROOT / "deployment.json").read_text())
model_directory = snapshot_download(
    repo_id=deployment["model_id"],
    revision=deployment["model_revision"],
    allow_patterns=["config.json", "model.safetensors"],
)
torch.set_num_threads(2)
model, config = load_classifier(model_directory)
model.to("cuda" if torch.cuda.is_available() else "cpu")


@spaces.GPU(duration=20)
def predict(image, top_k):
    return {
        "model_id": deployment["model_id"],
        "model_revision": deployment["model_revision"],
        "predictions": classify(model, config, image, int(top_k)),
    }


with gr.Blocks(title="Lab 1 — CIFAR-10") as demo:
    gr.Markdown("# CIFAR-10 image classifier\nResNet-18 fine-tuned in Google Colab.")
    gr.Markdown("Classes: " + ", ".join(config["labels"]))
    with gr.Tab("Predict"):
        image = gr.Image(type="pil", label="Image", image_mode="RGB")
        top_k = gr.Slider(1, len(config["labels"]), value=3, step=1, label="Top K")
        button = gr.Button("Predict", variant="primary")
        result = gr.JSON(label="Predictions")
        button.click(predict, [image, top_k], result, api_name="predict")
    with gr.Tab("Metrics"):
        summary = json.loads((ROOT / "assets" / "summary.json").read_text())
        gr.Markdown(f"Test accuracy: **{summary['test_accuracy']:.2%}**")
        gr.Dataframe(pd.read_csv(ROOT / "assets" / "per_class_metrics.csv"), interactive=False)
        gr.Image(str(ROOT / "assets" / "confusion_matrix.png"), label="Confusion matrix", interactive=False)
        gr.JSON(summary, label="Training and evaluation details")


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()