{-# OPTIONS_GHC -Wall -fno-warn-orphans #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances, TypeFamilies #-}
{-# LANGUAGE Safe #-}

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

A 'StateSpace' is an affine space where the associated vector space
has scalars that are instances of 'Fractional'.
If p is an instance of 'StateSpace', then the associated vectorspace
'Diff' p is intended to represent the space of (time) derivatives
of paths in p.

'StateSpace' is very similar to Conal Elliott's 'AffineSpace'.
-}

module Physics.Learn.StateSpace
    ( StateSpace(..)
    , (.-^)
    , Time
    , DifferentialEquation
    , InitialValueProblem
    , EvolutionMethod
    , SolutionMethod
    , stepSolution
    , eulerMethod
    )
    where

import Data.AdditiveGroup
    ( AdditiveGroup(..)
    )
import Data.VectorSpace
    ( VectorSpace(..)
    )
import Physics.Learn.Position
    ( Position
    , shiftPosition
    , displacement
    )
import Physics.Learn.CarrotVec
    ( Vec
    , (^*)
    , (^-^)
    )

infixl 6 .+^, .-^
infix  6 .-.

-- | An instance of 'StateSpace' is a data type that can serve as the state
--   of some system.  Alternatively, a 'StateSpace' is a collection of dependent
--   variables for a differential equation.
--   A 'StateSpace' has an associated vector space for the (time) derivatives
--   of the state.  The associated vector space is a linearized version of
--   the 'StateSpace'.
class (VectorSpace (Diff p), Fractional (Scalar (Diff p))) => StateSpace p where
  -- | Associated vector space
  type Diff p
  -- | Subtract points
  (.-.)  :: p -> p -> Diff p
  -- | Point plus vector
  (.+^)  :: p -> Diff p -> p

-- | The scalars of the associated vector space can be thought of as time intervals.
type Time p = Scalar (Diff p)

-- | Point minus vector
(.-^) :: StateSpace p => p -> Diff p -> p
p
p .-^ :: forall p. StateSpace p => p -> Diff p -> p
.-^ Diff p
v = p
p forall p. StateSpace p => p -> Diff p -> p
.+^ forall v. AdditiveGroup v => v -> v
negateV Diff p
v

instance StateSpace Double where
    type Diff Double = Double
    .-. :: Double -> Double -> Diff Double
(.-.) = (-)
    .+^ :: Double -> Diff Double -> Double
(.+^) = forall a. Num a => a -> a -> a
(+)

instance StateSpace Vec where
    type Diff Vec = Vec
    .-. :: Vec -> Vec -> Diff Vec
(.-.) = forall v. AdditiveGroup v => v -> v -> v
(^-^)
    .+^ :: Vec -> Diff Vec -> Vec
(.+^) = forall v. AdditiveGroup v => v -> v -> v
(^+^)

-- | Position is not a vector, but displacement (difference in position) is a vector.
instance StateSpace Position where
    type Diff Position = Vec
    .-. :: Position -> Position -> Diff Position
(.-.) = forall a b c. (a -> b -> c) -> b -> a -> c
flip Position -> Position -> Vec
displacement
    .+^ :: Position -> Diff Position -> Position
(.+^) = forall a b c. (a -> b -> c) -> b -> a -> c
flip Vec -> Position -> Position
shiftPosition

instance (StateSpace p, StateSpace q, Time p ~ Time q) => StateSpace (p,q) where
  type Diff (p,q)   = (Diff p, Diff q)
  (p
p,q
q) .-. :: (p, q) -> (p, q) -> Diff (p, q)
.-. (p
p',q
q') = (p
p forall p. StateSpace p => p -> p -> Diff p
.-. p
p', q
q forall p. StateSpace p => p -> p -> Diff p
.-. q
q')
  (p
p,q
q) .+^ :: (p, q) -> Diff (p, q) -> (p, q)
.+^ (Diff p
u,Diff q
v)   = (p
p forall p. StateSpace p => p -> Diff p -> p
.+^ Diff p
u, q
q forall p. StateSpace p => p -> Diff p -> p
.+^ Diff q
v)

instance (StateSpace p, StateSpace q, StateSpace r, Time p ~ Time q
         ,Time q ~ Time r) => StateSpace (p,q,r) where
  type Diff (p,q,r)      = (Diff p, Diff q, Diff r)
  (p
p,q
q,r
r) .-. :: (p, q, r) -> (p, q, r) -> Diff (p, q, r)
.-. (p
p',q
q',r
r') = (p
p forall p. StateSpace p => p -> p -> Diff p
.-. p
p', q
q forall p. StateSpace p => p -> p -> Diff p
.-. q
q', r
r forall p. StateSpace p => p -> p -> Diff p
.-. r
r')
  (p
p,q
q,r
r) .+^ :: (p, q, r) -> Diff (p, q, r) -> (p, q, r)
.+^ (Diff p
u,Diff q
v,Diff r
w)    = (p
p forall p. StateSpace p => p -> Diff p -> p
.+^ Diff p
u, q
q forall p. StateSpace p => p -> Diff p -> p
.+^ Diff q
v, r
r forall p. StateSpace p => p -> Diff p -> p
.+^ Diff r
w)

inf :: a -> [a]
inf :: forall a. a -> [a]
inf a
x = a
x forall a. a -> [a] -> [a]
: forall a. a -> [a]
inf a
x

instance AdditiveGroup v => AdditiveGroup [v] where
    zeroV :: [v]
zeroV   = forall a. a -> [a]
inf forall v. AdditiveGroup v => v
zeroV
    ^+^ :: [v] -> [v] -> [v]
(^+^)   = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall v. AdditiveGroup v => v -> v -> v
(^+^)
    negateV :: [v] -> [v]
negateV = forall a b. (a -> b) -> [a] -> [b]
map forall v. AdditiveGroup v => v -> v
negateV

instance VectorSpace v => VectorSpace [v] where
    type Scalar [v] = Scalar v
    Scalar [v]
c *^ :: Scalar [v] -> [v] -> [v]
*^ [v]
xs = [Scalar [v]
c forall v. VectorSpace v => Scalar v -> v -> v
*^ v
x | v
x <- [v]
xs]

instance StateSpace p => StateSpace [p] where
    type Diff [p] = [Diff p]
    .-. :: [p] -> [p] -> Diff [p]
(.-.) = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall p. StateSpace p => p -> p -> Diff p
(.-.)
    .+^ :: [p] -> Diff [p] -> [p]
(.+^) = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall p. StateSpace p => p -> Diff p -> p
(.+^)

-- | A differential equation expresses how the dependent variables (state)
--   change with the independent variable (time).
--   A differential equation is specified by giving the (time) derivative
--   of the state as a function of the state.
--   The (time) derivative of a state is an element of the associated vector space.
type DifferentialEquation state = state -> Diff state

-- | An initial value problem is a differential equation along with an initial state.
type InitialValueProblem state = (DifferentialEquation state, state)

-- | A (numerical) solution method is a way of converting
--   an initial value problem into a list of states (a solution).
--   The list of states need not be equally spaced in time.
type SolutionMethod state = InitialValueProblem state -> [state]

-- | An evolution method is a way of approximating the state
--   after advancing a finite interval in the independent
--   variable (time) from a given state.
type EvolutionMethod state
    = DifferentialEquation state  -- ^ differential equation
    -> Time state                 -- ^ time interval
    -> state                      -- ^ initial state
    -> state                      -- ^ evolved state

-- | Given an evolution method and a time step, return the solution method
--   which applies the evolution method repeatedly with with given time step.
--   The solution method returned will produce an infinite list of states.
stepSolution :: EvolutionMethod state -> Time state -> SolutionMethod state
stepSolution :: forall state.
EvolutionMethod state -> Time state -> SolutionMethod state
stepSolution EvolutionMethod state
ev Time state
dt (DifferentialEquation state
de, state
ic) = forall a. (a -> a) -> a -> [a]
iterate (EvolutionMethod state
ev DifferentialEquation state
de Time state
dt) state
ic

-- | The Euler method is the simplest evolution method.
--   It increments the state by the derivative times the time step.
eulerMethod :: StateSpace state => EvolutionMethod state
eulerMethod :: forall state. StateSpace state => EvolutionMethod state
eulerMethod DifferentialEquation state
de Scalar (Diff state)
dt state
st = state
st forall p. StateSpace p => p -> Diff p -> p
.+^ DifferentialEquation state
de state
st forall v s. (VectorSpace v, s ~ Scalar v) => v -> s -> v
^* Scalar (Diff state)
dt