'''
Function:     vert_strip

Arguments:    x_pixels - Number of pixels in the horizontal (azimuthal) direction.
              y_pixels - Number of pixels in the vertical (elevational) direction.
              rot_angle - Azimuthal location of the center of the vertical strip stimulus. (deg)
              strip_rad - Azimuthal radius of the vertical strip stimulus. (deg)
              freq_s - Number of equally spaced bars which comprise the moving strip stimulus.
              deg_per_ms - Downward velocity of the bars. (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
                              downward motion of bars within a vertical strip. 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:  Generates a binary (0,1) movie by first forming a set of bars on the sphere by truncating a sinusoid with
              spatial frequency freq_s/(180 deg) at 0, then moving these bars downward at a speed deg_per_ms. The width 
              of the strip (and the bars) is 2*strip_rad, and the strip is centered at azimuth rot_angle.

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


import numpy as np


def vert_strip(x_pixels,y_pixels,rot_angle,strip_rad,freq_s,deg_per_ms,T,dt): 
    
    time_steps = int(np.ceil(T/dt))+1
    rotated_image = np.zeros((y_pixels,x_pixels,time_steps))
    
    rad_pixels = int(np.floor(strip_rad/360.0*x_pixels)) # Number of pixels the strip spans width-wise.
    rot_bin = int(np.floor(rot_angle/180.0*y_pixels)) # The horizontal index of the bin on which the strip is centered.
    
    # Convert bar movement speed to radians per time step
    rad_per_step = deg_per_ms*np.pi/180*dt
    
    # Calculate the elevation angles (rad) of the pixel centers (array of length ypixels+1).
    phi_vals = np.linspace(0,np.pi,y_pixels+1,endpoint=True)
    # all elements except last plus all elements except first divided by 2 (array of length ypixels).
    phi_vals = (phi_vals[:-1] + phi_vals[1:])/2
    
    
    # Generate the movie by setting each column of the image which the stimulus strip spans to the binary values of the
    # truncated sinusoid.
    for i in range(time_steps):
        for j in range(-rad_pixels,rad_pixels+1):
            #negative indices in the x-dimension are accessed from the last one of the array counting backwards
            #set all y-pixels at once since phi_vals contains the elevation of the centers of each pixel
            rotated_image[:,rot_bin+j,i] = 1.0*(np.sin(2*freq_s*(phi_vals - rad_per_step*i)) > 0)
        
    return rotated_image
    


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'

    x_pixels = 360
    y_pixels = 180
    T = 5000
    dt = 1000/60
    rot_angle = -180
    strip_rad = 15
    freq_s = 8
    deg_per_ms = 0.5
        
    rot_im = vert_strip(x_pixels,y_pixels,rot_angle,strip_rad,freq_s,deg_per_ms,T,dt)
        
    fig = plt.figure(figsize=(5,5))
    ims = []
    for i in xrange(np.int64(T/dt)):
        ims.append((plt.imshow(rot_im[:,:,i],cmap='gray'),))
        
    ani = animation.ArtistAnimation(fig,ims,interval=dt,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/strip.mp4', writer = FFwriter, fps=60/2)
    
    #Uncomment this line to make sure the movie generation works
    #plt.show()