// Anti-Recoil Prediction Methods for EBBarrel
// Provides linear extrapolation, weighted regression, and Kalman filter prediction
//
// Key concept: The buffer stores samples within AntiRecoilBufferTime seconds.
// Samples within the last AntiRecoilDiscardTime seconds may be contaminated by recoil shock.
// All prediction algorithms work ONLY on the "safe" (oldest) portion of the buffer
// and extrapolate forward to the current time.

#include "EBBarrel.h"

// Returns the number of safe (non-contaminated) samples based on a time threshold.
// Samples whose timestamp >= (CurrentTime - DiscardTime) are considered potentially contaminated.
static int32 GetSafeCount(const TArray<FTimestampedTransform>& History, double CurrentTime, float DiscardTime)
{
	if (History.Num() == 0) return 0;
	double SafeCutoff = CurrentTime - FMath::Max(0.0f, DiscardTime);
	int32 SafeN = 0;
	for (int32 i = 0; i < History.Num(); i++)
	{
		if (History[i].Timestamp < SafeCutoff)
		{
			SafeN = i + 1;
		}
		else
		{
			break;
		}
	}
	return SafeN;
}

void UEBBarrel::TriggerDebugIMUShock()
{
	if (!DebugSimulateIMUShock) return;

	DebugIMUShockActive = true;
	DebugIMUShockLineCaptured = false; // Reset so the new shock replaces the previous yellow line
	DebugIMUShockStartTime = GetWorld()->GetTimeSeconds();

	// Generate a random shock direction (sharp upward + random lateral, simulating recoil kick)
	FVector RandomDir = FMath::VRand();
	// Bias upward to simulate typical recoil pattern
	RandomDir.Z = FMath::Abs(RandomDir.Z) * 2.0f;
	RandomDir.Normalize();

	DebugIMUShockAimOffset = RandomDir * FMath::DegreesToRadians(DebugIMUShockAngle);
	DebugIMUShockPosOffset = RandomDir * DebugIMUShockPosition;
}

void UEBBarrel::UpdateTransformHistory()
{
	if (AntiRecoilMode == EAntiRecoilMode::ARM_None)
	{
		return;
	}

	double CurrentTime = GetWorld()->GetTimeSeconds();

	FTimestampedTransform Sample;
	Sample.Timestamp = CurrentTime;
	Sample.Location = GetComponentTransform().GetLocation();
	Sample.Aim = GetComponentTransform().GetUnitAxis(EAxis::X);

	// Apply simulated IMU shock perturbation if active
	if (DebugSimulateIMUShock && DebugIMUShockActive)
	{
		float ElapsedShock = (float)(CurrentTime - DebugIMUShockStartTime);
		if (ElapsedShock < DebugIMUShockDuration)
		{
			// Decaying shock: intensity decreases over the shock duration
			float ShockAlpha = 1.0f - (ElapsedShock / DebugIMUShockDuration);
			// Add some high-frequency noise to simulate IMU vibration
			FVector FrameNoise = FMath::VRand() * 0.3f;

			// Perturb aim direction
			FVector AimPerturbation = (DebugIMUShockAimOffset + FrameNoise * FMath::DegreesToRadians(DebugIMUShockAngle)) * ShockAlpha;
			Sample.Aim = (Sample.Aim + AimPerturbation).GetSafeNormal();

			// Perturb position
			FVector PosPerturbation = (DebugIMUShockPosOffset + FrameNoise * DebugIMUShockPosition) * ShockAlpha;
			Sample.Location += PosPerturbation;
		}
		else
		{
			DebugIMUShockActive = false;
		}
	}

	TransformHistory.Add(Sample);

	// Trim buffer: remove samples older than AntiRecoilBufferTime
	double OldestAllowed = CurrentTime - FMath::Max(0.05f, AntiRecoilBufferTime);
	while (TransformHistory.Num() > 0 && TransformHistory[0].Timestamp < OldestAllowed)
	{
		TransformHistory.RemoveAt(0);
	}
}

