module NN.CLI(cli) where

import           Control.Applicative
import           Control.Lens          hiding ((<.>))
import           Control.Monad
import qualified Data.ByteString.Lazy  as BS
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 -> FilePath -> FilePath -> 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"

netPdf :: NetBuilder -> FilePath -> IO ()
netPdf net path = void $ visualize (parse net) & pdf path

torchCode :: NetBuilder -> FilePath -> IO ()
torchCode net path = do
  let Just code = parse net & Torch.backend
  writeFile path code

data Command = Caffe String String | Torch String | PDF String

opts :: Parser Command
opts = subparser (caffe <> torch <> pdf')
       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"
         pdf' = nc "pdf" (PDF <$> filename) "Generate a PDF 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")

run :: NetBuilder -> Command -> IO ()
run net (Caffe prototxtPath binaryToTextPath) = caffePrototxt net prototxtPath binaryToTextPath
run net (Torch path) = torchCode net path
run net (PDF path) = netPdf net path

cli :: NetBuilder -> IO ()
cli net = execParser (info (helper <*> opts) idm) >>= run net