"""AMOS Qwen3.5 Core AI adapter. No PyTorch neural model is loaded."""
import json,sys
from pathlib import Path
import numpy as np
import torch
from coreai.runtime import AIModel,NDArray,SpecializationOptions
from transformers import AutoTokenizer,Qwen2VLImageProcessorPil

class AmosDecoder:
 @classmethod
 async def load(cls,bundle):
  self=cls();self.root=Path(bundle).resolve();sys.path.insert(0,str(self.root/'serve'))
  import amos.decoder as decoder
  decoder.IMAGE_PIXELS=224*224
  self.manifest=json.loads((self.root/'manifest.json').read_text())
  self.codec=decoder.Codec(AutoTokenizer.from_pretrained(self.root/'model'),Qwen2VLImageProcessorPil.from_pretrained(self.root/'model',size={'shortest_edge':224*224,'longest_edge':224*224}),255)
  self.codec.echo=self.manifest['echo']
  self.temperature=decoder.calibration(self.root) or {};self.span=decoder.span_decoding(self.root)
  self.models=[];self.functions={}
  names=['vision','heads']
  if self.manifest.get('text_layout')=='layer-sharded':names+=['embedding','finish']+[r['asset'].removesuffix('.aimodel') for r in self.manifest['layers']]
  else:names+=['text']
  for name in names:
   m=await AIModel.load(self.root/(name+'.aimodel'),SpecializationOptions.default());self.models.append(m);self.functions[name]=m.load_function('main')
  return self

 async def raw(self,state,questions,images=()):
  from amos.decoder import expand
  from PIL import Image
  def resize(im):return im.convert('RGB').resize((224,224),Image.Resampling.BICUBIC)
  images=[resize(im) for im in images]
  questions=[{**q,'options':[resize(o) if isinstance(o,Image.Image) else o for o in q['options']]} if 'options' in q else q for q in questions]
  flat,owner=expand(questions)
  ids,at,extra=self.codec.encode_case(state,flat,images)
  n=len(ids);L=self.manifest['sequence_length'];d=self.manifest['hidden_size']
  if n>L:raise ValueError(f'Request needs {n} tokens; Core AI capacity is {L}.')
  args={'input_ids':np.full((1,L),self.codec.pad,np.int32),'attention_mask':np.zeros((1,L),np.int32),'position_ids':np.zeros((4,1,L),np.int32),'visual_embeddings':np.zeros((1,L,d),np.float32),'visual_mask':np.zeros((1,L,1),bool)}
  args['input_ids'][0,:n]=ids;args['attention_mask'][0,:n]=1;args['position_ids'][0,0,:n]=np.arange(n)
  image_at={s:(rows,cols) for s,rows,cols in extra.get('images',[])}
  features=[]
  if image_at:
   pv=extra['pixel_values'].numpy()
   for start in range(0,len(pv),196):
    out=await self.functions['vision']({'pixel_values':NDArray(pv[start:start+196].astype(np.float32))});features.append(out[self.manifest['vision_output']].numpy())
  cursor=0;offset=0;fi=0
  while cursor<n:
   if cursor in image_at:
    rows,cols=image_at[cursor]
    if (rows,cols)!=(7,7):raise ValueError('Only 224x224 image grids are supported')
    end=cursor+49;args['position_ids'][1,0,cursor:end]=offset;args['position_ids'][2,0,cursor:end]=offset+np.repeat(np.arange(7),7);args['position_ids'][3,0,cursor:end]=offset+np.tile(np.arange(7),7)
    args['visual_embeddings'][0,cursor:end]=features[fi];args['visual_mask'][0,cursor:end]=True;offset+=7;cursor=end;fi+=1
   else:
    args['position_ids'][1:,0,cursor]=offset;offset+=1;cursor+=1
  if self.manifest.get('text_layout')=='layer-sharded':
   result=await self.functions['embedding']({k:NDArray(v) for k,v in args.items() if k!='attention_mask'})
   h,cos,sin=[result[k] for k in self.manifest['embedding_outputs']]
   causal=(np.arange(L)[:,None]>=np.arange(L)[None,:])[None,None]&args['attention_mask'][:,None,None,:].astype(bool)
   for layer in self.manifest['layers']:
    result=await self.functions[layer['asset'].removesuffix('.aimodel')]({'hidden_states':h,'cos':cos,'sin':sin,'position_ids':NDArray(args['position_ids'][0].copy()),'attention_mask':NDArray(causal)})
    h=result[layer['output']]
   output=await self.functions['finish']({'hidden_states':h})
  else:output=await self.functions['text']({k:NDArray(v) for k,v in args.items()})
  hidden,logits=[output[k].numpy() for k in self.manifest['outputs']]
  head_results={}
  for i,(q,pos) in enumerate(zip(flat,at)):
   if q['type'] not in ['span','region']:continue
   first,offsets=extra['state'];state_hidden=np.zeros((L,d),np.float32);state_hidden[:len(offsets)]=hidden[0,first:first+len(offsets)]
   cells=np.zeros((49,d),np.float32)
   if extra.get('images'):
    start,rows,cols=extra['images'][0];cells=hidden[0,start:start+49].copy()
   if q['type']=='region' and not extra.get('images'):raise ValueError('Region requires a state image')
   ho=await self.functions['heads']({'question_hidden':NDArray(hidden[0,pos].copy()),'state_hidden':NDArray(state_hidden),'image_hidden':NDArray(cells)})
   span,region,box=[torch.from_numpy(ho[k].numpy().copy()) for k in self.manifest['head_outputs']]
   head_results[i]=span[:,:len(offsets)] if q['type']=='span' else torch.cat([region,box]) if self.manifest['has_box'] else region
  return {'flat':flat,'owner':owner,'at':at,'extra':extra,'logits':logits[0],'heads':head_results,'tokens':n,'hidden':hidden[0]}

 async def predict(self,state,questions,images=()):
  from amos.decoder import options,decode_head
  raw=await self.raw(state,questions,images);out=[None]*len(questions)
  for j,(q,k,pos) in enumerate(zip(raw['flat'],raw['owner'],raw['at'])):
   typ=questions[k]['type']
   if q['type'] in ['span','region']:
    out[k]=decode_head(q,raw['heads'][j],raw['extra'],state,self.manifest['keep'],self.span);continue
   logits=torch.tensor(raw['logits'][pos,:len(options(q))]);p=(logits/self.temperature.get(q['type'],1)).softmax(-1).tolist()
   if typ=='multi':out[k]=out[k] or {'type':typ,'probs':[]};out[k]['probs'].append(p[0])
   elif typ=='noul':out[k]={'type':typ,'p_yes':p[0]}
   else:
    out[k]={'type':typ,'probs':p,'choice':int(np.argmax(p))}
    if typ=='score':out[k]['expected']=sum(i*v for i,v in enumerate(p))
  return out

 async def request(self,req):
  from amos.serve import parse,answer
  from PIL import Image
  import base64,io
  def image(s):return Image.open(io.BytesIO(base64.b64decode(s.split(',')[-1]))).convert('RGB')
  state,ids,qs=parse(req);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]
  rs=await self.predict(state,qs,ims)
  return {'answers':{i:answer(q,r) for i,q,r in zip(ids,qs,rs)}}