PiecewisePolynomial.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

{-|
Copyright   : Predictable Network Solutions Ltd., 2020-2024
License     : BSD-3-Clause
Description : Distributions of completion times via piecewise polynomials.

@'DQ'@ is a probability distribution of completion time
using the numeric type @Rational@.
This type represents a mixed discrete / continuous probability distribution
where the continuous part is represented in terms of piecewise polynomials.
-}
module DeltaQ.PiecewisePolynomial
    ( -- * Type
      DQ
    , distribution
    , fromPositiveMeasure
    , unsafeFromPositiveMeasure

    -- * Operations
    , meetsQTA
    , Moments (..)
    , moments
    , onConditionOfSuccess
    , timeout
    , residualTime

    -- * Internal
    , complexity
    ) where

import Algebra.PartialOrd
    ( PartialOrd (..)
    )
import Data.Maybe
    ( fromMaybe
    )
import DeltaQ.Class
    ( DeltaQ (..)
    , Outcome (..)
    , ProbabilisticOutcome (..)
    , eventuallyFromMaybe
    )
import Control.DeepSeq
    ( NFData
    )
import Numeric.Function.Piecewise
    ( Piecewise
    )
import Numeric.Measure.Finite.Mixed
    ( Measure
    )
import Numeric.Polynomial.Simple
    ( Poly
    )
import Numeric.Probability.Moments
    ( Moments (..)
    )
import Numeric.Function.Piecewise.Polynomial
    ( integral
    )
import qualified Data.Function.Class as Function
import qualified Numeric.Function.Piecewise as Piecewise
import qualified Numeric.Measure.Finite.Mixed as Measure
import qualified Numeric.Measure.Probability as Prob
import qualified Numeric.Polynomial.Simple as Poly

{-----------------------------------------------------------------------------
    Type
------------------------------------------------------------------------------}
-- | Probability distribution of durations.
newtype DQ = DQ (Measure Rational)
    deriving (Eq, Show, NFData)

-- | Get the distribution function as piecewise function of polynomials.
distribution :: DQ -> Piecewise (Poly Rational)
distribution (DQ m) = Measure.distribution m

-- | Interpret a finite, signed 'Measure' as a probability distribution.
--
-- In order to admit an interpretation as probability, the measure needs
-- to be positive.
-- This condition is checked, and if it does not hold,
-- the function returns 'Nothing'.
fromPositiveMeasure :: Measure Rational -> Maybe DQ
fromPositiveMeasure m
    | Measure.isPositive m = Just (unsafeFromPositiveMeasure m)
    | otherwise = Nothing

-- | Interpret a finite, positive 'Measure' as a probability distribution.
--
-- /The precondition that the measure is positive is not checked!/
unsafeFromPositiveMeasure :: Measure Rational -> DQ
unsafeFromPositiveMeasure = DQ

-- | Helper function for lifting a binary operation on distribution functions.
onDistribution2
    :: (a ~ Rational)
    => String
    -> (Piecewise (Poly a) -> Piecewise (Poly a) -> Piecewise (Poly a))
    -> DQ -> DQ -> DQ
onDistribution2 err f (DQ mx) (DQ my) =
    DQ
    $ fromMaybe impossible
    $ Measure.fromDistribution
    $ f (Measure.distribution mx) (Measure.distribution my)
  where
    impossible = error $ "impossible: not a finite measure in " <> err

-- | Size of the representation of a probability distribution,
-- i.e. number of pieces of the piecewise function and degrees
-- of the polynomials.
--
-- This quantity is relevant to stating and analyzing
-- the asymptotic time complexity of operations.
complexity :: DQ -> Int
complexity (DQ m) = sum (map complexityOfPiece pieces)
  where
    pieces = Piecewise.toAscPieces $ Measure.distribution m
    complexityOfPiece = (+1) . max 0 . Poly.degree . snd

{-----------------------------------------------------------------------------
    Operations
------------------------------------------------------------------------------}
instance Outcome DQ where
    type Duration DQ = Rational

    never = DQ Measure.zero

    wait t = DQ $ Measure.dirac t

    sequentially (DQ mx) (DQ my) = DQ (Measure.convolve mx my)

    firstToFinish = onDistribution2 "firstToFinish" $ \x y -> x + y - x * y

    lastToFinish = onDistribution2 "lastToFinish" (*)

instance ProbabilisticOutcome DQ where
    type Probability DQ = Rational

    choice p = onDistribution2 "choice" $ \x y ->
        scale p x + scale (1 - p) y
      where
        scale = Piecewise.mapPieces . Poly.scale

