{-| 
    Module      : Vocoder.Audio
    Description : Frequency-domain filters
    Copyright   : (c) Marek Materzok, 2021
    License     : BSD2

This module allows easy frequency-domain processing on audio streams
created in @conduit-audio@.
-}
module Vocoder.Audio(
    VocoderAudioSource(..),
    concatenateV,
    sourceVocoder,
    sourceVocoderWithPhase,
    processAudio,
    processAudioWithPhase,
    processVocoderAudio
  ) where

import Data.Conduit
import Data.Conduit.Audio
import qualified Data.Conduit.Combinators as DCC
import Control.Applicative
import Control.Monad
import Vocoder
import Vocoder.Conduit
import Vocoder.Conduit.Filter
import Vocoder.Conduit.Frames
import qualified Data.Vector.Storable as V

data VocoderAudioSource m = VocoderAudioSource {
    VocoderAudioSource m
-> (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
sourceV :: (Frame, (ZipList Phase, ZipList Phase)) 
            -> ConduitT () Frame m (V.Vector Double, (ZipList Phase, ZipList Phase)),
    VocoderAudioSource m -> Rate
rateV :: Rate,
    VocoderAudioSource m -> Channels
channelsV :: Channels,
    VocoderAudioSource m -> Channels
framesV :: Frames,
    VocoderAudioSource m -> VocoderParams
vocoderParamsV :: VocoderParams
}

-- | Applies a conduit filter to an audio stream, producing a vocoder stream.
--   This allows to seamlessly concatenate audio streams for vocoder processing.
processVocoderAudio 
    :: Monad m
    => VocoderParams
    -> Filter m
    -> AudioSource m Double 
    -> VocoderAudioSource m
processVocoderAudio :: VocoderParams
-> Filter m -> AudioSource m Rate -> VocoderAudioSource m
processVocoderAudio VocoderParams
par Filter m
c AudioSource m Rate
src = ((Frame, (ZipList Frame, ZipList Frame))
 -> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame)))
-> Rate
-> Channels
-> Channels
-> VocoderParams
-> VocoderAudioSource m
forall (m :: * -> *).
((Frame, (ZipList Frame, ZipList Frame))
 -> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame)))
