"""
train.py  RouteRisk ML Pipeline
=================================
Responsibilities:
  1. Load flight data from database
  2. Pre-compute route distance lookup table
  3. Feature engineering
  4. Train/test split + imputation + SMOTE
  5. Train and evaluate multiple models
  6. Auto-select best model (F1 -> ROC-AUC)
  7. Save all artifacts to /models and /data

Run:
    python train.py

Outputs:
    data/route_distances.csv
    models/best_model.pkl
    models/num_imputer.pkl
    models/cat_imputer.pkl
    models/num_cols.pkl
    models/feature_columns.pkl
    models/label_encoders.pkl
    models/model_metadata.json
"""

import sys
import os
import json
import warnings
warnings.filterwarnings('ignore')

import pandas as pd
import numpy as np
import joblib

from pandas.tseries.holiday import USFederalHolidayCalendar
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import (
    classification_report, roc_auc_score,
    f1_score, precision_score, recall_score, accuracy_score,
)
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.dummy import DummyClassifier
from sklearn.impute import SimpleImputer
from imblearn.over_sampling import SMOTE
import xgboost as xgb

#  Database connection 
path_to_add = os.path.abspath('../..')
sys.path.insert(0, path_to_add)
from core.database import engine

#  Holiday calendar (wide coverage) 
_cal      = USFederalHolidayCalendar()
_holidays = _cal.holidays(start='2020-01-01', end='2030-12-31')

#  Global encoder store (populated inside build_features) 
label_encoders = {}

