'''
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
Text file with list of desired tdms files to analyze

USAGE:
python resampling_plotting.py <Registration.xml> <text file of tdms data to analyze>

Outputs:
Process data saved as .csv file
Animation saved as .mp4 video file
Graphs of raw and processed data saved in sub-directory (named with tdms filename)
    .png and .svg formats

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

'''

import numpy as np
from tdms_contents import parseTDMSfile
import os
import matplotlib.pyplot as plt
from mayavi import mlab
import xml.etree.ElementTree as ET
from tvtk.tools import visual
from tvtk.api import tvtk
from tvtk.common import configure_input
from traits.api import on_trait_change
import stl
from mayavi.modules.text import Text
import math
from argparse import Namespace
import moviepy.editor as mpy
import moviepy.video as mpv
import tdms_plotting
import pandas as pd
import sys
import ConfigParser
import zipfile
from PIL import Image

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 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 transformSTL_get_actor(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)
    return 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 main(args):

    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]

    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)]

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

    f = open(args[-1])
    tdms_list = f.read().splitlines()

    for tdms_name in tdms_list:
        # try:
        print '---------------'
        print tdms_name
        tdms_file = os.path.join(os.path.dirname(args[-1]), tdms_name)
        data = parseTDMSfile(tdms_file)

        if 'Patellofemoral' in tdms_file:
            reg_data_file = os.path.join(os.path.split(args[-2])[0], 'PatellofemoralJoint', os.path.split(tdms_file)[1][:-4] + 'npz')
            patella = True
        else:
            reg_data_file =  os.path.join(os.path.split(args[-2])[0], 'TibiofemoralJoint',
                               os.path.split(tdms_file)[1][:-4] + 'npz')
            patella = False
        print reg_data_file
        R = dict(np.load(reg_data_file))

        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"])
            if patella:
                R["T_PATOS_PAT"] = convert_left_2_right(R["T_PATOS_PAT"])

        M = np.array(data[u'State.Knee JCS'][u'Knee JCS Medial']) / 1000 + R['Offset_FEMORI_TIB'][0]
        P = np.array(data[u'State.Knee JCS'][u'Knee JCS Posterior']) / 1000 + R['Offset_FEMORI_TIB'][1]
        S = np.array(data[u'State.Knee JCS'][u'Knee JCS Superior']) / 1000 + R['Offset_FEMORI_TIB'][2]
        F = np.radians(np.array(data[u'State.Knee JCS'][u'Knee JCS Flexion'])) + R['Offset_FEMORI_TIB'][3]
        V = np.radians(np.array(data[u'State.Knee JCS'][u'Knee JCS Valgus'])) + R['Offset_FEMORI_TIB'][4]
        IR = np.radians(np.array(data[u'State.Knee JCS'][u'Knee JCS Internal Rotation'])) + R['Offset_FEMORI_TIB'][5]

        if patella:
            M_pat = np.array(data[u'State.Knee PTFJ'][u'Knee PTFJ Medial']) / 1000 + R['Offset_FEMORI_PAT'][0]
            P_pat = np.array(data[u'State.Knee PTFJ'][u'Knee PTFJ Posterior']) / 1000 + R['Offset_FEMORI_PAT'][1]
            S_pat = np.array(data[u'State.Knee PTFJ'][u'Knee PTFJ Superior']) / 1000 + R['Offset_FEMORI_PAT'][2]
            F_pat = np.radians(np.array(data[u'State.Knee PTFJ'][u'Knee PTFJ Flexion'])) + R['Offset_FEMORI_PAT'][3]
            V_pat = np.radians(np.array(data[u'State.Knee PTFJ'][u'Knee PTFJ Valgus'])) + R['Offset_FEMORI_PAT'][4]
            IR_pat = np.radians(np.array(data[u'State.Knee PTFJ'][u'Knee PTFJ Internal Rotation'])) + R['Offset_FEMORI_PAT'][5]

        if knee_side == 'Left':
            M*=-1
            V*=-1
            IR*=-1
            if patella == True:
                M_pat*=-1
                V_pat*=-1
                IR_pat*=-1

        v = mlab.figure()
        if patella:
            T_I_PAT = np.dot(R["T_I_PATOS"], R["T_PATOS_PAT"])
        T_I_TIB = np.dot(R["T_I_TIBOS"], R["T_TIBOS_TIB"])

        idx = tdms_plotting.tdms_contents([[tdms_file]])[0]
        df = pd.read_csv(tdms_file[:-4]+'csv')
        time = df['Extracted Time Points [s]']

        positions_TF = np.empty([len(IR[idx]), 4, 4])
        positions_PF = np.empty([len(IR[idx]), 4, 4])
        actors = []

        act = transformSTL_get_actor(v, bone_stls[0], np.identity(4), tvtk.Property(opacity=1, color=(1, 1, 1)))
        v.scene.add_actor(act)

        for ii in range(len(M[idx])):
            T_FEMORI_TIB = get_RB_transformation_matrix(M[idx][ii], P[idx][ii], S[idx][ii], F[idx][ii], V[idx][ii], IR[idx][
                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

            try:
                act = transformSTL_get_actor(v, bone_stls[1], positions_TF[ii], tvtk.Property(opacity=1, color=(1, 1, 1)))

                if patella:
                    T_FEMORI_PAT = get_RB_transformation_matrix(M_pat[ii], P_pat[ii], S_pat[ii], F_pat[ii], V_pat[ii], IR_pat[
                        ii])  # Transform between FEM CS and PAT CS Relative angles/positions_TF b/w bones
                    if np.isnan(T_FEMORI_PAT[0][0]):
                        actors.append([act])
                        pass
                    else:
                        positions_PF[ii] = np.dot(np.dot(T_I_TIB, np.dot(np.linalg.inv(T_FEMORI_TIB), T_FEMORI_PAT)),
                                                  np.linalg.inv(T_I_PAT))  # Define the positions_TF of the femur in the image CS

                        act_Pat = transformSTL_get_actor(v, bone_stls[2], positions_PF[ii],
                                                         tvtk.Property(opacity=1, color=(1, 1, 1)))
                        print positions_PF[ii]
                        actors.append([act, act_Pat])
                else:
                    actors.append([act])
            except:
                "Missing data at index {}".format(ii)
                pass

        def make_frame(t):
            engine = mlab.get_engine()
            scene = engine.current_scene
            timetext = Text()
            if 'actor_old' in globals():
                v.scene.remove_actor(actor_old[0])
                v.scene.remove_actor(actor_old[1])
                if patella == True:
                    v.scene.remove_actor(actor_old[1])

            t = int(t)
            for act_i in actors[t]:
                v.scene.add_actor(act_i)

            global actor_old

            timetext.text = 'Time = %f sec' % (float(time[t]))
            v.scene.add_actor(timetext.actor)
            actor_old = [actors[t], timetext.actor]

            if 'view' not in globals():
                global view
                view = mlab.view(distance='auto')
            else:
                mlab.view(*view)
            t = (t + 1) % len(time)
            return mlab.screenshot(antialiased=True)

        if len(time)>5:
            animation = mpy.VideoClip(make_frame, duration=len(time))
            animation =mpv.fx.all.accel_decel(animation, new_duration=5)
            animation.write_videofile(tdms_file[:-4] + 'mp4', fps=len(time)/5)
            mlab.close(all=True)
        else:
            ss = Image.fromarray(make_frame(1), 'RGB')
            ss.save(tdms_file[:-4] + 'png')


        # @mlab.show
        # @mlab.animate(delay=100)
        # @on_trait_change('scene.activated')
        # def anim():
        #     """Animate the b1 box."""
        #     t = 0
        #     engine = mlab.get_engine()
        #     scene = engine.current_scene
        #     timetext = Text()
        #     timetext.text = 'Time = %f sec' % (float(time[t]))
        #     v.scene.add_actor(timetext.actor)
        #     # timetext = mlab.text(0.7, 0.01, 'Time = %f sec' % (float(time[t]) / 1000.0), figure = scene)
        #     while 1:
        #         if 'actor_old' in globals():
        #             v.scene.remove_actor(actor_old[0])
        #             if patella == True:
        #                 v.scene.remove_actor(actor_old[1])
        #
        #         try:
        #             for act_i in actors[t]:
        #                 v.scene.add_actor(act_i)
        #
        #             global actor_old
        #             actor_old = [actors[t], timetext.actor]
        #         except:
        #             pass
        #         timetext.text = 'Time = %f sec' % (float(time[t]))
        #         t = (t + 1) % len(time)
        #
        #         yield
        #
        # anim()

        # del globals()['actor_old']
        # del globals()['view']
        # except:
        #     print "Error with tdms file: "
        #     print tdms_file


if __name__ == "__main__":

    main(sys.argv)