"""Lê os arquivos brutos mais recentes e grava o CSV longo do tema.

Uso:
    python3 scripts/tratar.py --tema economia

Saídas:
    dados/tratados/<tema>.csv              indicador,data,valor,unidade,fonte
    dados/tratados/<tema>_pares_paises.csv valores por país usados nas medianas
    seção "Extração" de cada ficha         arquivo bruto, data, primeira e última observação

Datas: AAAA (anual), AAAA-MM (mensal). O mês corrente, incompleto, fica fora;
séries anuais só têm anos completos; valores do FMI a partir do ano corrente são
projeções e ficam fora (D-017).
"""

import argparse
import csv
import re
import statistics
from collections import defaultdict

from comum import (HOJE, METODOLOGIA, RAIZ, TRATADOS, brutos_mais_recentes, consultas,
                   escrever_csv, ler_fichas, ler_json, ler_manifesto)
from coletar import nome_sidra

MES_CORRENTE = f"{HOJE.year:04d}-{HOJE.month:02d}"
ANO_CORRENTE = HOJE.year


# ---------------------------------------------------------------- leitores

def ler_sgs(codigo):
    """Devolve ({data: valor}, frequência, arquivos). Data em AAAA, AAAA-MM ou AAAA-MM-DD."""
    arquivos = brutos_mais_recentes("bcb_sgs", f"sgs{codigo}")
    if not arquivos:
        raise FileNotFoundError(f"sem bruto para SGS {codigo}; rode coletar.py")
    pontos = {}
    for arq in arquivos:
        for obs in ler_json(arq):
            d, m, a = obs["data"].split("/")
            if obs["valor"] in ("", None):
                continue
            pontos[f"{a}-{m}-{d}"] = float(obs["valor"])
    datas = list(pontos)
    if all(d.endswith("-01-01") for d in datas):
        freq, pontos = "anual", {d[:4]: v for d, v in pontos.items()}
    elif all(d.endswith("-01") for d in datas):
        freq, pontos = "mensal", {d[:7]: v for d, v in pontos.items()}
    else:
        freq = "diaria"
    return dict(sorted(pontos.items())), freq, arquivos


def ler_sidra(caminho):
    """Devolve {(codigo_variavel, periodo): valor} e os arquivos. Período: AAAA, AAAAMM ou AAAATT."""
    arquivos = brutos_mais_recentes("ibge_sidra", nome_sidra(caminho))
    if not arquivos:
        raise FileNotFoundError(f"sem bruto para SIDRA {caminho}; rode coletar.py")
    linhas = ler_json(arquivos[-1])
    cab = linhas[0]
    col_var = next(k[:-1] + "C" for k, v in cab.items() if v == "Variável")
    col_per = next(k[:-1] + "C" for k, v in cab.items() if v in ("Ano", "Mês", "Trimestre"))
    tipo_per = next(v for v in cab.values() if v in ("Ano", "Mês", "Trimestre"))
    pontos = {}
    for linha in linhas[1:]:
        try:
            pontos[(linha[col_var], linha[col_per])] = float(linha["V"])
        except ValueError:  # "...", "-", "X": sem dado
            continue
    return pontos, tipo_per, arquivos


def ler_sdmx(arquivo):
    """SDMX-JSON 3.0 do FMI. Devolve {(dim1, dim2, ...): {periodo: valor}} e os nomes das dimensões."""
    d = ler_json(arquivo)
    estrutura = d["data"]["structures"][0]
    dims = estrutura["dimensions"]["series"]
    periodos = [v["value"] for v in estrutura["dimensions"]["observation"][0]["values"]]
    series = {}
    for chave, s in d["data"]["dataSets"][0]["series"].items():
        idx = [int(i) for i in chave.split(":")]
        rotulo = tuple(dims[i]["values"][j]["id"] for i, j in enumerate(idx))
        obs = {}
        for k, v in s.get("observations", {}).items():
            if v and v[0] not in (None, ""):
                obs[periodos[int(k)]] = float(v[0])
        series[rotulo] = obs
    return series, [dim["id"] for dim in dims]


# ---------------------------------------------------------------- transformações

def fim_de_mes(diarios):
    """Série diária -> valor do último dia disponível de cada mês, até hoje."""
    mensal = {}
    for data, valor in sorted(diarios.items()):
        if data <= HOJE.isoformat():
            mensal[data[:7]] = valor
    return mensal


def sem_mes_corrente(mensal):
    return {d: v for d, v in mensal.items() if d < MES_CORRENTE}


def mes_anterior(aaaa_mm, n=1):
    a, m = int(aaaa_mm[:4]), int(aaaa_mm[5:])
    total = a * 12 + (m - 1) - n
    return f"{total // 12:04d}-{total % 12 + 1:02d}"


