"""
Módulo de proyección de defunciones sobre la serie mensual de SINADEF.

Método (el mismo esquema de las competencias M3/M4):

  1. Se arma la serie mensual desde 2017-01 hasta el mes de corte elegido.
  2. Cada configuración de cada familia de modelos se evalúa con *rolling
     origin*: se repite una ventana de test de 3 meses en varios orígenes
     hacia atrás y se promedian los errores. Con un solo test de 3 puntos la
     elección del ganador la decidiría el ruido.
  3. Gana el menor MAPE promedio. Si la serie tiene ceros el MAPE no es
     confiable y se desempata por MASE.
  4. El ganador se reentrena con TODA la serie y proyecta hasta diciembre.
  5. Los intervalos salen de los errores relativos observados en el backtest,
     así son comparables entre familias (un intervalo analítico de SARIMA no
     es comparable con uno de XGBoost).

Un naive estacional corre siempre como referencia: si ningún modelo le gana,
hay que decirlo en vez de presentar un ganador que no aporta nada.
"""

from __future__ import annotations

import logging
import warnings
from dataclasses import dataclass, field
from itertools import product
from typing import Callable

import numpy as np
import pandas as pd

warnings.filterwarnings("ignore")
logging.getLogger("prophet").setLevel(logging.ERROR)
logging.getLogger("cmdstanpy").setLevel(logging.ERROR)

VENTANA_TEST = 3   # meses por ventana de evaluación
N_ORIGENES = 6     # cuántas ventanas hacia atrás
ESTACION = 12

# Versión del algoritmo: forma parte de la clave del repositorio de proyecciones,
# así un cambio de método invalida lo cacheado con el método anterior.
#   3: rangos simétricos estimados con 24 ventanas del modelo ganador (antes:
#      cuantiles con signo de las 6 ventanas de la selección). En la validación
#      walk-forward de 25 meses (analisis/calibracion_bandas.py) el rango del
#      80% original cubrió el 52% de los casos; el nuevo, el 80%.
VERSION = 3
K_BANDAS = 24      # ventanas con que se estiman los rangos del ganador


# ─────────────────────────── Métricas ───────────────────────────

def mape(y, yhat) -> float:
    y, yhat = np.asarray(y, float), np.asarray(yhat, float)
    ok = y != 0
    if not ok.any():
        return float("inf")
    return float(np.mean(np.abs((y[ok] - yhat[ok]) / y[ok])) * 100)


def smape(y, yhat) -> float:
    y, yhat = np.asarray(y, float), np.asarray(yhat, float)
    den = (np.abs(y) + np.abs(yhat)) / 2
    ok = den != 0
    if not ok.any():
        return float("inf")
    return float(np.mean(np.abs(y[ok] - yhat[ok]) / den[ok]) * 100)


def mase(y, yhat, y_train) -> float:
    """Error relativo al naive estacional sobre el train. <1 = mejor que naive."""
    y, yhat, y_train = np.asarray(y, float), np.asarray(yhat, float), np.asarray(y_train, float)
    if len(y_train) <= ESTACION:
        return float("inf")
    escala = np.mean(np.abs(y_train[ESTACION:] - y_train[:-ESTACION]))
    if escala == 0:
        return float("inf")
    return float(np.mean(np.abs(y - yhat)) / escala)


def rmse(y, yhat) -> float:
    return float(np.sqrt(np.mean((np.asarray(y, float) - np.asarray(yhat, float)) ** 2)))


# ─────────────────────────── Modelos ───────────────────────────
# Cada familia expone fit_predict(y, h, **params) -> np.ndarray de h valores.
# Reciben una Serie de pandas con índice mensual y devuelven el pronóstico.


def _naive_estacional(y: pd.Series, h: int, **_) -> np.ndarray:
    """Repite los últimos 12 meses. La referencia a batir."""
    ult = y.values[-ESTACION:]
    return np.array([ult[i % ESTACION] for i in range(h)], float)


def _sarima(y: pd.Series, h: int, order=(1, 1, 1), seasonal_order=(0, 1, 1, 12), **_) -> np.ndarray:
    from statsmodels.tsa.statespace.sarimax import SARIMAX

    m = SARIMAX(
        y.astype(float),
        order=order,
        seasonal_order=seasonal_order,
        enforce_stationarity=False,
        enforce_invertibility=False,
    ).fit(disp=False, maxiter=200)
    return np.asarray(m.forecast(h), float)


