#!/usr/bin/env python3

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

from colorsys import hsv_to_rgb;

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('-E', '--extra-extended-operators', action='store_true', help='''
        In addition to the operators of '--extended-operators' introduce
        the ternary ('?:') operator.
        This operator is very powerful but dramatically increases
        simplification time by an nth degree.
    ''')

    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("-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",
    "?:": "ternary",
};

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

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

    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.extra_extended_operators:
        return set(standard + extended + ("?:", ));
    else:
        return set(standard);

LITERAL_ESCAPE = "\033[38;2;200;200;100m";
VARIABLE_ESCAPE = "\033[38;2;100;100;200m";

operator_colors = [];

for i in range(10):
    r, g, b = hsv_to_rgb(i / 10, 1.0, 0.8);
    r, g, b = [int(255 * x) for x in (r, g, b)];
    operator_colors.append(f"\033[38;2;{r};{g};{b}m");

RESET_ESCAPE = "\033[0m";

def pretty(exp, in_color = False, depth = 0):

    start_color, end_color = "", "";

    if in_color:
        start_color = operator_colors[depth % len(operator_colors)];

        end_color = RESET_ESCAPE;

    retval = "";

    match exp:
        case ("literal", x):
            retval = str(x);

            if in_color:
                retval = LITERAL_ESCAPE + retval + RESET_ESCAPE;

        case ("variable", x):
            retval = x

            if in_color:
                retval = VARIABLE_ESCAPE + retval + RESET_ESCAPE;

        case (op, inner):
            retval += start_color + "(" + end_color;

            retval += start_color + op + end_color;

            retval += pretty(inner, in_color, depth + 1);

            retval += start_color + ")" + end_color;

        case (op, left, right):
            retval += start_color + "(" + end_color;

            retval += pretty(left, in_color, depth + 1);

            retval += " " + start_color + op + end_color + " ";

            retval += pretty(right, in_color, depth + 1);

            retval += start_color + ")" + end_color;

        case ("?:", cond, left, right):
            retval += start_color + "(" + end_color;

            retval += pretty(cond, in_color, depth + 1);

            retval += " " + start_color + "?" + end_color + " ";

            retval += pretty(left, in_color, depth + 1);

            retval += " " + start_color + ":" + end_color + " ";

            retval += pretty(right, in_color, depth + 1);

            retval += start_color + ")" + end_color;

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

    return retval;


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 (_, cond, true, false):
                    return extract_variables(cond) + extract_variables(true) + extract_variables(false);
                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, # or
        '&&': lambda x, y: x & y, # and

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

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

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

    ternary_operators = {
        '?:': lambda x, y, z: (x & y) | (~x & z), # ternary
    };

    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);

        print([len(x) for x in todo]);

        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_cost = my_cost + 1;

                unary_truthtable = function(my_truthtable) & M;
                unary_expression = (name, my_expression);

                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()):
                    binary_cost = my_cost + 1 + costs[other_truthtable];

                    # x + y
                    binary_truthtable = function(my_truthtable, other_truthtable);
                    binary_truthtable = binary_truthtable & M
                    binary_expression = (name, my_expression, other_expression);

                    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);

                    consider(binary_truthtable, binary_expression, binary_cost);

        # consider ternary operators:
        for name, function in sorted(ternary_operators.items()):
            if name in available_operators:
                s = sorted(lookup.items());
                for a_truthtable, a_expression in s:
                    for b_truthtable, b_expression in s:
                        ternary_cost = 1 + my_cost + costs[a_truthtable] + costs[b_truthtable];

                        # x ? y : z
                        ternary_truthtable = function(my_truthtable, a_truthtable, b_truthtable) & M;
                        ternary_expression = (name, my_expression, a_expression, b_expression);

                        consider(ternary_truthtable, ternary_expression, ternary_cost);

                        # y ? x : z
                        ternary_truthtable = function(a_truthtable, my_truthtable, b_truthtable) & M;
                        ternary_expression = (name, a_expression, my_expression, b_expression);

                        consider(ternary_truthtable, ternary_expression, ternary_cost);

                        # y ? z : x
                        ternary_truthtable = function(a_truthtable, b_truthtable, my_truthtable) & M;
                        ternary_expression = (name, a_expression, b_expression, my_expression);

                        consider(ternary_truthtable, ternary_expression, ternary_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.");

            if "?:" in available_operators:
                print();

                print("You have selected ternary operators, so this WILL "
                      "take time. This will probably take weeks to run.");

            print();

            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 parse(text):
    class Tokenizer:
        def __init__(self, text):
            self.position = 0;
            self.text = text;
            self.tokenkind = None;
            self.tokendata = None;

        def next(self):
            # skip whitespace:
            while self.position < len(self.text) and self.text[self.position].isspace():
                self.position += 1;

            if self.position == len(self.text):
                self.tokenkind = "EOF";
                return;

            match self.text[self.position]:
                case "0" | "1":
                    self.tokenkind = "literal";
                    self.tokendata = self.text[self.position];
                    self.position += 1;

                case "w" | "x" | "y" | "z":
                    self.tokenkind = "variable";
                    self.tokendata = self.text[self.position];
                    self.position += 1;

                case "(" | ")" | "?" | ":":
                    self.tokenkind = self.text[self.position];
                    self.position += 1;

                case ">":
                    # could be ">" or ">=":
                    self.position += 1;

                    match self.text[self.position]:
                        case "=":
                            self.tokenkind = ">=";
                            self.position += 1;
                        case _:
                            self.tokenkind = ">"

                case "<":
                    # could be "<" or "<=":
                    self.position += 1;

                    match self.text[self.position]:
                        case "=":
                            self.tokenkind = "<=";
                            self.position += 1;
                        case _:
                            self.tokenkind = "<"

                case "=":
                    # could only be "=="
                    self.position += 1;

                    match self.text[self.position]:
                        case "=":
                            self.tokenkind = "==";
                            self.position += 1;

                        case _:
                            raise BaseException("expecting '=' after '='");

                case "!":
                    # could be "!" or "!|" or "&!" or "!="
                    self.position += 1;

                    match self.text[self.position]:
                        case "&":
                            self.tokenkind = "!&";
                            self.position += 1;
                        case "|":
                            self.tokenkind = "!|";
                            self.position += 1;
                        case "=":
                            self.tokenkind = "!=";
                            self.position += 1;
                        case _:
                            self.tokenkind = "!"

                case "&":
                    # could be "&&" or "&!"
                    self.position += 1;

                    match self.text[self.position]:
                        case "&":
                            self.tokenkind = "&&";
                            self.position += 1;
                        case "!":
                            self.tokenkind = "&!";
                            self.position += 1;
                        case _:
                            raise BaseException("expecting '&' or '!' after '&'");

                case "|":
                    # could be "||" or "|!"
                    self.position += 1;

                    match self.text[self.position]:
                        case "|":
                            self.tokenkind = "||";
                            self.position += 1;
                        case "!":
                            self.tokenkind = "|!";
                            self.position += 1;
                        case _:
                            raise BaseException("expecting '|' or '!' after '|'");

                case x:
                    raise BaseException(f'unknown token "{x}"!');

    def parse_primary(tokenizer):
        match tokenizer.tokenkind:
            case "literal":
                ast = ("literal", int(tokenizer.tokendata));
                tokenizer.next();

            case "variable":
                ast = ("variable", tokenizer.tokendata);
                tokenizer.next();

            case "(":
                tokenizer.next();

                ast = parse_root(tokenizer);

                if tokenizer.tokenkind != ")":
                    raise BaseException(f'expected token ")", found: {tokenizer.tokenkind}');

                tokenizer.next();

            case x:
                raise BaseException(f'unexpected token "{x}"!');

        return ast;

    def parse_prefix(tokenizer):
        if tokenizer.tokenkind == "!":
            tokenizer.next();
            return ("!", parse_prefix(tokenizer));
        else:
            return parse_primary(tokenizer);

    def parse_compares(tokenizer):
        left = parse_prefix(tokenizer);

        while tokenizer.tokenkind in ("<=", "<", ">", ">="):
            op = tokenizer.tokenkind;
            tokenizer.next();
            left = (op, left, parse_prefix(tokenizer));

        return left;

    def parse_equals(tokenizer):
        left = parse_compares(tokenizer);

        while tokenizer.tokenkind in ("==", "!="):
            op = tokenizer.tokenkind;
            tokenizer.next();
            left = (op, left, parse_compares(tokenizer));

        return left;

    def parse_ands(tokenizer):
        left = parse_equals(tokenizer);

        while tokenizer.tokenkind in ("&&", "&!", "!&"):
            op = tokenizer.tokenkind;
            tokenizer.next();
            left = (op, left, parse_equals(tokenizer));

        return left;

    def parse_ors(tokenizer):
        left = parse_ands(tokenizer);

        while tokenizer.tokenkind in ("||", "|!", "!|"):
            op = tokenizer.tokenkind;
            tokenizer.next();
            left = (op, left, parse_ands(tokenizer));

        return left;

    def parse_ternary(tokenizer):
        exp = parse_ors(tokenizer);

        if tokenizer.tokenkind in ("?", ):
            cond = exp;

            tokenizer.next();

            true = parse_ors(tokenizer);

            if tokenizer.tokenkind != ":":
                raise BaseException("expecting ':' in ternary expression, "
                    "found: '{tokenizer.tokenkind}'!");

            tokenizer.next();

            false = parse_ternary(tokenizer);

            return ("?:", cond, true, false);
        else:
            return exp;

    def parse_root(tokenizer):
        return parse_ternary(tokenizer);

    tokenizer = Tokenizer(text);

    tokenizer.next();

    root = parse_root(tokenizer);

    if tokenizer.tokenkind != "EOF":
        raise BaseException(
            f"reached end of expression. "
            f"Expecting EOF. found: {tokenizer.tokenkind}!");

    return root;

def evaluate(expr):
    match expr:
        case ("literal", 0): return 0;
        case ("literal", 1): return M;

        case ("variable", "w"): return W;
        case ("variable", "x"): return X;
        case ("variable", "y"): return Y;
        case ("variable", "z"): return Z;

        case ("!", inner):
            return ~evaluate(inner) & 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 ("|!", 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 ("?:", cond, left, right):
            cond = evaluate(cond);
            return (cond & evaluate(left)) | (~cond & evaluate(right));

        case _:
            print(f'expr = {expr}');
            assert(not "TODO");

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

    print();

    print("The py-parser struggles with deeply-nested parentheses.");
    print("I may have to just write my recursive-descent parser.");

    print();

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

        print();

    print(f'I can simplify any tree down to {max(cost.values())} operators or less.')
    print();

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

        line = line.strip();

        if not line: continue;

        try:
            parsed = parse(line);
        except BaseException as e:
            print(f"syntax error: {e}");
            continue;

        truthtable = evaluate(parsed);

        if truthtable in lookup:
            text = "";

            c = cost[truthtable];

            LITERAL_ESCAPE = "\033[38;2;200;200;100m";

            print(f'{cost[truthtable]}: {pretty(lookup[truthtable], args.color)}')

            print(text);
        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, args.color)}');
    elif args.command:
        try:
            parsed = parse(args.command);
        except BaseException as e:
            print(f"syntax error: {e}");
            return 1;

        truthtable = evaluate(parsed);

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

    return 0;

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