import os, time, warnings
import numpy as np
warnings.filterwarnings("ignore")

from preprocessing import load_diabetes, load_heart, load_sepsis
from models import train_all_models
from evaluation import (
    evaluate_all_models,
    plot_roc, plot_metrics_bar, plot_confusion_matrices,
    plot_feature_importance,
    plot_cross_disease_heatmap, plot_cross_disease_bar,
)

os.makedirs("results/figures", exist_ok=True)
os.makedirs("results/tables",  exist_ok=True)


def banner(title):
    print("\n" + "═" * 65)
    print(f"  {title}")
    print("═" * 65)


def run_pipeline(disease, loader, path):
    banner(f"DISEASE: {disease.upper()}")

    print("\n[1/4] Loading & preprocessing…")
    t0 = time.time()
    X_train, X_test, y_train, y_test, feats = loader(path)
    print(f"  Train: {X_train.shape}  |  Test: {X_test.shape}")
    print(f"  Class balance (train) → 0: {int((y_train==0).sum())}  "
          f"1: {int((y_train==1).sum())}")
    print(f"  Features: {feats}")
    print(f"  Preprocessing time: {time.time()-t0:.1f}s")

    print("\n[2/4] Training models (5-fold CV grid search)…")
    t0 = time.time()
    fitted, cv_scores = train_all_models(X_train, y_train)
    print(f"\n  Training complete in {time.time()-t0:.1f}s")
    print("  Cross-validation AUC scores:")
    for name, score in sorted(cv_scores.items(), key=lambda x: -x[1]):
        print(f"    {name:22s}: {score:.4f}")

    print("\n[3/4] Evaluating on held-out test set…")
    results_df = evaluate_all_models(fitted, X_test, y_test, disease)

    print("\n[4/4] Generating plots…")
    plot_roc(fitted, X_test, y_test, disease)
    plot_metrics_bar(results_df, disease)
    plot_confusion_matrices(fitted, X_test, y_test, disease)

    # Feature importance from best tree-based model
    best = results_df["AUC-ROC"].idxmax()
    print(f"  Best model: {best} (AUC = {results_df['AUC-ROC'].max():.4f})")
    for pref in ["Gradient Boosting", "Random Forest", best]:
        if pref in fitted and hasattr(fitted[pref], "feature_importances_"):
            plot_feature_importance(fitted[pref], feats, disease, pref)
            break

    return results_df



    all_results = {}

    for disease, loader, path in [
        ("Diabetes",      load_diabetes, "data/diabetes.csv"),
        ("Heart Disease", load_heart,    "data/heart_disease.csv"),
        ("Sepsis",        load_sepsis,   "data/sepsis.csv"),
    ]:
        if not os.path.exists(path):
            print(f"\n[ERROR] {path} not found.")
            continue
        all_results[disease] = run_pipeline(disease, loader, path)

    if len(all_results) > 1:
        banner("CROSS-DISEASE COMPARISON")
        summary = plot_cross_disease_heatmap(all_results)
        plot_cross_disease_bar(all_results)
        print("\n  Final AUC-ROC Summary:")
        print(summary.to_string())
        mean_auc = summary.mean(axis=1)
        best_model = mean_auc.idxmax()
        print(f"\n  Best overall model: {best_model} "
              f"(mean AUC = {mean_auc[best_model]:.4f})")

    banner("PIPELINE COMPLETE")
    print("\nOutputs saved to:")
    print("  ./results/figures/  — ROC curves, bar charts, confusion matrices, feature importance")
    print("  ./results/tables/   — per-disease and cross-disease CSVs")