"""gridlm: two small models trained on a thegrid fossil record.

  screen     encoder: [CLS] parent [SEP] child -> P(child is born >= 2 times in its epoch)
  generator  causal LM: [BOS] parent [SEP] child [EOS], trained on successful children only

Held-out design: rows are split by EPOCH (train = earliest ~80% of rows' epochs, val = next
~10%, test = last ~10%), so the test set is later, unseen runs of the colony. Test metrics are
also reported on "novel" children whose genome never appears in the training split.

Inputs are the two instruction sequences only. Excluded on purpose: mutation_type (recorded
only for successful origins, so it leaks the label) and the parent's epoch birth count
(includes births after the child appeared).

usage: python train.py data.jsonl.gz outdir [--threads N] [--epochs-screen E] [--epochs-gen E]
"""
from __future__ import annotations

import argparse
import gzip
import json
import math
import os
import random
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

MAXG = 64
SPECIALS = ["[PAD]", "[CLS]", "[SEP]", "[BOS]", "[EOS]", "[UNK]"]
PAD, CLS, SEP, BOS, EOS, UNK = range(6)
DEV = torch.device("cpu")
KEY_OPS = ["alloc", "copy", "fork", "harvest", "scan", "move", "ifnotdone", "jmpb", "salvage", "steal"]


# ---------------------------------------------------------------- data
def load(path):
    rows = []
    with gzip.open(path, "rt", encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            if 1 <= len(r["child"]) <= MAXG and 1 <= len(r["parent"]) <= MAXG:
                rows.append(r)
    return rows


def split_by_epoch(rows):
    epochs = sorted({r["epoch"] for r in rows})
    per = {e: 0 for e in epochs}
    for r in rows:
        per[r["epoch"]] += 1
    total, acc, cut1, cut2 = len(rows), 0, None, None
    for e in epochs:
        acc += per[e]
        if cut1 is None and acc >= 0.8 * total:
            cut1 = e
        if cut2 is None and acc >= 0.9 * total:
            cut2 = e
    tr = [r for r in rows if r["epoch"] <= cut1]
    va = [r for r in rows if cut1 < r["epoch"] <= cut2]
    te = [r for r in rows if r["epoch"] > cut2]
    return tr, va, te, (cut1, cut2)


def build_vocab(rows):
    ops = sorted({t for r in rows for t in r["parent"] + r["child"]})
    return SPECIALS + ops


def enc_pair(r, idx):
    p = [idx.get(t, UNK) for t in r["parent"]]
    c = [idx.get(t, UNK) for t in r["child"]]
    return p, c


def screen_tensor(rows, idx):
    L = 2 * MAXG + 2
    X = np.full((len(rows), L), PAD, dtype=np.int64)
    y = np.zeros(len(rows), dtype=np.float32)
    for i, r in enumerate(rows):
        p, c = enc_pair(r, idx)
        seq = [CLS] + p + [SEP] + c
        X[i, :len(seq)] = seq
        y[i] = r["label"]
    return torch.from_numpy(X), torch.from_numpy(y)


def lm_tensor(rows, idx):
    L = 2 * MAXG + 3
    X = np.full((len(rows), L), PAD, dtype=np.int64)
    M = np.zeros((len(rows), L), dtype=np.float32)   # 1 where the NEXT token is a child token or EOS
    for i, r in enumerate(rows):
        p, c = enc_pair(r, idx)
        seq = [BOS] + p + [SEP] + c + [EOS]
        X[i, :len(seq)] = seq
        start = 1 + len(p)                 # position of SEP; predicts child[0]
        M[i, start:start + len(c) + 1] = 1  # child tokens and EOS are the targets
    return torch.from_numpy(X), torch.from_numpy(M)


# ---------------------------------------------------------------- models
class Encoder(nn.Module):
    def __init__(self, V, d=128, layers=4, heads=4, ff=512, L=2 * MAXG + 2, causal=False):
        super().__init__()
        self.tok = nn.Embedding(V, d, padding_idx=PAD)
        self.pos = nn.Embedding(L + 1, d)
        layer = nn.TransformerEncoderLayer(d, heads, ff, dropout=0.1, batch_first=True, norm_first=True)
        self.body = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False)
        self.norm = nn.LayerNorm(d)
        self.causal = causal

    def forward(self, x):
        T = x.shape[1]
        h = self.tok(x) + self.pos(torch.arange(T, device=x.device))[None]
        mask = nn.Transformer.generate_square_subsequent_mask(T, device=x.device) if self.causal else None
        h = self.body(h, mask=mask, src_key_padding_mask=(x == PAD), is_causal=self.causal)
        return self.norm(h)


