content: {-# LANGUAGE DataKinds            #-}
  {-# LANGUAGE FlexibleContexts     #-}
  {-# LANGUAGE GADTs                #-}
  {-# LANGUAGE LambdaCase           #-}
  {-# LANGUAGE ScopedTypeVariables  #-}
  {-# LANGUAGE TupleSections        #-}
  {-# LANGUAGE TypeApplications     #-}
  {-# LANGUAGE TypeOperators        #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  module Data.Type.Vector.Util where
  
  import           Control.DeepSeq
  import           Data.Bifunctor
  import           Data.Distributive
  import           Data.Foldable
  import           Data.Monoid
  import           Data.Type.Combinator
  import           Data.Type.Conjunction
  import           Data.Type.Fin
  import           Data.Type.Nat
  import           Data.Type.Vector
  import           Type.Class.Known
  import           Type.Class.Witness
  import           Type.Family.Nat
  import           Type.Family.Nat.Util
  
  instance (Known Nat n, Distributive f) => Distributive (VecT n f) where
      distribute xs = vgen_ $ \i -> distribute $ index i <$> xs
  
  instance NFData (f a) => NFData (VecT n f a) where
      rnf = \case
        ØV      -> ()
        x :* xs -> x `deepseq` xs `deepseq` ()
  
  splitVec
      :: Nat n
      -> VecT (n + m) f a
      -> (VecT n f a, VecT m f a)
  splitVec = \case
      Z_   -> (ØV,)
      S_ n -> \case
        x :* xs -> first (x :*) (splitVec n xs)
  {-# INLINE splitVec #-}
  
  zipVecs
      :: (Traversable g, Applicative g, Known Nat n)
      => (VecT m g a -> b)
      -> VecT m g (VecT n g a)
      -> VecT n g b
  zipVecs = liftVec
  {-# INLINE zipVecs #-}
  
  liftVec
      :: (Applicative f, Traversable g)
      => (VecT m g a -> b)
      -> VecT m g (f a)
      -> f b
  liftVec f xs = f <$> sequenceA xs
  {-# INLINE liftVec #-}
  
  zipVecsD
      :: (Distributive g, Known Nat n)
      => (VecT m g a -> b)
      -> VecT m g (VecT n g a)
      -> VecT n g b
  zipVecsD = liftVecD
  {-# INLINE zipVecsD #-}
  
  liftVecD
      :: (Distributive f, Distributive g)
      => (VecT m g a -> b)
      -> VecT m g (f a)
      -> f b
  liftVecD f xs = f <$> distribute xs
  {-# INLINE liftVecD #-}
  
  curryV
      :: (VecT ('S n) f a -> b)
      -> f a
      -> VecT n f a
      -> b
  curryV f x xs = f (x :* xs)
  {-# INLINE curryV #-}
  
  curryV'
      :: (Vec ('S n) a -> b)
      -> a
      -> Vec n a
      -> b
  curryV' f x xs = f (I x :* xs)
  {-# INLINE curryV' #-}
  
  curryV2'
      :: (Vec N2 a -> b)
      -> a -> a -> b
  curryV2' f x y = f (I x :* I y :* ØV)
  {-# INLINE curryV2' #-}
  
  curryV3'
      :: (Vec N3 a -> b)
      -> a -> a -> a -> b
  curryV3' f x y z = f (I x :* I y :* I z :* ØV)
  {-# INLINE curryV3' #-}
  
  
  uncurryV
      :: (f a -> VecT n f a -> b)
      -> VecT ('S n) f a
      -> b
  uncurryV f = \case
      x :* xs -> f x xs
  {-# INLINE uncurryV #-}
  
  uncurryV'
      :: (a -> Vec n a -> b)
      -> Vec ('S n) a
      -> b
  uncurryV' f = \case
      I x :* xs -> f x xs
  {-# INLINE uncurryV' #-}
  
  append'
      :: VecT n f a
      -> VecT m f a
      -> VecT (n + m) f a
  append' = \case
      ØV -> id
      x :* xs -> (x :*) . append' xs
  {-# INLINE append' #-}
  
  vecFunc
      :: Known Nat n
      => (a -> Vec n b)
      -> Vec n (a -> b)
  vecFunc f = vgen_ (\i -> I $ index' i . f)
  {-# INLINE vecFunc #-}
  
  unVecFunc
      :: Vec n (a -> b)
      -> a
      -> Vec n b
  unVecFunc xs x = fmap ($ x) xs
  {-# INLINE unVecFunc #-}
  
  vgenA
      :: Applicative g
      => Nat n
      -> (Fin n -> g (f a))
      -> g (VecT n f a)
  vgenA = \case
    Z_   -> \_ -> pure ØV
    S_ n -> \f -> (:*) <$> f FZ <*> vgenA n (f . FS)
  {-# INLINE vgenA #-}
  
  uniformVec
      :: Eq (f a)
      => VecT ('S m) f a
      -> Maybe (f a)
  uniformVec = \case
      x :* xs | getAll (vfoldMap (All . (== x)) xs) -> Just x
              | otherwise                           -> Nothing
  {-# INLINE uniformVec #-}
  
  uncons'
      :: VecT ('S n) f a
      -> (f a, VecT n f a)
  uncons' (x :* xs) = (x, xs)
  
  len
      :: VecT n f a
      -> Nat n
  len = \case
      ØV      -> Z_
      _ :* xs -> S_ (len xs)
  
  select
      :: forall n f a. ()
      => VecT ('S n) f a
      -> VecT ('S n) (f :&: VecT n f) a
  select xs0 = go Z_ ØV (len xs0) xs0
    where
      go  :: forall m o. ()
          => Nat m
          -> VecT m f a
          -> Nat ('S o)
          -> VecT ('S o) f a
          -> VecT ('S o) (f :&: VecT (m + o) f) a
      go m xs = \case
        S_ Z_ -> \case
          y :* ØV -> (y :&: xs) :* ØV
                       \\ addZero m
        S_ o@(S_ p) -> \case
          y :* ys -> (y :&: (xs `append'` ys)) :* go (S_ m) (y :* xs) o ys
                       \\ succAssoc m p
  
  sumV
      :: Num a
      => Vec f a
      -> a
  sumV = \case
      ØV          -> 0
      xs@(_ :* _) -> foldl1' (+) xs
  {-# INLINE sumV #-}
  
  foldl1'
      :: (a -> a -> a)
      -> Vec ('S n) a
      -> a
  foldl1' f = \case
      I x :* ØV          -> x
      I x :* ys@(_ :* _) -> foldl' f x ys
  {-# INLINE foldl1' #-}
  
  reverse'
      :: forall n f a. ()
      => VecT n f a
      -> VecT n f a
  reverse' v0 = go Z_ ØV (known \\ v0) v0
    where
      go  :: forall m o. ()
          => Nat m
          -> VecT m f a
          -> Nat o
          -> VecT o f a
          -> VecT (m + o) f a
      go m xs = \case
        Z_   -> \case
          ØV      -> xs    \\ addZero m
        S_ o -> \case
          y :* ys -> go (S_ m) (y :* xs) o ys
              \\ succAssoc m o
  {-# INLINE reverse' #-}
  
  unzip'
      :: Vec n (a, b)
      -> (Vec n a, Vec n b)
  unzip' = \case
      ØV              -> (ØV, ØV)
      I (x,y) :* xsys -> let (xs, ys) = unzip' xsys
                         in  (I x :* xs, I y :* ys)
  