module Armax (applyArmax) where
 
-- Some aliases.
type Signal a = [a]                 -- stream
type History a = [a]                -- finite time-reversed stream
type Coef a = [a]                   -- finite time-reversed FIR coefs
type Predictor a = History a -> a   -- causal predictor


-- This computes an ARMAX SISO model, where the MA reuses the delay
-- line (history) used for the AR output feedback.
armax  :: Fractional a => (Predictor a, Predictor a) -> History a -> [a] -> [a]
armax (ma, ar) = s 
    where s hs (x:xs) = y:ys 
              where h   = ar hs + x      -- ARX
                    y   = ma hs'         -- MA
                    hs' = h:hs           -- update history
                    ys  = s hs' xs       -- tail

-- Predictors can be constructed from FIR filter coefficients.
fir :: (Num a) => Coef a -> History a -> a
fir as bs = foldr (+) 0 $ zipWith (*) as bs


applyArmax :: Fractional a => (Coef a, Coef a) -> History a -> Signal a -> Signal a
applyArmax (num, den) = armax (fir num, fir $ diffeq den) where
    -- Convert denominator polynomial d(z) to normalized form where
    -- d(0)=1 to obtain fir coefficients for a difference equation.
    diffeq (d0:ds) = let scale = (-1) / d0 in map (* scale) ds


-- test = take 10 $ applyArmax ([1],[1,-0.9]) [0,0..] (1:[0,0..])