{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_HADDOCK show-extensions #-}

-- |
-- Module      : Swarm
-- Description : Particle Swarm Optimisation
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental
module PSO.Swarm
  ( -- * Idea
    --
    -- | This module implements different variants of Particle Swarm
    -- Optimisation (PSO), so if you're completely unfamiliar with PSO, go read
    -- https://en.wikipedia.org/wiki/Particle_swarm_optimization first, it's a
    -- nice introduction.
    --
    -- Let there be \(M\in\mathbb{N}\) 'Bee's in the swarm. Bees live in a phase
    -- space (which is a vector space). What this means is that each bee is
    -- described by a phase \(\psi\). The two most common choices of the phase
    -- are:
    --
    -- * \((q, p)\in\mathcal{H}^2\) with \(q\) being the canonical coordinate
    -- and \(p\) being the canonical momentum.
    -- * \(|\psi\rangle\in\mathcal{H}^2\), i.e. a vector in Hilbert space.
    --
    -- While flying around, bees may pick up some information about the
    -- surrounding environment. We call the part of it that bees remember a
    -- local (or private) /guide/. It's a guide because it helps bees decide
    -- where to move next. The simplest and most common guide is 'BeeGuide'.
    --
    -- The whole point of having a swarm of bees rather than one bee is
    -- collective behavior. In case of PSO this is achieved by having a /global
    -- guide/. Each bee then uses this global guide, its own local guide, and
    -- its current phase to compute the next phase. This is represented by
    -- 'PhaseUpdater'.
    --
    -- The whole optimisation process then looks more or less like this:
    --
    --    1. __Initialisation__ (see 'mkSwarm' function).
    --
    --        * Generate initial phases \((\psi^{(0)}_i)_{i\in\{1\dots M\}}\).
    --        * Evaluate fitness in these phases, i.e. compute
    --          \((f(\psi^{(0)}_i))_{i\in\{1\dots M\}}\) (see 'mkBee' function).
    --        * Initialise local guides \((g^{(0)}_{l,i})_{i\in\{1\dots M\}}\)
    --          (see 'mkLocalG' function).
    --        * Construct bees using phases, fitness values, and guides (see
    --          'mkBee' function).
    --        * Initialise global guide \(g^{(0)}_g\) (see 'mkGlobalG'
    --          function).
    --        * Construct swarm using bees and global guide.
    --
    --    2. __Main Loop__. While the termination condition is not met, do the
    --    following (see 'updateSwarm' function):
    --
    --        * Update each bee (see 'updateBee' function), i.e.
    --
    --            * Using global guide \(g^{(n)}_g\), local guide
    --              \(g^{(n)}_{l,i}\), current phase \(\psi^{(n)}_i\), and
    --              \(f(\psi^{(n)}_i)\), come up with a new phase
    --              \(\psi^{(n+1)}_i\) (see 'PhaseUpdater').
    --            * Evaluate the fitness in this new phase, i.e. compute
    --              \(f(\psi^{(n+1)}_i)\).
    --            * Using \(\psi^{(n+1)}_i\) and \(f(\psi^{(n+1)}_i)\), compute
    --              the local guide in the next iteration \(g^{(n+1)}_i\) (see
    --              'updateLocalG' function).
    --
    --        * Update the global guide (see 'updateGlobalG' function).
    --

    -- * Basic types
    Bee(..)
  , CMState(..)
  , QMState(..)
  , LocalGuide(..)
  , BeeGuide(..)
  , GlobalGuide(..)
  , SwarmGuide(..)
  , Swarm(..)
  , PhaseUpdater(..)

    -- * Important functions
  , wpgUpdater
  -- , kickUpdater
  , deltaUpdater
  , absorbUpdater
  , mkCMState
  , mkQMState
  , optimiseND
  , optimiseNDFromList

    -- * Details
  , VelocityUpdater(..)
  , upScale
  , upLocal
  , upGlobal
  , upStep
  , updateBee
  , updateSwarm
  , mkBee
  , mkSwarm
  , fromPure
  , iterateNM
  , iterateWhileM
  , Absorbable(..)
  , DeltaWell(..)
  , Scalable(..)

    -- * Lenses
  , HasPos(..)
  , HasVal(..)
  , HasVar(..)
  , HasGuide(..)
  , HasState(..)
  , HasRun(..)
  , HasBees(..)
  , HasUpdater(..)
  , HasIteration(..)
  ) where

import Data.Complex
import Data.Ord(comparing)
import qualified Data.List as List
-- import Data.Monoid
import Data.Semigroup
import qualified Data.Vector.Generic as GV
import qualified Data.Vector.Storable as V

import Control.Arrow
-- import Control.Lens
import Lens.Micro
import Lens.Micro.TH
import Lens.Micro.Extras
import Control.DeepSeq
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.Reader

import qualified Foreign
import Foreign.Storable.Tuple

import GHC.Float (float2Double)
import GHC.Generics(Generic)

-- import qualified Numeric.LinearAlgebra as LA
-- import Numeric.LinearAlgebra.Devel(matrixFromVector, orderOf, MatrixOrder(..))
-- import Numeric.LinearAlgebra.Data(flatten, tr)
-- import qualified Numeric.LinearAlgebra.Devel as LADevel

import PSO.VectorSpace
import PSO.Random

-- | A bee in our swarm.
--
-- * @α@ is the type of a point in our phase space (i.e. position-velocity pair
-- in the classical case and wave function in the quantum mechanical)
-- * @β@ is the type of the local guide, i.e. \"memory\" of the bee.
-- * @r@ return type of the function we're minimising.
--
-- /Note:/ use 'state', 'guide', and 'val' lenses to access the fields.
data Bee α β r = Bee
  { _beeState :: !α
  , _beeGuide :: !β
  , _beeVal   :: !r
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''Bee

-- | Point in the classical phase space.
--
-- @χ@ -- vector in coordinate space. We always use time steps of 1, so physical
-- dimension of velocity is equal to the one of coordinate. The type of velocity
-- is thus also @χ@.
--
-- /Note:/ use 'pos' and 'vel' lenses to access the fields.
data CMState χ = CMState
  { _cmstatePos :: !χ
  , _cmstateVel :: !χ
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''CMState

-- | Point in our quantum mechanical phase space (i.e. Hilbert space).
--
-- /Note:/ use 'pos' lens to access the field.
data QMState χ = QMState
  { _qmstatePos :: !χ
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''QMState

-- | Rather than limiting ourselves to a single type of local guide (because who
-- knows which info we might want to keep track of), we define a type class.
-- @LocalGuide β α r@ means that @β@ is a local guide for a bee with phase of
-- type @α@ and fitness function returning @r@.
class LocalGuide β α r where
  -- | Initialisation of a guide. Given a phase \(\psi^{(0)}\) and its fitness
  -- \(f(\psi^{(0)})\), @mkLocalG@ creates local guide \(g^{(0)}_l\).
  mkLocalG     :: α -> r -> β
  -- | Update step. After computing the phase \(\psi^{(n+1)}\) (of type @α@) and
  -- its fitness \(f(\psi^{(n+1)})\) (of type @r@) in the @(n+1)@'st iteration,
  -- we use @updateLocalG@ to compute \(g^{(n+1)}_l\) given \(g^{(n)}_l\).
  updateLocalG :: α -> r -> β -> β

-- | Similar to 'LocalGuide' we define a class of global guides rather than a
-- single one. Here, @GlobalGuide γ β@ means that @γ@ can serve as a global
-- guide for local guides of type @β@.
class (LocalGuide β α r) => GlobalGuide γ β α r where
  -- | Initialisation of a guide. Given a collection of bees, @mkLocalG@
  -- initialises global guide \(g^{(0)}_g\).
  mkGlobalG     :: (Foldable t, Functor t) => t (Bee α β r) -> γ
  -- | Update step. After computing phases, values, and local guides of bees in
  -- the @(n+1)@'st iteration, we use @updateLocalG@ to compute \(g^{(n+1)}_g\)
  -- given \(g^{(n)}_g\).
  updateGlobalG :: (Foldable t, Functor t) => t (Bee α β r) -> γ -> γ

-- | A guide that remembers position and its corresponding fitness.
--
-- /Note:/ use 'pos' and 'val' lenses to access the fields.
data BeeGuide χ r = BeeGuide
  { _beeguidePos :: !χ -- ^ Position.
  , _beeguideVal :: !r -- ^ Fitness value.
  } deriving (Generic, Show)

makeLensesWith abbreviatedFields ''BeeGuide

instance (NFData χ, NFData r) => NFData (BeeGuide χ r)

-- | We compare guides by their fitness values.
instance Eq r => Eq (BeeGuide χ r) where
  x == y = (x ^. val) == (y ^. val)

-- | We compare guides by their fitness values.
instance Ord r => Ord (BeeGuide χ r) where
  compare = comparing (view val)

-- | 'BeeGuide' is the simplest and, probably, the most useful example of a
-- 'LocalGuide'. It keeps track of the best position and its fitness. That is,
-- if we've run @n@ iterations (calls to 'updateLocalG'), then the best position
-- is given by \(\operatorname{argmin}_{k\in\{0\dots n\}}f(\psi^{(k)}_i)\), and
-- its fitness is just \(\operatorname{min}_{k\in\{0\dots n\}}f(\psi^{(k)}_i)\).
instance (HasPos α χ, Ord r) => LocalGuide (BeeGuide χ r) α r where
  mkLocalG ψ fψ = BeeGuide (ψ ^. pos) fψ
  updateLocalG ψ fψ g
    | fψ < g ^. val = g & (pos .~ (ψ ^. pos))
                        & (val .~ fψ)
    | otherwise = g

-- | 'BeeGuide' can also serve as a 'GlobalGuide'. It keeps track of the overall
-- best position among all bees and iterations and its fitness. So after @n@
-- iterations, the best position is given by
-- \(\operatorname{argmin}_{k\in\{0\dots n\},i\in\{1\dots M\}}f(\psi^{(k)}_i)\),
-- and its fitness is
-- \(\operatorname{min}_{k\in\{0\dots n\},i\in\{1\dots M\}}f(\psi^{(k)}_i)\).
instance (LocalGuide (BeeGuide χ r) α r, Ord (BeeGuide χ r))
  => GlobalGuide (BeeGuide χ r) (BeeGuide χ r) α r where
    mkGlobalG = minimum . fmap (view guide)
    updateGlobalG xs _ = minimum . fmap (view guide) $ xs

-- | A guide that remembers
--
--    * Best position throughout the whole swarm.
--    * Fitness value at this position.
--    * Current iteration.
--    * Variance in the fitness values of the local guides.
--
-- /Note:/ use 'pos', 'val', 'iteration', and 'var' lenses to access the fields.
data SwarmGuide χ r = SwarmGuide
  { _swarmguidePos       :: !χ
  , _swarmguideVal       :: !r
  , _swarmguideIteration :: !Int
  , _swarmguideVar       :: () -- !r
  } deriving (Generic, Show)

makeLensesWith abbreviatedFields ''SwarmGuide

instance (NFData χ, NFData r) => NFData (SwarmGuide χ r)

-- | Calculates the mean value of a sequence.
mean :: (Fractional a, Foldable v) => v a -> a
mean xs = sum xs / fromIntegral (length xs)

-- | Calculates the variance in a sequence.
variance :: (Fractional a, Foldable v, Functor v) => v a -> a
variance xs = let m = mean xs in mean . fmap (\x -> (x - m)^^2) $ xs

-- | 'SwarmGuide' is another example of a 'GlobalGuide'. On top of keeping track
-- of the best position and its fitness value (what 'BeeGuide' does), we
-- remember the current iteration ('mkGlobalG' sets it to @0@, and each call to
-- 'updateGlobalG' increases the iteration by @1@). And finally, each iteration
-- we compute the variance \(\operatorname{Var}\{f(\psi^{(n)}_i)\}_i\) in the
-- fitness values of bees.
instance (LocalGuide (BeeGuide χ r) α r, Ord (BeeGuide χ r))
  => GlobalGuide (SwarmGuide χ r) (BeeGuide χ r) α r where
    mkGlobalG xs = SwarmGuide (x ^. pos) (x ^. val) 0 var
      where x   = minimum . fmap (view guide) $ xs
            var = () -- variance . fmap (view val) $ xs
    updateGlobalG xs g = mkGlobalG xs & iteration .~ (g ^. iteration + 1)

-- | Given global guide and bee in the @n@'th iteration, calculates the phase of
-- the bee in the @(n+1)@'st iteration. This calculation is done in some 'Monad'
-- @m@ which allows us to, for example, use random numbers when updating the
-- phase.
--
-- /Note:/ use 'run' lens to unpack the function.
newtype PhaseUpdater m γ β α r = PhaseUpdater
  { _phaseupdaterRun :: γ -> Bee α β r -> m α
  }

makeLensesWith abbreviatedFields ''PhaseUpdater

-- | 'PhaseUpdater's can be combined to build more difficult ones.
instance (Monad m) => Semigroup (PhaseUpdater m γ β α r) where
  (PhaseUpdater a) <> (PhaseUpdater b) = PhaseUpdater update
    where update guide bee = a guide bee >>= \ψ ->
            b guide (bee & state .~ ψ)

-- | Our swarm.
data Swarm m γ β α r = Swarm
  { _swarmBees      :: ![Bee α β r] -- ^ Collection of bees.
  , _swarmGuide     :: !γ -- ^ Global guide.
  , _swarmFunc      :: !(α -> m r) -- ^ Fitness function.
  , _swarmUpdater   :: !(PhaseUpdater m γ β α r) -- ^ Updater.
  }

makeLensesWith abbreviatedFields ''Swarm

-- | Updates the canonical momentum in the classical phase. Given a global guide
-- and a bee described by classical mechanics, returns the momentum of the bee
-- in the next iteration. 'VelocityUpdater's can be combined using '<>' operator
-- and afterwards converted to 'PhaseUpdater's using 'upStep' function.
newtype VelocityUpdater m γ χ r = VelocityUpdater
  { _velocityupdaterRun :: γ -> Bee (CMState χ) (BeeGuide χ r) r -> m χ
  }

makeLensesWith abbreviatedFields ''VelocityUpdater

-- | 'VelocityUpdater's can be combined to form more complex ones.
instance (Monad m) => Semigroup (VelocityUpdater m γ χ r) where
  (VelocityUpdater a) <> (VelocityUpdater b) = VelocityUpdater $ update
    where update guide bee =
            a guide bee >>= \v -> b guide (bee & (state . vel) .~ v)

-- | Given a constant \(\varphi_l\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = p^{(n)} + \varphi_l U(0, 1) \circ (g^{(n)}_l - q^{(n)}) \;,
-- \]
-- where \(\circ\) denotes the Hadamard product, \(U(0, 1)\) is a vector
-- whose elements are uniformly distributed in the \([0, 1)\) interval, and
-- \(g^{(n)}_l\) is local best position in the @n@'th iteration.
upLocal :: (RandomScalable m χ, VectorSpace χ λ) => λ -> VelocityUpdater m γ χ r
upLocal φl = VelocityUpdater $ upLocalImpl φl

upLocalImpl ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos α χ, HasVel α χ, HasPos β χ)
  => λ -> γ -> Bee α β r -> m χ
upLocalImpl φ _ bee =
  let δq = bee ^. guide . pos - bee ^. state . pos
      p  = bee ^. state . vel
   in liftM (scale φ) (randScale δq) >>= \δp -> return (p + δp)

-- | Given a constant \(\varphi_g\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = p^{(n)} + \varphi_g U(0, 1) \circ (g^{(n)}_g - q^{(n)}) \;,
-- \]
-- where \(\circ\) denotes the Hadamard product, \(U(0, 1)\) is a vector
-- whose elements are uniformly distributed in the \([0, 1)\) interval, and
-- \(g^{(n)}_g\) is global best position in the @n@'th iteration.
upGlobal ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => λ -> VelocityUpdater m γ χ r
upGlobal φg = VelocityUpdater $ upGlobalImpl φg

upGlobalImpl ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos α χ, HasVel α χ, HasPos γ χ)
  => λ -> γ -> Bee α β r -> m χ
upGlobalImpl φ g bee = do
  let δq = g ^. pos - bee ^. state . pos
      p  = bee ^. state . vel
   in liftM (scale φ) (randScale δq) >>= \δp -> return (p + δp)

-- | Given a constant \(\omega\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = \omega p^{(n)} \;.
-- \]
upScale :: (Monad m, Scalable λ χ) => λ -> VelocityUpdater m γ χ r
upScale ω = VelocityUpdater $ upScaleImpl ω

upScaleImpl :: (Monad m, Scalable λ χ, HasVel α χ) => λ -> γ -> Bee α β r -> m χ
upScaleImpl ω _ = return . scale ω . view (state . vel)

-- | Converts a 'VelocityUpdater' into a 'PhaseUpdater'. This is done by first
-- running the velocity updater to compute \(p^{(n+1)}\). \(q^{(n+1)}\) is then
-- calculated as \(q^{(n)} + p^{(n+1)}\).
upStep ::
     (Monad m, Num χ)
  => VelocityUpdater m γ χ r -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
upStep (VelocityUpdater upVel) = PhaseUpdater update
  where update guide bee = do
          δq <- upVel guide bee
          let x = bee ^. state
              q = x ^. pos
          return $ x & (vel .~ δq)
                     & (pos .~ q + δq)

-- | Creates the standard \"WPG\" updater. Given constants
-- \(\omega,\varphi_l,\varphi_g\), creates a 'PhaseUpdater' that updates the
-- phase according to
-- \[
--  \left\{
--  \begin{aligned}
--    p^{(n+1)} &= \omega p^{(n)}
--              + \varphi_l U_l(0, 1) \circ (g^{(n)}_l - q^{(n)})
--              + \varphi_g U_g(0, 1) \circ (g^{(n)}_g - q^{(n)}) \;, \\
--    q^{(n+1)} &= q^{(n)} + p^{(n+1)} \;,
--  \end{aligned}
--  \right.
-- \]
-- where \(\circ\) denotes the Hadamard product and \(U_l(0, 1)\) and
-- \(U_g(0,1)\) are vectors whose elements are uniformly distributed in the
-- \([0, 1)\) interval.
wpgUpdater ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => (λ, λ, λ) -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
wpgUpdater (ω, φl, φg) =
  upStep $ (upScale ω) <> (upLocal φl) <> (upGlobal φg)

-- | Creates the standard updater which uses constriction rather than inertia
-- parameter. Given constants
-- \(\kappa,\varphi_l,\varphi_g\), creates a 'PhaseUpdater' that updates the
-- phase according to
-- \[
--  \left\{
--  \begin{aligned}
--    p^{(n+1)} &= \kappa \left( p^{(n)}
--              + \varphi_l U_l(0, 1) \circ (g^{(n)}_l - q^{(n)})
--              + \varphi_g U_g(0, 1) \circ (g^{(n)}_g - q^{(n)}) \right) \;, \\
--    q^{(n+1)} &= q^{(n)} + p^{(n+1)} \;,
--  \end{aligned}
--  \right.
-- \]
-- where \(\circ\) denotes the Hadamard product and \(U_l(0, 1)\) and
-- \(U_g(0,1)\) are vectors whose elements are uniformly distributed in the
-- \([0, 1)\) interval.
constrictedUpdater ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => (λ, λ, λ) -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
constrictedUpdater (κ, φl, φg) =
  upStep $ (upLocal φl) <> (upGlobal φg) <> (upScale κ)

-- randomWalkUpdater :: (UniformDist p gen m)
--                   => (p, p) -> Updater m gen s g p r
-- randomWalkUpdater bounds = Updater $ f
--   where f _ _ = do
--           randGen <- ask
--           v' <- lift $ uniform bounds randGen
--           return v'
--
-- kickUpdater ::
--      ( RandomScalable p gen m
--      , UniformDist p gen m
--      , VectorSpace p a
--      , Scalable r p
--      , HasVar s r
--      , HasPos s p
--      , HasPos g p
--      , Floating r
--      , Ord r
--      , s ~ SwarmGuide p r
--      )
--   => (r, r) -> (p, p) -> (a, a, a) -> Updater m gen s g p r
-- kickUpdater cutoff bounds wpg = Updater $ kickUpdaterImpl cutoff bounds wpg
--
-- kickUpdaterImpl ::
--      ( RandomScalable p gen m
--      , UniformDist p gen m
--      , VectorSpace p a
--      , Scalable r p
--      , HasVar s r
--      , HasPos s p
--      , HasPos g p
--      , Floating r
--      , Ord r
--      )
--   => (r, r) -> (p, p) -> (a, a, a) -> s -> Bee p r g -> ReaderT gen m p
-- kickUpdaterImpl (cutoff, c) bounds wpg stats x = runUpdater updater stats x
--   where updater = if (stats ^. var) < cutoff
--                      then randomWalkUpdater bounds
--                           <> upScale (c * sqrt (stats ^.var))
--                      else standardUpdater wpg

-- | Update the position according to QDPSO (Quantum Delta well Particle Swarm
-- Optimisation).
class (Monad m) => DeltaWell m λ χ where
  upDeltaWell :: λ -> χ -> χ -> m χ

-- | 1D QDPSO updater. Given a constant \(\kappa\in\mathbb{R}\), center
-- \(g^{(n)}\in\mathbb{R}\) of the density distribution, and bees current
-- position \(q^{(n)}\in\mathbb{R}\), its position in the next iteration is
-- computed as
-- \[
--  q^{(n+1)} =
--    q^{(n)} + \sgn(2\cdot\operatorname{rand}(0,1) - 1) \kappa
--                \log(\frac{1}{\operatorname{rand}(0, 1)}) \cdot
--                  |q^{(n)} - g^{(n)}| \;.
-- \]
instance (Monad m, RealFloat λ, Ord λ, Randomisable m λ)
  => DeltaWell m λ λ where
    upDeltaWell κ p x = do
      u    <- random
      sign <- (\c -> signum (2 * c - 1)) <$> random
      return $ p + (sign * (log . recip $ u) / κ) * abs (x - p)

-- | Extending QDPSO to \(\mathbb{C}\). 1D real version is simply applied to
-- both real and imaginary parts.
instance (RealFloat λ, DeltaWell m λ λ)
  => DeltaWell m λ (Complex λ) where
    upDeltaWell κ p x =
      let f getter = uncurry (upDeltaWell κ) . over each getter
       in return (:+) `ap` f realPart (p, x)
                      `ap` f imagPart (p, x)

-- | Extending QDPSO to multidimensional case. 1D version is simply applied to
-- each component.
instance (Monad m, Foreign.Storable ξ, DeltaWell m λ ξ)
  => DeltaWell m λ (V.Vector ξ) where
    upDeltaWell κ p x = V.zipWithM (upDeltaWell κ) p x

{-
liftMatrix2M ::
     (Monad m, LA.Element α, LA.Element β, LA.Element γ,
      LA.Container V.Vector α, Num α,
      LA.Container V.Vector β, Num β,
      LA.Transposable (LA.Matrix α) (LA.Matrix α),
      LA.Transposable (LA.Matrix β) (LA.Matrix β))
  => (LA.Vector α -> LA.Vector β -> m (LA.Vector γ))
  -> LA.Matrix α -> LA.Matrix β -> m (LA.Matrix γ)
liftMatrix2M f m1@(LA.size -> (r,c)) m2
  | LA.size m1 /= LA.size m2 = error "nonconformant matrices in liftMatrix2M"
  | orderOf m1 == RowMajor =
      return (matrixFromVector RowMajor r c) `ap` f (flatten m1) (flatten m2)
  | otherwise =
      return (matrixFromVector ColumnMajor r c) `ap`
        f (flatten . tr $ m1) (flatten . tr $ m2)

instance (Monad m, Num ξ, LA.Container V.Vector ξ, LA.Transposable (LA.Matrix ξ)(LA.Matrix ξ), DeltaWell m λ ξ)
  => DeltaWell m λ (LA.Matrix ξ) where
    upDeltaWell κ = liftMatrix2M (V.zipWithM (upDeltaWell κ))
-}

-- | Creates a QDPSO updater.
deltaUpdater :: forall m g γ β χ λ r.
     ( Randomisable m λ
     , RealFloat λ
     , Ord λ
     , VectorSpace χ λ
     , DeltaWell m λ χ
     , HasPos β χ
     , HasPos γ χ
     )
  => λ -> γ -> Bee (QMState χ) β r -> m (QMState χ)
deltaUpdater κ g bee = do
  φ₁ <- random :: m λ
  φ₂ <- random :: m λ
  -- lift $ putStrLn $ "φ₁ = " ++ show φ₁ ++ ", φ₂ = " ++ show φ₂
  let center = (φ₁ / (φ₁ + φ₂)) `scale` (g ^. pos)
             + (φ₂ / (φ₁ + φ₂)) `scale` (bee ^. guide . pos)
  x <- upDeltaWell κ center (bee ^. state . pos)
  -- lift $ putStrLn $ "p - x = " ++ show (bee ^. state . pos - p)
  return $ (bee ^. state) & (pos .~ x)

data RelativePosition = Inside | LeftOf | RightOf

class RelativelyPositioned α β where
  getRelPos :: α -> β -> RelativePosition

instance (Ord λ) => RelativelyPositioned (λ, λ) λ where
  getRelPos (l, h) x
    | x < l     = LeftOf
    | x > h     = RightOf
    | otherwise = Inside

class Absorbable λ α where
  absorb :: (λ, λ) -> α -> α

class Reflectable λ α where
  reflect :: (λ, λ) -> α -> α

instance (Ord χ) => Absorbable χ χ where
  absorb (l, h) x
    | x < l     = l
    | x > h     = h
    | otherwise = x

instance (Ord λ) => Absorbable λ (Complex λ) where
  absorb bs = over each (absorb bs)

instance (Ord λ, Num λ)
  => Absorbable λ (λ, λ) where
    absorb (l, h) (q, p) = case getRelPos (l, h) q of
                             Inside  -> (q, p)
                             LeftOf  -> (l, 0)
                             RightOf -> (h, 0)

zipWithTuple :: (a -> b -> c) -> (a, a) -> (b, b) -> (c, c)
zipWithTuple f (x₁, x₂) (y₁, y₂) = (f x₁ y₁, f x₂ y₂)

instance (Ord λ, Num λ)
  => Absorbable λ (Complex λ, Complex λ) where
    absorb bounds ψ =
      let f getter = absorb bounds . over each getter
       in zipWithTuple (:+) (f realPart ψ) (f imagPart ψ)

instance (Absorbable λ (ξ, ξ), GV.Vector χ ξ, GV.Vector χ (ξ, ξ))
  => Absorbable λ (CMState (χ ξ)) where
    absorb bounds ψ =
      let (q, p) = GV.unzip . GV.map (absorb bounds)
                 $ GV.zip (ψ ^. pos) (ψ ^. vel)
       in ψ & (pos .~ q) & (vel .~ p)

instance (Absorbable λ ξ, GV.Vector χ ξ)
  => Absorbable λ (QMState (χ ξ)) where
    absorb bounds = pos %~ GV.map (absorb bounds)

absorbUpdater :: (Monad m, Absorbable λ α) => (λ, λ) -> PhaseUpdater m γ β α r
absorbUpdater bounds = PhaseUpdater $ upAbsorbImpl bounds

upAbsorbImpl :: (Monad m, Absorbable λ α) => (λ, λ) -> γ -> Bee α β r -> m α
upAbsorbImpl bounds _ = return . absorb bounds . view state

-- | Given a swarm and bee in the @n@'th iteration, returns a computation that
-- returns the bee in the @(n+1)@'st iteration.
updateBee ::
     (Monad m, LocalGuide β α r)
  => Swarm m γ β α r -> Bee α β r -> m (Bee α β r)
updateBee xs x = do
  ψ'  <- (xs ^. updater . run) (xs ^. guide) x
  fψ' <- (xs ^. func) ψ'
  return $ x & (state .~ ψ')
             & (val   .~ fψ')
             & (guide %~ updateLocalG ψ' fψ')

-- | Given a swarm in the @n@'th iteration, returns a computation that returns
-- the swarm in the @(n+1)@'st iteration.
updateSwarm :: (Monad m, GlobalGuide γ β α r)
            => Swarm m γ β α r -> m (Swarm m γ β α r)
updateSwarm swarm = do
  xs <- mapM (updateBee swarm) (swarm ^. bees)
  return $ swarm & (bees .~ xs)
                 & (guide %~ updateGlobalG xs)

-- | Applies \(f\) to \(x\) \(\max(0, n - 1)\) times and accumulates the
-- results. Returned list will have \(n\) elements.
iterateNM ::
     (Monad m)
  => Int -- ^ \(n\), number of times to apply \(f\)
  -> (a -> m a) -- ^ function \(f\) to apply
  -> a -- ^ initial value \(x\)
  -> m [a]
iterateNM n f x
  | n == 0 = return [x]
  | otherwise = do
    y <- f x
    ys <- iterateNM (n - 1) f y
    return $ x : ys

iterateWhileM ::
     (Monad m)
  => (a -> Bool)
  -> (a -> m a) -- ^ function \(f\) to apply
  -> a -- ^ initial value \(x\)
  -> m [a]
iterateWhileM predicate func x
  | predicate x = do
      y  <- func x
      ys <- iterateWhileM predicate func y
      return $ x : ys
  | otherwise   = return [x]

mkCMState :: (UniformDist m χ, Num χ) => (χ, χ) -> m (CMState χ)
mkCMState bounds = liftM (\q -> CMState q 0) $ uniform bounds

mkQMState :: (UniformDist m χ, Num χ) => (χ, χ) -> m (QMState χ)
mkQMState bounds = liftM (\q -> QMState q)   $ uniform bounds

-- | @mkSwarm updater initialiser inside func n@ creates a new 'Swarm' where
--
-- * @updater@ is the update strategy used to obtain velocities of bees in the
-- next iteration.
-- * @initialiser@ creates a new random 'Bee'. It is assumed that position of a
-- bee created using @initialiser@ will satisfy @inside@.
-- * Given a position \(x\in\mathcal{H}\) returns whether \(x\in\mathcal{V}\)
-- holds.
-- * @func@ is the function \(f\) we're trying to minimise.
-- * @n@ is number of bees in the newly created swarm.
mkSwarm ::
     (Monad m, GlobalGuide γ β α r)
  => m α
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> Int
  -> m (Swarm m γ β α r)
mkSwarm newState updater func n = (sequence . replicate n $ newState) >>=
  mapM (mkBee func) >>= \xs ->
    return $ Swarm xs (mkGlobalG xs) func updater

mkSwarmFromList ::
     (Monad m, GlobalGuide γ β α r)
  => [α]
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> m (Swarm m γ β α r)
mkSwarmFromList states updater func =
  mapM (mkBee func) states >>= \xs ->
    return $ Swarm xs (mkGlobalG xs) func updater

fromPure :: (Monad m, HasPos α χ) => (χ -> r) -> α -> m r
fromPure f ψ = return $ f (ψ ^. pos)
--
-- class Between a b where
--   isBetween :: (a, a) -> b -> Bool
--
-- instance (Ord a) => Between a a where
--   isBetween (l, h) x = l <= x && x <= h
--
-- instance {-# OVERLAPS #-} (Between a a) => Between (Complex a) (Complex a) where
--   isBetween (rl :+ il, rh :+ ih) (rx :+ ix) =
--     isBetween (rl, rh) rx && isBetween (il, ih) ix
--
-- instance {-# OVERLAPS #-} (Foreign.Storable a, Between a a)
--   => Between (V.Vector a) (V.Vector a) where
--     isBetween (low, high) x =
--       V.and $ V.zipWith3 (\l h x -> isBetween (l, h) x) low high x
--
-- -- | Creates a new 'Bee' with a random position inside a given 1D interval and
-- -- zero velocity.
-- randomBee1D ::
--      (Num p, UniformDist p gen m)
--   => (p, p) -- ^ @[low, high]@ interval. It is assumed that @low <= high@.
--   -> (p -> ReaderT gen m r) -- ^ Function we're minimising.
--   -> ReaderT gen m (Bee p r (BeeGuide p r))
-- randomBee1D bounds func = do
--   g <- ask
--   x <- lift $ uniform bounds g
--   fx <- func x
--   return $ Bee x 0 fx (BeeGuide x fx)

-- | Given a fitness function \(f\) and phase \(\psi^{(0)}\), creates a new
-- 'Bee'.
mkBee :: (Monad m, LocalGuide β α r)
  => (α -> m r) -> α -> m (Bee α β r)
mkBee f ψ = do
  fψ <- f ψ
  return $ Bee ψ (mkLocalG ψ fψ) fψ

-- -- | Simplified @mkSwarm@ for 1D pure functions.
-- -- simpleInit1D ::
-- --      ( Num a
-- --      , RandomScalable a gen m
-- --      , UniformDist a gen m
-- --      , Between a a
-- --      , Ord r
-- --      , s ~ g
-- --      , g ~ BeeGuide a r
-- --      )
-- --   => (a, a, a) -- ^ WPG parameters \((\omega,\varphi_l,\varphi_g)\).
-- --   -> (a, a) -- ^ Boundaries \([a, b]\).
-- --   -> (a -> r) -- ^ Function \(f\) we're minimising.
-- --   -> Int -- ^ Number of bees.
-- --   -> ReaderT gen m (Swarm m gen s g a r)
-- -- simpleInit1D wpg (low, high) func =
-- --   let updater = standardUpdater wpg
-- --       initialiser = randomBee1D (low, high)
-- --       boundaries = isBetween (low, high)
-- --   in mkSwarm updater initialiser boundaries (fromPure func)
--
--
-- unpackToV :: (Foreign.Storable a) => [(a, a)] -> (V.Vector a, V.Vector a)
-- unpackToV xs = let a = V.fromList . map fst $ xs
--                    b = V.fromList . map snd $ xs
--                 in (a, b)
--
-- -- | Simplified @mkSwarm@ for N-dimensional functions.
-- simpleInitND ::
--      ( Foreign.Storable a
--      , VectorSpace p a
--      , RandomScalable p gen m
--      , UniformDist p gen m
--      , Between p p
--      , Ord r
--      , p ~ V.Vector a
--      , GlobalGuide s g p r
--      )
--   => Updater m gen s g p r -- ^ PWG parameters \((\omega,\varphi_l,\varphi_g)\)
--   -> [(a, a)]
--   -> [(a, a)] -- ^ List of intervals \([a_i,b_i]\). \(\mathcal{V}\) is then
--   -- \(\prod_i [a_i,b_i]\).
--   -> (p -> r) -- ^ Function \(f\) we're minimising.
--   -> Int -- ^ Number of bees.
--   -> ReaderT gen m (Swarm m gen s g p r)
-- simpleInitND updater initBounds bounds func =
--   let initialiser = randomBeeND (unpackToV initBounds)
--       boundaries = isBetween (unpackToV bounds)
--   in mkSwarm updater initialiser boundaries (fromPure func)
--
-- -- optimise1D ::
-- --      ( Num a
-- --      , RandomScalable a gen m
-- --      , UniformDist a gen m
-- --      , Between a a
-- --      , Ord r
-- --      , s ~ g
-- --      , g ~ BeeGuide a r
-- --      )
-- --   => (a, a, a)
-- --   -> (a, a)
-- --   -> (a -> r)
-- --   -> Int
-- --   -> (Swarm m gen s g a r -> Bool)
-- --   -> ReaderT gen m [Swarm m gen s g a r]
-- -- optimise1D wpg bounds func n predicate = do
-- --   swarm <- simpleInit1D wpg bounds func n
-- --   iterateWhileM (not . predicate) updateSwarm swarm

optimiseND ::
     ( Monad m
     , HasPos α χ
     , GlobalGuide γ β α r
     )
  => m α
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> Int
  -> (Swarm m γ β α r -> Bool)
  -> m [Swarm m γ β α r]
optimiseND newState updater func n predicate = do
  swarm <- mkSwarm newState updater func n
  iterateWhileM (not . predicate) updateSwarm swarm

optimiseNDFromList ::
     ( Monad m
     , HasPos α χ
     , GlobalGuide γ β α r
     )
  => [α]
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> (Swarm m γ β α r -> Bool)
  -> (Swarm m γ β α r -> m ())
  -> m (Swarm m γ β α r)
optimiseNDFromList states updater func predicate process =
  do
    swarm <- mkSwarmFromList states updater func
    go swarm
  where
        go x
          | predicate x = return x
          | otherwise   = process x >> updateSwarm x >>= go


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE TypeApplications      #-}

module Main where

import Control.Monad.Primitive

import qualified Data.Vector.Sized as VS
--import Data.Vector.Algorithms.Intro as VAI

--import Neldermead
import Data.Vector.GDP.Sorted
import ObjRSimplex
import Numeric.LinearAlgebra.Static

import qualified GDP

simplexEx1 :: VSimplex R 2
simplexEx1 = VS.fromTuple (v1, v2, v3)
  where
    (v1 :: R 2) = vector [0.0, 3.0]
    (v2 :: R 2) = vector [1.0, 4.0]
    (v3 :: R 2) = vector [2.0, 0.0]

createObjSimplex :: (ObjRF n GDP.~~ fObjName)
  -> VSimplex R n
  -> ObjRSimplex fObjName n
createObjSimplex f = ObjRSimplex . VS.map (makeFObjRVtx f)

sortObjSimplex :: forall fObjName n m. PrimMonad m
  => (CmpObjRVtx fObjName n GDP.~~ RVtxObjCmp fObjName)
  -> ObjRSimplex fObjName n
  -> m (SortedObjRSimplex fObjName n)
sortObjSimplex cmp smplx = do
  v <- VS.thaw $ vtxs (smplx :: ObjRSimplex fObjName n)
  SortedObjRSimplex <$> sortByAndFreeze cmp v

main :: IO ()
main = do
  GDP.name norm_1 $ \f -> do
    let smplx = createObjSimplex f simplexEx1
        c = getRVtxObjCmp f
    (SortedObjRSimplex v) <- sortObjSimplex c smplx
    putStrLn . show . VS.map (GDP.the . getObjRVtx) $ GDP.the v


module Gosadenban
( main3

) where

-------------------------------------------------------------
--module
-------------------------------------------------------------

import Numeric.LinearAlgebra
import System.Random
import qualified Data.ByteString.Lazy as B
import Data.List
import Data.Time

import Tools
import Durstenfeld
import ForMNIST

------------------------------------------------

data ActivationFunction = ActivationFunction
    { func  :: Double -> Double ,
      func' :: Double -> Double ,
      description :: String
    }

data Params = Params
    { weight :: Matrix Double ,
      bias :: Matrix Double
    } deriving (Show)

data DParams = DParams
    { dweight :: Matrix Double ,
      dbias :: Matrix Double
    } deriving (Show)

data LayerProp = LayerProp
    { params :: Params ,
      af :: ActivationFunction
    }

data HyperProp = HyperProp
    { vw0 :: Matrix Double ,
      vb0 :: Matrix Double
    }

data Hyperparameter = Hyperparameter
    { learningrate :: Double ,
      alpha0 :: Double ,
      beta0  :: Double ,
      beta1  :: Double
    }

data Network = Network
    { layers :: [LayerProp] ,
      hprop  :: [HyperProp] ,
      hpara  :: Hyperparameter
    }

data PropagateLayer
    = PropagateLayer
        { inData  :: Matrix Double ,
          outData :: Matrix Double ,
          pAf :: ActivationFunction ,
          pParams :: Params
        }
    | PropagateSensorLayer
        { outData :: Matrix Double
        }

data BackPropagateLayer
    = BackPropagateLayer
        { inGrad  :: Matrix Double ,
          outGrad :: Matrix Double ,
          bpAf :: ActivationFunction ,
          bpParams :: Params ,
          dParams :: DParams
        }
    | BackPropagateLastLayer
        { outGrad :: Matrix Double
        }

data ErrorLayer
    = ErrorLayer
        { lastData :: Matrix Double ,
          teachData :: Matrix Double ,
          errorValue :: Double ,
          grad :: Matrix Double
        } deriving (Show)

-------------------------------------------------------------

propagate :: PropagateLayer -> LayerProp -> PropagateLayer
propagate dataJ layerK = PropagateLayer
    { inData  = x ,
      outData = y ,
      pAf     = functions ,
      pParams = p
    }
    where x = outData dataJ
          p = params layerK
          w = weight p
          b = bias p
          a = add b $ x Numeric.LinearAlgebra.<> w
          functions =  af layerK
          h = func functions
          y = cmap h a

propagateNetwork :: Matrix Double -> Network -> [PropagateLayer]
propagateNetwork input net = calcs
    where calcs = scanl propagate layer0 (layers net)
          layer0 = PropagateSensorLayer { outData = input}

-------------------------------------------------------------

backpropagate :: PropagateLayer -> BackPropagateLayer -> BackPropagateLayer
backpropagate pLayerJ bpLayerK = BackPropagateLayer
    { inGrad = indx ,
      outGrad = dx ,
      bpAf = af ,
      bpParams  = p ,
      dParams = dP
    }
    where af = pAf pLayerJ
          p = pParams pLayerJ
          h' = func' af
          indx = outGrad bpLayerK
          y = outData pLayerJ
          dL' = cmap h' y
          dL = indx * dL'
          db = dL
          x = inData pLayerJ
          w = weight p
          dw = (tr x) Numeric.LinearAlgebra.<> dL
          dx = dL Numeric.LinearAlgebra.<> (tr w)
          dP = DParams { dweight = dw, dbias = db}

backpropagateNetwork :: ErrorLayer -> [PropagateLayer] ->[BackPropagateLayer]
backpropagateNetwork el prL = tail ans1
    where ans1 = scanr backpropagate layerFin prL
          layerFin = BackPropagateLastLayer {outGrad = grad0}
          grad0 = grad el

-------------------------------------------------------------

--Softmax & cross entropy error
errorNetwork :: Matrix Double -> Matrix Double -> ErrorLayer
errorNetwork ld td = ErrorLayer
    { lastData = dataFromSoftmaxFunction ,
      teachData = td ,
      errorValue = e / len ,
      grad = g
    }
    where dataFromSoftmaxFunction = softmaxBatch' ld
          g = add dataFromSoftmaxFunction $ scale (-1) td
          y = map asRow $ toRows dataFromSoftmaxFunction
          t = map asRow $ toRows td
          e' = zipWith cee t y
          e = sum e'
          len = length' t

-------------------------------------------------------------

inference :: Matrix Double -> Int -> Double
inference x t = ans
    where (r, c) = maxIndex x
          ans = if c == t then 1.0 else 0.0

accuracyNet :: Network -> Matrix Double -> Int -> Double
accuracyNet net x t = a
    where ans' = propagateNetwork x net
          ans  = last $ map outData ans'
          a    = inference ans t

accuracyNetwork :: Network -> [Matrix Double] -> [Int] -> Double
accuracyNetwork net xs ts = sumans / lenans
    where f x t = accuracyNet net x t
          ans = zipWith f xs ts
          sumans = sum ans
          lenans = length' xs

----------------------------------------------------------------------
sgd :: Matrix Double -> Matrix Double -> Hyperparameter -> Matrix Double
sgd x dx hp = x - scale eta dx
    where eta = learningrate hp

update :: Hyperparameter -> Params -> DParams -> Params
update hp p dp = newParam
    where b  = bias p
          w  = weight p
          db = dbias dp
          dw = dweight dp
          new_b = sgd b db hp
          new_w = sgd w dw hp
          newParam = Params {bias = new_b, weight = new_w}

updateNetowrk :: Network -> [DParams] -> Network
updateNetowrk net dp = new_Network
    where layerList = layers net
          hp = hpara net
          hpp = hprop net
          paramList = map params layerList
          afList = map af layerList
          f = update hp
          new_Para = zipWith f paramList dp
          new_Layer = zipWith makeLayerProp new_Para afList
          new_Network = Network{layers = new_Layer, hpara = hp, hprop = hpp  }

momentum :: Matrix Double -> Matrix Double -> Hyperparameter -> Matrix Double -> ( Matrix Double, Matrix Double)
momentum x dx hp v = (x + new_v, new_v)
    where eta = learningrate hp
          alpha = alpha0 hp
          new_v = (scale alpha v) - scale eta dx

adagrad :: Matrix Double -> Matrix Double -> Hyperparameter -> Matrix Double -> ( Matrix Double, Matrix Double)
adagrad x dx hp v = (new_x, new_v')
    where eta = learningrate hp
          alpha = alpha0 hp
          new_v' = v + dx*dx
          new_v = cmap sqrt new_v'
          new_v'' = cmap devdev new_v
          new_x = x - scale eta ( dx/new_v'' )

devdev :: Double -> Double
devdev x =
    if x == 0
        then 1.0
        else x

update' :: Hyperparameter -> (HyperProp , Params) -> DParams -> (Params, HyperProp)
update' hp (hpp, p) dp = (newParam, newHpp)
    where b  = bias p
          w  = weight p
          db = dbias dp
          dw = dweight dp
          vw = vw0 hpp
          vb = vb0 hpp
          (new_w, new_vw) = momentum w dw hp vw
          (new_b, new_vb) = momentum b db hp vb
          newParam = Params {bias = new_b, weight = new_w}
          newHpp = HyperProp {vw0 = new_vw, vb0 = new_vb}

updateNetowrk' :: Network -> [DParams] -> Network
updateNetowrk' net dp = new_Network
    where layerList = layers net
          hp = hpara net
          hpp = hprop net
          paramList = map params layerList
          afList = map af layerList
          f = update' hp
          hppp = zip hpp paramList
          new_Para_And_Hpp = zipWith f hppp dp
          new_Para = map fst new_Para_And_Hpp
          new_Hpp  = map snd new_Para_And_Hpp
          new_Layer = zipWith makeLayerProp new_Para afList
          new_Network = Network{layers = new_Layer, hpara = hp, hprop = new_Hpp  }

----------------------------------------------------------------------

makeLayerProp :: Params -> ActivationFunction -> LayerProp
makeLayerProp p aF = LayerProp {params = p, af = aF }

----------------------------------------------------------------------

loadL :: B.ByteString -> [Matrix Double]
loadL xs = tL4
    where tLabel  = loadLabel xs
          f x = pickUpNum x tLabel
          tL1 = map f batch_List
          tL2 = map ( map oneHotEncoder ) tL1
          tL3 = map ( map flatten ) tL2
          tL4 = map fromRows tL3

loadI :: B.ByteString -> [Matrix Double]
loadI xs = tI4
    where tImage  = loadImage xs
          g x = pickUpNum x tImage
          tI1 = map g batch_List
          tI2 = map ( map (scale (1/255.0) )) tI1
          tI3 = map ( map flatten ) tI2
          tI4 = map fromRows tI3

----------------------------------------------------------------------

finPara :: Matrix Double -> Matrix Double -> Params
finPara b w = Params {bias = b, weight = w}

finLayer :: Params -> ActivationFunction -> LayerProp
finLayer p a = LayerProp {params = p, af = a }

----------------------------------------------------------------------

sigmoidFunction :: Floating a => a -> a
sigmoidFunction x = 1 / (1 + exp (-x) )

sigmoidFunctionBack :: Floating a => a -> a
sigmoidFunctionBack  x = x*(1-x)

sigmoidFunc :: ActivationFunction
sigmoidFunc = ActivationFunction
    { func  = sigmoidFunction,
      func' = sigmoidFunctionBack,
      description = "Sigmoid Function"
    }

reluFunction :: Double -> Double
reluFunction x
    | x <= 0 = 0.0
    | x > 0  = x

reluFunctionBack :: Double -> Double
reluFunctionBack x
    | x <= 0 = 0.0
    | x > 0  = 1.0

reluFunc :: ActivationFunction
reluFunc = ActivationFunction
    { func  = reluFunction,
      func' = reluFunctionBack,
      description = "Relu Function"
    }

identity :: ActivationFunction
identity = ActivationFunction
    { func  = id,
      func' = backOne,
      description = "id"
    }

backOne :: Double -> Double
backOne x = 1.0

-------------------------------------------------------------

softmax :: (Linear t c, Floating t, Container c t) => c t -> c t
softmax m = scale  (1/bunbo) newM
    where maxValue = maxElement m
          newM     = cmap  (\x -> exp $ x - maxValue ) m
          bunbo    = sumElements newM

softmaxBatch' :: Matrix Double -> Matrix Double
softmaxBatch' m  =  m1
    where list  = toRows m
          list1 = fmap softmax list
          m1 = fromRows list1

cee :: Matrix Double -> Matrix Double -> Double
cee t y = (-1.0)*e
    where delta = 0.0000001
          y1 = cmap (\x -> log (x + delta) ) y
          y2 = flatten y1
          y3 = diag y2
          e = sumElements (t Numeric.LinearAlgebra.<> y3)

----------------------------------------------------------------------

checkIt (e0, net) (x ,t) = (e, newNetwork)
    where ans' = propagateNetwork x net
          ans  = map outData ans'
          answer  = last ans
          answer2 = errorNetwork answer t
          e = errorValue answer2
          finalAnswer' = backpropagateNetwork answer2 ans'
          finalA = map dParams $ init finalAnswer'
          newNetwork = updateNetowrk' net finalA

checking net0 lt = ans
    where ans = scanl checkIt net0 lt

checking' net0 lt = ans
    where ans = foldl checkIt net0 lt

-------------------------------------------------------------
--Parameter
-------------------------------------------------------------

data_size = 10000
batch_size = 100
trials_num = 2000

input_size  = 784
hidden_size = 50
output_size = 10

------------------------------------------------

-------------------------------------------------------------
--Initial Condition (?)
-------------------------------------------------------------

nw1 = randomMatrix input_size  hidden_size (mkStdGen 101) :: (Matrix Double)
nw2 = randomMatrix hidden_size output_size (mkStdGen 102) :: (Matrix Double)

nb1 = matrix hidden_size $ take (batch_size * hidden_size) $ repeat 0.0
nb2 = matrix output_size $ take (batch_size * output_size) $ repeat 0.0

nv_w1 =  matrix hidden_size $ take (input_size  * hidden_size) $ repeat 0.0
nv_w2 =  matrix output_size $ take (hidden_size * output_size) $ repeat 0.0

nv_b1 =  matrix hidden_size $ take (batch_size * hidden_size) $ repeat 0.0
nv_b2 =  matrix output_size $ take (batch_size * output_size) $ repeat 0.0

------------------------------------------------

param1 = Params {bias = nb1, weight = nw1}
layer1 = LayerProp {params = param1, af = reluFunc}
hprop1 = HyperProp { vw0 = nv_w1, vb0 = nv_b1 }

param2 = Params {bias = nb2, weight = nw2}
layer2 = LayerProp {params = param2, af = identity }
hprop2  = HyperProp { vw0 = nv_w2, vb0 = nv_b2}

hyperParam = Hyperparameter {learningrate = 0.0001, alpha0 = 0.9, beta0 = 0.0, beta1 = 0.0}

networkA = Network {layers = [layer1,layer2], hpara = hyperParam, hprop = [hprop1,hprop2]}

--makeBatchList gen num_of_trials data_size batch_size
batch_List = makeBatchList 100000 trials_num data_size batch_size

main3 = do

    z <- getZonedTime
    print z

    x <- getCurrentTime

    contents1 <- B.readFile "MNIST/train-labels.idx1-ubyte"
    contents2 <- B.readFile "MNIST/train-images.idx3-ubyte"

    let tL4 = loadL contents1
        --tL4 = [[Batch X 10],[Batch X 10],...] :: [Matrix Double]

        tI4 = loadI contents2
        --tI4 = [[Batch X data_size],[Batch X data_size],...] :: [Matrix Double]

        test_Image_and_Label = zip tI4 tL4

        --e1 = checking' (0.0,networkA) test_Image_and_Label
        --e2 =  fst e1

        e1' =  checking (0.0,networkA) test_Image_and_Label
        e2' =  tail $ map fst e1'
        net1' = last $ map snd e1'

    --print e2
    writeFile "logs/e_mmt-0.0001.dat" $ encorderForDat2 e2'

    contents3 <- B.readFile "MNIST/t10k-labels.idx1-ubyte"
    contents4 <- B.readFile "MNIST/t10k-images.idx3-ubyte"

    let xs = loadImageTest contents4
        ts = loadLabelTest contents3
        --a = accuracyNet net1' xs ts
        --acc = accuracyNetwork net1' xs ts
        bs = map toRows $ map bias $ map params $ layers net1'
        bs' = map (foldl1 add) bs
        bsl = map recip $ map length' bs
        newbs = map asRow $ zipWith scale bsl bs'
        w = map weight $ map params $ layers net1'
        para = zipWith finPara newbs w
        acfunc = map af $ layers net1'
        newLayer = zipWith finLayer para acfunc
        finNetwork = Network {layers = newLayer, hpara = Hyperparameter {learningrate = 1.0, alpha0 = 0.0, beta0 = 0.0, beta1 = 0.0},hprop = [hprop1,hprop2]}

        acc = accuracyNetwork finNetwork xs ts

    print acc

    y <- getCurrentTime
    print$ diffUTCTime y x

----------------------------------------------------------------------
----------------------------------------------------------------------
----------------------------------------------------------------------
----------------------------------------------------------------------

loopCheck x net t = (e, finalA, newNetwork)
    where ans' = propagateNetwork x net
          ans  = map outData ans'
          answer  = last ans
          answer2 = errorNetwork answer t
          e = errorValue answer2
          finalAnswer' = backpropagateNetwork answer2 ans'
          finalA = map dParams $ init finalAnswer'
          newNetwork = updateNetowrk net finalA

loop _  _  _ 0 = []
loop x network t n = e : loop x newNetwork t (n-1)
    where  (e, finalA, newNetwork) = loopCheck x network t

randomList = randomChoice 10 10 (mkStdGen 10 )

----------------------------------------------------------------------


{-# LANGUAGE FunctionalDependencies #-}

module Distribution (
    Dist, pdf, logpdf, samplefrom,
    Uniform(..), Normal(..), Bernoulli(..),
    Exponential(..), Gamma(..), Poisson(..),
    Binomial(..), Beta(..),
    Prob, Weight, Weighted
) where

import Randomness ( Random(..), uniform, Prob, Weight, Weighted )
import Numeric.Log ( Log(..) )
import qualified Data.Random.Distribution.Uniform as R ( uniform )
import qualified Data.Random.Distribution.Normal as R ( normal )
import qualified Data.Random.Distribution.Bernoulli as R ( bernoulli )
import qualified Data.Random.Distribution.Binomial as R ( binomial )
import qualified Data.Random.Distribution.Poisson as R ( poisson )
import qualified Statistics.Distribution as S ( ContDistr(density, logDensity, quantile), DiscreteDistr(probability, logProbability) )
import qualified Statistics.Distribution.Uniform as S ( uniformDistr )
import qualified Statistics.Distribution.Normal as S ( normalDistr )
import qualified Statistics.Distribution.Binomial as S ( binomial )
import qualified Statistics.Distribution.Poisson as S ( poisson )
import qualified Statistics.Distribution.Gamma as S ( gammaDistr )
import qualified Statistics.Distribution.Exponential as S ( exponential )
import qualified Statistics.Distribution.Binomial as S ( binomial )
import qualified Statistics.Distribution.Beta as S ( betaDistr )
import qualified Statistics.Distribution.DiscreteUniform as S ( discreteUniformAB )

class Dist d a | d -> a where
    samplefrom :: d -> Random a

    pdf :: d -> a -> Double
    pdf d = exp . ln . logpdf d

    logpdf :: d -> a -> Weight
    logpdf d = Exp . log . pdf d

data Uniform = Uniform Double Double
-- data DiscreteUniform = DiscreteUniform Int Int
data Normal = Normal Double Double
data Bernoulli = Bernoulli Prob
data Exponential = Exponential Double
data Gamma = Gamma Double Double
data Poisson = Poisson Double
data Binomial = Binomial Int Prob
data Beta = Beta Double Double

instance Dist Uniform Double where
    pdf (Uniform a b) = S.density $ S.uniformDistr a b
    logpdf (Uniform a b) = Exp . S.logDensity (S.uniformDistr a b)
    samplefrom (Uniform a b) = do
        u <- uniform
        return $ a + u * (b - a)

{-
instance Dist DiscreteUniform Int where
    pdf (DiscreteUniform a b) = S.probability $ S.discreteUniformAB a b
    logpdf (DiscreteUniform a b) = Exp . S.logProbability (S.discreteUniformAB a b)
    samplefrom (DiscreteUniform a b) = uniform (a, b)
-}

instance Dist Normal Double where
    pdf (Normal mu sig) = S.density $ S.normalDistr mu sig
    logpdf (Normal mu sig) = Exp . S.logDensity (S.normalDistr mu sig)
    samplefrom (Normal mu sig) = S.quantile (S.normalDistr mu sig) <$> uniform

instance Dist Bernoulli Bool where
    pdf (Bernoulli p) b = if b then p else 1-p
    logpdf (Bernoulli p) b = if b then Exp $ log p else Exp $ log (1-p)
    samplefrom (Bernoulli p) = (<= p) <$> uniform

instance Dist Gamma Double where
    pdf (Gamma k theta) = S.density $ S.gammaDistr k theta
    logpdf (Gamma k theta) = Exp . S.logDensity (S.gammaDistr k theta)
    samplefrom (Gamma k theta) = S.quantile (S.gammaDistr k theta) <$> uniform

instance Dist Exponential Double where
    pdf (Exponential lambda) = S.density $ S.exponential lambda
    logpdf (Exponential lambda) = Exp . S.logDensity (S.exponential lambda)
    samplefrom (Exponential lambda) = S.quantile (S.exponential lambda) <$> uniform

instance Dist Beta Double where
    pdf (Beta a b) = S.density $ S.betaDistr a b
    logpdf (Beta a b) = Exp . S.logDensity (S.betaDistr a b)
    samplefrom (Beta a b) = do
        x <- samplefrom (Gamma a 1)
        y <- samplefrom (Gamma b 1)
        return $ x / (x+y)

instance Dist Poisson Int where
    pdf (Poisson lambda) = S.probability $ S.poisson lambda
    logpdf (Poisson lambda) = Exp . S.logProbability (S.poisson lambda)
    samplefrom (Poisson lambda) = step 0 1.0
        where
            l = exp (-lambda)
            step k p = do
                u <- uniform
                let p' = p*u
                if p' <= l then return k
                           else step (k+1) p'

instance Dist Binomial Int where
    pdf (Binomial n p) = S.probability $ S.binomial n p
    logpdf (Binomial n p) = Exp . S.logProbability (S.binomial n p)
    samplefrom d@(Binomial n p) | n <= 10 = sumBernoullis n 0
                                | otherwise = do
        let a = 1 + n `div` 2
        let b = 1 + n - a
        x <- samplefrom $ Beta (fromIntegral a) (fromIntegral b)
        if x >= p
            then samplefrom $ Binomial (a-1) (p/x)
            else do
                s <- samplefrom (Binomial (b-1) ((p-x)/(1-x)))
                return (s + a)
        where
            sumBernoullis 0 s = return s
            sumBernoullis n s = do
                u <- uniform
                if u <= p then sumBernoullis (n-1) (s+1) else sumBernoullis (n-1) s


{-# LANGUAGE Rank2Types, FlexibleContexts, UndecidableInstances, FlexibleInstances, ScopedTypeVariables, DataKinds #-}

module Lib.Utility.Mnist
(
    getMnistData, miniBatch,
    TrainingData
) where

import Lib.Utility.Mnist.IO
import qualified Numeric.LinearAlgebra as HMatrix
import qualified Numeric.LinearAlgebra.Static as Static
import Numeric.LinearAlgebra.Devel
import Control.Monad.ST
import Data.Array.ST
import Control.Arrow
import Data.Proxy
import GHC.TypeLits
import Control.Monad.Trans.Maybe
import System.Random

miniBatch :: forall n m. (KnownNat n, KnownNat m) => TrainingData n -> MaybeT IO (TrainingData m)
miniBatch (images, labels) = MaybeT $ do
    let allSize = fromIntegral $ natVal (Proxy :: Proxy n)
    let size = fromIntegral $ natVal (Proxy :: Proxy m)
    gen <- newStdGen
    let rs = HMatrix.idxs $ take size $ randomRs (0, allSize - 1) gen
    let extr = (HMatrix.Pos rs, HMatrix.All)
    return $ do
        images' <- (Static.create $ Static.unwrap images HMatrix.?? extr) :: Maybe (Static.L m 784)
        labels' <- (Static.create $ Static.unwrap labels HMatrix.?? extr) :: Maybe (Static.L m 10)
        return (images', labels')


{-# LANGUAGE FlexibleContexts #-}
module InterpolationAlgorithms
  (riemannNormCoord,estimateGrad)
  where

import           Control.Monad
import           Control.Monad.ST
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Devel

consInterpMat :: Int
              -> Int
              -> [Vector Double]
              -> Matrix Double
consInterpMat rowlen collen vecs = runST $ do
  m <- unsafeThawMatrix mat
  mapM_ (\(i, v) -> write' m i $ flatOuter v) vecs'
  unsafeFreezeMatrix m
  where
    flatOuter vec' = vjoin [1, vec', flatten . join outer $ vec'] :: Vector Double
    collen' = 1 + collen + (collen * collen)
    mat = konst 0 (rowlen, collen') :: Matrix Double
    vecs' = zip [0 :: Int ..] vecs
    write' m i v' = setMatrix m i 0 (asRow v')

estimateGrad :: Int
             -> Int
             -> Matrix Double
             -> Matrix Double
estimateGrad nums gradlen pcaMat = mat
  where
    mat = consInterpMat nums gradlen pcaMat'
    pcaMat' = map (subVector 1 gradlen . (<# getPCAEigVec pcaMat)) matRows
    matRows = toRows pcaMat

riemannNormCoord :: Matrix Double -> Matrix Double -> [Vector Double]
riemannNormCoord res mat = toColumns leastsquaresSol
  where
    leastsquaresSol = cmap (* (-0.5)) $ (res <\> mat) ?? (Pos $ idxs [1, 2], All)

getPCAEigVec :: Matrix Double -> Matrix Double
getPCAEigVec = snd . eigSH . trustSym . snd . meanCov


module Main where

import Numeric.LinearAlgebra.Data (loadMatrix)

import qualified Euler.Problem001
import qualified Euler.Problem002
import qualified Euler.Problem003
import qualified Euler.Problem004
import qualified Euler.Problem005
import qualified Euler.Problem006
import qualified Euler.Problem007
import qualified Euler.Problem008
import qualified Euler.Problem009
import qualified Euler.Problem010
import qualified Euler.Problem011
import qualified Euler.Problem012
import qualified Euler.Problem013
import qualified Euler.Problem014
import qualified Euler.Problem015
import qualified Euler.Problem016
import qualified Euler.Problem017
import qualified Euler.Problem018
import qualified Euler.Problem019
import qualified Euler.Problem020
import qualified Euler.Problem021
import qualified Euler.Problem022
import qualified Euler.Problem023
import qualified Euler.Problem024
import qualified Euler.Problem025
import qualified Euler.Problem026
import qualified Euler.Problem027
import qualified Euler.Problem028
import qualified Euler.Problem029
import qualified Euler.Problem030
import qualified Euler.Problem031

main :: IO ()
main = do
  label "001" $ Euler.Problem001.solution 1000
  label "002" $ Euler.Problem002.solution 4000000
  label "003" $ Euler.Problem003.solution 600851475143
  label "004" $ Euler.Problem004.solution 1000
  label "005" $ Euler.Problem005.solution 20
  label "006" $ Euler.Problem006.solution 100
  label "007" $ Euler.Problem007.solution 10001
  label "008" $ Euler.Problem008.solution 13
  label "009" $ Euler.Problem009.solution 1000
  label "010" $ Euler.Problem010.solution 2000000
  matrix011 <- loadMatrix "input/011.txt"
  label "011" $ Euler.Problem011.solution 4 matrix011
  label "012" $ Euler.Problem012.solution 500
  label "013" $ Euler.Problem013.solution 10
  label "014" $ Euler.Problem014.solution 1000000
  label "015" $ Euler.Problem015.solution 20
  label "016" $ Euler.Problem016.solution 1000
  label "017" $ Euler.Problem017.solution 1000
  file018 <- readFile "input/018.txt"
  label "018" $ Euler.Problem018.solution file018
  label "019" $ Euler.Problem019.solution (1901, 1, 1) (2000, 12, 31)
  label "020" $ Euler.Problem020.solution 100
  label "021" $ Euler.Problem021.solution 10000
  file022 <- readFile "input/022.txt"
  label "022" $ Euler.Problem022.solution file022
  label "023" $ Euler.Problem023.solution 28123
  label "024" $ Euler.Problem024.solution 1000000
  label "025" $ Euler.Problem025.solution 1000
  label "026" $ Euler.Problem026.solution (1000 :: Int)
  label "027" $ Euler.Problem027.solution (1000 :: Int)
  label "028" $ Euler.Problem028.solution (1001 :: Int)
  label "029" $ Euler.Problem029.solution (100 :: Int)
  label "030" $ Euler.Problem030.solution (5 :: Int)
  label "031" $ Euler.Problem031.solution (200 :: Int)

label :: Integral n => String -> n -> IO ()
label s i = do
  putStr $ "Problem " ++ s ++ " solution: "
  print $ toInteger i


{-|
Module      : FileParsing
Description : Definition of Parsers for csv-files
Copyright   : (c) Julian Kopka Larsen, 2015
Stability   : experimental

Various Functions that provide parsing of CSV files.
-}
module FileParsing (
    parseToArray,
    parseToMatrix,
    parseToDouble,
    normalizeColumns,
    p2NormArray,
    p2NormList,
    matrix2Array
) where

import Text.Parsec
import Data.Array
import Data.Packed.Vector
import Numeric.LinearAlgebra
import Numeric.Statistics

-- Csv parser
csvFile = endBy line eol
line = sepBy cell (char ',')
cell = many (noneOf ",\n")
eol = char '\n'

-- | Parse a File to a list of lists of elements.
parseToStrings :: String -> [[String]]
parseToStrings input = case parse csvFile "(unknown)" input of
                        Left e -> []
                        Right c -> c

-- | Read a csv string to List of Lists of 'Double'
parseToDouble :: String -> [[Double]]
parseToDouble input = map (map read) . tail $ parseToStrings input

-- | Read a csv string into a 'Matrix' of 'Double'
parseToMatrix :: String -> Matrix Double
parseToMatrix = fromLists . parseToDouble

-- | Convert a 'Matrix Double' into an 'Array' of 'Vector Double'
matrix2Array :: Matrix Double -> Array Int (Vector Double)
matrix2Array m = listArray (1,n) r
                 where   r = toRows m
                         n = length r

-- | Parse a csv 'String' into an 'Array' of 'Vector Double'
parseToArray :: String -> Array Int (Vector Double)
parseToArray = matrix2Array . parseToMatrix

-- | Normalize (alt. Studetize) the Collumns in a 'Matrix'
normalizeColumns :: Matrix Double -> Matrix Double
normalizeColumns m = fromColumns $ map studentize (toColumns m)

-- | Parse a csv 'String' into a Normalized 'Array' of 'Vector Double'
p2NormArray :: String -> Array Int (Vector Double)
p2NormArray = matrix2Array . normalizeColumns . parseToMatrix

-- | Parse a csv 'String' into a Normalized 'List' of 'Vector Double'
p2NormList :: String -> [Vector Double]
p2NormList = toRows . normalizeColumns . parseToMatrix


module Test.Statistics where

import Prelude
import qualified Statistics as S
import Test.QuickCheck

runListFold :: S.Statistics t a -> [t] -> a
runListFold = S.runFold foldl

prop_count :: [Bool] -> Property
prop_count l = runListFold (S.count (True ==) 0) l === (length . filter (True ==)) l

main :: IO ()
main = do
    quickCheck prop_count


{-# LANGUAGE CPP                                      #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TemplateHaskell                          #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeOperators                            #-}
{-# OPTIONS_GHC -fno-warn-orphans                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

#if MIN_VERSION_base(4,12,0)
{-# LANGUAGE NoStarIsType         #-}
#endif

import           Control.Monad
import           Data.Complex
import           Data.Function
import           Data.Functor.Identity
import           Data.Maybe
import           Data.Proxy
import           GHC.TypeLits
import           Hedgehog
import           Numeric.Natural
import           System.Exit
import           System.IO
import qualified Data.Vector.Generic                 as UVG
import qualified Data.Vector.Generic.Sized           as VG
import qualified Data.Vector.Sized                   as V
import qualified Data.Vector.Storable.Sized          as VS
import qualified Hedgehog.Gen                        as Gen
import qualified Hedgehog.Range                      as Range
import qualified Numeric.LinearAlgebra               as HU
import qualified Numeric.LinearAlgebra.Static        as H
import qualified Numeric.LinearAlgebra.Static.Vector as H

instance KnownNat n => Eq (H.R n) where
    (==) = (==) `on` H.extract

instance KnownNat n => Eq (H.C n) where
    (==) = (==) `on` H.extract

instance (KnownNat m, KnownNat n) => Eq (H.L m n) where
    (==) = (==) `on` H.extract

instance (KnownNat m, KnownNat n) => Eq (H.M m n) where
    (==) = (==) `on` H.extract

someNatVal' :: Natural -> SomeNat
someNatVal' = fromJust . someNatVal . fromIntegral

genDouble :: Gen H.ℝ
genDouble = Gen.double (Range.linearFracFrom 0 (-10) 10)

genComplex :: Gen H.ℂ
genComplex = (:+) <$> genDouble <*> genDouble

sizeRange :: Num a => Range a
sizeRange = Range.constant 5 10

validVector
    :: (UVG.Vector v a, KnownNat n, MonadTest m)
    => VG.Vector v n a
    -> m ()
validVector v = VG.length v === UVG.length (VG.fromSized v)

prop_rVec :: Property
prop_rVec = property $ do
    xs <- forAll $ Gen.list sizeRange genDouble
    SomeNat (Proxy :: Proxy n) <- pure $ someNatVal' (fromIntegral (length xs))
    tripping (H.vector @n xs)
             (VS.map (* 2) . H.rVec)
             (Identity . (/ 2) . H.vecR)

prop_rVec_konst :: Property
prop_rVec_konst = property $ do
    x <- forAll genDouble
    SomeNat (Proxy :: Proxy n) <- fmap someNatVal' . forAll
                                $ Gen.integral sizeRange
    validVector . H.rVec @n $ H.konst x

prop_vecR :: Property
prop_vecR = property $ do
    xs <- forAll $ Gen.list sizeRange genDouble
    VS.withSizedList xs $ \v ->
      tripping v ((* 2) . H.vecR)
                 (Identity . VS.map (/ 2) . H.rVec)

prop_cVec :: Property
prop_cVec = property $ do
    xs <- forAll $ Gen.list sizeRange genComplex
    SomeNat (Proxy :: Proxy n) <- pure $ someNatVal' (fromIntegral (length xs))
    tripping (fromJust . H.create $ HU.fromList xs)
             (VS.map (* 2) . H.cVec @n)
             (Identity . (/ 2) . H.vecC)

prop_cVec_konst :: Property
prop_cVec_konst = property $ do
    x <- forAll genComplex
    SomeNat (Proxy :: Proxy n) <- fmap someNatVal' . forAll
                                $ Gen.integral sizeRange
    validVector . H.cVec @n $ H.konst x

prop_vecC :: Property
prop_vecC = property $ do
    xs <- forAll $ Gen.list sizeRange genComplex
    VS.withSizedList xs $ \v ->
      tripping v ((* 2) . H.vecC)
                 (Identity . VS.map (/ 2) . H.cVec)

genMatList :: Gen a -> Gen (SomeNat, SomeNat, [[a]])
genMatList g = do
    m  <- Gen.integral sizeRange
    n  <- Gen.integral sizeRange
    xs <- (replicateM (fromIntegral m) . replicateM (fromIntegral n)) g
    return ( someNatVal' m
           , someNatVal' n
           , xs
           )

genMatKonst :: Gen a -> Gen (SomeNat, SomeNat, a)
genMatKonst g = do
    m <- Gen.integral sizeRange
    n <- Gen.integral sizeRange
    x <- g
    return ( someNatVal' m
           , someNatVal' n
           , x
           )

prop_lRows :: Property
prop_lRows = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , xs
     ) <- forAll $ genMatList genDouble
    tripping (fromJust . H.create $ HU.fromLists xs)
             (V.map (* 2) . H.lRows @m @n)
             (Identity . (/ 2) . H.rowsL)

prop_lRows_konst :: Property
prop_lRows_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genDouble
    validVector . H.lRows @m @n $ H.konst x

prop_rowsL :: Property
prop_rowsL = property $ do
    (_, SomeNat (Proxy :: Proxy n), xs)  <- forAll $ genMatList genDouble
    V.withSizedList xs $ \(v :: V.Vector m [Double]) ->
      tripping (V.map (fromJust . H.create . HU.fromList) v)
               ((* 2) . H.rowsL @m @n)
               (Identity . V.map (/ 2) . H.lRows)

prop_rowsL_konst :: Property
prop_rowsL_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genDouble
    let rs = V.replicate $ H.konst x
    validVector . H.lVec . H.rowsL @m @n $ rs
    tripping rs
        ((* 2) . H.rowsL @m @n)
        (Identity . V.map (/ 2) . H.lRows)

prop_lCols :: Property
prop_lCols = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , xs
     ) <- forAll $ genMatList genDouble
    tripping (fromJust . H.create $ HU.fromLists xs)
             (V.map (* 2) . H.lCols @m @n)
             (Identity . (/ 2) . H.colsL)

prop_lCols_konst :: Property
prop_lCols_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genDouble
    validVector . H.lCols @m @n $ H.konst x

prop_colsL :: Property
prop_colsL = property $ do
    (_, SomeNat (Proxy :: Proxy m), xs)  <- forAll $ genMatList genDouble
    V.withSizedList xs $ \(v :: V.Vector n [Double]) ->
      tripping (V.map (fromJust . H.create . HU.fromList) v)
               ((* 2) . H.colsL @m @n)
               (Identity . V.map (/ 2) . H.lCols)

prop_colsL_konst :: Property
prop_colsL_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genDouble
    let cs = V.replicate $ H.konst x
    validVector . H.lVec . H.colsL @m @n $ cs
    tripping cs
        ((* 2) . H.colsL @m @n)
        (Identity . V.map (/ 2) . H.lCols)

prop_vecL :: Property
prop_vecL = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , xs
     ) <- forAll $ genMatList genDouble
    let v :: VS.Vector (m * n) Double
        v = fromJust $ VS.fromList @_ @(m * n) (concat xs)
    tripping v ((* 2) . H.vecL @m @n)
              (Identity . VS.map (/ 2) . H.lVec)

prop_lVec :: Property
prop_lVec = property $ do
    ( SomeNat (pM@Proxy :: Proxy m)
     , SomeNat (pN@Proxy :: Proxy n)
     , xs
     ) <- forAll $ genMatList genDouble
    let m :: H.L m n
        m = fromJust
          . H.create
          . (fromIntegral (natVal pM) HU.>< fromIntegral (natVal pN))
          $ concat xs
    tripping m (VS.map (* 2) . H.lVec)
              (Identity . (/ 2) . H.vecL)

prop_lVec_konst :: Property
prop_lVec_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genDouble
    let xs = H.konst x
    validVector $ H.lVec @m @n xs
    tripping xs
        (VS.map (* 2) . H.lVec @m @n)
        (Identity . (/ 2) . H.vecL)

prop_mRows :: Property
prop_mRows = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , xs
     ) <- forAll $ genMatList genComplex
    tripping (fromJust . H.create $ HU.fromLists xs)
             (V.map (* 2) . H.mRows @m @n)
             (Identity . (/ 2) . H.rowsM)

prop_mRows_konst :: Property
prop_mRows_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genComplex
    validVector . H.mRows @m @n $ H.konst x

prop_rowsM :: Property
prop_rowsM = property $ do
    (_, SomeNat (Proxy :: Proxy n), xs)  <- forAll $ genMatList genComplex
    V.withSizedList xs $ \(v :: V.Vector m [Complex Double]) ->
      tripping (V.map (fromJust . H.create . HU.fromList) v)
               ((* 2) . H.rowsM @m @n)
               (Identity . V.map (/ 2) . H.mRows)

prop_rowsM_konst :: Property
prop_rowsM_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genComplex
    let rs = V.replicate $ H.konst x
    validVector . H.mVec . H.rowsM @m @n $ rs
    tripping rs
        ((* 2) . H.rowsM @m @n)
        (Identity . V.map (/ 2) . H.mRows)

prop_mCols :: Property
prop_mCols = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , xs
     ) <- forAll $ genMatList genComplex
    tripping (fromJust . H.create $ HU.fromLists xs)
             (V.map (* 2) . H.mCols @m @n)
             (Identity . (/ 2) . H.colsM)

prop_mCols_konst :: Property
prop_mCols_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genComplex
    validVector . H.mCols @m @n $ H.konst x

prop_colsM :: Property
prop_colsM = property $ do
    (_, SomeNat (Proxy :: Proxy m), xs)  <- forAll $ genMatList genComplex
    V.withSizedList xs $ \(v :: V.Vector n [Complex Double]) ->
      tripping (V.map (fromJust . H.create . HU.fromList) v)
               ((* 2) . H.colsM @m @n)
               (Identity . V.map (/ 2) . H.mCols)

prop_colsM_konst :: Property
prop_colsM_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genComplex
    let cs = V.replicate $ H.konst x
    validVector . H.mVec . H.colsM @m @n $ cs
    tripping cs
        ((* 2) . H.colsM @m @n)
        (Identity . V.map (/ 2) . H.mCols)

prop_vecM :: Property
prop_vecM = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , xs
     ) <- forAll $ genMatList genComplex
    let v :: VS.Vector (m * n) (Complex Double)
        v = fromJust $ VS.fromList @_ @(m * n) (concat xs)
    tripping v ((* 2) . H.vecM @m @n)
              (Identity . VS.map (/ 2) . H.mVec)

prop_mVec :: Property
prop_mVec = property $ do
    ( SomeNat (pM@Proxy :: Proxy m)
     , SomeNat (pN@Proxy :: Proxy n)
     , xs
     ) <- forAll $ genMatList genComplex
    let m :: H.M m n
        m = fromJust
          . H.create
          . (fromIntegral (natVal pM) HU.>< fromIntegral (natVal pN))
          $ concat xs
    tripping m (VS.map (* 2) . H.mVec)
               (Identity . (/ 2) . H.vecM)

prop_mVec_konst :: Property
prop_mVec_konst = property $ do
    ( SomeNat (Proxy :: Proxy m)
     , SomeNat (Proxy :: Proxy n)
     , x
     ) <- forAll $ genMatKonst genComplex
    let xs = H.konst x
    validVector $ H.mVec @m @n xs
    tripping xs
        (VS.map (* 2) . H.mVec @m @n)
        (Identity . (/ 2) . H.vecM)

main :: IO ()
main = do
  hSetBuffering stdout LineBuffering
  hSetBuffering stderr LineBuffering

  results <- checkParallel $$(discover)

  unless results exitFailure


module CCEA.XOR
  (
    xorFitness
  , xorBreeding
  , xorEA
  , xorVars
  ) where

import NN.NeuralNetwork
import RandomUtil.Random
import AI.HNN.FF.Network

import qualified Data.Map.Strict as Map
import Numeric.LinearAlgebra.HMatrix
import Data.List (sortOn, partition, mapAccumL)
import System.Random
import System.Random.Shuffle
import Data.Tuple
import Data.Traversable

import CCEA.EA

xorVars :: NNVars
xorVars = nnVars

nnVars :: NNVars
nnVars = Map.fromList [("numberInputs", 2)
                      ,("numberHidden", 2)
                      ,("numberOutputs",1)
                      ,("timesToTrain",1)
                      ,("learningRate",0.25)
                      ,("sigmoidOrTanh",0)
                      ,("randomLowerBound", (-2))
                      ,("randomUpperBound", 2)
                      ,("gaussianMean", 0)
                      ,("gaussianStdDev", 1)
                      ,("mutationRate", 0.5)]

xorFitness :: NN n => FitnessFunction n
xorFitness = xorFitnessFunction

xorFitnessFunction :: NN n => n -> Double
xorFitnessFunction net = let x1  = (get net $ fromList [0.0,0.0]) ! 0
                             x2  = (get net $ fromList [0.0,1.0]) ! 0
                             x3  = (get net $ fromList [1.0,0.0]) ! 0
                             x4  = (get net $ fromList [1.0,1.0]) ! 0
                             x1e = 1 - abs (0 - x1)
                             x2e = 1 - abs (1 - x2)
                             x3e = 1 - abs (1 - x3)
                             x4e = 1 - abs (0 - x4)
                         in (x1e + x2e + x3e + x4e) / 4

xorBreeding :: (RandomGen g, NN n) => BreedingStrategy n g
xorBreeding g = notElitist xorBreedingStrategy xorBreedingStrategy g

xorBreedingStrategy ::  (RandomGen g, NN n) => BreedingStrategy n g
xorBreedingStrategy = mutateWeights

xorEA :: RandomGen g => IO (EA Net g)
xorEA = do
  pop <- sequence $ map (\_ -> create nnVars) [1..100]
  return $ EA pop (tournament 2) xorFitness xorBreeding


module Day01 (solveDay01) where

import           Data.Char       (isDigit)
import           Data.Complex
import           Data.Foldable   (foldl')
import           Data.List.Split (splitOn)
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M

type Point   = Complex Double
type Instr   = String
type Command = (Instr, Point)

origin :: Point
origin = 0 :+ 0

commands :: [String] -> [Command]
commands = map parseCommand
  where
    parseCommand s = (dir, dist)
      where dir  = take 1 s
            dist = (read (takeWhile isDigit (drop 1 s)) :: Double) :+ 0

-- √-1
i :: Complex Double
i = 0 :+ 1

step :: (Point, Point) -> Command -> (Point, Point)
step (heading, point) ("R", dist) = (h', (h' * dist) + point)
  where h' = heading * conjugate i
step (heading, point) ("L", dist) = (h', (h' * dist) + point)
  where h' = heading * i
step _ _                          = error "Unrecognized command"

step' :: (Point, [Point]) -> Command -> (Point, [Point])
step' (heading, points) (cmd, dist)
  | cmd == "R" = let h' = heading * conjugate i
                 in (h', points ++ explodePoint (h' * dist))
  | cmd == "L" = let h' = heading * i
                 in (h', points ++ explodePoint (h' * dist))
  | otherwise  = error "Unrecognized command in `steps`"

processCommands' :: [Command] -> [Point]
processCommands' = snd . foldl' step' (north, [])
  where north = 0 :+ 1

processCommands :: [Command] -> Point
processCommands = snd . foldl' step (north, origin)
  where north = 0 :+ 1

calcAns :: Point -> Double
calcAns (x :+ y) = abs x + abs y

explodePoint :: Point -> [Point]
explodePoint (x :+ y)
  | x == 0  = replicate (abs . floor $ y) (0 :+ signum y)
  | y == 0  = replicate (abs . floor $ x) (signum x :+ 0)
  | otherwise = error "Invalid point provided to `explodePoint`"

walk :: [Point] -> Point
walk ps = go ps origin M.empty
  where
    go :: [Point] -> Point -> Map (Double, Double) Bool -> Point
    go []     _ _ = error "Did not encounter any places twice"
    go (s:ss) p m = case M.lookup (x, y) m of
                        Just _  -> p'
                        Nothing -> go ss p' m'
      where p'@(x :+ y) = p + s
            m'          = M.insert (x, y) True m

solveDay01 :: FilePath -> IO ()
solveDay01 pathToInput = do
  putStrLn "Solution for day one: "
  ls <- readFile pathToInput
  let cleanInput = splitOn ", " ls
  -- Part 1
  print . calcAns . processCommands . commands $ cleanInput
  -- Part 2
  print . calcAns . walk . processCommands' . commands $ cleanInput
  putStrLn ""


import Control.Monad (replicateM)
import System.Random.MWC (create, GenIO)
import System.Random.MWC.Distributions (beta, bernoulli, dirichlet)
import Numeric.LinearAlgebra.HMatrix
       (Matrix, matrix, Vector, vector, konst, fromRows, (!), toList, fromList, size)
import qualified Data.Vector.Storable as V

import Utils (multinomial)
import VB (lQ, calculateEQZ, calculateQZ, calculateEQmu, calculateQmu, expectationBeta, em)

generateSampleMatrix :: Double  -- ^ alpha parameter for Beta distribution
                     -> Double  -- ^ beta parameter for Beta distribution
                     -> Vector Double  -- ^ alpha parameter for Dirichlet distr
                     -> Int  -- ^ number of samples
                     -> Int  -- ^ vector size
                     -> Int  -- ^ Number of mixins
                     -> IO (Matrix Double, Vector Double)
generateSampleMatrix a b alpha n d k = do
    gen <- create
    pivT <- dirichlet (toList alpha) gen
    let piv = fromList pivT
    muK <-
        replicateM
            k
            (replicateM
                 d
                 (beta a b gen) >>=
             \m ->
                  return (vector m))
    let mu = fromRows muK
    genData <-
        replicateM
            n
            (generateVectors piv mu gen)
    return (fromRows genData, piv)
    where generateVectors :: Vector Double
                          -> Matrix Double
                          -> GenIO
                          -> IO (Vector Double)
          generateVectors piv mu gen = do
              zn <- multinomial piv gen
              let Just j = V.findIndex (== 1) zn
              xn <-
                  mapM
                      (\i ->
                            bernoulli
                                ((mu ! j) !
                                 i)
                                gen >>=
                            \b ->
                                 if b
                                     then return 1.0
                                     else return 0.0)
                      [0 .. d - 1]
              return (vector xn)

-- | L(q) function must increase
testLQ = do
  let k = 6
      n = 1000
      d = 3
      a = 0.6
      b = 1
      alpha = konst (1 / 6) k
  (mX, piV) <- generateSampleMatrix a b alpha n d k
  let test_k = 10
      test_alpha = konst 0.1 test_k
      test_a = matrix d [0.5 | _ <- [1..test_k*d]]
      test_b = matrix d [0.5 | _ <- [1..test_k*d]]
  gen <- create
  pivT <- dirichlet (toList test_alpha) gen
  let piv = fromList pivT
  let qZ = fromRows (replicate n piv)
  let eQz = calculateEQZ qZ
  let qMu = (test_a, test_b)
  let eQmu = calculateEQmu qMu
  print (lQ mX eQz eQmu)
  let qZ' = calculateQZ mX eQmu eQz
      eQz' = calculateEQZ qZ'
      qMu' = calculateQmu mX eQz' qMu
      eQmu' = calculateEQmu qMu'
  print (lQ mX eQz' eQmu')
  let qZ'' = calculateQZ mX eQmu' eQz'
      eQz'' = calculateEQZ qZ''
      qMu'' = calculateQmu mX eQz'' qMu'
      eQmu'' = calculateEQmu qMu''
  print (lQ mX eQz'' eQmu'')

-- | EM first step. Can't imagine.
testEM = do
  let k = 6
      n = 1000
      d = 3
      a = 0.6
      b = 1
      alpha = konst (1 / 6) k
  (mX, piV) <- generateSampleMatrix a b alpha n d k
  let test_k = 10
      test_alpha = konst 0.1 test_k
      test_a = 0.5
      test_b = 0.5
  gen <- create
  pivT <- dirichlet (toList test_alpha) gen
  let piv = fromList pivT
  print "Original pi vector"
  print piV
  print "Calculated pi vector"
  print (em mX piv test_alpha test_a test_b)

main :: IO ()
main = do
  putStrLn "Test suite not yet implemented"
  let k = 3
      n = 10
      d = 3
      a = 1
      b = 1
      alpha = konst 0.01 k
  d <- generateSampleMatrix a b alpha n d k
  print d
  putStrLn "Test L(q). Must increase."
  testLQ
  putStrLn "Test EM work or not :)"
  testEM
  return ()


{-# LANGUAGE CPP
           , TypeOperators
           , KindSignatures
           , TypeApplications
           , ScopedTypeVariables
           , AllowAmbiguousTypes
           , TypeFamilies
           , FlexibleContexts
           , DeriveFunctor
           , DeriveFoldable
           , DeriveTraversable
           , TemplateHaskell
           , RankNTypes
           , UndecidableInstances
           -- , DeriveAnyClass
           -- , GeneralizedNewtypeDeriving
           -- , DatatypeContexts
  #-}
{-# OPTIONS_GHC -Wall #-}

module Main where

import Control.Applicative
import Control.Monad (unless)
import Data.Complex
import Data.Traversable
import System.Exit (exitFailure)
import Test.QuickCheck (choose, vectorOf)  -- , elements, collect)
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.All (quickCheckAll)

dft :: [Complex Double] -> [Complex Double]
dft xs = [ sum [ x * exp((0.0 :+ (-1.0)) * 2 * pi / lenXs * fromIntegral(k * n))
          | (x, n) <- zip xs [0..]  ]
          | k <- [0..(length xs - 1)]  ]
 where lenXs = fromIntegral $ length xs

type Unop a = a -> a

dft' :: forall f a. ( Applicative f, Traversable f, Sized f
                    , RealFloat a) => Unop (f (Complex a))
dft' xs = omegas (size @f) `cross` xs

cross :: (Functor n, Foldable m, Applicative m, Num a) =>
  n (m a) -> m a -> n a  -- matrix * vector
mat `cross` vec = dot vec <$> mat

dot :: (Foldable f, Applicative f, Num a) =>
  f a -> f a -> a        -- dot product
u `dot` v = sum (liftA2 (*) u v)

newtype (g :. f) a = O (g (f a))
  deriving (Eq, Functor, Foldable, Traversable)

instance (Show a, Show (f a), Show (g String), Functor g) => Show ((g :. f) a) where
  show (O x) = show $ fmap show x

instance (Applicative f, Applicative g) => Applicative (g :. f) where
  pure = O . pure . pure
  O h <*> (O x) = O $ fmap (<*>) h <*> x

class Sized (f :: * -> *) where
    size :: Integer

instance (Sized f, Sized g) => Sized (g :. f) where
    size = size @f * size @g

class FFT f where
  type Reverse f :: * -> *
  fft :: RealFloat a => f (Complex a) -> Reverse f (Complex a)

-- Constraints filled in, by responding to compiler errors incrementally.
instance (Traversable g, Traversable (Reverse g),
          Applicative (Reverse g),
          FFT g,
          Traversable f,
          Applicative f, Applicative (Reverse f),
          FFT f, Sized f, Sized (Reverse g) ) =>
  FFT (g :. f) where
    type Reverse (g :. f) = Reverse f :. Reverse g
    fft = O . fft' . transpose . twiddle . fft' . unO
fft' :: ( Traversable g
        , Applicative (Reverse g)
        , FFT g
        , Traversable f
        , Applicative f
        , RealFloat a ) =>
  g (f (Complex a)) -> Reverse g (f (Complex a))
fft' = transpose . fmap fft . transpose

twiddle :: forall g f a. ( Applicative g, Traversable g, Sized g
                         , Applicative f, Traversable f, Sized f
                         , RealFloat a ) =>
  g (f (Complex a)) -> g (f (Complex a))
twiddle = (liftA2 . liftA2) (*) (omegas (size @(g :. f)))

omegas :: ( Applicative g, Traversable g
          , Applicative f, Traversable f
          , RealFloat a) =>
  Integer -> g (f (Complex a))
-- 'i' was unrecognized.
-- omegas n = powers <$> powers (exp (-i * 2 * pi / fromIntegral n))
omegas n = powers <$> powers (cis (2 * pi / fromIntegral n))

powers :: ( Applicative f, Traversable f
          , Fractional a) => a -> f a
-- I don't have time to digest Conal's "Generic Parallel Scan" talk, right now.
-- powers = fst . lscanAla Product . pure
-- But, I do need to accomodate all applicative functors, as opposed to just lists.
-- So, I'll try the Traversable approach.
powers w = fmap (/ w) . snd . mapAccumL (\x y -> (x * y, x * y)) 1 $ pure w

unO :: (g :. f) a -> g (f a)
unO (O x) = x

-- Copied from Conal's generic-fft library.
transpose :: (Traversable g, Applicative f) => g (f a) -> f (g a)
transpose = sequenceA

-- Quad
newtype Quad a = Quad ((a,a),(a,a))
  deriving (Show, Functor, Foldable, Traversable)
instance Applicative Quad where
  pure x = Quad ((x,x),(x,x))
  Quad ((g,h),(u,v)) <*> Quad ((w,x),(y,z)) = Quad ((g w, h x),(u y, v z))
instance Sized Quad where
  size = 4

-- Pair
data  Pair a = a :# a
  deriving ( Show, Eq, Functor, Foldable, Traversable )
instance Applicative Pair where
  pure x = x :# x
  g :# h <*> (x :# y) = g x :# h y
instance Sized Pair where
  size = 2
instance FFT Pair where
  type Reverse Pair = Pair
  fft (x :# y) = (x + y) :# (x - y)

-- Tree of Pairs
data Tree a = Leaf (Pair a)
            | Branch (Pair (Tree a))
  deriving (Show, Eq, Functor, Foldable, Traversable)

-- Tree building utility.
mkTree :: Integer -> [a] -> Tree a
mkTree n xs | n == 0    = Leaf (head xs :# head (tail xs))
            | otherwise = Branch (mkTree (n-1) (evens xs) :# mkTree (n-1) (odds xs))

-- Pair of Pair building utility.
mkPofP :: [a] -> Pair (Pair a)
-- mkPofP []        = error "Empty list provided!"
-- mkPofP [w]       = error "Wrong number of elements given!"
-- mkPofP [w,x]     = error "Wrong number of elements given!"
-- mkPofP [w,x,y]   = error "Wrong number of elements given!"
-- mkPofP [w,x,y,z] = O ((w :# x) :# (y :# z))
-- mkPofP xs        = O (unO (mkPofP (evens xs)) :# (unO (mkPofP (odds xs))))
mkPofP [w,x,y,z] = (w :# y) :# (x :# z)
-- mkPofP xs        = mkPofP (evens xs) :# mkPofP (odds xs)
mkPofP _         = error "Wrong number of elements given!"

evens :: [a] -> [a]
evens []     = []
evens (x:xs) = x : odds xs

odds :: [a] -> [a]
odds []     = []
odds (_:xs) = evens xs

newtype Id a = Id {unId :: a}
  deriving ( Show, Eq, Functor, Foldable, Traversable )

instance Applicative Id where
  pure = Id
  Id g <*> (Id x) = Id (g x)

instance Sized Id where
  size = 1

instance FFT Id where
  type Reverse Id = Id
  fft (Id x) = Id x

toList :: Foldable f => f a -> [a]
toList = foldr (:) []

assertion1 :: Bool
assertion1 = dft (map (:+ 0.0) [1.0, 0.0, 0.0, 0.0]) == [1.0 :+ 0.0, 1.0 :+ 0.0, 1.0 :+ 0.0, 1.0 :+ 0.0]

assertion2 :: Bool
assertion2 = dft (map (:+ 0.0) [1.0, 0.0, 0.0, 0.0]) == toList (dft' (fmap (:+ 0.0) (Quad ((1.0, 0.0), (0.0, 0.0)))))

assertion3 :: Bool
assertion3 = fft (fmap (:+ (0 :: Double)) (O ((1 :# 0) :# (0 :# 0)))) == dft' (fmap (:+ (0::Double)) (O ((1 :# 0) :# (0 :# 0))))

exeMain :: IO ()
exeMain = do
  print (powers (cis (2 * pi / 4)) :: Quad (Complex Float))
  print assertion1
  print assertion2
  print assertion3

-- QuickCheck types & propositions
newtype FFTTestVal  f g = FFTTestVal {
  getVal :: (g :. f) (Complex Double)
} -- deriving (Show)

instance (Show (f (Complex Double)), Show (g String), Functor g) =>
  Show (FFTTestVal f g) where
    show = show . getVal

instance (Applicative f, Applicative g) => Arbitrary (FFTTestVal f g) where
    arbitrary = do
        depth  <- choose (2, 5)
        -- rvals <- vectorOf (2^depth) (choose (-1.0, 1.0))
        -- ivals <- vectorOf (2^depth) (choose (-1.0, 1.0))
        -- let values = zipWith (:+) rvals ivals
        -- return $ FFTTestVal (fmap (:+ 0.0) (mkP depth))
        -- return $ FFTTestVal (tree2FuncComp $ mkTree depth values)
        return $ FFTTestVal $ O (pure (pure (depth :+ (0.0::Double))))

-- prop_fail = False  -- Test that QuickCheck is, actually, running.

prop_fft_test testVal = fft val == dft' val
  where types = testVal :: FFTTestVal Pair Pair
        val   = getVal testVal

-- Entry point for unit tests.
return []
runTests = $quickCheckAll

testMain :: IO ()
testMain = do
    -- allPass <- $quickCheckAll -- Run QuickCheck on all prop_ functions
    allPass <- runTests
    unless allPass exitFailure
    putStrLn "Success!"

-- This is a clunky, but portable, way to use the same Main module file
-- for both an application and for unit tests.
-- MAIN_FUNCTION is preprocessor macro set to exeMain or testMain.
-- That way we can use the same file for both an application and for tests.
-- #ifndef MAIN_FUNCTION
-- #define MAIN_FUNCTION exeMain
-- #endif
main :: IO ()
-- main = MAIN_FUNCTION
main = testMain


{-# LANGUAGE TypeApplications #-}

module Main where

import Control.Monad.State
import Control.Monad.Trans.Class
import Numeric.LinearAlgebra.HMatrix hiding (scale)
import qualified Graphics.Rendering.OpenGL as GL
import Foreign.Storable
import GLCode
import Sketch
import Points
import Lines

v2f :: Real a => a -> a -> GL.Vertex2 Float
v2f x y = GL.Vertex2 (realToFrac x) (realToFrac y)

translate :: (Fractional a, Storable a) => a -> a -> a -> Matrix a
translate x y z = (4 >< 4) [ 1.0, 0.0, 0.0, 0.0,
                             0.0, 1.0, 0.0, 0.0,
                             0.0, 0.0, 1.0, 0.0,
                             x  , y  , z  , 1.0 ]

scale :: (Fractional a, Storable a) => a -> a -> a -> Matrix a
scale x y z = (4 >< 4) [ x  , 0.0, 0.0, 0.0,
                         0.0, y  , 0.0, 0.0,
                         0.0, 0.0, z  , 0.0,
                         0.0, 0.0, 0.0, 1.0 ]

rotate :: (Fractional a, Storable a, Floating a) => a -> Matrix a
rotate theta = let c = cos theta
                   s = sin theta
               in (4 >< 4) [ c, -s, 0.0, 0.0,
                             s, c, 0.0, 0.0,
                             0.0, 0.0, 1.0, 0.0,
                             0.0, 0.0, 0.0, 1.0]

translateBy :: Float -> Float -> Float -> StateT (Matrix Float) (StateT World IO) ()
translateBy x y z = modify (translate x y z `mul`)

scaleBy :: Float -> StateT (Matrix Float) (StateT World IO) ()
scaleBy s = modify (scale s s s `mul`)

rotateBy :: Float -> StateT (Matrix Float) (StateT World IO) ()
rotateBy s = modify (rotate s `mul`)

main :: IO ()
main = mainLoopState 512 512 1 (ident @Float 4) $ \time -> do

        put (ident @Float 4)

        let t = 5*realToFrac time :: Float

        GL.clearColor GL.$= GL.Color4 0.0 0.0 0.0 1
        io $ GL.clear [GL.ColorBuffer]

        GL.lineSmooth GL.$= GL.Enabled
        GL.lineWidth GL.$= 2
        GL.hint GL.LineSmooth GL.$= GL.Nicest

        modify (rotate (0.1*time) `mul`)
        replicateM_ 23 $ do
            let r = 10+5*sin (0.71*time)
            rotateBy (pi/r)
            let s = sin (0.9*time)
            scaleBy s
            translateBy 0.01 0.01 0.0
            transform <- get

            lift $ setUniform "turtle"
                       "transform" transform
            lift $ drawLine "turtle"
                     "vPosition" [v2f (-1.0) 0.0, v2f 1.0 0.0]

            io GL.flush


{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Lib
import Numeric.LinearAlgebra
import Data.Traversable

train_1_path = "/Users/klinkhae/OSU_Coursework/Rob537/HW1/hw1_data/train1.csv"
train_2_path = "/Users/klinkhae/OSU_Coursework/Rob537/HW1/hw1_data/train2.csv"
test_1_path = "/Users/klinkhae/OSU_Coursework/Rob537/HW1/hw1_data/test1.csv"
test_2_path = "/Users/klinkhae/OSU_Coursework/Rob537/HW1/hw1_data/test2.csv"
test_3_path = "/Users/klinkhae/OSU_Coursework/Rob537/HW1/hw1_data/test3.csv"

lrs = [0.005, 0.01, 0.05, 0.1]
hus = [4, 8, 12, 16]
eps = [1000, 5000, 10000, 100000]
sig = [True, False]

main :: IO ()
main = do
  let exps = generateExps lrs eps hus sig

  putStrLn "All Experiments with training data set 1"
  runExpsIO train_1_path [train_1_path, test_1_path, test_2_path, test_3_path] exps

  putStrLn "All Experiments with training data set 2"
  runExpsIO train_2_path [train_2_path, test_1_path, test_2_path, test_3_path] exps