{-# LANGUAGE ViewPatterns #-}
--------------------------------------------------------------------------------
-- |
-- Module      : ArrayFire.Sparse
-- Copyright   : David Johnson (c) 2019-2026
-- License     : BSD3
-- Maintainer  : David Johnson <code@dmj.io>
-- Stability   : Experimental
-- Portability : GHC
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func.htm)
-- Functions to create and handle sparse arrays and matrix operations.
--
-- *Note*
-- Sparse functionality support was added to ArrayFire in v3.4.0.
--
-- >>> createSparseArray 3 3 (vector @Double 3 [1,2,3]) (vector @Int32 3 [0,1,2]) (vector @Int32 3 [0,1,2]) COO
--
--
--------------------------------------------------------------------------------
module ArrayFire.Sparse where

import Control.Exception (throw)

import ArrayFire.Exception
import ArrayFire.Types
import ArrayFire.FFI
import ArrayFire.Internal.Algorithm (af_any_true_all)
import ArrayFire.Internal.Sparse
import ArrayFire.Internal.Types
import Data.Int

-- | This function converts af::array of values, row indices and column indices into a sparse array.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__create.htm#ga42c5cf729a232c1cbbcfe0f664f3b986)
--
-- *Note*
-- This function only create references of these arrays into the sparse data structure and does not do deep copies.
--
-- >>> createSparseArray 3 3 (vector @Double 3 [1,2,3]) (vector @Int32 3 [0,1,2]) (vector @Int32 3 [0,1,2]) COO
--
createSparseArray
  :: (AFType a, Fractional a)
  => Int
  -- ^ is the number of rows in the dense matrix
  -> Int
  -- ^ is the number of columns in the dense matrix
  -> Array a
  -- ^ is the 'Array' containing the non-zero elements of the matrix
  -> Array Int32
  -- ^ is the row indices for the sparse array
  -> Array Int32
  -- ^ the column indices for the sparse array
  -> Storage
  -- ^ the storage format of the sparse array
  -> Array a
  -- ^ Sparse Array
createSparseArray :: forall a.
(AFType a, Fractional a) =>
Int
-> Int
-> Array a
-> Array Int32
-> Array Int32
-> Storage
-> Array a
createSparseArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> DimT
r) (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> DimT
c) Array a
arr1 Array Int32
arr2 Array Int32
arr3 Storage
s =
  Array a
-> Array Int32
-> Array Int32
-> (Ptr AFArray -> AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array a
forall a.
Array a
-> Array Int32
-> Array Int32
-> (Ptr AFArray -> AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array a
op3Int Array a
arr1 Array Int32
arr2 Array Int32
arr3 (\Ptr AFArray
p AFArray
ar1 AFArray
ar2 AFArray
ar3 -> Ptr AFArray
-> DimT
-> DimT
-> AFArray
-> AFArray
-> AFArray
-> AFStorage
-> IO AFErr
af_create_sparse_array Ptr AFArray
p DimT
r DimT
c AFArray
ar1 AFArray
ar2 AFArray
ar3 (Storage -> AFStorage
toStorage Storage
s))

-- af_err af_create_sparse_array_from_ptr(af_array *out, const dim_t nRows, const dim_t nCols, const dim_t nNZ, const void * const values, const int * const rowIdx, const int * const colIdx, const af_dtype type, const af_storage stype, const af_source src);

-- | This function converts a dense af_array into a sparse array.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__create.htm#ga52e3b2895cf9e9d697a06b4b44190d92)
--
-- *Note*
-- This function only create references of these arrays into the sparse data structure and does not do deep copies.
--
-- >>> createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR
-- ArrayFire Array
-- Storage Format : AF_STORAGE_CSR
-- [2 2 1 1]
-- ArrayFire Array: Values
-- [4 1 1 1]
--     1.0000     3.0000     2.0000     4.0000
-- ArrayFire Array: RowIdx
-- [3 1 1 1]
--          0          2          4
-- ArrayFire Array: ColIdx
-- [4 1 1 1]
--          0          1          0          1
--
createSparseArrayFromDense
  :: (AFType a, Fractional a)
  => Array a
  -- ^ is the source dense matrix
  -> Storage
  -- ^ is the storage format of the sparse array
  -> Array a
  -- ^ 'Array' for the sparse array with the given storage type
createSparseArrayFromDense :: forall a. (AFType a, Fractional a) => Array a -> Storage -> Array a
createSparseArrayFromDense Array a
a Storage
s
  -- Guard: converting an all-zero dense matrix (NNZ = 0) segfaults inside
  -- ArrayFire (observed on AF 3.8.2). Throw a proper AFException instead of
  -- crashing the process.
  | Double
nonZero Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0.0 =
      AFException -> Array a
forall a e. (HasCallStack, Exception e) => e -> a
throw (AFException -> Array a) -> AFException -> Array a
forall a b. (a -> b) -> a -> b
$ AFExceptionType -> Int -> String -> AFException
AFException AFExceptionType
SizeError Int
203
        String
"createSparseArrayFromDense: input has no non-zero elements; zero-NNZ sparse arrays crash the underlying ArrayFire library"
  | Bool
otherwise =
      Array a
a Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
x -> Ptr AFArray -> AFArray -> AFStorage -> IO AFErr
af_create_sparse_array_from_dense Ptr AFArray
p AFArray
x (Storage -> AFStorage
toStorage Storage
s))
  where
    (Double
nonZero, Double
_) = Array a
a Array a
-> (Ptr Double -> Ptr Double -> AFArray -> IO AFErr)
-> (Double, Double)
forall a b arr.
(Storable a, Storable b) =>
Array arr -> (Ptr a -> Ptr b -> AFArray -> IO AFErr) -> (a, b)
`infoFromArray2` Ptr Double -> Ptr Double -> AFArray -> IO AFErr
af_any_true_all :: (Double, Double)

-- | Convert an existing sparse array into a different storage format.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__convert__to.htm)
--
-- Converting storage formats is allowed between 'CSR', 'COO' and DENSE.
--
-- When converting to DENSE, a dense array is returned.
--
-- *Note*
-- 'CSC' is currently not supported.
--
-- >>> array = createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR
-- >>> array
-- ArrayFire Array
-- Storage Format : AF_STORAGE_CSR
-- [2 2 1 1]
-- ArrayFire Array: Values
-- [4 1 1 1]
--     1.0000
--     3.0000
--     2.0000
--     4.0000

