UnrealEngineUWP/Engine/Plugins/Animation/ControlRig/Source/ControlRigDeveloper/Private/AnimGraphNode_ControlRig.cpp

// Copyright Epic Games, Inc. All Rights Reserved.

#include "AnimGraphNode_ControlRig.h"
#include "Kismet2/CompilerResultsLog.h"
#include "DetailLayoutBuilder.h"
#include "DetailCategoryBuilder.h"
#include "Widgets/Text/STextBlock.h"
#include "Widgets/Input/SCheckBox.h"
#include "DetailWidgetRow.h"
#include "SVariableMappingWidget.h"
#include "ScopedTransaction.h"
#include "Kismet2/BlueprintEditorUtils.h"
#include "AnimationGraphSchema.h"
#include "ControlRigBlueprintGeneratedClass.h"
#include "ControlRigBlueprint.h"

#define LOCTEXT_NAMESPACE "AnimGraphNode_ControlRig"

UAnimGraphNode_ControlRig::UAnimGraphNode_ControlRig(const FObjectInitializer& ObjectInitializer)
	: Super(ObjectInitializer)
{
}

FText UAnimGraphNode_ControlRig::GetNodeTitle(ENodeTitleType::Type TitleType) const
{
	// display control rig here
	return LOCTEXT("AnimGraphNode_ControlRig_Title", "Control Rig");
}

FText UAnimGraphNode_ControlRig::GetTooltipText() const
{
	// display control rig here
	return LOCTEXT("AnimGraphNode_ControlRig_Tooltip", "Evaluates a control rig");
}

void UAnimGraphNode_ControlRig::CreateCustomPins(TArray<UEdGraphPin*>* OldPins)
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()
	
	// we do this to refresh input variables
	RebuildExposedProperties();
	
	// Grab the SKELETON class here as when we are reconstructed during during BP compilation
	// the full generated class is not yet present built.
	UClass* TargetClass = GetTargetSkeletonClass();

	if (!TargetClass)
	{
		// Nothing to search for properties
		return;
	}

	// Need the schema to extract pin types
	const UEdGraphSchema_K2* Schema = CastChecked<UEdGraphSchema_K2>(GetSchema());

	// Default anim schema for util funcions
	const UAnimationGraphSchema* AnimGraphDefaultSchema = GetDefault<UAnimationGraphSchema>();

#if WITH_EDITOR

	// sustain the current set of custom pins - we'll refrain from changing the node until post load is complete
	if (UControlRigBlueprintGeneratedClass* GeneratedClass = Cast<UControlRigBlueprintGeneratedClass>(GetTargetClass()))
	{
		if (UControlRigBlueprint* RigBlueprint = Cast<UControlRigBlueprint>(GeneratedClass->ClassGeneratedBy))
		{
			if (RigBlueprint->HasAllFlags(RF_NeedPostLoad))
			{
				if(OldPins)
				{
					for (UEdGraphPin* OldPin : *OldPins)
					{
						// do not rebuild sub pins as they will be treated after in UK2Node::RestoreSplitPins
						const bool bIsSubPin = OldPin->ParentPin && OldPin->ParentPin->bHidden;
						if (bIsSubPin)
						{
							continue;
						}
						
						bool bFound = false;
						for (UEdGraphPin* CurrentPin : Pins)
						{
							if (CurrentPin->GetFName() == OldPin->GetFName() &&
								CurrentPin->PinType == OldPin->PinType)
							{
								bFound = true;
								break;
							}
						}

						if (!bFound)
						{
							FName PropertyName = OldPin->GetFName();
							UEdGraphPin* NewPin = CreatePin(EEdGraphPinDirection::EGPD_Input, OldPin->PinType, PropertyName);
							NewPin->PinFriendlyName = OldPin->PinFriendlyName;
							
							// Newly created pin does not have an auto generated default value, so we need to generate one here
							// Missing the auto-gen default would cause UEdGraphPin::MovePersistentDataFromOldPin to override
							// the actual default value with the empty auto-gen default, causing BP compiler to complain
							// This is similar to how the following two functions create and initialize new pins, create first 
							// and then set an auto-gen default
							// FOptionalPinManager::CreateVisiblePins()
							// FAnimBlueprintNodeOptionalPinManager::PostInitNewPin()
							AnimGraphDefaultSchema->SetPinAutogeneratedDefaultValueBasedOnType(NewPin);

							AnimGraphDefaultSchema->ResetPinToAutogeneratedDefaultValue(NewPin, false);

							CustomizePinData(NewPin, PropertyName, INDEX_NONE);
						}
					}
				}
				return;
			}
		}
	}
