'''
Function:     bar_image


Arguments:    x_pixels - Number of pixels in the horizontal (azimuthal) direction.
              y_pixels - Number of pixels in the vertical (elevational) direction.
              bar_parameters - (optional) Parameter dictionary, expected variables and definitions below.
                                 
              Expected contents of the optional dictionary bar_parameters:
              
              num_bars - Number of randomly placed bars in the generated spherical image. If not specified, the default
                         value is 100.
              arc_length - Length (in degrees) of each randomly placed bar. If not specified, the default value is 30.
              arc_width - Width (in degrees) of each randomly placed bar. If not specified, the default value is 5.
              
           
Output:       image - A binary (0,1) image with 1's corresponding to "on" coordinates, generated by randomly placing
                      "arc bars" on the surface of a sphere. An arc bar has a specified width --- the width in degrees
                      of a cross-section of each bar --- and length --- the length in degrees of the arcs which span
                      the length of each bar.
           
           
Description:  Generates a spherical image with randomly distributed bars. The bars are placed by determining a random
              point on the surface of the sphere (selected uniformly) as a corner, along with a random direction for
              the length of the bar. Then, the corner point is translated to form one width arc. This width arc is then
              translated along the length, filling out a single "arc bar". This process is repeated for the specified
              number of bars without consideration of bar overlaps.


Authors:     James Trousdale - jamest212@gmail.com
	     * Adapted from MATLAB code by Sam Carroll.
'''


__all__ = ["bar_image"]

import numpy as np
from numpy.random import rand as rand

from numpy import sin as sin
from numpy import cos as cos

import sys

sys.path.append('../')
from utilities import axis_rot_mat,my_2d_hist


#                                            #
#                                            #
###### definition of function bar_image ######
#                                            #
#                                            #

