"""Supervised fine-tuning of the tactic policy: goal -> next tactic.

Usage: python train.py --data data/sft_train.jsonl [--data more.jsonl] --out models/NAME
       [--base PATH] [--n N] [--epochs E] [--lr LR] [--bs B] [--maxlen L] [--minutes M]
LoRA (rank --lora, default 64; 0 = full fine-tune) on a bf16 base, bf16 autocast on
Apple MPS; checkpoints are saved merged.  Loss is on the
tactic tokens only.  Logs to logs/NAME.train.log.
"""
import argparse, json, math, random, time
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
HERE = Path(__file__).resolve().parent
BASE = "Qwen/Qwen2.5-Coder-0.5B-Instruct"

ap = argparse.ArgumentParser()
ap.add_argument("--data", action="append", required=True)
ap.add_argument("--weights", action="append", type=int, help="repeat factor per data file")
ap.add_argument("--out", required=True); ap.add_argument("--base", default=BASE)
ap.add_argument("--n", type=int, default=0); ap.add_argument("--epochs", type=float, default=1)
ap.add_argument("--lr", type=float, default=5e-5); ap.add_argument("--bs", type=int, default=8)
ap.add_argument("--maxlen", type=int, default=512); ap.add_argument("--minutes", type=float, default=0)
ap.add_argument("--lora", type=int, default=64)
ap.add_argument("--valid", default=str(HERE / "data/sft_valid.jsonl"))
a = ap.parse_args()
out = HERE / "models" / a.out if not a.out.startswith("/") else Path(a.out)
out.mkdir(parents=True, exist_ok=True)
log = open(HERE / "logs" / f"{out.name}.train.log", "a")
def P(*x):
    s = time.strftime("%H:%M:%S ") + " ".join(map(str, x)); print(s, flush=True); log.write(s + "\n"); log.flush()

dev = "mps"
tok = AutoTokenizer.from_pretrained(a.base)
model = AutoModelForCausalLM.from_pretrained(a.base, dtype=torch.bfloat16).to(dev)
if a.lora:
    from peft import LoraConfig, get_peft_model, PeftModel
    if (Path(a.base) / "adapter_config.json").exists():
        pass  # bases are always merged checkpoints (see save below)
    model = get_peft_model(model, LoraConfig(r=a.lora, lora_alpha=2 * a.lora, lora_dropout=0.0,
                           target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]))
    for p_ in model.parameters():
        if p_.requires_grad: p_.data = p_.data.float()   # fp32 adapter weights
    inner = model.base_model.model
    inner.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    inner.enable_input_require_grads()
else:
    inner = model
model.config.use_cache = False
inner.model.embed_tokens.weight.requires_grad_(False)   # tied with the output head; keeps the vocabulary fixed

def lm_loss(ids, lab, att):
    """Cross-entropy on the tactic tokens only; the output head is applied only
    at those positions (the full-vocabulary logits would not fit in memory)."""
    h = inner.model(input_ids=ids, attention_mask=att).last_hidden_state[:, :-1]
    y = lab[:, 1:]
    sel = y != -100
    logits = inner.lm_head(h[sel]).float()
    return torch.nn.functional.cross_entropy(logits, y[sel])
EOS = tok.eos_token_id

def encode(r):
    p = tok(r["prompt"], add_special_tokens=False)["input_ids"]
    c = tok(r["completion"], add_special_tokens=False)["input_ids"] + [EOS]
    p = p[-(a.maxlen - len(c)):] if len(p) + len(c) > a.maxlen else p
    if len(c) >= a.maxlen: return None
    return p + c, [-100] * len(p) + c

rows = []
ws = a.weights or [1] * len(a.data)
for d, w in zip(a.data, ws):
    rs = [json.loads(l) for l in open(d)]
    rows += rs * w
random.Random(0).shuffle(rows)
if a.n: rows = rows[:a.n]
data = [e for e in map(encode, rows) if e]
valid = [e for e in map(encode, [json.loads(l) for l in open(a.valid)][:300]) if e]
P("train examples", len(data), "valid", len(valid), "base", a.base)

def batches(ds, bs, shuffle=True):
    idx = list(range(len(ds)))
    if shuffle: random.shuffle(idx)
    # bucket by length inside chunks of 50 batches to limit padding
    for i in range(0, len(idx), bs * 50):
        chunk = sorted(idx[i:i + bs * 50], key=lambda j: len(ds[j][0]))
        bl = [chunk[k:k + bs] for k in range(0, len(chunk), bs)]
        if shuffle: random.shuffle(bl)
        for b in bl: yield b

def collate(ds, b):
    m = max(len(ds[j][0]) for j in b)
    ids = torch.full((len(b), m), EOS); lab = torch.full((len(b), m), -100); att = torch.zeros((len(b), m), dtype=torch.long)
    for r, j in enumerate(b):
        x, y = ds[j]; ids[r, :len(x)] = torch.tensor(x); lab[r, :len(y)] = torch.tensor(y); att[r, :len(x)] = 1
    return ids.to(dev), lab.to(dev), att.to(dev)

@torch.no_grad()
def evaluate():
    model.eval(); tot, n = 0.0, 0
    for b in batches(valid, a.bs, False):
        ids, lab, att = collate(valid, b)
        with torch.autocast("mps", dtype=torch.bfloat16):
            tot += lm_loss(ids, lab, att).item(); n += 1
    model.train(); return tot / n

def save():
    """Save a merged, standalone checkpoint (base + adapter) in bf16."""
    if a.lora:
        import copy
        m = copy.deepcopy(model).merge_and_unload()
        m.to(torch.bfloat16).save_pretrained(out); del m
    else:
        model.save_pretrained(out)
    tok.save_pretrained(out)

steps = int(math.ceil(len(data) / a.bs) * a.epochs)
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=a.lr, weight_decay=0.0)
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1, (s + 1) / 100) * max(0.05, 1 - s / steps))
P("steps", steps, "valid loss before", round(evaluate(), 4))
t0 = time.time(); step = 0; run = 0.0; done = False
model.train()
while not done:
    for b in batches(data, a.bs):
        ids, lab, att = collate(data, b)
        with torch.autocast("mps", dtype=torch.bfloat16):
            loss = lm_loss(ids, lab, att)
        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; run = 0.98 * run + 0.02 * loss.item() if step > 1 else loss.item()
        if step % 50 == 0:
            el = time.time() - t0
            P(f"step {step}/{steps} loss {run:.4f} {el/step:.2f}s/step eta {(steps-step)*el/step/60:.0f}min")
        if step % 500 == 0:
            P("valid loss", round(evaluate(), 4)); save()
        if step >= steps or (a.minutes and time.time() - t0 > a.minutes * 60):
            done = True; break
P("final valid loss", round(evaluate(), 4), "steps", step)
save()
json.dump(vars(a) | {"steps_done": step, "examples": len(data)}, open(out / "train_args.json", "w"), indent=1)
P("saved", out)