"""Neural-only evaluation for three user-uploaded RAG/evidence datasets.
All ranking uses original MiniLM cosine / pretrained MS-MARCO cross encoder.
Graph is a fixed previously specified directed-link/title selection policy; gold is
never accessed by scoring/selection. QASC has TWO explicitly separate protocols.
"""
import os
os.environ['OPENBLAS_NUM_THREADS']='1'; os.environ['MKL_NUM_THREADS']='1'; os.environ['OMP_NUM_THREADS']='4'
import json, time, struct, collections, sys
from pathlib import Path
import numpy as np
import torch
from safetensors.torch import load_file
from minilm_cpu_pipeline import MiniLM, WordPiece, embed
from neural_crossencoder_only import load_crossencoder,PairTokenizer,score_pairs
from neural_evidence_graph_core import (hover,graph_arrays,graph_select,seeded_sample,doc,ROOT)
from parquet_lite import read_col,reconstruct
from parquet_probe import parse

K=(2,3,5,8)

def get_dense():
 m=MiniLM();sd=load_file(str(ROOT/'original_minilm/model.safetensors'))
 for key in ('embeddings.position_ids','pooler.dense.bias','pooler.dense.weight'):sd.pop(key)
 m.load_state_dict(sd,strict=True);m.eval();torch.set_num_threads(4)
 return m,WordPiece(str(ROOT/'vocab.txt'),length=128)

def embed_texts(texts,m,tok,batch=32,label=''):
 t=time.perf_counter();out=embed(texts,m,tok,batch_size=batch)
 print('DENSE_EMBED',label,len(texts),'seconds',round(time.perf_counter()-t,3),flush=True)
 return out,time.perf_counter()-t

def get_qasc():
 p=ROOT/'train-00000-of-00001.parquet';f=open(p,'rb');f.seek(-8,2)
 size=struct.unpack('<I',f.read(4))[0];f.seek(-8-size,2);meta=parse(f.read(size));rows=[]
 for rg in meta[4]:
  cols={}
  for c in rg[1]:
   key=tuple(c[3][3])
   if key not in [('id',),('question',),('fact1',),('fact2',)]:continue
   cols[key]=reconstruct(*read_col(f,c,{key:(1,0)}),1,0)
  rows.extend([{'id':cols[('id',)][i], 'question':cols[('question',)][i], 'facts':[cols[('fact1',)][i],cols[('fact2',)][i]]} for i in range(rg[3])])
 rows=[r for r in rows if all(x and str(x).strip() for x in r['facts']) and r['facts'][0]!=r['facts'][1]]
 facts=list(dict.fromkeys(z.strip() for r in rows for z in r['facts']))
 return rows,facts

def prepare_multihop(n,m,tok):
 t=time.perf_counter();corpus=json.load(open(ROOT/'corpus.json'));allq=json.load(open(ROOT/'MultiHopRAG.json'))
 eligible=[r for r in allq if len({e['url'] for e in r['evidence_list']})>=2]
 questions=seeded_sample(eligible,n)
 texts=[r['title']+'. '+r['body'][:1800] for r in corpus]
 de,tidx=embed_texts(texts,m,tok,label='MultiHop corpus')
 qe,tq=embed_texts([r['query'] for r in questions],m,tok,label='MultiHop queries')
 sim=qe@de.T;inds=np.argsort(-sim,axis=1,kind='stable')[:,:20]
 rows=[];sc=[]
 for i,r in enumerate(questions):
  docs=[doc(corpus[j]['title'],corpus[j]['body'],corpus[j]['url']) for j in inds[i]]
  rows.append({'qid':f'multihop:{i}', 'query':r['query'],'docs':docs,'gold':{e['url'] for e in r['evidence_list']},'pool':'Top-20 neural MiniLM retrieval from full 609-article corpus (no gold insertion)'})
  sc.append(sim[i,inds[i]])
 return rows,sc,{'index_seconds':tidx,'query_embedding_seconds':tq,'total_prep_seconds':time.perf_counter()-t,'index_size':len(corpus),'candidate_pool_size':20}

