"""Etapa 3: compara cada município com o esperado e aponta os sinais.

Método (padronização indireta, o padrão em epidemiologia):
1. Taxas de referência por sexo e faixa etária, calculadas para o Brasil e para cada UF.
2. Óbitos esperados de cada município = taxas de referência x população do município.
   Assim, um município não parece "pior" só porque tem mais idosos.
3. RMP (razão de mortalidade padronizada) = observados / esperados. 1,0 = igual à referência.
4. Suavização bayesiana empírica (Poisson-Gama): municípios pequenos têm RMP instável;
   a suavização puxa os valores extremos de pouca evidência para perto da média.
5. Teste de Poisson com correção para múltiplas comparações (FDR de Benjamini-Hochberg),
   porque ao testar 5.570 municípios alguns sairiam "significativos" só por acaso.

Um SINAL é um município acima do esperado para o próprio estado, com evidência forte:
  q < 0,01, probabilidade posterior de RMP > 1 acima de 99%, RMP suavizada >= 1,3 e
  pelo menos 10 óbitos observados.

Comparar com o próprio estado reduz o efeito de diferenças regionais na qualidade do
diagnóstico e do registro. Mesmo assim, sinal é hipótese, não conclusão.

Cada sinal recebe um contexto, porque dois vieses conhecidos inflam o número de óbitos:
- CAPITAL: diagnostica e registra melhor que o interior do mesmo estado;
- POLO de tratamento: pacientes de outras cidades acabam registrados como moradores de
  cidades com hospital de câncer. Medimos isso pelos óbitos por câncer de NÃO residentes
  ocorridos no município (>= 200 óbitos e >= 25% dos óbitos por câncer de residentes).
- ENTORNO DE CAPITAL: municípios a até 40 km de uma capital compartilham o acesso ao
  diagnóstico dela (ex.: Timon, vizinha de Teresina).
Sinais sem esses vieses são os PRIORITÁRIOS para investigação.

Uso: python analisar.py
"""
import json

import numpy as np
import pandas as pd
from scipy import stats

from config import ANOS, BRUTOS, GRUPOS, PROCESSADOS, RESULTADOS

ANOS_PESSOA = len(ANOS)  # população do Censo 2022 x 5 anos de observação
CRITERIO = {"q_max": 0.01, "prob_min": 0.99, "rmp_min": 1.3, "obitos_min": 10}
POLO = {"nao_residentes_min": 200, "razao_min": 0.25}
ENTORNO_KM = 40
CAPITAIS = {
    1100205, 1200401, 1302603, 1400100, 1501402, 1600303, 1721000, 2111300, 2211001,
    2304400, 2408102, 2507507, 2611606, 2704302, 2800308, 2927408, 3106200, 3205309,
    3304557, 3550308, 4106902, 4205407, 4314902, 5002704, 5103403, 5208707, 5300108,
}


def carregar():
    pop = pd.read_csv(BRUTOS / "populacao_censo2022.csv")
    pop["cod6"] = pop["cod_ibge"] // 10
    pop["uf"] = pop["cod_ibge"] // 100000
    pop["pessoa_anos"] = pop["pop"] * ANOS_PESSOA
    ob = pd.read_csv(PROCESSADOS / "obitos_agregados.csv")
    muns = json.loads((BRUTOS / "municipios.json").read_text())
    nomes = pd.DataFrame([{
        "cod_ibge": m["id"],
        "municipio": m["nome"],
        "uf_sigla": (m.get("microrregiao") or {}).get("mesorregiao", {}).get("UF", {}).get("sigla")
                    or (m.get("regiao-imediata") or {}).get("regiao-intermediaria", {}).get("UF", {}).get("sigla"),
        "regiao_imediata": (m.get("regiao-imediata") or {}).get("nome"),
    } for m in muns])
    return pop, ob, nomes


def centroides():
    """Centro aproximado (meio da caixa envolvente) de cada município, pela malha do IBGE."""
    geo = json.loads((BRUTOS / "malha_municipios.geojson").read_text())
    saida = {}
    for f in geo["features"]:
        g = f["geometry"]
        polys = g["coordinates"] if g["type"] == "MultiPolygon" else [g["coordinates"]]
        pts = [p for poly in polys for p in poly[0]]
        lon = (min(p[0] for p in pts) + max(p[0] for p in pts)) / 2
        lat = (min(p[1] for p in pts) + max(p[1] for p in pts)) / 2
        saida[int(f["properties"]["codarea"])] = (lat, lon)
    return saida


