module Main where

import Bisimulation (bisimulation2)
import Data.Partition (numBlocks)
import Data.Trie qualified as Trie
import Data.UnionFind
import DotParser (readDotFile)
import Mealy (MealyMachine (..), outputFunction, transitionFunction)
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.IO qualified as T
import System.Environment (getArgs)

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
    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] -> [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
    model <- readDotFile dotFile

    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]

    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).
-- Copar is almost twice as fast on ESM, but SplittingTree is faster on a
-- BRP benchmark. I guess, theoretically, Copar should be faster generally.
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 <- readDotFile dotFile
    putStr $ "file parsed, initial state = "
    T.putStrLn $ initialState m
    if copar
      then runCopar m
      else runSplittingTree m

  runCopar _ = error "no longer supported"
  --  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]

    (partition, _splittingTree) <- evalStateT (refine (\_ -> pure ()) outputFuns reverseFuns) (initialPRState states)
    putStrLn $ "Done" <> show (numBlocks partition)