import numpy as np
from copy import copy
import matplotlib.pyplot as plt


MAX_SCHEMATIC_SIZE = 32


def unpackbits(x, num_bits):
    xshape = list(x.shape)
    x = x.reshape([-1, 1])
    to_and = 2**np.arange(num_bits).reshape([1, num_bits])
    return (x & to_and).astype(bool).astype(int).reshape(xshape + [num_bits])


def generate_and_count(n, l, masks):
    counts = {}
    count = 0

    for _ in range(n):
        arr = np.random.choice([0, 1], l)
        for i in range(0, 2**l):
            temp = copy(arr)
            mask = masks[i]
            temp[np.where(mask)] = 2
            key = 0

            for j in range(l):
                key += temp[j] * 3 ** j

            if key not in counts:
                counts[key] = 1
                count += 1

    return count


def pop_func(l):
    n_pops, avg_counts = [], []
    masks = [unpackbits(np.array(i), l) for i in range(0, 2 ** l)]
    for n in range(1, 100):
        print('generate and count', n)
        n_pops.append(n)
        counts = []
        for _ in range(10):
            counts.append(generate_and_count(n, l, masks))
        avg_counts.append(np.mean(counts))

    for n in range(100, 251, 10):
        print('generate and count', n)
        n_pops.append(n)
        counts = []
        for _ in range(10):
            counts.append(generate_and_count(n, l, masks))
        avg_counts.append(np.mean(counts))

    for n in range(300, 500, 50):
        print('generate and count', n)
        n_pops.append(n)
        counts = []
        for _ in range(10):
            counts.append(generate_and_count(n, l, masks))
        avg_counts.append(np.mean(counts))

    for n in range(500, 1001, 100):
        print('generate and count', n)
        n_pops.append(n)
        counts = []
        for _ in range(10):
            counts.append(generate_and_count(n, l, masks))
        avg_counts.append(np.mean(counts))

    return n_pops, avg_counts


def main():
    n_pops, avg_counts = pop_func(9)
    plt.plot(n_pops, avg_counts)
    plt.show()


if __name__ == '__main__':
    main()