class Screen(nn.Module):
    def __init__(self, V, **kw):
        super().__init__()
        self.enc = Encoder(V, **kw)
        self.head = nn.Linear(self.enc.tok.embedding_dim, 1)

    def forward(self, x):
        return self.head(self.enc(x)[:, 0]).squeeze(-1)


class Generator(nn.Module):
    def __init__(self, V, **kw):
        super().__init__()
        self.enc = Encoder(V, L=2 * MAXG + 3, causal=True, **kw)
        self.head = nn.Linear(self.enc.tok.embedding_dim, V)

    def forward(self, x):
        return self.head(self.enc(x))


# ---------------------------------------------------------------- metrics
def auroc(scores, labels):
    s = np.asarray(scores, dtype=np.float64)
    y = np.asarray(labels, dtype=np.int64)
    pos, neg = y.sum(), len(y) - y.sum()
    if pos == 0 or neg == 0:
        return float("nan")
    order = np.argsort(s, kind="mergesort")
    ranks = np.empty(len(s))
    sorted_s = s[order]
    i = 0
    while i < len(s):                       # average ranks for ties
        j = i
        while j + 1 < len(s) and sorted_s[j + 1] == sorted_s[i]:
            j += 1
        ranks[order[i:j + 1]] = (i + j) / 2 + 1
        i = j + 1
    return float((ranks[y == 1].sum() - pos * (pos + 1) / 2) / (pos * neg))


def avg_precision(scores, labels):
    order = np.argsort(-np.asarray(scores))
    y = np.asarray(labels)[order]
    tp = np.cumsum(y)
    prec = tp / np.arange(1, len(y) + 1)
    return float((prec * y).sum() / max(1, y.sum()))


def precision_at(scores, labels, frac):
    k = max(1, int(len(scores) * frac))
    order = np.argsort(-np.asarray(scores))[:k]
    return float(np.asarray(labels)[order].mean())


def report(name, scores, labels):
    return {"model": name, "n": len(labels), "base_rate": float(np.mean(labels)),
            "auroc": auroc(scores, labels), "avg_precision": avg_precision(scores, labels),
            "precision_at_top10pct": precision_at(scores, labels, 0.10)}


# ---------------------------------------------------------------- baselines
def edit_distance(a, b):
    prev = list(range(len(b) + 1))
    for i, x in enumerate(a, 1):
        cur = [i]
        for j, yv in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (x != yv)))
        prev = cur
    return prev[-1]


def features(r):
    c, p = r["child"], r["parent"]
    f = [len(c), len(p), len(c) - len(p), edit_distance(p, c)]
    f += [c.count(op) for op in KEY_OPS]
    f.append(float(all(op in c for op in ("alloc", "copy", "fork"))))
    return f


def logistic(Xtr, ytr, Xte, iters=300):
    mu, sd = Xtr.mean(0), Xtr.std(0) + 1e-6
    A = torch.tensor((Xtr - mu) / sd, dtype=torch.float32)
    B = torch.tensor((Xte - mu) / sd, dtype=torch.float32)
    y = torch.tensor(ytr, dtype=torch.float32)
    w = torch.zeros(A.shape[1], requires_grad=True)
    b = torch.zeros(1, requires_grad=True)
    opt = torch.optim.LBFGS([w, b], max_iter=iters)

    def closure():
        opt.zero_grad()
        loss = F.binary_cross_entropy_with_logits(A @ w + b, y) + 1e-4 * (w * w).sum()
        loss.backward()
        return loss
    opt.step(closure)
    with torch.no_grad():
        return (B @ w + b).numpy()


