"""
predict.py — RouteRisk Inference Interface
==========================================
This module is the single entry point the frontend calls to get delay predictions.
It loads all trained artifacts from /models (produced by train.py) and exposes:

    predict_delay(flight_input: dict) -> dict

Prerequisites (must exist before calling predict_delay):
    data/route_distances.csv      ← produced by train.py
    models/best_model.pkl
    models/num_imputer.pkl
    models/feature_columns.pkl
    models/label_encoders.pkl
    models/num_cols.pkl
    models/model_metadata.json

Usage from frontend / Streamlit:
    from predict import predict_delay, WEATHER_CONDITIONS
    result = predict_delay({...})
"""

import os
import json
import joblib
import pandas as pd
import numpy as np
from pandas.tseries.holiday import USFederalHolidayCalendar

from sklearn.exceptions import InconsistentVersionWarning
import warnings
warnings.filterwarnings("ignore", category=InconsistentVersionWarning)

# ── Holiday calendar — loaded once at module import, not per call ─────────────
_cal      = USFederalHolidayCalendar()
_holidays = _cal.holidays(start='2020-01-01', end='2030-12-31')

# ── Route distance lookup — loaded once at module import ─────────────────────
_route_dist_df     = pd.read_csv("data/route_distances.csv")
_route_dist_lookup = _route_dist_df.set_index(['origin', 'dest'])['distance_miles'].to_dict()

# ── Route weather delay rate lookup — loaded once at module import ────────────
_route_weather_df      = pd.read_csv("data/route_weather_rates.csv")
_route_weather_lookup  = _route_weather_df.set_index('route_str')['route_weather_delay_rate'].to_dict()

# ── Weather condition map — frontend passes a label, model gets 0/1 ──────────
WEATHER_CONDITIONS = {
    "Clear":        0,
    "Partly Cloudy":0,
    "Rain":         1,
    "Thunderstorm": 1,
    "Snow / Ice":   1,
    "Fog / Low Vis":1,
}


# ─────────────────────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────────────────────

def get_distance(origin: str, dest: str) -> float:
    """
    Look up pre-computed route distance from the CSV.
    Fallback order: exact match → reversed direction → dataset median.
    """
    key         = (origin.upper(), dest.upper())
    reverse_key = (dest.upper(), origin.upper())

    if key in _route_dist_lookup:
        return float(_route_dist_lookup[key])
    elif reverse_key in _route_dist_lookup:
        return float(_route_dist_lookup[reverse_key])
    else:
        fallback = float(_route_dist_df['distance_miles'].median())
        print(f"⚠️  Route {origin}→{dest} not in lookup — using median {fallback:.0f} mi")
        return fallback


# ─────────────────────────────────────────────────────────────────────────────
# MAIN INTERFACE
# ─────────────────────────────────────────────────────────────────────────────

