-----------------------------------------------------------------------------
-- Copyright 2019, Advise-Me project team. This file is distributed under 
-- the terms of the Apache License 2.0. For more information, see the files
-- "LICENSE.txt" and "NOTICE.txt", which are included in the distribution.
-----------------------------------------------------------------------------
-- |
-- Maintainer  :  bastiaan.heeren@ou.nl
-- Stability   :  provisional
-- Portability :  portable (depends on ghc)
--
-----------------------------------------------------------------------------

module Util.List
   ( safeHead, returnOnJustM, groupings, cartProd
   , selections, safeSingleton, safeLast, appLast, modifyAt
   , index
   , concatMapM
   , groupBy'
   , orderBy
   ) where

import Data.Function (on)
import Data.List (elemIndex, sortBy, groupBy)
import Control.Arrow (second, (&&&))

safeHead :: [a] -> Maybe a
safeHead [] = Nothing
safeHead (x:_) = Just x

safeLast :: [a] -> Maybe a
safeLast = foldl (\_ a -> Just a) Nothing

appLast :: (a -> a) -> [a] -> [a]
appLast _ [] = []
appLast f [a] = [f a]
appLast f (x:xs) = x : appLast f xs

modifyAt :: Int -> (a -> a) -> [a] -> [a]
modifyAt _ _ [] = []
modifyAt n f (x:xs)
  | n == 0 = f x : xs
  | otherwise = x : modifyAt (n-1) f xs

safeSingleton :: [a] -> Maybe a
safeSingleton [a] = Just a
safeSingleton _ = Nothing

cartProd :: [a] -> [b] -> [(a,b)]
cartProd xs ys = do
  x <- xs
  y <- ys
  return (x,y)

selections :: [a] -> [(a,[a])]
selections [] = []
selections (x:xs) = (x,xs) : map (second (x:)) (selections xs)

returnOnJustM :: (Monad m, Show a) => (a -> m (Maybe b)) -> [a] -> m (Maybe (b, [a]))
returnOnJustM _ [] = return Nothing
returnOnJustM f input =
  returnAndRetain f [] input
  where
    returnAndRetain _ _ [] = return Nothing
    returnAndRetain f old new@(x:xs) = do
      fx <- f x
      case fx of
        Nothing -> returnAndRetain f (old ++ [x]) xs
        Just b  -> return $ Just (b, old ++ new)

groupings :: [a] -> [b] -> [[(a,b)]]
groupings [] _ = [[]]
groupings _ [] = [[]]
groupings (a:as) bs = concatMap (\(x,xs) -> map ((a,x):) $ groupings as xs) $ selections bs

-- | Find index of an element in a list.
index :: Eq a => [a] -> a -> Maybe Int
index list = flip elemIndex list

-- | Monadic 'concatMap'.
concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM f xs = concat <$> mapM f xs

-- | Alternative to 'Data.List.groupBy', where only adjacent elements are
-- compared - so the predicate does not assume transitivity.
groupBy' :: (a -> a -> Bool) -> [a] -> [[a]]
groupBy' _ [] = []
groupBy' _ [x] = [[x]]
groupBy' predicate (x:y:xs)
   | predicate x y = let sameGroup:diffGroups = groupBy' predicate (y:xs)
                in  (x:sameGroup):diffGroups
   | otherwise = [x] : groupBy' predicate (y:xs)



-- | Alternative to 'Data.List.groupBy', in which the list is sorted and grouped
-- on a certain attribute, and each group is labelled with said attribute.
orderBy :: (Ord b) => (a->b) -> [a] -> [(b,[a])]
orderBy attribute
   = map (fst . head &&& map snd)
   . groupBy ((==) `on` fst) 
   . sortBy (compare `on` fst)
   . map (attribute &&& id)