-- ArrayFire Array: RowIdx
-- [3 1 1 1]
--          0
--          2
--          4

-- ArrayFire Array: ColIdx
-- [4 1 1 1]
--          0
--          1
--          0
--          1
--
-- >>> sparseConvertTo array COO
-- ArrayFire Array
-- Storage Format : AF_STORAGE_COO
-- [2 2 1 1]
-- ArrayFire Array: Values
-- [4 1 1 1]
--     1.0000
--     2.0000
--     3.0000
--     4.0000

-- ArrayFire Array: RowIdx
-- [4 1 1 1]
--          0
--          1
--          0
--          1

-- ArrayFire Array: ColIdx
-- [4 1 1 1]
--          0
--          0
--          1
--          1
--

-- | Converts a sparse 'Array' from one storage format ('Storage') to another
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__convert__to.htm)
sparseConvertTo
  :: (AFType a, Fractional a)
  => Array a
  -- ^ is the source sparse matrix to be converted
  -> Storage
  -- ^ is the storage format of the output sparse array
  -> Array a
  -- ^ the sparse array with the given storage type
sparseConvertTo :: forall a. (AFType a, Fractional a) => Array a -> Storage -> Array a
sparseConvertTo Array a
a Storage
s =
  Array a
a Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
x -> Ptr AFArray -> AFArray -> AFStorage -> IO AFErr
af_sparse_convert_to Ptr AFArray
p AFArray
x (Storage -> AFStorage
toStorage Storage
s))

-- | Returns a dense array from a sparse input
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__dense.htm#ga80c3d8db78d537b74d9caebcf359b6a5)
--
-- Converts the sparse matrix into a dense matrix and returns it
--
-- >>> array = createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR
-- >>> array
-- ArrayFire Array
-- Storage Format : AF_STORAGE_CSR
-- [2 2 1 1]
-- ArrayFire Array: Values
-- [4 1 1 1]
--     1.0000
--     3.0000
--     2.0000
--     4.0000
--
-- ArrayFire Array: RowIdx
-- [3 1 1 1]
--          0
--          2
--          4
--
-- ArrayFire Array: ColIdx
-- [4 1 1 1]
--          0
--          1
--          0
--          1
--
sparseToDense
  :: (AFType a, Fractional a)
  => Array a
  -> Array a
