"""Calibrate from original BF16 and measure actual native W4A4 + BF16 correction."""
import os,json,time,math
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import load_file
from flashinfer import nvfp4_quantize,mm_fp4
from flashinfer.gemm import nvfp4_quantize_smooth,mm_nvfp4_svdquant
BASE=Path(os.environ['IMAGE21_BASE'])/'transformer'
OLD=Path(os.environ['IMAGE21_WORK_DIR'])
INDEX=json.loads((BASE/'diffusion_pytorch_model.safetensors.index.json').read_text())['weight_map']
def weight(name):
 with safe_open(BASE/INDEX[name+'.weight'],framework='pt',device='cuda') as f:return f.get_tensor(name+'.weight')
def split(x):
 g=torch.arange(len(x)//4,device=x.device)
 return x[g*4+g%4].contiguous(),x[g*4+(g+1)%4].contiguous()
def error(y,ref):return ((y.float()-ref.float()).square().mean()/ref.float().square().mean()).sqrt().item()
def prepare(w,x,smoothing,rank,method="svd",mse=False):
 amax=x.float().abs().amax(0).clamp_min(1e-5);wmax=w.float().abs().amax(0).clamp_min(1e-5)
 smooth=torch.ones_like(amax) if smoothing==0 else (amax**smoothing/wmax**(1-smoothing)).clamp(.01,100)
 smooth/=smooth.log().mean().exp()
 pre=(1/smooth).to(torch.bfloat16)
 ws=w.float()/pre.float()[None,:]
 if rank:
  torch.manual_seed(1234)
  if method=='svd':
   u,s,v=torch.svd_lowrank(ws,q=rank+32,niter=3)
   up=u[:,:rank]*s[:rank];down=v[:,:rank]
  elif method=='output':
   u,s,v=torch.svd_lowrank(ws@(x.float()*pre.float()).T,q=rank+32,niter=3)
   up=u[:,:rank];down=ws.T@up
  else:
   u,s,v=torch.svd_lowrank(x.float()*pre.float(),q=rank+32,niter=3)
   down=v[:,:rank];up=ws@down
  residual=ws-up@down.T
 else:residual=ws
 gx=(2688/((x.float()*pre.float()).abs().max()*1.2)).reshape(1)
 gw=(2688/residual.abs().max()).reshape(1)
 if mse:
  from weight_quant import quantize_weight
  h=(x.float()*pre.float()).square().mean(0)
  h=h+.01*h.mean()
  wq,wsc=quantize_weight(residual,gw,h)
 else:wq,wsc=nvfp4_quantize(residual.to(torch.bfloat16),gw,backend='cute-dsl')
 alpha=(1/(gx*gw)).reshape(1)
 result={'weight':wq.view(torch.uint8),'sf':wsc.view(torch.uint8).reshape(-1),'pre':pre,'gx':gx,'alpha':alpha}
 if rank:result.update(down=(pre.float()[:,None]*down).to(torch.bfloat16).contiguous(),up=(up/alpha).to(torch.bfloat16).contiguous())
 return result
def run(x,p):
 from dynamic_scale import scale
 gx,alpha,up=scale(x,p['pre'],p['gx'],p['alpha'],p['up']) if 'up' in p else (p['gx'],p['alpha'],None)
 q,sf=nvfp4_quantize_smooth(x,p['pre'],gx,backend='cute-dsl')
 if 'down' in p:
  d=x@p['down']
  return mm_nvfp4_svdquant(q,p['weight'],sf,p['sf'],alpha,d,up,backend='cute-dsl')
 return mm_fp4(q,p['weight'].T,sf,p['sf'],p['alpha'],backend='b12x')
@torch.inference_mode()
def main():
 rows=[]
 for block,leaf in [(0,'attn.to_q'),(0,'img_mlp.out'),(8,'attn.to_v'),(16,'img_mlp.proj'),(16,'img_mlp.gate_layer'),(24,'attn.to_out.0'),(31,'attn.to_k')]:
  name=f'transformer_blocks.{block}.{leaf}';w=weight(name)
  x=load_file(str(OLD/'calibration'/f'{name}.safetensors'))['activations'].cuda();fit,check=split(x)
  rf=fit@w.T;rc=check@w.T
  for smooth in [0.,.5,1.]:
   for rank in [0,128]:
    start=time.time();p=prepare(w,fit,smooth,rank);y=run(fit,p);yc=run(check,p)
    row={'name':name,'smooth':smooth,'rank':rank,'fit_nrmse':error(y,rf),'check_nrmse':error(yc,rc),'seconds':time.time()-start};rows.append(row);print(json.dumps(row),flush=True)
    Path('probe.json').write_text(json.dumps(rows,indent=2))
if __name__=='__main__':main()