// ===========================================================================
// Copyright 2018 Autodesk, Inc. All rights reserved.
//
// Use of this software is subject to the terms of the Autodesk license
// agreement provided at the time of installation or download, or which
// otherwise accompanies this software in either electronic or hard copy form.
// ===========================================================================
// 
// hikControlRigUtils.mel
//
// Description:
//		Utility methods for HIK Control Rigs. Should not contain
//		any UI code. All code that references UI or triggers UI
// 		updates should live in hikControlRigUI.mel or 
//		hikControlRigOperations.mel
//


///////////////////////////////////////////////////////////////////////
// Global identifiers
///////////////////////////////////////////////////////////////////////

global int $gHIKneedSyncOnSetKeyframe = 0; // used to determine if a HIK rig synck is required on next setKeyframe


///////////////////////////////////////////////////////////////////////
// Local methodes...
///////////////////////////////////////////////////////////////////////

// 
// Remove all unused multi's of type $attr from $node
//
proc cleanupUnusedMultis( string $node, string $attr )
{
    // List existing elements of the multi atttribute plug
    string $plug = $node + "." + $attr;
    string $connections[] = `listAttr -multi $plug`;

    // For each attribute ...
    for ( $c in $connections )
    {
        // Make sure it is the attribute is a multi ...
        string $p = $node + "." + $c;

        // ... that it is not of the form node.attr[0].subattr
        string $buf[];
        if ( tokenize( $p, ".", $buf ) > 2 )
            continue;

        // ... and that it is not connected to anything
        string $n[] = `listConnections -s 1 $p`;
        if ( size( $n ) != 0 )
            continue;

        // If all these conditions are met, remove it.
        removeMultiInstance  $p;
    }  
}

proc cleanupDestinationMultis( string $masterEff, string $attr, string $nodeType )
{
    string $masterEffPlug = $masterEff + "." + $attr;

    // Get a list of all nodes connected to the specified attribute on the master effector
    string $masterEffConn[] = `listConnections -s 0 -d 1 -p 1 -type $nodeType $masterEffPlug`;
    if(size($masterEffConn) == 0)
        return;

    // Tokenize the returned attribute name.
    string $buffer[];
    tokenize($masterEffConn[0], ".",$buffer);

    // If it is malformed (i.e. not of the form node.attr[index] ), exit
    if ( size($buffer) <= 1 )
        return;

    // Separate the index from the attribute name
    string $destNode      = $buffer[0];
    string $attrWithIndex = $buffer[1];
    tokenize($attrWithIndex, "[",$buffer);
    string $attrNoIndex = $buffer[0];

    // Remove unused multi attribute plugs
    cleanupUnusedMultis( $destNode, $attrNoIndex );
}

proc string[] getAuxEffectors( string $effector )
{
    return `listConnections -d 1 -type hikIKEffector $effector`;
}

proc string getPinningObject( string $pRigElement )
{
	// TODO Could be optimised for IK elems
	
	string $rig = hikGetControlRigFromRigElement( $pRigElement );
	string $pinningObject;
	
	if( $rig != "" )
	{
		string $effectors[] = `listConnections -d 0 -s 1 -type hikIKEffector $rig`;
		
		if( size( $effectors ) > 0 )
		{
		    string $pinPlug = $effectors[0] + ".pinT";
		    string $pinningObjects[]  = `listConnections -d 1 -s 0 -type HIKPinning2State $pinPlug`;
		    
		    if( size( $pinningObjects ) == 1 )
		    {
				$pinningObject = $pinningObjects[0];
			}
		}
	}
	return $pinningObject;
}


