```{-# OPTIONS_GHC -Wall #-}
{-# Language StandaloneDeriving #-}
{-# Language DeriveDataTypeable #-}

module SpatialMath ( module Xyz
, module Quat
--                   , Xyz(..)
--                   , Quat(..)
, Euler(..)
, euler321OfQuat
, euler321OfDcm
, quatOfEuler321
, dcmOfQuat
, dcmOfQuatB2A
, dcmOfEuler321
, quatOfDcm
, quatOfDcmB2A
, rotVecByDcm
, rotVecByDcmB2A
, rotVecByQuat
, rotVecByQuatB2A
, rotVecByEuler
, rotVecByEulerB2A
) where

import qualified Xyz
import qualified Quat
import Xyz ( Xyz(..) )
import Quat ( Quat(..) )

import Numeric.LinearAlgebra
import Foreign.Storable ( Storable )
import Data.Data ( Data )
import Data.Typeable ( Typeable1 )

data Euler a = Euler a a a deriving (Eq, Show)

deriving instance Typeable1 Euler
deriving instance Data a => Data (Euler a)

rotateXyzAboutX :: Floating a => Xyz a -> a -> Xyz a
rotateXyzAboutX (Xyz ax ay az) rotAngle = Xyz bx by bz
where
cosTheta = cos rotAngle
sinTheta = sin rotAngle

bx =  ax
by =  ay*cosTheta + az*sinTheta
bz = -ay*sinTheta + az*cosTheta

rotateXyzAboutY :: Floating a => Xyz a -> a -> Xyz a
rotateXyzAboutY (Xyz ax ay az) rotAngle = Xyz bx by bz
where
cosTheta = cos rotAngle
sinTheta = sin rotAngle

bx =  ax*cosTheta - az*sinTheta
by =  ay
bz =  ax*sinTheta + az*cosTheta

rotateXyzAboutZ :: Floating a => Xyz a -> a -> Xyz a
rotateXyzAboutZ (Xyz ax ay az) rotAngle = Xyz bx by bz
where
cosTheta = cos rotAngle
sinTheta = sin rotAngle

bx =  ax*cosTheta + ay*sinTheta
by = -ax*sinTheta + ay*cosTheta
bz =  az

euler321OfQuat :: RealFloat a => Quat a -> Euler a
euler321OfQuat (Quat q0 q1 q2 q3) = Euler yaw pitch roll
where
r11 = q0*q0 + q1*q1 - q2*q2 - q3*q3
r12 = 2.0*(q1*q2 + q0*q3)
mr13' = -2.0*(q1*q3 - q0*q2)
mr13 -- nan protect
| mr13' >  1 =  1
| mr13' < -1 = -1
| otherwise = mr13'
r23 = 2.0*(q2*q3 + q0*q1)
r33 = q0*q0 - q1*q1 - q2*q2 + q3*q3

yaw   = atan2 r12 r11
pitch = asin mr13
roll  = atan2 r23 r33

quatOfDcm :: (Storable a, RealFloat a) => Matrix a -> Quat a
quatOfDcm = quatOfEuler321 . euler321OfDcm

quatOfDcmB2A :: (Storable a, RealFloat a) => Matrix a -> Quat a
quatOfDcmB2A = Quat.inv . quatOfDcm

euler321OfDcm :: (RealFloat a, Storable a) => Matrix a -> Euler a
euler321OfDcm r = Euler yaw pitch roll
where
r11 = r @@> (0,0)
r12 = r @@> (0,1)
mr13' = -(r @@> (0,2))
mr13 -- nan protect
| mr13' >  1 =  1
| mr13' < -1 = -1
| otherwise = mr13'
r23 = r @@> (1,2)
r33 = r @@> (2,2)

yaw   = atan2 r12 r11
pitch = asin mr13
roll  = atan2 r23 r33

quatOfEuler321 :: (Floating a, Ord a) => Euler a -> Quat a
quatOfEuler321 (Euler yaw pitch roll) = Quat.normalize q
where
sr2 = sin \$ 0.5*roll
cr2 = cos \$ 0.5*roll
sp2 = sin \$ 0.5*pitch
cp2 = cos \$ 0.5*pitch
sy2 = sin \$ 0.5*yaw
cy2 = cos \$ 0.5*yaw
q0 = cr2*cp2*cy2 + sr2*sp2*sy2
q1 = sr2*cp2*cy2 - cr2*sp2*sy2
q2 = cr2*sp2*cy2 + sr2*cp2*sy2
q3 = cr2*cp2*sy2 - sr2*sp2*cy2

q' = Quat q0 q1 q2 q3

q
| q0 < 0 = negate q'
| otherwise = q'

dcmOfQuat :: (Num a, Element a) => Quat a -> Matrix a
dcmOfQuat (Quat q0 q1 q2 q3) = fromLists [ [r0, r1, r2]
, [r3, r4, r5]
, [r6, r7, r8]
]
where
-- 1st column
r0 = q0*q0 + q1*q1 - q2*q2 - q3*q3
r3 = 2*(q1*q2 - q0*q3)
r6 = 2*(q1*q3 + q0*q2)

-- 2nd column
r1 = 2*(q1*q2 + q0*q3)
r4 = q0*q0 - q1*q1 + q2*q2 - q3*q3
r7 = 2*(q2*q3 - q0*q1)

-- 3rd column
r2 = 2*(q1*q3 - q0*q2)
r5 = 2*(q2*q3 + q0*q1)
r8 = q0*q0 - q1*q1 - q2*q2 + q3*q3

dcmOfEuler321 :: (Floating a, Element a, Ord a) => Euler a -> Matrix a
dcmOfEuler321 = dcmOfQuat . quatOfEuler321

dcmOfQuatB2A :: (Num a, Element a) => Quat a -> Matrix a
dcmOfQuatB2A = dcmOfQuat . Quat.inv

-- | vec_b = R_a2b * vec_a
rotVecByDcm :: (Num a, Storable a) => Matrix a -> Xyz a -> Xyz a
rotVecByDcm dcm vec = Xyz.mult3x3ByXyz dcm vec

-- | vec_a = R_a2b^T * vec_b
rotVecByDcmB2A :: (Num a, Storable a) => Matrix a -> Xyz a -> Xyz a
rotVecByDcmB2A dcm vec = Xyz.mult3x3TransposeByXyz dcm vec

-- | vec_b = q_a2b * vec_a * q_a2b^(-1)
--   vec_b = R(q_a2b) * vec_a
rotVecByQuat :: (Num a, Element a) => Quat a -> Xyz a -> Xyz a
rotVecByQuat q = rotVecByDcm (dcmOfQuat q)

rotVecByQuatB2A :: (Num a, Element a) => Quat a -> Xyz a -> Xyz a
rotVecByQuatB2A q = rotVecByDcmB2A (dcmOfQuat q)

rotVecByEuler :: (Floating a, Element a, Ord a) => Euler a -> Xyz a -> Xyz a
rotVecByEuler = rotVecByDcm . dcmOfEuler321

rotVecByEulerB2A :: (Floating a, Element a, Ord a) => Euler a -> Xyz a -> Xyz a
rotVecByEulerB2A = rotVecByDcmB2A . dcmOfEuler321

-- void
-- get_wind_angles_from_v_bw_b(double * alpha, double * beta, double * airspeed, const xyz_t * const v_bw_b)
-- {
--   double airspeed_internal_memory;
--   double * airspeed_internal = &airspeed_internal_memory;
--
--   if (airspeed != NULL)
--     *airspeed = xyz_norm(v_bw_b) + 1e-12;
--
--   if (beta != NULL)
--   {
--     if (airspeed != NULL)
--       airspeed_internal = airspeed;
--     else
--       *airspeed_internal = xyz_norm(v_bw_b) + 1e-12;
--
--     *beta  =  asin ( v_bw_b->y / *airspeed_internal );
--   }
--
--   if (alpha != NULL)
--     *alpha =  atan2( v_bw_b->z, v_bw_b->x );
-- }
--
-- void
-- get_wind_angles( double * alpha,
--                  double * beta,
--                  double * airspeed,
--                  xyz_t * v_bw_b_out,
--                  const quat_t * const q_n2b,
--                  const xyz_t * const v_bn_b,
--                  const xyz_t * const v_wn_n)
-- {
--   xyz_t v_wn_b;
--   rot_vec_by_quat_a2b( &v_wn_b, q_n2b, v_wn_n);
--   xyz_t v_bw_b;
--   xyz_diff( &v_bw_b, v_bn_b, &v_wn_b);
--
--   get_wind_angles_from_v_bw_b( alpha, beta, airspeed, &v_bw_b );
--
--   if (v_bw_b_out != NULL)
--     xyz_memcpy( v_bw_b_out, &v_bw_b);
-- }
--
-- void
-- v_bw_b_from_wind_angles( xyz_t * v_bw_b, const double alpha, const double beta, const double airspeed)
-- {
--   v_bw_b->x = airspeed*cos(alpha)*cos(beta);
--   v_bw_b->y = airspeed*sin(beta);
--   v_bw_b->z = airspeed*cos(beta)*sin(alpha);
-- }
```