#!/usr/bin/env python3
"""40_treinar.py — Fases A (tokenizador) e B (gerador) do KHTST_RUST.

Uso: python3 scripts/40_treinar.py [--fase a|b|ab] [--teto_min N]
Protocolo (n-04 §4): A = nested dropout + commitment + GSQ/RCO Stiefel;
B = MLE exato sobre hierarquias congeladas + corrupção κ + MTP + JEPA + PCGrad.
Teto de tempo wall-clock (leitura honesta): para no limite e reporta."""
import argparse, json, os, sys, time

def carregar_sequencias(pasta, bpe, qual):
    import khtst_rust as K
    b = K.Bpe.carregar(os.path.join(pasta, "dados", "bpe.json"))
    seqs = []
    with open(os.path.join(pasta, "dados", f"shard_{qual}.jsonl")) as f:
        for linha in f:
            t = json.loads(linha)["texto"]
            ids = b.codificar(t)
            # janelas de 256 tokens (ctx) — deslocamento 128 (50% sobreposição)
            for i in range(0, max(1, len(ids) - 256), 128):
                seqs.append(ids[i:i + 256])
    return seqs

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--fase", default="ab", choices=["a", "b", "ab"])
    ap.add_argument("--teto_min", type=float, default=10.0)
    args = ap.parse_args()

    pasta = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.insert(0, pasta)
    os.makedirs(os.path.join(pasta, "resultados"), exist_ok=True)
    import khtst_rust as K

    cfg = K.Config(os.path.join(pasta, "configs", "base_rust.json"))
    bpe = K.Bpe.carregar(os.path.join(pasta, "dados", "bpe.json"))

    # ajusta passos ao teto de tempo: estimativa por benchmark rápido
    relatorio = {"fases": {}}

    if args.fase in ("a", "ab"):
        print("=== FASE A — tokenizador de conceitos (GSQ+RCO Stiefel) ===", flush=True)
        seqs = carregar_sequencias(pasta, bpe, "treino")
        print(f"sequências de treino: {len(seqs)} (ctx=256)", flush=True)
        cb = None
        t0 = time.time()
        saida = K.treinar_fase_a(
            cfg, seqs,
            caminho_telemetria=os.path.join(pasta, "resultados", "telemetria_fase_a.jsonl"),
            caminho_checkpoint=os.path.join(pasta, "resultados", "tokenizador_conceitos.json"),
        )
        dt = time.time() - t0
        r = saida["resumo"]
        relatorio["fases"]["a"] = r
        print(json.dumps(r, indent=1, ensure_ascii=False))
        print(f"fase A: {dt:.0f}s | Δperda {r['perda_inicial']:.3f}→{r['perda_final']:.3f} "
              f"| acc_recon {r['acc_recon']:.3f} | δ_ort_max {r['delta_ort_max']:.2e} "
              f"| perp_codebook {r['perplexidade_codebook']:.1f} "
              f"| monotonia_prefixo {r['monotonia_prefixo']:.2f} "
              f"| escore_aprendizado {r['escore_aprendizado']:.3f}", flush=True)

    if args.fase in ("b", "ab"):
        print("=== FASE B — gerador next-concept (MLE exato + κ + MTP + JEPA + PCGrad) ===", flush=True)
        seqs = carregar_sequencias(pasta, bpe, "treino")
        t0 = time.time()
        saida = K.treinar_fase_b(
            cfg, seqs,
            caminho_tokenizador=os.path.join(pasta, "resultados", "tokenizador_conceitos.json"),
            caminho_telemetria=os.path.join(pasta, "resultados", "telemetria_fase_b.jsonl"),
            caminho_checkpoint=os.path.join(pasta, "resultados", "gerador_conceitos.json"),
        )
        dt = time.time() - t0
        r = saida["resumo"]
        relatorio["fases"]["b"] = r
        print(json.dumps(r, indent=1, ensure_ascii=False))
        print(f"fase B: {dt:.0f}s | ΔCE {r['perda_inicial']:.3f}→{r['perda_final']:.3f} "
              f"| acc_conceitos {r['acc_conceitos']:.3f} "
              f"| escore_aprendizado {r['escore_aprendizado']:.3f}", flush=True)

    with open(os.path.join(pasta, "resultados", "resumo_treino.json"), "w") as f:
        json.dump(relatorio, f, indent=1, ensure_ascii=False)
    print("resumo salvo em resultados/resumo_treino.json")

if __name__ == "__main__":
    main()