import torch
import joblib
import gradio as gr
import matplotlib.pyplot as plt
from sklearn.multioutput import MultiOutputClassifier

# Allow loading MultiOutputClassifier
torch.serialization.add_safe_globals([MultiOutputClassifier])
                                    
# Load the Model and Dependencies
tfidf_vectorizer = joblib.load("tfidf_vectorizer.pkl")
category_encoder = joblib.load("category_encoder.pkl")
team_encoder = joblib.load("team_encoder.pkl")

# Load the model with full object support
multi_label_classifier = torch.load("multi_label_classifier.pth", weights_only=False)
# multi_label_classifier.eval()

# Dummy function to get keywords
def get_top_keywords_per_category(category, n=5):
    keywords_dict = {
        "UX Issue": ["mobile", "responsive", "alignment", "css", "layout"],
        "Backend Failure": ["API", "server", "timeout", "database", "error"],
        "UI Bug": ["button", "color", "CSS", "overlap", "scroll"],
    }
    return keywords_dict.get(category, ["No keywords found"])

# Function to predict and visualize results
def predict_with_visuals(phrase):
    text_features = tfidf_vectorizer.transform([phrase])
    predicted_labels = multi_label_classifier.predict(text_features)

    predicted_category = category_encoder.inverse_transform([predicted_labels[0][0]])[0]
    predicted_team = team_encoder.inverse_transform([predicted_labels[0][1]])[0]
    team_email = f"support@{predicted_team.replace(' ', '').lower()}.com"
    keywords = get_top_keywords_per_category(predicted_category, n=5)

    # Visualization: Pie chart for category distribution
    categories = ["UX Issue", "Backend Failure", "UI Bug"]
    category_counts = [1 if cat == predicted_category else 0 for cat in categories]

    fig, ax = plt.subplots()
    ax.pie(category_counts, labels=categories, autopct="%1.1f%%", colors=["#ff9999","#66b3ff","#99ff99"])
    ax.set_title("Predicted Category Distribution")

    # Styled output
    result = f"""
    <div style='font-size: 18px; font-family: Arial;'>
    <strong>📌 Predicted Category:</strong> <span style='color:blue;'>{predicted_category}</span><br>
    <strong>👨‍💻 Assigned Team:</strong> <span style='color:green;'>{predicted_team}</span><br>
    <strong>📧 Team Email:</strong> <span style='color:red;'>{team_email}</span><br>
    <strong>🔑 Top Keywords:</strong> <span style='color:purple;'>{', '.join(keywords)}</span>
    </div>
    """

    return result, fig

# Gradio Interface with Better UI
interface = gr.Interface(
    fn=predict_with_visuals,
    inputs=gr.Textbox(lines=2, placeholder="Enter defect description..."),
    outputs=["html", "plot"],
    title="🔎 AI Defect Ticket Classifier",
    description="Enter a defect description to predict its **Category, Assigned Team, and relevant Keywords**. Get a **visual breakdown** of the classification!",
    theme="compact"
)

# Launch App
interface.launch()