// Fill out your copyright notice in the Description page of Project Settings.


#include "MultiRotorDrone.h"
#include "Misc/DateTime.h"
#include <stdlib.h>
#include <math.h>
#include "Kismet/GameplayStatics.h"

// Sets default values
AMultiRotorDrone::AMultiRotorDrone()
{
	// Set this pawn to call Tick() every frame.  You can turn this off to improve performance if you don't need it.
	PrimaryActorTick.bCanEverTick = true;

}

// Called when the game starts or when spawned
void AMultiRotorDrone::BeginPlay()
{
	Super::BeginPlay();
	Kt = AIR_DENCITY * Rotors.Ct * powf(Rotors.Diameter, 4);
	Km = AIR_DENCITY * Rotors.Cq * powf(Rotors.Diameter, 5);
	minThrust = 0.5 * Mass * GRAVITY;

	maxThrustPerMotor = ((MaxRPM / 60) * (MaxRPM / 60)) * Kt;

	MaxThrust = maxThrustPerMotor * 4;

	SrefZ = sqrtf(Rotors.p1.X * Rotors.p1.X + Rotors.p1.Y * Rotors.p1.Y) * sqrtf(Rotors.p2.X * Rotors.p2.X + Rotors.p2.Y * Rotors.p2.Y);

	COMLocation = { -0.02, 0.0, -0.29 };

	if (Mode == GUIDED)
	{
		if (!use_true_state)
		{
			XYerrTol = 0.75;
		}
	}
}

void AMultiRotorDrone::SetTrueState(FTrueDroneState true_state)
{
	FVector prev_position = TrueState.Position;
	FVector prev_euler = TrueState.Orientation;
	FVector prev_velocity = TrueState.Velocity;
	TrueState.Position = true_state.Position * 0.01;
	TrueState.Position.Z *= -1.0;
	TrueState.Orientation = true_state.Orientation * (PI / 180);
	if (!TrueState.initialised)
	{
		if (use_true_state)
		{
			TargetPosition = TrueState.Position;
			TargetYaw = TrueState.Orientation.Z;
		}
		TrueState.initialised = true;
		return;
	}
	//UE_LOG(LogTemp, Warning, TEXT("Prev position %f %f %f"), prev_position.X, prev_position.Y, prev_position.Z)
	//UE_LOG(LogTemp, Warning, TEXT("Curr position %f %f %f"), TrueState.Position.X, TrueState.Position.Y, TrueState.Position.Z)
	TrueState.Velocity = (TrueState.Position - prev_position) / MjStep;
	TrueState.Acceleration = (TrueState.Velocity - prev_velocity) / MjStep;
	FVector EulerDt = (TrueState.Orientation - prev_euler) / MjStep;
	EulerDt2PQR(EulerDt);
}

void AMultiRotorDrone::SetCurrentState()
{
	FVector prev_position = State.Position;
	FVector prev_euler = State.Orientation;
	FVector prev_velocity = State.Velocity;
	if (use_true_state)
	{
		State = TrueState;
		return;
	}
	State.Position = EstState.position;
	State.Orientation = EstState.quat.GetEulerRad();
	State.Orientation.Z = State.Orientation.Z;
	State.Velocity = EstState.velocity;
	//UE_LOG(LogTemp, Display, TEXT("True orientation degrees %f %f %f"), true_state.Orientation.X, true_state.Orientation.Y, true_state.Orientation.Z);
	if (!State.initialised)
	{
		TargetPosition = State.Position;
		TargetYaw = State.Orientation.Z;
		State.initialised = true;
	}
	State.Acceleration = EstState.acceleration;
	PQR = EstState.pqr;
}

FDroneData AMultiRotorDrone::GetDronInfo()
{
	FDroneData drone_data;
	for (int32 i = 0; i < 4; i++)
	{
		drone_data.MotorsRPM[i] = RPM[i] * 60;
	}
	drone_data.Velocity = State.Velocity;
	drone_data.Position = State.Position;
	drone_data.Orientation = State.Orientation;

	return drone_data;
}

void AMultiRotorDrone::EulerDt2PQR(FVector& EulerDt)
{
	float r = State.Orientation.X;
	float p = State.Orientation.Y;
	float y = State.Orientation.Z;

	float srtp = sinf(r) * tanf(p);
	float crtp = cosf(r) * tanf(p);
	float srsp = sinf(r) / cosf(p);
	float crsp = cosf(r) / cosf(p);

	Matrix33 R(1, srtp, crtp, 0, cosf(p), -sinf(p), 0, srsp, crsp);

	PQR = R.Inverse() * EulerDt;
}

void AMultiRotorDrone::SetImuSensor(UIMUSensor* sensor)
{
	Imu_sensor = sensor;
	StateEstimator.RegisterBuffer(Sensor::GYRO, Imu_sensor->GyroDataBuffer);
	StateEstimator.RegisterBuffer(Sensor::ACCEL, Imu_sensor->AccelDataBuffer);
	StateEstimator.RegisterBuffer(Sensor::MAG, Imu_sensor->MagDataBuffer);
}

void  AMultiRotorDrone::SetGPSSensor(UGPSSensor* sensor)
{
	GPS_sensor = sensor;
	StateEstimator.RegisterBuffer(Sensor::GPS, GPS_sensor->DataBuffer);
}

void AMultiRotorDrone::SetBaroSensor(UBarometer* sensor)
{
	Baro_sensor = sensor;
	StateEstimator.RegisterBuffer(Sensor::BARO, Baro_sensor->DataBuffer);
}