proc getControlSetsAndPinningObj( string $IKs[], string $FKs[], string $controlSets[], string $pinningObj[]  )
{
	for( $ik in $IKs )
	{

        if( hikIsAuxiliary( $ik ) )
        {
            continue;
        }

		string $controlSet = hikGetControlRigFromRigElement( $ik );

		if( $controlSet != "" && !stringArrayContains( $controlSet, $controlSets ) )
		{
			$controlSets[ size($controlSets) ] = $controlSet;

			string $obj = getPinningObject($ik);  

			if( $obj != "" )
			{
				$pinningObj[ size($pinningObj) ] = $obj;
			}
		}
	}

	for( $fk in $FKs )
	{
		string $controlSet = hikGetControlRigFromRigElement( $fk );

		if( $controlSet != "" && !stringArrayContains( $controlSet, $controlSets ) )
		{
			$controlSets[ size($controlSets) ] = $controlSet;

			string $obj = getPinningObject($fk);  

			if( $obj != "" )
			{
				$pinningObj[ size($pinningObj) ] = $obj;
			}
		}
	}

}

proc string CreateSolverNode()
{
	string $solver = `createNode "HIKSolverNode"`;
	setAttr ($solver+".isHistoricallyInteresting") 0;
	return $solver;
}




///////////////////////////////////////////////////////////////////////
// Global methodes...
///////////////////////////////////////////////////////////////////////


global proc string[] hikGetRigIKNodes(string $character)
{
	int $i;
	string $lIKNodes[];
	int $lMatched = 0;
	string $curIKNode;
	string $lControlRig = hikGetControlRig($character);

	if($lControlRig != "")
	{
		for($i = 0; $i < GetHIKEffectorCount(); $i++)
		{
			$curIKNode = hikGetIKEffectorNode($character, $i);

			if( $curIKNode != "" )
				$lIKNodes[$lMatched++] = $curIKNode;
		}
	}

	return $lIKNodes;
}

global proc string[] hikGetRigFKNodes(string $character, int $includeRefNode)
{
	int $i;
	string $curFkNode;
	int $lMatched = 0;
	string $lFKNodes[];
	string $lControlRig = hikGetControlRig($character);
	if( $lControlRig != "" )
	{
		for($i = 0; $i < hikGetNodeCount(); $i++)
		{
			if( $i == 0 && $includeRefNode == 0 )
				continue;
				
			$curFkNode = hikGetSkNode($lControlRig, $i);

			if( $curFkNode != "" )
				$lFKNodes[$lMatched++] = $curFkNode;
		}
	}

	return $lFKNodes;
}


global proc string hikGetControlRig(string $pCharacter)
{
	if( size($pCharacter) && objExists( $pCharacter ) )
	{
		string $controlset[] = `listConnections -type "HIKControlSetNode" $pCharacter`;
		if( size($controlset) > 0 )
			return $controlset[0];
	}

	return "";
}

global proc string hikGetOrBuildSolverNode( string $pCharacter )
{
	// This add a solver to a created character if that character has no solver attached to it
	string $solver[] = hikGetSolverFromCharacter($pCharacter);

	if( size($solver) == 0 )
	{
		// No character, build one
		
		string $propState = hikGetProperty2StateFromCharacter($pCharacter);

		$solver[0] = CreateSolverNode();
	
		setAttr ($solver[0]+".InputStance" ) 1;

		hikConnectPropState($propState, $solver[0] );
		hikConnectCharacterDefinition($pCharacter, $solver[0] );
	}

	string $lState2SK[] = hikGetState2SKFromCharacter($pCharacter);
	
	if( size($lState2SK) == 0 )
	{
		$lState2SK[0] = hikBuildSkFromState($pCharacter);
	
		hikConnectCharacterStateObject($solver[0], $lState2SK[0]);
	}
	return $solver[0];
}







// Helper function to return effector ID
global proc int hikGetEffectorID( string $node )
{
	string $idPlug = $node + ".effectorID";
	int $effectorID = `getAttr $idPlug`;
	return $effectorID;
}

global proc string hikGetControlRigFromEffector(string $pEffector)
{	
	string $controlRigNode[] = `listConnections -s 0 -d 1 -type HIKControlSetNode $pEffector`;
	if(size($controlRigNode) > 0)
	{
		return $controlRigNode[0];
	}
	return "";
}

