content: #!/usr/bin/env stack
  -- stack --install-ghc runghc --resolver lts-20.26 --package backprop-0.2.6.5 --package random --package hmatrix-backprop-0.1.3.0 --package one-liner-instances --package split --package ghc-typelits-natnormalise --package ghc-typelits-knownnat --package hmatrix --package hmatrix-vector-sized --package microlens --package vector-sized --package transformers
  
  -- | Replace the second line with this one to have "./model.hs" open a ghci session
  -- stack --install-ghc ghci --resolver lts-20.26 --package backprop-0.2.6.5 --package random --package hmatrix-backprop-0.1.3.0 --package one-liner-instances --package split --package ghc-typelits-natnormalise --package ghc-typelits-knownnat --package hmatrix --package hmatrix-vector-sized --package microlens --package vector-sized --package transformers
  
  {-# LANGUAGE DataKinds                                #-}
  {-# LANGUAGE DeriveGeneric                            #-}
  {-# LANGUAGE FlexibleContexts                         #-}
  {-# LANGUAGE FlexibleInstances                        #-}
  {-# LANGUAGE GADTs                                    #-}
  {-# LANGUAGE LambdaCase                               #-}
  {-# LANGUAGE MultiParamTypeClasses                    #-}
  {-# LANGUAGE PartialTypeSignatures                    #-}
  {-# LANGUAGE PatternSynonyms                          #-}
  {-# LANGUAGE RankNTypes                               #-}
  {-# LANGUAGE ScopedTypeVariables                      #-}
  {-# LANGUAGE TypeApplications                         #-}
  {-# LANGUAGE TypeInType                               #-}
  {-# LANGUAGE TypeOperators                            #-}
  {-# LANGUAGE ViewPatterns                             #-}
  {-# OPTIONS_GHC -fno-warn-orphans                     #-}
  {-# OPTIONS_GHC -fno-warn-partial-type-signatures     #-}
  {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
  {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}
  {-# OPTIONS_GHC -fwarn-redundant-constraints          #-}
  
  import           Control.Monad.Trans.State             ( runState, state )
  import           Data.Bifunctor                        ( first )
  import           Data.List                             ( foldl', unfoldr )
  import           Data.List.Split                       ( chunksOf )
  import           Data.Tuple                            ( swap )
  import           GHC.Generics                          ( Generic )
  import           GHC.TypeNats                          ( KnownNat, type (<=), type (+) )
  import           Lens.Micro                            ( Lens )
  
  import           Numeric.Backprop                      ( Backprop, BVar, Reifies, pattern T2, W
                                                         , (^^.), auto, evalBP, evalBP2, gradBP, isoVar2, sequenceVar )
  import           Numeric.LinearAlgebra.Static.Backprop ( L, R, (&), (#), (#>), headTail, konst, sumElements )
  import           Numeric.LinearAlgebra.Static.Vector   ( lVec, rVec, vecL, vecR )
  import           Numeric.OneLiner                      ( gAbs, gDivide, gFromInteger, gFromRational
                                                         , gMinus, gNegate, gPlus, gRecip, gSignum, gTimes )
  import           System.Random                         ( Random(random,randomR), randomIO )
  
  import qualified Data.Vector.Storable.Sized            as SVS
  import qualified Numeric.LinearAlgebra                 as HU
  import qualified Numeric.LinearAlgebra.Static          as H
  import qualified Prelude.Backprop                      as B
  
  data a :& b = !a :& !b
    deriving (Show, Generic)
  infixr 2 :&
  
  type Model p a b = forall z. Reifies z W
                  => BVar z p
                  -> BVar z a
                  -> BVar z b
  
  linReg :: Model (Double :& Double) Double Double
  linReg (a :&& b) x = b * x + a
  
  squaredErrorGrad
      :: (Backprop p, Backprop b, Num b)
      => Model p a b      -- ^ Model
      -> a                -- ^ Observed input
      -> b                -- ^ Observed output
      -> p                -- ^ Parameter guess
      -> p                -- ^ Gradient
  squaredErrorGrad f x targ = gradBP $ \p ->
      (f p (auto x) - auto targ) ^ 2
  
  trainModel
      :: (Fractional p, Backprop p, Num b, Backprop b)
      => Model p a b      -- ^ model to train
      -> p                -- ^ initial parameter guess
      -> [(a,b)]          -- ^ list of observations
      -> p                -- ^ updated parameter guess
  trainModel f = foldl' $ \p (x,y) -> p - 0.1 * squaredErrorGrad f x y p
  
  trainModelIO
      :: (Fractional p, Backprop p, Num b, Backprop b, Random p)
      => Model p a b      -- ^ model to train
      -> [(a,b)]          -- ^ list of observations
      -> IO p             -- ^ parameter guess
  trainModelIO m xs = do
      p0 <- (/ 10) . subtract 0.5 <$> randomIO
      return $ trainModel m p0 xs
  
  testTrainLinReg :: IO (Double :& Double)
  testTrainLinReg = trainModelIO linReg (concat (replicate 1000 samps))
    where
      samps = [(1,1),(2,3),(3,5),(4,7),(5,9)]
  
  logistic :: Floating a => a -> a
  logistic x = 1 / (1 + exp (-x))
  
  feedForward
      :: (KnownNat i, KnownNat o)
      => Model (L o i :& R o) (R i) (R o)
  feedForward (w :&& b) x = w #> x + b
  
  feedForwardLog
      :: (KnownNat i, KnownNat o)
      => Model (L o i :& R o) (R i) (R o)
  feedForwardLog (w :&& b) x = logistic (w #> x + b)
  
  testTrainPerceptron :: IO [R 1]
  testTrainPerceptron = do
      trained <- trainModelIO feedForwardLog $ take 10000 (cycle samps)
      print trained
      return [ evalBP2 feedForwardLog trained r | (r, _) <- samps ]
    where
      samps = [ (H.vec2 0 0, 0)
              , (H.vec2 1 0, 0)
              , (H.vec2 0 1, 0)
              , (H.vec2 1 1, 1)
              ]
  
  logReg :: Model (Double :& Double) Double Double
  logReg ab = logistic . linReg ab
  
  feedForwardLog'
      :: (KnownNat i, KnownNat o)
      => Model (L o i :& R o) (R i) (R o)
  feedForwardLog' wb = logistic . feedForward wb
  
  softMax :: (Reifies z W, KnownNat n) => BVar z (R n) -> BVar z (R n)
  softMax x = konst (1 / sumElements expx) * expx
    where
      expx = exp x
  
  feedForwardSoftMax
      :: (KnownNat i, KnownNat o)
      => Model (L o i :& R o) (R i) (R o)
  feedForwardSoftMax wb = softMax . feedForward wb
  
  (<~)
      :: (Backprop p, Backprop q)
      => Model  p       b c
      -> Model       q  a b
      -> Model (p :& q) a c
  (f <~ g) (p :&& q) = f p . g q
  infixr 8 <~
  
  testTrainTwoLayer :: IO [R 1]
  testTrainTwoLayer = do
      trained <- trainModelIO model (take 10000 (cycle samps))
      print trained
      return [ evalBP2 model trained r | (r, _) <- samps ]
    where
      model :: Model _ (R 2) (R 1)
      model = feedForwardLog' @4 @1 <~ feedForwardLog' @2 @4
      samps = [ (H.vec2 0 0, 0)
              , (H.vec2 1 0, 1)
              , (H.vec2 0 1, 1)
              , (H.vec2 1 1, 0)
              ]
  
  type ModelS p s a b = forall z. Reifies z W
                     => BVar z p
                     -> BVar z a
                     -> BVar z s
                     -> (BVar z b, BVar z s)
  
  ar2 :: ModelS (Double :& (Double :& Double)) Double Double Double
  ar2 (c :&& (φ1 :&& φ2)) yLast yLastLast =
      ( c + φ1 * yLast + φ2 * yLastLast, yLast )
  
  fcrnn
      :: (KnownNat i, KnownNat o)
      => ModelS ((L o i :& L o o) :& R o) (R o) (R i) (R o)
  fcrnn ((wX :&& wS) :&& b) x s = ( y, logistic y )
    where
      y  = (wX #> x) + (wS #> s) + b
  
  (<*~*)
    :: (Backprop p, Backprop q, Backprop s, Backprop t)
      => ModelS  p        s       b c
      -> ModelS       q        t  a b
      -> ModelS (p :& q) (s :& t) a c
  (f <*~* g) (p :&& q) x (s :&& t) = (z, s' :&& t')
    where
      (y, t') = g q x t
      (z, s') = f p y s
  infixr 8 <*~*
  
  mapS
      :: (forall z. Reifies z W => BVar z b -> BVar z c)
      -> ModelS p s a b
      -> ModelS p s a c
  mapS f g p x = first f . g p x
  
  toS :: Model  p   a b
      -> ModelS p s a b
  toS f p x s = (f p x, s)
  
  (<*~)
    :: (Backprop p, Backprop q)
      => Model   p         b c
      -> ModelS       q  s a b
      -> ModelS (p :& q) s a c
  (f <*~ g) (p :&& q) x = first (f p) . g q x
  infixr 8 <*~
  
  unroll
      :: (Backprop a, Backprop b)
      => ModelS p s  a   b
      -> ModelS p s [a] [b]
  unroll f p xs s0 = swap $ B.mapAccumL f' s0 xs
    where
      -- we have to re-arrange the order of arguments and tuple a bit to
      -- match what `mapAccumL` expects
      f' s x = swap (f p x s)
  
  unrollLast
      :: (Backprop a, Backprop b)
      => ModelS p s  a  b
      -> ModelS p s [a] b
  unrollLast f = mapS (last . sequenceVar) (unroll f)
  
  unrollLast'
      :: Backprop a
      => ModelS p s  a  b
      -> ModelS p s [a] b
  unrollLast' f p xs s0 = foldl' go (undefined, s0) (sequenceVar xs)
    where
      go (_, s) x = f p x s
  
  
  fixState
      :: s
      -> ModelS p s a b
      -> Model  p   a b
  fixState s0 f p x = fst $ f p x (auto s0)
  
  zeroState
      :: Num s
      => ModelS p s a b
      -> Model  p   a b
  zeroState = fixState 0
  
  trainState
      :: (Backprop p, Backprop s)
      => ModelS  p    s  a b
      -> Model  (p :& s) a b
  trainState f (p :&& s) x = fst $ f p x s
  
  prime
      :: Foldable t
      => ModelS p s a b     -- ^ model
      -> p                  -- ^ parameterization
      -> s                  -- ^ initial state
      -> t a                -- ^ priming input
      -> s                  -- ^ primed state
  prime f p = foldl' $ evalBP2 (\s x -> snd $ f (auto p) x s)
  
  feedback
      :: (Backprop a, Backprop s)
      => ModelS p s a a     -- ^ model
      -> p                  -- ^ parameterization
      -> s                  -- ^ initial state
      -> a                  -- ^ initial input
      -> [a]                -- ^ inifinite feedback loop
  feedback f p s0 x0 = unfoldr go (x0, s0)
    where
      go (x, s) = Just (x, (y, s'))
        where
          -- 'T2' tuples up a pair of 'BVar's into a 'BVar' of a tuple
          (y, s') = evalBP (uncurry T2 . f (auto p) (auto x)) s
  
  testAR2 :: IO [Double]
  testAR2 = do
      trained <- trainModelIO model $ take 10000 samps
      let primed = prime    model0 trained 0      (take 19 series)
          output = feedback model0 trained primed (series !! 19)
      print trained
      return $ take 200 output
    where
      -- sine wave with period 25
      series :: [Double]
      series = [ sin (2 * pi * t / 25) | t <- [0..]              ]
      samps  = [ (init c, last c)      | c <- chunksOf 19 series ]
      model0 :: ModelS _ _ Double Double
      model0 = ar2
      model  :: Model  _   [Double] Double
      model  = zeroState $ unrollLast model0
  
  testRNN :: IO [R 1]
  testRNN = do
      trained <- trainModelIO model $ take 100000 samps
      let primed = prime   model0  trained 0      (take 19 series)
          output = feedback model0 trained primed (series !! 19)
      return $ take 200 output
    where
      -- sine wave with period 25
      series :: [H.R 1]
      series = [ H.konst (sin (2 * pi * t / 25)) | t <- [0..]              ]
      samps  = [ (init c, last c)                | c <- chunksOf 19 series ]
      model0 :: ModelS _ _ (R 1) (R 1)
      model0 = feedForward @30 @1
           <*~ mapS logistic (fcrnn @1 @30)
      model  :: Model  _   [R 1] (R 1)
      model  = zeroState $ unrollLast model0
  
  recurrently
      :: (Backprop a, Backprop b)
      => Model  p   (a :& b) b
      -> ModelS p b  a       b
  recurrently f p x yLast = (y, y)
    where
      y = f p (x :&& yLast)
  
  recurrentlyWith
      :: (Backprop a, Backprop b)
      => (forall z. Reifies z W => BVar z c -> BVar z b)
      -> Model  p   (a :& b) c
      -> ModelS p b  a       c
  recurrentlyWith store f p x yLast = (y, store y)
    where
      y = f p (x :&& yLast)
  
  ffOnSplit
      :: forall i o. (KnownNat i, KnownNat o)
      => Model _ (R i :& R o) (R o)
  ffOnSplit p (rI :&& rO) = feedForward p (rI # rO)
  
  fcrnn'
      :: (KnownNat i, KnownNat o)
      => ModelS _ (R o) (R i) (R o)
  fcrnn' = recurrentlyWith logistic (\p -> feedForward p . uncurryT (#))
  
  lagged
      :: (KnownNat n, 1 <= n)
      => Model  p       (R (n + 1)) b
      -> ModelS p (R n) Double      b
  lagged f p x xLasts = (y, xLasts')
    where
      fullLasts    = xLasts & x
      y            = f p fullLasts
      (_, xLasts') = headTail fullLasts
  
  ar :: (KnownNat n, 1 <= n)
     => ModelS _ (R n) Double Double
  ar = lagged (\p -> fst . headTail . feedForward @_ @1 p)
  
  ar2' :: ModelS _ (R 2) Double Double
  ar2' = ar @2
  
  main :: IO ()
  main = do
      putStrLn "Linear regression"
      print =<< testTrainLinReg
  
      putStrLn "Single layer perceptron learning AND"
      mapM_ print =<< testTrainPerceptron
  
      putStrLn "Two-layer ANN learning XOR"
      mapM_ print =<< testTrainTwoLayer
  
      putStrLn "Sine (AR2)"
      ar2Test <- testAR2
      mapM_ print (take 30 ar2Test)
      writeFile "ar2sin.dat" $ unlines (show <$> ar2Test)
  
      putStrLn "Sine (RNN)"
      rnnTest <- testRNN
      mapM_ print (take 30 rnnTest)
      writeFile "rnnsin.dat" $ unlines (show . HU.sumElements . H.extract <$> rnnTest)
  
  pattern (:&&) :: (Backprop a, Backprop b, Reifies z W)
                => BVar z a -> BVar z b -> BVar z (a :& b)
  pattern x :&& y <- (\xy -> (xy ^^. t1, xy ^^. t2)->(x, y))
    where
      (:&&) = isoVar2 (:&) (\case x :& y -> (x, y))
  {-# COMPLETE (:&&) #-}
  
  t1 :: Lens (a :& b) (a' :& b) a a'
  t1 f (x :& y) = (:& y) <$> f x
  
  t2 :: Lens (a :& b) (a :& b') b b'
  t2 f (x :& y) = (x :&) <$> f y
  
  instance (Num a, Num b) => Num (a :& b) where
      (+)         = gPlus
      (-)         = gMinus
      (*)         = gTimes
      negate      = gNegate
      abs         = gAbs
      signum      = gSignum
      fromInteger = gFromInteger
  
  instance (Fractional a, Fractional b) => Fractional (a :& b) where
      (/) = gDivide
      recip = gRecip
      fromRational = gFromRational
  
  instance (Random a, Random b) => Random (a :& b) where
      random g0 = (x :& y, g2)
        where
          (x, g1) = random g0
          (y, g2) = random g1
      randomR (x0 :& y0, x1 :& y1) g0 = (x :& y, g2)
        where
          (x, g1) = randomR (x0, x1) g0
          (y, g2) = randomR (y0, y1) g1
  
  instance (Backprop a, Backprop b) => Backprop (a :& b)
  
  uncurryT
      :: (Backprop a, Backprop b, Reifies z W)
      => (BVar z a -> BVar z b -> BVar z c)
      -> BVar z (a :& b)
      -> BVar z c
  uncurryT f x = f (x ^^. t1) (x ^^. t2)
  
  instance (KnownNat n, KnownNat m) => Random (L n m) where
      random = runState . fmap vecL $ SVS.replicateM (state random)
      randomR (xs,ys) = runState . fmap vecL $ SVS.zipWithM (curry (state . randomR))
          (lVec xs) (lVec ys)
  
  instance (KnownNat n) => Random (R n) where
      random = runState $ vecR <$> SVS.replicateM (state random)
      randomR (xs,ys) = runState . fmap vecR $ SVS.zipWithM (curry (state . randomR))
          (rVec xs) (rVec ys)
  
  