# ############################################################################
# Erin L Landguth
# UCDavis Physics 250: Modeling Chaos and Complexity
# Spring 2007 - Final Project
# Description: This program explores disease propagation by implementing the 
#	SIR model - 3D ODE System.
#
#	Graphs of S,I,R versus time steps are plotted.
# ############################################################################
#
# For a 3D ODE
#     dx/dt = f(x,y)
#     dy/dt = g(x,y)
# 	  dz/dt = h(x,y)
# See RK3DIntegrator() below for how the fourth-order Rungle-Kutta method 
#	integrates.
#
# The example used here is the following 3D ODE SIR model:
#	  dS/dt = -aSI
#	  dI/dt = aSI - bI
#	  dR/dt = bI
#	where S - susceptible, I - Infected, R - Removed
#	and a - infection/contact rate
#	    b - removal/recovery rate
#
#	Note: The fraction of the population will be viewed.  Thus N = S+I+R with
#		s=S/N, i=I/N, and r=R/N.  and then s+i+r = 1
#

# Import plotting routines
from pylab import *

# A SIR 3D ODE - return small s, i, and r.
def SDot(a,b,S,I,R):
	N=float(S+I+R)
	s=S/N
	i=I/N
	r=R/N
	return -a*s*i*N
	
def IDot(a,b,S,I,R):
	N=float(S+I+R)
	s=S/N
	i=I/N
	r=R/N
	return a*s*i*N-b*i
	
def RDot(a,b,S,I,R):
	N=float(S+I+R)
	s=S/N
	i=I/N
	r=R/N
	return b*i
	
# 3D Fourth-Order Runge-Kutta Integrator (for 4 parameters)
def RK3DIntegrator(a,b,x,y,z,f,g,h,dt):
	k1x = dt * f(a,b,x,y,z)
	k1y = dt * g(a,b,x,y,z)
	k1z = dt * h(a,b,x,y,z)
	k2x = dt * f(a,b,x + k1x / 2.0,y + k1y / 2.0,z + k1z / 2.0)
	k2y = dt * g(a,b,x + k1x / 2.0,y + k1y / 2.0,z + k1z / 2.0)
	k2z = dt * h(a,b,x + k1x / 2.0,y + k1y / 2.0,z + k1z / 2.0)
	k3x = dt * f(a,b,x + k2x / 2.0,y + k2y / 2.0,z + k2z / 2.0)
	k3y = dt * g(a,b,x + k2x / 2.0,y + k2y / 2.0,z + k2z / 2.0)
	k3z = dt * h(a,b,x + k2x / 2.0,y + k2y / 2.0,z + k2z / 2.0)
	k4x = dt * f(a,b,x + k3x,y + k3y,z + k3z)
	k4y = dt * g(a,b,x + k3x,y + k3y,z + k3z)
	k4z = dt * h(a,b,x + k3x,y + k3y,z + k3z)
	x = x + ( k1x + 2.0 * k2x + 2.0 * k3x + k4x ) / 6.0
	y = y + ( k1y + 2.0 * k2y + 2.0 * k3y + k4y ) / 6.0
	z = z + ( k1z + 2.0 * k2z + 2.0 * k3z + k4z ) / 6.0
	return x,y,z

# Simulation parameters
# Integration time step
dt = 0.01
#
# Control parameter for the SIR 3D ODE:
a = 0.5			# infection rate
b =	0.3			# death rate due to other causes
#
# Set up arrays of iterates for several different initial conditions
x = [16220.]
y = [164.]
z = [0.]
# But then rescale them
N=x[0]+y[0]+z[0]
x = [x[0]/N]
y = [y[0]/N]
z = [z[0]/N]

# Time
t  = [ 0.0]
# The number of time steps to integrate over
N = 10000

# The main loop that generates the orbit, storing the states
for n in range(0,N):
  # at each time step calculate new x(t) and y(t) and z(t)
  xtemp,ytemp,ztemp = RK3DIntegrator(a,b,x[n],y[n],z[n],SDot,IDot,RDot,dt)
  x.append(xtemp)
  y.append(ytemp)
  z.append(ztemp)
  t.append(t[n] + dt)

# Plot S,I,R vs. time
figure(1)
plot(t,x,'b',label='s(t)')
plot(t,y,'r',label='i(t)')
plot(t,z,'k',label='r(t)')
xlabel('time')
ylabel('populations')
title('The SIR Model for a = ' + str(a) + ' and b = ' + str(b))
legend()

# Use command below to save figure
#savefig('RK3D_SIR', dpi=600)

# Display the plot in a window
show()

