"""
Estimación del subregistro de homicidios en SINADEF.

Entrena un clasificador sobre las muertes violentas YA tipificadas (homicidio vs
resto) y lo aplica a las que quedaron como "NO SE CONOCE", para estimar cuántas
de ellas tienen perfil de homicidio.

Advertencia que acompaña a todo resultado de este script: el modelo aprende de
los casos que alguien decidió clasificar. Si un caso quedó sin tipificar porque
era ambiguo, pertenece a una población distinta de la de entrenamiento, y la
extrapolación deja de estar garantizada. Por eso el número que produce es un
escenario central, no una medición, y se reporta junto a las cotas sin supuestos.
"""
import re
import sys

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

DB = "/home/osalas/apps/sinadef/sinadef.duckdb"
CAUSAS = [f"DEBIDO_CAUSA_{c}" for c in "ABCDEF"]


def cargar():
    c = duckdb.connect(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,
                   DEPARTAMENTO_FALLECIMIENTO, {', '.join(CAUSAS)}
            FROM tbl_SINADEF
            WHERE MUERTE_VIOLENTA <> 'SIN REGISTRO'
        """).fetchdf()
    finally:
        c.close()
    return df


def preparar(df):
    # Texto de las seis causas concatenado: es la señal más rica del registro.
    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))
    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 main():
    df = preparar(cargar())
    etiquetados = df[df["MUERTE_VIOLENTA"] != "NO SE CONOCE"].copy()
    nsc = df[df["MUERTE_VIOLENTA"] == "NO SE CONOCE"].copy()
    etiquetados["y"] = (etiquetados["MUERTE_VIOLENTA"] == "HOMICIDIO").astype(int)

    print(f"etiquetados: {len(etiquetados):,}  (homicidios {etiquetados['y'].sum():,})")
    print(f"sin tipificar (NO SE CONOCE): {len(nsc):,}\n")

    CAT = ["NECROPSIA", "TIPO_LUGAR", "SEXO", "INSTITUCION", "TIPO_SEGURO",
           "TIPO_CDEF", "DEPARTAMENTO_FALLECIMIENTO"]
    NUM = ["edad_num"]

    pre = ColumnTransformer([
        ("txt", TfidfVectorizer(max_features=3000, ngram_range=(1, 2), min_df=5), "texto"),
        ("cat", OneHotEncoder(handle_unknown="ignore", min_frequency=20), CAT),
        ("num", "passthrough", NUM),
    ])

    from lightgbm import LGBMClassifier
    base = LGBMClassifier(n_estimators=400, learning_rate=0.06, num_leaves=63,
                          min_child_samples=30, subsample=0.9, colsample_bytree=0.8,
                          verbosity=-1, n_jobs=8)

    # Validación temporal: entrenar con el pasado y probar en el futuro, nunca
    # al azar. Un split aleatorio inflaría el desempeño porque las prácticas de
    # registro cambian con el tiempo.
    corte = "2025-01"
    tr = etiquetados[etiquetados["periodo"] < corte]
    te = etiquetados[etiquetados["periodo"] >= corte]
    print(f"VALIDACIÓN TEMPORAL  train<{corte}: {len(tr):,}   test: {len(te):,}")

    pipe = Pipeline([("pre", pre), ("clf", base)])
    pipe.fit(tr[["texto"] + CAT + NUM], tr["y"])
    p = pipe.predict_proba(te[["texto"] + CAT + NUM])[:, 1]
    print(f"  AUC  {roc_auc_score(te['y'], p):.4f}")
    print(f"  AP   {average_precision_score(te['y'], p):.4f}   (base {te['y'].mean():.3f})")
    print(f"  Brier {brier_score_loss(te['y'], p):.4f}")

    # Calibración: sin ella las probabilidades no se pueden sumar para estimar
    # un conteo esperado.
    cal = CalibratedClassifierCV(pipe, method="isotonic", cv=3)
    cal.fit(tr[["texto"] + CAT + NUM], tr["y"])
    pc = cal.predict_proba(te[["texto"] + CAT + NUM])[:, 1]
    print(f"  Brier calibrado {brier_score_loss(te['y'], pc):.4f}")
    print(f"  esperados vs reales en test: {pc.sum():.0f} vs {te['y'].sum()}"
          f"  (error {(pc.sum()/te['y'].sum()-1)*100:+.1f}%)\n")

    # Modelo final con todo lo etiquetado, aplicado a los no tipificados.
    final = CalibratedClassifierCV(Pipeline([("pre", pre), ("clf", base)]),
                                   method="isotonic", cv=3)
    final.fit(etiquetados[["texto"] + CAT + NUM], etiquetados["y"])
    nsc["p_hom"] = final.predict_proba(nsc[["texto"] + CAT + NUM])[:, 1]
    nsc[["FECHA", "ANIO", "DEPARTAMENTO_FALLECIMIENTO", "SEXO", "TIPO_LUGAR",
         "NECROPSIA", "texto", "p_hom"]].to_parquet(
        "/home/osalas/apps/sinadef/analisis/nsc_scored.parquet", index=False)

    print("ESTIMACIÓN POR AÑO (suma de probabilidades calibradas)")
    print(f"  {'año':6}{'homic.':>9}{'NSC':>7}{'NSC~hom':>9}{'total est.':>12}{'subregistro':>13}")
    c = duckdb.connect(DB, read_only=True)
    for a in sorted(nsc["ANIO"].unique()):
        reg = c.execute("SELECT COUNT(*) FROM tbl_SINADEF WHERE MUERTE_VIOLENTA='HOMICIDIO' AND ANIO=?",
                        [a]).fetchone()[0]
        sub = nsc[nsc["ANIO"] == a]
        esp = sub["p_hom"].sum()
        print(f"  {a:6}{reg:9,}{len(sub):7,}{esp:9.0f}{reg+esp:12,.0f}{esp/reg*100:12.0f}%")
    c.close()


if __name__ == "__main__":
    main()