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

#haven't implemented real data, this is just a simulation for 

class Simulation():
    
    def __init__(self):
        #SIMULATION PARAMETERS
        self.n=350  #number of individuals
        self.p_infected = 1  #initial percentage of infection (0-100%)
        self.r_contagious=2 #transmission radius (0-100)
        self.p_contagious= 10  #transmission rate (0-100%)
        self.p_isolation =75  #isolation rate (0-100%)
        self.t_infected=100   #recovery time in FRAMES (0-infinity) 
        self.d_rate = 5 #death rate (0-100%)
        self.reinfect_rate = 0.05 #re-infect rate (0-100%)
        self.infected=0
        self.dead=0
        self.recovered=0
        self.people=[]

    def instantiate(self):

        #for instantiating
        

        #Initialise loop. Creating all the individuals in random positions. Infecting some and putting some in quarantine.
        for i in range(self.n):
            p = Person(i,np.random.random()*100, np.random.random()*100,
                        np.random.random() * 100, np.random.random() * 100,
                        (np.random.random()+0.5)*100, self.t_infected, False)
                        #(index, xpos, ypos, xobj, yobj, velocity, infected_time, fixed=false)

            if np.random.random()<self.p_infected/100: #in percentage
                p.infect(0)
                self.infected+=1
            if np.random.random()<self.p_isolation/100: #in percentage
                p.fixed=True

            self.people.append(p)


    def graphics(self):

        self.instantiate()
        
        #this creates all the graphics
        fig = plt.figure(figsize=(18,9))
        ax = fig.add_subplot(1,2,1)
        cx = fig.add_subplot(1,2,2)
        ax.axis('off')
        cx.axis([0,1000,0,self.n]) #(Xmin,Xmax,Ymin,Ymax)
        scatt=ax.scatter([p.posx for p in self.people],
                [p.posy for p in self.people],c='blue',s=8) 
        box = plt.Rectangle((0,0),100,100,fill=False)
        ax.add_patch(box)
        cvst,=cx.plot(self.infected,color="red",label="Currently infected")
        rvst,=cx.plot(self.recovered,color="gray",label="Recovered")
        dvst,=cx.plot(self.dead,color="black", label="Dead")
        cx.legend(handles=[rvst,cvst,dvst])
        cx.set_xlabel("Time in days")
        cx.set_ylabel("People")


        ct=[self.infected]
        rt=[self.recovered]
        t=[0]
        dt=[self.dead]
        
        return scatt,cvst,ct,rvst,rt,dvst,dt,t,fig

    #function excecuted frame by frame
    #frame : time variable
    def update(self,frame,rt,ct,dt,t):
        
        concycle = 0 #susceptible
        recycle = 0  #recovered
        dead = 0  
        colors = []
        sizes = [8 for p in self.people]
        for p in self.people:
            #check how much time the person has been sick, such that if total time passes, it heals them
            p.check_contagious(frame)
            #check if person is dead, if so, fix them in place
            p.check_death()
            #animate the movement of each person
            p.update_pos(0,0)
            if p.retired:
                recycle+=1 #count the amount of recovered
                for per in self.people:
                #check for people around the sick individual and infect the ones within the
                #transmission radius given the probability
                    if per.index==p.index or per.infected or per.die:
                        pass
                    else:
                        d=p.get_dist(per.posx,per.posy)
                        if d< self.r_contagious:
                            if np.random.random() < self.reinfect_rate / 100:
                                per.infect(frame)
                                sizes[per.index]=40
                                p.fix_position()
            if p.die:
                dead+=1 #count the amount of dead people
            if p.infected:
                concycle=concycle+1 #count the amount of infected
                if p.fixed:
                    pass
                elif np.random.random() < self.d_rate/100:
                    p.kill()


                for per in self.people:
                #check for people around the sick individual and infect the ones within the
                #transmission radius given the probability
                    if per.index==p.index or per.infected or per.retired or per.die:
                        pass
                    else:
                        d=p.get_dist(per.posx,per.posy)
                        if d<self.r_contagious:
                            if np.random.random() < self.p_contagious / 100:
                                per.infect(frame)
                                sizes[per.index]=80
            
            colors.append(p.get_color()) #change dot color according to the person's status

        #update the plotting data
        ct.append(concycle)
        rt.append(recycle)
        dt.append(dead)
        t.append(frame)


        #transfer the data to the matplotlib graphics
        offsets=np.array([[p.posx for p in self.people],
                        [p.posy for p in self.people]])
        scatt.set_offsets(np.ndarray.transpose(offsets))
        scatt.set_color(colors)
        scatt.set_sizes(sizes)
        cvst.set_data(t,ct)
        rvst.set_data(t,rt)
        dvst.set_data(t,dt)

        return scatt,cvst,rvst,dvst


#run the animation indefinitely
if __name__ == "__main__":
    sim = Simulation()
    scatt,cvst,ct,rvst,rt,dvst,dt,t,fig = sim.graphics()
    animation = FuncAnimation(fig, sim.update, interval=25, fargs=(rt,ct,dt,t), blit=True)
    plt.show()