#!/usr/bin/env python3
"""49_suite_jepa.py — Suíte JEPA v0.5.0: 10 variantes + SonicMoE + ReplaySSM.

Cada variante recebe dados sintéticos HONESTOS com estrutura previsível da
modalidade; o critério de aprendizado é a QUEDA da perda preditiva com o
std das incorporações > 0,05 (sem colapso — Teo J-1). Provas em r-20.

  Image-JEPA  : imagem sintética (3,16,16) com blocos → patches 2D
  V-JEPA 2    : vídeo sintético (4,8,8) em movimento → tubos 3D
  Audio-JEPA : espectrograma (16,32) com harmônicos → bandas
  Point-JEPA : nuvem de pontos em anel → grupos FPS
  3D-JEPA    : volume (8,8,8) com esfera → voxels
  Brain-JEPA : sinais (4,32) tipo EEG → canal×tempo
  T-JEPA     : série temporal AR(1) multicanal → patches 1D
  LeJEPA     : edição = translação conhecida (instrução, não resultado)
  MC-JEPA    : modalidade A = sinal, B = versão não-linear de A
  BJEPA      : sequência periódica → causal + anti-causal + ciclo
  SonicMoE   : capacidade + balanceamento (Teos SM-1/SM-2)
  ReplaySSM  : estabilidade + replay (Teos R-1/R-2)
Saída: resultados/suite_jepa_v050.json
"""
import json, math, os, random, sys, time