def _ets(y: pd.Series, h: int, trend="add", seasonal="add", damped=False, **_) -> np.ndarray:
    from statsmodels.tsa.holtwinters import ExponentialSmoothing

    m = ExponentialSmoothing(
        y.astype(float),
        trend=trend,
        seasonal=seasonal,
        seasonal_periods=ESTACION,
        damped_trend=damped if trend else False,
        initialization_method="estimated",
    ).fit(optimized=True)
    return np.asarray(m.forecast(h), float)


def _theta(y: pd.Series, h: int, deseasonalize=True, **_) -> np.ndarray:
    from statsmodels.tsa.forecasting.theta import ThetaModel

    m = ThetaModel(y.astype(float), period=ESTACION, deseasonalize=deseasonalize).fit()
    return np.asarray(m.forecast(h), float)


def _prophet(
    y: pd.Series,
    h: int,
    changepoint_prior_scale=0.05,
    seasonality_mode="additive",
    **_,
) -> np.ndarray:
    from prophet import Prophet

    df = pd.DataFrame({"ds": y.index.to_timestamp(), "y": y.values.astype(float)})
    m = Prophet(
        changepoint_prior_scale=changepoint_prior_scale,
        seasonality_mode=seasonality_mode,
        yearly_seasonality=True,
        weekly_seasonality=False,
        daily_seasonality=False,
        interval_width=0.8,
    )
    m.fit(df)
    fut = m.make_future_dataframe(periods=h, freq="MS")
    return np.asarray(m.predict(fut)["yhat"].values[-h:], float)


def _features_ml(y: pd.Series, n_lags: int = 12) -> tuple[pd.DataFrame, pd.Series]:
    """Lags + estacionalidad de calendario + tendencia lineal."""
    df = pd.DataFrame({"y": y.values.astype(float)}, index=y.index)
    for L in range(1, n_lags + 1):
        df[f"lag{L}"] = df["y"].shift(L)
    df["mes"] = [p.month for p in y.index]
    df["t"] = np.arange(len(df))
    df = df.dropna()
    return df.drop(columns=["y"]), df["y"]


def _ml_recursivo(modelo, y: pd.Series, h: int, n_lags: int = 12) -> np.ndarray:
    """Pronóstico recursivo: cada mes predicho se realimenta como lag."""
    X, target = _features_ml(y, n_lags)
    modelo.fit(X, target)

    hist = list(y.values.astype(float))
    out, t0 = [], len(y)
    ultimo_mes = y.index[-1]
    for i in range(h):
        p = ultimo_mes + i + 1  # Period aritmético mensual
        fila = {f"lag{L}": hist[-L] for L in range(1, n_lags + 1)}
        fila["mes"] = p.month
        fila["t"] = t0 + i
        pred = float(modelo.predict(pd.DataFrame([fila])[X.columns])[0])
        out.append(pred)
        hist.append(pred)
    return np.array(out, float)


def _xgboost(y: pd.Series, h: int, max_depth=3, n_estimators=300, learning_rate=0.05, **_):
    from xgboost import XGBRegressor

    return _ml_recursivo(
        XGBRegressor(
            max_depth=max_depth,
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            subsample=0.9,
            colsample_bytree=0.9,
            reg_lambda=1.0,
            n_jobs=1,          # el paralelismo va por combinación, no dentro
            verbosity=0,
        ),
        y,
        h,
    )


@dataclass
class Familia:
    nombre: str
    fn: Callable
    grid: list[dict] = field(default_factory=list)


def _grid(**kw) -> list[dict]:
    claves = list(kw)
    return [dict(zip(claves, combo)) for combo in product(*kw.values())]


def familias() -> list[Familia]:
    return [
        Familia("SARIMA", _sarima, _grid(
            order=[(1, 1, 1), (2, 1, 1), (1, 1, 2), (0, 1, 1), (2, 1, 2)],
            seasonal_order=[(0, 1, 1, 12), (1, 1, 0, 12), (1, 1, 1, 12)],
        )),
        Familia("ETS", _ets, _grid(
            trend=["add", None],
            seasonal=["add", "mul"],
            damped=[True, False],
        )),
        Familia("Theta", _theta, _grid(deseasonalize=[True, False])),
        Familia("Prophet", _prophet, _grid(
            changepoint_prior_scale=[0.01, 0.05, 0.2, 0.5],
            seasonality_mode=["additive", "multiplicative"],
        )),
        Familia("XGBoost", _xgboost, _grid(
            max_depth=[2, 3, 4],
            n_estimators=[200, 400],
            learning_rate=[0.03, 0.1],
        )),
        Familia("Naive estacional", _naive_estacional, [{}]),
    ]


