"""
Modelos quantitativos do 2º turno presidencial 2026 (Lula x Flávio Bolsonaro).
Autor: Saulo Dutra. Convenções declaradas:
- "spread" = Flávio − Lula, em pontos percentuais (pp), sempre sobre votos VÁLIDOS
  (dois candidatos, brancos/nulos/indecisos excluídos), salvo quando dito "total".
- "share2" = Flávio / (Flávio + Lula): parcela de Flávio nos votos entre os dois.
- Datas ISO. Dias contados em dias corridos.
Rode: uv run --python .venv/Scripts/python.exe python modelos.py
"""
from __future__ import annotations

import json
import math
from collections import defaultdict
from dataclasses import dataclass, asdict
from itertools import combinations
from pathlib import Path

import numpy as np
import pandas as pd
from scipy import stats

RAIZ = Path(__file__).resolve().parent
DADOS = RAIZ / "dados"
PESQUISA = RAIZ.parent / "pesquisa"
SAIDA = RAIZ / "saida"
SITE_DATA = RAIZ.parent / "site" / "data"
SAIDA.mkdir(exist_ok=True)
SITE_DATA.mkdir(parents=True, exist_ok=True)

HOJE = pd.Timestamp("2026-10-06")
DATA_2T = pd.Timestamp("2026-10-25")
DIAS_RESTANTES = (DATA_2T - HOJE).days
RNG = np.random.default_rng(20261025)
N_SIM = 200_000


# ----------------------------------------------------------------------------
# utilidades
# ----------------------------------------------------------------------------
def validos(a: float, b: float) -> tuple[float, float]:
    """Renormaliza dois percentuais (sobre total) para válidos entre os dois."""
    s = a + b
    return (100 * a / s, 100 * b / s) if s else (float("nan"), float("nan"))


def spread_validos(flavio: float, lula: float) -> float:
    f, lu = validos(flavio, lula)
    return f - lu


def ler_csv(nome: str, pasta: Path = PESQUISA) -> pd.DataFrame | None:
    p = pasta / nome
    if not p.exists():
        return None
    try:
        return pd.read_csv(p)
    except Exception as e:  # noqa: BLE001
        print(f"[aviso] não consegui ler {p}: {e}")
        return None


def num(x) -> float:
    try:
        v = float(str(x).replace(",", ".").replace("%", "").strip())
        return v
    except (TypeError, ValueError):
        return float("nan")


# ----------------------------------------------------------------------------
# 1. Série de pesquisas de 2º turno e house effects (Jackman 2005)
# ----------------------------------------------------------------------------
@dataclass
class Poll2T:
    instituto: str
    fim: str
    amostra: float
    lula: float
    flavio: float
    spread: float  # válidos, Flávio − Lula
    fonte: str


def carregar_2t() -> list[Poll2T]:
    rows: list[Poll2T] = []
    df = ler_csv("pesquisas_2t.csv")
    if df is None or df.empty:
        df = pd.read_csv(DADOS / "ultimas_pesquisas_2t.csv")
        df["inicio"] = df["data_fim"]
        df = df.rename(columns={"data_fim": "fim"})
    for _, r in df.iterrows():
        lu, fl = num(r.get("lula")), num(r.get("flavio"))
        if math.isnan(lu) or math.isnan(fl):
            continue
        fim = str(r.get("fim") or r.get("data_fim"))
        rows.append(Poll2T(str(r["instituto"]).strip(), fim[:10], num(r.get("amostra")),
                           lu, fl, spread_validos(fl, lu), str(r.get("fonte", ""))))
    # dedup por (instituto, fim)
    seen, out = set(), []
    for p in sorted(rows, key=lambda p: p.fim):
        k = (p.instituto.lower(), p.fim)
        if k in seen:
            continue
        seen.add(k)
        out.append(p)
    return out


