'''
This script uses a registration XML to extract the necessary transformation matrices for data processing of OpenKnee Experimental Mechanics.

Inputs:
Registration xml with the appropriate locations of the segmented stls
Boolean to indicate whether or not to show registration figure: 1 show figure, anything else do not show figure

USAGE:
python convert_to_image_CS.py <Registration.xml> <boolean to show registration figure>

Outputs:
Saves .npz of transformation matrices and offsets. The .npz will be used in further analysis scripts since it is an easy format to import into Python. Saved in same directory as the registration xml.

Written By:
Erica Neumann, MS
Computational Biomodeling (CoBi) Core
Department of Biomedical Engineering, Cleveland Clinic
morrile2@ccf.org

'''

import os
import ConfigParser
import numpy as np
import math
from mayavi import mlab
import stl
import xml.etree.ElementTree as ET
from svd_FitHypersphere import fit_hypersphere
from tdms_contents import parseTDMSfile
import zipfile
import re
import sys

# To access any VTK object, we use 'tvtk', which is a Python wrapping of
# VTK replacing C++ setters and getters by Python properties and
# converting numpy arrays to VTK arrays when setting data.
from tvtk.api import tvtk
from tvtk.common import configure_input
from tvtk.tools import visual

def get_transformation_matrix(q1, q2, q3, q4, q5, q6):
    ''' Create transformation matrix to convert position information from world coordinate system to bone sensor coordinate system'''

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

    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)

    T[0, 3] = q1
    T[1, 3] = q2
    T[2, 3] = q3
    T[3, 3] = 1

    return T

def get_RB_transformation_matrix(q1, q2, q3, q4, q5, q6):
    ''' Create transformation matrix to convert position information from one RB coordinate system to another'''

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

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

    T[1, 0] = 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[1, 2] = -math.cos(q5) * math.sin(q4)

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

    T[0, 3] = q3*math.sin(q5)+q1
    T[1, 3] = -q3*math.sin(q4)*math.cos(q5)+q2*math.cos(q4)
    T[2, 3] = q3*math.cos(q4)*math.cos(q5)+q2*math.sin(q4)
    T[3, 3] = 1

    return T

def transform_to_Optotrak_CS(B1m, B1mr, n_pts):
    ''' Function to convert collected points (B1m) into bone optotrack coordinate system (B1mr)'''

    # Within this data set, there are 30 points, each with an x,y,and z coordinate
    # next we want to convert this data set into an array. The units for the axis
    # are all in m.

    # But first need to get rid of the quotations at the beginning and end of string.
    B1m = B1m.replace('"', '')
    B1msplit = B1m.split(" ")

    # Get rid of first two cells, and convert string into float
    corrected_list = B1msplit[2:]
    corrected_floatlist = map(float, corrected_list)

    ## matrix of the 30 points
    B1mCoord = np.asarray(corrected_floatlist)
    B1mCoord = B1mCoord.reshape(n_pts, -1)

    # Within this data set, there are 30 points, each with an x,y,z,roll, pitch and
    # yaw coordinate next we want to convert this data set into an array. The
    # units for the x,y, and z axis are in m and the roll, pitch, and yaw axis
    # are in rad.

    # But first need to get rid of the quotations at the beginning and end of string.
    B1mr = B1mr.replace('"', '')
    B1mrsplit = B1mr.split(" ")

    # Get rid of first two cells, and convert string into float
    corrected_list = B1mrsplit[2:]
    corrected_floatlist = map(float, corrected_list)

    ## matrix of the 30 points
    B1mrCoord = np.asarray(corrected_floatlist)
    B1mrCoord = B1mrCoord.reshape(n_pts, -1)

    ## Define coordinate transformations of data
    P1 = np.ones((4, 1))
    Coord1 = np.zeros((n_pts, 3))

    for i in range(0, n_pts):
        q1 = B1mrCoord[i, 0]
        q2 = B1mrCoord[i, 1]
        q3 = B1mrCoord[i, 2]
        q4 = B1mrCoord[i, 3]
        q5 = B1mrCoord[i, 4]
        q6 = B1mrCoord[i, 5]

        T1 = get_transformation_matrix(q1, q2, q3, q4, q5, q6)

        P1[0, 0] = B1mCoord[i, 0]
        P1[1, 0] = B1mCoord[i, 1]
        P1[2, 0] = B1mCoord[i, 2]

        invT1 = np.linalg.inv(T1)

        A = np.dot(invT1, P1) # Transform to bone Optotrak sensor coordinate system
        Coord1[i, 0] = A[0, 0]
        Coord1[i, 1] = A[1, 0]
        Coord1[i, 2] = A[2, 0]

    return Coord1, T1

