{-
See also: statistics, order-statistics, interpolation
-}
module Quantile where

import qualified Data.NonEmpty.Set as NESet
import qualified Data.NonEmpty.Class as NonEmptyC
import qualified Data.NonEmpty as NonEmpty
import Data.NonEmpty ((!:))

import qualified Algebra.RealRing as Real
import qualified Algebra.Field as Field
import NumericPrelude.Numeric
import NumericPrelude.Base
import Prelude ()


{- |
The function @discrete xs@ is the inverse cumulative distribution function
of the discrete distribution represented by @xs@.
It is a piecewise linear interpolation
between all numbers in @xs@ in ascending order.
In @discrete xs r@ the parameter @r@ must be between 0 and 1.
It is @discrete xs 0 == minimum xs@,
@discrete xs 1 == maximum xs@,
and @discrete xs 0 == median xs@.
-}
discrete :: (Field.C a, Real.C a) => NonEmpty.T [] a -> a -> a
discrete (NonEmpty.Cons x []) = const x
discrete xs =
   case NESet.fromList $ NonEmptyC.zip xs $ (0::Int) !: [1 ..] of
      set -> \r ->
         if r==one
           then fst $ NESet.findMax set
           else
            {-
            'Set.elemAt' would be faster than 'drop',
            but is only available since GHC-7.8.4.
            -}
            let (k, frac) = splitFraction $ fromIntegral (NESet.size set - 1) * r
            in  case drop k $ NonEmpty.flatten $ NESet.toAscList set of
                   (x0,_):(x1,_):_ -> x0*(one-frac) + x1*frac
                   _ -> error "Quantile.discrete: unexpected empty list"