import sys
import numpy as np
import os
import pandas as pd
try: import ConfigParser as cp
except: import configparser as cp
import copy
from lxml import etree as et
import matplotlib.pyplot as plt
import csv


def Tranform_axis_in_world(origin, axes):
    """given the origin and axes of a coordinate system, find the matrix to transfrom from world to that coordinate system"""

    axes = np.asarray(axes)

    RM = np.eye(4)
    RM[:3,3] = origin
    RM[:3, :3] = axes.T

    return RM

def Transform_A_in_B(origin_a, axes_a, origin_b, axes_b):

    A_in_world = Tranform_axis_in_world(origin_a, axes_a)
    B_in_world = Tranform_axis_in_world(origin_b, axes_b)

    world_in_B =  np.linalg.inv(B_in_world)

    A_in_B = np.matmul(world_in_B,A_in_world)

    return A_in_B


def find_model_offsets(ModelPropertiesXml):
    """ find the inital offsets in the tibiofemoral joint of the model"""

    # extract the origins and axes of the tibia and femur from the model properties file
    model_properties = et.parse(ModelPropertiesXml)
    ModelProperties = model_properties.getroot()
    Landmarks = ModelProperties.find("Landmarks")
    FMO = np.array(Landmarks.find("FMO").text.split(",")).astype(np.float)
    Xf_axis = np.array(Landmarks.find("Xf_axis").text.split(",")).astype(np.float)
    Yf_axis = np.array(Landmarks.find("Yf_axis").text.split(",")).astype(np.float)
    Zf_axis= np.array(Landmarks.find("Zf_axis").text.split(",")).astype(np.float)
    TBO =np.array( Landmarks.find("TBO").text.split(",")).astype(np.float)
    Xt_axis=np.array( Landmarks.find("Xt_axis").text.split(",")).astype(np.float)
    Yt_axis= np.array(Landmarks.find("Yt_axis").text.split(",")).astype(np.float)
    Zt_axis = np.array(Landmarks.find("Zt_axis").text.split(",")).astype(np.float)

    femur_axes = [Xf_axis,Yf_axis,Zf_axis]
    tibia_axes = [Xt_axis,Yt_axis,Zt_axis]

    # get the transformatrion matrix to get the tibia in the femur coordinate system
    T_in_F = Transform_A_in_B( TBO, tibia_axes, FMO, femur_axes) # this gives the transformation of the femur relative to the tibia coordnte system

    # extract the rotations and translations along the joints axes from the transformation matrix
    # use: https://simtk.org/plugins/moinmoin/openknee/Infrastructure/ExperimentationMechanics?action=AttachFile&do=view&target=Knee+Coordinate+Systems.pdf
    # page 6

    beta = np.arcsin(T_in_F[ 0, 2])
    alpha = np.arctan2(-T_in_F[ 1, 2], T_in_F[ 2, 2])
    gamma = np.arctan2(-T_in_F[ 0, 1], T_in_F[ 0, 0])

    ca = np.cos(alpha)
    sa = np.sin(alpha)
    cb = np.cos(beta)
    sb = np.sin(beta)

    b = (T_in_F[ 1, 3]*ca) + (T_in_F[ 2, 3]*sa)
    c = ((T_in_F[ 2, 3]*ca) - (T_in_F[ 1, 3]*sa))/ cb
    a = T_in_F[ 0, 3] - (c*sb)

    model_offsets = [a,b,c,alpha,beta,gamma]

    # convert alpha, beta, gamma to degrees
    model_offsets[3:6] = np.degrees(model_offsets[3:6])
    model_offsets = np.array(model_offsets)

    return model_offsets  


