#!/usr/bin/python3
"""
Python3 script to extract movement data from Xsens mvnx XML data structure

Execute in terminal as: python3 Extract-Xsens-MVN-Data_AllSegments.py

Output filename is generated automatically using:
    Current timestamp + input filename + option + segment . txt

Author: R. Muralikrishnan
19.07.2018
"""

import xml.etree.ElementTree as ET
import datetime

## Get the input file to process

try:
    V_input_file = input("\nEnter input filename: ")
    V_xmltree = ET.parse(V_input_file)
    V_root = V_xmltree.getroot()
    
    # https://stackoverflow.com/a/1912483
    # https://docs.python.org/3/library/xml.etree.elementtree.html#xml.etree.ElementTree.Element
            
except FileNotFoundError:
    print("\nSomething's wrong with the file name you entered.  Please check if the file exists.")

else:       
    # If there's no problem with opening the file, proceed further  
   
    ## Define the data options 
    
    D_options = { 1 : "orientation",
                  2 : "position",
                  3 : "velocity",
                  4 : "acceleration",
                  5 : "angularVelocity",
                  6 : "angularAcceleration",
                 # 7 : "sensorFreeAcceleration",
                 # 8 : "sensorMagneticField",
                 # 9 : "sensorOrientation"
                }
    
    ## Define the data dimension for each option
    # Ref: MVN User Manual Page 91 (Document Page 101/172)
    D_option_dimensions = { 1 : 4,     # "orientation" : 23 segments x 1 x 4 
                             2 : 3,     # "position",   : 23 segments x 1 x 3 ...
                             3 : 3,     # "velocity",              23 x 1 x 3
                             4 : 3,     # "acceleration",          23 x 1 x 3
                             5 : 3,     # "angularVelocity",       23 x 1 x 3
                             6 : 3      # "angularAcceleration",   23 x 1 x 3
                             # 7 : 3,   # "sensorFreeAcceleration",23 x 1 x 3
                             # 8 : 3,   # "sensorMagneticField",   23 x 1 x 3
                             # 9 : 4    # "sensorOrientation"      23 x 1 x 4
                            }
    
    ## Get the option for which data is desired 
    print("\nThe following data options are available.")
    
    for V_option_num, V_option_name in D_options.items():
        print(V_option_num, ":", V_option_name)
        
    try:
        V_desired_option = int(input("Enter the option number you would like in the output (1-6) : "))
        
        if not (V_desired_option >= 1 and V_desired_option <= 6):
            print("That's an invalid option!  Options can range from 1 to 6.")
            # Don't proceed further!
            V_proceed = False
            
        else:
        ### COMMENTED OUT for getting data for ALL segments!    
        #    ## Get the segment for which data is desired
        #    try:
        #        V_desired_segment = int(input("Enter the segment number for which to extract data (1-23): "))
                
        #        if not (V_desired_segment >= 1 and V_desired_segment <= 23):
        #            print("That's an invalid segment!  Segment numbers can range from 1 to 23.")
        #            # Don't proceed further!
        #            V_proceed = False
        #        else:
        #            # Everything ok with the input ... set the flag to proceed further
        #            V_proceed = True            
            V_proceed = True  ### ADDED this here for getting data for ALL segments!
            
        #    except:
        #        print("Something is wrong with the segment you entered.  Did you enter a number in the range of 1 to 23?")
        #        V_proceed = False
        
    except:
        print("Something is wrong with the option you entered.  Did you enter a number in the range of 1 to 9?")
        V_proceed = False
        
    # Get a time range to extract the data
    # If none, then all data will be extracted.
    try:
        print("\nYou can optionally choose a time range below to extract data.")
        print("Entering nothing at all here would extract data from all timepoints.")
        V_desired_time_begin = input("Enter the begin time (Inpoint) in the format HH:MM:SS:ms : ")
        if V_desired_time_begin.upper() != "":

            # From the HH:MM:SS:XX begin and end timepoints, work out the time in milliseconds for comparison later.
            # Note that, the movement data was recorded with 60 Hz sampling frequency ... that means,
            # the entries for XX above would run from 00 to 59 ... and not the whole gamut of 0 to 999 ms.
            # To get the ms value from it, we need (XX/60) * 1000.
            
            L_time_begin = V_desired_time_begin.split(":")
            V_time_begin_ms = int(L_time_begin[0]) * 60 * 60 * 1000 + int(L_time_begin[1]) * 60 * 1000 + int(L_time_begin[2]) * 1000 + 1000*int(L_time_begin[3])/60
 
            V_desired_time_end = input("Enter the end time (Outpoint) in the format HH:MM:SS:ms : ")
            
            L_time_end = V_desired_time_end.split(":")
            V_time_end_ms = int(L_time_end[0]) * 60 * 60 * 1000 + int(L_time_end[1]) * 60 * 1000 + int(L_time_end[2]) * 1000 + 1000*int(L_time_end[3])/60
                        
            V_extract_all_timepoints = False
            
        else:
            print("\nExtracting data from all available timepoints.")
            V_extract_all_timepoints = True
            
    except:
        print("\nSomething was wrong with the begin and/or end time you specified.") 
        print("Extracting data from all available timepoints instead!")
        V_extract_all_timepoints = True
    
    
    ## Proceed reading the data if the proceed flag is set to true above
    
    if V_proceed:
        V_timestamp = "Dummy"
        
        # Get the date and time: Refs:
        # https://stackoverflow.com/a/311655
        # https://stackoverflow.com/a/415519
        # V_date = str(datetime.date.today())
        # V_time = str(datetime.datetime.now().time())
        # https://stackoverflow.com/a/14229023
        V_current_time = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S__')
        
        V_output_file = V_current_time + V_input_file + "__" + D_options[V_desired_option].capitalize() + "_S_ALL.txt"
        #V_output_file = input("Enter an output file name. Abc.txt: ") 
     
        for V_child in V_xmltree.iter():
            # print(x.tag, x.text, x.attrib.items)
            
            ## Get the child's attributes
            for V_field, V_value in V_child.attrib.items():
                #if y in ["tc", "index"]:
                if V_field == "tc":
                    V_timestamp = V_value
                    # From the HH:MM:SS:ms timestamp, work out the time in milliseconds for comparison later.
                    L_timestamp = V_timestamp.split(":")
                    V_timestamp_ms = int(L_timestamp[0]) * 60 * 60 * 1000 + int(L_timestamp[1]) * 60 * 1000 + int(L_timestamp[2]) * 1000 + 1000*int(L_timestamp[3])/60         
                    #print(V_field, V_value, end = " ")
                    #pass

            ## Get the child's tag itself
            # Split the unclean tagline of the form "{http://...} orientation"
            # to be able to extract the tag alone.
            L_tag = V_child.tag.split("}")

            ## Get the child's text
            V_text = V_child.text
            
            L_segments = [x for x in range(1,24)]
            
            ## Extract data if the text is not an empty line
            if isinstance(V_text, str):
                V_text_stripped = V_text.strip()      # The whole data line...
                L_contents = V_text_stripped.split()  # ...split by space 
                       
                ## Select lines that contain data for the desired option
                #if len(L_tag) > 1 and L_tag[1] == "orientation" 
                if L_tag[1] == D_options[V_desired_option]:
                    #print(IndexTimeInfo, L_tag[1], L_contents[0])
                    
                    for V_desired_segment in L_segments:
                        ## Select the fields to print from the contents for desired segment
                        # Remember ... index begin-inclusive, end-exclusive!
                        # So for segment 1 1x3, we need L_contents[(1-1)*3 : 1*3] = [0:3] => 0, 1, 2
                        # .. for segment 2 1x3, we need L_contents[(2-1)*3 : 2*3] = [3:6] => 3, 4, 5
                        # .. and so on
                        V_index_begin = (V_desired_segment - 1) * D_option_dimensions[V_desired_option]
                        V_index_end = V_desired_segment * D_option_dimensions[V_desired_option]
    
                        try:
                            ## Print the data
                            # print(IndexTimeInfo, L_tag[1], L_contents[V_index_begin : V_index_end])
                            # Either all data is printed, or only certain time points as desired above.
                            if V_extract_all_timepoints or (V_timestamp_ms >= V_time_begin_ms and V_timestamp_ms <= V_time_end_ms):                        
                                # Print the data without listed representation to a file                       
                                # A * before the listname will unpack the list and returns the element in it!
                                with open(V_output_file, "a") as File_1:
                                    print(V_timestamp, V_timestamp_ms,
                                          L_tag[1],
                                          "Seg_" + str(V_desired_segment),
                                          * L_contents[V_index_begin : V_index_end], # Unpack list elements: https://stackoverflow.com/a/35119046
                                          sep = "\t",
                                          file = File_1) # Print to file: https://stackoverflow.com/a/9236236
    
                        except:
                            print("Something wrong with extracting/printing the data.\nPerhaps you're trying to extract non-existing fields???")

        else:
            try:
                with open(V_output_file, "r") as File_2:
                    # File is there and readable.
                    print("\nSuccessful.  Output written to: ", V_output_file, ".", sep = "")
                    
            except:
                print("\nOops. No output generated!\nPerhaps no data available for", D_options[V_desired_option].capitalize(), "in the input file (for the time range you specified).")