{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Scheme.Types where

import Control.Exception (Exception, throw)
import Control.Monad.Except (ExceptT)
import Data.Array (Array, elems)
import qualified Data.ByteString as BS
import Data.Complex (Complex ((:+)), imagPart, realPart)
import Data.IORef (IORef)
import Data.Ratio (denominator, numerator)
import Data.Text (Text)
import qualified Data.Text as T
import Data.Typeable (Typeable)
import System.IO (Handle)
import Text.Pretty.Simple (pPrint, pPrintLightBg)
import Text.Show.Functions ()

type Env = IORef [(Text, IORef SchemeVal)]

type SchemeResult = Either SchemeError

type IOSchemeResult = ExceptT SchemeError IO

extractValue :: Either a p -> p
extractValue (Right val) = val
extractValue _ = error "Called with Left"

data Number
  = Integer Integer
  | Real Double
  | Rational Rational
  | Complex (Complex Double)
  deriving (Show, Typeable, Eq)

showNum :: Number -> Text
showNum (Integer i) = T.pack $ show i
showNum (Real d) = T.pack $ show d
showNum (Rational r) = T.pack $ show (numerator r) <> "/" <> show (denominator r)
showNum (Complex p) = T.pack $ show (realPart p) <> "+" <> show (imagPart p) <> "i"

instance Ord Number where
  compare (Integer x) (Integer y) = compare x y
  compare (Integer x) (Real y) = compare (fromInteger x) y
  compare (Integer x) (Rational y) = compare (fromInteger x) y
  compare (Real x) (Integer y) = compare x (fromInteger y)
  compare (Rational x) (Integer y) = compare x (fromInteger y)
  compare (Real x) (Real y) = compare x y
  compare (Real x) (Rational y) = compare (toRational x) y
  compare (Rational x) (Real y) = compare x (toRational y)
  compare (Rational x) (Rational y) = compare x y
  compare _ _ = throw $ InvalidOperation "Cannot compare complex numbers"

instance Num Number where
  fromInteger x = Integer x

  (Integer x) + (Integer y) = Integer (x + y)
  (Integer x) + (Real y) = Real (fromInteger x + y)
  (Integer x) + (Rational y) = Rational (fromInteger x + y)
  (Integer x) + (Complex y) = Complex (fromInteger x + y)
  (Real x) + (Integer y) = Real (x + fromInteger y)
  (Rational x) + (Integer y) = Rational (x + fromInteger y)
  (Complex x) + (Integer y) = Complex (x + fromInteger y)
  (Real x) + (Real y) = Real (x + y)
  (Real x) + (Rational y) = Rational (toRational x + y)
  (Real x) + (Complex y) = Complex ((x :+ 0.0) + y)
  (Rational x) + (Real y) = Rational (x + toRational y)
  (Complex x) + (Real y) = Complex (x + (y :+ 0.0))
  (Rational x) + (Rational y) = Rational (x + y)
  (Complex x) + (Complex y) = Complex (x + y)
  _ + _ = throw $ InvalidOperation "Cannot add rationals and complex numbers"

  (Integer x) * (Integer y) = Integer (x * y)
  (Integer x) * (Real y) = Real (fromInteger x * y)
  (Integer x) * (Rational y) = Rational (fromInteger x * y)
  (Integer x) * (Complex y) = Complex (fromInteger x * y)
  (Real x) * (Integer y) = Real (x * fromInteger y)
  (Rational x) * (Integer y) = Rational (x * fromInteger y)
  (Complex x) * (Integer y) = Complex (x * fromInteger y)
  (Real x) * (Real y) = Real (x * y)
  (Real x) * (Rational y) = Rational (toRational x * y)
  (Real x) * (Complex y) = Complex ((x :+ 0.0) * y)
  (Rational x) * (Real y) = Rational (x * toRational y)
  (Complex x) * (Real y) = Complex (x * (y :+ 0.0))
  (Rational x) * (Rational y) = Rational (x * y)
  (Complex x) * (Complex y) = Complex (x * y)
  _ * _ = throw $ InvalidOperation "Cannot multiply rationals and complex numbers"

  (Integer x) - (Integer y) = Integer (x - y)
  (Integer x) - (Real y) = Real (fromInteger x - y)
  (Integer x) - (Rational y) = Rational (fromInteger x - y)
  (Integer x) - (Complex y) = Complex (fromInteger x - y)
  (Real x) - (Integer y) = Real (x - fromInteger y)
  (Rational x) - (Integer y) = Rational (x - fromInteger y)
  (Complex x) - (Integer y) = Complex (x - fromInteger y)
  (Real x) - (Real y) = Real (x - y)
  (Real x) - (Rational y) = Rational (toRational x - y)
  (Real x) - (Complex y) = Complex ((x :+ 0.0) - y)
  (Rational x) - (Real y) = Rational (x - toRational y)
  (Complex x) - (Real y) = Complex (x - (y :+ 0.0))
  (Rational x) - (Rational y) = Rational (x - y)
  (Complex x) - (Complex y) = Complex (x - y)
  _ - _ = throw $ InvalidOperation "Cannot subtract rationals and complex numbers"

  negate (Integer x) = Integer (negate x)
  negate (Real x) = Real (negate x)
  negate (Rational x) = Rational (negate x)
  negate (Complex x) = Complex (negate x)

  abs (Integer x) = Integer (abs x)
  abs (Real x) = Real (abs x)
  abs (Rational x) = Rational (abs x)
  abs (Complex x) = Complex (abs x)

  signum (Integer x) = Integer (signum x)
  signum (Real x) = Real (signum x)
  signum (Rational x) = Rational (signum x)
  signum (Complex x) = Complex (signum x)

instance Fractional Number where
  fromRational x = Rational x

  (Integer x) / (Integer y) = Rational (fromInteger x / fromInteger y)
  (Integer x) / (Real y) = Real (fromInteger x / y)
  (Integer x) / (Rational y) = Rational (fromInteger x / y)
  (Integer x) / (Complex y) = Complex (fromInteger x / y)
  (Real x) / (Integer y) = Real (x / fromInteger y)
  (Rational x) / (Integer y) = Rational (x / fromInteger y)
  (Complex x) / (Integer y) = Complex (x / fromInteger y)
  (Real x) / (Real y) = Real (x / y)
  (Real x) / (Rational y) = Rational (toRational x / y)
  (Real x) / (Complex y) = Complex ((x :+ 0.0) / y)
  (Rational x) / (Real y) = Rational (x / toRational y)
  (Complex x) / (Real y) = Complex (x / (y :+ 0.0))
  (Rational x) / (Rational y) = Rational (x / y)
  (Complex x) / (Complex y) = Complex (x / y)
  _ / _ = throw $ InvalidOperation "Cannot divide rationals and complex numbers"

  recip (Integer x) = Real (recip $ fromInteger x)
  recip (Real x) = Real (recip x)
  recip (Rational x) = Rational (recip x)
  recip (Complex x) = Complex (recip x)

data Fn = Fn {macro :: Bool, params :: [Text], vararg :: Maybe Text, body :: [SchemeVal], closure :: Env}

instance Show Fn where
  show _ = "<fn>"

data SchemeVal
  = List [SchemeVal]
  | PairList [SchemeVal] SchemeVal
  | Vector (Array Int SchemeVal)
  | Bytevector BS.ByteString
  | String Text
  | Character Char
  | Symbol Text
  | Boolean Bool
  | Number Number
  | Nil
  | Primitive ([SchemeVal] -> SchemeResult SchemeVal)
  | Fun Fn
  | IOFun ([SchemeVal] -> IOSchemeResult SchemeVal)
  | Port Handle
  deriving (Show, Typeable)

instance Eq SchemeVal where
  (==) (List x) (List y) = length x == length y && all (uncurry (==)) (zip x y)
  (==) (PairList xs x) (PairList ys y) = List (xs ++ [x]) == List (ys ++ [y])
  (==) (Vector x) (Vector y) = x == y
  (==) (Bytevector x) (Bytevector y) = x == y
  (==) (String x) (String y) = x == y
  (==) (Character x) (Character y) = x == y
  (==) (Symbol x) (Symbol y) = x == y
  (==) (Boolean x) (Boolean y) = x == y
  (==) (Number x) (Number y) = x == y
  (==) _ _ = False

unwordVals :: [SchemeVal] -> Text
unwordVals xs = T.unwords $ showVal <$> xs

showVal :: SchemeVal -> Text
showVal (List (Symbol "quote" : xs)) = "'" <> unwordVals xs
showVal (List (Symbol "quasiquote" : xs)) = "`" <> unwordVals xs
showVal (List (Symbol "unquote" : xs)) = "," <> unwordVals xs
showVal (List (Symbol "unquote-splicing" : xs)) = ",@" <> unwordVals xs
showVal (List contents) = "(" <> unwordVals contents <> ")"
showVal (PairList contents cdr) = "(" <> unwordVals contents <> " . " <> showVal cdr <> ")"
showVal (Vector vec) = T.pack $ "#(" <> unwords (map show $ elems vec) <> ")"
showVal (Bytevector vec) = T.pack $ "#u8(" <> unwords (map show $ BS.unpack vec) <> ")"
showVal (String s) = "\"" <> s <> "\""
showVal (Character a) = T.pack $ "#\\" <> [a]
showVal (Symbol s) = s
showVal (Boolean True) = "#t"
showVal (Boolean False) = "#f"
showVal (Number n) = showNum n
showVal Nil = "nil"
showVal Primitive {} = "<prim>"
showVal Fun {} = "<fun>"
showVal IOFun {} = "<io>"
showVal (Port _) = "<port>"

dumpAST :: SchemeVal -> IO ()
dumpAST = dumpAST' True

dumpAST' :: Bool -> SchemeVal -> IO ()
dumpAST' True = pPrint
dumpAST' False = pPrintLightBg

data SchemeError
  = Generic Text
  | ArgumentLengthMismatch Int [SchemeVal]
  | TypeMismatch Text SchemeVal
  | UnboundSymbol Text
  | ParserError Text
  | NotFunction SchemeVal
  | InvalidOperation Text
  | ReservedName Text
  | EmptyList
  deriving (Show)

instance Exception SchemeError

showError :: SchemeError -> Text
showError (TypeMismatch err vap) = "Invalid type: expected " <> err <> ", but found " <> T.pack (show vap)
showError (Generic err) = "Unexpectec error: " <> err
showError (UnboundSymbol sym) = "Unbound symbol: " <> sym
showError (ArgumentLengthMismatch ex act) = "Expected " <> T.pack (show ex) <> " but found " <> T.pack (show $ length act)
showError (ParserError err) = "Parsing error, could not parse input: " <> err
showError (NotFunction err) = "Attempt at calling " <> showVal err <> " as a function"
showError (InvalidOperation err) = "Invalid: " <> err
showError (ReservedName err) = "Cannot define : " <> err <> " because it is a reserved name"
showError EmptyList = "Empty list is not allowed"


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

import           Data.Complex
import           Graphics.SvgTree
import           Linear.V2
import           Reanimate
import           Reanimate.Ease
import           Codec.Picture

-- layer 3
main :: IO ()
main = reanimate $ setDuration 30 $ sceneAnimation $ do
    _ <- newSpriteSVG $ mkBackgroundPixel (PixelRGBA8 252 252 252 0xFF)
    play $ fourierA (fromToS 0 5)      -- Rotate 15 times
      # setDuration 50
      # signalA (reverseS . powerS 2 . reverseS) -- Start fast, end slow
      # pauseAtEnd 2
    play $ fourierA (constantS 0)       -- Don't rotate at all
      # setDuration 10
      # reverseA
      # signalA (powerS 2)                       -- Start slow, end fast
      # pauseAtEnd 2

-- layer 2
fourierA :: (Double -> Double) -> Animation
fourierA genPhi = animate $ \t ->
    let circles = setFourierLength (t*piFourierLen) piFourier
        coeffs = fourierCoefficients $ rotateFourier (genPhi t) circles
    in mkGroup
    [ drawCircles coeffs
    , withStrokeColor "green" $
      withStrokeLineJoin JoinRound $
      withFillOpacity 0 $
      withStrokeWidth (defaultStrokeWidth*2) $
      mkLinePath $ mkFourierOutline circles
    , let x :+ y = sum coeffs in
      translate x y $ withFillColor "red" $ mkCircle (defaultStrokeWidth*3)
    ]

drawCircles :: [Complex Double] -> SVG
drawCircles [] = mkGroup []
drawCircles ( x :+ y : xs) =
  translate x y $ drawCircles' xs

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

-- layer 1
newtype Fourier = Fourier {fourierCoefficients :: [Complex Double]}

piFourier :: Fourier
piFourier = mkFourier $ lineToPoints 500 $
  toLineCommands $ extractPath $ scale 15 $
  center $ latexAlign "\\pi"

piFourierLen :: Double
piFourierLen = sum $ map magnitude $ drop 1 $ take 500 $ fourierCoefficients piFourier

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

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

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

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

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


{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}

module Grenade.Layers.Internal.BLAS
    ( -- easy interface
      matXVec
    , outerV
    , checkVectors
    , unsafeMemCopyVectorFromTo
    , memCopyVectorFromTo
    , unsafeMemZero
    , memZero
      -- more complicated, but direct, function calls
    , BlasTranspose (..)
    , swapTranspose
    , dgemmUnsafe
    , dgemvUnsafe
    , dgerUnsafe
    ) where

import           Control.Monad
import           Data.IORef
import           Data.Proxy
import qualified Data.Vector.Storable         as V
import           Foreign                      (withForeignPtr)
import           Foreign.C.Types
import           Foreign.Ptr
import           Foreign.Storable             (sizeOf)
import           GHC.IO.Handle.Text           (memcpy)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra        as LA
import qualified Numeric.LinearAlgebra.Static as LAS
import           System.IO.Unsafe             (unsafePerformIO)

import           Grenade.Layers.Internal.CUDA
import           Grenade.Types
import           Grenade.Utils.Vector

import           Debug.Trace

#define USE_DGEMM_ONLY 0

-- | Computes vec2 <- mat * vec1 + beta * vec2.
matXVec :: BlasTranspose -> V.Vector RealNum -> V.Vector RealNum -> RealNum -> V.Vector RealNum -> IO (V.Vector RealNum)
matXVec trMat mat vec1 beta vec2 =
#if USE_DGEMM_ONLY
  dgemmUnsafe trMat BlasNoTranspose (m, k) (ay, 1) 1.0 mat vec1 beta vec2
# else
  dgemvUnsafe trMat (m, k) 1.0 mat vec1 beta vec2
#endif
  where
    ay = V.length vec1
    ax = V.length vec2
    (m, k) = swapTranspose trMat (ax, ay)

-- | Computes the outer product of two vectors: mat <- vec1 `outer` vec2
outerV :: V.Vector RealNum -> V.Vector RealNum -> IO (V.Vector RealNum)
outerV vec1 vec2 =
#if USE_DGEMM_ONLY
  dgemmUnsafe BlasNoTranspose BlasNoTranspose (o, 1) (1, i) 1.0 vec1 vec2 0 (createVectorUnsafe (i * o)) -- beta = 0 initialises the matrix
#else
  createVector (i * o) >>= memZero >>= dgerUnsafe (o, i) 1.0 vec1 vec2
#endif
  where o = V.length vec1
        i = V.length vec2

-- | Check two vectors if they are equal. For testing purposes.
checkVectors :: V.Vector RealNum -> V.Vector RealNum -> Bool
checkVectors v1 v2 = V.length v1 == V.length v2 && and (zipWith (==) (toStr v1) (toStr v2))
  where
    toStr :: V.Vector RealNum -> [String]
    toStr v = map (show . round . (*10^5)) $ V.toList v
{-# INLINE checkVectors #-}

-- | Newtype holding CINT for Transpose values.
newtype CBLAS_TRANSPOSET =
  CBLAS_TransposeT CInt
  deriving (Eq, Show)

-- | Transpose values
data BlasTranspose
  = BlasNoTranspose
  | BlasTranspose
  | BlasConjTranspose
  | BlasConjNoTranspose
  deriving (Eq, Show)

encodeTransposeIntBool :: BlasTranspose -> Int
encodeTransposeIntBool BlasNoTranspose     = 0
encodeTransposeIntBool BlasTranspose       = 1
encodeTransposeIntBool BlasConjTranspose   = 1
encodeTransposeIntBool BlasConjNoTranspose = 0
{-# INLINE encodeTransposeIntBool #-}

swapTranspose :: BlasTranspose -> (Int, Int) -> (Int, Int)
swapTranspose BlasNoTranspose x        = x
swapTranspose BlasTranspose (a, b)     = (b, a)
swapTranspose BlasConjNoTranspose x    = x
swapTranspose BlasConjTranspose (a, b) = (b, a)
{-# INLINE swapTranspose #-}

-- | Error text
mkDimText :: (Show a1, Show a2, Show a3, Show a4, Show a5, Show a6) => (a1, a2) -> (a3, a4) -> (a5, a6) -> String
mkDimText (ax, ay) (bx, by) (cx, cy) = "resulting dimensions: [" ++ show ax ++ "x" ++ show ay ++ "]*[" ++ show bx ++ "x" ++ show by ++ "]=[" ++ show cx ++ "x" ++ show cy ++ "]"

-- | Computes: C <- alpha*op( A )*op( B ) + beta*C, where op(X) may transpose the matrix X
--
-- dgemm, see http://www.netlib.org/lapack/explore-html/d1/d54/group__double__blas__level3_gaeda3cbd99c8fb834a60a6412878226e1.html for the documentation.
--
-- void cblas_dgemm (
--              const CBLAS_LAYOUT      layout,
--              const CBLAS_TRANSPOSE   TransA,
--              const CBLAS_TRANSPOSE   TransB,
--              const int       M,
--              const int       N,
--              const int       K,
--              const double    alpha,
--              const double *          A,
--              const int       lda,
--              const double *          B,
--              const int       ldb,
--              const double    beta,
--              double *        C,
--              const int       ldc
-- 	)
{-# NOINLINE dgemmUnsafe #-}
dgemmUnsafe :: BlasTranspose    -- ^ Transpose Matrix A
            -> BlasTranspose    -- ^ Transpose Matrix B
            -> (Int, Int)       -- ^ Rows and cols of A on entry (not transposed)
            -> (Int, Int)       -- ^ Rows and Cols of B on entry (not transposed)
            -> RealNum           -- ^ Alpha
            -> V.Vector RealNum  -- ^ A
            -> V.Vector RealNum  -- ^ B
            -> RealNum           -- ^ Beta
            -> V.Vector RealNum  -- ^ C
            -> IO (V.Vector RealNum)  -- ^ Return new C
dgemmUnsafe trA trB (axIn, ayIn) (bxIn, byIn) alpha matrixA matrixB beta matrixC
  | isBadGemm =
    error $!
    "bad dimension args to dgemmUnsafe: ax ay bx by cx cy: " ++
    show [ax, ay, bx, by, ax, by] ++ " matrix C length: " ++ show (V.length matrixC) ++ "\n\t" ++ mkDimText (ax, ay) (bx, by) (ax, by)
  | otherwise = do
      V.unsafeWith matrixA $ \aPtr' ->
        V.unsafeWith matrixB $ \bPtr' ->
          V.unsafeWith matrixC $ \cPtr' ->  do
#ifdef USE_FLOAT
            sgemm_direct
#else
            dgemm_direct
#endif
              (encodeTransposeIntBool trA) -- transpose A
              (encodeTransposeIntBool trB) -- transpose B
              (fromIntegral ax)     -- rows of C = rows of A transposed
              (fromIntegral by)     -- cols of C = cols of B transposed
              (fromIntegral ay)     -- k = cols of A transposed = rows of B transposed
              alpha
              aPtr'
              (fromIntegral axIn) -- LDA
              bPtr'
              (fromIntegral bxIn) -- LDB
              beta
              cPtr'
              (fromIntegral ax)   -- LDC
            return matrixC
  where
    (ax, ay) = swapTranspose trA (axIn, ayIn)
    (bx, by) = swapTranspose trB (bxIn, byIn)
    isBadGemm = minimum [ax, ay, bx, by] <= 0 || not (ax * by == V.length matrixC && ay == bx)

-- | Computes: Y <- alpha*op( A )*X + beta*Y, where op(A) may transpose the matrix A
--
-- dgemv, see http://www.netlib.org/lapack/explore-html/d7/d15/group__double__blas__level2_gadd421a107a488d524859b4a64c1901a9.html#gadd421a107a488d524859b4a64c1901a9
--
-- void cblas_dgemv 	(
--              const CBLAS_LAYOUT  	layout,
-- 		const CBLAS_TRANSPOSE  	TransA,
-- 		const int  	M,
-- 		const int  	N,
-- 		const double  	alpha,
-- 		const double *  	A,
-- 		const int  	lda,
-- 		const double *  	X,
-- 		const int  	incX,
-- 		const double  	beta,
-- 		double *  	Y,
-- 		const int  	incY
-- 	)
{-# NOINLINE dgemvUnsafe #-}
dgemvUnsafe :: BlasTranspose    -- ^ Transpose Matrix
            -> (Int, Int)       -- ^ rows and cols of A on entry (not transposed)
            -> RealNum           -- ^ Alpha
            -> V.Vector RealNum  -- ^ A
            -> V.Vector RealNum  -- ^ X
            -> RealNum           -- ^ Beta
            -> V.Vector RealNum  -- ^ C
            -> IO (V.Vector RealNum)  -- ^ Return new C
dgemvUnsafe trA (m, k) alpha matrixA vecX beta vecY
  | ax /= V.length vecY || ay /= V.length vecX =
    error $!
    "bad dimension args to dgemvUnsafe: ax ay (length vecX) (length vecY): " ++
    show [ax, ay, V.length vecX, V.length vecY] ++ " \n\t" ++ mkDimText (ax, ay) (V.length vecX, 1) (m, 1)
  | otherwise = do
      V.unsafeWith matrixA $ \aPtr' ->
        V.unsafeWith vecX $ \xPtr' ->
          V.unsafeWith vecY $ \yPtr' -> do
#ifdef USE_FLOAT
            sgemv_direct
#else
            dgemv_direct
#endif
              (encodeTransposeIntBool trA)      -- transpose A
              (fromIntegral m)
              (fromIntegral k)
              alpha
              aPtr'
              (fromIntegral m)
              xPtr'
              1
              beta
              yPtr'
              1
            return vecY
  where
    (ax, ay) = swapTranspose trA (m, k)

-- | Computes: A <- alpha*X*Y^T + A
--
-- dger, see http://www.netlib.org/lapack/explore-html/d7/d15/group__double__blas__level2_ga458222e01b4d348e9b52b9343d52f828.html#ga458222e01b4d348e9b52b9343d52f828
--
-- void cblas_dger 	(
--              const CBLAS_LAYOUT  	layout,
-- 		const int  	M,
-- 		const int  	N,
-- 		const double  	alpha,
-- 		const double *  	X,
-- 		const int  	incX,
-- 		const double *  	Y,
-- 		const int  	incY,
-- 		double *  	A,
-- 		const int  	lda
-- 	)
{-# NOINLINE dgerUnsafe #-}
dgerUnsafe :: (Int, Int)           -- ^ Dimensions of matrix A
           -> RealNum               -- ^ Alpha
           -> V.Vector RealNum      -- ^ X
           -> V.Vector RealNum      -- ^ C
           -> V.Vector RealNum      -- ^ A
           -> IO (V.Vector RealNum)  -- ^ Return new C
dgerUnsafe (ax, ay) alpha vecX vecY matrixA
  | ax /= len || ay /= V.length vecY =
    error $! "bad dimension args to dgerUnsafe: X Y ax ay: " ++ show [len, V.length vecY, ax, ay] ++ " \n\t" ++ mkDimText (ax, ay) (V.length vecX, 1) (V.length vecY, 1)
  | otherwise = do
      V.unsafeWith matrixA $ \aPtr' ->
        V.unsafeWith vecX $ \xPtr' ->
          V.unsafeWith vecY $ \yPtr' -> do
#ifdef USE_FLOAT
            sger_direct
#else
            dger_direct
#endif
              (fromIntegral ax)
              (fromIntegral ay)
              alpha
              xPtr'
              1
              yPtr'
              1
              aPtr'
              (fromIntegral ax)
            return matrixA
  where
    len = V.length vecX

-- |  Matrix mult for general dense matrices
type BLASGemmFunFFI scale el
  =  Int -- transpose A: 1, not transpose A: 0
  -> Int -- transpose B: 1, not transpose B: 0
  -> CInt -- m
  -> CInt -- n
  -> CInt -- k
  -> {- scal A * B -} scale  -- alpha
  -> {- Matrix A-} Ptr el    -- A
  -> CInt                    -- LDA
  -> {- B -} Ptr el
  -> CInt
  -> scale                   -- beta
  -> {- C -}  Ptr el
  -> CInt
  -> IO ()

foreign import ccall unsafe "dgemm_direct" dgemm_direct :: BLASGemmFunFFI Double Double
foreign import ccall unsafe "sgemm_direct" sgemm_direct :: BLASGemmFunFFI Float Float

-- |  Matrix mult for general dense matrices
type BLASGemvFunFFI scale el
  =  Int    -- transpose A: 1, not transpose A: 0
  -> CInt   -- m
  -> CInt   -- n
  -> scale  -- alpha
  -> Ptr el -- Matrix A
  -> CInt   -- LDA
  -> Ptr el
  -> CInt
  -> scale -- beta
  -> Ptr el
  -> CInt
  -> IO ()

foreign import ccall unsafe "dgemv_direct" dgemv_direct :: BLASGemvFunFFI Double Double
foreign import ccall unsafe "sgemv_direct" sgemv_direct :: BLASGemvFunFFI Float Float

type BlasGerxFunFFI scale el
  =  CInt
  -> CInt
  -> scale
  -> Ptr el
  -> CInt
  -> Ptr el
  -> CInt
  -> Ptr el
  -> CInt
  -> IO ()

foreign import ccall unsafe "dger_direct" dger_direct :: BlasGerxFunFFI Double Double
foreign import ccall unsafe "sger_direct" sger_direct :: BlasGerxFunFFI Float Float

-- toRows :: Int -> V.Vector Double -> [V.Vector Double]
-- toRows m vec = LA.toRows . reshapeF m . LA.vector . V.toList $ vec
--   where reshapeF r = LA.tr' . LA.reshape r

-- vec1 :: LAS.R 10
-- vec1 = LAS.vector [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.34868882831497655, 0.0, 1.4026932193043212e-2]

-- -- dEdy:
-- -- mmCheck : [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-4.8711473950841355e-2,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-1.9595481318788644e-3]
-- -- mm' : [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.24602759633121374,-1.2277801012906878e-2,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-9.89711206989971e-3,-4.939070836308897e-4,0.0,0.0]

-- vec3 :: V.Vector Double
-- vec3 = (
--   V.fromList [0.0, 0.0, 0.0, 0.0, (-0.13969898085418284)])

-- vec2 :: LAS.R 5
-- vec2 = LAS.vector [0.0, 0.0, 0.0, 0.0, -0.13969898085418284]

-- res = vec1 `LAS.outer` vec2

-- test =
--   toRows 10 $
--   outerV (LAS.extract vec1) (LAS.extract vec2) (V.replicate (LAS.size vec1 * LAS.size vec2) 10)


{-# LANGUAGE RankNTypes, BangPatterns, ScopedTypeVariables, FlexibleInstances #-}

module Math.Probably.MALA where

import qualified Math.Probably.PDF as PDF
import Math.Probably.Sampler
import qualified Math.Probably.PDF as PDF
import Control.Applicative

import Numeric.LinearAlgebra
import Text.Printf
import System.IO
import Data.Maybe

import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import qualified Data.Vector.Unboxed as U
import qualified Control.Monad.State.Strict as S
import qualified Data.Vector.Storable as VS

import Debug.Trace
import Data.IORef
import Control.Spoon

data MalaPar = MalaPar { mpXi :: !(Vector Double),
                         mpPi :: !Double,
                         mpGradLast :: !(Vector Double),
                         mpSigma :: !Double,
                         mpCount :: !Double,
                         mpAccept :: !Double,
                         mpFreezeSigma :: !Bool,
                         mpThin :: !Int,
                         mpMaxGradLen :: !Double } --,
--                         mpLastRatio :: !Double }
  deriving Show

trunc t v = scale (t/(max t (norm2 v))) v

mala1 :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         -> Bool
         -> MalaPar
         -> Sampler MalaPar
mala1 cov postgrad useCache
      (MalaPar xi piCached gradientiCached sigma tr tracc freeze thinn maxGrad) = do
  let (pi,gradienti) = if useCache
                          then (piCached, gradientiCached)
                          else let (p, gr_untr) = postgrad xi
                               in (p, trunc maxGrad gr_untr)
  let xstarMean = xi + scale (sigma/2) (cov `covMul` gradienti)
  xstar <- mvnSampler xstarMean sigma cov
  u <- unitSample
  let (!pstar, gradientStarUntrunc) = postgrad xstar
      gradientStar = trunc maxGrad gradientStarUntrunc
  let !revJumpMean = xstar + scale (sigma/2) (cov `covMul` gradientStar)
      ptop = mvnPDF revJumpMean sigma cov xi
      pbot = mvnPDF xstarMean sigma cov xstar
      ratio = exp $   pstar -pi + ptop - pbot
      tr' = max 1 tr
      sigmaNext = case () of
         _ | freeze -> sigma
      accept = tracc / tr
      freezeNext = freeze {-| freeze = True
                 | not freeze = tr > 100 && accept > 0.5 && accept < 0.6  -}
  if trace (show $ (tr, pstar ,pi , ratio, sigma)) $ u < ratio
     then return $ MalaPar xstar pstar (gradientStar)
                           (if freezeNext then sigma else (min 1.4 $ 1+kmala/tr')*sigma)
                           (tr+1) (tracc+1) freezeNext  thinn maxGrad
--                           sigma (tr+1) (tracc+1)
     else return $ MalaPar xi pi ( gradienti)
                           (if freezeNext then sigma else (max 0.7143 $ 1-kmala/tr')**1.3*sigma)
                           (tr+1) tracc freezeNext thinn maxGrad
--                           sigma (tr+1) tracc

runMalaMP :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> [(Double,Vector Double)] -> RIO (MalaPar, [(Double,Vector Double)])
runMalaMP cov  pdf nsam init xs0 = go nsam init xs0 where
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return (mpar, xs)
  go n y xs = do y1 <- sample $ mala1 cov pdf True y
--                 io $ do putStrLn $ show (mpCount y1, mpPi y1, mpSigma y1)
--                         hFlush stdout
                 let newChainRes = if mpThin y1 == 0 || round (mpCount y1) `mod` mpThin y1 ==0
                                      then let !xi = mpXi y1
                                               !pi = mpPi y1
                                               !more = (pi,xi)
                                           in more:xs
                                      else xs
                 go (n-1) y1 newChainRes

runMalaUntilBetter :: Covariance a => a ->(Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> RIO (MalaPar, [(Double,Vector Double)])
runMalaUntilBetter cov  pdf nsam init = go nsam init [] where
  pTarget = mpPi init
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return (mpar, xs)
  go n y xs = do y1 <- sample $ mala1 cov  pdf True y
                 io $ do putStrLn $ show (mpCount y1, mpPi y1, mpSigma y1)
                         hFlush stdout
                 let newChainRes = if mpThin y1 == 0 || round (mpCount y1) `mod` mpThin y1 ==0
                                      then (mpPi y1, mpXi y1):xs
                                      else xs
                 if mpPi y1 > pTarget
                    then return (y1, newChainRes)
                    else go (n-1) y1 $ newChainRes

{-runMalaMPaccept' :: Matrix Double -> (Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> RIO (MalaPar, [(Double,Vector Double)])
runMalaMPaccept'  cov pdf nsam init  = do
  stseed <- S.get
  res <- io $ do
   seedref <- newIORef stseed
   resRef <- newIORef Nil
   let go !mpar
        | (round $ mpAccept mpar) >= nsam
          = return mpar
        | otherwise
          = do seed <- readIORef seedref
               putStr ((show (round $ mpAccept mpar)) ++".") >> hFlush stdout
               let (!mpar1, !seed1) = unSam (mala1 cov pdf mpar) seed
               writeIORef seedref seed1
               modifyIORef resRef (Cons (Pair (mpPi mpar1) ( mpXi mpar1)))
               go mpar1
   mpar <- go init
   seed <- readIORef seedref
   reslist <- readIORef resRef
   return (seed, (mpar,reslist))
  S.put $ fst res
  let mp = fst $ snd res
  return $ (mp, map unPair $ unList $ snd $ snd res)

--  pi = pdf $ toList $ mpXi init -}

runMalaMPaccept :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         -> Int -> MalaPar ->  RIO (MalaPar, [(Double,Vector Double)])
runMalaMPaccept cov pdf nsam init = go init [] where
  go !mpar !xs
    | (round $ mpAccept mpar) >= nsam = do
         io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
         io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
         return (mpar, xs)
    | mpAccept mpar < 0.5 && mpCount mpar > 20 =
         return (mpar, [])
    | otherwise = do
         !mpar1 <- sample $ mala1 cov  pdf True mpar
         io $ putStr ((show (round $ mpAccept mpar)) ++".") >> hFlush stdout
         go mpar1 $ (mpPi mpar1, mpXi mpar1):xs

kmala = 5

{-runMalaRioESS ::  Matrix Double -> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Vector Double -> RIO [Vector Double]
runMalaRioESS cov pdf want_ess xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 grad0 sigma0 0 0 False
        nsam0 = want_ess*20
    io $ putStrLn $ "initial sigma = "++show (mpSigma mp0)
    (mp1, xs1) <- runMalaMP cov pdf nsam0 mp0 []
    let have_ess = min (mpAccept mp1) $ calcESSprim $ map snd xs1
    if have_ess > realToFrac want_ess
       then return $ map snd xs1
       else do let need_ess =  max 1 $ realToFrac want_ess - have_ess
                   samples_per_es = realToFrac nsam0/have_ess
                   to_do = round $ samples_per_es * need_ess
               (mp2, xs2) <- runMalaMP cov pdf to_do mp1 xs1
               return  $ map snd  xs2 -}

runMalaRioCodaESS ::  Covariance a => a-> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Double -> Vector Double -> RIO [Vector Double]
runMalaRioCodaESS cov pdf want_ess truncN xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 (trunc truncN grad0) sigma0 0 0 False 0 truncN
        nsam0 = want_ess*1
    let converged mp  xs = do
         let have_ess = min (mpAccept mp) $ calcESSprim $ map snd xs
         io $ putStrLn $ "ESS=" ++show have_ess
         if have_ess > realToFrac want_ess
            then return $ map snd xs
            else do let need_ess =  max 1 $ realToFrac want_ess - have_ess
                        samples_per_es = realToFrac nsam0/have_ess
                        to_do = round $ samples_per_es * need_ess
                    io $ putStrLn $ "running converged for "++show to_do
                    (mp2, xs2) <- runMalaMP cov pdf to_do (mp {mpFreezeSigma = True}) xs
                    io $ putStrLn $ "All done"
                    return $ map snd xs2
    let go mp  n xs = do
            (mp2, xs2) <- runMalaMP cov  pdf n mp []
            let testres = mannWhitneyUtest TwoTailed 0.05 (U.fromList $ map fst xs2)
                                            (U.fromList $ map fst xs)

            if testres/= Just NotSignificant
               then do io$ putStrLn $ "not converged: "++show testres++" at "++show (mpPi mp2)
                       go mp2  (round $ realToFrac n*2) xs2
               else do io$ putStrLn "converged!"
                       converged mp2  (xs2++xs)

    io $ putStrLn $ "initial sigma = "++show (mpSigma mp0)
    (mp1, xs1) <- runMalaMPaccept cov  pdf nsam0 mp0
    case xs1 of
            [] -> return []
            _ -> go mp1  (round $ mpCount mp1) xs1

{-    case (spoon (invlndet cov), mbCholSH cov) of
       (Just (covInv, (lndt,_)), Just covChol) -> go_rest (covInv, covChol)
       _ -> let cov' = PDF.posdefify cov in
            case (spoon (invlndet cov'), mbCholSH cov') of
              (Just (covInv, (lndt,_)), Just covChol) -> go_rest (covInv, covChol)
              _ ->  do io $ putStrLn "non-invertible covariance matrix"
                       return [] -}

runMalaRioSimple ::  Covariance a => a -> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Double -> Int -> Vector Double -> RIO [Vector Double]
runMalaRioSimple cov pdf samples maxGrad thinN xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 2.7e-4 --1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 (trunc maxGrad grad0) sigma0 0 0 False thinN maxGrad

    (mp2, xs2) <- runMalaMP cov  pdf samples mp0 []
    io $ putStrLn $ "All done"
    return $ map snd xs2

calcCovariance :: Vector Double ->
                  Vector Double ->
                  (Vector Double -> (Double,Vector Double)) ->
                  (Vector Double -> Double) ->
                  Either (Vector Double) (Matrix Double, Matrix Double, Matrix Double)
calcCovariance vinit vnear postgrad posterior = finalcov where
   ndim = dim vinit
   finalcov
     | ndim > 0 -- > 20000 --FIXME
        = Left $ calcFDindepVars vinit vnear posterior --Left $ iCov vinit --
     | otherwise
        = hessToCov (calcFDhess vinit vnear postgrad) Nothing

iCov v = VS.replicate (VS.length v) 1 -- in (m,m,m)

calcFDindepVars v v' post = trace ("FDVars = "++show (VS.take 10 vars))  $ vars where
   hv =  mapVector (*1e-4) v --mapVector (max 1e-9 .  abs) $ v - v'
   n = dim v
   postv = post v
   postPlus i = post $ v VS.// [(i,v @>i + hv @> i)]
   postMinus i = post $ v VS.// [(i,v @>i - hv@> i)]
   vars = buildVector n fvar
   fvar i = negate $ recip $ (postPlus i - 2*postv + postMinus i)/((hv @> i)*(hv @> i))

calcFDhess v v' postgrad = hess2 where
   grad =  snd . postgrad
   gradi i = (@>i) . grad
   hv = mapVector (max 1e-9 . abs) $ v - v'
   n = dim v
   gradv = grad v
   grads =  fromRows $ flip map [0..(n-1)] $ \i ->
               grad (v VS.// [(i,v @>i + hv @> i)])
   fhess (i,j) | i<j = 0
               | otherwise = (grads @@>(j,i) - (gradv @> i))
                                       /(2*(hv @> j)) +
                             (grads @@>(i,j)- (gradv @> j))
                                       /(2*(hv @> i))

--   hess3 = scale (recip $ realToFrac n) $ sum $ map outerSelf $ grads

   hess1 = buildMatrix n n fhess
   hess2 = buildMatrix n n $ \(i,j) -> if i>=j then hess1 @@> (i,j)
                                               else hess1 @@> (j,i)

outerSelf v= v `outer` v

hessToCov hess mOriginal =
   let mTryPosdefify = case mOriginal of
            Nothing -> hessToCov (PDF.posdefify hess) (Just hess)
            Just horiginal -> let ds =  mapVector (negate . recip)
                                         $ takeDiag horiginal
                              in trace ("USING DIAGS"++show (VS.take 10 ds)) $ Left $ ds
   in
   case spoon $ inv $ negate $ hess of
     Just cov -> case mbCholSH cov of
                   Just cholm ->  trace ("invert success:"++show (VS.take 10 $ takeDiag cov) ++ "det="++show (det cov)) $ Right (cov, negate hess, cholm)
                   Nothing -> trace ("chol fail") mTryPosdefify
     Nothing -> trace ("inv fail") mTryPosdefify
                  {- Just cov -> case mbCholSH cov of
                                Just cholm ->  Right (cov, negate hessToCov,cholm)
                                Nothing -> case mbCholSH $ PDF.posdefify cov of
                                            Just cholm -> Right (cov, negate hess,cholm)
                                            Nothing -> vars
                  Nothing -> vars -}

acceptSM ampar  | mpCount ampar == 0 = "0/0"
               | otherwise = printf "%.3g" (rate::Double) ++ " ("++show yes++"/"++show total++")" where
   rate = realToFrac (yes) / realToFrac (total)
   yes = mpAccept ampar
   total = mpCount ampar


module Tools
( randomMatrix
, makeBatchList
, encorderForDat2
, encorderForDat
, length'
, pickUpNum
, setElem
, oneHotEncoder
, oneHotEncoder'
) where

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

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

import Durstenfeld

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

-------------------------------------------------------------
--Tools
-------------------------------------------------------------

--重みの初期化用
--要素が0。0から1.0の間の数で、指定した行数、列数のMatrixを生成する。
randomMatrix :: RandomGen g => Int -> Int -> g -> Matrix Double
randomMatrix row col g = matrix col $ take (row*col) $ randomRs (0.0, 0.01) g

--ex) randomMatrix 3 4 (mkStdGen 100)

makeBatchList :: (Eq t, Num t) => Int -> t -> Int -> Int -> [[Int]]
makeBatchList _ 0 _ _ = []
makeBatchList g num ds bs = batch' : makeBatchList (g - 1) (num - 1) ds bs
    where data_size = ds
          batch_size = bs
          batch' = randomChoice data_size batch_size (mkStdGen g )

--数字のリストのリストをファイルに書き込むために、良い感じにいじっている。
encorderForDat :: Show a => [[a]] -> [Char]
encorderForDat xs = intercalate "\n" $  fmap (intercalate " " . fmap show) xs

--リストを書き込む用
encorderForDat2 :: Show a => [a] -> [Char]
encorderForDat2 xs = intercalate "\n" $  fmap show xs

length' :: [a] -> Double
length' [] = 0.0
length' (x:xs) = 1.0 + length' xs

pickUpNum :: [Int] -> [a] -> [a]
pickUpNum [] _ = []
pickUpNum (x:xs) list = number : pickUpNum xs list
    where number = list !! x

setElem1 :: Matrix R -> (Int, Int) -> Int -> Int -> R -> Matrix R
setElem1 m (row, col) maxrow maxcol newElem = newM
    where v  = flatten m
          num = row * maxcol + (col+1) --要素が何番目か
          v1 = subVector 0 (num-1) v
          v2 = subVector (num) (maxcol*maxrow-num) v
          v3 = vector [newElem]
          newV = vjoin [v1,v3,v2]
          newM = reshape maxcol newV

--〇
setElem :: Matrix R -> (Int, Int) -> R -> Matrix R
setElem m (row, col) newElem = setElem1 m (row, col) maxrow maxcol newElem
    where maxrow = rows m
          maxcol = cols m

--ラベル(0~9)をOne-Hot表現のMatrix(1行10列)に変換
--ex) 9 --> 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0
oneHotEncoder :: Int -> Matrix Double
oneHotEncoder num = setElem zeroM (0 ,num) 1.0
    where zeroM = row [0.0,0,0,0,0,0,0,0,0,0]

-- =============================================

--setElem1' :: Matrix R -> (Int, Int) -> Int -> Int -> R -> Matrix R
setElem' v col newElem = newV
    where v1 = subVector 0 (col-1) v
          v2 = subVector col (10-col) v
          v3 = vector [newElem]
          newV = vjoin [v1,v3,v2]

oneHotEncoder' :: Int -> Vector R
oneHotEncoder' num = setElem' zeroV (num+1) 1.0
    where zeroV = fromList [0.0,0,0,0,0,0,0,0,0,0]

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


{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Sam where

import Statistics.Distribution
import Statistics.Distribution.Normal (normalDistr)
import Statistics.Distribution.Beta (betaDistr)
import Statistics.Distribution.Gamma (gammaDistr)
import qualified Statistics.Distribution.Poisson as Poisson
import Data.List
import Data.Map (empty,lookup,insert,size,keys)
import Data.IORef
import Control.Monad
import Control.Monad.Extra
import System.IO.Unsafe
import Control.Monad.State.Lazy (State, state , put, get, runState, MonadTrans (lift))
import Debug.Trace
import Control.Monad.Trans.Writer
import Numeric.Log
import System.Random hiding (uniform)
import Data.Monoid (Product (Product))

{--
    This file defines
    * Two monads: Prob (for probabilities) and Meas (for unnormalized probabilities)
    * Two inference methods: lwis and mh
--}

-- A "Tree" is a lazy, infinitely wide and infinitely deep tree, labelled by Doubles
-- Our source of randomness will be a Tree, populated by uniform [0,1] choices for each label.
-- Often people would just use a list or stream instead of a tree.
-- But a tree allows us to be lazy about how far we are going all the time.
data Tree = Tree Double [Tree]

-- A probability distribution over a is a state transformer over trees
-- ie a function Tree -> (a , Tree)
-- The idea is that it uses up bits of the tree as it runs
newtype Prob a = Prob (State Tree a)

-- Two key things to do with trees:
-- Split tree splits a tree in two (bijectively)
-- Get the label at the head of the tree and discard the rest
splitTree :: Tree -> (Tree , Tree)
splitTree (Tree r (t : ts)) = (t , Tree r ts)

uniform :: Prob Double
uniform = Prob $
      do ~(Tree r (t:ts)) <- get
         put t
         return r

-- Probabilities for a monad.
-- Sequencing is done by splitting the tree
-- and using different bits for different computations.
instance Monad Prob where
  return a = Prob $ return a
  (Prob m) >>= f = Prob $
                        do g <- get
                           let (g1,g2) = splitTree g
                           put g1
                           x <- m
                           put g2
                           let (Prob m') = f x
                           m'
instance Functor Prob where fmap = liftM
instance Applicative Prob where {pure = return ; (<*>) = ap}

{-- An unnormalized measure is represented by a probability distribution over pairs of a weight and a result --}
newtype Meas a = Meas (WriterT (Product (Log Double)) Prob a)
  deriving(Functor, Applicative, Monad)

{-- The two key methods for Meas are sample (from a probability) and score (aka factor, weight) --}
score :: Double -> Meas ()
score r = Meas $ tell $ Product $ (Exp . log) $ (if r==0 then exp(-300) else r)

scorelog :: Double -> Meas ()
scorelog r = Meas $ tell $ Product $ Exp $ (if r==0 then exp(-300) else r)

sample :: Prob a -> Meas a
sample p = Meas $ lift p

{-- Preliminaries for the simulation methods. Generate a tree with uniform random labels
    This uses SPLIT to split a random seed --}
randomTree :: RandomGen g => g -> Tree
randomTree g = let (a,g') = random g in Tree a (randomTrees g')
randomTrees :: RandomGen g => g -> [Tree]
randomTrees g = let (g1,g2) = split g in (randomTree g1) : (randomTrees g2)

{-- Run prob runs a probability deterministically, given a source of randomness --}
runProb :: Prob a -> Tree -> a
runProb (Prob a) rs = fst $ runState a rs

normal :: Double -> Double -> Prob Double
normal m s = do quantile (normalDistr m s) <$> uniform

normalPdf :: Double -> Double -> Double -> Double
normalPdf m s = density $ normalDistr m s

exponential :: Double -> Prob Double
exponential rate =
  do x <- uniform
     return $ - (log x / rate)

expPdf :: Double -> Double -> Double
expPdf rate x = exp (-rate*x) * rate

gamma :: Double -> Double -> Prob Double
gamma a b = do
  quantile (gammaDistr a b) <$> uniform

beta :: Double -> Double -> Prob Double
beta a b = do
  quantile (betaDistr a b) <$> uniform

poisson :: Double -> Prob Integer
poisson lambda = do
  x <- uniform
  let cmf = scanl1 (+) $ map (probability $ Poisson.poisson lambda) [0,1..]
  let (Just n) = findIndex (> x) cmf
  return $ fromIntegral n

poissonPdf :: Double -> Integer -> Double
poissonPdf rate n = probability (Poisson.poisson rate) (fromIntegral n)

dirichlet :: [Double] -> Prob[Double]
dirichlet as = do
  xs <- mapM exponential as
  let s = Prelude.sum xs
  let ys = map (/ s) xs
  return ys

uniformbounded :: Double -> Double -> Prob Double
uniformbounded lower upper = do
  x <- uniform
  return $ (upper - lower) * x + lower

bernoulli :: Double -> Prob Bool
bernoulli r = do
  x <- uniform
  return $ x < r

{-
 uniform distribution on [0, ..., n-1]
-}
uniformdiscrete :: Int -> Prob Int
uniformdiscrete n =
  do
    let upper = fromIntegral n
    r <- uniformbounded 0 upper
    return $ floor r

{-- Categorical distribution: takes a list of k numbers that sum to 1,
    and returns a number between 0 and (k-1) --}
categorical :: [Double] -> Prob Int
categorical xs = do
  r <- uniform
  let (Just i) = findIndex (>r) $ tail $ scanl (+) 0 xs
  return i

{-- Stochastic memoization.
    We use unsafePerformIO to maintain
    a table of calls that have already been made.
    If a is finite, we could just sample all values of a in advance
    and avoid unsafePerformIO.
    If it is countably infinite, there probably also are implementation tricks.
--}
memoize :: Ord a => (a -> Prob b) -> Prob (a -> b)
memoize f =  Prob $ do g <- get
                       let ( Tree _ gs, g2) = splitTree g
                       put g2
                       return $ unsafePerformIO $ do
                                ref <- newIORef Data.Map.empty
                                return $ \x -> unsafePerformIO $ do
                                          m <- fmap (Data.Map.lookup x) (readIORef ref)
                                          case m of
                                              Just y -> return y
                                              Nothing -> do
                                                            let (Prob k) = f x
                                                            n <- readIORef ref
                                                            let (y,_) = runState k (gs !! (1 + size n))
                                                            modifyIORef' ref (Data.Map.insert x y)
                                                            return y

{-- Stochastic memoization for recursive functions.
    Applying 'memoize' to a recursively defined function only memoizes at the
    top-level: recursive calls are calls to the non-memoized function.
    'memrec' is an alternative implementation which resolves recursion and
    memoization at the same time, so that recursive calls are also memoized.
--}
memrec :: Ord a => Show a => ((a -> b) -> (a -> Prob b)) -> Prob (a -> b)
memrec f =
   Prob $ do
    g <- get
    let ( Tree _ gs, g2) = splitTree g
    put g2
    return $ unsafePerformIO $ do
                  ref <- newIORef Data.Map.empty
                  let memoized_fixpoint = \x -> unsafePerformIO $ do
                                m <- fmap (Data.Map.lookup x) (readIORef ref)
                                case m of
                                      Just y -> return y
                                      Nothing -> do
                                                  n <- readIORef ref
                                                  let fix = f memoized_fixpoint
                                                  let Prob k = fix x
                                                  let (y, _) = runState k (gs !! (1 + size n))
                                                  modifyIORef' ref (Data.Map.insert x y)
                                                  return y
                  return memoized_fixpoint


-- program: S.A.R.A.H. jam simulator
-- author: Maurício C. Antunes
-- e-mail: mauricio.antunes@gmail.com
-- license: public domain

module Main where

import Control.Applicative
import Prelude

import Data.Maybe
import Graphics.UI.Gtk
import Graphics.Rendering.Cairo
import Control.Monad
import Data.IORef
import Data.List
import Data.Time
import Data.Complex

-- Constants

accelerator = 0.7*carSize :: Double
brake = 10*accelerator:: Double
carSize = 2*pi/59 :: Double
responseTime = 0.24 :: Double
drawSide = 5/2 :: Double

-- A few conveniences

eventWindowSize = do
    dr <- eventWindow
    w <- liftIO $ drawWindowGetWidth dr
    h <- liftIO $ drawWindowGetHeight dr
    return $ if w*h > 1
        then (fromIntegral w, fromIntegral h)
        else (1,1)

eventPolarCoordinates = do
    (w,h) <- eventWindowSize
    (x,y) <- eventCoordinates
    let (origX, origY) = (w/2, h/2)
    let (scaleX, scaleY) = (drawSide/w, drawSide/h)
    let (x',y') = (scaleX*(x-origX), scaleY*(y-origY))
    let (radius,theta) = polar $ x' :+ y'
    return $ (radius,theta)

getAndSet :: a -> IO (IO a, a -> IO ())
getAndSet a = do
    ior <- newIORef a
    let get = readIORef ior
    let set = writeIORef ior
    return (get,set)

diffTime :: UTCTime -> UTCTime -> Double
diffTime = (realToFrac .) . diffUTCTime

moveToLineTo :: Double -> Double
 -> Double -> Double -> Render ()
moveToLineTo a b c d = moveTo a b >> lineTo c d

-- Car list handling

-- Each car is represented by a pair of Doubles. The first
-- Double is its position in a circular road, represented by
-- an angle. The second is its angular velocity. The general
-- idea behind the simulation is that in a list of cars each
-- one will try to keep a safe speed to avoid a crash in the
-- event of a sudden brake of the next car.

newCarList nCars = take nCars $ zip [0,2*pi/nCars'..] (repeat 0)
    where nCars' = fromIntegral nCars

-- This resizes car lists by copying or keeping those
-- at lower speeds.

newCarListFromList nCars [] = newCarListFromList nCars [(0,0)]
newCarListFromList nCars list = sortBy ((. fst).(compare . fst)) $
    take nCars $ cycle $ sortBy ((. snd).(compare . snd)) list

-- Safe speed for car, given data from itself and the next
-- and, possibly, a forced (by the user) jam. Speed changes
-- are limited by accelerator and brake maxima.

newSpeed dt jam (p1,s1) (p2,s2) = min cv $ max bv $ ds - br
    where
        pd = (p2-p1-carSize) - responseTime*(s2-s1)
        pj = maybe pd ((subtract $ carSize/2)
         . (until (>0) (+2*pi)) . (subtract p1)) jam
        dd = brake*(max 0 $ min pd pj)
        br = brake*responseTime
        ds = sqrt $ br^2 + 2*dd
        cv = s1 + accelerator*dt
        bv = s1 - brake*dt

-- Update positions and speeds based on a timestep and maybe
-- taking a forced congestion into account

updateCarList _ _ [] = []
updateCarList timestep jam list = zip newPositions' newSpeeds
    where
        fakeCar = (p+2*pi,s) where (p,s) = head list
        newSpeeds = zipWith ns list (tail list ++ [fakeCar])
            where ns = newSpeed timestep jam
        newPositions = zipWith3 mean fsts snds newSpeeds
            where
                mean a b c = a + timestep*(b+c)/2
                fsts = map fst list
                snds = map snd list
        newPositions' = map (subtract base) newPositions
        base = (*(2*pi)) $ fromIntegral $ floor $ (/ (2*pi)) $
            head newPositions

about = do
    ad <- aboutDialogNew
    set ad [ aboutDialogName := "S.A.R.A.H."
           , aboutDialogVersion := "1.0"
           , aboutDialogAuthors := ["Maurício C. Antunes "
                                ++ "<mauricio.antunes@gmail.com>"]
           , aboutDialogComments := "Software Automation of "
                                ++ "Road Automobile Headache"]
    dialogRun ad
    widgetDestroy ad

main :: IO ()
main = do

    initGUI

    mainWindow <- windowNew
    drawingArea <- drawingAreaNew

    (getTimeStamp,setTimeStamp) <- getCurrentTime >>= getAndSet
    (getCars,setCars) <- getAndSet $ newCarList 20
    (getJam,setJam) <- getAndSet Nothing
    (getTimeoutId,setTimeoutId) <- getAndSet Nothing

    -- If 'resume' is called, 'step' will be called at small
    -- timesteps to update car data. If 'pause' is called, 'step'
    -- calls are stoped.  'resume' is called at program startup,
    -- and then the pause button alternates 'resume' and 'pause'.

    let step = do
         time <- getCurrentTime
         dt <- getTimeStamp >>= return . (diffTime time)
         setTimeStamp time
         liftM2 (updateCarList dt) getJam getCars >>= setCars
    let pause = do
         maybe (return ()) timeoutRemove =<< getTimeoutId
         setTimeoutId Nothing
    let resume = do
         setTimeoutId . Just =<< flip timeoutAdd 33
          (step >> widgetQueueDraw drawingArea >> return True)
         getCurrentTime >>= setTimeStamp

    -- The elements of the graphic interface are the set of
    -- buttons, the scale to set the number of cars and the
    -- car track. They are named as 'buttons', 'howMany' and
    -- 'track'. Each of them contains other widgets inside, but
    -- there's no reason to expose their names to the main IO.

    buttons <- do

        qr <- buttonNewFromStock stockClear
        on qr buttonActivated $ do
            (liftM length) getCars >>= setCars . newCarList
            getCurrentTime >>= setTimeStamp
            widgetQueueDraw drawingArea

        qp <- toggleButtonNewWithLabel stockMediaPause
        buttonSetUseStock qp True
        on qp toggled $ do
            p <- toggleButtonGetActive qp
            case p of
                True -> pause
                False -> resume

        qa <- buttonNewFromStock stockAbout
        on qa buttonActivated $ about

        qq <- buttonNewFromStock stockQuit
        on qq buttonActivated (do
                       widgetDestroy mainWindow
                       mainQuit)

        bb <- hButtonBoxNew
        containerAdd bb qr
        containerAdd bb qp
        containerAdd bb qa
        containerAdd bb qq
        return bb

    howMany <- do

        sc <- vScaleNewWithRange 1 40 1
        after sc valueChanged $ do
            v <- liftM floor $ rangeGetValue sc
            c <- getCars
            setCars $ newCarListFromList v c
            widgetQueueDraw drawingArea

        scaleSetValuePos sc PosTop
        scaleSetDigits sc 0
--        rangeSetUpdatePolicy sc UpdateDiscontinuous
        rangeSetValue sc =<< liftM (fromIntegral . length) getCars

        al <- alignmentNew 0.5 0.5 0 1
        alignmentSetPadding al 15 15 15 15
        containerAdd al sc
        return al

    track <- do

        let dr = drawingArea
        widgetAddEvents dr [PointerMotionMask]

        on dr motionNotifyEvent $ do
            (r,t) <- eventPolarCoordinates
            liftIO $ if (0.8<r && r<1.2)
                then setJam (Just t)
                else setJam Nothing
            liftIO $ widgetQueueDraw dr
            return True

        on dr leaveNotifyEvent $ liftIO $
            setJam Nothing >> return True

        on dr draw $ do
            w <- liftIO $ (fromIntegral <$> widgetGetAllocatedWidth dr)
            h <- liftIO $ (fromIntegral <$> widgetGetAllocatedHeight dr)
            jam <- liftIO getJam
            cars <- liftIO getCars
            translate (w/2) (h/2)
            scale (w/drawSide) (h/drawSide)
            road2render jam cars
            -- return True

        af <- aspectFrameNew 0.5 0.5 (Just 1)
        frameSetShadowType af ShadowNone
        containerAdd af dr
        return af

    -- 'layout' is a widget that contains all interface elements
    -- properly arranged.

    layout <- do
        vb <- vBoxNew False 0
        hb <- hBoxNew False 0
        boxPackStart vb track PackGrow 0
        boxPackStart vb buttons PackNatural 0
        boxPackStart hb howMany PackNatural 0
        boxPackStart hb vb PackGrow 0
        return hb

    set mainWindow [ windowTitle := "S.A.R.A.H."
                   , windowDefaultWidth := 400
                   , windowDefaultHeight := 400 ]
    on mainWindow objectDestroy mainQuit
    containerAdd mainWindow layout
    widgetShowAll mainWindow

    resume

    mainGUI

-- As the name says, this takes road info, in the form of a
-- possible jam and a list of cars, and make it into a Cairo
-- render.  Road will have radius 1.

road2render :: Maybe Double -> [(Double,Double)] -> Render ()
road2render jam cars = do
    newPath
    setSourceRGB 0 0 0
    drawRoad
    when (isJust jam) drawJam
    setSourceRGBA 0 0 0 0.55
    let cars' = map fst cars
    let rotations = zipWith subtract (0:cars') cars'
    sequence_ $ map ((>> drawCar) . rotate) rotations
 where
    drawRoad = setLineWidth 0.01 >> setDash [2*pi/34,2*pi/34]
     (pi/34) >> arc 0.0 0.0 1.0 0.0 (2*pi) >> stroke
    drawJam = setLineWidth 0.005 >> setDash [0.03,0.02] 0.04 >>
     save >> rotate (fromJust jam) >> moveToLineTo 0.8 0 1.2
     0 >> stroke >> setDash [] 0 >> moveToLineTo 0.8 (-0.015)
     0.8 0.015 >> moveToLineTo 1.2 (-0.015) 1.2 0.015 >> stroke
     >> restore
    drawCar = arc 1 0 (carSize/2) 0 (2*pi) >> fill


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- |
-- Module      :  Test.Opcount
-- Copyright   :  (c) 2016-2020 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Test.Opcount (
    opCountTests,
    splitRadixOpcountTests,
    difOpcountTests,
    opcountRegressionTests
  ) where

import Control.Monad (mzero)
import Data.Complex
import Data.Typeable (Typeable)
import Test.HUnit (Assertion,
                   (@?=),
                   assertEqual)
import Test.Hspec

import qualified Data.FlagSet as FS

import Spiral (Config(..))
import Spiral.Config
import Spiral.Driver
import Spiral.Exp
import Spiral.FFT.CooleyTukey
import Spiral.Monad
import Spiral.OpCount
import Spiral.SPL
import Spiral.SPL.Run
import Spiral.Search
import Spiral.Search.FFTBreakdowns
import Spiral.Search.OpCount

import Test.Instances ()

opCountTests :: Spec
opCountTests =
    describe "Opcount" $ do
    splitRadixOpcountTests
    difOpcountTests

-- The number of multiplies and additions for spit radix decomposition of size n
-- when using three-multiply form of complex multiply. Taken from Table II of
-- Heideman and Burrus.
splitRadixOpcounts :: [(Int, Int, Int)]
splitRadixOpcounts = [ (4,    0,     16)
                     , (8,    4,     52)
                     , (16,   20,    148)
                     , (32,   68,    388)
                     , (64,   196,   964)
                     , (128,  516,   2308)
                     , (256,  1284,  5380)
                     --, (512,  3076,  12292)
                     --, (1024, 7172,  27652)
                     --, (2048, 16388, 61444)
                     --, (4096, 36868, 135172)
                     ]

splitRadixOpcountTests :: Spec
splitRadixOpcountTests =
    describe "Split radix opcounts" $
    sequence_ [mkTest n (muls + adds) | (n, muls, adds) <- splitRadixOpcounts]
  where
    mkTest :: Int -> Int -> Spec
    mkTest n nops =
      mkOpCountTest ("Split radix " ++ show n) fs nops (runSearch () f (DFT n))
      where
        fs :: [DynFlag]
        fs = [ StoreIntermediate
             , SplitComplex
             , CSE
             , Rewrite
             ]

    f :: (Typeable a, Typed a, MonadSpiral m)
      => SPL (Exp a)
      -> S s m (SPL (Exp a))
    f (F n w) = splitRadixBreakdown n w
    f _       = mzero

-- The DIF form should have the same operation count as split radix when the
-- DifRewrite flag is enabled.
difOpcountTests :: Spec
difOpcountTests =
    describe "DIF opcounts" $
    sequence_ [mkTest n (muls + adds) | (n, muls, adds) <- splitRadixOpcounts]
  where
    mkTest :: Int -> Int -> Spec
    mkTest n nops =
      mkOpCountTest ("DIF " ++ show n) fs nops (return $ dif n)
      where
        fs :: [DynFlag]
        fs = [ StoreIntermediate
             , SplitComplex
             , CSE
             , Rewrite
             , DifRewrite
             ]

mkOpCountTest :: String
              -> [DynFlag]
              -> Int
              -> Spiral (SPL (Exp (Complex Double)))
              -> Spec
mkOpCountTest desc fs nops gen =
    it desc $ do
    ops <- runSpiralWith mempty $ withOpcountFlags fs $ do
           f   <- gen
           toProgram "f" (Re f) >>= countProgramOps
    return $ mulOps ops + addOps ops @?= nops

withOpcountFlags :: MonadConfig m => [DynFlag] -> m a -> m a
withOpcountFlags fs =
    localConfig $ \env -> env { dynFlags  = FS.fromList fs
                              , maxUnroll = 256
                              }

opcountRegressionTests :: Int -> Spec
opcountRegressionTests max_size = do
    text <- runIO $ readFile rEGRESSION_FILE
    let opcounts = (map parseCSV . drop 1 . lines) text
    mapM_ test [(n,totalOps,mulOps,addOps) | [n,totalOps,mulOps,addOps] <- opcounts, n <= max_size]
  where
    rEGRESSION_FILE :: FilePath
    rEGRESSION_FILE = "benchmark/data/search-opcount.csv"

    parseCSV :: String -> [Int]
    parseCSV = map read . splitOn (== ',')

    test :: (Int,Int,Int,Int) -> Spec
    test (n, allOps0, mulOps0, addOps0) = it ("Regression(" ++ show n ++ ")") $ do
        ops <- runSpiralWith config $ do
               e    <- searchOpCount (Re (DFT n) :: SPL (Exp Double))
               prog <- toProgram "dft" e
               countProgramOps prog
        return $ checkOps (allOps ops, mulOps ops, addOps ops) (allOps0, mulOps0, addOps0)

    checkOps :: (Int, Int, Int) -> (Int, Int, Int) -> Assertion
    checkOps ops@(total, _, _) ops'@(total', _, _)
      | total < total' = assertEqual "IMPROVED!" ops' ops
      | otherwise      = ops @?= ops'

    config :: Config
    config = mempty { dynFlags  = FS.fromList fs
                    , maxUnroll = 256
                    }

    fs :: [DynFlag]
    fs = [ StoreIntermediate
         , SplitComplex
         , CSE
         , Rewrite
         ]

splitOn :: (Char -> Bool) -> String -> [String]
splitOn p s = case dropWhile p s of
                  "" -> []
                  s' -> w : splitOn p s''
                        where (w, s'') = break p s'


{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
{-# LANGUAGE FlexibleContexts #-}
module Neural.Network
  ( Network(..)
  , weights
  , biases
  , newNetwork
  , randomNetwork
  , feedforwards
  , feedforward
  )
where

import           System.Random
import           Foreign.Storable               ( Storable )
import           Control.Monad.State
import qualified Data.Semigroup                as Semigroup
import           Numeric.LinearAlgebra
import           GHC.Generics                   ( Generic )
import           Control.DeepSeq

import           Neural.Activation              ( ActivationFunction(..) )
import           Neural.Layer            hiding ( weights
                                                , biases
                                                )
import qualified Neural.Layer                  as Layer

-- | Simple feed-forward network.
newtype Network a = Network [Layer a]
    deriving (Show, Generic, NFData)

-- | Get the weights of the network.
weights :: Network a -> [Matrix a]
weights (Network layers) = map Layer.weights layers

-- | Get the biases of the network.
biases :: Network a -> [Vector a]
biases (Network layers) = map Layer.biases layers

-- | Construct a new network from a list of weigth matricws and bias vectors.
newNetwork :: [Matrix a] -> [Vector a] -> Network a
newNetwork ws bs = Network $ zipWith newLayer ws bs

-- | Construct a new random network given the sizes of the layers
-- Each bias and weight is given a random value normally
-- distributed with mean 0 and standard deviation 1.
randomNetwork
  :: (RandomGen g, Random a, Element a, Floating a)
  => [Int]    -- ^ The sizes of the layers
  -> State g (Network a)
randomNetwork sizes = do
  let sizes' = zip (init sizes) (tail sizes)
  layers <- mapM (uncurry randomLayer) sizes'
  return $ Network layers

-- | Propagate an input through the network and return every intermediate
-- activation and weighted input along the way.
feedforwards
  :: (Numeric a, Num (Vector a))
  => ActivationFunction a   -- ^ The activation function
  -> Vector a               -- ^ The network input
  -> Network a              -- ^ The network
  -> [(Vector a, Vector a)] -- ^ The activations of all the layers
feedforwards σ x (Network layers) = tail $ scanl (feed σ . fst) (x, x) layers

-- | Propagate an input through the network
-- and return the activation and weighted input of the last layer.
feedforward
  :: (Numeric a, Num (Vector a))
  => ActivationFunction a  -- ^ The activation function
  -> Vector a              -- ^ The network input
  -> Network a             -- ^ The network
  -> (Vector a, Vector a)  -- ^ The activations of the last layer
feedforward σ x = last . feedforwards σ x


{-# LANGUAGE BangPatterns, DeriveDataTypeable, RecordWildCards #-}
-- |
-- Module      : Criterion.Analysis
-- Copyright   : (c) 2009-2014 Bryan O'Sullivan
--
-- License     : BSD-style
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : GHC
--
-- Analysis code for benchmarks.

module Criterion.Analysis
    (
      Outliers(..)
    , OutlierEffect(..)
    , OutlierVariance(..)
    , SampleAnalysis(..)
    , analyseSample
    , scale
    , analyseMean
    , countOutliers
    , classifyOutliers
    , noteOutliers
    , outlierVariance
    , resolveAccessors
    , validateAccessors
    , regress
    ) where

-- Temporary: to support pre-AMP GHC 7.8.4:
import Data.Monoid

import Control.Arrow (second)
import Control.Monad (unless, when)
import Control.Monad.Reader (ask)
import Control.Monad.Trans
import Control.Monad.Trans.Except
import Criterion.IO.Printf (note, prolix)
import Criterion.Measurement (secs, threshold)
import Criterion.Monad (Criterion, getGen, getOverhead)
import Criterion.Types
import Data.Int (Int64)
import Data.Maybe (fromJust)
import Statistics.Function (sort)
import Statistics.Quantile (weightedAvg)
import Statistics.Regression (bootstrapRegress, olsRegress)
import Statistics.Resampling (resample)
import Statistics.Sample (mean)
import Statistics.Sample.KernelDensity (kde)
import Statistics.Types (Estimator(..), Sample)
import System.Random.MWC (GenIO)
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Statistics.Resampling.Bootstrap as B
import Prelude

-- | Classify outliers in a data set, using the boxplot technique.
classifyOutliers :: Sample -> Outliers
classifyOutliers sa = U.foldl' ((. outlier) . mappend) mempty ssa
    where outlier e = Outliers {
                        samplesSeen = 1
                      , lowSevere = if e <= loS && e < hiM then 1 else 0
                      , lowMild = if e > loS && e <= loM then 1 else 0
                      , highMild = if e >= hiM && e < hiS then 1 else 0
                      , highSevere = if e >= hiS && e > loM then 1 else 0
                      }
          !loS = q1 - (iqr * 3)
          !loM = q1 - (iqr * 1.5)
          !hiM = q3 + (iqr * 1.5)
          !hiS = q3 + (iqr * 3)
          q1   = weightedAvg 1 4 ssa
          q3   = weightedAvg 3 4 ssa
          ssa  = sort sa
          iqr  = q3 - q1

-- | Compute the extent to which outliers in the sample data affect
-- the sample mean and standard deviation.
outlierVariance :: B.Estimate  -- ^ Bootstrap estimate of sample mean.
                -> B.Estimate  -- ^ Bootstrap estimate of sample
                               --   standard deviation.
                -> Double      -- ^ Number of original iterations.
                -> OutlierVariance
outlierVariance µ σ a = OutlierVariance effect desc varOutMin
  where
    ( effect, desc ) | varOutMin < 0.01 = (Unaffected, "no")
                     | varOutMin < 0.1  = (Slight,     "slight")
                     | varOutMin < 0.5  = (Moderate,   "moderate")
                     | otherwise        = (Severe,     "severe")
    varOutMin = (minBy varOut 1 (minBy cMax 0 µgMin)) / σb2
    varOut c  = (ac / a) * (σb2 - ac * σg2) where ac = a - c
    σb        = B.estPoint σ
    µa        = B.estPoint µ / a
    µgMin     = µa / 2
    σg        = min (µgMin / 4) (σb / sqrt a)
    σg2       = σg * σg
    σb2       = σb * σb
    minBy f q r = min (f q) (f r)
    cMax x    = fromIntegral (floor (-2 * k0 / (k1 + sqrt det)) :: Int)
      where
        k1    = σb2 - a * σg2 + ad
        k0    = -a * ad
        ad    = a * d
        d     = k * k where k = µa - x
        det   = k1 * k1 - 4 * σg2 * k0

-- | Count the total number of outliers in a sample.
countOutliers :: Outliers -> Int64
countOutliers (Outliers _ a b c d) = a + b + c + d
{-# INLINE countOutliers #-}

-- | Display the mean of a 'Sample', and characterise the outliers
-- present in the sample.
analyseMean :: Sample
            -> Int              -- ^ Number of iterations used to
                                -- compute the sample.
            -> Criterion Double
analyseMean a iters = do
  let µ = mean a
  _ <- note "mean is %s (%d iterations)\n" (secs µ) iters
  noteOutliers . classifyOutliers $ a
  return µ

-- | Multiply the 'Estimate's in an analysis by the given value, using
-- 'B.scale'.
scale :: Double                 -- ^ Value to multiply by.
      -> SampleAnalysis -> SampleAnalysis
scale f s@SampleAnalysis{..} = s {
                                 anMean = B.scale f anMean
                               , anStdDev = B.scale f anStdDev
                               }

-- | Perform an analysis of a measurement.
analyseSample :: Int            -- ^ Experiment number.
              -> String         -- ^ Experiment name.
              -> V.Vector Measured -- ^ Sample data.
              -> ExceptT String Criterion Report
analyseSample i name meas = do
  Config{..} <- ask
  overhead <- lift getOverhead
  let ests      = [Mean,StdDev]
      -- The use of filter here throws away very-low-quality
      -- measurements when bootstrapping the mean and standard
      -- deviations.  Without this, the numbers look nonsensical when
      -- very brief actions are measured.
      stime     = measure (measTime . rescale) .
                  G.filter ((>= threshold) . measTime) . G.map fixTime .
                  G.tail $ meas
      fixTime m = m { measTime = measTime m - overhead / 2 }
      n         = G.length meas
      s         = G.length stime
  _ <- lift $ prolix "bootstrapping with %d of %d samples (%d%%)\n"
              s n ((s * 100) `quot` n)
  gen <- lift getGen
  rs <- mapM (\(ps,r) -> regress gen ps r meas) $
        ((["iters"],"time"):regressions)
  resamps <- liftIO $ resample gen ests resamples stime
  let [estMean,estStdDev] = B.bootstrapBCA confInterval stime ests resamps
      ov = outlierVariance estMean estStdDev (fromIntegral n)
      an = SampleAnalysis {
               anRegress    = rs
             , anOverhead   = overhead
             , anMean       = estMean
             , anStdDev     = estStdDev
             , anOutlierVar = ov
             }
  return Report {
      reportNumber   = i
    , reportName     = name
    , reportKeys     = measureKeys
    , reportMeasured = meas
    , reportAnalysis = an
    , reportOutliers = classifyOutliers stime
    , reportKDEs     = [uncurry (KDE "time") (kde 128 stime)]
    }

-- | Regress the given predictors against the responder.
--
-- Errors may be returned under various circumstances, such as invalid
-- names or lack of needed data.
--
-- See 'olsRegress' for details of the regression performed.
regress :: GenIO
        -> [String]             -- ^ Predictor names.
        -> String               -- ^ Responder name.
        -> V.Vector Measured
        -> ExceptT String Criterion Regression
regress gen predNames respName meas = do
  when (G.null meas) $
    throwE "no measurements"
  accs <- ExceptT . return $ validateAccessors predNames respName
  let unmeasured = [n | (n, Nothing) <- map (second ($ G.head meas)) accs]
  unless (null unmeasured) $
    throwE $ "no data available for " ++ renderNames unmeasured
  let (r:ps)      = map ((`measure` meas) . (fromJust .) . snd) accs
  Config{..} <- ask
  (coeffs,r2) <- liftIO $
                 bootstrapRegress gen resamples confInterval olsRegress ps r
  return Regression {
      regResponder = respName
    , regCoeffs    = Map.fromList (zip (predNames ++ ["y"]) (G.toList coeffs))
    , regRSquare   = r2
    }

singleton :: [a] -> Bool
singleton [_] = True
singleton _   = False

-- | Given a list of accessor names (see 'measureKeys'), return either
-- a mapping from accessor name to function or an error message if
-- any names are wrong.
resolveAccessors :: [String]
                 -> Either String [(String, Measured -> Maybe Double)]
resolveAccessors names =
  case unresolved of
    [] -> Right [(n, a) | (n, Just (a,_)) <- accessors]
    _  -> Left $ "unknown metric " ++ renderNames unresolved
  where
    unresolved = [n | (n, Nothing) <- accessors]
    accessors = flip map names $ \n -> (n, Map.lookup n measureAccessors)

-- | Given predictor and responder names, do some basic validation,
-- then hand back the relevant accessors.
validateAccessors :: [String]   -- ^ Predictor names.
                  -> String     -- ^ Responder name.
                  -> Either String [(String, Measured -> Maybe Double)]
validateAccessors predNames respName = do
  when (null predNames) $
    Left "no predictors specified"
  let names = respName:predNames
      dups = map head . filter (not . singleton) .
             List.group . List.sort $ names
  unless (null dups) $
    Left $ "duplicated metric " ++ renderNames dups
  resolveAccessors names

renderNames :: [String] -> String
renderNames = List.intercalate ", " . map show

-- | Display a report of the 'Outliers' present in a 'Sample'.
noteOutliers :: Outliers -> Criterion ()
noteOutliers o = do
  let frac n = (100::Double) * fromIntegral n / fromIntegral (samplesSeen o)
      check :: Int64 -> Double -> String -> Criterion ()
      check k t d = when (frac k > t) $
                    note "  %d (%.1g%%) %s\n" k (frac k) d
      outCount = countOutliers o
  when (outCount > 0) $ do
    _ <- note "found %d outliers among %d samples (%.1g%%)\n"
         outCount (samplesSeen o) (frac outCount)
    check (lowSevere o) 0 "low severe"
    check (lowMild o) 1 "low mild"
    check (highMild o) 1 "high mild"
    check (highSevere o) 0 "high severe"


{-# LANGUAGE NamedFieldPuns, TemplateHaskell #-}

module School.Train.Test.ForwardPass
( forwardPassTest ) where

import Conduit ((.|), yield, liftIO, sinkList, await)
import Data.Either (isLeft)
import Numeric.LinearAlgebra (R)
import School.TestUtils (doCost, fromRight, randomAffineParams, randomMatrix, testState, weight1)
import School.Train.AppTrain (AppTrain)
import School.Train.ForwardPass
import School.Train.TrainState (TrainState(..), def)
import School.Types.PingPong (pingPongSingleton, reversePingPong, toPingPong)
import School.Types.Slinky (Slinky(..))
import School.Unit.Affine (affine)
import School.Unit.CostFunction (CostFunction)
import School.Unit.RecLin (recLin)
import School.Unit.Unit (Unit(..))
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitBackward (BackwardStack)
import School.Unit.UnitParams (UnitParams(..))
import Test.Tasty (TestTree)
import Test.Tasty.QuickCheck hiding ((><))
import Test.Tasty.TH
import Test.QuickCheck.Monadic (assert, monadicIO)

weight :: CostFunction R (AppTrain R)
weight = weight1

prop_no_units :: Property
prop_no_units = monadicIO $ do
  let forward = forwardPass [] weight
  let source = yield ([], SNil)
  let pass = source .| forward .| sinkList
  result <- testState pass def
  assert $ isLeft result

prop_single_recLin :: (Positive Int) -> (Positive Int) -> Property
prop_single_recLin (Positive bSize) (Positive fSize) = monadicIO $ do
  let forward = forwardPass [recLin] weight
  input <- liftIO $ BatchActivation <$> (randomMatrix bSize fSize)
  let source = yield ([input], SNil)
  let pass = source .| forward .| await
  result <- testState pass def
  let out = apply recLin EmptyParams input
  let (cost, grad) = doCost weight out SNil
  let bParams = reversePingPong . paramList $ def
  let state = def { paramList = bParams }
  let stack = ([input], grad, cost)
  let check = Right $ (Just stack, state) :: Either String (Maybe (BackwardStack R), TrainState R)
  assert $ result == check

prop_aff_rl_aff_rl :: (Positive Int) -> (Positive Int) -> (Positive Int) -> (Positive Int) -> Property
prop_aff_rl_aff_rl (Positive b) (Positive f) (Positive h) (Positive o) = monadicIO $ do
  let units = [affine, recLin, affine, recLin]
  let forward = forwardPass units weight
  input <- liftIO $ BatchActivation <$> (randomMatrix b f)
  let source = yield ([input], SNil)
  let pass = source .| forward .| await
  params1 <- liftIO $ randomAffineParams f h
  params2 <- liftIO $ randomAffineParams h o
  let allParams = toPingPong [ params1
                             , EmptyParams
                             , params2
                             , EmptyParams
                             ]
  let paramList = fromRight (pingPongSingleton EmptyParams) allParams
  let initState = def { paramList }
  result <- testState pass initState
  let out1 = apply affine params1 input
  let out2 = apply recLin EmptyParams out1
  let out3 = apply affine params2 out2
  let out4 = apply recLin EmptyParams out3
  let (cost, grad) = doCost weight out4 SNil
  let bParams = reversePingPong paramList
  let state = def { paramList = bParams }
  let stack = ([out3, out2, out1, input], grad, cost)
  let check = Right $ (Just stack, state) :: Either String (Maybe (BackwardStack R), TrainState R)
  assert $ result == check

forwardPassTest:: TestTree
forwardPassTest = $(testGroupGenerator)