{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}

-- |
-- Module: Numeric.Eproc.ConfSeq
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Anytime-valid confidence sequence for the mean of bounded
-- observations.
--
-- For samples @x_t@ in @[lo, hi]@ with common conditional mean
--
--     @mu = E[x_t | F_{t-1}]   for all t@
--
-- (@F_{t-1}@ being the filtration generated by everything observed
-- strictly before time @t@; for i.i.d. samples this is just
-- @E[x]@), the running state yields a confidence interval @C_t@
-- after every observation, with time-uniform coverage:
--
--     @P(for all t, mu in C_t) >= 1 - alpha@
--
-- whenever @C_t@ is reported at all (see 'interval' for the empty
-- case). The guarantee holds uniformly over time, so the user may
-- inspect the interval after every observation and stop at any
-- data-dependent time -- optional stopping does not erode coverage.
--
-- The construction is the /hedged capital/ confidence sequence of
-- Waudby-Smith & Ramdas (2024), Theorem 3, evaluated over a finite
-- grid of candidate means. All arithmetic is carried out in
-- @[0, 1]@ coordinates internally; observations are mapped affinely
-- at the boundary. Each candidate @m@ runs a pair of betting
-- processes: a /positive-direction/ capital @K^+_t(m)@ wagering
-- that the mean exceeds @m@, and a /negative-direction/ capital
-- @K^-_t(m)@ wagering the reverse. The base bet is a single
-- predictable plug-in (their eq. (26)), computed once per update
-- from the running regularized mean and variance of the data and
-- shared by every candidate: it never depends on @m@, and only a
-- final truncation to @c \/ m@ (respectively @c \/ (1 - m)@), with
-- @c = 1\/2@, is candidate-specific. This @m@-freeness is what
-- makes the survivor set provably an interval (Theorem 3);
-- @m@-dependent bets can produce non-interval survivor sets (their
-- Section E.4), which is why this module does not use the library's
-- 'Numeric.Eproc.Common.Bettor' strategies.
--
-- A candidate @m@ is rejected once the max-hedge (@theta = 1\/2@)
-- capital @max(K^+_t(m), K^-_t(m)) \/ 2@ crosses @1 \/ alpha@.
-- Under the truth @m = mu@ each capital process is a nonnegative
-- supermartingale, the max is dominated by the convex combination
-- @(K^+ + K^-) \/ 2@, and Ville's inequality bounds the probability
-- that the truth is ever rejected by @alpha@. No multiplicity
-- correction across grid candidates is needed: coverage concerns
-- only the true mean's own test, and rejection of other candidates
-- merely tightens the interval.
--
-- Grid resolution is an accuracy\/cost knob. Interval endpoints are
-- quantized to the grid -- a @g@-point grid resolves them to within
-- @(hi - lo) \/ (g + 1)@ -- and per-update cost is @O(live
-- candidates)@, shrinking as evidence accumulates and candidates
-- are rejected.
--
-- == Example
--
-- Estimate the mean of a stream in @[0, 1]@ with empirical mean
-- @0.8@, at level @alpha = 0.05@ on a 100-point grid:
--
-- >>> let Right cfg = config 0.0 1.0 0.05 100
-- >>> let xs = concat (replicate 50 [1, 1, 0, 1, 1, 0, 1, 1, 1, 1])
-- >>> interval cfg (foldl' (update cfg) (initial cfg) xs)
-- Just (0.7326732673267327,0.8613861386138614)

module Numeric.Eproc.ConfSeq (
  -- * Confidence-sequence configuration and state
    Config
  , State
  , ConfigError(..)

  -- * Construction
  , config
  , initial

  -- * Streaming
  , update

  -- * Inspection
  , interval
  , samples
  ) where

import qualified Data.List as L
import GHC.Float (log1p)
import Numeric.Eproc.Common (ConfigError(..), finite)

-- types ----------------------------------------------------------------------

-- | Confidence-sequence configuration. Build with 'config'.
--
--   Carries the sample bounds, the significance level, the grid
--   size, and the precomputed per-candidate rejection threshold
--   @log(2 \/ alpha)@ along with the bet numerator
--   @2 log(2 \/ alpha)@.
data Config = Config {
    Config -> Double
cfg_lo         :: {-# UNPACK #-} !Double  -- ^ sample lower bound
  , Config -> Double
cfg_hi         :: {-# UNPACK #-} !Double  -- ^ sample upper bound
  , Config -> Double
cfg_alpha      :: {-# UNPACK #-} !Double  -- ^ significance level
  , Config -> Int
cfg_grid       :: {-# UNPACK #-} !Int     -- ^ grid size @g@
  , Config -> Double
cfg_log_thresh :: {-# UNPACK #-} !Double  -- ^ @log(2 \/ alpha)@
  , Config -> Double
cfg_bet_num    :: {-# UNPACK #-} !Double  -- ^ @2 log(2 \/ alpha)@
  }

-- | One live grid candidate: its grid index and the running
--   log-capitals of the positive- and negative-direction bets.
data Point = Point
  {-# UNPACK #-} !Int     -- grid index j
  {-# UNPACK #-} !Double  -- log K^+
  {-# UNPACK #-} !Double  -- log K^-

-- | Streaming confidence-sequence state. Construct with 'initial'
--   and fold observations through 'update'.
--
--   Carries the sample count, the shared plug-in bettor statistics
--   (regularized running sums in @[0, 1]@ coordinates), and the
--   live grid candidates. Rejected candidates are dropped
--   permanently, so the reported intervals are nested.
--
--   Invariant: 'initial' and 'update' construct the live list fully
--   forced -- no thunks in the spine or the elements -- so a 'State'
--   in WHNF is already in normal form.
data State = State {
    State -> Int
st_n        :: {-# UNPACK #-} !Int     -- ^ sample count
  , State -> Double
st_sum_y    :: {-# UNPACK #-} !Double  -- ^ @sum y_i@
  , State -> Double
st_sum_dev2 :: {-# UNPACK #-} !Double  -- ^ @sum (y_i - mu_i)^2@
  , State -> [Point]
st_live     :: ![Point]                -- ^ live grid candidates
  }

-- | WSR (2024) truncation level @c = 1\/2@. Bets are capped at
--   @c \/ m@ (positive direction) and @c \/ (1 - m)@ (negative
--   direction), keeping every capital factor at least @1 - c > 0@.
trunc_c :: Double
trunc_c :: Double
trunc_c = Double
0.5
{-# INLINE trunc_c #-}

-- construction ---------------------------------------------------------------

-- | Build a 'Config' for the confidence sequence.
--
--   The candidate means form the interior grid
--
--       @m_j = lo + (j \/ (g + 1)) * (hi - lo),   j = 1 .. g@
--
--   (endpoints excluded, so that in @[0, 1]@ coordinates the bet
--   truncations @c \/ m@ and @c \/ (1 - m)@ stay finite). The
--   per-candidate rejection threshold @log(2 \/ alpha)@ and the bet
--   numerator @2 log(2 \/ alpha)@ are precomputed.
--
--   Returns 'Left' with a 'ConfigError' on inputs that would leave
--   the mathematical regime: @alpha@ non-finite or outside
--   @(0, 1)@; @lo@ or @hi@ non-finite, or @lo >= hi@; or a grid
--   size below @1@.
--
--   >>> let Right cfg = config 0.0 1.0 0.05 100
config
  :: Double  -- ^ sample lower bound @lo@
  -> Double  -- ^ sample upper bound @hi@
  -> Double  -- ^ significance level @alpha@
  -> Int     -- ^ grid size @g@
  -> Either ConfigError Config
config :: Double -> Double -> Double -> Int -> Either ConfigError Config
config !Double
lo !Double
hi !Double
alpha !Int
g
  | Bool -> Bool
not (Double -> Bool
finite Double
alpha Bool -> Bool -> Bool
&& Double
alpha Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& Double
alpha Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1) =
      ConfigError -> Either ConfigError Config
forall a b. a -> Either a b
Left (Double -> ConfigError
InvalidAlpha Double
alpha)
  | Bool -> Bool
not (Double -> Bool
finite Double
lo Bool -> Bool -> Bool
&& Double -> Bool
finite Double
hi Bool -> Bool -> Bool
&& Double
lo Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
hi) =
      ConfigError -> Either ConfigError Config
forall a b. a -> Either a b
Left (Double -> Double -> ConfigError
InvalidBounds Double
lo Double
hi)
  | Int
g Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 =
      ConfigError -> Either ConfigError Config
forall a b. a -> Either a b
Left (Int -> ConfigError
InvalidGridSize Int
g)
  | Bool
otherwise = Config -> Either ConfigError Config
forall a b. b -> Either a b
Right Config {
        cfg_lo :: Double
cfg_lo         = Double
lo
      , cfg_hi :: Double
cfg_hi         = Double
hi
      , cfg_alpha :: Double
cfg_alpha      = Double
alpha
      , cfg_grid :: Int
cfg_grid       = Int
g
      , cfg_log_thresh :: Double
cfg_log_thresh = Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
alpha)
      , cfg_bet_num :: Double
cfg_bet_num    = Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
alpha)
      }
{-# INLINE config #-}

-- | The initial 'State' for a fresh confidence sequence.
--
--   Every grid candidate starts live with both log-capitals at @0@
--   (i.e., @K^+ = K^- = 1@); the shared bettor statistics start
--   from their regularized priors (@mu_0 = 1\/2@,
--   @sigma^2_0 = 1\/4@ in @[0, 1]@ coordinates).
--
--   >>> let s0 = initial cfg
initial :: Config -> State
initial :: Config -> State
initial Config{Double
Int
cfg_lo :: Config -> Double
cfg_hi :: Config -> Double
cfg_alpha :: Config -> Double
cfg_grid :: Config -> Int
cfg_log_thresh :: Config -> Double
cfg_bet_num :: Config -> Double
cfg_lo :: Double
cfg_hi :: Double
cfg_alpha :: Double
cfg_grid :: Int
cfg_log_thresh :: Double
cfg_bet_num :: Double
..} = State {
      st_n :: Int
st_n        = Int
0
    , st_sum_y :: Double
st_sum_y    = Double
0
    , st_sum_dev2 :: Double
st_sum_dev2 = Double
0
    , st_live :: [Point]
st_live     = Int -> [Point]
points Int
1
    }
  where
    -- built eagerly: the tail is forced before consing, so the
    -- whole list is in normal form on construction.
    points :: Int -> [Point]
points !Int
j
      | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
cfg_grid = []
      | Bool
otherwise    =
          let !p :: Point
p    = Int -> Double -> Double -> Point
Point Int
j Double
0 Double
0
              !rest :: [Point]
rest = Int -> [Point]
points (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          in  Point
p Point -> [Point] -> [Point]
forall a. a -> [a] -> [a]
: [Point]
rest
{-# INLINE initial #-}

-- streaming ------------------------------------------------------------------

-- | Fold one observation into the running 'State'.
--
--   Maps the observation to @[0, 1]@ coordinates via
--   @y = (x - lo) \/ (hi - lo)@ and computes the shared predictable
--   plug-in bet from the statistics accumulated through the
--   /previous/ step (Waudby-Smith & Ramdas (2024), eq. (26)):
--
--       @lambda_t = sqrt (2 log(2 \/ alpha)
--                     \/ (sigma^2_{t-1} * t * log(1 + t)))@
--
--   The bet is computed once and shared across all
--   live candidates -- its independence from @m@ is what keeps the
--   survivor set an interval. Each live candidate @m@ then updates
--   its pair of log-capitals with the truncated bets
--   @min lambda_t (c \/ m)@ and @min lambda_t (c \/ (1 - m))@,
--   with @c = 1\/2@, and
--   is dropped iff @max(log K^+, log K^-)@ has reached
--   @log(2 \/ alpha)@. Finally @y@ is folded into the shared
--   statistics, preserving predictability of the next bet.
--
--   /Precondition/: @x@ must lie in the @[lo, hi]@ interval given
--   to 'config'. The coverage guarantee of the sequence depends on
--   it. Out-of-range observations can drive a capital factor
--   negative, taking the construction out of the supermartingale
--   regime entirely; the function does not check for this.
--
--   >>> let s1 = update cfg s0 0.7
update :: Config -> State -> Double -> State
update :: Config -> State -> Double -> State
update Config{Double
Int
cfg_lo :: Config -> Double
cfg_hi :: Config -> Double
cfg_alpha :: Config -> Double
cfg_grid :: Config -> Int
cfg_log_thresh :: Config -> Double
cfg_bet_num :: Config -> Double
cfg_lo :: Double
cfg_hi :: Double
cfg_alpha :: Double
cfg_grid :: Int
cfg_log_thresh :: Double
cfg_bet_num :: Double
..} State{Double
Int
[Point]
st_n :: State -> Int
st_sum_y :: State -> Double
st_sum_dev2 :: State -> Double
st_live :: State -> [Point]
st_n :: Int
st_sum_y :: Double
st_sum_dev2 :: Double
st_live :: [Point]
..} !Double
x =
  let !y :: Double
y    = (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
cfg_lo) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
cfg_hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
cfg_lo)
      !t :: Int
t    = Int
st_n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      !td :: Double
td   = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
t
      !gp1 :: Double
gp1  = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
cfg_grid Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      -- sigma^2_{t-1} = (1/4 + sum_{i<=t-1} (y_i - mu_i)^2) / t
      !sig2 :: Double
sig2 = (Double
0.25 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
st_sum_dev2) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
td
      !lam :: Double
lam  = Double -> Double
forall a. Floating a => a -> a
sqrt (Double
cfg_bet_num Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
sig2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
td Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log1p Double
td))
      -- built eagerly, as in 'initial': the tail is forced before
      -- consing, so the new live list is in normal form on
      -- construction.
      go :: [Point] -> [Point]
go [] = []
      go (Point Int
j Double
lp Double
ln : [Point]
ps) =
        let !m :: Double
m    = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
gp1
            !d :: Double
d    = Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m
            !lp' :: Double
lp'  = Double
lp Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log1p (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
lam (Double
trunc_c Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
m) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d)
            !ln' :: Double
ln'  = Double
ln Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log1p (Double -> Double
forall a. Num a => a -> a
negate (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
lam (Double
trunc_c Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m)))
                                  Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
d)
            !rest :: [Point]
rest = [Point] -> [Point]
go [Point]
ps
        in  if Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
lp' Double
ln' Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
cfg_log_thresh
              then [Point]
rest
              else Int -> Double -> Double -> Point
Point Int
j Double
lp' Double
ln' Point -> [Point] -> [Point]
forall a. a -> [a] -> [a]
: [Point]
rest
      !live :: [Point]
live   = [Point] -> [Point]
go [Point]
st_live
      -- fold y into the shared statistics only now: the bet above
      -- used statistics through t-1, so predictability holds. the
      -- deviation at step t uses the current-inclusive mean mu_t.
      !sum_y' :: Double
sum_y' = Double
st_sum_y Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
y
      !mu :: Double
mu     = (Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sum_y') Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
td Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)
      !dev :: Double
dev    = Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu
      !dev2' :: Double
dev2'  = Double
st_sum_dev2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
dev Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
dev
  in  Int -> Double -> Double -> [Point] -> State
State Int
t Double
sum_y' Double
dev2' [Point]
live
{-# INLINE update #-}

-- inspection -----------------------------------------------------------------

-- | The current confidence interval, in the original @[lo, hi]@
--   coordinates.
--
--   The interval spans the surviving grid candidates, widened by
--   one grid step at each end (or clamped to @lo@ \/ @hi@ at the
--   grid's edges). The widening is what makes off-grid true means
--   safe: Theorem 3 guarantees the ideal continuum survivor set is
--   an interval, so its endpoints are bracketed by the nearest
--   /rejected/ grid candidates, and reporting those sentinels
--   yields a superset of the continuum interval. Whenever the
--   result is 'Just', it therefore covers the true mean uniformly
--   over time with probability at least @1 - alpha@ -- no
--   multiplicity correction across candidates is needed, since
--   coverage concerns only the true mean's own test.
--
--   'Nothing' means every grid candidate has been rejected: the
--   evidence has resolved the mean below the grid's resolution.
--   For a true mean lying exactly on the grid this has probability
--   at most @alpha@ (its own test must have rejected). For an
--   off-grid true mean it additionally occurs once the continuum
--   survivor interval shrinks inside a single grid cell -- a
--   quantization horizon far beyond the point where the reported
--   width is comparable to the grid spacing. Treat 'Nothing' as a
--   signal to rerun with a larger grid, not as an inference.
--
--   >>> interval cfg (initial cfg)
--   Just (0.0,1.0)
interval :: Config -> State -> Maybe (Double, Double)
interval :: Config -> State -> Maybe (Double, Double)
interval Config{Double
Int
cfg_lo :: Config -> Double
cfg_hi :: Config -> Double
cfg_alpha :: Config -> Double
cfg_grid :: Config -> Int
cfg_log_thresh :: Config -> Double
cfg_bet_num :: Config -> Double
cfg_lo :: Double
cfg_hi :: Double
cfg_alpha :: Double
cfg_grid :: Int
cfg_log_thresh :: Double
cfg_bet_num :: Double
..} State{Double
Int
[Point]
st_n :: State -> Int
st_sum_y :: State -> Double
st_sum_dev2 :: State -> Double
st_live :: State -> [Point]
st_n :: Int
st_sum_y :: Double
st_sum_dev2 :: Double
st_live :: [Point]
..} = case [Point]
st_live of
  []                  -> Maybe (Double, Double)
forall a. Maybe a
Nothing
  (Point Int
j0 Double
_ Double
_ : [Point]
ps) ->
    let !jmin :: Int
jmin = (Int -> Point -> Int) -> Int -> [Point] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' (\Int
acc (Point Int
j Double
_ Double
_) -> Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
acc Int
j) Int
j0 [Point]
ps
        !jmax :: Int
jmax = (Int -> Point -> Int) -> Int -> [Point] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' (\Int
acc (Point Int
j Double
_ Double
_) -> Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
acc Int
j) Int
j0 [Point]
ps
        !gp1 :: Double
gp1  = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
cfg_grid Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        !w :: Double
w    = Double
cfg_hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
cfg_lo
        !l :: Double
l | Int
jmin Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1        = Double
cfg_lo
           | Bool
otherwise        =
               Double
cfg_lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
jmin Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
gp1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
w
        !u :: Double
u | Int
jmax Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
cfg_grid = Double
cfg_hi
           | Bool
otherwise        =
               Double
cfg_lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
jmax Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
gp1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
w
    in  (Double, Double) -> Maybe (Double, Double)
forall a. a -> Maybe a
Just (Double
l, Double
u)
{-# INLINE interval #-}

-- | The number of samples consumed so far.
--
--   >>> samples s0
--   0
samples :: State -> Int
samples :: State -> Int
samples = State -> Int
st_n
{-# INLINE samples #-}