"""Native Show-o2 visual I/O modules, with the Qwen backbone supplied externally."""
import json, math
from pathlib import Path
import torch
from torch import nn
import torch.nn.functional as F
from accelerate import init_empty_weights
from safetensors.torch import load_file
from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
from native.modeling_siglip import SiglipEncoder
from native.modules import (PatchEmbed,RMSNorm,TimestepEmbedder,DiffusionHeadConfig,
                            ModulatedAttentionBlock,FinalLayer)

class VisualFlowComponents(nn.Module):
    def __init__(self,config,siglip_config):
        super().__init__()
        self.config=config
        p=config['patch_size'];c=config['image_latent_dim']
        d=config['hidden_size'];s=config['clip_latent_dim']
        self.image_embedder_und=PatchEmbed(patch_size=p,in_chans=c,embed_dim=s)
        self.image_embedder_gen=PatchEmbed(patch_size=p,in_chans=c,embed_dim=d)
        vision=SiglipVisionConfig(**siglip_config['vision_config'])
        vision._attn_implementation='sdpa'
        vision.num_hidden_layers-=1
        self.position_embedding=nn.Embedding((vision.image_size//vision.patch_size)**2,s)
        self.und_trans=SiglipEncoder(vision)
        self.fusion_proj=nn.Sequential(RMSNorm(s+d),nn.Linear(s+d,d),nn.GELU(),nn.Linear(d,d))
        head=DiffusionHeadConfig()
        self.time_embed=TimestepEmbedder(head.hidden_size)
        self.diff_proj=nn.Sequential(nn.Linear(d,head.hidden_size),nn.GELU(),nn.Linear(head.hidden_size,head.hidden_size))
        self.time_embed_proj=nn.Linear(head.hidden_size,d)
        self.diffusion_head_a=nn.ModuleList([ModulatedAttentionBlock(head,i) for i in range(config['num_diffusion_layers'])])
        self.diffusion_head_b=FinalLayer(head.hidden_size,p,c)

    def position_encoding(self,height,width):
        # Exact native misc.interpolate_pos_encoding arithmetic, patch_size=1.
        weight=self.position_embedding.weight
        side=int(math.sqrt(weight.shape[0]))
        value=weight.unsqueeze(0).reshape(1,side,side,weight.shape[1]).permute(0,3,1,2)
        value=F.interpolate(value,size=(height,width),mode='bicubic',align_corners=False)
        return value.permute(0,2,3,1).reshape(1,-1,weight.shape[1])

    def encode_latents(self,latents,t):
        """One image per batch item: normalized VAE BCHW -> time+visual LLM inputs."""
        assert latents.ndim==4 and latents.shape[1]==16
        p=self.config['patch_size'];h,w=latents.shape[-2:]
        assert h%p==0 and w%p==0
        und=self.image_embedder_und(latents)
        if (h//p,w//p)==(27,27):
            pos=self.position_embedding(torch.arange(729,device=latents.device).unsqueeze(0))
        else:
            pos=self.position_encoding(h//p,w//p)
        und=self.und_trans(und+pos,return_dict=True).last_hidden_state
        gen=self.image_embedder_gen(latents)
        fused=self.fusion_proj(torch.cat([und,gen],dim=-1))
        time=self.time_embed(t,latents.dtype)
        return torch.cat([self.time_embed_proj(time).unsqueeze(1),fused],dim=1)

    def predict_velocity(self,last_hidden_states,t,modality_positions,attention_mask=None):
        """External Qwen final states -> native flow prediction in patch-vector order.

        modality_positions includes each time token. Use the full sequence's
        native multimodal attention mask for integration with text/interleaving.
        """
        time=self.time_embed(t,last_hidden_states.dtype)
        x=self.diff_proj(last_hidden_states)
        ids=torch.arange(x.shape[1],device=x.device).unsqueeze(0)
        for layer in self.diffusion_head_a:
            x=layer(hidden_states=x,adaln_input=time,attention_mask=attention_mask,
                    position_ids=ids,modality_positions=modality_positions)[0]
        return self.diffusion_head_b(x,time,modality_positions)

    def velocity_image(self,patch_vectors,height,width):
        """Exact native single-image output order; input excludes time token."""
        p=self.config['patch_size'];c=self.config['image_latent_dim']
        b=patch_vectors.shape[0]
        x=patch_vectors.reshape(b,height//p,width//p,p,p,c)
        x=x.reshape(b,height*width,c).transpose(1,2)
        return x.reshape(b,c,height,width)

def load_model(path=None,dtype=torch.float32,device='cpu'):
    root=Path(path or Path(__file__).parent)
    with init_empty_weights():
        model=VisualFlowComponents(json.loads((root/'source_config.json').read_text()),json.loads((root/'siglip_config.json').read_text()))
    model.load_state_dict(load_file(root/'model.safetensors'),strict=True,assign=True)
    return model.to(device=device,dtype=dtype).eval()

def load_vae(path=None,dtype=torch.float32,device='cpu'):
    from native.wan21_vae import WanVAE
    root=Path(path or Path(__file__).parent)
    return WanVAE(z_dim=16,vae_pth=str(root/'vae/source.pth'),dtype=dtype,device=device)