"""
Proyección diaria de homicidios con corrección de rezago de registro.

El problema: SINADEF publica por fecha de ocurrencia, pero los certificados
llegan con retraso. Los últimos días siempre aparecen incompletos, así que un
modelo entrenado con ellos aprende una caída que no existe.

Enfoque, en tres piezas:

  1. **Modelo de conteo diario.** GLM binomial negativa con día de la semana,
     tendencia suave y estacionalidad anual (armónicos de Fourier). Es la
     familia estándar en vigilancia epidemiológica — el algoritmo de Farrington
     que usan los centros europeos de control de enfermedades es un GLM de esta
     misma forma. Binomial negativa y no Poisson porque los conteos están
     sobredispersos (varianza/media = 1.4).

  2. **Curva de rezago sin snapshots.** El modelo se ajusta SOLO con días ya
     consolidados. Luego se compara lo observado en los días recientes contra lo
     que el modelo esperaba: ese cociente, por antigüedad, es la fracción ya
     reportada. Es el mismo razonamiento del ajuste por truncamiento a la
     derecha que se usó para corregir el retraso de notificación en vigilancia
     de VIH (Brookmeyer & Damiano), pero estimando el rezago del propio ajuste
     en lugar de una fecha de reporte que aquí no existe.

  3. **Backtest honesto.** Origen móvil sobre los últimos 60 días, prediciendo
     el mismo horizonte que luego se proyecta, y evaluando también si los
     intervalos cubren lo que dicen cubrir.
"""

from __future__ import annotations

import datetime as dt
import warnings

import duckdb
import numpy as np
import pandas as pd
import statsmodels.api as sm

warnings.filterwarnings("ignore")

DB = "sinadef.duckdb"
DIAS = ["lunes", "martes", "miércoles", "jueves", "viernes", "sábado", "domingo"]

# Versión del método. Entra en la clave de caché: si cambia el algoritmo, los
# resultados guardados con la versión anterior dejan de servirse y se recalculan.
# Sin esto, una mejora del modelo quedaría enmascarada por el caché.
#   1 → imputación multiplicativa (obs / fracción)
#   2 → imputación aditiva μ·(1−f); un día sin registros ya no imputa cero
VERSION = 2

# Días finales que se consideran no consolidados y quedan fuera del ajuste.
# 21 días es holgado: la medición directa mostró que a ~7 días solo se ve el
# 79% del valor, y que a partir de ~3 semanas el crecimiento ya es marginal.
DIAS_NO_CONSOLIDADOS = 21


def serie_diaria(filtro_sql: str = "MUERTE_VIOLENTA='HOMICIDIO'") -> pd.DataFrame:
    c = duckdb.connect(DB, read_only=True)
    try:
        filas = c.execute(f"""
            SELECT FECHA AS f, COUNT(*) FILTER (WHERE {filtro_sql}) AS n
            FROM tbl_SINADEF WHERE FECHA >= '2017-01-01'
            GROUP BY 1 ORDER BY 1
        """).fetchall()
    finally:
        c.close()

    df = pd.DataFrame(filas, columns=["fecha", "n"])
    df["fecha"] = pd.to_datetime(df["fecha"])
    # Días sin ninguna defunción no aparecen como fila; se rellenan con 0.
    idx = pd.date_range(df["fecha"].min(), df["fecha"].max(), freq="D")
    df = df.set_index("fecha").reindex(idx, fill_value=0).rename_axis("fecha").reset_index()
    return df


