module ShiftedSignal where

{-
  Some of the basic routines are also present
  in NumericPrelude.LaurentPolynomial.
  But switching to NumericPrelude we risk usability in Hugs and Lhs2TeX.
-}

import qualified Signal

import qualified PowerSeries
import qualified Polynomial

import Useful(nest)
import qualified Useful
import qualified Data.List as List

import Signal ((^))
import Prelude hiding ((^))

import qualified Prelude

data T a = Cons {start :: Int, signal :: [a]}
   deriving (Show, Eq)

{-* Basic Operations -}

empty :: T a
empty = fromList []

fromScalar :: a -> T a
fromScalar x = fromList [x]

(!) :: Num a => T a -> Int -> a
(!) (Cons xt x) n =
   if n<xt
     then 0
     else head (drop (n-xt) x ++ [0])

fromList :: [a] -> T a
fromList = Cons 0

bounds :: T a -> (Int, Int)
bounds (Cons xt x) = (xt, xt + length x - 1)

translate :: Int -> T a -> T a
translate t (Cons xt x) = Cons (xt+t) x

amplify :: Num a => a -> T a -> T a
amplify k (Cons xt x) = Cons xt (Signal.amplify k x)

neg :: Num a => T a -> T a
neg (Cons xt x) = Cons xt (map negate x)

raise :: Num a => a -> T a -> T a
raise k (Cons xt x) = Cons xt (map (k+) x)

{- |
  Two signals may be different in structure but represent the same infinite signal.
  This function checks whether two signals represent the same infinite signal.
-}
equivalent :: Num a => T a -> T a -> Bool
equivalent xs ys =
   let (Cons xt x, Cons yt y) =
          if start xs <= start ys
            then (xs,ys)
            else (ys,xs)
       recurse (a:as) (b:bs) = a == b && recurse as bs
       recurse []     bs     = all (0==) bs
       recurse as     []     = all (0==) as
   in  recurse x (replicate (yt-xt) 0 ++ y)

{- |
  Check whether a signal is an impulse at time zero.
-}
isImpulse :: Num a => T a -> Bool
isImpulse (Cons xt x) =
   and (zipWith (\i xi -> xi==0 || i==0) [xt..] x)

superpose :: Num a => T a -> T a -> T a
superpose (Cons _ [])  y          = y
superpose  x          (Cons _ []) = x
superpose (Cons xt x) (Cons yt y) =
   if xt < yt
     then Cons xt (Signal.delaySuperpose (yt-xt) x y)
     else Cons yt (Signal.delaySuperpose (xt-yt) y x)

convolve :: Num a => T a -> T a -> T a
convolve (Cons xt x) (Cons yt y) =
   Cons (xt+yt) (Signal.convolve x y)

alternate :: Num a => T a -> T a
alternate (Cons xt x) =
   Cons xt (zipWith id (drop (xt `mod` 2) (cycle [id,negate])) x)

reverse :: T a -> T a
reverse (Cons xt x) =
   Cons (1 - xt - length x) (List.reverse x)

upSample :: Num a => Int -> T a -> T a
upSample k (Cons xt x) = Cons (xt*k) (Signal.upSample k x)

downSample :: Int -> T a -> T a
downSample k (Cons xt x) =
   {- by negating xt we can benefit
      of the floor rounding of div,
      thus we round until the next multiple of k -}
   let (xtd, xtm) = (-xt) `divMod` k
   in  Cons (-xtd) (Signal.downSample k (drop xtm x))


-- for sliceVert, see MathObj.LaurentPolynomial

sliceHoriz :: Int -> T a -> [T a]
sliceHoriz k (Cons xt x) =
   let (xtd, xtm) = (-xt) `divMod` k
       (x0s, x1s) = splitAt xtm (Useful.sliceHoriz k x)
   in  map (Cons (-xtd)  ) x1s ++
       map (Cons (-xtd-1)) x0s

{- the inverse of sliceHoriz, inefficient but comprehensible -}
interleave :: Num a => [T a] -> T a
interleave xs =
   foldl1 superpose
          (zipWith translate
                   [0..]
                   (map (upSample (length xs)) xs))