#endif

	// We'll track the names we encounter by removing from this list, if anything remains the properties
	// have been removed from the target class and we should remove them too
	TSet<FName> BeginExposableNames;
	TSet<FName> CurrentlyExposedNames;
	for(const FOptionalPinFromProperty& OptionalPin : CustomPinProperties)
	{
		BeginExposableNames.Add(OptionalPin.PropertyName);

		if(OptionalPin.bShowPin)
		{
			CurrentlyExposedNames.Add(OptionalPin.PropertyName);
		}
	}

	for (TPair<FName, FRigVMExternalVariable>& VariablePair : InputVariables)
	{
		FName PropertyName = VariablePair.Key;
		BeginExposableNames.Remove(PropertyName);

		if (CurrentlyExposedNames.Contains(PropertyName))
		{
			FRigVMExternalVariable& Variable = VariablePair.Value;

			FEdGraphPinType PinType = RigVMTypeUtils::PinTypeFromExternalVariable(Variable);
			if (!PinType.PinCategory.IsValid())
			{
				continue;
			}

			UEdGraphPin* NewPin = CreatePin(EEdGraphPinDirection::EGPD_Input, PinType, PropertyName);
			NewPin->PinFriendlyName = FText::FromName(PropertyName);

			// Newly created pin does not have an auto generated default value, so we need to generate one here
			// Missing the auto-gen default would cause UEdGraphPin::MovePersistentDataFromOldPin to override
			// the actual default value with the empty auto-gen default, causing BP compiler to complain
			// This is similar to how the following two functions create and initialize new pins, create first 
			// and then set an auto-gen default
			// FOptionalPinManager::CreateVisiblePins()
			// FAnimBlueprintNodeOptionalPinManager::PostInitNewPin()
			AnimGraphDefaultSchema->SetPinAutogeneratedDefaultValueBasedOnType(NewPin);

			// We cant interrogate CDO here as we may be mid-compile, so we can only really
			// reset to the autogenerated default.
			AnimGraphDefaultSchema->ResetPinToAutogeneratedDefaultValue(NewPin, false);

			// Extract default values from the Target Control Rig if possible
			// Memory could be null if Control Rig is compiling, so only do it if Memory is not null
			if (Variable.IsValid(false))
			{ 
				if (UControlRigBlueprintGeneratedClass* GeneratedClass = Cast<UControlRigBlueprintGeneratedClass>(GetTargetClass()))
				{
					for (TFieldIterator<FProperty> PropertyIt(GeneratedClass); PropertyIt; ++PropertyIt)
					{
						if (PropertyIt->GetFName() == PropertyName)
						{
							FString DefaultValue;
							
							// Variable.Memory here points to the corresponding property in the Control Rig BP CDO, it was initialized in UAnimGraphNode_ControlRig::RebuildExposedProperties 
							PropertyIt->ExportTextItem_Direct(DefaultValue, Variable.Memory, nullptr, nullptr, PPF_None, nullptr);
							
							if (!DefaultValue.IsEmpty())
							{
								AnimGraphDefaultSchema->TrySetDefaultValue(*NewPin, DefaultValue);
							}
						}
					}
				}
			} 
			
			// sustain the current set of custom pins - we'll refrain from changing the node until post load is complete 
			CustomizePinData(NewPin, PropertyName, INDEX_NONE);
		}
	}

	if(const UClass* ControlRigClass = GetTargetClass())
	{
		if(UControlRig* CDO = ControlRigClass->GetDefaultObject<UControlRig>())
		{
			if(const URigHierarchy* Hierarchy = CDO->GetHierarchy())
			{
				Hierarchy->ForEach<FRigControlElement>([&](FRigControlElement* ControlElement) -> bool
				{
					if (Hierarchy->IsAnimatable(ControlElement))
					{
						const FName ControlName = ControlElement->GetName();
						BeginExposableNames.Remove(ControlName);

						if (CurrentlyExposedNames.Contains(ControlName))
						{
							const FEdGraphPinType PinType = Hierarchy->GetControlPinType(ControlElement);
							if (!PinType.PinCategory.IsValid())
							{
								return true;
							}

							UEdGraphPin* NewPin = CreatePin(EEdGraphPinDirection::EGPD_Input, PinType, ControlName);
							NewPin->PinFriendlyName = FText::FromName(ControlElement->GetName());

							// Newly created pin does not have an auto generated default value, so we need to generate one here
							// Missing the auto-gen default would cause UEdGraphPin::MovePersistentDataFromOldPin to override
							// the actual default value with the empty auto-gen default, causing BP compiler to complain
							// This is similar to how the following two functions create and initialize new pins, create first 
							// and then set an auto-gen default
							// FOptionalPinManager::CreateVisiblePins()
							// FAnimBlueprintNodeOptionalPinManager::PostInitNewPin()
							AnimGraphDefaultSchema->SetPinAutogeneratedDefaultValueBasedOnType(NewPin);

							// We cant interrogate CDO here as we may be mid-compile, so we can only really
							// reset to the autogenerated default.
							AnimGraphDefaultSchema->ResetPinToAutogeneratedDefaultValue(NewPin, false);

							// Extract default values from the Target Control Rig if possible
							const FString DefaultValue = Hierarchy->GetControlPinDefaultValue(ControlElement, true);
							if(!DefaultValue.IsEmpty())
							{
								AnimGraphDefaultSchema->TrySetDefaultValue(*NewPin, DefaultValue);
							}
							
							// sustain the current set of custom pins - we'll refrain from changing the node until post load is complete 
							CustomizePinData(NewPin, ControlName, INDEX_NONE);
						}
					}
					
					return true;
				});
			}
		}
	}

	// Remove any properties that no longer exist on the target class
	for (FName& RemovedPropertyName : BeginExposableNames)
	{
		CustomPinProperties.RemoveAll([RemovedPropertyName](const FOptionalPinFromProperty& InOptionalPin) { return InOptionalPin.PropertyName == RemovedPropertyName; });
	}
}

