-- | Copyright: (c) 2024-2025 Joshua Moerman, Open Universiteit
-- SPDX-License-Identifier: EUPL-1.2
module HsiMethod where

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

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

data Mode = HSI | W
  deriving (Eq, Ord, Show, Read)

data HsiMethodOptions = HsiMethodOptions
  { filename :: FilePath
  , mode :: Mode
  }
  deriving Show

hsiMethodOptionsParser :: Parser HsiMethodOptions
hsiMethodOptionsParser =
  HsiMethodOptions
    <$> argument str (help "Filename to read (dot format)" <> metavar "FILE")
    <*> option auto (long "mode" <> help "Mode (HSI, W)" <> metavar "MODE" <> showDefault <> value HSI)

mainHsiMethod :: HsiMethodOptions -> CommonOptions -> IO ()
mainHsiMethod HsiMethodOptions{..} CommonOptions{..} = do
  let
    logging st x = when verbose (hPutStrLn stderr st >> hPrint stderr x >> hPutStrLn stderr "")

  logging "FILENAME" filename
  machine <- readDotFile filename

  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 (const (pure ())) outputFuns reverseFuns) (initialPRState states)

  logging "PARTITION" partition
  logging "TREE" splittingTree

  let
    siFor s = stateIdentifierFor s partition splittingTree
    wset = Trie.toList . foldr (Trie.union . siFor) Trie.empty $ states
    prefixes = stateCover (\s -> [((i, o), t) | i <- inputs, let (o, t) = behaviour s i]) initialState
    -- TODO: add middle transition(s)
    testSuite = case mode of
      HSI -> Trie.reducePrefixes [px <> sx | s <- states, let px = prefixes Map.! s, sx <- Trie.toList (siFor s)]
      W -> Trie.reducePrefixes [prefixes Map.! s <> sx | s <- states, sx <- wset]

  logging "HARMONISED STATE IDENTIFIERS" [(s, Trie.toList (siFor s)) | s <- states]
  logging "W-SET" wset
  logging "STATE COVER" prefixes

  mapM_ print testSuite