inference_engine.py

import os
import logging
import joblib
import json
import numpy as np
import torch
import pickle
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from textblob import TextBlob

# --- 1. SUPPRESS WARNINGS ---
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
logging.getLogger('tensorflow').setLevel(logging.FATAL)

MODEL_PATH = "models/mental_bert_model" 
META_MODEL_PATH = "models/model_meta_fusion.pkl" # Added Meta-Fusion Model Path

def load_config():
    with open("config.json", "r") as f:
        return json.load(f)

class ClinicalAI:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # --- LOAD NLP MODEL ---
        self.bert_model = None
        if os.path.exists(MODEL_PATH):
            try:
                self.tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
                self.bert_model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
                self.bert_model.to(self.device)
                self.bert_model.eval()
                self.labels = {0: 'ADHD', 1: 'Depression', 2: 'Anxiety', 3: 'Autism'}
            except Exception as e:
                print(f"Error loading Mental-BERT: {e}")
                
        # --- LOAD META-FUSION MODEL (.pkl) ---
        self.meta_model = None
        if os.path.exists(META_MODEL_PATH):
            try:
                self.meta_model = pickle.load(open(META_MODEL_PATH, "rb"))
            except Exception as e:
                print(f"Error loading Meta-Fusion Model: {e}")
        
        self.labels_list = ['ADHD', 'Depression', 'Anxiety', 'Autism']

    def _get_sentiment_score(self, text):
        return TextBlob(str(text)).sentiment.polarity

    def analyze_single_nlp_response(self, text, question_category=None):
        if not text or len(text) < 2:
            return {"diagnosis": "Insufficient Data", "confidence": 0.0, "probs": np.zeros(4), "analysis": "Insufficient data to analyze."}

        # --- GUARDRAIL 1: SENTIMENT & NEGATION CHECK ---
        blob = TextBlob(str(text))
        sentiment = blob.sentiment.polarity 
        text_lower = text.lower().strip()
        
        # --- NEW GUARDRAIL A: LEXICAL NEGATION CHECK (Fixes Keyword Bias) ---
        # If the answer is relatively short and contains explicit denial, override Mental-BERT
        is_short_answer = len(text_lower.split()) < 30
        
        if is_short_answer:
            # 1. Direct clear-cut denials
            if text_lower.startswith("no ") or text_lower.startswith("no,") or text_lower == "no":
                # Check to ensure they aren't saying "No, I am depressed"
                negative_confirmations = ["am depressed", "am anxious", "do worry", "struggle"]
                if not any(nc in text_lower for nc in negative_confirmations):
                    return {
                        "diagnosis": "No Significant Risk", "confidence": 0.95,
                        "probs": np.array([0.05, 0.05, 0.05, 0.05]),
                        "analysis": "No indicators detected (Patient explicitly denied symptom)."
                    }
                    
            # 2. Symptom-specific negations (e.g., "not worry", "did not struggle", "do not panic")
            symptom_negations = ["not struggle", "did not struggle", "don't struggle", 
                                 "not worry", "do not worry", "don't worry",
                                 "mind relaxes", "i am fine", "nothing like that"]
            
            if any(sn in text_lower for sn in symptom_negations):
                 return {
                    "diagnosis": "No Significant Risk", "confidence": 0.90,
                    "probs": np.array([0.05, 0.05, 0.05, 0.05]),
                    "analysis": "No indicators detected (Patient explicitly negated symptom triggers)."
                }
        
        # --- GUARDRAIL B: PHYSICAL/MEDICAL AILMENT CHECK ---
        # If user mentions physical pain, don't jump to psychiatric conclusions
        physical_keywords = ["pain", "ache", "arthritis", "ra", "joint", "migraine", "thyroid", "insomnia", "physical", "doctor", "disease", "illness", "injury", "headache"]
        cognitive_keywords = ["focus", "concentrate", "worry", "anxious", "sad", "depress", "fidget", "panic", "sleep", "tired", "energy", "attention", "memory", "forget", "stress"]
        
        has_physical = any(word in text_lower for word in physical_keywords)
        has_cognitive = any(word in text_lower for word in cognitive_keywords)
        
        # Guardrail triggers ONLY if it's purely a physical complaint with NO cognitive distress mentioned
        if has_physical and not has_cognitive and question_category == 'General':
             return {
                "diagnosis": "No Significant Risk",
                "confidence": 0.85, 
                "probs": np.array([0.01, 0.01, 0.01, 0.01]),
                "analysis": f"User described a physical condition without explicit cognitive distress markers."
            }

        negations = ["not", "don't", "do not", "never", "rarely", "no ", "stop", "quit", "manage", "fine", "good"]
        has_negation = any(neg in text_lower for neg in negations)

        # High positive sentiment = No Risk
        if sentiment > 0.25:
            return {
                "diagnosis": "No Significant Risk", "confidence": 0.95, 
                "probs": np.array([0.01, 0.01, 0.01, 0.01]),
                "analysis": f"User indicated no issues (Sentiment: {sentiment:.2f}, Conf: 95.0%)"
            }

        # Increased threshold to 0.05 so "I don't relax well" (which has 0.0 sentiment) bypasses this and goes to BERT.
        if has_negation and sentiment > 0.05:
             return {
                "diagnosis": "No Significant Risk", "confidence": 0.90, 
                "probs": np.array([0.01, 0.01, 0.01, 0.01]),
                "analysis": f"User indicated no issues (Negation detected, Sentiment: {sentiment:.2f})"
            }

        # --- LOGIC 3: CONTEXT AWARENESS ---
        if question_category and sentiment < 0.1:
            cat_map = {'A': 'ADHD', 'D': 'Depression', 'E': 'Anxiety', 'B': 'Autism', 'C': 'Social Communication Disorder'}
            mapped_diag = cat_map.get(question_category, None)
            
            if mapped_diag:
                probs = np.zeros(4)
                if mapped_diag == 'ADHD': probs[0] = 0.85
                elif mapped_diag == 'Depression': probs[1] = 0.85
                elif mapped_diag == 'Anxiety': probs[2] = 0.85
                elif mapped_diag == 'Autism': probs[3] = 0.85
                
                return {
                    "diagnosis": mapped_diag, "confidence": 0.85 + abs(sentiment)/2,
                    "probs": probs, "analysis": f"Detected {mapped_diag} indicators (Context-Aware)"
                }

        # --- GUARDRAIL 4: BERT MODEL INFERENCE (FALLBACK) ---
        if not self.bert_model: 
            return {"diagnosis": "Model Error", "confidence": 0.0, "probs": np.zeros(4), "analysis": "Model not loaded."}
        
        try:
            inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=128, padding=True).to(self.device)
            with torch.no_grad():
                outputs = self.bert_model(**inputs)
                probs = torch.nn.functional.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
            
            pred_id = np.argmax(probs)
            confidence = float(probs[pred_id])
            diagnosis = self.labels[pred_id]

            if confidence < 0.45:
                diagnosis = "No Significant Risk"
                confidence = 1.0 - confidence

            return {
                "diagnosis": diagnosis, "confidence": confidence, "probs": probs, 
                "analysis": f"AI Detected {diagnosis} ({confidence*100:.1f}%)"
            }
        except Exception as e:
            return {"diagnosis": "Error", "confidence": 0.0, "probs": np.zeros(4), "analysis": f"Processing error: {str(e)}"}

    def predict_symptom_category(self, text):
        if not self.bert_model: return None
        if TextBlob(text).sentiment.polarity > 0.2: return None

        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=128, padding=True).to(self.device)
        with torch.no_grad():
            outputs = self.bert_model(**inputs)
            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        
        pred_id = torch.argmax(probs).item()
        if probs[0][pred_id].item() > 0.4:
            return self.labels[pred_id]
        return None

    def run_meta_fusion(self, rf_probs, nlp_probs_list):
        """
        Fuses Questionnaire + NLP results using the saved Meta-Fusion .pkl model.
        Falls back to mathematical fusion if the model fails.
        """
        rf_vector = np.array([
            rf_probs.get('ADHD', 0.0), rf_probs.get('Depression', 0.0),
            rf_probs.get('Anxiety', 0.0), rf_probs.get('Autism', 0.0)
        ])
        
        if nlp_probs_list and len(nlp_probs_list) > 0:
            nlp_vector = np.mean(nlp_probs_list, axis=0)
            
            # --- INTEGRATING THE .PKL META-MODEL ---
            if self.meta_model is not None:
                try:
                    # Combine RF and NLP features into a single array for the model
                    combined_features = np.concatenate([rf_vector, nlp_vector]).reshape(1, -1)
                    final_vector = self.meta_model.predict_proba(combined_features)[0]
                    method = "Machine Learning Meta-Fusion (.pkl)"
                except Exception as e:
                    # Fallback to math if features don't match exactly
                    final_vector = (rf_vector * 0.6) + (nlp_vector * 0.4)
                    method = "Mathematical Meta-Fusion (Fallback)"
            else:
                final_vector = (rf_vector * 0.6) + (nlp_vector * 0.4)
                method = "Mathematical Meta-Fusion"
        else:
            final_vector = rf_vector
            method = "Questionnaire Only"

        max_prob = np.max(final_vector)
        pred_id = np.argmax(final_vector)
        
        if max_prob < 0.35:
            diagnosis = "No Significant Risk"
            confidence = 1.0 - max_prob 
        else:
            diagnosis = self.labels_list[pred_id]
            confidence = max_prob
        
        all_probs = {self.labels_list[i]: final_vector[i] for i in range(4)}

        return {
            "diagnosis": diagnosis,
            "confidence": confidence,
            "all_probs": all_probs,
            "method": method
        }

    def get_suggestions(self, diagnosis, severity):
        # Reads dynamically from config.json to support Admin Panel
        config = load_config()
        if diagnosis == "No Significant Risk":
            return ("🌟 Great News! Your screening suggests you are currently in a healthy cognitive state. "
                    "Keep up your good habits: maintain a balanced sleep schedule, eat well, and continue socializing.")

        diag_key = diagnosis.split()[0]
        if "Autism" in diagnosis: diag_key = "Autism"
        
        return config["suggestions"].get(diag_key, {}).get(severity, "Consult a specialist.")

