{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Statistics.Array
  ( median
  , minimum
  , maximum
  , quartiles
  , listDerivative
  , mad
  , bowleySkew
  ) where

import Prelude hiding (minimum,maximum)

import Control.Monad (forM_)
import Data.Int (Int64)
import Data.Primitive (Prim,PrimArray)
import Data.Primitive.Contiguous (index, write)
import Statistics.Array.Types (Quartiles(..),AscList(..))

import qualified Data.Primitive.Contiguous as Arr
import qualified Data.Primitive.Sort as Arr
import qualified Statistics.Array.Types as Asc

-- | @O(1)@ Median element.
median :: Prim a
  => AscList a -> a
{-# inlineable median #-}
{-# specialize median :: AscList Int -> Int #-}
{-# specialize median :: AscList Int64 -> Int64 #-}
{-# specialize median :: AscList Double -> Double #-}
median (AscList arr) = Arr.index arr (Arr.size arr `div` 2)

-- | @O(1)@ Minimum element.
minimum :: Prim a
  => AscList a -> a
{-# inlineable minimum #-}
{-# specialize minimum :: AscList Int -> Int #-}
{-# specialize minimum :: AscList Int64 -> Int64 #-}
{-# specialize minimum :: AscList Double -> Double #-}
minimum (AscList arr) = Arr.index arr 0

-- | @O(1)@ Maximum element.
maximum :: Prim a
  => AscList a -> a
{-# inlineable maximum #-}
{-# specialize maximum :: AscList Int -> Int #-}
{-# specialize maximum :: AscList Int64 -> Int64 #-}
{-# specialize maximum :: AscList Double -> Double #-}
maximum (AscList arr) = Arr.index arr (Arr.size arr - 1)

quartiles :: Prim a
  => AscList a -> Quartiles a
quartiles (AscList arr) = Quartiles
  { q1 = Arr.index arr (len `div` 4)
  , q2 = Arr.index arr (len `div` 2)
  , q3 = Arr.index arr (3 * len `div` 4)
  }
  where
  len = Arr.size arr

-- hmmm, I guess I'm going with Int64 here to match with Chronos
-- times will be stored in number of nanoseconds after unix epoch
-- which allows me times up to about the year 2262

-- so:
--   check skew of time diffs
--   check dispersion of time diffs (median absolute deviation)
--   check skew of data sizes
--   check dispersion of time diffs (median absolute deviation)
--   check size of data (at the mode)

-- | Median of absolute deviations.
-- This is a measure of dispersion that is more robust than standard deviation.
mad :: forall a.
     (Prim a, Ord a, Num a)
  => AscList a -> a
{-# inlineable mad #-}
{-# specialize mad :: AscList Int -> Int #-}
{-# specialize mad :: AscList Int64 -> Int64 #-}
{-# specialize mad :: AscList Double -> Double #-}
mad asc@(AscList arr) =
  let m = median asc
      devs = (\x -> abs (x - m)) `Arr.map` arr :: PrimArray a
      devsAsc = Arr.sort devs
   in median $ Asc.unsafeFromAscendingArray devsAsc

-- | Take the difference between adjacent elements. Resulting array has
-- length one less than the argument array.
listDerivative :: (Prim a, Num a)
  => AscList a -> PrimArray a
{-# inlineable listDerivative #-}
{-# specialize listDerivative :: AscList Int -> PrimArray Int #-}
{-# specialize listDerivative :: AscList Int64 -> PrimArray Int64 #-}
{-# specialize listDerivative :: AscList Double -> PrimArray Double #-}
listDerivative (AscList arr) = Arr.create $ do
  let len' = Arr.size arr - 1
  diffs <- Arr.new len'
  let go !dstIx !prevElement = if dstIx < len'
        then do
          let !currentElement = index arr (dstIx + 1)
          write diffs dstIx (currentElement - prevElement)
          go (dstIx + 1) currentElement
        else pure diffs
  go 0 (index arr 0)

bowleySkew :: (Integral a) => Quartiles a -> Double
{-# inline bowleySkew #-}
bowleySkew Quartiles{q1,q2,q3}
  -- the limit as q1 approaches q3 is of course dependant on the direction of approach in the 3d space {q1,q2,q3}
  -- realistically, I'm going to call this zero, since real data has no approach, and the dirac distribution has zero skew
  | q1 == q3 = 0.0
  -- if `q1 == q2`, then skew simplifies to `(q3 - q1) / (q3 - q1)`, or `1`
  | q2 == q1 = 1.0
  -- similarly, if `q2 == q3`, then skew simplifies to `-1`
  | q2 == q3 = -1.0
  -- I'm not sure if either of those situations are reliable indicators of skew, though
  | otherwise = (realToFrac $ q3 + q1 - 2*q2) / (realToFrac $ q3 - q1)


module NeuralNetwork where

import           Numeric.LinearAlgebra as LA
import           Text.Printf ( printf )
import           Numeric.Morpheus.Activation ( relu, reluGradient, sigmoid,
                    sigmoidGradient, tanh_, tanhGradient)

data Activation = Relu | Sigmoid | Tanh | Id

-- Neural network layer: weights, biases, and activation
data Layer a = Layer (Matrix a) (Matrix a) Activation

type NeuralNetwork a = [Layer a]

-- | Weight and bias gradients
data Gradients a = Gradients (Matrix a) (Matrix a)

getActivation :: Activation -> (Matrix Double -> Matrix Double)
getActivation Id = id
getActivation Sigmoid = sigmoid
getActivation Relu = relu
getActivation Tanh = tanh_

getActivation' :: Activation -> Matrix R -> Matrix R -> Matrix R
getActivation' Id = flip const
getActivation' Sigmoid = \x dY -> sigmoidGradient x * dY
getActivation' Relu = \x dY -> reluGradient x * dY
getActivation' Tanh = \x dY -> tanhGradient x * dY

bias' :: Matrix R -> Matrix R
bias' dY = cmap (/ m) r
  where
    r = matrix (cols dY) $ map sumElements (toColumns dY)
    m = fromIntegral $ rows dY

linearW' :: (Numeric t, Fractional t) => Matrix t -> Matrix t -> Matrix t
linearW' x dy = cmap (/ m) (tr' x LA.<> dy)
  where
    m = fromIntegral $ rows x

linearX' :: Numeric t => Matrix t -> Matrix t -> Matrix t
linearX' w dy = dy LA.<> tr' w

-- | Gradient descent optimization
optimize :: R -> Int -> [Layer R] -> (Matrix R, Matrix R) -> [Layer R]
optimize lr iterN net0 dataSet = last $ take iterN (iterate step net0)
  where
    step net = zipWith f net dW
        where
            (_, dW) = pass net dataSet

    f :: Layer R -> Gradients R -> Layer R
    f (Layer w b act) (Gradients dW dB) =
        Layer (w - lr `scale` dW) (b - lr `scale` dB) act

pass :: NeuralNetwork R -> (Matrix R, Matrix R) -> (Matrix R, [Gradients R])
pass net (x, tgt) = (pred, grads)
  where
    (_, pred, grads) = _pass x net

    _pass inp [] = (loss', pred, [])
      where
        pred = sigmoid inp
        loss' = pred - tgt

    _pass inp (Layer w b sact:layers) = (dX, pred, Gradients dW dB:t)
      where
        lin = (inp LA.<> w) + b
        y = getActivation sact lin

        (dZ, pred, t) = _pass y layers

        dY = getActivation' sact lin dZ
        dW = linearW' inp dY
        dB = bias' dY
        dX = linearX' w dY

data AdamParameters = AdamParameters { _beta1 :: R, _beta2 :: R, _epsilon :: R, _lr :: R }

adamParams :: AdamParameters
adamParams = AdamParameters { _beta1 = 0.9
                            , _beta2 = 0.999
                            , _epsilon = 1e-8
                            , _lr = 0.001
                            }

-- | Adam optimization
optimizeAdam :: AdamParameters -> Int -> [Layer R] -> (Matrix R, Matrix R) -> [Layer R]
optimizeAdam p iterN w0 dataSet = w
  where
    s0 = map zf w0
    v0 = map zf w0
    zf (Layer a b _) = (zerosLike a, zerosLike b)
    zerosLike m = matrix c (replicate (r*c) 0.0)
      where
        r = rows m
        c = cols m
    (w, _, _) = adam p iterN (w0, s0, v0) dataSet

adam :: AdamParameters -> Int -> ([Layer R], [(Matrix R, Matrix R)], [(Matrix R, Matrix R)]) -> (Matrix R, Matrix R) -> ([Layer R], [(Matrix R, Matrix R)], [(Matrix R, Matrix R)])
adam p@AdamParameters { _lr = lr, _beta1 = beta1, _beta2 = beta2, _epsilon = epsilon } iterN (w0, s0, v0) dataSet = last $ take iterN (iterate step (w0, s0, v0))
      where
        step (w, s, v) = (wN, sN, vN)
          where
            (_, dW) = pass w dataSet

            sN = zipWith f2 s dW
            vN = zipWith f3 v dW
            wN = zipWith3 f w vN sN

        f (Layer w_ b_ sf) (vW, vB) (sW, sB) =
           Layer (w_ - lr `scale` vW / (sqrt sW `addC` epsilon))
                 (b_ - lr `scale` vB / (sqrt sB `addC` epsilon))
                 sf

        addC m c = cmap (+ c) m

        f2 (sW, sB) (Gradients dW dB) =
          ( beta2 `scale` sW + (1 - beta2) `scale` (dW^2)
          , beta2 `scale` sB + (1 - beta2) `scale` (dB^2))

        f3 (vW, vB) (Gradients dW dB) =
          ( beta1 `scale` vW + (1 - beta1) `scale` dW
          , beta1 `scale` vB + (1 - beta1) `scale` dB)

genWeights :: (Int, Int) -> IO (Matrix Double, Matrix Double)
genWeights (nin, nout) = do
  w <- _genWeights (nin, nout)
  b <- _genWeights (1, nout)
  return (w, b)
    where
      _genWeights (nin, nout) = do
          let k = sqrt (1.0 / fromIntegral nin)
          w <- randn nin nout
          return (k `scale` w)

makeCircles :: Int -> R -> R -> IO (Matrix R, Matrix R)
makeCircles m factor noise = do
  let rand' n = scale (2 * pi) <$> rand n 1
      m1 = m `div` 2
      m2 = m - (m `div` 2)

  r1 <- rand' m1
  r2 <- rand' m2
  ns <- scale noise <$> randn m 2

  let outerX = cos r1
      outerY = sin r1
      innerX = scale factor $ cos r2
      innerY = scale factor $ sin r2
      -- Merge them all
      x = fromBlocks [[outerX, outerY], [innerX, innerY]]

      -- Labels
      y1 = m1 >< 1 $ repeat 0
      y2 = m2 >< 1 $ repeat 1
      y = y1 === y2

  return (x + ns, y)

makeSpirals :: Int -> R -> IO (Matrix R, Matrix R)
makeSpirals m noise = do
  r0 <- scale (780 * 2 * pi / 360) . sqrt <$> rand m 1
  d1x0 <- scale noise <$> rand m 1
  d1y0 <- scale noise <$> rand m 1

  let d1x = d1x0 - cos r0 * r0
  let d1y = d1y0 + sin r0 * r0

  let x = fromBlocks [[d1x, d1y], [- d1x, - d1y]] / 10.0
  let y1 = m >< 1 $ repeat 0
  let y2 = m >< 1 $ repeat 1
  let y = y1 === y2
  return (x, y)

accuracy :: [Layer R] -> (Matrix R, Matrix R) -> R
accuracy net (dta, tgt) = 100 * (1 - e / m)
  where
    pred = net `inferBinary` dta
    e = sumElements $ abs (tgt - pred)
    m = fromIntegral $ rows tgt

inferBinary :: NeuralNetwork R -> Matrix R -> Matrix R
inferBinary net dta =
  let pred = forward net dta
  in cmap (\a -> if a < 0.5 then 0 else 1) pred

forward :: NeuralNetwork R -> Matrix R -> Matrix R
forward net dta = fst $ pass net (dta, undefined)

genNetwork :: [Int] -> [Activation] -> IO [Layer R]
genNetwork nodes activations = do
    weights <- mapM genWeights nodes'
    return (zipWith (\(w, b) a -> Layer w b a) weights activations)
  where
    nodes' = zip nodes (tail nodes)

experiment1 :: IO ()
experiment1 = do
  trainSet <- makeCircles 200 0.6 0.1
  testSet <- makeCircles 100 0.6 0.1

  let (dta, tgt) = trainSet

  (w1_rand, b1_rand) <- genWeights (2, 128)
  (w2_rand, b2_rand) <- genWeights (128, 1)

  let net = [ Layer w1_rand b1_rand Relu
            , Layer w2_rand b2_rand Id ]

  -- -- Alternatively:
  -- net0 <- genNetwork [2, 128, 1] [Relu, Id]

  let epochs = 1000
      lr = 0.001  -- Learning rate
      net' = optimize lr epochs net trainSet
      netA = optimizeAdam adamParams epochs net trainSet

  putStrLn $ printf "Circles problem, 1 hidden layer of 128 neurons, %d epochs" epochs
  putStrLn "---"

  putStrLn $ printf "Training accuracy (gradient descent) %.1f" (net' `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy (gradient descent) %.1f\n" (net' `accuracy` testSet)

  putStrLn $ printf "Training accuracy (Adam) %.1f" (netA `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy (Adam) %.1f\n" (netA `accuracy` testSet)

  putStrLn ""

experiment2 :: IO ()
experiment2 = do
  trainSet <- makeSpirals 200 0.5
  testSet <- makeSpirals 100 0.5

  let (dta, tgt) = trainSet
      epochs = 700

  putStrLn $ printf "Spirals problem, Adam, %d epochs" epochs
  putStrLn "---"
  putStrLn "1 hidden layer, 128 neurons (513 parameters)"
  net0 <- genNetwork [2, 128, 1] [Relu, Id]
  let net0' = optimizeAdam adamParams epochs net0 trainSet

  putStrLn $ printf "Training accuracy %.1f" (net0' `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy %.1f\n" (net0' `accuracy` testSet)

  putStrLn "1 hidden layer, 512 neurons (2049 parameters)"
  net1 <- genNetwork [2, 512, 1] [Relu, Id]
  let net1' = optimizeAdam adamParams epochs net1 trainSet

  putStrLn $ printf "Training accuracy %.1f" (net1' `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy %.1f\n" (net1' `accuracy` testSet)

  putStrLn "3 hidden layers, 40, 25, and 10 neurons (1416 parameters)"
  net2 <- genNetwork [2, 40, 25, 10, 1] [Relu, Relu, Relu, Id]
  let net2' = optimizeAdam adamParams epochs net2 trainSet

  putStrLn $ printf "Training accuracy %.1f" (net2' `accuracy` trainSet)
  putStrLn $ printf "Validation accuracy %.1f\n" (net2' `accuracy` testSet)

main :: IO ()
main = experiment1 >> experiment2


{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}

module Day12 where

import Control.Monad.State (MonadState, State, execState)
import Data.ByteString qualified as B
import Data.Complex (Complex((:+)), cis)
import Data.Functor (($>))
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Void (Void)
import Flow ((.>))
import Optics
    ( Optic, A_Setter, Is, view, noPrefixFieldLabels, makeFieldLabelsWith )
import Optics.State (use)
import Optics.State.Operators ((%=))
import Text.Megaparsec qualified as P
import Text.Megaparsec.Char qualified as PC
import Text.Megaparsec.Char.Lexer qualified as PCL
-- import Text.Pretty.Simple (pPrint)

type Parser = P.Parsec Void T.Text

data NavAction
    = MoveNorth
    | MoveSouth
    | MoveEast
    | MoveWest
    | TurnLeft
    | TurnRight
    | MoveForward
    deriving (Bounded, Enum, Eq, Ord, Read, Show)

data NavInstr = MkNavInstr
    { navAction :: !NavAction
    , navValue :: !Int
    } deriving (Eq, Ord, Read, Show)

data ShipState = MkShipState
    { shipPhase :: !Double
    , shipPosition :: !(Complex Double)
    , wayPosition :: !(Complex Double)
    } deriving (Eq, Read, Show)

makeFieldLabelsWith noPrefixFieldLabels ''ShipState

main :: IO ()
main = do
    file <- readFileUtf8 "day-12/input.txt"
    case P.parse (parseNavInstrs <* P.eof) "day 12 input" file of
        Left err -> putStrLn $ P.errorBundlePretty err
        Right navInstrs -> do
            print $ part1 navInstrs
            print $ part2 navInstrs

part2 :: [NavInstr] -> Double
part2 = pilotShip2 .> view #shipPosition .> manhattanNorm

part1 :: [NavInstr] -> Double
part1 = pilotShip1 .> view #shipPosition .> manhattanNorm

pilotShip2 :: [NavInstr] -> ShipState
pilotShip2 instrs = execState (applyInstrs2 instrs) $
    MkShipState
        { shipPhase = 0
        , shipPosition = pure 0
        , wayPosition = 10 :+ 1
        }

pilotShip1 :: [NavInstr] -> ShipState
pilotShip1 instrs = execState (applyInstrs1 instrs) $
    MkShipState
        { shipPhase = 0
        , shipPosition = pure 0
        , wayPosition = pure 0
        }

applyInstrs1 :: [NavInstr] -> State ShipState ()
applyInstrs1 = foldr (\instr st -> applyInstr1 instr *> st) (pure ())

applyInstr1 :: NavInstr -> State ShipState ()
applyInstr1 (MkNavInstr action intVal) = case action of
    MoveNorth -> #shipPosition += (0 :+ val)
    MoveSouth -> #shipPosition -= (0 :+ val)
    MoveEast -> #shipPosition += (val :+ 0)
    MoveWest -> #shipPosition -= (val :+ 0)
    TurnLeft -> #shipPhase %= (+ val) .> (`modulo` 360)
    TurnRight -> #shipPhase %= subtract val .> (`modulo` 360)
    MoveForward -> do
        phase <- use #shipPhase
        #shipPosition += fmap (* val) (cisDegrees phase)
  where
    val = fromIntegral intVal

applyInstrs2 :: [NavInstr] -> State ShipState ()
applyInstrs2 = foldr (\instr st -> applyInstr2 instr *> st) (pure ())

applyInstr2 :: NavInstr -> State ShipState ()
applyInstr2 (MkNavInstr action intVal) = case action of
    MoveNorth -> #wayPosition += (0 :+ val)
    MoveSouth -> #wayPosition -= (0 :+ val)
    MoveEast -> #wayPosition += (val :+ 0)
    MoveWest -> #wayPosition -= (val :+ 0)
    TurnLeft -> #wayPosition *= cisDegrees val
    TurnRight -> #wayPosition *= cisDegrees (-val)
    MoveForward -> do
        wayPos <- use #wayPosition
        #shipPosition += fmap (* val) wayPos
  where
    val = fromIntegral intVal

cisDegrees :: Double -> Complex Double
cisDegrees degRaw
    | deg == 0 = 1 :+ 0
    | deg == 90 = 0 :+ 1
    | deg == 180 = (-1) :+ 0
    | deg == 270 = 0 :+ (-1)
    | otherwise = cis rad
  where
    deg = degRaw `modulo` 360
    rad = deg * pi / 180

modulo :: Double -> Double -> Double
modulo x m = x - floorDbl (x / m) * m
  where
    floorDbl = floor .> fromIntegral @Int

manhattanDist :: RealFloat a => Complex a -> Complex a -> a
manhattanDist u v = sum $ fmap abs $ v - u

manhattanNorm :: RealFloat a => Complex a -> a
manhattanNorm = manhattanDist (pure 0)

parseNavInstrs :: Parser [NavInstr]
parseNavInstrs = parseNavInstr `P.sepEndBy` PC.space

parseNavInstr :: Parser NavInstr
parseNavInstr = do
    action <- parseNavAction
    val <- PC.space *> PCL.decimal
    pure $ MkNavInstr action val

parseNavAction :: Parser NavAction
parseNavAction = P.choice
    [ PC.char 'N' $> MoveNorth
    , PC.char 'S' $> MoveSouth
    , PC.char 'E' $> MoveEast
    , PC.char 'W' $> MoveWest
    , PC.char 'L' $> TurnLeft
    , PC.char 'R' $> TurnRight
    , PC.char 'F' $> MoveForward
    ]

readFileUtf8 :: FilePath -> IO T.Text
readFileUtf8 path = TE.decodeUtf8 <$> B.readFile path

(+=)
    :: (Is k A_Setter, MonadState s m, Num a)
    => Optic k is s s a a
    -> a
    -> m ()
optic += x = optic %= (+ x)

(-=)
    :: (Is k A_Setter, MonadState s m, Num a)
    => Optic k is s s a a
    -> a
    -> m ()
optic -= x = optic %= subtract x

(*=)
    :: (Is k A_Setter, MonadState s m, Num a)
    => Optic k is s s a a
    -> a
    -> m ()
optic *= x = optic %= (* x)


{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE BlockArguments #-}
module AmpOp where

import Control.Monad.Fix

import Signal
import Circuit
import Data.Complex

type Input = Complex Double
type Output = Complex Double
type Value = Double
type Impedance = Signal (Complex Double)

dc12 :: Signal Input
dc12 = Signal $ const 12

dc14 :: Signal Input
dc14 = Signal $ const 14

dc100 :: Signal Input
dc100 = Signal $ const 100

sen :: Signal Input
sen = Signal $ \time -> 12 * sin ((pi / 2) * time)

outlet :: Signal Input
outlet = Signal $ \time -> 220 * sqrt 2 * sin (2 * pi * 60 * time)

ground :: Signal Input
ground = Signal $ const 0

data AmpOp = AmpOp { name :: String,
                     openLoopGain :: Signal (Complex Double),
                     manufacturer :: String
                }

data State =
  OneNode (Complex Double) |
  TwoNodes (Complex Double, Complex Double) |
  ThreeNodes (Complex Double, Complex Double, Complex Double) |
  FourNodes (Complex Double, Complex Double, Complex Double, Complex Double)

getOutput :: State -> Output
getOutput (OneNode x) = x
getOutput (TwoNodes (_,x)) = x
getOutput (ThreeNodes (_,_,x)) = x
getOutput (FourNodes (_,_,_,x)) = x

data SignalState =
  SignalOneNode (Signal (Complex Double)) |
  SignalTwoNodes (Signal (Complex Double), Signal (Complex Double)) |
  SignalThreeNodes (Signal (Complex Double), Signal (Complex Double), Signal (Complex Double)) |
  SignalFourNodes (Signal (Complex Double), Signal (Complex Double), Signal (Complex Double), Signal (Complex Double))
  deriving (Show)

getSignalOutput :: SignalState -> Signal Output
getSignalOutput (SignalOneNode x) = x
getSignalOutput (SignalTwoNodes (_,x)) = x
getSignalOutput (SignalThreeNodes (_,_,x)) = x
getSignalOutput (SignalFourNodes (_,_,_,x)) = x

instance Show AmpOp where
    show ap = Prelude.show "AmpOp => Name: " ++ Prelude.show (name ap)

lm741 :: AmpOp
lm741 = AmpOp "LM741" 200000 "Texas. Instruments"

r1 :: Impedance
r1 = 10000

r2 :: Impedance
r2 = 50000

r3 :: Impedance
r3 = 20000

r4 :: Impedance
r4 = 30000

eps = 0.000001

ampOpBuffer :: AmpOp -> SignalState -> Either String (Circuit Input SignalState)
ampOpBuffer model (SignalOneNode initial) =
    Right $ ($ initial) <$> mfix
    (\f -> Circuit $ \vIn -> Signal $ \time vOutOld ->
        let vOut = (vIn * openLoopGain model) / (1 + openLoopGain model)
        in if abs((realPart <$> vOut) - (realPart <$> vOutOld)) <= eps
            then SignalOneNode vOut
            else f vOut)
ampOpBuffer _ _ = Left "Wrong type for initial state!"

ampOpNonInverting :: AmpOp -> Impedance -> Impedance -> SignalState -> Either String (Circuit Input SignalState)
ampOpNonInverting model r1 r2 (SignalTwoNodes initial) =
    Right $ ($ initial) <$> mfix
    (\f -> Circuit $ \vIn -> Signal $ \time (vXOld, vOutOld) ->
        let vX = ((vIn * openLoopGain model) - vOutOld) / openLoopGain model
            vOut = (vXOld * (r1 + r2)) / r1
        in if abs((realPart <$> vOut) - (realPart <$> vOutOld)) <= eps && abs((realPart <$> vX) - (realPart <$> vXOld)) <= eps
            then SignalTwoNodes (vX, vOut)
            else f (vX, vOut))
ampOpNonInverting _ _ _ _ = Left "Wrong type for initial state!"

ampOpInverting :: AmpOp -> Impedance -> Impedance -> SignalState -> Either String (Circuit Input SignalState)
ampOpInverting model r1 r2 (SignalTwoNodes initial) =
    Right $ ($ initial) <$> mfix
    (\f -> Circuit $ \vIn -> Signal $ \time (vXOld, vOutOld) ->
        let vX = (- vOutOld) / openLoopGain model
            vOut = ((vXOld * (r1 + r2)) - (vIn * r2)) / r1
        in if abs((realPart <$> vOutOld) - (realPart <$> vOut)) <= eps && abs((realPart <$> vXOld) - (realPart <$> vX)) <= eps
            then SignalTwoNodes (vX, vOut)
            else f (vX, vOut))
ampOpInverting _ _ _ _ = Left "Wrong type for initial state!"

-- ampOpSpecial :: AmpOp -> Impedance -> Impedance -> Impedance -> Impedance -> SignalState -> Either String (Circuit Input SignalState)
-- ampOpSpecial model z1 z2 z3 z4 (SignalThreeNodes initial) =
--   Right $ ($ initial) <$> mfix
--   (\f -> Circuit $ \vIn -> Signal $ \time (vXOld, vKOld, vOutOld) ->
--       let vK = vOutOld + (vOutOld / openLoopGain model)
--           vX = ((z2 + z3) * vKOld) / z3
--           vOut = (((z2 * z4 + z1 * z4 + z1 * z2) * vXOld) - z1 * z4 * vKOld - z2 * z4 * vIn) / (z1 * z2)
--       in if abs((realPart <$> vOutOld) - (realPart <$> vOut)) <= eps
--           && abs((realPart <$> vKOld) - (realPart <$> vK)) <= eps
--           && abs((realPart <$> vXOld) - (realPart <$> vX)) <= eps
--           then SignalThreeNodes (vX, vK, vOut)
--           else f (vX, vK, vOut))

ampOpInvertingTest :: AmpOp -> Impedance -> Impedance -> SignalState -> Circuit Input SignalState
ampOpInvertingTest model r1 r2 (SignalTwoNodes initial) =
    ($ initial) <$> mfix
    (\f -> Circuit $ \vIn -> Signal $ \time (vXOld, vOutOld) ->
        let vX = (- vOutOld) / openLoopGain model
            vOut = ((vXOld * (r1 + r2)) - (vIn * r2)) / r1
        in if abs((realPart <$> vOutOld) - (realPart <$> vOut)) <= eps && abs((realPart <$> vXOld) - (realPart <$> vX)) <= eps
            then SignalTwoNodes (vX, vOut)
            else f (vX, vOut))

-- Calculating the output given a fixed input, to demonstrate the fix point only across the iteration axis
ampOpInverting' :: AmpOp -> Impedance -> Impedance -> Input -> State -> Either String Output
ampOpInverting' _ _ _ _ (OneNode _) = Left "Wrong type for initial state!"
ampOpInverting' model r1 r2 input (TwoNodes initial) = Right $ fix calculate initial
    where calculate f (vXOld, vOutOld) = if abs(realPart vOutOld - realPart vOut) <= 0.000001 && abs(realPart vXOld - realPart vX) <= 0.000001 then vOut else f (vX, vOut)
            where vX = (- vOutOld) / (openLoopGain model `at` 0)
                  vOut = (vXOld * (r1 + r2) `at` 0 - (input * r2 `at` 0)) / r1 `at` 0


-- might not work if module is not installed, possible fix is to place the file in
-- hnn-master root folder

import AI.HNN.FF.ComplexNetwork
import Numeric.LinearAlgebra
import Data.Complex
import System.Random.MWC
import Foreign.Storable           (Storable)

-- actually ignoring imaginary part, adding constant +0i to all real numbers
samples :: Samples (Complex Double)
samples = [
  (fromList [ 1:+0, 1:+0, 1:+0, 1:+0, 1:+0
            , 0:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 0:+0, 0:+0, 1:+0, 1:+0, 1:+0
            , 0:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 1:+0, 1:+0, 1:+0, 1:+0, 1:+0 ], fromList [1:+0]),   -- three

  (fromList [ 1:+0, 1:+0, 1:+0, 1:+0, 1:+0
            , 0:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 1:+0, 1:+0, 1:+0, 1:+0, 1:+0
            , 0:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 1:+0, 1:+0, 1:+0, 1:+0, 1:+0 ], fromList [1:+0]),   -- three

  (fromList [ 0:+0, 1:+0, 1:+0, 1:+0, 1:+0
            , 0:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 0:+0, 0:+0, 0:+0, 1:+0, 1:+0
            , 0:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 0:+0, 1:+0, 1:+0, 1:+0, 1:+0 ], fromList [1:+0]),   -- three

  (fromList [ 1:+0, 1:+0, 1:+0, 1:+0, 1:+0
            , 1:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 1:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 1:+0, 0:+0, 0:+0, 0:+0, 1:+0
            , 1:+0, 1:+0, 1:+0, 1:+0, 1:+0 ], fromList [0:+0]),   -- not a three

  (fromList [ 1:+0, 1:+0, 1:+0, 1:+0, 1:+0
            , 1:+0, 0:+0, 0:+0, 0:+0, 0:+0
            , 1:+0, 0:+0, 0:+0, 0:+0, 0:+0
            , 1:+0, 0:+0, 0:+0, 0:+0, 0:+0
            , 1:+0, 0:+0, 0:+0, 0:+0, 0:+0 ], fromList [0:+0]),   -- not a three

  (fromList [ 0:+0, 1:+0, 1:+0, 1:+0, 0:+0
            , 0:+0, 1:+0, 0:+0, 1:+0, 0:+0
            , 0:+0, 1:+0, 1:+0, 1:+0, 0:+0
            , 0:+0, 1:+0, 0:+0, 1:+0, 0:+0
            , 0:+0, 1:+0, 1:+0, 1:+0, 0:+0 ], fromList [0:+0]),   -- not a three

  (fromList [ 0:+0, 0:+0, 1:+0, 0:+0, 0:+0
            , 0:+0, 1:+0, 1:+0, 0:+0, 0:+0
            , 1:+0, 0:+0, 1:+0, 0:+0, 0:+0
            , 0:+0, 0:+0, 1:+0, 0:+0, 0:+0
            , 0:+0, 0:+0, 1:+0, 0:+0, 0:+0 ], fromList [0:+0]) ]  -- not a three

main :: IO ()
main = do
  n <- createComplexNetwork 25 [250] 1
  let n' = trainNTimes 100 0.5 complexSigmoid complexSigmoid' n samples
  mapM_ (putStrLn . show . output n' complexSigmoid . fst) samples
  putStrLn "-------------"
  putStrLn . show . output n' complexSigmoid $ testInput

  where testInput = fromList [ 0:+0, 0:+0, 1:+0, 1:+0, 1:+0
                             , 0:+0, 0:+0, 0:+0, 0:+0, 1:+0
                             , 0:+0, 0:+0, 1:+0, 1:+0, 1:+0
                             , 0:+0, 0:+0, 0:+0, 0:+0, 1:+0
                             , 0:+0, 0:+0, 1:+0, 1:+0, 1:+0 ]


module Main where

import Control.Applicative
import Graphics.Rendering.Chart.Backend.Diagrams (toFile)
import qualified Graphics.Rendering.Chart.Easy as Plt
import Kalman
import Numeric.LinearAlgebra
import Sim

vanderPol = toFile Plt.def "VanderPol.svg" $ do
  let x0 = [1.4, 0]
  let q = diagl (replicate 2 0.0001)
  let r = diagl (replicate 2 0.5)
  let s0 = Estimate (vector [0, 5]) (diagl (replicate 2 5))

  let plant = Plant fA2 id fH2 id
  let sim = take 500 $ simulate plant (vector x0) 0
  let measurements = map snd sim
  let trajectory = map fst sim

  let sys = System fA2 fH2 q r
  let ekfs = scanl (ekf sys) s0 measurements
  let ukfs = scanl (ukf sys) s0 measurements

  let t1 = linspace 500 (0, 50 :: Double)
  let toPoints v = oneList (toList t1) $ map ((\[x] -> x) . toList . subVector 1 1) v

  Plt.plot (Plt.line "Plant" [toPoints trajectory])
  Plt.plot (Plt.line "UKF" [toPoints $ map sX ukfs])
  Plt.plot (Plt.line "EKF" [toPoints $ map sX ekfs])
  Plt.plot (Plt.line "Measurements" [toPoints measurements])

-- Plt.plot (Plt.line "LPF" [toPoints (movingAverage measurements)])

fA2 :: (Floating a) => [a] -> [a]
fA2 [x1, x2] = [x1 - x2 * dt, x2 + (- u * (1 - x1 ** 2) * x2 + x1) * dt]
  where
    u = 0.2
    dt = 0.1

fH2 :: (Floating a) => [a] -> [a]
fH2 [x1, x2] = [x1, x2]

oneList :: [a] -> [b] -> [(a, b)]
oneList [] _ = []
oneList (x : xs) (y : ys) = (x, y) : oneList xs ys

main = vanderPol


-----------------------------------------------------------------------------
-- |
-- Module      :  ForSyDe.Atom.Skel.Vector.DSP
-- Copyright   :  (c) George Ungureanu, KTH/EECS/ESY 2019-2020
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  ugeorge@kth.se
-- Stability   :  experimental
-- Portability :  portable
--
-- This module exports a set of 'Vector' patterns commonly used in signal processing
-- designs.
-----------------------------------------------------------------------------
module ForSyDe.Atom.Skel.Vector.DSP where

import Data.Complex
import qualified Data.Number.FixedFunctions as F
import ForSyDe.Atom.MoC as MoC
import ForSyDe.Atom.Skel.Vector as V hiding (duals, unduals)
import ForSyDe.Atom.Skel.Vector.Matrix as M
import ForSyDe.Atom ((><))

-- | Return the Taylor window. The Taylor window allows for a selectable sidelobe
-- suppression with a minimum broadening. This window is commonly used in radar
-- processing [1].  Code inspired from a <https://github.com/scipy/scipy/pull/8032 pull request for SciPy>.
--
-- Reference:
--
-- #ref1# [1] W. Carrara, R. Goodman, and R. Majewski "Spotlight Synthetic Aperture
-- Radar: Signal Processing Algorithms" Pages 512-513, July 1995.
taylor :: (Eq a, Floating a)
       => Int
       -- ^ number of points in the output window.
       -> Int
       -- ^ Number of nearly constant level sidelobes adjacent to the mainlobe
       -> a
       -- ^ Desired peak sidelobe level in decibels (db) relative to the mainlobe
       -> Vector a
       -- ^ The window, with the center value normalized to one (the value one appears
       -- only if the number of samples is odd).
taylor n nbar level = V.farm11 (*scale) w
  where
    -- explicit conversions to floating
    bN   = fromIntegral n
    nBar = fromIntegral nbar
    ma   = map fromIntegral [1 .. nbar-1]
    -- calculate intermediate values
    b  = 10**((-level) / 20)
    a  = log(b + sqrt(b**2 - 1)) / pi
    s2 = nBar ** 2 / (a**2 + (nBar - 0.5)**2)
    -- functions for calculating coefficients
    fmcalc m  = let numer = (-1)**(m+1) * product[1.0 - m**2/s2/(a**2 + (j - 0.5)**2) | j <- ma]
                    denom = 2 * product[1 - m**2/j**2 | j <- ma, j /= m]
                in numer / denom
    ccalc m x = cos(2 * pi * x * (m - bN/2 + 1/2) / bN)
    wcalc m   = 2 * dotvv (vector $ map fmcalc ma) (vector $ map (ccalc m) ma) + 1
    -- calculate window coefficients
    w  = vector $ map (wcalc . fromIntegral) [0..n-1]
    -- normalize (Note that this is not described in the original text [1])
    scale = 1 / wcalc (bN - 1) / 2

-- (*^*) :: (RealFrac a, RealFrac b) => a -> b -> a
-- x *^* y = realToFrac $ realToFrac x ** realToFrac y

-- | Returns a Taylor window with default arguments: 4 sidelobes, and peak sidelobe level of -30dB
taylor' n = taylor 4 (-30)

-- | Calculates the dot product between two vectors.
dotvv :: Num a => Vector a -> Vector a -> a
dotvv a b
  | V.length a == V.length b = V.reduce (+) (V.farm21 (*) a b)
  | otherwise                = error "Vector sizes must match"

-- | Calculates the dot product between a vector and a matrix
dotvm :: Num a => Vector a -> Vector (Vector a) -> Vector a
dotvm = dotvm' (+) (*)

-- | Higher order version of 'dotvv'. Applies dot product on vectors of /structured/
-- /types/, e.g. 'Signal's. See 'dotvv'.
dotvv' :: Num a
     => ((a -> a -> a) -> f a -> f a -> f a)
     -- ^ higher-order function that can wrap the (*) operation.
     -> Vector (f a) -> Vector (f a) -> f a
dotvv' wrap a b
  | V.length a == V.length b = V.reduce (wrap (+)) $ V.farm21 (wrap (*)) a b
  | otherwise                = error "Vector sizes must match"

-- | Higher order version of 'dotvm'. Implements the template for a dot operation
-- between a vector and a matrix.
--
-- >>> let mA = vector [vector[1,-1,1],  vector[1,-1,-1],  vector[1,1,-1],  vector[1,1,1]]
-- >>> let y  = vector[1,0,0,0]
-- >>> dotVecMat (+) (*) mA y
-- <1,1,1,1>
dotvm' :: (b -> b -> b)     -- ^ kernel function for a row/column reduction, e.g. @(+)@ for dot product
       -> (a -> b -> b)     -- ^ binary operation for pair-wise elements, e.g. @(*)@ for dot product
       -> Vector a          -- ^ /length/ = @xa@
       -> Vector (Vector b) -- ^ /size/ = @(xa,ya)@
       -> Vector b          -- ^ /length/ = @ya@
dotvm' f g vs = V.reduce (V.farm21 f) . V.farm21 (\x -> V.farm11 (g x)) vs

-- | Higher order pattern implementing the template for a dot operation between a
-- vector and a matrix.
--
-- >>> let mA = vector [vector[1,-1,1],  vector[1,-1,-1],  vector[1,1,-1],  vector[1,1,1]]
-- >>> let y  = vector[1,0,0,0]
-- >>> dotVecMat (+) (*) mA y
-- <1,1,1,1>
dotmv' :: (a -> a -> a)
       -- ^ kernel function for a row/column reduction, e.g. @(+)@ for dot product
       -> (b -> a -> a)
       -- ^ binary operation for pair-wise elements, e.g. @(*)@ for dot product
       -> Vector (Vector b) -- ^ /size/ = @(xa,ya)@
       -> Vector a      -- ^ /length/ = @xa@
       -> Vector a      -- ^ /length/ = @xa@
dotmv' f g mA y = V.farm11 (\x -> V.reduce f $ V.farm21 g x y) mA

-- | Compute a Hanning window.
--
-- Inspired from <https://hackage.haskell.org/package/sdr-0.1.0.6/src/hs_sources/SDR/FilterDesign.hs>
hanning :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
hanning size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.5 * (1 - cos((2 * pi * i) / (n - 1)))

-- | Compute a Hamming window.
--
-- Inspired from <https://hackage.haskell.org/package/sdr-0.1.0.6/src/hs_sources/SDR/FilterDesign.hs>
hamming :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
hamming size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.54 - 0.46 * cos((2 * pi * i) / (n - 1))

-- | Compute a Blackman window.
--
-- Inspired from <https://hackage.haskell.org/package/sdr-0.1.0.6/src/hs_sources/SDR/FilterDesign.hs>
blackman :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
blackman size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.42 - 0.5 * cos((2 * pi * i) / (n - 1)) + 0.08 * cos((4 * pi * i) / (n - 1))

-- | Moving average filter (FIR) on numbers, applied in reverse order (more
-- optimized).
--
-- >>> let v = vector [0,0,0,0,0,1]
-- >>> let c = vector [1,2,3]
-- >>> fir c
-- <0,0,0,3,2,1>
fir :: Num a
    => Vector a  -- ^ vector of coefficients
    -> Vector a  -- ^ input vector of numbers; /size/ = @n@
    -> Vector a  -- ^ output vector of numbers; /size/ = @n@
fir coefs = V.reverse . V.farm11 applyFilter . tails . V.reverse
  where
    applyFilter = V.reduce (+) . V.farm21 (*) coefs

-- | Higher order version of 'fir'. Can create a systolic FIR process network. See
-- <ForSyDe-Atom.html#ungureanu20a [Ungureanu20a]> for a discussion on the relation
-- between 'fir' and 'fir''.
--
-- >>> let c = vector [1,2,3]
-- >>> let s = SY.signal [1,0,0,0,0,0,0,0]
-- >>> fir' (SY.comb21 (+)) (\c -> SY.comb11 (*c)) (SY.delay 0) c s
-- {3,2,1,0,0,0,0,0}
fir' :: (a -> a -> a)  -- ^ process/operation replacing '+'
     -> (c -> a -> a)  -- ^ process/operation replacing '*'
     -> (a -> a)       -- ^ delay process
     -> Vector c       -- ^ vector of coefficients
     -> a              -- ^ input signal/structure
     -> a              -- ^ output signal/structure
fir' plus times delay coefs =
  -- V.reduce plus . V.farm21 (\c -> times c) coefs . V.recur (V.fanoutn n delay <: id)
  V.reduce plus . V.farm21 times coefs . V.recuri (V.fanoutn n delay)
  where n = V.length coefs - 1

-- | generates the "twiddle" coefficients for a FFT network.
twiddles :: Floating a => Int -> V.Vector (Complex a)
twiddles bN = (bitrev . V.take (bN `div` 2)) (V.farm11 bW $ vector [0..])
  where bW x = (cis . negate) (-2 * pi * fromIntegral x / fromIntegral bN)

-- | Radix-2 decimation-in-frequecy Fast Fourier Transform, applied on numbers.  For
-- the difference between DIT- and DIF-FFT, see
-- <https://www.slideshare.net/chappidi_saritha/decimation-in-time-and-frequency>
fft :: RealFloat a
    => Int -- ^ number of FFT stages (@== log_2@ of /length/ of input vector)
    -> V.Vector (Complex a) -> V.Vector (Complex a)
fft k vs | n == 2^k = V.reverse $ bitrev $ (stage `V.pipe1` V.iterate k (*2) 2) vs
         | otherwise = error $ "n=" ++ show n ++ "; k=" ++ show k
  where
    stage   w = V.concat . V.farm21 segment (twiddles n) . V.group w
    segment t = (><) unduals . (><) (V.farm22 (butterfly t)) . duals
    n         = V.length vs        -- length of input
    -------------------------------------------------
    butterfly w x0 x1 = let t = w * x1 in (x0 + t, x0 - t) -- kernel function

-- | Higher order version of 'fft'. Takes the "butterfly" function as argument.
--
-- > butterfly w x0 x1 = let t = w * x1 in (x0 + t, x0 - t)
-- > fft' butterfly === fft
fft' :: Floating a
     => (Complex a -> a -> a -> (a, a)) -- ^ "butterfly" function.
     -> Int  -- ^ number of FFT stages (@== log_2@ of /length/ of input vector)
     -> V.Vector a -> V.Vector a
fft' butterfly k vs | n == 2^k = bitrev $ (stage `V.pipe1` (V.iterate k (*2) 2)) vs
  where
    stage   w = V.concat . V.farm21 segment (twiddles n) . V.group w
    segment t = (><) unduals . (><) (V.farm22 (butterfly t)) . duals
    n         = V.length vs        -- length of input

-- | splits a vector in two equal parts.
--
-- >>> duals $ vector [1,2,3,4,5,6,7]
-- (<1,2,3>,<4,5,6>)
duals    :: Vector a -> (Vector a, Vector a)
duals v  = (V.take k v, V.drop k v)
  where k = V.length v `div` 2

-- | concatenates a previously split vector. See also 'duals'
unduals  :: Vector a -> Vector a -> Vector a
unduals x y =  x <++> y

-- | performs a bit-reverse permutation.
--
-- <<fig/skel-vector-comm-bitrev.png>>
-- <<fig/skel-vector-comm-bitrev-net.png>>
--
-- >>> bitrev $ vector ["000","001","010","011","100","101","110","111"]
-- <"111","011","101","001","110","010","100","000">
bitrev (x:>Null) = V.unit x
bitrev xs        = bitrev (evens xs) <++> bitrev (odds xs)

-- duals   v = let k = length v `div` 2
--             in  S.farm22 (,) (take k v) (drop k v)

-- unduals = (<++>)

-- fft n = V.vector . map (fmap realToFrac) . Sh.fromVector . ShDSP.fft (2^n) . Sh.vector . map (fmap realToFrac) . V.fromVector


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeOperators #-}

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

{-|
Module      : Numeric.LinearAlgebra.MatrixFree.Static
Description : Compose matrix-free linear maps
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental

Enables the user to compose matrix-free `LinearMap` objects in the same way
that one would compose regular matrices.
-}
module Numeric.LinearAlgebra.MatrixFree.Combinators
    ( (===)
    , (|||)
    , blockMatrix
    , blockDiag
    , LinearMapList(..)
    , applyLinearMapList
    )
where

import           GHC.TypeLits
import           Numeric.LinearAlgebra          ( Transposable(..) )
import           Numeric.LinearAlgebra.Static   ( (#)
                                                , split
                                                , R
                                                )
import           Numeric.LinearAlgebra.MatrixFree
                                                ( LinearMap(..), (##>))
import           Numeric.LinearAlgebra.Static.RList (Sum, RList(..))

-- Stacking matrices either horizontally or vertically are linked through a
-- transpose, so the underlying functions are shared between the two
-- stacking cases. Only in a switched order.
stackLinearMaps :: (KnownNat m, KnownNat n)
                => LinearMap m p -> LinearMap n p -> R p -> R (m + n)
stackLinearMaps la lb v = la ##> v # lb ##> v

addLinearMaps :: (KnownNat n, KnownNat p)
              => LinearMap m n -> LinearMap m p -> R (n + p) -> R m
addLinearMaps la lb v = let (v1, v2) = split v in la ##> v1 + lb ##> v2

-- | Vertically concatenate two `LinearMap`s \(\begin{bmatrix}A\\B\end{bmatrix}\)
(===) :: (KnownNat m, KnownNat n, KnownNat p)
      => LinearMap m n -> LinearMap p n -> LinearMap (m + p) n
(===) la lb = LinearMap (stackLinearMaps la lb) (addLinearMaps (tr la) (tr lb))

infixl 2 ===

-- | Horizontally concatenate two `LinearMap`s \(\begin{bmatrix}A&B\end{bmatrix}\)
(|||) :: (KnownNat m, KnownNat n, KnownNat p)
      => LinearMap m n -> LinearMap m p -> LinearMap m (n + p)
(|||) la lb = tr (tr la === tr lb)

infixl 3 |||

-- | Generate a block matrix from four `LinearMap`s
-- \[
-- \begin{bmatrix}
--   A & B \\
--   C & D
-- \end{bmatrix}
-- \]
blockMatrix
    :: (KnownNat m, KnownNat n, KnownNat p, KnownNat q)
    => LinearMap m n             -- ^ Upper left matrix
    -> LinearMap m p             -- ^ Upper right matrix
    -> LinearMap q n             -- ^ Lower left matrix
    -> LinearMap q p             -- ^ Lower right matrix
    -> LinearMap (m + q) (n + p) -- ^ Resulting block matrix
blockMatrix a b c d = a ||| b === c ||| d

-- | A non-empty list of "LinearMap" matrices of different sizes.
data LinearMapList (ms :: [Nat]) (ns :: [Nat]) where
    LOne :: (KnownNat m, KnownNat n) => LinearMap m n -> LinearMapList '[m] '[n]
    (:-:) :: (KnownNat m, KnownNat n) => LinearMap m n -> LinearMapList ms ns
                                      -> LinearMapList (m ': ms) (n ': ns)

-- | Concatenate a "LinearMapList" of "LinearMap"s as a block diagonal.
-- Note that this has poor performance; it runs `split` on the
-- input vector O(n) times, where n is the number of entries in the list.
blockDiag :: LinearMapList ms ns -> LinearMap (Sum ms) (Sum ns)
blockDiag ls = case blockDiagWorker ls of
  LWorker lsW -> lsW

data LWorker m n where
  LWorker :: (KnownNat m, KnownNat n) => LinearMap m n -> LWorker m n

blockDiagWorker :: LinearMapList ms ns -> LWorker (Sum ms) (Sum ns)
blockDiagWorker (LOne l) = LWorker l
blockDiagWorker (l :-: ls) = case blockDiagWorker ls of
  LWorker lsW -> LWorker (diag2 l lsW)

-- Combine two `LinearMap`s on a diagonal without having to use
-- konst 0, which would just be a waste of processing time (sum a list and
-- multiply it by zero)
diag2 :: (KnownNat m, KnownNat n, KnownNat p, KnownNat q)
      => LinearMap m n
      -> LinearMap p q
      -> LinearMap (m + p) (n + q)
diag2 la lb = LinearMap (f la lb) (f (tr la) (tr lb))
  where
    f :: (KnownNat m', KnownNat n', KnownNat p', KnownNat q')
      => LinearMap m' n' -> LinearMap p' q' -> R (n' + q') -> R (m' + p')
    f l1 l2 v = let (v1, v2) = split v in
      (l1 ##> v1) # (l2 ##> v2)

-- | Apply a list of "LinearMap"s over a list of static vectors
applyLinearMapList :: LinearMapList ms ns -> RList ns -> RList ms
applyLinearMapList (LOne _) (_ ::: _) = error "This case is impossible"
applyLinearMapList (_ :-: _) (ROne _) = error "This case is impossible"
applyLinearMapList (LOne l) (ROne r) = ROne (l ##> r)
applyLinearMapList (l :-: ls) (r ::: rs) = (l ##> r) ::: applyLinearMapList ls rs


module Statistics.Information.Continuous.MutualInfo where

import Data.Matrix
import Numeric.SpecFunctions
import Statistics.Information.Utils.KNN
import Statistics.Information.Utils.Random
import System.Random

mi :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double -> Double
mi g k base xs ys | (k < nrows xs) && (nrows xs == nrows ys) =
  (-a - b + c + d) / (log $ fromIntegral base)
  where
    (g1, g2) = split g
    noisyxs = noisy g1 xs
    noisyys = noisy g2 ys
    joint = noisyxs <|> noisyys
    distances = knnDistances chebyshevDistance (toLists joint) k
    a = avgdigamma chebyshevDistance noisyxs distances
    b = avgdigamma chebyshevDistance noisyys distances
    c = digamma (fromIntegral k)
    d = digamma . fromIntegral . nrows $ xs

cmi :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
       Matrix Double -> Double
cmi g k base xs ys zs | (k < nrows xs) && (nrows xs == nrows ys) =
  (-a - b + c + d) / (log $ fromIntegral base)
  where
    [g1, g2, g3] = splitN 3 g
    noisyxs = noisy g1 xs
    noisyys = noisy g2 ys
    noisyzs = noisy g3 zs
    joint = noisyxs <|> noisyys <|> noisyzs
    distances = knnDistances chebyshevDistance (toLists joint) k
    a = avgdigamma chebyshevDistance (noisyxs <|> noisyzs) distances
    b = avgdigamma chebyshevDistance (noisyys <|> noisyzs) distances
    c = avgdigamma chebyshevDistance noisyzs distances
    d = digamma (fromIntegral k)


{-# LANGUAGE StandaloneDeriving #-}
module Examples.Quantum where

import LAoP.Utils
import LAoP.Matrix.Type
import Data.Complex
import Prelude hiding (id, (.))

deriving instance Ord a => Ord (Complex a)

xor :: (Bool, Bool) -> Bool
xor (False, b) = b
xor (True, b) = not b

ker :: Num e => Matrix e a b -> Matrix e a a
ker m = tr m . m

cnot :: Matrix (Complex Double) (Bool, Bool) (Bool, Bool)
cnot = kr fstM (fromF xor)

ccnot :: (Num e, Ord e) => Matrix e ((Bool, Bool), Bool) ((Bool, Bool), Bool)
ccnot = kr fstM (fromF f)
  where
    f = xor . tp (uncurry (&&)) id
    tp f g (a,b) = (f a, g b)

had :: Matrix (Complex Double) Bool Bool
had = (1/sqrt 2) .| fromLists [[1, 1], [1, -1]]

bell :: Matrix (Complex Double) (Bool, Bool) (Bool, Bool)
bell = cnot . (had >< iden)


module Main where

import qualified Numeric.LinearAlgebra as LA
import qualified MachineLearning.Types as T
import qualified MachineLearning as ML
import qualified MachineLearning.Classification.Binary as CB
import qualified MachineLearning.Classification.OneVsAll as OVA
import qualified MachineLearning.TerminalProgress as TP
import qualified MachineLearning.Optimization as Opt
import MachineLearning.MultiSvmClassifier
import MachineLearning.SoftmaxClassifier

featuresMapParameter = 2

-- processFeatures :: T.Matrix -> T.Matrix
processFeatures muSigma = ML.addBiasDimension . ML.featureNormalization muSigma . ML.mapFeatures featuresMapParameter

calcAccuracy :: (Model a) => a -> T.Matrix -> T.Vector -> T.Vector -> Double
calcAccuracy m x y theta = OVA.calcAccuracy y yPredicted
  where yPredicted = hypothesis m x theta

main = do
  putStrLn "\n=== Optical Recognition of Handwritten Digits Data Set ===\n"
  -- Step 1. Data loading.
  -- Step 1.1 Training Data loading.
  (x, y) <- pure ML.splitToXY <*> LA.loadMatrix "digits_classification/optdigits.tra"
  -- Step 1.1 Testing Data loading.
  (xTest, yTest) <- pure ML.splitToXY <*> LA.loadMatrix "digits_classification/optdigits.tes"

  -- Step 2. Outputs and features preprocessing.
  let numLabels = 10
      svm = MultiClass $ MultiSvm 1 numLabels
      softmax = MultiClass $ Softmax numLabels
      muSigma = ML.meanStddev (ML.mapFeatures featuresMapParameter x)
      x1 = processFeatures muSigma x
      xTest1 = processFeatures muSigma xTest
      initialTheta = LA.konst 0.1 (numLabels*(LA.cols x1))
  print $ LA.size x1
  print $ LA.size $ x1 LA.? [1..10]
  -- Step 3. Learning.
  putStrLn "Learning Multi SVM model..."
  (thetaSvm, optPathSvm) <- TP.learnWithProgressBar (Opt.minimize (Opt.BFGS2 0.01 0.7) svm 0.0001 5 (CB.L2 2) x1 y) initialTheta 20

  putStrLn "\nLearning MGD Softmax model"
  --(thetaSm, optPathSm) <- TP.learnWithProgressBar (Opt.minimize (Opt.BFGS2 0.01 0.7) softmax 0.000001 5 1 x1 y) initialTheta 20
  (thetaSm, optPathSm) <- TP.learnWithProgressBar (Opt.minimize (Opt.MinibatchGradientDescent 0 1024 0.01) svm 0.000001 200 (CB.L2 1) x1 y) initialTheta 20

  putStrLn "\nLearning GD Softmax model"
  (thetaSmGD, optPathSmGD) <- TP.learnWithProgressBar (Opt.minimize (Opt.GradientDescent 0.01) softmax 0.000001 50 (CB.L2 1) x1 y) initialTheta 20

  -- Step 4. Prediction and checking accuracy
  let accuracyTrainSvm = calcAccuracy svm x1 y thetaSvm
      accuracyTestSvm = calcAccuracy svm xTest1 yTest thetaSvm

  -- Step 5. Printing results.
  putStrLn $ "\nSVM: Number of iterations to learn: " ++ show (LA.rows optPathSvm)

  putStrLn $ "SVM: Accuracy on train set (%): " ++ show (accuracyTrainSvm*100)
  putStrLn $ "SVM: Accuracy on test set (%): " ++ show (accuracyTestSvm*100)

  let accuracyTrainSm = calcAccuracy softmax x1 y thetaSm
      accuracyTestSm = calcAccuracy softmax xTest1 yTest thetaSm

  -- Step 5. Printing results.
  putStrLn $ "\nSoftmax: Number of iterations to learn: " ++ show (LA.rows optPathSm)

  putStrLn $ "Softmax: Accuracy on train set (%): " ++ show (accuracyTrainSm*100)
  putStrLn $ "Softmax: Accuracy on test set (%): " ++ show (accuracyTestSm*100)

  let accuracyTrainSmGD = calcAccuracy softmax x1 y thetaSmGD
      accuracyTestSmGD = calcAccuracy softmax xTest1 yTest thetaSmGD

  -- Step 5. Printing results.
  putStrLn $ "\nGD Softmax: Number of iterations to learn: " ++ show (LA.rows optPathSmGD)

  putStrLn $ "GD Softmax: Accuracy on train set (%): " ++ show (accuracyTrainSmGD*100)
  putStrLn $ "GD Softmax: Accuracy on test set (%): " ++ show (accuracyTestSmGD*100)


{-# LANGUAGE BangPatterns , RankNTypes, GADTs, DataKinds #-}

{- | The 'Numerical.HBLAS.BLAS.Level1' module provides a fully general
yet type safe Level1 BLAS API.

When in doubt about the semantics of an operation,
consult your system's BLAS api documentation, or just read the documentation
for
<https://software.intel.com/sites/products/documentation/hpc/mkl/mklman/index.htm the Intel MKL BLAS distribution>

A few basic notes about how to invoke BLAS routines.

Many BLAS operations take one or more arguments of type 'Transpose'.
'Tranpose' has the following different constructors, which tell BLAS
routines what transformation to implicitly apply to an input matrix @mat@ with dimension @n x m@.

*  'NoTranspose' leaves the matrix @mat@ as is.

* 'Transpose' treats the @mat@ as being implicitly transposed, with dimension
    @m x n@. Entry @mat(i,j)@ being treated as actually being the entry
    @mat(j,i)@. For Real matrices this is also the matrix adjoint operation.
    ie @Tranpose(mat)(i,j)=mat(j,i)@

*  'ConjNoTranspose' will implicitly conjugate @mat@, which is a no op for Real ('Float' or 'Double') matrices, but for
'Complex Float' and 'Complex Double' matrices, a given matrix entry @mat(i,j)==x':+'y@
will be treated as actually being  @conjugate(mat)(i,j)=y':+'x@.

* 'ConjTranpose' will implicitly transpose and conjugate the input matrix.
ConjugateTranpose acts as matrix adjoint for both real and complex matrices.

The *gemm operations  work as follows (using 'sgemm' as an example):

* @'sgemm trLeft trRight alpha beta left right result'@, where @trLeft@ and @trRight@
are values of type 'Transpose' that respectively act on the matrices @left@ and @right@.

* the generalized matrix computation thusly formed can be viewed as being
@result = alpha * trLeft(left) * trRight(right) + beta * result@

the *gemv operations are akin to the *gemm operations, but with @right@ and @result@
being vectors rather than matrices.

the *trsv operations solve for @x@ in the equation @A x = y@ given @A@ and @y@.
The 'MatUpLo' argument determines if the matrix should be treated as upper or
lower triangular and 'MatDiag' determines if the triangular solver should treat
the diagonal of the matrix as being all 1's or not.  A general pattern of invocation
would be @'strsv' matuplo  tranposeMatA  matdiag  matrixA  xVector@.
A key detail to note is that the input vector is ALSO the result vector,
ie 'strsv' and friends updates the vector place.

-}

module Numerical.HBLAS.BLAS.Level1(
        sasum
        ,dasum
        ,scasum
        ,dzasum

        ,saxpy
        ,daxpy
        ,caxpy
        ,zaxpy

        ,scopy
        ,dcopy
        ,ccopy
        ,zcopy

        ,sdot
        ,ddot
        --,sdsdot
        ,dsdot

        ,cdotu
        ,cdotc
        ,zdotu
        ,zdotc

        ,snrm2
        ,dnrm2
        ,scnrm2
        ,dznrm2

        ,srot
        ,drot

        --,srotg
        --,drotg

        ,srotm
        ,drotm

        --,srotmg
        --,drotmg

        ,sscal
        ,dscal
        ,cscal
        ,zscal
        ,csscal
        ,zdscal

        ,sswap
        ,dswap
        ,cswap
        ,zswap

        ,isamax
        ,idamax
        ,icamax
        ,izamax

            ) where

import Numerical.HBLAS.UtilsFFI
import Numerical.HBLAS.BLAS.FFI.Level1
import Numerical.HBLAS.BLAS.Internal.Level1
import Control.Monad.Primitive
import Data.Complex

sasum :: PrimMonad m => AsumFun Float (PrimState m) m Float
sasum = asumAbstraction "sasum" cblas_sasum_safe cblas_sasum_unsafe

dasum :: PrimMonad m => AsumFun Double (PrimState m) m Double
dasum = asumAbstraction "dasum" cblas_dasum_safe cblas_dasum_unsafe

scasum :: PrimMonad m => AsumFun (Complex Float) (PrimState m) m Float
scasum = asumAbstraction "scasum" cblas_scasum_safe cblas_scasum_unsafe

dzasum :: PrimMonad m => AsumFun (Complex Double) (PrimState m) m Double
dzasum = asumAbstraction "dzasum" cblas_dzasum_safe cblas_dzasum_unsafe

saxpy :: PrimMonad m => AxpyFun Float (PrimState m) m
saxpy = axpyAbstraction "saxpy" cblas_saxpy_safe cblas_saxpy_unsafe (\x f -> f x)

daxpy :: PrimMonad m => AxpyFun Double (PrimState m) m
daxpy = axpyAbstraction "daxpy" cblas_daxpy_safe cblas_daxpy_unsafe (\x f -> f x)

caxpy :: PrimMonad m => AxpyFun (Complex Float) (PrimState m) m
caxpy = axpyAbstraction "caxpy" cblas_caxpy_safe cblas_caxpy_unsafe withRStorable_

zaxpy :: PrimMonad m => AxpyFun (Complex Double) (PrimState m) m
zaxpy = axpyAbstraction "zaxpy" cblas_zaxpy_safe cblas_zaxpy_unsafe withRStorable_

scopy :: PrimMonad m => CopyFun Float (PrimState m) m
scopy = copyAbstraction "scopy" cblas_scopy_safe cblas_scopy_unsafe

dcopy :: PrimMonad m => CopyFun Double (PrimState m) m
dcopy = copyAbstraction "dcopy" cblas_dcopy_safe cblas_dcopy_unsafe

ccopy :: PrimMonad m => CopyFun (Complex Float) (PrimState m) m
ccopy = copyAbstraction "ccopy" cblas_ccopy_safe cblas_ccopy_unsafe

zcopy :: PrimMonad m => CopyFun (Complex Double) (PrimState m) m
zcopy = copyAbstraction "zcopy" cblas_zcopy_safe cblas_zcopy_unsafe

sdot :: PrimMonad m => NoScalarDotFun Float (PrimState m) m Float
sdot = noScalarDotAbstraction "sdot" cblas_sdot_safe cblas_sdot_unsafe

ddot :: PrimMonad m => NoScalarDotFun Double (PrimState m) m Double
ddot = noScalarDotAbstraction "ddot" cblas_ddot_safe cblas_ddot_unsafe

--sdsdot :: PrimMonad m => ScalarDotFun Float (PrimState m) m Float
--sdsdot = scalarDotAbstraction "sdsdot" cblas_sdsdot_safe cblas_sdsdot_unsafe withRStorable withRStorable

dsdot :: PrimMonad m => NoScalarDotFun Float (PrimState m) m Double
dsdot = noScalarDotAbstraction "dsdot" cblas_dsdot_safe cblas_dsdot_unsafe

cdotu :: PrimMonad m => ComplexDotFun (Complex Float) (PrimState m) m
cdotu = complexDotAbstraction "cdotu" cblas_cdotu_safe cblas_cdotu_unsafe

cdotc :: PrimMonad m => ComplexDotFun (Complex Float) (PrimState m) m
cdotc = complexDotAbstraction "cdotc" cblas_cdotc_safe cblas_cdotc_unsafe

zdotu :: PrimMonad m => ComplexDotFun (Complex Double) (PrimState m) m
zdotu = complexDotAbstraction "zdotu" cblas_zdotu_safe cblas_zdotu_unsafe

zdotc :: PrimMonad m => ComplexDotFun (Complex Double) (PrimState m) m
zdotc = complexDotAbstraction "zdotc" cblas_zdotc_safe cblas_zdotc_unsafe

snrm2 :: PrimMonad m => Nrm2Fun Float (PrimState m) m Float
snrm2 = norm2Abstraction "snrm2" cblas_snrm2_safe cblas_snrm2_unsafe

dnrm2 :: PrimMonad m => Nrm2Fun Double (PrimState m) m Double
dnrm2 = norm2Abstraction "dnrm2" cblas_dnrm2_safe cblas_dnrm2_unsafe

scnrm2 :: PrimMonad m => Nrm2Fun (Complex Float) (PrimState m) m Float
scnrm2 = norm2Abstraction "scnrm2" cblas_scnrm2_safe cblas_scnrm2_unsafe

dznrm2 :: PrimMonad m => Nrm2Fun (Complex Double) (PrimState m) m Double
dznrm2 = norm2Abstraction "dznrm2" cblas_dznrm2_safe cblas_dznrm2_unsafe

srot :: PrimMonad m => RotFun Float (PrimState m) m
srot = rotAbstraction "srot" cblas_srot_safe cblas_srot_unsafe

drot :: PrimMonad m => RotFun Double (PrimState m) m
drot = rotAbstraction "drot" cblas_drot_safe cblas_drot_unsafe

--srotg :: PrimMonad m => RotgFun Float (PrimState m) m
--srotg = rotgAbstraction "srotg" cblas_srotg_safe cblas_srotg_unsafe

--drotg :: PrimMonad m => RotgFun Double (PrimState m) m
--drotg = rotgAbstraction "drotg" cblas_drotg_safe cblas_drotg_unsafe

srotm :: PrimMonad m => RotmFun Float (PrimState m) m
srotm = rotmAbstraction "srotm" cblas_srotm_safe cblas_srotm_unsafe

drotm :: PrimMonad m => RotmFun Double (PrimState m) m
drotm = rotmAbstraction "drotm" cblas_drotm_safe cblas_drotm_unsafe

--srotmg :: PrimMonad m => RotmgFun Float (PrimState m) m
--srotmg = rotmgAbstraction "srotmg" cblas_srotmg_safe cblas_srotmg_unsafe

--drotmg :: PrimMonad m => RotmgFun Double (PrimState m) m
--drotmg = rotmgAbstraction "drotmg" cblas_drotmg_safe cblas_drotmg_unsafe

sscal :: PrimMonad m => ScalFun Float Float (PrimState m) m
sscal = scalAbstraction "sscal" cblas_sscal_safe cblas_sscal_unsafe (\x f -> f x )

dscal :: PrimMonad m => ScalFun Double Double (PrimState m) m
dscal = scalAbstraction "dscal" cblas_dscal_safe cblas_dscal_unsafe (\x f -> f x )

cscal :: PrimMonad m => ScalFun (Complex Float) (Complex Float) (PrimState m) m
cscal = scalAbstraction "cscal" cblas_cscal_safe cblas_cscal_unsafe withRStorable_

zscal :: PrimMonad m => ScalFun (Complex Double) (Complex Double) (PrimState m) m
zscal = scalAbstraction "zscal" cblas_zscal_safe cblas_zscal_unsafe withRStorable_

csscal :: PrimMonad m => ScalFun Float (Complex Float) (PrimState m) m
csscal = scalAbstraction "csscal" cblas_csscal_safe cblas_csscal_unsafe (\x f -> f x )

zdscal :: PrimMonad m => ScalFun Double (Complex Double) (PrimState m) m
zdscal = scalAbstraction "zdscal" cblas_zdscal_safe cblas_zdscal_unsafe (\x f -> f x )

sswap :: PrimMonad m => SwapFun Float (PrimState m) m
sswap = swapAbstraction "sswap" cblas_sswap_safe cblas_sswap_unsafe

dswap :: PrimMonad m => SwapFun Double (PrimState m) m
dswap = swapAbstraction "dswap" cblas_dswap_safe cblas_dswap_unsafe

cswap :: PrimMonad m => SwapFun (Complex Float) (PrimState m) m
cswap = swapAbstraction "cswap" cblas_cswap_safe cblas_cswap_unsafe

zswap :: PrimMonad m => SwapFun (Complex Double) (PrimState m) m
zswap = swapAbstraction "zswap" cblas_zswap_safe cblas_zswap_unsafe

isamax :: PrimMonad m => IamaxFun Float (PrimState m) m
isamax = iamaxAbstraction "isamax" cblas_isamax_safe cblas_isamax_unsafe

idamax :: PrimMonad m => IamaxFun Double (PrimState m) m
idamax = iamaxAbstraction "idamax" cblas_idamax_safe cblas_idamax_unsafe

icamax :: PrimMonad m => IamaxFun (Complex Float) (PrimState m) m
icamax = iamaxAbstraction "icamax" cblas_icamax_safe cblas_icamax_unsafe

izamax :: PrimMonad m => IamaxFun (Complex Double)(PrimState m) m
izamax = iamaxAbstraction "izamax" cblas_izamax_safe cblas_izamax_unsafe

{-
isamin :: PrimMonad m => IaminFun Float (PrimState m) m
isamin = iaminAbstraction "isamin" cblas_isamin_safe cblas_isamin_unsafe

idamin :: PrimMonad m => IaminFun Double (PrimState m) m
idamin = iaminAbstraction "idamin" cblas_idamin_safe cblas_idamin_unsafe

icamin :: PrimMonad m => IaminFun (Complex Float) (PrimState m) m
icamin = iaminAbstraction "icamin" cblas_icamin_safe cblas_icamin_unsafe

izamin :: PrimMonad m => IaminFun (Complex Double)(PrimState m) m
izamin = iaminAbstraction "izamin" cblas_izamin_safe cblas_izamin_unsafe
-}


--
-- ActLayer: activate layer
--

module CNN.ActLayer (
  relu
, relu'
, step
, softmax
, activate
, deactivate
, reverseActFunc
, funcToNum
, numToFunc
) where

import Numeric.LinearAlgebra hiding (step)
import CNN.Algebra
import CNN.Image
import CNN.LayerType

{- |
step

  IN : matrix

  OUT: matrix

>>> let m = fromLists [[1.0, 0.0, (-1.1), (-0.5), 2.1]] :: Matrix R
>>> step m
(1><5)
 [ 1.0, 0.0, 0.0, 0.0, 1.0 ]

-}

step :: ActFunc
step = cmap f
  where
    f :: Double -> Double
    f a = if a > 0.0 then 1.0 else 0.0

{- |
relu

  IN : array of Double

  OUT: array of Double

>>> relu $ fromLists [[1.0, 0.0, (-1.1), (-0.5), 2.1]]
(1><5)
 [ 1.0, 0.0, 0.0, 0.0, 2.1 ]

-}

relu :: ActFunc
relu = cmap (`max` 0.0)

{- |
relu'

  IN : array of Double

  OUT: array of Double

>>> relu' $ fromLists [[1.0, 0.0, (-1.1), (-0.5), 2.1]]
(1><5)
 [ 1.0, 0.0, 0.0, 0.0, 1.0 ]

-}

relu' :: ActFunc
relu' = cmap (\x -> if x > 0.0 then 1.0 else 0.0)

{- |
softmax

  IN : array of Double

  OUT: normalized array of Double

>>> let a1 = softmax $ fromLists [[1.0, 1.0, 1.0]]
>>> (a1 `atIndex` (0,0)) * 3.0
1.0
>>> let a2 = softmax $ fromLists [[1.0, 3.0, 2.0]]
>>> (a2 `atIndex` (0,0)) < (a2 `atIndex` (0,1))
True
>>> (a2 `atIndex` (0,0)) < (a2 `atIndex` (0,2))
True
>>> (a2 `atIndex` (0,1)) > (a2 `atIndex` (0,2))
True
-}

softmax :: ActFunc
softmax as = mscale (1.0 / sume) es
  where
    amax = maxElement as
    es   = cmap (\x -> exp (x - amax)) as
    sume = sumElements es

{-
activate

  IN : activation function
       image

  OUT: calculated image

-}

activate :: ActFunc -> Image -> Image
activate f = map f

{- |
deactivate

  IN : activation function
       image
       difference from previous layer

  OUT: difference and updated layer

>>> fst $ deactivate relu [fromLists [[1.5,(-2.0)]]] [fromLists [[0.5,0.1]]]
[(1><2)
 [ 0.5, 0.0 ]]

-}

m0 :: Matrix R
m0 = fromLists [[0.0]]
m1 :: Matrix R
m1 = fromLists [[1.0]]

deactivate :: ActFunc -> Image -> Delta -> (Delta, Maybe Layer)
deactivate f im delta
  | c == m0   = (dl', Nothing)
  | c == m1   = ([] , Nothing)
  | otherwise = ([] , Nothing)
  where
    c = f m0
    f' = map relu' im
    dl' = zipWith (*) delta f'

reverseActFunc :: ActFunc -> Layer
reverseActFunc relu = ActLayer step

{- |
>>> funcToNum step
0
>>> funcToNum relu
1
-}

funcToNum :: ActFunc -> Int
funcToNum f = case f of
  step    -> 0
  relu    -> 1
  relu'   -> 2
  softmax -> 3
  _       -> error ("invalid function")

numToFunc :: Int -> ActFunc
numToFunc i = case i of
  0 -> step
  1 -> relu
  2 -> relu'
  3 -> softmax
  _ -> error ("invalid function number: " ++ (show i))


{-# LANGUAGE RebindableSyntax, TypeSynonymInstances #-}
module DSLsofMath.W08 where
import Prelude hiding (Num(..),(/),(^),Fractional(..),Floating(..),sum)
import DSLsofMath.W05
import qualified DSLsofMath.W06
import DSLsofMath.W06 (Taylor, integ, sinx, sinf, cosx, cosf)
import DSLsofMath.Algebra
import Data.Complex

expx :: Field a => PowerSeries a
expx = integ 1 expx

expf :: Field a => a -> a
expf = evalPS 100 expx

i :: Ring a => Complex a
i = zero :+ one

ex1 :: Field a => Complex a
ex1 = expf i

compScale :: Ring a => a -> PowerSeries a -> PowerSeries a
compScale c (Poly as) = Poly (zipWith (*) as (iterate (c*) 1))

type PSC a = PowerSeries (Complex a)
expix :: Field a => PSC a
expix = compScale i expx

cosxisinx :: Field a => PSC a
cosxisinx = cosx + Poly [i] * sinx

ex2, ex2' :: Field a => PSC a
ex2   = takePoly 8 expix
ex2'  = takePoly 8 cosxisinx

test2 :: Bool
test2 = ex2 == (ex2' :: PSC Rational)

check2 :: Bool
check2 = ex2 == coeff2

coeff2 ::  PSC Rational
coeff2 =   Poly [      1            :+ 0, {-"\qquad"-}  0 :+     1,
                   (-  1  /  2)     :+ 0, {-"\qquad"-}  0 :+ (-  1  /  6),
                       1  /  24     :+ 0, {-"\qquad"-}  0 :+     1  /  120,
                   (-  1  /  720)   :+ 0, {-"\qquad"-}  0 :+ (-  1  /  5040)]

ex2R :: Poly Rational
ex2R = takePoly 8 cosx

ex2I :: Poly Rational
ex2I = takePoly 8 sinx

derivT :: Taylor a -> Taylor a
derivT = tail

integT :: a -> Taylor a -> Taylor a
integT = (:)

val ::  Additive a => Taylor a  ->  a
val (a:_)   =   a
val _       =   zero

expT, sinT, cosT :: Transcendental a => Taylor a -> Taylor a
expT  as  = integT  (exp  (val as))  (expT as   * derivT as)
sinT  as  = integT  (sin  (val as))  (cosT as   * derivT as)
cosT  as  = integT  (cos  (val as))  (-sinT as  * derivT as)

instance Transcendental a => Transcendental (Taylor a) where
   pi   =  [pi]; exp  =  expT;  sin  =  sinT;  cos  =  cosT

instance Field a => MulGroup (Taylor a) where
  recip = recipStream

recipStream :: Field a => Taylor a -> Taylor a
recipStream [] = error "recipStream: divByZero"
recipStream (f0:f') = g
  where  g   = g0:g'
         g0  = recip f0
         g'  = negate (f' * g*g)

expTx :: Taylor REAL
expTx = expT [0,1]
testExpT = take 5 expTx == [1.0, 1.0, 1.0, 1.0, 1.0]

expix3 :: Taylor (Complex Double)
expix3 = exp (i * x)            where  i  = [0 :+ 1]; x  = [0,1]

ex2alt3 = take 10 expix3
testex2alt3 = ex2alt3 ==  take 10 (concat (repeat [1, i, -1, -i]))

cosxisinx3 :: Taylor (Complex Double)
cosxisinx3 = cos x + i * sin x  where  i  = [0 :+ 1];  x  = [0,1]

ex2'alt3 = take 10 cosxisinx3

fs :: Field a => PowerSeries a
fs = integ 1 fs'
   where  fs'   = integ 0 fs''
          fs''  = exp3x + 3 * fs' - 2 * fs
exp3x :: Field a => PowerSeries a
exp3x = compScale 3 expx


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

import Control.Monad
import Control.Monad.Random
import Data.List
import Data.Maybe
import Data.Singletons
import Data.Singletons.Prelude
import Data.Singletons.TypeLits
import Numeric.LinearAlgebra.Static
import System.Environment
import Text.Read

data Weights i o = W { wBiases :: !(R o)
                     , wNodes  :: !(L o i)
                     }                      -- an "o x i" layer

data Network :: Nat -> [Nat] -> Nat -> * where
    O     :: !(Weights i o)
          -> Network i '[] o
    (:&~) :: KnownNat h
          => !(Weights i h)
          -> !(Network h hs o)
          -> Network i (h ': hs) o
infixr 5 :&~

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

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

runLayer :: (KnownNat i, KnownNat o)
         => Weights i o
         -> R i
         -> R o
runLayer (W wB wN) v = wB + wN #> v

runNet :: (KnownNat i, KnownNat o)
       => Network i hs o
       -> R i
       -> R o
runNet = \case
   O w -> \(!v) -> logistic (runLayer w v)
   (w :&~ n') -> \(!v) -> let v' = logistic (runLayer w v)
                          in runNet n' v'

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

randomNet :: forall m i hs o. (MonadRandom m, KnownNat i, SingI hs, KnownNat o)
          => m (Network i hs o)
randomNet = go sing
  where
    go :: forall h hs'. KnownNat h
       => Sing hs'
       -> m (Network h hs' o)
    go = \case
        SNil            ->     O <$> randomWeights
        SNat `SCons` ss -> (:&~) <$> randomWeights <*> go ss

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

netTest :: MonadRandom m => Double -> Int -> m String
netTest rate n = do
    inps <- replicateM n $ do
      s <- getRandom
      return $ randomVector s Uniform * 2 - 1
    let outs = flip map inps $ \v ->
                 if v `inCircle` (fromRational 0.33, 0.33)
                      || v `inCircle` (fromRational (-0.33), 0.33)
                   then fromRational 1
                   else fromRational 0
    net0 :: Network 2 '[16, 8] 1 <- randomNet
    let trained = foldl' trainEach net0 (zip inps outs)
          where
            trainEach :: (KnownNat i, SingI hs, KnownNat o)
                      => Network i hs o
                      -> (R i, R o)
                      -> Network i hs o
            trainEach nt (i, o) = train rate i o nt

        outMat = [ [ render (norm_2 (runNet trained (vector [x / 25 - 1,y / 10 - 1])))
                   | x <- [0..50] ]
                 | y <- [0..20] ]
        render r | r <= 0.2  = ' '
                 | r <= 0.4  = '.'
                 | r <= 0.6  = '-'
                 | r <= 0.8  = '='
                 | otherwise = '#'

    return $ unlines outMat
  where
    inCircle :: KnownNat n => R n -> (R n, Double) -> Bool
    v `inCircle` (o, r) = norm_2 (v - o) <= r

main :: IO ()
main = do
    args <- getArgs
    let n    = readMaybe =<< (args !!? 0)
        rate = readMaybe =<< (args !!? 1)
    putStrLn "Training network..."
    putStrLn =<< evalRandIO (netTest (fromMaybe 0.25   rate)
                                     (fromMaybe 500000 n   )
                            )

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

-- * Exercises
--

popLayer :: Network i (h ': hs) o -> (Weights i h, Network h hs o)
popLayer (w :&~ n) = (w, n)

addNetWeights :: forall i hs o. (KnownNat i, KnownNat o)
              => Sing hs
              -> Network i hs o
              -> Network i hs o
              -> Network i hs o
addNetWeights = \case
    SNil -> \case
      O w1 -> \case
        O w2      -> O (addW w1 w2)
    SNat `SCons` ss -> \case
      w1 :&~ n1 -> \case
        w2 :&~ n2 -> addW w1 w2 :&~ addNetWeights ss n1 n2
  where
    addW :: (KnownNat n, KnownNat m)
         => Weights n m
         -> Weights n m
         -> Weights n m
    addW (W b1 w1) (W b2 w2) = W (b1 + b2) (w1 + w2)

hiddenSing :: Network i hs o -> Sing hs
hiddenSing = \case
    O _      -> SNil
    _ :&~ n' -> SNat `SCons` hiddenSing n'


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TypeFamilies #-}

module Math.Markov where

import Control.Monad.Random (getRandoms, Rand, RandomGen)
import Data.Counter (count)
import Data.Default (Default(..))
import Data.List (nub, sort)
import Data.List.Split (splitOn)
import qualified Data.Map as M
import Data.Maybe (fromMaybe)
import qualified Data.Set as S
import qualified Data.Vector as VB
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA

import Misc.Utils (ngrams)
import Math.Dist (DiscreteDist(..), IndexMap, JointDiscreteDistribution(..), Prob, Simulable(..), bisect, getProb, normalizeVec, trainDiscrete, valsToIndices)
import Math.Multiarray ((^!^), (^:^), MultiArray(..), multiarray)

-- MARKOV MODEL --

data MarkovModel v a = Markov { pi0 :: DiscreteDist v a, trans :: LA.Matrix Prob }
    deriving (Eq, Show)

rowStochasticMatrix :: [V.Vector Double] -> LA.Matrix Prob
rowStochasticMatrix rows = LA.fromRows $ normalizeVec <$> rows

-- construct a MarkovModel from an initial distribution and a transition matrix (entry i, j is probability of transition from i to j, so that the matrix is row-stochastic)
markov :: DiscreteDist v a -> [V.Vector Prob] -> MarkovModel v a
markov dist rows
    | length rows /= n               = error $ "transition matrix must have " ++ show n ++ " rows"
    | any (/= n) $ V.length <$> rows = error $ "all rows of transition matrix must have length " ++ show n
    | otherwise                      = Markov {pi0 = dist, trans = mat}
    where
        n = length $ vals dist
        mat = rowStochasticMatrix rows

instance (Enum v, Num v, Ord v, Ord a, Show a) => JointDiscreteDistribution MarkovModel v a where
    -- vars are always [0, 1, 2, ...]
    getVars :: MarkovModel v a -> [v]
    getVars _ = [0..]
    getJointProbEvent :: MarkovModel v a -> [[a]] -> Prob
    getJointProbEvent dist valss = exp $ getJointLogProbEvent dist valss
    getJointLogProb :: MarkovModel v a -> [a] -> Double
    getJointLogProb _ []                       = 0.0
    getJointLogProb (Markov {pi0, trans}) vals = sum logProbs
        where
            indices = valsToIndices pi0 vals
            transitionIndices = zip indices (tail indices)
            logProbs = log <$> ((probs pi0 V.! (head indices)) : (LA.atIndex trans <$> transitionIndices))
    -- marginal distribution is no longer a simple Markov chain
    marginalizeIndices = undefined
    -- infinite list of marginal distributions
    marginals :: MarkovModel v a -> [DiscreteDist v a]
    marginals (Markov {pi0, trans}) = dists
        where
            pis = normalizeVec <$> iterate (LA.<# trans) (probs pi0)
            vals' = vals pi0
            cdfs = V.scanl' (+) 0.0 <$> pis
            dists = [Discrete {var = i, vals = vals', probs = pmf, cdf = cdf'} | (i, pmf, cdf') <- zip3 [0..] pis cdfs]
    -- conditional distribution is no longer a simple Markov chain
    conditionOn = undefined
    -- distribution is infinite so can't be represented as a DiscreteDist
    forgetJoint = undefined

-- computes conditional distribution on a single variable given the nearest preceding and succeeding variable (if existent)
-- input is t and a pair ((k, val), (l, val')), and we compute P(X_t | X_{t - k} = val, X_{t + l} = val')
-- each subpair can also be Nothing, indicating the absence of such a condition
markovConditionOn :: (Integral v, Ord a) => MarkovModel v a -> v -> (Maybe (Int, a), Maybe (Int, a)) -> DiscreteDist v a
markovConditionOn (Markov {pi0, trans}) t pair = Discrete {var = t, vals = vals', probs = pmf', cdf = cdf'}
    where
        stepRow :: V.Vector Prob -> Int -> V.Vector Prob
        stepRow vec k = (iterate (LA.<# trans) vec) !! k
        stepCol :: V.Vector Prob -> Int -> V.Vector Prob
        stepCol vec k = (iterate (trans LA.#>) vec) !! k
        pis = iterate (LA.<# trans) (probs pi0)
        vals' = vals pi0
        n = VB.length vals'
        indVec :: Int -> V.Vector Prob
        indVec i = V.generate n (\j -> if (j == i) then 1.0 else 0.0)
        vecMul = V.zipWith (*)
        t' = (fromInteger $ toInteger t) :: Int
        valIdx = M.fromList $ zip (VB.toList (vals pi0)) [0..]
        val2Index = (valIdx M.!)
        pmf = case pair of
            (Nothing, Nothing)              -> pis !! t'
            (Just (k, val), Nothing)        -> stepRow (indVec $ val2Index val) k
            (Nothing, Just (l, val'))       -> stepCol (indVec $ val2Index val') l `vecMul` (pis !! t')
            (Just (k, val), Just (l, val')) -> stepCol (indVec $ val2Index val') l `vecMul` stepRow (indVec $ val2Index val) k
        pmf' = normalizeVec pmf
        cdf' = V.scanl' (+) 0.0 pmf'

-- trains a Markov model from data
-- smoothing constant is how much to add to each entry of transition matrix, and to each entry of initial vector
trainMarkovModel :: (Ord a, Show a) => Maybe [a] -> Double -> [[a]] -> MarkovModel Integer a
trainMarkovModel vals smooth chains = markov pi0 (LA.toRows trans)
    where
        allVals = fromMaybe (sort $ nub $ concat chains) vals
        chains' = filter (not . null) chains
        pi0 = trainDiscrete 0 (Just allVals) smooth (head <$> chains)
        indices = valsToIndices pi0 <$> chains
        transitionPairs = concat [zip inds (tail inds) | inds <- indices]
        n = length allVals
        addSmooth (key, val) = (key, val + smooth)
        trans = LA.assoc (n, n) smooth (addSmooth <$> (M.toList $ count transitionPairs))

-- samples a Markov model (infinitely)
markovSample :: (RandomGen g) => MarkovModel v a -> Rand g [a]
markovSample (Markov {pi0, trans}) = do
    rs <- getRandoms
    let indices = bisect (head rs) cdf0 : [bisect r (transCdfs VB.! i) | (i, r) <- zip indices (tail rs)]
    return $ (vs VB.!) <$> indices
    where
        vs = vals pi0
        cdf0 = cdf pi0
        transCdfs = VB.fromList $ (V.scanl' (+) 0.0) <$> (LA.toRows trans)

instance (b ~ [a]) => Simulable (MarkovModel v) a b where
    sample :: (RandomGen g) => MarkovModel v a -> Rand g [a]
    sample = markovSample

-- given a distribution on integers steps, creates a MarkovModel representing the random walk on the i.i.d. steps
-- must be bounded within a finite range
boundedIntegerRandomWalk :: DiscreteDist v Int -> (Int, Int) -> MarkovModel Integer Int
boundedIntegerRandomWalk dist (low, high) = markov pi0 rows
    where
        n = high - low + 1
        pi0 = trainDiscrete 0 (Just [low..high]) 0.0 [0]
        rows = [V.generate n (\j -> getProb dist (j + low - i)) | i <- [low..high]]

-- gets backward transition probabilities P(X_n = i | X_{n+1} = j)
backwardTransitions :: MarkovModel v a -> [V.Vector Prob]
backwardTransitions (Markov {trans}) = normalizeVec <$> (LA.toRows $ LA.tr trans)

-- TERMINATING MARKOV MODEL --

-- type for values in a terminating Markov chain
data Token a = Token a | Stop
    deriving (Eq, Ord, Show)

isToken :: Token a -> Bool
isToken (Token _) = True
isToken _         = False

fromToken :: Token a -> a
fromToken (Token t) = t
fromToken _         = error "invalid token"

-- converts elements to Tokens and appends a final Stop element
-- this Nothing value can be used as a terminating value for Markov chain generation
terminate :: [a] -> [Token a]
terminate = (++ [Stop]) . map Token

trainTerminatingMarkovModel :: (Ord a, Show a) => Maybe [a] -> Double -> [[a]] -> MarkovModel Integer (Token a)
trainTerminatingMarkovModel vals smooth chains = trainMarkovModel tokens smooth (terminate <$> chains)
    where tokens = (Stop :) . map Token <$> vals

instance {-# OVERLAPPING #-} (b ~ [a]) => Simulable (MarkovModel v) (Token a) b where
    -- override behavior so sampling terminates upon reaching Stop
    sample = fmap (map fromToken . takeWhile isToken) . markovSample

-- N-GRAM MODELS --

-- stores n, alphabet, mapping from alphabet to indices, and transition probability matrices for each Markovity level
data NgramModel v a = Ngrams Int (VB.Vector a) (IndexMap a) [MultiArray Prob]
    deriving (Eq, Show)

instance (Enum v, Num v, Ord v, Ord a, Show a) => JointDiscreteDistribution NgramModel v a where
    getVars :: NgramModel v a -> [v]
    getVars _ = [0..]
    getJointProbEvent :: NgramModel v a -> [[a]] -> Prob
    getJointProbEvent dist valss = exp $ getJointLogProbEvent dist valss
    getJointLogProb :: NgramModel v a -> [a] -> Double
    getJointLogProb (Ngrams n _ valIdx mats) vals = sum logprobs
        where
            indices = (valIdx M.!) <$> vals
            idxGrams = [take i indices | i <- [1..(n - 1)]] ++ ngrams n indices
            matSeq = init mats ++ repeat (last mats)
            logprobs = log <$> zipWith (^!^) matSeq idxGrams
    marginalizeIndices = undefined
    marginals = undefined
    conditionOn = undefined
    forgetJoint = undefined

trainNgramModel :: (Ord a, Show a) => Int -> Maybe [a] -> Double -> [[a]] -> NgramModel Integer a
trainNgramModel n alphabet smooth chains = Ngrams n (VB.fromList alphabet') valIdx transArrs
    where
        alphabet' = fromMaybe (sort $ nub $ concat chains) alphabet
        m = length alphabet'
        valIdx = M.fromList $ zip alphabet' [0..]  -- mapping from alphabet values to indices
        idxChains = map (valIdx M.!) <$> chains
        getTransitions :: Int -> MultiArray Prob
        getTransitions k = multiarray (replicate k m) (concat rows)
            where
                pairs = concat $ (map (splitAt (k - 1)) . ngrams k) <$> idxChains
                cumProds = (m ^) <$> reverse [0..(k - 2)]
                transitionPairs = [(sum $ zipWith (*) cumProds fromInds, head toInds) | (fromInds, toInds) <- pairs]
                addSmooth (key, val) = (key, val + smooth)
                trans = LA.assoc (m ^ (k - 1), m) smooth (addSmooth <$> (M.toList $ count transitionPairs))
                rows = V.toList . normalizeVec <$> LA.toRows trans
        transArrs = getTransitions <$> [1..n]

-- samples an n-gram model (infinitely)
ngramSample :: (RandomGen g) => NgramModel v a -> Rand g [a]
ngramSample (Ngrams n alphabet _ transArrs) = do
    rs <- getRandoms
    let indices = [] : [rotateInds arr prev r | (prev, arr, r) <- zip3 indices arrs rs]
    return [alphabet VB.! (last inds) | inds <- tail indices]
    where
        getTail :: [Int] -> [Int]
        getTail inds = if (length inds < n) then inds else tail inds
        rotateInds :: MultiArray Prob -> [Int] -> Double -> [Int]
        rotateInds arr inds r = indTail ++ [nextInd]
            where
                indTail = getTail inds
                row = entries $ arr ^:^ (((Just . pure) <$> indTail) ++ [Nothing])
                cdf = V.scanl' (+) 0.0 row
                nextInd = bisect r cdf
        arrs = transArrs ++ (repeat $ last transArrs)

instance (b ~ [a]) => Simulable (NgramModel v) a b where
    sample :: (RandomGen g) => NgramModel v a -> Rand g [a]
    sample = ngramSample

-- condSlice :: (Ord a) => NgramModel v a -> [a] -> [(Int, a, Prob)]
-- condSlice (Ngrams n vals valIdx transArrs) vals' = zip3 [0..] (VB.toList vals) (V.toList pdf)
--     where
--         lastN = reverse $ take (n - 1) $ reverse vals'
--         inds = (valIdx M.!) <$> lastN
--         arr = transArrs !! length inds
--         pdf = entries $ arr ^:^ (((Just . pure) <$> inds) ++ [Nothing])

-- TERMINATING N-GRAM MODELS --

trainTerminatingNgramModel :: (Ord a, Show a) => Int -> Maybe [a] -> Double -> [[a]] -> NgramModel Integer (Token a)
trainTerminatingNgramModel n alphabet smooth chains = trainNgramModel n toks smooth (terminate <$> chains)
    where toks = (Stop :) . map Token <$> alphabet

instance {-# OVERLAPPING #-} (b ~ [a]) => Simulable (NgramModel v) (Token a) b where
    -- override behavior so sampling terminates upon reaching Stop
    sample = fmap (map fromToken . takeWhile isToken) . ngramSample

-- TEXT N-GRAM MODELS --

data CharNgramParams = CharNgramParams {
      charAlphabet :: Maybe String,   -- alphabet to use
      charSmooth :: Double,           -- smoothing constant
      charLineSep :: String,          -- line separator string
      charTransform :: Char -> Char,  -- function to transform each character
      charFilter :: Bool              -- if True, filter out missing characters
}

instance Default CharNgramParams where
    def = CharNgramParams { charAlphabet = Nothing, charSmooth = 1e-8, charLineSep = "\n", charTransform = id, charFilter = True }

trainCharNgramModel :: Int -> CharNgramParams -> String -> NgramModel Integer (Token Char)
trainCharNgramModel n params text = trainTerminatingNgramModel n (charAlphabet params) (charSmooth params) lines''
    where
        lines = splitOn (charLineSep params) text
        lines' = map (charTransform params) <$> lines
        lines'' = case (charAlphabet params) of
            Just alphabet -> if (charFilter params)
                                then filter (`S.member` alphabetSet) <$> lines'
                                else lines'
                where alphabetSet = S.fromList alphabet
            Nothing       -> lines'


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Gelu
Description : Gaussian Error Linear Unit (GELU)
Copyright   : (c) Manuel Schneckenreither, 2020
License     : BSD2
Stability   : experimental

This module implements the Gaussian Error Linear Unit (GELU) activiation function. See

Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).

Available at: https://arxiv.org/pdf/1606.08415.pdf

As in the paper we simply let μ = 0 and σ = 1. Futher, we use the simplified and thus fast representation: x * sigmoid (1.702 * x)

-}
module Grenade.Layers.Gelu (
    Gelu (..)
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Serialize
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS

import           Grenade.Core

-- | A Gaussion Error Linear Unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
--
--   Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).
data Gelu = Gelu
  deriving (Generic, NFData, Show)

instance UpdateLayer Gelu where
  type Gradient Gelu = ()
  runUpdate _ _ _ = Gelu
  reduceGradient _ = ()

instance RandomLayer Gelu where
  createRandomWith _ _ = return Gelu

instance Serialize Gelu where
  put _ = return ()
  get = return Gelu

geluForwardFast :: Floating x => x -> x
geluForwardFast x = x / (e ** (-1.702 * x) + 1) -- = x * sigmoid (1.702 * x)
  where
    e = 2.71828

geluBackwardFast :: Floating x => x -> x
geluBackwardFast x = (e ** (1.702 * x) * (1 + e ** (1.702 * x) + 1.702 * x)) / (1 + e ** (1.702 * x)) ** 2
  where
    e = 2.71828

instance (KnownNat i) => Layer Gelu ('D1 i) ('D1 i) where
  type Tape Gelu ('D1 i) ('D1 i) = S ('D1 i)

  runForwards _ (S1D y) = (S1D y, S1D (gelu y))
    where
      gelu = LAS.dvmap geluForwardFast
  runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (gelu' y * dEdy))
    where
      gelu' = LAS.dvmap geluBackwardFast

instance (KnownNat i, KnownNat j) => Layer Gelu ('D2 i j) ('D2 i j) where
  type Tape Gelu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (gelu y))
    where
      gelu = LAS.dmmap geluForwardFast
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (gelu' y * dEdy))
    where
      gelu' = LAS.dmmap geluBackwardFast

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Gelu ('D3 i j k) ('D3 i j k) where

  type Tape Gelu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (gelu y))
    where
      gelu = LAS.dmmap geluForwardFast
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (gelu' y * dEdy))
    where
      gelu' = LAS.dmmap geluBackwardFast


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Poisson
-- Copyright : (c) 2009, 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Poisson distribution.  This is the discrete probability
-- distribution of a number of events occurring in a fixed interval if
-- these events occur with a known average rate, and occur
-- independently from each other within that interval.

module Statistics.Distribution.Poisson
    (
      PoissonDistribution
    -- * Constructors
    , poisson
    , poissonE
    -- * Accessors
    , poissonLambda
    -- * References
    -- $references
    ) where

import Control.Applicative
import Data.Data            (Data, Typeable)
import GHC.Generics         (Generic)
import Numeric.SpecFunctions (incompleteGamma,logFactorial)
import Numeric.MathFunctions.Constants (m_neg_inf)

import qualified Statistics.Distribution as D
import qualified Statistics.Distribution.Poisson.Internal as I
import Statistics.Internal

newtype PoissonDistribution = PD {
      poissonLambda :: Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show PoissonDistribution where
  showsPrec i (PD l) = defaultShow1 "poisson" l i
instance Read PoissonDistribution where
  readPrec = defaultReadPrecM1 "poisson" poissonE

instance D.Distribution PoissonDistribution where
    cumulative (PD lambda) x
      | x < 0        = 0
      | isInfinite x = 1
      | isNaN      x = error "Statistics.Distribution.Poisson.cumulative: NaN input"
      | otherwise    = 1 - incompleteGamma (fromIntegral (floor x + 1 :: Int)) lambda

instance D.DiscreteDistr PoissonDistribution where
    probability (PD lambda) x = I.probability lambda (fromIntegral x)
    logProbability (PD lambda) i
      | i < 0     = m_neg_inf
      | otherwise = log lambda * fromIntegral i - logFactorial i - lambda

instance D.Variance PoissonDistribution where
    variance = poissonLambda

instance D.Mean PoissonDistribution where
    mean = poissonLambda

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

instance D.MaybeVariance PoissonDistribution where
    maybeStdDev   = Just . D.stdDev

instance D.Entropy PoissonDistribution where
  entropy (PD lambda) = I.poissonEntropy lambda

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

-- | Create Poisson distribution.
poisson :: Double -> PoissonDistribution
poisson l = maybe (error $ errMsg l) id $ poissonE l

-- | Create Poisson distribution.
poissonE :: Double -> Maybe PoissonDistribution
poissonE l
  | l >=  0   = Just (PD l)
  | otherwise = Nothing

errMsg :: Double -> String
errMsg l = "Statistics.Distribution.Poisson.poisson: lambda must be non-negative. Got "
        ++ show l

-- $references
--
-- * Loader, C. (2000) Fast and Accurate Computation of Binomial
--   Probabilities. <http://projects.scipy.org/scipy/raw-attachment/ticket/620/loader2000Fast.pdf>
-- * Adell, J., Lekuona, A., and Yu, Y. (2010) Sharp Bounds on the
--   Entropy of the Poisson Law and Related Quantities
--   <http://arxiv.org/pdf/1001.2897.pdf>


{-# LANGUAGE ViewPatterns        #-}
--------------------------------------------------------------------------------
-- |
-- Module      : ArrayFire.BLAS
-- Copyright   : David Johnson (c) 2019-2020
-- License     : BSD3
-- Maintainer  : David Johnson <djohnson.m@gmail.com>
-- Stability   : Experimental
-- Portability : GHC
--
-- Basic Linear Algebra Subprograms (BLAS) API
--
-- @
-- main :: IO ()
-- main = print (matmul x y xProp yProp)
--  where
--     x,y :: Array Double
--     x = matrix (2,3) [[1,2],[3,4],[5,6]]
--     y = matrix (3,2) [[1,2,3],[4,5,6]]
--
--     xProp, yProp :: MatProp
--     xProp = None
--     yProp = None
-- @
-- @
--  ArrayFire Array
--  [2 2 1 1]
--     22.0000    49.0000
--     28.0000    64.0000
-- @
--------------------------------------------------------------------------------
module ArrayFire.BLAS where

import Data.Complex

import ArrayFire.FFI
import ArrayFire.Internal.BLAS
import ArrayFire.Internal.Types

-- | The following applies for Sparse-Dense matrix multiplication.
--
-- This function can be used with one sparse input. The sparse input must always be the lhs and the dense matrix must be rhs.
--
-- The sparse array can only be of 'CSR' format.
--
-- The returned array is always dense.
--
-- optLhs an only be one of AF_MAT_NONE, AF_MAT_TRANS, AF_MAT_CTRANS.
--
-- optRhs can only be AF_MAT_NONE.
--
-- >>> matmul (matrix @Double (2,2) [[1,2],[3,4]]) (matrix @Double (2,2) [[1,2],[3,4]]) None None
-- ArrayFire Array
-- [2 2 1 1]
--     7.0000    15.0000
--    10.0000    22.0000
matmul
  :: AFType a
  => Array a
  -- ^ 2D matrix of Array a, left-hand side
  -> Array a
  -- ^ 2D matrix of Array a, right-hand side
  -> MatProp
  -- ^ Left hand side matrix options
  -> MatProp
  -- ^ Right hand side matrix options
  -> Array a
  -- ^ Output of 'matmul'
matmul arr1 arr2 prop1 prop2 = do
  op2 arr1 arr2 (\p a b -> af_matmul p a b (toMatProp prop1) (toMatProp prop2))

-- | Scalar dot product between two vectors. Also referred to as the inner product.
--
-- >>> dot (vector @Double 10 [1..]) (vector @Double 10 [1..]) None None
-- ArrayFire Array
-- [1 1 1 1]
--   385.0000
dot
  :: AFType a
  => Array a
  -- ^ Left-hand side input
  -> Array a
  -- ^ Right-hand side input
  -> MatProp
  -- ^ Options for left-hand side. Currently only AF_MAT_NONE and AF_MAT_CONJ are supported.
  -> MatProp
  -- ^ Options for right-hand side. Currently only AF_MAT_NONE and AF_MAT_CONJ are supported.
  -> Array a
  -- ^ Output of 'dot'
dot arr1 arr2 prop1 prop2 =
  op2 arr1 arr2 (\p a b -> af_dot p a b (toMatProp prop1) (toMatProp prop2))

-- | Scalar dot product between two vectors. Also referred to as the inner product. Returns the result as a host scalar.
--
-- >>> dotAll (vector @Double 10 [1..]) (vector @Double 10 [1..]) None None
-- 385.0 :+ 0.0
dotAll
  :: AFType a
  => Array a
  -- ^ Left-hand side array
  -> Array a
  -- ^ Right-hand side array
  -> MatProp
  -- ^ Options for left-hand side. Currently only AF_MAT_NONE and AF_MAT_CONJ are supported.
  -> MatProp
  -- ^ Options for right-hand side. Currently only AF_MAT_NONE and AF_MAT_CONJ are supported.
  -> Complex Double
  -- ^ Real and imaginary component result
dotAll arr1 arr2 prop1 prop2 = do
  let (real,imag) =
        infoFromArray22 arr1 arr2 $ \a b c d ->
          af_dot_all a b c d (toMatProp prop1) (toMatProp prop2)
  real :+ imag

-- | Transposes a matrix.
--
-- >>> array = matrix @Double (2,3) [[2,3],[3,4],[5,6]]
-- >>> array
-- ArrayFire Array
-- [2 3 1 1]
--     2.0000     3.0000     5.0000
--     3.0000     4.0000     6.0000
--
-- >>> transpose array True
-- ArrayFire Array
-- [3 2 1 1]
--     2.0000     3.0000
--     3.0000     4.0000
--     5.0000     6.0000
--
transpose
  :: AFType a
  => Array a
  -- ^ Input matrix to be transposed
  -> Bool
  -- ^ Should perform conjugate transposition
  -> Array a
  -- ^ The transposed matrix
transpose arr1 (fromIntegral . fromEnum -> b) =
  arr1 `op1` (\x y -> af_transpose x y b)

-- | Transposes a matrix.
--
-- * Warning: This function mutates an array in-place, all subsequent references will be changed. Use carefully.
--
-- >>> array = matrix @Double (2,2) [[1,2],[3,4]]
-- >>> array
-- ArrayFire Array
-- [3 2 1 1]
--    1.0000     4.0000
--    2.0000     5.0000
--    3.0000     6.0000
--
-- >>> transposeInPlace array False
-- >>> array
-- ArrayFire Array
-- [2 2 1 1]
--    1.0000     2.0000
--    3.0000     4.0000
--
transposeInPlace
  :: AFType a
  => Array a
  -- ^ Input matrix to be transposed
  -> Bool
  -- ^ Should perform conjugate transposition
  -> IO ()
transposeInPlace arr (fromIntegral . fromEnum -> b) =
  arr `inPlace` (`af_transpose_inplace` b)


{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra.Util.Convolution
Copyright   :  (c) Alberto Ruiz 2012
License     :  GPL

Maintainer  :  Alberto Ruiz (aruiz at um dot es)
Stability   :  provisional

-}
-----------------------------------------------------------------------------

module Numeric.LinearAlgebra.Util.Convolution(
   corr, conv, corrMin,
   corr2, conv2, separable
) where

import Numeric.LinearAlgebra

vectSS :: Element t => Int -> Vector t -> Matrix t
vectSS n v = fromRows [ subVector k n v | k <- [0 .. dim v - n] ]

corr :: Product t => Vector t -- ^ kernel
                  -> Vector t -- ^ source
                  -> Vector t
{- ^ correlation

>>> corr (fromList[1,2,3]) (fromList [1..10])
fromList [14.0,20.0,26.0,32.0,38.0,44.0,50.0,56.0]

-}
corr ker v | dim ker <= dim v = vectSS (dim ker) v <> ker
           | otherwise = error $ "corr: dim kernel ("++show (dim ker)++") > dim vector ("++show (dim v)++")"

conv :: (Product t, Num t) => Vector t -> Vector t -> Vector t
{- ^ convolution ('corr' with reversed kernel and padded input, equivalent to polynomial product)

>>> conv (fromList[1,1]) (fromList [-1,1])
fromList [-1.0,0.0,1.0]

-}
conv ker v = corr ker' v'
  where
    ker' = (flatten.fliprl.asRow) ker
    v' | dim ker > 1 = join [z,v,z]
       | otherwise   = v
    z = constant 0 (dim ker -1)

corrMin :: (Container Vector t, RealElement t, Product t)
        => Vector t
        -> Vector t
        -> Vector t
-- ^ similar to 'corr', using 'min' instead of (*)
corrMin ker v = minEvery ss (asRow ker) <> ones
  where
    minEvery a b = cond a b a a b
    ss = vectSS (dim ker) v
    ones = konst' 1 (dim ker)

matSS :: Element t => Int -> Matrix t -> [Matrix t]
matSS dr m = map (reshape c) [ subVector (k*c) n v | k <- [0 .. r - dr] ]
  where
    v = flatten m
    c = cols m
    r = rows m
    n = dr*c

corr2 :: Product a => Matrix a -> Matrix a -> Matrix a
-- ^ 2D correlation
corr2 ker mat = dims
              . concatMap (map ((<.> ker') . flatten) . matSS c . trans)
              . matSS r $ mat
  where
    r = rows ker
    c = cols ker
    ker' = flatten (trans ker)
    rr = rows mat - r + 1
    rc = cols mat - c + 1
    dims | rr > 0 && rc > 0 = (rr >< rc)
         | otherwise = error $ "corr2: dim kernel ("++sz ker++") > dim matrix ("++sz mat++")"
    sz m = show (rows m)++"x"++show (cols m)

conv2 :: (Num a, Product a) => Matrix a -> Matrix a -> Matrix a
-- ^ 2D convolution
conv2 k m = corr2 (fliprl . flipud $ k) pm
  where
    pm | r == 0 && c == 0 = m
       | r == 0           = fromBlocks [[z3,m,z3]]
       |           c == 0 = fromBlocks [[z2],[m],[z2]]
       | otherwise        = fromBlocks [[z1,z2,z1]
                                       ,[z3, m,z3]
                                       ,[z1,z2,z1]]
    r = rows k - 1
    c = cols k - 1
    h = rows m
    w = cols m
    z1 = konst' 0 (r,c)
    z2 = konst' 0 (r,w)
    z3 = konst' 0 (h,c)

-- TODO: could be simplified using future empty arrays

separable :: Element t => (Vector t -> Vector t) -> Matrix t -> Matrix t
-- ^ matrix computation implemented as separated vector operations by rows and columns.
separable f = fromColumns . map f . toColumns . fromRows . map f . toRows