#model_name=alilkirav4bis_GameModel
#created with IK-Rigger v1.2
#BY http://steamcommunity.com/id/OMGTheresABearInMyOatmeal/
#This is just a modified version of valves' biped simple script.
ik_1= [u'bip_UpperArm_r', u'bip_UpperArm_l', u'bip_hip_l', u'bip_hip_r']
ik_2= [u'bip_lowerArm_r', u'bip_lowerArm_l', u'bip_knee_l', u'bip_knee_r']
ik_3= [u'bip_hand_r', u'bip_hand_l', u'bip_foot_l', u'bip_foot_r']
ik_elbow_name= [u'Elbow_r', u'Elbow_l', u'Knee_l', u'Knee_r']
ik_hand_name= [u'Hand_r', u'Hand_l', u'Foot_l', u'Foot_r']
ik_spinelist= [u'bip_spine_0', u'bip_spine_1', u'bip_spine_2', u'bip_spine_3']
rig_to_rig_parent= {u'Foot_l': u'rigRoot', u'Knee_r': u'rigRoot', u'Elbow_r': u'rigRoot', u'Hand_r': u'rigRoot', u'Hand_l': u'rigRoot', u'Knee_l': u'rigRoot', u'Elbow_l': u'rigRoot', u'Foot_r': u'rigRoot', u'spine_1': u'rigRoot'}
bone_to_rig_parent= {u'bip_neck': u'spine_4', u'bip_collar_l': u'spine_4', u'bip_collar_r': u'spine_4', u'bip_pelvis': u'spine_1', u'bip_head': u'Neck'}
rig_name_to_bone= {u'Hip': u'bip_pelvis', u'Head': u'bip_head', u'CollarL': u'bip_collar_l', u'CollarR': u'bip_collar_r', u'Neck': u'bip_neck'}
axis_offset=[[0.0, 10.0, 0.0], [2.0, 10.0, 0.0], [0.0, -10.0, 0.0], [0.0, -10.0, 0.0]]

dic={u'RigBody': [u'spine_4', u'spine_3', u'spine_2', u'spine_1', u'rig_root', u'Neck', u'Hip', u'Head'], u'RigArms': [u'Elbow_l', u'Elbow_r', u'Hand_l', u'Hand_r'], u'RigLegs': [u'Knee_r', u'Knee_l', u'Foot_r', u'Foot_l'], u'Face': [u'flex controls']}
colortext={}
parents=[u'Face', u'RigBody', u'RigArms', u'RigLegs']







import random,ast
import vs


#==================================================================================================
def AddValidObjectToList( objectList, obj ):
    if ( obj != None ): objectList.append( obj )
    
def invert_dict(d):
    return dict([(v, k) for k, v in d.iteritems()])




#==================================================================================================
def get_childeren(rig_name):
    temp=[]

    
    for elm in bone_to_rig_parent:
        
        if bone_to_rig_parent[elm]==rig_name:
            temp.append(elm)
            temp.extend(get_childeren( invert_dict(rig_name_to_bone)[elm]))
    return temp

    

#==================================================================================================
#gets string list of bone's rig name
def bone_to_rigname(bone_list):
    temp=[]
    tempdic=invert_dict(rig_name_to_bone)
    for elm in bone_list:
        temp.append(tempdic[elm])

    return temp

def get_list_of_fingers():
    fingers=[]
    for key in bone_to_rig_parent.keys():
        if  bone_to_rig_parent[key] in ik_hand_name:
            fingers.append(key)
            fingers.extend(get_childeren( invert_dict(rig_name_to_bone)[key]))

            

            
    return bone_to_rigname(fingers)


#==================================================================================================
def HideControlGroups( rig, rootGroup, *groupNames ):
    for name in groupNames:    
        group = rootGroup.FindChildByName( name, False )
        if ( group != None ):
            rig.HideControlGroup( group )

    

