"""
train.py — Train a ~50M-parameter GPT-style causal LM and export a quantized ONNX model.

Designed for a Kaggle Notebook with a T4 GPU. The resulting 4-bit ONNX file stays
well under 50 MB, so it can be uploaded as a single artifact and run client-side
in the browser via ONNX Runtime Web.

==============================================================================
  Kaggle Notebook setup
==============================================================================
  1. New Notebook  ->  Settings  ->  Accelerator: GPU T4 x1
  2. Upload your text corpus as a Kaggle dataset, with a file named `dataset.txt`
     inside it. The mount path is typically /kaggle/input/<dataset-name>/dataset.txt
  3. Add this script as a Notebook utility, OR paste it into a cell.
  4. Install runtime deps (only needed once per session):
         !pip install -q tokenizers onnxruntime onnx
  5. Run:
         !python train.py --dataset /kaggle/input/<dataset-name>/dataset.txt \
                          --out_dir /kaggle/working

==============================================================================
  Pipeline
==============================================================================
  1. Train a byte-level BPE tokenizer (vocab_size=8000) on dataset.txt
  2. Stream-encode dataset.txt -> uint16 binary shards (train.bin / val.bin)
  3. Train a 12-layer / 512-dim transformer (~46M params) with FP16 AMP
  4. Export to ONNX (opset 17, dynamic batch & sequence axes)
  5. Quantize weights to 4-bit via ONNX Runtime's MatMulNBits quantizer
  6. Print final model size and confirm it stays under 50 MB

==============================================================================
  Outputs (in out_dir)
==============================================================================
  - tokenizer.json      : byte-level BPE tokenizer (for browser-side encoding)
  - train.bin/val.bin   : tokenized binary shards (memmap-friendly)
  - model_best.pt       : best-val FP32 checkpoint
  - model_fp32.pt       : final FP32 checkpoint
  - model_fp32.onnx     : unquantized ONNX (for debugging)
  - model_q4.onnx       : 4-bit quantized ONNX   <-- the deployable artifact

  Expected file size:  ~25 MB for the 4-bit model.

==============================================================================
  Browser deployment
==============================================================================
  Use `onnxruntime-web` (npm) to load model_q4.onnx, and `@huggingface/tokenizers`
  (WASM) to load tokenizer.json. Minimal sketch:

      import * as ort from 'onnxruntime-web';
      import { Tokenizer } from '@huggingface/tokenizers';

      const session  = await ort.InferenceSession.create('./model_q4.onnx');
      const tokenizer = await Tokenizer.fromFile('./tokenizer.json');

      function generate(prompt, maxNew = 64) {
        let ids = tokenizer.encode(prompt).ids;
        for (let i = 0; i < maxNew; i++) {
          const ctx = ids.slice(-256);                       // crop to block_size
          const inp = new ort.Tensor('int64', BigInt64Array.from(ctx.map(BigInt)), [1, ctx.length]);
          const out = await session.run({ input_ids: inp });
          const logits = out.logits.data;                    // [1 * ctx.length * vocab]
          const next = sample(logits, ctx.length, vocabSize); // your sampler
          ids = ids.concat(next);
        }
        return tokenizer.decode(ids, true);
      }
"""

from __future__ import annotations

import argparse
import math
import os
import time
from dataclasses import asdict, dataclass
from typing import Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# ============================================================================
# Configuration
# ============================================================================

@dataclass
class Config:
    # ---- Tokenizer ----
    vocab_size: int = 8000          # will be overridden by actual tokenizer size

    # ---- Model architecture  (sums to ~46.1M params) ----
    n_layer:     int = 12
    n_head:      int = 8
    n_embd:      int = 512
    block_size:  int = 256          # max context length
    ffn_mult:    int = 4            # FFN hidden = n_embd * ffn_mult = 2048
    dropout:     float = 0.1
    bias:        bool = False       # GPT-2 / LLaMA convention: no biases

    # ---- Training ----
    batch_size:    int = 32
    grad_accum:    int = 4          # effective batch = 32 * 4 = 128
    learning_rate: float = 3e-4
    weight_decay:  float = 0.1
    max_iters:     int = 5000
    warmup_iters:  int = 200
    eval_interval: int = 500
    eval_iters:    int = 50
    grad_clip:     float = 1.0
    seed:          int = 1337

    # ---- Data split ----
    val_split:     float = 0.05