void UEBBarrel::ComputeAntiRecoilTransform()
{
	switch (AntiRecoilMode)
	{
	case EAntiRecoilMode::ARM_None:
		Aim = GetComponentTransform().GetUnitAxis(EAxis::X);
		Location = GetComponentTransform().GetLocation();
		break;

	case EAntiRecoilMode::ARM_Buffer:
		if (TransformHistory.Num() > 0)
		{
			Aim = TransformHistory[0].Aim;
			Location = TransformHistory[0].Location;
		}
		else
		{
			Aim = GetComponentTransform().GetUnitAxis(EAxis::X);
			Location = GetComponentTransform().GetLocation();
		}
		break;

	case EAntiRecoilMode::ARM_LinearExtrapolation:
		{
			int32 SafeN = GetSafeCount(TransformHistory, GetWorld()->GetTimeSeconds(), AntiRecoilDiscardTime);
			if (SafeN >= 2)
			{
				PredictLinearExtrapolation(GetWorld()->GetTimeSeconds(), Location, Aim);
			}
			else if (TransformHistory.Num() > 0)
			{
				Aim = TransformHistory[0].Aim;
				Location = TransformHistory[0].Location;
			}
			else
			{
				Aim = GetComponentTransform().GetUnitAxis(EAxis::X);
				Location = GetComponentTransform().GetLocation();
			}
		}
		break;

	case EAntiRecoilMode::ARM_WeightedRegression:
		{
			int32 SafeN = GetSafeCount(TransformHistory, GetWorld()->GetTimeSeconds(), AntiRecoilDiscardTime);
			if (SafeN >= 2)
			{
				PredictWeightedRegression(GetWorld()->GetTimeSeconds(), Location, Aim);
			}
			else if (TransformHistory.Num() > 0)
			{
				Aim = TransformHistory[0].Aim;
				Location = TransformHistory[0].Location;
			}
			else
			{
				Aim = GetComponentTransform().GetUnitAxis(EAxis::X);
				Location = GetComponentTransform().GetLocation();
			}
		}
		break;

	case EAntiRecoilMode::ARM_KalmanFilter:
		{
			int32 SafeN = GetSafeCount(TransformHistory, GetWorld()->GetTimeSeconds(), AntiRecoilDiscardTime);
			if (SafeN > 0)
			{
				// Feed only the latest SAFE sample to the Kalman filter
				const FTimestampedTransform& LatestSafe = TransformHistory[SafeN - 1];
				UpdateKalmanFilter(LatestSafe.Timestamp, LatestSafe.Location, LatestSafe.Aim);
				PredictKalmanFilter(GetWorld()->GetTimeSeconds(), Location, Aim);
			}
			else if (TransformHistory.Num() > 0)
			{
				Aim = TransformHistory[0].Aim;
				Location = TransformHistory[0].Location;
			}
			else
			{
				Aim = GetComponentTransform().GetUnitAxis(EAxis::X);
				Location = GetComponentTransform().GetLocation();
			}
		}
		break;
	}
}

// --- Linear Extrapolation ---
// Computes average velocity from ALL safe samples (consecutive differences),
// then extrapolates from the last safe sample to current time.

