#include "attitude.h"
#include "pid.h"
#include <math.h>

#define MAHONY_KP	2.5f
#define MAHONY_KI	0.01f

static Vector3 gyro_bias = {};

volatile uint8_t imu_update_flag = 0;
volatile uint8_t pid_update_flag = 0;

float angle_kp_pitch 	= 2.5f;
float angle_kp_roll 	= 2.5f;
float angle_kp_yaw 		= 2.0f;

pid_t pid_pitch = {.kp = 0.6f, .ki = 0.0f, .kd = 0.025f};
pid_t pid_roll 	= {.kp = 0.6f, .ki = 0.0f, .kd = 0.025f};
pid_t pid_yaw 	= {.kp = 0.6f, .ki = 0.0f, .kd = 0.0f};

/*int16_t desired_roll = 0;
int16_t desired_pitch = 0;

float desired_roll_rate = 0.0f;
float desired_pitch_rate = 0.0f;

float roll_rate_error = 0.0f;
float pitch_rate_error = 0.0f;
float yaw_rate_error = 0.0f;

float error_roll = 0.0f;
float error_pitch = 0.0f;
float error_yaw = 0.0f;*/

void attitude_init(attitude_t* att)
{
	att->orientation.w = 1.0f;
	att->orientation.x = 0.0f;
	att->orientation.y = 0.0f;
	att->orientation.z = 0.0f;
	
	att->gyro.x = 0.0f;
	att->gyro.y = 0.0f;
	att->gyro.z = 0.0f;
}

void attitude_estimator_update(attitude_t* att, const imu_scaled_t* imu)
{
	Vector3 gyro 	= {imu->gx, imu->gy, imu->gz};
	Vector3 accel 	= {imu->ax, imu->ay, imu->az};
	
	float norm = sqrtf(accel.x*accel.x + accel.y*accel.y + accel.z*accel.z);
	
	if (norm > 0.00000000001f)
	{
		accel.x /= norm;
		accel.y /= norm;
		accel.z /= norm;
		
		Vector3 g = gravity_from_quat(&att->orientation);
		
		Vector3 error = 
		{
			accel.y*g.z - accel.z*g.y,
			accel.z*g.x - accel.x*g.z,
			accel.x*g.y - accel.y*g.x
		};
		
		gyro_bias.x += MAHONY_KI * error.x * IMU_DT;
		gyro_bias.y += MAHONY_KI * error.y * IMU_DT;
		gyro_bias.z += MAHONY_KI * error.z * IMU_DT;
		
		gyro_bias.x = constrain(gyro_bias.x, -GYRO_BIAS_LIM, GYRO_BIAS_LIM);
		gyro_bias.y = constrain(gyro_bias.y, -GYRO_BIAS_LIM, GYRO_BIAS_LIM);
		gyro_bias.z = constrain(gyro_bias.z, -GYRO_BIAS_LIM, GYRO_BIAS_LIM);
		
		gyro.x += MAHONY_KP * error.x + gyro_bias.x;
		gyro.y += MAHONY_KP * error.y + gyro_bias.y;
		gyro.z += MAHONY_KP * error.z + gyro_bias.z;
	}
	
	att->gyro = gyro;
	
	Quaternion h = {gyro.x * DEG2RAD, gyro.y * DEG2RAD, gyro.z * DEG2RAD, 0};
	
	Quaternion dq = QuatProd(&att->orientation, &h);
	dq = QuatConstProd(&dq, 0.5f * IMU_DT);
	att->orientation = QuatSum(&att->orientation, &dq);
	
	att->orientation = QuatNormalize(&att->orientation, 1.0f);
}

void attitude_controller_update(control_channels_t* control, const rc_channels* rx, const attitude_t* att)
{
	Quaternion q_target = rx_to_quaternion(rx);
	Quaternion q_error 	= QuatGetError(&att->orientation, &q_target, true);
	
	Vector3 angle_error = 
	{
		2.0f * q_error.x,
		2.0f * q_error.y,
		2.0f * q_error.z
	};
	
	Vector3 desired_rate = 
	{
		angle_error.x * angle_kp_pitch,
		angle_error.y * angle_kp_roll,
		angle_error.z * angle_kp_yaw
	};
	
	desired_rate.x = constrain(desired_rate.x, -RATE_LIM, RATE_LIM);
	desired_rate.y = constrain(desired_rate.y, -RATE_LIM, RATE_LIM);
	desired_rate.z = constrain(desired_rate.z, -RATE_LIM, RATE_LIM);
	
	float pitch_error 	= desired_rate.x - att->gyro.x;
	float roll_error 	= desired_rate.y - att->gyro.y;
	float yaw_error 	= desired_rate.z - att->gyro.z;
	
	control->pitch 	= pid_update(&pid_pitch, pitch_error, att->gyro.x, IMU_DT);
	control->roll 	= pid_update(&pid_roll, roll_error, att->gyro.y, IMU_DT);
	control->yaw 	= pid_update(&pid_yaw, yaw_error, att->gyro.z, IMU_DT);
}

