// ===========================================================================
// 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.
// ===========================================================================

//<doc>
//<name dynamicConstraintMembership>
//<owner "Autodesk">
//
//<synopsis>
//		dynamicConstraintMemberShip( string $mode )
//
//<returns>
//		None.	
//
//<description>
//	This command is useful for editing and displaying constraint membership. It works off of the current selection.
//  If a constraint node is not selected it tries to guess at the desired one based on the selected
//  components. If the constraint to edit is ambiguous from the selection, then one will need to
//  include the constraint node in the selection. The mode determines how the selection is to affect
//  the constraint membership. It is equivalent to invoking the matching items on the nConstraint menu.
//
//<flags>
//	string $mode	Must be one of: "select", "replace", "add", "remove"
//
//<examples>
//  dynamicConstraintMembership "select";
//</doc>


proc int canSeeShape( string $shape )
{
	return( getAttr( $shape + ".visibility" ) && !getAttr( $shape + ".intermediateObject" ) );
}

proc string getComponentMesh( string $component )
{
	string $cons[] = `listConnections -sh 1 -type nBase $component`; 	
	if( size( $cons ) < 0 ){
		string $format = (uiRes("m_dynamicConstraintMembership.kNoInputNucleus"));
     	string $warnStr = `format -stringArg $component $format`;
		warning( $warnStr );
		return "";
	}
	if( nodeType( $cons[0] ) == "nParticle"){
		return $cons[0];
	}
	string $mesh = findTypeInHistory( $cons[0], "mesh", 1, 0 );		
	if( $mesh == "" ||  !canSeeShape( $mesh ) ){
		$mesh = findTypeInHistory( $cons[0], "mesh", 0, 1 );		
	}
	if( $mesh == "" ){
		string $format = (uiRes("m_dynamicConstraintMembership.kNoVisibleMesh"));
     	string $warnStr = `format -stringArg $cons[0] $format`;
		warning( $warnStr );
	}
	return $mesh;
}

proc string[] getSelectionArray( string $attr, float $indices[] )
{
	string $selArray[];
	float $index;
	for( $index in $indices ){
		$selArray[size($selArray)] = ($attr + "["+(int)$index+"]");	
	}
	return $selArray;
}

proc string getSetFromComponent( string $component )
{
	string $con[] = `listConnections ($component + ".componentGroupId")`;
	if( size( $con ) > 0 ){
		$con = `listConnections -type objectSet ($con[0] + ".message")`;
		if( size( $con ) > 0 ){
			return( $con[0] );
		}
	}
	return "";
}

