{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
module Data.Array.Nested.Internal.Ranked where

import Prelude hiding (mappend, mconcat)

import Control.DeepSeq (NFData)
import Control.Monad.ST
import Data.Array.RankedS qualified as S
import Data.Bifunctor (first)
import Data.Coerce (coerce)
import Data.Foldable (toList)
import Data.Kind (Type)
import Data.List.NonEmpty (NonEmpty)
import Data.Proxy
import Data.Type.Equality
import Data.Vector.Storable qualified as VS
import Foreign.Storable (Storable)
import GHC.Float qualified (log1p, expm1, log1pexp, log1mexp)
import GHC.TypeLits
import GHC.TypeNats qualified as TN

import Data.Array.Mixed.XArray (XArray(..))
import Data.Array.Mixed.XArray qualified as X
import Data.Array.Mixed.Internal.Arith
import Data.Array.Mixed.Lemmas
import Data.Array.Mixed.Permutation
import Data.Array.Mixed.Shape
import Data.Array.Mixed.Types
import Data.Array.Nested.Internal.Mixed
import Data.Array.Nested.Internal.Shape


-- | A rank-typed array: the number of dimensions of the array (its /rank/) is
-- represented on the type level as a 'Nat'.
--
-- 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).
--
-- 'Ranked' is a newtype around a 'Mixed' of 'Nothing's.
type Ranked :: Nat -> Type -> Type
newtype Ranked n a = Ranked (Mixed (Replicate n Nothing) a)
deriving instance Eq (Mixed (Replicate n Nothing) a) => Eq (Ranked n a)
deriving instance Ord (Mixed (Replicate n Nothing) a) => Ord (Ranked n a)
deriving instance NFData (Mixed (Replicate n Nothing) a) => NFData (Ranked n a)

instance (Show a, Elt a) => Show (Ranked n a) where
  showsPrec d arr = showParen (d > 10) $
    showString "rfromListLinear " . shows (toList (rshape arr)) . showString " "
      . shows (rtoListLinear arr)

-- 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 via (ShowViaToListLinear sh (Ranked n a)) instance (Show a, Elt a) => Show (Mixed sh (Ranked n a))

