src/tune_hyperparams.py

"""
Optuna hyperparameter tuning for LightGBM (GPU) on the duty-aware features.

Runs N trials of Bayesian optimization, each trial trains on 2018-2023
and evaluates on 2024 holdout. Saves best params + retrained final model.

Usage:
  python tune_hyperparams.py              # 50 trials, lgbm
  python tune_hyperparams.py --trials 100
  python tune_hyperparams.py --model xgb  # tune XGBoost instead
"""

import os, sys, argparse, json
import numpy as np
import pandas as pd
import optuna
import lightgbm as lgb
import xgboost as xgb
from sklearn.metrics import roc_auc_score, average_precision_score
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

optuna.logging.set_verbosity(optuna.logging.WARNING)

PROC_DIR = os.path.join(os.path.dirname(__file__), "..", "data", "processed")
PLOTS    = os.path.join(PROC_DIR, "plots")
os.makedirs(PLOTS, exist_ok=True)

BASE_FEATURES = [
    "A_weather_delay_rate", "A_weather_cancel_rate", "A_avg_weather_delay_min",
    "A_p75_weather_delay_min", "A_p95_weather_delay_min", "A_nas_delay_rate",
    "A_overall_weather_delay_rate", "A_overall_avg_weather_delay_min",
    "B_weather_delay_rate", "B_weather_cancel_rate", "B_avg_weather_delay_min",
    "B_p75_weather_delay_min", "B_p95_weather_delay_min", "B_nas_delay_rate",
    "B_overall_weather_delay_rate", "B_overall_avg_weather_delay_min",
    "pair_combined_weather_rate", "pair_max_weather_rate", "pair_min_weather_rate",
    "pair_weather_rate_sum", "pair_avg_weather_delay_min", "both_high_risk",
    "Month", "is_spring_summer", "median_turnaround_min",
]
DUTY_FEATURES = [
    "A_dep_hour_median", "A_late_dep_rate", "A_early_dep_rate",
    "A_avg_block_min", "A_late_aircraft_delay_rate", "A_avg_late_aircraft_min",
    "B_dep_hour_median", "B_late_dep_rate", "B_early_dep_rate",
    "B_avg_block_min", "B_late_aircraft_delay_rate", "B_avg_late_aircraft_min",
    "tight_connection_rate", "very_tight_rate",
    "cascade_risk", "total_duty_block_min", "duty_overrun_risk", "late_dep_sequence",
]


def load_data(use_duty: bool = True):
    path = os.path.join(PROC_DIR,
                        "sequence_features_duty.parquet" if use_duty
                        else "sequence_features.parquet")
    df = pd.read_parquet(path)
    season_cols = [c for c in df.columns if c.startswith("season_")]
    all_feats   = BASE_FEATURES + (DUTY_FEATURES if use_duty else []) + season_cols
    feat_cols   = [c for c in all_feats if c in df.columns]
    return df, feat_cols


def lgbm_objective(trial, X_train, y_train, X_val, y_val):
    params = {
        "n_estimators":     trial.suggest_int("n_estimators", 300, 2000),
        "learning_rate":    trial.suggest_float("learning_rate", 0.01, 0.15, log=True),
        "max_depth":        trial.suggest_int("max_depth", 4, 10),
        "num_leaves":       trial.suggest_int("num_leaves", 31, 255),
        "subsample":        trial.suggest_float("subsample", 0.5, 1.0),
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.5, 1.0),
        "min_child_samples":trial.suggest_int("min_child_samples", 10, 100),
        "reg_alpha":        trial.suggest_float("reg_alpha", 1e-4, 10.0, log=True),
        "reg_lambda":       trial.suggest_float("reg_lambda", 1e-4, 10.0, log=True),
        "class_weight": "balanced",
        "device": "cuda",
        "random_state": 42,
        "n_jobs": -1,
        "verbose": -1,
    }
    m = lgb.LGBMClassifier(**params)
    m.fit(X_train, y_train,
          eval_set=[(X_val, y_val)],
          callbacks=[lgb.early_stopping(40, verbose=False), lgb.log_evaluation(-1)])
    return roc_auc_score(y_val, m.predict_proba(X_val)[:, 1])


def xgb_objective(trial, X_train, y_train, X_val, y_val):
    neg, pos = (y_train == 0).sum(), (y_train == 1).sum()
    params = {
        "n_estimators":     trial.suggest_int("n_estimators", 300, 2000),
        "learning_rate":    trial.suggest_float("learning_rate", 0.01, 0.15, log=True),
        "max_depth":        trial.suggest_int("max_depth", 4, 10),
        "subsample":        trial.suggest_float("subsample", 0.5, 1.0),
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.5, 1.0),
        "min_child_weight": trial.suggest_int("min_child_weight", 1, 20),
        "gamma":            trial.suggest_float("gamma", 1e-4, 5.0, log=True),
        "reg_alpha":        trial.suggest_float("reg_alpha", 1e-4, 10.0, log=True),
        "reg_lambda":       trial.suggest_float("reg_lambda", 1e-4, 10.0, log=True),
        "scale_pos_weight": neg / pos,
        "tree_method": "hist",
        "device": "cuda",
        "eval_metric": "aucpr",
        "early_stopping_rounds": 40,
        "random_state": 42,
        "n_jobs": -1,
    }
    m = xgb.XGBClassifier(**params)
    m.fit(X_train, y_train, eval_set=[(X_val, y_val)], verbose=False)
    return roc_auc_score(y_val, m.predict_proba(X_val)[:, 1])


