{-# LANGUAGE PatternGuards, BangPatterns, TypeFamilies #-}
{-# LANGUAGE CPP #-}
#if __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Stream.Future.Skew
-- Copyright   :  (C) 2008-2015 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)
    , length    -- O(log n)
    , tail      -- O(1)
    , 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)
    , replicate -- O(log n)
    , insert    -- O(n)
    , insertBy
    , update
    , adjust    -- O(log n)
    , toFuture
    , singleton
    ) where

import Control.Applicative hiding (empty)
import Control.Comonad
import Data.Functor.Alt
import Data.Functor.Extend
#if MIN_VERSION_base(4,8,0)
import Prelude hiding (tail, drop, dropWhile, last, span, repeat, replicate, break)
import Data.Foldable (toList)
#else
import Data.Foldable
import Data.Traversable (Traversable, traverse)
import Prelude hiding (null, tail, drop, dropWhile, length, foldr, last, span, repeat, replicate, break)
#endif
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup hiding (Last)
#endif
import Data.Semigroup.Foldable
import Data.Semigroup.Traversable
#if MIN_VERSION_base(4,7,0)
import qualified GHC.Exts as Exts
#endif

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)
#if MIN_VERSION_base(4,8,0)
  length Tip{} = 1
  length (Bin n _ _ _) = n
  null _ = False
#endif

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 (a :< _) = extract a
  extract (Last a) = extract a

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 a0 = go a0 (Tip a0) where
    go :: a -> Complete a -> Future a
    go a as | ass <- bin a as as = as :< go a ass
  (<*>) = (<.>)

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
#if MIN_VERSION_base(4,8,0)
  length (Last t) = weight t
  length (t :< ts) = weight t + length ts
  null _ = False
#endif

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

-- | /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 #-}

#if !(MIN_VERSION_base(4,8,0))
-- | /O(log n)/.
length :: Future a -> Int
length (Last t) = weight t
length (t :< ts) = weight t + length ts
#endif

-- | /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 (<|) #-}

-- | /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 #-}

-- | /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 (extract 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)

#if MIN_VERSION_base(4,7,0)
instance Exts.IsList (Future a) where
  type Item (Future a) = a
  toList = Data.Foldable.toList
  fromList [] = error "fromList: empty list"
  fromList (x:xs) = go x xs
    where go a [] = singleton a
          go a (b:bs) = a <| go b bs
#else
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
#endif

toFuture :: [a] -> Maybe (Future a)
toFuture [] = Nothing
#if MIN_VERSION_base(4,7,0)
toFuture xs = Just (Exts.fromList xs)
#else
toFuture xs = Just (fromList xs)
#endif

-- /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