// This funcion gets a mask of the selected body part that can later be used
// with the hikBodyPart command. In body part mode it returns the selected
// parts. Otherwise a full body mask is returned.
//
global proc int hikGetBodyPartSelectionMask( string $controlSet, int $warnOnEmpty )
{

	int $mask = 0;

    if( `optionVar -query keyFullBody` == 2 )
    {
        // Body part mode
        string $cmd = "hikBodyPart -mask";
  
	    int $lastNodeId = GetHIKEffectorCount();
	    for($i = 0; $i < $lastNodeId; $i++)
	    {
		    string $hikNodeName = GetHIKEffectorName($i);		
		    string $attrName = $controlSet + "." + $hikNodeName;
		    string $nodes[] = `listConnections $attrName`;		
	
            if( size( $nodes )>0  )
            {
                string $selectedNodes[] = `ls -sl $nodes`;
            
                if( size ( $selectedNodes ) > 0 )
                {
                    $cmd += " -forEffector "+$i;
                }
		    }
	    }
	    $lastNodeId = hikGetNodeCount();
	    for($i = 0; $i < $lastNodeId; $i++)
	    {
		    string $hikNodeName = GetHIKNodeName($i);		
		    string $attrName = $controlSet + "." + $hikNodeName;
		    string $nodes[] = `listConnections $attrName`;		
	
            if( size( $nodes )>0  )
            {
                string $selectedNodes[] = `ls -sl $nodes`;
            
                if( size ( $selectedNodes ) > 0 )
                {
                    $cmd += " -forNode "+$i;
                }
		    }
	    }

        $mask = `eval $cmd`;

        if( $warnOnEmpty && $mask == 0 )
        {
            warning (uiRes("m_hikControlRigUtils.kNothingSelected"));
        }
    }
    else
    {
        $mask = `hikBodyPart -fullMask`;
    }

	return $mask;
}

global proc int hikIsRotateOnlyFK( int $nodeId )
{
    // Despite its name, this function works for both FK and SK elements
	if(    $nodeId == 0  /*Reference*/
	    || $nodeId == 1  /*Hips*/
	    || $nodeId == 49 /*Hips translation*/ )
	{
		return false;
	}
	else
	{
		return true;	
	}
}

global  proc string hikGetSolverFromNode(string $node)
{
	if ( ! attributeExists( "InputCharacterState", $node ) )
		return "";

	string $lConn[] = `listConnections -s on -d off ($node + ".InputCharacterState")`;
	for($item in $lConn)
	{
		if(`objectType -isa "HIKSolverNode" $item`)
		{
			return $item;
		}
		else
		{
			string $lSolver = hikGetSolverFromNode($item);
			if($lSolver != "")
			{
				return $lSolver;
			}
		}
	}
	return "";
}

global proc hikSetKeyframeCB() 
{
    global int $gHIKneedSyncOnSetKeyframe;

    if( $gHIKneedSyncOnSetKeyframe )
    {
        $gHIKneedSyncOnSetKeyframe = 0;

        int $keyingMode = `optionVar -query keyFullBody`;    
         
        if( $keyingMode != 1 && $keyingMode != 2 )
        {
		    return;
	    }
        
      
        string $selIK[] = `ls -sl -type hikIKEffector`;
        string $selFK[] = `ls -sl -type hikFKJoint`;
    
        if( size($selIK) > 0 || size($selFK) > 0  ) 
        {
		    string $controlSets[]; 
		    string $pinningObjects[];
		    getControlSetsAndPinningObj( $selIK, $selFK, $controlSets, $pinningObjects );		

		    for( $controlSet in $controlSets )
		    {
    			hikSyncControlRig( $controlSet );
	    	}		
        }
	}	
}

