content: {-# LANGUAGE DataKinds            #-}
  {-# LANGUAGE FlexibleContexts     #-}
  {-# LANGUAGE GADTs                #-}
  {-# LANGUAGE KindSignatures       #-}
  {-# LANGUAGE LambdaCase           #-}
  {-# LANGUAGE PolyKinds            #-}
  {-# LANGUAGE QuasiQuotes          #-}
  {-# LANGUAGE RankNTypes           #-}
  {-# LANGUAGE RecordWildCards      #-}
  {-# LANGUAGE ScopedTypeVariables  #-}
  {-# LANGUAGE TemplateHaskell      #-}
  {-# LANGUAGE TypeApplications     #-}
  {-# LANGUAGE TypeInType           #-}
  {-# LANGUAGE TypeOperators        #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  import           Control.Category
  import           Control.DeepSeq
  import           Control.Exception
  import           Control.Monad
  import           Data.Foldable
  import           Data.Kind
  import           Data.List hiding                      ((\\))
  import           Data.Maybe
  import           Data.Singletons
  import           Data.Monoid
  import           Data.Singletons.TypeLits
  import           Data.Time.Clock
  import           Options.Applicative
  import           Prelude hiding                        ((.), id)
  import           Statistics.Distribution.Uniform
  import           System.Random.MWC
  import           TensorOps.Backend.BTensor
  import           TensorOps.Backend.NTensor
  import           TensorOps.Learn.NeuralNet
  import           TensorOps.Learn.NeuralNet.FeedForward
  import           TensorOps.NatKind
  import           TensorOps.Types
  import           Text.PrettyPrint.ANSI.Leijen hiding   ((<>),(<$>))
  import           Text.Printf
  import           Type.Class.Higher.Util
  import qualified Data.String.Here                      as H
  import qualified TensorOps.Tensor                      as TT
  
  
  netTest
      :: forall k (t :: [k] -> Type).
       ( Tensor t
       , ElemT t ~ Double
       , NFData1 t
       , NFData (t '[FromNat 1])
       , NFData (t '[FromNat 2])
       )
      => Proxy t
      -> Double
      -> Int
      -> [Integer]
      -> GenIO
      -> IO String
  netTest _ rate n hs g = withSingI (sFromNat @k (SNat @1)) $
                          withSingI (sFromNat @k (SNat @2)) $ do
      ((inps,outs),tp) <- time $ do
        inps :: [t '[FromNat 2]] <- replicateM n (genRand (uniformDistr (-1) 1) g)
        let outs :: [t '[FromNat 1]]
            outs = flip map inps $ \v -> TT.konst $
                     if v `inCircle` (TT.konst 0.33, 0.33)
                          || v `inCircle` (TT.konst (-0.33), 0.33)
                       then 1
                       else 0
        evaluate . force $ (inps, outs)
      printf "Generated test points (%s)\n" (show tp)
      net0 :: Network t (FromNat 2) (FromNat 1)
              <- genNet (hs `zip` repeat actLogistic) actLogistic g
      let trained = foldl' trainEach net0 (zip inps outs)
            where
              trainEach :: SingI o
                        => Network t i o
                        -> (t '[i], t '[o])
                        -> Network t i o
              trainEach nt (i, o) = nt `deepseq` trainNetwork squaredError rate i o nt
      (trained', tn) <- time $ return trained
      printf "Network trained (%s)\n" (show tn)
      let outMat = [ [ render . TT.unScalar . join TT.dot . runNetwork trained' $
                         fromJust (TT.fromList [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
          :: SingI n
          => t '[n]
          -> (t '[n], Double)
          -> Bool
      v `inCircle` (o, r) = let d = TT.zip (-) v o
                            in  TT.unScalar (d `TT.dot` d) <= r**2
  
  
  data Opts = O { oRate    :: Double
                , oSamples :: Int
                , oNetwork :: [Integer]
                , oTests   :: [TestT]
                }
  
  opts :: Parser Opts
  opts = O <$> option auto
                 ( long "rate" <> short 'r' <> metavar "STEP"
                <> help "Neural network learning rate"
                <> value 1 <> showDefault
                 )
           <*> option auto
                 ( long "samps" <> short 's' <> metavar "COUNT"
                <> help "Number of samples to train the network on"
                <> value 50000 <> showDefault
                 )
           <*> option auto
                 ( long "layers" <> short 'l' <> metavar "LIST"
                <> help "List of hidden layer sizes"
                <> value [12,8] <> showDefault
                 )
           <*> (nub <$> (some (argument readBackend (metavar "BACKEND")))
                 <|> pure [TTBLAS TBHMat]
               )
  
  main :: IO ()
  main = withSystemRandom $ \g -> do
      O{..} <- execParser $ info (helper <*> opts)
          ( fullDesc
         <> header "tensor-ops-dots - train neural nets with tensor-ops"
         <> progDescDoc (Just d)
          )
  
      printf "rate: %f | samps: %d | layers: %s\n" oRate oSamples (show oNetwork)
  
      forM_ oTests $ \t -> do
        printf "Training %s network ...\n" (ttLong t)
        let tester :: Double -> Int -> [Integer] -> GenIO -> IO String
            tester = case t of
              TTNested TVList -> netTest (Proxy @NTensorL)
              TTNested TVVect -> netTest (Proxy @NTensorV)
              TTBLAS   TBHMat -> netTest (Proxy @(BTensorV (HMat Double)))
        putStrLn =<< tester oRate oSamples oNetwork g
    where
      d :: Doc
      d = string [H.here|
  Trains a feed-forward neural network on a 2D classifier using tensor-ops
  machinery, with the given backends.  (If none provided, backend defaults to 'b')
  |]
       <$$> mempty
       <$$> string "Backends:"
       <$$> vsep [ string $ printf "- %s: %s" (ttShort t) (ttLong t)
                 | t <- allTests ]
  
  time
      :: NFData a
      => IO a
      -> IO (a, NominalDiffTime)
  time x = do
      t1 <- getCurrentTime
      y  <- evaluate . force =<< x
      t2 <- getCurrentTime
      return (y, t2 `diffUTCTime` t1)
  
  
  data TestT = TTNested TestV
             | TTBLAS   TestB
             deriving (Show, Eq, Ord)
  data TestV = TVList
             | TVVect
             deriving (Show, Eq, Ord)
  data TestB = TBHMat
             deriving (Show, Eq, Ord)
  
  allTests :: [TestT]
  allTests = ([TTNested] <*> [TVList, TVVect]) ++ [TTBLAS TBHMat]
  
  readBackend :: ReadM TestT
  readBackend = eitherReader $ \s -> case s of
      "nl" -> Right $ TTNested TVList
      "nv" -> Right $ TTNested TVVect
      "b"  -> Right $ TTBLAS   TBHMat
      o    -> Left  $ "Unknown backend: " ++ o
  
  ttShort
      :: TestT
      -> String
  ttShort = \case
      TTNested v -> 'n' : tvShort v
      TTBLAS   b -> 'b' : tbShort b
    where
      tvShort = \case
        TVList -> "l"
        TVVect -> "v"
      tbShort = \case
        TBHMat -> ""
  
  ttLong
      :: TestT
      -> String
  ttLong = \case
      TTNested v -> printf "Nested (%s)" (tvLong v)
      TTBLAS   b -> printf "BLAS (%s)" (tbLong b)
    where
      tvLong = \case
        TVList -> "List"
        TVVect -> "Vector"
      tbLong = \case
        TBHMat -> "HMatrix"
  
  