module Autoencoder where

import AI.HNN.FF.Network (Network)
import Numeric.LinearAlgebra.HMatrix (Vector, Matrix)
import System.Exit
import qualified AI.HNN.FF.Network as Network
import qualified Data.Vector
import qualified Numeric.LinearAlgebra.HMatrix as Vector

newtype Encoder = Encoder (Matrix Double)
  deriving (Show, Read, Eq)

newtype Decoder = Decoder (Matrix Double)
  deriving (Show, Read, Eq)

type LayerCount = Int
type EpochCount = Int
type Autoencoder = (Encoder, Decoder)

generateAutoencoderIO :: [Vector Double] -> LayerCount -> EpochCount -> IO Autoencoder
generateAutoencoderIO v l e = do
  net <- Network.createNetwork (Vector.size $ head v) [l] (Vector.size $ head v)
  let Network.Network x = Network.trainNTimes e 0.8 Network.tanh Network.tanh' net (zipWith (,) v v)
  case Data.Vector.toList x of
    [enc, dec] -> return (Encoder enc, Decoder dec)
    _ -> error "colossal failure"

encode :: Encoder -> Vector Double -> Vector Double
encode (Encoder m) v = Vector.app m $ Vector.vjoin [v, 1]

decode :: Decoder -> Vector Double -> Vector Double
decode (Decoder m) v = Vector.app m $ Vector.vjoin [v, 1]


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- Module : Graphics.Color.Model.LCH
module Graphics.Color.Model.LCH
  ( LCH
  -- * Constructors for an LCH color model.
  , pattern ColorLCH
  , pattern ColorLCHA
  , Color
  , ColorModel(..)
  , lch2lxy
  , lxy2lch
  ) where

