#!/usr/bin/env python3
"""Original reference baselines. Stdlib; Unicode word tokens; BM25 k1=1.2,b=.75.
No labels are used for prediction. Stable ties by document/option order.
"""
import collections,json,math,pathlib,re,sys
TOK=lambda s:re.findall(r'[^\W_]+',s.lower(),re.UNICODE)
def rank_bm25(query,docs):
 toks=[TOK(d) for d in docs];freq=[collections.Counter(t) for t in toks];n=len(docs);avg=sum(map(len,toks))/n if n else 1;df=collections.Counter(w for t in toks for w in set(t));scores=[]
 for i,f in enumerate(freq):
  score=0
  for w in set(TOK(query)):
   if not f[w]:continue
   idf=math.log(1+(n-df[w]+.5)/(df[w]+.5));tf=f[w];score+=idf*tf*2.2/(tf+1.2*(.25+.75*len(toks[i])/avg))
  scores.append((score,i))
 return [i for score,i in sorted(scores,key=lambda x:(-x[0],x[1])) if score>0]
def aggregate(rows,queries):
 out={}
 for category in sorted(set(q['category'] for q in queries)):
  qs=[q for q in queries if q['category']==category];vals=[];false=0
  for q in qs:
   ranked=rows[q['id']];g=set(q['relevant']);positions=[i+1 for i,v in enumerate(ranked[:5]) if v in g]
   if g:vals.append((sum(x in g for x in ranked[:1])/len(g),sum(x in g for x in ranked[:5])/len(g),1/min(positions) if positions else 0))
   else:false+=bool(ranked)
  out[category]={'n':len(qs),'recall_at_1':sum(v[0] for v in vals)/len(vals) if vals else None,'recall_at_5':sum(v[1] for v in vals)/len(vals) if vals else None,'mrr_at_5':sum(v[2] for v in vals)/len(vals) if vals else None,'nonempty_unanswerable':false if not vals else None}
 return out
if __name__=='__main__':
 root=pathlib.Path(sys.argv[1]);f=json.load(open(root/'rag-fixtures.json'));docs=f['documents'];queries=f['queries'];bm={};exact={}
 for q in queries:
  bm[q['id']]=[docs[i]['id'] for i in rank_bm25(q['query'],[d['text'] for d in docs])[:5]]
  words=set(TOK(q['query']));exact[q['id']]=[d['id'] for d in docs if words<=set(TOK(d['text']))][:5]
 capy=json.load(open(root/'evidence/capy-retrieval.json'));out={'configuration':{'bm25_k1':1.2,'bm25_b':.75,'tokenization':'Unicode words, lowercase; no stemming or stopwords','capy_top_k':5,'no_llm':True},'corpus_documents':len(docs),'queries':len(queries),'methods':{'capy_recall':aggregate({x['id']:x['ranked_ids'] for x in capy},queries),'exact_token_and':aggregate(exact,queries),'bm25':aggregate(bm,queries)}}
 (root/'evidence/retrieval-summary.json').write_text(json.dumps(out,ensure_ascii=False,indent=2))
 (root/'evidence/baseline-rankings.json').write_text(json.dumps({'bm25':bm,'exact_token_and':exact},ensure_ascii=False,indent=2))
 # RuDecide skill-routing diagnostic, only public choice options and user message as input.
 rows=[json.loads(x) for x in (root/'track_b_applied.jsonl').read_text().splitlines()];rows=[x for x in rows if x['task']=='skills_unseen_catalogs'];pred=[]
 for row in rows:
  choices=row['question']['criteria'];keys=list(choices);rank=rank_bm25(row['state'],list(choices.values()));bmkey=keys[rank[0]] if rank else 'none';pred.append({'id':row['id'],'gold':row['answer'],'first_option':keys[0],'always_none':'none','bm25_or_none':bmkey,'has_lexical_match':bool(rank),'choice_count':len(keys)})
 summary={'n':len(pred),'split':'track_b_applied / skills_unseen_catalogs, complete 300 rows','learned_models_run':False,'accuracy':{m:sum(p[m]==p['gold'] for p in pred)/len(pred) for m in ['first_option','always_none','bm25_or_none']},'unmatched_query_count':sum(not p['has_lexical_match'] for p in pred),'none_gold_count':sum(p['gold']=='none' for p in pred),'methods':'Same tokenization/BM25; each choice description is a document, no training, no translation, fallback none only if all scores zero.'}
 (root/'evidence/routing-predictions.json').write_text(json.dumps(pred,ensure_ascii=False,indent=2));(root/'evidence/routing-summary.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2));print(json.dumps(out,ensure_ascii=False));print(json.dumps(summary,ensure_ascii=False))