# ─────────────────────────── Backtesting ───────────────────────────

def origenes(n: int, ventana: int, k: int) -> list[int]:
    """
    Índices de corte para rolling origin, del más reciente al más antiguo.
    Cada corte deja `ventana` meses como test y exige al menos 3 años de train.
    """
    minimo_train = ESTACION * 3
    cortes = []
    for i in range(k):
        corte = n - ventana - i
        if corte < minimo_train:
            break
        cortes.append(corte)
    return cortes


def evaluar(fn, params: dict, y: pd.Series, cortes: list[int], ventana: int) -> dict:
    """Corre una configuración en todos los orígenes y agrega sus errores."""
    mapes, mases, rmses, smapes, errores_rel = [], [], [], [], []

    for corte in cortes:
        train, test = y.iloc[:corte], y.iloc[corte : corte + ventana]
        if len(test) < ventana:
            continue
        try:
            pred = fn(train, ventana, **params)
        except Exception:
            return {"ok": False}
        if pred is None or len(pred) != ventana or not np.all(np.isfinite(pred)):
            return {"ok": False}

        pred = np.clip(pred, 0, None)  # un conteo de defunciones no es negativo
        mapes.append(mape(test.values, pred))
        smapes.append(smape(test.values, pred))
        mases.append(mase(test.values, pred, train.values))
        rmses.append(rmse(test.values, pred))
        with np.errstate(divide="ignore", invalid="ignore"):
            rel = (pred - test.values) / np.where(test.values == 0, np.nan, test.values)
        errores_rel.extend(rel[np.isfinite(rel)].tolist())

    if not mapes:
        return {"ok": False}

    return {
        "ok": True,
        "mape": float(np.mean(mapes)),
        "smape": float(np.mean(smapes)),
        "mase": float(np.mean(mases)),
        "rmse": float(np.mean(rmses)),
        "n_origenes": len(mapes),
        "errores_rel": errores_rel,
    }


# ─────────────────────────── Orquestación ───────────────────────────

def _bandas(errores_rel: list[float], central: np.ndarray) -> dict:
    """
    Intervalos empíricos a partir de los errores relativos del backtest.

    Se usa el error observado y no el intervalo analítico de cada modelo porque
    solo así las bandas son comparables entre familias: XGBoost no tiene uno y
    el de SARIMA asume supuestos que Prophet no cumple.
    """
    if len(errores_rel) < 8:
        return {}
    # Rango simétrico en escala logarítmica: se usa el TAMAÑO del error, no su
    # signo. Con el signo, un periodo en que el modelo se quedó corto (2021-2024,
    # homicidios al alza) corría todo el rango hacia arriba justo cuando la serie
    # se aplanó, y el valor real caía por debajo del límite inferior.
    tam = np.abs(np.log1p(np.asarray(errores_rel, dtype=float)))
    def rango(cobertura: float) -> tuple[list, list]:
        a = float(np.quantile(tam, cobertura))
        return (np.clip(central * np.exp(-a), 0, None).round(1).tolist(),
                (central * np.exp(a)).round(1).tolist())
    p10, p90 = rango(0.80)
    p25, p75 = rango(0.50)
    return {"p10": p10, "p90": p90, "p25": p25, "p75": p75}


