"""Source-byte, strict-state, native-feature/token and codec-path verification."""
import argparse,bisect,gc,hashlib,json,pathlib,struct,sys,time,io
import torch
import numpy as np
from PIL import Image
from component import build_model,load_model,tokenizer_view
class SourceRanges:
 def __init__(self,root,mapping):
  self.shards={};self.root=root;self.dcp=any('chunks' in e for e in mapping.values())
  if self.dcp:return
  for shard in {e['source_shard'] for e in mapping.values()}:
   parts=sorted((int(p.name.rsplit('_',1)[1]),p) for p in (root/'.download').glob(shard.replace('/','__')+'_*') if not p.name.endswith('.part'))
   assert parts
   self.shards[shard]=(parts,parts[0][0]-min(e['data_offsets'][0] for e in mapping.values() if e['source_shard']==shard))
 def read(self,e):
  if self.dcp:
   t=torch.empty(e['shape'],dtype=DTYPES[e['dtype']]);coverage=torch.zeros(e['shape'][0],dtype=torch.bool)
   for c in e['chunks']:
    raw=(self.root/'.download'/c['cache_file']).read_bytes();assert hashlib.sha256(raw).hexdigest()==c['serialized_sha256']
    part=torch.load(io.BytesIO(raw),map_location='cpu',weights_only=True);assert list(part.shape)==c['tensor_sizes'] and part.dtype==t.dtype
    a=c['tensor_offsets'][0];b=a+c['tensor_sizes'][0];assert not coverage[a:b].any();coverage[a:b]=True
    assert c['tensor_offsets'][1:]==[0]*(t.ndim-1) and c['tensor_sizes'][1:]==list(t.shape)[1:]
    t[tuple(slice(a,a+b) for a,b in zip(c['tensor_offsets'],c['tensor_sizes']))]=part
   assert coverage.all();return bytearray(t.view(torch.uint8).numpy().tobytes())
  parts,base=self.shards[e['source_shard']];starts=[p[0] for p in parts];a,b=e['data_offsets'];a+=base;b+=base;result=bytearray()
  while a<b:
   start,p=parts[bisect.bisect_right(starts,a)-1]
   with p.open('rb') as f:
    f.seek(a-start);data=f.read(min(b-a,p.stat().st_size-a+start))
   if not data:raise RuntimeError('Source gap')
   result.extend(data);a+=len(data)
  return result
DTYPES={'BF16':torch.bfloat16,'F32':torch.float32,'F16':torch.float16,'I64':torch.int64,'I32':torch.int32,'BOOL':torch.bool}
def tree(x):
 if isinstance(x,torch.Tensor):return x.detach().cpu().clone()
 if isinstance(x,dict):return {k:tree(v) for k,v in x.items()}
 if isinstance(x,(list,tuple)):return [tree(v) for v in x]
 return x

def compare(a,b):
 if isinstance(a,torch.Tensor):
  assert a.shape==b.shape
  if not a.is_floating_point():assert torch.equal(a,b);return 0.
  assert torch.isfinite(a).all() and torch.isfinite(b).all()
  torch.testing.assert_close(a,b,rtol=2e-4,atol=2e-4)
  return (a.float()-b.float()).abs().max().item()
 if isinstance(a,dict):return max([compare(a[k],b[k]) for k in a] or [0.])
 if isinstance(a,(list,tuple)):
  assert len(a)==len(b);return max([compare(x,y) for x,y in zip(a,b)] or [0.])
 assert a==b;return 0.

def image(h,w):
 y,x=np.mgrid[:h,:w];return Image.fromarray(np.stack([(x*7+y*3)%256,(x*2+y*11)%256,(x*13+y*5)%256],-1).astype('uint8'))
