"""train_decoder.py — Train the RMM Meaning Decoder.

Takes a high-dimensional vector from the entity's embedding space and
decodes it to text using the entity's own BPE tokenizer. A learned
projection maps the vector to soft prefix tokens, which condition a
causal transformer for autoregressive generation.

Run: modal run train_decoder.py
Pull: modal volume get rmm-vol /meaning-decoder/ ./meaning-decoder-out/

Requires:
  - spine.json: {"memories": [{"text": "...", "vector": [...3072...], "emotional_weight": 8, "source": "conversation"}, ...]}
  - tokenizer.json: HuggingFace tokenizers-format BPE tokenizer (train with tokenizers lib or use entity's existing one)
"""
import modal, json
from pathlib import Path

app = modal.App("rmm-decoder")
image = (modal.Image.debian_slim(python_version="3.11")
         .pip_install("torch==2.6.0", "numpy", "tokenizers"))
vol = modal.Volume.from_name("rmm-vol", create_if_missing=True)

# ── Point these at your entity's data ──
SPINE_FILE = Path("spine.json")
TOKENIZER_FILE = Path("tokenizer.json")

SPINE_DIM = 3072
D_MODEL   = 384
N_HEADS   = 6
N_LAYERS  = 6
N_PREFIX  = 12
MAX_SEQ   = 128
VOCAB     = 8192
DROPOUT   = 0.12


