src/fuse_layer.py

# =========================
# src/fuse_layer.py (FINAL CLEAN - NO OLD SCORE ✅)
# =========================
import json
import os
from glob import glob

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

load_dotenv()


# =========================
# MODEL RESOLUTION
# =========================
def _latest_chat_model_dir(models_dir: str) -> str:
    candidates = glob(os.path.join(models_dir, "chat_brain_*"))
    candidates = [d for d in candidates if os.path.isdir(d)]
    if not candidates:
        raise FileNotFoundError(f"No chat_brain_* folders found in {models_dir}")
    candidates.sort(key=lambda d: os.path.getmtime(d), reverse=True)
    return candidates[0]


def _resolve_model_dir(models_dir: str):
    tag = os.getenv("CHAT_MODEL_TAG", "").strip().lower()
    if tag:
        model_dir = os.path.join(models_dir, f"chat_brain_{tag}")
        if not os.path.exists(model_dir):
            raise FileNotFoundError(f"Model not found: {model_dir}")
        return tag, model_dir

    model_dir = _latest_chat_model_dir(models_dir)
    tag = os.path.basename(model_dir).replace("chat_brain_", "")
    return tag, model_dir


# =========================
# CHAT MODEL
# =========================
def batch_predict_proba(texts, model_dir, batch_size=32, max_len=64):
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
    model.eval()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)

    probs = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]

        enc = tokenizer(
            batch,
            truncation=True,
            padding=True,
            max_length=max_len,
            return_tensors="pt"
        ).to(device)

        with torch.no_grad():
            logits = model(**enc).logits
            p = torch.softmax(logits, dim=1)[:, 1].cpu().numpy()

        probs.append(p)

    return np.concatenate(probs)


# =========================
# LEVEL FUNCTIONS
# =========================
def prob_to_chat_level(p, thr_high=0.70, thr_med=0.40):
    if p >= thr_high:
        return "High"
    if p >= thr_med:
        return "Medium"
    return "Low"


def ensure_profile_levels(df):
    df = df.copy()
    if "risk_level_profile" not in df.columns:
        p70, p90 = np.percentile(df["risk_profile"], [70, 90])

        def level(x):
            if x >= p90:
                return "High"
            if x >= p70:
                return "Medium"
            return "Low"

        df["risk_level_profile"] = df["risk_profile"].apply(level)

    return df


def fuse_level(chat_level, profile_level, final_score,
               final_thr_high=0.65, final_thr_med=0.40):

    if chat_level == "High":
        return "High"

    if chat_level == "Medium" and profile_level == "High":
        return "High"

    if final_score >= final_thr_high:
        return "High"

    if final_score >= final_thr_med:
        return "Medium"

    return "Low"


# =========================
# FILE NAMING
# =========================
def _make_tagged_filename(name, tag):
    root, ext = os.path.splitext(name)
    if not ext:
        ext = ".csv"
    return f"{root}_{tag}{ext}"


# =========================
# MAIN FUSION
# =========================
def run_fusion(
    processed_dir="data/processed",
    models_dir="outputs/models",
    artifacts_dir="outputs/artifacts",
    text_input_csv="text_all_clean.csv",
    profile_csv="bd_profile_with_risk.csv",
    output_csv="fusion_final_output.csv",
    id_col="id",
    chat_thr_high=0.70,
    chat_thr_med=0.40,
    w_chat=0.60,
    w_prof=0.40,
    final_thr_high=0.65,
    final_thr_med=0.40,
    allow_index_join=True,
):

    # =========================
    # LOAD MODEL
    # =========================
    tag, model_dir = _resolve_model_dir(models_dir)
    print(f"\n✅ Using model: {tag}")

    text_path = os.path.join(processed_dir, text_input_csv)
    prof_path = os.path.join(processed_dir, profile_csv)

    df_text = pd.read_csv(text_path)
    df_prof = pd.read_csv(prof_path)

    df_prof = ensure_profile_levels(df_prof)

    # =========================
    # 1) CHAT PROBS
    # =========================
    df_text["chat_prob"] = batch_predict_proba(
        df_text["text"].astype(str).tolist(),
        model_dir
    )

    df_text["chat_level"] = df_text["chat_prob"].apply(
        lambda p: prob_to_chat_level(p, chat_thr_high, chat_thr_med)
    )

    # SAVE CHAT OUTPUT ✅
    chat_out = os.path.join(
        processed_dir,
        _make_tagged_filename("chat_with_probs.csv", tag)
    )
    df_text.to_csv(chat_out, index=False)
    print("✅ Saved:", chat_out)

    # =========================
    # 2) MERGE
    # =========================
    if id_col in df_text.columns and id_col in df_prof.columns:
        df = df_text.merge(
            df_prof[[id_col, "risk_profile", "risk_level_profile"]],
            on=id_col
        )
    else:
        n = min(len(df_text), len(df_prof))
        df = df_text.iloc[:n].copy()
        df["risk_profile"] = df_prof.iloc[:n]["risk_profile"].values
        df["risk_level_profile"] = df_prof.iloc[:n]["risk_level_profile"].values

    # =========================
    # 3) NORMALIZATION
    # =========================
    min_r = df["risk_profile"].min()
    max_r = df["risk_profile"].max()

    df["risk_profile_norm"] = (
        df["risk_profile"] - min_r
    ) / (max_r - min_r + 1e-8)

    # =========================
    # 4) FUSION (ONLY CLEAN VERSIONS)
    # =========================

    # NORMALIZED ✅
    df["final_score_norm"] = (
        w_chat * df["chat_prob"]
        + w_prof * df["risk_profile_norm"]
    )

    # ADAPTIVE 🔥 (BEST)
    pc = df["chat_prob"]
    pr = df["risk_profile_norm"]

    alpha = pc / (pc + pr + 1e-8)
    alpha = np.clip(alpha, 0.6, 0.95)  # 🔥 gives text dominance, prevents profile from ruining text recall

    df["final_score_adaptive"] = (
        alpha * pc + (1 - alpha) * pr
    )

    # =========================
    # FINAL OUTPUT
    # =========================
    df["final_risk_score"] = df["final_score_adaptive"]

    df["final_risk_level"] = df.apply(
        lambda row: fuse_level(
            row["chat_level"],
            row["risk_level_profile"],
            row["final_risk_score"],
            final_thr_high,
            final_thr_med
        ),
        axis=1
    )

    # =========================
    # SAVE FINAL FUSION ✅
    # =========================
    fusion_out = os.path.join(
        processed_dir,
        _make_tagged_filename(output_csv, tag)
    )

    df.to_csv(fusion_out, index=False)
    print("✅ Saved:", fusion_out)

    # =========================
    # META
    # =========================
    os.makedirs(artifacts_dir, exist_ok=True)

    meta = {
        "model": tag,
        "fusion": "adaptive + normalized",
        "weights": {"chat": w_chat, "profile": w_prof}
    }

    with open(os.path.join(artifacts_dir, f"fusion_meta_{tag}.json"), "w") as f:
        json.dump(meta, f, indent=2)

    return {
        "chat_csv": chat_out,
        "fusion_csv": fusion_out
    }