# mocap_rman.py
# The output rib is saved in the project/RIB_Archive directory. 
#
# Malcolm Kesson
# Milim Lee
# March 12 2017
  
from mocap_db import MoCapDB
import os
from maya_proj_utils import MayaProjUtils
import math
 
class MoCapRMan(MoCapDB):
    def __init__(self, name, datapath, scaling):
        self.name = name
        self.begin = 1
        self.end = 1
        self.utils = None
        try:
            self.utils = MayaProjUtils()
            self.archivepath = self.utils.getRIB_ArchivePath()
            if os.path.exists(self.archivepath) == False:
                os.mkdir(self.archivepath)
            self.begin = self.utils.getAnimationStart()
            self.end = self.utils.getAnimationEnd()
            self.scenename = self.utils.getSceneName()
        except:
            self.archivepath = os.path.dirname(datapath)
            self.scenename = 'untitled'
        MoCapDB.__init__(self, datapath, scaling)
        dataname = os.path.basename(datapath)
        self.dataname = dataname[:-4] # remove .txt
    # __________________________________________________
    def writePoints(self, trail, step, width, cache):
        if self.dataIsValid() == False:
            print('writePoints is valid')
            return ''
        current_frame = self.utils.getCurrentTime()
        begin,end = self.getDataRange(trail)
        ribpath = self.getRibPath('pnt',trail,step,width)
        # Reuse a previously computed rib archive...
        if self.canReuse(ribpath,cache):
            return ribpath
        f = open(ribpath, 'w')
        f.write(self.getRibHeader(begin,end))
        f.write('Points "P" [')
        for n in range(len(self.markers)):
            data = self.getMarkerData(n,begin,end,step)
            if len(data) == 0: # marker has no data for this frame
                continue
            count = 0
            for coord in data:
                if count == 0 or count % 15 == 0:
                    f.write('\n\t')
                f.write('%1.4f ' % coord )
                count += 1
        f.write('\n\t] "constantwidth" [%1.4f]\n' % width)
        f.close()
        return ribpath        
    # __________________________________________________
    # The curve type is catmull-rom because it guarantees the curve will pass
    # through the coordinates of the control vertices. However, the first and
    # last xyz coordinates must be repeated.
  
    def writeCurves(self, trail, step, width, cache):
        if self.dataIsValid() == False:
            return ''
        begin,end = self.getDataRange(trail)
        # We need at least 4 cvs for a valid curve
        data_range = end - begin
        if data_range < 4:
            end = begin + 4
        ribpath = self.getRibPath('cvs',trail,step,width)
        # Reuse a previously computed rib archive...
        if cache == 'Reuse' and os.path.exists(ribpath) == True:
            return ribpath
            
        f = open(ribpath, 'w')
        f.write(self.getRibHeader(begin,end))
        f.write('AttributeBegin\n')
        f.write('\tBasis "catmull-rom" 1 "catmull-rom" 1\n')
        for n in range(len(self.markers)):
            # data is a simple list of coordinates [x,y,z,x,y,z etc...]
            data = self.getMarkerData(n,begin,end,step)
            # The marker has no data for this frame, go to the next marker.
            if len(data) == 0: 
                continue            
            # A "catmull-rom" generally has the first and last cvs repeated,
            # hence, we add 2 to the number of cvs.
            numcvs = (len(data)/3) + 2
            f.write('\tCurves "cubic" [%d] "nonperiodic" "P" [' % numcvs)
            # Repeat the first CV
            f.write('%1.4f %1.4f %1.4f ' % (data[0],data[1],data[2]) )
            count = 2
            for coord in data:
                if count % 15 == 0:
                    f.write('\n\t\t')
                f.write('%1.4f ' % coord )
                count += 1
            # Repeat the last CV
            x,y,z = data[-3:]
            f.write('%1.4f %1.4f %1.4f ' % (x,y,z) )
            f.write('\n\t\t] "constantwidth" [%1.4f]\n' % width)
        f.write('AttributeEnd\n')
        f.close()
        return ribpath
  
    #making body shape skeleton
    def createSkeleton(self, trail, step, width, cache, isMixed=False):
        if self.dataIsValid() == False:
            return ''
        begin,end = self.getDataRange(trail)
        data_range = end - begin
        end = begin + 4
        if isMixed:
            ribpath = self.getRibPath('breakdown',trail,step,width)
        else:
            ribpath = self.getRibPath('cvs',trail,step,width)
        if cache == 'Reuse' and os.path.exists(ribpath) == True:
            return ribpath
        orders = [['RFHD', 'LFHD', 'LBHD', 'RBHD', 'RFHD'], #head to little spine
                ['RBHD', 'C7', 'RBAC', 'T10', 'RBWT', 'RFWT', 'RTHI', 'RKNE', 'RSHN', 'RANK', 'RHEE', 'RMT5', 'RTOE', 'RHEE'], #right spine to right foot
                ['LBHD', 'C7', 'T10', 'LBWT', 'LFWT', 'LTHI', 'LKNE', 'LSHN', 'LANK', 'LHEE', 'LMT5', 'LTOE', 'LHEE'], #left spine to left foot
                ['C7', 'CLAV', 'STRN', 'LFWT'], # belly
                ['LBWT', 'LSHO', 'LFWT'],
                ['RBWT', 'RSHO', 'RFWT'],
                ['STRN', 'RFWT'], # belly
                ['C7', 'LSHO', 'LUPA', 'LELB', 'LFRM', 'LWRB', 'LFIN', 'LWRA', 'LWRB'], #left arm
                ['C7', 'RSHO', 'RUPA', 'RELB', 'RFRM', 'RWRB', 'RWRA', 'RFIN', 'RWRB']] #right arm
        f = open(ribpath, 'w')
        f.write(self.getRibHeader(begin,end))
        f.write('AttributeBegin\n')
        for frame in range(begin, end, step):
            coordsList = self.getFrameData(frame)
            coordsDic = self.getMarkerFrameData(frame) #dictionary markername
  
            f.write('\tPoints "P" [')
            for n in range(len(coordsList)):
                f.write('%1.4f %1.4f %1.4f ' % (coordsList[n][0],coordsList[n][1],coordsList[n][2]))
            f.write('\n\t\t] "constantwidth" [%1.4f]\n' % (width*2))            
                        
            for n in range(len(orders)):
                count = 0
                for m in range(len(orders[n])):
                    if (orders[n][m] in coordsDic) and (len(coordsDic[orders[n][m]]) != 0):
                        count += 1
                f.write('\tCurves "linear" [%d] "nonperiodic" "P" [' % count)
                
                for m in range(len(orders[n])):
                    if (orders[n][m] in coordsDic) and (len(coordsDic[orders[n][m]]) != 0):
                        f.write('%1.4f %1.4f %1.4f ' % (coordsDic[orders[n][m]][0],coordsDic[orders[n][m]][1],coordsDic[orders[n][m]][2]))
                f.write('\n\t\t] "constantwidth" [%1.4f]\n' % width)
        
        f.write('AttributeEnd\n')
        f.close()
        return ribpath
  
    #write body shape with joker's face
    def createJoker(self, trail, step, width, cache, scale, myRibModelPath, isMixed=False):
        if self.dataIsValid() == False:
            return ''
        tail = 1
        begin,end = self.getDataRange(trail)
        current_frame = self.utils.getCurrentTime()        
        data_range = end - begin
        if isMixed:
            ribpath = self.getRibPath('breakdown',trail,step,width)
        else:
            ribpath = self.getRibPath('cvs',trail,step,width)
        if cache == 'Reuse' and os.path.exists(ribpath) == True:
            return ribpath
        orders = [['RFWT','RTHI', 'RKNE', 'RSHN', 'RANK', 'RHEE'], ['LFWT','LTHI', 'LKNE', 'LSHN', 'LANK', 'LHEE'], ['LUPA', 'LELB', 'LFRM', 'LWRB'], ['RUPA', 'RELB', 'RFRM', 'RWRB']]
                
        f = open(ribpath, 'w')
        f.write(self.getRibHeader(begin,end))
        f.write('AttributeBegin\n')
        for frame in range(begin, end, step):
            coordsList = self.getFrameData(frame)
            coordsDic = self.getMarkerFrameData(frame) #dictionary markername
  
            if (isMixed and (frame >= 214)) or (not isMixed):    
                f.write('\tBasis "catmull-rom" 1 "catmull-rom" 1\n')            
                for n in range(len(orders)):
                    count = 0
                    orderlength = len(orders[n])
                    for m in range(orderlength):
                        if (orders[n][m] in coordsDic) and (len(coordsDic[orders[n][m]]) != 0):
                            count += 1
                    f.write('\tCurves "cubic" [%d] "nonperiodic" "P" [' % (count+2))
                    
                    f.write('%1.4f %1.4f %1.4f ' % (coordsDic[orders[n][0]][0],coordsDic[orders[n][0]][1],coordsDic[orders[n][0]][2]))
                    for m in range(orderlength):
                        if (orders[n][m] in coordsDic) and (len(coordsDic[orders[n][m]]) != 0):
                            f.write('%1.4f %1.4f %1.4f ' % (coordsDic[orders[n][m]][0],coordsDic[orders[n][m]][1],coordsDic[orders[n][m]][2]))
                    f.write('%1.4f %1.4f %1.4f ' % (coordsDic[orders[n][orderlength-1]][0],coordsDic[orders[n][orderlength-1]][1],coordsDic[orders[n][orderlength-1]][2]))
                    f.write('\n\t\t] "constantwidth" [%1.4f]\n' % width)
            
            if frame == (end-1):
                #putting face!
                vecInputs = []
                vecInputs.extend([coordsDic["RFHD"], coordsDic["LFHD"], coordsDic["RBHD"], coordsDic["LBHD"]])
                facePos = self._getPointAvg(vecInputs)
                faceRotResult = self._rotateY(coordsDic["RBHD"], coordsDic["RFHD"])
                ribPath = myRibModelPath + "/clownFace.rib"
                faceScale = scale * 27
                self._writeTransform(f, facePos[0], facePos[1], facePos[2], 0.0 , faceRotResult[0], 0.0, faceScale, faceScale, faceScale, ribPath)
  
                #putting body!
                if (isMixed and (frame >= 107))  or (not isMixed):
                    vecInputs = []
                    vecInputs.extend([coordsDic["LSHO"], coordsDic["RSHO"], coordsDic["LFWT"], coordsDic["RFWT"]])
                    cardPos = self._getPointAvg(vecInputs)
                    vecInputs = []
                    vecInputs.extend([coordsDic["LSHO"], coordsDic["RSHO"]])                
                    newBodyRefpoint = self._getPointAvg(vecInputs)
                    cardRotResult = self._rotateY(coordsDic["T10"], newBodyRefpoint)
                    ribPath = myRibModelPath + "/card.rib"
                    bodyScale = scale * 37
                    self._writeTransform(f, cardPos[0], cardPos[1], cardPos[2], 0.0 , cardRotResult[0], 0.0, bodyScale, bodyScale, bodyScale, ribPath)                                
  
                #putting left hand!
                if (isMixed and (frame >= 143)) or (not isMixed):                
                    vecInputs = []
                    vecInputs.extend([coordsDic["LFIN"], coordsDic["LWRA"]])                
                    newLeftHandRefpoint = self._getPointAvg(vecInputs)
                    leftHandRotResult = self._aimY(coordsDic["LWRB"], newLeftHandRefpoint)
                    ribPath = myRibModelPath + "/rightHand.rib"
                    lHandScale = scale * 30
                    self._writeTransform(f, coordsDic["LWRB"][0], coordsDic["LWRB"][1], coordsDic["LWRB"][2], leftHandRotResult[0] , 0.0, leftHandRotResult[1], lHandScale, lHandScale, lHandScale, ribPath)                                                
                #putting right hand!
                    vecInputs = []
                    vecInputs.extend([coordsDic["RWRA"], coordsDic["RFIN"]])                
                    newRightHandRefpoint = self._getPointAvg(vecInputs)
                    rightHandRotResult = self._aimY(coordsDic["RFRM"], coordsDic["RFIN"])
                    ribPath = myRibModelPath + "/leftHand.rib"
                    rHandScale = scale * 30
                    self._writeTransform(f, coordsDic["RFRM"][0], coordsDic["RFRM"][1], coordsDic["RFRM"][2], rightHandRotResult[0] , 0.0, rightHandRotResult[1], rHandScale, rHandScale, rHandScale, ribPath)                        
    
                #putting left shoe!
                if (isMixed and (frame >= 179)) or (not isMixed):
                    vecInputs = []
                    vecInputs.extend([coordsDic["RMT5"], coordsDic["RTOE"]])                
                    newLeftShoeRefpoint = self._getPointAvg(vecInputs)
                    leftShoeRotResult = self._aimY(coordsDic["RHEE"], newLeftShoeRefpoint)
                    ribPath = myRibModelPath + "/leftShoe.rib"
                    lShoeScale = scale * 20
                    self._writeTransform(f, coordsDic["RHEE"][0], coordsDic["RHEE"][1], coordsDic["RHEE"][2], leftShoeRotResult[0] , 0.0, leftShoeRotResult[1], lShoeScale, lShoeScale, lShoeScale, ribPath)                        
                #putting right shoe!
                    vecInputs = []
                    vecInputs.extend([coordsDic["LMT5"], coordsDic["LTOE"]])                
                    newRightShoeRefpoint = self._getPointAvg(vecInputs)    
                    rightShoeRotResult = self._aimY(coordsDic["LHEE"], newRightShoeRefpoint)
                    ribPath = myRibModelPath + "/rightShoe.rib"
                    rShoeScale = scale * 20
                    self._writeTransform(f, coordsDic["LHEE"][0], coordsDic["LHEE"][1], coordsDic["LHEE"][2], rightShoeRotResult[0] , 0.0, rightShoeRotResult[1], rShoeScale, rShoeScale, rShoeScale, ribPath)        
  
        f.write('AttributeEnd\n')
        f.close()
        return ribpath
  
    # __________________________________________________
    def writeBlobby(self, trail, step, width, cache, volume=False, isMixed=False):
        if self.dataIsValid() == False:
            return ''
        begin,end = self.getDataRange(trail)
        geoname = 'blobS'
        if isMixed:
            geoname = 'breakdown'
        elif volume:
            geoname = 'blobV'
        ribpath = self.getRibPath(geoname,trail,step,width)
        # Reuse a previously computed rib archive...
        if cache == 'Reuse' and os.path.exists(ribpath) == True:
            return ribpath
        
        f = open(ribpath, 'w')
        f.write(self.getRibHeader(begin,end))
        alldata = []
        for n in range(len(self.markers)):
            data = self.getMarkerData(n,begin,end,step)
            if len(data) == 0:
                continue
            alldata.extend(data)
        numblobs = len(alldata)/3
        if volume:
            f.write('Blobby %d [8 \n' % numblobs)
        else:
            f.write('Blobby %d [\n' % numblobs)
        # Make each blob an ellipsoid and provide its array index.
        # The indices monotonously increment by 16.
        for n in range(numblobs):
            f.write('\t1001 %d\n' % (n * 16))
        # Add the blending code "0" and the number of blobs to blend.
        f.write('\t0 %d ' % + numblobs)
        # Specify the indices of all the blobs.
        for n in range(numblobs):
            f.write(' %d' % n)
        f.write(']\n\t[\n')
        for n in range(0, len(alldata), 3):
            x = alldata[n]
            y = alldata[n+1]
            z = alldata[n+2]
            f.write('\t')
            f.write('%1.4f 0 0 0  ' % (width *2))
            f.write('0 %1.4f 0 0  ' % (width*2))
            f.write('0 0 %1.4f 0  ' % (width*2))
            f.write('%1.4f %1.4f %1.4f 1\n' % (x,y,z))
        f.write('\t] [""]\n')
        f.close()
        return ribpath
  
        
    # __________________________________________________
    # Given getPadding(27), returns '0027'
    def getPadding(self, frame):
        return '%0*d' % (4, frame)
  
  
    # __________________________________________________
    def getRibPath(self,geoname,trail,step,geosize):
        frame = self.utils.getCurrentTime()
        pad = self.getPadding(frame) 
        size = '%1.3f' % geosize # Avoid excessive trailing digits
        size = size.replace('.', '')
        ribname = '%s_%s_%s_%s_%s.%s.rib' % (self.dataname, geoname,
                                          str(trail), str(step),
                                          size, pad)
        path = os.path.join(self.archivepath, self.scenename)
        if os.path.exists(path) == False:
            os.mkdir(path)
        path = os.path.join(path, self.name)
        if os.path.exists(path) == False:
            os.mkdir(path)
        fullpath = os.path.join(path, ribname)
        return fullpath
  
  
    #___________________________________________________________
    # Rotate x and z angles
    def _aimY(self, vec1, vec2): #arrow goes from vec1 to vec2
        xyLength = math.sqrt((vec2[0] - vec1[0]) * (vec2[0] - vec1[0]) + (vec2[1] - vec1[1]) * (vec2[1] - vec1[1]))
        vecLength = math.sqrt((vec2[0] - vec1[0]) * (vec2[0] - vec1[0]) + (vec2[1] - vec1[1]) * (vec2[1] - vec1[1]) + (vec2[2] - vec1[2]) * (vec2[2] - vec1[2]))
    
        #xyLength will be zero when vec is pointing along the +z or -z axis
        out = []
        if xyLength == 0:
            if (vec2[0] - vec1[0]) > 0:
                zAngle = math.radians(90)
            else:
                zAngle = math.randians(-90)
        else:
            zAngle = math.acos((vec2[1] - vec1[1])/xyLength)
    
        xAngle = math.acos(xyLength/vecLength)
    
        if (vec2[2] - vec1[2]) > 0:
            xAngle = xAngle
        else:
            xAngle = -xAngle
      
        if (vec2[0] - vec1[0]) > 0:
            zAngle = -zAngle
        else:
            zAngle = zAngle
            
        out.extend([math.degrees(xAngle), math.degrees(zAngle)])
        return out
  
  
    #______________________________________________________
    # Rotate y angle
    def _rotateY(self, vec1, vec2): #arrow goes from vec1 to vec2
        xLength = vec2[0] - vec1[0]
        xzLength = math.sqrt((vec2[0] - vec1[0]) * (vec2[0] - vec1[0]) + (vec2[2] - vec1[2]) * (vec2[2] - vec1[2]))
    
        out = []
        if xLength == 0:
            if (vec2[2] - vec1[2]) > 0:
                yAngle = math.radians(90)
            else:
                yAngle = math.randians(-90)
        else:
            yAngle = math.acos((vec2[0] - vec1[0])/xzLength)
    
        if (vec2[2] - vec1[2]) > 0:
            yAngle = -yAngle
        else:
            yAngle = yAngle
        out.append(math.degrees(yAngle))
        return out
  
  
    #______________________________________________________
    # Calculate average vector of points    
    def _getPointAvg(self, points): 
    
        length = 0
        sumX = 0
        sumY = 0
        sumZ = 0
        
        for pnt in points:
            if(len(pnt) > 0):
                sumX += pnt[0] 
                sumY += pnt[1]
                sumZ += pnt[2]
                length += 1
        
        out = []    
        resultX = sumX / length
        resultY = sumY / length
        resultZ = sumZ / length
        out.extend([resultX, resultY, resultZ])
        return out
  
        
    #______________________________________________________
    # Write transform commend in order.
    def _writeTransform(self, f, transX, transY, transZ, rotX, rotY, rotZ, scaleX, scaleY, scaleZ, ribPath):
        f.write('\tTransformBegin\n')
        f.write('\t\tTranslate %1.4f %1.4f %1.4f\n' %(transX, transY, transZ))
        f.write('\t\tRotate %1.4f %d %d %d\n' %(rotX, 1, 0, 0))
        f.write('\t\tRotate %1.4f %d %d %d\n' %(rotY, 0, 1, 0))
        f.write('\t\tRotate %1.4f %d %d %d\n' %(rotZ, 0, 0, 1))
        f.write('\t\tScale %1.4f %1.4f %1.4f\n' %(scaleX, scaleY, scaleZ))
        f.write('\t\tReadArchive "%s"\n' %(ribPath))
        f.write('\tTransformEnd\n')
  
  
    # __________________________________________________
    # Puts some useful information at the beginning of the rib archive file.    
    def getRibHeader(self,begin,end):
        x,y,z,X,Y,Z = self.getFrameBbox(end) #self.getBbox()
        rib =  '#bbox: %1.4f %1.4f %1.4f %1.4f %1.4f %1.4f \n' % (x,y,z,X,Y,Z)
        rib += '# Data begin/end: %d to %d\n' % (begin,end)
        rib += '# Total number of mocap markers %d\n' % len(self.markers)
        rib += '# Total number of mocap frames %d\n' % self.frames
        return rib
  
    # __________________________________________________    
    def getDataRange(self,trail):
        # Set a couple of reasonable default values for the data_begin/end.
        # self.begin and self.end are "read" from the "Frame Range" settings 
        # of Maya's Render Setting window - see the constructor.
        data_begin = self.begin - trail
        data_end = self.end
        if data_begin < 0:
            data_begin = 0;
        # Use Maya's values...
        if self.utils != None:
            current = self.utils.getCurrentTime()
            data_end = current
            if (current - trail) < 1:
                data_begin = 1
            else:
                data_begin = current - trail    
        return [data_begin,data_end]
  
    # __________________________________________________    
    def dataIsValid(self):
        if len(self.getBbox()) == 0:
            return False
        return True
  
    # __________________________________________________    
    def canReuse(self,ribpath,cache):
        if cache == 'Reuse' or cache == 'reuse' and os.path.exists(ribpath) == True:
            return True
        return False