#!/usr/bin/env python3
"""Stage 2A: QLoRA-style LoRA distillation of Fable 5 traces into GLM-5.2-FP8.

Trains LoRA adapters on attention projections (late band) of the FROZEN block-FP8
base, using the custom differentiable FP8 autograd Function (fp8_diff_patch).

Driven by a wall-clock budget + step cap (device_map=auto runs the 753B model as a
serial pipeline, so throughput is limited). Saves adapter incrementally.
"""
import os, sys, time, json, math, argparse, random
os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
sys.path.insert(0, "/workspace")
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
import fp8_diff_patch

def log(*a): print(f"[{time.strftime('%H:%M:%S')}]", *a, flush=True)

ap = argparse.ArgumentParser()
ap.add_argument("--model", default="/workspace/glm52-fp8")
ap.add_argument("--data", default="/workspace/fable5-chatml.jsonl")
ap.add_argument("--out", default="/workspace/checkpoints/v2-fable5")
ap.add_argument("--rank", type=int, default=64)
ap.add_argument("--alpha", type=int, default=128)
ap.add_argument("--lora_min_layer", type=int, default=60)
ap.add_argument("--max_seq_len", type=int, default=1024)
ap.add_argument("--lr", type=float, default=2e-5)
ap.add_argument("--grad_accum", type=int, default=8)
ap.add_argument("--epochs", type=int, default=1)
ap.add_argument("--max_steps", type=int, default=0, help="0 = unlimited (use time budget)")
ap.add_argument("--time_budget_min", type=float, default=99999.0)
ap.add_argument("--save_every", type=int, default=50)
ap.add_argument("--log_every", type=int, default=1)
ap.add_argument("--max_examples", type=int, default=0, help="0 = all")
ap.add_argument("--warmup", type=int, default=10)
ap.add_argument("--ablation_coeff", type=float, default=0.5, help="Ablation hook coefficient")
ap.add_argument("--ablation_layers", type=str, default="61,62,63,64,65", help="Comma-separated layer indices for ablation hooks")
ap.add_argument("--seed", type=int, default=42)
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
random.seed(args.seed); torch.manual_seed(args.seed)

ATTN = ["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj", "o_proj"]

log("installing fp8 differentiable patch"); fp8_diff_patch.install()
log("loading tokenizer")
tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
if tok.pad_token is None: tok.pad_token = tok.eos_token

log("loading model ...")
t0 = time.time()
maxmem = {i: "128GiB" for i in range(8)}
model = AutoModelForCausalLM.from_pretrained(
    args.model, trust_remote_code=True, dtype=torch.bfloat16,
    device_map="auto", max_memory=maxmem)
log(f"model loaded in {time.time()-t0:.0f}s")
dm = getattr(model, "hf_device_map", {})
bad = [k for k,v in dm.items() if str(v) in ("cpu","disk") or "meta" in str(v)]
assert not bad, f"OFFLOAD DETECTED (would break backward): {bad[:5]}"
model.config.use_cache = False

# ========== ABLATION HOOKS (refusal direction subtraction during training) ==========
import json as _json
_pca = torch.load("/workspace/refusal_pca.pt", map_location="cpu")
_layer_indices = _json.load(open("/workspace/layer_indices.json"))
_ABLATION_LAYERS = [int(x) for x in args.ablation_layers.split(",")]
_ABLATION_COEFF = args.ablation_coeff
log(f"Installing ablation hooks on layers {_ABLATION_LAYERS} with coeff {_ABLATION_COEFF}")

def _make_ablation_hook(layer_idx):
    comps = _pca[layer_idx][:2]  # top 2 PCA components [2, 6144]
    def hook_fn(module, input, output):
        if isinstance(output, tuple):
            hs = output[0]
        else:
            hs = output
        for comp in comps:
            c = comp.to(device=hs.device, dtype=hs.dtype)
            proj = (hs * c).sum(dim=-1, keepdim=True)
            hs = hs - _ABLATION_COEFF * proj * c
        if isinstance(output, tuple):
            return (hs,) + output[1:]
        return hs
    return hook_fn

_base = model.model if hasattr(model, "model") else model
for _L in _ABLATION_LAYERS:
    _base.layers[_L].register_forward_hook(_make_ablation_hook(_L))
log(f"Ablation hooks installed on {len(_ABLATION_LAYERS)} layers")
# ========== END ABLATION HOOKS ==========


targets = []
for n,_ in model.named_modules():
    if any(n.endswith(a) for a in ATTN):
        p = n.split(".")
        try: li = int(p[p.index("layers")+1])
        except Exception: continue
        if li >= args.lora_min_layer: targets.append(n)
log(f"LoRA targets: {len(targets)} (layers>={args.lora_min_layer})")
lcfg = LoraConfig(r=args.rank, lora_alpha=args.alpha, lora_dropout=0.05, bias="none",
                  task_type="CAUSAL_LM", target_modules=targets)
model = get_peft_model(model, lcfg)
model.print_trainable_parameters()
model.train()

emb_dev = model.get_input_embeddings().weight.device
# loss is produced on the last layer's device; CrossEntropy handled by model when labels passed

# ---------- data ----------
def _ids(messages, add_gen):
    r = tok.apply_chat_template(messages, tokenize=True,
                                add_generation_prompt=add_gen)
    # transformers 5.x returns a BatchEncoding dict; older returns a list
    if hasattr(r, "get") or isinstance(r, dict):
        r = r["input_ids"]
    if r and isinstance(r[0], (list, tuple)):  # batched
        r = r[0]
    return list(r)