def house_effects(polls: list[Poll2T], janela_dias: int = 21) -> dict:
    """
    Modelo aditivo (Jackman 2005): spread_it = mu(t) + delta_i + e_it.
    mu(t): média móvel ponderada por recência (kernel gaussiano, meia-vida ~10 dias) e por sqrt(n).
    delta_i: efeito de casa = média dos resíduos de cada instituto, com shrinkage bayesiano
             (prior N(0, 2²), pelo tamanho típico de house effects de ~1-3 pp na literatura).
    Iteramos mu/delta até convergir (backfitting).
    """
    if len(polls) < 3:
        return {"consenso": [], "house": {}, "nota": "poucas pesquisas"}
    t = np.array([(pd.Timestamp(p.fim) - HOJE).days for p in polls], dtype=float)
    y = np.array([p.spread for p in polls])
    n = np.array([p.amostra if not math.isnan(p.amostra) else 2000 for p in polls])
    inst = np.array([p.instituto for p in polls])
    w_n = np.sqrt(n / 2000.0)
    delta = defaultdict(float)
    tau2_prior = 2.0**2
    grade = np.arange(t.min(), 1.0, 1.0)  # até hoje (t=0)

    def mu_em(tt: np.ndarray, yy: np.ndarray, ww: np.ndarray, hl: float = 10.0) -> np.ndarray:
        lam = math.log(2) / hl
        out = np.empty_like(tt)
        for i, g in enumerate(tt):
            k = np.exp(-lam * np.abs(t - g)) * ww
            out[i] = np.sum(k * yy) / np.sum(k) if k.sum() > 0 else np.nan
        return out

    for _ in range(50):
        y_adj = y - np.array([delta[i] for i in inst])
        mu_t = mu_em(t, y_adj, w_n)
        resid = y - mu_t
        novo = {}
        for i in set(inst):
            m = inst == i
            k = m.sum()
            s2 = max(np.var(resid[m], ddof=0) if k > 1 else 4.0, 1.0)
            # posterior mean com prior N(0,tau2) e verossimilhança N(delta, s2/k)
            novo[i] = (resid[m].mean() * k / s2) / (k / s2 + 1 / tau2_prior)
        # identificação: média ponderada dos house effects = 0
        mean_d = np.mean(list(novo.values()))
        novo = {i: v - mean_d for i, v in novo.items()}
        if max(abs(novo[i] - delta[i]) for i in novo) < 1e-4:
            delta = defaultdict(float, novo)
            break
        delta = defaultdict(float, novo)

    y_adj = y - np.array([delta[i] for i in inst])
    mu_grade = mu_em(grade, y_adj, w_n)
    # incerteza do consenso: erro-padrão efetivo via n efetivo do kernel + dispersão residual
    lam = math.log(2) / 10.0
    se = []
    for g in grade:
        k = np.exp(-lam * np.abs(t - g)) * w_n
        neff = (k.sum() ** 2) / np.sum(k**2) if k.sum() > 0 else 1
        s_res = float(np.sqrt(np.mean((y_adj - mu_em(t, y_adj, w_n)) ** 2)))
        se.append(max(s_res, 1.5) / math.sqrt(max(neff, 1)))
    hoje = float(mu_grade[-1]) if len(mu_grade) else float("nan")
    return {
        "consenso": [{"dia": int(d), "data": (HOJE + pd.Timedelta(days=int(d))).strftime("%Y-%m-%d"),
                      "spread": round(float(m), 2), "se": round(float(s), 2)}
                     for d, m, s in zip(grade, mu_grade, se) if not math.isnan(m)],
        "spread_hoje": round(hoje, 2),
        "se_hoje": round(float(se[-1]), 2) if se else None,
        "house": {str(i): round(float(delta[i]), 2) for i in sorted(set(inst))},
        "n_pesquisas": len(polls),
        "pesquisas": [asdict(p) for p in polls],
    }


