import numpy as np
import sys
import os


def read_AvgNtargets(filename,N_episodes):
	xs = np.zeros((N_episodes))
	ys = np.zeros((N_episodes))
	f = open(filename,'r')
	lines = f.readlines()
	f.close()
	k = 0
	for line in lines:
		p=line.split()
		xs[k] = int(p[0])
		ys[k] = float(p[1])
		k+=1
	return [xs,ys]
	
def read_Ntargets(filename,N_runs,N_episodes):
	xs = np.zeros((N_episodes,N_runs))
	f = open(filename,'r')
	lines = f.readlines()
	f.close()
	for line in lines:
		p=line.split()
		k = int(p[0])
		ep = int(p[1])
		xs[ep][k] = float(p[2])
	return xs
	
def read_Ntargets_direct(filename,N_runs,N_episodes):
	xs = np.zeros((N_runs,N_episodes))
	f = open(filename,'r')
	lines = f.readlines()
	f.close()
	for line in lines:
		p=line.split()
		k = int(p[0])
		ep = int(p[1])
		xs[k][ep] = float(p[2])
	return xs
	
def read_AvgPvalues(filename,dt,N_times):
	wts = np.zeros((N_times))
	p0s = np.zeros((N_times))
	p1s = np.zeros((N_times))
	f = open(filename,'r')
	lines = f.readlines()
	f.close()
	for line in lines:
		p=line.split()
		k = int(p[0])
		wts[k] = (1+k)*dt
		p0s[k] = float(p[1])
		p1s[k] = float(p[2])
	return [wts,p0s,p1s]

def compute_AvgPvalues_from_AvgH(filename,dt,N_times):
	Ns = 2*N_times
	Hs = np.zeros((Ns,2))
	f = open(filename,'r')
	lines = f.readlines()
	f.close()
	for line in lines:
		p=line.split()
		s = int(p[0])
		a = int(p[1])
		Hs[s][a] = float(p[2])
	
	wts = np.zeros((N_times))
	p0s = np.zeros((N_times))
	p1s = np.zeros((N_times))
	for t in range(N_times):
		wts[t] = (1+t)*dt
		p0s[t] = Hs[t][1]/(Hs[t][1]+Hs[t][0])
		p1s[t] = Hs[t+Ns/2][1]/(Hs[t+Ns/2][1]+Hs[t+Ns/2][0])
	return [wts,p0s,p1s]

def read_AvgPvalues_bunches(filename,dt,N_times,N_bunches,N_episodes):
	wts = np.zeros((N_times))
	p0s = np.zeros((N_times))
	p1s = np.zeros((N_times))
	for bunch in range(N_bunches):
		f = open(filename+str(bunch)+'_episode'+str(N_episodes)+'.dat','r')
		lines = f.readlines()
		f.close()
		for line in lines:
			p=line.split()
			k = int(p[0])
			wts[k] = (1+k)*dt
			p0s[k] += float(p[1])
			p1s[k] += float(p[2])
	p0s /= N_bunches
	p1s /= N_bunches
	return [wts,p0s,p1s]

def compute_AvgPvalues_from_AvgH_bunches(filename,dt,N_times,N_bunches,N_episodes):
	Ns = 2*N_times
	Hs = np.zeros((Ns,2))
	for bunch in range(N_bunches):
		f = open(filename+str(bunch)+'_episode'+str(N_episodes)+'.dat','r')
		lines = f.readlines()
		f.close()
		for line in lines:
			p=line.split()
			s = int(p[0])
			a = int(p[1])
			Hs[s][a] += float(p[2])
	Hs /= N_bunches
	
	wts = np.zeros((N_times))
	p0s = np.zeros((N_times))
	p1s = np.zeros((N_times))
	for t in range(N_times):
		wts[t] = (1+t)*dt
		p0s[t] = Hs[t][1]/(Hs[t][1]+Hs[t][0])
		p1s[t] = Hs[t+Ns/2][1]/(Hs[t+Ns/2][1]+Hs[t+Ns/2][0])
	return [wts,p0s,p1s]

def read_AvgPvaluesEpisode(filename,dt,N_times):
	AvgPvaluesEpisode = []
	f = open(filename,'r')
	lines = f.readlines()
	f.close()
	wts = np.zeros((N_times))
	p0s = np.zeros((N_times))
	p1s = np.zeros((N_times))
	ep0 = 0
	for line in lines:
		p=line.split()
		ep = int(p[0])
		if (ep > ep0):
			AvgPvaluesEpisode.append([wts,p0s,p1s])
			wts = np.zeros((N_times))
			p0s = np.zeros((N_times))
			p1s = np.zeros((N_times))
			ep0 = ep
		k = int(p[1])
		wts[k] = (1+k)*dt
		p0s[k] = float(p[2])
		p1s[k] = float(p[3])
	AvgPvaluesEpisode.append([wts,p0s,p1s])
	return AvgPvaluesEpisode
	
