module Data.Series.Generic.View (
    -- * Accessing a single element
    (!),
    at,
    iat,

    -- * Bulk access
    select,
    slice,
    selectWhere,
    selectSubset,
    Selection,

    -- * Resizing
    require,
    requireWith,
    filter,
    filterWithKey,
    catMaybes,
    dropIndex,

    -- * Creating and accessing ranges
    Range(..),
    to,
    from,
    upto,
) where


import           Data.Series.Index      ( Index )
import qualified Data.Series.Index      as Index
import qualified Data.Series.Index.Internal as Index.Internal
import           Data.Maybe             ( fromJust, isJust )
import           Data.Series.Generic.Definition ( Series(..) )
import qualified Data.Series.Generic.Definition as G
import           Data.Set               ( Set )
import qualified Data.Set               as Set
import qualified Data.Vector            as Boxed
import           Data.Vector.Generic    ( Vector )
import qualified Data.Vector.Generic    as Vector

import           Prelude                hiding ( filter )

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

infixr 9 `to` -- Ensure that @to@ binds strongest
infixl 1 `select` 


-- | \(O(1)\). Extract a single value from a series, by index. 
-- An exception is thrown if the index is out-of-bounds.
--
-- A safer alternative is @iat@, which returns 'Nothing' if the index is
-- out-of-bounds.
(!) :: Vector v a => Series v k a -> Int -> a
(MkSeries Index k
_ v a
vs) ! :: forall (v :: * -> *) a k. Vector v a => Series v k a -> Int -> a
! Int
ix = v a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
(Vector.!) v a
vs Int
ix


-- | \(O(\log n)\). Extract a single value from a series, by key.
at :: (Vector v a, Ord k) => Series v k a -> k -> Maybe a
at :: forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> k -> Maybe a
at (MkSeries Index k
ks v a
vs) k
k = do
    Int
