{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# 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 'StaticShapeX'
  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, KnownShapeX(..), StaticShapeX(..), 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


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 -> StaticShapeX (Replicate m Nothing)
    go SZ = ZSX
    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

ixAppSplit :: Proxy sh' -> StaticShapeX sh -> IIxX (sh ++ sh') -> (IIxX sh, IIxX sh')
ixAppSplit _ ZSX idx = (ZIX, idx)
ixAppSplit p (_ :$@ ssh) (i :.@ idx) = first (i :.@) (ixAppSplit p ssh idx)
ixAppSplit p (_ :$? ssh) (i :.? idx) = first (i :.?) (ixAppSplit 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


-- | 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, an empty array may have
-- elements with nonempty sizes, but then the whole array is still empty.

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

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 mirrorring '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)

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 !(IIxX sh2) !(MixedVecs s (sh1 ++ sh2) a)


-- | Tree giving the shape of every array component.
type family ShapeTree a where
  ShapeTree (Primitive _) = ()
  ShapeTree Int = ()
  ShapeTree Double = ()
  ShapeTree () = ()

  ShapeTree (a, b) = (ShapeTree a, ShapeTree b)
  ShapeTree (Mixed sh a) = (IIxX sh, ShapeTree a)
  ShapeTree (Ranked n a) = (IIxR n, ShapeTree a)
  ShapeTree (Shaped sh a) = (IIxS 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 -> IIxX 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.
  mfromList :: forall n sh. KnownShapeX (n : sh) => NonEmpty (Mixed sh a) -> Mixed (n : sh) a

  mtoList :: 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 ====== --
  -- Remember I said that this module needed better management of exports?

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

  mshapeTree :: a -> ShapeTree a

  mshapeTreeZero :: Proxy 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 :: IIxX 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 :: IIxX 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' => IIxX (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 :: IIxX 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)
  mfromList l = M_Primitive (X.fromList knownShapeX (coerce (toList l)))
  mtoList (M_Primitive arr) = coerce (X.toList 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.generate sh (error $ "memptyArray Int: shape was not empty (" ++ show sh ++ ")"))
  mshapeTree _ = ()
  mshapeTreeZero _ = ()
  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'
    => IIxX (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

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)
  mfromList l = M_Tup2 (mfromList ((\(M_Tup2 x _) -> x) <$> l))
                       (mfromList ((\(M_Tup2 _ y) -> y) <$> l))
  mtoList (M_Tup2 a b) = zipWith M_Tup2 (mtoList a) (mtoList 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)
  mshapeTreeZero _ = (mshapeTreeZero (Proxy @a), mshapeTreeZero (Proxy @b))
  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) -> IIxX sh
  mshape (M_Nest arr)
    | Dict <- X.lemAppKnownShapeX (knownShapeX @sh) (knownShapeX @sh')
    = fst (ixAppSplit (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

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

  mtoList (M_Nest arr) = coerce (mtoList 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.ixAppend sh (X.zeroIxX (knownShapeX @sh'))))

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

  mshapeTreeZero _ = (X.zeroIxX (knownShapeX @sh'), mshapeTreeZero (Proxy @a))

  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.ixAppend sh (mshape example))
                                                 (mindex example (X.zeroIxX (knownShapeX @sh')))
    where
      sh' = mshape example

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

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

  mvecsWritePartial :: forall sh1 sh2 s. KnownShapeX sh2
                    => IIxX (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.ixAppend sh12 sh') idx arr vecs

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


-- | Check whether a given shape corresponds on the statically-known components of the shape.
checkBounds :: IIxX sh' -> StaticShapeX sh' -> Bool
checkBounds ZIX ZSX = True
checkBounds (n :.@ sh') (n' :$@ ssh') = n == fromIntegral (fromSNat n') && checkBounds sh' ssh'
checkBounds (_ :.? sh') (() :$? ssh') = checkBounds sh' ssh'

-- | 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) => IIxX sh -> (IIxX sh -> a) -> Mixed sh a
mgenerate sh f
  -- TODO: Do we need this checkBounds check elsewhere as well?
  | not (checkBounds sh (knownShapeX @sh)) =
      error $ "mgenerate: Shape " ++ show sh ++ " not valid for shape type " ++ show (knownShapeX @sh)
  -- If the shape is empty, there is no first element, so we should not try to
  -- generate it.
  | X.shapeSize sh == 0 = memptyArray sh
  | otherwise =
      let firstidx = X.zeroIxX' sh
          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_ (tail (X.enumShape sh)) $ \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

mfromVector :: forall sh a. (KnownShapeX sh, Storable a) => IIxX sh -> VS.Vector a -> Mixed sh (Primitive a)
mfromVector sh v
  | not (checkBounds sh (knownShapeX @sh)) =
      error $ "mfromVector: Shape " ++ show sh ++ " not valid for shape type " ++ show (knownShapeX @sh)
  | otherwise =
      M_Primitive (X.fromVector sh v)

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

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

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

mconstantP :: forall sh a. (KnownShapeX sh, Storable a) => IIxX sh -> a -> Mixed sh (Primitive a)
mconstantP sh x
  | not (checkBounds sh (knownShapeX @sh)) =
      error $ "mconstant: Shape " ++ show sh ++ " not valid for shape type " ++ show (knownShapeX @sh)
  | otherwise =
      M_Primitive (X.constant sh x)

-- | This 'Coercible' constraint holds for the scalar types for which 'Mixed'
-- is defined.
mconstant :: forall sh a. (KnownShapeX sh, Storable a, Coercible (Mixed sh (Primitive a)) (Mixed sh a))
          => IIxX sh -> a -> Mixed sh a
mconstant sh x = coerce (mconstantP sh x)

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

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"

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)

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

  mtoList :: forall m sh. Mixed (m : sh) (Ranked n a) -> [Mixed sh (Ranked n a)]
  mtoList (M_Ranked arr)
    | Dict <- lemKnownReplicate (Proxy @n)
    = coerce @[Mixed sh (Mixed (Replicate n 'Nothing) a)] @[Mixed sh (Ranked n a)] (mtoList 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. IIxX 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 ixCvtXR (mshapeTree arr)

  mshapeTreeZero _ = (zeroIxR (inatSing @n), mshapeTreeZero (Proxy @a))

  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. IIxX 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'
                    => IIxX (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. IIxX 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 SShape sh where
  ShNil :: SShape '[]
  ShCons :: SNat n -> SShape sh -> SShape (n : sh)
deriving instance Show (SShape sh)
infixr 5 `ShCons`

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

sshapeKnown :: SShape sh -> Dict KnownShape sh
sshapeKnown ShNil = Dict
sshapeKnown (ShCons 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 :: SShape sh' -> StaticShapeX (MapJust sh')
    go ShNil = ZSX
    go (ShCons 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 :: SShape sh1' -> MapJust (sh1' ++ sh2) :~: MapJust sh1' ++ MapJust sh2
    go ShNil = Refl
    go (ShCons _ 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)

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

  mtoList :: forall n sh'. Mixed (n : sh') (Shaped sh a) -> [Mixed sh' (Shaped sh a)]
  mtoList (M_Shaped arr)
    | Dict <- lemKnownMapJust (Proxy @sh)
    = coerce @[Mixed sh' (Mixed (MapJust sh) a)] @[Mixed sh' (Shaped sh a)] (mtoList 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'. IIxX 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 (ixCvtXS (knownShape @sh)) (mshapeTree arr)

  mshapeTreeZero _ = (zeroIxS (knownShape @sh), mshapeTreeZero (Proxy @a))

  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. IIxX 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
                    => IIxX (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. IIxX 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"

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)
infixr 3 :::

deriving stock instance Functor (ListR n)

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 instance Functor (IxR n)

instance Foldable (IxR n) where
  foldr f z (IxR l) = foldr f z l

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 sh) = case sh of
  i ::: sh' -> Just (UnconsIxRRes (IxR sh') i)
  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 instance Functor (ShR n)

instance Foldable (ShR n) where
  foldr f z (ShR l) = foldr f z l

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 sh) = case sh of
  i ::: sh' -> Just (UnconsShRRes (ShR sh') i)
  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

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

shapeSizeR :: IIxR n -> Int
shapeSizeR ZIR = 1
shapeSizeR (n :.: sh) = n * shapeSizeR sh


rshape :: forall n a. (KnownINat n, Elt a) => Ranked n a -> IIxR n
rshape (Ranked arr)
  | Dict <- lemKnownReplicate (Proxy @n)
  , Refl <- lemRankReplicate (Proxy @n)
  = ixCvtXR (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 => IIxR n -> (IIxR n -> a) -> Ranked n a
rgenerate sh f
  | Dict <- knownIxR sh
  , Dict <- lemKnownReplicate (Proxy @n)
  , Refl <- lemRankReplicate (Proxy @n)
  = Ranked (mgenerate (ixCvtRX 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)

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

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)

rfromVector :: forall n a. (KnownINat n, Storable a) => IIxR n -> VS.Vector a -> Ranked n (Primitive a)
rfromVector sh v
  | Dict <- lemKnownReplicate (Proxy @n)
  = Ranked (mfromVector (ixCvtRX sh) v)

rfromList :: forall n a. (KnownINat n, Elt a) => NonEmpty (Ranked n a) -> Ranked (S n) a
rfromList l
  | Dict <- lemKnownReplicate (Proxy @n)
  = Ranked (mfromList ((\(Ranked x) -> x) <$> l))

rfromList1 :: Elt a => NonEmpty a -> Ranked I1 a
rfromList1 = Ranked . mfromList . fmap mscalar

rtoList :: Elt a => Ranked (S n) a -> [Ranked n a]
rtoList (Ranked arr) = coerce (mtoList 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) => IIxR n -> a -> Ranked n (Primitive a)
rconstantP sh x
  | Dict <- lemKnownReplicate (Proxy @n)
  = Ranked (mconstantP (ixCvtRX sh) x)

rconstant :: forall n a. (KnownINat n, Storable a, Coercible (Mixed (Replicate n Nothing) (Primitive a)) (Mixed (Replicate n Nothing) a))
          => IIxR n -> a -> Ranked n a
rconstant sh x = coerce (rconstantP sh x)

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


-- ====== 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)

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)
infixr 3 ::$

deriving stock instance Functor (ListS sh)

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 instance Functor (IxS sh)

instance Foldable (IxS sh) where
  foldr f z (IxS l) = foldr f z l

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 shl) = case shl of
  i ::$ shl' -> Just (UnconsIxSRes (IxS shl') i)
  ZS -> Nothing

type IIxS sh = IxS sh Int

type role ShS nominal representational
type ShS :: [Nat] -> Type -> Type
newtype ShS sh i = ShS (ListS sh i)
  deriving (Show, Eq, Ord)

deriving newtype instance Functor (ShS sh)

instance Foldable (ShS sh) where
  foldr f z (ShS l) = foldr f z l

pattern ZSS :: forall sh i. () => sh ~ '[] => ShS sh i
pattern ZSS = ShS ZS

pattern (:$$)
  :: forall {sh1} {i}.
     forall n sh. (n : sh ~ sh1)
  => i -> ShS sh i -> ShS sh1 i
pattern i :$$ shl <- (unconsShS -> Just (UnconsShSRes shl i))
  where i :$$ (ShS shl) = ShS (i ::$ shl)
{-# COMPLETE ZSS, (:$$) #-}
infixr 3 :$$

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

zeroIxS :: SShape sh -> IIxS sh
zeroIxS ShNil = ZIS
zeroIxS (ShCons _ sh) = 0 :.$ zeroIxS sh

cvtSShapeIxS :: SShape sh -> IIxS sh
cvtSShapeIxS ShNil = ZIS
cvtSShapeIxS (ShCons n sh) = fromIntegral (fromSNat n) :.$ cvtSShapeIxS sh

ixCvtXS :: SShape sh -> IIxX (MapJust sh) -> IIxS sh
ixCvtXS ShNil ZIX = ZIS
ixCvtXS (ShCons _ sh) (n :.@ idx) = n :.$ ixCvtXS sh idx

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

shapeSizeS :: IIxS sh -> Int
shapeSizeS ZIS = 1
shapeSizeS (n :.$ sh) = 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 -> IIxS sh
sshape _ = cvtSShapeIxS (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 (ixCvtSX (cvtSShapeIxS (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)

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

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)

sfromVector :: forall sh a. (KnownShape sh, Storable a) => VS.Vector a -> Shaped sh (Primitive a)
sfromVector v
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mfromVector (ixCvtSX (cvtSShapeIxS (knownShape @sh))) v)

sfromList :: forall n sh a. (KnownNat n, KnownShape sh, Elt a)
          => NonEmpty (Shaped sh a) -> Shaped (n : sh) a
sfromList l
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mfromList ((\(Shaped x) -> x) <$> l))

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

stoList :: Elt a => Shaped (n : sh) a -> [Shaped sh a]
stoList (Shaped arr) = coerce (mtoList 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 (ixCvtSX (cvtSShapeIxS (knownShape @sh))) x)

sconstant :: forall sh a. (KnownShape sh, Storable a, Coercible (Mixed (MapJust sh) (Primitive a)) (Mixed (MapJust sh) a))
          => a -> Shaped sh a
sconstant x = coerce (sconstantP @sh x)

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