#==================================================================================================
# Compute the direction from boneA to boneB
#==================================================================================================
def ComputeVectorBetweenBones( boneA, boneB, scaleFactor ):
    
    vPosA = vs.Vector( 0, 0, 0 )
    boneA.GetAbsPosition( vPosA )
    
    vPosB = vs.Vector( 0, 0, 0 )
    boneB.GetAbsPosition( vPosB )
    
    vDir = vs.Vector( 0, 0, 0 )
    vs.mathlib.VectorSubtract( vPosB, vPosA, vDir )
    vDir.NormalizeInPlace()
    
    vScaledDir = vs.Vector( 0, 0, 0 )
    vs.mathlib.VectorScale( vDir, scaleFactor, vScaledDir )    
    
    return vScaledDir
    
   
#==================================================================================================
# Build a simple ik rig for the currently selected animation set
#==================================================================================================
def BuildRig(helpoption):
    
    # Get the currently selected animation set and shot
    shot = sfm.GetCurrentShot()
    animSet = sfm.GetCurrentAnimationSet()
    gameModel = animSet.gameModel
    rootGroup = animSet.GetRootControlGroup()
    
    # Start the biped rig to which all of the controls and constraints will be added
    rig = sfm.BeginRig( "rig_biped_" + animSet.GetName()+str(random.randint(1,500)) );
    if ( rig == None ):
        return
    
    # Change the operation mode to passthrough so changes chan be made temporarily
    sfm.SetOperationMode( "Pass" )
    
    # Move everything into the reference pose
    sfm.SelectAll()
    sfm.SetReferencePose()
    
    #==============================================================================================
    # Find the dag nodes for all of the bones in the model which will be used by the script
    #==============================================================================================
    finger_group=get_list_of_fingers()
    boneRoot      = sfmUtils.FindFirstDag( [ "RootTransform" ], True )




    boneupper=[]
    ##[handlename:bone]
    boneelbow=[]
    bonehand=[]
    for i in range(len(ik_1)):

        
        boneupper.append(sfmUtils.FindFirstDag( [ ik_1[i] ], True ))
        boneelbow.append(sfmUtils.FindFirstDag( [ ik_2[i] ], True ))
        bonehand.append(sfmUtils.FindFirstDag( [ ik_3[i] ], True ))


    bone_spine=[]
    for i in range(len(ik_spinelist)):
        bone_spine.append(sfmUtils.FindFirstDag( [ ik_spinelist[i] ], True ))
        

           