ix <- k -> Index k -> Maybe Int
forall k. Ord k => k -> Index k -> Maybe Int
Index.lookupIndex k
k Index k
ks
    a -> Maybe a
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
Vector.unsafeIndex v a
vs Int
ix 
{-# INLINE at #-}


-- | \(O(1)\). Extract a single value from a series, by index.
iat :: Vector v a => Series v k a -> Int -> Maybe a
iat :: forall (v :: * -> *) a k.
Vector v a =>
Series v k a -> Int -> Maybe a
iat (MkSeries Index k
_ v a
vs) =  v a -> Int -> Maybe a
forall (v :: * -> *) a. Vector v a => v a -> Int -> Maybe a
(Vector.!?) v a
vs
{-# INLINE iat #-}


-- | require a series with a new index.
-- Contrary to 'select', all keys in @'Set' k@ will be present in the re-indexed series.
require :: (Vector v a, Vector v (Maybe a), Ord k) 
        => Series v k a -> Index k -> Series v k (Maybe a)
{-# INLINE require #-}
require :: forall (v :: * -> *) a k.
(Vector v a, Vector v (Maybe a), Ord k) =>
Series v k a -> Index k -> Series v k (Maybe a)
require = (k -> Maybe a)
-> (a -> Maybe a)
-> Series v k a
-> Index k
-> Series v k (Maybe a)
forall (v :: * -> *) a b k.
(Vector v a, Vector v b, Ord k) =>
(k -> b) -> (a -> b) -> Series v k a -> Index k -> Series v k b
requireWith (Maybe a -> k -> Maybe a
forall a b. a -> b -> a
const Maybe a
forall a. Maybe a
Nothing) a -> Maybe a
forall a. a -> Maybe a
Just


-- | Generalization of 'require', which maps missing keys to values.
-- This is particularly useful for 'Vector' instances which don't support 'Maybe', like "Data.Vector.Unboxed".
requireWith :: (Vector v a, Vector v b, Ord k)
            => (k -> b)  -- ^ Function to apply to keys which are missing from the input series, but required in the input index
            -> (a -> b)  -- ^ Function to apply to values which are in the input series and input index.
            -> Series v k a 
            -> Index k 
            -> Series v k b
{-# INLINE requireWith #-}
requireWith :: forall (v :: * -> *) a b k.
(Vector v a, Vector v b, Ord k) =>
(k -> b) -> (a -> b) -> Series v k a -> Index k -> Series v k b
requireWith k -> b
replacement a -> b
f Series v k a
xs Index k
ss 
    = let existingKeys :: Index k
existingKeys = Series v k a -> Index k
forall {k1} (v :: k1 -> *) k2 (a :: k1). Series v k2 a -> Index k2
index Series v k a
xs Index k -> Index k -> Index k
forall k. Ord k => Index k -> Index k -> Index k
`Index.intersection` Index k
ss
          newKeys :: Index k
newKeys      = Index k
ss Index k -> Index k -> Index k
forall k. Ord k => Index k -> Index k -> Index k
`Index.difference` Index k
existingKeys
       in (a -> b) -> Series v k a -> Series v k b
forall (v :: * -> *) a b k.
(Vector v a, Vector v b) =>
(a -> b) -> Series v k a -> Series v k b
G.map a -> b
f (Series v k a
xs Series v k a -> Index k -> Series v k a
forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
`selectSubset` Index k
existingKeys) Series v k b -> Series v k b -> Series v k b
forall a. Semigroup a => a -> a -> a
<> Index k -> v b -> Series v k b
forall {k} (v :: k -> *) k1 (a :: k).
Index k1 -> v a -> Series v k1 a
MkSeries Index k
newKeys (Int -> [b] -> v b
forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
Vector.fromListN (Index k -> Int
forall k. Index k -> Int
Index.size Index k
newKeys) (k -> b
replacement (k -> b) -> [k] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Index k -> [k]
forall k. Index k -> [k]
Index.toAscList Index k
newKeys))


-- | Drop the index of a series by replacing it with an @Int@-based index. Values will
-- be indexed from 0.
dropIndex :: Series v k a -> Series v Int a
{-# INLINE dropIndex #-}
dropIndex :: forall {k} (v :: k -> *) k (a :: k). Series v k a -> Series v Int a
dropIndex (MkSeries Index k
ks v a
vs) = Index Int -> v a -> Series v Int a
forall {k} (v :: k -> *) k1 (a :: k).
Index k1 -> v a -> Series v k1 a
MkSeries ([Int] -> Index Int
forall k. [k] -> Index k
Index.Internal.fromDistinctAscList [Int
0..Index k -> Int
forall k. Index k -> Int
Index.size Index k
ks Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) v a
vs


-- | Filter elements. Only elements for which the predicate is @True@ are kept. 
-- Notice that the filtering is done on the values, not on the keys; see 'filterWithKey'
-- to filter while taking keys into account.
filter :: (Vector v a, Vector v Int, Ord k) 
       => (a -> Bool) -> Series v k a -> Series v k a
{-# INLINE filter #-}
filter :: forall (v :: * -> *) a k.
(Vector v a, Vector v Int, Ord k) =>
(a -> Bool) -> Series v k a -> Series v k a
filter a -> Bool
predicate xs :: Series v k a
xs@(MkSeries Index k
ks v a
vs) 
    = let indicesToKeep :: v Int
indicesToKeep = (a -> Bool) -> v a -> v Int
forall (v :: * -> *) a.
(Vector v a, Vector v Int) =>
(a -> Bool) -> v a -> v Int
Vector.findIndices a -> Bool
predicate v a
vs
          keysToKeep :: Index k
keysToKeep = [k] -> Index k
forall k. [k] -> Index k
Index.Internal.fromDistinctAscList [Int -> Index k -> k
forall k. HasCallStack => Int -> Index k -> k
Index.Internal.elemAt Int
ix Index k
ks | Int
ix <- v Int -> [Int]
forall (v :: * -> *) a. Vector v a => v a -> [a]
Vector.toList v Int
indicesToKeep]
       in Series v k a
xs Series v k a -> Index k -> Series v k a
forall (s :: * -> *) (v :: * -> *) a k.
(Selection s, Vector v a, Ord k) =>
Series v k a -> s k -> Series v k a
forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
`select` Index k
keysToKeep


-- | Filter elements, taking into account the corresponding key. Only elements for which 
-- the predicate is @True@ are kept. 
filterWithKey :: (Vector v a, Vector v Int, Vector v Bool, Ord k) 
              => (k -> a -> Bool) 
              -> Series v k a 
              -> Series v k a
{-# INLINE filterWithKey #-}
filterWithKey :: forall (v :: * -> *) a k.
(Vector v a, Vector v Int, Vector v Bool, Ord k) =>
(k -> a -> Bool) -> Series v k a -> Series v k a
filterWithKey k -> a -> Bool
predicate Series v k a
xs = Series v k a
xs Series v k a -> Series v k Bool -> Series v k a
forall (v :: * -> *) a k.
(Vector v a, Vector v Int, Vector v Bool, Ord k) =>
Series v k a -> Series v k Bool -> Series v k a
`selectWhere` (k -> a -> Bool) -> Series v k a -> Series v k Bool
forall (v :: * -> *) a b k.
(Vector v a, Vector v b) =>
(k -> a -> b) -> Series v k a -> Series v k b
G.mapWithKey k -> a -> Bool
predicate Series v k a
xs


-- | \(O(n)\) Only keep elements which are @'Just' v@. 
catMaybes :: (Vector v a, Vector v (Maybe a), Vector v Int, Ord k) 
       => Series v k (Maybe a) -> Series v k a
{-# INLINE catMaybes #-}
catMaybes :: forall (v :: * -> *) a k.
(Vector v a, Vector v (Maybe a), Vector v Int, Ord k) =>
Series v k (Maybe a) -> Series v k a
catMaybes = (Maybe a -> a) -> Series v k (Maybe a) -> Series v k a
forall (v :: * -> *) a b k.
(Vector v a, Vector v b) =>
(a -> b) -> Series v k a -> Series v k b
G.map Maybe a -> a
forall a. HasCallStack => Maybe a -> a
fromJust (Series v k (Maybe a) -> Series v k a)
-> (Series v k (Maybe a) -> Series v k (Maybe a))
-> Series v k (Maybe a)
-> Series v k a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe a -> Bool) -> Series v k (Maybe a) -> Series v k (Maybe a)
forall (v :: * -> *) a k.
(Vector v a, Vector v Int, Ord k) =>
(a -> Bool) -> Series v k a -> Series v k a
filter Maybe a -> Bool
forall a. Maybe a -> Bool
isJust


-- | Datatype representing an /inclusive/ range of keys, which can either be bounded
-- or unbounded. The canonical ways to construct a 'Range' are to use 'to', 'from', and 'upto':
--
-- >>> 'a' `to` 'z'
-- Range (from 'a' to 'z')
-- >>> from 'd'
-- Range (from 'd')
-- >>> upto 'q'
-- Range (up to 'q')
--
-- A 'Range' can be used to efficiently select a sub-series with 'select'.
data Range k 
    = BoundedRange k k
    | From k
    | UpTo k
    deriving (Range k -> Range k -> Bool
(Range k -> Range k -> Bool)
-> (Range k -> Range k -> Bool) -> Eq (Range k)
forall k. Eq k => Range k -> Range k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall k. Eq k => Range k -> Range k -> Bool
== :: Range k -> Range k -> Bool
$c/= :: forall k. Eq k => Range k -> Range k -> Bool
/= :: Range k -> Range k -> Bool
Eq)


instance Show k => Show (Range k) where
    show :: Range k -> String
    show :: Range k -> String
show (BoundedRange k
start k
stop) = [String] -> String
forall a. Monoid a => [a] -> a
mconcat [String
"Range (from ", k -> String
forall a. Show a => a -> String
show k
start, String
" to ", k -> String
forall a. Show a => a -> String
show k
stop, String
")"]
    show (From k
start) = [String] -> String
forall a. Monoid a => [a] -> a
mconcat [String
"Range (from ", k -> String
forall a. Show a => a -> String
show k
start, String
")"]
    show (UpTo k
stop) = [String] -> String
forall a. Monoid a => [a] -> a
mconcat [String
"Range (up to ", k -> String
forall a. Show a => a -> String
show k
stop, String
")"]


-- | Find the keys which are in range. In case of an empty 'Series',
-- the returned value is 'Nothing'.
keysInRange :: Ord k => Series v k a -> Range k -> Maybe (k, k)
{-# INLINE keysInRange #-}
keysInRange :: forall {k} k (v :: k -> *) (a :: k).
Ord k =>
Series v k a -> Range k -> Maybe (k, k)
keysInRange (MkSeries Index k
ks v a
_) Range k
rng
    = let inrange :: Set k
inrange = Range k -> Set k
inRange Range k
rng
       in if Set k -> Bool
forall a. Set a -> Bool
Set.null Set k
inrange 
            then Maybe (k, k)
forall a. Maybe a
Nothing
            else (k, k) -> Maybe (k, k)
forall a. a -> Maybe a
Just (Set k -> k
forall a. Set a -> a
Set.findMin Set k
inrange, Set k -> k
forall a. Set a -> a
Set.findMax Set k
inrange)
    where
        inRange :: Range k -> Set k
inRange (BoundedRange k
start k
stop)  = (k -> Bool) -> Set k -> Set k
forall a. (a -> Bool) -> Set a -> Set a
Set.takeWhileAntitone (k -> k -> Bool
forall a. Ord a => a -> a -> Bool
<= k
stop) 
                                           (Set k -> Set k) -> Set k -> Set k
forall a b. (a -> b) -> a -> b
$ (k -> Bool) -> Set k -> Set k
forall a. (a -> Bool) -> Set a -> Set a
Set.dropWhileAntitone (k -> k -> Bool
forall a. Ord a => a -> a -> Bool
< k
start) (Set k -> Set k) -> Set k -> Set k
forall a b. (a -> b) -> a -> b
$ Index k -> Set k
forall k. Index k -> Set k
Index.toSet Index k
ks
        inRange (From k
start)               = (k -> Bool) -> Set k -> Set k
forall a. (a -> Bool) -> Set a -> Set a
Set.dropWhileAntitone (k -> k -> Bool
forall a. Ord a => a -> a -> Bool
< k
start) (Set k -> Set k) -> Set k -> Set k
forall a b. (a -> b) -> a -> b
$ Index k -> Set k
forall k. Index k -> Set k
Index.toSet Index k
ks
        inRange (UpTo k
stop)                = (k -> Bool) -> Set k -> Set k
forall a. (a -> Bool) -> Set a -> Set a
Set.takeWhileAntitone (k -> k -> Bool
forall a. Ord a => a -> a -> Bool
<= k
stop) (Set k -> Set k) -> Set k -> Set k
forall a b. (a -> b) -> a -> b
$ Index k -> Set k
forall k. Index k -> Set k
Index.toSet Index k
ks


-- | Create a bounded 'Range' which can be used for slicing. This function
-- is expected to be used in conjunction with 'select'.
--
-- For unbound ranges, see 'from' and 'upto'.
to :: Ord k => k -> k -> Range k
to :: forall k. Ord k => k -> k -> Range k
to k
k1 k
k2 = k -> k -> Range k
forall k. k -> k -> Range k
BoundedRange (k -> k -> k
forall a. Ord a => a -> a -> a
min k
k1 k
k2) (k -> k -> k
forall a. Ord a => a -> a -> a
max k
k1 k
k2)


-- | Create an unbounded 'Range' which can be used for slicing. 
-- This function is expected to be used in conjunction with 'select'. 
--
-- For bound ranges, see 'to'.
from :: k -> Range k
from :: forall k. k -> Range k
from = k -> Range k
forall k. k -> Range k
From


-- | Create an unbounded 'Range' which can be used for slicing. This function
-- is expected to be used in conjunction with 'select'. 
--
-- For bound ranges, see 'to'.
upto :: k -> Range k
upto :: forall k. k -> Range k
upto = k -> Range k
forall k. k -> Range k
UpTo


-- | Class for datatypes which can be used to select sub-series using 'select'.
--
-- There are two use-cases for 'select':
--
--  * Bulk random-access (selecting from an 'Index' of keys);
--  * Bulk ordered access (selecting from a 'Range' of keys).
--
-- See the documentation for 'select'.
class Selection s where
    -- | Select a subseries. There are two main ways to do this.
    --
    -- The first way to do this is to select a sub-series based on keys:
    --
    -- >>> let xs = Series.fromList [('a', 10::Int), ('b', 20), ('c', 30), ('d', 40)]
    -- >>> xs `select` Index.fromList ['a', 'd']
    -- index | values
    -- ----- | ------
    --   'a' |     10
    --   'd' |     40
    --
    -- The second way to select a sub-series is to select all keys in a range:
    --
    -- >>> xs `select` 'b' `to` 'c'
    -- index | values
    -- ----- | ------
    --   'b' |     20
    --   'c' |     30
    --
    -- Such ranges can also be unbounded. (i.e. all keys smaller or larger than some key), like so:
    --
    -- >>> xs `select` upto 'c'
    -- index | values
    -- ----- | ------
    --   'a' |     10
    --   'b' |     20
    --   'c' |     30
    -- >>> xs `select` from 'c'
    -- index | values
    -- ----- | ------
    --   'c' |     30
    --   'd' |     40
    --
    -- Note that with 'select', you'll always get a sub-series; if you ask for a key which is not
    -- in the series, it'll be ignored:
    --
    -- >>> xs `select` Index.fromList ['a', 'd', 'e']
    -- index | values
    -- ----- | ------
    --   'a' |     10
    --   'd' |     40
    --
    -- See 'require' if you want to ensure that all keys are present.
    select :: (Vector v a, Ord k) => Series v k a -> s k -> Series v k a


instance Selection Index where
    -- | Select all keys in 'Index' from a series. Keys which are not
    -- in the series are ignored.
    select :: (Vector v a, Ord k) => Series v k a -> Index k -> Series v k a
    {-# INLINE select #-}
    select :: forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
select Series v k a
xs Index k
ss
        = let selectedKeys :: Index k
selectedKeys = Series v k a -> Index k
forall {k1} (v :: k1 -> *) k2 (a :: k1). Series v k2 a -> Index k2
index Series v k a
xs Index k -> Index k -> Index k
forall k. Ord k => Index k -> Index k -> Index k
`Index.intersection` Index k
ss
            -- Surprisingly, using `Vector.backpermute` does not
            -- perform as well as `Vector.map (Vector.unsafeIndex vs)`
            -- for large Series
           in Series v k a
xs Series v k a -> Index k -> Series v k a
forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
`selectSubset` Index k
selectedKeys

-- | Selecting a sub-series from a 'Set' is a convenience
-- function. Internally, the 'Set' is converted to an index first.
instance Selection Set where
    select :: (Vector v a, Ord k) => Series v k a -> Set k -> Series v k a
    {-# INLINE select #-}
    select :: forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Set k -> Series v k a
select Series v k a
xs = Series v k a -> Index k -> Series v k a
forall (s :: * -> *) (v :: * -> *) a k.
(Selection s, Vector v a, Ord k) =>
Series v k a -> s k -> Series v k a
forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
select Series v k a
xs (Index k -> Series v k a)
-> (Set k -> Index k) -> Set k -> Series v k a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set k -> Index k
forall k. Set k -> Index k
Index.fromSet


-- | Selecting a sub-series from a list is a convenience
-- function. Internally, the list is converted to an index first.
instance Selection [] where
    select :: (Vector v a, Ord k) => Series v k a -> [k] -> Series v k a
    {-# INLINE select #-}
    select :: forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> [k] -> Series v k a
select Series v k a
xs = Series v k a -> Index k -> Series v k a
forall (s :: * -> *) (v :: * -> *) a k.
(Selection s, Vector v a, Ord k) =>
Series v k a -> s k -> Series v k a
forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
select Series v k a
xs (Index k -> Series v k a)
-> ([k] -> Index k) -> [k] -> Series v k a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [k] -> Index k
forall k. Ord k => [k] -> Index k
Index.fromList


-- | Selecting a sub-series based on a @Range@ is most performant.
-- Constructing a @Range@ is most convenient using the 'to' function.
instance Selection Range where
    select :: (Vector v a, Ord k) => Series v k a -> Range k -> Series v k a
    {-# INLINE select #-}
    select :: forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Range k -> Series v k a
select Series v k a
series Range k
rng = case Series v k a -> Range k -> Maybe (k, k)
forall {k} k (v :: k -> *) (a :: k).
Ord k =>
Series v k a -> Range k -> Maybe (k, k)
keysInRange Series v k a
series Range k
rng of 
        Maybe (k, k)
Nothing              -> Series v k a
forall a. Monoid a => a
mempty
        Just (k
kstart, k
kstop) -> let indexOf :: Series v k a -> k -> Int
indexOf Series v k a
xs k
k = k -> Index k -> Int
forall k. (HasCallStack, Ord k) => k -> Index k -> Int
Index.Internal.findIndex k
k (Series v k a -> Index k
forall {k1} (v :: k1 -> *) k2 (a :: k1). Series v k2 a -> Index k2
index Series v k a
xs)
                                 in Int -> Int -> Series v k a -> Series v k a
forall (v :: * -> *) a k.
Vector v a =>
Int -> Int -> Series v k a -> Series v k a
slice (Series v k a
series Series v k a -> k -> Int
forall {k1} {k} {v :: k1 -> *} {a :: k1}.
Ord k =>
Series v k a -> k -> Int
`indexOf` k
kstart) (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Series v k a -> k -> Int
forall {k1} {k} {v :: k1 -> *} {a :: k1}.
Ord k =>
Series v k a -> k -> Int
indexOf Series v k a
series k
kstop) Series v k a
series


-- | Select a sub-series from a series matching a condition.
selectWhere :: (Vector v a, Vector v Int, Vector v Bool, Ord k) => Series v k a -> Series v k Bool -> Series v k a
{-# INLINE selectWhere #-}
selectWhere :: forall (v :: * -> *) a k.
(Vector v a, Vector v Int, Vector v Bool, Ord k) =>
Series v k a -> Series v k Bool -> Series v k a
selectWhere Series v k a
xs Series v k Bool
ys = Series v k a
xs Series v k a -> Index k -> Series v k a
forall (s :: * -> *) (v :: * -> *) a k.
(Selection s, Vector v a, Ord k) =>
Series v k a -> s k -> Series v k a
forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
`select` Set k -> Index k
forall k. Set k -> Index k
Index.fromSet Set k
keysWhereTrue
    where
        (MkSeries Index k
_ v Bool
cond) = Series v k Bool
ys Series v k Bool -> Index k -> Series v k Bool
forall (s :: * -> *) (v :: * -> *) a k.
(Selection s, Vector v a, Ord k) =>
Series v k a -> s k -> Series v k a
forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
`select` Series v k a -> Index k
forall {k1} (v :: k1 -> *) k2 (a :: k1). Series v k2 a -> Index k2
index Series v k a
xs
        whereValuesAreTrue :: Set Int
whereValuesAreTrue = [Int] -> Set Int
forall a. [a] -> Set a
Set.fromDistinctAscList ([Int] -> Set Int) -> [Int] -> Set Int
forall a b. (a -> b) -> a -> b
$ v Int -> [Int]
forall (v :: * -> *) a. Vector v a => v a -> [a]
Vector.toList ((Bool -> Bool) -> v Bool -> v Int
forall (v :: * -> *) a.
(Vector v a, Vector v Int) =>
(a -> Bool) -> v a -> v Int
Vector.findIndices Bool -> Bool
forall a. a -> a
id v Bool
cond)
        keysWhereTrue :: Set k
keysWhereTrue = (Int -> k) -> Set Int -> Set k
forall a b. (a -> b) -> Set a -> Set b
Set.mapMonotonic (Int -> Index k -> k
forall k. HasCallStack => Int -> Index k -> k
`Index.Internal.elemAt` Series v k a -> Index k
forall {k1} (v :: k1 -> *) k2 (a :: k1). Series v k2 a -> Index k2
index Series v k a
xs) Set Int
whereValuesAreTrue


-- | Implementation of `select` where the selection keys are known
-- to be a subset of the series. This precondition is NOT checked.
--
-- This is a performance optimization and therefore is not normally exposed.
selectSubset :: (Vector v a, Ord k) => Series v k a -> Index k -> Series v k a
{-# INLINE selectSubset #-}
selectSubset :: forall (v :: * -> *) a k.
(Vector v a, Ord k) =>
Series v k a -> Index k -> Series v k a
selectSubset (MkSeries Index k
ks v a
vs) Index k
ss
    -- TODO: 
    --   Is it possible to scan over the series once
    --   while filtering away on keys? Initial attempts did not lead
    --   to performance improvements, but I can't imagine that calling
    --   `Index.Internal.findIndex` repeatedly is efficient
    = 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
ss (v a -> Series v k a) -> v a -> Series v k a
forall a b. (a -> b) -> a -> b
$ Vector a -> v a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
Boxed.convert
                  (Vector a -> v a) -> Vector a -> v a
forall a b. (a -> b) -> a -> b
$ (k -> a) -> Vector k -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
Boxed.map (v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
Vector.unsafeIndex v a
vs (Int -> a) -> (k -> Int) -> k -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k -> Index k -> Int
forall k. (HasCallStack, Ord k) => k -> Index k -> Int
`Index.Internal.findIndex` Index k
ks))
                  (Vector k -> Vector a) -> Vector k -> Vector a
forall a b. (a -> b) -> a -> b
$ Index k -> Vector k
forall (v :: * -> *) k. Vector v k => Index k -> v k
Index.toAscVector Index k
ss


-- | Yield a subseries based on indices. The end index is not included.
slice :: Vector v a
      => Int -- ^ Start index
      -> Int -- ^ End index, which is not included
      -> Series v k a 
      -> Series v k a
{-# INLINE slice #-}
slice :: forall (v :: * -> *) a k.
Vector v a =>
Int -> Int -> Series v k a -> Series v k a
slice Int
start Int
stop (MkSeries Index k
ks v a
vs) 
    = let stop' :: Int
stop' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
Vector.length v a
vs) Int
stop
    in MkSeries { index :: Index k
index  = Int -> Index k -> Index k
forall k. Int -> Index k -> Index k
Index.take (Int
stop' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start) (Index k -> Index k) -> Index k -> Index k
forall a b. (a -> b) -> a -> b
$ Int -> Index k -> Index k
forall k. Int -> Index k -> Index k
Index.drop Int
start Index k
ks
                , values :: v a
values = Int -> Int -> v a -> v a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
Int -> Int -> v a -> v a
Vector.slice Int
start (Int
stop' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start) v a
vs
                }