// Called every frame
void AMultiRotorDrone::Tick(float DeltaTime)
{
	Super::Tick(DeltaTime);
	//ProcessIMUData();

	if (!StateEstimator.IsInitialised())
	{
		StateEstimator.InitialiseFilter();
	}
	else
	{
		StateEstimator.Update();
		StateEstimator.GetState(EstState);
		SetCurrentState();
		ControllerUpdate();
	}

	UE_LOG(LogTemp, Warning, TEXT("EKF estimated position: %f, %f, %f!"), EstState.position.X, EstState.position.Y, EstState.position.Z);
	UE_LOG(LogTemp, Warning, TEXT("EKF estimated velocity: %f, %f, %f!"), EstState.velocity.X, EstState.velocity.Y, EstState.velocity.Z);
	FVector Euler = EstState.quat.GetEulerRad();
	UE_LOG(LogTemp, Warning, TEXT("EKF estimated orientation: %f, %f, %f!"), Euler.X, Euler.Y, Euler.Z);

	UE_LOG(LogTemp, Warning, TEXT("True position: %f, %f, %f!"), TrueState.Position.X, TrueState.Position.Y, TrueState.Position.Z);
	UE_LOG(LogTemp, Warning, TEXT("True velocity: %f, %f, %f!"), TrueState.Velocity.X, TrueState.Velocity.Y, TrueState.Velocity.Z);
	UE_LOG(LogTemp, Warning, TEXT("True orientation: %f, %f, %f!"), TrueState.Orientation.X, TrueState.Orientation.Y, TrueState.Orientation.Z);
}

// Called to bind functionality to input
void AMultiRotorDrone::SetupPlayerInputComponent(UInputComponent* PlayerInputComponent)
{
	Super::SetupPlayerInputComponent(PlayerInputComponent);
	//PlayerInputComponent->BindAxis(TEXT("height_increase"), this, &AMultiRotorDrone::IncreaseTargetHeight);
}

void AMultiRotorDrone::IncreaseDecreaseTargetHeight(float value)
{
	float dt = UGameplayStatics::GetWorldDeltaSeconds(this);
	if (Mode == GUIDED)
	{
		LandingStarted = false;
		TargetPosition.Z -= MaxVelocity * dt * value;
		//GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::White, FString::SanitizeFloat(TargetPosition.Z));
	}
	else if (Mode == MANUAL)
	{
		if (KeepAltitude) KeepAltitude = false;
		if (!AltitudeManualControl) AltitudeManualControl = true;
		CurrentThrottle += ThrottleRate * dt * value;
		if (CurrentThrottle > PWM_max_us)
		{
			CurrentThrottle = PWM_max_us;
		}

		if (CurrentThrottle < PWM_min_us)
		{
			CurrentThrottle = PWM_min_us;
		}

		//GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::White, FString::SanitizeFloat(CurrentThrottle));
	}
	//UE_LOG(LogTemp, Display, TEXT("Current Height Desired: %f %f"), TargetPosition.Z, value);
}

void AMultiRotorDrone::HoldAltitude()
{
	//TargetPosition.Z = State.Position.Z;
	//TargetPosition.X = State.Position.X;
	//TargetPosition.Y = State.Position.Y;
	//KeepAltitude = true;
	//KeepXYPosition = true;
	CurrentThrottle = 1500;
	pid.Xvel_int = 0.0;
	pid.Yvel_int = 0.0;
}

void AMultiRotorDrone::MoveLeftRight(float value)
{
	float dt = UGameplayStatics::GetWorldDeltaSeconds(this);
	LandingStarted = false;
	if (Mode == GUIDED)
	{
		TargetPosition.Y += MaxVelocity * dt * value;
		//GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::White, FString::SanitizeFloat(TargetPosition.Y));
	}
	else if (Mode == MANUAL)
	{
		if (KeepXYPosition) KeepXYPosition = false;
		if (!XYmanualControl) XYmanualControl = true;
		CurrentRollPWM += RollPitchPWMRate * dt * value;
		if (CurrentRollPWM > PWM_max_us)
		{
			CurrentRollPWM = PWM_max_us;
		}

		if (CurrentRollPWM < PWM_min_us)
		{
			CurrentRollPWM = PWM_min_us;
		}

		//GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::White, FString::SanitizeFloat(CurrentRollPWM));
	}
}

void AMultiRotorDrone::MoveForwardBackward(float value)
{
	LandingStarted = false;
	float dt = UGameplayStatics::GetWorldDeltaSeconds(this);
	if (Mode == GUIDED)
	{
		TargetPosition.X += MaxVelocity * dt * -value;
		//GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::White, FString::SanitizeFloat(TargetPosition.X));
	}
	else if (Mode == MANUAL)
	{
		if (KeepXYPosition) KeepXYPosition = false;
		CurrentPitchPWM += RollPitchPWMRate * dt * value;
		if (CurrentPitchPWM > PWM_max_us)
		{
			CurrentPitchPWM = PWM_max_us;
		}

		if (CurrentPitchPWM < PWM_min_us)
		{
			CurrentPitchPWM = PWM_min_us;
		}
	}
}

void AMultiRotorDrone::HoldXYPosition(bool X, bool Y)
{
	//TargetPosition.Z = State.Position.Z;
	//TargetPosition.X = State.Position.X;
	//TargetPosition.Y = State.Position.Y;
	//KeepAltitude = true;
	KeepXYPosition = true;
	if (X) CurrentPitchPWM = 1500;
	if (Y) CurrentRollPWM = 1500;
	pid.Xvel_int = 0.0;
	pid.Yvel_int = 0.0;
}

void AMultiRotorDrone::UpdateYaw(float value)
{
	if (Mode == MANUAL && KeepXYPosition)
	{
		KeepXYPosition = false;
	}
	float dt = UGameplayStatics::GetWorldDeltaSeconds(this);
	TargetYaw = wrap_PI(TargetYaw + YawRate * dt * value);
	//GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::White, FString::SanitizeFloat(TargetYaw));
}