# ============================================================================
# Model (GPT-style)
# ============================================================================

class CausalSelfAttention(nn.Module):
    """Multi-head causal self-attention with fused QKV projection.

    Manual attention (instead of F.scaled_dot_product_attention) so that the
    ONNX export stays portable across ONNX Runtime Web builds.
    """

    def __init__(self, cfg: Config):
        super().__init__()
        assert cfg.n_embd % cfg.n_head == 0, "n_embd must be divisible by n_head"
        self.n_head  = cfg.n_head
        self.n_embd  = cfg.n_embd
        self.head_dim = cfg.n_embd // cfg.n_head
        self.dropout  = cfg.dropout

        # Fused QKV
        self.c_attn  = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=cfg.bias)
        self.c_proj  = nn.Linear(cfg.n_embd,     cfg.n_embd, bias=cfg.bias)
        self.resid_dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=2)
        # (B, T, C) -> (B, nh, T, hd)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        # Manual scaled-dot-product attention with causal mask
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
        att = att.masked_fill(~causal_mask, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = F.dropout(att, p=self.dropout, training=self.training)
        y = att @ v                                    # (B, nh, T, hd)

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.c_proj(y))


class MLP(nn.Module):
    """Position-wise feed-forward: Linear -> GELU -> Linear."""

    def __init__(self, cfg: Config):
        super().__init__()
        hidden = cfg.n_embd * cfg.ffn_mult
        self.c_fc    = nn.Linear(cfg.n_embd, hidden, bias=cfg.bias)
        self.c_proj  = nn.Linear(hidden, cfg.n_embd, bias=cfg.bias)
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.c_fc(x)
        x = F.gelu(x, approximate="tanh")
        x = self.c_proj(x)
        return self.dropout(x)


class Block(nn.Module):
    """Pre-norm transformer block (GPT-2 style)."""

    def __init__(self, cfg: Config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.attn = CausalSelfAttention(cfg)
        self.ln_2 = nn.LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.mlp  = MLP(cfg)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp (self.ln_2(x))
        return x


class GPT(nn.Module):
    """GPT-style decoder-only transformer.

    Untied input/output embeddings on purpose — this bumps the parameter count
    from ~30M (tied) to ~46M (untied), which is what we want.
    """

    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.wte    = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.wpe    = nn.Embedding(cfg.block_size, cfg.n_embd)
        self.drop   = nn.Dropout(cfg.dropout)
        self.h      = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f   = nn.LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)

        # Init all weights
        self.apply(self._init_weights)
        # Scale residual projections by 1/sqrt(2*n_layer)
        for pn, p in self.named_parameters():
            if pn.endswith("c_proj.weight"):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer))

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def num_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())

    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
        B, T = idx.shape
        assert T <= self.cfg.block_size, f"sequence {T} > block_size {self.cfg.block_size}"
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device).unsqueeze(0)
        x = self.drop(self.wte(idx) + self.wpe(pos))
        for block in self.h:
            x = block(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,
            )
        return logits, loss


class GPTForExport(nn.Module):
    """Thin wrapper that returns only logits — used for ONNX export."""

    def __init__(self, gpt: GPT):
        super().__init__()
        self.gpt = gpt

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        logits, _ = self.gpt(input_ids)
        return logits


# ============================================================================
# Tokenizer
# ============================================================================

