import gradio as gr
from transformers import AutoTokenizer
import torch
import matplotlib.pyplot as plt
from model import EnergySmellsDetector
from config import SMELLS, BEST_THRESHOLD

TOKENIZER = "microsoft/graphcodebert-base"
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
model = EnergySmellsDetector.load_model_from_hf()

def get_predictions(code_snippet):
    inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True)
    with torch.no_grad():
        probs = model(**inputs)[0].cpu().numpy().flatten()  # Model output is already sigmoid applied
        rounded_logits = (probs > BEST_THRESHOLD).astype(int)

    # Prepare results in a dictionary
    results = {label: {"Detected": bool(pred), "Confidence": round(prob * 100, 2)}
               for label, pred, prob in zip(SMELLS, rounded_logits, probs)}

    return results, plot_bar_chart(results)


def plot_bar_chart(results):
    labels = list(results.keys())
    confidences = [results[label]["Confidence"] for label in labels]

    plt.figure(figsize=(8, 4))
    plt.barh(labels, confidences, color=['green' if results[label]["Detected"] else 'red' for label in labels])
    plt.xlabel("Confidence (%)")
    plt.xlim(0, 100)
    plt.title(f"Energy Smells Probabilities - Threshold: {BEST_THRESHOLD}")
    plt.gca().invert_yaxis()  # Invert y-axis for better readability
    plt.tight_layout()
    img_path = "confidence_chart.png"
    plt.savefig(img_path)
    plt.close()
    
    return img_path


textbox = gr.Textbox(label="Enter your code snippet", placeholder="Paste your code here...")
title = "Energy Smells Detector"
description = "Analyze your code for potential energy smells. The model detects 9 different energy inefficiencies in your code."

gr.Interface(
    title=title,
    description=description,
    inputs=textbox,
    fn=get_predictions,
    outputs=[
        gr.Json(label="Detection Results"),
        gr.Image(label="Confidence Bar Chart")
    ],
).launch()