module Main where

import DotParser
import Mealy
import MyLib

import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.State.Strict
import Control.Monad (forM_, when)
import Control.Monad.ST (runST)
import Copar.Algorithm (refine)
import Copar.Functors.Polynomial (Polynomial)
import Data.Map.Strict qualified as Map
import Data.Maybe (mapMaybe)
import Data.Partition (isRefinementOf, numBlocks)
import Data.Proxy (Proxy(..))
import Data.Semigroup (Arg(..))
import Data.Set qualified as Set
import Data.List.Ordered (nubSort)
import System.Environment
import Text.Megaparsec

{-
  Hacked together, you can view the result with:

    tred relation.dot | dot -Tpng -G"rankdir=BT" > relation.png

  tred is the graphviz tool to remove transitive edges. And the rankdir
  attribute flips the graph upside down.
-}
main :: IO ()
main = do
  -- Read dot file
  [dotFile] <- getArgs
  print dotFile
  transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile

  -- convert to mealy
  let machine = convertToMealy transitions

  -- print some basic info
  putStrLn $ (show . length $ states machine) <> " states, " <> (show . length $ inputs machine) <> " inputs and " <> (show . length $ outputs machine) <> " outputs"
  putStrLn "Small sample:"
  print . take 4 . states $ machine
  print . take 4 . inputs $ machine
  print . take 4 . outputs $ machine
  putStrLn ""

  -- DEBUG OUTPUT
  -- forM_ (states machine) (\s -> do
  --   print s
  --   forM_ (inputs machine) (\i -> do
  --     putStr "  "
  --     let (o, t) = behaviour machine s i
  --     putStrLn $ "--" <> (show i) <> "/" <> (show o) <> "->" <> (show t)
  --     )
  --   )

  let printPartition p = putStrLn $ "  number of states = " <> show (numBlocks p)

  -- Minimise input, so we know the actual number of states
  printPartition (runST $ refine (Proxy @Polynomial) (mealyMachineToEncoding machine) True)
  putStrLn ""

  -- Then compute each projection
  let minimiseProjection o = runST $ refine (Proxy @Polynomial) (mealyMachineToEncoding (project machine o)) True
      outs = outputs machine
      projections = Map.fromSet minimiseProjection (Set.fromList outs)

  -- Print number of states of each projection
  forM_ (Map.assocs projections) (\(o, partition) -> do
    print o
    printPartition partition
    putStrLn ""
    )

  -- Consider all pairs
  let keys = Map.argSet projections
      pairs = Set.cartesianProduct keys keys
      distinctPairs = Set.filter (\(Arg o1 _, Arg o2 _) -> o1 < o2) pairs -- bit inefficient

  -- Check refinement relations for all pairs
  (equiv, rel) <- flip execStateT (Map.empty, []) $ do
    forM_ (Map.assocs projections) (\(o1, b1) -> do
      (repr, _) <- get
      if o1 `Map.member` repr
        then do
          liftIO . putStrLn $ "skip " <> (show o1) <> " |-> " <> (show (repr Map.! o1))
        else do
          liftIO $ print o1
          forM_ (Map.assocs projections) (\(o2, b2) -> do
            (repr0, _) <- get
            when (o1 < o2 && o2 `Map.notMember` repr0) $ do
              case (isRefinementOf b1 b2, isRefinementOf b2 b1) of
                (True, True) -> do
                  (repr, ls) <- get
                  put (Map.insert o2 o1 repr, ls)
                (True, False) -> do
                  (repr, ls) <- get
                  put (repr, (o1, o2):ls)
                (False, True) -> do
                  (repr, ls) <- get
                  put (repr, (o2, o1):ls)
                (False, False) -> return ()

              -- liftIO $ putStr " vs. "
              -- liftIO $ print o2
            )
      )

  putStrLn ""
  putStrLn "Equivalences"
  forM_ (Map.assocs equiv) (\(o2, o1) -> do
    putStrLn $ "  " <> (show o2) <> " == " <> (show o1)
    )

  let cleanRel = [(Map.findWithDefault o1 o1 equiv, Map.findWithDefault o2 o2 equiv) | (o1, o2) <- rel]
  putStrLn ""
  putStrLn "Relation (smaller points to bigger)"
  forM_ (nubSort cleanRel) (\(o1, o2) -> do
    putStrLn $ "  " <> (show o2) <> " -> " <> (show o1)
    )

  return ()