"""Resume cada série por mandato: herdou, entregou e média, com e sem 2020–21.

Uso:
    python3 scripts/resumir.py --tema economia
    python3 scripts/resumir.py --tema politica_externa --entrada politica_externa_comercio.csv

Lê dados/tratados/<tema>.csv (ou o arquivo de --entrada) e grava, ao lado, o mesmo nome
com o sufixo _mandatos. Séries de acumulado parcial (com "_parcial" no id) não são resumidas.
Regras na metodologia, seção 3, e nas decisões D-006 e D-007:

- herdou: último período antes da posse (ano anterior ou dezembro do ano anterior; em séries
  com lacunas, a última observação até quatro anos antes da posse);
- entregou: último período do mandato, ou o último disponível (parcial);
- média: segue o campo `media_mandato` da ficha
    geometrica  taxa média composta dos valores anuais
    aritmetica  média simples dos valores anuais
    variacao    como aritmetica, mas o número principal é entregou − herdou
- séries mensais viram anuais pelo campo `anualizacao` (dezembro ou media);
- anos incompletos ficam fora das médias anuais.
"""

import argparse
import math
from collections import defaultdict

from comum import (HOJE, TRATADOS, anos_do_mandato, escrever_csv, ler_csv, ler_fichas,
                   ler_mandatos, meses_do_mandato)

ANOS_PANDEMIA = {"2020", "2021"}
SEM_RESUMO = {"ipca_numero_indice", "pib_variacao_real_acumulado_no_ano"}


def ficha_da_serie(serie, fichas):
    if serie in fichas:
        return fichas[serie]
    base = serie.split("__")[0]
    if base in fichas:
        return fichas[base]
    for prefixo, ficha_id in (("divida_bruta_gg_", "divida_bruta_gg"), ("ipca_", "ipca_12m"),
                              ("pib_per_capita_", "pib_per_capita_variacao_real")):
        if serie.startswith(prefixo):
            return fichas[ficha_id]
    return None


def tipo_de_media(serie, ficha):
    if "__brasil_menos_mediana" in serie or serie == "ipca_meta":
        return "aritmetica"
    if serie == "pib_per_capita_indice":
        return "indice"
    return ficha.get("media_mandato", "aritmetica")


def media(valores, tipo):
    if not valores:
        return None
    if tipo == "geometrica":
        produto = math.prod(1 + v / 100 for v in valores)
        return (produto ** (1 / len(valores)) - 1) * 100
    return sum(valores) / len(valores)


def anuais_de_mensal(pontos, anualizacao):
    """{AAAA-MM: v} -> {AAAA: v}, só anos completos."""
    por_ano = defaultdict(dict)
    for data, v in pontos.items():
        por_ano[data[:4]][data[5:]] = v
    anuais = {}
    for ano, meses in por_ano.items():
        if anualizacao == "dezembro" and "12" in meses:
            anuais[ano] = meses["12"]
        elif anualizacao == "media" and len(meses) == 12:
            anuais[ano] = sum(meses.values()) / 12
    return anuais


