#!/usr/bin/env python3

import os;
import readline
import argparse
import pickle;
import sys;

def parse_args(argv):
    parser = argparse.ArgumentParser(prog = '4-variable-simplifier');

    parser.add_argument('-c', '--command', help = '''
        Provide a boolean expression to simpliy on the command-line.
        Without this option, boolean expressions are read stdin.
    ''')

    parser.add_argument('-e', '--extended-operators', action='store_true', help='''
        Use more than just the standard three boolean operators (not, or, and).
        This options allows use of: nor, nand, xor, nxor, andn, orn.
    ''')

    parser.add_argument('-o', '--custom-operators', help='''
        Pick and choose which operators the simplifier can use. Comma seperated.
        You can refer to them using their names or their symbols.
    ''')

    parser.add_argument('-E', '--experimental-operators', help='''
        For two boolean variables, there are four possible inputs into a
        binary operator. For each of those four an operator could return true
        or false. That would mean in theory there could exist up to 16
        theoretical operators with different behaviours. Even with the
        'extended operators' option, we're only introducing 9 possible
        operators, that's less that 60% of the possibilities!
        Some of these would of course be useless, like a binary operator that
        always produces 'false' regardless of the input, for instance, but
        other may be extremely useful for acting in place of much more complex
        trees of operators. In theory, the simplifier would produce the
        smallest most compact expression trees possible.

        Of course, no one would implement something like this and dare
        to anger the logic gods...
    ''')

    parser.add_argument("-n", '--use-names', help='''
        Print the operator's names in the expression tree, rather than the
        symbols.
    ''')

    parser.add_argument('-p', "--printout", action = 'store_true', help='''
        Print out all simplified expressions for the given expressions and
        quit.
    ''')

    parser.add_argument("--simplifications-file", \
        default = ".simplifications.bin", help='''
        Path to pickle file that stores the cached simplifications.
    ''')

    parser.add_argument("-y", "--yes", action = 'store_true', help='''
        Do not ask before generating simplification cache.
    ''')

    parser.add_argument("-q", "--quiet", action = 'store_true', help='''
        Do not print anything but the simplification. Assumes '--yes'.
    ''')

    parser.add_argument("-C", "--color", default = 'auto', \
        choices = ["off", "on", "auto"], help='''
        Select whether to use terminal colors. Default is to use if it is
        supported. use '--color=on' to force it on. '--color=off' otherwise.
    ''')

    return parser.parse_args(argv[1:]);

symbol_to_name = {
    "!": "not",
    "||": "or",
    "&&": "and",
    "!|": "nor",
    "!&": "nand",
    "!=": "xor",
    "==": "nxor",
    "|!": "orn",
    "&!": "andn",
};

def determine_available_operators(args):
    standard = ("!", "||", "&&");

    extended = ("!|", "!&", "!=", "==", "|!", "&!");

    experimental = (
        "<0000>",
        "<0001>",
        "<0010>",
        "<0011>",
        "<0100>",
        "<0101>",
        "<0110>",
        "<0111>",
        "<1000>",
        "<1001>",
        "<1010>",
        "<1011>",
        "<1100>",
        "<1101>",
        "<1110>",
        "<1111>");

    if args.extended_operators:
        return set(standard + extended);
    elif args.custom_operators:
        available_operators = set();

        for o in args.custom_operators.split(","):
            if o in symbol_to_name:
                available_operators.add(o);
            elif o in symbol_to_name.values():
                for k, v in symbol_to_name.items():
                    if v == o:
                        available_operators.add(k);
            else:
                raise BaseException(f"not an operator: '{o}'");

        return available_operators;
    elif args.experimental_operators:
        return set(experimental);
    else:
        return set(standard);

def pretty(exp):
    match exp:
        case ("literal", x):
            return str(x);

        case ("variable", x):
            return x

        case (op, inner):
            return f"({op}{pretty(inner)})"

        case (op, left, right):
            return f"({pretty(left)} {op} {pretty(right)})"

        case _:
            print(exp);
            assert(not "TODO");

W = 0b0101_0101_0101_0101
X = 0b0011_0011_0011_0011
Y = 0b0000_1111_0000_1111
Z = 0b0000_0000_1111_1111

