#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
09_big_suite.py - the large four-domain task set behind `07_domain_eval.py run --suite big --hard-n N` (N tasks per domain, up to the pool size).
Public benchmarks, English only, every answer checked by code; the tasks of a smaller N are the FIRST N of a bigger one (nested), so a quick run
(--hard-n 60) is a prefix of the full run (--hard-n 300) and the two can be pooled:
  code    HumanEval+ (164) and MBPP sanitized test (257), 2:3 mix; the model writes the function, the unit tests run on it
  math    MATH-500 integer answers of level 2-5 (278) and AIME (90, AI-MO/aimo-validation-aime), 3:1 mix; final line "#### <integer>"
  reason  MMLU-Pro (240, 14 subjects round-robin, 10 options) and BIG-Bench Hard (120, 8 tasks), 2:1 mix; final line "Answer: <x>"
  en      RACE-high reading comprehension (120), CoLA grammar acceptability (120, half acceptable), and 120 writing prompts with machine-checked
          constraints (word counts, paragraphs, bullets, case, keywords, start/end text, quotes, title, commas, P.S.), 1:1:1 mix
Needs `pip install datasets` and network access; the model-written code is EXECUTED (see 07).  Selftest (same needs):
  python 09_big_suite.py selftest          reference answers pass, garbage fails, every constraint checker passes its own good/bad example
  python 09_big_suite.py validate          runs EVERY reference solution (421 code tasks) and prints the ids to put in BAD_CODE