# ----------------------------------------------------------------------------
# 2. Viés histórico das pesquisas (prior do erro)
# ----------------------------------------------------------------------------
def vies_historico() -> dict:
    """
    erro_spread = (poll_a − poll_b) − (result_a − result_b), com a = candidato de direita.
    Convenção: erro NEGATIVO = pesquisas subestimaram a direita.
    """
    df = ler_csv("erro_pesquisas.csv")
    base = []
    if df is not None and not df.empty:
        for _, r in df.iterrows():
            # Canônico: e = spread_poll(direita − esquerda) − spread_result(direita − esquerda).
            # No CSV da pesquisa, cand_a é a ESQUERDA (Dilma/Haddad/Lula) e cand_b a direita,
            # então o sinal precisa ser invertido. Detectamos pelo nome de cand_b.
            pa, ra, pb, rb = (num(r.get(c)) for c in ("poll_a", "result_a", "poll_b", "result_b"))
            if any(math.isnan(v) for v in (pa, pb, ra, rb)):
                continue
            direita_e_b = any(t in str(r.get("cand_b", "")).lower() for t in ("bolsonaro", "aécio", "aecio", "serra", "alckmin"))
            e = ((pb - pa) - (rb - ra)) if direita_e_b else ((pa - pb) - (ra - rb))
            turno = str(r.get("turno")).strip().upper().replace("º", "").replace("T", "")
            turno = f"{turno}T" if turno in ("1", "2") else str(r.get("turno"))
            base.append({"ano": int(num(r.get("ano"))), "turno": turno,
                         "instituto": str(r.get("instituto")), "erro_spread": round(e, 2),
                         "campo_subestimado": str(r.get("campo_subestimado", ""))})
    # fallback mínimo e verificável (válidos): 2026-1T das três grandes
    if not base:
        res = 47.03 - 45.16
        for inst, fl, lu in (("Datafolha", 42, 45), ("Quaest", 45, 46), ("AtlasIntel", 43.8, 46.7)):
            base.append({"ano": 2026, "turno": "1T", "instituto": inst,
                         "erro_spread": round((fl - lu) - res, 2), "campo_subestimado": "direita"})
    e = np.array([b["erro_spread"] for b in base])
    por_ano = defaultdict(list)
    for b in base:
        por_ano[(b["ano"], b["turno"])].append(b["erro_spread"])
    media_por_eleicao = {f"{k[0]}-{k[1]}": round(float(np.mean(v)), 2) for k, v in sorted(por_ano.items())}
    # Prior do viés para o 2T de 2026 (um ponto por eleição, não por instituto, para não
    # superponderar anos com mais institutos). Erros de 2º turno pesam 1,0; erros de 1º turno
    # pesam 0,5, porque no 2T o quadro de dois candidatos reduz o erro de alocação de
    # indecisos (ver erro 2022: 1T ≈ −7 pp contra 2T ≈ −2 pp nas grandes casas).
    chaves = list(media_por_eleicao)
    pontos = np.array([media_por_eleicao[k] for k in chaves])
    # pesos: 2º turno 1,0; o 1º turno de 2026 0,75 (mesmo ciclo, mesmas casas, mesmos candidatos);
    # 1º turno de anos anteriores 0,25 (quadro multicandidato, erro de alocação de indecisos)
    pesos = np.array([1.0 if k.endswith("2T") else (0.75 if k.startswith("2026") else 0.25) for k in chaves])
    if len(pontos):
        m_w = float(np.average(pontos, weights=pesos))
        v_w = float(np.average((pontos - m_w) ** 2, weights=pesos)) if len(pontos) > 1 else 9.0
        sd_w = max(math.sqrt(v_w), 1.5)  # piso de 1,5 pp: nunca fingir certeza sobre o viés
    else:
        m_w, sd_w = 0.0, 3.0
    return {
        "registros": base,
        "media_por_eleicao": media_por_eleicao,
        "prior_vies": {"media": round(m_w, 2), "desvio": round(sd_w, 2), "n_eleicoes": int(len(pontos)),
                       "regra": "média ponderada dos erros médios por eleição (2T peso 1; 1T de 2026 peso 0,75; 1T anteriores 0,25); sd com piso 1,5 pp"},
        "rmse_total": round(float(np.sqrt(np.mean(e**2))), 2) if len(e) else None,
        "pct_subestimou_direita": round(100 * float(np.mean(e < 0)), 1) if len(e) else None,
    }


