#!/usr/bin/env python3
"""
Final merge ("v9") that produced the leaderboard submission (Codabench ID #138, 64.2482).

It is a deterministic weighted vote over the saved raw passes in training/passes/.
No model is called and no labels are used. This is the exact logic that was run, cleaned only
to read from the repo layout.

    python training/ensemble_v9.py                       # -> predictions_v9_rebuilt.jsonl
    cmp predictions_v9_rebuilt.jsonl predictions.jsonl   # identical

Each source is a dict {(q_no, subject, category, level): "1".."4"} (answers in ORIGINAL option order):
  p1  : Pass 1  think, 2048 tok, original order, all 2238 questions
  p2  : Pass 2  think, 4096 tok, options rotated by 1  (hard + questions where CoT != Direct)
  p3  : Pass 3  think, 4096 tok, system prompt with Sanskrit/Pali/technical-term hint (uncertain set)
  p4  : Pass 4  think, 4096 tok, options rotated by 2  (uncertain set)
  p5  : Pass 5  think, 4096 tok, options rotated by 3  (uncertain set)
  mac : repeat of the Pass 1 configuration on a second machine (Apple Metal), disputed + hard questions
  cot : 4-shot English CoT baseline (think=False)
  dr  : 4-shot direct-answer baseline (think=False)
Rotated passes (p2, p4, p5) are ignored on questions whose options refer to option positions
("ඉහත", "above", "1 සහ 2", ...) because rotation changes the meaning of such options.
"""
import argparse
import json
import re
from collections import Counter
from pathlib import Path

POS_RE = re.compile(r"ඉහත|above|\b[1-4]\s*(?:හා|සහ|හෝ|,|and|or)\s*[1-4]\b", re.I)

PASS_FILES = {
    "p1": "pass1_think_2048_orig.jsonl",
    "p2": "pass2_think_4096_rot1.jsonl",
    "p3": "pass3_think_4096_hint.jsonl",
    "p4": "pass4_think_4096_rot2.jsonl",
    "p5": "pass5_think_4096_rot3.jsonl",
    "mac": "mac_think_2048_orig.jsonl",
    "cot": "baseline_cot_4shot.jsonl",
    "dr": "baseline_direct_4shot.jsonl",
}


def key_of(r):
    return (r["q_no"], r["subject"], r["category"], r["level"])


def load_map(path):
    m = {}
    with Path(path).open(encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            if line.strip():
                r = json.loads(line)
                k = key_of(r)
                answer = str(r["answer"])
                if k in m or answer not in {"1", "2", "3", "4"}:
                    raise ValueError(f"Duplicate key or invalid answer: {path}:{line_no}")
                m[k] = answer
    return m


def merge_one(k, pos, p1, p2, p3, p4, p5, mac, cot, dr):
    """Weighted vote for one question. k = (q_no, subject, category, level)."""
    lvl = k[3]
    v = Counter()
    if lvl == "hard":
        v[p1[k]] += 1.15
        if k in mac: v[mac[k]] += 1.10
        if k in p3:  v[p3[k]]  += 1.10
        for m in (p2, p4, p5):
            if k in m and k not in pos: v[m[k]] += 1.10
        if k in cot: v[cot[k]] += 0.35
        if k in dr:  v[dr[k]]  += 0.35
    elif lvl == "medium":
        v[p1[k]] += 1.25
        if k in mac: v[mac[k]] += 1.05
        if k in p3:  v[p3[k]]  += 1.10
        for m in (p2, p4, p5):
            if k in m and k not in pos: v[m[k]] += 1.05
        if k in cot: v[cot[k]] += 0.35
        if k in dr:  v[dr[k]]  += 0.35
    else:  # easy: Pass 1 anchored at 2.40, so only a margin >= 2 supermajority can overturn it
        v[p1[k]] += 2.40
        if k in mac: v[mac[k]] += 1.0
        if k in p3:  v[p3[k]]  += 1.0
        for m in (p2, p4, p5):
            if k in m and k not in pos: v[m[k]] += 0.95
        if k in cot: v[cot[k]] += 0.30
        if k in dr:  v[dr[k]]  += 0.30
    return v.most_common(1)[0][0]


def positional_keys(questions):
    """Keys of questions whose option text refers to option positions (input-only regex)."""
    pos = set()
    for q in questions:
        if POS_RE.search(" ".join(map(str, q["choices"]))):
            pos.add((q["q_no"], q["subject"], q["category"], q["level"]))
    return pos


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--passes", default=str(Path(__file__).parent / "passes"))
    ap.add_argument("--data", default=None,
                    help="dir with easy/medium/hard/*.json (needed only for the positional-option filter)")
    ap.add_argument("--out", default="predictions_v9_rebuilt.jsonl")
    a = ap.parse_args()

    if Path(a.out).resolve() == (Path(__file__).resolve().parents[1] / "predictions.jsonl"):
        raise ValueError("Refusing to overwrite the submitted predictions.jsonl; choose a different --out.")
    maps = {n: load_map(Path(a.passes) / f) for n, f in PASS_FILES.items()}
    ref = [json.loads(l) for l in open(Path(a.passes) / PASS_FILES["p1"], encoding="utf-8") if l.strip()]

    pos = set()
    if a.data:
        qs = []
        for lvl in ("easy", "medium", "hard"):
            for f in sorted((Path(a.data) / lvl).glob("*.json")):
                for q in json.load(open(f, encoding="utf-8")):
                    q["level"] = lvl
                    qs.append(q)
        pos = positional_keys(qs)
    else:
        pos_file = Path(a.passes) / "positional_option_keys.json"
        pos = {tuple(x) for x in json.load(open(pos_file, encoding="utf-8"))}

    with open(a.out, "w", encoding="utf-8") as f:
        for q in ref:
            k = key_of(q)
            ans = merge_one(k, pos, maps["p1"], maps["p2"], maps["p3"], maps["p4"], maps["p5"],
                            maps["mac"], maps["cot"], maps["dr"])
            f.write(json.dumps({"q_no": q["q_no"], "subject": q["subject"], "category": q["category"],
                                "level": q["level"], "answer": ans}, ensure_ascii=False) + "\n")
    print("wrote", a.out)


if __name__ == "__main__":
    main()