proc selectAppendConstraintMembers( string $constraint )
{
	string $components[] = `listConnections -type nComponent $constraint`;
	if( size($components) < 1 ){
		string $format = (uiRes("m_dynamicConstraintMembership.kNoConnectComponent"));
     	string $warnStr = `format -stringArg $constraint $format`;
		warning( $warnStr );
		return;
	}
	string $component;	
	for( $component in $components ){
		string $mesh 	= getComponentMesh( $component );
		if( $mesh == "" ){
			continue;
		}
		
		int $elements 	= getAttr( $component + ".elements" );
		if( $elements == 0 ){
			// If the constraint uses sets then select the set components
			// Note that if the input mesh is not visible we need to select
			// the set components on the output mesh, so we use subsitute
			// to replace the component members with the output mesh. 
			// The component sets are always maintained on the input mesh, however,
			// in order to avoid dependancy graph loop conditions.
			string $set = getSetFromComponent( $component );
			if( $set != "" ){
				string $setMembers[] = `sets -q $set`;
				int $numMembers = size( $setMembers );
				if( $numMembers > 0 ){
					string $buffer[];	
					int $numTokens = `tokenize $setMembers[0] "." $buffer`;		
					if( $numTokens > 1 ){
						string $setObj = $buffer[0];
						int $i;
						for( $i = 0; $i < $numMembers; $i++){
							$setMembers[$i] = `substitute $setObj $setMembers[$i] $mesh`;
						}
						select -add $setMembers;
					}
				}
				continue;
			}
		}
		
		int $type 		= getAttr( $component + ".componentType");
		float $indices[]; // needs to be float for getAttr
		if( $elements == 0 ){
			$indices = getAttr( $component + ".componentIndices" );
		}
		string $attr;
		int $totalElements[];
		$totalElements[0] = 0;
		if( $type == 0 ){  // none
			continue;
		} else if( $type == 6 ){ // object
			select -add $mesh;
			continue;
		} else if( $type == 2 ){ // point
			if( nodeType( $mesh ) == "nParticle" ){
				$attr = ($mesh + ".pt");
				if( $elements == 2 ){
					$totalElements[0] = getAttr ( $mesh +".count");
				}
			} else {
				$attr = ($mesh + ".vtx");
				if( $elements == 2 ){
					$totalElements = `polyEvaluate -v $mesh`;
				}
			}
		} else if( $type == 3 ){ // edge
			$attr = ($mesh + ".e");
			if( $elements == 2 ){
				$totalElements = `polyEvaluate -e $mesh`;
			}
		} else if( $type == 4 ){ // face
			$attr = ($mesh + ".f");
			if( $elements == 2 ){
				$totalElements = `polyEvaluate -f $mesh`;
			}
		} else {
			string $fmt = (uiRes("m_dynamicConstraintMembership.kUnsupportedType"));
			warning( `format -s $component -s $type $fmt`);
			continue;
		}
		if( $elements == 2 ){
			if( $totalElements[0] > 0 ){
				select -add ($attr + "[0:" + ($totalElements[0]-1) + "]"); 
			}
		} else {
			string $items[] = getSelectionArray($attr, $indices);
			select -add $items;
		}
	}
}

proc selectMembers()
{
	string $constraints[] = `ls -sl -dag -type dynamicConstraint`;

	if( size( $constraints ) < 1 ){
		warning( (uiRes("m_dynamicConstraintMembership.kNoConstraints")));
		return;
	}
	select -clear;
	string $constraint;
	for( $constraint in $constraints ){
		selectAppendConstraintMembers( $constraint );
	}
}


proc int inComponentList( string $component, int $item )
{
	float $indices[] = getAttr( $component + ".componentIndices" );
	float $indice;
	for( $indice in $indices ){
		if( $item == (int)$indice ){
			return true;
		}
	}
	return false;
}

proc string getNObjectFromMesh( string $mesh )
{
	return( findTypeInHistory( $mesh, "nBase", 1,1 ) );
}



proc string [] getNConstraintComponents( string $constraint )
{
	string $con[]= `listConnections -type nComponent ($constraint + ".componentIds")`;
	return $con;
}

