module Main where

import Bisimulation (bisimulation2)
import DotParser
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 Data.Maybe (mapMaybe)
import System.Environment
import Text.Megaparsec

debugOutput :: Bool
debugOutput = False

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

main :: IO ()
main = do
  [dotFile] <- getArgs
  print dotFile
  transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile

  let
    machine = convertToMealy transitions
    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