# ----------------------------------------------------------------------------
# 3. Transferência de votos (matriz com incerteza Beta)
# ----------------------------------------------------------------------------
def transferencia(n_sim: int = 50_000) -> dict:
    r = json.loads((DADOS / "resultado_1t_2026.json").read_text(encoding="utf-8"))
    v = r["validos"]
    # Premissas de transferência (parcela que vai p/ Flávio | Lula | abstém/branco), com
    # concentração Beta. Rótulo: ESTIMATIVA do autor, ancorada em: apoio declarado (Caiado,
    # Zema → Flávio), perfil do eleitor (Cury e Santos: direita/antipetista; Samara: esquerda).
    # Se o arquivo transferencia_votos.json trouxer "segunda opção" medida, ela substitui.
    # Premissas = MÉDIA das "segundas opções" medidas no cenário de 2º turno (votos totais do
    # subgrupo, Flávio | Lula | branco-nulo-abstém): Quaest 02-03/10/2026 (BR-02197) e
    # Datafolha 29/09-01/10/2026 (BR-08039), via G1. Cury: 34/23/41 e 44/30/24 -> 39/26,5/34,5.
    # Renan Santos: 60/11/27 e 54/28/16 -> 57/19,5/23,5. Caiado: 43/19/36 e 46/30/22 -> 44,5/24,5/31
    # (medido ANTES do apoio declarado em 06/10; o endosso é upside não modelado). Zema: Datafolha
    # agregado set/2026 48/25/27. Samara (UP) e "outros" (PSTU/PCB/PCO/DC): ESTIMATIVA do autor
    # (eleitorado de esquerda). Fontes em pesquisa/transferencia_votos.md, seção 1.2.
    premissas = {
        "cury": (0.39, 0.265, 0.345), "santos": (0.57, 0.195, 0.235), "caiado": (0.445, 0.245, 0.31),
        "zema": (0.48, 0.25, 0.27), "samara": (0.05, 0.70, 0.25), "outros": (0.10, 0.60, 0.30),
    }
    rotulo = {"cury": "medido (Quaest+Datafolha)", "santos": "medido (Quaest+Datafolha)",
              "caiado": "medido (Quaest+Datafolha), pré-apoio", "zema": "medido (Datafolha set/26)",
              "samara": "estimativa", "outros": "estimativa"}
    medido = {}
    tj = PESQUISA / "transferencia_votos.json"
    if tj.exists():
        try:
            d = json.loads(tj.read_text(encoding="utf-8"))
            medido = d.get("segunda_opcao", {}) or d.get("transferencia", {}) or {}
        except Exception:  # noqa: BLE001
            medido = {}
    for k, val in medido.items():
        k2 = k.lower()
        if k2 in premissas and isinstance(val, dict):
            f, lu = num(val.get("flavio")), num(val.get("lula"))
            if not math.isnan(f) and not math.isnan(lu):
                f, lu = f / 100 if f > 1 else f, lu / 100 if lu > 1 else lu
                premissas[k2] = (f, lu, max(0.0, 1 - f - lu))
    conc = 40.0  # concentração da Beta/Dirichlet: ~±8 pp de incerteza em cada parcela
    base_f, base_l = v["flavio"]["pct"], v["lula"]["pct"]
    # retenção dos próprios eleitores (abstenção diferencial): Dirichlet também
    ret = {"flavio": 0.965, "lula": 0.965}
    sims_f, sims_l = np.zeros(n_sim), np.zeros(n_sim)
    sims_f += base_f * RNG.beta(ret["flavio"] * 60, (1 - ret["flavio"]) * 60, n_sim)
    sims_l += base_l * RNG.beta(ret["lula"] * 60, (1 - ret["lula"]) * 60, n_sim)
    # choque de mobilização diferencial (abstenção assimétrica entre os campos): ±1,2 pp (1 sd)
    choque = RNG.normal(0, 1.2, n_sim)
    sims_f += choque / 2
    sims_l -= choque / 2
    detalhe = {}
    for cand, (pf, pl, pa) in premissas.items():
        pct = v[cand]["pct"]
        draws = RNG.dirichlet(np.array([pf, pl, pa]) * conc, n_sim)
        sims_f += pct * draws[:, 0]
        sims_l += pct * draws[:, 1]
        detalhe[cand] = {"pct_1t": pct, "para_flavio": pf, "para_lula": pl, "abstem": pa, "origem": rotulo.get(cand, "estimativa"),
                         "esperado_flavio_pp": round(pct * pf, 2), "esperado_lula_pp": round(pct * pl, 2)}
    share = 100 * sims_f / (sims_f + sims_l)
    spread = 2 * share - 100
    return {
        "premissas": detalhe,
        "nota": "Cury, Renan Santos, Caiado e Zema usam a segunda opção medida por Quaest e Datafolha (média); Samara e outros são estimativa do autor.",
        "share_flavio": {"media": round(float(share.mean()), 2), "p05": round(float(np.percentile(share, 5)), 2),
                         "p95": round(float(np.percentile(share, 95)), 2)},
        "spread": {"media": round(float(spread.mean()), 2), "p05": round(float(np.percentile(spread, 5)), 2),
                   "p95": round(float(np.percentile(spread, 95)), 2)},
        "p_flavio": round(float(np.mean(spread > 0)), 4),
        "pool_transferivel_pp": round(sum(v[c]["pct"] for c in premissas), 2),
    }