def proyectar(
    serie: pd.Series,
    h: int,
    n_jobs: int = 8,
    ventana: int = VENTANA_TEST,
    k_origenes: int = N_ORIGENES,
    progreso: Callable[[dict], None] | None = None,
) -> dict:
    """
    Corre el grid completo con rolling origin, elige el mejor y proyecta h meses.

    `progreso` recibe el estado parcial para que la UI pueda mostrar en qué va.
    """
    from joblib import Parallel, delayed

    n = len(serie)
    cortes = origenes(n, ventana, k_origenes)
    if not cortes:
        raise ValueError(
            f"Serie demasiado corta: {n} meses. Se necesitan al menos "
            f"{ESTACION * 3 + ventana} para validar."
        )

    fams = familias()
    tareas = [(f.nombre, f.fn, p) for f in fams for p in f.grid]
    total = len(tareas)

    def avisar(**kw):
        if progreso:
            progreso(kw)

    avisar(
        fase="evaluando",
        total=total,
        completadas=0,
        mensaje=f"Evaluando {total} configuraciones en {len(cortes)} ventanas de {ventana} meses",
    )

    # El paralelismo va por configuración: cada tarea es un fit corto y
    # secuencial, así se aprovechan los núcleos sin sobresuscribir la máquina.
    resultados = Parallel(n_jobs=n_jobs, backend="loky", batch_size=1)(
        delayed(evaluar)(fn, params, serie, cortes, ventana) for _, fn, params in tareas
    )

    evaluadas = []
    for (nombre, _, params), r in zip(tareas, resultados):
        if r.get("ok"):
            evaluadas.append({"modelo": nombre, "params": params, **r})

    if not evaluadas:
        raise RuntimeError("Ninguna configuración pudo ajustarse a esta serie")

    avisar(fase="seleccionando", total=total, completadas=total,
           mensaje=f"{len(evaluadas)} de {total} configuraciones ajustaron correctamente")

    # Si la serie tiene ceros el MAPE deja de ser confiable y manda el MASE.
    hay_ceros = bool((serie.values == 0).any())
    criterio = "mase" if hay_ceros else "mape"
    evaluadas.sort(key=lambda r: (r[criterio], r["mase"]))
    mejor = evaluadas[0]

    naive = next((r for r in evaluadas if r["modelo"] == "Naive estacional"), None)

    # Reentrenamiento con toda la serie y proyección.
    avisar(fase="proyectando", total=total, completadas=total,
           mensaje=f"Reentrenando {mejor['modelo']} con la serie completa y proyectando {h} meses")

    fn_mejor = next(f.fn for f in fams if f.nombre == mejor["modelo"])
    central = np.clip(fn_mejor(serie, h, **mejor["params"]), 0, None)

    # Los rangos se estiman con más ventanas que la selección: con solo 6
    # ventanas solapadas el rango refleja el sesgo de los últimos meses y no la
    # distribución del error, y queda demasiado angosto.
    errores_bandas = mejor["errores_rel"]
    cortes_b = origenes(n, ventana, max(K_BANDAS, k_origenes))
    if len(cortes_b) > len(cortes):
        eb = evaluar(fn_mejor, mejor["params"], serie, cortes_b, ventana)
        if eb.get("ok"):
            errores_bandas = eb["errores_rel"]

    ultimo = serie.index[-1]
    meses = [str(ultimo + i + 1) for i in range(h)]

    ranking = [
        {
            "modelo": r["modelo"],
            "params": {k: str(v) for k, v in r["params"].items()},
            "mape": round(r["mape"], 2),
            "smape": round(r["smape"], 2),
            "mase": round(r["mase"], 3),
            "rmse": round(r["rmse"], 1),
        }
        for r in evaluadas[:15]
    ]

    # Mejor configuración de cada familia, para comparar familias y no solo runs.
    por_familia = {}
    for r in evaluadas:
        if r["modelo"] not in por_familia:
            por_familia[r["modelo"]] = {
                "modelo": r["modelo"],
                "mape": round(r["mape"], 2),
                "mase": round(r["mase"], 3),
                "rmse": round(r["rmse"], 1),
                "params": {k: str(v) for k, v in r["params"].items()},
            }

    return {
        "meses": meses,
        "central": central.round(1).tolist(),
        "bandas": _bandas(errores_bandas, central),
        "algo_v": VERSION,
        # Errores relativos (pred - real) / real con que se arma el rango; se
        # devuelven para poder auditar la calibración sin repetir el cálculo.
        "errores_bandas": [round(e, 4) for e in errores_bandas],
        "mejor": {
            "modelo": mejor["modelo"],
            "params": {k: str(v) for k, v in mejor["params"].items()},
            "mape": round(mejor["mape"], 2),
            "smape": round(mejor["smape"], 2),
            "mase": round(mejor["mase"], 3),
            "rmse": round(mejor["rmse"], 1),
        },
        "naive": (
            {"mape": round(naive["mape"], 2), "mase": round(naive["mase"], 3)} if naive else None
        ),
        "gana_al_naive": bool(naive and mejor["mase"] < naive["mase"]),
        "criterio": criterio,
        "hay_ceros": hay_ceros,
        "ranking": ranking,
        "por_familia": sorted(por_familia.values(), key=lambda r: r[criterio]),
        "validacion": {
            "ventana_meses": ventana,
            "n_origenes": len(cortes),
            "meses_evaluados": len(cortes) * ventana,
            "primer_origen": str(serie.index[cortes[-1]]),
            "ultimo_origen": str(serie.index[cortes[0]]),
        },
        "configuraciones_probadas": total,
        "configuraciones_validas": len(evaluadas),
    }