"""Reproduce the AMOS mini Core AI export from pinned source files."""
import argparse,json,sys
from pathlib import Path
import torch
from coreai_torch import TorchConverter,get_decomp_table

PINNED_DEPENDENCIES={'jhu-clsp/ettin-encoder-400m':'7662476d60abb071a5bd319c9f3074f3072c062d','google/siglip2-base-patch16-224':'75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2'}
def pin_dependencies():
 from huggingface_hub import hf_hub_download
 import amos.packing,amos.vision
 def pinned(repo,filename,**kwargs):
  return hf_hub_download(repo,filename,revision=PINNED_DEPENDENCIES[repo],**kwargs)
 amos.packing.hf_hub_download=pinned;amos.vision.hf_hub_download=pinned

class Text(torch.nn.Module):
 def __init__(self,m):
  super().__init__(); self.m=m
 def forward(self,input_ids,position_ids,segment_ids,option_ids,visual_embeddings,visual_mask):
  m=self.m
  emb=torch.where(visual_mask,visual_embeddings,m.encoder.tok_embeddings(input_ids))
  h=m.encoder(input_ids,position_ids,segment_ids,option_ids,inputs_embeds=emb)
  return m.ptr_q(h),m.ptr_o(h),m.span_q(h),m.span_k(h),m.region_q(h),m.region_k(h),m.multi_bias

class Vision(torch.nn.Module):
 def __init__(self,m):
  super().__init__(); self.vision=m.vision; self.projector=m.projector
 def forward(self,pixel_values):
  v=self.vision
  x=v.patch_embedding(pixel_values).flatten(2).transpose(1,2)+v.position_embedding
  for b in v.blocks:x=b(x)
  return self.projector(v.post_layernorm(x))

def save(m,args,path):
 ep=torch.export.export(m,args).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('--source',default='upstream/amos-1-mini');p.add_argument('--output',default='bundles/amos-1-mini-CoreAI');p.add_argument('--length',type=int,default=512);a=p.parse_args()
 source=Path(a.source).resolve();out=Path(a.output);out.mkdir(parents=True,exist_ok=True)
 sys.path.insert(0,str(source/'serve'))
 from amos.model import load_run
 pin_dependencies()
 torch.set_num_threads(4)
 m,codec=load_run(source);m.eval()
 args=(torch.full((1,a.length),codec.pad,dtype=torch.int32),torch.arange(a.length,dtype=torch.int32)[None],torch.zeros(1,a.length,dtype=torch.int32),torch.zeros(1,a.length,dtype=torch.int32),torch.zeros(1,a.length,m.encoder.config.hidden_size),torch.zeros(1,a.length,1,dtype=torch.bool))
 print('Exporting text',flush=True)
 outputs=save(Text(m).eval(),args,out/'text.aimodel')
 print('Exporting vision',flush=True)
 vo=save(Vision(m).eval(),(torch.zeros(1,3,224,224),),out/'vision.aimodel')
 (out/'manifest.json').write_text(json.dumps({'sequence_length':a.length,'hidden_size':m.encoder.config.hidden_size,'outputs':outputs,'vision_output':vo[0],'precision':'float32'},indent=2))
 print('Saved',flush=True)
if __name__=='__main__':main()