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

-- |
-- Module: Numeric.Eproc.Bernoulli
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- One-sided Bernoulli rate anytime-valid test.
--
-- For samples @x_t@ in @{0, 1}@, tests @H_0: E[x] <= p_0@ against
-- @H_1: E[x] > p_0@.
--
-- A single wealth process is run:
--
--     @W_n = prod_{i=1..n} (1 + lambda_i * (x_i - p_0))@
--
-- where each per-step bet @lambda_i@ is chosen predictably (from
-- data observed strictly before step @i@) and clipped to
-- @[0, lambda_max]@ so that the wealth factor stays nonnegative for
-- every admissible observation. Under @H_0@ the wealth process is
-- a nonnegative supermartingale, so by Ville's inequality the
-- probability of @W_n@ ever crossing @1 \/ alpha@ is at most
-- @alpha@, regardless of when the user decides to stop streaming
-- samples.
--
-- Unlike "Numeric.Eproc.Bounded", the alternative here is one-sided,
-- so a single wealth process suffices and no Bonferroni adjustment
-- is needed -- the rejection threshold is @log(1 \/ alpha)@.
--
-- == Example
--
-- Test @H_0: E[x] <= 0.05@ at level @alpha = 1e-3@ against a stream
-- with empirical rate @~0.5@:
--
-- >>> let cfg = config 1.0e-3 0.05 Newton
-- >>> let xs  = take 200 (cycle [True, False])
-- >>> decide cfg (foldl' (update cfg) (initial cfg) xs)
-- Reject

module Numeric.Eproc.Bernoulli (
  -- * Test configuration and state
    Config
  , State
  , Verdict(..)

  -- * Bettor strategies
  , Bettor(..)

  -- * Construction
  , config
  , initial

  -- * Streaming
  , update
  , decide

  -- * Inspection
  , log_wealth
  , samples
  ) where

import Numeric.Eproc.Common (Bettor(..), Verdict(..))

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

-- here, the centred observation @z_t@ referenced in
-- "Numeric.Eproc.Common" is @x_t - p_0@; the safe-bet ceiling
-- @lambda_max@ is derived from @p_0@ (see 'config').

-- bettor state. one constructor per 'Bettor' alternative; the
-- constructor used in a given 'State' matches the 'Bettor' chosen in
-- the enclosing 'Config'.
data BetState =
    SFixed
  | SAdaptive
      {-# UNPACK #-} !Double  -- sum of z (centred observation)
      {-# UNPACK #-} !Double  -- sum of z^2 (for online variance)
      {-# UNPACK #-} !Int     -- count
  | SNewton
      {-# UNPACK #-} !Double  -- current bet lambda
      {-# UNPACK #-} !Double  -- running sum of per-step squared gradients

-- | Bernoulli rate test configuration. Build with 'config'.
--
--   Carries the bettor strategy, the baseline rate, the significance
--   level, the precomputed log-wealth rejection threshold, and the
--   safe-bet ceiling derived from @p_0@.
data Config = Config {
    -- ^ bettor strategy
    Config -> Bettor
cfg_bettor     :: !Bettor
    -- ^ safe-bet ceiling
  , Config -> Double
cfg_lam_max    :: {-# UNPACK #-} !Double
    -- ^ baseline rate @p_0@
  , Config -> Double
cfg_p0         :: {-# UNPACK #-} !Double
    -- ^ significance level @alpha@
  , Config -> Double
cfg_alpha      :: {-# UNPACK #-} !Double
    -- ^ rejection threshold @log(1 \/ alpha)@
  , Config -> Double
cfg_log_thresh :: {-# UNPACK #-} !Double
  }

-- | Streaming test state. Construct with 'initial' and fold
--   observations through 'update'.
--
--   Carries the sample count, running log-wealth, and whatever
--   per-step state the chosen 'Bettor' needs.
data State = State {
    State -> Int
st_n     :: {-# UNPACK #-} !Int       -- ^ sample count
  , State -> Double
st_log_w :: {-# UNPACK #-} !Double    -- ^ running log-wealth
  , State -> BetState
st_bet   :: !BetState                 -- ^ bettor state
  }

-- internal -------------------------------------------------------------------

-- per-bettor initial state.
init_bet :: Bettor -> BetState
init_bet :: Bettor -> BetState
init_bet Bettor
b = case Bettor
b of
  Fixed Double
_  -> BetState
SFixed
  Bettor
Adaptive -> Double -> Double -> Int -> BetState
SAdaptive Double
0 Double
0 Int
0
  Bettor
Newton   -> Double -> Double -> BetState
SNewton Double
0 Double
1.0e-6  -- small acc seed avoids div-by-zero
{-# INLINE init_bet #-}

-- compute the next bet 'lambda' from the bettor and its current
-- state. for Adaptive we form a Kelly-style plug-in from the running
-- sample mean and variance; for Newton the bet is just the last
-- lambda chosen by the Newton step (updated during 'step_bet').
bet_lambda :: Bettor -> Double -> BetState -> Double
bet_lambda :: Bettor -> Double -> BetState -> Double
bet_lambda Bettor
b !Double
lam_max !BetState
s = case Bettor
b of
  Fixed Double
lam -> Double
lam
  Bettor
Adaptive -> case BetState
s of
    SAdaptive !Double
sm !Double
sm2 !Int
n
      | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0    -> Double
0
      | Bool
otherwise ->
          let !nd :: Double
nd  = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
              !mu :: Double
mu  = Double
sm Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd
              !mu2 :: Double
mu2 = Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
mu
              !var :: Double
var = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double
sm2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
nd Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu2)
              !den :: Double
den = Double
var Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
mu2
              !raw :: Double
raw = if Double
den Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
mu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
den
          in  Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
lam_max Double
raw)
    BetState
_ -> Double
0
  Bettor
Newton -> case BetState
s of
    SNewton !Double
lam Double
_ -> Double
lam
    BetState
_              -> Double
0
{-# INLINE bet_lambda #-}

-- update bettor state with newly observed centred value 'z'. for
-- Adaptive this is just accumulating sums; for Newton we take one
-- Newton step on the per-step log-wealth loss '-log(1 + lambda * z)',
-- accumulating squared gradients for adaptive scaling.
step_bet :: Bettor -> Double -> BetState -> Double -> BetState
step_bet :: Bettor -> Double -> BetState -> Double -> BetState
step_bet Bettor
b !Double
lam_max !BetState
s !Double
z = case Bettor
b of
  Fixed Double
_ -> BetState
SFixed
  Bettor
Adaptive -> case BetState
s of
    SAdaptive !Double
sm !Double
sm2 !Int
n -> Double -> Double -> Int -> BetState
SAdaptive (Double
sm Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
z) (Double
sm2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    BetState
_                     -> Double -> Double -> Int -> BetState
SAdaptive Double
z (Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z) Int
1
  Bettor
Newton -> case BetState
s of
    SNewton !Double
lam !Double
acc ->
      let !denom :: Double
denom = Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
lam Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z
          !g :: Double
g     = if Double
denom Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double -> Double
forall a. Num a => a -> a
negate Double
z Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
denom
          !acc' :: Double
acc'  = Double
acc Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
g Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
g
          !lam' :: Double
lam'  = Double
lam Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
g Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
acc'
          !clp :: Double
clp   = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
lam_max Double
lam')
      in  Double -> Double -> BetState
SNewton Double
clp Double
acc'
    BetState
_ -> Double -> Double -> BetState
SNewton Double
0 Double
1.0e-6
{-# INLINE step_bet #-}

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

-- | Build a 'Config' for the Bernoulli rate test.
--
--   The safe-bet ceiling @lambda_max@ is set so that the wealth
--   factor @1 + lambda * (x - p_0)@ stays nonnegative for both
--   @x = 0@ and @x = 1@. The binding constraint is @x = 0@, which
--   requires @lambda <= 1 \/ p_0@; the ceiling stored is half this
--   to leave numerical margin -- the WSR safety recommendation.
--
--   @p_0@ must lie strictly in @(0, 1)@ and @alpha@ strictly in
--   @(0, 1)@. The degenerate case @p_0 = 0@ would make @lambda_max@
--   infinite (any divergence would reject immediately and the test
--   becomes uninteresting); the caller is expected to pass a small
--   positive baseline.
--
--   >>> let cfg = config 1.0e-3 0.05 Newton
config
  :: Double  -- ^ significance level @alpha@, in @(0, 1)@
  -> Double  -- ^ baseline rate @p_0@, in @(0, 1)@
  -> Bettor  -- ^ bettor strategy
  -> Config
config :: Double -> Double -> Bettor -> Config
config !Double
alpha !Double
p0 !Bettor
b = Config {
    cfg_bettor :: Bettor
cfg_bettor     = Bettor
b
  , cfg_lam_max :: Double
cfg_lam_max    = Double
0.5 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
p0
  , cfg_p0 :: Double
cfg_p0         = Double
p0
  , cfg_alpha :: Double
cfg_alpha      = Double
alpha
  , cfg_log_thresh :: Double
cfg_log_thresh = Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
alpha)
  }
{-# INLINE config #-}

-- | The initial 'State' for a fresh streaming test.
--
--   Log-wealth starts at @0@ (i.e., wealth @1@) and the bettor
--   starts in the per-strategy initial state appropriate for the
--   'Bettor' chosen in the 'Config'.
--
--   >>> let s0 = initial cfg
initial :: Config -> State
initial :: Config -> State
initial Config{Double
Bettor
cfg_bettor :: Config -> Bettor
cfg_lam_max :: Config -> Double
cfg_p0 :: Config -> Double
cfg_alpha :: Config -> Double
cfg_log_thresh :: Config -> Double
cfg_bettor :: Bettor
cfg_lam_max :: Double
cfg_p0 :: Double
cfg_alpha :: Double
cfg_log_thresh :: Double
..} = State {
    st_n :: Int
st_n     = Int
0
  , st_log_w :: Double
st_log_w = Double
0
  , st_bet :: BetState
st_bet   = Bettor -> BetState
init_bet Bettor
cfg_bettor
  }
{-# INLINE initial #-}

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

-- | Fold one observation into the running 'State'.
--
--   @True@ means @x_t = 1@ (the event of interest occurred -- e.g.,
--   two readings diverged); @False@ means @x_t = 0@ (they matched).
--   The caller decides what \"matched\" means at the application
--   level.
--
--   Computes the centred observation @z = x - p_0@, queries the
--   bettor for its predictable bet, accumulates log-wealth via
--
--       @log_w' = log_w + log (1 + lambda * z)@
--
--   and then steps the bettor state given the newly observed @z@.
--
--   >>> let s1 = update cfg s0 True
update :: Config -> State -> Bool -> State
update :: Config -> State -> Bool -> State
update Config{Double
Bettor
cfg_bettor :: Config -> Bettor
cfg_lam_max :: Config -> Double
cfg_p0 :: Config -> Double
cfg_alpha :: Config -> Double
cfg_log_thresh :: Config -> Double
cfg_bettor :: Bettor
cfg_lam_max :: Double
cfg_p0 :: Double
cfg_alpha :: Double
cfg_log_thresh :: Double
..} State{Double
Int
BetState
st_n :: State -> Int
st_log_w :: State -> Double
st_bet :: State -> BetState
st_n :: Int
st_log_w :: Double
st_bet :: BetState
..} !Bool
x =
  let !xd :: Double
xd     = if Bool
x then Double
1 else Double
0
      !z :: Double
z      = Double
xd Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
cfg_p0
      !lam :: Double
lam    = Bettor -> Double -> BetState -> Double
bet_lambda Bettor
cfg_bettor Double
cfg_lam_max BetState
st_bet
      !fac :: Double
fac    = Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
lam Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z
      !logw' :: Double
logw'  = Double
st_log_w Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log Double
fac
      !s' :: BetState
s'     = Bettor -> Double -> BetState -> Double -> BetState
step_bet Bettor
cfg_bettor Double
cfg_lam_max BetState
st_bet Double
z
  in  Int -> Double -> BetState -> State
State (Int
st_n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double
logw' BetState
s'
{-# INLINE update #-}

-- | Compute the current 'Verdict' from the running 'State'.
--
--   'Reject' iff log-wealth has crossed the threshold
--   @log(1 \/ alpha)@; equivalently, wealth has exceeded
--   @1 \/ alpha@. Under @H_0@, by Ville's inequality, the
--   probability of this ever happening is at most @alpha@ -- and
--   crucially this bound holds at /every/ sample size
--   simultaneously, so the user is free to peek at the verdict as
--   often as they like and stop on the first 'Reject'.
--
--   >>> decide cfg s0
--   Continue
decide :: Config -> State -> Verdict
decide :: Config -> State -> Verdict
decide Config{Double
Bettor
cfg_bettor :: Config -> Bettor
cfg_lam_max :: Config -> Double
cfg_p0 :: Config -> Double
cfg_alpha :: Config -> Double
cfg_log_thresh :: Config -> Double
cfg_bettor :: Bettor
cfg_lam_max :: Double
cfg_p0 :: Double
cfg_alpha :: Double
cfg_log_thresh :: Double
..} State{Double
Int
BetState
st_n :: State -> Int
st_log_w :: State -> Double
st_bet :: State -> BetState
st_n :: Int
st_log_w :: Double
st_bet :: BetState
..}
  | Double
st_log_w Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
cfg_log_thresh = Verdict
Reject
  | Bool
otherwise                  = Verdict
Continue
{-# INLINE decide #-}

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

-- | The current log-wealth.
--
--   This is the natural \"test statistic\": it is monotone (in
--   expectation under @H_1@) in the evidence against @H_0@
--   accumulated so far, and the test rejects exactly when it crosses
--   @log(1 \/ alpha)@.
--
--   >>> log_wealth s0
--   0.0
log_wealth :: State -> Double
log_wealth :: State -> Double
log_wealth = State -> Double
st_log_w
{-# INLINE log_wealth #-}

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