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

module Main where

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 System.Environment

extraChecks :: Bool
extraChecks = False

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

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

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

  -- 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 strategy MergerStats{..}
        | numberOfComponents <= 4 = Stop
        | otherwise = Continue

  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"

  -- OUTPUT
  ---------
  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 $ "\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 ("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

-- * 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 _ _ _ -> IO _
printBasics outputFuns reverseFuns MealyMachine{..} = do
  putStrLn $ (show . length $ states) <> " states, " <> (show . length $ inputs) <> " inputs and " <> (show . length $ outputs) <> " outputs"
  when extraChecks $ 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)