"""
Clasificador de tipo de muerte violenta — entrenamiento y persistencia.

Entrena sobre las muertes violentas YA tipificadas para distinguir HOMICIDIO del
resto, valida con dos esquemas complementarios, y guarda el modelo entrenado
para reutilizarlo sin volver a ajustarlo.

Dos validaciones, porque miden cosas distintas:

  * **CV estratificada (5 pliegues)**: mide si el patrón es estable a lo largo de
    toda la serie. Responde "¿existe un patrón aprendible?".
  * **Validación temporal (rolling origin)**: entrena solo con el pasado y evalúa
    en el futuro. Responde "¿sirve para clasificar casos nuevos?". Es la que
    importa para aplicar el modelo a casos sin tipificar, porque las prácticas de
    registro cambian con los años y una CV aleatoria las mezcla.

El desempeño clave no es el AUC sino el **error de conteo**: si la suma de
probabilidades calibradas reproduce el número real de homicidios en un tramo que
el modelo no vio, entonces esa suma sirve para estimar cuántos hay entre los
casos sin tipificar.

Uso:
    python entrenar_modelo.py            # entrena, valida y guarda
    python entrenar_modelo.py --aplicar  # carga el modelo guardado y puntúa
"""
from __future__ import annotations

import argparse
import json
import re
import sys
import time
from pathlib import Path

import duckdb
import joblib
import numpy as np
import pandas as pd
from lightgbm import LGBMClassifier
from sklearn.calibration import CalibratedClassifierCV
from sklearn.compose import ColumnTransformer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import average_precision_score, brier_score_loss, roc_auc_score
from sklearn.model_selection import StratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder

BASE = Path(__file__).resolve().parent
DB = BASE.parent / "sinadef.duckdb"
MODELO = BASE / "modelos" / "clasificador_homicidio.joblib"
META = BASE / "modelos" / "metadata.json"

CAUSAS = [f"DEBIDO_CAUSA_{c}" for c in "ABCDEF"]
CAT = ["NECROPSIA", "TIPO_LUGAR", "SEXO", "INSTITUCION", "TIPO_SEGURO",
       "TIPO_CDEF", "DEPARTAMENTO_FALLECIMIENTO", "ESTADO_CIVIL",
       "NIVEL_DE_INSTRUCCION"]
NUM = ["edad_num"]
FEATS = ["texto"] + CAT + NUM


def cargar() -> pd.DataFrame:
    c = duckdb.connect(str(DB), read_only=True)
    try:
        df = c.execute(f"""
            SELECT FECHA, ANIO, MES, MUERTE_VIOLENTA, NECROPSIA, TIPO_LUGAR, SEXO,
                   EDAD, TIEMPO_EDAD, INSTITUCION, TIPO_SEGURO, TIPO_CDEF,
                   ESTADO_CIVIL, NIVEL_DE_INSTRUCCION,
                   DEPARTAMENTO_FALLECIMIENTO, PROVINCIA_FALLECIMIENTO,
                   DISTRITO_FALLECIMIENTO, {', '.join(CAUSAS)}
            FROM tbl_SINADEF WHERE MUERTE_VIOLENTA <> 'SIN REGISTRO'
        """).fetchdf()
    finally:
        c.close()

    # El texto de las seis causas es la señal más rica del certificado. Se quita
    # "SIN REGISTRO", que es relleno y no información.
    txt = df[CAUSAS].fillna("").agg(" ".join, axis=1).str.upper()
    txt = txt.str.replace("SIN REGISTRO", " ", regex=False)
    df["texto"] = txt.apply(lambda s: re.sub(r"[^A-ZÁÉÍÓÚÑ ]", " ", s).strip())

    df["edad_num"] = pd.to_numeric(df["EDAD"], errors="coerce")
    df.loc[df["TIEMPO_EDAD"] != "AÑOS", "edad_num"] = 0
    df["edad_num"] = df["edad_num"].fillna(-1)
    df["periodo"] = df["ANIO"].astype(str) + "-" + df["MES"].astype(str)
    return df


def construir() -> Pipeline:
    pre = ColumnTransformer([
        ("txt", TfidfVectorizer(max_features=4000, ngram_range=(1, 2), min_df=5,
                                sublinear_tf=True), "texto"),
        ("cat", OneHotEncoder(handle_unknown="ignore", min_frequency=20), CAT),
        ("num", "passthrough", NUM),
    ])
    clf = LGBMClassifier(n_estimators=500, learning_rate=0.05, num_leaves=63,
                         min_child_samples=30, subsample=0.9, subsample_freq=1,
                         colsample_bytree=0.8, reg_lambda=1.0,
                         verbosity=-1, n_jobs=8, random_state=42)
    return Pipeline([("pre", pre), ("clf", clf)])


def metricas(y, p) -> dict:
    return {
        "auc": float(roc_auc_score(y, p)),
        "ap": float(average_precision_score(y, p)),
        "brier": float(brier_score_loss(y, p)),
        "esperados": float(p.sum()),
        "reales": int(y.sum()),
        "error_conteo_pct": float((p.sum() / y.sum() - 1) * 100) if y.sum() else np.nan,
    }


