#Functions
from random import *
from pylab import plot,draw
from numpy import arange,sign,where
def ColorMap(person):		#Defines colors based on person's belief
	belief = person.belief
	if belief == 1:
		return 'y'
	if belief >= .75 and belief < 1:
		return 'r'
	elif belief < .75 and belief >= .5:
		return 'g'
	elif belief < .5 and belief > 0:
		return 'b'
	elif belief == 0:
		return 'k'
def initmap(people):		#Creats initial plot categories for each color
	xy=[[[],[]],[[],[]],[[],[]],[[],[]],[[],[]]]
	for i in people:
		color = ColorMap(i)
		if color == 'y':
			xy[0][0].append(i.position[0])
			xy[0][1].append(i.position[1])
		elif color == 'r':
			xy[1][0].append(i.position[0])
			xy[1][1].append(i.position[1])
		elif color == 'g':
			xy[2][0].append(i.position[0])
			xy[2][1].append(i.position[1])
		elif color == 'b':
			xy[3][0].append(i.position[0])
			xy[3][1].append(i.position[1])
		elif color == 'k':
			xy[4][0].append(i.position[0])
			xy[4][1].append(i.position[1])
	plottemp0, = plot(xy[0][0],xy[0][1],'yo')
	plottemp1, = plot(xy[1][0],xy[1][1],'rD')
	plottemp2, = plot(xy[2][0],xy[2][1],'gH')
	plottemp3, = plot(xy[3][0],xy[3][1],'bp')
	plottemp4, = plot(xy[4][0],xy[4][1],'ks')
	return plottemp0,plottemp1,plottemp2,plottemp3,plottemp4
def mapupdate(people):			#Updates positions to replace plot data
	xy=[[[],[]],[[],[]],[[],[]],[[],[]],[[],[]]]
	for i in people:
		color = ColorMap(i)
		if color == 'y':
			xy[0][0].append(i.position[0])
			xy[0][1].append(i.position[1])
		elif color == 'r':
			xy[1][0].append(i.position[0])
			xy[1][1].append(i.position[1])
		elif color == 'g':
			xy[2][0].append(i.position[0])
			xy[2][1].append(i.position[1])
		elif color == 'b':
			xy[3][0].append(i.position[0])
			xy[3][1].append(i.position[1])
		elif color == 'k':
			xy[4][0].append(i.position[0])
			xy[4][1].append(i.position[1])
	return xy

def positions(people):	#Returns postions of all people as array
	x = [i.position[0] for i in people]
	y = [i.position[1] for i in people]
	return x,y

def addbelief(belief1,belief2,percent):	#Adds belief between people and makes sure it never exceeds 1  Multiples by percent to allow variable transfer
	belief2 += (belief1 * percent)
	if belief2 > 1:
		belief2 = 1
	elif belief2 < 0:
		belief2 = 0
	return belief2

def newpos(position,max,n):		#Calculates new position for person, assuming person moves a random amount up to 1
	for i in range(2):
		position[i]+=choice(arange(-max[i]/n,max[i]/n,.0001*max[i]/n))
		if abs(position[i]) >= max[i]:
			position[i] = -position[i]
			position[i] -= 2*sign(position[i])*(abs(position[i])-max[i])
	return position

def phone(person1,person2,mindist):		#Calculates addbelief for phone tool
	if person2.iswitness():		#Only adds for one person
		return person2
	else:
		person2.belief = addbelief(person1.belief,person2.belief,.75)
		return person2

def radio(person,people,mindist):		#Calculates add belief for large radius around person
	j=0
	for i in people:
		if i.iswitness():
			continue
		else:
			if person.distance(i)<=3*mindist:
				people[j].belief = addbelief(person.belief,i.belief,.33)
		j+=1
	return people
	
def tv(person,people,mindist):		#Same, but even larger radius
	j=0
	for i in people:
		if i.iswitness():
			continue
		else:
			if person.distance(i)<=6*mindist:
				people[j].belief = addbelief(person.belief,i.belief,.2)
		j+=1
	return people

def internet(person,people,mindist):	#Even larger radius
	j=0
	for i in people:
		if i.iswitness():
			continue
		else:
			if person.distance(i)<=10*mindist:
				people[j].belief = addbelief(person.belief,i.belief,.1)
		j+=1
	return people

def toolprob(tools):	#Calculates which tool based on tool prob.
	tool = choice(arange(0.,1.01,.01))
	diff = []
	order = [1,2,0,3,4]
	for i in tools.values():
		diff.append(abs(tool-i))
	index = order[list(diff).index(min(diff))]
	index = tools.keys()[index]
	return tools[index],index

def whichtool(person):	#Returns which tool to be used or none for no tool
	phoneprob = .75
	radioprob = .33
	tvprob = .2
	internetprob = .1
	tools = {'none':1.,phone:phoneprob,radio:radioprob,tv:tvprob,internet:internetprob}
	return toolprob(tools)
