content: {-# LANGUAGE ConstraintKinds        #-}
  {-# LANGUAGE DataKinds              #-}
  {-# LANGUAGE DeriveFoldable         #-}
  {-# LANGUAGE DeriveFunctor          #-}
  {-# LANGUAGE DeriveTraversable      #-}
  {-# LANGUAGE FlexibleContexts       #-}
  {-# LANGUAGE FlexibleInstances      #-}
  {-# LANGUAGE FunctionalDependencies #-}
  {-# LANGUAGE GADTs                  #-}
  {-# LANGUAGE InstanceSigs           #-}
  {-# LANGUAGE KindSignatures         #-}
  {-# LANGUAGE LambdaCase             #-}
  {-# LANGUAGE MultiParamTypeClasses  #-}
  {-# LANGUAGE PolyKinds              #-}
  {-# LANGUAGE ScopedTypeVariables    #-}
  {-# LANGUAGE StandaloneDeriving     #-}
  {-# LANGUAGE TypeApplications       #-}
  {-# LANGUAGE TypeFamilies           #-}
  {-# LANGUAGE TypeFamilyDependencies #-}
  {-# LANGUAGE TypeInType             #-}
  {-# LANGUAGE TypeOperators          #-}
  {-# LANGUAGE UndecidableInstances   #-}
  
  module Data.Nested
    ( Vec(..)
    , Nesting(..)
    , Nesting1(..), nesting1Every
    , Nested
    , genNested, genNestedA
    , indexNested, indexNested'
    , transpose
    , transpose'
    , gmul'
    , diagNV
    , joinNested
    , mapNVecSlices
    , nIxRows
    , vGen, vIFoldMap, itraverseNested
    , liftNested
    , unScalar, unNest, unVector
    , sumRowsNested
    ) where
  
  import           Control.Applicative
  import           Control.DeepSeq
  import           Data.Distributive
  import           Data.Foldable
  import           Data.Kind
  import           Data.List.Util
  import           Data.Monoid
  import           Data.Singletons
  import           Data.Singletons.Prelude.List hiding (Length, Reverse, (%:++), sReverse)
  import           Data.Type.Combinator
  import           Data.Type.Combinator.Util
  import           Data.Type.Index
  import           Data.Type.Length                    as TCL
  import           Data.Type.Product as TCP hiding     (toList)
  import           Data.Type.Sing
  import           Data.Type.SnocProd
  import           Data.Type.Uniform
  import           TensorOps.NatKind
  import           Type.Class.Witness
  import           Type.Family.List
  import           Type.Family.List.Util
  import qualified Data.Singletons.TypeLits            as GT
  import qualified Data.Type.Nat                       as TCN
  import qualified Data.Type.Vector                    as TCV
  import qualified Data.Type.Vector.Util               as TCV
  import qualified Data.Vector.Sized                   as VS
  
  data Uncons :: (k -> Type -> Type) -> k -> Type -> Type where
      UNil  :: Uncons v (FromNat 0) a
      UCons :: !(Sing n) -> !a -> !(v n a) -> Uncons v (Succ n) a
  
  class NatKind k => Vec (v :: k -> Type -> Type) where
      vHead   :: p j -> v (Succ j) a -> a
      vTail   :: v (Succ j) a -> v j a
      vGenA   :: Applicative f => Sing j -> (IndexN k j -> f a) -> f (v j a)
      vIndex  :: IndexN k j -> v j a -> a
      vUncons :: Sing j -> v j a -> Uncons v j a
      vEmpty  :: v (FromNat 0) a
      vCons   :: a -> v j a -> v (Succ j) a
      vITraverse
          :: Applicative f
          => (IndexN k j -> a -> f b)
          -> v j a
          -> f (v j b)
  
  vGen
      :: Vec (v :: k -> Type -> Type)
      => Sing j
      -> (IndexN k j -> a)
      -> v j a
  vGen s f = getI $ vGenA s (I . f)
  {-# INLINE vGen #-}
  
  vIFoldMap
      :: (Monoid m, Vec v)
      => (IndexN k j -> a -> m)
      -> v j a
      -> m
  vIFoldMap f = getConst . vITraverse (\i -> Const . f i)
  
  instance Vec (Flip2 VS.VectorT I) where
      vHead _ = getI . VS.head . getFlip2
      {-# INLINE vHead #-}
      vTail = Flip2 . VS.tail . getFlip2
      {-# INLINE vTail #-}
      vGenA = \case
        GT.SNat -> fmap Flip2 . VS.generateA . (fmap I .)
      {-# INLINE vGenA #-}
      vIndex i = (VS.!! i) . getFlip2
      {-# INLINE vIndex #-}
      vUncons = \case
        GT.SNat -> \case
          Flip2 xs -> case VS.uncons xs of
            VS.VNil           -> UNil
            VS.VCons (I y) ys -> UCons sing y (Flip2 ys)
      {-# INLINE vUncons #-}
      vEmpty = Flip2 VS.empty
      {-# INLINE vEmpty #-}
      vCons x (Flip2 xs) = Flip2 (VS.cons (I x) xs)
      {-# INLINE vCons #-}
      vITraverse f (Flip2 xs) = Flip2 <$> VS.itraverse (\i (I x) -> I <$> f i x) xs
      {-# INLINE vITraverse #-}
  
  instance Vec (Flip2 TCV.VecT I) where
      vHead _ = getI . TCV.head' . getFlip2
      {-# INLINE vHead #-}
      vTail = Flip2 . TCV.tail' . getFlip2
      {-# INLINE vTail #-}
      vGenA = \case
        SN n -> \f -> Flip2 <$> TCV.vgenA n (fmap I . f)
      {-# INLINE vGenA #-}
      vIndex i = TCV.index' i . getFlip2
      {-# INLINE vIndex #-}
      vUncons = \case
        SN TCN.Z_ -> \case
          Flip2 TCV.ØV -> UNil
        SN (TCN.S_ n) -> \case
          Flip2 (I x TCV.:* xs) -> UCons (SN n) x (Flip2 xs)
      {-# INLINE vUncons #-}
      vEmpty = Flip2 TCV.ØV
      {-# INLINE vEmpty #-}
      vCons x (Flip2 xs) = Flip2 (I x TCV.:* xs)
      {-# INLINE vCons #-}
      vITraverse f (Flip2 xs) = Flip2 <$> TCV.itraverse (\i (I x) -> I <$> f i x) xs
      {-# INLINE vITraverse #-}
  
  class Nesting (w :: k -> Type) (c :: j -> Constraint) (v :: k -> j -> j) where
      nesting :: w i -> c a :- c (v i a)
  
  class Nesting1 (w :: k -> Type) (c :: j -> Constraint) (v :: k -> j) where
      nesting1 :: w a -> Wit (c (v a))
  
  instance Nesting w NFData (Flip2 VS.VectorT I) where
      nesting _ = Sub Wit
      {-# INLINE nesting #-}
  instance Nesting w Show (Flip2 VS.VectorT I) where
      nesting _ = Sub Wit
      {-# INLINE nesting #-}
  instance Functor f      => Nesting1 w    Functor      (Flip2 VS.VectorT f) where
      nesting1 _ = Wit
      {-# INLINE nesting1 #-}
  instance Applicative f  => Nesting1 Sing Applicative  (Flip2 VS.VectorT f) where
      nesting1 GT.SNat = Wit
      {-# INLINE nesting1 #-}
  instance Foldable f     => Nesting1 w    Foldable     (Flip2 VS.VectorT f) where
      nesting1 _ = Wit
      {-# INLINE nesting1 #-}
  instance Traversable f  => Nesting1 w    Traversable  (Flip2 VS.VectorT f) where
      nesting1 _ = Wit
      {-# INLINE nesting1 #-}
  instance Distributive f => Nesting1 Sing Distributive (Flip2 VS.VectorT f) where
      nesting1 GT.SNat = Wit
      {-# INLINE nesting1 #-}
  
  
  instance Nesting w NFData (Flip2 TCV.VecT I) where
      nesting _ = Sub Wit
      {-# INLINE nesting #-}
  instance Nesting w Show (Flip2 TCV.VecT I) where
      nesting _ = Sub Wit
      {-# INLINE nesting #-}
  instance Functor f      => Nesting1 w    Functor      (Flip2 TCV.VecT f) where
      nesting1 _ = Wit
      {-# INLINE nesting1 #-}
  instance Applicative f  => Nesting1 Sing Applicative  (Flip2 TCV.VecT f) where
      nesting1 (SN n) = Wit \\ n
      {-# INLINE nesting1 #-}
  instance Foldable f     => Nesting1 w    Foldable     (Flip2 TCV.VecT f) where
      nesting1 _ = Wit
      {-# INLINE nesting1 #-}
  instance Traversable f  => Nesting1 w    Traversable  (Flip2 TCV.VecT f) where
      nesting1 _ = Wit
      {-# INLINE nesting1 #-}
  instance Distributive f => Nesting1 Sing Distributive (Flip2 TCV.VecT f) where
      nesting1 (SN n) = Wit \\ n
      {-# INLINE nesting1 #-}
  
  nesting1Every
      :: forall p w c v as. Nesting1 w c v
      => p v
      -> Prod w as
      -> Wit (Every c (v <$> as))
  nesting1Every p = \case
      Ø   -> Wit
      (w :: w a) :< (ws :: Prod w as')
          -> Wit  \\ (nesting1 w :: Wit (c (v a)))
                  \\ (nesting1Every p ws :: Wit (Every c (v <$> as')))
  {-# INLINE nesting1Every #-}
  
  data Nested :: (k -> Type -> Type) -> [k] -> Type -> Type where
      NØ :: !a                     -> Nested v '[]       a
      NS :: !(v j (Nested v js a)) -> Nested v (j ': js) a
  
  instance (NFData a, Nesting Proxy NFData v) => NFData (Nested v js a) where
      rnf = \case
        NØ x  -> deepseq x  ()
        NS (xs :: v j (Nested v ks a))
              -> deepseq xs ()
                   \\ (nesting Proxy :: NFData (Nested v ks a) :- NFData (v j (Nested v ks a)))
      {-# INLINE rnf #-}
  
  instance (Num a, Applicative (Nested v js)) => Num (Nested v js a) where
      (+)         = liftA2 (+)
      {-# INLINE (+) #-}
      (*)         = liftA2 (*)
      {-# INLINE (*) #-}
      (-)         = liftA2 (-)
      {-# INLINE (-) #-}
      negate      = fmap negate
      {-# INLINE negate #-}
      abs         = fmap abs
      {-# INLINE abs #-}
      signum      = fmap signum
      {-# INLINE signum #-}
      fromInteger = pure . fromInteger
      {-# INLINE fromInteger #-}
  
  instance Nesting1 Proxy Functor v => Functor (Nested v js) where
      fmap f = \case
        NØ x  -> NØ (f x)
        NS (xs :: v j (Nested v ks a))
              -> NS $ (fmap.fmap) f xs
                        \\ (nesting1 Proxy :: Wit (Functor (v j)))
  
  instance (SingI js, Nesting1 Sing Applicative v, Nesting1 Proxy Functor v) => Applicative (Nested v js) where
      pure :: forall a. a -> Nested v js a
      pure x = go sing
        where
          go  :: Sing ks
              -> Nested v ks a
          go = \case
            SNil     -> NØ x
            (s :: Sing k) `SCons` ss -> NS (pure (go ss))
                      \\ (nesting1 s :: Wit (Applicative (v k)))
      {-# INLINE pure #-}
      (<*>) :: forall a b. Nested v js (a -> b) -> Nested v js a -> Nested v js b
      (<*>) = go sing
        where
          go  :: Sing ks
              -> Nested v ks (a -> b)
              -> Nested v ks a
              -> Nested v ks b
          go = \case
            SNil -> \case
              NØ f -> \case
                NØ x -> NØ (f x)
            (s :: Sing k) `SCons` ss -> \case
              NS fs -> \case
                NS xs -> NS $ liftA2 (go ss) fs xs
                                \\ (nesting1 s :: Wit (Applicative (v k)))
      {-# INLINE (<*>) #-}
  
  instance Nesting1 Proxy Foldable v => Foldable (Nested v js) where
      foldMap f = \case
        NØ x  -> f x
        NS (xs :: v j (Nested v ks a))
              -> (foldMap . foldMap) f xs
                   \\ (nesting1 Proxy :: Wit (Foldable (v j)))
  
  instance (Nesting1 Proxy Functor v, Nesting1 Proxy Foldable v, Nesting1 Proxy Traversable v) => Traversable (Nested v js) where
      traverse f = \case
        NØ x  -> NØ <$> f x
        NS (xs :: v j (Nested v ks a))
              -> NS <$> (traverse . traverse) f xs
                   \\ (nesting1 Proxy :: Wit (Traversable (v j)))
  
  instance (Vec v, SingI js, Nesting1 Proxy Functor v) => Distributive (Nested v js) where
      distribute
          :: forall f a. Functor f
          => f (Nested v js a)
          -> Nested v js (f a)
      distribute xs = genNested sing $ \i -> indexNested i <$> xs
      {-# INLINE distribute #-}
      -- distribute = flip go sing
      --   where
      --     go  :: f (Nested v ks a)
      --         -> Sing ks
      --         -> Nested v ks (f a)
      --     go xs = \case
      --       SNil         -> NØ $ unScalar <$> xs
      --       s `SCons` ss -> NS . vGen s $ \i ->
      --                         go (fmap (indexNested' (i :< Ø)) xs) ss
  
  -- TODO: rewrite rules?  lazy pattern matches?
  nHead
      :: forall v p j js a. Vec v
      => p j
      -> Nested v (Succ j ': js) a
      -> Nested v js a
  nHead p = \case
    NS xs -> vHead p xs
  {-# INLINE nHead #-}
  
  nTail
      :: Vec v
      => Nested v (Succ j ': js) a
      -> Nested v (j ': js) a
  nTail = \case
    NS xs -> NS $ vTail xs
  {-# INLINE nTail #-}
  
  unScalar
      :: Nested v '[] a
      -> a
  unScalar = \case
    NØ x -> x
  {-# INLINE unScalar #-}
  
  unNest
      :: Nested v (j ': js) a
      -> v j (Nested v js a)
  unNest = \case
    NS xs -> xs
  {-# INLINE unNest #-}
  
  unVector
      :: Functor (v j)
      => Nested v '[j] a
      -> v j a
  unVector = \case
      NS xs -> unScalar <$> xs
  {-# INLINE unVector #-}
  
  nVector
      :: Functor (v j)
      => v j a
      -> Nested v '[j] a
  nVector = NS . fmap NØ
  {-# INLINE nVector #-}
  
  genNested
      :: Vec (v :: k -> Type -> Type)
      => Sing ns
      -> (Prod (IndexN k) ns -> a)
      -> Nested v ns a
  genNested s f = getI $ genNestedA s (I . f)
  {-# INLINE genNested #-}
  
  genNestedA
      :: (Vec (v :: k -> Type -> Type), Applicative f)
      => Sing ns
      -> (Prod (IndexN k) ns -> f a)
      -> f (Nested v ns a)
  genNestedA = \case
      SNil         -> \f -> NØ <$> f Ø
      s `SCons` ss -> \f -> NS <$> vGenA s (\i -> genNestedA ss (f . (i :<)))
  
  indexNested
      :: Vec (v :: k -> Type -> Type)
      => Prod (IndexN k) ns
      -> Nested v ns a
      -> a
  indexNested = \case
      Ø -> \case
        NØ x  -> x
      i :< is -> \case
        NS xs -> indexNested is (vIndex i xs)
  -- indexNested i = unScalar . indexNested' i
  --                   \\ appendNil (prodLength i)
  
  indexNested'
      :: Vec (v :: k -> Type -> Type)
      => Prod (IndexN k) ms
      -> Nested v (ms ++ ns) a
      -> Nested v ns a
  indexNested' = \case
      Ø -> id
      i :< is -> \case
        NS xs -> indexNested' is (vIndex i xs)
  
  joinNested
      :: forall v ns ms a. Nesting1 Proxy Functor v
      => Nested v ns (Nested v ms a)
      -> Nested v (ns ++ ms) a
  joinNested = \case
      NØ x  -> x
      NS (xs :: v j (Nested v js (Nested v ms a))) ->
        NS $ fmap joinNested xs
          \\ (nesting1 Proxy :: Wit (Functor (v j)))
  
  mapNVecSlices
      :: forall v ns ms a b. Nesting1 Proxy Functor v
      => (Nested v ms a -> b)
      -> Length ns
      -> Nested v (ns ++ ms) a
      -> Nested v ns b
  mapNVecSlices f = \case
      LZ -> NØ . f
      LS l -> \case
        NS (xs :: v j (Nested v js a)) ->
          NS $ mapNVecSlices f l <$> xs
            \\ (nesting1 Proxy :: Wit (Functor (v j)))
  
  diagNV'
      :: forall v n ns a. (Vec v, Nesting1 Proxy Functor v)
      => Sing n
      -> Nested v (n ': n ': ns) a
      -> Nested v (n ': ns) a
  diagNV' s = \case
    NS (xs :: v n (Nested v (n ': ns) a)) -> case vUncons s xs of
      UNil          -> NS vEmpty
      UCons (s' :: Sing n')
            (y :: Nested v (n ': ns) a)
            (ys :: v n' (Nested v (n ': ns) a)) ->
        case nesting1 Proxy :: Wit (Functor (v n')) of
          Wit -> case diagNV' s' (NS (nTail <$> ys)) of
            NS zs -> NS $ vCons (nHead s' y) zs
  
  diagNV
      :: (Vec v, Nesting1 Proxy Functor v)
      => Sing n
      -> Uniform n ms
      -> Nested v (n ': n ': ms) a
      -> Nested v '[n] a
  diagNV s = \case
      UØ   -> diagNV' s
      US u -> diagNV s u . diagNV' s
  
  itraverseNested
      :: forall k (v :: k -> Type -> Type) (ns :: [k]) a b f. (Applicative f, Vec v)
      => (Prod (IndexN k) ns -> a -> f b)
      -> Nested v ns a
      -> f (Nested v ns b)
  itraverseNested f = \case
      NØ x  -> NØ <$> f Ø x
      NS xs -> NS <$> vITraverse (\i -> itraverseNested (\is -> f (i :< is))) xs
  
  gmul'
      :: forall ms os ns v a.
       ( Nesting1 Proxy Functor      v
       , Nesting1 Sing  Applicative  v
       , SingI ns
       , Num a
       , Vec v
       )
      => Length ms
      -> Length os
      -> Length ns
      -> Nested v (ms         ++ os) a
      -> Nested v (Reverse os ++ ns) a
      -> Nested v (ms         ++ ns) a
  gmul' lM _ _ x y = joinNested $ mapNVecSlices f lM x
    where
      f   :: Nested v os a
          -> Nested v ns a
      f = getSum
        . getConst
        -- . itraverseNested (\i x' -> Const . Sum $ fmap (x' *) (indexNested' (prodReverse' i) y))
        . itraverseNested (\i x' -> Const . Sum $ fmap (x' *) (indexNested' (TCP.reverse' i) y))
  {-# INLINE gmul' #-}
  
  -- | Transpose by iteratively sequencing/distributing layers
  transpose
      :: forall v os a.
       ( Nesting1 Proxy Functor      v
       , Nesting1 Proxy Foldable     v
       , Nesting1 Proxy Traversable  v
       , Nesting1 Sing  Distributive v
       )
      => Sing os
      -> Nested v os a
      -> Nested v (Reverse os) a
  transpose s = transposeHelp (snocProd (singProd s))
  {-# INLINE transpose #-}
  
  transposeHelp
      :: forall v os a.
       ( Nesting1 Proxy Functor      v
       , Nesting1 Proxy Foldable     v
       , Nesting1 Proxy Traversable  v
       , Nesting1 Sing  Distributive v
       )
      => SnocProd Sing os
      -> Nested v os a
      -> Nested v (Reverse os) a
  transposeHelp = \case
      ØS -> \case
        NØ x -> NØ x
      (sOs' :: SnocProd Sing os') :& (sO :: Sing o) ->
        (\\ (nesting1 Proxy :: Wit (Functor      (v o)))) $
        (\\ (nesting1 sO    :: Wit (Distributive (v o)))) $ \x ->
          let lOs'  :: Length os'
              lOs'  = snocProdLength sOs'
              x' :: Nested v os' (v o a)
              x' = mapNVecSlices unVector lOs' x
                     \\ appendSnoc lOs' sO
              xT :: Nested v (Reverse os') (v o a)
              xT = transposeHelp sOs' x'
              y :: v o (Nested v (Reverse os') a)
              y = distribute xT
              y' :: Nested v '[o] (Nested v (Reverse os') a)
              y' = nVector y
          in  joinNested y'
                \\ snocReverse lOs' sO
  
  -- | Transpose by populating a new 'Nested' from scratch
  transpose'
      :: Vec v
      => Length os
      -> Sing (Reverse os)
      -> Nested v os a
      -> Nested v (Reverse os) a
  transpose' l sR x = genNested sR $ \i -> indexNested (TCP.reverse' i) x
                        \\ reverseReverse l
  {-# INLINE transpose' #-}
  
  nIxRows
      :: forall k (v :: k -> Type -> Type) ns ms a b f. (Nesting1 Proxy Functor v, Applicative f, Vec v)
      => Length ns
      -> (Prod (IndexN k) ns -> Nested v ms a -> f b)
      -> Nested v (ns ++ ms) a
      -> f (Nested v ns b)
  nIxRows = \case
      LZ   -> \f -> fmap NØ . f Ø
      LS l -> \f -> \case
        NS (xs :: v j (Nested v js a)) ->
          fmap NS . vITraverse (\i -> nIxRows l (\is ys -> f (i :< is) ys)) $ xs
  
  liftNested
      :: Distributive (Nested v ns)
      => (TCV.Vec n a -> a)
      -> TCV.Vec n (Nested v ns a)
      -> Nested v ns a
  liftNested = TCV.liftVecD
  {-# INLINE liftNested #-}
  
  sumRowsNested
      :: forall v n ns a.
       ( Foldable (v n)
       , Num a
       , SingI ns
       , Nesting1 Proxy Functor v
       , Nesting1 Sing Applicative v
       )
      => Nested v (n ': ns) a
      -> Nested v ns a
  sumRowsNested (NS xs) = sum' (toList xs)
  