#!/usr/bin/env python3
"""Document classification with an open LLM over an OpenAI-compatible API,
scored on 20 Newsgroups (bydate test split) or any folder-per-class tree.
Conditions: zero-shot one document per request (zero1); zero-shot ten per
request (zero10); few-shot, two examples per class, ten per request (few10).

  PRL_KEY=prlc_... MODE=zero10 N=7532 RPM=20 IN_FLIGHT=2 python3 classify_20news.py

Written for https://paraloncloud.com/resources/document-classification-open-llm-api
"""
import json, os, re, sys, time, random, threading, queue, requests
from collections import Counter, defaultdict
API="https://paraloncloud.com/v1/chat/completions"; MODEL=os.environ.get("MODEL","qwen3.8-27b"); KEY=os.environ["PRL_KEY"]
MODE=os.environ.get("MODE","zero1")          # zero1 | zero10 | few10
N=int(os.environ.get("N","2000")); IN_FLIGHT=int(os.environ.get("IN_FLIGHT","2")); RPM=int(os.environ.get("RPM","20"))   # free key: 20 rpm
ROOT=os.environ.get("ROOT","20news-bydate-test"); TRAIN=os.environ.get("TRAIN","20news-bydate-train")
PRICE_IN,PRICE_OUT=0.12,1.70   # USD per 1M tokens, qwen3.8-27b list price
CLASSES=sorted(os.listdir(ROOT))
def clean(txt):
    body=txt.split("\n\n",1)[1] if "\n\n" in txt else txt      # drop the header block
    lines=[l for l in body.split("\n") if not l.lstrip().startswith(">") and not re.match(r"^\S+ writes:$", l.strip()) and "wrote:" not in l[:80]]
    body="\n".join(lines); body=body.split("\n--\n")[0]
    words=body.split(); return " ".join(words[:400])
def load(split,n_per_class,seed):
    rnd=random.Random(seed); docs=[]
    for c in CLASSES:
        files=sorted(os.listdir(f"{split}/{c}")); rnd.shuffle(files)
        for f in files[:n_per_class]:
            t=clean(open(f"{split}/{c}/{f}",encoding="latin-1").read())
            if len(t.split())>=20: docs.append({"id":f"{c}/{f}","label":c,"text":t})
    rnd.shuffle(docs); return docs