def matriz_diseño(fechas: pd.Series, t0: pd.Timestamp, armonicos: int = 2) -> pd.DataFrame:
    """Día de la semana + tendencia + estacionalidad anual."""
    # El índice se descarta: si `fechas` llega como slice conserva el índice
    # original y pandas alinearía por él, dejando la matriz entera en NaN.
    fechas = pd.Series(pd.to_datetime(fechas).to_numpy())
    X = pd.DataFrame(index=range(len(fechas)))
    dow = fechas.dt.dayofweek

    # Lunes es la categoría de referencia; el resto entra como indicadoras.
    for k in range(1, 7):
        X[f"dow_{k}"] = (dow == k).astype(float)

    t = (fechas - t0).dt.days.to_numpy(dtype=float)
    X["t"] = t / 365.25                    # tendencia lineal, en años
    X["t2"] = (t / 365.25) ** 2            # curvatura: el nivel no crece recto

    # Estacionalidad anual como senos y cosenos: más estable que 11 indicadoras
    # de mes cuando parte del año tiene pocos datos.
    for k in range(1, armonicos + 1):
        X[f"sin{k}"] = np.sin(2 * np.pi * k * t / 365.25)
        X[f"cos{k}"] = np.cos(2 * np.pi * k * t / 365.25)

    return X


def ajustar(df: pd.DataFrame, t0: pd.Timestamp):
    """GLM binomial negativa. Alpha se estima por perfil de verosimilitud."""
    X = sm.add_constant(matriz_diseño(df["fecha"], t0), has_constant="add")
    y = df["n"].to_numpy(dtype=float)

    mejor, mejor_ll = None, -np.inf
    for alpha in (0.01, 0.03, 0.05, 0.08, 0.12, 0.2, 0.3):
        try:
            m = sm.GLM(y, X, family=sm.families.NegativeBinomial(alpha=alpha)).fit()
            if m.llf > mejor_ll:
                mejor, mejor_ll, mejor_alpha = m, m.llf, alpha
        except Exception:
            continue
    if mejor is None:
        raise RuntimeError("No se pudo ajustar el GLM")
    mejor.alpha_elegido = mejor_alpha
    return mejor


def predecir(modelo, fechas: pd.Series, t0: pd.Timestamp) -> np.ndarray:
    X = sm.add_constant(matriz_diseño(fechas, t0), has_constant="add")
    return np.asarray(modelo.predict(X), dtype=float)


def curva_rezago(df: pd.DataFrame, modelo, t0: pd.Timestamp, corte: pd.Timestamp) -> pd.DataFrame:
    """
    Fracción ya reportada según la antigüedad del día.

    Para cada día reciente se compara lo observado contra lo que el modelo
    —ajustado solo con días consolidados— dice que debería haber. El cociente
    estima qué proporción del total ya entró al sistema.
    """
    recientes = df[df["fecha"] > corte - pd.Timedelta(days=DIAS_NO_CONSOLIDADOS + 14)].copy()
    recientes["esperado"] = predecir(modelo, recientes["fecha"], t0)
    recientes["antiguedad"] = (corte - recientes["fecha"]).dt.days
    recientes = recientes[recientes["antiguedad"] >= 0]

    # Se suaviza con una ventana móvil: día a día el cociente es muy ruidoso.
    recientes = recientes.sort_values("antiguedad")
    obs = recientes["n"].to_numpy(dtype=float)
    esp = recientes["esperado"].to_numpy(dtype=float)
    frac = []
    for i in range(len(recientes)):
        lo, hi = max(0, i - 2), min(len(recientes), i + 3)
        s_esp = esp[lo:hi].sum()
        frac.append(obs[lo:hi].sum() / s_esp if s_esp > 0 else np.nan)

    recientes["fraccion_reportada"] = np.clip(frac, 0.01, 1.0)
    # La completitud no puede bajar cuando el dato envejece.
    vals = recientes["fraccion_reportada"].to_numpy()
    for i in range(len(vals) - 2, -1, -1):
        vals[i] = min(vals[i], 1.0)
    recientes["fraccion_reportada"] = vals
    return recientes[["fecha", "antiguedad", "n", "esperado", "fraccion_reportada"]]


