import numpy as np
import sys
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt

def radial_stagger(num_hst, width, height, hst_z, towerheight, R1, dsep=0., field='polar', savedir='.', plot=False):
	'''
	Ref. (Collado and Guallar, 2012), Campo: Generation of regular heliostat field.

	Generate a rather large field

	Arguements:
	num_hst    : int, number of heliostats
	width      : float, mirror width (m)
	height     : float, mirror height (m)
	hst_z      : float, the vertical location of each heliostat (m)
	towerheight: float, tower height (m)
	R1         : float, distance from the first row to the bottom of the tower (0, 0, 0)
	dsep       : float, separation distance (m)
	field      : str, 'polar-half' or 'surround-half' or 'polar' or 'surround' field, 
		               the 'half' option is for simulation a symmetric field
	savedir    : str, directory of saving the pos_and_aiming.csv

	Return:
	pos_and_aiming: (n, 7) array, position, focal length and aiming point of each generated heliostat
	'''

	# heliostat diagonal distantce
	DH=np.sqrt(height**2+width**2) 

	# distance between contiguous helistat center on the X and Y plane
	DM=DH+dsep

	# minimum radial increment
	delta_Rmin=0.866*DM

	# number of heliostats in the first row
	Nhel1 =int(2.*np.pi*R1/DM)


	# the total number of zones (estimated)
	#Nzones=int(np.log(5.44*3*(num_hst/az_rim*np.pi)/Nhel1**2+1)/np.log(4))+1

	X=np.array([])
	Y=np.array([])

	num=0
	i=0
	sys.stderr.write('DM '+repr(DM)+'\n')
	sys.stderr.write('dRm'+repr(delta_Rmin)+'\n')

	if field[-4:]=='half':

		while num<num_hst/2:
		    Nrows= int((2.**(i))*Nhel1/5.44)
		    Nhel=(2**(i))*Nhel1
		    R=Nhel/2./np.pi*DM
		    delta_az=2.*np.pi/Nhel
		    sys.stderr.write('\nzone %d\n' % (i,))
		    sys.stderr.write("R %f\n"%(R,))
		    sys.stderr.write("rows %d\n"%(Nrows,))
		    sys.stderr.write("hst %d\n"%(Nhel,))
		    sys.stderr.write("daz %f\n"%(delta_az,))

		    for row in range(Nrows):

		        r=R+float(row)*delta_Rmin

		        for nh in range(Nhel):

		            if num<num_hst/2:   
		       
		                if row%2==0:
		                    # the odd row
		                    azimuth=delta_az/2.+float(nh)*delta_az

		                else:
		                    # the even row
		                    azimuth=float(nh)*delta_az  
		         
		                xx=r*np.sin(azimuth)
		                yy=r*np.cos(azimuth)

		                if field=='polar-half':
		                    if (xx>0 and yy>0):  
		                        X=np.append(X, xx)
		                        Y=np.append(Y, yy)
		                        num+=1
		                elif field=='surround-half':
		                    if (yy>0):  
		                        X=np.append(X, xx)
		                        Y=np.append(Y, yy)
		                        num+=1
		    i+=1

	else:
		while num<num_hst:
		    Nrows= int((2.**(i))*Nhel1/5.44)
		    Nhel=(2**(i))*Nhel1
		    R=Nhel/2./np.pi*DM
		    delta_az=2.*np.pi/Nhel
		    
		    sys.stderr.write("\nzone %d"%(i,))
		    sys.stderr.write("R %f\n"%(R,))
		    sys.stderr.write("rows %d\n"%(Nrows,))
		    sys.stderr.write("hst %d\n"%(Nhel,))
		    sys.stderr.write("daz %f\n"%(delta_az,))		

		    for row in range(Nrows):

		        r=R+float(row)*delta_Rmin

		        for nh in range(Nhel):
		      
		            if num<num_hst:

		                if row%2==0:
		                    # the odd row
		                    azimuth=delta_az/2.+float(nh)*delta_az

		                else:
		                    # the even row
		                    azimuth=float(nh)*delta_az  
		         
		                xx=r*np.sin(azimuth)
		                yy=r*np.cos(azimuth)
		     
		                if field=='polar':
		                    if (yy>0):  
		                        X=np.append(X, xx)
		                        Y=np.append(Y, yy)
		                        num+=1
		                elif field=='surround':
		                        X=np.append(X, xx)
		                        Y=np.append(Y, yy)
		                        num+=1
		    i+=1

	sys.stderr.write("\ntotal hst %d\n"%(num,))
	hstpos=np.zeros(num*3).reshape(num, 3)
	hstpos[:, 0]=X
	hstpos[:, 1]=Y
	hstpos[:,2]=hst_z

	aim_x=np.zeros(num)
	aim_y=np.zeros(num)
	aim_z=np.ones(num)*towerheight

	foc=np.sqrt((X-aim_x)**2+(Y-aim_y)**2+(hstpos[:,2]-aim_z)**2)

	pos_and_aiming=np.append(X, (Y, hstpos[:,2], foc, aim_x, aim_y, aim_z))
	title=np.array(['x', 'y', 'z', 'foc', 'aim x', 'aim y', 'aim z', 'm', 'm', 'm', 'm', 'm', 'm', 'm'])
	pos_and_aiming=pos_and_aiming.reshape(7,int(len(pos_and_aiming)/7))
	pos_and_aiming=np.append(title, pos_and_aiming.T)
	pos_and_aiming=pos_and_aiming.reshape(int(len(pos_and_aiming)/7), 7)

	np.savetxt('%s/pos_and_aiming.csv'%savedir, pos_and_aiming, fmt='%s', delimiter=',')

	if plot:
		fts=24
		plt.figure(dpi=100.,figsize=(12,9))
		plt.plot(X, Y, '.')
		#plt.xlim(-1000, 1000)
		#plt.ylim(-1000, 1000)
		plt.xticks(fontsize=fts)
		plt.yticks(fontsize=fts)
		plt.xlabel('x (m)', fontsize=fts)
		plt.ylabel('y (m)', fontsize=fts)
		plt.savefig(savedir+'/field_layout.png', bbox_inches='tight')
		plt.close()

	return pos_and_aiming


if __name__=='__main__':
    
    pos_and_aim=radial_stagger(num_hst=4000, width=12.3, height=9.75, hst_z=1., towerheight=100, R1=87.5, dsep=0., field='polar', savedir='.', plot=True)

        
