import math
import numpy as np

def decode(x, N):
    index = 0
    output = np.zeros((N))
    while x > 0 and index < N:
        output[index] = x & 0b1
        x >>= 1
        index += 1
    return output

def hamming_distance(a, b):
    return np.sum(np.logical_xor(a, b))

def xor(x, bits):
    return np.sum(x[:bits]) % 2

def compute_pyramids(N):
    num_orders = max(int(N / 2), 1)
    pyramids = np.zeros((num_orders, N, N)).astype(np.int32)
    for i in range(2, N):
        for j in range(1, i):
            pyramids[0][i][j] = j
    for order in range(1, num_orders):
        # build out the first column
        acc = 0
        for i in range(order * 2 + 2, N):
            acc += pyramids[order - 1][i - 2][1]
            pyramids[order][i][1] = acc
        # accumulate the first column and place it on the diagonal(s)
        for k in range(0, int(order / 2) + 1):
            acc = 0
            for i in range(order * 2 + 2, N):
                acc += pyramids[order][i][1]
                pyramids[order][i][i - 1 - 2 * k] = acc
        # for odd, copy the first column to the first diagonal
        if order % 2 == 1:
            k += 1
            for i in range(order * 2 + 2, N):
                pyramids[order][i][i - 1 - 2 * k] = pyramids[order][i][1]
        # integrate under the diagonal
        inset = 1
        for j in reversed(range(2, N - 2 * k - 2)):
            acc = pyramids[order][N - inset - 1][j]
            for i in range(N - inset, N):
                acc += pyramids[order - 1][i - 2][j]
                pyramids[order][i][j] = acc
            if order * 2 + 2 < N - inset:
                inset += 1
    return pyramids

def compute_pyramids_full(N):
    num_orders = max(int(N / 2), 1)
    pyramids = np.zeros((num_orders, N, N)).astype(np.int32)
    # 1st order can be filled in as multiplication and forms the base case
    for i in range(0, N):
        for j in range(0, i + 1):
            pyramids[0][i][j] = (i - j + 1) * (j + 1)
    for order in range(1, num_orders):
        offset = order * 2

        # fill in the LHS and diagonal
        for i in range(0, N - offset):
            value = math.comb(2 * (order + 1) + i - 1, i)
            pyramids[order][i + offset][0] = value
            # mirror
            pyramids[order][i + offset][i + offset] = value

        # accumulate along the diagonals
        for i in range(1, N):
            value = pyramids[order][i][0]
            acc = value
            for j in range(1, N - i):
                value += acc
                pyramids[order][i + j][j] = value
                acc += pyramids[order - 1][i + j - 1][j - 1]

    return pyramids

def get_total_band_count_2(distance, band_distance, N):
    if band_distance % 2 == 1:
        return 0
    order = int(band_distance / 2) - 1
    if order < 0:
        return 0
    if distance < order + 1:
        return 0
    if distance > N - order - 1:
        return 0
    order_root = math.factorial(2 * (order + 1)) / math.factorial(order + 1) ** 2
    scale = math.comb(N - (order + 1) * 2, distance - order - 1)
    value = math.comb(2 * (order + 1) + N - 2 * (order + 1), N - 2 * (order + 1))
    return order_root * scale * value

def get_incoherent_band_count_2(pyramids, distance, band_distance, k, N):
    if k == 0 or k == N or band_distance % 2 == 1:
        return 0
    order = int(band_distance / 2) - 1
    if order < 0:
        return 0
    if distance < order + 1:
        return 0
    if distance > N - order - 1:
        return 0
    order_root = math.factorial(2 * (order + 1)) / math.factorial(order + 1) ** 2
    scale = math.comb(N - (order + 1) * 2, distance - order - 1)
    value = pyramids[order][N - 2][k - 1]
    return order_root * scale * value

    # pyramid = pyramids[order]
    # offset = (N - 1 - order) - distance
    # multiplier = pyramid[2 * order + offset][2 * order + 1 + offset]
    # row = N - offset
    # column = k
    # value = pyramid[row][column]
    # return multiplier * value

def get_incoherent_band_count(pyramids, distance, band_distance, k, N):
    if k == 0 or k == N or band_distance % 2 == 1:
        return 0
    order = int(band_distance / 2) - 1
    if order < 0:
        return 0
    if distance < order + 1:
        return 0
    if distance > N - order - 1:
        return 0
    if distance < k:
        distance = N - distance
        k = N - k
    pyramid = pyramids[order]
    offset = (N - 1 - order) - distance
    multiplier = pyramid[2 * order + 2 + offset][2 * order + 1 + offset]
    row = N - offset
    column = k
    value = pyramid[row][column]
    return multiplier * value

