#include "LPAnimChecks.h"

#include "Animation/AnimData/IAnimationDataModel.h"
#include "Animation/AnimSequence.h"
#include "Engine/SkeletalMesh.h"
#include "PhysicsEngine/BodySetup.h"
#include "PhysicsEngine/PhysicsAsset.h"
#include "PhysicsEngine/SkeletalBodySetup.h"
#include "Rendering/SkeletalMeshLODModel.h"
#include "Rendering/SkeletalMeshModel.h"

namespace
{
	// Os du tronc dont les capsules forment le volume teste (roles du profil).
	const TCHAR* const TrunkRoles[] = { TEXT("hips"), TEXT("pelvis"), TEXT("spine") };

	// Seuils internes de detection (pas des seuils d'acceptation, fixes dans LPAnimEval).
	constexpr float ContactHeight = 3.f;     // cm : pied source considere au sol
	constexpr float PlantedMove = 1.f;       // cm par image : pied source immobile
	constexpr float StraightLeg = 15.f;      // degres : jambe presque tendue
	constexpr float JumpFactor = 1.5f;       // une rotation d'image peut depasser la source de 50 %
	constexpr float JumpMargin = 2.f;        // ... plus 2 degres

	int32 BoneOf(const FLPPose& Pose, const TCHAR* Role, int32 Index = 0)
	{
		const FName Name = LPPoseMeasure::RoleBone(Pose, Role, Index);
		return Name.IsNone() ? INDEX_NONE : Pose.Names.IndexOfByKey(Name);
	}

	FVector At(const FLPPose& Pose, int32 Bone)
	{
		return Pose.Component.IsValidIndex(Bone) ? Pose.Component[Bone].GetLocation() : FVector::ZeroVector;
	}

	// Avant du corps : perpendiculaire horizontale a l'axe des hanches.
	FVector ForwardOf(const FLPPose& Pose)
	{
		FVector Side = At(Pose, BoneOf(Pose, TEXT("thigh_l"))) - At(Pose, BoneOf(Pose, TEXT("thigh_r")));
		Side.Z = 0.f;
		return FVector::CrossProduct(Side.GetSafeNormal(), FVector::UpVector).GetSafeNormal();
	}

	float Flexion(const FLPPose& Pose, int32 A, int32 B, int32 C)
	{
		const FVector Upper = (At(Pose, B) - At(Pose, A)).GetSafeNormal();
		const FVector Lower = (At(Pose, C) - At(Pose, B)).GetSafeNormal();
		return FMath::RadiansToDegrees(FMath::Acos(FMath::Clamp(FVector::DotProduct(Upper, Lower), -1.f, 1.f)));
	}

	// Genou derriere la ligne hanche-cheville (cm, positif = hyperextension), jambe presque tendue seulement.
	float KneeBehind(const FLPPose& Pose, int32 Hip, int32 Knee, int32 Ankle)
	{
		if (Flexion(Pose, Hip, Knee, Ankle) > StraightLeg)
		{
			return 0.f;
		}
		const FVector Line = (At(Pose, Ankle) - At(Pose, Hip)).GetSafeNormal();
		const FVector Offset = (At(Pose, Knee) - At(Pose, Hip)) - Line * FVector::DotProduct(At(Pose, Knee) - At(Pose, Hip), Line);
		return FMath::Max(0.f, -static_cast<float>(FVector::DotProduct(Offset, ForwardOf(Pose))));
	}

	// Repli sans actif physique : une capsule par segment de la colonne (bassin -> base du cou), rayon pris sur
	// la peau du maillage (sommets dont l'os dominant est du tronc), au 30e centile des distances a l'axe :
	// prudent, les flancs plus larges que la profondeur ne declenchent pas de faux positifs.
	constexpr float SkinPercentile = 0.30f;

