module Utils.ARC
  ( ARCTick
  , ARC, initARC
  , arcAlterF, lookupARC, insertARC
  , ARCHandle, initARCHandle, cachedARC, cachedARC'
  , lookupARCHandle
  , readARCHandle
  , arcRecentSize, arcFrequentSize, arcGhostRecentSize, arcGhostFrequentSize
  , getARCRecentWeight, getARCFrequentWeight
  , describeARC
  ) where

import ClassyPrelude

import Data.OrdPSQ (OrdPSQ)
import qualified Data.OrdPSQ as OrdPSQ

import Control.Lens

-- https://web.archive.org/web/20210115184012/https://dbs.uni-leipzig.de/file/ARC.pdf
-- https://jaspervdj.be/posts/2015-02-24-lru-cache.html


newtype ARCTick = ARCTick { _getARCTick :: Word64 }
  deriving (Eq, Ord, Show, Typeable)
  deriving newtype (NFData)

makeLenses ''ARCTick

data ARC k w v = ARC
  { arcRecent, arcFrequent :: !(OrdPSQ k ARCTick (v, w))
  , arcGhostRecent, arcGhostFrequent :: !(OrdPSQ k ARCTick ())
  , arcRecentWeight, arcFrequentWeight :: !w
  , arcTargetRecent, arcMaximumWeight :: !w
  , arcMaximumGhost :: !Int
  }

instance (NFData k, NFData w, NFData v) => NFData (ARC k w v) where
  rnf ARC{..} = rnf arcRecent
          `seq` rnf arcFrequent
          `seq` rnf arcGhostRecent
          `seq` rnf arcGhostFrequent
          `seq` rnf arcRecentWeight
          `seq` rnf arcFrequentWeight
          `seq` rnf arcTargetRecent
          `seq` rnf arcMaximumWeight
          `seq` rnf arcMaximumGhost

describeARC :: Show w
            => ARC k w v
            -> String
describeARC ARC{..} = intercalate ", "
  [ "arcRecent: "         <> show (OrdPSQ.size arcRecent)
  , "arcFrequent: "       <> show (OrdPSQ.size arcFrequent)
  , "arcGhostRecent: "    <> show (OrdPSQ.size arcGhostRecent)
  , "arcGhostFrequent: "  <> show (OrdPSQ.size arcGhostFrequent)
  , "arcRecentWeight: "   <> show arcRecentWeight
  , "arcFrequentWeight: " <> show arcFrequentWeight
  , "arcTargetRecent: "   <> show arcTargetRecent
  , "arcMaximumWeight: "  <> show arcMaximumWeight
  , "arcMaximumGhost: "   <> show arcMaximumGhost
  ]

arcRecentSize, arcFrequentSize, arcGhostRecentSize, arcGhostFrequentSize :: ARC k w v -> Int
arcRecentSize = OrdPSQ.size . arcRecent
arcFrequentSize = OrdPSQ.size . arcFrequent
arcGhostRecentSize = OrdPSQ.size . arcGhostRecent
arcGhostFrequentSize = OrdPSQ.size . arcGhostFrequent

getARCRecentWeight, getARCFrequentWeight :: ARC k w v -> w
getARCRecentWeight = arcRecentWeight
getARCFrequentWeight = arcFrequentWeight

initialARCTick :: ARCTick
initialARCTick = ARCTick 0

initARC :: forall k w v.
           Integral w
        => Int -- ^ @arcMaximumGhost@
        -> w -- ^ @arcMaximumWeight@
        -> (ARC k w v, ARCTick)
initARC arcMaximumGhost arcMaximumWeight
  | arcMaximumWeight < 0 = error "initARC given negative maximum weight"
  | arcMaximumGhost < 0 = error "initARC given negative maximum ghost size"
  | otherwise = (, initialARCTick) ARC
    { arcRecent = OrdPSQ.empty
    , arcFrequent = OrdPSQ.empty
    , arcGhostRecent = OrdPSQ.empty
    , arcGhostFrequent = OrdPSQ.empty
    , arcRecentWeight = 0
    , arcFrequentWeight = 0
    , arcMaximumWeight
    , arcTargetRecent = 0
    , arcMaximumGhost
    }


infixl 6 |-
(|-) :: (Num a, Ord a) => a -> a -> a
(|-) m s
  | s >= m = 0
  | otherwise = m - s

  
arcAlterF :: forall f k w v.
             ( Ord k
             , Functor f
             , Integral w
             )
          => k
          -> (Maybe (v, w) -> f (Maybe (v, w)))
          -> ARC k w v
          -> ARCTick -> f (ARC k w v, ARCTick)
