{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE ViewPatterns      #-}
module Tests.Transform
    (
      tests
    ) where

import Data.Bits ((.&.), shiftL)
import Data.Complex (Complex((:+)))
import Data.Functor ((<$>))
import Numeric.Sum (kbn, sumVector)
import Statistics.Function (within)
import Statistics.Transform (CD, dct, fft, idct, ifft)
import Test.Framework (Test, testGroup)
import Test.Framework.Providers.QuickCheck2 (testProperty)
import Test.QuickCheck (Positive(..), Arbitrary(..), Gen, choose, vectorOf, counterexample)
import Test.QuickCheck.Property (Property(..))
import Tests.Helpers (testAssertion)
import Text.Printf (printf)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U


tests :: Test
tests = testGroup "fft" [
          testProperty "t_impulse"        t_impulse
        , testProperty "t_impulse_offset" t_impulse_offset
        , testProperty "ifft . fft = id"  (t_fftInverse $ ifft . fft)
        , testProperty "fft . ifft = id"  (t_fftInverse $ fft . ifft)
        , testProperty "idct . dct = id [up to scale]"
            (t_fftInverse (\v -> U.map (/ (2 * fromIntegral (U.length v))) $ idct $ dct v))
        , testProperty "dct . idct = id [up to scale]"
            (t_fftInverse (\v -> U.map (/ (2 * fromIntegral (U.length v))) $ idct $ dct v))
          -- Exact small size DCT
          -- 1
        , testDCT [1] $ [2]
          -- 2
        , testDCT [1,0] $ map (*2) [1, cos (pi/4)   ]
        , testDCT [0,1] $ map (*2) [1, cos (3*pi/4) ]
          -- 4
        , testDCT [1,0,0,0] $ map (*2) [1, cos(  pi/8), cos( 2*pi/8), cos( 3*pi/8)]
        , testDCT [0,1,0,0] $ map (*2) [1, cos(3*pi/8), cos( 6*pi/8), cos( 9*pi/8)]
        , testDCT [0,0,1,0] $ map (*2) [1, cos(5*pi/8), cos(10*pi/8), cos(15*pi/8)]
        , testDCT [0,0,0,1] $ map (*2) [1, cos(7*pi/8), cos(14*pi/8), cos(21*pi/8)]
          -- Exact small size IDCT
          -- 1
        , testIDCT [1] [1]
          -- 2
        , testIDCT [1,0]            [1,         1          ]
        , testIDCT [0,1] $ map (*2) [cos(pi/4), cos(3*pi/4)]
          -- 4
        , testIDCT [1,0,0,0]            [1,            1,            1,            1            ]
        , testIDCT [0,1,0,0] $ map (*2) [cos(   pi/8), cos( 3*pi/8), cos( 5*pi/8), cos( 7*pi/8) ]
        , testIDCT [0,0,1,0] $ map (*2) [cos( 2*pi/8), cos( 6*pi/8), cos(10*pi/8), cos(14*pi/8) ]
        , testIDCT [0,0,0,1] $ map (*2) [cos( 3*pi/8), cos( 9*pi/8), cos(15*pi/8), cos(21*pi/8) ]
        ]

-- A single real-valued impulse at the beginning of an otherwise zero
-- vector should be replicated in every real component of the result,
-- and all the imaginary components should be zero.
t_impulse :: Double -> Positive Int -> Bool
t_impulse k (Positive m) = G.all (c_near i) (fft v)
  where v = i `G.cons` G.replicate (n-1) 0
        i = k :+ 0
        n = 1 `shiftL` (m .&. 6)

-- If a real-valued impulse is offset from the beginning of an
-- otherwise zero vector, the sum-of-squares of each component of the
-- result should equal the square of the impulse.
t_impulse_offset :: Double -> Positive Int -> Positive Int -> Bool
t_impulse_offset k (Positive x) (Positive m) = G.all ok (fft v)
  where v = G.concat [G.replicate xn 0, G.singleton i, G.replicate (n-xn-1) 0]
        ok (re :+ im) = within ulps (re*re + im*im) (k*k)
        i  = k :+ 0
        xn = x `rem` n
        n  = 1 `shiftL` (m .&. 6)

-- Test that (ifft . fft ≈ id)
--
-- Approximate equality here is tricky. Smaller values of vector tend
-- to have large relative error. Thus we should test that vectors as
-- whole are approximate equal.
t_fftInverse :: (HasNorm (U.Vector a), U.Unbox a, Num a, Show a, Arbitrary a)
             => (U.Vector a -> U.Vector a) -> Property
t_fftInverse roundtrip = MkProperty $ do
  x <- genFftVector
  let n  = G.length x
      x' = roundtrip x
      d  = G.zipWith (-) x x'
      nd = vectorNorm d
      nx = vectorNorm x
  unProperty
     $ counterexample "Original vector"
     $ counterexample (show x )
     $ counterexample "Transformed one"
     $ counterexample (show x')
     $ counterexample (printf "Length = %i" n)
     $ counterexample (printf "|x - x'| / |x| = %.6g" (nd / nx))
     $ nd <= 3e-14 * nx

-- Test discrete cosine transform
testDCT :: [Double] -> [Double] -> Test
testDCT (U.fromList -> vec) (U.fromList -> res)
  = testAssertion ("DCT test for " ++ show vec)
  $ vecEqual 3e-14 (dct vec) res

-- Test inverse discrete cosine transform
testIDCT :: [Double] -> [Double] -> Test
testIDCT (U.fromList -> vec) (U.fromList -> res)
  = testAssertion ("IDCT test for " ++ show vec)
  $ vecEqual 3e-14 (idct vec) res



----------------------------------------------------------------

-- With an error tolerance of 8 ULPs, a million QuickCheck tests are
-- likely to all succeed. With a tolerance of 7, we fail around the
-- half million mark.
ulps :: Int
ulps = 8

c_near :: CD -> CD -> Bool
c_near (a :+ b) (c :+ d) = within ulps a c && within ulps b d

-- Arbitrary vector for FFT od DCT
genFftVector :: (U.Unbox a, Arbitrary a) => Gen (U.Vector a)
genFftVector = do
  n <- (2^)  <$> choose (1,9::Int)    -- Size of vector
  G.fromList <$> vectorOf n arbitrary -- Vector to transform

-- Ad-hoc type class for calculation of vector norm
class HasNorm a where
  vectorNorm :: a -> Double

instance HasNorm (U.Vector Double) where
  vectorNorm = sqrt . sumVector kbn . U.map (\x -> x*x)

instance HasNorm (U.Vector CD) where
  vectorNorm = sqrt . sumVector kbn . U.map (\(x :+ y) -> x*x + y*y)

-- Approximate equality for vectors
vecEqual :: Double -> U.Vector Double -> U.Vector Double -> Bool
vecEqual ε v u
  = vectorNorm (U.zipWith (-) v u) < ε * vectorNorm v