	TArray<FLPCheckCapsule> SkinCapsules(const USkeletalMesh* Mesh, const FLPPose& BodyT, const TSet<int32>& Trunk)
	{
		TArray<FLPCheckCapsule> Capsules;
		const FSkeletalMeshModel* Imported = Mesh ? Mesh->GetImportedModel() : nullptr;
		if (!Imported || Imported->LODModels.Num() == 0 || BodyT.Rest.Num() != BodyT.Names.Num())
		{
			return Capsules;
		}
		const FReferenceSkeleton& MeshRef = Mesh->GetRefSkeleton();
		TArray<int32> Nodes;
		Nodes.Add(BoneOf(BodyT, TEXT("pelvis")));
		for (int32 Index = 0; ; ++Index)
		{
			const int32 Bone = BoneOf(BodyT, TEXT("spine"), Index);
			if (Bone == INDEX_NONE)
			{
				break;
			}
			Nodes.Add(Bone);
		}
		Nodes.Add(BoneOf(BodyT, TEXT("neck")));
		Nodes.RemoveAll([](int32 Bone) { return Bone == INDEX_NONE; });
		struct FSegment { int32 Bone; FVector A; FVector B; TArray<float> Distances; };
		TArray<FSegment> Segments;
		for (int32 Index = 0; Index + 1 < Nodes.Num(); ++Index)
		{
			const FVector A = BodyT.Rest[Nodes[Index]].GetLocation();
			const FVector B = BodyT.Rest[Nodes[Index + 1]].GetLocation();
			if (FVector::Dist(A, B) > 2.f)
			{
				Segments.Add({ Nodes[Index], A, B, {} });
			}
		}
		for (const FSkelMeshSection& Section : Imported->LODModels[0].Sections)
		{
			for (const FSoftSkinVertex& Vertex : Section.SoftVertices)
			{
				int32 Best = 0;
				for (int32 Influence = 1; Influence < MAX_TOTAL_INFLUENCES; ++Influence)
				{
					if (Vertex.InfluenceWeights[Influence] > Vertex.InfluenceWeights[Best])
					{
						Best = Influence;
					}
				}
				if (!Section.BoneMap.IsValidIndex(Vertex.InfluenceBones[Best]))
				{
					continue;
				}
				const int32 Bone = BodyT.Names.IndexOfByKey(MeshRef.GetBoneName(Section.BoneMap[Vertex.InfluenceBones[Best]]));
				if (!Trunk.Contains(Bone))
				{
					continue;
				}
				const FVector P(Vertex.Position);
				for (FSegment& Segment : Segments)
				{
					const FVector AB = Segment.B - Segment.A;
					const float T = static_cast<float>(FVector::DotProduct(P - Segment.A, AB) / AB.SizeSquared());
					if (T >= 0.f && T <= 1.f)
					{
						Segment.Distances.Add(static_cast<float>(FVector::Dist(P, Segment.A + AB * T)));
						break;
					}
				}
			}
		}
		for (FSegment& Segment : Segments)
		{
			if (Segment.Distances.Num() < 20)
			{
				continue;
			}
			Segment.Distances.Sort();
			const FTransform& Rest = BodyT.Rest[Segment.Bone];
			FLPCheckCapsule& Capsule = Capsules.AddDefaulted_GetRef();
			Capsule.Bone = Segment.Bone;
			Capsule.Center = Rest.InverseTransformPosition((Segment.A + Segment.B) * 0.5f);
			Capsule.Axis = Rest.InverseTransformVectorNoScale((Segment.B - Segment.A).GetSafeNormal());
			Capsule.HalfLength = static_cast<float>(FVector::Dist(Segment.A, Segment.B)) * 0.5f;
			Capsule.Radius = Segment.Distances[FMath::FloorToInt(Segment.Distances.Num() * SkinPercentile)];
		}
		return Capsules;
	}

	float SegmentDistance(const FVector& P, const FVector& A, const FVector& B)
	{
		const FVector AB = B - A;
		const float T = AB.IsNearlyZero() ? 0.f : FMath::Clamp(static_cast<float>(FVector::DotProduct(P - A, AB) / AB.SizeSquared()), 0.f, 1.f);
		return static_cast<float>(FVector::Dist(P, A + AB * T));
	}
}

namespace LPAnimChecks
{
	TArray<FLPCheckCapsule> TrunkCapsules(const USkeletalMesh* Mesh, const FLPPose& BodyT)
	{
		TArray<FLPCheckCapsule> Capsules;
		const UPhysicsAsset* Physics = Mesh ? Mesh->GetPhysicsAsset() : nullptr;
		TSet<int32> Trunk;
		for (const TCHAR* Role : TrunkRoles)
		{
			for (int32 Index = 0; ; ++Index)
			{
				const int32 Bone = BoneOf(BodyT, Role, Index);
				if (Bone == INDEX_NONE)
				{
					break;
				}
				Trunk.Add(Bone);
			}
		}
		if (!Physics)
		{
			return SkinCapsules(Mesh, BodyT, Trunk);
		}
		for (const TObjectPtr<USkeletalBodySetup>& Setup : Physics->SkeletalBodySetups)
		{
			const int32 Bone = Setup ? BodyT.Names.IndexOfByKey(Setup->BoneName) : INDEX_NONE;
			if (Bone == INDEX_NONE || !Trunk.Contains(Bone))
			{
				continue;
			}
			for (const FKSphylElem& Sphyl : Setup->AggGeom.SphylElems)
			{
				FLPCheckCapsule& Capsule = Capsules.AddDefaulted_GetRef();
				Capsule.Bone = Bone;
				Capsule.Center = Sphyl.Center;
				Capsule.Axis = Sphyl.Rotation.Quaternion().GetAxisZ();
				Capsule.HalfLength = Sphyl.Length * 0.5f;
				Capsule.Radius = Sphyl.Radius;
			}
			for (const FKSphereElem& Sphere : Setup->AggGeom.SphereElems)
			{
				FLPCheckCapsule& Capsule = Capsules.AddDefaulted_GetRef();
				Capsule.Bone = Bone;
				Capsule.Center = Sphere.Center;
				Capsule.Radius = Sphere.Radius;
			}
		}
		return Capsules.Num() > 0 ? Capsules : SkinCapsules(Mesh, BodyT, Trunk);
	}

