{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
module Data.Array.Comfort.Shape.Extra
(
Simplex(..),
) where
import qualified Data.Array.Comfort.Shape as Shape
import qualified Type.Data.Num.Unary as Unary
import Type.Data.Num (integralFromProxy)
import Type.Base.Proxy (Proxy(Proxy))
import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.FixedLength as FL
import qualified Data.NonEmpty as NonEmpty
import qualified Data.List.HT as ListHT
import Control.Applicative ((<$>))
data Simplex d size =
Simplex {
Simplex d size -> UnaryProxy d
simplexDimension :: UnaryProxy d,
Simplex d size -> size
simplexSize :: size
} deriving (Int -> Simplex d size -> ShowS
[Simplex d size] -> ShowS
Simplex d size -> String
(Int -> Simplex d size -> ShowS)
-> (Simplex d size -> String)
-> ([Simplex d size] -> ShowS)
-> Show (Simplex d size)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall d size.
(Natural d, Show size) =>
Int -> Simplex d size -> ShowS
forall d size. (Natural d, Show size) => [Simplex d size] -> ShowS
forall d size. (Natural d, Show size) => Simplex d size -> String
showList :: [Simplex d size] -> ShowS
$cshowList :: forall d size. (Natural d, Show size) => [Simplex d size] -> ShowS
show :: Simplex d size -> String
$cshow :: forall d size. (Natural d, Show size) => Simplex d size -> String
showsPrec :: Int -> Simplex d size -> ShowS
$cshowsPrec :: forall d size.
(Natural d, Show size) =>
Int -> Simplex d size -> ShowS
Show)
type UnaryProxy d = Proxy (Unary.Un d)
binomials :: Integral a => a -> [a]
binomials :: a -> [a]
binomials a
n =
(a -> (a, a) -> a) -> a -> [(a, a)] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (\a
acc (a
num,a
den) -> a -> a -> a
forall a. Integral a => a -> a -> a
div (a
acca -> a -> a
forall a. Num a => a -> a -> a
*a
num) a
den) a
1 ([a] -> [a] -> [(a, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a
n, a -> a
forall a. Enum a => a -> a
pred a
n ..] [a
1..a
n])
simplexLayoutSize :: Integral i => Int -> i -> i
simplexLayoutSize :: Int -> i -> i
simplexLayoutSize Int
d i
n =
case Int -> [i] -> [i]
forall a. Int -> [a] -> [a]
drop Int
d ([i] -> [i]) -> [i] -> [i]
forall a b. (a -> b) -> a -> b
$ i -> [i]
forall a. Integral a => a -> [a]
binomials i
n of
[] -> i
0
i
m:[i]
_ -> i
m
instance (Unary.Natural d, Shape.C size) => Shape.C (Simplex d size) where
size :: Simplex d size -> Int
size (Simplex UnaryProxy d
d size
sz) =
Int -> Int -> Int
forall i. Integral i => Int -> i -> i
simplexLayoutSize (UnaryProxy d -> Int
forall x y. (Integer x, Num y) => Proxy x -> y
integralFromProxy UnaryProxy d
d) (size -> Int
forall sh. C sh => sh -> Int
Shape.size size
sz)
headSingletonFromProxy ::
(Unary.Natural d) => UnaryProxy d -> Unary.HeadSingleton d
headSingletonFromProxy :: UnaryProxy d -> HeadSingleton d
headSingletonFromProxy UnaryProxy d
Proxy = HeadSingleton d
forall n. Natural n => HeadSingleton n
Unary.headSingleton
predHeadSingleton :: Unary.HeadSingleton (Unary.Succ d) -> UnaryProxy d
predHeadSingleton :: HeadSingleton (Succ d) -> UnaryProxy d
predHeadSingleton HeadSingleton (Succ d)
Unary.Succ = UnaryProxy d
forall a. Proxy a
Proxy
simplexIndices :: (Unary.Natural d) => UnaryProxy d -> [a] -> [FL.T d a]
simplexIndices :: UnaryProxy d -> [a] -> [T d a]
simplexIndices UnaryProxy d
d =
case UnaryProxy d -> HeadSingleton d
forall d. Natural d => UnaryProxy d -> HeadSingleton d
headSingletonFromProxy UnaryProxy d
d of
HeadSingleton d
Unary.Zero -> [T Zero a] -> [a] -> [T Zero a]
forall a b. a -> b -> a
const [T Zero a
forall a. T Zero a
FL.end]
m :: HeadSingleton d
m@HeadSingleton d
Unary.Succ -> \[a]
as -> do
(a
a,[a]
ts) <- [a] -> [[a]] -> [(a, [a])]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
as ([[a]] -> [(a, [a])]) -> [[a]] -> [(a, [a])]
forall a b. (a -> b) -> a -> b
$ T [] [a] -> [[a]]
forall (f :: * -> *) a. T f a -> f a
NonEmpty.tail (T [] [a] -> [[a]]) -> T [] [a] -> [[a]]
forall a b. (a -> b) -> a -> b
$ [a] -> T [] [a]
forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Cons g, Empty g) =>
f a -> T f (g a)
NonEmpty.tails [a]
as
(a
a a -> T n1 a -> T (Succ n1) a
forall a n. a -> T n a -> T (Succ n) a
FL.!:) (T n1 a -> T (Succ n1) a) -> [T n1 a] -> [T (Succ n1) a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> UnaryProxy n1 -> [a] -> [T n1 a]
forall d a. Natural d => UnaryProxy d -> [a] -> [T d a]
simplexIndices (HeadSingleton (Succ n1) -> UnaryProxy n1
forall d. HeadSingleton (Succ d) -> UnaryProxy d
predHeadSingleton HeadSingleton d
HeadSingleton (Succ n1)
m) [a]
ts
instance
(Unary.Natural d, Shape.Indexed size) =>
Shape.Indexed (Simplex d size) where
type Index (Simplex d size) = FL.T d (Shape.Index size)
indices :: Simplex d size -> [Index (Simplex d size)]
indices (Simplex UnaryProxy d
d size
sz) = UnaryProxy d -> [Index size] -> [T d (Index size)]
forall d a. Natural d => UnaryProxy d -> [a] -> [T d a]
simplexIndices UnaryProxy d
d ([Index size] -> [T d (Index size)])
-> [Index size] -> [T d (Index size)]
forall a b. (a -> b) -> a -> b
$ size -> [Index size]
forall sh. Indexed sh => sh -> [Index sh]
Shape.indices size
sz
inBounds :: Simplex d size -> Index (Simplex d size) -> Bool
inBounds (Simplex UnaryProxy d
_d size
sz) Index (Simplex d size)
ix =
let ixs :: [Index size]
ixs = T d (Index size) -> [Index size]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Fold.toList Index (Simplex d size)
T d (Index size)
ix
getOffset :: Index size -> Int
getOffset = size -> Index size -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset size
sz
in (Index size -> Bool) -> [Index size] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (size -> Index size -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds size
sz) [Index size]
ixs Bool -> Bool -> Bool
&&
[Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Int -> Int -> Bool) -> [Int] -> [Bool]
forall a b. (a -> a -> b) -> [a] -> [b]
ListHT.mapAdjacent Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ([Int] -> [Bool]) -> [Int] -> [Bool]
forall a b. (a -> b) -> a -> b
$ (Index size -> Int) -> [Index size] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Index size -> Int
getOffset [Index size]
ixs)
unifiedSizeOffset :: Simplex d size -> (Int, Index (Simplex d size) -> Result check Int)
unifiedSizeOffset (Simplex UnaryProxy d
d size
sz) =
let (Int
n, Index size -> Result check Int
getOffset) = size -> (Int, Index size -> Result check Int)
forall sh check.
(Indexed sh, Checking check) =>
sh -> (Int, Index sh -> Result check Int)
Shape.unifiedSizeOffset size
sz
dInt :: Int
dInt = UnaryProxy d -> Int
forall x y. (Integer x, Num y) => Proxy x -> y
integralFromProxy UnaryProxy d
d
in (Int -> Int -> Int
forall i. Integral i => Int -> i -> i
simplexLayoutSize Int
dInt Int
n,
\Index (Simplex d size)
ixs -> do
[Int]
ks <- (Index size -> Result check Int)
-> [Index size] -> Result check [Int]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
Trav.traverse Index size -> Result check Int
getOffset ([Index size] -> Result check [Int])
-> [Index size] -> Result check [Int]
forall a b. (a -> b) -> a -> b
$ T d (Index size) -> [Index size]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Fold.toList Index (Simplex d size)
T d (Index size)
ixs
Int -> Result check Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Result check Int) -> Int -> Result check Int
forall a b. (a -> b) -> a -> b
$
Int -> Int -> Int
forall i. Integral i => Int -> i -> i
simplexLayoutSize Int
dInt Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
Int -> Int -> Int
forall a. Num a => a -> a -> a
-
[Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall i. Integral i => Int -> i -> i
simplexLayoutSize
((Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate Int -> Int
forall a. Enum a => a -> a
pred Int
dInt) ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
-) [Int]
ks)))
instance
(Unary.Natural d, Shape.InvIndexed size) =>
Shape.InvIndexed (Simplex d size) where
unifiedIndexFromOffset :: Simplex d size -> Int -> Result check (Index (Simplex d size))
unifiedIndexFromOffset (Simplex UnaryProxy d
d size
sh) Int
k =
let n :: Int
n = size -> Int
forall sh. C sh => sh -> Int
Shape.size size
sh in
let dInt :: Int
dInt = UnaryProxy d -> Int
forall x y. (Integer x, Num y) => Proxy x -> y
integralFromProxy UnaryProxy d
d in
T d (Result check (Index size)) -> Result check (T d (Index size))
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
Trav.sequenceA (T d (Result check (Index size))
-> Result check (T d (Index size)))
-> T d (Result check (Index size))
-> Result check (T d (Index size))
forall a b. (a -> b) -> a -> b
$ ((Int, Int), T d (Result check (Index size)))
-> T d (Result check (Index size))
forall a b. (a, b) -> b
snd (((Int, Int), T d (Result check (Index size)))
-> T d (Result check (Index size)))
-> ((Int, Int), T d (Result check (Index size)))
-> T d (Result check (Index size))
forall a b. (a -> b) -> a -> b
$
((Int, Int) -> Int -> ((Int, Int), Result check (Index size)))
-> (Int, Int)
-> T d Int
-> ((Int, Int), T d (Result check (Index size)))
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
Trav.mapAccumL
(\(Int
a,Int
k0) Int
m ->
case ((Int, Int) -> Bool) -> [(Int, Int)] -> [(Int, Int)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0) (Int -> Bool) -> ((Int, Int) -> Int) -> (Int, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> Int
forall a b. (a, b) -> b
snd) ([(Int, Int)] -> [(Int, Int)]) -> [(Int, Int)] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$
(Int -> (Int, Int)) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
bi -> (Int
bi, Int
k0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int -> Int
forall i. Integral i => Int -> i -> i
simplexLayoutSize Int
m (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
biInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))) ([Int] -> [(Int, Int)]) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$
(Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
n) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+) Int
a of
[] -> String -> ((Int, Int), Result check (Index size))
forall a. HasCallStack => String -> a
error String
"unifiedIndexFromOffset: offset out of range"
(Int
b,Int
k1):[(Int, Int)]
_ -> ((Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, Int
k1), size -> Int -> Result check (Index size)
forall sh check.
(InvIndexed sh, Checking check) =>
sh -> Int -> Result check (Index sh)
Shape.unifiedIndexFromOffset size
sh Int
b))
(Int
0, Int -> Int -> Int
forall i. Integral i => Int -> i -> i
simplexLayoutSize Int
dInt Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k)
(T (T d) Int -> T d Int
forall (f :: * -> *) a. Traversable f => T f a -> f a
NonEmpty.init (T (T d) Int -> T d Int) -> T (T d) Int -> T d Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Int -> T d Int -> T (T d) Int
forall (f :: * -> *) b a.
Traversable f =>
(b -> a -> b) -> b -> f a -> T f b
NonEmpty.scanl (-) Int
dInt (T d Int -> T (T d) Int) -> T d Int -> T (T d) Int
forall a b. (a -> b) -> a -> b
$ Int -> T d Int
forall n a. Natural n => a -> T n a
FL.repeat Int
1)