module Main where

import Bisimulation (bisimulation2)
import Data.UnionFind
import DotParser (convertToMealy, parseTransFull)
import Mealy (MealyMachine (..), outputFunction, transitionFunction)
import Data.Partition (numBlocks)
import SplittingTree (PRState (..), getPartition, initialPRState, refine)
import StateIdentifiers (stateIdentifierFor)
import Data.Trie qualified as Trie

import Control.Monad.Trans.State (execStateT)
import Data.List qualified as List
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust, mapMaybe)
import Data.Set qualified as Set
import MealyRefine
import System.Environment (getArgs)
import Text.Megaparsec (parseMaybe)

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

mainHSI :: [String] -> IO ()
mainHSI args = case args of
  [dotFile] -> run dotFile
  _ -> putStrLn "Please provide a dot file"
 where
  run dotFile = do
    print dotFile
    transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile

    -- convert to mealy
    let
      MealyMachine{..} = convertToMealy transitions
      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]

    PRState{..} <- execStateT (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] -> [i] -> MealyMachine s i o -> MealyMachine (s, s) i o
interleavingComposition alph1 alph2 m =
  MealyMachine
    { states = error "states should not be necessary"
    , inputs = alph1 ++ alph2
    , outputs = error "outputs should not be necessary"
    , behaviour = \(s1, s2) i ->
        case Map.lookup i alphLookup of
          Just False -> let (o, t) = behaviour m s1 i in (o, (t, s2))
          Just True -> let (o, t) = behaviour m s2 i in (o, (s1, t))
          Nothing -> error "symbol not in either alphabet"
    , initialState = (initialState m, initialState m)
    }
 where
  alphLookup = Map.fromList ([(a, False) | a <- alph1] ++ [(a, True) | a <- alph2])

mainInputDecomp :: [String] -> IO ()
mainInputDecomp args = case args of
  [dotFile] -> run dotFile
  _ -> putStrLn "Please provide a dot file"
 where
  run dotFile = do
    print dotFile
    transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile

    let model = convertToMealy transitions
        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]

    print $ length (states model)
    print $ length (inputs model)
    print $ length (outputs model)

    putStrLn "Dependent pairs:"
    print $ length dependentPairs

    let dps = Set.fromList dependentPairs
        dpsFun i j = i == j || (i, j) `Set.member` dps || (j, i) `Set.member` dps
        trans =
          null [(i, j, k) | i <- inputs model, j <- inputs model, dpsFun i j, k <- inputs model, dpsFun j k, not (dpsFun i k)]

    putStrLn "Transitive?"
    print trans

    let closure = foldr (uncurry equate) empty dependentPairs
        step [] = Nothing
        step ls@(i : _) = Just (List.partition (\j -> equivalent i j closure) ls)
        classes = List.unfoldr step (inputs model)

    mapM_ print classes
    case length classes of
      0 -> putStrLn "ERROR"
      1 -> putStrLn "INDECOMPOSABLE"
      n -> putStrLn ("MAYBE DECOMPOSABLE: " ++ show n ++ " classes")

-- Used to determine whether Copar is faster than SplittingTree (it is).
mainRefine :: [String] -> IO ()
mainRefine args = case args of
  [dotFile, copar] -> run dotFile (read copar)
  _ -> putStrLn "Please provide a dot file and Boolean"
 where
  run dotFile copar = do
    m <- convertToMealy . mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile
    putStrLn $ "file parsed, initial state = " <> initialState m
    if copar
      then runCopar m
      else runSplittingTree m

  runCopar m =
    let printPartition p = putStrLn $ "Done " <> show (numBlocks p)
     in printPartition (refineMealy (mealyMachineToEncoding m))

  runSplittingTree MealyMachine{..} = do
    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 mm = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s mm]

    PRState{..} <- execStateT (refine (\_ -> pure ()) outputFuns reverseFuns) (initialPRState states)
    putStrLn $ "Done" <> show (Map.size (getPartition partition))