# ---------------------------------------------------------------- training
def batches(n, bs, shuffle, rng, lengths=None):
    """Shuffled batches; with `lengths`, rows are sorted by length inside chunks of 50
    batches so each batch can be trimmed to its own longest row."""
    idx = list(range(n))
    if shuffle:
        rng.shuffle(idx)
    if lengths is not None:
        chunk = bs * 50
        idx = [i for c in range(0, n, chunk) for i in sorted(idx[c:c + chunk], key=lambda j: lengths[j])]
        groups = [idx[i:i + bs] for i in range(0, n, bs)]
        if shuffle:
            rng.shuffle(groups)
        yield from groups
        return
    for i in range(0, n, bs):
        yield idx[i:i + bs]


def trim(x, *others):
    """Cut right padding to the longest row in the batch (sequences are left-aligned)."""
    L = max(1, int((x != PAD).sum(1).max()))
    return (x[:, :L],) + tuple(o[:, :L] for o in others)


@torch.no_grad()
def screen_scores(model, X, bs=4096):
    model.eval()
    return np.concatenate([torch.sigmoid(model(trim(X[i:i + bs])[0].to(DEV))).cpu().numpy()
                           for i in range(0, len(X), bs)])


@torch.no_grad()
def lm_logprob(model, X, M, bs=2048):
    """Mean log-prob per target token of each row's child (+EOS) given its parent."""
    model.eval()
    out = []
    for i in range(0, len(X), bs):
        x, m = trim(X[i:i + bs], M[i:i + bs])
        x, m = x.to(DEV), m.to(DEV)
        logp = F.log_softmax(model(x[:, :-1]), -1)
        tgt = x[:, 1:]
        lp = logp.gather(-1, tgt.unsqueeze(-1)).squeeze(-1) * m[:, :-1]
        out.append((lp.sum(1) / m[:, :-1].sum(1).clamp(min=1)).cpu().numpy())
    return np.concatenate(out)


def train_screen(Xtr, ytr, Xva, yva, V, epochs, log, rng):
    model = Screen(V).to(DEV)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
    pos_w = ((1 - ytr.mean()) / ytr.mean()).detach().to(DEV)
    best, best_state = -1, None
    for ep in range(epochs):
        model.train()
        t0, tot = time.time(), 0.0
        lengths = (Xtr != PAD).sum(1).tolist()
        for b in batches(len(Xtr), 512, True, rng, lengths):
            logits = model(trim(Xtr[b])[0].to(DEV))
            loss = F.binary_cross_entropy_with_logits(logits, ytr[b].to(DEV), pos_weight=pos_w)
            opt.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            tot += loss.item() * len(b)
        va = auroc(screen_scores(model, Xva), yva.numpy())
        log(f"screen epoch {ep + 1}/{epochs} loss {tot / len(Xtr):.4f} val_auroc {va:.4f} ({time.time() - t0:.0f}s)")
        if va > best:
            best, best_state = va, {k: v.clone() for k, v in model.state_dict().items()}
    model.load_state_dict(best_state)
    return model, best


