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=500):
        self.width = width
        self.height = height

        self.humandens = humandens
        self.mosquitodens = mosquitodens
        self.immunepct = immunepct
        self.mosqinfpct = mosqinfpct
        self.hm_infpct = hm_infpct
        self.mh_infpct = mh_infpct
        self.hinfdiepct = hinfdiepct
        self.mhungrypct = mhungrypct
        self.humandiepct = humandiepct
        self.mosqdiepct = mosqdiepct
        self.mosqnetdens = mosqnetdens

        self.time_steps = time_steps

        self.humans = self.gen_humans()
        self.mosquitos = self.gen_mosquitos()

    def recycle_human(self):
        # 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

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

    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 _ in range(self.time_steps):
            self.step()
            self.draw()

    def step(self):
        # dingen
        pass

    def draw(self):
        # this function draws the humans
        colors = matplotlib.colors.ListedColormap(
            ["black", "green", "red", "yellow"])
        bounds = [Human.DEAD, Human.HEALTHY, Human.INFECTED, Human.IMMUNE]
        norm = matplotlib.colors.BoundaryNorm(bounds, colors.N)

        plt.imshow(self.humans, cmap=colors, norm=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__":
    plt.ion()
    model = Model()
    while True:
        model.draw()