#   Project: Socail Sensitivity Model
#   File name: SensitivityModel.py
#   Description:  Model of social sensitivity and migration
#     
#   @author Sean Roberts: sean.g.roberts@gmail.com
#  
#     
#   @see The GNU Public License (GPL)
#  /
# / 
#   This program is free software; you can redistribute it and/or modify 
#   it under the terms of the GNU General Public License as published by 
#   the Free Software Foundation; either version 2 of the License, or 
#   (at your option) any later version.
#   
#   This program is distributed in the hope that it will be useful, but 
#   WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY 
#   or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License 
#   for more details.
#   
#   You should have received a copy of the GNU General Public License along 
#   with this program; if not, see <http://www.gnu.org/licenses/>.
#  /

# Explanation at http://wp.me/pZBGv-zD
# uses ffmpeg - http://ffmpeg.org/ 
# and matplotlib - http://matplotlib.sourceforge.net/
# it assumes there is a folder called 'images' in the same directory as this file (to store temporary images)

# there are three classes:
# Agent - an agent
# Area - a local area that contains many agents has a certain amount of resources
# World - a collection of Areas 

import math
import random
from matplotlib import pyplot as P
import numpy as np
import itertools

class Agent:
	def __init__(self,g):
		self.sensitive = g
		
	
class Area:
	def __init__(self):
		self.agents = []
		self.resource = 0.0
	def addAgent(self,a):
		self.agents.append(a)
	def removeAgent(self,a):
		try:
			self.remove(a)
		except:
			pass
	def removeRandomAgent(self):
		del self.agents[random.randint(0,len(self.agents))]
	def removeAgentProb(self,probSensitive,probNonSensitive,support,deathThreshold,maxCapacity,resourceThresh):
		# remove an agent based on their probability of 
		if len(self.agents)>0:
			thresholds = self.getThresholds(probSensitive,probNonSensitive,support,resourceThresh)
	#		print thresholds
			while len([x for x in thresholds if x < deathThreshold])>0 or len(self.agents)>maxCapacity:
				# need to remove agents
				# remove one at with probability proportionate to their thresholds
				if len(self.agents)<maxCapacity:
					print "death"
				Xthresh = [1.0/x for x in thresholds] # high score is less probable to be picked
				r = random.uniform(0,sum(Xthresh))
				total = 0
				searching = True
				for i in range(len(Xthresh)):
					total += Xthresh[i]
					if r < total and searching:
						del self.agents[i]
						searching = False
				# re-evaluate thresholds
				thresholds = self.getThresholds(probSensitive,probNonSensitive,support,resourceThresh)
				#print thresholds
	#		print "---"


	def getThresholds(self,probSensitive,probNonSensitive,support,resourceThresh):
		resourceShare = self.resource
		if self.resource>0 & len(self.agents)>0:
			resourceShare = self.resource/float(len(self.agents))
		if resourceShare > resourceThresh:
		# plenty of food, everyone has the same coping level
			thresholds = [probNonSensitive for x in self.agents]
		else:
			# when things are tough, difference comes out.
			thresholds = []
			for x in self.agents:
				if x.sensitive:
					thresholds.append(probSensitive)
				else:
					thresholds.append(probNonSensitive)
		# threshold is likelihood of surviving given gene, resources and support
		thresholds = [(x * resourceShare * float(support)) for x in thresholds]
		return thresholds
		
		
	def setResource(self,x):
		self.resource = x

	def mate(self,driftProb):
		a = []
		for x in self.agents:
			a.append(x)
		pairs = zip(a[:len(a)/2],a[len(a)/2:])
		for p in pairs:
			inherit = random.choice([p[0].sensitive,p[1].sensitive])
			if random.random()<driftProb:
				inherit =  not inherit
			self.agents.append(Agent(inherit))

class World:
	def __init__(self,size,StartAgent):
		self.map = [[]]
		self.size = size
		for x in range(size):
			self.map.append([])
			for y in range(size):
				self.map[x].append(Area())
		for x in range(StartAgent):
#			a = Agent(x % 2==0)  # start with 50/50 mix
#			a = Agent(random.choice([True,False])) # start with random mix
			a = Agent(True) 			# start with sensitive agents
#			a = Agent(False) 			# start with Non-sensitive agents
#			self.map[size-2][size-2].addAgent(a)  # start in high-resource
#			self.map[1][1].addAgent(a) # start in low-resource
			self.map[size-5][(size/2)].addAgent(a) 
			self.map[size-5-1][(size/2)].addAgent(Agent(True)) 
			self.map[size-5][(size/2)-1].addAgent(Agent(True)) 
			self.map[size-5][(size/2)+1].addAgent(Agent(True))
			