def km(a, b):
    lat1, lon1, lat2, lon2 = map(np.radians, (a[0], a[1], b[0], b[1]))
    h = np.sin((lat2 - lat1) / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin((lon2 - lon1) / 2) ** 2
    return 6371 * 2 * np.arcsin(np.sqrt(h))


def bh_fdr(p):
    """Valores q de Benjamini-Hochberg."""
    p = np.asarray(p, dtype=float)
    n = len(p)
    ordem = np.argsort(p)
    q = p[ordem] * n / np.arange(1, n + 1)
    q = np.minimum.accumulate(q[::-1])[::-1]
    saida = np.empty(n)
    saida[ordem] = np.minimum(q, 1)
    return saida


def bayes_empirico(o, e):
    """Suavização Poisson-Gama com priori estimada por momentos (Marshall, 1991)."""
    o, e = np.asarray(o, float), np.asarray(e, float)
    ok = e > 0
    m = o[ok].sum() / e[ok].sum()
    theta = np.where(ok, o / np.where(ok, e, 1), m)
    s2 = (e[ok] * (theta[ok] - m) ** 2).sum() / e[ok].sum()
    var = max(s2 - m / e[ok].mean(), 1e-4 * m * m)
    alfa, beta = m * m / var, m / var
    post_media = (o + alfa) / (e + beta)
    prob_acima = stats.gamma.sf(1.0, a=o + alfa, scale=1 / (e + beta))
    return post_media, prob_acima, {"media_priori": m, "var_priori": var, "alfa": alfa, "beta": beta}


def comparar(obs, pop_g, ob_g, nivel):
    """RMP de cada município contra a referência do Brasil (nivel=None) ou da própria UF."""
    chaves = ["sexo", "faixa"] + ([nivel] if nivel else [])
    ref_ob = ob_g.groupby(chaves)["obitos"].sum()
    ref_pa = pop_g.groupby(chaves)["pessoa_anos"].sum()
    taxa = (ref_ob / ref_pa).fillna(0).rename("taxa").reset_index()
    esp = pop_g.merge(taxa, on=chaves, how="left")
    esp["esperados"] = esp["pessoa_anos"] * esp["taxa"].fillna(0)
    esperados = esp.groupby("cod6")["esperados"].sum()
    df = pd.DataFrame({"observados": obs}).join(esperados, how="outer").fillna(0)
    o, e = df["observados"].to_numpy(), df["esperados"].to_numpy()
    df["rmp"] = np.where(e > 0, o / np.where(e > 0, e, 1), np.nan)
    df["ic95_inf"] = np.where(o > 0, stats.chi2.ppf(0.025, 2 * o) / 2, 0) / np.where(e > 0, e, np.nan)
    df["ic95_sup"] = stats.chi2.ppf(0.975, 2 * (o + 1)) / 2 / np.where(e > 0, e, np.nan)
    df["p_acima"] = np.where(e > 0, stats.poisson.sf(o - 1, e), 1.0)
    df["q_acima"] = bh_fdr(df["p_acima"])
    df["rmp_suav"], df["prob_acima"], df.attrs["priori"] = bayes_empirico(o, e)
    return df


def preparar_base(pop, ob, nomes, entorno_km=ENTORNO_KM, polo=POLO):
    """Atributos de cada município: qualidade do registro, capital, entorno e polo."""
    base = pop[["cod6", "cod_ibge", "uf"]].drop_duplicates().merge(nomes, on="cod_ibge", how="left")
    base = base.set_index("cod6")
    # Qualidade do registro: % de óbitos com causa mal definida (capítulo R da CID-10)
    aux = ob[ob["grupo"].isin(["todas_causas", "mal_definidas"])].pivot_table(
        index="cod6", columns="grupo", values="obitos", aggfunc="sum", fill_value=0)
    base["obitos_total"] = aux["todas_causas"].reindex(base.index).fillna(0).astype(int)
    base["pct_mal_definidas"] = (100 * aux["mal_definidas"] / aux["todas_causas"]).reindex(base.index).round(1)
    base["populacao"] = pop.groupby("cod6")["pop"].sum()
    base["capital"] = base["cod_ibge"].isin(CAPITAIS)
    cent = centroides()
    caps = [cent[c] for c in CAPITAIS if c in cent]
    base["entorno_capital"] = [
        (not cap) and cod in cent and min(km(cent[cod], c) for c in caps) <= entorno_km
        for cod, cap in zip(base["cod_ibge"], base["capital"])
    ]
    nao_res = ob[ob["grupo"] == "cancer_nao_residentes"].groupby("cod6")["obitos"].sum()
    res_cancer = ob[ob["grupo"] == "todos"].groupby("cod6")["obitos"].sum()
    base["cancer_nao_residentes"] = nao_res.reindex(base.index).fillna(0).astype(int)
    base["polo"] = (base["cancer_nao_residentes"] >= polo["nao_residentes_min"]) & (
        base["cancer_nao_residentes"] >= polo["razao_min"] * res_cancer.reindex(base.index).fillna(0))
    return base


def classificar(df):
    df["sinal"] = (
        (df["q_acima_uf"] < CRITERIO["q_max"])
        & (df["prob_acima_uf"] > CRITERIO["prob_min"])
        & (df["rmp_suav_uf"] >= CRITERIO["rmp_min"])
        & (df["observados"] >= CRITERIO["obitos_min"])
    )
    df["contexto"] = np.select(
        [~df["sinal"], df["capital"], df["polo"], df["entorno_capital"]],
        ["", "capital", "polo de tratamento", "entorno de capital"], default="prioritário")
    return df


def calcular(pop, ob, base, fdr_global=False, verbose=False):
    """RMP, suavização, testes e sinais para todos os grupos.

    fdr_global=False: correção FDR dentro de cada tipo de câncer (5.570 testes por grupo).
    fdr_global=True:  correção sobre todos os testes de todos os grupos juntos.
    """
    ob = ob[ob["cod6"].isin(base.index)]
    linhas, prioris = [], []
    for grupo, info in GRUPOS.items():
        ob_g = ob[ob["grupo"] == grupo]
        pop_g = pop if info["sexo"] is None else pop[pop["sexo"] == info["sexo"]]
        ob_g = ob_g if info["sexo"] is None else ob_g[ob_g["sexo"] == info["sexo"]]
        obs = ob_g.groupby("cod6")["obitos"].sum().reindex(base.index).fillna(0)
        br = comparar(obs, pop_g, ob_g, None)
        uf = comparar(obs, pop_g, ob_g, "uf")
        for ref, r in (("br", br), ("uf", uf)):
            prioris.append({"grupo": grupo, "referencia": ref, **r.attrs["priori"]})
        df = base[["cod_ibge", "municipio", "uf_sigla", "regiao_imediata", "populacao",
                   "pct_mal_definidas", "capital", "entorno_capital", "polo"]].copy()
        df["grupo"] = grupo
        df["observados"] = obs.astype(int)
        for pref, r in (("br", br), ("uf", uf)):
            for c in ["esperados", "rmp", "ic95_inf", "ic95_sup", "p_acima", "q_acima", "rmp_suav", "prob_acima"]:
                df[f"{c}_{pref}"] = r[c].reindex(df.index)
        linhas.append(df.reset_index())
    res = pd.concat(linhas, ignore_index=True)
    if fdr_global:
        res["q_acima_uf"] = bh_fdr(res["p_acima_uf"])
        res["q_acima_br"] = bh_fdr(res["p_acima_br"])
    res = classificar(res)
    # Sinais vizinhos: outros sinais do mesmo câncer na mesma região imediata do IBGE
    viz = res[res["sinal"]].groupby(["grupo", "regiao_imediata"]).size()
    chave = list(zip(res["grupo"], res["regiao_imediata"]))
    res["sinais_na_regiao"] = np.where(res["sinal"], [viz.get(k, 0) for k in chave], 1).astype(int) - 1
    res["sinais_na_regiao"] = res["sinais_na_regiao"].clip(lower=0)
    if verbose:
        for grupo, info in GRUPOS.items():
            d = res[res["grupo"] == grupo]
            print(f"{info['nome']:32s} óbitos={int(d['observados'].sum()):>8,}  sinais={int(d['sinal'].sum()):>4}"
                  f"  prioritários={int((d['contexto'] == 'prioritário').sum()):>3}")
    return res, pd.DataFrame(prioris)


def main():
    RESULTADOS.mkdir(parents=True, exist_ok=True)
    pop, ob, nomes = carregar()
    ob["uf"] = ob["cod6"] // 10000
    base = preparar_base(pop, ob, nomes)
    fora = set(ob["cod6"]) - set(base.index) - set(ob[ob["grupo"] == "cancer_nao_residentes"]["cod6"])
    if fora:
        n = ob[ob["cod6"].isin(fora) & (ob["grupo"] == "todas_causas")]["obitos"].sum()
        print(f"Aviso: {len(fora)} códigos de município do SIM fora do Censo 2022 ({n} óbitos) ignorados.")

    res, prioris = calcular(pop, ob, base, verbose=True)
    res.to_csv(RESULTADOS / "municipios_rmp.csv", index=False, float_format="%.5g")
    prioris.to_csv(RESULTADOS / "parametros_suavizacao.csv", index=False, float_format="%.6g")

    sinais = res[res["sinal"]].copy()
    sinais["nome_grupo"] = sinais["grupo"].map(lambda g: GRUPOS[g]["nome"])
    sinais["ordem_ctx"] = (sinais["contexto"] != "prioritário").astype(int)
    sinais = sinais.sort_values(["ordem_ctx", "grupo", "rmp_suav_uf"], ascending=[True, True, False])
    cols = ["contexto", "nome_grupo", "municipio", "uf_sigla", "regiao_imediata", "cod_ibge",
            "populacao", "observados", "esperados_uf", "rmp_uf", "ic95_inf_uf", "ic95_sup_uf",
            "rmp_suav_uf", "q_acima_uf", "rmp_br", "rmp_suav_br", "sinais_na_regiao", "pct_mal_definidas"]
    sinais[cols].to_csv(RESULTADOS / "sinais.csv", index=False, float_format="%.4g")

    # Visão por estado: RMP de cada UF contra o Brasil
    ufs = res.groupby(["grupo", "uf_sigla"])[["observados", "esperados_br"]].sum().reset_index()
    ufs["rmp_br"] = ufs["observados"] / ufs["esperados_br"]
    ufs.to_csv(RESULTADOS / "ufs_rmp.csv", index=False, float_format="%.4g")

    print(f"\nTotal de sinais: {len(sinais)} (em {sinais['cod_ibge'].nunique()} municípios)")
    print(sinais["contexto"].value_counts().to_string())
    return res, sinais


if __name__ == "__main__":
    main()