"""Verify every original logical tensor hash and run the complete latent encoder."""
import hashlib,json,sys,struct
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from component import load_model

root=Path(__file__).parent
torch.set_num_threads(4)
if len(sys.argv)>1:
    original_file=Path(sys.argv[1])
    with original_file.open('rb') as f:source_sha=hashlib.file_digest(f,'sha256').hexdigest()
    assert source_sha=='9d55ca6e9c4bd404dfac552ad65dae673274b28aeffe8f3f00bd7047a1db465e'
    original=torch.load(original_file,weights_only=True,map_location='cpu')
    if 'state_dict' in original:original=original['state_dict']
    save_file({k:v.contiguous() for k,v in original.items()},str(root/'model.safetensors'))
    del original
expected=json.loads((root/'tensor_sha256.json').read_text())
with safe_open(root/'model.safetensors',framework='pt') as f:
    assert set(f.keys())==set(expected) and len(expected)==150
    for k in f.keys():
        t=f.get_tensor(k)
        assert hashlib.sha256(t.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()).hexdigest()==expected[k],k
model=load_model(root)
assert len(model.state_dict())==150
with torch.inference_mode():
    x=torch.linspace(-1,1,16*32*32).reshape(1,16,32,32)
    pooled,tokens=model(x)
    assert tuple(tokens.shape)==(1,256,768) and tuple(pooled.shape)==(1,768)
    assert torch.isfinite(tokens).all() and torch.isfinite(pooled).all()
    assert torch.equal(pooled,model.ln_post(tokens.mean(1))@model.proj)
    bp,bt=model(torch.cat([x,x*.5]))
    assert torch.equal(bt[0],tokens[0]) and torch.isfinite(bp).all()
with (root/'model.safetensors').open('rb') as f:sha=hashlib.file_digest(f,'sha256').hexdigest()
report={'all_150_previously_validated_tensor_hashes_match':True,'strict_load':True,'native_complete_encoder_forward':True,'native_pooling_exact':True,'batch_independence':True,'sha256':sha,'bytes':(root/'model.safetensors').stat().st_size,'device':'cpu','dtype':'float32'}
(root/'recovery_verification.json').write_text(json.dumps(report,indent=2)+'\n')
with (root/'model.safetensors').open('rb') as f:
    n=struct.unpack('<Q',f.read(8))[0];header=json.loads(f.read(n))
payload=sum(v['data_offsets'][1]-v['data_offsets'][0] for k,v in header.items() if k!='__metadata__')
ex=json.loads((root/'extraction.json').read_text());ex.update(sha256=sha,bytes=report['bytes'],tensor_payload_bytes=payload);(root/'extraction.json').write_text(json.dumps(ex,indent=2)+'\n')
print('OPENVISION PUBLICATION COPY VERIFIED',sha,flush=True)