// ===========================================================================
// 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.
// ===========================================================================
// 
// hikRigVisibilityUtils.mel
//
// Description:
//		Utility methods Related to the pakaging of character and retargetter in HIK
//



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

proc string createContainerWithIcon( string $name, string $nodes[], string $icon )
{
	// check if there is an existing container for these nodes
	string $container;
	
	for( $node in $nodes )
	{
		$container = `container -q -fc {$node}`;
		
		if( $container != "" )
		{
			break;
		}
	}

	
	if( $container == "" )
	{
		$container = `container -name $name`;		
	}
	else
	{
		//Container exists already, just rename it.
		if($container != $name)
		{
			string $str = "renaming " + $container + " to " + $name + "\n";
			print($str);
			
			string $tmpBuffer[];
			$tmpBuffer = `ls $container`;
			if(size($tmpBuffer) > 0)
			{
				rename $container $name;
				$container = $name;
			}
		}
	}
	
	if( $icon != "" )
	{
		setAttr -type "string" ($container+".iconName") $icon;
	}
	
	
	container -e -f -addNode $nodes $container;
		
	return $container;
}



proc string controlRigPackage( string $pCtrlRigNode )
{
	string $rigContainer;
	string $ikContainer;
	string $fkContainer;
	
	
	if( `objExists $pCtrlRigNode` && (`nodeType $pCtrlRigNode`=="HIKControlSetNode" ) )
	{
		
		string $characterNode[] = `listConnections -type HIKCharacterNode -s true -d false ($pCtrlRigNode+".InputCharacterDefinition")`;
		string $ikElements[]    = `listConnections -type hikIKEffector -s true -d false $pCtrlRigNode`;
		string $fkElements[]    = `listConnections -type hikFKJoint -s true -d false $pCtrlRigNode`;
		string $refElement[]    = `listConnections -s true -d false ( $pCtrlRigNode + ".Reference" )`;
		
		// Remove reference node if defined as an hikFKJoint
		//
		if( size( $refElement ) > 0 )
			$fkElements = stringArrayRemove( { $refElement[0] }, $fkElements );
		
		string $characterName = $pCtrlRigNode;

		if( size( $characterNode ) > 0  )
		{
			$characterName = $characterNode[0];
		}
				
		$rigContainer = createContainerWithIcon( ($characterName+"_Rig"), {$pCtrlRigNode}, "HIKcharacter.png" );
		 
		if( size( $ikElements ) > 0 )
		{
			// $ikContainer = `container -name ($characterName+"_Ik") -addNode $ikElements`;
			$ikContainer = createContainerWithIcon( ($characterName+"_Ik"), $ikElements, "HIKik.png" );
			container -e -f -addNode {$ikContainer} $rigContainer;
		}

		if( size( $fkElements ) > 0 )
		{
			// $fkContainer = `container -name ($characterName+"_Fk") -addNode $fkElements`;
			$fkContainer = createContainerWithIcon( ($characterName+"_Fk"), $fkElements, "HIKskel.png" );
			container -e -f -addNode {$fkContainer} $rigContainer;
		}
		
		// Add control rig reference to rig container
		//
		if( size( $refElement ) > 0 )
			container -e -f -addNode { $refElement[0] } $rigContainer;
	}
	return $rigContainer;
}

