program_generation

#!/usr/bin/python
import argparse
import random
import subprocess

a_expressions = ["int", "id", "+", "-", "*", "input"]
b_expressions = ["true", "false", "<", "=", "not", "and", "or"]
statements = ["if", "while", "output", "ass", "skip", "semi"]
vars = []
max_int = 10
max_bexpr_depth = 2
max_aexpr_depth = 2
max_depth = 4
tab = "   "


def rand_inst(arr):
    return arr[random.randint(0, len(arr) - 1)]


def gen_if(cond, then_stmt, else_stmt, indent):
    return f"{indent*tab}if {cond} then {{\n{then_stmt}\n{(indent) * tab}}} else {{\n{else_stmt}\n{(indent)*tab}}}"


def gen_while(cond, do_stmt, indent):
    return f"{indent*tab}while ({cond}) do {{\n{do_stmt}\n{(indent)*tab}}}"


def gen_output(aexpr, indent):
    return f"{indent*tab}output {aexpr}"


def gen_ass(var, aexpr, indent):
    return f"{indent*tab}{var} := {aexpr}"


def gen_var_dec():
    new_var = f"v{len(vars)}"
    vars.append(new_var)
    return f"var {new_var}"


def gen_skip(indent):
    return f"{indent*tab}skip"


def gen_semi(first_stmt, second_stmt):
    return f"{first_stmt};\n{second_stmt}"


def gen_op(op, lhs, rhs):
    return f"{lhs} {op} {rhs}"


def gen_input():
    return "input"


def random_var():
    index = random.randint(0, len(vars) - 1)
    if index >= 0:
        return vars[index]
    else:
        raise Exception("Program had no variables to use")


def random_aexpr(max_depth):
    if (max_depth > max_aexpr_depth):
        max_depth = max_aexpr_depth

    if (max_depth <= 1):
        return random.randint(0, max_int)

    aexpr = rand_inst(a_expressions)
    match aexpr:
        case "int":
            return random.randint(0, max_int)
        case "id":
            return random_var()
        case "+" | "-" | "*":
            return f"({gen_op(aexpr, random_aexpr(max_depth - 1),
                              random_aexpr(max_depth - 1))})"
        case "input":
            return gen_input()
        case _:
            raise Exception("Error, generation of aexpr failed")


def random_bexpr(max_depth):
    if (max_depth > max_bexpr_depth):
        max_depth = max_bexpr_depth

    if (max_depth <= 1):
        if (random.randint(0, 1) == 0):
            return "true"
        else:
            return "false"

    bexpr = rand_inst(b_expressions)

    match bexpr:
        case "true":
            return "true"
        case "false":
            return "false"
        case "not":
            return f"not({random_bexpr(max_depth - 1)})"
        case "or" | "and":
            return f"({gen_op(bexpr, random_bexpr(max_depth - 1),
                              random_bexpr(max_depth - 1))})"
        case "<" | "=":
            return f"({gen_op(bexpr, random_aexpr(max_depth - 1),
                              random_aexpr(max_depth - 1))})"
        case _:
            raise Exception("Error, generation of bexpr failed")


def random_stmt(max_depth, indent):
    if (max_depth <= 1):
        match  random.randint(0, 2):
            case 0:
                return gen_skip(indent)
            case 1:
                return gen_ass(random_var(), random_aexpr(1), indent)

    stmt = rand_inst(statements)

    match stmt:
        case "skip":
            return gen_skip(indent)
        case "if":
            return gen_if(random_bexpr(max_depth - 1),
                          random_stmt(max_depth - 1, indent + 1),
                          random_stmt(max_depth - 1, indent + 1), indent)
        case "while":
            return gen_while(random_bexpr(max_depth - 1),
                             random_stmt(max_depth - 1, indent + 1), indent)
        case "output":
            return gen_output(random_aexpr(max_depth - 1), indent)
        case "ass":
            return gen_ass(random_var(), random_aexpr(max_depth - 1), indent)
        case "semi":
            return gen_semi(random_stmt(max_depth - 1, indent),
                            random_stmt(max_depth - 1, indent))
        case _:
            raise Exception("Error, generation of stmt failed")