def train_tokenizer(dataset_path: str, vocab_size: int, out_dir: str):
    """Train a byte-level BPE tokenizer using the HuggingFace `tokenizers` lib."""
    from tokenizers import Tokenizer
    from tokenizers.models import BPE
    from tokenizers.trainers import BpeTrainer
    from tokenizers.pre_tokenizers import ByteLevel
    from tokenizers.decoders import ByteLevel as ByteLevelDecoder

    tokenizer = Tokenizer(BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False, use_regex=True)
    tokenizer.decoder = ByteLevelDecoder()

    trainer = BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=["<pad>", "<bos>", "<eos>", "<unk>"],
        initial_alphabet=ByteLevel.alphabet(),
        show_progress=True,
    )
    # `tokenizer.train` streams from disk — fine for multi-GB datasets.
    tokenizer.train([dataset_path], trainer)

    pad_id = tokenizer.token_to_id("<pad>")
    if pad_id is not None:
        tokenizer.enable_padding(pad_id=pad_id, pad_token="<pad>")

    out_path = os.path.join(out_dir, "tokenizer.json")
    tokenizer.save(out_path)
    print(f"[tokenizer] saved -> {out_path}  (vocab_size={tokenizer.get_vocab_size()})")
    return tokenizer


def load_tokenizer(out_dir: str):
    from tokenizers import Tokenizer
    return Tokenizer.from_file(os.path.join(out_dir, "tokenizer.json"))


# ============================================================================
# Data pipeline
# ============================================================================

def encode_dataset(tokenizer, dataset_path: str, out_path: str, chunk_chars: int = 5_000_000) -> int:
    """Stream-encode dataset.txt to a binary file of uint16 token ids.

    uint16 is enough for vocab_size up to 65535 — well above our 8000.
    """
    written = 0
    with open(dataset_path, "r", encoding="utf-8", errors="ignore") as fin, \
         open(out_path, "wb") as fout:
        while True:
            chunk = fin.read(chunk_chars)
            if not chunk:
                break
            enc = tokenizer.encode(chunk)
            arr = np.asarray(enc.ids, dtype=np.uint16)
            arr.tofile(fout)
            written += arr.size
    print(f"[data] encoded {written:,} tokens -> {out_path}")
    return written


def load_binary(path: str) -> np.ndarray:
    """Load tokenized data as a memory-mapped numpy array (no full copy)."""
    n_bytes = os.path.getsize(path)
    n = n_bytes // np.dtype(np.uint16).itemsize
    return np.memmap(path, dtype=np.uint16, mode="r", shape=(n,))


def split_train_val(data: np.ndarray, val_split: float) -> Tuple[np.ndarray, np.ndarray]:
    n = len(data)
    n_val = int(n * val_split)
    return data[: n - n_val], data[n - n_val :]


def get_batch(data: np.ndarray, block_size: int, batch_size: int, device: str):
    """Sample a random batch of (x, y) pairs for next-token prediction."""
    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
    x = torch.stack([torch.from_numpy(data[i : i + block_size].astype(np.int64))      for i in ix])
    y = torch.stack([torch.from_numpy(data[i + 1 : i + 1 + block_size].astype(np.int64)) for i in ix])
    if device == "cuda":
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y


# ============================================================================
# Training loop
# ============================================================================

def get_lr(it: int, cfg: Config) -> float:
    """Linear warmup -> cosine decay to 10% of base LR."""
    if it < cfg.warmup_iters:
        return cfg.learning_rate * (it + 1) / cfg.warmup_iters
    if it > cfg.max_iters:
        return cfg.learning_rate * 0.1
    decay_ratio = (it - cfg.warmup_iters) / max(cfg.max_iters - cfg.warmup_iters, 1)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return cfg.learning_rate * coeff


