{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ExistentialQuantification #-}

module Main (main) where

import           Data.Vector                    ( Vector )
import qualified Data.Vector                   as V
import qualified Data.Text.Lazy.IO             as T
import           Data.Text                      ( Text )
import qualified Data.Text.Lazy.Builder        as Build
import qualified Options.Applicative           as Options
import qualified Text.Megaparsec               as Mega
import qualified Text.Megaparsec.Char          as Mega
import qualified Text.Megaparsec.Char.Lexer    as Mega
import           Data.Void
import           System.Random
import           System.IO
import           Numeric

import           Types
import           Generator
import           Output
import           Probability

data SomeMonoid = forall m. SomeMonoid (MonoidType m)

data Opts = Opts
  { optMonoid :: SomeMonoid
  , optStates :: Int
  , optSymbols :: SymbolSpec
  , optZeroFrequency :: Probability
  , optRandomState :: Maybe StdGen
  , optDifferentValues :: Maybe Int
  }

readMonoid :: Options.ReadM SomeMonoid
readMonoid = Options.eitherReader $ \case
  "Z,max"    -> Right (SomeMonoid MaxInt)
  "Word,or"  -> Right (SomeMonoid OrWord)
  "powerset" -> Right (SomeMonoid Powerset)
  _ ->
    Left
      $  "Unknown monoid type. Valid values are '"
      <> "Z,max"
      <> "', '"
      <> "Word,or"
      <> "' and '"
      <> "powerset"
      <> "'"

readSymbols :: Options.ReadM SymbolSpec
readSymbols = Options.maybeReader (Mega.parseMaybe parser)
  where
    parser :: Mega.Parsec Void String (Vector Int)
    parser = V.fromList <$> Mega.decimal `Mega.sepBy` (Mega.string ",")

readCount :: String -> Either String Int
readCount input = case readDec input of
  [(0, "")] -> Left "Count must be >0"
  [(x, "")] -> Right x
  _ -> Left "Count not parse Number"

parseOpts :: Options.Parser Opts
parseOpts =
  Opts
    <$> Options.option
          readMonoid
          (Options.long "monoid" <> Options.metavar "MONOID" <> Options.help
            (  "Monoid to use. Valid valies are '"
            <> "Z,max"
            <> "', '"
            <> "Word,or"
            <> "' and '"
            <> "powerset"
            <> "'"
            )
          )
    <*> Options.option
          Options.auto
          (Options.long "states" <> Options.metavar "NUM" <> Options.help
            "Number of states to generate"
          )
    <*> Options.option
          readSymbols
          (  Options.long "symbols"
          <> Options.metavar "NUM,NUM,..."
          <> Options.help
               "Comma separated list of symbols per arity. E.g. 2,0,1 means two symbols with arity 0, non with arity 1 and one with arity two"
          )
    <*> Options.option
          (Options.eitherReader readProbability)
          (  Options.long "zero-frequency"
          <> Options.showDefault
          <> Options.value (Probability 7 1)
          <> Options.metavar "FREQ"
          <> Options.help
               "Frequency of edges with zero weight as number between 0 and 1."
          )
    <*> Options.optional
          (Options.option Options.auto (Options.long "random-state"))
    <*> Options.optional
          (Options.option (Options.eitherReader readCount) (Options.long "different-values" <> Options.metavar "NUM" <> Options.help "Maximal number of differnt monoid values to generate"))

withSpec :: Opts -> (forall m . WTASpec m -> x) -> x
withSpec opts f = case optMonoid opts of
  SomeMonoid m -> f WTASpec { monoid     = m
                            , numStates  = optStates opts
                            , numSymbols = optSymbols opts
                            }

main :: IO ()
main = do
  let optSpec =
        Options.info (parseOpts Options.<**> Options.helper) (Options.fullDesc)
  opts <- Options.execParser optSpec

  case optRandomState opts of
    Nothing -> return ()
    Just x  -> setStdGen x

  withSpec opts $ \spec -> do
    randGen <- getStdGen
    wta <- runGenerator (GeneratorConfig spec (optZeroFrequency opts) (optDifferentValues opts)) genWTA
    putStrLn $ "# Random state for this automaton: '" <> show randGen <> "'"
    T.putStr (Build.toLazyText (buildWTA wta))