global proc string[] hikGetFkNodesFromControlRig( string $controlSet )
{
	string $fkNodes[];
	int $arraySize = 0;	

	int $lastNodeId = hikGetNodeCount();
	for($i = 0; $i < $lastNodeId; $i++)
	{
		string $hikNodeName = GetHIKNodeName($i);		
		string $attrName = $controlSet + "." + $hikNodeName;
		string $node[] = `listConnections $attrName`;		
		
		if( size($node) > 0 )
		{
			$fkNodes[$arraySize] = $node[0];
			$arraySize += 1;
		}
	}
	return $fkNodes;
}

global proc string hikGetBaseEffector( string $effector )
{
	string $plugs[] = `listConnections -d 1 -s 0 -plugs on ($effector+".ControlSet")`;
	string $strippedPlug[];

	if( size( $plugs ) != 1 )
	{
		return "";
	}
	
	tokenize( $plugs[0], "[", $strippedPlug );
	
	string $dstPlug = $strippedPlug[0]+"[0]";
	
	string $lsrc[] = `listConnections -d 0 -s 1 $dstPlug`;

	return $lsrc[0];
}

global proc int hikIsAuxiliary( string $pEffector )
{
	// returns 1 if this is an auxiliary
	// Othersiwe if the ffector is valid you can assume it is a pivot

	int $isAuxEffector = 0;
	
	if( `nodeType $pEffector` == "hikIKEffector" )
	{
	    $isAuxEffector = `getAttr ($pEffector + ".auxEffector")`;
	}

	return $isAuxEffector;	
}

global proc string[] hikGetIkNodesFromControlRig( string $controlSet, int $includeAux )
{
	string $ikNodes[];
	int $arraySize = 0;	

	int $lastNodeId = GetHIKEffectorCount();
	for($i = 0; $i < $lastNodeId; $i++)
	{
		string $hikNodeName = GetHIKEffectorName($i);		
		string $attrName = $controlSet + "." + $hikNodeName;
		string $nodes[] = `listConnections $attrName`;		
		
		for( $node in $nodes )
		{
			$ikNodes[$arraySize] = $node;
			$arraySize += 1;

			if( !$includeAux )
			{
				break;				
			}
		}
	}
	return $ikNodes;
}

global proc string hikGetControlRigFromRigElement( string $pRigElement )
{
	string $rigs[] = `listConnections -d 1 -s 0 -type HIKControlSetNode $pRigElement`;

	if( size( $rigs ) == 1 )
	{
		return $rigs[0];
	}
	else
	{
		return "";
	}
}

global proc string  hikGetIKWriter(string $ctrlset)
{
	string $lIKWriters[] = `ls -type HIKState2Effector`;
	for( $lIKWriterIter in $lIKWriters )
	{
		string $lIKWriterIterAttr = $lIKWriterIter + ".InputEffectorState";
		string $lEffectorFromCharacters[] = `listConnections $lIKWriterIterAttr`;
		if(size($lEffectorFromCharacters))
		{
			string $lCharacter1[] = `listConnections -d true -s true -type HIKCharacterNode $lEffectorFromCharacters[0]`;
			string $lCharacter2[] = `listConnections -d true -s true -type HIKCharacterNode $ctrlset`;
			
			// Test predicate for IK writers:
			//	- Must be driven by the same HIKCharacterNode
			if(size($lCharacter1) && size($lCharacter2) && isSameObject( $lCharacter1[0], $lCharacter2[0]) == 1)
			{
				return $lIKWriterIter;
			}
		}		
	}	
	return "";
	
}

global proc string hikGetSrcOnEffector( string $plug )
{
	string $src[] = `listConnections -s 1 -d 0 -type "hikIKEffector" -plugs 1 $plug`;

	if( size( $src ) > 0 )
	{				
		return $src[0]; 
	}

	return $plug;
}


