module LStarMain where

import Bisimulation (bisimulation2)
import DotParser (readDotFile)
import LStar
import Mealy

import Control.Monad (when)
import Control.Monad.Trans.Class
import Control.Monad.Trans.State.Strict
import Data.Map.Strict qualified as Map
import Options.Applicative

data LStarOptions = LStarOptions
  { filename :: FilePath
  , debugOutput :: Bool
  }
  deriving Show

lStarOptionsParser :: Parser LStarOptions
lStarOptionsParser =
  LStarOptions
    <$> argument str (help "Filename to read (dot format)" <> metavar "FILE")
    <*> switch (long "verbose" <> short 'v' <> help "Enable extra debugging output")

mainLStar :: LStarOptions -> IO ()
mainLStar LStarOptions{..} = do
  let
    dotFile = filename

  print dotFile
  machine <- readDotFile dotFile

  let
    alphabet = inputs machine
    tInit = initialState machine
    tOut s i = fst (behaviour machine s i)
    tTrans s i = snd (behaviour machine s i)

    mq0 = semanticsForState machine (initialState machine)
    mq = countingMQ (\w -> when debugOutput (print w) >> return (mq0 w))

    loop table = do
      lift $ putStrLn "Making the table closed and consistent"
      (table2, b) <- makeClosedAndConsistentA mq table
      let (hInit, size, hTransMap, hOutMap) = createHypothesis table2
          hTrans s i = hTransMap Map.! (s, i)
          hOut s i = hOutMap Map.! (s, i)
          res = bisimulation2 alphabet tOut tTrans tInit hOut hTrans hInit

      lift $ putStrLn (if b then "Table changed" else "Table did not changed")
      lift $ putStrLn (show size <> " states")

      case res of
        Nothing -> do
          lift $ putStrLn "Done learning!"
          return size
        Just cex -> do
          lift $ putStrLn "CEX:" >> print cex
          table3 <- processCounterexampleA cex mq table2
          loop table3

    learner = do
      table <- initialiseA alphabet mq
      loop table

  (a, b) <- runStateT learner 0

  putStrLn $ "Size:  " <> show a
  putStrLn $ "MQs:   " <> show b

semanticsForState :: MealyMachine s i o -> s -> [i] -> o
semanticsForState _ _ [] = error ""
semanticsForState MealyMachine{..} q [a] = fst (behaviour q a)
semanticsForState m@MealyMachine{..} q (a : w) = semanticsForState m (snd (behaviour q a)) w