##replace bone name string key with sfm bone obj
    for bone in bone_to_rig_parent.keys():
        temp=sfmUtils.FindFirstDag( [ bone ], True )
        bone_to_rig_parent[temp]=bone_to_rig_parent.pop(bone)
        
        for key in rig_name_to_bone.keys():
            if rig_name_to_bone[key] == bone:
                rig_name_to_bone[key]=temp



    #==============================================================================================
    # Create the rig handles and constrain them to existing bones
    #==============================================================================================
    rigRoot    = sfmUtils.CreateConstrainedHandle( "rig_root",     boneRoot,    bCreateControls=False )
 #maybe not needed   rig_bone["rigRoot"]=rigRoot
    
    for key in bone_to_rig_parent.keys():
  
        if bone_to_rig_parent[key] == "rigRoot":
            
            bone_to_rig_parent[key]=rigRoot
            
    for key in rig_to_rig_parent.keys():
  
        if rig_to_rig_parent[key] == "rigRoot":
            
            rig_to_rig_parent[key]=rigRoot
    
    rig_to_rig_Helper={}
    rig_hand=[]
    for i in range(len(ik_3)):
        rig_hand.append(sfmUtils.CreateConstrainedHandle( ik_hand_name[i],     bonehand[i],    bCreateControls=False ))
            # Create a helper handle which will remain constrained to the each foot position that can be used for parenting.
        temp= sfmUtils.CreateConstrainedHandle( ("rig_footHelper"+str(i)), bonehand[i], bCreateControls=False )
        rig_to_rig_Helper[rig_hand[-1]]=temp

        for key in bone_to_rig_parent.keys():
            if bone_to_rig_parent[key] == ik_hand_name[i]:
                bone_to_rig_parent[key]=rig_hand[-1]


        for key in rig_to_rig_parent.keys():
            if rig_to_rig_parent[key] == ik_hand_name[i]:
                rig_to_rig_parent[key]=rig_hand[-1]
            if key == ik_hand_name[i]:
                rig_to_rig_parent[rig_hand[-1]]=rig_to_rig_parent.pop(key)

                
    
    rig_spine=[]
    for i in range(len(bone_spine)):
        rig_spine.append(sfmUtils.CreateConstrainedHandle( "spine_"+str(i+1),     bone_spine[i],    bCreateControls=False ))
        
        a=("spine_"+str(i+1))

        for key in bone_to_rig_parent.keys():
            if bone_to_rig_parent[key] == a:
                bone_to_rig_parent[key]=rig_spine[-1]


        for key in rig_to_rig_parent.keys():
            if rig_to_rig_parent[key] == a:
                rig_to_rig_parent[key]=rig_spine[-1]
            if key == a:
                rig_to_rig_parent[rig_spine[-1]]=rig_to_rig_parent.pop(key)



        
 ######other            
              
    rig_bones=[]
    sfm_bones=[]
    bonerig_to_parent={}