def predict_delay(flight_input: dict) -> dict:
    """
    Predict whether a flight will be delayed (arrival ≥ 15 min late).

    Parameters
    ----------
    flight_input : dict
        Required keys:
            carrier      (str)  — IATA carrier code, e.g. 'DL', 'UA', 'WN'
            origin       (str)  — Origin airport IATA code, e.g. 'ATL'
            dest         (str)  — Destination airport IATA code, e.g. 'LAX'
            crs_dep_time (int)  — Scheduled departure time as HHMM integer, e.g. 900 = 9:00 AM
            fl_date      (str)  — Flight date in 'YYYY-MM-DD' format, e.g. '2024-12-24'

        Optional keys (defaults applied if omitted):
            dep_delay              (float) — departure delay in minutes at gate;
                                             defaults to 0 (pre-booking: assumes on-time departure).
                                             Pass actual value for real-time at-gate prediction,
                                             e.g. 45.0 if the flight is already 45 min late.
            had_weather_delay      (int, 0 or 1) — 0 = clear, 1 = adverse weather
                                                    Use WEATHER_CONDITIONS dict to convert
                                                    a UI label to this value.
            route_weather_delay_rate (float)     — historical weather delay rate for this route;
                                                   auto-looked up from route_weather_rates.csv
            distance               (float)       — route distance in miles;
                                                   auto-looked up from route_distances.csv if omitted

    Returns
    -------
    dict with keys:
        prediction              (str)   — "DELAYED" or "ON TIME"
        delay_probability       (float) — probability of delay, e.g. 0.7321
        confidence              (str)   — human-readable confidence, e.g. "73.2%"
        distance_miles          (float) — route distance used in prediction
        is_holiday_week         (bool)  — True if flight is within 3 days of a US federal holiday
        days_to_holiday         (int)   — days to nearest US federal holiday
        is_peak_season          (bool)  — True if month is Jun/Jul/Aug/Dec
        weather_flag            (int)   — 0 (clear) or 1 (adverse) as passed in
        model_used              (str)   — name of the model that produced this prediction

    Example
    -------
    >>> result = predict_delay({
    ...     'carrier':      'DL',
    ...     'origin':       'ATL',
    ...     'dest':         'LAX',
    ...     'crs_dep_time': 900,
    ...     'fl_date':      '2024-07-04',
    ... })
    >>> print(result['prediction'], result['delay_probability'])
    """

    # ── Load artifacts ────────────────────────────────────────────────────────
    model     = joblib.load("models/best_model.pkl")
    num_imp   = joblib.load("models/num_imputer.pkl")
    feat_cols = joblib.load("models/feature_columns.pkl")
    les       = joblib.load("models/label_encoders.pkl")
    num_cols_ = joblib.load("models/num_cols.pkl")

    with open("models/model_metadata.json") as f:
        meta = json.load(f)

    # ── Apply defaults for optional fields ───────────────────────────────────
    flight_input = flight_input.copy()                      # don't mutate caller's dict
    flight_input.setdefault('dep_delay', 0.0)               # 0 = pre-booking / on-time assumption
    flight_input.setdefault('had_weather_delay', 0)
    route_key = f"{flight_input['origin'].upper()}_{flight_input['dest'].upper()}"
    flight_input.setdefault(
        'route_weather_delay_rate',
        _route_weather_lookup.get(route_key, 0.0)
    )
    if 'distance' not in flight_input:
        flight_input['distance'] = get_distance(
            flight_input['origin'], flight_input['dest']
        )

    df_input = pd.DataFrame([flight_input])

    # ── Date-derived features (fl_date is the single source of truth) ────────
    fl_date = pd.to_datetime(flight_input['fl_date'])
    df_input['day_of_week']  = fl_date.dayofweek
    df_input['day_of_month'] = fl_date.day
    df_input['month']        = fl_date.month
    df_input['year']         = fl_date.year

    days_to_holiday              = int(min(abs((_holidays - fl_date).days)))
    df_input['days_to_holiday']  = days_to_holiday
    df_input['is_holiday_week']  = int(days_to_holiday <= 3)
    df_input['is_peak_season']   = int(fl_date.month in [6, 7, 8, 12])

    # ── Departure period bucket ───────────────────────────────────────────────
    df_input['dep_hour']   = df_input['crs_dep_time'] // 100
    df_input['dep_period'] = pd.cut(
        df_input['dep_hour'],
        bins=[0, 6, 12, 17, 21, 24],
        labels=['red_eye', 'morning', 'afternoon', 'evening', 'night'],
        right=False,
    )

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

    # ── Drop raw columns the model never saw ─────────────────────────────────
    df_input.drop(
        columns=[c for c in ['fl_date', 'dep_hour', 'crs_dep_time'] if c in df_input.columns],
        inplace=True,
    )

    # ── Encode categoricals using saved LabelEncoders ────────────────────────
    # Unseen labels (airport / carrier not in training data) → -1
    for col, le in les.items():
        if col in df_input.columns:
            df_input[col] = df_input[col].map(
                lambda x, _le=le: int(_le.transform([x])[0]) if x in _le.classes_ else -1
            )

    # ── Align columns to exact training order ────────────────────────────────
    df_input = df_input.reindex(columns=feat_cols, fill_value=0)

    # ── Impute numeric columns (after encoding so all cols are numeric) ───────
    valid_num_cols = [c for c in num_cols_ if c in df_input.columns]
    df_input[valid_num_cols] = num_imp.transform(df_input[valid_num_cols])

    # ── Predict ───────────────────────────────────────────────────────────────
    prob  = float(model.predict_proba(df_input)[0][1])
    label = int(prob >= 0.5)

    return {
        "prediction":              "DELAYED" if label else "ON TIME",
        "delay_probability":       round(prob, 4),
        "confidence":              f"{max(prob, 1 - prob) * 100:.1f}%",
        "distance_miles":          flight_input['distance'],
        "dep_delay":               flight_input['dep_delay'],
        "is_holiday_week":         bool(days_to_holiday <= 3),
        "days_to_holiday":         days_to_holiday,
        "is_peak_season":          bool(fl_date.month in [6, 7, 8, 12]),
        "weather_flag":            flight_input['had_weather_delay'],
        "model_used":              meta["best_model_name"],
    }


