import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from concurrent.futures import ThreadPoolExecutor

# ------------------------------
# 1) Load Models
# ------------------------------

# Mental Model
mental_model_name = "ShubhamCoder01/stress-detection-model"
mental_tokenizer = AutoTokenizer.from_pretrained(mental_model_name)
mental_model = AutoModelForSequenceClassification.from_pretrained(mental_model_name)

# Physical / Emotion Model
physical_model_name = "ShubhamCoder01/General-Emotion-detect"
physical_tokenizer = AutoTokenizer.from_pretrained(physical_model_name)
physical_model = AutoModelForSequenceClassification.from_pretrained(physical_model_name)
physical_id2label = {0: 'Joy', 1: 'Sadness', 2: 'Anger', 3: 'Love', 4: 'Surprise', 5: 'Fear'}

# Polarity Model
polarity_model_name = "ShubhamCoder01/college-context-polarity-detect"
polarity_tokenizer = AutoTokenizer.from_pretrained(polarity_model_name)
polarity_model = AutoModelForSequenceClassification.from_pretrained(polarity_model_name)
polarity_id2label = {0: 'NEGATIVE', 1: 'POSITIVE', 2:"NORMAL"}

# ------------------------------
# 2) Fuzzy Stress Score Maps
# ------------------------------
mental_map = {"Anxiety": 0.85, "Bipolar": 0.80, "Depression": 0.92,
              "Normal": 0.0, "Personality disorder": 0.78,
              "Stress": 0.95, "Suicidal": 1.0}

polarity_map = {"NEGATIVE": 0.75, "POSITIVE": 0.0, "NORMAL": 0.30}

emotion_map = {"Joy": 0.0, "Sadness": 0.88, "Anger": 0.70,
               "Love": 0.10, "Surprise": 0.25, "Fear": 0.95}

# Conditional probabilities (Mental -> Polarity / Emotion)
cond_p = {
    "Anxiety": {"NEGATIVE": 0.8, "POSITIVE": 0.1, "NORMAL": 0.5},
    "Bipolar": {"NEGATIVE": 0.7, "POSITIVE": 0.2, "NORMAL": 0.5},
    "Depression": {"NEGATIVE": 0.9, "POSITIVE": 0.05, "NORMAL": 0.4},
    "Normal": {"NEGATIVE": 0.2, "POSITIVE": 0.6, "NORMAL": 0.8},
    "Personality disorder": {"NEGATIVE": 0.6, "POSITIVE": 0.2, "NORMAL": 0.5},
    "Stress": {"NEGATIVE": 0.85, "POSITIVE": 0.05, "NORMAL": 0.3},
    "Suicidal": {"NEGATIVE": 1.0, "POSITIVE": 0.0, "NORMAL": 0.2}
}

cond_e = {
    "Anxiety": {"Joy":0.1,"Sadness":0.85,"Anger":0.7,"Love":0.1,"Surprise":0.2,"Fear":0.9},
    "Bipolar": {"Joy":0.2,"Sadness":0.75,"Anger":0.6,"Love":0.2,"Surprise":0.3,"Fear":0.8},
    "Depression": {"Joy":0.0,"Sadness":0.95,"Anger":0.7,"Love":0.05,"Surprise":0.1,"Fear":0.95},
    "Normal": {"Joy":0.8,"Sadness":0.1,"Anger":0.1,"Love":0.7,"Surprise":0.3,"Fear":0.2},
    "Personality disorder": {"Joy":0.2,"Sadness":0.8,"Anger":0.6,"Love":0.2,"Surprise":0.3,"Fear":0.7},
    "Stress": {"Joy":0.0,"Sadness":0.88,"Anger":0.7,"Love":0.1,"Surprise":0.25,"Fear":0.95},
    "Suicidal": {"Joy":0.0,"Sadness":1.0,"Anger":0.8,"Love":0.0,"Surprise":0.1,"Fear":1.0}
}

# Model weights
W_M, W_P, W_E = 0.75, 0.15, 0.10

# ------------------------------
# 3) Prediction Functions
# ------------------------------
def predict_mental(text):
    inputs = mental_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = mental_model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    idx = torch.argmax(probs).item()
    return mental_model.config.id2label[idx], probs[0][idx].item()

def predict_physical(text):
    inputs = physical_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = physical_model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    idx = torch.argmax(probs).item()
    return physical_id2label[idx], probs[0][idx].item()

def predict_polarity(text):
    inputs = polarity_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = polarity_model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    idx = torch.argmax(probs).item()
    return polarity_id2label[idx], probs[0][idx].item()

# ------------------------------
# 4) Fuzzy Stress Score Calculation
# ------------------------------
def calculate_stress_score(mental_label, mental_conf,
                           polarity_label, polarity_conf,
                           emotion_label, emotion_conf):
    adj_M = mental_conf * mental_map[mental_label]
    adj_P = polarity_conf * cond_p[mental_label][polarity_label] * polarity_map[polarity_label]
    adj_E = emotion_conf * cond_e[mental_label][emotion_label] * emotion_map[emotion_label]
    stress_score = W_M * adj_M + W_P * adj_P + W_E * adj_E
    return stress_score

# ------------------------------
# 5) Combined Prediction
# ------------------------------
def predict_all(text):
    # Parallel model inference
    with ThreadPoolExecutor() as executor:
        future_mental = executor.submit(predict_mental, text)
        future_physical = executor.submit(predict_physical, text)
        future_polarity = executor.submit(predict_polarity, text)
        mental_label, mental_conf = future_mental.result()
        physical_label, physical_conf = future_physical.result()
        polarity_label, polarity_conf = future_polarity.result()

    stress_score = calculate_stress_score(
        mental_label, mental_conf,
        polarity_label, polarity_conf,
        physical_label, physical_conf
    )

    # Low confidence flag
    low_confidence = mental_conf < 0.25 or physical_conf < 0.5 or polarity_conf < 0.5

    return {
        "Mental": mental_label,
        "Mental_conf": round(mental_conf,4),
        "Physical": physical_label,
        "Physical_conf": round(physical_conf,4),
        "Polarity": polarity_label,
        "Polarity_conf": round(polarity_conf,4),
        "Stress_Score": round(stress_score,4),
        "Low_Confidence_Flag": low_confidence
    }

# ------------------------------
# 6) Gradio Interface
# ------------------------------
iface = gr.Interface(
    fn=predict_all,
    inputs=gr.Textbox(lines=3, placeholder="Enter a sentence..."),
    outputs="json",
    title="Student Stress Detection (Fuzzy Probabilistic)",
    description="Predict Mental, Emotion, Polarity and Fuzzy Stress Score using parallel model inference."
)

iface.launch(share=True)