"""
Calculate the invariant model for the supplied position data.
"""

import numpy as np
from matplotlib import pyplot as plt
import scipy.interpolate as ip
import invariants_py.reparameterization as reparam
import invariants_py.calculate_invariants.opti_calculate_vector_invariants_position as ocp_tolerance
import invariants_py.calculate_invariants.opti_calculate_vector_invariants_position_mf as ocp_movingframe
import invariants_py.calculate_invariants.opti_calculate_vector_invariants_position_mj as ocp_minimumjerk
import csv as csv 
import os

def read_csv_file(fn, skipheader):
    '''
    Read a csv file (typically generated by an Orocos reporter) and
    stores its contents in a numpy array.
    For demonstrations, the first column (in the file and in the matrix) corresponds to time, the
    other columns to the different degrees of freedom.
    NOTE skipheader corresponds to the number of rows that belong to the header
    '''
    data = []
    header = []
    i=0
    with open(fn,'r') as csvfile:
        freader = csv.reader(csvfile,delimiter=' ',skipinitialspace=True)
        for row in freader:
            i=i+1
            if i > skipheader:
                #print '=>',', '.join(row)
                if len(row)>1:
                    if row[-1]== '':
                        data.append( [float(e) for e in row[:-1] ])
                    else:
                        data.append( [float(e) for e in row])
            else:
                header.append(row)
    return np.array( data ), header

def load_contour_data(data_location): 
    # Load contour and preprocess position data
    data,header = read_csv_file(data_location,1)
    position_data = np.array(data) # extract unique values px, py, pz
    #position_data[:,3] = np.mean(position_data[:,3]) # make Z-coordinate constant to ensure it's planar
    return position_data, header

if __name__ == '__main__':

    """Select contour"""
    contour_number = 1 # {0 : Rob's contour, 1 : Sabine's contour, 2 : Murtaza's contour} 

    """Select optimization"""
    optimization = ocp_tolerance # {ocp_tolerance, ocp_movingframe, ocp_minimumjerk}

    """Parameters of the calculation """    
    window_nb_samples = 20 # number of samples in window
    window_length = 0.05 # length of the window in [m], default 0.05m
    model_nb_samples = 250 # number of samples in invariant model
    
    # Parameters of the optimization problems
    weight_measurements = 1 # weight of the measurements in the optimization problem
    weight_regularization = 10**-12 # weight of the regularization term in the optimization problem
    weight_movingframes = 10**-6 # weight of the moving frames in the optimization problem

    """Load segmented data of chosen contour"""
    # set the correct paths
    script_dir = os.path.dirname(os.path.realpath(__file__))
    data_location = script_dir + '/data/segmented_2Dcontour_Sabine_Aug5.txt'
    position_data, header = load_contour_data(data_location) # load data
    
    # Plotting the data
    plt.figure(figsize=(8,3))
    plt.axis('equal')
    plt.plot(position_data[:,1],position_data[:,2],'.-')    
        
    """
    Reparameterize position data to path length (aka progress, arc length)
    """
    trajectory,time_profile,arclength,nb_samples,stepsize = reparam.reparameterize_positiontrajectory_arclength(position_data[:,1:],N=model_nb_samples)
    
    plt.figure(figsize=(8,3))
    plt.axis('equal')
    plt.plot(trajectory[:,0],trajectory[:,1],'.-')
 
    """
    Calculate model invariants
    """
    # specify optimization problem symbolically
    if optimization == ocp_minimumjerk:
        FS_calculation_problem = ocp_minimumjerk.OCP_calc_pos(window_len=nb_samples, w_pos = weight_measurements, w_regul = weight_regularization, geometric = True, weight_movingframes = weight_movingframes)
    elif optimization == ocp_movingframe:
        FS_calculation_problem = ocp_movingframe.OCP_calc_pos(window_len=nb_samples, planar_task = True, bool_unsigned_invariants = True, w_pos = 10, w_deriv = (10**-6)*np.array([1.0, 1.0, 1.0]), w_abs = (10**-10)*np.array([1.0, 1.0]), geometric=True)
    elif optimization == ocp_tolerance:
        FS_calculation_problem = ocp_tolerance.OCP_calc_pos(window_len=nb_samples, bool_unsigned_invariants = False, rms_error_traj = 0.001, geometric = True)

    # calculate invariants given measurements
    invariants, calculate_trajectory, movingframes = FS_calculation_problem.calculate_invariants(trajectory,stepsize)

    # figures
    plt.figure(figsize=(14,6))
    plt.subplot(2,2,1)
    plt.plot(trajectory[:,0],trajectory[:,1],'.-')
    plt.plot(calculate_trajectory[:,0],calculate_trajectory[:,1],'.-')
    plt.title('Trajectory [m]')

    plt.subplot(2,2,3)
    plt.plot(arclength,invariants[:,0])
    plt.plot(0,0)
    plt.title('Velocity [m/-]')

    plt.subplot(2,2,2)
    plt.plot(arclength,invariants[:,1])
    plt.plot(0,0)
    plt.title('Curvature [rad/-]')

    plt.subplot(2,2,4)
    plt.plot(arclength,invariants[:,2])
    plt.plot(0,1)
    plt.title('Torsion [rad/-]')

    plt.show()
    

