content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  import Control.DeepSeq
  import Criterion.Main
  import Criterion.Types
  import Data.Char
  import Data.Functor.Identity
  import Data.Time
  import qualified Data.Vector as V
  import GHC.Generics (Generic)
  import GHC.TypeLits
  import Lens.Micro
  import Lens.Micro.TH
  import Numeric.Backprop
  import Numeric.Backprop.Class
  import qualified Numeric.LinearAlgebra as HM
  import Numeric.LinearAlgebra.Static
  import System.Directory
  import qualified System.Random.MWC as MWC
  
  type family HKD f a where
    HKD Identity a = a
    HKD f a = f a
  
  data Layer' i o f
    = Layer
    { _lWeights :: !(HKD f (L o i))
    , _lBiases :: !(HKD f (R o))
    }
    deriving (Generic)
  
  type Layer i o = Layer' i o Identity
  
  deriving instance (KnownNat i, KnownNat o) => Show (Layer i o)
  instance NFData (Layer i o)
  
  makeLenses ''Layer'
  
  data Network' i h1 h2 o f
    = Net
    { _nLayer1 :: !(HKD f (Layer i h1))
    , _nLayer2 :: !(HKD f (Layer h1 h2))
    , _nLayer3 :: !(HKD f (Layer h2 o))
    }
    deriving (Generic)
  
  type Network i h1 h2 o = Network' i h1 h2 o Identity
  
  deriving instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Show (Network i h1 h2 o)
  instance NFData (Network i h1 h2 o)
  
  makeLenses ''Network'
  
  main :: IO ()
  main = do
    g <-
      MWC.initialize
        . V.fromList
        . map (fromIntegral . ord)
        $ "hello world"
    test0 <- MWC.uniformR @(R 784, R 10) ((0, 0), (1, 1)) g
    net0 <- MWC.uniformR @(Network 784 300 100 10) (-0.5, 0.5) g
    t <- getZonedTime
    let tstr = formatTime defaultTimeLocale "%Y%m%d-%H%M%S" t
    createDirectoryIfMissing True "bench-results"
    defaultMainWith
      defaultConfig
        { reportFile = Just $ "bench-results/mnist-bench_" ++ tstr ++ ".html"
        , timeLimit = 10
        }
      [ bgroup
          "gradient"
          [ let runTest x y = gradNetManual x y net0
             in bench "manual" $ nf (uncurry runTest) test0
          , let runTest x y = gradBP (netErr x y) net0
             in bench "bp-lens" $ nf (uncurry runTest) test0
          , let runTest x y = gradBP (netErrHKD x y) net0
             in bench "bp-hkd" $ nf (uncurry runTest) test0
          , let runTest x y = gradBP (\n' -> netErrHybrid n' y x) net0
             in bench "hybrid" $ nf (uncurry runTest) test0
          ]
      , bgroup
          "descent"
          [ let runTest x y = trainStepManual 0.02 x y net0
             in bench "manual" $ nf (uncurry runTest) test0
          , let runTest x y = trainStep 0.02 x y net0
             in bench "bp-lens" $ nf (uncurry runTest) test0
          , let runTest x y = trainStepHKD 0.02 x y net0
             in bench "bp-hkd" $ nf (uncurry runTest) test0
          , let runTest x y = trainStepHybrid 0.02 x y net0
             in bench "hybrid" $ nf (uncurry runTest) test0
          ]
      , bgroup
          "run"
          [ let runTest = runNetManual net0
             in bench "manual" $ nf runTest (fst test0)
          , let runTest x = evalBP (`runNetwork` x) net0
             in bench "bp-lens" $ nf runTest (fst test0)
          , let runTest x = evalBP (`runNetworkHKD` x) net0
             in bench "bp-hkd" $ nf runTest (fst test0)
          , let runTest x = evalBP (`runNetHybrid` x) net0
             in bench "hybrid" $ nf runTest (fst test0)
          ]
      ]
  
  -- ------------------------------
  -- - "Backprop" Lens Mode       -
  -- ------------------------------
  
  runLayer ::
    (KnownNat i, KnownNat o, Reifies s W) =>
    BVar s (Layer i o) ->
    BVar s (R i) ->
    BVar s (R o)
  runLayer l x = (l ^^. lWeights) #>! x + (l ^^. lBiases)
  {-# INLINE runLayer #-}
  
  softMax :: (KnownNat n, Reifies s W) => BVar s (R n) -> BVar s (R n)
  softMax x = konst' (1 / sumElements' expx) * expx
    where
      expx = exp x
  {-# INLINE softMax #-}
  
  runNetwork ::
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) =>
    BVar s (Network i h1 h2 o) ->
    R i ->
    BVar s (R o)
  runNetwork n =
    softMax
      . runLayer (n ^^. nLayer3)
      . logistic
      . runLayer (n ^^. nLayer2)
      . logistic
      . runLayer (n ^^. nLayer1)
      . auto
  {-# INLINE runNetwork #-}
  
  crossEntropy ::
    (KnownNat n, Reifies s W) =>
    R n ->
    BVar s (R n) ->
    BVar s Double
  crossEntropy t r = negate $ log r <.>! auto t
  {-# INLINE crossEntropy #-}
  
  netErr ::
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) =>
    R i ->
    R o ->
    BVar s (Network i h1 h2 o) ->
    BVar s Double
  netErr x t n = crossEntropy t (runNetwork n x)
  {-# INLINE netErr #-}
  
  trainStep ::
    forall i h1 h2 o.
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) =>
    Double ->
    R i ->
    R o ->
    Network i h1 h2 o ->
    Network i h1 h2 o
  trainStep r !x !t !n = n - realToFrac r * gradBP (netErr x t) n
  {-# INLINE trainStep #-}
  
  -- ------------------------------
  -- - "Backprop" HKD Mode        -
  -- ------------------------------
  
  runLayerHKD ::
    (KnownNat i, KnownNat o, Reifies s W) =>
    BVar s (Layer i o) ->
    BVar s (R i) ->
    BVar s (R o)
  runLayerHKD (splitBV -> Layer w b) x = w #>! x + b
  {-# INLINE runLayerHKD #-}
  
  runNetworkHKD ::
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) =>
    BVar s (Network i h1 h2 o) ->
    R i ->
    BVar s (R o)
  runNetworkHKD (splitBV -> Net l1 l2 l3) =
    softMax
      . runLayerHKD l3
      . logistic
      . runLayerHKD l2
      . logistic
      . runLayerHKD l1
      . auto
  {-# INLINE runNetworkHKD #-}
  
  netErrHKD ::
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) =>
    R i ->
    R o ->
    BVar s (Network i h1 h2 o) ->
    BVar s Double
  netErrHKD x t n = crossEntropy t (runNetworkHKD n x)
  {-# INLINE netErrHKD #-}
  
  trainStepHKD ::
    forall i h1 h2 o.
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) =>
    Double ->
    R i ->
    R o ->
    Network i h1 h2 o ->
    Network i h1 h2 o
  trainStepHKD r !x !t !n = n - realToFrac r * gradBP (netErrHKD x t) n
  {-# INLINE trainStepHKD #-}
  
  -- ------------------------------
  -- - "Manual" Mode              -
  -- ------------------------------
  
  runLayerManual ::
    (KnownNat i, KnownNat o) =>
    Layer i o ->
    R i ->
    R o
  runLayerManual l x = (l ^. lWeights) #> x + (l ^. lBiases)
  {-# INLINE runLayerManual #-}
  
  softMaxManual :: KnownNat n => R n -> R n
  softMaxManual x = konst (1 / sumElements expx) * expx
    where
      expx = exp x
  {-# INLINE softMaxManual #-}
  
  runNetManual ::
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) =>
    Network i h1 h2 o ->
    R i ->
    R o
  runNetManual n =
    softMaxManual
      . runLayerManual (n ^. nLayer3)
      . logistic
      . runLayerManual (n ^. nLayer2)
      . logistic
      . runLayerManual (n ^. nLayer1)
  {-# INLINE runNetManual #-}
  
  gradNetManual ::
    forall i h1 h2 o.
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) =>
    R i ->
    R o ->
    Network i h1 h2 o ->
    Network i h1 h2 o
  gradNetManual x t (Net (Layer w1 b1) (Layer w2 b2) (Layer w3 b3)) =
    let y1 = w1 #> x
        z1 = y1 + b1
        x2 = logistic z1
        y2 = w2 #> x2
        z2 = y2 + b2
        x3 = logistic z2
        y3 = w3 #> x3
        z3 = y3 + b3
        o0 = exp z3
        o1 = HM.sumElements (extract o0)
        o2 = o0 / konst o1
        -- o3 = - (log o2 <.> t)
        dEdO3 = 1
        dEdO2 = - (dEdO3 * t / o2)
        dEdO1 = -((dEdO2 <.> o0) / (o1 ** 2))
        dEdO0 = konst dEdO1 + dEdO2 / konst o1
        dEdZ3 = dEdO0 * o0
        dEdY3 = dEdZ3
        dEdX3 = tr w3 #> dEdY3
        dEdZ2 = dEdX3 * (x3 * (1 - x3))
        dEdY2 = dEdZ2
        dEdX2 = tr w2 #> dEdY2
        dEdZ1 = dEdX2 * (x2 * (1 - x2))
        dEdY1 = dEdZ1
        dEdB3 = dEdZ3
        dEdW3 = dEdY3 `outer` x3
        dEdB2 = dEdZ2
        dEdW2 = dEdY2 `outer` x2
        dEdB1 = dEdZ1
        dEdW1 = dEdY1 `outer` x
     in Net (Layer dEdW1 dEdB1) (Layer dEdW2 dEdB2) (Layer dEdW3 dEdB3)
  {-# INLINE gradNetManual #-}
  
  trainStepManual ::
    forall i h1 h2 o.
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) =>
    Double ->
    R i ->
    R o ->
    Network i h1 h2 o ->
    Network i h1 h2 o
  trainStepManual r !x !t !n =
    let gN = gradNetManual x t n
     in n - (realToFrac r * gN)
  
  -- ------------------------------
  -- - "Hybrid" Mode              -
  -- ------------------------------
  
  layerOp :: (KnownNat i, KnownNat o) => Op '[Layer i o, R i] (R o)
  layerOp = op2 $ \(Layer w b) x ->
    ( w #> x + b
    , \g -> (Layer (g `outer` x) g, tr w #> g)
    )
  {-# INLINE layerOp #-}
  
  logisticOp ::
    Floating a =>
    Op '[a] a
  logisticOp = op1 $ \x ->
    let lx = logistic x
     in (lx, \g -> lx * (1 - lx) * g)
  {-# INLINE logisticOp #-}
  
  softMaxOp ::
    KnownNat n =>
    Op '[R n] (R n)
  softMaxOp = op1 $ \x ->
    let expx = exp x
        tot = sumElements expx
        invtot = 1 / tot
        res = konst invtot * expx
     in ( res
        , \g -> res - konst (invtot ** 2) * exp (2 * x) * g
        )
  {-# INLINE softMaxOp #-}
  
  softMaxCrossEntropyOp ::
    KnownNat n =>
    R n ->
    Op '[R n] Double
  softMaxCrossEntropyOp targ = op1 $ \x ->
    let expx = exp x
        sm = konst (1 / sumElements expx) * expx
        ce = negate $ log sm <.> targ
     in ( ce
        , \g -> (sm - targ) * konst g
        )
  {-# INLINE softMaxCrossEntropyOp #-}
  
  runNetHybrid ::
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) =>
    BVar s (Network i h1 h2 o) ->
    R i ->
    BVar s (R o)
  runNetHybrid n =
    liftOp1 softMaxOp
      . liftOp2 layerOp (n ^^. nLayer3)
      . liftOp1 logisticOp
      . liftOp2 layerOp (n ^^. nLayer2)
      . liftOp1 logisticOp
      . liftOp2 layerOp (n ^^. nLayer1)
      . auto
  {-# INLINE runNetHybrid #-}
  
  netErrHybrid ::
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W) =>
    BVar s (Network i h1 h2 o) ->
    R o ->
    R i ->
    BVar s Double
  netErrHybrid n t =
    liftOp1 (softMaxCrossEntropyOp t)
      . liftOp2 layerOp (n ^^. nLayer3)
      . liftOp1 logisticOp
      . liftOp2 layerOp (n ^^. nLayer2)
      . liftOp1 logisticOp
      . liftOp2 layerOp (n ^^. nLayer1)
      . auto
  {-# INLINE netErrHybrid #-}
  
  trainStepHybrid ::
    forall i h1 h2 o.
    (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) =>
    Double ->
    R i ->
    R o ->
    Network i h1 h2 o ->
    Network i h1 h2 o
  trainStepHybrid r !x !t !n =
    let gN = gradBP (\n' -> netErrHybrid n' t x) n
     in n - (realToFrac r * gN)
  {-# INLINE trainStepHybrid #-}
  
  -- ------------------------------
  -- - Operations                 -
  -- ------------------------------
  
  infixr 8 #>!
  (#>!) ::
    (KnownNat m, KnownNat n, Reifies s W) =>
    BVar s (L m n) ->
    BVar s (R n) ->
    BVar s (R m)
  (#>!) = liftOp2 . op2 $ \m v ->
    (m #> v, \g -> (g `outer` v, tr m #> g))
  {-# INLINE (#>!) #-}
  
  infixr 8 <.>!
  (<.>!) ::
    (KnownNat n, Reifies s W) =>
    BVar s (R n) ->
    BVar s (R n) ->
    BVar s Double
  (<.>!) = liftOp2 . op2 $ \x y ->
    ( x <.> y
    , \g -> (konst g * y, x * konst g)
    )
  {-# INLINE (<.>!) #-}
  
  konst' ::
    (KnownNat n, Reifies s W) =>
    BVar s Double ->
    BVar s (R n)
  konst' = liftOp1 . op1 $ \c -> (konst c, HM.sumElements . extract)
  {-# INLINE konst' #-}
  
  sumElements :: KnownNat n => R n -> Double
  sumElements = HM.sumElements . extract
  {-# INLINE sumElements #-}
  
  sumElements' ::
    (KnownNat n, Reifies s W) =>
    BVar s (R n) ->
    BVar s Double
  sumElements' = liftOp1 . op1 $ \x -> (sumElements x, konst)
  {-# INLINE sumElements' #-}
  
  logistic :: Floating a => a -> a
  logistic x = 1 / (1 + exp (-x))
  {-# INLINE logistic #-}
  
  -- ------------------------------
  -- - Instances                  -
  -- ------------------------------
  
  instance (KnownNat i, KnownNat o) => Num (Layer i o) where
    Layer w1 b1 + Layer w2 b2 = Layer (w1 + w2) (b1 + b2)
    Layer w1 b1 - Layer w2 b2 = Layer (w1 - w2) (b1 - b2)
    Layer w1 b1 * Layer w2 b2 = Layer (w1 * w2) (b1 * b2)
    abs (Layer w b) = Layer (abs w) (abs b)
    signum (Layer w b) = Layer (signum w) (signum b)
    negate (Layer w b) = Layer (negate w) (negate b)
    fromInteger x = Layer (fromInteger x) (fromInteger x)
  
  instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Num (Network i h1 h2 o) where
    Net a b c + Net d e f = Net (a + d) (b + e) (c + f)
    Net a b c - Net d e f = Net (a - d) (b - e) (c - f)
    Net a b c * Net d e f = Net (a * d) (b * e) (c * f)
    abs (Net a b c) = Net (abs a) (abs b) (abs c)
    signum (Net a b c) = Net (signum a) (signum b) (signum c)
    negate (Net a b c) = Net (negate a) (negate b) (negate c)
    fromInteger x = Net (fromInteger x) (fromInteger x) (fromInteger x)
  
  instance (KnownNat i, KnownNat o) => Fractional (Layer i o) where
    Layer w1 b1 / Layer w2 b2 = Layer (w1 / w2) (b1 / b2)
    recip (Layer w b) = Layer (recip w) (recip b)
    fromRational x = Layer (fromRational x) (fromRational x)
  
  instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Fractional (Network i h1 h2 o) where
    Net a b c / Net d e f = Net (a / d) (b / e) (c / f)
    recip (Net a b c) = Net (recip a) (recip b) (recip c)
    fromRational x = Net (fromRational x) (fromRational x) (fromRational x)
  
  instance KnownNat n => MWC.Variate (R n) where
    uniform g = randomVector <$> MWC.uniform g <*> pure Uniform
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g
  
  instance (KnownNat m, KnownNat n) => MWC.Variate (L m n) where
    uniform g = uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g
  
  instance (KnownNat i, KnownNat o) => MWC.Variate (Layer i o) where
    uniform g = Layer <$> MWC.uniform g <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g
  
  instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => MWC.Variate (Network i h1 h2 o) where
    uniform g = Net <$> MWC.uniform g <*> MWC.uniform g <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g
  
  instance Backprop (R n) where
    zero = zeroNum
    add = addNum
    one = oneNum
  
  instance (KnownNat n, KnownNat m) => Backprop (L m n) where
    zero = zeroNum
    add = addNum
    one = oneNum
  
  instance (KnownNat i, KnownNat o) => Backprop (Layer i o)
  instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Backprop (Network i h1 h2 o)
  