module Option where

import qualified Feature
import qualified Parameters as Params
import qualified HiddenMarkovModel as HMM

import qualified Distribution.Verbosity as Verbosity
import qualified Distribution.ReadE as ReadE
import Distribution.Verbosity (Verbosity)

import qualified Options.Applicative as OP

import qualified SignalProcessingMethods as SPMethods
import qualified Data.StorableVector.Lazy as SVL
import qualified Sound.SoxLib as SoxLib

import qualified System.Path.PartClass as PathClass
import qualified System.Path as Path
import qualified System.IO as IO
import Text.Printf (printf, )

import Control.Monad (when, )
import Control.Applicative (pure, (<*>), (<$>), (<|>), )
import Control.Applicative.HT (liftA5, )

import qualified Data.Map as Map
import qualified Data.List as List
import Data.Monoid ((<>), )



data Flags =
   Flags {
      verbosity :: Verbosity,
      chunkSize :: SVL.ChunkSize,
      sampleRate :: Maybe SoxLib.Rate
   }
   deriving (Show)


optionVerbosity :: OP.ReadM Verbosity
optionVerbosity =
   OP.eitherReader $ ReadE.runReadE Verbosity.flagToVerbosity

-- cf. Distribution.Simple.Utils
noticeAction :: Flags -> IO () -> IO ()
noticeAction flags emit =
   when (verbosity flags >= Verbosity.normal) emit

notice :: Flags -> String -> IO ()
notice flags msg = noticeAction flags (putStrLn msg)

infoAction :: Flags -> IO () -> IO ()
infoAction flags emit =
   when (verbosity flags >= Verbosity.verbose) emit

infoMsg :: Flags -> String -> IO ()
infoMsg flags msg = infoAction flags (putStrLn msg)

warn :: Flags -> String -> IO ()
warn flags msg =
   when (verbosity flags >= Verbosity.normal) $ do
      IO.hFlush IO.stdout
      IO.hPutStr IO.stderr msg


parseFlags :: OP.Parser Flags
parseFlags =
   pure Flags
    <*> OP.option optionVerbosity
           ( OP.value Verbosity.normal
          <> OP.short 'v'
          <> OP.long "verbose"
          <> OP.metavar "0..3"
          <> OP.help "verbosity" )
    <*> OP.option (fmap SVL.chunkSize OP.auto)
           ( OP.value (SVL.chunkSize 65536) -- crashes SoX' FLAC writer
          <> OP.long "chunksize"
          <> OP.metavar "NUMSAMPLES"
          <> OP.help "size of processing units" )
    <*> OP.option (fmap Just OP.auto)
           ( OP.value Nothing
          <> OP.long "samplerate"
          <> OP.metavar "HERTZ"
          <> OP.help "override sample rate of input files" )

model :: OP.Parser (IO Feature.HMM)
model =
   OP.option (Feature.readHMM <$> path) $
      OP.long "model" <>
      OP.metavar "PATH" <>
      OP.value (return Feature.hmmHardwired) <>
      OP.help "CSV file containing HMM parameters"

path :: (PathClass.FileDir fd) => OP.ReadM (Path.AbsRel fd)
path = OP.eitherReader Path.parse

numStates :: OP.Parser Int
numStates =
   OP.option
         (OP.eitherReader $ \str ->
            case reads str of
               [(n, "")] ->
                  if n>0
                    then Right n
                    else Left "not positive"
               _ -> Left "not an integer") $
      OP.long "numstates" <>
      OP.metavar "NUMBER" <>
      OP.help "number of states for Hidden Markov Model" <>
      OP.value 6


data TrainingFlags =
   TrainingFlags {
      trainingSignalProcessing :: SPMethods.T,
      trainingFeature :: Feature.Class,
      trainingConvergence :: HMM.Convergence,
      trainingMLPack :: Bool,
      trainingPlot :: Bool
   }

formatBool :: Bool -> String
formatBool b = if b then "yes" else "no"

