#include "LPAnimTwistCommandlet.h"

#include "Animation/AnimData/IAnimationDataController.h"
#include "Animation/AnimData/IAnimationDataModel.h"
#include "Animation/AnimSequence.h"
#include "Animation/Skeleton.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "Dom/JsonObject.h"
#include "HAL/FileManager.h"
#include "LPAnimMethodSelect.h"
#include "LPPoseMeasure.h"
#include "LPSkeletonProfile.h"
#include "Misc/FileHelper.h"
#include "Misc/PackageName.h"
#include "Misc/Paths.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonWriter.h"
#include "UObject/Package.h"
#include "UObject/SavePackage.h"

DEFINE_LOG_CATEGORY_STATIC(LogLPAnimTwist, Log, All);

namespace
{
	// Tolerances du controle de conservation : un os hors des paires ne doit pas bouger.
	constexpr float KeepDegrees = 0.01f;
	constexpr float KeepCentimeters = 0.01f;

	struct FTwistPair
	{
		int32 Bend = INDEX_NONE;
		int32 Twist = INDEX_NONE;
		int32 AxisChild = INDEX_NONE;
		// Axe du segment dans le repere du parent en pose T ; rotation locale de l'os Bend en pose T ;
		// rotation locale de l'os Twist (par rapport a l'os Bend) en pose T.
		FVector AxisParent = FVector::ZeroVector;
		FQuat BendLocalT = FQuat::Identity;
		FQuat TwistLocalT = FQuat::Identity;
	};

	struct FClipResult
	{
		FString Clip;
		int32 Keys = 0;
		float TwistMoved = 0.f;
		float KeepAngle = 0.f;
		float KeepDistance = 0.f;
		bool bSaved = false;
		FString Error;
	};

	FString ClipPath(const FString& Dir, const FString& Name)
	{
		return FString::Printf(TEXT("%s/%s.%s"), *Dir, *Name, *Name);
	}

	// Gabarits : "<Corps>:<chemin d'asset>,..." (-Templates). Un corps absent prend <BodyClipDir>/<Prefixe><RefClip>.
	TMap<FString, FString> ParseTemplates(const FString& Text)
	{
		TMap<FString, FString> Templates;
		TArray<FString> Items;
		Text.ParseIntoArray(Items, TEXT(","), true);
		for (const FString& Item : Items)
		{
			FString Name;
			FString Path;
			if (Item.Split(TEXT(":"), &Name, &Path))
			{
				Templates.Add(Name.TrimStartAndEnd(), Path.TrimStartAndEnd());
			}
		}
		return Templates;
	}

	TArray<FTwistPair> PairsFor(const FLPPose& BodyT)
	{
		TArray<FTwistPair> Pairs;
		const FLPSkeletonProfile* Profile = LPSkeletonProfile::Find(BodyT.Profile);
		if (!Profile)
		{
			return Pairs;
		}
		struct FRoles { const TCHAR* Bend; const TCHAR* Twist; const TCHAR* Axis; };
		const FRoles Roles[] = {
			{ TEXT("upperarm"), TEXT("upperarm_twist"), TEXT("lowerarm") },
			{ TEXT("lowerarm"), TEXT("lowerarm_twist"), TEXT("hand") },
			{ TEXT("thigh"), TEXT("thigh_twist"), TEXT("calf") },
		};
		for (const TCHAR* Side : { TEXT("_l"), TEXT("_r") })
		{
			for (const FRoles& Role : Roles)
			{
				FTwistPair Pair;
				Pair.Bend = BodyT.Names.IndexOfByKey(Profile->Bone(FString(Role.Bend) + Side));
				Pair.Twist = BodyT.Names.IndexOfByKey(Profile->Bone(FString(Role.Twist) + Side));
				Pair.AxisChild = BodyT.Names.IndexOfByKey(Profile->Bone(FString(Role.Axis) + Side));
				if (Pair.Bend == INDEX_NONE || Pair.Twist == INDEX_NONE || Pair.AxisChild == INDEX_NONE)
				{
					continue;
				}
				const int32 Parent = BodyT.Parents[Pair.Bend];
				if (Parent == INDEX_NONE)
				{
					continue;
				}
				const FQuat ParentT = BodyT.Component[Parent].GetRotation();
				const FVector AxisT = (BodyT.Component[Pair.AxisChild].GetLocation() - BodyT.Component[Pair.Bend].GetLocation()).GetSafeNormal();
				Pair.AxisParent = ParentT.UnrotateVector(AxisT);
				Pair.BendLocalT = ParentT.Inverse() * BodyT.Component[Pair.Bend].GetRotation();
				Pair.TwistLocalT = BodyT.Component[Pair.Bend].GetRotation().Inverse() * BodyT.Component[Pair.Twist].GetRotation();
				if (!AxisT.IsNearlyZero())
				{
					Pairs.Add(Pair);
				}
			}
		}
		return Pairs;
	}

