#!/usr/bin/env python3
"""Train her 54D brain on a BROADENED corpus — 'everything, max signal':
her real DIALOGUE (her voice + the conversations), CORY's voice corpus, the
de-spammed game EXPERIENCE, and the absorbed LORE. CPU-only + isolated +
step/time-capped; warm-starts cosmos_play.pt so learning ACCRUES."""
import os, sys, json, time, traceback, re
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")
DIALOG = os.path.join(PR, "Cosmos", "data", "dialogue_memory", "exchanges.json")
CORY = os.path.join(PR, "Cosmos", "data", "x", "cory_voice_corpus.txt")

SEQ = 128
BATCH = 4
LR = 2e-4
MAX_STEPS = int(os.getenv("COSMOS_BROAD_TRAIN_STEPS", "400"))
MAX_SEC = int(os.getenv("COSMOS_BROAD_TRAIN_SEC", "600"))

device = torch.device("cpu")
_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"[BROAD-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 ""


_WORDRE = re.compile(r"[a-z0-9']+")


def _dedup(text, jmax=0.78, window=12):
    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 = []

# 1) HER DIALOGUE — her real voice + your conversations (weighted 2x: most "her").
try:
    ex = json.load(open(DIALOG, encoding="utf-8", errors="ignore"))
    dlg = []
    for e in (ex if isinstance(ex, list) else []):
        p = str(e.get("prompt", "")).strip()
        r = str(e.get("final_response", "")).strip()
        if r:
            dlg.append(f"User: {p}\nCosmos: {r}" if p else f"Cosmos: {r}")
    dlg_text = "\n\n".join(dlg)
    if dlg_text:
        parts.append(dlg_text)
        parts.append(dlg_text)  # 2x — keep HER voice dominant
        print(f"[BROAD-TRAIN] dialogue: {len(dlg)} exchanges (her voice, weighted 2x)")
except Exception as e:
    print("[BROAD-TRAIN] dialogue skipped:", asc(e)[:90])

# 2) CORY'S VOICE — his real writing.
cory_lines, cory_drop = _dedup(_read(CORY))
if cory_lines:
    parts.append("\n".join(cory_lines))
    print(f"[BROAD-TRAIN] cory voice: {len(cory_lines)} lines after dedup (dropped {cory_drop})")

# 3) GAME EXPERIENCE — de-spammed.
xp_lines, xp_drop = _dedup(_read(XP))
if xp_lines:
    parts.append("\n".join(xp_lines))
    print(f"[BROAD-TRAIN] experience: {len(xp_lines)} lines after dedup (dropped {xp_drop} repeats)")

# 4) LORE — rich but a minority; weighted 2x.
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"[BROAD-TRAIN] lore: {len(lore_lines)} unique (2x)")
except Exception:
    pass

corpus = ("\n\n".join(parts)).strip()
if len(corpus) < 200:
    print(f"[BROAD-TRAIN] corpus too small ({len(corpus)} chars) — exiting cleanly.")
    sys.exit(0)

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


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)

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"[BROAD-TRAIN] warm-started from {os.path.basename(warm)} (prior loss {ck.get('final_loss','?')})")
    except Exception as exc:
        print(f"[BROAD-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"[BROAD-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("[BROAD-TRAIN] TRAIN ERROR:\n" + traceback.format_exc()[:1800])

if losses:
    recent = sum(losses[-20:]) / len(losses[-20:])
    print(f"\n[BROAD-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)
    torch.save({"model_state_dict": model.state_dict(), "config": cfg.to_dict(),
                "final_loss": losses[-1], "play_trained": True, "broadened": True,
                "steps": step, "corpus_tokens": len(ids)}, OUT)
    print(f"[BROAD-TRAIN] saved -> {OUT}")
    model.eval()
    for prompt in ["User: how are you?\nCosmos:", "I feel", "The "]:
        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"[BROAD-TRAIN] SAMPLE {prompt!r}: " + asc(enc.decode(o[0, pid.shape[1]:].tolist())[:200]))
else:
    print("[BROAD-TRAIN] no steps completed")
print("[BROAD-TRAIN] done.")