{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

import Control.Monad.Random
import Data.Binary                  as B
import Data.Singletons
import Data.Singletons.Prelude
import Data.Singletons.TypeLits
import GHC.Generics                 (Generic)
import Numeric.LinearAlgebra.Static

-- * Weights
--

data Weights i o = W { wBiases :: !(R o)
                     , wNodes  :: !(L o i)
                     }
  deriving (Show, Generic)

instance (KnownNat i, KnownNat o) => Binary (Weights i o)

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

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

-- * Network
--

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 :&~

deriving instance (KnownNat i, KnownNat o) => Show (Network i hs o)

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

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

hiddenStruct :: Network i hs o -> [Integer]
hiddenStruct = \case
    O _    -> []
    _ :&~ (n' :: Network h hs' o)
           -> natVal (Proxy @h)
            : hiddenStruct n'

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

randomNet :: forall m i hs o. (MonadRandom m, KnownNat i, SingI hs, KnownNat o)
          => m (Network i hs o)
randomNet = randomNet' sing

putNet :: (KnownNat i, KnownNat o)
       => Network i hs o
       -> Put
putNet = \case
    O w     -> put w
    w :&~ n -> put w *> putNet n

getNet :: forall i hs o. (KnownNat i, KnownNat o)
       => Sing hs
       -> Get (Network i hs o)
getNet = \case
    SNil            ->     O <$> get
    SNat `SCons` ss -> (:&~) <$> get <*> getNet ss

instance (KnownNat i, SingI hs, KnownNat o) => Binary (Network i hs o) where
    put = putNet
    get = getNet sing

-- * OpaqueNet
--

data OpaqueNet :: Nat -> Nat -> * where
    ONet :: Network i hs o -> OpaqueNet i o

runOpaqueNet :: (KnownNat i, KnownNat o)
             => OpaqueNet i o
             -> R i
             -> R o
runOpaqueNet (ONet n) x = runNet n x

numHiddens :: OpaqueNet i o -> Int
numHiddens (ONet n) = go n
  where
    go :: Network i hs o -> Int
    go = \case
        O _      -> 0
        _ :&~ n' -> 1 + go n'

randomONet :: (MonadRandom m, KnownNat i, KnownNat o)
           => [Integer]
           -> m (OpaqueNet i o)
randomONet hs = case toSing hs of
                  SomeSing ss -> ONet <$> randomNet' ss

putONet :: (KnownNat i, KnownNat o)
        => OpaqueNet i o
        -> Put
putONet (ONet net) = do
    put (hiddenStruct net)
    putNet net

getONet :: (KnownNat i, KnownNat o)
        => Get (OpaqueNet i o)
getONet = do
    hs <- get
    withSomeSing hs $ \ss ->
      ONet <$> getNet ss

instance (KnownNat i, KnownNat o) => Binary (OpaqueNet i o) where
    put = putONet
    get = getONet

-- * OpaqueNet'
--

type OpaqueNet' i o r = (forall hs. Network i hs o -> r) -> r

oNet' :: Network i hs o -> OpaqueNet' i o r
oNet' n = \f -> f n

runOpaqueNet' :: (KnownNat i, KnownNat o)
              => OpaqueNet' i o (R o)
              -> R i
              -> R o
runOpaqueNet' oN x = oN (\n -> runNet n x)
--            :: ((forall hs. Network i hs o -> R o) -> R o)
--            -> R i
--            -> R o

numHiddens' :: OpaqueNet' i o Int -> Int
numHiddens' oN = oN go
  where
    go :: Network i hs o -> Int
    go = \case
        O _      -> 0
        _ :&~ n' -> 1 + go n'
--          :: ((forall hs. Network i hs o -> Int) -> Int)
--          -> Int

withRandomONet' :: (MonadRandom m, KnownNat i, KnownNat o)
                => [Integer]
                -> (forall hs. Network i hs o -> m r)
                -> m r
--         aka, => [Integer]
--              -> OpaqueNet' i o (m r)
withRandomONet' hs f = withSomeSing hs $ \ss -> do
                         net <- randomNet' ss
                         f net

putONet' :: (KnownNat i, KnownNat o)
         => OpaqueNet' i o Put
         -> Put
putONet' oN = oN $ \net -> do
                      put (hiddenStruct net)
                      putNet net

getONet' :: (KnownNat i, KnownNat o)
         => (forall hs. Network i hs o -> Get r)
         -> Get r
--  aka, => OpaqueNet' i o (Get r)
getONet' f = do
    hs <- get :: Get [Integer]
    withSomeSing hs $ \ss -> do
      n <- getNet ss
      f n

main :: IO ()
main = do
    putStrLn "What hidden layer structure do you want?"
    hs <- readLn
    n  <- randomONet hs
    case n of
      ONet (net :: Network 10 hs 3) -> do
        print net
        -- blah blah stuff with our dynamically generated net

main' :: IO ()
main' = do
    putStrLn "What hidden layer structure do you want?"
    hs <- readLn
    withRandomONet' hs $ \(net :: Network 10 hs 3) -> do
      print net
      -- blah blah stuff with our dynamically generated net

-- * SomeNet
--

data SomeNet where
    SNet :: (KnownNat i, KnownNat o)
         => Network i hs o
         -> SomeNet

withONet :: SomeNet
         -> (forall i o. (KnownNat i, KnownNat o) => OpaqueNet i o -> r)
         -> r
withONet (SNet n) f = f (ONet n)

randomSNet :: forall m. MonadRandom m
           => Integer
           -> [Integer]
           -> Integer
           -> m SomeNet
randomSNet i hs o =
    withSomeSing i  $ \(SNat :: Sing (i  :: Nat  )) ->
    withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
    withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) ->
      SNet <$> (randomNet' ss :: m (Network i hs o))

withRandomSNet' :: forall m r. MonadRandom m
                => Integer
                -> [Integer]
                -> Integer
                -> (forall i hs o. (KnownNat i, KnownNat o) => Network i hs o -> m r)
                -> m r
withRandomSNet' i hs o f =
    withSomeSing i  $ \(SNat :: Sing (i  :: Nat  )) ->
    withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
    withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) -> do
      n <- randomNet' ss :: m (Network i hs o)
      f n

instance Binary SomeNet where
    put (SNet (net :: Network i hs o)) = do
        put $ natVal (Proxy @i)
        put $ hiddenStruct net
        put $ natVal (Proxy @o)
        putNet net
    get = do
        i  <- get :: Get Integer
        hs <- get :: Get [Integer]
        o  <- get :: Get Integer
        withSomeSing i    $ \(SNat :: Sing (i  :: Nat  )) ->
          withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
          withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) -> do
            n <- getNet ss :: Get (Network i hs o)
            return (SNet n)

