import joblib
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import recall_score, fbeta_score, roc_auc_score

MODEL_PATH = "best_model_v2_calibrated.joblib"
CSV_PATH = "nhanes_sleep_extended_v2.csv"

TARGET = "Sleep_Risk"

FEATURES = [
    "Age",
    "Gender",
    "BMI",
    "BP_SYS",
    "BP_DIA",
    "Phys_Activity_Days",
    "DPQ_Score",
    "Smoking_Indicator",
    "Alcohol_Feature",
    "Diabetes_Indicator",
    "Cycle",
]

def main():
    df = pd.read_csv(CSV_PATH)

    # drop rows with missing values in required columns
    df = df.dropna(subset=FEATURES + [TARGET]).copy()

    X = df[FEATURES]
    y = df[TARGET].astype(int)

    # Train/test split (keep same logic as training if possible)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.20, random_state=42, stratify=y
    )

    bundle = joblib.load(MODEL_PATH)
    model = bundle.get("model") if isinstance(bundle, dict) and "model" in bundle else bundle
    calibrator = bundle.get("calibrator") if isinstance(bundle, dict) and "calibrator" in bundle else None

    # Get probabilities
    p_test_raw = model.predict_proba(X_test)[:, 1]
    if calibrator is not None:
        p_test = calibrator.predict(p_test_raw.reshape(-1, 1))
    else:
        p_test = p_test_raw

    # AUC (threshold-free)
    auc = roc_auc_score(y_test, p_test)
    print(f"TEST AUC: {auc:.4f}")

    # Sweep thresholds
    thresholds = np.round(np.arange(0.05, 0.90, 0.01), 2)

    best_f2 = (-1, None)
    th_recall_085 = None

    for th in thresholds:
        pred = (p_test >= th).astype(int)
        rec = recall_score(y_test, pred)
        f2 = fbeta_score(y_test, pred, beta=2)

        if f2 > best_f2[0]:
            best_f2 = (f2, th)

        if th_recall_085 is None and rec >= 0.85:
            th_recall_085 = th

    print(f"Best threshold by F2: {best_f2[1]} (F2={best_f2[0]:.4f})")
    if th_recall_085 is not None:
        print(f"Threshold achieving recall ≥ 0.85 : {th_recall_085}")
    else:
        print("No threshold achieved recall ≥ 0.85 in tested range.")

if __name__ == "__main__":
    main()