import os
import numpy as np
import pandas as pd

from IPython import embed

def main():
    df = pd.read_csv(os.path.join("derived_data", "figure2_baseline_properties.csv"), sep=";", index_col=0)
    # calculate velocities for the cells in the front as a function of
    # their position-difference to cells behind them
    receptor_pos = df.receptor_pos_absolute.values
    phase_time = (df.phase_shifted/(2*np.pi) * df.eod_period).values * 1000 # ms
    receptor_pos_sort, phase_time_sort = list(zip(*sorted(list(zip(receptor_pos, phase_time)))))
    receptor_pos_sort = np.asarray(receptor_pos_sort)
    phase_time_sort = np.asarray(phase_time_sort)

    start_cell_num = 10
    min_cell_dist = 5  # mm

    shape = (start_cell_num, receptor_pos.shape[0])
    pos_diff = np.ones(shape) * np.nan
    time_diff = np.ones(shape) * np.nan
    velocity = np.ones(shape) * np.nan
    for i, (start_pos, start_time) in enumerate(zip(receptor_pos_sort[:start_cell_num], phase_time_sort[:start_cell_num])):
        for j, (end_pos, end_time) in enumerate(zip(
                receptor_pos_sort[i+min_cell_dist:],
                phase_time_sort[i+min_cell_dist:])):

            if (end_time-start_time) < 0: # disregard negative velocities (since its within-group variance)
                continue

            pos_diff[i,j] = end_pos-start_pos
            time_diff[i,j] = end_time-start_time
            velocity[i,j] = pos_diff[i,j]/time_diff[i,j]

    # binning for average
    binWidth = 4
    flat_pos_diff = pos_diff[np.where(np.isfinite(pos_diff))]
    flat_time_diff = time_diff[np.where(np.isfinite(time_diff))]
    flat_velocity = velocity[np.where(np.isfinite(velocity))]
    pos_diff_bins = np.arange(np.floor(flat_pos_diff.min().min()), np.ceil(flat_pos_diff.max().max())+binWidth, binWidth)
    pos_diff_centers = pos_diff_bins[:-1]+binWidth/2
    avg_velocity = np.zeros(pos_diff_centers.shape[0])

    for i, (left, right) in enumerate(zip(pos_diff_bins[:-1], pos_diff_bins[1:])):
        avg_velocity[i] = np.median(flat_velocity[(flat_pos_diff >= left) & (flat_pos_diff < right)])    

    neighbor_num = 5
    neighbor_velocities = np.ones((receptor_pos_sort.shape[0]-neighbor_num, neighbor_num)) * np.nan
    for i, (start_pos, start_time) in enumerate(zip(
            receptor_pos_sort[:-neighbor_num],
            phase_time_sort[:-neighbor_num])):

        j = 0
        for end_pos, end_time in  zip(receptor_pos_sort[i:], phase_time_sort[i:]):
            if (end_time-start_time) <= 0:
                continue
            neighbor_velocities[i,j] = (end_pos-start_pos)/(end_time-start_time)
            j += 1
            if j >= neighbor_num:
                break

    np.savez_compressed(os.path.join("derived_data", "suppfig5_velocities.npz"), position_differences=pos_diff, velocities=velocity,
                        average_velocities=avg_velocity, position_difference_centers=pos_diff_centers)