{-# LANGUAGE FlexibleInstances, FlexibleContexts #-}
import Test.QuickCheck
import Data.Array.CArray
import Data.Complex
import Math.FFT
import Foreign.Storable
import Text.Printf
import System.Environment (getArgs)
import System.IO
import System.Random

instance Arbitrary (Complex Double) where
    arbitrary = do
        r <- arbitrary
        i <- arbitrary
        return $ r :+ i
    coarbitrary = error "no coarbitrary for Complex"

instance (IArray CArray e, Arbitrary e) => Arbitrary (CArray Int e) where
    arbitrary = do
        u <- choose (1,100)
        es <- vector (u+1)
        return $ listArray (0,u) es
    coarbitrary = error "no coarbitrary for CArray"

instance (IArray CArray e, Arbitrary e) => Arbitrary (CArray (Int,Int) e) where
    arbitrary = do
        u0 <- choose (1,30)
        u1 <- choose (1,30)
        es <- vector ((u0 + 1) * (u1 + 1))
        return $ listArray ((0,0),(u0,u1)) es
    coarbitrary = error "no coarbitrary for CArray"

instance (IArray CArray e, Arbitrary e) => Arbitrary (CArray (Int,Int,Int) e) where
    arbitrary = do
        u0 <- choose (1,20)
        u1 <- choose (1,20)
        u2 <- choose (1,20)
        es <- vector ((u0 + 1) * (u1 + 1) * (u2 + 1))
        return $ listArray ((0,0,0),(u0,u1,u2)) es
    coarbitrary = error "no coarbitrary for CArray"


-- about :: (Ix i, FFTWFloat e) => CArray i e -> CArray i e -> Bool
about x y = small $ normSup (liftArray2 (-) x y) / (1 + normSup (liftArray2 (+) x y))
    where small a = a < 1e-15

partAbout a b = about a (slice ba ba b)
    where ba = bounds a

aboutIdem f x = f x `about` x

prop_dft     = aboutIdem $ idft . dft
prop_dftRC a = aboutIdem ((if odd (shape a !! 0) then dftCRO else dftCR) . dftRC) a
prop_dftRC_dft a = partAbout (dftRC a) (dft . amap (:+0) $ a)
prop_dht_idem a = aboutIdem (amap (/ fromIntegral (shape a !! 0)) . dht . dht) a


prop_dft2     = aboutIdem $ idft . dft
prop_dft22    = aboutIdem $ idftN [0,1] . dftN [0,1]
prop_dft22'   = aboutIdem $ idftN [1,0] . dftN [1,0]

prop_dftRC2 a = aboutIdem ((if odd (shape a !! 0) then dftCRO else dftCR) . dftRC) a
prop_dftRC_dft2 a = partAbout (dftRC a) (dft . amap (:+0) $ a)
prop_dftRC_dft22 a = partAbout (dftRCN [0,1] a) (dftN [0,1] . amap (:+0) $ a)
prop_dht_idem2 a = aboutIdem (amap (/ fromIntegral (shape a !! 0)) . dht . dht) a

prop_dft3     = aboutIdem $ idft . dft
prop_dft32    = aboutIdem $ idftN [0,1] . dftN [0,1]
prop_dft32'   = aboutIdem $ idftN [1,0] . dftN [1,0]
prop_dft33    = aboutIdem $ idftN [0,1,2] . dftN [0,1,2]
prop_dft33'   = aboutIdem $ idftN [0,2,1] . dftN [0,2,1]
prop_dft33''  = aboutIdem $ idftN [2,0,1] . dftN [2,0,1]

c_tests :: [(String, CArray Int (Complex Double) -> Bool)]
c_tests = [ ("dft idem 1D" , prop_dft)
          ]

c_tests2 :: [(String, CArray (Int,Int) (Complex Double) -> Bool)]
c_tests2 = [ ("dft idem 2D" , prop_dft2)
           , ("dft idem 2D/2" , prop_dft22)
           , ("dft idem 2D/2'" , prop_dft22')
          ]

c_tests3 :: [(String, CArray (Int,Int,Int) (Complex Double) -> Bool)]
c_tests3 = [ ("dft idem 3D" , prop_dft3)
           , ("dft idem 3D/2" , prop_dft32)
           , ("dft idem 3D/2'" , prop_dft32')
           , ("dft idem 3D/3" , prop_dft33)
           , ("dft idem 3D/3'" , prop_dft33')
           , ("dft idem 3D/3''" , prop_dft33'')
          ]

r_tests :: [(String, CArray Int Double -> Bool)]
r_tests = [ ("dftRC/CR idem  1D" , prop_dftRC)
          , ("dftRC dft 1D" , prop_dftRC_dft)
          , ("dht idem 1D" , prop_dht_idem)
          ]

r_tests2 :: [(String, CArray (Int,Int) Double -> Bool)]
r_tests2 = [ ("dftRC/CR idem  2D" , prop_dftRC2)
           , ("dftRC dft 2D" , prop_dftRC_dft2)
           , ("dftRC dft 2D/2" , prop_dftRC_dft22)
           , ("dht idem 2D" , prop_dht_idem2)
          ]

main = do
    x <- getArgs
    let n = if null x then 20 else read . head $ x
        conf = Config { configMaxTest = n
                      , configMaxFail = 1000
                      , configSize = (+ 3) . (`div` 2)
                      , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s]
                      }
    mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) c_tests
    mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) r_tests
    mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) c_tests2
    mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) r_tests2
    mapM_ (\(s,a) -> printf "%-25s: " s >> check conf a) c_tests3