propInterleave :: Num a => Int -> T a -> Bool
propInterleave n x =
   interleave (sliceHoriz n x) == x

refine :: (Num a) => Int -> T a -> T a -> T a
refine k mask = convolve mask . upSample k

multiRefine :: (Num a) => Int -> Int -> T a -> T a -> T a
multiRefine n k mask = nest n (refine k mask)




{-* Lifting scheme -}

convolveAccumulate2 :: Num a => T a -> T a -> T a -> T a
convolveAccumulate2 s x =
   superpose (convolve (upSample 2 s) x)

liftingStep2 :: Num a => T a -> (T a, T a) -> (T a, T a)
liftingStep2 s (x,y) =
   (convolveAccumulate2 s x y, x)

lifting2 :: Num a => [T a] -> (T a, T a) -> (T a, T a)
lifting2 = flip (foldl (flip liftingStep2))

reverseLifting2 :: Num a => [T a] -> [T a]
reverseLifting2 = map neg . List.reverse


propReverseLifting2 :: Num a => [T a] -> (T a, T a) -> Bool
propReverseLifting2 ss fb =
   lifting2 (reverseLifting2 ss) (lifting2 ss fb) == fb


convolveAccumulate :: Num a => T a -> T a -> T a -> T a
convolveAccumulate s x =
   superpose (convolve s x)

noUpsampleLiftingStep2 :: Num a => T a -> (T a, T a) -> (T a, T a)
noUpsampleLiftingStep2 s (x,y) =
   (convolveAccumulate s x y, x)

noUpsampleLifting2 :: Num a => [T a] -> (T a, T a) -> (T a, T a)
noUpsampleLifting2 = flip (foldl (flip noUpsampleLiftingStep2))




{-* Biorthogonal CDF and orthogonal Daubechies wavelets -}


{-** Filters -}

lazyWavelet :: Num a => (T a, T a)
lazyWavelet = (Cons 0 [1], Cons 1 [1])

{- coefficients of orthogonal Daubechies generator masks
   without smoothness factors -}
daubechiesGeneratorPureList :: Floating a => Int -> [a]
daubechiesGeneratorPureList 0 = [1]
daubechiesGeneratorPureList 1 = [1]
daubechiesGeneratorPureList 2 =
   let w = pi/12
   in  map (sqrt 2 *) [cos w, - sin w]
daubechiesGeneratorPureList 4 =
   map ((8 * sqrt 2) *)
       [2.3037781330694596e-1,
       -2.066646826819176e-1,
        7.52726188069721e-2, 
       -1.05974017849973e-2]
daubechiesGeneratorPureList _ =
   error "daubechiesGeneratorPureList: order yet unsupported"

daubechiesGeneratorMask :: Floating a => Int -> T a
daubechiesGeneratorMask n =
   let mask = fromList (daubechiesGeneratorPureList n)
   in  nest n (convolve (fromList [1/2,1/2])) mask

orthogonalLowPassToHighPass,
   orthogonalHighPassToLowPass :: Num a => T a -> T a
orthogonalLowPassToHighPass =
   translate 1 . ShiftedSignal.reverse . alternate

orthogonalHighPassToLowPass =
   alternate . ShiftedSignal.reverse . translate (-1)

propOrthogonalLowHighPass :: Num a => T a -> Bool
propOrthogonalLowHighPass x =
   orthogonalHighPassToLowPass (orthogonalLowPassToHighPass x) == x  &&
   orthogonalLowPassToHighPass (orthogonalHighPassToLowPass x) == x


translationPrimalToDual, translationDualToPrimal :: Int
translationPrimalToDual =  1
translationDualToPrimal = -1

crossLowPassToHighPass, crossHighPassToLowPass ::
   Num a => Int -> T a -> T a
crossLowPassToHighPass translation =
   translate translation . alternate

crossHighPassToLowPass translation =
   alternate . translate translation

biorthogonalComplement ::
   Num a => Int -> (T a, T a) -> (T a, T a)
biorthogonalComplement translation (lp,hp) =
   (crossHighPassToLowPass translation hp,
    crossLowPassToHighPass translation lp)

primalToDual, dualToPrimal :: Num a => (T a, T a) -> (T a, T a)
primalToDual = biorthogonalComplement translationPrimalToDual
dualToPrimal = biorthogonalComplement translationDualToPrimal




