oracle7 / root /wander-train /code_organ_v2.py
Add files using upload-large-folder tool
ea2261d verified 10 days ago
31.4 kB
"""code_organ_v2 - v1 (MBPP-frozen b2e52a1d) + STRICT mode for a single example: every move named by the
request. Programs as specified free moves over a verified function library, checked by execution; every
verified program joins the library (the stored final forms - continuous learning without retraining).
REQUEST   description + input/output examples (asserts) -> function name, argument values and types, expected outputs,
          constants the description names ("at least 4 characters", "the character 'l'")
MOVES     typed library functions. The TYPE SYSTEM is the specified structure: a function applies only where its
          argument types connect; within that the choice is free. The description's words rank the moves.
SEARCH    bottom-up enumeration with OBSERVATIONAL EQUIVALENCE: programs that behave identically on the examples are
          kept once, so compositions of 3-4 moves stay tractable.
CHECK     the rendered source is executed against every example; nothing is released that did not run and pass.
STORE     a verified program becomes a library function (its description words are its keywords, its signature its
          type); later requests compose it - measured as the library-growth curve.
Only library templates are ever evaluated; no text from the request is executed."""
from __future__ import annotations

import ast
import itertools
import json
import math
import re
from collections import Counter

# ----------------------------------------------------------------------------------------------------------- types
def tyof(v):
    if isinstance(v, bool):
        return "bool"
    if isinstance(v, int):
        return "int"
    if isinstance(v, float):
        return "float"
    if isinstance(v, str):
        return "str"
    if isinstance(v, list):
        return "list"
    if isinstance(v, tuple):
        return "tuple"
    if isinstance(v, dict):
        return "dict"
    if isinstance(v, set):
        return "set"
    return "any"


GROUP = {"num": {"int", "float"}, "seq": {"list", "tuple", "str"}, "coll": {"list", "tuple", "set"},
         "any": {"bool", "int", "float", "str", "list", "tuple", "dict", "set", "any"}}


def fits(t, want):
    return t == want or t in GROUP.get(want, ())


# ------------------------------------------------------------------------------------------------------ primitives
class Prim:
    __slots__ = ("name", "ins", "out", "fn", "src", "kw")

    def __init__(self, name, ins, out, fn, src, kw):
        self.name, self.ins, self.out, self.fn, self.src = name, ins, out, fn, src
        self.kw = {w.lower().rstrip("s") for w in kw.split()}   # v2: normalised exactly like the request words


def _big(x):
    if isinstance(x, int) and abs(x) > 10 ** 7:
        raise ValueError("too large")
    return x


def _is_prime(n):
    if n < 2:
        return False
    return all(n % d for d in range(2, int(n ** 0.5) + 1))


def _uniq(xs):
    out = []
    for x in xs:
        if x not in out:
            out.append(x)
    return out


def _flatten(xs):
    out = []
    for x in xs:
        out += list(x) if isinstance(x, (list, tuple)) else [x]
    return out