proc replaceMembers(string $constraint, string $sel[])
{
	string $components[] = getNConstraintComponents( $constraint );
	string $component, $item, $info[2];
	int $i;
	for( $component in $components ){
		int $cType = getAttr( $component + ".componentType");
		if( $cType == 6 ){ // surface
			continue; // for now... implement surface replacement later
		}
		int $elements = getAttr( $component + ".elements");
		int $newIndices[] = {};
		string $mesh = "";
		for( $item in $sel ){
			$info[0] = $item;
			$info[1] = 0;
			int $compType = itemComponentInfo( $item, $info );
			if( $compType == 7 ){
				$compType = 2; // particles are defined as vertices on component
			}
			if( $compType != $cType ){
				continue;
			}
			if( $mesh != "" ){
				if( $mesh != $info[0] ){
					continue;
				}
			} else {
				// determine if this mesh is the one for this component
				string $nObj = getNObjectFromMesh( $info[0] );
				if( $nObj == "" ){
					continue;
				}
				string $cons[] = `listConnections -type nComponent ($nObj + ".nucleusId")`; 	
				if( stringArrayContains( $component, $cons ) ){
					$mesh = $info[0];
				} else {
					continue;
				}
			}	
			// this is a valid element for the current component
			int $index = $info[1];
			$newIndices[size($newIndices)] = $index;
		}
		int $numIndices = size( $newIndices );
		if( $numIndices < 1 ){
			continue; // Unclear if we should delete a component when the new selection does not select any elements
					// For now we simply leave the component unchanged. This allows one to do a replace without
					// needing to select elements of all component of the constraint.
		}

		string $set = getSetFromComponent( $component );
		if( $set != "" ){
			string $nObj = getNObjectFromMesh( $mesh );
			string $baseItem;
			if( $cType == 2 ){ // point
				$baseItem = ".vtx";
			} else if( $cType == 3 ){ // edge
				$baseItem = ".e";
			} else if( $cType == 4 ){ // face
				$baseItem = ".f";
			} else {
				string $fmt = (uiRes("m_dynamicConstraintMembership.kReplaceUnsupportedCTypeWrn"));
				warning(`format -s $cType $fmt`);
				continue;
			}
			string $inputMesh = findTypeInHistory( $nObj, "mesh", 0, 1 );
			if( $inputMesh == "" ){
				warning( (uiRes("m_dynamicConstraintMembership.kReplaceNoInputMeshWrn")));
				continue;
			}

			int $intObj = getAttr( $inputMesh + ".intermediateObject");
			if( $intObj ){
				setAttr( $inputMesh + ".intermediateObject") false;
			}
			$baseItem = ($inputMesh + $baseItem);
			string $setElements[] = `sets -q $set`;

			// remove all but the first element, otherwise the set automatically
			// breaks needed connections
			for( $i = size( $setElements ) -1; $i > 0; $i-- ){
				sets -remove $set $setElements[$i];
			}
			int $firstAlreadyInSet = false;
			for( $i = 0; $i < $numIndices; $i++ ){
				string $setItem = ($baseItem + "[" + $newIndices[$i] + "]");
				if( $setItem == $setElements[0] ){
					$firstAlreadyInSet = true;
				} else {
					sets -addElement $set $setItem;
				}
			}	
			if( !$firstAlreadyInSet ){
				sets -remove $set $setElements[0];
			}
			if( $intObj ){
				setAttr( $inputMesh + ".intermediateObject") true;
			}
		} else {
			if( $elements != 0 ){
				setAttr ( $component + ".elements") 0;
			}
			string $indAttr = ($component + ".componentIndices");
			int $oldNumIndices = `getAttr -size $indAttr`;
			for( $i = 0; $i < $numIndices; $i++ ){
				setAttr ($indAttr + "[" + $i + "]") $newIndices[$i];
			}
			while( $i < $oldNumIndices ){
				removeMultiInstance -b true ($indAttr + "[" + $i + "]");
				// if shrunk then unlink old attrs at end of multi 
				$i++;
			}
		}
	}
}

