{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE Safe #-}

{- | 
Module      :  Physics.Learn.Mechanics
Copyright   :  (c) Scott N. Walck 2014-2019
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

Newton's second law and all that
-}

module Physics.Learn.Mechanics
    ( TheTime
    , TimeStep
    , Velocity
    -- * Simple one-particle state
    , SimpleState
    , SimpleAccelerationFunction
    , simpleStateDeriv
    , simpleRungeKuttaStep
    -- * One-particle state
    , St(..)
    , DSt(..)
    , OneParticleSystemState
    , OneParticleAccelerationFunction
    , oneParticleStateDeriv
    , oneParticleRungeKuttaStep
    , oneParticleRungeKuttaSolution
    -- * Two-particle state
    , TwoParticleSystemState
    , TwoParticleAccelerationFunction
    , twoParticleStateDeriv
    , twoParticleRungeKuttaStep
    -- * Many-particle state
    , ManyParticleSystemState
    , ManyParticleAccelerationFunction
    , manyParticleStateDeriv
    , manyParticleRungeKuttaStep
    )
    where

import Data.VectorSpace
    ( AdditiveGroup(..)
    , VectorSpace(..)
    )
import Physics.Learn.StateSpace
    ( StateSpace(..)
    , Diff
    , DifferentialEquation
    )
import Physics.Learn.RungeKutta
    ( rungeKutta4
    , integrateSystem
    )
import Physics.Learn.Position
    ( Position
    )
import Physics.Learn.CarrotVec
    ( Vec
    )

-- | Time (in s).
type TheTime = Double

-- | A time step (in s).
type TimeStep = Double

-- | Velocity of a particle (in m/s).
type Velocity = Vec

-------------------------------
-- Simple one-particle state --
-------------------------------

-- | A simple one-particle state,
--   to get started quickly with mechanics of one particle.
type SimpleState = (TheTime,Position,Velocity)

-- | An acceleration function gives the particle's acceleration as
--   a function of the particle's state.
--   The specification of this function is what makes one single-particle
--   mechanics problem different from another.
--   In order to write this function, add all of the forces
--   that act on the particle, and divide this net force by the particle's mass.
--   (Newton's second law).
type SimpleAccelerationFunction = SimpleState -> Vec

-- | Time derivative of state for a single particle
--   with a constant mass.
simpleStateDeriv :: SimpleAccelerationFunction  -- ^ acceleration function for the particle
                 -> DifferentialEquation SimpleState  -- ^ differential equation
simpleStateDeriv a (t, r, v) = (1, v, a(t, r, v))

-- | Single Runge-Kutta step
simpleRungeKuttaStep :: SimpleAccelerationFunction  -- ^ acceleration function for the particle
                     -> TimeStep  -- ^ time step
                     -> SimpleState  -- ^ initial state
                     -> SimpleState  -- ^ state after one time step
simpleRungeKuttaStep = rungeKutta4 . simpleStateDeriv

------------------------
-- One-particle state --
------------------------

-- | The state of a single particle is given by
--   the position of the particle and the velocity of the particle.
data St = St { position :: Position
             , velocity :: Velocity
             }
          deriving (Show)

-- | The associated vector space for the
--   state of a single particle.
data DSt = DSt Vec Vec
           deriving (Show)

instance AdditiveGroup DSt where
    zeroV = DSt zeroV zeroV
    negateV (DSt dr dv) = DSt (negateV dr) (negateV dv)
    DSt dr1 dv1 ^+^ DSt dr2 dv2 = DSt (dr1 ^+^ dr2) (dv1 ^+^ dv2)

instance VectorSpace DSt where
    type Scalar DSt = Double
    c *^ DSt dr dv = DSt (c*^dr) (c*^dv)

instance StateSpace St where
    type Diff St = DSt
    St r1 v1 .-. St r2 v2  = DSt (r1 .-. r2) (v1 .-. v2)
    St r1 v1 .+^ DSt dr dv = St (r1 .+^ dr) (v1 .+^ dv)

-- | The state of a system of one particle is given by the current time,
--   the position of the particle, and the velocity of the particle.
--   Including time in the state like this allows us to
--   have time-dependent forces.
type OneParticleSystemState = (TheTime,St)

-- | An acceleration function gives the particle's acceleration as
--   a function of the particle's state.
type OneParticleAccelerationFunction = OneParticleSystemState -> Vec

-- | Time derivative of state for a single particle
--   with a constant mass.
oneParticleStateDeriv :: OneParticleAccelerationFunction  -- ^ acceleration function for the particle
                      -> DifferentialEquation OneParticleSystemState  -- ^ differential equation
oneParticleStateDeriv a st@(_t, St _r v) = (1, DSt v (a st))

-- | Single Runge-Kutta step
oneParticleRungeKuttaStep :: OneParticleAccelerationFunction  -- ^ acceleration function for the particle
                          -> TimeStep  -- ^ time step
                          -> OneParticleSystemState  -- ^ initial state
                          -> OneParticleSystemState  -- ^ state after one time step
oneParticleRungeKuttaStep = rungeKutta4 . oneParticleStateDeriv

-- | List of system states
oneParticleRungeKuttaSolution :: OneParticleAccelerationFunction  -- ^ acceleration function for the particle
                              -> TimeStep  -- ^ time step
                              -> OneParticleSystemState  -- ^ initial state
                              -> [OneParticleSystemState]  -- ^ state after one time step
oneParticleRungeKuttaSolution = integrateSystem . oneParticleStateDeriv

------------------------
-- Two-particle state --
------------------------

-- | The state of a system of two particles is given by the current time,
--   the position and velocity of particle 1,
--   and the position and velocity of particle 2.
type TwoParticleSystemState = (TheTime,St,St)

-- | An acceleration function gives a pair of accelerations
--   (one for particle 1, one for particle 2) as
--   a function of the system's state.
type TwoParticleAccelerationFunction = TwoParticleSystemState -> (Vec,Vec)

-- | Time derivative of state for two particles
--   with constant mass.
twoParticleStateDeriv :: TwoParticleAccelerationFunction  -- ^ acceleration function for two particles
                      -> DifferentialEquation TwoParticleSystemState  -- ^ differential equation
twoParticleStateDeriv af2 st2@(_t, St _r1 v1, St _r2 v2) = (1, DSt v1 a1, DSt v2 a2)
    where
      (a1,a2) = af2 st2

-- | Single Runge-Kutta step for two-particle system
twoParticleRungeKuttaStep :: TwoParticleAccelerationFunction  -- ^ acceleration function
                          -> TimeStep  -- ^ time step
                          -> TwoParticleSystemState  -- ^ initial state
                          -> TwoParticleSystemState  -- ^ state after one time step
twoParticleRungeKuttaStep = rungeKutta4 . twoParticleStateDeriv

-------------------------
-- Many-particle state --
-------------------------

-- | The state of a system of many particles is given by the current time
--   and a list of one-particle states.
type ManyParticleSystemState = (TheTime,[St])

-- | An acceleration function gives a list of accelerations
--   (one for each particle) as
--   a function of the system's state.
type ManyParticleAccelerationFunction = ManyParticleSystemState -> [Vec]

-- | Time derivative of state for many particles
--   with constant mass.
manyParticleStateDeriv :: ManyParticleAccelerationFunction  -- ^ acceleration function for many particles
                       -> DifferentialEquation ManyParticleSystemState  -- ^ differential equation
manyParticleStateDeriv af st@(_t, sts) = (1, [DSt v a | (v,a) <- zip vs as])
    where
      vs = map velocity sts
      as = af st

-- | Single Runge-Kutta step for many-particle system
manyParticleRungeKuttaStep :: ManyParticleAccelerationFunction  -- ^ acceleration function
                           -> TimeStep  -- ^ time step
                           -> ManyParticleSystemState  -- ^ initial state
                           -> ManyParticleSystemState  -- ^ state after one time step
manyParticleRungeKuttaStep = rungeKutta4 . manyParticleStateDeriv



-- Can we automatically incorporate Newton's third law?