"""Fase 2: geração de hipóteses para o câncer de esôfago (PLANO-fase2.md), revista após uma revisão técnica assistida por IA.

Principal (como no plano, com a regra de ausentes revista em exposicoes.py):
A. Poisson por variável (óbitos 2015–2024, offset = log dos esperados do próprio estado), razão de taxas (RT) por
   desvio-padrão da variável transformada (log(1 + x) nas variáveis de produção), erros-padrão sanduíche agrupados
   por região imediata; Benjamini–Hochberg entre as 10 variáveis. Só municípios com dado na variável.
B. Modelo conjunto com as 10 variáveis e a % de causas mal definidas (casos completos).
C. Desfechos de comparação: pulmão e todos os cânceres exceto esôfago. A diferença entre a RT do esôfago e a do
   desfecho de comparação é testada formalmente num modelo empilhado (mesmos municípios, mesma variância agrupada).
D. Agrupamentos versus vizinhos, com Benjamini–Hochberg entre as 10 variáveis.
E. Os 16 municípios versus a mediana do estado.

Sensibilidades (não pré-especificadas, todas rotuladas): efeitos fixos de estado nos modelos A e B; regra antiga de
ausentes (ausente = 0); lenha só da extração; modelo conjunto sem lenha e sem erva-mate extraída (mais municípios);
variáveis históricas (população rural em 2000 e 2010, poço ou nascente em 2010); análise por macrorregião (F).

Saídas: resultados/fase2_*.csv
Uso: python hipoteses.py   (depois de exposicoes.py, varredura.py e validar.py)
"""
import json

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

from analisar import bh_fdr, centroides
from config import BRUTOS, RESULTADOS
from varredura import esperados_validacao

VARIAVEIS = ["fumo", "mate_cultivo", "mate_extracao", "cafe", "lenha", "agua_poco", "agua_nascente", "rural",
             "pib_pc", "altitude"]
EXTRAS = ["lenha_extracao", "rural_2000", "rural_2010", "agua_poco_nascente_2010"]
LOG1P = ["fumo", "mate_cultivo", "mate_extracao", "cafe", "lenha", "lenha_extracao"]
DESFECHOS = {"esofago": "Esôfago", "pulmao": "Pulmão (comparação)", "outros": "Outros cânceres, exceto esôfago (comparação)"}
RAIO_VIZINHOS_KM = 100
REGIOES = {"Sul (PR, SC, RS)": ["PR", "SC", "RS"], "Minas Gerais": ["MG"]}


def codigos_regiao():
    """Código oficial da região imediata do IBGE por município (o nome não basta: há nomes repetidos entre UFs,
    como Itabaiana/PB e Itabaiana/SE, Valença/BA e Valença/RJ)."""
    muns = json.loads((BRUTOS / "municipios.json").read_text())
    return {m["id"]: m["regiao-imediata"]["id"] for m in muns}


def desfecho(grupo, rmp):
    """Observados e esperados (estado) em 2015–2024 por município (cod6)."""
    d = rmp[rmp["grupo"] == grupo].set_index("cod6")[["observados", "esperados_uf"]]
    v = esperados_validacao(grupo)
    return pd.DataFrame({"obs": d["observados"] + v["obs_val"].reindex(d.index).fillna(0),
                         "esp": d["esperados_uf"] + v["esp_val"].reindex(d.index).fillna(0)})


def poisson_agrupado(y, offset, X, grupos):
    """Poisson com offset (IRLS) e covariância sanduíche agrupada. Devolve coeficientes e covariância."""
    beta = np.linalg.lstsq(X, np.log((y + 0.5) / np.exp(offset)), rcond=None)[0]
    for _ in range(200):
        mu = np.exp(np.clip(offset + X @ beta, -50, 50))
        passo = np.linalg.solve(X.T @ (X * mu[:, None]), X.T @ (y - mu))
        beta += passo
        if np.abs(passo).max() < 1e-9:
            break
    mu = np.exp(offset + X @ beta)
    pao = np.linalg.pinv(X.T @ (X * mu[:, None]))
    ug = pd.DataFrame(X * (y - mu)[:, None]).groupby(grupos).sum().to_numpy()
    g = ug.shape[0]
    return beta, pao @ (ug.T @ ug * g / (g - 1)) @ pao


def dummies_uf(uf):
    return pd.get_dummies(uf, drop_first=True, dtype=float).to_numpy()


