"""Prompt set for comparing two servers that hold the same model.

Every graded prompt has one right answer that a program can check.
The open prompts are not graded; a person reads them side by side.
"""
import datetime
import json
import math
import os
import random

CODE_RULE = ('Reply with one Python code block and nothing else. '
             'Use only the standard library.')
ANSWER_RULE = ('Work briefly, then give the result on the last line as '
               '"ANSWER: " followed by the answer only.')
JSON_RULE = 'Reply with the JSON only, no code fence and no other text.'

# Each code task: id, task text, a reference solution, and a test.
CODE = [
    ('roman', 'Write roman_to_int(s) that converts a Roman numeral such '
     'as "MCMXCIV" to an integer.', '''
def roman_to_int(s):
    v = dict(I=1, V=5, X=10, L=50, C=100, D=500, M=1000)
    t = 0
    for a, b in zip(s, s[1:] + 'I'):
        t += -v[a] if v[a] < v[b] else v[a]
    return t
''', '''
assert roman_to_int("MCMXCIV") == 1994
assert roman_to_int("XLII") == 42
assert roman_to_int("III") == 3
assert roman_to_int("MMMCMXCIX") == 3999
assert roman_to_int("CDXLIV") == 444
'''),
    ('intervals', 'Write merge_intervals(xs) that takes a list of '
     '[start, end] pairs in any order and returns the merged, sorted '
     'list of pairs. Intervals that touch, such as [1, 2] and [2, 3], '
     'merge.', '''
def merge_intervals(xs):
    out = []
    for a, b in sorted(xs):
        if out and a <= out[-1][1]:
            out[-1][1] = max(out[-1][1], b)
        else:
            out.append([a, b])
    return out
''', '''
f = lambda xs: [list(p) for p in merge_intervals(xs)]
assert f([[8, 10], [1, 3], [2, 6], [15, 18]]) == [[1, 6], [8, 10], [15, 18]]
assert f([[1, 2], [2, 3]]) == [[1, 3]]
assert f([]) == []
assert f([[1, 10], [2, 3], [4, 5]]) == [[1, 10]]
assert f([[5, 6], [1, 2]]) == [[1, 2], [5, 6]]
'''),
    ('lcs', 'Write lcs_length(a, b) that returns the length of the '
     'longest common subsequence of two strings.', '''
def lcs_length(a, b):
    row = [0] * (len(b) + 1)
    for x in a:
        prev = 0
        for j, y in enumerate(b, 1):
            prev, row[j] = row[j], (prev + 1 if x == y
                                    else max(row[j], row[j - 1]))
    return row[-1]
''', '''
assert lcs_length("ABCBDAB", "BDCABA") == 4
assert lcs_length("", "abc") == 0
assert lcs_length("abc", "abc") == 3
assert lcs_length("abc", "def") == 0
assert lcs_length("AGGTAB", "GXTXAYB") == 4
'''),
    ('duration', 'Write parse_duration(s) that converts a string such '
     'as "1h30m15s", "45m" or "2d4h" to seconds. Units are d, h, m, s, '
     'each optional, always in that order. Raise ValueError for an '
     'empty or malformed string.', '''
import re
def parse_duration(s):
    m = re.fullmatch(r'(?:(\\d+)d)?(?:(\\d+)h)?(?:(\\d+)m)?(?:(\\d+)s)?', s)
    if not s or not m:
        raise ValueError(s)
    d, h, mi, se = (int(x or 0) for x in m.groups())
    return ((d * 24 + h) * 60 + mi) * 60 + se
''', '''
assert parse_duration("1h30m15s") == 5415
assert parse_duration("45m") == 2700
assert parse_duration("2d4h") == 187200
assert parse_duration("90s") == 90
for bad in ("", "abc", "5x", "1m1h"):
    try:
        parse_duration(bad)
    except ValueError:
        pass
    else:
        raise AssertionError(bad)
'''),
    ('toposort', 'Write topo_sort(graph) where graph maps each node (a '
     'string) to the list of nodes that must come after it. Every node '
     'appears as a key. Return the order that is lexicographically '
     'smallest among valid orders. Raise ValueError if there is a '
     'cycle.', '''
import heapq
def topo_sort(graph):
    deg = {n: 0 for n in graph}
    for n in graph:
        for m in graph[n]:
            deg[m] += 1
    heap = [n for n in graph if not deg[n]]
    heapq.heapify(heap)
    out = []
    while heap:
        n = heapq.heappop(heap)
        out.append(n)
        for m in graph[n]:
            deg[m] -= 1
            if not deg[m]:
                heapq.heappush(heap, m)
    if len(out) != len(graph):
        raise ValueError('cycle')
    return out
''', '''
g = {"d": [], "b": ["d"], "a": ["d"], "c": ["a"]}
assert list(topo_sort(g)) == ["b", "c", "a", "d"]
assert list(topo_sort({})) == []
assert list(topo_sort({"x": ["y"], "y": []})) == ["x", "y"]
try:
    topo_sort({"a": ["b"], "b": ["a"]})
except ValueError:
    pass
else:
    raise AssertionError("cycle")
'''),
    ('base64', 'Write b64encode(data) that takes bytes and returns the '
     'standard padded Base64 text as a str. Do not import base64 or '
     'binascii.', '''
def b64encode(data):
    abc = ('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
           '0123456789+/')
    out = []
    for i in range(0, len(data), 3):
        c = data[i:i + 3]
        n = int.from_bytes(c + bytes(3 - len(c)), 'big')
        s = ''.join(abc[n >> k & 63] for k in (18, 12, 6, 0))
        out.append(s[:len(c) + 1] + '=' * (3 - len(c)))
    return ''.join(out)
''', '''
import base64 as _b
import re, sys
src = open(sys.argv[0]).read().split("# TESTS")[0]
assert not re.search(r"(import|from)\\s+(base64|binascii)", src)
for d in (b"", b"f", b"fo", b"foo", b"foob", bytes(range(256))):
    assert b64encode(d) == _b.b64encode(d).decode(), d
'''),
    ('weekday', 'Write weekday(year, month, day) that returns the day '
     'of the week for a Gregorian date as an integer, Monday 0 to '
     'Sunday 6. Do not import datetime, calendar or time.', '''
def weekday(year, month, day):
    if month < 3:
        month += 12
        year -= 1
    k, j = year % 100, year // 100
    h = (day + 13 * (month + 1) // 5 + k + k // 4 + j // 4 + 5 * j) % 7
    return (h + 5) % 7
''', '''
import datetime as _d
import re, sys
src = open(sys.argv[0]).read().split("# TESTS")[0]
assert not re.search(r"(import|from)\\s+(datetime|calendar|time)", src)
for y, m, d in ((2026, 11, 11), (2000, 2, 29), (1999, 12, 31),
                (2024, 1, 1), (1900, 3, 1), (2100, 2, 28)):
    assert weekday(y, m, d) == _d.date(y, m, d).weekday(), (y, m, d)
'''),
    ('rle', 'Write rle_encode(s) and rle_decode(s). Encoding turns '
     '"aaabcc" into "3a1b2c": each run becomes its length then the '
     'character. The input to encode never contains digits. '
     'rle_decode reverses it; run lengths may have several digits.', '''
import itertools, re
def rle_encode(s):
    return ''.join('%d%s' % (len(list(g)), c)
                   for c, g in itertools.groupby(s))
def rle_decode(s):
    return ''.join(c * int(n) for n, c in re.findall(r'(\\d+)(\\D)', s))
''', '''
assert rle_encode("aaabcc") == "3a1b2c"
assert rle_encode("") == ""
assert rle_encode("z" * 12 + "y") == "12z1y"
assert rle_decode("12z1y") == "z" * 12 + "y"
assert rle_decode("") == ""
for s in ("hello world", "aAaA", "  x  "):
    assert rle_decode(rle_encode(s)) == s
'''),
    ('rpn', 'Write eval_rpn(tokens) that evaluates a list of string '
     'tokens in reverse Polish notation with the operators + - * /. '
     'Operands are integers and may be negative. Division truncates '
     'toward zero. Return an int.', '''
def eval_rpn(tokens):
    st = []
    for t in tokens:
        if t in '+-*/' and len(t) == 1:
            b, a = st.pop(), st.pop()
            if t == '+': st.append(a + b)
            elif t == '-': st.append(a - b)
            elif t == '*': st.append(a * b)
            else:
                q = abs(a) // abs(b)
                st.append(q if (a < 0) == (b < 0) else -q)
        else:
            st.append(int(t))
    return st[0]
''', '''
assert eval_rpn(["2", "1", "+", "3", "*"]) == 9
assert eval_rpn(["4", "13", "5", "/", "+"]) == 6
assert eval_rpn(["-7", "2", "/"]) == -3
assert eval_rpn(["7", "-2", "/"]) == -3
assert eval_rpn(["10", "6", "9", "3", "+", "-11", "*", "/", "*", "17",
                 "+", "5", "+"]) == 22
assert eval_rpn(["-5"]) == -5
'''),
    ('lru', 'Write a class LRUCache with __init__(self, capacity), '
     'get(self, key) returning the value or -1, and '
     'put(self, key, value). When the cache is over capacity, put '
     'evicts the least recently used key. Both get and put count as '
     'use.', '''
import collections
class LRUCache:
    def __init__(self, capacity):
        self.n, self.d = capacity, collections.OrderedDict()
    def get(self, key):
        if key not in self.d:
            return -1
        self.d.move_to_end(key)
        return self.d[key]
    def put(self, key, value):
        self.d[key] = value
        self.d.move_to_end(key)
        if len(self.d) > self.n:
            self.d.popitem(last=False)
''', '''
c = LRUCache(2)
c.put(1, 1); c.put(2, 2)
assert c.get(1) == 1
c.put(3, 3)
assert c.get(2) == -1
c.put(4, 4)
assert c.get(1) == -1 and c.get(3) == 3 and c.get(4) == 4
c.put(3, 30); c.put(5, 5)
assert c.get(4) == -1 and c.get(3) == 30 and c.get(5) == 5
'''),
    ('brackets', 'Write balanced(s) that returns True when every '
     'bracket in s among ()[]{} is closed in the right order, and '
     'False otherwise. Other characters are ignored.', '''
def balanced(s):
    pair, st = {')': '(', ']': '[', '}': '{'}, []
    for c in s:
        if c in '([{':
            st.append(c)
        elif c in pair and (not st or st.pop() != pair[c]):
            return False
    return not st
''', '''
assert balanced("") is True
assert balanced("a(b[c]{d})e") is True
assert balanced("([)]") is False
assert balanced("((") is False
assert balanced("}") is False
assert balanced("{[()()]}") is True
'''),
    ('spiral', 'Write spiral(matrix) that returns the elements of a '
     'rectangular list of lists in clockwise spiral order starting at '
     'the top left. An empty matrix gives an empty list.', '''
def spiral(matrix):
    m, out = [list(r) for r in matrix], []
    while m and m[0]:
        out += m.pop(0)
        m = [list(r) for r in zip(*m)][::-1]
    return out
''', '''
assert list(spiral([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) == \\
    [1, 2, 3, 6, 9, 8, 7, 4, 5]
assert list(spiral([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])) == \\
    [1, 2, 3, 4, 8, 12, 11, 10, 9, 5, 6, 7]
assert list(spiral([])) == []
assert list(spiral([[1], [2], [3]])) == [1, 2, 3]
assert list(spiral([[7]])) == [7]
'''),
]