float AMultiRotorDrone::GetJ(FVector PropLoc, float RPS, float D)
{
	FQuaternion orient_quat;
	orient_quat.InitializeFromEulerAngles(State.Orientation.X, State.Orientation.Y, State.Orientation.Z);
	Matrix33 mT = orient_quat.GetMatrix().Transposed();
	FVector Vb = mT * (State.Velocity * -1);
	FVector localAeroVel = Vb + FVector::CrossProduct(PQR, PropLoc);
	FVector direction(0, 0, -1);
	float lVeln = FVector::DotProduct(localAeroVel, direction);
	FVector vVeln = direction * lVeln;
	float V = vVeln.Size();
	float J;
	if (RPS > 0.01)
	{
		J = V / (RPS * D);
	}
	else
	{
		J = V / D;
	}

	return J;
}

void AMultiRotorDrone::AltitudeController(float& target_height, float& AltCmd)
{
	float current_height = State.Position.Z;;
	float vert_velocity = State.Velocity.Z;

	float max_rps = MaxRPM / 60;
	float min_rps = MinRPM / 60;
	float up_limit = 0.5 * 4 * max_rps * max_rps * Kt;
	//float up_limit = 1.5 * Mass * GRAVITY;
	float low_limit = 0.1 * Mass * GRAVITY;

	if (Mode == GUIDED)
	{
		UE_LOG(LogTemp, Display, TEXT("Target height %f"), target_height);
		float height_error = target_height - current_height;
		if (!use_true_state && fabs(height_error) <= 0.5)
		{
			height_error = 0.0;
		}

		float target_vel = height_error / MjStep;
		float tAcc = pid.Kpz * height_error;
		float dEr = (height_error - pid.prev_XYZ_error.Z) / MjStep;
		pid.prev_XYZ_error.Z = height_error;
		float VelError = target_vel - State.Velocity.Z;
		//float tAcc = tVel - pid.Kvz * State.Velocity.Z;
		//if (current_height > target_height)
		//{
			tAcc -= pid.Kvz * State.Velocity.Z;
		//}
		//float tAcc = tVel + pid.Kd_z * dEr;
		AltCmd = -tAcc + GRAVITY * Mass;
		//AltCmd = -tAcc * Mass;
		//UE_LOG(LogTemp, Display, TEXT("Target Z position %f"), TargetPosition.Z);
		//UE_LOG(LogTemp, Display, TEXT("Current Z position %f"), State.Position.Z);
	}
	else
	{
		float ratio = (CurrentThrottle - PWM_min_us) / (PWM_max_us - PWM_min_us);
		AltCmd = ratio * MaxThrust;
		float target_vel = -MaxVelocity + 2 * MaxVelocity * ratio;
		UE_LOG(LogTemp, Display, TEXT("Target Z velocity %f"), target_vel);
		float VelError = -target_vel - State.Velocity.Z;
		float accel_ff = VelError / MjStep;
		pid.Zvel_int += VelError * MjStep;
		//AltCmd = -pid.Kd_z * VelError - pid.Ki_z * pid.Zvel_int - accel_ff;
		AltCmd = -pid.Kd_z * VelError - pid.Ki_z * pid.Zvel_int;
		AltCmd *= Mass;
		//GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::White, FString::SanitizeFloat(target_vel));
	}

	AltCmd = fmax(fmin(AltCmd, up_limit), low_limit);
	ThrustCmd = AltCmd;
	UE_LOG(LogTemp, Display, TEXT("Alt CMD %f"), AltCmd);
}

void AMultiRotorDrone::YawController(float& target_yaw, float& tau_yaw)
{
	//UE_LOG(LogTemp, Display, TEXT("Target yaw %f"), TargetYaw);
	float yaw = State.Orientation.Z;
	float r = PQR.Z;
	float er = wrap_PI(target_yaw - yaw);

	tau_yaw = pid.Kp_yaw * er - pid.Kd_yaw * r;
	float max_rps = MaxRPM / 60;
	float min_rps = MinRPM / 60;
	float up_limit = 2 * max_rps * max_rps * Km - 2 * min_rps * min_rps * Km;
	float low_limit = -up_limit;
	tau_yaw = fmax(fmin(tau_yaw, up_limit), low_limit);
	//UE_LOG(LogTemp, Display, TEXT("Tau yaw %f"), tau_yaw);
}

