'''
Function:     rotate_sphere

Arguments:    image - The image to be rotated, a two-dimensional array of size (y_pixels,x_pixels).
		      rot_angle - Horizontal rotation axis. Rotation about this axis induces apparent clockwise rotation of the
		      image. (deg)
	      deg_per_ms - Speed of the clockwise rotation about the horizontal axis at the angle rot_angle. (deg/ms)
	      T - Total movie duration. (ms)
	      dt - Time step used for generating the movie. Inverse of the sampling rate. (ms)
		
Output:       rotated_image - A 3-dimensional array of size (y_pixels,x_pixels,T/dt+1) containing the movie generated by
			      rotating the given image about the horizontal axis at angle rot_angle, clockwise at speed
			      deg_per_ms. The last axis corresponds to time (i.e., the frames of the movie), sampled
			      at equally spaced time points dt apart over the interval [0,T].
		
Description:  Accepts a spherical image, and generates a movie by rotating the image about the specified horizontal axis
	      in the clockwise direction, at the specified rotational speed. Considering a specific frame of the movie,
	      the value of a pixel in the frame is found by determining the spherical coordinates of the frame center,
	      then inverting the rotation to determine the time t=0 location of that pixel center. The pixel value of 
	      the frame under consideration is set by the t=0 value of the pixel in which the inverse-rotated pixel 
	      center lies. In other words, the time t=t_1 value of a pixel P is set to the value of the pixel in the
	      original image (the time t=0 frame of the movie) which, under the total rotation over the interval
	      [0,t_1], would lie nearest to the pixel P. 
			
	      For binary images, an alternative method would be to push forward "on" pixels only, which would be 
	      significantly faster. However, a push-forward method could be susceptible to tearing around the poles due 
	      to differential pixel density, and the method does not generalize well to non-binary images.

Authors:     James Trousdale - jamest212@gmail.com
'''


import numpy as np

from numpy import sin as sin
from numpy import cos as cos
from numpy import arctan2 as arctan2
from numpy import arccos as arccos


#                                                #
#                                                #
###### definition of function rotate_sphere ######
#                                                #
#                                                #


def rotate_sphere(image,rot_angle,deg_per_ms,T,dt):
	time_steps = int(np.ceil(T/dt)+1)
	rotated_image = np.zeros(np.append(np.shape(image),time_steps))
	
	rad_per_step = deg_per_ms*np.pi/180.0*dt; # Radians of rotation per time step
	
	# The vectors theta_vals and phi_vals contain the azimuth and elevation of the pixel centers, respectively.
	theta_vals = np.linspace(0,2*np.pi,np.size(image,1)+1,endpoint=True)
	theta_vals = (theta_vals[:-1] + theta_vals[1:])/2
	phi_vals = np.linspace(0,np.pi,np.size(image,0)+1,endpoint=True)
	phi_vals = (phi_vals[:-1] + phi_vals[1:])/2
	
	# Compute the horizontal and vertical dimensions of pixels in degrees
	dth = theta_vals[1]-theta_vals[0]
	dphi = phi_vals[1]-phi_vals[0]
	
	# Pre-calculate the rectangular coordinates of every pixel center.
	pixel_rect_coords = np.zeros((3,np.prod(np.size(image))))
	pixel_rect_coords[0,:] = np.reshape(np.outer(sin(phi_vals),cos(theta_vals)),np.prod(np.size(image)))
	pixel_rect_coords[1,:] = np.reshape(np.outer(sin(phi_vals),sin(theta_vals)),np.prod(np.size(image)))
	pixel_rect_coords[2,:] = np.reshape(np.outer(cos(phi_vals),np.ones(len(theta_vals))),np.prod(np.size(image)))
	
	u = -np.array([cos(rot_angle*np.pi/180),sin(rot_angle*np.pi/180),0]) # Unit vector in direction of rotation axis
		
	reshaped_image = np.reshape(image,np.prod(np.shape(image))) # Vectorized image

	
	# As described above, at each step, the value of every pixel is determined by rotating that pixel backwards in time
	# in order to determine its time t=0 position. Then, the value is read out from the original image.
	for i in range(time_steps):
		
		th = -rad_per_step*(i+1) # Inverse of the total rotation up to the current time.
		
		# Calculate the matrix rotation which yields the inverse rotation when applied to a rectangular coordinate.
		rot_mat = [[cos(th)+u[0]**2*(1-cos(th)),u[0]*u[1]*(1-cos(th))-u[2]*sin(th),u[0]*u[2]*(1-cos(th))+u[1]*sin(th)],
		[u[1]*u[0]*(1-cos(th))+u[2]*sin(th),cos(th)+u[1]**2*(1-cos(th)),u[1]*u[2]*(1-cos(th))-u[0]*sin(th)],
		[u[2]*u[0]*(1-cos(th))-u[1]*sin(th),u[2]*u[1]*(1-cos(th))+u[0]*sin(th),cos(th)+u[2]**2*(1-cos(th))]]
		
		# Calculate the rectangular coordinates at time t=0 of every pixel.
		xyz = np.dot(rot_mat,pixel_rect_coords)
		
		# Snap the rectangular coordinates computed above to corresponding spherical pixels.
		theta_inds = np.int64(np.floor(np.mod(arctan2(xyz[1,:],xyz[0,:]),2*np.pi)/dth))
		phi_inds = np.int64(np.floor(arccos(xyz[2,:])/dphi))
		
		# Compute the indices of the spherical pixels for the vectorized image, and set the current frame by the values
		# of the rotated pixels in the original image.
		inds = np.int64(np.size(image,1)*np.mod(phi_inds,len(phi_inds)) + theta_inds)
		rotated_image[:,:,i] = np.reshape(reshaped_image[inds],np.shape(image))
		
	return rotated_image
	

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


if __name__ == "__main__":
	import matplotlib.pyplot as plt
	import matplotlib.animation as animation
	
	#NOTE: this line is specific to the ffmpeg installed using MacPorts on Mac OS X and 
	#may need to be commented or changed depending on your specific installation
	#of ffmpeg
	plt.rcParams['animation.ffmpeg_path'] = '/opt/local/bin/ffmpeg'
	
	from bar_image import *
	from natural_image import *
	from white_image import *
	
	x_pixels = 360
	y_pixels = 180
	rot_angle = 45
	deg_per_ms = .5
	t = 5000
	dt = 1000/60
	bar_parameters = {'num_bars':25,'arc_length':40,'arc_width':5}
	
	# Note: Uncomment one of the following 3 lines to generate movies from rotations of different image types
	#image = natural_image(x_pixels=360,y_pixels=180,image_lib_path='~/Research/projects/fly/images/')
	#image = white_image(x_pixels=360,y_pixels=180,upsample=4)
	image = bar_image(x_pixels,y_pixels,bar_parameters)
	
	rotated_image = rotate_sphere(image,rot_angle,deg_per_ms,t,dt)
	
	
	fig = plt.figure(figsize=(5,5))
	ims = []
	
	for i in xrange(np.int64(t/dt)):
		ims.append((plt.imshow(rotated_image[:,:,i],cmap='gray'),))
		
	ani = animation.ArtistAnimation(fig,ims,interval=dt*10,repeat_delay=3000,blit=True)
	
	FFwriter = animation.FFMpegWriter()
	
	#comment this line if you encounter an issue saving the animation and uncomment
	#the one below to check the basic animation generation code works.
	ani.save('output/rotation_movie.mp4', writer = FFwriter, fps=60/2)
	
	#Uncomment this line to make sure the movie generation works
        #plt.show()