switch :: Bool -> String -> String -> OP.Parser Bool
switch deflt name helpMsg =
   OP.flag' False (OP.long ("no-" ++ name))
   <|>
   OP.flag deflt True
      (OP.long name <>
       OP.help (printf "%s (default: %s)" helpMsg (formatBool deflt)))

feature :: Params.T -> OP.Parser (IO Feature.Class)
feature params =
   let lookupFeature name =
         case Map.lookup name Feature.dictionaryMerged of
            Just feat -> return feat
            Nothing ->
               ioError $ userError $ unlines $
                  printf "unknown feature \"%s\"" name :
                  "known features:" :
                  Map.keys Feature.dictionaryMerged
   in  OP.option (fmap lookupFeature OP.str) $
         OP.long "feature" <>
         OP.metavar "NAME" <>
         OP.value
            (return $ Feature.lowRateSqrt $ Params.featureSampleRate params) <>
         OP.help
            ("one of " ++
             List.intercalate ", " (Map.keys Feature.dictionaryMerged))

plot :: OP.Parser Bool
plot =
   OP.switch $
      OP.long "plot" <>
      OP.help "show feature vectors grouped by model states"

mlpack :: OP.Parser Bool
mlpack =
   OP.switch $
      OP.long "mlpack" <>
      OP.help "emit files that let you perform training using mlpack"

trainingFlags ::
   OP.Parser SPMethods.T -> Params.T -> OP.Parser (IO TrainingFlags)
trainingFlags signalProcessing params =
   liftA5
      (\sigProc lookupFeature cvg mlp plt -> do
         feat <- lookupFeature
         return $ TrainingFlags sigProc feat cvg mlp plt)
      signalProcessing (feature params) HMM.convergenceOptions mlpack plot


type Commands = OP.Mod OP.CommandFields

simpleAction :: String -> String -> OP.Parser a -> Commands a
simpleAction name helpText act =
   OP.command name $ OP.info (OP.helper <*> act) (OP.progDesc helpText)

withInOutPaths ::
   (PathClass.FileDir fd0, PathClass.FileDir fd1) =>
   (Flags -> Path.AbsRel fd0 -> Path.AbsRel fd1 -> IO ()) ->
   (FilePath -> FilePath -> Flags -> IO ())
withInOutPaths act =
   \inputStr outputStr flags -> do
      let ioExc = either (ioError . userError) return
      input <- ioExc $ Path.parse inputStr
      output <- ioExc $ Path.parse outputStr
      act flags input output

transferActionApp ::
   (PathClass.FileDir fd0, PathClass.FileDir fd1) =>
   String -> String ->
   OP.Parser
      (Flags -> Path.AbsRel fd0 -> Path.AbsRel fd1 -> IO ()) ->
   Commands (Flags -> IO ())
transferActionApp name msg parser =
   simpleAction name msg $
   OP.liftA3 withInOutPaths parser
      (OP.strArgument (OP.metavar "SRC"))
      (OP.strArgument (OP.metavar "DST"))

transferAction ::
   (PathClass.FileDir fd0, PathClass.FileDir fd1) =>
   String -> String ->
   (Flags -> Path.AbsRel fd0 -> Path.AbsRel fd1 -> IO ()) ->
   Commands (Flags -> IO ())
transferAction name msg = transferActionApp name msg . pure

multiAction ::
   (PathClass.FileDir fd0, PathClass.FileDir fd1) =>
   String -> String ->
   OP.Parser (Bool -> [Path.AbsRel fd0] -> Path.AbsRel fd1 -> a) ->
   Commands a
multiAction name msg parse =
   simpleAction name msg $
   parse
      <*> plot
      <*> OP.some (OP.argument path (OP.metavar "SRC"))
      <*> OP.option path (OP.long "output" <> OP.metavar "DST")


info :: OP.Parser a -> OP.ParserInfo a
info parser =
   OP.info
      (OP.helper <*> parser)
      (OP.fullDesc <> OP.progDesc "Classify sounds of xenopus laevis")