@torch.no_grad()
def estimate_loss(model: GPT, train_data, val_data, cfg: Config, device: str) -> dict:
    out = {}
    model.eval()
    use_amp = device == "cuda"
    for split, data in [("train", train_data), ("val", val_data)]:
        losses = torch.zeros(cfg.eval_iters)
        for k in range(cfg.eval_iters):
            x, y = get_batch(data, cfg.block_size, cfg.batch_size, device)
            with autocast_ctx(use_amp):
                _, loss = model(x, y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out


def autocast_ctx(enabled: bool):
    """Wrapper so we can call `with autocast_ctx(True):` regardless of API version.

    Different PyTorch versions expose autocast in different places with different
    signatures. We try each known shape in turn and catch both ImportError (module
    doesn't exist) and TypeError (signature mismatch).
    """
    if not enabled:
        import contextlib
        return contextlib.nullcontext()
    # Try the new torch.amp API first (PyTorch 2.0+)
    try:
        from torch.amp import autocast as _ac
        try:
            return _ac(device_type="cuda", dtype=torch.float16)
        except TypeError:
            # Older 2.x where the kwarg differs
            return _ac(dtype=torch.float16)
    except (ImportError, TypeError):
        pass
    # Fall back to the legacy torch.cuda.amp API
    try:
        from torch.cuda.amp import autocast as _ac_legacy
        return _ac_legacy(dtype=torch.float16)
    except (ImportError, TypeError, AttributeError):
        pass
    # Last resort: no autocast at all
    import contextlib
    return contextlib.nullcontext()


def make_grad_scaler(enabled: bool):
    """Create a GradScaler if possible; return None if AMP is unavailable.

    Same defensive pattern as autocast_ctx — try each known signature and
    gracefully degrade to plain FP32 if none work.
    """
    if not enabled:
        return None
    # PyTorch 2.0+ new API
    try:
        from torch.amp import GradScaler as _GS
        try:
            return _GS(device_type="cuda")
        except TypeError:
            # The class exists but doesn't accept device_type (intermediate versions)
            return _GS()
    except (ImportError, TypeError):
        pass
    # Legacy API
    try:
        from torch.cuda.amp import GradScaler as _GS_legacy
        return _GS_legacy()
    except (ImportError, TypeError, AttributeError):
        pass
    # No AMP available — train in plain FP32
    print("[amp] GradScaler unavailable; training in plain FP32 (slower but fine).")
    return None


def train_model(model: GPT, train_data, val_data, cfg: Config, device: str, out_dir: str):
    # Optimizer — try `fused=True` (PyTorch 2.x), fall back if unsupported
    try:
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=cfg.learning_rate,
            betas=(0.9, 0.95),
            weight_decay=cfg.weight_decay,
            eps=1e-8,
            fused=(device == "cuda"),
        )
    except TypeError:
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=cfg.learning_rate,
            betas=(0.9, 0.95),
            weight_decay=cfg.weight_decay,
            eps=1e-8,
        )

    scaler = make_grad_scaler(device == "cuda")

    best_val = float("inf")
    t0 = time.time()
    iter_num = 0

    while iter_num < cfg.max_iters:
        # LR schedule
        lr = get_lr(iter_num, cfg)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        # Gradient accumulation
        optimizer.zero_grad(set_to_none=True)
        accum_loss = 0.0
        for _ in range(cfg.grad_accum):
            x, y = get_batch(train_data, cfg.block_size, cfg.batch_size, device)
            with autocast_ctx(device == "cuda"):
                _, loss = model(x, y)
            loss = loss / cfg.grad_accum
            if scaler is not None:
                scaler.scale(loss).backward()
            else:
                loss.backward()
            accum_loss += loss.item()

        if scaler is not None:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
            scaler.step(optimizer)
            scaler.update()
        else:
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
            optimizer.step()

        # Logging
        if iter_num % 100 == 0:
            elapsed = time.time() - t0
            total_tokens = (iter_num + 1) * cfg.batch_size * cfg.block_size * cfg.grad_accum
            tps = total_tokens / max(elapsed, 1e-9)
            print(f"[train] iter {iter_num:5d}/{cfg.max_iters} "
                  f"| loss {accum_loss:.4f} | lr {lr:.2e} "
                  f"| t {elapsed:6.1f}s | {tps:>10,.0f} tok/s", flush=True)

        # Periodic eval
        if iter_num > 0 and iter_num % cfg.eval_interval == 0:
            losses = estimate_loss(model, train_data, val_data, cfg, device)
            print(f"  -> eval  train {losses['train']:.4f}   val {losses['val']:.4f}", flush=True)
            if losses["val"] < best_val:
                best_val = losses["val"]
                torch.save(
                    {"model_state": model.state_dict(), "cfg": asdict(cfg),
                     "iter": iter_num, "val_loss": best_val},
                    os.path.join(out_dir, "model_best.pt"),
                )
                print(f"  -> saved best (val={best_val:.4f})", flush=True)

        iter_num += 1

    # Final checkpoint
    torch.save(
        {"model_state": model.state_dict(), "cfg": asdict(cfg),
         "iter": iter_num, "val_loss": best_val},
        os.path.join(out_dir, "model_fp32.pt"),
    )
    print(f"[train] done. best_val={best_val:.4f}")


