import argparse, yaml, random, os, json, math
from pathlib import Path

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--config",default="configs/lamm_tiny.yaml")
    args=ap.parse_args()
    cfg=yaml.safe_load(open(args.config))

    try:
        import torch
        from lamm.model import LAMMModel
        from lamm.data import stream_hf, batched, format_sft
        from lamm.tokenizer import ByteTokenizer
    except Exception as e:
        raise SystemExit("Install training dependencies: pip install -e '.[train]'") from e

    seed=cfg["seed"]
    random.seed(seed); torch.manual_seed(seed)
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu")

    tok=ByteTokenizer()
    model=LAMMModel(
        vocab_size=tok.vocab_size,
        d_model=cfg["model"]["d_model"],
        n_heads=cfg["model"]["n_heads"],
        d_ff=cfg["model"]["d_ff"],
        n_layers=cfg["model"]["n_layers"],
        lam=cfg["model"]["state_lambda"],
    ).to(device)

    opt=torch.optim.AdamW(
        model.parameters(),
        lr=cfg["training"]["learning_rate"],
        weight_decay=cfg["training"]["weight_decay"]
    )

    out=Path(cfg["training"]["output_dir"])
    out.mkdir(parents=True,exist_ok=True)
    (out/"config.json").write_text(json.dumps(cfg,indent=2))

    stream=stream_hf(
        cfg["data"]["dataset"],
        cfg["data"]["split"],
        cfg["data"].get("source")
    )

    max_len=cfg["model"]["max_length"]
    batch_size=cfg["training"]["batch_size"]
    grad_acc=cfg["training"]["gradient_accumulation"]
    max_steps=cfg["training"]["max_steps"]
    ckpt_every=cfg["training"]["checkpoint_every"]

    def make_batch(records):
        seqs=[tok.encode(format_sft(r),max_len) for r in records]
        L=max(len(x) for x in seqs)
        L=min(L,max_len)
        ids=torch.full((len(seqs),L),tok.PAD,dtype=torch.long)
        labels=torch.full_like(ids,-100)
        for i,s in enumerate(seqs):
            s=s[:L]
            ids[i,:len(s)]=torch.tensor(s)
            # Prompt masking is approximated at the serialized boundary:
            # train on the complete solution serialization while excluding PAD.
            labels[i,:len(s)]=torch.tensor(s)
        return ids.to(device), labels.to(device)

    model.train()
    running=0.0
    step=0
    seen=0

    for batch in batched(stream,batch_size):
        ids,labels=make_batch(batch)
        outp=model(ids,labels)
        loss=outp["loss"]/grad_acc
        loss.backward()
        running += float(loss.item())
        seen += len(batch)

        if (step+1)%grad_acc==0:
            torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)
            opt.step(); opt.zero_grad(set_to_none=True)

        step += 1
        if step%25==0:
            print(json.dumps({
                "step":step,
                "loss":running/max(1,min(25,step)),
                "seen":seen,
                "device":str(device)
            }))
            running=0.0

        if step%ckpt_every==0:
            ckpt=out/f"step-{step}.pt"
            torch.save({
                "model":model.state_dict(),
                "optimizer":opt.state_dict(),
                "step":step,
                "seen":seen
            },ckpt)
            print("checkpoint",ckpt)

        if step>=max_steps:
            break

    ckpt=out/"final.pt"
    torch.save({
        "model":model.state_dict(),
        "optimizer":opt.state_dict(),
        "step":step,
        "seen":seen
    },ckpt)
    print(json.dumps({
        "status":"TRAINING_COMPLETE",
        "steps":step,
        "samples_seen":seen,
        "checkpoint":str(ckpt),
        "device":str(device)
    }))

if __name__=="__main__":
    main()