def transform_sphere_points(B1m, B1mr):

    Coord1, T1 = transform_to_Optotrak_CS(B1m, B1mr, 30)

    ACoord1 = Coord1[0:10, 0:3]
    BCoord1 = Coord1[10:20, 0:3]
    CCoord1 = Coord1[20:30, 0:3]


    # Sphere fit for Rigid Body Collected Points (m), all three spheres. Set to NAN if points were not digitized.
    NAN = float('nan')
    try:
        B1mSphereA = fit_hypersphere(ACoord1, method="Pratt")
    except:
        B1mSphereA = [(0), (NAN, NAN, NAN)]

    try:
        B1mSphereB = fit_hypersphere(BCoord1, method="Pratt")
    except:
        B1mSphereB = [(0), (NAN, NAN, NAN)]

    try:
        B1mSphereC = fit_hypersphere(CCoord1, method="Pratt")
    except:
        B1mSphereC = [(0), (NAN, NAN, NAN)]

    return B1mSphereA, B1mSphereB, B1mSphereC

def get_stl_points(file):
    reader = tvtk.STLReader()
    reader.file_name = file
    reader.update()

    stl_points = reader.output.points.data.to_array()

    return stl_points

def SVD_3D(x, y):

    '''Defines transformation matrix for x defined in y (T_y_x).
    https://igl.ethz.ch/projects/ARAP/svd_rot.pdf was referenced for formulation of this function.'''

    assert len(x) == len(y)

    # Compute centroid of each dataset
    x_bar = np.mean(x, axis=0)
    y_bar = np.mean(y, axis=0)

    # Compute the centered vectors with d x n dimensions (d is the number of dimensions and n is the number of points)
    A = (x - x_bar).transpose()
    B = (y - y_bar).transpose()

    # Compute the covariance matrix
    C = np.dot(B, A.transpose())

    # Compute singular value decomposition
    P, D, QT = np.linalg.svd(C, full_matrices=True)

    R = np.dot(np.dot(P, np.diag([1, 1, np.linalg.det(np.dot(P, QT))])), QT)

    d = y_bar.transpose() - np.dot(R, x_bar.transpose())

    # Define transformation matrix for x defined in y (T_y_x)
    T = np.zeros((4, 4))
    T[0:3, 0:3] = R
    T[3, 3] = 1
    T[0:3, 3] = d

    return T

def add_STL_2_maya(v, m, c = (1,1,1)):
    reader = tvtk.STLReader()
    reader.file_name = m
    reader.update()

    mapper2 = tvtk.PolyDataMapper()
    configure_input(mapper2, reader.output)

    prop = tvtk.Property(opacity=.3, color=c)

    actor2 = tvtk.Actor(mapper=mapper2, property=prop)

    v.scene.add_actor(actor2)

def convertTOarray(data, r):

    data = data.replace("", '')
    data = data.split(" ")
    corrected_list = data[r:len(data)]
    corrected_list[-1] = corrected_list[-1][0:-1]
    data_float = map(float, corrected_list)

    return np.asarray(data_float)

def transformSTL_add_2_maya(v, stl_name, A, prop):
    '''Transformation matrix needs to be in m and rad'''
    stl_file = stl.Mesh.from_file(stl_name, calculate_normals=True)

    A[0][3] = A[0][3] * 1000
    A[1][3] = A[1][3] * 1000
    A[2][3] = A[2][3] * 1000

    stl_file.transform(A)
    stl_file.save('test_1.stl')
    reader2 = tvtk.STLReader()
    reader2.file_name = 'test_1.stl'
    reader2.update()

    mapper2 = tvtk.PolyDataMapper()
    configure_input(mapper2, reader2.output)

    # Add ultrasound probe to Mayavi scene
    actor2 = tvtk.Actor(mapper=mapper2, property=prop)
    v.scene.add_actor(actor2)

