#!/usr/bin/env python3

import random
import time
from matplotlib import pyplot

HUMAN_DEAD = 0
HUMAN_HEALTHY = 1
HUMAN_INFECTED = 2
HUMAN_IMMUNE = 3

MOSQ_POSITION = 0
MOSQ_INFECTED = 1
MOSQ_HUNGRY = 2


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

        # The initial density of humans
        self.humandens = humandens

        # The initial density mosquitos
        self.mosquitodens = mosquitodens

        # The chance that a human is immune
        self.immumepct = immumepct

        # The chance that a mosquito is born infected
        self.mosqinfpct = mosqinfpct

        # The chance that a human infects a mosquito
        self.hm_infpct = hm_infpct

        # The chance that a mosquito infects a human
        self.mh_infpct = mh_infpct

        # The chance that a mosquito gets hungry
        self.mhungrypct = mhungrypct

        # The chance that an infected human dies
        self.hinfdiepct = hinfdiepct

        # The chance that an uninfected/immume human dies
        self.humandiepct = humandiepct

        # The chance that a mosquito dies
        self.mosqdiepct = mosqdiepct

        # The density of mosquito nets
        self.mosqnetdens = mosqnetdens

        maxhumans = round(self.humandens * self.width * self.height)
        maxmosquitos = round(self.mosquitodens * self.width * self.height)
        maxmosqnets = round(self.mosqnetdens * self.width * self.height)

        # Human: 0=dead, 1=healthy, 2=infected, 3=immume
        # Mosquito: [(x, y), infected, hungry]
        self.humans = [[HUMAN_DEAD for _ in range(width)]
                       for _ in range(height)]
        self.mosquitos = []
        self.mosqnets = [[False]*width for _ in range(height)]
        for _ in range(maxhumans):
            self.addhuman()
        for _ in range(maxmosquitos):
            self.addmosquito(random.uniform(0, 1) < mosqinfpct)
        for _ in range(maxmosqnets):
            self.addmosqnet()

    def addhuman(self):
        x, y = random.choice([(x, y)
                              for x in range(self.width) for y in range(self.height)
                              if self.humans[y][x] == HUMAN_DEAD])

        if random.uniform(0, 1) < self.immumepct:
            self.humans[y][x] = HUMAN_IMMUNE
        else:
            self.humans[y][x] = HUMAN_HEALTHY

    def addmosquito(self, infected=False):
        height = len(self.humans)
        width = len(self.humans[0])
        self.mosquitos.append([(random.randint(0, width-1),
                                random.randint(0, height-1)), infected, False])

    def addmosqnet(self):
        x, y = random.choice([(x, y)
                              for x in range(self.width) for y in range(self.height)
                              if self.mosqnets[y][x] == False])
        self.mosqnets[y][x] = True

    def move(self, pos):
        x, y = pos
        possibs = [(x, y) for x, y
                   in [(x-1, y), (x, y-1), (x+1, y), (x, y+1)]
                   if x > -1 and x < self.width
                   and y > -1 and y < self.height
                   and not self.mosqnets[y][x]]
        try:
            return random.choice(possibs)
        except IndexError:
            return (x, y)

    def getstates(self):
        humanhealthy = sum(1 for row in self.humans for human in row if human == HUMAN_HEALTHY)
        humaninfected = sum(1 for row in self.humans for human in row if human == HUMAN_INFECTED)
        humanimmune = sum(1 for row in self.humans for human in row if human == HUMAN_IMMUNE)
        mosqinfected = sum(1 for mosq in self.mosquitos if mosq[MOSQ_INFECTED])
        mosqhungry = sum(1 for mosq in self.mosquitos if mosq[MOSQ_HUNGRY])
        return humanhealthy, humaninfected, humanimmune, mosqinfected, mosqhungry

    def step(self):
        # Copy the human values
        humans = [human_row[:] for human_row in self.humans]

        # Now I am become death, the destroyer of worlds
        killed = 0
        for human_row in humans:
            for i, human in enumerate(human_row):
                if human == HUMAN_INFECTED \
                   and random.uniform(0, 1) < self.hinfdiepct:
                    human_row[i] = HUMAN_DEAD
                    killed += 1
                elif human != HUMAN_DEAD \
                     and random.uniform(0, 1) < self.humandiepct:
                    human_row[i] = HUMAN_DEAD
                    killed += 1

        mosqkilled = 0
        mosquitos = []
        for pos, infected, hungry in self.mosquitos:
            # Randomly move the mosquito to a new position
            (x, y) = self.move(pos)

            if hungry and self.humans[y][x]:
                # If a mosquito is hungry and in the same cell as a mosquito,
                # human, it'll bite
                hungry = False
                human = self.humans[y][x]
                if human == HUMAN_INFECTED:
                    # If the human is infected, become infected
                    infected = random.uniform(0, 1) < self.hm_infpct

                if infected and human == HUMAN_HEALTHY and \
                   random.uniform(0, 1) < self.mh_infpct:
                    # If the human is healthy and not immune, there's a
                    # mh_infpct chance it will become infected.
                    humans[y][x] = HUMAN_INFECTED

            elif not hungry:
                # If the mosquito is not hungry, there's a mhungrypct
                # chance it'll become hungry
                hungry = random.uniform(0, 1) < self.mhungrypct

            # There's a mosqdiepct% chance that a mosq will die
            if random.uniform(0, 1) < self.mosqdiepct:
                mosqkilled += 1
            else:
                mosquitos.append(((x, y), infected, hungry))

        # Send all new values to the model
        self.mosquitos = mosquitos
        self.humans = humans

        for _ in range(killed):
            # Add a human for each one that had died
            self.addhuman()
        for _ in range(mosqkilled):
            # Add a mosqito for each one that had died
            self.addmosquito()

    @staticmethod
    def drawhuman(human):
        if human == HUMAN_DEAD:
            return (31, 31, 31)
        elif human == HUMAN_HEALTHY:
            return (0, 127, 0)
        elif human == HUMAN_INFECTED:
            return (255, 0, 0)
        elif human == HUMAN_IMMUNE:
            return (0, 255, 0)

    @staticmethod
    def drawmosq(mosq):
        _, hungry, infected = mosq
        if hungry and infected:
            return (255, 61, 0)
        elif hungry:  # and not infected
            return 191, 128, 0
        elif infected:  # and not hungry
            return (127, 0, 0)
        else:  # neither
            return (0, 0, 223)

    def printlegend(self):
        print('Humans (large blocks):    ', end='')
        print('\033[38;2;%i;%i;%imDead\033[0m' %
              self.drawhuman(HUMAN_DEAD), end=', ')
        print('\033[38;2;%i;%i;%imHealthy\033[0m' %
              self.drawhuman(HUMAN_HEALTHY), end=', ')
        print('\033[38;2;%i;%i;%imInfected\033[0m' %
              self.drawhuman(HUMAN_INFECTED), end=', ')
        print('\033[38;2;%i;%i;%imImmune\033[0m' %
              self.drawhuman(HUMAN_IMMUNE))

        print('Mosquitos (small blocks): ', end='')
        print('\033[38;2;%i;%i;%imHungry and Infected\033[0m' %
              self.drawmosq(((0, 0), True, True)), end=', ')
        print('\033[38;2;%i;%i;%imHungry\033[0m' %
              self.drawmosq(((0, 0), True, False)), end=', ')
        print('\033[38;2;%i;%i;%imInfected\033[0m' %
              self.drawmosq(((0, 0), False, True)), end=', ')
        print('\033[38;2;%i;%i;%imNeither\033[0m' %
              self.drawmosq(((0, 0), False, False)))

        print('Mosquito nets:            /\\')

    def print(self):
        data = [[self.drawhuman(human) for human in row]
                for row in self.humans]

        # Draw the mosquitos
        mosqs = [[None] * self.width for _ in range(self.height)]
        for mosq in self.mosquitos:
            x, y = mosq[0]
            mosqs[y][x] = self.drawmosq(mosq)

        print('\033[2J\033[0;0H' + '\033[0m\n'.join(''.join([
              '\033[48;2;%i;%i;%im' % data[row][x] +
              ('\033[38;2;%i;%i;%im\u2590\u258c' % mosqs[row][x] if mosqs[row][x]
              else '\033[38;2;255;255;255m/\\' if self.mosqnets[row][x]
              else '  ')
              for x in range(self.width)])
              for row in range(0, self.height)) + '\033[0m')

        self.printlegend()

    def gui(self, printsteps=1):
        try:
            fps = 15
            while True:
                prevtime = time.time()
                for _ in range(printsteps):
                    self.step()
                self.print()
                time.sleep(max(0, (1/fps) - (time.time()-prevtime)))
        except KeyboardInterrupt:
            return


