
import numpy as np
import random
from collections import namedtuple


def generate_prob_matrix(n):
    matrix = np.random.rand(n, n)

    for i in range(n):
        matrix[i][i] = 0

    for i in range(n):
        matrix[i] = (1/np.sum(matrix[i]))*matrix[i]

    return matrix


def categorical(p):
    return np.random.choice(len(p), 1, p=p)[0]


Drone = namedtuple('Drone', 'speed probability')
Site = namedtuple('Site', 'location')


class System:

    def __init__(self, sites, drones):
        self.sites = {}
        self.drones = {}
        n = len(sites)

        for i, drone in enumerate(drones):
            self.drones[i] = drone

        for i, site in enumerate(sites):
            self.sites[i] = site

        distance = np.zeros([n, n])
        for i in range(n):
            for j in range(n):
                if i < j:
                    x = np.subtract(sites[i], sites[j])
                    d = np.linalg.norm(x)
                    distance[i][j] = d
                    distance[j][i] = d
        self.distance = distance

    def get_site(self, site_id):
        return self.sites[site_id]

    def get_drone(self, drone_id):
        return self.drones[drone_id]

    def compute_path_distance(self, path):
        n = len(path)
        d = 0
        for i in range(n - 1):
            d += self.distance[path[i]][path[i + 1]]
        return d

    def compute_path_time(self, path, drone_id):
        d = self.compute_path_distance(path)
        return d/self.get_drone(drone_id).speed

    def generate_path_of_length(self, length, drone_id):
        path = []
        P = self.get_drone(drone_id).probability
        num_sites = len(self.sites)
        s = categorical([1/num_sites]*num_sites)
        path.append(s)
        site = s
        for i in range(length):
            site = categorical(P[site])
            path.append(site)
        return path

    def generate_path(self, s, t, drone_id):
        path = [s]
        P = self.get_drone(drone_id).probability
        site = categorical(P[s])
        path.append(site)
        while site != t:
            site = categorical(P[site])
            path.append(site)
        return path

    @staticmethod
    def generate_random_system(n, k):

        locations = np.random.rand(n, 2)
        sites = []
        for i in locations:
            sites.append(Site(i))

        drones = []
        for i in range(k):
            speed = abs(random.random())
            probability = generate_prob_matrix(n)
            drones.append(Drone(speed, probability))

        return System(sites, drones)


def _compute_arrival_times(path, drone_id, sites, speed):
    arrival_times = []
    t = 0
    for i in range(len(path) - 1):
        t += system.compute_path_time(path[i:i+2], drone_id=drone_id)
        arrival_times.append((drone_id, path[i], path[i+1], t))
    return arrival_times


def _generate_arrival_times(system, num_drones, length):
    arrival_times = [[] for _ in range(len(system.sites))]

    events = []
    for i in range(system):
        pass
        events.extend(compute_arrival_times(path, i))

    def get_key(item):
        return item[3]

    events = sorted(events, key=get_key)

    for event in events:
        drone_id = event[0]
        site_id = event[2]
        time = event[3]
        arrival_times[site_id].append((drone_id, time))

    return arrival_times


def compute_cost(system, n):
    arrival_times = generate_arrival_times(system, n)
    interarrival_times = [[] for _ in range(len(system.sites))]
    for i in range(len(arrival_times)):
        arrivals = arrival_times[i]
        for j in range(len(arrivals) - 1):
            interarrival_times[i].append(arrivals[j+1][1] - arrivals[j][1])

    interarrival_avgs = [compute_average(i) for i in interarrival_times]
    return max(interarrival_avgs)


def compute_average(data):
    return (1/len(data))*sum(data)

[((23, 8, 23, 48), 'collections.namedtuple', 'namedtuple', ({(23, 19, 23, 26): '"""Drone"""', (23, 28, 23, 47): '"""speed probability"""'}, {}), "('Drone', 'speed probability')", False, 'from collections import namedtuple\n'), ((24, 7, 24, 37), 'collections.namedtuple', 'namedtuple', ({(24, 18, 24, 24): '"""Site"""', (24, 26, 24, 36): '"""location"""'}, {}), "('Site', 'location')", False, 'from collections import namedtuple\n'), ((8, 13, 8, 33), 'numpy.random.rand', 'np.random.rand', ({(8, 28, 8, 29): 'n', (8, 31, 8, 32): 'n'}, {}), '(n, n)', True, 'import numpy as np\n'), ((40, 19, 40, 35), 'numpy.zeros', 'np.zeros', ({(40, 28, 40, 34): '[n, n]'}, {}), '([n, n])', True, 'import numpy as np\n'), ((92, 20, 92, 40), 'numpy.random.rand', 'np.random.rand', ({(92, 35, 92, 36): 'n', (92, 38, 92, 39): '2'}, {}), '(n, 2)', True, 'import numpy as np\n'), ((14, 23, 14, 40), 'numpy.sum', 'np.sum', ({(14, 30, 14, 39): 'matrix[i]'}, {}), '(matrix[i])', True, 'import numpy as np\n'), ((99, 24, 99, 39), 'random.random', 'random.random', ({}, {}), '()', False, 'import random\n'), ((44, 24, 44, 55), 'numpy.subtract', 'np.subtract', ({(44, 36, 44, 44): 'sites[i]', (44, 46, 44, 54): 'sites[j]'}, {}), '(sites[i], sites[j])', True, 'import numpy as np\n'), ((45, 24, 45, 41), 'numpy.linalg.norm', 'np.linalg.norm', ({(45, 39, 45, 40): 'x'}, {}), '(x)', True, 'import numpy as np\n')]