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

import Bisimulation (bisimulation2)
import CommonOptions
import Data.Partition (Block (..), numBlocks)
import Data.UnionFind (empty, equate, equivalent)
import DotParser (readDotFile)
import Mealy (MealyMachine (..), outputFunction, transitionFunction)
import MealyRefine (refineMealy)

import Control.Monad (mapAndUnzipM)
import Data.List qualified as List
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust)
import Data.Set qualified as Set
import Data.Text qualified as T
import Options.Applicative
import System.Exit (exitFailure, exitSuccess)
import System.FilePath ((</>))

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

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

-- Interleaving composition of restriction to subalphabets
-- precondition: alph1 and alph2 have no common elements
interleavingComposition :: Ord i => [[i]] -> MealyMachine s i o -> MealyMachine [s] i o
interleavingComposition alphs m =
  MealyMachine
    { states = error "states should not be necessary"
    , inputs = concat alphs
    , outputs = outputs m -- or some subset?
    , behaviour = \ss i ->
        case Map.lookup i alphLookup of
          Just idx ->
            let (o, t) = behaviour m (ss !! idx) i
                (p, _ : s) = splitAt idx ss
             in (o, p ++ [t] ++ s)
          Nothing -> error "symbol not in either alphabet"
    , initialState = replicate (length alphs) (initialState m)
    }
 where
  alphLookup = Map.fromList [(i, idx) | (alph, idx) <- zip alphs [0 ..], i <- alph]

-- mainInputDecomp :: [String] -> IO ()
-- mainInputDecomp args = case args of
--   [dotFile] -> putStrLn ("reading " <> dotFile) >> run dotFile
--   _ -> putStrLn "Please provide a dot file"
--  where
mainDecomposeInput :: DecomposeInputOptions -> CommonOptions -> IO ()
mainDecomposeInput DecomposeInputOptions{..} CommonOptions{..} = do
  let
    dotFile = filename
    report s = appendFile (logDirectory </> "hs-decompose-input.txt") (dotFile <> "\t" <> s <> "\n")
    witness s = appendFile (resultsDirectory </> "hs-decompose-input-witnesses.txt") (dotFile <> "\n" <> s <> "\n\n")

  report "START-INPUT-DECOMP"
  model <- readDotFile dotFile
  let
    inputSizes = [length (f model) | f <- [states, inputs, outputs]]

  report $ "INPUT" <> "\t" <> show inputSizes
  putStrLn $ "[states, inputs, outputs] = " <> show inputSizes

  let
    composition i j = interleavingComposition [[i], [j]] model
    bisim i j =
      let compo = composition i j
       in bisimulation2
            [i, j]
            (outputFunction model)
            (transitionFunction model)
            (initialState model)
            (outputFunction compo)
            (transitionFunction compo)
            (initialState compo)
    dependent i j = isJust $ bisim i j
    dependentPairs = [(i, j) | i <- inputs model, j <- inputs model, j > i, dependent i j]

  -- The relation dependentPairs is typically not transitive.
  let closure = foldr (uncurry equate) Data.UnionFind.empty dependentPairs
      step _ [] = Nothing
      step clos ls@(i : _) = Just (List.partition (\j -> equivalent i j clos) ls)
      classes = List.unfoldr (step closure) (inputs model)

  case length classes of
    0 -> do
      report "ERROR"
      exitFailure
    1 -> do
      report "INDECOMPOSABLE"
      putStrLn "INDECOMPOSABLE"
      exitSuccess
    n -> do
      report $ "MAYBE DECOMPOSABLE" <> "\t" <> show n
      putStrLn ("MAYBE DECOMPOSABLE: " ++ show n ++ " classes")

  let
    loop currentClosure currentClasses = do
      let
        bigCompo = interleavingComposition currentClasses model
        result = bisimulation2 (inputs model) (outputFunction model) (transitionFunction model) (initialState model) (outputFunction bigCompo) (transitionFunction bigCompo) (initialState bigCompo)

      -- print currentClasses
      print result

      case result of
        Nothing -> return currentClasses
        Just cex -> do
          let
            -- The counterexample is always the shortest. I am not completely
            -- confident that all pairs should be then consideren dependent.
            newDependentPairs = zip cex (tail cex)
            newClosure = foldr (uncurry equate) currentClosure newDependentPairs
            newClasses = List.unfoldr (step newClosure) (inputs model)
          loop newClosure newClasses

  finalClasses <- loop closure classes

  let
    reachability initial next = go [initial] Set.empty
     where
      go [] visited = visited
      go (s : todo) visited
        | Set.member s visited = go todo visited
        | otherwise = go (todo ++ next s) (Set.insert s visited)
    minimiseComponent cls =
      let
        reachableStates = reachability (initialState model) (\s -> [snd (behaviour model s i) | i <- cls])
        machine = MealyMachine{states = Set.toList reachableStates, inputs = cls, outputs = outputs model, behaviour = behaviour model, initialState = initialState model}
        partition = refineMealy machine
       in
        partition
    processClass cls = do
      let (Block p) = numBlocks . minimiseComponent $ cls
      putStr (show cls) >> putStr " of size " >> print p
      return (p, p * length cls)

  putStrLn $ "Final classes " <> show (length finalClasses)
  (sizes, transitions) <- mapAndUnzipM processClass finalClasses

  witness $ T.unpack . T.unlines . fmap T.unwords $ classes
  report $ "DECOMPOSABLE" <> "\t" <> show sizes <> "\t" <> show transitions

  putStrLn "DECOMPOSABLE"
  putStrLn $ "Total size = " <> show (sum sizes)
  putStrLn $ "Total transitions = " <> show (sum transitions)
  exitSuccess