# ============================================================================
# ONNX export & quantization
# ============================================================================

def export_onnx(model: GPT, cfg: Config, out_dir: str) -> str:
    """Export model to ONNX with dynamic batch and sequence axes.

    Export happens on CPU so the resulting graph is device-agnostic.
    """
    model.eval()
    wrapper = GPTForExport(model).to("cpu")
    dummy = torch.zeros((1, cfg.block_size), dtype=torch.long)
    out_path = os.path.join(out_dir, "model_fp32.onnx")

    print(f"[onnx] exporting -> {out_path} (opset 17, dynamic axes)")
    torch.onnx.export(
        wrapper,
        dummy,
        out_path,
        opset_version=17,
        input_names=["input_ids"],
        output_names=["logits"],
        dynamic_axes={
            "input_ids": {0: "batch", 1: "sequence"},
            "logits":    {0: "batch", 1: "sequence"},
        },
        do_constant_folding=True,
    )
    size_mb = os.path.getsize(out_path) / (1024 * 1024)
    print(f"[onnx] FP32 size: {size_mb:.1f} MB")
    return out_path


def quantize_onnx(onnx_path: str, bits: int, out_dir: str) -> str:
    """Quantize ONNX weights.

    - bits=4: MatMulNBits quantizer (recommended; ORT 1.17+). Output uses the
              `MatMulNBits` operator natively supported by ONNX Runtime Web.
    - bits=8: dynamic QInt8 quantization. Smaller quality hit but ~2x larger
              file (close to the 50 MB ceiling).
    """
    out_path = os.path.join(out_dir, f"model_q{bits}.onnx")

    if bits == 4:
        from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer
        from onnxruntime.quantization.shape_inference import quant_pre_process

        # Pre-process: shape inference + optimization (improves quantization quality)
        preproc_path = onnx_path.replace(".onnx", "_preproc.onnx")
        try:
            quant_pre_process(
                input_path=onnx_path,
                output_path=preproc_path,
                skip_symbolic_shape=False,
                skip_onnx_shape=False,
                skip_optimization=False,
            )
            input_for_quant = preproc_path
        except Exception as e:
            print(f"[quant] pre-process failed ({e}); falling back to raw model")
            input_for_quant = onnx_path

        print(f"[quant] 4-bit via MatMulNBits (symmetric, accuracy_level=4)")
        try:
            quantizer = MatMulNBitsQuantizer(
                model_input=input_for_quant,
                model_output=out_path,
                bits=4,
                use_symmetric=True,
                accuracy_level=4,
            )
        except TypeError:
            # Older ORT API without accuracy_level
            quantizer = MatMulNBitsQuantizer(
                model_input=input_for_quant,
                model_output=out_path,
                bits=4,
                use_symmetric=True,
            )
        quantizer.process()

        if os.path.exists(preproc_path):
            os.remove(preproc_path)

    elif bits == 8:
        from onnxruntime.quantization.quantize import quantize_dynamic, QuantType
        print(f"[quant] 8-bit via dynamic QInt8 (per-channel)")
        quantize_dynamic(
            model_input=onnx_path,
            model_output=out_path,
            weight_type=QuantType.QInt8,
            op_types_to_quantize=["MatMul", "Gemm"],
            per_channel=True,
            reduce_range=False,
        )
    else:
        raise ValueError(f"Unsupported bits={bits}. Use 4 or 8.")

    size_mb = os.path.getsize(out_path) / (1024 * 1024)
    print(f"[quant] {bits}-bit ONNX size: {size_mb:.2f} MB")
    return out_path


