#!/usr/bin/env stack
-- stack runghc --resolver lts-16 --package ad --package hmatrix-vector-sized-0.1.3.0 --package vector-sized --package hmatrix --package comonad --package free

-- | Source file accompanying
-- https://blog.jle.im/entry/hamiltonian-dynamics-in-haskell.html

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns        #-}

import           Data.Foldable
import           Data.Maybe
import           GHC.TypeLits
import           Numeric.AD
import           Numeric.LinearAlgebra.Static
import           Numeric.LinearAlgebra.Static.Vector
import qualified Control.Comonad                     as C
import qualified Control.Comonad.Cofree              as C
import qualified Data.Vector.Sized                   as V

-- | A @'System' m n@ represents a system parameterized by @n@ generalized
-- coordinates, moving in an "underlying" @m@-dimensional cartesian
-- coordinate system.
data System m n = System
    { sysInertia       :: R m                         -- ^ 'm' vector
    , sysCoords        :: R n -> R m                  -- ^ f
    , sysJacobian      :: R n -> L m n                -- ^ J_f
    , sysHessian       :: R n -> V.Vector n (L m n)   -- ^ H_f
    , sysPotential     :: R n -> Double               -- ^ U
    , sysPotentialGrad :: R n -> R n                  -- ^ grad U
    }

-- | A system's position in configuration space
data Config n = Config
    { confPositions  :: R n
    , confVelocities :: R n
    }
  deriving Show

-- | A system's position in phase space
data Phase n = Phase
    { phasePositions :: R n
    , phaseMomenta   :: R n
    }
  deriving Show

-- | Given a position in @n@-dimensional generalized coordinates, return
-- the position in the underlying @m@-dimensional cartesian coordinate
-- system.
underlyingPosition
    :: System m n
    -> R n
    -> R m
underlyingPosition = sysCoords

-- | Give the system's generalized momenum, given its position in
-- configuration space.
momenta
    :: (KnownNat n, KnownNat m)
    => System m n
    -> Config n
    -> R n
momenta s (Config q v) = tr j #> mHat #> j #> v
  where
    j    = sysJacobian s q
    mHat = diag (sysInertia s)

-- | Convert a position in configuration space to a position in phase space
toPhase
    :: (KnownNat n, KnownNat m)
    => System m n
    -> Config n
    -> Phase n
toPhase s c = Phase (confPositions c) (momenta s c)

-- | Shift around the Hessian into the shape expected
tr2 :: (KnownNat m, KnownNat n)
    => V.Vector m (L n n)
    -> V.Vector n (L m n)
tr2 = fmap rowsL . traverse lRows
{-# INLINE tr2 #-}

-- | Vector of vectors to matrix
vec2l
    :: KnownNat n
    => V.Vector m (V.Vector n Double)
    -> L m n
vec2l = rowsL . fmap gvecR
{-# INLINE vec2l #-}

-- | Make a system given its inertias, coordinate functions, and potential
-- energy function
mkSystem
    :: (KnownNat m, KnownNat n)
    => R m
    -> (forall a. RealFloat a => V.Vector n a -> V.Vector m a)
    -> (forall a. RealFloat a => V.Vector n a -> a)
    -> System m n
mkSystem m f u = System
                    -- < convert from      | actual thing | convert to >
    { sysInertia       =                     m
    , sysCoords        =        gvecR      . f            . grVec
    , sysJacobian      = tr   . vec2l      . jacobianT f  . grVec
    , sysHessian       = tr2  . fmap vec2l . hessianF f   . grVec
    , sysPotential     =                     u            . grVec
    , sysPotentialGrad =        gvecR      . grad u       . grVec
    }


-- | Equations of motion for a system at a given position in phase space
hamilEqns
    :: (KnownNat n, KnownNat m)
    => System m n
    -> Phase n
    -> (R n, R n)       -- dq/dt and dp/dt
hamilEqns s (Phase q p) = (dqdt, dpdt)
  where
    j       = sysJacobian s q
    trj     = tr j
    mHat    = diag (sysInertia s)
    kHat    = trj `mul` mHat `mul` j
    kHatInv = inv kHat
    dqdt    = kHatInv #> p
    dpdt    = gvecR bigUglyThing - sysPotentialGrad s q
      where
        bigUglyThing =
          fmap (\j2 -> -p <.> kHatInv #> trj #> mHat #> j2 #> kHatInv #> p)
               (sysHessian s q)

-- | Step a system's position through phase space using Euler integration
stepEuler
    :: (KnownNat n, KnownNat m)
    => System m n       -- ^ the system
    -> Double           -- ^ dt
    -> Phase n          -- ^ q(t) and p(t)
    -> Phase n          -- ^ q(t + dt) and p(t + dt)
stepEuler s dt ph@(Phase q p) = Phase (q + konst dt * dq) (p + konst dt * dp)
  where
    (dq, dp) = hamilEqns s ph

-- | Iterate Euler's method to repeatedly step a system through phase space
runSystem
    :: (KnownNat n, KnownNat m)
    => System m n       -- ^ the system
    -> Double           -- ^ dt
    -> Phase n          -- ^ initial phase
    -> [Phase n]        -- ^ progression of the system using Euler integration
runSystem s dt = go
  where
    go p0 = p0 : go (stepEuler s dt p0)

-- | A simple particle in 2D cartesian space under gravity
simpleSystem :: System 2 2
simpleSystem = mkSystem (vec2 5 5) id pot
  where
    -- potential energy of a gravity field
    -- U(x,y) = 9.8 * y
    pot :: RealFloat a => V.Vector 2 a -> a
    pot xy = 9.8 * (xy `V.index` 1)

-- | An initial position in configuration space, representing a particle
-- at <0,0> with initial velocity <1,3>
simpleConfig0 :: Config 2
simpleConfig0 = Config
    { confPositions  = vec2 0 0
    , confVelocities = vec2 1 3
    }

simpleMain :: IO ()
simpleMain =
    mapM_ (disp 2 . phasePositions)  -- position with 2 digits of precision
  . take 25                          -- 25 steps
  $ runSystem simpleSystem 0.1 (toPhase simpleSystem simpleConfig0)


-- | A pendulum system, parameterized by its angle clockwise from
-- equilibrium
pendulum :: System 2 1
pendulum = mkSystem (vec2 5 5) coords pot      -- 5kg particle
  where
    -- <x,y> = <-0.5 sin(theta), -0.5 cos(theta)>
    -- pendulum of length 0.25
    coords :: RealFloat a => V.Vector 1 a -> V.Vector 2 a
    coords (V.head->theta) = V.fromTuple (- 0.25 * sin theta, - 0.25 * cos theta)
    -- potential energy of gravity field
    -- U(x,y) = 9.8 * y
    pot :: RealFloat a => V.Vector 1 a -> a
    pot q = 9.8 * (coords q `V.index` 1)

-- | An initial pendulum position in configuration space, representing
-- a pendulum at equilibrium with initial angular velocity 0.1 rad/sec
-- clockwise
pendulumConfig0 :: Config 1
pendulumConfig0 = Config
    { confPositions  = 0
    , confVelocities = 0.1
    }

pendulumMain :: IO ()
pendulumMain =
    mapM_ (disp 3 . phasePositions)  -- position with 2 digits of precision
  . take 25                          -- 25 steps
  $ runSystem pendulum 0.1 (toPhase pendulum pendulumConfig0)

main :: IO ()
main = pendulumMain

