#!/usr/bin/env python3
"""Written for https://paraloncloud.com/resources/semantic-search-open-embeddings-api
Needs: pip install numpy scikit-learn requests; PRL_KEY=prlc_... in the environment;
support_tickets_990.jsonl and queries_300.json from the same /files/guides/ folder.

Semantic search over 990 support tickets with an open embedding model over an
OpenAI-compatible API, scored by intent agreement of the retrieved tickets.

Three measurements:
  A. leave-one-out: each ticket as the query against the other 989 -- does the
     nearest ticket share its intent? (precision@1, @5), with and without the
     E5 "query: " / "passage: " prefixes, against a TF-IDF keyword baseline.
  B. 300 short search queries in five languages (from gen_queries.py), run
     against the English corpus: precision@5 by language, embeddings vs TF-IDF.
  C. throughput and cost of embedding the corpus through the API.
"""
import json, os, sys, time, requests, numpy as np
from collections import defaultdict
from sklearn.feature_extraction.text import TfidfVectorizer
API="https://paraloncloud.com/v1/embeddings"; KEY=os.environ["PRL_KEY"]; MODEL=os.environ.get("MODEL","multilingual-e5-small")
PRICE_IN=0.02; BATCH=64
rows=[json.loads(l) for l in open("support_tickets_990.jsonl")]
docs=[r["customer_message"] for r in rows]; intents=[r["intent"] for r in rows]
stats={"tokens":0,"requests":0,"seconds":0.0}
def embed(texts):
    out=[]
    for i in range(0,len(texts),BATCH):
        t=time.time(); r=requests.post(API,headers={"Authorization":f"Bearer {KEY}"},json={"model":MODEL,"input":texts[i:i+BATCH]},timeout=120); dt=time.time()-t
        r.raise_for_status(); d=r.json(); stats["tokens"]+=d["usage"]["prompt_tokens"]; stats["requests"]+=1; stats["seconds"]+=dt
        out.extend(e["embedding"] for e in sorted(d["data"],key=lambda e:e["index"]))
    m=np.array(out,dtype=np.float32); return m/np.linalg.norm(m,axis=1,keepdims=True)
def prec_at(sims, qintents, k, exclude_self=False):
    hits1=hits5=0
    for i,row in enumerate(sims):
        if exclude_self: row=row.copy(); row[i]=-1
        top=np.argsort(-row)[:k]
        hits1+= intents[top[0]]==qintents[i]; hits5+= sum(intents[j]==qintents[i] for j in top)/k
    return round(100*hits1/len(sims),1), round(100*hits5/len(sims),1)
res={}
# --- C + A: corpus with and without prefix
t0=time.time(); P=embed(["passage: "+d for d in docs]); corpus_s=time.time()-t0
res["corpus"]={"docs":len(docs),"tokens":stats["tokens"],"seconds":round(corpus_s,1),"docs_per_s":round(len(docs)/corpus_s,1),"cost_usd":round(stats["tokens"]*PRICE_IN/1e6,5),"requests":stats["requests"]}
Q=embed(["query: "+d for d in docs])
res["A_prefixed"]=dict(zip(("p1","p5"),prec_at(Q@P.T,intents,5,True)))
NP=embed(docs)
res["A_noprefix"]=dict(zip(("p1","p5"),prec_at(NP@NP.T,intents,5,True)))
tf=TfidfVectorizer(sublinear_tf=True,ngram_range=(1,2),min_df=1).fit(docs); T=tf.transform(docs); TS=(T@T.T).toarray()
res["A_tfidf"]=dict(zip(("p1","p5"),prec_at(TS,intents,5,True)))
# --- B: multilingual short queries
qs=json.load(open("queries.json")); qtexts=[q["query"] for q in qs]; qint=[q["intent"] for q in qs]
QE=embed(["query: "+t for t in qtexts]); S=QE@P.T
TQ=(tf.transform(qtexts)@T.T).toarray()
bylang=defaultdict(lambda:{"n":0,"emb1":0,"emb5":0.0,"tf1":0,"tf5":0.0})
for i,q in enumerate(qs):
    b=bylang[q["lang"]]; b["n"]+=1
    top=np.argsort(-S[i])[:5]; b["emb1"]+= intents[top[0]]==qint[i]; b["emb5"]+= sum(intents[j]==qint[i] for j in top)/5
    topt=np.argsort(-TQ[i])[:5]; b["tf1"]+= intents[topt[0]]==qint[i]; b["tf5"]+= sum(intents[j]==qint[i] for j in topt)/5
res["B_by_lang"]={l:{"n":b["n"],"emb_p1":round(100*b["emb1"]/b["n"],1),"emb_p5":round(100*b["emb5"]/b["n"],1),"tfidf_p1":round(100*b["tf1"]/b["n"],1),"tfidf_p5":round(100*b["tf5"]/b["n"],1)} for l,b in bylang.items()}
# worst intents for the multilingual queries
byint=defaultdict(lambda:[0,0])
for i,q in enumerate(qs):
    top=np.argsort(-S[i])[:5]; byint[q["intent"]][0]+=sum(intents[j]==qint[i] for j in top)/5; byint[q["intent"]][1]+=1
res["B_by_intent_p5"]={k:round(100*v[0]/v[1],1) for k,v in sorted(byint.items(),key=lambda kv:kv[1][0]/kv[1][1])}
# examples: three queries with their top hit
ex=[]
for i in (0,7,14,21,28):
    top=int(np.argmax(S[i])); ex.append({"query":qtexts[i],"lang":qs[i]["lang"],"intent":qint[i],"top_intent":intents[top],"top_text":docs[top][:140],"sim":round(float(S[i][top]),3)})
res["examples"]=ex
res["totals"]={"tokens":stats["tokens"],"requests":stats["requests"],"cost_usd":round(stats["tokens"]*PRICE_IN/1e6,5),"api_seconds":round(stats["seconds"],1)}
print(json.dumps(res,indent=1,ensure_ascii=False)); json.dump(res,open("results.json","w"),indent=1,ensure_ascii=False)