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

-- |
-- Module: Numeric.Eproc.Mixture
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Uniform convex mixture of e-processes.
--
-- Given @K@ component e-processes @E^1_t, ..., E^K_t@ adapted to a
-- common filtration -- each testing (its facet of) a shared null
-- @H_0@ -- their arithmetic mean
--
--     @M_t = (E^1_t + ... + E^K_t) \/ K@
--
-- is itself an e-process with @M_0 = 1@: convex combinations
-- preserve the nonnegative-supermartingale property. By Ville's
-- inequality @P(sup_t M_t >= 1 \/ alpha) <= alpha@ under @H_0@, so a
-- level-@alpha@ test of the /combined/ null rejects when
-- @sup_t log(E^1_t + ... + E^K_t)@ crosses @log(K \/ alpha)@ -- no
-- Bonferroni correction, and strictly more powerful than one, since
-- the sum dominates the max. Use a mixture when the alternative has
-- several qualitatively different faces (a location shift, a shape
-- change, a rare-outlier channel, ...) and you want a single test
-- with power against their union.
--
-- This module does not own or update the components: they may be
-- heterogeneous (different test modules, different observation
-- transformations), so the caller steps each component itself and
-- feeds 'update' the vector of their current log e-values, as
-- reported by each module's @log_evalue@ accessor, one entry per
-- component in a fixed order.
--
-- Two preconditions are the caller's responsibility, and the
-- type-I guarantee depends on both:
--
--   1. Each entry must be the current log e-value of a genuine
--      e-process for @H_0@, and all components must be adapted to
--      the same filtration and stepped in lockstep -- 'update' is
--      called exactly once per underlying observation, after all
--      components have absorbed it.
--
--   2. The vector must have exactly the @K@ entries declared in
--      'config', always in the same order.
--
-- The rejection latch is kept on the supremum of the /mixture's/
-- log-wealth. Latching (or summing) per-component suprema instead
-- would combine peaks attained at different times -- a quantity
-- that can exceed anything the mixture ever reached, silently
-- inflating the effective alpha. Ville's inequality bounds the
-- mixture's own supremum; that is the only sound latch, and it is
-- the one this module maintains.
--
-- == Example
--
-- Combine a sign test and a magnitude test running against the same
-- stream of differences @d_t@ (the shape used for two-channel
-- symmetry testing):
--
-- >>> import qualified Numeric.Eproc.Bernoulli.TwoSided as Sign
-- >>> import qualified Numeric.Eproc.Bounded as Magn
-- >>> import qualified Numeric.Eproc.Mixture as Mix
-- >>> let Right sc = Sign.config 0.5 1.0e-3 Sign.Newton
-- >>> let Right mc = Magn.config 0.0 (-1.0) 1.0 1.0e-3 Magn.Newton
-- >>> let Right xc = Mix.config 2 1.0e-3
-- >>> :{
-- let step (s, m, x) d =
--       let s' = Sign.update sc s (d > 0)
--           m' = Magn.update mc m d
--       in  (s', m', Mix.update xc x
--                      [Sign.log_evalue s', Magn.log_evalue m'])
-- :}
-- >>> let ds = take 400 (cycle [0.6, 0.7, -0.2, 0.8])
-- >>> let z0 = (Sign.initial sc, Magn.initial mc, Mix.initial xc)
-- >>> let (_, _, xf) = foldl' step z0 ds
-- >>> Mix.decide xc xf
-- Reject
-- >>> Mix.p_value xc xf
-- 9.482234479673792e-34

module Numeric.Eproc.Mixture (
  -- * Mixture configuration and state
    Config
  , State
  , Verdict(..)
  , ConfigError(..)

  -- * Construction
  , config
  , initial

  -- * Streaming
  , update
  , decide

  -- * Inspection
  , log_wealth
  , log_wealth_sup
  , log_evalue
  , log_evalue_sup
  , p_value
  , samples
  ) where

import qualified Data.List as L
import Numeric.Eproc.Common (Verdict(..), ConfigError(..), finite)

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

