"""MS MARCO MiniLM cross-encoder, loaded from uploaded weights. CPU, offline."""
import os,unicodedata,re,time
import numpy as np
import torch
from torch import nn
from safetensors.torch import load_file
from minilm_cpu_pipeline import BertEmb,Enc,WordPiece
from neural_evidence_graph_core import ROOT

class BertCrossEncoder(nn.Module):
 def __init__(self):
  super().__init__();self.embeddings=BertEmb();self.encoder=Enc()
  self.pooler=nn.Linear(384,384);self.classifier=nn.Linear(384,1)
 def forward(self,ids,mask,types):
  pos=self.embeddings.position_ids[:,:ids.shape[1]]
  h=self.embeddings.LayerNorm(self.embeddings.word_embeddings(ids)+self.embeddings.position_embeddings(pos)+self.embeddings.token_type_embeddings(types))
  h=self.encoder(h,mask[:,None,None,:].bool())
  return self.classifier(torch.tanh(self.pooler(h[:,0,:]))).squeeze(-1)

def load_crossencoder():
 torch.set_num_threads(min(4,os.cpu_count() or 4))
 try:torch.set_num_interop_threads(1)
 except RuntimeError:pass
 m=BertCrossEncoder();raw=load_file(str(ROOT/'model.safetensors'),device='cpu');clean={}
 for k,v in raw.items():
  if k=='bert.embeddings.position_ids':continue
  if k.startswith('bert.pooler.dense.'):k='pooler.'+k[len('bert.pooler.dense.'):]
  elif k.startswith('bert.'):k=k[5:]
  clean[k]=v
 m.load_state_dict(clean,strict=True);m.eval()
 print('Loaded original pretrained MS MARCO model, tensors',len(clean),flush=True)
 return m

class PairTokenizer(WordPiece):
 def __init__(self,length=256):super().__init__(str(ROOT/'vocab.txt'),length)
 def tokens(self,text):
  text=''.join(ch for ch in unicodedata.normalize('NFD',text.lower()) if unicodedata.category(ch)!='Mn' and (ch.isspace() or unicodedata.category(ch)[0]!='C'))
  out=[]
  for x in re.findall(r'[^\W_]+|_|[^\w\s]',text,flags=re.UNICODE):out.extend(self.wp(x))
  return out
 def pair(self,q,d):
  a=self.tokens(q);b=self.tokens(d);cap=self.max_length-3
  while len(a)+len(b)>cap:
   if len(b)>len(a):b.pop()
   else:a.pop()
  ids=[101]+a+[102]+b+[102]
  tt=[0]*(len(a)+2)+[1]*(len(b)+1)
  return ids,tt

@torch.inference_mode()
def score_pairs(m,tok,pairs,batch=20):
 out=np.empty(len(pairs),np.float32);times=[]
 for start in range(0,len(pairs),batch):
  p=pairs[start:start+batch];inputs=[tok.pair(q,d) for q,d in p]
  L=max(len(x[0]) for x in inputs)
  ids=np.zeros((len(p),L),np.int64);tt=np.zeros_like(ids);mask=np.zeros_like(ids)
  for i,(u,v) in enumerate(inputs):ids[i,:len(u)]=u;tt[i,:len(v)]=v;mask[i,:len(u)]=1
  t=time.perf_counter();out[start:start+len(p)]=m(torch.from_numpy(ids),torch.from_numpy(mask),torch.from_numpy(tt)).numpy();times.append(time.perf_counter()-t)
  if start and start%(batch*35)==0:print('scored',start,'/',len(pairs),flush=True)
 return out,{'total_forward_s':float(sum(times)),'forward_ms_per_pair':float(sum(times)*1000/max(1,len(pairs))),'batches':len(times)}