"""
Clasificador multiclase del tipo de muerte violenta.

A diferencia del binario (homicidio / no homicidio), asigna a cada caso sin
tipificar una probabilidad para cada uno de los cinco tipos que usa SINADEF.
Responde cuántos de los 18,664 sin tipificar tienen un certificado lo bastante
explícito como para asignarles tipo con confianza, y a qué categoría.

El modelo aprende de las etiquetas que SINADEF ya asignó: reproduce esas
decisiones, no las valida. Un caso clasificado con alta confianza es un caso
*documentalmente parecido* a los que SINADEF etiquetó así.
"""
from __future__ import annotations

import sys
import time
from pathlib import Path

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 classification_report, confusion_matrix
from sklearn.model_selection import StratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder

BASE = Path(__file__).resolve().parent
sys.path.insert(0, str(BASE))
from features_v2 import cargar, preparar  # noqa: E402

CAT = ["NECROPSIA", "TIPO_LUGAR", "SEXO", "INSTITUCION", "TIPO_SEGURO", "TIPO_CDEF",
       "ESTADO_CIVIL", "NIVEL_DE_INSTRUCCION", "DEPARTAMENTO_FALLECIMIENTO"]
NUM = ["edad_num", "dow", "n_causas", "largo_txt"]
COLS = ["txt_a", "txt_resto"] + CAT + NUM


def construir():
    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(objective="multiclass", n_estimators=350, learning_rate=0.07,
                         num_leaves=48, min_child_samples=30, subsample=0.9,
                         subsample_freq=1, colsample_bytree=0.8, verbosity=-1,
                         n_jobs=8, random_state=42)
    return Pipeline([("pre", pre), ("clf", clf)])


def main():
    t0 = time.time()
    df = preparar(cargar())
    lab = df[df.MUERTE_VIOLENTA != "NO SE CONOCE"].copy().reset_index(drop=True)
    nsc = df[df.MUERTE_VIOLENTA == "NO SE CONOCE"].copy().reset_index(drop=True)
    print(f"tipificados {len(lab):,} · sin tipificar {len(nsc):,}\n", flush=True)

    y = lab["MUERTE_VIOLENTA"].to_numpy()

    print("── VALIDACIÓN CRUZADA (3 pliegues) — desempeño POR CLASE ──", flush=True)
    oof = np.empty(len(lab), dtype=object)
    for k, (itr, ite) in enumerate(StratifiedKFold(3, shuffle=True, random_state=42).split(lab[COLS], y), 1):
        m = construir().fit(lab.iloc[itr][COLS], y[itr])
        oof[ite] = m.predict(lab.iloc[ite][COLS])
        print(f"   pliegue {k}/3 listo", flush=True)
    print()
    print(classification_report(y, oof, digits=3, zero_division=0))

    print("── MATRIZ DE CONFUSIÓN (fila = real, columna = predicho) ──")
    clases = sorted(set(y))
    cm = confusion_matrix(y, oof, labels=clases)
    w = max(len(c) for c in clases)
    print(" " * (w + 2) + "".join(f"{c[:11]:>13}" for c in clases))
    for i, c in enumerate(clases):
        print(f"  {c:{w}}" + "".join(f"{v:13,}" for v in cm[i]))

    print("\n── MODELO FINAL aplicado a los sin tipificar ──", flush=True)
    final = construir().fit(lab[COLS], y)
    proba = final.predict_proba(nsc[COLS])
    cf = list(final.named_steps["clf"].classes_)
    nsc["pred"] = [cf[i] for i in proba.argmax(1)]
    nsc["conf"] = proba.max(1)
    for i, c in enumerate(cf):
        nsc[f"p_{c}"] = proba[:, i]
    nsc.to_parquet(BASE / "nsc_multiclase.parquet", index=False)

    print(f"\n  {'umbral':24}{'casos':>9}{'% de 18,664':>14}")
    for u in (0.5, 0.6, 0.7, 0.8, 0.9, 0.95):
        n = int((nsc.conf >= u).sum())
        print(f"  confianza >= {u:<11.2f}{n:9,}{n/len(nsc)*100:13.1f}%")

    print(f"\n  {'clase asignada':24}{'total':>9}{'conf>=0.8':>12}{'conf>=0.9':>12}")
    for c in cf:
        s = nsc[nsc.pred == c]
        print(f"  {c:24}{len(s):9,}{int((s.conf>=0.8).sum()):12,}{int((s.conf>=0.9).sum()):12,}")
    print(f"\n  tiempo {time.time()-t0:.0f}s  →  {BASE/'nsc_multiclase.parquet'}")


if __name__ == "__main__":
    main()