def read_AvgPvaluesEpisode_bunches(filename,dt,N_times,N_bunches):
	for bunch in range(N_bunches):
		AvgPvaluesEpisode = []
		f = open(filename+str(bunch)+'.dat','r')
		lines = f.readlines()
		f.close()
		wts = np.zeros((N_times))
		p0s = np.zeros((N_times))
		p1s = np.zeros((N_times))
		ep0 = 0
		for line in lines:
			p=line.split()
			ep = int(p[0])
			if (ep > ep0):
				AvgPvaluesEpisode.append([wts,p0s,p1s])
				wts = np.zeros((N_times))
				p0s = np.zeros((N_times))
				p1s = np.zeros((N_times))
				ep0 = ep
			k = int(p[1])
			wts[k] = (1+k)*dt
			p0s[k] = float(p[2])
			p1s[k] = float(p[3])
		AvgPvaluesEpisode.append([wts,p0s,p1s])
		
		if (bunch==0):
			AvgPvaluesEpisodeTOT = AvgPvaluesEpisode
		else:
			for l in range(len(AvgPvaluesEpisode)):
				AvgPvaluesEpisodeTOT[l][1] += AvgPvaluesEpisode[l][1]
				AvgPvaluesEpisodeTOT[l][2] += AvgPvaluesEpisode[l][2]
	for l in range(len(AvgPvaluesEpisodeTOT)):
		AvgPvaluesEpisodeTOT[l][1]/=N_bunches
		AvgPvaluesEpisodeTOT[l][2]/=N_bunches
	return AvgPvaluesEpisodeTOT
	
def read_AvgTargetTimes(filename,N_episodes):
	xs = np.zeros((N_episodes),dtype=int)
	ys = np.zeros((N_episodes))
	f = open(filename,'r')
	lines = f.readlines()
	f.close()
	k=-1
	for line in lines:
		k+=1
		p=line.split()
		xs[k] = int(p[0])
		ys[k] = float(p[1])
	return xs,ys
	
def read_AvgTargetTimes_bunches(filename,N_bunches,N_episodes):
	xs = np.zeros((N_episodes),dtype=int)
	ys = np.zeros((N_episodes))
	for bunch in range(N_bunches):
		f = open(filename+str(bunch)+'.dat','r')
		lines = f.readlines()
		f.close()
		k=-1
		for line in lines:
			k+=1
			p=line.split()
			xs[k] += int(p[0])
			ys[k] += float(p[1])
	xs /= N_bunches
	ys /= N_bunches
	return xs,ys
	
def read_TargetTimes(filename,N_runs,N_episodes):
	xs = np.zeros((N_episodes,N_runs))				# average time to find the target within a single episode. Organized in a good way for the box plot
	f = open(filename,'r')
	lines = f.readlines()
	f.close()
	for line in lines:
		p=line.split()
		k = int(p[0])
		ep = int(p[1])
		xs[ep][k] = float(p[2])
	return xs
	
def read_TargetTimes_bunches(filename,N_bunches,N_runs,N_episodes):
	xs = np.zeros((N_episodes,N_runs*N_bunches))				# average time to find the target within a single episode. Organized in a good way for the box plot
	for bunch in range(N_bunches):
		f = open(filename+str(bunch)+'.dat','r')
		lines = f.readlines()
		f.close()
		for line in lines:
			p=line.split()
			k = int(p[0])
			ep = int(p[1])
			xs[ep][k] = float(p[2])
	return xs
	
def read_parameters(filename):
	f=open(filename,'r')
	lines = f.readlines()
	f.close()
	#0 f.write('# units\n')
	#1 f.write('L = %11.8f (length)\n' % L)
	L=float(lines[1].split()[2])
	#2 f.write('tau = %11.8f (time)\n' % tau)
	tau=float(lines[2].split()[2])
	#3 f.write('# free parameters\n')
	#4 f.write('Pe = %11.8f (Peclet)\n' % Pe)
	Pe=float(lines[4].split()[2])
	#5 f.write('ell = %11.8f (persistence)\n' % ell)
	ell=float(lines[5].split()[2])
	#6 f.write('R = %11.8f (radius of the target)\n' % R_target)
	R_target=float(lines[6].split()[2])
	#7 f.write('# other parameters\n')
	#8 f.write('dt = %11.8f (integration time step)\n' % dt)
	dt=float(lines[8].split()[2])
	#9 f.write('# PS parameters\n')
	#10 f.write('N runs = %10d (number of independent runs)\n' % N_runs)
	N_runs=int(lines[10].split()[2])
	#11 f.write('N episodes = %6d (number of episodes for each run)\n' % N_episodes)
	N_episodes=int(lines[11].split()[2])
	#12 f.write('T = %11.8f (duration of a single episode)\n' % time_single_episode)
	time_single_episode=float(lines[12].split()[2])
	#13 f.write('PT = %11.8f (max duration of a phase)\n' % max_phase_duration)
	max_phase_duration=float(lines[13].split()[2])
	#14 f.write('rew = %11.8f (reward when a target is found)\n' % reward)
	reward=float(lines[14].split()[2])
	#15 f.write('gamma = %11.8f (gamma)\n' % gamma)
	eta=float(lines[15].split()[2])
	#16 f.write('eta = %11.8f (eta)\n' % eta)
	eta=float(lines[16].split()[2])
	f.close()
	return [L,tau,Pe,ell,R_target,dt,N_runs,N_episodes,time_single_episode,max_phase_duration,reward,gamma,eta]
	
