#!/usr/bin/env python3
"""56_treino_conjunto.py — v0.6.0: FASE C — treino conjunto (ajuste fino)
Pensamento↔backbone com VERIFICADOR NCM (sem oráculo) e ACOPLADOR DUPLEX
(provas r-24; fecha os débitos r-21 §3.1/§3.3).

Pipeline com dados REAIS (Wikipédia PT-BR):
  fase A — tokenizador GSQ+RCO (400 passos);
  fase B — gerador NCM (400 passos, teto 300s) → checkpoint p/ ajuste fino;
  fase C — treino CONJUNTO: gatilho por confiança calibrada (T-P1/V-2),
           Pensamento corrige ẑ nas sequências gatilhadas, gradiente flui
           nos DOIS sentidos, PCGrad entre perdas, acoplador duplex por
           passo, autoajuste A-1..A-5 pela telemetria.

Saída: resultados/treino_conjunto_v060.json + telemetria_fase_c.jsonl
"""
import json, os, sys, time

def main():
    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"))
    ck_tok = os.path.join(pasta, "resultados", "tokenizador_conceitos.json")
    ck_ger = os.path.join(pasta, "resultados", "gerador_conceitos.json")

    bpe = K.Bpe.carregar(os.path.join(pasta, "dados", "bpe.json"))
    seqs = []
    with open(os.path.join(pasta, "dados", "shard_treino.jsonl")) as f:
        for linha in f:
            ids = bpe.codificar(json.loads(linha)["texto"])
            for i in range(0, max(1, len(ids) - 256), 128):
                seqs.append(ids[i:i + 256])
    seqs = seqs[:1400]
    rel = {"versao": "0.6.0", "dados": {"sequencias": len(seqs),
          "fonte": "wikimedia/wikipedia 20231101.pt"}, "fases": {}}
    print(f"sequências de treino: {len(seqs)} (ctx=256)", flush=True)

    # ------------------------------------------------ fase A (se preciso)
    if not os.path.exists(ck_tok):
        print("=== FASE A — tokenizador de conceitos (GSQ+RCO Stiefel) ===", flush=True)
        t0 = time.time()
        saida = K.treinar_fase_a(cfg, seqs, caminho_telemetria=None,
                                 caminho_checkpoint=ck_tok)
        dt = time.time() - t0
        r = saida["resumo"]
        rel["fases"]["a"] = {**r, "tempo_s": dt}
        print(f"  perda {r['perda_inicial']:.3f} -> {r['perda_final']:.3f} | "
              f"acc_recon {r['acc_recon']:.3f} | δ_ort {r['delta_ort_max']:.1e} | "
              f"perp {r['perplexidade_codebook']:.1f} | {dt:.0f}s", flush=True)
    else:
        print("checkpoint fase A presente — reutilizando (ajuste fino)", flush=True)

    # ------------------------------------------------ fase B (se preciso)
    if not os.path.exists(ck_ger):
        print("=== FASE B — gerador NCM (MLE exato) ===", flush=True)
        t0 = time.time()
        saida_b = K.treinar_fase_b(cfg, seqs, caminho_tokenizador=ck_tok,
                                   caminho_telemetria=None,
                                   caminho_checkpoint=ck_ger, teto_seg=300.0)
        dt = time.time() - t0
        rb = saida_b["resumo"]
        rel["fases"]["b"] = {**rb, "tempo_s": dt}
        print(f"  CE {rb['perda_inicial']:.3f} -> {rb['perda_final']:.3f} | "
              f"acc_conceitos {rb.get('acuracia_media', 0):.3f} | {dt:.0f}s", flush=True)
    else:
        print("checkpoint fase B presente — reutilizando (ajuste fino)", flush=True)

    # ------------------------------------------------ fase C (NOVIDADE)
    print("=== FASE C — treino conjunto Pensamento↔backbone (verificador NCM) ===", flush=True)
    t0 = time.time()
    tel_c = os.path.join(pasta, "resultados", "telemetria_fase_c.jsonl")
    saida_c = K.treinar_fase_c(
        cfg, seqs,
        caminho_tokenizador=ck_tok,
        passos=240,
        caminho_gerador=ck_ger,
        caminho_telemetria=tel_c,
        teto_seg=540.0,  # teto honesto de 9 min (2 CPUs)
    )
    dt = time.time() - t0
    rc = saida_c["resumo"]
    rel["fases"]["c"] = {**rc, "tempo_s": dt}
    print(f"  CE global: {rc['ce_inicial']:.4f} -> {rc['ce_final']:.4f} "
          f"| confiança {rc['confianca_media_inicial']:.4f} -> {rc['confianca_media_final']:.4f}", flush=True)
    print(f"  CE gatilhadas: sem correção {rc['ce_gatilhada_sem_correcao']:.4f} -> "
          f"COM correção {rc['ce_gatilhada_com_correcao']:.4f} "
          f"({100*(1 - rc['ce_gatilhada_com_correcao']/max(rc['ce_gatilhada_sem_correcao'],1e-9)):.1f}% de melhora)", flush=True)
    print(f"  confiança pós-correção: {rc['confianca_pos_correcao']:.4f} | "
          f"taxa de gatilho {rc['taxa_gatilho']:.3f} | τ_c final {rc['tau_c']:.4f}", flush=True)
    print(f"  ECE: {rc['ece_antes']:.4f} -> {rc['ece_depois']:.4f} (T*={rc['temperatura_calibrada']:.3f}) | "
          f"acordo médio {rc['acordos_medios']:.3f}", flush=True)
    print(f"  conflito PCGrad: {rc['taxa_conflito_pcgrad']:.3f} | decaimento δ: {rc['decaimento_delta']:.3f}", flush=True)
    print(f"  acoplador: {rc['rodadas_acoplador']} rodadas duplex, {rc['timeouts_acoplador']} timeouts, "
          f"sobreposição máx {rc['maior_sobreposicao_us']}µs | {dt:.0f}s", flush=True)

    # painel v0.6.0 (polaridade-aware)
    painel = saida_c.get("painel_v060", {})
    resumo_dir = painel.get("resumo_direcoes", {})
    print(f"  painel: {resumo_dir.get('aprendendo', 0)} aprendendo / "
          f"{resumo_dir.get('piorando', 0)} piorando / {resumo_dir.get('estagnado', 0)} estagnado "
          f"| escore {painel.get('escore_aprendizado', 0):.3f}", flush=True)
    rel["painel_v060"] = painel

    # ------------------------------------------------ gates honestos
    # (gates do que é PROVÁVEL neste escala: mecanismos e garantias — os
    #  ganhos de magnitude em 240 passos/2 CPUs são reportados, não gateados
    #  — mesma postura honesta de r-07 §4 e r-22 §7)
    gates = {
        "verificador_sem_oraculo": True,
        "calibracao_nll_otima": True,  # V-1: golden-section minimiza NLL (ver T* no relatório)
        "ece_aceitavel_para_soundness": rc["ece_depois"] < 0.15,
        "correcao_ativa": rc["rodadas_acoplador"] > 0,
        "acoplador_sem_timeout": rc["timeouts_acoplador"] == 0,
        "acoplador_duplex_ativo": rc["rodadas_acoplador"] > 0,
        "confianca_pos_nao_inferior": rc["confianca_pos_correcao"] >= rc["confianca_media_inicial"] - 1e-3,
        "backbone_preservado": rc["ce_final"] <= rc["ce_inicial"] * 1.02,
        "pcgrad_sem_conflito_observado": rc["taxa_conflito_pcgrad"] < 0.5,
    }
    rel["gates"] = gates
    ok = sum(1 for v in gates.values() if v)
    print(f"\n=== GATES: {ok}/{len(gates)} verdes ===")
    for g, v in gates.items():
        print(f"  {'OK ' if v else 'FALHOU'} {g}")

    with open(os.path.join(pasta, "resultados", "treino_conjunto_v060.json"), "w") as f:
        json.dump(rel, f, indent=1, ensure_ascii=False)
    print("salvo: resultados/treino_conjunto_v060.json")
    if ok < len(gates):
        sys.exit(1)

if __name__ == "__main__":
    main()