module CollectionTools ( module CollectionTools, foldr, replicate )
    where

-- Ports of functions formerly found in module List.

import Prelude ()
import Collections
import Stream
import QuickCheck hiding ( Config(size) )

-- ** Functions on Collections ** --

-- | Cross-product of two collections.  For Sequences it is in
-- Cantor-order, that is [(a!!1, b!!1), (a!!2, b!!1), (a!!1, b!!2),
-- (a!!3, b!!1), (a!!2, b!!2), ...]
cross_with_spec (+) a b = [ (a!!(i-j+1)) + (b!!j)
			    | i <- [1..], j <- [1..i]
			  , a `has_index` (i-j+1)
			  , b `has_index` j
			  ]
-- The implementation should do the same, but also terminate ;-)
-- (And it shouldn't call 'has_index' which is so expensive on
-- Streams.)

{-- For ordered structures, it does the same as:
o_cross_with (+) a b = join_apply (\ x -> join_apply (\y -> (x, y)
						     ) b
				  ) a
--}
{-- Algorithm with pattern matching:

TODO: instead of all this variants I'ld rather have a proper recursion
invariant for the helper function.  (Phil Wadler's "Prettier Printer"
paper gives a good expression of what I mean.)

We use Streams as the intermediate data structure, first to avoid
ambiguity, second because the needed operations are just Nil, Cons,
'front_view', and 'apply'.

cross a Empty = empty
cross a b     = cro Nil a
    where
    heads ss = [ h | (h:<_) <- ss ]
    tails ss = [ t | (_:<t) <- ss ]

    cro Empty (Empty) = empty
    cro ss    (Empty) = heads ss <+> cro (          tails ss) []
    cro ss    (a:<as) = heads ss <+> cro (Cons ns $ tails ss) as
                      where
		      ns = apply (\x->(a, x)) b 

-- Infinite-only version:

cross_with :: (Sequence seq a, Sequence seq' b, Sequence seq'' c) =>
	      (a -> b -> c) -> seq a -> seq' b -> seq'' c
cross_with (*) a b = cw empty a
    where
    cw ss a = apply first ss ++ cw (ns <: apply but_first ss) (but_first a)
	where
	ns = apply (first a*) b
--}

cross_with :: (Collection c1 a, Collection c2 b, Collection c3 c) =>
	      (a -> b -> c) -> c1 a -> c2 b -> c3 c
cross_with (*) as bs | is_empty bs = empty
		     | otherwise   = cw Nil as
    where
    join_nonempty :: ( Collection coll a, Collection c (coll a)
		     , Collection cool x ) =>
		     (coll a -> x) -> c (coll a) -> cool x
    join_nonempty f = join_apply (\s -> if is_empty s then empty
				                      else single$ f s)
    cw ss as = let heads = join_nonempty first     ss
		   tails = join_nonempty but_first ss
               in case (front_view ss, front_view as) of
        (Nothing, Nothing     ) -> empty
	(_      , Nothing     ) -> heads <+> cw tails empty
	(_      , Just (a, as)) -> heads <+> cw (Cons ns tails) as
	    where
	    ns = apply (a*) bstream
    bstream = stream bs

cross :: (Collection c1 a, Collection c2 b, Collection c3 (a, b)) =>
	 c1 a -> c2 b -> c3 (a, b)
cross = cross_with (,)
-- cross_with f a b = apply (uncurry f) $ cross a b


partition_spec coll p = (filter p coll, filter (not.p) coll)

partition coll p = fold empty f (+)
    where 
    f x | p x       = (single x, empty)
	| otherwise = (empty, single x)
    (a, b) + (x, y) = (a<+>x, b<+>y)


-- Unlike ancient 'inits' and 'tails' our 'fronts' and 'backs' always
-- return the shortest segment first.  I suspect ancient 'tails' did
-- only the inverse to achieve sharing, which democratic sequences
-- achieve in any case...

prop_fronts_front xs = forAll (elements $ fronts xs) (\ ys ->
                       ys  ==  xs `front` length ys  )
 
prop_fronts_increasing_length xs
    = strictly_increasing $ (apply size $ fronts xs :: [Int])
    -- Also checks, that no segment occurs twice.

fronts coll | is_empty coll = single empty
	    | otherwise     = fronts (but_last coll) >: coll

backs coll | is_empty coll = single empty
	   | otherwise     = backs (but_first coll) >: coll


unfold f x = case f x of
	     (Empty)    -> empty
	     (Single x) -> single x
	     (s :+: l)  -> unfold f s <+> unfold f l


-- ** Functions working from the right / back. ** --

-- foldr f e coll = fold id f (.) coll e

foldr1 f coll = foldr f (last coll) (but_last coll)

unfoldr g x = case g x of
	      (Nothing)      -> empty
	      (Just (f, bf)) -> f <: unfoldr g bf


-- scanr_spec f e coll = reverse $ map (foldr f e) (backs coll)

scanr (*) e coll = case back_view coll of
		     (Nothing)      -> empty
		     (Just (bl, l)) -> scanr (*) (l*e) bl
		   >: e

scanr1 (*) coll = scanr (*) (last coll) (but_last coll)

mapAccumR f s coll = case back_view coll of
		     (Nothing)      -> (s, empty)
		     (Just (xs, x)) -> (s'', ys >: y)
			 where (s', y) = f s x
			       (s'', ys) = mapAccumL f s' xs

-- Coding 'scanr' and 'mapAccumR' backwards gives them the same nice
-- properties as 'scanl' and 'mapAccumL': The result is constructed
-- lazily, instead of "all at once when you look at its first
-- element". 

-- ** And functions working from the left / front. ** --

foldl f e coll = fold id (flip f) (flip (.)) coll e

foldl1 f coll = foldl f (first coll) (but_first coll)

unfoldl f x = case f x of
	      (Nothing)      -> empty
	      (Just (bl, l)) -> unfoldl f bl >: l

scanl (*) e coll = e <: case front_view coll of
 		        (Nothing)      -> empty
			(Just (f, bf)) -> scanl (*) (e*f) bf

scanl1 (*) coll = scanl (*) (first coll) (but_first coll)

mapAccumL f s coll = case front_view coll of
		     (Nothing)      -> (s, empty)
		     (Just (x, xs)) -> (s'', y <: ys)
			 where (s', y) = f s x
			       (s'', ys) = mapAccumL f s' xs

{-- TODO

nub_by, intersect_by, union_by, maximum_by, minimum_by

zip family

--}

-- This function is only needed for backwards compatibility:
--   1. with list comprehensions (or pattern matching) of Haskell'98
--   2. to interface with Modules that don't yet use Collections.
to_list :: Collection coll a => coll a -> [a]
to_list = foldr (:) []
{-# RULES "to_list"    convert = to_list #-}
{-# RULES "to_list/id" to_list = id      #-}


-- Normal conversion from a list (using fold) will use the traversal
-- mode of the list, and consequently create the result element-wise
-- from a tip.  
-- single a ++ (single b ++ (single c ++ .. (single z ++ empty) .. ))

-- If the target structure is balanced the following function will
-- "balance the parenthesis'" in that expression, such that the
-- balancing algorithm will have no work to to.  The asymptotic
-- running time drops from NlogN to N.

balanced_from_list :: Collection coll a => [a] -> coll a
balanced_from_list xs = fst $ fl (size xs) xs
    where
    fl_spec n xs = (convert $ take n xs, drop n xs)
    -- The implementation of 'fl' avoids the copying done by 'take'.
    -- Only a constant factor, since the number of elements that would
    -- be copied is: n/2 + 2*n/4 + 4*n/8 + ... = n
    fl 0 xs = (empty,             xs)
    fl 1 xs = (single (first xs), but_first xs)
    fl n xs = (a <+> b,           rb)
	where
	n2 = n `div` 2
	(a, ra) = fl n2     xs
	(b, rb) = fl (n-n2) ra
{-# RULES "balanced_from_list"    convert = balanced_from_list #-}
{-# RULES "balanced_from_list/id" balanced_from_list = id      #-}

-- ** Functions on Sequences ** --

intersperse_spec sep seq = first seq ++ concat_apply (single sep >:)

intersperse sep seq = case split seq of
		       (Empty) -> empty
		       (Single x) -> single x
		       (s :+: l)  -> s' ++ single sep ++ l'
			   where l' = intersperse sep l
				 s' = intersperse sep s

group_by p coll = case front_view coll of
		  (Nothing)      -> empty
		  (Just (f, bf)) -> (f<:ff) <: group_by p bf'
		      where (ff, bf') = span (p f) ff

group coll = group_by (==) coll


-- ** helper functions ** --

strictly_increasing coll = and $ coll * but_first coll
    where (*) :: (Ord a, Sequence seq a, Sequence seq Bool) =>
		 seq a -> seq a -> seq Bool
	  (*) = zip_with (<)

-- increasing coll = and $ zip_with (<=) coll (but_first coll)