{- Cf. Modula3's DaubechiesWavelet.FilterPureAbsSqr -}
daubechiesFilterPureAbsSqr :: Fractional a => Int -> T a
daubechiesFilterPureAbsSqr n =
   alternate (binomialWaveletMask n)

binomialWaveletMask :: Fractional a => Int -> T a
binomialWaveletMask n =
   (if mod n 2 == 0 then id else neg)
      (binomialWaveletMaskPure n)

binomialWaveletMaskNeg :: Fractional a => Int -> T a
binomialWaveletMaskNeg n =
   (if mod n 2 == 0 then neg else id)
      (binomialWaveletMaskPure n)

binomialWaveletMaskPure :: Fractional a => Int -> T a
binomialWaveletMaskPure n =
   evalSignalPolynomial
      (map fromInteger (take n (integrationPower n))) hatMask

hatMask :: Fractional a => T a
hatMask = Cons (-1) [1/4, 1/2, 1/4]

evalSignalPolynomial :: Num a => [a] -> T a -> T a
evalSignalPolynomial series w =
   Polynomial.hornerRing
      empty (superpose . fromScalar) (convolve w) series



{- Compute the series power (1-x)^-n,
   in contrast to that Polynomial.powerSeries only works for fractional types. -}
integrationPower, integrationPower' :: Int -> [Integer]
integrationPower n =
   let nl = fromIntegral n
   in  scanl (\x k -> x*(k+nl) `div` (k+1)) 1 [(0::Integer)..]

integrationPower' n =
   nest n (scanl1 (+)) (1 : repeat 0)

propIntegrationPower :: Int -> Int -> Bool
propIntegrationPower n m =
   take m (integrationPower n) == take m (integrationPower' n)


generatorMask :: Fractional a => Int -> T a
generatorMask n =
   nest n (convolve (Cons 0 [1/2, 1/2])) (Cons 0 [1])

waveletMask :: Fractional a => Int -> Int -> T a
waveletMask n m =
   let mask = nest m (convolve (Cons 0 [1/2, -1/2]))
                     (waveletMaskNoVan n m)
   in  translate (1 - cdfBaseOrder n m) mask

{- normalised such that (sum (signal (alternate (waveletMask' n m))) == 1) -}
waveletMask' :: Fractional a => Int -> Int -> T a
waveletMask' n m =
   let mask = nest m (convolve (Cons 0 [1/2, -1/2]))
                     (waveletMaskNoVan' n m)
   in  translate (1 - cdfBaseOrder n m) mask

{- The amplification asserts that lifting leads to the lazy wavelet 
   normalised to (1), (0,1). -}
waveletMaskNoVan :: Fractional a => Int -> Int -> T a
waveletMaskNoVan n m =
   amplify 2 (binomialWaveletMask (cdfBaseOrder n m))

waveletMaskNoVan' :: Fractional a => Int -> Int -> T a
waveletMaskNoVan' n m =
   binomialWaveletMaskNeg (cdfBaseOrder n m)

cdfBaseOrder :: Int -> Int -> Int
cdfBaseOrder n m =
   let (mn2, modu) = (m+n) `divMod` 2
   in  if modu==0
       then mn2
       else error ("generator order " ++ show n ++
                   " and wavelet order " ++ show m ++
                   " have different parity")

{- Check perfect reconstruction property
   with the determinant of a polyphase matrix. -}
reconstructionPP :: Num a => T a -> T a -> T a
reconstructionPP h g =
   let [he,ho] = sliceHoriz 2 h
       [ge,go] = sliceHoriz 2 g
   in  superpose (convolve he go)
                 (neg (convolve ge ho))

cdfReconstructionPP :: Fractional a => Int -> Int -> T a
cdfReconstructionPP n m =
   reconstructionPP (generatorMask n) (waveletMask n m)

propCDFReconstructionPP :: Int -> Int -> Bool
propCDFReconstructionPP n m =
   case filter (0/=) (signal (cdfReconstructionPP n m :: T Rational)) of
      [_] -> True
      _   -> False

{- Check perfect reconstruction property
   with the determinant of a modulation matrix. -}
reconstructionMM :: Num a => T a -> T a -> T a
reconstructionMM h g =
   superpose (convolve (alternate h) g)
             (neg (convolve h (alternate g)))

cdfReconstructionMM :: Fractional a => Int -> Int -> T a
cdfReconstructionMM n m =
   reconstructionMM (generatorMask n) (waveletMask n m)

propCDFReconstructionMM :: Int -> Int -> Bool
propCDFReconstructionMM n m =
   case filter (0/=) (signal (cdfReconstructionMM n m :: T Rational)) of
      [_] -> True
      _   -> False



{-** Explicit lifting decomposition of CDF -}

{-|
  Lifting sequence for the CDF-n,0 filter bank.
-}
cdfLiftingSeq :: Fractional a => Int -> [T a]
cdfLiftingSeq n =
   if even n
     then cdfLiftingSeqEven n
     else cdfLiftingSeqOdd  n

cdfLiftingSeqEven :: Fractional a => Int -> [T a]
cdfLiftingSeqEven n =
   let ms = tail (map fromInteger (iterate (2+) 0))
       n2 = fromIntegral n ^ 2
       a = scanl (\ai m -> recip ((n2 - m^2) * ai))
                 (recip (fromIntegral n)) ms
       b = iterate (2+) 1
   in  take (div n 2)
            (zipWith amplify
                (zipWith (*) a b)
                (cycle [Cons 0 [1,1], Cons (-1) [1,1]]))

cdfLiftingSeqOdd :: Fractional a => Int -> [T a]
cdfLiftingSeqOdd n =
   let ms = map fromInteger (iterate (2+) 1)
       n2 = fromIntegral n ^ 2
       a = scanl (\ai m -> recip ((n2 - m^2) * ai))
                 (recip (fromIntegral n)) ms
       b = iterate (2+) 1
   in  take (div (n+1) 2)
            (zipWith amplify a
                (Cons 1 [1] : zipWith Cons (cycle [-1,0])
                                      (map (\bi -> [bi,bi+2]) b)))


{- because in every lifting step the filters are swapped
   the sign of the polyphase determinant changes. -}
propCDFLiftingEven :: Int -> Bool
propCDFLiftingEven n2 =
   (amplify (2^n2) (generatorMask (2*n2) :: T Rational),
    amplify (2^(n2-1)) (waveletMask (2*n2) 0)) ==
   lifting2 (cdfLiftingSeqEven (2*n2)) lazyWavelet

{- If the lifting steps are reversed,
   then two slices of the generator can be reconstructed
   without computing the highpass filter. -}
propCDFLiftingEvenSep :: Int -> Bool
propCDFLiftingEvenSep n2 =
   let n = 2*n2
   in  noUpsampleLifting2
         (List.reverse (cdfLiftingSeq n))
         (fromScalar 1, fromScalar (0::Rational)) ==
       (\[x,y] -> (x,y))
         (sliceHoriz 2 (amplify (2^n2)
              (translate (mod n2 2 - n2) (generatorMask n :: T Rational))))

cdfLiftingReconstruct :: Fractional a => [T a]
cdfLiftingReconstruct =
   let w    = Cons (-1) [1,2,1]
       sigPoly =
              foldr (\c x -> superpose (fromScalar c) (convolve w x))
                    (fromList [])
       n  = 1.375
       zs = map sigPoly [
               [1],
               [n-2, 1],
               [(n-2)*(n-4), 3*(n-4), 3],
               [(n-2)*(n-4)*(n-6), 6*(n-4)*(n-6), 15*(n-6), 15]
            ]
       reconNegMM g h =
          superpose (convolve (alternate h) g)
                    (convolve h (alternate g))
   in  zipWith reconNegMM zs (tail zs)


{- the wavelet masks seems to have alternating signs
   such that (waveletMask n m) with the same n but different m
   can not always be lifted into each other. -}
{-| For n and m compute the lifting filter which lifts
    waveletMask n (mod m 2) to waveletMask n m. -}
lastLiftingDiff :: Fractional a => Int -> Int -> T a
lastLiftingDiff n m =
   waveletMask n m `superpose`
   amplify ((-1) ^ (1 + div m 2)) (waveletMask n (mod m 2))

lastLiftingStep0Series :: Fractional a => Int -> Int -> [a]
lastLiftingStep0Series n m =
   signal (lastLiftingDiff n m) `PowerSeries.divide` signal (generatorMask n)

propLastLiftingStep0Finite :: Int -> Int -> Int -> Bool
propLastLiftingStep0Finite n m k =
   checkFiniteSeries (2*m-1) k (lastLiftingStep0Series n m :: [Rational])

lastLiftingStep0 :: Fractional a => Int -> Int -> T a
lastLiftingStep0 n m =
   amplify ((-1) ^ div n 2)
           (Cons (1-m) (take (2*m-1) (lastLiftingStep0Series n m)))

propLastLiftingStep0Upsampled :: Int -> Int -> Bool
propLastLiftingStep0Upsampled n m =
   checkUpsampled (lastLiftingStep0 n m :: T Rational)





lastLiftingStep1Series :: Int -> Int -> [Integer]
lastLiftingStep1Series n m =
   let nu  = (n + mod n 2) `div` 2
       mn2 = cdfBaseOrder n m
       wp0 = drop nu  (integrationPower nu)
       wp1 = drop mn2 (integrationPower mn2)
   in  wp0 `Polynomial.sub`
          nest (mn2-nu) (0:)
               (nest (mn2-nu) (Polynomial.mul [1,-1]) wp1)

lastLiftingStep1Polynomial :: Int -> Int -> [Integer]
lastLiftingStep1Polynomial n m =
   take (m - mod n 2) (lastLiftingStep1Series n m)

propLastLiftingStep1Finite :: Int -> Int -> Int -> Bool
propLastLiftingStep1Finite n m k =
   checkFiniteSeries (m - mod n 2) k (lastLiftingStep1Series n m)

lastLiftingStep1 :: Fractional a => Int -> Int -> T a
lastLiftingStep1 n m =
   let x = amplify ((-1) ^ div (n+1) 2)
              (evalSignalPolynomial (map fromIntegral
                 (lastLiftingStep1Polynomial n m)) hatMask)
   in  if even n
         then x
         else convolve (Cons (-1) [-1/4,0,1/4]) x

propLastLiftingStep1Upsampled :: Int -> Int -> Bool
propLastLiftingStep1Upsampled n m =
   checkUpsampled (lastLiftingStep1 n m :: T Rational)



propLastLiftingStep :: Int -> Int -> Bool
propLastLiftingStep n m =
   lastLiftingStep0 n m == (lastLiftingStep1 n m :: T Rational)




{- Check if a series aborts at position 'n'.
   Of course it can only check a finite number 'k' of subsequent coefficients
   for being zero. -}
checkFiniteSeries :: (Eq a, Num a) => Int -> Int -> [a] -> Bool
checkFiniteSeries n k x  =  take k (drop n x) == replicate k 0

{- check if a signal has zeros at every even index
   (This differs from purely upsampled signals where every odd index is zero.) -}
checkUpsampled :: (Eq a, Num a) => T a -> Bool
checkUpsampled  =  all (0==) . signal . downSample 2


{- list of possible orders for CDF filter banks -}
cdfOrders :: [(Int,Int)]
cdfOrders =
   concatMap (\sum' -> map (\n -> (n,sum'-n)) [1 .. (sum'-1)]) [2,4 ..]


{-
  GNUPlot.plotPath [] (signal (nest 4 (refine 2 (generatorMask 4)) (waveletMask 4 10)))
-}


{-
  Using the power series for fractional powers it should be possible
  to generate
   1. the lowpass of fractional CDF banks
   2. the highpass by computing the co-recursive polynomial division
      which is given in the CDF paper.
      In contrast to the integer case this division should not abort.
  The problem is that the right hand sight contains $\alternate{w}$
  and the left hand side $w$.
  We must replace $w$ by $1-w$ in the series
  which requires the computation of series values. :-(
  Maybe we should determine the solution in terms of $v$ with $v = w - 1/2$.
  This doesn't lead to a solution because then the pure power $w$
  does no longer exist which shifts the first coefficient of $\alternate{g}$ away.
-}