module Main where

import DotParser
import Mealy
import MealyRefine
import Partition

import Control.Monad (forM_, when, replicateM)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.State.Strict
import Data.Function (on)
import Data.List (minimumBy, maximum, sort, intercalate)
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)
    )


  -- Now let's combine components to minimise the total size
  let totalSize = sum (fmap (numBlocks . snd) topMods)

  putStrLn ""
  putStrLn $ "num = " <> show (length topMods) <> ", size = " <> show totalSize

  let score ps p3 = numBlocks p3 - sum (fmap numBlocks ps)
      combine ops = let os = fmap fst ops
                        ps = fmap snd ops
                        p3 = foldr1 commonRefinement ps
                     in ((os, p3), score ps p3)
      isSortedOn f ls = and $ zipWith (\a b -> f a < f b) ls (drop 1 ls)
      allCombs n projs = fmap combine . filter (isSortedOn fst) $ replicateM n projs
      minComb n projs = minimumBy (compare `on` snd) (allCombs n projs)

  let loop 1 = return ()
      loop n = do
        (projmap, currentSize) <- get
        -- liftIO . print . fmap numBlocks . Map.elems $ projmap
        let ((os, p3), gain) = minComb n (Map.assocs projmap)
            o3 = intercalate "x" os
            newSize = currentSize + gain
            newProjmap = Map.insert o3 p3 . Map.withoutKeys projmap $ Set.fromList os
        liftIO . putStrLn $ show o3 <> " -> num = " <> show (Map.size newProjmap) <> ", size = " <> show newSize <> ", max = " <> show (maximum . fmap numBlocks . Map.elems $ newProjmap)
        put (newProjmap, newSize)
        if Map.size newProjmap < n
          then loop (Map.size newProjmap)
          else loop n

  putStrLn "2"
  _ <- execStateT (loop 2) (Map.fromList topMods, totalSize)

  return ()