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

static const float PI = 3.14159265359f;
static const float TO_DEG = 180.0f/PI;
static const float TO_RAD = PI/180.0f;

struct Quaternion 
{
  float w, x, y, z;
};

static Quaternion qCurrent = { 1, 0, 0, 0 };

static const float period = 0.01f; // 100Hz

static bool isFirst = true;

inline void vecNormalize(Vec3& v)
{
  float n = sqrtf(v.x * v.x + v.y * v.y + v.z * v.z);
  if (n > 1e-9f) 
  {
    v.x /= n;
    v.y /= n;
    v.z /= n;
  }
}

inline Vec3 vecCross(const Vec3& a, const Vec3& b)
{
  return 
  {
    a.y * b.z - a.z * b.y,
    a.z * b.x - a.x * b.z,
    a.x * b.y - a.y * b.x
  };
}

inline void normalizeQuaternion(Quaternion& q)
{
  float norm = sqrtf(q.w * q.w + q.x * q.x + q.y * q.y + q.z * q.z);
  if (norm > 1e-12f) 
  {
    q.w /= norm;
    q.x /= norm;
    q.y /= norm;
    q.z /= norm;
  }
}

static Quaternion quaternionMultiply(const Quaternion& q1, const Quaternion& q2)
{
  Quaternion r;
  r.w = q1.w * q2.w - q1.x * q2.x - q1.y * q2.y - q1.z * q2.z;
  r.x = q1.w * q2.x + q1.x * q2.w + q1.y * q2.z - q1.z * q2.y;
  r.y = q1.w * q2.y - q1.x * q2.z + q1.y * q2.w + q1.z * q2.x;
  r.z = q1.w * q2.z + q1.x * q2.y - q1.y * q2.x + q1.z * q2.w;
  return r;
}

static Quaternion rotateVectorByQuaternion(const Quaternion& q, const Vec3& In)
{
  Quaternion r = quaternionMultiply(quaternionMultiply(q, Quaternion{ 0, In.x, In.y, In.z }), Quaternion{ q.w, -q.x, -q.y, -q.z });
  
  return r;
}

static Vec3 backRotateVectorByQuaternion(const Quaternion& q, const Vec3& In)
{
  Quaternion r = quaternionMultiply(quaternionMultiply(Quaternion{q.w, -q.x, -q.y, -q.z}, Quaternion{ 0, In.x, In.y, In.z }), q);
  
  return { r.x, r.y, r.z };
}

static Quaternion backRotateVectorByQuaternion2(const Quaternion& q, const Quaternion& In)
{
  Quaternion r = quaternionMultiply(quaternionMultiply(Quaternion{q.w, -q.x, -q.y, -q.z}, In), q);
  
  return r;
}

static Quaternion createYawQuaternion(float angle)
{
  Quaternion q;
  
  q.w = cosf(angle);
  q.x = 0.0f;
  q.y = 0.0f;
  q.z = sinf(angle);
  
  return q;
}

static Quaternion AccQuaternion(Quaternion& current, Vec3 a, float gravity, Vec3 xyz, float& error)
{
  float acc = sqrtf(a.x*a.x+a.y*a.y+a.z*a.z);
  if(acc>(1.0f+gravity) || acc<(1.0f-gravity)) return current;
  
  vecNormalize(a);
  
  float x=a.x-xyz.x, y=a.y-xyz.y, z=a.z-xyz.z;
  
  error=sqrtf(x*x+y*y+z*z)/10.0f;
  
  Vec3 g { 0, 0, 1 };
  Vec3 axis = vecCross(g, a);
  float w = 1 + (g.x * a.x + g.y * a.y + g.z * a.z);
  Quaternion q = { w, axis.x, axis.y, axis.z };
  normalizeQuaternion(q);
  
  Quaternion qYaw {current.w, 0, 0, current.z};
  
  return quaternionMultiply(q, qYaw); // ??????????? ?????? ?? Z
}

static Quaternion GyroQuaternion(Quaternion& current, float wx, float wy, float wz)
{
  Quaternion Mapp = current;
  Quaternion Spd { 0, wx, wy, wz };
  Quaternion aq = quaternionMultiply(Spd, Mapp);
  
  Mapp.w -= 0.5f * aq.w;
  Mapp.x -= 0.5f * aq.x;
  Mapp.y -= 0.5f * aq.y;
  Mapp.z -= 0.5f * aq.z;
  
  normalizeQuaternion(Mapp);
  return Mapp;
}

