{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Avoid reverse" #-}
module Main where

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

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

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

myWriteFile :: FilePath -> String -> IO ()
myWriteFile filename = writeFile ("results/" ++ filename)

{-
  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
  machine <- convertToMealy . mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile

  -- 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 for equivalent partitions, so that we skip redundant work.
  let
    preord p1 p2 = toPreorder (comparePartitions p1 p2)
    (equiv, uniqPartitions) = equivalenceClasses preord projections

  putStrLn ""
  putStrLn "Representatives"
  print . fmap fst $ uniqPartitions

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

  -- Then we compare each pair of partitions. We only keep the finest
  -- partitions, since the coarse ones don't provide value to us.
  let
    (topMods, downSets) = maximalElements preord uniqPartitions
    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
    )

  -- Then we try to combine paritions, so that we don't end up with
  -- too many components. (Which would be too big to be useful.)
  let strategy MergerStats{..}
        | numberOfComponents <= 4 = Stop
        | otherwise = Continue

  projmap <- heuristicMerger topMods strategy

  -- Now we are going to output the components we found.
  let
    equivInv = converseRelation equiv
    projmapN = zip projmap [1 :: Int ..]
    action ((os, p), componentIdx) = do
      let
        name = intercalate "x" os
        osWithRel = concat $ os : [Map.findWithDefault [] o downSets | o <- os]
        osWithRelAndEquiv = concat $ osWithRel : [Map.findWithDefault [] o equivInv | o <- osWithRel]
        componentOutputs = Set.fromList osWithRelAndEquiv
        proj = projectToComponent (`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)

      do
        let
          filename = "partition_" <> show componentIdx <> ".dot"
          idx2State = Map.map head . converseRelation $ state2idx
          stateBlocks = fmap (fmap (idx2State Map.!)) . Partition.toBlocks $ partition
          content = unlines . fmap unwords $ stateBlocks

        putStrLn $ "Output (partition) in file " <> filename
        myWriteFile filename content

      do
        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
          filename1 = "component_" <> show componentIdx <> ".dot"
          content1 = toString . mealyToDot name $ initialFirst

          -- So far so good, `initialFirst` could serve as our output
          -- But we do one more optimisation on the machine
          -- We remove inputs, on which the machine does nothing
          deadInputs0 = Map.fromListWith (++) . fmap (\(s, i, o, t) -> (i, [(s, o, t)])) $ initialFirst
          deadInputs = Map.keysSet . Map.filter (all (\(s, o, t) -> s == t && isNothing o)) $ deadInputs0
          result = filter (\(_, i, _, _) -> i `Set.notMember` deadInputs) initialFirst

          -- Convert to a file
          filename2 = "component_reduced_" <> show componentIdx <> ".dot"
          content2 = toString . mealyToDot name $ result

        putStrLn $ "Output (reduced machine) in file " <> filename1
        myWriteFile filename1 content1

        putStrLn $ "Dead inputs = " <> show (Set.size deadInputs)

        putStrLn $ "Output (reduced machine) in file " <> filename2
        myWriteFile filename2 content2

  mapM_ action projmapN