import sys

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# set seed for reproducibility
np.random.seed(100)


def init_population(size, age_mean=45, age_max=105, bounds_x=None, bounds_y=None):

    # Mean Age Source
    # https://www.statista.com/statistics/624303/average-age-of-the-population-in-germany/

    '''initialized the population for the simulation

    the population matrix for this simulation has the following columns:

    0 : unique ID
    1 : current x coordinate
    2 : current y coordinate
    3 : current heading in x direction
    4 : current heading in y direction
    5 : current speed
    6 : current state (0=healthy, 1=sick, 2=immune, 3=dead)
    7 : age
    8 : infected_since (frame the person got infected)
    9 : recovery vector (used in determining when someone recovers or dies)

    Keyword arguments
    -----------------
    pop_size : int
        the size of the population

    mean_age : int
        the mean age of the population. Age affects mortality chances

    max_age : int
        the max age of the population

    xbounds : 2d array
        lower and upper bounds of x axis

    ybounds : 2d array
        lower and upper bounds of y axis
    '''

    # initialize population matrix
    if bounds_x is None:
        bounds_x = [0, 1]
    if bounds_y is None:
        bounds_y = [0, 1]
    population = np.zeros((size, 10))

    # initalize unique IDs
    population[:, 0] = [x for x in range(size)]

    # initialize random coordinates
    population[:, 1] = np.random.uniform(low=bounds_x[0] + 0.05, high=bounds_x[1] - 0.05,
                                         size=(size,))
    population[:, 2] = np.random.uniform(low=bounds_y[0] + 0.05, high=bounds_y[1] - 0.05,
                                         size=(size,))

    # initialize random headings -1 to 1
    population[:, 3] = np.random.normal(loc=0, scale=1 / 3,
                                        size=(size,))
    population[:, 4] = np.random.normal(loc=0, scale=1 / 3,
                                        size=(size,))

    # initialize random speeds
    population[:, 5] = np.random.normal(0.01, 0.01 / 3)

    # initalize ages
    std_age = (age_max - age_mean) / 3
    population[:, 7] = np.int32(np.random.normal(loc=age_mean,
                                                 scale=std_age,
                                                 size=(size,)))

    population[:, 7] = np.clip(population[:, 7], a_min=0,
                               a_max=age_max)  # clip those younger than 0 years

    # build recovery_vector
    # TODO: make risks age dependent
    population[:, 9] = np.random.normal(loc=0.5, scale=0.5 / 3, size=(size,))

    return population


def update_positions(population):
    '''update positions of all people

    Uses heading and speed to update all positions for
    the next time step

    Keyword arguments
    -----------------
    population : ndarray
        the array numpy containing all the population information
    '''

    # update positions
    # x
    population[:, 1] = population[:, 1] + (population[:, 3] * population[:, 5])
    # y
    population[:, 2] = population[:, 2] + (population[:, 4] * population[:, 5])

    return population


def out_of_bounds(population, xbounds, ybounds):
    '''checks which people are about to go out of bounds and corrects


    '''
    # update headings and positions where out of bounds
    # update x heading
    # determine number of elements that need to be updated
    shp = population[:, 3][(population[:, 1] <= xbounds[:, 0]) &
                           (population[:, 3] < 0)].shape
    population[:, 3][(population[:, 1] <= xbounds[:, 0]) &
                     (population[:, 3] < 0)] = np.random.normal(loc=0.5,
                                                                scale=0.5 / 3,
                                                                size=shp)

    shp = population[:, 3][(population[:, 1] >= xbounds[:, 1]) &
                           (population[:, 3] > 0)].shape
    population[:, 3][(population[:, 1] >= xbounds[:, 1]) &
                     (population[:, 3] > 0)] = -np.random.normal(loc=0.5,
                                                                 scale=0.5 / 3,
                                                                 size=shp)

    # update y heading
    shp = population[:, 4][(population[:, 2] <= ybounds[:, 0]) &
                           (population[:, 4] < 0)].shape
    population[:, 4][(population[:, 2] <= ybounds[:, 0]) &
                     (population[:, 4] < 0)] = np.random.normal(loc=0.5,
                                                                scale=0.5 / 3,
                                                                size=shp)

    shp = population[:, 4][(population[:, 2] >= ybounds[:, 1]) &
                           (population[:, 4] > 0)].shape
    population[:, 4][(population[:, 2] >= ybounds[:, 1]) &
                     (population[:, 4] > 0)] = -np.random.normal(loc=0.5,
                                                                 scale=0.5 / 3,
                                                                 size=shp)

    return population