static Quaternion nlerp(const Quaternion& q1, const Quaternion& q2, float alpha)
{

  float dot = q1.w * q2.w + q1.x * q2.x + q1.y * q2.y + q1.z * q2.z;
  
  Quaternion q2_ = q2;
  if (dot < 0)
  {
    q2_.w = -q2.w;
    q2_.x = -q2.x;
    q2_.y = -q2.y;
    q2_.z = -q2.z;
  }
  
  Quaternion r;
  r.w = (1.0f - alpha) * q1.w + alpha * q2_.w;
  r.x = (1.0f - alpha) * q1.x + alpha * q2_.x;
  r.y = (1.0f - alpha) * q1.y + alpha * q2_.y;
  r.z = (1.0f - alpha) * q1.z + alpha * q2_.z;
  
  normalizeQuaternion(r);
  
  return r;
}

inline float GetAngle(float a1, float a2, float az)
{
  if(a2 == 0.0f && az == 0.0f) 
  {
    if(a1>0.0f) return 90.0f;
    if(a1<0.0f) return -90.0f;
    return 0.0f;
  }
  
  return atanf(a1/sqrtf(a2*a2+az*az)) * TO_DEG;
}

static Vec3 quaternionToPitchRollYaw(const Quaternion& q, Vec3& cosXYZ)
{
  Quaternion pry=rotateVectorByQuaternion(q, {0, 0, 1});
  
  float yaw = 2.0f * atan2f(q.z, q.w) * TO_DEG;
  if(yaw<0.0f) yaw=360.0f+yaw;
  
  cosXYZ = {pry.x, pry.y, pry.z};
  
  return // Sovereign orientation 
  {
    GetAngle(pry.y, pry.x, pry.z),  // Pitch
    GetAngle(-pry.x, pry.y, pry.z), // Roll
    yaw                             // Yaw
  };
}

static void addMagneto(Quaternion& q, Vec3 mag, float alpha, const float shift)
{
  static Quaternion yq = createYawQuaternion(shift*TO_RAD);
  
  vecNormalize(mag);
  
  Quaternion mQ={0, mag.x, mag.y, mag.z};
  
  Quaternion mW = backRotateVectorByQuaternion2(q, mQ);
  
  mW=quaternionMultiply(mW, yq); // Shifting the axes to the true north
  
  float gamma = mW.x * mW.x + mW.y * mW.y;
  float beta = sqrtf(gamma + mW.x * sqrtf(gamma));
  
  Quaternion mD
  {
    beta / (sqrtf(2.0f * gamma)),
    0.0f,
    0.0f,
    mW.y / (sqrtf(2.0f) * beta),
  };
  
  mD.w = (1.0f - alpha) + alpha * mD.w;
  mD.z = alpha * mD.z;
  
  if(mD.w!=mD.w || mD.x!=mD.x || mD.y!=mD.y || mD.z!=mD.z) return;
  
  q=quaternionMultiply(q, mD);
  normalizeQuaternion(q);
}

ORI WorkAccGyroMag(const Vec3 acc, const Vec3 gyr, const Vec3 mag, const float mag_shift, const float alpha)
{
    float wx = gyr.x * 500.0f / 30000.0f * TO_RAD * period;
  float wy = gyr.y * 500.0f / 30000.0f * TO_RAD * period;
  float wz = gyr.z * 500.0f / 30000.0f * TO_RAD * period;

  Vec3 aB = acc;
  
  Vec3 cosXYZ;
  Vec3 pry=quaternionToPitchRollYaw(qCurrent, cosXYZ);
  
  static float error=0;
 // Quaternion qAcc = AccQuaternion(qCurrent, aB, 0.05f, cosXYZ, error); // Tolerance for gravity deviation 5%
  
  qCurrent = GyroQuaternion(qCurrent, wx, wy, wz);
  
  if(error>0.01f) error=0.01f;
  if(error<alpha) error=alpha;
  
  /*
  Quaternion qFused = nlerp(qCurrent, qAcc, error);
  
  if (isFirst)
  {
    qFused = qAcc;
    isFirst = false;
  }
  
  qCurrent = qFused;
  */
  //addMagneto(qCurrent, mag, alpha, mag_shift);
  
  Vec3 ine=backRotateVectorByQuaternion(qCurrent, aB);
  
  return
  {
    cosXYZ.x, cosXYZ.y, cosXYZ.z,
    pry.x, pry.y, pry.z,
    ine.x, ine.y, ine.z,
  };
}

Vec3 RotateToZ(const Vec3 vec, bool Rev)
{
  Quaternion v { 0, vec.x, vec.y, vec.z };
  Quaternion q = { qCurrent.w, 0, 0, Rev ? -qCurrent.z : qCurrent.z };
  normalizeQuaternion(q);
  
  q = quaternionMultiply(quaternionMultiply(v, q), q); // ??????????? ?????? ?? Z
  
  return { q.x, q.y, q.z };
}

void GoZero()
{
  qCurrent= {1.0f, 0.0f, 0.0f, 0.0f};

}
