import torch
import numpy as np
import joblib
import json
import gradio as gr
from transformers import AutoTokenizer, AutoModel
import spaces

LABEL_ORDER = [
    "Desktop & Mobile & Web Development",
    "Cybersecurity",
    "AI / Machine Learning / Data Science",
    "Infrastructure (DevOps, Cloud, Databases, Networking)",
    "Clinical Diagnosis, Treatment & Surgery",
    "Medication & Pharmacology",
    "Mental Health (Clinical)",
    "Healthcare Organizations, System, Hospitals",
    "Nutrition",
    "Payments & Personal Budgeting",
    "Banking",
    "Investment, Markets & Cryptocurrency",
    "Corporate Accounting",
    "Physics & Mathematics",
    "Chemistry",
    "Biology",
    "Team Sports",
    "Individual Sports",
    "Fitness & Training",
    "Civil, Structural & Architecture",
    "Mechanical & Electrical Engineering",
    "Family & Relationships",
    "Personal Growth & Reflection",
    "Travel",
    "Marketing & Sales",
    "Entrepreneurship & Startups",
    "Management & Strategy & Human Resources",
    "Criminal & Civil Law",
    "Labor, Family & Contract Law",
    "Corporate, Regulatory & International Law",
    "General Law (misc.)",
    "Game",
    "Film",
    "Music",
    "Literature",
    "Painting",
]

HIERARCHY = {
    "Technology & Programming": [
        "Desktop & Mobile & Web Development",
        "Cybersecurity",
        "AI / Machine Learning / Data Science",
        "Infrastructure (DevOps, Cloud, Databases, Networking)",
    ],
    "Medical": [
        "Clinical Diagnosis, Treatment & Surgery",
        "Medication & Pharmacology",
        "Mental Health (Clinical)",
        "Healthcare Organizations, System, Hospitals",
        "Nutrition",
    ],
    "Finance": [
        "Payments & Personal Budgeting",
        "Banking",
        "Investment, Markets & Cryptocurrency",
        "Corporate Accounting",
    ],
    "Science": [
        "Physics & Mathematics",
        "Chemistry",
        "Biology",
    ],
    "Sports": [
        "Team Sports",
        "Individual Sports",
        "Fitness & Training",
    ],
    "Engineering": [
        "Civil, Structural & Architecture",
        "Mechanical & Electrical Engineering",
    ],
    "Personal": [
        "Family & Relationships",
        "Personal Growth & Reflection",
        "Travel",
    ],
    "Business": [
        "Marketing & Sales",
        "Entrepreneurship & Startups",
        "Management & Strategy & Human Resources",
    ],
    "Law": [
        "Criminal & Civil Law",
        "Labor, Family & Contract Law",
        "Corporate, Regulatory & International Law",
        "General Law (misc.)",
    ],
    "Art": [
        "Game",
        "Film",
        "Music",
        "Literature",
        "Painting",
    ],
}

CHILD_TO_PARENT = {}
for parent, children in HIERARCHY.items():
    for child in children:
        CHILD_TO_PARENT[child] = parent

id_to_label = {i: label for i, label in enumerate(LABEL_ORDER)}

MODEL_NAME = "BAAI/bge-small-en-v1.5"
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading embedding model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME).to(device)
model.eval()

print("Loading trained classifier...")
classifier = joblib.load("classifier.joblib")


def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
        input_mask_expanded.sum(1), min=1e-9
    )


def get_embedding(texts):
    encoded_input = tokenizer(
        texts, padding=True, truncation=True, max_length=512, return_tensors="pt"
    ).to(device)
    with torch.no_grad():
        model_output = model(**encoded_input)
    embeddings = mean_pooling(model_output, encoded_input["attention_mask"])
    embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
    return embeddings.cpu().numpy()


print("Model ready!")


