"""Pretrained MS MARCO MiniLM-L6 cross-encoder CPU-only held-out evaluation.
Inputs are user-uploaded safetensors/tokenizer and existing HotpotQA / 2Wiki candidate pools.
A fixed 200/600 split (seed 20261030) is reused from previous experiments.
No gold labels used for the cross-encoder, and none fed to selection.
"""
import os,time,json,argparse,statistics,unicodedata,re
os.environ['OPENBLAS_NUM_THREADS']='1';os.environ['OMP_NUM_THREADS']='4'
from pathlib import Path
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from safetensors.torch import load_file
from minilm_cpu_pipeline import BertEmb,Enc,WordPiece
from hotpot_subset_eval import load_data as hotpot_data
from wiki_subset_eval import load_data as wiki_data
from graph_bridge_eval import prep_links,graph_select

ROOT=Path('/mnt/data');RNG=20261030

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,token_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(token_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)); torch.set_num_interop_threads(1)
 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
 missing,unexpected=m.load_state_dict(clean,strict=True)
 assert not missing and not unexpected,(missing,unexpected)
 m.eval();print('Loaded pretrained MS MARCO cross-encoder: 106 tensors, all parameters match',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=16,limit=None):
 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(f' scored {start}/{len(pairs)} pairs',flush=True)
 return out,{'total_forward_s':float(sum(times)),'forward_ms_per_pair':float(sum(times)*1000/max(len(pairs),1)),'batches':len(times)}

def coverage(scores,items,k):
 y=np.zeros((len(items),10),np.uint8)
 for i,r in enumerate(items):
  gold={x.casefold().strip() for x,_ in r['supporting_facts']}
  for j,(title,_) in enumerate(r['context']):y[i,j]=int(title.casefold().strip() in gold)
 orders=np.argsort(-scores,axis=1,kind='stable')
 return np.asarray([int(y[i,orders[i,:k]].sum()==y[i].sum() and y[i].sum()>0) for i in range(len(items))]),orders

def run(ds,m,tok,batch=16,max_q=None):
 start=time.perf_counter()
 items,_=(hotpot_data() if ds=='hotpot' else wiki_data())
 test=items[200:800];n=len(test)
 if max_q is not None:n=min(n,max_q);test=test[:n]
 pairs=[(r['question'],title+'. '+ ' '.join(sentences)) for r in test for title,sentences in r['context']]
 print(f'{ds}: scoring {n} heldout questions, {len(pairs)} pairs, seq_limit={tok.max_length}',flush=True)
 ts=time.perf_counter();ranks,perf=score_pairs(m,tok,pairs,batch=batch)
 ts=time.perf_counter()-ts; scores=ranks.reshape(n,10)
 emb=np.load(ROOT/('hotpot_subset_out/hotpot_minilm_800.npz' if ds=='hotpot' else 'wiki_subset_out/wiki_minilm_800.npz'))
 Q=emb['q'][200:200+n];D=emb['docs'][200:200+n];Fq=emb['facets'][200:200+n]
 cosine=np.einsum('nd,nkd->nk',Q,D)
 base3,b_order=coverage(cosine,test,3)
 report=json.load(open(ROOT/('hotpot_subset_out/graph_results.json' if ds=='hotpot' else 'wiki_subset_out/graph_results.json')))
 link,match=prep_links(test);docsim=np.einsum('nkd,njd->nkj',D,D); facsim=np.einsum('nmd,nkd->nmk',Fq,D)
 gorders=np.asarray([graph_select(cosine[i],docsim[i],facsim[i],link[i],match[i],k=5,**report['parameters']) for i in range(n)])
 truth=np.zeros((n,10),np.uint8)
 for i,r in enumerate(test):
  gold={x.casefold().strip() for x,_ in r['supporting_facts']}
  for j,(title,_) in enumerate(r['context']):truth[i,j]=int(title.casefold().strip() in gold)
 cross_order=np.argsort(-scores,axis=1,kind='stable')
 def hit(order,k):return np.array([int(truth[i,order[i,:k]].sum()==truth[i].sum() and truth[i].sum()>0) for i in range(n)],np.int8)
 methods={'cosine':b_order,'cross_encoder':cross_order,'evidence_graph':gorders}
 metrics={name:{str(k):float(hit(order,k).mean()) for k in (1,2,3,4,5)} for name,order in methods.items()}
 rng=np.random.default_rng(9027);boots=rng.integers(0,n,size=(3000,n));cis={}
 a=hit(cross_order,3)
 for name,order in methods.items():
  if name=='cross_encoder':continue
  delta=a-hit(order,3)
  cis[name]={'difference':float(delta.mean()),'ci95':np.quantile(delta[boots].mean(axis=1),[.025,.975]).tolist(),'wins':int((delta>0).sum()),'losses':int((delta<0).sum()),'ties':int((delta==0).sum())}
 result={'dataset':ds,'questions':n,'pairs':n*10,'model':'cross-encoder/ms-marco-MiniLM-L6-v2 pretrained original weights','max_wordpieces':tok.max_length,'batch_size':batch,'metrics_complete_gold_doc_coverage':metrics,'cross_encoder_minus':cis,'timing':{**perf,'total_scoring_seconds_with_tokenization':ts,'scoring_ms_per_question':ts*1000/n,'total_run_s':time.perf_counter()-start},'protocol':'Existing seeded held-out questions 200:800 (200 prior tune questions excluded). Same ten candidates. Gold labels used only to measure coverage. Graph selector retains previously tuned parameters.'}
 out=ROOT/f'crossencoder_{ds}_{tok.max_length}_{n}_results.json';out.write_text(json.dumps(result,indent=2))
 np.savez_compressed(ROOT/f'crossencoder_{ds}_{tok.max_length}_{n}_scores.npz',scores=scores,orders=cross_order)
 print('RESULT',json.dumps(result,indent=2),flush=True)
 return result

if __name__=='__main__':
 ap=argparse.ArgumentParser();ap.add_argument('--datasets',nargs='+',default=['hotpot','wiki']);ap.add_argument('--max-length',type=int,default=256);ap.add_argument('--batch-size',type=int,default=16);ap.add_argument('--max-queries',type=int,default=None)
 args=ap.parse_args();model=load_crossencoder();tokenizer=PairTokenizer(length=args.max_length)
 for ds in args.datasets:run(ds,model,tokenizer,args.batch_size,args.max_queries)