import math
from statistics import median

def count_one_bits(n):
    return bin(n).count("1")

def compute_distance(a, b):
    distance = count_one_bits(a ^ b)
    # return 1 / (8 ** distance)
    return 1 / (2 ** distance)

def xor_n(n):
    return count_one_bits(n) % 2

def compute_distances(N):
    return [[compute_distance(i, j) for j in range(N)] for i in range(N)]

def compute_nn_probabilities(i, knowns, distances):
    total = 0.0
    total_zero = 0.0
    total_one = 0.0
    for known in knowns:
        j = known[0]
        distance = distances[i][j]
        total += distance
        if known[1] == 0:
            total_zero += distance
        else:
            total_one += distance
    p_zero = total_zero / total
    p_one = total_one / total
    return (p_zero, p_one)

def compute_est_coherence(i, knowns, coherences, distances):
    total = 0.0
    coherence = 0.0
    for known in knowns:
        j = known[0]
        distance = distances[i][j]
        total += distance
        coherence += distance * coherences[j]
    return coherence / total

def compute_est_coherences(N, knowns, distances):
    nn_probabilities = [None for i in range(N)]
    est_coherences = [None for i in range(N)]

    # for known in knowns:
    #     i = known[0]
    #     nn_probabilities[i] = compute_nn_probabilities(i, knowns, distances)
    for known in knowns:
        i = known[0]
        nn_probabilities[i] = (1.0 - known[1], 1.0 * known[1])
    
    for i in range(len(nn_probabilities)):
        if not nn_probabilities[i] is None:
            continue
        nn_probabilities[i] = compute_nn_probabilities(i, knowns, distances)

    print(nn_probabilities)
        
    for i in range(len(nn_probabilities)):
        total = 0.0
        coherence = 0.0
        p_i = nn_probabilities[i]
        for j in range(len(nn_probabilities)):
            if i == j:
                continue
            p_j = nn_probabilities[j]
            distance = distances[i][j]
            total += distance
            coherence += (p_i[0] * p_j[0] + p_i[1] * p_j[1]) * distance
        # print(coherence, total)
        est_coherences[i] = coherence / total

    # for known in knowns:
    #     i = known[0]
    #     est_coherences[i] = nn_probabilities[i][known[1]]

    # for i in range(len(est_coherences)):
    #     if not est_coherences[i] is None:
    #         continue
    #     est_coherences[i] = compute_est_coherence(i, knowns, est_coherences, distances)

    # print(est_coherences)

    return est_coherences

def score(coherences):
    # while len(coherences) > 1:
    #     coherences = [(coherences[i] + coherences[i + 1]) / 2 for i in range(0, len(coherences), 2)]
    # return coherences[0]

    # return median(coherences)
    return sum(coherences) / len(coherences)

def xor_by_index(knowns, index):
    mask = 1 << index
    knowns = knowns[:]
    for i in range(len(knowns)):
        known = knowns[i]
        if known[0] & mask:
            knowns[i] = (known[0], known[1] ^ 1)
    return knowns

def main():
    n = 3
    N = 2 ** n
    distances = compute_distances(N)

    knowns = [(i, xor_n(i)) for i in [
        0, 3, 4, 5, 7
        # 3, 5, 6, 10, 12, 14
        # 1, 3, 7, 10, 14, 15
        # 0, 3, 5, 6, 10, 11, 12
        # 0, 3, 5, 6, 10, 11, 12, 24, 30
    ]]
    print(knowns)
    print()

    # knowns = [
    #     (1, 1),
    #     (3, 0),
    #     (7, 1),
    #     (10, 0),
    #     (14, 1),
    #     (15, 0)
    # ]

    # knowns = [
    #     (0, 0),
    #     (3, 0),
    #     (4, 1),
    #     (5, 0),
    #     (7, 1)
    # ]

    # knowns = [
    #     (0, 0),
    #     (1, 1),
    #     (2, 1),
    #     (3, 0),
    #     (4, 1),
    #     (5, 0),
    #     (6, 0),
    #     (7, 1)
    # ]

    coherences = compute_est_coherences(N, knowns, distances)
    best_coherence = score(coherences)
    print(best_coherence)

    while best_coherence < 1.0:
        print()
        # print(knowns)
        # print()
        best_index = -1
        for i in range(0, n):
            coherences = compute_est_coherences(N, xor_by_index(knowns, i), distances)
            coherence = score(coherences)
            print(coherence)
            if coherence > best_coherence:
                best_coherence = coherence
                best_index = i
        if best_index < 0:
            break
        knowns = xor_by_index(knowns, best_index)


if __name__ == "__main__":
    main()