    {-# LANGUAGE FlexibleContexts #-}
    -- |
    -- Module    : Statistics.Autocorrelation
    -- Copyright : (c) 2009 Bryan O'Sullivan
    -- License   : BSD3
    --
    -- Maintainer  : bos@serpentine.com
    -- Stability   : experimental
    -- Portability : portable
    --
    -- Functions for computing autocovariance and autocorrelation of a
    -- sample.

    module Statistics.Autocorrelation
        (
          autocovariance
        , autocorrelation
        ) where

    import Prelude hiding (sum)
    import Statistics.Function (square)
    import Statistics.Sample (mean)
    import Statistics.Sample.Internal (sum)
    import qualified Data.Vector.Generic as G

    -- | Compute the autocovariance of a sample, i.e. the covariance of
    -- the sample against a shifted version of itself.
    autocovariance :: (G.Vector v Double, G.Vector v Int) => v Double -> v Double
    autocovariance a = G.map f . G.enumFromTo 0 $ l-2
      where
        f k = sum (G.zipWith (*) (G.take (l-k) c) (G.slice k (l-k) c))
              / fromIntegral l
        c   = G.map (subtract (mean a)) a
        l   = G.length a

    -- | Compute the autocorrelation function of a sample, and the upper
    -- and lower bounds of confidence intervals for each element.
    --
    -- /Note/: The calculation of the 95% confidence interval assumes a
    -- stationary Gaussian process.
    autocorrelation :: (G.Vector v Double, G.Vector v Int) => v Double -> (v Double, v Double, v Double)
    autocorrelation a = (r, ci (-), ci (+))
      where
        r           = G.map (/ G.head c) c
          where c   = autocovariance a
        dllse       = G.map f . G.scanl1 (+) . G.map square $ r
          where f v = 1.96 * sqrt ((v * 2 + 1) / l)
        l           = fromIntegral (G.length a)
        ci f        = G.cons 1 . G.tail . G.map (f (-1/l)) $ dllse