"""Reproduz todas as análises do Cledon a partir dos dados públicos e confere os resultados com os publicados.

Roda, na ordem, as etapas descritas no README (download, agregação, análises das Fases 1 a 2b e 1c), confere se os
arquivos baixados são os mesmos usados na publicação (SHA-256) e compara cada tabela de resultados com a referência.

Uso (dentro da pasta com o código, com as dependências de requirements.txt instaladas):
  python reproduzir.py --referencia zenodo        # baixa os resultados publicados (Zenodo v1.7) e reproduz tudo
  python reproduzir.py --referencia resultados.zip  # o mesmo, com o zip de resultados já baixado (ou uma pasta)
  python reproduzir.py                            # sem referência: guarda uma cópia de resultados/ e compara com ela
  python reproduzir.py --fases fase1,validacao    # só algumas fases (as anteriores precisam já ter rodado)
  python reproduzir.py --conferir --referencia zenodo  # não roda nada; só compara resultados/ com a referência
  python reproduzir.py --relatorios               # também gera as páginas HTML (mapa, relatórios, preprint)
  python reproduzir.py --fases relatorios         # só as páginas HTML, a partir de resultados/ já prontos
  python reproduzir.py --listar                   # mostra as etapas sem executar

Tempo: cerca de 1 hora num computador comum, mais os downloads (~5 GB, conforme a rede do DATASUS).
Saída: dados/reproducao/<data-hora>/ com o registro de cada etapa (log.txt) e o relatório (relatorio.md).
O DATASUS republica arquivos corrigidos; quando um arquivo baixado difere do usado na publicação, o relatório avisa,
e diferenças nos resultados que dependem dele são esperadas.
"""
import argparse
import hashlib
import io
import os
import shutil
import subprocess
import sys
import time
import urllib.request
import zipfile
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd

RAIZ = Path(__file__).resolve().parent
sys.path.insert(0, str(RAIZ))
from config import BRUTOS, DADOS, RESULTADOS  # noqa: E402

ZENODO_DOI = "10.5281/zenodo.24195057"
ZENODO_RESULTADOS = "https://zenodo.org/records/24195057/files/cledon-fase1-resultados-v1.7.zip?download=1"

# (fase, descrição, comando). A ordem importa: cada fase usa saídas das anteriores.
ETAPAS = [
    ("fase1", "Download: municípios, malhas, Censo 2022, estimativas, SIM 2020–2024",
     ["baixar.py", "municipios", "populacao", "estimativas", "sim"]),
    ("fase1", "Agregação dos óbitos 2020–2024", ["preparar.py"]),
    ("fase1", "RMP, suavização, testes e sinais", ["analisar.py"]),
    ("fase1", "Análises de sensibilidade (~15 min)", ["sensibilidade.py"]),
    ("validacao", "Download: Censo 2010 e SIM 2015–2019 e 2020 em CSV (OpenDataSUS)",
     ["baixar.py", "validacao-csv", "equivalencia"]),
    ("validacao", "Agregação dos óbitos 2015–2019", ["preparar.py", "validacao"]),
    ("validacao", "Validação em período independente", ["validar.py"]),
    ("validacao", "Equivalência entre FTP e OpenDataSUS (2020)", ["equivalencia_fontes.py"]),
    ("validacao", "Tamanho efetivo do teste (complementar)", ["validar_complementar.py"]),
    ("fase1b", "Download: Registro Hospitalar de Câncer (INCA), com e sem SP", ["baixar.py", "rhc", "rhc-sp"]),
    ("fase1b", "Óbitos por ano, 2015–2024", ["contar_anual.py"]),
    ("fase1b", "Varredura espacial de Kulldorff (~10 min)", ["varredura.py"]),
    ("fase1b", "Registro hospitalar: razão de incidência proporcional", ["rhc.py"]),
    ("fase1b", "Registro hospitalar com SP (sensibilidade)", ["rhc.py", "com-sp"]),
    ("fase1b", "Tendências anuais", ["tendencias.py"]),
    ("fase2", "Exposições municipais (IBGE, altitude)", ["exposicoes.py"]),
    ("fase2", "Regressões ecológicas", ["hipoteses.py"]),
    ("fase2b", "Exposições de Minas Gerais (IBGE, CNES)", ["exposicoes_2b.py"]),
    ("fase2b", "Análises de Minas Gerais", ["hipoteses_2b.py"]),
    ("fase1c", "Download: SIM 1996–2014", ["baixar.py", "historico"]),
    ("fase1c", "Série histórica 1996–2014", ["historico.py"]),
]
RELATORIOS = [("relatorios", f"Gera {s}", [s]) for s in
              ("gerar_mapa.py", "gerar_relatorio.py", "gerar_fase1b.py", "gerar_fase2.py", "gerar_fase2b.py",
               "gerar_fase1c.py", "gerar_entenda.py", "gerar_preprint.py")]
