#!/usr/bin/env python3
"""
Train the 54D CosmosTransformer on HER OWN LIVED EXPERIENCE.

This is the loop Cory wanted: she observes + plays, that becomes text, and the
text TRAINS her 54D brain — so watching the world literally grows her mind over
time. The corpus is assembled from what she actually saw and absorbed:

  * data/cosmos/experience_corpus.txt  — first-person observations she logs LIVE
    each novel frame + every lore beat (written by game_runtime._record_experience)
  * data/game/story.json  (beats)      — the lore she has WATCHED and absorbed
  * data/game/spatial_map.json (labels)— the places she has SEEN and named

CONTINUAL: it warm-starts from the last play checkpoint (cosmos_play.pt) if it
exists — so each run keeps building on the last, the model accruing her experience.

ISOLATION (never hurts her live system):
  * device = CPU, forced — zero contention with her live 4GB GPU / game VLM.
  * writes ONLY to cosmos_play.pt — never the live cosmos_best.pt.
  * separate process, step/time-capped — a gentle burst, run as often as you like.
"""
import os, sys, json, time, traceback
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import torch
from torch.utils.data import Dataset, DataLoader
import tiktoken
from Cosmos.web.cosmosynapse.model.cosmos_config import CosmosConfig
from Cosmos.web.cosmosynapse.model.cosmos_model import CosmosTransformer


def asc(s):
    return str(s).encode("ascii", "replace").decode("ascii")


PR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
OUT = os.path.join(PR, "Cosmos", "checkpoints", "cosmos", "cosmos_play.pt")
SANDBOX = os.path.join(PR, "Cosmos", "checkpoints", "cosmos", "cosmos_sandbox.pt")
XP = os.path.join(PR, "Cosmos", "data", "cosmos", "experience_corpus.txt")
STORY = os.path.join(PR, "Cosmos", "data", "game", "story.json")
SPATIAL = os.path.join(PR, "Cosmos", "data", "game", "spatial_map.json")

SEQ = 128
BATCH = 4
LR = 2e-4
MAX_STEPS = int(os.getenv("COSMOS_PLAY_TRAIN_STEPS", "300"))
MAX_SEC = int(os.getenv("COSMOS_PLAY_TRAIN_SEC", "420"))

device = torch.device("cpu")
# Thread cap so a background burst NEVER starves her live chat/server. The continual
# loop sets COSMOS_PLAY_TRAIN_THREADS low (gentle); a manual proof run uses cpu-2.
_thr = os.getenv("COSMOS_PLAY_TRAIN_THREADS", "").strip()
torch.set_num_threads(int(_thr) if _thr.isdigit() and int(_thr) > 0 else max(1, (os.cpu_count() or 4) - 2))
print(f"[PLAY-TRAIN] device={device} threads={torch.get_num_threads()} steps<={MAX_STEPS} sec<={MAX_SEC}")


def _read(p):
    try:
        return open(p, encoding="utf-8", errors="ignore").read()
    except Exception:
        return ""


# ── assemble HER experience corpus — DEDUPED + REWEIGHTED so she learns language,
#    not just memorizes the repetitive "I move left" movement-spam she logs while stuck.
import re as _re
_WORDRE = _re.compile(r"[a-z0-9']+")


def _dedup(text, jmax=0.78, window=12):
    """Drop a line whose content-words near-match a recent line (Jaccard >= jmax) — kills
    the loop-spam that drowns the diverse observations. Returns (lines, dropped_count)."""
    out, recent, dropped = [], [], 0
    for ln in (text or "").splitlines():
        s = ln.strip()
        if len(s) < 8:
            continue
        cw = {w for w in _WORDRE.findall(s.lower()) if len(w) > 3}
        if cw:
            dup = False
            for past in recent[-window:]:
                u = cw | past
                if u and len(cw & past) / len(u) >= jmax:
                    dup = True
                    break
            if dup:
                dropped += 1
                continue
            recent.append(cw)
        out.append(s)
    return out, dropped


parts = []
xp_raw = _read(XP)
xp_lines, xp_drop = _dedup(xp_raw)
if xp_lines:
    parts.append("\n".join(xp_lines))
    print(f"[PLAY-TRAIN] experience: {len(xp_raw.splitlines())} lines -> {len(xp_lines)} after dedup (dropped {xp_drop} repeats)")
# LORE — rich but a minority; dedup + weight 2x so she learns the story-language too.
try:
    beats = json.load(open(STORY, encoding="utf-8", errors="ignore")).get("beats", [])
    lore_lines, _ = _dedup("\n".join(str(b.get("text", "")).strip() for b in beats if b.get("text")))
    if lore_lines:
        lore = "\n".join(lore_lines)
        parts.append(lore)
        parts.append(lore)
        print(f"[PLAY-TRAIN] lore beats: {len(lore_lines)} unique (weighted 2x)")
