{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

module Test.Prelude.Stencil (

  test_stencil,

) where

import Prelude                                                  as P
import Data.Label
import Data.Maybe
import Data.Typeable
import Test.QuickCheck                                          ( Arbitrary )
import Test.HUnit                                               ( (@?=) )
import Test.Framework
import Test.Framework.Providers.QuickCheck2
import Test.Framework.Providers.HUnit

import Config
import QuickCheck.Arbitrary.Array                               ()

import Data.Array.Accelerate                                    as A
import Data.Array.Accelerate.IO                                 as A
import Data.Array.Accelerate.Examples.Internal                  as A
import Data.Array.Unboxed                                       as IArray hiding ( Array )
import qualified Data.Array.IArray                              as IArray


-- TODO:
--
--  * Tests for boundary conditions: Mirror and Wrap
--  * Higher dimensional stencils
--

--
-- Stencil ---------------------------------------------------------------------
--

test_stencil :: Backend -> Config -> Test
test_stencil backend opt = testGroup "stencil" $ catMaybes
  [ testElt configInt8   (undefined :: Int8)
  , testElt configInt16  (undefined :: Int16)
  , testElt configInt32  (undefined :: Int32)
  , testElt configInt64  (undefined :: Int64)
  , testElt configWord8  (undefined :: Word8)
  , testElt configWord16 (undefined :: Word16)
  , testElt configWord32 (undefined :: Word32)
  , testElt configWord64 (undefined :: Word64)
  , testElt configFloat  (undefined :: Float)
  , testElt configDouble (undefined :: Double)
  , testBoundary
  ]
  where
    testElt :: forall a. (P.Num a, A.Num a, Similar a, Arbitrary a, IArray UArray a)
            => (Config :-> Bool)
            -> a
            -> Maybe Test
    testElt ok _
      | P.not (get ok opt)      = Nothing
      | otherwise               = Just $ testGroup (show (typeOf (undefined :: a)))
          [ testProperty "1D"                   (test_stencil1D  :: Array DIM1 a -> Property)
          , testProperty "2D 3x3 dense"         (test_stencil2D1 :: Array DIM2 a -> Property)
          , testProperty "2D 3x3 cross"         (test_stencil2D2 :: Array DIM2 a -> Property)
          , testProperty "2D non-symmetric"     (test_stencil2D3 :: Array DIM2 (a,a) -> Property)
          ]

    -- 1D Stencil
    --
    test_stencil1D :: (P.Num a, A.Num a, Similar a, IArray UArray a) => Vector a -> Property
    test_stencil1D vec = toList (acc vec) ~?= elems (ref (toIArray vec))
      where
        pattern (x,y,z) = x + z - 2 * y

        acc xs = run1 backend (stencil pattern clamp) xs

        ref :: (P.Num e, IArray UArray e) => UArray Int e -> UArray Int e
        ref xs =
          let (minx,maxx)   = bounds xs
              clamp x       = Right (minx `P.max` x `P.min` maxx)
          in
          stencil1DRef pattern clamp xs

    -- 2D Stencil
    --
    test_stencil2D1 :: (P.Num a, A.Num a, Similar a, IArray UArray a) => Array DIM2 a -> Property
    test_stencil2D1 vec = toList (acc vec) ~?= elems (ref (toIArray vec))
      where
        pattern ( (t1, t2, t3)
                , (l , m,  r )
                , (b1, b2, b3)
                )
                = (t1 + t2 + t3 - l + 4*m - r - b1 - b2 - b3)

        acc xs = run1 backend (stencil pattern (function (const 0))) xs

        ref :: (P.Num a, IArray UArray a) => UArray (Int,Int) a -> UArray (Int,Int) a
        ref xs =
          let
              sh                    = bounds xs
              constant ix
                | inRange sh ix     = Right ix
                | otherwise         = Left 0

          in
          stencil2DRef pattern constant xs


    test_stencil2D2 :: (P.Num a, A.Num a, Similar a, IArray UArray a) => Array DIM2 a -> Property
    test_stencil2D2 vec = toList (acc vec) ~?= elems (ref (toIArray vec))
      where
        pattern ( (_, t, _)
                , (l, m, r)
                , (_, b, _)
                )
                = (t + l + r + b - 4 * m)

        acc xs =
          let pattern' :: A.Num a => Stencil3x3 a -> Exp a
              pattern' = pattern
          in
          run1 backend (stencil pattern' clamp) xs

        ref :: (P.Num a, IArray UArray a) => UArray (Int,Int) a -> UArray (Int,Int) a
        ref xs =
          let ((minu,minv),(maxu,maxv)) = bounds xs
              clamp (u,v) = Right (minu `P.max` u `P.min` maxu
                                  ,minv `P.max` v `P.min` maxv)
          in
          stencil2DRef pattern clamp xs

    test_stencil2D3 :: (P.Num a, A.Num a, Similar a) => Array DIM2 (a,a) -> Property
    test_stencil2D3 vec = toList (acc vec) ~?= elems (ref (toIArray vec))
      where
        pattern :: forall a. A.Num a => Stencil3x3 (a,a) -> Exp a
        pattern ( (_, _, _) , (x, _, _) , (y, _, z))
          = let (x1,x2) = unlift x
                (y1,y2) = unlift y
                (z1,z2) = unlift z
            in
            x1 - y2 + y1 - z2 + z1 - x2

        pattern' ( (_, _, _) , (x, _, _) , (y, _, z))
          = let (x1,x2) = x
                (y1,y2) = y
                (z1,z2) = z
            in
            x1 - y2 + y1 - z2 + z1 - x2

        acc xs = run1 backend (stencil pattern (function (\_ -> constant (0,0)))) xs

        ref :: P.Num a => IArray.Array (Int,Int) (a,a) -> IArray.Array (Int,Int) a
        ref xs =
          let
              sh                    = bounds xs
              constant ix
                | inRange sh ix     = Right ix
                | otherwise         = Left (0,0)
          in
          stencil2DRef pattern' constant xs

    -- If the constant boundary is not properly implemented,
    -- then this will lead to a segmentation fault.
    testBoundary :: Maybe Test
    testBoundary = Just . testCase "boundary segfault" $ do
      let f ((x,_,_,_,_),_,_,_,_) = x
          b = function (const 0)
          s = stencil (f::Stencil5x5 Int -> Exp Int) b (A.fill (lift (Z:.1:.1000000 :: DIM2)) (0::Exp Int))
          a = run backend s
      indexArray a (Z:.0:.0) @?= 0

--
-- Reference implementation
--
stencil1DRef
    :: (IArray array a, IArray array b)
    => ((a,a,a) -> b)
    -> (Int -> Either a Int)
    -> array Int a
    -> array Int b
stencil1DRef pattern boundary xs =
  let
      indexAt ix = case boundary ix of
        Left e          -> e
        Right ix'       -> xs IArray.! ix'

      f ix = let x = indexAt (ix-1)
                 y = indexAt ix
                 z = indexAt (ix+1)
             in
             pattern (x,y,z)
   in
  array (bounds xs) [(ix, f ix) | ix <- indices xs]


stencil2DRef :: (IArray array a, IArray array b)
    => (((a,a,a), (a,a,a), (a,a,a)) -> b)
    -> ((Int,Int) -> Either a (Int,Int))
    -> array (Int,Int) a
    -> array (Int,Int) b
stencil2DRef pattern boundary xs =
  let
      indexAt ix = case boundary ix of
        Left e          -> e
        Right ix'       -> xs IArray.! ix'

      f (y,x) = let t1 = indexAt (y-1,x-1)
                    t2 = indexAt (y-1,x  )
                    t3 = indexAt (y-1,x+1)
                    l  = indexAt (y,  x-1)
                    m  = indexAt (y,  x  )
                    r  = indexAt (y,  x+1)
                    b1 = indexAt (y+1,x-1)
                    b2 = indexAt (y+1,x  )
                    b3 = indexAt (y+1,x+1)
                in
                pattern ((t1,t2,t3), (l,m,r), (b1,b2,b3))
  in
  array (bounds xs) [(ix, f ix) | ix <- indices xs]