{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{- |
Can be an alternative to the @enumset@ package.
-}
module Data.Array.Comfort.Bool (
   Array,
   shape,
   reshape,
   mapShape,

   fromList,
   toList,
   fromSet,
   toSet,

   member,
   union,
   difference,
   intersection,
   ) where

import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.Array.Comfort.Check as Check

import qualified Data.IntSet as IntSet
import qualified Data.Set as Set
import qualified Data.List as List
import Data.IntSet (IntSet)
import Data.Set (Set)


data Array sh =
   Array {
      forall sh. Array sh -> sh
shape_ :: sh,
      forall sh. Array sh -> IntSet
_intSet :: IntSet
   }


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

reshape :: (Shape.C sh0, Shape.C sh1) => sh1 -> Array sh0 -> Array sh1
reshape :: forall sh0 sh1. (C sh0, C sh1) => sh1 -> Array sh0 -> Array sh1
reshape = String
-> (Array sh0 -> sh0)
-> (sh1 -> Array sh0 -> Array sh1)
-> sh1
-> Array sh0
-> Array sh1
forall sh0 sh1 array0 array1.
(C sh0, C sh1) =>
String
-> (array0 -> sh0)
-> (sh1 -> array0 -> array1)
-> sh1
-> array0
-> array1
Check.reshape String
"Storable" Array sh0 -> sh0
forall sh. Array sh -> sh
shape (\sh1
sh Array sh0
arr -> Array sh0
arr{shape_ :: sh1
shape_ = sh1
sh})

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



fromList :: (Shape.Indexed sh) => sh -> [Shape.Index sh] -> Array sh
fromList :: forall sh. Indexed sh => sh -> [Index sh] -> Array sh
fromList sh
sh = sh -> IntSet -> Array sh
forall sh. sh -> IntSet -> Array sh
Array sh
sh (IntSet -> Array sh)
-> ([Index sh] -> IntSet) -> [Index sh] -> Array sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Key] -> IntSet
IntSet.fromList ([Key] -> IntSet) -> ([Index sh] -> [Key]) -> [Index sh] -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Index sh -> Key) -> [Index sh] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
List.map (sh -> Index sh -> Key
forall sh. Indexed sh => sh -> Index sh -> Key
Shape.offset sh
sh)

toList :: (Shape.InvIndexed sh) => Array sh -> [Shape.Index sh]
toList :: forall sh. InvIndexed sh => Array sh -> [Index sh]
toList (Array sh
sh IntSet
set) = (Key -> Index sh) -> [Key] -> [Index sh]
forall a b. (a -> b) -> [a] -> [b]
map (sh -> Key -> Index sh
forall sh. InvIndexed sh => sh -> Key -> Index sh
Shape.indexFromOffset sh
sh) ([Key] -> [Index sh]) -> [Key] -> [Index sh]
forall a b. (a -> b) -> a -> b
$ IntSet -> [Key]
IntSet.toList IntSet
set

fromSet ::
   (Shape.Indexed sh, Shape.Index sh ~ ix, Ord ix) => sh -> Set ix -> Array sh
fromSet :: forall sh ix.
(Indexed sh, Index sh ~ ix, Ord ix) =>
sh -> Set ix -> Array sh
fromSet sh
sh = sh -> [Index sh] -> Array sh
forall sh. Indexed sh => sh -> [Index sh] -> Array sh
fromList sh
sh ([ix] -> Array sh) -> (Set ix -> [ix]) -> Set ix -> Array sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set ix -> [ix]
forall a. Set a -> [a]
Set.toList

toSet ::
   (Shape.InvIndexed sh, Shape.Index sh ~ ix, Ord ix) => Array sh -> Set ix
toSet :: forall sh ix.
(InvIndexed sh, Index sh ~ ix, Ord ix) =>
Array sh -> Set ix
toSet = [ix] -> Set ix
forall a. Ord a => [a] -> Set a
Set.fromList ([ix] -> Set ix) -> (Array sh -> [ix]) -> Array sh -> Set ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh -> [ix]
Array sh -> [Index sh]
forall sh. InvIndexed sh => Array sh -> [Index sh]
toList


errorArray :: String -> String -> a
errorArray :: forall a. String -> String -> a
errorArray String
name String
msg =
   String -> a
forall a. HasCallStack => String -> a
error (String
"Array.Comfort.Bool." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
msg)


member :: (Shape.Indexed sh) => Shape.Index sh -> Array sh -> Bool
member :: forall sh. Indexed sh => Index sh -> Array sh -> Bool
member Index sh
ix (Array sh
sh IntSet
set) = Key -> IntSet -> Bool
IntSet.member (sh -> Index sh -> Key
forall sh. Indexed sh => sh -> Index sh -> Key
Shape.offset sh
sh Index sh
ix) IntSet
set


lift2 :: (Shape.Indexed sh, Eq sh) =>
   String -> (IntSet -> IntSet -> IntSet) ->
   Array sh -> Array sh -> Array sh
lift2 :: forall sh.
(Indexed sh, Eq sh) =>
String
-> (IntSet -> IntSet -> IntSet) -> Array sh -> Array sh -> Array sh
lift2 String
name IntSet -> IntSet -> IntSet
op (Array sh
shA IntSet
setA) (Array sh
shB IntSet
setB) =
   if sh
shA sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
== sh
shB
      then sh -> IntSet -> Array sh
forall sh. sh -> IntSet -> Array sh
Array sh
shA (IntSet -> Array sh) -> IntSet -> Array sh
forall a b. (a -> b) -> a -> b
$ IntSet -> IntSet -> IntSet
op IntSet
setA IntSet
setB
      else String -> String -> Array sh
forall a. String -> String -> a
errorArray String
name String
"shapes mismatch"

union :: (Shape.Indexed sh, Eq sh) => Array sh -> Array sh -> Array sh
union :: forall sh. (Indexed sh, Eq sh) => Array sh -> Array sh -> Array sh
union = String
-> (IntSet -> IntSet -> IntSet) -> Array sh -> Array sh -> Array sh
forall sh.
(Indexed sh, Eq sh) =>
String
-> (IntSet -> IntSet -> IntSet) -> Array sh -> Array sh -> Array sh
lift2 String
"union" IntSet -> IntSet -> IntSet
IntSet.union

intersection :: (Shape.Indexed sh, Eq sh) => Array sh -> Array sh -> Array sh
intersection :: forall sh. (Indexed sh, Eq sh) => Array sh -> Array sh -> Array sh
intersection = String
-> (IntSet -> IntSet -> IntSet) -> Array sh -> Array sh -> Array sh
forall sh.
(Indexed sh, Eq sh) =>
String
-> (IntSet -> IntSet -> IntSet) -> Array sh -> Array sh -> Array sh
lift2 String
"intersection" IntSet -> IntSet -> IntSet
IntSet.intersection

difference :: (Shape.Indexed sh, Eq sh) => Array sh -> Array sh -> Array sh
difference :: forall sh. (Indexed sh, Eq sh) => Array sh -> Array sh -> Array sh
difference = String
-> (IntSet -> IntSet -> IntSet) -> Array sh -> Array sh -> Array sh
forall sh.
(Indexed sh, Eq sh) =>
String
-> (IntSet -> IntSet -> IntSet) -> Array sh -> Array sh -> Array sh
lift2 String
"difference" IntSet -> IntSet -> IntSet
IntSet.difference