FASES = list(dict.fromkeys(f for f, _, _ in ETAPAS))

# Tolerâncias da comparação numérica
RTOL, ATOL = 1e-6, 1e-9          # igual (diferenças de arredondamento de ponto flutuante)
RTOL_ARRED = 1e-3                # só no último algarismo gravado (alguns CSVs têm 4 algarismos significativos)


def sha256(caminho):
    h = hashlib.sha256()
    with open(caminho, "rb") as f:
        while bloco := f.read(1 << 20):
            h.update(bloco)
    return h.hexdigest()


class Registro:
    def __init__(self, pasta):
        self.pasta = pasta
        self.log = open(pasta / "log.txt", "a", encoding="utf-8")

    def __call__(self, msg=""):
        print(msg, flush=True)
        self.log.write(msg + "\n")
        self.log.flush()


def preparar_referencia(arg, pasta, reg):
    """Devolve a pasta com os CSVs de referência (do Zenodo, de um zip, de uma pasta ou cópia dos atuais)."""
    destino = pasta / "referencia"
    if arg is None:
        csvs = sorted(RESULTADOS.glob("*.csv"))
        if not csvs:
            reg("Sem referência: resultados/ está vazia. Use --referencia zenodo para comparar com a publicação.")
            return None
        destino.mkdir()
        for c in csvs:
            shutil.copy2(c, destino / c.name)
        reg(f"Referência: cópia dos {len(csvs)} CSVs atuais de resultados/ (feita antes de rodar).")
        return destino
    if arg == "zenodo":
        reg(f"Referência: resultados publicados no Zenodo (DOI {ZENODO_DOI}); baixando...")
        with urllib.request.urlopen(ZENODO_RESULTADOS, timeout=300) as r:
            dados = r.read()
        origem = zipfile.ZipFile(io.BytesIO(dados))
    else:
        p = Path(arg).expanduser().resolve()
        if p.is_dir():
            destino.mkdir()
            for c in sorted(p.rglob("*.csv")):
                shutil.copy2(c, destino / c.name)
            reg(f"Referência: pasta {p}.")
            return destino
        origem = zipfile.ZipFile(p)
        reg(f"Referência: arquivo {p.name}.")
    destino.mkdir()
    n = 0
    for nome in origem.namelist():
        if nome.endswith(".csv") and not nome.startswith("__MACOSX"):
            (destino / Path(nome).name).write_bytes(origem.read(nome))
            n += 1
    reg(f"  {n} CSVs de referência.")
    return destino


def rodar(etapas, reg):
    for k, (fase, desc, cmd) in enumerate(etapas, 1):
        reg(f"\n[{k}/{len(etapas)}] {fase} · {desc}\n    $ python {' '.join(cmd)}")
        t0 = time.time()
        # Os relatórios citam o DOI da versão publicada (na publicação, ele vem da variável CLEDON_DOI).
        env = dict(os.environ, CLEDON_DOI=os.environ.get("CLEDON_DOI", ZENODO_DOI))
        proc = subprocess.Popen([sys.executable, *cmd], cwd=RAIZ, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                text=True, encoding="utf-8", errors="replace", env=env)
        ultimas = []
        for linha in proc.stdout:
            reg.log.write("    " + linha)
            ultimas = (ultimas + [linha.rstrip()])[-15:]
        proc.wait()
        reg.log.flush()
        dur = time.time() - t0
        if proc.returncode != 0:
            reg("    FALHOU (código %d) após %.0f s. Últimas linhas:" % (proc.returncode, dur))
            for l in ultimas:
                reg("      " + l)
            reg("    Se foi falha de rede do DATASUS ou do IBGE, rode de novo: arquivos já baixados são reaproveitados.")
            return False
        reg("    ok (%.0f s)" % dur)
    return True


def conferir_entradas(ref, reg):
    """Compara o SHA-256 dos arquivos baixados com os registrados na publicação."""
    arq = (ref / "hashes_entradas.csv") if ref else None
    if not arq or not arq.exists():
        reg("Sem hashes_entradas.csv na referência; conferência das entradas não realizada.")
        return None
    h = pd.read_csv(arq)
    linhas = []
    for r in h.itertuples():
        p = BRUTOS / r.arquivo
        estado = "ausente" if not p.exists() else ("igual" if sha256(p) == r.sha256 else "diferente")
        linhas.append({"uso": r.uso, "arquivo": r.arquivo, "estado": estado})
    df = pd.DataFrame(linhas)
    resumo = df.pivot_table(index="uso", columns="estado", values="arquivo", aggfunc="count", fill_value=0)
    return df, resumo


