    -- |
    -- Module      : Foundation.Collection.Zippable
    -- License     : BSD-style
    -- Maintainer  : foundation
    -- Stability   : experimental
    -- Portability : portable
    --
    -- Common functions (e. g. 'zip', 'zipWith') that are useful for combining
    -- multiple collections.
    --
    {-# LANGUAGE FlexibleInstances     #-}
    {-# LANGUAGE MultiParamTypeClasses #-}
    {-# LANGUAGE StandaloneDeriving    #-}
    {-# LANGUAGE GeneralizedNewtypeDeriving #-}
    module Foundation.Collection.Zippable
        ( BoxedZippable(..)
        , Zippable(..)
        ) where

    import qualified Basement.UArray as UV
    import qualified Basement.BoxedArray as BA
    import qualified Basement.String as S
    import           Foundation.Collection.Element
    import           Foundation.Collection.Sequential
    import           Basement.Compat.Base
    import           Basement.Types.AsciiString(AsciiString(..))
    import qualified Prelude
    import           GHC.ST

    class Sequential col => Zippable col where

      -- | 'zipWith' generalises 'zip' by zipping with the function given as the
      --   first argument, instead of a tupling function. For example, @'zipWith' (+)@
      --   is applied to two collections to produce the collection of corresponding
      --   sums.
      zipWith :: (Sequential a, Sequential b)
              => (Element a -> Element b -> Element col)
              -> a -> b -> col
      zipWith f a b = go f (toList a, toList b)
        where go f' = maybe mempty (\(x, xs) -> uncurry2 f' x `cons` go f' xs) . uncons2

      -- | Like 'zipWith', but works with 3 collections.
      zipWith3 :: (Sequential a, Sequential b, Sequential c)
               => (Element a -> Element b -> Element c -> Element col)
               -> a -> b -> c -> col
      zipWith3 f a b c = go f (toList a, toList b, toList c)
        where go f' = maybe mempty (\(x, xs) -> uncurry3 f' x `cons` go f' xs) . uncons3

      -- | Like 'zipWith', but works with 4 collections.
      zipWith4 :: (Sequential a, Sequential b, Sequential c, Sequential d)
               => (Element a -> Element b -> Element c -> Element d -> Element col)
               -> a -> b -> c -> d -> col
      zipWith4 fn a b c d = go fn (toList a, toList b, toList c, toList d)
        where go f' = maybe mempty (\(x, xs) -> uncurry4 f' x `cons` go f' xs) . uncons4

      -- | Like 'zipWith', but works with 5 collections.
      zipWith5 :: (Sequential a, Sequential b, Sequential c, Sequential d, Sequential e)
               => (Element a -> Element b -> Element c -> Element d -> Element e -> Element col)
               -> a -> b -> c -> d -> e -> col
      zipWith5 fn a b c d e = go fn (toList a, toList b, toList c, toList d, toList e)
        where go f' = maybe mempty (\(x, xs) -> uncurry5 f' x `cons` go f' xs) . uncons5

      -- | Like 'zipWith', but works with 6 collections.
      zipWith6 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e
                  , Sequential f)
               => (Element a -> Element b -> Element c -> Element d -> Element e -> Element f -> Element col)
               -> a -> b -> c -> d -> e -> f -> col
      zipWith6 fn a b c d e f = go fn (toList a, toList b, toList c, toList d, toList e, toList f)
        where go f' = maybe mempty (\(x, xs) -> uncurry6 f' x `cons` go f' xs) . uncons6

      -- | Like 'zipWith', but works with 7 collections.
      zipWith7 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e
                  , Sequential f, Sequential g )
               => (Element a -> Element b -> Element c -> Element d -> Element e -> Element f -> Element g -> Element col)
               -> a -> b -> c -> d -> e -> f -> g -> col
      zipWith7 fn a b c d e f g = go fn (toList a, toList b, toList c, toList d, toList e, toList f, toList g)
        where go f' = maybe mempty (\(x, xs) -> uncurry7 f' x `cons` go f' xs) . uncons7

    instance Zippable [c]

    instance UV.PrimType ty => Zippable (UV.UArray ty) where
      zipWith f as bs = runST $ UV.builderBuild_ 64 $ go f (toList as) (toList bs)
        where
          go _  []       _        = return ()
          go _  _        []       = return ()
          go f' (a':as') (b':bs') = UV.builderAppend (f' a' b') >> go f' as' bs'

    instance Zippable (BA.Array ty) where
      zipWith f as bs = runST $ BA.builderBuild_ 64 $ go f (toList as) (toList bs)
        where
          go _  []       _        = return ()
          go _  _        []       = return ()
          go f' (a':as') (b':bs') = BA.builderAppend (f' a' b') >> go f' as' bs'

    instance Zippable S.String where
      zipWith f as bs = runST $ S.builderBuild_ 64 $ go f (toList as) (toList bs)
        where
          go _  []       _        = return ()
          go _  _        []       = return ()
          go f' (a':as') (b':bs') = S.builderAppend (f' a' b') >> go f' as' bs'

    deriving instance Zippable AsciiString

    class Zippable col => BoxedZippable col where

      -- | 'zip' takes two collections and returns a collections of corresponding
      --   pairs. If one input collection is short, excess elements of the longer
      --   collection are discarded.
      zip :: ( Sequential a, Sequential b
             , Element col ~ (Element a, Element b) )
          => a -> b -> col
      zip = zipWith (,)

      -- | Like 'zip', but works with 3 collections.
      zip3 :: ( Sequential a, Sequential b, Sequential c
              , Element col ~ (Element a, Element b, Element c) )
           => a -> b -> c -> col
      zip3 = zipWith3 (,,)

      -- | Like 'zip', but works with 4 collections.
      zip4 :: ( Sequential a, Sequential b, Sequential c, Sequential d
              , Element col ~ (Element a, Element b, Element c, Element d) )
           => a -> b -> c -> d -> col
      zip4 = zipWith4 (,,,)

      -- | Like 'zip', but works with 5 collections.
      zip5 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e
              , Element col ~ (Element a, Element b, Element c, Element d, Element e) )
           => a -> b -> c -> d -> e -> col
      zip5 = zipWith5 (,,,,)

      -- | Like 'zip', but works with 6 collections.
      zip6 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e, Sequential f
              , Element col ~ (Element a, Element b, Element c, Element d, Element e, Element f) )
           => a -> b -> c -> d -> e -> f -> col
      zip6 = zipWith6 (,,,,,)

      -- | Like 'zip', but works with 7 collections.
      zip7 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e, Sequential f, Sequential g
              , Element col ~ (Element a, Element b, Element c, Element d, Element e, Element f, Element g) )
           => a -> b -> c -> d -> e -> f -> g -> col
      zip7 = zipWith7 (,,,,,,)

      -- | 'unzip' transforms a collection of pairs into a collection of first
      --   components and a collection of second components.
      unzip :: (Sequential a, Sequential b, Element col ~ (Element a, Element b))
            => col -> (a, b)
      unzip = go . toList
        where go []          = (mempty, mempty)
              go ((a, b):xs) =
                  let (as, bs) = go xs
                  in (a `cons` as, b `cons` bs)

      -- | Like 'unzip', but works on a collection of 3-element tuples.
      unzip3 :: ( Sequential a, Sequential b, Sequential c
                , Element col ~ (Element a, Element b, Element c) )
             => col -> (a, b, c)
      unzip3 = go . toList
        where go [] = (mempty, mempty, mempty)
              go ((a, b, c):xs) =
                  let (as, bs, cs) = go xs
                  in (a `cons` as, b `cons` bs, c `cons` cs)

      -- | Like 'unzip', but works on a collection of 4-element tuples.
      unzip4 :: ( Sequential a, Sequential b, Sequential c, Sequential d
                , Element col ~ (Element a, Element b, Element c, Element d) )
             => col -> (a, b, c, d)
      unzip4 = go . toList
        where go [] = (mempty, mempty, mempty, mempty)
              go ((a, b, c, d):xs) =
                  let (as, bs, cs, ds) = go xs
                  in (a `cons` as, b `cons` bs, c `cons` cs, d `cons` ds)

      -- | Like 'unzip', but works on a collection of 5-element tuples.
      unzip5 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e
                , Element col ~ (Element a, Element b, Element c, Element d, Element e) )
             => col -> (a, b, c, d, e)
      unzip5 = go . toList
        where go [] = (mempty, mempty, mempty, mempty, mempty)
              go ((a, b, c, d, e):xs) =
                  let (as, bs, cs, ds, es) = go xs
                  in (a `cons` as, b `cons` bs, c `cons` cs, d `cons` ds, e `cons` es)

      -- | Like 'unzip', but works on a collection of 6-element tuples.
      unzip6 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e, Sequential f
                , Element col ~ (Element a, Element b, Element c, Element d, Element e, Element f) )
             => col -> (a, b, c, d, e, f)
      unzip6 = go . toList
        where go [] = (mempty, mempty, mempty, mempty, mempty, mempty)
              go ((a, b, c, d, e, f):xs) =
                  let (as, bs, cs, ds, es, fs) = go xs
                  in (a `cons` as, b `cons` bs, c `cons` cs, d `cons` ds, e `cons` es, f `cons` fs)

      -- | Like 'unzip', but works on a collection of 7-element tuples.
      unzip7 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e, Sequential f, Sequential g
                , Element col ~ (Element a, Element b, Element c, Element d, Element e, Element f, Element g) )
            => col -> (a, b, c, d, e, f, g)
      unzip7 = go . toList
        where go [] = (mempty, mempty, mempty, mempty, mempty, mempty, mempty)
              go ((a, b, c, d, e, f, g):xs) =
                  let (as, bs, cs, ds, es, fs, gs) = go xs
                  in (a `cons` as, b `cons` bs, c `cons` cs, d `cons` ds, e `cons` es, f `cons` fs, g `cons` gs)

    instance BoxedZippable [a]
    instance BoxedZippable (BA.Array ty)

    -- * Tuple helper functions

    uncons2 :: (Sequential a, Sequential b) => (a, b) -> Maybe ((Element a, Element b), (a, b))
    uncons2 xs = let (as, bs) = xs
                 in do (a', as') <- uncons as
                       (b', bs') <- uncons bs
                       return ((a', b'), (as', bs'))

    uncons3 :: (Sequential a, Sequential b, Sequential c)
            => (a, b, c)
            -> Maybe ((Element a, Element b, Element c), (a, b, c))
    uncons3 xs = let (as, bs, cs) = xs
                 in do (a', as') <- uncons as
                       (b', bs') <- uncons bs
                       (c', cs') <- uncons cs
                       return ((a', b', c'), (as', bs', cs'))

    uncons4 :: (Sequential a, Sequential b, Sequential c, Sequential d)
            => (a, b, c, d)
            -> Maybe ( (Element a, Element b, Element c, Element d)
                     , (a, b, c, d) )
    uncons4 xs = let (as, bs, cs, ds) = xs
                 in do (a', as') <- uncons as
                       (b', bs') <- uncons bs
                       (c', cs') <- uncons cs
                       (d', ds') <- uncons ds
                       return ((a', b', c', d'), (as', bs', cs', ds'))

    uncons5 :: (Sequential a, Sequential b, Sequential c, Sequential d, Sequential e)
            => (a, b, c, d, e)
            -> Maybe ( (Element a, Element b, Element c, Element d, Element e)
                     , (a, b, c, d, e) )
    uncons5 xs = let (as, bs, cs, ds, es) = xs
                 in do (a', as') <- uncons as
                       (b', bs') <- uncons bs
                       (c', cs') <- uncons cs
                       (d', ds') <- uncons ds
                       (e', es') <- uncons es
                       return ((a', b', c', d', e'), (as', bs', cs', ds', es'))

    uncons6 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e
               , Sequential f )
            => (a, b, c, d, e, f)
            -> Maybe ( (Element a, Element b, Element c, Element d, Element e, Element f)
                     , (a, b, c, d, e, f) )
    uncons6 xs = let (as, bs, cs, ds, es, fs) = xs
                 in do (a', as') <- uncons as
                       (b', bs') <- uncons bs
                       (c', cs') <- uncons cs
                       (d', ds') <- uncons ds
                       (e', es') <- uncons es
                       (f', fs') <- uncons fs
                       return ((a', b', c', d', e', f'), (as', bs', cs', ds', es', fs'))

    uncons7 :: ( Sequential a, Sequential b, Sequential c, Sequential d, Sequential e
               , Sequential f, Sequential g )
            => (a, b, c, d, e, f, g)
            -> Maybe ( (Element a, Element b, Element c, Element d, Element e, Element f
                     , Element g)
                     , (a, b, c, d, e, f, g) )
    uncons7 xs = let (as, bs, cs, ds, es, fs, gs) = xs
                 in do (a', as') <- uncons as
                       (b', bs') <- uncons bs
                       (c', cs') <- uncons cs
                       (d', ds') <- uncons ds
                       (e', es') <- uncons es
                       (f', fs') <- uncons fs
                       (g', gs') <- uncons gs
                       return ( (a', b', c', d', e', f', g')
                              , (as', bs', cs', ds', es', fs', gs') )

    uncurry2 :: (a -> b -> c) -> (a, b) -> c
    uncurry2 = Prelude.uncurry

    uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d
    uncurry3 fn (a, b, c) = fn a b c

    uncurry4 :: (a -> b -> c -> d -> g) -> (a, b, c, d) -> g
    uncurry4 fn (a, b, c, d) = fn a b c d

    uncurry5 :: (a -> b -> c -> d -> e -> f) -> (a, b, c, d, e) -> f
    uncurry5 fn (a, b, c, d, e) = fn a b c d e

    uncurry6 :: (a -> b -> c -> d -> e -> f -> g) -> (a, b, c, d, e, f) -> g
    uncurry6 fn (a, b, c, d, e, f) = fn a b c d e f

    uncurry7 :: (a -> b -> c -> d -> e -> f -> g -> h) -> (a, b, c, d, e, f, g) -> h
    uncurry7 fn (a, b, c, d, e, f, g) = fn a b c d e f g