{-# LANGUAGE PatternGuards, BangPatterns #-}
{-# LANGUAGE CPP #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Stream.Future.Skew
-- Copyright   :  (C) 2008-2011 Edward Kmett,
--                (C) 2004 Dave Menendez
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Anticausal streams implemented as non-empty skew binary random access lists
--
-- The Applicative zips streams, but since these are potentially infinite
-- this is stricter than would be desired. You almost always want
------------------------------------------------------------------------------


module Data.Stream.Future.Skew
    ( Future(..)
    , (<|)      -- O(1)
    , cons
    , length    -- O(log n)
    , head      -- O(1)
    , tail      -- O(1)
    , tails
    , last      -- O(log n)
    , uncons    -- O(1)
    , index     -- O(log n)
    , drop      -- O(log n)
    , dropWhile -- O(n)
    , indexed
    , from
    , break
    , span
    , split     -- O(log n)
    , splitW    -- O(log n)
    , repeat
    , replicate -- O(log n)
    , insert    -- O(n)
    , insertBy
    , update
    , adjust    -- O(log n)
    , fromList
    , toFuture
    ) where

import Control.Applicative hiding (empty)
import Control.Comonad
import Data.Functor.Alt
import Data.Functor.Extend
import Data.Foldable hiding (toList)
import Data.Traversable (Traversable, traverse)
import Data.Semigroup hiding (Last)
import Data.Semigroup.Foldable
import Data.Semigroup.Traversable
import Prelude hiding (null, head, tail, drop, dropWhile, length, foldr, last, span, repeat, replicate, break)

infixr 5 :<, <|

data Complete a
    = Tip a
    | Bin {-# UNPACK #-} !Int a !(Complete a) !(Complete a)
    deriving Show

instance Functor Complete where
  fmap f (Tip a) = Tip (f a)
  fmap f (Bin w a l r) = Bin w (f a) (fmap f l) (fmap f r)

instance Extend Complete where
  extended = extend

instance Comonad Complete where
  extend f w@Tip {} = Tip (f w)
  extend f w@(Bin n _ l r) = Bin n (f w) (extend f l) (extend f r)
  extract (Tip a) = a
  extract (Bin _ a _ _) = a

instance Foldable Complete where
  foldMap f (Tip a) = f a
  foldMap f (Bin _ a l r) = f a `mappend` foldMap f l `mappend` foldMap f r
  foldr f z (Tip a) = f a z
  foldr f z (Bin _ a l r) = f a (foldr f (foldr f z r) l)

instance Foldable1 Complete where
  foldMap1 f (Tip a) = f a
  foldMap1 f (Bin _ a l r) = f a <> foldMap1 f l <> foldMap1 f r

instance Traversable Complete where
  traverse f (Tip a) = Tip <$> f a
  traverse f (Bin n a l r) = Bin n <$> f a <*> traverse f l <*> traverse f r

instance Traversable1 Complete where
  traverse1 f (Tip a) = Tip <$> f a
  traverse1 f (Bin n a l r) = Bin n <$> f a <.> traverse1 f l <.> traverse1 f r

bin :: a -> Complete a -> Complete a -> Complete a
bin a l r = Bin (1 + weight l + weight r) a l r
{-# INLINE bin #-}

weight :: Complete a -> Int
weight Tip{} = 1
weight (Bin w _ _ _) = w
{-# INLINE weight #-}

-- A future is a non-empty skew binary random access list of nodes.
-- The last node, however, is allowed to contain fewer values.
data Future a
  = Last !(Complete a)
  | !(Complete a) :< Future a
--  deriving Show


instance Show a => Show (Future a) where
  showsPrec d as = showParen (d >= 10) $
    showString "fromList " . showsPrec 11 (toList as)

instance Functor Future where
  fmap f (t :< ts) = fmap f t :< fmap f ts
  fmap f (Last t) = Last (fmap f t)

instance Extend Future where
  extended = extend

instance Comonad Future where
  extend g (Last t)  = Last (extendTree g t Last)
  extend g (t :< ts) = extendTree g t (:< ts) :< extend g ts
  extract = head

extendTree :: (Future a -> b) -> Complete a -> (Complete a -> Future a) -> Complete b
extendTree g w@Tip{}         f = Tip (g (f w))
extendTree g w@(Bin n _ l r) f = Bin n (g (f w)) (extendTree g l (:< f r))  (extendTree g r f)

instance Apply Future where
  Last (Tip f)         <.> as                   = singleton (f (extract as))
  fs                   <.> Last (Tip a)         = singleton (extract fs a)
  Last (Bin _ f lf rf) <.> Last (Bin _ a la ra) = f a <| (lf :< Last rf  <.> la :< Last ra )
  Last (Bin _ f lf rf) <.> Bin _ a la ra :< as  = f a <| (lf :< Last rf  <.> la :< ra :< as)
  Last (Bin _ f lf rf) <.> Tip a :< as          = f a <| (lf :< Last rf  <.> as            )
  Bin _ f lf rf :< fs  <.> Last (Bin _ a la ra) = f a <| (lf :< rf :< fs <.> la :< Last ra )
  Bin _ f lf rf :< fs  <.> Tip a :< as          = f a <| (lf :< rf :< fs <.> as            )
  Bin _ f lf rf :< fs  <.> Bin _ a la ra :< as  = f a <| (lf :< rf :< fs <.> la :< ra :< as)
  Tip f :< fs          <.> Tip a :< as          = f a <| (fs             <.> as            )
  Tip f :< fs          <.> Bin _ a la ra :< as  = f a <| (fs             <.> la :< ra :< as)
  Tip f :< fs          <.> Last (Bin _ a la ra) = f a <| (fs             <.> la :< Last ra )

instance ComonadApply Future where
  (<@>) = (<.>)

instance Applicative Future where
  pure = repeat
  (<*>) = (<.>)

instance Alt Future where
  as <!> bs = foldr (<|) bs as

instance Foldable Future where
  foldMap f (t :< ts) = foldMap f t `mappend` foldMap f ts
  foldMap f (Last t) = foldMap f t
  foldr f z (t :< ts) = foldr f (foldr f z ts) t
  foldr f z (Last t) = foldr f z t

toList :: Future a -> [a]
toList = foldr (:) []

instance Foldable1 Future where
  foldMap1 f (t :< ts) = foldMap1 f t <> foldMap1 f ts
  foldMap1 f (Last t) = foldMap1 f t

instance Traversable Future where
  traverse f (t :< ts) = (:<) <$> traverse f t <*> traverse f ts
  traverse f (Last t) = Last <$> traverse f t

instance Traversable1 Future where
  traverse1 f (t :< ts) = (:<) <$> traverse1 f t <.> traverse1 f ts
  traverse1 f (Last t) = Last <$> traverse1 f t

repeat :: a -> Future a
repeat a0 = go a0 (Tip a0)
    where
      go :: a -> Complete a -> Future a
      go a as | ass <- bin a as as = as :< go a ass
{-# INLINE repeat #-}

-- | /O(log n)/
replicate :: Int -> a -> Future a
replicate n a
  | n <= 0    = error "replicate: non-positive argument"
  | otherwise = go 1 n a (Tip a) (\ _ r -> r)
  where
  -- invariants:
  -- tb is a complete tree of i nodes all equal to b
  -- 1 <= i = 2^m-1 <= j
  -- k accepts r such that 0 <= r < i
  go :: Int -> Int -> b -> Complete b -> (Int -> Future b -> r) -> r
  go !i !j b tb k
    | j >= i2p1 = go i2p1 j b (Bin i2p1 b tb tb) k'
    | j >= i2   = k (j - i2) (tb :< Last tb)
    | otherwise = k (j - i) (Last tb)
    where
      i2 = i * 2
      i2p1 = i2 + 1
      k' r xs
        | r >= i2   = k (r - i2) (tb :< tb :< xs)
        | r >= i    = k (r - i) (tb :< xs)
        | otherwise = k r xs
{-# INLINE replicate #-}

mapWithIndex :: (Int -> a -> b) -> Future a -> Future b
mapWithIndex f0 as0 = spine f0 0 as0
  where
    spine f m (Last as) = Last (tree f m as)
    spine f m (a :< as) = tree f m a :< spine f (m + weight a) as
    tree f m (Tip a) = Tip (f m a)
    tree f m (Bin n a l r) = Bin n (f m a) (tree f (m + 1) l) (tree f (m + 1 + weight l) r)

indexed :: Future a -> Future (Int, a)
indexed = mapWithIndex (,)
{-# INLINE indexed #-}

from :: Num a => a -> Future a
from a = mapWithIndex ((+) . fromIntegral) (pure a)
{-# INLINE from #-}

-- | /O(1)/
singleton :: a -> Future a
singleton a = Last (Tip a)
{-# INLINE singleton #-}

-- | /O(log n)/.
length :: Future a -> Int
length (Last t) = weight t
length (t :< ts) = weight t + length ts

-- | /O(1)/ cons
(<|) :: a -> Future a -> Future a
a <| (l :< Last r)
  | weight l == weight r = Last (bin a l r)
a <| (l :< r :< as)
  | weight l == weight r = bin a l r :< as
a <| as = Tip a :< as
{-# INLINE (<|) #-}


cons :: a -> Future a -> Future a
cons = (<|)
{-# INLINE cons #-}

-- | /O(1)/
head :: Future a -> a
head (a :< _) = extract a
head (Last a) = extract a
{-# INLINE head #-}

-- | /O(1)/.
tail :: Future a -> Maybe (Future a)
tail (Tip{} :< ts) = Just ts
tail (Bin _ _ l r :< ts) = Just (l :< r :< ts)
tail (Last Tip{}) = Nothing
tail (Last (Bin _ _ l r)) = Just (l :< Last r)
{-# INLINE tail #-}

tails :: Future a -> Future (Future a)
tails = duplicate
{-# INLINE tails #-}

-- | /O(log n)/.
last :: Future a -> a
last (_ :< as) = last as
last (Last as) = go as
  where go (Tip a) = a
        go (Bin _ _ _ r) = go r

-- | /O(1)/.
uncons :: Future a -> (a, Maybe (Future a))
uncons (Last (Tip a))       = (a, Nothing)
uncons (Last (Bin _ a l r)) = (a, Just (l :< Last r))
uncons (Tip a       :< as)  = (a, Just as)
uncons (Bin _ a l r :< as)  = (a, Just (l :< r :< as))
{-# INLINE uncons #-}

-- | /O(log n)/.
index :: Int -> Future a -> a
index i (Last t)
  | i < weight t = indexComplete i t
  | otherwise    = error "index: out of range"
index i (t :< ts)
  | i < w     = indexComplete i t
  | otherwise = index (i - w) ts
  where w = weight t

indexComplete :: Int -> Complete a -> a
indexComplete 0 (Tip a) = a
indexComplete i (Bin w a l r)
  | i == 0    = a
  | i <= w'   = indexComplete (i-1) l
  | otherwise = indexComplete (i-1-w') r
  where w' = div w 2
indexComplete _ _ = error "index: index out of range"

-- | /O(log n)/.
drop :: Int -> Future a -> Maybe (Future a)
drop 0 ts = Just ts
drop i (t :< ts) = case compare i w of
  LT -> Just (dropComplete i t (:< ts))
  EQ -> Just ts
  GT -> drop (i - w) ts
  where w = weight t
drop i (Last t)
  | i < w     = Just (dropComplete i t Last)
  | otherwise = Nothing
  where w = weight t

dropComplete :: Int -> Complete a -> (Complete a -> Future a) -> Future a
dropComplete 0 t f             = f t
dropComplete 1 (Bin _ _ l r) f = l :< f r
dropComplete i (Bin w _ l r) f = case compare (i - 1) w' of
  LT -> dropComplete (i-1) l (:< f r)
  EQ -> f r
  GT -> dropComplete (i-1-w') r f
  where w' = div w 2
dropComplete _ _ _ = error "drop: index out of range"

-- /O(n)/.
dropWhile :: (a -> Bool) -> Future a -> Maybe (Future a)
dropWhile p as
  | p (head as) = tail as >>= dropWhile p
  | otherwise = Just as

-- /O(n)/
span :: (a -> Bool) -> Future a -> ([a], Maybe (Future a))
span p aas = case uncons aas of
  (a, Just as) | p a, (ts, fs) <- span p as -> (a:ts, fs)
  (a, Nothing) | p a                        -> ([a], Nothing)
  (_, _)                                    -> ([], Just aas)

-- /O(n)/
break :: (a -> Bool) -> Future a -> ([a], Maybe (Future a))
break p = span (not . p)

-- /(O(n), O(log n))/ split at _some_ edge where function goes from False to True.
-- best used with a monotonic function
split :: (a -> Bool) -> Future a -> ([a], Maybe (Future a))
split p l@(Last a)
  | p (extract a)  = ([], Just l)
  | otherwise      = splitComplete p a Last
split p (a :< as)
  | p (extract as) = splitComplete p a (:< as)
  | (ts, fs) <- split p as = (foldr (:) ts a, fs)

-- for use when we know the split occurs within a given tree
splitComplete :: (a -> Bool) -> Complete a -> (Complete a -> Future a) -> ([a], Maybe (Future a))
splitComplete p t@(Tip a) f
  | p a       = ([], Just (f t))
  | otherwise = ([a], Nothing)
splitComplete p t@(Bin _ a l r) f
  | p a                                               = ([], Just (f t))
  | p (extract r), (ts, fs) <- splitComplete p l (:< f r) = (a:ts, fs)
  |                (ts, fs) <- splitComplete p r f        = (a:foldr (:) ts l, fs)

-- /(O(n), O(log n))/ split at _some_ edge where function goes from False to True.
-- best used with a monotonic function
--
-- > splitW p xs = (map extract &&& fmap (fmap extract)) . split p . duplicate
splitW :: (Future a -> Bool) -> Future a -> ([a], Maybe (Future a))
splitW p l@(Last a)
  | p l       = ([], Just l)
  | otherwise = splitCompleteW p a Last
splitW p (a :< as)
  | p as                    = splitCompleteW p a (:< as)
  | (ts, fs) <- splitW p as = (foldr (:) ts a, fs)

-- for use when we know the split occurs within a given tree
splitCompleteW :: (Future a -> Bool) -> Complete a -> (Complete a -> Future a) -> ([a], Maybe (Future a))
splitCompleteW p t@(Tip a) f
  | w <- f t, p w = ([], Just w)
  | otherwise = ([a], Nothing)
splitCompleteW p t@(Bin _ a l r) f
  | w <- f t, p w                                    = ([], Just w)
  | w <- f r, p w, (ts, fs) <- splitCompleteW p l (:< w) = (a:ts, fs)
  |                (ts, fs) <- splitCompleteW p r f      = (a:foldr (:) ts l, fs)

fromList :: [a] -> Future a
fromList [] = error "fromList: empty list"
fromList (x:xs) = go x xs
  where go a [] = singleton a
        go a (b:bs) = a <| go b bs

toFuture :: [a] -> Maybe (Future a)
toFuture [] = Nothing
toFuture xs = Just (fromList xs)

-- /O(n)/
insert :: Ord a => a -> Future a -> Future a
insert a as = case split (a<=) as of
    (_, Nothing)  -> foldr (<|) (singleton a) as
    (ts, Just as') -> foldr (<|) (a <| as') ts

-- /O(n)/. Finds the split in O(log n), but then has to recons
insertBy :: (a -> a -> Ordering) -> a -> Future a -> Future a
insertBy cmp a as = case split (\b -> cmp a b <= EQ) as of
    (_, Nothing)  -> foldr (<|) (singleton a) as
    (ts, Just as') -> foldr (<|) (a <| as') ts

-- /O(log n)/ Change the value of the nth entry in the future
adjust :: Int -> (a -> a) -> Future a -> Future a
adjust !n f d@(Last a)
  | n < weight a = Last (adjustComplete n f a)
  | otherwise = d
adjust !n f (a :< as)
  | n < w = adjustComplete n f a :< as
  | otherwise = a :< adjust (n - w) f as
  where w = weight a

adjustComplete :: Int -> (a -> a) -> Complete a -> Complete a
adjustComplete 0 f (Tip a) = Tip (f a)
adjustComplete _ _ t@Tip{} = t
adjustComplete n f (Bin m a l r)
  | n == 0 = Bin m (f a) l r
  | n < w = Bin m a (adjustComplete (n - 1) f l) r
  | otherwise = Bin m a l (adjustComplete (n - 1 - w) f r)
  where w = weight l

update :: Int -> a -> Future a -> Future a
update n = adjust n . const