"""Topic 6: train the Bengio-style N-gram MLP language model on the full WikiText-103 as a batch job.

This is Topic 5's notebook turned into a script. Every setting lives in DEFAULTS below and can be
overridden on the command line, e.g.

    python train_ngram.py --n 5 --work_dir runs/wt103_n5 --max_steps 2000

Re-running with the same --work_dir automatically resumes from <work_dir>/latest.pt, so a job that
hits its walltime can simply be resubmitted.
"""
import argparse
import json
import math
import os
import time
from collections import Counter

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import BatchSampler, DataLoader, Dataset, RandomSampler, SequentialSampler

DEFAULTS = {
    # where files live
    "data_dir": "/project/scv/examples/machine_learning/tutorials/LLM/data/datasets/WikiText103/wikitext-103",
    "work_dir": "runs/wt103",      # checkpoints, metrics log and the exported model go here
    "cache_dir": "",               # tokenized data cache; "" means <work_dir>/cache
    # data
    "train_mb": 0,                 # MB of wiki.train.tokens to use; 0 = the whole file (~539 MB)
    "vocab_size": 32_768,          # most frequent words kept, incl. <unk> and <eos>
    # model
    "n": 4,                        # N of the N-gram: the previous n-1 words predict the next word
    "n_embd": 128,                 # embedding dimension per word
    "n_hidden": 1024,              # hidden layer width
    # training
    "batch_size": 1024,
    "lr": 1e-3,                    # peak learning rate (AdamW)
    "min_lr": 1e-4,                # final learning rate at the end of the cosine decay
    "warmup_steps": 2000,
    "weight_decay": 0.0,
    "max_steps": 300_000,
    "eval_interval": 5000,         # evaluate on the full validation set and checkpoint this often
    "log_interval": 500,
    "num_workers": -1,             # DataLoader workers; -1 = (cores given by the batch system) - 1
    "amp": 1,                      # 1 = bfloat16 mixed precision on GPUs that support it
    "time_limit_hours": 11.5,      # stop cleanly (with a checkpoint) before the job's walltime
    "seed": 1337,
}
# settings that define the data/model; they must not change when resuming a run
SHAPE_KEYS = ["train_mb", "vocab_size", "n", "n_embd", "n_hidden"]


def get_config():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    for key, value in DEFAULTS.items():
        parser.add_argument(f"--{key}", type=type(value), default=value)
    parser.add_argument("--fresh", action="store_true", help="ignore an existing checkpoint and start over")
    cfg = vars(parser.parse_args())
    if not cfg["cache_dir"]:
        cfg["cache_dir"] = os.path.join(cfg["work_dir"], "cache")
    if cfg["num_workers"] < 0:
        cfg["num_workers"] = max(int(os.environ.get("NSLOTS", 1)) - 1, 0)
    return cfg


# ----------------------------------------------------------------------------- data

def iter_lines(path, max_bytes=0):
    """Yield the words of each non-empty line, plus an <eos> token, reading at most ~max_bytes."""
    read = 0
    with open(path, encoding="utf-8") as f:
        for line in f:
            read += len(line)
            if max_bytes and read > max_bytes:
                break
            words = line.split()
            if words:
                yield words + ["<eos>"]


def build_cache(cfg):
    """Tokenize the dataset once into <cache_dir>/{vocab.json, train.npy, valid.npy, test.npy}."""
    tag = f"{cfg['train_mb'] or 'all'}MB_v{cfg['vocab_size']}"
    cache = os.path.join(cfg["cache_dir"], tag)
    vocab_file = os.path.join(cache, "vocab.json")
    if os.path.exists(vocab_file):
        print(f"using cached data in {cache}", flush=True)
        return cache

    os.makedirs(cache, exist_ok=True)
    t0 = time.time()
    train_path = os.path.join(cfg["data_dir"], "wiki.train.tokens")
    counts = Counter()
    for words in iter_lines(train_path, cfg["train_mb"] * 1_000_000):
        counts.update(words)
    specials = ["<unk>", "<eos>"]
    itos = specials + [w for w, _ in counts.most_common() if w not in specials][: cfg["vocab_size"] - len(specials)]
    stoi = {w: i for i, w in enumerate(itos)}
    print(f"counted {sum(counts.values()):,} words, {len(counts):,} distinct, in {time.time() - t0:.0f} s", flush=True)

    dtype = np.uint16 if len(itos) <= 2**16 else np.int32   # half the disk/RAM of int32
    for split in ["train", "valid", "test"]:
        max_bytes = cfg["train_mb"] * 1_000_000 if split == "train" else 0
        ids = [stoi.get(w, 0) for words in iter_lines(os.path.join(cfg["data_dir"], f"wiki.{split}.tokens"), max_bytes)
               for w in words]
        np.save(os.path.join(cache, f"{split}.npy"), np.array(ids, dtype=dtype))
        print(f"  {split}: {len(ids):,} tokens", flush=True)
        del ids
    json.dump(itos, open(vocab_file, "w"))                   # written last: marks the cache complete
    print(f"tokenized in {time.time() - t0:.0f} s", flush=True)
    return cache