def resumo(b, cov, k):
    ep = np.sqrt(cov[k, k])
    return {"rt_por_dp": np.exp(b[k]), "ic95_inf": np.exp(b[k] - 1.96 * ep), "ic95_sup": np.exp(b[k] + 1.96 * ep),
            "p": 2 * stats.norm.sf(abs(b[k] / ep))}


def univariada(dados, z, info, v, efeito_uf=False, sel=None):
    d = pd.concat([dados, z[v].rename("x"), info[["regiao_imediata", "uf_sigla"]]], axis=1).dropna()
    d = d[d["esp"] > 0]
    if sel is not None:
        d = d[d.index.isin(sel)]
    cols = [np.ones(len(d)), d["x"].to_numpy()]
    X = np.column_stack(cols + ([dummies_uf(d["uf_sigla"])] if efeito_uf else []))
    b, cov = poisson_agrupado(d["obs"].to_numpy(float), np.log(d["esp"].to_numpy()), X, d["regiao_imediata"].to_numpy())
    return {"municipios": len(d), **resumo(b, cov, 1)}


def conjunta(dados, z, info, cols, efeito_uf=False, sel=None):
    d = pd.concat([dados, z[cols], info[["regiao_imediata", "uf_sigla"]]], axis=1).dropna()
    d = d[d["esp"] > 0]
    if sel is not None:
        d = d[d.index.isin(sel)]
    cols = [c for c in cols if d[c].std() > 0]
    X = np.column_stack([np.ones(len(d))] + [d[c].to_numpy() for c in cols]
                        + ([dummies_uf(d["uf_sigla"])] if efeito_uf else []))
    b, cov = poisson_agrupado(d["obs"].to_numpy(float), np.log(d["esp"].to_numpy()), X, d["regiao_imediata"].to_numpy())
    return [{"variavel": c, "municipios": len(d), **resumo(b, cov, k)} for k, c in enumerate(cols, 1)]


def diferenca(dados_a, dados_b, z, info, v):
    """Modelo empilhado (esôfago e comparação nos mesmos municípios): testa se as RT diferem."""
    base = pd.concat([z[v].rename("x"), info["regiao_imediata"]], axis=1).dropna()
    partes = []
    for k, dd in enumerate((dados_a, dados_b)):
        d = dd.join(base, how="inner")
        d = d[d["esp"] > 0].copy()
        d["k"] = k
        partes.append(d)
    comuns = partes[0].index.intersection(partes[1].index)
    d = pd.concat([p.loc[comuns] for p in partes])
    a = (d["k"] == 0).to_numpy(float)
    X = np.column_stack([a, 1 - a, d["x"].to_numpy() * a, d["x"].to_numpy() * (1 - a)])
    b, cov = poisson_agrupado(d["obs"].to_numpy(float), np.log(d["esp"].to_numpy()), X, d["regiao_imediata"].to_numpy())
    dif = b[2] - b[3]
    ep = np.sqrt(cov[2, 2] + cov[3, 3] - 2 * cov[2, 3])
    return {"razao_das_rt": np.exp(dif), "ic95_inf": np.exp(dif - 1.96 * ep), "ic95_sup": np.exp(dif + 1.96 * ep),
            "p_diferenca": 2 * stats.norm.sf(abs(dif / ep))}


