"""Self-aggregation module: builds an items file in which each problem is followed by the written-up final parts
(the text after </think>) of K independent attempts. The model then reviews the attempts and gives one answer. Ids,
tasks and gold answers stay the same, so the file runs through gen.py / posthoc_coc.py / grade.py unchanged.

Also writes, per id, the weighted-vote answer of the same K attempts (field "vote"), so the aggregator can be compared
with the vote on exactly the same attempts, and used only where the vote is not unanimous ("unanimous").
usage: agg_items.py --recs results/m180_fb4/mini.jsonl --items data/coc/math180.jsonl --out data/agg/m180.jsonl
"""
import argparse
import collections
import json
import os
import random
import re
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from grade import last_boxed  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--recs", required=True)
ap.add_argument("--items", nargs="+", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--k", type=int, default=4)
ap.add_argument("--max-chars", type=int, default=2000)
a = ap.parse_args()

items = {}
for p in a.items:
    for l in open(p, encoding="utf-8"):
        r = json.loads(l)
        items[r["id"]] = r
by = collections.defaultdict(list)
for l in open(a.recs, encoding="utf-8"):
    r = json.loads(l)
    if r["id"] in items:
        by[r["id"]].append(r)

MC = re.compile(r"answer is:?\s*\(?\**\s*([A-J])\b")


def key(task, fin):
    if task == "mmlupro":
        m = MC.findall(fin)
        return m[-1] if m else None
    b = last_boxed(fin)
    return b.strip().replace(" ", "") if b else None


def clip(s):
    s = s.strip()
    if len(s) <= a.max_chars:
        return s
    return s[: a.max_chars // 3] + "\n[...]\n" + s[-(a.max_chars - a.max_chars // 3):]


TAIL = {"math": "Put your final answer within \\boxed{}.",
        "mmlupro": "End your response with 'The answer is (X)' where X is the letter of the correct option."}

n = 0
os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
with open(a.out, "w", encoding="utf-8") as fo:
    for pid, g in by.items():
        it = items[pid]
        g = sorted(g, key=lambda r: r["sample"])[: a.k]
        fins = [r["completion"].split("</think>", 1)[1] if "</think>" in r["completion"] else "" for r in g]
        votes = collections.defaultdict(float)
        for r, f in zip(g, fins):
            k_ = key(it["task"], f)
            if k_:
                votes[k_] += (0.25 if r.get("module") == "fallback" else 1.0) / max(r["n_tokens"], 200) ** 0.5
        keys = [key(it["task"], f) for f in fins]
        order = list(range(len(g)))
        random.Random(pid).shuffle(order)
        att = "\n\n".join(f"### Attempt {i + 1}\n{clip(fins[j]) or '(no final write-up)'}" for i, j in enumerate(order))
        prompt = (f"{it['prompt']}\n\n---\n\n{len(g)} independent attempts at this problem are shown below. They can "
                  f"contain mistakes, and their final answers can disagree.\n\n{att}\n\n---\n\nReview the attempts "
                  "critically. Find the points where they disagree, check the disputed steps yourself, and decide "
                  "which answer is correct; if all of them are wrong, solve the problem yourself. Do not just count "
                  "votes: an answer given by fewer attempts can be the right one. "
                  + TAIL["mmlupro" if it["task"] == "mmlupro" else "math"])
        fo.write(json.dumps({**it, "prompt": prompt, "vote": max(votes, key=votes.get) if votes else None,
                             "unanimous": len(set(keys)) == 1 and keys[0] is not None}, ensure_ascii=False) + "\n")
        n += 1
print(n, "items written")