-- | Mixture configuration. Build with 'config'.
--
--   Carries the component count @K@, the significance level, the
--   precomputed rejection threshold @log(K \/ alpha)@, and @log K@
--   (the mixture log-wealth of a fresh state).
data Config = Config {
    -- ^ component count @K@
    Config -> Int
cfg_k          :: {-# UNPACK #-} !Int
    -- ^ significance level @alpha@
  , Config -> Double
cfg_alpha      :: {-# UNPACK #-} !Double
    -- ^ rejection threshold @log(K \/ alpha)@
  , Config -> Double
cfg_log_thresh :: {-# UNPACK #-} !Double
    -- ^ @log K@
  , Config -> Double
cfg_log_k      :: {-# UNPACK #-} !Double
  }

-- | Streaming mixture state. Construct with 'initial' and fold
--   per-step component log e-value vectors through 'update'.
--
--   Tracks the current mixture log-wealth @log(sum_i E^i_t)@ and
--   its latched supremum, which is what 'decide' tests against the
--   rejection threshold.
data State = State {
    State -> Int
st_n           :: {-# UNPACK #-} !Int     -- ^ update count
  , State -> Double
st_log_sum     :: {-# UNPACK #-} !Double  -- ^ log(sum_i E^i)
  , State -> Double
st_sup_log_sum :: {-# UNPACK #-} !Double  -- ^ sup of the above
  }

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

-- | Build a 'Config' for a @K@-component uniform mixture at level
--   @alpha@.
--
--   The rejection threshold is precomputed as @log(K \/ alpha)@:
--   the mixture @M_t = (sum_i E^i_t) \/ K@ crosses @1 \/ alpha@
--   exactly when the sum crosses @K \/ alpha@.
--
--   Returns 'Left' with a 'ConfigError' on inputs outside the
--   mathematical regime: @K < 1@, or @alpha@ non-finite or outside
--   @(0, 1)@.
--
--   >>> let Right cfg = config 4 1.0e-3
config
  :: Int     -- ^ component count @K@
  -> Double  -- ^ significance level @alpha@
  -> Either ConfigError Config
config :: Int -> Double -> Either ConfigError Config
config !Int
k !Double
alpha
  | Int
k 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
InvalidComponentCount Int
k)
  | 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
otherwise =
      let !kd :: Double
kd = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k
      in  Config -> Either ConfigError Config
forall a b. b -> Either a b
Right Config {
              cfg_k :: Int
cfg_k          = Int
k
            , cfg_alpha :: Double
cfg_alpha      = Double
alpha
            , cfg_log_thresh :: Double
cfg_log_thresh = Double -> Double
forall a. Floating a => a -> a
log (Double
kd Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
alpha)
            , cfg_log_k :: Double
cfg_log_k      = Double -> Double
forall a. Floating a => a -> a
log Double
kd
            }
{-# INLINE config #-}

-- | The initial 'State' for a fresh mixture.
--
--   Every component starts at e-value @1@, so the mixture log-sum
--   (and its supremum) starts at @log K@.
--
--   >>> let s0 = initial cfg
initial :: Config -> State
initial :: Config -> State
initial Config{Double
Int
cfg_k :: Config -> Int
cfg_alpha :: Config -> Double
cfg_log_thresh :: Config -> Double
cfg_log_k :: Config -> Double
cfg_k :: Int
cfg_alpha :: Double
cfg_log_thresh :: Double
cfg_log_k :: Double
..} = State {
    st_n :: Int
st_n           = Int
0
  , st_log_sum :: Double
st_log_sum     = Double
cfg_log_k
  , st_sup_log_sum :: Double
st_sup_log_sum = Double
cfg_log_k
  }
{-# INLINE initial #-}

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

-- | Fold one step's component log e-values into the running
--   'State': computes the current mixture log-sum via a numerically
--   stable log-sum-exp and latches its supremum.
--
--   /Preconditions/ (documented in the module header, unchecked
--   here): the vector holds exactly the @K@ log e-values of
--   components adapted to a common filtration, in a fixed order,
--   with 'update' called once per underlying observation. The
--   degenerate empty vector leaves the state unchanged.
--
--   >>> let s1 = update cfg s0 [0.1, -0.2, 0.0, 0.4]
update :: Config -> State -> [Double] -> State
update :: Config -> State -> [Double] -> State
update Config
_ st :: State
st@State{Double
Int
st_n :: State -> Int
st_log_sum :: State -> Double
st_sup_log_sum :: State -> Double
st_n :: Int
st_log_sum :: Double
st_sup_log_sum :: Double
..} [Double]
les = case [Double]
les of
  []       -> State
st
  (Double
l : [Double]
ls) ->
    let !m :: Double
m = (Double -> Double -> Double) -> Double -> [Double] -> Double
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
l [Double]
ls
        !s :: Double
s = (Double -> Double -> Double) -> Double -> [Double] -> Double
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' (\ !Double
acc Double
v -> Double
acc Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
exp (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m)) Double
0 [Double]
les
        -- all components at e-value zero: the mixture log-sum is
        -- -Infinity, and (m +) would poison it into NaN.
        !cur :: Double
cur | Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
m Bool -> Bool -> Bool
&& Double
m Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 = Double
m
             | Bool
otherwise             = Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log Double
s
    in  State {
            st_n :: Int
st_n           = Int
st_n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          , st_log_sum :: Double
st_log_sum     = Double
cur
          , st_sup_log_sum :: Double
st_sup_log_sum = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
st_sup_log_sum Double
cur
          }
{-# INLINE update #-}

-- | Compute the current 'Verdict' from the running 'State'.
--
--   'Reject' iff the supremum-so-far of @log(sum_i E^i_t)@ has ever
--   crossed @log(K \/ alpha)@ -- equivalently, the mixture
--   e-process @M_t@ has exceeded @1 \/ alpha@ at some point in the
--   stream so far. Under the combined @H_0@, by Ville's inequality,
--   the probability of this ever happening is at most @alpha@,
--   simultaneously over all sample sizes: peek and stop freely.
--
--   >>> decide cfg s0
--   Continue
decide :: Config -> State -> Verdict
decide :: Config -> State -> Verdict
decide Config{Double
Int
cfg_k :: Config -> Int
cfg_alpha :: Config -> Double
cfg_log_thresh :: Config -> Double
cfg_log_k :: Config -> Double
cfg_k :: Int
cfg_alpha :: Double
cfg_log_thresh :: Double
cfg_log_k :: Double
..} State{Double
Int
st_n :: State -> Int
st_log_sum :: State -> Double
st_sup_log_sum :: State -> Double
st_n :: Int
st_log_sum :: Double
st_sup_log_sum :: Double
..}
  | Double
st_sup_log_sum 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 mixture log-wealth @log(sum_i E^i_t)@, before
--   normalization by @K@. Not monotone; bounded above by
--   'log_wealth_sup'. Starts at @log K@.
--
--   >>> log_wealth s0
--   1.3862943611198906
log_wealth :: State -> Double
log_wealth :: State -> Double
log_wealth = State -> Double
st_log_sum
{-# INLINE log_wealth #-}

-- | The supremum-so-far of @log(sum_i E^i_t)@. Monotone
--   nondecreasing; 'decide' rejects exactly when it crosses
--   @log(K \/ alpha)@. Starts at @log K@.
--
--   >>> log_wealth_sup s0
--   1.3862943611198906
log_wealth_sup :: State -> Double
log_wealth_sup :: State -> Double
log_wealth_sup = State -> Double
st_sup_log_sum
{-# INLINE log_wealth_sup #-}

-- | The current log e-value of the mixture: the log of
--   @M_t = (sum_i E^i_t) \/ K@, i.e. 'log_wealth' minus @log K@,
--   normalized so a fresh state sits at @0@. This is itself a
--   component-shaped quantity: mixtures nest, so it can in turn be
--   fed to an outer mixture. Not monotone; bounded above by
--   'log_evalue_sup'.
--
--   >>> log_evalue s0
--   0.0
log_evalue :: Config -> State -> Double
log_evalue :: Config -> State -> Double
log_evalue Config{Double
Int
cfg_k :: Config -> Int
cfg_alpha :: Config -> Double
cfg_log_thresh :: Config -> Double
cfg_log_k :: Config -> Double
cfg_k :: Int
cfg_alpha :: Double
cfg_log_thresh :: Double
cfg_log_k :: Double
..} State{Double
Int
st_n :: State -> Int
st_log_sum :: State -> Double
st_sup_log_sum :: State -> Double
st_n :: Int
st_log_sum :: Double
st_sup_log_sum :: Double
..} = Double
st_log_sum Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
cfg_log_k
{-# INLINE log_evalue #-}

-- | The supremum-so-far of the log e-value: 'log_wealth_sup' minus
--   @log K@. Monotone nondecreasing, starting at @0@; 'decide'
--   rejects exactly when it crosses @log(1 \/ alpha)@.
--
--   >>> log_evalue_sup s0
--   0.0
log_evalue_sup :: Config -> State -> Double
log_evalue_sup :: Config -> State -> Double
log_evalue_sup Config{Double
Int
cfg_k :: Config -> Int
cfg_alpha :: Config -> Double
cfg_log_thresh :: Config -> Double
cfg_log_k :: Config -> Double
cfg_k :: Int
cfg_alpha :: Double
cfg_log_thresh :: Double
cfg_log_k :: Double
..} State{Double
Int
st_n :: State -> Int
st_log_sum :: State -> Double
st_sup_log_sum :: State -> Double
st_n :: Int
st_log_sum :: Double
st_sup_log_sum :: Double
..} = Double
st_sup_log_sum Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
cfg_log_k
{-# INLINE log_evalue_sup #-}

-- | The anytime-valid p-value: the reciprocal of the largest
--   mixture e-value attained so far. Monotone nonincreasing; under
--   the combined @H_0@, @P(exists t: p_t <= alpha) <= alpha@ for
--   every @alpha@ simultaneously. 'decide' returns 'Reject' exactly
--   when this value has reached the configured @alpha@ or below.
--
--   >>> p_value cfg s0
--   1.0
p_value :: Config -> State -> Double
p_value :: Config -> State -> Double
p_value Config
cfg State
s = Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
1 (Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double
forall a. Num a => a -> a
negate (Config -> State -> Double
log_evalue_sup Config
cfg State
s)))
{-# INLINE p_value #-}

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