#!/usr/bin/env python
"""Reproducible GPU/CPU scaling benchmark for Hilbert Operator Intelligence."""
from __future__ import annotations
import argparse, json, os, random, time
from pathlib import Path
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, Subset, TensorDataset
ROOT=Path(__file__).resolve().parents[1]
import sys; sys.path.insert(0,str(ROOT))
from hilbert_mi.scalable import build_model

def seed_all(s): random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s)
def synthetic(n, classes, size, seed):
    g=torch.Generator().manual_seed(seed); y=torch.randint(classes,(n,),generator=g); x=torch.randn(n,3,size,size,generator=g)*.15
    # deterministic class-dependent spectral/spatial signal
    for i in range(n):
        k=int(y[i]); x[i,:, (k*3)%size::10, :] += .6; x[i,:, :, (k*5)%size::11] += .35
    return TensorDataset(x,y)
def datasets(name, root, ntrain, ntest, size, seed):
    if name=="synthetic": return synthetic(ntrain,10,size,seed), synthetic(ntest,10,size,seed+1),10
    try:
        import torchvision.transforms as T; import torchvision.datasets as D
    except ImportError as e: raise SystemExit("Install benchmark extras: pip install -e '.[benchmark]'") from e
    tf=T.Compose([T.Resize((size,size)),T.ToTensor()])
    cls={"cifar10":D.CIFAR10,"cifar100":D.CIFAR100}[name]; nc=10 if name=="cifar10" else 100
    tr=cls(root=root,train=True,download=True,transform=tf); te=cls(root=root,train=False,download=True,transform=tf)
    if ntrain and ntrain<len(tr): tr=Subset(tr,range(ntrain))
    if ntest and ntest<len(te): te=Subset(te,range(ntest))
    return tr,te,nc
def evaluate(m,loader,dev):
    m.eval(); ok=n=0; loss=0.; ce=nn.CrossEntropyLoss(reduction="sum")
    with torch.no_grad():
      for x,y in loader:
        x,y=x.to(dev),y.to(dev); z=m(x); loss+=ce(z,y).item(); ok+=(z.argmax(1)==y).sum().item(); n+=len(y)
    return {"accuracy":ok/n,"nll":loss/n}
def main():
    p=argparse.ArgumentParser(); p.add_argument("--config",required=True); p.add_argument("--dataset",default="cifar10",choices=["cifar10","cifar100","synthetic"]); p.add_argument("--epochs",type=int,default=None); p.add_argument("--batch-size",type=int,default=None); p.add_argument("--train-samples",type=int,default=0); p.add_argument("--test-samples",type=int,default=0); p.add_argument("--image-size",type=int,default=32); p.add_argument("--seed",type=int,default=7); p.add_argument("--output",default="results/run.json"); p.add_argument("--amp",action="store_true"); a=p.parse_args()
    cfg=json.load(open(a.config)); seed_all(a.seed); dev=torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tr,te,nc=datasets(a.dataset,str(ROOT/"data"),a.train_samples,a.test_samples,a.image_size,a.seed)
    bs=a.batch_size or cfg.get("batch_size",128); epochs=a.epochs or cfg.get("epochs",20)
    dl=DataLoader(tr,batch_size=bs,shuffle=True,num_workers=cfg.get("workers",2),pin_memory=dev.type=="cuda"); tl=DataLoader(te,batch_size=bs,shuffle=False,num_workers=cfg.get("workers",2),pin_memory=dev.type=="cuda")
    model=build_model(cfg,nc).to(dev); opt=torch.optim.AdamW(model.parameters(),lr=cfg.get("lr",3e-4),weight_decay=cfg.get("weight_decay",1e-4)); ce=nn.CrossEntropyLoss(); scaler=torch.amp.GradScaler("cuda",enabled=a.amp and dev.type=="cuda")
    t0=time.perf_counter(); seen=0
    for ep in range(epochs):
      model.train()
      for x,y in dl:
        x,y=x.to(dev,non_blocking=True),y.to(dev,non_blocking=True); opt.zero_grad(set_to_none=True)
        with torch.autocast(device_type=dev.type,enabled=a.amp and dev.type=="cuda"):
          loss=ce(model(x),y)
        scaler.scale(loss).backward(); scaler.step(opt); scaler.update(); seen+=len(y)
    if dev.type=="cuda": torch.cuda.synchronize()
    sec=time.perf_counter()-t0; met=evaluate(model,tl,dev)
    result={"benchmark":"hoi-image-v1","dataset":a.dataset,"config":Path(a.config).stem,"seed":a.seed,"epochs":epochs,"train_samples":len(tr),"test_samples":len(te),"image_size":a.image_size,"params":model.param_count,"device":str(dev),"gpu":torch.cuda.get_device_name(0) if dev.type=="cuda" else None,"torch":torch.__version__,"amp":bool(a.amp and dev.type=="cuda"),"train_seconds":sec,"samples_per_second":seen/sec,"peak_gpu_memory_gb":torch.cuda.max_memory_allocated()/1e9 if dev.type=="cuda" else 0.0,**met}
    out=ROOT/a.output; out.parent.mkdir(parents=True,exist_ok=True); out.write_text(json.dumps(result,indent=2)+"\n"); print(json.dumps(result,indent=2))
if __name__=="__main__": main()