	// Traite un clip : cles locales -> composant, swing / twist, nouvelles cles locales des os modifies.
	FClipResult ProcessClip(UAnimSequence* Sequence, const FLPPose& BodyT, const TArray<FTwistPair>& Pairs, bool bDryRun)
	{
		FClipResult Result;
		Result.Clip = Sequence->GetName();
		const USkeleton* Skeleton = Sequence->GetSkeleton();
		const IAnimationDataModel* Model = Sequence->GetDataModel();
		if (!Skeleton || !Model)
		{
			Result.Error = TEXT("squelette ou donnees absents");
			return Result;
		}
		const FReferenceSkeleton& Ref = Skeleton->GetReferenceSkeleton();
		const int32 Num = Ref.GetNum();
		if (Num != BodyT.Names.Num())
		{
			Result.Error = TEXT("squelette different du corps de reference");
			return Result;
		}
		const int32 Keys = Model->GetNumberOfKeys();
		Result.Keys = Keys;
		TArray<TArray<FTransform>> Local;
		Local.SetNum(Num);
		for (int32 Bone = 0; Bone < Num; ++Bone)
		{
			const FName Name = Ref.GetBoneName(Bone);
			if (Model->IsValidBoneTrackName(Name))
			{
				Model->GetBoneTrackTransforms(Name, Local[Bone]);
			}
			if (Local[Bone].Num() != Keys)
			{
				Local[Bone].Init(Ref.GetRefBonePose()[Bone], Keys);
			}
		}
		TSet<int32> Changed;
		TArray<TArray<FTransform>> NewLocal = Local;
		for (int32 Key = 0; Key < Keys; ++Key)
		{
			TArray<FTransform> Comp;
			Comp.SetNum(Num);
			for (int32 Bone = 0; Bone < Num; ++Bone)
			{
				const int32 Parent = Ref.GetParentIndex(Bone);
				Comp[Bone] = Parent == INDEX_NONE ? Local[Bone][Key] : Local[Bone][Key] * Comp[Parent];
			}
			for (const FTwistPair& Pair : Pairs)
			{
				// Rotation complete du membre : portee par l'os Twist (sa rotation locale de pose T retiree), que le
				// clip sorte du retargeter ou ait deja ete traite. Decomposition dans le repere du parent (AP4).
				const int32 Parent = Ref.GetParentIndex(Pair.Bend);
				const FQuat ParentNow = Comp[Parent].GetRotation();
				const FQuat Full = Comp[Pair.Twist].GetRotation() * Pair.TwistLocalT.Inverse();
				const FQuat Delta = ((ParentNow.Inverse() * Full) * Pair.BendLocalT.Inverse()).GetNormalized();
				FQuat Swing;
				FQuat Twist;
				Delta.ToSwingTwist(Pair.AxisParent, Swing, Twist);
				const FQuat Before = Comp[Pair.Bend].GetRotation();
				FTransform BendNew = Comp[Pair.Bend];
				BendNew.SetRotation((ParentNow * Swing * Pair.BendLocalT).GetNormalized());
				Result.TwistMoved = FMath::Max(Result.TwistMoved, FMath::RadiansToDegrees(Before.AngularDistance(BendNew.GetRotation())));
				FTransform& BendLocal = NewLocal[Pair.Bend][Key];
				BendLocal.SetRotation((ParentNow.Inverse() * BendNew.GetRotation()).GetNormalized());
				Changed.Add(Pair.Bend);
				// Enfants directs de l'os Bend (dont l'os Twist) : transformation composant conservee.
				for (int32 Child = 0; Child < Num; ++Child)
				{
					if (Ref.GetParentIndex(Child) != Pair.Bend)
					{
						continue;
					}
					FTransform& ChildLocal = NewLocal[Child][Key];
					ChildLocal.SetRotation((BendNew.GetRotation().Inverse() * Comp[Child].GetRotation()).GetNormalized());
					ChildLocal.SetTranslation(BendNew.GetRotation().UnrotateVector(Comp[Child].GetLocation() - BendNew.GetLocation()));
					Changed.Add(Child);
				}
			}
			// Controle : toute la chaine recomposee avec les nouvelles cles, hors os Bend, doit etre inchangee.
			TArray<FTransform> Check;
			Check.SetNum(Num);
			for (int32 Bone = 0; Bone < Num; ++Bone)
			{
				const int32 Parent = Ref.GetParentIndex(Bone);
				Check[Bone] = Parent == INDEX_NONE ? NewLocal[Bone][Key] : NewLocal[Bone][Key] * Check[Parent];
				bool bBend = false;
				for (const FTwistPair& Pair : Pairs)
				{
					bBend |= Pair.Bend == Bone;
				}
				if (!bBend)
				{
					Result.KeepAngle = FMath::Max(Result.KeepAngle, FMath::RadiansToDegrees(Check[Bone].GetRotation().AngularDistance(Comp[Bone].GetRotation())));
					Result.KeepDistance = FMath::Max(Result.KeepDistance, static_cast<float>(FVector::Dist(Check[Bone].GetLocation(), Comp[Bone].GetLocation())));
				}
			}
		}
		if (Result.KeepAngle > KeepDegrees || Result.KeepDistance > KeepCentimeters)
		{
			Result.Error = FString::Printf(TEXT("conservation hors tolerance : %.4f deg, %.4f cm"), Result.KeepAngle, Result.KeepDistance);
			return Result;
		}
		if (bDryRun)
		{
			return Result;
		}
		IAnimationDataController& Controller = Sequence->GetController();
		Controller.OpenBracket(FText::FromString(TEXT("Repartition de la torsion Bend / Twist (AP3)")), false);
		for (const int32 Bone : Changed)
		{
			const FName Name = Ref.GetBoneName(Bone);
			if (!Model->IsValidBoneTrackName(Name))
			{
				Controller.AddBoneCurve(Name, false);
			}
			TArray<FVector> Positions;
			TArray<FQuat> Rotations;
			TArray<FVector> Scales;
			for (const FTransform& Xform : NewLocal[Bone])
			{
				Positions.Add(Xform.GetTranslation());
				Rotations.Add(Xform.GetRotation());
				Scales.Add(Xform.GetScale3D());
			}
			Controller.SetBoneTrackKeys(Name, Positions, Rotations, Scales, false);
		}
		Controller.CloseBracket(false);
		UPackage* Package = Sequence->GetPackage();
		Package->MarkPackageDirty();
		FSavePackageArgs Args;
		Args.TopLevelFlags = RF_Public | RF_Standalone;
		const FString File = FPackageName::LongPackageNameToFilename(Package->GetName(), FPackageName::GetAssetPackageExtension());
		Result.bSaved = UPackage::SavePackage(Package, Sequence, *File, Args);
		if (!Result.bSaved)
		{
			Result.Error = TEXT("sauvegarde impossible");
		}
		return Result;
	}
}

