"""Export FP32 AMOS medium one component at a time to bound conversion memory."""
import gc,json,sys
from pathlib import Path
import torch
from convert_decoder import save,Heads
from qwen_export_helpers import DecisionVision

class Embedding(torch.nn.Module):
 def __init__(self,decoder):super().__init__();self.embed_tokens=decoder.embed_tokens;self.rotary=decoder.rotary_emb
 def forward(self,input_ids,position_ids,visual_embeddings,visual_mask):
  h=torch.where(visual_mask,visual_embeddings,self.embed_tokens(input_ids))
  cos,sin=self.rotary(h,position_ids[1:]);return h,cos,sin

class Layer(torch.nn.Module):
 def __init__(self,layer):super().__init__();self.layer=layer
 def forward(self,hidden_states,cos,sin,position_ids,attention_mask):
  return self.layer(hidden_states,position_embeddings=(cos,sin),position_ids=position_ids,attention_mask=attention_mask if self.layer.block_type=='full_attention' else None,use_cache=False)

class Finish(torch.nn.Module):
 def __init__(self,model,letters):
  super().__init__();self.norm=model.model.language_model.norm
  self.register_buffer('readout',model.get_output_embeddings().weight[letters].detach().float().clone())
 def forward(self,hidden_states):
  h=self.norm(hidden_states);return h,h@self.readout.T

def main():
 source=Path('upstream/amos-1-medium');root=Path('bundles/amos-1-medium-CoreAI');root.mkdir(exist_ok=True,parents=True)
 sys.path.insert(0,str((source/'serve').resolve()))
 from amos.decoder import load_model,load_heads
 torch.set_num_threads(4)
 m,codec=load_model(str(source/'model'),'cpu',dtype=torch.bfloat16);m.eval();heads=load_heads(source/'model',m,codec)
 from stable_delta import install
 install()
 decoder=m.model.language_model;L=512;d=decoder.config.hidden_size
 manifest={'sequence_length':L,'hidden_size':d,'precision':'float32','vision_precision':'float32','echo':codec.echo,'kinds':heads.kinds,'keep':heads.keep,'has_box':heads.box is not None,'image_side':224,'noncausal':False,'text_layout':'layer-sharded','layers':[]}
 path=root/'embedding.aimodel'
 decoder.embed_tokens.float();emb=Embedding(decoder).eval()
 args=(torch.ones(1,L,dtype=torch.int32),torch.arange(L,dtype=torch.int32).view(1,1,-1).expand(4,1,-1),torch.zeros(1,L,d),torch.zeros(1,L,1,dtype=torch.bool))
 with torch.no_grad():_,cos,sin=emb(*args)
 print('Embedding',flush=True);manifest['embedding_outputs']=save(emb,args,path)
 decoder.embed_tokens.to('meta');del emb;gc.collect()
 layer_args=(torch.zeros(1,L,d),cos,sin,torch.arange(L,dtype=torch.int32)[None],torch.ones(1,1,L,L,dtype=torch.bool).tril())
 for i,layer in enumerate(decoder.layers):
  name=f'layer-{i:02}.aimodel';print(name,flush=True)
  layer.float().eval();outputs=save(Layer(layer).eval(),layer_args,root/name)
  manifest['layers'].append({'asset':name,'output':outputs[0],'kind':layer.block_type})
  layer.to('meta');gc.collect()
 print('Finish',flush=True);decoder.norm.float();manifest['outputs']=save(Finish(m,codec.letters).eval(),(torch.zeros(1,L,d),),root/'finish.aimodel')
 m.get_output_embeddings().to('meta');decoder.norm.to('meta');gc.collect()
 print('Vision',flush=True);m.model.visual.float();manifest['vision_output']=save(DecisionVision(m.model.visual).eval(),(torch.zeros(196,1536),),root/'vision.aimodel')[0]
 m.model.visual.to('meta');gc.collect()
 print('Heads',flush=True);manifest['head_outputs']=save(Heads(heads).eval(),(torch.zeros(d),torch.zeros(L,d),torch.zeros(49,d)),root/'heads.aimodel')
 (root/'manifest.json').write_text(json.dumps(manifest,indent=2));print('Saved FP32 sharded model',flush=True)
if __name__=='__main__':main()