{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2014
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-----------------------------------------------------------------------------

module Numeric.AD.Rank1.Newton
  (
  -- * Newton's Method (Forward)
    findZero
  , inverse
  , fixedPoint
  , extremum
  -- * Gradient Ascent/Descent (Kahn)
  , gradientDescent
  , gradientAscent
  ) where

import Prelude hiding (all, mapM)
import Data.Foldable (all)
import Data.Traversable
import Numeric.AD.Mode
import Numeric.AD.Rank1.Forward (Forward, diff, diff')
import Numeric.AD.Rank1.Kahn as Kahn (Kahn, gradWith')
import Numeric.AD.Internal.On

-- $setup
-- >>> import Data.Complex

-- | The 'findZero' function finds a zero of a scalar function using
-- Newton's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Examples:
--
-- >>> take 10 $ findZero (\x->x^2-4) 1
-- [1.0,2.5,2.05,2.000609756097561,2.0000000929222947,2.000000000000002,2.0]
--
-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1)
-- 0.0 :+ 1.0
findZero :: (Fractional a, Eq a) => (Forward a -> Forward a) -> a -> [a]
findZero f = go where
  go x = x : if x == xn then [] else go xn where
    (y,y') = diff' f x
    xn = x - y/y'
{-# INLINE findZero #-}

-- | The 'inverse' function inverts a scalar function using
-- Newton's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes
-- constant ("it converges"), no further elements are returned.
--
-- Example:
--
-- >>> last $ take 10 $ inverse sqrt 1 (sqrt 10)
-- 10.0
inverse :: (Fractional a, Eq a) => (Forward a -> Forward a) -> a -> a -> [a]
inverse f x0 y = findZero (\x -> f x - auto y) x0
{-# INLINE inverse  #-}

-- | The 'fixedPoint' function find a fixedpoint of a scalar
-- function using Newton's method; its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- If the stream becomes constant ("it converges"), no further
-- elements are returned.
--
-- >>> last $ take 10 $ fixedPoint cos 1
-- 0.7390851332151607
fixedPoint :: (Fractional a, Eq a) => (Forward a -> Forward a) -> a -> [a]
fixedPoint f = findZero (\x -> f x - x)
{-# INLINE fixedPoint #-}

-- | The 'extremum' function finds an extremum of a scalar
-- function using Newton's method; produces a stream of increasingly
-- accurate results.  (Modulo the usual caveats.) If the stream
-- becomes constant ("it converges"), no further elements are returned.
--
-- >>> last $ take 10 $ extremum cos 1
-- 0.0
extremum :: (Fractional a, Eq a) => (On (Forward (Forward a)) -> On (Forward (Forward a))) -> a -> [a]
extremum f = findZero (diff (off . f . On))
{-# INLINE extremum #-}

-- | The 'gradientDescent' function performs a multivariate
-- optimization, based on the naive-gradient-descent in the file
-- @stalingrad\/examples\/flow-tests\/pre-saddle-1a.vlad@ from the
-- VLAD compiler Stalingrad sources.  Its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- It uses reverse mode automatic differentiation to compute the gradient.
gradientDescent :: (Traversable f, Fractional a, Ord a) => (f (Kahn a) -> Kahn a) -> f a -> [f a]
gradientDescent f x0 = go x0 fx0 xgx0 0.1 (0 :: Int)
  where
    (fx0, xgx0) = Kahn.gradWith' (,) f x0
    go x fx xgx !eta !i
      | eta == 0     = [] -- step size is 0
      | fx1 > fx     = go x fx xgx (eta/2) 0 -- we stepped too far
      | zeroGrad xgx = [] -- gradient is 0
      | otherwise    = x1 : if i == 10
                            then go x1 fx1 xgx1 (eta*2) 0
                            else go x1 fx1 xgx1 eta (i+1)
      where
        zeroGrad = all (\(_,g) -> g == 0)
        x1 = fmap (\(xi,gxi) -> xi - eta * gxi) xgx
        (fx1, xgx1) = Kahn.gradWith' (,) f x1
{-# INLINE gradientDescent #-}

-- | Perform a gradient descent using reverse mode automatic differentiation to compute the gradient.
gradientAscent :: (Traversable f, Fractional a, Ord a) => (f (Kahn a) -> Kahn a) -> f a -> [f a]
gradientAscent f = gradientDescent (negate . f)
{-# INLINE gradientAscent #-}


{-# LANGUAGE ScopedTypeVariables, TypeFamilies, TupleSections #-}

module Fuml.Supervised.Accuracy where

import Fuml.Core
import Numeric.LinearAlgebra
import Data.List (partition, transpose)

accuracy :: (Eq o) => [(Vector Double, o)] -> Predict p o -> Double
accuracy theData model =
  let isCorrect (v,y) = predict model v == y
      corrects = map isCorrect theData
  in realToFrac (length $ filter id corrects) / realToFrac (length corrects)

rmse :: [(Vector Double, Double)] -> Predict p Double -> Double
rmse theData model =
  let getDev (v,y) = let d = y - predict model v in d*d
      devs = map getDev theData
  in sqrt $ sum devs / realToFrac (length devs)

splitData :: Int -> Int -> [a] -> ([a], [a])
splitData nfolds fold xs
  = let xfs = zip xs (cycle $ [0..(nfolds-1)])
        (zins, zouts) = partition ((==fold) . snd) xfs
    in (map fst zins,map fst zouts)

crossValidate :: Monad m
               => Int
               -> [(Vector Double, o)]
               -> ( [(Vector Double, o)] -> Predict p o1 -> Double )
               -> Supervisor m o p o1
               -> Maybe Int
               -> m Double
crossValidate nfolds theData metric super (Just fold) =
  let (testD, trainD) = splitData nfolds fold theData
      fit = runSupervisor super Nothing trainD
  in fmap (metric testD) fit
crossValidate nfolds theData metric super Nothing =
  let folds = [0..(nfolds-1)]
      accs = mapM (crossValidate nfolds theData metric super . Just) folds
  in fmap avg accs

gridSearch  ::  Monad m
               => Int
               -> [(Vector Double, o)]
               -> ( [(Vector Double, o)] -> Predict p o1 -> Double )
               -> (hyper -> Supervisor m o p o1)
               -> [hyper]
               -> Maybe Int
               -> m [(hyper,Double)]
gridSearch nfolds theData metric hsuper hypers (Just fold) =
  let (testD, trainD) = splitData nfolds fold theData
      getAcc h = fmap ((h,) . metric testD) $ runSupervisor (hsuper h) Nothing $ trainD
  in mapM getAcc hypers
gridSearch nfolds theData metric hsuper hypers Nothing = do
  let folds = [0..(nfolds-1)]
  accs <- mapM (gridSearch nfolds theData metric hsuper hypers . Just) folds
  let f accs = (fst $ head accs, avg $ map snd accs)
  return $ map f $ transpose accs

avg xs = sum xs / realToFrac (length xs)


import           Control.DeepSeq
import           Data.Graph.Inductive
import           Data.Graph.Inductive.PatriciaTree as GP
import           GraphBuilder
import           InterpolationAlgorithms
import           Numeric.LinearAlgebra             (size)
import           Plots
import           System.Environment
import           System.Random                     (randomRIO)
import           Utils

main :: IO ()
main = do
  args <- getArgs
  file <- checkFile args
  let nbds = checkArgs args

  matrix' <- getReqMatrix file

  let (rsize,dims) = size matrix'

{-
The purpose of this algorithm is to find the Riemann normal
coordinates (a coordinate system permissible on Riemann manifolds) of
the tangent space at this point. This is built up by first finding the
K-nearest neighbors to this point, and then applying some linear
dimensionality reduction step to this collection of points (also
called a neighborhood). We perform principal component analysis,
however, other linear dimensionality reduction alogrithms are
permissible. The variable *BASEPOINT* is a randomly chosen point from
the dataset.
-}
  basePoint <- randomRIO (0,rsize-1)

  print $ "The base point is point " ++ show basePoint

{-
Find the M-nearest neighbors about each point (hereby referred to as
the "representative point" with respect to its neighbors) in the
dataset and determine the edges (their necessary existence and weight)
between each member of a neighorhood set and their shared
"representative" point as the Euclidean distance between the
representative point and every point in the neighborhood set.
-}

{-The variable *GCTX*, is a list of contexts for the graph. Essentially,
each context contains the nodes leaving and entering a given node. You
may wish to consult the FGL (Functional Graph Library) documentation
for more information. The variable *POINTS* is a list of the K nearest
neighbors from *BASEPOINT*, a randomnly chosen point in the
dataset. The variable *INDS* is a list of the indices of the points
contained in *POINTS*.
-}
  let (gctx,points,inds) = theContextsAndDists nbds basePoint matrix'

      -- Build the graph from the contexts *GCTX*!
      graph = force . foldr (joinContext . force) (empty :: GP.Gr Double Double) $ gctx

      -- Deterimne the sets of shortest distances from points *POINTS*
      graphShortDists = shortDistMat rsize $ shortestDists inds graph

      {- This matrix is used in the interpolation step for finding the Riemann
          normal coordinates. Please refer to the paper in the repository for
          more information.
          -}
      res = estimateGrad (snd nbds) (dims-1) points
      -- The Riemann normal coordinates.
      res' = riemannNormCoord res graphShortDists

  -- If the graph is connected, then perform and save the results of the analysis.
  if isConnected graph
     then plotAndSave2D' file rsize basePoint nbds res'
     else error "The graph is not connected."


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

{-|
Module      : Numeric.Proximal.Static
Description : Solve proximal problems using FISTA
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.Proximal.Static
    ( fista
    , fistaCost
    , indicatorBox
    , l1Cost
    , l1Prox
    , lsqLipschitz
    , proxIndicatorBox
    , quadraticLipschitz
    , residuals
    )
where

import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static
import           Numeric.LinearAlgebra.Static.Util
import           Data.Maybe
import           Data.Tuple.HT                  ( fst3 )

------------------------------------------------------------------
--  Lipschitz calculations on actual matrixes by max eigenvalue --
------------------------------------------------------------------

-- TODO: Implement solution for extracting k largest eigenvalues directly.

-- Note that these are not particularly efficient as each eigenvalue must be extracted and then searched through to find the maximum.
-- | Find the Lipschitz constant for a least squares problem.
lsqLipschitz :: (KnownNat m, KnownNat n) => L m n -> Double
lsqLipschitz a = foldR1 max . eigenvalues $ mTm a

-- | Find the Lipschitz constant for a quadratic program.
quadraticLipschitz :: (KnownNat n) => Sym n -> Double
quadraticLipschitz = foldR1 max . eigenvalues

---------------------------------------------------------
--                  Proximal Functions                 --
---------------------------------------------------------

-- | Returns the cost function for the box indicator function.
-- I(l, u) = {0 if l < xᵢ < u ∀ xᵢ ∈ x
--            ∞ otherwise}
indicatorBox :: Double -> Double -> (forall n . KnownNat n => R n -> Double)
indicatorBox l u =
    foldR (\a b -> a + (if (b > u) || (b < l) then 1 / 0 else 0)) 0

-- | Proximal operator for the box indicator function
proxIndicatorBox :: Double -> Double -> (forall n . KnownNat n => R n -> R n)
proxIndicatorBox l u = dvmap bound
  where
    bound x = case (x < l, x > u) of
        (True, _   ) -> l
        (_   , True) -> u
        _            -> x

-- | Fast Iterative Soft Thresholding Algorithm
-- Used to solve problems of the form F(x) = f(x) + g(x), where
-- f is a smooth convex function with grad F being L-Lipschitz continuous and
-- g is a (potentially non-smooth) convex function. Common problems solved
-- with FISTA include min ||Ax - b|| st. x \in C or min ||Ax - b|| + ||x||_1
fista
    :: (KnownNat n)
    => (R n -> R n)    -- ^ gradient of f
    -> (R n -> R n)    -- ^ proximal operator of g
    -> Double          -- ^ Lipschitz constant
    -> Maybe (R n)     -- ^ Initial guess at a solution
    -> [R n]           -- ^ Output cost function and optimal solution over time
fista grad_f prox_g lipschitz (fromMaybe (konst 0) -> x0) = map fst3
    $ iterate go (x0, x0, 1)
  where
        -- Sometimes update is called p_L
    gradient_step = \x -> x - konst (1 / lipschitz) * grad_f x
    prox_step     = prox_g
    update        = prox_step . gradient_step

    -- Algorithm here
    go (x, y, t) =
        let
            x_update = update y
            t_update = (1 + sqrt (1 + 4 * t ^^ (2 :: Int))) / 2
            y_update =
                x_update + konst ((t - 1) / t_update) * (x_update - x)
        in
            (x_update, y_update, t_update)

-- | Residuals for a list of vectors.
residuals :: forall n . KnownNat n => [R n] -> [Double]
residuals x = zipWith (\y yn -> norm_2 (y - yn)) x (tail x)

-----------------------------------------------------------------
--                        Cost Functions                       --
-----------------------------------------------------------------

-- | Takes the most smooth convex function \(f\) and the convex (potentially
-- non-smooth) function \(g\) and calculates the cost of a given input.
fistaCost :: (R n -> Double) -> (R n -> Double) -> R n -> Double
fistaCost f g x = f x + g x

-- | Cost function for an ℓ₁ term
l1Cost :: (KnownNat n) => Double -> R n -> Double
l1Cost lambda = (lambda *) . norm_1

-- | Proximal equation a the ℓ₁ term
l1Prox :: (KnownNat n) => Double -> R n -> R n
l1Prox lambda = dvmap f where f xi = signum xi * max (abs xi - lambda) 0


-- |
-- Module    : Statistics.Resampling.Bootstrap
-- Copyright : (c) 2009, 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The bootstrap method for statistical inference.

module Statistics.Resampling.Bootstrap
    ( bootstrapBCA
    -- * References
    -- $references
    ) where

import           Data.Vector.Generic ((!))
import qualified Data.Vector.Unboxed as U

import Statistics.Distribution (cumulative, quantile)
import Statistics.Distribution.Normal
import Statistics.Resampling (Bootstrap(..), jackknife)
import Statistics.Sample (mean)
import Statistics.Types (Sample, CL, Estimate, ConfInt, estimateFromInterval,
                         estimateFromErr, CL, significanceLevel)
import qualified Statistics.Resampling as R

data T = {-# UNPACK #-} !Double :< {-# UNPACK #-} !Double
infixl 2 :<

-- | Bias-corrected accelerated (BCA) bootstrap. This adjusts for both
--   bias and skewness in the resampled distribution.
--
--   BCA algorithm is described in ch. 5 of Davison, Hinkley "Confidence
--   intervals" in section 5.3 "Percentile method"
bootstrapBCA
  :: CL Double       -- ^ Confidence level
  -> Sample          -- ^ Full data sample
  -> [(R.Estimator, Bootstrap U.Vector Double)]
  -- ^ Estimates obtained from resampled data and estimator used for
  --   this.
  -> [Estimate ConfInt Double]
bootstrapBCA confidenceLevel sample resampledData
  = map e resampledData
  where
    e (est, Bootstrap pt resample)
      | U.length sample == 1 || isInfinite bias =
          estimateFromErr      pt (0,0) confidenceLevel
      | otherwise =
          estimateFromInterval pt (resample ! lo, resample ! hi) confidenceLevel
      where
        -- Quantile estimates for given CL
        lo    = max (cumn a1) 0
          where a1 = bias + b1 / (1 - accel * b1)
                b1 = bias + z1
        hi    = min (cumn a2) (ni - 1)
          where a2 = bias + b2 / (1 - accel * b2)
                b2 = bias - z1
        -- Number of resamples
        ni    = U.length resample
        n     = fromIntegral ni
        -- Corrections
        z1    = quantile standard (significanceLevel confidenceLevel / 2)
        cumn  = round . (*n) . cumulative standard
        bias  = quantile standard (probN / n)
          where probN = fromIntegral . U.length . U.filter (<pt) $ resample
        accel = sumCubes / (6 * (sumSquares ** 1.5))
          where (sumSquares :< sumCubes) = U.foldl' f (0 :< 0) jack
                f (s :< c) j = s + d2 :< c + d2 * d
                    where d  = jackMean - j
                          d2 = d * d
                jackMean     = mean jack
        jack  = jackknife est sample

-- $references
--
-- * Davison, A.C; Hinkley, D.V. (1997) Bootstrap methods and their
--   application. <http://statwww.epfl.ch/davison/BMA/>


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric, OverloadedStrings #-}
-- |
-- Module    : Statistics.Distribution.DiscreteUniform
-- Copyright : (c) 2016 André Szabolcs Szelp
-- License   : BSD3
--
-- Maintainer  : a.sz.szelp@gmail.com
-- Stability   : experimental
-- Portability : portable
--
-- The discrete uniform distribution. There are two parametrizations of
-- this distribution. First is the probability distribution on an
-- inclusive interval {1, ..., n}. This is parametrized with n only,
-- where p_1, ..., p_n = 1/n. ('discreteUniform').
--
-- The second parametrization is the uniform distribution on {a, ..., b} with
-- probabilities p_a, ..., p_b = 1/(a-b+1). This is parametrized with
-- /a/ and /b/. ('discreteUniformAB')

module Statistics.Distribution.DiscreteUniform
    (
      DiscreteUniform
    -- * Constructors
    , discreteUniform
    , discreteUniformAB
    -- * Accessors
    , rangeFrom
    , rangeTo
    ) where

import Control.Applicative (empty)
import Data.Aeson   (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary  (Binary(..))
import Data.Data    (Data, Typeable)
import System.Random.Stateful (uniformRM)
import GHC.Generics (Generic)

import qualified Statistics.Distribution as D
import Statistics.Internal

-- | The discrete uniform distribution.
data DiscreteUniform = U {
      rangeFrom  :: {-# UNPACK #-} !Int
    -- ^ /a/, the lower bound of the support {a, ..., b}
    , rangeTo    :: {-# UNPACK #-} !Int
    -- ^ /b/, the upper bound of the support {a, ..., b}
    } deriving (Eq, Typeable, Data, Generic)

instance Show DiscreteUniform where
  showsPrec i (U a b) = defaultShow2 "discreteUniformAB" a b i
instance Read DiscreteUniform where
  readPrec = defaultReadPrecM2 "discreteUniformAB" (\a b -> Just (discreteUniformAB a b))

instance ToJSON   DiscreteUniform
instance FromJSON DiscreteUniform where
  parseJSON (Object v) = do
    a <- v .: "uniformA"
    b <- v .: "uniformB"
    return $ discreteUniformAB a b
  parseJSON _ = empty

instance Binary DiscreteUniform where
  put (U a b) = put a >> put b
  get         = discreteUniformAB <$> get <*> get

instance D.Distribution DiscreteUniform where
  cumulative (U a b) x
    | x < fromIntegral a = 0
    | x > fromIntegral b = 1
    | otherwise = fromIntegral (floor x - a + 1) / fromIntegral (b - a + 1)

instance D.DiscreteDistr DiscreteUniform where
  probability (U a b) k
    | k >= a && k <= b = 1 / fromIntegral (b - a + 1)
    | otherwise        = 0

instance D.Mean DiscreteUniform where
  mean (U a b) = fromIntegral (a+b)/2

instance D.Variance DiscreteUniform where
  variance (U a b) = (fromIntegral (b - a + 1)^(2::Int) - 1) / 12

instance D.MaybeMean DiscreteUniform where
  maybeMean = Just . D.mean

instance D.MaybeVariance DiscreteUniform where
  maybeStdDev   = Just . D.stdDev
  maybeVariance = Just . D.variance

instance D.Entropy DiscreteUniform where
  entropy (U a b) = log $ fromIntegral $ b - a + 1

instance D.MaybeEntropy DiscreteUniform where
  maybeEntropy = Just . D.entropy

instance D.ContGen DiscreteUniform where
  genContVar d = fmap fromIntegral . D.genDiscreteVar d

instance D.DiscreteGen DiscreteUniform where
  genDiscreteVar (U a b) = uniformRM (a,b)

-- | Construct discrete uniform distribution on support {1, ..., n}.
--   Range /n/ must be >0.
discreteUniform :: Int             -- ^ Range
                -> DiscreteUniform
discreteUniform n
  | n < 1     = error $ msg ++ "range must be > 0. Got " ++ show n
  | otherwise = U 1 n
  where msg = "Statistics.Distribution.DiscreteUniform.discreteUniform: "

-- | Construct discrete uniform distribution on support {a, ..., b}.
discreteUniformAB :: Int             -- ^ Lower boundary (inclusive)
                  -> Int             -- ^ Upper boundary (inclusive)
                  -> DiscreteUniform
discreteUniformAB a b
  | b < a     = U b a
  | otherwise = U a b


{- |
Raw, unnormalized multi-dimensional versions of the transforms in @fftw@.

Note that the forwards and backwards transforms of this module are not actually
inverses.  For example, @run idft (run dft v) /= v@ in general.

For more information on the individual transforms, see
<http://www.fftw.org/fftw3_doc/What-FFTW-Really-Computes.html>.

@since 0.2
-}

module Numeric.FFT.Vector.Unnormalized.Multi
  (
        -- * Creating and executing 'Plan's
        run,
        plan,
        execute,
        -- * Complex-to-complex transforms
        dft,
        idft,
        -- * Real-to-complex transforms
        dftR2C,
        dftC2R,
  ) where

import Numeric.FFT.Vector.Base
import Foreign
import Foreign.C
import Data.Complex

#include <fftw3.h>

-- | Whether the complex fft is forwards or backwards.
type CDirection = CInt

foreign import ccall unsafe fftw_plan_dft
    :: CInt -> Ptr CInt -> Ptr (Complex Double) -> Ptr (Complex Double)
        -> CDirection -> CFlags -> IO (Ptr CPlan)

foreign import ccall unsafe fftw_plan_dft_r2c
    :: CInt -> Ptr CInt -> Ptr Double -> Ptr (Complex Double) -> CFlags
        -> IO (Ptr CPlan)

foreign import ccall unsafe fftw_plan_dft_c2r
    :: CInt -> Ptr CInt -> Ptr (Complex Double) -> Ptr Double -> CFlags
        -> IO (Ptr CPlan)

dftND :: CDirection -> TransformND (Complex Double) (Complex Double)
dftND d = TransformND
  { inputSizeND = id
  , outputSizeND = id
  , creationSizeFromInputND = id
  , makePlanND = \rk dims a b -> withPlanner . fftw_plan_dft rk dims a b d
  , normalizationND = const id
  }

-- | A forward discrete Fourier transform.  The output and input sizes are the same (@n@).
dft :: TransformND (Complex Double) (Complex Double)
dft = dftND (#const FFTW_FORWARD)

-- | A backward discrete Fourier transform.  The output and input sizes are the same (@n@).
idft :: TransformND (Complex Double) (Complex Double)
idft = dftND (#const FFTW_BACKWARD)

-- | A forward discrete Fourier transform with real data.  If the input size is @n0 * ... * nk@,
-- the output size will be @n0 * ... * nk \`div\` 2 + 1@.
dftR2C :: TransformND Double (Complex Double)
dftR2C = TransformND {
              inputSizeND = id,
              outputSizeND = \n -> n `div` 2 + 1,
              creationSizeFromInputND = id,
              makePlanND = \rk dims a b -> withPlanner . fftw_plan_dft_r2c rk dims a b,
              normalizationND = const id
          }

-- | A backward discrete Fourier transform which produces real data.
--
-- This 'Transform' behaves differently than the others:
--
--  - Calling @plan dftC2R n@ creates a 'Plan' whose /output/ size is @n@, and whose
--    /input/ size is @n \`div\` 2 + 1@.
--
--  - If @length v == n@, then @length (run dftC2R v) == 2*(n-1)@.
dftC2R :: TransformND (Complex Double) Double
dftC2R = TransformND {
            inputSizeND = \n -> n `div` 2 + 1,
            outputSizeND = id,
            creationSizeFromInputND = \n -> 2 * (n-1),
            makePlanND = \rk dims a b -> withPlanner . fftw_plan_dft_c2r rk dims a b,
            normalizationND = const id
        }


module Polynomial (multiply, multiplyWithIntCoeffs, Poly(Poly), Coefficients) where

import Data.Complex
import DFT (dft, idft)
import Data.List (take, zip, repeat)

type Coefficients = [Complex Float]
data Poly = Poly Coefficients deriving (Eq)

instance Show Poly where
  show (Poly p) = show' p 0
    where show' [] degree = ""
          show' (x:xs) 0 = "(" ++ show (round (realPart x)) ++ " + " ++ ("(" ++ show (round (imagPart x)) ++ ")") ++ "i" ++ ")" ++ " + " ++ show' xs 1
          show' [x] degree = "(" ++ show (round (realPart x)) ++ " + " ++ ("(" ++ show (round (imagPart x)) ++ ")") ++ "i" ++ ")" ++ "x^" ++ show degree
          show' (x:xs) degree = "(" ++ show (round (realPart x)) ++ " + " ++ ("(" ++ show (round (imagPart x)) ++ ")") ++ "i" ++ ")" ++ "x^" ++ show degree ++ " + " ++ show' xs (degree + 1)

multiplyWithIntCoeffs :: Poly -> Poly -> [Int]
multiplyWithIntCoeffs p q = map round $ map (realPart) $ coeffs
  where (Poly coeffs) = multiply p q

multiply :: Poly -> Poly -> Poly
multiply (Poly p) (Poly q) = Poly (trimPoly $ idft (dft p' |*| dft q'))
  where (p', q') = pad p q
        trimPoly = take (length p + length q - 1)

pad :: Coefficients -> Coefficients -> (Coefficients, Coefficients)
pad p q = (p', q')
  where maxDegree = closestPowerOfTwo (length p + length q + 2)
        p' = p ++ (take (maxDegree - (length p)) (repeat 0))
        q' = q ++ (take (maxDegree - (length q)) (repeat 0))

(|*|) :: Coefficients -> Coefficients -> Coefficients
(|*|) = hadamard

hadamard :: Coefficients -> Coefficients -> Coefficients
hadamard p q = [ a * b | (a, b) <- zip p q ]

closestPowerOfTwo :: Int -> Int
closestPowerOfTwo x = 2 ^ ceiling (log (fromIntegral x) / log 2)


{-|
Module      : HABQTlib.RandomStates

Generation of random pure (from Haar measure) and mixed states (from measure
induced by partial tracing of purified states).
-}
module HABQTlib.RandomStates
  ( genPureSV
  , genDM
  ) where

import HABQTlib.Data
import Numeric.LinearAlgebra
  ( norm_2
  , randn
  , scalar
  , sumElements
  , takeDiag
  , toComplex
  , tr
  )
import qualified Numeric.LinearAlgebra as LA

-- | Generate a random mixed state of specified rank.
genDM :: Dim -> Rank -> IO DensityMatrix
genDM dim r = do
  r1 <- randn dim r
  r2 <- randn dim r
  let a = toComplex (r1, r2)
  let h = a LA.<> tr a
  let hTr = sumElements $ takeDiag h
  return . DensityMatrix $ h / scalar hTr

-- | Generate a random pure state from Hilbert space of given dimension.
genPureSV :: Dim -> IO PureStateVector
genPureSV dim = do
  r1 <- randn dim 1
  r2 <- randn dim 1
  let sv = toComplex (r1, r2)
  return . PureStateVector $ sv / toComplex (scalar (norm_2 sv), 0)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}

-- TODO: This module is in the wrong module (these are not linear programs)
{-|
Module      : Numeric.LinearProgramming.L1.Static
Description : Solve l1 programs with static vectors/matrices.
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.LinearProgramming.L1.Static (
  l1Solve,
  l1SolveGT,
  l1SolveO,
  lInfSolveO,
  l1SolveU
)

where

import           Data.Maybe (fromJust)
import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.LinearProgramming.L1 as L1

-- Internal only. Could lead to runtime exceptions except in this case we trust
-- the output from the solvers is the correct dimensions.
mkL :: KnownNat n => LA.Vector ℝ -> R n
mkL = fromJust . create

-- | Solution in the L_1 norm, with L_1 regularization, of a linear system @Ax=b@.
--
-- @argmin_x  λ||x||_1 + ||Ax-b||_1@
l1Solve :: (KnownNat m, KnownNat n) =>
           Double -- ^ λ
        -> L m n  -- ^ A
        -> R m    -- ^ b
        -> R n    -- ^ x
l1Solve λ (extract -> a) (extract -> b) = mkL $ L1.l1Solve λ a b

-- | Solution in the L_1 norm, with L_1 regularization, of a system of linear inequalities @Ax>=b@.
--
-- @argmin_x  λ||x||_1 + ||step(b-Ax)||_1@
l1SolveGT :: (KnownNat m, KnownNat n) =>
             Double -- ^ λ
          -> L m n  -- ^ A
          -> R m    -- ^ b
          -> R n    -- ^ x
l1SolveGT λ (extract -> a) (extract -> b) = mkL $ L1.l1SolveGT λ a b

-- | L_1 solution of overconstrained system Ax=b.
--
-- @argmin_x ||Ax-b||_1@
l1SolveO :: (KnownNat m, KnownNat n) =>
            L m n -- ^ A
         -> R m   -- ^ b
         -> R n   -- ^ x
l1SolveO (extract -> a) (extract -> b) = mkL $ L1.l1SolveO a b

-- | L_inf solution of overconstrained system Ax=b.
--
-- @argmin_x ||Ax-b||_inf@
lInfSolveO :: (KnownNat m, KnownNat n) =>
              L m n -- ^ A
           -> R m   -- ^ b
           -> R n   -- ^ x
lInfSolveO (extract -> a) (extract -> b) = mkL $ L1.lInfSolveO a b

-- | L1 solution of underconstrained linear system Ax=b.
--
-- @argmin_x ||x||_1 such that Ax=b@
l1SolveU :: (KnownNat m, KnownNat n) =>
            L m n -- ^ A
         -> R m   -- ^ b
         -> R n   -- ^ x
l1SolveU (extract -> a) (extract -> b) = mkL $ L1.l1SolveU a b


{-# LANGUAGE RecordWildCards #-}

module Trade.Algorithm.Drawdown where

import qualified Data.List as List
import qualified Data.Vector as Vec

import qualified Statistics.Sample as Sample

{-

dds :: [Double] -> [(Int, Double)]
dds xs =
  let go n a [] = [(n, a)]
      go n 0 (x:xs) = go n x xs
      go n a (0:xs) = go (n+1) a xs
      go n a (x:xs) =
        case a*x > 0 of
          True -> go (n+1) (a+x) xs
          False -> let ys = go 1 x xs in (n, a):ys
  in go 1 0 (zipWith (-) (tail xs) xs)

drawdownSeries :: [Double] -> [(Double, Double)]
drawdownSeries us@(u:_) =
  let go acc [] = [acc]
      go (a, b) ((x, y):us) =
        let vs = go (a+fromIntegral x, b+y) us
        in (a, b) : vs
  in go (0, u) (dds us)

drawdowns :: [Double] -> [(Int, Double)]
drawdowns = filter ((<0) . snd) . dds

numberOfDrawdowns :: [Double] -> Int
numberOfDrawdowns = length . drawdowns

maximumDrawdown :: [Double] -> Double
maximumDrawdown = minimum . map snd . dds

wins :: [Double] -> [(Int, Double)]
wins = filter ((>0) . snd) . dds

numberOfWins :: [Double] -> Int
numberOfWins = length . wins

maximumWin :: [Double] -> Double
maximumWin = maximum . map snd . dds

data DrawdownStatistics = DrawdownStatistics {
  numOfWins :: !Int
  , maxWin :: !Double
  , meanWin :: !Double
  , stdDevWin :: !Double
  , numOfDrawdowns :: !Int
  , maxDrawdown :: !Double
  , meanDrawdown :: !Double
  , stdDevDrawdown :: !Double
  } deriving (Show)

drawdownStats :: [Double] -> DrawdownStatistics
drawdownStats xs =
  let ds = map snd (dds xs)
      wins = filter (>0) ds
      vwins = Vec.fromList wins
      losses = filter (<0) ds
      vlosses = Vec.fromList losses
  in DrawdownStatistics {
    numOfWins = length wins
    , maxWin = maximum wins
    , meanWin = Sample.mean vwins
    , stdDevWin = Sample.stdDev vwins
    , numOfDrawdowns = length losses
    , maxDrawdown = minimum losses
    , meanDrawdown = Sample.mean vlosses
    , stdDevDrawdown = Sample.stdDev vlosses
    }

reportDrawdownStats :: [Double] -> HtmlIO
reportDrawdownStats xs =
  let DrawdownStatistics{..} = drawdownStats xs
  in Rep.vtable $
     ["Number of wins", show numOfWins]
     : ["Win max", show maxWin]
     : ["Win mean", show meanWin]
     : ["Win stddev", show stdDevWin]
     : ["----------"]
     : ["Number of drawdowns", show numOfDrawdowns]
     : ["Drawdown max", show maxDrawdown]
     : ["Drawdown mean", show meanDrawdown]
     : ["Drawdown stddev", show stdDevDrawdown]
     : []

-}


#!/usr/bin/env stack
-- stack runghc --package reanimate
{-# LANGUAGE OverloadedStrings #-}
module Main (main) where

import           Data.Complex
import qualified Data.Text           as T
import           Graphics.SvgTree
import           Linear.V2
import           Reanimate

main :: IO ()
main = reanimate $ pauseAtEnd 2
  fourierAnimation_

sWidth :: Double
sWidth = 0.02

piFourier :: Fourier
piFourier = mkFourier piPoints

piPoints :: [RPoint]
piPoints = lineToPoints 500 $
  toLineCommands $ extractPath $ scale 10 $ center $ latexAlign "\\pi"

fourierAnimation_ :: Animation
fourierAnimation_ = mkAnimation 50 $ \t ->
    let fLength = t
        circles = setFourierLength (fLength*maxLength) piFourier
        maxLength = sum $ map magnitude $ take 499 $ drop 1 $ fourierCoefficients piFourier
        phi = fromToS 0 15 t
    in mkGroup
    [ mkBackground "black"
    , drawCircles $ fourierCoefficients $ rotateFourier phi circles
    , withStrokeColor "green" $
      withFillOpacity 0 $
      mkLinePath $ mkFourierOutline circles
    , withFillColor "white" $
      translate (-screenWidth/16*7) (screenHeight/16*7) $
      latex $ T.pack $ "Circles: " ++ show (length $ fourierCoefficients circles)
    ]

newtype Fourier = Fourier {fourierCoefficients :: [Complex Double]}

pointAtFourier :: Fourier -> Complex Double
pointAtFourier = sum . fourierCoefficients

mkFourier :: [RPoint] -> Fourier
mkFourier points = Fourier $ findCoefficient 0 :
    concat [ [findCoefficient n, findCoefficient (-n)] | n <- [1..] ]
  where
    findCoefficient :: Int -> Complex Double
    findCoefficient n =
        sum [ toComplex point * exp (negate (fromIntegral n) * 2 *pi * i*t) * deltaT
            | (idx, point) <- zip [0::Int ..] points, let t = fromIntegral idx/nPoints ]
    i = 0 :+ 1
    toComplex (V2 x y) = x :+ y
    deltaT = recip nPoints
    nPoints = fromIntegral (length points)

-- setFourierCircles :: Double -> Fourier -> Fourier
-- setFourierCircles n _ | n < 1 = error "Invalid argument. Need at least one circle."
-- setFourierCircles n (Fourier coeffs) =
--     Fourier $ take iCircles coeffs ++ [coeffs!!iCircles * realToFrac fCircle]
--   where
--     (iCircles, fCircle) = divMod' n 1

setFourierLength :: Double -> Fourier -> Fourier
setFourierLength _ (Fourier []) = Fourier []
setFourierLength len0 (Fourier (first:lst)) = Fourier $ first : worker len0 lst
  where
    worker _len [] = []
    worker len (c:cs) =
      if magnitude c < len
        then c : worker (len - magnitude c) cs
        else [c * realToFrac (len / magnitude c)]

rotateFourier :: Double -> Fourier -> Fourier
rotateFourier phi (Fourier coeffs) =
    Fourier $ worker coeffs (0::Integer)
  where
    worker [] _ = []
    worker (x:rest) 0 = x : worker rest 1
    worker [left] n = worker [left,0] n
    worker (left:right:rest) n =
      let n' = fromIntegral n in
      left * exp (negate n' * 2 * pi * i * phi') :
      right * exp (n' * 2 * pi * i * phi') :
      worker rest (n+1)
    i = 0 :+ 1
    -- n = length coeffs `div` 2
    phi' = realToFrac phi

drawCircles :: [Complex Double] -> SVG
drawCircles circles = mkGroup
    [ worker circles
    , withStrokeWidth sWidth $
      withStrokeColor "white" $
      withStrokeLineJoin JoinRound $
      withFillOpacity 0 $
      mkLinePath [ (x, y) | x :+ y <- scanl (+) 0 circles ] ]
  where
    worker [] = None
    worker (x :+ y : rest) =
      let radius = sqrt(x*x+y*y) in
      mkGroup
      [ withStrokeWidth sWidth $
        withStrokeColor "dimgrey" $
        withFillOpacity 0 $
        mkCircle radius
      , translate x y $ worker rest ]

mkFourierOutline :: Fourier -> [(Double, Double)]
mkFourierOutline fourier =
    [ (x, y)
    | idx <- [0 .. granularity]
    , let x :+ y = pointAtFourier $ rotateFourier (idx/granularity) fourier
    ]
  where
    granularity = 500


module ISRES where

import Numeric.LinearAlgebra
import Numeric.NLOPT

isres01 :: Int -> Either Result Solution
isres01 n =
  let objf x = x `dot` x + 22 -- define objective
      stop = ObjectiveRelativeTolerance 1e-12 :| [] -- define stopping criterion
      algorithm = ISRES objf [] [] (SeedValue 22) -- specify algorithm
      lowerbounds = fromList $ replicate n (-10) -- specify bounds
      upperbounds = fromList $ replicate n (10) -- specify bounds
      problem = GlobalProblem lowerbounds upperbounds stop algorithm
      x0 = fromList $ replicate n 5 -- specify initial guess
   in minimizeGlobal problem x0

isres02 :: Int -> Either Result Solution
isres02 n =
  let objf x = - norm_1 x
      stop = ObjectiveRelativeTolerance 1e-12 :| [] -- define stopping criterion
      algorithm = ISRES objf [] [eqCon] (SeedValue 22) -- specify algorithm
      lowerbounds = fromList $ replicate n 0 -- specify bounds
      upperbounds = fromList $ replicate n 10 -- specify bounds
      eqCon =
        EqualityConstraint
          { eqConstraintFunctions = Scalar (\x -> x ! 1 * (x ! 0 - 5)),
            {-               Scalar
                            ( \x ->
                                let x1 = subVector 0 (n `div` 2) x
                                    x2 = subVector ((n `div` 2) + 1) n x
                                 in x1 `dot` (x2 - upperbounds)
                            ),
             -}
            eqConstraintTolerance = 1e-12
          }

      problem = GlobalProblem lowerbounds upperbounds stop algorithm
      x0 = fromList $ replicate n 5 -- specify initial guess
   in minimizeGlobal problem x0

isres03 :: Int -> Either Result Solution
isres03 n =
  if n < 3
    then error "n must be at least 3"
    else
      let objf x = - norm_1 x
          stop = ObjectiveRelativeTolerance 1e-12 :| [] -- define stopping criterion
          algorithm = ISRES objf [] [eqCon] (SeedValue 22) -- specify algorithm
          lowerbounds = fromList $ replicate n 0 -- specify bounds
          upperbounds = fromList $ replicate n 10 -- specify bounds
          eqCon =
            EqualityConstraint
              { eqConstraintFunctions =
                  Vector
                    2
                    ( \x _ ->
                        fromList
                          [ x ! 1 * (x ! 0 - 10),
                            x ! 2 * (x ! 1 - (10 - x ! 0))
                          ]
                    ),
                eqConstraintTolerance = 1e-12
              }

          problem = GlobalProblem lowerbounds upperbounds stop algorithm
          x0 = fromList $ replicate n 0 -- specify initial guess
       in minimizeGlobal problem x0

{-
See https://www.mat.univie.ac.at/~neum/glopt/coconut/Benchmark/Library1_new_v1.html

Alkyl problem in AMPL format:
var x2 := 1.745, >= 0, <= 2;
var x3 := 1.2, >= 0, <= 1.6;
var x4 := 1.1, >= 0, <= 1.2;
var x5 := 3.048, >= 0, <= 5;
var x6 := 1.974, >= 0, <= 2;
var x7 := 0.893, >= 0.85, <= 0.93;
var x8 := 0.928, >= 0.9, <= 0.95;
var x9 := 8, >= 3, <= 12;
var x10 := 3.6, >= 1.2, <= 4;
var x11 := 1.45, >= 1.45, <= 1.62;
var x12 := 1, >= 0.99, <= 1.01010101010101;
var x13 := 1, >= 0.99, <= 1.01010101010101;
var x14 := 1, >= 0.9, <= 1.11111111111111;
var x15 := 1, >= 0.99, <= 1.01010101010101;

minimize obj:  - 6.3*x5*x8 + 5.04*x2 + 0.35*x3 + x4 + 3.36*x6;

subject to
e2:  - 0.819672131147541*x2 + x5 - 0.819672131147541*x6 = 0;
e3: 0.98*x4 - x7*(0.01*x5*x10 + x4) = 0;
e4:  - x2*x9 + 10*x3 + x6 = 0;
e5: x5*x12 - x2*(1.12 + 0.13167*x9 - 0.0067*x9*x9) = 0;
e6: x8*x13 - 0.01*(1.098*x9 - 0.038*x9*x9) - 0.325*x7 - 0.57425 = 0;
e7: x10*x14 + 22.2*x11 - 35.82 = 0;
e8: x11*x15 - 3*x8 + 1.33 = 0;
-}
alkyl :: [Double] -> Either Result Solution
alkyl x0 =
  let lb = fromList [0, 0, 0, 0, 0, 0.85, 0.9, 3, 1.2, 1.45, 0.99, 0.99, 0.9, 0.99]
      ub = fromList [2, 1.6, 1.2, 5, 2, 0.93, 0.95, 12, 4, 1.62, 1.01010101010101, 1.01010101010101, 1.11111111111111, 1.01010101010101]
      objf x = - 6.3 * x ! 3 * x ! 6 + 5.04 * x ! 0 + 0.35 * x ! 1 + x ! 2 + 3.36 * x ! 4
      stop = ObjectiveRelativeTolerance 1e-12 :| []
      eqCons =
        [ EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> -0.819672131147541 * x ! 0 + x ! 3 - 0.819672131147541 * x ! 4),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> 0.98 * x ! 2 - x ! 5 * (0.01 * x ! 3 * x ! 8 + x ! 2)),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> - x ! 0 * x ! 7 + 10 * x ! 1 + x ! 4),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> x ! 3 * x ! 10 - x ! 0 * (1.12 + 0.13167 * x ! 7 - 0.0067 * x ! 7 * x ! 7)),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> x ! 6 * x ! 11 - 0.01 * (1.098 * x ! 7 - 0.038 * x ! 7 * x ! 7) - 0.325 * x ! 5 - 0.57425),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> x ! 8 * x ! 12 + 22.2 * x ! 9 - 35.82),
              eqConstraintTolerance = 1e-10
            },
          EqualityConstraint
            { eqConstraintFunctions = Scalar (\x -> x ! 9 * x ! 13 - 3 * x ! 6 + 1.33),
              eqConstraintTolerance = 1e-10
            }
        ]
      algorithm = ISRES objf [] eqCons (SeedValue 22)
      problem = GlobalProblem lb ub stop algorithm
   in minimizeGlobal problem (fromList x0)


{-# LANGUAGE Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Immutable dense matrices.

module Numeric.LinearAlgebra.Matrix (
    -- * Immutable matrices
    Matrix,
    dim,

    -- * Matrix construction
    fromList,
    fromCols,
    fromRows,
    constant,
    zero,

    -- * Accessing matrices
    at,
    indices,
    elems,
    assocs,

    -- * Incremental matrix updates
    update,
    unsafeUpdate,
    accum,
    unsafeAccum,

    -- * Derived matrices
    map,
    zipWith,

    -- * Matrix views
    slice,
    takeRows,
    dropRows,
    splitRowsAt,
    takeCols,
    dropCols,
    splitColsAt,

    -- * Matrix rows and columns
    col,
    cols,
    row,
    rows,

    -- * Matrix diagonals
    diag,

    -- * Conversions to vectors
    toVector,

    -- * Conversions from vectors
    fromVector,
    fromCol,
    fromRow,

    -- * Matrix math operations
    shiftDiag,
    shiftDiagWithScale,
    add,
    addWithScale,
    sub,
    scale,
    scaleRows,
    scaleCols,
    negate,
    conjugate,

    -- * Linear algebra
    trans,
    conjTrans,
    rank1Update,

    -- ** Matrix-Vector multiplication
    mulVector,
    mulVectorWithScale,
    addMulVectorWithScales,

    -- ** Matrix-Matrix multiplication
    mulMatrix,
    mulMatrixWithScale,
    addMulMatrixWithScales,

    -- * Conversions between foreign pointers
    unsafeFromForeignPtr,
    unsafeToForeignPtr,

    -- * Mutable interface
    module Numeric.LinearAlgebra.Matrix.ST,

    -- * Hermitian views
    module Numeric.LinearAlgebra.Matrix.Herm,

    -- * Triangular views
    module Numeric.LinearAlgebra.Matrix.Tri,

    -- * Cholesky factorizations
    module Numeric.LinearAlgebra.Matrix.Cholesky,

    -- * Eigenvalues and eigenvectors
    module Numeric.LinearAlgebra.Matrix.Eigen,

    -- * Basic multivariate statistics
    module Numeric.LinearAlgebra.Matrix.Statistics,

    ) where

import Prelude hiding ( map, zipWith, negate, )
import Control.Monad( zipWithM_ )
import Control.Monad.ST( ST )
import Foreign( Storable )
import Text.Printf( printf )

import Foreign.BLAS( BLAS1, BLAS2, BLAS3, Trans(..) )
import Foreign.VMath( VNum )

import Numeric.LinearAlgebra.Matrix.Base
import Numeric.LinearAlgebra.Matrix.Herm
import Numeric.LinearAlgebra.Matrix.Tri
import Numeric.LinearAlgebra.Matrix.ST
import Numeric.LinearAlgebra.Matrix.Cholesky
import Numeric.LinearAlgebra.Matrix.Eigen
import Numeric.LinearAlgebra.Matrix.Statistics
import Numeric.LinearAlgebra.Vector( Vector )
import qualified Numeric.LinearAlgebra.Vector as V

infixl 7 `scale`, `scaleRows`, `scaleCols`
infixl 6 `add`, `shiftDiag`, `sub`

-- | Create a matrix of the given dimension with the given vectors as
-- columns.
fromCols :: (Storable e) => (Int,Int) -> [Vector e] -> Matrix e
fromCols mn cs = create $ do
    a <- new_ mn
    withColsM a $ \cs' -> zipWithM_ V.copyTo cs' cs
    return a

-- | Create a matrix of the given dimension with the given vectors as
-- rows.
fromRows :: (Storable e) => (Int,Int) -> [Vector e] -> Matrix e
fromRows (m,n) rs = create $ do
    a <- new_ (m,n)
    sequence_ [ setRow a i r | (i,r) <- zip [ 0..m-1 ] rs ]
    return a

-- | Get the given row of the matrix.
row :: (BLAS1 e) => Matrix e -> Int -> Vector e
row a i
    | i < 0 || i >= m = error $
        printf ("row <matrix with dim (%d,%d)> %d:"
                ++ " index out of range") m n i
    | otherwise =
        unsafeRow a i
  where
    (m,n) = dim a
{-# INLINE row #-}

-- | Version of 'row' that doesn't range-check indices.
unsafeRow :: (BLAS1 e) => Matrix e -> Int -> Vector e
unsafeRow a@(Matrix v _ n lda) i
    | lda == 1 = v
    | otherwise = V.create $ do
        r <- V.new_ n
        unsafeRowTo r a i
        return r

-- | Get a list of the rows in a matrix
rows :: (BLAS1 e) => Matrix e -> [Vector e]
rows a = [ unsafeRow a i | i <- [ 0..m-1 ] ]
  where
    (m,_) = dim a

-- | Get the diagonal of the matrix.
diag :: (Storable e) => Matrix e -> Vector e
diag a = V.create $ do
    x <- V.new_ mn
    diagTo x a
    return x
  where
    (m,n) = dim a
    mn = min m n

-- | @shiftDiag d a@ returns @diag(d) + a@.
shiftDiag :: (BLAS1 e) => Vector e -> Matrix e -> Matrix e
shiftDiag s a = create $ do
    a' <- newCopy a
    shiftDiagM_ s a'
    return a'

-- | @shiftDiagWithScale alpha d a@ returns @alpha * diag(d) + a@.
shiftDiagWithScale :: (BLAS1 e) => e -> Vector e -> Matrix e -> Matrix e
shiftDiagWithScale e s a = create $ do
    a' <- newCopy a
    shiftDiagWithScaleM_ e s a'
    return a'

-- | @add a b@ returns @a + b@.
add :: (VNum e) => Matrix e -> Matrix e -> Matrix e
add = result2 addTo

-- | @sub a b@ returns @a - b@.
sub :: (VNum e) => Matrix e -> Matrix e -> Matrix e
sub = result2 subTo

-- | @scale k a@ returns @k * a@.
scale :: (BLAS1 e) => e -> Matrix e -> Matrix e
scale k a = create $ do
    a' <- newCopy a
    scaleM_ k a'
    return a'

-- | @addWithScale alpha x y@ returns @alpha * x + y@.
addWithScale :: (BLAS1 e) => e -> Matrix e -> Matrix e -> Matrix e
addWithScale alpha x y = create $ do
    y' <- newCopy y
    addWithScaleM_ alpha x y'
    return y'

-- | @scaleRows s a@ returns @diag(s) * a@.
scaleRows :: (BLAS1 e) => Vector e -> Matrix e -> Matrix e
scaleRows s a = create $ do
    a' <- newCopy a
    scaleRowsM_ s a'
    return a'

-- | @scaleCols s a@ returns @a * diag(s)@.
scaleCols :: (BLAS1 e) => Vector e -> Matrix e -> Matrix e
scaleCols s a = create $ do
    a' <- newCopy a
    scaleColsM_ s a'
    return a'

-- | @negate a@ returns @-a@.
negate :: (VNum e) => Matrix e -> Matrix e
negate = result negateTo

-- | @conjugate a@ returns @conjugate(a)@.
conjugate :: (VNum e) => Matrix e -> Matrix e
conjugate = result conjugateTo

-- | @trans a@ retunrs @trans(a)@.
trans :: (BLAS1 e)
      => Matrix e
      -> Matrix e
trans a = let
    (m,n) = dim a
    in create $ do
        a' <- new_ (n,m)
        transTo a' a
        return a'

-- | @conjTrans a@ retunrs @conj(trans(a))@.
conjTrans :: (BLAS1 e)
          => Matrix e
          -> Matrix e
conjTrans a = let
    (m,n) = dim a
    in create $ do
        a' <- new_ (n,m)
        conjTransTo a' a
        return a'

-- | @rank1Update alpha x y a@ returns @alpha * x * y^H + a@.
rank1Update :: (BLAS2 e)
            => e
            -> Vector e
            -> Vector e
            -> Matrix e
            -> Matrix e
rank1Update alpha x y a =
    create $ do
        a' <- newCopy a
        rank1UpdateM_ alpha x y a'
        return a'

-- | @mulVector transa a x@
-- returns @op(a) * x@, where @op(a)@ is determined by @transa@.
mulVector :: (BLAS2 e)
          => Trans -> Matrix e
          -> Vector e
          -> Vector e
mulVector transa a x = let
    m = case transa of NoTrans -> (fst . dim) a
                       _       -> (snd . dim) a
    in V.create $ do
        y <- V.new_ m
        mulVectorTo y transa a x
        return y

-- | @mulVectorWithScale alpha transa a x@
-- retunrs @alpha * op(a) * x@, where @op(a)@ is determined by @transa@.
mulVectorWithScale :: (BLAS2 e)
                   => e
                   -> Trans -> Matrix e
                   -> Vector e
                   -> Vector e
mulVectorWithScale alpha transa a x = let
    m = case transa of NoTrans -> (fst . dim) a
                       _       -> (snd . dim) a
    in V.create $ do
        y <- V.new_ m
        mulVectorWithScaleTo y alpha transa a x
        return y

-- | @addMulVectorWithScales alpha transa a x beta y@
-- returns @alpha * op(a) * x + beta * y@, where @op(a)@ is
-- determined by @transa@.
addMulVectorWithScales :: (BLAS2 e)
                       => e
                       -> Trans -> Matrix e
                       -> Vector e
                       -> e
                       -> Vector e
                       -> Vector e
addMulVectorWithScales alpha transa a x beta y =
    V.create $ do
        y' <- V.newCopy y
        addMulVectorWithScalesM_ alpha transa a x beta y'
        return y'

-- | @mulMatrix transa a transb b@
-- returns @op(a) * op(b)@, where @op(a)@ and @op(b)@ are determined
-- by @transa@ and @transb@.
mulMatrix :: (BLAS3 e)
          => Trans -> Matrix e
          -> Trans -> Matrix e
          -> Matrix e
mulMatrix transa a transb b = let
    m = case transa of NoTrans -> (fst . dim) a
                       _       -> (snd . dim) a
    n = case transb of NoTrans -> (snd . dim) b
                       _       -> (fst . dim) b
    in create $ do
        c <- new_ (m,n)
        mulMatrixTo c transa a transb b
        return c

-- | @mulMatrixWithScale alpha transa a transb b@
-- returns @alpha * op(a) * op(b)@, where @op(a)@ and @op(b)@ are determined
-- by @transa@ and @transb@.
mulMatrixWithScale :: (BLAS3 e)
                   => e
                   -> Trans -> Matrix e
                   -> Trans -> Matrix e
                   -> Matrix e
mulMatrixWithScale alpha transa a transb b = let
    m = case transa of NoTrans -> (fst . dim) a
                       _       -> (snd . dim) a
    n = case transb of NoTrans -> (snd . dim) b
                       _       -> (fst . dim) b
    in create $ do
        c <- new_ (m,n)
        mulMatrixWithScaleTo c alpha transa a transb b
        return c

-- | @addMulMatrixWithScales alpha transa a transb b beta c@
-- returns @alpha * op(a) * op(b) + beta * c@, where @op(a)@ and
-- @op(b)@ are determined by @transa@ and @transb@.
addMulMatrixWithScales :: (BLAS3 e)
                       => e
                       -> Trans -> Matrix e
                       -> Trans -> Matrix e
                       -> e
                       -> Matrix e
                       -> Matrix e
addMulMatrixWithScales alpha transa a transb b beta c =
    create $ do
        c' <- newCopy c
        addMulMatrixWithScalesM_ alpha transa a transb b beta c'
        return c'

result :: (Storable e, Storable f)
       => (forall s . STMatrix s f -> Matrix e -> ST s a)
       -> Matrix e
       -> Matrix f
result f a = create $ newResult f a
{-# INLINE result #-}

result2 :: (Storable e, Storable f, Storable g)
        => (forall s . STMatrix s g -> Matrix e -> Matrix f -> ST s a)
        -> Matrix e
        -> Matrix f
        -> Matrix g
result2 f a1 a2 = create $ newResult2 f a1 a2
{-# INLINE result2 #-}

newResult :: (RMatrix m, Storable e, Storable f)
          => (STMatrix s f -> m e -> ST s a)
          -> m e
          -> ST s (STMatrix s f)
newResult f a = do
    mn <- getDim a
    c <- new_ mn
    _ <- f c a
    return c
{-# INLINE newResult #-}

newResult2 :: (RMatrix m1, RMatrix m2, Storable e, Storable f, Storable g)
           => (STMatrix s g -> m1 e -> m2 f -> ST s a)
           -> m1 e
           -> m2 f
           -> ST s (STMatrix s g)
newResult2 f a1 a2 = do
    mn <- getDim a1
    c <- new_ mn
    _ <- f c a1 a2
    return c
{-# INLINE newResult2 #-}


{-# LANGUAGE BangPatterns #-}
{- |
Module      : ML.NN
Description : Simple neural network implementation, for education purposes only.
-}
module ML.NN
       (
           -- * Data types
           ActivationFunction  -- Re-export from NN.ActivationFunction.
       ,   Layer(..)
       ,   Network(..)

           -- * Network initialization
       ,   randLayer
       ,   randNetwork

           -- * Running networks
       ,   feedForward
       ,   runNetwork
       ,   feedForwardSet
       ,   runNetworkSet
       ) where

import           ML.NN.ActivationFunction (ActivationFunction)

import           Control.Monad (forM, replicateM)
import           Control.Monad.Random (Rand, liftRand)
import           Data.Foldable (foldl')
import           Data.Random.Normal (normal)
import           Numeric.LinearAlgebra (Matrix, R, Vector, (><), (<>), (#>), cmap,
                                        fromColumns, vector)
import           System.Random (RandomGen)

-- ==================================================
-- Types.
-- ==================================================

-- | A layer in a neural network is a bias vector and weight matrix.
--
-- Let /n/ be the number of neurons in the layer and /m/ be the number
-- of inputs to the layer. Then 'layerBiases' is a vector in Rⁿ and
-- 'layerWeights' is an /nxm/ real matrix.
data Layer = Layer
             { layerBiases  :: {-# UNPACK #-} !(Vector R)
                                              -- ^ A vector in Rⁿ representing the neuron biases.
             , layerWeights :: {-# UNPACK #-} !(Matrix R)
                                              -- ^ An /nxm/ real matrix of the neuron weights.
             } deriving (Show, Read)

-- | A network is a list of layers.
data Network = Network
               { networkLayers :: [Layer]
               } deriving (Show, Read)

-- ==================================================
-- Network initialization.
-- ==================================================

-- | Generate a random neural network given the size of each layer.
--
-- For example, @[3,2,4]@ will generate a neural network with 3 input
-- neurons, 1 hidden layer with 2 neurons and 4 output neurons.
randNetwork :: RandomGen g => [Int] -> Rand g Network
randNetwork sizes = Network <$> forM dims (uncurry randLayer)
    where dims = adjacentPairs sizes

-- | Return a list of pairs of adjacent elements in a list.
--
-- Example: adjacentPairs [1,2,3] ==> [(1,2), (2,3)]
adjacentPairs :: [a] -> [(a,a)]
adjacentPairs xs = zip (init xs) (tail xs)

-- | Generate a randomly initialized layer in a neural network.
randLayer :: RandomGen g
             => Int           -- ^ Number of inputs into the layer.
             -> Int           -- ^ Number of neurons in the layer.
             -> Rand g Layer  -- ^ Randomly generated layer.
randLayer numInputs numNeurons = do
    bias <- vector <$> replicateM numNeurons (liftRand normal)
    weights <- (numNeurons><numInputs) <$> replicateM (numNeurons*numInputs) (liftRand normal)
    return $ Layer bias weights

-- ==================================================
-- Running networks.
-- ==================================================

-- | Feed the output from a previous layer to the next layer.
feedForward :: ActivationFunction  -- ^ Neuron activation function.
            -> Vector R            -- ^ Input to the layer.
            -> Layer               -- ^ Layer to run.
            -> Vector R            -- ^ Output from the layer.
feedForward act x (Layer biases weights) = cmap act z
    where z = (weights #> x) + biases

-- | Run a neural network.
runNetwork :: Network             -- ^ Network to run.
           -> ActivationFunction  -- ^ Neuron activation function.
           -> Vector R            -- ^ Network input.
           -> Vector R            -- ^ Network output.
runNetwork (Network layers) act input = foldl' (feedForward act) input layers

feedForwardSet :: ActivationFunction -> Matrix R -> Layer -> Matrix R
feedForwardSet act xM (Layer biases weights) = cmap act z
    where z = (weights <> xM) + fromColumns [biases]

runNetworkSet :: Network -> ActivationFunction -> Matrix R -> Matrix R
runNetworkSet (Network layers) act input = foldl' (feedForwardSet act) input layers


module Neural
  ( Layer(Layer)
  , Network
  , evalLayer
  , evalLayerM
  , evalNetwork
  , evalNetworkM
  , evalOutputError
  , evalOutputErrorM
  , evalLayerError
  , evalLayerErrorM
  , evalNetworkError
  , evalNetworkErrorM
  , gradientDescent
  , gradientDescentM
  , sumColumns
  , vecToMat
  )
where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data

-- Dependent types would be great to represent some of these constraints.

-- A layer in the neural network. The weights are represented by a matrix with
-- 1 row for each neuron, and 1 column for each weight. The biases are represented
-- by a vector. The number of biases must match the number of weights.
data Layer = Layer (Matrix R) (Vector R) deriving Show

-- A network is a list of layers, where each layer has the same number of weights as
-- the previous one has neurons. Otherwise, evalNetwork will fail.
type Network = [Layer]

-- A utility function that converts a vector to a row matrix
vecToMat :: Vector R -> Matrix R
vecToMat = fromRows . (:[])

-- A utility function that sums the colums in a matrix to produce a row vector
sumColumns :: Matrix R -> Vector R
sumColumns m = (vector . replicate (rows m) $ 1) <# m

-- A sigmoid function which is at 0.5 for the input 0, and tends to 1 as the input
-- increases, and to 0 as the input decreases.
sigmoid :: Floating f => f -> f
sigmoid z = 1 / (1 + exp (-z))

-- Differential of the sigmoid fuunction
sigmoid' :: Floating f => f -> f
sigmoid' z = sigmoid z * (1 - sigmoid z)

-- Evaluates a layer of the neural network given a set of inputs. The number of inputs
-- must match the number of weights in the input network or this function will fail
-- with an exception. Returns the outputs as well as the weighted inputs.
evalLayer :: Layer -> Vector R -> (Vector R, Vector R)
evalLayer layer inputs = (flatten output, flatten weightedInput)
  where
    (output, weightedInput) = evalLayerM layer (fromRows [inputs])

-- Evaluates a layer of the neural network given a set of inputs. Contrary to evalLayer,
-- the set of inputs is provided as a matrix where each row is a vector of inputs as
-- given in evalLayer.
evalLayerM :: Layer -> Matrix R -> (Matrix R, Matrix R)
evalLayerM (Layer weights biases) inputs = (sigmoid weightedInputs, weightedInputs)
  where
    inputRows = rows inputs
    biasMatrix = fromColumns . replicate inputRows $ biases
    weightedInputs = tr' $ (weights <> tr' inputs) + biasMatrix

-- Feed forward the given function (e.g. evalLayer) over the network
feedForward :: (Layer -> a -> (a, a)) -> a -> Network -> [(a, a)]
feedForward f inputs network = tail . reverse $ feedForward' f inputs network
  where
    feedForward' f inputs network = foldl (\acc@((a,_):_) l -> (f l a) : acc) [(inputs, undefined)] network

-- Evaluates a neural network by folding a set of inputs through each layer.
-- The number of inputs must match the number of weights in the first layer,
-- and the number of elements in the output vector depends on the number of neurons
-- in the final layer of the network. Returns the output and weighted input of each
-- layer in the network.
evalNetwork :: Network -> Vector R -> [(Vector R, Vector R)]
evalNetwork network inputs = feedForward evalLayer inputs network

-- Evaluates a neural network. Contrary to evalNetwork, the input is taken in the same form
-- as evalLayerM, i.e. a matrix where the rows are a single input vector.
evalNetworkM :: Network -> Matrix R -> [(Matrix R, Matrix R)]
evalNetworkM network inputs = feedForward evalLayerM inputs network

-- Evaluates the error in the output layer of a network
evalOutputError :: Vector R -> Vector R -> Vector R -> Vector R -> Vector R
evalOutputError output weightedInput inputs training = (output - training) * sigmoid' weightedInput

-- Evaluates the error in the output layer of a network. Input is a matrix where rows are inputs.
evalOutputErrorM :: Matrix R -> Matrix R -> Matrix R -> Matrix R -> Matrix R
evalOutputErrorM output weightedInput inputs training = (output - training) * sigmoid' weightedInput

-- Evaluates the error in a layer of a network given the weighted cost, and the next layer's weight
-- and error.
evalLayerError :: Vector R -> Matrix R -> Vector R -> Vector R
evalLayerError weightedCost nextLayerWeights nextLayerError =
  ((tr' nextLayerWeights) #> nextLayerError) * (sigmoid' weightedCost)

-- Evaluates the error in a layer like evalLayerError, but takes a matrix of inputs as rows.
evalLayerErrorM :: Matrix R -> Matrix R -> Matrix R -> Matrix R
evalLayerErrorM weightedCost nextLayerWeights nextLayerError =
  (nextLayerError <> nextLayerWeights) * (sigmoid' weightedCost)

-- Evaluates the error for each layer in the network by means of backpropogation.
evalNetworkError :: Network -> [(Vector R, Vector R)] -> Vector R -> Vector R -> [Vector R]
evalNetworkError network networkValues inputs training =
  map flatten $ evalNetworkErrorM network networkValuesM (vecToMat inputs) (vecToMat training)
    where
      networkValuesM = map (\(a,b) -> (vecToMat a, vecToMat b)) networkValues

-- Evaluates the error for each layer in the network by means of backpropogation.
evalNetworkErrorM :: Network -> [(Matrix R, Matrix R)] -> Matrix R -> Matrix R -> [Matrix R]
evalNetworkErrorM network networkValues inputs training = map fst backPropogation
  where
    -- The layers of the network in reverse order
    reverseNetwork = reverse network
    -- The activations of the layers of the network in reverse order
    reverseLayerVals = reverse networkValues
    -- (reverse layers, reverse activations)
    reverseLayerData = zip reverseNetwork reverseLayerVals
    -- The activations of the output layer of the network
    outputLayer = tail reverseLayerVals
    -- The weights of the output layer
    Layer outputWeights _ = head reverseNetwork
    -- The activations and weighted inputs of the output layer
    (outputValues, outputLayerWeightedInput) = head reverseLayerVals
    -- The error of the output layer
    outputError = evalOutputErrorM outputValues outputLayerWeightedInput inputs training
    -- An accumulator for backpropogation, started with the output layer's error
    acc = [(outputError, outputWeights)]
    -- The backpropogation algorithhm (apply evalOutputError to each layer in reverse order)
    backPropogate = \acc@((loe,lw):_) (Layer w _, (_,z)) -> (evalLayerErrorM z lw loe, w):acc
    backPropogation = foldl backPropogate acc (tail reverseLayerData)

-- Gradient descent function. Given a network, its inputs, its activations, and its error,
-- produces a new network corrected through gradient descent.
gradientDescent :: Network -> Vector R -> [Vector R] -> [Vector R] -> R -> Network
gradientDescent network input activations error learnRate =
  gradientDescentM network (vecToMat input) (map vecToMat activations) (map vecToMat error) learnRate

-- Gradient descent function. Given a network, its inputs, its activations, and its error,
-- produces a new network corrected through gradient descent.
gradientDescentM :: Network -> Matrix R -> [Matrix R] -> [Matrix R] -> R -> Network
gradientDescentM network input activations error learnRate =
  map gradientDescent' layerData
    where
      -- Get the activations from the previous layer of neurons
      -- by offsetting the activations with the input at the start
      -- and discarding the final element
      previousActivations = input : (init activations)
      -- Layers zipped with their errors and the previous layer's activation
      layerData = zip3 network error previousActivations
      -- Gradient descent for one layer
      gradientDescent' (Layer weight bias, error, previousActivation) =
        -- New layer is evaluated by calculating the gradient for the weight and bias and
        -- subtracting gradient * learnRate from the current value
        Layer
          (weight - (weightGradients * realToFrac learnRate))
          (bias - (biasGradients * realToFrac learnRate))
            where
              -- Calculate gradient for each rate
              -- Calculated using the kronecker product, a generalisation of tensor product to matrices
              --weightGradients = kronecker previousActivation (tr' error)
              --weightGradients = (previousActivation <> (tr' error)) -- / (fromIntegral $ cols previousActivation)
              weightGradients = (tr' error <> previousActivation) / (fromIntegral . cols $ previousActivation)
              -- Sum the bias gradient components for each neuron (the columns)
              biasGradients = sumColumns error / (fromIntegral . size $ bias)


{-# LANGUAGE FlexibleContexts #-}

module Spectral(layout,default_layout) where

import Numeric.LinearAlgebra.HMatrix
import Numeric.LinearAlgebra.Data

degree_matrix m =
  let a = foldl1 (+) (toRows m)
      b = takeDiag m
      in
    a - b

compute_laplacian adj_matrix node_weights =
  let (n,m) = size adj_matrix
      weighted_adjency = adj_matrix
      mat_D = diag $ degree_matrix weighted_adjency
      mat_L = mat_D - weighted_adjency
      in
      (mat_L,mat_D)

symmetrize m = m + tr m

layout adjacent_matrix node_weights =
  let sym_adjacent_matrix = symmetrize adjacent_matrix
      (mat_L,mat_D) = compute_laplacian sym_adjacent_matrix node_weights
      eigenVectors = toColumns $ snd $ geigSH' mat_L mat_D
      v2 = eigenVectors !! 1
      v3 = eigenVectors !! 2
  in
    [(v2  ! i, v3 ! i) | i <- [0..size v2-1] ]

default_layout adjacent_matrix =
  let (n,_) = size adjacent_matrix
      ones = matrix n [1 | i <- [1..n]]
  in
    layout adjacent_matrix ones


module Numeric.FFT.Utils
       ( omega, slicevecs, slicemvecs, primes, isPrime
       , allFactors, factors
       , primitiveRoot, invModN, log2, isPow2, dupperm, (%.%)
       , compositions, makeComp, multisetPerms
       , backpermuteM
       ) where

import qualified Control.Monad               as CM
import           Control.Monad.ST
import           Data.Bits
import           Data.Complex
import           Data.List                   (nub)
import qualified Data.List                   as L
import qualified Data.Vector                 as V
import           Data.Vector.Unboxed
import qualified Data.Vector.Unboxed.Mutable as MV
import           Prelude                     hiding (all, concatMap, dropWhile,
                                              enumFromTo, filter, head, length,
                                              map, maximum, null, reverse)
import qualified Prelude                     as P

import           Numeric.FFT.Types

-- | Roots of unity.
omega :: Int -> Complex Double
omega n = cis (2 * pi / fromIntegral n)

-- | Slice a vector @v@ into equally sized parts, each of length @m@.
slicevecs :: Int -> VCD -> VVCD
slicevecs m v = V.map (\i -> slice (i * m) m v) $
                V.enumFromN 0 (length v `div` m)

-- | Slice a mutable vector @v@ into equally sized parts, each of
-- length @m@.
slicemvecs :: Int -> MVCD a -> VMVCD a
slicemvecs m v = V.map (\i -> MV.slice (i * m) m v) $
                 V.enumFromN 0 (MV.length v `div` m)

-- | Determine primitive roots modulo n.
--
-- From Wikipedia (https://en.wikipedia.org/wiki/Primitive_root_modulo_n):
--
-- No simple general formula to compute primitive roots modulo n is
-- known.  There are however methods to locate a primitive root that
-- are faster than simply trying out all candidates.  If the
-- multiplicative order of a number m modulo n is equal to phi(n) (the
-- order of Z_n^x), then it is a primitive root.  In fact the converse
-- is true: if m is a primitive root modulo n, then the multiplicative
-- order of m is phi(n).  We can use this to test for primitive roots.
-- [Here, phi(n) is Euler's totient function, and Z_n^x is the
-- multiplicative group of integers modulo n.]
--
-- First, compute phi(n).  Then determine the different prime factors
-- of phi(n), say p1, ..., pk.  Now, for every element m of Z_n^x,
-- compute
--
--    m^(phi(n) / pi) mod n   for i = 1, ..., k
--
-- using a fast algorithm for modular exponentiation such as
-- exponentiation by squaring.  A number m for which these k results
-- are all different from 1 is a primitive root.
--
-- [In our case, n is restricted to being prime, and phi(p) = p - 1
-- for prime p.]
--
primitiveRoot :: Int -> Int
primitiveRoot p
  | isPrime p =
    let tot = p - 1
        -- ^ Euler's totient function for prime values.
        totpows = map (tot `div`) $ fromList $ nub $ toList $ allFactors tot
        -- ^ Powers to check.
        check n = all (/=1) $ map (expt p n) totpows
        -- ^ All powers are different from 1 => primitive root.
    in fromIntegral $ head $ dropWhile (not . check) $ fromList [1..p-1]
  | otherwise = error "Attempt to take primitive root of non-prime value"

-- | Fast exponentation modulo n by squaring.
expt :: Int -> Int -> Int -> Int
expt n b pow = fromIntegral $ go pow
  where bb = fromIntegral b
        nb = fromIntegral n
        go :: Int -> Integer
        go p
          | p == 0 = 1
          | p `mod` 2 == 1 = (bb * go (p - 1)) `mod` nb
          | otherwise = let h = go (p `div` 2) in (h * h) `mod` nb

-- | Find inverse element in multiplicative integer group modulo n.
invModN :: Int -> Int -> Int
invModN n g = head $ filter (\iv -> (g * iv) `mod` n == 1) $ enumFromTo 1 (n-1)

-- | Prime sieve from Haskell wiki.
primes :: Integral a => [a]
primes = 2 : primes'
  where primes' = sieve [3, 5 ..] 9 primes'
        sieve (x:xs) q ps@ ~(p:t)
          | x < q     = x : sieve xs q ps
          | otherwise =     sieve [n | n <- xs, rem n p /= 0] (P.head t^(2::Int)) t

-- | Naive primality testing.
isPrime :: Integral a => a -> Bool
isPrime n = n `P.elem` P.takeWhile (<= n) primes

-- | Simple prime factorisation.
allFactors :: (Integral a, Unbox a) => a -> Vector a
allFactors n = fromList $ go n primes
  where go cur pss@(p:ps)
          | cur == p         = [p]
          | cur `mod` p == 0 = p : go (cur `div` p) pss
          | otherwise        = go cur ps

-- | Simple prime factorisation: small factors only; largest/last
-- factor picked out as "special".
factors :: (Integral a, Unbox a) => a -> (a, Vector a)
factors n = let (lst, rest) = go n primes in (lst, fromList rest)
  where go cur pss@(p:ps)
          | cur == p         = (p, [])
          | cur `mod` p == 0 = let (lst, rest) = go (cur `div` p) pss
                               in (lst, p : rest)
          | otherwise        = go cur ps

-- | Base-2 logarithm.
log2 :: Int -> Int
log2 1 = 0
log2 n = 1 + log2 (n `div` 2)

-- | Check for powers of two.
isPow2 :: Int -> Bool
isPow2 1 = True
isPow2 n
  | n `mod` 2 == 0 = isPow2 $ n `div` 2
  | otherwise      = False

-- | Duplicate a sub-permutation to fill a given vector length.
dupperm :: Int -> VI -> VI
dupperm n p =
  let sublen = length p
      shift di = map (+(sublen * di)) p
  in concatMap shift $ enumFromN 0 (n `div` sublen)

-- | Composition of permutations.
(%.%) :: VI -> VI -> VI
p1 %.% p2 = backpermute p2 p1

-- | Generate all compositions of a given integer.
compositions :: Int -> V.Vector (Vector Int)
compositions 0 = V.empty
compositions n = let fs = allFactors n
                 in V.reverse $ V.map (makeComp fs) $ V.enumFromN 0 (2^(n-1))

-- | Generate a single composition of a given integer.
makeComp :: Vector Int -> Int -> Vector Int
makeComp fs i = fromList $ foldOps (toList fs) $ makeOps (length fs) i
  where foldOps :: [Int] -> [Bool] -> [Int]
        foldOps (f:fs) ops = go f fs ops
          where go acc [] [] = [acc]
                go acc (f:fs) (op:ops) = if op then go (acc * f) fs ops
                                         else acc : go f fs ops
        makeOps :: Int -> Int -> [Bool]
        makeOps n i = P.replicate (n - 1 - P.length bs) False P.++ bs
          where bs = P.dropWhile not $ P.reverse $
                     P.map (testBit i) [0..finiteBitSize i-1]

-- | Generate all distinct permutations of a multiset in lexicographic
-- order.
multisetPerms :: Vector Int -> [Vector Int]
multisetPerms idp = sidp : L.unfoldr step sidp
  where sidp = fromList $ L.sort $ toList idp
        step v = case permStep v of
          Nothing -> Nothing
          Just p  -> Just (p, p)
        permStep :: Vector Int -> Maybe (Vector Int)
        permStep v =
          if null ks
          then Nothing
          else let k = maximum ks
                   ls = filter (\i -> v ! k < v ! i) $ enumFromN 0 n
                   l = maximum ls
               in Just $ revEnd k (swap k l)
          where n = length v
                ks = filter (\i -> v ! i < v ! (i+1)) $ enumFromN 0 (n-1)
                swap a b = generate n $ \i ->
                  if i == a then v ! b
                  else if i == b then v ! a
                       else v ! i
                revEnd f vv = generate n $ \i ->
                  if i <= f
                  then vv ! i
                  else vv ! (n - i + f)

-- | ST monad version of vector permutation.
backpermuteM :: Int -> VI -> MVCD s -> MVCD s -> ST s ()
backpermuteM n perm vin vout =
  CM.forM_ [0..n-1] $ \i -> do
    idx <- indexM perm i
    x <- MV.unsafeRead vin idx
    MV.unsafeWrite vout i x


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Strict #-}
module Pinwheel.Base where

import           Data.Complex

-- {-# INLINE fourierMellin #-}
-- fourierMellin :: (Eq e, RealFloat e) => e -> Int -> Int -> (e, e) -> Complex e
-- fourierMellin sigma angularFreq radialFreq (!x, !y) =
--   if x == 0 && y == 0
--     then 0
--     else (x :+ y) ** (fromIntegral (-angularFreq) :+ 0) *
--          ((x ^ 2 + y ^ 2) :+ 0) **
--          (((fromIntegral angularFreq - 2 + sigma) :+ (-fromIntegral radialFreq)) /
--           2)

{-# INLINE fourierMellin #-}
fourierMellin :: (Eq e, RealFloat e) => e -> Int -> Int -> (e, e) -> Complex e
fourierMellin sigma angularFreq radialFreq (!x, !y) =
  if x == 0 && y == 0
    then 0
    else let r = sqrt $ x ^ 2 + y ^ 2
             theta = atan2 y x
         in ((r :+ 0) ** ((sigma - 1) :+ fromIntegral (-radialFreq))) *
            (cis $ fromIntegral (-angularFreq) * theta)

{-# INLINE fourierMellinInv #-}
fourierMellinInv :: (Eq e, RealFloat e) => e -> Int -> Int -> (e, e) -> Complex e
fourierMellinInv sigma angularFreq radialFreq (!x, !y) =
  if x == 0 && y == 0
    then 0
    else (x :+ y) ** (fromIntegral angularFreq :+ 0) *
         ((x ^ 2 + y ^ 2) :+ 0) **
         (((-(fromIntegral angularFreq) - sigma) :+ fromIntegral radialFreq) /
          2)

{-# INLINE fourierMellinInvPeriod #-}
fourierMellinInvPeriod ::
     (Eq e, RealFloat e) => e -> e -> Int -> Int -> (e, e) -> Complex e
fourierMellinInvPeriod sigma periodEnv angularFreq radialFreq (!x, !y) =
  if x == 0 && y == 0
    then 0
    else (x :+ y) ** (fromIntegral angularFreq :+ 0) *
         ((x ^ 2 + y ^ 2) :+ 0) **
         (((-(fromIntegral angularFreq) - sigma) :+ 2 * pi * fromIntegral radialFreq / (log periodEnv)) /
          2)


{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE KindSignatures             #-}

module Main where

import Control.Monad
import Control.Monad.Random
import Data.List
import Data.Maybe
import Numeric.LinearAlgebra
import System.Environment
import Text.Read

{-
 Weights, Network & neural net implementation by Justin Le (not mine!)

 Tutorial by Justin Le:

 https://blog.jle.im/entry/practical-dependent-types-in-haskell-1.html
-}
data Weights = W { wBiases :: !(Vector Double)  -- n
                 , wNodes  :: !(Matrix Double)  -- n x m
                 }                              -- "m to n" layer
 deriving Show

data Network :: * where
    O     :: !Weights
          -> Network
    (:&~) :: !Weights
          -> !Network
          -> Network
  deriving Show -- ?

infixr 5 :&~

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

logistic' :: Floating a => a -> a
logistic' x = logix * (1 - logix)
  where
    logix = logistic x

runLayer :: Weights -> Vector Double -> Vector Double
runLayer (W wB wN) v = wB + wN #> v

runNet :: Network -> Vector Double -> Vector Double
runNet (O w)      !v = logistic (runLayer w v)
runNet (w :&~ n') !v = let v' = logistic (runLayer w v)
                       in  runNet n' v'

randomWeights :: MonadRandom m => Int -> Int -> m Weights
randomWeights i o = do
    seed1 :: Int <- getRandom
    seed2 :: Int <- getRandom
    let wB = randomVector  seed1 Uniform o * 2 - 1
        wN = uniformSample seed2 o (replicate i (-1, 1))
    return $ W wB wN

randomNet :: MonadRandom m => Int -> [Int] -> Int -> m Network
randomNet i []     o =     O <$> randomWeights i o
randomNet i (h:hs) o = (:&~) <$> randomWeights i h <*> randomNet h hs o

train :: Double           -- ^ learning rate
      -> Vector Double    -- ^ input vector
      -> Vector Double    -- ^ target vector
      -> Network          -- ^ network to train
      -> Network
train rate x0 target = fst . go x0
  where
    go :: Vector Double    -- ^ input vector
       -> Network          -- ^ network to train
       -> (Network, Vector Double)
    -- handle the output layer
    go !x (O w@(W wB wN))
        = let y    = runLayer w x
              o    = logistic y
              -- the gradient (how much y affects the error)
              --   (logistic' is the derivative of logistic)
              -- Here, E := (o - target) ^ 2 is implied
              dEdy = logistic' y * (o - target)
              -- new bias weights and node weights
              wB'  = wB - scale rate dEdy
              wN'  = wN - scale rate (dEdy `outer` x)
              w'   = W wB' wN'
              -- bundle of derivatives for next step
              dWs  = tr wN #> dEdy
          in  (O w', dWs)
    -- handle the inner layers
    go !x (w@(W wB wN) :&~ n)
        = let y          = runLayer w x
              o          = logistic y
              -- get dWs', bundle of derivatives from rest of the net
              (n', dWs') = go o n
              -- the gradient (how much y affects the error)
              dEdy       = logistic' y * dWs'
              -- new bias weights and node weights
              wB'  = wB - scale rate dEdy
              wN'  = wN - scale rate (dEdy `outer` x)
              w'   = W wB' wN'
              -- bundle of derivatives for next step
              dWs  = tr wN #> dEdy
          in  (w' :&~ n', dWs)

{-
 Inserting my own training exercise here.
 This is inspired by the language-model task of predicting the following
 element in a 4-gram, based upon the last three elements available.
 The task: train a randomly generated system of weights to recognize
 the following pattern & predict the following item according to it:

 0 1 2 0 1 2 0 1 2 ..
-}

inf_pattern :: [Int]
inf_pattern = cycle [0..2]

fin_pttrn_len :: Int
fin_pttrn_len = 99999

{-
 Forming training data from the inf_pattern.
 The training data will be a list of pairs, each pair containing
 * a training input, 3-element vector;
 * a target output, 1-element vector.
-}
training_data :: [(Vector Double,  -- 3 * 1
                   Vector Double)] -- 1 * 1
training_data =
  let fin_pattern = take fin_pttrn_len inf_pattern
      list_pairs =
        [ (take 3 sub_pattern, [sub_pattern !! 3])
            | j <- [0..(length fin_pattern - 4)],
              let sub_pattern = drop j fin_pattern ]
  in map (\(inp, out) -> (vector $ fromIntegral <$> inp,
                          vector $ fromIntegral <$> out))
         list_pairs

-- Actually, if we were really being realistic about this then we would
-- make our inputs and outputs work according to some sort of one-hot encoding.
-- But never mind.

-- 1. Get an initial random network. We'll pick a basic enough size:
myRandNet :: MonadRandom m => m Network
myRandNet = randomNet 3 [4, 5] 1

-- 2. How can I get a pre-training prediction, then iteratively
--    train on the next available sample, followed by getting the
--    revised prediction?

-- An auxiliary function for "sequencing together"
-- training a network & printing info about it.
foldTraining :: Double              -- training rate
             -> Network             -- untrained network
             -> [(Vector Double,    -- pairs of training input
                  Vector Double)]   --          and training target output
             -> IO Network          -- train steps & print steps alternate in IO
foldTraining rate net pairs =
  foldl (\io_net sample@(inp, tgt) -> do
            n <- io_net
            putStrLn $ "Will train on: " ++ show inp
            let prediction = runNet n inp
            putStrLn $ "Predicted following item: " ++ show prediction
            -- now trying 10 rounds each per data sample
            let n' = last $
                     take 10 $
                     iterate (train rate inp tgt) n
                          -- ^ :: Network -> Network
            return n')
        (return net)
        pairs

-- We can come back later & modify this to give back the resulting
-- trained network, for use elsewhere (i.e. recurrent NN's!)
displayTraining :: Double              -- training rate
                -> [(Vector Double,    -- pairs of training input
                     Vector Double)]   --          and training target output
                -> IO ()
displayTraining rate pairs = do
  -- pluck out a random net
  randNet <- evalRandIO myRandNet
  -- show what it predicts for 1st sample *before* any training is done
  let inp_0 = fst $ head training_data
  let tgt_0 = snd $ head training_data
  let predict_0 = runNet randNet inp_0
  putStrLn $ "The untrained network predicts that following: "
             ++ show inp_0 ++ ", we will get: " ++ show predict_0
  -- now, alternate showing predictions & performing training on some input
  foldTraining rate randNet pairs
  return ()

{-
 Remark.
 The network is obviously learning *something*, because early predictions
 (which are basically noise) eventually give way to predicting minuscule
 numbers to follow every instance of "2", and numbers infinitesimally
 less than "1" following every instance of "0". But learning that every
 instance of "1" is followed by a "2" seems to take much longer! I haven't
 witnessed it yet despite using up to 10K training sample, with 1 learning
 round per sample.

 What if I try more rounds? It doesn't seem to help as far as training the
  network to detect the "2 is followed by 0" part of the pattern.
-}

{-
 RECURRENT NEURAL NETWORKS.
 We study a simple case of these, Elman network a.k.a. simple RNN.
 * for ea. training sample, we simply take *one element* of the training
   data at a particular point in the sequence (rather than the window-base
   n-gram approach of before.)
 * for the training step at current time t, we rely upon:
    a) the training input at current time t, AND
    b) the activation from previous time (t-1)
        (so we have to initialize a "blank activation" for use
         at the very first step, t := 1; this is h := the zero vector.)
    c) therefore, in a simple model with just *two* sets of weights
       (besides the one which handles the previous activation), the
       back-propagation due to comparing current output vs. current target
       will back-propagate a bundle of "nudges" dh_t given what W2_t
       is at the current time;
       this in turn must back-prop *weight deltas* to
       BOTH U_t (which applies to the previous activation)
       AND W1_t (which applies to the current training input.)
-}

-- RNN ~ Weights x Network.
-- We will model this using a separate collection of weights,
-- U, that we simply carry along while remembering to train.
data RNN = MkRNN {
  wts :: Weights,   -- <- U: h-to-h
  net :: Network }  -- <- (W: i-to-h) :&~ (V: h-to-o)

runRNN :: RNN
       -> Vector Double -- input vector
       -> Vector Double -- previous-step activation
       -> Vector Double -- prediction
runRNN rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN))) x_t h_t_1 =
    let z_u   = (uB + uN #> h_t_1) + -- weighted previous activation
                (wB + wN #> x_t)     -- weighted current training input
        h_t   = logistic z_u         -- current activation.
                                     -- *Need to include somewhere in output!*
        z_h   = (vB + vN #> h_t)     -- next layer of weights
        y_t   = logistic z_h         -- current final activation i.e. output
    in  y_t  -- ^ softmax usually, for tagging problems & similar
             --   (multiple way classification)

{-
 The below is *not* actually a full implementation of the simple RNN
 described at the beginning of Jurafsky Ch. 9, because it does not
 perform the overall training in the two-pass algorithm which is
 explained there.
 In particular, the first pass is more or less all here, but the
 second pass (Backpropgataion Through Time) is *not* taking place right now.
 The idea is actually to *wait* (I think?) until all of y_1, ..., y_l have been
 attempted, *then* (in phase #2) work backwards & successive apply gradient
 terms to V, U, W, U, W, U, W ... (? sequence of matrices being adjusted),
 propagating error back through h_l, h_(l-1), h_(l-2) ... , h_1 *all in one phase*!

 The idea is that, instead of just nudging the matrix entries of U (the recurrent
 part) once per training sample, you would actually wait until the end, and *then*
 have steps iteratively unfolding along the lines of:

 dh_t       = tr vN #> dEdz
 dh_(t-1)   = tr uN #> dEdh_t
 dh_(t-2)   = tr uN #> dEdh_(t-1)
 dh_(t-3)   = tr uN #> dEdh_(t-2)
 ...

-}

trainRNN :: RNN -- ~ U x (W :&~ O V) :
                -- the  graph for this neural network is
                -- as follows:
                {-
                  h_t-1 -> U -> h_t -> V -> y_t
                              ^
                  x     -> W /
                -}
          -> Double        -- training rate
          -> Vector Double -- activation, h_t-1
          -> Vector Double -- training input, x_t
          -> Vector Double -- training target output, tgt_t
          -> (Vector Double, RNN) -- the new activation, h_t, and
                                  -- all three weight-sets of the RNN trained
-- -- don't write recursions, just do everything by hand
trainRNN rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN))) rate h_t_1 x_t tgt_t =
  -- compute the model output
  let z_u   = (uB + uN #> h_t_1) + -- weighted previous activation
              (wB + wN #> x_t)     -- weighted current training input
      h_t   = logistic z_u         -- current activation.
                                   -- *Need to include somewhere in output!*
      z_h   = (vB + vN #> h_t)     -- next layer of weights
      y_t   = logistic z_h         -- current final activation i.e. output
      -- Here, E = (y_t - tgt_t) ^ 2 is implied.
      dEdz = logistic' z_h * (y_t - tgt_t) -- this vector reports, for each
                                            -- component of the final activation (y_t)_j,
                                            -- how much a unit "bump" to it will impact the
                                            -- matching component (E_t)_j of the error,
                                            -- which is not explicitly calculated here.
      -- Bump the preceding-weights bias-vector vB in the opposite direction
      -- of dEdy, & (per component) in proportion with the impact that its change
      -- has on the error E. This is (- scale rate dEdy) (we scale by learning rate)
      vB'   = vB - scale rate dEdz
      -- Bump every matrix entry in vN by the appropriate amount.
      -- Row #1's matrix-entry bumps are given by transposing the layer input (h_t)'s
      -- entries by (- scale rate dEdy)_1, then transposing (so that connections from the
      -- 1st component of h_t will be reinforced appropriately.)
      -- Do this for all entries of "d_vB / d_y = -scale rate dEdy", and you will get
      vN'   = vN + ((- scale rate dEdz) `outer` h_t)
      -- The above two terms vB', vN' specified how to "bump" the entries of the V-layer
      -- bias & weights matrix, given what the layer input h_t currently is.
      -- Now, figure out how to "bump" the entries of the V-layer input h_t,
      -- given what the current V-layer weights vB / vN are.
      -- In fact, only vN will play any role.
      dh_t  = tr vN #> dEdz -- should it be (- scale rate dEdy)?
      -- I think that the learning rate factor is omitted when back-propagating
      -- a new error to the preceding layer, because otherwise the n-th layer back
      -- would get a factor of (rate ^ n) -> 0 as n increases, causing only the
      -- very few final layers to really be modified.

      -- Now, backpropagate this desired change in h_t to train the wts in U, W.
      -- Just as we previously had      "dy := - scale rate dEdy"
      -- so also will we now have       "dEdh := logistic' h_t * dh_t"
      --                                "dh := - scale rate dEdh"
      -- Train both U (whose input was h_t_1) and W (whose input was x_t).
      -- The activation, h_t, on these linear maps' composite result is now
      -- the vector being nudged, because it was the input of the final layer.
      dEdh  = logistic' h_t * dh_t

      uB'   = uB - scale rate dEdh
      uN'   = uN + ((- scale rate dEdh) `outer` h_t_1)
      wB'   = wB - scale rate dEdh
      wN'   = wN + ((- scale rate dEdh) `outer` x_t)  -- x_t *not* h_t / h_t_1
  in  ( h_t,                  -- the new activation
        MkRNN                 -- the new RNN
          (W uB' uN')
          ((W wB' wN') :&~
            O (W vB' vN')))

-- Sample function for using trainRNN one time.
f = do
    rnn_0 <- evalRandIO (randRNN 1 2 5)
    let h_rnn' =
          trainRNN rnn_0 0.25
                   (vector $ fromIntegral <$> [0,0,0,0,0])
                   (vector $ fromIntegral <$> [0])
                   (vector $ fromIntegral <$> [1,4])
    return h_rnn'

-- Need function to product a random RNN.
-- The relevant dimensions need to be:

{-
h_t-1 -> U -> h_t -> V -> y_t
            ^
x     -> W /

For x (input vector)  - dim(x): m
    h (activation)    - dim(h): h
    y (output vector) - dim(y): n
we have that the matrix sizes should be
  W:  h * m
  U:  h * h
  V:  n * h

rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN)))
-}
randRNN :: MonadRandom m => Int   -- input vector dim
                         -> Int   -- output vector dim
                         -> Int   -- hidden vector dim
                         -> m RNN
randRNN i o h = do
  u <- randomWeights h h
  w <- randomWeights i h
  v <- randomWeights h o
  return $ MkRNN u (w :&~ O v)

{-
 A typical form of training data to take will be
 a list [(inp, outp)] with inp  :: Vector Double
                           outp :: Vector Double
-}

foldRNNTraining :: Double
                -> (Vector Double, RNN)    -- initial hidden state & RNN
                -> [(Vector Double,        -- training data
                     Vector Double)]
                -> IO (Vector Double, RNN) -- final hidden state & trained RNN
foldRNNTraining rate (h_0, rnn_0) pairs =
  foldl (\io p@(x, tgt) -> do
              (h, rnn) <- io
              putStrLn $ "Current prediction given: " ++ show x
                         ++ " is: " ++ show (runRNN rnn x h)
              let h_rnn' = trainRNN rnn rate h x tgt
              return h_rnn')
          (return (h_0, rnn_0))
          pairs

displayRNNTraining :: Double            -- training rate
                   -- -> Vector Double     -- initial activation vector h_0
                   -> [(Vector Double,  -- lists of training inputs
                        Vector Double)] -- paired with target outputs
                   -> IO ()
displayRNNTraining rate pairs = do
  -- get a random RNN
  randRNN <- evalRandIO (randRNN 1 1 5)
  -- we will always use an initial activation h_0 of zeroes
  let h_0 = vector $ fromIntegral <$> (take 5 $ repeat 0)
  foldRNNTraining rate (h_0, randRNN) pairs
  return ()

-- no n-gram window this time, just inputs & outputs one at a time each.
rnn_training_data :: [(Vector Double,  -- 1 * 1
                       Vector Double)] -- 1 * 1
rnn_training_data =
  let fin_pattern = take fin_pttrn_len inf_pattern
      list_pairs =
        [ (take 1 sub_pattern, [sub_pattern !! 1])
            | j <- [0..(length fin_pattern - 2)],
              let sub_pattern = drop j fin_pattern ]
  in map (\(inp, out) -> (vector $ fromIntegral <$> inp,
                          vector $ fromIntegral <$> out))
         list_pairs

-- Retry the simple RNN training using softmax.
-- But at the same time, the 1-vectors in the training data
-- [(inp, out)] pairs should become one-hot encoding 3-vectors.
-- Likewise, we will want 'randRNN 3 3 5' (5 the size of the hidden layer.)
softmax :: Vector Double -> Vector Double
softmax v =
  let vexp = exp v :: Vector Double
  in  vexp / (scalar $ sumElements vexp)

-- Note that the derivative of (vector) softmax is matrix-valued
softmax' :: Vector Double -> Matrix Double
softmax' v =
  let s = softmax v
  in  diag s - s `outer` s

runRNN_softmax :: RNN
               -> Vector Double -- input vector
               -> Vector Double -- previous-step activation
               -> Vector Double -- prediction
runRNN_softmax rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN))) x_t h_t_1 =
    let z_u   = (uB + uN #> h_t_1) +
                (wB + wN #> x_t)
        h_t   = logistic z_u
        z_h   = (vB + vN #> h_t)
        y_t   = softmax z_h
    in  y_t  -- ^ softmax usually, for tagging problems & similar
             --   (multiple way classification)

trainRNN_softmax  :: RNN
                  -> Double
                  -> Vector Double
                  -> Vector Double
                  -> Vector Double
                  -> (Vector Double, RNN)
trainRNN_softmax rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN))) rate h_t_1 x_t tgt_t =
  let z_u   = (uB + uN #> h_t_1) +
              (wB + wN #> x_t)
      h_t   = logistic z_u
      z_h   = (vB + vN #> h_t)
      y_t   = softmax z_h                   -- <-
      dEdz  = softmax' z_h #> (y_t - tgt_t)  -- <-
      vB'   = vB - scale rate dEdz
      vN'   = vN + ((- scale rate dEdz) `outer` h_t)
      dh_t  = tr vN #> dEdz
      dEdh  = logistic' h_t * dh_t

      uB'   = uB - scale rate dEdh
      uN'   = uN + ((- scale rate dEdh) `outer` h_t_1)
      wB'   = wB - scale rate dEdh
      wN'   = wN + ((- scale rate dEdh) `outer` x_t)
  in  ( h_t,
        MkRNN
          (W uB' uN')
          ((W wB' wN') :&~
            O (W vB' vN')))

foldRNNTraining_softmax :: Double
                        -> (Vector Double, RNN)    -- initial hidden state & RNN
                        -> [(Vector Double,        -- training data
                             Vector Double)]
                        -> IO (Vector Double, RNN) -- final hidden state & trained RNN
foldRNNTraining_softmax rate (h_0, rnn_0) pairs =
  foldl (\io p@(x, tgt) -> do
              (h, rnn) <- io
              putStrLn $ "Current prediction given: " ++ show x
                         ++ " is: " ++ show (runRNN_softmax rnn x h)
              let h_rnn' = trainRNN_softmax rnn rate h x tgt
              return h_rnn')
          (return (h_0, rnn_0))
          pairs

displayRNNTraining_softmax :: Double
                           -> [(Vector Double,  -- lists of training inputs
                                Vector Double)] -- paired with target outputs
                           -> IO ()
displayRNNTraining_softmax rate pairs = do
  -- get a random RNN
  randRNN <- evalRandIO (randRNN 3 3 7)
  -- we will always use an initial activation h_0 of zeroes
  let h_0 = vector $ fromIntegral <$> (take 7 $ repeat 0)
  foldRNNTraining_softmax rate (h_0, randRNN) pairs
  return ()

--
to_onehot :: [Int] -> Vector Double
to_onehot m = case m of
  [0] -> vector $ fromIntegral <$> [1, 0, 0]
  [1] -> vector $ fromIntegral <$> [0, 1, 0]
  [2] -> vector $ fromIntegral <$> [0, 0, 1]
  -- so that to_onehot won't be partial. Not using this branch
  _   -> vector $ fromIntegral <$> [0, 0, 0]

rnn_training_data_onehot :: [(Vector Double,  -- 3 * 1
                              Vector Double)] -- 3 * 1
rnn_training_data_onehot =
  let fin_pattern = take fin_pttrn_len inf_pattern
      list_pairs =
        [ (take 1 sub_pattern, [sub_pattern !! 1])
            | j <- [0..(length fin_pattern - 2)],
              let sub_pattern = drop j fin_pattern ]
  in map (\(inp, out) -> (to_onehot inp, to_onehot out))
         list_pairs

{-
 To see these examples in action:

 > displayTraining 0.25 training_data
 > displayRNNTraining 0.25 rnn_training_data
 > displayRNNTraining_softmax 0.25 rnn_training_data_onehot

-}

(!!?) :: [a] -> Int -> Maybe a
xs !!? i = listToMaybe (drop i xs)

main :: IO ()
main = return ()