{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE MonoLocalBinds #-}
{-|
   Given an e-graph representing expressions of our language, we might want to
   extract, out of all expressions represented by some equivalence class, /the best/
   expression (according to a 'CostFunction') represented by that class

   The function 'extractBest' allows us to do exactly that: get the best
   expression represented in an e-class of an e-graph given a 'CostFunction'
 -}
module Data.Equality.Extraction
  (
  -- * Extraction
    extractBest

  -- * Cost
  , CostFunction
  , depthCost
  ) where

import qualified Data.Set as S
import qualified Data.IntMap.Strict as IM

import Data.Equality.Graph.Internal (EGraph(classes))
import Data.Equality.Utils
import Data.Equality.Graph

-- vvvv and necessarily all the best sub-expressions from children equilalence classes

-- | Extract the /best/ expression from an equivalence class according to a
-- 'CostFunction'
--
-- @
-- (i, egr) = ...
--    i <- represent expr
--            ...
--
-- bestExpr = extractBest egr 'depthCost' i
-- @
--
-- For a real example you might want to check out the source code of 'Data.Equality.Saturation.equalitySaturation''
extractBest :: forall anl lang cost
             . (Language lang, Ord cost)
            => EGraph anl lang        -- ^ The e-graph out of which we are extracting an expression
            -> CostFunction lang cost -- ^ The cost function to define /best/
            -> ClassId                -- ^ The e-class from which we'll extract the expression
            -> Fix lang               -- ^ The resulting /best/ expression, in its fixed point form.
extractBest egr cost (flip find egr -> i) = 

    -- Use `egg`s strategy of find costs for all possible classes and then just
    -- picking up the best from the target e-class.  In practice this shouldn't
    -- find the cost of unused nodes because the "topmost" e-class will be the
    -- target, and all sub-classes must be calculated?
    let allCosts = findCosts (classes egr) mempty

     in case findBest i allCosts of
        Just (CostWithExpr (_,n)) -> n
        Nothing    -> error $ "Couldn't find a best node for e-class " <> show i

  where

    -- | Find the lowest cost of all e-classes in an e-graph in an extraction
    findCosts :: ClassIdMap (EClass anl lang) -> ClassIdMap (CostWithExpr lang cost) -> ClassIdMap (CostWithExpr lang cost)
    findCosts eclasses current =

      let (modified, updated) = IM.foldlWithKey f (False, current) eclasses

          {-# INLINE f #-}
          f :: (Bool, ClassIdMap (CostWithExpr lang cost)) -> Int -> EClass anl lang -> (Bool, ClassIdMap (CostWithExpr lang cost))
          f acc@(_, beingUpdated) i' EClass{eClassNodes = nodes} =
                let
                    currentCost = IM.lookup i' beingUpdated

                    newCost = S.foldl' (\c n -> case (c, nodeTotalCost beingUpdated n) of
                                                  (Nothing, Nothing) -> Nothing
                                                  (Nothing, Just nc) -> Just nc
                                                  (Just oc, Nothing) -> Just oc
                                                  (Just oc, Just nc) -> Just (oc `min` nc)
                                       ) Nothing nodes
                    -- Current cost + get lowest cost and corresponding node of an e-class if possible
                 in case (currentCost, newCost) of

                    (Nothing, Just new) -> (True, IM.insert i' new beingUpdated)

                    (Just (CostWithExpr old), Just (CostWithExpr new))
                      | fst new < fst old -> (True, IM.insert i' (CostWithExpr new) beingUpdated)

                    _ -> acc

        -- If any class was modified, loop
       in if modified
            then findCosts eclasses updated
            else updated

    -- | Get the total cost of a node in an e-graph if possible at this stage of
    -- the extraction
    --
    -- For a node to have a cost, all its (canonical) sub-classes have a cost and
    -- an associated better expression. We return the constructed best expression
    -- with its cost
    nodeTotalCost :: Traversable lang => ClassIdMap (CostWithExpr lang cost) -> ENode lang -> Maybe (CostWithExpr lang cost)
    nodeTotalCost m (Node n) = do
        expr <- traverse ((`IM.lookup` m) . flip find egr) n
        return $ CostWithExpr (cost (fst . unCWE <$> expr), Fix $ snd . unCWE <$> expr)
    {-# INLINE nodeTotalCost #-}
{-# INLINABLE extractBest #-}

-- | A cost function is used to attribute a cost to representations in the
-- e-graph and to extract the best one.
--
-- The cost function is polymorphic over the type used for the cost, however
-- @cost@ must instance 'Ord' in order for the defined 'CostFunction' to
-- fulfill its purpose. That's why we have an @Ord cost@ constraint in
-- 'Data.Equality.Saturation.equalitySaturation' and 'extractBest'
--
-- === Example
-- @
-- symCost :: Expr Int -> Int
-- symCost = \case
--     BinOp Integral e1 e2 -> e1 + e2 + 20000
--     BinOp Diff e1 e2 -> e1 + e2 + 500
--     BinOp x e1 e2 -> e1 + e2 + 3
--     UnOp x e1 -> e1 + 30
--     Sym _ -> 1
--     Const _ -> 1
-- @
type CostFunction l cost = l cost -> cost

-- | Simple cost function: the deeper the expression, the bigger the cost
depthCost :: Language l => CostFunction l Int
depthCost = (+1) . sum
{-# INLINE depthCost #-}

-- | Find the current best node and its cost in an equivalence class given only the class and the current extraction
-- This is not necessarily the best node in the e-graph, only the best in the current extraction state
findBest :: ClassId -> ClassIdMap (CostWithExpr lang a) -> Maybe (CostWithExpr lang a)
findBest = IM.lookup
{-# INLINE findBest #-}

newtype CostWithExpr lang a = CostWithExpr { unCWE :: (a, Fix lang) }

instance Eq a => Eq (CostWithExpr lang a) where
  (==) (CostWithExpr (a,_)) (CostWithExpr (b,_)) = a == b
  {-# INLINE (==) #-}

instance Ord a => Ord (CostWithExpr lang a) where
  compare (CostWithExpr (a,_)) (CostWithExpr (b,_)) = a `compare` b
  {-# INLINE compare #-}