module NumericPrelude.List where

import Data.List (unfoldr, genericReplicate)
import NumericPrelude.Condition (toMaybe)

{- * Slice lists -}


{-| keep every k-th value from the list

   Since these implementations check for the end of lists,
   they may fail in fixpoint computations on infinite lists. -}
sieve, sieve', sieve'', sieve''' :: Int -> [a] -> [a]
sieve k =
   unfoldr (\xs -> toMaybe (not (null xs)) (head xs, drop k xs))

sieve' k = map head . sliceVert k

-- this one works only on finite lists
sieve'' k x = map (x!!) [0,k..(length x-1)]

sieve''' k = map head . takeWhile (not . null) . iterate (drop k)

{- sliceHoriz is faster than sliceHoriz' but consumes slightly more memory
   (although it needs no swapping) -}
sliceHoriz, sliceHoriz' :: Int -> [a] -> [[a]]
sliceHoriz n =
   map (sieve n) . take n . iterate (drop 1)

sliceHoriz' n =
   foldr (\x ys -> let y = last ys in takeMatch ys ((x:y):ys)) (replicate n [])


sliceVert, sliceVert' :: Int -> [a] -> [[a]]
sliceVert n =
   map (take n) . takeWhile (not . null) . iterate (drop n)
      {- takeWhile must be performed before (map take)
         in order to handle (n==0) correctly -}

sliceVert' n =
   unfoldr (\x -> toMaybe (not (null x)) (splitAt n x))


{- * Use lists as counters -}

{- | Make a list as long as another one -}
{-# INLINE takeMatch #-}
takeMatch :: [b] -> [a] -> [a]
takeMatch = flip (zipWith const)

{-# INLINE dropMatch #-}
dropMatch :: [b] -> [a] -> [a]
dropMatch (_:xs) (_:ys) = dropMatch xs ys
dropMatch _ ys = ys

{-# INLINE splitAtMatch #-}
splitAtMatch :: [b] -> [a] -> ([a],[a])
splitAtMatch (_:ns) (x:xs) =
   let (as,bs) = splitAtMatch ns xs
   in  (x:as,bs)
splitAtMatch _ [] = ([],[])
splitAtMatch [] xs = ([],xs)

{-# INLINE replicateMatch #-}
replicateMatch :: [a] -> b -> [b]
replicateMatch xs y =
   takeMatch xs (repeat y)

{- |
Compare the length of two lists over different types.
For finite lists it is equivalent to (compare (length xs) (length ys))
but more efficient.
-}
{-# INLINE compareLength #-}
compareLength :: [a] -> [b] -> Ordering
compareLength (_:xs) (_:ys) = compareLength xs ys
compareLength []     []     = EQ
compareLength (_:_)  []     = GT
compareLength []     (_:_)  = LT





{- * Zip lists -}

{- | zip two lists using an arbitrary function, the shorter list is padded -}
{-# INLINE zipWithPad #-}
zipWithPad :: a               {-^ padding value -}
           -> (a -> a -> b)   {-^ function applied to corresponding elements of the lists -}
           -> [a]
           -> [a]
           -> [b]
zipWithPad z f =
   let aux l []          = map (\x -> f x z) l
       aux [] l          = map (\y -> f z y) l
       aux (x:xs) (y:ys) = f x y : aux xs ys
   in  aux

{-# INLINE zipWithOverlap #-}
zipWithOverlap :: (a -> c) -> (b -> c) -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithOverlap fa fb fab =
   let aux (x:xs) (y:ys) = fab x y : aux xs ys
       aux xs [] = map fa xs
       aux [] ys = map fb ys
   in  aux

{- | Zip two lists which must be of the same length.
    This is checked only lazily, that is unequal lengths are detected only
    if the list is evaluated completely.
    But it is more strict than @zipWithPad undefined f@
    since the latter one may succeed on unequal length list if @f@ is lazy. -}
zipWithMatch
   :: (a -> b -> c)   {-^ function applied to corresponding elements of the lists -}
   -> [a]
   -> [b]
   -> [c]
zipWithMatch f =
   let aux (x:xs) (y:ys) = f x y : aux xs ys
       aux []     []     = []
       aux _      _      = error "zipWith: lists must have the same length"
   in  aux

{-# INLINE zipNeighborsWith #-}
zipNeighborsWith :: (a -> a -> a) -> [a] -> [a]
zipNeighborsWith f xs = zipWith f xs (drop 1 xs)



{- * Lists of lists -}

{- |
Transform

@
 [[00,01,02,...],          [[00],
  [10,11,12,...],   -->     [10,01],
  [20,21,22,...],           [20,11,02],
  ...]                      ...]
@

With @concat . shear@ you can perform a Cantor diagonalization,
that is an enumeration of all elements of the sub-lists
where each element is reachable within a finite number of steps.
It is also useful for polynomial multiplication (convolution).
-}
shear :: [[a]] -> [[a]]
shear xs@(_:_) =
   let (y:ys,zs) = unzip (map (splitAt 1) xs)
       zipConc (a:as) (b:bs) = (a++b) : zipConc as bs
       zipConc [] bs = bs
       zipConc as [] = as
   in  y : zipConc ys (shear (dropWhileRev null zs))
              {- Dropping trailing empty lists is necessary,
                 otherwise finite lists are filled with empty lists. -}
shear [] = []

{- |
Transform

@
 [[00,01,02,...],          [[00],
  [10,11,12,...],   -->     [01,10],
  [20,21,22,...],           [02,11,20],
  ...]                      ...]
@

It's like 'shear' but the order of elements in the sub list is reversed.
Its implementation seems to be more efficient than that of 'shear'.
If the order does not matter, better choose 'shearTranspose'.
-}
shearTranspose :: [[a]] -> [[a]]
shearTranspose =
   let -- zipCons is like zipWith (:) keep lists which are too long
       zipCons (x:xs) (y:ys) = (x:y) : zipCons xs ys
       zipCons [] ys = ys
       zipCons xs [] = map (:[]) xs
       aux (x:xs) yss = [x] : zipCons xs yss
       aux [] yss = []:yss
   in  foldr aux []


{- |
Operate on each combination of elements of the first and the second list.
In contrast to the list instance of 'Monad.liftM2'
in holds the results in a list of lists.
It holds
@concat (outerProduct f xs ys)  ==  liftM2 f xs ys@
-}
outerProduct :: (a -> b -> c) -> [a] -> [b] -> [[c]]
outerProduct f xs ys = map (flip map ys . f) xs



{- * Various -}

{-# INLINE partitionMaybe #-}
partitionMaybe :: (a -> Maybe b) -> [a] -> ([b], [a])
partitionMaybe f =
   foldr (\x ~(y,z) -> case f x of
             Just x' -> (x' : y, z)
             Nothing -> (y, x : z)) ([],[])

{- |
It holds @splitLast xs == (init xs, last xs)@,
but 'splitLast' is more efficient
if the last element is accessed after the initial ones,
because it avoids memoizing list.
-}
{-# INLINE splitLast #-}
splitLast :: [a] -> ([a], a)
splitLast [] = error "splitLast: empty list"
splitLast [x] = ([], x)
splitLast (x:xs) =
   let (xs', lastx) = splitLast xs in (x:xs', lastx)

propSplitLast :: Eq a => [a] -> Bool
propSplitLast xs =
   splitLast xs  ==  (init xs, last xs)

{- |
Remove the longest suffix of elements satisfying p.
In contrast to 'reverse . dropWhile p . reverse'
this works for infinite lists, too.
-}
{-# INLINE dropWhileRev #-}
dropWhileRev :: (a -> Bool) -> [a] -> [a]
dropWhileRev p =
   foldr (\x xs -> if p x && null xs then [] else x:xs) []


{- |
Apply a function to the last element of a list.
If the list is empty, nothing changes.
-}
{-# INLINE mapLast #-}
mapLast :: (a -> a) -> [a] -> [a]
mapLast f =
   let recurse []     = [] -- behaviour as needed in powerBasis
              -- error "mapLast: empty list"
       recurse (x:[]) = f x : []
       recurse (x:xs) = x : recurse xs
   in  recurse

{-# INLINE padLeft #-}
padLeft :: a -> Int -> [a] -> [a]
padLeft  c n xs = replicate (n - length xs) c ++ xs


{-# INLINE padRight #-}
padRight :: a -> Int -> [a] -> [a]
padRight c n xs = xs ++ replicate (n - length xs) c


{- |
@reduceRepeated@ is an auxiliary function that,
for an associative operation @op@,
computes the same value as

  @reduceRepeated op a0 a n = foldr op a0 (genericReplicate n a)@

but applies "op" O(log n) times and works for large n.
-}

{-# INLINE reduceRepeated #-}
{-# INLINE reduceRepeatedSlow #-}
reduceRepeated, reduceRepeatedSlow ::
   (a -> a -> a) -> a -> a -> Integer -> a
reduceRepeated _  a0 _ 0 = a0
reduceRepeated op a0 a n =
   if even n
     then reduceRepeated op a0 (op a a) (div n 2)
     else reduceRepeated op (op a0 a) (op a a) (div n 2)

reduceRepeatedSlow op a0 a n =
   foldr op a0 (genericReplicate n a)


{- |
For an associative operation @op@ this computes
   @iterateAssoc op a = iterate (op a) a@
but it is even faster than @map (reduceRepeated op a a) [0..]@
since it shares temporary results.

The idea is:
From the list @map (reduceRepeated op a a) [0,(2*n)..]@
we compute the list @map (reduceRepeated op a a) [0,n..]@,
and iterate that until @n==1@.
-}
iterateAssoc, iterateLeaky :: (a -> a -> a) -> a -> [a]
iterateAssoc op a =
   foldr (\pow xs -> pow : concatMap (\x -> [x, op x pow]) xs)
         undefined (iterate (\x -> op x x) a)

{- |
This is equal to 'iterateAssoc'.
The idea is the following:
The list we search is the fixpoint of the function:
"Square all elements of the list,
then spread it and fill the holes with successive numbers
of their left neighbour."
This also preserves log n applications per value.
However it has a space leak,
because for the value with index @n@
all elements starting at @div n 2@ must be kept.
-}
iterateLeaky op x =
   let merge (a:as) b = a : merge b as
       merge _ _ = error "iterateLeaky: an empty list cannot occur"
       sqrs = map (\y -> op y y) z
       z = x : merge sqrs (map (op x) sqrs)
   in  z