{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeAbstractions #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-|
TODO:
* We should be more consistent in whether functions take a 'StaticShX'
  argument or a 'KnownShapeX' constraint.

* Document the choice of using 'INat' for ranks and 'Nat' for shapes. Point
  being that we need to do induction over the former, but the latter need to be
  able to get large.

-}

module Data.Array.Nested.Internal where

import Prelude hiding (mappend)

import Control.Monad (forM_, when)
import Control.Monad.ST
import qualified Data.Array.RankedS as S
import Data.Bifunctor (first)
import Data.Coerce (coerce, Coercible)
import Data.Foldable (toList)
import Data.Kind
import Data.List.NonEmpty (NonEmpty)
import Data.Proxy
import Data.Type.Equality
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import Foreign.Storable (Storable)
import GHC.TypeLits

import Data.Array.Mixed (XArray, IxX(..), IIxX, ShX(..), IShX, KnownShapeX(..), StaticShX(..), type (++), pattern GHC_SNat)
import qualified Data.Array.Mixed as X
import Data.INat


-- Invariant in the API
-- ====================
--
-- In the underlying XArray, there is some shape for elements of an empty
-- array. For example, for this array:
--
--   arr :: Ranked I3 (Ranked I2 Int, Ranked I1 Float)
--   rshape arr == 0 :.: 0 :.: 0 :.: ZIR
--
-- the two underlying XArrays have a shape, and those shapes might be anything.
-- The invariant is that these element shapes are unobservable in the API.
-- (This is possible because you ought to not be able to get to such an element
-- without indexing out of bounds.)
--
-- Note, though, that the converse situation may arise: the outer array might
-- be nonempty but then the inner arrays might. This is fine, an invariant only
-- applies if the _outer_ array is empty.
--
-- TODO: can we enforce that the elements of an empty (nested) array have
-- all-zero shape?
--   -> no, because mlift and also any kind of internals probing from outsiders


-- Primitive element types
-- =======================
--
-- There are a few primitive element types; arrays containing elements of such
-- type are a newtype over an XArray, which it itself a newtype over a Vector.
-- Unfortunately, the setup of the library requires us to list these primitive
-- element types multiple times; to aid in extending the list, all these lists
-- have been marked with [PRIMITIVE ELEMENT TYPES LIST].


type family Replicate n a where
  Replicate Z a = '[]
  Replicate (S n) a = a : Replicate n a

type family MapJust l where
  MapJust '[] = '[]
  MapJust (x : xs) = Just x : MapJust xs

lemKnownReplicate :: forall n. KnownINat n => Proxy n -> Dict KnownShapeX (Replicate n Nothing)
lemKnownReplicate _ = X.lemKnownShapeX (go (inatSing @n))
  where
    go :: SINat m -> StaticShX (Replicate m Nothing)
    go SZ = ZKSX
    go (SS n) = () :!$? go n

lemRankReplicate :: forall n. KnownINat n => Proxy n -> X.Rank (Replicate n (Nothing @Nat)) :~: n
lemRankReplicate _ = go (inatSing @n)
  where
    go :: SINat m -> X.Rank (Replicate m (Nothing @Nat)) :~: m
    go SZ = Refl
    go (SS n) | Refl <- go n = Refl

lemReplicatePlusApp :: forall n m a. KnownINat n => Proxy n -> Proxy m -> Proxy a
                    -> Replicate (n +! m) a :~: Replicate n a ++ Replicate m a
lemReplicatePlusApp _ _ _ = go (inatSing @n)
  where
    go :: SINat n' -> Replicate (n' +! m) a :~: Replicate n' a ++ Replicate m a
    go SZ = Refl
    go (SS n) | Refl <- go n = Refl

shAppSplit :: Proxy sh' -> StaticShX sh -> IShX (sh ++ sh') -> (IShX sh, IShX sh')
shAppSplit _ ZKSX idx = (ZSX, idx)
shAppSplit p (_ :!$@ ssh) (i :$@ idx) = first (i :$@) (shAppSplit p ssh idx)
shAppSplit p (_ :!$? ssh) (i :$? idx) = first (i :$?) (shAppSplit p ssh idx)


-- | Wrapper type used as a tag to attach instances on. The instances on arrays
-- of @'Primitive' a@ are more polymorphic than the direct instances for arrays
-- of scalars; this means that if @orthotope@ supports an element type @T@ that
-- this library does not (directly), it may just work if you use an array of
-- @'Primitive' T@ instead.
newtype Primitive a = Primitive a

-- | Element types that are primitive; arrays of these types are just a newtype
-- wrapper over an array.
class PrimElt a where
  fromPrimitive :: Mixed sh (Primitive a) -> Mixed sh a
  toPrimitive :: Mixed sh a -> Mixed sh (Primitive a)

  default fromPrimitive :: Coercible (Mixed sh a) (Mixed sh (Primitive a)) => Mixed sh (Primitive a) -> Mixed sh a
  fromPrimitive = coerce

  default toPrimitive :: Coercible (Mixed sh (Primitive a)) (Mixed sh a) => Mixed sh a -> Mixed sh (Primitive a)
  toPrimitive = coerce

-- [PRIMITIVE ELEMENT TYPES LIST]
instance PrimElt Int
instance PrimElt Double
instance PrimElt ()


-- | Mixed arrays: some dimensions are size-typed, some are not. Distributes
-- over product-typed elements using a data family so that the full array is
-- always in struct-of-arrays format.
--
-- Built on top of 'XArray' which is built on top of @orthotope@, meaning that
-- dimension permutations (e.g. 'mtranspose') are typically free.
--
-- Many of the methods for working on 'Mixed' arrays come from the 'Elt' type
-- class.
type Mixed :: [Maybe Nat] -> Type -> Type
data family Mixed sh a
-- NOTE: When opening up the Mixed abstraction, you might see dimension sizes
-- that you're not supposed to see. In particular, you might see (nonempty)
-- sizes of the elements of an empty array, which is information that should
-- ostensibly not exist; the full array is still empty.

newtype instance Mixed sh (Primitive a) = M_Primitive (XArray sh a)
  deriving (Show)

-- [PRIMITIVE ELEMENT TYPES LIST]
newtype instance Mixed sh Int = M_Int (XArray sh Int)
  deriving (Show)
newtype instance Mixed sh Double = M_Double (XArray sh Double)
  deriving (Show)
newtype instance Mixed sh () = M_Nil (XArray sh ())  -- no content, orthotope optimises this (via Vector)
  deriving (Show)
-- etc.

data instance Mixed sh (a, b) = M_Tup2 !(Mixed sh a) !(Mixed sh b)
deriving instance (Show (Mixed sh a), Show (Mixed sh b)) => Show (Mixed sh (a, b))
-- etc.

newtype instance Mixed sh1 (Mixed sh2 a) = M_Nest (Mixed (sh1 ++ sh2) a)
deriving instance Show (Mixed (sh1 ++ sh2) a) => Show (Mixed sh1 (Mixed sh2 a))


-- | Internal helper data family mirroring 'Mixed' that consists of mutable
-- vectors instead of 'XArray's.
type MixedVecs :: Type -> [Maybe Nat] -> Type -> Type
data family MixedVecs s sh a

newtype instance MixedVecs s sh (Primitive a) = MV_Primitive (VS.MVector s a)

-- [PRIMITIVE ELEMENT TYPES LIST]
newtype instance MixedVecs s sh Int = MV_Int (VS.MVector s Int)
newtype instance MixedVecs s sh Double = MV_Double (VS.MVector s Double)
newtype instance MixedVecs s sh () = MV_Nil (VS.MVector s ())  -- no content, MVector optimises this
-- etc.

data instance MixedVecs s sh (a, b) = MV_Tup2 !(MixedVecs s sh a) !(MixedVecs s sh b)
-- etc.

data instance MixedVecs s sh1 (Mixed sh2 a) = MV_Nest !(IShX sh2) !(MixedVecs s (sh1 ++ sh2) a)


-- | Tree giving the shape of every array component.
type family ShapeTree a where
  ShapeTree (Primitive _) = ()
  -- [PRIMITIVE ELEMENT TYPES LIST]
  ShapeTree Int = ()
  ShapeTree Double = ()
  ShapeTree () = ()

  ShapeTree (a, b) = (ShapeTree a, ShapeTree b)
  ShapeTree (Mixed sh a) = (IShX sh, ShapeTree a)
  ShapeTree (Ranked n a) = (IShR n, ShapeTree a)
  ShapeTree (Shaped sh a) = (ShS sh, ShapeTree a)


-- | Allowable scalar types in a mixed array, and by extension in a 'Ranked' or
-- 'Shaped' array. Note the polymorphic instance for 'Elt' of @'Primitive'
-- a@; see the documentation for 'Primitive' for more details.
class Elt a where
  -- ====== PUBLIC METHODS ====== --

  mshape :: KnownShapeX sh => Mixed sh a -> IShX sh
  mindex :: Mixed sh a -> IIxX sh -> a
  mindexPartial :: forall sh sh'. Mixed (sh ++ sh') a -> IIxX sh -> Mixed sh' a
  mscalar :: a -> Mixed '[] a

  -- | All arrays in the list, even subarrays inside @a@, must have the same
  -- shape; if they do not, a runtime error will be thrown. See the
  -- documentation of 'mgenerate' for more information about this restriction.
  -- Furthermore, the length of the list must correspond with @n@: if @n@ is
  -- @Just m@ and @m@ does not equal the length of the list, a runtime error is
  -- thrown.
  --
  -- If you want a single-dimensional array from your list, map 'mscalar'
  -- first.
  mfromList1 :: forall n sh. KnownShapeX (n : sh) => NonEmpty (Mixed sh a) -> Mixed (n : sh) a

  mtoList1 :: Mixed (n : sh) a -> [Mixed sh a]

  -- | Note: this library makes no particular guarantees about the shapes of
  -- arrays "inside" an empty array. With 'mlift' and 'mlift2' you can see the
  -- full 'XArray' and as such you can distinguish different empty arrays by
  -- the "shapes" of their elements. This information is meaningless, so you
  -- should not use it.
  mlift :: forall sh1 sh2. KnownShapeX sh2
        => (forall sh' b. (KnownShapeX sh', Storable b) => Proxy sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
        -> Mixed sh1 a -> Mixed sh2 a

  -- | See the documentation for 'mlift'.
  mlift2 :: forall sh1 sh2 sh3. (KnownShapeX sh2, KnownShapeX sh3)
         => (forall sh' b. (KnownShapeX sh', Storable b) => Proxy sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b -> XArray (sh3 ++ sh') b)
         -> Mixed sh1 a -> Mixed sh2 a -> Mixed sh3 a

  -- ====== PRIVATE METHODS ====== --

  -- | Create an empty array. The given shape must have size zero; this may or may not be checked.
  memptyArray :: IShX sh -> Mixed sh a

  mshapeTree :: a -> ShapeTree a

  mshapeTreeEq :: Proxy a -> ShapeTree a -> ShapeTree a -> Bool

  mshapeTreeEmpty :: Proxy a -> ShapeTree a -> Bool

  mshowShapeTree :: Proxy a -> ShapeTree a -> String

  -- | Create uninitialised vectors for this array type, given the shape of
  -- this vector and an example for the contents.
  mvecsUnsafeNew :: IShX sh -> a -> ST s (MixedVecs s sh a)

  mvecsNewEmpty :: Proxy a -> ST s (MixedVecs s sh a)

  -- | Given the shape of this array, an index and a value, write the value at
  -- that index in the vectors.
  mvecsWrite :: IShX sh -> IIxX sh -> a -> MixedVecs s sh a -> ST s ()

  -- | Given the shape of this array, an index and a value, write the value at
  -- that index in the vectors.
  mvecsWritePartial :: KnownShapeX sh' => IShX (sh ++ sh') -> IIxX sh -> Mixed sh' a -> MixedVecs s (sh ++ sh') a -> ST s ()

  -- | Given the shape of this array, finalise the vectors into 'XArray's.
  mvecsFreeze :: IShX sh -> MixedVecs s sh a -> ST s (Mixed sh a)


