# FlockingAnalysis.py
# Run as FlockingAnalysis.py filename timestep

# Open up a data file of flock movements (called filename) and plot graphs of the differences between a given birds
# position/velocity and that of the CM of the flock as a function of time. 

# determine filename and timestep
import sys
filename = sys.argv[1]
timestep = float(sys.argv[2])

# import modules
from numpy import *
from pylab import *
from FlockingGetPositions import *

# Get the position data for the Original Flock
Data = GetPositions(filename)
flock_positions_x = Data[0]
flock_positions_y = Data[1]
bird_positions_x = Data[2]
bird_positions_y = Data[3]
N = Data[4]
num_timesteps = Data[5]


# Now use this position data to determine velocity data for the CM of the flock as well as individual birds

# Create Empty lists
flock_velocities_x = []
flock_velocities_y = []
bird_velocities_x = []
bird_velocities_y = []

for n in xrange(N):
    bird_velocities_x.append([])
    bird_velocities_y.append([])


# Loop through the Position Data to fill in lists 
for t in xrange(num_timesteps - 1):
    flock_velocities_x.append( (flock_positions_x[t+1] - flock_positions_x[t])/timestep )
    flock_velocities_y.append( (flock_positions_y[t+1] - flock_positions_y[t])/timestep )
    for n in xrange (N):
        bird_velocities_x[n].append ( (bird_positions_x[n][t+1] - bird_positions_x[n][t])/timestep ) 
        bird_velocities_y[n].append ( (bird_positions_y[n][t+1] - bird_positions_y[n][t])/timestep )


# Now determine the differences in position and velocity between the CM of the flock and the individual birds as a function of time

# Create Empty lists
x_position_differences = []
y_position_differences = []
x_velocity_differences = []
y_velocity_differences = []

for n in xrange(N):
    x_position_differences.append([])
    y_position_differences.append([])
    x_velocity_differences.append([])
    y_velocity_differences.append([])

# Fill in the lists
for t in xrange(num_timesteps):
    for n in xrange(N):
        x_position_differences[n].append( flock_positions_x[t] - bird_positions_x[n][t] )
        y_position_differences[n].append( flock_positions_y[t] - bird_positions_y[n][t] )
        if t < num_timesteps - 1: # one less data point for velocites than positions
              x_velocity_differences[n].append( flock_velocities_x[t] - bird_velocities_x[n][t] )
              y_velocity_differences[n].append( flock_velocities_y[t] - bird_velocities_y[n][t] )



# Also create a set of data containing the Smoothed (time averaged) velocity differences
# Note smooth factor should divide evenly into total number of velocity data points
smooth_factor = 500

# create the empty lists
x_velocity_differences_smoothed = []
y_velocity_differences_smoothed = []
for n in xrange(N):
    x_velocity_differences_smoothed.append([])
    y_velocity_differences_smoothed.append([])

# fill in the data
for n in xrange(N):
    xvs = 0.0
    yvs = 0.0
    
    for t in xrange(num_timesteps - 1):
        if mod(t+1,smooth_factor) != 0: 
           xvs = xvs + x_velocity_differences[n][t]
           yvs = yvs + y_velocity_differences[n][t]
        if mod(t+1,smooth_factor) == 0:
           x_velocity_differences_smoothed[n].append(xvs/smooth_factor)
           xvs = 0.0
           y_velocity_differences_smoothed[n].append(yvs/smooth_factor)
           yvs = 0.0


# create a list of times to plot our position data against
Tposition = []
for t in xrange(num_timesteps):
    Tposition.append(t*timestep)

# create a list of times to plot our velocity data against, one data point shorter
Tvelocity = []
for t in xrange(num_timesteps - 1):
    Tvelocity.append(t*timestep)

# Create a time series to ploot smoothed velocity data against
L = len(y_velocity_differences_smoothed[7]) #7 is arbitrary
T = smooth_factor*timestep*arange(L)


# Finally Plot some of this data


# Plot the differences in position between the individual birds and the CM of the flock as a function of time

# The x_position_differences
figure(1,(8,6))
title('X Position Differences')
xlabel('time')
ylabel('X Flock - X Bird')
for n in xrange(N):
    plot(Tposition, x_position_differences[n])

savefig(filename + 'XPositionDifferences', dpi=600)

# The y_position_differences
figure(2,(8,6))
title('Y Position Differences')
xlabel('time')
ylabel('Y Flock - Y Bird')
for n in xrange(N):
    plot(Tposition, y_position_differences[n])

savefig(filename + 'YPositionDifferences', dpi=600) 



# Plot the differences in velocities between the individual birds and the CM of the flock as a function of time

# The x_velocity_differences
figure(3,(8,6))
title('X Velocity Differences')
xlabel('time')
ylabel('Vx Flock - Vx Bird')
for n in xrange(N):
    plot(Tvelocity, x_velocity_differences[n])

savefig(filename + 'XVelocityDifferences', dpi=600)

# The y_velocity_differences
figure(4,(8,6))
title('Y Velocity Differences')
xlabel('time')
ylabel('Vy Flock - Vy Bird')
for n in xrange(N):
    plot(Tvelocity, y_velocity_differences[n]) 

savefig(filename + 'YVelocityDifferences', dpi=600)


# Plot the SMOOTHED differences in velocities between the individual birds and the CM of the flock as a function of time

# The smoothed x_velocity_differences
figure(5,(8,6))
title('X Velocity Differences Smoothed')
xlabel('time')
ylabel('Vx Flock - Vx Bird')
for n in xrange(N):
    plot(T, x_velocity_differences_smoothed[n])

savefig(filename + 'XVelocityDifferencesSmoothed', dpi=600)

# The smoothed y_velocity_differences
figure(6,(8,6))
title('Y Velocity Differences Smoothed')
xlabel('time')
ylabel('Vy Flock - Vy Bird')
for n in xrange(N):
    plot(T, y_velocity_differences_smoothed[n]) 

savefig(filename + 'YVelocityDifferencesSmoothed', dpi=600)

show()