def prepare_hover(n,m,tok):
 t=time.perf_counter();rows=hover(n)
 # Exact 10 candidate passages; labels sourced from corresponding HoVer claims.
 alltexts=[r['query'] for r in rows]+[d['title']+'. '+d['body'] for r in rows for d in r['docs']]
 em,sec=embed_texts(alltexts,m,tok,label='HoVer queries+passages')
 Q=em[:len(rows)]; Ds=np.split(em[len(rows):],np.cumsum([len(r['docs']) for r in rows])[:-1]);sc=[Ds[i]@Q[i] for i in range(len(rows))]
 return rows,sc,{'embedding_seconds':sec,'prep_seconds':time.perf_counter()-t,'candidate_pool_size':10,'eligibility':'HoVer dev claims with HotpotQA dev passages available and all annotated support titles present; conditional subset'}

def prepare_qasc(n,m,tok):
 t=time.perf_counter();allrows,facts=get_qasc();qs=seeded_sample(allrows,n)
 fe,tidx=embed_texts(facts,m,tok,label='QASC all distinct supporting fact sentences')
 qe,tq=embed_texts([r['question'] for r in qs],m,tok,label='QASC query')
 sim=qe@fe.T;inds=np.argsort(-sim,axis=1,kind='stable')[:,:60]
 rows_native=[];scores_native=[];rows_ctrl=[];scores_ctrl=[]
 for i,r in enumerate(qs):
  gold={str(x).strip() for x in r['facts']}
  docs=[doc('',facts[int(j)],facts[int(j)]) for j in inds[i,:20]]
  rows_native.append({'qid':r['id'],'query':r['question'],'docs':docs,'gold':gold,'pool':'Top-20 MiniLM from all unique QASC supporting fact sentences, without gold insertion'})
  scores_native.append(sim[i,inds[i,:20]])
  neg=[int(j) for j in inds[i] if facts[int(j)] not in gold][:18]
  if len(neg)!=18:continue
  chosen=list(gold)+[facts[j] for j in neg]
  perm=np.random.default_rng(i+301).permutation(len(chosen))
  chosen=[chosen[int(j)] for j in perm]
  look={f:j for j,f in enumerate(facts)}
  rows_ctrl.append({'qid':r['id'],'query':r['question'],'docs':[doc('',z,z) for z in chosen],'gold':gold,'pool':'CONTROLLED: 2 gold facts deliberately inserted + 18 MiniLM hard negatives (oracle=100%)'})
  scores_ctrl.append(np.array([sim[i,look[z]] for z in chosen],np.float32))
 return (rows_native,scores_native),(rows_ctrl,scores_ctrl),{'index_seconds':tidx,'query_embedding_seconds':tq,'total_prep_seconds':time.perf_counter()-t,'fact_corpus_size':len(facts),'eligible_questions':len(allrows),'candidate_pool_size':20}