def run_case(m,root,case):
 kind=json.loads((root/'config.json').read_text()).get('component',root.name)
 if kind=='pixelumm_pixel_io':
  px,tubes,hs,t=case['input'];return {'image_und':m.pixel_patch_embed(px),'image_gen':m.pixel_patch_embed_mot_gen(px),'video_gen':m.pixel_video_patch_embed_mot_gen(tubes),'image_head':m.unpatchify_head(hs),'video_head':m.video_unpatchify_head(hs),'time_mlp_inactive_in_release':m.time_embedder(t)}
 if kind=='sensenova_pixel_io':
  px,g,hs,t=case['input'];return {'understanding':m.encode_pixels(px,g),'generation':m.encode_pixels(px,g,True),'pixel_head':m.decode_hidden(hs,(2,3)),'time':m.fm_modules['timestep_embedder'](t),'noise':m.fm_modules['noise_scale_embedder'](t)}
 if kind=='mingtok':
  x=case['input'];o=m(x);pix=m.forward_pixel_decoder(o['x_norm_patchtokens']);sem=m.forward_feature_decoder_wo_cache(o['latent']*m.scaling_factor+m.mean)['x_norm_patchtokens']
  return {'latent':o['latent'],'semantic':o['x_norm_patchtokens'],'reconstruction':pix,'semantic_from_latent':sem}
 if kind=='uniar_bsq':
  px,g=case['input'];bits=m(px,grid_thw=g,bsq_only=True,direct_indices=True)
  joined=torch.cat([bits[0]]+bits[1],dim=-1)
  assert torch.all((joined==0)|(joined==1))
  features=m.bsq_indices_to_features(joined,multiscale_idx='all',with_merger=True)
  direct=m(px,grid_thw=g)
  compare(features,direct)
  return {'bits':joined,'decoded_features':features,'direct_features':direct}
 t=tokenizer_view(m,root);ims=case['input'];inp=t.image_processor(ims);hs=t.visual(inp['pixel_values'],inp['image_grid_thw']);ids=t._quantize(hs,inp['image_grid_thw'])
 actual=t.encode_batch(ims)
 assert actual==[x.tolist() for x in ids.split(inp['image_grid_thw'].prod(-1).tolist())]
 assert ((ids>=0)&(ids<t.codebook_size)).all()
 # Retain inverse VQ embedding + post-quant projection; no pixel renderer is claimed.
 codes=torch.nn.functional.normalize(t.vqmodel.quantize.embedding.weight,dim=-1)[ids]
 inverse=torch.nn.functional.linear(codes,t.vqmodel.post_quant_conv.weight[:,:,0,0],t.vqmodel.post_quant_conv.bias)
 return {'hidden':hs,'ids':ids,'inverse_vq_features':inverse,'api_token_ids':actual}