proc string solverPackage( string $pSolverNode )
{
	string $solverContainer;
	
	
	if( `objExists $pSolverNode` && (`nodeType $pSolverNode`=="HIKSolverNode" ) )
	{
		
		string $characterNode[] = `listConnections -type HIKCharacterNode -s true -d false ($pSolverNode+".InputCharacterDefinition")`;
		
		string $ikInNodes[] = `listConnections -type HIKPinning2State -s true -d false ($pSolverNode+".InputEffectorState")`;
		string $fkInNodes[]  = `listConnections -type HIKFK2State -s true -d false ($pSolverNode+".InputCharacterState")`;
		string $skOutNodes[] = `listConnections -type HIKState2SK -s false -d true ($pSolverNode+".OutputCharacterState")`;
		string $ikOutNodes[]  = `listConnections -type HIKEffectorFromCharacter -s false -d true ($pSolverNode+".OutputCharacterState")`;
		string $fkOutNodes[] = `listConnections -type HIKState2FK -s false -d true ($pSolverNode+".OutputCharacterState")`;
		string $propertyStateNodes[] = `listConnections -type HIKProperty2State -s true -d false ($pSolverNode+".InputPropertySetState")`;
		
		// Go deeper where it should
		
		if( size( $ikInNodes ) == 1 )
		{
			string $extraNodes[] = `listConnections -type HIKEffector2State -s true -d false ($ikInNodes[0]+".InputEffectorState")`;

			$ikInNodes = stringArrayCatenate( $ikInNodes, $extraNodes );
		} 
		
		if( size( $ikOutNodes ) == 1 )
		{
			string $extraNodes[] = `listConnections -type HIKState2Effector -s false -d true ($ikOutNodes[0]+".OutputEffectorState")`;

			$ikOutNodes = stringArrayCatenate( $ikOutNodes, $extraNodes );
			
			if( size( $extraNodes ) == 1 )
			{
				$extraNodes = `listConnections -type ComputeLocal -s false -d true $extraNodes[0]`;
				$ikOutNodes = stringArrayCatenate( $ikOutNodes, $extraNodes );			
			}
		} 
		

		if( size( $fkOutNodes ) == 1 )
		{
			string $extraNodes[] = `listConnections -type ComputeLocal -s false -d true $fkOutNodes[0]`;
			$fkOutNodes = stringArrayCatenate( $fkOutNodes, $extraNodes );			
		} 
		
		// Add nodes to the container and create subcontainers when needed
		string $characterName = $pSolverNode;

		if( size( $characterNode ) > 0  )
		{
			$characterName = $characterNode[0];
		}

						
		$solverContainer = createContainerWithIcon( ($characterName+"_Solver"), {$pSolverNode}, "HIKsolver.png" );
		 		 
		if( size( $ikInNodes ) > 0 )
		{
			// string $ikInContainer = `container -name ($characterName+"_IkIn") -addNode $ikInNodes`;
			string $ikInContainer = createContainerWithIcon( ($characterName+"_IkIn"), $ikInNodes, "HIKconcentrator.png" );
			container -e -f -addNode {$ikInContainer} $solverContainer;
		}

		if( size( $fkInNodes ) > 0 )
		{
			string $fkInContainer = createContainerWithIcon( ($characterName+"_FkIn"), $fkInNodes, "HIKconcentrator.png" );
			container -e -f -addNode {$fkInContainer} $solverContainer;
		}

		if( size( $skOutNodes ) > 0 )
		{
			string $skOutContainer = createContainerWithIcon( ($characterName+"_SkOut"), $skOutNodes, "HIKdeconcentrator.png" );
			container -e -f -addNode {$skOutContainer} $solverContainer;
		}

		if( size( $ikOutNodes ) > 0 )
		{
			string $ikOutContainer = createContainerWithIcon( ($characterName+"_IkOut"), $ikOutNodes, "HIKdeconcentrator.png" );
			container -e -f -addNode {$ikOutContainer} $solverContainer;
		}

		if( size( $fkOutNodes ) > 0 )
		{
			string $fkOutContainer = createContainerWithIcon( ($characterName+"_fkOut"), $fkOutNodes, "HIKdeconcentrator.png" );
			container -e -f -addNode {$fkOutContainer} $solverContainer;
		}

		if( size( $propertyStateNodes ) > 0 )
		{
			container -e -f -addNode $propertyStateNodes $solverContainer;
		}
	}
	return $solverContainer;
}


proc string findBindPose( string $pCharacterNode )
{
	string $bindPoseNode;
	
	string $skElements[] = `listConnections -type joint -s true -d false $pCharacterNode`;

	if( size( $skElements ) > 0 )
	{
		string $bindPoses[] = `listConnections -type dagPose -s false -d true ($skElements[0]+".bindPose")`;
		
		if( size( $bindPoses ) > 0 )
		{
			$bindPoseNode = $bindPoses[0];
		}
	}
	return $bindPoseNode;
}

proc string[] findSkNodesFromBindPose( string $pBindPose ) 
{
	string $skElements[] = `listConnections -type joint -s true -d false ($pBindPose+".worldMatrix")`;

	return $skElements;
}

proc string[] findIntermediateSkChilds( string $currentLeaf ) 
{
	string $lChilds[];
	
	// Look for child modelst that arent characterised but that could be part of the skeleton
	// Look only at joints here because we do not want to end-up with unrelated modesl here
	string $lImmediateChilds[] = `listRelatives -pa -type joint -children $currentLeaf`;
	
	
	for( $child in $lImmediateChilds )
	{		
		$lChilds[ size($lChilds) ] = $child;
		$lChilds = stringArrayCatenate( $lChilds, findIntermediateSkChilds( $child )  );		
	}
	return $lChilds;
}