# ############################################################################
# Erin L Landguth
# UCDavis Physics 250: Modeling Chaos and Complexity
# Spring 2007 - Final Project
# Description: This program explores disease propagation across space and
#	time by implementing Cellular Automata theory.
#
#	Each grid value transitions between 3 state values based on rules; 
#		S - Susceptible (2), I - Infected (1), R - Removed (0).
#
#	An animated plot is produced to watch the phase transitions.
#
#	Graphs of S,I,R versus time steps are plotted.
#
#	The user can choose which initial conditions to run.
# ############################################################################

# Import Modules
from numpy import *					# General commands and functions
from pylab import *					# Animated plots
from Scientific.Statistics import *	# Statistic calculations
from numpy.random import *				# Random/statistic calculations
from InfectedGrids import *			# Functions to initialize grids
from CARules import *					# CA Update function rules
from Numeric import *					# Well, just in case.
from Tkinter import *					# Need some buttons

# Some module checks
try:
	import pygame
	import pygame.surfarray as surfarray
except ImportError:
	raise ImportError, "PyGame required."
	
# ########################################################
# Create animated function - Spatial configuration display
# ########################################################
def SpDisplayState(state,N,M,surface,CellSize):
	ConfigRect = pygame.Rect(0,0,M*CellSize,N*CellSize)
	surface.fill(SColor,ConfigRect)
	for i in range(N):
		for j in range(M):
			if state[i][j] == 1:
				surface.fill(IColor,[j*CellSize,i*CellSize,CellSize,CellSize])
			if state[i][j] == 0:
				surface.fill(RColor,[j*CellSize,i*CellSize,CellSize,CellSize])
	
# Set up initial grid configuration - user defined.
print "SIR CA Model"
print "Which initial grid configuration do you want to use?"
print "1. Random Infected"
print "2. Center Infected"
print "3. Patchy Infected"
print "4. Your own text file"
# Error check for number 1 or 2 or 3 or 4
while True:
	number = raw_input('Enter the number: ')
	if (number=='1' or number=='2' or number=='3' or number=='4'):
		break
	elif number=='':
		number='1'
		break
	else:
		print "Invalid Entry.", number, "is not a choice."

# Call each initial grid function from InfectedGrids.py module
if (number=='1'):
	print "\n"
	print "Decide the grid dimensions:"
	N=raw_input("N = ")
	N=float(N)
	M=raw_input("M = ")
	M=float(M)
	print "How much % initial infected cells do you want?"
	intinfected=raw_input("Initial Infected = ")
	intinfected=float(intinfected)
	state = RandomInfected(N,M,intinfected)

elif (number=='2'):
	print "\n"
	print "Decide the grid dimensions:"
	N=raw_input("N = ")
	N=float(N)
	M=N
	print "How much % initial infected cells do you want?"
	intinfected=raw_input("Initial Infected = ")
	intinfected=float(intinfected)
	state = CenterInfected(N,intinfected)
	
elif (number=='3'):
	print "\n"
	print "Decide the grid dimensions:"
	N=raw_input("N = ")
	N=float(N)
	M=N
	print "How much % initial infected cells do you want?"
	intinfected=raw_input("Initial Infected = ")
	intinfected=float(intinfected)
	print "How many patches do you want?"
	patchnumber=raw_input("Number of Patches = ")
	patchnumber=float(patchnumber)	
	state = PatchyInfected(N,intinfected,patchnumber)	

else:
	print "\n"
	print "What is the name of your file?"
	print "Make sure file is in same directory as this program."
	print "Make sure to add .txt to file name."
	file=raw_input("File Name = ")
	state,M,N = UserInfected(file)

# Ask user to define time for simulation.
print "\n"
print "How long do you want simulation to run?"
Time = raw_input("Total Time = ")
Time = int(Time)

# Count initial 2's, 1's, and 0's to plot later.
intScount = 0
intIcount = 0
intRcount = 0
for i in range(N):
	for j in range(M):
		if state[i][j]==2:
			intScount = intScount + 1
		if state[i][j]==1:
			intIcount = intIcount + 1
		if state[i][j]==0:
			intRcount = intRcount + 1
# And then put these into an array to append to later
S = [float(intScount)]
I = [float(intIcount)]
R = [float(intRcount)]

# Prompt user for which kind of rules to use
print "\n"
print "Which update rules do you want to use?"
print "1. Fixed Probability"
print "2. Probability based on neighbors"
# Error check for number 1 or 2
while True:
	rulenumber = raw_input('Enter the number: ')
	if (rulenumber=='1' or rulenumber=='2'):
		break
	else:
		print "Invalid Entry.", rulenumber, "is not a choice."