def _num(s):
    return pd.to_numeric(s, errors="coerce")


def comparar_tabela(a, b):
    """Compara dois CSVs (novo a, referência b). Devolve (estado, detalhe)."""
    if a.read_bytes() == b.read_bytes():
        return "idêntico", ""
    try:
        x, y = pd.read_csv(a, low_memory=False), pd.read_csv(b, low_memory=False)
    except Exception as e:  # noqa: BLE001
        return "diferente", f"não foi possível ler: {e}"
    if list(x.columns) != list(y.columns):
        faltam, sobram = set(y.columns) - set(x.columns), set(x.columns) - set(y.columns)
        return "diferente", f"colunas diferentes (faltam {sorted(faltam)}, novas {sorted(sobram)})"
    if len(x) != len(y):
        return "diferente", f"{len(x)} linhas, referência {len(y)}"

    def medir(x, y):
        pior, col_pior, texto = 0.0, "", []
        for c in x.columns:
            nx, ny = _num(x[c]), _num(y[c])
            numerica = nx.notna().sum() == x[c].notna().sum() and ny.notna().sum() == y[c].notna().sum()
            if numerica and (nx.notna().any() or ny.notna().any()):
                vx, vy = nx.to_numpy(float), ny.to_numpy(float)
                if not np.array_equal(np.isnan(vx), np.isnan(vy)):
                    texto.append(c)
                    continue
                ok = ~np.isnan(vx)
                if ok.any():
                    dif = np.abs(vx[ok] - vy[ok]) / (np.abs(vy[ok]) + ATOL)
                    dif[np.abs(vx[ok] - vy[ok]) <= ATOL] = 0
                    if dif.max() > pior:
                        pior, col_pior = float(dif.max()), c
            elif not x[c].astype(str).equals(y[c].astype(str)):
                texto.append(c)
        return pior, col_pior, texto

    pior, col, texto = medir(x, y)
    reordenado = False
    if texto:  # talvez só a ordem das linhas tenha mudado
        chave = list(x.columns)
        xs = x.astype(str).sort_values(chave).index
        ys = y.astype(str).sort_values(chave).index
        p2, c2, t2 = medir(x.loc[xs].reset_index(drop=True), y.loc[ys].reset_index(drop=True))
        if not t2:
            pior, col, texto, reordenado = p2, c2, [], True
    if texto:
        return "diferente", f"valores não numéricos diferentes em {texto[:5]}"
    ordem = "; linhas em outra ordem" if reordenado else ""
    if pior <= RTOL:
        return "igual (arredondamento)", (f"maior diferença relativa {pior:.1e}" if pior
                                          else "mesmos valores, outra formatação") + ordem
    if pior <= RTOL_ARRED:
        return "igual no último algarismo", f"maior diferença relativa {pior:.1e} em {col}" + ordem
    return "diferente", f"maior diferença relativa {pior:.1e} em {col}"


def conferir_resultados(ref, reg):
    linhas = []
    for b in sorted(ref.glob("*.csv")):
        if b.name == "hashes_entradas.csv":
            continue  # registro das entradas, conferido à parte
        a = RESULTADOS / b.name
        if not a.exists():
            linhas.append({"tabela": b.name, "estado": "ausente", "detalhe": "não foi gerada"})
            continue
        estado, detalhe = comparar_tabela(a, b)
        linhas.append({"tabela": b.name, "estado": estado, "detalhe": detalhe})
    return pd.DataFrame(linhas)


