src/predict.py

import json
import os

import joblib
import numpy as np
import pandas as pd
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


def _prob_to_level(p: float, thr_high: float, thr_med: float) -> str:
    if p >= thr_high:
        return "High"
    if p >= thr_med:
        return "Medium"
    return "Low"


def _fuse_level(chat_level: str, prof_level: str, s: float, thr_high: float, thr_med: float) -> str:
    # Safety-first overrides (reduce false negatives)
    if chat_level == "High":
        return "High"
    if chat_level == "Medium" and prof_level == "High":
        return "High"

    if s >= thr_high:
        return "High"
    if s >= thr_med:
        return "Medium"
    return "Low"


class SuicideRiskPredictor:
    """
    Real-time predictor for ONE user input:
      - text (string)
      - profile (dict of structured features matching your BD columns)
    """

    def __init__(
        self,
        model_name: str = "banglabert",
        models_dir: str = "outputs/models",
        artifacts_dir: str = "outputs/artifacts",
        fusion_meta_path: str = None,
        profile_meta_path: str = "outputs/artifacts/profile_risk_meta.joblib",
    ):
        self.model_name = model_name

        # ---- Load fusion config (weights/thresholds) ----
        if not fusion_meta_path:
            fusion_meta_path = os.path.join(artifacts_dir, f"fusion_meta_{model_name}.json")

        if os.path.exists(fusion_meta_path):
            with open(fusion_meta_path, "r", encoding="utf-8") as f:
                meta = json.load(f)
        else:
            meta = {}

        # Safely extract weights (with fallback)
        weights = meta.get("weights", {})
        self.w_chat = float(weights.get("w_chat", weights.get("chat", 0.60)))
        self.w_prof = float(weights.get("w_prof", weights.get("profile", 0.40)))

        # Safely extract thresholds
        chat_thrs = meta.get("chat_thresholds", {"high": 0.70, "medium": 0.40})
        self.chat_thr_high = float(chat_thrs.get("high", 0.70))
        self.chat_thr_med = float(chat_thrs.get("medium", 0.40))

        final_thrs = meta.get("final_thresholds", {"high": 0.65, "medium": 0.40})
        self.final_thr_high = float(final_thrs.get("high", 0.65))
        self.final_thr_med = float(final_thrs.get("medium", 0.40))

        self.max_len = int(meta.get("max_len", 64))

        # ---- Chat model ----
        self.chat_model_dir = os.path.join(models_dir, f"chat_brain_{model_name}")
        self.tokenizer = AutoTokenizer.from_pretrained(self.chat_model_dir)
        self.chat_model = AutoModelForSequenceClassification.from_pretrained(self.chat_model_dir)
        self.chat_model.eval()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.chat_model.to(self.device)

        # ---- Profile artifacts ----
        self.pre = joblib.load(os.path.join(artifacts_dir, "profile_preprocessor.joblib"))
        self.gmm = joblib.load(os.path.join(artifacts_dir, "gmm.joblib"))
        self.kde = joblib.load(os.path.join(artifacts_dir, "kde.joblib"))
        self.profile_meta = joblib.load(profile_meta_path) if os.path.exists(profile_meta_path) else None

        # calibration percentiles for risk conversion
        if self.profile_meta:
            self.low_gmm, self.high_gmm = self.profile_meta["calibration_ll_percentiles_gmm"]
            self.low_kde, self.high_kde = self.profile_meta["calibration_ll_percentiles_kde"]
            self.p70 = self.profile_meta["p70"]
            self.p90 = self.profile_meta["p90"]
            self.kde_bandwidth = self.profile_meta["kde_bandwidth"]
            self.w_gmm = self.profile_meta["w_gmm"]
            self.w_kde = self.profile_meta["w_kde"]
        else:
            # fallback if meta not found
            self.low_gmm, self.high_gmm = (-50.0, -5.0)
            self.low_kde, self.high_kde = (-50.0, -5.0)
            self.p70, self.p90 = (0.70, 0.90)
            self.w_gmm, self.w_kde = (0.6, 0.4)

    def _chat_prob(self, text: str) -> float:
        enc = self.tokenizer(
            [str(text)],
            truncation=True,
            padding=True,
            max_length=self.max_len,
            return_tensors="pt",
        ).to(self.device)

        with torch.no_grad():
            logits = self.chat_model(**enc).logits
            p = torch.softmax(logits, dim=1)[:, 1].detach().cpu().numpy()[0]
        return float(p)

    def _ll_to_risk(self, ll: np.ndarray, low: float, high: float) -> np.ndarray:
        norm = np.clip((ll - low) / (high - low + 1e-9), 0, 1)
        return 1 - norm

    def _profile_risk(self, profile_dict: dict) -> tuple[float, str]:
        """
        profile_dict keys should match your BD structured columns AFTER preprocessing.
        Missing keys are allowed (will become NaN and get imputed).
        """
        expected_cols = [
            "age_group", "age", "gender", "profession_group", "religion",
            "hometown", "reason", "reason_description", "time", "temperature",
            "feels_like", "temp_min", "temp_max", "air_pressure", "air_humidity",
            "wind_speed", "wind_deg", "clouds_sky", "weather_main", "weather_description"
        ]
        
        # Populate missing columns with NaN
        filled_dict = {}
        for col in expected_cols:
            filled_dict[col] = profile_dict.get(col, np.nan)

        df = pd.DataFrame([filled_dict])
        X = self.pre.transform(df)

        ll_gmm = self.gmm.score_samples(X)
        ll_kde = self.kde.score_samples(X)

        risk_gmm = self._ll_to_risk(ll_gmm, self.low_gmm, self.high_gmm)[0]
        risk_kde = self._ll_to_risk(ll_kde, self.low_kde, self.high_kde)[0]

        risk_profile = float(self.w_gmm * risk_gmm + self.w_kde * risk_kde)

        # level using stored percentiles
        if risk_profile >= self.p90:
            level = "High"
        elif risk_profile >= self.p70:
            level = "Medium"
        else:
            level = "Low"
        return risk_profile, level

    def predict_one(self, text: str, profile: dict) -> dict:
        chat_prob = self._chat_prob(text)
        chat_level = _prob_to_level(chat_prob, self.chat_thr_high, self.chat_thr_med)

        risk_profile, risk_level_profile = self._profile_risk(profile)

        final_risk_score = float(self.w_chat * chat_prob + self.w_prof * risk_profile)
        final_risk_level = _fuse_level(
            chat_level, risk_level_profile, final_risk_score,
            self.final_thr_high, self.final_thr_med
        )

        return {
            "chat_prob": chat_prob,
            "chat_level": chat_level,
            "risk_profile": risk_profile,
            "risk_level_profile": risk_level_profile,
            "final_risk_score": final_risk_score,
            "final_risk_level": final_risk_level,
        }