{-# LANGUAGE DeriveGeneric #-}
module DataSketches.Quantiles.RelativeErrorQuantile.Internal where

import Data.Primitive (MutVar, readMutVar)
import Data.Word
import qualified Data.Vector as Vector
import GHC.Generics
import System.Random.MWC (Gen)

import DataSketches.Core.Snapshot
import DataSketches.Quantiles.RelativeErrorQuantile.Types
import DataSketches.Core.Internal.URef (URef, readURef)
import DataSketches.Quantiles.RelativeErrorQuantile.Internal.Auxiliary (ReqAuxiliary)
import DataSketches.Quantiles.RelativeErrorQuantile.Internal.Compactor (ReqCompactor)
import Control.DeepSeq (NFData, rnf)
import Control.Monad.Primitive (PrimMonad (PrimState))
import qualified DataSketches.Quantiles.RelativeErrorQuantile.Internal.Compactor as Compactor
import qualified DataSketches.Quantiles.RelativeErrorQuantile.Internal.DoubleBuffer as DoubleBuffer
import qualified DataSketches.Quantiles.RelativeErrorQuantile.Internal.Auxiliary as Auxiliary
import Control.Exception (Exception)


{- |
This Relative Error Quantiles Sketch is the Haskell implementation based on the paper
"Relative Error Streaming Quantiles", https://arxiv.org/abs/2004.01668, and loosely derived from
a Python prototype written by Pavel Vesely, ported from the Java equivalent.

This implementation differs from the algorithm described in the paper in the following:

The algorithm requires no upper bound on the stream length.
Instead, each relative-compactor counts the number of compaction operations performed
so far (via variable state). Initially, the relative-compactor starts with INIT_NUMBER_OF_SECTIONS.
Each time the number of compactions (variable state) exceeds 2^{numSections - 1}, we double
numSections. Note that after merging the sketch with another one variable state may not correspond
to the number of compactions performed at a particular level, however, since the state variable
never exceeds the number of compactions, the guarantees of the sketch remain valid.

The size of each section (variable k and sectionSize in the code and parameter k in
the paper) is initialized with a value set by the user via variable k.
When the number of sections doubles, we decrease sectionSize by a factor of sqrt(2).
This is applied at each level separately. Thus, when we double the number of sections, the
nominal compactor size increases by a factor of approx. sqrt(2) (+/- rounding).

The merge operation here does not perform "special compactions", which are used in the paper
to allow for a tight mathematical analysis of the sketch.

This implementation provides a number of capabilities not discussed in the paper or provided
in the Python prototype.

The Python prototype only implemented high accuracy for low ranks. This implementation
provides the user with the ability to choose either high rank accuracy or low rank accuracy at
the time of sketch construction.

- The Python prototype only implemented a comparison criterion of "<". This implementation
allows the user to switch back and forth between the "<=" criterion and the "<=" criterion.
-}
data ReqSketch s = ReqSketch
  { ReqSketch s -> Word32
k :: !Word32
  , ReqSketch s -> RankAccuracy
rankAccuracySetting :: !RankAccuracy
  , ReqSketch s -> Criterion
criterion :: !Criterion
  , ReqSketch s -> Gen s
sketchRng :: {-# UNPACK #-} !(Gen s)
  , ReqSketch s -> URef s Word64
totalN :: {-# UNPACK #-} !(URef s Word64)
  , ReqSketch s -> URef s Double
minValue :: {-# UNPACK #-} !(URef s Double)
  , ReqSketch s -> URef s Double
maxValue :: {-# UNPACK #-} !(URef s Double)
  , ReqSketch s -> URef s Double
sumValue :: {-# UNPACK #-} !(URef s Double)
  , ReqSketch s -> URef s Int
retainedItems :: {-# UNPACK #-} !(URef s Int)
  , ReqSketch s -> URef s Int
maxNominalCapacitiesSize :: {-# UNPACK #-} !(URef s Int)
  , ReqSketch s -> MutVar s (Maybe ReqAuxiliary)
aux :: {-# UNPACK #-} !(MutVar s (Maybe ReqAuxiliary))
  , ReqSketch s -> MutVar s (Vector (ReqCompactor s))
compactors :: {-# UNPACK #-} !(MutVar s (Vector.Vector (ReqCompactor s)))
  } deriving ((forall x. ReqSketch s -> Rep (ReqSketch s) x)
-> (forall x. Rep (ReqSketch s) x -> ReqSketch s)
-> Generic (ReqSketch s)
forall x. Rep (ReqSketch s) x -> ReqSketch s
forall x. ReqSketch s -> Rep (ReqSketch s) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall s x. Rep (ReqSketch s) x -> ReqSketch s
forall s x. ReqSketch s -> Rep (ReqSketch s) x
$cto :: forall s x. Rep (ReqSketch s) x -> ReqSketch s
$cfrom :: forall s x. ReqSketch s -> Rep (ReqSketch s) x
Generic)

instance NFData (ReqSketch s) where
  rnf :: ReqSketch s -> ()
rnf !ReqSketch s
rs = ()

data ReqSketchSnapshot = ReqSketchSnapshot
    { ReqSketchSnapshot -> RankAccuracy
snapshotRankAccuracySetting :: !RankAccuracy
    , ReqSketchSnapshot -> Criterion
snapshotCriterion :: !Criterion
    , ReqSketchSnapshot -> Word64
snapshotTotalN :: !Word64
    , ReqSketchSnapshot -> Double
snapshotMinValue :: !Double
    , ReqSketchSnapshot -> Double
snapshotMaxValue :: !Double
    , ReqSketchSnapshot -> Int
snapshotRetainedItems :: !Int
    , ReqSketchSnapshot -> Int
snapshotMaxNominalCapacitiesSize :: !Int
    -- , aux :: !(MutVar s (Maybe ()))
    , ReqSketchSnapshot -> Vector (Snapshot ReqCompactor)
snapshotCompactors :: !(Vector.Vector (Snapshot ReqCompactor))
    } deriving Int -> ReqSketchSnapshot -> ShowS
[ReqSketchSnapshot] -> ShowS
ReqSketchSnapshot -> String
(Int -> ReqSketchSnapshot -> ShowS)
-> (ReqSketchSnapshot -> String)
-> ([ReqSketchSnapshot] -> ShowS)
-> Show ReqSketchSnapshot
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ReqSketchSnapshot] -> ShowS
$cshowList :: [ReqSketchSnapshot] -> ShowS
show :: ReqSketchSnapshot -> String
$cshow :: ReqSketchSnapshot -> String
showsPrec :: Int -> ReqSketchSnapshot -> ShowS
$cshowsPrec :: Int -> ReqSketchSnapshot -> ShowS
Show

instance TakeSnapshot ReqSketch where
  type Snapshot ReqSketch = ReqSketchSnapshot 
  takeSnapshot :: ReqSketch (PrimState m) -> m (Snapshot ReqSketch)
takeSnapshot ReqSketch{Word32
Gen (PrimState m)
MutVar (PrimState m) (Maybe ReqAuxiliary)
MutVar (PrimState m) (Vector (ReqCompactor (PrimState m)))
URef (PrimState m) Double
URef (PrimState m) Int
URef (PrimState m) Word64
RankAccuracy
Criterion
compactors :: MutVar (PrimState m) (Vector (ReqCompactor (PrimState m)))
aux :: MutVar (PrimState m) (Maybe ReqAuxiliary)
maxNominalCapacitiesSize :: URef (PrimState m) Int
retainedItems :: URef (PrimState m) Int
sumValue :: URef (PrimState m) Double
maxValue :: URef (PrimState m) Double
minValue :: URef (PrimState m) Double
totalN :: URef (PrimState m) Word64
sketchRng :: Gen (PrimState m)
criterion :: Criterion
rankAccuracySetting :: RankAccuracy
k :: Word32
compactors :: forall s. ReqSketch s -> MutVar s (Vector (ReqCompactor s))
aux :: forall s. ReqSketch s -> MutVar s (Maybe ReqAuxiliary)
maxNominalCapacitiesSize :: forall s. ReqSketch s -> URef s Int
retainedItems :: forall s. ReqSketch s -> URef s Int
sumValue :: forall s. ReqSketch s -> URef s Double
maxValue :: forall s. ReqSketch s -> URef s Double
minValue :: forall s. ReqSketch s -> URef s Double
totalN :: forall s. ReqSketch s -> URef s Word64
sketchRng :: forall s. ReqSketch s -> Gen s
criterion :: forall s. ReqSketch s -> Criterion
rankAccuracySetting :: forall s. ReqSketch s -> RankAccuracy
k :: forall s. ReqSketch s -> Word32
..} = RankAccuracy
-> Criterion
-> Word64
-> Double
-> Double
-> Int
-> Int
-> Vector (Snapshot ReqCompactor)
-> ReqSketchSnapshot
ReqSketchSnapshot RankAccuracy
rankAccuracySetting Criterion
criterion
    (Word64
 -> Double
 -> Double
 -> Int
 -> Int
 -> Vector ReqCompactorSnapshot
 -> ReqSketchSnapshot)
-> m Word64
-> m (Double
      -> Double
      -> Int
      -> Int
      -> Vector ReqCompactorSnapshot
      -> ReqSketchSnapshot)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> URef (PrimState m) Word64 -> m Word64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> m a
readURef URef (PrimState m) Word64
totalN
    m (Double
   -> Double
   -> Int
   -> Int
   -> Vector ReqCompactorSnapshot
   -> ReqSketchSnapshot)
-> m Double
-> m (Double
      -> Int -> Int -> Vector ReqCompactorSnapshot -> ReqSketchSnapshot)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> URef (PrimState m) Double -> m Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> m a
readURef URef (PrimState m) Double
minValue
    m (Double
   -> Int -> Int -> Vector ReqCompactorSnapshot -> ReqSketchSnapshot)
-> m Double
-> m (Int
      -> Int -> Vector ReqCompactorSnapshot -> ReqSketchSnapshot)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> URef (PrimState m) Double -> m Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> m a
readURef URef (PrimState m) Double
maxValue
    m (Int -> Int -> Vector ReqCompactorSnapshot -> ReqSketchSnapshot)
-> m Int
-> m (Int -> Vector ReqCompactorSnapshot -> ReqSketchSnapshot)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> URef (PrimState m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> m a
readURef URef (PrimState m) Int
retainedItems
    m (Int -> Vector ReqCompactorSnapshot -> ReqSketchSnapshot)
-> m Int -> m (Vector ReqCompactorSnapshot -> ReqSketchSnapshot)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> URef (PrimState m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> m a
readURef URef (PrimState m) Int
maxNominalCapacitiesSize
    m (Vector ReqCompactorSnapshot -> ReqSketchSnapshot)
-> m (Vector ReqCompactorSnapshot) -> m ReqSketchSnapshot
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (MutVar (PrimState m) (Vector (ReqCompactor (PrimState m)))
-> m (Vector (ReqCompactor (PrimState m)))
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (Vector (ReqCompactor (PrimState m)))
compactors m (Vector (ReqCompactor (PrimState m)))
-> (Vector (ReqCompactor (PrimState m))
    -> m (Vector ReqCompactorSnapshot))
-> m (Vector ReqCompactorSnapshot)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ReqCompactor (PrimState m) -> m ReqCompactorSnapshot)
-> Vector (ReqCompactor (PrimState m))
-> m (Vector ReqCompactorSnapshot)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ReqCompactor (PrimState m) -> m ReqCompactorSnapshot
forall (a :: * -> *) (m :: * -> *).
(TakeSnapshot a, PrimMonad m) =>
a (PrimState m) -> m (Snapshot a)
takeSnapshot)

getCompactors :: PrimMonad m => ReqSketch (PrimState m) -> m (Vector.Vector (ReqCompactor (PrimState m)))
getCompactors :: ReqSketch (PrimState m) -> m (Vector (ReqCompactor (PrimState m)))
getCompactors = MutVar (PrimState m) (Vector (ReqCompactor (PrimState m)))
-> m (Vector (ReqCompactor (PrimState m)))
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar (MutVar (PrimState m) (Vector (ReqCompactor (PrimState m)))
 -> m (Vector (ReqCompactor (PrimState m))))
-> (ReqSketch (PrimState m)
    -> MutVar (PrimState m) (Vector (ReqCompactor (PrimState m))))
-> ReqSketch (PrimState m)
-> m (Vector (ReqCompactor (PrimState m)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ReqSketch (PrimState m)
-> MutVar (PrimState m) (Vector (ReqCompactor (PrimState m)))
forall s. ReqSketch s -> MutVar s (Vector (ReqCompactor s))
compactors

computeTotalRetainedItems :: PrimMonad m => ReqSketch (PrimState m) -> m Int
computeTotalRetainedItems :: ReqSketch (PrimState m) -> m Int
computeTotalRetainedItems ReqSketch (PrimState m)
this = do
  Vector (ReqCompactor (PrimState m))
compactors <- ReqSketch (PrimState m) -> m (Vector (ReqCompactor (PrimState m)))
forall (m :: * -> *).
PrimMonad m =>
ReqSketch (PrimState m) -> m (Vector (ReqCompactor (PrimState m)))
getCompactors ReqSketch (PrimState m)
this
  (Int -> ReqCompactor (PrimState m) -> m Int)
-> Int -> Vector (ReqCompactor (PrimState m)) -> m Int
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Vector b -> m a
Vector.foldM Int -> ReqCompactor (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
Int -> ReqCompactor (PrimState m) -> m Int
countBuffer Int
0 Vector (ReqCompactor (PrimState m))
compactors
  where
    countBuffer :: Int -> ReqCompactor (PrimState m) -> m Int
countBuffer Int
acc ReqCompactor (PrimState m)
compactor = do
      DoubleBuffer (PrimState m)
buff <- ReqCompactor (PrimState m) -> m (DoubleBuffer (PrimState m))
forall (m :: * -> *).
PrimMonad m =>
ReqCompactor (PrimState m) -> m (DoubleBuffer (PrimState m))
Compactor.getBuffer ReqCompactor (PrimState m)
compactor
      Int
buffSize <- DoubleBuffer (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
DoubleBuffer (PrimState m) -> m Int
DoubleBuffer.getCount DoubleBuffer (PrimState m)
buff
      Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
buffSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
acc

retainedItemCount :: PrimMonad m => ReqSketch (PrimState m) -> m Int
retainedItemCount :: ReqSketch (PrimState m) -> m Int
retainedItemCount = URef (PrimState m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> m a
readURef (URef (PrimState m) Int -> m Int)
-> (ReqSketch (PrimState m) -> URef (PrimState m) Int)
-> ReqSketch (PrimState m)
-> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ReqSketch (PrimState m) -> URef (PrimState m) Int
forall s. ReqSketch s -> URef s Int
retainedItems

-- | Get the total number of items inserted into the sketch
count :: PrimMonad m => ReqSketch (PrimState m) -> m Word64
count :: ReqSketch (PrimState m) -> m Word64
count = URef (PrimState m) Word64 -> m Word64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> m a
readURef (URef (PrimState m) Word64 -> m Word64)
-> (ReqSketch (PrimState m) -> URef (PrimState m) Word64)
-> ReqSketch (PrimState m)
-> m Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ReqSketch (PrimState m) -> URef (PrimState m) Word64
forall s. ReqSketch s -> URef s Word64
totalN

mkAuxiliaryFromReqSketch :: PrimMonad m => ReqSketch (PrimState m) -> m ReqAuxiliary
mkAuxiliaryFromReqSketch :: ReqSketch (PrimState m) -> m ReqAuxiliary
mkAuxiliaryFromReqSketch ReqSketch (PrimState m)
this = do
  Word64
total <- ReqSketch (PrimState m) -> m Word64
forall (m :: * -> *).
PrimMonad m =>
ReqSketch (PrimState m) -> m Word64
count ReqSketch (PrimState m)
this
  Int
retainedItems <- ReqSketch (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
ReqSketch (PrimState m) -> m Int
retainedItemCount ReqSketch (PrimState m)
this
  Vector (ReqCompactor (PrimState m))
compactors <- ReqSketch (PrimState m) -> m (Vector (ReqCompactor (PrimState m)))
forall (m :: * -> *).
PrimMonad m =>
ReqSketch (PrimState m) -> m (Vector (ReqCompactor (PrimState m)))
getCompactors ReqSketch (PrimState m)
this
  RankAccuracy
-> Word64
-> Int
-> Vector (ReqCompactor (PrimState m))
-> m ReqAuxiliary
forall (m :: * -> *) s.
(PrimMonad m, s ~ PrimState m) =>
RankAccuracy
-> Word64 -> Int -> Vector (ReqCompactor s) -> m ReqAuxiliary
Auxiliary.mkAuxiliary (ReqSketch (PrimState m) -> RankAccuracy
forall s. ReqSketch s -> RankAccuracy
rankAccuracySetting ReqSketch (PrimState m)
this) Word64
total Int
retainedItems Vector (ReqCompactor (PrimState m))
compactors

data CumulativeDistributionInvariants
  = CumulativeDistributionInvariantsSplitsAreEmpty
  | CumulativeDistributionInvariantsSplitsAreNotFinite
  | CumulativeDistributionInvariantsSplitsAreNotUniqueAndMontonicallyIncreasing
  deriving (Int -> CumulativeDistributionInvariants -> ShowS
[CumulativeDistributionInvariants] -> ShowS
CumulativeDistributionInvariants -> String
(Int -> CumulativeDistributionInvariants -> ShowS)
-> (CumulativeDistributionInvariants -> String)
-> ([CumulativeDistributionInvariants] -> ShowS)
-> Show CumulativeDistributionInvariants
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CumulativeDistributionInvariants] -> ShowS
$cshowList :: [CumulativeDistributionInvariants] -> ShowS
show :: CumulativeDistributionInvariants -> String
$cshow :: CumulativeDistributionInvariants -> String
showsPrec :: Int -> CumulativeDistributionInvariants -> ShowS
$cshowsPrec :: Int -> CumulativeDistributionInvariants -> ShowS
Show, CumulativeDistributionInvariants
-> CumulativeDistributionInvariants -> Bool
(CumulativeDistributionInvariants
 -> CumulativeDistributionInvariants -> Bool)
-> (CumulativeDistributionInvariants
    -> CumulativeDistributionInvariants -> Bool)
-> Eq CumulativeDistributionInvariants
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CumulativeDistributionInvariants
-> CumulativeDistributionInvariants -> Bool
$c/= :: CumulativeDistributionInvariants
-> CumulativeDistributionInvariants -> Bool
== :: CumulativeDistributionInvariants
-> CumulativeDistributionInvariants -> Bool
$c== :: CumulativeDistributionInvariants
-> CumulativeDistributionInvariants -> Bool
Eq)

instance Exception CumulativeDistributionInvariants