# This software may be modified and distributed under the terms of the
# BSD-3-Clause license.
# Authors: Giulio Romualdi and Nicola Piga

#!/usr/bin/env python3

import numpy as np
import argparse
import csv
import matplotlib
import matplotlib.pyplot as plt

def read_data(file_name):
    data = []

    with open(file_name, newline='') as csv_data:
        for row in csv_data:
            try:
                data.append([float(num_string.rstrip()) for num_string in row.split(sep = ", ") if num_string != ''])
            except:
                pass

    return np.array(data)

def times(signal, i0 = 0):
    times = signal.copy()
    counter = i0 * (1.0 / 100.0)
    dt = 1.0 / 100.0
    for i in range(len(times)):
        times[i] = counter
        counter = counter + dt

    return times

def plot_data(axes, x_data, y_data, color, linewidth = 1.0):

    col = []
    if color == "blue":
        col = plt.cm.Blues(0.8)
    elif color == "red":
        col = plt.cm.Reds(0.8)

    return axes.plot(x_data, y_data, color = col, linewidth = linewidth)

def plot_and_save(data):
    fig, ax = plt.subplots(1,2)

    channel_0 = data[:, 0]
    channel_1 = data[:, 1]

    plot_0, = plot_data(ax[0], times(channel_0), channel_0, "red")
    plot_1, = plot_data(ax[0], times(channel_1), channel_1, "blue")

    legend = [plot_0, plot_1]
    legend = fig.legend(legend, labels = ["Reference", "Actual"], ncol = 2, loc = "upper center", frameon=False)

    for line in legend.get_lines():
        line.set_linewidth(2.0)

    label_font_size = 10
    ax[0].grid()
    ax[0].set_xlabel('time [s]', fontsize = label_font_size)
    ax[0].set_ylabel('position [rad]', fontsize = label_font_size)
    ax[0].set_title('joint position', fontsize = label_font_size)

    error = channel_1 - channel_0
    plot_error, = ax[1].plot(times(channel_0), error, color="green")
    ax[1].grid()
    ax[1].set_xlabel('time [s]', fontsize=label_font_size)
    ax[1].set_ylabel('error [rad]', fontsize=label_font_size)
    ax[1].set_title('Error Plot', fontsize=label_font_size)

    figure = plt.gcf()
    plt.tight_layout()
    plt.savefig("./figure.png", bbox_inches='tight', dpi = 150)

    # plt.show()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='plot_dataset.py is a python script useful to plot a dataset generated by JointControlModule and save it.')
    parser.add_argument('--dataset', type = str, required = True, help='Name of the dataset')
    args = parser.parse_args()

    data = read_data(args.dataset)

    plot_and_save(data)
