src/streamlit_app.py

"""
Flight Delay Prediction - Streamlit App
Dataset: Kaggle 2015 Flight Delays (flights.csv) o datos sinteticos de demo.
Arquitectura: entrenamiento on-demand guardado en st.session_state
(sin depender de archivos .joblib persistidos en disco -> evita problemas
de Spaces que reinician o no conservan outputs/ entre despliegues).
"""

import io
import warnings

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import streamlit as st

from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    roc_auc_score, confusion_matrix, RocCurveDisplay,
)
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

warnings.filterwarnings("ignore")

st.set_page_config(
    page_title="Prediccion de Retrasos de Vuelos",
    page_icon="\u2708\ufe0f",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown("""
<style>
    .main-header { font-size: 2.4rem; font-weight: 800; color: #1F3864; margin-bottom: 0.2rem; }
    .sub-header { font-size: 1rem; color: #555; margin-bottom: 1.5rem; }
    .section-title {
        font-size: 1.3rem; font-weight: 700; color: #2c3e50;
        border-bottom: 2px solid #2E75B6; padding-bottom: 4px; margin-bottom: 1rem;
    }
    .info-box {
        background: #eaf4fb; border: 1px solid #aed6f1; border-radius: 6px;
        padding: 0.8rem 1rem; margin-bottom: 1rem; font-size: 0.9rem;
    }
    .stTabs [data-baseweb="tab"] { font-weight: 600; }
</style>
""", unsafe_allow_html=True)

AIRLINES = ["AA", "DL", "UA", "WN", "B6", "AS", "NK", "F9"]
AIRPORTS = ["ATL", "LAX", "ORD", "DFW", "JFK", "SFO", "MIA", "DEN", "SEA", "BOS"]
DOW_LABELS = {1: "Lunes", 2: "Martes", 3: "Miercoles", 4: "Jueves", 5: "Viernes", 6: "Sabado", 7: "Domingo"}
FEATURES = [
    "AIRLINE", "ORIGIN_AIRPORT", "DESTINATION_AIRPORT", "ROUTE",
    "MONTH", "DAY_OF_WEEK", "DEP_HOUR", "DISTANCE", "IS_WEEKEND",
    "ROUTE_FREQUENCY", "AIRLINE_AVG_DELAY", "ORIGIN_AVG_DELAY", "DEST_AVG_DELAY",
]


@st.cache_data
def generate_synthetic_data(n_samples: int = 8000, seed: int = 42) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    month = rng.integers(1, 13, n_samples)
    day = rng.integers(1, 29, n_samples)
    day_of_week = rng.integers(1, 8, n_samples)
    airline = rng.choice(AIRLINES, n_samples)
    origin = rng.choice(AIRPORTS, n_samples)
    dest = rng.choice(AIRPORTS, n_samples)
    sched_hour = rng.integers(5, 24, n_samples)
    distance = rng.integers(150, 2700, n_samples)

    p = (
        0.12
        + 0.012 * (sched_hour >= 17)
        + 0.035 * np.isin(month, [6, 7, 12])
        + 0.03 * np.isin(origin, ["ORD", "JFK", "LAX"])
        + 0.02 * np.isin(airline, ["WN", "B6", "F9"])
        + 0.01 * (day_of_week >= 5)
    )
    delayed = rng.binomial(1, np.clip(p, 0.02, 0.6))
    dep_delay = np.where(delayed == 1, rng.gamma(2.2, 18, n_samples) + 16, rng.normal(0, 7, n_samples))
    dep_delay = np.round(dep_delay).astype(int)
    scheduled_departure = sched_hour * 100 + rng.choice([0, 5, 10, 15, 20, 30, 45, 50], n_samples)

    return pd.DataFrame({
        "YEAR": 2015, "MONTH": month, "DAY": day, "DAY_OF_WEEK": day_of_week,
        "AIRLINE": airline, "ORIGIN_AIRPORT": origin, "DESTINATION_AIRPORT": dest,
        "SCHEDULED_DEPARTURE": scheduled_departure, "DEPARTURE_DELAY": dep_delay,
        "DISTANCE": distance, "CANCELLED": 0, "DIVERTED": 0,
    })


@st.cache_data
def load_uploaded(file_bytes) -> pd.DataFrame:
    return pd.read_csv(io.BytesIO(file_bytes), low_memory=False)


def clean_data(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    if "CANCELLED" in df.columns:
        df = df[df["CANCELLED"] == 0]
    if "DIVERTED" in df.columns:
        df = df[df["DIVERTED"] == 0]
    df = df.drop_duplicates()

    df["SCHEDULED_DEPARTURE"] = pd.to_numeric(df["SCHEDULED_DEPARTURE"], errors="coerce")
    df["DEPARTURE_DELAY"] = pd.to_numeric(df["DEPARTURE_DELAY"], errors="coerce")
    df = df[df["DEPARTURE_DELAY"].between(-60, 600)]

    for col in ["AIRLINE", "ORIGIN_AIRPORT", "DESTINATION_AIRPORT"]:
        df[col] = df[col].astype(str).str.upper().str.strip()
    df = df.dropna(subset=["AIRLINE", "ORIGIN_AIRPORT", "DESTINATION_AIRPORT", "DEPARTURE_DELAY", "SCHEDULED_DEPARTURE"])
    return df


def feature_engineer(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["DELAYED"] = (df["DEPARTURE_DELAY"] > 15).astype(int)
    df["DEP_HOUR"] = (df["SCHEDULED_DEPARTURE"] // 100).clip(0, 23).astype(int)
    df["IS_WEEKEND"] = df["DAY_OF_WEEK"].isin([6, 7]).astype(int)
    df["ROUTE"] = df["ORIGIN_AIRPORT"] + "_" + df["DESTINATION_AIRPORT"]
    return df


def compute_stats(df: pd.DataFrame) -> dict:
    return {
        "route_freq": df.groupby("ROUTE").size(),
        "airline_avg_delay": df.groupby("AIRLINE")["DEPARTURE_DELAY"].mean(),
        "origin_avg_delay": df.groupby("ORIGIN_AIRPORT")["DEPARTURE_DELAY"].mean(),
        "dest_avg_delay": df.groupby("DESTINATION_AIRPORT")["DEPARTURE_DELAY"].mean(),
    }


def apply_stats(df: pd.DataFrame, stats: dict) -> pd.DataFrame:
    df = df.copy()
    df["ROUTE_FREQUENCY"] = df["ROUTE"].map(stats["route_freq"]).fillna(1)
    df["AIRLINE_AVG_DELAY"] = df["AIRLINE"].map(stats["airline_avg_delay"]).fillna(0)
    df["ORIGIN_AVG_DELAY"] = df["ORIGIN_AIRPORT"].map(stats["origin_avg_delay"]).fillna(0)
    df["DEST_AVG_DELAY"] = df["DESTINATION_AIRPORT"].map(stats["dest_avg_delay"]).fillna(0)
    return df


def build_preprocessor(X: pd.DataFrame) -> ColumnTransformer:
    cat_cols = [c for c in ["AIRLINE", "ORIGIN_AIRPORT", "DESTINATION_AIRPORT", "ROUTE"] if c in X.columns]
    num_cols = [c for c in FEATURES if c not in cat_cols]
    try:
        encoder = OneHotEncoder(handle_unknown="ignore", sparse_output=False)
    except TypeError:
        encoder = OneHotEncoder(handle_unknown="ignore", sparse=False)
    return ColumnTransformer([
        ("num", Pipeline([("imputer", SimpleImputer(strategy="median")), ("scaler", StandardScaler())]), num_cols),
        ("cat", Pipeline([("imputer", SimpleImputer(strategy="most_frequent")), ("onehot", encoder)]), cat_cols),
    ])


def train_models(raw_df, model_names, n_estimators, test_size):
    clean = clean_data(raw_df)
    feat = feature_engineer(clean)

    train_idx, test_idx = train_test_split(
        feat.index, test_size=test_size, random_state=42, stratify=feat["DELAYED"]
    )
    train_df, test_df = feat.loc[train_idx], feat.loc[test_idx]

    stats = compute_stats(train_df)
    train_df = apply_stats(train_df, stats)
    test_df = apply_stats(test_df, stats)

    X_train, X_test = train_df[FEATURES], test_df[FEATURES]
    y_train, y_test = train_df["DELAYED"], test_df["DELAYED"]

    preprocessor = build_preprocessor(X_train)
    model_map = {}
    if "Bosque Aleatorio" in model_names:
        model_map["Bosque Aleatorio"] = RandomForestClassifier(
            n_estimators=n_estimators, max_depth=10, class_weight="balanced", random_state=42, n_jobs=-1
        )
    if "Gradient Boosting" in model_names:
        model_map["Gradient Boosting"] = GradientBoostingClassifier(
            n_estimators=min(n_estimators, 100), random_state=42
        )
    if "Regresion Logistica" in model_names:
        model_map["Regresion Logistica"] = LogisticRegression(max_iter=1000, class_weight="balanced")

    results = {}
    for name, clf in model_map.items():
        pipe = Pipeline([("preprocess", preprocessor), ("model", clf)])
        pipe.fit(X_train, y_train)
        y_pred = pipe.predict(X_test)
        y_prob = pipe.predict_proba(X_test)[:, 1]
        results[name] = {
            "pipe": pipe,
            "acc": accuracy_score(y_test, y_pred),
            "precision": precision_score(y_test, y_pred, zero_division=0),
            "recall": recall_score(y_test, y_pred, zero_division=0),
            "f1": f1_score(y_test, y_pred, zero_division=0),
            "roc_auc": roc_auc_score(y_test, y_prob),
            "cm": confusion_matrix(y_test, y_pred),
            "y_test": y_test, "y_prob": y_prob,
        }

    fi = None
    if "Bosque Aleatorio" in results:
        rf_pipe = results["Bosque Aleatorio"]["pipe"]
        cat_names = rf_pipe.named_steps["preprocess"].named_transformers_["cat"].named_steps["onehot"].get_feature_names_out()
        num_cols = [c for c in FEATURES if c not in ["AIRLINE", "ORIGIN_AIRPORT", "DESTINATION_AIRPORT", "ROUTE"]]
        all_names = list(num_cols) + list(cat_names)
        importances = rf_pipe.named_steps["model"].feature_importances_
        fi = pd.Series(importances, index=all_names).sort_values(ascending=False).head(15)

    return results, stats, fi


def predict_flight(pipe, stats, airline, origin, dest, month, day_of_week, dep_hour, distance):
    route = f"{origin}_{dest}"
    row = pd.DataFrame([{
        "AIRLINE": airline, "ORIGIN_AIRPORT": origin, "DESTINATION_AIRPORT": dest, "ROUTE": route,
        "MONTH": month, "DAY_OF_WEEK": day_of_week, "DEP_HOUR": dep_hour, "DISTANCE": distance,
        "IS_WEEKEND": int(day_of_week in (6, 7)),
        "ROUTE_FREQUENCY": stats["route_freq"].get(route, 1),
        "AIRLINE_AVG_DELAY": stats["airline_avg_delay"].get(airline, 0.0),
        "ORIGIN_AVG_DELAY": stats["origin_avg_delay"].get(origin, 0.0),
        "DEST_AVG_DELAY": stats["dest_avg_delay"].get(dest, 0.0),
    }])[FEATURES]
    return float(pipe.predict_proba(row)[0, 1])


with st.sidebar:
    st.markdown("## \u2699\ufe0f Configuracion")

    st.markdown("### \U0001F4C2 Fuente de datos")
    uploaded = st.file_uploader(
        "Sube flights.csv (opcional)", type=["csv"],
        help="Sube el CSV real de Kaggle (2015 Flight Delays). Si no, se usa data sintetica de demo."
    )
    n_synthetic = st.slider("Muestras sinteticas (si no subes archivo)", 2000, 30000, 8000, step=1000)

    st.markdown("---")
    st.markdown("### \U0001F916 Modelos a entrenar")
    model_choices = st.multiselect(
        "Selecciona clasificadores",
        ["Bosque Aleatorio", "Gradient Boosting", "Regresion Logistica"],
        default=["Bosque Aleatorio", "Regresion Logistica"],
    )

    st.markdown("### \U0001F522 Parametros de entrenamiento")
    test_frac = st.slider("Fraccion de test", 0.1, 0.4, 0.25, step=0.05)
    n_trees = st.slider("Arboles (RF / GB)", 20, 200, 80, step=10)

    st.markdown("---")
    train_btn = st.button("\U0001F680 Entrenar modelos", width='stretch', type="primary")

    st.markdown("---")
    st.markdown("### \U0001F52E Predictor de retraso")
    pred_airline = st.selectbox("Aerolinea", AIRLINES)
    pred_origin = st.selectbox("Aeropuerto origen", AIRPORTS, index=0)
    pred_dest = st.selectbox("Aeropuerto destino", AIRPORTS, index=1)
    pred_month = st.slider("Mes", 1, 12, 6)
    pred_dow_label = st.selectbox("Dia de la semana", list(DOW_LABELS.values()), index=4)
    pred_hour = st.slider("Hora de salida (24h)", 0, 23, 17)
    pred_distance = st.slider("Distancia (millas)", 100, 3000, 900, step=50)
    predict_btn = st.button("\U0001F50D Predecir retraso", width='stretch')

if uploaded:
    raw_df = load_uploaded(uploaded.read())
    data_source = "CSV subido"
else:
    raw_df = generate_synthetic_data(n_synthetic)
    data_source = f"Datos sinteticos de demo ({n_synthetic:,} muestras)"

clean_df = clean_data(raw_df)
feat_df = feature_engineer(clean_df)

st.markdown('<p class="main-header">\u2708\ufe0f Prediccion de Retrasos de Vuelos</p>', unsafe_allow_html=True)
st.markdown(
    f'<p class="sub-header">Machine Learning sobre datos de vuelos &nbsp;|&nbsp; '
    f'Datos: <b>{data_source}</b> &nbsp;|&nbsp; {len(raw_df):,} registros</p>',
    unsafe_allow_html=True,
)

for key in ["results", "stats", "fi"]:
    if key not in st.session_state:
        st.session_state[key] = None

if train_btn:
    if not model_choices:
        st.warning("Selecciona al menos un modelo.")
    else:
        with st.spinner("Entrenando modelos... puede tardar un momento"):
            res, stats, fi = train_models(raw_df, model_choices, n_trees, test_frac)
        st.session_state.results = res
        st.session_state.stats = stats
        st.session_state.fi = fi
        st.success("Modelos entrenados correctamente.")

tab1, tab2, tab3, tab4 = st.tabs([
    "\U0001F4CA Analisis exploratorio", "\U0001F9E0 Resultados del modelo",
    "\U0001F4C8 Importancia de variables", "\U0001F52E Predictor",
])

with tab1:
    st.markdown('<p class="section-title">Analisis Exploratorio</p>', unsafe_allow_html=True)

    c1, c2, c3, c4 = st.columns(4)
    c1.metric("Vuelos totales", f"{len(raw_df):,}")
    c2.metric("Tasa de retraso", f"{feat_df['DELAYED'].mean()*100:.1f}%")
    c3.metric("Aerolineas", feat_df["AIRLINE"].nunique())
    c4.metric("Aeropuertos", feat_df["ORIGIN_AIRPORT"].nunique())

    st.markdown("---")
    col_a, col_b = st.columns(2)

    with col_a:
        st.markdown("**Tasa de retraso por aerolinea**")
        by_airline = feat_df.groupby("AIRLINE")["DELAYED"].mean().sort_values(ascending=False)
        fig, ax = plt.subplots(figsize=(6, 4))
        colors = plt.cm.Blues_r(np.linspace(0.2, 0.8, len(by_airline)))
        by_airline.plot(kind="bar", ax=ax, color=colors)
        ax.set_ylabel("Tasa de retraso"); ax.tick_params(axis="x", rotation=0)
        fig.tight_layout(); st.pyplot(fig); plt.close(fig)

    with col_b:
        st.markdown("**Tasa de retraso por aeropuerto de origen**")
        by_airport = feat_df.groupby("ORIGIN_AIRPORT")["DELAYED"].mean().sort_values(ascending=False).head(10)
        fig, ax = plt.subplots(figsize=(6, 4))
        by_airport.plot(kind="bar", ax=ax, color="#F0A500", alpha=0.85)
        ax.set_ylabel("Tasa de retraso"); ax.tick_params(axis="x", rotation=45)
        fig.tight_layout(); st.pyplot(fig); plt.close(fig)

    col_c, col_d = st.columns(2)
    with col_c:
        st.markdown("**Tasa de retraso por hora de salida**")
        by_hour = feat_df.groupby("DEP_HOUR")["DELAYED"].mean()
        fig, ax = plt.subplots(figsize=(6, 3.5))
        ax.plot(by_hour.index, by_hour.values, color="#1F3864", linewidth=2.5, marker="o", markersize=4)
        ax.fill_between(by_hour.index, by_hour.values, alpha=0.15, color="#1F3864")
        ax.set_xlabel("Hora (24h)"); ax.set_ylabel("Tasa de retraso")
        fig.tight_layout(); st.pyplot(fig); plt.close(fig)

    with col_d:
        st.markdown("**Tasa de retraso por mes**")
        by_month = feat_df.groupby("MONTH")["DELAYED"].mean()
        fig, ax = plt.subplots(figsize=(6, 3.5))
        ax.bar(by_month.index, by_month.values, color="#2E75B6", alpha=0.85)
        ax.set_xlabel("Mes"); ax.set_ylabel("Tasa de retraso")
        fig.tight_layout(); st.pyplot(fig); plt.close(fig)

    st.markdown("**Distribucion de minutos de retraso**")
    fig, ax = plt.subplots(figsize=(12, 3))
    ax.hist(clean_df["DEPARTURE_DELAY"].dropna(), bins=60, color="#1F3864", alpha=0.8)
    ax.set_xlabel("Minutos de retraso"); ax.set_ylabel("Cantidad de vuelos")
    fig.tight_layout(); st.pyplot(fig); plt.close(fig)

    st.markdown("**Mapa de calor: Hora x Dia de la semana**")
    pivot_data = feat_df.groupby(["DAY_OF_WEEK", "DEP_HOUR"])["DELAYED"].mean().unstack(fill_value=0)
    fig, ax = plt.subplots(figsize=(14, 4))
    sns.heatmap(pivot_data, ax=ax, cmap="YlOrRd", linewidths=0.3, cbar_kws={"label": "Tasa de retraso"})
    ax.set_xlabel("Hora del dia"); ax.set_ylabel("Dia de la semana (1=Lun)")
    fig.tight_layout(); st.pyplot(fig); plt.close(fig)

    with st.expander("Vista previa de los datos"):
        st.dataframe(raw_df.head(200), width='stretch')

with tab2:
    st.markdown('<p class="section-title">Evaluacion de Modelos</p>', unsafe_allow_html=True)

    if st.session_state.results is None:
        st.markdown(
            '<div class="info-box">Configura tus modelos en la barra lateral y presiona '
            '<b>Entrenar modelos</b> para ver resultados aqui.</div>',
            unsafe_allow_html=True,
        )
    else:
        results = st.session_state.results

        st.markdown("**Comparacion de metricas**")
        metrics_df = pd.DataFrame({
            name: {"Accuracy": r["acc"], "Precision": r["precision"], "Recall": r["recall"],
                   "F1": r["f1"], "ROC AUC": r["roc_auc"]}
            for name, r in results.items()
        }).T.round(4)
        st.dataframe(metrics_df, width='stretch')

        fig, ax = plt.subplots(figsize=(max(4, len(results) * 2.5), 4))
        bars = ax.bar(metrics_df.index, metrics_df["ROC AUC"], color=["#1F3864", "#2E75B6", "#F0A500"][:len(results)], alpha=0.85)
        ax.axhline(0.5, color="gray", linestyle="--", linewidth=1, label="Azar (0.5)")
        ax.set_ylim(0, 1); ax.set_ylabel("ROC AUC"); ax.legend()
        for bar, val in zip(bars, metrics_df["ROC AUC"]):
            ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01, f"{val:.3f}", ha="center", fontweight="bold")
        fig.tight_layout(); st.pyplot(fig); plt.close(fig)

        fig, ax = plt.subplots(figsize=(7, 5))
        for name, r in results.items():
            RocCurveDisplay.from_predictions(r["y_test"], r["y_prob"], name=name, ax=ax)
        ax.set_title("Curvas ROC"); fig.tight_layout(); st.pyplot(fig); plt.close(fig)

        for name, r in results.items():
            with st.expander(f"{name} - ROC AUC: {r['roc_auc']:.4f}"):
                fig, ax = plt.subplots(figsize=(5, 4))
                sns.heatmap(r["cm"], annot=True, fmt="d", cmap="Blues",
                            xticklabels=["No retraso", "Retraso"], yticklabels=["No retraso", "Retraso"], ax=ax)
                ax.set_xlabel("Prediccion"); ax.set_ylabel("Real")
                fig.tight_layout(); st.pyplot(fig); plt.close(fig)

with tab3:
    st.markdown('<p class="section-title">Importancia de Variables</p>', unsafe_allow_html=True)

    if st.session_state.fi is None:
        st.markdown(
            '<div class="info-box">Entrena un modelo <b>Bosque Aleatorio</b> para ver la importancia de variables.</div>',
            unsafe_allow_html=True,
        )
    else:
        fi = st.session_state.fi
        col_fi1, col_fi2 = st.columns([2, 1])
        with col_fi1:
            fig, ax = plt.subplots(figsize=(8, 5))
            colors = plt.cm.RdYlGn(np.linspace(0.2, 0.8, len(fi)))[::-1]
            fi.plot(kind="barh", ax=ax, color=colors)
            ax.set_xlabel("Importancia"); ax.set_title("Importancia de variables - Bosque Aleatorio")
            ax.invert_yaxis()
            fig.tight_layout(); st.pyplot(fig); plt.close(fig)
        with col_fi2:
            st.dataframe(fi.reset_index().rename(columns={"index": "Variable", 0: "Importancia"}).round(4), width='stretch')

with tab4:
    st.markdown('<p class="section-title">Predictor de Retraso de Vuelo</p>', unsafe_allow_html=True)
    st.markdown(
        '<div class="info-box">Configura los datos del vuelo en la barra lateral, entrena un modelo, '
        'luego presiona <b>Predecir retraso</b>.</div>',
        unsafe_allow_html=True,
    )

    if predict_btn:
        if st.session_state.results is None:
            st.warning("Primero entrena al menos un modelo.")
        elif pred_origin == pred_dest:
            st.warning("El aeropuerto de origen y destino no pueden ser el mismo.")
        else:
            results = st.session_state.results
            stats = st.session_state.stats
            day_of_week = [k for k, v in DOW_LABELS.items() if v == pred_dow_label][0]

            st.markdown("### Predicciones")
            for name, r in results.items():
                prob = predict_flight(
                    r["pipe"], stats, pred_airline, pred_origin, pred_dest,
                    pred_month, day_of_week, pred_hour, pred_distance,
                )
                pct = prob * 100
                if pct < 20:
                    nivel, color = "Bajo riesgo", "#27ae60"
                elif pct < 45:
                    nivel, color = "Riesgo moderado", "#F0A500"
                else:
                    nivel, color = "Alto riesgo", "#c0392b"

                st.markdown(f"**{name}** -> `{pct:.1f}%` de probabilidad de retraso - "
                            f"<span style='color:{color}; font-weight:700'>{nivel}</span>", unsafe_allow_html=True)
                st.progress(min(int(pct), 100))

st.markdown("---")
st.markdown(
    "<small>Prediccion de retrasos de vuelos. Dataset: Kaggle 2015 Flight Delays "
    "(o datos sinteticos de demo si no se sube CSV)</small>",
    unsafe_allow_html=True,
)