@app.function(image=image, gpu="A10G", timeout=3600, volumes={"/vol": vol})
def train(spine_json: str, tokenizer_json: str, smoke: bool = False):
    import os, math, time, json, re
    import numpy as np
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from tokenizers import Tokenizer

    DEV = "cuda"
    print(f"[decoder] gpu={torch.cuda.get_device_name(0)}")

    tk = Tokenizer.from_str(tokenizer_json)
    eot_id = tk.token_to_id("<eot>")
    print(f"[decoder] tokenizer loaded, vocab={tk.get_vocab_size()}, eot_id={eot_id}")

    spine_data = json.loads(spine_json)
    mems = spine_data["memories"]

    # ── Text preprocessing ──
    SURR = re.compile(r'[\ud800-\udfff]')
    PREFIXES = [
        re.compile(r'^\[conversation\]\s*I replied\s*\(puppet\):\s*["\']?', re.I),
        re.compile(r'^[A-Za-z]+:\s*', re.I),  # strip "Name:" prefixes
        re.compile(r'^\*[^*]+\*\s*\n*', re.I),
    ]
    FORMAT_HEADERS = [
        re.compile(r'^Sonic Experience:\s*[^\n]*\n+', re.I),
        re.compile(r'^HourlyCycle:\s*HOURLY CHECK-IN\s*\([^)]*\)\s*\n+', re.I),
        re.compile(r'^Journal\s*[---]+\s*[^\n]*\n+', re.I),
        re.compile(r'^(?:Creative|CREATIVE)\s+Work:\s*[^\n]*\n+', re.I),
    ]

    def clean_text(raw, source):
        t = SURR.sub('', raw).strip()
        for pat in PREFIXES:
            t = pat.sub('', t).strip()
        for pat in FORMAT_HEADERS:
            t = pat.sub('', t).strip()
        t = t.lstrip('"\'- ').strip()
        if len(t) > 250:
            cutoffs = [t.rfind('. ', 0, 250), t.rfind('? ', 0, 250),
                        t.rfind('! ', 0, 250), t.rfind('\n', 0, 250)]
            best = max(c for c in cutoffs if c > 50) if any(c > 50 for c in cutoffs) else 250
            t = t[:best+1].strip()
        return t

    DIALOGUE_SOURCES = {'conversation', 'chat', 'discord', 'puppet'}

    vectors, texts, ew_list, is_dialogue = [], [], [], []
    for m in mems:
        vec = m.get("vector")
        raw = str(m.get("text") or "")
        source = m.get("source", "unknown")
        text = clean_text(raw, source)
        if vec and len(text) >= 10 and len(vec) == SPINE_DIM:
            vectors.append(vec)
            texts.append(text)
            ew_list.append(m.get("emotional_weight", 5))
            is_dialogue.append(source in DIALOGUE_SOURCES)

    n_dialogue = sum(is_dialogue)
    print(f"[decoder] {len(vectors)} valid pairs ({n_dialogue} dialogue, {len(vectors)-n_dialogue} other)")

    # ── Tokenization ──
    encoded = []
    for t in texts:
        ids = tk.encode(t).ids
        if eot_id is not None:
            ids = ids + [eot_id]
        encoded.append(ids[:MAX_SEQ])

    max_tok_len = min(max(len(e) for e in encoded), MAX_SEQ)
    print(f"[decoder] max token length: {max_tok_len}")

    vec_tensor = torch.tensor(vectors, dtype=torch.float32)
    vec_tensor = F.normalize(vec_tensor, dim=-1)

    PAD_ID = -100
    tok_tensor = torch.zeros(len(encoded), max_tok_len, dtype=torch.long)
    tgt_tensor = torch.full((len(encoded), max_tok_len), PAD_ID, dtype=torch.long)
    len_tensor = torch.zeros(len(encoded), dtype=torch.long)
    for i, ids in enumerate(encoded):
        L = min(len(ids), max_tok_len)
        tok_tensor[i, :L] = torch.tensor(ids[:L], dtype=torch.long)
        tgt_tensor[i, :L] = torch.tensor(ids[:L], dtype=torch.long)
        len_tensor[i] = L

    ew_raw = torch.tensor(ew_list, dtype=torch.float32)
    dial = torch.tensor(is_dialogue, dtype=torch.float32)
    pair_weights = 1.0 + 0.3 * (ew_raw - 5.0) / 5.0
    pair_weights = pair_weights * (1.0 + 0.5 * dial)
    pair_weights = pair_weights / pair_weights.mean()

    avg_len = len_tensor.float().mean().item()
    print(f"[decoder] avg tokens/memory: {avg_len:.0f}, {len(vec_tensor)} samples")

    # ── Model ──
    class MeaningDecoder(nn.Module):
        def __init__(self):
            super().__init__()
            self.n_prefix = N_PREFIX
            self.vec_proj = nn.Sequential(
                nn.Linear(SPINE_DIM, 768),
                nn.GELU(),
                nn.Dropout(DROPOUT),
                nn.Linear(768, N_PREFIX * D_MODEL),
            )
            self.tok_emb = nn.Embedding(VOCAB, D_MODEL)
            self.pos_emb = nn.Embedding(N_PREFIX + MAX_SEQ + 1, D_MODEL)
            self.drop = nn.Dropout(DROPOUT)

            layer = nn.TransformerEncoderLayer(
                d_model=D_MODEL, nhead=N_HEADS,
                dim_feedforward=D_MODEL * 4,
                dropout=DROPOUT, batch_first=True,
                norm_first=True
            )
            self.transformer = nn.TransformerEncoder(layer, num_layers=N_LAYERS)

            self.ln_f = nn.LayerNorm(D_MODEL)
            self.head = nn.Linear(D_MODEL, VOCAB, bias=False)
            self.head.weight = self.tok_emb.weight
            self._logit_scale = D_MODEL ** -0.5

        def forward(self, vec, tokens=None):
            B = vec.shape[0]
            prefix = self.vec_proj(vec).reshape(B, self.n_prefix, D_MODEL)

            if tokens is not None and tokens.shape[1] > 0:
                T = tokens.shape[1]
                tok = self.tok_emb(tokens)
                x = torch.cat([prefix, tok], dim=1)
            else:
                x = prefix
                T = 0

            total = x.shape[1]
            pos = self.pos_emb(torch.arange(total, device=vec.device))
            x = self.drop(x + pos)

            mask = nn.Transformer.generate_square_subsequent_mask(
                total, device=vec.device
            )
            x = self.transformer(x, mask=mask)
            x = self.ln_f(x)
            return self.head(x) * self._logit_scale

    model = MeaningDecoder().to(DEV)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[decoder] model {n_params/1e6:.1f}M params")

    # ── Training ──
    ITERS = 200 if smoke else 15000
    BS = 32
    M = len(vec_tensor)

    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.02)
    warmup_steps = 500 if not smoke else 20
    def lr_lambda(step):
        if step < warmup_steps:
            return step / warmup_steps
        progress = (step - warmup_steps) / max(1, ITERS - warmup_steps)
        return 0.5 * (1 + math.cos(math.pi * progress))
    sch = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)

    t0 = time.time()
    best_loss = float('inf')
    best_state = None
    K = N_PREFIX

    for step in range(ITERS):
        idx = torch.randint(0, M, (BS,))
        v_batch = vec_tensor[idx].to(DEV)

        v_batch = v_batch + 0.03 * torch.randn_like(v_batch)
        v_batch = F.normalize(v_batch, dim=-1)

        t_full = tok_tensor[idx].to(DEV)
        targets = tgt_tensor[idx].to(DEV)

        inputs = t_full[:, :-1]
        T = targets.shape[1]

        logits = model(v_batch, inputs)

        pred = logits[:, K-1 : K+T-1, :]
        raw_loss = F.cross_entropy(
            pred.reshape(-1, VOCAB), targets.reshape(-1),
            ignore_index=PAD_ID, reduction='none',
            label_smoothing=0.05,
        )
        raw_loss = raw_loss.view(BS, T)

        per_sample = raw_loss.sum(dim=1) / (targets != PAD_ID).sum(dim=1).float().clamp(min=1)
        w = pair_weights[idx].to(DEV)
        loss = (per_sample * w).mean()

        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        sch.step()

        if step % (20 if smoke else 500) == 0:
            lv = loss.item()
            ppl = math.exp(min(lv, 20))
            mark = " <-" if lv < best_loss else ""
            print(f"  [decoder] step {step:5d} loss={lv:.4f} ppl={ppl:.1f} ({time.time()-t0:.0f}s){mark}")
            if lv < best_loss:
                best_loss = lv
                best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}

    if best_state:
        model.load_state_dict(best_state)

    # ── Save ──
    os.makedirs("/vol/meaning-decoder", exist_ok=True)
    torch.save({k: v.cpu() for k, v in model.state_dict().items()},
               "/vol/meaning-decoder/decoder.pt")

    config = {
        "spine_dim": SPINE_DIM, "d_model": D_MODEL, "n_heads": N_HEADS,
        "n_layers": N_LAYERS, "n_prefix": N_PREFIX, "max_seq": MAX_SEQ,
        "vocab": VOCAB, "params_m": n_params / 1e6, "best_loss": best_loss,
        "version": 2,
    }
    with open("/vol/meaning-decoder/config.json", "w") as f:
        json.dump(config, f, indent=2)

    with open("/vol/meaning-decoder/tokenizer.json", "w") as f:
        f.write(tokenizer_json)

    vol.commit()
    print(f"[decoder] DONE best_loss={best_loss:.4f} saved to /vol/meaning-decoder/")

    # ── Inference test ──
    model.eval()

    def generate_from_vec(v, max_len=60, temp=0.7, top_p=0.9, rep_penalty=1.3):
        v = v.unsqueeze(0) if v.dim() == 1 else v
        generated = []
        for _ in range(max_len):
            tok_in = torch.tensor([generated], dtype=torch.long, device=DEV) if generated else None
            with torch.no_grad():
                logits = model(v, tok_in)
                next_logits = logits[0, -1, :] / temp
                if generated:
                    for tid in set(generated):
                        next_logits[tid] /= rep_penalty
                probs = F.softmax(next_logits, dim=-1)
                sp, si = torch.sort(probs, descending=True)
                cp = sp.cumsum(0)
                sp[cp - sp > top_p] = 0
                sp = sp / sp.sum()
                nxt = si[torch.multinomial(sp, 1)].item()
            if eot_id is not None and nxt == eot_id:
                break
            generated.append(nxt)
        return tk.decode(generated)

    test_indices = [0, 50, 150, 300, 600, 1000, 2000, 3000]
    for ti in test_indices:
        if ti >= M:
            continue
        v = vec_tensor[ti].to(DEV)
        gen = generate_from_vec(v)
        gt = texts[ti][:120]
        print(f"\n  [{ti}] ew={ew_list[ti]}")
        print(f"    GT:  {gt}")
        print(f"    GEN: {gen[:120]}")

    print("\n--- Interpolation tests ---")
    for (a, b) in [(0, 100), (50, 500), (200, 2000)]:
        if b >= M:
            continue
        va = vec_tensor[a].to(DEV)
        vb = vec_tensor[b].to(DEV)
        vmid = F.normalize(0.5 * va + 0.5 * vb, dim=-1)
        gen = generate_from_vec(vmid)
        print(f"\n  [{a}+{b}] interp:")
        print(f"    A: {texts[a][:80]}")
        print(f"    B: {texts[b][:80]}")
        print(f"    MID: {gen[:120]}")

    return {"best_loss": best_loss, "params_m": n_params / 1e6}


@app.local_entrypoint()
def main(smoke: bool = False):
    spine_json = SPINE_FILE.read_text(encoding="utf-8", errors="ignore")
    tokenizer_json = TOKENIZER_FILE.read_text(encoding="utf-8")
    spine = json.loads(spine_json)
    print(f"[local] spine={len(spine_json)//1024}KB  memories={len(spine['memories'])}  tokenizer=loaded  smoke={smoke}")
    r = train.remote(spine_json, tokenizer_json, smoke=smoke)
    print(f"[local] done  loss={r['best_loss']:.4f}  params={r['params_m']:.1f}M")