import gradio as gr
import torch
import torch.nn as nn
import json
from transformers import AutoTokenizer, AutoConfig, AutoModel

# ----------------------------
# Configuration / Meta
# ----------------------------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_LENGTH = 256
WEIGHTS_PATH = "best_model.pth"  # your model file
TOKENIZER_PATH = "tokenizer"
META_PATH = "meta.json"

# Load meta.json
with open(META_PATH, "r") as f:
    meta = json.load(f)
TARGET_COLS = meta.get("target_cols", ["anger","fear","joy","sadness","surprise"])
MAX_LENGTH = meta.get("max_length", MAX_LENGTH)
BACKBONE = meta.get("backbone", "microsoft/deberta-v3-base")

# ----------------------------
# Load tokenizer
# ----------------------------
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)

# ----------------------------
# Define model class
# ----------------------------
class EmotionClassifier(nn.Module):
    def __init__(self, model_name, num_labels=5, dropout_rate=0.3):
        super().__init__()
        self.config = AutoConfig.from_pretrained(model_name)
        self.transformer = AutoModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(dropout_rate)
        self.classifier = nn.Linear(self.config.hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
        cls = outputs.last_hidden_state[:, 0, :]
        cls = self.dropout(cls)
        logits = self.classifier(cls)
        return logits

# ----------------------------
# Load model
# ----------------------------
model = EmotionClassifier(BACKBONE, num_labels=len(TARGET_COLS))
state = torch.load(WEIGHTS_PATH, map_location=DEVICE)
# Strip 'module.' prefix if present
new_state = {k.replace("module.", ""): v for k,v in state.items()}
model.load_state_dict(new_state, strict=False)
model.to(DEVICE)
model.eval()

# ----------------------------
# Prediction function
# ----------------------------
def predict_text(text, threshold=0.5):
    inputs = tokenizer(
        text,
        truncation=True,
        padding=True,
        max_length=MAX_LENGTH,
        return_tensors="pt"
    )
    inputs = {k:v.to(DEVICE) for k,v in inputs.items()}
    
    with torch.no_grad():
        logits = model(**inputs)
    
    probs = torch.sigmoid(logits).cpu().numpy()[0]
    
    # Percentages
    percentages = {label: round(float(p)*100,2) for label,p in zip(TARGET_COLS, probs)}
    
    # Binary predictions
    binary = {label: int(p>threshold) for label,p in zip(TARGET_COLS, probs)}
    
    # Sort probabilities descending for display
    percentages_sorted = dict(sorted(percentages.items(), key=lambda x: x[1], reverse=True))
    
    return percentages_sorted, binary

# ----------------------------
# Gradio Interface
# ----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Multi-label Emotion Classifier")
    txt = gr.Textbox(lines=4, placeholder="Enter text here...", label="Input Text")
    thresh = gr.Slider(minimum=0.1, maximum=0.9, value=0.5, step=0.01, label="Threshold")
    btn = gr.Button("Predict")
    probs_out = gr.Label(label="Emotion Probabilities (%)")
    binary_out = gr.JSON(label="Binary Predictions (0/1)")
    
    btn.click(predict_text, inputs=[txt, thresh], outputs=[probs_out, binary_out])

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