void UAnimGraphNode_ControlRig::ValidateAnimNodeDuringCompilation(USkeleton* ForSkeleton, FCompilerResultsLog& MessageLog)
{
	if (UClass* TargetClass = GetTargetClass())
	{
		if (UControlRigBlueprint* Blueprint = Cast<UControlRigBlueprint>(TargetClass->ClassGeneratedBy))
		{
			URigHierarchy* Hierarchy = Blueprint->Hierarchy;
			if (ForSkeleton)
			{
				const FReferenceSkeleton& ReferenceSkeleton = ForSkeleton->GetReferenceSkeleton();
				const TArray<FMeshBoneInfo>& BoneInfos = ReferenceSkeleton.GetRefBoneInfo();

				for (const FMeshBoneInfo& BoneInfo : BoneInfos)
				{
					const FRigElementKey BoneKey(BoneInfo.Name, ERigElementType::Bone);
					if (Hierarchy->Contains(BoneKey))
					{
						const FRigElementKey ParentKey = Hierarchy->GetFirstParent(BoneKey);

						FName DesiredParentName = NAME_None;
						if (BoneInfo.ParentIndex != INDEX_NONE)
						{
							DesiredParentName = BoneInfos[BoneInfo.ParentIndex].Name;
						}

						if (DesiredParentName != ParentKey.Name)
						{
							FString Message = FString::Printf(TEXT("@@ - Hierarchy discrepancy for bone '%s' - different parents on Control Rig vs SkeletalMesh."), *BoneInfo.Name.ToString());
							MessageLog.Warning(*Message, this);
						}
					}
				}
			}
		}
	}

	Super::ValidateAnimNodeDuringCompilation(ForSkeleton, MessageLog);
}

