module Data.Array.Comfort.Boxed (
   Array,
   shape,
   reshape,
   mapShape,
   accessMaybe, (!),
   Array.toList,
   Array.fromList,
   Array.vectorFromList,
   toAssociations,
   fromMap,
   toMap,
   fromTuple,
   toTuple,
   fromRecord,
   toRecord,
   fromContainer,
   toContainer,
   indices,
   Array.replicate,

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

   pick,
   Array.append,
   Array.take, Array.drop,
   Array.takeLeft, Array.takeRight, Array.split,
   Array.takeCenter,
   ) where

import qualified Data.Array.Comfort.Boxed.Unchecked as Array
import qualified Data.Array.Comfort.Container as Container
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.Boxed.Unchecked (Array(Array))

import qualified Data.Primitive.Array as Prim

import qualified Control.Monad.Primitive as PrimM
import qualified Control.Monad.Trans.State as MS
import Control.Monad.ST (runST)
import Control.Applicative ((<$>))

import qualified Data.Foldable as Fold
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Map (Map)
import Data.Set (Set)
import Data.Traversable (Traversable, traverse)
import Data.Foldable (forM_)
import Data.Either.HT (maybeRight)
import Data.Maybe (fromMaybe)

import Prelude hiding (zipWith, replicate)


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

reshape :: (Shape.C sh0, Shape.C sh1) => sh1 -> Array sh0 a -> Array sh1 a
reshape :: forall sh0 sh1 a.
(C sh0, C sh1) =>
sh1 -> Array sh0 a -> Array sh1 a
reshape = forall sh0 sh1 array0 array1.
(C sh0, C sh1) =>
String
-> (array0 -> sh0)
-> (sh1 -> array0 -> array1)
-> sh1
-> array0
-> array1
Check.reshape String
"Boxed" forall sh a. Array sh a -> sh
shape forall sh1 sh0 a. sh1 -> Array sh0 a -> Array sh1 a
Array.reshape

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


indices :: (Shape.Indexed sh) => sh -> Array.Array sh (Shape.Index sh)
indices :: forall sh. Indexed sh => sh -> Array sh (Index sh)
indices sh
sh = forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList sh
sh forall a b. (a -> b) -> a -> b
$ forall sh. Indexed sh => sh -> [Index sh]
Shape.indices sh
sh

fromMap :: (Ord k) => Map k a -> Array (Set k) a
fromMap :: forall k a. Ord k => Map k a -> Array (Set k) a
fromMap Map k a
m = forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList (forall k a. Map k a -> Set k
Map.keysSet Map k a
m) (forall k a. Map k a -> [a]
Map.elems Map k a
m)

toMap :: (Ord k) => Array (Set k) a -> Map k a
toMap :: forall k a. Ord k => Array (Set k) a -> Map k a
toMap Array (Set k) a
arr = forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a. Set a -> [a]
Set.toAscList forall a b. (a -> b) -> a -> b
$ forall sh a. Array sh a -> sh
shape Array (Set k) a
arr) (forall sh a. C sh => Array sh a -> [a]
Array.toList Array (Set k) a
arr)

fromTuple ::
   (TupleShape.NestedTuple tuple) =>
   Shape.DataTuple tuple a -> Array (Shape.NestedTuple ixtype tuple) a
fromTuple :: forall tuple a ixtype.
NestedTuple tuple =>
DataTuple tuple a -> Array (NestedTuple ixtype tuple) a
fromTuple DataTuple tuple a
tuple =
   case forall s a. State s a -> s -> a
MS.evalState (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
TupleShape.decons DataTuple tuple a
tuple) (Int -> Element
Shape.Element Int
0) of
      (tuple
sh, [a]
xs) -> forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList (forall ixtype tuple. tuple -> NestedTuple ixtype tuple
Shape.NestedTuple tuple
sh) [a]
xs

toTuple ::
   (TupleShape.NestedTuple tuple) =>
   Array (Shape.NestedTuple ixtype tuple) a -> Shape.DataTuple tuple a
toTuple :: forall tuple ixtype a.
NestedTuple tuple =>
Array (NestedTuple ixtype tuple) a -> DataTuple tuple a
toTuple Array (NestedTuple ixtype tuple) a
arr =
   forall s a. State s a -> s -> a
