
import plotly.graph_objects as go
import numpy as np
import random
from plotly.subplots import make_subplots

steps = 500 #
day = 5 # means day = 10*steps
agents = 300
temp = 0.1
agents_radius = 2

height = 100 # height of the square
width = 100 # width of the square

x0 = [random.uniform(0,width) for i in range(agents)]
y0 = [random.uniform(0,height) for i in range(agents)]

x = x0
y = y0

# colors
color_susceptible = 'MediumSeaGreen'
color_exposed = 'pink'
color_infected = 'red'
color_recovered = 'gray'

# create list of agents
agents_list = [color_infected] # 1 infected
for i in range(1,agents): # agents-1 susceptibles
    agents_list.append(color_susceptible)

# numbers at the beginning
number_susceptible = agents-1
number_exposed = 1
number_infected = 1
number_recovered = 0

#create lists for counting days of different states
daycounting_list = [0 for i in range(agents)]
daycounting_list[0] = 10

# Create two lists, fill both with random numbers
x_speeds = [temp*random.uniform(-10,10) for i in range(agents)]
y_speeds = [temp*random.uniform(-10,10) for i in range(agents)]

# Only every second agent is moving
'''
for i in range(0,agents,2):
    x_speeds[i] = 0
    y_speeds[i] = 0
'''

#lists for graphics
susceptible_series = []
infected_series = []
recovered_series = []
frames = []

for step in range(steps):

    for k in range(agents):

        if x[k] >= width or x[k] <= 0:
            x_speeds[k] *= -1
        x[k] = x[k] + x_speeds[k]

        if y[k] >= height or y[k] <= 0:
            y_speeds[k] *= -1
        y[k] = y[k] + y_speeds[k]


    for dot1 in range(agents):
        for dot2 in range(dot1,agents):
            # proove the distance between the agents
            if abs(x[dot1]-x[dot2]) <= agents_radius and abs(y[dot1]-y[dot2]) <= agents_radius:
                # proove if one of them is infected
                if agents_list[dot1] == color_infected and agents_list[dot2] == color_susceptible: # and random.random() < 0.3: # 45% probability of getting infected
                    agents_list[dot2] = color_exposed
                    number_exposed += 1
                    number_susceptible -= 1
                    daycounting_list[dot2] = 1
                elif agents_list[dot1] == color_susceptible and agents_list[dot2] == color_infected: # and random.random() < 0.3:
                    agents_list[dot1] = color_exposed
                    number_exposed += 1
                    number_susceptible -= 1
                    daycounting_list[dot1] = 1

        # +1 day in daycounting_list if a day passed and the agent is exposed
        if daycounting_list[dot1] >= 1 and daycounting_list[dot1] < 4*day:
            daycounting_list[dot1] += 1
        if daycounting_list[dot1] >= 4*day and daycounting_list[dot1] < 14*day:
            if daycounting_list[dot1] == 4*day:
                number_infected += 1
                number_exposed -= 1
            daycounting_list[dot1] += 1
            agents_list[dot1] = color_infected
        if daycounting_list[dot1] == 14*day:
            daycounting_list[dot1] = 15*day
            agents_list[dot1] = color_recovered
            number_recovered += 1
            number_infected -= 1

    infected_series.append(number_infected)
    susceptible_series.append(number_susceptible)
    recovered_series.append(number_recovered)

    frames.append(go.Frame(data=[go.Scatter(x=x, y=y, \
        mode = "markers", \
        marker=dict(size = 10, color=agents_list, \
        )),
        go.Scatter(y = infected_series),go.Scatter(y = susceptible_series),go.Scatter(y = recovered_series)] ))


# ======= DARSTELLUNG =======

fig = make_subplots(
    rows=2, cols=1, row_heights=[4, 2], \
    subplot_titles=('Agents Simulation', 'Data Visualisation'))

# Add square
fig.add_shape(type="rect",
    x0=-agents_radius, y0=-agents_radius, x1=width+agents_radius, y1=height+agents_radius,
    line=dict(color="black")
)

fig.add_trace(go.Scatter(x = x0, y = y0, name = 'Agents', mode = "markers", \
marker=dict(size = 10, color=color_susceptible, \
)), row=1, col=1)  # trace of index 1

fig.add_trace(go.Scatter(y=[0], marker_color = "red", marker_size=1, name = 'Infected'), row=2, col=1)  #trace of index 2
fig.add_trace(go.Scatter(y=[0], marker_color = "MediumSeaGreen", marker_size=1, name = 'Susceptible'), row=2, col=1)  #trace of index 3
fig.add_trace(go.Scatter(y=[0], marker_color = "gray", marker_size=1, name = 'Recovered'), row=2, col=1)  #trace of index 4

fig.frames=frames

button = dict(label='Start',
             method='animate',
             args=[None, dict(frame=dict(duration=50, redraw=False),
                              transition=dict(duration=0),
                              fromcurrent=True,
                              mode='immediate')])
fig.update_layout(updatemenus=[dict(type='buttons',
                              showactive=False,
                              y=0,
                              x=1.05,
                              xanchor='left',
                              yanchor='bottom',
                              buttons=[button] )],
                  width=800, height=500)

# Set axes properties
fig.update_layout(xaxis1_showgrid = False, xaxis1_showticklabels = False, yaxis1_showgrid = False, yaxis1_showticklabels = False, \
                  xaxis2_range=[0, steps], xaxis2_autorange=False,
                  xaxis2_gridcolor = 'lightgrey', xaxis2_linecolor = 'black',
                  yaxis2_range=[0, agents + 1], yaxis2_autorange=False,
                  yaxis2_gridcolor = 'lightgrey', yaxis2_linecolor = 'black')

fig.update_layout(autosize=False, width=1000, height=1000,
                  paper_bgcolor="white", plot_bgcolor = 'white')

fig.show()