def change_kinematics_reporting(kinematics_expt):
    """flip the channels which are reported opposite to the model"""
    
    # changing the names and directions of the channels to align with our JCS definitions
    
    # read csv
    df = pd.read_csv(kinematics_expt)
  
    # experiment= "Posterior"
    # model= anterior translation (need to flip)
    df.rename({'Knee JCS Posterior [mm]' : 'Knee JCS Anterior Translation [mm]'}, axis = "columns", inplace = True)
    df['Knee JCS Anterior Translation [mm]'] = df['Knee JCS Anterior Translation [mm]'].apply(lambda x: float(x)* -1)
    
    # experiment  =  "Flexion"
    # model = extension (need to flip)
    df.rename({'Knee JCS Flexion [deg]' : 'Knee JCS Extension Rotation [deg]'}, axis = "columns", inplace = True)
    df['Knee JCS Extension Rotation [deg]'] = df['Knee JCS Extension Rotation [deg]'].apply(lambda x: float(x)* -1)
  
    # experiment = "Internal Rotation"
    # model = External (need to flip)
    df.rename({'Knee JCS Internal Rotation [deg]' : 'Knee JCS External Rotation [deg]'}, axis = "columns", inplace = True)
    df['Knee JCS External Rotation [deg]'] = df['Knee JCS External Rotation [deg]'].apply(lambda x: float(x)* -1)
     
    # change header to be more descriptive, data remains same
    df.rename({'Knee JCS Medial [mm]' : 'Knee JCS Medial Translation [mm]'}, axis = "columns", inplace = True)
    df.rename({'Knee JCS Superior [mm]' : 'Knee JCS Superior Translation [mm]'}, axis = "columns", inplace = True)
    df.rename({'Knee JCS Valgus [deg]' : 'Knee JCS Abduction Rotation [deg]'}, axis = "columns", inplace = True)
    
    print(df['Knee JCS External Rotation [deg]'])
    # writing into the file
    df.to_csv(kinematics_expt, index=False)


def change_kinetics_reporting(kinetics_expt):
    """flip the channels which are reported opposite to the model."""

    # changing the names and directios of the channels so the force descriptions are clearly defined in our coordinate systems
    
    
     # read csv
    df = pd.read_csv(kinetics_expt)
    
    #experiment = "Lateral Drawer"
    #model = medial (need to flip)
    df.rename({'JCS Load Lateral Drawer [N]' : 'External Tibia_x Load [N]'}, axis = "columns", inplace = True)
    df['External Tibia_x Load [N]'] = df['External Tibia_x Load [N]'].apply(lambda x: float(x)* -1)
    
    #experiment = "Distraction"
    #model = superior (need to flip - distractiong is a force in the inferior direction)
    df.rename({'JCS Load Distraction [N]' : 'External Tibia_z Load [N]'}, axis = "columns", inplace = True)
    df['External Tibia_z Load [N]'] = df['External Tibia_z Load [N]'].apply(lambda x: float(x)* -1)
    
    # experiment = "Varus Torque"
    # model = valgus (need to flip)
    df.rename({'JCS Load Varus Torque [Nm]' : 'External Tibia_y Moment [Nm]'}, axis = "columns", inplace = True)
    df['External Tibia_y Moment [Nm]'] = df['External Tibia_y Moment [Nm]'].apply(lambda x: float(x)* -1)
   
    # change header to be more descriptive, data remains same
    df.rename({'JCS Load Anterior Drawer [N]' : 'External Tibia_y Load [N]'}, axis = "columns", inplace = True)
    df.rename({'JCS Load Extension Torque [Nm]' : 'External Tibia_x Moment [Nm]'}, axis = "columns", inplace = True)
    df.rename({'JCS Load External Rotation Torque [Nm]' : 'External Tibia_z Moment [Nm]'}, axis = "columns", inplace = True)
    
    # writing into the file
    df.to_csv(kinetics_expt, index=False)
  
         
def T_tib_in_fem(a, b, c, alpha, beta, gamma):
    # this transformation matrix will give the position of the tibia in the femur coordinate system for each data point

    T_fem_tib = np.zeros((len(a),4,4))

    ca = np.cos(alpha)
    cb = np.cos(beta)
    cg = np.cos(gamma)
    sa = np.sin(alpha)
    sb = np.sin(beta)
    sg = np.sin(gamma)
    
   
    T_fem_tib[:,0,0] = np.multiply(cb,cg)
    T_fem_tib[:, 0, 1] = np.multiply(-cb,sg)
    T_fem_tib[:, 0, 2] = sb
    T_fem_tib[:, 0, 3] = np.multiply(c,sb) + a

    T_fem_tib[:, 1, 0] = np.multiply(np.multiply(sa,sb),cg) + np.multiply(ca,sg)
    T_fem_tib[:, 1, 1] = -np.multiply(np.multiply(sa,sb),sg) + np.multiply(ca, cg)
    T_fem_tib[:, 1, 2] = -np.multiply(sa, cb)
    T_fem_tib[:, 1, 3] = -np.multiply(c,np.multiply(sa, cb))+ np.multiply(b,ca)

    T_fem_tib[:, 2, 0] = -np.multiply(np.multiply(ca, sb), cg) + np.multiply(sa,sg)
    T_fem_tib[:, 2, 1] = np.multiply(np.multiply(ca,sb),sg) + np.multiply(sa,cg)
    T_fem_tib[:, 2, 2] = np.multiply(ca, cb)
    T_fem_tib[:, 2, 3] = np.multiply(c,np.multiply(ca, cb))+ np.multiply(b,sa)

    T_fem_tib[:, 3, 3] = 1.0

    return T_fem_tib


