{-| 
    Module      : Vocoder.Conduit.Frames
    Description : Frame processing
    Copyright   : (c) Marek Materzok, 2021
    License     : BSD2
-}
{-# LANGUAGE BangPatterns, FlexibleContexts #-}
module Vocoder.Conduit.Frames (
    framesOfE,
    genFramesOfE,
    sumFramesE
    ) where

import Control.Arrow
import Data.Conduit
import Data.MonoTraversable
import Data.Maybe(fromMaybe)
import qualified Data.Sequences as Seq

-- | Splits a chunked input stream into overlapping frames of constant size
--   suitable for STFT processing.
framesOfE :: (Monad m, Seq.IsSequence seq) => Seq.Index seq -> Seq.Index seq -> ConduitT seq seq m ()
framesOfE :: Index seq -> Index seq -> ConduitT seq seq m ()
framesOfE Index seq
chunkSize Index seq
hopSize = Index seq -> Index seq -> seq -> ConduitT seq seq m seq
forall (m :: * -> *) seq.
(Monad m, IsSequence seq) =>
Index seq -> Index seq -> seq -> ConduitT seq seq m seq
genFramesOfE Index seq
chunkSize Index seq
hopSize ([Element seq] -> seq
forall seq. IsSequence seq => [Element seq] -> seq
Seq.fromList []) ConduitT seq seq m seq
-> ConduitT seq seq m () -> ConduitT seq seq m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> ConduitT seq seq m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | More general version of framesOfE, suitable for processing multiple inputs.
genFramesOfE :: (Monad m, Seq.IsSequence seq) => Seq.Index seq -> Seq.Index seq -> seq -> ConduitT seq seq m seq
genFramesOfE :: Index seq -> Index seq -> seq -> ConduitT seq seq m seq
genFramesOfE Index seq
chunkSize Index seq
hopSize seq
q = do
    Maybe seq
mnextv <- ConduitT seq seq m (Maybe seq)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await
    case Maybe seq
mnextv of
        Maybe seq
Nothing -> seq -> ConduitT seq seq m seq
forall (m :: * -> *) a. Monad m => a -> m a
return seq
q
        Just seq
nextv -> do
            let newBuf :: seq
newBuf = seq
q seq -> seq -> seq
forall a. Monoid a => a -> a -> a
`mappend` seq
nextv
            let newBufLen :: Index seq
newBufLen = seq -> Index seq
forall seq. IsSequence seq => seq -> Index seq
Seq.lengthIndex seq
newBuf
            (seq -> ConduitT seq seq m ()) -> [seq] -> ConduitT seq seq m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ seq -> ConduitT seq seq m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield [Index seq -> seq -> seq
forall seq. IsSequence seq => Index seq -> seq -> seq
Seq.take Index seq
chunkSize (seq -> seq) -> seq -> seq
forall a b. (a -> b) -> a -> b
$ Index seq -> seq -> seq
forall seq. IsSequence seq => Index seq -> seq -> seq
Seq.drop Index seq
k seq
newBuf
                        | Index seq
k <- [Index seq
0, Index seq
hopSize .. Index seq
newBufLen Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
- Index seq
chunkSize]]
            let dropcnt :: Index seq
dropcnt = ((Index seq
newBufLen Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
- Index seq
chunkSize) Index seq -> Index seq -> Index seq
forall a. Integral a => a -> a -> a
`div` Index seq
hopSize) Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
* Index seq
hopSize Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
+ Index seq
hopSize
            let q' :: seq
q' = Index seq -> seq -> seq
forall seq. IsSequence seq => Index seq -> seq -> seq
Seq.drop Index seq
dropcnt seq
newBuf
            Index seq -> Index seq -> seq -> ConduitT seq seq m seq
forall (m :: * -> *) seq.
(Monad m, IsSequence seq) =>
Index seq -> Index seq -> seq -> ConduitT seq seq m seq
genFramesOfE Index seq
chunkSize Index seq
hopSize seq
q'

-- | Builds a chunked output stream from a stream of overlapping frames.
sumFramesE :: (Monad m, Seq.IsSequence seq, Num (Element seq)) => Seq.Index seq -> Seq.Index seq -> ConduitT seq seq m ()
sumFramesE :: Index seq -> Index seq -> ConduitT seq seq m ()
sumFramesE Index seq
chunkSize Index seq
hopSize = Index seq -> [(Index seq, seq)] -> ConduitT seq seq m ()
forall (m :: * -> *).
Monad m =>
Index seq -> [(Index seq, seq)] -> ConduitT seq seq m ()
process Index seq
0 []
    where
        ith :: Index seq -> (Index seq, seq) -> Element seq
ith Index seq
i (Index seq
n, seq
c0) = Element seq -> Maybe (Element seq) -> Element seq
forall a. a -> Maybe a -> a
fromMaybe Element seq
0 (Maybe (Element seq) -> Element seq)
-> Maybe (Element seq) -> Element seq
forall a b. (a -> b) -> a -> b
$ seq -> Index seq -> Maybe (Element seq)
forall seq.
IsSequence seq =>
seq -> Index seq -> Maybe (Element seq)
Seq.index seq
c0 (Index seq
i Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
- Index seq
n)
        publish :: t (Index seq, seq) -> ConduitT i seq m ()
publish t (Index seq, seq)
q = seq -> ConduitT i seq m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield (seq -> ConduitT i seq m ()) -> seq -> ConduitT i seq m ()
forall a b. (a -> b) -> a -> b
$ [Element seq] -> seq
forall seq. IsSequence seq => [Element seq] -> seq
Seq.fromList ([Element seq] -> seq) -> [Element seq] -> seq
forall a b. (a -> b) -> a -> b
$ (Index seq -> Element seq) -> [Index seq] -> [Element seq]
forall a b. (a -> b) -> [a] -> [b]
map (\Index seq
i -> t (Element seq) -> Element seq
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (t (Element seq) -> Element seq) -> t (Element seq) -> Element seq
forall a b. (a -> b) -> a -> b
$ ((Index seq, seq) -> Element seq)
-> t (Index seq, seq) -> t (Element seq)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Index seq -> (Index seq, seq) -> Element seq
forall seq.
(IsSequence seq, Num (Element seq)) =>
Index seq -> (Index seq, seq) -> Element seq
ith Index seq
i) t (Index seq, seq)
q) [Index seq
0 .. Index seq
chunkSizeIndex seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
-Index seq
1]
        publishRest :: [(Index seq, seq)] -> ConduitT i seq m ()
