{-# LANGUAGE ApplicativeDo          #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DeriveGeneric          #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE ViewPatterns           #-}

import           Control.Arrow                         ((&&&))
import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.State.Strict
import           Control.Monad.Trans.Writer.Strict
import           Data.Bifunctor
import           Data.Either.Validation
import           Data.Finite
import           Data.Foldable
import           Data.IDX
import           Data.Kind
import           Data.List
import           Data.Monoid
import           Data.Profunctor
import           Data.Proxy
import           Data.Singletons
import           Data.Singletons.Prelude               (Sing(..))
import           Data.Singletons.TypeLits
import           Data.String
import           Data.Time.Clock
import           Data.Type.Combinator
import           Data.Type.Equality
import           Data.Type.Vector
import           GHC.Generics                          (Generic)
import           GHC.TypeLits                          as TL
import           GHC.TypeLits.Compare
import           Options.Applicative hiding            (ParserResult(..))
import           Statistics.Distribution
import           Statistics.Distribution.Uniform
import           System.Directory
import           System.FilePath
import           System.Random.MWC
import           System.Random.MWC.Distributions
import           TensorOps.Backend.BTensor
import           TensorOps.Learn.NeuralNet
import           TensorOps.Learn.NeuralNet.FeedForward
import           TensorOps.Types hiding                ((&&&))
import           Text.Printf
import           Type.Class.Higher.Util
import           Type.Family.Nat
import qualified Codec.Compression.GZip                as GZ
import qualified Control.Foldl                         as F
import qualified Data.ByteString.Lazy                  as BS
import qualified Data.Map.Strict                       as M
import qualified Data.Vector                           as V
import qualified Data.Vector.Unboxed                   as VU
import qualified Network.HTTP.Simple                   as HTTP
import qualified TensorOps.Tensor                      as TT
import qualified Text.PrettyPrint.Boxes                as B

mnistBase :: String
mnistBase = "http://yann.lecun.com/exdb/mnist"

mnistFiles :: Matrix '[N2, N2] FilePath
mnistFiles = ("train-images-idx3-ubyte" :+ "train-labels-idx1-ubyte" :+ ØV)
          :* ("t10k-images-idx3-ubyte"  :+ "t10k-labels-idx1-ubyte"  :+ ØV)
          :* ØV

data Opts = O { oRate        :: Double
              , oLayers      :: [Integer]
              , oBatch       :: Integer
              , oDataDir     :: FilePath
              , oNoConfusion :: Bool
              , oWhite       :: Bool
              , oInduce      :: Maybe (Finite 10)
              }
    deriving (Show, Eq, Generic)

opts :: Parser Opts
opts = O <$> option auto
               ( long "rate" <> short 'r' <> metavar "STEP"
              <> help "Neural network learning rate"
              <> value 0.02 <> showDefault
               )
         <*> option auto
               ( long "layers" <> short 'l' <> metavar "LIST"
              <> help "List of hidden layer sizes"
              <> value [300,100] <> showDefault
               )
         <*> option auto
               ( long "batch" <> short 'b' <> metavar "AMOUNT"
              <> help "Training batch size"
              <> value 1000 <> showDefault
               )
         <*> strOption
               ( long "data" <> short 'd' <> metavar "PATH"
              <> help "Directory to store/cache MNIST data files"
              <> value "data/mnist" <> showDefaultWith id
               )
         <*> switch
               ( long "noconfusion" <> short 'c'
              <> help "Disable confusion matrix validation and only display % error every batch"
               )
         <*> switch
               ( long "white" <> short 'w'
              <> help "Train with an eleventh \"white noise\" class to train network on negative results"
               )
         <*> optional (
               option readFin
                 ( long "induce" <> short 'i' <> metavar "DIGIT"
                <> help ("Every batch, attempt to induce an image "
                      ++ "of the given digit with the trained network"
                        )
                 )
               )
  where
    readFin :: forall n. KnownNat n => ReadM (Finite n)
    readFin = do
        i <- auto
        case packFinite i of
          Nothing -> readerError $
            printf "Number %d out of range (%d)" i (natVal (Proxy @n) - 1)
          Just x  -> return x

main :: IO ()
main = do
    O{..} <- execParser $ info (helper <*> opts)
        ( fullDesc
       <> header "tensor-ops-mnist - train neural nets on MNIST data set"
       <> progDesc (unlines ["Simple test of tensor-ops tensors (with hmatrix"
                            ,"backend) on MNIST classification challenge"
                            ]
                   )
        )

    mnistDat <- loadData oDataDir
    putStrLn "Loaded data."

    withSomeSing oWhite $ \(w :: Sing (w :: Bool)) ->
      withKnownNat (nOut w) $
        case Proxy %<=? Proxy :: 1 :<=? NOut w of
          LE Refl -> case Proxy %<=? Proxy :: 10 :<=? NOut w of
            LE Refl ->
              learn w (Proxy @(BTensorV HMatD)) mnistDat
                oRate oLayers oBatch oInduce oNoConfusion
            NLE _ -> error "impossible"
          NLE _ -> error "impossible"

loadData
    :: FilePath
    -> IO (Vec N2 [(Int, VU.Vector Int)])
loadData dataDir = do
    createDirectoryIfMissing True dataDir

    printf "Loading data from %s\n" dataDir
    mnistBs <- forM mnistFiles $ \u -> do
      BS.readFile (dataDir </> u) `catch` \(_ :: IOException) -> do
        printf "'%s' not found; downloading from %s ...\n" u mnistBase
        b <- GZ.decompress . HTTP.getResponseBody
               <$> HTTP.httpLBS (fromString (mnistBase </> u -<.> "gz"))
        BS.writeFile (dataDir </> u) b
        return b

    let f   :: Vec N2 FilePath
            -> Vec N2 BS.ByteString
            -> I (Validation [String] (FilePath, FilePath, IDXData, IDXLabels))
        f (I uIm :* I uLb :* ØV) (I bIm :* I bLb :* ØV) = I $ do
            im <- maybe (Failure [printf "Could not decode image %s." uIm]) Success
                    $ decodeIDX bIm
            lb <- maybe (Failure [printf "Could not decode labels %s." uLb]) Success
                    $ decodeIDXLabels bLb
            pure (uIm, uLb, im, lb)

        mnistDat' :: Either [String] (Vec N2 [(Int, VU.Vector Int)])
        mnistDat' = do
          imlb <- validationToEither . sequenceA $ vap f mnistFiles mnistBs
          forM imlb $ \(uIm, uLb, im, lb) ->
            maybe (Left [printf "Could not combine %s and %s." uIm uLb]) Right
              $ labeledIntData lb im

    evaluate . force
        =<< either (ioError . userError . unlines) return mnistDat'

processDat
    :: forall (n :: Nat) (l :: Nat) t.
     ( Fractional (ElemT t)
     , KnownNat n
     , KnownNat l
     , Tensor t
     )
    => (Int, VU.Vector Int)
    -> Either String (t '[n], (t '[l], Finite l))
