{-# LANGUAGE NoImplicitPrelude #-}

module Precursor.Structure.Foldable
  ( -- * Folds
    Foldable
  , fold
  , foldMap
  , foldr
  , foldr'
  , foldl
  , foldlLazy
  , foldl1
  , foldr1
  , toList
  , null
  , length
  , elem
  , minimum
  , maximum
  , sum
  , product
  , head
  , last
  , (!!)
    -- ** Specialized folds
  , concat
  , concatMap
  , and
  , or
  , any
  , all
  , maximumBy
  , minimumBy
  , foldr2
  , foldr3
    -- ** Searches
  , notElem
  , find
  ) where

import           Data.Foldable              hiding (foldl, foldl1, foldr1,
                                             maximum, maximumBy, minimum,
                                             minimumBy, product, sum)
import qualified Data.Foldable
import           Precursor.Algebra.Eq
import           Precursor.Algebra.Monoid
import           Precursor.Algebra.Ord
import           Precursor.Algebra.Ring
import           Precursor.Algebra.Semiring
import           Precursor.Coerce
import           Precursor.Control.Category
import           Precursor.Data.Bool
import           Precursor.Data.Maybe
import           Precursor.Function
import           Precursor.Numeric.Num

-- $setup
-- >>> import Precursor.Data.List
-- >>> import Precursor.Algebra.Ord
-- >>> import Test.QuickCheck

-- | Left-associative fold of a structure.
--
-- In the case of lists, 'foldl', when applied to a binary
-- operator, a starting value (typically the left-identity of the operator),
-- and a list, reduces the list using the binary operator, from left to
-- right:
--
-- > foldl f z [x1, x2, ..., xn] == (...((z `f` x1) `f` x2) `f`...) `f` xn
--
-- Note that to produce the outermost application of the operator the
-- entire input list must be traversed. This means that 'foldl' will
-- diverge if given an infinite list.
--
-- This version is strict, in contrast to Prelude's 'Data.Foldable.foldl'.
--
-- This ensures that each step of the fold is forced to weak head normal
-- form before being applied, avoiding the collection of thunks that would
-- otherwise occur. This is often what you want to strictly reduce a finite
-- list to a single, monolithic result (e.g. 'length').
--
-- For a lazy version, use 'foldlLazy'.
--
-- For a general 'Foldable' structure this should be semantically identical
-- to,
--
-- @foldl f z = 'List.foldl' f z . 'toList'@
--
foldl :: Foldable t => (b -> a -> b) -> b -> t a -> b
foldl = foldl'

-- | Left-associative fold of a structure but with lazy application of
-- the operator.
--
-- Note that if you want an efficient left-fold, you probably want to
-- use 'foldl' instead of 'foldlLazy'. The reason for this is that latter does
-- not force the "inner" results (e.g. @z `f` x1@ in the above example)
-- before applying them to the operator (e.g. to @(`f` x2)@). This results
-- in a thunk chain @O(n)@ elements long, which then must be evaluated from
-- the outside-in.
--
-- For a general 'Foldable' structure this should be semantically identical
-- to,
--
-- @foldlLazy f z = 'List.foldl' f z . 'toList'@
foldlLazy :: Foldable t => (b -> a -> b) -> b -> t a -> b
foldlLazy = Data.Foldable.foldl

-- | A variant of 'foldr' that has no base case,
-- and returns 'Nothing' for empty structures.
--
-- @'foldr1' f = 'List.foldr1' f . 'toList'@
foldr1 :: Foldable t => (a -> a -> a) -> t a -> Maybe a
foldr1 f = foldr (\x -> Just . maybe x (f x)) Nothing

-- | A variant of 'foldl' that has no base case,
-- and returns 'Nothing' for empty structures.
--
-- @'foldl1' f = 'List.foldl1' f . 'toList'@
foldl1 :: Foldable t => (a -> a -> a) -> t a -> Maybe a
foldl1 f = foldl g Nothing where
  g Nothing   x = Just x
  g (Just xs) x = Just (f xs x)

-- | A Scott-ish encoding of a zip. Possibly very inefficient.
newtype ScottZip a b =
  ScottZip (a -> (ScottZip a b -> b) -> b)

-- | Fold over two 'Foldable's at once.
--
-- prop> zip xs ys === foldr2 (\x y zs -> (x,y) : zs) [] xs ys
foldr2 :: (Foldable f, Foldable g) => (a -> b -> c -> c) -> c -> f a -> g b -> c
foldr2 c i xs = foldr f (const i) xs . ScottZip #. foldr g (\_ _ -> i) where
  g e2 r2 e1 r1 = c e1 e2 (coerce r1 r2)
  f e r (ScottZip x) = x e r

-- | Fold over three 'Foldable's at once.
--
-- prop> zip3 ws xs ys === foldr3 (\w x y zs -> (w,x,y) : zs) [] ws xs ys
foldr3 :: (Foldable f, Foldable g, Foldable h)
       => (a -> b -> c -> d -> d)
       -> d -> f a -> g b -> h c -> d
foldr3 c i xs ys =
  foldr f (const i) xs . ScottZip . foldr2 g (\_ _ -> i) ys where
    g e2 e3 r2 e1 r1 = c e1 e2 e3 (coerce r1 r2)
    f e r (ScottZip x) = x e r


newtype Max a = Max {getMax :: Maybe a}
newtype Min a = Min {getMin :: Maybe a}

instance Ord a => Monoid (Max a) where
  mempty = Max Nothing

  {-# INLINE mappend #-}
  m `mappend` Max Nothing = m
  Max Nothing `mappend` n = n
  (Max m@(Just x)) `mappend` (Max n@(Just y))
    | x >= y    = Max m
    | otherwise = Max n

instance Ord a => Monoid (Min a) where
  mempty = Min Nothing

  {-# INLINE mappend #-}
  m `mappend` Min Nothing = m
  Min Nothing `mappend` n = n
  (Min m@(Just x)) `mappend` (Min n@(Just y))
    | x <= y    = Min m
    | otherwise = Min n

-- | The largest element of a structure with respect to the given comparison
-- function. Returns 'Nothing' on an empty input.
--
-- prop> maximum (xs :: [Int]) === maximumBy compare xs
maximumBy :: Foldable t => (a -> a -> Ordering) -> t a -> Maybe a
maximumBy cmp = foldr1 max'
  where max' x y = case cmp x y of
                        LT -> y
                        _  -> x

-- | The least element of a structure with respect to the given comparison
-- function. Returns 'Nothing' on an empty input.
--
-- prop> minimum (xs :: [Int]) === minimumBy compare xs
minimumBy :: Foldable t => (a -> a -> Ordering) -> t a -> Maybe a
minimumBy cmp = foldr1 min'
  where min' x y = case cmp x y of
                        GT -> y
                        _  -> x

-- | The largest element of a structure. Returns 'Nothing' on empty
-- structures.
--
-- >>> maximum [1,2,3]
-- Just 3
-- >>> maximum []
-- Nothing
maximum :: (Ord a, Foldable t) => t a -> Maybe a
maximum =
    getMax . foldMap (Max #. (Just :: a -> Maybe a))

-- | The least element of a structure. Returns 'Nothing' on empty
-- structures.
--
-- >>> minimum [1,2,3]
-- Just 1
-- >>> minimum []
-- Nothing
minimum :: (Ord a, Foldable t) => t a -> Maybe a
minimum =
    getMin . foldMap (Min #. (Just :: a -> Maybe a))

-- | The 'sum' function computes the sum of the numbers of a structure.
--
-- prop> sum (xs :: [Integer]) === foldl (+) 0 xs
sum :: (Semiring a, Foldable t) => t a -> a
sum = foldl' (+) zero
{-# INLINE sum #-}

-- | The 'product' function computes the product of the numbers of a
-- structure.
--
-- prop> product (xs :: [Integer]) === foldl (*) 1 xs
product :: (Semiring a, Foldable t) => t a -> a
product = foldl' (*) one
{-# INLINE product #-}

-- | The first element of a structure, or 'Nothing' if it's empty.
--
-- >>> head [1,2,3]
-- Just 1
-- >>> head []
-- Nothing
--
-- prop> head xs === last (reverse xs)
head :: Foldable t => t a -> Maybe a
head = foldr1 const

-- | The last element of a structure, or 'Nothing' if it's empty.
--
-- >>> last [1,2,3]
-- Just 3
-- >>> last []
-- Nothing
--
-- prop> last xs === head (reverse xs)
last :: Foldable t => t a -> Maybe a
last = foldl1 (flip const)

infixl 9 !!
-- | Index (subscript) operator, starting from 0. Returns 'Nothing' for
-- out-of-range indices.
--
-- >>> [1,2,3] !! 0
-- Just 1
-- >>> [1,2,3] !! (-1)
-- Nothing
-- >>> [1,2,3] !! 3
-- Nothing
(!!) :: Foldable f => f a -> Int -> Maybe a
(!!) _ m | m < 0 = Nothing
(!!) xs m = foldr f b xs m where
  b = const Nothing
  f e _ 0 = Just e
  f _ a n = a (n-1)