import matplotlib.pyplot as plt 
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,
                 immumepct=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(a != Human.DEAD))
        deaths = numpy.random.rand(len(humans)) < self.humandiepct
        grid[humans[deaths]] = Human.DEAD

        for x, y in humans:
            if numpy.random.rand() < self.humandiepct:
                grid[x, y] = 
        
        # Pick a random, unpopulated spot
        x, y = random.choice(np.transpose(np.where(a == Human.DEAD)))
        
    def gen_humans(self):
        p_dead = 1 - self.humandens
        p_immune = self.humandens * self.immunepct
        p_healthy = p_dead - p_immune

        return np.random.choice((Human.DEAD, Human.HEALTHY, Human.IMMUNE),
                                size=(self.width, self.height),
                                p=(p_dead, p_immune, p_healthy))

    def gen_mosquitos(self):
        mosquitos = []

        count = self.width * self.height * self.mosquitodens
        for _ in count:



    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
        plt.imshow(self.humans)
        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()

    for i in range(500):
        img = np.random.randint(0, 3, (10000, 10000))
        plt.imshow(img)
        plt.pause(0.0001)
        plt.clf()