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):
        
        plt.ion()
        self.width = width
        self.height = height

        # 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 infected
        self.mosqinfpct = mosqinfpct
        # Chance for a human to be infected by a mosquito bite
        self.hm_infpct = hm_infpct
        # Chance for a mosquito to be infected from biting an infected human
        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.humans = self.gen_humans()
        self.mosquitos = self.gen_mosquitos()
        self.init_draw()
       
    def init_draw(self):
        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):
        # Get all living humans
        humans = np.transpose(np.where(self.humans != Human.DEAD))

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

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

        # Pick a random, unpopulated spot
        x, y = random.choice(np.transpose(np.where(self.humans == Human.DEAD)),
                             size=np.count_nonzero(deaths))


    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.humans == Human.INFECTED))

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

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

    def gen_humans(self):
        # 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):
        count = self.width * self.height * self.mosquitodens

    def run(self):
        for t in range(self.time_steps):
            self.step()
            self.draw(t)

    def step(self):
        self.do_malaria()

    def draw(self, t):
        # this function draws the humans
        plt.title("t={}".format(t))
        plt.imshow(self.humans, cmap=self.colors, norm=self.norm)

        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


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


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