void UAnimGraphNode_ControlRig::RebuildExposedProperties()
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	if (UControlRigBlueprintGeneratedClass* TargetClass = Cast<UControlRigBlueprintGeneratedClass>(GetTargetClass()))
	{
		if (UControlRigBlueprint* RigBlueprint = Cast<UControlRigBlueprint>(TargetClass->ClassGeneratedBy))
		{
			if (RigBlueprint->HasAllFlags(RF_NeedPostLoad))
			{
				return;
			}
		}
	}

	TSet<FName> OldOptionalPinNames;
	TSet<FName> OldExposedPinNames;
	for(const FOptionalPinFromProperty& OptionalPin : CustomPinProperties)
	{
		OldOptionalPinNames.Add(OptionalPin.PropertyName);

		if(OptionalPin.bShowPin)
		{
			OldExposedPinNames.Add(OptionalPin.PropertyName);
		}
	}
	CustomPinProperties.Empty();

	// go through exposed properties, and clean up
	GetVariables(true, InputVariables);
	// we still update OUtputvariables
	// we don't want output to be exposed
	GetVariables(false, OutputVariables);

	// clear IO variables that don't exist anymore
	auto ClearInvalidMapping = [](const TMap<FName, FRigVMExternalVariable>& InVariables, TMap<FName, FName>& InOutMapping)
	{
		TArray<FName> KeyArray;
		InOutMapping.GenerateKeyArray(KeyArray);

		for (int32 Index=0; Index<KeyArray.Num(); ++Index)
		{
			// if this input doesn't exist anymore
			if (!InVariables.Contains(KeyArray[Index]))
			{
				InOutMapping.Remove(KeyArray[Index]);
			}
		}
	};

	ClearInvalidMapping(InputVariables, Node.InputMapping);
	ClearInvalidMapping(OutputVariables, Node.OutputMapping);

	auto MakeOptionalPin = [&OldExposedPinNames](FName InPinName)-> FOptionalPinFromProperty
	{
		FOptionalPinFromProperty OptionalPin;
		OptionalPin.PropertyName = InPinName;
		OptionalPin.bShowPin = OldExposedPinNames.Contains(InPinName);
		OptionalPin.bCanToggleVisibility = true;
		OptionalPin.bIsOverrideEnabled = false;
		return OptionalPin;
	};
	
	for (auto Iter = InputVariables.CreateConstIterator(); Iter; ++Iter)
	{
		CustomPinProperties.Add(MakeOptionalPin(Iter.Key()));
	}

	// also add all of the controls
	if(const UClass* TargetClass = GetTargetClass())
	{
		if(UControlRig* CDO = TargetClass->GetDefaultObject<UControlRig>())
		{
			if(const URigHierarchy* Hierarchy = CDO->GetHierarchy())
			{
				Hierarchy->ForEach<FRigControlElement>([&](const FRigControlElement* ControlElement) -> bool
				{
					if (Hierarchy->IsAnimatable(ControlElement))
					{
						const FName ControlName = ControlElement->GetName();
						CustomPinProperties.Add(MakeOptionalPin(ControlName));
					}
					return true;
				});
			}
		}
	}
}

bool UAnimGraphNode_ControlRig::IsInputProperty(const FName& PropertyName) const
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	// this is true for both input variables and controls
	return InputVariables.Contains(PropertyName) || !OutputVariables.Contains(PropertyName);
}

FRigControlElement* UAnimGraphNode_ControlRig::FindControlElement(const FName& InControlName) const
{
	if(const UClass* ControlRigClass = GetTargetClass())
	{
		if(UControlRig* CDO = ControlRigClass->GetDefaultObject<UControlRig>())
		{
			if(const URigHierarchy* Hierarchy = CDO->GetHierarchy())
			{
				return (FRigControlElement*)Hierarchy->Find<FRigControlElement>(FRigElementKey(InControlName, ERigElementType::Control));
			}
		}
	}
	return nullptr;
}

bool UAnimGraphNode_ControlRig::IsAvailableToMapToCurve(const FName& PropertyName, bool bInput) const
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	// find if input or output
	// ensure it could convert to float
	const FRigVMExternalVariable* Variable = (bInput) ? InputVariables.Find(PropertyName) : OutputVariables.Find(PropertyName);
	if (Variable)
	{
		return Variable->TypeName == TEXT("float");
	}

	if(const FRigControlElement* ControlElement = FindControlElement(PropertyName))
	{
		return ControlElement->Settings.ControlType == ERigControlType::Float;
	}

	return ensure(false);
}

bool UAnimGraphNode_ControlRig::IsPropertyExposeEnabled(FName PropertyName) const
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	// if known exposable, and and if it hasn't been exposed yet
	if (CustomPinProperties.ContainsByPredicate([PropertyName](const FOptionalPinFromProperty& InOptionalPin){ return InOptionalPin.PropertyName == PropertyName; }))
	{
		return IsInputProperty(PropertyName);
	}

	return false;
}

