src/streamlit_app.py

from __future__ import annotations

import os
import warnings
from pathlib import Path
from typing import Dict, Tuple

import joblib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.base import clone
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
from sklearn.inspection import permutation_importance

warnings.filterwarnings("ignore")

RANDOM_STATE = 42
DATA_DIR = Path("/app/data")
OUTPUT_DIR = Path("outputs")
OUTPUT_DIR.mkdir(exist_ok=True)


def create_demo_dataset(n: int = 8000) -> pd.DataFrame:
    """Crea un dataset demo si no existen los CSV reales de Kaggle."""
    rng = np.random.default_rng(RANDOM_STATE)
    airlines = np.array(["AA", "DL", "UA", "WN", "B6", "AS", "NK", "F9"])
    airports = np.array(["ATL", "LAX", "ORD", "DFW", "JFK", "SFO", "MIA", "DEN", "SEA", "BOS"])
    month = rng.integers(1, 13, n)
    day = rng.integers(1, 29, n)
    day_of_week = rng.integers(1, 8, n)
    airline = rng.choice(airlines, n)
    origin = rng.choice(airports, n)
    dest = rng.choice(airports, n)
    sched_hour = rng.integers(5, 24, n)
    distance = rng.integers(150, 2700, n)

    # Probabilidad con patrones realistas: tarde, diciembre/verano, aeropuertos congestionados.
    base = 0.12
    p = (
        base
        + 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) + 16, rng.normal(0, 7, n))
    dep_delay = np.round(dep_delay).astype(int)

    scheduled_departure = sched_hour * 100 + rng.choice([0, 5, 10, 15, 20, 30, 45, 50], n)
    actual_departure = scheduled_departure + dep_delay

    df = pd.DataFrame(
        {
            "YEAR": 2015,
            "MONTH": month,
            "DAY": day,
            "DAY_OF_WEEK": day_of_week,
            "AIRLINE": airline,
            "FLIGHT_NUMBER": rng.integers(1, 9000, n),
            "ORIGIN_AIRPORT": origin,
            "DESTINATION_AIRPORT": dest,
            "SCHEDULED_DEPARTURE": scheduled_departure,
            "DEPARTURE_TIME": actual_departure,
            "DEPARTURE_DELAY": dep_delay,
            "SCHEDULED_ARRIVAL": scheduled_departure + rng.integers(60, 320, n),
            "ARRIVAL_DELAY": dep_delay + rng.normal(0, 12, n).round().astype(int),
            "DISTANCE": distance,
        }
    )
    return df


def load_data(sample_size: int = 80_000) -> pd.DataFrame:
    flights_path = DATA_DIR / "flights.csv"
    if flights_path.exists():
        print("Cargando data/flights.csv...")
        # Leer columnas necesarias para ahorrar memoria.
        usecols = [
            "YEAR", "MONTH", "DAY", "DAY_OF_WEEK", "AIRLINE", "FLIGHT_NUMBER",
            "ORIGIN_AIRPORT", "DESTINATION_AIRPORT", "SCHEDULED_DEPARTURE",
            "DEPARTURE_TIME", "DEPARTURE_DELAY", "SCHEDULED_ARRIVAL",
            "ARRIVAL_DELAY", "DISTANCE", "CANCELLED", "DIVERTED"
        ]
        df = pd.read_csv(flights_path, usecols=lambda c: c in usecols, low_memory=False)
        # Dataset completo es grande. Para clase/proyecto, muestra estratificada simple.
        if len(df) > sample_size:
            df = df.sample(sample_size, random_state=RANDOM_STATE)
    else:
        print("No se encontró data/flights.csv. Usando dataset demo.")
        df = create_demo_dataset()
    return df


def understand_data(df: pd.DataFrame) -> None:
    summary = pd.DataFrame({
        "variable": df.columns,
        "tipo": [str(df[c].dtype) for c in df.columns],
        "valores_faltantes": [df[c].isna().sum() for c in df.columns],
        "porcentaje_faltante": [df[c].isna().mean() * 100 for c in df.columns],
        "valores_unicos": [df[c].nunique(dropna=True) for c in df.columns],
    })
    summary.to_csv(OUTPUT_DIR / "data_understanding_summary.csv", index=False)
    print("Filas:", len(df), "Columnas:", df.shape[1])
    print(summary)


