import numpy as np
import matplotlib.pyplot as plt

def flip(n, index):
    return n ^ (1 << index)

def distance(i, j):
    return bin(i ^ j).count('1')

def matrix_system_with_two_knowns(p, q, N):
    S = 2 ** N
    mat = np.zeros((S, S))
    val = np.zeros(S)
    for i in range(0, S):
        if i == p:
            mat[i][i] = 1.0
            val[i] = 1.0
        elif i == q:
            mat[i][i] = 1.0
        else:
            mat[i][i] = -1.0
            for j in range(0, N):
                mat[i][flip(i,j)] = 1.0 / N
    return (mat, val)

def main():
    final_values = []
    final_x = []
    final_y = []

    for N in range(11, 12):
        print(N)
        S = 2 ** N
        distances = np.zeros((S, S))
        for i in range(0, S):
            for j in range(0, S):
                distances[i][j] = distance(i,j)

        # final_values = []
        # final_basis = []
        visited_distances = set()
        for p in range(0, S):
            for q in range(p + 1, S):
                pq_distance = distances[p, q]
                if pq_distance in visited_distances:
                    continue
                visited_distances.add(pq_distance)
                (mat, val) = matrix_system_with_two_knowns(p, q, N)
                solution = np.linalg.inv(mat).dot(val)
                for i in range(0, len(solution)):
                    final_x.append(distances[i, p] / N)
                    final_y.append(distances[i, q] / N)
                    final_values.append(solution[i])

            # values = list(set(solution))
            # values.sort()
            # if len(values) <= 1:
            #     continue
            # basis = [1.0 * i / (len(values) - 1) for i in range(len(values))]

            # final_values.extend(values)
            # final_basis.extend(basis)

    # fig, ax = plt.subplots()
    # ax.scatter(final_values, final_basis)

    # print(np.linalg.lstsq((final_x, final_y), final_values))

    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    ax.scatter(final_x, final_y, final_values)

    ax.grid(True)
    plt.show()

if __name__ == "__main__":
    main()