import hashlib,json,os,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:
    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('1190_original_tensor_payloads',len(expected)==1190)
m=load_model(root,os.environ['UNILIP_DECODER_SOURCE'])
check('strict_all_1190_entries',len(m.state_dict())==1190)
check('flow_scaling_factor',m.unilip_factor==10.6)
check('last_six_Qwen_layers',len(m.model.llm_connector.layers)==6 and m.model.llm_connector.layers[0].self_attn.layer_idx==18)
with torch.inference_mode():
    for size in [112,448]:
        x=torch.linspace(-1,1,3*size*size).reshape(1,3,size,size)
        features=m.image_features(x);g=size//28
        check(f'{size}px_image_features',tuple(features.shape)==(1,g*g,896) and torch.isfinite(features).all().item())
        raw=m.model.vision_tower(x).last_hidden_state[:,1:].reshape(1,size//14,size//14,1024)
        # Independent indexing oracle for the original InternVL shuffle order:
        # channels are [inner-height, inner-width, feature], not pixel_unshuffle order.
        shuffled=raw.reshape(1,g,2,g,2,1024).permute(0,1,3,2,4,5).reshape(1,g*g,4096)
        direct=m.model.multi_modal_projector(shuffled)
        check(f'{size}px_shuffle_order_exact',torch.equal(features,direct))
        z=m.model.vae_decoder.clip_down(features)
        check(f'{size}px_native_latent',tuple(z.shape)==(1,32,g,g) and torch.isfinite(z).all().item())
        rgb=m.decode(z)
        direct_rgb=m.model.vae_decoder(features)
        check(f'{size}px_native_decoder_exact',torch.equal(rgb,direct_rgb))
        check(f'{size}px_finite_RGB',tuple(rgb.shape)==(1,3,g*32,g*32) and torch.isfinite(rgb).all().item(),shape=list(rgb.shape))
    hidden=torch.linspace(-0.25,0.25,8*896).reshape(1,8,896)
    mask=torch.ones(1,8,dtype=torch.long)
    cond=m.condition(hidden,mask)
    direct=m.model.llm_connector(inputs_embeds=hidden,attention_mask=torch.zeros(1,1,8,8),position_ids=torch.arange(8)[None],use_cache=False).last_hidden_state
    check('bidirectional_connector_exact',torch.equal(cond,m.model.projector(direct)) and tuple(cond.shape)==(1,8,2304))
    check('trained_query_shape',tuple(m.model.latent_queries.shape)==(1,256,896) and torch.isfinite(m.model.latent_queries).all().item())
    z=torch.linspace(-1,1,32*4*4).reshape(1,32,4,4)
    y=m.model.dit(z,timestep=torch.tensor([500.]),encoder_hidden_states=cond,encoder_attention_mask=mask,return_dict=False)[0]
    check('complete_28layer_SANA_flow',tuple(y.shape)==z.shape and torch.isfinite(y).all().item())
    flow=m.model.vae_decoder.clip_down(features)*m.unilip_factor
    check('flow_scale_roundtrip',torch.allclose(flow/m.unilip_factor,m.model.vae_decoder.clip_down(features),atol=1e-6,rtol=1e-6))
report={'strict_load':True,'source_tensor_bytes_equal':True,'checks':checks,'smoke':{'passed':True,'device':'cpu','dtype':'float32'},'elapsed_seconds':time.monotonic()-started,'limitations':['No original Qwen language model or LM head is included; connector inputs in smoke tests are synthetic hidden states, not a text-to-image quality evaluation.','Original vision encoder, original trained custom decoder class, final six Qwen layers, projections, learned queries and full SANA DiT are loaded strictly.','The official custom decoder source is supplied externally by pinned SHA256 because its repository has no root code license.','112px/448px vision fixtures and 128px/512px RGB decode; CPU float32 tests do not establish CUDA FlashAttention numerical parity.']}
(root/'validation.json').write_text(json.dumps(report,indent=2)+'\n')
print('UNILIP VALIDATED',len(checks),flush=True)