luastd/Quaternion.lua

---@class Quaternion
---@field x number
---@field y number
---@field z number
---@field w number
local Quaternion = {}
Quaternion.__index = Quaternion


function math.sign(x)
    return x < 0 and -1 or 1
  end

---@param x number
---@param y number
---@param z number
---@param w number
---@return Quaternion
function Quaternion.new(x, y, z, w)
  return setmetatable({ x = x or 0, y = y or 0, z = z or 0, w = w or 1 }, Quaternion)
end

---@return Quaternion
function Quaternion:copy()
  return Quaternion.new(self.x, self.y, self.z, self.w)
end

---@param axis Vector3
---@param angle number
---@return Quaternion
function Quaternion.fromAxisAngle(axis, angle)
  axis = axis:normalize()
  local halfAngle = angle * 0.5
  local sinHalf = math.sin(halfAngle)
  return Quaternion.new(
      axis.x * sinHalf,
      axis.y * sinHalf,
      axis.z * sinHalf,
      math.cos(halfAngle)
  )
end

---@return string
function Quaternion:__tostring()
  return string.format("Quaternion(%.3f, %.3f, %.3f, %.3f)", self.x, self.y, self.z, self.w)
end

---@return number
function Quaternion:length()
  return math.sqrt(self.x^2 + self.y^2 + self.z^2 + self.w^2)
end

---@return Quaternion
function Quaternion:normalize()
  local len = self:length()
  if len == 0 then return Quaternion.new(0, 0, 0, 1) end
  return Quaternion.new(self.x / len, self.y / len, self.z / len, self.w / len)
end

---@param q Quaternion
---@return Quaternion
function Quaternion:mul(q)
  return Quaternion.new(
    self.w * q.x + self.x * q.w + self.y * q.z - self.z * q.y,
    self.w * q.y - self.x * q.z + self.y * q.w + self.z * q.x,
    self.w * q.z + self.x * q.y - self.y * q.x + self.z * q.w,
    self.w * q.w - self.x * q.x - self.y * q.y - self.z * q.z
  )
end

---@return table
function Quaternion:toMatrix()
  local x, y, z, w = self.x, self.y, self.z, self.w
  return {
    {1 - 2*y^2 - 2*z^2, 2*x*y - 2*w*z,     2*x*z + 2*w*y},
    {2*x*y + 2*w*z,     1 - 2*x^2 - 2*z^2, 2*y*z - 2*w*x},
    {2*x*z - 2*w*y,     2*y*z + 2*w*x,     1 - 2*x^2 - 2*y^2}
  }
end

---@return number
---@return number
---@return number
function Quaternion:toEuler()
  local x, y, z, w = self.x, self.y, self.z, self.w

  local siny_cosp = 2 * (w * y + x * z)
  local cosy_cosp = 1 - 2 * (y^2 + z^2)
  local yaw = math.atan2(siny_cosp, cosy_cosp)

  local sinp = 2 * (w * x - y * z)
  local pitch = math.abs(sinp) >= 1 and (math.pi / 2) * math.sign(sinp) or math.asin(sinp)

  local sinr_cosp = 2 * (w * z + x * y)
  local cosr_cosp = 1 - 2 * (z^2 + x^2)
  local roll = math.atan2(sinr_cosp, cosr_cosp)

  return roll, pitch, yaw
end


---@return table
function Quaternion:toTable()
  return { x = self.x, y = self.y, z = self.z, w = self.w }
end

---@param t table
---@return Quaternion
function Quaternion.fromTable(t)
  return Quaternion.new(t.x, t.y, t.z, t.w)
end


---@param rx number
---@param ry number
---@param rz number
---@return Quaternion
function Quaternion.fromEuler(rx, ry, rz)
  local cy = math.cos(rz * 0.5)
  local sy = math.sin(rz * 0.5)
  local cp = math.cos(ry * 0.5)
  local sp = math.sin(ry * 0.5)
  local cr = math.cos(rx * 0.5)
  local sr = math.sin(rx * 0.5)

  return Quaternion.new(
    sr * cp * cy - cr * sp * sy,
    cr * sp * cy + sr * cp * sy,
    cr * cp * sy - sr * sp * cy,
    cr * cp * cy + sr * sp * sy
  )
end


---@param v Vector3 | Vec3Like
---@return table
function Quaternion:rotateVector(v)
  local qvec = { x = self.x, y = self.y, z = self.z }
  local uv = {
    x = qvec.y * v.z - qvec.z * v.y,
    y = qvec.z * v.x - qvec.x * v.z,
    z = qvec.x * v.y - qvec.y * v.x
  }
  local uuv = {
    x = qvec.y * uv.z - qvec.z * uv.y,
    y = qvec.z * uv.x - qvec.x * uv.z,
    z = qvec.x * uv.y - qvec.y * uv.x
  }

  return {
    x = v.x + 2 * (self.w * uv.x + uuv.x),
    y = v.y + 2 * (self.w * uv.y + uuv.y),
    z = v.z + 2 * (self.w * uv.z + uuv.z)
  }
end

---@return Quaternion
function Quaternion:conjugate()
  return Quaternion.new(-self.x, -self.y, -self.z, self.w)
end

return Quaternion