# ----------------------------------------------------------------------------
# 4. Histórico 1T → 2T (prior estrutural) e regressão
# ----------------------------------------------------------------------------
def historico_1t_2t() -> dict:
    df = ler_csv("historico_2t.csv")
    if df is None or df.empty or "pct_lider_1t" not in df.columns:
        df = pd.read_csv(DADOS / "historico_2t_base.csv")
    df = df.copy()
    for c in ("pct_lider_1t", "pct_segundo_1t", "pct_vencedor_2t", "pct_perdedor_2t"):
        df[c] = df[c].map(num)
    df = df.dropna(subset=["pct_lider_1t", "pct_segundo_1t", "pct_vencedor_2t"])
    # margem do líder do 1T sobre o 2º, nos dois turnos (válidos entre os dois no 1T)
    df["m1"] = 100 * (df.pct_lider_1t - df.pct_segundo_1t) / (df.pct_lider_1t + df.pct_segundo_1t)
    lider_venceu = df.apply(lambda r: str(r.get("vencedor_2t", "")).strip().lower() == str(r.get("lider_1t", "")).strip().lower(), axis=1)
    df["m2"] = np.where(lider_venceu, df.pct_vencedor_2t - df.pct_perdedor_2t, -(df.pct_vencedor_2t - df.pct_perdedor_2t))
    x, y = df.m1.values, df.m2.values
    reg = stats.linregress(x, y)
    m1_2026 = 100 * (47.03 - 45.16) / (47.03 + 45.16)
    pred = reg.intercept + reg.slope * m1_2026
    resid_sd = float(np.sqrt(np.sum((y - (reg.intercept + reg.slope * x)) ** 2) / max(len(x) - 2, 1)))
    # Laplace: P(líder vence) com n vitórias em n eleições → (n+1)/(n+2)
    n = int(len(df))
    k = int(lider_venceu.sum())
    return {
        "tabela": [{"ano": int(r.ano), "lider_1t": r.lider_1t, "m1": round(float(r.m1), 2), "m2": round(float(r.m2), 2),
                    "lider_venceu": bool(v)} for (_, r), v in zip(df.iterrows(), lider_venceu)],
        "regressao": {"slope": round(float(reg.slope), 3), "intercept": round(float(reg.intercept), 3),
                      "r2": round(float(reg.rvalue**2), 3), "resid_sd": round(resid_sd, 2), "n": n},
        "m1_2026": round(float(m1_2026), 2),
        "pred_m2_2026": round(float(pred), 2),
        "laplace_p_lider": round((k + 1) / (n + 2), 3),
        "lider_venceu_k_de_n": f"{k}/{n}",
    }


