{-# LANGUAGE UndecidableInstances #-}

module Telescope.Data.DataCube where

import Data.Kind
import Data.Massiv.Array as M hiding (mapM)
import Data.Proxy
import GHC.TypeLits (natVal)
import Telescope.Data.Array (AxesIndex (..))
import Telescope.Data.Axes (Axes, Major (Row))


-- Results ------------------------------------------------------------------------------

newtype DataCube (as :: [Type]) f = DataCube
  { forall (as :: [*]) f. DataCube as f -> Array D (IndexOf as) f
array :: Array D (IndexOf as) f
  }


instance (Index (IndexOf as), Eq f) => Eq (DataCube as f) where
  DataCube Array D (IndexOf as) f
arr == :: DataCube as f -> DataCube as f -> Bool
== DataCube Array D (IndexOf as) f
arr2 = Array D (IndexOf as) f
arr Array D (IndexOf as) f -> Array D (IndexOf as) f -> Bool
forall a. Eq a => a -> a -> Bool
== Array D (IndexOf as) f
arr2


instance (Ragged L (IndexOf as) f, Show f) => Show (DataCube as f) where
  show :: DataCube as f -> String
show (DataCube Array D (IndexOf as) f
a) = Array D (IndexOf as) f -> String
forall a. Show a => a -> String
show Array D (IndexOf as) f
a


class HasIndex (as :: [Type]) where
  type IndexOf as :: Type


instance HasIndex '[] where
  type IndexOf '[] = Ix0
instance HasIndex '[a] where
  type IndexOf '[a] = Ix1
instance HasIndex '[a, b] where
  type IndexOf '[a, b] = Ix2
instance HasIndex '[a, b, c] where
  type IndexOf '[a, b, c] = Ix3
instance HasIndex '[a, b, c, d] where
  type IndexOf '[a, b, c, d] = Ix4
instance HasIndex '[a, b, c, d, e] where
  type IndexOf '[a, b, c, d, e] = Ix5


outerList
  :: forall a as f
   . (Lower (IndexOf (a : as)) ~ IndexOf as, Index (IndexOf as), Index (IndexOf (a : as)))
  => DataCube (a : as) f
  -> [DataCube as f]
outerList :: forall a (as :: [*]) f.
(Lower (IndexOf (a : as)) ~ IndexOf as, Index (IndexOf as),
 Index (IndexOf (a : as))) =>
DataCube (a : as) f -> [DataCube as f]
outerList (DataCube Array D (IndexOf (a : as)) f
a) = (Array D (Lower (IndexOf (a : as))) f -> [DataCube as f])
-> Array D (IndexOf (a : as)) f -> [DataCube as f]
forall ix r e m.
(Index ix, Index (Lower ix), Source r e, Monoid m) =>
(Array r (Lower ix) e -> m) -> Array r ix e -> m
foldOuterSlice Array D (Lower (IndexOf (a : as))) f -> [DataCube as f]
Array D (IndexOf as) f -> [DataCube as f]
row Array D (IndexOf (a : as)) f
a
 where
  row :: Array D (IndexOf as) f -> [DataCube as f]
  row :: Array D (IndexOf as) f -> [DataCube as f]
row Array D (IndexOf as) f
r = [Array D (IndexOf as) f -> DataCube as f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube Array D (IndexOf as) f
r]


transposeMajor
  :: (IndexOf (a : b : xs) ~ IndexOf (b : a : xs), Index (Lower (IndexOf (b : a : xs))), Index (IndexOf (b : a : xs)))
  => DataCube (a : b : xs) f
  -> DataCube (b : a : xs) f
transposeMajor :: forall a b (xs :: [*]) f.
(IndexOf (a : b : xs) ~ IndexOf (b : a : xs),
 Index (Lower (IndexOf (b : a : xs))),
 Index (IndexOf (b : a : xs))) =>
DataCube (a : b : xs) f -> DataCube (b : a : xs) f
transposeMajor (DataCube Array D (IndexOf (a : b : xs)) f
arr) = Array D (IndexOf (b : a : xs)) f -> DataCube (b : a : xs) f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube (Array D (IndexOf (b : a : xs)) f -> DataCube (b : a : xs) f)
-> Array D (IndexOf (b : a : xs)) f -> DataCube (b : a : xs) f
forall a b. (a -> b) -> a -> b
$ Array D (IndexOf (b : a : xs)) f
-> Array D (IndexOf (b : a : xs)) f
forall r ix e.
(Index (Lower ix), Index ix, Source r e) =>
Array r ix e -> Array D ix e
transposeInner Array D (IndexOf (a : b : xs)) f
Array D (IndexOf (b : a : xs)) f
arr


transposeMinor4
  :: DataCube [a, b, c, d] f
  -> DataCube [a, b, d, c] f
transposeMinor4 :: forall a b c d f.
DataCube '[a, b, c, d] f -> DataCube '[a, b, d, c] f
transposeMinor4 (DataCube Array D (IndexOf '[a, b, c, d]) f
arr) = Array D (IndexOf '[a, b, d, c]) f -> DataCube '[a, b, d, c] f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube (Array D (IndexOf '[a, b, d, c]) f -> DataCube '[a, b, d, c] f)
-> Array D (IndexOf '[a, b, d, c]) f -> DataCube '[a, b, d, c] f
forall a b. (a -> b) -> a -> b
$ Array D Ix4 f -> Array D Ix4 f
forall r ix e.
(Index (Lower ix), Index ix, Source r e) =>
Array r ix e -> Array D ix e
transposeOuter Array D Ix4 f
Array D (IndexOf '[a, b, c, d]) f
arr


transposeMinor3
  :: DataCube [a, b, c] f
  -> DataCube [a, c, b] f
transposeMinor3 :: forall a b c f. DataCube '[a, b, c] f -> DataCube '[a, c, b] f
transposeMinor3 (DataCube Array D (IndexOf '[a, b, c]) f
arr) = Array D (IndexOf '[a, c, b]) f -> DataCube '[a, c, b] f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube (Array D (IndexOf '[a, c, b]) f -> DataCube '[a, c, b] f)
-> Array D (IndexOf '[a, c, b]) f -> DataCube '[a, c, b] f
forall a b. (a -> b) -> a -> b
$ Array D (IxN 3) f -> Array D (IxN 3) f
forall r ix e.
(Index (Lower ix), Index ix, Source r e) =>
Array r ix e -> Array D ix e
transposeOuter Array D (IxN 3) f
Array D (IndexOf '[a, b, c]) f
arr


-- Slice along the 1st major dimension
sliceM0
  :: ( Lower (IndexOf (a : xs)) ~ IndexOf xs
     , Index (IndexOf xs)
     , Index (IndexOf (a : xs))
     )
  => Int
  -> DataCube (a : xs) f
  -> DataCube xs f
sliceM0 :: forall a (xs :: [*]) f.
(Lower (IndexOf (a : xs)) ~ IndexOf xs, Index (IndexOf xs),
 Index (IndexOf (a : xs))) =>
Int -> DataCube (a : xs) f -> DataCube xs f
sliceM0 Int
a (DataCube Array D (IndexOf (a : xs)) f
arr) = Array D (IndexOf xs) f -> DataCube xs f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube (Array D (IndexOf (a : xs)) f
arr Array D (IndexOf (a : xs)) f
-> Int -> Array D (Lower (IndexOf (a : xs))) f
forall r ix e.
(HasCallStack, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> Array r (Lower ix) e
!> Int
a)


-- Slice along the 2nd major dimension
sliceM1
  :: forall a b xs f
   . ( Lower (IndexOf (a : b : xs)) ~ IndexOf (a : xs)
     , Index (IndexOf (a : xs))
     , Index (IndexOf (a : b : xs))
     )
  => Int
  -> DataCube (a : b : xs) f
  -> DataCube (a : xs) f
sliceM1 :: forall a b (xs :: [*]) f.
(Lower (IndexOf (a : b : xs)) ~ IndexOf (a : xs),
 Index (IndexOf (a : xs)), Index (IndexOf (a : b : xs))) =>
Int -> DataCube (a : b : xs) f -> DataCube (a : xs) f
sliceM1 Int
b (DataCube Array D (IndexOf (a : b : xs)) f
arr) =
  let dims :: Int
dims = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Integer
natVal @(Dimensions (IndexOf (a : b : xs))) Proxy (Dimensions (IndexOf (a : b : xs)))
forall {k} (t :: k). Proxy t
Proxy
   in Array D (IndexOf (a : xs)) f -> DataCube (a : xs) f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube (Array D (IndexOf (a : xs)) f -> DataCube (a : xs) f)
-> Array D (IndexOf (a : xs)) f -> DataCube (a : xs) f
forall a b. (a -> b) -> a -> b
$ Array D (IndexOf (a : b : xs)) f
arr Array D (IndexOf (a : b : xs)) f
-> (Dim, Int) -> Array D (Lower (IndexOf (a : b : xs))) f
forall r ix e.
(HasCallStack, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> (Dim, Int) -> Array D (Lower ix) e
<!> (Int -> Dim
Dim (Int
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1), Int
b)


-- Slice along the 3rd major dimension
sliceM2
  :: forall a b c xs f
   . ( Lower (IndexOf (a : b : c : xs)) ~ IndexOf (a : b : xs)
     , Index (IndexOf (a : b : xs))
     , Index (IndexOf (a : b : c : xs))
     )
  => Int
  -> DataCube (a : b : c : xs) f
  -> DataCube (a : b : xs) f
sliceM2 :: forall a b c (xs :: [*]) f.
(Lower (IndexOf (a : b : c : xs)) ~ IndexOf (a : b : xs),
 Index (IndexOf (a : b : xs)), Index (IndexOf (a : b : c : xs))) =>
Int -> DataCube (a : b : c : xs) f -> DataCube (a : b : xs) f
sliceM2 Int
c (DataCube Array D (IndexOf (a : b : c : xs)) f
arr) =
  let dims :: Int
dims = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Integer
natVal @(Dimensions (IndexOf (a : b : c : xs))) Proxy (Dimensions (IndexOf (a : b : c : xs)))
forall {k} (t :: k). Proxy t
Proxy
   in Array D (IndexOf (a : b : xs)) f -> DataCube (a : b : xs) f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube (Array D (IndexOf (a : b : xs)) f -> DataCube (a : b : xs) f)
-> Array D (IndexOf (a : b : xs)) f -> DataCube (a : b : xs) f
forall a b. (a -> b) -> a -> b
$ Array D (IndexOf (a : b : c : xs)) f
arr Array D (IndexOf (a : b : c : xs)) f
-> (Dim, Int) -> Array D (Lower (IndexOf (a : b : c : xs))) f
forall r ix e.
(HasCallStack, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> (Dim, Int) -> Array D (Lower ix) e
<!> (Int -> Dim
Dim (Int
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2), Int
c)


splitM0
  :: forall a xs f m
   . ( Index (IndexOf (a : xs))
     , MonadThrow m
     )
  => Int
  -> DataCube (a : xs) f
  -> m (DataCube (a : xs) f, DataCube (a : xs) f)
splitM0 :: forall a (xs :: [*]) f (m :: * -> *).
(Index (IndexOf (a : xs)), MonadThrow m) =>
Int
-> DataCube (a : xs) f
-> m (DataCube (a : xs) f, DataCube (a : xs) f)
splitM0 Int
a (DataCube Array D (IndexOf (a : xs)) f
arr) = do
  let dims :: Int
dims = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Integer
natVal @(Dimensions (IndexOf (a : xs))) Proxy (Dimensions (IndexOf (a : xs)))
forall {k} (t :: k). Proxy t
Proxy
  (Array D (IndexOf (a : xs)) f
arr1, Array D (IndexOf (a : xs)) f
arr2) <- Dim
-> Int
-> Array D (IndexOf (a : xs)) f
-> m (Array D (IndexOf (a : xs)) f, Array D (IndexOf (a : xs)) f)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Dim -> Int -> Array r ix e -> m (Array D ix e, Array D ix e)
M.splitAtM (Int -> Dim
Dim Int
dims) Int
a Array D (IndexOf (a : xs)) f
arr
  (DataCube (a : xs) f, DataCube (a : xs) f)
-> m (DataCube (a : xs) f, DataCube (a : xs) f)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D (IndexOf (a : xs)) f -> DataCube (a : xs) f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube Array D (IndexOf (a : xs)) f
arr1, Array D (IndexOf (a : xs)) f -> DataCube (a : xs) f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube Array D (IndexOf (a : xs)) f
arr2)


splitM1
  :: forall a b xs f m
   . ( Index (IndexOf (a : xs))
     , Index (IndexOf (a : b : xs))
     , MonadThrow m
     )
  => Int
  -> DataCube (a : b : xs) f
  -> m (DataCube (a : b : xs) f, DataCube (a : b : xs) f)
splitM1 :: forall a b (xs :: [*]) f (m :: * -> *).
(Index (IndexOf (a : xs)), Index (IndexOf (a : b : xs)),
 MonadThrow m) =>
Int
-> DataCube (a : b : xs) f
-> m (DataCube (a : b : xs) f, DataCube (a : b : xs) f)
splitM1 Int
b (DataCube Array D (IndexOf (a : b : xs)) f
arr) = do
  let dims :: Int
dims = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Integer
natVal @(Dimensions (IndexOf (a : xs))) Proxy (Dimensions (IndexOf (a : xs)))
forall {k} (t :: k). Proxy t
Proxy
  (Array D (IndexOf (a : b : xs)) f
arr1, Array D (IndexOf (a : b : xs)) f
arr2) <- Dim
-> Int
-> Array D (IndexOf (a : b : xs)) f
-> m (Array D (IndexOf (a : b : xs)) f,
      Array D (IndexOf (a : b : xs)) f)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Dim -> Int -> Array r ix e -> m (Array D ix e, Array D ix e)
M.splitAtM (Int -> Dim
Dim Int
dims) Int
b Array D (IndexOf (a : b : xs)) f
arr
  (DataCube (a : b : xs) f, DataCube (a : b : xs) f)
-> m (DataCube (a : b : xs) f, DataCube (a : b : xs) f)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D (IndexOf (a : b : xs)) f -> DataCube (a : b : xs) f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube Array D (IndexOf (a : b : xs)) f
arr1, Array D (IndexOf (a : b : xs)) f -> DataCube (a : b : xs) f
forall (as :: [*]) f. Array D (IndexOf as) f -> DataCube as f
DataCube Array D (IndexOf (a : b : xs)) f
arr2)


dataCubeAxes :: (Index (IndexOf as), AxesIndex (IndexOf as)) => DataCube as f -> Axes Row
dataCubeAxes :: forall (as :: [*]) f.
(Index (IndexOf as), AxesIndex (IndexOf as)) =>
DataCube as f -> Axes 'Row
dataCubeAxes (DataCube Array D (IndexOf as) f
arr) =
  let Sz IndexOf as
ix = Array D (IndexOf as) f -> Sz (IndexOf as)
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
M.size Array D (IndexOf as) f
arr
   in IndexOf as -> Axes 'Row
forall ix. AxesIndex ix => ix -> Axes 'Row
indexAxes IndexOf as
ix