global proc string[] hikGetFKWriterAndSolverNode(string $ctrlset)
{
	string $result[];
	clear( $result );
	string $lFKWriters[] = `ls -type HIKState2FK`;
	for( $lFKWriterIter in $lFKWriters )
	{
		string $lCharacter1[] = `listConnections -d true -s true -type HIKCharacterNode $lFKWriterIter`;
		string $lCharacter2[] = `listConnections -d true -s true -type HIKCharacterNode $ctrlset`;
	
		// Test predicate for FK writers:
		//	- Must be driven by the same HIKCharacterNode
		//	- Input PGX must be connected.
		if(size($lCharacter1) && size($lCharacter2) && isSameObject( $lCharacter1[0], $lCharacter2[0]) == 1)
		{
			string $solverNode = hikGetSolverFromNode( $lFKWriterIter );
			if ( size( $solverNode ) > 0 ) // Make sure we have the right HIKState2FK node
			{
				$result[0] = $lFKWriterIter;
				$result[1] = $solverNode;
				return $result;
			}
		}
	}	
	return $result;
	
}

global proc hikSyncControlRig( string $pControlSet )
{
    global int $gHIKneedSyncOnSetKeyframe;

	string $lFKWriterAndSolverNode[];
	$lFKWriterAndSolverNode = hikGetFKWriterAndSolverNode( $pControlSet );
	$lIKWriter = hikGetIKWriter($pControlSet);
	// Sync FK/IK on SK
	if( size($lFKWriterAndSolverNode) && size($lIKWriter) )
	{		
		hikFKIKSync( $pControlSet, $lFKWriterAndSolverNode[0], $lFKWriterAndSolverNode[1], $lIKWriter, false);	
	}
    
    $gHIKneedSyncOnSetKeyframe = 0;

}

//
// 
//  Procedure Name:
//		hikManipStart
//
//  Description Name;
//	This function is called when the translation or rotation manipulor begin
//
//  Input Value:
//		None 
//
//  Output Value:
//		None
// 
global proc hikManipStart( int $translate, int $rotate ) 
{
    int $keyingMode = `optionVar -query keyFullBody`;    
          
    if( $keyingMode != 1 && $keyingMode != 2 )
    {
        return;
    }
	
    undoInfo -openChunk;

          
    string $selIK[] = `ls -sl -type hikIKEffector`;
    string $selFK[] = `ls -sl -type hikFKJoint`;
    
    if( size($selIK) > 0 || size($selFK) > 0  ) 
    {
		string $controlSets[]; 
		string $pinningObjects[];
		getControlSetsAndPinningObj( $selIK, $selFK, $controlSets, $pinningObjects );		
    
		// do a synch at the beginning of the manip
		// This ensure tht the manipulated Rig is align with the visual represenaion 
		// that the user tries to manipulate
    	for( $controlSet in $controlSets )
		{
			hikSyncControlRig( $controlSet );
		}		
		
        if( size( $controlSets  ) > 0 )
        {

            string $cmd = "hikManip -start";
             
            if( $keyingMode == 2 )
            {
                $cmd += " -bp";
            }
            else
            {
                $cmd += " -fb";
            }

            if( $translate )
            {
                $cmd += " -t";
            }
            if( $rotate )
            {
                $cmd += " -r";
            }

		    for( $controlSet in $controlSets )
		    {
                $cmd += " \""+$controlSet+"\"";
            }

            eval $cmd;
        }
       
	}
}

global proc hikManipStop() 
{

   int $keyingMode = `optionVar -query keyFullBody`;    
         
    if( $keyingMode != 1 && $keyingMode != 2 )
    {
		return;
	}
        
      
    string $selIK[] = `ls -sl -type hikIKEffector`;
    string $selFK[] = `ls -sl -type hikFKJoint`;
    
    if( size($selIK) > 0 || size($selFK) > 0  ) 
    {
		string $controlSets[]; 
		string $pinningObjects[];
		getControlSetsAndPinningObj( $selIK, $selFK, $controlSets, $pinningObjects );		

		for( $controlSet in $controlSets )
		{
			hikSyncControlRig( $controlSet );
		}		

	    // First thing to do is to push all the pinning value, because
	    // we gonna override them

        if( size( $controlSets ) > 0 )
        {
            hikManip -stop $controlSets;
        }
	}	

    undoInfo -closeChunk;
}