def md_tabela(df):
    cab = "| " + " | ".join(df.columns) + " |\n|" + "---|" * len(df.columns) + "\n"
    return cab + "".join("| " + " | ".join(str(v) for v in r) + " |\n" for r in df.itertuples(index=False))


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--referencia", help="'zenodo', um zip de resultados ou uma pasta com os CSVs publicados")
    ap.add_argument("--fases", help="fases a rodar, separadas por vírgula: " + ",".join(FASES)
                    + " (ou 'relatorios', só as páginas HTML)")
    ap.add_argument("--conferir", action="store_true", help="não roda as análises; só compara")
    ap.add_argument("--relatorios", action="store_true", help="gera também as páginas HTML")
    ap.add_argument("--listar", action="store_true", help="mostra as etapas e sai")
    a = ap.parse_args()

    fases = FASES if not a.fases else [f.strip() for f in a.fases.split(",")]
    invalidas = [f for f in fases if f not in FASES + ["relatorios"]]
    if invalidas:
        sys.exit(f"Fases desconhecidas: {invalidas}. Use: {', '.join(FASES)}")
    a.relatorios = a.relatorios or "relatorios" in fases
    fases = [f for f in fases if f != "relatorios"]
    etapas = [e for e in ETAPAS if e[0] in fases] + (RELATORIOS if a.relatorios else [])
    if a.listar:
        for k, (f, d, c) in enumerate(etapas, 1):
            print(f"{k:2d}. [{f}] {d}\n      python {' '.join(c)}")
        return 0

    pasta = DADOS / "reproducao" / datetime.now().strftime("%Y-%m-%d_%H%M%S")
    pasta.mkdir(parents=True)
    reg = Registro(pasta)
    reg(f"Cledon · reprodução das análises · {datetime.now():%d/%m/%Y %H:%M}")
    reg(f"Python {sys.version.split()[0]} · pandas {pd.__version__} · numpy {np.__version__} · pasta {RAIZ}")
    livre = shutil.disk_usage(RAIZ).free / 1e9
    if not a.conferir and livre < 8:
        reg(f"Atenção: só {livre:.0f} GB livres; os dados brutos ocupam cerca de 5 GB.")
    RESULTADOS.mkdir(exist_ok=True)
    ref = preparar_referencia(a.referencia, pasta, reg)

    t0 = time.time()
    ok = True
    if not a.conferir:
        reg(f"\nEtapas: {len(etapas)} ({', '.join(fases)}{', relatórios' if a.relatorios else ''})")
        ok = rodar(etapas, reg)
    total = (time.time() - t0) / 60

    rel = [f"# Reprodução das análises do Cledon\n",
           f"- Data: {datetime.now():%d/%m/%Y %H:%M}; duração {total:.0f} min",
           f"- Python {sys.version.split()[0]}, pandas {pd.__version__}, numpy {np.__version__}",
           f"- Fases: {', '.join(fases)}" + (" (só conferência, nada foi rodado)" if a.conferir else ""),
           f"- Referência: {a.referencia or 'cópia dos resultados anteriores'}",
           f"- Execução: {'concluída' if ok else 'INTERROMPIDA por falha (ver log.txt)'}\n"]
    codigo = 0 if ok else 2
    if ref:
        ent = conferir_entradas(ref, reg)
        if ent is not None:
            df_e, resumo = ent
            rel += ["## Arquivos de entrada (SHA-256 contra os da publicação)\n", md_tabela(resumo.reset_index()), ""]
            dif = df_e[df_e["estado"] == "diferente"]
            if len(dif):
                rel += ["Arquivos diferentes dos usados na publicação (o órgão pode ter republicado versões corrigidas; "
                        "resultados que dependem deles podem mudar):\n"] + [f"- {r}" for r in dif["arquivo"]] + [""]
        res = conferir_resultados(ref, reg)
        cont = res["estado"].value_counts()
        rel += ["## Tabelas de resultados\n", md_tabela(cont.rename_axis("estado").reset_index(name="tabelas")), "",
                md_tabela(res), ""]
        problemas = res[res["estado"].isin(["diferente", "ausente"])]
        if len(problemas) and ok:
            codigo = 1
        reg("\n" + "=" * 70)
        reg("RESULTADO DA CONFERÊNCIA")
        if ent is not None:
            reg("Entradas (arquivos baixados x publicação):")
            reg(resumo.to_string())
        reg("\nTabelas de resultados:")
        reg(cont.to_string())
        for r in problemas.itertuples():
            reg(f"  {r.estado}: {r.tabela} — {r.detalhe}")
        conclusao = ("Todas as tabelas reproduzidas" if not len(problemas) and ok
                     else f"{len(problemas)} tabela(s) com diferença ou ausentes" if ok else "Execução interrompida")
        reg(f"\n{conclusao}.")
        rel.append(f"**Conclusão:** {conclusao}.\n")
    rel.append("Critérios: 'idêntico' = mesmo arquivo byte a byte; 'igual (arredondamento)' = diferença relativa "
               f"≤ {RTOL:g}; 'igual no último algarismo' = ≤ {RTOL_ARRED:g}, compatível com o arredondamento gravado; "
               "'diferente' = acima disso.")
    (pasta / "relatorio.md").write_text("\n".join(rel), encoding="utf-8")
    reg(f"\nRelatório: {pasta / 'relatorio.md'}\nRegistro completo: {pasta / 'log.txt'}")
    return codigo


if __name__ == "__main__":
    sys.exit(main())