Info.hs

-----------------------------------------------------------------------------
-- |
-- Module      :  Algorithm.EqSat.Info
-- Copyright   :  (c) Fabricio Olivetti 2021 - 2024
-- License     :  BSD3
-- Maintainer  :  fabricio.olivetti@gmail.com
-- Stability   :  experimental
-- Portability :
--
-- Functions related to info/data calculation in Equality Graph data structure
-- Heavily based on hegg (https://github.com/alt-romes/hegg by alt-romes)
--
-----------------------------------------------------------------------------

module Algorithm.EqSat.Info where

import Control.Lens ( over )
import Control.Monad --(forM, forM_, when, foldM, void)
import Control.Monad.State
import Data.AEq (AEq ((~==)))
import Data.IntMap (IntMap) -- , delete, empty, insert, toList)
import qualified Data.IntMap as IntMap
import Data.Map (Map)
import qualified Data.Map as Map
import Data.SRTree
import Data.SRTree.Eval (evalFun, evalOp, PVector)
import Data.HashSet (HashSet)
import qualified Data.HashSet as Set
import qualified Data.IntSet as IntSet
import Algorithm.EqSat.Egraph
import Data.AEq (AEq ((~==)))
import Algorithm.EqSat.Queries

import Data.Maybe
import qualified Data.Set as TrueSet
import Data.Sequence (Seq(..), (><))

import Debug.Trace

-- * Data related functions 

-- | join data from two e-classes
-- TODO: instead of folding, just do not apply rules
-- list of values instead of single value
joinData :: EClassData -> EClassData -> EClassData
joinData (EData c1 b1 cn1 fit1 dl1 p1 sz1) (EData c2 b2 cn2 fit2 dl2 p2 sz2) =
  --EData (min c1 c2) b (combineConsts cn1 cn2) (minMaybe fit1 fit2) (bestParam p1 p2 fit1 fit2) (min sz1 sz2)
  EData (min c1 c2) (choose b1 b2) (choose cn1 cn2) (maxMaybe fit1 fit2) (choose dl1 dl2) (choose p1 p2) (choose sz1 sz2)
  where
    isFst = c1 <= c2
    choose x y = if isFst then x else y
    chooseF x y = if maxIsFst then x else y

    maxIsFst = case (fit1, fit2) of
                 (Nothing, Nothing) -> True
                 (Nothing,  Just f) -> False
                 (Just f , Nothing) -> True
                 (Just f1, Just f2) -> f1 >= f2

    maxMaybe Nothing x = x
    maxMaybe x Nothing = x
    maxMaybe x y       = max x y

    bestParam Nothing x _ _ = x
    bestParam x Nothing _ _ = x
    bestParam x y (Just f1) (Just f2) = if f1 >= f2 then x else y

    b = if c1 <= c2 then b1 else b2
    combineConsts (ConstVal x) (ConstVal y)
      | abs (x-y) < 1e-7   = ConstVal $ (x+y)/2
      | isNaN x || isInfinite x = ConstVal y 
      | isNaN y || isInfinite y = ConstVal x
      | isNaN x && isNaN y = ConstVal x
      | x ~== y = ConstVal $ (x+y)/2
      | abs (x / y) < 1 + 1e-6 || abs (y / x) < 1 + 1e-6 = ConstVal $ min x y
      | isInfinite x && isInfinite y = ConstVal x
      | isInfinite x && isNaN y = ConstVal y
      | isNaN x && isInfinite y = ConstVal x
      | otherwise          = error $ "Combining different values: " <> show x <> " " <> show y <> " " <> show (x/y)
    combineConsts (ParamIx ix) (ParamIx iy) = ParamIx (min ix iy)
    combineConsts NotConst x = x
    combineConsts x NotConst = x
    combineConsts (ParamIx ix) (ConstVal x) = ConstVal x
    combineConsts (ConstVal x) (ParamIx ix) = ConstVal x -- p - p = 0
    combineConsts x y = error (show x <> " " <> show y)