-- Arrays of scalars are basically just arrays of scalars.
instance Storable a => Elt (Primitive a) where
  mshape (M_Primitive a) = X.shape a
  mindex (M_Primitive a) i = Primitive (X.index a i)
  mindexPartial (M_Primitive a) i = M_Primitive (X.indexPartial a i)
  mscalar (Primitive x) = M_Primitive (X.scalar x)
  mfromList1 l = M_Primitive (X.fromList1 knownShapeX (coerce (toList l)))
  mtoList1 (M_Primitive arr) = coerce (X.toList1 arr)

  mlift :: forall sh1 sh2.
           (Proxy '[] -> XArray (sh1 ++ '[]) a -> XArray (sh2 ++ '[]) a)
        -> Mixed sh1 (Primitive a) -> Mixed sh2 (Primitive a)
  mlift f (M_Primitive a)
    | Refl <- X.lemAppNil @sh1
    , Refl <- X.lemAppNil @sh2
    = M_Primitive (f Proxy a)

  mlift2 :: forall sh1 sh2 sh3.
            (Proxy '[] -> XArray (sh1 ++ '[]) a -> XArray (sh2 ++ '[]) a -> XArray (sh3 ++ '[]) a)
         -> Mixed sh1 (Primitive a) -> Mixed sh2 (Primitive a) -> Mixed sh3 (Primitive a)
  mlift2 f (M_Primitive a) (M_Primitive b)
    | Refl <- X.lemAppNil @sh1
    , Refl <- X.lemAppNil @sh2
    , Refl <- X.lemAppNil @sh3
    = M_Primitive (f Proxy a b)

  memptyArray sh = M_Primitive (X.empty sh)
  mshapeTree _ = ()
  mshapeTreeEq _ () () = True
  mshapeTreeEmpty _ () = False
  mshowShapeTree _ () = "()"
  mvecsUnsafeNew sh _ = MV_Primitive <$> VSM.unsafeNew (X.shapeSize sh)
  mvecsNewEmpty _ = MV_Primitive <$> VSM.unsafeNew 0
  mvecsWrite sh i (Primitive x) (MV_Primitive v) = VSM.write v (X.toLinearIdx sh i) x

  -- TODO: this use of toVector is suboptimal
  mvecsWritePartial
    :: forall sh' sh s. KnownShapeX sh'
    => IShX (sh ++ sh') -> IIxX sh -> Mixed sh' (Primitive a) -> MixedVecs s (sh ++ sh') (Primitive a) -> ST s ()
  mvecsWritePartial sh i (M_Primitive arr) (MV_Primitive v) = do
    let offset = X.toLinearIdx sh (X.ixAppend i (X.zeroIxX' (X.shape arr)))
    VS.copy (VSM.slice offset (X.shapeSize (X.shape arr)) v) (X.toVector arr)

  mvecsFreeze sh (MV_Primitive v) = M_Primitive . X.fromVector sh <$> VS.freeze v

-- [PRIMITIVE ELEMENT TYPES LIST]
deriving via Primitive Int instance Elt Int
deriving via Primitive Double instance Elt Double
deriving via Primitive () instance Elt ()

-- Arrays of pairs are pairs of arrays.
instance (Elt a, Elt b) => Elt (a, b) where
  mshape (M_Tup2 a _) = mshape a
  mindex (M_Tup2 a b) i = (mindex a i, mindex b i)
  mindexPartial (M_Tup2 a b) i = M_Tup2 (mindexPartial a i) (mindexPartial b i)
  mscalar (x, y) = M_Tup2 (mscalar x) (mscalar y)
  mfromList1 l = M_Tup2 (mfromList1 ((\(M_Tup2 x _) -> x) <$> l))
                        (mfromList1 ((\(M_Tup2 _ y) -> y) <$> l))
  mtoList1 (M_Tup2 a b) = zipWith M_Tup2 (mtoList1 a) (mtoList1 b)
  mlift f (M_Tup2 a b) = M_Tup2 (mlift f a) (mlift f b)
  mlift2 f (M_Tup2 a b) (M_Tup2 x y) = M_Tup2 (mlift2 f a x) (mlift2 f b y)

  memptyArray sh = M_Tup2 (memptyArray sh) (memptyArray sh)
  mshapeTree (x, y) = (mshapeTree x, mshapeTree y)
  mshapeTreeEq _ (t1, t2) (t1', t2') = mshapeTreeEq (Proxy @a) t1 t1' && mshapeTreeEq (Proxy @b) t2 t2'
  mshapeTreeEmpty _ (t1, t2) = mshapeTreeEmpty (Proxy @a) t1 && mshapeTreeEmpty (Proxy @b) t2
  mshowShapeTree _ (t1, t2) = "(" ++ mshowShapeTree (Proxy @a) t1 ++ ", " ++ mshowShapeTree (Proxy @b) t2 ++ ")"
  mvecsUnsafeNew sh (x, y) = MV_Tup2 <$> mvecsUnsafeNew sh x <*> mvecsUnsafeNew sh y
  mvecsNewEmpty _ = MV_Tup2 <$> mvecsNewEmpty (Proxy @a) <*> mvecsNewEmpty (Proxy @b)
  mvecsWrite sh i (x, y) (MV_Tup2 a b) = do
    mvecsWrite sh i x a
    mvecsWrite sh i y b
  mvecsWritePartial sh i (M_Tup2 x y) (MV_Tup2 a b) = do
    mvecsWritePartial sh i x a
    mvecsWritePartial sh i y b
  mvecsFreeze sh (MV_Tup2 a b) = M_Tup2 <$> mvecsFreeze sh a <*> mvecsFreeze sh b

-- Arrays of arrays are just arrays, but with more dimensions.
instance (Elt a, KnownShapeX sh') => Elt (Mixed sh' a) where
  -- TODO: this is quadratic in the nesting depth because it repeatedly
  -- truncates the shape vector to one a little shorter. Fix with a
  -- moverlongShape method, a prefix of which is mshape.
  mshape :: forall sh. KnownShapeX sh => Mixed sh (Mixed sh' a) -> IShX sh
  mshape (M_Nest arr)
    | Dict <- X.lemAppKnownShapeX (knownShapeX @sh) (knownShapeX @sh')
    = fst (shAppSplit (Proxy @sh') (knownShapeX @sh) (mshape arr))

  mindex (M_Nest arr) i = mindexPartial arr i

  mindexPartial :: forall sh1 sh2.
                   Mixed (sh1 ++ sh2) (Mixed sh' a) -> IIxX sh1 -> Mixed sh2 (Mixed sh' a)
  mindexPartial (M_Nest arr) i
    | Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @sh2) (Proxy @sh')
    = M_Nest (mindexPartial @a @sh1 @(sh2 ++ sh') arr i)

  mscalar = M_Nest

  mfromList1 :: forall n sh. KnownShapeX (n : sh)
             => NonEmpty (Mixed sh (Mixed sh' a)) -> Mixed (n : sh) (Mixed sh' a)
  mfromList1 l
    | Dict <- X.lemKnownShapeX (X.ssxAppend (knownShapeX @(n : sh)) (knownShapeX @sh'))
    = M_Nest (mfromList1 (coerce l))

  mtoList1 (M_Nest arr) = coerce (mtoList1 arr)

  mlift :: forall sh1 sh2. KnownShapeX sh2
        => (forall shT b. (KnownShapeX shT, Storable b) => Proxy shT -> XArray (sh1 ++ shT) b -> XArray (sh2 ++ shT) b)
        -> Mixed sh1 (Mixed sh' a) -> Mixed sh2 (Mixed sh' a)
  mlift f (M_Nest arr)
    | Dict <- X.lemKnownShapeX (X.ssxAppend (knownShapeX @sh2) (knownShapeX @sh'))
    = M_Nest (mlift f' arr)
    where
      f' :: forall shT b. (KnownShapeX shT, Storable b) => Proxy shT -> XArray ((sh1 ++ sh') ++ shT) b -> XArray ((sh2 ++ sh') ++ shT) b
      f' _
        | Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @sh') (Proxy @shT)
        , Refl <- X.lemAppAssoc (Proxy @sh2) (Proxy @sh') (Proxy @shT)
        , Dict <- X.lemKnownShapeX (X.ssxAppend (knownShapeX @sh') (knownShapeX @shT))
        = f (Proxy @(sh' ++ shT))

  mlift2 :: forall sh1 sh2 sh3. (KnownShapeX sh2, KnownShapeX sh3)
         => (forall shT b. (KnownShapeX shT, Storable b) => Proxy shT -> XArray (sh1 ++ shT) b -> XArray (sh2 ++ shT) b -> XArray (sh3 ++ shT) b)
         -> Mixed sh1 (Mixed sh' a) -> Mixed sh2 (Mixed sh' a) -> Mixed sh3 (Mixed sh' a)
  mlift2 f (M_Nest arr1) (M_Nest arr2)
    | Dict <- X.lemKnownShapeX (X.ssxAppend (knownShapeX @sh2) (knownShapeX @sh'))
    , Dict <- X.lemKnownShapeX (X.ssxAppend (knownShapeX @sh3) (knownShapeX @sh'))
    = M_Nest (mlift2 f' arr1 arr2)
    where
      f' :: forall shT b. (KnownShapeX shT, Storable b) => Proxy shT -> XArray ((sh1 ++ sh') ++ shT) b -> XArray ((sh2 ++ sh') ++ shT) b -> XArray ((sh3 ++ sh') ++ shT) b
      f' _
        | Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @sh') (Proxy @shT)
        , Refl <- X.lemAppAssoc (Proxy @sh2) (Proxy @sh') (Proxy @shT)
        , Refl <- X.lemAppAssoc (Proxy @sh3) (Proxy @sh') (Proxy @shT)
        , Dict <- X.lemKnownShapeX (X.ssxAppend (knownShapeX @sh') (knownShapeX @shT))
        = f (Proxy @(sh' ++ shT))

  memptyArray sh = M_Nest (memptyArray (X.shAppend sh (X.completeShXzeros (knownShapeX @sh'))))

  mshapeTree arr = (mshape arr, mshapeTree (mindex arr (X.zeroIxX (knownShapeX @sh'))))

  mshapeTreeEq _ (sh1, t1) (sh2, t2) = sh1 == sh2 && mshapeTreeEq (Proxy @a) t1 t2

  mshapeTreeEmpty _ (sh, t) = X.shapeSize sh == 0 && mshapeTreeEmpty (Proxy @a) t

  mshowShapeTree _ (sh, t) = "(" ++ show sh ++ ", " ++ mshowShapeTree (Proxy @a) t ++ ")"

  mvecsUnsafeNew sh example
    | X.shapeSize sh' == 0 = mvecsNewEmpty (Proxy @(Mixed sh' a))
    | otherwise = MV_Nest sh' <$> mvecsUnsafeNew (X.shAppend sh (mshape example))
                                                 (mindex example (X.zeroIxX (knownShapeX @sh')))
    where
      sh' = mshape example

  mvecsNewEmpty _ = MV_Nest (X.completeShXzeros (knownShapeX @sh')) <$> mvecsNewEmpty (Proxy @a)

  mvecsWrite sh idx val (MV_Nest sh' vecs) = mvecsWritePartial (X.shAppend sh sh') idx val vecs

  mvecsWritePartial :: forall sh1 sh2 s. KnownShapeX sh2
                    => IShX (sh1 ++ sh2) -> IIxX sh1 -> Mixed sh2 (Mixed sh' a)
                    -> MixedVecs s (sh1 ++ sh2) (Mixed sh' a)
                    -> ST s ()
  mvecsWritePartial sh12 idx (M_Nest arr) (MV_Nest sh' vecs)
    | Dict <- X.lemKnownShapeX (X.ssxAppend (knownShapeX @sh2) (knownShapeX @sh'))
    , Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @sh2) (Proxy @sh')
    = mvecsWritePartial @a @(sh2 ++ sh') @sh1 (X.shAppend sh12 sh') idx arr vecs

  mvecsFreeze sh (MV_Nest sh' vecs) = M_Nest <$> mvecsFreeze (X.shAppend sh sh') vecs


-- | Create an array given a size and a function that computes the element at a
-- given index.
--
-- __WARNING__: It is required that every @a@ returned by the argument to
-- 'mgenerate' has the same shape. For example, the following will throw a
-- runtime error:
--
-- > foo :: Mixed [Nothing] (Mixed [Nothing] Double)
-- > foo = mgenerate (10 :.: ZIR) $ \(i :.: ZIR) ->
-- >         mgenerate (i :.: ZIR) $ \(j :.: ZIR) ->
-- >           ...
--
-- because the size of the inner 'mgenerate' is not always the same (it depends
-- on @i@). Nested arrays in @ox-arrays@ are always stored fully flattened, so
-- the entire hierarchy (after distributing out tuples) must be a rectangular
-- array. The type of 'mgenerate' allows this requirement to be broken very
-- easily, hence the runtime check.
mgenerate :: forall sh a. (KnownShapeX sh, Elt a) => IShX sh -> (IIxX sh -> a) -> Mixed sh a
mgenerate sh f = case X.enumShape sh of
  [] -> memptyArray sh
  firstidx : restidxs ->
    let firstelem = f (X.zeroIxX' sh)
        shapetree = mshapeTree firstelem
    in if mshapeTreeEmpty (Proxy @a) shapetree
         then memptyArray sh
         else runST $ do
                vecs <- mvecsUnsafeNew sh firstelem
                mvecsWrite sh firstidx firstelem vecs
                -- TODO: This is likely fine if @a@ is big, but if @a@ is a
                -- scalar this array copying inefficient. Should improve this.
                forM_ restidxs $ \idx -> do
                  let val = f idx
                  when (not (mshapeTreeEq (Proxy @a) (mshapeTree val) shapetree)) $
                    error "Data.Array.Nested mgenerate: generated values do not have equal shapes"
                  mvecsWrite sh idx val vecs
                mvecsFreeze sh vecs

mtranspose :: forall sh a. (KnownShapeX sh, Elt a) => [Int] -> Mixed sh a -> Mixed sh a
mtranspose perm =
  mlift (\(Proxy @sh') -> X.rerankTop (knownShapeX @sh) (knownShapeX @sh) (knownShapeX @sh')
                            (X.transpose perm))

mappend :: forall n m sh a. (KnownShapeX sh, KnownShapeX (n : sh), KnownShapeX (m : sh), KnownShapeX (X.AddMaybe n m : sh), Elt a)
        => Mixed (n : sh) a -> Mixed (m : sh) a -> Mixed (X.AddMaybe n m : sh) a
mappend = mlift2 go
  where go :: forall sh' b. (KnownShapeX sh', Storable b)
           => Proxy sh' -> XArray (n : sh ++ sh') b -> XArray (m : sh ++ sh') b -> XArray (X.AddMaybe n m : sh ++ sh') b
        go Proxy | Dict <- X.lemAppKnownShapeX (knownShapeX @sh) (knownShapeX @sh') = X.append

mfromVectorP :: forall sh a. (KnownShapeX sh, Storable a) => IShX sh -> VS.Vector a -> Mixed sh (Primitive a)
mfromVectorP sh v = M_Primitive (X.fromVector sh v)

mfromVector :: forall sh a. (KnownShapeX sh, Storable a, PrimElt a) => IShX sh -> VS.Vector a -> Mixed sh a
mfromVector sh v = fromPrimitive (mfromVectorP sh v)

mtoVectorP :: Storable a => Mixed sh (Primitive a) -> VS.Vector a
mtoVectorP (M_Primitive v) = X.toVector v

mtoVector :: (Storable a, PrimElt a) => Mixed sh a -> VS.Vector a
mtoVector arr = mtoVectorP (coerce toPrimitive arr)

mfromList :: (KnownShapeX '[n], Elt a) => NonEmpty a -> Mixed '[n] a
mfromList = mfromList1 . fmap mscalar

mtoList :: Elt a => Mixed '[n] a -> [a]
mtoList = map munScalar . mtoList1

munScalar :: Elt a => Mixed '[] a -> a
munScalar arr = mindex arr ZIX

mconstantP :: forall sh a. (KnownShapeX sh, Storable a) => IShX sh -> a -> Mixed sh (Primitive a)
mconstantP sh x = M_Primitive (X.constant sh x)

mconstant :: forall sh a. (KnownShapeX sh, Storable a, PrimElt a)
          => IShX sh -> a -> Mixed sh a
mconstant sh x = fromPrimitive (mconstantP sh x)

mslice :: (KnownShapeX sh, Elt a) => [(Int, Int)] -> Mixed sh a -> Mixed sh a
mslice ivs = mlift $ \_ -> X.slice ivs

mrev1 :: (KnownShapeX (n : sh), Elt a) => Mixed (n : sh) a -> Mixed (n : sh) a
mrev1 = mlift $ \_ -> X.rev1

mreshape :: forall sh sh' a. (KnownShapeX sh, KnownShapeX sh', Elt a)
         => IShX sh' -> Mixed sh a -> Mixed sh' a
mreshape sh' = mlift $ \(_ :: Proxy shIn) ->
                        X.reshapePartial (knownShapeX @sh) (knownShapeX @shIn) sh'

mliftPrim :: (KnownShapeX sh, Storable a)
          => (a -> a)
          -> Mixed sh (Primitive a) -> Mixed sh (Primitive a)
mliftPrim f (M_Primitive (X.XArray arr)) = M_Primitive (X.XArray (S.mapA f arr))

mliftPrim2 :: (KnownShapeX sh, Storable a)
           => (a -> a -> a)
           -> Mixed sh (Primitive a) -> Mixed sh (Primitive a) -> Mixed sh (Primitive a)
mliftPrim2 f (M_Primitive (X.XArray arr1)) (M_Primitive (X.XArray arr2)) =
  M_Primitive (X.XArray (S.zipWithA f arr1 arr2))

instance (KnownShapeX sh, Storable a, Num a) => Num (Mixed sh (Primitive a)) where
  (+) = mliftPrim2 (+)
  (-) = mliftPrim2 (-)
  (*) = mliftPrim2 (*)
  negate = mliftPrim negate
  abs = mliftPrim abs
  signum = mliftPrim signum
  fromInteger n =
    case X.ssxToShape' (knownShapeX @sh) of
      Just sh -> M_Primitive (X.constant sh (fromInteger n))
      Nothing -> error "Data.Array.Nested.fromIntegral: \
                       \Unknown components in shape, use explicit mconstant"

-- [PRIMITIVE ELEMENT TYPES LIST] (really, a partial list of just the numeric types)
deriving via Mixed sh (Primitive Int) instance KnownShapeX sh => Num (Mixed sh Int)
deriving via Mixed sh (Primitive Double) instance KnownShapeX sh => Num (Mixed sh Double)


-- | A rank-typed array: the number of dimensions of the array (its /rank/) is
-- represented on the type level as a 'INat'.
--
-- Valid elements of a ranked arrays are described by the 'Elt' type class.
-- Because 'Ranked' itself is also an instance of 'Elt', nested arrays are
-- supported (and are represented as a single, flattened, struct-of-arrays
-- array internally).
--
-- Note that this 'INat' is not a "GHC.TypeLits" natural, because we want a
-- type-level natural that supports induction.
--
-- 'Ranked' is a newtype around a 'Mixed' of 'Nothing's.
type Ranked :: INat -> Type -> Type
newtype Ranked n a = Ranked (Mixed (Replicate n Nothing) a)
deriving instance Show (Mixed (Replicate n Nothing) a) => Show (Ranked n a)

-- | A shape-typed array: the full shape of the array (the sizes of its
-- dimensions) is represented on the type level as a list of 'Nat's. Note that
-- these are "GHC.TypeLits" naturals, because we do not need induction over
-- them and we want very large arrays to be possible.
--
-- Like for 'Ranked', the valid elements are described by the 'Elt' type class,
-- and 'Shaped' itself is again an instance of 'Elt' as well.
--
-- 'Shaped' is a newtype around a 'Mixed' of 'Just's.
type Shaped :: [Nat] -> Type -> Type
newtype Shaped sh a = Shaped (Mixed (MapJust sh) a)
deriving instance Show (Mixed (MapJust sh) a) => Show (Shaped sh a)

-- just unwrap the newtype and defer to the general instance for nested arrays
newtype instance Mixed sh (Ranked n   a) = M_Ranked (Mixed sh (Mixed (Replicate n Nothing) a))
deriving instance Show (Mixed sh (Mixed (Replicate n Nothing) a)) => Show (Mixed sh (Ranked n a))
newtype instance Mixed sh (Shaped sh' a) = M_Shaped (Mixed sh (Mixed (MapJust sh'        ) a))
deriving instance Show (Mixed sh (Mixed (MapJust sh'        ) a)) => Show (Mixed sh (Shaped sh' a))

newtype instance MixedVecs s sh (Ranked n   a) = MV_Ranked (MixedVecs s sh (Mixed (Replicate n Nothing) a))
newtype instance MixedVecs s sh (Shaped sh' a) = MV_Shaped (MixedVecs s sh (Mixed (MapJust sh'        ) a))


-- 'Ranked' and 'Shaped' can already be used at the top level of an array nest;
-- these instances allow them to also be used as elements of arrays, thus
-- making them first-class in the API.
instance (Elt a, KnownINat n) => Elt (Ranked n a) where
  mshape (M_Ranked arr) | Dict <- lemKnownReplicate (Proxy @n) = mshape arr
  mindex (M_Ranked arr) i | Dict <- lemKnownReplicate (Proxy @n) = Ranked (mindex arr i)

  mindexPartial :: forall sh sh'. Mixed (sh ++ sh') (Ranked n a) -> IIxX sh -> Mixed sh' (Ranked n a)
  mindexPartial (M_Ranked arr) i
    | Dict <- lemKnownReplicate (Proxy @n)
    = coerce @(Mixed sh' (Mixed (Replicate n Nothing) a)) @(Mixed sh' (Ranked n a)) $
        mindexPartial arr i

  mscalar (Ranked x) = M_Ranked (M_Nest x)

  mfromList1 :: forall m sh. KnownShapeX (m : sh)
             => NonEmpty (Mixed sh (Ranked n a)) -> Mixed (m : sh) (Ranked n a)
  mfromList1 l
    | Dict <- lemKnownReplicate (Proxy @n)
    = M_Ranked (mfromList1 (coerce l))

  mtoList1 :: forall m sh. Mixed (m : sh) (Ranked n a) -> [Mixed sh (Ranked n a)]
  mtoList1 (M_Ranked arr)
    | Dict <- lemKnownReplicate (Proxy @n)
    = coerce @[Mixed sh (Mixed (Replicate n 'Nothing) a)] @[Mixed sh (Ranked n a)] (mtoList1 arr)

  mlift :: forall sh1 sh2. KnownShapeX sh2
        => (forall sh' b. (KnownShapeX sh', Storable b) => Proxy sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
        -> Mixed sh1 (Ranked n a) -> Mixed sh2 (Ranked n a)
  mlift f (M_Ranked arr)
    | Dict <- lemKnownReplicate (Proxy @n)
    = coerce @(Mixed sh2 (Mixed (Replicate n Nothing) a)) @(Mixed sh2 (Ranked n a)) $
        mlift f arr

  mlift2 :: forall sh1 sh2 sh3. (KnownShapeX sh2, KnownShapeX sh3)
         => (forall sh' b. (KnownShapeX sh', Storable b) => Proxy sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b -> XArray (sh3 ++ sh') b)
         -> Mixed sh1 (Ranked n a) -> Mixed sh2 (Ranked n a) -> Mixed sh3 (Ranked n a)
  mlift2 f (M_Ranked arr1) (M_Ranked arr2)
    | Dict <- lemKnownReplicate (Proxy @n)
    = coerce @(Mixed sh3 (Mixed (Replicate n Nothing) a)) @(Mixed sh3 (Ranked n a)) $
        mlift2 f arr1 arr2

  memptyArray :: forall sh. IShX sh -> Mixed sh (Ranked n a)
  memptyArray i
    | Dict <- lemKnownReplicate (Proxy @n)
    = coerce @(Mixed sh (Mixed (Replicate n Nothing) a)) @(Mixed sh (Ranked n a)) $
        memptyArray i

  mshapeTree (Ranked arr)
    | Refl <- lemRankReplicate (Proxy @n)
    , Dict <- lemKnownReplicate (Proxy @n)
    = first shCvtXR (mshapeTree arr)

  mshapeTreeEq _ (sh1, t1) (sh2, t2) = sh1 == sh2 && mshapeTreeEq (Proxy @a) t1 t2

  mshapeTreeEmpty _ (sh, t) = shapeSizeR sh == 0 && mshapeTreeEmpty (Proxy @a) t

  mshowShapeTree _ (sh, t) = "(" ++ show sh ++ ", " ++ mshowShapeTree (Proxy @a) t ++ ")"

  mvecsUnsafeNew idx (Ranked arr)
    | Dict <- lemKnownReplicate (Proxy @n)
    = MV_Ranked <$> mvecsUnsafeNew idx arr

  mvecsNewEmpty _
    | Dict <- lemKnownReplicate (Proxy @n)
    = MV_Ranked <$> mvecsNewEmpty (Proxy @(Mixed (Replicate n Nothing) a))

  mvecsWrite :: forall sh s. IShX sh -> IIxX sh -> Ranked n a -> MixedVecs s sh (Ranked n a) -> ST s ()
  mvecsWrite sh idx (Ranked arr) vecs
    | Dict <- lemKnownReplicate (Proxy @n)
    = mvecsWrite sh idx arr
        (coerce @(MixedVecs s sh (Ranked n a)) @(MixedVecs s sh (Mixed (Replicate n Nothing) a))
           vecs)

  mvecsWritePartial :: forall sh sh' s. KnownShapeX sh'
                    => IShX (sh ++ sh') -> IIxX sh -> Mixed sh' (Ranked n a)
                    -> MixedVecs s (sh ++ sh') (Ranked n a)
                    -> ST s ()
  mvecsWritePartial sh idx arr vecs
    | Dict <- lemKnownReplicate (Proxy @n)
    = mvecsWritePartial sh idx
        (coerce @(Mixed sh' (Ranked n a))
                @(Mixed sh' (Mixed (Replicate n Nothing) a))
           arr)
        (coerce @(MixedVecs s (sh ++ sh') (Ranked n a))
                @(MixedVecs s (sh ++ sh') (Mixed (Replicate n Nothing) a))
           vecs)

  mvecsFreeze :: forall sh s. IShX sh -> MixedVecs s sh (Ranked n a) -> ST s (Mixed sh (Ranked n a))
  mvecsFreeze sh vecs
    | Dict <- lemKnownReplicate (Proxy @n)
    = coerce @(Mixed sh (Mixed (Replicate n Nothing) a))
             @(Mixed sh (Ranked n a))
        <$> mvecsFreeze sh
              (coerce @(MixedVecs s sh (Ranked n a))
                      @(MixedVecs s sh (Mixed (Replicate n Nothing) a))
                      vecs)


-- | The shape of a shape-typed array given as a list of 'SNat' values.
data ShS sh where
  ZSS :: ShS '[]
  (:$$) :: forall n sh. SNat n -> ShS sh -> ShS (n : sh)
deriving instance Show (ShS sh)
deriving instance Eq (ShS sh)
deriving instance Ord (ShS sh)
infixr 3 :$$

-- | A statically-known shape of a shape-typed array.
class KnownShape sh where knownShape :: ShS sh
instance KnownShape '[] where knownShape = ZSS
instance (KnownNat n, KnownShape sh) => KnownShape (n : sh) where knownShape = natSing :$$ knownShape

sshapeKnown :: ShS sh -> Dict KnownShape sh
sshapeKnown ZSS = Dict
sshapeKnown (GHC_SNat :$$ sh) | Dict <- sshapeKnown sh = Dict

lemKnownMapJust :: forall sh. KnownShape sh => Proxy sh -> Dict KnownShapeX (MapJust sh)
lemKnownMapJust _ = X.lemKnownShapeX (go (knownShape @sh))
  where
    go :: ShS sh' -> StaticShX (MapJust sh')
    go ZSS = ZKSX
    go (n :$$ sh) = n :!$@ go sh

lemMapJustPlusApp :: forall sh1 sh2. KnownShape sh1 => Proxy sh1 -> Proxy sh2
                  -> MapJust (sh1 ++ sh2) :~: MapJust sh1 ++ MapJust sh2
lemMapJustPlusApp _ _ = go (knownShape @sh1)
  where
    go :: ShS sh1' -> MapJust (sh1' ++ sh2) :~: MapJust sh1' ++ MapJust sh2
    go ZSS = Refl
    go (_ :$$ sh) | Refl <- go sh = Refl

instance (Elt a, KnownShape sh) => Elt (Shaped sh a) where
  mshape (M_Shaped arr) | Dict <- lemKnownMapJust (Proxy @sh) = mshape arr
  mindex (M_Shaped arr) i | Dict <- lemKnownMapJust (Proxy @sh) = Shaped (mindex arr i)

  mindexPartial :: forall sh1 sh2. Mixed (sh1 ++ sh2) (Shaped sh a) -> IIxX sh1 -> Mixed sh2 (Shaped sh a)
  mindexPartial (M_Shaped arr) i
    | Dict <- lemKnownMapJust (Proxy @sh)
    = coerce @(Mixed sh2 (Mixed (MapJust sh) a)) @(Mixed sh2 (Shaped sh a)) $
        mindexPartial arr i

  mscalar (Shaped x) = M_Shaped (M_Nest x)

  mfromList1 :: forall n sh'. KnownShapeX (n : sh')
             => NonEmpty (Mixed sh' (Shaped sh a)) -> Mixed (n : sh') (Shaped sh a)
  mfromList1 l
    | Dict <- lemKnownMapJust (Proxy @sh)
    = M_Shaped (mfromList1 (coerce l))

  mtoList1 :: forall n sh'. Mixed (n : sh') (Shaped sh a) -> [Mixed sh' (Shaped sh a)]
  mtoList1 (M_Shaped arr)
    | Dict <- lemKnownMapJust (Proxy @sh)
    = coerce @[Mixed sh' (Mixed (MapJust sh) a)] @[Mixed sh' (Shaped sh a)] (mtoList1 arr)

  mlift :: forall sh1 sh2. KnownShapeX sh2
        => (forall sh' b. (KnownShapeX sh', Storable b) => Proxy sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
        -> Mixed sh1 (Shaped sh a) -> Mixed sh2 (Shaped sh a)
  mlift f (M_Shaped arr)
    | Dict <- lemKnownMapJust (Proxy @sh)
    = coerce @(Mixed sh2 (Mixed (MapJust sh) a)) @(Mixed sh2 (Shaped sh a)) $
        mlift f arr

  mlift2 :: forall sh1 sh2 sh3. (KnownShapeX sh2, KnownShapeX sh3)
         => (forall sh' b. (KnownShapeX sh', Storable b) => Proxy sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b -> XArray (sh3 ++ sh') b)
         -> Mixed sh1 (Shaped sh a) -> Mixed sh2 (Shaped sh a) -> Mixed sh3 (Shaped sh a)
  mlift2 f (M_Shaped arr1) (M_Shaped arr2)
    | Dict <- lemKnownMapJust (Proxy @sh)
    = coerce @(Mixed sh3 (Mixed (MapJust sh) a)) @(Mixed sh3 (Shaped sh a)) $
        mlift2 f arr1 arr2

  memptyArray :: forall sh'. IShX sh' -> Mixed sh' (Shaped sh a)
  memptyArray i
    | Dict <- lemKnownMapJust (Proxy @sh)
    = coerce @(Mixed sh' (Mixed (MapJust sh) a)) @(Mixed sh' (Shaped sh a)) $
        memptyArray i

  mshapeTree (Shaped arr)
    | Dict <- lemKnownMapJust (Proxy @sh)
    = first (shCvtXS (knownShape @sh)) (mshapeTree arr)

  mshapeTreeEq _ (sh1, t1) (sh2, t2) = sh1 == sh2 && mshapeTreeEq (Proxy @a) t1 t2

  mshapeTreeEmpty _ (sh, t) = shapeSizeS sh == 0 && mshapeTreeEmpty (Proxy @a) t

  mshowShapeTree _ (sh, t) = "(" ++ show sh ++ ", " ++ mshowShapeTree (Proxy @a) t ++ ")"

  mvecsUnsafeNew idx (Shaped arr)
    | Dict <- lemKnownMapJust (Proxy @sh)
    = MV_Shaped <$> mvecsUnsafeNew idx arr

  mvecsNewEmpty _
    | Dict <- lemKnownMapJust (Proxy @sh)
    = MV_Shaped <$> mvecsNewEmpty (Proxy @(Mixed (MapJust sh) a))

  mvecsWrite :: forall sh' s. IShX sh' -> IIxX sh' -> Shaped sh a -> MixedVecs s sh' (Shaped sh a) -> ST s ()
  mvecsWrite sh idx (Shaped arr) vecs
    | Dict <- lemKnownMapJust (Proxy @sh)
    = mvecsWrite sh idx arr
        (coerce @(MixedVecs s sh' (Shaped sh a)) @(MixedVecs s sh' (Mixed (MapJust sh) a))
           vecs)

  mvecsWritePartial :: forall sh1 sh2 s. KnownShapeX sh2
                    => IShX (sh1 ++ sh2) -> IIxX sh1 -> Mixed sh2 (Shaped sh a)
                    -> MixedVecs s (sh1 ++ sh2) (Shaped sh a)
                    -> ST s ()
  mvecsWritePartial sh idx arr vecs
    | Dict <- lemKnownMapJust (Proxy @sh)
    = mvecsWritePartial sh idx
        (coerce @(Mixed sh2 (Shaped sh a))
                @(Mixed sh2 (Mixed (MapJust sh) a))
           arr)
        (coerce @(MixedVecs s (sh1 ++ sh2) (Shaped sh a))
                @(MixedVecs s (sh1 ++ sh2) (Mixed (MapJust sh) a))
           vecs)

  mvecsFreeze :: forall sh' s. IShX sh' -> MixedVecs s sh' (Shaped sh a) -> ST s (Mixed sh' (Shaped sh a))
  mvecsFreeze sh vecs
    | Dict <- lemKnownMapJust (Proxy @sh)
    = coerce @(Mixed sh' (Mixed (MapJust sh) a))
             @(Mixed sh' (Shaped sh a))
        <$> mvecsFreeze sh
              (coerce @(MixedVecs s sh' (Shaped sh a))
                      @(MixedVecs s sh' (Mixed (MapJust sh) a))
                      vecs)


-- Utility functions to satisfy the type checker sometimes

rewriteMixed :: sh1 :~: sh2 -> Mixed sh1 a -> Mixed sh2 a
rewriteMixed Refl x = x


-- ====== API OF RANKED ARRAYS ====== --

arithPromoteRanked :: forall n a. KnownINat n
                   => (forall sh. KnownShapeX sh => Mixed sh a -> Mixed sh a)
                   -> Ranked n a -> Ranked n a
arithPromoteRanked | Dict <- lemKnownReplicate (Proxy @n) = coerce

arithPromoteRanked2 :: forall n a. KnownINat n
                    => (forall sh. KnownShapeX sh => Mixed sh a -> Mixed sh a -> Mixed sh a)
                    -> Ranked n a -> Ranked n a -> Ranked n a
arithPromoteRanked2 | Dict <- lemKnownReplicate (Proxy @n) = coerce

instance (KnownINat n, Storable a, Num a) => Num (Ranked n (Primitive a)) where
  (+) = arithPromoteRanked2 (+)
  (-) = arithPromoteRanked2 (-)
  (*) = arithPromoteRanked2 (*)
  negate = arithPromoteRanked negate
  abs = arithPromoteRanked abs
  signum = arithPromoteRanked signum
  fromInteger n = case inatSing @n of
                    SZ -> Ranked (M_Primitive (X.scalar (fromInteger n)))
                    SS _ -> error "Data.Array.Nested.fromIntegral(Ranked): \
                                  \Rank non-zero, use explicit mconstant"

-- [PRIMITIVE ELEMENT TYPES LIST] (really, a partial list of just the numeric types)
deriving via Ranked n (Primitive Int) instance KnownINat n => Num (Ranked n Int)
deriving via Ranked n (Primitive Double) instance KnownINat n => Num (Ranked n Double)

type role ListR nominal representational
type ListR :: INat -> Type -> Type
data ListR n i where
  ZR :: ListR Z i
  (:::) :: forall n {i}. i -> ListR n i -> ListR (S n) i
deriving instance Show i => Show (ListR n i)
deriving instance Eq i => Eq (ListR n i)
deriving instance Ord i => Ord (ListR n i)
deriving instance Functor (ListR n)
infixr 3 :::

instance Foldable (ListR n) where
  foldr f z l = foldr f z (listRToList l)

listRToList :: ListR n i -> [i]
listRToList ZR = []
listRToList (i ::: is) = i : listRToList is

knownListR :: ListR n i -> Dict KnownINat n
knownListR ZR = Dict
knownListR (_ ::: l) | Dict <- knownListR l = Dict

-- | An index into a rank-typed array.
type role IxR nominal representational
type IxR :: INat -> Type -> Type
newtype IxR n i = IxR (ListR n i)
  deriving (Show, Eq, Ord)
  deriving newtype (Functor, Foldable)

pattern ZIR :: forall n i. () => n ~ Z => IxR n i
pattern ZIR = IxR ZR

pattern (:.:)
  :: forall {n1} {i}.
     forall n. (S n ~ n1)
  => i -> IxR n i -> IxR n1 i
pattern i :.: sh <- (unconsIxR -> Just (UnconsIxRRes sh i))
  where i :.: IxR sh = IxR (i ::: sh)
{-# COMPLETE ZIR, (:.:) #-}
infixr 3 :.:

data UnconsIxRRes i n1 =
  forall n. ((S n) ~ n1) => UnconsIxRRes (IxR n i) i
unconsIxR :: IxR n1 i -> Maybe (UnconsIxRRes i n1)
unconsIxR (IxR (i ::: sh')) = Just (UnconsIxRRes (IxR sh') i)
unconsIxR (IxR ZR) = Nothing

type IIxR n = IxR n Int

knownIxR :: IxR n i -> Dict KnownINat n
knownIxR (IxR sh) = knownListR sh

type role ShR nominal representational
type ShR :: INat -> Type -> Type
newtype ShR n i = ShR (ListR n i)
  deriving (Show, Eq, Ord)
  deriving newtype (Functor, Foldable)

type IShR n = ShR n Int

pattern ZSR :: forall n i. () => n ~ Z => ShR n i
pattern ZSR = ShR ZR

pattern (:$:)
  :: forall {n1} {i}.
     forall n. (S n ~ n1)
  => i -> ShR n i -> ShR n1 i
pattern i :$: sh <- (unconsShR -> Just (UnconsShRRes sh i))
  where i :$: (ShR sh) = ShR (i ::: sh)
{-# COMPLETE ZSR, (:$:) #-}
infixr 3 :$:

data UnconsShRRes i n1 =
  forall n. S n ~ n1 => UnconsShRRes (ShR n i) i
unconsShR :: ShR n1 i -> Maybe (UnconsShRRes i n1)
unconsShR (ShR (i ::: sh')) = Just (UnconsShRRes (ShR sh') i)
unconsShR (ShR ZR) = Nothing

knownShR :: ShR n i -> Dict KnownINat n
knownShR (ShR sh) = knownListR sh

zeroIxR :: SINat n -> IIxR n
zeroIxR SZ = ZIR
zeroIxR (SS n) = 0 :.: zeroIxR n

ixCvtXR :: IIxX sh -> IIxR (X.Rank sh)
ixCvtXR ZIX = ZIR
ixCvtXR (n :.@ idx) = n :.: ixCvtXR idx
ixCvtXR (n :.? idx) = n :.: ixCvtXR idx

shCvtXR :: IShX sh -> IShR (X.Rank sh)
shCvtXR ZSX = ZSR
shCvtXR (n :$@ idx) = X.fromSNat' n :$: shCvtXR idx
shCvtXR (n :$? idx) = n :$: shCvtXR idx

ixCvtRX :: IIxR n -> IIxX (Replicate n Nothing)
ixCvtRX ZIR = ZIX
ixCvtRX (n :.: idx) = n :.? ixCvtRX idx

shCvtRX :: IShR n -> IShX (Replicate n Nothing)
shCvtRX ZSR = ZSX
shCvtRX (n :$: idx) = n :$? shCvtRX idx

shapeSizeR :: IShR n -> Int
shapeSizeR ZSR = 1
shapeSizeR (n :$: sh) = n * shapeSizeR sh


rshape :: forall n a. (KnownINat n, Elt a) => Ranked n a -> IShR n
rshape (Ranked arr)
  | Dict <- lemKnownReplicate (Proxy @n)
  , Refl <- lemRankReplicate (Proxy @n)
  = shCvtXR (mshape arr)

rindex :: Elt a => Ranked n a -> IIxR n -> a
rindex (Ranked arr) idx = mindex arr (ixCvtRX idx)

rindexPartial :: forall n m a. (KnownINat n, Elt a) => Ranked (n +! m) a -> IIxR n -> Ranked m a
rindexPartial (Ranked arr) idx =
  Ranked (mindexPartial @a @(Replicate n Nothing) @(Replicate m Nothing)
            (rewriteMixed (lemReplicatePlusApp (Proxy @n) (Proxy @m) (Proxy @Nothing)) arr)
            (ixCvtRX idx))

-- | __WARNING__: All values returned from the function must have equal shape.
-- See the documentation of 'mgenerate' for more details.
rgenerate :: forall n a. Elt a => IShR n -> (IIxR n -> a) -> Ranked n a
rgenerate sh f
  | Dict <- knownShR sh
  , Dict <- lemKnownReplicate (Proxy @n)
  , Refl <- lemRankReplicate (Proxy @n)
  = Ranked (mgenerate (shCvtRX sh) (f . ixCvtXR))

-- | See the documentation of 'mlift'.
rlift :: forall n1 n2 a. (KnownINat n2, Elt a)
      => (forall sh' b. KnownShapeX sh' => Proxy sh' -> XArray (Replicate n1 Nothing ++ sh') b -> XArray (Replicate n2 Nothing ++ sh') b)
      -> Ranked n1 a -> Ranked n2 a
rlift f (Ranked arr)
  | Dict <- lemKnownReplicate (Proxy @n2)
  = Ranked (mlift f arr)

rsumOuter1P :: forall n a.
               (Storable a, Num a, KnownINat n)
            => Ranked (S n) (Primitive a) -> Ranked n (Primitive a)
rsumOuter1P (Ranked arr)
  | Dict <- lemKnownReplicate (Proxy @n)
  = Ranked
    . coerce @(XArray (Replicate n 'Nothing) a) @(Mixed (Replicate n 'Nothing) (Primitive a))
    . X.sumOuter (() :!$? ZKSX) (knownShapeX @(Replicate n Nothing))
    . coerce @(Mixed (Replicate (S n) Nothing) (Primitive a)) @(XArray (Replicate (S n) Nothing) a)
    $ arr

rsumOuter1 :: forall n a.
              (Storable a, Num a, PrimElt a, KnownINat n)
           => Ranked (S n) a -> Ranked n a
rsumOuter1 = coerce fromPrimitive . rsumOuter1P @n @a . coerce toPrimitive

rtranspose :: forall n a. (KnownINat n, Elt a) => [Int] -> Ranked n a -> Ranked n a
rtranspose perm (Ranked arr)
  | Dict <- lemKnownReplicate (Proxy @n)
  = Ranked (mtranspose perm arr)

rappend :: forall n a. (KnownINat n, Elt a)
        => Ranked (S n) a -> Ranked (S n) a -> Ranked (S n) a
rappend | Dict <- lemKnownReplicate (Proxy @n) = coerce mappend

rscalar :: Elt a => a -> Ranked I0 a
rscalar x = Ranked (mscalar x)

rfromVectorP :: forall n a. (KnownINat n, Storable a) => IShR n -> VS.Vector a -> Ranked n (Primitive a)
rfromVectorP sh v
  | Dict <- lemKnownReplicate (Proxy @n)
  = Ranked (mfromVectorP (shCvtRX sh) v)

rfromVector :: forall n a. (KnownINat n, Storable a, PrimElt a) => IShR n -> VS.Vector a -> Ranked n a
rfromVector sh v = coerce fromPrimitive (rfromVectorP sh v)

rtoVectorP :: Storable a => Ranked n (Primitive a) -> VS.Vector a
rtoVectorP = coerce mtoVectorP

rtoVector :: (Storable a, PrimElt a) => Ranked n a -> VS.Vector a
rtoVector = coerce mtoVector

rfromList1 :: forall n a. (KnownINat n, Elt a) => NonEmpty (Ranked n a) -> Ranked (S n) a
rfromList1 l
  | Dict <- lemKnownReplicate (Proxy @n)
  = Ranked (mfromList1 (coerce l))

rfromList :: Elt a => NonEmpty a -> Ranked I1 a
rfromList = Ranked . mfromList1 . fmap mscalar

rtoList :: Elt a => Ranked (S n) a -> [Ranked n a]
rtoList (Ranked arr) = coerce (mtoList1 arr)

rtoList1 :: Elt a => Ranked I1 a -> [a]
rtoList1 = map runScalar . rtoList

runScalar :: Elt a => Ranked I0 a -> a
runScalar arr = rindex arr ZIR

rconstantP :: forall n a. (KnownINat n, Storable a) => IShR n -> a -> Ranked n (Primitive a)
rconstantP sh x
  | Dict <- lemKnownReplicate (Proxy @n)
  = Ranked (mconstantP (shCvtRX sh) x)

rconstant :: forall n a. (KnownINat n, Storable a, PrimElt a)
          => IShR n -> a -> Ranked n a
rconstant sh x = coerce fromPrimitive (rconstantP sh x)

rslice :: (KnownINat n, Elt a) => [(Int, Int)] -> Ranked n a -> Ranked n a
rslice ivs = rlift $ \_ -> X.slice ivs

rrev1 :: (KnownINat n, Elt a) => Ranked (S n) a -> Ranked (S n) a
rrev1 = rlift $ \_ -> X.rev1

rreshape :: forall n n' a. (KnownINat n, KnownINat n', Elt a)
         => IShR n' -> Ranked n a -> Ranked n' a
rreshape sh' (Ranked arr)
  | Dict <- lemKnownReplicate (Proxy @n)
  , Dict <- lemKnownReplicate (Proxy @n')
  = Ranked (mreshape (shCvtRX sh') arr)


-- ====== API OF SHAPED ARRAYS ====== --

arithPromoteShaped :: forall sh a. KnownShape sh
                   => (forall shx. KnownShapeX shx => Mixed shx a -> Mixed shx a)
                   -> Shaped sh a -> Shaped sh a
arithPromoteShaped | Dict <- lemKnownMapJust (Proxy @sh) = coerce

arithPromoteShaped2 :: forall sh a. KnownShape sh
                    => (forall shx. KnownShapeX shx => Mixed shx a -> Mixed shx a -> Mixed shx a)
                    -> Shaped sh a -> Shaped sh a -> Shaped sh a
arithPromoteShaped2 | Dict <- lemKnownMapJust (Proxy @sh) = coerce

instance (KnownShape sh, Storable a, Num a) => Num (Shaped sh (Primitive a)) where
  (+) = arithPromoteShaped2 (+)
  (-) = arithPromoteShaped2 (-)
  (*) = arithPromoteShaped2 (*)
  negate = arithPromoteShaped negate
  abs = arithPromoteShaped abs
  signum = arithPromoteShaped signum
  fromInteger n = sconstantP (fromInteger n)

-- [PRIMITIVE ELEMENT TYPES LIST] (really, a partial list of just the numeric types)
deriving via Shaped sh (Primitive Int) instance KnownShape sh => Num (Shaped sh Int)
deriving via Shaped sh (Primitive Double) instance KnownShape sh => Num (Shaped sh Double)

type role ListS nominal representational
type ListS :: [Nat] -> Type -> Type
data ListS sh i where
  ZS :: ListS '[] i
  (::$) :: forall n sh {i}. i -> ListS sh i -> ListS (n : sh) i
deriving instance Show i => Show (ListS sh i)
deriving instance Eq i => Eq (ListS sh i)
deriving instance Ord i => Ord (ListS sh i)
deriving instance Functor (ListS sh)
infixr 3 ::$

instance Foldable (ListS sh) where
  foldr f z l = foldr f z (listSToList l)

listSToList :: ListS sh i -> [i]
listSToList ZS = []
listSToList (i ::$ is) = i : listSToList is

-- | An index into a shape-typed array.
--
-- For convenience, this contains regular 'Int's instead of bounded integers
-- (traditionally called \"@Fin@\"). Note that because the shape of a
-- shape-typed array is known statically, you can also retrieve the array shape
-- from a 'KnownShape' dictionary.
type role IxS nominal representational
type IxS :: [Nat] -> Type -> Type
newtype IxS sh i = IxS (ListS sh i)
  deriving (Show, Eq, Ord)
  deriving newtype (Functor, Foldable)

pattern ZIS :: forall sh i. () => sh ~ '[] => IxS sh i
pattern ZIS = IxS ZS

pattern (:.$)
  :: forall {sh1} {i}.
     forall n sh. (n : sh ~ sh1)
  => i -> IxS sh i -> IxS sh1 i
pattern i :.$ shl <- (unconsIxS -> Just (UnconsIxSRes shl i))
  where i :.$ IxS shl = IxS (i ::$ shl)
{-# COMPLETE ZIS, (:.$) #-}
infixr 3 :.$

data UnconsIxSRes i sh1 =
  forall n sh. (n : sh ~ sh1) => UnconsIxSRes (IxS sh i) i
unconsIxS :: IxS sh1 i -> Maybe (UnconsIxSRes i sh1)
unconsIxS (IxS (i ::$ shl')) = Just (UnconsIxSRes (IxS shl') i)
unconsIxS (IxS ZS) = Nothing

type IIxS sh = IxS sh Int

data UnconsShSRes sh1 =
  forall n sh. (n : sh ~ sh1) => UnconsShSRes (ShS sh) (SNat n)
unconsShS :: ShS sh1 -> Maybe (UnconsShSRes sh1)
unconsShS (i :$$ shl') = Just (UnconsShSRes shl' i)
unconsShS ZSS = Nothing

zeroIxS :: ShS sh -> IIxS sh
zeroIxS ZSS = ZIS
zeroIxS (_ :$$ sh) = 0 :.$ zeroIxS sh

ixCvtXS :: ShS sh -> IIxX (MapJust sh) -> IIxS sh
ixCvtXS ZSS ZIX = ZIS
ixCvtXS (_ :$$ sh) (n :.@ idx) = n :.$ ixCvtXS sh idx

shCvtXS :: ShS sh -> IShX (MapJust sh) -> ShS sh
shCvtXS ZSS ZSX = ZSS
shCvtXS (_ :$$ sh) (n :$@ idx) = n :$$ shCvtXS sh idx

ixCvtSX :: IIxS sh -> IIxX (MapJust sh)
ixCvtSX ZIS = ZIX
ixCvtSX (n :.$ sh) = n :.@ ixCvtSX sh

shCvtSX :: ShS sh -> IShX (MapJust sh)
shCvtSX ZSS = ZSX
shCvtSX (n :$$ sh) = n :$@ shCvtSX sh

shapeSizeS :: ShS sh -> Int
shapeSizeS ZSS = 1
shapeSizeS (n :$$ sh) = X.fromSNat' n * shapeSizeS sh


-- | This does not touch the passed array, all information comes from 'KnownShape'.
sshape :: forall sh a. (KnownShape sh, Elt a) => Shaped sh a -> ShS sh
sshape _ = knownShape @sh

sindex :: Elt a => Shaped sh a -> IIxS sh -> a
sindex (Shaped arr) idx = mindex arr (ixCvtSX idx)

sindexPartial :: forall sh1 sh2 a. (KnownShape sh1, Elt a) => Shaped (sh1 ++ sh2) a -> IIxS sh1 -> Shaped sh2 a
sindexPartial (Shaped arr) idx =
  Shaped (mindexPartial @a @(MapJust sh1) @(MapJust sh2)
            (rewriteMixed (lemMapJustPlusApp (Proxy @sh1) (Proxy @sh2)) arr)
            (ixCvtSX idx))

-- | __WARNING__: All values returned from the function must have equal shape.
-- See the documentation of 'mgenerate' for more details.
sgenerate :: forall sh a. (KnownShape sh, Elt a) => (IIxS sh -> a) -> Shaped sh a
sgenerate f
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mgenerate (shCvtSX (knownShape @sh)) (f . ixCvtXS (knownShape @sh)))

-- | See the documentation of 'mlift'.
slift :: forall sh1 sh2 a. (KnownShape sh2, Elt a)
      => (forall sh' b. KnownShapeX sh' => Proxy sh' -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b)
      -> Shaped sh1 a -> Shaped sh2 a
slift f (Shaped arr)
  | Dict <- lemKnownMapJust (Proxy @sh2)
  = Shaped (mlift f arr)

ssumOuter1P :: forall sh n a.
               (Storable a, Num a, KnownNat n, KnownShape sh)
            => Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a)
ssumOuter1P (Shaped arr)
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped
    . coerce @(XArray (MapJust sh) a) @(Mixed (MapJust sh) (Primitive a))
    . X.sumOuter (natSing @n :!$@ ZKSX) (knownShapeX @(MapJust sh))
    . coerce @(Mixed (Just n : MapJust sh) (Primitive a)) @(XArray (Just n : MapJust sh) a)
    $ arr

ssumOuter1 :: forall sh n a.
              (Storable a, Num a, PrimElt a, KnownNat n, KnownShape sh)
           => Shaped (n : sh) a -> Shaped sh a
ssumOuter1 = coerce fromPrimitive . ssumOuter1P @sh @n @a . coerce toPrimitive

stranspose :: forall sh a. (KnownShape sh, Elt a) => [Int] -> Shaped sh a -> Shaped sh a
stranspose perm (Shaped arr)
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mtranspose perm arr)

sappend :: forall n m sh a. (KnownNat n, KnownNat m, KnownShape sh, Elt a)
        => Shaped (n : sh) a -> Shaped (m : sh) a -> Shaped (n + m : sh) a
sappend | Dict <- lemKnownMapJust (Proxy @sh) = coerce mappend

sscalar :: Elt a => a -> Shaped '[] a
sscalar x = Shaped (mscalar x)

sfromVectorP :: forall sh a. (KnownShape sh, Storable a) => VS.Vector a -> Shaped sh (Primitive a)
sfromVectorP v
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mfromVectorP (shCvtSX (knownShape @sh)) v)

sfromVector :: forall sh a. (KnownShape sh, Storable a, PrimElt a) => VS.Vector a -> Shaped sh a
sfromVector v = coerce fromPrimitive (sfromVectorP @sh @a v)

stoVectorP :: Storable a => Shaped sh (Primitive a) -> VS.Vector a
stoVectorP = coerce mtoVectorP

stoVector :: (Storable a, PrimElt a) => Shaped sh a -> VS.Vector a
stoVector = coerce mtoVector

sfromList1 :: forall n sh a. (KnownNat n, KnownShape sh, Elt a)
           => NonEmpty (Shaped sh a) -> Shaped (n : sh) a
sfromList1 l
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mfromList1 (coerce l))

sfromList :: (KnownNat n, Elt a) => NonEmpty a -> Shaped '[n] a
sfromList = Shaped . mfromList1 . fmap mscalar

stoList :: Elt a => Shaped (n : sh) a -> [Shaped sh a]
stoList (Shaped arr) = coerce (mtoList1 arr)

stoList1 :: Elt a => Shaped '[n] a -> [a]
stoList1 = map sunScalar . stoList

sunScalar :: Elt a => Shaped '[] a -> a
sunScalar arr = sindex arr ZIS

sconstantP :: forall sh a. (KnownShape sh, Storable a) => a -> Shaped sh (Primitive a)
sconstantP x
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mconstantP (shCvtSX (knownShape @sh)) x)

sconstant :: forall sh a. (KnownShape sh, Storable a, PrimElt a)
          => a -> Shaped sh a
sconstant x = coerce fromPrimitive (sconstantP @sh x)

sslice :: (KnownShape sh, Elt a) => [(Int, Int)] -> Shaped sh a -> Shaped sh a
sslice ivs = slift $ \_ -> X.slice ivs

srev1 :: (KnownNat n, KnownShape sh, Elt a) => Shaped (n : sh) a -> Shaped (n : sh) a
srev1 = slift $ \_ -> X.rev1

sreshape :: forall sh sh' a. (KnownShape sh, KnownShape sh', Elt a)
         => ShS sh' -> Shaped sh a -> Shaped sh' a
sreshape sh' (Shaped arr)
  | Dict <- lemKnownMapJust (Proxy @sh)
  , Dict <- lemKnownMapJust (Proxy @sh')
  = Shaped (mreshape (shCvtSX sh') arr)