SENTENCE = 'strawberry raspberry blackberry cranberry'
NUMBERS = [41, 7, 93, 12, 58, 66, 3, 29, 77, 35, 50]
MATH = [
    ('arith', 'Compute 17 * 23 + 19 * 31 - 144.', 17 * 23 + 19 * 31 - 144),
    ('modpow', 'What is 2 to the power 20, modulo 97?', pow(2, 20, 97)),
    ('primes', 'What is the sum of all prime numbers below 60?',
     sum(n for n in range(2, 60) if all(n % d for d in range(2, n)))),
    ('gcd', 'What is the greatest common divisor of 1071 and 462?',
     math.gcd(1071, 462)),
    ('zeros', 'How many zeros are at the end of 100 factorial?',
     len(str(math.factorial(100))) - len(str(math.factorial(100))
                                         .rstrip('0'))),
    ('hex', 'Convert hexadecimal 3F7A to decimal.', 0x3F7A),
    ('days', 'How many days are there from 10 February 2024 to 1 March '
     '2025, counting the end date but not the start date?',
     (datetime.date(2025, 3, 1) - datetime.date(2024, 2, 10)).days),
    ('height', 'Ann is taller than Ben. Cy is shorter than Ben. Dee is '
     'taller than Ann. Who is the second tallest? Answer with the name.',
     'Ann'),
    ('speed', 'A car drives 150 km at 60 km/h and then 150 km at 100 '
     'km/h. What is its average speed for the whole trip in km/h?', 75),
    ('divisible', 'How many integers from 1 to 1000 inclusive are '
     'divisible by 3 or by 5 but not by both?',
     sum((n % 3 == 0) != (n % 5 == 0) for n in range(1, 1001))),
    ('bits', 'How many 1 bits are in the binary form of 2026?',
     bin(2026).count('1')),
    ('workers', 'Three workers finish a job in 12 days. At the same '
     'rate, how many days do nine workers need?', 4),
    ('letters', 'How many times does the letter r appear in "%s"?'
     % SENTENCE, SENTENCE.count('r')),
    ('median', 'What is the median of %s?' % NUMBERS,
     sorted(NUMBERS)[len(NUMBERS) // 2]),
]

CSV = 'id,name,score\n7,Ada,91\n3,Linus,78\n12,Grace,85'
WORDS = ['pear', 'Fig', 'apple', 'cherry', 'banana', 'date', 'elder',
         'grape', 'Apricot']
FORMAT = [
    ('extract', 'From this text make a JSON object with exactly the '
     'keys "name", "age", "city" and "languages" (a list, in the order '
     'given). Age is a number.\n\nText: Maren Holt, who turned 34 last '
     'spring, moved to Tromso and writes Rust, Zig and C for a living.',
     {'name': 'Maren Holt', 'age': 34, 'city': 'Tromso',
      'languages': ['Rust', 'Zig', 'C']}),
    ('csv', 'Convert this CSV to a JSON array of objects, one per row, '
     'in the same order. "id" and "score" are numbers.\n\n' + CSV,
     [{'id': 7, 'name': 'Ada', 'score': 91},
      {'id': 3, 'name': 'Linus', 'score': 78},
      {'id': 12, 'name': 'Grace', 'score': 85}]),
    ('fib', 'Give the first 15 Fibonacci numbers, starting 1, 1, as a '
     'JSON array of numbers.',
     [1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610]),
    ('sort', 'Sort these words alphabetically, ignoring case but '
     'keeping each word as written, and give a JSON array of strings: '
     + ', '.join(WORDS) + '.', sorted(WORDS, key=str.lower)),
    ('dates', 'Give a JSON object mapping each event to its date in '
     'ISO 8601 form (YYYY-MM-DD). Keys are "launch", "review" and '
     '"close".\n\nThe launch is on March 5th, 2021; the review is on '
     'the 17th of November 2022; the close is on 1 Jan 2024.',
     {'launch': '2021-03-05', 'review': '2022-11-17',
      'close': '2024-01-01'}),
    ('nested', 'Give a JSON object with the key "matrix" holding a 3 by '
     '3 identity matrix as a list of lists of numbers, and the key '
     '"trace" holding its trace.',
     {'matrix': [[1, 0, 0], [0, 1, 0], [0, 0, 1]], 'trace': 3}),
]

OPEN = [
    ('handshake', 'Explain to a junior engineer why TCP opens a '
     'connection with a three-way handshake and not two. At most 200 '
     'words.'),
    ('review', 'Review this function and list its problems, most '
     'serious first.\n\n```python\ndef find_user(db, name, cache={}):\n'
     '    if name in cache:\n        return cache[name]\n'
     '    q = "SELECT * FROM users WHERE name = \'" + name + "\'"\n'
     '    row = db.execute(q).fetchone()\n    cache[name] = row\n'
     '    return row\n```'),
    ('mutex', 'Explain the difference between a mutex and a semaphore '
     'with one concrete example of each. At most 200 words.'),
    ('lsm', 'Compare B-trees and LSM-trees for a write-heavy workload. '
     'At most 200 words.'),
    ('commit', 'Write a Git commit message, an imperative subject and '
     'a short body, for this change.\n\n```diff\n-    timeout = 30\n'
     '+    timeout = int(os.environ.get("FETCH_TIMEOUT", "30"))\n'
     '     resp = session.get(url, timeout=timeout)\n'
     '+    resp.raise_for_status()\n```'),
]

FIRST = ['Marta', 'Jonas', 'Aiko', 'Pavel', 'Leila', 'Tomas', 'Greta',
         'Samir', 'Noor', 'Emil', 'Rosa', 'Viktor', 'Hana', 'Diego',
         'Freya', 'Kofi']
LAST = ['Voss', 'Lindqvist', 'Tanaka', 'Novak', 'Haddad', 'Berg',
        'Moreau', 'Okafor', 'Silva', 'Keller', 'Duarte', 'Ivanov',
        'Larsen', 'Costa', 'Weber', 'Mensah']
CITY = ['Lyon', 'Porto', 'Graz', 'Turku', 'Ghent', 'Malmo', 'Bilbao',
        'Split', 'Leeds', 'Bergen', 'Padua', 'Brno']
STATUS = ['shipped', 'pending', 'returned', 'delivered']


def ledger(lines, seed):
    """Return a synthetic order ledger and four questions about it."""
    rng = random.Random(seed)
    ids = rng.sample(range(10000, 100000), lines)
    rows = [dict(id='ORD-%d' % i, who='%s %s' % (rng.choice(FIRST),
                                                 rng.choice(LAST)),
                 city=rng.choice(CITY), amount=rng.randrange(100, 9900),
                 status=rng.choice(STATUS)) for i in ids]
    a, b, c, d, e = (lines * k // 100 for k in (23, 41, 58, 72, 87))
    rows[b]['who'] = 'Ilse Quarrington'
    text = '\n'.join('%(id)s | customer: %(who)s | city: %(city)s | '
                     'amount: %(amount)d EUR | status: %(status)s' % r
                     for r in rows)
    head = ('Below is an order ledger, one order per line. Answer the '
            'question after it from the ledger alone.\n\n' + text
            + '\n\nQuestion: ')
    tail = (' Give the result on the last line as "ANSWER: " followed '
            'by the answer only.')
    return [
        ('amount', head + 'What is the amount in EUR of order %s?'
         % rows[a]['id'] + tail, rows[a]['amount']),
        ('customer', head + 'Which order ID belongs to the customer '
         'Ilse Quarrington?' + tail, rows[b]['id']),
        ('next', head + 'Which city is on the line immediately after '
         'order %s?' % rows[c]['id'] + tail, rows[c + 1]['city']),
        ('sum', head + 'What is the sum in EUR of the amounts of orders '
         '%s and %s?' % (rows[d]['id'], rows[e]['id']) + tail,
         rows[d]['amount'] + rows[e]['amount']),
    ]


# Ethereum facts with one fixed answer each.
ETH = [
    ('slots', 'How many slots are in one epoch of the Ethereum beacon '
     'chain?', 32),
    ('slot-time', 'How many seconds long is one slot on Ethereum '
     'mainnet?', 12),
    ('transfer-gas', 'What is the intrinsic gas cost of a plain ETH '
     'transfer with empty calldata?', 21000),
    ('selector', 'What is the 4-byte function selector of the ERC-20 '
     'function transfer(address,uint256)? Answer in lowercase hex '
     'without the 0x prefix.', 'a9059cbb'),
    ('blob-eip', 'What is the number of the EIP that introduced blob '
     'transactions (proto-danksharding)?', 4844),
    ('blob-elements', 'How many field elements does one EIP-4844 blob '
     'hold?', 4096),
    ('basefee', 'In EIP-1559, what is the value of '
     'BASE_FEE_MAX_CHANGE_DENOMINATOR?', 8),
    ('replay-eip', 'What is the number of the EIP that added the chain '
     'ID to transaction signatures for replay protection?', 155),
    ('code-size', 'What is the maximum deployed contract code size in '
     'bytes set by EIP-170?', 24576),
    ('sync-committee', 'How many validators are in a beacon chain sync '
     'committee on mainnet?', 512),
    ('sync-period', 'How many epochs does one sync committee period '
     'last on mainnet?', 256),
    ('blob-type', 'What is the transaction type number of an EIP-4844 '
     'blob transaction?', 3),
    ('empty-hash', 'What are the first four bytes of the Keccak-256 '
     'hash of the empty byte string? Answer in lowercase hex without '
     'the 0x prefix.', 'c5d24601'),
]


def prompts():
    """Return every prompt as a dict, in the order they are sent."""
    out = []
    for name, task, ref, test in CODE:
        out.append(dict(id='code-' + name, kind='code', max_tokens=1024,
                        prompt=task + ' ' + CODE_RULE, ref=ref, test=test))
    for name, task, want in MATH:
        out.append(dict(id='math-' + name, kind='answer', max_tokens=768,
                        prompt=task + ' ' + ANSWER_RULE, want=want))
    for name, task, want in FORMAT:
        out.append(dict(id='json-' + name, kind='json', max_tokens=512,
                        prompt=task + ' ' + JSON_RULE, want=want))
    for size, lines, seed in (('30k', 900, 11), ('100k', 3000, 12)):
        for name, prompt, want in ledger(lines, seed):
            out.append(dict(id='find-%s-%s' % (size, name), kind='answer',
                            max_tokens=256, prompt=prompt, want=want))
    for name, task, want in ETH:
        out.append(dict(id='eth-fact-' + name, kind='answer',
                        max_tokens=256, prompt=task + ' ' + ANSWER_RULE,
                        want=want))
    # Private report prompts, built by eth_prompts.py; optional.
    private = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                           'eth-prompts.json')
    if os.path.exists(private):
        out += json.load(open(private))
    for name, task in OPEN:
        out.append(dict(id='open-' + name, kind='open', max_tokens=700,
                        prompt=task))
    return out


if __name__ == '__main__':
    for p in prompts():
        print(p['id'], p['kind'], len(p['prompt']),
              json.dumps(p.get('want')))