"""Fase 1b, parte C: tendências anuais (2015–2024) da mortalidade por câncer de esôfago.

Plano registrado em PLANO-fase1b.md (antes de rodar). Resumo:
- óbitos por ano, município, sexo e faixa (contar_anual.py, mesmas regras de preparar.py);
- população por ano: interpolação linear entre os Censos 2010 e 2022 para 2015–2021; Censo 2022 para
  2022–2024 (municípios criados depois de 2010 usam 2022 em todos os anos);
- esperados em relação ao próprio estado NO MESMO ANO (a tendência é relativa ao estado);
- unidades: os 16 municípios com sinal (por região e juntos) e os agrupamentos de esôfago de varredura.py;
- RMP por ano e regressão de Poisson (log dos esperados como offset): a inclinação estima a variação
  anual da RMP. Descritivo, sem limiar de decisão.

Saídas: resultados/tendencias_anuais.csv e resultados/tendencias_inclinacao.csv
Uso: python tendencias.py   (depois de contar_anual.py e varredura.py)
"""
import numpy as np
import pandas as pd

from config import BRUTOS, PROCESSADOS, RESULTADOS

ANOS = list(range(2015, 2025))
GRUPO = "esofago"
REGIOES = {"Minas Gerais": ["MG"], "Sul (PR, SC, RS)": ["PR", "SC", "RS"], "Bahia": ["BA"]}


def populacao_anual():
    p22 = pd.read_csv(BRUTOS / "populacao_censo2022.csv")
    p10 = pd.read_csv(BRUTOS / "populacao_censo2010.csv").rename(columns={"pop": "pop10"})
    p = p22.merge(p10, on=["cod_ibge", "sexo", "faixa"], how="left")
    sem = p["pop10"].isna()
    partes = []
    for ano in ANOS:
        f = min(ano, 2022) - 2010
        q = p[["cod_ibge", "sexo", "faixa"]].copy()
        q["ano"] = ano
        q["pop"] = np.where(sem, p["pop"], p["pop10"] + f / 12 * (p["pop"] - p["pop10"])).clip(min=0)
        partes.append(q)
    pop = pd.concat(partes, ignore_index=True)
    pop["cod6"] = pop["cod_ibge"] // 10
    pop["uf"] = pop["cod_ibge"] // 100000
    return pop


def esperados(pop, ob):
    """Observados e esperados (taxas do estado no mesmo ano, por sexo e faixa) por município e ano."""
    ob = ob[ob["cod6"].isin(set(pop["cod6"]))].copy()
    ob["uf"] = ob["cod6"] // 10000
    taxa = (ob.groupby(["ano", "uf", "sexo", "faixa"])["obitos"].sum()
            / pop.groupby(["ano", "uf", "sexo", "faixa"])["pop"].sum()).fillna(0).rename("taxa")
    p = pop.merge(taxa, left_on=["ano", "uf", "sexo", "faixa"], right_index=True, how="left")
    p["esp"] = p["pop"] * p["taxa"].fillna(0)
    e = p.groupby(["cod6", "ano"])["esp"].sum()
    o = ob.groupby(["cod6", "ano"])["obitos"].sum().reindex(e.index).fillna(0)
    return pd.DataFrame({"obs": o, "esp": e}).reset_index()


def poisson_inclinacao(anos, obs, esp):
    """Regressão de Poisson log(μ) = log(esp) + a + b·(ano − média); Newton-Raphson. Devolve b e EP."""
    x = np.column_stack([np.ones(len(anos)), anos - np.mean(anos)])
    beta = np.array([np.log(obs.sum() / esp.sum()), 0.0])
    for _ in range(50):
        mu = esp * np.exp(x @ beta)
        grad = x.T @ (obs - mu)
        hess = x.T @ (x * mu[:, None])
        passo = np.linalg.solve(hess, grad)
        beta += passo
        if np.abs(passo).max() < 1e-10:
            break
    mu = esp * np.exp(x @ beta)
    cov = np.linalg.inv(x.T @ (x * mu[:, None]))
    # sobredispersão (Pearson): o EP é inflado quando há mais variação que a de Poisson
    phi = max(1.0, float(((obs - mu) ** 2 / mu).sum() / (len(obs) - 2)))
    return beta[1], np.sqrt(cov[1, 1] * phi), phi


