{-# LANGUAGE OverloadedStrings #-}

module Main where

import Data.Partition
import Data.Preorder
import DotParser (readDotFile)
import DotWriter
import Mealy
import MealyRefine
import Merger

import Data.Bifunctor
import Data.List (sortOn)
import Data.List.Ordered (nubSort)
import Data.Map.Strict qualified as Map
import Data.Maybe (isNothing)
import Data.Set qualified as Set
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Data.Text.Lazy.IO qualified as TL
import Data.Tuple (swap)
import Debug.Trace (traceMarkerIO)
import System.Environment

-- | This functions inverts a map. In the new map the values are lists.
converseRelation :: Ord b => Map.Map a b -> Map.Map b [a]
converseRelation = Map.fromListWith (++) . fmap (second pure . swap) . Map.assocs

main :: IO ()
main = do
  -- Read dot file
  ls <- getArgs
  let dotFile = case ls of
        [x] -> x
        _ -> error "Please provide exactly one argument (filepath of dot file)"

  traceMarkerIO "read input"

  putStr "reading " >> putStrLn dotFile
  machine <- readDotFile dotFile

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

  traceMarkerIO "start minimisation"

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

  let
    outputFuns = [(i, fun) | i <- inputs machine, let fun s = fst (behaviour machine s i)]
    reverseTransitionMaps i = Map.fromListWith (++) [(t, [s]) | s <- states machine, let t = snd (behaviour machine s i)]
    reverseFuns = [(i, fun) | i <- inputs machine, let mm = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s mm]

  -- Minimise input, so we know the actual number of states
  printPartition (refineFuns outputFuns reverseFuns (states machine))
  putStrLn ""

  traceMarkerIO "start minimisation for each component"

  -- Then compute each projection
  let
    outs = outputs machine
    mappedOutputFuns o = [(i, (o ==) . f) | (i, f) <- outputFuns]
    projections = [(o, refineFuns (mappedOutputFuns o) reverseFuns (states machine)) | o <- outs]

  -- Print number of states of each projection
  mapM_
    ( \(o, partition) -> do
        T.putStr o
        putStr " -> "
        printPartition partition
    )
    projections

  traceMarkerIO "component equivalences"

  -- First we check for equivalent partitions, so that we skip redundant work.
  let
    (equiv, uniqPartitions) = equivalenceClasses comparePartitions projections

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

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

  traceMarkerIO "component lattice"

  -- 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 comparePartitions uniqPartitions
    foo (a, b) = (numBlocks b, a)

  putStrLn ""
  putStrLn "Top modules"
  mapM_
    ( \(b, o) -> do
        putStrLn $ "  " <> show o <> " has size " <> show b
    )
    (sortOn (negate . fst) . fmap foo $ topMods)

  traceMarkerIO "heuristic merger"

  -- 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

  print projmap

  traceMarkerIO "output"

  -- 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 = T.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 proj

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

      do
        let
          filename = "partition_" <> show componentIdx <> ".dot"
          content = T.unlines . fmap T.unwords . toBlocks $ p

        putStrLn $ "Output (partition) in file " <> filename
        T.writeFile ("results/" <> 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 = (Map.!) (getPartition p)
          -- 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
        TL.writeFile ("results/" <> filename1) content1

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

        putStrLn $ "Output (reduced machine) in file " <> filename2
        TL.writeFile ("results/" <> filename2) content2

  mapM_ action projmapN