"""
Retrieve and make a 2D spacetime plot of soliton detected edge data

"""
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import griddata
import os


path =  'data_path' # path to the data folder
root_folder_name = 'cropped_path'    # name of the cropped image data folder

camera_time_step = 1./120   # inverse of the camera frame rate
cm_per_pixel = 0.  # number of cm per pixel in the cropped image: use an image of the meter stick to get this.

edge_data_dir = path + '/' + root_folder_name + '_edge_data'

file_list = os.listdir(edge_data_dir)  
file_list = sorted([name for name in file_list if (name.endswith('.txt') and not name.startswith('.'))]) 

n_time_steps = len(file_list)

x = np.array([])
y = np.array([])
z = np.array([])

# load all the scattered data into big arrays

for i in range(n_time_steps):
    (y_data, x_data, z_data) = np.loadtxt(edge_data_dir+'/'+file_list[i])

    x = np.append(x, x_data)  # position along the tank, in pixels
    y = np.append(y, y_data)  # frame time in time steps
    z = np.append(z, z_data)  # edge detected surface position, in pixels

x = cm_per_pixel * x   # convert pixel units to cm
z = cm_per_pixel * z

y = camera_time_step * y  # convert timeteps to seconds

# define the grid to interpolate the data to
xmax = max(x)
ymax = max(y)

xi, yi = np.mgrid[0:xmax:300j, 0:ymax:200j]  # complex argument sets the number of points in the grid

# interpolate the data onto a regular grid
zi = griddata((x, y), z, (xi,yi), method='cubic')

num_contours = 30  # 20

# contour plot the interpolated data
plt.contourf(xi, yi, zi, num_contours, cmap=plt.cm.viridis)
plt.xlabel('Horizontal Position (cm)')
plt.ylabel('Time (seconds)')
plt.title(root_folder_name)
plt.colorbar() # draw a colorbar (units are cm)
plt.axes().set_aspect(1.0/plt.axes().get_data_ratio()*0.6)
plt.savefig(path+'/'+root_folder_name+'_spacetime.pdf')    

plt.show()