from pylab import *
from SIRZ import *

def variable_graph():
	Ival = []
	Pfinal = []
	Zfinal = []
	for i in arange(0., 1., 0.01):
		Model = SIRZ(0.5, 0.2, 0.05, 0.01, 0.5, i)
		Ival.append(Model.I)
		traj = RK3DIntegrator(Model, 1., 500.)
		Pfinal.append(Model.P)
		Zfinal.append(Model.Z)
	
	figure(1)
	xlabel('Inital Infected Population')
	ylabel('Final Percent of Initial Population')
	plot(Ival, Pfinal, 'bo', label = 'Population (S+I+R)')
	plot(Ival, Zfinal, 'ko', label = 'Zombies')
	legend()
	show()
	#savefig("VARIABLE_PHI2", dpi = 600)
	
def SIR_Model():
	Model = SIRZ(2., 0.05, 0.2, 0., 0., 1., 0.0002)
	traj = RK3DIntegrator(Model, .1, 50.)
	
	figure(1)
	title('SIR Model')
	xlabel('Time (Days)')
	ylabel('Percent of Initial Population')
	plot(traj[5],traj[0], 'bo', label = 'Seceptible')
	plot(traj[5],traj[1], 'go', label = 'Infected')
	plot(traj[5],traj[2], 'ro', label = 'Recoverd')
	plot(traj[5],traj[3], 'ko', label = 'Dead')
	plot(traj[5],traj[4], 'b-', label = 'Total Population')
	legend()



def SIRZ_Model():
	Model = SIRZ(0.5, 0.2, 0.05, 0., 0., 0.0002)
	traj = RK3DIntegrator(Model, 1., 250.)
	
	figure(1)
	title('SIR Model + Fatalites')
	xlabel('Time (Days)')
	ylabel('Percent of Initial Population')
	plot(traj[5],traj[0], 'bo', label = 'Seceptible')
	plot(traj[5],traj[1], 'go', label = 'Infected')
	plot(traj[5],traj[2], 'ro', label = 'Recoverd')
	plot(traj[5],traj[3], 'ko', label = 'Killed')
	legend()

	show()