module Main where

import Data.Trie qualified as Trie
import DotParser (readDotFile)
import Mealy (MealyMachine (..))
import SplittingTree (initialPRState, refine)
import StateIdentifiers (stateIdentifierFor)

import Control.Monad.Trans.State (evalStateT)
import Data.Map.Strict qualified as Map
import System.Environment (getArgs)

main :: IO ()
main = do
  args <- getArgs
  case args of
    ("HSI" : ls) -> mainHSI ls
    _ -> putStrLn "Please provide one of [HSI, InputDecomp]"

mainHSI :: [String] -> IO ()
mainHSI args = case args of
  [dotFile] -> run dotFile
  _ -> putStrLn "Please provide a dot file"
 where
  run dotFile = do
    print dotFile
    machine <- readDotFile dotFile

    -- convert to mealy
    let
      MealyMachine{..} = machine
      outputFuns = [(i, fun) | i <- inputs, let fun s = fst (behaviour s i)]
      reverseTransitionMaps i = Map.fromListWith (++) [(t, [s]) | s <- states, let t = snd (behaviour s i)]
      reverseFuns = [(i, fun) | i <- inputs, let m = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s m]

    (partition, splittingTree) <- evalStateT (refine print outputFuns reverseFuns) (initialPRState states)

    putStrLn "\nPARTITION"
    print partition

    putStrLn "\nTREE"
    print splittingTree

    let
      siFor s = stateIdentifierFor s partition splittingTree

    putStrLn "\nHARMONISED STATE IDENTIFIERS"
    sis <- mapM (\s -> let si = siFor s in print (Trie.toList si) >> return si) states

    putStrLn "\nW-SET"
    print (Trie.toList . foldr Trie.union Trie.empty $ sis)