void AMultiRotorDrone::RollPitchCmdController(float& Xref, float& Yref, float& roll_cmd, float& pitch_cmd)
{
	float posX, posY;
	float velX, velY;

	posX = State.Position.X;
	posY = State.Position.Y;

	velX = State.Velocity.X;
	velY = State.Velocity.Y;

	Matrix33 yaw_rot;
	yaw_rot.FromEuler(0.0, 0.0, State.Orientation.Z);

	float max_roll = MaxRoll;
	float max_pitch = MaxPitch;
	FVector XYvel(velX, velY, 0);
	FVector Kp(pid.Kpx, pid.Kpy, 0);
	FVector Kd = { -pid.Kd_x, pid.Kd_y, 0.0 };
	FVector Kv(pid.Kvx, pid.Kvy, 0);
	FVector rp_cmd;
	if (Mode == GUIDED)
	{

		FVector PosXY(posX, posY, 0);
		FVector XYref(Xref, Yref, 0);
		//FVector Ki(-pid.I_xy, pid.I_xy, 0.0);
		bool insideX = false;
		bool insideY = false;
		FVector Error = yaw_rot * (XYref - PosXY);
		if (fabs(Error.Y) <= XYerrTol)
		{
			insideY = true;
			if (fabsf(XYvel.Y) <= 0.5)
			{
				Error.Y = 0.0f;
			}
		}
		if (fabs(Error.X) <= XYerrTol)
		{
			insideX = true;
			if (fabsf(XYvel.X) <= 0.5)
			{
				Error.X = 0.0f;
			}
		}

		float err_norm = Error.Size();
		float vel_norm = XYvel.Size();
		//UE_LOG(LogTemp, Display, TEXT("Vel size %f"), vel_norm);
		if (!insideX || !insideY)
		{
			if (vel_norm >= 1.0)
			{
				FVector hit_vector = (XYvel / vel_norm) * err_norm;
				FVector hit_dist_vector = Error - hit_vector;
				if (hit_dist_vector.Size() < XYerrTol)
				{
					//UE_LOG(LogTemp, Display, TEXT("Error X Y %f %f"), Error.X, Error.Y);
					Error = hit_vector;
					//UE_LOG(LogTemp, Display, TEXT("Error X Y %f %f"), Error.X, Error.Y);
				}
			}
			rp_cmd = Kp * Error + Kv * XYvel;
		}
		else
		{
			if (vel_norm > 0.5)
			{
				rp_cmd = Kv * XYvel;
			}
			else
			{
				rp_cmd = { 0.0, 0.0, 0.0 };
			}
		}
		
		//UE_LOG(LogTemp, Display, TEXT("RP command X Y %f %f"), rp_cmd.X, rp_cmd.Y);
		//UE_LOG(LogTemp, Display, TEXT("RP command X Y %f %f"), rp_cmd.X, rp_cmd.Y);
		/*
		float xy_tol = XYerrTol;
		if (insideX && insideY)
		{
			if (abs(velX) < 0.5 && abs(velY) < 0.5)
			{
				Error.Y = 0.0f;
				Error.X = 0.0f;
			}
			else
			{
				float md = fmin(fabsf(Error.Y), fabsf(Error.X));
				xy_tol = md / 2;
				if (fabsf(velX) < 1.0 && fabsf(velY) < 1.0)
				{
					max_roll = PI / 10;
					max_pitch = PI / 10;
				}
			}
		}
		*/
		
		float XYclose_radius = 5.0;

		//rp_cmd = Kp * Error + Kv * XYvel;
		FVector d_error = (Error - pid.prev_XYZ_error) / MjStep;
		pid.prev_XYZ_error.X = Error.X;
		pid.prev_XYZ_error.Y = Error.Y;
		//rp_cmd = Kp * Error + Kd * d_error;
		
		//FVector rp_cmd = Kp * Error;

		roll_cmd = fmax(fmin(rp_cmd.Y, max_roll), -max_roll);
		pitch_cmd = fmax(fmin(rp_cmd.X, max_pitch), -max_pitch);
		UE_LOG(LogTemp, Display, TEXT("Target X position %f"), TargetPosition.X);
        //UE_LOG(LogTemp, Display, TEXT("Current X position %f"), State.Position.X);
		UE_LOG(LogTemp, Display, TEXT("Target Y position %f"), TargetPosition.Y);
		//UE_LOG(LogTemp, Display, TEXT("Current Y position %f"), State.Position.Y);
	}
	else
	{
		float bx_c, by_c;
		if (KeepXYPosition && XYvel.Size2D() > 0.1)
		{
			rp_cmd = Kv * XYvel;
			//rp_cmd.X = fmax(-MaxAcceleration, fmin(rp_cmd.X, MaxAcceleration));
		    //rp_cmd.Y = fmax(-MaxAcceleration, fmin(rp_cmd.Y, MaxAcceleration));
		    //float c = ThrustCmd / Mass;

			//pitch_cmd = atan2f(rp_cmd.X, c);
			//roll_cmd = atan2f(rp_cmd.Y, c);
			roll_cmd = fmax(-MaxRoll, fmin(rp_cmd.Y, MaxRoll));
		    pitch_cmd = fmax(-MaxPitch, fmin(rp_cmd.X, MaxPitch));
		}
		else
		{
			if (KeepXYPosition)
			{
				KeepXYPosition = false;
			}
			float roll_ratio = (CurrentRollPWM - PWM_min_us) / (PWM_max_us - PWM_min_us);
			float pitch_ratio = (CurrentPitchPWM - PWM_min_us) / (PWM_max_us - PWM_min_us);

			roll_cmd = -MaxRoll + roll_ratio * 2 * MaxRoll;
			pitch_cmd = -MaxPitch + pitch_ratio * 2 * MaxPitch;
		}

		FQuaternion roll_pitch_quat;
		roll_pitch_quat.InitializeFromEulerAngles(roll_cmd, pitch_cmd, 0);
		Matrix33 roll_pitch_mat = roll_pitch_quat.GetMatrix();
		FVector z_axis(0, 0, 1);
		z_axis = roll_pitch_mat * z_axis;
		bx_c = -z_axis.X;
		by_c = -z_axis.Y;
		//UE_LOG(LogTemp, Display, TEXT("bx by commanded %f %f "), bx_c, by_c);

		FVector euler = State.Orientation;
		FQuaternion quat;
		quat.InitializeFromEulerAngles(euler.X, euler.Y, 0);
		Matrix33 mat = quat.GetMatrix();
		//Matrix33 mat;
		//mat.FromEuler(euler.X, euler.Y, 0);
		float bx = mat(0, 2);
		float by = mat(1, 2);
		//UE_LOG(LogTemp, Display, TEXT("Current bx by %f %f"), bx, by);

		float bx_dot = pid2.Kp_roll * (bx_c - bx);
		float by_dot = pid2.Kp_pitch * (by_c - by);

		Matrix33 mat2(mat(1, 0), -mat(0, 0), 0.0, mat(1, 1), -mat(0, 1), 0.0, 0.0, 0.0, 0.0);
		mat2 = mat2 * (1 / mat(2, 2));
		FVector b(bx_dot, by_dot, 0.0);
		FVector pq_c = mat2 * b;
		float p_c = pq_c.X;
		float q_c = pq_c.Y;
		//UE_LOG(LogTemp, Display, TEXT("Commanded p q : %f %f"), p_c, q_c);
		p_c = fmax(-MaxRPrate, fmin(p_c, MaxRPrate));
		q_c = fmax(-MaxRPrate, fmin(q_c, MaxRPrate));

		float p = PQR.X;
		float q = PQR.Y;
		float r = PQR.Z;

		roll_cmd = pid2.Kp_p * (p_c - p);
		pitch_cmd = pid2.Kp_q * (q_c - q);
	}
}