# ============================================================================
# Main
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Train a ~50M GPT and export a quantized ONNX model."
    )
    parser.add_argument("--dataset",    default="/kaggle/input/dataset.txt",
                        help="Path to plain-text dataset (will also try ./dataset.txt)")
    parser.add_argument("--out_dir",    default="/kaggle/working",
                        help="Output directory for artifacts")
    # Training overrides
    parser.add_argument("--max_iters",  type=int, default=5000)
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--grad_accum", type=int, default=4)
    parser.add_argument("--block_size", type=int, default=256)
    parser.add_argument("--vocab_size", type=int, default=8000)
    parser.add_argument("--n_layer",    type=int, default=12)
    parser.add_argument("--n_embd",     type=int, default=512)
    parser.add_argument("--n_head",     type=int, default=8)
    parser.add_argument("--lr",         type=float, default=3e-4)
    parser.add_argument("--seed",       type=int, default=1337)
    parser.add_argument("--quant_bits", type=int, default=4, choices=[4, 8])
    parser.add_argument("--skip_train", action="store_true",
                        help="Skip training; load model_fp32.pt and re-export/quantize")
    parser.add_argument("--force_cpu", action="store_true",
                        help="Force CPU mode even if CUDA is available (useful for P100 or "
                             "other GPUs unsupported by the installed PyTorch)")
    # `parse_known_args` silently ignores unrecognized CLI args. This matters
    # in Kaggle/Colab, where Jupyter injects `-f /root/.local/.../kernel-*.json`
    # into sys.argv — plain `parse_args()` would crash with SystemExit(2).
    args, _unknown = parser.parse_known_args()

    # ---- Resolve dataset path ----
    # Kaggle mounts attached datasets at /kaggle/input/<dataset-slug>/ — NOT
    # /kaggle/input/datasets/<user>/<slug>/. The slug is usually the dataset
    # title lowercased with dashes/spaces removed. If the user's --dataset
    # path doesn't exist, we:
    #   (a) try ./dataset.txt,
    #   (b) search /kaggle/input, /content, and cwd for any *.txt file,
    #   (c) if nothing matches, list every file under /kaggle/input so the
    #       user can see exactly what's mounted and copy the right path.
    candidate_paths = [args.dataset, os.path.join(os.getcwd(), "dataset.txt")]
    found = next((p for p in candidate_paths if p and os.path.exists(p)), None)

    # Acceptable filenames when auto-discovering (case-insensitive)
    ACCEPTABLE_NAMES = {"dataset.txt", "dataset1.txt", "corpus.txt", "data.txt", "train.txt"}
    ACCEPTABLE_SUFFIXES = (".txt", ".md", ".csv")

    if not found:
        search_roots = [d for d in ("/kaggle/input", "/content", "/content/data", ".")
                        if os.path.isdir(d)]
        for root in search_roots:
            for dirpath, _dirs, files in os.walk(root):
                # Prefer files with known names first
                for fname in files:
                    if fname.lower() in ACCEPTABLE_NAMES:
                        found = os.path.join(dirpath, fname)
                        break
                if found:
                    break
                # Fall back to any .txt file
                for fname in files:
                    if fname.lower().endswith(ACCEPTABLE_SUFFIXES):
                        found = os.path.join(dirpath, fname)
                        break
                if found:
                    break
            if found:
                break

    if not found:
        # Build a helpful listing of everything actually mounted under /kaggle/input
        listing = []
        for root in ("/kaggle/input", "/content"):
            if not os.path.isdir(root):
                continue
            listing.append(f"\n  Contents of {root}:")
            for dirpath, dirs, files in os.walk(root):
                rel = os.path.relpath(dirpath, root)
                for f in files:
                    full = os.path.join(dirpath, f)
                    try:
                        size_mb = os.path.getsize(full) / (1024 * 1024)
                    except OSError:
                        size_mb = 0
                    listing.append(f"    {full}  ({size_mb:.2f} MB)")
            if not any(os.scandir(root)):
                listing.append(f"    (empty — nothing mounted under {root})")
        listing_str = "\n".join(listing) if listing else "  (no /kaggle/input or /content directory found)"

        raise FileNotFoundError(
            f"No text dataset found.\n"
            f"  Tried explicit path : {args.dataset}\n"
            f"  Tried cwd fallback  : {os.path.join(os.getcwd(), 'dataset.txt')}\n"
            f"  Searched under      : {search_roots}\n"
            f"  Acceptable filenames: {sorted(ACCEPTABLE_NAMES)} (or any *.txt)\n"
            f"\nWhat's actually on disk:{listing_str}\n"
            f"\nFix:\n"
            f"  Copy the full path of your dataset file from the listing above and run:\n"
            f"    !python train.py --dataset '/kaggle/input/<exact-path-from-listing>'\n"
            f"\n  If nothing is listed, you haven't attached a Kaggle Dataset yet.\n"
            f"  In the Kaggle sidebar (right) → 'Add Data' → search your dataset → add it."
        )
    args.dataset = found
    print(f"[data] using dataset: {args.dataset}")

    os.makedirs(args.out_dir, exist_ok=True)

    cfg = Config(
        vocab_size=args.vocab_size,
        n_layer=args.n_layer,
        n_head=args.n_head,
        n_embd=args.n_embd,
        block_size=args.block_size,
        batch_size=args.batch_size,
        grad_accum=args.grad_accum,
        learning_rate=args.lr,
        max_iters=args.max_iters,
        seed=args.seed,
    )

    torch.manual_seed(cfg.seed)
    np.random.seed(cfg.seed)

    # ---- Device selection with GPU compatibility check ----
    # Some GPUs (notably Tesla P100 = sm_60, Kepler = sm_35/37) are no longer
    # supported by modern PyTorch. PyTorch will print warnings but then crash
    # on the first CUDA op. We detect this proactively and either auto-fall-back
    # to CPU or honor --force_cpu.
    device = "cpu"
    gpu_info = ""
    if not args.force_cpu and torch.cuda.is_available():
        try:
            cap = torch.cuda.get_device_capability(0)  # (major, minor) e.g. (6, 0) for P100
            gpu_name = torch.cuda.get_device_name(0)
            gpu_info = f"  ({gpu_name}, sm_{cap[0]}{cap[1]})"
            # PyTorch's compiled-in minimum is sm_70 (Volta). Anything below
            # will produce wrong results or crash at runtime.
            if cap[0] < 7:
                print("=" * 70)
                print("WARNING: GPU compatibility issue detected")
                print(f"  GPU: {gpu_name} (compute capability {cap[0]}.{cap[1]})")
                print(f"  This PyTorch build requires sm_70+ (Volta/Turing/Ampere/...).")
                print(f"  Running on this GPU will crash or produce garbage.")
                print(f"  Options:")
                print(f"    1. Switch the Kaggle accelerator to 'GPU T4 x1' (recommended)")
                print(f"       Settings → Accelerator → GPU T4 x1")
                print(f"    2. Force CPU mode (much slower but works):")
                print(f"       !python train.py --force_cpu --out_dir /kaggle/working")
                print(f"  Auto-falling back to CPU for now.")
                print("=" * 70)
                device = "cpu"
            else:
                device = "cuda"
        except Exception as e:
            print(f"[device] CUDA check failed ({e}); falling back to CPU.")
            device = "cpu"

    if args.force_cpu and torch.cuda.is_available():
        print("[device] --force_cpu set; using CPU despite CUDA being available.")

    print("=" * 70)
    print(f"Device : {device}" + (f"  ({torch.cuda.get_device_name(0)})" if device == "cuda" else gpu_info))
    print(f"Config : {asdict(cfg)}")
    print("=" * 70)

    # 1. Tokenizer ------------------------------------------------------------
    tok_path = os.path.join(args.out_dir, "tokenizer.json")
    if os.path.exists(tok_path):
        print("[1/6] Loading existing tokenizer...")
        tokenizer = load_tokenizer(args.out_dir)
    else:
        print("[1/6] Training tokenizer...")
        tokenizer = train_tokenizer(args.dataset, cfg.vocab_size, args.out_dir)
    cfg.vocab_size = tokenizer.get_vocab_size()  # actual size may differ slightly
    print(f"      vocab_size = {cfg.vocab_size}")

    # 2. Tokenize dataset -----------------------------------------------------
    train_bin = os.path.join(args.out_dir, "train.bin")
    val_bin   = os.path.join(args.out_dir, "val.bin")
    if not (os.path.exists(train_bin) and os.path.exists(val_bin)):
        print("[2/6] Tokenizing dataset...")
        all_bin = os.path.join(args.out_dir, "all.bin")
        encode_dataset(tokenizer, args.dataset, all_bin)
        all_data = load_binary(all_bin)
        train_data, val_data = split_train_val(all_data, cfg.val_split)
        np.asarray(train_data).tofile(train_bin)
        np.asarray(val_data).tofile(val_bin)
        os.remove(all_bin)
    else:
        print("[2/6] Loading pre-tokenized shards...")
    train_data = load_binary(train_bin)
    val_data   = load_binary(val_bin)
    print(f"      train tokens = {len(train_data):,}    val tokens = {len(val_data):,}")

    # 3. Build model ----------------------------------------------------------
    print("[3/6] Building model...")
    model = GPT(cfg).to(device)
    n_params = model.num_parameters()
    print(f"      params = {n_params:,}  (~{n_params/1e6:.2f}M)")

    # 4. Train ----------------------------------------------------------------
    if args.skip_train:
        ckpt = os.path.join(args.out_dir, "model_fp32.pt")
        if not os.path.exists(ckpt):
            ckpt = os.path.join(args.out_dir, "model_best.pt")
        print(f"[4/6] skip_train=True; loading {ckpt}")
        model.load_state_dict(torch.load(ckpt, map_location=device)["model_state"])
    else:
        print("[4/6] Training...")
        train_model(model, train_data, val_data, cfg, device, args.out_dir)

    # 5. Export ONNX ----------------------------------------------------------
    print("[5/6] Exporting ONNX...")
    onnx_fp32 = export_onnx(model, cfg, args.out_dir)

    # 6. Quantize -------------------------------------------------------------
    print(f"[6/6] Quantizing to {args.quant_bits}-bit...")
    quant_path = quantize_onnx(onnx_fp32, args.quant_bits, args.out_dir)

    # ---- Final report -------------------------------------------------------
    final_mb = os.path.getsize(quant_path) / (1024 * 1024)
    print("=" * 70)
    print(f"Final model : {quant_path}")
    print(f"Final size  : {final_mb:.2f} MB")
    if final_mb < 50:
        print(f"OK          : Under 50 MB upload limit (margin {50 - final_mb:.2f} MB)")
    else:
        print(f"WARN        : Over 50 MB by {final_mb - 50:.2f} MB")
        print("              Mitigation: lower --vocab_size / --n_layer, or use --quant_bits 4")
    print(f"Tokenizer   : {os.path.join(args.out_dir, 'tokenizer.json')}")
    print(f"Deploy with : ONNX Runtime Web  (npm: onnxruntime-web)")
    print("=" * 70)


if __name__ == "__main__":
    main()