{-# LANGUAGE RecordWildCards #-}

{-# LANGUAGE ScopedTypeVariables #-}

{-# LANGUAGE PolyKinds #-}

{-# LANGUAGE DataKinds #-}

{-# LANGUAGE AllowAmbiguousTypes #-}

{-# LANGUAGE FlexibleInstances #-}

{-# LANGUAGE TypeApplications #-}

{-# LANGUAGE ViewPatterns #-}

{-# LANGUAGE KindSignatures #-}

-- |

module ArrayFire.Array where

import Control.Exception

import Control.Monad

import Data.Proxy

import Data.Vector.Storable hiding (mapM_, take, concat, concatMap)

import qualified Data.Vector.Storable as V

import Foreign.ForeignPtr

import Foreign.Marshal hiding (void)

import Foreign.Ptr

import Foreign.Storable

import System.IO.Unsafe

import ArrayFire.Exception

import ArrayFire.FFI

import ArrayFire.Util

import ArrayFire.Internal.Array

import ArrayFire.Internal.Defines

import ArrayFire.Internal.Types

-- | Smart constructor for creating a scalar 'Array'

scalar :: AFType a => a -> Array a

scalar x = mkArray [1] [x]

-- | Smart constructor for creating a vector 'Array'

vector :: AFType a => Int -> [a] -> Array a

vector n = mkArray [n] . take n

-- | Smart constructor for creating a matrix 'Array'

matrix :: AFType a => (Int,Int) -> [[a]] -> Array a

matrix (x,y)

= mkArray [x,y]

. concat

. take y

. fmap (take x)

-- | Smart constructor for creating a cubic 'Array'

cube :: AFType a => (Int,Int,Int) -> [[[a]]] -> Array a

cube (x,y,z)

= mkArray [x,y,z]

. concat

. fmap concat

. take z

. fmap (take y)

. (fmap . fmap . take) x

-- | Smart constructor for creating a tensor 'Array'

tensor :: AFType a => (Int, Int,Int,Int) -> [[[[a]]]] -> Array a

tensor (w,x,y,z)

= mkArray [w,x,y,z]

. concat

. fmap concat

. (fmap . fmap) concat

. take z

. (fmap . take) y

. (fmap . fmap . take) x

. (fmap . fmap . fmap . take) w

-- | Internal function for 'Array' construction

mkArray

:: forall array

. AFType array

=> [Int]

-- ^ Dimensions

-> [array]

-- ^ Array elements

-> Array array

-- ^ Returned array

{-# NOINLINE mkArray #-}

mkArray dims xs =

unsafePerformIO $ do

when (Prelude.length (take size xs) < size) $ do

let msg = "Invalid elements provided. "

<> "Expected "

<> show size

<> " elements received "

<> show (Prelude.length xs)

throwIO (AFException SizeError 203 msg)

dataPtr <- castPtr <$> newArray (Prelude.take size xs)

let ndims = fromIntegral (Prelude.length dims)

alloca $ \arrayPtr -> do

zeroOutArray arrayPtr

dimsPtr <- newArray (DimT . fromIntegral <$> dims)

throwAFError =<< af_create_array arrayPtr dataPtr ndims dimsPtr dType

free dataPtr >> free dimsPtr

arr <- peek arrayPtr

Array <$> newForeignPtr af_release_array_finalizer arr

where

size = Prelude.product dims

dType = afType (Proxy @array)

-- | Copies an 'Array' to a new 'Array'

copyArray

:: AFType a

=> Array a

-- ^ 'Array' to be copied

-> Array a

-- ^ Newly copied 'Array'

copyArray = (`op1` af_copy_array)

-- | Retains an 'Array', increases reference count

retainArray

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Array a

retainArray =

(`op1` af_retain_array)

-- | Retrieves 'Array' reference count

getDataRefCount

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Int

-- ^ Reference count

getDataRefCount =

fromIntegral . (`infoFromArray` af_get_data_ref_count)

-- | Should manual evaluation occur

setManualEvalFlag

:: Bool

-- ^ Whether or not to perform manual evaluation

-> IO ()

setManualEvalFlag (fromIntegral . fromEnum -> b) =

afCall (af_set_manual_eval_flag b)

-- | Retrieve manual evaluation status

getManualEvalFlag

:: IO Bool

getManualEvalFlag =

toEnum . fromIntegral <$> afCall1 af_get_manual_eval_flag

-- | Retrieve element count

getElements

:: AFType a

=> Array a

-- ^ Input 'Array'

-> Int

-- ^ Count of elements in 'Array'

getElements a =

fromIntegral (a `infoFromArray` af_get_elements)

-- | Retrieve type of 'Array'

getType

:: AFType a

=> Array a

-> AFDType

getType a = fromAFType (a `infoFromArray` af_get_type)

-- | Retrieves dimensions of 'Array'

getDims

:: AFType a

=> Array a

-> (Int,Int,Int,Int)

getDims arr = do

let (a,b,c,d) = arr `infoFromArray4` af_get_dims

(fromIntegral a, fromIntegral b, fromIntegral c, fromIntegral d)

-- | Retrieves number of dimensions in 'Array'

getNumDims

:: AFType a

=> Array a

-> Int

getNumDims = fromIntegral . (`infoFromArray` af_get_numdims)

-- | Checks if an 'Array' is empty

isEmpty

:: AFType a

=> Array a

-> Bool

isEmpty a = toEnum . fromIntegral $ (a `infoFromArray` af_is_empty)

-- | Checks if an 'Array' is a scalar (contains only one element)

isScalar

:: AFType a

=> Array a

-> Bool

isScalar a = toEnum . fromIntegral $ (a `infoFromArray` af_is_scalar)

-- | Checks if an 'Array' is row-oriented

isRow

:: AFType a

=> Array a

-> Bool

isRow a = toEnum . fromIntegral $ (a `infoFromArray` af_is_row)

-- | Checks if an 'Array' is a column-oriented

isColumn

:: AFType a

=> Array a

-> Bool

isColumn a = toEnum . fromIntegral $ (a `infoFromArray` af_is_column)

-- | Checks if an 'Array' is a vector

isVector

:: AFType a

=> Array a

-> Bool

isVector a = toEnum . fromIntegral $ (a `infoFromArray` af_is_vector)

-- | Checks if an 'Array' is a Complex

isComplex

:: AFType a

=> Array a

-> Bool

isComplex a = toEnum . fromIntegral $ (a `infoFromArray` af_is_complex)

-- | Checks if an 'Array' is Real

isReal

:: AFType a

=> Array a

-> Bool

isReal a = toEnum . fromIntegral $ (a `infoFromArray` af_is_real)

-- | Checks if an 'Array' is 'Double'

isDouble

:: AFType a

=> Array a

-> Bool

isDouble a = toEnum . fromIntegral $ (a `infoFromArray` af_is_double)

-- | Checks if an 'Array' is 'Float'

isSingle

:: AFType a

=> Array a

-> Bool

isSingle a = toEnum . fromIntegral $ (a `infoFromArray` af_is_single)

-- | Checks if an 'Array' is 'Double', 'Float', Complex 'Double', or Complex 'Float'

isRealFloating

:: AFType a

=> Array a

-> Bool

isRealFloating a = toEnum . fromIntegral $ (a `infoFromArray` af_is_realfloating)

-- | Checks if an 'Array' is 'Double' or 'Float'

isFloating

:: AFType a

=> Array a

-> Bool

isFloating a = toEnum . fromIntegral $ (a `infoFromArray` af_is_floating)

-- | Checks if an 'Array' is of type Int16, Int32, or Int64

isInteger

:: AFType a

=> Array a

-> Bool

isInteger a = toEnum . fromIntegral $ (a `infoFromArray` af_is_integer)

-- | Checks if an 'Array' is of type CBool

isBool

:: AFType a

=> Array a

-> Bool

isBool a = toEnum . fromIntegral $ (a `infoFromArray` af_is_bool)

-- | Checks if an 'Array' is sparse

isSparse

:: AFType a

=> Array a

-> Bool

isSparse a = toEnum . fromIntegral $ (a `infoFromArray` af_is_sparse)

-- | Converts an 'Array' to a 'Storable' 'Vector'

toVector :: forall a . AFType a => Array a -> Vector a

toVector arr@(Array fptr) = do

unsafePerformIO . mask_ . withForeignPtr fptr $ \arrPtr -> do

let len = getElements arr

size = len * getSizeOf (Proxy @a)

ptr <- mallocBytes (len * size)

throwAFError =<< af_get_data_ptr (castPtr ptr) arrPtr

newFptr <- newForeignPtr finalizerFree ptr

pure $ unsafeFromForeignPtr0 newFptr len

-- | Converts an 'Array' to [a]

toList :: forall a . AFType a => Array a -> [a]

toList = V.toList . toVector

-- | Retrieves single scalar value from an 'Array'

getScalar :: forall a b . (Storable a, AFType b) => Array b -> a

getScalar (Array fptr) =

unsafePerformIO . mask_ . withForeignPtr fptr $ \arrPtr -> do

alloca $ \ptr -> do

throwAFError =<< af_get_scalar (castPtr ptr) arrPtr

peek ptr