"""(3) Overfit/memorization test: ATOM recipe on a ~2 KB text, repeated passes.
usage: overfit_atom.py NAME {fresh|ckpt} MINUTES [LR]
Mirrors tools/run_atom_native.py per-step logic: AdamW(wd=.01), clip 1.0, stabilize_dynamics,
efference every 10, atom flush every 64, last_atom_dropout .5, field loss .05, reset per pass.
Eval: teacher-forced CE on the text + greedy regeneration from the first 32 bytes
(generation path with train-like commit labels, i.e. the decode bug neutralised, and original)."""
import json, sys, time
from pathlib import Path
import torch, torch.nn.functional as F
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import CKPT, ROOT, load_model
from src.atom_native import AtomNativeModel
from src.io.atomizer import Atomizer
OUT = Path(__file__).resolve().parent
TEXT = OUT / "memorize_2k.txt"


def make_text():
    raw = (ROOT / "data/corpus_fr_medium.txt").read_bytes()[:2300]
    cut = raw[:2048].rfind(b"\n\n") + 2
    TEXT.write_bytes(raw[:cut])


def fresh():
    m = AtomNativeModel(d_model=32, n_modes=32, n_atoms_max=64, max_payload_bytes=1,
                        atomizer=Atomizer(max_span_bytes=1), field_max_rms=3.0, enable_merge=False,
                        field_loss_weight=0.05, field_obligatory_hard=True, last_atom_readout=True, field_leak=0.3)
    return m


def flush(m):  # same as trainer --atom-flush-every 64 (AtomSet is unbounded; aggregation is O(n^2)/tick)
    m.core.atoms = type(m.core.atoms)(); m.core.energy_history = []


@torch.no_grad()
def tf_ce(m, data):
    m.eval(); m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(data, reset=True)
    ls = []
    for i in range(len(pk) - 1):
        ls.append(float(m.transition_loss(pk[i], pk[i + 1])[1]["byte_loss"]))
        if (i + 1) % 64 == 0: flush(m)
    m.train(); return sum(ls) / len(ls)


@torch.no_grad()
def regen(m, data, train_labels=True, n=None):
    m.eval(); orig = Atomizer._infer_commit_boundary
    if train_labels:
        Atomizer._infer_commit_boundary = classmethod(lambda cls, p: "newline" if p[-1] in (10, 13) else "max_span")
    try:
        m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(data[:32], reset=True)
        for q in pk[:-1]: m.forward_packet(q)
        cur = pk[-1]; out = bytearray(data[:32]); n = n or len(data)
        while len(out) < n:
            if len(out) % 64 == 0: flush(m)
            b = int(m.forward_packet(cur)["surface"]["byte_logits"][0].argmax()); out.append(b)
            cur = m.atomizer.packet_from_payload(bytes([b]))
    finally:
        Atomizer._infer_commit_boundary = orig
    m.train(); same = next((i for i, (u, v) in enumerate(zip(out, data)) if u != v), len(data))
    return same, bytes(out[:260]).decode("utf-8", "replace")


def main():
    name, init, minutes = sys.argv[1], sys.argv[2], float(sys.argv[3])
    lr = float(sys.argv[4]) if len(sys.argv) > 4 else 3e-4
    torch.set_num_threads(1); torch.manual_seed(20260913)
    if not TEXT.exists(): make_text()
    data = TEXT.read_bytes()
    m = fresh() if init == "fresh" else load_model(CKPT)
    m.surface.payload_enabled = False; m.surface.last_atom_dropout = 0.5; m.field_loss_weight = 0.05; m._efference_every = 10
    m.train()
    opt = torch.optim.AdamW([p for p in m.parameters() if p.requires_grad], lr=lr, weight_decay=0.01)
    log = []; t0 = time.time(); step = 0; ep = 0
    r0 = {"pass": 0, "tf_ce": tf_ce(m, data)}; log.append(r0); print(json.dumps(r0), flush=True)
    while time.time() - t0 < minutes * 60:
        m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(data, reset=True); ep += 1
        for i in range(len(pk) - 1):
            step += 1; opt.zero_grad(set_to_none=True)
            if step % 10 == 0: loss, info = m.efference_loss(pk[i], pk[i + 1])
            else: loss, info = m.transition_loss(pk[i], pk[i + 1])
            loss.backward(); torch.nn.utils.clip_grad_norm_([p for p in m.parameters() if p.requires_grad], 1.0); opt.step()
            if step % 64 == 0:
                m.core.atoms = type(m.core.atoms)(); m.core.energy_history = []
            m.stabilize_dynamics_parameters()
        if ep % 5 == 0:
            r = {"pass": ep, "step": step, "min": round((time.time() - t0) / 60, 1), "tf_ce": tf_ce(m, data)}
            log.append(r); print(json.dumps(r), flush=True)
    s1, h1 = regen(m, data, True); s2, h2 = regen(m, data, False)
    res = {"name": name, "init": init, "lr": lr, "text_bytes": len(data), "passes": ep, "steps": step,
           "final_tf_ce": tf_ce(m, data), "regen_prefix_match_train_labels": s1, "regen_head_train_labels": h1,
           "regen_prefix_match_orig_labels": s2, "regen_head_orig_labels": h2, "log": log}
    (OUT / f"overfit_atom_{name}.json").write_text(json.dumps(res, indent=1, ensure_ascii=False))
    print(json.dumps({k: v for k, v in res.items() if k != "log"}, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()