# Call rule update function from CARules.py
if (rulenumber=='1'):
	print "\n"
	print "What is the infection rate?"
	a=raw_input("a = ")
	a=float(a)
	print "\n"
	print "What is the recovery rate?"
	b=raw_input("b = ")
	b=float(b)
	
else:
	print "\n"
	print "What is the infectivity of the disease?"
	K=raw_input("K = ")
	print "\n"
	print "What is the recovery rate?"
	b=raw_input("b = ")
	b=float(b)
	
# Set defaults for the screen display mode
CellSize = 4
size = (CellSize*M,CellSize*N)
pygame.init()
SColor = 0, 0, 255			# Blue
IColor = 255, 0, 0			# Red
RColor = 0, 0, 0 			# Black

# Get display surface
screen = pygame.display.set_mode(size)
pygame.display.set_caption('2D SIR Cellular Automaton Simulator')
SpDisplayState(state,N,M,screen,CellSize)
pygame.display.flip()
# Create RGB array whose elements refer to screen pixels
# Strangeness: sptmdiag is no longer used, but if this line
#	is removed, then display updates slow down considerably!
sptmdiag = surfarray.pixels3d(screen)

# Capture Initial configuration for saving figure uncomment
matshow(state)
#title('Random Infected Initial Configuration')
#title('Center Infected Initial Configuration')
title('Random Patch Infected Initial Configuration')
show()

# ######################
# Iterate Update process
# ######################
# count 2's, 1's, and 0's at each time step
# Also display results with pygame
def GO():
	global state
	for i in range(Time):

		# In CAFixedUpdate function call
		if (rulenumber=='1'):
		
			# Not quite sure how this works
			for event in pygame.event.get():
				if event.type == pygame.QUIT: sys.exit()
		
			# Iterate state
			state = CAFixedUpdate(state,a,b)
		
			# Display state
			SpDisplayState(state,N,M,screen,CellSize)
			pygame.display.flip()
		
			# Count initialize
			Scount = 0
			Icount = 0
			Rcount = 0
		
			# Step through new matrix counting up 2's, 1's, and 0's
			for i in range(N):
				for j in range(M):
					if state[i][j]==2:
						Scount = Scount + 1
					if state[i][j]==1:
						Icount = Icount + 1
					if state[i][j]==0:
						Rcount = Rcount + 1
					
			# Append the new counts to S, I, R arrays
			S.append(float(Scount))
			I.append(float(Icount))
			R.append(float(Rcount))
		
		# In CAStochasticUpdate function call
		if (rulenumber=='2'):
	
			# Not quite sure how this works
			for event in pygame.event.get():
				if event.type == pygame.QUIT: sys.exit()
		
			# Iterate state
			state = CAStochasticUpdate(state,K,b)
		
			# Display state
			SpDisplayState(state,N,M,screen,CellSize)
			pygame.display.flip()
		
			# Count initialize
			Scount = 0
			Icount = 0
			Rcount = 0
		
			# Step through new matrix counting up 2's, 1's, and 0's
			for i in range(N):
				for j in range(M):
					if state[i][j]==2:
						Scount = Scount + 1
					if state[i][j]==1:
						Icount = Icount + 1
					if state[i][j]==0:
						Rcount = Rcount + 1
					
			# Append the new counts to S, I, R array
			S.append(float(Scount))
			I.append(float(Icount))
			R.append(float(Rcount))		

# #######
# Button!
# #######
root = Tk()
bGO = Button(root,text='GO',command=GO)
bGO.pack(side="right")
root.mainloop()

# #######################################
# Plotting Commands for S,I,R versus Time
# #######################################

# Create time array
time = arange(0,Time+1,1)

# Rescale S,I,and R 
TotalPopulation = N*M
SNorm=[]
INorm=[]
RNorm=[]
for i in range(len(S)):
	SNorm.append(S[i]/TotalPopulation)
	INorm.append(I[i]/TotalPopulation)
	RNorm.append(R[i]/TotalPopulation)
	
# Plot S,I,R vs. time
figure(1)
plot(time,SNorm,'b',label='s(t)')
plot(time,INorm,'r',label='i(t)')
plot(time,RNorm,'k',label='r(t)')
xlabel('time')
ylabel('populations')
if (rulenumber=='1'):
	title('The SIR CA Model with fixed probabilities a = ' + str(a) + ' and b = ' + str(b))
if (rulenumber=='2'):
	title('The SIR CA Model with neighborhood influence for K = ' + str(K) + ' and b = ' + str(b))
legend()

# Display the plot in a window 
show()

