"""Full fine-tune of GPT-2 on tokenling chat/diary data.

Loss is applied only to the character spans each row marks as trained (the
pet's words and tool calls, and diary bodies), never to owner lines, headers
or app-provided {tool results}.
"""
import argparse
import collections
import json
import math
import random
import re
import sys
import time

import torch
from transformers import AutoTokenizer, GPT2LMHeadModel

import gen_data as G


def encode(tok, row, max_len):
    text = row["text"]
    enc = tok(text, return_offsets_mapping=True)
    ids, offs = enc["input_ids"], enc["offset_mapping"]
    trained = bytearray(len(text) + 1)
    for a, b in row["spans"]:
        trained[a:b] = b"\x01" * (b - a)
    labels = [i if any(trained[s:e]) else -100 for i, (s, e) in zip(ids, offs)]
    if row["kind"] == "diary":  # learn to stop after the entry
        ids.append(tok.eos_token_id)
        labels.append(tok.eos_token_id)
    return ids[:max_len], labels[:max_len]


def split_chat(row):
    hdr, body = row["text"].split("---\n", 1)
    meta = dict(l.split(": ", 1) for l in hdr.splitlines()[1:] if ": " in l)
    return hdr + "---\n", body, meta["pet"].split(" · ")[0], meta["owner"]


def with_history(row, donor):
    """Prepend a donor chat's lines (names swapped, untrained) as earlier
    history, so the pet learns to stay in character with long contexts."""
    hdr, body, pet, owner = split_chat(row)
    _, dbody, dpet, downer = split_chat(donor)
    swap = {dpet: pet, downer: owner}
    dbody = re.sub(r"\b(%s)\b" % "|".join(map(re.escape, swap)), lambda m: swap[m.group(1)], dbody)
    shift = len(dbody)
    return {**row, "text": hdr + dbody + body,
            "spans": [[a + shift, b + shift] for a, b in row["spans"]]}


def augment(rows, frac):
    chats = [r for r in rows if r["kind"] == "chat"]
    # The nearest history must end in the recipient's starting sleep state;
    # otherwise an untrained donor nap teaches an awake reply while asleep.
    pools = {False: [], True: []}
    for r in chats:
        hdr, body, pet, _ = split_chat(r)
        asleep = f", {pet} is asleep" in hdr
        for ln in body.splitlines():
            if ln == f"* {pet} woke up *": asleep = False
            elif ln.startswith(pet + ": ") and "[sleep]" in ln: asleep = True
        pools[asleep].append(r)
    out = []
    for r in rows:
        out.append(r)
        if r["kind"] == "chat" and random.random() < frac:
            aug = r
            hdr, _, pet, _ = split_chat(r)
            asleep = f", {pet} is asleep" in hdr
            for _ in range(random.choice([1, 2])):
                candidates = pools[asleep]
                if not candidates:
                    break
                d = random.choice(candidates)
                aug = with_history(aug, d)
                hdr, _, pet, _ = split_chat(d)
                asleep = f", {pet} is asleep" in hdr
            out.append(aug)
    return out


OFFER = re.compile(r"^\* .+ (offered .+|to play catch) \*$")


def late_answers(rows, frac, rng):
    """Move the [yes]/[no] of some offer answers to the end of the line.
    Teachers almost always answer first, so a pet that sampled an emote
    before answering learned that the answer was skipped and ate nothing.
    The line keeps its length, so the stored spans still fit."""
    out, moved = [], 0
    for r in rows:
        if r["kind"] != "chat":
            out.append(r)
            continue
        lines = r["text"].split("\n")
        for i in range(1, len(lines)):
            if not OFFER.match(lines[i - 1]) or "{" in lines[i] or rng.random() >= frac:
                continue
            m = re.search(r"( \[(yes|no)\]|^[^:]+: (\[(yes|no)\]) )", lines[i])
            if not m:
                continue
            tag = "[" + (m.group(2) or m.group(4)) + "]"
            if m.group(1).startswith(" "):
                line = lines[i][:m.start()] + lines[i][m.end():]
            else:  # answered first: "Mochi: [yes] nom" -> "Mochi: nom [yes]"
                line = lines[i][:m.start(3)] + lines[i][m.end():]
            if line + " " + tag != lines[i]:
                lines[i] = line + " " + tag
                moved += 1
        out.append({**r, "text": "\n".join(lines)})
    print(f"late answers: moved {moved}")
    return out


def batches(data, bs, shuffle):
    idx = list(range(len(data)))
    if shuffle:
        random.shuffle(idx)
    for i in range(0, len(idx), bs):
        yield [data[j] for j in idx[i:i + bs]]