def entrenar():
    t0 = time.time()
    df = cargar()
    lab = df[df["MUERTE_VIOLENTA"] != "NO SE CONOCE"].copy()
    lab["y"] = (lab["MUERTE_VIOLENTA"] == "HOMICIDIO").astype(int)
    print(f"Muertes violentas tipificadas : {len(lab):,}")
    print(f"  de las cuales HOMICIDIO     : {lab['y'].sum():,} ({lab['y'].mean()*100:.1f}%)")
    print(f"Sin tipificar (NO SE CONOCE)  : {(df['MUERTE_VIOLENTA']=='NO SE CONOCE').sum():,}\n")

    X, y = lab[FEATS], lab["y"].to_numpy()

    print("── VALIDACIÓN CRUZADA ESTRATIFICADA (5 pliegues) ──")
    print("   ¿existe un patrón aprendible y estable?\n")
    print(f"   {'pliegue':9}{'AUC':>8}{'AP':>8}{'Brier':>9}{'esperados':>11}{'reales':>8}{'error':>9}")
    cv, filas = StratifiedKFold(5, shuffle=True, random_state=42), []
    for i, (itr, ite) in enumerate(cv.split(X, y), 1):
        m = construir().fit(X.iloc[itr], y[itr])
        p = m.predict_proba(X.iloc[ite])[:, 1]
        r = metricas(y[ite], p)
        filas.append(r)
        print(f"   {i:<9}{r['auc']:8.4f}{r['ap']:8.4f}{r['brier']:9.4f}"
              f"{r['esperados']:11.0f}{r['reales']:8d}{r['error_conteo_pct']:+8.1f}%")
    med = {k: float(np.mean([f[k] for f in filas])) for k in ("auc", "ap", "brier")}
    sd = {k: float(np.std([f[k] for f in filas])) for k in ("auc", "ap", "brier")}
    print(f"\n   media  AUC {med['auc']:.4f} ±{sd['auc']:.4f} · "
          f"AP {med['ap']:.4f} ±{sd['ap']:.4f} · Brier {med['brier']:.4f}")

    print("\n── VALIDACIÓN TEMPORAL (rolling origin) ──")
    print("   ¿sirve para clasificar casos nuevos, entrenando solo con el pasado?\n")
    print(f"   {'entrena hasta':15}{'test':10}{'AUC':>8}{'esperados':>11}{'reales':>8}{'error':>9}")
    temporales = []
    for corte in ("2022-01", "2023-01", "2024-01", "2025-01", "2026-01"):
        tr, te = lab[lab["periodo"] < corte], lab[lab["periodo"] >= corte]
        if len(te) < 500:
            continue
        m = construir().fit(tr[FEATS], tr["y"])
        p = m.predict_proba(te[FEATS])[:, 1]
        r = metricas(te["y"].to_numpy(), p)
        r["corte"] = corte
        temporales.append(r)
        print(f"   {corte:15}{len(te):<10,}{r['auc']:8.4f}"
              f"{r['esperados']:11.0f}{r['reales']:8d}{r['error_conteo_pct']:+8.1f}%")

    print("\n── MODELO FINAL ──")
    print("   Se reentrena con TODO lo tipificado y se calibra por isotónica:")
    print("   sin calibrar, las probabilidades no se pueden sumar para estimar conteos.")
    final = CalibratedClassifierCV(construir(), method="isotonic", cv=3)
    final.fit(X, y)
    MODELO.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump({"modelo": final, "feats": FEATS, "cat": CAT, "num": NUM,
                 "causas": CAUSAS}, MODELO, compress=3)

    meta = {
        "entrenado": time.strftime("%Y-%m-%d %H:%M:%S"),
        "n_entrenamiento": int(len(lab)),
        "n_homicidios": int(lab["y"].sum()),
        "n_sin_tipificar": int((df["MUERTE_VIOLENTA"] == "NO SE CONOCE").sum()),
        "cv_media": med, "cv_desv": sd,
        "temporal": temporales,
        "segundos": round(time.time() - t0, 1),
    }
    META.write_text(json.dumps(meta, indent=2, ensure_ascii=False))
    print(f"\n   modelo  → {MODELO}  ({MODELO.stat().st_size/1e6:.1f} MB)")
    print(f"   metadata → {META}")
    print(f"   tiempo total: {meta['segundos']}s")


def aplicar():
    if not MODELO.exists():
        sys.exit(f"No existe {MODELO}. Corre primero sin --aplicar.")
    bundle = joblib.load(MODELO)
    df = cargar()
    nsc = df[df["MUERTE_VIOLENTA"] == "NO SE CONOCE"].copy()
    nsc["p_hom"] = bundle["modelo"].predict_proba(nsc[bundle["feats"]])[:, 1]
    salida = BASE / "nsc_scored.parquet"
    nsc.drop(columns=CAUSAS).to_parquet(salida, index=False)
    print(f"Puntuados {len(nsc):,} casos sin tipificar → {salida}")
    print(f"  suma de probabilidades (homicidios esperados): {nsc['p_hom'].sum():,.0f}")
    print(f"  con p>=0.5: {(nsc['p_hom']>=0.5).sum():,}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--aplicar", action="store_true",
                    help="carga el modelo guardado y puntúa los casos sin tipificar")
    a = ap.parse_args()
    aplicar() if a.aplicar else entrenar()