"""Fase 1b: óbitos por ANO, município de residência, sexo, faixa etária e grupo (2015–2024).

Reaproveita exatamente as regras de preparar.py (mesmas exclusões e grupos); só guarda o ano,
que vem do nome do arquivo (um arquivo por UF e ano no FTP, ou um CSV nacional por ano).
2020–2024: arquivos .dbc do FTP do DATASUS; 2015–2019: CSVs do OpenDataSUS (ver validação).
Saída: dados/processados/obitos_anuais.csv. Confere que a soma bate com obitos_agregados*.csv.

Uso: python contar_anual.py
"""
import re
from collections import Counter
from concurrent.futures import ProcessPoolExecutor

import pandas as pd

from config import ANOS, ANOS_VALIDACAO, BRUTOS, PROCESSADOS, UFS
from preparar import processar, processar_csv


def contar(arquivo):
    ano = int(re.search(r"(20\d\d)", arquivo.name).group(1))
    funcao = processar_csv if arquivo.suffix == ".csv" else processar
    nome, cont, descartes = funcao(arquivo)
    return ano, nome, cont


def main():
    arquivos = [BRUTOS / "sim" / f"DO{uf}{ano}.dbc" for ano in ANOS for uf in UFS]
    arquivos += [BRUTOS / "sim_csv" / f"Mortalidade_Geral_{ano}.csv" for ano in ANOS_VALIDACAO]
    total = Counter()
    with ProcessPoolExecutor() as ex:
        for n, (ano, nome, cont) in enumerate(ex.map(contar, arquivos), 1):
            for (m, s, f, g), k in cont.items():
                if g != "cancer_nao_residentes":
                    total[(ano, m, s, f, g)] += k
            print(f"  {nome} ({n}/{len(arquivos)})", flush=True)
    df = pd.DataFrame([(*k, v) for k, v in total.items()],
                      columns=["ano", "cod6", "sexo", "faixa", "grupo", "obitos"])
    df.to_csv(PROCESSADOS / "obitos_anuais.csv", index=False)

    # Conferência com as bases agregadas usadas na análise principal e na validação
    for anos, sufixo in [(ANOS, ""), (ANOS_VALIDACAO, f"_{ANOS_VALIDACAO[0]}_{ANOS_VALIDACAO[-1]}")]:
        ag = pd.read_csv(PROCESSADOS / f"obitos_agregados{sufixo}.csv")
        ag = ag[ag["grupo"] != "cancer_nao_residentes"].groupby("grupo")["obitos"].sum()
        an = df[df["ano"].isin(anos)].groupby("grupo")["obitos"].sum()
        dif = (ag - an.reindex(ag.index).fillna(0)).abs().sum()
        print(f"{anos[0]}–{anos[-1]}: diferença com obitos_agregados{sufixo}.csv = {dif:.0f}")
    print(df.groupby(["ano"])["obitos"].sum().to_string())


if __name__ == "__main__":
    main()