#!/usr/bin/env python3
"""57_moe_redirecionavel.py — v0.6.0: MoEs de codificação/decodificação
REDIRECIONÁVEIS com Mecanismo de Atenção próprio independente (paralelo)
para as variantes JEPA e Kohonen SOM (provas r-25).

Demos (com dados REAIS da Wikipédia PT-BR):
  1. M-R1..M-R5: neutralidade no nascimento, invariância sob
     redirecionamento, paralelismo medido (2 threads), water-filling;
  2. JEPA com MoE — ablação de 3 braços (pequena / densa / moe) em patches
     de texto REAL (a "densa" isola o ganho do roteamento);
  3. SOM 4D refinado por MoE com alocação por QE (water-filling vs
     uniforme) + redirecionamento p/ decodificador;
  4. AcopladorDuplo: protocolo duplex Pensamento↔backbone (AC-1/AC-2).

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

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(), "demos": {}}

    # ---------- dados REAIS: texto da Wikipédia PT-BR → patches ----------
    with open(os.path.join(pasta, "dados", "shard_validacao.jsonl")) as f:
        texto = json.loads(f.readline())["texto"]
    # normaliza para [0,1] por código de caractere (canal único)
    serie = [min(ord(c), 255) / 255.0 for c in texto[:2048] if c.strip()]
    patches = nr.patches_serie(serie, 1, len(serie), 16)  # (n, 16) reais
    print(f"dados reais: {len(patches)} patches de 16 dims "
          f"({len(texto[:2048])} caracteres da Wikipédia PT-BR)", flush=True)
    res["dados_reais"] = {"n_patches": len(patches), "d_patch": 16,
                          "fonte": "wikimedia/wikipedia 20231101.pt (validação)"}

    # -------------------------------------------- 1: M-R1..M-R5
    print("\n=== MoE redirecionável: provas M-R1..M-R5 ===", flush=True)
    t0 = time.time()
    m = nr.demo_moe_redirecionavel(d=16, n=96, semente=42)
    dt = time.time() - t0
    print(f"  M-R2 neutralidade no nascimento: gate={m['valor_gate_no_nascimento']}, dif={m['neutralidade_dif']:.2e}")
    print(f"  M-R1 capacidade: {m['carga_antes']} == {m['carga_depois']} | perda aux {m['perda_aux_antes']:.4f} ~ {m['perda_aux_depois']:.4f}")
    print(f"  M-R3 paralelismo: t_par={m['t_paralelo_s']*1000:.1f}ms < t_seq={m['t_sequencial_s']*1000:.1f}ms (speedup {m['speedup']:.2f}x)")
    print(f"  M-R5 water-filling: cota {m['cota_wf']:.4f} <= uniforme {m['cota_uniforme']:.4f} | alocacao {m['alocacao_water_filling']}")
    res["demos"]["provas_mr"] = {**m, "tempo_s": dt}

    # -------------------------------------------- 2: JEPA com MoE (3 braços)
    print("\n=== JEPA com MoE redirecionável (3 braços, dados reais) ===", flush=True)
    bracos = {}
    for modo in ("pequena", "densa", "moe"):
        t0 = time.time()
        r = nr.jepa_com_moe(patches, modo, passos=90, frac_mascara=0.45, semente=7)
        dt = time.time() - t0
        bracos[modo] = r
        print(f"  {modo:8s}: perda preditiva {r['perda_pred_inicial']:.4f} -> {r['perda_pred_final']:.4f} "
              f"(queda {r['queda_abs']:.4f}, {dt:.1f}s, std_emb {r['std_emb']:.3f})")
    moe_melhor_densa = bracos["moe"]["perda_pred_final"] < bracos["densa"]["perda_pred_final"]
    moe_melhor_pequena = bracos["moe"]["perda_pred_final"] < bracos["pequena"]["perda_pred_final"]
    res["demos"]["jepa_moe"] = {
        "bracos": bracos,
        "moe_melhor_que_densa": moe_melhor_densa,
        "moe_melhor_que_pequena": moe_melhor_pequena,
        "gate_jepa_moe_aprende": bracos["moe"]["aprendeu"],
        "gate_jepa_moe_vs_densa": moe_melhor_densa,
        "gate_jepa_moe_vs_pequena": moe_melhor_pequena,
    }

    # -------------------------------------------- 3: SOM refinado por MoE
    print("\n=== SOM 4D refinado por MoE (water-filling + redirect) ===", flush=True)
    t0 = time.time()
    patches_f64 = [[float(x) for x in p] for p in patches]
    s = nr.som_com_moe(patches_f64, lado=2, passos=90, semente=42)
    dt = time.time() - t0
    print(f"  QE: {s['qe_antes']:.4f} -> {s['qe_depois_water_filling']:.4f} "
          f"(water-filling, {s['reducao_qe_pct']:.1f}% | uniforme: {s['qe_depois_uniforme']:.4f})")
    print(f"  carga por célula (water-filling ∝ QE): {s['carga_por_celula']}")
    print(f"  MSE do decodificador pós-redirect: {s['mse_decodificador_pos_redirect']:.6f} "
          f"(vs QE² {s['qe_antes']**2:.6f})")
    res["demos"]["som_moe"] = {**s, "tempo_s": dt}

    # -------------------------------------------- 4: AcopladorDuplo
    print("\n=== AcopladorDuplo: protocolo duplex Pensamento↔backbone ===", flush=True)
    t0 = time.time()
    a = nr.demo_acoplador_duplo(rodadas=120, timeout_ms=1500)
    dt = time.time() - t0
    print(f"  rodadas: {a['rodadas']} | timeouts: {a['timeouts']} | tempo {a['tempo_total_ms']:.1f}ms")
    print(f"  duplex simultâneo: A→B {a['envios_a_para_b']} == B→A {a['envios_b_para_a']} | "
          f"sobreposição máx {a['maior_sobreposicao_us']}µs")
    res["demos"]["acoplador_duplo"] = {**a, "tempo_s": dt}

    # ------------------------------------------------------------------ resumo
    gates = []
    for nome, sd in res["demos"].items():
        gates += [(f"{nome}.{g[5:]}", val) for g, val in sd.items() if g.startswith("gate_")]
        if isinstance(sd.get("gates"), dict):
            gates += [(f"{nome}.{g[5:]}", val) for g, val in sd["gates"].items() if g.startswith("gate_")]
    ok = sum(1 for _, val in gates if val)
    res["resumo"] = {"gates_totais": len(gates), "gates_ok": ok, "todos_verdes": ok == len(gates)}
    print(f"\n=== RESUMO: {ok}/{len(gates)} gates verdes ===")

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

if __name__ == "__main__":
    main()