Quaternion rx_to_quaternion(const rc_channels* rx)
{
	float pitch = int_mapping(rx->rc_pitch, -500, 500, -45, 45) * DEG2RAD;
	float roll 	= int_mapping(rx->rc_roll, -500, 500, -45, 45) * DEG2RAD;
	float yaw 	= 0;
	
	Vector3 pry = {pitch, roll, yaw};
	
	return QuatCreateFromEuler(&pry);
}

Vector3 gravity_from_quat(const Quaternion* q)
{
	Vector3 g;
	
	g.x = 2 * (q->x*q->z - q->w*q->y);
	g.y = 2 * (q->w*q->x + q->y*q->z);
	g.z = q->w*q->w - q->x*q->x - q->y*q->y + q->z*q->z;
	
	return g;
}

void TIM6_DAC_IRQHandler()
{
	if (TIM6->SR & TIM_SR_UIF)
	{
		TIM6->SR &= ~TIM_SR_UIF;
		imu_update_flag = 1;
		pid_update_flag = 1;
	}
}

void attitude_update(attitude_t* attitude, imu_scaled_t* imu)
{
	if (!imu_update_flag)
		return;
	
	imu_update_flag = 0;
	
	imu_read_scaled(imu);
	attitude_estimator_update(attitude, imu);
	
	// complementary_filter_update(attitude, imu);
	// yaw_rate_update(attitude, imu);
}

void attitude_pid_update(control_channels_t* control, const rc_channels* rx, const attitude_t* att)
{
	if (!pid_update_flag)
		return;
	pid_update_flag = 0;
	attitude_controller_update(control, rx, att);
}

float constrain(float x, float min, float max)
{
	if (x < min) x = min; else if (x > max) x = max;
	return x;
}

/*void complementary_filter_update(attitude_t* att, const imu_scaled_t* imu)
{
	static float roll_acc;
	static float pitch_acc;
	
	roll_acc = accel_roll_deg(imu);
	pitch_acc = accel_pitch_deg(imu);
	
	integrate_gyro_roll_deg(&att->roll, imu);
	integrate_gyro_pitch_deg(&att->pitch, imu);
	
	att->roll = CF_ALPHA * att->roll + (1 - CF_ALPHA) * roll_acc;
	att->pitch = CF_ALPHA * att->pitch + (1 - CF_ALPHA) * pitch_acc;
}*/

/*void yaw_rate_update(attitude_t* attitude, imu_scaled_t* imu)
{
	attitude->yaw_rate = imu->gz;
}

void attitude_pid_update(control_channels_t* control, 
                         const rc_channels* rx, 
                         const attitude_t* att,
                         const imu_scaled_t* imu)
{
	if (pid_update_flag)
	{
		pid_update_flag = 0;
		
		desired_roll  = int_mapping(rx->rc_roll, -500, 500, -45, 45);
		desired_pitch = int_mapping(rx->rc_pitch, -500, 500, -45, 45);
		
		desired_roll_rate  = angle_kp_roll  * (desired_roll - att->roll);
		desired_pitch_rate = angle_kp_pitch * (desired_pitch - att->pitch);
		
		if (desired_roll_rate > 200) desired_roll_rate = 200;
		if (desired_roll_rate < -200) desired_roll_rate = -200;
		
		if (desired_pitch_rate > 200)  desired_pitch_rate = 200;
		if (desired_pitch_rate < -200) desired_pitch_rate = -200;
		
		roll_rate_error = desired_roll_rate - imu->gy;
		pitch_rate_error = desired_pitch_rate - imu->gx;
		yaw_rate_error = - imu->gz;
		
		control->roll  = pid_update(&pid_roll, roll_rate_error, imu->gy, IMU_DT);
		control->pitch = pid_update(&pid_pitch, pitch_rate_error, imu->gx, IMU_DT);
		control->yaw   = pid_update(&pid_yaw, yaw_rate_error, imu->gz, IMU_DT);
	}
}*/

float accel_roll_deg(const imu_scaled_t* imu) {
	// right-left
	return atan2f(imu->ax, sqrtf(imu->ay * imu->ay + imu->az * imu->az)) * 180.0f / PI;
}

float accel_pitch_deg(const imu_scaled_t* imu)
{
	// forward-backward
	return atan2f(-imu->ay, sqrtf(imu->ax * imu->ax + imu->az * imu->az)) * 180.0f / PI;
}

void integrate_gyro_roll_deg(float* roll, const imu_scaled_t* imu)
{
	// right-left
	*roll += imu->gy * IMU_DT;
}

void integrate_gyro_pitch_deg(float* pitch, const imu_scaled_t* imu)
{
	// forward-backward
	*pitch += imu->gx * IMU_DT;
}

void integrate_gyro_yaw_deg(float* yaw, const imu_scaled_t* imu)
{
	// forward-backward
	*yaw += imu->gz * IMU_DT;
}