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

{- | 
Module      :  Physics.Learn.StateSpace
Copyright   :  (c) Scott N. Walck 2014
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
    , TimeDerivative
    )
    where

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

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

-- | A 'StateSpace' has an associated vector space, the vectors of which
--   can be multiplied or divided by scalars.
--   An example would be the set of positions of a particle.
--   Position is not a vector, but displacement (difference in position) is a vector.
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 .-^ v = p .+^ negateV v

instance StateSpace Double where
    type Diff Double = Double
    (.-.) = (-)
    (.+^) = (+)

instance StateSpace Vec where
    type Diff Vec = Vec
    (.-.) = (^-^)
    (.+^) = (^+^)

instance StateSpace Position where
    type Diff Position = Vec
    (.-.) = flip displacement
    (.+^) = flip shiftPosition

instance (StateSpace p, StateSpace q, Time p ~ Time q) => StateSpace (p,q) where
  type Diff (p,q)   = (Diff p, Diff q)
  (p,q) .-. (p',q') = (p .-. p', q .-. q')
  (p,q) .+^ (u,v)   = (p .+^ u, 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,q,r) .-. (p',q',r') = (p .-. p', q .-. q', r .-. r')
  (p,q,r) .+^ (u,v,w)    = (p .+^ u, q .+^ v, r .+^ w)

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

instance AdditiveGroup v => AdditiveGroup [v] where
    zeroV   = inf zeroV
    (^+^)   = zipWith (^+^)
    negateV = map negateV

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

instance StateSpace p => StateSpace [p] where
    type Diff [p] = [Diff p]
    (.-.) = zipWith (.-.)
    (.+^) = zipWith (.+^)

-- | The time derivative of a state is an element of the associated vector space.
type TimeDerivative state = state -> Diff state

{-
class HasTimeDerivative state where
    timeDeriv :: state -> Diff state
-}