def main():
    pop = populacao_anual()
    ob = pd.read_csv(PROCESSADOS / "obitos_anuais.csv")
    ob = ob[ob["grupo"] == GRUPO]
    oe = esperados(pop, ob)

    sim = pd.read_csv(RESULTADOS / "municipios_rmp.csv")
    sinais = sim[(sim["grupo"] == GRUPO) & sim["sinal"] & (sim["contexto"] == "prioritário")]
    unidades = {f"16 sinais — {reg}": sinais[sinais["uf_sigla"].isin(ufs)]["cod6"] for reg, ufs in REGIOES.items()}
    unidades["16 sinais — todos"] = sinais["cod6"]
    for _, r in sinais.sort_values(["uf_sigla", "municipio"]).iterrows():
        unidades[f"município — {r['municipio']} ({r['uf_sigla']})"] = [r["cod6"]]
    ag = pd.read_csv(RESULTADOS / "varredura_agrupamentos.csv")
    mb = pd.read_csv(RESULTADOS / "varredura_municipios.csv")
    for _, a in ag[ag["grupo"] == GRUPO].iterrows():
        cods = mb[(mb["grupo"] == GRUPO) & (mb["ordem"] == a["ordem"])]["cod_ibge"] // 10
        unidades[f"agrupamento {a['ordem']:02d} — {a['centro']} ({a['uf_centro']})"] = cods
    todos_ag = mb[mb["grupo"] == GRUPO]["cod_ibge"] // 10
    unidades["agrupamentos — todos"] = todos_ag

    anuais, incl = [], []
    for nome, cods in unidades.items():
        d = oe[oe["cod6"].isin(set(cods))].groupby("ano")[["obs", "esp"]].sum().reindex(ANOS)
        d["rmp"] = d["obs"] / d["esp"]
        for ano, r in d.iterrows():
            anuais.append({"unidade": nome, "ano": ano, "obitos": int(r["obs"]), "esperados": r["esp"], "rmp": r["rmp"]})
        b, ep, phi = poisson_inclinacao(np.array(ANOS, float), d["obs"].to_numpy(float), d["esp"].to_numpy(float))
        p1 = d.loc[2015:2019]
        p2 = d.loc[2020:2024]
        incl.append({"unidade": nome, "municipios": len(set(cods)),
                     "rmp_2015_19": p1["obs"].sum() / p1["esp"].sum(), "rmp_2020_24": p2["obs"].sum() / p2["esp"].sum(),
                     "variacao_anual_pct": 100 * (np.exp(b) - 1),
                     "ic95_inf_pct": 100 * (np.exp(b - 1.96 * ep) - 1), "ic95_sup_pct": 100 * (np.exp(b + 1.96 * ep) - 1),
                     "sobredispersao": phi, "anos_com_rmp_acima_de_1": int((d["rmp"] > 1).sum())})
    anuais, incl = pd.DataFrame(anuais), pd.DataFrame(incl)
    anuais.to_csv(RESULTADOS / "tendencias_anuais.csv", index=False, float_format="%.4g")
    incl.to_csv(RESULTADOS / "tendencias_inclinacao.csv", index=False, float_format="%.4g")

    pd.set_option("display.width", 250, "display.max_columns", 30)
    print(anuais[anuais["unidade"].str.startswith("16")].pivot(index="unidade", columns="ano", values="rmp").round(2).to_string())
    print()
    print(incl.round(2).to_string(index=False))


if __name__ == "__main__":
    main()