def kinetics_tibia_to_femur(kinetics_expt, kinematics_expt):
    """ convert the kinetics channel to report loads applied to femur in the tibia coordinate system.
    if a loading channel is given, the 'time' in the group will be updated too"""

    # initially, loads are reported as external tibia loads in the tibia coordinate system.
    # (a 'lateral' load is really along the tibia x axis not the JCS lateral axis)
    # tibia_kinetics_data_at_tibia_origin = np.asarray(kinetics_group.data)
    # to report the external forces applied to the femur, we need to invert the forces applied to the tibia
    # femur_kinetics_data_at_tibia_origin = -tibia_kinetics_data_at_tibia_origin
    
    df = pd.read_csv(kinetics_expt)
    # pd.DataFrame(df.values * -1, columns=df.columns, index=df.index) # may need to adjust for header
    arr = df.to_numpy()
    arr = arr * -1 
    
    # now we need to translate the forces and moments so they are applied at the femur origin instead of the tibia origin

    # find a,b,c,alpha,beta,gamma
    # note: the kinematics data was given in deg, mm. need to convert to rad, m
    # kinematics_data = np.asarray(kinematics_group.data)
    df1 =  pd.read_csv(kinematics_expt)
    arr1 = df1.to_numpy() 
    a = arr1[:,2]/1000.0
    b = arr1[:,3]/1000.0
    c = arr1[:,4]/1000.0
    
  
    alpha = np.radians(arr1[:,5])
    beta = np.radians(arr1[:,6])
    gamma = np.radians(arr1[:,7])

    # find the transformation of tibia in femur coordinte for each time point
    T = T_tib_in_fem(a, b, c, alpha, beta, gamma)

    # invert to get the position of femur in tibia CS
    T_fem_in_tib = np.linalg.inv(T)

    # vector from tibia origin to femur origin in tibia coordinate system at each time point
    vec_fmo = T_fem_in_tib[:, 0:3, 3] # units m
    vec_fmo = vec_fmo.T # transpose so it will be in the same shape as the data ie (axis, time point)
    # print(vec_fmo)
    # the moments at the femur origin are the moments at the tibia origin plus the forces at the tibia origin
    # cross with the moment arm (vector from femur origin to tibia origin in tibia cs, so negative of vec_fmo)
    
    loads = arr[:,2:5] 
    torques = arr[:,5:8] 
   
    # print(loads)
    # print(torques)
    # load_moments = np.multiply(loads, vec_fmo) # units same as intial kinetics torques
     # load_moments = np.cross(-vec_fmo.T, loads.T).T
    load_moments = np.cross(-vec_fmo.T, loads)
    torques = torques + load_moments # units same as intial kinetics torques

    # store all the results back in the data, channels
 
    df["External Tibia_x Load [N]"] = loads[:,0]
    df["External Tibia_y Load [N]"] = loads[:,1]
    df["External Tibia_z Load [N]"] = loads[:,2]
    df["External Tibia_x Moment [Nm]"] = torques[:,0]
    df["External Tibia_y Moment [Nm]"] = torques[:,1]
    df["External Tibia_z Moment [Nm]"] = torques[:,2]

    # change header for all 
  
    df.rename({'External Tibia_x Load [N]' : 'External Femur_x Load [N]'}, axis = "columns", inplace = True)
    df.rename({'External Tibia_y Load [N]' : 'External Femur_y Load [N]'}, axis = "columns", inplace = True)
    df.rename({'External Tibia_z Load [N]' : 'External Femur_z Load [N]'}, axis = "columns", inplace = True)
    df.rename({'External Tibia_x Moment [Nm]' : 'External Femur_x Moment [Nm]'}, axis = "columns", inplace = True)
    df.rename({'External Tibia_y Moment [Nm]' : 'External Femur_y Moment [Nm]'}, axis = "columns", inplace = True)
    df.rename({'External Tibia_z Moment [Nm]' : 'External Femur_z Moment [Nm]'}, axis = "columns", inplace = True)
  
 
   # writing into the file with updated csv name 
    filename = os.path.basename(kinetics_expt)
    loading_id = filename[0:15]
    # filename = os.path.join(path, filename)
    df.to_csv(loading_id +'_kinetics_in_TibiaCS.csv', index=False)
     