instance DeltaQ DQ where

    uniform a = DQ . Measure.uniform a

    successWithin (DQ m) = Function.eval (Measure.distribution m)

    failure (DQ m) = 1 - Measure.total m

    quantile (DQ m) p =
        eventuallyFromMaybe
        $ quantileFromMonotone (Measure.distribution m) p

    earliest (DQ m) = eventuallyFromMaybe $ fmap fst $ Measure.support m

    deadline (DQ m)= eventuallyFromMaybe $ fmap snd $ Measure.support m

-- | Partial order of cumulative distribution functions.
--
-- @'leq' x y@ holds if and only if for all completion times @t@,
-- the probability to succeed within the time @t@
-- is always larger (or equal) for @x@ compared to @y@.
-- In other words, @x@ has a higher probability of completing faster.
--
-- > x `leq` y  <=>  ∀ t. successWithin x t >= successWithin y t
instance PartialOrd DQ where
    m1 `leq` m2 =
        all isNonNegativeOnSegment
        $ toSegments
        $ distribution m1 - distribution m2

{-----------------------------------------------------------------------------
    Operations
    Helper functions
------------------------------------------------------------------------------}
-- | Helper type for segements of a piecewise functions.
data Segment a b
    = Jump a (b,b)
    | Polynomial (a,a) (b,b) (Poly a)
    | End a b
    deriving (Eq, Show)

-- | Helper function that elaborates a piecewise function
-- into a list of segments.
toSegments :: (a ~ Rational) => Piecewise (Poly a) -> [Segment a a]
toSegments = goJump 0 . Piecewise.toAscPieces
  where
    goJump _ [] = []
    goJump prev ((x1, o) : xos)
        | y0 < y1   = Jump x1 (y0, y1) : nexts
        | otherwise = nexts
      where
        y1 = Poly.eval o x1
        y0 = Poly.eval prev x1
        nexts = goPoly x1 y1 o xos

    goPoly x1 y1 o [] =
        End x1 y1 : goJump o []
    goPoly x1 y1 o xos@((x2, _) : _) =
        Polynomial (x1, x2) (y1, Poly.eval o x2) o : goJump o xos

{-----------------------------------------------------------------------------
    Operations
    quantile
------------------------------------------------------------------------------}
-- | Compute a quantile from a monotonically increasing function.
quantileFromMonotone :: (a ~ Rational) => Piecewise (Poly a) -> a -> Maybe a
quantileFromMonotone pieces = findInSegments segments
  where
    segments = toSegments pieces

    findInSegments _ 0
        = Just 0
    findInSegments [] _
        = Nothing
    findInSegments (Jump x1 (y1, y2) : xys) y
        | y1 <= y && y < y2 = Just x1
        | otherwise         = findInSegments xys y
    findInSegments (Polynomial (x1, x2) (y1, y2) o : xys) y
        | y1 == y           = Just x1
        | y1 <  y && y < y2 = Poly.root precision y (x1, x2) o
        | otherwise = findInSegments xys y
    findInSegments (End x1 y1 : _) y
        | y1 == y   = Just x1
        | otherwise = Nothing

precision :: Rational
precision = 1 / 10^(10 :: Integer)

{-----------------------------------------------------------------------------
    Operations
    meetsQTA
------------------------------------------------------------------------------}
-- | Test whether the given probability distribution of completion times
-- is equal to or better than a given
-- __quantitative timeliness agreement__ (QTA).
--
-- Synonym for `leq` of the partial order,
--
-- > p `meetsQTA` qta  =  p `leq` qta
meetsQTA :: DQ -> DQ -> Bool
meetsQTA = leq

isNonNegativeOnSegment :: (a ~ Rational) => Segment a a -> Bool
isNonNegativeOnSegment (Jump _ (y1, y2)) =
    y1 >= 0 && y2 >= 0
isNonNegativeOnSegment (Polynomial (x1, x2) _ poly) =
    compareToZero == Just GT || compareToZero == Just EQ
  where
    compareToZero = Poly.compareToZero (x1, x2, poly)
isNonNegativeOnSegment (End _ y) =
    y >= 0

