#include "irs.h"
#include <math.h>

#define MAHONY_KP 		2.5f
#define MAHONY_KI 		0.01f

#define ACC_ALPHA		0.01f
#define ACC_MAX_ERROR	0.2f

static Vector3 gyro_bias = {0};

void IRS_init(IRS* irs)
{
    irs->q.w = 1.0f;
    irs->q.x = 0.0f;
    irs->q.y = 0.0f;
    irs->q.z = 0.0f;
}


void IRS_update(IRS* irs, const Vector3* gyro_in, const Vector3* accel_in, float dt)
{
	Quaternion qConjugate = QuatConjugate(&irs->q);
	
	// gyro update
  
	irs->gyro = *gyro_in;
	
	Vector3 gyro = 
	{
		gyro_in->x * dt * DEG2RAD,
		gyro_in->y * dt * DEG2RAD,
		gyro_in->z * dt * DEG2RAD
	};
	
	Quaternion g = {gyro.x, gyro.y, gyro.z, 0};
	g = QuatProd(&irs->q, &g);
	g = QuatConstProd(&g, 0.5f);
	irs->q = QuatSum(&irs->q, &g);
	irs->q = QuatNormalize(&irs->q, 1.0f);
	
	Quaternion q0010 = {0.0f, 0.0f, 1.0f, 0.0f};
	Quaternion conj0010prod = QuatProd(&qConjugate, &q0010);
	Quaternion prtilts = QuatProd(&conj0010prod, &irs->q);
	
	irs->oriPRT.x = prtilts.x;
	irs->oriPRT.y = prtilts.y;
	irs->oriPRT.z = prtilts.z;
	
    // /gyro update
	
	// accel update
	
	Vector3 accel = {accel_in->x, accel_in->y, -accel_in->z};
	
	irs->accel = accel;
	
	restoreQuat(irs, &accel);
	
	
	
	
	// /accel update
	
	
	
	/*Vector3 accel = {accel_in->x, accel_in->y, -accel_in->z};
	
	float acc_len = sqrtf(accel.x*accel.x + accel.y*accel.y + accel.z*accel.z);
	
	if (acc_len > 1e-6f)
	{
		Vector3 acc_norm =
		{
			accel.x / acc_len,
			accel.y / acc_len,
			accel.z / acc_len
		};
		
		float dyn = fabsf(acc_len - 1.0f);
		
		float trust;
		if (dyn > ACC_MAX_ERROR)
			trust = 0.0f;
		else
			trust = 1.0f - (dyn / ACC_MAX_ERROR);
		
		float gain = ACC_ALPHA * trust;
		
		if (gain > 1e-5f)
		{
			Vector3 est = IRS_getGravity(&q);
			
			Vector3 error = 
			{
				acc_norm.y*est.z - acc_norm.z*est.y,
				acc_norm.z*est.x - acc_norm.x*est.z,
				acc_norm.x*est.y - acc_norm.y*est.x
			};
			
			Quaternion corretrion = 
			{
				error.x * gain * 0.5f,
				error.y * gain * 0.5f,
				error.z * gain * 0.5f,
				1.0f
			};
			
			q = QuatProd(&corretrion, &q);
			q = QuatNormalize(&q, 1.0f);
		}
	}
	
	irs->q = q;
	irs->gyro = gyro;
	irs->accel = accel;*/
}

void restoreQuat(IRS* irs, const Vector3* accel)
{
	float len = lengthV3(accel);
	
	static float quat_acc_alpha = 0.03f;
	static float quat_acc_max = 0.02f;
	
	float dyn = fabsf(len - 1.0f);
	if (dyn > quat_acc_max) dyn = 1.0f; else dyn /= quat_acc_max;
	
	float gain = quat_acc_alpha * (1.0f - dyn);
	
	if (gain < 0.0001f) return;
	
	Vector3 acc = normalizeV3(accel, 1.0f);
	
	Vector3 est;
	est.x = 2.0f * (irs->q.x * irs->q.z - irs->q.w * irs->q.y);
	est.y = 2.0f * (irs->q.w * irs->q.x - irs->q.y * irs->q.z);
	est.z = irs->q.w * irs->q.w - irs->q.x * irs->q.x - irs->q.y * irs->q.y + irs->q.z * irs->q.z;
	
	Vector3 cross = Cross(&acc, &est);
	float dot = DotV3(&acc, &est);
	
	if (dot < 0.0f)
	{
		float error_len = lengthV3(&cross);
		if (error_len < 0.001f) {
			cross.x = 1.0f;
			cross.y = 0.0f;
			cross.z = 0.0f;
		}
		else
		{
			cross = constProdV3(&cross, 1.0f / error_len);
		}
	}
	
	Vector3 axis = constProdV3(&cross, gain * 0.5f);
	
	Quaternion correction =
	{
		axis.x,
		axis.y,
		axis.z,
		1.0f
	};
	
	irs->q = QuatProd(&irs->q, &correction);
	irs->q = QuatNormalize(&irs->q, 1.0f);
}

void setAccelShift(IRS* irs, const float roll, const float pitch, const float yaw)
{
	float h_roll = (roll * DEG2RAD) / 2.0f;
	float h_pitch = (pitch * DEG2RAD) / 2.0f;
	float h_yaw = (yaw * DEG2RAD) / 2.0f;
	
	Quaternion q_roll = {0.0f, sinf(h_roll), 0.0f, cosf(h_roll)};
	Quaternion q_pitch = {sinf(h_pitch), 0.0f, 0.0f, cosf(h_pitch)};
	Quaternion q_yaw = {0.0f, 0.0f, sinf(h_yaw), cosf(h_yaw)};
	
	Quaternion prProd = QuatProd(&q_pitch, &q_roll);
	Quaternion pryProd = QuatProd(&prProd, &q_yaw);
	
	irs->shiftAccelPRY = pryProd;
}

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

    return g;
}