def assess(name,rows,dense_scores,cross_scores,loadtime,extra):
 tgraph=time.perf_counter();details=[]
 counts={method:{str(k):[] for k in K} for method in ('dense_cosine','crossencoder','dense_graph','crossencoder_graph')}
 oracle={str(k):[] for k in K};av=[];links=[];sizecnt=collections.Counter()
 for i,r in enumerate(rows):
  gold=r['gold'];docs=r['docs'];ss=np.asarray(dense_scores[i],dtype=np.float32);cc=np.asarray(cross_scores[i],dtype=np.float32)
  assert len(ss)==len(cc)==len(docs)
  link,match=graph_arrays(r);links.append(bool(link.any()))
  avail=len(gold&{d['id'] for d in docs});av.append(avail==len(gold));sizecnt[len(gold)]+=1
  orders={
   'dense_cosine':np.argsort(-ss,kind='stable').tolist(),
   'crossencoder':np.argsort(-cc,kind='stable').tolist(),
   'dense_graph':graph_select(ss,link,match,max(K)),
   'crossencoder_graph':graph_select(cc,link,match,max(K)),
  }
  record={'id':r['qid'],'gold_count':len(gold),'pool_gold_count':avail,'orders':orders,'hits':{}}
  for k in K:
   oracle[str(k)].append(int(avail==len(gold) and len(gold)<=k))
   for method,order in orders.items():
    hit=int(gold.issubset({docs[j]['id'] for j in order[:k]}))
    counts[method][str(k)].append(hit)
    record['hits'].setdefault(method,{})[str(k)]=hit
  details.append(record)
 graphseconds=time.perf_counter()-tgraph
 delta=np.array(counts['crossencoder_graph']['3'],int)-np.array(counts['crossencoder']['3'],int)
 if len(delta):
  gen=np.random.default_rng(31);ci=np.percentile(delta[gen.integers(0,len(delta),(2000,len(delta)))].mean(axis=1)*100,[2.5,97.5])
 else:ci=[float('nan')]*2
 res={'dataset':name,'questions':len(rows),'candidate_count_total':sum(len(r['docs']) for r in rows),'candidate_pool':rows[0]['pool'] if rows else '',
   'gold_support_count_distribution':dict(sizecnt),'gold_available_fraction':round(float(np.mean(av))*100,2),'gold_available_count':sum(av), 'graph_link_fraction':round(float(np.mean(links))*100,2),
   'oracle_complete_coverage_pct':{k:round(float(np.mean(v))*100,2) for k,v in oracle.items()},
   'complete_coverage_pct':{met:{k:round(float(np.mean(val))*100,2) for k,val in d.items()} for met,d in counts.items()},
   'crossencoder_graph_minus_crossencoder_at3_pp':round(float(delta.mean()*100),2),'ci95_pp':[round(float(z),2) for z in ci],
   'paired_wins_losses_ties_at3':{'wins':int((delta>0).sum()),'losses':int((delta<0).sum()),'ties':int((delta==0).sum())},
   'timing_seconds':{**extra,**loadtime,'graph_selection_total_s':round(graphseconds,3),'graph_selection_ms_per_question':round(graphseconds*1000/max(1,len(rows)),2)},
   'models':{'dense':'user uploaded all-MiniLM-L6-v2', 'crossencoder':'user uploaded cross-encoder/ms-marco-MiniLM-L6-v2'},
   'protocol':'MiniLM embeddings at max 128 WordPiece for dense. Cross-encoder truncates paired query+passage at 256 WordPiece. Both score exact same candidate pools. Fixed evidence graph .4 minmax relevance + .5 title in already-selected passage + .2 query-title score. No TF-IDF anywhere in candidate construction or scoring. Gold used solely for pool audit and evaluation.'}
 (ROOT/f'{name}_neural_only_results.json').write_text(json.dumps(res,indent=2))
 (ROOT/f'{name}_neural_only_per_question.json').write_text(json.dumps(details))
 print('RESULT',name,json.dumps(res),flush=True)
 return res

def evaluate(name,rows,denses,embtimes,ce,tok):
 pairs=[(r['query'],(d['title']+'. ' if d['title'] else '')+d['body']) for r in rows for d in r['docs']]
 print('CROSS_START',name,'n',len(rows),'pairs',len(pairs),flush=True)
 t=time.perf_counter();cross,perf=score_pairs(ce,tok,pairs,batch=20)
 tsec=time.perf_counter()-t
 lengths=np.cumsum([len(r['docs']) for r in rows])[:-1]
 scores=np.split(cross,lengths)
 report=assess(name,rows,denses,scores,{'crossencoder_scoring_seconds':round(tsec,3),'crossencoder_scoring_ms_per_question':round(tsec*1000/len(rows),2),'crossencoder_forward_seconds':round(perf['total_forward_s'],3)},embtimes)
 return report

if __name__=='__main__':
 n=int(sys.argv[1]) if len(sys.argv)>1 else 300
 ds=sys.argv[2:] or ['hover','multihop','qasc']
 dense,tok=get_dense();ce=load_crossencoder();ctok=PairTokenizer(length=256)
 for task in ds:
  if task=='hover':
   rows,scores,times=prepare_hover(n,dense,tok)
   evaluate('hover',rows,scores,times,ce,ctok)
  elif task=='multihop':
   rows,scores,times=prepare_multihop(n,dense,tok)
   evaluate('multihop',rows,scores,times,ce,ctok)
  elif task=='qasc':
   (rows,scores),(ctrl,cscores),times=prepare_qasc(n,dense,tok)
   evaluate('qasc_neural_retrieval',rows,scores,times,ce,ctok)
   evaluate('qasc_controlled_pool',ctrl,cscores,{'cached_fact_embeddings_used':True},ce,ctok)
  else:raise ValueError(task)