{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}

module DecomposeOutput where

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

import Control.Monad (when)
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 Options.Applicative
import System.FilePath ((</>))

data DecomposeOutputOptions = DecomposeOutputOptions
  { filename :: FilePath
  , numComponents :: Int
  }
  deriving Show

decomposeOutputOptionsParser :: Parser DecomposeOutputOptions
decomposeOutputOptionsParser =
  DecomposeOutputOptions
    <$> argument str (help "Filename to read (dot format)" <> metavar "FILE")
    <*> option auto (long "components" <> short 'c' <> help "Number of components" <> metavar "NUM" <> showDefault <> value 2)
      <**> helper

mainDecomposeOutput :: DecomposeOutputOptions -> CommonOptions -> IO ()
mainDecomposeOutput DecomposeOutputOptions{..} CommonOptions{..} = do
  let
    report s = appendFile (logDirectory </> "hs-decompose-output-hs.txt") (filename <> "\t" <> s <> "\n")

  -- READING INPUT
  ----------------
  putStrLn $ "reading " <> filename
  machine <- readDotFile filename

  -- PREPROCESSING
  ----------------
  let (outputFuns, reverseFuns) = preprocess machine
  printBasics outputFuns reverseFuns machine extraChecks

  -- MINIMISING EACH COMPONENT
  ----------------------------
  let mappedOutputFuns o = [(i, (o ==) . f) | (i, f) <- outputFuns]
      projections = [(o, refineFuns (mappedOutputFuns o) reverseFuns (states machine)) | o <- outputs machine]

  putStrLn $ "\nComponents " <> show (length (outputs machine))
  mapM_ (\(o, p) -> putStr "  " >> T.putStr o >> putStr " has size " >> print (numBlocks p)) projections

  -- REDUCING NUMBER OF COMPONENTS
  -- by checking which partitions are equivalent
  ----------------------------------------------
  let (equiv, uniqPartitions) = equivalenceClasses comparePartitions projections

  putStrLn $ "\nRepresentatives " <> show (length uniqPartitions)
  print . fmap fst $ uniqPartitions

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

  -- COMPUTING THE LATTICE OF COMPONENTS
  -- 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)
    sortedTopMods = sortOn (negate . fst) . fmap foo $ topMods

  putStrLn $ "\nTop modules " <> show (length topMods)
  mapM_ (\(b, o) -> putStr "  " >> T.putStr o >> putStr " has size " >> print b) sortedTopMods

  -- HEURISTIC MERGING
  -- 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
    numStrategy current
      | numberOfComponents current <= numComponents = StopWith (componentPartitions current)
      | otherwise = Continue
    prevStrategy current = case previous current of
      Just prev -> if totalSize prev < totalSize current then StopWith (componentPartitions prev) else Continue
      _ -> Continue
    strategy c = case prevStrategy c of
      StopWith x -> StopWith x
      Continue -> numStrategy c

  putStrLn "\nHeuristic merging"
  projmap <- heuristicMerger topMods strategy

  putStrLn "\nDone"
  putStrLn $ "  components: " <> show (length projmap)
  putStrLn $ "  sizes: " <> show (fmap (numBlocks . snd) projmap)
  putStrLn "Start writing output files"

  report $ "PAR-BIT-DECOMP" <> "\t" <> show (length (states machine)) <> "\t" <> show (length (inputs machine)) <> "\t" <> show (length (outputs machine)) <> "\t" <> show (length projmap) <> "\t" <> show (sum (fmap (numBlocks . snd) projmap)) <> "\t" <> show (fmap (numBlocks . snd) projmap)

  -- OUTPUT
  ---------
  let
    equivInv = converseRelation equiv
    projmapN = zip projmap [1 :: Int ..]
    processComponent ((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 $ "\nComponent " <> show os
      when extraChecks (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 (resultsDirectory </> 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 (resultsDirectory </> filename1) content1

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

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

  mapM_ processComponent projmapN

-- * Helper functions

-- | Computes the predecessors of each state.
preprocess :: _ => MealyMachine _ _ _ -> _
preprocess MealyMachine{..} = (outputFuns, reverseFuns)
 where
  outputFuns = [(i, fun) | i <- inputs, let fun s = fst (behaviour s i)]
  reverseTransitionMaps i = Map.fromListWith (++) [(t, [s]) | s <- states, let t = snd (behaviour s i)]
  reverseFuns = [(i, fun) | i <- inputs, let mm = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s mm]

-- | Prints basic info.
printBasics :: _ => _ -> _ -> MealyMachine _ _ _ -> Bool -> IO _
printBasics outputFuns reverseFuns MealyMachine{..} extraChecksEnabled = do
  putStrLn $ (show . length $ states) <> " states, " <> (show . length $ inputs) <> " inputs and " <> (show . length $ outputs) <> " outputs"
  when extraChecksEnabled $ do
    printPartition (refineFuns outputFuns reverseFuns states)
    putStrLn ""

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

-- | Prints the number of blocks.
printPartition :: Partition s -> IO ()
printPartition p = putStrLn $ "number of states = " <> show (numBlocks p)