void AMultiRotorDrone::AttitudeController(float& roll_cmd, float& pitch_cmd, float& roll_tau, float& pitch_tau)
{
	if (Mode == GUIDED)
	{
		FVector pr(State.Orientation.Y, State.Orientation.X, 0.0);
		FVector pr_ref(pitch_cmd, roll_cmd, 0.0);
		FVector qp(PQR.Y, PQR.X, 0.0);
		//UE_LOG(LogTemp, Display, TEXT("PQR %f %f %f"), PQR.X, PQR.Y, PQR.Z);
		FVector Kp(pid.Kpp, pid.Kpr, 0.0);
		FVector Kd(pid.Kdp, pid.Kdr, 0.0);

		FVector pr_error = pr_ref - pr;
		//UE_LOG(LogTemp, Display, TEXT("PR error %f %f"), pr_error.X, pr_error.Y);
		FVector prev_int = pid.pr_error_int;
		pid.pr_error_int += pr_error * MjStep;
		pid.pr_error_int -= prev_int * pid.AntiWU;

		//FVector error_dt = (pr_error - pid.prev_pr_error) / MjStep;
		pid.prev_pr_error = pr_error;

		FVector tau_pr = Kp * pr_error + pid.I_pr * pid.pr_error_int - Kd * qp;
		//FVector tau_pr = Kp * pr_error + pid.I_pr * pid.pr_error_int + Kd * error_dt;
		pitch_tau = tau_pr.X;
		roll_tau = tau_pr.Y;

		//pitch_tau = pitch_cmd;
		//roll_tau = roll_cmd;
	}
	else
	{
		pitch_tau = pitch_cmd;
		roll_tau = roll_cmd;
	}
	//UE_LOG(LogTemp, Display, TEXT("Roll tau Pitch tau raw %f %f "), roll_tau, pitch_tau);

	float max_rps = MaxRPM / 60;
	float min_rps = MinRPM / 60;

	float up_limit = 2 * max_rps * max_rps * Kt - 2 * min_rps * min_rps * Kt;
	float low_limit = -up_limit;

	pitch_tau = fmax(fmin(pitch_tau, up_limit), low_limit);
	roll_tau = fmax(fmin(roll_tau, up_limit), low_limit);

	//UE_LOG(LogTemp, Display, TEXT("Roll tau Pitch tau %f %f "), roll_tau, pitch_tau);
}

void AMultiRotorDrone::Mixer(float& tot_thrust, float& roll_tau, float& pitch_tau, float& yaw_tau)
{
	//UE_LOG(LogTemp, Display, TEXT("Total thrust %f"), tot_thrust);

	float xf = fabs(Rotors.p1.X);
	float yf = fabs(Rotors.p1.Y);
	float xr = fabs(Rotors.p2.X);
	float yr = fabs(Rotors.p2.Y);

	//float w1sq = pitch_tau / (2 * Kt * (xf + xr)) - roll_tau / (2 * Kt * (yf + yr)) + (tot_thrust * xr) / (2 * (xf + xr)) + (yaw_tau * yr) / (2 * Km * (yf + yr));
	float w1sq = pitch_tau / (2 * Kt * (xf + xr)) - roll_tau / (2 * Kt * (yf + yr)) + (tot_thrust * xr) / (2 * Kt * (xf + xr)) + (yaw_tau * yr) / (2 * Km * (yf + yr));
	//float w2sq = roll_tau / (2 * Kt * (yf + yr)) - pitch_tau / (2 * Kt * (xf + xr)) + (tot_thrust * xf) / (2 * (xf + xr)) + (yaw_tau * yf) / (2 * Km * (yf + yr));
	float w2sq = roll_tau / (2 * Kt * (yf + yr)) - pitch_tau / (2 * Kt * (xf + xr)) + (tot_thrust * xf) / (2 * Kt * (xf + xr)) + (yaw_tau * yf) / (2 * Km * (yf + yr));
	//float w3sq = pitch_tau / (2 * Kt * (xf + xr)) + roll_tau / (2 * Kt * (yf + yr)) + (tot_thrust * xr) / (2 * (xf + xr)) - (yaw_tau * yr) / (2 * Km * (yf + yr));
	float w3sq = pitch_tau / (2 * Kt * (xf + xr)) + roll_tau / (2 * Kt * (yf + yr)) + (tot_thrust * xr) / (2 * Kt * (xf + xr)) - (yaw_tau * yr) / (2 * Km * (yf + yr));
	//float w4sq = (tot_thrust * xf) / (2 * (xf + xr)) - roll_tau / (2 * Kt * (yf + yr)) - pitch_tau / (2 * Kt * (xf + xr)) - (yaw_tau * yf) / (2 * Km * (yf + yr));
	float w4sq = (tot_thrust * xf) / (2 * Kt * (xf + xr)) - roll_tau / (2 * Kt * (yf + yr)) - pitch_tau / (2 * Kt * (xf + xr)) - (yaw_tau * yf) / (2 * Km * (yf + yr));
	//UE_LOG(LogTemp, Display, TEXT("Computed omegas %f %f %f %f"), w1sq, w2sq, w3sq, w4sq);

	ComputedRPM[0] = w1sq;
	ComputedRPM[1] = w2sq;
	ComputedRPM[2] = w3sq;
	ComputedRPM[3] = w4sq;

	float min_level = (0.05 * Mass * GRAVITY) / Kt;
	float max_level = (MaxRPM / 60) * (MaxRPM / 60);
	float min_wsq = ComputedRPM[0];
	for (int32 i = 1; i < 4; i++)
	{
		//if (ComputedRPM[i] < min_wsq) min_wsq = ComputedRPM[i];
		if (ComputedRPM[i] < 0) ComputedRPM[i] = 0;
	}

	//if (min_wsq < min_level)
	//{
		//float diff = min_level - min_wsq;
		//for (int32 i = 0; i < 4; i++)
		//{
			//ComputedRPM[i] += diff;
		//}
	//}

	//for (int32 i = 0; i < 4; i++)
	//{
		//if (ComputedRPM[i] > max_level) ComputedRPM[i] = max_level;
	//}

	if (w1sq >= 0.1) RPM[0] = sqrtf(w1sq);
	if (w2sq > 0.1) RPM[1] = sqrtf(w2sq);
	if (w3sq > 0.1) RPM[2] = -sqrtf(w3sq);
	if (w4sq > 0.1) RPM[3] = -sqrtf(w4sq);

	ComputeForcesAndMoments();
}

