module HsiMethod 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 Options.Applicative

newtype HsiMethodOptions = HsiMethodOptions
  { filename :: FilePath
  }
  deriving Show

hsiMethodOptionsParser :: Parser HsiMethodOptions
hsiMethodOptionsParser =
  HsiMethodOptions
    <$> argument str (help "Filename to read (dot format)" <> metavar "FILE")

mainHsiMethod :: HsiMethodOptions -> IO ()
mainHsiMethod HsiMethodOptions{..} = do
  let dotFile = filename
  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)