module Main where

import DotParser
import DotWriter
import Mealy
import MealyRefine
import Merger
import Partition

import Control.Monad (forM_, when)
import Control.Monad.Trans.State.Strict
import Data.Bifunctor
import Data.List (sort, sortOn, intercalate)
import Data.List.Ordered (nubSort)
import Data.Map.Strict qualified as Map
import Data.Maybe (mapMaybe)
import Data.Set qualified as Set
import Data.Tuple (swap)
import System.Environment
import Text.Megaparsec

converseRelation :: (Ord a, Ord b) => Map.Map a b -> Map.Map b [a]
converseRelation m = Map.fromListWith (++) . fmap (second pure . swap) . Map.assocs $ m

{-
  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, state2idx) = 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

  let equivInv = converseRelation equiv
      relMap = Map.fromListWith (++) . fmap (second pure) $ rel
      projmapN = zip projmap [1..]

  forM_ projmapN (\((os, p), i) -> do
    let name = intercalate "x" os
        filename = "component" <> show i <> ".dot"

        osWithRel = concat $ os:[Map.findWithDefault [] o relMap | o <- os]
        osWithRelAndEquiv = concat $ osWithRel:[Map.findWithDefault [] o equivInv | o <- osWithRel]
        componentOutputs = Set.fromList osWithRelAndEquiv
        proj = projectToComponent (flip Set.member componentOutputs) machine
        -- Sanity check: compute partition again
        partition = refineMealy . mealyMachineToEncoding $ proj

    putStrLn $ ""
    putStrLn $ "Component " <> show os
    putStrLn $ "Correct? " <> show (comparePartitions p partition)
    putStrLn $ "Size = " <> show (numBlocks p)
    putStrLn $ "Output in file " <> filename

    let MealyMachine{..} = proj
        -- We enumerate all transitions in the full automaton
        transitions = [(s, i, o, t) | s <- states, i <- inputs, let (o, t) = behaviour s i]
        -- This is the quotient map, from state to block
        state2block = blockOfState p . (state2idx Map.!)
        -- We apply this to each transition, and then nubSort the duplicates away
        transitionsBlocks = nubSort [(state2block s, i, o, state2block t) | (s, i, o, t) <- transitions]
        -- The initial state should be first
        initialBlock = state2block initialState
        -- Sorting on "/= initialBlock" puts the initialBlock in front
        initialFirst = sortOn (\(s,_,_,_) -> s /= initialBlock) transitionsBlocks
        -- Convert to a file
        content = toString . mealyToDot name $ initialFirst

    writeFile filename content
    )

  return ()