import numpy as np
import pandas as pd
import os

from sklearn.model_selection import train_test_split, RandomizedSearchCV

import lightgbm as lgb
from scipy.stats import randint as sp_randint, uniform as sp_uniform

from sklearn.model_selection import StratifiedKFold
from sklearn.base import clone


from sklearn.metrics import (
    roc_auc_score, recall_score, precision_score, accuracy_score, 
    confusion_matrix, fbeta_score, precision_recall_curve, roc_curve
)

/home/shezin/.local/lib/python3.8/site-packages/pandas/core/computation/expressions.py:20: UserWarning: Pandas requires version '2.7.3' or newer of 'numexpr' (version '2.7.1' currently installed).
  from pandas.core.computation.check import NUMEXPR_INSTALLED
/usr/local/lib/python3.8/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.24.4
  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"

# Load cleaned dataframe
df = pd.read_pickle("../data/prepared/NSDUH_2023_prepared.pkl")  # <- change filename

# Binary target and weights
TARGET = "IRAMDEYR"
WEIGHT = "ANALWT2_C"

df[TARGET] = (df[TARGET] == 1).astype(int)
y = df[TARGET]
w = df[WEIGHT]

# Load tiers
from tiering_preparation import TIERS

# ----- choose which tier to run -----
tier_name = "Tier_3_Personalized"      # "Tier_1_Basic" or "Tier_2_Clinical" or "Tier_3_Personalized"
feature_list = TIERS[tier_name]

X = df[feature_list]

X_train, X_test, y_train, y_test, w_train, w_test = train_test_split(
    X, y, w,
    test_size=0.2,
    stratify=y,
    random_state=42,
)

1. LightGBM base model and hyperparameter search space

In this step, we define:

1. A base LightGBM classifier that knows our problem is binary and class-imbalanced, and
2. A set of hyperparameter distributions that RandomizedSearchCV will sample from when searching for good models.

# Base LightGBM model
base_lgb = lgb.LGBMClassifier(
    objective="binary",
    boosting_type="gbdt",
    class_weight="balanced",
    random_state=42,
)

# Parameter distributions for random search
param_dist = {
    "num_leaves":       sp_randint(31, 128),
    "max_depth":        sp_randint(3, 12),
    "learning_rate":    sp_uniform(0.005, 0.05),   # 0.01–0.10
    "n_estimators":     sp_randint(600, 1500),
    "min_child_samples": sp_randint(5, 45),
    "subsample":        sp_uniform(0.6, 0.4),     # 0.6–1.0
    "colsample_bytree": sp_uniform(0.6, 0.4),     # 0.6–1.0
}

2. Hyperparameter tuning using RandomizedSearchCV (Recall-optimized model)

In this step, we run a RandomizedSearchCV to find the LightGBM hyperparameters that maximize recall during 5-fold cross-validation.

The goal here is to create a model that catches as many MDE cases as possible, which matches our clinical objective of minimizing false negatives (important for screening and triage).

Why Randomized Search?

Instead of exhaustively trying every possible combination (GridSearch), RandomizedSearchCV:

• Samples random combinations from the parameter distributions we defined earlier
• Allows exploring a broad hyperparameter space
• Is much faster than GridSearch
• Often finds equally good (or better) results with fewer evaluations

Why optimize for recall?

• MDE cases are relatively rare in the population
• Missing a positive case is clinically costly
• High recall ensures that we correctly identify the majority of true MDE cases
• Later, we will tune thresholds to improve NPV even further

The code below runs the search and returns the best model based on mean recall across folds.

random_search_recall = RandomizedSearchCV(
    estimator=base_lgb,
    param_distributions=param_dist,
    n_iter=40,
    scoring="recall",
    cv=5,
    n_jobs=-1,
    random_state=42,
    verbose=1,
)

random_search_recall.fit(X_train, y_train, sample_weight=w_train)
best_lgb_recall = random_search_recall.best_estimator_

from joblib import dump
model_path = f"../models/lightgbm_{tier_name.lower()}_recall.joblib"
dump(best_lgb_recall, model_path)

print("Best params (Recall optimized):")
print(random_search_recall.best_params_)

Fitting 5 folds for each of 40 candidates, totalling 200 fits
Best params (Recall optimized):
{'colsample_bytree': 0.6028265220878869, 'learning_rate': 0.006153121252070788, 'max_depth': 5, 'min_child_samples': 7, 'n_estimators': 1084, 'num_leaves': 81, 'subsample': 0.8721230154351118}

3. (NOT USING) Hyperparameter tuning using RandomizedSearchCV (AUC-optimized model)

In this step, we run a second hyperparameter search - this time optimizing for ROC AUC instead of recall.

Why optimize AUC separately?

While the recall-optimized model is best for catching as many MDE cases as possible, the AUC-optimized model helps us:

• Identify the model that best separates positive vs negative cases overall
• Improve the ranking quality of predicted probabilities
• Ensure that the model produces meaningful risk scores, which is important for threshold tuning later
• Compare two complementary optimization strategies (Recall vs AUC), which strengthens our study methodology

