-- | This module provides a general interface for working with mechanical systems.
module Goal.Simulation.Physics.Configuration
    ( -- * Configurations
    -- ** Types
      Generalized (Generalized)
    , PhaseSpace
    , GeneralizedVelocity
    , GeneralizedAcceleration
    , GeneralizedInertia
    -- ** Accessors
    , body
    , position
    , velocity
    , momentum
    -- * ForceFields
    -- ** Classes
    , ForceField (force)
    , Conservative (potentialEnergy)
    , vectorField
    , mechanicalEnergy
    -- ** Types
    , Gravity (Gravity)
    , earthGravity
    , Damping (Damping)
    -- * Util
    , periodic
    , revolutions
    , sliceVectorField
    , pairToPoint
    , pointToPair
    ) where

--- Imports ---


-- Goal --

import Goal.Geometry

-- Qualified --

import qualified Data.Vector.Storable as C


--- Configuration Space ---


-- Charts --

-- | Generalized coordinates of a mechanical system.
data Generalized = Generalized

-- Manifolds --

-- | The 'Tangent' space of a mechanical system in 'Generalized' coordinates is the space of 'GeneralizedVelocity's.
type GeneralizedVelocity m = Tangent Generalized m

-- | The second order 'Tangent' space of a mechanical system is the space of 'GeneralizedAcceleration's.
type GeneralizedAcceleration m = Tangent Partials (GeneralizedVelocity m)

-- | The tangent 'Bundle' of a dynamical system is also known as the
-- 'PhaseSpace'. The "state" of a mechanical system is typically understood to
-- be an element of the 'PhaseSpace'.
type PhaseSpace m = Bundle Generalized m

-- | The 'Riemannian' metric on a mechanical system is known as the 'GeneralizedInertia'.
type GeneralizedInertia m = Tensor (GeneralizedVelocity m) (GeneralizedVelocity m)

-- Functions --

-- | Returns the underlying mechanical 'Manifold' of an element of the 'PhaseSpace'.
body :: Manifold m => Partials :#: PhaseSpace m -> m
body = removeBundle . manifold

-- | Returns the 'position' coordinates of a mechanical system.
position :: Manifold m => Partials :#: PhaseSpace m -> Generalized :#: m
position = projectTangent . bundleToTangent

-- | Returns the 'velocity' coordinates of a mechanical system.
velocity :: Manifold m => Partials :#: PhaseSpace m -> Partials :#: GeneralizedVelocity m
velocity = bundleToTangent

-- | Returns the 'momentum' coordinates of a mechanical system.
momentum :: Manifold m => Differentials :#: PhaseSpace m -> Differentials :#: GeneralizedVelocity m
momentum = bundleToTangent

kineticEnergy :: Riemannian Generalized m => Partials :#: GeneralizedVelocity m -> Double
kineticEnergy dq = 0.5 * (flat dq <.> dq)

-- Force fields --

-- | Gravitational force.
newtype Gravity = Gravity Double

earthGravity :: Gravity
earthGravity = Gravity 9.80665

newtype Damping = Damping Double

-- | A 'ForceField' describes how to attach a force vector (an element of the
-- dual space of generalized accelerations) to every point in the phase space.
-- Note that a 'ForceField' is not necessarily 'Conservative'.
class Manifold m => ForceField f m where
    force :: f -> Partials :#: PhaseSpace m -> Differentials :#: GeneralizedAcceleration m

-- | A 'Conservative' force depends only on 'position's and can be described as the gradient of a 'potentialEnergy'.
class ForceField f m => Conservative f m where
    potentialEnergy :: f -> Generalized :#: m -> Double

-- | The 'vectorField' function takes a 'ForceField' on a mechanical system and converts it into the appropriate element of the 'Tangent' space of the 'PhaseSpace'.
vectorField :: (Riemannian Generalized m, ForceField f m)
    => f
    -> Partials :#: PhaseSpace m
    -> Partials :#: Tangent Partials (PhaseSpace m)
