"""Etapa 3b: testes de sensibilidade dos sinais.

Verifica o quanto os resultados mudam quando se trocam escolhas do método:
  A. raio do "entorno de capital" (20, 40 e 60 km);
  B. critério de "polo de tratamento" (mais brando e mais rígido);
  C. correção FDR sobre todos os tipos de câncer juntos, em vez de dentro de cada tipo;
  D. população anual (estimativas do IBGE 2020, 2021 e 2024; Censo 2022; 2023 interpolado)
     no lugar da população fixa de 2022;
  E. estabilidade temporal dos sinais de esôfago (2020–2022 vs 2023–2024). Não é validação
     independente: os municípios foram selecionados com o período completo.

Uso: python sensibilidade.py   (depois de analisar.py)
"""
import os
import tempfile
from collections import Counter
from concurrent.futures import ProcessPoolExecutor

import datasus_dbc
import dbfread
import pandas as pd

from analisar import POLO, ENTORNO_KM, calcular, carregar, preparar_base
from config import ANOS, BRUTOS, GRUPOS, RESULTADOS, UFS


def resumo(res, nome):
    s = res[res["sinal"]]
    p = s[s["contexto"] == "prioritário"]
    esof = sorted(p[p["grupo"] == "esofago"]["cod_ibge"])
    return {"cenario": nome, "sinais": len(s), "prioritarios": len(p),
            "esofago_prioritarios": len(esof), "_esof": esof,
            "_pares": set(zip(p["grupo"], p["cod_ibge"]))}


def pop_anual(pop):
    """Reescala a população por idade e sexo de 2022 pelo total de cada ano."""
    est = pd.read_csv(BRUTOS / "populacao_estimativas_anuais.csv").pivot(index="cod_ibge", columns="ano", values="pop")
    censo = pop.groupby("cod_ibge")["pop"].sum()
    est[2022] = censo
    est[2023] = (est[2022] + est[2024]) / 2
    fator = (est[[2020, 2021, 2022, 2023, 2024]].sum(axis=1) / est[2022]).rename("fator")
    p = pop.merge(fator, left_on="cod_ibge", right_index=True, how="left")
    p["pessoa_anos"] = p["pop"] * p["fator"].fillna(5)
    return p.drop(columns="fator")


def contar_c15(arq):
    """Óbitos por câncer de esôfago (C15) por município de residência e ano, de um arquivo do SIM."""
    ano, c = int(arq.stem[-4:]), Counter()
    fd, tmp = tempfile.mkstemp(suffix=".dbf")
    os.close(fd)
    try:
        datasus_dbc.decompress(str(arq), tmp)
        for r in dbfread.DBF(tmp, encoding="latin-1", char_decode_errors="replace"):
            if (r.get("CAUSABAS") or "").strip().upper()[:3] == "C15":
                m = (r.get("CODMUNRES") or "").strip()
                if m.isdigit():
                    c[(int(m), ano)] += 1
    finally:
        os.remove(tmp)
    return c


def estabilidade_esofago(ref, cenarios):
    """E. Excesso por período e presença em cada cenário, para os sinais prioritários de esôfago."""
    arqs = [BRUTOS / "sim" / f"DO{uf}{a}.dbc" for a in ANOS for uf in UFS]
    tot = Counter()
    with ProcessPoolExecutor() as ex:
        for c in ex.map(contar_c15, arqs):
            tot.update(c)
    por_ano = pd.DataFrame([(m, a, n) for (m, a), n in tot.items()], columns=["cod6", "ano", "obitos"])
    e = ref[(ref["grupo"] == "esofago") & (ref["contexto"] == "prioritário")].copy()
    e["cod6"] = e["cod_ibge"] // 10
    piv = por_ano[por_ano["cod6"].isin(e["cod6"])].pivot_table(index="cod6", columns="ano", values="obitos", fill_value=0)
    e = e.set_index("cod6").join(piv).reset_index()
    o1, o2 = e[[2020, 2021, 2022]].sum(axis=1), e[[2023, 2024]].sum(axis=1)
    e["rmp_br_2020_22"] = o1 / (e["esperados_br"] * 3 / 5)
    e["rmp_br_2023_24"] = o2 / (e["esperados_br"] * 2 / 5)
    e["maior_ano_pct"] = e[ANOS].max(axis=1) / e["observados"]
    for c in cenarios[1:]:
        e[c["cenario"][:2].strip() + " " + c["cenario"][3:40]] = e["cod_ibge"].isin(c["_esof"])
    cols_cen = [col for col in e.columns if isinstance(col, str) and col[:1] in "ABCD" and col[1:2] in (".", " ")]
    e["em_todos_os_cenarios"] = e[cols_cen].all(axis=1)
    e.to_csv(RESULTADOS / "estabilidade_esofago.csv", index=False, float_format="%.4g")
    print(f"Esôfago: {int(e['em_todos_os_cenarios'].sum())} de {len(e)} prioritários em todos os cenários")
    return e


def main():
    pop, ob, nomes = carregar()
    ob["uf"] = ob["cod6"] // 10000
    cenarios = []

    base = preparar_base(pop, ob, nomes)
    ref, _ = calcular(pop, ob, base)
    principal = resumo(ref, f"Principal (entorno {ENTORNO_KM} km, polo ≥200 e ≥25%, FDR por tipo, pop. 2022 fixa)")
    cenarios.append(principal)

    for km_ in (20, 60):
        b = preparar_base(pop, ob, nomes, entorno_km=km_)
        r, _ = calcular(pop, ob, b)
        cenarios.append(resumo(r, f"A. Entorno de capital = {km_} km"))

    for nome, polo in (("B. Polo mais brando (≥100 óbitos e ≥15%)", {"nao_residentes_min": 100, "razao_min": 0.15}),
                       ("B. Polo mais rígido (≥400 óbitos e ≥40%)", {"nao_residentes_min": 400, "razao_min": 0.40})):
        b = preparar_base(pop, ob, nomes, polo=polo)
        r, _ = calcular(pop, ob, b)
        cenarios.append(resumo(r, nome))

    r, _ = calcular(pop, ob, base, fdr_global=True)
    cenarios.append(resumo(r, "C. FDR sobre os 13 tipos de câncer juntos"))

    r, _ = calcular(pop_anual(pop), ob, base)
    cenarios.append(resumo(r, "D. População anual (IBGE 2020, 2021, 2024; Censo 2022; 2023 interpolado)"))

    linhas = []
    for c in cenarios:
        comuns = len(c["_pares"] & principal["_pares"])
        linhas.append({
            "cenario": c["cenario"], "sinais": c["sinais"], "prioritarios": c["prioritarios"],
            "prioritarios_em_comum_com_principal": comuns,
            "esofago_prioritarios": c["esofago_prioritarios"],
            "esofago_em_comum_com_principal": len(set(c["_esof"]) & set(principal["_esof"])),
        })
    tab = pd.DataFrame(linhas)
    tab.to_csv(RESULTADOS / "sensibilidade.csv", index=False)
    estabilidade_esofago(ref, cenarios)
    pd.set_option("display.width", 250, "display.max_colwidth", 90)
    print(tab.to_string(index=False))
    return tab


if __name__ == "__main__":
    main()