#creates rigs thens copys the rig object to other dic that has it
##might need to optimize this part
    for key in rig_name_to_bone.keys():
        rig_bones.append(sfmUtils.CreateConstrainedHandle(key ,     rig_name_to_bone[key],    bCreateControls=False ))

        sfm_bones.append(rig_name_to_bone[key])

       # for x in bone_to_rig_parent.keys():
        if bone_to_rig_parent[rig_name_to_bone[key]]== key:
            print("yesssssss")
            bone_to_rig_parent[rig_name_to_bone[key]]=rig_bones[-1]

        for i in rig_to_rig_parent.keys():
            if rig_to_rig_parent[i] == key:
                rig_to_rig_parent[i]=rig_bones[-1]

        bonerig_to_parent[rig_bones[-1]]=bone_to_rig_parent[rig_name_to_bone[key]]

        
        
    for x in bonerig_to_parent.keys():
        if bonerig_to_parent[x] in rig_name_to_bone:  
            dex=sfm_bones.index(rig_name_to_bone[bonerig_to_parent[x]])
            bonerig_to_parent[x]=rig_bones[dex]



        
    
    # Use the direction from the heel to the toe to compute the knee offsets, 
    # this makes the knee offset indpendent of the inital orientation of the model.
    #vKneeOffsetR = ComputeVectorBetweenBones( boneFootR, boneToeR, 10 )
    #vKneeOffset = ComputeVectorBetweenBones( bonehand[0], bonehand[0], 10 )
    rig_elbows=[]
    
    
    for i in range(len(ik_2)):
        vKneeOffset = vs.Vector( axis_offset[i][0], axis_offset[i][1], axis_offset[i][2] )
        rig_elbows.append((sfmUtils.CreateOffsetHandle( ik_elbow_name[i], boneelbow[i], vKneeOffset,  bCreateControls=False )))
        
        for key in bone_to_rig_parent.keys():
            if bone_to_rig_parent[key] == ik_elbow_name[i]:
                bone_to_rig_parent[key]=rig_elbows[-1]
                
        for key in rig_to_rig_parent.keys():
            if rig_to_rig_parent[key] == ik_elbow_name[i]:
                rig_to_rig_parent[key]=rig_elbows[-1]
            if key == ik_elbow_name[i]:
                rig_to_rig_parent[rig_elbows[-1]]=rig_to_rig_parent.pop(key)












    # Create a list of all of the rig dags
    allRigHandles = [rigRoot]
    
    allRigHandles.extend(rig_hand+rig_elbows+rig_spine+rig_bones+rig_to_rig_Helper.keys())
    #==============================================================================================
    # Generate the world space logs for the rig handles and remove the constraints
    #==============================================================================================

    
    sfm.ClearSelection()
    sfmUtils.SelectDagList( allRigHandles )
    sfm.GenerateSamples()
    sfm.RemoveConstraints() 

    
    #==============================================================================================
    # Build the rig handle hierarchy
    #==============================================================================================
    #sfmUtils.ParentMaintainWorld( rigPelvis,        rigRoot )

    #for i in range(len (rig_hand)):
      #  sfmUtils.ParentMaintainWorld( rig_hand[i],        rigRoot )
    #    sfmUtils.ParentMaintainWorld( rig_elbows[i],        rig_hand[i] )
    


    for i in range(len(rig_spine)-1):
        sfmUtils.ParentMaintainWorld( rig_spine[i+1],        rig_spine[i] )




    for key in rig_to_rig_parent.keys():

        sfmUtils.ParentMaintainWorld( key,        rig_to_rig_parent[key] )
        
    for key in bonerig_to_parent.keys():

        if bonerig_to_parent[key]in rig_to_rig_Helper and  helpoption: #for bones parented to hands and foot i.e fingers,toes
            sfmUtils.ParentMaintainWorld( key,        rig_to_rig_Helper[bonerig_to_parent[key]] )
            continue

            
        sfmUtils.ParentMaintainWorld( key,        bonerig_to_parent[key] )
        




    # Set the defaults of the rig transforms to the current locations. Defaults are stored in local
    # space, so while the parent operation tries to preserve default values it is cleaner to just
    # set them once the final hierarchy is constructed.
    sfm.SetDefault()
    
    
    
    #==============================================================================================
    # Create constraints to drive the bone transforms using the rig handles
    #==============================================================================================
    
    # The following bones are simply constrained directly to a rig handle
    sfmUtils.CreatePointOrientConstraint( rigRoot,      boneRoot        )
    for i in range(len(rig_spine)):
        sfmUtils.CreatePointOrientConstraint( rig_spine[i],      bone_spine[i]        )


    i=0
    for key in rig_name_to_bone.keys():
        if key in finger_group:
            
            CreateOrientConstraint(rig_bones[i],      rig_name_to_bone[key]       )

        else:
          sfmUtils.CreatePointOrientConstraint(rig_bones[i],      rig_name_to_bone[key]       )  
        i+=1






    # Create ik constraints for the arms and legs that will control the rotation of the hip / knee and 
    # upper arm / elbow joints based on the position of the foot and hand respectively.

    for i in range(len(ik_1)):
        sfmUtils.BuildArmLeg( rig_elbows[i], rig_hand[i],      boneupper[i],  bonehand[i], True )
    



    
    #==============================================================================================
    # Create handles for the important attachment points 
    #==============================================================================================    
    attachmentGroup = rootGroup.CreateControlGroup( "Attachments" )  
    attachmentGroup.SetVisible( False )
    
    sfmUtils.CreateAttachmentHandleInGroup( "pvt_heel_R",       attachmentGroup )
    sfmUtils.CreateAttachmentHandleInGroup( "pvt_toe_R",        attachmentGroup )
    sfmUtils.CreateAttachmentHandleInGroup( "pvt_outerFoot_R",  attachmentGroup )
    sfmUtils.CreateAttachmentHandleInGroup( "pvt_innerFoot_R",  attachmentGroup )
    
    sfmUtils.CreateAttachmentHandleInGroup( "pvt_heel_L",       attachmentGroup )
    sfmUtils.CreateAttachmentHandleInGroup( "pvt_toe_L",        attachmentGroup )
    sfmUtils.CreateAttachmentHandleInGroup( "pvt_outerFoot_L",  attachmentGroup )
    sfmUtils.CreateAttachmentHandleInGroup( "pvt_innerFoot_L",  attachmentGroup )
    

    #==============================================================================================
    # Re-organize the selection groups
    #==============================================================================================
    topLevelColor = vs.Color( 0, 128, 255, 255 )
    grouphold={}
    HideControlGroups( rig, rootGroup, "Body", "Arms", "Legs", "Root" )
    face = rootGroup.FindChildByName( "Face", False )
    rootGroup.MoveChildToBottom( face )
    for i in parents:
        grouphold[i]=(rootGroup.CreateControlGroup(str(i)))
        rootGroup.MoveChildToBottom( grouphold[i] )
        
        if i in colortext.keys():


            rgb=ast.literal_eval(colortext[i].strip("'"))

            
    
            col = vs.Color( rgb[0],rgb[1],rgb[2],rgb[3])
            
            
            grouphold[i].SetGroupColor( col, False )
        else:
            grouphold[i].SetGroupColor( topLevelColor, False )
        for key in dic:
            if key == i:#2nd level parent
                for child in dic[key]:
                    if child in dic: ##if child also a parent
                        grouphold[child]=(rootGroup.CreateControlGroup(str(child)))
                        #print(grouphold[child])
                        grouphold[i].AddChild(grouphold[child])
                        if child in colortext.keys():
                            rgb=ast.literal_eval(colortext[child].strip("'"))
                            col = vs.Color( rgb[0],rgb[1],rgb[2],rgb[3])
                            grouphold[child].SetGroupColor( col, False )
                         
                        
                   # for rig in allRigHandles:
                    #    if child in str(rig):
                            #print(child)
                     #       sfmUtils.AddDagControlsToGroup(  grouphold[i], rig  )
                     
        
                        for subkey in dic[child]:##3rd level parent 
                            if subkey in dic:
                                
                                grouphold[subkey]=(rootGroup.CreateControlGroup(str(subkey)))
                                grouphold[child].AddChild(grouphold[subkey])


                                if subkey in colortext.keys():
                                    rgb=ast.literal_eval(colortext[subkey].strip("'"))
                                    col = vs.Color( rgb[0],rgb[1],rgb[2],rgb[3])
                                    grouphold[subkey].SetGroupColor( col, False )

    #adds rig contorols to each parent
    for i in grouphold:
        for key in dic:
            if key == i:
                for child in dic[key]:

                    for rig in allRigHandles:
                        if child in str(rig):
                            #print(child)
                            sfmUtils.AddDagControlsToGroup(  grouphold[i], rig  )





    rigHelpersGroup = rootGroup.CreateControlGroup( "RigHelpers" )
    rigHelpersGroup.SetVisible( False )
    rigHelpersGroup.SetSnappable( False )

    #sfmUtils.AddDagControlsToGroup( RigArms, rigRoot  )
 
        
    
    
    


    for i in rig_to_rig_Helper.values():
        sfmUtils.AddDagControlsToGroup( rigHelpersGroup,i)
   # sfmUtils.MoveControlGroup( "rig_footRoll_L", rigLegsGroup, LeftLegGroup )
    #sfmUtils.MoveControlGroup( "rig_footRoll_R", rigLegsGroup, RightLegGroup )



    
    #Re-order the groups
    other=rootGroup.FindChildByName( "Other", False )
    unknown=rootGroup.FindChildByName( "Unknown", False )
    fingersGroup = rootGroup.FindChildByName( "Fingers", False )

    rootGroup.MoveChildToBottom( fingersGroup )
    rootGroup.MoveChildToBottom( other )
    rootGroup.MoveChildToBottom( unknown )
    
    # End the rig definition
    sfm.EndRig()
    return
    
#==================================================================================================
# Script entry
#==================================================================================================

# Construct the rig for the selected animation set
sfm.EndRig()



BuildRig(False);