MS.evalState
      (forall shape a.
ElementTuple shape =>
shape -> State [a] (DataTuple shape a)
TupleShape.cons forall a b. (a -> b) -> a -> b
$ forall ixtype tuple. NestedTuple ixtype tuple -> tuple
Shape.getNestedTuple forall a b. (a -> b) -> a -> b
$ forall sh a. Array sh a -> sh
shape Array (NestedTuple ixtype tuple) a
arr)
      (forall sh a. C sh => Array sh a -> [a]
Array.toList Array (NestedTuple ixtype tuple) a
arr)

fromRecord ::
   (Traversable f) =>
   f a -> Array (Shape.Record f) a
fromRecord :: forall (f :: * -> *) a. Traversable f => f a -> Array (Record f) a
fromRecord f a
xs =
   forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList
      (forall (f :: * -> *). f Element -> Record f
Shape.Record forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> a
MS.evalState (Int -> Element
Shape.Element Int
0) forall a b. (a -> b) -> a -> b
$
       forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall a b. a -> b -> a
const State Element Element
TupleShape.next) f a
xs)
      (forall (t :: * -> *) a. Foldable t => t a -> [a]
Fold.toList f a
xs)

toRecord ::
   (Traversable f) =>
   Array (Shape.Record f) a -> f a
toRecord :: forall (f :: * -> *) a. Traversable f => Array (Record f) a -> f a
toRecord Array (Record f) a
arr =
   forall s a. State s a -> s -> a
MS.evalState
      (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall a b. a -> b -> a
const forall a. State [a] a
TupleShape.get) forall a b. (a -> b) -> a -> b
$
       (\(Shape.Record f Element
record) -> f Element
record) forall a b. (a -> b) -> a -> b
$ forall sh a. Array sh a -> sh
shape Array (Record f) a
arr)
      (forall sh a. C sh => Array sh a -> [a]
Array.toList Array (Record f) a
arr)

fromContainer :: (Container.C f) => f a -> Array (Container.Shape f) a
fromContainer :: forall (f :: * -> *) a. C f => f a -> Array (Shape f) a
fromContainer f a
xs = forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList (forall (f :: * -> *) a. C f => f a -> Shape f
Container.toShape f a
xs) (forall (t :: * -> *) a. Foldable t => t a -> [a]
Fold.toList f a
xs)

toContainer :: (Container.C f) => Array (Container.Shape f) a -> f a
toContainer :: forall (f :: * -> *) a. C f => Array (Shape f) a -> f a
toContainer Array (Shape f) a
arr = forall (f :: * -> *) a. C f => Shape f -> [a] -> f a
Container.fromList (forall sh a. Array sh a -> sh
Array.shape Array (Shape f) a
arr) (forall sh a. C sh => Array sh a -> [a]
Array.toList Array (Shape f) a
arr)


infixl 9 !

(!) :: (Shape.Indexed sh) => Array sh a -> Shape.Index sh -> a
! :: forall sh a. Indexed sh => Array sh a -> Index sh -> a
(!) Array sh a
arr Index sh
ix =
   forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => String -> a
error String
"Array.Comfort.Boxed.!: index out of bounds") forall a b. (a -> b) -> a -> b
$
   forall sh a. Indexed sh => Array sh a -> Index sh -> Maybe a
accessMaybe Array sh a
arr Index sh
ix

accessMaybe :: (Shape.Indexed sh) => Array sh a -> Shape.Index sh -> Maybe a
accessMaybe :: forall sh a. Indexed sh => Array sh a -> Index sh -> Maybe a
accessMaybe (Array sh
sh Array a
arr) Index sh
ix =
   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Array a -> Int -> a
Prim.indexArray Array a
arr) forall a b. (a -> b) -> a -> b
$ forall a b. Either a b -> Maybe b
maybeRight forall a b. (a -> b) -> a -> b
$
   forall a. Result Checked a -> Either String a
Shape.getChecked forall a b. (a -> b) -> a -> b
$ forall sh check.
(Indexed sh, Checking check) =>
sh -> Index sh -> Result check Int
Shape.unifiedOffset sh
sh Index sh
ix