def backtest(df: pd.DataFrame, t0: pd.Timestamp, corte: pd.Timestamp,
             horizonte: int = 5, n_origenes: int = 12, paso: int = 5,
             meses_ventana: int = 24) -> dict:
    """
    Origen móvil sobre los últimos días consolidados.

    Se evalúa solo contra días ya maduros: comparar una predicción contra un día
    incompleto mediría el rezago, no la calidad del modelo.

    Entrena con la misma ventana móvil que el modelo final; si el backtest usara
    toda la serie mediría un modelo distinto del que después proyecta, y sus
    métricas no aplicarían al resultado.
    """
    limite = corte - pd.Timedelta(days=DIAS_NO_CONSOLIDADOS)
    consolidado = df[df["fecha"] <= limite].reset_index(drop=True)

    errores, dentro, filas = [], [], []
    for i in range(n_origenes):
        fin_train = len(consolidado) - horizonte - i * paso
        if fin_train < 400:
            break
        train = consolidado.iloc[:fin_train]
        if meses_ventana:
            desde = train["fecha"].iloc[-1] - pd.DateOffset(months=meses_ventana)
            train = train[train["fecha"] >= desde]
        if len(train) < 200:
            continue
        test = consolidado.iloc[fin_train:fin_train + horizonte]
        try:
            m = ajustar(train, t0)
        except Exception:
            continue
        pred = predecir(m, test["fecha"], t0)
        real = test["n"].to_numpy(dtype=float)

        # Intervalo al 80% de la binomial negativa alrededor de la media.
        alpha = m.alpha_elegido
        var = pred + alpha * pred ** 2
        lo = np.maximum(0, pred - 1.2816 * np.sqrt(var))
        hi = pred + 1.2816 * np.sqrt(var)

        errores.append(np.abs(real - pred))
        dentro.append(((real >= lo) & (real <= hi)).astype(float))
        filas.append({
            "origen": str(train["fecha"].iloc[-1].date()),
            "real": int(real.sum()),
            "predicho": float(pred.sum()),
            "error_pct": float((pred.sum() - real.sum()) / real.sum() * 100) if real.sum() else np.nan,
        })

    e = np.concatenate(errores) if errores else np.array([np.nan])
    d = np.concatenate(dentro) if dentro else np.array([np.nan])
    tot_real = np.array([f["real"] for f in filas], dtype=float)
    tot_pred = np.array([f["predicho"] for f in filas], dtype=float)

    return {
        "n_origenes": len(filas),
        "mae_diario": float(np.mean(e)),
        "rmse_diario": float(np.sqrt(np.mean(e ** 2))),
        "cobertura_80": float(np.mean(d) * 100),
        "mape_ventana": float(np.mean(np.abs(tot_pred - tot_real) / tot_real) * 100) if len(tot_real) else np.nan,
        "sesgo_ventana": float(np.mean((tot_pred - tot_real) / tot_real) * 100) if len(tot_real) else np.nan,
        "detalle": filas,
    }


# ─────────────────────── Proyección de cierre de mes ───────────────────────

def curva_isotonica(rez: pd.DataFrame) -> dict[int, float]:
    """
    Fracción reportada por antigüedad, forzada a ser monótona.

    El cociente observado/esperado día a día es muy ruidoso; la regresión
    isotónica impone lo único que sabemos con certeza: un dato no puede
    "desregistrarse" al envejecer.
    """
    from sklearn.isotonic import IsotonicRegression

    r = rez[(rez["antiguedad"] >= 1) & (rez["antiguedad"] <= 35)].sort_values("antiguedad")
    if len(r) < 5:
        return {}
    ratio = (r["n"] / r["esperado"].clip(lower=0.1)).clip(0, 1.5)
    iso = IsotonicRegression(y_min=0.25, y_max=1.0, increasing=True)
    return dict(zip(r["antiguedad"].astype(int), iso.fit_transform(r["antiguedad"], ratio)))