def build_example(messages):
    """Return (input_ids, labels) with only assistant tokens supervised."""
    prefix_ids = _ids(messages[:-1], True)
    full_ids = _ids(messages, False)
    if len(full_ids) <= len(prefix_ids):
        return None
    labels = [-100]*len(prefix_ids) + full_ids[len(prefix_ids):]
    # LEFT-truncate: keep the TAIL so the assistant completion (the supervised
    # span) is always preserved; long user prefixes get clipped from the front.
    if len(full_ids) > args.max_seq_len:
        full_ids = full_ids[-args.max_seq_len:]
        labels = labels[-args.max_seq_len:]
    if all(l == -100 for l in labels):  # assistant span longer than max_seq_len
        return None
    return full_ids, labels

log("loading + templating data ...")
examples = []
with open(args.data) as f:
    lines = f.readlines()
random.shuffle(lines)
if args.max_examples: lines = lines[:args.max_examples]
skipped = 0
for line in lines:
    try:
        msgs = json.loads(line)["messages"]
        ex = build_example(msgs)
        if ex is None: skipped += 1; continue
        examples.append(ex)
    except Exception:
        skipped += 1
log(f"prepared {len(examples)} examples (skipped {skipped})")

# ---------- optimizer ----------
trainable = [p for p in model.parameters() if p.requires_grad]
opt = torch.optim.AdamW(trainable, lr=args.lr, betas=(0.9,0.95), weight_decay=0.0)
total_microbatches = len(examples) * args.epochs
steps_per_epoch = math.ceil(len(examples)/args.grad_accum)
planned_steps = steps_per_epoch * args.epochs
if args.max_steps: planned_steps = min(planned_steps, args.max_steps)
def lr_at(step):
    if step < args.warmup: return args.lr * (step+1)/args.warmup
    prog = (step-args.warmup)/max(1, planned_steps-args.warmup)
    return args.lr * 0.5*(1+math.cos(math.pi*min(1.0,prog)))

cfg = dict(model=args.model, data=args.data, rank=args.rank, alpha=args.alpha,
           lora_min_layer=args.lora_min_layer, target_modules=targets,
           max_seq_len=args.max_seq_len, lr=args.lr, grad_accum=args.grad_accum,
           epochs=args.epochs, warmup=args.warmup, planned_steps=planned_steps,
           n_examples=len(examples), lora_dropout=0.05, seed=args.seed,
           method="LoRA on frozen block-FP8 base via custom FP8 autograd Function",
           optimizer="AdamW betas(0.9,0.95) wd0", schedule="cosine")
json.dump(cfg, open(os.path.join(args.out,"training_config.json"),"w"), indent=2)

def save_ckpt(tag, extra=None):
    d = os.path.join(args.out, tag)
    model.save_pretrained(d)   # saves only adapter (PEFT)
    if extra: json.dump(extra, open(os.path.join(d,"train_state.json"),"w"), indent=2)
    log(f"saved checkpoint -> {d}")

# ---------- train ----------
log(f"START training: {len(examples)} ex, accum {args.grad_accum}, "
    f"planned_steps {planned_steps}, budget {args.time_budget_min}min")
loss_hist = []
step = 0; microbatch = 0; running = 0.0; t_start = time.time()
opt.zero_grad(set_to_none=True)
stop = False
for epoch in range(args.epochs):
    if stop: break
    random.shuffle(examples)
    for ids, labels in examples:
        ids_t = torch.tensor([ids], device=emb_dev)
        lab_t = torch.tensor([labels], device=emb_dev)
        out = model(input_ids=ids_t, labels=lab_t)
        loss = out.loss / args.grad_accum
        loss.backward()
        running += out.loss.item()
        microbatch += 1
        if microbatch % args.grad_accum == 0:
            torch.nn.utils.clip_grad_norm_(trainable, 1.0)
            for g in opt.param_groups: g["lr"] = lr_at(step)
            opt.step(); opt.zero_grad(set_to_none=True)
            step += 1
            avg = running/args.grad_accum; running = 0.0
            if step % args.log_every == 0:
                el = time.time()-t_start
                loss_hist.append({"step":step,"loss":round(avg,4),
                                  "lr":lr_at(step),"elapsed_s":round(el,1)})
                json.dump(loss_hist, open(os.path.join(args.out,"loss_log.json"),"w"))
                log(f"step {step}/{planned_steps} loss {avg:.4f} "
                    f"lr {lr_at(step):.2e} elapsed {el/60:.1f}min "
                    f"({el/step:.1f}s/step)")
            if step % args.save_every == 0:
                save_ckpt(f"step-{step}", {"step":step,"loss":avg,"loss_hist":loss_hist[-5:]})
            if args.max_steps and step >= args.max_steps: stop=True; break
            if (time.time()-t_start)/60 >= args.time_budget_min:
                log("TIME BUDGET reached -> stopping"); stop=True; break

# flush any partial accumulation
if microbatch % args.grad_accum != 0:
    torch.nn.utils.clip_grad_norm_(trainable, 1.0); opt.step(); opt.zero_grad(set_to_none=True)

save_ckpt("final", {"final_step":step,"loss_hist":loss_hist})
final_loss = loss_hist[-1]["loss"] if loss_hist else None
summary = dict(completed_steps=step, planned_steps=planned_steps,
               final_loss=final_loss, first_loss=(loss_hist[0]["loss"] if loss_hist else None),
               elapsed_min=round((time.time()-t_start)/60,1),
               examples_seen=microbatch, n_examples=len(examples))
json.dump(summary, open(os.path.join(args.out,"training_summary.json"),"w"), indent=2)
log("DONE", json.dumps(summary))
print("TRAIN_DONE", flush=True)