"""Exact extraction checks and official visual forward parity with an echo LLM stub."""
import ast,hashlib,json,time
from pathlib import Path
from types import SimpleNamespace
import torch
from torch import nn
torch.set_num_threads(4)
from einops import rearrange
from safetensors import safe_open
from component import load_model,load_vae

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)

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():
        v=f.get_tensor(k)
        assert hashlib.sha256(v.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()).hexdigest()==hashes[k],k
check('571_source_tensor_payloads',len(hashes)==571)
m=load_model(root)
check('native_visual_strict_load',len(m.state_dict())==571)
check('complete_semantic_and_flow_layers',len(m.und_trans.layers)==26 and len(m.diffusion_head_a)==10)

# Execute the unchanged original forward/unpatchify AST against the same weights.
# The only substitute is the deliberately absent language backbone: it echoes
# supplied embeddings, so this verifies the visual boundary without claiming VQA.
ns={'torch':torch,'nn':nn,'rearrange':rearrange,'torch_int':int}
misc=ast.parse((root/'native/misc.py').read_text())
pos=next(n for n in misc.body if isinstance(n,ast.FunctionDef) and n.name=='interpolate_pos_encoding')
exec(compile(ast.fix_missing_locations(ast.Module(body=[pos],type_ignores=[])),'native/misc.py','exec'),ns)
tree=ast.parse((root/'native/modeling_showo2_qwen2_5.py').read_text())
cls=next(n for n in tree.body if isinstance(n,ast.ClassDef) and n.name=='Showo2Qwen2_5')
methods=[n for n in cls.body if isinstance(n,ast.FunctionDef) and n.name in ['forward','unpatchify']]
for fn in methods:exec(compile(ast.fix_missing_locations(ast.Module(body=[fn],type_ignores=[])),'native/modeling_showo2_qwen2_5.py','exec'),ns)
class EchoBackbone(nn.Module):
    def __init__(self,d):
        super().__init__();self.d=d;self.model=SimpleNamespace(embed_tokens=self.embed)
    def embed(self,tokens):return torch.zeros(*tokens.shape,self.d,device=tokens.device)
    def forward(self,inputs_embeds,**kwargs):
        return {'logits':inputs_embeds[...,:1],'hidden_states':[inputs_embeds]}
class Reference(nn.Module):pass
Reference.forward=ns['forward'];Reference.unpatchify=ns['unpatchify']
ref=Reference();ref.config=SimpleNamespace(**m.config)
for name,child in m.named_children():setattr(ref,name,child)
ref.showo=EchoBackbone(m.config['hidden_size']);ref.eval()
with torch.inference_mode():
    for h,w in [(8,8),(8,12)]:
        latent=torch.linspace(-1,1,16*h*w).reshape(1,16,h,w)
        t=torch.tensor([0.5]);length=h*w//4+1
        ref.image_position_ids=torch.arange(length-1).unsqueeze(0)
        mask=torch.zeros(1,1,length,length)
        positions=torch.tensor([[[0,length]]])
        inputs=m.encode_latents(latent,t)
        check(f'{h}x{w}_dual_path_fusion',tuple(inputs.shape)==(1,length,1536) and torch.isfinite(inputs).all().item())
        patches=m.predict_velocity(inputs,t,positions,attention_mask=mask)
        velocity=m.velocity_image(patches[:,1:],h,w)
        official=ref(text_tokens=torch.zeros(1,length,dtype=torch.long),image_latents=latent,t=t,attention_mask=mask,modality_positions=positions,device='cpu')[1]
        check(f'{h}x{w}_native_forward_parity',torch.equal(velocity,official),max_abs=(velocity-official).abs().max().item())
        check(f'{h}x{w}_finite_velocity',tuple(velocity.shape)==(1,16,h,w) and torch.isfinite(velocity).all().item())
    check('fixed_27x27_position_path',torch.equal(m.position_encoding(27,27),m.position_embedding.weight.unsqueeze(0)))

ex=json.loads((root/'vae/extraction.json').read_text())
with (root/'vae/source.pth').open('rb') as f:digest=hashlib.file_digest(f,'sha256').hexdigest()
check('Wan21_official_LFS_SHA256',digest==ex['source_file_lfs_sha256'])
vae=load_vae(root)
original=torch.load(root/'vae/source.pth',map_location='cpu',weights_only=True,mmap=True)
check('Wan21_native_strict_source_state',set(vae.model.state_dict())==set(original) and all(torch.equal(v,original[k]) for k,v in vae.model.state_dict().items()),tensors=len(original))
del original
with torch.inference_mode():
    for frames,h,w in [(1,64,64),(5,32,32)]:
        x=torch.linspace(-1,1,3*frames*h*w).reshape(1,3,frames,h,w)
        mu,logvar,features=vae.model.encode(x,vae.scale)
        check(f'Wan21_{frames}frame_encode',tuple(mu.shape)==(1,16,1+(frames-1)//4,h//8,w//8) and torch.isfinite(mu).all().item(),shape=list(mu.shape))
        rgb=vae.model.decode(mu,vae.scale)
        check(f'Wan21_{frames}frame_decode',tuple(rgb.shape)==tuple(x.shape) and torch.isfinite(rgb).all().item(),shape=list(rgb.shape))
        if frames==1:
            embedded=m.encode_latents(mu[:,:,0],torch.tensor([0.5]))
            check('RGB_VAE_visual_boundary',tuple(embedded.shape)==(1,17,1536) and torch.isfinite(embedded).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':['Visual forward parity uses unchanged official source methods with an echo language-backbone stub; it does not validate the excluded Qwen model, language-conditioned generation, image quality or VQA.','Visual adapter smoke coverage is single-image, small square/rectangular latent grids. Wan2.1 codec checks cover one and five frames. No full video flow head or complete sampler benchmark.','Native output token ordering and trained time-token convention are preserved. Full-model integration requires the proper multimodal sequence mask and actual final Qwen hidden states.']}
(root/'validation.json').write_text(json.dumps(report,indent=2)+'\n')
print('SHOWO VALIDATED',len(checks),flush=True)