import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
import numpy as np
import math
import os

def get_transformation_matrix(q4, q5, q6):
    ''' Transform from optotrak global coordinates to optotrak position sensor coordinates'''

    T = np.zeros((3, 3))

    T[0, 0] = math.cos(q6) * math.cos(q5)
    T[1, 0] = math.sin(q6) * math.cos(q5)
    T[2, 0] = -math.sin(q5)

    T[0, 1] = math.cos(q6) * math.sin(q5) * math.sin(q4) - math.sin(q6) * math.cos(q4)
    T[1, 1] = math.sin(q6) * math.sin(q5) * math.sin(q4) + math.cos(q6) * math.cos(q4)
    T[2, 1] = math.cos(q5) * math.sin(q4)

    T[0, 2] = math.cos(q6) * math.sin(q5) * math.cos(q4) + math.sin(q6) * math.sin(q4)
    T[1, 2] = math.sin(q6) * math.sin(q5) * math.cos(q4) - math.cos(q6) * math.sin(q4)
    T[2, 2] = math.cos(q5) * math.cos(q4)

    return T


subj_ID = 'CMULTIS008-1'
root_path = os.path.join('/home/morrile2/Documents/MULTIS Data/MULTIS_invitro', subj_ID, 'Registration')
segments = ['UpperLeg', 'UpperArm', 'LowerLeg', 'LowerArm']
segments = ['UpperLeg']

text = open(os.path.join(root_path, subj_ID + 'orientCheck.txt'), 'w')

for seg in segments:
    filename = subj_ID + '_'+ ''.join([c for c in seg if c.isupper()]) + '_US_CT.xml'
    filepath = os.path.join(root_path, filename)

    doc = ET.parse(filepath)

    root = doc.getroot()

    regions = ['CentralAnterior', 'CentralPosterior']
    for reg in regions:
        print subj_ID + ' ' + seg + ' ' + reg
        text.write(subj_ID + ' ' + seg + ' ' + reg + '\n')
        ind = root.find(seg).find(reg).find('Indentation')
        locs = ind.findall('USPosition')

        z_axis = np.array([0,0,1])
        roll = []
        pitch = []
        yaw = []
        vect = []
        x = []
        y = []
        z = []

        for pos in locs:
            roll.append(np.degrees(float(pos.find("roll").get('value'))))
            pitch.append(np.degrees(float(pos.find("pitch").get('value'))))
            yaw.append(np.degrees(float(pos.find("yaw").get('value'))))
            x.append(float(pos.find("x").get("value")))
            y.append(float(pos.find("y").get("value")))
            z.append(float(pos.find("z").get("value")))

        x = np.array(x)*1000
        y = np.array(y)*1000
        z = np.array(z)*1000

        for i in range(len(roll)):
            if roll[i]-roll[0] > 180:
                roll[i] = roll[i]-360
            elif roll[i]-roll[0] < -180:
                roll[i] = roll[i] + 360
            if pitch[i]-pitch[0] > 180:
                pitch[i] = pitch[i]-360
            elif pitch[i] - pitch[0] < -180:
                pitch[i] = pitch[i] + 360
            if yaw[i]-yaw[0] > 180:
                yaw[i] = yaw[i]-360
            elif yaw[i] - yaw[0] < -180:
                yaw[i] = yaw[i] + 360

        plt.figure()
        plt.plot(roll-roll[0], marker = '*')
        plt.plot(pitch-pitch[0], marker = '*')
        plt.plot(yaw-yaw[0], marker = '*')
        plt.legend(['Roll', 'Pitch', 'Yaw'])
        plt.title(subj_ID + ' ' + seg + ' ' + reg)

        text.write('Roll %f degrees \n' %(max((abs(roll - roll[0])))))
        text.write('Pitch %f degrees \n' %(max((abs(pitch-pitch[0])))))
        text.write('Yaw %f degrees \n' %(max((abs(yaw-yaw[0])))))
        text.write('============================================\n')
        print 'Roll', max((abs(roll - roll[0])))
        print 'Pitch', max((abs(pitch-pitch[0])))
        print 'Yaw', max((abs(yaw-yaw[0])))
        print '============================================'

        # plt.savefig(filepath[0:-4]+'_'+''.join([c for c in reg if c.isupper()])+'_orientCheck.png')

        plt.figure()
        plt.plot(x-x[0], marker='*')
        plt.plot(y-y[0], marker='*')
        plt.plot(z-z[0], marker='*')
        plt.legend(['x', 'y', 'z'])
        plt.title(subj_ID + ' ' + seg + ' ' + reg)
        plt.show()

text.close()