publishRest [(Index seq, seq)]
q | [(Index seq, seq)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Index seq, seq)]
q    = () -> ConduitT i seq m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                      | Bool
otherwise = [(Index seq, seq)] -> ConduitT i seq m ()
forall (m :: * -> *) (t :: * -> *) i.
(Monad m, Foldable t, Functor t) =>
t (Index seq, seq) -> ConduitT i seq m ()
publish [(Index seq, seq)]
q ConduitT i seq m () -> ConduitT i seq m () -> ConduitT i seq m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [(Index seq, seq)] -> ConduitT i seq m ()
publishRest ([(Index seq, seq)] -> [(Index seq, seq)]
nextq [(Index seq, seq)]
q)
        nextq :: [(Index seq, seq)] -> [(Index seq, seq)]
nextq [(Index seq, seq)]
q = ((Index seq, seq) -> (Index seq, seq))
-> [(Index seq, seq)] -> [(Index seq, seq)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
+ (-Index seq
chunkSize)) (Index seq -> Index seq)
-> (seq -> seq) -> (Index seq, seq) -> (Index seq, seq)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** seq -> seq
forall a. a -> a
id) ([(Index seq, seq)] -> [(Index seq, seq)])
-> [(Index seq, seq)] -> [(Index seq, seq)]
forall a b. (a -> b) -> a -> b
$ ((Index seq, seq) -> Bool)
-> [(Index seq, seq)] -> [(Index seq, seq)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (\(Index seq
n, seq
c) -> seq -> Index seq
forall seq. IsSequence seq => seq -> Index seq
Seq.lengthIndex seq
c Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
+ Index seq
n Index seq -> Index seq -> Bool
forall a. Ord a => a -> a -> Bool
<= Index seq
chunkSize) [(Index seq, seq)]
q
        process2 :: Index seq -> [(Index seq, seq)] -> ConduitT seq seq m ()
process2 Index seq
sofar [(Index seq, seq)]
q
            | Index seq
sofar Index seq -> Index seq -> Bool
forall a. Ord a => a -> a -> Bool
>= Index seq
chunkSize = do
                [(Index seq, seq)] -> ConduitT seq seq m ()
forall (m :: * -> *) (t :: * -> *) i.
(Monad m, Foldable t, Functor t) =>
t (Index seq, seq) -> ConduitT i seq m ()
publish [(Index seq, seq)]
q
                Index seq -> [(Index seq, seq)] -> ConduitT seq seq m ()
process2 (Index seq
sofar Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
- Index seq
chunkSize) ([(Index seq, seq)] -> ConduitT seq seq m ())
-> [(Index seq, seq)] -> ConduitT seq seq m ()
forall a b. (a -> b) -> a -> b
$ [(Index seq, seq)] -> [(Index seq, seq)]
nextq [(Index seq, seq)]
q
            | Bool
otherwise = Index seq -> [(Index seq, seq)] -> ConduitT seq seq m ()
process (Index seq
sofar Index seq -> Index seq -> Index seq
forall a. Num a => a -> a -> a
+ Index seq
hopSize) [(Index seq, seq)]
q
        process :: Index seq -> [(Index seq, seq)] -> ConduitT seq seq m ()
process Index seq
sofar [(Index seq, seq)]
q = do
            Maybe seq
next <- ConduitT seq seq m (Maybe seq)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await
            case Maybe seq
next of
                Maybe seq
Nothing -> [(Index seq, seq)] -> ConduitT seq seq m ()
forall (m :: * -> *) i.
Monad m =>
[(Index seq, seq)] -> ConduitT i seq m ()
publishRest [(Index seq, seq)]
q
                Just seq
next' -> Index seq -> [(Index seq, seq)] -> ConduitT seq seq m ()
process2 Index seq
sofar ([(Index seq, seq)]
q [(Index seq, seq)] -> [(Index seq, seq)] -> [(Index seq, seq)]
forall a. [a] -> [a] -> [a]
++ [(Index seq
sofar, seq
next')])