packages feed

dunning-t-digest-0.1.0.0: src/Data/Sketch/TDigest.hs

-- |
-- Module      : Data.Sketch.TDigest
-- Description : Dunning t-digest for online quantile estimation
-- Copyright   : (c) Nadia Yvette Chambers, 2025
-- License     : BSD-3-Clause
-- Maintainer  : nadia.yvette.chambers@gmail.com
-- Stability   : experimental
--
-- A pure functional implementation of the Dunning t-digest data structure,
-- using the merging digest variant with the \(K_1\) (arcsine) scale function.
-- The t-digest provides streaming, mergeable, memory-bounded approximation
-- of quantile (percentile) queries with high accuracy in the tails.
--
-- == Background
--
-- The /streaming quantile problem/ asks: given a (possibly unbounded) stream
-- of real-valued observations, answer queries of the form "what is the value
-- at the \(q\)-th quantile?" using bounded memory.
-- Munro & Paterson (1980) established that exact selection from a stream of
-- \(n\) elements requires \(\Omega(n)\) space in the comparison model
-- (<https://doi.org/10.1016/0304-3975(80)90061-4>), so any sub-linear space
-- algorithm must accept approximation.  Greenwald & Khanna (2001) gave the
-- first \(\varepsilon\)-approximate streaming quantile summary with space
-- \(O\!\bigl(\frac{1}{\varepsilon}\log(\varepsilon n)\bigr)\)
-- (<https://doi.org/10.1145/375663.375670>), guaranteeing uniform error across
-- all quantiles.  The t-digest takes a different approach: it trades uniform
-- guarantees for much higher accuracy in the extreme tails (\(q \approx 0\) or
-- \(q \approx 1\)), which is the regime most relevant to SLA monitoring,
-- anomaly detection, and financial risk measurement.
--
-- == The t-digest
--
-- The t-digest, introduced by Ted Dunning
-- (<https://doi.org/10.1016/j.simpa.2020.100049>; see also Dunning & Ertl,
-- <https://arxiv.org/abs/1902.04023>), represents an empirical distribution as
-- an ordered sequence of /centroids/ \((m_i, w_i)\), where \(m_i\) is a
-- weighted mean and \(w_i\) is a count of observations.  Centroids are kept
-- sorted by mean.  The key idea is to use a /scale function/ \(k(q, \delta)\)
-- that maps the quantile axis \([0, 1]\) to a "scale space" in which uniform
-- spacing corresponds to the desired non-uniform resolution in quantile space.
--
-- This module implements the /merging digest/ variant with the \(K_1\)
-- (arcsine) scale function:
--
-- \[
--   k(q, \delta) \;=\; \frac{\delta}{2\pi}\,\arcsin(2q - 1)
-- \]
--
-- The \(K_1\) function has infinite derivative at \(q = 0\) and \(q = 1\),
-- meaning it allocates proportionally more centroids near the tails.  Its
-- inverse is:
--
-- \[
--   q(k, \delta) \;=\; \frac{1 + \sin\!\bigl(\frac{2\pi k}{\delta}\bigr)}{2}
-- \]
--
-- A new observation may be merged into an existing centroid \(i\) only if the
-- resulting centroid would satisfy the /size constraint/:
--
-- \[
--   k\!\bigl(q_{\mathrm{upper}},\, \delta\bigr) \;-\; k\!\bigl(q_{\mathrm{lower}},\, \delta\bigr) \;\le\; 1
-- \]
--
-- where \(q_{\mathrm{lower}}\) and \(q_{\mathrm{upper}}\) are the quantile
-- boundaries of the (proposed) merged centroid.  This constraint ensures that
-- centroids near \(q = 0\) and \(q = 1\) remain small (even singletons),
-- while centroids near the median may absorb many observations.
--
-- == Space bounds
--
-- The number of centroids in a t-digest is bounded by \(O(\delta)\)
-- /regardless/ of the number of observations \(n\).  Specifically, the integer
-- range of the scale function is
-- \(\lceil k(0,\delta)\rceil \ldots \lfloor k(1,\delta)\rfloor =
-- \lceil -\delta/2\rceil \ldots \lfloor \delta/2\rfloor\),
-- giving at most \(\delta + 1\) unit intervals and therefore at most
-- \(\delta + 1\) centroids after compression.  In practice the compression
-- threshold is set to \(3\delta\) centroids (before triggering a compress
-- pass), so the working-set size is at most \(3\delta\) centroids.  With the
-- default \(\delta = 100\), this means at most 300 centroids regardless of
-- whether the stream contains \(10^3\) or \(10^{12}\) observations.
--
-- == Implementation: finger trees with a four-component measure
--
-- This module stores centroids in a @'Data.FingerTree.FingerTree'@ from the
-- @fingertree@ package, as described by Hinze & Paterson (2006)
-- (<https://doi.org/10.1017/S0956796805005769>).  Finger trees support
-- amortised \(O(\log n)\) split and concatenation, and \(O(1)\) access to
-- extremal elements, making them well suited for the sorted-centroid
-- representation.
--
-- The monoidal measure carried by the tree has four components:
--
-- 1. @mWeight@ \(= \sum w_i\): cumulative weight, enabling split-by-weight
--    for quantile queries.
-- 2. @mCount@ \(= |\{i\}|\): centroid count, enabling \(O(1)\)
--    'centroidCount'.
-- 3. @mMaxMean@ \(= \max\{m_i\}\): maximum mean over the subtree, enabling
--    split-by-mean for insertion and CDF queries.  Because centroids are
--    stored in sorted order, @mMaxMean@ is monotone over prefixes.
-- 4. @mMeanWeightSum@ \(= \sum m_i w_i\): the sum of products of mean and
--    weight.  This enables \(O(1)\) computation of the merged mean of any
--    contiguous chunk: \(\bar{m} = \texttt{mMeanWeightSum} /
--    \texttt{mWeight}\).  This is the key to achieving \(O(\delta \log n)\)
--    compression: each of the \(O(\delta)\) chunks produced by splitting at
--    scale-function unit boundaries can be collapsed into a single centroid
--    without traversing its elements.
--
-- == Companion implementations: array-backed 2-3-4 trees
--
-- Twenty-two mutable implementations in this project (in C, C++, Rust, Go,
-- Zig, Java, C#, and others) use array-backed 2-3-4 trees instead of
-- finger trees.  The 2-3-4 tree is a B-tree of order 4 (Bayer & McCreight,
-- 1972; <https://doi.org/10.1007/BF00288683>), isomorphic to a red-black tree
-- via the correspondence established by Guibas & Sedgewick (1978)
-- (<https://doi.org/10.1109/SFCS.1978.3>; see also Sedgewick, 2008,
-- <https://sedgewick.io/wp-content/themes/flavor/papers/2008LLRB.pdf>),
-- provides worst-case \(O(\log n)\) insertion, deletion, and search with
-- excellent cache locality when nodes are packed into a flat array.  This is
-- particularly important for robustness at very fine-grained quantile queries
-- (e.g., \(q = 0.9999\)) where the tail centroids that determine accuracy
-- must be located quickly and updated with minimal overhead.  The array-backed
-- layout avoids pointer-chasing and improves branch-prediction behaviour,
-- yielding 2--5\(\times\) speedups in practice over pointer-based trees.
--
-- == Quick start
--
-- @
-- import Data.Sketch.TDigest
-- import Data.List ('Data.List.foldl\'')
--
-- main :: IO ()
-- main = do
--   let td = 'Data.List.foldl\'' (flip 'add') 'empty' [1.0 .. 10000.0]
--   print ('quantile' 0.99 td)   -- Just ~9900.5
--   print ('cdf' 5000.0 td)      -- Just ~0.5
-- @
module Data.Sketch.TDigest
  ( -- * Types
    TDigest,
    Centroid (..),

    -- * Construction
    empty,
    emptyWith,

    -- * Insertion
    add,
    addWeighted,

    -- * Compression
    compress,

    -- * Queries
    quantile,
    cdf,

    -- * Merging
    merge,

    -- * Accessors
    totalWeight,
    centroidCount,
    centroidList,
    getDelta,
    getMin,
    getMax,

    -- * Reconstruction
    fromComponents,
  )
where

import Data.FingerTree (FingerTree, Measured (..), ViewL (..), ViewR (..), (<|), (|>))
import qualified Data.FingerTree as FT

-- ---------------------------------------------------------------------------
-- Measure (monoidal annotation for the finger tree)
-- ---------------------------------------------------------------------------

-- | Monoidal measure carried by every internal node of the finger tree.
--
-- Following Hinze & Paterson (2006)
-- (<https://doi.org/10.1017/S0956796805005769>), a finger tree is
-- parameterised by a monoid whose cached values enable efficient splitting.
-- The t-digest requires /four/ independent capabilities from the tree, so the
-- measure is a four-component product monoid:
--
-- * @mWeight@ — cumulative weight \(\sum w_i\).  Used by 'quantile' to
--   split the tree at a target cumulative weight in \(O(\log n)\).
--
-- * @mCount@ — number of centroids \(|\{i\}|\).  Provides \(O(1)\)
--   'centroidCount' and is used during quantile interpolation to detect
--   boundary centroids.
--
-- * @mMaxMean@ — maximum centroid mean \(\max\{m_i\}\) in the subtree.
--   Because centroids are sorted by mean, this value is monotone over
--   prefixes, enabling 'FT.split' by mean value for insertion ('addWeighted')
--   and CDF queries ('cdf').
--
-- * @mMeanWeightSum@ — the sum \(\sum m_i w_i\).  Combined with @mWeight@,
--   this allows the weighted mean of any contiguous subtree to be computed in
--   \(O(1)\): \(\bar{m} = \texttt{mMeanWeightSum}\,/\,\texttt{mWeight}\).
--   This is the critical component that makes 'compress' run in
--   \(O(\delta \log n)\) rather than \(O(n)\): each chunk produced by
--   splitting at \(K_1\) unit boundaries is collapsed into a single centroid
--   without iterating over its elements.
data Measure = Measure
  { mWeight :: {-# UNPACK #-} !Double,
    mCount :: {-# UNPACK #-} !Int,
    mMaxMean :: {-# UNPACK #-} !Double,
    mMeanWeightSum :: {-# UNPACK #-} !Double
  }
  deriving (Show)

instance Semigroup Measure where
  (Measure w1 c1 mm1 mws1) <> (Measure w2 c2 mm2 mws2) =
    Measure (w1 + w2) (c1 + c2) (max mm1 mm2) (mws1 + mws2)

instance Monoid Measure where
  mempty = Measure 0 0 (-(1 / 0)) 0

-- ---------------------------------------------------------------------------
-- Types
-- ---------------------------------------------------------------------------

-- | A single centroid in the t-digest, representing a cluster of nearby
-- values by their weighted mean and total weight.
--
-- In the t-digest framework (Dunning, 2021;
-- <https://doi.org/10.1016/j.simpa.2020.100049>), the empirical distribution
-- is approximated by an ordered sequence of centroids \((m_i, w_i)\).  When a
-- new observation \(x\) with weight \(w\) is merged into an existing centroid
-- \((m_i, w_i)\), the weighted mean update rule is applied:
--
-- \[
--   m_i' \;=\; \frac{m_i \, w_i \;+\; x \, w}{w_i + w},
--   \qquad
--   w_i' \;=\; w_i + w
-- \]
--
-- This is the standard incremental weighted mean, which is exact in
-- floating-point arithmetic up to the usual rounding.  Note that the centroid
-- does /not/ store individual observations — only the summary statistics
-- \((m_i, w_i)\) are retained, which is what gives the t-digest its bounded
-- space.
data Centroid = Centroid
  { -- | Weighted mean of all values merged into this centroid.
    cMean :: {-# UNPACK #-} !Double,
    -- | Total weight (count) of values in this centroid.  For unweighted
    -- streams, this is simply the number of observations that have been
    -- merged into this centroid.
    cWeight :: {-# UNPACK #-} !Double
  }
  deriving (Show)

instance Measured Measure Centroid where
  measure c = Measure (cWeight c) 1 (cMean c) (cMean c * cWeight c)

-- | The t-digest data structure for online quantile estimation.
--
-- Internally, a t'TDigest' consists of:
--
-- * A 'FingerTree' of t'Centroid's, sorted by mean.  The tree carries the
--   four-component @Measure@ described above, enabling \(O(\log n)\) split
--   operations by both mean and cumulative weight.
--
-- * Cached metadata: the total weight \(N = \sum w_i\), the global minimum
--   and maximum of all observed values, the compression parameter \(\delta\),
--   and the compression threshold \(3\delta\).
--
-- __Invariants:__
--
-- 1. Centroids are sorted in non-decreasing order of 'cMean'.
-- 2. @tdTotalWeight@ equals @mWeight (measure tdCentroids)@ and equals the
--    sum of all 'cWeight' values.
-- 3. @tdMin@ \(\le m_1\) and @tdMax@ \(\ge m_k\) (where \(k\) is the number
--    of centroids), with equality in the singleton case.
-- 4. After 'compress', every centroid satisfies the \(K_1\) size constraint:
--    \(k(q_{\mathrm{upper}}, \delta) - k(q_{\mathrm{lower}}, \delta) \le 1\),
--    where \(q_{\mathrm{lower}}\) and \(q_{\mathrm{upper}}\) are the
--    normalised cumulative weight boundaries of the centroid.
-- 5. The centroid count never exceeds \(3\delta\) for sustained periods;
--    insertions that push the count above this threshold trigger an automatic
--    'compress' pass.
data TDigest = TDigest
  { tdCentroids :: !(FingerTree Measure Centroid),
    tdTotalWeight :: !Double,
    tdMin :: !Double,
    tdMax :: !Double,
    tdDelta :: !Double,
    tdMaxCentroids :: {-# UNPACK #-} !Int
  }
  deriving (Show)

-- ---------------------------------------------------------------------------
-- Construction
-- ---------------------------------------------------------------------------

-- | Create an empty t-digest with the default compression parameter
-- \(\delta = 100\).
--
-- This is a good starting point for most applications.  With \(\delta = 100\),
-- the digest will use at most 300 centroids (the compression threshold is
-- \(3\delta\)), occupying roughly 4.8 KB of centroid data.  Empirically, this
-- yields quantile errors below \(10^{-4}\) at the median and below
-- \(10^{-6}\) for \(q < 0.01\) or \(q > 0.99\)
-- (Dunning & Ertl, 2019; <https://arxiv.org/abs/1902.04023>).
empty :: TDigest
empty = emptyWith 100

-- | Create an empty t-digest with a given compression parameter \(\delta\).
--
-- The compression parameter controls the trade-off between accuracy and space:
--
-- * __Larger \(\delta\)__ (e.g., 200–500) means more centroids are retained,
--   giving higher accuracy — especially at extreme quantiles — at the cost of
--   more memory and slower queries.
-- * __Smaller \(\delta\)__ (e.g., 20–50) means fewer centroids, saving
--   memory but increasing quantile estimation error.
--
-- The maximum number of centroids after compression is \(\delta + 1\)
-- (one per integer unit in the range of \(K_1\)), and the compression
-- threshold (the point at which automatic compression is triggered during
-- insertion) is set to \(\lceil 3\delta \rceil\).  Typical values used in
-- production systems are \(\delta \in [50, 300]\).
--
-- Setting \(\delta \le 0\) is not meaningful and will result in a digest that
-- compresses aggressively to zero or one centroid.
emptyWith :: Double -> TDigest
emptyWith delta =
  TDigest
    { tdCentroids = FT.empty,
      tdTotalWeight = 0,
      tdMin = 1 / 0,
      tdMax = -(1 / 0),
      tdDelta = delta,
      tdMaxCentroids = ceiling (delta * 3)
    }

-- ---------------------------------------------------------------------------
-- Scale function K_1
-- ---------------------------------------------------------------------------

-- | The \(K_1\) (arcsine) scale function:
--
-- \[
--   k(q, \delta) \;=\; \frac{\delta}{2\pi}\,\arcsin(2q - 1)
-- \]
--
-- This function maps the quantile domain \([0, 1]\) to the "scale space"
-- \([-\delta/2,\; \delta/2]\).  Its derivative
-- \(k'(q) = \frac{\delta}{\pi\sqrt{q(1-q)}}\) diverges at \(q = 0\) and
-- \(q = 1\), causing centroids near the tails to be allocated much more
-- finely than centroids near the median — which is the defining feature of
-- the t-digest's accuracy profile.
kScale :: Double -> Double -> Double
kScale delta q = (delta / (2 * pi)) * asin (2 * q - 1)

-- | Inverse of the \(K_1\) scale function:
--
-- \[
--   q(k, \delta) \;=\; \frac{1 + \sin\!\bigl(\frac{2\pi k}{\delta}\bigr)}{2}
-- \]
--
-- Used during 'compress' to compute the quantile boundaries corresponding to
-- integer scale-function values, i.e., the boundaries of the unit intervals
-- in scale space.
kScaleInv :: Double -> Double -> Double
kScaleInv delta k = (1 + sin (2 * pi * k / delta)) / 2

-- ---------------------------------------------------------------------------
-- FingerTree helpers
-- ---------------------------------------------------------------------------

ftToList :: FingerTree Measure Centroid -> [Centroid]
ftToList ft = case FT.viewl ft of
  EmptyL -> []
  x :< rest -> x : ftToList rest

splitByMean :: Double -> FingerTree Measure Centroid -> (FingerTree Measure Centroid, FingerTree Measure Centroid)
splitByMean x = FT.split (\m -> mMaxMean m >= x)

-- ---------------------------------------------------------------------------
-- Adding values
-- ---------------------------------------------------------------------------

-- | Add a single value with weight 1 to the digest.
--
-- \(O(\log n)\) amortised, where \(n\) is the number of centroids.
-- Equivalent to @'addWeighted' x 1@.
add :: Double -> TDigest -> TDigest
add x = addWeighted x 1

-- | Add a value \(x\) with a given weight \(w\) to the digest.
--
-- The algorithm proceeds as follows:
--
-- 1. __Split__ the finger tree at the insertion point using
--    @'FT.split' (\m -> mMaxMean m >= x)@, yielding a left subtree (all
--    centroids with mean \(< x\)) and a right subtree (mean \(\ge x\)).
--    This is \(O(\log n)\) by the finger tree split theorem
--    (Hinze & Paterson, 2006; <https://doi.org/10.1017/S0956796805005769>).
--
-- 2. __Find nearest neighbour:__ examine the rightmost centroid of the left
--    subtree and the leftmost centroid of the right subtree.  For each
--    candidate neighbour \((m_i, w_i)\), compute the proposed merged weight
--    \(w_i + w\) and check the \(K_1\) scale-function constraint:
--
--    \[
--      k\!\bigl(q_{\mathrm{upper}},\, \delta\bigr)
--      \;-\; k\!\bigl(q_{\mathrm{lower}},\, \delta\bigr)
--      \;\le\; 1
--    \]
--
--    where \(q_{\mathrm{lower}}\) and \(q_{\mathrm{upper}}\) are the
--    normalised cumulative weight boundaries of the proposed merged centroid.
--
-- 3. __Merge or insert:__ if one or both neighbours can absorb the new value,
--    merge with the /closer/ one (by distance \(|m_i - x|\)) using the
--    weighted mean update rule.  If neither can absorb it (because doing so
--    would violate the size constraint), insert a new singleton centroid
--    \((x, w)\) into the tree.
--
-- 4. __Auto-compress:__ if the centroid count exceeds the threshold
--    \(3\delta\), trigger a 'compress' pass.
--
-- The overall amortised cost is \(O(\log n)\), dominated by the finger tree
-- split and concatenation.
addWeighted :: Double -> Double -> TDigest -> TDigest
addWeighted x w td =
  let n = tdTotalWeight td + w
      newMin = min x (tdMin td)
      newMax = max x (tdMax td)
      delta = tdDelta td
      cs = tdCentroids td
      newC = Centroid x w
      td' =
        if FT.null cs
          then
            td
              { tdCentroids = FT.singleton newC,
                tdTotalWeight = n,
                tdMin = newMin,
                tdMax = newMax
              }
          else
            let (left, right) = splitByMean x cs
                leftWeight = mWeight (FT.measure left)
                result = tryMergeNeighbor delta n leftWeight left right newC
             in td
                  { tdCentroids = result,
                    tdTotalWeight = n,
                    tdMin = newMin,
                    tdMax = newMax
                  }
   in if mCount (FT.measure (tdCentroids td')) > tdMaxCentroids td'
        then compress td'
        else td'

-- | Try to merge with nearest neighbor; insert if neither allows merging.
tryMergeNeighbor ::
  Double ->
  Double ->
  Double ->
  FingerTree Measure Centroid ->
  FingerTree Measure Centroid ->
  Centroid ->
  FingerTree Measure Centroid
tryMergeNeighbor delta n leftWeight left right newC =
  let x = cMean newC
      k = kScale delta

      leftNeighbor = case FT.viewr left of
        EmptyR -> Nothing
        leftRest :> lc ->
          let cumBefore = mWeight (FT.measure leftRest)
              proposed = cWeight lc + cWeight newC
              q0 = cumBefore / n
              q1 = (cumBefore + proposed) / n
              canMerge = k q1 - k q0 <= 1.0
              dist = abs (cMean lc - x)
           in if canMerge then Just (leftRest, lc, dist) else Nothing

      rightNeighbor = case FT.viewl right of
        EmptyL -> Nothing
        rc :< rightRest ->
          let proposed = cWeight rc + cWeight newC
              q0 = leftWeight / n
              q1 = (leftWeight + proposed) / n
              canMerge = k q1 - k q0 <= 1.0
              dist = abs (cMean rc - x)
           in if canMerge then Just (rightRest, rc, dist) else Nothing
   in case (leftNeighbor, rightNeighbor) of
        (Just (leftRest, lc, ldist), Just (rightRest, rc, rdist))
          | ldist <= rdist ->
              (leftRest |> mergeCentroid lc newC) FT.>< right
          | otherwise ->
              left FT.>< (mergeCentroid rc newC <| rightRest)
        (Just (leftRest, lc, _), Nothing) ->
          (leftRest |> mergeCentroid lc newC) FT.>< right
        (Nothing, Just (rightRest, rc, _)) ->
          left FT.>< (mergeCentroid rc newC <| rightRest)
        (Nothing, Nothing) ->
          left FT.>< (newC <| right)

-- | Merge two centroids using the weighted mean update rule:
--
-- \[
--   m' = \frac{m_a \, w_a + m_b \, w_b}{w_a + w_b},
--   \qquad
--   w' = w_a + w_b
-- \]
mergeCentroid :: Centroid -> Centroid -> Centroid
mergeCentroid a b =
  let w = cWeight a + cWeight b
      m = (cMean a * cWeight a + cMean b * cWeight b) / w
   in Centroid m w

-- ---------------------------------------------------------------------------
-- Compression (split-based greedy merge)
-- ---------------------------------------------------------------------------

-- | Compress the digest by merging centroids that fall within the same
-- \(K_1\) scale-function unit interval.
--
-- The compression algorithm works as follows:
--
-- 1. Compute the integer range of the \(K_1\) scale function:
--    \(j_{\min} = \lceil k(0, \delta) \rceil = \lceil -\delta/2 \rceil\) and
--    \(j_{\max} = \lfloor k(1, \delta) \rfloor = \lfloor \delta/2 \rfloor\).
--
-- 2. For each integer \(j \in \{j_{\min}+1, \ldots, j_{\max}\}\), compute the
--    cumulative weight boundary \(b_j = k^{-1}(j, \delta) \cdot N\), where
--    \(N\) is the total weight.
--
-- 3. Split the finger tree at each boundary \(b_j\) by cumulative weight
--    (using @'FT.split' (\m -> mWeight m > b_j)@), yielding \(O(\delta)\)
--    contiguous chunks.
--
-- 4. Collapse each chunk into a single centroid using the @mMeanWeightSum@
--    and @mWeight@ components of the monoidal measure:
--    \(\bar{m} = \texttt{mMeanWeightSum}\,/\,\texttt{mWeight}\).  This is
--    \(O(1)\) per chunk — no traversal of individual centroids is needed.
--
-- __Complexity:__ \(O(\delta \log n)\), because there are \(O(\delta)\) split
-- operations, each costing \(O(\log n)\) where \(n\) is the pre-compression
-- centroid count.  After compression, the centroid count is at most
-- \(\delta + 1\).
compress :: TDigest -> TDigest
compress td
  | cnt <= 1 = td
  | otherwise =
      let n = tdTotalWeight td
          delta = tdDelta td
          cs = tdCentroids td
          -- K1 range: k(0) = -delta/2, k(1) = delta/2
          -- Integer unit boundaries from ceil(k(0)) to floor(k(1))
          kMin = kScale delta 0 -- = -delta/2
          kMax = kScale delta 1 -- = +delta/2
          jMin = ceiling kMin :: Int
          jMax = floor kMax :: Int
          -- Build boundaries: q values at each integer k-value
          boundaries = [kScaleInv delta (fromIntegral j) * n | j <- [jMin + 1 .. jMax]]
          -- Split-and-merge at each boundary
          merged = splitMerge boundaries cs
       in td {tdCentroids = merged}
  where
    cnt = mCount (FT.measure (tdCentroids td))

-- | Split a finger tree at cumulative weight boundaries and merge each
-- chunk into a single centroid.  This is the inner loop of 'compress'.
--
-- The function walks through the list of weight boundaries, performing
-- an @'FT.split'@ at each one.  Each resulting chunk (a contiguous sub-tree
-- of centroids whose combined weight falls within a single \(K_1\) unit
-- interval) is collapsed via 'mergeChunk' into a single centroid and appended
-- to the accumulator.
splitMerge :: [Double] -> FingerTree Measure Centroid -> FingerTree Measure Centroid
splitMerge boundaries tree = go boundaries tree FT.empty
  where
    go [] remaining acc =
      -- Last chunk: everything remaining
      case mergeChunk remaining of
        Nothing -> acc
        Just c -> acc |> c
    go (b : bs) remaining acc =
      let (chunk, rest) = FT.split (\m -> mWeight m > b) remaining
       in case mergeChunk chunk of
            Nothing -> go bs rest acc
            Just c -> go bs rest (acc |> c)

-- | Merge all centroids in a finger tree chunk into a single centroid
-- using the monoidal measure.  Runs in \(O(1)\) — no traversal of
-- individual centroids is needed, because the measure already caches
-- \(\sum w_i\) and \(\sum m_i w_i\).
mergeChunk :: FingerTree Measure Centroid -> Maybe Centroid
mergeChunk ft
  | w == 0 = Nothing
  | otherwise = Just (Centroid (mws / w) w)
  where
    m = FT.measure ft
    w = mWeight m
    mws = mMeanWeightSum m

-- ---------------------------------------------------------------------------
-- Quantile estimation
-- ---------------------------------------------------------------------------

-- | Estimate the value at quantile \(q\) (\(0 \le q \le 1\)).
--
-- The algorithm uses an interpolation scheme that treats each centroid as
-- representing a point mass at its mean, spread uniformly over a weight
-- interval centred at the centroid's cumulative midpoint.  Between
-- consecutive centroid midpoints, the estimated quantile function is linearly
-- interpolated:
--
-- \[
--   \hat{x}(q) \;=\; m_i + \frac{q \cdot N - \mathrm{mid}_i}
--   {\mathrm{mid}_{i+1} - \mathrm{mid}_i} \cdot (m_{i+1} - m_i)
-- \]
--
-- where \(\mathrm{mid}_i = \sum_{j<i} w_j + w_i/2\) is the cumulative
-- midpoint of centroid \(i\), and \(N = \sum w_j\).
--
-- __Boundary handling:__ for the leftmost centroid, if \(q \cdot N\) falls
-- below \(w_1 / 2\), the function interpolates between the global minimum
-- (@tdMin@) and \(m_1\).  Symmetrically, for the rightmost centroid, it
-- interpolates between \(m_k\) and the global maximum (@tdMax@).  This
-- ensures that 'quantile' returns @tdMin@ at \(q = 0\) and @tdMax@ at
-- \(q = 1\).
--
-- __Complexity:__ \(O(\log n)\) via @'FT.split'@ on cumulative weight,
-- followed by a constant amount of local interpolation work.
--
-- Returns 'Nothing' if the digest is empty.
quantile :: Double -> TDigest -> Maybe Double
quantile q td
  | numCentroids == 0 = Nothing
  | numCentroids == 1 =
      case FT.viewl cs of
        c :< _ -> Just (cMean c)
        EmptyL -> Nothing
  | otherwise = Just (findQuantile (clamp 0 1 q))
  where
    cs = tdCentroids td
    n = tdTotalWeight td
    mn = tdMin td
    mx = tdMax td
    numCentroids = mCount (FT.measure cs)

    findQuantile :: Double -> Double
    findQuantile q' =
      let target = q' * n
          (left, right) = FT.split (\m -> mWeight m > target) cs
          leftWeight = mWeight (FT.measure left)
          leftCount = mCount (FT.measure left)
       in case FT.viewl right of
            EmptyL ->
              case FT.viewr left of
                _ :> lastC -> interpolateRight lastC (leftWeight - cWeight lastC) target
                EmptyR -> mx
            cur :< rightRest ->
              interpolateAt leftCount leftWeight cur left rightRest target

    interpolateAt :: Int -> Double -> Centroid -> FingerTree Measure Centroid -> FingerTree Measure Centroid -> Double -> Double
    interpolateAt i cumulative c left rest target
      | i == 0 && target < cWeight c / 2 =
          if cWeight c == 1
            then mn
            else mn + (cMean c - mn) * (target / (cWeight c / 2))
      | i == numCentroids - 1 =
          if target > n - cWeight c / 2
            then
              if cWeight c == 1
                then mx
                else
                  let remaining = n - cWeight c / 2
                   in cMean c + (mx - cMean c) * ((target - remaining) / (cWeight c / 2))
            else cMean c
      | otherwise =
          let mid = cumulative + cWeight c / 2
           in case FT.viewl rest of
                nextC :< _ ->
                  let nextMid = cumulative + cWeight c + cWeight nextC / 2
                   in if target <= nextMid
                        then
                          let frac =
                                if nextMid == mid
                                  then 0.5
                                  else (target - mid) / (nextMid - mid)
                           in cMean c + frac * (cMean nextC - cMean c)
                        else
                          let newLeft = left FT.>< FT.singleton c
                           in interpolateAt (i + 1) (cumulative + cWeight c) nextC newLeft (ftTail rest) target
                EmptyL -> cMean c

    interpolateRight :: Centroid -> Double -> Double -> Double
    interpolateRight c _cumulative target =
      if target > n - cWeight c / 2
        then
          if cWeight c == 1
            then mx
            else
              let remaining = n - cWeight c / 2
               in cMean c + (mx - cMean c) * ((target - remaining) / (cWeight c / 2))
        else cMean c

    ftTail :: FingerTree Measure Centroid -> FingerTree Measure Centroid
    ftTail ft = case FT.viewl ft of
      EmptyL -> FT.empty
      _ :< r -> r

-- ---------------------------------------------------------------------------
-- CDF estimation
-- ---------------------------------------------------------------------------

-- | Estimate the cumulative distribution function (CDF) at value \(x\),
-- i.e., the fraction of the distribution that lies at or below \(x\).
--
-- The CDF is estimated by piecewise-linear interpolation between centroid
-- midpoints.  For a query point \(x\) falling between the means of
-- consecutive centroids \(m_i\) and \(m_{i+1}\), the estimated CDF is:
--
-- \[
--   \hat{F}(x) \;=\; \frac{1}{N}\left(
--     \mathrm{mid}_i + \frac{x - m_i}{m_{i+1} - m_i}
--     \cdot (\mathrm{mid}_{i+1} - \mathrm{mid}_i)
--   \right)
-- \]
--
-- where \(\mathrm{mid}_i = \sum_{j<i} w_j + w_i/2\).
--
-- __Boundary handling:__ if \(x \le \texttt{tdMin}\) the function returns 0;
-- if \(x \ge \texttt{tdMax}\) it returns 1.  For \(x\) below the first
-- centroid mean or above the last, the function interpolates between the
-- global extreme and the nearest centroid mean, mirroring the boundary
-- treatment in 'quantile'.
--
-- __Complexity:__ \(O(\log n)\) via @'FT.split'@ on the @mMaxMean@ component
-- of the monoidal measure, which locates the pair of centroids straddling
-- the query point without scanning.
--
-- Returns 'Nothing' if the digest is empty.
cdf :: Double -> TDigest -> Maybe Double
cdf x td
  | numCentroids == 0 = Nothing
  | x <= mn = Just 0
  | x >= mx = Just 1
  | otherwise = Just (findCdf x)
  where
    cs = tdCentroids td
    n = tdTotalWeight td
    mn = tdMin td
    mx = tdMax td
    numCentroids = mCount (FT.measure cs)

    findCdf :: Double -> Double
    findCdf x' =
      let (left, right) = splitByMean x' cs
       in case (FT.viewr left, FT.viewl right) of
            (EmptyR, rc :< _) ->
              cdfAtFirst rc x'
            (_, EmptyL) ->
              case FT.viewr left of
                lRest :> lc ->
                  cdfAtLast lc (mWeight (FT.measure lRest)) x'
                EmptyR -> 1.0
            (lRest :> lc, rc :< _) ->
              let lcCum = mWeight (FT.measure lRest)
                  lcIdx = mCount (FT.measure lRest)
                  rcIdx = mCount (FT.measure left)
               in if x' <= cMean lc
                    then
                      if lcIdx == 0
                        then cdfAtFirst lc x'
                        else case FT.viewr lRest of
                          llRest :> llc ->
                            cdfBetween llc (mWeight (FT.measure llRest)) lc lcCum x'
                          EmptyR -> cdfAtFirst lc x'
                    else
                      if rcIdx == numCentroids - 1 && x' > cMean rc
                        then cdfAtLast rc (mWeight (FT.measure left)) x'
                        else cdfBetween lc lcCum rc (mWeight (FT.measure left)) x'

    cdfAtFirst :: Centroid -> Double -> Double
    cdfAtFirst c x'
      | x' < cMean c =
          let innerW = cWeight c / 2
              frac =
                if cMean c == mn
                  then 1.0
                  else (x' - mn) / (cMean c - mn)
           in (innerW * frac) / n
      | otherwise = (cWeight c / 2) / n

    cdfAtLast :: Centroid -> Double -> Double -> Double
    cdfAtLast c cumBefore x'
      | x' > cMean c =
          let halfW = cWeight c / 2
              rightW = n - cumBefore - halfW
              frac =
                if mx == cMean c
                  then 0.0
                  else (x' - cMean c) / (mx - cMean c)
           in (cumBefore + halfW + rightW * frac) / n
      | otherwise = (cumBefore + cWeight c / 2) / n

    cdfBetween :: Centroid -> Double -> Centroid -> Double -> Double -> Double
    cdfBetween lc lcCum rc rcCum x'
      | x' <= cMean lc = (lcCum + cWeight lc / 2) / n
      | x' >= cMean rc = (rcCum + cWeight rc / 2) / n
      | otherwise =
          let lMid = lcCum + cWeight lc / 2
              rMid = rcCum + cWeight rc / 2
              frac =
                if cMean lc == cMean rc
                  then 0.5
                  else (x' - cMean lc) / (cMean rc - cMean lc)
           in (lMid + frac * (rMid - lMid)) / n

-- ---------------------------------------------------------------------------
-- Merge
-- ---------------------------------------------------------------------------

-- | Merge two t-digests into one, preserving accuracy.
--
-- The merge operation inserts every centroid of the second digest into the
-- first (using 'addWeighted' with the centroid's mean and weight), then
-- applies 'compress' to restore the \(K_1\) size invariant.
--
-- This is the standard approach for combining digests computed on
-- disjoint data partitions, enabling distributed and parallel quantile
-- estimation.  In a MapReduce-style pipeline, each mapper builds a local
-- t'TDigest' and the reducer merges them with 'merge'.  Because 'compress'
-- enforces the same \(O(\delta)\) centroid bound, the merged result has
-- the same space footprint as a single-stream digest.
--
-- See Dunning (2021), Section 4.3 (<https://doi.org/10.1016/j.simpa.2020.100049>)
-- for a discussion of mergeability and its applications.
merge :: TDigest -> TDigest -> TDigest
merge td other =
  let otherCs = ftToList (tdCentroids other)
      combined = foldl' (\d c -> addWeighted (cMean c) (cWeight c) d) td otherCs
   in compress combined

-- ---------------------------------------------------------------------------
-- Queries
-- ---------------------------------------------------------------------------

-- | Return the total weight of all values added to the digest.
--
-- This is \(O(1)\), as the total weight is cached in the t'TDigest' record.
-- For an unweighted stream, this equals the number of observations.
totalWeight :: TDigest -> Double
totalWeight = tdTotalWeight

-- | Return the number of centroids currently stored in the digest.
--
-- This is \(O(1)\) via the @mCount@ component of the finger tree's monoidal
-- measure.  The count is always at most \(3\delta\) (and at most
-- \(\delta + 1\) immediately after 'compress').
centroidCount :: TDigest -> Int
centroidCount = mCount . FT.measure . tdCentroids

-- ---------------------------------------------------------------------------
-- Utility
-- ---------------------------------------------------------------------------

clamp :: Double -> Double -> Double -> Double
clamp lo hi x
  | x < lo = lo
  | x > hi = hi
  | otherwise = x

-- ---------------------------------------------------------------------------
-- Additional accessors (for Mutable interop)
-- ---------------------------------------------------------------------------

-- | Return the list of centroids in sorted order (by mean).
--
-- Useful for serialisation, interoperability with mutable implementations,
-- debugging, and visualisation of the digest's internal distribution.  The
-- list is produced by an in-order traversal of the finger tree in
-- \(O(n)\).
centroidList :: TDigest -> [Centroid]
centroidList = ftToList . tdCentroids

-- | Return the compression parameter \(\delta\).
--
-- This is needed for serialisation and for reconstructing a digest with
-- 'fromComponents'.
getDelta :: TDigest -> Double
getDelta = tdDelta

-- | Return the minimum observed value.
--
-- The global minimum is tracked separately from the centroids because the
-- first centroid's mean may be larger than the minimum (if multiple values
-- have been merged into it).  The minimum is used for boundary interpolation
-- in 'quantile' and 'cdf' at \(q \to 0\) and \(x \to \min\).
getMin :: TDigest -> Double
getMin = tdMin

-- | Return the maximum observed value.
--
-- Symmetric to 'getMin': the global maximum is used for boundary
-- interpolation in 'quantile' and 'cdf' at \(q \to 1\) and \(x \to \max\).
getMax :: TDigest -> Double
getMax = tdMax

-- | Reconstruct a t-digest from its serialised components: a list of
-- centroids (which /must/ be in non-decreasing order of mean), the total
-- weight, the global minimum and maximum, and the compression parameter
-- \(\delta\).
--
-- This function trusts the caller to provide correctly sorted centroids and
-- consistent metadata.  It is intended for deserialisation and for
-- transferring digests between this pure implementation and the mutable
-- array-backed implementations in other languages.  No validation or
-- re-compression is performed.
--
-- __Usage example:__
--
-- @
-- let cs = 'centroidList' td
--     tw = 'totalWeight' td
--     mn = 'getMin' td
--     mx = 'getMax' td
--     d  = 'getDelta' td
--     td' = 'fromComponents' cs tw mn mx d
-- -- td' is equivalent to td
-- @
fromComponents :: [Centroid] -> Double -> Double -> Double -> Double -> TDigest
fromComponents cs tw mn mx delta =
  TDigest
    { tdCentroids = FT.fromList cs,
      tdTotalWeight = tw,
      tdMin = mn,
      tdMax = mx,
      tdDelta = delta,
      tdMaxCentroids = ceiling (delta * 3)
    }