def gen_fully_random_program(max_loc):
    global max_depth
    global tab
    main_stmt = f"fun main() {{\n{tab}"
    loc = 1

    number_of_vars = 5
    if max_loc < 100:
        number_of_vars = 10
    elif max_loc < 500:
        number_of_vars = 20
    else:
        number_of_vars = 30

    for i in range(number_of_vars):
        main_stmt += gen_var_dec() + "; "
        loc += 1

    main_stmt += "\n"
    while loc < max_loc:
        new_stmt = f"{random_stmt(max_depth, 1)};\n"
        main_stmt += new_stmt
        loc += new_stmt.count("\n")

    return main_stmt + "\n}"


def gen_partial_random_program(max_loc):
    # Template blocks where generated by chat gpt
    TEMPLATE_BLOCKS = [

        # Constant folding and propagation
        "x := 2 + 3 + 5;\ny := x + 1 - 2;\noutput y;",

        # Constant folding with nested expressions
        "a := 10;\nb := a + 0;\nc := b * 1;\noutput c;",

        # Dead code elimination - if false branch
        "if false then {\n  x := 42;\n} else {\n    skip;\n};",

        # Dead code elimination - redundant else
        "if true then {\n    y := 1;\n} else {\n    y := 999;\n};\noutput y;",

        # Unused variable
        "a := 5;\nb := a + 3;\n// b is never used",

        # Redundant assignment
        "x := 7;\nx := x;\nx := x + 0;\noutput x;",

        # Assignment overwritten before use
        "z := 1;\nz := 2;\noutput z;",

        # Constant equality check
        "x := 3;\nif x = 3 then {\n    output 1;\n} else {\n    output 0;\n};",

        # Unreachable loop
        "x := 10;\nwhile (x < 0) do {\n    x := x + 1;\n};",

        # Dead code inside branches
        "if false then {\n    y := 100;\n    output y;\n} else {\n    skip;\n};",

        # Multiple constant uses
        "a := 4;\nb := a + 4 + a;\nc := b * 2;\noutput c;",

        # Constant comparison short-circuiting
        "if 3 < 5 then {\n    x := 1;\n} else {\n    x := 2;\n};\noutput x;",

        # Redundant operations in sequence
        "x := 0;\nx := x + 1;\nx := x + 0;\nx := x - 0;\noutput x;",

        # Conditional overwrite
        "x := 5;\nif true then {\n    x := 6;\n} else {\n    x := 7;\n};\noutput x;"

        # Factorial
        "y := input;\nx := 1;\nwhile 1 < x do {\ny:=x * y;\n x:=x - 1\n};\noutput y;"

        # For loop
        "i := 1;\nwhile i < input do {i := i + 1};\noutput i;"
    ]

    main_stmt = "var x; var y; var z; var a; var b; var c; var i;\n"

    loc = 0
    while loc < max_loc:
        next_inst = rand_inst(TEMPLATE_BLOCKS) + "\n"
        loc += next_inst.count("\n")
        main_stmt += next_inst

    return f"fun main() {{\n{main_stmt}\n}}"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-loc", "--line-of-code", dest="loc", type=int)
    parser.add_argument("-f", "--fully-random", dest="fully_random", type=bool)
    parser.add_argument("-p", "--partial-random",
                        dest="partial_random", type=bool)
    args = parser.parse_args()

    max_loc = args.loc

    if args.partial_random:
        generated = gen_partial_random_program(max_loc)
    elif args.fully_random:
        generated = gen_fully_random_program(max_loc)
    else:
        print("Error, invalid mode.")
        return

    # Write to file
    path = "generated/generated.while"
    with open(path, "w") as f:
        f.write(generated)

    print(f"Generated program saved to {path}")

    try:
        subprocess.run(["./build/while", "-s", "-p",
                       path, "-l" "Debug"], check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running analyzer: {e}")


if __name__ == "__main__":
    main()