proc addMembers(string $constraint, string $sel[])
{
	string $components[] = getNConstraintComponents( $constraint );
	string $component, $item, $info[2];
	int $i;
	for( $component in $components ){
		int $cType = getAttr( $component + ".componentType");
		if( $cType == 6 ){ // surface
			continue; 
		}
		int $elements = getAttr( $component + ".elements");
		int $newIndices[] = {};
		string $mesh = "";
		for( $item in $sel ){
			$info[0] = $item;
			$info[1] = 0;
			int $compType = itemComponentInfo( $item, $info );
			if( $compType == 7 ){
				$compType = 2; // particles are defined as vertices on component
			}
			if( $compType != $cType ){
				continue;
			}
			if( $mesh != "" ){
				if( $mesh != $info[0] ){
					continue;
				}
			} else {
				// determine if this mesh is the one for this component
				string $nObj = getNObjectFromMesh( $info[0] );
				if( $nObj == "" ){
					continue;
				}
				string $cons[] = `listConnections -type nComponent ($nObj + ".nucleusId")`; 	
				if( stringArrayContains( $component, $cons ) ){
					$mesh = $info[0];
				} else {
					continue;
				}
			}	
			// this is a valid element for the current component
			int $index = $info[1];
			$newIndices[size($newIndices)] = $index;
		}
		int $numIndices = size( $newIndices );
		if( $numIndices < 1 ){
			continue;  // nothing to add
		}

		string $set = getSetFromComponent( $component );
		if( $set != "" ){
			string $nObj = getNObjectFromMesh( $mesh );
			string $baseItem;
			if( $cType == 2 ){ // point
				$baseItem = ".vtx";
			} else if( $cType == 3 ){ // edge
				$baseItem = ".e";
			} else if( $cType == 4 ){ // face
				$baseItem = ".f";
			} else {
				string $fmt = (uiRes("m_dynamicConstraintMembership.kAddUnsupportedCTypeWrn"));
				warning(`format -s $cType $fmt`);
				continue;
			}
			string $inputMesh = findTypeInHistory( $nObj, "mesh", 0, 1 );
			if( $inputMesh == "" ){
				warning( (uiRes("m_dynamicConstraintMembership.kAddNoInputMeshWrn")));
				continue;
			}
			$baseItem = ($inputMesh + $baseItem);
			for( $i = 0; $i < $numIndices; $i++ ){
				string $setItem = ($baseItem + "[" + $newIndices[$i] + "]");
				sets -addElement $set $setItem;
			}	
		} else {
			string $indAttr = ($component + ".componentIndices");
			float $oldIndices[] = `getAttr $indAttr`;
			int $oldNumIndices = size( $oldIndices );
			int $j;
			int $attrIndex = $oldNumIndices;
			for( $i = 0; $i < $numIndices; $i++ ){
				int $inList = false;
				int $newInd = $newIndices[$i];
				for( $j = 0; $j < $oldNumIndices; $j++ ){
					if( $newInd == (int)$oldIndices[$j] ){
						$inList = true;
						break;
					}
				} 
				if( !$inList ){
					setAttr ($indAttr + "[" + $attrIndex + "]") $newInd;
					$attrIndex++;
				}
			}
		}
	}

}

