package filter

import (
	"math"
)

type MadgwickFilter struct {
	Quaternion Quaternion
	Beta       float64 // git gut bwtween 0.01..0.1
}

type Quaternion struct {
	W, X, Y, Z float64
}

func (q *Quaternion) Normalize() {
	mag := math.Sqrt(q.W*q.W + q.X*q.X + q.Y*q.Y + q.Z*q.Z)
	q.W /= mag
	q.X /= mag
	q.Y /= mag
	q.Z /= mag
}

func (f *MadgwickFilter) Update(gx, gy, gz, ax, ay, az float64, dt float64) {
	q := f.Quaternion

	norm := math.Sqrt(ax*ax + ay*ay + az*az)
	if norm == 0 {
		return
	}
	ax /= norm
	ay /= norm
	az /= norm

	// get gravity vector
	vx := 2 * (q.X*q.Z - q.W*q.Y)
	vy := 2 * (q.W*q.X + q.Y*q.Z)
	vz := q.W*q.W - q.X*q.X - q.Y*q.Y + q.Z*q.Z

	ex := (ay*vz - az*vy)
	ey := (az*vx - ax*vz)
	ez := (ax*vy - ay*vx)

	gx += f.Beta * ex
	gy += f.Beta * ey
	gz += f.Beta * ez

	// integrate over rate of change
	halfdt := 0.5 * dt
	dq := Quaternion{
		W: 0,
		X: gx * halfdt,
		Y: gy * halfdt,
		Z: gz * halfdt,
	}
	qDot := dq.Multiply(q)

	q.W += qDot.W
	q.X += qDot.X
	q.Y += qDot.Y
	q.Z += qDot.Z
	q.Normalize()

	f.Quaternion = q
}

func (q Quaternion) Multiply(r Quaternion) Quaternion {
	return Quaternion{
		W: q.W*r.W - q.X*r.X - q.Y*r.Y - q.Z*r.Z,
		X: q.W*r.X + q.X*r.W + q.Y*r.Z - q.Z*r.Y,
		Y: q.W*r.Y - q.X*r.Z + q.Y*r.W + q.Z*r.X,
		Z: q.W*r.Z + q.X*r.Y - q.Y*r.X + q.Z*r.W,
	}
}

func (q Quaternion) ToEuler() (roll, pitch, yaw float64) {
	sinr_cosp := 2 * (q.W*q.X + q.Y*q.Z)
	cosr_cosp := 1 - 2*(q.X*q.X+q.Y*q.Y)
	roll = math.Atan2(sinr_cosp, cosr_cosp)

	sinp := 2 * (q.W*q.Y - q.Z*q.X)
	if math.Abs(sinp) >= 1 {
		pitch = math.Copysign(math.Pi/2, sinp)
	} else {
		pitch = math.Asin(sinp)
	}

	siny_cosp := 2 * (q.W*q.Z + q.X*q.Y)
	cosy_cosp := 1 - 2*(q.Y*q.Y+q.Z*q.Z)
	yaw = math.Atan2(siny_cosp, cosy_cosp)

	return
}