def bar_image(x_pixels,y_pixels,bar_parameters={}):
	image = np.zeros((y_pixels,x_pixels))
	
	# Ensure that the height and width of pixels agree (this could probably be generalized without much trouble).
	if (np.float(x_pixels)/360 != np.float(y_pixels)/180):
		print 'X resolution and Y resolution do not match in generation of a bar image. Exiting...\n\n'
		exit(1)
	
	# Check if bar image parameters are given --- if not, set to default values.
	if 'num_bars' in bar_parameters:
		num_bars = bar_parameters['num_bars']
	else:
		num_bars = 100
	if 'arc_length' in bar_parameters:
		arc_length = bar_parameters['arc_length']
	else:
		arc_length = 30
	if 'arc_width' in bar_parameters:
		arc_width = bar_parameters['arc_width']
	else:
		arc_width = 5
	
	# Convert length and width to radians.
	arc_length = arc_length*np.pi/180
	arc_width = arc_width*np.pi/180
	
	# Determine the azimuth and elevation of pixel edges.
	theta_edges = np.linspace(0,360,np.size(image,1)+1,endpoint=True)*np.pi/180
	phi_edges = np.linspace(0,180,np.size(image,0)+1,endpoint=True)*np.pi/180
			
	# Rotation resolution for generating bars - rotate ~ half the horizontal width of a pixel each time to ensure that
	# bars are completely filled.
	D = 360.0/x_pixels/2*np.pi/360
	
	# Compute the number of rotational steps taken in spanning the width and length of each bar.
	n_width_steps = np.int64(np.ceil(arc_width/D))
	n_length_steps = np.int64(np.ceil(arc_length/D))
	

	# Place bars on the image
	for _ in range(num_bars):
		# Uniformly sample a point on the surface of the unit sphere to be a corner of the bar being generated.
		theta_init = 2*np.pi*rand()
		phi_init = np.arccos(2*rand()-1)
		xyz_init = np.array([sin(phi_init)*cos(theta_init),sin(phi_init)*sin(theta_init),cos(phi_init)])
				
		# Compute the axis around which we will rotate the bar length-wise. This is randomly selected to be a vector
		# lying in the plane perpendicular to the vector pointing at the bar corner generated above. By selecting such
		# a vector, we guarantee that a rotation by the specified arc length actually generates bars with that length.
		xyz_l_axis = np.zeros(3)
		xyz_l_axis[1] = rand()
		xyz_l_axis[2] = rand()
		xyz_l_axis[0] = -(xyz_l_axis[1]*xyz_init[1]+xyz_l_axis[2]*xyz_init[2])/xyz_init[0]
		xyz_l_axis = xyz_l_axis/np.linalg.norm(xyz_l_axis)

		# Similarly, select the axis about which we rotate to form the width of the bar. Given that a corner of the bar
		# and the length are already determined, the only degree of freedom remaining for the vector directed along the
		# axis of rotation to form the bar width is its length. Thus, we initially set the z component to 1 and solve
		# for the x,y components before normalizing to a unit vector.
		xyz_w_axis = np.zeros(3)
		xyz_w_axis[2] = 1
		xyz_w_axis[:2] = np.linalg.solve([[xyz_init[0],xyz_init[1]],[xyz_l_axis[0],xyz_l_axis[1]]],
									[-xyz_init[2]*xyz_w_axis[2],-xyz_l_axis[2]*xyz_w_axis[2]])
		xyz_w_axis = xyz_w_axis/np.linalg.norm(xyz_w_axis)
		
		# The procedure will be to first rotate the corner to give the rectangular coordinates of one width-wise
		# cross-section of the bar.
		xyz_bar = np.zeros((3,n_width_steps))
		xyz_bar[:,0] = xyz_init;
		
		rot_mat_w = axis_rot_mat(xyz_w_axis,D) # Rotational matrix for the width-wise rotation.

		# Rotate the initial corner along the width of the bar, recording the rectangular coordinates of the covered
		# points.
		for j in range(1,n_width_steps):
			xyz_bar[:,j] = np.dot(rot_mat_w,xyz_bar[:,j-1])
			
		theta = np.zeros(n_width_steps*n_length_steps)
		phi = np.zeros(n_width_steps*n_length_steps)
		
		# Determine the spherical coordinates of the "on" pixels corresponding to the generated width cross-section of
		# the bar.
		theta[0:n_width_steps] = np.mod(np.arctan2(xyz_bar[1,:],xyz_bar[0,:]),2*np.pi)
		phi[0:n_width_steps] = np.arccos(xyz_bar[2,:])
		
		rot_mat_l = axis_rot_mat(xyz_l_axis,D) # Rotational matrix for the length-wise rotation.
			
		# Rotate the width cross-section along the length of the bar, recording spherical coordinates of the covered
		# points.
		for j in range(1,n_length_steps):
			xyz_bar = np.dot(rot_mat_l,xyz_bar)
			theta[(j*n_width_steps):((j+1)*n_width_steps)] = np.mod(np.arctan2(xyz_bar[1,:],xyz_bar[0,:]),2*np.pi)
			phi[(j*n_width_steps):((j+1)*n_width_steps)] = np.arccos(xyz_bar[2,:])
			
		# Increment pixel values corresponding to the points covering the bar by 1.
		image += my_2d_hist(theta,phi,theta_edges,phi_edges)

	# Many pixels may have values greater than 1, owing to overlaps greater than 1. Set these pixels to 1.
	image[image > 1] = 1
	
	return image


#                                                #
#                                                #
###### default code executed when bar_image ###### 
###### is called without parameters         ######
#                                                #
#                                                #


if __name__ == "__main__":
	import matplotlib.pyplot as plt
	
	x_pixels = 360
	y_pixels = 180
	bar_parameters = {'num_bars':25,'arc_length':40,'arc_width':5}
	image = bar_image(x_pixels,y_pixels,bar_parameters)
		
	plt.imshow(image,cmap='gray')
	plt.show()
