#!/usr/bin/env python3
"""58_autoajuste_iterativo.py — v0.6.0: LOOP DE APRIMORAMENTO contínuo
(análise de telemetria → autoajuste de parâmetros → seleção de arquitetura
por UCB1 → repetição até convergência marginal; provas r-26).

Mecanismo (dados REAIS da Wikipédia PT-BR):
  - BRAÇOS de arquitetura: {mlp-pequena, moe-e4, moe-e8-redirect} — o
    redirecionável troca papéis cod↔dec a cada rodada (M-R1/M-R4);
  - UCB1 (A-4) escolhe o braço: r̄_a + √(2 ln T / n_a) — arrependimento
    sublinear; recompensa = melhora relativa da perda preditiva − custo de
    tempo normalizado;
  - AUTOAJUSTE por telemetria (A-1..A-5): veredito Theil–Sen da série do
    braço regula a fração de máscara (estagnado ⇒ explora mais;
    divergindo ⇒ explora menos) — parâmetro REGULADO por sinal medido;
  - PARADA MARGINAL (A-2 no nível meta): para quando a melhora marginal do
    melhor braço em 3 avaliações < 2%;
  - AUDITORIA: cada decisão (escolha UCB1, regulação, parada) registrada.

Saída: resultados/autoajuste_iterativo_v060.json
"""
import json, math, os, time

def ucb1_escolher(stats, t):
    """A-4: braço com maior r̄ + √(2 ln T / n); braços não-jogados primeiro."""
    for a, (n, _) in stats.items():
        if n == 0:
            return a
    ln_t = math.log(max(t, 2))
    melhor, melhor_v = None, -1e18
    for a, (n, soma) in stats.items():
        v = soma / n + math.sqrt(2.0 * ln_t / n)
        if v > melhor_v:
            melhor_v, melhor = v, a
    return melhor