ULPAnimTwistCommandlet::ULPAnimTwistCommandlet()
{
	IsClient = false;
	IsEditor = true;
	IsServer = false;
	LogToConsole = true;
}

int32 ULPAnimTwistCommandlet::Main(const FString& Params)
{
	FString ClipDir = TEXT("/Game/Quaternius/UAL");
	FString RefClip = TEXT("A_TPose");
	FString BodiesText = TEXT("Bridget:B_,Michael:M_");
	FString ClipsText;
	FString Out = FPaths::Combine(FPaths::ProjectSavedDir(), TEXT("AnimPipeline"), TEXT("quaternius"));
	FParse::Value(*Params, TEXT("ClipDir="), ClipDir);
	FParse::Value(*Params, TEXT("RefClip="), RefClip);
	FParse::Value(*Params, TEXT("Bodies="), BodiesText, false);
	FParse::Value(*Params, TEXT("Clips="), ClipsText, false);
	FParse::Value(*Params, TEXT("Out="), Out);
	// AP7 : clips convertis (prefixes) dans un autre dossier que les clips source ; gabarits explicites par corps.
	FString BodyClipDir;
	FString TemplatesText;
	FParse::Value(*Params, TEXT("BodyClipDir="), BodyClipDir);
	if (BodyClipDir.IsEmpty())
	{
		BodyClipDir = ClipDir;
	}
	FParse::Value(*Params, TEXT("Templates="), TemplatesText, false);
	const TMap<FString, FString> Templates = ParseTemplates(TemplatesText);
	const auto TemplatePath = [&](const FString& BodyName, const FString& Prefix)
	{
		const FString* Path = Templates.Find(BodyName);
		return Path ? *Path : ClipPath(BodyClipDir, Prefix + RefClip);
	};
	const bool bDryRun = FParse::Param(*Params, TEXT("DryRun"));
	IFileManager::Get().MakeDirectory(*Out, true);

	// Methode de transformation (AP7) : le calque de reference (pose A ou T) en depend.
	FString MethodError;
	if (!LPAnimMethodSelect::Apply(Params, ClipPath(ClipDir, RefClip), MethodError))
	{
		UE_LOG(LogLPAnimTwist, Error, TEXT("TORSION %s"), *MethodError);
		return 1;
	}

	FLPPose Reference;
	if (!LPPoseMeasure::Sample(ClipPath(ClipDir, RefClip), 0.f, Reference))
	{
		UE_LOG(LogLPAnimTwist, Error, TEXT("TORSION reference absente %s"), *ClipPath(ClipDir, RefClip));
		return 1;
	}
	IAssetRegistry& Registry = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(TEXT("AssetRegistry")).Get();
	Registry.ScanPathsSynchronous({ BodyClipDir }, true);
	TArray<FAssetData> Assets;
	Registry.GetAssetsByPath(FName(*BodyClipDir), Assets, false);
	TArray<FString> Wanted;
	ClipsText.ParseIntoArray(Wanted, TEXT(","), true);

	TArray<FString> BodyItems;
	BodiesText.ParseIntoArray(BodyItems, TEXT(","), true);
	int32 Failures = 0;
	TSharedRef<FJsonObject> Report = MakeShared<FJsonObject>();
	Report->SetBoolField(TEXT("essai_a_blanc"), bDryRun);
	Report->SetStringField(TEXT("dossier_clips_corps"), BodyClipDir);
	TArray<TSharedPtr<FJsonValue>> BodyReports;
	for (const FString& Item : BodyItems)
	{
		FString BodyName;
		FString Prefix;
		if (!Item.Split(TEXT(":"), &BodyName, &Prefix))
		{
			continue;
		}
		FLPPose Template;
		FLPPose RefT;
		FLPPose BodyT;
		if (!LPPoseMeasure::Sample(TemplatePath(BodyName, Prefix), 0.f, Template)
			|| !LPPoseMeasure::TPoses(Reference, Template, RefT, BodyT))
		{
			UE_LOG(LogLPAnimTwist, Error, TEXT("TORSION %s : pose T impossible"), *BodyName);
			++Failures;
			continue;
		}
		const TArray<FTwistPair> Pairs = PairsFor(BodyT);
		UE_LOG(LogLPAnimTwist, Display, TEXT("TORSION %s profil %s : %d paires Bend / Twist"), *BodyName, *BodyT.Profile, Pairs.Num());
		TArray<TSharedPtr<FJsonValue>> ClipReports;
		float WorstMoved = 0.f;
		for (const FAssetData& Asset : Assets)
		{
			const FString Name = Asset.AssetName.ToString();
			if (Asset.AssetClassPath != UAnimSequence::StaticClass()->GetClassPathName() || !Name.StartsWith(Prefix)
				|| (Wanted.Num() > 0 && !Wanted.Contains(Name.RightChop(Prefix.Len()))))
			{
				continue;
			}
			UAnimSequence* Sequence = LoadObject<UAnimSequence>(nullptr, *ClipPath(BodyClipDir, Name));
			if (!Sequence)
			{
				continue;
			}
			const FClipResult Result = ProcessClip(Sequence, BodyT, Pairs, bDryRun);
			WorstMoved = FMath::Max(WorstMoved, Result.TwistMoved);
			TSharedRef<FJsonObject> Clip = MakeShared<FJsonObject>();
			Clip->SetStringField(TEXT("clip"), Result.Clip);
			Clip->SetNumberField(TEXT("cles"), Result.Keys);
			Clip->SetNumberField(TEXT("torsion_deplacee_deg"), Result.TwistMoved);
			Clip->SetNumberField(TEXT("conservation_deg"), Result.KeepAngle);
			Clip->SetNumberField(TEXT("conservation_cm"), Result.KeepDistance);
			Clip->SetBoolField(TEXT("sauve"), Result.bSaved);
			Clip->SetStringField(TEXT("erreur"), Result.Error);
			ClipReports.Add(MakeShared<FJsonValueObject>(Clip));
			if (!Result.Error.IsEmpty())
			{
				++Failures;
				UE_LOG(LogLPAnimTwist, Warning, TEXT("TORSION echec %s : %s"), *Result.Clip, *Result.Error);
			}
			else
			{
				UE_LOG(LogLPAnimTwist, Display, TEXT("TORSION %s %d cles, torsion deplacee %.1f deg, conservation %.5f deg %.5f cm%s"),
					*Result.Clip, Result.Keys, Result.TwistMoved, Result.KeepAngle, Result.KeepDistance, bDryRun ? TEXT(" (essai)") : TEXT(""));
			}
		}
		TSharedRef<FJsonObject> Body = MakeShared<FJsonObject>();
		Body->SetStringField(TEXT("corps"), BodyName);
		Body->SetNumberField(TEXT("paires"), Pairs.Num());
		Body->SetNumberField(TEXT("torsion_deplacee_max_deg"), WorstMoved);
		Body->SetArrayField(TEXT("clips"), ClipReports);
		BodyReports.Add(MakeShared<FJsonValueObject>(Body));
	}
	Report->SetArrayField(TEXT("corps"), BodyReports);
	Report->SetNumberField(TEXT("echecs"), Failures);
	Report->SetBoolField(TEXT("valide"), Failures == 0);
	FString Text;
	FJsonSerializer::Serialize(Report, TJsonWriterFactory<>::Create(&Text));
	FFileHelper::SaveStringToFile(Text, *FPaths::Combine(Out, TEXT("rapport_torsion.json")), FFileHelper::EEncodingOptions::ForceUTF8WithoutBOM);
	UE_LOG(LogLPAnimTwist, Display, TEXT("TORSION resultat %s : %d echec(s)"), Failures == 0 ? TEXT("VALIDE") : TEXT("REFUSE"), Failures);
	return Failures == 0 ? 0 : 1;
}