content: {-# LANGUAGE AllowAmbiguousTypes, CPP, ConstraintKinds, DataKinds, FlexibleContexts, FlexibleInstances, GADTs, MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedLabels, OverloadedStrings, PartialTypeSignatures, PolyKinds, QuasiQuotes, RankNTypes, ScopedTypeVariables     #-}
  {-# LANGUAGE TypeApplications, TypeFamilyDependencies, TypeInType, TypeOperators, UndecidableInstances                               #-}
  {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver #-}
  {-# OPTIONS_GHC -Wno-partial-type-signatures -fconstraint-solver-iterations=10 #-}
  
  {-| Talk to Tensorboard. You will need to run it manually with
  'tensorboard --logdir <yourdir>'
  
  It will occasionally update on its own.  We have fairly comprehensive
  Tensorboard support and integration with showing images, plots, and graphs.
  -}
  
  module Torch.Tensorboard where
  import           Control.Monad
  import           Control.Monad.Logger
  import           Data.Binary.Put
  import           Data.Bits
  import qualified Data.ByteString              as B
  import qualified Data.ByteString              as BS
  import qualified Data.ByteString.Internal     as BS
  import qualified Data.ByteString.Lazy         as BL
  import           Data.Coerce
  import           Data.IORef
  import qualified Data.Map                     as M
  import           Data.ProtoLens
  import           Data.ProtoLens.Labels        ()
  import           Data.Singletons
  import           Data.Singletons.TypeLits
  import           Data.Text                    (Text)
  import qualified Data.Text                    as T
  import qualified Data.Text.Encoding           as T
  import           Data.Time.Clock.POSIX
  import           Data.Time.Format
  import           Data.Time.LocalTime
  import           Data.Vector.Storable         (Vector)
  import qualified Data.Vector.Storable         as V
  import           Foreign.C.Types
  import           Foreign.ForeignPtr
  import           Foreign.Ptr
  import           GHC.Float
  import           GHC.Int
  import           GHC.Word
  import qualified Graphics.Matplotlib          as M
  import           Lens.Micro
  import           Network.HostName
  import qualified Statistics.Quantile          as S
  import           System.Directory
  import           System.FilePath
  import           System.IO
  import qualified Tensorboard.Proto.Attributes as PA
  import qualified Tensorboard.Proto.Event      as PE
  import qualified Tensorboard.Proto.Graph      as PG
  import qualified Tensorboard.Proto.Summary    as PS
  import qualified Tensorboard.Proto.Tensor     as PT
  import qualified Torch.C.Tensor               as C
  import qualified Torch.C.Types                as C
  import qualified Torch.C.Variable             as C
  import           Torch.Images
  import           Torch.Internal.CRC32C
  import           Torch.Tensor                 as Tensor
  import           Torch.Types                  as Tensor
  import           Torch.Visualization
  
  -- | `SummaryWriter` keeps track of some basic information so we can talk to
  -- Tensorboard. You'll be passing this around everywhere.
  data SummaryWriter = SummaryWriter { runName      :: Text
                                     , runDirectory :: Text
                                     , runFilename  :: Text
                                     , lastStep     :: IORef Int
                                     }
  
  -- | Make a new summary writer for a given directory and run name. We will also
  -- add some information to the run time like the current host and local timestamp.
  summaryWriter :: Text -> Text -> IO SummaryWriter
  summaryWriter directory runname = do
    now <- getCurrentTime
    tz <- getCurrentTimeZone
    let time = formatTime defaultTimeLocale "%Y-%m-%d_%H:%M:%S" (utcToZonedTime tz now)
    let timeUnix = formatTime defaultTimeLocale "%s" (utcToZonedTime tz now)
    host <- T.pack <$> getHostName
    let name = runname <> "_" <> host <> "_" <> T.pack time
    let location = T.unpack directory </> T.unpack name
    createDirectoryIfMissing True location
    step <- newIORef 0
    pure $ SummaryWriter { runName = name
                         , runDirectory = T.pack location
                         , runFilename = "events.haskelltorch." <> T.pack timeUnix <> ".tfevents." <> host
                         , lastStep = step
                         }
  
  -- | Advance the writer to the next step. This helps you keep track of the
  -- iteration number without needing a manual counter.
  nextStep :: SummaryWriter -> IO SummaryWriter
  nextStep r = do
    modifyIORef' (lastStep r) (+ 1)
    pure r
  
  -- | Explicitly set the step counter.
  setStep :: SummaryWriter -> Int -> IO SummaryWriter
  setStep sw s = do
    writeIORef (lastStep sw) s
    pure sw
  
  -- | Add a scalar to Tensorboard with a given name and value.
  addScalar :: Real a => SummaryWriter -> Text -> a -> IO ()
  addScalar summary name value =
    writeEvent summary (EventSummary name $ SummaryValue $ fromRational $ toRational value)
  
  -- | Add a histogram to Tensorboard, we do some internal work to make sure the
  -- number of bins and their sizes make sense.
  addHistogram ::
    (TensorConstraints ty 'KCpu sz) =>
    SummaryWriter -> Text -> Tensor ty 'KCpu sz -> IO ()
  addHistogram summary name tensor = do
    bins  <- histogramBinsAuto <$> V.map toDoubleC <$> toVector tensor
    withSomeSing (fromIntegral bins) (\ n@(SNat :: SNat m) -> do
      hist_ <- toVector =<< histc @m tensor 0 0
      let hist = V.map toDoubleC hist_
      min   <- toDouble <$> (fromScalar =<< Tensor.min tensor)
      max   <- toDouble <$> (fromScalar =<< Tensor.max tensor)
      let step = (max-min)/2
      let binLimits = V.generate bins (\i -> step * fromIntegral i)
      sum     <- fromScalar =<< Tensor.sum tensor
      sumSq <- fromScalar =<< Tensor.sum =<< Tensor.pow tensor =<< toScalar 2
      writeEvent summary
        (EventSummary name
         (SummaryHistogram
           { shMin = min
           , shMax = max
           , shNum = fromIntegral $ numel tensor
           , shSum = toDouble sum
           , shSumSq = toDouble sumSq
           , shBucketLimit = binLimits
           , shBucket = hist })))
  
  -- | Add an RGB image tensor to the Tensorboard.
  addImageRGB :: (KnownNat szw, KnownNat szh, SingI ty)
              => SummaryWriter -> Text -> Tensor ty 'KCpu '[3, szh, szw] -> IO ()
  addImageRGB summary name image = do
      (bytes, len) <- V.unsafeToForeignPtr0 <$> rgbTensorToVector image
      (h,w,c,b) <- pure $ case V.toList $ size image of
         [3,h,w] -> (h,w,3,bytes)
         -- TODO Move this to compile time
         _       -> error $ "Don't know how to display a tensor with these dimensions as an RGB image " ++ show (size image)
      writeEvent summary
        (EventSummary name
         (SummaryImage
          { siHeight     = fromIntegral h
          , siWidth      = fromIntegral w
          , siColorspace = c
          , siBytes      = BS.PS b 0 len}))
  
  -- | Add a greyscale image tensor to the Tensorboard.
  addImageGrey :: (KnownNat szw, KnownNat szh, SingI ty)
               => SummaryWriter -> Text -> Tensor ty 'KCpu '[1, szh, szw] -> IO ()
  addImageGrey summary name image = do
      (bytes, len) <- V.unsafeToForeignPtr0 <$> greyTensorToVector image
      (h,w,c,b) <- pure $ case V.toList $ size image of
         [1,h,w] -> (h,w,1,bytes)
         -- TODO Move this to compile time
         _       -> error $ "Don't know how to display a tensor with these dimensions as an Greyscale image " ++ show (size image)
      writeEvent summary
        (EventSummary name
         (SummaryImage
          { siHeight     = fromIntegral h
          , siWidth      = fromIntegral w
          , siColorspace = c
          , siBytes      = BS.PS b 0 len}))
  
  -- | Add a tensor that contains a grid of RGB images, like a batch of RGB images
  -- with the leading dimension being the number of images. Grids are padded with
  -- the given Size, constructed with the @Size@ datatype. They are padded with a
  -- given constant, usually 0.
  addImageRGBGrid ::
    forall (imagesPerRow :: Nat) (padding :: Nat) (rows :: Nat) (szh :: Nat) (szw :: Nat) (nr :: Nat) (ty :: TensorType).
      _
    => SummaryWriter
    -> Text
    -> Size imagesPerRow
    -> Padding padding
    -> TensorTyToHs ty
    -> Tensor ty 'KCpu '[nr, 3, szh, szw]
    -> IO ()
  addImageRGBGrid summary name sz padding padValue images = do
    grid <- makeRGBGrid sz padding padValue images
    addImageRGB summary name grid
  
  -- | Add a matplotlib plot to Tensorboard. This uses the Haskell matplotlib
  -- bindings from the `matplotlib` package,
  -- `https://github.com/abarbu/matplotlib-haskell`.
  addPlot :: SummaryWriter -> Text -> M.Matplotlib -> IO ()
  addPlot summary name mplot = do
    (fname, handle) <- openTempFile "/tmp/" "plot.png"
    hClose handle
    M.file fname mplot
    (h,w,_) <- fileImageProperties (T.pack fname)
    withSomeSing (fromIntegral h)
      (\ szh@(SNat :: SNat szh) ->
          withSomeSing (fromIntegral w)
            (\ szw@(SNat :: SNat szw) -> do
                tensor <- typed @TByte <$> readRGBTensorFromFile @szh @szw (T.pack fname)
                addImageRGB summary name tensor))
    fileExists <- doesFileExist fname
    when fileExists (removeFile fname)
  
  -- | Add the trace of a computation to Tensorboard as a graph.
  -- TODO Support naming nodes, more type information, more device information
  addGraph :: SummaryWriter -> Text -> ForeignPtr C.CTracingState -> IO ()
  addGraph summary name trace = do
    (inputs, nodes, outputs, ret, block) <-
      withForeignPtr trace C.tracing_state_graph
    inputs' <- mapM (\i -> do
                      b <- cbool <$> C.check_value_tensor i
                      na <- T.pack <$> C.value_name i
                      if b then do
                        sz <- C.value_sizes i
                        ty <- C.value_scalar_type i
                        pure (na, Just sz, Just ty)
                        else
                        pure (na, Nothing, Nothing))
              $ V.toList inputs
    nodes' <- mapM (\(n,nr) -> do
                     ins <- C.node_inputs n
                     outs <- C.node_outputs n
                     kind <- C.node_kind n
                     b <- C.node_has_attribute n "value"
                     attr <- if cbool b then do
                              ak <- C.node_attribute_kind n "value"
                              case ak of
                                C.AttributeKindFloat -> do
                                  f <- C.node_get_attribute_float n "value"
                                  pure $ Just $ defMessage & #f .~ coerce f
                                C.AttributeKindInt -> do
                                  i <- C.node_get_attribute_int n "value"
                                  pure $ Just $ defMessage & #i .~ fromIntegral i
                                C.AttributeKindString -> do
                                  s <- C.node_get_attribute_string n "value"
                                  pure $ Just $ defMessage & #s .~ T.encodeUtf8 (T.pack s)
                                C.AttributeKindTensor -> do
                                  tensor <- C.node_get_attribute_tensor n "value"
                                  ty <- C.getType tensor
                                  tensor' <- C.contiguous_mm tensor (fromIntegral $ fromEnum C.MemoryFormatContiguous)
                                  ptr <- newForeignPtr_ =<< castPtr <$> C.data_ptr tensor'
                                  len <- fromIntegral <$> C.numel tensor'
                                  let get :: V.Storable a => [a]
                                      get = V.toList $ V.unsafeFromForeignPtr0 (castForeignPtr ptr) len
                                  pure $ Just $ defMessage & #tensor .~
                                    case ty of
                                      C.ScalarTypeByte ->
                                        defMessage & #dtype .~ PS.DT_UINT8
                                                   & #intVal .~ map (fromIntegral :: Word8 -> Int32) get
                                      C.ScalarTypeChar ->
                                        defMessage & #dtype .~ PS.DT_INT8
                                                   & #intVal .~ map (fromIntegral :: Int8 -> Int32) get
                                      C.ScalarTypeShort ->
                                        defMessage & #dtype .~ PS.DT_INT16
                                                   & #intVal .~ map (fromIntegral :: Int16 -> Int32) get
                                      C.ScalarTypeInt ->
                                        defMessage & #dtype .~ PS.DT_INT32
                                                   & #intVal .~ get
                                      C.ScalarTypeLong ->
                                        defMessage & #dtype .~ PS.DT_INT64
                                                   & #int64Val .~ get
                                      C.ScalarTypeHalf ->
                                        defMessage & #dtype .~ PS.DT_HALF
                                                   & #halfVal .~ get
                                      C.ScalarTypeFloat ->
                                        defMessage & #dtype .~ PS.DT_FLOAT
                                                   & #floatVal .~ get
                                      C.ScalarTypeDouble ->
                                        defMessage & #dtype .~ PS.DT_DOUBLE
                                                   & #doubleVal .~ get
                                      C.ScalarTypeUndefined ->
                                        defMessage & #dtype .~ PS.DT_INVALID
                                _ -> pure Nothing
                       else pure Nothing
                     ins' <- mapM (\v -> do
                                  n <- C.value_name v
                                  t <- C.value_type v
                                  ct <- convertCType t
                                  pure $ TempGraphValue v n t ct)
                            $ V.toList ins
                     outs' <- mapM (\v -> do
                                  n <- C.value_name v
                                  t <- C.value_type v
                                  ct <- convertCType t
                                  pure $ TempGraphValue v n t ct)
                             $ V.toList outs
                     pure (n, TempGraphNode nr kind ins' outs' attr)) $ zip (V.toList nodes) [0..]
    let nodesMap = M.fromList nodes'
    let insToNodes  = M.fromListWith (++)
                    $ concatMap (\(n, TempGraphNode nr kinds ins outs attr) ->
                               map (\(TempGraphValue v _ _ _) -> (v,[n])) ins) nodes'
    let outsToNodes = M.fromListWith (++)
                    $ concatMap (\(n, TempGraphNode nr kinds ins outs attr) ->
                               map (\(TempGraphValue v _ _ _) -> (v,[n])) outs) nodes'
    let nodeName nr kind = T.replace "prim::" "" $ T.replace "aten::" "" $ T.pack kind <> "_"<> T.pack (show nr)
    writeEvent summary
      (EventGraph name
        (map (\(name, size, ty) ->
                 case (size, ty) of
                   (Just sz, Just t) ->
                     defMessage & #name .~ "input_" <> name
                                & #op .~ "Input"
                                & #attr .~ M.fromList [("_output_shapes",
                                                        defMessage & #list .~
                                                         (defMessage & #shape .~
                                                           [defMessage & #dim .~
                                                            map (\x -> defMessage & #size .~ x) (V.toList sz)]))]
                   (_, _) ->
                     defMessage & #name .~ "input_" <> name
                                & #op .~ "Input")
           inputs'
         ++
         map (\(node,(TempGraphNode nr kind ins outs attr)) ->
                 defMessage & #name .~ nodeName nr kind
                            & #op .~ T.pack kind
                            & #device .~ T.pack ""
                            & #input .~ map (\(TempGraphValue v n t ct) ->
                                                  case M.lookup v outsToNodes of
                                                    Nothing -> T.pack ("input_" ++ n)
                                                    Just [outNode] -> case M.lookup outNode nodesMap of
                                                                     Just (TempGraphNode nr kind _ _ _) -> nodeName nr kind
                                                                     Nothing                            -> error "Can't find node") ins
                            & #attr .~ M.fromList([("_output_shapes",
                                                     defMessage
                                                     & #list .~ (defMessage
                                                                 & #shape .~
                                                                    (map (\(TempGraphValue v n t ct) ->
                                                                             case ct of
                                                                               TraceTypeSimple _ _ ->
                                                                                 defMessage & #dim .~ [defMessage & #size .~ 1]
                                                                               TraceTypeTensor _ (Just sz) _ _ ->
                                                                                 defMessage & #dim .~ (map (\x -> defMessage & #size .~ x)
                                                                                                           (V.toList sz))
                                                                               _ -> defMessage)
                                                                      outs)))]
                                                  ++ (case attr of
                                                       Nothing -> []
                                                       Just a  -> [("value",a)])))
             nodes'))
  
  -- | A convenience function, add the summary statistics and histograms for a
  -- tensor.
  addTensorSummary :: (Num (TensorTyToHs ty), Num (TensorTyToHsC ty),
                       V.Storable (TensorTyToHs ty), V.Storable (TensorTyToHsC ty),
                       SingI sz, SingI ty) =>
                     SummaryWriter -> Text -> Tensor ty 'KCpu sz -> IO ()
  addTensorSummary summary name tensor = do
    mean   <- fromScalar =<< Tensor.mean tensor
    median <- fromScalar =<< Tensor.median tensor
    var    <- fromScalar =<< Tensor.var tensor True
    min    <- fromScalar =<< Tensor.min tensor
    max    <- fromScalar =<< Tensor.max tensor
    addScalar    summary ("summaries/" <> name <> "/mean")   $ toDouble mean
    addScalar    summary ("summaries/" <> name <> "/median") $ toDouble median
    addScalar    summary ("summaries/" <> name <> "/stddev") $ Prelude.sqrt $ toDouble var
    addScalar    summary ("summaries/" <> name <> "/min")    $ toDouble min
    addScalar    summary ("summaries/" <> name <> "/max")    $ toDouble max
    addHistogram summary ("summaries/" <> name <> "/histogram") tensor
  
  -- * Basic heuristics
  
  -- | The numpy auto histogram mode
  histogramBinsAuto :: Vector Double -> Int
  histogramBinsAuto x | V.length x < 1000 = histogramBinsSturges x
                      | otherwise = histogramBinsFD x
  
  -- | The Freedman-Diaconis histogram bin estimator.
  -- bins = 2 \frac{IQR}{n^{1/3}}
  histogramBinsFD :: Vector Double -> Int
  histogramBinsFD v =
    ceiling $ (V.maximum v - V.minimum v) /
                 (2 * (S.weightedAvg 3 4 v - S.weightedAvg 1 4 v)
                    / (fromIntegral (V.length v) ** (1/3)))
  
  -- | Sturges's histogram bin estimator.
  histogramBinsSturges :: Vector Double -> Int
  histogramBinsSturges v = Prelude.round $ logBase 2 (fromIntegral $ V.length v) + 1
  
  -- * Internal
  
  -- | This is internal, you never need to use it manually. It's what we fill out
  -- with the various functions in this module and then serialize for Tensorboard.
  data Summary = SummaryValue Double
               | SummaryHistogram { shMin         :: Double
                                  , shMax         :: Double
                                  , shNum         :: Double
                                  , shSum         :: Double
                                  , shSumSq       :: Double
                                  , shBucketLimit :: Vector Double
                                  , shBucket      :: Vector Double }
               | SummaryImage { siHeight     :: Int32
                              , siWidth      :: Int32
                              , siColorspace :: Int32
                              , siBytes      :: BS.ByteString }
  
  -- | This is internal.
  data Node = Node { nName   :: Text
                   , nOp     :: Text
                   , nInput  :: [Text]
                   , nDevice :: Text
                   }
  
  -- | Internal. Used to capture the various kinds of entities that Tensorboard
  -- can understand.
  data Event = EventVersion
             | EventGraph Text [PG.NodeDef]
             | EventSummary Text Summary
             | EventMessage LogLevel Text
             | EventSessionLog
             | EventMetadata
             | EventMetaGraph
  
  -- | Internal
  data TempGraphValue = TempGraphValue (Ptr C.CJitValue) String (Ptr C.CType) TraceType
    deriving (Show)
  -- | Internal
  data TempGraphNode  = TempGraphNode Int String [TempGraphValue] [TempGraphValue] (Maybe PA.AttrValue)
    deriving (Show)
  
  -- | Internal
  eventToProto :: Int -> Event -> IO PE.Event
  eventToProto step event = do
    (time :: Double) <- realToFrac <$> getPOSIXTime
    let d = case event of
              EventVersion -> defMessage & #fileVersion .~ "brain.Event:2"
              EventMessage lvl txt ->
                defMessage & #logMessage .~ (defMessage & #level .~ (case lvl of
                                                          LevelDebug         -> PE.LogMessage'DEBUG
                                                          LevelInfo          -> PE.LogMessage'INFO
                                                          LevelWarn          -> PE.LogMessage'WARN
                                                          LevelError         -> PE.LogMessage'ERROR
                                                          LevelOther "fatal" -> PE.LogMessage'FATAL
                                                          LevelOther _       -> PE.LogMessage'UNKNOWN)
                                          & #message .~ txt)
              EventSummary tag (SummaryValue val) ->
                defMessage & #summary .~
                (defMessage & #value .~
                  [defMessage & #tag .~ tag
                  -- TODO Metadata!
                       & #simpleValue .~ double2Float val])
              EventSummary tag sh@SummaryHistogram {} ->
                defMessage & #summary .~
                (defMessage & #value .~
                  [defMessage & #tag .~ tag
                  -- TODO Metadata!
                       & #histo .~
                       (defMessage
                            & #min .~ shMin sh
                            & #max .~ shMax sh
                            & #num .~ shNum sh
                            & #sum   .~ shSum sh
                            & #sumSquares .~ shSumSq sh
                            & #bucketLimit .~ V.toList (shBucketLimit sh)
                            & #bucket .~ V.toList (shBucket sh))])
              EventSummary tag sh@SummaryImage{} ->
                defMessage & #summary .~
                (defMessage & #value .~
                  [defMessage & #tag .~ tag
                  -- TODO Metadata!
                       & #image .~
                       (defMessage
                            & #height             .~ siHeight sh
                            & #width              .~ siWidth sh
                            & #colorspace         .~ siColorspace sh
                            & #encodedImageString .~ siBytes sh)])
              EventGraph tag nodes ->
                defMessage & #graphDef .~
                  encodeMessage ((defMessage :: PG.GraphDef)
                                  & #versions .~ (defMessage & #producer .~ 1
                                                             & #minConsumer .~ 1
                                                             & #badConsumers .~ [])
                                  & #version .~ 3
                                  & #node .~ nodes)
    pure $ d & #wallTime .~ time
             & #step .~ fromIntegral step
  
  -- | Internal
  encodeEvent :: PE.Event -> B.ByteString
  encodeEvent ev = BL.toStrict $ runPut $ do
    putByteString len
    putWord32le (maskedCRC32 len)
    putByteString msg
    putWord32le (maskedCRC32 msg)
    where msg = encodeMessage ev
          len = BL.toStrict $ runPut $ putWord64le (fromIntegral $ B.length msg)
  
  -- | Internal
  writeEvent :: SummaryWriter -> Event -> IO ()
  writeEvent r e = do
    step <- readIORef $ lastStep r
    let fname = T.unpack (runDirectory r) </> (T.unpack $ runFilename r)
    b <- doesFileExist fname
    unless b $ do
      bs <- encodeEvent <$> eventToProto step EventVersion
      B.writeFile fname bs
    bs <- encodeEvent <$> eventToProto step e
    B.appendFile fname bs
  
  -- | Internal
  maskedCRC32 :: BS.ByteString -> Word32
  maskedCRC32 bs = (crc32c bs `rotateR` 15) + 0xa282ead8
  