if __name__ == '__main__':

    # 
    # STEP 1 — LOAD DATA
    # 
    print("=" * 60)
    print("STEP 1 — Loading data from database")
    print("=" * 60)

    # Set SAMPLE_SIZE to None for a full production training run
    SAMPLE_SIZE = 500_000

    if SAMPLE_SIZE:
        # Sampling pushed into SQL — never loads the full table into memory
        query = f"""
            SELECT *
            FROM flights
            WHERE origin IS NOT NULL
                AND dest IS NOT NULL
                AND carrier IS NOT NULL
            ORDER BY RANDOM()
            LIMIT {SAMPLE_SIZE}
        """
    else:
        query = """
            SELECT *
            FROM flights
            WHERE origin IS NOT NULL
                AND dest IS NOT NULL
                AND carrier IS NOT NULL
        """

    df = pd.read_sql(query, engine)
    print(f"  Loaded {len(df):,} rows, {df.shape[1]} columns")
    print(f"  Class distribution:\n{df.groupby('arr_del15').size().to_string()}")


    # 
    # STEP 2 — PRE-COMPUTE ROUTE DISTANCE LOOKUP
    # 
    print("\n" + "=" * 60)
    print("STEP 2 — Pre-computing route distance lookup")
    print("=" * 60)

    os.makedirs("data", exist_ok=True)

    if os.path.exists("data/route_distances.csv"):
        print("   data/route_distances.csv already exists — skipping")
    else:
        print("  CSV not found — loading full table to compute route distances ...")
        df_all = pd.read_sql(
            """
            SELECT origin, dest, distance
            FROM flights
            WHERE origin IS NOT NULL
                AND dest IS NOT NULL
                AND distance IS NOT NULL
            """,
            engine,
        )
        route_distances = (
            df_all
            .groupby(['origin', 'dest'])['distance']
            .max()
            .reset_index()
            .rename(columns={'distance': 'distance_miles'})
        )
        route_distances.to_csv("data/route_distances.csv", index=False)
        print(f"  {len(route_distances):,} unique routes saved -> data/route_distances.csv")
        del df_all   # free memory immediately


    # 
    # STEP 3 — FEATURE ENGINEERING
    # 
    print("\n" + "=" * 60)
    print("STEP 3 — Feature engineering")
    print("=" * 60)

    def build_features(df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()

        # Date-derived features
        df['fl_date']      = pd.to_datetime(df['fl_date'])
        df['day_of_week']  = df['fl_date'].dt.dayofweek   # 0=Mon, 6=Sun
        df['day_of_month'] = df['fl_date'].dt.day

        # Holiday proximity (requires actual date — not just month/year)
        df['days_to_holiday'] = df['fl_date'].apply(
            lambda d: int(min(abs((_holidays - d).days)))
        )
        df['is_holiday_week'] = (df['days_to_holiday'] <= 3).astype(int)

        # Peak travel season: June/July/August/December
        df['is_peak_season'] = df['month'].isin([6, 7, 8, 12]).astype(int)

        # Departure time  named bucket (avoids treating 800 as numeric)
        df['dep_hour']   = df['crs_dep_time'] // 100
        df['dep_period'] = pd.cut(
            df['dep_hour'],
            bins=[0, 6, 12, 17, 21, 24],
            labels=['red_eye', 'morning', 'afternoon', 'evening', 'night'],
            right=False,
        )

        # Route identifier
        df['route'] = df['origin'] + '_' + df['dest']

        # Weather flag: 1 if any weather delay was recorded
        df['had_weather_delay'] = (
            pd.to_numeric(df['weather_delay'], errors='coerce').fillna(0) > 0
        ).astype(int)

        # Drop leaky / post-flight / raw columns
        drop_cols = [
            'id', 'fl_date', 'flight_num', 'tail_number',
            'arr_delay',                                      # direct leakage of target
            'carrier_delay', 'weather_delay', 'nas_delay',   # replaced by engineered flags
            'security_delay', 'late_aircraft_delay',
            'cancelled', 'diverted',
            'crs_dep_time', 'dep_hour',
        ]
        df.drop(columns=[c for c in drop_cols if c in df.columns], inplace=True)

        # Label-encode categoricals (store each encoder for inference reuse)
        for col in ['carrier', 'origin', 'dest', 'route', 'dep_period']:
            le = LabelEncoder()
            df[col] = le.fit_transform(df[col].astype(str))
            label_encoders[col] = le

        return df


    df_model = build_features(df)
    X = df_model.drop(columns=['arr_del15'])
    y = df_model['arr_del15']

    print(f"  Features ({len(X.columns)}): {X.columns.tolist()}")
    print(f"  Class distribution:\n{y.value_counts().to_string()}")


    # 
    # STEP 4 — TRAIN / TEST SPLIT + IMPUTATION
    # 
    print("\n" + "=" * 60)
    print("STEP 4 — Train/test split and imputation")
    print("=" * 60)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    print(f"  Train: {len(X_train):,} rows | Test: {len(X_test):,} rows")

    # Compute route weather delay rate on train split only.
    # Decode route integers -> strings (e.g. 'ORD_LAX') so route_weather_rates.csv
    # has string keys that predict_delay can look up by name, not integer codes.
    route_le = label_encoders['route']
    route_strings_train = pd.Series(
        route_le.inverse_transform(X_train['route'].astype(int)),
        index=X_train.index,
        name='route_str',
    )
    route_weather_rate = (
        pd.concat([route_strings_train, X_train['had_weather_delay']], axis=1)
        .groupby('route_str')['had_weather_delay']
        .mean()
        .rename('route_weather_delay_rate')
    )
    X_train['route_weather_delay_rate'] = route_strings_train.map(route_weather_rate)

    route_strings_test = pd.Series(
        route_le.inverse_transform(X_test['route'].astype(int)),
        index=X_test.index,
    )
    X_test['route_weather_delay_rate'] = route_strings_test.map(route_weather_rate).fillna(0)

    route_weather_rate.reset_index().to_csv("data/route_weather_rates.csv", index=False)
    print(f"  Route weather rates computed on train: {len(route_weather_rate)} routes")

    # Route + month weather delay rate (for frontend weather pre-selection)
    route_month_weather = (
        pd.concat([route_strings_train, X_train[['month', 'had_weather_delay']]], axis=1)
        .groupby(['route_str', 'month'])['had_weather_delay']
        .mean()
        .rename('weather_delay_rate')
        .reset_index()
    )
    route_month_weather.to_csv("data/route_month_weather.csv", index=False)
    print(f"  Route-month weather rates: {len(route_month_weather)} rows")

    # Identify numeric vs categorical columns after encoding
    num_cols = X_train.select_dtypes(include=[np.number]).columns.tolist()
    cat_cols = X_train.select_dtypes(exclude=[np.number]).columns.tolist()

    # Fit imputers on train only; transform both splits
    num_imputer = SimpleImputer(strategy='median')
    X_train[num_cols] = num_imputer.fit_transform(X_train[num_cols])
    X_test[num_cols]  = num_imputer.transform(X_test[num_cols])

    cat_imputer = None
    if cat_cols:
        cat_imputer = SimpleImputer(strategy='most_frequent')
        X_train[cat_cols] = cat_imputer.fit_transform(X_train[cat_cols])
        X_test[cat_cols]  = cat_imputer.transform(X_test[cat_cols])

    print(f"  NaNs in X_train after imputation: {X_train.isna().sum().sum()}")
    print(f"  NaNs in X_test  after imputation: {X_test.isna().sum().sum()}")


    # 
    # STEP 5 — SMOTE (handle class imbalance)
    # 
    print("\n" + "=" * 60)
    print("STEP 5 — SMOTE oversampling")
    print("=" * 60)

    if y_train.nunique() > 1:
        smote = SMOTE(random_state=42)
        X_train_bal, y_train_bal = smote.fit_resample(X_train, y_train)
        print(f"  After SMOTE: {pd.Series(y_train_bal).value_counts().to_dict()}")
    else:
        X_train_bal, y_train_bal = X_train, y_train
        print("  SMOTE skipped: only one class present")


    # 
    # STEP 6 — TRAIN MODELS
    # 
    print("\n" + "=" * 60)
    print("STEP 6 — Training models")
    print("=" * 60)

    pos_weight = (
        (y_train == 0).sum() / (y_train == 1).sum()
        if y_train.nunique() > 1 else 1
    )

    models = {
        "Baseline (majority)": DummyClassifier(strategy="most_frequent"),
        "Logistic Regression": LogisticRegression(max_iter=1000, class_weight='balanced'),
        "Random Forest":       RandomForestClassifier(
                                n_estimators=200, class_weight='balanced',
                                random_state=42, n_jobs=-1),
        "Gradient Boosting":   GradientBoostingClassifier(
                                n_estimators=200, learning_rate=0.05,
                                max_depth=4, random_state=42),
        "XGBoost":             xgb.XGBClassifier(
                                n_estimators=200, learning_rate=0.05,
                                max_depth=4, use_label_encoder=False,
                                eval_metric='logloss', random_state=42,
                                scale_pos_weight=pos_weight),
    }

    results = {}
    for name, model in models.items():
        print(f"\n  Training: {name} ...")
        model.fit(X_train_bal, y_train_bal)
        y_pred  = model.predict(X_test)
        y_proba = model.predict_proba(X_test)[:, 1] if hasattr(model, 'predict_proba') else None
        auc     = roc_auc_score(y_test, y_proba) if (y_proba is not None and y_test.nunique() > 1) else None
        results[name] = {'model': model, 'y_pred': y_pred, 'auc': auc}
        print(classification_report(y_test, y_pred, zero_division=0))
        if auc:
            print(f"  ROC-AUC: {auc:.4f}")


    # 
    # STEP 7 — EVALUATION SUMMARY
    # 
    print("\n" + "=" * 60)
    print("STEP 7 — Evaluation summary")
    print("=" * 60)

    summary_rows = []
    for name, res in results.items():
        summary_rows.append({
            'Model':     name,
            'Accuracy':  accuracy_score(y_test, res['y_pred']),
            'Precision': precision_score(y_test, res['y_pred'], zero_division=0),
            'Recall':    recall_score(y_test, res['y_pred'], zero_division=0),
            'F1':        f1_score(y_test, res['y_pred'], zero_division=0),
            'ROC-AUC':   res['auc'],
        })

    summary_df = pd.DataFrame(summary_rows).set_index('Model').round(4)
    print("\nMODEL COMPARISON")
    print(summary_df.to_string())


    # 
    # STEP 8 — AUTO-SELECT BEST MODEL AND SAVE ARTIFACTS
    # 
    print("\n" + "=" * 60)
    print("STEP 8 — Auto-select best model and save artifacts")
    print("=" * 60)

    # Rank by F1 first, ROC-AUC as tiebreaker; exclude dummy baseline
    ranked = (
        summary_df[summary_df.index != "Baseline (majority)"]
        .sort_values(['F1', 'ROC-AUC'], ascending=False)
    )
    best_model_name = ranked.index[0]
    best_model      = results[best_model_name]['model']
    best_metrics    = summary_df.loc[best_model_name]

    print(f"\n   Best model: {best_model_name}")
    print(f"     F1:       {best_metrics['F1']:.4f}")
    print(f"     ROC-AUC:  {best_metrics['ROC-AUC']:.4f}")
    print(f"     Accuracy: {best_metrics['Accuracy']:.4f}")

    os.makedirs("models", exist_ok=True)

    joblib.dump(best_model,         "models/best_model.pkl")
    joblib.dump(num_imputer,        "models/num_imputer.pkl")
    joblib.dump(X_train.columns.tolist(), "models/feature_columns.pkl")
    joblib.dump(label_encoders,     "models/label_encoders.pkl")
    joblib.dump(num_cols,           "models/num_cols.pkl")

    if cat_imputer is not None:
        joblib.dump(cat_imputer, "models/cat_imputer.pkl")

    metadata = {
        "best_model_name": best_model_name,
        "metrics":         best_metrics.to_dict(),
        "trained_on_rows": len(X),
        "features":        X.columns.tolist(),
        "trained_at":      pd.Timestamp.now().isoformat(),
    }
    with open("models/model_metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    print("\n   Artifacts saved:")
    for fname in sorted(os.listdir("models")):
        size = os.path.getsize(f"models/{fname}")
        print(f"      models/{fname}  ({size/1024:.1f} KB)")

    print("\n   Metadata written to models/model_metadata.json")
    print(json.dumps(metadata, indent=2))
    print("\nTraining complete.")