1import math
  2
  3class Candy:
  4    def __init__(self, mass, uranium):
  5        self.mass = mass
  6        self.uranium = uranium
  7
  8    def get_uranium_quantity(self):
  9        return self.mass * self.uranium
 10
 11    def get_mass(self):
 12        return self.mass
 13
 14
 15class Person:
 16    def __init__(self, position):
 17        self.position = tuple(position)
 18
 19    def get_position(self):
 20        return self.position
 21    
 22    def set_position(self, new_pos):
 23        self.position = new_pos
 24
 25
 26class Kid(Person):
 27
 28    max_uran = 20
 29
 30    def __init__(self, position, initiative):
 31        super().__init__(position)
 32        self.initiative = initiative
 33        self.uranium_quantity = 0
 34        #this is a variable for later, which is useful in class FluxCapacitor
 35        self.unvisited_hosts = []
 36        
 37    def get_initiative(self):
 38        return self.initiative
 39
 40    def add_candy(self, candy):
 41        self.uranium_quantity += candy.get_uranium_quantity()
 42
 43    def is_critical(self):
 44        if self.uranium_quantity > self.max_uran:
 45            return True
 46        return False
 47
 48
 49class Host(Person):
 50
 51    def __init__(self, position, candies):
 52        super().__init__(position)
 53        self.candies = [Candy(mass, uranium) for mass, uranium in candies]
 54
 55    def remove_candy(self, func):
 56        if len(self.candies) == 0: 
 57            return None
 58        else:
 59            chosen_candy = func(self.candies)
 60            self.candies.remove(chosen_candy)
 61            return chosen_candy
 62
 63
 64class FluxCapacitor:
 65
 66    def __init__(self, participants):
 67        self.participants = participants
 68
 69    #This is a help function to find the indices 
 70    #of duplicates in the distances from kid to host list
 71    def indices_dup_dist(distances):
 72        list_duplicates = set()
 73        unique_dict = {}
 74        for index, elem in enumerate(distances):
 75            if elem not in unique_dict:
 76                unique_dict[elem] = index
 77            else:
 78                list_duplicates.add(unique_dict[elem])
 79                list_duplicates.add(index)
 80        return list(list_duplicates)
 81
 82    def get_victim(self):
 83        victims = set()
 84        self.participants = list(self.participants)
 85        #print(self.participants)
 86        #this for loop assigns to the kid objects a list of unvisited hosts
 87        for kid in self.participants:
 88            if isinstance(kid, Kid):
 89                kid.unvisited_hosts = [x for x in self.participants
 90                                       if isinstance(x, Host)]
 91
 92        while not victims:
 93            #this for loop sends every kid to its respective host
 94            for kid in self.participants:
 95                if isinstance(kid, Kid):
 96                    hosts = [x for x in kid.unvisited_hosts 
 97                             if isinstance(x, Host)]
 98                    hosts_coord = [x.position for x in hosts]
 99                    distances = [math.sqrt((x[0] - kid.position[0]) ** 2 + 
100                                 (x[1] - kid.position[1]) ** 2) 
101                                 for x in hosts_coord]
102                    closest_host = []
103                    
104                    if len(distances) > len(set(distances)):
105                        indices = FluxCapacitor.indices_dup_dist(distances)
106                        x_coord = []
107                        y_coord = []
108                        for index in indices:
109                            x_coord.append(hosts_coord[index][0])
110                            y_coord.append(hosts_coord[index][1])
111                        if len(set(x_coord)) > 1:
112                            least_x_val = x_coord.index(min(x_coord))
113                            closest_host = hosts[least_x_val]
114                        else:
115                            least_y_val = y_coord.index(min(y_coord))
116                            closest_host = hosts[least_y_val]
117                        kid.set_position(closest_host.position)
118                        kid.unvisited_hosts.remove(closest_host)
119                        continue
120                    
121                    closest_host = hosts[distances.index(min(distances))]
122                    kid.set_position(closest_host.position)
123                    kid.unvisited_hosts.remove(closest_host)
124
125            #this loop gives every kid candies from the hosts
126            for host in self.participants:
127                if isinstance(host, Host):
128                    if len(host.candies) > 0:
129                        kids = [x for x in self.participants
130                            if isinstance(x, Kid)]
131                        kids_at_host = []
132                        for kid in kids:
133                            if kid.position == host.position:
134                                kids_at_host.append(kid)
135                        kid_init = []
136
137                        for kid in kids_at_host:
138                            kid_init.append(kid.initiative)
139                        
140                        while len(kid_init) > 0 and len(host.candies) > 0:
141                            kid_index = kid_init.index(max(kid_init))
142                            kid = kids_at_host[kid_index]
143                            kid.add_candy(host.remove_candy(candy_max_mass))
144                            kid_init.pop(kid_index)
145                            
146            after_loop_kids = [x for x in self.participants if
147                               isinstance(x, Kid)]
148
149            kids_ur = [x.is_critical() for x in after_loop_kids]
150
151            #here we add the victims, if there are any, in a set
152            for index, boolean_val in enumerate(kids_ur):
153                if boolean_val == True:
154                    victims.add(after_loop_kids[index])
155
156            if len(victims) > 0:
157                break
158            
159            for kid in after_loop_kids:
160                if not kid.unvisited_hosts:
161                    return None
162                continue
163        
164        return victims
165
166
167def candy_max_mass(treats):
168    new_list = [x.get_mass() for x in treats]
169    max_index = new_list.index(max(new_list))
170    return treats[max_index]