from classes import *
from functions import *
from pylab import ion,plot,axis,ioff,draw,figlegend,hold,close,figure,show,title,xlabel,ylabel
"""
Bad News Propagation Simulator

Considers an initial full witness at origin and sees what happens as witness contacts others.
Currently only 5 pathways for propagation.
Currently each with set probabilities.
Currently unstable.  Plots crash if selected.
Very slow.  Essentially won't work for > 100 people
"""
if __name__== '__main__':
	seed()	#Seed random number generator
	ion()	#Turn on interactive mode
	ans = ' '
	while ans != 'y' and ans != 'n':		#Initializing variables
		ans = raw_input('Do you wish to enter initial conditions (y or n)? ')
	if ans == 'y':
		# while ans != '1' and ans != '2':
			# ans = raw_input('Do you want to see how long until the system has full witnesses (1)\n \
# or run for a specific amount of time (2) (Enter 1 or 2)?\n \
# Warning:  1st option may take a long, long time:  ')		#Was for while loop which loops until system is filled with full witnesses
		# if ans == '1':									#Doesn't work.
			# t = 0					#N is number of time steps
			# dt = .1
		# else:
		t =	float(input('How long do you wish to run this simulation? '))
		dt = float(input('What size time step? '))	#n is number of people
		N = int(t/dt)
		n = int(input('How many people in this system? '))		
		xmax = float(input('What is your max x value? '))
		ymax = float(input('What is your max y value? '))
		belief = -1
		while belief < 0 or belief > 1:
			belief = float(input('Enter initial witness belief (number in [0,1]: '))
	elif ans == 'n':
		print 'Using default values.  Runs for specified time'
		t =	0
		dt = .1
		n = 100		#Initializing variables
		N = 600	#n is number of people
				#N is number of time steps
		xmax = 10.
		ymax = 10.
		belief = 1.
		print 't =', dt*N
		print '# of people =', n
		print 'Max x =', xmax
		print 'Max y =', ymax
		print 'Initial witness belief =', belief
	xstep = 2*xmax/10000
	ystep = 2*ymax/10000
	xran = arange(-xmax,xmax+xstep,xstep)	#Max and min x and y
	yran = arange(-ymax,ymax+ystep,ystep)	#Maximum area is square
	people = [person(choice(xran),choice(yran),0.) for i in range(n)]	#Makes list of person with random positions
	witnesscount = 0
	count = 0
	totalcount = []
	for i in range(n):
		if (people[i].position == array([0.,0.])).all():
			people[i].belief = belief #Makes people at origin the initial witnesses
			count+=1
		if people[i].belief == 1:
			witnesscount+=1
	if count == 0:
		for i in range(3):
			people[i] = person(0.,0.,belief)  #Makes specific number of initial witnesses if not at origin found.
	totalcount.append(witnesscount)
	mindist = sqrt(xmax**2+ymax**2)/n		#Defines min distance
	figure(1,(10,10))
	plot0,plot1,plot2,plot3,plot4 = initmap(people)
	axis([-xmax,xmax,-ymax,ymax])
	figlegend((plot0,plot1,plot2,plot3,plot4),('= 1','>= .75, < 1','< .75, >= .5','< .5, > 0','== 0'),'upper right')
	title('%i People spreading the news over time'%n)
	xlabel('x')
	ylabel('y')
	time = [0]
	phonecount = [[0],[0]]
	othercount = [[[0],[0],[0]],[[0],[0],[0]]]
	randomrange = concatenate((arange(-1.,-.499,.001),arange(.5,1.001,.001)))
	# Time/while loop start
	# if ans != '1':
	time,totalcount,phonecount,othercount = timeloop(dt,N,n,people,mindist,xmax,ymax,randomrange,totalcount,phonecount,othercount,time,plot0,plot1,plot2,plot3,plot4)
	# else:
		# time,totalcount,phonecount,othercount = whileloop(dt,n,people,mindist,xmax,ymax,randomrange,totalcount,phonecount,othercount,time,plot0,plot1,plot2,plot3,plot4)
	ioff()
	figure(2)
	hold(True)
	plot(time,totalcount,'c--')
	title('# of Full witnesses over time')
	xlabel('time')
	ylabel('# of full witnesses')
	# figure(3)
	# plot(phonecount[1],phonecount[0],'k--')
	# title('# of phones uses over time')
	# xlabel('time')
	# ylabel('# of phone uses')
	# figure(4)
	# plot(othercount[1][0],othercount[0][0],'r--')
	# title('# of radio uses over time')
	# xlabel('time')
	# ylabel('# of radio uses')
	# figure(5)
	# plot(othercount[1][1],othercount[0][1],'g--')
	# title('# of tv uses over time')
	# xlabel('time')
	# ylabel('# of tv uses')
	# figure(6)
	# plot(othercount[1][2],othercount[0][2],'b--')
	# title('# of internet uses over time')
	# xlabel('time')
	# ylabel('# of internet uses')
	show()