import numpy as np
import subprocess
import os 
import itertools
import numpy as np


def _generate_arrays(depth, max_size, max_scalar):
    arr = []
    if depth == 0:
        # Generate both scalars and arrays at leaf level
        if max_size > 0: 
            arr.append(list(range(max_size)))
        if max_scalar > 0: 
            arr.extend(list(range(max_scalar)))
    else:
        # Generate both scalars and arrays at each level
        for sizes in itertools.product(range(max_size), repeat=max_size):
            for sub_arrays in itertools.product(*[generate_arrays(depth - 1, size, max_scalar) for size in sizes]):
                arr.append(list(sub_arrays)) # array
        if max_scalar > 0: 
            arr.extend(list(range(max_scalar)))
    for p in itertools.permutations(arr):
        yield list(p)


def generate_random_arrays(depth, max_size, max_scalar, scalar_types=['%d','%s']):
    # Base case: depth 0, return a random scalar
    if depth == 0:
        size = np.random.randint(max_scalar+1)
        arr = np.random.choice(scalar_types) if size == 0 else ['%d'] * size
    else:
        # Decide randomly whether to generate a deeper array or stop at a scalar
        if np.random.choice([True, False]):
            # Generate a deeper array
            size = np.random.randint(max_size + 1)  # Random size, including 0
            arr = [generate_random_arrays(depth-1, max_size, max_scalar) for _ in range(size)]
        else:
            # Stop at a scalar
            size = np.random.randint(max_scalar+1)
            arr = np.random.choice(scalar_types) if size == 0 else ['%d'] * size
    return arr


def run_test(test_case):
    # Prepare the command and the input
    cmd = [os.environ['KG_PATH'], "-e", test_case]
    
    # Run the command with the input and capture the output
    result = subprocess.run(cmd, text=True, capture_output=True)
    
    # Return the result
    return result.stdout


def fill_template(s):
    params = []
    i = 0
    d = 0
    while(i < len(s)): 
        if s[i] == '%':
           i += 1
           if s[i] == 'd':
               params.append(d)
           elif s[i] == 's':
               params.append('"'+chr(ord('a')+(d % 26))+'"')
           d += 1
        i += 1 
    s = s % tuple(params)
    return s


def gen_tests(max_test_cases=10000, max_depth=3, max_size=3, max_scalar=3):
    # Prepare all possible test cases
    seen = set()
    while len(seen) < max_test_cases: 
        array = generate_random_arrays(np.random.randint(max_depth), max_size, max_scalar)
        # Convert numpy array to Klong array string representation
        array_str = (array if isinstance(array, str) else str(array)).replace(",","").replace("'",'"').replace('"%d"', '%d').replace('"%s"','%s')
        test_case = ",/" + array_str
        if test_case in seen:
            continue
        seen.add(test_case)

    seen = list(seen)
    seen.sort(key=lambda x: (len(x),x))

    test_cases = set()
    for test_case in seen:
        test_case = fill_template(test_case)
        pretty_case = test_case.replace('"','""')
        output = run_test(test_case).strip()
        s = f"""t("{pretty_case}"; {test_case}; {output})"""
        test_cases.add(s)

    # Print the test strings
    test_cases = list(test_cases)
    test_cases.sort(key=lambda x: (len(x),x))
    for o in test_cases:
        print(o)



if __name__ == '__main__':
    print("""
:"DO NOT MODIFY: GENERATED BY gen_join_over.py"

err::0
wl::{.w(x);.p("")}
fail::{err::1;.d("failed: ");.p(x);.d("expected: ");wl(z);.d("got: ");wl(y);[]}
t::{:[~y~z;fail(x;y;z);[]]}

""")
    gen_tests()

    print("""
:[err;[];.p("ok!")]
""")

