import lzma
import sys
import os
import itertools

from booleantools import generate_function


def is_good_zhegalkin(function_number, args_count):
    return all(map(lambda x: len(x) <= 3, generate_function(function_number, args_count).listform))


def get_rank_file_object(args_count, rank_ind):
    base_name = "base_{}_rank_{}.txt".format(args_count, rank_ind)
    if os.path.exists(base_name):
        return open(base_name, "rb")
    elif os.path.exists(base_name + ".xz"):
        return lzma.LZMAFile(base_name + ".xz", "rb")
    else:
        return None


def is_function_exists(function_number, file_object):
    file_object.seek(function_number // 8)
    result_byte = ord(file_object.read(1))
    return (result_byte >> (7 - function_number % 8)) % 2 == 1


def function_values_to_number(values):
    assert set(values).issubset({"0", "1"})
    res = 0
    for ch in values:
        res *= 2
        if ch == "1":
            res += 1
    return res


def count_lens(args_count, white_map_path):
    bit_counts = bytes(bin(x).count("1") for x in range(256))
    len_to_count = {}
    white_map = None
    if white_map_path:
        white_map = open(white_map_path, "rb").read()
        print(f"Total whitelisted functions: {sum(map(lambda x: bit_counts[x], white_map))}")
    for rank_ind in range(1, 10):
        file_object = get_rank_file_object(args_count, rank_ind)
        if not file_object:
            break
        src_data = file_object.read()
        if white_map:
            assert len(src_data) == len(white_map)
            src_data = bytes(x & y for x,y in zip(src_data, white_map))
        len_to_count[rank_ind] = sum(map(lambda x: bit_counts[x], src_data))
    total = sum(len_to_count.values())
    for rank_ind, count in len_to_count.items():
        print("len {} count {} share {}".format(rank_ind, count, 1.0 * count / total))
    print("average: {}".format(sum([key * value for key, value in len_to_count.items()]) / total))



def get_function_len(args_count, function_number):
    for rank_ind in range(1, 10):
        file_object = get_rank_file_object(args_count, rank_ind)
        if not file_object:
            assert False
        if is_function_exists(function_number, file_object):
            return rank_ind
    assert False


def count_symmetricals(args_count):
    if args_count <= 4:
        func_strings = list(open(f"./build/functions_strings_{args_count}.txt").read().split("\n")[:-1])
    else:
        func_strings = []
    print(f"Total {len(func_strings)} representations")
    functions_to_check = set()
    for layer_values in itertools.product([0, 1], repeat=args_count + 1):
        func_values = []
        for argument_number in range(2 ** args_count):
            if layer_values[bin(argument_number).count("1")]:
                func_values.append("1")
            else:
                func_values.append("0")
        function_number = function_values_to_number("".join(func_values))
        functions_to_check.add(function_number)

    function_ranks = {}
    for fn in functions_to_check:
        rank = get_function_len(args_count, fn)
        if rank not in function_ranks:
            function_ranks[rank] = []
        function_ranks[rank].append(fn)
    zhegalkin_functions = {key: list(filter(lambda x: is_good_zhegalkin(x, args_count), value)) for key,value in function_ranks.items()}
    for processed_function_ranks in [function_ranks, zhegalkin_functions]:
        if processed_function_ranks is zhegalkin_functions:
            print("Со степенью <= 3:")

        len_to_count = {key: len(value) for key, value in processed_function_ranks.items()}
        total = sum(len_to_count.values())
        for rank_ind, count in sorted(len_to_count.items(), key=lambda x: x[0]):
            print("len {} count {} share {}".format(rank_ind, count, 1.0 * count / total))
        print(
            "average: {}".format(
                sum([1.0 * key * value for key, value in len_to_count.items()]) / total
            )
        )


    print(r"\begin{center}")
    print(r"\begin{longtable}{| l| r | p{11cm}|}")
    print(r"\hline")
    if args_count <= 4:
        print(r"Длина функции & Номер & ПСПФ\\")
    else:
        print(r"Длина функции & Номер & Полином Жегалкина\\")

    print(r"\hline")
    print(r"\endhead")
    print(r"\hline \multicolumn{3}{r}{\textit{Продолжение на следующей странице}} \\")
    print(r"\endfoot")
    print(r"\hline")
    print(r"\endlastfoot")
    for rank in sorted(function_ranks.keys()):
        for function_number in sorted(function_ranks[rank]):

            if args_count <= 4:
                function_representation = func_strings[function_number]
            else:
                function_representation = generate_function(function_number, args_count).tex_str().strip()

            print(f"{rank} & {function_number} & ${function_representation}$\\\\")
        if rank != max(function_ranks.keys()):
            print(r"\hline")
    print(r"\hline")
    print(r"\end{longtable}")
    print(r"\addtocounter{table}{-1}")
    print(r"\end{center}")


def main():
    if sys.argv[1] == "count_lens":
        white_map_path = None
        if len(sys.argv) > 3:
            white_map_path = sys.argv[3]
        count_lens(sys.argv[2], white_map_path)
        return
    elif sys.argv[1] == "count_symmetricals":
        count_symmetricals(int(sys.argv[2]))
        return
    args_count = {
        8: 3,
        16: 4,
        32: 5,
    }[len(sys.argv[1])]
    function_number = function_values_to_number(sys.argv[1])
    print("Function number is {}".format(function_number))
    for rank_ind in range(1, 10):
        file_object = get_rank_file_object(args_count, rank_ind)
        if not file_object:
            print("No results file for {} arguments rank {}".format(args_count, rank_ind))
            sys.exit(1)
        if is_function_exists(function_number, file_object):
            print("Result len is {}".format(rank_ind))
            return


if __name__ == "__main__":
    main()