def proyectar_mes(
    filtro_sql: str = "MUERTE_VIOLENTA='HOMICIDIO'",
    mes: str | None = None,
    meses_ventana: int = 24,
    horizonte_backtest: int = 5,
    n_origenes: int = 12,
    progreso=None,
) -> dict:
    """
    Proyecta el cierre de `mes` (formato 'YYYY-MM'; por defecto, el último con
    datos) combinando tres piezas: lo ya observado, el rezago pendiente de los
    días transcurridos y la predicción de los días que faltan.
    """
    def avisar(**kw):
        if progreso:
            progreso(kw)

    avisar(fase="datos", mensaje="Construyendo la serie diaria")
    df = serie_diaria(filtro_sql)
    if df["n"].sum() < 400:
        raise ValueError(
            f"La serie filtrada tiene solo {int(df['n'].sum())} casos en total. "
            "Se necesitan al menos 400 para ajustar un modelo diario fiable."
        )

    t0 = df["fecha"].min()
    objetivo = pd.Period(mes, freq="M") if mes else pd.Period(df["fecha"].max(), freq="M")
    ini_mes = pd.Timestamp(objetivo.start_time)
    fin_mes = pd.Timestamp(objetivo.end_time).normalize()

    # El corte es el día siguiente al último con datos: hasta ahí se observa.
    corte = df["fecha"].max() + pd.Timedelta(days=1)
    limite = corte - pd.Timedelta(days=DIAS_NO_CONSOLIDADOS)

    cons = df[df["fecha"] <= limite]
    if len(cons) < 400:
        raise ValueError("No hay suficientes días consolidados para entrenar")

    train = cons[cons["fecha"] >= cons["fecha"].iloc[-1] - pd.DateOffset(months=meses_ventana)]
    if train["n"].sum() < 200:
        raise ValueError(
            f"La ventana de {meses_ventana} meses solo tiene {int(train['n'].sum())} casos. "
            "Amplía la ventana o usa un filtro menos restrictivo."
        )

    avisar(fase="ajuste", mensaje=f"Ajustando GLM sobre {len(train)} días ({meses_ventana} meses)")
    modelo = ajustar(train, t0)
    alpha = modelo.alpha_elegido

    avisar(fase="backtest", mensaje=f"Validando con {n_origenes} orígenes móviles")
    bt = backtest(df, t0, corte, horizonte_backtest, n_origenes, meses_ventana=meses_ventana)

    avisar(fase="rezago", mensaje="Estimando la curva de rezago de registro")
    rez = curva_rezago(df, modelo, t0, corte)
    frac = curva_isotonica(rez)

    # --- Días del mes ya transcurridos
    obs = df[(df["fecha"] >= ini_mes) & (df["fecha"] < corte) & (df["fecha"] <= fin_mes)].copy()
    obs["antiguedad"] = (corte - obs["fecha"]).dt.days
    obs["fraccion"] = obs["antiguedad"].map(lambda a: min(max(frac.get(int(a), 1.0), 0.05), 1.0))
    obs["esperado"] = predecir(modelo, obs["fecha"], t0)

    # La imputación es aditiva, no un cociente observado/fracción. Dividir falla
    # justo donde más importa: un día recién ocurrido puede tener 0 registros, y
    # 0 dividido por cualquier fracción sigue siendo 0 — el día quedaría como si
    # no hubiera pasado nada. Lo que falta por registrar se estima con el modelo:
    # si se ha registrado la fracción f, lo pendiente tiene media μ·(1−f).
    obs["pendiente"] = obs["esperado"] * (1 - obs["fraccion"])
    obs["corregido"] = obs["n"] + obs["pendiente"]

    observado = int(obs["n"].sum())
    corregido = float(obs["corregido"].sum())

    # --- Días que faltan para cerrar el mes
    futuros = pd.date_range(max(corte, ini_mes), fin_mes, freq="D")
    if len(futuros):
        pred = predecir(modelo, pd.Series(futuros), t0)
        var_fut = float(np.sum(pred + alpha * pred**2))
    else:
        pred, var_fut = np.array([]), 0.0

    # Incertidumbre de lo imputado. Tiene dos fuentes: la variabilidad propia de
    # los casos pendientes (binomial negativa sobre μ·(1−f)) y el error de la
    # curva de rezago, que se estima con pocas observaciones. La segunda domina,
    # y se toma como un 30% de lo imputado para no maquillarla.
    pend = float(obs["pendiente"].sum())
    var_pend = float(np.sum(obs["pendiente"] + alpha * obs["pendiente"] ** 2))
    sd_rezago = float(np.sqrt(var_pend + (pend * 0.30) ** 2))
    total = corregido + float(pred.sum())
    sd = float(np.sqrt(var_fut + sd_rezago**2))

    # --- Mismo mes en años anteriores, para dar escala
    hist = []
    for a in range(objetivo.year - 5, objetivo.year):
        p = pd.Period(f"{a}-{objetivo.month:02d}", freq="M")
        v = df[(df["fecha"] >= p.start_time) & (df["fecha"] <= p.end_time)]["n"].sum()
        if v > 0:
            hist.append({"anio": str(a), "n": int(v)})

    dias = []
    for _, r in obs.iterrows():
        dias.append({
            "fecha": r["fecha"].strftime("%Y-%m-%d"),
            "dow": int(r["fecha"].dayofweek),
            "observado": int(r["n"]),
            "corregido": round(float(r["corregido"]), 2),
            "pendiente": round(float(r["pendiente"]), 2),
            "esperado": round(float(r["esperado"]), 2),
            "fraccion": round(float(r["fraccion"]), 3),
            "tipo": "observado",
        })
    for f, p in zip(futuros, pred):
        v = float(p + alpha * p**2)
        dias.append({
            "fecha": f.strftime("%Y-%m-%d"),
            "dow": int(f.dayofweek),
            "proyectado": round(float(p), 2),
            "lo": round(max(0.0, p - 1.2816 * np.sqrt(v)), 2),
            "hi": round(float(p + 1.2816 * np.sqrt(v)), 2),
            "tipo": "proyectado",
        })

    efectos = {}
    for k in range(1, 7):
        if f"dow_{k}" in modelo.params:
            efectos[DIAS[k]] = {
                "factor": round(float(np.exp(modelo.params[f"dow_{k}"])), 3),
                "pct": round(float((np.exp(modelo.params[f"dow_{k}"]) - 1) * 100), 1),
                "p": round(float(modelo.pvalues[f"dow_{k}"]), 4),
            }
    efectos[DIAS[0]] = {"factor": 1.0, "pct": 0.0, "p": None, "referencia": True}

    return {
        "mes": str(objetivo),
        "corte": corte.strftime("%Y-%m-%d"),
        "ultimo_dia_con_datos": df["fecha"].max().strftime("%Y-%m-%d"),
        "dias_transcurridos": len(obs),
        "dias_faltantes": len(futuros),
        "observado": observado,
        "rezago_pendiente": round(corregido - observado, 1),
        "proyectado_faltantes": round(float(pred.sum()), 1),
        "total": round(total),
        # Hacia afuera: redondear un intervalo "al más cercano" lo estrecha y le
        # quita parte de la cobertura que declara.
        "lo80": int(np.floor(total - 1.2816 * sd)),
        "hi80": int(np.ceil(total + 1.2816 * sd)),
        "dias": dias,
        "curva_rezago": [
            {"antiguedad": int(a), "fraccion": round(float(f), 3)}
            for a, f in sorted(frac.items()) if a <= 30
        ],
        "efectos_dow": efectos,
        "tendencia_anual_pct": round(float((np.exp(modelo.params["t"]) - 1) * 100), 1),
        "backtest": {k: (round(v, 2) if isinstance(v, float) else v)
                     for k, v in bt.items() if k != "detalle"},
        "modelo": {
            "familia": "Binomial negativa (log)",
            "alpha": alpha,
            "dias_entrenamiento": len(train),
            "ventana_meses": meses_ventana,
            "desde": train["fecha"].iloc[0].strftime("%Y-%m-%d"),
            "hasta": train["fecha"].iloc[-1].strftime("%Y-%m-%d"),
        },
        "historico_mismo_mes": hist,
    }