from SIRZ import *
from pylab import *
from Dialog import Dialog
from Tkinter import *

class SIRZ:
	def __init__(self, B, g, a, phi1, phi2, I):
		self.B = B
		self.a = a
		self.g = g
		self.phi = phi1*phi2
		self.phi1 = phi1
		self.phi2 = phi2
		self.S = 1.
		self.I = I
		self.R = 0.
		self.Z = 0.
		self.P = self.S +self.I
		
	def dsdt(self):
		return -(self.B * self.S * self.I) - self.phi*(self.S/(self.R+self.S))*self.Z
	def didt(self):
		return self.B * self.S * self.I - self.I*(self.g + self.a)
	def drdt(self):
		return self.g*self.I - self.phi*(self.R/(self.R+self.S))*self.Z
	def dzdt(self):
		return self.a*self.I - self.Z*self.phi1*(1.-self.phi2)



class EntryField(Frame):

	def __init__(self, master, text, inital_value):
		Frame.__init__(self, master)
		Label(self, text=text).pack(side=LEFT)
		self.value = DoubleVar(value=inital_value)
		Entry(self,textvariable=self.value).pack(side=LEFT, fill=X)
	
	def get(self):
		try: 
			self.value.get()
		except ValueError:
			InputError(self)
		else: 
			return self.value.get()
		
	def report(self):
		print self.get()
		
class StartButton(Frame):
	def __init__(self, master, button_list):
		Frame.__init__(self, master, bd=2, relief=SUNKEN)
		for button, action in button_list:
			Button(self, text=button, command=action).pack(side=LEFT)

class InputError(Dialog):

	def __init__(self, master):
		Dialog.__init__(self, master, title = 'Input Error',
			text = 'Input must be Positive Numbers',
			bitmap = 'warning', default = 0,
			strings = ('Cancel',))

class MainWindow(Frame):

	def __init__(self, master):
		Frame.__init__(self, master)
		Label(self, text="").pack(side=TOP)
		self.Beta = EntryField(self, "Beta infection rate: ", 0.5)
		self.Beta.pack(side=TOP, fill=X)
		
		Label(self, text="").pack(side=TOP)
		self.Gamma = EntryField(self, "Gamma recovory rate: ", 0.2)
		self.Gamma.pack(side=TOP, fill=X)
		
		Label(self, text="").pack(side=TOP)
		self.Alpha = EntryField(self, "Alpha death rate: ", 0.05)
		self.Alpha.pack(side=TOP, fill=X)
		
		Label(self, text="").pack(side=TOP)
		self.Phi1 = EntryField(self, "Phi(1) encounter rate: ", 0.01)
		self.Phi1.pack(side=TOP, fill=X)
		
		Label(self, text="").pack(side=TOP)
		self.Phi2 = EntryField(self, "Phi(2) zombie strength: ", 0.5)
		self.Phi2.pack(side=TOP, fill=X)

		Label(self, text="").pack(side=TOP)
		self.I = EntryField(self, "Inital Infected (Percent): ",0.0002)
		self.I.pack(side=TOP, fill=X)
		
		Label(self, text="").pack(side=TOP)
		self.Days = EntryField(self, "Time (Days): ",250)
		self.Days.pack(side=TOP, fill=X)

		StartButton(self, [('Graph', self.graph)]).pack(side=BOTTOM, fill=X)
			
	def graph(self):
		B = self.Beta.get()
		g = self.Gamma.get()
		a = self.Alpha.get()
		phi1 = self.Phi1.get()
		phi2 = self.Phi2.get()
		I = self.I.get()
		Days = self.Days.get()

		Model = SIRZ(B, g, a, phi1, phi2, I)
		traj = RK3DIntegrator(Model, 1., Days)
	
		figure(1)
		title('SIRZ 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 = 'Zombies')
		plot(traj[5],traj[4], 'b-', label = 'Total Population')
		legend()
		show()

mw = MainWindow(None)
mw.pack()
mw.mainloop()