import tdsmParserMultis
import numpy as np
import os
import matplotlib.pyplot as plt

"""""
    This script was written to read any given number of tdms files and plot all collected forces and moments.
    To Use: Change directory, and check the group and channel names to ensure they are the same as below **
"""""
Directory = "../PilotStudy/Data/"
Files = os.listdir(Directory)


dataDict = {}
for files in Files:
        if files.endswith('.tdms'):
            data = tdsmParserMultis.parseTDMSfile(Directory + files)
            Fx = np.array(data[u'data'][u'Load Cell_Fx']) # ** check group and channel names
            Fy = np.array(data[u'data'][u'Load Cell_Fy'])
            Fz = np.array(data[u'data'][u'Load Cell_Fz'])
            Mx = np.array(data[u'data'][u'Load Cell_Mx'])
            My = np.array(data[u'data'][u'Load Cell_My'])
            Mz = np.array(data[u'data'][u'Load Cell_Mz'])

            dataDict[files] = [Fx, Fy, Fz, Mx, My, Mz]
        else:
            continue

check = True
for filename in dataDict:
    fig, (ax1, ax2) = plt.subplots(nrows=2, figsize=(10,10))

    if check:
        for data in range(6):
            _toPlot = dataDict[filename][data]

            color = ['red', 'orange', 'purple', 'green', 'black', 'blue']
            label = ['Fx', 'Fy', 'Fz', 'Mx', 'My', 'Mz']

            if data < 3:
                ax1.plot(_toPlot, color=color[data], lw=1, alpha=0.8, label='${}$'.format(label[data]))
                ax1.set_xlabel("time(ms)")
                ax1.set_ylabel("Force (N)")
                ax1.set_title(filename)
                leg1 = ax1.legend(loc='upper right', prop={'size': 10}, borderpad=0.1, handlelength=5)
                leg1.get_frame().set_alpha(0.8)

            if data >= 3:
                ax2.plot(_toPlot, color=color[data], lw=1, alpha=0.8, label='${}$'.format(label[data]))
                ax2.set_xlabel("time(ms)")
                ax2.set_ylabel("Torque (Nm)")
                ax2.set_title(filename)
                leg2 = ax2.legend(loc='upper right', prop={'size': 10}, borderpad=0.1, handlelength=5)
                leg2.get_frame().set_alpha(0.8)


            mngr = plt.get_current_fig_manager()
            mngr.window.setGeometry(100,100,1000,1000)
            plt.tight_layout()

        check = True
    plt.savefig(filename + '.png')
plt.show()