def main():
    import khtst_rust.nucleo_rust as nr

    pasta = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    res = {"versao": "0.5.0", "variantes": {}}
    rng = random.Random(42)

    def relatorio(nome, r, t0):
        # critério honesto: a perda PREDITIVA (sem regularizador VC) cai E o
        # std das incorporações fica > 0,05 (sem colapso — Teo J-1)
        pi = r.get("perda_pred_inicial", r["perda_inicial"])
        pf = r.get("perda_pred_final", r["perda_final"])
        if pi is None: pi = r["perda_inicial"]
        if pf is None: pf = r["perda_final"]
        queda = 100.0 * (1 - pf / pi) if (pi and pi > 0) else 0.0
        dt = time.time() - t0
        ok = (pf is not None and pi is not None and pf < pi) and r.get("colapso_evitado", True)
        print(f"  {nome:12s} perda_pred {pi:.4f} -> {pf:.4f} "
              f"(-{queda:.1f}%) | std_emb {r.get('std_emb_final', 0):.3f} | {'OK' if ok else 'FALHOU'} | {dt:.1f}s", flush=True)
        res["variantes"][nome] = {**r, "queda_pct": queda, "tempo_s": dt, "gate_aprendeu": ok}

    cfg = json.dumps({"epocas": 300, "d_emb": 24, "lr": 2e-3, "m": 0.99, "gamma_vc": 0.3, "semente": 42})

    # ------------------------------------------------------------- Image-JEPA
    print("=== Image-JEPA (patches 2D) ===", flush=True)
    # imagem rica: gradiente diagonal + círculo + faixa (estrutura previsível
    # mas não-trivial — 8x8 patches de 2x2)
    img = [0.0] * (3 * 16 * 16)
    for y in range(16):
        for x in range(16):
            d2 = (x - 8) ** 2 + (y - 8) ** 2
            circ = 1.0 if d2 < 18 else 0.0
            faixa = 0.5 if 12 <= x < 15 else 0.0
            grad = 0.03 * (x + y)
            for c in range(3):
                img[(c * 16 + y) * 16 + x] = 0.15 + circ + faixa + grad + rng.gauss(0, 0.04) + 0.08 * c
    t0 = time.time()
    patches = nr.patches_2d([float(v) for v in img], 3, 16, 16, 2, 2)
    r = nr.treinar_jepa([[float(x) for x in p] for p in patches], [8, 8], "image", cfg)
    relatorio("image", r, t0)

    # ------------------------------------------------------------- V-JEPA 2
    print("=== V-JEPA 2 (tubos 3D) ===", flush=True)
    vid = []
    for t in range(4):
        quadro = []
        for y in range(8):
            for x in range(8):
                cx = 2 + t  # objeto em movimento
                v = 1.0 if (cx <= x < cx + 3 and 2 <= y < 6) else 0.1
                quadro.append(v + rng.gauss(0, 0.05))
        vid.append(quadro)
    # forma (T, H, W) achatada: adaptador de tubos via patches_2d por quadro
    # não exposto — usamos a API de série por quadro: tubo = patches do quadro
    t0 = time.time()
    tubos = []
    for quadro in vid:
        tp = nr.patches_2d([float(v) for v in quadro], 1, 8, 8, 4, 4)
        tubos.append([[float(x) for x in p] for p in tp])
    # tubos (T'=4, H'=2, W'=2): cada tubo = bloco 4x4 do quadro t
    patches_v = [[float(v) for v in tubos[t][k]]
                 for t in range(4) for k in range(len(tubos[t]))]
    r = nr.treinar_jepa(patches_v, [4, 2, 2], "vjepa2", cfg)
    relatorio("vjepa2", r, t0)

    # ------------------------------------------------------------- Audio-JEPA
    print("=== Audio-JEPA (bandas de espectrograma) ===", flush=True)
    spec = []
    for f in range(16):
        linha = []
        for t in range(32):
            v = math.sin(2 * math.pi * (2 if f < 6 else 8) * t / 32.0) * (1 if f < 6 else 0.5)
            linha.append(v + rng.gauss(0, 0.1))
        spec.append(linha)
    # patches do espectrograma: grade (F'=16/4, T'=32/8) via patches_2d
    t0 = time.time()
    flat = [float(v) for linha in spec for v in linha]
    patches_a = nr.patches_2d(flat, 1, 16, 32, 4, 8)  # (4, 4) = 16 patches
    r = nr.treinar_jepa([[float(x) for x in p] for p in patches_a], [4, 4], "audio", cfg)
    relatorio("audio", r, t0)

    # ------------------------------------------------------------- Point-JEPA
    print("=== Point-JEPA (grupos FPS) ===", flush=True)
    pc = []
    for i in range(90):
        a = 2 * math.pi * i / 90
        pc.append([math.cos(a) * 3, math.sin(a) * 3, rng.gauss(0, 0.2)])
    for i in range(60):  # segundo anel
        a = 2 * math.pi * i / 60
        pc.append([math.cos(a) * 1 + 5, math.sin(a) * 1, rng.gauss(0, 0.2)])
    t0 = time.time()
    grupos = nr.grupos_pontos([[float(v) for v in p] for p in pc], 12, 6, 42)
    r = nr.treinar_jepa(grupos, [12], "point", cfg)
    relatorio("point", r, t0)

    # ------------------------------------------------------------- 3D-JEPA
    print("=== 3D-JEPA (voxels) ===", flush=True)
    vol = []
    for x in range(8):
        plano = []
        for y in range(8):
            for z in range(8):
                d2 = (x - 4) ** 2 + (y - 4) ** 2 + (z - 4) ** 2
                v = 1.0 if d2 < 9 else 0.1
                plano.append(v + rng.gauss(0, 0.05))
        vol.append(plano)
    t0 = time.time()
    patches_vx = []
    for plano in vol:
        tp = nr.patches_2d([float(v) for v in plano], 1, 8, 8, 4, 4)
        patches_vx.append([[float(x) for x in p] for p in tp])
    # voxel-patches: bloco 4x4 do plano x (grade X'=8, H'=2, W'=2)
    patches_3d = [[float(v) for v in patches_vx[x][k]]
                  for x in range(8) for k in range(len(patches_vx[x]))]
    r = nr.treinar_jepa(patches_3d, [8, 2, 2], "jepa3d", cfg)
    relatorio("jepa3d", r, t0)

    # ------------------------------------------------------------- Brain-JEPA
    print("=== Brain-JEPA (canal×tempo) ===", flush=True)
    sinais = []
    for ch in range(4):
        fase = ch * 0.5
        sinais.append([math.sin(2 * math.pi * (1 + ch * 0.25) * t / 32.0 + fase) + rng.gauss(0, 0.15) for t in range(64)])
    t0 = time.time()
    flat = [float(v) for s in sinais for v in s]
    patches_b = nr.patches_serie(flat, 4, 64, 8)
    r = nr.treinar_jepa([[float(x) for x in p] for p in patches_b], [4, 8], "brain", cfg)
    relatorio("brain", r, t0)

    # ------------------------------------------------------------- T-JEPA
    print("=== T-JEPA (série temporal canal-independente) ===", flush=True)
    series = []
    for ch in range(6):
        x, a = 0.0, 0.6 + 0.05 * ch
        linha = []
        for t in range(96):
            x = a * x + rng.gauss(0, 0.3) + 0.2 * math.sin(t / 8.0)
            linha.append(x)
        series.append(linha)
    t0 = time.time()
    flat = [float(v) for s in series for v in s]
    patches_t = nr.patches_serie(flat, 6, 96, 8)
    r = nr.treinar_jepa([[float(x) for x in p] for p in patches_t], [6, 12], "tjepa", cfg)
    relatorio("tjepa", r, t0)

    # ------------------------------------------------------------- LeJEPA
    print("=== LeJEPA (edição-condicionado) ===", flush=True)
    dados_l, edits_l, editados_l = [], [], []
    for _ in range(120):
        x = [rng.gauss(0, 1) for _ in range(6)]
        e = [rng.uniform(-2, 2), rng.uniform(-2, 2)]  # instrução: deslocamento
        y = [xj + e[0] + 0.3 * e[1] * xj for xj in x]  # edição aplicada
        dados_l.append([float(v) for v in x])
        edits_l.append([float(v) for v in e])
        editados_l.append([float(v) for v in y])
    t0 = time.time()
    r = nr.treinar_lejepa(dados_l, edits_l, editados_l, cfg)
    relatorio("lejepa", r, t0)

    # ------------------------------------------------------------- MC-JEPA
    print("=== MC-JEPA (cruz-modal) ===", flush=True)
    a_m, b_m = [], []
    for _ in range(120):
        s = [rng.gauss(0, 1) for _ in range(6)]
        a_m.append([float(v) for v in s])
        b_m.append([float(math.tanh(v) * 0.8 + 0.2 * s[0]) for v in s])  # B = f(A)
    t0 = time.time()
    r = nr.treinar_mcjepa(a_m, b_m, cfg)
    relatorio("mcjepa", r, t0)

    # ------------------------------------------------------------- BJEPA
    print("=== BJEPA (bidirecional com ciclo) ===", flush=True)
    seq = []
    for t in range(48):
        seq.append([math.sin(2 * math.pi * t / 12.0) + 0.5 * math.cos(2 * math.pi * t / 6.0) + rng.gauss(0, 0.05) for _ in range(4)])
    t0 = time.time()
    r = nr.treinar_bjepa([[float(v) for v in s] for s in seq], json.dumps({"epocas": 300, "d_emb": 24, "lr": 2e-3, "gamma_vc": 0.3}))
    relatorio("bjepa", r, t0)

    # ------------------------------------------------------------- SonicMoE
    print("=== SonicMoE (IO/tile-aware MoE) ===", flush=True)
    t0 = time.time()
    moe = nr.demo_sonic_moe(128, 16, 4, 2, 40)
    dt = time.time() - t0
    ok_moe = moe["capacidade_respeitada"] and moe["perda_aux"] > 0.5
    print(f"  carga: {moe['carga_por_especialista']} | aux: {moe['perda_aux']:.3f} | descartados: {moe['descartados']} | {'OK' if ok_moe else 'FALHOU'}")
    res["variantes"]["sonicmoe"] = {**moe, "tempo_s": dt, "gate_aprendeu": ok_moe}

    # ------------------------------------------------------------- ReplaySSM
    print("=== ReplaySSM (SSM com replay) ===", flush=True)
    t0 = time.time()
    ssm = nr.demo_replay_ssm(40, 0.5)
    dt = time.time() - t0
    ok_ssm = ssm["magnitudes_finitas"]
    print(f"  magnitudes finitas pós-troca: {ssm['magnitudes_finitas']} (max {ssm['magnitude_max_pos_troca']:.3f}) | {'OK' if ok_ssm else 'FALHOU'}")
    res["variantes"]["replayssm"] = {**ssm, "tempo_s": dt, "gate_aprendeu": ok_ssm}

    # ---------------------------------------------------------------- resumo
    gates = [(n, v.get("gate_aprendeu", False)) for n, v in res["variantes"].items()]
    ok = sum(1 for _, v in gates if v)
    res["resumo"] = {"variantes_totais": len(gates), "variantes_ok": ok,
                     "todas_verdes": ok == len(gates)}
    print(f"\n=== RESUMO: {ok}/{len(gates)} variantes com aprendizado comprovado ===")

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

if __name__ == "__main__":
    main()