module Main where

import Bisimulation (bisimulation2)
import Data.Partition (Block (..), numBlocks)
import Data.Trie qualified as Trie
import Data.UnionFind (empty, equate, equivalent)
import DotParser (readDotFile)
import Mealy (MealyMachine (..), outputFunction, transitionFunction)
import MealyRefine (refineMealy)
import SplittingTree (initialPRState, refine)
import StateIdentifiers (stateIdentifierFor)

import Control.Monad.Trans.State (evalStateT)
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 System.Environment (getArgs)
import System.Exit (exitFailure, exitSuccess)

main :: IO ()
main = do
  args <- getArgs
  case args of
    ("HSI" : ls) -> mainHSI ls
    ("InputDecomp" : ls) -> mainInputDecomp ls
    _ -> putStrLn "Please provide one of [HSI, InputDecomp]"

mainHSI :: [String] -> IO ()
mainHSI args = case args of
  [dotFile] -> run dotFile
  _ -> putStrLn "Please provide a dot file"
 where
  run dotFile = do
    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)

-- Interleaving composition of restriction to subalphabets
-- precondigiotn: 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
  run dotFile = do
    let
      report str = appendFile "results/log.txt" (dotFile <> "\t" <> str <> "\n")
      witness str = appendFile "results/witnesses.txt" (dotFile <> "\n" <> str <> "\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) 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
      action cls = do
        let (Block p) = numBlocks . minimiseComponent $ cls
        putStr (show cls) >> putStr " of size " >> putStrLn (show p)
        return (p, p * length cls)

    putStrLn $ "Final classes " <> show (length finalClasses)
    (sizes, transitions) <- unzip <$> mapM action 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