"""Fase 2b: baixa e monta as variáveis definidas em PLANO-fase2b.md (registrado antes dos dados).

IBGE/SIDRA: aguardente (Censo Agropecuário 2006, tabela 1724; 2017, tabela 6960), milho (PAM 1612, 2005–2010),
água de poço ou nascente em 1991 (tabela 179) e 2000 (tabela 1436), população rural em 1991 (tabela 202).
CNES (FTP do DATASUS): endoscópios digestivos (equipamento 0333) e hospitais habilitados em oncologia (CACON/UNACON).

Segurança dos dados: os arquivos do CNES descrevem estabelecimentos, não pacientes, mas trazem campos como CPF/CNPJ do
estabelecimento. Só são lidos os campos necessários (CNES, CODUFMUN, TIPEQUIP, CODEQUIP, QT_EXIST, SGRUPHAB); o arquivo
descompactado é temporário; o arquivo original é apagado depois de extraídos esses campos (o SHA-256 fica registrado em
dados/brutos/fase2b/cnes_hashes.csv); a saída tem só contagens por município.
Mesmas regras de ausentes da Fase 2: "-" = 0; "...", "X" e ausência na resposta = ausente.

Saídas: dados/brutos/fase2b/ (respostas e arquivos originais, fora do git) e resultados/fase2b_exposicoes.csv
Uso: python exposicoes_2b.py
"""
import hashlib
import json
import os
import tempfile
import time
import urllib.request

import datasus_dbc
import dbfread
import numpy as np
import pandas as pd
import requests

from analisar import centroides
from hipoteses import codigos_regiao
from config import BRUTOS, RESULTADOS, UFS

PASTA = BRUTOS / "fase2b"
CONSULTAS = {
    "aguardente_2006": "t/1724/n6/all/v/1986,1987/p/2006/c12528/111701/c218/0/c12517/0/c12763/0/c12764/0/c12617/0",
    "aguardente_2017": "t/6960/n6/all/v/10094/p/2017/c829/46302/c12528/111701/c218/46502/c12517/113601",
    "milho": "t/1612/n6/all/v/216/p/2005-2010/c81/2711",
    "agua_1991": "t/179/n6/all/v/96/p/1991/c1/0/c61/0,1032,1036",
    "agua_2000": "t/1436/n6/all/v/96/p/2000/c1/0/c61/0,92854",
    "rural_1991": "t/202/n6/all/v/93/p/1991/c2/0/c1/0,2",
}
FTP_CNES = "ftp://ftp.datasus.gov.br/dissemin/publicos/CNES/200508_/Dados/{tipo}/{tipo}{uf}{aamm}.dbc"
ONCOLOGIA = {"1701", "1702", "1703", "1706", "1707", "1708", "1709", "1712", "1713"}
CAMPOS_EQ = ["CNES", "CODUFMUN", "TIPEQUIP", "CODEQUIP", "QT_EXIST"]
CAMPOS_HB = ["CNES", "CODUFMUN", "SGRUPHAB"]


def sidra(nome):
    destino = PASTA / f"{nome}.json"
    if not destino.exists():
        for tentativa in range(4):
            try:
                r = requests.get(f"https://apisidra.ibge.gov.br/values/{CONSULTAS[nome]}", timeout=600)
                r.raise_for_status()
                destino.write_text(r.text)
                break
            except requests.RequestException:
                time.sleep(10 * (tentativa + 1))
        else:
            raise RuntimeError(f"falhou: {nome}")
    df = pd.DataFrame(json.loads(destino.read_text())[1:])
    df["cod_ibge"] = df["D1C"].astype(int)
    v = df["V"].str.strip()
    df["x"] = np.where(v == "-", 0.0, pd.to_numeric(v, errors="coerce"))   # "-" = 0; "...", "X" = ausente
    return df


def coluna(df, codigo):
    return next(c for c in df.columns if c.startswith("D") and c.endswith("C") and codigo in set(df[c]))


def cnes(tipo, aamm, campos):
    """Lê os arquivos do CNES de todas as UFs, só com os campos necessários (cache em CSV agregável)."""
    destino = PASTA / f"cnes_{tipo}_{aamm}.csv"
    if destino.exists():
        return pd.read_csv(destino, dtype=str)
    partes = []
    for uf in UFS:
        dbc = PASTA / f"{tipo}{uf}{aamm}.dbc"
        if not dbc.exists():
            for tentativa in range(4):
                try:
                    with urllib.request.urlopen(FTP_CNES.format(tipo=tipo, uf=uf, aamm=aamm), timeout=300) as r, open(dbc, "wb") as f:
                        f.write(r.read())
                    break
                except OSError:
                    time.sleep(10 * (tentativa + 1))
            else:
                raise RuntimeError(f"CNES {tipo}{uf}{aamm}")
        fd, tmp = tempfile.mkstemp(suffix=".dbf")
        os.close(fd)
        try:
            datasus_dbc.decompress(str(dbc), tmp)
            t = dbfread.DBF(tmp, encoding="latin-1", char_decode_errors="replace")
            partes.append(pd.DataFrame([{k: (r.get(k) or "") for k in campos} for r in t]))
        finally:
            os.remove(tmp)                      # nada descompactado fica em disco
        # registra o SHA-256 do arquivo original (para reprodução) e apaga-o: só os campos necessários ficam guardados
        h = hashlib.sha256(dbc.read_bytes()).hexdigest()
        with open(PASTA / "cnes_hashes.csv", "a") as f:
            f.write(f"{dbc.name},{dbc.stat().st_size},{h}\n")
        dbc.unlink()
        print(f"  CNES {tipo} {uf} {aamm}: ok", flush=True)
    df = pd.concat(partes, ignore_index=True)
    df.to_csv(destino, index=False)             # só os campos necessários
    return df