def collate(rows, pad):
    n = max(len(r[0]) for r in rows)
    ids = torch.full((len(rows), n), pad)
    lab = torch.full((len(rows), n), -100)
    att = torch.zeros((len(rows), n), dtype=torch.long)
    for k, (i, l) in enumerate(rows):
        ids[k, :len(i)] = torch.tensor(i)
        lab[k, :len(l)] = torch.tensor(l)
        att[k, :len(i)] = 1
    return ids, lab, att


@torch.no_grad()
def evaluate(model, data, bs, pad):
    model.eval()
    tot = cnt = 0.0
    for rows in batches(data, bs, False):
        ids, lab, att = (t.cuda() for t in collate(rows, pad))
        with torch.autocast("cuda", dtype=torch.bfloat16):
            out = model(input_ids=ids, attention_mask=att, labels=lab)
        n = (lab[:, 1:] != -100).sum().item()
        tot += out.loss.item() * n
        cnt += n
    model.train()
    return tot / max(cnt, 1)


def fresh_rows(path, drop):
    """Every row is re-checked against the current data rules first: rows
    written for an older app (care events in the owner's mouth, catch without
    an offer) would teach behavior the app no longer has."""
    rows, stale = [], collections.Counter()
    for line in open(path):
        row = json.loads(line)
        try:
            G.recheck(row)
            rows.append(row)
        except ValueError as e:
            stale[str(e).split(":")[0]] += 1
    if stale:
        for k, v in stale.most_common():
            print(f"{v:6}  stale: {k}")
        if not drop:
            sys.exit(f"{sum(stale.values())} rows fail the current rules; fix the data or pass --drop-stale")
    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True)
    ap.add_argument("--base", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--epochs", type=float, default=3)
    ap.add_argument("--lr", type=float, default=6e-5)
    ap.add_argument("--bs", type=int, default=16)
    ap.add_argument("--max-len", type=int, default=1024)
    ap.add_argument("--history-aug", type=float, default=0.25)
    ap.add_argument("--late-answer", type=float, default=0.3,
                    help="share of offer answers whose [yes]/[no] moves to the end of the line")
    ap.add_argument("--drop-stale", action="store_true",
                    help="drop rows the current rules reject instead of refusing to train")
    args = ap.parse_args()
    random.seed(0)
    torch.manual_seed(0)

    tok = AutoTokenizer.from_pretrained(args.base)
    rows = fresh_rows(args.data, args.drop_stale)
    random.shuffle(rows)
    n_eval = max(64, len(rows) // 50)
    evals = [encode(tok, r, args.max_len) for r in rows[:n_eval]]
    # A separate generator, so the held-out split and history draws do not move.
    fit = late_answers(rows[n_eval:], args.late_answer, random.Random(1))
    train = [encode(tok, r, args.max_len) for r in augment(fit, args.history_aug)]
    print(f"train={len(train)} eval={len(evals)} "
          f"tokens={sum(len(d[0]) for d in train)} trained={sum(sum(x != -100 for x in d[1]) for d in train)}")

    model = GPT2LMHeadModel.from_pretrained(args.base, torch_dtype=torch.float32).cuda()
    model.gradient_checkpointing_enable()
    model.config.use_cache = False
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.05)
    steps = math.ceil(len(train) / args.bs * args.epochs)
    warm = max(20, steps // 30)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min(1, (s + 1) / warm) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, s / steps)))))
    print(f"steps={steps} eval_loss0={evaluate(model, evals, args.bs, tok.eos_token_id):.4f}", flush=True)

    step, t0, best = 0, time.time(), 1e9
    while step < steps:
        for rows in batches(train, args.bs, True):
            ids, lab, att = (t.cuda() for t in collate(rows, tok.eos_token_id))
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss = model(input_ids=ids, attention_mask=att, labels=lab).loss
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            opt.zero_grad(set_to_none=True)
            step += 1
            if step % 25 == 0:
                print(f"step {step}/{steps} loss {loss.item():.4f} lr {sched.get_last_lr()[0]:.2e} "
                      f"{time.time() - t0:.0f}s", flush=True)
            if step % 200 == 0 or step == steps:
                ev = evaluate(model, evals, args.bs, tok.eos_token_id)
                print(f"eval step {step}: {ev:.4f}", flush=True)
                if ev < best:
                    best = ev
                    model.config.use_cache = True
                    model.save_pretrained(args.out, safe_serialization=True)
                    tok.save_pretrained(args.out)
                    model.config.use_cache = False
            if step >= steps:
                break
    print(f"best eval {best:.4f}")


if __name__ == "__main__":
    main()