'''
This script uses a registration XML to read bone and ligament STLs to plot ligament lengths and animate their insertion points from Open Knee Experimental Mechanics.

Inputs:
Registration xml with the appropriate locations of the segmented stls

USAGE:
python ligament_analysis.py <Registration.xml> <text file of csv data to analyze>

Outputs:
Plots of ligaments lengths saved as .png

Written By:
Rodrigo Lopez-Navarro
Department of Biomedical Engineering, Cleveland Clinic
'''
import numpy as np
import os
import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET
from tvtk.api import tvtk
from tvtk.common import configure_input
from traits.api import on_trait_change
import stl
from stl import mesh
from mayavi.modules.text import Text
import math
from argparse import Namespace
import pandas as pd
import sys
import csv
import vtk
import ConfigParser

def transformSTL(stl_name, A):
    '''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()

def insertion_SD(bone_polydata, lig_polydata,mult):
    '''Calculates insertion point from signed distance function'''
    implicitPolyDataDistance = vtk.vtkImplicitPolyDataDistance()
    bone_data = bone_polydata.GetOutput()
    implicitPolyDataDistance.SetInput(bone_data)
    lig_data = lig_polydata.GetOutput()

    #Get points from PolyData
    points = vtk.vtkPoints()
    a = np.zeros(3)
    for i in range(lig_data.GetNumberOfPoints()):
        lig_data.GetPoint(i, a)
        points.InsertNextPoint(a[0], a[1], a[2])

    # Add distances to each point
    signedDistances = vtk.vtkFloatArray()
    signedDistances.SetNumberOfComponents(1)
    signedDistances.SetName("SignedDistances")
    dist = []
    proximity_nodes = []
    np.asarray(proximity_nodes)

    # Evaluate the signed distance function at all of the grid points
    for pointId in range(points.GetNumberOfPoints()):
        p = points.GetPoint(pointId)
        signedDistance = implicitPolyDataDistance.EvaluateFunction(p)
        signedDistances.InsertNextValue(signedDistance)
        dist.append(signedDistance)
    np.asarray(dist)
    thresh = mult*np.ptp(dist)
    for i in range(points.GetNumberOfPoints()):
        if dist[i] < thresh:
            points.GetPoint(i, a)
            proximity_nodes.append(a)

    #Find centroid of proximity points and find closest point on the bone
    c = np.average(proximity_nodes, axis=0)
    cellId = vtk.reference(0)
    subId = vtk.reference(0)
    d = vtk.reference(0.0)
    cellLocator = vtk.vtkCellLocator()
    cellLocator.SetDataSet(bone_data)
    cellLocator.BuildLocator()
    ins_pt = [0, 0, 0]
    cellId = vtk.reference(0)
    subId = vtk.reference(0)
    d = vtk.reference(0.0)
    cellLocator.FindClosestPoint(c,ins_pt,cellId,subId,d)

    return ins_pt

def insertion_OBB(bone_polydata, lig_polydata):
    '''Calculates insertion point from OBB of ligament'''
    bone_data = bone_polydata.GetOutput()
    lig_data = lig_polydata.GetOutput()

    points = vtk.vtkPoints()
    a = np.zeros(3)
    for i in range(lig_data.GetNumberOfPoints()):
        lig_data.GetPoint(i, a)
        points.InsertNextPoint(a[0], a[1], a[2])

    #Compute OBB
    obbTree = vtk.vtkOBBTree()
    obbTree.SetDataSet(lig_data)
    obbTree.BuildLocator()
    corner = [0, 0, 0]
    max = [0, 0, 0]
    mid = [0, 0, 0]
    min = [0, 0, 0]
    size = [0, 0, 0]
    obbTree.ComputeOBB(points, corner, max, mid, min, size)
    corner = np.array(corner)
    max = np.array(max)
    mid = np.array(mid)
    min = np.array(min)
    c1 = corner
    c2 = corner + mid
    c3 = corner + min
    c4 = corner + min + mid
    c5 = corner + max
    c6 = corner + max + mid
    c7 = corner + max + min
    c8 = corner + max + min + mid
    plA = [c1, c2, c3, c4]
    plB = [c5, c6, c7, c8]
    midpA = np.average(plA, axis=0)
    midpB = np.average(plB, axis=0)

    #Find which center of box faces at the ends of ligament is closest to the bone
    cellId = vtk.reference(0)
    subId = vtk.reference(0)
    dx = vtk.reference(0.0)
    dy = vtk.reference(0.0)
    cellLocator = vtk.vtkCellLocator()
    cellLocator.SetDataSet(bone_data)
    cellLocator.BuildLocator()
    x = [0, 0, 0]
    y = [0, 0, 0]
    cellId = vtk.reference(0)
    subId = vtk.reference(0)
    dx = vtk.reference(0.0)
    dy = vtk.reference(0.0)
    cellLocator.FindClosestPoint(midpA,x,cellId,subId,dx)
    cellLocator.FindClosestPoint(midpB,y,cellId,subId,dy)
    if dx < dy:
        ins_pt = x
    else:
        ins_pt = y

    return ins_pt

def convert_left_2_right(T):
    T[0,1] *= -1
    T[0,2] *= -1
    T[0,3] *= -1
    T[1,0] *= -1
    T[2,0] *= -1
    return T

def get_RB_transformation_matrix(q1, q2, q3, q4, q5, q6):
    ''' Convert relative positions into transformation matrix'''

    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 insertion_point(bone_stl, ligament_stl,mult):
    '''Reads STL and computes SD and OBB insertion points'''
    reader1 = vtk.vtkSTLReader()
    reader1.SetFileName(bone_stl)
    reader1.Update()

    reader2 = vtk.vtkSTLReader()
    reader2.SetFileName(ligament_stl)
    reader2.Update()

    x = insertion_SD(reader1,reader2,mult)
    y = insertion_OBB(reader1,reader2)
    insertionpoint_sd = np.array([x[0],x[1],x[2],1])
    insertionpoint_OBB = np.array([y[0], y[1], y[2], 1])
    return insertionpoint_sd, insertionpoint_OBB

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 methodcheck(method, SD, OBB):
    if method == 'SD':
        point = SD
    elif method == 'OBB':
        point = OBB
    else:
        raise IOError, "Incorrect method selection... Check method selection xml"
    return point
def main(args):
    input_xml = ET.parse(args[-1])
    root = input_xml.getroot()

    M_A_ACL = root.find("Femur").find("ACL").find("file").text
    M_A_PCL = root.find("Femur").find("PCL").find("file").text
    M_A_MCL = root.find("Femur").find("MCL").find("file").text
    M_A_LCL = root.find("Femur").find("LCL").find("file").text

    M_B_ACL = root.find("Tibia").find("ACL").find("file").text
    M_B_PCL = root.find("Tibia").find("PCL").find("file").text
    M_B_MCL = root.find("Tibia").find("MCL").find("file").text
    M_B_LCL = root.find("Tibia").find("LCL").find("file").text

    input_xml1 = ET.parse(args[-3])
    root = input_xml1.getroot()

    knee_dir = root.find("Knee_directory").text
    knee_id = os.path.split(knee_dir)[1]
    test_prot = root.find("Test_protocol").text

    femur_stl_name = root.find("Bones").find("Femur").find("file").text
    tibia_stl_name = root.find("Bones").find("Tibia").find("file").text
    fibula_stl_name = root.find("Bones").find("Fibula").find("file").text
    patella_stl_name = root.find("Bones").find("Patella").find("file").text
    ACL_stl_name = root.find("Ligaments").find("ACL").find("file").text
    PCL_stl_name = root.find("Ligaments").find("PCL").find("file").text
    MCL_stl_name = root.find("Ligaments").find("MCL").find("file").text
    LCL_stl_name = root.find("Ligaments").find("LCL").find("file").text

    bone_stls = [os.path.join(knee_dir, 'mri-'+knee_id, femur_stl_name), os.path.join(knee_dir, 'mri-'+knee_id, tibia_stl_name), os.path.join(knee_dir, 'mri-'+knee_id, patella_stl_name), os.path.join(knee_dir, 'mri-'+knee_id, fibula_stl_name)]
    ligament_stls = [os.path.join(knee_dir, 'mri-' + knee_id, ACL_stl_name), os.path.join(knee_dir, 'mri-' + knee_id, PCL_stl_name), os.path.join(knee_dir, 'mri-' + knee_id, MCL_stl_name), os.path.join(knee_dir, 'mri-' + knee_id, LCL_stl_name)]

    knee_side = root.find("Knee_side").text

    data_folder = os.path.join(os.path.dirname(args[-2]), test_prot, 'Data', "Ligament Analysis")
    if not os.path.exists(data_folder):
        os.makedirs(data_folder)

    R = dict(np.load(args[-3][:-4]+'.npz'))

    #Computer Insertion points, A and B are insertion points of ligament at each bone
    #Last input is the "proximity" threshold as a percentage. Ie. Points which are below X% of maximum distance between ligament and bone will be selected
    A_ACL_SD, A_ACL_OBB = insertion_point(bone_stls[0], ligament_stls[0], 0.08)
    B_ACL_SD, B_ACL_OBB = insertion_point(bone_stls[1], ligament_stls[0], 0.04)
    A_PCL_SD, A_PCL_OBB = insertion_point(bone_stls[0], ligament_stls[1], 0.05)
    B_PCL_SD, B_PCL_OBB = insertion_point(bone_stls[1], ligament_stls[1], 0.05)
    A_MCL_SD, A_MCL_OBB = insertion_point(bone_stls[0], ligament_stls[2], 0.08)
    B_MCL_SD, B_MCL_OBB = insertion_point(bone_stls[1], ligament_stls[2], 0.05)
    A_LCL_SD, A_LCL_OBB = insertion_point(bone_stls[0], ligament_stls[3], 0.05)
    B_LCL_SD, B_LCL_OBB = insertion_point(bone_stls[3], ligament_stls[3], 0.05)

    A_ACL = methodcheck(M_A_ACL, A_ACL_SD, A_ACL_OBB)
    B_ACL = methodcheck(M_B_ACL, B_ACL_SD, B_ACL_OBB)
    A_PCL = methodcheck(M_A_PCL, A_PCL_SD, A_PCL_OBB)
    B_PCL = methodcheck(M_B_PCL, B_PCL_SD, B_PCL_OBB)
    A_MCL = methodcheck(M_A_MCL, A_MCL_SD, A_MCL_OBB)
    B_MCL = methodcheck(M_B_MCL, B_MCL_SD, B_MCL_OBB)
    A_LCL = methodcheck(M_A_LCL, A_LCL_SD, A_LCL_OBB)
    B_LCL = methodcheck(M_B_LCL, B_LCL_SD, B_LCL_OBB)

    print 'Femur attachment'
    print "ACL: {}" .format(A_ACL)
    print "PCL: {}" .format(A_PCL)
    print "MCL: {}" .format(A_MCL)
    print "LCL: {}" .format(A_LCL)

    print 'Tibia attachment'
    print "ACL: {}" .format(B_ACL)
    print "PCL: {}" .format(B_PCL)
    print "MCL: {}" .format(B_MCL)
    print "LCL: {}" .format(B_LCL)

    if knee_side == "Left":
        R['T_TIBOS_TIB'] = convert_left_2_right(R["T_TIBOS_TIB"])
        R["T_FEMOS_FEMORI"] = convert_left_2_right(R["T_FEMOS_FEMORI"])
        R["T_PATOS_PAT"] = convert_left_2_right(R["T_PATOS_PAT"])

    f = open(args[-2])
    csv_name = f.read().replace('\n', '')
    csv_file = os.path.join(os.path.dirname(args[-2]), csv_name)
    f.close()
    df = pd.read_csv(csv_file)
    M = df['Knee JCS Medial  [mm]']
    P = df['Knee JCS Posterior  [mm]']
    S = df['Knee JCS Superior  [mm]']
    F = df['Knee JCS Flexion  [deg]']
    V = df['Knee JCS Valgus  [deg]']
    IR = df['Knee JCS Internal Rotation  [deg]']
    time = df['Extracted Time Points [s]']

    config = ConfigParser.RawConfigParser()
    if not config.read(
            os.path.join(os.path.dirname(args[-2]), csv_name.split('/')[0], 'Configuration', 'State.cfg')):
        raise IOError, "Cannot load configuration file... Check path."
    Offset_FEMORI_TIB = convertTOarray(config.get('Knee JCS', 'Position Offset (m,rad)'), 1)

    if knee_side == "Left":
        M = -1 * np.array(M) / 1000 - Offset_FEMORI_TIB[0]
        P = np.array(P) / 1000 + Offset_FEMORI_TIB[1]
        S = np.array(S) / 1000 + Offset_FEMORI_TIB[2]
        F = np.radians(np.array(F)) + Offset_FEMORI_TIB[3]
        V = np.radians(-1 * np.array(V)) - Offset_FEMORI_TIB[4]
        IR = np.radians(-1 * np.array(IR)) - Offset_FEMORI_TIB[5]

    else:
        M = np.array(M) / 1000 - Offset_FEMORI_TIB[0]
        P = np.array(P) / 1000 + Offset_FEMORI_TIB[1]
        S = np.array(S) / 1000 + Offset_FEMORI_TIB[2]
        F = np.radians(np.array(F)) + Offset_FEMORI_TIB[3]
        V = np.radians(np.array(V)) - Offset_FEMORI_TIB[4]
        IR = np.radians(np.array(IR)) - Offset_FEMORI_TIB[5]

    T_I_TIB = np.dot(R["T_I_TIBOS"], R["T_TIBOS_TIB"])

    positions_TF = np.empty([np.size(IR), 4, 4])

    num = np.size(M)
    length_ACL = np.zeros(num)
    length_PCL = np.zeros(num)
    length_MCL = np.zeros(num)
    length_LCL = np.zeros(num)

    for ii in range(num):
        T_FEMORI_TIB = get_RB_transformation_matrix(M[ii], P[ii], S[ii], F[ii], V[ii],
                                                    IR[
                                                        ii])  # Transform between FEM CS and TIB CS Relative angles/positions b/w bones

        positions_TF[ii] = np.dot(T_I_TIB, np.linalg.inv(
            np.dot(np.dot(R["T_I_FEMOS"], R["T_FEMOS_FEMORI"]),
                   T_FEMORI_TIB)))  # Define the positions of the femur in the image CS

        # IMPORTANT!!!! This function converts the positions_TF transformation matrix into millimeters instead of meters. Therefore this line must run prior                to taking the dot product for the point transformation since the points are given in mm!
        transformSTL(bone_stls[0], positions_TF[ii])

        ## Change HERE whether to use OBB or SD for length calculation
        # ACL points transformations and animation process
        A_CS_ACL = np.dot(positions_TF[ii], A_ACL)
        length_ACL[ii] = np.sqrt(np.sum((A_CS_ACL - B_ACL) ** 2))

        # PCL points transformations and animation process
        A_CS_PCL = np.dot(positions_TF[ii], A_PCL)
        length_PCL[ii] = np.sqrt(np.sum((A_CS_PCL - B_PCL) ** 2))

        # MCL points transformations and animation process
        A_CS_MCL = np.dot(positions_TF[ii], A_MCL)
        length_MCL[ii] = np.sqrt(np.sum((A_CS_MCL - B_MCL) ** 2))

        # LCL points transformations and animation process
        A_CS_LCL = np.dot(positions_TF[ii], A_LCL)
        length_LCL[ii] = np.sqrt(np.sum((A_CS_LCL - B_LCL) ** 2))

    ACL_DL = np.subtract(length_ACL,np.average(length_ACL))
    PCL_DL = np.subtract(length_PCL,np.average(length_PCL))
    MCL_DL = np.subtract(length_MCL,np.average(length_MCL))
    LCL_DL = np.subtract(length_LCL,np.average(length_LCL))

    plt.figure()
    plt.plot(time, ACL_DL)
    plt.plot(time, PCL_DL)
    plt.plot(time, MCL_DL)
    plt.plot(time, LCL_DL)
    plt.legend(['ACL', 'PCL', 'MCL', 'LCL'], loc='upper right')
    plt.ylabel('Length change [mm]')
    plt.xlabel('Time [s]')
    plt.grid(True)
    plt.suptitle(csv_name.split('/')[0])
    plt.savefig(data_folder + '/' + 'Ligament_length_change.png')

    plt.figure()

    plt.subplot(211)
    plt.plot(time, length_ACL)
    plt.ylabel('Length [mm]')
    plt.xlabel('Time [s]')
    plt.legend(['ACL'], loc='upper right')
    plt.grid(True)

    plt.subplot(212)
    plt.plot(time, length_PCL)
    plt.ylabel('Length [mm]')
    plt.xlabel('Time [s]')
    plt.legend(['PCL'], loc='upper right')
    plt.grid(True)
    plt.suptitle(csv_name.split('/')[0])
    plt.savefig(data_folder + '/' + 'AP_length_absolute.png')

    plt.figure()

    plt.subplot(211)
    plt.plot(time, length_MCL)
    plt.ylabel('Length [mm]')
    plt.xlabel('Time [s]')
    plt.legend(['MCL'], loc='upper right')
    plt.grid(True)

    plt.subplot(212)
    plt.plot(time, length_LCL)
    plt.ylabel('Length [mm]')
    plt.xlabel('Time [s]')
    plt.legend(['LCL'], loc='upper right')
    plt.grid(True)
    plt.suptitle(csv_name.split('/')[0])
    plt.savefig(data_folder + '/' + 'ML_length_absolute.png')

if __name__ == "__main__":

    main(sys.argv)

    # main(['callfunction', '/home/lopezr3/Documents/CC/Registration/oks003_registration_01.xml', '/home/lopezr3/Documents/CC/oks003/joint_mechanics-oks003/TibiofemoralJoint/KinematicsKinetics/oks003_TF_csv.txt', '/home/lopezr3/Documents/CC/oks003/joint_mechanics-oks003/TibiofemoralJoint/KinematicsKinetics/oks003_method_selection.xml'])