def clean_data(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    initial_rows = len(df)

    df = df.drop_duplicates()
    if "CANCELLED" in df.columns:
        df = df[df["CANCELLED"] == 0]
    if "DIVERTED" in df.columns:
        df = df[df["DIVERTED"] == 0]

    # Crear fecha real.
    df["FLIGHT_DATE"] = pd.to_datetime(
        dict(year=df["YEAR"], month=df["MONTH"], day=df["DAY"]), errors="coerce"
    )

    # Convertir horas HHMM a hora.
    for col in ["SCHEDULED_DEPARTURE", "DEPARTURE_TIME", "SCHEDULED_ARRIVAL"]:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")
            df[col + "_HOUR"] = (df[col] // 100).clip(0, 23)
            df[col + "_MINUTE"] = (df[col] % 100).clip(0, 59)

    # Corregir retrasos extremos que pueden ser errores de captura.
    if "DEPARTURE_DELAY" in df.columns:
        df["DEPARTURE_DELAY"] = pd.to_numeric(df["DEPARTURE_DELAY"], errors="coerce")
        df = df[df["DEPARTURE_DELAY"].between(-60, 600) | df["DEPARTURE_DELAY"].isna()]

    # Columnas categóricas a texto.
    for col in ["AIRLINE", "ORIGIN_AIRPORT", "DESTINATION_AIRPORT"]:
        df[col] = df[col].astype(str).str.upper().str.strip()
        df.loc[df[col].isin(["NAN", "NONE", ""]), col] = np.nan

    quality = pd.DataFrame({
        "metrica": ["filas_iniciales", "filas_finales", "duplicados_eliminados_o_cancelados"],
        "valor": [initial_rows, len(df), initial_rows - len(df)],
    })
    quality.to_csv(OUTPUT_DIR / "data_quality_summary.csv", index=False)
    df.to_csv(OUTPUT_DIR / "clean_flights.csv", index=False)
    return df


def add_basic_features(df: pd.DataFrame) -> pd.DataFrame:
    """Crea variables base sin estadísticas agregadas (evita fuga de datos en el modelo)."""
    df = df.copy()
    df["DELAYED"] = (df["DEPARTURE_DELAY"] > 15).astype(int)
    if "SCHEDULED_DEPARTURE_HOUR" in df.columns:
        dep_hour = df["SCHEDULED_DEPARTURE_HOUR"]
    else:
        dep_hour = (pd.to_numeric(df["SCHEDULED_DEPARTURE"], errors="coerce") // 100).clip(0, 23)
    df["DEP_HOUR"] = dep_hour.fillna(0).astype(int)
    df["IS_WEEKEND"] = df["DAY_OF_WEEK"].isin([6, 7]).astype(int)
    df["ROUTE"] = df["ORIGIN_AIRPORT"].astype(str) + "_" + df["DESTINATION_AIRPORT"].astype(str)
    return df


def compute_aggregate_stats(df: pd.DataFrame) -> Dict[str, pd.Series | float]:
    """Calcula estadísticas históricas solo a partir del conjunto de referencia (train)."""
    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(),
        "distance_median": float(df["DISTANCE"].median()) if "DISTANCE" in df.columns else 700.0,
    }


def apply_aggregate_features(df: pd.DataFrame, stats: Dict[str, pd.Series | float]) -> 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 add_features(df: pd.DataFrame) -> pd.DataFrame:
    """Dataset para EDA/dashboard. Los agregados usan todo el histórico (solo visualización)."""
    df = add_basic_features(df)
    stats = compute_aggregate_stats(df)
    df = apply_aggregate_features(df, stats)
    df.to_csv(OUTPUT_DIR / "model_ready_flights.csv", index=False)
    return df


def delay_rate(group: pd.DataFrame, col: str, top: int | None = None) -> pd.DataFrame:
    out = group.groupby(col).agg(
        flights=("DELAYED", "size"),
        delay_rate=("DELAYED", "mean"),
        avg_delay=("DEPARTURE_DELAY", "mean"),
    ).reset_index()
    out = out.sort_values("delay_rate", ascending=False)
    if top:
        out = out.head(top)
    return out


def plot_bar(df: pd.DataFrame, x: str, y: str, title: str, file_name: str, rotation: int = 45) -> None:
    plt.figure(figsize=(10, 5))
    plt.bar(df[x].astype(str), df[y])
    plt.title(title)
    plt.xlabel(x)
    plt.ylabel(y)
    plt.xticks(rotation=rotation, ha="right")
    plt.tight_layout()
    plt.savefig(OUTPUT_DIR / file_name, dpi=150)
    plt.close()


def exploratory_analysis(df: pd.DataFrame) -> None:
    airline = delay_rate(df, "AIRLINE")
    airport = delay_rate(df, "ORIGIN_AIRPORT", top=20)
    month = delay_rate(df, "MONTH")
    hour = delay_rate(df, "DEP_HOUR")
    dow = delay_rate(df, "DAY_OF_WEEK")

    airline.to_csv(OUTPUT_DIR / "eda_delay_rate_by_airline.csv", index=False)
    airport.to_csv(OUTPUT_DIR / "eda_delay_rate_by_airport.csv", index=False)
    month.to_csv(OUTPUT_DIR / "eda_delay_rate_by_month.csv", index=False)
    hour.to_csv(OUTPUT_DIR / "eda_delay_rate_by_departure_hour.csv", index=False)
    dow.to_csv(OUTPUT_DIR / "eda_delay_rate_by_day_of_week.csv", index=False)

    plot_bar(airline, "AIRLINE", "delay_rate", "Tasa de retrasos por aerolínea", "delay_rate_airline.png")
    plot_bar(airport, "ORIGIN_AIRPORT", "delay_rate", "Tasa de retrasos por aeropuerto origen", "delay_rate_airport.png")
    plot_bar(month, "MONTH", "delay_rate", "Tasa de retrasos por mes", "delay_rate_month.png", 0)
    plot_bar(hour, "DEP_HOUR", "delay_rate", "Tasa de retrasos por hora de salida", "delay_rate_hour.png", 0)

    plt.figure(figsize=(10, 5))
    plt.hist(df["DEPARTURE_DELAY"].dropna(), bins=60)
    plt.title("Distribución de minutos de retraso en salida")
    plt.xlabel("Minutos de retraso")
    plt.ylabel("Cantidad de vuelos")
    plt.tight_layout()
    plt.savefig(OUTPUT_DIR / "delay_distribution.png", dpi=150)
    plt.close()


def build_preprocessor(X: pd.DataFrame) -> Tuple[ColumnTransformer, list, list]:
    categorical_features = ["AIRLINE", "ORIGIN_AIRPORT", "DESTINATION_AIRPORT", "ROUTE"]
    numeric_features = [
        "MONTH", "DAY_OF_WEEK", "DEP_HOUR", "DISTANCE", "IS_WEEKEND",
        "ROUTE_FREQUENCY", "AIRLINE_AVG_DELAY", "ORIGIN_AVG_DELAY", "DEST_AVG_DELAY"
    ]
    categorical_features = [c for c in categorical_features if c in X.columns]
    numeric_features = [c for c in numeric_features if c in X.columns]

    try:
        encoder = OneHotEncoder(handle_unknown="ignore", sparse_output=False)
    except TypeError:
        encoder = OneHotEncoder(handle_unknown="ignore", sparse=False)

    preprocessor = ColumnTransformer(
        transformers=[
            ("num", Pipeline([("imputer", SimpleImputer(strategy="median")), ("scaler", StandardScaler())]), numeric_features),
            ("cat", Pipeline([("imputer", SimpleImputer(strategy="most_frequent")), ("onehot", encoder)]), categorical_features),
        ],
        remainder="drop",
    )
    return preprocessor, numeric_features, categorical_features


def train_and_evaluate(df: pd.DataFrame) -> Dict[str, Pipeline]:
    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"
    ]

    base_df = add_basic_features(df)
    base_df = base_df.dropna(subset=["DELAYED", "AIRLINE", "ORIGIN_AIRPORT", "DESTINATION_AIRPORT"])
    if base_df["DELAYED"].nunique() < 2:
        raise ValueError("La variable objetivo DELAYED debe tener al menos dos clases.")

    train_idx, test_idx = train_test_split(
        base_df.index,
        test_size=0.25,
        random_state=RANDOM_STATE,
        stratify=base_df["DELAYED"],
    )

    train_base = base_df.loc[train_idx]
    test_base = base_df.loc[test_idx]

    stats = compute_aggregate_stats(train_base)
    train_df = apply_aggregate_features(train_base, stats)
    test_df = apply_aggregate_features(test_base, stats)

    features = [c for c in features if c in train_df.columns]
    X_train = train_df[features]
    X_test = test_df[features]
    y_train = train_df["DELAYED"]
    y_test = test_df["DELAYED"]

    preprocessor, _, _ = build_preprocessor(X_train)
    models = {
        "Regresion Logistica": LogisticRegression(max_iter=1000, class_weight="balanced"),
        "Bosque Aleatorio": RandomForestClassifier(
            n_estimators=80, max_depth=10, random_state=RANDOM_STATE,
            class_weight="balanced", n_jobs=1,
        ),
        "Gradient Boosting": GradientBoostingClassifier(random_state=RANDOM_STATE),
    }

    trained_models: Dict[str, Pipeline] = {}
    rows = []
    roc_data = {}

    for name, clf in models.items():
        pipe = Pipeline([("preprocess", clone(preprocessor)), ("model", clf)])
        pipe.fit(X_train, y_train)
        y_pred = pipe.predict(X_test)
        y_prob = pipe.predict_proba(X_test)[:, 1]
        rows.append({
            "model": name,
            "accuracy": 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),
        })
        trained_models[name] = pipe
        roc_data[name] = (y_pred, y_prob)
        print(name, rows[-1])

    metrics = pd.DataFrame(rows).sort_values("roc_auc", ascending=False)
    metrics.to_csv(OUTPUT_DIR / "model_metrics.csv", index=False)

    best_name = metrics.iloc[0]["model"]
    best_model = trained_models[best_name]
    y_pred, y_prob = roc_data[best_name]

    cm = confusion_matrix(y_test, y_pred)
    pd.DataFrame(
        cm,
        index=["Real_No_Retraso", "Real_Retraso"],
        columns=["Pred_No_Retraso", "Pred_Retraso"],
    ).to_csv(OUTPUT_DIR / "confusion_matrix.csv")

    plt.figure(figsize=(6, 5))
    plt.imshow(cm)
    plt.title(f"Matriz de confusión - {best_name}")
    plt.xlabel("Predicción")
    plt.ylabel("Real")
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            plt.text(j, i, cm[i, j], ha="center", va="center")
    plt.tight_layout()
    plt.savefig(OUTPUT_DIR / "confusion_matrix.png", dpi=150)
    plt.close()

    plt.figure(figsize=(7, 5))
    for name, (_, prob) in roc_data.items():
        RocCurveDisplay.from_predictions(y_test, prob, name=name)
    plt.title("Curvas ROC por modelo")
    plt.tight_layout()
    plt.savefig(OUTPUT_DIR / "roc_curve.png", dpi=150)
    plt.close()

    result = permutation_importance(
        best_model, X_test, y_test, n_repeats=3, random_state=RANDOM_STATE,
        scoring="roc_auc", n_jobs=1,
    )
    imp = pd.DataFrame({
        "feature": X_test.columns,
        "importance": result.importances_mean,
    }).sort_values("importance", ascending=False)
    imp.to_csv(OUTPUT_DIR / "feature_importance.csv", index=False)
    plot_bar(imp.head(10), "feature", "importance", "Top 10 variables importantes", "feature_importance.png")

    joblib.dump(best_model, OUTPUT_DIR / "best_model.joblib")
    joblib.dump(features, OUTPUT_DIR / "feature_columns.joblib")
    joblib.dump(stats, OUTPUT_DIR / "inference_stats.joblib")
    X_test.to_csv(OUTPUT_DIR / "X_test.csv", index=False)
    y_test.to_csv(OUTPUT_DIR / "y_test.csv", index=False)
    return trained_models