void AMultiRotorDrone::ComputeForcesAndMoments()
{

	for (int32 i = 0; i < 4; i++)
	{
		float rps_sq = ComputedRPM[i];
		float rps = sqrtf(rps_sq);
		FVector prop_loc;
		switch (i)
		{
		case 0:
			prop_loc = Rotors.p1;
			break;
		case 1:
			prop_loc = Rotors.p2;
			break;
		case 2:
			prop_loc = Rotors.p3;
			break;
		case 3:
			prop_loc = Rotors.p4;
			break;
		}

		//float J = GetJ(prop_loc, rps, Rotors.Diameter);
		//float Ct = Rotors.get_value(Rotors.cThrust, J);
		//float Cq = Rotors.get_value(Rotors.cTorque, J);
		//UE_LOG(LogTemp, Display, TEXT("Prop %d J Ct Cq %f %f %f"), i+1, J, Ct, Cq);
		//Kt = AIR_DENCITY * Ct * powf(Rotors.Diameter, 4);
		//Km = AIR_DENCITY * Cq * powf(Rotors.Diameter, 5);

		float force = Kt * rps_sq;
		ForcesToApply[i] = force;
		MomentsToApply[i] = Km * rps_sq;
		if (i == 0 || i == 1)
		{
			MomentsToApply[i] *= -1;
		}
	}
	UE_LOG(LogTemp, Display, TEXT("Forces %f %f %f %f"), ForcesToApply[0], ForcesToApply[1], ForcesToApply[2], ForcesToApply[3]);
	UE_LOG(LogTemp, Display, TEXT("Moments %f %f %f %f"), MomentsToApply[0], MomentsToApply[1], MomentsToApply[2], MomentsToApply[3]);
}

void AMultiRotorDrone::ComputedDragForce()
{
	FQuaternion Orient;
	Orient.InitializeFromEulerAngles(State.Orientation.X, State.Orientation.Y, State.Orientation.Z);
	Matrix33 Tl2b = Orient.GetMatrix().Transposed();
	FVector Vb = Tl2b * State.Velocity;
	double Qbar = 0.5 * AIR_DENCITY * (Vb.X * Vb.X + Vb.Y * Vb.Y + Vb.Z * Vb.Z);

	double QbarX = 0.5 * AIR_DENCITY * Vb.X * Vb.X;
	double QbarY = 0.5 * AIR_DENCITY * Vb.Y * Vb.Y;
	double QbarZ = 0.5 * AIR_DENCITY * Vb.Z * Vb.Z;

	//Vb.Normalize();

	int sign = Vb.X >= 0 ? -1 : 1;
	ComputedDrag[0] = sign * QbarX * SrefX * CdXY;

	sign = Vb.Y >= 0 ? -1 : 1;
	ComputedDrag[1] = sign * QbarY * SrefY * CdXY;

	sign = Vb.Z >= 0 ? 1 : -1;
	ComputedDrag[2] = sign * QbarZ * SrefZ * SrefZ * CdZ;

	FVector drag(ComputedDrag[0], ComputedDrag[1], -ComputedDrag[2]);
	FVector wdrag = Tl2b.Transposed() * drag;
	//UE_LOG(LogTemp, Display, TEXT("Computed Drag %f %f %f"), wdrag[0], wdrag[1], wdrag[2]);
}

void AMultiRotorDrone::ControllerUpdate()
{
	float total_thrust;
	if (LandingStarted)
	{
		Landing();
	}
	float target_height = TargetPosition.Z;
	AltitudeController(target_height, total_thrust);

	float target_yaw = TargetYaw;
	//float tau_yaw;
	float yaw_cmd;
	YawController(target_yaw, yaw_cmd);

	float roll_cmd = 0;
	float pitch_cmd = 0;
	float xref = TargetPosition.X;
	float yref = TargetPosition.Y;
	RollPitchCmdController(xref, yref, roll_cmd, pitch_cmd);
	//UE_LOG(LogTemp, Display, TEXT("Roll Pitch commands %f %f"), target_roll, target_pitch);
	float tau_roll, tau_pitch;
	AttitudeController(roll_cmd, pitch_cmd, tau_roll, tau_pitch);
	//BodyRateController(roll_cmd, pitch_cmd, yaw_cmd, tau_roll, tau_pitch, tau_yaw);

	Mixer(total_thrust, tau_roll, tau_pitch, yaw_cmd);

	//UE_LOG(LogTemp, Display, TEXT("Commanded Position %f %f %f"), TargetPosition.X, TargetPosition.Y, TargetPosition.Z);
	//UE_LOG(LogTemp, Display, TEXT("Current Position %f %f %f"), State.Position.X, State.Position.Y, State.Position.Z);

	//ComputedDragForce();

	//UE_LOG(LogTemp, Display, TEXT("Motors commands %f %f %f %f"), ComputedRPM[0], ComputedRPM[1], ComputedRPM[2], ComputedRPM[3]);
}