# ----------------------------------------------------------------------------
# 5. Monte Carlo do 2º turno: três modelos + média bayesiana de modelos
# ----------------------------------------------------------------------------
def monte_carlo(he: dict, vb: dict, tr: dict, hist: dict) -> dict:
    n = N_SIM
    dias = max(DIAS_RESTANTES, 1)
    # drift de campanha: passeio aleatório com sd anual-equivalente calibrado em ~0,45 pp/√dia
    # (mudança típica do spread entre a 1ª e a última semana de 2T no Brasil 2014/2018/2022: 2–4 pp)
    sd_drift = 0.45 * math.sqrt(dias)
    out = {}
    # --- Modelo A: consenso das pesquisas + viés histórico + drift
    mu = he.get("spread_hoje", float("nan"))
    se = he.get("se_hoje") or 2.0
    if not math.isnan(mu):
        vies = RNG.normal(vb["prior_vies"]["media"], vb["prior_vies"]["desvio"], n)
        # erro NEGATIVO = subestimou direita → resultado = poll − erro
        a = RNG.normal(mu, se, n) - vies + RNG.normal(0, sd_drift, n)
        out["A_pesquisas_vies"] = a
    # --- Modelo B: transferência mecânica do 1T
    b = RNG.normal(tr["spread"]["media"], (tr["spread"]["p95"] - tr["spread"]["p05"]) / 3.29, n) + RNG.normal(0, sd_drift * 0.6, n)
    out["B_transferencia"] = b
    # --- Modelo C: estrutural histórico 1T→2T
    c = RNG.normal(hist["pred_m2_2026"], hist["regressao"]["resid_sd"], n)
    out["C_historico"] = c
    pesos = {"A_pesquisas_vies": 0.45, "B_transferencia": 0.35, "C_historico": 0.20}
    pesos = {k: v for k, v in pesos.items() if k in out}
    s = sum(pesos.values())
    pesos = {k: v / s for k, v in pesos.items()}
    idx = RNG.choice(list(pesos), size=n, p=list(pesos.values()))
    comb = np.select([idx == k for k in pesos], [out[k] for k in pesos])
    out["Combinado"] = comb

    def resumo(x: np.ndarray) -> dict:
        return {"p_flavio": round(float(np.mean(x > 0)), 4), "media": round(float(x.mean()), 2),
                "p05": round(float(np.percentile(x, 5)), 2), "p50": round(float(np.percentile(x, 50)), 2),
                "p95": round(float(np.percentile(x, 95)), 2), "sd": round(float(x.std()), 2)}

    bins = np.arange(-20, 20.5, 0.5)
    h, _ = np.histogram(comb, bins=bins, density=True)
    return {
        "modelos": {k: resumo(v) for k, v in out.items()},
        "pesos": pesos,
        "dias_restantes": dias, "sd_drift": round(sd_drift, 2),
        "hist": {"bins": [round(float(b), 2) for b in bins[:-1]], "dens": [round(float(x), 5) for x in h]},
        "n_sim": n,
        "share_flavio_combinado": resumo((comb + 100) / 2),
    }


# ----------------------------------------------------------------------------
# 6. Atualização bayesiana sequencial (para quando entrarem Datafolha 08/10, Atlas 09/10...)
# ----------------------------------------------------------------------------
def bayes_sequencial(hist: dict, vb: dict, polls: list[Poll2T]) -> dict:
    """
    Prior: spread ~ N(pred_m2_2026, resid_sd²) do modelo estrutural (o que as pesquisas "deveriam"
    medir). Cada pesquisa de 2T entra como verossimilhança sobre o spread MEDIDO:
    N(spread_poll, se_amostral² + house²), com se_amostral = 2·100·√(0,25/n) e house = 1,5 pp
    (dispersão típica entre casas). Passo conjugado: precisões somam, média pondera por precisão.
    O viés compartilhado (erro sistemático comum a todas as casas) NÃO encolhe com mais pesquisas:
    entra uma única vez no fim — média − viés_esperado, variância + tau². É isso que separa este
    modelo de "somar pesquisas até a incerteza sumir".
    """
    m, s2 = hist["pred_m2_2026"], hist["regressao"]["resid_sd"] ** 2
    tau2 = vb["prior_vies"]["desvio"] ** 2
    vies = vb["prior_vies"]["media"]
    house2 = 1.5**2

    def passo(etapa: str, obs: float | None) -> dict:
        sd_tot = math.sqrt(s2 + tau2)
        return {"etapa": etapa, "obs_spread": None if obs is None else round(obs, 2),
                "media_medida": round(m, 2), "media": round(m - vies, 2), "sd": round(sd_tot, 2),
                "p_flavio": round(float(1 - stats.norm.cdf(0, m - vies, sd_tot)), 4)}

    passos = [passo("prior (estrutural 1T→2T)", None)]
    for p in sorted(polls, key=lambda p: p.fim):
        nn = p.amostra if not math.isnan(p.amostra) else 2000
        se2 = (2 * 100 * math.sqrt(0.25 / nn)) ** 2 + house2
        prec = 1 / s2 + 1 / se2
        m = (m / s2 + p.spread / se2) / prec
        s2 = 1 / prec
        passos.append(passo(f"{p.instituto} {p.fim}", p.spread))
    return {"passos": passos, "posterior": passos[-1], "vies_aplicado": vies, "tau": round(math.sqrt(tau2), 2)}


