"""Fase 2b: análises A a E de PLANO-fase2b.md (Minas Gerais; Brasil como secundária).

Reaproveita os modelos da Fase 2 (hipoteses.py): Poisson com os esperados do estado como offset e erros-padrão
sanduíche agrupados por região imediata. Em Minas Gerais, as variáveis são padronizadas pelo desvio-padrão de MG;
no Brasil, pelo desvio-padrão nacional.

Saídas: resultados/fase2b_*.csv
Uso: python hipoteses_2b.py   (depois de exposicoes_2b.py e hipoteses.py)
"""
import numpy as np
import pandas as pd

from analisar import bh_fdr
from config import RESULTADOS
from hipoteses import codigos_regiao, conjunta, desfecho, diferenca, univariada

NOVAS = ["aguardente", "milho", "agua_pn_1991", "agua_pn_2000", "rural_1991", "endoscopia", "dist_oncologia"]
SENS = ["aguardente_2017", "aguardente_produz", "endoscopia_2020"]
LOG1P = ["aguardente", "aguardente_2017", "milho", "endoscopia", "endoscopia_2020"]
FASE2 = ["rural", "pib_pc", "altitude"]


BINARIAS = ["aguardente_produz"]   # entram como 0/1, sem padronização: a RT compara produtores e não produtores


def z(df, cols):
    t = df[cols].copy()
    for c in cols:
        if c in LOG1P:
            t[c] = np.log1p(t[c])
    zz = (t - t.mean()) / t.std()
    for c in cols:
        if c in BINARIAS:
            zz[c] = df[c]
    return zz