proc string[] findIntermediateSkNodes( string $sk[] )
{
	string $intermediateNodes[];
	string $notLeafs[];
	
	for( $bone in $sk )
	{
		string $lCurrentChain[];
		string $lParent = firstParentOf( $bone );
		
		while( $lParent != "" )
		{
			for( $otherBone in $sk )
			{
				if( isSameObject( $otherBone, $lParent ) )
				{
					$lParent = "";
					
					$notLeafs[ size($notLeafs) ] = $otherBone;
					break;
				}
			}
			if( $lParent != "" )
			{
				stringArrayInsertAtIndex( 0, $lCurrentChain, $lParent );
				$lParent = firstParentOf( $lParent );
			}
			else
			{
				$intermediateNodes = stringArrayCatenate( $intermediateNodes, $lCurrentChain );
			}
		}
	}
	
	// Find the childs of the leafs that we may want to keep
	string $leafs[] = stringArrayRemove( $notLeafs, $sk );			
	
	for( $leaf in $leafs )
	{
		$intermediateNodes = stringArrayCatenate( $intermediateNodes, findIntermediateSkChilds( $leaf )  );	
	}
	
	return $intermediateNodes;
}

proc string findSrcCountainer( string $plug )
{
	string $container;
	string $srcNode[] = `listConnections -type HIKCharacterNode -s true -d false $plug`;
	
	if( size( $srcNode ) > 0 )
	{
		$container = `container -q -fc $srcNode`;
	}
	
	return $container;
} 


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


global proc string hikRetargeterPackage( string $pRetargeterNode )
{
	string $retargeterContainer;
		
	if( `objExists $pRetargeterNode` && (`nodeType $pRetargeterNode`=="HIKRetargeterNode" ) )
	{	
		string $srcPropertySetState[] = `listConnections -type HIKProperty2State -s true -d false ($pRetargeterNode+".InputSrcPropertySetState")`;
		string $dstPropertySetState[] = `listConnections -type HIKProperty2State -s true -d false ($pRetargeterNode+".InputDstPropertySetState")`;
		string $inSk2State[]		  = `listConnections -type HIKSK2State -s true -d false ($pRetargeterNode+".InputCharacterState")`;
		
		string $ikOutNodes[] = `listConnections -type HIKEffectorFromCharacter -s false -d true ($pRetargeterNode+".OutputCharacterState")`;
		string $fkOutNodes[] = `listConnections -type HIKState2FK -s false -d true ($pRetargeterNode+".OutputCharacterState")`;
		string $skOutNode[]  = `listConnections -type HIKState2SK -s false -d true ($pRetargeterNode+".OutputCharacterState")`;

		string $inputCharacterContainer = findSrcCountainer( $pRetargeterNode+".InputCharacterDefinitionSrc" );
		string $outputCharacterContainer = findSrcCountainer( $pRetargeterNode+".InputCharacterDefinitionDst" );		
		

		if( size( $ikOutNodes ) == 1 )
		{
			string $extraNodes[] = `listConnections -type HIKState2Effector -s false -d true ($ikOutNodes[0]+".OutputEffectorState")`;

			$ikOutNodes = stringArrayCatenate( $ikOutNodes, $extraNodes );
			
			if( size( $extraNodes ) == 1 )
			{
				$extraNodes = `listConnections -type ComputeLocal -s false -d true $extraNodes[0]`;
				$ikOutNodes = stringArrayCatenate( $ikOutNodes, $extraNodes );			
			}
		} 
		

		if( size( $fkOutNodes ) == 1 )
		{
			string $extraNodes[] = `listConnections -type ComputeLocal -s false -d true $fkOutNodes[0]`;
			$fkOutNodes = stringArrayCatenate( $fkOutNodes, $extraNodes );			
		} 
		
		// Add nodes to the container and create subcontainers when needed
		string $retargeterName = $pRetargeterNode;
						
		$retargeterContainer = createContainerWithIcon( ($retargeterName+"_Retargeter"), {$pRetargeterNode}, "HIKretargeter.png" );

		if( $inputCharacterContainer == "" )
		{
			$inputCharacterContainer = $retargeterContainer;
		}

		if( $outputCharacterContainer == "" )
		{
			$outputCharacterContainer = $retargeterContainer;
		}

		if( size( $srcPropertySetState ) > 0 )
		{
			container -e -f -addNode $srcPropertySetState $retargeterContainer;
		}
		if( size( $dstPropertySetState ) > 0 )
		{
			container -e -f -addNode $dstPropertySetState $retargeterContainer;
		}
		if( size( $inSk2State ) > 0 )
		{
			container -e -f -addNode $inSk2State $inputCharacterContainer;
		}

		if( size( $ikOutNodes ) > 0 )
		{
			string $ikOutContainer = createContainerWithIcon( ($retargeterName+"_IkOut"), $ikOutNodes, "HIKdeconcentrator.png" );
			container -e -f -addNode {$ikOutContainer} $outputCharacterContainer;
		}

		if( size( $fkOutNodes ) > 0 )
		{
			string $fkOutContainer = createContainerWithIcon( ($retargeterName+"_fkOut"), $fkOutNodes, "HIKdeconcentrator.png" );
			container -e -f -addNode {$fkOutContainer} $outputCharacterContainer;
		}
		if( size( $skOutNode ) > 0 )
		{
			container -e -f -addNode $skOutNode $outputCharacterContainer;
		}
		
		
		
		
	}
	return $retargeterContainer;
}


