--------------------------------------------------------------------------------
-- |
-- Module      :  Data.List.Util
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--------------------------------------------------------------------------------
module Data.List.Util where

import           Data.Bifunctor
import           Data.Ext
import qualified Data.Foldable as F
import qualified Data.List as List
import           Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NonEmpty
import           Data.List.Zipper (allNonEmptyNexts, extractNext)
import qualified Data.List.Zipper as Zipper
import           Data.Maybe
import           Data.Ord (comparing)

--------------------------------------------------------------------------------

-- | Given an input list, computes all lists in which just one element is missing.
--
-- >>> mapM_ print $ leaveOutOne [1..5]
-- (1,[2,3,4,5])
-- (2,[1,3,4,5])
-- (3,[1,2,4,5])
-- (4,[1,2,3,5])
-- (5,[1,2,3,4])
-- >>> leaveOutOne []
-- []
-- >>> leaveOutOne [1]
-- [(1,[])]
leaveOutOne    :: [a] -> [(a,[a])]
leaveOutOne :: [a] -> [(a, [a])]
leaveOutOne [a]
xs = (Zipper a -> [a]) -> (a, Zipper a) -> (a, [a])
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Zipper a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList ((a, Zipper a) -> (a, [a]))
-> (Zipper a -> (a, Zipper a)) -> Zipper a -> (a, [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (a, Zipper a) -> (a, Zipper a)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (a, Zipper a) -> (a, Zipper a))
-> (Zipper a -> Maybe (a, Zipper a)) -> Zipper a -> (a, Zipper a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Zipper a -> Maybe (a, Zipper a)
forall a. Zipper a -> Maybe (a, Zipper a)
extractNext
              (Zipper a -> (a, [a])) -> [Zipper a] -> [(a, [a])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Zipper a -> [Zipper a]
forall a. Zipper a -> [Zipper a]
allNonEmptyNexts ([a] -> Zipper a
forall a. [a] -> Zipper a
Zipper.fromList [a]
xs)


--------------------------------------------------------------------------------
-- * Improved functions for minima and maxima

-- | Safe variant of Prelude.minimum.
--
-- >>> minimum1 [] :: Maybe ()
-- Nothing
-- >>> minimum1 [1,2,3]
-- Just 1
minimum1 :: Ord a => [a] -> Maybe a
minimum1 :: [a] -> Maybe a
minimum1 = (a -> a -> Ordering) -> [a] -> Maybe a
forall a. (a -> a -> Ordering) -> [a] -> Maybe a
minimum1By a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

-- | Safe variant of Prelude.maximum.
--
-- >>> maximum1 [] :: Maybe ()
-- Nothing
-- >>> maximum1 [1,2,3]
-- Just 3
maximum1 :: Ord a => [a] -> Maybe a
maximum1 :: [a] -> Maybe a
maximum1 = (a -> a -> Ordering) -> [a] -> Maybe a
forall a. (a -> a -> Ordering) -> [a] -> Maybe a
minimum1By ((a -> a -> Ordering) -> a -> a -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare)

-- | Total variant of Data.List.minimumBy.
--
-- >>> minimum1By (comparing abs) [] :: Maybe Int
-- Nothing
-- >>> minimum1By (comparing abs) [1,-2,3]
-- Just 1
minimum1By     :: (a -> a -> Ordering) -> [a] -> Maybe a
minimum1By :: (a -> a -> Ordering) -> [a] -> Maybe a
minimum1By a -> a -> Ordering
cmp = \case
  [] -> Maybe a
forall a. Maybe a
Nothing
  [a]
xs -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ (a -> a -> Ordering) -> [a] -> a
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
List.minimumBy a -> a -> Ordering
cmp [a]
xs

-- | Computes all minima by comparing some property.
--
-- >>> minimaOn (max 2) [1,2,3,4,5,-1]
-- [-1,2,1]
minimaOn   :: Ord b => (a -> b) -> [a] -> [a]
minimaOn :: (a -> b) -> [a] -> [a]
minimaOn a -> b
f = (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
minimaBy ((a -> b) -> a -> a -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing a -> b
f)

-- | Computes all minima.
--
-- >>> minimaBy (comparing abs) [1,2,3,2,1,-1]
-- [-1,1,1]
minimaBy     :: (a -> a -> Ordering) -> [a] -> [a]
minimaBy :: (a -> a -> Ordering) -> [a] -> [a]
minimaBy a -> a -> Ordering
cmp = \case
  []     -> []
  (a
x:[a]
xs) -> NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NonEmpty.toList (NonEmpty a -> [a]) -> NonEmpty a -> [a]
forall a b. (a -> b) -> a -> b
$ (NonEmpty a -> a -> NonEmpty a) -> NonEmpty a -> [a] -> NonEmpty a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (\mins :: NonEmpty a
mins@(a
m:|[a]
_) a
y -> case a
m a -> a -> Ordering
`cmp` a
y of
                                                               Ordering
LT -> NonEmpty a
mins
                                                               Ordering
EQ -> a
y a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
NonEmpty.<| NonEmpty a
mins
                                                               Ordering
GT -> a
ya -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:|[]
                                          ) (a
xa -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:|[]) [a]
xs

-- | extracts all minima from the list. The result consists of the
-- list of minima, and all remaining points. Both lists are returned
-- in the order in which they occur in the input.
--
-- >>> extractMinimaBy compare [1,2,3,0,1,2,3,0,1,2,0,2]
-- [0,0,0] :+ [2,3,1,2,3,1,2,1,2]
extractMinimaBy     :: (a -> a -> Ordering) -> [a] -> [a] :+ [a]
extractMinimaBy :: (a -> a -> Ordering) -> [a] -> [a] :+ [a]
extractMinimaBy a -> a -> Ordering
cmp = \case
  []     -> [] [a] -> [a] -> [a] :+ [a]
forall core extra. core -> extra -> core :+ extra
:+ []
  (a
x:[a]
xs) -> (NonEmpty a -> [a]) -> (NonEmpty a :+ [a]) -> [a] :+ [a]
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NonEmpty.toList ((NonEmpty a :+ [a]) -> [a] :+ [a])
-> (NonEmpty a :+ [a]) -> [a] :+ [a]
forall a b. (a -> b) -> a -> b
$ (a -> (NonEmpty a :+ [a]) -> NonEmpty a :+ [a])
-> (NonEmpty a :+ [a]) -> [a] -> NonEmpty a :+ [a]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
y (mins :: NonEmpty a
mins@(a
m:|[a]
_) :+ [a]
rest) ->
                                             case a
m a -> a -> Ordering
`cmp` a
y of
                                               Ordering
LT -> NonEmpty a
mins NonEmpty a -> [a] -> NonEmpty a :+ [a]
forall core extra. core -> extra -> core :+ extra
:+ a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
rest
                                               Ordering
EQ -> (a
y a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
NonEmpty.<| NonEmpty a
mins) NonEmpty a -> [a] -> NonEmpty a :+ [a]
forall core extra. core -> extra -> core :+ extra
:+ [a]
rest
                                               Ordering
GT -> (a
ya -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:|[]) NonEmpty a -> [a] -> NonEmpty a :+ [a]
forall core extra. core -> extra -> core :+ extra
:+ NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NonEmpty.toList NonEmpty a
mins [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a]
rest
                                          ) ((a
xa -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:|[]) NonEmpty a -> [a] -> NonEmpty a :+ [a]
forall core extra. core -> extra -> core :+ extra
:+ []) [a]
xs
  -- TODO: This is actually a good scenario for testing how much slower :+ is compared
  -- to doing nothing. i..e compare minimaBy and extractMinimaBy
  -- note that I'm using foldr here, and foldl' before


--------------------------------------------------------------------------------
-- * Partitioning and Grouping

-- | Given a function f, partitions the list into three lists
-- (lts,eqs,gts) such that:
--
-- - f x == LT for all x in lts
-- - f x == EQ for all x in eqs
-- - f x == gt for all x in gts
--
-- >>> partition3 (compare 4) [0,1,2,2,3,4,5,5,6,6,7,7,7,7,7,8]
-- ([5,5,6,6,7,7,7,7,7,8],[4],[0,1,2,2,3])
--
partition3   :: Foldable f => (a -> Ordering) -> f a -> ([a],[a],[a])
partition3 :: (a -> Ordering) -> f a -> ([a], [a], [a])
partition3 a -> Ordering
f = (a -> ([a], [a], [a]) -> ([a], [a], [a]))
-> ([a], [a], [a]) -> f a -> ([a], [a], [a])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> ([a], [a], [a]) -> ([a], [a], [a])
g ([],[],[])
  where
    g :: a -> ([a], [a], [a]) -> ([a], [a], [a])
g a
x ([a]
lts,[a]
eqs,[a]
gts) = case a -> Ordering
f a
x of
                          Ordering
LT -> (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
lts,   [a]
eqs,  [a]
gts)
                          Ordering
EQ -> (  [a]
lts, a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
eqs,  [a]
gts)
                          Ordering
GT -> (  [a]
lts,   [a]
eqs,a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
gts)

-- | A version of groupBy that uses the given Ordering to group
-- consecutive Equal items
--
-- >>> groupBy' compare [0,1,2,2,3,4,5,5,6,6,7,7,7,7,7,8]
-- [0 :| [],1 :| [],2 :| [2],3 :| [],4 :| [],5 :| [5],6 :| [6],7 :| [7,7,7,7],8 :| []]
groupBy'     :: (a -> a -> Ordering) -> [a] -> [NonEmpty a]
groupBy' :: (a -> a -> Ordering) -> [a] -> [NonEmpty a]
groupBy' a -> a -> Ordering
cmp = [a] -> [NonEmpty a]
go
  where
    go :: [a] -> [NonEmpty a]
go = \case
      []       -> []
      (a
x:[a]
xs)   -> let ([a]
pref,[a]
rest) = (a -> Bool) -> [a] -> ([a], [a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
List.span (\a
y -> a
x a -> a -> Ordering
`cmp` a
y Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ) [a]
xs
                  in (a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
pref) NonEmpty a -> [NonEmpty a] -> [NonEmpty a]
forall a. a -> [a] -> [a]
: [a] -> [NonEmpty a]
go [a]
rest