def resumir_serie(serie, pontos, ficha, mandato):
    tipo = tipo_de_media(serie, ficha)
    mensal = len(next(iter(pontos))) == 7
    primeiro_disponivel = min(pontos)
    em_curso = mandato["fim"] > HOJE

    if mensal:
        periodos = meses_do_mandato(mandato)
        a, m = int(periodos[0][:4]), int(periodos[0][5:])
        herdou_periodo = f"{a - 1:04d}-12" if m == 1 else f"{a:04d}-{m - 1:02d}"
        anuais = anuais_de_mensal(pontos, ficha.get("anualizacao", "dezembro"))
    else:
        periodos = [str(a) for a in anos_do_mandato(mandato)]
        herdou_periodo = str(int(periodos[0]) - 1)
        if herdou_periodo not in pontos:
            # séries bienais ou trienais (IDEB, PISA) ou com ano faltando: vale a última
            # observação antes da posse, até quatro anos antes
            anteriores = [p for p in pontos if int(periodos[0]) - 4 <= int(p) < int(periodos[0])]
            if anteriores:
                herdou_periodo = max(anteriores)
        anuais = pontos

    no_mandato = [p for p in periodos if p in pontos]
    if not no_mandato:
        return None
    entregou_periodo = no_mandato[-1]
    anos = [str(a) for a in anos_do_mandato(mandato)] if mensal else periodos
    if mensal and not anos:  # mandatos curtos (Dilma 2): usa os meses
        anos = sorted({p[:4] for p in periodos})

    if mensal and ficha.get("anualizacao") == "media":
        # média ponderada pelo tempo: todos os meses disponíveis do mandato
        base_media = {p: pontos[p] for p in no_mandato}
        chave_ano = lambda p: p[:4]
        rotulo = f"{no_mandato[0]} a {no_mandato[-1]} (meses)"
    else:
        base_media = {a: anuais[a] for a in anos if a in anuais}
        chave_ano = lambda p: p
        rotulo = f"{min(base_media)} a {max(base_media)}" if base_media else ""

    valores = list(base_media.values())
    sem_pandemia = [v for p, v in base_media.items() if chave_ano(p) not in ANOS_PANDEMIA]
    if tipo == "indice":
        valores = sem_pandemia = []

    herdou = pontos.get(herdou_periodo)
    entregou = pontos[entregou_periodo]
    acumulado = None
    if tipo == "geometrica" and base_media:
        acumulado = (math.prod(1 + v / 100 for v in base_media.values()) - 1) * 100
    if tipo == "indice" and herdou:
        acumulado = (entregou / herdou - 1) * 100

    notas = []
    parcial = em_curso or entregou_periodo != periodos[-1]
    if primeiro_disponivel > periodos[0]:  # cobertura incompleta: sem média, para não enganar
        notas.append(f"série começa em {primeiro_disponivel}")
        parcial = True
        valores = sem_pandemia = []
        acumulado = None
    if em_curso:
        notas.append("mandato em curso")

    def r(x):
        return "" if x is None else round(x, 4)

    return {
        "herdou_periodo": herdou_periodo if herdou is not None else "",
        "herdou": r(herdou),
        "entregou_periodo": entregou_periodo,
        "entregou": r(entregou),
        "entregou_menos_herdou": r(entregou - herdou) if herdou is not None else "",
        "media": r(media(valores, tipo)),
        "media_sem_2020_21": r(media(sem_pandemia, tipo)) if len(sem_pandemia) != len(valores) else "",
        "acumulado_no_mandato": r(acumulado),
        "tipo_media": tipo,
        "periodos_na_media": rotulo,
        "parcial": "sim" if parcial else "",
        "observacao": "; ".join(notas),
    }


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--tema", required=True)
    ap.add_argument("--entrada", help="CSV em dados/tratados/ (padrão: <tema>.csv)")
    args = ap.parse_args()

    entrada = TRATADOS / (args.entrada or f"{args.tema}.csv")
    fichas = {f["id"]: f for f in ler_fichas(args.tema)}
    series = defaultdict(dict)
    unidades = {}
    for linha in ler_csv(entrada):
        series[linha["indicador"]][linha["data"]] = float(linha["valor"])
        unidades[linha["indicador"]] = linha["unidade"]

    n_paises = defaultdict(dict)
    pares = TRATADOS / f"{entrada.stem}_pares_paises.csv"
    if pares.exists():
        for linha in ler_csv(pares):
            if "__n_" in linha["indicador"]:
                n_paises[linha["indicador"]][linha["data"]] = int(float(linha["valor"]))

    mandatos = [m for m in ler_mandatos() if m["papel"] != "referencia"]
    campos = ["indicador", "mandato", "herdou_periodo", "herdou", "entregou_periodo", "entregou",
              "entregou_menos_herdou", "media", "media_sem_2020_21", "acumulado_no_mandato",
              "tipo_media", "periodos_na_media", "parcial", "unidade", "observacao"]
    linhas = []
    for serie in sorted(series):
        ficha = ficha_da_serie(serie, fichas)
        if serie in SEM_RESUMO or "_parcial" in serie or ficha is None:
            continue
        for m in mandatos:
            res = resumir_serie(serie, series[serie], ficha, m)
            if res is None:
                continue
            if "__mediana_" in serie and "menos" not in serie:
                base, grupo = serie.split("__mediana_")
                ns = [n for a, n in n_paises.get(f"{base}__n_{grupo}", {}).items()
                      if a in [str(x) for x in anos_do_mandato(m)]]
                if ns:
                    nota = f"mediana de {min(ns)} a {max(ns)} países" if min(ns) != max(ns) else f"mediana de {ns[0]} países"
                    res["observacao"] = "; ".join(x for x in (res["observacao"], nota) if x)
            res.update(indicador=serie, mandato=m["id"], unidade=unidades[serie])
            linhas.append([res[c] for c in campos])

    destino = entrada.with_name(entrada.stem + "_mandatos.csv")
    escrever_csv(destino, campos, linhas)
    print(f"{len(linhas)} linhas -> {destino.relative_to(destino.parents[2])}")


if __name__ == "__main__":
    main()