-- |
-- Module      :  Data.List.Nth
-- Copyright   :  (c) OleksandrZhabenko 2019
-- License     :  MIT
--
-- Maintainer  :  olexandr543@yahoo.com
--
-- Various additional operations on lists.
--

module Data.List.Nth 
  (
    -- * Operations on lists to take a part of a list
       takeNth
       , takeWithFirst
       , dropNth
       , dropWithFirst
    -- * Operation to apply a function that creates an inner list to an element of the outer list   
       , bGroups
  ) where

-- | Function that takes a list containing elements consequently up to n occurencies of @p a = True@ in @map p [a]@
takeNth :: Int -> (a -> Bool) -> [a] -> [a]
takeNth n p xs | n <= 0 = []
               | otherwise = 
  takeWithFirst (not . p) xs ++ takeNth (n - 1) p (dropWithFirst p xs) 

-- | Function that drops a list containing elements consequently up to n occurencies of @p a = True@ in @map p [a]@
dropNth :: Int -> (a -> Bool) -> [a] -> [a]
dropNth n p xs | n <= 0 = xs
               | otherwise = 
  last . take n . tail . iterate (dropWithFirst p) $ xs
  
-- | Function that applies additional function 'f' to @a@ if @p a = True@
bGroups :: (a -> Bool) -> (a -> [a]) -> [a] -> [a]
bGroups p f = concatMap (\x -> if p x then f x else [x])

-- | Function to take elements of the list after the first occurence of @p a = True@ in @map p [a]@ excluding the element which results in the first occurance
dropWithFirst :: (a -> Bool) -> ([a] -> [a])
dropWithFirst p = fst . foldr f v
  where
    f x (ys,xs) = (if p x then ys else xs,x:xs)
    v = ([],[])

-- | Function to take elements of the list till the first occurence of @p a = True@ in 'map p [a]' including the element which results in the first occurance
takeWithFirst :: (a -> Bool) -> ([a] -> [a])
takeWithFirst p = fst . foldr f v
  where
    f x (ys,xs) = (if p x then x:ys else [x],x:xs)
    v = ([],[])