def frange(start, end, steps):
    stepsize = (end - start) / (steps-1)
    for n in range(steps):
        yield start + stepsize * n


def statesvstime(tmax=40000):
    model = Model(64, 64)
    tr = list(range(tmax))
    results = [None]*tmax
    for t in tr:
        print(t/tmax, end='\r')
        model.step()
        results[t] = model.getstates()
    humanhealthy, humaninfected, humanimmune, mosqinfected, mosqhungry = zip(*results)

    pyplot.plot(tr, humanhealthy, label='Healthy humans')
    pyplot.plot(tr, humaninfected, label='Infected humans')
    pyplot.plot(tr, humanimmune, label='Immune humans')
    pyplot.plot(tr, mosqinfected, label='Infected mosquitos')
    pyplot.plot(tr, mosqhungry, label='Hungry Mosquitos')
    pyplot.legend()
    pyplot.xlabel('time (steps)')
    pyplot.ylabel('amount')
    pyplot.show()


def netvsimmune(tmin=10000, denscnt=200):
    dr = list(frange(0, 1, denscnt))
    results = [None]*denscnt
    for i, dens in enumerate(dr):
        print(dens, end='\r')
        model = Model(16, 16, mosqnetdens=dens)
        for t in range(tmin):
            model.step()
        results[i] = sum(1 for row in model.humans for human in row if human == HUMAN_IMMUNE)

    pyplot.plot(dr, results, label='Immunity')
    pyplot.legend()
    pyplot.xlabel('Mosquito net density')
    pyplot.ylabel('Immune humans')
    pyplot.show()


def main():
    model = Model()
    model.gui()

    statesvstime()
    netvsimmune()


if __name__ == '__main__':
    random.seed()
    main()