def mediana_por_ano(valores_por_pais, paises, minimo=None):
    """{pais: {ano: v}} -> ({ano: mediana}, {ano: n}).

    A mediana de um ano só é publicada se pelo menos metade dos países do grupo tiver dado
    naquele ano (D-022); com menos, o grupo comparado seria outro."""
    minimo = minimo if minimo is not None else -(-len(paises) // 2)
    anos = sorted({a for p in paises for a in valores_por_pais.get(p, {})})
    med, n = {}, {}
    for ano in anos:
        vals = [valores_por_pais[p][ano] for p in paises if ano in valores_por_pais.get(p, {})]
        if vals and len(vals) >= minimo:
            med[ano] = statistics.median(vals)
            n[ano] = len(vals)
    return med, n


# ---------------------------------------------------------------- por indicador

class Saida:
    def __init__(self):
        self.linhas = []
        self.paises = []
        self.extracao = defaultdict(list)  # id da ficha -> [(série, arquivos, dados)]

    def add(self, ficha, serie_id, dados, unidade=None, fonte=None, arquivos=(), casas=4):
        unidade = unidade or ficha["unidade"]
        fonte = fonte or ficha["codigo_serie"].split(";")[0]
        for data, valor in sorted(dados.items()):
            self.linhas.append([serie_id, data, round(valor, casas), unidade, fonte])
        self.extracao[ficha["id"]].append((serie_id, list(arquivos), dados))


def tratar_sgs_simples(ficha, saida):
    for codigo in consultas(ficha):
        pontos, freq, arqs = ler_sgs(codigo)
        if freq == "diaria":
            pontos = fim_de_mes(pontos)
        if freq != "anual":
            pontos = sem_mes_corrente(pontos)
        serie_id = ficha["id"]
        if ficha["id"] == "divida_bruta_gg" and codigo == "4537":
            serie_id = "divida_bruta_gg_metodologia_antiga"
        if ficha.get("transformacao", "").startswith("multiplicar por -1"):
            pontos = {d: -v for d, v in pontos.items()}
        saida.add(ficha, serie_id, pontos, fonte=f"BCB/SGS {codigo}", arquivos=arqs)


def tratar_pib(ficha, saida):
    pontos, _, arqs = ler_sidra(consultas(ficha)[0])
    serie = {p: v for (var, p), v in pontos.items() if var == "6563"}
    anual = {p[:4]: v for p, v in serie.items() if p.endswith("04")}
    saida.add(ficha, ficha["id"], anual, fonte="IBGE/SIDRA 5932", arquivos=arqs)
    ultimo = max(serie)
    if not ultimo.endswith("04"):  # ano corrente, acumulado até o último trimestre divulgado
        mes_fim = f"{ultimo[:4]}-{int(ultimo[4:]) * 3:02d}"
        saida.add(ficha, "pib_variacao_real_acumulado_no_ano", {mes_fim: serie[ultimo]},
                  fonte="IBGE/SIDRA 5932 (acumulado no ano, parcial)", arquivos=arqs)
    return anual


def tratar_pib_per_capita(ficha, saida, pib_anual):
    c6784, c6579 = consultas(ficha)
    p6784, _, arq1 = ler_sidra(c6784)
    p6579, _, arq2 = ler_sidra(c6579)
    var = {a: v for (vv, a), v in p6784.items() if vv == "9814"}
    pop = {a: v * 1000 for (vv, a), v in p6784.items() if vv == "93"}  # mil pessoas -> pessoas
    saida.add(ficha, ficha["id"], var, fonte="IBGE/SIDRA 6784", arquivos=arq1)

    # D-016: anos sem Conta Nacional anual, a partir do PIB trimestral e da população estimada
    pop_estimada = {a: v for (vv, a), v in p6579.items() if vv == "9324"}
    preliminar = {}
    ano = int(max(var)) + 1
    while str(ano) in pib_anual and str(ano) in pop_estimada:
        pop_ant = pop.get(str(ano - 1)) or pop_estimada.get(str(ano - 1))
        g_pop = pop_estimada[str(ano)] / pop_ant - 1
        preliminar[str(ano)] = round(((1 + pib_anual[str(ano)] / 100) / (1 + g_pop) - 1) * 100, 1)
        pop[str(ano)] = pop_estimada[str(ano)]
        ano += 1
    if preliminar:
        saida.add(ficha, ficha["id"], preliminar,
                  fonte="cálculo preliminar sobre IBGE/SIDRA 5932 e 6579 (D-016)", arquivos=arq2)

    completo = {**var, **preliminar}
    indice, nivel = {}, 100.0
    indice["2002"] = nivel
    for a in sorted(completo):
        if int(a) > 2002:
            nivel *= 1 + completo[a] / 100
            indice[a] = nivel
    saida.add(ficha, "pib_per_capita_indice", indice, unidade="índice (2002 = 100)",
              fonte="cálculo do projeto sobre pib_per_capita_variacao_real", arquivos=[], casas=2)


def tratar_ipca(ficha, saida):
    pontos, _, arqs = ler_sidra(consultas(ficha)[0])

    def mensal(variavel):
        return sem_mes_corrente({f"{p[:4]}-{p[4:]}": v for (var, p), v in pontos.items()
                                 if var == variavel and p[:4] >= "2000"})

    ipca12 = mensal("2265")
    saida.add(ficha, ficha["id"], ipca12, fonte="IBGE/SIDRA 1737", arquivos=arqs)
    saida.add(ficha, "ipca_numero_indice", mensal("2266"), unidade="índice (dez/1993 = 100)",
              fonte="IBGE/SIDRA 1737", arquivos=arqs)
    meta, _, arqs_meta = ler_sgs(consultas(ficha, "consulta_extra_bcb_sgs")[0])
    saida.add(ficha, "ipca_meta", {a: v for a, v in meta.items() if a >= "2000"},
              unidade="% (meta do CMN)", fonte="BCB/SGS 13521", arquivos=arqs_meta)
    return ipca12


def tratar_juro_real(ficha, saida, ipca12):
    selic_mes, _, arqs = ler_sgs(consultas(ficha)[0])
    selic_mes = sem_mes_corrente(selic_mes)
    real = {}
    for mes in selic_mes:
        janela = [mes_anterior(mes, k) for k in range(12)]
        if all(m in selic_mes for m in janela) and mes in ipca12:
            fator = 1.0
            for m in janela:
                fator *= 1 + selic_mes[m] / 100
            real[mes] = (fator / (1 + ipca12[mes] / 100) - 1) * 100
    saida.add(ficha, ficha["id"], real, fonte="cálculo sobre BCB/SGS 4390 e IBGE/SIDRA 1737", arquivos=arqs)


def tratar_pcps(ficha, saida):
    arqs = brutos_mais_recentes("fmi_pcps", "pcps_" + consultas(ficha)[0].replace(".", "-"))
    series, _ = ler_sdmx(arqs[-1])
    (obs,) = series.values()
    mensal = {p.replace("-M", "-"): v for p, v in obs.items() if "-M" in p and p[:4] >= "2000"}
    saida.add(ficha, ficha["id"], sem_mes_corrente(mensal), fonte="FMI/PCPS PALLFNF", arquivos=arqs)


def tratar_fred(ficha, saida):
    serie = consultas(ficha)[0]
    arqs = brutos_mais_recentes("fred", f"fred_{serie}")
    with open(arqs[-1], encoding="utf-8") as f:
        mensal = {l["observation_date"][:7]: float(l[serie]) for l in csv.DictReader(f)
                  if l[serie] not in ("", ".") and l["observation_date"] >= "2000"}
    saida.add(ficha, ficha["id"], sem_mes_corrente(mensal), fonte=f"FRED {serie}", arquivos=arqs)


def ler_grupos():
    grupos = defaultdict(list)
    with open(METODOLOGIA / "pares.csv", encoding="utf-8") as f:
        for linha in csv.DictReader(f):
            grupos[linha["grupo"]].append(linha["iso3"])
    return grupos


def tratar_weo(ficha, saida):
    indicador = consultas(ficha)[0]
    arqs = brutos_mais_recentes("fmi_weo", f"weo_{indicador}")
    series, dims = ler_sdmx(arqs[-1])
    i_pais = dims.index("COUNTRY")
    por_pais = {}
    for rotulo, obs in series.items():
        anuais = {a: v for a, v in obs.items() if re.fullmatch(r"\d{4}", a) and "2000" <= a < str(ANO_CORRENTE)}
        por_pais[rotulo[i_pais]] = anuais

    if indicador == "NGDPRPPPPC":  # nível -> variação anual
        por_pais = {p: {a: (v / obs[str(int(a) - 1)] - 1) * 100 for a, v in obs.items()
                        if str(int(a) - 1) in obs}
                    for p, obs in por_pais.items()}

    fonte = f"FMI/WEO {indicador}"
    base = ficha["id"]
    brasil = por_pais.get("BRA", {})
    saida.add(ficha, f"{base}__brasil", brasil, fonte=fonte, arquivos=arqs)
    for grupo, paises in ler_grupos().items():
        med, n = mediana_por_ano(por_pais, paises)
        saida.add(ficha, f"{base}__mediana_{grupo}", med, fonte=f"{fonte}, mediana de {grupo}", arquivos=arqs)
        dif = {a: brasil[a] - med[a] for a in brasil if a in med}
        saida.add(ficha, f"{base}__brasil_menos_mediana_{grupo}", dif,
                  unidade="p.p.", fonte=f"{fonte}, Brasil menos mediana de {grupo}", arquivos=arqs)
        for a in sorted(n):
            saida.paises.append([f"{base}__n_{grupo}", grupo, a, n[a], "países", fonte])
    for agregado in consultas(ficha, "agregados"):
        if agregado == "G001":
            saida.add(ficha, f"{base}__mundo", por_pais.get("G001", {}), fonte=f"{fonte}, mundo (G001)",
                      arquivos=arqs)
    for pais, obs in sorted(por_pais.items()):
        for a, v in sorted(obs.items()):
            saida.paises.append([base, pais, a, round(v, 4), ficha["unidade"], fonte])


# ---------------------------------------------------------------- fichas

def atualizar_extracao(ficha, itens):
    """Reescreve o bloco entre os marcadores de extração da ficha."""
    manifestos = {}
    linhas = ["| Série | Arquivo bruto | Extraído em | Observações | Primeira | Última |",
              "|---|---|---|---|---|---|"]
    for serie_id, arquivos, dados in itens:
        nomes, datas = [], []
        for arq in arquivos:
            fonte = arq.parent.name
            manifestos.setdefault(fonte, ler_manifesto(fonte))
            nomes.append(f"`{arq.relative_to(RAIZ)}`")
            datas.append(manifestos[fonte].get(arq.name, {}).get("data_extracao", "?")[:16].replace("T", " "))
        datas_txt = ", ".join(sorted(set(datas))) or "(derivada)"
        nomes_txt = "<br>".join(nomes) or "(derivada)"
        if dados:
            linhas.append(f"| `{serie_id}` | {nomes_txt} | {datas_txt} | {len(dados)} | {min(dados)} | {max(dados)} |")
        else:
            linhas.append(f"| `{serie_id}` | {nomes_txt} | {datas_txt} | 0 | | |")
    bloco = "\n".join(linhas)
    caminho = ficha["_arquivo"]
    texto = caminho.read_text(encoding="utf-8")
    novo = re.sub(r"(<!-- extracao:inicio -->\n).*?(\n<!-- extracao:fim -->)",
                  lambda m: m.group(1) + bloco + m.group(2), texto, flags=re.S)
    caminho.write_text(novo, encoding="utf-8")


# ---------------------------------------------------------------- main

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--tema", required=True)
    args = ap.parse_args()

    proprio = RAIZ / "scripts" / f"coletar_{args.tema}.py"
    if proprio.exists():
        raise SystemExit(f"O tema {args.tema} tem script próprio: rode python3 scripts/{proprio.name} --so-tratar")
    fichas = {f["id"]: f for f in ler_fichas(args.tema)}
    saida = Saida()
    feitos = set()

    def rodar(ficha_id, funcao, *extra):
        feitos.add(ficha_id)
        return funcao(fichas[ficha_id], saida, *extra)

    # ordem importa onde um indicador usa outro
    if "pib_variacao_real" in fichas:
        pib_anual = rodar("pib_variacao_real", tratar_pib)
        if "pib_per_capita_variacao_real" in fichas:
            rodar("pib_per_capita_variacao_real", tratar_pib_per_capita, pib_anual)
    if "ipca_12m" in fichas:
        ipca12 = rodar("ipca_12m", tratar_ipca)
        if "juro_real_ex_post" in fichas:
            rodar("juro_real_ex_post", tratar_juro_real, ipca12)

    for ficha_id, ficha in fichas.items():
        if ficha_id in feitos:
            continue
        api = ficha.get("api")
        if api == "bcb_sgs":
            rodar(ficha_id, tratar_sgs_simples)
        elif api == "fmi_weo":
            rodar(ficha_id, tratar_weo)
        elif api == "fmi_pcps":
            rodar(ficha_id, tratar_pcps)
        elif api == "fred":
            rodar(ficha_id, tratar_fred)
        else:
            print(f"! sem tratamento para {ficha_id} (api {api})")

    saida.linhas.sort(key=lambda l: (l[0], l[1]))
    escrever_csv(TRATADOS / f"{args.tema}.csv", ["indicador", "data", "valor", "unidade", "fonte"], saida.linhas)
    if saida.paises:
        escrever_csv(TRATADOS / f"{args.tema}_pares_paises.csv",
                     ["indicador", "pais", "data", "valor", "unidade", "fonte"], saida.paises)
    for ficha_id, itens in saida.extracao.items():
        atualizar_extracao(fichas[ficha_id], itens)

    series = sorted({l[0] for l in saida.linhas})
    print(f"{len(saida.linhas)} linhas, {len(series)} séries -> dados/tratados/{args.tema}.csv")
    for s in series:
        datas = [l[1] for l in saida.linhas if l[0] == s]
        print(f"  {s:62s} {min(datas)} a {max(datas)} ({len(datas)})")


if __name__ == "__main__":
    main()