vectorField f qdq = fromCoordinates (Tangent qdq) $ coordinates (velocity qdq) C.++ coordinates (sharp $ force f qdq)

mechanicalEnergy :: (Riemannian Generalized m, Conservative f m)
    => f
    -> Partials :#: PhaseSpace m
    -> (Double, Double)
mechanicalEnergy f qdq = (kineticEnergy $ velocity qdq, potentialEnergy f $ position qdq)


--- Functions ---


periodic :: Manifold m => [Bool] -> (Partials :#: PhaseSpace m) -> (Partials :#: PhaseSpace m)
periodic bls qdq =
    let q = position qdq
        q' = fromList (manifold q) $ clipper <$> zip bls (listCoordinates q)
        dq' = fromCoordinates (Tangent q') . coordinates $ velocity qdq
     in tangentToBundle dq'
       where clipper (True,x) = x - 2 * pi * fromIntegral (revolutions x)
             clipper (False,x) = x

revolutions :: Double -> Int
revolutions x
    | x >= pi = floor $ (x + pi) / (2*pi)
    | x <= -pi = ceiling $ (x - pi) / (2*pi)
    | otherwise = 0

sliceVectorField :: (Riemannian Generalized m, ForceField f m)
    => Double
    -> Int
    -> f
    -> Partials :#: PhaseSpace m
    -> (Double,Double)
    -> (Double,Double)
sliceVectorField scl i f qdq =
    pointToPair scl i . vectorField f . pairToPoint i qdq

pairToPoint :: Manifold m => Int -> Partials :#: PhaseSpace m -> (Double, Double) -> Partials :#: PhaseSpace m
pairToPoint n qdq (x,dx) =
    let (hxs,_:txs) = splitAt n . listCoordinates $ position qdq
        (hdxs,_:tdxs) = splitAt n . listCoordinates $ velocity qdq
     in fromList (manifold qdq) $ hxs ++ x:txs ++ hdxs ++ dx:tdxs

pointToPair :: Manifold m => Double -> Int -> (c :#: m) -> (Double, Double)
pointToPair scl n dqddq =
     (scl * coordinate n dqddq, scl * coordinate (n + div 2 (dimension $ manifold dqddq)) dqddq)

--- Instances ---


instance (ForceField f m, ForceField g m) => ForceField (f,g) m where
    force (f,g) qdq = force f qdq  <+> force g qdq

instance Manifold m => ForceField (Partials :#: PhaseSpace m -> Differentials :#: GeneralizedAcceleration m) m where
    force f = f

-- Damping --

instance Manifold m => ForceField Damping m where
    force (Damping c) qdq =
        let dq = velocity qdq
         in fromCoordinates (Tangent dq) . coordinates $ alterCoordinates (negate . (*c)) dq


--- Graveyard ---

{-
traversableToBundle :: (T.Traversable f, Manifold (f m), Manifold m) => f (d :#: Bundle c m) -> d :#: Bundle c (f m)
traversableToBundle qdqt =
    let pbt = Bundle $ removeBundle . manifold <$> qdqt
        (qs,qds) = unzip . F.toList $ cleanPoint <$> qdqt
     in fromCoordinates pbt $ C.concat qs C.++ C.concat qds
       where cleanPoint pbp = C.splitAt (dimension $ manifold pbp) $ coordinates pbp

bundleToTraversable :: (T.Traversable f, Manifold (f m), Manifold m) => d :#: Bundle c (f m) -> f (d :#: Bundle c m)
bundleToTraversable pbtp =
    let pbt = removeBundle $ manifold pbtp
     in snd . T.mapAccumL seperatePlanarLink (C.splitAt (dimension pbt) $ coordinates pbtp) $ pbt
      where seperatePlanarLink (qs,qds) pb =
                let (hqs,tqs) = C.splitAt (dimension pb) qs
                    (hqds,tqds) = C.splitAt (dimension pb) qds
                 in ((tqs,tqds), fromCoordinates (Bundle pb) $ hqs C.++ hqds)
-}