def get_total_band_count(pyramids, distance, band_distance, N):
    if band_distance % 2 == 1:
        return 0
    order = int(band_distance / 2) - 1
    if order < 0:
        return 0
    if distance < order + 1:
        return 0
    if distance > N - order - 1:
        return 0
    pyramid = pyramids[order]
    offset = (N - 1 - order) - distance
    length = N + 1 - 2 * (order + 1)
    a = pyramid[2 * order + 2 + offset][2 * order + 1 + offset]
    b = pyramid[2 * order + 2 + (length - offset - 1)][2 * order + 1 + (length - offset - 1)]
    return a * b

# def compute_band_distances(pyramids, N):
#     num_orders = max(int(N / 2), 1)
#     incoherent_bands = np.zeros((N + 1, N + 1, N + 1)).astype(np.int32)
#     for order in range(0, num_orders):
#         band_distance = (order + 1) * 2
#         for k in range(1, N):
            

#     for k in range(0, N + 1):
#         for distance in range()

def main():
    # N = 8
    # print(compute_pyramids_full(N))
    # total_distances = np.zeros((N + 1, N + 1)).astype(np.int32)
    # for i in range(0, N + 1):
    #     for j in range(0, N + 1):
    #         total_distances[i][j] = get_total_band_count_2(i, j, N)
    # print(total_distances)
    # return

    max_N = 8
    orders = [np.zeros((max_N + 1, max_N + 1)).astype(np.int32) for _ in range(0, max_N)]

    print('Attempting discrete solution...')
    pyramids = compute_pyramids_full(max_N + 1)

    for N in range(max_N, max_N + 1):
    # for N in range(2, max_N + 1):
        print('=============================')
        print('N@', N)
        print('Generating points...')
        points = []
        for i in range(0, 2 ** N):
            points.append(decode(i, N))
        
        print('Computing bands...')
        bands = [[] for _ in range(0, N + 1)]
        for i in range(1, len(points)):
            distance = hamming_distance(points[0], points[i])
            bands[distance].append(points[i])

        print('Computing band distances...')
        incoherent_distances = np.zeros((N + 1, N + 1))
        for k in range(0, N + 1):
            print('k@', k)
            # print(k, '================================')
            x_a = points[0]
            y_a = xor(x_a, k)
            incoherent_bands = np.zeros((N + 1, N + 1)).astype(np.int32)
            precomputed_incoherent_bands = np.zeros((N + 1, N + 1)).astype(np.int32)
            total_bands = np.zeros((N + 1, N + 1)).astype(np.int32)
            precomputed_total_bands = np.zeros((N + 1, N + 1)).astype(np.int32)
            for distance in range(0, N + 1):
                band = bands[distance]
                for x_b in band:
                    y_b = xor(x_b, k)
                    if y_a != y_b:
                        incoherent_distances[k][distance] += 1
                
                if len(band) < 2:
                    continue
                for band_origin in range(0, len(band)):
                    x_p = band[band_origin]
                    # print(x_p)
                    y_p = xor(x_p, k)
                    for i in range(0, len(band)):
                        if i == band_origin:
                            continue
                        x_q = band[i]
                        y_q = xor(x_q, k)
                        band_distance = hamming_distance(x_p, x_q)
                        total_bands[distance][band_distance] += 1
                        if y_p != y_q:
                            incoherent_bands[distance][band_distance] += 1
                for band_distance in range(0, N + 1):
                    precomputed_incoherent_bands[distance][band_distance] = get_incoherent_band_count_2(pyramids, distance, band_distance, k, N)
                    precomputed_total_bands[distance][band_distance] = get_total_band_count_2(distance, band_distance, N)
            # print(incoherent_bands)
            for order in range(0, int(N / 2)):
                root = math.factorial(2 * (order + 1)) / math.factorial(order + 1) ** 2
                index = order * 2 + 2
                orders[order][N][k] = incoherent_bands[-2 - order][index] / root

            print(incoherent_bands)
            print(precomputed_incoherent_bands)
            print(total_bands)
            print(precomputed_total_bands)
            # print(total_bands)
                    # print(distance, hamming_distance(x_p, x_q), y_p, y_q)
    # for i in range(0, len(orders)):
    #     print(orders[i])
    #     # print(pyramids[i])
    #     print('========================================')
        # print(incoherent_distances)
    # print(bands)
    
if __name__ == "__main__":
    main()