"""Etapa 2: transforma as declarações de óbito em contagens agregadas.

Para cada município de residência, sexo, faixa etária e grupo de causa, conta os óbitos
de 2020 a 2024. Nenhum dado individual sai desta etapa: o resultado é só contagem.

Uso: python preparar.py
"""
import os
import tempfile
from collections import Counter
from concurrent.futures import ProcessPoolExecutor

import datasus_dbc
import dbfread
import pandas as pd

from config import ANOS, AUXILIARES, BRUTOS, FAIXAS, GRUPOS, PROCESSADOS, UFS

# Faixas de CID de todos os grupos, avaliadas pelos 3 primeiros caracteres da causa básica
REGRAS = [(g, ini, fim, d["sexo"]) for g, d in GRUPOS.items() for ini, fim in d["cid"]]
REGRAS += [(g, ini, fim, None) for g, faixas in AUXILIARES.items() for ini, fim in faixas]


def idade_em_anos(campo):
    """Campo IDADE do SIM: 1º dígito é a unidade (0–3 = menos de 1 ano, 4 = anos, 5 = 100+)."""
    if not campo or len(campo) != 3 or not campo.isdigit():
        return None
    unidade, valor = campo[0], int(campo[1:])
    if unidade in "0123":
        return 0
    if unidade == "4":
        return valor
    if unidade == "5":
        return 100 + valor
    return None


def processar(arquivo):
    cont = Counter()
    descartes = Counter()
    fd, tmp = tempfile.mkstemp(suffix=".dbf")
    os.close(fd)
    try:
        datasus_dbc.decompress(str(arquivo), tmp)
        for reg in dbfread.DBF(tmp, encoding="latin-1", char_decode_errors="replace"):
            cid = (reg.get("CAUSABAS") or "").strip().upper()[:3]
            mun = (reg.get("CODMUNRES") or "").strip()
            ocor = (reg.get("CODMUNOCOR") or "").strip()
            # Polos de tratamento: óbitos por câncer de não residentes ocorridos no município
            if "C00" <= cid <= "C97" and len(ocor) == 6 and ocor.isdigit() and ocor != mun:
                cont[(int(ocor), 0, "todas", "cancer_nao_residentes")] += 1
            sexo = (reg.get("SEXO") or "").strip()
            idade = idade_em_anos((reg.get("IDADE") or "").strip())
            if len(mun) != 6 or not mun.isdigit() or mun.endswith("0000"):
                descartes["municipio_ignorado"] += 1
                continue
            if sexo not in ("1", "2"):
                descartes["sexo_ignorado"] += 1
                continue
            if idade is None:
                descartes["idade_ignorada"] += 1
                continue
            faixa = FAIXAS[min(idade // 5, len(FAIXAS) - 1)]
            s = int(sexo)
            for grupo, ini, fim, sexo_grupo in REGRAS:
                if ini <= cid <= fim and (sexo_grupo is None or sexo_grupo == s):
                    cont[(int(mun), s, faixa, grupo)] += 1
    finally:
        os.remove(tmp)
    return arquivo.name, cont, descartes


if __name__ == "__main__":
    PROCESSADOS.mkdir(parents=True, exist_ok=True)
    arquivos = [BRUTOS / "sim" / f"DO{uf}{ano}.dbc" for ano in ANOS for uf in UFS]
    faltando = [a.name for a in arquivos if not a.exists()]
    if faltando:
        raise SystemExit(f"Faltam arquivos do SIM (rode baixar.py): {faltando[:5]}...")

    total = Counter()
    descartes = Counter()
    with ProcessPoolExecutor() as ex:
        for n, (nome, cont, desc) in enumerate(ex.map(processar, arquivos), 1):
            total.update(cont)
            descartes.update(desc)
            print(f"  {nome} ({n}/{len(arquivos)})", flush=True)

    df = pd.DataFrame(
        [(m, s, f, g, n) for (m, s, f, g), n in total.items()],
        columns=["cod6", "sexo", "faixa", "grupo", "obitos"],
    )
    df.to_csv(PROCESSADOS / "obitos_agregados.csv", index=False)
    resumo = df.groupby("grupo")["obitos"].sum().sort_values(ascending=False)
    print("\nÓbitos 2020–2024 por grupo:")
    print(resumo.to_string())
    print("\nRegistros descartados:", dict(descartes))
    pd.Series(descartes).to_csv(PROCESSADOS / "descartes.csv", header=["registros"])