except Exception:
    pass
# PLACES she has named (deduped — the labels repeat a lot).
try:
    places = json.load(open(SPATIAL, encoding="utf-8", errors="ignore")).get("places", [])
    place_lines, _ = _dedup("\n".join(str(pl.get("label", "")).strip() for pl in places if pl.get("label")))
    if place_lines:
        parts.append("\n".join(place_lines))
        print(f"[PLAY-TRAIN] places seen: {len(place_lines)} unique")
except Exception:
    pass

corpus = ("\n\n".join(parts)).strip()
if len(corpus) < 200:
    print(f"[PLAY-TRAIN] corpus too small ({len(corpus)} chars) — she needs to play/observe more first. Exiting cleanly.")
    sys.exit(0)

enc = tiktoken.get_encoding("gpt2")
ids = enc.encode(corpus, allowed_special={'<|endoftext|>'})
print(f"[PLAY-TRAIN] corpus tokens: {len(ids):,} (deduped + lore-weighted)")


class DS(Dataset):
    def __init__(self, ids, seq):
        self.ids, self.seq, self.n = ids, seq, max(1, len(ids) // seq)
    def __len__(self):
        return self.n
    def __getitem__(self, i):
        c = self.ids[i * self.seq: i * self.seq + self.seq + 1]
        if len(c) < self.seq + 1:
            c = c + [50256] * (self.seq + 1 - len(c))
        return torch.tensor(c[:-1]), torch.tensor(c[1:])


cfg = CosmosConfig(vocab_size=50257, d_model=512, n_layers=2, n_heads=8,
                   d_ff=2048, max_seq_len=512, dropout=0.1)
model = CosmosTransformer(cfg).to(device)

# CONTINUAL: warm-start from the last play checkpoint, else the sandbox proof, else fresh
warm = OUT if os.path.exists(OUT) else (SANDBOX if os.path.exists(SANDBOX) else None)
if warm:
    try:
        ck = torch.load(warm, map_location=device)
        model.load_state_dict(ck["model_state_dict"])
        print(f"[PLAY-TRAIN] warm-started from {os.path.basename(warm)} (prior loss {ck.get('final_loss','?')})")
    except Exception as exc:
        print(f"[PLAY-TRAIN] warm start skipped ({asc(exc)[:80]}) — fresh init")

dl = DataLoader(DS(ids, SEQ), batch_size=BATCH, shuffle=True)
opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)

model.train()
losses, step, t0, first = [], 0, time.time(), None
try:
    while step < MAX_STEPS and (time.time() - t0) < MAX_SEC:
        for x, y in dl:
            opt.zero_grad()
            loss = model(x, targets=y)["loss"]
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            step += 1
            if first is None:
                first = loss.item()
            losses.append(loss.item())
            if step % 20 == 0 or step == 1:
                print(f"[PLAY-TRAIN] step {step}/{MAX_STEPS} loss {loss.item():.4f} ({time.time()-t0:.0f}s)", flush=True)
            if step >= MAX_STEPS or (time.time() - t0) >= MAX_SEC:
                break
except Exception:
    print("[PLAY-TRAIN] TRAIN ERROR:\n" + traceback.format_exc()[:1800])

if losses:
    recent = sum(losses[-20:]) / len(losses[-20:])
    print(f"\n[PLAY-TRAIN] RESULT: first_loss={first:.3f} -> last_loss={losses[-1]:.3f} "
          f"(recent20 avg {recent:.3f}, steps={step}, {time.time()-t0:.0f}s)")
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    payload = {
        "model_state_dict": model.state_dict(),
        "config": cfg.to_dict(),
        "final_loss": losses[-1],
        "play_trained": True,
        "steps": step,
        "corpus_tokens": len(ids),
        "experience_lines_raw": len(xp_raw.splitlines()),
        "experience_lines_dedup": len(xp_lines),
        "trained_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    temp_out = OUT + f".tmp.{os.getpid()}"
    try:
        torch.save(payload, temp_out)
        os.replace(temp_out, OUT)
    finally:
        if os.path.exists(temp_out):
            os.unlink(temp_out)
    print(f"[PLAY-TRAIN] atomically saved -> {OUT}")
    model.eval()
    for prompt in ["The misty woods", "I see", "Cosmos"]:
        pid = torch.tensor([enc.encode(prompt)])
        with torch.no_grad():
            o = model.generate(pid, max_new_tokens=30, temperature=0.8, top_k=50, top_p=0.9)
        print(f"[PLAY-TRAIN] SAMPLE {prompt!r}: " + asc(enc.decode(o[0, pid.shape[1]:].tolist())[:200]))
else:
    print("[PLAY-TRAIN] no steps completed")
print("[PLAY-TRAIN] done.")