-- | The Map module provides tools for developing function space 'Manifold's.
-- A map is a 'Manifold' where the 'Point's of the Manifold represent
-- parametric functions between 'Manifold's. The defining feature of 'Map's is
-- that they have a particular 'Domain' and 'Codomain', which themselves are
-- 'Manifold's.

module Goal.Simulation.Optimization (
    -- * Mean Squared Error
      meanSquaredError
    -- * Cauchy Sequences
    , cauchyLimit
    , cauchySequence
    -- * Gradient Pursuit
    , stochasticGradientDescent
    , stochasticGradientAscent
    , stochasticVanillaGradientDescent
    , stochasticVanillaGradientAscent
    , boundedStochasticVanillaGradientDescent
    , boundedStochasticVanillaGradientAscent
    -- * Least Squares
    , designMatrix
    , leastSquares
    , leastSquares0
    -- ** Newton
    , newtonStep
    , newtonSequence
    -- ** Gauss Newton
    , gaussNewtonStep
    ) where

--- Imports ---

import Prelude hiding (map,minimum,maximum)

-- Goal --

import Goal.Core
import Goal.Geometry
import Goal.Probability

import Goal.Simulation.Mealy


--- Stochastic Pursuit ---


type StochasticPursuit x c m = Mealy x (c :#: m)


--- Gradient Descent ---

stochasticGradientAscent :: (Riemannian c m, Manifold m)
    => Double -- ^ Step size
    -> (c :#: m -> x -> forall s . RandST s (Differentials :#: Tangent c m)) -- ^ Gradient calculator
    -> (c :#: m) -- ^ The initial point
    -> RandST r (StochasticPursuit x c m) -- ^ The gradient ascent
stochasticGradientAscent eps0 f0' = accumulateRandomFunction (accumulator eps0 f0')
    where accumulator eps f' x cm = do
              dcm <- f' cm x
              let cm' = gradientStep eps $ sharp dcm
              return (cm',cm')

boundedStochasticVanillaGradientAscent :: Manifold m
    => Double -- ^ Step size
    -> Double -- ^ Gradient Rejection Bound
    -> (c :#: m -> x -> forall s . RandST s (Differentials :#: Tangent c m)) -- ^ Gradient calculator
    -> (c :#: m) -- ^ The initial point
    -> RandST r (StochasticPursuit x c m) -- ^ The gradient ascent
boundedStochasticVanillaGradientAscent eps0 bnd0 f0' = accumulateRandomFunction (accumulator eps0 bnd0 f0')
    where accumulator eps bnd f' x cm = do
              dcm <- f' cm x
              let cm' = if maximum (abs <$> listCoordinates dcm) < bnd
                            then gradientStep eps $ breakChart dcm
                            else trace "Ping!" cm
              return (cm',cm')

stochasticVanillaGradientAscent :: Manifold m
    => Double -- ^ Step size
    -> (c :#: m -> x -> forall s . RandST s (Differentials :#: Tangent c m)) -- ^ Gradient calculator
    -> (c :#: m) -- ^ The initial point
    -> RandST r (StochasticPursuit x c m) -- ^ The gradient ascent
stochasticVanillaGradientAscent eps0 f0' = accumulateRandomFunction (accumulator eps0 f0')
    where accumulator eps f' x cm = do
              dcm <- f' cm x
              let cm' = gradientStep eps $ breakChart dcm
              return (cm',cm')

stochasticGradientDescent :: (Riemannian c m, Manifold m)
    => Double -- ^ Step size
    -> (c :#: m -> x -> forall s . RandST s (Differentials :#: Tangent c m)) -- ^ Gradient calculator
    -> (c :#: m) -- ^ The initial point
    -> RandST r (StochasticPursuit x c m) -- ^ The gradient ascent
stochasticGradientDescent eps = stochasticGradientAscent (-eps)

boundedStochasticVanillaGradientDescent :: Manifold m
    => Double -- ^ Step size
    -> Double -- ^ Gradient Rejection Bound
    -> (c :#: m -> x -> forall s . RandST s (Differentials :#: Tangent c m)) -- ^ Gradient calculator
    -> (c :#: m) -- ^ The initial point
    -> RandST r (StochasticPursuit x c m) -- ^ The gradient ascent
boundedStochasticVanillaGradientDescent eps = boundedStochasticVanillaGradientAscent (-eps)

stochasticVanillaGradientDescent :: Manifold m
    => Double -- ^ Step size
    -> (c :#: m -> x -> forall s . RandST s (Differentials :#: Tangent c m)) -- ^ Gradient calculator
    -> (c :#: m) -- ^ The initial point
    -> RandST r (StochasticPursuit x c m) -- ^ The gradient ascent
stochasticVanillaGradientDescent eps = stochasticVanillaGradientAscent (-eps)


--- Mean Squared Error --


meanSquaredError
    :: ((c :#: m) -> [x] -> [Double]) -- ^ Error Function
    -> (c :#: m) -- ^ Current Point
    -> [x] -- ^ Sample points
    -> Double -- ^ Mean squared error
meanSquaredError err p xs =
    let rsdls = err p xs
     in (*0.5) . mean $ (**2) <$> rsdls


--- Cauchy Sequences ---


cauchyLimit :: Manifold m => Int -> Double -> [c :#: m] -> c :#: m
-- | Attempts to calculate the limit of a sequence. This finds the iterate with a
-- sufficiently small 'Euclidean' distance from the previous iterate, or returns
-- the nth iterate.
cauchyLimit n eps ps = last $ cauchySequence n eps ps

cauchySequence :: Manifold m => Int -> Double -> [c :#: m] -> [c :#: m]
cauchySequence n eps ps =
    let ps' = take n ps
        pps' = takeWhile taker . zip ps' $ tail ps'
     in head ps : map snd pps'
       where taker (p1,p2) = let p = alterChart Cartesian $ p1 <-> p2 in eps < sqrt (p <.> p)


-- Least Squares --


designMatrix :: Manifold m => [c :#: m] -> Function (Dual c) Cartesian :#: Tensor Euclidean m
-- | A glorified fromRows operation.
designMatrix rws = matrixTranspose $ coordinateTransform rws

leastSquares :: Manifold m => [c :#: m] -> [Double] -> Dual c :#: m
leastSquares xs ys =
    let mtx = designMatrix xs
        mtxt = matrixTranspose mtx
        prj = matrixInverse (mtxt <#> mtx) <#> mtxt
     in prj >.> euclideanPoint ys

leastSquares0 :: Manifold m => (Function c Cartesian :#: Tensor Euclidean m) -> [Double] -> c :#: m
leastSquares0 mtx ys =
    let mtxt = matrixTranspose mtx
        prj = matrixInverse (mtxt <#> mtx) <#> mtxt
     in prj >.> euclideanPoint ys

-- Newton --

newtonStep :: Manifold m
    => Double -- ^ Step size
    -> (Differentials :#: Tangent c m) -- ^ Derivatives
    -> (Function Partials Differentials :#: Tensor (Tangent c m) (Tangent c m)) -- ^ Hessian
    -> (c :#: m) -- ^ Step
newtonStep eps f' f'' = gradientStep (-eps) $ matrixInverse f'' >.> f'

newtonSequence :: Manifold m
    => Double -- ^ Step Size
    -> (c :#: m -> Differentials :#: Tangent c m) -- ^ Derivatives
    -> (c :#: m -> Function Partials Differentials :#: Tensor (Tangent c m) (Tangent c m)) -- ^ Hessian
    -> (c :#: m) -- ^ Initial point
    -> [c :#: m] -- ^ Newton sequence
newtonSequence eps f' f'' = iterate iterator
  where iterator p = newtonStep eps (f' p) (f'' p)


-- Gauss Newton --

gaussNewtonStep :: Manifold m => Double -> [Double] -> [Differentials :#: Tangent c m] -> c :#: m
gaussNewtonStep eps rs grds = gradientStep (-eps) $ leastSquares0 (designMatrix grds) rs


--- Graveyard ---


{-
gaussNewtonPursuit :: Manifold m
    => Double -- ^ Damping Factor
    -> (c :#: m -> [x] -> [Double]) -- ^ Residual Function
    -> (c :#: m -> [x] -> [Differentials :#: Tangent c m]) -- ^ Residual Differential
    -> (c :#: m) -- ^ Initial guess
    -> StochasticPursuit x c m -- ^ Pursuit
gaussNewtonPursuit dmp residual residuald = accumulateFunction accumulator
  where accumulator xs cm =
            let cm' = gaussNewtonStep dmp (residual cm xs) (residuald cm xs)
             in (cm', cm')
             -}