#!/usr/bin/env python3
import collections,hashlib,json,math,pathlib,re,sys
from baselines import TOK,aggregate
root=pathlib.Path(sys.argv[1]);f=json.load(open(root/'scifact-fixtures.json'));docs=f['documents'];toks=[TOK(d['text']) for d in docs];freq=[collections.Counter(t) for t in toks];df=collections.Counter(w for t in toks for w in set(t));n=len(docs);avg=sum(map(len,toks))/n
rankings={}
for q in f['queries']:
 words=set(TOK(q['query']));scores=[]
 for i,t in enumerate(freq):
  score=0
  for w in words:
   tf=t.get(w,0)
   if tf:score+=math.log(1+(n-df[w]+.5)/(df[w]+.5))*tf*2.2/(tf+1.2*(.25+.75*len(toks[i])/avg))
  if score>0:scores.append((score,i))
 rankings[q['id']]=[docs[i]['id'] for _,i in sorted(scores,key=lambda x:(-x[0],x[1]))[:5]]
capy={x['id']:x['ranked_ids'] for x in json.load(open(root/'evidence/capy-scifact.json'))}
assert len(capy)==50 and set(capy)==set(rankings)
summary={'corpus_documents':len(docs),'selected_queries':len(f['queries']),'scope':'retrieval only; full SciFact corpus; selected 50 of 300 official test queries; not full BEIR score','methods':{'capy_recall':aggregate(capy,f['queries']),'bm25':aggregate(rankings,f['queries'])}}
(root/'evidence/scifact-summary.json').write_text(json.dumps(summary,indent=2));(root/'evidence/scifact-baseline-rankings.json').write_text(json.dumps(rankings,indent=2))
rows=json.load(open(root/'evidence/routing-predictions.json'));routing=json.load(open(root/'evidence/routing-summary.json'));routing['counts']={}
for m in ['first_option','always_none','bm25_or_none']:
 routing['counts'][m]={'correct':sum(x[m]==x['gold'] for x in rows),'unnecessary_selection_on_none':sum(x['gold']=='none' and x[m]!='none' for x in rows),'missed_required_selection':sum(x['gold']!='none' and x[m]=='none' for x in rows),'wrong_non_none_selection':sum(x['gold']!='none' and x[m]!='none' and x[m]!=x['gold'] for x in rows)}
(root/'evidence/routing-summary.json').write_text(json.dumps(routing,indent=2))
# Recompute synthetic aggregate independently from raw rankings without aggregate().
qdata=json.load(open(root/'rag-fixtures.json'));raw=json.load(open(root/'evidence/capy-retrieval.json'));qby={q['id']:q for q in qdata['queries']};check={}
for group in ['exact','inflection','paraphrase']:
 rr=[x for x in raw if qby[x['id']]['category']==group];correct1=sum(bool(x['ranked_ids']) and x['ranked_ids'][0] in qby[x['id']]['relevant'] for x in rr);correct5=sum(bool(set(x['ranked_ids'])&set(qby[x['id']]['relevant'])) for x in rr);check[group]={'n':len(rr),'hits1':correct1,'hits5':correct5}
(root/'evidence/independent-count-check.json').write_text(json.dumps(check,indent=2));print(json.dumps(summary));print(json.dumps(routing['counts']));print(json.dumps(check))