# random_search_auc = RandomizedSearchCV(
#     estimator=base_lgb,
#     param_distributions=param_dist,
#     n_iter=40,
#     scoring="roc_auc",
#     cv=5,
#     n_jobs=-1,
#     random_state=42,
#     verbose=1,
# )

# random_search_auc.fit(X_train, y_train, sample_weight=w_train)
# best_lgb_auc = random_search_auc.best_estimator_

# print("Best params (AUC optimized):")
# print(random_search_auc.best_params_)

4. Evaluation helper function

def evaluate_model(y_true, preds, probs):
    tn, fp, fn, tp = confusion_matrix(y_true, preds).ravel()

    specificity = tn / (tn + fp)
    sensitivity = recall_score(y_true, preds)  # same as recall
    ppv = precision_score(y_true, preds)
    npv = tn / (tn + fn)

    return {
        "Accuracy": accuracy_score(y_true, preds),
        "Sensitivity": sensitivity,
        "Specificity": specificity,
        "PPV": ppv,
        "NPV": npv,
        "AUC": roc_auc_score(y_true, probs),
        "F1": fbeta_score(y_true, preds, beta=1),
        "F2": fbeta_score(y_true, preds, beta=2),
    }

5. Saving CV fold wise Results, Mean and SD (only on Train data) - threshold=0.5

# 5-fold CV table on TRAIN set (no test leakage)
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

metrics_order = ["Accuracy", "Sensitivity", "Specificity", "PPV", "NPV", "AUC", "F1", "F2"]
fold_rows = []

for fold_i, (tr_idx, val_idx) in enumerate(skf.split(X_train, y_train), start=1):
    X_tr, X_val = X_train.iloc[tr_idx], X_train.iloc[val_idx]
    y_tr, y_val = y_train.iloc[tr_idx], y_train.iloc[val_idx]
    w_tr = w_train.iloc[tr_idx]

    model = clone(best_lgb_recall)  
    model.fit(X_tr, y_tr, sample_weight=w_tr)

    proba = model.predict_proba(X_val)[:, 1]
    pred  = (proba >= 0.5).astype(int)

    m = evaluate_model(y_val, pred, proba)
    m["Fold"] = f"Fold {fold_i}"
    fold_rows.append(m)

# Table: rows=Metric, cols=Fold1..Fold5
fold_df = pd.DataFrame(fold_rows).set_index("Fold")[metrics_order].T
fold_df["Mean"] = fold_df.mean(axis=1)
fold_df["Std Dev"] = fold_df.std(axis=1)

# Adds tier column 
fold_df.insert(0, "Tier", tier_name)

# Metric as a column
fold_df = fold_df.reset_index().rename(columns={"index": "Metric"})

# Save
output_csv = "../results/lightgbm_5fold_threshold0.50_by_tier.csv"

if os.path.exists(output_csv):
    fold_df.to_csv(output_csv, mode="a", header=False, index=False)
else:
    fold_df.to_csv(output_csv, index=False)

print(f"Saved fold table to: {output_csv}")
fold_df

Saved fold table to: ../results/lightgbm_5fold_threshold0.50_by_tier.csv

6. Evaluate Recall-optimized and AUC-optimized models on the test set (threshold = 0.5)

After performing two separate hyperparameter searches one maximizing Recall and one maximizing AUC we now evaluate both tuned models on the held-out test set.

The goal is to compare how each model behaves using the default decision threshold of 0.5, before we perform threshold tuning.

# Recall-optimized model
proba_recall = best_lgb_recall.predict_proba(X_test)[:, 1]
pred_recall  = (proba_recall >= 0.5).astype(int)
metrics_recall = evaluate_model(y_test, pred_recall, proba_recall)

print("Recall-optimized (threshold 0.5):")
print(metrics_recall)

# # AUC-optimized model
# proba_auc = best_lgb_auc.predict_proba(X_test)[:, 1]
# pred_auc  = (proba_auc >= 0.5).astype(int)
# metrics_auc = evaluate_model(y_test, pred_auc, proba_auc)

# print("\nAUC-optimized (threshold 0.5):")
# print(metrics_auc)

Recall-optimized (threshold 0.5):
{'Accuracy': 0.8542000457770657, 'Sensitivity': 0.8648648648648649, 'Specificity': 0.8528233621914976, 'PPV': 0.4313529705441837, 'NPV': 0.9799554565701559, 'AUC': 0.9327786243408417, 'F1': 0.575616255829447, 'F2': 0.720120020003334}

6.1. Saving metrics results of TEST DATA to /results/lightgbm_testset_metrics_by_tier.csv

tier_name = "Tier_3_Personalized"   # change for desired tier Tier_1_Basic / Tier_2_Clinical / Tier_3_Personalized

test_df = pd.DataFrame({
    "Metric": list(metrics_recall.keys()),
    "Tier":   [tier_name] * len(metrics_recall),
    "Value":  list(metrics_recall.values())
})


test_df["Value"] = test_df["Value"].astype(float).round(6)

output_csv = "../results/lightgbm_testset_metrics_by_tier.csv"