processDat (l,d) = (,) <$> x <*> y
  where
    x :: Either String (t '[n])
    x = maybe (Left (printf "Bad input vector (Expected %d, got %d)" n (VU.length d))) Right
      . TT.fromList . map ((/255) . fromIntegral)
      $ VU.toList d
    y :: Either String (t '[l], Finite l)
    y = maybe (Left (printf "Label out of range (Got %d, expected [0,%d) )" l')) Right
      . fmap (TT.oneHot 1 0 &&& id)
      $ packFinite (fromIntegral l)
    n :: Integer
    n = natVal (Proxy @n)
    l' :: Integer
    l' = natVal (Proxy @l)

type family NOut (w :: Bool) = (n :: Nat) | n -> w where
    NOut 'False = 10
    NOut 'True  = 11

nOut
    :: Sing w
    -> Sing (NOut w)
nOut = \case
    SFalse -> SNat @10
    STrue  -> SNat @11

learn
    :: forall (t :: [Nat] -> Type) (w :: Bool) (o :: Nat).
     ( Tensor t
     , RealFloat (ElemT t)
     , NFData1 t
     , NFData (t '[784])
     , NFData (t '[o])
     , o ~ NOut w
     , 10 TL.<= o
     , 1  TL.<= o
     , KnownNat o
     )
    => Sing w
    -> Proxy t
    -> Vec N2 [(Int, VU.Vector Int)]
    -> Double
    -> [Integer]
    -> Integer
    -> Maybe (Finite 10)
    -> Bool
    -> IO ()
learn w _ dat rate layers (fromIntegral->batch) ind nc =
      withSystemRandom $ \g -> do
    dat' <- either (ioError . userError . unlines) return
          . validationToEither
          . (traverse . traverse) processDat'
          $ dat

    dat'' <- evaluate $ force dat'
    putStrLn "Data processed."

    let tXY, vXY :: [(t '[784], (t '[o], Finite o))]
        (tXY, vXY) = case dat'' of
                       I t :* I v :* ØV -> (t,v)

    net0 :: Network t 784 o
            <- genNet (layers `zip` repeat (actMap logistic)) actSoftmax g

    printf "rate: %f | batch: %d | layers: %s\n" rate batch (show layers)
    when (fromSing w) $
      putStrLn "white noise class enabled"
    forM_ ind $ \i ->
      printf "inducing: %d\n" (getFinite i)

    trainEpochs tXY ((map . second) snd vXY) g net0
  where
    processDat'
        :: (Int, VU.Vector Int)
        -> Validation [String] (t '[784], (t '[o], Finite o))
    processDat' = eitherToValidation . first (:[]) . processDat
    noiseClass :: t '[11]
    noiseClass = TT.oneHot 1 0 (finite 10)
    noiseFin   :: Finite 11
    noiseFin   = finite 10
    ind' :: Maybe (t '[o])
    ind' = TT.oneHot 1 0 . weakenN <$> ind
    trainEpochs
        :: [(t '[784], (t '[o], Finite o))]
        -> [(t '[784], Finite o)]
        -> GenIO
        -> Network t 784 o
        -> IO ()
    trainEpochs (V.fromList->tr) vd g = trainEpoch 1
      where
        trainEpoch
            :: Integer
            -> Network t 784 o
            -> IO ()
        trainEpoch e nt0 = do
            printf "[Epoch %d]\n" e
            tr' <- case w of
                SFalse -> return tr
                STrue  -> do
                  extr <- V.replicateM (V.length tr `div` 10) $ do
                            v <- genRand (uniformDistr 0 1) g
                            α <- realToFrac <$> genContVar (uniformDistr 0 1) g
                            return $ scaleT α v
                  return $ tr <> fmap (, (noiseClass, noiseFin)) extr

            queue <- evaluate . force =<< uniformShuffle tr' g

            printf "Training on %d samples in batches of %d ...\n" (V.length tr') batch

            nt1 <- trainBatch 1 queue nt0
            trainEpoch (succ e) nt1
          where
            trainBatch
                :: Integer
                -> V.Vector (t '[784], (t '[o], Finite o))
                -> Network t 784 o
                -> IO (Network t 784 o)
            trainBatch b (V.splitAt batch->(xs,xss)) nt
                | V.null xs = return nt
                | otherwise = do
              printf "Batch %d ...\n" b
              (nt', t) <- time . return $ trainAll nt ((fmap . second) fst xs)
              printf "Trained on %d samples in %s\n" (V.length xs) (show t)
              vd' :: [(t '[784], Finite o)] <- case w of
                  SFalse ->
                    return vd
                  STrue  -> do
                    extr <- replicateM (length vd `div` 10) $ do
                              v <- genRand (uniformDistr 0 1) g
                              α <- realToFrac <$> genContVar (uniformDistr 0 1) g
                              return $ scaleT α v
                    return $ vd <> fmap (, noiseFin) extr
              let tscore = F.fold (validate nt') ((fmap . second) snd xs)
                  vconf  = F.fold (confusion nt') vd'
                  confmat = ( B.vcat B.left
                                  [ B.text (printf "[%d]" (getFinite i))
                                  | i <- range :: [Finite o]
                                  ] B.<+>
                            )
                          . B.hsep 1 B.top
                          . map (\(_,r) -> B.vcat B.right
                                         . map (\(_, c) -> B.text (show c))
                                         $ M.toList r
                                )
                          $ M.toList vconf
                  vscore :: Double
                  vscore
                    | nc        = F.fold (validate nt') vd'
                    | otherwise = (/ fromIntegral (length vd')) . sum
                                . map (\(i, as) -> fromIntegral $ as M.! i)
                                $ M.toList vconf
              printf "Training:   %.2f%% error\n" ((1 - tscore) * 100)
              printf "Validation: %.2f%% error\n" ((1 - vscore) * 100)
              B.printBox confmat
              forM_ ind' $ \i -> do
                x0 <- genRand (uniformDistr 0 0.05) g
                let x1 = induceNum nt' i 1 5000 x0
                    y1 = intercalate "/"
                       . map (printf "%.2f" . realToFrac @_ @Double)
                       . TT.toList
                       $ runNetwork nt' x1
                putStrLn (renderOut x1)
                putStrLn y1

              trainBatch (succ b) xss nt'
        validate
            :: Network t 784 o
            -> F.Fold (t '[784], Finite o) Double
        validate n = (\s l -> fromIntegral s / fromIntegral l)
                 <$> lmap (uncurry f) F.sum
                 <*> F.length
          where
            f :: t '[784] -> Finite o -> Int
            f x r | TT.argMax (runNetwork n x) == r = 1
                  | otherwise                       = 0
        confusion
            :: Network t 784 o
            -> F.Fold (t '[784], Finite o) (M.Map (Finite o) (M.Map (Finite o) Integer))
        confusion n = F.Fold (\m (x, r) ->
                                let y = TT.argMax (runNetwork n x)
                                in  M.insertWith (M.unionWith (+)) y (M.singleton r 1) m
                             )
                             (M.fromList [ (i, M.fromList [ (j, 0) | j <- range ] )
                                         | i <- range
                                         ]
                               )
                             id
    trainAll
        :: Foldable f
        => Network t 784 o
        -> f (t '[784], t '[o])
        -> Network t 784 o
    trainAll = foldl' $ \nt (i,o) -> nt `deepseq`
        trainNetwork crossEntropy rate' i o nt
    rate' :: ElemT t
    rate' = realToFrac rate
    induceNum
        :: Network t 784 o
        -> t '[o]
        -> ElemT t
        -> Int
        -> t '[784]
        -> t '[784]
    induceNum n t r = go
      where
        go i x
          | i == 0    = x
          | otherwise = let x' = induceNetwork crossEntropy r t n x
                        in  x' `deepseq` go (i - 1) x'

time
    :: NFData a
    => IO a
    -> IO (a, NominalDiffTime)
time x = do
    t1 <- getCurrentTime
    y  <- evaluate . force =<< x
    t2 <- getCurrentTime
    return (y, t2 `diffUTCTime` t1)

renderOut
    :: forall t. (Tensor t, Real (ElemT t))
    => t '[784]
    -> String
renderOut = intercalate "\n"
          . ($ [])
          . appEndo
          . execWriter
          . execStateT (replicateM 28 go)
          . TT.toList
  where
    go :: StateT [ElemT t] (Writer (Endo [String])) ()
    go = do
        x <- state $ splitAt 28
        lift . tell . Endo . (++) . (:[]) $
            map (render . realToFrac)
          . concatMap (\y -> [y,y])
          $ x
    render :: Double -> Char
    render r | r <= 0.2  = ' '
             | r <= 0.4  = '.'
             | r <= 0.8  = '-'
             | r <= 1.9  = '='
             | otherwise = '#'

range
    :: forall n. KnownNat n
    => [Finite n]
range = finite <$> [0 .. natVal (Proxy @n) - 1]