def othercounter(othercount,index1,index2,index3,time):	#Counts uses for tools other than phone or none
	if othercount[index1+1][index2][index3]!=time:
		othercount[index1+1][index2].append(time)
		othercount[index1][index2].append(0)
	elif othercount[index1+1][index2][index3]==time:
		othercount[index1][index2][index3]+=1
	return othercount

def timeloop(dt,N,n,people,mindist,xmax,ymax,randomrange,totalcount,phonecount,othercount,time,plot0,plot1,plot2,plot3,plot4):
	for t in range(N):
		witnesscount = 0
		for i in range(n):
				for j in range(n):
					if people[j].iswitness() != 1:
						if j!=i:
							# tool,function = whichtool(people[i])	#Finds out what to use
							# if tool == 1.:
							if people[i].distance(people[j])<=mindist:	#Picks method of add belief depending on tool
								people[i].belief=addbelief(people[i].belief,people[j].belief,choice(randomrange)) # Allowing negative for possibility of disbelief
							# elif tool == .75:
								# k=i
								# while k==i:
									# k = randrange(0,n,1)
								# people[i] = function(people[i],people[k],mindist)
								# if phonecount[1][t-1]!=t*dt:
									# phonecount[1].append(t*dt)
									# phonecount[0].append(0)
								# elif phonecount[1][t-1]==t*dt:
									# phonecount[0][t-1]+=1
							# else:
								# if function != 'none':
									# people = function(people[i],people,mindist)
									# if function.__name__=='radio':
										# othercount = othercounter(othercount,0,0,t-1,t*dt)
									# if function.__name__=='tv':
										# othercount = othercounter(othercount,0,1,t-1,t*dt)
									# if function.__name__=='internet':
										# othercount = othercounter(othercount,0,2,t-1,t*dt)
													# Attempting to make it animated
		for i in range(n):
			people[i].position = newpos(people[i].position,[xmax,ymax],n)	#Calculates new positions for all the people and end time loop
			if people[i].belief == 1:
				witnesscount+=1
		xy = mapupdate(people)		#Wasn't able to turn off autoscale, but assume limits of area are at the max x and y
		plot0.set_data(xy[0])
		plot1.set_data(xy[1])
		plot2.set_data(xy[2])
		plot3.set_data(xy[3])
		plot4.set_data(xy[4])
		draw()
		totalcount.append(witnesscount)
		time.append((t+1)*dt)
		if witnesscount == len(people):
			print 'Full witnesses at t = %f'% ((t+1)*dt)
			break
	return time,totalcount,phonecount,othercount

# def whileloop(dt,n,people,xmax,ymax,mindist,randomrange,totalcount,phonecount,othercount,time,plot0,plot1,plot2,plot3,plot4):
	# witnesscount = 0
	# t = 0
	# while witnesscount != len(people):
		# witnesscount = 0
		# for i in range(n):
				# for j in range(n):
					# if people[j].iswitness() != 1:
						# if j!=i:
							# tool,function = whichtool(people[i])	#Finds out what to use
							# if tool == 1.:
								# if people[i].distance(people[j])<=mindist:	#Picks method of add belief depending on tool
									# people[i].belief=addbelief(people[i].belief,people[j].belief,choice(randomrange)) # Allowing negative for possibility of disbelief
							# elif tool == .75:
								# k=i
								# while k==i:
									# k = randrange(0,n,1)
								# people[i] = function(people[i],people[k],mindist)
								# if phonecount[1][t-1]!=t*dt:
									# phonecount[1].append(t*dt)
									# phonecount[0].append(0)
								# elif phonecount[1][t-1]==t*dt:
									# phonecount[0][t-1]+=1
							# else:
								# if function != 'none':
									# people = function(people[i],people,mindist)
									# if function.__name__=='radio':
										# othercount = othercounter(othercount,0,0,t-1,t*dt)
									# if function.__name__=='tv':
										# othercount = othercounter(othercount,0,1,t-1,t*dt)
									# if function.__name__=='internet':
										# othercount = othercounter(othercount,0,2,t-1,t*dt)
													# Attempting to make it animated
		# for i in range(n):
			# people[i].position = newpos(people[i].position,[xmax,ymax],n)	#Calculates new positions for all the people and end time loop
			# if people[i].belief == 1:
				# witnesscount+=1
		# xy = mapupdate(people)		#Wasn't able to turn off autoscale, but assume limits of area are at the max x and y
		# plot0.set_data(xy[0])
		# plot1.set_data(xy[1])
		# plot2.set_data(xy[2])
		# plot3.set_data(xy[3])
		# plot4.set_data(xy[4])
		# draw()
		# totalcount.append(witnesscount)
		# time.append((t+1)*dt)
		# if witnesscount == len(people):
			# print 'Full witnesses at t = %f'% ((t+1)*dt)
			# break
		# t+=1
	# return time,totalcount,phonecount,othercount