"""
============================================================
Multi-Model Classifier - Ensemble Approach
============================================================
استخدام 4 نماذج (Primary, Secondary, Tertiary الجاهز، و Emotion).
"""
import os
import re
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from transformers import AutoTokenizer, AutoModel, pipeline
import numpy as np
from typing import List, Dict
from huggingface_hub import hf_hub_download

class MentalHealthDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            str(self.texts[idx]),
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(),
            'attention_mask': encoding['attention_mask'].squeeze(),
            'labels': torch.tensor(self.labels[idx], dtype=torch.long)
        }


class MentalHealthClassifier(nn.Module):
    def __init__(self, model_name, n_classes, dropout=0.3):
        super().__init__()
        self.model_name = model_name
        try:
            self.bert = AutoModel.from_pretrained(model_name)
        except Exception as e:
            print(f"⚠ Failed loading {model_name}: {e}. Falling back to distilbert-base-uncased")
            self.bert = AutoModel.from_pretrained("distilbert-base-uncased")
            self.model_name = "distilbert-base-uncased"

        self.drop = nn.Dropout(dropout)
        self.fc = nn.Linear(self.bert.config.hidden_size, n_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = outputs.last_hidden_state[:, 0, :]
        dropped = self.drop(pooled)
        return self.fc(dropped)


class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0, class_weights=None):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.class_weights = class_weights

    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(inputs, targets, weight=self.class_weights, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()


class EnsembleDiagnoser:
    EMOTION_TO_LABEL_BIAS = {
        'sadness': {'Depression': 0.4, 'Suicide': 0.2},
        'fear': {'Anxiety/Stress': 0.4, 'PTSD': 0.3},
        'anger': {'Bipolar': 0.2, 'PTSD': 0.2},
        'disgust': {'Depression': 0.2},
        'joy': {'Normal': 0.2, 'Bipolar': 0.4},
        'surprise': {'Normal': 0.1, 'Bipolar': 0.2},
        'neutral': {'Normal': 0.3},
    }

    def __init__(self, models_config: dict, labels: List[str], device=None):
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.labels = labels
        self.label2idx = {lbl: i for i, lbl in enumerate(labels)}
        self.idx2label = {i: lbl for i, lbl in enumerate(labels)}
        self.models_config = models_config

        self.primary_model = None
        self.primary_tokenizer = None
        self.secondary_model = None
        self.secondary_tokenizer = None
        self.emotion_pipe = None
        self.tertiary_pipe = None

    def load_pretrained(self, primary_path: str = None, secondary_path: str = None):
        REPO_ID = "Ziad9022/Mentallico-Weights"

        # ====== Primary Model ======
        primary_cfg = self.models_config['primary']
        if primary_path:
            if not os.path.exists(primary_path):
                filename = os.path.basename(primary_path)
                print(f"📥 Downloading primary weights '{filename}' from Hub...")
                primary_path = hf_hub_download(repo_id=REPO_ID, filename=filename)

            ckpt = torch.load(primary_path, map_location=self.device, weights_only=False)
            model_name = ckpt.get('model_name', primary_cfg['name'])
            self.primary_model = MentalHealthClassifier(model_name, len(self.labels))
            self.primary_model.load_state_dict(ckpt['model_state_dict'])
            self.primary_model.to(self.device).eval()
            self.primary_tokenizer = AutoTokenizer.from_pretrained(model_name)
            print(f"✓ Loaded primary model from Hub cache: {primary_path}")

        # ====== Secondary Model ======
        secondary_cfg = self.models_config['secondary']
        if secondary_path:
            if not os.path.exists(secondary_path):
                filename = os.path.basename(secondary_path)
                print(f"📥 Downloading secondary weights '{filename}' from Hub...")
                secondary_path = hf_hub_download(repo_id=REPO_ID, filename=filename)

            ckpt = torch.load(secondary_path, map_location=self.device, weights_only=False)
            model_name = ckpt.get('model_name', secondary_cfg['name'])
            self.secondary_model = MentalHealthClassifier(model_name, len(self.labels))
            self.secondary_model.load_state_dict(ckpt['model_state_dict'])
            self.secondary_model.to(self.device).eval()
            self.secondary_tokenizer = AutoTokenizer.from_pretrained(model_name)
            print(f"✓ Loaded secondary model from Hub cache: {secondary_path}")

        device_id = 0 if self.device.type == 'cuda' else -1

        # Emotion Pipeline
        try:
            emotion_name = self.models_config['emotion']['name']
            self.emotion_pipe = pipeline("text-classification", model=emotion_name, top_k=None, device=device_id)
            print(f"✓ Loaded emotion pipeline: {emotion_name}")
        except Exception as e:
            print(f"⚠ Could not load emotion pipeline: {e}")
            self.emotion_pipe = None

        # Tertiary Pipeline
        try:
            tertiary_name = self.models_config['tertiary']['name']
            self.tertiary_pipe = pipeline("text-classification", model=tertiary_name, top_k=None, device=device_id)
            print(f"✓ Loaded tertiary pipeline: {tertiary_name}")
        except Exception as e:
            print(f"⚠ Could not load tertiary pipeline: {e}")
            self.tertiary_pipe = None

    @torch.no_grad()
    def _predict_single(self, model, tokenizer, text: str, max_len=128):
        if model is None: return None
        encoding = tokenizer(str(text), max_length=max_len, padding='max_length', truncation=True, return_tensors='pt')
        input_ids = encoding['input_ids'].to(self.device)
        attention_mask = encoding['attention_mask'].to(self.device)
        logits = model(input_ids, attention_mask)
        probs = F.softmax(logits, dim=1).cpu().numpy()[0]
        return probs

    def _emotion_signal(self, text: str) -> np.ndarray:
        if self.emotion_pipe is None: return np.zeros(len(self.labels))
        try:
            results = self.emotion_pipe(text[:512])
            if isinstance(results, list) and len(results) > 0:
                if isinstance(results[0], list): results = results[0]
            bias = np.zeros(len(self.labels))
            for r in results:
                emotion, score = r['label'].lower(), r['score']
                if emotion in self.EMOTION_TO_LABEL_BIAS:
                    for lbl, weight in self.EMOTION_TO_LABEL_BIAS[emotion].items():
                        if lbl in self.label2idx:
                            bias[self.label2idx[lbl]] += score * weight
            if bias.sum() > 0: bias = bias / bias.sum()
            return bias
        except:
            return np.zeros(len(self.labels))

    def _tertiary_signal(self, text: str) -> np.ndarray:
        if self.tertiary_pipe is None: return np.zeros(len(self.labels))
        try:
            results = self.tertiary_pipe(text[:512])
            if isinstance(results, list) and len(results) > 0:
                if isinstance(results[0], list): results = results[0]
            
            bias = np.zeros(len(self.labels))
            mapping = {
                'depression': 'Depression', 'anxiety': 'Anxiety/Stress',
                'stress': 'Anxiety/Stress', 'bipolar': 'Bipolar',
                'ptsd': 'PTSD', 'suicide': 'Suicide', 'normal': 'Normal'
            }
            for r in results:
                lbl, score = r['label'].lower(), r['score']
                for key, mapped_lbl in mapping.items():
                    if key in lbl and mapped_lbl in self.label2idx:
                        bias[self.label2idx[mapped_lbl]] += score
            if bias.sum() > 0: bias = bias / bias.sum()
            return bias
        except:
            return np.zeros(len(self.labels))

    def predict(self, text: str, return_details=False) -> dict:
        individual_probs = {}
        
        individual_probs['primary'] = self._predict_single(self.primary_model, self.primary_tokenizer, text) if self.primary_model else None
        individual_probs['secondary'] = self._predict_single(self.secondary_model, self.secondary_tokenizer, text) if self.secondary_model else None
        individual_probs['tertiary'] = self._tertiary_signal(text)
        individual_probs['emotion_bias'] = self._emotion_signal(text)

        final_probs = np.zeros(len(self.labels))
        total_weight = 0.0

        if individual_probs['primary'] is not None:
            w = self.models_config['primary']['weight']
            final_probs += w * individual_probs['primary']
            total_weight += w

        if individual_probs['secondary'] is not None:
            w = self.models_config['secondary']['weight']
            final_probs += w * individual_probs['secondary']
            total_weight += w

        if individual_probs['tertiary'].sum() > 0:
            w = self.models_config['tertiary']['weight']
            final_probs += w * individual_probs['tertiary']
            total_weight += w

        if individual_probs['emotion_bias'].sum() > 0:
            w = self.models_config['emotion']['weight']
            final_probs += w * individual_probs['emotion_bias']
            total_weight += w

        if total_weight > 0:
            final_probs = final_probs / total_weight

        if final_probs.sum() == 0:
            return {'label': 'Normal', 'confidence': 0.0, 'probabilities': dict(zip(self.labels, [1.0/len(self.labels)]*len(self.labels))), 'error': 'No models loaded'}

        # ========================================================
        # 💉 CLINICAL HEURISTIC BOOSTING (To fix Class Imbalance)
        # ========================================================
        text_lower = text.lower()
        
        # 1. PTSD Booster
        ptsd_anchors = ['flashback', 'nightmare', 'accident', 'trauma', 'panic attack']
        if any(word in text_lower for word in ptsd_anchors):
            if 'PTSD' in self.label2idx:
                final_probs[self.label2idx['PTSD']] *= 5
                
        # 2. Bipolar Booster
        bipolar_anchors = ['superpower', 'endless energy', 'racing', 'billionaire', "haven't slept"]
        if any(word in text_lower for word in bipolar_anchors):
            if 'Bipolar' in self.label2idx:
                final_probs[self.label2idx['Bipolar']] *= 3.5

        # إعادة الحساب بعد التضخيم
        if final_probs.sum() > 0:
            final_probs = final_probs / final_probs.sum()
            
        pred_idx = int(np.argmax(final_probs))
        pred_label = self.idx2label[pred_idx]
        confidence = float(final_probs[pred_idx])
        probabilities_dict = dict(zip(self.labels, final_probs.tolist()))
        # ========================================================

        # ========================================================
        # 🚨 START OF EMERGENCY SAFETY OVERRIDE (SUICIDE PROTOCOL)
        # ========================================================
        suicide_keywords = [
            r'\bsuicide\b', r'\bkill myself\b', r'\bend my life\b', r'\bdie\b', r'\bno hope\b',
            'انتحر', 'اموت نفسي', 'انهي حياتي', 'اقتل نفسي'
        ]
        
        has_suicide_keywords = any(re.search(word, text_lower) for word in suicide_keywords)
        suicide_prob = probabilities_dict.get('Suicide', 0.0)
        
        is_suicide_probable = suicide_prob > 0.35 

        if has_suicide_keywords or is_suicide_probable:
            pred_label = 'Suicide'
            confidence = max(suicide_prob, 0.99)
            probabilities_dict['Suicide'] = confidence
        # ========================================================
        # 🚨 END OF EMERGENCY SAFETY OVERRIDE
        # ========================================================

        result = {
            'label': pred_label,
            'confidence': confidence,
            'probabilities': probabilities_dict,
        }
        if return_details:
            result['details'] = {k: (v.tolist() if isinstance(v, np.ndarray) else v) for k, v in individual_probs.items()}
        return result