-- | Unchecked precondition: item weights are always less than `arcMaximumWeight`
arcAlterF k f oldARC@ARC{..} now
  | later <= initialARCTick = uncurry (arcAlterF k f) $ initARC arcMaximumGhost arcMaximumWeight
  | otherwise = (, later) <$> if
    | Just (_p, x@(_, w), arcFrequent') <- OrdPSQ.deleteView k arcFrequent
      -> f (Just x) <&> \(fromMaybe x -> x'@(_, w'))
           -> let (arcFrequent'', arcFrequentWeight'', arcGhostFrequent') = evictToSize (arcMaximumWeight |- arcTargetRecent |- w') arcFrequent' (arcFrequentWeight - w) arcGhostFrequent
               in oldARC
                    { arcFrequent = OrdPSQ.insert k now x' arcFrequent''
                    , arcFrequentWeight = arcFrequentWeight'' + w'
                    , arcGhostFrequent = arcGhostFrequent'
                    }
    | Just (_p, x@(_, w), arcRecent') <- OrdPSQ.deleteView k arcRecent
      -> f (Just x) <&> \(fromMaybe x -> x'@(_, w'))
           -> let (arcFrequent', arcFrequentWeight', arcGhostFrequent') = evictToSize (arcMaximumWeight |- arcTargetRecent |- w') arcFrequent arcFrequentWeight arcGhostFrequent
               in oldARC
                    { arcRecent = arcRecent'
                    , arcRecentWeight = arcRecentWeight - w
                    , arcFrequent = OrdPSQ.insert k now x' arcFrequent'
                    , arcFrequentWeight = arcFrequentWeight' + w'
                    , arcGhostFrequent = arcGhostFrequent'
                    }
    | Just (_p, (), arcGhostRecent') <- OrdPSQ.deleteView k arcGhostRecent
      -> f Nothing <&> \case
           Nothing -> oldARC
             { arcGhostRecent = OrdPSQ.insert k now () arcGhostRecent'
             }
           Just x@(_, w)
             -> let arcTargetRecent' = min arcMaximumWeight $ arcTargetRecent + max avgWeight (round $ toRational (OrdPSQ.size arcGhostFrequent) / toRational (OrdPSQ.size arcGhostRecent) * toRational avgWeight)
                    (arcFrequent', arcFrequentWeight', arcGhostFrequent') = evictToSize (arcMaximumWeight |- arcTargetRecent' |- w) arcFrequent arcFrequentWeight arcGhostFrequent
                    (arcRecent', arcRecentWeight', arcGhostRecent'') = evictToSize (max arcTargetRecent' $ arcMaximumWeight |- arcFrequentWeight' |- w) arcRecent arcRecentWeight arcGhostRecent'
                 in oldARC
                      { arcRecent = arcRecent'
                      , arcFrequent = OrdPSQ.insert k now x arcFrequent'
                      , arcGhostRecent = arcGhostRecent''
                      , arcGhostFrequent = arcGhostFrequent'
                      , arcRecentWeight = arcRecentWeight'
                      , arcFrequentWeight = arcFrequentWeight' + w
                      , arcTargetRecent = arcTargetRecent'
                      }
    | Just (_p, (), arcGhostFrequent') <- OrdPSQ.deleteView k arcGhostFrequent
      -> f Nothing <&> \case
           Nothing -> oldARC
             { arcGhostFrequent = OrdPSQ.insert k now () arcGhostFrequent'
             }
           Just x@(_, w)
             -> let arcTargetRecent' = arcTargetRecent |- max avgWeight (round $ toRational (OrdPSQ.size arcGhostRecent) / toRational (OrdPSQ.size arcGhostFrequent) * toRational avgWeight)
                    (arcFrequent', arcFrequentWeight', arcGhostFrequent'') = evictToSize (arcMaximumWeight |- arcTargetRecent' |- w) arcFrequent arcFrequentWeight arcGhostFrequent'
                    (arcRecent', arcRecentWeight', arcGhostRecent') = evictToSize (max arcTargetRecent' $ arcMaximumWeight |- arcFrequentWeight' |- w) arcRecent arcRecentWeight arcGhostRecent
                 in oldARC
                      { arcRecent = arcRecent'
                      , arcFrequent = OrdPSQ.insert k now x arcFrequent'
                      , arcGhostRecent = arcGhostRecent'
                      , arcGhostFrequent = arcGhostFrequent''
                      , arcRecentWeight = arcRecentWeight'
                      , arcFrequentWeight = arcFrequentWeight' + w
                      , arcTargetRecent = arcTargetRecent'
                      }
    | otherwise -> f Nothing <&> \case
        Nothing -> oldARC
          { arcGhostRecent = OrdPSQ.insert k now () $ evictGhostToCount arcGhostRecent
          }
        Just x@(_, w)
          -> let (arcRecent', arcRecentWeight', arcGhostRecent') = evictToSize (max arcTargetRecent (arcMaximumWeight |- arcFrequentWeight) |- w) arcRecent arcRecentWeight arcGhostRecent
              in oldARC
                   { arcRecent = OrdPSQ.insert k now x arcRecent'
                   , arcRecentWeight = arcRecentWeight' + w
                   , arcGhostRecent = arcGhostRecent'
                   }
  where
    avgWeight = round $ toRational (arcRecentWeight + arcFrequentWeight) / toRational (OrdPSQ.size arcFrequent + OrdPSQ.size arcRecent)
    
    later :: ARCTick
    later = over getARCTick succ now
    
    evictToSize :: w -> OrdPSQ k ARCTick (v, w) -> w -> OrdPSQ k ARCTick () -> (OrdPSQ k ARCTick (v, w), w, OrdPSQ k ARCTick ())
    evictToSize tSize c cSize ghostC
      | cSize <= tSize = (c, cSize, ghostC)
      | Just (k', p', (_, w'), c') <- OrdPSQ.minView c = evictToSize tSize c' (cSize - w') . evictGhostToCount $ OrdPSQ.insert k' p' () ghostC
      | otherwise = error "evictToSize: cannot reach required size through eviction"

    evictGhostToCount :: OrdPSQ k ARCTick () -> OrdPSQ k ARCTick ()
    evictGhostToCount c
      | OrdPSQ.size c <= arcMaximumGhost = c
      | Just (_, _, _, c') <- OrdPSQ.minView c = evictGhostToCount c'
      | otherwise = error "evictGhostToCount: cannot reach required count through eviction"

lookupARC :: forall k w v.
             ( Ord k
             , Integral w
             )
          => k
          -> (ARC k w v, ARCTick)
          -> Maybe (v, w)
lookupARC k = getConst . uncurry (arcAlterF k Const)

insertARC :: forall k w v.
             ( Ord k
             , Integral w
             )
          => k
          -> Maybe (v, w)
          -> ARC k w v
          -> ARCTick -> (ARC k w v, ARCTick)
insertARC k newVal = (runIdentity .) . arcAlterF k (const $ pure newVal)


newtype ARCHandle k w v = ARCHandle { _getARCHandle :: IORef (ARC k w v, ARCTick) }
  deriving (Eq, Typeable)

initARCHandle :: forall k w v m.
                 ( MonadIO m
                 , Integral w
                 )
              => Int -- ^ @arcMaximumGhost@
              -> w -- ^ @arcMaximumWeight@
              -> m (ARCHandle k w v)
initARCHandle maxGhost maxWeight = fmap ARCHandle . newIORef $ initARC maxGhost maxWeight

cachedARC' :: forall k w v m.
              ( MonadIO m
              , Ord k
              , Integral w
              , NFData k, NFData w, NFData v
              )
           => ARCHandle k w v
           -> k
           -> (Maybe (v, w) -> m (Maybe (v, w)))
           -> m (Maybe v)
cachedARC' (ARCHandle arcVar) k f = do
  oldVal <- lookupARC k <$> readIORef arcVar
  newVal <- f oldVal
  modifyIORef' arcVar $ force . uncurry (insertARC k newVal)
  -- Using `modifyIORef'` instead of `atomicModifyIORef'` might very
  -- well drop newer values computed during the update.
  --
  -- Currently we accept that to reduce lock contention.
  --
  -- Another alternative would be to use "optimistic locking",
  -- i.e. read the current value of `arcVar`, compute an updated
  -- version, and write it back atomically iff the `ARCTick` hasn't
  -- changed.
  --
  -- This was not implemented to avoid the overhead and contention
  -- likely associated with the atomic transaction required for the
  -- "compare and swap"
  return $ view _1 <$> newVal

cachedARC :: forall k w v m.
             ( MonadIO m
             , Ord k
             , Integral w
             , NFData k, NFData w, NFData v
             )
          => ARCHandle k w v
          -> k
          -> (Maybe (v, w) -> m (v, w))
          -> m v
cachedARC h k f = fromMaybe (error "cachedARC: cachedARC' returned Nothing") <$> cachedARC' h k (fmap Just . f)
  
lookupARCHandle :: forall k w v m.
                   ( MonadIO m
                   , Ord k
                   , Integral w
                   )
                => ARCHandle k w v
                -> k
                -> m (Maybe (v, w))
lookupARCHandle (ARCHandle arcVar) k = lookupARC k <$> readIORef arcVar


readARCHandle :: MonadIO m
              => ARCHandle k w v
              -> m (ARC k w v, ARCTick)
readARCHandle (ARCHandle arcVar) = readIORef arcVar