{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
module Data.Vec.Quaternion
    ( Quaternion(..)
    , qconj
    , realq
    , imagq
    , qscale
    , qrotate
    , qpow
    , slerp
    , toAxisAngle
    , fromAxisAngle
    , toRotationMatrix
    , qintegrate
    ) where

import Prelude hiding (head,tail,foldr,foldl,map,zipWith,last)
import Data.Vec

-- TODO: use Packed
newtype Quaternion a = Q (Vec4 a)

deriving instance Access N0 a (Quaternion a)
deriving instance Access N1 a (Quaternion a)
deriving instance Access N2 a (Quaternion a)
deriving instance Access N3 a (Quaternion a)
instance Head (Quaternion a) a where
    head (Q v) = head v
instance Tail (Quaternion a) (Vec3 a) where
    tail (Q v) = tail v
instance Last (Quaternion a) a where
    last (Q v) = last v
instance Fold (Quaternion a) a where
    fold f (Q v) = fold f v
    foldl f a (Q v) = foldl f a v
    foldr f a (Q v) = foldr f a v
instance Map a b (Quaternion a) (Quaternion b) where
    map f (Q v) = Q (map f v)
instance ZipWith a b c (Quaternion a) (Quaternion b) (Quaternion c) where
    zipWith f (Q u) (Q v) = Q (zipWith f u v)

deriving instance Show a => Show (Quaternion a)
deriving instance Eq a => Eq (Quaternion a)

qconj :: Num a => Quaternion a -> Quaternion a
qconj (Q (s :. v)) = Q (s :. negate v)

-- | Construct a quaternion with the given real part, and zero imaginary parts.
realq :: Num a => a -> Quaternion a
realq a = Q (a :. vec 0)

-- | Construct a quaternion with the given imaginary parts, and zero real part.
imagq :: Num a => Vec3 a -> Quaternion a
imagq w = Q (0 :. w)

-- | @qscale a q@ is equivalent to, but more efficient than, @realq a * q@.
qscale :: Num a => a -> Quaternion a -> Quaternion a
qscale a (Q q) = Q (vec a * q)

-- | Unlike the 'Vec' instance for 'Num', multiplication is defined as the quaternion multiplication operator, not element-by-element. 'fromInteger' creates a real-valued quaternion (that is, one with zero imaginary parts).
-- 'abs' returns the real-valued quaternion equal to the argument's Euclidean norm. 'signum' is the same as 'normalize' (NOT an element-by-element signum). In this way, the identity @q == abs q * signum q@ is satisfied for all but the zero quaternion, which of course becomes @NaN@ when normalized.
instance (Num a, Floating a) => Num (Quaternion a) where
    Q a + Q b = Q (a+b)
    Q a - Q b = Q (a-b)
    Q (as:.av) * Q (bs:.bv) = Q ((as*bs - dot av bv) :. (vec as*bv + vec bs*av + cross av bv))
    negate (Q q) = Q (negate q)
    fromInteger n = Q (fromInteger n :. vec 0)
    abs q = Q (norm q :. vec 0)
    signum = normalize

-- | '(/)' is defined as right division, @q * r^-1@. Note that this is not equivalent to left division, @r^-1 * q@.
-- 'fromRational' produces a real-valued quaternion.
instance Floating a => Fractional (Quaternion a) where
    recip q = qscale (recip (dot q q)) (qconj q)
    -- Right division. Not equivalent to left division.
    q / r = q * recip r
    fromRational t = Q (fromRational t :. vec 0)

-- | All functions in 'Floating' are defined as appropriate over the quaternions, and do not operate element-by-element.
-- 'pi' is a real-valued quaternion.
instance RealFloat a => Floating (Quaternion a) where
    pi = Q (pi :. vec 0)
    exp (Q (s:.v)) = Q ((exp s * cos nv) :. (v * vec (exp s * sin_recip nv))) where nv = norm v
    log q@(Q (s:.v)) = Q (log (norm q) :. (normalize v * vec (acos (s / norm q))))
    q ** p = exp (log q * p)
    -- TODO: use a more efficient formula for sqrt
    sqrt q = q ** 0.5
    -- TODO: for sin, cos, sinh, cosh, use the Taylor series expansions when nv is close to zero.
    sin (Q (s:.v)) = Q ((sin s * cosh nv) :. (v * vec (cos s * sinh_recip nv))) where nv = norm v
    cos (Q (s:.v)) = Q ((cos s * cosh nv) :. (v * vec (- sin s * sinh_recip nv))) where nv = norm v
    sinh (Q (s:.v)) = Q ((sinh s * cos nv) :. (v * vec (cosh s * sin_recip nv))) where nv = norm v
    cosh (Q (s:.v)) = Q ((cosh s * cos nv) :. (v * vec (sinh s * sin_recip nv))) where nv = norm v
    tan q = sin q / cos q
    tanh q = sin q / cos q
    asinh q = log (q + sqrt (q*q + 1))
    acosh q = log (q + sqrt (q*q - 1))
    atanh q = qscale 0.5 (log (1+q) - log (1-q))
    asin q@(Q (_:.v)) = -u*asinh(q*u) where u = Q (0:.normalize v)
    acos q@(Q (_:.v)) = -u*acosh(q)   where u = Q (0:.normalize v)
    atan q@(Q (_:.v)) = -u*atanh(q*u) where u = Q (0:.normalize v)

-- | Rotation of a 3-dimensional vector by a quaternion.
qrotate :: Floating a => Quaternion a -> Vec3 a -> Vec3 a
qrotate q v = tail (q * Q (0:.v) * qconj q)

-- | Real power of a quaternion. @qpow q p@ is equivalent to, but more efficient than, @q ** 'realq' p@.
qpow :: RealFloat a => Quaternion a -> a -> Quaternion a
qpow q p = exp (qscale p (log q))

-- | Spherical linear interpolation of quaternion rotations.
--slerp q0 _  0 = q0 -- otherwise qpow gives NaN
slerp :: RealFloat a => Quaternion a -> Quaternion a -> a -> Quaternion a
slerp q0 q1 t = qpow (q1 / q0) t * q0

-- | Convert from a unit quaternion to an axis-angle representation.
toAxisAngle :: RealFloat a => Quaternion a -> (a, Vec3 a)
toAxisAngle q = if | w > 1       -> toAxisAngle (normalize q)
                   | s < epsilon -> (0, 1:.0:.0:.())
                   | otherwise   -> (2.0*acos w, vec (recip s) * tail q)
                where w = head q
                      s = sqrt (1-w*w)

-- | Convert from an axis-angle representation of rotation to a unit quaternion (if the axis is a unit vector).
fromAxisAngle :: Floating a => Vec3 a -> a -> Quaternion a
fromAxisAngle axis angle = Q (cos (0.5*angle) :. (vec (sin (0.5*angle)) * axis))

toRotationMatrix :: Floating a => Quaternion a -> Vec3 (Vec3 a)
-- TODO: use more efficient formula
toRotationMatrix q = (qrotate q (1:.0:.0) :. qrotate q (0:.1:.0) :. qrotate q (0:.0:.1) :. ())
-- XXX: why doesn't this one work?
{-
toRotationMatrix (Q (w:.x:.y:.z:.())) =
    ( ((1-2*y*y-z*z) :. (2*x*y-2*w*z) :. (2*x*z+2*w*y) :. ())
    :.((2*x*y+2*w*z) :. (1-x*x-z*z)   :. (2*y*z+2*w*x) :. ())
    :.((2*x*z-2*w*y) :. (2*y*z+2*w*x) :. (1-x*x-y*y) :. ()):.())
    -}

-- | "Euler integration": apply an angular velocity to a quaternion. The angular velocity is a vector parallel to the axis of rotation, with magnitude equal to the speed times the time step.
qintegrate :: (RealFloat a, Ord a) => Vec3 a -> Quaternion a -> Quaternion a
qintegrate w q = dq * q
    where dq = Q (c :. (vec s * theta))
          mThetaSq = normSq theta
          mTheta = sqrt mThetaSq
          theta = vec 0.5 * w
          -- Taylor approximation is used when mTheta is close to zero (sin x / x becomes unstable)
          (c,s) = if mThetaSq*mThetaSq/24 < epsilon
                     then (1 - mThetaSq/2, 1 - mThetaSq/6)
                     else (cos mTheta, sin mTheta / mTheta)

--
-- numerics junk
--

-- | Machine epsilon. (Requires @ScopedTypeVariables@)
epsilon :: forall a. RealFloat a => a
epsilon = encodeFloat 1 (fromIntegral $ 1-floatDigits (undefined :: a))

eps120_4 :: RealFloat a => a
eps120_4 = (120*epsilon) ** (1/4)

-- | @sin_recip x == sin x / x@, even when x is close to zero.
sin_recip :: RealFloat a => a -> a
sin_recip  x = if abs x < eps120_4 then 1 - x*x/6 else sin  x / x
-- | @sinh_recip x == sinh x / x@, even when x is close to zero.
sinh_recip :: RealFloat a => a -> a
sinh_recip x = if abs x < eps120_4 then 1 + x*x/6 else sinh x / x

