{-# LANGUAGE RecordWildCards #-}

module System.Hardware.Acquisition.Adaptor.Filter
	( adaptor
	, movingAverage
	, geometricAverage
	, exponentialAverage
	) where

import Prelude
import Control.Monad.Error
import Data.Maybe
import Data.Traversable
import Data.Foldable (toList)
import Data.IORef
import Data.Sequence (Seq, (<|))
import qualified Data.Sequence as Seq

import System.Hardware.Acquisition as Acquisition

{-|
	IIR (infinite impulse response) filter, modelled by:

		&#x2211; out[n] &#xd7; a[n] = &#x2211; in[n] &#xd7; b[n]

	where in[0] is the most recent input, and out[0] is the current output.

	To model FIR (finite impulse response) filters, simply give an empty
	list for the a[_] feedback coefficients.
-}
adaptor
	:: MonadIO m
	=> [Value] {-^ a[_], feedback filter coefficients -}
	-> [Value] {-^ b[_], feedforward filter coefficients -}
	-> Adaptor m
adaptor [] bs c = adaptor [1.0] bs c
adaptor as@(a0:at) bs c = liftIO $ do
	let backwardOrder = length at
	let forwardOrder = length bs

	state <- newIORef (Whatever, Seq.empty, Seq.empty)
	let convolute :: MonadIO m => Sample -> m Sample
	    convolute s = liftIO (atomicModifyIORef state modify) where
		modify (m, xs, ys) = (,)
				(measurement s, xs', ys')
				s { value = y0 } where
			(xs1, ys1) = case measurement s == m of
				True -> (xs, ys)
				False -> (Seq.empty, Seq.empty)
			xs' = Seq.take forwardOrder (value s <| xs1)
			y0 = (sum (zipWith (*) bs (toList xs'))
				+ sum (zipWith (*) at (toList ys1))) / a0
			ys' = Seq.take backwardOrder (y0 <| ys1)

	let getSample :: (Error e, MonadError e m, MonadIO m) => m Sample
	    getSample = Acquisition.getSample c >>= convolute

	let close :: MonadIO m => m ()
	    close = return ()

	return Channel {..}

{-| FIR: moving window average. -}
movingAverage
	:: MonadIO m
	=> Int {-^ Number of samples to average over -}
	-> Adaptor m
movingAverage n = adaptor [] (replicate n (1.0 / fromIntegral n))

{-|
	FIR: n-tap geometrically-decreasing weighted average, almost.
	Except that we ensure the coefficients sum up to unity.
-}
geometricAverage
	:: MonadIO m
	=> Value
	-> Int
	-> Adaptor m
geometricAverage base n = adaptor [] coeffs where
	decay = replicate n (1.0 - 1.0 / base)
	unit = recip . sum $ scanl (*) 1.0 decay
	coeffs = scanl (*) unit decay

{-| IIR: expotentially-decreasing weighted average. -}
exponentialAverage
	:: MonadIO m
	=> Value
	-> Adaptor m
exponentialAverage weight = adaptor [1.0, 1.0 - weight] [weight]