def train_generator(Xtr, Mtr, Xva, Mva, V, epochs, log, rng):
    model = Generator(V).to(DEV)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
    best, best_state = -1e9, None
    for ep in range(epochs):
        model.train()
        t0, tot, cnt = time.time(), 0.0, 0.0
        lengths = (Xtr != PAD).sum(1).tolist()
        for b in batches(len(Xtr), 256, True, rng, lengths):
            x, m = trim(Xtr[b], Mtr[b])
            x, m = x.to(DEV), m.to(DEV)
            logits = model(x[:, :-1])
            nll = F.cross_entropy(logits.reshape(-1, V), x[:, 1:].reshape(-1), reduction="none")
            nll = (nll.view(x.shape[0], -1) * m[:, :-1]).sum() / m[:, :-1].sum()
            opt.zero_grad()
            nll.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            tot += nll.item()
            cnt += 1
        va = float(lm_logprob(model, Xva, Mva).mean())
        log(f"generator epoch {ep + 1}/{epochs} nll {tot / cnt:.4f} val_logprob/token {va:.4f} ({time.time() - t0:.0f}s)")
        if va > best:
            best, best_state = va, {k: v.clone() for k, v in model.state_dict().items()}
    model.load_state_dict(best_state)
    return model, best


@torch.no_grad()
def sample_children(model, parents_ids, vocab, n, temperature=0.9, rng_seed=0):
    torch.manual_seed(rng_seed)
    model.eval()
    out = []
    for p in parents_ids[:n]:
        seq = [BOS] + p + [SEP]
        child = []
        for _ in range(MAXG + 1):
            logits = model(torch.tensor([seq], device=DEV))[0, -1] / temperature
            logits[[PAD, CLS, SEP, BOS, UNK]] = -1e9
            tok = int(torch.multinomial(F.softmax(logits, -1), 1))
            if tok == EOS:
                break
            child.append(tok)
            seq.append(tok)
        out.append([vocab[t] for t in child])
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("data")
    ap.add_argument("outdir")
    ap.add_argument("--threads", type=int, default=os.cpu_count())
    ap.add_argument("--epochs-screen", type=int, default=4)
    ap.add_argument("--epochs-gen", type=int, default=6)
    ap.add_argument("--max-train", type=int, default=0, help="cap training rows (0 = all), for a smoke run")
    ap.add_argument("--samples", type=int, default=500)
    ap.add_argument("--device", default="cpu")
    a = ap.parse_args()
    global DEV
    DEV = torch.device(a.device)
    torch.set_num_threads(a.threads)
    os.makedirs(a.outdir, exist_ok=True)
    logf = open(os.path.join(a.outdir, "train.log"), "a", encoding="utf-8")

    def log(msg):
        line = f"{time.strftime('%H:%M:%S')} {msg}"
        print(line, flush=True)
        logf.write(line + "\n")
        logf.flush()

    random.seed(0)
    np.random.seed(0)
    torch.manual_seed(0)
    rng = random.Random(0)

    rows = load(a.data)
    tr, va, te, cuts = split_by_epoch(rows)
    if a.max_train:
        rng.shuffle(tr)
        tr = tr[:a.max_train]
    vocab = build_vocab(tr)
    idx = {t: i for i, t in enumerate(vocab)}
    train_children = {r["child_id"] for r in tr}
    novel = np.array([r["child_id"] not in train_children for r in te])
    log(f"rows {len(rows)} train {len(tr)} val {len(va)} test {len(te)} (novel test {int(novel.sum())}); "
        f"epoch cuts {cuts}; vocab {len(vocab)}; train base rate {np.mean([r['label'] for r in tr]):.3f}")

    # baselines
    yte = np.array([r["label"] for r in te])
    ftr = np.array([features(r) for r in tr[:200000]], dtype=np.float64)
    fte = np.array([features(r) for r in te], dtype=np.float64)
    results = {"split": {"epoch_cuts": cuts, "train": len(tr), "val": len(va), "test": len(te),
                         "novel_test": int(novel.sum())}, "test": [], "test_novel": []}
    lr_scores = logistic(ftr, np.array([r["label"] for r in tr[:200000]]), fte)
    rule = fte[:, -1]
    for name, s in (("baseline: has alloc+copy+fork", rule), ("baseline: logistic on 15 hand features", lr_scores)):
        results["test"].append(report(name, s, yte))
        results["test_novel"].append(report(name, s[novel], yte[novel]))
        log(f"{name}: test auroc {results['test'][-1]['auroc']:.4f}")

    # screen
    Xtr, ytr = screen_tensor(tr, idx)
    Xva, yva = screen_tensor(va, idx)
    Xte, _ = screen_tensor(te, idx)
    screen, best = train_screen(Xtr, ytr, Xva, yva, len(vocab), a.epochs_screen, log, rng)
    s = screen_scores(screen, Xte)
    results["test"].append(report("gridlm screen", s, yte))
    results["test_novel"].append(report("gridlm screen", s[novel], yte[novel]))
    log(f"screen: test auroc {results['test'][-1]['auroc']:.4f} (novel {results['test_novel'][-1]['auroc']:.4f})")
    torch.save({k: v.cpu() for k, v in screen.state_dict().items()}, os.path.join(a.outdir, "screen.pt"))

    # generator: successful children only
    trp = [r for r in tr if r["label"] == 1]
    vap = [r for r in va if r["label"] == 1]
    Gtr, Mtr = lm_tensor(trp, idx)
    Gva, Mva = lm_tensor(vap, idx)
    gen, _ = train_generator(Gtr, Mtr, Gva, Mva, len(vocab), a.epochs_gen, log, rng)
    torch.save({k: v.cpu() for k, v in gen.state_dict().items()}, os.path.join(a.outdir, "generator.pt"))
    Gte, Mte = lm_tensor(te, idx)
    lp = lm_logprob(gen, Gte, Mte)
    results["test"].append(report("gridlm generator log p(child|parent)", lp, yte))
    results["test_novel"].append(report("gridlm generator log p(child|parent)", lp[novel], yte[novel]))
    log(f"generator ranking: test auroc {results['test'][-1]['auroc']:.4f}")

    # samples: validity, novelty, and the screen's opinion (model judging model: indicative only)
    test_parents = [enc_pair(r, idx)[0] for r in te]
    rng.shuffle(test_parents)
    kids = sample_children(gen, test_parents, vocab, a.samples)
    parents_txt = [[vocab[t] for t in p] for p in test_parents[:a.samples]]
    valid = [1 <= len(k) <= MAXG for k in kids]
    afc = [all(op in k for op in ("alloc", "copy", "fork")) for k in kids]
    same = [k == p for k, p in zip(kids, parents_txt)]
    fake = [{"parent": p, "child": k, "label": 0, "child_id": ""} for p, k in zip(parents_txt, kids) if 1 <= len(k) <= MAXG]
    Xs, _ = screen_tensor(fake, idx)
    gen_scores = screen_scores(screen, Xs) if len(fake) else np.array([])
    real_scores = s[: len(fake)]
    results["samples"] = {
        "n": len(kids), "valid_length": float(np.mean(valid)), "has_alloc_copy_fork": float(np.mean(afc)),
        "identical_to_parent": float(np.mean(same)),
        "screen_mean_p_generated": float(gen_scores.mean()) if len(gen_scores) else None,
        "screen_mean_p_real_test_children": float(real_scores.mean()),
        "note": "screen scores on generated children are one model judging another; only a live colony is ground truth",
    }
    with open(os.path.join(a.outdir, "samples.jsonl"), "w", encoding="utf-8") as f:
        for p, k in zip(parents_txt, kids):
            f.write(json.dumps({"parent": p, "child": k}) + "\n")
    results["vocab"] = vocab
    results["config"] = {"d_model": 128, "layers": 4, "heads": 4, "ff": 512, "max_genome": MAXG,
                         "label": "child genome born >= 2 times in its epoch",
                         "screen_params": sum(p.numel() for p in screen.parameters()),
                         "generator_params": sum(p.numel() for p in gen.parameters())}
    with open(os.path.join(a.outdir, "results.json"), "w", encoding="utf-8") as f:
        json.dump(results, f, indent=1)
    with open(os.path.join(a.outdir, "vocab.json"), "w", encoding="utf-8") as f:
        json.dump(vocab, f)
    log("done: " + json.dumps(results["samples"]))


if __name__ == "__main__":
    main()