def write_recommendations(df: pd.DataFrame) -> None:
    top_airline = delay_rate(df, "AIRLINE").head(3)
    top_airport = delay_rate(df, "ORIGIN_AIRPORT").head(3)
    top_month = delay_rate(df, "MONTH").head(3)
    top_hour = delay_rate(df, "DEP_HOUR").head(3)

    text = f"""
RECOMENDACIONES EMPRESARIALES

Para aerolíneas:
- Revisar procesos operacionales en las aerolíneas con mayor tasa de retraso: {', '.join(top_airline['AIRLINE'].astype(str))}.
- Optimizar rotación de tripulación, mantenimiento preventivo y tiempos de conexión en rutas frecuentes.
- Usar alertas tempranas para vuelos en horarios de mayor riesgo.

Para aeropuertos:
- Aumentar personal y recursos en los aeropuertos de origen con mayor tasa de retraso: {', '.join(top_airport['ORIGIN_AIRPORT'].astype(str))}.
- Reforzar ramp agents, gates, seguridad y coordinación ATC en los meses críticos: {', '.join(top_month['MONTH'].astype(str))}.
- Priorizar recursos en las horas con mayor riesgo: {', '.join(top_hour['DEP_HOUR'].astype(str))}.

Para viajeros:
- Preferir vuelos en la mañana cuando sea posible.
- Evitar conexiones cortas en aeropuertos con alta tasa de retrasos.
- En meses de alta demanda, elegir itinerarios con más tiempo de conexión.
"""
    (OUTPUT_DIR / "business_recommendations.txt").write_text(text, encoding="utf-8")


def main() -> None:
    raw = load_data()
    understand_data(raw)
    clean = clean_data(raw)
    feat = add_features(clean)
    exploratory_analysis(feat)
    print("Distribución de clases:")
    print(feat["DELAYED"].value_counts(normalize=True).rename("proportion"))
    feat["DELAYED"].value_counts().to_csv(OUTPUT_DIR / "class_distribution.csv")
    train_and_evaluate(clean)
    write_recommendations(feat)
    print("Proyecto completado. Revisa la carpeta outputs/.")


if __name__ == "__main__":
    main()