# ####################
# Some stats print out
# ####################
print "The maximum infected individuals = ",max(I)
print "At time step = "
for i in range(len(I)):
	maxi=max(I)
	if(I[i]==maxi):
		print i
		
# Capture End configuration for saving figure uncomment
matshow(state,cmap=cm.hot)
#title('Random Infected after 10 Time Steps')
#title('Center Infected after 10 Time Steps')
title('Random Patch after 10 Time Steps')
show()

# ############################################################################
# Erin L Landguth
# UCDavis Physics 250: Modeling Chaos and Complexity
# Spring 2007 - Final Project
# Description: This module includes the functions to define the CA update rules.
#	1. CAFixedUpdate: a fixed probability to transition between each state
#		a - infection rate
#		b - recovery rate
# 	2. CAStochasticUpdate: looking around the neighborhood and updating
#		K - infectivity of disease
#		b - recovery rate
# Create state transition rules.  Create in functions defined above.
#		1. S -> I
#			a. with fixed probability a=0.5 
#			b. with probability (1 - e^-Kv)
#				K = infectivity of the disease is [0,1] float K = 1.0
#				v = count of I cells around S. [0,8] integers
#		2. I -> R
#			with fixed cure probability (recovery rate) b=0.3
# ############################################################################

# Import Modules
from numpy import *					# General commands and functions
from Scientific.Statistics import *	# Statistic calculations
from numpy.random import *				# Random/statistic calculations
from Numeric import *					# Well, just in case.

# Update matrix based on fixed probabilities.
def CAFixedUpdate(initialmatrix,a,b):

	# Make sure a and b are float numbers
	a = float(a)
	b = float(b)
	
	# Loop through initial state matrix
	for i in range(shape(initialmatrix)[0]):
		for j in range(shape(initialmatrix)[1]):
		
			# Generate random number
			arandomnumber = rand()
			brandomnumber = rand()
		
			# Change each 2 position to 1 if arandomnumber < a (S -> I)
			if (initialmatrix[i][j] == 2 and arandomnumber <= a):
				initialmatrix[i][j] = 1
				
			# Change each 1 position to 0 if brandomnumber < b (I -> R)
			elif (initialmatrix[i][j] == 1 and brandomnumber <= b):
				initialmatrix[i][j] = 0
				
			# Else leave the same (Removed population)
			else:
				initialmatrix[i][j] = initialmatrix[i][j]
	
	# Return next state matrix
	return initialmatrix

# Update matrix based on stachastic elements.
def CAStochasticUpdate(initialmatrix,K,b):
	
	# Make sure K and b are float numbers
	K = float(K)
	b = float(b)
	
	#Extract Dimensions
	N = shape(initialmatrix)[0]
	M = shape(initialmatrix)[1]
	
	# Loop through initial state matrix
	for i in range(N):
		for j in range(M):
			
			# Generage random numbers
			brandomnumber = rand()
			toIrandomnumber = rand()
			
			# Count up neighborhood I cells around S cell
			#	Create spatially homogenous lattice - use modulus.
			Icount = 0
			if (initialmatrix[i][j] == 2):
				if initialmatrix[i%N][(j+1)%M] == 1:
					Icount = Icount + 1
				if initialmatrix[(i+1)%N][(j+1)%M] == 1:
					Icount = Icount + 1
				if initialmatrix[(i+1)%N][j%M] == 1:
					Icount = Icount + 1
				if initialmatrix[(i+1)%N][(j-1)%M] == 1:
					Icount = Icount + 1
				if initialmatrix[i%N][(j-1)%M] == 1:
					Icount = Icount + 1
				if initialmatrix[(i-1)%N][(j-1)%M] == 1:
					Icount = Icount + 1
				if initialmatrix[(i-1)%N][j%M] == 1:
					Icount = Icount + 1
				if initialmatrix[(i-1)%N][(j+1)%M] == 1:
					Icount = Icount + 1
				
				# Generate probability
				toIprobability = float(1.0 - e**(-K*Icount))
				
				# And change each 2 position to 1 based on if < toIprobability.
				if (toIrandomnumber <= toIprobability):
					initialmatrix[i][j] = 1
				
			# Change each 1 position to 0 if brandomnumber < b (I -> R)
			elif (initialmatrix[i][j] == 1 and brandomnumber <= b):
				initialmatrix[i][j] = 0
			
			# Else leave matrix alone
			else:
				initialmatrix[i][j] = initialmatrix[i][j]
				
	# Return next matrix
	return initialmatrix

# ############################################################################
# Erin L Landguth
# UCDavis Physics 250: Modeling Chaos and Complexity
# Spring 2007 - Final Project
# Description: This module includes the functions to populate the grid space.
#	1. RandomInfected: NxM with % cells infected
# 	2. CenterInfected: % center cells of the NxN grid infected
#	3. PatchyInfected: NxN grid with random patches
#	4. User inputted file
# ############################################################################

