import tdsmParserMultis
import dicom
import numpy as np
import matplotlib.pyplot as plt
import SimpleITK as sitk
import warnings
import xml.etree.ElementTree as ET
from XMLparser import getAcceptance
import peakutils
import os
import math

from scipy import stats

warnings.simplefilter('ignore', np.RankWarning)
warnings.simplefilter('ignore', RuntimeWarning)

dir = '/Users/schimmt/Documents/Multis/studies/InstrumentedUltrasound/dat/Validation/MULTIS103-1'
tdmsDirectory = dir + '/Data/'

tdms_list= sorted(os.listdir(tdmsDirectory))
tdms_files = []
for files in tdms_list:
    if files.endswith('.tdms'):
        tdms_files.append(files)
cnt = 0
a = open(dir + '/' + 'resultant.txt', 'wb')
fig1, ax1 = plt.subplots(nrows=1, ncols=1, figsize=(18, 9.5))
resultant = []
for tdms_filename in tdms_files:

    tdms_file = tdmsDirectory + tdms_filename

    data = tdsmParserMultis.parseTDMSfile(tdms_file)
    Fx = np.array(data[u'State.6-DOF Load'][u'6-DOF Load Fx'])
    Fy = np.array(data[u'State.6-DOF Load'][u'6-DOF Load Fy'])
    Fz = np.array(data[u'State.6-DOF Load'][u'6-DOF Load Fz'])
    Fxx = np.array(data[u'Sensor.Load Cell'][u'Load Cell_Fx'])  # Fx and Fz reversed
    Fyy = np.array(data[u'Sensor.Load Cell'][u'Load Cell_Fy'])
    Fzz = np.array(data[u'Sensor.Load Cell'][u'Load Cell_Fz'])
    Mx = np.array(data[u'State.6-DOF Load'][u'6-DOF Load Mx'])
    My = np.array(data[u'State.6-DOF Load'][u'6-DOF Load My'])
    Mz = np.array(data[u'State.6-DOF Load'][u'6-DOF Load Mz'])
    pulse = np.array(data[u'Sensor.Run Number Pulse Train'][u'Run Number Pulse Train'])

    int1 = 1000
    int2 = 2000
    int3 = 6000
    int4 = 7000

    dataList = [Fx, Fy, Fz, Fxx, Fyy, Fzz, Mx, My, Mz]
    dataNames = ['Fx', 'Fy', 'Fz', 'Fxx', 'Fyy', 'Fzz', 'Mx', 'My', 'Mz']
    i = 0
    intdataList1 = []
    for data in dataList:
        FT1 = data[int1:int2]
        FT2 = data[int3:int4]

        avgFT1 = np.average(FT1)
        avgFT2 = np.average(FT2)

        diff = avgFT1 - avgFT2
        intdataList1.append([dataNames[i],avgFT1,avgFT2, diff])

        i += 1

    # print intdataList1
    i = 0
    for forces in intdataList1:
        # print forces[0]

        if forces[0] == 'Fx':
            ax1.plot(cnt, forces[-1], color='red', marker='o')
            fx = forces[-1]
        if forces[0] == 'Fy':
            ax1.plot(cnt, forces[-1], color='blue', marker='o')
            fy = forces[-1]
        if forces[0] == 'Fz':
            ax1.plot(cnt, forces[-1], color='green', marker='o')
            fz = forces[-1]
        # else:
        #     ax1.plot(cnt, forces[-1], color='purple', marker='o')
        # plt.axhline(y=4.82, color='c', linestyle='-')
        plt.axhline(y=-4.82, color='c', linestyle='-', alpha=0.5)
        plt.axhline(y=-0, color='c', linestyle='-', alpha=0.5)
        ax1.set_xlim(-.5, len(tdms_files))

        i += 1

    fx = intdataList1[0][-1]
    fy = intdataList1[1][-1]
    fz = intdataList1[2][-1]

    mag = np.sqrt((fx**2)+(fy**2)+(fz**2))
    resultant.append(mag)
    a.write("%s\n" % (str(round(mag,4))))

    print mag


    ax1.scatter(cnt+.2,-mag, color='black', marker='*')

    # plt.show()
    # exit()
    # fig, (ax2, ax3) = plt.subplots(nrows=2, ncols=1, figsize=(18, 9.5))
    #
    # ax2.plot(Fx, label="$Fx$", color='blue',linestyle='-')
    # ax2.plot(Fy, label="$Fy$", color='red', linestyle='-')
    # ax2.plot(Fz, label="$Fz$", color='green', linestyle='-')
    # # ax2.plot(Fxx, label="$Fx$", color='blue', linestyle='--', alpha=.5)
    # # ax2.plot(Fyy, label="$Fy$", color='red', linestyle--', alpha=.5)
    # # ax2.plot(Fzz, label="$Fz$", color='green', linestyle='--', alpha=.5)
    # ax3.plot(Mx, label="$Mx$", color='blue', linestyle='-')
    # ax3.plot(My, label="$My$", color='red', linestyle='-')
    # ax3.plot(Mz, label="$Mz$", color='green', linestyle='-')
    #
    # leg1 = ax2.legend(loc='upper right', prop={'size': 14}, borderpad=0.2, handlelength=3)
    # leg1 = ax3.legend(loc='upper right', prop={'size': 14}, borderpad=0.2, handlelength=3)
    #
    # ax2.set_ylim(-5.5, 1)
    # ax3.set_ylim(-.02, .02)
    #plt.savefig('/Users/schimmt/Documents/Multis/studies/InstrumentedUltrasound/dat/Validation/PNGs/' + tdms_filename[0:3] + '_FT' + 'png')

    #

    cnt += 1

ax1.plot(-2,0,color='red', marker='o', label='Fx')
ax1.plot(-2,0,color='blue', marker='o', label='Fy')
ax1.plot(-2,0,color='green', marker='o', label='Fz')
ax1.plot(-2,0,color='black', marker='o', label='Resultant')
ax1.set_ylabel('Force (N)')
ax1.set_xlabel('Trial #')
leg1 = ax1.legend(loc='center right', prop={'size': 14}, borderpad=0.2, handlelength=3)

plt.savefig('/Users/schimmt/Documents/Multis/studies/InstrumentedUltrasound/dat/Validation/PNGs/' + '_FT')
a.write("%s +/-" % (str(round(np.average(resultant), 2))))
a.write("%s\n" % (str(round(np.std(resultant), 2))))

# plt.show()