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

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)
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)
    for k in f.keys():
        v=f.get_tensor(k)
        assert hashlib.sha256(v.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()).hexdigest()==expected[k],k
check('684_source_tensor_payloads',len(expected)==684)
m=load_model(root)
check('native_strict_load',len(m.state_dict())==684)
check('original_VAE_scale',m.vae_scale==0.2325)
with torch.inference_mode():
    for h,w in [(64,64),(64,96)]:
        x=torch.linspace(-1,1,3*h*w).reshape(1,3,h,w)
        z=m.encode(x)
        native=m.vae.encode(x).mode()*m.vae_scale
        native=native.permute(0,2,3,1).contiguous()
        check(f'{h}x{w}_native_encode_exact',torch.equal(z,native) and tuple(z.shape)==(1,h//16,w//16,16),shape=list(z.shape))
        rgb=m.decode(z)
        direct=m.vae.decode((z/m.vae_scale).permute(0,3,1,2).contiguous())
        check(f'{h}x{w}_native_decode_exact',torch.equal(rgb,direct))
        check(f'{h}x{w}_finite_reconstruction',rgb.shape==x.shape and torch.isfinite(rgb).all().item(),mse=(rgb-x).square().mean().item())
        tokens=z.flatten(1,2);n=tokens.shape[1]
        mask=torch.zeros(1,n);mask[:,::2]=1
        encoded=m.mar.forward_mae_encoder(tokens,mask,m.mar.fake_latent,image_shape=(h//16,w//16))
        check(f'{h}x{w}_masked_encoder_geometry',encoded.shape[1]==m.mar.buffer_size+n//2 and torch.isfinite(encoded).all().item(),shape=list(encoded.shape))
        decoded=m.mar.forward_mae_decoder(encoded,mask,image_shape=(h//16,w//16))
        check(f'{h}x{w}_masked_decoder_geometry',tuple(decoded.shape)==(1,n,m.mar.decoder_embed_dim) and torch.isfinite(decoded).all().item())
        projected=m.proj_in(encoded)
        restored=m.proj_out(projected)
        check(f'{h}x{w}_language_projections',projected.shape[-1]==896 and restored.shape==encoded.shape and torch.isfinite(restored).all().item())
    c=decoded[0,:2]
    prediction=m.mar.diffloss.net(torch.zeros(2,16),torch.tensor([500,500]),c)
    check('diffusion_prediction_and_variance',tuple(prediction.shape)==(2,32) and torch.isfinite(prediction).all().item())
    torch.manual_seed(721)
    sampled=m.mar.diffloss.sample(c,cfg=1.0)
    check('native_100step_sampling_cpu',m.mar.diffloss.gen_diffusion.num_timesteps==100 and tuple(sampled.shape)==(2,16) and torch.isfinite(sampled).all().item())
    torch.manual_seed(721)
    guided=m.mar.diffloss.sample(c,cfg=2.0)
    check('native_CFG_sampling_cpu',tuple(guided.shape)==(2,16) and torch.isfinite(guided).all().item())
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 64px square/rectangular VAE/MAR fixtures and two-token native diffusion sampling; not image quality or downstream VQA benchmarks.','The trained Qwen model is excluded; complete language-conditioned generation still requires its hidden states and the official masked generation schedule.','Sampling tensor allocations follow the input device instead of hardcoded CUDA; native probability distributions and sampler equations are retained.']}
(root/'validation.json').write_text(json.dumps(report,indent=2)+'\n')
print('HARMON VALIDATED',len(checks),flush=True)