1import math
  2
  3
  4class Candy:
  5
  6    def __init__(self, mass, uranium):
  7        """Initializer."""
  8        self._mass = mass
  9        self._uranium = uranium
 10
 11    def get_uranium_quantity(self):
 12        """Calculate and return the quantity of uranium in grams."""
 13        return self._mass * self._uranium
 14
 15    def get_mass(self):
 16        """Return the mass of the candy."""
 17        return self._mass
 18
 19
 20class Person:
 21
 22    def __init__(self, position):
 23        """Initializer."""
 24        self._position = position
 25
 26    def get_position(self):
 27        """Return the position (coordinates) of the person."""
 28        return self._position
 29
 30    def set_position(self, new_position):
 31        """Change the position to the new coordinates: new_position."""
 32        self._position = new_position
 33
 34
 35class Kid(Person):
 36
 37    CRITICAL_RADIATION = 20
 38
 39    def __init__(self, position, initiative):
 40        """Initializer."""
 41        super().__init__(position)
 42        self._initiative = initiative
 43        self._current_radiation = 0
 44        self._candy_bag = []
 45        self.visited_hosts = []
 46
 47    def get_initiative(self):
 48        """Return the value of the initiative."""
 49        return self._initiative
 50
 51    def add_candy(self, candy):
 52        """Get Candy object instance and add its radiation to the total current radiation sum."""
 53        self._current_radiation += candy.get_uranium_quantity()
 54        self._candy_bag.append(candy)
 55
 56    def is_critical(self):
 57        """Return bool value of whether the accumulated candy radiation has reached critical amounts."""
 58        return self._current_radiation > self.CRITICAL_RADIATION
 59
 60
 61class Host(Person):
 62
 63    def __init__(self, position, candies):
 64        """Initializer."""
 65        super().__init__(position)
 66        self.candy_bag = [Candy(*candy) for candy in candies]
 67        self.sort_candy_bag_based_on_mass()
 68
 69    def sort_candy_bag_based_on_mass(self):
 70        """Sort the candy bag so that candies go from the lowest mass to the highest mass."""
 71        self.candy_bag.sort(key=lambda candy: candy.get_mass())
 72
 73    def remove_candy(self, choose_candy_to_give):
 74        """Remove a candy from the candy bag. Which one is determined by the passed argument - function."""
 75        if self.candy_bag:
 76            candy_to_remove = choose_candy_to_give(self.candy_bag)
 77            self.candy_bag.remove(candy_to_remove)
 78            return candy_to_remove
 79        return None
 80
 81
 82class FluxCapacitor:
 83
 84    def __init__(self, participants):
 85        """Initializer."""
 86        self._participants = participants
 87        self._victims = set()
 88        self._kids = self._get_kids()
 89        self._hosts = self._get_hosts()
 90
 91    def _get_kids(self):
 92        """Return set of kids."""
 93        return {kid for kid in self._participants if isinstance(kid, Kid)}
 94
 95    def _get_hosts(self):
 96        """Return set of hosts."""
 97        return {host for host in self._participants if isinstance(host, Host)}
 98
 99    def _get_next_host_to_visit(self, kid):
100        """Find the next host for each left kid to visit based on shortest distance."""
101        list_distances_hosts = []
102        for curr_host in self._hosts:
103            if curr_host not in kid.visited_hosts:
104                curr_dist = math.dist(kid.get_position(), curr_host.get_position())
105                dist_x_y = (curr_dist, curr_host.get_position()[0], curr_host.get_position()[1], curr_host)
106                list_distances_hosts.append(dist_x_y)
107
108        list_distances_hosts.sort()
109        return list_distances_hosts[0][-1]
110
111    def _get_kids_that_will_visit(self, host):
112        """Get list of kids that will visit the current host. Sort list based on initiative."""
113        if self._kids:
114            kids_to_visit = []
115            for kid in self._kids:
116                if kid not in self._victims:
117                    if self._get_next_host_to_visit(kid) is host:
118                        kids_to_visit.append(kid)
119            kids_to_visit.sort(key=lambda curr_kid: curr_kid.get_initiative(), reverse=True)
120            return kids_to_visit
121        return None
122
123    @staticmethod
124    def choose_candy_from_host(candy_bag):
125        """Get the chunkiest candy from the host's bag. Candies in the bag are previously sorted by mass."""
126        return candy_bag[-1]
127
128    def _give_poisonous_candy(self, host):
129        """Attempt to give poisonous candy to the innocent children. If any kid dies, add it to the pile of victims."""
130        if self._get_kids_that_will_visit(host):
131            visitation_happened = False
132            for kid in self._get_kids_that_will_visit(host):
133                if host.candy_bag:
134                    visitation_happened = True
135                    candy_to_be_given = host.remove_candy(self.choose_candy_from_host)
136                    kid.add_candy(candy_to_be_given)
137                    kid.visited_hosts.append(host)
138                    kid.set_position(host.get_position())
139                    if kid.is_critical():
140                        self._victims.add(kid)
141            return visitation_happened
142        return False
143
144    def get_victim(self):
145        """Iterate visitations to distribute the candy. Return set of the first victims or None if all survived."""
146        happy_ending = False
147        while not self._victims and not happy_ending:
148            happy_ending = True
149            for host in self._hosts:
150                if self._give_poisonous_candy(host):
151                    # No happy ending until all the candy has been given
152                    happy_ending = False
153        return None if happy_ending else self._victims