ECheckBoxState UAnimGraphNode_ControlRig::IsPropertyExposed(FName PropertyName) const
{
	if (CustomPinProperties.ContainsByPredicate([PropertyName](const FOptionalPinFromProperty& InOptionalPin){ return InOptionalPin.bShowPin && InOptionalPin.PropertyName == PropertyName; }))
	{
		return ECheckBoxState::Checked;
	}

	return ECheckBoxState::Unchecked;
}

void UAnimGraphNode_ControlRig::OnPropertyExposeCheckboxChanged(ECheckBoxState NewState, FName PropertyName)
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	FOptionalPinFromProperty* FoundPin = CustomPinProperties.FindByPredicate([PropertyName](const FOptionalPinFromProperty& InOptionalPin)
	{
		return InOptionalPin.PropertyName == PropertyName;
	});
	
	if (FoundPin)
	{
		FoundPin->bShowPin = !FoundPin->bShowPin;

		FGuardValue_Bitfield(bDisableOrphanPinSaving, true);
		ReconstructNode();
		
		// see if any of my child has the mapping, and clear them
		if (NewState == ECheckBoxState::Checked)
		{
			FScopedTransaction Transaction(LOCTEXT("PropertyExposedChanged", "Expose Property to Pin"));
			Modify();

			FBlueprintEditorUtils::MarkBlueprintAsStructurallyModified(GetBlueprint());

			bool bInput = IsInputProperty(PropertyName);
			// if checked, we clear mapping
			// and unclear all children
			Node.SetIOMapping(bInput, PropertyName, NAME_None); 
		}
	}
}