# --- THE LOGIC TREE FUNCTION ---
def get_logic_driven_questions(flags):
    """
    STRICT LOGIC TREE (DSM-5 Rules). Reads dynamically from config.json.
    """
    config = load_config()
    nlp_qs = config.get("nlp_questions", {})

    active_sections = set()
    
    # 1. ADHD Logic
    if flags['S1']:
        active_sections.add('A') 
        active_sections.add('E') # Comorbidity rule

    # 2. Depression Logic
    if flags['S2']:
        active_sections.add('D') 
        active_sections.add('E') # Comorbidity rule

    # 3. Anxiety Logic
    if flags['S3']:
        active_sections.add('E')

    # 4. Differential Diagnosis (ASD vs SPCD)
    if flags['S4']:
        if flags['S5']:
            active_sections.add('B') # ASD
            if 'C' in active_sections: active_sections.remove('C')
        else:
            active_sections.add('C') # SPCD
            if 'B' in active_sections: active_sections.remove('B')
    
    # 5. Standalone S5 (Rigidity without Social) -> ASD indicators
    elif flags['S5']: 
        active_sections.add('B')

    # GENERATE QUESTIONS WITH CATEGORY METADATA
    questions = []
    
    if 'A' in active_sections:
        questions.append({'text': nlp_qs.get('A', "Describe a time recently when you had to do a boring task. Did you struggle to start?"), 'cat': 'A'})
    
    if 'D' in active_sections:
        questions.append({'text': nlp_qs.get('D', "Do you still enjoy your favorite hobbies, or does everything feel like too much effort?"), 'cat': 'D'})
    
    if 'E' in active_sections:
        questions.append({'text': nlp_qs.get('E', "When you have a quiet moment, does your mind relax, or do you constantly worry about the future?"), 'cat': 'E'})
    
    if 'B' in active_sections: # ASD
        questions.append({'text': nlp_qs.get('B', "How do you react when your daily routine is suddenly changed or if you are in a loud, crowded place?"), 'cat': 'B'})
    
    if 'C' in active_sections: # SPCD
        questions.append({'text': nlp_qs.get('C', "Do people ever tell you that you are blunt or rude when you don't mean to be?"), 'cat': 'C'})

    # FALLBACK: HEALTHY PATIENT CHECK
    if not questions:
        questions.append({'text': nlp_qs.get('General', "Please describe your general mood and mental state over the past two weeks. Do you feel happy and energetic?"), 'cat': 'General'})
        
    return questions