## A single purpose script to parse motive tracker version 1.5 csv export data.
## Rigid bodies represented as GH planes using Rhino Script Syntax.
## Josh Bard, Carnegie Mellon University 01.20.17    

import rhinoscriptsyntax as rs

def frameDataFromCSV(CSVPath):
    """ creates a structured list of MOCAP frame data from a single 
        Motive Tracker session exported to CSV
        Input: path to CSV file
        Output: a list of frames with the following structure:
                [timestamp, for each rb [id,[x,y,z],[roll,pitch,yaw] or 
                [None](for dropped rb), num dropped rbs in frame]"""
    frames = []
    
    with open (CSVPath) as MocapData: # open csv file and close when finished
        for line in MocapData.readlines(): 
            rawLine = line.split(',')
            
            # determine total num of rbs tracked during session found in 
            # default "info" file header
            if (rawLine[0] == "info" and rawLine[1] == "rigidbodycount"):
                rb_Total = int(rawLine[2]) 

           # construct data at each frame
            elif (rawLine[0] == "frame"): # Motive convention for labeling rows
                rb_Count = int(rawLine[3]) # num rbs tracked in frame
                rb_Dropped = rb_Total - rb_Count # num dropped rbs in frame
                
                # initialize single frame
                frame = []
                frame.append(float(rawLine[2])) # add timestamp
                frame.extend([[None] for _ in range(rb_Total)]) # dropped rbs remain None
                frame.append(rb_Dropped) 
                
                # for each tracked rigid body compile frame data
                for i in range (rb_Count):
                    fos = i*11 # fixed offset per rigid body def
                    id = int(rawLine [4 + fos]) # motive assigned id
                    pos = [float(x) for x in rawLine[5+fos:8+fos]] # [x,y,z]
                    orient = [float(x) for x in rawLine[12+fos:15+fos]] #[roll,pitch,yaw]
                    frame[id] = [id]+[pos] + [orient]
                frames.append(frame)
    return frames 


def movePlane(x,y,z,scale = True):
    """move and optionally scale a generic plane
       Note: Motive units set to meters. scale if rhino units set to mm"""
    basePlane = rs.WorldXYPlane()
    movePoint = rs.AddPoint(x,y,z)
    origin = (0,0,0)
    if scale == True:
        scale = (1000,1000,1000)
        movePoint = rs.ScaleObject(movePoint,origin,scale)
    movePlane = rs.MovePlane(basePlane,movePoint)
    return movePlane


def rotatePlane (x,y,z,rz,rx,ry,scale):
    """Rotate base plane about rz,ry,rx(yaw,pitch,roll)"""
    basePlane = movePlane(x,y,z,scale)
    rzPlane = rs.RotatePlane(basePlane,rz,basePlane.ZAxis)
    ryPlane = rs.RotatePlane(rzPlane,ry,rzPlane.YAxis)
    rotatePlane = rs.RotatePlane(ryPlane,rx,ryPlane.XAxis)
    return rotatePlane

def genRbPlane(rb_data,scale):
    if (rb_data[0] != None):
        x,y,z = rb_data[1][0],rb_data[1][1],rb_data[1][2]
        rz,ry,rx = rb_data[2][0],rb_data[2][1],rb_data[2][2]
        rb_plane = rotatePlane(x,y,z,rz,rx,ry,scale)
        return rb_plane
    else: return None

def compileRbPlanes(frame_data,scale):
    rb_num = (len(frame_data[0])- 2)
    rb_planes = [[] for _ in range(rb_num)]
    
    for frame in frame_data:
        for rb in range(rb_num):
            rb_plane = genRbPlane(frame[rb+1],scale)
            rb_planes[rb].append(rb_plane)         
    return rb_planes           

def list_to_tree(input, none_and_holes=True, source=[0]):
    """Transforms nestings of lists or tuples to a Grasshopper DataTree
       from https://gist.github.com/piac/ef91ac83cb5ee92a1294"""
    from Grasshopper import DataTree as Tree
    from Grasshopper.Kernel.Data import GH_Path as Path
    from System import Array
    def proc(input,tree,track):
        path = Path(Array[int](track))
        if len(input) == 0 and none_and_holes: tree.EnsurePath(path); return
        for i,item in enumerate(input):
            if hasattr(item, '__iter__'): #if list or tuple
                track.append(i); proc(item,tree,track); track.pop()
            else:
                if none_and_holes: tree.Insert(item,path,i)
                elif item is not None: tree.Add(item,path)
            
    if input is not None: t=Tree[object]();proc(input,t,source[:]);return t

def printMOCAPSessionReport (frame_data):        
    frame_num = len(frame_data)
    rb_num = (len(frame_data[0])- 2)
    droppedRbCount = 0
    droppedFrameCount = 0
    captureDuration = frame_data[-1][0]
    
    for frame in frame_data:
        droppedRbCount += frame[-1]
        if frame[-1] == rb_num: droppedFrameCount +=1
    per_droppedFrames = droppedFrameCount/frame_num
    per_droppedRbs = droppedRbCount/(rb_num * frame_num)
    report = (
               'There were %d rigid bodies tracked in the scene\n'
               'There were %d frames captured during this session\n'
               'There were %d (%.3f%%) dropped rigid bodies during the session\n'
               'There were %d (%.3f%%) dropped frames during the session\n'
               'capture duration %.2f seconds'
                ) %(rb_num,frame_num,droppedRbCount,per_droppedRbs,
                    droppedFrameCount,per_droppedFrames,captureDuration)
    return report

################################################################################
## Generate data for GH Output
frame_data = frameDataFromCSV(CSVPath)
print (printMOCAPSessionReport(frame_data))
rb_planes = list_to_tree(compileRbPlanes(frame_data,scale))