void UAnimGraphNode_ControlRig::CustomizeDetails(IDetailLayoutBuilder& DetailBuilder)
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	Super::CustomizeDetails(DetailBuilder);

	// We dont allow multi-select here
	if (DetailBuilder.GetSelectedObjects().Num() > 1)
	{
		return;
	}

	// input/output exposure feature START
	RebuildExposedProperties();

	IDetailCategoryBuilder& InputCategoryBuilder = DetailBuilder.EditCategory(FName(TEXT("Input")));

	FDetailWidgetRow& InputWidgetRow = InputCategoryBuilder.AddCustomRow(FText::FromString(TEXT("Input")));
	InputWidgetRow.WholeRowContent()
	[
		SNew(SVariableMappingWidget)
		.OnVariableMappingChanged(FOnVariableMappingChanged::CreateUObject(this, &UAnimGraphNode_ControlRig::OnVariableMappingChanged, true))
		.OnGetVariableMapping(FOnGetVariableMapping::CreateUObject(this, &UAnimGraphNode_ControlRig::GetVariableMapping, true))
		.OnGetAvailableMapping(FOnGetAvailableMapping::CreateUObject(this, &UAnimGraphNode_ControlRig::GetAvailableMapping, true))
		.OnCreateVariableMapping(FOnCreateVariableMapping::CreateUObject(this, &UAnimGraphNode_ControlRig::CreateVariableMapping, true))
		.OnVariableOptionAvailable(FOnVarOptionAvailable::CreateUObject(this, &UAnimGraphNode_ControlRig::IsAvailableToMapToCurve, true))
		.OnPinGetCheckState(FOnPinGetCheckState::CreateUObject(this, &UAnimGraphNode_ControlRig::IsPropertyExposed))
		.OnPinCheckStateChanged(FOnPinCheckStateChanged::CreateUObject(this, &UAnimGraphNode_ControlRig::OnPropertyExposeCheckboxChanged))
		.OnPinIsEnabledCheckState(FOnPinIsCheckEnabled::CreateUObject(this, &UAnimGraphNode_ControlRig::IsPropertyExposeEnabled))
	];

	IDetailCategoryBuilder& OutputCategoryBuilder = DetailBuilder.EditCategory(FName(TEXT("Output")));
	FDetailWidgetRow& OutputWidgetRow = OutputCategoryBuilder.AddCustomRow(FText::FromString(TEXT("Output")));
	OutputWidgetRow.WholeRowContent()
	[
		SNew(SVariableMappingWidget)
		.OnVariableMappingChanged(FOnVariableMappingChanged::CreateUObject(this, &UAnimGraphNode_ControlRig::OnVariableMappingChanged, false))
		.OnGetVariableMapping(FOnGetVariableMapping::CreateUObject(this, &UAnimGraphNode_ControlRig::GetVariableMapping, false))
		.OnGetAvailableMapping(FOnGetAvailableMapping::CreateUObject(this, &UAnimGraphNode_ControlRig::GetAvailableMapping, false))
		.OnCreateVariableMapping(FOnCreateVariableMapping::CreateUObject(this, &UAnimGraphNode_ControlRig::CreateVariableMapping, false))
		.OnVariableOptionAvailable(FOnVarOptionAvailable::CreateUObject(this, &UAnimGraphNode_ControlRig::IsAvailableToMapToCurve, false))
		.OnPinGetCheckState(FOnPinGetCheckState::CreateUObject(this, &UAnimGraphNode_ControlRig::IsPropertyExposed))
		.OnPinCheckStateChanged(FOnPinCheckStateChanged::CreateUObject(this, &UAnimGraphNode_ControlRig::OnPropertyExposeCheckboxChanged))
		.OnPinIsEnabledCheckState(FOnPinIsCheckEnabled::CreateUObject(this, &UAnimGraphNode_ControlRig::IsPropertyExposeEnabled))
	];

	TSharedRef<IPropertyHandle> ClassHandle = DetailBuilder.GetProperty(TEXT("Node.ControlRigClass"), GetClass());
	if (ClassHandle->IsValidHandle())
	{
		ClassHandle->SetOnPropertyValueChanged(FSimpleDelegate::CreateUObject(this, &UAnimGraphNode_ControlRig::OnInstanceClassChanged, &DetailBuilder));
	}

	// input/output exposure feature END

	// alpha property blending support START
	TSharedRef<IPropertyHandle> NodeHandle = DetailBuilder.GetProperty(FName(TEXT("Node")), GetClass());

	if (Node.AlphaInputType != EAnimAlphaInputType::Bool)
	{
		DetailBuilder.HideProperty(NodeHandle->GetChildHandle(GET_MEMBER_NAME_CHECKED(FAnimNode_ControlRig, bAlphaBoolEnabled)));
		DetailBuilder.HideProperty(NodeHandle->GetChildHandle(GET_MEMBER_NAME_CHECKED(FAnimNode_ControlRig, AlphaBoolBlend)));
	}

	if (Node.AlphaInputType != EAnimAlphaInputType::Float)
	{
		DetailBuilder.HideProperty(NodeHandle->GetChildHandle(GET_MEMBER_NAME_CHECKED(FAnimNode_ControlRig, Alpha)));
		DetailBuilder.HideProperty(NodeHandle->GetChildHandle(GET_MEMBER_NAME_CHECKED(FAnimNode_ControlRig, AlphaScaleBias)));
	}

	if (Node.AlphaInputType != EAnimAlphaInputType::Curve)
	{
		DetailBuilder.HideProperty(NodeHandle->GetChildHandle(GET_MEMBER_NAME_CHECKED(FAnimNode_ControlRig, AlphaCurveName)));
	}

	if ((Node.AlphaInputType != EAnimAlphaInputType::Float)
		&& (Node.AlphaInputType != EAnimAlphaInputType::Curve))
	{
		DetailBuilder.HideProperty(NodeHandle->GetChildHandle(GET_MEMBER_NAME_CHECKED(FAnimNode_ControlRig, AlphaScaleBiasClamp)));
	}
	// alpha property blending support END
}

void UAnimGraphNode_ControlRig::GetVariables(bool bInput, TMap<FName, FRigVMExternalVariable>& OutVariables) const
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	OutVariables.Reset();

	if (UControlRigBlueprintGeneratedClass* TargetClass = Cast<UControlRigBlueprintGeneratedClass>(GetTargetClass()))
	{
		if (UControlRigBlueprint* RigBlueprint = Cast<UControlRigBlueprint>(TargetClass->ClassGeneratedBy))
		{
			//RigBlueprint->CleanupVariables();
			UControlRig* ControlRig = TargetClass->GetDefaultObject<UControlRig>();
			if (ControlRig)
			{
				const TArray<FRigVMExternalVariable>& PublicVariables = ControlRig->GetPublicVariables();
				for (const FRigVMExternalVariable& PublicVariable : PublicVariables)
				{
					if (!bInput || !PublicVariable.bIsReadOnly)
					{
						OutVariables.Add(PublicVariable.Name, PublicVariable);
					}
				}
			}
		}
	}
}

