from __future__ import annotations

import argparse
import time

import torch
import torch.nn.functional as F

from audio_dit import AudioDiT

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--batch-sizes", type=int, nargs="+", default=[64, 128, 256])
    ap.add_argument("--x-res", type=int, default=384)
    ap.add_argument("--y-res", type=int, default=256)
    ap.add_argument("--warmup", type=int, default=10)
    ap.add_argument("--iters", type=int, default=40)
    args = ap.parse_args()
    dev = "cuda"
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    model = AudioDiT(x_res=args.x_res, y_res=args.y_res).to(dev)
    opt = torch.optim.AdamW(model.parameters(), lr=2e-4)
    print(f"[bench] {model.num_params()/1e6:.2f}M params, grid {model.grid_h}x{model.grid_w} "
          f"({model.grid_h*model.grid_w} tokens) at x_res={args.x_res} y_res={args.y_res}", flush=True)

    for B in args.batch_sizes:
        try:
            torch.cuda.reset_peak_memory_stats()
            x1 = torch.randn(B, 1, args.y_res, args.x_res, device=dev)
            seq = torch.randn(B, 32, 768, device=dev)
            pool = torch.randn(B, 512, device=dev)

            def step():
                x0 = torch.randn_like(x1)
                t = torch.sigmoid(torch.randn(B, device=dev))
                tb = t.view(-1, 1, 1, 1)
                xt = (1 - tb) * x0 + tb * x1
                target = x1 - x0
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    v = model(xt, t, seq, pool)
                    loss = F.mse_loss(v.float(), target)
                opt.zero_grad(set_to_none=True)
                loss.backward()
                opt.step()

            for _ in range(args.warmup):
                step()
            torch.cuda.synchronize()
            t0 = time.time()
            for _ in range(args.iters):
                step()
            torch.cuda.synchronize()
            el = time.time() - t0
            sps = args.iters / el
            peak = torch.cuda.max_memory_allocated() / 2**30
            print(f"[bench] batch={B:<5} {sps:6.2f} steps/s  {sps*B:8.1f} samples/s  "
                  f"peak {peak:5.1f} GiB  90k-step ETA {90000/sps/3600:5.2f} h", flush=True)
        except torch.cuda.OutOfMemoryError:
            print(f"[bench] batch={B:<5} OOM", flush=True)
            torch.cuda.empty_cache()

if __name__ == "__main__":
    main()