newtype instance MixedVecs s sh (Ranked n a) = MV_Ranked (MixedVecs s sh (Mixed (Replicate n Nothing) 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 => Elt (Ranked n a) where
  mshape (M_Ranked arr) = mshape arr
  mindex (M_Ranked arr) i = 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 =
    coerce @(Mixed sh' (Mixed (Replicate n Nothing) a)) @(Mixed sh' (Ranked n a)) $
        mindexPartial arr i

  mscalar (Ranked x) = M_Ranked (M_Nest ZSX x)

  mfromListOuter :: forall sh. NonEmpty (Mixed sh (Ranked n a)) -> Mixed (Nothing : sh) (Ranked n a)
  mfromListOuter l = M_Ranked (mfromListOuter (coerce l))

  mtoListOuter :: forall m sh. Mixed (m : sh) (Ranked n a) -> [Mixed sh (Ranked n a)]
  mtoListOuter (M_Ranked arr) =
    coerce @[Mixed sh (Mixed (Replicate n 'Nothing) a)] @[Mixed sh (Ranked n a)] (mtoListOuter arr)

  mlift :: forall sh1 sh2.
           StaticShX sh2
        -> (forall sh' b. Storable b => StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
        -> Mixed sh1 (Ranked n a) -> Mixed sh2 (Ranked n a)
  mlift ssh2 f (M_Ranked arr) =
    coerce @(Mixed sh2 (Mixed (Replicate n Nothing) a)) @(Mixed sh2 (Ranked n a)) $
      mlift ssh2 f arr

  mlift2 :: forall sh1 sh2 sh3.
            StaticShX sh3
         -> (forall sh' b. Storable b => StaticShX 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 ssh3 f (M_Ranked arr1) (M_Ranked arr2) =
    coerce @(Mixed sh3 (Mixed (Replicate n Nothing) a)) @(Mixed sh3 (Ranked n a)) $
      mlift2 ssh3 f arr1 arr2

  mliftL :: forall sh1 sh2.
            StaticShX sh2
         -> (forall sh' b. Storable b => StaticShX sh' -> NonEmpty (XArray (sh1 ++ sh') b) -> NonEmpty (XArray (sh2 ++ sh') b))
         -> NonEmpty (Mixed sh1 (Ranked n a)) -> NonEmpty (Mixed sh2 (Ranked n a))
  mliftL ssh2 f l =
    coerce @(NonEmpty (Mixed sh2 (Mixed (Replicate n Nothing) a)))
           @(NonEmpty (Mixed sh2 (Ranked n a))) $
      mliftL ssh2 f (coerce l)

  mcast ssh1 sh2 psh' (M_Ranked arr) = M_Ranked (mcast ssh1 sh2 psh' arr)

  mtranspose perm (M_Ranked arr) = M_Ranked (mtranspose perm arr)

  mconcat l = M_Ranked (mconcat (coerce l))

  type ShapeTree (Ranked n a) = (IShR n, ShapeTree a)

  mshapeTree (Ranked arr) = first shCvtXR' (mshapeTree arr)

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

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

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

  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 =
    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.
                       IShX (sh ++ sh') -> IIxX sh -> Mixed sh' (Ranked n a)
                    -> MixedVecs s (sh ++ sh') (Ranked n a)
                    -> ST s ()
  mvecsWritePartial sh idx arr vecs =
    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 =
    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)

instance (KnownNat n, KnownElt a) => KnownElt (Ranked n a) where
  memptyArrayUnsafe :: forall sh. IShX sh -> Mixed sh (Ranked n a)
  memptyArrayUnsafe i
    | Dict <- lemKnownReplicate (SNat @n)
    = coerce @(Mixed sh (Mixed (Replicate n Nothing) a)) @(Mixed sh (Ranked n a)) $
        memptyArrayUnsafe i

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

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


arithPromoteRanked :: forall n a b.
                      (forall sh. Mixed sh a -> Mixed sh b)
                   -> Ranked n a -> Ranked n b
arithPromoteRanked = coerce

arithPromoteRanked2 :: forall n a b c.
                       (forall sh. Mixed sh a -> Mixed sh b -> Mixed sh c)
                    -> Ranked n a -> Ranked n b -> Ranked n c
arithPromoteRanked2 = coerce

instance (NumElt a, PrimElt a, Num a) => Num (Ranked n a) where
  (+) = arithPromoteRanked2 (+)
  (-) = arithPromoteRanked2 (-)
  (*) = arithPromoteRanked2 (*)
  negate = arithPromoteRanked negate
  abs = arithPromoteRanked abs
  signum = arithPromoteRanked signum
  fromInteger = error "Ranked(fromInteger): Cannot implement fromInteger, use rreplicateScal"

instance (FloatElt a, NumElt a, PrimElt a, Num a) => Fractional (Ranked n a) where
  fromRational _ = error "Data.Array.Nested.fromRational: No singletons available, use explicit rreplicateScal"
  recip = arithPromoteRanked recip
  (/) = arithPromoteRanked2 (/)

instance (FloatElt a, NumElt a, PrimElt a, Num a) => Floating (Ranked n a) where
  pi = error "Data.Array.Nested.pi: No singletons available, use explicit rreplicateScal"
  exp = arithPromoteRanked exp
  log = arithPromoteRanked log
  sqrt = arithPromoteRanked sqrt
  (**) = arithPromoteRanked2 (**)
  logBase = arithPromoteRanked2 logBase
  sin = arithPromoteRanked sin
  cos = arithPromoteRanked cos
  tan = arithPromoteRanked tan
  asin = arithPromoteRanked asin
  acos = arithPromoteRanked acos
  atan = arithPromoteRanked atan
  sinh = arithPromoteRanked sinh
  cosh = arithPromoteRanked cosh
  tanh = arithPromoteRanked tanh
  asinh = arithPromoteRanked asinh
  acosh = arithPromoteRanked acosh
  atanh = arithPromoteRanked atanh
  log1p = arithPromoteRanked GHC.Float.log1p
  expm1 = arithPromoteRanked GHC.Float.expm1
  log1pexp = arithPromoteRanked GHC.Float.log1pexp
  log1mexp = arithPromoteRanked GHC.Float.log1mexp


remptyArray :: KnownElt a => Ranked 1 a
remptyArray = mtoRanked (memptyArray ZSX)

rshape :: Elt a => Ranked n a -> IShR n
rshape (Ranked arr) = shCvtXR' (mshape arr)

rrank :: Elt a => Ranked n a -> SNat n
rrank = shrRank . rshape

-- | The total number of elements in the array.
rsize :: Elt a => Ranked n a -> Int
rsize = shrSize . rshape

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

rindexPartial :: forall n m a. Elt a => Ranked (n + m) a -> IIxR n -> Ranked m a
rindexPartial (Ranked arr) idx =
  Ranked (mindexPartial @a @(Replicate n Nothing) @(Replicate m Nothing)
            (castWith (subst2 (lemReplicatePlusApp (ixrRank idx) (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. KnownElt a => IShR n -> (IIxR n -> a) -> Ranked n a
rgenerate sh f
  | sn@SNat <- shrRank sh
  , Dict <- lemKnownReplicate sn
  , Refl <- lemRankReplicate sn
  = Ranked (mgenerate (shCvtRX sh) (f . ixCvtXR))

-- | See the documentation of 'mlift'.
rlift :: forall n1 n2 a. Elt a
      => SNat n2
      -> (forall sh' b. Storable b => StaticShX sh' -> XArray (Replicate n1 Nothing ++ sh') b -> XArray (Replicate n2 Nothing ++ sh') b)
      -> Ranked n1 a -> Ranked n2 a
rlift sn2 f (Ranked arr) = Ranked (mlift (ssxFromSNat sn2) f arr)

-- | See the documentation of 'mlift2'.
rlift2 :: forall n1 n2 n3 a. Elt a
       => SNat n3
       -> (forall sh' b. Storable b => StaticShX sh' -> XArray (Replicate n1 Nothing ++ sh') b -> XArray (Replicate n2 Nothing ++ sh') b -> XArray (Replicate n3 Nothing ++ sh') b)
       -> Ranked n1 a -> Ranked n2 a -> Ranked n3 a
rlift2 sn3 f (Ranked arr1) (Ranked arr2) = Ranked (mlift2 (ssxFromSNat sn3) f arr1 arr2)

rsumOuter1P :: forall n a.
               (Storable a, NumElt a)
            => Ranked (n + 1) (Primitive a) -> Ranked n (Primitive a)
rsumOuter1P (Ranked arr)
  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
  = Ranked (msumOuter1P arr)

rsumOuter1 :: forall n a. (NumElt a, PrimElt a)
           => Ranked (n + 1) a -> Ranked n a
rsumOuter1 = rfromPrimitive . rsumOuter1P . rtoPrimitive

rsumAllPrim :: (PrimElt a, NumElt a) => Ranked n a -> a
rsumAllPrim (Ranked arr) = msumAllPrim arr

rtranspose :: forall n a. Elt a => PermR -> Ranked n a -> Ranked n a
rtranspose perm arr
  | sn@SNat <- rrank arr
  , Dict <- lemKnownReplicate sn
  , length perm <= fromIntegral (natVal (Proxy @n))
  = rlift sn
          (\ssh' -> X.transposeUntyped (natSing @n) ssh' perm)
          arr
  | otherwise
  = error "Data.Array.Nested.rtranspose: Permutation longer than rank of array"

rconcat :: forall n a. Elt a => NonEmpty (Ranked (n + 1) a) -> Ranked (n + 1) a
rconcat
  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
  = coerce mconcat

rappend :: forall n a. Elt a
        => Ranked (n + 1) a -> Ranked (n + 1) a -> Ranked (n + 1) a
rappend arr1 arr2
  | sn@SNat <- rrank arr1
  , Dict <- lemKnownReplicate sn
  , Refl <- lemReplicateSucc @(Nothing @Nat) @n
  = coerce (mappend @Nothing @Nothing @(Replicate n Nothing))
      arr1 arr2

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

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

rfromVector :: forall n a. PrimElt a => IShR n -> VS.Vector a -> Ranked n a
rfromVector sh v = rfromPrimitive (rfromVectorP sh v)

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

rtoVector :: PrimElt a => Ranked n a -> VS.Vector a
rtoVector = coerce mtoVector

rfromListOuter :: forall n a. Elt a => NonEmpty (Ranked n a) -> Ranked (n + 1) a
rfromListOuter l
  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
  = Ranked (mfromListOuter (coerce l :: NonEmpty (Mixed (Replicate n Nothing) a)))

rfromList1 :: Elt a => NonEmpty a -> Ranked 1 a
rfromList1 l = Ranked (mfromList1 l)

rfromList1Prim :: PrimElt a => [a] -> Ranked 1 a
rfromList1Prim l = Ranked (mfromList1Prim l)

rtoListOuter :: forall n a. Elt a => Ranked (n + 1) a -> [Ranked n a]
rtoListOuter (Ranked arr)
  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
  = coerce (mtoListOuter @a @Nothing @(Replicate n Nothing) arr)

rtoList1 :: Elt a => Ranked 1 a -> [a]
rtoList1 = map runScalar . rtoListOuter

rfromListPrim :: PrimElt a => [a] -> Ranked 1 a
rfromListPrim l =
  let ssh = SUnknown () :!% ZKX
      xarr = X.fromList1 ssh l
  in Ranked $ fromPrimitive $ M_Primitive (X.shape ssh xarr) xarr

rfromListPrimLinear :: PrimElt a => IShR n -> [a] -> Ranked n a
rfromListPrimLinear sh l =
  let M_Primitive _ xarr = toPrimitive (mfromListPrim l)
  in Ranked $ fromPrimitive $ M_Primitive (shCvtRX sh) (X.reshape (SUnknown () :!% ZKX) (shCvtRX sh) xarr)

rfromListLinear :: forall n a. Elt a => IShR n -> NonEmpty a -> Ranked n a
rfromListLinear sh l = rreshape sh (rfromList1 l)

rtoListLinear :: Elt a => Ranked n a -> [a]
rtoListLinear (Ranked arr) = mtoListLinear arr

rfromOrthotope :: PrimElt a => SNat n -> S.Array n a -> Ranked n a
rfromOrthotope sn arr
  | Refl <- lemRankReplicate sn
  = let xarr = XArray arr
    in Ranked (fromPrimitive (M_Primitive (X.shape (ssxFromSNat sn) xarr) xarr))

rtoOrthotope :: PrimElt a => Ranked n a -> S.Array n a
rtoOrthotope (rtoPrimitive -> Ranked (M_Primitive sh (XArray arr)))
  | Refl <- lemRankReplicate (shrRank $ shCvtXR' sh)
  = arr

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

rnest :: forall n m a. Elt a => SNat n -> Ranked (n + m) a -> Ranked n (Ranked m a)
rnest n arr
  | Refl <- lemReplicatePlusApp n (Proxy @m) (Proxy @(Nothing @Nat))
  = coerce (mnest (ssxFromSNat n) (coerce arr))

runNest :: forall n m a. Elt a => Ranked n (Ranked m a) -> Ranked (n + m) a
runNest rarr@(Ranked (M_Ranked (M_Nest _ arr)))
  | Refl <- lemReplicatePlusApp (rrank rarr) (Proxy @m) (Proxy @(Nothing @Nat))
  = Ranked arr

rzip :: Ranked n a -> Ranked n b -> Ranked n (a, b)
rzip = coerce mzip

runzip :: Ranked n (a, b) -> (Ranked n a, Ranked n b)
runzip = coerce munzip

rrerankP :: forall n1 n2 n a b. (Storable a, Storable b)
         => SNat n -> IShR n2
         -> (Ranked n1 (Primitive a) -> Ranked n2 (Primitive b))
         -> Ranked (n + n1) (Primitive a) -> Ranked (n + n2) (Primitive b)
rrerankP sn sh2 f (Ranked arr)
  | Refl <- lemReplicatePlusApp sn (Proxy @n1) (Proxy @(Nothing @Nat))
  , Refl <- lemReplicatePlusApp sn (Proxy @n2) (Proxy @(Nothing @Nat))
  = Ranked (mrerankP (ssxFromSNat sn) (shCvtRX sh2)
                     (\a -> let Ranked r = f (Ranked a) in r)
                     arr)

-- | If there is a zero-sized dimension in the @n@-prefix of the shape of the
-- input array, then there is no way to deduce the full shape of the output
-- array (more precisely, the @n2@ part): that could only come from calling
-- @f@, and there are no subarrays to call @f@ on. @orthotope@ errors out in
-- this case; we choose to fill the @n2@ part of the output shape with zeros.
--
-- For example, if:
--
-- @
-- arr :: Ranked 5 Int   -- of shape [3, 0, 4, 2, 21]
-- f :: Ranked 2 Int -> Ranked 3 Float
-- @
--
-- then:
--
-- @
-- rrerank _ _ _ f arr :: Ranked 5 Float
-- @
--
-- and this result will have shape @[3, 0, 4, 0, 0, 0]@. Note that the
-- "reranked" part (the last 3 entries) are zero; we don't know if @f@ intended
-- to return an array with shape all-0 here (it probably didn't), but there is
-- no better number to put here absent a subarray of the input to pass to @f@.
rrerank :: forall n1 n2 n a b. (PrimElt a, PrimElt b)
        => SNat n -> IShR n2
        -> (Ranked n1 a -> Ranked n2 b)
        -> Ranked (n + n1) a -> Ranked (n + n2) b
rrerank sn sh2 f (rtoPrimitive -> arr) =
  rfromPrimitive $ rrerankP sn sh2 (rtoPrimitive . f . rfromPrimitive) arr

rreplicate :: forall n m a. Elt a
           => IShR n -> Ranked m a -> Ranked (n + m) a
rreplicate sh (Ranked arr)
  | Refl <- lemReplicatePlusApp (shrRank sh) (Proxy @m) (Proxy @(Nothing @Nat))
  = Ranked (mreplicate (shCvtRX sh) arr)

rreplicateScalP :: forall n a. Storable a => IShR n -> a -> Ranked n (Primitive a)
rreplicateScalP sh x
  | Dict <- lemKnownReplicate (shrRank sh)
  = Ranked (mreplicateScalP (shCvtRX sh) x)

rreplicateScal :: forall n a. PrimElt a
               => IShR n -> a -> Ranked n a
rreplicateScal sh x = rfromPrimitive (rreplicateScalP sh x)

rslice :: forall n a. Elt a => Int -> Int -> Ranked (n + 1) a -> Ranked (n + 1) a
rslice i n arr
  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
  = rlift (rrank arr)
          (\_ -> X.sliceU i n)
          arr

rrev1 :: forall n a. Elt a => Ranked (n + 1) a -> Ranked (n + 1) a
rrev1 arr =
  rlift (rrank arr)
        (\(_ :: StaticShX sh') ->
          case lemReplicateSucc @(Nothing @Nat) @n of
            Refl -> X.rev1 @Nothing @(Replicate n Nothing ++ sh'))
        arr

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

rflatten :: Elt a => Ranked n a -> Ranked 1 a
rflatten (Ranked arr) = mtoRanked (mflatten arr)

riota :: (Enum a, PrimElt a, Elt a) => Int -> Ranked 1 a
riota n = TN.withSomeSNat (fromIntegral n) $ mtoRanked . miota

-- | Throws if the array is empty.
rminIndexPrim :: (PrimElt a, NumElt a) => Ranked n a -> IIxR n
rminIndexPrim rarr@(Ranked arr)
  | Refl <- lemRankReplicate (rrank (rtoPrimitive rarr))
  = ixCvtXR (mminIndexPrim arr)

-- | Throws if the array is empty.
rmaxIndexPrim :: (PrimElt a, NumElt a) => Ranked n a -> IIxR n
rmaxIndexPrim rarr@(Ranked arr)
  | Refl <- lemRankReplicate (rrank (rtoPrimitive rarr))
  = ixCvtXR (mmaxIndexPrim arr)

rdot1Inner :: forall n a. (PrimElt a, NumElt a) => Ranked (n + 1) a -> Ranked (n + 1) a -> Ranked n a
rdot1Inner arr1 arr2
  | SNat <- rrank arr1
  , Refl <- lemReplicatePlusApp (SNat @n) (Proxy @1) (Proxy @(Nothing @Nat))
  = coerce (mdot1Inner (Proxy @(Nothing @Nat))) arr1 arr2

-- | This has a temporary, suboptimal implementation in terms of 'mflatten'.
-- Prefer 'rdot1Inner' if applicable.
rdot :: (PrimElt a, NumElt a) => Ranked n a -> Ranked n a -> a
rdot = coerce mdot

rtoXArrayPrimP :: Ranked n (Primitive a) -> (IShR n, XArray (Replicate n Nothing) a)
rtoXArrayPrimP (Ranked arr) = first shCvtXR' (mtoXArrayPrimP arr)

rtoXArrayPrim :: PrimElt a => Ranked n a -> (IShR n, XArray (Replicate n Nothing) a)
rtoXArrayPrim (Ranked arr) = first shCvtXR' (mtoXArrayPrim arr)

rfromXArrayPrimP :: SNat n -> XArray (Replicate n Nothing) a -> Ranked n (Primitive a)
rfromXArrayPrimP sn arr = Ranked (mfromXArrayPrimP (ssxFromShape (X.shape (ssxFromSNat sn) arr)) arr)

rfromXArrayPrim :: PrimElt a => SNat n -> XArray (Replicate n Nothing) a -> Ranked n a
rfromXArrayPrim sn arr = Ranked (mfromXArrayPrim (ssxFromShape (X.shape (ssxFromSNat sn) arr)) arr)

rfromPrimitive :: PrimElt a => Ranked n (Primitive a) -> Ranked n a
rfromPrimitive (Ranked arr) = Ranked (fromPrimitive arr)

rtoPrimitive :: PrimElt a => Ranked n a -> Ranked n (Primitive a)
rtoPrimitive (Ranked arr) = Ranked (toPrimitive arr)

mtoRanked :: forall sh a. Elt a => Mixed sh a -> Ranked (Rank sh) a
mtoRanked arr
  | Refl <- lemAppNil @sh
  , Refl <- lemAppNil @(Replicate (Rank sh) (Nothing @Nat))
  , Refl <- lemRankReplicate (shxRank (mshape arr))
  = Ranked (mcast (ssxFromShape (mshape arr)) (convSh (mshape arr)) (Proxy @'[]) arr)
  where
    convSh :: IShX sh' -> IShX (Replicate (Rank sh') Nothing)
    convSh ZSX = ZSX
    convSh (smn :$% (sh :: IShX sh'T))
      | Refl <- lemReplicateSucc @(Nothing @Nat) @(Rank sh'T)
      = SUnknown (fromSMayNat' smn) :$% convSh sh

rtoMixed :: forall n a. Ranked n a -> Mixed (Replicate n Nothing) a
rtoMixed (Ranked arr) = arr