import sys,pathlib,json,torch,time,importlib
key='unicom_vision_compression';root=pathlib.Path(__file__).resolve().parent;sys.path.insert(0,str(root.resolve()));torch.set_num_threads(4);torch.manual_seed(7)
import component as c
checks=[]
def check(name,ok,**extra):
 assert ok,(name,extra)
 checks.append(dict(name=name,passed=True,**extra));print(name,extra,flush=True)
def finite(name,t):check(name,torch.isfinite(t).all().item(),shape=list(t.shape))
t0=time.time();print('LOAD',key,flush=True);m=c.load_model(root);check('native_strict_load',True,tensors=len(m.state_dict()))
with torch.inference_mode():
 if key=='internlumina_u2_visual_io':
  ids=torch.randint(0,4096,(1,3,8));a=c.embed_codes(m,ids);finite('eight_codebook_embedding',a);h=torch.randn(3,2048);b=c.codebook_logits(m,h,ids[0]);finite('eight_codebook_teacher_logits',b);check('logit_shape',b.shape==(3,8,4096))
  # Causality: changing a later code may not influence earlier codebook logits.
  ids2=ids[0].clone();ids2[:,5]=(ids2[:,5]+1)%4096;d=c.codebook_logits(m,h,ids2);err=(b[:,:6]-d[:,:6]).abs().max().item();check('codebook_causality',err<1e-5,max_abs_error=err)
 elif key=='unicom_vision_compression':
  patches=torch.randn(1,16,768);shape=torch.tensor([[4,4]]);a,z=c.unicom_features(m,patches,shape,torch.ones(1,16));finite('siglip_features',a);finite('compressed_64',z);d=m.compression.decompress(z);finite('decompressed_1152',d);check('compression_shapes',a.shape==(1,16,1152) and z.shape==(1,16,64) and d.shape==a.shape)
  t=torch.tensor([0.5]);te=m.time_embed(t);flow=m.patch_embed(z,te)[0];finite('latent_to_backbone',flow);v=m.final_layer(flow,te);finite('backbone_to_latent',v);check('flow_shapes',flow.shape==(1,16,3584) and v.shape==z.shape)
  aligned=c.unicom_align(m,torch.randn(1,16,1216));finite('explicit_1216_aligner',aligned)
  try:c.unicom_align(m,a)
  except ValueError:check('upstream_1152_vs_1216_guard',True)
  else:raise AssertionError('dimension mismatch must be explicit')
 elif key=='internvlu_vision_bridge':
  x=torch.randn(1,3,448,448);a=c.internvlu_features(m,x);finite('vision_pixelshuffle_bridge_448',a);check('merged_token_shape',a.shape==(1,256,2048))
 elif key=='hidream_o1_vision_pixel_io':
  grid=torch.tensor([[1,4,4]]);patches=torch.randn(16,1536);a=m.model.visual(patches,grid_thw=grid);a=a[0] if isinstance(a,tuple) else a;finite('qwen_visual_4x4',a)
  z=torch.randn(1,4,3072);p=m.model.x_embedder(z);t=m.model.t_embedder1(torch.tensor([0.5]));out=m.model.final_layer2(p);finite('pixel_input',p);finite('time_condition',t);finite('pixel_output',out);check('patch32_shapes',p.shape==(1,4,4096) and out.shape==z.shape)
 elif key=='ming_flash_omni2_vision_io':
  for grid in [torch.tensor([[1,4,4]]),torch.tensor([[2,4,4]])]:
   patches=torch.randn(int(grid.prod()),1536);a=c.ming_features(m,patches,grid);finite('vision_bridge_grid_'+str(grid.tolist()),a);check('normalized_features_'+str(grid.tolist()),torch.allclose(a.norm(dim=-1),torch.ones(a.shape[0]),atol=1e-5))
 elif key=='lance_vision':
  for grid in [torch.tensor([[1,4,4]]),torch.tensor([[2,4,4]])]:
   patches=torch.randn(int(grid.prod()),1176);a=m(patches,grid_thw=grid);finite('vit_grid_'+str(grid.tolist()),a);check('merged_tokens_'+str(grid.tolist()),a.shape==(int(grid.prod())//4,2048))
 else:raise ValueError(key)
report=dict(strict_load=True,source_tensor_bytes_equal=json.loads((root/'extraction.json').read_text())['source_tensor_bytes_equal'],checks=checks,smoke=dict(passed=True,device='cpu',dtype='float32'),elapsed_seconds=time.time()-t0,limitations=['Component tests only; full multimodal backbone and generation quality not tested.'])
if key=='unicom_vision_compression':report['limitations'].append('Published native direct SigLIP→aligner helper has 1152 vs 1216 dimension mismatch; wrapper refuses implicit adaptation. Pixel reconstruction requires separate decoder and VAE.')
(root/'validation.json').write_text(json.dumps(report,indent=2)+'\n');print('VALIDATED',key,report['elapsed_seconds'],flush=True)