DQuaternion.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}

#ifdef TH_LIFT
{-# LANGUAGE DeriveLift #-}
{-# LANGUAGE DerivingStrategies #-}
#endif

module Geomancy.DQuaternion
  ( DQuaternion
  , dquaternion
  , withDQuaternion

  , axisAngle
  , intrinsic
  , extrinsic
  , rotate
  , rotatePoint
  , rotationBetween
  , lookAtUp

  , (^*)
  , (^/)
  , slerp

  , conjugate
  , norm
  , quadrance
  , dot
  , normalize
  , qNaN

  , fromQuaternion
  , toQuaternion
  ) where

import Control.DeepSeq (NFData(rnf))
import Foreign (Storable(..), castPtr)
import Foreign.Ptr.Diff (peekDiffOff, pokeDiffOff)
import GHC.Float (double2Float, float2Double)

#ifdef TH_LIFT
import Language.Haskell.TH.Syntax (Lift)
#endif

#ifdef SERIALISE
import Codec.Serialise (Serialise(..))
import qualified Codec.Serialise.Encoding as Serialise
import qualified Codec.Serialise.Decoding as Serialise
#endif

import Graphics.Gl.Block (Block(..))
import Geomancy.DVec3 (DVec3, dvec3, withDVec3)
import Geomancy.Quaternion (Quaternion, quaternion, withQuaternion)

import qualified Geomancy.DVec3 as DVec3

data DQuaternion = DQuaternion
  {-# UNPACK #-} !Double
  {-# UNPACK #-} !Double
  {-# UNPACK #-} !Double
  {-# UNPACK #-} !Double
  deriving (Eq, Ord, Show)
#ifdef TH_LIFT
  deriving Lift
#endif

{-# INLINE dquaternion #-}
dquaternion :: Double -> Double -> Double -> Double -> DQuaternion
dquaternion = DQuaternion

{-# INLINE withDQuaternion #-}
withDQuaternion
  :: DQuaternion
  -> (Double -> Double -> Double -> Double -> r)
  -> r
withDQuaternion (DQuaternion a b c d) f = f a b c d

{-# INLINE fromQuaternion #-}
fromQuaternion :: Quaternion -> DQuaternion
fromQuaternion q =
  withQuaternion q \a b c d ->
    dquaternion (float2Double a) (float2Double b) (float2Double c) (float2Double d)

{-# INLINE toQuaternion #-}
toQuaternion :: DQuaternion -> Quaternion
toQuaternion q =
  withDQuaternion q \a b c d ->
    quaternion (double2Float a) (double2Float b) (double2Float c) (double2Float d)

{-# INLINE (^*) #-}
(^*) :: DQuaternion -> Double -> DQuaternion
DQuaternion a b c d ^* x =
  DQuaternion
    (a * x)
    (b * x)
    (c * x)
    (d * x)

{-# INLINE (^/) #-}
(^/) :: DQuaternion -> Double -> DQuaternion
DQuaternion a b c d ^/ x =
  DQuaternion
    (a / x)
    (b / x)
    (c / x)
    (d / x)

slerp :: DQuaternion -> DQuaternion -> Double -> DQuaternion
slerp q p t
  | 1.0 - cosphi < 1e-12 =
      q
  | otherwise =
      ( (q   ^* sin ((1 - t) * phi)) +
         f p ^* sin (t * phi)
      ) ^/ sin phi
  where
    phi = acos cosphi

    (cosphi, f) =
      if dqp < 0 then
        (-dqp, negate)
      else
        (dqp, id)

    dqp = dot q p

{-# INLINE conjugate #-}
conjugate :: DQuaternion -> DQuaternion
conjugate (DQuaternion e x y z) = DQuaternion e (-x) (-y) (-z)

{-# INLINE norm #-}
norm :: DQuaternion -> Double
norm = sqrt . quadrance

{-# INLINE quadrance #-}
quadrance :: DQuaternion -> Double
quadrance q = dot q q

{-# INLINE dot #-}
dot :: DQuaternion -> DQuaternion -> Double
dot (DQuaternion a b c d) (DQuaternion e f g h) =
  a * e +
  b * f +
  c * g +
  d * h

{-# INLINE normalize #-}
normalize :: DQuaternion -> DQuaternion
normalize v =
  if nearZero q || nearZero (1-q) then
    v
  else
    let
      DQuaternion e i j k = v
    in
      DQuaternion (e / l) (i / l) (j / l) (k / l)

  where
    q = dot v v
    l = sqrt q

    nearZero a = abs a <= 1e-12

instance NFData DQuaternion where
  rnf DQuaternion{} = ()

instance Num DQuaternion where
  {-# INLINE (+) #-}
  DQuaternion a b c d + DQuaternion e f g h =
    DQuaternion
      (a + e)
      (b + f)
      (c + g)
      (d + h)

  {-# INLINE (-) #-}
  DQuaternion a b c d - DQuaternion e f g h =
    DQuaternion
      (a - e)
      (b - f)
      (c - g)
      (d - h)

  {-# INLINE (*) #-}
  DQuaternion as ax ay az * DQuaternion bs bx by bz =
    withDVec3 v \x y z ->
      DQuaternion s x y z
    where
      s = as * bs - DVec3.dot v1 v2
      v = DVec3.cross v1 v2 + v2 DVec3.^* as + v1 DVec3.^* bs
      v1 = dvec3 ax ay az
      v2 = dvec3 bx by bz

  {-# INLINE fromInteger #-}
  fromInteger x = DQuaternion (fromInteger x) 0 0 0

  {-# INLINE abs #-}
  abs z = DQuaternion (norm z) 0 0 0

  {-# INLINE signum #-}
  signum q@(DQuaternion e i j k)
    | m == 0 =
        q
    | not (isInfinite m || isNaN m) =
        DQuaternion (e * misqrt) (i * misqrt) (j * misqrt) (k * misqrt)
    | any isNaN [e, i, j, k] = qNaN
    | not (ii || ij || ik) = DQuaternion 1 0 0 0
    | not (ie || ij || ik) = DQuaternion 0 1 0 0
    | not (ie || ii || ik) = DQuaternion 0 0 1 0
    | not (ie || ii || ij) = DQuaternion 0 0 0 1
    | otherwise = qNaN
    where
      m = quadrance q
      misqrt = recip (sqrt m)

      ie = isInfinite e
      ii = isInfinite i
      ij = isInfinite j
      ik = isInfinite k

{-# INLINE qNaN #-}
qNaN :: DQuaternion
qNaN = DQuaternion dNaN dNaN dNaN dNaN
  where
    dNaN = 0/0

instance Storable DQuaternion where
  {-# INLINE sizeOf #-}
  sizeOf _ = 32

  {-# INLINE alignment #-}
  alignment _ = 32

  {-# INLINE poke #-}
  poke ptr (DQuaternion a b c d) = do
    poke ptr' a
    pokeElemOff ptr' 1 b
    pokeElemOff ptr' 2 c
    pokeElemOff ptr' 3 d
    where
      ptr' = castPtr ptr

  {-# INLINE peek #-}
  peek ptr = DQuaternion
    <$> peek ptr'
    <*> peekElemOff ptr' 1
    <*> peekElemOff ptr' 2
    <*> peekElemOff ptr' 3
    where
      ptr' = castPtr ptr

{-# INLINE axisAngle #-}
axisAngle :: DVec3 -> Double -> DQuaternion
axisAngle axis rads =
  withDVec3 (DVec3.normalize axis DVec3.^* sin half) $
    dquaternion (cos half)
  where
    half = rads / 2

intrinsic :: Double -> Double -> Double -> DQuaternion
intrinsic roll pitch yaw =
  dquaternion
    (cx * cy * cz + sx * sy * sz)
    (sx * cy * cz + cx * sy * sz)
    (cx * sy * cz - sx * cy * sz)
    (cx * cy * sz - sx * sy * cz)
  where
    cx = cos (pitch * 0.5)
    sx = sin (pitch * 0.5)
    cy = cos (yaw * 0.5)
    sy = sin (yaw * 0.5)
    cz = cos (roll * 0.5)
    sz = sin (roll * 0.5)

extrinsic :: Double -> Double -> Double -> DQuaternion
extrinsic heading elevation tilt =
  axisAngle (dvec3 0 0 1) tilt *
  axisAngle (dvec3 1 0 0) elevation *
  axisAngle (dvec3 0 1 0) heading

{-# INLINE rotate #-}
rotate :: DQuaternion -> DVec3 -> DVec3
rotate (DQuaternion w x y z) v =
  withDVec3 v \vx vy vz ->
    let
      x2 = x + x
      y2 = y + y
      z2 = z + z
      tx = y2 * vz - z2 * vy
      ty = z2 * vx - x2 * vz
      tz = x2 * vy - y2 * vx
    in
      dvec3
        (vx + w * tx + y * tz - z * ty)
        (vy + w * ty + z * tx - x * tz)
        (vz + w * tz + x * ty - y * tx)

{-# INLINE rotatePoint #-}
rotatePoint :: DQuaternion -> DVec3 -> DVec3 -> DVec3
rotatePoint q origin point =
  origin + rotate q (point - origin)

rotationBetween :: DVec3 -> DVec3 -> DQuaternion
rotationBetween v1 v2 = axisAngle axis angle
  where
    axis = DVec3.cross v1 v2
    angle = acos cosAngle
    cosAngle =
      max (-1) . min 1 $
        DVec3.dot (DVec3.normalize v1) (DVec3.normalize v2)

lookAtUp :: DVec3 -> DVec3 -> DVec3 -> DQuaternion
lookAtUp src dst up = rot2 * rot1
  where
    dir3 = dst - src

    rot1 = rotationBetween (dvec3 0 0 1) dir3

    rot2 = rotationBetween newUp fixedUp

    newUp = rotate rot1 up
    fixedUp = DVec3.cross (DVec3.cross dir3 up) dir3

instance Block DQuaternion where
  type PackedSize DQuaternion = 32
  alignment140 _  = 32
  sizeOf140       = sizeOfPacked
  alignment430    = alignment140
  sizeOf430       = sizeOf140
  isStruct _      = True
  read140     = peekDiffOff
  write140    = pokeDiffOff
  read430     = read140
  write430    = write140
  readPacked  = read140
  writePacked = write140
  {-# INLINE alignment140 #-}
  {-# INLINE sizeOf140 #-}
  {-# INLINE alignment430 #-}
  {-# INLINE sizeOf430 #-}
  {-# INLINE isStruct #-}
  {-# INLINE read140 #-}
  {-# INLINE write140 #-}
  {-# INLINE read430 #-}
  {-# INLINE write430 #-}
  {-# INLINE readPacked #-}
  {-# INLINE writePacked #-}

#ifdef SERIALISE
instance Serialise DQuaternion where
  encode v = withDQuaternion v \x y z w ->
    Serialise.encodeListLen 4 <> encode x <> encode y <> encode z <> encode w
  decode =
    Serialise.decodeListLenOf 4 *> (dquaternion <$> decode <*> decode <*> decode <*> decode)
#endif