#			self.map[(size/2)+1][(size/2)+1].addAgent(Agent(x % 2==0)) # start in middle
#			self.map[(size/2)-1][(size/2)+1].addAgent(Agent(x % 2==0)) # start in middle
#			self.map[(size/2)+1][(size/2)-1].addAgent(Agent(x % 2==0)) # start in middle
			
	def death(self,maxCapacity,probS,probNonS,resourceThresh):
		for x in range(self.size):
			for y in range(self.size):
				res = self.resources(x,y)
				support = self.getSupport(x,y,maxCapacity)
				self.map[x][y].setResource(res)
				self.map[x][y].removeAgentProb(probS,probNonS,support,deathThreshold,maxCapacity,resourceThresh)
				
				
# 				while len(self.map[x][y].agents) > float(maxCapacity)*res:
# 					#print "DEATH\t",x,"\t",y,"\t",res,"\t",support
# 					self.map[x][y].removeAgentProb(probS,probNonS,support,deathThreshold)
	def getSupport(self,x,y,maxCapacity):
		# number of people in surrounding area
		numOthers = 0
		for i in itertools.product([1,0,-1],repeat=2):
			try:
				numOthers += len(self.map[x+i[0]][y+i[1]].agents)
			except:
				pass
 		ret = float(numOthers) / float(maxCapacity * 9)
 		if ret >0.99:
 			ret = 0.99
 		return ret
#		return numOthers	

	def mate(self,driftProb):
		for x in range(self.size):
			for y in range(self.size):
				self.map[x][y].mate(driftProb)
	def migrate(self,migProb,maxMigDist):
		for x in range(self.size):
			for y in range(self.size):
				for a in self.map[x][y].agents:
					acopy = a
					if random.random()<migProb:
						# migrate!
						self.map[x][y].removeAgent(a)
						# choose new destination by brute force
						newY = y+random.choice([1,0,-1,2,-2])
						newX = x+random.choice([1,0,-1,2,-2])
						# check the new destination is inside the box
						while newX<0 or newX>=self.size or newY<0 or newY>=self.size:
							newY = y+random.choice([i-maxMigDist for i in range((2*maxMigDist)+1)])
							newX = x+random.choice([i-maxMigDist for i in range((2*maxMigDist)+1)])
					#	print "MIGRATE ",x,y,newX,newY
						self.map[newX][newY].addAgent(acopy)
	
	def resources(self, x,y):
		# random resources
		#return 0.7 + (0.3*random.random())
		
		# an area's resources are based on latitude and longitude (x,y)	
		return math.pow((float(x+1)/float(self.size)),2) * math.pow((float(y+1)/float(self.size)),2)

###################
# Model Parameters  
###################
size = 60  # length of x axis
world = World(size,10)

maxCapacity = 10 # capacity of single area
years = 200# number of years to run

driftProb = 0.01 # chance of mutation
migProb = 0.05 # chance of migration
maxMigrationDistance = 1 # maximum distance an agent can migrate

sensitiveThresh = 0.6 # probability sensitive agent will cope under stress
nonSensitiveThresh = 0.9 #  probability non-sensitive agent will cope under stress, 
							#must be <1


resourceThresh = 0.5/maxCapacity
deathThreshold = 0.003  # this has to be fiddled a bit!

ticks = 2  # report every ticks years
tickNumber = 0 

# Run the model!
for year in range(years):
#	if year == 10:
#		deathThreshold = 0.01 # raise deathThreshold

	# death
	world.death(maxCapacity,sensitiveThresh,nonSensitiveThresh,resourceThresh)
	# mate
	world.mate(driftProb)			
	# migrate
	world.migrate(migProb,maxMigrationDistance)
	sLevels = []
	for x in range(size):
#		print "\n",
		for y in range(size):
			slevel = -1
			if len(world.map[x][y].agents)>0:
				# slevel is percentage of area that is sensitive
				slevel = float(sum([aa.sensitive for aa in world.map[x][y].agents]))/float(len(world.map[x][y].agents))
				# slevel is num of senstive / maxCapacity
				# this ensures small numbers in a small pop don't have a big effect
				#slevel = float(sum([aa.sensitive for aa in world.map[x][y].agents]))/float(len(world.map[x][y].agents))
			#print year,x,y,slevel
#			print slevel,"\t",
			sLevels.append(slevel)
#	print "\n\n\n"
	if year % ticks == 0:
		print year,sum([sum([len(y.agents) for y in x]) for x in world.map])
		data = np.array(sLevels).reshape(size,size)
		P.clf() # clear the figure so we're not redrawing on top
		# levels so non-populated areas are white, 
		# red = sensitive
		# blue = non-sensitive
		# green = mix
		cs = P.contourf(data,levels = [0,0.33,0.66,1])
		fname = 'images/tmp%04d'%tickNumber
		tickNumber +=1
		P.savefig(fname, dpi=100)
#	P.show()	

import os
os.system('ffmpeg -f image2 -i images/tmp%04d.png video.mpg')
os.system('rm images/*.png')