import ipaddress
from random import randint
from random import shuffle
import time


#############################################################
class BinaryTrie():
    def __init__(self):
        self.MASK = (1 << 31) - 1
        self.nodes = [0 for _ in range(64)]
        self.empty_nodes = [63 - i for i in range(63)]

    def __sizeof__(self):
        return self.nodes.__sizeof__() + self.empty_nodes.__sizeof__()

    def _get_left(self, value):
        return (value >> 1) & self.MASK

    def _get_right(self, value):
        return value >> (31 + 1)

    def _get_type(self, value):
        return value & 1

    def _create_node(self):
        if len(self.empty_nodes) == 0:
            self.empty_nodes = [31 + len(self.nodes) - i for i in range(32)]
            self.nodes += [0 for _ in range(32)]
        index = self.empty_nodes.pop()
        return index

    def _remove_nodes(self, node):
        if node == 0:
            return
        self.empty_nodes.append(node)
        left_node = self._get_left(self.nodes[node])
        self._remove_nodes(left_node)
        right_node = self._get_right(self.nodes[node])
        self._remove_nodes(right_node)
        self.nodes[node] = 0

    def add(self, address):
        value, size = self.change_format(address)
        current_node = 0
        for i in range(size):
            bit = 31 - i
            if self._get_type(self.nodes[current_node]) == 1:  # leaf
                return
            if (value >> bit) & 1 == 0:
                current_left = self._get_left(self.nodes[current_node])
                if current_left == 0:
                    current_left = self._create_node()
                    self.nodes[current_node] |= current_left << 1
                current_node = current_left
            else:
                current_right = self._get_right(self.nodes[current_node])
                if current_right == 0:
                    current_right = self._create_node()
                    self.nodes[current_node] |= current_right << (31 + 1)
                current_node = current_right

        current_left = self._get_left(self.nodes[current_node])  # delete children if a vertex added
        current_right = self._get_right(self.nodes[current_node])

        self._remove_nodes(current_left)
        self._remove_nodes(current_right)

        self.nodes[current_node] = 1  # leaf

    def __contains__(self, address):
        value, size = self.change_format(address)
        current_node = 0
        for i in range(size):
            bit = 31 - i
            if self._get_type(self.nodes[current_node]) == 1:
                return True
            if (value >> bit) & 1 == 0:
                current_left = self._get_left(self.nodes[current_node])
                if current_left == 0:
                    return False
                current_node = current_left
            else:
                current_right = self._get_right(self.nodes[current_node])
                if current_right == 0:
                    return False
                current_node = current_right
        return self._get_type(self.nodes[current_node]) == 1

    def change_format(self, address):
        if '/' in address:
            address, routing_prefix = address.split('/')
            size = 32 - int(routing_prefix)
        else:
            size = 0

        address = list(map(int, address.split('.')))

        value = 0
        for item in address:
            value <<= 8
            value ^= item

        value >>= size
        value <<= size
        return value, 32 - size


##############################################################
def add_set(subnet):
    try:
        addresses_set.add(ipaddress.ip_network(subnet))
    except Exception:
        subnet = ipaddress.ip_interface(subnet)
        addresses_set.add(subnet.network)


def contains_set(address):
    address = ipaddress.ip_address(address)
    for subnet in addresses_set:
        if address in subnet:
            return True
    return False


def add_list(subnet):
    try:
        addresses_list.append(ipaddress.ip_network(subnet))
    except Exception:
        subnet = ipaddress.ip_interface(subnet)
        addresses_list.append(subnet.network)


def contains_list(address):
    address = ipaddress.ip_address(address)
    for subnet in addresses_list:
        if address in subnet:
            return True
    return False


#### TEST PART ####

def get_rand_address(N):
    mask = [str(randint(1, 255)) + '.' + str(randint(1, 255)) + '.' + str(randint(1, 255)) + '.' + str(randint(1, 255))
            for i in range(N)]
    submask = [mask[i] + '/' + str(j) for i in range(N) for j in range(8, 32)]
    shuffle(submask)
    return submask


