{-# LANGUAGE GADTs               #-}
{-# LANGUAGE NoImplicitPrelude   #-}
{-# LANGUAGE PatternGuards       #-}
{-# LANGUAGE QuasiQuotes         #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE ViewPatterns        #-}
-- |
-- Module      : Data.Array.Accelerate.CUDA.CodeGen.Arithmetic
-- Copyright   : [2008..2014] Manuel M T Chakravarty, Gabriele Keller
--               [2009..2014] Trevor L. McDonell
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.CUDA.CodeGen.Arithmetic
  where

-- friends
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.CUDA.CodeGen.Base
import Data.Array.Accelerate.CUDA.CodeGen.Constant
import Data.Array.Accelerate.CUDA.CodeGen.Monad
import Data.Array.Accelerate.CUDA.CodeGen.Type

-- libraries
import Prelude                                          ( String, ($), (++), (-), undefined, otherwise )
import Data.Bits                                        ( finiteBitSize )
import Control.Monad.State.Strict
import Language.C
import Language.C.Quote.CUDA
import Foreign.Storable                                 ( sizeOf )


-- Operations from Num
-- ===================

add :: Exp -> Exp -> Exp
add x y = [cexp| $exp:x + $exp:y |]

sub :: Exp -> Exp -> Exp
sub x y = [cexp| $exp:x - $exp:y |]

mul :: Exp -> Exp -> Exp
mul x y = [cexp| $exp:x * $exp:y |]

negate :: Exp -> Exp
negate x = [cexp| - $exp:x |]

abs :: forall a. NumType a -> Exp -> Exp
abs (FloatingNumType t) x
  = mathf t "fabs" [x]

abs (IntegralNumType t) x
  | signedIntegralNum t
  , IntegralDict <- integralDict t
  = case sizeOf (undefined::a) of
      8 -> ccall "llabs" [x]
      _ -> ccall "abs" [x]

  | otherwise
  = x

signum :: NumType a -> Exp -> Gen Exp
signum (IntegralNumType t) x
  | IntegralDict <- integralDict t
  , unsignedIntegralNum t
  = return [cexp| $exp:x > $exp:(integral t 0) |]

  | IntegralDict <- integralDict t
  = do x' <- bind (typeOf t) x
       return [cexp| ($exp:x' > $exp:(integral t 0)) - ($exp:x' < $exp:(integral t 0)) |]

signum (FloatingNumType t) x
  | FloatingDict <- floatingDict t
  = do x' <- bind (typeOf t) x
       return [cexp| $exp:x' == $exp:(floating t 0)
                       ? $exp:(floating t 0)
                       : $exp:(mathf t "copysign" [floating t 1, x']) |]


-- Operators from Integral & Bits
-- ==============================

quot :: Exp -> Exp -> Exp
quot x y = [cexp| $exp:x / $exp:y |]

rem :: Exp -> Exp -> Exp
rem x y = [cexp| $exp:x % $exp:y |]

quotRem :: IntegralType a -> Exp -> Exp -> Gen (Exp,Exp)
quotRem (typeOf -> t') x y = do
  x' <- bind t' x
  y' <- bind t' y
  q  <- bind t' (x' `quot` y')
  r  <- bind t' (x' `sub` (y' `mul` q))
  return (q, r)

idiv :: IntegralType a -> Exp -> Exp -> Gen Exp
idiv t x y
  | unsignedIntegralNum t
  = return (x `quot` y)

  | IntegralDict <- integralDict t
  , zero         <- integral t 0
  , one          <- integral t 1
  = do
      x' <- bind (typeOf t) x
      y' <- bind (typeOf t) y
      return $
        cases [ ((x' `gt` zero) `land` (y' `lt` zero), ((x' `sub` one) `quot` y') `sub` one)
              , ((x' `lt` zero) `land` (y' `gt` zero), ((x' `add` one) `quot` y') `sub` one)
              ]
              (x' `quot` y')

mod :: IntegralType a -> Exp -> Exp -> Gen Exp
mod t x y
  | unsignedIntegralNum t
  = return (x `rem` y)

  | IntegralDict <- integralDict t
  , zero         <- integral t 0
  = do
       x' <- bind (typeOf t) x
       y' <- bind (typeOf t) y
       r  <- bind (typeOf t) (x' `rem` y')
       return $
         ((((x' `gt` zero) `land` (y' `lt` zero)) `lor` ((x' `lt` zero) `land` (y' `gt` zero)))
          ?: ( r `neq` zero ?: ( r `add` y', zero )
             , r ))

divMod :: IntegralType a -> Exp -> Exp -> Gen (Exp, Exp)
divMod t x y | IntegralDict <- integralDict t = do
  x'    <- bind (typeOf t) x
  y'    <- bind (typeOf t) y
  (q,r) <- quotRem t x' y'

  sr    <- signum (IntegralNumType t) r
  sy'   <- signum (IntegralNumType t) y'

  -- Somewhat awful way to inject an ifThenElse statement
  vd    <- lift fresh
  vm    <- lift fresh
  modify (\st -> st { localBindings = [citem| $ty:(typeOf t) $id:vd, $id:vm; |] : localBindings st })
  modify (\st -> st { localBindings = [citem| if ( $exp:(sr `eq` negate sy') ) {
                                                  $id:vd = $exp:(q `sub` integral t 1);
                                                  $id:vm = $exp:(r `add` y') ;
                                              } else {
                                                  $id:vd = $exp:q;
                                                  $id:vm = $exp:r;
                                              } |] : localBindings st })
  return ( cvar vd, cvar vm )


band :: Exp -> Exp -> Exp
band x y = [cexp| $exp:x & $exp:y |]

bor :: Exp -> Exp -> Exp
bor x y = [cexp| $exp:x | $exp:y |]

xor :: Exp -> Exp -> Exp
xor x y = [cexp| $exp:x ^ $exp:y |]

bnot :: Exp -> Exp
bnot x = [cexp| ~ $exp:x |]

shiftL :: Exp -> Exp -> Exp
shiftL x i = [cexp| $exp:x << $exp:i |]

-- Arithmetic right shift (unchecked)
--
shiftRA :: Exp -> Exp -> Exp
shiftRA x i = [cexp| $exp:x >> $exp:i |]

-- Logical right shift (unchecked)
--
shiftRL :: IntegralType a -> Exp -> Exp -> Exp
shiftRL ty x i =
  let int  = typeOf (integralType :: IntegralType Int)
      word = typeOf (integralType :: IntegralType Word)
  in
  case ty of
    TypeInt{}    -> [cexp| ($ty:int)        (($ty:word)           $exp:x >> $exp:i) |]
    TypeInt8{}   -> [cexp| (typename Int8)  ((typename Word8)     $exp:x >> $exp:i) |]
    TypeInt16{}  -> [cexp| (typename Int16) ((typename Word16)    $exp:x >> $exp:i) |]
    TypeInt32{}  -> [cexp| (typename Int32) ((typename Word32)    $exp:x >> $exp:i) |]
    TypeInt64{}  -> [cexp| (typename Int64) ((typename Word64)    $exp:x >> $exp:i) |]
    TypeCShort{} -> [cexp| (short)          ((unsigned short)     $exp:x >> $exp:i) |]
    TypeCInt{}   -> [cexp| (int)            ((unsigned int)       $exp:x >> $exp:i) |]
    TypeCLong{}  -> [cexp| (long)           ((unsigned long)      $exp:x >> $exp:i) |]
    TypeCLLong{} -> [cexp| (long long)      ((unsigned long long) $exp:x >> $exp:i) |]

    -- unsigned types use arithmetic shift
    _            -> $internalCheck "shiftRL" "unhandled signed type" (unsignedIntegralNum ty) (shiftRA x i)


rotateL :: forall a. IntegralType a -> Exp -> Exp -> Gen Exp
rotateL t x i | IntegralDict <- integralDict t = do
  let int  = integralType :: IntegralType Int
      wsib = finiteBitSize (undefined::a)
  --
  x' <- bind (typeOf t)    x
  i' <- bind (typeOf int) (i `band` integral int (wsib - 1))
  return $ (x' `shiftL` i') `bor` (shiftRL t x' (integral int wsib `sub` i'))

rotateR :: IntegralType a -> Exp -> Exp -> Gen Exp
rotateR t x i = rotateL t x (negate i)


-- Operators from Fractional & Floating
-- ====================================

fdiv :: Exp -> Exp -> Exp
fdiv x y = [cexp| $exp:x / $exp:y |]

recip :: FloatingType a -> Exp -> Exp
recip t x | FloatingDict <- floatingDict t = fdiv (floating t 1) x

sin :: FloatingType a -> Exp -> Exp
sin t x = mathf t "sin" [x]

cos :: FloatingType a -> Exp -> Exp
cos t x = mathf t "cos" [x]

tan :: FloatingType a -> Exp -> Exp
tan t x = mathf t "tan" [x]

asin :: FloatingType a -> Exp -> Exp
asin t x = mathf t "asin" [x]

acos :: FloatingType a -> Exp -> Exp
acos t x = mathf t "acos" [x]

atan :: FloatingType a -> Exp -> Exp
atan t x = mathf t "atan" [x]

sinh :: FloatingType a -> Exp -> Exp
sinh t x = mathf t "sinh" [x]

cosh :: FloatingType a -> Exp -> Exp
cosh t x = mathf t "cosh" [x]

tanh :: FloatingType a -> Exp -> Exp
tanh t x = mathf t "tanh" [x]

asinh :: FloatingType a -> Exp -> Exp
asinh t x = mathf t "asinh" [x]

acosh :: FloatingType a -> Exp -> Exp
acosh t x = mathf t "acosh" [x]

atanh :: FloatingType a -> Exp -> Exp
atanh t x = mathf t "atanh" [x]

exp :: FloatingType a -> Exp -> Exp
exp t x = mathf t "exp" [x]

sqrt :: FloatingType a -> Exp -> Exp
sqrt t x = mathf t "sqrt" [x]

pow :: FloatingType a -> Exp -> Exp -> Exp
pow t x y = mathf t "pow" [x,y]

log :: FloatingType a -> Exp -> Exp
log t x = mathf t "log" [x]

logBase :: FloatingType a -> Exp -> Exp -> Exp
logBase t x y = log t y `fdiv` log t x


-- Operators from RealFrac
-- =======================

trunc :: FloatingType a -> IntegralType b -> Exp -> Exp
trunc ta tb x = cast tb $ mathf ta "trunc" [x]

round :: FloatingType a -> IntegralType b -> Exp -> Exp
round ta tb x = cast tb $ mathf ta "round" [x]

floor :: FloatingType a -> IntegralType b -> Exp -> Exp
floor ta tb x = cast tb $ mathf ta "floor" [x]

ceiling :: FloatingType a -> IntegralType b -> Exp -> Exp
ceiling ta tb x = cast tb $ mathf ta "ceil" [x]


-- Operators from RealFloat
-- ========================

atan2 :: FloatingType a -> Exp -> Exp -> Exp
atan2 t x y = mathf t "atan2" [x, y]

isNaN :: Exp -> Exp
isNaN x = ccall "isnan" [x]


-- Relational and equality operators
-- =================================

lt :: Exp -> Exp -> Exp
lt x y = [cexp| $exp:x < $exp:y |]

gt :: Exp -> Exp -> Exp
gt x y = [cexp| $exp:x > $exp:y |]

leq  :: Exp -> Exp -> Exp
leq x y = [cexp| $exp:x <= $exp:y |]

geq :: Exp -> Exp -> Exp
geq x y = [cexp| $exp:x >= $exp:y |]

eq :: Exp -> Exp -> Exp
eq x y = [cexp| $exp:x == $exp:y |]

neq :: Exp -> Exp -> Exp
neq x y = [cexp| $exp:x != $exp:y |]

max :: ScalarType a -> Exp -> Exp -> Exp
max (NonNumScalarType _) x y =
  let t = scalarType :: ScalarType Int32
  in  max t (cast t x) (cast t y)

max (NumScalarType (IntegralNumType _)) x y = ccall   "max"  [x,y]
max (NumScalarType (FloatingNumType t)) x y = mathf t "fmax" [x,y]

min :: ScalarType a -> Exp -> Exp -> Exp
min (NonNumScalarType _) x y =
  let t = scalarType :: ScalarType Int32
  in  min t (cast t x) (cast t y)

min (NumScalarType (IntegralNumType _)) x y = ccall   "min"  [x,y]
min (NumScalarType (FloatingNumType t)) x y = mathf t "fmin" [x,y]


-- Logical operators
-- =================

land :: Exp -> Exp -> Exp
land x y = [cexp| $exp:x && $exp:y |]

lor :: Exp -> Exp -> Exp
lor x y = [cexp| $exp:x || $exp:y |]

lnot :: Exp -> Exp
lnot x = [cexp| ! $exp:x |]


-- Type Conversions
-- ================

ord :: Exp -> Exp
ord = cast (scalarType :: ScalarType Int)

chr :: Exp -> Exp
chr = cast (scalarType :: ScalarType Char)

boolToInt :: Exp -> Exp
boolToInt = cast (scalarType :: ScalarType Int)

fromIntegral :: IntegralType a -> NumType b -> Exp -> Exp
fromIntegral _ tb = cast tb


-- Helpers
-- =======

cast :: TypeOf a => a -> Exp -> Exp
cast t x = [cexp| ($ty:(typeOf t)) $exp:x |]

mathf :: forall t. FloatingType t -> String -> [Exp] -> Exp
mathf ty f args | FloatingDict <- floatingDict ty =
  let
      fun = f ++ case sizeOf (undefined :: t) of
                   4  -> "f"
                   8  -> []
                   16 -> "l"        -- long double
                   _  -> $internalError "mathf" "unsupported floating point size"
  in
  ccall fun args


infix 0 ?:
(?:) :: Exp -> (Exp, Exp) -> Exp
(?:) p (t,e) = [cexp| $exp:p ? $exp:t : $exp:e |]

cases :: [(Exp, Exp)] -> Exp -> Exp
cases []           def = def
cases ((p,b):rest) def = p ?: (b, cases rest def)