def padronizar(exp):
    t = exp.copy()
    for v in LOG1P:
        if v in t:
            t[v] = np.log1p(t[v])
    return t, (t - t.mean()) / t.std()


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

    def ler(nome):
        e = pd.read_csv(RESULTADOS / nome).set_index("cod_ibge")
        e.index = e.index // 10
        return e.reindex(info.index)

    exp, exp0 = ler("fase2_exposicoes.csv"), ler("fase2_exposicoes_ausente_zero.csv")
    tr, zs = padronizar(exp[VARIAVEIS + EXTRAS])
    _, zs0 = padronizar(exp0[VARIAVEIS])
    zs["mal_definidas"] = (info["pct_mal_definidas"] - info["pct_mal_definidas"].mean()) / info["pct_mal_definidas"].std()
    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()}

    # A e C: univariadas (principal) e com efeitos fixos de estado (sensibilidade)
    uni = []
    for g, rotulo in DESFECHOS.items():
        for v in VARIAVEIS + EXTRAS:
            for fe in (False, True):
                uni.append({"desfecho": g, "rotulo": rotulo, "variavel": v, "efeito_fixo_uf": fe,
                            **univariada(dados[g], zs, info, v, efeito_uf=fe)})
    uni = pd.DataFrame(uni)
    uni["q_bh"] = np.nan
    principal = uni["variavel"].isin(VARIAVEIS)
    for _, idx in uni[principal].groupby(["desfecho", "efeito_fixo_uf"]).groups.items():
        uni.loc[idx, "q_bh"] = bh_fdr(uni.loc[idx, "p"].to_numpy())

    # C: teste formal da diferença esôfago versus comparação
    dif = pd.DataFrame([{"comparacao": comp, "variavel": v, **diferenca(dados["esofago"], dados[comp], zs, info, v)}
                        for comp in ("pulmao", "outros") for v in VARIAVEIS])
    dif["q_bh"] = dif.groupby("comparacao")["p_diferenca"].transform(lambda p: bh_fdr(p.to_numpy()))

    # Sensibilidade: regra antiga (ausente = 0)
    sens0 = pd.DataFrame([{"variavel": v, **univariada(dados["esofago"], zs0, info, v)} for v in VARIAVEIS])

    # B: modelos conjuntos
    conj = []
    for nome, cols in (("todas_as_10", VARIAVEIS + ["mal_definidas"]),
                       ("sem_lenha_e_mate_extraido",
                        [v for v in VARIAVEIS if v not in ("lenha", "mate_extracao")] + ["mal_definidas"])):
        for fe in (False, True):
            for r in conjunta(dados["esofago"], zs, info, cols, efeito_uf=fe):
                conj.append({"modelo": nome, "efeito_fixo_uf": fe, **r})
    conj = pd.DataFrame(conj)

    # D: agrupamentos versus vizinhos
    mb = pd.read_csv(RESULTADOS / "varredura_municipios.csv")
    mb = mb[mb["grupo"] == "esofago"]
    ag = pd.read_csv(RESULTADOS / "varredura_agrupamentos.csv")
    ag = ag[ag["grupo"] == "esofago"].set_index("ordem")
    cent = centroides()
    cods = info.index.to_numpy()
    lat = np.radians([cent.get(info.loc[c, "cod_ibge"], (np.nan, np.nan))[0] for c in cods])
    lon = np.radians([cent.get(info.loc[c, "cod_ibge"], (np.nan, np.nan))[1] for c in cods])
    em_agrup = set(mb["cod_ibge"] // 10)
    pos = {c: i for i, c in enumerate(cods)}
    dp = tr[VARIAVEIS].std()
    comp = []
    for ordem, grupo_m in mb.groupby("ordem"):
        dentro = list(grupo_m["cod_ibge"] // 10)
        idx = [pos[c] for c in dentro]
        h = (np.sin((lat[:, None] - lat[idx][None, :]) / 2) ** 2
             + np.cos(lat[:, None]) * np.cos(lat[idx][None, :]) * np.sin((lon[:, None] - lon[idx][None, :]) / 2) ** 2)
        dist = (6371 * 2 * np.arcsin(np.sqrt(h))).min(axis=1)
        ufs = set(info.loc[dentro, "uf_sigla"])
        viz = [c for c, dd in zip(cods, dist) if dd <= RAIO_VIZINHOS_KM and c not in em_agrup
               and info.loc[c, "uf_sigla"] in ufs]
        for v in VARIAVEIS:
            a, b_ = tr.loc[dentro, v].dropna(), tr.loc[viz, v].dropna()
            comp.append({"ordem": ordem, "centro": ag.loc[ordem, "centro"], "uf": ag.loc[ordem, "uf_centro"],
                         "municipios_dentro_com_dado": len(a), "vizinhos_com_dado": len(b_), "variavel": v,
                         "media_dentro": a.mean(), "media_vizinhos": b_.mean(),
                         "diferenca_padronizada": (a.mean() - b_.mean()) / dp[v] if len(a) and len(b_) else np.nan})
    comp = pd.DataFrame(comp)
    res = []
    for v, d in comp.groupby("variavel"):
        x = d["diferenca_padronizada"].dropna()
        p_, n_ = int((x > 1e-12).sum()), int((x < -1e-12).sum())
        res.append({"variavel": v, "agrupamentos_com_dado": len(x), "mediana_diferenca_padronizada": x.median(),
                    "positivos": p_, "negativos": n_, "empates": len(x) - p_ - n_,
                    "p_teste_do_sinal": stats.binomtest(p_, p_ + n_, 0.5).pvalue if p_ + n_ else np.nan})
    res = pd.DataFrame(res)
    res["q_bh"] = bh_fdr(res["p_teste_do_sinal"].fillna(1).to_numpy())
    res = res.sort_values("mediana_diferenca_padronizada", ascending=False)

    # E: os 16 municípios versus a mediana do estado
    sinais = rmp[(rmp["grupo"] == "esofago") & rmp["sinal"] & (rmp["contexto"] == "prioritário")]["cod6"]
    bruto = exp[VARIAVEIS + ["agua_poco_nascente_2010", "rural_2000"]]
    med_uf = bruto.groupby(info["uf_sigla"]).median()
    e16 = []
    for c in sinais:
        uf = info.loc[c, "uf_sigla"]
        linha = {"municipio": info.loc[c, "municipio"], "uf": uf}
        for v in bruto.columns:
            linha[v] = bruto.loc[c, v]
            linha[f"{v}_mediana_uf"] = med_uf.loc[uf, v]
        e16.append(linha)
    e16 = pd.DataFrame(e16)

    # Histórico: correlação entre as medidas de 2022 e as anteriores
    hist = pd.DataFrame([{"par": f"{a} x {b}", "spearman": exp[[a, b]].dropna().corr(method="spearman").iloc[0, 1]}
                         for a, b in (("rural", "rural_2010"), ("rural", "rural_2000"),
                                      ("agua_poco", "agua_poco_nascente_2010"), ("agua_nascente", "agua_poco_nascente_2010"))])

    # F: por macrorregião (não pré-especificado): univariadas e conjunto sem lenha
    reg = []
    for nome, ufs in REGIOES.items():
        sel = info.index[info["uf_sigla"].isin(ufs)]
        for v in VARIAVEIS:
            if (exp.loc[sel, v] > 0).sum() < 20:
                continue
            for g in ("esofago", "pulmao"):
                reg.append({"regiao": nome, "modelo": "univariada", "desfecho": g, "variavel": v,
                            **univariada(dados[g], zs, info, v, sel=sel)})
        cols = [v for v in VARIAVEIS if v != "lenha" and (exp.loc[sel, v] > 0).sum() >= 20] + ["mal_definidas"]
        for r in conjunta(dados["esofago"], zs, info, cols, sel=sel):
            reg.append({"regiao": nome, "modelo": "conjunto", "desfecho": "esofago", **r})
    reg = pd.DataFrame(reg)

    saidas = {"fase2_univariada": uni, "fase2_diferenca_esofago_comparacao": dif, "fase2_sens_ausente_zero": sens0,
              "fase2_conjunto": conj, "fase2_agrupamentos_vs_vizinhos": comp, "fase2_agrupamentos_resumo": res,
              "fase2_16_municipios": e16, "fase2_historico_correlacoes": hist, "fase2_por_regiao_nao_preespecificado": reg}
    for nome, df in saidas.items():
        df.to_csv(RESULTADOS / f"{nome}.csv", index=False, float_format="%.4g")
    tr.corr(method="spearman").to_csv(RESULTADOS / "fase2_correlacoes.csv", float_format="%.3f")
    for antigo in ("fase2_multivariada.csv", "fase2_por_regiao_conjunto_nao_preespecificado.csv"):
        (RESULTADOS / antigo).unlink(missing_ok=True)

    pd.set_option("display.width", 250, "display.max_columns", 40)
    e = uni[uni["desfecho"] == "esofago"]
    print("A. Esôfago, univariadas (colunas: sem / com efeitos fixos de UF)")
    print(e.pivot(index="variavel", columns="efeito_fixo_uf", values="rt_por_dp").join(
        e[~e["efeito_fixo_uf"]].set_index("variavel")[["municipios", "ic95_inf", "ic95_sup", "q_bh"]]).round(3).to_string())
    print("\nC. RT por desfecho (sem efeito fixo)")
    print(uni[~uni["efeito_fixo_uf"]].pivot(index="variavel", columns="desfecho", values="rt_por_dp").round(3).to_string())
    print("\nC. Diferença esôfago versus comparação")
    print(dif.round(4).to_string(index=False))
    print("\nSensibilidade ausente = 0 (esôfago)")
    print(sens0.round(3).to_string(index=False))
    print("\nB. Conjuntos")
    print(conj.pivot_table(index="variavel", columns=["modelo", "efeito_fixo_uf"], values="rt_por_dp").round(3).to_string())
    print(conj.groupby(["modelo", "efeito_fixo_uf"])["municipios"].first().to_string())
    print("\nD. Agrupamentos versus vizinhos")
    print(res.round(4).to_string(index=False))
    print("\nHistórico")
    print(hist.round(3).to_string(index=False))
    print("\nF. Por região")
    print(reg.pivot_table(index="variavel", columns=["regiao", "modelo", "desfecho"], values="rt_por_dp").round(3).to_string())


if __name__ == "__main__":
    main()