def process_data(kinematics_expt, Kinetics_expt, offsets_expt, ModelPropertiesXml):
    # this is for laxity processing only, will add passive flexion processing later
    
    # Apply expt offsets to kinematics data

    df = pd.read_csv(kinematics_expt)
    df_offsets = pd.read_csv(offsets_expt)
 
    print(df_offsets)
    print(df['Knee JCS Internal Rotation [deg]'])
    
    sum_column1 = df['Knee JCS Medial [mm]'].apply(lambda x: float(x) + df_offsets['Knee JCS Medial [mm]'])
    sum_column2 = df['Knee JCS Posterior [mm]'].apply(lambda x: float(x) + df_offsets['Knee JCS Posterior [mm]'])
    sum_column3 = df['Knee JCS Superior [mm]'].apply(lambda x: float(x) + df_offsets['Knee JCS Superior [mm]'])
    sum_column4 = df['Knee JCS Flexion [deg]'].apply(lambda x: float(x) +  df_offsets['Knee JCS Flexion [deg]'])
    sum_column5 = df['Knee JCS Valgus [deg]'].apply(lambda x: float(x) + df_offsets['Knee JCS Valgus [deg]'])
    sum_column6 = df['Knee JCS Internal Rotation [deg]'].apply(lambda x: float(x) + df_offsets['Knee JCS Internal Rotation [deg]'])
    
    
    df["Knee JCS Medial [mm]"] = sum_column1
    df["Knee JCS Posterior [mm]"] = sum_column2
    df["Knee JCS Superior [mm]"] = sum_column3
    df["Knee JCS Flexion [deg]"] = sum_column4
    df["Knee JCS Valgus [deg]"] = sum_column5
    df["Knee JCS Internal Rotation [deg]"] = sum_column6
    
    print(df['Knee JCS Internal Rotation [deg]'])
    df.to_csv(kinematics_expt, index=False)
     
    # =========
    
    # report the axes in the same direction as the model reporting :
    # positive directions are - extension, adduction (varus), external
    change_kinematics_reporting(kinematics_expt)
    change_kinetics_reporting(kinetics_expt)

    # ==========
    
    # report the kinetics as external loads applied to femur in tibia coordinate system
    kinetics_tibia_to_femur(kinetics_expt, kinematics_expt)
    
    # ==========

    # apply model offsets (subtract) - Note this is done AFTER changing the signs of the data to register with model outputs.
    model_offsets = find_model_offsets(ModelPropertiesXml)
    
    print (model_offsets)
    df1 = pd.read_csv(kinematics_expt)
    
    
    diff_column1 = df1['Knee JCS Medial Translation [mm]'].apply(lambda x: float(x) - model_offsets[0])
    diff_column2 = df1['Knee JCS Anterior Translation [mm]'].apply(lambda x: float(x) - model_offsets[1])
    diff_column3 = df1['Knee JCS Superior Translation [mm]'].apply(lambda x: float(x) - model_offsets[2])
    diff_column4 = df1['Knee JCS Extension Rotation [deg]'].apply(lambda x: float(x) - model_offsets[3])
    diff_column5 = df1['Knee JCS Abduction Rotation [deg]'].apply(lambda x: float(x) - model_offsets[4])
    diff_column6 = df1['Knee JCS External Rotation [deg]'].apply(lambda x: float(x) - model_offsets[5])
     
    df1["Knee JCS Medial Translation [mm]"] = diff_column1
    df1["Knee JCS Anterior Translation [mm]"] = diff_column2
    df1["Knee JCS Superior Translation [mm]"] = diff_column3
    df1["Knee JCS Extension Rotation [deg]"] = diff_column4
    df1["Knee JCS Abduction Rotation [deg]"] = diff_column5
    df1["Knee JCS External Rotation [deg]"] = diff_column6
    
    
    print(df1['Knee JCS External Rotation [deg]'])
    # ===========
    
    # save kinematics file with updated name
    
    filename = os.path.basename(kinematics_expt)
    loading_id = filename[0:15]
    df1.to_csv(loading_id +'_kinematics_in_JCS.csv', index=False)
    


if __name__ == "__main__":
    
    ModelPropertiesXml ='C:/oks/app/KneeHub/test/ModelProperties.xml'
    offsets_expt = 'C:/oks/app/KneeHub/test/kinematic_offsets.csv'
    
    
    
    kinematics_expt ='C:/oks/app/KneeHub/test/Laxity_0deg_VV2_kinematics_in_JCS_experiment.csv'
    kinetics_expt = 'C:/oks/app/KneeHub/test/Laxity_0deg_VV2_TibiaKinetics_in_TibiaCS_experiment.csv'
    
    
    process_data(kinematics_expt, kinetics_expt, offsets_expt, ModelPropertiesXml)

# update for multile files
# plot data
# save png files along with csv