sparseToDense :: forall a. (AFType a, Fractional a) => Array a -> Array a
sparseToDense = (Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` Ptr AFArray -> AFArray -> IO AFErr
af_sparse_to_dense)

-- | Returns reference to components of the input sparse array.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__info.htm#gae6b553df80e21c174d374e82d8505ba5)
--
-- Returns reference to values, row indices, column indices and storage format of an input sparse array
--
-- >>> (values, rows, cols, storage) = sparseGetInfo $ createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR
-- >>> values
-- ArrayFire Array
-- [4 1 1 1]
--     1.0000
--     3.0000
--     2.0000
--     4.0000
--
-- >>> cols
-- ArrayFire Array
-- [3 1 1 1]
--          0
--          2
--          4
--
-- >>> rows
-- ArrayFire Array
-- [4 1 1 1]
--          0
--          1
--          0
--          1
--
-- >>> storage
-- CSR
--
sparseGetInfo
  :: (AFType a, Fractional a)
  => Array a
  -> (Array a, Array a, Array a, Storage)
sparseGetInfo :: forall a.
(AFType a, Fractional a) =>
Array a -> (Array a, Array a, Array a, Storage)
sparseGetInfo Array a
x = do
  let (Array a
a,Array a
b,Array a
c,AFStorage
d) = Array a
x Array a
-> (Ptr AFArray
    -> Ptr AFArray
    -> Ptr AFArray
    -> Ptr AFStorage
    -> AFArray
    -> IO AFErr)
-> (Array a, Array a, Array a, AFStorage)
forall b a.
Storable b =>
Array a
-> (Ptr AFArray
    -> Ptr AFArray -> Ptr AFArray -> Ptr b -> AFArray -> IO AFErr)
-> (Array a, Array a, Array a, b)
`op3p1` Ptr AFArray
-> Ptr AFArray
-> Ptr AFArray
-> Ptr AFStorage
-> AFArray
-> IO AFErr
af_sparse_get_info
  (Array a
a,Array a
b,Array a
c,AFStorage -> Storage
fromStorage AFStorage
d)

-- | Returns reference to the values component of the sparse array.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__values.htm)
--
-- Returns reference to the values component of the sparse array.
-- Values is the 'Array' containing the non-zero elements of the dense matrix.
--
-- >>> sparseGetValues (createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR)
-- ArrayFire Array
-- [4 1 1 1]
--     1.0000
--     3.0000
--     2.0000
--     4.0000
--
sparseGetValues
  :: (AFType a, Fractional a)
  => Array a
  -> Array a
sparseGetValues :: forall a. (AFType a, Fractional a) => Array a -> Array a
sparseGetValues = (Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` Ptr AFArray -> AFArray -> IO AFErr
af_sparse_get_values)

-- | Returns reference to the row indices component of the sparse array. More...
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__row__idx.htm)
--
-- Returns reference to the row indices component of the sparse array.
-- Row indices is the 'Array' containing the row indices of the sparse array
-- (for 'CSR' storage these are the compressed row offsets, of length
-- rows + 1).
--
-- >>> sparseGetRowIdx (createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR)
-- ArrayFire Array
-- [3 1 1 1]
--          0
--          2
--          4
--
sparseGetRowIdx
  :: (AFType a, Fractional a)
  => Array a
  -> Array a
sparseGetRowIdx :: forall a. (AFType a, Fractional a) => Array a -> Array a
sparseGetRowIdx = (Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` Ptr AFArray -> AFArray -> IO AFErr
af_sparse_get_row_idx)

-- | Returns reference to the column indices component of the sparse array. More...
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__col__idx.htm)
--
-- Returns reference to the column indices component of the sparse array.
-- Column indices is the 'Array' containing the column indices of the sparse array.
--
-- >>> sparseGetColIdx (createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR)
-- ArrayFire Array
-- [4 1 1 1]
--          0
--          1
--          0
--          1
--
sparseGetColIdx
  :: (AFType a, Fractional a)
  => Array a
  -> Array a
sparseGetColIdx :: forall a. (AFType a, Fractional a) => Array a -> Array a
sparseGetColIdx = (Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` Ptr AFArray -> AFArray -> IO AFErr
af_sparse_get_col_idx)

-- | Returns the storage type of a sparse array.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__storage.htm)
--
-- Returns the number of non zero elements in the sparse array.
-- This is always equal to the size of the values array.
--
-- >>> sparseGetStorage $ createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR
-- CSR
--
sparseGetStorage
  :: (AFType a, Fractional a)
  => Array a
  -> Storage
sparseGetStorage :: forall a. (AFType a, Fractional a) => Array a -> Storage
sparseGetStorage Array a
a =
  AFStorage -> Storage
fromStorage (Array a
a Array a -> (Ptr AFStorage -> AFArray -> IO AFErr) -> AFStorage
forall a b.
Storable a =>
Array b -> (Ptr a -> AFArray -> IO AFErr) -> a
`infoFromArray` Ptr AFStorage -> AFArray -> IO AFErr
af_sparse_get_storage)

-- | Returns the number of non zero elements in the sparse array. More...
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__sparse__func__nnz.htm#ga0c1ad61d829c02a280c28820eb91f03e)
--
-- Returns the number of non zero elements in the sparse array.
-- This is always equal to the size of the values array.
--
-- >>> sparseGetNNZ $ createSparseArrayFromDense (matrix @Double (2,2) [[1,2],[3,4]]) CSR
-- 4
--
sparseGetNNZ
  :: (AFType a, Fractional a)
  => Array a
  -> Int
sparseGetNNZ :: forall a. (AFType a, Fractional a) => Array a -> Int
sparseGetNNZ Array a
a =
  DimT -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Array a
a Array a -> (Ptr DimT -> AFArray -> IO AFErr) -> DimT
forall a b.
Storable a =>
Array b -> (Ptr a -> AFArray -> IO AFErr) -> a
`infoFromArray` Ptr DimT -> AFArray -> IO AFErr
af_sparse_get_nnz)