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

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

-- TODO: use common options
verbose :: Bool
verbose = False

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
  when verbose (print 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)

  when verbose $ do
    putStrLn "\nPARTITION"
    print partition

    putStrLn "\nTREE"
    print splittingTree

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

  when verbose $ do
    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)

    putStrLn "\nSTATE COVER"
    print prefixes

    putStrLn "\nTEST SUITE"

  mapM_ print testSuite