zipWith ::
   (Shape.C sh, Eq sh) =>
   (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith :: forall sh a b c.
(C sh, Eq sh) =>
(a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith a -> b -> c
f Array sh a
a Array sh b
b =
   if forall sh a. Array sh a -> sh
shape Array sh a
a forall a. Eq a => a -> a -> Bool
== forall sh a. Array sh a -> sh
shape Array sh b
b
      then forall sh a b c.
C sh =>
(a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
Array.zipWith a -> b -> c
f Array sh a
a Array sh b
b
      else forall a. HasCallStack => String -> a
error String
"zipWith: shapes mismatch"


(//) ::
   (Shape.Indexed sh) => Array sh a -> [(Shape.Index sh, a)] -> Array sh a
// :: forall sh a.
Indexed sh =>
Array sh a -> [(Index sh, a)] -> Array sh a
(//) (Array sh
sh Array a
arr) [(Index sh, a)]
xs = forall a. (forall s. ST s a) -> a
runST (do
   MutableArray s a
marr <- forall (m :: * -> *) a.
PrimMonad m =>
Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
Prim.thawArray Array a
arr Int
0 (forall sh. C sh => sh -> Int
Shape.size sh
sh)
   forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Index sh, a)]
xs forall a b. (a -> b) -> a -> b
$ \(Index sh
ix,a
a) -> forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Prim.writeArray MutableArray s a
marr (forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh
sh Index sh
ix) a
a
   forall sh a. sh -> Array a -> Array sh a
Array sh
sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Prim.unsafeFreezeArray MutableArray s a
marr)

accumulate ::
   (Shape.Indexed sh) =>
   (a -> b -> a) -> Array sh a -> [(Shape.Index sh, b)] -> Array sh a
accumulate :: forall sh a b.
Indexed sh =>
(a -> b -> a) -> Array sh a -> [(Index sh, b)] -> Array sh a
accumulate a -> b -> a
f (Array sh
sh Array a
arr) [(Index sh, b)]
xs = forall a. (forall s. ST s a) -> a
runST (do
   MutableArray s a
marr <- forall (m :: * -> *) a.
PrimMonad m =>
Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
Prim.thawArray Array a
arr Int
0 (forall sh. C sh => sh -> Int
Shape.size sh
sh)
   forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Index sh, b)]
xs forall a b. (a -> b) -> a -> b
$ \(Index sh
ix,b
b) -> forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> (a -> a) -> m ()
updateArray MutableArray s a
marr (forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh
sh Index sh
ix) forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> a
f b
b
   forall sh a. sh -> Array a -> Array sh a
Array sh
sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Prim.unsafeFreezeArray MutableArray s a
marr)

updateArray ::
   PrimM.PrimMonad m =>
   Prim.MutableArray (PrimM.PrimState m) a -> Int -> (a -> a) -> m ()
updateArray :: forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> (a -> a) -> m ()
updateArray MutableArray (PrimState m) a
marr Int
k a -> a
f = forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Prim.writeArray MutableArray (PrimState m) a
marr Int
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
Prim.readArray MutableArray (PrimState m) a
marr Int
k

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

fromAssociations ::
   (Shape.Indexed sh) => a -> sh -> [(Shape.Index sh, a)] -> Array sh a
fromAssociations :: forall sh a. Indexed sh => a -> sh -> [(Index sh, a)] -> Array sh a
fromAssociations a
a sh
sh [(Index sh, a)]
xs = forall a. (forall s. ST s a) -> a
runST (do
   MutableArray s a
marr <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
Prim.newArray (forall sh. C sh => sh -> Int
Shape.size sh
sh) a
a
   forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Index sh, a)]
xs forall a b. (a -> b) -> a -> b
$ \(Index sh
ix,a
x) -> forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Prim.writeArray MutableArray s a
marr (forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh
sh Index sh
ix) a
x
   forall sh a. sh -> Array a -> Array sh a
Array sh
sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Prim.unsafeFreezeArray MutableArray s a
marr)



pick ::
   (Shape.Indexed sh0, Shape.C sh1) =>
   Array (sh0,sh1) a -> Shape.Index sh0 -> Array sh1 a
pick :: forall sh0 sh1 a.
(Indexed sh0, C sh1) =>
Array (sh0, sh1) a -> Index sh0 -> Array sh1 a
pick (Array (sh0
sh0,sh1
sh1) Array a
x) Index sh0
ix0 =
   forall sh a. sh -> Array a -> Array sh a
Array sh1
sh1 forall a b. (a -> b) -> a -> b
$
   let k :: Int
k = forall sh. C sh => sh -> Int
Shape.size sh1
sh1
   in forall a. Array a -> Int -> Int -> Array a
Prim.cloneArray Array a
x (forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh0
sh0 Index sh0
ix0 forall a. Num a => a -> a -> a
* Int
k) Int
k