void UAnimGraphNode_ControlRig::OnVariableMappingChanged(const FName& PathName, const FName& Curve, bool bInput)
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	FScopedTransaction Transaction(LOCTEXT("VariableMappingChanged", "Change Variable Mapping"));
	Modify();

	FBlueprintEditorUtils::MarkBlueprintAsStructurallyModified(GetBlueprint());

	// @todo: this is not enough when we start breaking down struct
	Node.SetIOMapping(bInput, PathName, Curve);
}

FName UAnimGraphNode_ControlRig::GetVariableMapping(const FName& PathName, bool bInput)
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	// @todo: this is not enough when we start breaking down struct
	return Node.GetIOMapping(bInput, PathName);
}

void UAnimGraphNode_ControlRig::GetAvailableMapping(const FName& PathName, TArray<FName>& OutArray, bool bInput)
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	UAnimBlueprint* AnimBP = CastChecked<UAnimBlueprint>(GetBlueprint());
	USkeleton* TargetSkeleton = AnimBP->TargetSkeleton;
	OutArray.Reset();
	if (TargetSkeleton)
	{
		const FSmartNameMapping* CurveMapping = TargetSkeleton->GetSmartNameContainer(USkeleton::AnimCurveMappingName);
		CurveMapping->FillNameArray(OutArray);
		
		// also add all controls
		if(const UClass* ControlRigClass = GetTargetClass())
		{
			if(UControlRig* CDO = ControlRigClass->GetDefaultObject<UControlRig>())
			{
				if(const URigHierarchy* Hierarchy = CDO->GetHierarchy())
				{
					Hierarchy->ForEach<FRigControlElement>([&](FRigControlElement* ControlElement) -> bool
					{
						if (Hierarchy->IsAnimatable(ControlElement))
						{
							OutArray.Add(ControlElement->GetName());
						}
						return true;
					});
				}
			}
		}

		// we want to exclude anything that has been mapped already
		TArray<FName> InputMapping, OutputMapping;
		Node.InputMapping.GenerateValueArray(InputMapping);
		Node.OutputMapping.GenerateValueArray(OutputMapping);

		// I have to remove Input/Output Curves from OutArray
		for (int32 Index = 0; Index < OutArray.Num(); ++Index)
		{
			const FName& Item = OutArray[Index];

			if (InputMapping.Contains(Item))
			{
				OutArray.RemoveAt(Index);
				--Index;
			}
			else if (OutputMapping.Contains(Item))
			{
				OutArray.RemoveAt(Index);
				--Index;
			}
		}
	}
}

void UAnimGraphNode_ControlRig::CreateVariableMapping(const FString& FilteredText, TArray< TSharedPtr<FVariableMappingInfo> >& OutArray, bool bInput)
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	// should have latest
	OutArray.Reset();

	bool bDoFiltering = !FilteredText.IsEmpty();
	const TMap<FName, FRigVMExternalVariable>& Variables = (bInput)? InputVariables : OutputVariables;
	for (auto Iter = Variables.CreateConstIterator(); Iter; ++Iter)
	{
		const FName& Name = Iter.Key();
		const FString& DisplayName = Name.ToString();

		const FString MappedName = GetVariableMapping(Name, bInput).ToString();
		// make sure it doesn't fit any of them
		if (!bDoFiltering || 
			(DisplayName.Contains(FilteredText) || MappedName.Contains(FilteredText)))
		{
			TSharedRef<FVariableMappingInfo> Item = FVariableMappingInfo::Make(Iter.Key());

			FVariableMappingInfo& ItemRaw = Item.Get();
			FString PathString = ItemRaw.GetPathName().ToString();
			FString DisplayString = ItemRaw.GetDisplayName();

			OutArray.Add(Item);
		}
	}

	if(bInput)
	{
		// add all controls
		if(const UClass* ControlRigClass = GetTargetClass())
		{
			if(UControlRig* CDO = ControlRigClass->GetDefaultObject<UControlRig>())
			{
				if(const URigHierarchy* Hierarchy = CDO->GetHierarchy())
				{
					Hierarchy->ForEach<FRigControlElement>([&](FRigControlElement* ControlElement) -> bool
					{
						if (Hierarchy->IsAnimatable(ControlElement))
						{
							const FName ControlName = ControlElement->GetName();
							const FString& DisplayName = ControlName.ToString();

							if (!bDoFiltering || DisplayName.Contains(FilteredText))
							{
								TSharedRef<FVariableMappingInfo> Item = FVariableMappingInfo::Make(ControlName);
								OutArray.Add(Item);
							}
						}
						return true;
					});
				}
			}
		}
	}
}