"""
import random
import re
import sys
from fractions import Fraction as Fr

HOW = "\n\nThink step by step, then give the final answer on the last line as: #### <integer>"
LETTERS = "ABCDEFGHIJ"
BAD_CODE = {"code_he32"}      # code task ids whose own reference solution fails its tests (`validate`: HumanEval/32 find_zero is float-sensitive)


def mix(pools, pattern):
    """The pools interleaved by the digits of `pattern` ("00101" = two from pool 0, one from pool 1, ...) until every pool is used up."""
    pos, out, i = [0] * len(pools), [], 0
    while any(pos[k] < len(pools[k]) for k in range(len(pools))) and i < 10 ** 6:
        k = int(pattern[i % len(pattern)])
        i += 1
        if pos[k] < len(pools[k]):
            out.append(pools[k][pos[k]])
            pos[k] += 1
    return out


def letter(L, k):
    """check: the last 'Answer: X' of the response is the letter L (X among the first k letters)."""
    pat = rf"Answer:[\s*(]*([{LETTERS[:k]}])\b"
    return lambda resp: (re.findall(pat, resp) or [None])[-1] == L


def token(target):
    """check: the last 'Answer: X' of the response has the single word/number X equal to the target (case-insensitive, '(A)' = 'a')."""
    want = re.sub(r"[^a-z0-9\-]", "", target.casefold())
    return lambda resp: bool((got := re.findall(r"Answer:[\s*(\[]*([A-Za-z0-9\-]+)", resp))) and got[-1].casefold() == want


# ---------------------------------------------------------------------------------------------------------------------------------- code
def pool_code(rnd, run_code, code_of):
    from datasets import load_dataset
    he, mb = [], []
    for r in load_dataset("evalplus/humanevalplus", split="test"):
        tests = r["test"] + f"\n\ncheck({r['entry_point']})"
        he.append((f"code_he{r['task_id'].split('/')[-1]}", "code",
                   "Complete this Python function. Reply with one Python code block containing the whole function (and any imports it needs).\n\n"
                   f"```python\n{r['prompt']}```",
                   (lambda t: lambda resp: run_code(code_of(resp), t, 20))(tests), f"```python\n{r['prompt']}{r['canonical_solution']}\n```"))
    for r in load_dataset("google-research-datasets/mbpp", "sanitized", split="test"):
        tests = "\n".join(r["test_imports"]) + "\n" + "\n".join(r["test_list"])
        mb.append((f"code_mbpp{r['task_id']}", "code",
                   "Write a Python function for the task below. Reply with one Python code block containing the whole function (and any imports it "
                   f"needs).\n\nTask: {r['prompt']}\n\nYour function must pass these tests (they fix its name and signature):\n" + "\n".join(r["test_list"]),
                   (lambda t: lambda resp: run_code(code_of(resp), t, 20))(tests), f"```python\n{r['code']}\n```"))
    he, mb = [t for t in he if t[0] not in BAD_CODE], [t for t in mb if t[0] not in BAD_CODE]
    rnd.shuffle(he)
    rnd.shuffle(mb)
    return mix([he, mb], "00111")       # 2 HumanEval+ : 3 MBPP per cycle of five


# ---------------------------------------------------------------------------------------------------------------------------------- math
def pool_math(rnd, num):
    from datasets import load_dataset
    ma, ai = [], []
    for r in load_dataset("HuggingFaceH4/MATH-500", split="test"):
        if re.fullmatch(r"-?\d+", str(r["answer"]).strip()) and int(r["level"]) >= 2:
            ans = Fr(str(r["answer"]).strip())
            ma.append((f"math_{r['unique_id'].replace('/', '_').removesuffix('.json')}", "math", r["problem"] + HOW, num(ans), f"#### {ans}"))
    for r in load_dataset("AI-MO/aimo-validation-aime", split="train"):
        if re.fullmatch(r"-?\d+", str(r["answer"]).strip()):
            ans = Fr(str(r["answer"]).strip())
            ai.append((f"math_aime{r['id']}", "math", r["problem"] + HOW, num(ans), f"#### {ans}"))
    rnd.shuffle(ma)
    rnd.shuffle(ai)
    return mix([ma, ai], "0001")                  # 3 MATH-500 : 1 AIME


# ---------------------------------------------------------------------------------------------------------------------------------- reasoning
BBH = {       # task -> the closing instruction that names the answer format
    "logical_deduction_five_objects": "Answer: <option letter>", "tracking_shuffled_objects_five_objects": "Answer: <option letter>",
    "date_understanding": "Answer: <option letter>", "web_of_lies": "Answer: Yes or Answer: No", "boolean_expressions": "Answer: True or Answer: False",
    "navigate": "Answer: Yes or Answer: No", "formal_fallacies": "Answer: valid or Answer: invalid", "object_counting": "Answer: <integer>"}


def pool_reason(rnd):
    from datasets import load_dataset
    by = {}
    for r in load_dataset("TIGER-Lab/MMLU-Pro", split="test"):
        by.setdefault(r["category"], []).append(r)
    cats = sorted(by)
    for c in cats:
        rnd.shuffle(by[c])
    mm = []
    for i in range(18):                                       # round-robin over the 14 subjects, the subject order shuffled each round
        order = cats[:]
        rnd.shuffle(order)
        for c in order:
            r = by[c][i]
            opts = "\n".join(f"{LETTERS[j]}. {o}" for j, o in enumerate(r["options"]))
            mm.append((f"reason_mmlu{r['question_id']}", "reason", f"{r['question']}\n\n{opts}\n\nThink step by step, then end with exactly: Answer: <letter>",
                       letter(r["answer"], 10), f"Answer: {r['answer']}"))
    mm = mm[:240]
    bb = []
    for task, fmt in BBH.items():
        rows = list(load_dataset("lukaemon/bbh", task, split="test"))
        rnd.shuffle(rows)
        bb.append([(f"reason_bbh_{task}_{k}", "reason", f"{r['input']}\n\nThink step by step, then end with exactly: {fmt}", token(r["target"]),
                    f"Answer: {r['target'].strip('()')}") for k, r in enumerate(rows[:15])])
    return mix([mm, mix(bb, "01234567")], "001")             # 2 MMLU-Pro : 1 BBH


# ---------------------------------------------------------------------------------------------------------------------------------- English
def _words(r):
    return len(re.findall(r"\w+", r))


C = {}      # constraint name -> (tag, make); make(rnd) -> (rule sentence, check, a response that passes, a response that fails)
CLASH = [("lowercase", "end_phrase"), ("lowercase", "start_word"), ("lowercase", "postscript"), ("uppercase", "end_phrase"),
         ("uppercase", "start_word"), ("wrap_quotes", "end_phrase"), ("wrap_quotes", "start_word"), ("wrap_quotes", "postscript"),
         ("end_phrase", "postscript"), ("title", "start_word")]
KEYS = ("river", "lantern", "courage", "market", "silver", "window", "winter", "bridge")
BANNED = ("very", "really", "good", "nice", "thing")
ENDS = ("Is there anything else I can help with?", "That is all for now.", "Thank you for reading.")
STARTS = ("Today", "Imagine", "Remember", "Consider")


def cons(tag):
    def deco(f):
        C[f.__name__] = (tag, f)
        return f
    return deco


@cons("len")
def words_min(rnd):
    n = rnd.choice((80, 120, 150, 200))
    return f"Your response must contain at least {n} words.", lambda r: _words(r) >= n, "word " * n, "word " * (n - 1)


@cons("len")
def words_max(rnd):
    n = rnd.choice((60, 80, 100, 120))
    return f"Your response must contain fewer than {n} words.", lambda r: _words(r) < n, "word " * (n - 1), "word " * n


@cons("struct")
def paragraphs(rnd):
    n = rnd.choice((2, 3, 4))
    return (f"Write exactly {n} paragraphs, separated from each other by a line containing only ***.",
            lambda r: len(p := re.split(r"(?m)^[ \t]*\*\*\*[ \t]*$", r.strip())) == n and all(x.strip() for x in p),
            "\n***\n".join(["para"] * n), "\n***\n".join(["para"] * (n + 1)))


@cons("struct")
def bullets(rnd):
    n = rnd.choice((3, 4, 5))
    return (f"Use exactly {n} bullet points in markdown format, each line starting with \"* \", and no other bullets.",
            lambda r: sum(1 for l in r.splitlines() if re.match(r"\s*\* ", l)) == n, "\n".join(["* item"] * n), "\n".join(["* item"] * (n + 1)))


@cons("kw")
def keywords(rnd):
    a, b = rnd.sample(KEYS, 2)
    return (f"Include the words \"{a}\" and \"{b}\".", lambda r: all(re.search(rf"\b{k}\b", r, re.I) for k in (a, b)), f"{a} and {b}", f"only {a}")


@cons("forbid")
def forbid(rnd):
    w = rnd.choice(BANNED)
    return f"Do not use the word \"{w}\" anywhere.", lambda r: not re.search(rf"\b{w}\b", r, re.I), "fine text", f"a {w} text"


@cons("case")
def lowercase(rnd):
    return "Your entire response must be in lowercase letters, with no capital letters at all.", lambda r: r == r.lower(), "all lower", "Not lower"


@cons("case")
def uppercase(rnd):
    return "Your entire response must be in ALL CAPITAL LETTERS.", lambda r: r == r.upper(), "ALL UPPER", "Not upper"


@cons("end")
def end_phrase(rnd):
    p = rnd.choice(ENDS)
    return f"Finish your response with the exact phrase \"{p}\" and add nothing after it.", lambda r: r.strip().endswith(p), f"text {p}", f"{p} text"


@cons("start")
def start_word(rnd):
    w = rnd.choice(STARTS)
    return f"Begin your response with the word \"{w}\".", lambda r: bool(re.match(rf"\W*{w}\b", r.strip())), f"{w} is a day", f"Then {w}"


@cons("wrap")
def wrap_quotes(rnd):
    return ("Wrap your entire response in double quotation marks.", lambda r: len(r.strip()) > 1 and r.strip()[0] == '"' and r.strip()[-1] == '"',
            '"text"', "text")


@cons("title")
def title(rnd):
    return "Include a title wrapped in double angular brackets, like <<title>>.", lambda r: bool(re.search(r"<<[^<>\n]+>>", r)), "<<A title>> body", "A title body"


@cons("punct")
def no_commas(rnd):
    return "Do not use any commas in your response.", lambda r: "," not in r, "no comma", "a, b"


@cons("ps")
def postscript(rnd):
    return "At the end, add a postscript that starts with P.S.", lambda r: bool(re.search(r"(?m)^\s*P\.S\.", r)), "text\nP.S. more", "text"


@cons("letter")
def letter_z(rnd):
    n = rnd.choice((3, 5))
    return f"The letter \"z\" must appear at least {n} times.", lambda r: r.lower().count("z") >= n, "z" * n, "z" * (n - 1)


TOPICS = ("a rainy morning in a small town", "why sleep matters", "the history of tea", "a lost umbrella", "learning to swim as an adult", "a city without cars",
          "the first day at a new job", "a village market", "the smell of fresh bread", "an old lighthouse", "a long train journey", "keeping a diary",
          "a robot that learns to paint", "the quiet of a library", "planting a vegetable garden", "a broken clock", "moving to another country",
          "the last bus home", "a mountain hike", "a neighbour who plays the violin", "saving money as a student", "a power cut at night",
          "the invention of the bicycle", "a street musician", "cooking for a big family", "an unexpected letter", "the sound of the sea",
          "a school sports day", "a very small apartment", "the first snow of winter")
FORMS = ("short story", "product description", "email to a colleague", "blog post", "poem", "short speech", "review", "letter to a friend",
         "news summary", "explanation for a ten-year-old")


def pool_en(rnd):
    from datasets import load_dataset
    race = list(load_dataset("ehovy/race", "high", split="test"))
    rnd.shuffle(race)
    ra = []
    for k, r in enumerate(race[:120]):          # several questions share one passage id, so the position is part of the task id
        opts = "\n".join(f"{LETTERS[j]}. {o}" for j, o in enumerate(r["options"]))
        ra.append((f"en_race_{k:03d}_{r['example_id'].replace('/', '_').replace('.txt', '')}", "en",
                   f"Read the passage, then answer the question.\n\nPassage:\n{r['article']}\n\nQuestion: {r['question']}\n\n{opts}\n\n"
                   "Think briefly, then end with exactly: Answer: <letter>", letter(r["answer"], 4), f"Answer: {r['answer']}"))
    cola = list(load_dataset("nyu-mll/glue", "cola", split="validation"))
    yes, no = [r for r in cola if r["label"] == 1], [r for r in cola if r["label"] == 0]
    rnd.shuffle(yes)
    rnd.shuffle(no)
    co = []
    for r in mix([yes[:60], no[:60]], "01"):
        co.append((f"en_cola_{r['idx']}", "en", "Is this sentence grammatically acceptable English, as a careful native speaker would judge it?\n\n"
                   f"Sentence: {r['sentence']}\n\nGive at most two sentences of reasoning, then end with exactly: Answer: Yes or Answer: No",
                   token("Yes" if r["label"] == 1 else "No"), f"Answer: {'Yes' if r['label'] == 1 else 'No'}"))
    wr, names = [], sorted(C)
    for i in range(120):
        for _ in range(200):
            pick = rnd.sample(names, 2 if rnd.random() < 0.7 else 3)
            if len({C[n][0] for n in pick}) == len(pick) and not any(a in pick and b in pick for a, b in CLASH):
                break
        parts = [C[n][1](rnd) for n in pick]
        wr.append((f"en_if{i:03d}", "en", f"Write a {rnd.choice(FORMS)} about {rnd.choice(TOPICS)}. Follow every rule below.\n" + "\n".join(f"- {p[0]}" for p in parts),
                   (lambda cs: lambda resp: all(c(resp) for c in cs))([p[1] for p in parts]), None))
    return mix([ra, co, wr], "012")


# ---------------------------------------------------------------------------------------------------------------------------------- the suite
def build_big(seed, n, run_code, code_of, num):
    """[(id, category, prompt, check(response) -> bool, reference response or None)]: the first n tasks of each pool (nested for any n)."""
    T = []
    for name, pool in (("code", lambda r: pool_code(r, run_code, code_of)), ("math", lambda r: pool_math(r, num)), ("reason", pool_reason), ("en", pool_en)):
        T += pool(random.Random(f"{seed}big{name}"))[:n]
    assert len({t[0] for t in T}) == len(T), "duplicate task ids"
    return T


def constraint_unit_tests():
    bad = []
    for name, (tag, make) in C.items():
        for s in range(8):
            _, chk, good, worse = make(random.Random(s))
            if not chk(good) or chk(worse):
                bad.append((name, s))
    return bad


if __name__ == "__main__" and sys.argv[1:] in (["selftest"], ["validate"]):
    import importlib
    sys.path.insert(0, __import__("pathlib").Path(__file__).resolve().parent.as_posix())
    m = importlib.import_module("07_domain_eval")
    if sys.argv[1] == "selftest":
        assert not constraint_unit_tests(), constraint_unit_tests()
        T = build_big(1, 12, m.run_code, m.code_of, m.num)
        ref = [t for t in T if t[4] is not None]
        bad = [t[0] for t in ref if not t[3](t[4])]
        junk = [t[0] for t in T if t[4] is not None and t[3]("zzz 987654321")]      # the constraint tasks have no reference and garbage may satisfy a rule
        per = {c: sum(t[1] == c for t in T) for c in ("code", "math", "reason", "en")}
        assert not bad and not junk and per == {c: 12 for c in per} and len({t[0] for t in T}) == len(T), (bad, junk, per)
        print(f"selftest OK: {len(T)} tasks {per}, {len(ref)} references pass, garbage fails, {len(C)} constraint checkers pass their own examples")
    else:
        bad = [t[0] for t in pool_code(random.Random("v"), m.run_code, m.code_of) if not t[3](t[4])]
        print(f"{len(bad)} reference solutions fail: BAD_CODE = {set(bad)!r}")