from random import randint
import time


## Trie part

class Trie:
    class Node:
        def __init__(self, data):
            self.data = data
            self.next_nodes = dict()
            self.routing_prefixes = set()
            self.isTerminal = False

    def __init__(self):
        self.root_node = Trie.Node(0)

    def add_address(self, address):
        # address = str('132.123.123.123/231')
        if '/' in address:
            address, routing_prefix = address.split('/')
        else:
            address, routing_prefix = address, None

        address = address.split('.')

        current_node = self.root_node
        for value in address:
            if value not in current_node.next_nodes:
                current_node.next_nodes[value] = Trie.Node(value)
            current_node = current_node.next_nodes[value]

        if routing_prefix is None:
            current_node.isTerminal = True
        else:
            current_node.routing_prefixes.add(routing_prefix)

    def __contains__(self, address):
        if '/' in address:
            address, routing_prefix = address.split('/')
        else:
            address, routing_prefix = address, None

        address = address.split('.')
        if len(address) != 4:
            return False

        try:
            address = list(map(int, address))
        except ValueError:
            return False

        current_node = self.root_node
        for value in address:
            if value not in current_node.next_nodes:
                return False
            current_node = current_node.next_nodes[value]

        if routing_prefix is None:
            return current_node.isTerminal
        else:
            return routing_prefix in current_node.routing_prefixes

    def __dfs__(self, node, prefix, result):
        if node.isTerminal:
            result.append(prefix)
        for routing_prefix in node.routing_prefixes:
            result.append(prefix + '/' + routing_prefix)

        if len(prefix):
            prefix += '.'
        for index in node.next_nodes:
            self.__dfs__(node.next_nodes[index], prefix + str(index), result)

    def get_all_addresses(self):
        result = []
        self.__dfs__(self.root_node, '', result)
        return result

    def __dfs_size__(self, node):
        ans = node.next_nodes.__sizeof__()
        ans += node.routing_prefixes.__sizeof__()
        ans += node.data.__sizeof__()

        for index in node.next_nodes:
            ans += self.__dfs_size__(node.next_nodes[index])
        return ans

    def __sizeof__(self):
        return self.__dfs_size__(self.root_node)


## Dict part

def add_address(address):
    if '/' in address:
        address, routing_prefix = address.split('/')
    else:
        routing_prefix = None
    addresses_dict.setdefault(address, set()).add(routing_prefix)


def contains(address):
    if '/' in address:
        address, routing_prefix = address.split('/')
    else:
        routing_prefix = None
    if address not in addresses_dict:
        return False
    return routing_prefix in addresses_dict[address]


## Make tests
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(i) for i in range(N) for j in range(32)]
    return submask


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

    start = time.time()
    for address in addresses:
        trie.add_address(address)
    time_trie = time.time() - start

    start = time.time()
    for address in addresses:
        add_address(address)
    time_dict = time.time() - start

    start = time.time()
    for address in addresses:
        addresses_set.add(address)
    time_set = time.time() - start

    return [time_trie, time_dict, time_set]


def test_contains(N):
    addresses = get_rand_address(N)
    start = time.time()
    for address in addresses:
        check = trie.__contains__(address)
    time_trie = time.time() - start
    start = time.time()
    for address in addresses:
        check = contains(address)
    time_dict = time.time() - start
    start = time.time()
    for address in addresses:
        check = address in addresses_set
    time_set = time.time() - start
    return [time_trie, time_dict, time_set]


###########

trie = Trie()
addresses_dict = dict()
addresses_set = set()

N = 3000
time_add = test_add(N)
print(N)
print('Add to trie: ', 'time = ', time_add[0] * 1000, 'ms')
print('Add to dict: ', 'time = ', time_add[1] * 1000, 'ms')
print('Add to set: ', 'time = ', time_add[2] * 1000, 'ms')

time_contains = test_contains(N)
print('Contains in trie: ', 'time = ', time_contains[0] * 1000, 'ms')
print('Contains in dict: ', 'time = ', time_contains[1] * 1000, 'ms')
print('Contains in set: ', 'time = ', time_contains[2] * 1000, 'ms')

print('Trie size in bytes', trie.__sizeof__())
print('Dict size in bytes', addresses_dict.__sizeof__())
print('Set size in bytes', addresses_set.__sizeof__())