# ----------------------------------------------------------------------------
# 7. Poder de voto no Congresso (Banzhaf e Shapley-Shubik) — teoria nomeada
# ----------------------------------------------------------------------------
def indices_poder(pesos: dict[str, int], quota: int, max_shapley: int = 14) -> dict:
    nomes = list(pesos)
    w = np.array([pesos[n] for n in nomes])
    k = len(nomes)
    # Banzhaf: contagem de coalizões em que o jogador é pivô (enumeração 2^k; k<=16 ok)
    swings = np.zeros(k)
    for mask in range(1 << k):
        tot = sum(w[i] for i in range(k) if mask >> i & 1)
        if tot >= quota:
            for i in range(k):
                if mask >> i & 1 and tot - w[i] < quota:
                    swings[i] += 1
    banzhaf = swings / swings.sum() if swings.sum() else swings
    res = {"quota": quota, "banzhaf": {n: round(float(b), 4) for n, b in zip(nomes, banzhaf)}}
    if k <= max_shapley:
        # Shapley-Shubik por fórmula de coalizões: phi_i = Σ_S |S|!(k−|S|−1)!/k! [S∪i vence, S perde]
        phi = np.zeros(k)
        fact = [math.factorial(j) for j in range(k + 1)]
        for i in range(k):
            outros = [j for j in range(k) if j != i]
            for r in range(len(outros) + 1):
                for S in combinations(outros, r):
                    t = sum(w[j] for j in S)
                    if t < quota <= t + w[i]:
                        phi[i] += fact[r] * fact[k - r - 1]
        phi /= fact[k]
        res["shapley"] = {n: round(float(p), 4) for n, p in zip(nomes, phi)}
    return res


def congresso() -> dict:
    base = json.loads((DADOS / "congresso_base.json").read_text(encoding="utf-8"))
    out = {}
    sen = base["senado"]
    cam = base["camara"]
    # Câmara: agrupa partidos pequenos para manter <= 14 jogadores no Shapley
    partidos = dict(cam["partidos_2027"])
    grandes = {k: v for k, v in partidos.items() if v >= 10}
    pequenos = sum(v for k, v in partidos.items() if v < 10)
    grandes["Outros(<10)"] = pequenos
    out["camara"] = {"cadeiras": partidos, "poder": indices_poder(grandes, cam["maioria_absoluta"]),
                     "poder_pec": indices_poder(grandes, cam["pec"]),
                     "limiares": {k: cam[k] for k in ("total", "maioria_absoluta", "pec", "dois_tercos")}}
    # Senado: composição total por partido só vem do arquivo senado_2027.csv; usamos blocos + eleitos
    df = ler_csv("senado_2027.csv")
    sen_part = None
    if df is not None and "partido" in df.columns and len(df) >= 70:
        sen_part = df["partido"].astype(str).str.strip().value_counts().to_dict()
    out["senado"] = {"eleitos_2026": sen["eleitos_2026_por_partido"], "blocos_2027": sen["blocos_2027"],
                     "limiares": {k: sen[k] for k in ("total", "maioria_absoluta", "pec", "dois_tercos")},
                     "composicao_total": sen_part}
    if sen_part:
        g = {k: v for k, v in sen_part.items() if v >= 3}
        g["Outros(<3)"] = sum(v for k, v in sen_part.items() if v < 3)
        out["senado"]["poder"] = indices_poder(g, sen["maioria_absoluta"])
        out["senado"]["poder_pec"] = indices_poder(g, sen["pec"])
    b = sen["blocos_2027"]
    out["senado"]["distancias"] = {
        "direita_para_maioria": sen["maioria_absoluta"] - b["direita"],
        "direita_para_pec": sen["pec"] - b["direita"],
        "direita_para_dois_tercos": sen["dois_tercos"] - b["direita"],
        "esquerda_para_bloquear_pec": (sen["total"] - sen["pec"] + 1) - b["esquerda"],
    }
    return out


