from __future__ import print_function, unicode_literals

import lzma
import sys
import os

import itertools


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):
    bit_counts = bytes(bin(x).count("1") for x in range(256))
    len_to_count = {}
    for rank_ind in range(1, 10):
        file_object = get_rank_file_object(args_count, rank_ind)
        if not file_object:
            break
        len_to_count[rank_ind] = sum(map(lambda x: bit_counts[x], file_object.read()))
    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 count_symmetricals(args_count):
    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)
        print("".join(func_values), function_number)

    function_ranks = {}
    for rank_ind in range(1, 10):
        if not functions_to_check:
            break
        file_object = get_rank_file_object(args_count, rank_ind)
        if not file_object:
            break
        function_ranks[rank_ind] = []
        for fn in list(functions_to_check):
            if is_function_exists(fn, file_object):
                function_ranks[rank_ind].append(fn)
                functions_to_check.remove(fn)
    print(function_ranks)
    len_to_count = {key: len(value) for key, value in function_ranks.items()}
    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([1.0 * key * value for key, value in len_to_count.items()]) / total
        )
    )


def main():
    if sys.argv[1] == "count_lens":
        count_lens(sys.argv[2])
        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()