# ─────────────────────────────────────────────────────────────────────────────
# FRONTEND USAGE EXAMPLES
# Run this file directly to verify the interface works end-to-end:
#     python predict.py
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":

    def _print_result(label: str, result: dict):
        icon = "DELAYED" if result["prediction"] == "DELAYED" else "ON TIME"
        print(f"\n{label}")
        print(f"  {icon}  —  {result['delay_probability']*100:.1f}% delay probability  (confidence {result['confidence']})")
        print(f"  Distance:        {result['distance_miles']:.0f} mi")
        print(f"  Dep delay in:    {result['dep_delay']} min")
        print(f"  Holiday week:    {result['is_holiday_week']}  ({result['days_to_holiday']} days to nearest holiday)")
        print(f"  Peak season:     {result['is_peak_season']}")
        print(f"  Weather flag:    {result['weather_flag']}")
        print(f"  Model:           {result['model_used']}")

    # ── Example 1: Christmas Eve + snow, pre-booking (no dep_delay) ──────────
    _print_result(
        "Example 1 — Christmas Eve, ORD->LAX, snowstorm, pre-booking",
        predict_delay({
            "carrier":           "UA",
            "origin":            "ORD",
            "dest":              "LAX",
            "crs_dep_time":      1800,
            "fl_date":           "2024-12-24",
            "had_weather_delay": WEATHER_CONDITIONS["Snow / Ice"],
            # dep_delay omitted — defaults to 0 (pre-booking mode)
        }),
    )

    # ── Example 2: Same flight, at-gate, already 45 min late ─────────────────
    _print_result(
        "Example 2 — Christmas Eve, ORD->LAX, snowstorm, at-gate 45 min late",
        predict_delay({
            "carrier":           "UA",
            "origin":            "ORD",
            "dest":              "LAX",
            "crs_dep_time":      1800,
            "fl_date":           "2024-12-24",
            "had_weather_delay": WEATHER_CONDITIONS["Snow / Ice"],
            "dep_delay":         45.0,   # flight is already 45 min delayed at gate
        }),
    )

    # ── Example 3: At-gate, early departure (negative dep_delay) ─────────────
    _print_result(
        "Example 3 — Regular Tuesday, ATL->JFK, departing 5 min early",
        predict_delay({
            "carrier":      "DL",
            "origin":       "ATL",
            "dest":         "JFK",
            "crs_dep_time": 900,
            "fl_date":      "2024-03-12",
            "dep_delay":    -5.0,   # 5 minutes early
        }),
    )

    # ── Example 4: dep_delay sweep — same flight, increasing gate delay ───────
    print("\nExample 4 — ORD->LAX July, how delay probability grows with gate delay")
    for dep_delay in [0, 15, 30, 60, 120]:
        r = predict_delay({
            "carrier":      "UA",
            "origin":       "ORD",
            "dest":         "LAX",
            "crs_dep_time": 1200,
            "fl_date":      "2024-07-15",
            "dep_delay":    float(dep_delay),
        })
        print(f"  dep_delay={dep_delay:3d} min:  {r['prediction']:8s}  {r['delay_probability']*100:.1f}%")

    # ── Example 5: Thanksgiving → holiday week flag ───────────────────────────
    _print_result(
        "Example 5 — Thanksgiving, LAS->ORD, pre-booking",
        predict_delay({
            "carrier":      "WN",
            "origin":       "LAS",
            "dest":         "ORD",
            "crs_dep_time": 1200,
            "fl_date":      "2025-11-27",
        }),
    )

    # ── Example 6: Carrier comparison on same route/date ─────────────────────
    print("\nExample 6 — ATL->JFK, same morning in July, 3 carriers compared")
    for carrier in ["DL", "AA", "B6"]:
        r = predict_delay({
            "carrier":      carrier,
            "origin":       "ATL",
            "dest":         "JFK",
            "crs_dep_time": 900,
            "fl_date":      "2024-07-15",
        })
        print(f"  {carrier}:  {r['prediction']:8s}  {r['delay_probability']*100:.1f}%")

    # ── Example 7: Seasonality sweep — same flight across all 12 months ───────
    print("\nExample 7 — ORD->MIA, noon departure, every month of 2024")
    for month in range(1, 13):
        fl_date = f"2024-{month:02d}-15"
        r = predict_delay({
            "carrier":      "UA",
            "origin":       "ORD",
            "dest":         "MIA",
            "crs_dep_time": 1200,
            "fl_date":      fl_date,
        })
        print(f"  {fl_date}:  {r['prediction']:8s}  {r['delay_probability']*100:.1f}%")