module Main where

import DotParser
import Mealy
import MealyRefine
import Merger
import Partition

import Control.Monad (forM_, when)
import Control.Monad.Trans.State.Strict
import Data.List (sort)
import Data.Map.Strict qualified as Map
import Data.Maybe (mapMaybe)
import Data.Set qualified as Set
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

  -- -- 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 (refineMealy (mealyMachineToEncoding machine))
  putStrLn ""

  -- Then compute each projection
  -- I did some manual preprocessing, these are the only interesting bits
  let -- outs = ["10", "10-O9", "2.2", "3.0", "3.1", "3.10", "3.12", "3.13", "3.14", "3.16", "3.17", "3.18", "3.19", "3.2", "3.20", "3.21", "3.3", "3.4", "3.6", "3.7", "3.8", "3.9", "5.0", "5.1", "5.12", "5.13", "5.17", "5.2", "5.21", "5.23", "5.6", "5.7", "5.8", "5.9", "quiescence"]
      outs = outputs machine
      projections0 = allProjections machine outs
      projections = zip outs $ fmap refineMealy projections0

  -- Print number of states of each projection
  forM_ projections (\(o, partition) -> do
    putStr $ o <> " -> "
    printPartition partition
    )

  -- First we check eqiuvalent partitions, so that we only work on one
  -- item in each equivalence class. This could be merged with the next
  -- phase of checking refinement, and that would be faster. But this is
  -- simpler.
  let checkRelsFor o1 p1 =
        forM_ projections (\(o2, p2) -> do
          (repr, ls) <- get
          -- We skip if o2 is equivelent to an earlier o
          when (o1 < o2 && o2 `Map.notMember` repr) $ do
            case isEquivalent p1 p2 of
              -- Equivalent: let o2 point to o1
              True -> put (Map.insert o2 o1 repr, ls)
              False -> return ()
          )
      checkAllRels projs =
        forM_ projs (\(o1, p1) -> do
          -- First we check if o1 is equivalent to an earlier o
          -- If so, we skip it. Else, we add it to the unique
          -- components and compare to all others.
          (repr, ls) <- get
          when (o1 `Map.notMember` repr) $ do
            put (repr, (o1, p1):ls)
            checkRelsFor o1 p1
          )
      (equiv, uniqPartitions) = execState (checkAllRels projections) (Map.empty, [])

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

  -- Then we compare each pair of partitions. If one is a coarsening of
  -- another, we can skip it later on. That is to say: we only want the
  -- finest partitions.
  let compareAll partitions =
        forM_ partitions (\(o1, b1) ->
          forM_ partitions (\(o2, b2) ->
            when (o1 < o2) $ do
              ls <- get
              case comparePartitions b1 b2 of
                Equivalent -> error "cannot happen"
                Refinement -> put $ (o1, o2):ls
                Coarsening -> put $ (o2, o1):ls
                Incomparable -> return ()
            )
          )
      rel = execState (compareAll uniqPartitions) []

  putStrLn ""
  putStrLn "Relation, coarser points to finer (bigger)"
  forM_ rel (\(o1, o2) -> do
    putStrLn $ "  " <> (show o2) <> " -> " <> (show o1)
    )

  -- Get rid of the coarser partitions
  let lowElements = Set.fromList . fmap snd $ rel
      allElements = Set.fromList . fmap fst $ uniqPartitions
      topElements = Set.difference allElements lowElements
      mods = Map.fromList uniqPartitions -- could be a lazy map
      topMods = Map.assocs $ Map.restrictKeys mods topElements
      foo (a, b) = (numBlocks b, a)

  putStrLn ""
  putStrLn "Top modules"
  forM_ (reverse . sort . fmap foo $ topMods) (\(b, o) -> do
    putStrLn $ "  " <> (show o) <> " has size " <> (show b)
    )

  let strategy MergerStats{..}
        | numberOfComponents <= 4 = Stop
        | otherwise = Continue

  projmap <- heuristicMerger topMods strategy

  print . fmap fst $ projmap

  return ()