def calc_heading(trajec):
    trajec.heading_norollover = floris_math.remove_angular_rollover(np.arctan2(trajec.velocities[:,1], trajec.velocities[:,0]), 3)
    ## kalman
    
    data = trajec.heading_norollover.reshape([len(trajec.heading_norollover),1])
    ss = 3 # state size
    os = 1 # observation size
    F = np.array([   [1,1,0], # process update
                     [0,1,1],
                     [0,0,1]],
                    dtype=np.float)
    H = np.array([   [1,0,0]], # observation matrix
                    dtype=np.float)
    Q = np.eye(ss) # process noise
    Q[0,0] = .01
    Q[1,1] = .01
    Q[2,2] = .01
    R = 1*np.eye(os) # observation noise
    
    initx = np.array([data[0,0], data[1,0]-data[0,0], 0], dtype=np.float)
    initv = 0*np.eye(ss)
    xsmooth,Vsmooth = kalman_math.kalman_smoother(data, F, H, Q, R, initx, initv, plot=False)

    trajec.heading_norollover_smooth = xsmooth[:,0]
    trajec.heading_smooth_diff = xsmooth[:,1]*trajec.fps
    
    trajec.heading = floris_math.fix_angular_rollover(trajec.heading_norollover)
    trajec.heading_smooth = floris_math.fix_angular_rollover(trajec.heading_norollover_smooth)
def calc_heading_for_axes(trajec, axis='xy'):
    if axis == 'xy':
        axes = [1,0]
    elif axis == 'xz':
        axes = [2,0]
    elif axis == 'yz':
        axes = [2,1]
    
        
    heading_norollover_for_axes = floris_math.remove_angular_rollover(np.arctan2(trajec.velocities[:,axes[0]], trajec.velocities[:,axes[1]]), 3)
    ## kalman
    
    data = heading_norollover_for_axes.reshape([len(heading_norollover_for_axes),1])
    ss = 3 # state size
    os = 1 # observation size
    F = np.array([   [1,1,0], # process update
                     [0,1,1],
                     [0,0,1]],
                    dtype=np.float)
    H = np.array([   [1,0,0]], # observation matrix
                    dtype=np.float)
    Q = np.eye(ss) # process noise
    Q[0,0] = .01
    Q[1,1] = .01
    Q[2,2] = .01
    R = 1*np.eye(os) # observation noise
    
    initx = np.array([data[0,0], data[1,0]-data[0,0], 0], dtype=np.float)
    initv = 0*np.eye(ss)
    xsmooth,Vsmooth = kalman_math.kalman_smoother(data, F, H, Q, R, initx, initv, plot=False)

    heading_norollover_smooth_for_axes = xsmooth[:,0]
    heading_smooth_diff_for_axes = xsmooth[:,1]*trajec.fps
    heading_for_axes = floris_math.fix_angular_rollover(heading_norollover_for_axes)
    heading_smooth_for_axes = floris_math.fix_angular_rollover(heading_norollover_smooth_for_axes)
    
    if axis == 'xy':
        trajec.heading_norollover_smooth_xy = heading_norollover_smooth_for_axes
        trajec.heading_smooth_diff_xy = heading_smooth_diff_for_axes
        trajec.heading_xy = heading_for_axes
        trajec.heading_smooth_xy = heading_smooth_for_axes
    if axis == 'xz':
        trajec.heading_norollover_smooth_xz = heading_norollover_smooth_for_axes
        trajec.heading_smooth_diff_xz = heading_smooth_diff_for_axes
        trajec.heading_xz = heading_for_axes
        trajec.heading_smooth_xz = heading_smooth_for_axes
    if axis == 'yz':
        trajec.heading_norollover_smooth_yz = heading_norollover_smooth_for_axes
        trajec.heading_smooth_diff_yz = heading_smooth_diff_for_axes
        trajec.heading_yz = heading_for_axes
        trajec.heading_smooth_yz = heading_smooth_for_axes
def calc_heading_from_velocities(velocities):
    heading_norollover = floris_math.remove_angular_rollover(np.arctan2(velocities[:,1], velocities[:,0]), 3)
    return heading_norollover
def calc_saccades_z(trajec, threshold_lo=200, threshold_hi=100000000, min_angle=10, plot=False):
    # thresholds in deg/sec
    # min_angle = minimum angle necessary to count as a saccade, deg

    fps = 100.
    
    heading_norollover_for_axes = floris_math.remove_angular_rollover(np.arctan2(trajec.speed_xy[:], trajec.velocities[:,2]), 3)
    ## kalman
    
    data = heading_norollover_for_axes.reshape([len(heading_norollover_for_axes),1])
    ss = 3 # state size
    os = 1 # observation size
    F = np.array([   [1,1,0], # process update
                     [0,1,1],
                     [0,0,1]],
                    dtype=np.float)
    H = np.array([   [1,0,0]], # observation matrix
                    dtype=np.float)
    Q = np.eye(ss) # process noise
    Q[0,0] = .01
    Q[1,1] = .01
    Q[2,2] = .01
    R = 1*np.eye(os) # observation noise
    
    initx = np.array([data[0,0], data[1,0]-data[0,0], 0], dtype=np.float)
    initv = 0*np.eye(ss)
    xsmooth,Vsmooth = kalman_math.kalman_smoother(data, F, H, Q, R, initx, initv, plot=False)

    heading_norollover_smooth_for_axes = xsmooth[:,0]
    heading_smooth_diff_for_axes = xsmooth[:,1]*trajec.fps
    heading_for_axes = floris_math.fix_angular_rollover(heading_norollover_for_axes)
    heading_smooth_for_axes = floris_math.fix_angular_rollover(heading_norollover_smooth_for_axes)
    
    trajec.heading_altitude_smooth = heading_smooth_for_axes
    trajec.heading_altitude_smooth_diff = heading_smooth_diff_for_axes
    
    ## saccades
    
    possible_saccade_array = (np.abs(trajec.heading_altitude_smooth_diff)*180/np.pi > threshold_lo)*(np.abs(trajec.heading_altitude_smooth_diff)*180/np.pi < threshold_hi)
    possible_saccades = nim.find_blobs(possible_saccade_array, [3,100])
    
    if len(possible_saccades) == 1:
        if np.sum(possible_saccades[0]) == 0:
            possible_saccades = []
    
    trajec.saccades_z = []
    
    
    if len(possible_saccades) > 0:
        for sac in possible_saccades:
            indices = np.where(sac==1)[0].tolist()
            if len(indices) > 0:
                # expand saccade range to make sure we get full turn
                #lo = np.max([indices[0]-5, 0])
                #hi = np.min([indices[-1]+5, len(trajec.speed)-2])
                #new_indices = np.arange(lo, hi).tolist()
                #tmp = np.where( np.abs(trajec.heading_diff_window[new_indices])*180/np.pi > 350 )[0].tolist()
                #indices = np.array(new_indices)[ tmp ].tolist()
                angle_of_saccade = np.abs(get_angle_of_saccade_z(trajec, indices)*180./np.pi)
                mean_speed = np.mean(trajec.speed[indices])
                if len(indices) > 3 and angle_of_saccade > 10: # and mean_speed > 0.005:
                    trajec.saccades_z.append(indices)

    if plot:
        
        fig = plt.figure()
        ax = fig.add_subplot(111)
        
        ax.plot(trajec.positions[:,0], trajec.positions[:,2], color='black')
        for sac in trajec.saccades_z:
            ax.plot(trajec.positions[sac,0], trajec.positions[sac,2], color='red')