package calculator

import (
	"math"

	"gonum.org/v1/gonum/mat"
)

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

// ToRotationMatrix converts a quaternion to a rotation matrix.
func (q Quaternion) ToRotationMatrix() *mat.Dense {
	w, x, y, z := q.W, q.X, q.Y, q.Z

	return mat.NewDense(3, 3, []float64{
		1 - 2*y*y - 2*z*z, 2*x*y - 2*w*z, 2*x*z + 2*w*y,
		2*x*y + 2*w*z, 1 - 2*x*x - 2*z*z, 2*y*z - 2*w*x,
		2*x*z - 2*w*y, 2*y*z + 2*w*x, 1 - 2*x*x - 2*y*y,
	})
}

// EulerToQuaternion converts Euler angles (yaw, pitch, roll) to a quaternion.
// Assuming the order of rotations is ZYX (yaw-pitch-roll).
func EulerToQuaternion(yaw, pitch, roll float64) Quaternion {
	halfYaw := yaw / 2
	halfPitch := pitch / 2
	halfRoll := roll / 2

	cy := math.Cos(halfYaw)
	sy := math.Sin(halfYaw)
	cp := math.Cos(halfPitch)
	sp := math.Sin(halfPitch)
	cr := math.Cos(halfRoll)
	sr := math.Sin(halfRoll)

	q := Quaternion{
		W: cy*cp*cr + sy*sp*sr,
		X: cy*cp*sr - sy*sp*cr,
		Y: sy*cp*sr + cy*sp*cr,
		Z: sy*cp*cr - cy*sp*sr,
	}

	return q
}

// EulerToRotMatrix converts Euler angles (roll, pitch, yaw) to a rotation matrix.
func EulerToRotMatrix(phi, theta, psi float64) *mat.Dense {
	// Calculate individual rotation matrices
	Rx := mat.NewDense(3, 3, []float64{
		1, 0, 0,
		0, math.Cos(phi), -math.Sin(phi),
		0, math.Sin(phi), math.Cos(phi),
	})

	Ry := mat.NewDense(3, 3, []float64{
		math.Cos(theta), 0, math.Sin(theta),
		0, 1, 0,
		-math.Sin(theta), 0, math.Cos(theta),
	})

	Rz := mat.NewDense(3, 3, []float64{
		math.Cos(psi), -math.Sin(psi), 0,
		math.Sin(psi), math.Cos(psi), 0,
		0, 0, 1,
	})

	// R = Rz * Ry * Rx
	RyRx := mat.NewDense(3, 3, nil)
	RyRx.Mul(Ry, Rx)

	R := mat.NewDense(3, 3, nil)
	R.Mul(Rz, RyRx)

	return R
}

// RotMatrixToQuaternion converts a rotation matrix to a quaternion.
func RotMatrixToQuaternion(R *mat.Dense) Quaternion {
	m := R.RawMatrix().Data

	trace := m[0] + m[4] + m[8]
	var q Quaternion

	if trace > 0 {
		S := 0.5 / math.Sqrt(trace+1.0)
		q.W = 0.25 / S
		q.X = (m[7] - m[5]) * S
		q.Y = (m[2] - m[6]) * S
		q.Z = (m[3] - m[1]) * S
	} else {
		if m[0] > m[4] && m[0] > m[8] {
			S := 2.0 * math.Sqrt(1.0+m[0]-m[4]-m[8])
			q.W = (m[7] - m[5]) / S
			q.X = 0.25 * S
			q.Y = (m[1] + m[3]) / S
			q.Z = (m[2] + m[6]) / S
		} else if m[4] > m[8] {
			S := 2.0 * math.Sqrt(1.0+m[4]-m[0]-m[8])
			q.W = (m[2] - m[6]) / S
			q.X = (m[1] + m[3]) / S
			q.Y = 0.25 * S
			q.Z = (m[5] + m[7]) / S
		} else {
			S := 2.0 * math.Sqrt(1.0+m[8]-m[0]-m[4])
			q.W = (m[3] - m[1]) / S
			q.X = (m[2] + m[6]) / S
			q.Y = (m[5] + m[7]) / S
			q.Z = 0.25 * S
		}
	}

	return q
}

func RotMatrixToEuler(R *mat.Dense) (yaw, pitch, roll float64) {
	m := R.RawMatrix().Data

	if m[6] < 1 {
		if m[6] > -1 {
			pitch = math.Asin(-m[6])
			roll = math.Atan2(m[7], m[8])
			yaw = math.Atan2(m[3], m[0])
		} else {
			pitch = math.Pi / 2
			roll = -math.Atan2(-m[1], m[4])
			yaw = 0
		}
	} else {
		pitch = -math.Pi / 2
		roll = math.Atan2(-m[1], m[4])
		yaw = 0
	}

	return yaw, pitch, roll
}