proc removeMembers(string $constraint, string $sel[])
{
	string $components[] = getNConstraintComponents( $constraint );
	string $component, $item, $info[2];
	int $i;
	for( $component in $components ){
		int $cType = getAttr( $component + ".componentType");
		if( $cType == 6 ){ // surface
			continue; // for now... implement surface remove later
		}
		int $elements = getAttr( $component + ".elements");
		int $newIndices[] = {};
		string $mesh = "";
		for( $item in $sel ){
			$info[0] = $item;
			$info[1] = 0;
			int $compType = itemComponentInfo( $item, $info );
			if( $compType == 7 ){
				$compType = 2; // particles are defined as vertices on component
			}
			if( $compType != $cType ){
				continue;
			}
			if( $mesh != "" ){
				if( $mesh != $info[0] ){
					continue;
				}
			} else {
				// determine if this mesh is the one for this component
				string $nObj = getNObjectFromMesh( $info[0] );
				if( $nObj == "" ){
					continue;
				}
				string $cons[] = `listConnections -type nComponent ($nObj + ".nucleusId")`; 	
				if( stringArrayContains( $component, $cons ) ){
					$mesh = $info[0];
				} else {
					continue;
				}
			}	
			int $index = $info[1];
/*
			if( !inComponentList( $component, $index ) ){
				continue;
			}
*/
			// this is a valid element for the current component
			$newIndices[size($newIndices)] = $index;
		}
		int $numIndices = size( $newIndices );
		if( $numIndices < 1 ){
			continue; // nothing to remove
		}
		string $set = getSetFromComponent( $component );

		if( $set != "" ){
			string $nObj = getNObjectFromMesh( $mesh );
			string $baseItem;
			if( $cType == 2 ){ // point
				$baseItem = ".vtx";
			} else if( $cType == 3 ){ // edge
				$baseItem = ".e";
			} else if( $cType == 4 ){ // face
				$baseItem = ".f";
			} else {
				string $fmt = (uiRes("m_dynamicConstraintMembership.kRemoveUnsupportedCTypeWrn"));
				warning(`format -s $cType $fmt`);
				continue;
			}
			string $inputMesh = findTypeInHistory( $nObj, "mesh", 0, 1 );
			if( $inputMesh == "" ){
				warning( (uiRes("m_dynamicConstraintMembership.kRemoveNoInputMeshWrn")));
				continue;
			}
			$baseItem = ($inputMesh + $baseItem);
			string $setObjs[] = `sets -q -no $set`;
			$setObjs = `ls -type mesh $setObjs`;
			int $numSetObj = size($setObjs);
			int $iObj[];

			// we need to temporarily turn intermediateObject off so that "ls -flatten"
			// works properly on the set elements
			for( $i = 0; $i < $numSetObj; $i++ ){
				string $setI = $setObjs[$i] + ".intermediateObject";
				$iObj[$i] = getAttr( $setI );
				if( $iObj[$i] ){
					setAttr $setI false;
				}
			}
			
			for( $i = 0; $i < $numIndices; $i++ ){
				string $setEl[] = `sets -q $set`;
				string $setElements[] = `ls -fl $setEl`;
				int $numEl = size( $setElements );
				if( $numEl < 2 ){
					warning( (uiRes("m_dynamicConstraintMembership.kRemoveSetWrn")));
					break;
				}
				string $setItem = ($baseItem + "[" + $newIndices[$i] + "]");
				sets -remove $set $setItem;
			}	
			for( $i = 0; $i < $numSetObj; $i++ ){
				string $setI = $setObjs[$i] + ".intermediateObject";
				if( $iObj[$i] ){
					setAttr $setI true;
				}
			}
		} else {
			string $indAttr = ($component + ".componentIndices");
			float $oldIndices[];
			if( $elements == 0 ){
				$oldIndices = `getAttr $indAttr`;
			} else {
				int $totalElements[];
				$totalElements[0] = 0;
				if( $cType == 2 ){ 		  // point		
					$totalElements = `polyEvaluate -v $mesh`;
				} else if( $cType == 3 ){ // edge		
					$totalElements = `polyEvaluate -e $mesh`;
				} else if( $cType == 4 ){ // face		
					$totalElements = `polyEvaluate -f $mesh`;
				}
				for( $i = 0; $i < $totalElements[0]; $i++ ){
					$oldIndices[$i] = (float) $i;
				}
				setAttr ( $component + ".elements") 0;
			}
			int $oldNumIndices = size( $oldIndices );
			int $newList[];	
			int $newListSize = 0;
			int $j;
			for( $i = 0; $i < $oldNumIndices; $i++ ){
				int $oldInd = (int)$oldIndices[$i];
				int $deleteItem = false;
				for( $j = 0; $j < $numIndices; $j++ ){
					if( $newIndices[$j] == $oldInd ){
						$deleteItem = true;
						break;
					}
				}	
				if( !$deleteItem ){
					$newList[$newListSize] = $oldInd;	
					$newListSize++;
				}
			}

			for( $i = 0; $i < $newListSize; $i++ ){
				setAttr ($indAttr + "[" + $i + "]") $newList[$i];
			}
			while( $i < $oldNumIndices ){
				removeMultiInstance -b true ($indAttr + "[" + $i + "]");
				// unlink extra attrs at end of multi 
				$i++;
			}
		}
	}

}

global proc dynamicConstraintMembership( string $cmd )
{
	if( $cmd == "select"  ){
		selectMembers();
	} else {
		string $constraint = getNConstraintToEdit( $cmd, 1 );
		if ( $constraint != "" ) {
			string $selected[] = `ls -sl -flatten`;
			$selected = stringArrayRemoveDuplicates( $selected );		
			
			if ($cmd == "replace" ){
				replaceMembers($constraint, $selected );
			} else if ($cmd == "add"     ){
				addMembers($constraint, $selected);
			} else if ($cmd == "remove"  ){
				removeMembers($constraint, $selected);
			}
			string $cons[] = `listConnections -type "nucleus" ($constraint + ".evalCurrent[0]")`;
			if( size($cons) > 0 ){
				getAttr ($cons[0] + ".forceDynamics"); // force update of connections
				refresh -f;
			}
		}
	}
}