global proc string hikCharacterPackage( string $pCharacterNode )
{
	string $characterContainer;
	
	if( `objExists $pCharacterNode` && (`nodeType $pCharacterNode`=="HIKCharacterNode" ) )
	{
		string $characterName = $pCharacterNode;
		
		string $controlSetNode[] = `listConnections -type HIKControlSetNode -s false -d true ($pCharacterNode+".OutputCharacterDefinition")`;
		string $solverNode[] = `listConnections -type HIKSolverNode -s false -d true ($pCharacterNode+".OutputCharacterDefinition")`;
		string $referenceNode[] = `listConnections -type transform -s true -d false ($pCharacterNode+".Reference")`;
		string $lBindPose = findBindPose( $pCharacterNode );
		string $skElements[];
		
		if( $lBindPose != "" )
		{
			// If there is a bind pose use it to identify the bones because it is more precise
			$skElements = findSkNodesFromBindPose( $lBindPose );
		}
		else
		{
			$skElements = `listConnections -type transform -s true -d false $pCharacterNode`;
			string $extraNodes[] = findIntermediateSkNodes( $skElements );
			
			$skElements = stringArrayCatenate( $skElements, $extraNodes );
		}
				
		$characterContainer = createContainerWithIcon( ($characterName+"_Character"), {$pCharacterNode}, "HIKcharacter.png" );

		if( size( $referenceNode ) > 0 )
		{			
			container -e -f -addNode $referenceNode $characterContainer;
			$skElements = stringArrayRemove( $referenceNode, $skElements );			

			string $referenceLocator[] = `listRelatives -pa -type locator -children $referenceNode[0]`;
			
			if( size( $referenceLocator ) > 0 )
			{			
				container -e -f -addNode $referenceLocator $characterContainer;
				$skElements = stringArrayRemove( $referenceLocator, $skElements );			
			}
		}
		
		if( $lBindPose != "" )
		{
			container -e -f -addNode {$lBindPose} $characterContainer;			
		}
		 
		if( size( $skElements ) > 0 )
		{					
			string $skContainer = createContainerWithIcon( ($characterName+"_Skeleton"), $skElements, "HIKskel.png" );
			container -e -f -addNode {$skContainer} $characterContainer;		
		}

		if( size( $controlSetNode ) > 0 )
		{
			// Assuming only one controlSet per character
		
			string $rigContainer = controlRigPackage( $controlSetNode[0] );
			
			if( $rigContainer != "" )
			{
				container -e -f -addNode {$rigContainer} $characterContainer;
			}
		}
		
		if( size( $solverNode ) > 0 )
		{
			// Assuming only one solver per character		
			string $solverContainer = solverPackage( $solverNode[0] );
			
			if( $solverContainer != "" )
			{
				container -e -f -addNode {$solverContainer} $characterContainer;
			}
		}				
	}
	return $characterContainer;
}

global proc hikCharactersPackageAll()
{
	string $characters[] = `ls -type HIKCharacterNode`;

	for( $character in $characters )
	{
		hikCharacterPackage( $character );
	}
}

global proc hikRetargetersPackageAll()
{
	string $retargeters[] = `ls -type HIKRetargeterNode`;

	for( $retargeter in $retargeters )
	{
		hikRetargeterPackage( $retargeter );
	}

}


global proc hikPackageAll()
{
	hikCharactersPackageAll();
	hikRetargetersPackageAll();	
}


global proc hikPackage(string $type)
// Description:
//		Package up either "Current Character", "Current Character Retargeters",
//		"All Characters" or "All Character Retargeters"
//
{
    string $lower = tolower($type);
    if( $lower == "current character") {
        hikCharacterPackage(hikGetCurrentCharacter());
    }
    else if( $lower == "current character retargeters") {
        string $rets[ ] = hikGetRetargeterFromCharacter( hikGetCurrentCharacter( ) ); 
        for( $r in $rets ) 
            hikRetargeterPackage( $r );
    }
    else if( $lower == "all characters") {
        hikCharactersPackageAll();
    }
    else if( $lower == "all character retargeters") {
        hikRetargetersPackageAll();
    }
    else {
        warning( (uiRes("m_hikPackageUtils.kHIKPackageInvalidType")) );
    }		
		
}