#!/usr/bin/env python3
"""51_treino_v050.py — Treino com dados reais (Wikipédia PT-BR) e telemetria v0.5.0.

Recoleta os shards (foram apagados na limpeza de estados — regeneráveis por
design), treina as fases A (tokenizador GSQ+RCO) e B (gerador NCM) com tetos
de tempo honestos, e adiciona o painel de telemetria v0.5.0 (detecção de
aprendizado Theil–Sen sobre as séries do JSONL).

Saída: resultados/treino_v050.json + telemetria_fase_a/b.jsonl
"""
import json, math, os, sys, time

def coletar(pasta, max_artigos=400, max_bytes=4_000_000):
    """Streaming controlado (tetos duros)."""
    from datasets import load_dataset
    ds = load_dataset("wikimedia/wikipedia", "20231101.pt", split="train", streaming=True)
    textos, total_bytes = [], 0
    for ex in ds:
        t = (ex.get("text") or "").strip()
        if len(t) < 200:
            continue
        textos.append(t)
        total_bytes += len(t.encode("utf-8"))
        if len(textos) % 100 == 0:
            print(f"  ... {len(textos)} artigos, {total_bytes/1e6:.2f} MB", flush=True)
        if len(textos) >= max_artigos or total_bytes >= max_bytes:
            break
    def limpa(t):
        partes = [p.strip() for p in t.split("\n") if len(p.strip()) > 80]
        return " ".join(partes[:6])[:1800]
    textos = [limpa(t) for t in textos]
    n_val = max(1, len(textos) // 10)
    with open(os.path.join(pasta, "dados", "shard_validacao.jsonl"), "w") as f:
        for t in textos[:n_val]:
            f.write(json.dumps({"texto": t}, ensure_ascii=False) + "\n")
    with open(os.path.join(pasta, "dados", "shard_treino.jsonl"), "w") as f:
        for t in textos[n_val:]:
            f.write(json.dumps({"texto": t}, ensure_ascii=False) + "\n")
    print(f"coletado: {len(textos)} artigos, {total_bytes/1e6:.2f} MB")

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

def painel_do_jsonl(caminho):
    """Painel v0.5.0: séries do JSONL → detecção Theil–Sen por métrica."""
    series = {}
    if not os.path.exists(caminho):
        return series
    with open(caminho) as f:
        for linha in f:
            try:
                ev = json.loads(linha)
            except json.JSONDecodeError:
                continue
            dados = ev.get("dados", {})
            if ev.get("tipo") == "dump":
                # dump final: {"nome_serie": [[passo, valor], ...]}
                for k, v in dados.items():
                    if isinstance(v, list):
                        ys = [par[1] for par in v if isinstance(par, list) and len(par) == 2]
                        if len(ys) >= 2:
                            series[k] = ys
            else:
                for k, v in dados.items():
                    if isinstance(v, (int, float)) and not isinstance(v, bool):
                        series.setdefault(k, []).append(float(v))
    painel = {}
    for k, ys in series.items():
        if len(ys) >= 8:
            med, vered = theil_sen(ys)
            painel[k] = {"inclinacao": med, "veredito": vered, "n": len(ys),
                         "primeiro": ys[0], "ultimo": ys[-1]}
    return painel

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"))

    # 1) dados (se ausentes — regeneráveis)
    if not os.path.exists(os.path.join(pasta, "dados", "shard_treino.jsonl")):
        print("=== Coleta de dados reais (Wikipédia PT-BR, streaming com tetos) ===", flush=True)
        coletar(pasta)
    else:
        print("shards presentes — reutilizando")

    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])
    print(f"sequências de treino: {len(seqs)} (ctx=256)", flush=True)
    # limita para teto de tempo honesto (~6 min de fase A + ~5 min de B)
    seqs = seqs[:1400]

    rel = {"versao": "0.5.0", "dados": {"artigos": "ver shard", "sequencias": len(seqs)}, "fases": {}}

    # 2) fase A
    print("=== FASE A — tokenizador de conceitos (GSQ+RCO Stiefel) ===", flush=True)
    t0 = time.time()
    tel_a = os.path.join(pasta, "resultados", "telemetria_fase_a.jsonl")
    saida = K.treinar_fase_a(
        cfg, seqs, caminho_telemetria=tel_a,
        caminho_checkpoint=os.path.join(pasta, "resultados", "tokenizador_conceitos.json"),
    )
    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} | acc_recon {r['acc_recon']:.3f} "
          f"| δ_ort {r['delta_ort_max']:.1e} | perp_codebook {r['perplexidade_codebook']:.1f} "
          f"| escore {r['escore_aprendizado']:.3f} | {dt:.0f}s", flush=True)

    # 3) fase B
    print("=== FASE B — gerador NCM (MLE exato, hierarquias congeladas) ===", flush=True)
    t0 = time.time()
    tel_b = os.path.join(pasta, "resultados", "telemetria_fase_b.jsonl")
    saida_b = K.treinar_fase_b(
        cfg, seqs,
        caminho_tokenizador=os.path.join(pasta, "resultados", "tokenizador_conceitos.json"),
        caminho_telemetria=tel_b,
        # checkpoint TEMPORÁRIO para a suíte V10 (apagado na limpeza final;
        # regenerável rodando este script)
        caminho_checkpoint=os.path.join(pasta, "resultados", "gerador_conceitos.json"),
        teto_seg=300.0,           # teto honesto de 5 min
    )
    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} | acc_conceitos {rb.get('acuracia_media', 0):.3f} "
          f"| {dt:.0f}s", flush=True)

    # 4) painel de telemetria v0.5.0 (Theil–Sen por série)
    print("=== Painel de telemetria v0.5.0 (detecção Theil–Sen) ===", flush=True)
    painel_a = painel_do_jsonl(tel_a)
    painel_b = painel_do_jsonl(tel_b)
    for nome, p in (("fase A", painel_a), ("fase B", painel_b)):
        for k, v in sorted(p.items()):
            print(f"  [{nome}] {k}: {v['primeiro']:.4f} -> {v['ultimo']:.4f} ({v['veredito']}, n={v['n']})")
    rel["telemetria_v050"] = {"fase_a": painel_a, "fase_b": painel_b}

    # gates de aprendizado honestos
    gates = {
        "perda_a_cai": r["perda_final"] < r["perda_inicial"],
        "ortogonalidade_a": r["delta_ort_max"] < 1e-5,
        "codebook_sem_colapso": r["perplexidade_codebook"] > 8.0,
        "ce_b_cai": rb["perda_final"] < rb["perda_inicial"],
        "deteccao_a_aprendendo": painel_a.get("perda", {}).get("veredito") == "aprendendo",
    }
    if painel_b:
        gates["deteccao_b_aprendendo"] = painel_b.get("perda", {}).get("veredito") == "aprendendo"
    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_v050.json"), "w") as f:
        json.dump(rel, f, indent=1, ensure_ascii=False)
    print("salvo: resultados/treino_v050.json")
    if ok < len(gates):
        sys.exit(1)

if __name__ == "__main__":
    main()