def theil_sen(ys):
    if len(ys) < 3:
        return 0.0, "amostra_insuficiente"
    incl = []
    for i in range(len(ys)):
        for j in range(i + 1, len(ys)):
            incl.append((ys[j] - ys[i]) / (j - i))
    incl.sort()
    med = incl[len(incl) // 2]
    m = sum(ys) / len(ys)
    rel = med / max(abs(m), 1e-9)
    return med, ("aprendendo" if rel < -1e-4 else ("divergindo" if rel > 1e-4 else "estagnado"))

def main():
    import khtst_rust.nucleo_rust as nr

    pasta = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    res = {"versao": "0.6.0", "modulo": nr.versao(), "auditoria": [], "rodadas": []}

    # ---------- dados REAIS ----------
    with open(os.path.join(pasta, "dados", "shard_validacao.jsonl")) as f:
        textos = [json.loads(l)["texto"] for l in f][:3]
    serie = []
    for t in textos:
        serie += [min(ord(c), 255) / 255.0 for c in t[:1600] if c.strip()]
    patches = nr.patches_serie(serie, 1, len(serie), 16)
    print(f"dados reais: {len(patches)} patches ({len(serie)} caracteres, 3 artigos)", flush=True)
    res["dados_reais"] = {"n_patches": len(patches), "fonte": "wikimedia/wikipedia 20231101.pt"}

    # ---------- estado do loop ----------
    bracos = {
        "mlp-pequena": {"modo": "pequena", "frac": 0.45, "serie_perda": [], "tempo_ref": None},
        "moe-e4": {"modo": "moe", "frac": 0.45, "serie_perda": [], "tempo_ref": None, "nota": "4 experts/papel"},
        "moe-e8-redirect": {"modo": "moe", "frac": 0.45, "serie_perda": [], "tempo_ref": None, "nota": "8 experts + redirect por rodada"},
    }
    stats = {a: [0, 0.0] for a in bracos}  # n, soma_recompensa
    t_global = 0
    melhor_melhora, melhor_braco = 0.0, None
    teto_rodadas = 14
    parou_por_margem = False

    print("\n=== LOOP: UCB1 (A-4) + regulação por telemetria (Theil–Sen) + parada marginal (A-2) ===", flush=True)
    for rodada in range(teto_rodadas):
        braco = ucb1_escolher({a: tuple(v) for a, v in stats.items()}, t_global)
        est = bracos[braco]
        # regulação pré-rodada: veredito Theil–Sen (≥3 pontos) ou variação
        # local (2 pontos) da SÉRIE do braço — parâmetro REGULADO por sinal
        reg = None
        if len(est["serie_perda"]) >= 3:
            _, vered = theil_sen(est["serie_perda"])
            if vered == "estagnado" and est["frac"] < 0.75:
                est["frac"] = min(0.75, est["frac"] + 0.10)
                reg = f"estagnado ⇒ frac_mascara ↑ para {est['frac']:.2f} (explora mais)"
            elif vered == "divergindo" and est["frac"] > 0.25:
                est["frac"] = max(0.25, est["frac"] - 0.10)
                reg = f"divergindo ⇒ frac_mascara ↓ para {est['frac']:.2f} (estabiliza)"
        elif len(est["serie_perda"]) == 2:
            a_, b_ = est["serie_perda"]
            if b_ > a_ * 1.02 and est["frac"] > 0.25:
                est["frac"] = max(0.25, est["frac"] - 0.10)
                reg = f"piora local ({100*(b_/a_-1):.1f}%) ⇒ frac_mascara ↓ para {est['frac']:.2f}"
            elif abs(b_ - a_) < 0.01 * abs(a_) and est["frac"] < 0.75:
                est["frac"] = min(0.75, est["frac"] + 0.10)
                reg = f"estagnação local ⇒ frac_mascara ↑ para {est['frac']:.2f}"
        if reg:
            res["auditoria"].append({"rodada": rodada, "tipo": "regulacao", "braco": braco, "acao": reg})
        # treina o braço (passos moderados p/ orçamento de tempo)
        t0 = time.time()
        r = nr.jepa_com_moe(patches, est["modo"], passos=70, frac_mascara=est["frac"], semente=1000 + rodada)
        dt = time.time() - t0
        melhora = (r["perda_pred_inicial"] - r["perda_pred_final"]) / max(r["perda_pred_inicial"], 1e-9)
        if est["tempo_ref"] is None:
            est["tempo_ref"] = dt
        custo = 0.5 * (dt / max(est["tempo_ref"], 1e-9) - 1.0)  # custo RELATIVO ao 1º do braço
        recompensa = max(0.0, melhora - max(0.0, custo))
        stats[braco][0] += 1
        stats[braco][1] += recompensa
        t_global += 1
        est["serie_perda"].append(r["perda_pred_final"])
        if melhora > melhor_melhora:
            melhor_melhora, melhor_braco = melhora, braco
        res["rodadas"].append({
            "rodada": rodada, "braco": braco, "modo": est["modo"], "frac_mascara": est["frac"],
            "perda_inicial": r["perda_pred_inicial"], "perda_final": r["perda_pred_final"],
            "melhora_rel": melhora, "recompensa": recompensa, "tempo_s": dt,
            "std_emb": r["std_emb"], "regulacao": reg,
        })
        print(f"  r{rodada:02d} {braco:16s} frac={est['frac']:.2f} perda {r['perda_pred_inicial']:.4f} -> "
              f"{r['perda_pred_final']:.4f} (melhora {100*melhora:.1f}%, {dt:.1f}s)"
              + (f" | REG: {reg}" if reg else ""), flush=True)
        # parada marginal (A-2 meta): só após 10 jogadas globais (a regulação
        # precisa de espaço p/ agir); melhor braço sem melhora marginal em 3 avaliações
        serie_mb = bracos[melhor_braco]["serie_perda"]
        if t_global >= 10 and stats[melhor_braco][0] >= 3 and len(serie_mb) >= 3:
            janelas = serie_mb[-3:]
            margem = (janelas[0] - janelas[-1]) / max(abs(janelas[0]), 1e-9)
            if margem < 0.02:
                res["auditoria"].append({
                    "rodada": rodada, "tipo": "parada_marginal", "braco": melhor_braco,
                    "acao": f"margem {100*margem:.2f}% < 2% em 3 avaliações — convergiu (A-2 meta)"})
                parou_por_margem = True
                print(f"  ⇒ PARADA MARGINAL: {melhor_braco} convergiu (margem {100*margem:.2f}% < 2%)", flush=True)
                break

    # ---------- consolidação ----------
    medias = {a: (stats[a][1] / stats[a][0] if stats[a][0] else 0.0) for a in stats}
    vencedor = max(medias, key=medias.get)
    base = bracos["mlp-pequena"]["serie_perda"]
    final_moe = bracos[vencedor]["serie_perda"]
    gap = (sum(base[-3:]) / max(len(base[-3:]), 1) - sum(final_moe[-3:]) / max(len(final_moe[-3:]), 1)) \
        / max(sum(base[-3:]) / max(len(base[-3:]), 1), 1e-9)
    # arrependimento empírico do UCB1 (auditoria — A-4)
    melhor_media = max(medias.values())
    arrepend = sum(stats[a][0] * (melhor_media - medias[a]) for a in stats) / max(t_global, 1)
    print(f"\n=== CONSOLIDAÇÃO ===")
    print(f"  vencedor UCB1: {vencedor} (recompensa média {medias[vencedor]:.4f})")
    print(f"  gap final vs baseline (mlp-pequena): {100*gap:.1f}% de perda preditiva menor")
    print(f"  arrependimento médio UCB1: {arrepend:.4f} (sublinear — Auer)")
    print(f"  rodadas executadas: {len(res['rodadas'])} | parada marginal: {parou_por_margem}")

    # ---------- autoajuste A-1..A-5 (demo auditada — r-21 §3.5 fechado) ----------
    print("\n=== Autoajuste A-1..A-5 (τ_c, B, N, g, UCB1) ===", flush=True)
    auto = nr.demo_autoajuste(semente=42)
    print(f"  τ_c={auto['tau_c_final']:.3f} | N={auto['n_tentativas_final']} | "
          f"B={auto['orcamento_b_final']} | g={auto['gate_final']:.1f} | "
          f"ajustes auditados: {auto['ajustes_auditados']}")
    print(f"  UCB1: melhor braço {auto['ucb1_melhor_braco']} | arrependimento médio {auto['ucb1_arrependimento_medio']:.4f}")

    # ---------- gates ----------
    gates = {
        "gate_ucb1_encontrou_moe": vencedor.startswith("moe"),
        "gate_gap_positivo_vs_baseline": gap > 0.0,
        "gate_arrependimento_pequeno": arrepend < 0.15,
        "gate_regulacao_ocorreu": any(a["tipo"] == "regulacao" for a in res["auditoria"]),
        "gate_parada_ou_teto": parou_por_margem or len(res["rodadas"]) == teto_rodadas,
        "gate_melhor_aprendeu": melhor_melhora > 0.0,
    }
    gates.update({g[5:]: v for g, v in auto["gates"].items()})
    res["consolidacao"] = {
        "vencedor": vencedor, "medias_recompensa": medias, "gap_vs_baseline_pct": 100 * gap,
        "arrependimento_medio": arrepend, "parou_por_margem": parou_por_margem,
        "melhor_melhora_rel": melhor_melhora,
    }
    res["autoajuste_demo"] = auto
    res["gates"] = gates
    ok = sum(1 for v in gates.values() if v)
    res["resumo"] = {"gates_totais": len(gates), "gates_ok": ok, "todos_verdes": ok == len(gates)}
    print(f"\n=== RESUMO: {ok}/{len(gates)} gates verdes ===")
    for g, v in gates.items():
        print(f"  {'OK ' if v else 'FALHOU'} {g}")

    os.makedirs(os.path.join(pasta, "resultados"), exist_ok=True)
    with open(os.path.join(pasta, "resultados", "autoajuste_iterativo_v060.json"), "w") as f:
        json.dump(res, f, indent=1, ensure_ascii=False)
    print("salvo: resultados/autoajuste_iterativo_v060.json")
    if not res["resumo"]["todos_verdes"]:
        raise SystemExit(1)

if __name__ == "__main__":
    main()