# ----------------------------------------------------------------------------
# 8. Estados: swing 2022 → 2026 (1T)
# ----------------------------------------------------------------------------
def estados() -> dict:
    a = ler_csv("resultado_1t_2026_uf.csv")
    b = ler_csv("resultado_2022_uf.csv")
    if a is None or b is None:
        return {"nota": "arquivos por UF ainda não disponíveis"}
    a = a.copy()
    b = b.copy()
    a["uf"] = a["uf"].astype(str).str.upper().str.strip()
    b["uf"] = b["uf"].astype(str).str.upper().str.strip()
    m = a.merge(b, on="uf", how="inner")
    for c in ("flavio", "lula", "lula_1t", "bolsonaro_1t", "lula_2t", "bolsonaro_2t", "eleitores"):
        if c in m.columns:
            m[c] = m[c].map(num)
    m["spread_2026"] = m.apply(lambda r: spread_validos(r.flavio, r.lula), axis=1)
    m["spread_2022_1t"] = m.apply(lambda r: spread_validos(r.bolsonaro_1t, r.lula_1t), axis=1)
    m["spread_2022_2t"] = m.apply(lambda r: spread_validos(r.bolsonaro_2t, r.lula_2t), axis=1) if "bolsonaro_2t" in m else np.nan
    m["swing"] = m.spread_2026 - m.spread_2022_1t
    cols = ["uf", "eleitores", "flavio", "lula", "spread_2026", "spread_2022_1t", "spread_2022_2t", "swing"]
    cols = [c for c in cols if c in m.columns]
    tab = m[cols].sort_values("swing", ascending=False)
    w = m.eleitores.fillna(0).values if "eleitores" in m else np.ones(len(m))
    return {
        "tabela": json.loads(tab.round(2).to_json(orient="records")),
        "swing_medio_ponderado": round(float(np.average(m.swing.fillna(0), weights=w)) if w.sum() else float(m.swing.mean()), 2),
        "estados_viraram": json.loads(m[(m.spread_2022_1t < 0) & (m.spread_2026 > 0)][["uf"]].to_json(orient="records")),
        "n": int(len(m)),
    }


# ----------------------------------------------------------------------------
# main
# ----------------------------------------------------------------------------
def limpo(o):
    """Troca NaN/inf por None e tipos numpy por nativos para gerar JSON válido."""
    if isinstance(o, dict):
        return {str(k): limpo(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)):
        return [limpo(v) for v in o]
    if isinstance(o, np.integer):
        return int(o)
    if isinstance(o, (np.floating, float)):
        return None if (math.isnan(o) or math.isinf(o)) else float(o)
    if isinstance(o, np.bool_):
        return bool(o)
    return o


def main() -> None:
    polls = carregar_2t()
    he = house_effects(polls)
    vb = vies_historico()
    tr = transferencia()
    hist = historico_1t_2t()
    mc = monte_carlo(he, vb, tr, hist)
    # bayes: usa só pesquisas de 2T com campo encerrado em 30/09 ou depois
    recentes = [p for p in polls if p.fim >= "2026-09-30"]
    by = bayes_sequencial(hist, vb, recentes)
    cg = congresso()
    uf = estados()
    resultado = {
        "gerado_em": pd.Timestamp.now().strftime("%Y-%m-%d %H:%M"),
        "hoje": HOJE.strftime("%Y-%m-%d"), "data_2t": DATA_2T.strftime("%Y-%m-%d"), "dias_restantes": DIAS_RESTANTES,
        "resultado_1t": json.loads((DADOS / "resultado_1t_2026.json").read_text(encoding="utf-8")),
        "house_effects": he, "vies_historico": vb, "transferencia": tr, "historico_1t_2t": hist,
        "monte_carlo": mc, "bayes": by, "congresso": cg, "estados": uf,
    }
    resultado = limpo(resultado)
    (SAIDA / "resultados.json").write_text(json.dumps(resultado, ensure_ascii=False, indent=1, allow_nan=False), encoding="utf-8")
    (SITE_DATA / "resultados.json").write_text(json.dumps(resultado, ensure_ascii=False, allow_nan=False), encoding="utf-8")
    print("pesquisas 2T:", he.get("n_pesquisas"), "| spread hoje:", he.get("spread_hoje"), "±", he.get("se_hoje"))
    print("house effects:", he.get("house"))
    print("prior viés:", vb["prior_vies"], "| media por eleição:", vb["media_por_eleicao"])
    print("transferência: spread", tr["spread"], "P(F)=", tr["p_flavio"])
    print("histórico: m1=", hist["m1_2026"], "pred m2=", hist["pred_m2_2026"], "reg=", hist["regressao"], "laplace=", hist["laplace_p_lider"])
    for k, v in mc["modelos"].items():
        print(f"MC {k}: P(F)={v['p_flavio']} média={v['media']} [{v['p05']}, {v['p95']}]")
    print("bayes posterior:", by["posterior"])
    print("senado distâncias:", cg["senado"]["distancias"])
    print("estados:", uf.get("n"), uf.get("swing_medio_ponderado"))


if __name__ == "__main__":
    main()