/*
void AMultiRotorDrone::CorrectForces()
{
	float minForce = ForcesToApply[0];

	for (int32 i = 1; i < 4; i++)
	{
		if (ForcesToApply[i] < 0)
		{
			minForce = ForcesToApply[i];
			//ForcesToApply[i] = 0.0;
		}
	}
	//float min_per_motor = minThrust * 0.25;

	float min_per_motor = 0;
	if (minForce < min_per_motor)
	{
		float diff = min_per_motor - minForce;
		for (int32 i = 0; i < 4; i++)
		{
			ForcesToApply[i] += diff;
		}
	}
}
*/

void AMultiRotorDrone::Landing()
{
	static bool initialized = false;
	static float prev_height = 0;
	static int height_counter = 0;
	if (!initialized)
	{
		GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::White, TEXT("LANDING STARTED"));
		prev_height = State.Position.Z;
		initialized = true;
	}

	float height = State.Position.Z;
	if (fabsf(height - prev_height) <= 0.001)
	{
		height_counter += 1;
		if (height_counter >= 1000)
		{
			GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::Green, TEXT("LANDING COMPLETED"));
			height_counter = 0;
			LandingStarted = false;
			initialized = false;
			TakeOffStarted = false;
			OnGround = true;
		}
	}

	float landing_speed = 3;
	if (TargetPosition.Z < 0)
	{
		TargetPosition.Z += landing_speed * MjStep;
	}

	prev_height = height;
}

//CONTROLLER 2

void AMultiRotorDrone::LateralController(float& x, float& dx, float& ddx, float& y, float& dy, float& ddy, float& c, float& bx_c, float& by_c)
{
	float dt = UGameplayStatics::GetWorldDeltaSeconds(this);

	float X = State.Position.X;
	float Y = State.Position.Y;
	float dX = State.Velocity.X;
	float dY = State.Velocity.Y;

	FQuaternion yaw_quat;
	yaw_quat.InitializeFromEulerAngles(0, 0, State.Orientation.Z);
	Matrix33 yaw_rot = yaw_quat.GetMatrix();

	float x_err = x - X;
	float dx_ff = x_err / dt;
	float ddx_ff = dx_ff / dt;


	float dx_err = dx_ff - dX;

	float y_err = y - Y;
	float dy_ff = y_err / dt;
	float ddy_ff = dy_ff / dt;

	float dy_err = dy_ff - dY;

	//FVector pos_error(x_err, y_err, 0);
	//FVector vel_error(dx_err, dy_err, 0);
	//pos_error = yaw_rot * pos_error;
	//vel_error = yaw_rot * vel_error;
	//x_err = pos_error.X;
	//y_err = pos_error.Y;
	//dx_err = vel_error.X;
	//dy_err = vel_error.Y;

	float ddx_cmd = pid2.xKp * x_err + pid2.xKd * dx_ff + ddx_ff;
	ddx_cmd = fmax(fmin(ddx_cmd, 5), -5);
	float ddy_cmd = pid2.yKp * y_err + pid2.yKd * dy_ff + ddy_ff;
	ddy_cmd = fmax(fmin(ddy_cmd, 5), -5);

	//float z_thrust = c * Mass;
	//float x_thrust = fabs(ddx_cmd * Mass);
	//float y_thrust = fabs(ddy_cmd * Mass);

	//FVector tot_thrust(x_thrust, y_thrust, z_thrust);
	//double thrust = tot_thrust.Size();
	//if (thrust > maxThrustPerMotor * 4)
	//{
		//tot_thrust.Normalize();
		//tot_thrust *= (maxThrustPerMotor * 4);
		//float thrust_margin = maxThrustPerMotor * 4 - z_thrust;
		//x_thrust = thrust * (x_thrust / tot_thrust.Size2D());
		//y_thrust = thrust_margin * (y_thrust / (x_thrust + y_thrust));
		//ddx_cmd = tot_thrust.X / Mass;
		//ddy_cmd = tot_thrust.Y / Mass;
		//c = tot_thrust.Z / Mass;
	//}

	//int sign_x = ddx_cmd >= 0 ? 1 : -1;
	//int sign_y = ddy_cmd >= 0 ? 1 : -1;

	//ddx_cmd = (x_thrust * sign_x) / Mass;
	//ddy_cmd = (y_thrust * sign_y) / Mass;

	UE_LOG(LogTemp, Display, TEXT("X Y accel commands: %f %f"), ddx_cmd, ddy_cmd);

	bx_c = -ddx_cmd / c;
	by_c = -ddy_cmd / c;
}

void AMultiRotorDrone::RollPitchController(float& bx_c, float& by_c, float& p_c, float& q_c)
{
	FVector euler = State.Orientation;
	UE_LOG(LogTemp, Display, TEXT("Carrent orientation: %f %f %f"), euler.X, euler.Y, euler.Z);
	FQuaternion quat;
	quat.InitializeFromEulerAngles(euler.X, euler.Y, euler.Z);
	Matrix33 mat = quat.GetMatrix();
	float bx = mat(0, 2);
	float by = mat(1, 2);

	float bx_dot = pid2.Kp_roll * (bx_c - bx);
	float by_dot = pid2.Kp_pitch * (by_c - by);

	Matrix33 mat2(mat(1, 0), -mat(0, 0), 0.0, mat(1, 1), -mat(0, 1), 0.0, 0.0, 0.0, 0.0);
	mat2 = mat2 * (1 / mat(2, 2));
	FVector b(bx_dot, by_dot, 0.0);
	FVector pq_c = mat2 * b;
	float max_tilt = PI / 4;
	p_c = pq_c.X;
	q_c = pq_c.Y;

	UE_LOG(LogTemp, Display, TEXT("Commanded p q : %f %f"), p_c, q_c);
}