def main(root):
 torch.set_num_threads(1);torch.manual_seed(17);start=time.time();mapping=json.loads((root/'tensor_mapping.json').read_text());ranges=SourceRanges(root,mapping)
 with (root/'model.safetensors').open('rb') as f:
  n=struct.unpack('<Q',f.read(8))[0];header=json.loads(f.read(n));base=8+n
  for k,e in mapping.items():
   h=header[k];assert h['shape']==e['shape'] and h['dtype']==e['dtype'];a,b=h['data_offsets'];f.seek(base+a);assert f.read(b-a)==ranges.read(e),k
 print('ALL SOURCE BYTES MATCH',flush=True)
 # Original native classes, independent direct source-range loading, eager reference attention.
 ref=build_model(root,attention='eager');sd=ref.state_dict();assert set(sd)==set(mapping),(set(sd)-set(mapping),set(mapping)-set(sd))
 with torch.no_grad():
  for k,e in mapping.items():sd[k].copy_(torch.frombuffer(ranges.read(e),dtype=DTYPES[e['dtype']]).reshape(e['shape']))
 del sd;gc.collect();kind=json.loads((root/'config.json').read_text()).get('component',root.name)
 if kind=='pixelumm_pixel_io':cases=[{'name':name,'input':(torch.randn(n,768),torch.randn(n,3072),torch.randn(n,4096),torch.tensor([0.,.5,1.]))} for name,n in [('single_patch',1),('six_patches',6),('multi_frame_tubes',24)]]
 elif kind=='sensenova_pixel_io':
  cases=[]
  for name,grid in [('square',[[4,4]]),('rectangle',[[4,6]]),('packed_two_images',[[4,4],[4,6]])]:
   g=torch.tensor(grid);cases.append({'name':name,'input':(torch.randn(int(g.prod(-1).sum()),768),g,torch.randn(1,6,4096),torch.tensor([0.,.5,1.]))})
 elif kind=='mingtok':cases=[{'name':f'square_{s}','input':torch.linspace(-1,1,3*s*s).reshape(1,3,s,s)} for s in [64,128]]
 elif kind=='uniar_bsq':
  vc=ref.config;cases=[]
  for name,grid in [('square',[[1,4,4]]),('rectangle',[[1,4,6]]),('packed_two_images',[[1,4,4],[1,4,6]])]:
   g=torch.tensor(grid);px=torch.randn(int(g.prod(-1).sum()),3*vc.temporal_patch_size*vc.patch_size**2,dtype=torch.bfloat16);cases.append({'name':name,'input':(px,g)})
 else:cases=[{'name':'square_rgb','input':[image(64,64)]},{'name':'rectangle_rgb','input':[image(64,96)]},{'name':'packed_two_images','input':[image(64,64),image(64,96)]}]
 results=[]
 with torch.inference_mode():
  for c in cases:results.append(tree(run_case(ref,root,c)));print('reference '+c['name'],flush=True)
 del ref;gc.collect();exported=load_model(root);checks=[]
 with torch.inference_mode():
  for c,r in zip(cases,results):
   o=tree(run_case(exported,root,c));err=compare(o,r);checks.append({'case':c['name'],'max_abs_error':err,'passed':True});print('PASS '+c['name']+' '+str(err),flush=True)
 # A useful-size native input smoke, separate from equivalence cases.
 if kind=='pixelumm_pixel_io':smoke={'name':'256_image_and_video_tube_io','input':(torch.randn(256,768),torch.randn(256,3072),torch.randn(256,4096),torch.tensor([.5]))}
 elif kind=='sensenova_pixel_io':smoke={'name':'256_native_pixel_io','input':(torch.randn(256,768),torch.tensor([[16,16]]),torch.randn(1,6,4096),torch.tensor([.5]))}
 elif kind=='mingtok':smoke={'name':'512_pixel_reconstruction','input':torch.from_numpy(np.asarray(image(512,512)).copy()).permute(2,0,1)[None].float()/127.5-1}
 elif kind=='llada_image_tokenizer':smoke={'name':'256_rgb_tokenization','input':[image(256,256)]}
 else:
  g=torch.tensor([[1,28,28]]);smoke={'name':'448_native_patch_encoding','input':(torch.randn(784,1536,dtype=torch.bfloat16),g)}
 with torch.inference_mode():out=tree(run_case(exported,root,smoke))
 compare(out,out)  # Recursive finite check, including smoke outputs.
 shapes={k:list(v.shape) for k,v in out.items() if isinstance(v,torch.Tensor)}
 if kind=='mingtok':
  arr=((out['reconstruction'][0].permute(1,2,0)+1)*127.5).round().clamp(0,255).byte().numpy();Image.fromarray(arr).save(root/'reconstruction_smoke.png');image(512,512).save(root/'input_smoke.png')
 report={'source_tensor_bytes_equal':True,'strict_load':True,'reference':'Pinned upstream native classes with all original source-range tensors loaded directly. Exported explicit loader uses the assembled checkpoint. Eager reference vs SDPA export where supported.','checks':checks,'smoke':{'case':smoke['name'],'shapes':shapes,'passed':True},'dtype':'BF16' if kind=='uniar_bsq' else 'FP32','torch':torch.__version__,'transformers':__import__('transformers').__version__,'tolerance':{'rtol':2e-4,'atol':2e-4},'seed':17,'elapsed_seconds':round(time.time()-start,2),'scope':'Component encoding/quantization/features/reconstruction as available; no language backbone, full text-to-image generation or downstream accuracy evaluation.'}
 (root/'validation.json').write_text(json.dumps(report,indent=2)+'\n');print('VALIDATED '+root.name,flush=True)
if __name__=='__main__':
 p=argparse.ArgumentParser();p.add_argument('root',type=pathlib.Path);main(p.parse_args().root)