module Data.Array.Comfort.Storable (
   Array,
   shape,
   reshape,
   mapShape,

   accessMaybe, (!),
   Array.toList,
   Array.vectorFromList,
   toAssociations,
   fromList,
   fromMap, toMap,
   fromIntMap, toIntMap,
   fromTuple, toTuple,
   fromRecord, toRecord,
   fromContainer,
   toContainer,
   sample,
   replicate,
   fromBoxed,
   toBoxed,
   fromStorableVector,
   toStorableVector,
   fromBlockArray1,
   fromBlockArray2,
   fromNonEmptyBlockArray2,

   Array.map,
   Array.mapWithIndex,
   zipWith,
   (//),
   accumulate,
   fromAssociations,

   pick,
   toRowArray,
   fromRowArray,
   Array.singleton,
   Array.append,
   Array.take, Array.drop,
   Array.takeLeft, Array.takeRight, Array.split,
   Array.takeCenter,

   Array.sum, Array.product,
   minimum, argMinimum,
   maximum, argMaximum,
   limits,
   Array.foldl,
   foldl1,
   foldMap,
   ) where

import qualified Data.Array.Comfort.Storable.Mutable.Unchecked as MutArrayNC
import qualified Data.Array.Comfort.Storable.Mutable.Private as MutArrayPriv
import qualified Data.Array.Comfort.Storable.Mutable as MutArray
import qualified Data.Array.Comfort.Storable.Unchecked as Array
import qualified Data.Array.Comfort.Storable.Memory as Memory
import qualified Data.Array.Comfort.Container as Container
import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Check as Check
import qualified Data.Array.Comfort.Shape.Tuple as TupleShape
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable.Unchecked (Array(Array))
import Data.Array.Comfort.Shape ((::+)((::+)))

import System.IO.Unsafe (unsafePerformIO)
import Foreign.Marshal.Array (copyArray, advancePtr)
import Foreign.Storable (Storable)
import Foreign.ForeignPtr (withForeignPtr)

import qualified Control.Monad.Trans.State as MS
import Control.Monad.ST (runST)

import qualified Data.StorableVector as SV
import qualified Data.StorableVector.Base as SVB
import qualified Data.IntMap as IntMap
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import qualified Data.Tuple.Strict as StrictTuple
import Data.IntMap (IntMap)
import Data.IntSet (IntSet)
import Data.Map (Map)
import Data.Set (Set)
import Data.Foldable (forM_)
import Data.Either.HT (maybeRight)
import Data.Tuple.HT (mapPair)
import Data.Semigroup
         (Semigroup, (<>), Min(Min,getMin), Max(Max,getMax), Arg(Arg))

import Prelude2010 hiding (map, zipWith, foldl1, minimum, maximum, replicate)
import Prelude ()


{- $setup
>>> import qualified DocTest.Data.Array.Comfort.Boxed.Unchecked
>>>                                              as TestBoxedArray
>>> import qualified Data.Array.Comfort.Boxed as BoxedArray
>>> import qualified Data.Array.Comfort.Storable as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>> import Data.Array.Comfort.Storable (Array, (!))
>>>
>>> import qualified Test.QuickCheck as QC
>>> import Test.ChasingBottoms.IsBottom (isBottom)
>>>
>>> import Control.Monad (replicateM)
>>> import Control.Applicative ((<$>), (<*>))
>>>
>>> import qualified Data.Map as Map
>>> import qualified Data.Set as Set
>>> import Data.Map (Map)
>>> import Data.Function.HT (Id)
>>> import Data.Complex (Complex((:+)))
>>> import Data.Tuple.HT (swap)
>>> import Data.Word (Word16)
>>>
>>> import Foreign.Storable (Storable)
>>>
>>> type ShapeInt = Shape.ZeroBased Int
>>> type X = Shape.Element
>>>
>>> shapeInt :: Int -> ShapeInt
>>> shapeInt = Shape.ZeroBased
>>>
>>> genArray :: QC.Gen (Array ShapeInt Word16)
>>> genArray = Array.vectorFromList <$> QC.arbitrary
>>>
>>> genArray2 :: QC.Gen (Array (ShapeInt,ShapeInt) Word16)
>>> genArray2 = do
>>>    xs <- QC.arbitrary
>>>    let n = length xs
>>>    (k,m) <-
>>>       if n == 0
>>>          then QC.elements [(,) 0, flip (,) 0] <*> QC.choose (1,20)
>>>          else fmap (\m -> (div n m, m)) $ QC.choose (1,n)
>>>    return $ Array.fromList (Shape.ZeroBased k, Shape.ZeroBased m) xs
>>>
>>> genArrayForShape :: (Shape.C shape) => shape -> QC.Gen (Array shape Word16)
>>> genArrayForShape sh =
>>>    Array.fromList sh <$> replicateM (Shape.size sh) QC.arbitrary
>>>
>>> genNonEmptyArray2 :: QC.Gen (Array (ShapeInt,ShapeInt) Word16)
>>> genNonEmptyArray2 = do
>>>    xs <- QC.getNonEmpty <$> QC.arbitrary
>>>    let n = length xs
>>>    m <- QC.choose (1,n)
>>>    return $ Array.fromList (Shape.ZeroBased (div n m), Shape.ZeroBased m) xs
>>>
>>> infix 4 ==?
>>> (==?) :: a -> a -> (a,a)
>>> (==?) = (,)
>>>
>>> forAllNonEmpty :: (Eq b) => (Array ShapeInt Word16 -> (b,b)) -> QC.Property
>>> forAllNonEmpty f =
>>>    QC.forAll genArray $ \xs ->
>>>    case f xs of
>>>       (resultArray,resultList) ->
>>>          if Array.shape xs == Shape.ZeroBased 0
>>>             then isBottom resultArray
>>>             else resultArray == resultList
>>>
>>>
>>> transpose ::
>>>    (Shape.Indexed sh0, Shape.Indexed sh1, Storable a) =>
>>>    Array (sh0,sh1) a -> Array (sh1,sh0) a
>>> transpose a = Array.sample (swap $ Array.shape a) (\(i,j) -> a!(j,i))
-}


shape :: Array sh a -> sh
shape = Array.shape

reshape :: (Shape.C sh0, Shape.C sh1) => sh1 -> Array sh0 a -> Array sh1 a
reshape = Check.reshape "Storable" shape Array.reshape

mapShape ::
   (Shape.C sh0, Shape.C sh1) => (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape f arr = reshape (f $ shape arr) arr


{- |
>>> Array.fromList (shapeInt 5) ['a'..]
StorableArray.fromList (ZeroBased {zeroBasedSize = 5}) "abcde"
-}
fromList :: (Shape.C sh, Storable a) => sh -> [a] -> Array sh a
fromList sh arr = runST (MutArrayNC.unsafeFreeze =<< MutArray.fromList sh arr)

fromMap :: (Ord k, Storable a) => Map k a -> Array (Set k) a
fromMap m = fromList (Map.keysSet m) (Map.elems m)

toMap :: (Ord k, Storable a) => Array (Set k) a -> Map k a
toMap = Map.fromAscList . toAssociations

fromIntMap :: (Storable a) => IntMap a -> Array IntSet a
fromIntMap m = fromList (IntMap.keysSet m) (IntMap.elems m)

toIntMap :: (Storable a) => Array IntSet a -> IntMap a
toIntMap = IntMap.fromAscList . toAssociations

{- |
>>> Array.fromTuple ('a',('b','c')) :: Array (Shape.NestedTuple Shape.TupleIndex (X,(X,X))) Char
StorableArray.fromList (NestedTuple {getNestedTuple = (Element 0,(Element 1,Element 2))}) "abc"

>>> :{
   let arr :: Array (Shape.NestedTuple Shape.TupleAccessor (X,(X,X))) Char
       arr = Array.fromTuple ('a',('b','c'))
   in (arr ! fst, arr ! (fst.snd))
:}
('a','b')
-}
fromTuple ::
   (TupleShape.NestedTuple tuple, Storable a) =>
   Shape.DataTuple tuple a -> Array (Shape.NestedTuple ixtype tuple) a
fromTuple tuple =
   case MS.evalState (TupleShape.decons tuple) (Shape.Element 0) of
      (sh, xs) -> fromList (Shape.NestedTuple sh) xs

toTuple ::
   (TupleShape.NestedTuple tuple, Storable a) =>
   Array (Shape.NestedTuple ixtype tuple) a -> Shape.DataTuple tuple a
toTuple arr =
   MS.evalState
      (TupleShape.cons $ Shape.getNestedTuple $ shape arr)
      (Array.toList arr)

{- |
>>> :{
   let arr = Array.fromRecord ('a' :+ 'b') in
   let (real:+imag) = Shape.indexRecordFromShape $ Array.shape arr in
   (arr ! real, arr ! imag)
:}
('a','b')
-}
fromRecord ::
   (Trav.Traversable f, Storable a) =>
   f a -> Array (Shape.Record f) a
fromRecord xs =
   fromList
      (Shape.Record $ flip MS.evalState (Shape.Element 0) $
       Trav.traverse (const TupleShape.next) xs)
      (Fold.toList xs)

toRecord ::
   (Trav.Traversable f, Storable a) =>
   Array (Shape.Record f) a -> f a
toRecord arr =
   MS.evalState
      (Trav.traverse (const TupleShape.get) $
       (\(Shape.Record record) -> record) $ shape arr)
      (Array.toList arr)

fromContainer ::
   (Container.C f, Storable a) => f a -> Array (Container.Shape f) a
fromContainer xs = fromList (Container.toShape xs) (Fold.toList xs)

toContainer ::
   (Container.C f, Storable a) => Array (Container.Shape f) a -> f a
toContainer arr = Container.fromList (Array.shape arr) (Array.toList arr)

sample ::
   (Shape.Indexed sh, Storable a) => sh -> (Shape.Index sh -> a) -> Array sh a
sample sh f = Array.fromList sh $ List.map f $ Shape.indices sh

replicate ::
   (Shape.Indexed sh, Storable a) =>
   sh -> a -> Array sh a
replicate sh a = runST (MutArrayNC.unsafeFreeze =<< MutArray.new sh a)


fromBoxed :: (Shape.C sh, Storable a) => BoxedArray.Array sh a -> Array sh a
fromBoxed arr = Array.fromList (BoxedArray.shape arr) $ BoxedArray.toList arr

toBoxed :: (Shape.C sh, Storable a) => Array sh a -> BoxedArray.Array sh a
toBoxed arr = BoxedArray.fromList (Array.shape arr) $ Array.toList arr


fromStorableVector ::
   (Storable a) => SVB.Vector a -> Array (Shape.ZeroBased Int) a
fromStorableVector xs =
   case SVB.toForeignPtr xs of
      (fptr,0,n) -> Array (Shape.ZeroBased n) fptr
      (fptr,s,n) ->
         Array.takeRight $
         Array (Shape.ZeroBased s ::+ Shape.ZeroBased n) fptr

toStorableVector :: (Shape.C sh, Storable a) => Array sh a -> SVB.Vector a
toStorableVector (Array sh fptr) =
   SVB.fromForeignPtr fptr $ Shape.size sh


{- |
>>> :{
   Array.fromBlockArray1 $ BoxedArray.fromList Set.empty [] :: Array (Map Char ShapeInt) Word16
:}
StorableArray.fromList (fromList []) []

>>> :{
   let block n a = Array.replicate (shapeInt n) (a::Word16) in
   Array.fromBlockArray1 $
   BoxedArray.fromList (Set.fromList "ABC") [block 2 0, block 3 1, block 5 2]
:}
StorableArray.fromList (fromList [('A',ZeroBased {... 2}),('B',ZeroBased {... 3}),('C',ZeroBased {... 5})]) [0,0,1,1,1,2,2,2,2,2]
-}
fromBlockArray1 ::
   (Ord k, Shape.C shape, Storable a) =>
   BoxedArray.Array (Set k) (Array shape a) -> Array (Map k shape) a
fromBlockArray1 a =
   reshape (BoxedArray.toMap $ fmap Array.shape a) $
   fromStorableVector $ SV.concat $
   List.map toStorableVector $ BoxedArray.toList a

{- |
Only the outer @BoxedArray@ need to be non-empty.

>>> :{
   let shapeR0 = shapeInt 2; shapeR1 = shapeInt 3 in
   let shapeC0 = shapeInt 3; shapeC1 = shapeInt 2 in
   let block sh a = Array.replicate sh (a::Word16) in
   Array.fromBlockArray2
      (Map.singleton 'A' shapeR0 <> Map.singleton 'B' shapeR1)
      (Map.singleton '1' shapeC0 <> Map.singleton '2' shapeC1) $
   BoxedArray.fromList (Set.fromList "AB", Set.fromList "12")
      [block (shapeR0,shapeC0) 0, block (shapeR0,shapeC1) 1,
       block (shapeR1,shapeC0) 2, block (shapeR1,shapeC1) 3]
:}
StorableArray.fromList (fromList [('A',ZeroBased {... 2}),('B',ZeroBased {... 3})],fromList [('1',ZeroBased {... 3}),('2',ZeroBased {... 2})]) [0,0,0,1,1,0,0,0,1,1,2,2,2,3,3,2,2,2,3,3,2,2,2,3,3]

prop> :{
   QC.forAll genArray2 $ \blockA1 ->
   QC.forAll genArray2 $ \blockB2 ->
   let shapeR0 = fst $ Array.shape blockA1 in
   let shapeC0 = snd $ Array.shape blockA1 in
   let shapeR1 = fst $ Array.shape blockB2 in
   let shapeC1 = snd $ Array.shape blockB2 in
   QC.forAll (genArrayForShape (shapeR0, shapeC1)) $ \blockA2 ->
   QC.forAll (genArrayForShape (shapeR1, shapeC0)) $ \blockB1 ->
   let blocked =
         BoxedArray.fromList (Set.fromList "AB", Set.fromList "12")
            [blockA1, blockA2, blockB1, blockB2] in

   transpose (Array.fromNonEmptyBlockArray2 blocked)
   QC.===
   Array.fromNonEmptyBlockArray2
      (TestBoxedArray.transpose (fmap transpose blocked))
:}

prop> :{
   QC.forAll genArray2 $ \blockA1 ->
   QC.forAll genArray2 $ \blockB2 ->
   QC.forAll genArray2 $ \blockC3 ->
   let shapeR0 = fst $ Array.shape blockA1 in
   let shapeC0 = snd $ Array.shape blockA1 in
   let shapeR1 = fst $ Array.shape blockB2 in
   let shapeC1 = snd $ Array.shape blockB2 in
   let shapeR2 = fst $ Array.shape blockC3 in
   let shapeC2 = snd $ Array.shape blockC3 in
   QC.forAll (genArrayForShape (shapeR0, shapeC1)) $ \blockA2 ->
   QC.forAll (genArrayForShape (shapeR0, shapeC2)) $ \blockA3 ->
   QC.forAll (genArrayForShape (shapeR1, shapeC0)) $ \blockB1 ->
   QC.forAll (genArrayForShape (shapeR1, shapeC2)) $ \blockB3 ->
   QC.forAll (genArrayForShape (shapeR2, shapeC0)) $ \blockC1 ->
   QC.forAll (genArrayForShape (shapeR2, shapeC1)) $ \blockC2 ->
   let blocked =
         BoxedArray.fromList (Set.fromList "ABC", Set.fromList "123")
            [blockA1, blockA2, blockA3,
             blockB1, blockB2, blockB3,
             blockC1, blockC2, blockC3] in

   transpose (Array.fromNonEmptyBlockArray2 blocked)
   QC.===
   Array.fromNonEmptyBlockArray2
      (TestBoxedArray.transpose (fmap transpose blocked))
:}
-}
fromNonEmptyBlockArray2 ::
   (Ord row,    Shape.C height, Eq height) =>
   (Ord column, Shape.C width,  Eq width) =>
   (Storable a) =>
   BoxedArray.Array (Set row, Set column) (Array (height, width) a) ->
   Array (Map row height, Map column width) a
fromNonEmptyBlockArray2 arr =
   let shapes = List.map Array.shape $ BoxedArray.toList arr in
   let width = Set.size $ snd $ BoxedArray.shape arr in
   let (rowIxs, columnIxs) =
         mapPair (Set.toAscList, Set.toAscList) $ BoxedArray.shape arr in
   case (ListHT.sieve width shapes, take width shapes) of
      (leftColumn@(_:_), topRow@(_:_)) ->
         fromBlockArray2
            (Map.fromList $ List.zip rowIxs $ List.map fst leftColumn)
            (Map.fromList $ List.zip columnIxs $ List.map snd topRow)
            arr
      _ -> errorArray "fromNonEmptyBlockArray2" "empty array"

{- |
Explicit parameters for the shape of the result matrix
allow for working with arrays of zero rows or columns.

>>> :{
   (id :: Id (array (height, Map Char ShapeInt) Word16)) $
   Array.fromBlockArray2
      (Map.singleton 'A' (shapeInt 2) <> Map.singleton 'B' (shapeInt 3))
      Map.empty $
   BoxedArray.fromList (Set.fromList "AB", Set.empty) []
:}
StorableArray.fromList (fromList [('A',ZeroBased {... 2}),('B',ZeroBased {... 3})],fromList []) []

prop> :{
   QC.forAll genArray2 $ \block ->
   let height = Map.singleton 'A' $ fst $ Array.shape block in
   let width  = Map.singleton '1' $ snd $ Array.shape block in

   Array.reshape (height,width) block
   QC.===
   Array.fromBlockArray2 height width
      (BoxedArray.replicate (Set.singleton 'A', Set.singleton '1') block)
:}
-}
fromBlockArray2 ::
   (Ord row,    Shape.C height, Eq height) =>
   (Ord column, Shape.C width,  Eq width) =>
   (Storable a) =>
   Map row height -> Map column width ->
   BoxedArray.Array (Set row, Set column) (Array (height, width) a) ->
   Array (Map row height, Map column width) a
fromBlockArray2 height width =
   Array.reshape (height, width) . fromStorableVector .
   SV.concat . List.concat . List.concatMap List.transpose .
   ListHT.sliceVertical (Map.size width) . BoxedArray.toList .
   BoxedArray.zipWith
      (\(h,w) block ->
         if (h,w) == Array.shape block
            then toRowSlices block
            else errorArray "fromBlockArray2" "block shapes mismatch")
      (BoxedArray.cartesian
         (BoxedArray.fromMap height) (BoxedArray.fromMap width))
{-
[[[111,111],[222,222]],[[333,333],[444,444]]]
  |
  v
[111,222,111,222,333,444,333,444]
-}

toRowSlices ::
   (Shape.C sh0, Shape.C sh1, Storable a) =>
   Array (sh0, sh1) a -> [SV.Vector a]
toRowSlices arr =
   SV.sliceVertical (Shape.size $ snd $ shape arr) $ toStorableVector arr


toAssociations ::
   (Shape.Indexed sh, Storable a) => Array sh a -> [(Shape.Index sh, a)]
toAssociations arr = zip (Shape.indices $ shape arr) (Array.toList arr)


errorArray :: String -> String -> a
errorArray name msg =
   error ("Array.Comfort.Storable." ++ name ++ ": " ++ msg)

infixl 9 !

(!) :: (Shape.Indexed sh, Storable a) => Array sh a -> Shape.Index sh -> a
(!) arr = either (errorArray "!") id . accessEither arr

accessMaybe ::
   (Shape.Indexed sh, Storable a) => Array sh a -> Shape.Index sh -> Maybe a
accessMaybe arr = maybeRight . accessEither arr

accessEither ::
   (Shape.Indexed sh, Storable a) =>
   Array sh a -> Shape.Index sh -> Either String a
accessEither arr ix = runST (do
   marr <- MutArrayNC.unsafeThaw arr
   case MutArrayPriv.readEither marr ix of
      Right access -> fmap Right access
      Left msg -> return $ Left msg)
--   for GHC>=7.8: Trav.sequenceA $ MutArrayPriv.readEither marr ix)


zipWith ::
   (Shape.C sh, Eq sh, Storable a, Storable b, Storable c) =>
   (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith f a b =
   if shape a == shape b
      then Array.zipWith f a b
      else errorArray "zipWith" "shapes mismatch"

(//) ::
   (Shape.Indexed sh, Storable a) =>
   Array sh a -> [(Shape.Index sh, a)] -> Array sh a
(//) arr xs = runST (do
   marr <- MutArray.thaw arr
   forM_ xs $ uncurry $ MutArray.write marr
   MutArrayNC.unsafeFreeze marr)

accumulate ::
   (Shape.Indexed sh, Storable a) =>
   (a -> b -> a) -> Array sh a -> [(Shape.Index sh, b)] -> Array sh a
accumulate f arr xs = runST (do
   marr <- MutArray.thaw arr
   forM_ xs $ \(ix,b) -> MutArray.update marr ix $ flip f b
   MutArrayNC.unsafeFreeze marr)

fromAssociations ::
   (Shape.Indexed sh, Storable a) =>
   a -> sh -> [(Shape.Index sh, a)] -> Array sh a
fromAssociations a sh xs = runST (do
   marr <- MutArray.new sh a
   forM_ xs $ uncurry $ MutArray.write marr
   MutArrayNC.unsafeFreeze marr)


{- |
prop> QC.forAll genNonEmptyArray2 $ \xs -> QC.forAll (QC.elements $ Shape.indices $ Array.shape xs) $ \(ix0,ix1) -> Array.pick xs ix0 ! ix1 == xs!(ix0,ix1)
-}
pick ::
   (Shape.Indexed sh0, Shape.C sh1, Storable a) =>
   Array (sh0,sh1) a -> Shape.Index sh0 -> Array sh1 a
pick (Array (sh0,sh1) x) ix0 =
   Array.unsafeCreateWithSize sh1 $ \k yPtr ->
   withForeignPtr x $ \xPtr ->
      copyArray yPtr (advancePtr xPtr (Shape.offset sh0 ix0 * k)) k

toRowArray ::
   (Shape.Indexed sh0, Shape.C sh1, Storable a) =>
   Array (sh0,sh1) a -> BoxedArray.Array sh0 (Array sh1 a)
toRowArray x = fmap (pick x) $ BoxedArray.indices $ fst $ Array.shape x

{- |
It is a checked error if a row width differs from the result array width.

prop> QC.forAll genArray2 $ \xs -> xs == Array.fromRowArray (snd $ Array.shape xs) (Array.toRowArray xs)
-}
fromRowArray ::
   (Shape.C sh0, Shape.C sh1, Eq sh1, Storable a) =>
   sh1 -> BoxedArray.Array sh0 (Array sh1 a) -> Array (sh0,sh1) a
fromRowArray sh1 x =
   Array.unsafeCreate (BoxedArray.shape x, sh1) $ \yPtr ->
   let k = Shape.size sh1 in
   forM_ (zip [0,k..] (BoxedArray.toList x)) $ \(j, Array sh1i row) ->
   if sh1 == sh1i
      then withForeignPtr row $ \xPtr -> copyArray (advancePtr yPtr j) xPtr k
      else errorArray "fromRowArray" "mismatching row width"


{- |
It is a checked error if the vector is empty.

prop> forAllNonEmpty $ \xs -> Array.minimum xs ==? minimum (Array.toList xs)
-}
minimum :: (Shape.C sh, Storable a, Ord a) => Array sh a -> a
minimum = foldl1 min

{- |
It is a checked error if the vector is empty.

prop> forAllNonEmpty $ \xs -> Array.maximum xs ==? maximum (Array.toList xs)
-}
maximum :: (Shape.C sh, Storable a, Ord a) => Array sh a -> a
maximum = foldl1 max

{-# INLINE foldl1 #-}
foldl1 :: (Shape.C sh, Storable a) => (a -> a -> a) -> Array sh a -> a
foldl1 op (Array sh x) = unsafePerformIO $
   withForeignPtr x $ \xPtr ->
      Memory.foldl1 (const id) op (Shape.size sh) xPtr 1

{- |
prop> forAllNonEmpty $ \xs -> Array.limits xs ==? (Array.minimum xs, Array.maximum xs)
-}
limits :: (Shape.C sh, Storable a, Ord a) => Array sh a -> (a,a)
limits = StrictTuple.mapPair (getMin, getMax) . foldMap (\x -> (Min x, Max x))

{-# INLINE foldMap #-}
foldMap ::
   (Shape.C sh, Storable a, Ord a, Semigroup m) => (a -> m) -> Array sh a -> m
foldMap f (Array sh x) = unsafePerformIO $
   withForeignPtr x $ \xPtr ->
      Memory.foldl1 (const f) (<>) (Shape.size sh) xPtr 1


argMinimum, argMaximum ::
   (Shape.InvIndexed sh, Storable a, Ord a) =>
   Array sh a -> (Shape.Index sh, a)
argMinimum xs = unArg xs $ getMin $ foldMapWithIndex (\k x -> Min (Arg x k)) xs
argMaximum xs = unArg xs $ getMax $ foldMapWithIndex (\k x -> Max (Arg x k)) xs

unArg ::
   (Shape.InvIndexed sh) => Array sh a -> Arg a Int -> (Shape.Index sh, a)
unArg xs (Arg x k) = (Shape.indexFromOffset (Array.shape xs) k, x)

{-# INLINE foldMapWithIndex #-}
foldMapWithIndex ::
   (Shape.C sh, Storable a, Semigroup m) => (Int -> a -> m) -> Array sh a -> m
foldMapWithIndex f (Array sh x) = unsafePerformIO $
   withForeignPtr x $ \xPtr -> Memory.foldl1 f (<>) (Shape.size sh) xPtr 1