"""Pinned native visual components. No text backbone or full generation pipeline."""
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
KEY = 'unicom_vision_compression'

def linear_from(state,name):
 w=state[name+'.weight'];return nn.Linear(w.shape[1],w.shape[0],bias=name+'.bias' in state)

def build(config,s):
 m=nn.Module()
 if KEY=='internlumina_u2_visual_io':
  from native.multicodebook import MultiCodebookARHead
  m.model=nn.Module();m.model.img_embeddings=nn.ModuleList([nn.Embedding(*s[f'model.img_embeddings.{i}.weight'].shape) for i in range(8)])
  m.model.img_embeddings_proj=linear_from(s,'model.img_embeddings_proj')
  m.lm_img_head=nn.ModuleList([linear_from(s,f'lm_img_head.{i}') for i in range(8)])
  m.multi_codebook_head=MultiCodebookARHead(config['hidden_size'],8,channel_depth=6,channel_heads=8)
 elif KEY=='unicom_vision_compression':
  from native.siglip2 import Siglip2VisionTransformer,LightProjector
  from native.compression import AttentionCompressionNet
  from native.flow_interfaces import TimestepEmbedder,LinearPatchEmbed,LinearFinalLayer
  m.vision_model=Siglip2VisionTransformer(config['vit']);m.vision_aligner=LightProjector(config['vit_aligner']);m.compression=AttentionCompressionNet(config['siglip_original_dim'],config['siglip_latent_dim'])
  h=config['hidden_size'];c=config['siglip_latent_dim'];m.timestep_emb=TimestepEmbedder(h);m.time_embed=TimestepEmbedder(h);m.patch_embed=LinearPatchEmbed(c,h);m.final_layer=LinearFinalLayer(h,c)
 elif KEY=='internvlu_vision_bridge':
  from native.modeling_intern_vit import InternVisionModel
  from native.configuration_intern_vit import InternVisionConfig
  c=InternVisionConfig(**config['vision_config']);c.use_flash_attn=False;m.vision_model=InternVisionModel(c)
  n=s['mlp1.0.weight'].shape[0];m.mlp1=nn.Sequential(nn.LayerNorm(n),linear_from(s,'mlp1.1'),nn.GELU(),linear_from(s,'mlp1.3'))
  m.special_token_embedding=nn.Embedding(*s['special_token_embedding.weight'].shape)
 elif KEY=='hidream_o1_vision_pixel_io':
  from native.qwen3_vl_transformers import Qwen3VLVisionModel
  from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLVisionConfig
  from native.pixel_interfaces import BottleneckPatchEmbed,FinalLayer,TimestepEmbedder
  c=Qwen3VLVisionConfig(**config['vision_config']);c._attn_implementation='eager';m.model=nn.Module();m.model.visual=Qwen3VLVisionModel(c)
  h=config['text_config']['hidden_size'];b=s['model.x_embedder.proj1.weight'].shape[0];m.model.x_embedder=BottleneckPatchEmbed(None,32,3,b,h);m.model.t_embedder1=TimestepEmbedder(None,h);m.model.final_layer2=FinalLayer(None,h,32,3)
 elif KEY=='ming_flash_omni2_vision_io':
  from native.qwen3_moe_vit import Qwen3MoeVisionTransformer,Qwen3VLMoeVisionConfig
  c=Qwen3VLMoeVisionConfig(**config['vision_config']);c._attn_implementation='eager';m.vision=Qwen3MoeVisionTransformer(c)
  m.linear_proj=nn.Sequential(linear_from(s,'linear_proj.0'),nn.GELU(),linear_from(s,'linear_proj.2'))
 elif KEY=='lance_vision':
  from native.qwen2_5_vl_vit import Qwen2_5_VisionTransformerPretrainedModel
  from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig
  c=Qwen2_5_VLVisionConfig(**config);c._attn_implementation='eager';m=Qwen2_5_VisionTransformerPretrainedModel(c)
 else:raise ValueError(KEY)
 return m

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

def embed_codes(model,codes):
 """InternLumina-U2 per-position 8-codebook input. Model vocab also contains mask IDs."""
 if KEY!='internlumina_u2_visual_io':raise ValueError(KEY)
 if codes.shape[-1]!=8:raise ValueError('Expected 8 codebooks in original order')
 return model.model.img_embeddings_proj(torch.cat([model.model.img_embeddings[i](codes[...,i]) for i in range(8)],dim=-1))

def codebook_logits(model,hidden,teacher_codes):
 """Teacher-forced head, without the diffusion backbone. hidden: [positions,2048]."""
 if KEY!='internlumina_u2_visual_io':raise ValueError(KEY)
 e=torch.stack([model.model.img_embeddings[i](teacher_codes[:,i]) for i in range(7)],dim=1)
 x=model.multi_codebook_head(hidden,e)
 return torch.stack([model.lm_img_head[i](x[:,i]) for i in range(8)],dim=1)

def unicom_features(model,pixel_patches,spatial_shapes,attention_mask=None):
 """Return 1152-d semantic features and 64-d compressed features, with native NaFlex layout."""
 if KEY!='unicom_vision_compression':raise ValueError(KEY)
 y=model.vision_model(pixel_patches,attention_mask=attention_mask,spatial_shapes=spatial_shapes).last_hidden_state
 return y,model.compression.compress(y)

def unicom_align(model,features):
 """Requires the checkpoint's 1216-d aligner input. No speculative feature concatenation."""
 if features.shape[-1]!=model.vision_aligner.layers[0].in_features:
  raise ValueError('Upstream checkpoint aligner expects 1216 features, while raw SigLIP emits 1152. Supply the correct upstream adapter input explicitly.')
 return model.vision_aligner(features)

def internvlu_features(model,pixels):
 if KEY!='internvlu_vision_bridge':raise ValueError(KEY)
 from native.pixel_shuffle import pixel_shuffle_v2
 y=model.vision_model(pixel_values=pixels,return_dict=True).last_hidden_state[:,1:]
 return model.mlp1(pixel_shuffle_v2(None,y,scale_factor=0.5,patch_aspect_ratio=1.0))

def ming_features(model,patches,grid_thw):
 if KEY!='ming_flash_omni2_vision_io':raise ValueError(KEY)
 y=model.vision(patches,grid_thw=grid_thw)
 if isinstance(y,tuple):y=y[0]
 return F.normalize(model.linear_proj(y),dim=-1)