class WindowDataset(Dataset):
    """Example i is (ids[i : i+n-1], ids[i+n-1]).

    Unlike Topic 5, __getitem__ receives a whole *list* of indices (we pass a BatchSampler to the
    DataLoader), so one vectorized NumPy lookup builds a full batch instead of 1024 tiny ones.
    """

    def __init__(self, ids, context):
        self.ids, self.context = ids, context

    def __len__(self):
        return len(self.ids) - self.context

    def __getitem__(self, idx):
        idx = np.asarray(idx)[:, None] + np.arange(self.context + 1)
        chunk = torch.from_numpy(self.ids[idx].astype(np.int64))
        return chunk[:, :-1], chunk[:, -1]


def make_loader(ids, cfg, shuffle, seed=0):
    ds = WindowDataset(ids, cfg["n"] - 1)
    if shuffle:
        sampler = RandomSampler(ds, generator=torch.Generator().manual_seed(seed))
    else:
        sampler = SequentialSampler(ds)
    # batch_size=None: the BatchSampler already makes batches, so the DataLoader must not batch again
    return DataLoader(ds, sampler=BatchSampler(sampler, cfg["batch_size"], drop_last=shuffle), batch_size=None,
                      num_workers=cfg["num_workers"], pin_memory=torch.cuda.is_available(),
                      persistent_workers=False)


# ----------------------------------------------------------------------------- model

class NGramMLP(nn.Module):
    """Bengio et al. (2003): embed the n-1 context words, concatenate, tanh hidden layer, softmax output."""

    def __init__(self, vocab_size, n, n_embd, n_hidden):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, n_embd)
        self.hidden = nn.Linear((n - 1) * n_embd, n_hidden)
        self.out = nn.Linear(n_hidden, vocab_size)

    def forward(self, x):                       # x: (B, n-1) word ids
        h = torch.tanh(self.hidden(self.emb(x).flatten(1)))
        return self.out(h)                      # (B, vocab_size) logits


def build_model(cfg):
    return NGramMLP(cfg["vocab_size"], cfg["n"], cfg["n_embd"], cfg["n_hidden"])


def lr_at(step, cfg):
    """Linear warmup, then cosine decay from lr to min_lr over max_steps."""
    if step < cfg["warmup_steps"]:
        return cfg["lr"] * (step + 1) / cfg["warmup_steps"]
    progress = min((step - cfg["warmup_steps"]) / max(cfg["max_steps"] - cfg["warmup_steps"], 1), 1.0)
    return cfg["min_lr"] + 0.5 * (cfg["lr"] - cfg["min_lr"]) * (1 + math.cos(math.pi * progress))


@torch.no_grad()
def evaluate(model, loader, device, autocast):
    """Average cross-entropy over every window of a split (exact, not a sample)."""
    model.eval()
    total, count = 0.0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)
        with autocast():
            logits = model(xb)
        total += F.cross_entropy(logits.float(), yb, reduction="sum").item()
        count += len(yb)
    model.train()
    return total / count


# ----------------------------------------------------------------------------- main