def test_add(N):
    addresses = get_rand_address(N)
    times = [[] for _ in range(3)]
    memories = [[] for _ in range(3)]

    for i in range(0, N, 100):
        current_addresses = [addresses[i + j] for j in range(100)]
        start = time.time()
        for address in current_addresses:
            add_set(address)
        time_set = time.time() - start
        times[0].append(time_set + (times[0][-1] if len(times[0]) else 0))
        memories[0].append(addresses_set.__sizeof__())

        start = time.time()
        for address in current_addresses:
            add_list(address)
        time_list = time.time() - start
        times[1].append(time_list + (times[1][-1] if len(times[1]) else 0))
        memories[1].append(addresses_list.__sizeof__())

        start = time.time()
        for address in current_addresses:
            binaryTrie.add(address)
        time_SUPERTREE = time.time() - start
        times[2].append(time_SUPERTREE + (times[2][-1] if len(times[2]) else 0))
        memories[2].append(binaryTrie.__sizeof__())

    return [times, memories]


check1 = []
check2 = []
check3 = []


def test_contains(N):
    addresses = get_rand_address(N)

    for i in range(len(addresses)):
        if '/' in addresses[i]:
            pref, suf = addresses[i].split('/')
            addresses[i] = pref

    times = [[] for _ in range(3)]
    for i in range(0, N, 100):
        current_addresses = [addresses[i + j] for j in range(100)]

        start = time.time()
        for address in current_addresses:
            check1.append(contains_set(address))
        time_set = time.time() - start
        times[0].append(time_set + (times[0][-1] if len(times[0]) else 0))

        start = time.time()
        for address in current_addresses:
            check2.append(contains_list(address))
        time_list = time.time() - start
        start = time.time()
        times[1].append(time_list + (times[1][-1] if len(times[1]) else 0))

        for address in current_addresses:
            check3.append(binaryTrie.__contains__(address))
        time_SUPERTREE = time.time() - start
        times[2].append(time_SUPERTREE + (times[2][-1] if len(times[2]) else 0))

    return times


addresses_set = set()
addresses_list = []
binaryTrie = BinaryTrie()
import matplotlib.pyplot as plt

N = 100000
print(N)
times_add, memories = test_add(N)
for i in range(len(times_add)):
    for j in range(len(times_add[0])):
        times_add[i][j] *= 1000
print('test_add done')
ox = [i for i in range(0, N, 100)]

fig, axes = plt.subplots(nrows=1, ncols=1)
axes.plot(ox, times_add[0], color='#00dd00', label="Set")
axes.plot(ox, times_add[1], color='#dd0000', label="List")
axes.plot(ox, times_add[2], color='#0000dd', label="Trie")
axes.legend(loc="upper left")
axes.set_xlabel("Количество подсетей")
axes.set_ylabel("Время работы add")
fig.set_size_inches(10, 7)
plt.show()

fig, axes = plt.subplots(nrows=1, ncols=1)
axes.plot(ox, memories[0], color='#00dd00', label="Set")
axes.plot(ox, memories[1], color='#dd0000', label="List")
axes.plot(ox, memories[2], color='#0000dd', label="Trie")
axes.legend(loc="upper left")
axes.set_xlabel("Количество подсетей")
axes.set_ylabel('Memory')
fig.set_size_inches(10, 7)
plt.show()

N = 1000
ox = [i for i in range(0, N, 100)]
times_contains = test_contains(N)
for i in range(len(times_contains)):
    for j in range(len(times_contains[0])):
        times_contains[i][j] *= 1000
print('test contains done')
fig, axes = plt.subplots(nrows=1, ncols=1)
axes.plot(ox, times_contains[0], color='#00dd00', label="Set")
axes.plot(ox, times_contains[1], color='#dd0000', label="List")
axes.plot(ox, times_contains[2], color='#0000dd', label="Trie")
axes.legend(loc="upper left")
axes.set_xlabel("Количество подсетей")
axes.set_ylabel("Время работы contains")
fig.set_size_inches(10, 7)
plt.show()