"""Compare response LLMs on mix-matched replies.
  gen <tag> prompts.json out.json   (needs a vLLM server at LLM_BASE; full Responder path: mix -> n=3 candidates -> validator -> regen)
  bt  out1.json [out2.json ...]     (MADLAD-400-3B (Apache-2.0) back-translates the non-English spans of user + reply into English)
  judge out1.json ...               (server LLM rates coherence/relevance of the back-translated reply, 1-5)
  summary out1.json ..."""
import sys, json, re, os, statistics as st
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
TAG = re.compile(r"<(eng|ibo|yor|hau|pcm)>(.*?)</\1>", re.S)

def add_native(P):
    """+4 native WAXAL test transcripts per ibo/yor/hau (monolingual user turns)"""
    L = [json.loads(l) for l in open("/scratch/duplex/asr_eval/list.jsonl")]
    for lg in ("ibo", "yor", "hau"):
        rs = [r for r in L if r["lang"] == lg and len(r["text"].split()) >= 6][:4]
        P += [dict(id=f"n_{lg}_{i}", kind=f"{lg} native", text=r["text"]) for i, r in enumerate(rs)]
    return P

def gen(tag, pf, out):
    from respond import Responder
    import mix as M
    R = Responder(); res = []
    for p in add_native(json.load(open(pf))):
        R.style = None
        r = R.reply(p["text"])
        spans = TAG.findall(r["reply"]); glot = []
        for l, t in spans:
            if l in ("ibo", "yor", "hau") and len(t.split()) >= 2:
                s = M.tagger("glotlid")._lid_scores(t); glot.append(max(s, key=s.get) == l)
        res.append(dict(p, model=tag, **{k: r[k] for k in ("target", "reply", "reply_mix", "tvd", "how", "first_cand_ok", "any_first_round_ok", "t_llm", "rounds")}, glot_ok=glot))
        print(json.dumps({"id": p["id"], "reply": r["reply"], "how": r["how"], "tvd": r["tvd"]}, ensure_ascii=False), flush=True)
    json.dump(res, open(out, "w"), ensure_ascii=False, indent=1)

def bt(files):
    import torch
    from transformers import T5ForConditionalGeneration, T5Tokenizer
    import mix as M
    tok = T5Tokenizer.from_pretrained("google/madlad400-3b-mt"); m = T5ForConditionalGeneration.from_pretrained("google/madlad400-3b-mt", torch_dtype=torch.bfloat16).cuda().eval()
    def tr(s):
        ids = tok("<2en> " + s, return_tensors="pt").input_ids.cuda()
        with torch.no_grad(): o = m.generate(ids, max_new_tokens=3 * len(s.split()) + 16, num_beams=2, no_repeat_ngram_size=4)
        return tok.decode(o[0], skip_special_tokens=True)
    def en(tagged):
        return " ".join(t if l in ("eng", "pcm") else tr(t) for l, t in TAG.findall(tagged))
    for f in files:
        R = json.load(open(f))
        for r in R:
            r["user_bt"] = en(M.tagger("glotlid").tag(r["text"])); r["reply_bt"] = en(r["reply"])
        json.dump(R, open(f, "w"), ensure_ascii=False, indent=1); print("bt", f, flush=True)

def judge(files):
    from llm import LLM
    J = LLM()
    for f in files:
        R = json.load(open(f))
        for r in R:
            q = (f"A user said (English gloss): \"{r['user_bt']}\"\nThe assistant replied (English back-translation): \"{r['reply_bt']}\"\n"
                 "Rate from 1 to 5 how coherent, sensible and relevant the reply is as a spoken assistant response (5 = fully sensible and on-topic, "
                 "1 = nonsense or off-topic). Output only the digit.")
            o, _ = J.generate([{"role": "user", "content": q}], n=1, temperature=0.0, max_tokens=4)
            m = re.search(r"[1-5]", o[0]); r["judge"] = int(m.group()) if m else None
        json.dump(R, open(f, "w"), ensure_ascii=False, indent=1); print("judged", f, flush=True)

def summary(files):
    rows = []
    for f in files:
        R = json.load(open(f)); tag = R[0]["model"]
        def grp(pred):
            S = [r for r in R if pred(r)]
            if not S: return None
            g = [x for r in S for x in r["glot_ok"]]
            j = [r["judge"] for r in S if r.get("judge")]
            return dict(n=len(S), first_cand_ok=round(sum(r["first_cand_ok"] for r in S) / len(S), 2), regen=round(sum(r["how"] != "first_round" for r in S) / len(S), 2),
                        fallback=round(sum(r["how"] in ("best_failed", "plain_matrix") for r in S) / len(S), 2), tvd=round(st.mean(r["tvd"] or 0 for r in S), 3),
                        glot_ok=round(sum(g) / len(g), 2) if g else None, judge=round(st.mean(j), 2) if j else None, llm_s=round(st.median(r["t_llm"] for r in S), 3))
        rows.append(dict(model=tag, all=grp(lambda r: True), mixed=grp(lambda r: not r["id"].startswith("n_") and not r["id"].startswith("e_")),
                         native=grp(lambda r: r["id"].startswith("n_"))))
    print(json.dumps(rows, indent=1)); return rows

if __name__ == "__main__":
    c = sys.argv[1]
    if c == "gen": gen(sys.argv[2], sys.argv[3], sys.argv[4])
    elif c == "bt": bt(sys.argv[2:])
    elif c == "judge": judge(sys.argv[2:])
    else: json.dump(summary(sys.argv[2:]), open("/scratch/duplex/llm_summary.json", "w"), indent=1)