void UEBBarrel::PredictLinearExtrapolation(double CurrentTime, FVector& OutLocation, FVector& OutAim) const
{
	const int32 SafeN = GetSafeCount(TransformHistory, GetWorld()->GetTimeSeconds(), AntiRecoilDiscardTime);
	if (SafeN < 2)
	{
		OutLocation = TransformHistory[0].Location;
		OutAim = TransformHistory[0].Aim;
		return;
	}

	// Compute average velocity from consecutive safe sample differences
	FVector AvgLinearVelocity = FVector::ZeroVector;
	FVector AvgAimDelta = FVector::ZeroVector;
	double TotalDt = 0.0;
	int32 ValidPairs = 0;

	for (int32 i = 1; i < SafeN; i++)
	{
		double dt = TransformHistory[i].Timestamp - TransformHistory[i - 1].Timestamp;
		if (dt > SMALL_NUMBER)
		{
			AvgLinearVelocity += (TransformHistory[i].Location - TransformHistory[i - 1].Location) / dt;
			AvgAimDelta += (TransformHistory[i].Aim - TransformHistory[i - 1].Aim) / dt;
			TotalDt += dt;
			ValidPairs++;
		}
	}

	if (ValidPairs == 0)
	{
		OutLocation = TransformHistory[SafeN - 1].Location;
		OutAim = TransformHistory[SafeN - 1].Aim;
		return;
	}

	AvgLinearVelocity /= (double)ValidPairs;
	AvgAimDelta /= (double)ValidPairs;

	// Extrapolate from the last SAFE sample to current time
	const FTimestampedTransform& LastSafe = TransformHistory[SafeN - 1];
	double ExtrapolationTime = CurrentTime - LastSafe.Timestamp;

	OutLocation = LastSafe.Location + AvgLinearVelocity * ExtrapolationTime;

	// Angular extrapolation using quaternion slerp
	// Use first and last safe samples for rotation direction
	const FTimestampedTransform& FirstSafe = TransformHistory[0];
	double SafeDeltaT = LastSafe.Timestamp - FirstSafe.Timestamp;

	if (SafeDeltaT > SMALL_NUMBER)
	{
		FQuat FirstQuat = FRotationMatrix::MakeFromX(FirstSafe.Aim).ToQuat();
		FQuat LastQuat = FRotationMatrix::MakeFromX(LastSafe.Aim).ToQuat();

		double TotalAlpha = ExtrapolationTime / SafeDeltaT;
		FQuat PredictedQuat = FQuat::Slerp(FirstQuat, LastQuat, 1.0 + TotalAlpha);

		OutAim = PredictedQuat.GetForwardVector().GetSafeNormal();
		if (OutAim.IsNearlyZero())
		{
			OutAim = LastSafe.Aim;
		}
	}
	else
	{
		OutAim = LastSafe.Aim;
	}
}

// --- Weighted Linear Regression ---
// Fits a weighted least-squares line through only the SAFE samples, extrapolates to current time.
// More recent safe samples get higher weight.

void UEBBarrel::PredictWeightedRegression(double CurrentTime, FVector& OutLocation, FVector& OutAim) const
{
	const int32 SafeN = GetSafeCount(TransformHistory, GetWorld()->GetTimeSeconds(), AntiRecoilDiscardTime);
	if (SafeN < 2)
	{
		OutLocation = TransformHistory[0].Location;
		OutAim = TransformHistory[0].Aim;
		return;
	}

	// Use timestamps relative to the first safe sample to avoid precision issues
	double T0 = TransformHistory[0].Timestamp;

	// Weighted linear regression: y = a + b * t
	// Weights: linearly increasing (1, 2, 3, ..., SafeN)
	double SumW = 0.0;
	double SumWT = 0.0;
	double SumWTT = 0.0;
	FVector SumWY = FVector::ZeroVector;
	FVector SumWTY = FVector::ZeroVector;
	FVector SumWAim = FVector::ZeroVector;
	FVector SumWTAim = FVector::ZeroVector;

	for (int32 i = 0; i < SafeN; i++)
	{
		double w = FMath::Pow((double)(i + 1), (double)RegressionWeightExponent); // Weight curve controlled by exponent
		double t = TransformHistory[i].Timestamp - T0;

		SumW += w;
		SumWT += w * t;
		SumWTT += w * t * t;
		SumWY += TransformHistory[i].Location * w;
		SumWTY += TransformHistory[i].Location * (w * t);
		SumWAim += TransformHistory[i].Aim * w;
		SumWTAim += TransformHistory[i].Aim * (w * t);
	}

	double Det = SumW * SumWTT - SumWT * SumWT;
	if (FMath::Abs(Det) <= SMALL_NUMBER)
	{
		OutLocation = TransformHistory[SafeN - 1].Location;
		OutAim = TransformHistory[SafeN - 1].Aim;
		return;
	}

	// Solve for position: intercept (a) and slope (b)
	FVector PosIntercept = (SumWY * SumWTT - SumWTY * SumWT) / Det;
	FVector PosSlope = (SumWTY * SumW - SumWY * SumWT) / Det;

	// Solve for aim: intercept (a) and slope (b)
	FVector AimIntercept = (SumWAim * SumWTT - SumWTAim * SumWT) / Det;
	FVector AimSlope = (SumWTAim * SumW - SumWAim * SumWT) / Det;

	// Extrapolate to current time
	double TPred = CurrentTime - T0;
	OutLocation = PosIntercept + PosSlope * TPred;

	FVector PredAim = AimIntercept + AimSlope * TPred;
	OutAim = PredAim.GetSafeNormal();
	if (OutAim.IsNearlyZero())
	{
		OutAim = TransformHistory[SafeN - 1].Aim;
	}
}

