from bokeh.plotting import figure
from bokeh.io import curdoc
from bokeh.layouts import row
from bokeh.models.widgets import Slider
from bokeh.models import ColumnDataSource
import pandas


class Person(object):
    """ Person class represents a person in the population """

    def __init__(self,
                 infected=False,
                 vaccinated=False,
                 x=0,
                 y=0,
                 infectedByx=0,
                 infectedByy=0,
                 stepInfected=0):
        self.infected = False
        self.vaccinated = False
        self.x = 0
        self.y = 0
        self.infectedByx = 0
        self.infectedByy = 0
        self.stepInfected = 0


def initialisePlot():
    plot = figure(plot_height=400, plot_width=400)
    plot.xgrid.grid_line_color = None
    plot.ygrid.grid_line_color = None
    plot.axis.visible = False
    plot.toolbar.logo = None
    plot.toolbar_location = None

    return plot


def dataChange(step):
    for i in range(100):
        current = dataFrame.iloc[i]
        if current.stepInfected < step:
            plot.circle(current.x, current.y, size=2, color='red')
            plot.line([current.x, current.infectedByx],
                      [current.y, current.infectedByy],
                      color='red')


def update_data(attrname, old, new):
    # get the current value
    step = round(stepSlider.value)
    dataChange(step)


# Set up data
dataFrame = pandas.read_csv('DIRECTORY TO THE CSV FILE HERE')
source = ColumnDataSource(dataFrame)
steps = 9
circleSize = 9

# Set up plot
plot = initialisePlot()
plot.circle('x', 'y', source=source, size=circleSize, color='blue')

stepSlider = Slider(title="Step", value=0, start=0, end=10, step=1)
stepSlider.on_change('value', update_data)

curdoc().add_root(row(plot, stepSlider, width=400))
curdoc().title = "Sliders"