"""Export AMOS Qwen3.5 bidirectional decisions and trained quote/box heads."""
import argparse,json,math,sys
from pathlib import Path
import torch
from coreai_torch import TorchConverter,get_decomp_table
from qwen_export_helpers import DecisionVision

class Text(torch.nn.Module):
 def __init__(self,model,letters):
  super().__init__();self.language_model=model.model.language_model;self.noncausal=bool(getattr(model.config,'amos_noncausal',False))
  self.register_buffer('readout',model.get_output_embeddings().weight[letters].detach().float().clone())
 def forward(self,input_ids,attention_mask,position_ids,visual_embeddings,visual_mask):
  emb=self.language_model.embed_tokens(input_ids)
  emb=torch.where(visual_mask,visual_embeddings.to(emb.dtype),emb)
  mask=attention_mask[:,None,None,:].bool()
  if not self.noncausal:
   width=input_ids.shape[1];idx=torch.arange(width,device=input_ids.device);mask=mask&(idx[:,None]>=idx[None,:])[None,None]
  h=self.language_model(inputs_embeds=emb,attention_mask={'full_attention':mask,'linear_attention':None},position_ids=position_ids,use_cache=False).last_hidden_state
  return h.float(),h.float()@self.readout.T

class Heads(torch.nn.Module):
 def __init__(self,heads):super().__init__();self.heads=heads
 def forward(self,question_hidden,state_hidden,image_hidden):
  h=self.heads;d=question_hidden.shape[0];dtype=h.span_q.weight.dtype
  q=question_hidden.to(dtype);s=state_hidden.to(dtype);im=image_hidden.to(dtype)
  span=torch.einsum('cd,lcd->cl',h.span_q(q).view(2,d),h.span_k(s).view(-1,2,d))/math.sqrt(d)
  region=h.region_k(im)@h.region_q(q)/math.sqrt(d)
  box=h.box(q,im,region,7,7) if h.box is not None else torch.zeros(4)
  return span.float(),region.float(),box.float()

def save(m,args,path):
 with torch.no_grad():ep=torch.export.export(m,args,strict=False).run_decompositions(get_decomp_table())
 program=TorchConverter().add_exported_program(ep).to_coreai();program.save_asset(path)
 return list(ep.graph_signature.user_outputs)

def main():
 p=argparse.ArgumentParser();p.add_argument('name');p.add_argument('--length',type=int,default=512);p.add_argument('--precision',default='bfloat16');a=p.parse_args()
 source=Path('upstream')/a.name;out=Path('bundles')/(a.name+'-CoreAI');out.mkdir(parents=True,exist_ok=True)
 sys.path.insert(0,str((source/'serve').resolve()))
 from amos.decoder import load_model,load_heads
 torch.set_num_threads(4);model,codec=load_model(str(source/'model'),'cpu',dtype=getattr(torch,a.precision));model.eval();heads=load_heads(source/'model',model,codec)
 d=model.config.text_config.hidden_size;L=a.length
 inputs=(torch.full((1,L),codec.pad,dtype=torch.int32),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))
 print('Exporting text',flush=True);outputs=save(Text(model,codec.letters).eval(),inputs,out/'text.aimodel')
 print('Exporting vision',flush=True);visual=model.model.visual;vo=save(DecisionVision(visual).eval(),(torch.zeros(196,1536),),out/'vision.aimodel')
 print('Exporting heads',flush=True);ho=save(Heads(heads).eval(),(torch.zeros(d),torch.zeros(L,d),torch.zeros(49,d)),out/'heads.aimodel')
 (out/'manifest.json').write_text(json.dumps({'sequence_length':L,'hidden_size':d,'outputs':outputs,'vision_output':vo[0],'head_outputs':ho,'precision':a.precision,'echo':codec.echo,'kinds':heads.kinds,'keep':heads.keep,'has_box':heads.box is not None,'image_side':224,'noncausal':bool(getattr(model.config,'amos_noncausal',False))},indent=2))
 print('Saved',flush=True)
if __name__=='__main__':main()