docs=load(ROOT, 10**6 if N>=7532 else max(1,N//20+2), 11)[:N]
fewshot=""
if MODE=="few10":
    ex=load(TRAIN,2,3); by=defaultdict(list)
    for d in ex: by[d["label"]].append(" ".join(d["text"].split()[:60]))
    fewshot="Examples of each category:\n"+"\n".join(f"[{c}] {by[c][0]} ... | [{c}] {by[c][1]} ..." for c in CLASSES if len(by[c])>=2)+"\n\n"
SYSTEM=("You classify Usenet posts from 1993 into exactly one of these 20 newsgroups: "+", ".join(CLASSES)+
        ". Judge by topic, not by tone. Answer with the category id only, via the schema.")
per=1 if MODE=="zero1" else 10
def schema(k):
    item={"type":"string","enum":CLASSES}
    if k==1: return {"type":"object","properties":{"category":item},"required":["category"]}
    return {"type":"object","properties":{"categories":{"type":"array","minItems":k,"maxItems":k,"items":item}},"required":["categories"]}
lock=threading.Lock(); stats={"req":0,"err":0,"in":0,"out":0,"lat":[]}; preds={}
def one(batch):
    if per==1: user=fewshot+"Post:\n"+batch[0]["text"]
    else: user=fewshot+"Classify each of the following posts, in order. Return one category per post.\n\n"+"\n\n".join(f"### Post {i+1}\n{d['text']}" for i,d in enumerate(batch))
    body={"model":MODEL,"messages":[{"role":"system","content":SYSTEM},{"role":"user","content":user}],
          "response_format":{"type":"json_schema","json_schema":{"name":"cls","schema":schema(len(batch))}},
          "chat_template_kwargs":{"enable_thinking":False},"temperature":0,"max_tokens":400}
    t0=time.time()
    try: r=requests.post(API,headers={"Authorization":f"Bearer {KEY}"},json=body,timeout=240)
    except Exception as e:
        with lock: stats["req"]+=1; stats["err"]+=1
        print("EXC",e,file=sys.stderr); return
    lat=time.time()-t0
    with lock: stats["req"]+=1; stats["lat"].append(lat)
    if r.status_code!=200:
        with lock: stats["err"]+=1
        print("HTTP",r.status_code,r.text[:150],file=sys.stderr); return
    d=r.json(); u=d.get("usage",{})
    if d["choices"][0].get("finish_reason")=="length":
        with lock: stats["err"]+=1
        print("TRUNCATED: raise max_tokens",file=sys.stderr); return
    with lock: stats["in"]+=u.get("prompt_tokens",0); stats["out"]+=u.get("completion_tokens",0)
    try:
        o=json.loads(d["choices"][0]["message"]["content"]); cats=[o["category"]] if per==1 else o["categories"]
    except Exception as e:
        with lock: stats["err"]+=1
        print("PARSE",e,file=sys.stderr); return
    with lock:
        for doc,c in zip(batch,cats): preds[doc["id"]]=c
q=queue.Queue()
for i in range(0,len(docs),per): q.put(docs[i:i+per])
start=time.time()
def worker():
    while True:
        try: b=q.get_nowait()
        except queue.Empty: return
        with lock: ahead=stats["req"]+IN_FLIGHT-((time.time()-start)/60*RPM+1)
        if ahead>0: time.sleep(min(5,ahead*60/RPM))
        one(b)
th=[threading.Thread(target=worker) for _ in range(IN_FLIGHT)]; [t.start() for t in th]; [t.join() for t in th]; wall=time.time()-start
scored=[(d["label"],preds.get(d["id"])) for d in docs]; n=len(scored); correct=sum(1 for a,b in scored if a==b); missing=sum(1 for a,b in scored if b is None)
tp=Counter(); fp=Counter(); fn=Counter()
for a,b in scored:
    if a==b: tp[a]+=1
    else: fn[a]+=1; fp[b]+=1
f1s=[]
for c in CLASSES:
    p=tp[c]/max(1,tp[c]+fp[c]); r_=tp[c]/max(1,tp[c]+fn[c]); f1s.append(2*p*r_/max(1e-9,p+r_))
per_class={c:round(tp[c]/max(1,tp[c]+fn[c])*100,1) for c in CLASSES}
conf=Counter((a,b) for a,b in scored if a!=b and b)
lat=sorted(stats["lat"]); pct=lambda x: round(lat[min(len(lat)-1,int(len(lat)*x))],2) if lat else None
m={"mode":MODE,"docs":n,"missing":missing,"accuracy_pct":round(100*correct/n,2),"macro_f1_pct":round(100*sum(f1s)/len(f1s),2),
   "requests":stats["req"],"errors":stats["err"],"tokens_in":stats["in"],"tokens_out":stats["out"],
   "tokens_in_per_doc":round(stats["in"]/n,1),"tokens_out_per_doc":round(stats["out"]/n,2),
   "latency_p50_s":pct(0.5),"latency_p90_s":pct(0.9),"wall_s":round(wall,1),"docs_per_hour":round(n/wall*3600),
   "cost_usd":round((stats["in"]*PRICE_IN+stats["out"]*PRICE_OUT)/1e6,4),"cost_per_1k_docs_usd":round((stats["in"]*PRICE_IN+stats["out"]*PRICE_OUT)/1e6/n*1000,4),
   "per_class_recall_pct":per_class,"top_confusions":[(f"{a} -> {b}",k) for (a,b),k in conf.most_common(8)]}
print(json.dumps(m,indent=1)); json.dump(m,open(f"results_{MODE}_{n}.json","w"),indent=1)
json.dump(preds,open(f"predictions_{MODE}_{n}.json","w"))   # per-document labels, for your own confusion analysis