"""
Modelo v2: todas las variables de SINADEF + ingeniería de atributos, con
medición de importancia y selección.

Sobre los códigos CIE-10: se tratan aparte del resto. Un código X85–Y09
(agresiones) lo asigna un codificador que ya decidió que la muerte fue un
homicidio, así que usarlo para predecir homicidio es en parte circular. Se
entrenan dos modelos —con y sin CIE— para medir cuánto del desempeño viene de
esa vía y cuánto de la información clínica genuina.
"""
from __future__ import annotations

import re
import time
from pathlib import Path

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

BASE = Path(__file__).resolve().parent
DB = BASE.parent / "sinadef.duckdb"
CAUSAS = [f"DEBIDO_CAUSA_{c}" for c in "ABCDEF"]
CIEX = [f"CAUSA_{c}_CIEX" for c in "ABCDEF"]


def cargar():
    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, ETNIA, PAIS_DOMICILIO,
                   DEPARTAMENTO_FALLECIMIENTO, PROVINCIA_FALLECIMIENTO, DISTRITO_FALLECIMIENTO,
                   DEPARTAMENTO_DOMICILIO, PROVINCIA_DOMICILIO, DISTRITO_DOMICILIO,
                   {', '.join(CAUSAS)}, {', '.join(CIEX)}
            FROM tbl_SINADEF WHERE MUERTE_VIOLENTA <> 'SIN REGISTRO'
        """).fetchdf()
    finally:
        c.close()
    return df


def limpiar(s):
    return re.sub(r"[^A-ZÁÉÍÓÚÑ ]", " ", str(s).upper().replace("SIN REGISTRO", " ")).strip()


def preparar(df):
    # La causa A es la causa básica de muerte; las siguientes son antecedentes.
    # Se separa para que el modelo pueda ponderarlas distinto.
    df["txt_a"] = df["DEBIDO_CAUSA_A"].map(limpiar)
    df["txt_resto"] = df[CAUSAS[1:]].fillna("").agg(" ".join, axis=1).map(limpiar)

    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)

    f = pd.to_datetime(df["FECHA"], errors="coerce")
    df["dow"] = f.dt.dayofweek.fillna(-1)          # el domingo concentra +22% de homicidios
    df["finde"] = (df["dow"] >= 5).astype(int)
    df["mes_num"] = pd.to_numeric(df["MES"], errors="coerce").fillna(-1)

    # Murió lejos de casa: un homicidio en vía pública lejos del domicilio tiene
    # un perfil distinto al de una muerte en el propio distrito.
    df["fuera_distrito"] = (df["DISTRITO_FALLECIMIENTO"] != df["DISTRITO_DOMICILIO"]).astype(int)
    df["fuera_depto"] = (df["DEPARTAMENTO_FALLECIMIENTO"] != df["DEPARTAMENTO_DOMICILIO"]).astype(int)
    df["extranjero"] = (~df["PAIS_DOMICILIO"].isin(["PERU", "SIN REGISTRO"])).astype(int)

    # Cuán completo está el certificado: un documento parco puede indicar
    # tanto un caso simple como uno que nadie quiso detallar.
    df["n_causas"] = (df[CAUSAS].notna() & (df[CAUSAS] != "SIN REGISTRO")).sum(axis=1)
    df["n_ciex"] = (df[CIEX].notna() & (df[CIEX] != "SIN REGISTRO")).sum(axis=1)
    df["largo_txt"] = (df["txt_a"] + " " + df["txt_resto"]).str.len()

    # Primer carácter del CIE-10: el capítulo (X, Y, S, T…) es lo informativo.
    for col in CIEX:
        df[col + "_cap"] = df[col].fillna("").str[:1].replace("", "NA")
    return df


CAT_BASE = ["NECROPSIA", "TIPO_LUGAR", "SEXO", "INSTITUCION", "TIPO_SEGURO", "TIPO_CDEF",
            "ESTADO_CIVIL", "NIVEL_DE_INSTRUCCION", "ETNIA",
            "DEPARTAMENTO_FALLECIMIENTO", "PROVINCIA_FALLECIMIENTO"]
NUM_BASE = ["edad_num", "dow", "finde", "mes_num", "fuera_distrito", "fuera_depto",
            "extranjero", "n_causas", "largo_txt"]
CAT_CIEX = [c + "_cap" for c in CIEX]


def pipe(usar_ciex: bool):
    cat = CAT_BASE + (CAT_CIEX if usar_ciex else [])
    num = NUM_BASE + (["n_ciex"] if usar_ciex else [])
    pre = ColumnTransformer([
        ("txt_a", TfidfVectorizer(max_features=2500, ngram_range=(1, 2), min_df=5,
                                  sublinear_tf=True), "txt_a"),
        ("txt_r", TfidfVectorizer(max_features=2500, ngram_range=(1, 2), min_df=5,
                                  sublinear_tf=True), "txt_resto"),
        ("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)]), cat, num


def main():
    df = preparar(cargar())
    lab = df[df["MUERTE_VIOLENTA"] != "NO SE CONOCE"].copy()
    lab["y"] = (lab["MUERTE_VIOLENTA"] == "HOMICIDIO").astype(int)
    lab["periodo"] = lab["ANIO"].astype(str) + "-" + lab["MES"].astype(str)
    tr, te = lab[lab["periodo"] < "2025-01"], lab[lab["periodo"] >= "2025-01"]
    print(f"train {len(tr):,}  ·  test {len(te):,}\n")

    print(f"{'variante':34}{'AUC':>9}{'AP':>9}{'err. conteo':>13}")
    print("─" * 66)
    res = {}
    for nombre, ciex in [("v1: sin CIE-10 (solo clínico)", False),
                         ("v2: con CIE-10", True)]:
        p_, cat, num = pipe(ciex)
        cols = ["txt_a", "txt_resto"] + cat + num
        t0 = time.time()
        p_.fit(tr[cols], tr["y"])
        pr = p_.predict_proba(te[cols])[:, 1]
        auc, ap = roc_auc_score(te["y"], pr), average_precision_score(te["y"], pr)
        err = (pr.sum() / te["y"].sum() - 1) * 100
        res[nombre] = (p_, cols, auc)
        print(f"{nombre:34}{auc:9.4f}{ap:9.4f}{err:+12.1f}%   ({time.time()-t0:.0f}s)")

    # Importancia por ganancia, agrupada por variable de origen: las categóricas
    # se abren en muchas columnas y hay que sumarlas para comparar de igual a igual.
    print("\nIMPORTANCIA (ganancia agregada por variable de origen) — modelo con CIE-10")
    p_, cols, _ = res["v2: con CIE-10"]
    pre, clf = p_.named_steps["pre"], p_.named_steps["clf"]
    nombres = pre.get_feature_names_out()
    imp = clf.booster_.feature_importance(importance_type="gain")
    agg = {}
    for n, g in zip(nombres, imp):
        if n.startswith("txt_a__"):       k = "TEXTO causa A (básica)"
        elif n.startswith("txt_r__"):     k = "TEXTO causas B-F"
        elif n.startswith("cat__"):       k = n.split("__")[1].rsplit("_", 1)[0] if "_cap" in n else n.split("__")[1].split("_")[0]
        else:                             k = n.split("__")[-1]
        agg[k] = agg.get(k, 0) + g
    tot = sum(agg.values())
    for k, v in sorted(agg.items(), key=lambda x: -x[1])[:18]:
        print(f"  {k:34}{v/tot*100:7.2f}%")


if __name__ == "__main__":
    main()