M = 0b1111_1111_1111_1111

def calculate_simplifications(args, available_operators):
    lookup = dict() # truthtable -> expression

    costs = dict() # truthtable -> cost

    todo = [set() for _ in range(100)] # indexed by cost, set of truthtables.

    todo_count = [0]; # I have to wrap in an array because python is dumb.

    def prequeue(truthtable, expression, cost):
        lookup[truthtable] = expression;
        costs[truthtable] = cost;
        todo[cost].add(truthtable);
        todo_count[0] += 1;

# literals:
    prequeue(0b1111_1111_1111_1111, ("literal", 1), cost = 1);
    prequeue(0b0000_0000_0000_0000, ("literal", 0), cost = 1);

    # variables:
    prequeue(W, ("variable", "w"), cost = 0);
    prequeue(X, ("variable", "x"), cost = 0);
    prequeue(Y, ("variable", "y"), cost = 0);
    prequeue(Z, ("variable", "z"), cost = 0);

    # Completely unnecessary critera, alphabetical variables is more
    # aesthetically pleasing. Thanks Benson.
    def is_aesthetically_better(this, that):
        def extract_variables(exp):
            match exp:
                case ("literal", x):
                    return ();
                case ("variable", x):
                    return (x, )
                case (_, inner):
                    return extract_variables(inner);
                case (_, left, right):
                    return extract_variables(left) + extract_variables(right);
                case _:
                    print(exp);
                    assert(not "TODO");

        return extract_variables(this) < extract_variables(that);

    unary_operators = {
        '!': lambda x: ~x,
    }

    binary_operators = {
        '||': lambda x, y: (x | y) & M, # or
        '&&': lambda x, y: (x & y) & M, # and

        '!|': lambda x, y: ~(x | y) & M, # nor
        '!&': lambda x, y: ~(x & y) & M, # nand

        '&!': lambda x, y: (~x & y) & M, # andn
        '|!': lambda x, y: (~x | y) & M, # orn

        '!=': lambda x, y: (x ^ y) & M, # xor
        '==': lambda x, y: (~(x ^ y)) & M, # nxor

        # these were found by running the simplifier with -o xor,and,or,not
        # 1100 - x
        # 1010 - y
        "<0000>": lambda x, y: 0, #                                                                                                         0
        "<0001>": lambda x, y: 0, #                                                                               ((x == 0) && (y == 0)) || 0
        "<0010>": lambda x, y: 0, #                                                     ((x == 0) && (y == 1)) ||                           0
        "<0011>": lambda x, y: 0, #                                                     ((x == 0) && (y == 1)) || ((x == 0) && (y == 0)) || 0
        "<0100>": lambda x, y: 0, #                           ((x == 1) && (y == 0)) ||                                                     0
        "<0101>": lambda x, y: 0, #                           ((x == 1) && (y == 0)) ||                           ((x == 0) && (y == 0)) || 0
        "<0110>": lambda x, y: 0, #                           ((x == 1) && (y == 0)) || ((x == 0) && (y == 1)) ||                           0
        "<0111>": lambda x, y: 0, #                           ((x == 1) && (y == 0)) || ((x == 0) && (y == 1)) || ((x == 0) && (y == 0)) || 0
        "<1000>": lambda x, y: 0, # ((x == 1) && (y == 1)) ||                                                                               0
        "<1001>": lambda x, y: 0, # ((x == 1) && (y == 1)) ||                                                     ((x == 0) && (y == 0)) || 0
        "<1010>": lambda x, y: 0, # ((x == 1) && (y == 1)) ||                           ((x == 0) && (y == 1)) ||                           0
        "<1011>": lambda x, y: 0, # ((x == 1) && (y == 1)) ||                           ((x == 0) && (y == 1)) || ((x == 0) && (y == 0)) || 0
        "<1100>": lambda x, y: 0, # ((x == 1) && (y == 1)) || ((x == 1) && (y == 0)) ||                                                     0
        "<1101>": lambda x, y: 0, # ((x == 1) && (y == 1)) || ((x == 1) && (y == 0)) ||                           ((x == 0) && (y == 0)) || 0
        "<1110>": lambda x, y: 0, # ((x == 1) && (y == 1)) || ((x == 1) && (y == 0)) || ((x == 0) && (y == 1)) ||                           0
        "<1111>": lambda x, y: M, # ((x == 1) && (y == 1)) || ((x == 1) && (y == 0)) || ((x == 0) && (y == 1)) || ((x == 0) && (y == 0)) || 0
    }

    def print_status():
        numerator = 65536 - todo_count[0];
        denominator = 65536

        line = f'{numerator} of 65536';
        line += f' ({numerator / denominator * 100:.2f}%):'
        line += f' [{my_cost}]';
        line += f' {pretty(my_expression)}';

        print(line);

    min_cost = 0;

    while todo_count[0]:
        truthtables = todo[min_cost];

        if not truthtables:
            min_cost += 1;
            continue;

        assert(todo_count[0] <= 65536);

        my_truthtable = min(truthtables);
        truthtables.discard(my_truthtable);
        my_cost = min_cost;
        my_expression = lookup[my_truthtable];
        todo_count[0] -= 1;

        if not args.quiet:
            print_status();

        def consider(new_truthtable, new_expression, new_cost):
            if new_truthtable not in costs:
                todo[new_cost].add(new_truthtable);

                costs[new_truthtable] = new_cost
                lookup[new_truthtable] = new_expression

                todo_count[0] += 1;
            elif new_cost < costs[new_truthtable] or \
                    (new_cost == costs[new_truthtable] and \
                        is_aesthetically_better(new_expression, \
                            lookup[new_truthtable])):
                current_cost = costs[new_truthtable];

                assert(new_cost >= min_cost);

                todo[current_cost].discard(new_truthtable);
                todo[new_cost].add(new_truthtable);

                costs[new_truthtable] = new_cost
                lookup[new_truthtable] = new_expression

        # consider unary operators:
        for name, function in sorted(unary_operators.items()):
            if name in available_operators:
                unary_truthtable = function(my_truthtable) & M;
                unary_expression = (name, my_expression);
                unary_cost = my_cost + 1;

                consider(unary_truthtable, unary_expression, unary_cost);

        # consider binary operators:
        for name, function in sorted(binary_operators.items()):
            if name in available_operators:
                for other_truthtable, other_expression in sorted(lookup.items()):
                    # x + y
                    binary_truthtable = function(my_truthtable, other_truthtable);
                    binary_truthtable = binary_truthtable & M
                    binary_expression = (name, my_expression, other_expression);
                    binary_cost = my_cost + 1 + costs[other_truthtable];

                    consider(binary_truthtable, binary_expression, binary_cost);

                    # y + x
                    binary_truthtable = function(other_truthtable, my_truthtable);
                    binary_truthtable = binary_truthtable & M
                    binary_expression = (name, other_expression, my_expression);
                    binary_cost = my_cost + 1 + costs[other_truthtable];

                    consider(binary_truthtable, binary_expression, binary_cost);

    return costs, lookup

