"""
ml_classifier.py
------------------
A scikit-learn RandomForestClassifier that provides a *secondary* ML-based
risk-level signal, trained here on synthetically generated labeled examples.

IMPORTANT RESEARCH NOTE: This synthetic training set encodes the same clinical
logic used elsewhere (higher questionnaire scores + more negative-affect
language => higher risk) purely so the pipeline has a working, inspectable
scikit-learn model end-to-end, as required by the "AI Models" spec. Before
any real research or clinical use, this model MUST be retrained on properly
consented, IRB-approved, labeled participant data. Ship this file as a
placeholder/demo model only — never present its output as clinically
validated.

Feature vector (order matters):
    [phq9_normalized, gad7_normalized, pss10_normalized,
     avg_sentiment, total_depression_markers, total_anxiety_markers,
     total_stress_markers, avg_first_person_ratio, avg_absolutist_count]
"""

import random
from typing import Dict, List, Tuple

import numpy as np
from sklearn.ensemble import RandomForestClassifier

FEATURE_NAMES = [
    "phq9_normalized",
    "gad7_normalized",
    "pss10_normalized",
    "avg_sentiment",
    "total_depression_markers",
    "total_anxiety_markers",
    "total_stress_markers",
    "avg_first_person_ratio",
    "avg_absolutist_count",
]

RISK_LABELS = ["Low", "Moderate", "High"]


def _synthesize_training_data(n: int = 1500, seed: int = 42) -> Tuple[np.ndarray, np.ndarray]:
    rng = random.Random(seed)
    X: List[List[float]] = []
    y: List[int] = []

    for _ in range(n):
        # Sample a "true" latent risk tier, then generate correlated features.
        tier = rng.choices([0, 1, 2], weights=[0.5, 0.3, 0.2])[0]  # Low/Mod/High
        base = {0: 15, 1: 45, 2: 75}[tier]
        noise = lambda spread: max(0, min(100, base + rng.gauss(0, spread)))

        phq9 = noise(15)
        gad7 = noise(15)
        pss10 = noise(15)
        sentiment = max(-1.0, min(1.0, (50 - base) / 60 + rng.gauss(0, 0.2)))
        dep_markers = max(0, int(rng.gauss(base / 20, 1.5)))
        anx_markers = max(0, int(rng.gauss(base / 22, 1.5)))
        stress_markers = max(0, int(rng.gauss(base / 25, 1.5)))
        first_person = max(0.0, min(0.3, base / 400 + rng.gauss(0, 0.02)))
        absolutist = max(0, rng.gauss(base / 30, 1.0))

        X.append([phq9, gad7, pss10, sentiment, dep_markers, anx_markers,
                  stress_markers, first_person, absolutist])
        y.append(tier)

    return np.array(X), np.array(y)


class MLRiskClassifier:
    """Thin wrapper around a RandomForestClassifier trained on synthetic data
    at import time. Exposes predictions plus feature-importance based
    explanations."""

    def __init__(self):
        X, y = _synthesize_training_data()
        self.model = RandomForestClassifier(
            n_estimators=200, max_depth=6, random_state=42
        )
        self.model.fit(X, y)

    def predict(self, feature_vector: Dict[str, float]) -> Dict:
        x = np.array([[feature_vector.get(name, 0.0) for name in FEATURE_NAMES]])
        proba = self.model.predict_proba(x)[0]
        pred_idx = int(np.argmax(proba))

        importances = self.model.feature_importances_
        contributions = sorted(
            zip(FEATURE_NAMES, importances),
            key=lambda t: t[1],
            reverse=True,
        )[:5]

        return {
            "predicted_risk": RISK_LABELS[pred_idx],
            "class_probabilities": {
                RISK_LABELS[i]: round(float(p), 3) for i, p in enumerate(proba)
            },
            "top_contributing_features": [
                {"feature": name, "importance": round(float(imp), 3)}
                for name, imp in contributions
            ],
            "model_disclaimer": (
                "This ML signal is produced by a demonstration model trained "
                "on synthetic data and is NOT clinically validated. It is a "
                "secondary signal only; the primary risk determination uses "
                "the validated PHQ-9/GAD-7/PSS-10 scoring rules."
            ),
        }