// NOTE: The following method is also called from hikcharactertoolcmd.cpp,
//       if you are about to modify it's name make sure to update the 
//       cpp file too.
//
// Remove all auxiliary pivot effectors. To make sure things are "wired-up"
// correctly in the future, cleanup unused multi attribute plugs
// in various nodes used by HIK.
// 
global  proc hikCleanupAuxPivots( string $parentEff )
{
    // make sure the parent's name is not empty
    if ( size( $parentEff ) == 0 )
        return;

    // Delete all child effectors
    string $effectors[] = getAuxEffectors( $parentEff );
    delete $effectors;
    
    // Delete dynamic attribs added via scripts
    cleanupDestinationMultis( $parentEff, "ControlSet", "HIKControlSetNode" );
}

// Get FK Nodes of a Rig
global proc string[] hikGetRigFkNodes(string $pCharacter, int $includeRefNode )
{
	int $i;
	string $curFkNode;
	int $lMatched = 0;
	string $lFKNodes[];
	string $lControlRig = hikGetControlRig($pCharacter);
	if( $lControlRig != "" )
	{
		for($i = 0; $i < hikGetNodeCount(); $i++)
		{
			if( $i == 0 && $includeRefNode == 0 )
				continue;
				
			$curFkNode = hikGetSkNode($lControlRig, $i);

			if( $curFkNode != "" )
				$lFKNodes[$lMatched++] = $curFkNode;
		}
	}

	return $lFKNodes;
}

// Get IK Effector Node assigned to a Rig Definition field
global proc string hikGetIKEffectorNode(string $pCharacter, int $nodeId)
{
	string $tmpNode[];
	string $curNode = "";
	string $nodeName = GetHIKEffectorName($nodeId);
	string $lControlRig = hikGetControlRig($pCharacter);
	if($lControlRig != "" && $nodeName != "")
	{
		string $tmpCharObjectName = $lControlRig + "." + $nodeName + "[0]";

		if (size(`listConnections $tmpCharObjectName`)>0)
		{
			$tmpNode = `listConnections $tmpCharObjectName`;
			$curNode = $tmpNode[0];
		}
	}

	return $curNode;
}

global
proc string[ ] hikGetSelectedEffectorsFromCharacter( string $character, int $baseEffectorsOnly )
{
	string $effectors[ ];
	if( size( $character ) )
	{
		string $ctrlRig = hikGetControlRig( $character );
		if( size( $ctrlRig ) )
		{
			string $selectedEffectors[ ] = `ls -sl -type hikIKEffector`;
			
			string $effector;
			for( $effector in $selectedEffectors )
			{
				// If selected effector is part of the specified character
				// append the base effector to the list
				//
				if( $ctrlRig == hikGetControlRigFromEffector( $effector ) )
					$effectors[ size( $effectors ) ] = ( $baseEffectorsOnly == 1 ? hikGetBaseEffector( $effector ) : $effector );
			}
		}
	}
	
	return $effectors;
}




///////////////////////////////////////////////////////////////////////
// Script commands
///////////////////////////////////////////////////////////////////////

// those are executed at the source of the file
manipMoveContext -e -preDragCommand ("hikManipStart 1 0") hikIKEffector Move;
manipMoveContext -e -postDragCommand hikManipStop hikIKEffector Move;
manipRotateContext -e -preDragCommand ("hikManipStart 0 1") ("hikIKEffector,hikFKJoint") Rotate;
manipRotateContext -e -postDragCommand hikManipStop ("hikIKEffector,hikFKJoint") Rotate;