src/train_chat_brain.py

import json
import os

import numpy as np
import pandas as pd
import torch
from datasets import Dataset
from dotenv import load_dotenv
from sklearn.metrics import (
    classification_report,
    f1_score,
    recall_score,
    accuracy_score,
    confusion_matrix,
)
from sklearn.model_selection import train_test_split
from sklearn.utils.class_weight import compute_class_weight
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
)
from sklearn.metrics import (
    precision_score,
    roc_auc_score,
    roc_curve,
    precision_recall_curve,
    auc,
)

from .model_registry import MODEL_REGISTRY
import matplotlib.pyplot as plt
import seaborn as sns

load_dotenv()


def tune_threshold(probs, y, min_recall=0.90):
    best = None
    for thr in np.linspace(0.05, 0.95, 181):
        preds = (probs >= thr).astype(int)
        rec = recall_score(y, preds, zero_division=0)
        f1 = f1_score(y, preds, zero_division=0)

        if rec >= min_recall:
            if best is None or f1 > best["f1"]:
                best = {"thr": float(thr), "recall": float(rec), "f1": float(f1)}

    if best is None:
        best = {"thr": 0.5, "recall": 0.0, "f1": 0.0}

    return best

def train_chat_brain(
    processed_dir="data/processed",
    models_dir="outputs/models",
    reports_dir="outputs/reports",
    max_length=128,   # 🔥 increased from 64
    epochs=3,
    lr=1e-5,
    bs_train=16,
    bs_eval=32,
    seed=42,
):
    os.makedirs(models_dir, exist_ok=True)
    os.makedirs(reports_dir, exist_ok=True)

    model_tag = os.getenv("CHAT_MODEL_TAG", "xlmr")
    model_name = MODEL_REGISTRY[model_tag]

    print(f"\n🚀 Training Model: {model_name}")

    # Load data
    df = pd.read_csv(os.path.join(processed_dir, "text_all_clean.csv"))
    df = df.dropna(subset=["text"]).reset_index(drop=True)

    # ✅ Stratified split
    train_df, temp_df = train_test_split(
        df, test_size=0.3, stratify=df["label"], random_state=seed
    )
    val_df, test_df = train_test_split(
        temp_df, test_size=0.5, stratify=temp_df["label"], random_state=seed
    )

    print("Train/Val/Test:", train_df.shape, val_df.shape, test_df.shape)

    # Tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    def tokenize(batch):
        return tokenizer(
            batch["text"],
            padding="max_length",
            truncation=True,
            max_length=max_length,
        )

    train_ds = Dataset.from_pandas(train_df[["text", "label"]]).map(tokenize, batched=True)
    val_ds = Dataset.from_pandas(val_df[["text", "label"]]).map(tokenize, batched=True)
    test_ds = Dataset.from_pandas(test_df[["text", "label"]]).map(tokenize, batched=True)

    train_ds.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
    val_ds.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
    test_ds.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

    # ✅ Proper class weights
    y_train = train_df["label"].values
    weights = compute_class_weight(class_weight="balanced", classes=np.array([0, 1]), y=y_train)
    class_weights = torch.tensor(weights, dtype=torch.float)

    print("Class weights:", class_weights.tolist())

    # Model
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=2,
        hidden_dropout_prob=0.3,
        attention_probs_dropout_prob=0.3,
    )

    # Weighted Trainer
    class WeightedTrainer(Trainer):
        def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
            labels = inputs["labels"]
            outputs = model(**inputs)
            logits = outputs.logits

            loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights.to(logits.device))
            loss = loss_fn(logits, labels)

            return (loss, outputs) if return_outputs else loss

    # Metrics
    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        preds = np.argmax(logits, axis=1)

        return {
            "accuracy": accuracy_score(labels, preds),
            "f1": f1_score(labels, preds),
            "recall": recall_score(labels, preds),
        }

    args = TrainingArguments(
        output_dir=os.path.join(models_dir, f"chat_brain_{model_tag}"),
        eval_strategy="epoch",
        save_strategy="no",
        learning_rate=lr,
        per_device_train_batch_size=bs_train,
        per_device_eval_batch_size=bs_eval,
        num_train_epochs=epochs,
        logging_steps=50,
        report_to="none",
        fp16=torch.cuda.is_available(),
        seed=seed,
    )

    trainer = WeightedTrainer(
        model=model,
        args=args,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        compute_metrics=compute_metrics,
    )

    # Train
    trainer.train()

    # Predictions
    def get_probs(ds):
        pred = trainer.predict(ds)
        logits = pred.predictions
        probs = torch.softmax(torch.tensor(logits), dim=1).numpy()[:, 1]
        labels = pred.label_ids
        return probs, labels

    val_probs, val_y = get_probs(val_ds)
    test_probs, test_y = get_probs(test_ds)

    # Threshold tuning
    best = tune_threshold(val_probs, val_y, min_recall=0.90)
    thr = best["thr"]

    test_pred = (test_probs >= thr).astype(int)

    # ===================== FINAL METRICS =====================
    report = classification_report(test_y, test_pred, digits=4)
    cm = confusion_matrix(test_y, test_pred)

    print("\n✅ Best Threshold:", best)
    print("\n📊 Classification Report:\n", report)
    print("\n📊 Confusion Matrix:\n", cm)

    # ===================== SAVE REPORT =====================
    report_path = os.path.join(reports_dir, f"{model_tag}_report.txt")

    with open(report_path, "w", encoding="utf-8") as f:
        f.write(f"Model: {model_name}\n\n")
        f.write("Best Threshold:\n")
        f.write(json.dumps(best, indent=2))
        f.write("\n\nClassification Report:\n")
        f.write(report)
        f.write("\n\nConfusion Matrix:\n")
        f.write(np.array2string(cm))

    print(f"\n✅ Report saved at: {report_path}")

    # ===================== CONFUSION MATRIX PLOT =====================
    plt.figure(figsize=(6, 5))

    sns.heatmap(
        cm,
        annot=True,
        fmt="d",
        cmap="Blues",
        xticklabels=["Non-Suicidal", "Suicidal"],
        yticklabels=["Non-Suicidal", "Suicidal"]
    )

    plt.xlabel("Predicted")
    plt.ylabel("Actual")
    plt.title(f"Confusion Matrix - {model_tag}")

    cm_path = os.path.join(reports_dir, f"{model_tag}_confusion_matrix.png")
    plt.savefig(cm_path)
    plt.close()

    print(f"✅ Confusion matrix saved at: {cm_path}")


    # ===================== EXTRA METRICS =====================

    accuracy = accuracy_score(test_y, test_pred)
    precision = precision_score(test_y, test_pred)
    recall = recall_score(test_y, test_pred)
    f1 = f1_score(test_y, test_pred)

    # AUC (uses probabilities, not labels)
    roc_auc = roc_auc_score(test_y, test_probs)

    # Specificity
    tn, fp, fn, tp = cm.ravel()
    specificity = tn / (tn + fp)

    # Balanced Accuracy
    balanced_acc = (recall + specificity) / 2

    metrics_dict = {
        "Accuracy": accuracy,
        "Precision": precision,
        "Recall": recall,
        "F1 Score": f1,
        "AUC": roc_auc,
        "Specificity": specificity,
        "Balanced Acc": balanced_acc,
    }

    print("\n📊 All Metrics:\n", metrics_dict)
    # ===================== METRICS BAR CHART =====================