void UAnimGraphNode_ControlRig::PostEditChangeProperty(FPropertyChangedEvent& PropertyChangedEvent)
{
	DECLARE_SCOPE_HIERARCHICAL_COUNTER_FUNC()

	Super::PostEditChangeProperty(PropertyChangedEvent);

	bool bRequiresNodeReconstruct = false;
	FProperty* ChangedProperty = PropertyChangedEvent.Property;

	if (ChangedProperty)
	{
		if (ChangedProperty->GetFName() == GET_MEMBER_NAME_CHECKED(FAnimNode_ControlRig, ControlRigClass))
		{
			bRequiresNodeReconstruct = true;
			RebuildExposedProperties();
		}

		if (ChangedProperty->GetFName() == GET_MEMBER_NAME_STRING_CHECKED(FAnimNode_ControlRig, AlphaInputType))
		{
			FScopedTransaction Transaction(LOCTEXT("ChangeAlphaInputType", "Change Alpha Input Type"));
			Modify();

			// Break links to pins going away
			for (int32 PinIndex = 0; PinIndex < Pins.Num(); ++PinIndex)
			{
				UEdGraphPin* Pin = Pins[PinIndex];
				if (Pin->PinName == GET_MEMBER_NAME_STRING_CHECKED(FAnimNode_ControlRig, Alpha))
				{
					if (Node.AlphaInputType != EAnimAlphaInputType::Float)
					{
						Pin->BreakAllPinLinks();
						PropertyBindings.Remove(Pin->PinName);
					}
				}
				else if (Pin->PinName == GET_MEMBER_NAME_STRING_CHECKED(FAnimNode_ControlRig, bAlphaBoolEnabled))
				{
					if (Node.AlphaInputType != EAnimAlphaInputType::Bool)
					{
						Pin->BreakAllPinLinks();
						PropertyBindings.Remove(Pin->PinName);
					}
				}
				else if (Pin->PinName == GET_MEMBER_NAME_STRING_CHECKED(FAnimNode_ControlRig, AlphaCurveName))
				{
					if (Node.AlphaInputType != EAnimAlphaInputType::Curve)
					{
						Pin->BreakAllPinLinks();
						PropertyBindings.Remove(Pin->PinName);
					}
				}
			}

			bRequiresNodeReconstruct = true;
		}
	}

	if (bRequiresNodeReconstruct)
	{
		ReconstructNode();
		FBlueprintEditorUtils::MarkBlueprintAsStructurallyModified(GetBlueprint());
	}
}

void UAnimGraphNode_ControlRig::CustomizePinData(UEdGraphPin* Pin, FName SourcePropertyName, int32 ArrayIndex) const
{
	Super::CustomizePinData(Pin, SourcePropertyName, ArrayIndex);

	if (Pin->PinName == GET_MEMBER_NAME_STRING_CHECKED(FAnimNode_ControlRig, Alpha))
	{
		Pin->bHidden = (Node.AlphaInputType != EAnimAlphaInputType::Float);

		if (!Pin->bHidden)
		{
			Pin->PinFriendlyName = Node.AlphaScaleBias.GetFriendlyName(Node.AlphaScaleBiasClamp.GetFriendlyName(Pin->PinFriendlyName));
		}
	}

	if (Pin->PinName == GET_MEMBER_NAME_STRING_CHECKED(FAnimNode_ControlRig, bAlphaBoolEnabled))
	{
		Pin->bHidden = (Node.AlphaInputType != EAnimAlphaInputType::Bool);
	}

	if (Pin->PinName == GET_MEMBER_NAME_STRING_CHECKED(FAnimNode_ControlRig, AlphaCurveName))
	{
		Pin->bHidden = (Node.AlphaInputType != EAnimAlphaInputType::Curve);

		if (!Pin->bHidden)
		{
			Pin->PinFriendlyName = Node.AlphaScaleBiasClamp.GetFriendlyName(Pin->PinFriendlyName);
		}
	}
}
#undef LOCTEXT_NAMESPACE