"""CPU-only evidence-completeness benchmark for HotpotQA distractor setting.
Uses frozen local MiniLM, no gold evidence in selection. Dataset's supporting facts evaluate chosen titles.
"""
import os,time,re,json,sys,hashlib
os.environ['OPENBLAS_NUM_THREADS']='1';os.environ['OMP_NUM_THREADS']='4'
from pathlib import Path
import numpy as np
from minilm_cpu_pipeline import load_model,WordPiece,embed
DATA=Path('/mnt/data/hotpot_dev_distractor_v1.json');OUT=Path('/mnt/data/hotpot_subset_out');OUT.mkdir(exist_ok=True)
N_TUNE=200;N_TEST=600;RNG=20261030


def facets(q):
 q=q.strip().rstrip('?')
 # Query-provided pieces only; no gold evidence or topic annotations.
 p=re.split(r'(?i)\b(?:and|or|whose|which|who|that|where|while|whereas)\b|[,;:]\s*',q)
 p=[x.strip() for x in p if len(x.split())>=2]
 if len(p)<2:
  w=q.split();m=max(2,len(w)//2)
  p=[' '.join(w[:m]),' '.join(w[m:])] if len(w)>=5 else [q,q]
 return (p[:3]+[q]*3)[:3]

def safe_key(x):return re.sub(r'\s+',' ',x.strip().lower())

def load_data():
 data=[r for r in json.load(open(DATA)) if len(r['context'])==10]; rng=np.random.default_rng(RNG)
 # Sample from all subtypes without using labels for feature engineering.
 ids=rng.permutation(len(data))[:N_TUNE+N_TEST]
 items=[data[int(i)] for i in ids]
 counts={}
 for r in items:counts[r['type']]=counts.get(r['type'],0)+1
 print('sample',len(items),'types',counts,flush=True)
 return items,ids

def encode_data(items):
 path=OUT/'hotpot_minilm_800.npz'
 identity=hashlib.sha256('|'.join(r['_id'] for r in items).encode()).hexdigest()[:16]
 if path.exists():
  dat=np.load(path)
  if str(dat['identity'])==identity:
   print('loaded cached embeddings',flush=True)
   return dat['q'],dat['docs'],dat['facets']
 model=load_model();tok=WordPiece(length=128)
 queries=[r['question'] for r in items]
 docs=[title+'. '+ ' '.join(sentences) for r in items for title,sentences in r['context']]
 ftexts=[t for r in items for t in facets(r['question'])]
 print('encoding',len(queries),'queries',len(docs),'docs',len(ftexts),'aspects',flush=True)
 st=time.perf_counter();q=embed(queries,model,tok,batch_size=24,callback=lambda i,n,t:print('q',i,n,round(t,1),flush=True) if i>0 else None)
 print('query done',round(time.perf_counter()-st,1),flush=True)
 d=embed(docs,model,tok,batch_size=24,callback=lambda i,n,t:print('docs',i,n,round(t,1),flush=True) if i>0 else None).reshape(len(items),10,384)
 print('documents done',round(time.perf_counter()-st,1),flush=True)
 f=embed(ftexts,model,tok,batch_size=24,callback=lambda i,n,t:print('facets',i,n,round(t,1),flush=True) if i>0 else None).reshape(len(items),3,384)
 print('facets done',round(time.perf_counter()-st,1),flush=True)
 np.savez_compressed(path,identity=identity,q=q,docs=d,facets=f)
 return q,d,f

def greedy(rel,dc,fac,method,k=5,alpha=.6,beta=.2,threshold=.35):
 # Single query score vector, 10x10 document similarity, facet-doc sim 3x10.
 K=len(rel);sel=[];remain=set(range(K));
 cfac=np.clip(fac-threshold,0,None)
 mx=np.zeros(len(cfac));mxcluster=np.zeros(K)
 for t in range(k):
  best=(-1e9,None)
  for j in sorted(remain):
   if method=='cosine':v=rel[j]
   elif method=='mmr':v=rel[j] if not sel else alpha*rel[j]-(1-alpha)*max(0,max(dc[j,x] for x in sel))
   elif method=='coverage':v=alpha*rel[j]+(1-alpha)*np.maximum(cfac[:,j]-mx,0).mean()
   elif method=='facility':
    weights=np.clip(rel,0,None)**2
    v=alpha*rel[j]+(1-alpha)*np.dot(weights,np.maximum(np.clip(dc[:,j],0,None)-mxcluster,0))/max(weights.sum(),1e-8)
   elif method=='bridge':
    bridge=max((dc[j,x] for x in sel),default=0)
    v=alpha*rel[j]+(1-alpha)*bridge-beta*max(0,bridge-.8) if sel else rel[j]
   elif method=='combined':
    gain=np.maximum(cfac[:,j]-mx,0).mean()
    # Complementary candidate information, plus query-conditioned facet coverage.
    weights=np.clip(rel,0,None)**2
    facility=np.dot(weights,np.maximum(np.clip(dc[:,j],0,None)-mxcluster,0))/max(weights.sum(),1e-8)
    v=alpha*rel[j]+(1-alpha)*(beta*gain+(1-beta)*facility)
   else:raise ValueError(method)
   if v>best[0]:best=(v,j)
  j=best[1];sel.append(j);remain.remove(j)
  mx=np.maximum(mx,cfac[:,j]);mxcluster=np.maximum(mxcluster,np.clip(dc[:,j],0,None))
 return sel

def eval_one(items,scores,docsim,facsim,method,params,inds,ks=(2,3,4,5)):
 out=np.zeros((len(inds),len(ks)),float); selected=[];goldcounts=[];omiss=[]
 for ii,i in enumerate(inds):
  r=items[i]; gold=set(safe_key(t) for t,_ in r['supporting_facts']);titles=[safe_key(x[0]) for x in r['context']]
  goldcounts.append(len(gold));possible=gold.issubset(set(titles));omiss.append(possible)
  order=greedy(scores[i],docsim[i],facsim[i],method,k=max(ks),**params)
  selected.append(order)
  for jj,k in enumerate(ks):out[ii,jj]=int(possible and gold.issubset(set(titles[x] for x in order[:k])))
 return out,np.asarray(selected),np.asarray(goldcounts),np.asarray(omiss)

def main():
 items,orig_idx=load_data();q,d,f=encode_data(items)
 sims=np.einsum('nd,nkd->nk',q,d).astype('float32');dcs=np.einsum('nkd,njd->nkj',d,d).astype('float32');fs=np.einsum('nmd,nkd->nmk',f,d).astype('float32')
 tunes=np.arange(N_TUNE);tests=np.arange(N_TUNE,N_TUNE+N_TEST)
 param_grid={
   'cosine':[{}],
   'mmr':[{'alpha':a} for a in [.3,.5,.7,.85,.95]],
   'coverage':[{'alpha':a,'threshold':t} for a in [.1,.3,.5,.7,.9] for t in [.2,.35,.5,.65]],
   'facility':[{'alpha':a} for a in [.2,.4,.6,.8,.95]],
   'bridge':[{'alpha':a,'beta':b} for a in [.4,.6,.8] for b in [0,.2,.4]],
   'combined':[{'alpha':a,'beta':b,'threshold':t} for a in [.2,.4,.6,.8] for b in [.25,.5,.75] for t in [.3,.5]],
 }
 best={};val={}
 for method,grid in param_grid.items():
  st=time.perf_counter();records=[]
  for params in grid:
   arr,_,_,_=eval_one(items,sims,dcs,fs,method,params,tunes)
   records.append((float(arr[:,1].mean()),float(arr[:,2].mean()),params))
  records.sort(key=lambda x:(x[0],x[1]),reverse=True)
  best[method]=records[0][2];val[method]=records[0][0]
  print('TUNED',method,'@3',records[0][0],best[method],'seconds',round(time.perf_counter()-st,1),flush=True)
 results={};test_arrays={};selections={}
 for method,params in best.items():
  st=time.perf_counter();arr,idx,ngold,oracle=eval_one(items,sims,dcs,fs,method,params,tests)
  ms=(time.perf_counter()-st)/len(tests)*1000
  results[method]={'parameters':params,'validation_complete_evidence_at3':val[method],'heldout_complete_evidence_at_k':{str(k):float(arr[:,j].mean()) for j,k in enumerate((2,3,4,5))},'selection_cpu_ms_per_question':ms,'gold_title_count_distribution':{str(x):int((ngold==x).sum()) for x in np.unique(ngold)},'oracle_coverage_rate':float(oracle.mean())}
  test_arrays[method]=arr;selections[method]=idx
  print('TEST',method,results[method]['heldout_complete_evidence_at_k'],'ms',round(ms,3),flush=True)
 base=test_arrays['cosine'];rng=np.random.default_rng(RNG)
 bs=rng.integers(len(tests),size=(2000,len(tests)))
 for method in best:
  delta=test_arrays[method][:,1]-base[:,1]
  ci=np.quantile(delta[bs].mean(axis=1),[.025,.975]).tolist()
  results[method]['delta_at3_vs_cosine']=float(delta.mean());results[method]['delta_at3_ci95']=ci
  results[method]['paired_win_loss_at3']={'wins':int((delta>0).sum()),'losses':int((delta<0).sum()),'ties':int((delta==0).sum())}
 ids=[items[i]['_id'] for i in tests];details=[]
 for j,i in enumerate(tests):
  r=items[i]; details.append({'id':r['_id'],'question':r['question'],'answer':r['answer'],'type':r['type'],'gold_support_titles':sorted(set(t for t,_ in r['supporting_facts'])),'cosine': [r['context'][v][0] for v in selections['cosine'][j]],'coverage':[r['context'][v][0] for v in selections['coverage'][j]],'combined':[r['context'][v][0] for v in selections['combined'][j]]})
 result={'dataset':'HotpotQA dev distractor','source_file':str(DATA),'total_dev':7405,'sample_count':len(items),'tune_count':N_TUNE,'heldout_count':N_TEST,'random_seed':RNG,'protocol':'Each question has own ten provided paragraphs; gold supporting passage titles are evaluation only; frozen MiniLM encoded whole query, each candidate passage, and three heuristic query facets; hyperparameters chosen on 250 dev cases; tested on separate 750 dev cases. No generative LLM was invoked.','results':results}
 (OUT/'hotpot_results.json').write_text(json.dumps(result,indent=2));(OUT/'hotpot_case_studies.json').write_text(json.dumps(details,indent=2));print('DONE',flush=True)
if __name__=='__main__':main()