-- | Calculate e-node data (constant values and cost)
makeAnalysis :: Monad m => CostFun -> ENode -> EGraphST m EClassData
makeAnalysis costFun enode =
  do consts <- calculateConsts enode
     enode' <- canonize enode
     cost   <- calculateCost costFun enode'
     sz <- sum <$> mapM (\ecId -> gets (_size . _info . (IntMap.! ecId) . _eClass)) (childrenOf enode')
     pure $ EData cost enode' consts Nothing Nothing [] (sz+1)

getChildrenMinHeight :: Monad m => ENode -> EGraphST m Int
getChildrenMinHeight enode = do
  let children = childrenOf enode
      minimum' [] = 0
      minimum' xs = minimum xs
  minimum' <$> mapM (\ec -> gets (_height . (IntMap.! ec) . _eClass)) children

-- | update the heights of each e-class
-- won't work if there's no root
calculateHeights :: Monad m => EGraphST m ()
calculateHeights =
  do queue   <- findRootClasses
     classes <- gets (Prelude.map fst . IntMap.toList . _eClass)
     let nClasses = length classes
     forM_ classes (setHeight nClasses) -- set all heights to max possible height (number of e-classes)
     forM_ queue (setHeight 0)          -- set root e-classes height to zero
     go queue (TrueSet.fromList queue) 1    -- next height is 1
  where
    setHeight x eId' =
      do eId <- canonical eId'
         ec <- getEClass eId
         let ec' = over height (const x) ec
         modify' $ over eClass (IntMap.insert eId ec')

    setMinHeight x eId' = -- set height to the minimum between current and x
      do eId <- canonical eId'
         h <- _height <$> getEClass eId
         setHeight (min h x) eId

    getChildrenEC :: Monad m => EClassId -> EGraphST m [EClassId]
    getChildrenEC ec' = do ec <- canonical ec'
                           gets (concatMap childrenOf' . _eNodes . (IntMap.! ec) . _eClass)

    childrenOf' (_, -1, -1, _) = []
    childrenOf' (_, e1, -1, _) = [e1]
    childrenOf' (_, e1, e2, _) = [e1, e2]

    go [] _    _ = pure ()
    go qs tabu h =
      do childrenOf <- (TrueSet.\\ tabu) . TrueSet.fromList . concat <$> forM qs getChildrenEC -- rerieve all unvisited children
         let childrenL = TrueSet.toList childrenOf
         forM_ childrenL (setMinHeight h) -- set the height of the children as the minimum between current and h
         go childrenL (TrueSet.union tabu childrenOf) (h+1) -- move one breadth search style

-- | calculates the cost of a node
calculateCost :: Monad m => CostFun -> SRTree EClassId -> EGraphST m Cost
calculateCost f t =
  do let cs = childrenOf t
     costs <- traverse (fmap (_cost . _info) . getEClass) cs
     pure . f $ replaceChildren costs t

-- | check whether an e-node evaluates to a const
calculateConsts :: Monad m => SRTree EClassId -> EGraphST m Consts
calculateConsts t =
  do let cs = childrenOf t
     eg <- get
     consts <- traverse (fmap (_consts . _info) . getEClass) cs
     case combineConsts $ replaceChildren consts t of
          ConstVal x | isNaN x -> pure (ConstVal x)
          a -> pure a

combineConsts :: SRTree Consts -> Consts
combineConsts (Const x)    = ConstVal x
combineConsts (Param ix)   = ParamIx ix
combineConsts (Var _)      = NotConst
combineConsts (Uni f t)    = case t of
                              ConstVal x -> ConstVal $ evalFun f x
                              --ParamIx  x -> ParamIx x
                              _          -> t
combineConsts (Bin op l r) = evalOp' l r
  where
    evalOp' (ParamIx ix) (ParamIx iy) = ParamIx (min ix iy)
    evalOp' (ConstVal x) (ConstVal y) = ConstVal $ evalOp op x y
    evalOp' _            _            = NotConst

insertFitness :: Monad m => EClassId -> Double -> [PVector] -> EGraphST m ()
insertFitness eId' fit params = do
  eId <- canonical eId'
  tree <- getBestExpr' eId
  let p = fromIntegral (length params)
  let f_compl = countNodes tree * log (countUniqueTokens tree) + p * (log (2 * pi * exp(1 - log 3)) - log p) / 2.0
  ec <- gets ((IntMap.! eId) . _eClass)
  let oldFit  = _fitness . _info $ ec
  --when (oldFit < Just fit) $ do
  let newInfo = (_info ec){_fitness = Just fit, _theta = params}
      newEc   = ec{_info = newInfo}
      sz = _size newInfo
  modify' $ over eClass (IntMap.insert eId newEc)
  if (isNothing oldFit)
    then modify' $ over (eDB . unevaluated) (IntSet.delete eId)
                 . over (eDB . fitRangeDB) (insertRange eId fit)
                 . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty)
                 . over (eDB . dlRangeDB) (insertRange eId f_compl)
    else modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId (fromJust oldFit))

insertDL :: Monad m => EClassId -> Double -> EGraphST m ()
insertDL eId fit' = do
  let fit = negate fit'
  ec <- gets ((IntMap.! eId) . _eClass)
  let sz = _size . _info $ ec
      newInfo = (_info ec){_dl = Just fit'}
      newEc   = ec{_info=newInfo}
  modify' $ over eClass (IntMap.insert eId newEc)
  modify' $ over (eDB . dlRangeDB) (insertRange eId fit)
          . over (eDB . sizeDLDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty)

-- | TODO: remove from here gets the best expression given the default cost function
getBestExpr' :: Monad m => EClassId -> EGraphST m (Fix SRTree)
getBestExpr' eid = do eid' <- canonical eid
                      best <- gets (_best . _info . (IntMap.! eid') . _eClass)
                      childs <- mapM getBestExpr' $ childrenOf best
                      pure . Fix $ replaceChildren childs best