void AMultiRotorDrone::BodyRateController(float& p_c, float& q_c, float& r_c, float& u_bar_p, float& u_bar_q, float& u_bar_r)
{
	float p = PQR.X;
	float q = PQR.Y;
	float r = PQR.Z;

	u_bar_p = pid2.Kp_p * (p_c - p);
	u_bar_q = pid2.Kp_q * (q_c - q);
	//u_bar_r = pid2.Kp_r * (r_c - r);

	UE_LOG(LogTemp, Display, TEXT("Commanded body rates : %f %f %f"), u_bar_p, u_bar_q, u_bar_r);

	float max_rps = MaxRPM / 60;
	float min_rps = MinRPM / 60;
	float pq_up_limit = 2 * max_rps * max_rps * Kt - 2 * min_rps * min_rps * Kt;
	float pq_low_limit = -pq_up_limit;

	float r_up_limit = 2 * max_rps * max_rps * Km - 2 * min_rps * min_rps * Km;
	float r_low_limit = -r_up_limit;

	u_bar_p = fmax(pq_low_limit, fmin(u_bar_p, pq_up_limit));
	u_bar_q = fmax(pq_low_limit, fmin(u_bar_p, pq_up_limit));
	//u_bar_r = fmax(r_low_limit, fmin(u_bar_p, r_up_limit));
}

void AMultiRotorDrone::HeadingController(float& yaw_c, float& yaw_tau)
{
	float yaw = State.Orientation.Z;

	yaw_tau = pid2.Kp_yaw * (yaw_c - yaw);
}

void AMultiRotorDrone::HeightController(float& z_t, float& dz_t, float& ddz_t, float& c)
{
	float Z = State.Position.Z;
	float dZ = State.Velocity.Z;
	float dt = UGameplayStatics::GetWorldDeltaSeconds(this);
	float z_error = z_t - Z;
	float dz_ff = z_error / dt;
	float ddz_ff = dz_ff / dt;

	float u_1_bar = pid2.zKp * z_error + pid2.zKd * (dz_ff - dZ) + ddz_ff;

	FQuaternion quat;
	quat.InitializeFromEulerAngles(State.Orientation.X, State.Orientation.Y, State.Orientation.Z);
	Matrix33 mat = quat.GetMatrix();

	float b_z = mat(2, 2);
	c = (-u_1_bar - GRAVITY) / b_z;
	if ((c + GRAVITY) * Mass > maxThrustPerMotor * 4)
	{
		c = ((maxThrustPerMotor * 4) / Mass) - GRAVITY;
	}

	//UE_LOG(LogTemp, Display, TEXT("Target thrust vs maxThrust %f %f"), c * Mass, maxThrustPerMotor * 4);
}

void AMultiRotorDrone::Controller2Update()
{
	UE_LOG(LogTemp, Display, TEXT("Target Height %f"), TargetPosition.Z);
	float dt = UGameplayStatics::GetWorldDeltaSeconds(this);

	float c = 0;
	float target_height = TargetPosition.Z;
	float dz_ff = MaxVelocity;
	float ddz_ff = MaxVelocity / dt;
	HeightController(target_height, dz_ff, ddz_ff, c);

	float bx_c = 0;
	float by_c = 0;
	float target_x = TargetPosition.X;
	float dx_ff = MaxVelocity;
	float ddx_ff = MaxVelocity / dt;
	float target_y = TargetPosition.Y;
	float dy_ff = MaxVelocity;
	float ddy_ff = MaxVelocity / dt;
	LateralController(target_x, dx_ff, ddx_ff, target_y, dy_ff, ddy_ff, c, bx_c, by_c);

	float p_c, q_c;
	RollPitchController(bx_c, by_c, p_c, q_c);

	float yaw_target = TargetYaw;
	float r_c;
	HeadingController(yaw_target, r_c);

	float u_bar_p, u_bar_q, u_bar_r;
	BodyRateController(p_c, q_c, r_c, u_bar_p, u_bar_q, u_bar_r);

	//float omega_c = c / Kt;
	//float omega_p = u_bar_p / Kt;
	//float omega_q = u_bar_q / Kt;
	//float omega_r = u_bar_r / Km;

	ComputedRPM[0] = (0.25 * (c - u_bar_p + u_bar_q + u_bar_r)) / Kt;
	ComputedRPM[1] = (0.25 * (c + u_bar_p - u_bar_q + u_bar_r)) / Kt;
	ComputedRPM[2] = (0.25 * (c + u_bar_p + u_bar_q - u_bar_r)) / Kt;
	ComputedRPM[3] = (0.25 * (c - u_bar_p - u_bar_q - u_bar_r)) / Kt;

	//float w1sq = 0.25 * (omega_c - omega_p + omega_q + omega_r);
	//float w2sq = 0.25 * (omega_c + omega_p - omega_q + omega_r);
	//float w3sq = 0.25 * (omega_c + omega_p + omega_q - omega_r);
	//float w4sq = 0.25 * (omega_c - omega_p - omega_q - omega_r);

	//ComputedRPM[0] = w1sq;
	//ComputedRPM[1] = w2sq;
	//ComputedRPM[2] = w3sq;
	//ComputedRPM[3] = w4sq;
	/*
	float min_rps_sq_value = 0;
	float min_rps_sq = min_rps_sq_value;
	for (int32 i = 0; i < 4; i++)
	{
		if (ComputedRPM[i] < min_rps_sq)
		{
			min_rps_sq = ComputedRPM[i];
		}
	}
	float diff = min_rps_sq_value - min_rps_sq;
	for (int32 i = 0; i < 4; i++)
	{
		ComputedRPM[i] += diff;
	}
	*/
	//UE_LOG(LogTemp, Display, TEXT("Motors commands %f %f %f %f"), ComputedRPM[0], ComputedRPM[1], ComputedRPM[2], ComputedRPM[3]);

	ComputeForcesAndMoments();
}