import joblib
import pandas as pd
from flask import Flask, request, jsonify

# --- Configuration ---
MODEL_FILE = 'pmad_prediction_model.joblib'
# These MUST match the order from your data_model_generator.py
MODEL_FEATURES = [
    'PPD_History', 'Anxiety_History', 'Birth_Trauma', 'Support_Score',
    'Sentiment_Score', 'Keyword_Guilt_Freq', 'Keyword_Worry_Freq',
    'Keyword_Intrusive_Freq', 'Keyword_Delusion_Freq', 
    'Avg_Sleep_Duration_hrs', 'Sleep_Fragmentation_Score', 
    'Avg_Heart_Rate_Variability', 'Activity_Steps_Avg'
]

# --- Load Model ---
try:
    model = joblib.load(MODEL_FILE)
    print(f"✅ Model '{MODEL_FILE}' loaded successfully.")
except FileNotFoundError:
    print(f"❌ ERROR: Model file '{MODEL_FILE}' not found. Did you upload it?")
    model = None
except Exception as e:
    print(f"❌ ERROR: Could not load model: {e}")
    model = None

# --- Init Flask App ---
app = Flask(__name__)

# --- Helper function to analyze the journal ---
def analyze_journal(journal_text):
    text = journal_text.lower()
    analysis = {
        'Sentiment_Score': 0.0, # Neutral
        'Keyword_Guilt_Freq': text.count('guilt') + text.count('failing') + text.count('bad mom'),
        'Keyword_Worry_Freq': text.count('worry') + text.count('anxious') + text.count('scared'),
        'Keyword_Intrusive_Freq': text.count('intrusive') + text.count('harm'),
        'Keyword_Delusion_Freq': text.count('delusion') + text.count('not real')
    }
    
    # Simple sentiment
    negative_words = analysis['Keyword_Guilt_Freq'] + analysis['Keyword_Worry_Freq']
    if negative_words > 2:
        analysis['Sentiment_Score'] = -0.7
    elif negative_words > 0:
        analysis['Sentiment_Score'] = -0.3
    
    if text.count('happy') > 0 or text.count('great') > 0:
         analysis['Sentiment_Score'] = 0.5
         
    return analysis

@app.route("/")
def home():
    return "Maia PMAD Model API (Complex) is running."

@app.route("/predict", methods=['POST'])
def predict():
    if model is None:
        return jsonify({"error": "Model is not loaded"}), 500

    try:
        # 1. Get data from Flutter app
        data = request.json
        print(f"Received data: {data}")

        # 2. Get the features we *can* get from the app
        # We use the app's 'mood' (1-5) to help set sentiment
        app_mood = data.get('mood', 3)
        app_sleep = data.get('sleep', 7.0)
        app_journal = data.get('journal', '')

        # 3. Analyze the journal to get keyword/sentiment features
        journal_analysis = analyze_journal(app_journal)
        
        # Override sentiment if mood is very low
        if app_mood <= 2 and journal_analysis['Sentiment_Score'] > -0.5:
            journal_analysis['Sentiment_Score'] = -0.5

        # 4. **INFER MISSING FEATURES**
        # We create a full 13-feature dictionary, setting defaults
        # for features the app doesn't collect.
        input_data = {
            # --- Inferred from Journal ---
            'Sentiment_Score': journal_analysis['Sentiment_Score'],
            'Keyword_Guilt_Freq': journal_analysis['Keyword_Guilt_Freq'],
            'Keyword_Worry_Freq': journal_analysis['Keyword_Worry_Freq'],
            'Keyword_Intrusive_Freq': journal_analysis['Keyword_Intrusive_Freq'],
            'Keyword_Delusion_Freq': journal_analysis['Keyword_Delusion_Freq'],
            
            # --- From App ---
            'Avg_Sleep_Duration_hrs': app_sleep,
            
            # --- Faked Defaults (Set to "average" or "neutral") ---
            'PPD_History': 0,         # Assume no history
            'Anxiety_History': 0,     # Assume no history
            'Birth_Trauma': 0,        # Assume no trauma
            'Support_Score': 3.0,     # Assume average support
            'Sleep_Fragmentation_Score': 0.5, # Assume average
            'Avg_Heart_Rate_Variability': 50.0, # Assume average
            'Activity_Steps_Avg': 4000  # Assume average
        }
        
        # 5. Create the DataFrame in the *exact* order
        df = pd.DataFrame([input_data], columns=MODEL_FEATURES)
        print(f"Model input vector:\n{df.to_string()}")

        # 6. Make prediction
        prediction_proba = model.predict_proba(df)
        
        # 7. Format the response
        class_map = {0: 'Healthy', 1: 'PPD', 2: 'PPA', 3: 'P-OCD', 4: 'Severe/PPP/PTSD'}
        probabilities = prediction_proba[0]
        
        output = {
            class_map[i]: prob for i, prob in enumerate(probabilities)
        }
        
        print(f"Prediction: {output}")
        return jsonify(output)

    except Exception as e:
        print(f"❌ Error during prediction: {e}")
        return jsonify({"error": str(e)}), 400

if __name__ == '__main__':
    # HF Spaces uses port 7860
    app.run(host='0.0.0.0', port=7860)