{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
-- | 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 StateCover.Simultaneous (simultaneousStateCover)
import StateIdentifiers (stateIdentifierFor)

import Control.Monad (when, unless)
import Control.Monad.Trans.State (evalStateT)
import Data.Map.Strict qualified as Map
import Options.Applicative
import System.IO
import System.Exit (exitSuccess)
import Data.Functor.Identity (runIdentity)

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

data HsiMethodOptions = HsiMethodOptions
  { filenames :: [FilePath]
  , mode :: Mode
  }
  deriving Show

hsiMethodOptionsParser :: Parser HsiMethodOptions
hsiMethodOptionsParser =
  HsiMethodOptions
    <$> many (argument str (help "Filename(s) 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" filenames

  when (length filenames > 1) $ do
    compositionalTesting filenames
    exitSuccess

  machine <- readDotFile (head filenames)

  let (partition, splittingTree) = mealyToHsi machine

  logging "PARTITION" partition
  logging "TREE" splittingTree

  let
    MealyMachine{..} = machine
    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

  logging "STATS" (length testSuite, sum (fmap length testSuite))

mealyToHsi :: _ -> _
mealyToHsi MealyMachine{..} =
  let
    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]
    ignore = const (pure ()) -- no debug info
  in
    runIdentity $ evalStateT (refine ignore outputFuns reverseFuns) (initialPRState states)



compositionalTesting :: [FilePath] -> IO ()
compositionalTesting filenames = do
  machines <- mapM readDotFile filenames

  unless (all (\m -> inputs m == inputs (head machines)) machines) $ do
    putStrLn "Warning: different intput sets, currently poorly implemented"
    return ()

  let
    inps = inputs (head machines)
    separatePrefixes m = stateCover (\s -> [((i, o), t) | i <- inputs m, let (o, t) = behaviour m s i]) (initialState m)
    hsi m = let (partition, splittingTree) = mealyToHsi m in \s -> stateIdentifierFor s partition splittingTree
    separateTestSuite m =
      let ps = separatePrefixes m
          ws = hsi m
       in Trie.reducePrefixes [px <> sx | s <- states m, let px = ps Map.! s, sx <- Trie.toList (ws s)]
    testSuiteUnion = Trie.reducePrefixes (concatMap separateTestSuite machines)

    simultaneousPrefixes = simultaneousStateCover [\s i -> Just (snd (behaviour m s i)) | m <- machines] inps (fmap initialState machines)
    simultaneousTestSuite m ps =
      let ws = hsi m
       in Trie.reducePrefixes [px <> sx | s <- states m, let px = ps Map.! s, sx <- Trie.toList (ws s)]
    testSuiteUnion2 = Trie.reducePrefixes . concat $ zipWith simultaneousTestSuite machines simultaneousPrefixes


  -- middle part of test not yet implemented

  putStrLn "\nStats"
  putStrLn $ "Sep words: " <> show (length testSuiteUnion)
  putStrLn $ "Sep symbols: " <> show (sum (fmap length testSuiteUnion))
  putStrLn $ "Com words: " <> show (length testSuiteUnion2)
  putStrLn $ "Com symbols: " <> show (sum (fmap length testSuiteUnion2))