def Arrow_From_A_to_B(x1, y1, z1, x2, y2, z2, c):
    ar1=visual.arrow(x=x1, y=y1, z=z1, color = c)
    ar1.length_cone=0.4

    arrow_length=np.sqrt((x2-x1)**2+(y2-y1)**2+(z2-z1)**2)
    ar1.actor.scale=[arrow_length, arrow_length, arrow_length]
    ar1.pos = ar1.pos/arrow_length
    ar1.axis = [x2-x1, y2-y1, z2-z1]
    return ar1

def get_xyzrpw(T):
    x = T[0,3]
    y = T[1,3]
    z = T[2,3]
    r = math.atan2(T[2, 1], T[2, 2])
    p = math.atan2(-T[2,0], math.sqrt(T[2,1]**2+T[2,2]**2))
    w = math.atan2(T[1, 0], T[0, 0])
    return x, y, z, r, p, w

def get_line(A, B):
    length = math.sqrt((A[0]-B[0])**2+(A[1]-B[1])**2+(A[2]-B[2])**2)
    vector = (A - B)
    return float(length), vector[0], vector[1], vector[2]

def calculate_registration_error(f, points_1, points_2, labels):

    dist_points_1 = []
    dist_points_2 = []
    diff = []
    angles_diff = []
    diff_percent = []

    for p1 in range(len(points_1)-1):
        for p2 in range(p1+1, len(points_1)):
            dist_points_1.append([labels[p1][:5], labels[p2][:5], get_line(points_1[p1][1], points_1[p2][1])])
            dist_points_2.append([labels[p1][:5], labels[p2][:5], get_line(points_2[p1][1], points_2[p2][1])])
            diff.append(dist_points_1[-1][2][0]-dist_points_2[-1][2][0])
            diff_percent.append(diff[-1]/((dist_points_1[-1][2][0]+dist_points_2[-1][2][0])/2)*100)
            v1 = dist_points_1[-1][2][1:]
            v2 = dist_points_2[-1][2][1:]
            angles_diff.append(np.degrees(np.arccos(np.dot(v1,v2)/(np.linalg.norm(v1)*np.linalg.norm(v2)))))

    # Digitized sphere radii and coordinates
    f.write("\n")
    f.write("Digitized Points\n")
    f.write("\n")
    f.write("\tSphere Fitting\n")
    f.write("\tMarker\t\t\tRadius (mm)\t\tCoordinates [x,y,z] (mm)\n")

    for i, l in enumerate(points_1):
        f.write("\t" + labels[i][:5] + "\t\t\t%0.2f\t\t\t%0.4F, %0.4f, %0.4f\n" % (
        float(l[0]*1000), float(l[1][0]), float(l[1][1]), float(l[1][2])))

    f.write("\n")
    f.write("\tDistances between markers\n")
    f.write("\tMarker 1\t\tMarker 2\t\tDistance (mm)\n")
    for l in dist_points_1:
        f.write("\t" + l[0] + "\t\t\t" + l[1] + "\t\t\t%0.2f\n" % ((l[2][0])))

    # STL radii and Coordinates
    f.write("\n")
    f.write("STL Spheres - MRI\n")
    f.write("\n")
    f.write("\tSphere Fitting\n")
    f.write("\tMarker\t\t\tRadius (mm)\t\tCoordinates [x,y,z] (mm)\n")

    for i, l in enumerate(points_2):
        f.write("\t" + labels[i][:5] + "\t\t\t%0.2f\t\t\t%0.4F, %0.4f, %0.4f\n" % (
        float(l[0]), float(l[1][0]), float(l[1][1]), float(l[1][2])))

    f.write("\n")
    f.write("\tDistances between markers\n")
    f.write("\tMarker 1\t\tMarker 2\t\tDistance (mm)\n")
    for l in dist_points_2:
        f.write("\t" + l[0] + "\t\t\t" + l[1] + "\t\t\t%0.2f\n" % ((l[2][0])))

    f.write("\n")
    f.write("Differences between digitized and MRI registration\n")
    f.write("\n")
    f.write("Marker 1\tMarker 2\tDistance Diff\t\t\tAngle(deg)\n")
    for i in range(len(diff)):
        f.write(dist_points_1[i][0] + "\t\t" + dist_points_1[i][1] + "\t\t%0.2f mm [%0.2f %%]\t\t%0.2f\n" % (
            float(diff[i]), float(diff_percent[i]), angles_diff[i]))

    f.write("\n")
    f.write("Average \t\t\t%0.2f mm [%0.2f %%]\t\t%0.2f\n" % (
    np.average(np.abs(diff)), np.average(np.abs(diff_percent)), np.average(np.abs(angles_diff))))
    f.write("Maximum \t\t\t%0.2f mm [%0.2f %%]\t\t%0.2f\n" % (
    np.max(np.abs(diff)), np.max(np.abs(diff_percent)), np.max(np.abs(angles_diff))))

    f.close()