{-----------------------------------------------------------------------------
    Operations
    Moments
------------------------------------------------------------------------------}
-- | For a given 'DQ', compute
--
-- * the success probability, and
--
-- * the 'Moments' of the conditional probability (distribution)
--   given that the outcome has succeeded.
moments :: DQ -> (Rational, Moments Rational)
moments (DQ m)
    | success == 0 =
        (0, Moments{mean = 0, variance = 0, skewness = 0, kurtosis = 1})
    | otherwise =
        (success, Prob.moments conditional)
  where
    success = Measure.total m
    conditional = Prob.unsafeFromMeasure $ Measure.scale (1/success) m

{-----------------------------------------------------------------------------
    Operations
    Conditional on success
------------------------------------------------------------------------------}
-- | Compute the probability (distribution) on the condition
-- that the outcome completes in a finite amount of time.
--
-- Put differently: Rescale the cumulative distribution function
-- such that the probability of failure becomes zero.
--
-- Return 'Nothing' if the outcome never completes, 'Just' otherwise.
-- 
-- > onConditionOfSuccess never = Nothing
-- > ∀x. failure x < 1  ==>  failure (onConditionOfSuccess x) = Just 0
-- > ∀x. failure x < 1  ==>
-- >   (snd . moments <$> onConditionOfSuccess x) = Just (snd (moments x))
onConditionOfSuccess :: DQ -> Maybe DQ
onConditionOfSuccess (DQ m)
    | success == 0 = Nothing
    | otherwise    = Just $ DQ conditional
  where
    success = Measure.total m
    conditional = Measure.scale (1/success) m

{-----------------------------------------------------------------------------
    Operations
    Timeout
------------------------------------------------------------------------------}
-- | Given a duration @dt@,
-- decompose a probability distribution of completion times
-- into two conditional probability distributions —
-- one distribution conditional on finishing within @dt@,
-- and the other distribution conditional on finishing after @dt@.
--
-- >   (within, p, after) = timeout dt o
-- >
-- > implies
-- >
-- >   o = choice p within (wait dt .>>. after)
-- >   p = successWithin o dt
-- >
-- >   successWithin within dt = 1
-- >     if p > 0
--
-- We define the corner cases as follows:
--
-- * @successWithin o dt = 0@ implies @before = never@.
--
-- * @successWithin o dt = 1@ implies @after = never@.
--
timeout :: Duration DQ -> DQ -> (DQ, Probability DQ, DQ)
timeout dt (DQ m) = (beforeCond, p, afterCond)
  where
    before = Measure.beforeOrAt dt m
    after = Measure.translate (- dt) $ Measure.after dt m
    p = Measure.total before
    beforeCond
        | p == 0 = never
        | otherwise = DQ $ Measure.scale (1/p) before
    afterCond
        | p == 1 = never
        | otherwise = DQ $ Measure.scale (1/(1-p)) after

{-----------------------------------------------------------------------------
    Operations
    Residual Time
------------------------------------------------------------------------------}
{- |
    For a renewal process,
    `residualTime` computes the time needed to wait until the next renewal,
    given that we have already waited an indefinite/random amount of time.

    The input argument is the probability distribution of completion times
    of the renewal process.
    The result is the residual (equilibrium) distribution of waiting times.

    The residual time is computed as
    $G(t) = (1/mu) * integral from 0 to t of (1 - F(s)) ds$.
    See also
    <https://en.wikipedia.org/wiki/Residual_time#Probability_distribution_of_the_residual_time>.

    Note that the input distribution must be a valid probability distribution 
    (i.e. total probability = 1), otherwise the result is `Nothing`.
    
    If it has zero mean, then it must be a Dirac distribution at 0; considering limits,
    the residual time is then also a Dirac distribution at 0.
-}
residualTime :: DQ -> Maybe DQ
residualTime f
    | success /= 1 = Nothing
    | mu == 0      = Just f
    | otherwise    = unsafeFromPositiveMeasure <$> Measure.fromDistribution g
        where
            (success, moments') = moments f
            mu = mean moments'
            df = distribution f  -- unpack the DQ to the piecewise polynomial representation
            oneMinusF =
                Piecewise.fromAscPieces
                $ oneFromZero
                $ Piecewise.toAscPieces
                $ Piecewise.mapPieces (Poly.constant 1 -) df
            g = Piecewise.mapPieces (Poly.scale (1/mu)) $ integral oneMinusF

-- a Piecewise is zero before its first piece, but 1 - F is 1 on [0, x1),
-- so add that piece explicitly if necessary
oneFromZero :: (Num a, Ord a) => [(a, Poly a)] -> [(a, Poly a)]
oneFromZero ps@((x, _) : _)
    | x > 0 = (0, Poly.constant 1) : ps
oneFromZero ps = ps