# Import Modules
from numpy import *
from Scientific.Statistics import *
from numpy.random import *

# Let's try NxM grid with user% Infected, the rest Susceptible, and 0% Removed
# 	Create as a function and randomly place user% Infected cells.
#	Return the NxM array.  Note, this does not have to be a square
def RandomInfected(N,M,initialinfect):

	# Turn initialinfect into a float number
	initialinfect = float(initialinfect)
	
	# Create a zero matrix that is MxN dimensions
	initialstate = zeros((int(N),int(M)))
	
	# Turn all cells into Susceptible cells
	initialstate[:][:] = 2	
	
	# Randomly choose user specified initial infected number of cells	
	totalinfected = int((initialinfect/100)*N*M)
	randomxy = random((totalinfected,2))		
	
	# Fill initial state grid with infected cells
	for i in range(totalinfected):		
		initialstate[int(randomxy[i][0]*N)][int(randomxy[i][1]*M)] = 1
	
	# Return NxM array that is filled with 1s and 2s
	return initialstate
	
# Create an NxN grid with infected cells clustered in the middle - roughly
def CenterInfected(N,initialinfect):
	
	# Turn initialinfect into a float number
	initialinfect = float(initialinfect)
	
	# Create a zero matrix that is MxN dimensions
	initialstate = zeros((int(N),int(N)))
	
	# Turn all cells into Susceptible cells
	initialstate[:][:] = 2	
	
	# Find total cells infected
	totalinfected = int((initialinfect/100)*N**2)
	
	# Begin to center these cells in a cluster
	squareinfected = int(sqrt(totalinfected))
	centerbegin = (N-squareinfected)/2
	
	# Fill initial state grid with infected cells
	for i in range(squareinfected):
		for j in range(squareinfected):
			initialstate[centerbegin+i][centerbegin+j]=1
	
	# Return NxN array that is filled with 1s and 2s
	return initialstate
	
# Create and NxN grid with random patches
# Overlap allowed.
def PatchyInfected(N,initialinfect,patchnumber):
	
	# Turn initialinfect into a float number
	initialinfect = float(initialinfect)
	
	# Create a zero matrix that is MxN dimensions
	initialstate = zeros((int(N),int(N)))
	
	# Turn all cells into Susceptible cells
	initialstate[:][:] = 2	
	
	# Find total cells infected
	totalinfected = int((initialinfect/100)*N**2)
	
	# Divide total cells infected into number of patches
	patchsize = totalinfected/float(patchnumber)
	
	# Begin to cluster the patchsize
	squareinfected = int(sqrt(patchsize))
		
	# Randomly choose start location for each patch
	randomxy = random((patchnumber,2))	
	
	# Fill initial state grid with infected clusters
	for i in range(int(patchnumber)):
		for j in range(squareinfected):
			for k in range(squareinfected):
				initialstate[((int(randomxy[i][0]*N))+j)%N][((int(randomxy[i][1]*N))+k)%N] = 1
	
	# Return NxN array that is filled with 1s and 2s
	return initialstate
	
# Create function to read in user inputted file
#	Returns matrix, N, and M
def UserInfected(file):
	
	# Open file given for reading
	inputfile = open(file,'r')
	
	# Read lines from the file
	lines = inputfile.readlines()
	
	# Close the file
	inputfile.close()
	
	# Extract the column dimension
	for i in lines:
		thisline = i.split()
		col_dim = len(thisline)
	
	# Extract the row dimension
	row_dim = len(lines)
	
	# Create a matrix of zeros to be filled 	
	matrix = zeros((row_dim,col_dim),float)
	
	# Create an empty matrix to append to 
	x = []
	
	# Split up each line in file and append to empty matrix, x 
	for i in lines:
		thisline = i.split()
		x.append(thisline)
	
	# Turn x into and array that is shape row_dim x col_dim
	x = array(x)
		
	# Fill up matrix with float value of array x
	for j in range(row_dim):
		for k in range(col_dim):
			matrix[j][k] = float(x[j][k])	
	
	# Return the matrix
	return matrix, col_dim, row_dim
	
# Define general function that user can specify
def InitialGrid(number,N,M,initialinfect,file):
	if number == '1':
		return RandomInfected(N,M,initialinfect)
	elif number == '2':
		return CenterInfected(N,initialinfect)
	elif number == '3':
		return PatchyInfected(N,initialinfect,patchnumber)
	else:
		return UserInfected(file)