// --- Simplified Kalman Filter ---
// Maintains state estimate [position, velocity, aim, angular_velocity]
// Only fed with SAFE (non-contaminated) measurements, predicts forward to current time

void UEBBarrel::UpdateKalmanFilter(double CurrentTime, const FVector& MeasuredLocation, const FVector& MeasuredAim)
{
	if (!KalmanInitialized)
	{
		KalmanPosition = MeasuredLocation;
		KalmanVelocity = FVector::ZeroVector;
		KalmanAim = MeasuredAim;
		KalmanAngularVelocity = FVector::ZeroVector;
		KalmanPosVariance = 1.0f;
		KalmanVelVariance = 1.0f;
		KalmanAimVariance = 1.0f;
		KalmanAngVelVariance = 1.0f;
		KalmanInitialized = true;
		KalmanLastTimestamp = CurrentTime;
		return;
	}

	float dt = (float)(CurrentTime - KalmanLastTimestamp);
	if (dt <= SMALL_NUMBER)
	{
		return;
	}

	KalmanLastTimestamp = CurrentTime;

	float Q = KalmanProcessNoise;
	float R = KalmanMeasurementNoise;

	// --- Position / Velocity ---
	// Predict step
	FVector PredPos = KalmanPosition + KalmanVelocity * dt;
	FVector PredVel = KalmanVelocity;
	float PredPosVar = KalmanPosVariance + KalmanVelVariance * dt * dt + Q * dt;
	float PredVelVar = KalmanVelVariance + Q * dt;

	// Update step (measurement = MeasuredLocation)
	float KGainPos = PredPosVar / (PredPosVar + R);
	KalmanPosition = PredPos + (MeasuredLocation - PredPos) * KGainPos;
	KalmanPosVariance = (1.0f - KGainPos) * PredPosVar;

	// Update velocity estimate from innovation
	FVector VelInnovation = (MeasuredLocation - PredPos) / dt;
	float KGainVel = PredVelVar / (PredVelVar + R / (dt * dt));
	KalmanVelocity = PredVel + VelInnovation * KGainVel;
	KalmanVelVariance = (1.0f - KGainVel) * PredVelVar;

	// --- Aim / Angular Velocity ---
	FVector PredAim = KalmanAim + KalmanAngularVelocity * dt;
	FVector PredAngVel = KalmanAngularVelocity;
	float PredAimVar = KalmanAimVariance + KalmanAngVelVariance * dt * dt + Q * dt;
	float PredAngVelVar = KalmanAngVelVariance + Q * dt;

	float KGainAim = PredAimVar / (PredAimVar + R);
	KalmanAim = PredAim + (MeasuredAim - PredAim) * KGainAim;
	KalmanAimVariance = (1.0f - KGainAim) * PredAimVar;

	FVector AngVelInnovation = (MeasuredAim - PredAim) / dt;
	float KGainAngVel = PredAngVelVar / (PredAngVelVar + R / (dt * dt));
	KalmanAngularVelocity = PredAngVel + AngVelInnovation * KGainAngVel;
	KalmanAngVelVariance = (1.0f - KGainAngVel) * PredAngVelVar;
}

void UEBBarrel::PredictKalmanFilter(double CurrentTime, FVector& OutLocation, FVector& OutAim) const
{
	if (!KalmanInitialized)
	{
		OutLocation = GetComponentTransform().GetLocation();
		OutAim = GetComponentTransform().GetUnitAxis(EAxis::X);
		return;
	}

	// Extrapolate from Kalman state to current time
	float dt = (float)(CurrentTime - KalmanLastTimestamp);

	OutLocation = KalmanPosition + KalmanVelocity * dt;

	FVector PredAim = KalmanAim + KalmanAngularVelocity * dt;
	OutAim = PredAim.GetSafeNormal();
	if (OutAim.IsNearlyZero())
	{
		OutAim = KalmanAim.GetSafeNormal();
	}
}
