import os
import argparse
import json

# We WANT the ZeroGPU shim active here, so do NOT pop SPACES_ZERO_GPU.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces          # MUST come before torch
import torch

# Importing train.py pops SPACES_ZERO_GPU, but spaces + torch are already
# initialised with the GPU shim above, so the decoration below still routes
# through the scheduler. Reuse its model / data helpers.
from train import TinyLM, build_vocab, make_batch


@spaces.GPU(duration=1800)
def run(data_path, epochs, batch_size, block_size, n_embed, n_layer, n_head, lr):
    device = "cuda"
    text = open(data_path, encoding="utf-8", errors="ignore").read()
    stoi, itos = build_vocab(text)
    model = TinyLM(len(stoi), n_embed, n_layer, n_head, block_size).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)

    model.train()
    last = None
    for epoch in range(epochs):
        x, y = make_batch(stoi, text, block_size, batch_size, device)
        logits, loss = model(x, tgt=y)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        last = float(loss.item())
        if epoch % 25 == 0 or epoch == epochs - 1:
            print(f"epoch {epoch}/{epochs}  loss {last:.4f}", flush=True)

    torch.save({"model": model.cpu().state_dict(), "vocab_size": len(stoi),
                "itos": {str(k): v for k, v in itos.items()},
                "cfg": {"n_embed": n_embed, "n_layer": n_layer, "n_head": n_head,
                        "block_size": block_size}},
               "cyber_model.pt")
    print("saved cyber_model.pt")
    return f"done loss={last:.4f}"


if __name__ == "__main__":
    p = argparse.ArgumentParser(description="Train TinyLM on ZeroGPU.")
    p.add_argument("--data", required=True)
    p.add_argument("--epochs", type=int, default=500)
    p.add_argument("--batch_size", type=int, default=32)
    p.add_argument("--block_size", type=int, default=128)
    p.add_argument("--n_embed", type=int, default=128)
    p.add_argument("--n_layer", type=int, default=4)
    p.add_argument("--n_head", type=int, default=4)
    p.add_argument("--lr", type=float, default=1e-3)
    a = p.parse_args()
    print(run(a.data, a.epochs, a.batch_size, a.block_size,
              a.n_embed, a.n_layer, a.n_head, a.lr))