module Data.Series.Generic.Scans (
    postscanl,
    prescanl,

    -- * Filling missing data

    forwardFill,
) where

import           Data.Series.Generic.Definition ( Series(..) )

import           Data.Vector.Generic            ( Vector )
import qualified Data.Vector.Generic            as Vector    

-- $setup

-- >>> import qualified Data.Series.Generic ( Series )

-- >>> import qualified Data.Series.Generic as Series

-- >>> import qualified Data.Series.Index as Index


-- | \(O(n)\) Left-to-right postscan.

--

-- >>> import qualified Data.Vector as V 

-- >>> let xs = Series.fromList (zip [0..] [1,2,3,4]) :: Series V.Vector Int Int

-- >>> xs

-- index | values

-- ----- | ------

--     0 |      1

--     1 |      2

--     2 |      3

--     3 |      4

-- >>> postscanl (+) 0 xs

-- index | values

-- ----- | ------

--     0 |      1

--     1 |      3

--     2 |      6

--     3 |     10

postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> Series v k b -> Series v k a
{-# INLINABLE postscanl #-}
postscanl :: forall (v :: * -> *) a b k.
(Vector v a, Vector v b) =>
(a -> b -> a) -> a -> Series v k b -> Series v k a
postscanl a -> b -> a
f a
s (MkSeries Index k
ix v b
vs) = Index k -> v a -> Series v k a
forall {k} (v :: k -> *) k1 (a :: k).
Index k1 -> v a -> Series v k1 a
MkSeries Index k
ix (v a -> Series v k a) -> v a -> Series v k a
forall a b. (a -> b) -> a -> b
$ (a -> b -> a) -> a -> v b -> v a
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b -> a) -> a -> v b -> v a
Vector.postscanl a -> b -> a
f a
s v b
vs


-- | \(O(n)\) Left-to-right prescan.

--

-- >>> import qualified Data.Vector as V 

-- >>> let xs = Series.fromList (zip [0..] [1,2,3,4]) :: Series V.Vector Int Int

-- >>> xs

-- index | values

-- ----- | ------

--     0 |      1

--     1 |      2

--     2 |      3

--     3 |      4

-- >>> prescanl (+) 0 xs

-- index | values

-- ----- | ------

--     0 |      0

--     1 |      1

--     2 |      3

--     3 |      6

prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> Series v k b -> Series v k a
{-# INLINABLE prescanl #-}
prescanl :: forall (v :: * -> *) a b k.
(Vector v a, Vector v b) =>
(a -> b -> a) -> a -> Series v k b -> Series v k a
prescanl a -> b -> a
f a
s (MkSeries Index k
ix v b
vs) = Index k -> v a -> Series v k a
forall {k} (v :: k -> *) k1 (a :: k).
Index k1 -> v a -> Series v k1 a
MkSeries Index k
ix (v a -> Series v k a) -> v a -> Series v k a
forall a b. (a -> b) -> a -> b
$ (a -> b -> a) -> a -> v b -> v a
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b -> a) -> a -> v b -> v a
Vector.prescanl a -> b -> a
f a
s v b
vs


-- | \(O(n)\) Replace all instances of 'Nothing' with the last previous

-- value which was not 'Nothing'.

--

-- >>> import qualified Data.Vector as V 

-- >>> let xs = Series.fromList (zip [0..] [Just 1, Just 2,Nothing, Just 3]) :: Series V.Vector Int (Maybe Int)

-- >>> xs

-- index |  values

-- ----- |  ------

--     0 |  Just 1

--     1 |  Just 2

--     2 | Nothing

--     3 |  Just 3

-- >>> forwardFill 0 xs

-- index | values

-- ----- | ------

--     0 |      1

--     1 |      2

--     2 |      2

--     3 |      3

--

-- If the first entry of the series is missing, the first input to 'forwardFill' will be used:

--

-- >>> let ys = Series.fromList (zip [0..] [Nothing, Just 2,Nothing, Just 3]) :: Series V.Vector Int (Maybe Int)

-- >>> ys

-- index |  values

-- ----- |  ------

--     0 | Nothing

--     1 |  Just 2

--     2 | Nothing

--     3 |  Just 3

-- >>> forwardFill 0 ys

-- index | values

-- ----- | ------

--     0 |      0

--     1 |      2

--     2 |      2

--     3 |      3

forwardFill :: (Vector v a, Vector v (Maybe a))
            => a -- ^ Until the first non-'Nothing' is found, 'Nothing' will be filled with this value.

            -> Series v k (Maybe a)
            -> Series v k a
{-# INLINABLE forwardFill #-}
forwardFill :: forall (v :: * -> *) a k.
(Vector v a, Vector v (Maybe a)) =>
a -> Series v k (Maybe a) -> Series v k a
forwardFill = (a -> Maybe a -> a) -> a -> Series v k (Maybe a) -> Series v k a
forall (v :: * -> *) a b k.
(Vector v a, Vector v b) =>
(a -> b -> a) -> a -> Series v k b -> Series v k a
postscanl a -> Maybe a -> a
forall a. a -> Maybe a -> a
go
    where
        go :: a -> Maybe a -> a
        go :: forall a. a -> Maybe a -> a
go a
lastValid Maybe a
Nothing = a
lastValid
        go a
_        (Just a
v) = a
v