pathname = "simplifications.bin"

def get_simplifications(args, available_operators):
    available_operators = tuple(sorted(available_operators));

    try:
        with open(pathname, "rb") as stream:
            cache = pickle.load(stream);
    except FileNotFoundError:
        cache = dict();

    if available_operators not in cache:
        if not args.quiet:
            print("Oh! looks like you're running this for the first time");
            print("I'll have to build up my cache of simplifications");
            print("This may take a while.");
            print("I'll only have to do this once.");

            if not args.yes:
                print();
                input("any input to start:");
                print();

        bundle = calculate_simplifications(args, available_operators);

        if not args.quiet:
            print();
            print("done!");
            print();

        cache[available_operators] = bundle;

        with open(pathname, "wb") as stream:
            pickle.dump(cache, stream);

        if not args.quiet:
            print();
            print("saved!");
            print();

    return cache[available_operators];

def create_parser():
    import pyparsing as pp

    from pyparsing import Forward, Suppress, Keyword, Group, ZeroOrMore, Optional, Literal

    root = Forward()

    literal = pp.Word('01');

    variable = pp.Word('wxyz')

    highest = literal | variable | (Suppress('(') + root + Suppress(')'));

    prefix_expression = Group(Literal("!") + highest) | highest

    relational_expression = \
          Group(prefix_expression + Literal('<=')  + prefix_expression) \
        | Group(prefix_expression + Literal('>=')  + prefix_expression) \
        | Group(prefix_expression + Literal('>') + prefix_expression) \
        | Group(prefix_expression + Literal('<') + prefix_expression) \
        | Group(prefix_expression + Literal('==') + prefix_expression) \
        | Group(prefix_expression + Literal('!=') + prefix_expression) \
        | prefix_expression;

    logical_and_expression = Forward();
    logical_and_expression <<= \
          Group(relational_expression + Literal('&&') + logical_and_expression) \
        | Group(relational_expression + Literal('!&') + logical_and_expression) \
        | Group(relational_expression + Literal('&!') + logical_and_expression) \
        | relational_expression;

    logical_or_expression = Forward()
    logical_or_expression <<= \
          Group(logical_and_expression + Literal('||') + logical_or_expression) \
        | Group(logical_and_expression + Literal('!|') + logical_or_expression) \
        | Group(logical_and_expression + Literal('|!') + logical_or_expression) \
        | logical_and_expression;

    root <<= logical_or_expression;

    return root;