def main(args):

    """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    '''Read in information from registration xml'''
    """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    show_plot = int(args[-1])
    input_xml = ET.parse(args[-2])
    root = input_xml.getroot()

    knee_dir = root.find("Knee_directory").text
    knee_id = os.path.split(knee_dir)[1]
    patella_reg_dir = os.path.join(knee_dir, 'joint_mechanics-' + knee_id, 'TibiofemoralJoint', 'KinematicsKinetics', 'Configuration',
                 'Patella Registration')

    knee_side = root.find("Knee_side").text
    femur_stl_name = root.find("Bones").find("Femur").find("file").text
    tibia_stl_name = root.find("Bones").find("Tibia").find("file").text
    patella_stl_name = root.find("Bones").find("Patella").find("file").text
    bone_stls = [os.path.join(knee_dir, 'geometry-'+knee_id, tibia_stl_name), os.path.join(knee_dir, 'geometry-'+knee_id, femur_stl_name), os.path.join(knee_dir, 'geometry-'+knee_id, patella_stl_name)]

    """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    '''Perform patella registration - this is only needs to be completed once'''
    """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    error_report_pat = False
    # Extract the collected divot points for Patella registration in state file (reference optotrak sensor and optotrak world CS)
    config_pat = ConfigParser.RawConfigParser()
    if not config_pat.read(os.path.join(patella_reg_dir, 'State.cfg')):
        print os.path.join(patella_reg_dir, 'State.cfg')
        raise IOError, "Cannot load PF configuration file... Check path."

    # Extract divot points digitized in the reference optotrak coordinate system
    PAT_divot_pts = config_pat.get('Patella MRI Marker', 'Collected Points Rigid Body 2 (m)')
    PAT_divot_sens = config_pat.get('Patella MRI Marker', 'Collected Points Rigid Body 2 Position Sensor (m,rad)')
    PAT_divots_REFOPT, T_REFOPT_divots = transform_to_Optotrak_CS(PAT_divot_pts, PAT_divot_sens, 12)

    # Read CAD patella registration points
    CAD_file = open(os.path.join(patella_reg_dir, 'CAD_PT_DIMENSIONS.txt'), 'r')
    lines = CAD_file.readlines()
    PAT_divot_pts_CAD = []
    for i in range(9, 21, 1):
        PAT_divot_pts_CAD.append(np.array(re.split(', | --', lines[i])[0:3], dtype=float))
    PAT_divot_pts_CAD = np.array(PAT_divot_pts_CAD) / 1000  # convert to meters

    # Get centers of spheres for patella in CAD CS
    PAT_spheres_CAD = []
    for i in range(6, 9, 1):
        PAT_spheres_CAD.append(np.array(re.split(', | --', lines[i])[0:3], dtype=float))
    PAT_spheres_CAD = np.array(PAT_spheres_CAD) / 1000  # convert to meters

    # Extract patella sphere STL file names from XML
    PAT_marker_stls = []
    for e in root.find("Registration").find("PTR").iter():
        if '.stl' in e.text:
            PAT_marker_stls.append(os.path.join(knee_dir, 'geometry-' + knee_id, e.text))

    # Find center point and radius of each fiducial sphere (STL) and plot in Mayavi window.

    # Patella registration marker is digitized in a clockwise manner starting from the superior position.
    # The order is therefore different for left and right knees.
    if knee_side == "Right":
        PAT_digitizing_order = ['PTR-S', 'PTR-M', 'PTR-L']
    elif knee_side == 'Left':
        PAT_digitizing_order = ['PTR-S', 'PTR-L', 'PTR-M']

    PAT_stl_spheres = [None] * 3
    for m in PAT_marker_stls:
        r, pos = fit_hypersphere(get_stl_points(m))
        idx = PAT_digitizing_order.index(os.path.split(m)[1][11:16])
        pos = np.array(pos)  # convert to meters
        PAT_stl_spheres[idx] = r, pos

    # Extract data collected for registration between patella sensor and reference marker.
    tdms_registration = os.path.join(patella_reg_dir, 'Registration marker relative placement.tdms')
    data = parseTDMSfile(tdms_registration)

    data_pat = {}
    data_ref = {}
    for key, value in data[u'data'].items():
        if 'Patella Sensor' in key:
            data_pat[key[-1]] = np.average(value)
        elif 'Reference Marker' in key:
            data_ref[key[-1]] = np.average(value)

    # Convert to m and radians because (.tdms) file is in mm and degrees
    T_OPTWORLD_PATOS = get_transformation_matrix(
        *[data_pat[k] / 1000 for k in ('x', 'y', 'z')] + [np.radians(data_pat[k]) for k in ('r', 'p', 'w')])
    T_OPTWORLD_REFOPT = get_transformation_matrix(
        *[data_ref[k] / 1000 for k in ('x', 'y', 'z')] + [np.radians(data_ref[k]) for k in ('r', 'p', 'w')])

    # Define several more transformations
    T_REFOPT_CAD = SVD_3D(PAT_divot_pts_CAD, PAT_divots_REFOPT)
    T_REFOPT_PATOS = np.dot(np.linalg.inv(T_OPTWORLD_REFOPT), T_OPTWORLD_PATOS)
    T_I_CAD = SVD_3D(PAT_spheres_CAD, [x[1]/1000 for x in PAT_stl_spheres])

    T_I_PATOS = np.dot(np.dot(T_I_CAD, np.linalg.inv(T_REFOPT_CAD)), T_REFOPT_PATOS)

    """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    '''Registration for tibiofemoral and patellofemoral joint data'''
    """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    # Get all processed trials for the tibiofemoral joint
    processed_data_list_TF = []
    for root_kd, dirs, files in os.walk(os.path.join(knee_dir, 'joint_mechanics-'+knee_id, 'TibiofemoralJoint', 'KinematicsKinetics')):
        for file in files:
            if file.endswith('processed.tdms'):
                processed_data_list_TF.append(os.path.join(root_kd, file))
    processed_data_list_TF.sort()

    # Get all processed trials for the patellofemoral joint
    processed_data_list_PF = []
    for root_kd, dirs, files in os.walk(os.path.join(knee_dir, 'joint_mechanics-'+knee_id, 'PatellofemoralJoint', 'KinematicsKinetics')):
        for file in files:
            if file.endswith('processed.tdms'):
                processed_data_list_PF.append(os.path.join(root_kd, file))
    processed_data_list_PF.sort()

    # Loop through to perform registration for each processed tdms file
    for tdms_num, tdms_file in enumerate(processed_data_list_TF+processed_data_list_PF):
        test_dir = os.path.dirname(os.path.dirname(tdms_file))
        print test_dir

        # Check for zipped project file within configuration directory and unzip if found
        for item in os.listdir(os.path.join(test_dir, 'Configuration')):  # loop through items in dir
            if item.endswith('.sVprj'):  # check for ".zip" extension
                file_name = os.path.join(test_dir, 'Configuration', item)  # get full path of files
                zip_ref = zipfile.ZipFile(file_name)  # create zipfile object
                zip_ref.extractall(os.path.join(test_dir, 'Configuration'))  # extract file to dir
                zip_ref.close()  # close file

        # Read in state file to extract the transformation matrices and digitized points
        config = ConfigParser.RawConfigParser()
        if not config.read(os.path.join(test_dir, 'Configuration', 'State.cfg')):
            raise IOError, "Cannot load configuration file... Check path."

        # Transformation matrices for tibiofemoral joint coordinate systems within the state configuration file
        T_TIBOS_TIB = convertTOarray(config.get('Knee JCS', 'T_Sensor2_RB2'), 2).reshape(4,-1)
        T_FEMOS_FEMORI = convertTOarray(config.get('Knee JCS', 'T_Sensor1_RB1'), 2).reshape(4,-1)
        if "Patellofemoral" in tdms_file:
            T_PATOS_PAT = convertTOarray(config.get('Knee PTFJ', 'T_Sensor2_RB2'), 2).reshape(4,-1)

        # Extract offsets for both coordinate systems
        Offset_FEMORI_TIB = convertTOarray(config.get('Knee JCS', 'Position Offset (m,rad)'), 1)
        if "Patellofemoral" in tdms_file:
            Offset_FEMORI_PAT = convertTOarray(config.get('Knee PTFJ', 'Position Offset (m,rad)'), 1)

        try:
            if "Patellofemoral" not in tdms_file:
                TIB_fid_pts = config.get('MRI Fiducial Sphere Positions', 'Collected Points Rigid Body 2 (m)')
                TIB_fid_sens = config.get('MRI Fiducial Sphere Positions', 'Collected Points Rigid Body 2 Position Sensor (m,rad)')

                # Extract the collected points and sensor positions_TF for the Femur fiducial markers
                FEM_fid_pts = config.get('MRI Fiducial Sphere Positions', 'Collected Points Rigid Body 1 (m)')
                FEM_fid_sens = config.get('MRI Fiducial Sphere Positions', 'Collected Points Rigid Body 1 Position Sensor (m,rad)')
        except:
            print("...........Missing MRI Fiducial Spheres...........")
            pass

        # Get digitized sphere centers
        spheres_tib = transform_sphere_points(TIB_fid_pts, TIB_fid_sens)
        spheres_fem = transform_sphere_points(FEM_fid_pts, FEM_fid_sens)

        # Extract the order of the digitized spheres to ensure the STLs are read in correctly
        digitizing_order = []
        file = open(os.path.join(knee_dir, 'joint_mechanics-'+knee_id, 'TibiofemoralJoint', 'KinematicsKinetics', 'Configuration','ReadMe_Notes.txt'), 'r')
        lines = file.readlines()

        digitizing_order.append('TBR-' + lines[3].split(' ')[-1].replace('\n', '').replace('\r', ''))
        digitizing_order.append('TBR-' + lines[4].split(' ')[-1].replace('\n', '').replace('\r', ''))
        digitizing_order.append('TBR-' + lines[5].split(' ')[-1].replace('\n', '').replace('\r', ''))
        digitizing_order.append('FMR-' + lines[8].split(' ')[-1].replace('\n', '').replace('\r', ''))
        digitizing_order.append('FMR-' + lines[9].split(' ')[-1].replace('\n', '').replace('\r', ''))
        digitizing_order.append('FMR-' + lines[10].split(' ')[-1].replace('\n', '').replace('\r', ''))

        # Extract fiducial sphere STL file names from XML
        marker_stls = []
        for e in root.find("Registration").find("FMR").iter():
            if '.stl' in e.text:
                marker_stls.append([os.path.join(knee_dir, 'geometry-'+knee_id, e.text), 'FMR-' + e.tag])

        for e in root.find("Registration").find("TBR").iter():
            if '.stl' in e.text:
                marker_stls.append([os.path.join(knee_dir, 'geometry-'+knee_id, e.text), 'TBR-' + e.tag])

        # Find center point and radius of each fiducial sphere (STL) and plot in Mayavi window.
        stl_spheres = [None] * 6
        for m, t in marker_stls:
            r, pos = fit_hypersphere(get_stl_points(m))
            idx = digitizing_order.index(t)
            stl_spheres[idx] = r, pos

        # Format sphere positions_TF from digitizing and STLs into Numpy Arrays for SVD_3D function
        positions_TF_DIG = np.array(map(list, zip(*(spheres_tib+spheres_fem)))[1])
        positions_TF_STL = np.array(map(list, zip(*stl_spheres))[1])

        # Calculate the transformation matrices using SVD method
        T_I_TIBOS = SVD_3D(np.array(positions_TF_DIG)[0:3], np.array(positions_TF_STL)[0:3]/1000)
        T_I_FEMOS = SVD_3D(np.array(positions_TF_DIG)[3:6], np.array(positions_TF_STL)[3:6]/1000)

        # Extract the digitized patella bone fiducial markers
        if "Patellofemoral" in tdms_file:
            PAT_bone_pts, T_bone_pts = transform_to_Optotrak_CS(
                config.get('Knee PTFJ', 'Collected Points Rigid Body 2 (m)'),
                config.get('Knee PTFJ', 'Collected Points Rigid Body 2 Position Sensor (m,rad)'), 4)

        """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
        '''Registration error calculation'''
        """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
        # Transform digitized sphere centers to image coordinate system
        digitized_spheres_transformed = map(list, spheres_tib+spheres_fem)
        positions_TF_DIG_I = np.empty_like(positions_TF_DIG)

        for i, s in enumerate(digitized_spheres_transformed[0:3]):
            sphere = np.ones(4)
            sphere[0:3] = s[1]
            sphere = np.dot(T_I_TIBOS, sphere)
            positions_TF_DIG_I[i] = sphere[0:3]
            digitized_spheres_transformed[i][1] = sphere[0:3]*1000

        for i, s in enumerate(digitized_spheres_transformed[3:6]):
            sphere = np.ones(4)
            sphere[0:3] = s[1]
            sphere = np.dot(T_I_FEMOS, sphere)
            positions_TF_DIG_I[i + 3] = sphere[0:3]
            digitized_spheres_transformed[i+3][1] = sphere[0:3]*1000

        # Calculate registration errors
        if tdms_num ==0:
            f = open(os.path.join(os.path.dirname(args[-2]), knee_id+'_registration_errors_TIB.txt'), 'w')
            calculate_registration_error(f, digitized_spheres_transformed[0:3], stl_spheres[0:3], digitizing_order[0:3])

            f = open(os.path.join(os.path.dirname(args[-2]), knee_id+'_registration_errors_FEM.txt'), 'w')
            calculate_registration_error(f, digitized_spheres_transformed[3:6], stl_spheres[3:6], digitizing_order[3:6])

        if error_report_pat == False and "Patellofemoral" in tdms_file:
            CAD_spheres_transformed = []
            for i, p in enumerate(PAT_spheres_CAD):
                CAD_spheres_transformed.append([.006, np.dot(T_I_CAD, np.concatenate((p, [1])))[0:3]*1000])

            f = open(os.path.join(os.path.dirname(args[-2]), knee_id+'_registration_errors_PAT.txt'), 'w')
            f.write("** Note that sphere fitting was not performed for the patella since spheres were not physically digitized. The relationship within the CAD coordinate system was utilized to perform the transformation. In addition the radius reported for the digitized spheres was taken from the CAD model of the patella registration marker.**")
            calculate_registration_error(f, CAD_spheres_transformed, PAT_stl_spheres, PAT_digitizing_order)
            error_report_pat = True

        """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
        '''Plotting'''
        """"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
        if show_plot == 1:
            # Plotting into mayavi window (all transformed to image coordinate system)
            v = mlab.figure()

            # Add marker_stls (tibia and femur) to Mayavi window (WHITE)
            for m in marker_stls:
                transformSTL_add_2_maya(v, m[0], np.identity(4), tvtk.Property(opacity=0.8))

            # Transform digitized spheres (tibia and femur) into mri imaging coordinate system (I) and add points to Mayavi window.
            colors = [(1,0,0), (0,1,0), (0,0,1)]

            for i, s in enumerate(positions_TF_DIG_I[0:3]):
                sphere = np.ones(4)
                sphere[0:3] = s

                mlab.points3d(sphere[0] * 1000, sphere[1] * 1000, sphere[2] * 1000, 1, scale_factor=20, figure=v, opacity=0.1,
                          color=colors[i], mode='sphere')

            for i, s in enumerate(positions_TF_DIG_I[3:6]):
                sphere = np.ones(4)
                sphere[0:3] = s

                mlab.points3d(sphere[0] * 1000, sphere[1] * 1000, sphere[2] * 1000, 1, scale_factor=20, figure=v, opacity=0.1,
                          color=colors[i], mode='sphere')


            # Plot Femur, Tibia, and Patella STLs from neutral position
            # Tibia is bone stl #0, femur #1
            transformSTL_add_2_maya(v, bone_stls[0], np.identity(4), tvtk.Property(opacity=0.8))
            transformSTL_add_2_maya(v, bone_stls[1], np.identity(4), tvtk.Property(opacity=0.8))
            transformSTL_add_2_maya(v, bone_stls[2], np.identity(4), tvtk.Property(opacity=0.8))

            if "Patellofemoral" in tdms_file:

                # Fiducial bone markers (color based on location)
                colors = [(0, 0, 0), (1, 0, 0), (1, 1, 0), (0, 0, 1)]
                for i, p in enumerate(PAT_bone_pts):
                    p = np.dot(T_I_PATOS, np.append(p, [1]))
                    mlab.points3d(p[0] * 1000, p[1] * 1000, p[2] * 1000, color=colors[i], figure=v)

                colors = [(0, 0, 0), (1, 0, 0), (1, 1, 0)]
                for m in PAT_marker_stls:
                    idx = PAT_digitizing_order.index(os.path.split(m)[1][11:16])
                    add_STL_2_maya(v, m, c=colors[idx])

                # Patella sphere centroids from CAD (WHITE)
                for p in CAD_spheres_transformed:
                    p = p[1]
                    mlab.points3d(p[0], p[1], p[2], figure=v)

                # Patella divot points from CAD (RED)
                for p in PAT_divot_pts_CAD:
                    p = np.dot(T_I_CAD, np.concatenate((p, [1]))) * 1000
                    mlab.points3d(p[0], p[1], p[2], figure=v, color=(1, 0, 0))

                # Patella divot points from reference optotrak (GREEN)
                for p in PAT_divots_REFOPT:
                    p = np.dot(np.dot(T_I_CAD, np.linalg.inv(T_REFOPT_CAD)), np.concatenate((p, [1]))) * 1000
                    mlab.points3d(p[0], p[1], p[2], figure=v, color=(0, 1, 0))

                # Patella sphere centroids from segmentation (BLUE)
                for p in [x[1] for x in PAT_stl_spheres]:
                    mlab.points3d(p[0], p[1], p[2], figure=v, color=(0, 0, 1))
            mlab.show()

        if not os.path.exists(os.path.join(knee_dir, os.path.split(args[-2])[0], "PatellofemoralJoint")):
            os.makedirs(os.path.join(knee_dir, os.path.split(args[-2])[0], "PatellofemoralJoint"))
        if not os.path.exists(os.path.join(knee_dir, os.path.split(args[-2])[0], "TibiofemoralJoint")):
            os.makedirs(os.path.join(knee_dir, os.path.split(args[-2])[0], "TibiofemoralJoint"))

        if "Patellofemoral" in tdms_file:
            npz_name = os.path.join(knee_dir, os.path.split(args[-2])[0], "PatellofemoralJoint", os.path.split(tdms_file)[1][:-5])
            saved_data = np.savez(npz_name, T_TIBOS_TIB= T_TIBOS_TIB, T_PATOS_PAT=T_PATOS_PAT, T_I_FEMOS=T_I_FEMOS, T_I_TIBOS=T_I_TIBOS, T_FEMOS_FEMORI=T_FEMOS_FEMORI, T_I_PATOS=T_I_PATOS, Offset_FEMORI_TIB=Offset_FEMORI_TIB, Offset_FEMORI_PAT=Offset_FEMORI_PAT)
        else:
            npz_name = os.path.join(knee_dir, os.path.split(args[-2])[0], "TibiofemoralJoint", os.path.split(tdms_file)[1][:-5])
            saved_data = np.savez(npz_name, T_TIBOS_TIB= T_TIBOS_TIB, T_I_FEMOS=T_I_FEMOS, T_I_TIBOS=T_I_TIBOS, T_FEMOS_FEMORI=T_FEMOS_FEMORI, Offset_FEMORI_TIB=Offset_FEMORI_TIB)

if __name__ == "__main__":

    main(sys.argv)