content: {-# LANGUAGE TemplateHaskell #-}
  module NN.CLI where
  
  import           Control.Applicative
  import           Control.Lens          hiding ((<.>))
  import           Control.Monad
  import qualified Data.ByteString.Lazy  as BS
  import           Data.GraphViz         (GraphvizOutput (..))
  import           NN.Backend.Caffe      as Caffe
  import           NN.Backend.Torch      as Torch
  import           NN.DSL
  import           NN.Passes
  import           NN.Visualize
  import           Options.Applicative   hiding ((&))
  import           System.Exit
  import           System.FilePath.Posix
  import           System.Process
  import           Text.ProtocolBuffers  as P
  
  caffePrototxt :: NetBuilder a -> FilePath -> String -> IO ()
  caffePrototxt net prototxtPath binaryToText' = do
    parse net & Caffe.middleEnd & Caffe.backend & messagePut & BS.writeFile binaryPath
    rawSystem binaryToText' [binaryPath, prototxtPath] >>= exitWith
      where
        binaryPath = prototxtPath <.> "protobinary"
  
  graph :: NetBuilder a -> GraphvizOutput -> FilePath -> IO ()
  graph net format path = void $ visualize (parse net) & render format path
  
  torchCode :: NetBuilder a -> FilePath -> IO ()
  torchCode net path = do
    let Just code = parse net & Torch.backend
    writeFile path code
  
  data Command = Caffe { _fileName, _binaryToText :: String}
               | Torch { _fileName :: String }
               | Visualize { _fileName :: String, _format :: String }
  makeLenses ''Command
  makePrisms ''Command
  
  opts :: Parser Command
  opts = subparser (caffe <> torch <> graph')
         where
           nc name parser desc = command name (info (helper <*> parser) (progDesc desc))
           caffe = nc "caffe" (Caffe <$> fileName' <*> binaryToText') "Generate a Caffe .prototxt to run with `caffe train --model=<>"
           torch = nc "torch" (Torch <$> fileName') "Generate Lua code to be `require`'d into an existing Torch script"
           graph' = nc "visualize" (Visualize <$> fileName' <*> vf) "Generate an image visualizing the model's connectivity"
           fileName' = strOption (long "output" <> help "Write output to FILE" <> metavar "FILE")
           binaryToText' = strOption (long "binary_to_text"
                                     <> help "Path to binary_to_text.py BINARY"
                                     <> showDefault
                                     <> metavar "BINARY"
                                     <> value "./binary_to_text.py")
           vf = strOption (long "format" <> help "pdf or png")
  
  
  run :: NetBuilder a -> Command -> IO ()
  run net (Caffe prototxtPath binaryToTextPath) = caffePrototxt net prototxtPath binaryToTextPath
  run net (Torch path) = torchCode net path
  run net (Visualize path format) = graph net (fmtTy format) path
      where
        fmtTy "png" = Png
        fmtTy "pdf" = Pdf
        fmtTy _ = error "Must be pdf or png"
  
  cli :: NetBuilder a -> IO ()
  cli net = execParser (info (helper <*> opts) idm) >>= run net
  