-----------------------------------------------------------------------------
--
-- Module      :  Data.List.Util
-- Copyright   :  (c) 2012-16 Brian W Bush
-- License     :  MIT
--
-- Maintainer  :  Brian W Bush <consult@brianwbush.info>
-- Stability   :  Unstable
-- Portability :  Portable
--
-- | Miscellaneous functions for manipulating lists.
--
-----------------------------------------------------------------------------


{-# LANGUAGE Safe #-}


module Data.List.Util (
-- * Subsets
  hasSubset
, disjoint
, replaceByFst
, noDuplicates
, sameElements
, notDuplicatedIn
, deleteOn
, maybeList
-- * Permutations
, elemPermutation
-- * Grouping and sorting
, groupOn
, nubOn
, sortOn
, sortedGroups
, sortedGroupsOn
, regroup
, regroupBy
-- * Bounded enumerations.
, domain
-- * Permutations.
, powerSet
, powerSetPermutations
, uniquePowerSetPermutations
, suffixes
-- * Monads.
, nest
-- * Mapping.
, map2
, zipf
-- * Lookups.
, lookupWithDefault
, lookupWithDefaultFunction
-- * Padding, stripping, and pruning.
, replicateHead
, padHead
, padTail
, chunk
, stripHead
, discardEnds
, removeDuplicates
, removeDuplicatesBy
, extract
, rollback
-- * Splitting.
, splitAtMatches
, splitAround
) where


import Control.Arrow ((&&&))
import Control.Monad (filterM)
import Data.Function (on)
import Data.List ((\\), deleteBy, deleteFirstsBy, elemIndex, groupBy, intersect, nub, nubBy, permutations, sort, sortBy, tails)
import Data.List.Split (chunksOf, wordsBy)
import Data.Maybe (fromMaybe)


-- | Test for containment.
hasSubset :: Eq a
          => [a]  -- ^ The potential superset.
          -> [a]  -- ^ The potential subset
          -> Bool -- ^ Whether the first set has the second as a subset.
hasSubset x y = null $ y \\ x


-- | Test for disjointness.
disjoint :: Eq a
         => [a]  -- ^ The first set.
         -> [a]  -- ^ The second set.
         -> Bool -- ^ Whether the sets are disjoint.
disjoint x y = null $ x `intersect` y


-- | Replace the first occurrence in a lookup table.
replaceByFst :: Eq a
             => (a, b)   -- ^ The replacement key and value.
             -> [(a, b)] -- ^ The lookup table.
             -> [(a, b)] -- ^ The updated table.
replaceByFst x = (x :) . deleteBy ((==) `on` fst) x


-- | Test for duplicates.
noDuplicates :: Eq a
             => [a]  -- ^ The list.
             -> Bool -- ^ Whether the list does not contain duplicates.
noDuplicates x = length x == length (nub x)


-- | Test for same elements.
sameElements :: Ord a
             => [a]  -- ^ The first list.
             -> [a]  -- ^ The second list.
             -> Bool -- ^ Whether the lists have the same elements.
sameElements x y = sort x == sort y


-- | Test for membership of an element in a list.
notDuplicatedIn :: Eq b
                => (a -> b) -- ^ The equality test.
                -> a        -- ^ The value.
                -> [a]      -- ^ The list.
                -> Bool     -- ^ Whether the value is in the list.
notDuplicatedIn f x ys = f x `notElem` map f ys


-- | Delete using a function to extract values.
deleteOn :: Eq b => (a -> b) -> a -> [a] -> [a]
deleteOn f = deleteBy ((==) `on` f)


-- | Lift a non-empty list.
maybeList :: [a] -> Maybe [a]
maybeList [] = Nothing
maybeList xs = Just xs


-- | Find the permutation of one list relative to another.
elemPermutation :: Eq a => [a] -> [a] -> Maybe [Int]
elemPermutation = mapM . flip elemIndex


-- | Group using a function to extract values.
groupOn :: Eq b => (a -> b) -> [a] -> [[a]]
groupOn f = groupBy ((==) `on` f) 


-- | Nub using a function to extract values.
nubOn :: Eq b => (a -> b) -> [a] -> [a]
nubOn f = nubBy ((==) `on` f)


-- | Sort using a function to extract values.
sortOn :: Ord b => (a -> b) -> [a] -> [a]
sortOn f = sortBy (compare `on` f)


-- | Sort and group.
sortedGroups :: Ord b => [(b, c)] -> [(b, [c])]
sortedGroups = sortedGroupsOn fst snd 


-- | Sort and group using a function to extract values.
sortedGroupsOn :: Ord b => (a -> b) -> (a -> c) -> [a] -> [(b, [c])]
sortedGroupsOn f g =
  fmap ((f . head) &&& fmap g)
    . groupOn f
    . sortOn f


-- | Ordered list of values in a bounded enumeration.
domain :: (Bounded a, Enum a) => [a]
domain = [minBound..maxBound]


-- | Generate a power set.  The algorithm used here is from <http://evan-tech.livejournal.com/220036.html>.
powerSet :: [a] -> [[a]]
powerSet = filterM (const domain)


-- | Generate a power set's permutations.
powerSetPermutations :: [a] -> [[a]]
powerSetPermutations = concatMap permutations . powerSet


-- | Generate a power set's unique permutations.
uniquePowerSetPermutations :: Eq a => [a] -> [[a]]
uniquePowerSetPermutations = nub . powerSetPermutations


-- | A list of suffixes of another list.
suffixes :: [a] -> [[a]]
suffixes = init . tails


-- | Move a monad inside a list.
nest :: Monad m => [a] -> [m a]
nest = map return


-- | Map at the second level.
map2 :: (a -> b) -> [[a]] -> [[b]]
map2 = map . map


-- | The 'zipf' function applies a list of functions to corresponding elements in a list.
zipf :: [a -> b] -- ^ The list of functions to be applied
     -> [a]      -- ^ The list of elements to which the functions will be applied
     -> [b]      -- ^ The result of the function applications
zipf (f:fs) (x:xs) = f x : zipf fs xs
zipf _ _ = []


-- | Look up a value in an association list.
lookupWithDefault :: Eq a =>
     b        -- ^ The default value.
  -> a        -- ^ The key.
  -> [(a, b)] -- ^ The associations.
  -> b        -- ^ The value for the key, or the default if the key is not present.
lookupWithDefault y = lookupWithDefaultFunction (const y)


-- | Look up a value in an association list.
lookupWithDefaultFunction :: Eq a =>
     (a -> b) -- ^ The function for generating the default value.
  -> a        -- ^ The key.
  -> [(a, b)] -- ^ The associations.
  -> b        -- ^ The value for the key, or the default if the key is not present.
lookupWithDefaultFunction f x = fromMaybe (f x) . lookup x


-- | Relicating the head of a list.
replicateHead :: Int -- ^ Number of times to replicate the head.
              -> [a] -- ^ The list.
              -> [a] -- ^ The list with the replicas of the head.
replicateHead n =
  uncurry (++)
    . (replicate n . head &&& id)


-- | Pad the head of a list.
padHead ::
     Int -- ^ Length of result.
  -> a   -- ^ Item to use for padding.
  -> [a] -- ^ The list.
  -> [a] -- ^ The padded list.
padHead n y xs =
  replicate (n - length xs) y ++ xs


-- | Pad the tail of a list.
padTail ::
     Int -- ^ Length of the result.
  -> a   -- ^ Item to use for padding.
  -> [a] -- ^ The list.
  -> [a] -- ^ THe padded list.
padTail n y xs =
  xs ++ replicate (n - length xs) y


-- | Break a list into equally sized chunks.
chunk ::
     Int   -- ^ Length of the chunks.
  -> [a]   -- ^ The list.
  -> [[a]] -- ^ The chunked list.
chunk = chunksOf
{-# DEPRECATED chunk "Use 'Data.List.Split.chunksOf' instead." #-}


-- | Remove leading items from a list.
stripHead :: Eq a =>
     a   -- ^ The item to remove from the start of list.
  -> [a] -- ^ The list.
  -> [a] -- ^ The list without the leading items.
stripHead x y@(ye : ys)
  | x == ye   = stripHead x ys
  | otherwise = y
stripHead _ [] = []


-- | The 'discard' function removes the first and last elements from a list. 
discardEnds :: [a] -- ^ The list
            -> [a] -- ^ The interior of the list
discardEnds = init . tail


-- | Remove duplicates from a list, maintaining its order.
removeDuplicates :: Eq a => [a] -> [a]
removeDuplicates = removeDuplicatesBy (==)


-- | Remove duplicates from a list, maintaining its order.
removeDuplicatesBy ::
     (a -> a -> Bool) -- ^ Equality test function.
  -> [a]              -- ^ The list.
  -> [a]              -- ^ The list with duplicates removed.
removeDuplicatesBy equals x =
  let
    duplicates = nubBy equals (deleteFirstsBy equals x (nubBy equals x))
  in
    deleteFirstsBy equals (nubBy equals x) duplicates


-- | Extract one element from a list.
extract :: (a -> Bool) -> [a] -> ([a], Maybe a)
extract _ [] = ([], Nothing)
extract p x =
  let
    extract' a @ (_, _, Just _) = a
    extract' a @ (_, [], Nothing) = a
    extract' (y', ze : zs, Nothing)
      | p ze = (y', zs, Just ze)
      | otherwise = extract' (ze : y', zs, Nothing)
    (y, z, w) = extract' ([], x, Nothing)
  in
    (rollback y z, w)


-- | Reverse a first list and add it to a second one.
rollback ::
      [a] -- ^ The list to be reversed and prepended.
   -> [a] -- ^ The list to be appended.
   -> [a] -- ^ The resulting list
rollback = flip (foldl (flip (:)))


-- | Split up a list at every match of a particular item.
splitAtMatches :: Eq a =>
     a     -- ^ The separatrix.
  -> [a]   -- ^ The list.
  -> [[a]] -- ^ The split list.
splitAtMatches = wordsBy . (==)
{-# DEPRECATED splitAtMatches "Use 'Data.List.Split.wordsBy' instead." #-}


-- | Split a list around a particular element.
splitAround :: Int -> [a] -> ([a], a, [a])
splitAround n x =
  let
    (y, z) = splitAt (n + 1) x
  in
    (init y, last y, z)


-- | Sort and regroup elements of an association using its keys.
regroup :: Ord a => [(a, b)] -> [(a, [b])]
regroup = regroupBy id


-- | Sort and regroup elements of an association using its keys.
regroupBy :: Ord a =>
     ([b] -> c) -- ^ The function for summarizing the second element of pairs.
  -> [(a, b)]   -- ^ The associations.
  -> [(a, c)]   -- ^ The regrouped associations.
regroupBy f =
  let
    sorter (x, _) (y, _) = compare x y
    grouper x y = fst x == fst y
    crusher ys = (fst $ head ys, f $ map snd ys)
  in
    map crusher . groupBy grouper . sortBy sorter