import hashlib, json, time
from pathlib import Path
import torch
from importlib.metadata import version
torch.set_num_threads(4)
from safetensors import safe_open
from component import load_model

root=Path(__file__).parent
started=time.monotonic(); checks=[]
def check(name,ok,**details):
    assert ok,(name,details)
    checks.append({'name':name,'passed':True,**details});print(name,details,flush=True)

hashes=json.loads((root/'tensor_sha256.json').read_text())
with safe_open(root/'model.safetensors',framework='pt') as f:
    assert set(f.keys())==set(hashes)
    for k in f.keys():
        t=f.get_tensor(k)
        assert hashlib.sha256(t.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()).hexdigest()==hashes[k],k
check('1335_source_tensor_payloads',len(hashes)==1335)
m=load_model(root)
check('native_strict_load',len(m.state_dict())==1335)
assert type(m.text_encoder).__module__=='native.open_clip.transformer'
check('eight_4096_by_8_codebooks',len(m.quantizer.codebooks)==8 and all(tuple(c.codebook.weight.shape)==(4096,8) for c in m.quantizer.codebooks))
with torch.inference_mode():
    x=torch.linspace(-1,1,3*256*256).reshape(1,3,256,256)
    ids=m.img_to_idx(x)
    check('token_geometry',tuple(ids.shape)==(1,8,256),shape=list(ids.shape))
    check('per_codebook_ID_bounds',ids.min().item()>=0 and ids.max().item()<4096)
    q=m.quantizer.idx_to_f(ids)
    canonical=torch.cat([c.codebook(ids[:,i]) for i,c in enumerate(m.quantizer.codebooks)],dim=-1)
    check('normalized_lookup_exact',torch.equal(q,canonical))
    rgb=m.idx_to_img(ids)
    direct=m.decoder(m.post_quant_proj(canonical)).clamp(-1,1)
    check('decode_ID_lookup_exact',torch.equal(rgb,direct))
    check('finite_RGB_reconstruction',rgb.shape==x.shape and torch.isfinite(rgb).all().item(),mse=(rgb-x).square().mean().item())
    reconstructed=m.img_to_reconstructed_img(x)
    check('native_quantizer_reconstruction_bound',torch.allclose(reconstructed,rgb,atol=2e-5,rtol=2e-5),max_abs=(rgb-reconstructed).abs().max().item())
    visual=m.encode_image(x,normalize=True)
    text=torch.zeros((1,m.context_length),dtype=torch.long)
    text[0,0]=49406;text[0,1]=49407
    textual=m.encode_text(text,normalize=True)
    check('semantic_image_features',tuple(visual.shape)==(1,768) and torch.isfinite(visual).all().item() and torch.allclose(visual.norm(dim=-1),torch.ones(1),atol=1e-5))
    check('semantic_text_features',tuple(textual.shape)==(1,768) and torch.isfinite(textual).all().item() and torch.allclose(textual.norm(dim=-1),torch.ones(1),atol=1e-5))
report={'strict_load':True,'source_tensor_bytes_equal':True,'checks':checks,'smoke':{'passed':True,'device':'cpu','dtype':'float32','resolution':[256,256]},'environment':{p:version(p) for p in ['torch','torchvision','timm','safetensors','typed-argument-parser','pytz']},'text_implementation':'Unchanged official project-vendored native.open_clip.transformer.TextTransformer, same configured default constructor branch. External OpenCLIP is not imported.','elapsed_seconds':time.monotonic()-started,'limitations':['Synthetic native codec and semantic feature fixtures; no downstream image quality, retrieval, VQA or text-to-image benchmark.','This is the complete trained tokenizer mapping including its semantic text tower; no language generator is included.','The eight codebooks each have 4096 entries of width 8. The sum of entries is 32768; IDs are local to each codebook.']}
(root/'validation.json').write_text(json.dumps(report,indent=2)+'\n')
print('UNITOK VALIDATED',len(checks),flush=True)