def main():
    cfg = get_config()
    t_start = time.time()
    torch.manual_seed(cfg["seed"])
    device = "cuda" if torch.cuda.is_available() else "cpu"
    use_amp = bool(cfg["amp"]) and device == "cuda" and torch.cuda.is_bf16_supported()
    autocast = lambda: torch.autocast(device_type=device, dtype=torch.bfloat16, enabled=use_amp)
    torch.backends.cuda.matmul.allow_tf32 = True             # faster fp32 matmuls on Ampere+ GPUs
    print(f"device {device} ({torch.cuda.get_device_name() if device == 'cuda' else 'cpu'}), "
          f"bf16 autocast {use_amp}, {cfg['num_workers']} loader workers", flush=True)

    os.makedirs(cfg["work_dir"], exist_ok=True)
    latest_path = os.path.join(cfg["work_dir"], "latest.pt")
    best_path = os.path.join(cfg["work_dir"], "best.pt")
    metrics_path = os.path.join(cfg["work_dir"], "metrics.jsonl")

    # --- data
    cache = build_cache(cfg)
    itos = json.load(open(os.path.join(cache, "vocab.json")))
    cfg["vocab_size"] = len(itos)
    splits = {s: np.load(os.path.join(cache, f"{s}.npy"), mmap_mode="r") for s in ["train", "valid", "test"]}
    val_loader = make_loader(splits["valid"], cfg, shuffle=False)

    # --- model, optimizer, and (if present) the checkpoint to resume from
    model = build_model(cfg).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
    step, best_val = 0, float("inf")
    if os.path.exists(latest_path) and not cfg["fresh"]:
        ckpt = torch.load(latest_path, map_location=device, weights_only=True)
        changed = [k for k in SHAPE_KEYS if ckpt["config"][k] != cfg[k]]
        if changed:
            raise SystemExit(f"{latest_path} was made with different {changed}; use another --work_dir or --fresh")
        model.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        step, best_val = ckpt["step"], ckpt["best_val"]
        print(f"resumed from {latest_path} at step {step:,} (best val loss {best_val:.4f})", flush=True)
    json.dump(cfg, open(os.path.join(cfg["work_dir"], "config.json"), "w"), indent=2)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"{n_params:,} parameters | {len(splits['train']):,} training tokens | "
          f"{len(splits['train']) // cfg['batch_size']:,} steps per epoch", flush=True)

    def save(path):
        torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), "config": cfg,
                    "itos": itos, "step": step, "best_val": best_val}, path + ".tmp")
        os.replace(path + ".tmp", path)                      # atomic: a killed job never leaves half a file

    def log(record):
        with open(metrics_path, "a") as f:
            f.write(json.dumps(record) + "\n")
        print(" | ".join(f"{k} {v:.4g}" if isinstance(v, float) else f"{k} {v}" for k, v in record.items()), flush=True)

    # --- training loop (a new shuffle seed per epoch, so a resumed run doesn't repeat the same order)
    epoch = step // max(len(splits["train"]) // cfg["batch_size"], 1)
    batches = iter(make_loader(splits["train"], cfg, shuffle=True, seed=cfg["seed"] + epoch))
    model.train()
    t0, tokens_seen, out_of_time = time.time(), 0, False
    while step < cfg["max_steps"]:
        try:
            xb, yb = next(batches)
        except StopIteration:
            epoch += 1
            batches = iter(make_loader(splits["train"], cfg, shuffle=True, seed=cfg["seed"] + epoch))
            xb, yb = next(batches)
        xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)

        for group in optimizer.param_groups:
            group["lr"] = lr_at(step, cfg)
        with autocast():
            loss = F.cross_entropy(model(xb).float(), yb)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
        step += 1
        tokens_seen += len(yb)

        if step % cfg["log_interval"] == 0:
            dt = time.time() - t0
            log({"step": step, "epoch": epoch, "train_loss": loss.item(), "lr": optimizer.param_groups[0]["lr"],
                 "ms_per_step": 1000 * dt / cfg["log_interval"], "tokens_per_s": tokens_seen / dt})
            t0, tokens_seen = time.time(), 0

        out_of_time = (time.time() - t_start) / 3600 > cfg["time_limit_hours"]
        if step % cfg["eval_interval"] == 0 or step == cfg["max_steps"] or out_of_time:
            val_loss = evaluate(model, val_loader, device, autocast)
            is_best = val_loss < best_val
            best_val = min(best_val, val_loss)
            save(latest_path)
            if is_best:
                save(best_path)
            log({"step": step, "val_loss": val_loss, "val_ppl": math.exp(val_loss), "best": is_best})
            t0, tokens_seen = time.time(), 0
        if out_of_time:
            print(f"time limit of {cfg['time_limit_hours']} h reached at step {step:,}: resubmit to continue", flush=True)
            return

    # --- final: test-set score of the best model, and a slim export for inference (no optimizer state)
    ckpt = torch.load(best_path, map_location=device, weights_only=True)
    model.load_state_dict(ckpt["model"])
    test_loader = make_loader(splits["test"], cfg, shuffle=False)
    test_loss = evaluate(model, test_loader, device, autocast)
    log({"step": ckpt["step"], "final_val_loss": ckpt["best_val"], "test_loss": test_loss, "test_ppl": math.exp(test_loss)})
    torch.save({"model": model.state_dict(), "config": cfg, "itos": itos, "step": ckpt["step"],
                "val_loss": ckpt["best_val"], "test_loss": test_loss}, os.path.join(cfg["work_dir"], "model.pt"))
    print(f"done in {(time.time() - t_start) / 3600:.2f} h; exported {os.path.join(cfg['work_dir'], 'model.pt')}", flush=True)


if __name__ == "__main__":
    main()