def run_tuning(model_type: str, n_trials: int, use_duty: bool):
    print(f"\nLoading {'duty-enriched' if use_duty else 'base'} features...")
    df, feat_cols = load_data(use_duty)
    print(f"Dataset: {df.shape}  |  Features: {len(feat_cols)}")

    train_mask = df["Year"] < 2024
    val_mask   = df["Year"] == 2024
    X_train = df[train_mask][feat_cols].astype(float)
    y_train = df[train_mask]["target"].astype(int)
    X_val   = df[val_mask][feat_cols].astype(float)
    y_val   = df[val_mask]["target"].astype(int)

    print(f"Train: {len(X_train):,}  Val: {len(X_val):,}")
    print(f"Model: {model_type.upper()}  |  Trials: {n_trials}\n")

    objective = (lgbm_objective if model_type == "lgbm" else xgb_objective)

    study = optuna.create_study(direction="maximize",
                                sampler=optuna.samplers.TPESampler(seed=42))
    study.optimize(
        lambda trial: objective(trial, X_train, y_train, X_val, y_val),
        n_trials=n_trials,
        show_progress_bar=True,
    )

    print(f"\nBest AUC: {study.best_value:.4f}")
    print(f"Best params: {json.dumps(study.best_params, indent=2)}")

    # Save best params
    params_path = os.path.join(PROC_DIR, f"best_params_{model_type}.json")
    with open(params_path, "w") as f:
        json.dump({"best_auc": study.best_value, "params": study.best_params}, f, indent=2)
    print(f"Params saved → {params_path}")

    # Retrain final model with best params on full train set
    print("\nRetraining final model with best params...")
    best = study.best_params

    if model_type == "lgbm":
        final = lgb.LGBMClassifier(
            **best, class_weight="balanced",
            device="cuda", random_state=42, n_jobs=-1, verbose=-1,
        )
        final.fit(X_train, y_train,
                  eval_set=[(X_val, y_val)],
                  callbacks=[lgb.early_stopping(50, verbose=False), lgb.log_evaluation(-1)])
        model_path = os.path.join(PROC_DIR, f"lgbm_tuned_model.txt")
        final.booster_.save_model(model_path)
    else:
        neg, pos = (y_train == 0).sum(), (y_train == 1).sum()
        final = xgb.XGBClassifier(
            **best, scale_pos_weight=neg/pos,
            tree_method="hist", device="cuda",
            eval_metric="aucpr", early_stopping_rounds=50,
            random_state=42, n_jobs=-1,
        )
        final.fit(X_train, y_train, eval_set=[(X_val, y_val)], verbose=False)
        model_path = os.path.join(PROC_DIR, f"xgb_tuned_model.json")
        final.save_model(model_path)

    proba = final.predict_proba(X_val)[:, 1]
    final_auc = roc_auc_score(y_val, proba)
    final_ap  = average_precision_score(y_val, proba)
    print(f"\nFinal tuned model — ROC-AUC: {final_auc:.4f}  AP: {final_ap:.4f}")
    print(f"Model saved → {model_path}")

    # Plot: optimization history + param importances
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # Trial history
    ax = axes[0]
    values = [t.value for t in study.trials if t.value is not None]
    best_so_far = [max(values[:i+1]) for i in range(len(values))]
    ax.plot(values, alpha=0.4, color="steelblue", label="Trial AUC")
    ax.plot(best_so_far, color="red", linewidth=2, label="Best so far")
    ax.axhline(final_auc, color="green", linestyle="--", label=f"Final tuned={final_auc:.4f}")
    ax.set_xlabel("Trial")
    ax.set_ylabel("ROC-AUC (2024 holdout)")
    ax.set_title(f"Optuna optimization history\n{model_type.upper()} — {n_trials} trials")
    ax.legend()
    ax.grid(alpha=0.3)

    # Param importance
    ax = axes[1]
    try:
        importances = optuna.importance.get_param_importances(study)
        params_sorted = list(importances.keys())[:10]
        vals_sorted   = [importances[p] for p in params_sorted]
        ax.barh(params_sorted[::-1], vals_sorted[::-1], color="steelblue")
        ax.set_xlabel("Importance")
        ax.set_title("Hyperparameter importance\n(Fanova method)")
        ax.grid(axis="x", alpha=0.3)
    except Exception:
        ax.text(0.5, 0.5, "Param importance\nnot available", ha="center", va="center",
                transform=ax.transAxes)

    plt.tight_layout()
    plot_path = os.path.join(PLOTS, f"optuna_{model_type}_tuning.png")
    plt.savefig(plot_path, dpi=150)
    plt.close()
    print(f"Plot saved → {plot_path}")

    return study, final_auc


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--trials", type=int, default=50)
    parser.add_argument("--model",  choices=["lgbm", "xgb", "both"], default="lgbm")
    parser.add_argument("--no-duty", action="store_true",
                        help="Use base features only (no duty enrichment)")
    args = parser.parse_args()

    use_duty = not args.no_duty
    models   = ["lgbm", "xgb"] if args.model == "both" else [args.model]

    results = {}
    for m in models:
        _, auc = run_tuning(m, args.trials, use_duty)
        results[m] = auc

    print("\n" + "="*50)
    print("TUNING SUMMARY")
    print("="*50)
    for m, auc in results.items():
        print(f"  {m.upper()}: {auc:.4f}")


if __name__ == "__main__":
    main()