import sys
from nptdms import TdmsFile

def parseTDMSfile(tdmsFileName):
    """
    This function was modified from tdms_contents.py  which is located in the Open Knee(s) repository.

    Original author:
        Ahmet Erdemir
        Computational Biomodeling (CoBi) Core
        Department of Biomedical Engineering
        Lerner Research Institute
        Cleveland Clinic
        Cleveland, OH
        erdemira@ccf.org

    This function will parse the given .tdms filen and return a dictionary of values.
    This dictionary will contain several entries.

    Keys and descriptions of the entries can be found in the OpenKnee(s) documentation.

    As of January 2016, the documentation for these values can be found in the OpenKnee(s) wiki page under Infrastructure/ExperimentationMechanics in a file called *2013CB-031-002.B simVITRO Data File Contents_Open Knee.pdf*
    http://wiki.simtk.org/openknee/Infrastructure/ExperimentationMechanics?action=AttachFile&do=get&target=2013CB-031-002.B+simVITRO+Data+File+Contents_Open+Knee.pdf

    :param tdmsFileName: string, The file name of the .tdms file
    :return: dictionary, A dictionary which contains the parsed values from the .tdms file

    :Note: the module *nptdms* must be installed for this function to run.

    :Example:

    .. code-block:: python

        import tdmsParser
        import matplotlib.pyplot as plt

        # The .tdms file name
        tdmsFileName = '../oks001/joint_mechanics-oks001/TibiofemoralJoint/KinematicsKinetics/005_passive flexion_optimized/Data/005_passive flexion_optimized_main_processed.tdms'
        # parse the file
        data = parseTDMSfile(tdmsFileName)

        # Plot flexion against internal rotation.
        flexion = data['State.JCS']['JCS Flexion']
        internalRotaiton = data['State.JCS']['JCS Internal Rotation']

        fig, ax1 = plt.subplots(nrows=1)
        ax1.plot(flexion, internalRotaiton, color='r')
        plt.show()
    """
    tdmsFile = TdmsFile(tdmsFileName) # Read the tdms file with the TdmsFile function
    groups = tdmsFile.groups() # Get the groups from the tdmsFile
    groupDict = {}
    # Iterate over the groups and get the channels from the groups
    for i, group in enumerate(groups):
        groupDict[group] = {}
        channels = tdmsFile.group_channels(group) # Get the channels from the group
        # Iterate over the channels, and check if the channel dictionary has the key 'NI_ChannelName'.
        # If the dictionary does have that key, then return the data from that channel.
        for channel in channels:
            channelData = _parseChannel(channel, targetKey='NI_ChannelName')
            # If None is returned, then go to the next channel.
            if channelData is None:
                continue
            # Else, add the channelData to the group dictionary
            else:
                groupDict[group].update(channelData)
    return groupDict

# This function will
def _parseChannel(channel, targetKey = 'NI_ChannelName'):
    channelDict = {} # Initialize the channel dictionary.
    # Try to parse the data from the channel.
    try:
        channelLabel = channel.properties[targetKey] # define the label for the channel.
        channelDict[channelLabel] = channel.data # Put the channel data in channelDict. The channel label is the key.
        return channelDict
    # If there is an exception, return None
    except:
        return None

if __name__ == '__main__':
    tdmsFileName = '/home/will/Projects/MCLS_CSU/MCLS_002/OpenKnee/oks001/joint_mechanics-oks001/TibiofemoralJoint/KinematicsKinetics/005_passive flexion_optimized/Data/005_passive flexion_optimized_main_processed.tdms'
    data = parseTDMSfile(tdmsFileName)

    import matplotlib.pyplot as plt
    flex = data['Kinematics.JCS.Actual']['Internal Rotation Angle']
    flexActual = data['Kinematics.JCS.Actual']['Flexion Angle']
    # loadSup = data['State.JCS Load']['JCS Load Distraction']

    xVals = range(len(flex))

    fig, ax1 = plt.subplots(nrows=1)
    ax1.plot(xVals, flex, color='r')

    plt.show()