# ===================== METRICS BAR CHART (WITH LABELS) =====================

    plt.figure(figsize=(8, 5))

    names = list(metrics_dict.keys())
    values = list(metrics_dict.values())

    bars = plt.bar(names, values)

    plt.ylim(0, 1)
    plt.title(f"Model Performance - {model_tag}")
    plt.xticks(rotation=30)

    # 🔥 Add value labels on top of each bar
    for bar in bars:
        height = bar.get_height()
        plt.text(
            bar.get_x() + bar.get_width() / 2,
            height + 0.02,                      # slight offset above bar
            f"{height:.3f}",                    # format value
            ha='center',
            va='bottom',
            fontsize=9
        )

    plt.tight_layout()

    bar_path = os.path.join(reports_dir, f"{model_tag}_metrics_bar.png")
    plt.savefig(bar_path)
    plt.close()

    print(f"✅ Metrics bar chart saved at: {bar_path}")

    # ===================== ROC CURVE =====================

    fpr, tpr, _ = roc_curve(test_y, test_probs)

    plt.figure()
    plt.plot(fpr, tpr, label=f"AUC = {roc_auc:.4f}")
    plt.plot([0, 1], [0, 1], linestyle="--")

    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title(f"ROC Curve - {model_tag}")
    plt.legend()

    roc_path = os.path.join(reports_dir, f"{model_tag}_roc_curve.png")
    plt.savefig(roc_path)
    plt.close()

    print(f"✅ ROC curve saved at: {roc_path}")

    # ===================== PRECISION-RECALL CURVE =====================

    precisions, recalls, _ = precision_recall_curve(test_y, test_probs)
    pr_auc = auc(recalls, precisions)

    plt.figure()
    plt.plot(recalls, precisions, label=f"PR AUC = {pr_auc:.4f}")

    plt.xlabel("Recall")
    plt.ylabel("Precision")
    plt.title(f"Precision-Recall Curve - {model_tag}")
    plt.legend()

    pr_path = os.path.join(reports_dir, f"{model_tag}_pr_curve.png")
    plt.savefig(pr_path)
    plt.close()

    print(f"✅ PR curve saved at: {pr_path}")

    # ===================== SAVE MODEL =====================
    model_dir = os.path.join(models_dir, f"chat_brain_{model_tag}")
    trainer.save_model(model_dir)
    tokenizer.save_pretrained(model_dir)

    with open(os.path.join(model_dir, "threshold.json"), "w") as f:
        json.dump(best, f, indent=2)

    return best