import matplotlib.pyplot as plt
import matplotlib.colors
import numpy as np
import random
from collections import defaultdict
from enum import IntEnum
from sys import argv
from matplotlib.patches import Patch


class Model:
    def __init__(self, width=50, height=50, humandens=0.15, mosquitodens=0.10,
                 immunepct=0.01, 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.00, time_steps=50000,
                 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.infected_trend = np.zeros(time_steps)
        self.immune_trend = np.zeros(time_steps)

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

        # statistics
        self.stats = defaultdict(int)

        if self.graphical:
            self.init_draw()

    def init_draw(self):
        plt.ion()
        plt.rcParams.update({'font.size': 18})
        self.colors = matplotlib.colors.ListedColormap(
            ["white", "green", "red", "yellow"])

    def make_babies(self, n):
        if n == 0:
            return

        self.stats["humans born"] += n

        births = np.transpose(random.sample(
            list(np.transpose(np.where(self.grid == Human.DEAD))), n))

        self.grid[births[0], births[1]] = \
            np.random.choice((Human.HEALTHY, Human.IMMUNE), size=n,
                             p=(1 - self.immunepct, self.immunepct))

        # Randomly distribute a net
        nets = births.T[np.random.rand(len(births.T)) < self.mosqnetdens].T
        self.nets[nets[0], nets[1]] = True

    def recycle_human(self):
        """
        Determine if a human dies of natural causes and then replace them by a
        new human.
        """
        # Find living humans, determine if they die, and if so, kill them
        humans = np.transpose(np.where(self.grid != Human.DEAD))
        deaths = np.random.rand(len(humans)) < self.humandiepct

        locations = humans[deaths].T
        self.grid[locations[0], locations[1]] = Human.DEAD
        self.nets[locations[0], locations[1]] = False

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

        # Replace the dead humans
        self.make_babies(death_count)

    def recycle_mosquito(self):
        """
        Determine if a mosquito dies of natural causes and then replace it by
        a new mosquito.
        """
        indices = np.random.rand(len(self.mosquitos)) < self.mosqdiepct
        deaths = self.mosquitos[indices]
        for mosq in deaths:
            mosq.hungry = False
            mosq.infected = np.random.rand() < self.mosqinfpct
            mosq.x = np.random.randint(0, self.width)
            mosq.y = np.random.randint(0, self.height)

        self.stats["mosquito deaths"] += len(deaths)

    def do_malaria(self):
        """
        This function determines who of the infected dies from their illness
        """
        # Find infected humans, determine if they die, and if so, kill them
        infected = np.transpose(np.where(self.grid == Human.INFECTED))
        deaths = np.random.rand(len(infected)) < self.hinfdiepct

        locs = infected[deaths].T
        self.grid[locs[0], locs[1]] = Human.DEAD
        self.nets[locs[0], locs[1]] = False

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

        # Replace the dead humans
        self.make_babies(death_count)

    def feed(self):
        """
        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:
                continue

            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:
                mos.infected = True
                self.stats["mosquitos infected"] += 1

    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)])[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 np.array(mosquitos)

    def gen_nets(self):
        """
        Generates the grid of nets
        """
        humans = np.transpose(np.where(self.grid != Human.DEAD))
        positions = humans[np.random.choice(
            len(humans), size=round(self.mosqnetdens * len(humans)))].T

        grid = np.zeros((self.width, self.height), dtype=bool)
        grid[positions[0], positions[1]] = True
        return grid

    def run(self):
        """
        This functions runs the simulation
        """
        # Actual simulation runs inside try except to catch keyboard interrupts
        # and always print stats
        try:
            for t in range(self.time_steps):
                self.step()

                total = np.count_nonzero(self.grid != Human.DEAD)
                infected = np.count_nonzero(self.grid == Human.INFECTED) / \
                    total
                immune = np.count_nonzero(self.grid == Human.IMMUNE) / total

                self.infected_trend[t] = infected
                self.immune_trend[t] = immune
                print(f"t={t}, infected={infected:.2f}, immune={immune:.2f}",
                      end='\r')

                if self.graphical and t % 1000 == 0:
                    self.draw(t)
        except KeyboardInterrupt:
            pass

        print()
        self.compile_stats()
        self.print_stats()
        self.draw_graph(t)

    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"]
        self.stats["final immunity"] = self.immune_trend[-1]

        self.stats["net count"] = len(np.where(self.nets)[0])

    def print_stats(self):
        """
        Prints the gathered statistics from the simulation
        """
        for stat, value in sorted(self.stats.items()):
            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()
        # check if mosquitoes die
        self.recycle_mosquito()
        # 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
        """
        plt.clf()

        plt.title("t={}".format(t))

        # draw the grid
        plt.imshow(self.grid, cmap=self.colors)

        # 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.y, mos.x, mos.get_color()+mos.get_shape())

        # draw the legend
        dead_patch = Patch(color="green", label="Healthy human")
        immune_patch = Patch(color="yellow", label="Immune human")
        infected_patch = Patch(color="red", label="Infected human")
        plt.legend(handles=[dead_patch, immune_patch, infected_patch],
                   loc=9, bbox_to_anchor=(0.5, -0.03), ncol=5)

        plt.pause(0.0001)

    def draw_graph(self, t):
        plt.ioff()
        plt.clf()

        plt.plot(np.arange(t), self.infected_trend[:t],
                 label="infected humans")
        plt.plot(np.arange(t), self.immune_trend[:t],
                 label="immune humans")

        plt.legend()

        plt.xlabel("timestep")
        plt.ylabel("prevalence")

        plt.show()


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__":
    try:
        graphical = argv[1] != "-ng"
    except IndexError:
        graphical = True

    model = Model(graphical=graphical)
    model.run()