-> Rate
-> Channels
-> Channels
-> VocoderParams
-> VocoderAudioSource m
VocoderAudioSource (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
newSource (AudioSource m Rate -> Rate
forall (m :: * -> *) a. AudioSource m a -> Rate
rate AudioSource m Rate
src) (AudioSource m Rate -> Channels
forall (m :: * -> *) a. AudioSource m a -> Channels
channels AudioSource m Rate
src) (AudioSource m Rate -> Channels
forall (m :: * -> *) a. AudioSource m a -> Channels
frames AudioSource m Rate
src) VocoderParams
par where
    freqStep :: Rate
freqStep = AudioSource m Rate -> Rate
forall (m :: * -> *) a. AudioSource m a -> Rate
rate AudioSource m Rate
src Rate -> Rate -> Rate
forall a. Fractional a => a -> a -> a
/ Channels -> Rate
forall a b. (Integral a, Num b) => a -> b
fromIntegral (VocoderParams -> Channels
vocFrameLength VocoderParams
par)
    newSource :: (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
newSource (Frame
q, (ZipList Frame, ZipList Frame)
ps) = 
        (AudioSource m Rate -> Source m Frame
forall (m :: * -> *) a. AudioSource m a -> Source m (Vector a)
source AudioSource m Rate
src Source m Frame
-> ConduitM Frame Frame m Frame -> ConduitM () Frame m Frame
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| Index Frame -> Index Frame -> Frame -> ConduitM Frame Frame m Frame
forall (m :: * -> *) seq.
(Monad m, IsSequence seq) =>
Index seq -> Index seq -> seq -> ConduitT seq seq m seq
genFramesOfE (VocoderParams -> Channels
vocInputFrameLength VocoderParams
par Channels -> Channels -> Channels
forall a. Num a => a -> a -> a
* AudioSource m Rate -> Channels
forall (m :: * -> *) a. AudioSource m a -> Channels
channels AudioSource m Rate
src) (VocoderParams -> Channels
vocHopSize VocoderParams
par Channels -> Channels -> Channels
forall a. Num a => a -> a -> a
* AudioSource m Rate -> Channels
forall (m :: * -> *) a. AudioSource m a -> Channels
channels AudioSource m Rate
src) Frame
q)
        ConduitM () Frame m Frame
-> ConduitT Frame (ZipList Frame) m (ZipList Frame, ZipList Frame)
-> ConduitT
     () (ZipList Frame) m (Frame, (ZipList Frame, ZipList Frame))
forall (m :: * -> *) a b r1 c r2.
Monad m =>
ConduitT a b m r1 -> ConduitT b c m r2 -> ConduitT a c m (r1, r2)
`fuseBoth`
        ((Frame -> ZipList Frame) -> ConduitT Frame (ZipList Frame) m ()
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
DCC.map ([Frame] -> ZipList Frame
forall a. [a] -> ZipList a
ZipList ([Frame] -> ZipList Frame)
-> (Frame -> [Frame]) -> Frame -> ZipList Frame
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Channels -> Frame -> [Frame]
forall a. Storable a => Channels -> Vector a -> [Vector a]
deinterleave (AudioSource m Rate -> Channels
forall (m :: * -> *) a. AudioSource m a -> Channels
channels AudioSource m Rate
src)) ConduitT Frame (ZipList Frame) m ()
-> ConduitM
     (ZipList Frame) (ZipList Frame) m (ZipList Frame, ZipList Frame)
-> ConduitT Frame (ZipList Frame) m (ZipList Frame, ZipList Frame)
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| (((), (ZipList Frame, ZipList Frame))
-> (ZipList Frame, ZipList Frame)
forall a b. (a, b) -> b
snd (((), (ZipList Frame, ZipList Frame))
 -> (ZipList Frame, ZipList Frame))
-> ConduitT
     (ZipList Frame)
     (ZipList Frame)
     m
     ((), (ZipList Frame, ZipList Frame))
-> ConduitM
     (ZipList Frame) (ZipList Frame) m (ZipList Frame, ZipList Frame)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VocoderParams
-> (ZipList Frame, ZipList Frame)
-> ConduitT (ZipList STFTFrame) (ZipList STFTFrame) m ()
-> ConduitT
     (ZipList Frame)
     (ZipList Frame)
     m
     ((), (ZipList Frame, ZipList Frame))
forall (f :: * -> *) (m :: * -> *) r.
(Applicative f, Monad m) =>
VocoderParams
-> (f Frame, f Frame)
-> ConduitT (f STFTFrame) (f STFTFrame) m r
-> ConduitT (f Frame) (f Frame) m (r, (f Frame, f Frame))
processFramesF VocoderParams
par (ZipList Frame, ZipList Frame)
ps (Filter m
-> Rate -> ConduitT (ZipList STFTFrame) (ZipList STFTFrame) m ()
forall (m :: * -> *).
Filter m
-> forall (f :: * -> *).
   Traversable f =>
   Rate -> ConduitT (f STFTFrame) (f STFTFrame) m ()
runFilter Filter m
c Rate
freqStep)))
        ConduitT
  () (ZipList Frame) m (Frame, (ZipList Frame, ZipList Frame))
-> Conduit (ZipList Frame) m Frame
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
forall (m :: * -> *) a b r c.
Monad m =>
ConduitT a b m r -> Conduit b m c -> ConduitT a c m r
`fuseUpstream`
        (ZipList Frame -> Frame) -> Conduit (ZipList Frame) m Frame
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
DCC.map ([Frame] -> Frame
forall a. Storable a => [Vector a] -> Vector a
interleave ([Frame] -> Frame)
-> (ZipList Frame -> [Frame]) -> ZipList Frame -> Frame
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ZipList Frame -> [Frame]
forall a. ZipList a -> [a]
getZipList)

-- | Connects the end of the first vocoder source to the beginning of the second. 
--   The two sources must have the same sample rate, channel count, vocoder hop size
--   and frame length.
concatenateV :: Monad m 
             => VocoderAudioSource m
             -> VocoderAudioSource m
             -> VocoderAudioSource m
concatenateV :: VocoderAudioSource m
-> VocoderAudioSource m -> VocoderAudioSource m
concatenateV VocoderAudioSource m
src1 VocoderAudioSource m
src2 
    | VocoderAudioSource m -> Rate
forall (m :: * -> *). VocoderAudioSource m -> Rate
rateV VocoderAudioSource m
src1          Rate -> Rate -> Bool
forall a. Eq a => a -> a -> Bool
/= VocoderAudioSource m -> Rate
forall (m :: * -> *). VocoderAudioSource m -> Rate
rateV VocoderAudioSource m
src2         = [Char] -> VocoderAudioSource m
forall a. HasCallStack => [Char] -> a
error [Char]
"Vocoder.Audio.concatenateV: mismatched rates"
    | VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
channelsV VocoderAudioSource m
src1      Channels -> Channels -> Bool
forall a. Eq a => a -> a -> Bool
/= VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
channelsV VocoderAudioSource m
src2     = [Char] -> VocoderAudioSource m
forall a. HasCallStack => [Char] -> a
error [Char]
"Vocoder.Audio.concatenateV: mismatched channels"
    | VocoderParams -> Channels
vocHopSize VocoderParams
par1     Channels -> Channels -> Bool
forall a. Eq a => a -> a -> Bool
/= VocoderParams -> Channels
vocHopSize VocoderParams
par2    = [Char] -> VocoderAudioSource m
forall a. HasCallStack => [Char] -> a
error [Char]
"Vocoder.Audio.concatenateV: mismatched hop size"
    | VocoderParams -> Channels
vocFrameLength VocoderParams
par1 Channels -> Channels -> Bool
forall a. Eq a => a -> a -> Bool
/= VocoderParams -> Channels
vocFrameLength VocoderParams
par2 = [Char] -> VocoderAudioSource m
forall a. HasCallStack => [Char] -> a
error [Char]
"Vocoder.Audio.concatenateV: mismatched frame length"
    | Bool
otherwise = ((Frame, (ZipList Frame, ZipList Frame))
 -> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame)))
