Quaternion.hs

-- | Quaternion transforms for three-dimensional Euclidean space
module Mini.Linear.Quaternion (
  -- * Type
  Quaternion (Quaternion),

  -- * Construction
  axisAngle,
  identity,
  zero,

  -- * Operations
  (%*%),
  (%+%),
  (%-%),
  (~*%),
  conjugate,
  inverse,
  norm,

  -- * Rotation
  rotate,

  -- * Interpolation
  slerp,

  -- * Conversion
  toMatrix,
) where

import Control.Applicative (
  liftA2,
 )
import Control.Monad (
  ap,
  liftM,
 )
import Mini.Hash.Class (
  Hashable,
  toBytes,
 )
import Mini.Linear.Approx (
  Approx,
  (~=),
 )
import qualified Mini.Linear.Matrix as M (
  identity,
 )
import Mini.Linear.Space (
  R1,
  R2,
  R3,
  R4,
  V3,
  V4 (V4),
  w,
  x,
  xy,
  xyz,
  xyzw,
  y,
  z,
 )
import Mini.Linear.Vector (
  cross,
  dot,
  (^+^),
  (~*^),
 )
import qualified Mini.Linear.Vector as V (
  norm,
  zero,
 )
import Mini.Optics.Lens (
  over,
  set,
  view,
 )
import Mini.Random.Class (
  Random,
  random,
 )
import Prelude (
  Applicative,
  Eq,
  Floating,
  Foldable,
  Functor,
  Monad,
  Monoid,
  Num,
  Ord,
  Semigroup,
  Show,
  Traversable,
  acos,
  and,
  concatMap,
  cos,
  fmap,
  foldr,
  mempty,
  negate,
  pure,
  recip,
  sin,
  sum,
  traverse,
  ($),
  (*),
  (+),
  (-),
  (.),
  (/),
  (<$>),
  (<*>),
  (<>),
  (>>=),
 )

-- Type

-- | Wrapper whose 'xyz'-components are imaginary and 'w'-component is real
newtype Quaternion a = Quaternion (V4 a)
  deriving (Show, Eq, Ord)

instance R1 Quaternion where
  x f (Quaternion v) =
    (\x' -> Quaternion $ set x x' v) <$> f (view x v)

instance R2 Quaternion where
  xy f (Quaternion v) =
    (\xy' -> Quaternion $ set xy xy' v) <$> f (view xy v)

instance R3 Quaternion where
  xyz f (Quaternion v) =
    (\xyz' -> Quaternion $ set xyz xyz' v) <$> f (view xyz v)

instance R4 Quaternion where
  xyzw f (Quaternion v) = Quaternion <$> f v

instance Functor Quaternion where
  fmap = liftM

instance Applicative Quaternion where
  pure = Quaternion . pure
  (<*>) = ap

instance Monad Quaternion where
  m >>= k =
    Quaternion
      ( V4
          (view x . k $ view x m)
          (view y . k $ view y m)
          (view z . k $ view z m)
          (view w . k $ view w m)
      )

instance Foldable Quaternion where
  foldr f b (Quaternion v) = foldr f b v

instance Traversable Quaternion where
  traverse f (Quaternion v) = Quaternion <$> traverse f v

instance (Semigroup a) => Semigroup (Quaternion a) where
  (<>) = liftA2 (<>)

instance (Monoid a) => Monoid (Quaternion a) where
  mempty = pure mempty

instance (Approx a) => Approx (Quaternion a) where
  p ~= q = and $ liftA2 (~=) p q

instance (Hashable a) => Hashable (Quaternion a) where
  toBytes = concatMap toBytes

instance (Random a) => Random (Quaternion a) where
  random g = let (v, g') = random g in (Quaternion v, g')

-- Construction

-- | Rotation around an axis /v/ a number of radians (assumes @v ~= hat v@)
axisAngle :: (Floating a) => V3 a -> a -> Quaternion a
axisAngle v a =
  let v' = sin (a / 2) ~*^ v
   in Quaternion $ V4 (view x v') (view y v') (view z v') (cos (a / 2))

-- | Multiplicative identity quaternion
identity :: (Num a) => Quaternion a
identity = set w 1 zero

-- | Additive identity quaternion
zero :: (Num a) => Quaternion a
zero = Quaternion V.zero

-- Operations

infixl 7 %*%

-- | Quaternion-quaternion multiplication
(%*%) :: (Num a) => Quaternion a -> Quaternion a -> Quaternion a
q %*% r =
  let qv = view xyz q
      qw = view w q
      rv = view xyz r
      rw = view w r
      v' = (qv `cross` rv) ^+^ (rw ~*^ qv) ^+^ (qw ~*^ rv)
      w' = (qw * rw - qv `dot` rv)
   in Quaternion $
        V4
          (view x v')
          (view y v')
          (view z v')
          w'

infixl 6 %+%

-- | Quaternion-quaternion addition
(%+%) :: (Num a) => Quaternion a -> Quaternion a -> Quaternion a
(%+%) = liftA2 (+)

infixl 6 %-%

-- | Quaternion-quaternion subtraction
(%-%) :: (Num a) => Quaternion a -> Quaternion a -> Quaternion a
(%-%) = liftA2 (-)

infixl 7 ~*%

-- | Scalar-quaternion multiplication
(~*%) :: (Num a) => a -> Quaternion a -> Quaternion a
s ~*% q = over w (s *) $ over xyz (s ~*^) q

-- | Conjugate of a quaternion
conjugate :: (Num a) => Quaternion a -> Quaternion a
conjugate = over xyz (fmap negate)

-- | Multiplicative inverse of a quaternion /q/ (assumes @not $ q ~= zero@)
inverse :: (Floating a) => Quaternion a -> Quaternion a
inverse q = recip (norm q * norm q) ~*% conjugate q

-- | Norm of a quaternion
norm :: (Floating a) => Quaternion a -> a
norm = V.norm . view xyzw

-- Rotation

-- | Apply a rotation /r/ to a vector (assumes /r/ is a unit quaternion)
rotate :: (Floating a) => Quaternion a -> V3 a -> V3 a
rotate r v =
  view xyz $
    r
      %*% Quaternion (V4 (view x v) (view y v) (view z v) 0)
      %*% conjugate r

-- Interpolation

-- | Spherical linear interpolation s.t. @slerp q r 0 == q && slerp q r 1 == r@
slerp :: (Floating a) => Quaternion a -> Quaternion a -> a -> Quaternion a
slerp q r t =
  let phi = acos . sum $ liftA2 (*) q r
      q' = (sin (phi * (1 - t)) / sin phi) ~*% q
      r' = (sin (phi * t) / sin phi) ~*% r
   in q' %+% r'

-- Conversion

-- | Convert a rotation /r/ to matrix form (assumes /r/ is a unit quaternion)
toMatrix :: (Num a) => Quaternion a -> V4 (V4 a)
toMatrix (Quaternion (V4 qx qy qz qw)) =
  set (x . x) (1 - 2 * (qy * qy + qz * qz))
    . set (x . y) (2 * (qx * qy + qw * qz))
    . set (x . z) (2 * (qx * qz - qw * qy))
    . set (y . x) (2 * (qx * qy - qw * qz))
    . set (y . y) (1 - 2 * (qx * qx + qz * qz))
    . set (y . z) (2 * (qy * qz + qw * qx))
    . set (z . x) (2 * (qx * qz + qw * qy))
    . set (z . y) (2 * (qy * qz - qw * qx))
    . set (z . z) (1 - 2 * (qx * qx + qy * qy))
    $ M.identity