"""Hardware measurement only; random input is NOT accuracy evaluation."""
import argparse,json,time,statistics
from pathlib import Path
import torch
from irstd.models import TinyUNet
from irstd.core import candidate_pool

def main():
 p=argparse.ArgumentParser();p.add_argument('--out',required=True);p.add_argument('--size',type=int,default=256)
 p.add_argument('--device',default='cuda');p.add_argument('--iterations',type=int,default=100);a=p.parse_args()
 if a.iterations<1 or a.size<2:p.error('Invalid benchmark dimensions')
 model=TinyUNet().to(a.device).eval();x=torch.randn(1,1,a.size,a.size,device=a.device)
 cuda=str(a.device).startswith('cuda')
 def sync():
  if cuda:torch.cuda.synchronize()
 if cuda:torch.cuda.reset_peak_memory_stats()
 times=[];full=[]
 with torch.no_grad():
  for i in range(20+a.iterations):
   sync();start=time.perf_counter();y=model(x).sigmoid();sync();middle=time.perf_counter()
   candidate_pool(y[0,0].cpu().numpy());end=time.perf_counter()
   if i>=20:times.append((middle-start)*1000);full.append((end-start)*1000)
 # Exact conv MAC estimate via hooks, measured for declared resolution.
 macs=[0]
 def hook(m,inp,out):
  macs[0]+=out.numel()*(m.in_channels//m.groups)*m.kernel_size[0]*m.kernel_size[1]
 handles=[m.register_forward_hook(hook) for m in model.modules() if isinstance(m,torch.nn.Conv2d)]
 with torch.no_grad():model(x)
 for h in handles:h.remove()
 Path(a.out).write_text(json.dumps(dict(device=a.device,gpu=torch.cuda.get_device_name() if cuda else None,
  input=[1,1,a.size,a.size],dtype='float32',warmup=20,iterations=a.iterations,
  parameters=sum(p.numel() for p in model.parameters()),conv_MACs=macs[0],conv_FLOPs_2_per_MAC=2*macs[0],
  model_latency_median_ms=statistics.median(times),end_to_end_median_ms=statistics.median(full),
  peak_allocated_bytes=torch.cuda.max_memory_allocated() if cuda else None,
  excludes='Disk IO, candidate MLP, graph solver; FLOPs exclude interpolation/activations/pooling',
  accuracy='not measured; random input'),indent=2))
if __name__=='__main__':main()