import hashlib,json,time
from pathlib import Path
import torch
torch.set_num_threads(4)
from safetensors import safe_open
from component import load_model,load_projector
from native.rqvaesiglip.modules import Upsample

root=Path(__file__).parent;checks=[];started=time.monotonic()
def check(name,ok,**details):
    assert ok,(name,details)
    checks.append({'name':name,'passed':True,**details});print(name,details,flush=True)
for name,path in [('vision',root/'model.safetensors'),('projector',root/'projector/model.safetensors')]:
    ex=json.loads((path.parent/'extraction.json').read_text())
    with path.open('rb') as f:digest=hashlib.file_digest(f,'sha256').hexdigest()
    check(name+'_official_LFS_SHA256',digest==ex['source_file_lfs_sha256'])
m=load_model(root);projector=load_projector(root)
check('native_vision_strict_load',len(m.state_dict())==1107)
check('native_projector_strict_load',len(projector.state_dict())==4)
for state,path in [(m.state_dict(),root/'model.safetensors'),(projector.state_dict(),root/'projector/model.safetensors')]:
    with safe_open(path,framework='pt') as f:
        assert set(f.keys())==set(state)
        assert all(torch.equal(state[k],f.get_tensor(k)) for k in f.keys())
check('all_1111_native_state_values',True)
codec=m.rqvaesiglip;q=codec.quantizer
check('four_depths_shared_codebook',len(q.codebooks)==4 and all(c is q.codebooks[0] for c in q.codebooks))
check('16384_codes_plus_padding_row',tuple(q.codebooks[0].weight.shape)==(16385,1024))
with torch.inference_mode():
    bf=torch.linspace(-1,1,32*4).reshape(1,32,2,2).bfloat16()
    expected_bf=torch.nn.functional.interpolate(bf.float(),scale_factor=2.0,mode='nearest').bfloat16()
    check('original_bfloat16_upsample_parity',torch.equal(Upsample(32,False)(bf),expected_bf))
    x=torch.linspace(-1,1,3*256*256).reshape(1,3,256,256)
    ids,z=codec.encode_image(x)
    check('native_encode_grid',tuple(ids.shape)==(1,16,16,4) and tuple(z.shape)==(1,16,16,1024))
    check('code_ID_bounds_exclude_padding',ids.min().item()>=0 and ids.max().item()<16384)
    canonical=q.embed_code(ids)
    check('straight_through_lookup_bound',torch.allclose(z,canonical,atol=2e-5,rtol=2e-5),max_abs=(z-canonical).abs().max().item())
    depth,_=codec.get_code_emb_with_depth(ids)
    check('depth_lookup_sum_exact',torch.equal(depth.sum(dim=-2),canonical))
    rgb=codec.decode(canonical)
    direct=codec.decoder(codec.post_quant_conv(canonical.permute(0,3,1,2).contiguous()))
    check('native_decode_lookup_exact',torch.equal(rgb,direct))
    check('finite_256px_RGB',tuple(rgb.shape)==tuple(x.shape) and torch.isfinite(rgb).all().item(),mse=(rgb-x).square().mean().item())
    features=projector(canonical.flatten(1,2))
    check('language_projector_geometry',tuple(features.shape)==(1,256,4096) and torch.isfinite(features).all().item())
    flat_ids=ids.reshape(1,256,4)[:,:2]
    logits=m.rqtransformer(features[:,:2],flat_ids,model_aux=codec)
    check('native_depth_teacher_forcing',tuple(logits.shape)==(1,2,4,16384) and torch.isfinite(logits).all().item())
    torch.manual_seed(721)
    generated,gids=m.rqtransformer.generate(features[:,:1].repeat(2,1,1),model_aux=codec,cfg=3.0)
    check('native_four_depth_generation',tuple(gids.shape)==(2,1,4) and gids.min().item()>=0 and gids.max().item()<16384)
    by_depth,_=codec.get_code_emb_with_depth(gids)
    check('generated_feature_lookup_exact',torch.equal(generated,by_depth.cumsum(dim=-2)[:,:,-1,:]))
    text=codec.siglip_model.get_text_features(input_ids=torch.zeros(1,64,dtype=torch.long))
    check('retained_trained_text_tower',text.shape[0]==1 and torch.isfinite(text).all().item(),shape=list(text.shape))
report={'strict_load':True,'source_tensor_bytes_equal':True,'checks':checks,'smoke':{'passed':True,'device':'cpu','dtype':'float32'},'elapsed_seconds':time.monotonic()-started,'limitations':['Synthetic fixed 256px native image fixture and short depth-transformer sequences, not image-quality or VQA benchmarks.','The separate Llama backbone and language tokenizer are excluded. Depth generation tests use visual-projector fixtures at that external language-state boundary.','Four residual depths share one learned codebook; the 16385th stored row is padding, not a valid generated visual ID.']}
(root/'validation.json').write_text(json.dumps(report,indent=2)+'\n')
print('VILAU VALIDATED',len(checks),flush=True)