"""Core AI neural execution; upstream packing/decoding uses PyTorch tensors only."""
import json,math,sys
from pathlib import Path
import numpy as np
import torch
from coreai.runtime import AIModel,NDArray,SpecializationOptions

class AmosMini:
 @classmethod
 async def load(cls,bundle):
  self=cls();self.root=Path(bundle).resolve();sys.path.insert(0,str(self.root/'serve'))
  from amos.packing import Codec
  cfg=json.loads((self.root/'config.json').read_text())
  # Codec's loader accepts a locally pinned tokenizer through this scoped patch.
  import amos.packing as packing
  old=packing.hf_hub_download
  packing.hf_hub_download=lambda repo,filename: str(self.root/filename)
  try:self.codec=Codec(cfg['repo'],max_state=cfg['max_state'],digits=cfg.get('digits',False),score_rank=cfg.get('score_rank',False))
  finally:packing.hf_hub_download=old
  self.manifest=json.loads((self.root/'manifest.json').read_text())
  self.text=await AIModel.load(self.root/'text.aimodel',SpecializationOptions.default());self.tf=self.text.load_function('main')
  self.vision=await AIModel.load(self.root/'vision.aimodel',SpecializationOptions.default());self.vf=self.vision.load_function('main')
  return self

 async def raw(self,state,questions,images=()):
  from amos.packing import pack,collate
  from amos.vision import preprocess
  feats=[]
  async def feature(im):
   pix=preprocess(im)[None].numpy()
   out=await self.vf({'pixel_values':NDArray(pix)})
   return torch.from_numpy(out[self.manifest['vision_output']].numpy()[0].copy())
  # Pack expects pre-projection 196-cell features. Instead use placeholders;
  # inject the Core AI projected 49-cell vectors directly into embedding slots.
  for im in images:feats.append(await feature(im))
  qs=[]
  for q in questions:
   q=dict(q)
   if 'options' in q:
    opts=[]
    for o in q['options']:
     if isinstance(o,str):opts.append(o)
     else:opts.append(torch.zeros(196,768));feats.append(await feature(o))
    q['options']=opts
   qs.append(q)
  p=pack(self.codec,state,qs,[torch.zeros(196,768) for _ in images],strict=True)
  n=len(p);L=self.manifest['sequence_length'];d=self.manifest['hidden_size']
  if n>L:raise ValueError(f'Request has {n} tokens; Core AI capacity is {L}.')
  b=collate([p],self.codec.pad)
  args={}
  for src,dst,pad in [('ids','input_ids',self.codec.pad),('pos','position_ids',0),('seg','segment_ids',-1),('opt','option_ids',0)]:
   arr=np.full((1,L),pad,dtype=np.int32);arr[0,:n]=getattr(p,src);args[dst]=NDArray(arr)
  ve=np.zeros((1,L,d),np.float32);vm=np.zeros((1,L,1),bool)
  for at,f in zip(p.image_at,feats):ve[0,at:at+49]=f.numpy();vm[0,at:at+49]=True
  args['visual_embeddings']=NDArray(ve);args['visual_mask']=NDArray(vm)
  result=await self.tf(args)
  pq,po,sq,sk,rq,rk,bias=[torch.from_numpy(result[k].numpy().copy()) for k in self.manifest['outputs']]
  Q=len(p.decide);O=len(b['o_q']);out={'opt':torch.full((Q,max(1,max(p.n_options or [0]))), -torch.inf)}
  for i in range(O):
   valid=b['option_of'][0]==i
   pooled=po[0,:valid.shape[0]][valid].mean(0)
   qi=int(b['o_q'][i]);rank=int(b['o_rank'][i]);at=p.decide[qi]
   out['opt'][qi,rank]=(pq[0,at]*pooled).sum()/math.sqrt(d)+(bias if p.types[qi]=='multi' else 0)
  spans=[i for i,t in enumerate(p.types) if t=='span'];regions=[i for i,t in enumerate(p.types) if t=='region']
  out['span_q']=torch.tensor(spans);out['region_q']=torch.tensor(regions)
  if spans:
   q=sq[0,p.decide].view(Q,2,d)[spans];k=sk[0,:b['ids'].shape[1]].view(-1,2,d)
   out['span']=torch.einsum('qcd,lcd->qcl',q,k)/math.sqrt(d)
   out['span'].masked_fill_(~b['state_mask'][:,None,:],-torch.inf)
  if regions:
   out['region']=rq[0,[p.decide[i] for i in regions]]@rk[0,:b['ids'].shape[1]].T/math.sqrt(d)
   out['region'].masked_fill_(~b['region_mask'],-torch.inf)
  return out,[p]

 async def predict(self,state,questions,images=()):
  from amos.model import decode
  out,ps=await self.raw(state,questions,images)
  return decode(out,ps)[0]

 async def request(self,req):
  from amos.serve import parse,answer
  from PIL import Image
  import base64,io
  state,ids,qs=parse(req)
  def image(s):return Image.open(io.BytesIO(base64.b64decode(s.split(',')[-1])))
  ims=[image(s) for s in req.get('images',[])]
  qs=[{**q,'options':[image(o['image']) if isinstance(o,dict) else o for o in q['options']]} if 'options' in q else q for q in qs]
  results=await self.predict(state,qs,ims)
  return {'answers':{i:answer(q,r) for i,q,r in zip(ids,qs,results)}}