CSS = """
<style>
.prob-table { width: 100%; border-collapse: collapse; font-size: 13px; }
.prob-table td { padding: 4px 0; vertical-align: middle; white-space: nowrap; }
.prob-table .rank { color: #888; width: 30px; text-align: right; padding-right: 8px; }
.prob-table .name { font-size: 12px; color: #ccc; white-space: nowrap; }
.prob-table .pct { width: 60px; text-align: right; padding-right: 10px; font-weight: 600; font-variant-numeric: tabular-nums; }
.prob-table .bar-cell { width: 100%; padding-left: 4px; }
.prob-bar-bg { background: #1e1e2e; border-radius: 4px; height: 20px; width: 100%; overflow: hidden; }
.prob-bar-fill { height: 100%; border-radius: 4px; transition: width 0.3s; }
.top-prediction { margin-bottom: 18px; padding: 14px 18px; background: #1a1a2e; border-radius: 8px; border-left: 4px solid #7c3aed; }
.top-prediction .label { font-size: 13px; color: #aaa; margin-bottom: 2px; }
.top-prediction .value { font-size: 20px; font-weight: 700; color: #fff; }
.top-prediction .sub { font-size: 13px; color: #9ca3af; margin-top: 4px; }
.section-title { font-size: 14px; font-weight: 600; color: #9ca3af; margin: 16px 0 8px 0; text-transform: uppercase; letter-spacing: 0.5px; }
</style>
"""


def _bar_color(pct):
    if pct >= 50:
        return "#7c3aed"
    if pct >= 20:
        return "#6366f1"
    if pct >= 5:
        return "#818cf8"
    if pct >= 1:
        return "#a5b4fc"
    return "#c7d2fe"

@spaces.GPU
def predict(text):
    if not text or not text.strip():
        return {}

    embedding = get_embedding([text])
    probs = classifier.predict_proba(embedding)[0]
    confidences = {id_to_label[i]: float(probs[i]) for i in range(len(probs))}
    return confidences



EXAMPLES = [
    ["Can you help me build a responsive navbar with HTML, CSS, and JavaScript that works on mobile devices?"],
    ["What are the best practices for protecting a web application against SQL injection attacks?"],
    ["Explain the difference between supervised and unsupervised learning in machine learning."],
    ["How do I set up a Kubernetes cluster on AWS EKS for deploying microservices?"],
    ["What are the common symptoms of type 2 diabetes and how is it diagnosed?"],
    ["Can you explain how SSRIs work for treating depression and what side effects to expect?"],
    ["What's the difference between a traditional IRA and a Roth IRA for retirement savings?"],
    ["How does compound interest work and what's the best way to start investing with $500?"],
    ["Explain the Schrödinger equation and how it applies to quantum mechanics."],
    ["What are the rules of soccer and how does the offside rule work?"],
    ["How do I create a workout plan for building muscle as a beginner?"],
    ["What are the key principles of structural engineering when designing a multi-story building?"],
    ["How do I handle conflicts with my partner in a healthy and constructive way?"],
    ["What are some strategies for improving my public speaking skills?"],
    ["What should I pack and plan for a two-week trip to Southeast Asia?"],
    ["How do I create an effective social media marketing campaign for my small business?"],
    ["What are the essential steps to validate a startup idea before building a product?"],
    ["How should I structure my sales team for a B2B SaaS company?"],
    ["What are my rights if my employer fires me without cause in California?"],
    ["How does copyright law protect original creative works?"],
    ["Can you recommend some indie games with great storytelling and emotional depth?"],
    ["What cinematography techniques make a thriller movie more suspenseful?"],
    ["How do I write a catchy chorus for a pop song?"],
    ["What are some must-read classic novels from the 20th century?"],
    ["What techniques should I use to paint realistic portraits with oil paints?"],
]

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        label="Input Text",
        placeholder="Enter text to classify...",
        lines=5,
    ),
    outputs=gr.Label(num_top_classes=len(LABEL_ORDER)),
    title="LLM-Prompts Topic Classifier",
    description="Classify text into one of 36 topics across 10 categories using BGE embeddings and logistic regression.",
    examples=EXAMPLES,
)

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