def main():
    rmp = pd.read_csv(RESULTADOS / "municipios_rmp.csv")
    info = rmp.drop_duplicates("cod6").set_index("cod6")[["cod_ibge", "municipio", "uf_sigla", "regiao_imediata",
                                                         "pct_mal_definidas"]]
    info["regiao_imediata"] = info["cod_ibge"].map(codigos_regiao())   # código oficial, não o nome
    e2 = pd.read_csv(RESULTADOS / "fase2_exposicoes.csv").set_index("cod_ibge")
    e2b = pd.read_csv(RESULTADOS / "fase2b_exposicoes.csv").set_index("cod_ibge")
    ex = e2[FASE2].join(e2b[NOVAS + SENS])
    ex.index = ex.index // 10
    ex = ex.reindex(info.index)
    ex["mal_definidas"] = info["pct_mal_definidas"]
    dados = {"esofago": desfecho("esofago", rmp), "pulmao": desfecho("pulmao", rmp)}
    todos = desfecho("todos", rmp)
    dados["outros"] = todos - dados["esofago"].reindex(todos.index).fillna(0)
    dados = {k: v.reindex(info.index) for k, v in dados.items()}

    mg = info.index[info["uf_sigla"] == "MG"]
    cols = NOVAS + SENS + FASE2 + ["mal_definidas"]
    zmg = z(ex.loc[mg], cols).reindex(info.index)
    zbr = z(ex, cols)

    # A e C: Minas Gerais, univariadas (esôfago e comparação) e diferença formal
    uni = []
    for g in dados:
        for v in NOVAS + SENS:
            uni.append({"ambito": "MG", "desfecho": g, "variavel": v, **univariada(dados[g], zmg, info, v, sel=mg)})
    uni = pd.DataFrame(uni)
    uni["q_bh"] = np.nan
    for g, idx in uni[uni["variavel"].isin(NOVAS)].groupby("desfecho").groups.items():
        uni.loc[idx, "q_bh"] = bh_fdr(uni.loc[idx, "p"].to_numpy())
    info_mg = info.loc[mg]
    dif = []
    for comp in ("pulmao", "outros"):
        for v in NOVAS:
            dif.append({"comparacao": comp, "variavel": v,
                        **diferenca(dados["esofago"].loc[mg], dados[comp].loc[mg], zmg.loc[mg], info_mg, v)})
    dif = pd.DataFrame(dif)
    dif["q_bh"] = dif.groupby("comparacao")["p_diferenca"].transform(lambda p: bh_fdr(p.to_numpy()))

    # B: cada variável nova acrescentada ao modelo da Fase 2 (ruralidade, PIB, altitude), sem e com causas mal definidas
    ajust = []
    for v in NOVAS + SENS:
        for extra in ([], ["mal_definidas"]):
            r = conjunta(dados["esofago"], zmg, info, FASE2 + extra + [v], sel=mg)
            linha = next(x for x in r if x["variavel"] == v)
            ajust.append({"variavel": v, "ajuste": "fase2" + ("+mal_definidas" if extra else ""), **linha})
    ajust = pd.DataFrame(ajust)

    # D: Brasil, sem e com efeitos fixos de estado
    br = []
    for v in NOVAS:
        for fe in (False, True):
            br.append({"ambito": "Brasil", "variavel": v, "efeito_fixo_uf": fe,
                       **univariada(dados["esofago"], zbr, info, v, efeito_uf=fe)})
    br = pd.DataFrame(br)

    # Descritivo: os 9 municípios de MG com sinal e as áreas dos agrupamentos de MG versus a mediana de MG
    sinais = rmp[(rmp["grupo"] == "esofago") & rmp["sinal"] & (rmp["contexto"] == "prioritário")]
    s_mg = sinais[sinais["uf_sigla"] == "MG"]["cod6"]
    mb = pd.read_csv(RESULTADOS / "varredura_municipios.csv")
    ag_mg = set(mb[(mb["grupo"] == "esofago") & (mb["uf_sigla"] == "MG")]["cod_ibge"] // 10)
    grupos = {"MG": list(mg), "9_sinais_MG": list(s_mg), "agrupamentos_MG": list(ag_mg & set(mg)),
              "resto_MG": [c for c in mg if c not in ag_mg]}
    desc = pd.DataFrame({f"{k}_{m}": (ex.loc[v, NOVAS].median() if m == "mediana" else ex.loc[v, NOVAS].notna().sum())
                         for k, v in grupos.items() for m in ("mediana", "n")})
    desc["total_9_sinais_MG"] = len(s_mg)
    corr = ex.loc[mg, NOVAS + FASE2].corr(method="spearman")

    for nome, df in {"fase2b_mg_univariada": uni, "fase2b_mg_diferenca": dif, "fase2b_mg_ajustada": ajust,
                     "fase2b_brasil": br}.items():
        df.to_csv(RESULTADOS / f"{nome}.csv", index=False, float_format="%.4g")
    desc.to_csv(RESULTADOS / "fase2b_mg_descritivo.csv", float_format="%.4g")
    corr.to_csv(RESULTADOS / "fase2b_mg_correlacoes.csv", float_format="%.3f")

    pd.set_option("display.width", 250, "display.max_columns", 30)
    print("A. MG, RT por DP (esôfago e comparação)")
    print(uni.pivot(index="variavel", columns="desfecho", values="rt_por_dp").round(3).to_string())
    print(uni[uni["desfecho"] == "esofago"][["variavel", "municipios", "rt_por_dp", "ic95_inf", "ic95_sup", "p", "q_bh"]]
          .round(4).to_string(index=False))
    print("\nC. Diferença esôfago versus comparação (MG)")
    print(dif.round(4).to_string(index=False))
    print("\nB. Ajustada pela Fase 2 (MG)")
    print(ajust[["variavel", "ajuste", "municipios", "rt_por_dp", "ic95_inf", "ic95_sup", "p"]].round(4).to_string(index=False))
    print("\nD. Brasil")
    print(br.round(4).to_string(index=False))
    print("\nDescritivo MG")
    print(desc.round(3).to_string())
    print("\nCorrelações MG")
    print(corr.round(2).to_string())


if __name__ == "__main__":
    main()