if os.path.exists(output_csv):
    test_df.to_csv(output_csv, mode="a", header=False, index=False)
else:
    test_df.to_csv(output_csv, index=False)

print(f"Saved TEST metrics for {tier_name} to {output_csv}")
test_df

Saved TEST metrics for Tier_3_Personalized to ../results/lightgbm_testset_metrics_by_tier.csv

7. Threshold tuning to optimize Sensitivity and NPV for screening - on selected model (Recall or AUC optimized)

# Choose which tuned model to threshold-tune
final_model = best_lgb_recall      # best_lgb_recall or best_lgb_auc - Only using best_lgb_recall

# Get probabilities on test set
y_proba_test = final_model.predict_proba(X_test)[:, 1]

thresholds = np.arange(0.1, 0.91, 0.05)

rows = []
for t in thresholds:
    preds_t = (y_proba_test >= t).astype(int)
    m = evaluate_model(y_test, preds_t, y_proba_test)
    m["threshold"] = t
    rows.append(m)

metrics_df = pd.DataFrame(rows)
metrics_df_sorted = metrics_df.sort_values("Sensitivity", ascending=False)

metrics_df_sorted.head(10)

8. Picking Best Threshold, prioritizing Recall, NPV - if both higher then prioritizing Specificity, F2

# BEST threshold row: prioritize Recall, NPV
screening_candidates = metrics_df[
    (metrics_df["Sensitivity"] >= 0.90) &
    (metrics_df["NPV"] >= 0.95)
]

if screening_candidates.empty:
    best_row = metrics_df.sort_values(
        ["Sensitivity", "NPV", "Specificity"],
        ascending=[False, False, False]
        ).iloc[0]
else:
    best_row = screening_candidates.sort_values(
        by=["Specificity", "F2"],
        ascending=[False, False]
        ).iloc[0]


best_threshold = float(best_row["threshold"])
print("Best threshold selected:", best_threshold)
print("Metrics at best threshold:")
print(best_row.to_dict())

Best threshold selected: 0.3500000000000001
Metrics at best threshold:
{'Accuracy': 0.8148317692835889, 'Sensitivity': 0.9129129129129129, 'Specificity': 0.8021708231037602, 'PPV': 0.3733115022513303, 'NPV': 0.9861795075456712, 'AUC': 0.9327786243408417, 'F1': 0.5299244625217896, 'F2': 0.7081845006988662, 'threshold': 0.3500000000000001}

# Row 1: threshold = 0.50 
row_05 = metrics_recall.copy()
row_05["threshold"] = 0.50
row_05["tier"] = tier_name
row_05["tuned_for"] = "recall"
row_05["n_features"] = len(feature_list)

# Row 2: best threshold from sweep
row_best = best_row.to_dict()
row_best["tier"] = tier_name
row_best["tuned_for"] = "recall"
row_best["n_features"] = len(feature_list)

two_rows_df = pd.DataFrame([row_05, row_best])

two_rows_df

9. Saving Results csv to results folder - two rows for each tier

output_csv = "../results/lightgbm_two_rows_per_tier.csv"

# Nice column order (optional)
front_cols = ["tier", "tuned_for", "n_features", "threshold"]
other_cols = [c for c in two_rows_df.columns if c not in front_cols]
two_rows_df = two_rows_df[front_cols + other_cols]

if os.path.exists(output_csv):
    two_rows_df.to_csv(output_csv, mode="a", header=False, index=False)
else:
    two_rows_df.to_csv(output_csv, index=False)

print(f"Appended 2 rows to: {output_csv}")

Appended 2 rows to: ../results/lightgbm_two_rows_per_tier.csv

10. (NOT USING) Saving results to lightgbm_all_tiers_results.csv

Note: For all three tiers, before running please change the tier_name in 3rd code block to desired tier. Also change final_model in the above code block to desired tuned mode, whether recall or auc Saving for both tuned model for recall and AUC

# # Explicitly record which model was used
# metrics_df_sorted["tier"] = tier_name
# metrics_df_sorted["model"] = "LightGBM"
# metrics_df_sorted["tuned_for"] = "recall"   # <-- IMPORTANT to change recall or auc tuned model results
# metrics_df_sorted["n_features"] = len(feature_list)

# cols = ["tier", "model", "tuned_for", "n_features", "threshold"] + [
#     c for c in metrics_df_sorted.columns
#     if c not in ["tier", "model", "tuned_for", "n_features", "threshold"]
# ]
# metrics_df_sorted = metrics_df_sorted[cols]

# # Output file
# output_csv = "../results/lightgbm_all_tiers_results.csv"

# # Append if exists, otherwise create
# if os.path.exists(output_csv):
#     metrics_df_sorted.to_csv(output_csv, mode="a", header=False, index=False)
# else:
#     metrics_df_sorted.to_csv(output_csv, index=False)

# # blank row after each tuned model, readability
# blank_row = pd.DataFrame(
#     [{col: "" for col in metrics_df_sorted.columns}]
# )

# blank_row.to_csv(output_csv, mode="a", header=False, index=False)

# print(f"Results appended to {output_csv}")