def create_trajectory (Parameters,Pvalues,x_target,y_target):
	L = 1.0							# box size
	tau = 1.0						# typical time required by a passive particle to cross the box
	# ... it follows
	D = L*L/(4*tau)					# diffusion coefficient for the passive particle (kept the same also for the ABP phase)
	
	# other parameters
	dt = 0.0001						# time step
	R_target = Parameters[4]		# radius of the target (circular) (it must be such that the typical step size due to diffusion is smaller than the target)
	
	# ABP parameters
	Pe = Parameters[2]				# Peclet number = v tau / L
	v = Pe * L / tau				# self-propulsion velocity (it must be such that typically the particle do not meet the target)
	ell = Parameters[3]				# persistence length  = v/(D_theta * L)
	D_theta = v/ell					# rotational diffusion coefficient
	
	sigma = sqrt(2*D*dt)
	sigma_theta = sqrt(2*D_theta*dt)
	vdt = v*dt
	
	x = x_target + 0.95 * R_target * np.cos(2*np.pi*random_uniform())
	y = y_target + 0.95 * R_target * np.sin(2*np.pi*random_uniform())
	s = 0
	
	
	traj = []
	
	
	#~ a_time_max = 5
	#~ Ns=2
	
	#~ action_times = define_action_times(dt)
	#~ print 'action_times = ',action_times
	
	#~ # initial particle position
	#~ x = 0.2
	#~ y = 0.75
	#~ theta = np.pi * 4./5.
	
	
	#~ s = 0
	#~ #                            p   a   p   a   a   a   p   a   p   a   a   a   p   p  a     p   p   p   p   a   p    a   p   p
	#~ actions =       np.asarray([ 2,  9, 10,  9,  5,  5,  9,  9,  7, 10,  5,  5, 10,  5])  # from 1 to 10 
	#~ angles = np.pi*np.asarray([0.0,2.2,0.0,5.7,5.5,5.5,0.0,2.3,0.0,2.1,2.2,2.2,0.0,0.0])
	#~ seeds = np.asarray([         1,  3,  4,  1,  1,  7,  4,  3,  4,  4,  4,  4, 20, 20]  )
	
	#~ xstot=[]
	#~ ystot=[]
	#~ sstot=[]
	#~ xstot2=[]
	#~ ystot2=[]
	#~ xstot3=[]
	#~ ystot3=[]
	#~ xstot4=[]
	#~ ystot4=[]
	#~ xstot5=[]
	#~ ystot5=[]
	
	#~ for k in range(len(actions)):
		
		#~ if (k==17):
			#~ x = 0.45
			#~ y = 0.15
		
		#~ seed(seeds[k])
		#~ a = actions[k]-1
		#~ # apply action
		#~ a_time = a % a_time_max
		#~ a_type = a / a_time_max
		#~ s_prime = s
		#~ if (a_type == 1): s_prime = (s+1) % Ns
		#~ if (s_prime == 1): theta = angles[k]
		
		#~ xs=[x]
		#~ ys=[y]
		#~ ss=[s_prime]
		#~ xs2=[x-L]
		#~ ys2=[y]
		#~ xs5=[x+L]
		#~ ys5=[y]
		#~ xs3=[x-L]
		#~ ys3=[y+L]
		#~ xs4=[x]
		#~ ys4=[y+L]
		
		#~ print action_times[a_time]+1
		#~ for t in range(action_times[a_time]+1):
			#~ # time evolution depending on new particle type
			#~ if (s_prime == 0):
				#~ x += sigma * gauss(0, 1)
				#~ y += sigma * gauss(0, 1)
			#~ else:
				#~ x += vdt*np.cos(theta) + sigma * gauss(0, 1)
				#~ y += vdt*np.sin(theta) + sigma * gauss(0, 1)
				#~ theta += sigma_theta * gauss(0, 1)
			#~ xs.append(x)
			#~ ys.append(y)
			#~ ss.append(s_prime)
			#~ xs2.append(x-L)
			#~ ys2.append(y)
			#~ xs5.append(x+L)
			#~ ys5.append(y)
			#~ xs3.append(x-L)
			#~ ys3.append(y+L)
			#~ xs4.append(x)
			#~ ys4.append(y+L)
		#~ xstot.append(xs)
		#~ ystot.append(ys)
		#~ sstot.append(ss)
		#~ xstot2.append(xs2)
		#~ ystot2.append(ys2)
		#~ xstot3.append(xs3)
		#~ ystot3.append(ys3)
		#~ xstot4.append(xs4)
		#~ ystot4.append(ys4)
		#~ xstot5.append(xs5)
		#~ ystot5.append(ys5)
		#~ # update s
		#~ s = s_prime
		
		#~ traj=[xstot,ystot,sstot,xstot2,ystot2,xstot3,ystot3,xstot4,ystot4,xstot5,ystot5]

	return traj