def main():
    PASTA.mkdir(parents=True, exist_ok=True)
    geo = json.loads((BRUTOS / "malha_municipios.geojson").read_text())
    base = pd.read_csv(RESULTADOS / "fase2_exposicoes.csv").set_index("cod_ibge")[["area_km2", "pop_2010", "pop_2022"]]
    pop_ref = base["pop_2010"].fillna(base["pop_2022"])
    rmp = pd.read_csv(RESULTADOS / "municipios_rmp.csv").drop_duplicates("cod_ibge").set_index("cod_ibge")
    base["regiao_imediata"] = pd.Series(codigos_regiao()).reindex(base.index)   # código oficial (há nomes repetidos)
    base["uf"] = rmp["uf_sigla"].reindex(base.index)

    # Aguardente (mil litros) por mil habitantes de 2010
    a = sidra("aguardente_2006")
    # duas parcelas (matéria-prima própria e adquirida): se alguma está sob sigilo ("X"), o total fica ausente
    a06 = a.groupby("cod_ibge")["x"].sum(min_count=2)
    base["aguardente"] = 1000 * 1000 * a06.reindex(base.index) / pop_ref          # litros por mil hab.
    # sensibilidade (não pré-especificada): há produção? "X" indica produção sob sigilo, portanto > 0
    v = a["V"].str.strip()
    a["produz"] = (v == "X") | (pd.to_numeric(v, errors="coerce") > 0)
    base["aguardente_produz"] = a.groupby("cod_ibge")["produz"].any().astype(float).reindex(base.index)
    a = sidra("aguardente_2017")
    base["aguardente_2017"] = 1000 * 1000 * a.set_index("cod_ibge")["x"].reindex(base.index) / base["pop_2022"]

    # Milho: área colhida média 2005–2010 como % da área municipal
    m = sidra("milho")
    m["ano"] = m[next(c for c in m.columns if c.startswith("D") and c.endswith("C") and set(m[c]) <= {str(x) for x in range(2005, 2011)})]
    milho = m.pivot_table(index="cod_ibge", columns="ano", values="x", aggfunc="first").reindex(base.index).mean(axis=1)
    base["milho"] = 100 * milho / (base["area_km2"] * 100)

    # Água de poço ou nascente em 1991 e 2000; população rural em 1991
    for nome, total, cats, col in (("agua_1991", "0", ["1032", "1036"], "agua_pn_1991"), ("agua_2000", "0", ["92854"], "agua_pn_2000")):
        d = sidra(nome)
        p = d.pivot_table(index="cod_ibge", columns=coluna(d, cats[0]), values="x", aggfunc="sum")
        base[col] = (100 * p[cats].sum(axis=1, min_count=1) / p[total]).reindex(base.index)
    r = sidra("rural_1991")
    p = r.pivot_table(index="cod_ibge", columns=coluna(r, "2"), values="x", aggfunc="sum")
    base["rural_1991"] = (100 * p["2"] / p["0"]).reindex(base.index)

    # CNES: endoscópios digestivos por 100 mil hab. na região imediata; distância até oncologia
    cent = centroides()
    for aamm, sufixo in (("1512", ""), ("2012", "_2020")):
        eq = cnes("EQ", aamm, CAMPOS_EQ)
        eq = eq[(eq["TIPEQUIP"] == "3") & (eq["CODEQUIP"] == "33")]
        eq["cod6"] = eq["CODUFMUN"].str[:6].astype(int)
        n_mun = eq.groupby("cod6")["QT_EXIST"].apply(lambda s: pd.to_numeric(s, errors="coerce").sum())
        por_mun = pd.Series(base.index // 10, index=base.index).map(n_mun).fillna(0)
        reg = pd.DataFrame({"n": por_mun, "pop": pop_ref, "regiao": base["regiao_imediata"]}).groupby("regiao")[["n", "pop"]].sum()
        base[f"endoscopia{sufixo}"] = base["regiao_imediata"].map(1e5 * reg["n"] / reg["pop"])
        if sufixo == "":
            hb = cnes("HB", aamm, CAMPOS_HB)
            onco = set(hb[hb["SGRUPHAB"].isin(ONCOLOGIA)]["CODUFMUN"].str[:6].astype(int))
            alvo = [c for c in base.index if c // 10 in onco and c in cent]
            la, lo = np.radians([cent[c][0] for c in alvo]), np.radians([cent[c][1] for c in alvo])
            dist = {}
            for c in base.index:
                if c not in cent:
                    continue
                y, x = np.radians(cent[c])
                h = np.sin((la - y) / 2) ** 2 + np.cos(y) * np.cos(la) * np.sin((lo - x) / 2) ** 2
                dist[c] = float((6371 * 2 * np.arcsin(np.sqrt(h))).min())
            base["dist_oncologia"] = pd.Series(dist).reindex(base.index)
            base["municipios_com_oncologia_2015"] = len(alvo)

    cols = ["aguardente", "aguardente_2017", "aguardente_produz", "milho", "agua_pn_1991", "agua_pn_2000", "rural_1991", "endoscopia",
            "endoscopia_2020", "dist_oncologia"]
    base[cols + ["regiao_imediata", "uf", "municipios_com_oncologia_2015"]].to_csv(RESULTADOS / "fase2b_exposicoes.csv", float_format="%.6g")
    pd.set_option("display.width", 200)
    print(base[cols].describe().T.round(3).to_string())
    print("Ausentes:", base[cols].isna().sum().to_dict())
    print("MG:", base.loc[base["uf"] == "MG", cols].describe().T[["count", "mean", "50%", "max"]].round(3).to_string())


if __name__ == "__main__":
    main()