P_ = [
    # numbers
    Prim("add", ("num", "num"), "num", lambda a, b: a + b, "({0} + {1})", "add sum plus total"),
    Prim("sub", ("num", "num"), "num", lambda a, b: a - b, "({0} - {1})", "subtract difference minus"),
    Prim("mul", ("num", "num"), "num", lambda a, b: _big(a * b), "({0} * {1})", "multiply product times"),
    Prim("fdiv", ("num", "num"), "num", lambda a, b: a // b, "({0} // {1})", "divide quotient integer"),
    Prim("div", ("num", "num"), "float", lambda a, b: a / b, "({0} / {1})", "divide ratio average"),
    Prim("mod", ("num", "num"), "num", lambda a, b: a % b, "({0} % {1})", "remainder modulo divisible"),
    Prim("pow", ("num", "num"), "num", lambda a, b: _big(a ** b) if b < 64 else 1 / 0, "({0} ** {1})", "power square cube exponent"),
    Prim("sq", ("num",), "num", lambda a: a * a, "({0} ** 2)", "square squares"),
    Prim("cube", ("num",), "num", lambda a: a ** 3, "({0} ** 3)", "cube cubes"),
    Prim("abs", ("num",), "num", abs, "abs({0})", "absolute"),
    Prim("neg", ("num",), "num", lambda a: -a, "(-{0})", "negative negate"),
    Prim("half", ("num",), "float", lambda a: a / 2, "({0} / 2)", "half area triangle"),
    Prim("sqrt", ("num",), "float", math.sqrt, "math.sqrt({0})", "root square"),
    Prim("isqrt", ("int",), "int", math.isqrt, "math.isqrt({0})", "root square integer"),
    Prim("fact", ("int",), "int", lambda n: math.factorial(n) if 0 <= n <= 30 else 1 / 0, "math.factorial({0})", "factorial"),
    Prim("gcd", ("int", "int"), "int", math.gcd, "math.gcd({0}, {1})", "gcd greatest common divisor"),
    Prim("lcm", ("int", "int"), "int", lambda a, b: a * b // math.gcd(a, b) if a and b else 0, "({0} * {1} // math.gcd({0}, {1}))", "lcm least common multiple"),
    Prim("isprime", ("int",), "bool", _is_prime, "(({0}) > 1 and all(({0}) % d for d in range(2, int(({0}) ** 0.5) + 1)))", "prime"),
    Prim("even", ("int",), "bool", lambda n: n % 2 == 0, "({0} % 2 == 0)", "even"),
    Prim("odd", ("int",), "bool", lambda n: n % 2 == 1, "({0} % 2 == 1)", "odd"),
    Prim("digits", ("int",), "list", lambda n: [int(c) for c in str(abs(n))], "[int(c) for c in str(abs({0}))]", "digits digit"),
    Prim("revnum", ("int",), "int", lambda n: int(str(abs(n))[::-1]), "int(str(abs({0}))[::-1])", "reverse number"),
    Prim("bin", ("int",), "str", lambda n: bin(n)[2:], "bin({0})[2:]", "binary"),
    Prim("frombin", ("str",), "int", lambda s: int(s, 2), "int({0}, 2)", "binary decimal convert"),
    Prim("popcount", ("int",), "int", lambda n: bin(n).count("1"), "bin({0}).count('1')", "set bits count ones"),
    Prim("xor", ("int", "int"), "int", lambda a, b: a ^ b, "({0} ^ {1})", "xor bit differ"),
    Prim("ispow2", ("int",), "bool", lambda n: n > 0 and n & (n - 1) == 0, "({0} > 0 and {0} & ({0} - 1) == 0)", "power two bit"),
    Prim("eq", ("any", "any"), "bool", lambda a, b: a == b, "({0} == {1})", "equal same"),
    Prim("lt", ("num", "num"), "bool", lambda a, b: a < b, "({0} < {1})", "less smaller"),
    Prim("gt", ("num", "num"), "bool", lambda a, b: a > b, "({0} > {1})", "greater larger"),
    Prim("not", ("bool",), "bool", lambda a: not a, "(not {0})", "not"),
    Prim("and", ("bool", "bool"), "bool", lambda a, b: a and b, "({0} and {1})", "and both"),
    Prim("or", ("bool", "bool"), "bool", lambda a, b: a or b, "({0} or {1})", "or either"),
    Prim("maxn", ("num", "num"), "num", max, "max({0}, {1})", "maximum largest bigger"),
    Prim("minn", ("num", "num"), "num", min, "min({0}, {1})", "minimum smallest"),
    Prim("round2", ("num",), "float", lambda a: round(a, 2), "round({0}, 2)", "round"),
    Prim("int", ("any",), "int", int, "int({0})", "integer convert"),
    Prim("str", ("any",), "str", str, "str({0})", "string convert"),
    Prim("float", ("num",), "float", float, "float({0})", "float"),
    Prim("tri", ("int",), "int", lambda n: n * (n + 1) // 2, "({0} * ({0} + 1) // 2)", "triangular sum natural"),
    Prim("pi_r2", ("num",), "float", lambda r: math.pi * r * r, "(math.pi * {0} * {0})", "area circle"),
    Prim("circ", ("num",), "float", lambda r: 2 * math.pi * r, "(2 * math.pi * {0})", "circumference perimeter circle"),
    # sequences
    Prim("len", ("seq",), "int", len, "len({0})", "length count number size"),
    Prim("sum", ("coll",), "num", lambda xs: sum(xs), "sum({0})", "sum total add"),
    Prim("max", ("coll",), "any", max, "max({0})", "maximum largest biggest max longest"),
    Prim("min", ("coll",), "any", min, "min({0})", "minimum smallest min"),
    Prim("sorted", ("coll",), "list", sorted, "sorted({0})", "sort sorted ascending order"),
    Prim("sorted_desc", ("coll",), "list", lambda xs: sorted(xs, reverse=True), "sorted({0}, reverse=True)", "sort descending largest"),
    Prim("rev", ("seq",), "seq", lambda xs: xs[::-1], "{0}[::-1]", "reverse reversed"),
    Prim("first", ("seq",), "any", lambda xs: xs[0], "{0}[0]", "first"),
    Prim("last", ("seq",), "any", lambda xs: xs[-1], "{0}[-1]", "last"),
    Prim("take", ("seq", "int"), "seq", lambda xs, n: xs[:n], "{0}[:{1}]", "first n smallest largest take"),
    Prim("drop", ("seq", "int"), "seq", lambda xs, n: xs[n:], "{0}[{1}:]", "remove first skip"),
    Prim("at", ("seq", "int"), "any", lambda xs, i: xs[i], "{0}[{1}]", "nth index element second third"),
    Prim("tolist", ("any",), "list", list, "list({0})", "list convert"),
    Prim("totuple", ("any",), "tuple", tuple, "tuple({0})", "tuple convert"),
    Prim("toset", ("coll",), "set", set, "set({0})", "unique distinct set"),
    Prim("uniq", ("coll",), "list", _uniq, "list(dict.fromkeys({0}))", "unique distinct duplicates remove"),
    Prim("inter", ("coll", "coll"), "tuple", lambda a, b: tuple(sorted(set(a) & set(b))), "tuple(sorted(set({0}) & set({1})))", "common shared similar intersection"),
    Prim("union", ("coll", "coll"), "tuple", lambda a, b: tuple(sorted(set(a) | set(b))), "tuple(sorted(set({0}) | set({1})))", "union combine"),
    Prim("diff", ("coll", "coll"), "list", lambda a, b: [x for x in a if x not in b], "[x for x in {0} if x not in {1}]", "difference not in"),
    Prim("concat", ("seq", "seq"), "seq", lambda a, b: a + b, "({0} + {1})", "concatenate join append combine merge"),
    Prim("count", ("seq", "any"), "int", lambda xs, x: xs.count(x), "{0}.count({1})", "count occurrences frequency"),
    Prim("index", ("seq", "any"), "int", lambda xs, x: xs.index(x), "{0}.index({1})", "index position find"),
    Prim("contains", ("any", "seq"), "bool", lambda x, xs: x in xs, "({0} in {1})", "contains present exists check"),
    Prim("allsame", ("coll",), "bool", lambda xs: len(set(xs)) <= 1, "(len(set({0})) <= 1)", "same equal all identical"),
    Prim("alluniq", ("coll",), "bool", lambda xs: len(set(xs)) == len(xs), "(len(set({0})) == len(list({0})))", "unique distinct duplicate"),
    Prim("issorted", ("coll",), "bool", lambda xs: list(xs) == sorted(xs), "(list({0}) == sorted({0}))", "sorted order check"),
    Prim("prod", ("coll",), "num", lambda xs: _big(math.prod(xs)), "math.prod({0})", "product multiply"),
    Prim("mean", ("coll",), "float", lambda xs: sum(xs) / len(xs), "(sum({0}) / len({0}))", "average mean"),
    Prim("map_sq", ("coll",), "list", lambda xs: [x * x for x in xs], "[x * x for x in {0}]", "square squares each"),
    Prim("map_cube", ("coll",), "list", lambda xs: [x ** 3 for x in xs], "[x ** 3 for x in {0}]", "cube cubes each"),
    Prim("map_neg", ("coll",), "list", lambda xs: [-x for x in xs], "[-x for x in {0}]", "negative"),
    Prim("map_abs", ("coll",), "list", lambda xs: [abs(x) for x in xs], "[abs(x) for x in {0}]", "absolute"),
    Prim("map_str", ("coll",), "list", lambda xs: [str(x) for x in xs], "[str(x) for x in {0}]", "string convert"),
    Prim("map_int", ("coll",), "list", lambda xs: [int(x) for x in xs], "[int(x) for x in {0}]", "integer convert"),
    Prim("map_len", ("coll",), "list", lambda xs: [len(x) for x in xs], "[len(x) for x in {0}]", "length each"),
    Prim("map_first", ("coll",), "list", lambda xs: [x[0] for x in xs], "[x[0] for x in {0}]", "first each"),
    Prim("map_last", ("coll",), "list", lambda xs: [x[-1] for x in xs], "[x[-1] for x in {0}]", "last each"),
    Prim("map_sum", ("coll",), "list", lambda xs: [sum(x) for x in xs], "[sum(x) for x in {0}]", "sum each"),
    Prim("map_add", ("coll", "num"), "list", lambda xs, k: [x + k for x in xs], "[x + {1} for x in {0}]", "add each"),
    Prim("map_mul", ("coll", "num"), "list", lambda xs, k: [x * k for x in xs], "[x * {1} for x in {0}]", "multiply each double doubled twice"),
    Prim("zipadd", ("coll", "coll"), "list", lambda a, b: [x + y for x, y in zip(a, b)], "[x + y for x, y in zip({0}, {1})]", "add two lists elementwise"),
    Prim("zipmul", ("coll", "coll"), "list", lambda a, b: [x * y for x, y in zip(a, b)], "[x * y for x, y in zip({0}, {1})]", "multiply two lists"),
    Prim("zipsub", ("coll", "coll"), "list", lambda a, b: [x - y for x, y in zip(a, b)], "[x - y for x, y in zip({0}, {1})]", "subtract two lists"),
    Prim("zipdiv", ("coll", "coll"), "list", lambda a, b: [x / y for x, y in zip(a, b)], "[x / y for x, y in zip({0}, {1})]", "divide two lists"),
    Prim("zip", ("coll", "coll"), "list", lambda a, b: list(zip(a, b)), "list(zip({0}, {1}))", "pair zip"),
    Prim("f_even", ("coll",), "list", lambda xs: [x for x in xs if x % 2 == 0], "[x for x in {0} if x % 2 == 0]", "even"),
    Prim("f_odd", ("coll",), "list", lambda xs: [x for x in xs if x % 2 == 1], "[x for x in {0} if x % 2 == 1]", "odd"),
    Prim("f_pos", ("coll",), "list", lambda xs: [x for x in xs if x > 0], "[x for x in {0} if x > 0]", "positive"),
    Prim("f_neg", ("coll",), "list", lambda xs: [x for x in xs if x < 0], "[x for x in {0} if x < 0]", "negative"),
    Prim("f_gt", ("coll", "num"), "list", lambda xs, k: [x for x in xs if x > k], "[x for x in {0} if x > {1}]", "greater than"),
    Prim("f_lenge", ("coll", "int"), "list", lambda xs, k: [x for x in xs if len(x) >= k], "[x for x in {0} if len(x) >= {1}]", "at least long characters"),
    Prim("f_neq", ("coll", "any"), "list", lambda xs, k: [x for x in xs if x != k], "[x for x in {0} if x != {1}]", "remove all"),
    Prim("cumsum", ("coll",), "list", lambda xs: list(itertools.accumulate(xs)), "list(itertools.accumulate({0}))", "cumulative running sum"),
    Prim("flatten", ("coll",), "list", _flatten, "[y for x in {0} for y in (x if isinstance(x, (list, tuple)) else [x])]", "flatten nested"),
    Prim("mostcommon", ("coll",), "any", lambda xs: Counter(xs).most_common(1)[0][0], "collections.Counter({0}).most_common(1)[0][0]", "most common frequent maximum occurrences"),
    Prim("freq", ("coll",), "dict", lambda xs: dict(Counter(xs)), "dict(collections.Counter({0}))", "frequency count each"),
    Prim("sortkey1", ("coll",), "list", lambda xs: sorted(xs, key=lambda t: t[1]), "sorted({0}, key=lambda t: t[1])", "sort second element"),
    Prim("maxkey_len", ("coll",), "any", lambda xs: max(xs, key=len), "max({0}, key=len)", "longest maximum length"),
    Prim("minkey_len", ("coll",), "any", lambda xs: min(xs, key=len), "min({0}, key=len)", "shortest minimum length"),
    Prim("range1", ("int",), "list", lambda n: list(range(1, n + 1)) if n <= 10 ** 5 else 1 / 0, "list(range(1, {0} + 1))", "first n natural numbers range"),
    Prim("range0", ("int",), "list", lambda n: list(range(n)) if n <= 10 ** 5 else 1 / 0, "list(range({0}))", "range"),
    # strings
    Prim("upper", ("str",), "str", str.upper, "{0}.upper()", "upper uppercase capital"),
    Prim("lower", ("str",), "str", str.lower, "{0}.lower()", "lower lowercase"),
    Prim("title", ("str",), "str", str.title, "{0}.title()", "title capitalize first letter each word"),
    Prim("capitalize", ("str",), "str", str.capitalize, "{0}.capitalize()", "capitalize"),
    Prim("swapcase", ("str",), "str", str.swapcase, "{0}.swapcase()", "toggle swap case"),
    Prim("strip", ("str",), "str", str.strip, "{0}.strip()", "strip trim whitespace spaces"),
    Prim("nospace", ("str",), "str", lambda s: s.replace(" ", ""), "{0}.replace(' ', '')", "remove spaces whitespace"),
    Prim("words", ("str",), "list", str.split, "{0}.split()", "words split"),
    Prim("splitby", ("str", "str"), "list", lambda s, d: s.split(d), "{0}.split({1})", "split delimiter"),
    Prim("joinsp", ("coll",), "str", lambda xs: " ".join(xs), "' '.join({0})", "join words"),
    Prim("join", ("coll",), "str", lambda xs: "".join(xs), "''.join({0})", "join concatenate characters"),
    Prim("replace", ("str", "str", "str"), "str", lambda s, a, b: s.replace(a, b), "{0}.replace({1}, {2})", "replace"),
    Prim("remove", ("str", "str"), "str", lambda s, a: s.replace(a, ""), "{0}.replace({1}, '')", "remove delete"),
    Prim("remove1", ("str", "str"), "str", lambda s, a: s.replace(a, "", 1), "{0}.replace({1}, '', 1)", "remove first occurrence"),
    Prim("removelast", ("str", "str"), "str", lambda s, a: (s[:s.rfind(a)] + s[s.rfind(a) + 1:]) if a in s else s, "(({0}[:{0}.rfind({1})] + {0}[{0}.rfind({1}) + 1:]) if {1} in {0} else {0})", "remove last occurrence"),
    Prim("chars", ("str",), "list", list, "list({0})", "characters chars"),
    Prim("isdigit", ("str",), "bool", str.isdigit, "{0}.isdigit()", "digit numeric integer"),
    Prim("isalpha", ("str",), "bool", str.isalpha, "{0}.isalpha()", "alphabet letters"),
    Prim("isupper", ("str",), "bool", str.isupper, "{0}.isupper()", "upper"),
    Prim("islower", ("str",), "bool", str.islower, "{0}.islower()", "lower"),
    Prim("startswith", ("str", "str"), "bool", str.startswith, "{0}.startswith({1})", "starts begins"),
    Prim("endswith", ("str", "str"), "bool", str.endswith, "{0}.endswith({1})", "ends"),
    Prim("find", ("str", "str"), "int", str.find, "{0}.find({1})", "find position index occurrence"),
    Prim("f_upper", ("str",), "str", lambda s: "".join(c for c in s if c.isupper()), "''.join(c for c in {0} if c.isupper())", "uppercase characters extract"),
    Prim("f_lower", ("str",), "str", lambda s: "".join(c for c in s if c.islower()), "''.join(c for c in {0} if c.islower())", "lowercase characters remove"),
    Prim("f_digit", ("str",), "str", lambda s: "".join(c for c in s if c.isdigit()), "''.join(c for c in {0} if c.isdigit())", "digits numbers extract"),
    Prim("f_alpha", ("str",), "str", lambda s: "".join(c for c in s if c.isalpha()), "''.join(c for c in {0} if c.isalpha())", "letters remove non alphabet"),
    Prim("f_alnum", ("str",), "str", lambda s: "".join(c for c in s if c.isalnum()), "''.join(c for c in {0} if c.isalnum())", "alphanumeric remove special characters"),
    Prim("vowels", ("str",), "int", lambda s: sum(c in "aeiouAEIOU" for c in s), "sum(c in 'aeiouAEIOU' for c in {0})", "vowels count"),
    Prim("ispal", ("seq",), "bool", lambda s: s == s[::-1], "({0} == {0}[::-1])", "palindrome"),
    Prim("sortchars", ("str",), "str", lambda s: "".join(sorted(s)), "''.join(sorted({0}))", "sort characters"),
    Prim("rev_words", ("str",), "str", lambda s: " ".join(s.split()[::-1]), "' '.join({0}.split()[::-1])", "reverse words"),
    Prim("ord", ("str",), "int", lambda s: ord(s) if len(s) == 1 else 1 / 0, "ord({0})", "ascii value character"),
    Prim("chr", ("int",), "str", lambda n: chr(n) if 0 <= n < 1114111 else 1 / 0, "chr({0})", "character ascii"),
    Prim("strlen_words", ("str",), "list", lambda s: [len(w) for w in s.split()], "[len(w) for w in {0}.split()]", "length words"),
    Prim("firstchar", ("str",), "str", lambda s: s[0], "{0}[0]", "first character"),
    Prim("lastchar", ("str",), "str", lambda s: s[-1], "{0}[-1]", "last character"),
    Prim("upper_first_last", ("str",), "str", lambda s: " ".join((w[0].upper() + w[1:-1] + w[-1].upper()) if len(w) > 1 else w.upper() for w in s.split()), "' '.join((w[0].upper() + w[1:-1] + w[-1].upper()) if len(w) > 1 else w.upper() for w in {0}.split())", "capitalize first last letters each word"),
]
# v2 library entries added from the gate panel's FAILURE CLASSES (general-purpose operations; checked on MBPP dev)
P_ += [
    Prim("c2f", ("num",), "float", lambda c: c * 9 / 5 + 32, "({0} * 9 / 5 + 32)", "celsius fahrenheit convert temperature"),
    Prim("f2c", ("num",), "float", lambda f: (f - 32) * 5 / 9, "(({0} - 32) * 5 / 9)", "fahrenheit celsius convert temperature"),
    Prim("isleap", ("int",), "bool", lambda y: y % 4 == 0 and (y % 100 != 0 or y % 400 == 0),
         "({0} % 4 == 0 and ({0} % 100 != 0 or {0} % 400 == 0))", "leap year"),
    Prim("joincomma", ("coll",), "str", lambda xs: ",".join(xs), "','.join({0})", "join comma separated"),
    Prim("nwords", ("str",), "int", lambda s: len(s.split()), "len({0}.split())", "count words number"),
]
IMPORTS = "import math\nimport itertools\nimport collections\n"


# --------------------------------------------------------------------------------------------------------- request
def parse_tests(tests, setup=""):
    """'assert f(a, b) == v' -> (fname, [(args, expected)], wrapper-src or None)"""
    fname, cases, wrap = None, [], None
    for t in tests:
        try:
            node = ast.parse(t.strip()).body[0]
        except SyntaxError:
            return None
        if not isinstance(node, ast.Assert):
            return None
        test = node.test
        if isinstance(test, ast.Call) or (isinstance(test, ast.UnaryOp) and isinstance(test.operand, ast.Call)):
            neg = isinstance(test, ast.UnaryOp)
            call = test.operand if neg else test
            exp = not neg
            lhs = call
        elif isinstance(test, ast.Compare) and len(test.ops) == 1 and isinstance(test.ops[0], ast.Eq):
            lhs = test.left
            try:
                exp = ast.literal_eval(test.comparators[0])
            except Exception:  # noqa: BLE001
                return None
        else:
            return None
        call, w = lhs, None
        if isinstance(lhs, ast.Call) and isinstance(lhs.func, ast.Name) and lhs.func.id in ("set", "sorted", "list", "tuple", "round", "abs") and lhs.args and isinstance(lhs.args[0], ast.Call):
            w, call = lhs.func.id, lhs.args[0]
        if not (isinstance(call, ast.Call) and isinstance(call.func, ast.Name)):
            return None
        try:
            args = [ast.literal_eval(a) for a in call.args]
        except Exception:  # noqa: BLE001
            return None
        if fname and call.func.id != fname:
            return None
        fname, wrap = call.func.id, w or wrap
        cases.append((args, exp))
    return (fname, cases, wrap) if fname and cases else None


WRAP = {None: lambda v: v, "set": set, "sorted": sorted, "list": list, "tuple": tuple, "round": round, "abs": abs}


ORDINAL = {"first": 0, "second": 1, "third": 2, "fourth": 3, "last": -1, "double": 2, "doubled": 2, "twice": 2,
           "triple": 3, "tripled": 3, "half": 2, "halve": 2, "halved": 2}


def consts_of(desc, cases, strict=False):
    """constants a program may use. STRICT (one example): only what the request itself names - its numbers, its
    quoted characters and its ordinals (first=0, second=1, last=-1) - so no positional coincidence is reachable"""
    text = re.split(r"\bso\b", desc or "")[0] if strict else (desc or "")
    named = [int(n) for n in re.findall(r"\b(\d{1,4})\b", text)]
    named += re.findall(r"'([^']{1,3})'", text)
    named += [v for w, v in ORDINAL.items() if re.search(r"\b%s\b" % w, text, re.I)]
    out = named if strict else [0, 1, 2, -1, "", " "] + named
    return list(dict.fromkeys(out))


GENERIC = {"write", "function", "python", "find", "given", "number", "numbers", "integer", "list", "string", "element",
           "elements", "value", "values", "check", "whether", "two", "the", "and", "from", "all", "using", "which", "that",
           "with", "into", "for", "each", "array", "tuple", "item", "items", "get", "given", "not", "out", "sequence"}


def words_of(desc):
    return {w.lower().rstrip("s") for w in re.findall(r"[A-Za-z]+", desc or "")
            if len(w) >= 3 and w.lower() not in GENERIC and w.lower().rstrip("s") not in GENERIC}


# ----------------------------------------------------------------------------------------------------------- search
class CodeOrgan:
    def __init__(self, library=None, max_size=4, cap=60000):
        self.prims = list(P_) + list(library or [])
        self.max_size, self.cap = max_size, cap
        self.why = ""

    def learn(self, name, prog, arity_types, out_t, desc):
        """STORE: a verified program becomes a library function; later requests compose it"""
        fnc = prog["fn"]
        src = "%s(%s)" % (name, ", ".join("{%d}" % i for i in range(len(arity_types))))
        p = Prim("lib_" + name, tuple(arity_types), out_t, lambda *vals: fnc(vals), src, " ".join(words_of(desc)))
        self.prims.append(p)
        self.library_src = getattr(self, "library_src", {})
        self.library_src[name] = prog["def"].replace(IMPORTS, "")
        return p

    def solve(self, desc, tests, setup="", visible=None, strict=False):
        self.why = ""
        pt = parse_tests(tests, setup)
        if not pt:
            self.why = "the examples are not simple asserts"
            return None
        fname, cases, wrap = pt
        vis = cases if visible is None else cases[:visible]
        if not vis:
            return None
        W = WRAP[wrap]
        nargs = len(vis[0][0])
        if any(len(a) != nargs for a, _ in cases):
            self.why = "varying arity"
            return None
        target = [W(e) if wrap else e for _, e in vis]
        kw = words_of(desc)
        # the request SPECIFIES the moves: rank by the description's words; keep the whole library but try relevant first
        # a stored program is a move only with >= 2 examples, and never inside its own definition; with ONE example
        # (strict) only the base library composes - forms verified on one example are too weak to build on from one more
        pool = [p for p in self.prims if p.name != "lib_" + fname and not (strict and p.name.startswith("lib_"))]
        prims = sorted(pool, key=lambda p: -len(p.kw & kw))
        # terminals: arguments and constants; a term = (src, type, outputs-tuple, fn)
        terms_by_size = {1: []}
        seen = set()

        def add(size, src, outs, fnc, used=frozenset()):
            key = repr(outs)
            try:
                match = [W(o) if wrap else o for o in outs] == target
            except Exception:  # noqa: BLE001
                match = False
            if match:
                # the request SPECIFIES the moves: a fit made only of moves the description never names is a
                # coincidence of three examples, not an answer - it is not accepted and does not claim the signature.
                # A fit must also READ its inputs: every argument is used, except an 'n' that only restates a length.
                reads = all(argn[i] in re.findall(r"a\d+", src) or restates_len[i] for i in range(nargs))
                # STRICT (one example): EVERY operation must be named by the request - the description specifies the
                # program and the single example only checks it
                named = all(pr.kw & kw for pr in used) if strict else any(pr.kw & kw for pr in used)
                if used and reads and named:
                    return (src, fnc)
                self.coincidences = getattr(self, "coincidences", 0) + 1
                return None
            if key in seen or len(key) > 1500 or len(seen) > self.cap:
                return None                               # bounded memory: huge values and past-budget terms are not kept
            seen.add(key)
            ty = tyof(outs[0])
            if any(tyof(o) != ty for o in outs):
                ty = "any"
            terms_by_size.setdefault(size, []).append((src, ty, outs, fnc, used))
            return None

        argn = ["a%d" % i for i in range(nargs)]
        restates_len = [all(isinstance(a[i], int) and any(isinstance(a[j], (list, tuple, str)) and len(a[j]) == a[i]
                                                         for j in range(nargs) if j != i) for a, _ in vis)
                        for i in range(nargs)]
        for i in range(nargs):
            hit = add(1, argn[i], tuple(a[i] for a, _ in vis), (lambda i: lambda xs: xs[i])(i))
            if hit:
                return self._emit(fname, argn, hit, cases, wrap, desc)
        for c in consts_of(desc, vis, strict):
            add(1, repr(c), tuple(c for _ in vis), (lambda c: lambda xs: c)(c))
        n_terms = 0
        best = None
        for size in range(2, self.max_size + 1):
            for p in prims:
                k = len(p.ins)
                # child sizes summing to size-1
                for parts in _compositions(size - 1, k):
                    pools = [terms_by_size.get(s, []) for s in parts]
                    pools = [[t for t in pool if fits(t[1], want) or t[1] == "any"] for pool, want in zip(pools, p.ins)]
                    if any(not pool for pool in pools):
                        continue
                    for combo in itertools.product(*pools):
                        n_terms += 1
                        if n_terms > self.cap:
                            self.why = "search budget exhausted"
                            return None
                        try:
                            outs = tuple(p.fn(*[t[2][j] for t in combo]) for j in range(len(vis)))
                        except Exception:  # noqa: BLE001
                            continue
                        if any(_oversize(o) for o in outs):
                            continue                      # bounded memory: never represent a huge intermediate value
                        src = p.src.format(*[t[0] for t in combo])
                        fnc = (lambda p, combo: lambda xs: p.fn(*[t[3](xs) for t in combo]))(p, combo)
                        used = frozenset([p]).union(*[t[4] for t in combo])
                        hit = add(size, src, outs, fnc, used)
                        if hit:
                            if not strict:
                                return self._emit(fname, argn, hit, cases, wrap, desc)
                            # COVERAGE: how much of what the request says the program's moves account for, jointly
                            named = len(set().union(*[pr.kw for pr in used]) & kw)
                            if best is None or named > best[0]:
                                best = (named, hit)
        if strict and best:                               # v2 STRICT: the fit covering most of the request, any size
            return self._emit(fname, argn, best[1], cases, wrap, desc)
        self.why = "no program of size <= %d fits the examples" % self.max_size
        return None

    def _emit(self, fname, argn, hit, cases, wrap, desc):
        src, fnc = hit
        lib = getattr(self, "library_src", {})
        used = [n for n in lib if re.search(r"\b%s\(" % re.escape(n), src)]
        deps = "".join(lib[n] for n in used)             # stored forms this program composes
        body = "%s%sdef %s(%s):\n    return %s\n" % (IMPORTS, deps, fname, ", ".join(argn), src)
        # CHECK: execute the RENDERED source (not the in-memory closure) against every example given
        ns = {}
        exec(compile(body, "<organ>", "exec"), ns)  # noqa: S102 - composed only of library templates
        W = WRAP[wrap]
        passed = 0
        for a, e in cases:
            try:
                out = ns[fname](*a)
                ok = (W(out) if wrap else out) == e
            except Exception:  # noqa: BLE001
                ok = False
            passed += ok
        return {"def": body, "fname": fname, "expr": src, "passed": passed, "cases": len(cases), "fn": fnc,
                "arity": len(argn), "reused": used}


def _oversize(o):
    if isinstance(o, bool):
        return False
    if isinstance(o, int):
        return abs(o) > 10 ** 12
    if isinstance(o, str):
        return len(o) > 400
    if isinstance(o, (list, tuple, set, dict)):
        return len(o) > 200 or any(_oversize(x) for x in list(o)[:50])
    return False


def _compositions(total, k):
    if k == 1:
        yield (total,)
        return
    for first in range(1, total - k + 2):
        for rest in _compositions(total - first, k - 1):
            yield (first,) + rest