	float SelfTest(const TArray<FLPCheckCapsule>& Trunk, const FLPPose& BodyT)
	{
		float Worst = Trunk.Num() > 0 ? 0.f : -1.f;
		for (const FLPCheckCapsule& Capsule : Trunk)
		{
			const FTransform& Bone = BodyT.Component[Capsule.Bone];
			const FVector Center = Bone.TransformPosition(Capsule.Center);
			const FVector Axis = Bone.TransformVectorNoScale(Capsule.Axis) * Capsule.HalfLength;
			// L'articulation de l'os porteur est une extremite du segment : profondeur egale au rayon si le centre
			// et l'axe ont ete bien ramenes dans le repere de l'os.
			const float Depth = Capsule.Radius - SegmentDistance(Bone.GetLocation(), Center - Axis, Center + Axis);
			Worst = FMath::Max(Worst, FMath::Abs(Depth - Capsule.Radius));
		}
		return Worst;
	}

	FLPClipChecks CheckClip(const FString& RefPath, const FString& BodyPath, const FLPPose& RefT, const FLPPose& BodyT,
		const TArray<FLPCheckCapsule>& Trunk)
	{
		FLPClipChecks Out;
		Out.Clip = FPaths::GetBaseFilename(BodyPath);
		const UAnimSequence* Anim = LoadObject<UAnimSequence>(nullptr, *BodyPath);
		const IAnimationDataModel* Model = Anim ? Anim->GetDataModel() : nullptr;
		if (!Model)
		{
			Out.Error = TEXT("clip converti absent");
			return Out;
		}
		const FFrameRate Rate = Model->GetFrameRate();
		Out.Frames = Model->GetNumberOfKeys();

		// Mise a l'echelle source -> corps : hauteur de hanche des deux poses T.
		const int32 RefHips = BoneOf(RefT, TEXT("hips"));
		const int32 BodyHips = BoneOf(BodyT, TEXT("hips"));
		const float Scale = RefHips != INDEX_NONE && BodyHips != INDEX_NONE && At(RefT, RefHips).Z > 1.f
			? static_cast<float>(At(BodyT, BodyHips).Z / At(RefT, RefHips).Z) : 1.f;

		// Points de contact (pieds, orteils) et articulations (coude, genou), par role.
		const TCHAR* const ContactRoles[] = { TEXT("foot_l"), TEXT("toe_l"), TEXT("foot_r"), TEXT("toe_r") };
		struct FJoint { const TCHAR* A; const TCHAR* B; const TCHAR* C; bool bKnee; };
		const FJoint Joints[] = {
			{ TEXT("thigh_l"), TEXT("calf_l"), TEXT("foot_l"), true },
			{ TEXT("thigh_r"), TEXT("calf_r"), TEXT("foot_r"), true },
			{ TEXT("upperarm_l"), TEXT("lowerarm_l"), TEXT("hand_l"), false },
			{ TEXT("upperarm_r"), TEXT("lowerarm_r"), TEXT("hand_r"), false },
		};
		const TCHAR* const JumpRoles[] = {
			TEXT("hips"), TEXT("head"), TEXT("clavicle_l"), TEXT("upperarm_l"), TEXT("lowerarm_l"), TEXT("hand_l"),
			TEXT("clavicle_r"), TEXT("upperarm_r"), TEXT("lowerarm_r"), TEXT("hand_r"),
			TEXT("thigh_l"), TEXT("calf_l"), TEXT("foot_l"), TEXT("thigh_r"), TEXT("calf_r"), TEXT("foot_r")
		};
		const TCHAR* const ProbeRoles[] = { TEXT("hand_l"), TEXT("lowerarm_l"), TEXT("hand_r"), TEXT("lowerarm_r") };

		FLPPose PrevRef;
		FLPPose PrevBody;
		for (int32 Key = 0; Key < Out.Frames; ++Key)
		{
			const float Time = static_cast<float>(Rate.AsSeconds(FFrameTime(Key)));
			FLPPose Ref;
			FLPPose Body;
			if (!LPPoseMeasure::Sample(RefPath, Time, Ref) || !LPPoseMeasure::Sample(BodyPath, Time, Body))
			{
				Out.Error = FString::Printf(TEXT("image %d illisible"), Key);
				return Out;
			}
			for (const TCHAR* Role : ContactRoles)
			{
				const int32 R = BoneOf(Ref, Role);
				const int32 B = BoneOf(Body, Role);
				if (R == INDEX_NONE || B == INDEX_NONE)
				{
					continue;
				}
				// Hauteur au-dessus de la pose T (pieds a plat), source a l'echelle du corps.
				const float RefHeight = static_cast<float>(At(Ref, R).Z - At(RefT, R).Z) * Scale;
				const float BodyHeight = static_cast<float>(At(Body, B).Z - At(BodyT, B).Z);
				const float Below = FMath::Max(0.f, -BodyHeight) - FMath::Max(0.f, -RefHeight);
				if (Below > Out.Floor)
				{
					Out.Floor = Below;
					Out.FloorBone = Body.Names[B].ToString();
				}
				if (Key > 0)
				{
					const FVector RefMove = (At(Ref, R) - At(PrevRef, R)) * Scale;
					const FVector BodyMove = At(Body, B) - At(PrevBody, B);
					const bool bPlanted = RefHeight < ContactHeight && FVector(RefMove.X, RefMove.Y, 0.f).Size() < PlantedMove;
					const float Slide = bPlanted ? static_cast<float>(FVector(BodyMove.X - RefMove.X, BodyMove.Y - RefMove.Y, 0.f).Size()) : 0.f;
					if (Slide > Out.Slide)
					{
						Out.Slide = Slide;
						Out.SlideBone = Body.Names[B].ToString();
					}
				}
			}
			for (const FJoint& Joint : Joints)
			{
				const int32 RA = BoneOf(Ref, Joint.A), RB = BoneOf(Ref, Joint.B), RC = BoneOf(Ref, Joint.C);
				const int32 BA = BoneOf(Body, Joint.A), BB = BoneOf(Body, Joint.B), BC = BoneOf(Body, Joint.C);
				if (RA == INDEX_NONE || RB == INDEX_NONE || RC == INDEX_NONE || BA == INDEX_NONE || BB == INDEX_NONE || BC == INDEX_NONE)
				{
					continue;
				}
				const float Gap = FMath::Abs(Flexion(Body, BA, BB, BC) - Flexion(Ref, RA, RB, RC));
				if (Gap > Out.FlexGap)
				{
					Out.FlexGap = Gap;
					Out.FlexJoint = Body.Names[BB].ToString();
				}
				if (Joint.bKnee)
				{
					const float Behind = KneeBehind(Body, BA, BB, BC) - KneeBehind(Ref, RA, RB, RC) * Scale;
					if (Behind > Out.KneeBack)
					{
						Out.KneeBack = Behind;
						Out.KneeBone = Body.Names[BB].ToString();
					}
				}
			}
			if (Key > 0)
			{
				for (const TCHAR* Role : JumpRoles)
				{
					const int32 R = BoneOf(Ref, Role);
					const int32 B = BoneOf(Body, Role);
					if (R == INDEX_NONE || B == INDEX_NONE)
					{
						continue;
					}
					const float RefTurn = FMath::RadiansToDegrees(Ref.Component[R].GetRotation().AngularDistance(PrevRef.Component[R].GetRotation()));
					const float BodyTurn = FMath::RadiansToDegrees(Body.Component[B].GetRotation().AngularDistance(PrevBody.Component[B].GetRotation()));
					const float Excess = BodyTurn - (RefTurn * JumpFactor + JumpMargin);
					if (Excess > Out.Jump)
					{
						Out.Jump = Excess;
						Out.JumpBone = Body.Names[B].ToString();
					}
				}
			}
			for (const TCHAR* Role : ProbeRoles)
			{
				const int32 B = BoneOf(Body, Role);
				if (B == INDEX_NONE)
				{
					continue;
				}
				const FVector Point = At(Body, B);
				for (const FLPCheckCapsule& Capsule : Trunk)
				{
					const FTransform& Bone = Body.Component[Capsule.Bone];
					const FVector Center = Bone.TransformPosition(Capsule.Center);
					const FVector Axis = Bone.TransformVectorNoScale(Capsule.Axis) * Capsule.HalfLength;
					const float Depth = Capsule.Radius - SegmentDistance(Point, Center - Axis, Center + Axis);
					if (Depth > Out.Penetration)
					{
						Out.Penetration = Depth;
						Out.PenetrationBone = Body.Names[B].ToString();
					}
				}
			}
			PrevRef = MoveTemp(Ref);
			PrevBody = MoveTemp(Body);
		}
		return Out;
	}
}