def update_randoms(population, heading_update_chance=0.02,
                   speed_update_chance=0.02):
    '''updates random states such as heading and speed'''

    # randomly update heading
    # x
    update = np.random.random(size=(pop_size,))
    shp = update[update <= heading_update_chance].shape
    population[:, 3][update <= heading_update_chance] = np.random.normal(loc=0,
                                                                         scale=1 / 3,
                                                                         size=shp)
    # y
    update = np.random.random(size=(pop_size,))
    shp = update[update <= heading_update_chance].shape
    population[:, 4][update <= heading_update_chance] = np.random.normal(loc=0,
                                                                         scale=1 / 3,
                                                                         size=shp)
    # randomize speeds
    update = np.random.random(size=(pop_size,))
    shp = update[update <= heading_update_chance].shape
    population[:, 5][update <= heading_update_chance] = np.random.normal(loc=0.01,
                                                                         scale=0.01 / 3,
                                                                         size=shp)
    return population


def infect(population, infection_range, infection_chance, frame):
    '''finds new infections


    Keyword arguments
    -----------------


    '''

    # find new infections
    infected_previous_step = population[population[:, 6] == 1]

    new_infections = []

    # if less than half are infected, slice based on infected (to speed up computation)
    if len(infected_previous_step) < (pop_size // 2):
        for patient in infected_previous_step:
            # define infection zone for patient
            infection_zone = [patient[1] - infection_range, patient[2] - infection_range,
                              patient[1] + infection_range, patient[2] + infection_range]

            # find healthy people surrounding infected patient
            indices = np.int32(population[:, 0][(infection_zone[0] < population[:, 1]) &
                                                (population[:, 1] < infection_zone[2]) &
                                                (infection_zone[1] < population[:, 2]) &
                                                (population[:, 2] < infection_zone[3]) &
                                                (population[:, 6] == 0)])
            for idx in indices:
                # roll die to see if healthy person will be infected
                if np.random.random() < infection_chance:
                    population[idx][6] = 1
                    population[idx][8] = frame
                    new_infections.append(idx)

    else:
        # if more than half are infected slice based in healthy people (to speed up computation)
        healthy_previous_step = population[population[:, 6] == 0]
        sick_previous_step = population[population[:, 6] == 1]
        for person in healthy_previous_step:
            # define infecftion range around healthy person
            infection_zone = [person[1] - infection_range, person[2] - infection_range,
                              person[1] + infection_range, person[2] + infection_range]

            if person[6] == 0:  # if person is not already infected, find if infected are nearby
                # find infected nearby healthy person
                poplen = len(sick_previous_step[:, 6][(infection_zone[0] < sick_previous_step[:, 1]) &
                                                      (sick_previous_step[:, 1] < infection_zone[2]) &
                                                      (infection_zone[1] < sick_previous_step[:, 2]) &
                                                      (sick_previous_step[:, 2] < infection_zone[3]) &
                                                      (sick_previous_step[:, 6] == 1)])

                if poplen > 0:
                    if np.random.random() < (infection_chance * poplen):
                        # roll die to see if healthy person will be infected
                        population[np.int32(person[0])][6] = 1
                        population[np.int32(person[0])][8] = frame
                        new_infections.append(np.int32(person[0]))

    if len(new_infections) > 0:
        print('at timestep %i these people got sick: %s' % (frame, new_infections))

    return population


def recover_or_die(population, frame, recovery_duration, mortality_chance):
    '''see whether to recover or die

    '''

    # find sick people
    sick_people = population[population[:, 6] == 1]

    # define vector of how long everyone has been sick
    illness_duration_vector = frame - sick_people[:, 8]

    recovery_odds_vector = (illness_duration_vector - recovery_duration[0]) / np.ptp(recovery_duration)
    recovery_odds_vector = np.clip(recovery_odds_vector, a_min=0, a_max=None)

    # update states of sick people
    indices = sick_people[:, 0][recovery_odds_vector >= sick_people[:, 9]]

    cured = []
    died = []

    # decide whether to die or recover
    for idx in indices:
        if np.random.random() <= mortality_chance:
            # die
            sick_people[:, 6][sick_people[:, 0] == idx] = 3
            died.append(np.int32(sick_people[sick_people[:, 0] == idx][:, 0][0]))
        else:
            # recover (become immune)
            sick_people[:, 6][sick_people[:, 0] == idx] = 2
            cured.append(np.int32(sick_people[sick_people[:, 0] == idx][:, 0][0]))

    if len(died) > 0:
        print('at timestep %i these people died: %s' % (frame, died))
    if len(cured) > 0:
        print('at timestep %i these people recovered: %s' % (frame, cured))

    # put array back into population
    population[population[:, 6] == 1] = sick_people

    return population


def update(frame, population, infection_range=0.01, infection_chance=0.03,
           recovery_duration=(200, 500), mortality_chance=0.02,
           xbounds=[0.02, 0.98], ybounds=[0.02, 0.98], wander_range=0.05,
           visualise=True, infected_plot=[]):
    # add one infection to jumpstart
    if frame == 50:
        population[0][6] = 1
        population[0][8] = 75

    # update out of bounds
    # define bounds arrays
    _xbounds = np.array([[xbounds[0] + 0.02, xbounds[1] - 0.02]] * len(population))
    _ybounds = np.array([[ybounds[0] + 0.02, ybounds[1] - 0.02]] * len(population))
    population = out_of_bounds(population, _xbounds, _ybounds)

    # update randoms
    population = update_randoms(population)

    # for dead ones: set speed and heading to 0
    population[:, 3:5][population[:, 6] == 3] = 0

    # update positions
    population = update_positions(population)

    # find new infections
    population = infect(population, infection_range, infection_chance, frame)
    infected_plot.append(len(population[population[:, 6] == 1]))

    # recover and die
    population = recover_or_die(population, frame, recovery_duration, mortality_chance)

    if visualise:
        # construct plot and visualise
        spec = fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5, 2])
        ax1.clear()
        ax2.clear()

        ax1.set_xlim(xbounds[0] - 0.02, xbounds[1] + 0.02)
        ax1.set_ylim(ybounds[0] - 0.02, ybounds[1] + 0.02)

        healthy = population[population[:, 6] == 0][:, 1:3]
        ax1.scatter(healthy[:, 0], healthy[:, 1], color='gray', s=2, label='healthy')

        infected = population[population[:, 6] == 1][:, 1:3]
        ax1.scatter(infected[:, 0], infected[:, 1], color='red', s=2, label='infected')

        immune = population[population[:, 6] == 2][:, 1:3]
        ax1.scatter(immune[:, 0], immune[:, 1], color='green', s=2, label='immune')

        fatalities = population[population[:, 6] == 3][:, 1:3]
        ax1.scatter(fatalities[:, 0], fatalities[:, 1], color='black', s=2, label='fatalities')

        # add text descriptors
        ax1.text(xbounds[0],
                 ybounds[1] + 0.03,
                 'timestep: %i, total: %i, healthy: %i infected: %i immune: %i fatalities: %i' % (frame,
                                                                                                  len(population),
                                                                                                  len(healthy),
                                                                                                  len(infected),
                                                                                                  len(immune),
                                                                                                  len(fatalities)),
                 fontsize=6)

        ax2.set_title('number of infected')
        ax2.text(0, pop_size * 0.05,
                 'https://github.com/paulvangentcom/python-corona-simulation',
                 fontsize=6, alpha=0.5)
        # ax2.set_xlim(0, simulation_steps)
        ax2.set_ylim(0, pop_size + 100)
        ax2.plot(infected_plot, color='gray')

        plt.savefig('render/%s.png' % frame)


if __name__ == '__main__':
    # set simulation parameters
    pop_size = 2000
    simulation_steps = 10000
    xbounds = [0, 1]
    ybounds = [0, 1]

    population = init_population(pop_size)

    # define figure
    fig = plt.figure(figsize=(5, 7))
    spec = fig.add_gridspec(ncols=1, nrows=2, height_ratios=[5, 2])

    ax1 = fig.add_subplot(spec[0, 0])
    plt.title('infection simulation')
    plt.xlim(xbounds[0], xbounds[1])
    plt.ylim(ybounds[0], ybounds[1])

    ax2 = fig.add_subplot(spec[1, 0])
    ax2.set_title('number of infected')
    ax2.set_xlim(0, simulation_steps)
    ax2.set_ylim(0, pop_size + 100)

    # start animation loop through matplotlib visualisation
    animation = FuncAnimation(fig, update, fargs=(population,), frames=simulation_steps, interval=33)
    plt.show()

    # alternatively dry run simulation without visualising
    # for i in range(simulation_steps):
    # update(i, population, visualise=False)
    # sys.stdout.write('\r')
    # sys.stdout.write('%i: %i/%i' %(i, len(infected), pop_size))