# src/core_model/predict.py
import os
import torch
import torch.nn.functional as F
from transformers import XLMRobertaTokenizer, XLMRobertaForSequenceClassification

class MindGuardPredictor:
    def __init__(self):
        # 1. Identify the Model Hub ID
        self.model_id = "MohitRajput45/mindguard-xlmr"
        
        # 2. Identify the EXACT path inside that Repo
        # Based on your HF screenshot, this is the full path to the files
        hf_subfolder = "artifacts/xlmr_weights/final_mindguard_model"
        
        # 3. Load the components from the Hub
        try:
            print(f"DEBUG: Attempting to load model from {self.model_id} in {hf_subfolder}...")
            self.tokenizer = XLMRobertaTokenizer.from_pretrained(
                self.model_id, 
                subfolder=hf_subfolder
            )
            self.model = XLMRobertaForSequenceClassification.from_pretrained(
                self.model_id, 
                subfolder=hf_subfolder
            )
            print("✅ Model and Tokenizer loaded successfully from Hub!")
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            raise e # Stop the app if the model fails to load
        
        # 3. Emotion Mapping (Mathematical ID to English Word)
        self.emotion_map = {
            0: 'Anxiety', 1: 'Bipolar', 2: 'Depression', 3: 'Normal', 
            4: 'Personality disorder', 5: 'Stress', 6: 'Suicidal', 7: 'admiration', 
            8: 'amusement', 9: 'anger', 10: 'annoyance', 11: 'approval', 
            12: 'caring', 13: 'confusion', 14: 'curiosity', 15: 'desire', 
            16: 'disappointment', 17: 'disapproval', 18: 'disgust', 19: 'embarrassment', 
            20: 'excitement', 21: 'fear', 22: 'gratitude', 23: 'grief', 
            24: 'joy', 25: 'love', 26: 'nervousness', 27: 'neutral', 
            28: 'optimism', 29: 'pride', 30: 'realization', 31: 'relief', 
            32: 'remorse', 33: 'sadness', 34: 'surprise'
        }
        
        # 4. Device Setup (GPU vs CPU)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()

    def determine_risk_level(self, emotion):
        """Categorizes emotions into clinical risk buckets."""
        emotion = emotion.lower()
        high_risk = ['panic', 'severe anxiety', 'depression', 'grief', 'suicidal', 'personality disorder']
        medium_risk = ['stress', 'anxiety', 'anger', 'burnout', 'fear', 'nervousness']
        
        if emotion in high_risk:
            return "High"
        elif emotion in medium_risk:
            return "Medium"
        else:
            return "Low"

    def predict(self, text):
        """The core engine: Text -> Tensor -> Prediction."""
        # Tokenize input
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=128, padding=True)
        inputs = {key: val.to(self.device) for key, val in inputs.items()}

        # Run inference
        with torch.no_grad(): 
            outputs = self.model(**inputs)
            logits = outputs.logits

        # Convert math to percentages
        probabilities = F.softmax(logits, dim=-1)
        confidence_score, predicted_class_id = torch.max(probabilities, dim=-1)

        # Map ID back to English
        class_id_number = predicted_class_id.item()
        predicted_label = self.emotion_map.get(class_id_number, "Unknown")
        risk_level = self.determine_risk_level(predicted_label)

        return {
            "text": text,
            "emotion": predicted_label,
            "confidence": round(confidence_score.item() * 100, 2),
            "risk_level": risk_level
        }

# --- Standard Testing Block ---
if __name__ == "__main__":
    predictor = MindGuardPredictor()
    sample_text = "I have a massive presentation tomorrow and my chest is tight."
    result = predictor.predict(sample_text)
    print("\n--- Prediction Results ---")
    print(f"Emotion: {result['emotion']} ({result['confidence']}%)")
    print(f"Risk Level: {result['risk_level']}")