module Merger where

import Partition

import Control.Monad (replicateM)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.State.Strict
import Data.Bifunctor (first)
import Data.List (minimumBy)
import Data.Map.Strict qualified as Map
import Data.Ord (comparing)
import Data.Set qualified as Set

data MergerStats = MergerStats
  { numberOfComponents :: Int
  , maximalComponent :: Int
  , totalSize :: Int
  }
  deriving (Eq, Ord, Read, Show)

data MergerAction = Stop | Continue
  deriving (Eq, Ord, Read, Show)

type MergerStrategy = MergerStats -> MergerAction

heuristicMerger :: Ord o => [(o, Partition)] -> MergerStrategy -> IO [([o], Partition)]
heuristicMerger components strategy = do
  projmap <- evalStateT (loop 2) (Map.fromList (fmap (first pure) components))
  return $ Map.assocs projmap
  where
    score ps p3 = numBlocks p3 - sum (fmap numBlocks ps)
    combine ops = let os = fmap fst ops
                      ps = fmap snd ops
                      p3 = foldr1 commonRefinement ps
                    in ((os, p3), score ps p3)
    isSortedOn f ls = and $ zipWith (\a b -> f a < f b) ls (drop 1 ls)
    allCombs n projs = fmap combine . filter (isSortedOn fst) $ replicateM n projs
    minComb n projs = minimumBy (comparing snd) (allCombs n projs)
    safeStrategy ms@MergerStats{..}
      | numberOfComponents <= 1 = Stop
      | otherwise = strategy ms

    loop n = do
      projmap <- get
      let numberOfComponents = Map.size projmap
          componentSizes = fmap numBlocks . Map.elems $ projmap
          maximalComponent = maximum componentSizes
          totalSize = sum componentSizes
          ms = MergerStats{..}
      liftIO . print $ ms
      case safeStrategy ms of
        Stop -> return projmap
        Continue -> do
          let ((os, p3), _) = minComb n (Map.assocs projmap)
              o3 = mconcat os
              newProjmap = Map.insert o3 p3 . Map.withoutKeys projmap $ Set.fromList os
          put newProjmap
          if Map.size newProjmap < n
            then loop (Map.size newProjmap)
            else loop n