import Data.Complex ( Complex(..), polar, mkPolar )
import Data.Fixed ( mod' )
import Foreign.Storable
import Graphics.Color.Model.Internal

-----------
--- LCH ---
-----------

-- | CIEL*C*H color model, representing a cylindrical reparameterization
--   of CIEL*a*b* or CIEL*u*v*.
data LCH

-- | `LCH` color model
newtype instance Color LCH e = LCH (V3 e)

-- | Constructor for @LCH@.
pattern ColorLCH :: e -> e -> e -> Color LCH e
pattern ColorLCH l c h = LCH (V3 l c h)
{-# COMPLETE ColorLCH #-}

-- | Constructor for @LCH@ with alpha channel.
pattern ColorLCHA :: e -> e -> e -> e -> Color (Alpha LCH) e
pattern ColorLCHA l c h a = Alpha (ColorLCH l c h) a
{-# COMPLETE ColorLCHA #-}

-- | `LCH` color model
deriving instance Eq e => Eq (Color LCH e)
-- | `LCH` color model
deriving instance Ord e => Ord (Color LCH e)
-- | `LCH` color model
deriving instance Functor (Color LCH)
-- | `LCH` color model
deriving instance Applicative (Color LCH)
-- | `LCH` color model
deriving instance Foldable (Color LCH)
-- | `LCH` color model
deriving instance Traversable (Color LCH)
-- | `LCH` color model
deriving instance Storable e => Storable (Color LCH e)

-- | `LCH` color model
instance Elevator e => Show (Color LCH e) where
  showsPrec _ = showsColorModel

-- | `LCH` color model
instance Elevator e => ColorModel LCH e where
  type Components LCH e = (e, e, e)
  toComponents (ColorLCH l c h) = (l, c, h)
  {-# INLINE toComponents #-}
  fromComponents (l, c, h) = ColorLCH l c h
  {-# INLINE fromComponents #-}

lch2lxy :: Color LCH Double -> Components LCH Double
lch2lxy (ColorLCH l c h) = (l, x, y)
 where
  !h' = h * pi / 180
  (x :+ y) = mkPolar c h'
{-# INLINE lch2lxy #-}

lxy2lch :: Components LCH Double -> Color LCH Double
lxy2lch (l, x, y) = ColorLCH l c h
 where
  (c,h') = polar (x :+ y)
  !h = (h' * 180 / pi) `mod'` 360
{-# INLINE lxy2lch #-}


module Main where

import Criterion.Main
import Data.Packed.Matrix as M
import Numeric.LinearAlgebra.Algorithms   as H  -- hmatrix
import Numeric.LinearAlgebra.Matrix       as LA -- lin-alg
import Numeric.LinearAlgebra.Matrix.Mat44       -- lin-alg

main = do
  let m = (4><4)[1,0,0,0
                ,0,1,0,0
                ,0,0,1,0
                ,0,0,0,1 :: Double]
  let n = Mat44 1 0 0 0
                0 1 0 0
                0 0 1 0
                0 0 0 (1 :: Double)
  putStrLn $ "hmatrix/det:" ++ show (H.det m)
  putStrLn $ "haray/det:" ++ show (LA.det n)
  defaultMain
    [ bgroup "det" [ bench "hmatrix" $ whnf H.det m
                   , bench "haray" $ whnf LA.det n]
    , bgroup "inv" [ bench "hmatrix" $ whnf H.inv m
                   , bench "haray" $ whnf LA.inv44 n]
    ]


{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE TypeOperators       #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
module Test.Grenade.Layers.Internal.Convolution where

import           Grenade.Layers.Internal.Convolution

import           Numeric.LinearAlgebra hiding (uniformSample, konst, (===))

import           Hedgehog
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import qualified Test.Grenade.Layers.Internal.Reference as Reference
import           Test.Hedgehog.Compat

prop_im2col_col2im_symmetrical_with_kernel_stride =
  let factors n = [x | x <- [1..n], n `mod` x == 0]
  in  property $ do
        height   <- forAll $ choose 2 100
        width    <- forAll $ choose 2 100
        kernel_h <- forAll $ (height `div`)    <$> Gen.element (factors height)
        kernel_w <- forAll $ (width `div`)     <$> Gen.element (factors width)
        input    <- forAll $ (height >< width) <$> Gen.list (Range.singleton $ height * width) (Gen.realFloat $ Range.linearFracFrom 0 (-100) 100)

        let stride_h = kernel_h
        let stride_w = kernel_w
        let out      = col2im kernel_h kernel_w stride_h stride_w height width . im2col kernel_h kernel_w stride_h stride_w $ input
        input === out

prop_im2col_col2im_behaves_as_reference =
  let ok extent kernel = [stride | stride <- [1..extent], (extent - kernel) `mod` stride == 0]
  in  property $ do
        height   <- forAll (choose 2 100)
        width    <- forAll (choose 2 100)
        kernel_h <- forAll (choose 2 (height - 1))
        kernel_w <- forAll (choose 2 (width - 1))
        stride_h <- forAll (Gen.element (ok height kernel_h))
        stride_w <- forAll (Gen.element (ok width kernel_w))
        input    <- forAll ((height >< width) <$> Gen.list (Range.singleton $ height * width) (Gen.realFloat $ Range.linearFracFrom 0 (-100) 100))

        let outFast       = im2col kernel_h kernel_w stride_h stride_w input
        let retFast       = col2im kernel_h kernel_w stride_h stride_w height width outFast

        let outReference  = Reference.im2col kernel_h kernel_w stride_h stride_w  input
        let retReference  = Reference.col2im kernel_h kernel_w stride_h stride_w height width outReference

        outFast === outReference
        retFast === retReference

tests :: IO Bool
tests = checkParallel $$(discover)


{-# LANGUAGE BangPatterns #-}
module Main where

import qualified Data.Vector                   as V
import qualified Data.Vector.Mutable           as VM
import qualified Data.Vector.Unboxed           as VU
import qualified Data.Vector.Unboxed.Mutable   as VUM
import qualified Data.Vector.Generic           as VG
import qualified Data.Vector.Generic.Mutable   as VGM
import qualified Data.ByteString.Char8         as BS
import qualified Data.ByteString.Lazy.Char8    as BSL
import qualified Data.ByteString.Builder       as BSB
import           Data.Char                      ( isSpace )
import           Data.List                      ( unfoldr )
import           Data.Function
import           Data.Complex
import           Control.Monad
import           Control.Monad.Primitive
import           Control.Monad.ST

main :: IO ()
main = do
  n   <- readLn :: IO Int
  abs <-
    VU.replicateM n
    $   (\v -> (v VU.! 0, v VU.! 1))
    .   VU.unfoldr (BS.readInt . BS.dropWhile isSpace)
    <$> BS.getLine
  let v = uncurry multiplyCT $ VU.unzip abs
  print 0
  VU.mapM_ print v

-- O(n^2) multiply polynomials
naiveMultiply :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
naiveMultiply f g = VU.create $ do
  let len = VU.length f + VU.length g - 1
  h <- VUM.new len
  flip VU.imapM_ f $ \i x ->
    flip VU.imapM_ g $ \j y -> VUM.modify h (\z -> z + x * y) (i + j)
  return h

-- O(n log n) multiply polynomials by FFT （Immutable Vector)
-- 859 ms in atc001-c
-- c.f. https://qiita.com/ageprocpp/items/0d63d4ed80de4a35fe79
multiply :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
{-# INLINE multiply #-}
multiply !f !g =
  let !len = VU.length f + VU.length g - 1
      !sz  = fix (\rec !k -> if k >= len then k else rec (k * 2)) 1
      !f'  = VU.map fromIntegral f
        VU.++ VU.replicate (sz - VU.length f) (0 :: Complex Double)
      !g' = VU.map fromIntegral g
        VU.++ VU.replicate (sz - VU.length g) (0 :: Complex Double)
  in  VU.map (round . magnitude . (/ fromIntegral sz))
        . VU.take len
        . dft (-1)
        $ VU.zipWith (*) (dft 1 f') (dft 1 g')

-- DFT algorithm
-- if sign == 1 then DFT
-- if sign == -1 then inverse DFT
dft :: Double -> VU.Vector (Complex Double) -> VU.Vector (Complex Double)
{-# INLINE dft #-}
dft !sign !f
  | VU.length f == 1
  = f
  | otherwise
  = let !n    = VU.length f
        !nd2  = n `div` 2
        !f0   = dft sign $ VU.map (f VU.!) $ VU.enumFromStepN 0 2 nd2
        !f1   = dft sign $ VU.map (f VU.!) $ VU.enumFromStepN 1 2 nd2
        -- f0 = dft sign $ VU.ifilter (\i _ -> i `mod` 2 == 0) f
        -- f1 = dft sign $ VU.ifilter (\i _ -> i `mod` 2 == 1) f
        !zeta = cis (sign * 2 * acos (-1) / fromIntegral n) :: Complex Double
    in  VU.imap (\i z -> (f0 VU.! (i `mod` nd2)) + z * (f1 VU.! (i `mod` nd2)))
          $ VU.iterateN n (* zeta) 1

-- O(n log n) multiply polynomials by FFT
-- c.f. https://qiita.com/ageprocpp/items/0d63d4ed80de4a35fe79
-- multiply :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
-- multiply f g = runST $ do
--   let len = VU.length f + VU.length g - 1
--   let sz = fix (\rec k -> if k >= len then k else rec (k*2)) 1
--   tf <- VU.thaw (VU.map fromIntegral f :: VU.Vector (Complex Double))
--   nf <- VUM.grow tf (sz - VU.length f)
--   tg <- VU.thaw (VU.map fromIntegral g :: VU.Vector (Complex Double))
--   ng <- VUM.grow tg (sz - VU.length g)
--   dft nf 1 sz
--   dft ng 1 sz
--   forM_ [0..sz-1] $ \i -> do
--     y <- VUM.read ng i
--     VUM.modify nf (* y) i
--   dft nf (-1) sz
--   ret <- VU.freeze nf
--   return $ VU.map ((`div` sz) . round . realPart) $ VU.take len ret

-- DFT algorithm
-- if sign == 1 then DFT
-- if sign == -1 then inverse DFT
-- dft :: PrimMonad m => VUM.MVector (PrimState m) (Complex Double) -> Double -> Int -> m ()
-- {-# INLINE dft #-}
-- dft f sign 1 = return ()
-- dft f sign n = do
--   let nd2 = n `div` 2
--   f0 <- VUM.new nd2
--   f1 <- VUM.new nd2
--   forM_ [0..nd2-1] $ \i -> do
--     x0 <- VUM.read f (2*i)
--     VUM.write f0 i x0
--     x1 <- VUM.read f (2*i+1)
--     VUM.write f1 i x1
--   dft f0 sign nd2
--   dft f1 sign nd2
--   let zeta = cis (sign * 2 * acos (-1) / fromIntegral n)
--   let zetas = VU.iterateN n (*zeta) 1
--   flip VU.imapM_ zetas $ \i z -> do
--     x0 <- VUM.read f0 (i `mod` nd2)
--     x1 <- VUM.read f1 (i `mod` nd2)
--     VUM.write f i (x0 + z * x1)

-- O(n log n) multiply polynomials by Cooley-Tukey FFT （Mutable Vector)
-- 448 ms in atc001-c
-- c.f. http://wwwa.pikara.ne.jp/okojisan/stockham/cooley-tukey.html

-- This function computes polynomial coefficient of (f * g)(x)
-- f & g are polynomial coefficient array of f(x) & g(x)
multiplyCT :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
multiplyCT f g = runST $ do
  let len = VU.length f + VU.length g - 1
  let sz  = fix (\rec k -> if k >= len then k else rec (k * 2)) 1
  tf <- VU.thaw (VU.map fromIntegral f :: VU.Vector (Complex Double))
  nf <- VUM.grow tf (sz - VU.length f)
  tg <- VU.thaw (VU.map fromIntegral g :: VU.Vector (Complex Double))
  ng <- VUM.grow tg (sz - VU.length g)
  fft sz nf
  fft sz ng
  forM_ [0 .. sz - 1] $ \i -> do
    y <- VUM.read ng i
    VUM.modify nf (* y) i
  ifft sz nf
  ret <- VU.freeze nf
  return $ VU.map (round . realPart) $ VU.take len ret
{-# INLINE multiplyCT #-}

fft
  :: PrimMonad m
  => Int -- length
  -> VUM.MVector (PrimState m) (Complex Double) -- target data of FFT
  -> m ()
fft n = cooleyTukey n 1 0 0
{-# INLINE fft #-}

ifft
  :: PrimMonad m
  => Int -- length
  -> VUM.MVector (PrimState m) (Complex Double) -- target data of FFT
  -> m ()
ifft n x = do
  forM_ [0 .. n - 1] $ VUM.modify x conjugate
  cooleyTukey n 1 0 0 x
  forM_ [0 .. n - 1] $ VUM.modify x ((/ fromIntegral n) . conjugate)
{-# INLINE ifft #-}

cooleyTukey
  :: PrimMonad m
  => Int -- length
  -> Int -- stride
  -> Int -- start position of block
  -> Int -- temp var for bit reverse position in start position of block
  -> VUM.MVector (PrimState m) (Complex Double) -- target data of FFT
  -> m ()
cooleyTukey n s q d x
  | n > 1 = do
    let m  = n `div` 2
    let θ0 = 2 * pi / fromIntegral n
    forM_ [0 .. m - 1] $ \p -> do
      let wp = cos (fromIntegral p * θ0) :+ (-sin (fromIntegral p * θ0))
      a <- VUM.read x (q + p + 0)
      b <- VUM.read x (q + p + m)
      VUM.write x (q + p + 0) (a + b)
      VUM.write x (q + p + m) ((a - b) * wp)
    cooleyTukey m (2 * s) (q + 0) (d + 0) x
    cooleyTukey m (2 * s) (q + m) (d + s) x
  | q > d = VUM.swap x q d
  | otherwise = return ()
{-# INLINE cooleyTukey #-}


module TwoMoon.Load where

  import Neural
  import Numeric.LinearAlgebra

  -- Parse two-moon dataset into a list of examples (pairs of a list of feature and a list of labels)
  loadExamples :: IO [Example]
  loadExamples = do
    twomoon <- loadMatrix "data/two_moon.txt"
    let x = vector `map` toLists (twomoon ¿ [0, 1])
    let gt = mkOneHot `map` toList (flatten (twomoon ¿ [2]))
    return $ zip x gt
      where
        mkOneHot :: Double -> Vector Double
        mkOneHot x = if x == 0.0 then vector [1.0, 0.0] else vector [0.0, 1.0]


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

module Main where

import Codec.Compression.GZip (decompress)
import Numeric.LinearAlgebra as V
import qualified Data.ByteString.Lazy as BS
import qualified Data.ByteString as B
import Data.Bifunctor
import Data.Functor
import Data.Word8
import qualified Data.Vector.Storable as V
import Control.Monad
import Control.Arrow
import Control.Applicative
import Pipes as P
import qualified Pipes.Prelude as PP
import System.Random
import Parser
import Data.Network
import Data.Network.FullyConnected as FC
import Algorithm.Training
import Algorithm.GradientDescent
import Data.Network.Activation
import Data.Network.Utils (makeBatchP, repeatM, squeeze, unzipP, integersP, toGaussian, averageP, accuracyBatch)

-- import System.Random
-- import Data.Binary.Get
-- import Control.Monad.Except as E
-- import qualified Data.Vector as V

main = do
  putStrLn "Extracting MNIST dataset"
  s <- decompress <$> BS.readFile "train-images-idx3-ubyte.gz" -- images data lazy bytestring
  l <- decompress <$> BS.readFile "train-labels-idx1-ubyte.gz" -- labels data lazy bytestring

  s' <- decompress <$> BS.readFile "t10k-images-idx3-ubyte.gz" -- images data lazy bytestring
  l' <- decompress <$> BS.readFile "t10k-labels-idx1-ubyte.gz" -- labels data lazy bytestring

  let mnisttraining = parseMNISTSet s l
  case mnisttraining of
    Right ((nrows, ncols), nsamples, samples) -> do
      putStrLn "MNIST training dataset succesfully parsed!"
      putStrLn $ show nsamples ++ " training samples of " ++ show nrows ++ "x" ++ show ncols ++ " labeled images."
      -- let s = take 100 samples
      -- forM_ s $ \(im, l) -> do
      --   putStrLn $ renderImage (nrows, ncols) im
      --   putStrLn $ "Target: " ++ show l
      --   putStrLn $ "Label: " ++ show (V.maxIndex l)
      let mnisttest = parseMNISTSet s' l'
      case mnisttest of
        Right ((nrows', ncols'), nsamples', samples') -> do
          putStrLn "MNIST test dataset succesfully parsed!"
          putStrLn $ show nsamples' ++ " test samples of " ++ show nrows' ++ "x" ++ show ncols' ++ " labeled images."

          let batchsize = 20
              -- Sample producer
              trainingSamplesP = each samples :: Producer MNISTSample IO ()
              (imageP, labelP) = unzipP trainingSamplesP :: (Producer Image IO (), Producer Label IO ())
              -- Batch sample producer(pairs of matrix)
              trainingBatches = PP.zip (imageP >-> (makeBatchP batchsize)) (labelP  >-> (makeBatchP batchsize))
              testSamplesP    = each samples' :: Producer MNISTSample IO ()
              testBatches     = uncurry PP.zip $ bimap (>-> makeBatchP batchsize) (>-> makeBatchP batchsize) (unzipP testSamplesP)
          g <- getStdGen
          -- let xavier = fmap (\(FC w b af) -> FC (scale (sqrt $ recip $ fromIntegral $ (uncurry (+)) (size w)) w) b af)
          let !(model :: FeedForwardFC Batch) = (makeRandomNetwork g id id 784 [(1000, logistic)] (10, softmaxMat))
              -- !model' = fmap (\(FC w b af) -> let (u,_,_) = svd w in (FC u b af)) model
              lf = mse
              lrf = const 0.05 -- \n -> exp (-(1.0 + fromIntegral n))
              tolerance = 0.0001
              test m = do
                (i,t) <- await
                let r = runNetwork m i
                    loss = evaluate lf (t, fst r)
                    a    = accuracyBatch (fst r, t)
                yield ((toDouble loss), a)
              reporter (i,t) ts = do
                putStrLn $ show ts
              interval = 10
              stopper = (uncurry (liftM2 (<|>)) . ((closeEnough tolerance) &&& (longEnough 100)))
              loop m = do -- the actual training steps
                s@(sid, (i,t)) <- await
                -- lift $ putStrLn $ "Input " ++ show (fst s)
                r <- lift $ runEffect $ (stepP m lf 0 lrf s) >-> reportP interval (reporter (i,t)) >->
                  (descendP (stopper))
                -- lift $ putStrLn ""
                yield r
                loop (tsmodel r)

              trainingChain m = PP.zip (integersP 0) trainingBatches >-> (PP.take 100) >-> (loop m)
              -- Train over all training samples, then test with testing dataset, for each iteration
              -- One iteration = one epoch
              loop' m n = do
                lift $ putStrLn $ "Epoch: " ++ show n
                (Just r) <- lift $ PP.last (trainingChain m)
                yield r
                let (lossP, accuracyP) = unzipP $ testBatches >-> test (tsmodel r)
                (loss, accuracy) <- lift $ liftM2 (,) (averageP lossP) (averageP accuracyP)
                lift $ putStrLn $ "Test loss: " ++ show (fst loss)
                lift $ putStrLn $ "Test accuracy: " ++ show (fst accuracy)
                loop' (tsmodel r) (succ n)

          runEffect $ for ((loop' model 0) >-> (PP.take 100)) $ \ts -> do
            lift $ putStrLn ""
            lift $ print ts

      putStrLn "Done!"
    Left er -> putStrLn er


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}

module Backprop.Learn.Model.Function (
  -- * Types
    ParamFunc(..)
  , ParamFuncP, pattern PFP, runParamFuncP
  , FixedFunc, pattern FF, runFixedFunc
  , paramMap, learnParam
  , idLearn
  -- ** Combinators
  , dimapPF, lmapPF, rmapPF, compPF, parPF
  -- *** Chain
  , (.-), nilPF, onlyPF
  -- * Activation functions
  -- | See <https://en.wikipedia.org/wiki/Activation_function>
  --
  -- ** Maps
  -- *** Unparameterized
  , step
  , logistic
  , softsign
  , reLU
  , softPlus
  , bentIdentity
  , siLU
  , softExponential
  , sinc
  , gaussian
  , tanh
  , atan
  , sin
  , vmap
  , vmap'
  -- *** Parameterized
  , liftUniform
  , isru
  , preLU
  , sreLU
  , sreLUPFP
  , eLU
  , isrLU
  , apl
  , aplPFP
  -- ** Mixing
  , softMax
  , maxout
  , kSparse
  ) where

import           Backprop.Learn.Model.Class
import           Control.Category
import           Data.Foldable
import           Data.Proxy
import           Data.Type.Length
import           Data.Type.Mayb hiding                        (type (<$>))
import           Data.Type.Tuple hiding                       (T2(..), T3(..))
import           Data.Typeable
import           GHC.TypeNats
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop hiding (tr)
import           Prelude hiding                               ((.), id)
import           Type.Class.Known
import           Type.Family.List
import qualified Data.Type.Tuple                              as T
import qualified Data.Vector.Sized                            as SV
import qualified Data.Vector.Storable.Sized                   as SVS
import qualified Numeric.LinearAlgebra                        as HU
import qualified Numeric.LinearAlgebra.Static                 as H
import qualified Numeric.LinearAlgebra.Static.Vector          as H

-- | An unparameterized function.  Has a 'Category' instance.
--
-- A @'FixedFunc' a b@ essentially the same as a:
--
-- @
-- forall s. 'Reifies' s 'W' => 'BVar' s a -> 'BVar' s b
-- @
--
-- And the 'FF' pattern (and 'runFixedFunc' extractor) witness this.
--
-- It is usually better to just use the functions directly, with
-- combinators like 'Dimap', 'LMap', 'RMap', 'dimapPF', 'lmapPF', 'rmapPF',
-- etc.  This is just provided to let you work nicely with 'ParamFunc'
-- combinators.
type FixedFunc = ParamFunc 'Nothing

-- | 'FF' and 'runFixedFunc' witness the fact that a @'FixedFunc' a b@ is
-- just a function from @'BVar' s a@ to @'BVar' s b@.
pattern FF :: (forall s. Reifies s W => BVar s a -> BVar s b) -> FixedFunc a b
pattern FF { runFixedFunc } <- (getFF->runFixedFunc) where
    FF f = PF (const f)
{-# COMPLETE FF #-}

getFF :: forall a b. FixedFunc a b -> (forall s. Reifies s W => BVar s a -> BVar s b)
getFF ff = runParamFunc ff N_

-- | Identity model, useful for using with other combinators.
idLearn :: FixedFunc a a
idLearn = FF id

-- | A @'ParamFunc' p a b@ is a parameterized function from @a@ to @b@,
-- potentially with trainable parameter @p@.
--
-- A utility wrapper for a deterministic and stateless model.
newtype ParamFunc p a b =
    PF { runParamFunc :: forall s. Reifies s W => Mayb (BVar s) p -> BVar s a -> BVar s b
       }
  deriving (Typeable)

instance ( Num a, Num b, MaybeC Num p, KnownMayb p)
      => Learn a b (ParamFunc p a b) where
    type LParamMaybe (ParamFunc p a b) = p

    runLearn p = stateless . runParamFunc p

instance Category (ParamFunc p) where
    id = PF (const id)
    f . g = PF $ \p -> runParamFunc f p
                     . runParamFunc g p

-- | Convenient type synonym for a 'ParamFunc' with parameters.
--
-- Mostly made to be easy to construct/deconstruct with 'PFP' and
-- 'runParamFuncP'.
type ParamFuncP p = ParamFunc ('Just p)

pattern PFP :: (forall s. Reifies s W => BVar s p -> BVar s a -> BVar s b)
            -> ParamFuncP p a b
pattern PFP { runParamFuncP } <- (getPFP->(getWF->runParamFuncP))
  where
    PFP f = PF $ \(J_ p) -> f p
{-# COMPLETE PFP #-}

newtype WrapFun p a b = WF { getWF :: forall s. Reifies s W => BVar s p -> BVar s a -> BVar s b }

getPFP :: ParamFuncP p a b -> WrapFun p a b
getPFP pf = WF (\case p -> runParamFunc pf (J_ p))

-- | Create a 'ParamFunc' from any instance of 'Learn' that does not have
-- state.
learnParam
    :: forall l a b. (Learn a b l, NoState l)
    => l
    -> ParamFunc (LParamMaybe l) a b
learnParam l = PF (runLearnStateless l)

-- | 'ParamFuncP' taking a singleton list; meant to be used with '.-'
onlyPF
    :: forall p a b. (KnownMayb p, MaybeC Backprop p)
    => ParamFunc p a b
    -> ParamFuncP (T (MaybeToList p)) a b
onlyPF f = PFP $ \ps -> case knownMayb @p of
                          N_   -> runParamFunc f N_
                          J_ _ -> runParamFunc f (J_ (isoVar tOnly onlyT ps))

-- | Compose two 'ParamFuncP's on lists.
(.-)
    :: forall ps qs a b c. (ListC (Backprop <$> ps), ListC (Backprop <$> qs), Known Length ps)
    => ParamFuncP (T ps) b c
    -> ParamFuncP (T qs) a b
    -> ParamFuncP (T (ps ++ qs)) a c
f .- g = PFP $ \ps -> runParamFuncP f (ps ^^. tTake @ps @qs known)
                    . runParamFuncP g (ps ^^. tDrop @ps @qs known)
infixr 9 .-

-- | The identity of '.-'
nilPF :: ParamFuncP (T '[]) a a
nilPF = id

-- | Pre- and post-compose a 'ParamFunc'
dimapPF
    :: (forall s. Reifies s W => BVar s a -> BVar s b)
    -> (forall s. Reifies s W => BVar s c -> BVar s d)
    -> ParamFunc p b c
    -> ParamFunc p a d
dimapPF f g h = PF $ \p -> g . runParamFunc h p . f

-- | Precompose a 'ParamFunc'
lmapPF
    :: (forall s. Reifies s W => BVar s a -> BVar s b)
    -> ParamFunc p b c
    -> ParamFunc p a c
lmapPF f = dimapPF f id

-- | Postcompose a 'ParamFunc'
rmapPF
    :: (forall s. Reifies s W => BVar s b -> BVar s c)
    -> ParamFunc p a b
    -> ParamFunc p a c
rmapPF = dimapPF id

-- | Compose two 'ParamFunc's sequentially
compPF
    :: forall p q a b c. (MaybeC Backprop p, MaybeC Backprop q, KnownMayb p, KnownMayb q)
    => ParamFunc p a b
    -> ParamFunc q b c
    -> ParamFunc (TupMaybe p q) a c
compPF f g = PF $ \pq ->
    let (p, q) = splitTupMaybe @_ @p @q (\(T2B p' q') -> (p', q')) pq
    in  runParamFunc g q . runParamFunc f p

-- | Compose two 'ParamFunc's in parallel
parPF
    :: forall p q a b c d.
     ( MaybeC Backprop p
     , MaybeC Backprop q
     , KnownMayb p
     , KnownMayb q
     , Backprop a
     , Backprop b
     , Backprop c
     , Backprop d
     )
    => ParamFunc p a c
    -> ParamFunc q b d
    -> ParamFunc (TupMaybe p q) (T.T2 a b) (T.T2 c d)
parPF f g = PF $ \pq (T2B x y) ->
    let (p, q) = splitTupMaybe @_ @p @q (\(T2B p' q') -> (p', q')) pq
    in  T2B (runParamFunc f p x)
            (runParamFunc g q y)

-- TODO: replace all of these with manual ops?

-- | Softmax normalizer
softMax
    :: (KnownNat i, Reifies s W)
    => BVar s (R i)
    -> BVar s (R i)
softMax x = expx / konst (norm_1V expx)
  where
    expx = exp x

-- | Logistic function
--
-- \[
-- \sigma(x) = \frac{1}{1 + e^{-x}}
-- \]
logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

-- | Binary step / heaviside step
--
-- To use with vectors ('R'), use 'vmap''.
--
-- \[
-- \begin{cases}
-- 0 & \text{for } x < 0 \\
-- 1 & \text{for } x \ge 0
-- \end{cases}
-- \]
step :: (Ord a, Num a) => a -> a
step x | x < 0     = 0
       | otherwise = 1

-- | Softsign activation function
--
-- \[
-- \frac{x}{1 + \lvert x \rvert}
-- \]
softsign :: Fractional a => a -> a
softsign x = x / (1 + abs x)

-- | Inverse square root unit
--
-- \[
-- \frac{x}{\sqrt{1 + \alpha x^2}}
-- \]
--
-- See 'liftUniform' to make this compatible with 'PFP'.
--
-- You can also just use this after partially applying it, to fix the
-- parameter (and not have it trained).
isru
    :: Floating a
    => a        -- ^ α (scaling parameter)
    -> a        -- ^ x
    -> a
isru α x = x / sqrt (1 + α * x * x)

-- | Rectified linear unit.
--
-- To use with vectors ('R'), use 'vmap''.
--
--
-- \[
-- \max(0,x)
-- \]
--
-- @
-- 'reLU' = 'preLU' 0
-- @
--
reLU :: (Num a, Ord a) => a -> a
reLU x | x < 0     = 0
       | otherwise = x

-- | Parametric rectified linear unit
--
-- To use with vectors ('R'), use 'vmap''.
--
-- If scaling parameter is a fixed (and not learned) parameter, this is
-- typically called a leaky recitified linear unit (typically with
-- α = 0.01).
--
-- To use as a learned parameter:
--
-- @
-- 'vmap' . 'preLU' :: 'BVar' s Double -> 'BVar' s ('R' n) -> BVar s (R n)
-- @
--
-- This can be give directly to 'PFP'.
--
-- To fix the paramater ("leaky"), just partially apply a parameter:
--
-- @
-- 'preLU' 0.01           :: 'BVar' s ('R' n) -> BVar s (R n)
-- preLU ('realToFrac' α) :: BVar s (R n) -> BVar s (R n)
-- @
--
-- See also 'rreLU'.
--
-- \[
-- \begin{cases}
-- \alpha x & \text{for } x < 0 \\
-- x        & \text{for } x \ge 0
-- \end{cases}
-- \]
preLU
    :: (Num a, Ord a)
    => a        -- ^ α (scaling parameter)
    -> a        -- ^ x
    -> a
preLU α x | x < 0     = α * x
          | otherwise = x

-- | Exponential linear unit
--
-- To use with vectors ('R'), use 'vmap''.
--
-- To use as a learned parameter:
--
-- @
-- 'vmap' . 'eLU' :: 'BVar' s Double -> 'BVar' s ('R' n) -> BVar s (R n)
-- @
--
-- This can be give directly to 'PFP'.
--
-- To fix the paramater, just partially apply a parameter:
--
-- @
-- 'vmap'' ('eLU' 0.01) :: 'BVar' s ('R' n) -> BVar s (R n)
-- @
--
-- \[
-- \begin{cases}
-- \alpha (e^x - 1) & \text{for } x < 0 \\
-- x                & \text{for } x \ge 0
-- \end{cases}
-- \]
eLU :: (Floating a, Ord a)
    => a        -- ^ α (scaling parameter)
    -> a        -- ^ x
    -> a
eLU α x | x < 0     = α * (exp x - 1)
        | otherwise = x

-- | S-shaped rectified linear activiation unit
--
-- See 'sreLUPFP' for an uncurried and uniformly lifted version usable with
-- 'PFP'.
--
-- \[
-- \begin{cases}
-- t_l + a_l (x - t_l) & \text{for } x \le t_l  \\
-- x                   & \text{for } t_l < x < t_r  \\
-- t_r + a_r (x - t_r) & \text{for } x \ge t_r
-- \end{cases}
-- \]
sreLU
    :: (Num a, Ord a)
    => a        -- ^ @t_l@
    -> a        -- ^ @a_l@
    -> a        -- ^ @t_r@
    -> a        -- ^ @a_l@
    -> a        -- ^ @x@
    -> a
sreLU tl al tr ar x
    | x < tl    = tl + al * (x - tl)
    | x > tr    = tr + ar * (x - tr)
    | otherwise = x

-- | An uncurried and uniformly lifted version of 'sreLU' directly usable
-- with 'PFP'.
sreLUPFP
    :: (KnownNat n, Reifies s W)
    => BVar s (T.T2 (T.T2 Double Double) (T.T2 Double Double))
    -> BVar s (R n)
    -> BVar s (R n)
sreLUPFP (T2B (T2B tl al) (T2B tr ar)) = vmap (sreLU tl al tr ar)

-- | Inverse square root linear unit
--
-- To use with vectors ('R'), use 'vmap''.
--
-- \[
-- \begin{cases}
-- \frac{x}{1 + \alpha x^2} & \text{for } x < 0 \\
-- x                        & \text{for } x \ge 0
-- \end{cases}
-- \]
isrLU
    :: (Floating a, Ord a)
    => a        -- ^ α
    -> a        -- ^ x
    -> a
isrLU α x
    | x < 0     = x / sqrt (1 + α * x * x)
    | otherwise = x

-- | Adaptive piecewise linear activation unit
--
-- See 'aplPFP' for an uncurried version usable with 'PFP'.
--
-- \[
-- \max(0, x_i) + \sum_j^M a_i^j \max(0, -x_i + b_i^j)
-- \]
apl :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (L n m)     -- ^ a
    -> BVar s (L n m)     -- ^ b
    -> BVar s (R m)       -- ^ x
    -> BVar s (R m)
apl as bs x = vmap' (max 0) x
            + sum (toRows (as * (bs - fromRows (SV.replicate x))))

-- | 'apl' uncurried, to be directly usable with 'PFP'.
aplPFP
    :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (T.T2 (L n m) (L n m))
    -> BVar s (R m)
    -> BVar s (R m)
aplPFP (T2B a b) = apl a b

-- | SoftPlus
--
-- \[
-- \ln(1 + e^x)
-- \]
softPlus :: Floating a => a -> a
softPlus x = log (1 + exp x)

-- | Bent identity
--
-- \[
-- \frac{\sqrt{x^2 + 1} - 1}{2} + x
-- \]
bentIdentity :: Floating a => a -> a
bentIdentity x = (sqrt (x * x + 1) - 1) / 2 + x

-- | Sigmoid-weighted linear unit.  Multiply 'logistic' by its input.
--
-- \[
-- x \sigma(x)
-- \]
siLU :: Floating a => a -> a
siLU x = x * logistic x

-- | SoftExponential
--
-- To use with vectors ('R'), use 'vmap''.
--
-- \[
-- \begin{cases}
-- - \frac{\ln(1 - \alpha (x + \alpha))}{\alpha} & \text{for } \alpha < 0  \\
-- x                                             & \text{for } \alpha = 0  \\
-- \frac{e^{\alpha x} - 1}{\alpha} + \alpha      & \text{for } \alpha > 0
-- \end{cases}
-- \]
softExponential
    :: (Floating a, Ord a)
    => a            -- ^ α
    -> a            -- ^ x
    -> a
softExponential α x = case compare α x of
    LT -> - log (1 - α * (x + α)) / α
    EQ -> x
    GT -> (exp (α * x) - 1) / α + α

-- | Sinc
--
-- \[
-- \begin{cases}
-- 1                 & \text{for } x = 0  \\
-- \frac{\sin(x)}{x} & \text{for } x \ne 0
-- \end{cases}
-- \]
sinc :: (Floating a, Eq a) => a -> a
sinc 0 = 1
sinc x = sin x / x

-- | Gaussian
--
-- \[
-- e^{-x^2}
-- \]
gaussian :: Floating a => a -> a
gaussian x = exp (- (x * x))

-- | Maximum of vector.
--
-- Compare to 'norm_InfV', which gives the maximum absolute value.
maxout :: (KnownNat n, Reifies s W) => BVar s (R n) -> BVar s Double
maxout = liftOp1 . op1 $ \x ->
    let n = HU.maxElement . H.extract $ x
    in  ( n
        , \d -> H.vecR . SVS.map (\e -> if e == n then d else 0) . H.rVec $ x
        )

-- | Usable with functions like '*', 'isru', etc. to turn them into a form
-- usable with 'PFP':
--
-- @
-- 'liftUniform' ('*')  :: 'BVar' s 'Double' -> BVar s ('R' n) -> BVar s (R n)
-- liftUniform 'isru' :: BVar s Double -> BVar s (R n) -> BVar s (R n)
-- @
--
-- Basically turns a parmaeterized function on individual elements of
-- into one that shares the same parameter across all elements of the
-- vector.
liftUniform
    :: (Reifies s W, KnownNat n)
    => (BVar s (R n) -> r)
    -> BVar s Double
    -> r
liftUniform f = f . konst

-- | Utility function to make a 'ParamFunc' that maps a parameterized
-- function over every item in the 'R'.  The parameter is shared across the
-- entire map, and trained.
paramMap
    :: KnownNat i
    => (forall s. Reifies s W => Mayb (BVar s) p -> BVar s Double -> BVar s Double)
    -> ParamFunc p (R i) (R i)
paramMap f = PF (vmap . f)

-- -- TODO: vmap can do better.

-- | Keep only the top @k@ values, and zero out all of the rest.
--
-- Useful for postcomposing in between layers (with a logistic function
-- before) to encourage the number of "activated" neurons is kept to be
-- around @k@.  Used in k-Sprase autoencoders (see 'KAutoencoder').
--
-- <http://www.ericlwilkinson.com/blog/2014/11/19/deep-learning-sparse-autoencoders>
kSparse
    :: forall n s. (Reifies s W, KnownNat n)
    => Int                      -- ^ number of items to keep
    -> BVar s (R n)
    -> BVar s (R n)
kSparse k = liftOp1 . op1 $ \xs ->
    let xsSort = HU.sortVector (H.extract xs)
        thresh = xsSort `HU.atIndex` (n - k)
        mask   = H.dvmap (\x -> if x >= thresh then 1 else 0) xs
    in  ( H.dvmap (\x -> if x >= thresh then x else 0) xs
        , (* mask)
        )
  where
    n = fromIntegral $ natVal (Proxy @n)


{-# LANGUAGE FlexibleContexts #-}
module FourierPinwheel.Filtering where

import           Data.Array.Repa            as R
import           Data.Complex
import           Data.List                  as L
import           Data.Vector.Generic        as VG
import           Data.Vector.Storable       as VS
import           Data.Vector.Unboxed        as VU
import           DFT.Plan
import           Filter.Utils
import           Numeric.GSL.Special.Bessel
import           Utils.List

{-# INLINE idealLowPassFilter #-}
idealLowPassFilter ::
     (VG.Vector vector (Complex Double))
  => Double
  -> Double
  -> Int
  -> Int
  -> vector (Complex Double)
idealLowPassFilter radius period numRhoFreqs numRFreqs =
  let centerRhoFreq = div numRhoFreqs 2
      centerRFreq = div numRFreqs 2
      sinc x =
        if x == 0
          then 1
          else sin (pi * x) / (pi * x)
   in VG.convert .
      toUnboxed . computeS . fromFunction (Z :. numRFreqs :. numRhoFreqs) $ \(Z :. i' :. j') ->
        let i = i' - centerRFreq
            j = j' - centerRhoFreq
         in -- if i == (j)
            --   then
          (2 * radius) / period *
                   sinc (fromIntegral j  / period * (2 * radius)) :+
                   0
              -- else 0

{-# INLINE applyFilterRho #-}
applyFilterRho ::
     DFTPlan
  -> VU.Vector (Complex Double)
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
applyFilterRho plan filterF arr' = do
  let (Z :. numRFreqs :. numThetaFreqs :. numRhoFreqs :. numPhiFreqs) =
        extent arr'
      arr =
        computeUnboxedS $
        R.backpermute
          (Z :. numRFreqs :. numRhoFreqs :. numThetaFreqs :. numPhiFreqs)
          (\(Z :. r :. rho :. theta :. phi) -> (Z :. r :. theta :. rho :. phi))
          arr'
  arrF <-
    fmap (fromUnboxed (extent arr) . VG.convert) .
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreqs, numRhoFreqs, numThetaFreqs, numPhiFreqs]
         [0, 1]) .
    VG.convert . toUnboxed $
    arr
  let convolvedArrF =
        computeUnboxedS .
        R.traverse2
          arrF
          (fromUnboxed (Z :. numRFreqs :. numRhoFreqs) filterF)
          const $ \fArr fFilter idx@(Z :. r :. rho :. theta :. phi) ->
          fArr idx * fFilter (Z :. r :. rho)
  fmap
    (computeS .
     R.backpermute
       (Z :. numRFreqs :. numThetaFreqs :. numRhoFreqs :. numPhiFreqs)
       (\(Z :. r :. theta :. rho :. phi) -> (Z :. r :. rho :. theta :. phi)) .
     fromUnboxed (extent arr) . VG.convert) .
    dftExecute
      plan
      (DFTPlanID
         IDFT1DG
         [numRFreqs, numRhoFreqs, numThetaFreqs, numPhiFreqs]
         [0, 1]) .
    VG.convert . toUnboxed $
    convolvedArrF

hollowCoefficients ::
     DFTPlan
  -> Double
  -> Double
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
hollowCoefficients plan radius period coefficients = do
  let (Z :. numRFreqs :. _ :. numRhoFreqs :. _) = extent coefficients
      lpf = idealLowPassFilter (log radius) (log period) numRhoFreqs numRFreqs
      filter =
        VG.convert .
        toUnboxed .
        computeS . makeFilter2D . fromUnboxed (Z :. numRFreqs :. numRhoFreqs) $
        lpf
  filterF <-
    VG.convert <$>
    dftExecute plan (DFTPlanID DFT1DG [numRFreqs, numRhoFreqs] [0, 1]) filter
  -- computeS . R.zipWith (-) coefficients <$>
  applyFilterRho plan filterF coefficients


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}

module Neural(
             Example,
             initNetwork,
             train,
             predict
             ) where

  import Numeric.LinearAlgebra
  import System.Random
  import Debug.Trace
  import Data.List

  import Tool

  type Example = (Vector Double, Vector Double)

  -- Fully connected neural network described in ADT
  data Network where
    -- Start layer : data layer. consists of just one data vector. (homogeneous)
    StartLayer :: Vector Double -> Network
    -- Hidden layer : consists of three elements
    -- 1. network subset that comes before this layer
    -- 2. weight matrix of size ((curr hidden layer dimension), (prev layer output dimension+1)).
    -- 3. current forward values (homogeneous)
    HiddenLayer :: Network -> Matrix Double -> Vector Double -> Network
    -- End layer : softmax layer. consists of three elements
    -- 1. network subset that comes before this layer
    -- 2. current softmax value
    -- 3. current loss of training (not used)
    EndLayer :: Network -> Vector Double -> Double -> Network
    deriving (Show)

  {- Activations and loss -}

  softmax :: Vector Double -> Vector Double
  softmax v = normalizer `seq` ((\val -> val / normalizer) `cmap` (exp v))
    where normalizer = sumElements (exp v)

  crossEntropy :: Vector Double -> Vector Double -> Double
  crossEntropy t y = -(sumElements $ (t * (log y)))

  activation x = tanh x

  {- Derivation computations -}

  {-
   - jaNodeHidden is a matrix of derivatives
   -
   - dy1/dx1 dy1/dx2 ... dy1/dxr
   -          ...
   - dyk/dx1 dyk/dx2 ... dyk/dxr
   -
   - where yi = activation (sum (weight * xj)), dimension of x is r, dimension of y is k.
   -}
  jaNodeHidden :: Matrix Double -> Vector Double -> Matrix Double
  jaNodeHidden weights y = weights * (repeatCols (cols weights) jaActivation)
    where
      jaActivation = 1 - (y^2)

  {-
   - jaWeightHidden is a matrix of derivatives
   -
   - dy1/dw11 dy1/dw12 ... dy1/dw1r
   -          ...
   - dyk/dwk1 dyk/dwk2 ... dyk/dwkr
   -
   - where yi = activation (sum (wij * xj)), dimension of x is r, dimension of y is k.
   -}
  jaWeightHidden :: Vector Double -> Vector Double -> Matrix Double
  jaWeightHidden x y = jaActivation `outer` jaLinear
    where
      jaActivation = 1 - (y^2)
      jaLinear = x

  -- Derivatives of softmax-crossentropy loss function
  jaLoss t y = y - t

  {- Network handling -}

  -- Initialize network with random weights and zero values
  -- The first argument is dimension of hidden layers.
  --   0th element represents dimension of the last hidden layer.
  --   0th element should be equal to the number of classes
  -- The second argument is dimension of input layer.
  -- length of hiddenDims should be larger than or equal to 1
  initNetwork :: [Int] -> Int -> IO Network
  initNetwork hiddenDims inputDim = do
    inner <- initInnerNetwork hiddenDims inputDim
    return $ EndLayer
      inner
      (zeros $ head hiddenDims)
      0.0
      where
        initInnerNetwork :: [Int] -> Int -> IO Network
        initInnerNetwork hiddenDims inputDim =
          if 1 == length hiddenDims then do
            randWeight <- randmat (-0.5, 0.5) (head hiddenDims) (inputDim+1) -- col+1 for bias
            return $ HiddenLayer
              (StartLayer (homo $ ones inputDim))
              randWeight
              (homo $ zeros (head hiddenDims)) -- append 1 for bias
          else do
            randWeight <- randmat (-0.5, 0.5) (head hiddenDims) (head (tail hiddenDims) + 1) -- col+1 for bias
            prevLayer <- (initInnerNetwork (tail hiddenDims) inputDim)
            return $ HiddenLayer
              prevLayer
              randWeight
              (homo $ zeros (last hiddenDims)) -- append 1 for bias

  -- Get a vector value of the last layer of the network
  lastLayerOf :: Network -> Vector Double
  lastLayerOf network = case network of
    StartLayer d -> d
    HiddenLayer _ _ value -> value
    EndLayer _ value _ -> value

  -- Train FC neural network 'epoch' times using given examples
  train :: Int -> Double -> [Example] -> Network -> Network
  train epoch learningRate examples network = examples `seq` ((iterate (trainEpoch learningRate examples) network)!!epoch)

  -- Train FC neural network with examples
  trainEpoch :: Double -> [Example] -> Network -> Network
  trainEpoch learningRate examples network =
    foldl' (\acc (input, gt) ->
      ((backward learningRate gt) . forward . (feedData input)) acc
      ) network examples

  -- Predict softmax values of network
  predict :: Network -> Vector Double -> Vector Double
  predict network input = (lastLayerOf . forward . (feedData input)) network

  -- Prepare Network with input
  feedData :: Vector Double -> Network -> Network
  feedData input network = input `seq` case network of
    StartLayer d -> StartLayer $ homo input
    HiddenLayer prevLayer weights value -> HiddenLayer (feedData input prevLayer) weights value
    EndLayer prevLayer value loss -> EndLayer (feedData input prevLayer) value loss

  -- Forward data to update values of Network. Leaves weights of layers unchanged.
  -- Use this for prediction
  forward :: Network -> Network
  forward network = case network of
    StartLayer d -> StartLayer d
    HiddenLayer prevLayer weights _ ->
      value `seq` HiddenLayer prevResult weights value
        where prevResult = forward prevLayer
              value = (homo (activation (weights #> lastLayerOf prevResult)))
    EndLayer prevLayer _ loss ->
      lastVector `seq` EndLayer prevResult (softmax $ invhomo lastVector) loss -- remove the last element since it is meaningless
        where prevResult = forward prevLayer
              lastVector = lastLayerOf prevResult

  -- Backpropagate errors to update weights of Network. Leaves values of layers unchanged
  -- Uses stochastic gradient descent
  backward :: Double -> Vector Double -> Network -> Network
  backward learningRate gt (EndLayer prevLayer value _) =
    EndLayer (innerBackward prevLayer lastDiff) value loss
      where
        {-
         - lastDiff is a vector of partial derivatives
         -
         - [dL/dY1 ... dL/dYk]
         -
         - where last softmax layer is k-dimension.
         -}
        lastDiff = jaLoss gt value
        loss = crossEntropy gt value

        innerBackward :: Network -> Vector Double -> Network
        innerBackward network nextDiffs = case network of
          StartLayer d -> StartLayer d
          HiddenLayer prevLayer weights values -> newWeights `seq` HiddenLayer (innerBackward prevLayer currDiff) newWeights values
            where
              {-
               - currDiff is a vector of partial derivatives
               -
               - [dL/dX1 ... dL/dXr]
               -
               - where X is the output of previous layer
               -}
              currDiff = (invhomo . flatten) $ ((jaNodeHidden weights (invhomo values)) * (repeatCols (cols weights) nextDiffs)) ? [0]
              newWeights = weightDiff `seq` (weights - ((\w -> w*learningRate) `cmap` weightDiff))
                where
                  {-
                   - weightDiff is a matrix filled with partial derivatives
                   -
                   -  dL/dW11  dL/dW12 ...  dL/dW1r
                   -               ...
                   -  dL/dWk1  dL/dWk2 ...  dL/dWkr
                   -
                   - where this hidden layer have 'k' nodes and previous layer has 'r' nodes
                   -}
                  weightDiff = pjacobian `seq` (pjacobian * (repeatCols (cols pjacobian) nextDiffs))
                    where
                      {-
                       - pjacobian is a matrix filled with partial derivatives (p is for pseudo)
                       -
                       -  dY1/dW11  dY1/dW12 ...  dY1/dW1r
                       -               ...
                       -  dYk/dWk1  dYk/dWk2 ...  dYk/dWkr
                       -
                       - where this hidden layer have 'k' nodes and previous layer has 'r' nodes
                       -}
                      pjacobian = let x = lastLayerOf prevLayer in x `seq` jaWeightHidden x (invhomo values)


{-# LANGUAGE TemplateHaskell #-}
module NewTTRS.Law where

import Control.Lens
import Data.Array.Unboxed
import Data.List
import Data.Monoid
import Statistics.Distribution (complCumulative, cumulative)
import Statistics.Distribution.Normal (normalDistr)

import NewTTRS.Outcome

data Law = Law { lawRaw :: !(UArray Int Double)
               , lawMean, lawStddev :: !Double }

data LawUpdate = LawUpdate
  { playerLaw , opponentLaw :: Law
  , updateOutcome :: Outcome
  }

-- | Law assigned to unrated players
defaultLaw :: Law
defaultLaw = normalLaw 1800 450

-- | The list of discrete scores characterized by a law
omega :: [Int]
omega = [0,10..3600]

-- | The parameter used to characterize upset probability.
alpha :: Double
alpha = 0.0148540595817432

-- | Generate a normalized law from a list of probabilities
-- which correspond to the elements of 'omega'.
lawFromList :: [Double] -> Law
lawFromList xs = Law (listArray (0,360) normalized) mean (sqrt variance)
  where
  normalized = fmap (/ sum xs) xs
  mean     = sum [ fromIntegral p       * x | (p,x) <- zip omega normalized ]
  variance = sum [ fromIntegral (p * p) * x | (p,x) <- zip omega normalized ]
           - mean * mean

-- | Find the probability that the score is in the range of
-- [i-5,i+5] given a law.
lawAt :: Law -> Int -> Double
lawAt law i = lawRaw law ! (i `div` 10)
{-# INLINE lawAt #-}

-- | Probability of an upset given the difference in two ratings.
upsetProbability :: Int -> Double
upsetProbability d = recip (1 + exp (alpha * fromIntegral d))

-- | Perform a bayesian inference given a player's law, the opponent's law
-- and the outcome of the matches between the two.
lawUpdate ::
  Law {- ^ Player's law to be updated -} ->
  Law {- ^ Opponent's law -} ->
  Outcome {- ^ Player's outcome against opponent -} ->
  Law
lawUpdate lp lq outcome
    = lawFromList
      [ sum [ upsetProbability (q - p) ^ w
            * upsetProbability (p - q) ^ l
            * lawAt lq q
            | q <- omega ]
      * lawAt lp p
      | p <- omega
      ]
  where
  w = view outcomeWins outcome
  l = view outcomeLosses outcome

-- | Perform the inverse computation of 'lawUpdate'
lawUnupdate ::
  Law {- ^ Player's law to be updated -} ->
  Law {- ^ Opponent's law -} ->
  Outcome {- ^ Player's outcome against opponent -} ->
  Law
lawUnupdate lp lq outcome
    = lawFromList
      [ lawAt lp p
      / sum [ upsetProbability (q - p) ^ w
            * upsetProbability (p - q) ^ l
            * lawAt lq q
            | q <- omega ]
      | p <- omega
      ]
  where
  w = view outcomeWins outcome
  l = view outcomeLosses outcome

-- | Chance that a player with the first law will defeat a player with the second law
chanceToWin :: Law -> Law -> Double
chanceToWin lp lq =
  sum [ upsetProbability (q - p)
      * lawAt lp p
      * lawAt lq q
      | q <- omega
      , p <- omega
      ]

-- | Generate a discretized normal law given a mean and standard deviation.
normalLaw ::
   Double {- ^ mean -} ->
   Double {- ^ standard deviation -} ->
   Law
normalLaw mean stddev = lawFromList $ mkNormal 0 3600 mean stddev

-- | Generate a discrete approximation of a normal distribution.
mkNormal ::
  Int    {- ^ lower bound -} ->
  Int    {- ^ upper bound -} ->
  Double {- ^ mean -} ->
  Double {- ^ standard deviation -} ->
  [Double]
mkNormal lo hi mean stddev =
    [ c (lo + 5)  ]
 ++ [ c (p+10) - c p | p <- [lo+5, lo+15 .. hi - 15]]
 ++ [ c' (hi - 5) ]
  where
  distr = normalDistr mean stddev
  c     = cumulative      distr . fromIntegral
  c'    = complCumulative distr . fromIntegral

-- | Degrade a law given a certain number of days. When days is
-- less than 1, no degrading is done.
timeEffect :: Int -> Law -> Law
timeEffect days law
  | days <= 0 = law
  | otherwise = lawFromList
              $ sum [ lawAt law x * timeAt y | x <- omega, y <- [-3600,-3590.. -x]]
              : [ sum [ lawAt law x * timeAt (r - x) | x <- omega ]
                | r <- omega \\ [0,3600] ]
             ++ [ sum [lawAt law x * timeAt y | x <- omega, y <- [3600-x,3610-x..3600]] ]
  where
  timeArray :: UArray Int Double
  timeArray = listArray (-360,360) $ mkNormal (-3600) 3600 0 (70 * sqrt (fromIntegral days / 365))
  timeAt y = timeArray ! (y `div` 10)

-- | Compute a metric of the law used to order players
lawScore :: Law -> Double
lawScore law = lawMean law - 2 * lawStddev law

-- | Return the raw law discrete approximations
lawElems :: Law -> [Double]
lawElems = elems . lawRaw


{-# LANGUAGE TemplateHaskell #-}
module Myelin.PyNN.Types.Synapse where

import Control.Lens
import Data.Aeson.TH
import Numeric.LinearAlgebra

type ScalarWeight = Float
type Probability = Float

{--|

ProjectionType corresponds to the different connectors available
in pynn. Some of the connectors are backend specific.

AllToAll,
FixedProbability,
DistanceDependentProbability,
FixedNumberPre,
FixedNumberPost,
OneToOne,
SmallWorld,
FromList,
FromFile

Not all have them been defined below. In addition there
are backend specific connectors.

--}
data ProjectionType =
    AllToAll {
        _weight :: Matrix Float
    }
    | OneToOne {
        _weight :: Matrix Float
    }
    | FixedNumberPost {
        _n :: Int,
        _weight ::  Matrix Float
    }
    | FixedNumberPre {
        _n :: Int,
        _weight ::  Matrix Float
    }
    | FromList {
        _weights :: [(Int, Int, Float)]
    }
    | FixedProbability {
        _probability :: Probability
    }
    deriving (Eq, Show)

--deriveJSON defaultOptions ''ProjectionType
makeLenses ''ProjectionType
makePrisms ''ProjectionType


{-# LANGUAGE NoImplicitPrelude #-}
module Algebra.Additive (
    -- * Class
    C,
    zero,
    (+), (-),
    negate, subtract,

    -- * Complex functions
    sum, sum1,
    sumNestedAssociative,
    sumNestedCommutative,

    -- * Instance definition helpers
    elementAdd, elementSub, elementNeg,
    (<*>.+), (<*>.-), (<*>.-$),

    -- * Instances for atomic types
    propAssociative,
    propCommutative,
    propIdentity,
    propInverse,
  ) where

import qualified Algebra.Laws as Laws

import Data.Int  (Int,  Int8,  Int16,  Int32,  Int64,  )
import Data.Word (Word, Word8, Word16, Word32, Word64, )

import qualified NumericPrelude.Elementwise as Elem
import Control.Applicative (Applicative(pure, (<*>)), )
import Data.Tuple.HT (fst3, snd3, thd3, )
import qualified Data.List.Match as Match

import qualified Data.Complex as Complex98
import qualified Data.Ratio as Ratio98
import qualified Prelude as P
import Prelude (Integer, Float, Double, fromInteger, )
import NumericPrelude.Base

infixl 6  +, -

{- |
Additive a encapsulates the notion of a commutative group, specified
by the following laws:

@
          a + b === b + a
    (a + b) + c === a + (b + c)
       zero + a === a
   a + negate a === 0
@

Typical examples include integers, dollars, and vectors.

Minimal definition: '+', 'zero', and ('negate' or '(-)')
-}

class C a where
    {-# MINIMAL zero, (+), ((-) | negate) #-}
    -- | zero element of the vector space
    zero     :: a
    -- | add and subtract elements
    (+), (-) :: a -> a -> a
    -- | inverse with respect to '+'
    negate   :: a -> a

    {-# INLINE negate #-}
    negate a = zero - a
    {-# INLINE (-) #-}
    a - b    = a + negate b

{- |
'subtract' is @(-)@ with swapped operand order.
This is the operand order which will be needed in most cases
of partial application.
-}
subtract :: C a => a -> a -> a
subtract = flip (-)

{- |
Sum up all elements of a list.
An empty list yields zero.

This function is inappropriate for number types like Peano.
Maybe we should make 'sum' a method of Additive.
This would also make 'lengthLeft' and 'lengthRight' superfluous.
-}
sum :: (C a) => [a] -> a
sum = foldl (+) zero

{- |
Sum up all elements of a non-empty list.
This avoids including a zero which is useful for types
where no universal zero is available.
-}
sum1 :: (C a) => [a] -> a
sum1 = foldl1 (+)

{- |
Sum the operands in an order,
such that the dependencies are minimized.
Does this have a measurably effect on speed?

Requires associativity.
-}
sumNestedAssociative :: (C a) => [a] -> a
sumNestedAssociative [] = zero
sumNestedAssociative [x] = x
sumNestedAssociative xs = sumNestedAssociative (sum2 xs)

{-
Make sure that the last entries in the list
are equally often part of an addition.
Maybe this can reduce rounding errors.
The list that sum2 computes is a breadth-first-flattened binary tree.

Requires associativity and commutativity.
-}
sumNestedCommutative :: (C a) => [a] -> a
sumNestedCommutative [] = zero
sumNestedCommutative xs@(_:rs) =
   let ys = xs ++ Match.take rs (sum2 ys)
   in  last ys

_sumNestedCommutative :: (C a) => [a] -> a
_sumNestedCommutative [] = zero
_sumNestedCommutative xs@(_:rs) =
   let ys = xs ++ take (length rs) (sum2 ys)
   in  last ys

{-
[a,b,c, a+b,c+(a+b)]
[a,b,c,d, a+b,c+d,(a+b)+(c+d)]
[a,b,c,d,e, a+b,c+d,e+(a+b),(c+d)+e+(a+b)]
[a,b,c,d,e,f, a+b,c+d,e+f,(a+b)+(c+d),(e+f)+((a+b)+(c+d))]
-}

sum2 :: (C a) => [a] -> [a]
sum2 (x:y:rest) = (x+y) : sum2 rest
sum2 xs = xs

{- |
Instead of baking the add operation into the element function,
we could use higher rank types
and pass a generic @uncurry (+)@ to the run function.
We do not do so in order to stay Haskell 98
at least for parts of NumericPrelude.
-}
{-# INLINE elementAdd #-}
elementAdd ::
   (C x) =>
   (v -> x) -> Elem.T (v,v) x
elementAdd f =
   Elem.element (\(x,y) -> f x + f y)

{-# INLINE elementSub #-}
elementSub ::
   (C x) =>
   (v -> x) -> Elem.T (v,v) x
elementSub f =
   Elem.element (\(x,y) -> f x - f y)

{-# INLINE elementNeg #-}
elementNeg ::
   (C x) =>
   (v -> x) -> Elem.T v x
elementNeg f =
   Elem.element (negate . f)

-- like <*>
infixl 4 <*>.+, <*>.-, <*>.-$

{- |
> addPair :: (Additive.C a, Additive.C b) => (a,b) -> (a,b) -> (a,b)
> addPair = Elem.run2 $ Elem.with (,) <*>.+  fst <*>.+  snd
-}
{-# INLINE (<*>.+) #-}
(<*>.+) ::
   (C x) =>
   Elem.T (v,v) (x -> a) -> (v -> x) -> Elem.T (v,v) a
(<*>.+) f acc =
   f <*> elementAdd acc

{-# INLINE (<*>.-) #-}
(<*>.-) ::
   (C x) =>
   Elem.T (v,v) (x -> a) -> (v -> x) -> Elem.T (v,v) a
(<*>.-) f acc =
   f <*> elementSub acc

{-# INLINE (<*>.-$) #-}
(<*>.-$) ::
   (C x) =>
   Elem.T v (x -> a) -> (v -> x) -> Elem.T v a
(<*>.-$) f acc =
   f <*> elementNeg acc

-- * Instances for atomic types

instance C Integer where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Float   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Double  where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int     where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int8    where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int16   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int32   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Int64   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word    where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word8   where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word16  where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word32  where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

instance C Word64  where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = P.fromInteger 0
   negate = P.negate
   (+)    = (P.+)
   (-)    = (P.-)

-- * Instances for composed types

instance (C v0, C v1) => C (v0, v1) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = (,) zero zero
   (+)    = Elem.run2 $ pure (,) <*>.+  fst <*>.+  snd
   (-)    = Elem.run2 $ pure (,) <*>.-  fst <*>.-  snd
   negate = Elem.run  $ pure (,) <*>.-$ fst <*>.-$ snd

instance (C v0, C v1, C v2) => C (v0, v1, v2) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero   = (,,) zero zero zero
   (+)    = Elem.run2 $ pure (,,) <*>.+  fst3 <*>.+  snd3 <*>.+  thd3
   (-)    = Elem.run2 $ pure (,,) <*>.-  fst3 <*>.-  snd3 <*>.-  thd3
   negate = Elem.run  $ pure (,,) <*>.-$ fst3 <*>.-$ snd3 <*>.-$ thd3

instance (C v) => C [v] where
   zero   = []
   negate = map negate
   (+) (x:xs) (y:ys) = (+) x y : (+) xs ys
   (+) xs     []     = xs
   (+) []     ys     = ys
   (-) (x:xs) (y:ys) = (-) x y : (-) xs ys
   (-) xs     []     = xs
   (-) []     ys     = negate ys

instance (C v) => C (b -> v) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero       _ = zero
   (+)    f g x = (+) (f x) (g x)
   (-)    f g x = (-) (f x) (g x)
   negate f   x = negate (f x)

-- * Properties

propAssociative :: (Eq a, C a) => a -> a -> a -> Bool
propCommutative :: (Eq a, C a) => a -> a -> Bool
propIdentity    :: (Eq a, C a) => a -> Bool
propInverse     :: (Eq a, C a) => a -> Bool

propCommutative  =  Laws.commutative (+)
propAssociative  =  Laws.associative (+)
propIdentity     =  Laws.identity (+) zero
propInverse      =  Laws.inverse (+) negate zero

-- legacy

instance (P.Integral a) => C (Ratio98.Ratio a) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero                =  P.fromInteger 0
   (+)                 =  (P.+)
   (-)                 =  (P.-)
   negate              =  P.negate

instance (P.RealFloat a) => C (Complex98.Complex a) where
   {-# INLINE zero #-}
   {-# INLINE negate #-}
   {-# INLINE (+) #-}
   {-# INLINE (-) #-}
   zero                =  P.fromInteger 0
   (+)                 =  (P.+)
   (-)                 =  (P.-)
   negate              =  P.negate


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif

---------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Simple matrix operation for low-dimensional primitives.
---------------------------------------------------------------------------
module Linear.Matrix
  ( (!*!), (!+!), (!-!), (!*), (*!), (!!*), (*!!), (!!/)
  , column
  , adjoint
  , M22, M23, M24, M32, M33, M34, M42, M43, M44
  , m33_to_m44, m43_to_m44
  , det22, det33, det44, inv22, inv33, inv44
  , identity
  , Trace(..)
  , translation
  , transpose
  , fromQuaternion
  , mkTransformation
  , mkTransformationMat
  , _m22, _m23, _m24
  , _m32, _m33, _m34
  , _m42, _m43, _m44
#if MIN_VERSION_base(4,8,0)
  , lu
  , luFinite
  , forwardSub
  , forwardSubFinite
  , backwardSub
  , backwardSubFinite
  , luSolve
  , luSolveFinite
  , luInv
  , luInvFinite
  , luDet
  , luDetFinite
#endif
  ) where

#if __GLASGOW_HASKELL__ < 710
import Control.Applicative
#endif
import Control.Lens hiding (index)
import Control.Lens.Internal.Context
import Data.Distributive
import Data.Foldable as Foldable
import Data.Functor.Rep
import Linear.Quaternion
import Linear.V2
import Linear.V3
import Linear.V4
import Linear.Vector
import Linear.Conjugate
import Linear.Trace

#if MIN_VERSION_base(4,8,0)
import GHC.TypeLits
import Linear.V
#endif

-- $setup
-- >>> import Control.Lens hiding (index)
-- >>> import Data.Complex (Complex (..))
-- >>> import Linear.V2
-- >>> import Linear.V3
-- >>> import Linear.V
-- >>> import Data.IntMap
-- >>> import Debug.SimpleReflect.Vars

-- | This is a generalization of 'Control.Lens.inside' to work over any corepresentable 'Functor'.
--
-- @
-- 'column' :: 'Representable' f => 'Lens' s t a b -> 'Lens' (f s) (f t) (f a) (f b)
-- @
--
-- In practice it is used to access a column of a matrix.
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) ^._x
-- V3 1 2 3
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) ^.column _x
-- V2 1 4
column :: Representable f => LensLike (Context a b) s t a b -> Lens (f s) (f t) (f a) (f b)
column l f es = o <$> f i where
   go = l (Context id)
   i = tabulate $ \ e -> ipos $ go (index es e)
   o eb = tabulate $ \ e -> ipeek (index eb e) (go (index es e))

infixl 7 !*!
-- | Matrix product. This can compute any combination of sparse and dense multiplication.
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) !*! V3 (V2 1 2) (V2 3 4) (V2 4 5)
-- V2 (V2 19 25) (V2 43 58)
--
-- >>> V2 (fromList [(1,2)]) (fromList [(2,3)]) !*! fromList [(1,V3 0 0 1), (2, V3 0 0 5)]
-- V2 (V3 0 0 2) (V3 0 0 15)
(!*!) :: (Functor m, Foldable t, Additive t, Additive n, Num a) => m (t a) -> t (n a) -> m (n a)
f !*! g = fmap (\ f' -> Foldable.foldl' (^+^) zero $ liftI2 (*^) f' g) f

infixl 6 !+!
-- | Entry-wise matrix addition.
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) !+! V2 (V3 7 8 9) (V3 1 2 3)
-- V2 (V3 8 10 12) (V3 5 7 9)
(!+!) :: (Additive m, Additive n, Num a) => m (n a) -> m (n a) -> m (n a)
as !+! bs = liftU2 (^+^) as bs

infixl 6 !-!
-- | Entry-wise matrix subtraction.
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) !-! V2 (V3 7 8 9) (V3 1 2 3)
-- V2 (V3 (-6) (-6) (-6)) (V3 3 3 3)
(!-!) :: (Additive m, Additive n, Num a) => m (n a) -> m (n a) -> m (n a)
as !-! bs = liftU2 (^-^) as bs

infixl 7 !*
-- | Matrix * column vector
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) !* V3 7 8 9
-- V2 50 122
(!*) :: (Functor m, Foldable r, Additive r, Num a) => m (r a) -> r a -> m a
m !* v = fmap (\r -> Foldable.sum $ liftI2 (*) r v) m

infixl 7 *!
-- | Row vector * matrix
--
-- >>> V2 1 2 *! V2 (V3 3 4 5) (V3 6 7 8)
-- V3 15 18 21

-- (*!) :: (Metric r, Additive n, Num a) => r a -> r (n a) -> n a
-- f *! g = dot f <$> distribute g

(*!) :: (Num a, Foldable t, Additive f, Additive t) => t a -> t (f a) -> f a
f *! g = sumV $ liftI2 (*^) f g

infixl 7 *!!
-- | Scalar-matrix product
--
-- >>> 5 *!! V2 (V2 1 2) (V2 3 4)
-- V2 (V2 5 10) (V2 15 20)
(*!!) :: (Functor m, Functor r, Num a) => a -> m (r a) -> m (r a)
s *!! m = fmap (s *^) m
{-# INLINE (*!!) #-}

infixl 7 !!*
-- | Matrix-scalar product
--
-- >>> V2 (V2 1 2) (V2 3 4) !!* 5
-- V2 (V2 5 10) (V2 15 20)
(!!*) :: (Functor m, Functor r, Num a) => m (r a) -> a -> m (r a)
(!!*) = flip (*!!)
{-# INLINE (!!*) #-}

infixl 7 !!/
-- | Matrix-scalar division
(!!/) :: (Functor m, Functor r, Fractional a) => m (r a) -> a -> m (r a)
m !!/ s = fmap (^/ s) m
{-# INLINE (!!/) #-}

-- | Hermitian conjugate or conjugate transpose
--
-- >>> adjoint (V2 (V2 (1 :+ 2) (3 :+ 4)) (V2 (5 :+ 6) (7 :+ 8)))
-- V2 (V2 (1.0 :+ (-2.0)) (5.0 :+ (-6.0))) (V2 (3.0 :+ (-4.0)) (7.0 :+ (-8.0)))
adjoint :: (Functor m, Distributive n, Conjugate a) => m (n a) -> n (m a)
adjoint = collect (fmap conjugate)
{-# INLINE adjoint #-}

-- * Matrices
--
-- Matrices use a row-major representation.

-- | A 2x2 matrix with row-major representation
type M22 a = V2 (V2 a)
-- | A 2x3 matrix with row-major representation
type M23 a = V2 (V3 a)
-- | A 2x4 matrix with row-major representation
type M24 a = V2 (V4 a)
-- | A 3x2 matrix with row-major representation
type M32 a = V3 (V2 a)
-- | A 3x3 matrix with row-major representation
type M33 a = V3 (V3 a)
-- | A 3x4 matrix with row-major representation
type M34 a = V3 (V4 a)
-- | A 4x2 matrix with row-major representation
type M42 a = V4 (V2 a)
-- | A 4x3 matrix with row-major representation
type M43 a = V4 (V3 a)
-- | A 4x4 matrix with row-major representation
type M44 a = V4 (V4 a)

-- | Build a rotation matrix from a unit 'Quaternion'.
fromQuaternion :: Num a => Quaternion a -> M33 a
fromQuaternion (Quaternion w (V3 x y z)) =
  V3 (V3 (1-2*(y2+z2)) (2*(xy-zw)) (2*(xz+yw)))
     (V3 (2*(xy+zw)) (1-2*(x2+z2)) (2*(yz-xw)))
     (V3 (2*(xz-yw)) (2*(yz+xw)) (1-2*(x2+y2)))
  where x2 = x*x
        y2 = y*y
        z2 = z*z
        xy = x*y
        xz = x*z
        xw = x*w
        yz = y*z
        yw = y*w
        zw = z*w
{-# INLINE fromQuaternion #-}

-- | Build a transformation matrix from a rotation matrix and a
-- translation vector.
mkTransformationMat :: Num a => M33 a -> V3 a -> M44 a
mkTransformationMat (V3 r1 r2 r3) (V3 tx ty tz) =
  V4 (snoc3 r1 tx) (snoc3 r2 ty) (snoc3 r3 tz) (V4 0 0 0 1)
  where snoc3 (V3 x y z) = V4 x y z
{-# INLINE mkTransformationMat #-}

-- |Build a transformation matrix from a rotation expressed as a
-- 'Quaternion' and a translation vector.
mkTransformation :: Num a => Quaternion a -> V3 a -> M44 a
mkTransformation = mkTransformationMat . fromQuaternion
{-# INLINE mkTransformation #-}

-- | Convert from a 4x3 matrix to a 4x4 matrix, extending it with the @[ 0 0 0 1 ]@ column vector
m43_to_m44 :: Num a => M43 a -> M44 a
m43_to_m44
  (V4 (V3 a b c)
      (V3 d e f)
      (V3 g h i)
      (V3 j k l)) =
  V4 (V4 a b c 0)
     (V4 d e f 0)
     (V4 g h i 0)
     (V4 j k l 1)

-- | Convert a 3x3 matrix to a 4x4 matrix extending it with 0's in the new row and column.
m33_to_m44 :: Num a => M33 a -> M44 a
m33_to_m44 (V3 r1 r2 r3) = V4 (vector r1) (vector r2) (vector r3) (point 0)

-- |The identity matrix for any dimension vector.
--
-- >>> identity :: M44 Int
-- V4 (V4 1 0 0 0) (V4 0 1 0 0) (V4 0 0 1 0) (V4 0 0 0 1)
-- >>> identity :: V3 (V3 Int)
-- V3 (V3 1 0 0) (V3 0 1 0) (V3 0 0 1)
identity :: (Num a, Traversable t, Applicative t) => t (t a)
identity = scaled (pure 1)

-- |Extract the translation vector (first three entries of the last
-- column) from a 3x4 or 4x4 matrix.
translation :: (Representable t, R3 t, R4 v) => Lens' (t (v a)) (V3 a)
translation = column _w._xyz
{-
translation f rs = aux <$> f (view _w <$> view _xyz rs)
 where aux (V3 x y z) = (_x._w .~ x) . (_y._w .~ y) . (_z._w .~ z) $ rs

-- translation :: (R3 t, R4 v, Functor f, Functor t) => (V3 a -> f (V3 a)) -> t (v a) -> f (t a)
-- translation = (. fmap (^._w)) . _xyz where
--   x ^. l = getConst (l Const x)
-}

-- |Extract a 2x2 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m22 :: (Representable t, R2 t, R2 v) => Lens' (t (v a)) (M22 a)
_m22 = column _xy._xy

-- |Extract a 2x3 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m23 :: (Representable t, R2 t, R3 v) => Lens' (t (v a)) (M23 a)
_m23 = column _xyz._xy

-- |Extract a 2x4 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m24 :: (Representable t, R2 t, R4 v) => Lens' (t (v a)) (M24 a)
_m24 = column _xyzw._xy

-- |Extract a 3x2 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m32 :: (Representable t, R3 t, R2 v) => Lens' (t (v a)) (M32 a)
_m32 = column _xy._xyz

-- |Extract a 3x3 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m33 :: (Representable t, R3 t, R3 v) => Lens' (t (v a)) (M33 a)
_m33 = column _xyz._xyz

-- |Extract a 3x4 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m34 :: (Representable t, R3 t, R4 v) => Lens' (t (v a)) (M34 a)
_m34 = column _xyzw._xyz

-- |Extract a 4x2 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m42 :: (Representable t, R4 t, R2 v) => Lens' (t (v a)) (M42 a)
_m42 = column _xy._xyzw

-- |Extract a 4x3 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m43 :: (Representable t, R4 t, R3 v) => Lens' (t (v a)) (M43 a)
_m43 = column _xyz._xyzw

-- |Extract a 4x4 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m44 :: (Representable t, R4 t, R4 v) => Lens' (t (v a)) (M44 a)
_m44 = column _xyzw._xyzw

-- |2x2 matrix determinant.
--
-- >>> det22 (V2 (V2 a b) (V2 c d))
-- a * d - b * c
det22 :: Num a => M22 a -> a
det22 (V2 (V2 a b) (V2 c d)) = a * d - b * c
{-# INLINE det22 #-}

-- |3x3 matrix determinant.
--
-- >>> det33 (V3 (V3 a b c) (V3 d e f) (V3 g h i))
-- a * (e * i - f * h) - d * (b * i - c * h) + g * (b * f - c * e)
det33 :: Num a => M33 a -> a
det33 (V3 (V3 a b c)
          (V3 d e f)
          (V3 g h i)) = a * (e*i-f*h) - d * (b*i-c*h) + g * (b*f-c*e)
{-# INLINE det33 #-}

-- |4x4 matrix determinant.
det44 :: Num a => M44 a -> a
det44 (V4 (V4 i00 i01 i02 i03)
          (V4 i10 i11 i12 i13)
          (V4 i20 i21 i22 i23)
          (V4 i30 i31 i32 i33)) =
  let
    s0 = i00 * i11 - i10 * i01
    s1 = i00 * i12 - i10 * i02
    s2 = i00 * i13 - i10 * i03
    s3 = i01 * i12 - i11 * i02
    s4 = i01 * i13 - i11 * i03
    s5 = i02 * i13 - i12 * i03

    c5 = i22 * i33 - i32 * i23
    c4 = i21 * i33 - i31 * i23
    c3 = i21 * i32 - i31 * i22
    c2 = i20 * i33 - i30 * i23
    c1 = i20 * i32 - i30 * i22
    c0 = i20 * i31 - i30 * i21
  in s0 * c5 - s1 * c4 + s2 * c3 + s3 * c2 - s4 * c1 + s5 * c0
{-# INLINE det44 #-}

-- |2x2 matrix inverse.
--
-- >>> inv22 $ V2 (V2 1 2) (V2 3 4)
-- V2 (V2 (-2.0) 1.0) (V2 1.5 (-0.5))
inv22 :: Fractional a => M22 a -> M22 a
inv22 m@(V2 (V2 a b) (V2 c d)) = (1 / det) *!! V2 (V2 d (-b)) (V2 (-c) a)
  where det = det22 m
{-# INLINE inv22 #-}

-- |3x3 matrix inverse.
--
-- >>> inv33 $ V3 (V3 1 2 4) (V3 4 2 2) (V3 1 1 1)
-- V3 (V3 0.0 0.5 (-1.0)) (V3 (-0.5) (-0.75) 3.5) (V3 0.5 0.25 (-1.5))
inv33 :: Fractional a => M33 a -> M33 a
inv33 m@(V3 (V3 a b c)
            (V3 d e f)
            (V3 g h i))
  = (1 / det) *!! V3 (V3 a' b' c')
                     (V3 d' e' f')
                     (V3 g' h' i')
  where a' = cofactor (e,f,h,i)
        b' = cofactor (c,b,i,h)
        c' = cofactor (b,c,e,f)
        d' = cofactor (f,d,i,g)
        e' = cofactor (a,c,g,i)
        f' = cofactor (c,a,f,d)
        g' = cofactor (d,e,g,h)
        h' = cofactor (b,a,h,g)
        i' = cofactor (a,b,d,e)
        cofactor (q,r,s,t) = det22 (V2 (V2 q r) (V2 s t))
        det = det33 m
{-# INLINE inv33 #-}

-- | 'transpose' is just an alias for 'distribute'
--
-- > transpose (V3 (V2 1 2) (V2 3 4) (V2 5 6))
-- V2 (V3 1 3 5) (V3 2 4 6)
transpose :: (Distributive g, Functor f) => f (g a) -> g (f a)
transpose = distribute
{-# INLINE transpose #-}

-- |4x4 matrix inverse.
inv44 :: Fractional a => M44 a -> M44 a
inv44 (V4 (V4 i00 i01 i02 i03)
          (V4 i10 i11 i12 i13)
          (V4 i20 i21 i22 i23)
          (V4 i30 i31 i32 i33)) =
  let s0 = i00 * i11 - i10 * i01
      s1 = i00 * i12 - i10 * i02
      s2 = i00 * i13 - i10 * i03
      s3 = i01 * i12 - i11 * i02
      s4 = i01 * i13 - i11 * i03
      s5 = i02 * i13 - i12 * i03
      c5 = i22 * i33 - i32 * i23
      c4 = i21 * i33 - i31 * i23
      c3 = i21 * i32 - i31 * i22
      c2 = i20 * i33 - i30 * i23
      c1 = i20 * i32 - i30 * i22
      c0 = i20 * i31 - i30 * i21
      det = s0 * c5 - s1 * c4 + s2 * c3 + s3 * c2 - s4 * c1 + s5 * c0
      invDet = recip det
  in invDet *!! V4 (V4 (i11 * c5 - i12 * c4 + i13 * c3)
                       (-i01 * c5 + i02 * c4 - i03 * c3)
                       (i31 * s5 - i32 * s4 + i33 * s3)
                       (-i21 * s5 + i22 * s4 - i23 * s3))
                   (V4 (-i10 * c5 + i12 * c2 - i13 * c1)
                       (i00 * c5 - i02 * c2 + i03 * c1)
                       (-i30 * s5 + i32 * s2 - i33 * s1)
                       (i20 * s5 - i22 * s2 + i23 * s1))
                   (V4 (i10 * c4 - i11 * c2 + i13 * c0)
                       (-i00 * c4 + i01 * c2 - i03 * c0)
                       (i30 * s4 - i31 * s2 + i33 * s0)
                       (-i20 * s4 + i21 * s2 - i23 * s0))
                   (V4 (-i10 * c3 + i11 * c1 - i12 * c0)
                       (i00 * c3 - i01 * c1 + i02 * c0)
                       (-i30 * s3 + i31 * s1 - i32 * s0)
                       (i20 * s3 - i21 * s1 + i22 * s0))
{-# INLINE inv44 #-}

#if MIN_VERSION_base(4,8,0)
-- | Compute the (L, U) decomposition of a square matrix using Crout's
--   algorithm. The 'Index' of the vectors must be 'Integral'.
lu :: ( Num a
      , Fractional a
      , Foldable m
      , Traversable m
      , Applicative m
      , Additive m
      , Ixed (m a)
      , Ixed (m (m a))
      , i ~ Index (m a)
      , i ~ Index (m (m a))
      , Eq i
      , Integral i
      , a ~ IxValue (m a)
      , m a ~ IxValue (m (m a))
      , Num (m a)
      )
   => m (m a)
   -> (m (m a), m (m a))
lu a =
    let n = fromIntegral (length a)
        initU = identity
        initL = zero
        buildLVal !i !j !l !u =
            let go !k !s
                    | k == j = s
                    | otherwise = go (k+1)
                                     ( s
                                      + ( (l ^?! ix i ^?! ix k)
                                        * (u ^?! ix k ^?! ix j)
                                        )
                                      )
                s' = go 0 0
            in l & (ix i . ix j) .~ ((a ^?! ix i ^?! ix j) - s')
        buildL !i !j !l !u
            | i == n = l
            | otherwise = buildL (i+1) j (buildLVal i j l u) u
        buildUVal !i !j !l !u =
            let go !k !s
                    | k == j = s
                    | otherwise = go (k+1)
                                     ( s
                                     + ( (l ^?! ix j ^?! ix k)
                                       * (u ^?! ix k ^?! ix i)
                                       )
                                     )
                s' = go 0 0
            in u & (ix j . ix i) .~ ( ((a ^?! ix j ^?! ix i) - s')
                                    / (l ^?! ix j ^?! ix j)
                                    )
        buildU !i !j !l !u
            | i == n = u
            | otherwise = buildU (i+1) j l (buildUVal i j l u)
        buildLU !j !l !u
            | j == n = (l, u)
            | otherwise =
                let l' = buildL j j l u
                    u' = buildU j j l' u
                in buildLU (j+1) l' u'
    in buildLU 0 initL initU

-- | Compute the (L, U) decomposition of a square matrix using Crout's
--   algorithm, using the vector's 'Finite' instance to provide an index.
luFinite :: ( Num a
            , Fractional a
            , Functor m
            , Finite m
            , n ~ Size m
            , KnownNat n
            , Num (m a)
            )
         => m (m a)
         -> (m (m a), m (m a))
luFinite a =
    bimap (fmap fromV . fromV)
          (fmap fromV . fromV)
          (lu (fmap toV (toV a)))

-- | Solve a linear system with a lower-triangular matrix of coefficients with
--   forwards substitution.
forwardSub :: ( Num a
              , Fractional a
              , Foldable m
              , Additive m
              , Ixed (m a)
              , Ixed (m (m a))
              , i ~ Index (m a)
              , i ~ Index (m (m a))
              , Eq i
              , Ord i
              , Integral i
              , a ~ IxValue (m a)
              , m a ~ IxValue (m (m a))
              )
           => m (m a)
           -> m a
           -> m a
forwardSub a b =
    let n = fromIntegral (length b)
        initX = zero
        coeff !i !j !s !x
            | j == i = s
            | otherwise = coeff i (j+1) (s + ((a ^?! ix i ^?! ix j) * (x ^?! ix j))) x
        go !i !x
            | i == n = x
            | otherwise = go (i + 1) (x & ix i .~ ( ((b ^?! ix i) - coeff i 0 0 x)
                                                  / (a ^?! ix i ^?! ix i)
                                                  ))
    in go 0 initX

-- | Solve a linear system with a lower-triangular matrix of coefficients with
--   forwards substitution, using the vector's 'Finite' instance to provide an
--   index.
forwardSubFinite :: ( Num a
                    , Fractional a
                    , Foldable m
                    , n ~ Size m
                    , KnownNat n
                    , Additive m
                    , Finite m
                    )
                 => m (m a)
                 -> m a
                 -> m a
forwardSubFinite a b = fromV (forwardSub (fmap toV (toV a)) (toV b))

-- | Solve a linear system with an upper-triangular matrix of coefficients with
--   backwards substitution.
backwardSub :: ( Num a
               , Fractional a
               , Foldable m
               , Additive m
               , Ixed (m a)
               , Ixed (m (m a))
               , i ~ Index (m a)
               , i ~ Index (m (m a))
               , Eq i
               , Ord i
               , Integral i
               , a ~ IxValue (m a)
               , m a ~ IxValue (m (m a))
               )
            => m (m a)
            -> m a
            -> m a
backwardSub a b =
    let n = fromIntegral (length b)
        initX = zero
        coeff !i !j !s !x
            | j == n = s
            | otherwise = coeff i
                                (j+1)
                                (s + ((a ^?! ix i ^?! ix j) * (x ^?! ix j)))
                                x
        go !i !x
            | i < 0 = x
            | otherwise = go (i-1)
                             (x & ix i .~ ( ((b ^?! ix i) - coeff i (i+1) 0 x)
                                          / (a ^?! ix i ^?! ix i)
                                          ))
    in go (n-1) initX

-- | Solve a linear system with an upper-triangular matrix of coefficients with
--   backwards substitution, using the vector's 'Finite' instance to provide an
--   index.
backwardSubFinite :: ( Num a
                     , Fractional a
                     , Foldable m
                     , n ~ Size m
                     , KnownNat n
                     , Additive m
                     , Finite m
                     )
                  => m (m a)
                  -> m a
                  -> m a
backwardSubFinite a b = fromV (backwardSub (fmap toV (toV a)) (toV b))

-- | Solve a linear system with LU decomposition.
luSolve :: ( Num a
           , Fractional a
           , Foldable m
           , Traversable m
           , Applicative m
           , Additive m
           , Ixed (m a)
           , Ixed (m (m a))
           , i ~ Index (m a)
           , i ~ Index (m (m a))
           , Eq i
           , Integral i
           , a ~ IxValue (m a)
           , m a ~ IxValue (m (m a))
           , Num (m a)
           )
        => m (m a)
        -> m a
        -> m a
luSolve a b =
    let (l, u) = lu a
    in backwardSub u (forwardSub l b)

-- | Solve a linear system with LU decomposition, using the vector's 'Finite'
--   instance to provide an index.
luSolveFinite :: ( Num a
                 , Fractional a
                 , Functor m
                 , Finite m
                 , n ~ Size m
                 , KnownNat n
                 , Num (m a)
                 )
              => m (m a)
              -> m a
              -> m a
luSolveFinite a b = fromV (luSolve (fmap toV (toV a)) (toV b))

-- | Invert a matrix with LU decomposition.
luInv :: ( Num a
         , Fractional a
         , Foldable m
         , Traversable m
         , Applicative m
         , Additive m
         , Distributive m
         , Ixed (m a)
         , Ixed (m (m a))
         , i ~ Index (m a)
         , i ~ Index (m (m a))
         , Eq i
         , Integral i
         , a ~ IxValue (m a)
         , m a ~ IxValue (m (m a))
         , Num (m a)
         )
      => m (m a)
      -> m (m a)
luInv a =
    let n = fromIntegral (length a)
        initA' = zero
        (l, u) = lu a
        go !i !a'
            | i == n = a'
            | otherwise = let e   = zero & ix i .~ 1
                              a'r = backwardSub u (forwardSub l e)
                          in go (i+1) (a' & ix i .~ a'r)
    in transpose (go 0 initA')

-- | Invert a matrix with LU decomposition, using the vector's 'Finite' instance
--   to provide an index.
luInvFinite :: ( Num a
               , Fractional a
               , Functor m
               , Finite m
               , n ~ Size m
               , KnownNat n
               , Num (m a)
               )
            => m (m a)
            -> m (m a)
luInvFinite a = fmap fromV (fromV (luInv (fmap toV (toV a))))

-- | Compute the determinant of a matrix using LU decomposition.
luDet :: ( Num a
         , Fractional a
         , Foldable m
         , Traversable m
         , Applicative m
         , Additive m
         , Trace m
         , Ixed (m a)
         , Ixed (m (m a))
         , i ~ Index (m a)
         , i ~ Index (m (m a))
         , Eq i
         , Integral i
         , a ~ IxValue (m a)
         , m a ~ IxValue (m (m a))
         , Num (m a)
         )
      => m (m a)
      -> a
luDet a =
    let (l, u) = lu a
        p      = Foldable.foldl (*) 1
    in (p (diagonal l)) * (p (diagonal u))

-- | Compute the determinant of a matrix using LU decomposition, using the
--   vector's 'Finite' instance to provide an index.
luDetFinite :: ( Num a
               , Fractional a
               , Functor m
               , Finite m
               , n ~ Size m
               , KnownNat n
               , Num (m a)
               )
            => m (m a)
            -> a
luDetFinite = luDet . fmap toV . toV
#endif


{-# LANGUAGE FlexibleContexts, TypeFamilies #-}

module Neural.LayerSpawners where

import Control.Monad
import Control.Parallel.Strategies
import Numeric.LinearAlgebra

import Neural.FF
import Neural.Util
import Numeric.Extra

maxPooling :: (Int, Int) -> [Matrix Double] -> [Matrix Double]
maxPooling (m_pool, n_pool) = map (fromLists . (map (map maxElement)) . (toBlocksEvery m_pool n_pool))

-- propagates the gradient through the pooling.
-- the maximum element in each $1 block of $2 is set to its corresponding gradient, and all other elements are set to zero.
maxPoolingBP :: (Int, Int) -> [Matrix Double] -> [Matrix Double] -> [Matrix Double]
maxPoolingBP (m_pool, n_pool) = zipWith (\nb_k x_k -> fromBlocks $ zipWith (zipWith demax)
                                                     (toLists nb_k) (toBlocksEvery m_pool n_pool x_k))
   where demax nb__ x__ = accum (konst 0.0 (m_pool,n_pool) :: Matrix Double) (+) [(maxIndex x__, nb__)]

-- the default layer opts
layerOpts :: FFLayerOptions
layerOpts = FFLayerOptions {
  layerDepth = 1,
  layerActivationFn = map (cmap sigma),
  layerActivationFnDeriv = map (cmap sigmaPrime),
  pooling = (\_ -> id), -- default pooling and poolingBP don't affect input ∨ gradient
  poolingBP = (\_ nb _ -> nb),
  poolSize = (0,0)
  }

-- A fully connected layer is just one big slice that connects to all slices in the prev. layer -- there's only
-- one matrix in weights & biases each. For simplicity I'm using head to pull out these matrices.
-- all the matrices in the layer below are concatenated to fully connect the net .
-- For layer j:
-- z_j(y_i) = (w_j . y_i) + b_j
-- y_j(z_j) = sigma(z_j(y_i))
-- ∇b_j = δ_j = ((tr w_k) . δ_k) * dy_j/dz_j
-- ∇w_j = δ_j . (tr y_i)

fullyConnectedLayer :: (Int, Int) -> (Int, Int) -> FFLayerOptions -> FFLayerOptions -> IO (FFLayer)
fullyConnectedLayer (m_i,n_i) (m_j,n_j) layerOpts_i layerOpts_j = do
  let depth_i = layerDepth layerOpts_i
      yFnDeriv_i = layerActivationFnDeriv layerOpts_i
      yFn_j = layerActivationFn layerOpts_j

  w <- randn (m_j * n_j) (m_i * n_i * depth_i)
  b <- randn (m_j * n_j) 1

  return (FFLayer {
             weights = [w],
             biases = [b],

             zFn = (\x_j (w_j, b_j) ->
                     [reshape n_j $ flatten
                      $ (head w_j <> (asColumn $ fullyFlatten x_j)) + (head b_j)]),
             yFn = yFn_j,

             nablaW = (\nabla_b_j x_j -> [(flatten $ head nabla_b_j) <.||> (fullyFlatten x_j)]),
             nablaB = (\a _ -> a),
             backprop = (\nabla_b_j z_i w_j ->
                          head $ toBlocksEvery m_i n_i
                          $ (reshape ((cols $ head z_i) * depth_i) $ tr (head w_j) #||> (flatten $ head nabla_b_j))
                          * (head $ yFnDeriv_i [fromBlocks [z_i]]))
             })

-- TODO: build pooling into the conv layer, then delete pooling layer. Don't be an ass

-- None of k conv slices in the same layer interact -- leaving out k subscript for convenience
-- let (d,d) be the size of the kernel window
-- z_i[x,y] = w_i.(x_i[x..x+d,y..y+d]) + b
-- y_i = sigma(z_i)
--
--
-- NOTES:
-- bias, weights are the same for each neuron in a slice
-- weights and bias form a kernel which is applied across input, the output of the kernel forming
--     a new matrix of size (input - d, input - d)

convLayer :: (Int, Int) -> (Int, Int) -> FFLayerOptions -> FFLayerOptions -> IO (FFLayer)
convLayer (m_i,n_i) (m_j,n_j) layerOpts_i layerOpts_j = do
  let depth_i = layerDepth layerOpts_i
      depth_j = layerDepth layerOpts_j
      yFnDeriv_i = layerActivationFnDeriv layerOpts_i
      yFn_j = layerActivationFn layerOpts_j
      pooling_j = pooling layerOpts_j
      poolingBP_j = poolingBP layerOpts_j
      (m_pool,n_pool) = poolSize layerOpts_j

  w <- replicateM depth_j $ randn m_j n_j
  b <- replicateM depth_j $ randn (m_i - m_j + 1) (n_i - n_j + 1)

  return (FFLayer {
             weights = w,
             biases = b,

             zFn = (\x_j (w_j, b_j) ->
                     let zf w_ b_ x_ = (corr2 w_ (fliprl . flipud $ x_)) + b_
                     in pooling_j (m_pool,n_pool)
                        $ zipWith (\w_ b_ -> (sum $ map (zf w_ b_) x_j) / (fromIntegral depth_j)) w_j b_j),
             yFn = yFn_j,

             nablaW = (\nabla_b_j x_j ->
                        let nwf nabla_b_ x_ = corr2 nabla_b_ x_
                        in if depth_i == 1
                           then parMap rdeepseq (\nabla_b_ -> nwf nabla_b_ (head x_j)) nabla_b_j
                           else parZipWith rdeepseq nwf nabla_b_j x_j),
             nablaB = poolingBP_j (m_pool,n_pool),
             backprop = (\nabla_b_j z_i w_j ->
                          zipWith3 (\nabla_b_ zp_ w_ -> (conv2 nabla_b_ w_) * zp_) nabla_b_j (yFnDeriv_i z_i) w_j)
             })


module Numeric.FFT.Types
       ( VCD, MVCD, VVCD, VMVCD, VVVCD, VI
       , Direction (..), Plan (..), BaseTransform (..)
       ) where

import           Data.Complex
import qualified Data.Vector                 as V
import           Data.Vector.Unboxed
import qualified Data.Vector.Unboxed.Mutable as MV

-- | Some useful type synonyms.
type VCD = Vector (Complex Double)
type MVCD s = MV.MVector s (Complex Double)
type VVCD = V.Vector VCD
type VVVCD = V.Vector VVCD
type VMVCD a = V.Vector (MVCD a)
type VI = Vector Int

-- | Transform direction: 'Forward' is the normal FFT, 'Inverse' is
-- inverse FFT.
data Direction = Forward | Inverse deriving (Eq, Show)

-- | A FFT plan.  This depends only on the problem size and can be
-- pre-computed and reused to transform (and inverse transform) any
-- number of vectors of the given size.
data Plan = Plan { plDLInfo  :: V.Vector (Int, Int, VVVCD, VVVCD)
                   -- ^ Size information and diagonal matrix entries
                   -- for Danielson-Lanczos recursive decomposition of
                   -- problem size.
                 , plPermute :: Maybe VI
                   -- ^ Input vector permutation to use before base
                   -- transformation and recursive Danielson-Lanczos
                   -- composition.
                 , plBase    :: BaseTransform
                   -- ^ Base transformation used for each sub-vector
                   -- before performing recursive Danielson-Lanczos
                   -- steps to form the full FFT result.
                 } deriving (Eq, Show)

-- | A "base transform" used at the "bottom" of the recursive
-- Cooley-Tukey decomposition of the input problem size: either a
-- simple DFT, a special hard-coded small problem size case, or a
-- Rader prime-length FFT invocation.
data BaseTransform = SpecialBase { baseSize :: Int }
                     -- ^ Hard-coded small-size base transform.
                   | DFTBase { baseSize :: Int
                             , dftWsFwd :: VCD
                             , dftWsInv :: VCD }
                     -- ^ Simple DFT base transform, giving problem
                     -- size and powers of roots of unity needed for
                     -- transform.
                   | RaderBase { baseSize      :: Int
                               , raderOutPerm  :: VI
                               , raderBFwd     :: VCD
                               , raderBInv     :: VCD
                               , raderConvSize :: Int
                               , raderConvPlan :: Plan }
                     -- ^ Prime-length Rader FFT base transform,
                     -- giving problem size, output index permutation
                     -- (the input index permutation is folded into
                     -- the main input permutation of the full
                     -- transform), pre-transformed Rader b sequence
                     -- for forward and inverse problems, the (padded
                     -- or not) problem size for Rader sequence
                     -- convolution and a sub-plan (either of size
                     -- baseSize-1 or the next larger power of two)
                     -- for computing the Rader convolution.
                     deriving (Eq, Show)


{-# LANGUAGE FlexibleContexts #-}

module Matrix.State
  (
    replaceMatrixElement
  , replaceVector
  , replaceElement
  , getElement
  , applyFMatrixElement
  , MatrixIndex
  , VectorIndex
  , MatrixMutation
  , VectorMutation
  , module Numeric.LinearAlgebra.HMatrix
  ) where

import Numeric.LinearAlgebra.HMatrix hiding (corr)

type MatrixIndex      = (MatrixRowIndex, VectorIndex)
type VectorIndex      = Int
type MatrixRowIndex   = Int
type MatrixMutation a = Matrix a -> Matrix a
type VectorMutation a = Vector a -> Vector a

getElement :: (Element a, Indexable (Vector a) a) => Matrix a -> MatrixIndex -> a
getElement m (i,j) = m ! i ! j

applyFMatrixElement :: (Element a, Indexable (Vector a) a)
  => (a -> a) -> MatrixIndex -> MatrixMutation a
applyFMatrixElement f i m = replaceMatrixElement i (f (getElement m i)) m

replaceMatrixElement :: (Element a) => MatrixIndex -> a -> MatrixMutation a
replaceMatrixElement (i,j) e m = replaceVector i v' m
  where
    v  = m ! i
    v' = replaceElement j e v

replaceVector :: (Element a) => MatrixRowIndex -> Vector a -> MatrixMutation a
replaceVector i v m = if not $ inBoundsMatrixRow i m
                      then m
                      else fromRows $ rowsAbove ++ v : rowsBelow
  where
    rowsAbove = toRows $ takeRows i       m
    rowsBelow = toRows $ dropRows (i + 1) m

replaceElement :: (Element a) => VectorIndex -> a -> VectorMutation a
replaceElement i e v = if not $ inBounds i v
                       then v
                       else fromList $ elemsBefore ++ e : elemsAfter
  where
    elemsBefore = take  i      $ lV
    elemsAfter  = drop (i + 1) $ lV
    lV         = toList v

inBounds :: (Element a) => VectorIndex -> Vector a -> Bool
inBounds i v = (i >= 0) && (i < (length $ toList v))

inBoundsMatrixRow :: MatrixRowIndex -> Matrix a -> Bool
inBoundsMatrixRow i m = (i >= 0) && (i < (rows m))


module MeasurementTests
  ( testMeasurements
  ) where

import Data.Maybe (fromJust)
import qualified Data.Vector as V
import HABQTlib.Data
import HABQTlib.MeasurementProcessing
import Numeric.LinearAlgebra (Complex(..), (><))
import qualified Numeric.LinearAlgebra as LA
import StateGenTests (arbPropPure)
import qualified Test.QuickCheck as QC
import Test.QuickCheck (Property, (.&&.))
import TestHelpers

povmValidProp :: PurePOVM -> Property
povmValidProp wsts = normedContents .&&. sumsToUnity
  where
    dim :: Int
    dim =
      (\(WeighedPureStateVector (_, PureStateVector sv)) -> LA.rows sv) $
      V.head wsts
    normedContents =
      QC.conjoin . V.toList $
      V.map (\(WeighedPureStateVector (_, sv)) -> arbPropPure sv) wsts
    msum :: LA.Matrix (LA.Complex Double)
    msum =
      V.foldl'
        (\acc (WeighedPureStateVector (w, PureStateVector sv)) ->
           acc + LA.scalar (w :+ 0) * (sv LA.<> LA.tr sv))
        ((dim >< dim) (repeat (0 :+ 0)))
        wsts
    sumsToUnity = LA.norm_Frob (msum - LA.ident dim) <= 1e-6

povmProbabilityNorm :: PurePOVM -> DensityMatrix -> Property
povmProbabilityNorm povm dm =
  QC.property $ abs (1 - V.sum (measurementProbs povm dm)) <= 1e-8

blochHelper :: Double -> Double -> SingleQbParam
blochHelper = (,)

coordProp :: Property
coordProp =
  let prop sv = fidelity sv sv' > 0.99999
        where
          sv' = blochAnglesToSV . fromJust . svToAngles $ sv
   in QC.forAll (QC.resize 2 QC.arbitrary) prop

testBlochCoords :: IO ()
testBlochCoords = do
  putStrLn "Testing Bloch coordinate transofrmations:"
  QC.quickCheckWith QC.stdArgs {QC.maxSuccess = 1000} coordProp

testAntipodal1QbPOVM :: IO ()
testAntipodal1QbPOVM = do
  putStrLn "Testing antipodal single qubit POVM:"
  QC.quickCheckWith QC.stdArgs {QC.maxSuccess = 1000} $
    ((povmValidProp . mkAntipodalPOVM) .) . blochHelper

genProdPOVM :: Int -> QC.Gen PurePOVM
genProdPOVM qbnum = do
  paramVs' <- QC.vectorOf qbnum QC.arbitrary
  let paramVs = fmap (uncurry blochHelper) paramVs'
      svs = map mkAntipodalPOVM paramVs
  return $ productPOVM svs

testProductPOVM :: IO ()
testProductPOVM = do
  putStrLn "Testing productPOVMs"
  let gen = do
        qbnum <- QC.suchThat QC.arbitrary (\n -> n >= 1 && n <= 6)
        let dim :: Dim
            dim = 2 ^ qbnum
        rank <- QC.suchThat QC.arbitrary (\n -> n <= dim && n > 0)
        povm <- genProdPOVM qbnum
        dm <- arbDM dim rank
        return (povm, dm)
  QC.quickCheckWith QC.stdArgs {QC.maxSuccess = 10000} $
    QC.forAll
      gen
      (\(povm, dm) -> povmValidProp povm .&&. povmProbabilityNorm povm dm)

testMeasurements :: IO ()
testMeasurements = do
  testBlochCoords
  testAntipodal1QbPOVM
  testProductPOVM


module Main where

import Numeric.LinearAlgebra
import Control.Arrow
import Common
import Forward
import BackProp
import AutoEncoder
import ActivFunc
import Other
import Mnist

main :: IO ()
main = do
  tm <- parseCsvToMatrixR "data/mnist_test_100.csv"
  let (ty, tx) = tr *** tr $ mnistRead 10 tm
  m <- parseCsvToMatrixR "data/mnist_train_600.csv"
  let (y, x) = tr *** tr $ mnistRead 10 m
  ws <- genWeights [784, 10]

  -- let pws = preTrains 0.5 1000 ramps x ws
  nws <- last . take 1000 $ iterateM (sgdMethod 100 (x, y) $ backPropClassification 0.5 ramps) ws

  -- putStrLn "not trained outputs"
  -- print . tr . classesToLabels . tr . forwardClassification rampC ws $ tx
  -- TODO: use function
  let o = tr . classesToLabels . tr . forwardClassification rampC nws $ tx
  let t = tr . classesToLabels . tr $ ty
  let l = head . toLists $ t - o
  let a = length l
  let s = length . filter (==0) $ l
  print $ fromIntegral s / fromIntegral a

  -- putStrLn "pretrainined outputs"
  -- print . classesToLabels . tr . forwardClassification rampC npws $ x