import matplotlib.pyplot as plt
import matplotlib.colors
import numpy as np
import random
from enum import IntEnum


class Model:
    def __init__(self, width=32, height=32, humandens=0.15, mosquitodens=0.10,
                 immunepct=0.1, mosqinfpct=0.1, hm_infpct=0.5, mh_infpct=0.5,
                 hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-6,
                 mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000,
                 graphical=True):
        
        self.width = width
        self.height = height

        # Determines if the simulation should be graphical
        self.graphical = graphical
        # The percentage of tiles that start as humans
        self.humandens = humandens
        # The percentage of tiles that contain mosquitos
        self.mosquitodens = mosquitodens
        # Percentage of humans that are immune
        self.immunepct = immunepct
        # Chance for a mosquito to be infectuous
        self.mosqinfpct = mosqinfpct
        # Chance for a mosquito to be infected by a human
        self.hm_infpct = hm_infpct
        # Chance for a human to be infected from a mosquito
        self.mh_infpct = mh_infpct
        # Chance that an infected human dies
        self.hinfdiepct = hinfdiepct
        # Chance for a mosquito to be hungry
        self.mhungrypct = mhungrypct
        # Chance for human to die from random causes
        self.humandiepct = humandiepct
        # Chance for a mosquito to die
        self.mosqdiepct = mosqdiepct
        # Percentage of tiles that contain mosquito nets
        self.mosqnetdens = mosqnetdens
        # The number of timesteps to run te simulation for
        self.time_steps = time_steps

        self.grid = self.gen_humans()
        self.mosquitos = self.gen_mosquitos()
        self.nets = self.gen_nets()

        # statistics
        self.stats = {
            "natural deaths": 0,
            "malaria deaths": 0,
            "total deaths": 0,
            "mosquitos fed": 0,
            "humans infected": 0,
            "mosquitos infected": 0,
            "net count": 0
        }

        if self.graphical:
            self.init_draw()
       
    def init_draw(self):
        plt.ion()
        self.colors = matplotlib.colors.ListedColormap(
            ["black", "green", "red", "yellow"])
        bounds = [Human.DEAD, Human.HEALTHY, Human.INFECTED, Human.IMMUNE]
        self.norm = matplotlib.colors.BoundaryNorm(bounds, self.colors.N)

    def recycle_human(self):
        """
        Determine if a human dies of natural causes and then replace them by a new human
        """
        # Get all living humans
        humans = np.transpose(np.where(self.grid != Human.DEAD))

        # Get a mask of humans to kill
        deaths = np.random.rand(len(humans)) < self.humandiepct

        # Kill them.
        self.grid[humans[deaths][:, 0], humans[deaths][:, 1]] = Human.DEAD

        # get num humans after killing
        humans_survive = len(np.transpose(np.where(self.grid != Human.DEAD)))

        death_count = len(humans) - humans_survive

        self.stats["natural deaths"] += death_count

        # Pick a random, unpopulated spot
        births = np.array(random.sample(list(np.transpose(np.where(self.grid == Human.DEAD))),
                                        death_count))

        # Deliver the newborns
        for birth in births:
            self.grid[birth[0]][birth[1]] = np.random.choice([Human.HEALTHY, Human.IMMUNE], 
                                                                p=[1-self.immunepct, self.immunepct])

    def do_malaria(self):
        """
        This function determines who of the infected dies from their illness
        """
        # Get all infected humans
        infected = np.transpose(np.where(self.grid == Human.INFECTED))

        # Decide which infected people die
        deaths = np.random.rand(len(infected)) < self.hinfdiepct

        # Now let's kill them
        self.grid[infected[deaths][:, 0], infected[deaths][:, 1]] = Human.DEAD

        self.stats["malaria deaths"] += len(np.where(deaths)[0])

    def feed(self):
        #TODO: dit refactoren?
        """
        Feed the mosquitos that want to and can be fed
        """
        for mos in self.mosquitos:
            if not mos.hungry:
                continue
            # state of current place on the grid where mosquito lives
            state = self.grid[mos.x, mos.y]

            if state != Human.DEAD:
                self.stats["mosquitos fed"] += 1
                mos.hungry = False
                
                # check if healthy human needs to be infected or mosquito becomes infected from eating
                if state == Human.HEALTHY and mos.infected and random.uniform(0, 1) < self.mh_infpct:
                    self.grid[mos.x, mos.y] = Human.INFECTED
                    self.stats["humans infected"] += 1
                elif state == Human.INFECTED and not mos.infected and random.uniform(0, 1) < self.hm_infpct:
                    self.stats["mosquitos infected"] += 1
                    mos.infected = True

    def determine_hunger(self):
        """
        Determines which mosquitos should get hungry
        """
        for mos in self.mosquitos:
            mos.hungry = not mos.hungry and random.uniform(0, 1) < self.mhungrypct
    
    def get_movementbox(self, x: int, y: int):
        """
        Returns indices of a moore neighbourhood around the given index
        """

        x_min = (x - 1)
        x_max = (x + 1)
        y_min = (y - 1)
        y_max = (y + 1)

        indices = [(i % self.width, j % self.height) for i in range(x_min, x_max + 1) 
                                                     for j in range(y_min, y_max + 1)]

        # remove current location from the indices
        indices.remove((x, y))

        return np.array(indices)

    def move_mosquitos(self):
        """
        Move the mosquitos to a new location, checks for mosquito nets
        """
        for mosq in self.mosquitos:
            # get the movement box for every mosquito
            movement = self.get_movementbox(mosq.x, mosq.y)

            # check for nets, and thus legal locations to go for the mosquito
            legal_moves = np.where(self.nets[tuple(movement.T)] == False)[0]
            
            # choose random new position
            new_pos = random.choice(legal_moves)

            mosq.x = movement[new_pos][0]
            mosq.y = movement[new_pos][1]

    def gen_humans(self):
        """
        Fill the grid with humans that can either be healthy or infected
        """
        # Calculate the probabilities
        p_dead = 1 - self.humandens
        p_immune = self.humandens * self.immunepct
        p_healthy = self.humandens - p_immune

        # Create the grid with humans.
        return np.random.choice((Human.DEAD, Human.HEALTHY, Human.IMMUNE),
                                size=(self.width, self.height),
                                p=(p_dead, p_healthy, p_immune))

    def gen_mosquitos(self):
        """
        Generate the list of mosquitos
        """

        mosquitos = []

        count = int(self.width * self.height * self.mosquitodens)

        # generate random x and y coordinates
        xs = np.random.randint(0, self.width, count)
        ys = np.random.randint(0, self.height, count)
        coords = list(zip(xs, ys))

        # generate the mosquitos
        for coord in coords:
            # determine if the mosquito is infected
            infected = random.uniform(0, 1) < self.mosqinfpct
            # determine if the mosquito starts out hungry
            hungry = random.uniform(0, 1) < self.mhungrypct
            mosquitos.append(Mosquito(coord[0], coord[1], infected, hungry))

        return mosquitos

    def gen_nets(self):
        """
        Generates the grid of nets
        """

        return np.random.choice([False, True], 
                                p=[1-self.mosqnetdens, self.mosqnetdens],
                                size=(self.width, self.height))

    def run(self):
        """
        This functions runs the simulation
        """
        print(chr(27) + "[2J")
        # actual simulation runs inside try except to catch keyboard interrupts and always print stats
        try:
            for t in range(self.time_steps):
                print("Simulating timestep: {}".format(t), end='\r')
                self.step()
                if self.graphical:
                    self.draw(t)
        except KeyboardInterrupt:
            pass
 
        print(chr(27) + "[2J")
        self.compile_stats()
        self.print_stats()

    def compile_stats(self):
        """
        Compiles a comprehensive list of statistics of the simulation
        """
        self.stats["total deaths"] = self.stats["malaria deaths"] + self.stats["natural deaths"]
       
        # print(np.where(self.nets))
        self.stats["net count"] = len(np.where(self.nets)[0])

    def print_stats(self):
        """
        Prints the gathered statistics from the simulation
        """
        for stat in self.stats:
            print(f"{stat}: {self.stats[stat]}")

    def step(self):
        """
        Step through a timestep of the simulation
        """
        # check who dies from malaria
        self.do_malaria()
        # check if people die from other causes
        self.recycle_human()
        # move mosquitos
        self.move_mosquitos()
        # feed hungry mosquitos
        self.feed()
        # make mosquitos hungry again
        self.determine_hunger()
        
    def draw(self, t: int):
        """
        Draws the grid of humans, tents and mosquitos
        """
        # this function draws the humans
        plt.title("t={}".format(t))
        # draw the grid
        plt.imshow(self.grid, cmap=self.colors, norm=self.norm)
        
        # draw nets
        net_locations = np.where(self.nets)
        plt.plot(net_locations[0], net_locations[1], 'w^')

        # draw mosquitos
        for mos in self.mosquitos:
            plt.plot(mos.x, mos.y, mos.get_color()+mos.get_shape())

        plt.pause(0.0001)
        plt.clf()


class Mosquito:
    def __init__(self, x: int, y: int, infected: bool, hungry: bool):
        self.x = x
        self.y = y

        self.infected = infected
        self.hungry = hungry
    
    def get_color(self):
        # returns the color for drawing, red if infected blue otherwise
        return "r" if self.infected else "b"

    def get_shape(self):
        # return the shape for drawing, o if hungry + otherwise
        return "o" if self.hungry else "+"


class Human(IntEnum):
    DEAD = 0
    HEALTHY = 1
    INFECTED = 2
    IMMUNE = 3


if __name__ == "__main__":
    model = Model(graphical=True)
    model.run()