def parse(parser, text):
    return parser.parseString(text, parseAll = True).asList()[0]

def evaluate(expr):
    match expr:
        case "0":
            return 0;
        case "1":
            return M;
        case "w":
            return W;
        case "x":
            return X;
        case "y":
            return Y;
        case "z":
            return Z;
        case ("!", subexp):
            return ~evaluate(subexp) & M;
        case (left, "||", right):
            return evaluate(left) | evaluate(right);
        case (left, "!|", right):
            return ~(evaluate(left) | evaluate(right)) & M;
        case (left, "|!", right):
            return (~evaluate(left) | evaluate(right)) & M;
        case (left, "&&", right):
            return evaluate(left) & evaluate(right);
        case (left, "!&", right):
            return ~(evaluate(left) & evaluate(right)) & M;
        case (left, "&!", right):
            return (~evaluate(left) & evaluate(right)) & M;
        case (left, "<", right):
            return (~evaluate(left) & evaluate(right)) & M;
        case (left, "<=", right):
            return (~evaluate(left) | evaluate(right)) & M;
        case (left, ">", right):
            return (evaluate(left) & ~evaluate(right)) & M;
        case (left, ">=", right):
            return (evaluate(left) | ~evaluate(right)) & M;
        case (left, "==", right):
            return ~(evaluate(left) ^ evaluate(right)) & M;
        case (left, "!=", right):
            return evaluate(left) ^ evaluate(right);
        case _:
            print(exp);
            assert(not "TODO");

def repl(args, cost, lookup):
    import readline;
    print("Give a boolean expression using the variables 'w', 'x', 'y' and 'z'.");
    print("Operations: not ('!'), or ('||'), and ('&&').""");

    print();

    if args.extended_operators:
        print("Extended operations: nor ('!|'), orn ('|!'), nand ('!&'), ");
        print("andn ('&!'), xor ('!='), and nxor ('==')");

        print();

    # print(f'I can do anything in {max(cost.values())} operations.')
    # print();

    parser = create_parser();

    while True:
        try:
            line = input(">>> ");
        except EOFError:
            return;

        line = line.strip();

        if not line: continue;

        truthtable = evaluate(parse(parser, line));

        if truthtable in lookup:
            print(f'{cost[truthtable]}: {pretty(lookup[truthtable])}')
        else:
            print('unreachable.')

def main(args):
    available_operators = determine_available_operators(args);

    match args.color:
        case 'off': args.color = False;
        case 'on': args.color = True;
        case 'auto': args.color = os.isatty(1);

    cost, lookup = get_simplifications(args, available_operators);

    if args.printout:
        for truthtable, exp in sorted(lookup.items()):
            print(f'{truthtable:016b}: {pretty(exp)}');
    elif args.command:
        truthtable = evaluate(parse(create_parser(), args.command));

        if truthtable in lookup:
            print(pretty(lookup[truthtable]));
        else:
            print('unreachable.')
    else:
        repl(args, cost, lookup);

    return 0;

exit(main(parse_args(sys.argv)))