"""
Day-1 throughput test: forward+backward tokens/sec for a preset on this machine.
Use it to size the token budget before any real run.

    python bench.py --config hoard_small --seq_len 1024 --batch 16 --iters 5
"""
import argparse, time
import mlx.core as mx
import mlx.nn as nn
import configs
from model import HOARD


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="hoard_small")
    ap.add_argument("--seq_len", type=int, default=1024)
    ap.add_argument("--batch", type=int, default=16)
    ap.add_argument("--iters", type=int, default=5)
    ap.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"])
    a = ap.parse_args()

    cfg = configs.get(a.config)
    model = HOARD(cfg)
    if a.dtype == "bfloat16":
        model.set_dtype(mx.bfloat16)
    c = model.param_counts()
    print(f"{a.config}: total {c['total']/1e6:.1f}M, active/token/loop {c['active_per_token_per_loop']/1e6:.1f}M, loops {cfg.n_loops}")

    def loss_fn(model, x):
        logits = model(x[:, :-1])
        return nn.losses.cross_entropy(logits.astype(mx.float32), x[:, 1:], reduction="mean")

    grad_fn = nn.value_and_grad(model, loss_fn)
    x = mx.random.randint(0, cfg.vocab_size, (a.batch, a.seq_len + 1))
    # warmup
    l, g = grad_fn(model, x); mx.eval(l, g)
    t0 = time.time()
    for _ in range(a.iters):
        l, g = grad_fn(model, x); mx.eval(l, g)
    dt = (time.time() - t0) / a.iters
    toks = a.batch * a.seq_len
    print(f"fwd+bwd: {dt*1000:.0f} ms/step, {toks/dt:,.0f} tok/s  "
          f"-> 1B tokens ~ {1e9/(toks/dt)/3600:.1f} h on this machine")
    print(f"peak memory: {mx.get_peak_memory()/2**30:.2f} GB")


if __name__ == "__main__":
    main()