{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE InstanceSigs          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
module Grenade.Layers.FullyConnected (
    FullyConnected (..)
  , FullyConnected' (..)
  , randomFullyConnected
  ) where

import           Control.DeepSeq
import           GHC.Generics                    (Generic)
import           GHC.TypeLits

import           System.Random.MWC               hiding (create)
import           Control.Monad.Primitive         (PrimBase, PrimState)

import           Data.Proxy
import           Data.Serialize

import qualified Numeric.LinearAlgebra           as LA
import           Numeric.LinearAlgebra.Static

import           Grenade.Core

import           Grenade.Layers.Internal.Update

-- | A basic fully connected (or inner product) neural network layer.
data FullyConnected i o = FullyConnected
                        !(FullyConnected' i o)   -- Neuron weights
                        !(FullyConnected' i o)   -- Neuron momentum
                        deriving (Generic, NFData)

data FullyConnected' i o = FullyConnected'
                         !(R o)   -- Bias
                         !(L o i) -- Activations
                        deriving (Generic, NFData)

instance Show (FullyConnected i o) where
  show FullyConnected {} = "FullyConnected"

instance (KnownNat i, KnownNat o) => UpdateLayer (FullyConnected i o) where
  type Gradient (FullyConnected i o) = (FullyConnected' i o)

  runUpdate LearningParameters {..} (FullyConnected (FullyConnected' oldBias oldActivations) (FullyConnected' oldBiasMomentum oldMomentum)) (FullyConnected' biasGradient activationGradient) =
    let (newBias, newBiasMomentum)    = descendVector learningRate learningMomentum learningRegulariser oldBias biasGradient oldBiasMomentum
        (newActivations, newMomentum) = descendMatrix learningRate learningMomentum learningRegulariser oldActivations activationGradient oldMomentum
    in FullyConnected (FullyConnected' newBias newActivations) (FullyConnected' newBiasMomentum newMomentum)

instance (KnownNat i, KnownNat o) => Layer (FullyConnected i o) ('D1 i) ('D1 o) where
  type Tape (FullyConnected i o) ('D1 i) ('D1 o) = R i
  -- Do a matrix vector multiplication and return the result.
  runForwards (FullyConnected (FullyConnected' wB wN) _) (S1D v) = (v, S1D (wB + wN #> v))

  -- Run a backpropogation step for a full connected layer.
  runBackwards (FullyConnected (FullyConnected' _ wN) _) x (S1D dEdy) =
          let wB'  = dEdy
              mm'  = dEdy `outer` x
              -- calcluate derivatives for next step
              dWs  = tr wN #> dEdy
          in  (FullyConnected' wB' mm', S1D dWs)

instance (KnownNat i, KnownNat o) => Serialize (FullyConnected i o) where
  put (FullyConnected (FullyConnected' b w) _) = do
    putListOf put . LA.toList . extract $ b
    putListOf put . LA.toList . LA.flatten . extract $ w

  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy i)
      b     <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
      k     <- maybe (fail "Vector of incorrect size") return . create . LA.reshape f . LA.fromList =<< getListOf get
      let bm = konst 0
      let mm = konst 0
      return $ FullyConnected (FullyConnected' b k) (FullyConnected' bm mm)

instance (KnownNat i, KnownNat o, KnownNat (i*o)) => RandomLayer (FullyConnected i o) where

  createRandomWith = randomFullyConnected

randomFullyConnected :: forall m i o . (PrimBase m, KnownNat i, KnownNat o, KnownNat (i*o))
                     => WeightInitMethod -> Gen (PrimState m) -> m (FullyConnected i o)
randomFullyConnected m gen = do
  wN <- getRandomMatrix i o m gen
  wB <- getRandomVector i o m gen
  let bm = konst 0
      mm = konst 0
  return $ FullyConnected (FullyConnected' wB wN) (FullyConnected' bm mm)
  where i = natVal (Proxy :: Proxy i)
        o = natVal (Proxy :: Proxy o)

-------------------- GNum instances --------------------

instance (KnownNat i, KnownNat o) => GNum (FullyConnected i o) where
  s |* FullyConnected i o = FullyConnected (s |* i) o
  FullyConnected i o |+ FullyConnected i2 o2 = FullyConnected (0.5 |* (i |+ i2)) (0.5 |* (o |+ o2))
  gFromRational r = FullyConnected (gFromRational r) (gFromRational 0)

instance (KnownNat i, KnownNat o) => GNum (FullyConnected' i o) where
  s |* FullyConnected' i o = FullyConnected' (fromRational s * i) o
  FullyConnected' i o |+ FullyConnected' i2 o2 = FullyConnected' (0.5 * (i + i2)) (0.5 * (o + o2))
  gFromRational r = FullyConnected' (fromRational r) (fromRational r)


module Main (
    main,
    StatisticsSuite (..),
    runSuite
) where

import Prelude
import Statistics (runFold)
import qualified TextStatistics as TS
import Control.Applicative
import Data.Text.Lazy as T
import qualified Data.Text.Lazy.Encoding as E
import qualified Data.ByteString.Lazy as B

data StatisticsSuite = StatisticsSuite {
    wordCount :: Int,
    lineCount :: Int,
    averageWordLength :: Maybe Float,
    mostCommonLetter :: Maybe Char
} deriving (Eq, Ord, Show)

statisticsSuite :: TS.TextStatistics StatisticsSuite
statisticsSuite = StatisticsSuite <$> TS.wordCount <*> TS.lineCount <*> TS.averageWordLength <*> TS.mostCommonLetter

runSuite :: T.Text -> StatisticsSuite
runSuite = runFold foldl' statisticsSuite

main :: IO ()
main = do
    content <- fmap E.decodeUtf8 B.getContents
    print $ runSuite $ content


{-# LANGUAGE ApplicativeDo          #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DeriveGeneric          #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE ViewPatterns           #-}

import           Control.Arrow                         ((&&&))
import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.State.Strict
import           Control.Monad.Trans.Writer.Strict
import           Data.Bifunctor
import           Data.Either.Validation
import           Data.Finite
import           Data.Foldable
import           Data.IDX
import           Data.Kind
import           Data.List
import           Data.Monoid
import           Data.Profunctor
import           Data.Proxy
import           Data.Singletons
import           Data.Singletons.Prelude               (Sing(..))
import           Data.Singletons.TypeLits
import           Data.String
import           Data.Time.Clock
import           Data.Type.Combinator
import           Data.Type.Equality
import           Data.Type.Vector
import           GHC.Generics                          (Generic)
import           GHC.TypeLits                          as TL
import           GHC.TypeLits.Compare
import           Options.Applicative hiding            (ParserResult(..))
import           Statistics.Distribution
import           Statistics.Distribution.Uniform
import           System.Directory
import           System.FilePath
import           System.Random.MWC
import           System.Random.MWC.Distributions
import           TensorOps.Backend.BTensor
import           TensorOps.Learn.NeuralNet
import           TensorOps.Learn.NeuralNet.FeedForward
import           TensorOps.Types hiding                ((&&&))
import           Text.Printf
import           Type.Class.Higher.Util
import           Type.Family.Nat
import qualified Codec.Compression.GZip                as GZ
import qualified Control.Foldl                         as F
import qualified Data.ByteString.Lazy                  as BS
import qualified Data.Map.Strict                       as M
import qualified Data.Vector                           as V
import qualified Data.Vector.Unboxed                   as VU
import qualified Network.HTTP.Simple                   as HTTP
import qualified TensorOps.Tensor                      as TT
import qualified Text.PrettyPrint.Boxes                as B

mnistBase :: String
mnistBase = "http://yann.lecun.com/exdb/mnist"

mnistFiles :: Matrix '[N2, N2] FilePath
mnistFiles = ("train-images-idx3-ubyte" :+ "train-labels-idx1-ubyte" :+ ØV)
          :* ("t10k-images-idx3-ubyte"  :+ "t10k-labels-idx1-ubyte"  :+ ØV)
          :* ØV

data Opts = O { oRate        :: Double
              , oLayers      :: [Integer]
              , oBatch       :: Integer
              , oDataDir     :: FilePath
              , oNoConfusion :: Bool
              , oWhite       :: Bool
              , oInduce      :: Maybe (Finite 10)
              }
    deriving (Show, Eq, Generic)

opts :: Parser Opts
opts = O <$> option auto
               ( long "rate" <> short 'r' <> metavar "STEP"
              <> help "Neural network learning rate"
              <> value 0.02 <> showDefault
               )
         <*> option auto
               ( long "layers" <> short 'l' <> metavar "LIST"
              <> help "List of hidden layer sizes"
              <> value [300,100] <> showDefault
               )
         <*> option auto
               ( long "batch" <> short 'b' <> metavar "AMOUNT"
              <> help "Training batch size"
              <> value 1000 <> showDefault
               )
         <*> strOption
               ( long "data" <> short 'd' <> metavar "PATH"
              <> help "Directory to store/cache MNIST data files"
              <> value "data/mnist" <> showDefaultWith id
               )
         <*> switch
               ( long "noconfusion" <> short 'c'
              <> help "Disable confusion matrix validation and only display % error every batch"
               )
         <*> switch
               ( long "white" <> short 'w'
              <> help "Train with an eleventh \"white noise\" class to train network on negative results"
               )
         <*> optional (
               option readFin
                 ( long "induce" <> short 'i' <> metavar "DIGIT"
                <> help ("Every batch, attempt to induce an image "
                      ++ "of the given digit with the trained network"
                        )
                 )
               )
  where
    readFin :: forall n. KnownNat n => ReadM (Finite n)
    readFin = do
        i <- auto
        case packFinite i of
          Nothing -> readerError $
            printf "Number %d out of range (%d)" i (natVal (Proxy @n) - 1)
          Just x  -> return x

main :: IO ()
main = do
    O{..} <- execParser $ info (helper <*> opts)
        ( fullDesc
       <> header "tensor-ops-mnist - train neural nets on MNIST data set"
       <> progDesc (unlines ["Simple test of tensor-ops tensors (with hmatrix"
                            ,"backend) on MNIST classification challenge"
                            ]
                   )
        )

    mnistDat <- loadData oDataDir
    putStrLn "Loaded data."

    withSomeSing oWhite $ \(w :: Sing (w :: Bool)) ->
      withKnownNat (nOut w) $
        case Proxy %<=? Proxy :: 1 :<=? NOut w of
          LE Refl -> case Proxy %<=? Proxy :: 10 :<=? NOut w of
            LE Refl ->
              learn w (Proxy @(BTensorV HMatD)) mnistDat
                oRate oLayers oBatch oInduce oNoConfusion
            NLE _ -> error "impossible"
          NLE _ -> error "impossible"

loadData
    :: FilePath
    -> IO (Vec N2 [(Int, VU.Vector Int)])
loadData dataDir = do
    createDirectoryIfMissing True dataDir

    printf "Loading data from %s\n" dataDir
    mnistBs <- forM mnistFiles $ \u -> do
      BS.readFile (dataDir </> u) `catch` \(_ :: IOException) -> do
        printf "'%s' not found; downloading from %s ...\n" u mnistBase
        b <- GZ.decompress . HTTP.getResponseBody
               <$> HTTP.httpLBS (fromString (mnistBase </> u -<.> "gz"))
        BS.writeFile (dataDir </> u) b
        return b

    let f   :: Vec N2 FilePath
            -> Vec N2 BS.ByteString
            -> I (Validation [String] (FilePath, FilePath, IDXData, IDXLabels))
        f (I uIm :* I uLb :* ØV) (I bIm :* I bLb :* ØV) = I $ do
            im <- maybe (Failure [printf "Could not decode image %s." uIm]) Success
                    $ decodeIDX bIm
            lb <- maybe (Failure [printf "Could not decode labels %s." uLb]) Success
                    $ decodeIDXLabels bLb
            pure (uIm, uLb, im, lb)

        mnistDat' :: Either [String] (Vec N2 [(Int, VU.Vector Int)])
        mnistDat' = do
          imlb <- validationToEither . sequenceA $ vap f mnistFiles mnistBs
          forM imlb $ \(uIm, uLb, im, lb) ->
            maybe (Left [printf "Could not combine %s and %s." uIm uLb]) Right
              $ labeledIntData lb im

    evaluate . force
        =<< either (ioError . userError . unlines) return mnistDat'

processDat
    :: forall (n :: Nat) (l :: Nat) t.
     ( Fractional (ElemT t)
     , KnownNat n
     , KnownNat l
     , Tensor t
     )
    => (Int, VU.Vector Int)
    -> Either String (t '[n], (t '[l], Finite l))
processDat (l,d) = (,) <$> x <*> y
  where
    x :: Either String (t '[n])
    x = maybe (Left (printf "Bad input vector (Expected %d, got %d)" n (VU.length d))) Right
      . TT.fromList . map ((/255) . fromIntegral)
      $ VU.toList d
    y :: Either String (t '[l], Finite l)
    y = maybe (Left (printf "Label out of range (Got %d, expected [0,%d) )" l')) Right
      . fmap (TT.oneHot 1 0 &&& id)
      $ packFinite (fromIntegral l)
    n :: Integer
    n = natVal (Proxy @n)
    l' :: Integer
    l' = natVal (Proxy @l)

type family NOut (w :: Bool) = (n :: Nat) | n -> w where
    NOut 'False = 10
    NOut 'True  = 11

nOut
    :: Sing w
    -> Sing (NOut w)
nOut = \case
    SFalse -> SNat @10
    STrue  -> SNat @11

learn
    :: forall (t :: [Nat] -> Type) (w :: Bool) (o :: Nat).
     ( Tensor t
     , RealFloat (ElemT t)
     , NFData1 t
     , NFData (t '[784])
     , NFData (t '[o])
     , o ~ NOut w
     , 10 TL.<= o
     , 1  TL.<= o
     , KnownNat o
     )
    => Sing w
    -> Proxy t
    -> Vec N2 [(Int, VU.Vector Int)]
    -> Double
    -> [Integer]
    -> Integer
    -> Maybe (Finite 10)
    -> Bool
    -> IO ()
learn w _ dat rate layers (fromIntegral->batch) ind nc =
      withSystemRandom $ \g -> do
    dat' <- either (ioError . userError . unlines) return
          . validationToEither
          . (traverse . traverse) processDat'
          $ dat

    dat'' <- evaluate $ force dat'
    putStrLn "Data processed."

    let tXY, vXY :: [(t '[784], (t '[o], Finite o))]
        (tXY, vXY) = case dat'' of
                       I t :* I v :* ØV -> (t,v)

    net0 :: Network t 784 o
            <- genNet (layers `zip` repeat (actMap logistic)) actSoftmax g

    printf "rate: %f | batch: %d | layers: %s\n" rate batch (show layers)
    when (fromSing w) $
      putStrLn "white noise class enabled"
    forM_ ind $ \i ->
      printf "inducing: %d\n" (getFinite i)

    trainEpochs tXY ((map . second) snd vXY) g net0
  where
    processDat'
        :: (Int, VU.Vector Int)
        -> Validation [String] (t '[784], (t '[o], Finite o))
    processDat' = eitherToValidation . first (:[]) . processDat
    noiseClass :: t '[11]
    noiseClass = TT.oneHot 1 0 (finite 10)
    noiseFin   :: Finite 11
    noiseFin   = finite 10
    ind' :: Maybe (t '[o])
    ind' = TT.oneHot 1 0 . weakenN <$> ind
    trainEpochs
        :: [(t '[784], (t '[o], Finite o))]
        -> [(t '[784], Finite o)]
        -> GenIO
        -> Network t 784 o
        -> IO ()
    trainEpochs (V.fromList->tr) vd g = trainEpoch 1
      where
        trainEpoch
            :: Integer
            -> Network t 784 o
            -> IO ()
        trainEpoch e nt0 = do
            printf "[Epoch %d]\n" e
            tr' <- case w of
                SFalse -> return tr
                STrue  -> do
                  extr <- V.replicateM (V.length tr `div` 10) $ do
                            v <- genRand (uniformDistr 0 1) g
                            α <- realToFrac <$> genContVar (uniformDistr 0 1) g
                            return $ scaleT α v
                  return $ tr <> fmap (, (noiseClass, noiseFin)) extr

            queue <- evaluate . force =<< uniformShuffle tr' g

            printf "Training on %d samples in batches of %d ...\n" (V.length tr') batch

            nt1 <- trainBatch 1 queue nt0
            trainEpoch (succ e) nt1
          where
            trainBatch
                :: Integer
                -> V.Vector (t '[784], (t '[o], Finite o))
                -> Network t 784 o
                -> IO (Network t 784 o)
            trainBatch b (V.splitAt batch->(xs,xss)) nt
                | V.null xs = return nt
                | otherwise = do
              printf "Batch %d ...\n" b
              (nt', t) <- time . return $ trainAll nt ((fmap . second) fst xs)
              printf "Trained on %d samples in %s\n" (V.length xs) (show t)
              vd' :: [(t '[784], Finite o)] <- case w of
                  SFalse ->
                    return vd
                  STrue  -> do
                    extr <- replicateM (length vd `div` 10) $ do
                              v <- genRand (uniformDistr 0 1) g
                              α <- realToFrac <$> genContVar (uniformDistr 0 1) g
                              return $ scaleT α v
                    return $ vd <> fmap (, noiseFin) extr
              let tscore = F.fold (validate nt') ((fmap . second) snd xs)
                  vconf  = F.fold (confusion nt') vd'
                  confmat = ( B.vcat B.left
                                  [ B.text (printf "[%d]" (getFinite i))
                                  | i <- range :: [Finite o]
                                  ] B.<+>
                            )
                          . B.hsep 1 B.top
                          . map (\(_,r) -> B.vcat B.right
                                         . map (\(_, c) -> B.text (show c))
                                         $ M.toList r
                                )
                          $ M.toList vconf
                  vscore :: Double
                  vscore
                    | nc        = F.fold (validate nt') vd'
                    | otherwise = (/ fromIntegral (length vd')) . sum
                                . map (\(i, as) -> fromIntegral $ as M.! i)
                                $ M.toList vconf
              printf "Training:   %.2f%% error\n" ((1 - tscore) * 100)
              printf "Validation: %.2f%% error\n" ((1 - vscore) * 100)
              B.printBox confmat
              forM_ ind' $ \i -> do
                x0 <- genRand (uniformDistr 0 0.05) g
                let x1 = induceNum nt' i 1 5000 x0
                    y1 = intercalate "/"
                       . map (printf "%.2f" . realToFrac @_ @Double)
                       . TT.toList
                       $ runNetwork nt' x1
                putStrLn (renderOut x1)
                putStrLn y1

              trainBatch (succ b) xss nt'
        validate
            :: Network t 784 o
            -> F.Fold (t '[784], Finite o) Double
        validate n = (\s l -> fromIntegral s / fromIntegral l)
                 <$> lmap (uncurry f) F.sum
                 <*> F.length
          where
            f :: t '[784] -> Finite o -> Int
            f x r | TT.argMax (runNetwork n x) == r = 1
                  | otherwise                       = 0
        confusion
            :: Network t 784 o
            -> F.Fold (t '[784], Finite o) (M.Map (Finite o) (M.Map (Finite o) Integer))
        confusion n = F.Fold (\m (x, r) ->
                                let y = TT.argMax (runNetwork n x)
                                in  M.insertWith (M.unionWith (+)) y (M.singleton r 1) m
                             )
                             (M.fromList [ (i, M.fromList [ (j, 0) | j <- range ] )
                                         | i <- range
                                         ]
                               )
                             id
    trainAll
        :: Foldable f
        => Network t 784 o
        -> f (t '[784], t '[o])
        -> Network t 784 o
    trainAll = foldl' $ \nt (i,o) -> nt `deepseq`
        trainNetwork crossEntropy rate' i o nt
    rate' :: ElemT t
    rate' = realToFrac rate
    induceNum
        :: Network t 784 o
        -> t '[o]
        -> ElemT t
        -> Int
        -> t '[784]
        -> t '[784]
    induceNum n t r = go
      where
        go i x
          | i == 0    = x
          | otherwise = let x' = induceNetwork crossEntropy r t n x
                        in  x' `deepseq` go (i - 1) x'

time
    :: NFData a
    => IO a
    -> IO (a, NominalDiffTime)
time x = do
    t1 <- getCurrentTime
    y  <- evaluate . force =<< x
    t2 <- getCurrentTime
    return (y, t2 `diffUTCTime` t1)

renderOut
    :: forall t. (Tensor t, Real (ElemT t))
    => t '[784]
    -> String
renderOut = intercalate "\n"
          . ($ [])
          . appEndo
          . execWriter
          . execStateT (replicateM 28 go)
          . TT.toList
  where
    go :: StateT [ElemT t] (Writer (Endo [String])) ()
    go = do
        x <- state $ splitAt 28
        lift . tell . Endo . (++) . (:[]) $
            map (render . realToFrac)
          . concatMap (\y -> [y,y])
          $ x
    render :: Double -> Char
    render r | r <= 0.2  = ' '
             | r <= 0.4  = '.'
             | r <= 0.8  = '-'
             | r <= 1.9  = '='
             | otherwise = '#'

range
    :: forall n. KnownNat n
    => [Finite n]
range = finite <$> [0 .. natVal (Proxy @n) - 1]


--
-- Propagation : train supervised dataset and judge images
--

module CNN.Propagation (
  trainLayers
, update
, evaluateLayers
, judgeImage
, forwardProp
) where

import Control.Monad
import Data.List (foldl')
import Debug.Trace
import Numeric.LinearAlgebra

import CNN.Algebra
import CNN.Image
import CNN.Layer
import CNN.LayerType

{-
trainLayers

  IN : layers
       reversed layers (for backward prop)
       trainig data

  OUT: difference of layers

-}

trainLayers :: [Layer] -> [Layer] -> Trainer -> [Maybe Layer]
trainLayers [] _ (i, c) = []
trainLayers ls rls (i, c) = dls
  where
    (y, op') = judgeImage ls i
    d = [reshape (size c) (y `vsub` c)]
    (_, dls) = backwardProp (zip (tail op') rls) (d, [])

{-
update

  IN : learning rate
       differences of layers of each training data
       original layers

  OUT: updated layers

-}

update :: Double -> [Layer] -> [[Maybe Layer]] -> [Layer]
update lr [] _  = []
update lr ls [] = ls
update lr (l:ls) (dl:dls) = updateLayer lr l dl : update lr ls dls

{-
evaluateLayers

  IN : layers
       trainer data

  OUT: result (output)
       ratio of correct answer

-}

{- 過去の実装
evaluate :: [Layer] -> [Trainer] -> [(Vector R, Double)]
evaluate _ [] = []
evaluate ls (s:ss) = (op, rt) : evaluate ls ss
  where
    op = flatten $ head (head $ forwardProp ls [fst s])
    rt = (snd s) <.> op
-}

evaluateLayers :: [Layer] -> Double -> Trainer -> Double
evaluateLayers ls rr (i, c) = rr + c <.> y
  where
    (y, op) = judgeImage ls i

judgeImage :: [Layer] -> Image -> (Vector R, [Image])
judgeImage ls im = (flatten $ head $ head y, op)
  where
    (y, op) = splitAt 1 $ foldl' forwardProp' [im] ls

---

{-
fowardProp

  IN : layers

  OUT: image list
       list of output of each layer

-}

forwardProp :: [Layer] -> [Image] -> [Image]
forwardProp [] is = is
forwardProp (l:ls) (i:is) = forwardProp ls (forwardLayer l i ++ is)

forwardProp' :: [Image] -> Layer -> [Image]
forwardProp' (i:is) l = forwardLayer l i ++ is

{- |
backwardProp

  IN : list of image and layer pair
       delta from previous step
       difference of layers

  OUT: delta
       difference of layers

-}

backwardProp :: [(Image, Layer)] -> (Delta, [Maybe Layer])
             -> (Delta, [Maybe Layer])
backwardProp [] (_, ls) = ([], ls)
backwardProp ((im,l):ols) (d, ls) = backwardProp ols (d', l':ls)
  where
    (d', l') = backwardLayer l im d

{-
selectLayer :: Layer -> Bool
selectLayer (ConvLayer _ _)   = False
selectLayer (FullConnLayer _) = True
selectLayer _                 = False
-}


{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
module Q.Options.ImpliedVol.Surface
  (
    Surface(..)
  , totalVarKT
  , fwdTotalVarKT
  , volKT)

where

import qualified Data.Map                                 as M
import qualified Data.Map.Strict                          as M
import           Data.Maybe                               (fromJust)
import           Numeric.LinearAlgebra                    hiding (maxElement,
                                                           minElement)
import qualified Q.Options.Bachelier                      as Bacherlier
import qualified Q.Options.Black76                        as B76
import           Q.Options.ImpliedVol
import           Q.Options.ImpliedVol.InterpolatingSmile
import           Q.Options.ImpliedVol.StrikeInterpolation
import           Q.Options.ImpliedVol.TimeInterpolation
import           Q.Options.ImpliedVol.TimeSlice
import           Q.SortedVector
import           Q.Types

-- | Implied volatility surface where the strikes are in the space of 'k' and
-- implied volatility time slice is 'v'.
data Surface v k = Surface
  {
    surfaceSpot              :: Spot                   -- ^ Spot.
  , surfaceTenors            :: SortedVector YearFrac  -- ^ Ordered list of tenors.
  , surfaceForwardCurve      :: YearFrac -> Forward    -- ^ The forward curve.
  , surfaceDiscountCurve     :: YearFrac -> DF         -- ^ The discount curve.
  , surfaceAtmTotalVar       :: YearFrac -> TotalVar   -- ^ A spline of the at the money total variance.
  , surfaceVols              :: M.Map YearFrac v       -- ^ Map from tenor to 'TimeSlice'
  , surfaceTimeInterpolation :: TimeInterpolation      -- ^ Method of interpolation between tenors.
  , surfaceType              :: VolType                -- ^ The type of surface.
  }

totalVarKT :: (StrikeSpace k, TimeSlice v k) => Surface v k -> Strike -> YearFrac -> TotalVar
totalVarKT surface@Surface{..} k t | t <= minElement surfaceTenors =
                                       extrapolateTotalVarFrom (minElement surfaceTenors) surface k t
                                   | t >= maxElement surfaceTenors =
                                       extrapolateTotalVarFrom (maxElement surfaceTenors) surface k t
                                   | otherwise =
                                       timeInterpolate surfaceTimeInterpolation surface k t

volKT :: (StrikeSpace k, TimeSlice v k) => Surface v k -> Strike -> YearFrac ->  Vol
volKT surface k t = totalVarToVol (totalVarKT surface k t) t

fwdTotalVarKT :: ( StrikeSpace k, TimeSlice v k) => Surface v k -> Strike -> YearFrac -> Strike -> YearFrac -> TotalVar
fwdTotalVarKT surface@Surface{..} k1 t1 k2 t2 = TotalVar $ (totalVarKT2 - totalVarKT1) / (unYearFrac t2 - unYearFrac t1)
  where (TotalVar totalVarKT1) = totalVarKT surface k1 t1
        (TotalVar totalVarKT2) = totalVarKT surface k2 t2

class StrikeSpace k where
  strikeSpaceToCash :: k -> YearFrac -> Spot -> Forward -> Vol -> VolShift -> Strike
  cashToStrikeSpace :: Strike -> YearFrac -> Spot -> Forward -> Vol -> VolShift -> k

instance StrikeSpace Strike where
  strikeSpaceToCash x _ _ _ _ _ = x
  cashToStrikeSpace k _ _ _ _ _ = k

instance StrikeSpace AbsRelStrike where
  strikeSpaceToCash (AbsRel x) _ _ (Forward f) _ _ = Strike $ x + f
  cashToStrikeSpace (Strike k) _ _ (Forward f) _ _ = AbsRel $ k - f

instance StrikeSpace LogRelStrike where
  strikeSpaceToCash (LogRel x) _ _ (Forward f) _ _ = Strike $ f * exp x
  cashToStrikeSpace (Strike k) _ _ (Forward f) _ _ = LogRel $ log $ k - f

instance StrikeSpace MoneynessForwardStrike  where
  strikeSpaceToCash (MoneynessForward x) (YearFrac t) _ (Forward f) (Vol atmVol) _ =
    Strike $ x * sqrt t * atmVol + f

  cashToStrikeSpace (Strike k) (YearFrac t) _ (Forward f) (Vol atmVol) _ =
    MoneynessForward $ (k - f) / (atmVol * sqrt t)

instance StrikeSpace LogMoneynessForwardStrike  where
  strikeSpaceToCash (LogMoneynessForward x) (YearFrac t) _ (Forward f) (Vol atmVol) (VolShift slnShift) =
    Strike $ (f + slnShift) * exp (x * (sqrt t) * atmVol) - slnShift

  cashToStrikeSpace (Strike k) (YearFrac t)  _ (Forward f) (Vol atmVol) (VolShift slnShift) =
    LogMoneynessForward $ (log ((k - slnShift) / (f + slnShift))) / (atmVol * sqrt t)

instance StrikeSpace LogMoneynessSpotStrike  where
  strikeSpaceToCash (LogMoneynessSpot x) (YearFrac t) (Spot s) _ (Vol atmVol) (VolShift slnShift) =
    Strike $ (s + slnShift) * exp (x * (sqrt t) * atmVol) - slnShift

  cashToStrikeSpace (Strike k) (YearFrac t)  (Spot s) _ (Vol atmVol) (VolShift slnShift) =
    LogMoneynessSpot $ (log ((k - slnShift) / (s + slnShift))) / (atmVol * sqrt t)

instance StrikeSpace MoneynessSpotStrike  where
  strikeSpaceToCash (MoneynessSpot x) (YearFrac t) (Spot s) _ (Vol atmVol) _ =
    Strike $ x * (sqrt t) * atmVol + s

  cashToStrikeSpace (Strike k) (YearFrac t) (Spot s) _ (Vol atmVol) _ =
    MoneynessSpot $ (k - s) / (atmVol * sqrt t)

slnShift Surface{..} = case surfaceType of
  (ShiftedLogNormal shift) -> shift
  _                        -> VolShift 0
extrapolateTotalVarFrom :: forall v k. (StrikeSpace k, TimeSlice v k) => YearFrac -> Surface v k -> Strike -> YearFrac -> TotalVar
extrapolateTotalVarFrom t0 surface@Surface{..} k t = let
  f0          = surfaceForwardCurve t0
  atmVol0     = totalVarToVol (surfaceAtmTotalVar t0) t0
  f           = surfaceForwardCurve t
  spot        = surfaceSpot
  atmTotalVar = surfaceAtmTotalVar t
  atmVol      = totalVarToVol atmTotalVar t
  x           = cashToStrikeSpace k t spot f atmVol (slnShift surface)::k
  k'          = strikeSpaceToCash x t0 spot f0 atmVol0 (slnShift surface)
  x'          = cashToStrikeSpace k' t0 spot f0 atmVol (slnShift surface)::k
  in totalVar (surfaceVols M.! t0) x'

timeInterpolate :: forall v k. (StrikeSpace k, TimeSlice v k) => TimeInterpolation -> Surface v k -> Strike -> YearFrac -> TotalVar
timeInterpolate Gatheral surface@Surface{..} k11 t =
  let (t1, smile1) = fromJust $ M.lookupLE t surfaceVols
      (t2, smile2) = fromJust $ M.lookupGE t surfaceVols
      (TotalVar thetaT)  = surfaceAtmTotalVar t
      (TotalVar thetaT1) = surfaceAtmTotalVar t1
      (TotalVar thetaT2) = surfaceAtmTotalVar t2
      alphaT  = (sqrt thetaT2 - sqrt thetaT1 ) / (sqrt thetaT1 - sqrt thetaT)
      atmVol  = totalVarToVol (TotalVar thetaT) t
      (Forward f)  = surfaceForwardCurve t
      (Forward f1) = surfaceForwardCurve t1
      (Forward f2) = surfaceForwardCurve t2
      df  = surfaceDiscountCurve t
      df1 = surfaceDiscountCurve t1
      df2 = surfaceDiscountCurve t2
      k1      = k11 $*$ (f1 / f)
      k2      = k11 $*$ (f2 / f)
      x1      = cashToStrikeSpace k1 t1 surfaceSpot (Forward f1) atmVol (slnShift surface)::k
      x2      = cashToStrikeSpace k2 t2 surfaceSpot (Forward f2) atmVol (slnShift surface)::k
      vol1 = totalVarToVol (totalVar smile1 x1) t1
      vol2 = totalVarToVol (totalVar smile2  x2) t2
      (Premium premium1) = vPremium $ euOption surfaceType k1 f1 df1 t1 vol1
      (Premium premium2) = vPremium $ euOption surfaceType k1 f2 df2 t1 vol1
      premiumT  = Premium $ (alphaT * premium1 $/$ k1 + (1- alphaT)* premium2 $/$ k2) $*$ k11
      x1 :: k
      x2 :: k
  in if surfaceType == Normal then
       volToTotalVar (euImpliedVol Normal Call (Forward f) k11 t df premiumT) t
     else
       volToTotalVar (euImpliedVol LogNormal Call (Forward f) k11 t df premiumT) t

timeInterpolate LinearInVol surface@Surface{..} k t =
  let f = surfaceForwardCurve t
      atmVol = totalVarToVol (surfaceAtmTotalVar t) t
      x      = cashToStrikeSpace k t surfaceSpot f atmVol (slnShift surface)::k
      (t1, smile1) = fromJust $ M.lookupLE t surfaceVols
      (t2, smile2) = fromJust $ M.lookupGE t surfaceVols
      (Vol vol1)         = totalVarToVol (totalVar smile1 x) t1
      (Vol vol2)         = totalVarToVol (totalVar smile2 x) t2
  in volToTotalVar (Vol $ linearInterpolate (t1, vol1) (t2, vol2) t) t

timeInterpolate LinearInTotalVar surface@Surface{..} k t =
  let f = surfaceForwardCurve t
      atmVol = totalVarToVol (surfaceAtmTotalVar t) t
      x      = cashToStrikeSpace k t surfaceSpot f atmVol (slnShift surface)::k
      (t1, smile1) = fromJust $ M.lookupLE t surfaceVols
      (t2, smile2) = fromJust $ M.lookupGE t surfaceVols
      (TotalVar tv1)         = totalVar smile1 x
      (TotalVar tv2)         = totalVar smile2 x
  in TotalVar $ linearInterpolate (t1, tv1) (t2, tv2) t

euOption Normal k f (DF df) (YearFrac t) vol =
  let r = Rate $ -(log df) / t
      bacherlier = Bacherlier.Bachelier (Forward f) r vol
  in Bacherlier.eucall bacherlier (YearFrac t) k

euOption _ k f df t vol =
  let b76 = B76.Black76 (Forward f) df t vol
  in B76.eucall b76 k

linearInterpolate (YearFrac t1, v1) (YearFrac t2, v2) (YearFrac t) =
  v1 + (v2 - v1)*(t - t1) / (t2 - t1)


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE IncoherentInstances  #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Vector.Image.FFT (
  fft1d
, fft2d
, fft3d
, Direction(..)
) where

import Prelude hiding(map)
import qualified Prelude as P
import qualified Data.List as L
import Data.Vector.Storable hiding(forM_,map,replicate,fromList,toList,(!?),(!))
import Data.Maybe
import Numeric.LinearAlgebra.Data hiding (fromList,toList,(!))
import Numeric.LinearAlgebra.HMatrix hiding (fromList,toList,(!))
import qualified Data.Vector.Storable as V
import Foreign
import GHC.Real
import Foreign.C.String
import Foreign.C.Types
import qualified Control.Monad as C
import Data.Vector.Image.Color.RGB
import System.IO.Unsafe(unsafePerformIO)

foreign import ccall unsafe "fft1d" c_fft1d :: Ptr (Complex Double) -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "fft2d" c_fft2d :: Ptr (Complex Double) -> CInt -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "fft3d" c_fft3d :: Ptr (Complex Double) -> CInt -> CInt -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "&fftw_free" c_p_freeComplexDouble :: FunPtr(Ptr (Complex Double) -> IO ())

data Direction = Forward | Backword

fft1d :: Vector (Complex Double) -> Direction -> Vector (Complex Double)
fft1d vec dir = unsafePerformIO $ do
  let n = V.length vec
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft1d v (fromIntegral n) 1
      Backword -> c_fft1d v (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 n

fft2d :: Vector (Complex Double) -> Int -> Direction -> Vector (Complex Double)
fft2d vec m dir = unsafePerformIO $ do
  let nm = V.length vec
  let n = nm `div` m
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft2d v (fromIntegral m) (fromIntegral n) 1
      Backword -> c_fft2d v (fromIntegral m) (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 nm

fft3d :: Vector (Complex Double) -> Int -> Int -> Direction -> Vector (Complex Double)
fft3d vec l m dir = unsafePerformIO $ do
  let lmn = V.length vec
  let n = lmn `div` (l*m)
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft3d v (fromIntegral l) (fromIntegral m) (fromIntegral n) 1
      Backword -> c_fft3d v (fromIntegral l) (fromIntegral m) (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 lmn


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE BangPatterns, ScopedTypeVariables #-}
-- |
-- Module    : Statistics.Distribution
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Type classes for probability distributions

module Statistics.Distribution
    (
      -- * Type classes
      Distribution(..)
    , DiscreteDistr(..)
    , ContDistr(..)
      -- ** Distribution statistics
    , MaybeMean(..)
    , Mean(..)
    , MaybeVariance(..)
    , Variance(..)
    , MaybeEntropy(..)
    , Entropy(..)
    , FromSample(..)
      -- ** Random number generation
    , ContGen(..)
    , DiscreteGen(..)
    , genContinuous
    , genContinous
      -- * Helper functions
    , findRoot
    , sumProbabilities
    ) where

import Control.Applicative ((<$>), Applicative(..))
import Control.Monad.Primitive (PrimMonad,PrimState)
import Prelude hiding (sum)
import Statistics.Function (square)
import Statistics.Sample.Internal (sum)
import System.Random.MWC (Gen, uniform)
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G

-- | Type class common to all distributions. Only c.d.f. could be
-- defined for both discrete and continuous distributions.
class Distribution d where
    -- | Cumulative distribution function.  The probability that a
    -- random variable /X/ is less or equal than /x/,
    -- i.e. P(/X/≤/x/). Cumulative should be defined for
    -- infinities as well:
    --
    -- > cumulative d +∞ = 1
    -- > cumulative d -∞ = 0
    cumulative :: d -> Double -> Double

    -- | One's complement of cumulative distribution:
    --
    -- > complCumulative d x = 1 - cumulative d x
    --
    -- It's useful when one is interested in P(/X/>/x/) and
    -- expression on the right side begin to lose precision. This
    -- function have default implementation but implementors are
    -- encouraged to provide more precise implementation.
    complCumulative :: d -> Double -> Double
    complCumulative d x = 1 - cumulative d x

-- | Discrete probability distribution.
class Distribution  d => DiscreteDistr d where
    -- | Probability of n-th outcome.
    probability :: d -> Int -> Double
    probability d = exp . logProbability d

    -- | Logarithm of probability of n-th outcome
    logProbability :: d -> Int -> Double
    logProbability d = log . probability d

-- | Continuous probability distribution.
--
--   Minimal complete definition is 'quantile' and either 'density' or
--   'logDensity'.
class Distribution d => ContDistr d where
    -- | Probability density function. Probability that random
    -- variable /X/ lies in the infinitesimal interval
    -- [/x/,/x+/δ/x/) equal to /density(x)/⋅δ/x/
    density :: d -> Double -> Double
    density d = exp . logDensity d

    -- | Inverse of the cumulative distribution function. The value
    -- /x/ for which P(/X/≤/x/) = /p/. If probability is outside
    -- of [0,1] range function should call 'error'
    quantile :: d -> Double -> Double

    -- | 1-complement of @quantile@:
    --
    -- > complQuantile x ≡ quantile (1 - x)
    complQuantile :: d -> Double -> Double
    complQuantile d x = quantile d (1 - x)

    -- | Natural logarithm of density.
    logDensity :: d -> Double -> Double
    logDensity d = log . density d

-- | Type class for distributions with mean. 'maybeMean' should return
--   'Nothing' if it's undefined for current value of data
class Distribution d => MaybeMean d where
    maybeMean :: d -> Maybe Double

-- | Type class for distributions with mean. If a distribution has
--   finite mean for all valid values of parameters it should be
--   instance of this type class.
class MaybeMean d => Mean d where
    mean :: d -> Double

-- | Type class for distributions with variance. If variance is
--   undefined for some parameter values both 'maybeVariance' and
--   'maybeStdDev' should return Nothing.
--
--   Minimal complete definition is 'maybeVariance' or 'maybeStdDev'
class MaybeMean d => MaybeVariance d where
    maybeVariance :: d -> Maybe Double
    maybeVariance d = (*) <$> x <*> x where x = maybeStdDev d
    maybeStdDev   :: d -> Maybe Double
    maybeStdDev = fmap sqrt . maybeVariance

-- | Type class for distributions with variance. If distribution have
--   finite variance for all valid parameter values it should be
--   instance of this type class.
--
--   Minimal complete definition is 'variance' or 'stdDev'
class (Mean d, MaybeVariance d) => Variance d where
    variance :: d -> Double
    variance d = square (stdDev d)
    stdDev   :: d -> Double
    stdDev = sqrt . variance

-- | Type class for distributions with entropy, meaning Shannon entropy
--   in the case of a discrete distribution, or differential entropy in the
--   case of a continuous one.  'maybeEntropy' should return 'Nothing' if
--   entropy is undefined for the chosen parameter values.
class (Distribution d) => MaybeEntropy d where
  -- | Returns the entropy of a distribution, in nats, if such is defined.
  maybeEntropy :: d -> Maybe Double

-- | Type class for distributions with entropy, meaning Shannon
--   entropy in the case of a discrete distribution, or differential
--   entropy in the case of a continuous one.  If the distribution has
--   well-defined entropy for all valid parameter values then it
--   should be an instance of this type class.
class (MaybeEntropy d) => Entropy d where
  -- | Returns the entropy of a distribution, in nats.
  entropy :: d -> Double

-- | Generate discrete random variates which have given
--   distribution.
class Distribution d => ContGen d where
  genContVar :: PrimMonad m => d -> Gen (PrimState m) -> m Double

-- | Generate discrete random variates which have given
--   distribution. 'ContGen' is superclass because it's always possible
--   to generate real-valued variates from integer values
class (DiscreteDistr d, ContGen d) => DiscreteGen d where
  genDiscreteVar :: PrimMonad m => d -> Gen (PrimState m) -> m Int

-- | Estimate distribution from sample. First parameter in sample is
--   distribution type and second is element type.
class FromSample d a where
  -- | Estimate distribution from sample. Returns nothing is there's
  --   not enough data to estimate or sample clearly doesn't come from
  --   distribution in question. For example if there's negative
  --   samples in exponential distribution.
  fromSample :: G.Vector v a => v a -> Maybe d

-- | Generate variates from continuous distribution using inverse
--   transform rule.
genContinuous :: (ContDistr d, PrimMonad m) => d -> Gen (PrimState m) -> m Double
genContinuous d gen = do
  x <- uniform gen
  return $! quantile d x

-- | Backwards compatibility with genContinuous.
genContinous :: (ContDistr d, PrimMonad m) => d -> Gen (PrimState m) -> m Double
genContinous = genContinuous
{-# DEPRECATED genContinous "Use genContinuous" #-}

data P = P {-# UNPACK #-} !Double {-# UNPACK #-} !Double

-- | Approximate the value of /X/ for which P(/x/>/X/)=/p/.
--
-- This method uses a combination of Newton-Raphson iteration and
-- bisection with the given guess as a starting point.  The upper and
-- lower bounds specify the interval in which the probability
-- distribution reaches the value /p/.
findRoot :: ContDistr d =>
            d                   -- ^ Distribution
         -> Double              -- ^ Probability /p/
         -> Double              -- ^ Initial guess
         -> Double              -- ^ Lower bound on interval
         -> Double              -- ^ Upper bound on interval
         -> Double
findRoot d prob = loop 0 1
  where
    loop !(i::Int) !dx !x !lo !hi
      | abs dx <= accuracy || i >= maxIters = x
      | otherwise                           = loop (i+1) dx'' x'' lo' hi'
      where
        err                   = cumulative d x - prob
        P lo' hi' | err < 0   = P x hi
                  | otherwise = P lo x
        pdf                   = density d x
        P dx' x' | pdf /= 0   = P (err / pdf) (x - dx)
                 | otherwise  = P dx x
        P dx'' x''
            | x' < lo' || x' > hi' || pdf == 0 = let y = (lo' + hi') / 2
                                                 in  P (y-x) y
            | otherwise                        = P dx' x'
    accuracy = 1e-15
    maxIters = 150

-- | Sum probabilities in inclusive interval.
sumProbabilities :: DiscreteDistr d => d -> Int -> Int -> Double
sumProbabilities d low hi =
  -- Return value is forced to be less than 1 to guard against roundoff errors.
  -- ATTENTION! this check should be removed for testing or it could mask bugs.
  min 1 . sum . U.map (probability d) $ U.enumFromTo low hi


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

-- |
-- Module       : AI.HNN.FF.Network
-- Copyright    : (c) 2012 Alp Mestanogullari
-- License      : BSD3
-- Maintainer   : alpmestan@gmail.com
-- Stability    : experimental
-- Portability  : GHC
--
-- An implementation of feed-forward neural networks in pure Haskell.
--
-- It uses weight matrices between each layer to represent the connections between neurons from
-- a layer to the next and exports only the useful bits for a user of the library.
--
-- Here is an example of using this module to create a feed-forward neural network with 2 inputs,
-- 2 neurons in a hidden layer and one neuron in the output layer, with random weights, and compute
-- its output for [1,2] using the sigmoid function for activation for all the neurons.
--
-- > import AI.HNN.FF.Network
-- > import Numeric.LinearAlgebra
-- >
-- > main = do
-- >   n <- createNetwork 2 [2] 1 :: IO (Network Double)
-- >   print $ output n sigmoid (fromList [1, 1])
--
-- /Note/: Here, I create a @Network Double@, but you can replace 'Double' with any number type
-- that implements the appropriate typeclasses you can see in the signatures of this module.
-- Having your number type implement the @Floating@ typeclass too is a good idea, since that's what most of the
-- common activation functions require.
--
-- /Note 2/: You can also give some precise weights to initialize the neural network with, with
-- 'fromWeightMatrices'. You can also restore a neural network you had saved using 'loadNetwork'.
--
-- Here is an example of how to train a neural network to learn the XOR function.
-- ( for reference: XOR(0, 0) = 0, XOR(0, 1) = 1, XOR(1, 0) = 1, XOR(1, 1) = 0 )
--
-- First, let's import hnn's feedforward neural net module, and hmatrix's vector types.
--
-- > import AI.HNN.FF.Network
-- > import Numeric.LinearAlgebra
--
-- Now, we will specify our training set (what the net should try to learn).
--
-- > samples :: Samples Double
-- > samples = [ fromList [0, 0] --> fromList [0]
-- >           , fromList [0, 1] --> fromList [1]
-- >           , fromList [1, 0] --> fromList [1]
-- >           , fromList [1, 1] --> fromList [0]
-- >           ]
--
-- You can see that this is basically a list of pairs of vectors, the first vector being
-- the input given to the network, the second one being the expected output. Of course,
-- this imply working on a neural network with 2 inputs, and a single neuron on the output layer. Then,
-- let's create one!
--
-- > main = do
-- >   net <- createNetwork 2 [2] 1
--
-- You may have noticed we haven't specified a signature this time, unlike in the earlier snippet.
-- Since we gave a signature to samples, specifying we're working with 'Double' numbers, and since
-- we are going to tie 'net' and 'samples' by a call to a learning function, GHC will gladly figure out
-- that 'net' is working with 'Double'.
--
-- Now, it's time to train our champion. But first, let's see how bad he is now. The weights are most likely
-- not close to those that will give a good result for simulating XOR. Let's compute the output of the net on
-- the input vectors of our samples, using 'tanh' as the activation function.
--
-- >   mapM_ (print . output net tanh . fst) samples
--
-- Ok, you've tested this, and it gives terrible results. Let's fix this by letting 'trainNTimes' teach our neural net
-- how to behave. Since we're using 'tanh' as our activation function, we will tell it to the training function,
-- and also specify its derivative.
--
-- >   let smartNet = trainNTimes 1000 0.8 tanh tanh' net samples
--
-- So, this tiny piece of code will run the backpropagation algorithm on the samples 1000 times, with a learning rate
-- of 0.8. The learning rate is basically how strongly we should modify the weights when we try to correct the error the net makes
-- on our samples. The bigger it is, the more the weights are going to change significantly. Depending on the case, it can be good,
-- but sometimes it can make the backprop algorithm oscillate around good weight values without actually getting to them.
-- You usually want to test several values and see which ones get you the nicest neural net, which generalizes well to samples
-- that are not in the training set while giving decent results on the training set.
--
-- Now, let's see how that worked out for us:
--
-- >   mapM_ (print . output smartNet tanh . fst) samples
--
-- You could even save that neural network's weights to a file, so that you don't need to train it again in the future, using 'saveNetwork':
--
-- >   saveNetwork "smartNet.nn" smartNet
--
-- Please note that 'saveNetwork' is just a wrapper around zlib compression + serialization using the binary package.
-- AI.HNN.FF.Network also provides a 'Data.Binary.Binary' instance for 'Network', which means you can also simply use
-- 'Data.Binary.encode' and 'Data.Binary.decode' to have your own saving/restoring routines, or to simply get a bytestring
-- we can send over the network, for example.
--
-- Here's a run of the program we described on my machine (with the timing): first set of
-- fromList's is the output of the initial neural network, the second one is the output of
-- 'smartNet' :-)
--
-- > fromList [0.574915179613429]
-- > fromList [0.767589097192215]
-- > fromList [0.7277396607146663]
-- > fromList [0.8227114080561128]
-- > ------------------
-- > fromList [6.763498312099933e-2]
-- > fromList [0.9775186355284375]
-- > fromList [0.9350823296850516]
-- > fromList [-4.400205702560454e-2]
-- >
-- > real    0m0.365s
-- > user    0m0.072s
-- > sys     0m0.016s
--
-- Rejoyce! Feel free to play around with the library and report any bug, feature request and whatnot to us on
-- our github repository <https://github.com/alpmestan/hnn/issues> using the appropriate tags. Also, you can
-- see the simple program we studied here with pretty colors at <https://github.com/alpmestan/hnn/blob/master/examples/ff/xor.hs>
-- and other ones at <https://github.com/alpmestan/hnn/tree/master/examples/ff>.

module AI.HNN.FF.Network
    (
    -- * Types
      Network(..)
    , ActivationFunction
    , ActivationFunctionDerivative
    , Sample
    , Samples
    , (-->)

    -- * Creating a neural network
    , createNetwork
    , fromWeightMatrices

    -- * Computing a neural network's output
    , output
    , tanh
    , tanh'
    , sigmoid
    , sigmoid'

    -- * Training a neural network
    , trainUntil
    , trainNTimes
    , trainUntilErrorBelow
    , quadError

    -- * Loading and saving a neural network
    , loadNetwork
    , saveNetwork
    ) where

import Codec.Compression.Zlib     (compress, decompress)
import Data.Binary                (Binary(..), encode, decode)
import Data.List                  (foldl')
import Foreign.Storable           (Storable)
import qualified Data.ByteString.Lazy  as B
import qualified Data.Vector           as V

import System.Random.MWC
import Numeric.LinearAlgebra.HMatrix hiding (corr)
import Data.Functor ((<$>))

-- | Our feed-forward neural network type. Note the 'Binary' instance, which means you can use
--   'encode' and 'decode' in case you need to serialize your neural nets somewhere else than
--   in a file (e.g over the network)
newtype Network a = Network
                 { matrices   :: V.Vector (Matrix a) -- ^ the weight matrices
                 } deriving Show

instance (Element a, Binary a) => Binary (Network a) where
  put (Network ms) = put . V.toList $ ms
  get = (Network . V.fromList) `fmap` get

-- | The type of an activation function, mostly used for clarity in signatures
type ActivationFunction a = a -> a

-- | The type of an activation function's derivative, mostly used for clarity in signatures
type ActivationFunctionDerivative a = a -> a

-- | The following creates a neural network with 'n' inputs and if 'l' is [n1, n2, ...]
--   the net will have n1 neurons on the first layer, n2 neurons on the second, and so on
--   ending with k neurons on the output layer, with random weight matrices as a courtesy of
-- 'System.Random.MWC.uniformVector'.
--
-- > createNetwork n l k
createNetwork :: (Variate a, Storable a) => Int -> [Int] -> Int -> IO (Network a)
createNetwork nInputs hiddens nOutputs =
  fmap Network $ withSystemRandom . asGenST $ \gen -> go gen dimensions V.empty
  where
        go _ [] !ms         = return ms
        go gen ((!n,!m):ds) ms = do
          !mat <- randomMat n m gen
          go gen ds (ms `V.snoc` mat)
        randomMat n m g = reshape m `fmap` uniformVector g (n*m)
        dimensions      = zip (hiddens ++ [nOutputs]) $
                              (nInputs+1 : hiddens)
{-# INLINE createNetwork #-}

-- | Creates a neural network with exactly the weight matrices given as input here.
--   We don't check that the numbers of rows/columns are compatible, etc.
fromWeightMatrices :: Storable a => V.Vector (Matrix a) -> Network a
fromWeightMatrices ws = Network ws
{-# INLINE fromWeightMatrices #-}

-- The `join [input, 1]' trick  below is a courtesy of Alberto Ruiz
-- <http://dis.um.es/~alberto/>. Per his words:
--
-- "The idea is that the constant input in the first layer can be automatically transferred to the following layers
-- by the learning algorithm (by setting the weights of a neuron to 1,0,0,0,...). This allows for a simpler
-- implementation and in my experiments those networks are able to easily solve non linearly separable problems."

-- | Computes the output of the network on the given input vector with the given activation function
output :: (Floating (Vector a), Numeric a, Storable a, Num (Vector a)) => Network a -> ActivationFunction a -> Vector a -> Vector a
output (Network{..}) act input = V.foldl' f (vjoin [input, 1]) matrices
  where f !inp m = cmap act $ m #> inp
{-# INLINE output #-}

-- | Computes and keeps the output of all the layers of the neural network with the given activation function
outputs :: (Floating (Vector a), Numeric a, Storable a, Num (Vector a)) => Network a -> ActivationFunction a -> Vector a -> V.Vector (Vector a)
outputs (Network{..}) act input = V.scanl f (vjoin [input, 1]) matrices
  where f !inp m = cmap act $ m #> inp
{-# INLINE outputs #-}

deltas :: (Floating (Vector a), Floating a, Numeric a, Container Vector a, Num (Vector a)) => Network a -> ActivationFunctionDerivative a -> V.Vector (Vector a) -> Vector a -> V.Vector (Matrix a)
deltas (Network{..}) act' os expected = V.zipWith outer (V.tail ds) (V.init os)
  where !dl = (V.last os - expected) * (deriv $ V.last os)
        !ds = V.scanr f dl (V.zip os matrices)
        f (!o, m) !del = deriv o * (tr m #> del)
        deriv = cmap act'
{-# INLINE deltas #-}

updateNetwork :: (Floating (Vector a), Floating a, Numeric a, Storable a, Num (Vector a), Container Vector a) => a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Sample a -> Network a
updateNetwork alpha act act' n@(Network{..}) (input, expectedOutput) = Network $ V.zipWith (+) matrices corr
    where !xs = outputs n act input
          !ds = deltas n act' xs expectedOutput
          !corr = V.map (scale (-alpha)) ds
{-# INLINE updateNetwork #-}

-- | Input vector and expected output vector
type Sample a = (Vector a, Vector a)

-- | List of 'Sample's
type Samples a = [Sample a]

-- | Handy operator to describe your learning set, avoiding unnecessary parentheses. It's just a synonym for '(,)'.
--   Generally you'll load your learning set from a file, a database or something like that, but it can be nice for
--   quickly playing with hnn or for simple problems where you manually specify your learning set.
--   That is, instead of writing:
--
-- > samples :: Samples Double
-- > samples = [ (fromList [0, 0], fromList [0])
-- >           , (fromList [0, 1], fromList [1])
-- >           , (fromList [1, 0], fromList [1])
-- >           , (fromList [1, 1], fromList [0])
-- >           ]
--
--   You can write:
--
-- > samples :: Samples Double
-- > samples = [ fromList [0, 0] --> fromList [0]
-- >           , fromList [0, 1] --> fromList [1]
-- >           , fromList [1, 0] --> fromList [1]
-- >           , fromList [1, 1] --> fromList [0]
-- >           ]
(-->) :: Vector a -> Vector a -> Sample a
(-->) = (,)

backpropOnce :: (Floating (Vector a), Floating a, Numeric a, Num (Vector a), Container Vector a) => a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Samples a -> Network a
backpropOnce rate act act' n samples = foldl' (updateNetwork rate act act') n samples
{-# INLINE backpropOnce #-}

-- | Generic training function.
--
-- The first argument is a predicate that will tell the backpropagation algorithm when to stop.
-- The first argument to the predicate is the epoch, i.e the number of times the backprop has been
-- executed on the samples. The second argument is /the current network/, and the third is the list of samples.
-- You can thus combine these arguments to create your own criterion.
--
-- For example, if you want to stop learning either when the network's quadratic error on the samples,
-- using the tanh function, is below 0.01, or after 1000 epochs, whichever comes first, you could
-- use the following predicate:
--
-- > pred epochs net samples = if epochs == 1000 then True else quadError tanh net samples < 0.01
--
-- You could even use 'Debug.Trace.trace' to print the error, to see how the error evolves while it's learning,
-- or redirect this to a file from your shell in order to generate a pretty graphics and what not.
--
-- The second argument (after the predicate) is the learning rate. Then come the activation function you want,
-- its derivative, the initial neural network, and your training set.
-- Note that we provide 'trainNTimes' and 'trainUntilErrorBelow' for common use cases.
trainUntil :: (Floating (Vector a), Floating a, Numeric a, Num (Vector a), Container Vector a) => (Int -> Network a -> Samples a -> Bool) -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Samples a -> Network a
trainUntil pr learningRate act act' net samples = go net 0
  where go n !k | pr k n samples = n
                | otherwise      = case backpropOnce learningRate act act' n samples of
                                    n' -> go n' (k+1)
{-# INLINE trainUntil #-}

-- | Trains the neural network with backpropagation the number of times specified by the 'Int' argument,
-- using the given learning rate (second argument).
trainNTimes :: (Floating (Vector a), Floating a, Numeric a, Num (Vector a), Container Vector a) => Int -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Samples a -> Network a
trainNTimes n = trainUntil (\k _ _ -> k > n)
{-# INLINE trainNTimes #-}

-- | Quadratic error on the given training set using the given activation function. Useful to create
-- your own predicates for 'trainUntil'.
quadError :: (Floating (Vector a), Floating a, Fractional (RealOf a), Normed (Vector a), Numeric a) => ActivationFunction a -> Network a -> Samples a -> RealOf a
quadError act net samples = realToFrac $ foldl' (\err (inp, out) -> err + (norm_2 $ output net act inp - out)) 0 samples
{-# INLINE quadError #-}

-- | Trains the neural network until the quadratic error ('quadError') comes below the given value (first argument),
-- using the given learning rate (second argument).
--
-- /Note/: this can loop pretty much forever when you're using a bad architecture for the problem, or inappropriate activation
-- functions.
trainUntilErrorBelow :: (Floating (Vector a), Floating a, Numeric a, Normed (Vector a), Ord a, Container Vector a, Num (RealOf a), a ~ RealOf a, Show a) => a -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Samples a -> Network a
trainUntilErrorBelow x rate act = trainUntil (\_ n s -> quadError act n s < x) rate act
{-# INLINE trainUntilErrorBelow #-}

-- | The sigmoid function:  1 / (1 + exp (-x))
sigmoid :: Floating a => a -> a
sigmoid !x = 1 / (1 + exp (-x))
{-# INLINE sigmoid #-}

-- | Derivative of the sigmoid function: sigmoid x * (1 - sigmoid x)
sigmoid' :: Floating a => a -> a
sigmoid' !x = case sigmoid x of
  s -> s * (1 - s)
{-# INLINE sigmoid' #-}

-- | Derivative of the 'tanh' function from the Prelude.
tanh' :: Floating a => a -> a
tanh' !x = case tanh x of
  s -> 1 - s**2
{-# INLINE tanh' #-}

-- | Loading a neural network from a file (uses zlib compression on top of serialization using the binary package).
--   Will throw an exception if the file isn't there.
loadNetwork :: (Storable a, Element a, Binary a) => FilePath -> IO (Network a)
loadNetwork fp = decode . decompress <$> B.readFile fp
{-# INLINE loadNetwork #-}

-- | Saving a neural network to a file (uses zlib compression on top of serialization using the binary package).
saveNetwork :: (Storable a, Element a, Binary a) => FilePath -> Network a -> IO ()
saveNetwork fp = B.writeFile fp . compress . encode
{-# INLINE saveNetwork #-}


-- `PolyKinds` ensures the derived HasRep instances are fully polymorphic. In future, we could try
-- to make this explicit in `deriveHasRep`.
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | The only module that should be needed by code that is in the dependency graph of calls to
--  `Categorifier.Categorify.expression`. If the module actually /calls/
--  `Categorifier.Categorify.expression`, then it will also need to
--   @import qualified "Categorifier.Categorify" as Categorify@.
module Categorifier.Client
  ( HasRep (..),
    Rep,
    deriveHasRep,

    -- * testing
    abstReturn,
    reprReturn,
  )
where

import Categorifier.Client.Internal (HasRep (..), Rep, deriveHasRep)
import Data.Complex (Complex)
import Data.Functor.Compose (Compose)
import Data.Functor.Identity (Identity)
import qualified Data.Functor.Product as Functor
import qualified Data.Functor.Sum as Functor
import Data.List.NonEmpty (NonEmpty)
import Data.Proxy (Proxy (..))
import Data.Ratio (Ratio)
import qualified Data.Semigroup as Semigroup
import qualified GHC.Generics as Generic

-- | This property should be true for every `HasRep` instance. `r` should be a classifying object,
--   and the first argument a comparison (e.g., `Bool` and `==` or `Hedgehog.Property` and
--  `Hedgehog.===`) to make it easy to test.
abstReturn :: forall a r. HasRep a => Proxy a -> (Rep a -> Rep a -> r) -> Rep a -> r
abstReturn Proxy eq = eq <$> id <*> repr . abst @a

-- | This property should be true for every `HasRep` instance. `r` should be a classifying object,
--   and the first argument a comparison (e.g., `Bool` and `==` or `Hedgehog.Property` and
--  `Hedgehog.===`) to make it easy to test.
reprReturn :: HasRep a => (a -> a -> r) -> a -> r
reprReturn eq = eq <$> id <*> abst . repr

deriveHasRep ''[]
deriveHasRep ''(,,)
deriveHasRep ''(,,,)
deriveHasRep ''(,,,,)
deriveHasRep ''(,,,,,)
deriveHasRep ''(,,,,,,)
deriveHasRep ''(,,,,,,,)
deriveHasRep ''(,,,,,,,,)
deriveHasRep ''(,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,,,,,,,,)
deriveHasRep ''(,,,,,,,,,,,,,,,,,,,,,,,,,,)
deriveHasRep ''Complex
deriveHasRep ''Compose
deriveHasRep ''Functor.Product
deriveHasRep ''Functor.Sum
deriveHasRep ''(Generic.:+:)
deriveHasRep ''(Generic.:*:)
deriveHasRep ''(Generic.:.:)
deriveHasRep ''Generic.K1
deriveHasRep ''Generic.M1
deriveHasRep ''Generic.Par1
deriveHasRep ''Generic.U1
deriveHasRep ''Identity
deriveHasRep ''Maybe
deriveHasRep ''NonEmpty
deriveHasRep ''Ordering
deriveHasRep ''Proxy
deriveHasRep ''Ratio
deriveHasRep ''Semigroup.All
deriveHasRep ''Semigroup.Any
deriveHasRep ''Semigroup.Dual
deriveHasRep ''Semigroup.Endo
deriveHasRep ''Semigroup.Max
deriveHasRep ''Semigroup.Min
deriveHasRep ''Semigroup.Product
deriveHasRep ''Semigroup.Sum


module Complex (
    Complex((:+)), realPart, imagPart, conjugate,
    mkPolar, cis, polar, magnitude, phase
  ) where

import Data.Complex


{-|
Module: MachineLearning.NeuralNetwork.Topology
Description: Neural Network's Topology
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX

Neural Network's Topology
-}

module MachineLearning.NeuralNetwork.Topology
(
  Topology
  , LossFunc(..)
  , makeTopology
  , loss
  , propagateForward
  , propagateBackward
  , numberOutputs
  , initializeTheta
  , initializeThetaIO
  , initializeThetaM
  , flatten
  , unflatten
)

where

import Control.Monad (zipWithM)
import Data.List (foldl')
import qualified Control.Monad.Random as RndM
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Utils (listOfTuplesToList)
import MachineLearning.NeuralNetwork.Layer (Layer(..), Cache(..))
import MachineLearning.NeuralNetwork.Regularization (Regularization, forwardReg, backwardReg)

-- | Loss function's type.
-- Takes x, weights and y.
type LossFunc = Matrix -> Matrix -> R

-- | Neural network topology has at least 2 elements: numver of input and number of outputs.
-- And sizes of hidden layers between 2 elements.
data Topology = Topology [(Int, Int)] [Layer] LossFunc

-- | Makes Neural Network's Topology.
-- Takes number of inputs, list of hidden layers, output layer and loss function.
makeTopology :: Int -> [Layer] -> Layer -> LossFunc -> Topology
makeTopology nInputs hiddenLayers outputLayer lossFunc =
  let layers = hiddenLayers ++ [outputLayer]
      layerSizes = nInputs : (map lUnits layers)
      sizes = getThetaSizes layerSizes
  in Topology sizes layers lossFunc

-- | Calculates loss for the given topology.
-- Takes topology, regularization, x, weights, y.
loss :: Topology -> Regularization -> Matrix -> [(Matrix, Matrix)] -> Matrix -> R
loss (Topology _ _ lf) reg x weights y =
  let lossValue = lf x y
      regValue = forwardReg reg weights
  in lossValue + regValue

-- | Implementation of forward propagation algorithm.
propagateForward :: Topology -> Matrix -> [(Matrix, Matrix)] -> (Matrix, [Cache])
propagateForward (Topology _ layers _) x thetaList = foldl' f (x, []) $ zip thetaList layers
  where f (a, cs) (theta, hl) =
          let (a', cache) = forwardPass hl a theta
          in (a', cache:cs)

-- | Makes one forward step for the given layer.
forwardPass :: Layer -> Matrix -> (Matrix, Matrix) -> (Matrix, Cache)
forwardPass layer a (b, w) = (a', Cache z a w)
  where z = lForward layer a b w
        a' = lActivation layer z

-- | Implementation of backward propagation algorithm.
propagateBackward :: Topology -> Regularization -> Matrix -> [Cache] -> Matrix -> [(Matrix, Matrix)]
propagateBackward (Topology _ layers _) reg scores (cache:cacheList) y = gradientList
  where cache' = Cache scores (cacheX cache) (cacheW cache)
        cacheList' = cache':cacheList
        gradientList = snd $ foldl' f (y, []) $ zip cacheList' $ reverse layers
        f (da, grads) (cache, hl) =
          let (da', db, dw) = backwardPass hl reg da cache
          in (da', (db, dw):grads)

-- | Makes one backward step for the given layer.
backwardPass :: Layer -> Regularization -> Matrix -> Cache -> (Matrix, Matrix, Matrix)
backwardPass layer reg da cache = (da', db, dw')
  where delta = lActivationGradient layer (cacheZ cache) da
        (da', db, dw) = lBackward layer delta cache
        dw' = dw + (backwardReg reg (cacheW cache))

-- | Returns number of outputs of the given topology.
numberOutputs :: Topology -> Int
numberOutputs (Topology nnt _ _) = fst $ last nnt

-- | Returns dimensions of weight matrices for given neural network topology
getThetaSizes :: [Int] -> [(Int, Int)]
getThetaSizes nn = zipWith (\r c -> (r, c)) (tail nn) nn

-- | Create and initialize weights vector with random values
-- for given neural network topology.
-- Takes a seed to initialize generator of random numbers as a first parameter.
initializeTheta :: Int -> Topology -> Vector
initializeTheta seed topology = RndM.evalRand (initializeThetaM topology) gen
  where gen = RndM.mkStdGen seed

-- | Create and initialize weights vector with random values
-- for given neural network topology inside IO Monad.
initializeThetaIO :: Topology -> IO Vector
initializeThetaIO = RndM.evalRandIO . initializeThetaM

-- | Create and initialize weights vector with random values
-- for given neural network topology inside RandomMonad.
initializeThetaM :: RndM.RandomGen g => Topology -> RndM.Rand g Vector
initializeThetaM topology = flatten <$> initializeThetaListM topology

-- | Create and initialize list of weights matrices with random values
-- for given neural network topology.
initializeThetaListM :: RndM.RandomGen g => Topology -> RndM.Rand g [(Matrix, Matrix)]
initializeThetaListM (Topology sizes layers _) = zipWithM lInitializeThetaM layers sizes

-- | Flatten list of matrices into vector.
flatten :: [(Matrix, Matrix)] -> Vector
flatten ms = V.concat $ map LA.flatten $ listOfTuplesToList ms

-- | Unflatten vector into list of matrices for given neural network topology.
unflatten :: Topology -> Vector -> [(Matrix, Matrix)]
unflatten (Topology sizes _ _) v =
  let offsets = reverse $ foldl' (\os (r, c) -> (r+r*c + head os):os) [0] (init sizes)
      ms = zipWith (\o (r, c) -> (LA.reshape r (slice o r), LA.reshape c (slice (o+r) (r*c)))) offsets sizes
      slice o n = V.slice o n v
  in ms


{-# LANGUAGE BangPatterns #-}
module STCPinwheel where

import           Control.Monad                  as M
import           Control.Monad.Parallel         as MP
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Binary
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           Filter.Utils
import           FokkerPlanck
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2D
import           FourierPinwheel
import           Image.IO
import           Pinwheel.FourierSeries2D
import           STC                            hiding (convolve)
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.Time

main = do
  args@(deviceIDsStr:numPointsStr:deltaStr:thresholdStr:numPointsReconStr:deltaReconStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:tauStr:numR2FreqStr:periodR2Str:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:initDistStr:initScaleStr:histFilePath:stdR2Str:stdStr:numBatchR2Str:numBatchR2FreqsStr:numBatchOriStr:batchSizeStr:sStr:radiusStr:deltaTStr:numThreadStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoints = read numPointsStr :: Int
      delta = read deltaStr :: Double
      threshold = read thresholdStr :: Double
      numPointsRecon = read numPointsReconStr :: Int
      deltaRecon = read deltaReconStr :: Double
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      tau = read tauStr :: Double
      numR2Freq = read numR2FreqStr :: Int
      periodR2 = read periodR2Str :: Double
      phiFreq = read phiFreqsStr :: Int
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Int
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Int
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Int
      scaleFreqs = L.map fromIntegral [-scaleFreq .. scaleFreq]
      initScale = read initScaleStr :: Double
      initDist = read initDistStr :: [(Double, Double, Double, Double)]
      initPoints = L.map (\(x, y, t, s) -> Point x y t s) initDist
      initSource = [L.head initPoints]
      initSink = [L.last initPoints]
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCPinwheel"
      stdR2 = read stdR2Str :: Double
      std = read stdStr :: Double
      numBatchR2 = read numBatchR2Str :: Int
      numBatchR2Freqs = read numBatchR2FreqsStr :: Int
      numBatchOri = read numBatchOriStr :: Int
      batchSize = read batchSizeStr :: Int
      s = read sStr :: Double
      radius = read radiusStr :: Double
      deltaT = read deltaTStr :: Double
      periodEnvelope = periodR2 ^ 2 * 2
      maxScale = sqrt periodEnvelope
  removePathForcibly folderPath
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath

  hist <-
    if flag
      then do printCurrentTime "read from files..."
              decodeFile histFilePath
      else do printCurrentTime "Start computing coefficients..."
              initialise []
              devs <- M.mapM device deviceIDs
              ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
              ptxs <- M.mapM createTargetFromContext ctxs
              runMonteCarloFourierCoefficientsGPU
                deviceIDs
                numThread
                1000000
                500000
                thetaSigma
                scaleSigma
                0
                s
                tau
                deltaT
                phiFreqs
                rhoFreqs
                thetaFreqs
                scaleFreqs
                maxScale
                histFilePath
      -- else sampleCartesian
      --        histFilePath
      --        folderPath
      --        ptxs
      --        numPoints
      --        periodEnvelope
      --        delta
      --        numOrientation
      --        initScale
      --        thetaSigma
      --        tau
      --        threshold
      --        s
      --        phiFreq
      --        rhoFreq
      --        thetaFreq
      --        scaleFreq
  printCurrentTime "Done"
  printCurrentTime "Start Convloution.."
  plan <-
    makePlanDiscrete
      folderPath
      emptyPlan
      numPointsRecon
      numPointsRecon
      numR2Freq
      (2 * thetaFreq + 1)
      (2 * scaleFreq + 1)
      (2 * phiFreq + 1)
      (2 * rhoFreq + 1)
  let coefficients = getNormalizedHistogramArr $ hist
  -- coefficients <- hollowCoefficients plan radius periodEnvelope coefficients'
      -- coefficients = getNormalizedHistogramArr $ hist
        -- normalizeFreqArr' std phiFreqs rhoFreqs .
  harmonicsArray <-
    createHarmonics
      numR2Freq
      phiFreq
      rhoFreq
      thetaFreq
      scaleFreq
      (-s)
      periodR2
      periodEnvelope
      coefficients
      -- harmonicsArray1 =
      --   createHarmonics
      --     numR2Freq
      --     phiFreq
      --     rhoFreq
      --     thetaFreq
      --     scaleFreq
      --     (-1.25)
      --     periodR2
      --     periodEnvelope
      --     coefficients
  let !initialSourceDistribution =
        computeInitialDistributionFourierPinwheel
          numR2Freq
          periodR2
          periodEnvelope
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          initSource
      !initialSinkDistribution =
        computeInitialDistributionFourierPinwheel
          numR2Freq
          periodR2
          periodEnvelope
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          initSink
      -- harmonicsArray' =
      --   centerHollow numR2Freq $
      --   pinwheelFourierCoefficientsAnatical
      --     numR2Freq
      --     phiFreq
      --     rhoFreq
      --     thetaFreq
      --     scaleFreq
      --     (-s)
      --     periodR2
      --     periodEnvelope
      -- harmonicsArray1' =
      --   -- L.map (centerHollowVector numR2Freq) $
      --   pinwheelFourierCoefficientsAnaticalList1
      --     numR2Freq
      --     phiFreq
      --     rhoFreq
      --     thetaFreq
      --     scaleFreq
      --     (-s)
      --     periodR2
      --     periodEnvelope
      -- harmonicsArray2' =
      --   -- L.map (centerHollowVector numR2Freq) $
      --   pinwheelFourierCoefficientsAnaticalList2
      --     numR2Freq
      --     phiFreq
      --     rhoFreq
      --     thetaFreq
      --     scaleFreq
      --     (-s)
      --     periodR2
      --     periodEnvelope
      -- initialSourceDistribution =
      --   computeInitialDistributionFull'
      --     numR2Freq
      --     periodR2
      --     phiFreq
      --     rhoFreq
      --     initSource
      -- initialSourceDistribution =
      --   computeInitialDistributionPowerMethodPinwheelBasis'
      --     numR2Freq
      --     (-s)
      --     periodR2
      --     phiFreq
      --     rhoFreq
      --     initSource
      -- initialSinkDistribution =
      --   computeInitialDistributionFull'
      --     numR2Freq
      --     periodR2
      --     phiFreq
      --     rhoFreq
      --     initSink
  -- source <-
  --   convoluvePinhweelBasisTest'
  --     coefficients
  --     harmonicsArray1'
  --     harmonicsArray2'
  --     initialSourceDistribution
  -- sink <-
  --   timeReversal' <$>
  --   convoluvePinhweelBasisTest'
  --     coefficients
  --     harmonicsArray1'
  --     harmonicsArray2'
  --     initialSinkDistribution
  -- let source =
  --       convoluvePinhweelBasis'
  --         coefficients
  --         harmonicsArray'
  --         initialSourceDistribution
  --     sink =
  --       timeReversal' $
  --       convoluvePinhweelBasis'
  --         coefficients
  --         harmonicsArray'
  --         initialSinkDistribution
  --     sourceMat =
  --       A.transpose . A.use . fromDFTArray . getDFTArrayVector $ source
  --     sinkMat = A.transpose . A.use . fromDFTArray . getDFTArrayVector $ sink-
  source <- convolve harmonicsArray initialSourceDistribution
  -- source1 <- convolve harmonicsArray1 initialSourceDistribution
  let numLogPolarFreq = ((2 * scaleFreq + 1) * (2 * thetaFreq + 1))
      -- sourceMat = A.transpose . toMatrixAcc $ source
      -- sourceMat1 = A.transpose . toMatrixAcc $ source1
      -- sourceMat2 =
      --   A.transpose . toMatrixAcc $
      --   parZipWithFPArray (VS.zipWith (-)) source source1
  -- sourceR2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     numLogPolarFreq
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     sourceMat
  -- plotImageRepa (folderPath </> "SourcePower.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   sourceR2
  sourceR2 <- plotFPArray plan (folderPath </> "SourcePower.png") source
  -- source1R2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     numLogPolarFreq
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     sourceMat1
  -- plotImageRepa (folderPath </> "SourcePower1.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   source1R2
  -- source2R2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     numLogPolarFreq
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     sourceMat2
  -- plotImageRepa (folderPath </> "SourcePower2.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   source2R2
  sink <- convolve harmonicsArray initialSinkDistribution
  let sinkMat = A.transpose . toMatrixAcc $ sink
  -- completion <- completionField' plan source (timeReversal' sink)
  -- let completionMat = createCuMat . getDFTArrayVector $ completion
  -- print completionMat
  sinkR2 <- plotFPArray plan (folderPath </> "SinkPower.png") sink
  -- sinkR2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     numLogPolarFreq
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     sinkMat
  -- plotImageRepa (folderPath </> "SinkPower.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   sinkR2
  -- completionR2 <-
  --   computeFourierSeriesR2Stream
  --     deviceIDs
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     [transposeCuMat completionMat]
  printCurrentTime "Start ploting.."
  -- plotDFTArrayPower
  --   (folderPath </> "CoefficientsPower.png")
  --   numR2Freq
  --   numR2Freq
  --   source
  -- plotImageRepaComplex (folderPath </> "Coefficients.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) .
  --   VU.convert . L.foldl1' (VS.zipWith (+)) . getDFTArrayVector $
  --   source
  -- plotImageRepaComplex (folderPath </> "Source.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   toUnboxed . sumS . rotate3D $
  --   sourceR2
  completionR2 <-
    completionFieldRepa
      plan
      (R.reshape
         (Z :. (L.length scaleFreqs) :. (L.length thetaFreqs) :. numPointsRecon :.
          numPointsRecon) $
       sourceR2)
      (timeReversalRepa thetaFreqs .
       R.reshape
         (Z :. (L.length scaleFreqs) :. (L.length thetaFreqs) :. numPointsRecon :.
          numPointsRecon) $
       sinkR2)
  plotImageRepa (folderPath </> "CompletionPower.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    completionR2
  printCurrentTime "Done"
  -- sparse
  let harmonicsArraySparse =
        computeHarmonicsArraySparse
          numPointsRecon
          deltaRecon
          numPointsRecon
          deltaRecon
          phiFreqs
          rhoFreqs
          thetaFreqs
          scaleFreqs
          -- (2 * pi)
          (log (periodR2 / 4))
          64
          s
      init =
        computeInitialDistribution'
          numPointsRecon
          numPointsRecon
          phiFreqs
          rhoFreqs
          0
          initSource
  harmonicsArrayDFT <-
    fmap (listArray (bounds harmonicsArraySparse)) .
    dftExecuteBatchP
      plan
      (DFTPlanID DFT1DG [numPointsRecon, numPointsRecon] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    harmonicsArraySparse
  sourceSparse <- convolve' Source plan coefficients harmonicsArrayDFT init
  plotDFTArrayPower
    (folderPath </> "SourcePower_Sparse.png")
    numPointsRecon
    numPointsRecon
    sourceSparse


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, ConstraintKinds, ExistentialQuantification, GADTs, RankNTypes, MultiParamTypeClasses, RankNTypes, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : First attempt at approximate counting
-- | Creator: Xiao Ling
-- | Created: 12/08/2015
-- | TODO   : test standard deviation of alpha, beta, and final version
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module Morris (

    morris
  , morris'

  , morris1
  , morris1'

  ) where

import Control.Monad.Random
import Control.Monad.Random.Class

import Data.Conduit
import Data.List.Split
import qualified Data.Conduit.List as Cl

import Core
import Statistics

{-----------------------------------------------------------------------------
    I. Morris Algorithm list of counter
------------------------------------------------------------------------------}

---- * Count the number of items in `as` to within `eps` of actual
---- * with confidence `delta`
morris' :: EpsDelta -> Batch a IO Counter
morris' ed = morris ed `using` evalRandIO

morris :: MonadRandom m => EpsDelta -> Streaming a m Counter
morris (ED e d) = medianOfMeans $ go cs
  where cs            = replicate (t*m) 0
        t             = round $ 1/(e^2*d)  :: Int
        m             = round . log $ 1/d
        go cs         = (\c -> 2^(round c) - 1) `ffmap` Cl.foldM (\cs _ -> traverse incr cs) cs
        ffmap         = fmap . fmap
        medianOfMeans = fmap median' . (fmap . fmap) mean' . fmap (chunksOf t)

{-----------------------------------------------------------------------------
    II. Morris Algorithm list of list of counter
------------------------------------------------------------------------------}

---- * Count the number of items in `as` to within `eps` of actual
---- * with confidence `delta`
morris1' :: EpsDelta -> Batch a IO Counter
morris1' ed = morris1 ed `using` evalRandIO

-- * Run on stream inputs `as` for t independent trials for `t = 1/eps^2 * d`,
-- * and `m` times in parallel, for `m = log(1/d)` and take the median
morris1 :: MonadRandom m => EpsDelta -> Streaming a m Counter
morris1 (ED e d) = medianOfMeans $ go ccs
      where
        medianOfMeans = fmap median' . (fmap . fmap) mean'
        ccs           = replicate m $ replicate t 0
        t             = round $ 1/(e^2*d)
        m             = round . log $ 1/d
        go ccs        = (\x -> 2^(round x) - 1) `fffmap` Cl.foldM (\xs _ -> incrs' xs) ccs
        fffmap        = fmap . fmap . fmap

-- * given a list of list of counters toss a coin for each counter and incr
-- * this can be flattened
incrs' :: MonadRandom m => [[Counter]] -> m [[Counter]]
incrs' = sequence . fmap (sequence . fmap incr)

{-----------------------------------------------------------------------------
    III. Utils
------------------------------------------------------------------------------}

-- * Increment a counter `x` with probability 1/2^x
incr :: MonadRandom m => Counter -> m Counter
incr x = do
  h <- toss . coin $ 0.5^(round x)
  return $ if isHead h then (seq () succ x) else seq () x

mean', median' :: (Floating a, Ord a, RealFrac a) => [a] -> Float
mean'   = fromIntegral . round . mean
median' = fromIntegral . round . median


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE IncoherentInstances  #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Vector.Image.Color.RGB (
  RGB(..)
, ColorSpace(..)
, RGB8
) where

import Prelude hiding(map)
import qualified Prelude as P
import qualified Data.List as L
import Data.Vector.Storable hiding(forM_,map,replicate,fromList,toList,(!?),(!))
import Data.Maybe
import Numeric.LinearAlgebra.Data hiding (fromList,toList,(!))
import Numeric.LinearAlgebra.HMatrix hiding (fromList,toList,(!))
import qualified Data.Vector.Storable as V
import Foreign
import GHC.Real
import Foreign.C.String
import Foreign.C.Types
import qualified Control.Monad as C

data RGB a = RGB {
   rgb_r :: a
 , rgb_g :: a
 , rgb_b :: a
} deriving (Show,Eq,Ord)

type RGB8 = RGB Word8

class ColorSpace color b where
  toRGB :: color b -> RGB b
  fromRGB :: RGB b -> color b
  toList :: color b -> [b]
  fromList :: [b] -> color b
  dim :: color b -> Int
  dummyElem :: color b -> b
  dummyElem _ = error "do not eval this"

instance Functor RGB where
  fmap func (RGB r g b) = RGB (func r) (func g) (func b)

op :: (ColorSpace c a,Num a) => (a -> a -> a) -> c a -> c a -> c a
op func rgb rgb2 =
  let rgb' = toList rgb
      rgb2' = toList rgb2
      in fromList $ P.map (uncurry func) $ zip rgb' rgb2'

fmap' :: (ColorSpace c a, ColorSpace c b)
     => (a -> b)
     -> c a
     -> c b
fmap' func val = fromList $ P.map func $ toList val

instance (ColorSpace c a,Num a) => Num (c a) where
  (+) = op (+)
  (-) = op (-)
  (*) = op (*)
  negate = fmap' negate
  abs = fmap' abs
  signum = fmap' signum
  fromInteger v = fromList $ replicate (dim dummyColor) v'
    where
      v' = fromIntegral v
      dummyColor = error "do not eval this" :: (c a)

instance Enum (RGB Int) where
  succ rgb =  rgb + RGB 1 1 1
  pred rgb =  rgb - RGB 1 1 1
  toEnum i = fromIntegral i
  fromEnum (RGB r g b) = r
  enumFrom rgb = rgb : (enumFrom $ succ rgb)
  enumFromThen rgb rgb1 = rgb : (enumFromThen rgb1 (2 * rgb1 - rgb))

instance Real (RGB Int) where
  toRational (RGB r g b) = (fromIntegral $ (r + 2*g + b) `div` 4) % 1

instance Integral (RGB Int) where
  quot = op quot
  rem = op rem
  div = op div
  mod = op mod
  quotRem a b = (quot a b,rem a b)
  divMod a b = (div a b,mod a b)
  toInteger (RGB r g b) = fromIntegral $ (r + 2*g + b) `div` 4

instance (Storable a) => Element a
instance (ColorSpace c a,Storable a) => Element (c a)

instance (ColorSpace c a,Storable a) => Storable (c a) where
  sizeOf    d = dim d * sizeOf (dummyElem d)
  alignment d = dim d * sizeOf (dummyElem d)
  peek ptr = do
    v <- C.forM [0..((dim d)-1)] $ \idx -> do
      peekByteOff ptr (sizeOf (dummyElem d) * idx)
    return $ fromList v
    where
      d = error "do not eval this" :: c a
  poke ptr val = do
    C.forM_ (zip [0..] (toList val)) $ \(idx,v) -> do
      pokeByteOff ptr (sizeOf (dummyElem d) * idx) v
    where
      d = error "do not eval this" :: c a

instance ColorSpace RGB a where
  toRGB = id
  fromRGB = id
  toList (RGB r g b) = [r,g,b]
  fromList [r,g,b] = (RGB r g b)
  fromList val = error $ "List length must be 3 but " L.++  show (L.length val)
  dim _ = 3


{-# LANGUAGE FlexibleInstances #-}
import Criterion.Main
import PCA.QModels
import PCA.Distribution
import Numeric.LinearAlgebra.Data (Vector, Matrix, fromRows, toRows)
import qualified Numeric.LinearAlgebra.HMatrix as H

import Control.DeepSeq

instance NFData (Distr Double Double) where
  rnf a = seq a ()

instance NFData (Distr (Vector Double) (Vector Double)) where
  rnf a = seq a ()

instance NFData (Distr (Vector Double) (Matrix Double)) where
  rnf a = seq a ()

instance NFData (Distr (Matrix Double) (Matrix Double)) where
  rnf a = seq a ()

prepareData n d = do
  trainT <- H.rand n d
  let tau = initQtau
      mu = initQmu d
      alpha = initQalpha d
      w = initQW d
      sigmaX = calcSigmaX tau w
  return (trainT, tau, mu, alpha, w, sigmaX)

foolCalcMX
    :: Matrix Double
    -> Distr Double Double
    -> Matrix Double
    -> Distr (Matrix Double) (Matrix Double)
    -> Distr (Vector Double) (Matrix Double)
    -> Matrix Double
foolCalcMX trainT tau sigmaX w mu =
    fromRows
        (map
             (\r ->
                   coef H.#> (r - mean mu))
             (toRows trainT))
  where
    coef = mean tau `H.scale` sigmaX H.<> H.tr (mean w)

main =
    defaultMain
        [ env (prepareData 100 3) $
          \small ->
               bgroup
                   "Small 100x3"
                   [ bench "Fool" $
                     whnf
                         (\(trainT,tau,mu,alpha,w,sigmaX) ->
                               foolCalcMX trainT tau sigmaX w mu)
                         small
                   , bench "Matrix" $
                     whnf
                         (\(trainT,tau,mu,alpha,w,sigmaX) ->
                               calcMX trainT tau sigmaX w mu)
                         small]
        , env (prepareData 10000 30) $
          \big ->
               bgroup
                   "Big 10000x30"
                   [ bench "Fool" $
                     whnf
                         (\(trainT,tau,mu,alpha,w,sigmaX) ->
                               foolCalcMX trainT tau sigmaX w mu)
                         big
                   , bench "Matrix" $
                     whnf
                         (\(trainT,tau,mu,alpha,w,sigmaX) ->
                               calcMX trainT tau sigmaX w mu)
                         big]]