-> Rate
-> Channels
-> Channels
-> VocoderParams
-> VocoderAudioSource m
forall (m :: * -> *).
((Frame, (ZipList Frame, ZipList Frame))
 -> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame)))
-> Rate
-> Channels
-> Channels
-> VocoderParams
-> VocoderAudioSource m
VocoderAudioSource (VocoderAudioSource m
-> (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
forall (m :: * -> *).
VocoderAudioSource m
-> (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
sourceV VocoderAudioSource m
src1 ((Frame, (ZipList Frame, ZipList Frame))
 -> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame)))
-> ((Frame, (ZipList Frame, ZipList Frame))
    -> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame)))
-> (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> VocoderAudioSource m
-> (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
forall (m :: * -> *).
VocoderAudioSource m
-> (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
sourceV VocoderAudioSource m
src2) (VocoderAudioSource m -> Rate
forall (m :: * -> *). VocoderAudioSource m -> Rate
rateV VocoderAudioSource m
src1) (VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
channelsV VocoderAudioSource m
src1) (VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
framesV VocoderAudioSource m
src1 Channels -> Channels -> Channels
forall a. Num a => a -> a -> a
+ VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
framesV VocoderAudioSource m
src2) (VocoderAudioSource m -> VocoderParams
forall (m :: * -> *). VocoderAudioSource m -> VocoderParams
vocoderParamsV VocoderAudioSource m
src1)
    where
    par1 :: VocoderParams
par1 = VocoderAudioSource m -> VocoderParams
forall (m :: * -> *). VocoderAudioSource m -> VocoderParams
vocoderParamsV VocoderAudioSource m
src1
    par2 :: VocoderParams
par2 = VocoderAudioSource m -> VocoderParams
forall (m :: * -> *). VocoderAudioSource m -> VocoderParams
vocoderParamsV VocoderAudioSource m
src2

-- | Creates an audio source from a vocoder source.
sourceVocoder :: Monad m 
              => VocoderAudioSource m
              -> AudioSource m Double
sourceVocoder :: VocoderAudioSource m -> AudioSource m Rate
sourceVocoder VocoderAudioSource m
src = Frame -> VocoderAudioSource m -> AudioSource m Rate
forall (m :: * -> *).
Monad m =>
Frame -> VocoderAudioSource m -> AudioSource m Rate
sourceVocoderWithPhase (VocoderParams -> Frame
zeroPhase (VocoderParams -> Frame) -> VocoderParams -> Frame
forall a b. (a -> b) -> a -> b
$ VocoderAudioSource m -> VocoderParams
forall (m :: * -> *). VocoderAudioSource m -> VocoderParams
vocoderParamsV VocoderAudioSource m
src) VocoderAudioSource m
src

-- | Creates an audio source from a vocoder source, with initial phase provided.
sourceVocoderWithPhase 
    :: Monad m 
    => Phase
    -> VocoderAudioSource m
    -> AudioSource m Double
sourceVocoderWithPhase :: Frame -> VocoderAudioSource m -> AudioSource m Rate
sourceVocoderWithPhase Frame
iphs VocoderAudioSource m
src = Source m Frame
-> Rate -> Channels -> Channels -> AudioSource m Rate
forall (m :: * -> *) a.
Source m (Vector a)
-> Rate -> Channels -> Channels -> AudioSource m a
AudioSource Source m Frame
newSource (VocoderAudioSource m -> Rate
forall (m :: * -> *). VocoderAudioSource m -> Rate
rateV VocoderAudioSource m
src) (VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
channelsV VocoderAudioSource m
src) (VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
framesV VocoderAudioSource m
src)
    where
    par :: VocoderParams
par = VocoderAudioSource m -> VocoderParams
forall (m :: * -> *). VocoderAudioSource m -> VocoderParams
vocoderParamsV VocoderAudioSource m
src
    phs :: ZipList Frame
phs = [Frame] -> ZipList Frame
forall a. [a] -> ZipList a
ZipList ([Frame] -> ZipList Frame) -> [Frame] -> ZipList Frame
forall a b. (a -> b) -> a -> b
$ Channels -> Frame -> [Frame]
forall a. Channels -> a -> [a]
replicate (VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
channelsV VocoderAudioSource m
src) Frame
iphs
    newSource :: Source m Frame
newSource = (VocoderAudioSource m
-> (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
forall (m :: * -> *).
VocoderAudioSource m
-> (Frame, (ZipList Frame, ZipList Frame))
-> ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
sourceV VocoderAudioSource m
src (Frame
forall a. Storable a => Vector a
V.empty, (ZipList Frame
phs, ZipList Frame
phs)) ConduitT () Frame m (Frame, (ZipList Frame, ZipList Frame))
-> Source m Frame -> Source m Frame
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> Source m Frame
forall (m :: * -> *) a. Monad m => a -> m a
return ())
             Source m Frame -> ConduitM Frame Frame m () -> Source m Frame
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| Index Frame -> Index Frame -> ConduitM Frame Frame m ()
forall (m :: * -> *) seq.
(Monad m, IsSequence seq, Num (Element seq)) =>
Index seq -> Index seq -> ConduitT seq seq m ()
sumFramesE (Channels
chunkSize Channels -> Channels -> Channels
forall a. Num a => a -> a -> a
* VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
channelsV VocoderAudioSource m
src) (VocoderParams -> Channels
vocHopSize VocoderParams
par Channels -> Channels -> Channels
forall a. Num a => a -> a -> a
* VocoderAudioSource m -> Channels
forall (m :: * -> *). VocoderAudioSource m -> Channels
channelsV VocoderAudioSource m
src)

-- | Applies a conduit filter to an audio stream.
processAudio 
    :: Monad m
    => VocoderParams
    -> Filter m
    -> AudioSource m Double 
    -> AudioSource m Double
processAudio :: VocoderParams
-> Filter m -> AudioSource m Rate -> AudioSource m Rate
processAudio VocoderParams
par = VocoderParams
-> Frame -> Filter m -> AudioSource m Rate -> AudioSource m Rate
forall (m :: * -> *).
Monad m =>
VocoderParams
-> Frame -> Filter m -> AudioSource m Rate -> AudioSource m Rate
processAudioWithPhase VocoderParams
par (VocoderParams -> Frame
zeroPhase VocoderParams
par)

-- | Applies a conduit filter to an audio stream, with initial phase provided.
processAudioWithPhase
    :: Monad m
    => VocoderParams
    -> Phase
    -> Filter m
    -> AudioSource m Double 
    -> AudioSource m Double
processAudioWithPhase :: VocoderParams
-> Frame -> Filter m -> AudioSource m Rate -> AudioSource m Rate
processAudioWithPhase VocoderParams
par Frame
iphs Filter m
c AudioSource m Rate
src = Frame -> VocoderAudioSource m -> AudioSource m Rate
forall (m :: * -> *).
Monad m =>
Frame -> VocoderAudioSource m -> AudioSource m Rate
sourceVocoderWithPhase Frame
iphs (VocoderAudioSource m -> AudioSource m Rate)
-> VocoderAudioSource m -> AudioSource m Rate
forall a b. (a -> b) -> a -> b
$ VocoderParams
-> Filter m -> AudioSource m Rate -> VocoderAudioSource m
forall (m :: * -> *).
Monad m =>
VocoderParams
-> Filter m -> AudioSource m Rate -> VocoderAudioSource m
processVocoderAudio VocoderParams
par Filter m
c AudioSource m Rate
src