{-| 
    Module      : Vocoder.Conduit
    Description : Phase vocoder in Conduit
    Copyright   : (c) Marek Materzok, 2021
    License     : BSD2

This module wraps phase vocoder algorithms for use in Conduit.
Two variants are provided, one for single channel processing,
and another for processing multiple channels synchronously.
-}
module Vocoder.Conduit(
      -- * Single-channel functions
      volumeFix,
      analysis,
      synthesis,
      processFrames,
      -- * Multi-channel functions
      volumeFixF,
      analysisF,
      synthesisF,
      processFramesF
    ) where

import Data.Conduit
import qualified Data.Conduit.List as DCL
import qualified Data.List.NonEmpty as DLN
import qualified Data.Vector.Storable as V
import Control.Arrow
import Vocoder

-- | Corrects for volume change introduced by STFT processing.
volumeFix :: Monad m => VocoderParams -> ConduitT STFTFrame STFTFrame m ()
volumeFix :: VocoderParams -> ConduitT STFTFrame STFTFrame m ()
volumeFix VocoderParams
par = (STFTFrame -> STFTFrame) -> ConduitT STFTFrame STFTFrame m ()
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
DCL.map ((STFTFrame -> STFTFrame) -> ConduitT STFTFrame STFTFrame m ())
-> (STFTFrame -> STFTFrame) -> ConduitT STFTFrame STFTFrame m ()
forall a b. (a -> b) -> a -> b
$ (Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
V.map (Double -> Double -> Double
forall a. Num a => a -> a -> a
* VocoderParams -> Double
volumeCoeff VocoderParams
par) (Vector Double -> Vector Double)
-> (Vector Double -> Vector Double) -> STFTFrame -> STFTFrame
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Vector Double -> Vector Double
forall a. a -> a
id

-- | Perform the phase vocoder analysis phase.
analysis :: Monad m => VocoderParams -> Phase -> ConduitT Frame STFTFrame m Phase
analysis :: VocoderParams
-> Vector Double
-> ConduitT (Vector Double) STFTFrame m (Vector Double)
analysis VocoderParams
par Vector Double
ph = (Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> Vector Double
-> ConduitT (Vector Double) STFTFrame m (Vector Double)
forall (m :: * -> *) a s b.
Monad m =>
(a -> s -> (s, b)) -> s -> ConduitT a b m s
DCL.mapAccum ((Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> Vector Double -> Vector Double -> (Vector Double, STFTFrame)
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Vector Double -> Vector Double -> (Vector Double, STFTFrame))
 -> Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> (Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> Vector Double
-> Vector Double
-> (Vector Double, STFTFrame)
forall a b. (a -> b) -> a -> b
$ VocoderParams
-> Vector Double -> Vector Double -> (Vector Double, STFTFrame)
analysisBlock VocoderParams
par) Vector Double
ph

-- | Perform the phase vocoder synthesis phase.
synthesis :: Monad m => VocoderParams -> Phase -> ConduitT STFTFrame Frame m Phase
synthesis :: VocoderParams
-> Vector Double
-> ConduitT STFTFrame (Vector Double) m (Vector Double)
synthesis VocoderParams
par Vector Double
ph = (STFTFrame -> Vector Double -> STFTFrame)
-> Vector Double
-> ConduitT STFTFrame (Vector Double) m (Vector Double)
forall (m :: * -> *) a s b.
Monad m =>
(a -> s -> (s, b)) -> s -> ConduitT a b m s
DCL.mapAccum ((Vector Double -> STFTFrame -> STFTFrame)
-> STFTFrame -> Vector Double -> STFTFrame
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Vector Double -> STFTFrame -> STFTFrame)
 -> STFTFrame -> Vector Double -> STFTFrame)
-> (Vector Double -> STFTFrame -> STFTFrame)
-> STFTFrame
-> Vector Double
-> STFTFrame
forall a b. (a -> b) -> a -> b
$ VocoderParams -> Vector Double -> STFTFrame -> STFTFrame
synthesisBlock VocoderParams
par) Vector Double
ph

-- | Perform frequency domain processing.
processFrames :: Monad m => VocoderParams -> (Phase, Phase) -> ConduitT STFTFrame STFTFrame m r -> ConduitT Frame Frame m (r, (Phase, Phase))
processFrames :: VocoderParams
-> STFTFrame
-> ConduitT STFTFrame STFTFrame m r
-> ConduitT (Vector Double) (Vector Double) m (r, STFTFrame)
processFrames VocoderParams
par (Vector Double
p1, Vector Double
p2) ConduitT STFTFrame STFTFrame m r
c = (\((Vector Double
p1', r
r), Vector Double
p2') -> (r
r, (Vector Double
p1', Vector Double
p2'))) (((Vector Double, r), Vector Double) -> (r, STFTFrame))
-> ConduitT
     (Vector Double)
     (Vector Double)
     m
     ((Vector Double, r), Vector Double)
-> ConduitT (Vector Double) (Vector Double) m (r, STFTFrame)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VocoderParams
-> Vector Double
-> ConduitT (Vector Double) STFTFrame m (Vector Double)
forall (m :: * -> *).
Monad m =>
VocoderParams
-> Vector Double
-> ConduitT (Vector Double) STFTFrame m (Vector Double)
analysis VocoderParams
par Vector Double
p1 ConduitT (Vector Double) STFTFrame m (Vector Double)
-> ConduitT STFTFrame STFTFrame m r
-> ConduitT (Vector Double) STFTFrame m (Vector Double, r)
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` (VocoderParams -> ConduitT STFTFrame STFTFrame m ()
forall (m :: * -> *).
Monad m =>
VocoderParams -> ConduitT STFTFrame STFTFrame m ()
volumeFix VocoderParams
par ConduitT STFTFrame STFTFrame m ()
-> ConduitT STFTFrame STFTFrame m r
-> ConduitT STFTFrame STFTFrame m r
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitT STFTFrame STFTFrame m r
c) ConduitT (Vector Double) STFTFrame m (Vector Double, r)
-> ConduitT STFTFrame (Vector Double) m (Vector Double)
-> ConduitT
     (Vector Double)
     (Vector Double)
     m
     ((Vector Double, r), Vector Double)
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` VocoderParams
-> Vector Double
-> ConduitT STFTFrame (Vector Double) m (Vector Double)
forall (m :: * -> *).
Monad m =>
VocoderParams
-> Vector Double
-> ConduitT STFTFrame (Vector Double) m (Vector Double)
synthesis VocoderParams
par Vector Double
p2

app_help :: Applicative f => (a -> s -> (s, b)) -> f a -> f s -> (f s, f b)
app_help :: (a -> s -> (s, b)) -> f a -> f s -> (f s, f b)
app_help a -> s -> (s, b)
f f a
a f s
b = f (s, b) -> (f s, f b)
forall (f :: * -> *) a b. Functor f => f (a, b) -> (f a, f b)
DLN.unzip (f (s, b) -> (f s, f b)) -> f (s, b) -> (f s, f b)
forall a b. (a -> b) -> a -> b
$ ((a, s) -> (s, b)) -> f (a, s) -> f (s, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> s -> (s, b)) -> (a, s) -> (s, b)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> s -> (s, b)
f) ((,) (a -> s -> (a, s)) -> f a -> f (s -> (a, s))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
a f (s -> (a, s)) -> f s -> f (a, s)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f s
b)

-- | Corrects for volume change introduced by STFT processing. 
volumeFixF :: (Applicative f, Monad m) => VocoderParams -> ConduitT (f STFTFrame) (f STFTFrame) m ()
volumeFixF :: VocoderParams -> ConduitT (f STFTFrame) (f STFTFrame) m ()
volumeFixF VocoderParams
par = (f STFTFrame -> f STFTFrame)
-> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
DCL.map ((f STFTFrame -> f STFTFrame)
 -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> (f STFTFrame -> f STFTFrame)
-> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall a b. (a -> b) -> a -> b
$ (STFTFrame -> STFTFrame) -> f STFTFrame -> f STFTFrame
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((STFTFrame -> STFTFrame) -> f STFTFrame -> f STFTFrame)
-> (STFTFrame -> STFTFrame) -> f STFTFrame -> f STFTFrame
forall a b. (a -> b) -> a -> b
$ (Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
V.map (Double -> Double -> Double
forall a. Num a => a -> a -> a
* VocoderParams -> Double
volumeCoeff VocoderParams
par) (Vector Double -> Vector Double)
-> (Vector Double -> Vector Double) -> STFTFrame -> STFTFrame
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Vector Double -> Vector Double
forall a. a -> a
id

-- | Perform the phase vocoder analysis phase.
analysisF :: (Applicative f, Monad m) => VocoderParams -> f Phase -> ConduitT (f Frame) (f STFTFrame) m (f Phase)
analysisF :: VocoderParams
-> f (Vector Double)
-> ConduitT (f (Vector Double)) (f STFTFrame) m (f (Vector Double))
analysisF VocoderParams
par f (Vector Double)
ph = (f (Vector Double)
 -> f (Vector Double) -> (f (Vector Double), f STFTFrame))
-> f (Vector Double)
-> ConduitT (f (Vector Double)) (f STFTFrame) m (f (Vector Double))
forall (m :: * -> *) a s b.
Monad m =>
(a -> s -> (s, b)) -> s -> ConduitT a b m s
DCL.mapAccum ((Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> f (Vector Double)
-> f (Vector Double)
-> (f (Vector Double), f STFTFrame)
forall (f :: * -> *) a s b.
Applicative f =>
(a -> s -> (s, b)) -> f a -> f s -> (f s, f b)
app_help ((Vector Double -> Vector Double -> (Vector Double, STFTFrame))
 -> f (Vector Double)
 -> f (Vector Double)
 -> (f (Vector Double), f STFTFrame))
-> (Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> f (Vector Double)
-> f (Vector Double)
-> (f (Vector Double), f STFTFrame)
forall a b. (a -> b) -> a -> b
$ (Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> Vector Double -> Vector Double -> (Vector Double, STFTFrame)
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Vector Double -> Vector Double -> (Vector Double, STFTFrame))
 -> Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> (Vector Double -> Vector Double -> (Vector Double, STFTFrame))
-> Vector Double
-> Vector Double
-> (Vector Double, STFTFrame)
forall a b. (a -> b) -> a -> b
$ VocoderParams
-> Vector Double -> Vector Double -> (Vector Double, STFTFrame)
analysisBlock VocoderParams
par) f (Vector Double)
ph

-- | Perform the phase vocoder synthesis phase.
synthesisF :: (Applicative f, Monad m) => VocoderParams -> f Phase -> ConduitT (f STFTFrame) (f Frame) m (f Phase)
synthesisF :: VocoderParams
-> f (Vector Double)
-> ConduitT (f STFTFrame) (f (Vector Double)) m (f (Vector Double))
synthesisF VocoderParams
par f (Vector Double)
ph = (f STFTFrame
 -> f (Vector Double) -> (f (Vector Double), f (Vector Double)))
-> f (Vector Double)
-> ConduitT (f STFTFrame) (f (Vector Double)) m (f (Vector Double))
forall (m :: * -> *) a s b.
Monad m =>
(a -> s -> (s, b)) -> s -> ConduitT a b m s
DCL.mapAccum ((STFTFrame -> Vector Double -> STFTFrame)
-> f STFTFrame
-> f (Vector Double)
-> (f (Vector Double), f (Vector Double))
forall (f :: * -> *) a s b.
Applicative f =>
(a -> s -> (s, b)) -> f a -> f s -> (f s, f b)
app_help ((STFTFrame -> Vector Double -> STFTFrame)
 -> f STFTFrame
 -> f (Vector Double)
 -> (f (Vector Double), f (Vector Double)))
-> (STFTFrame -> Vector Double -> STFTFrame)
-> f STFTFrame
-> f (Vector Double)
-> (f (Vector Double), f (Vector Double))
forall a b. (a -> b) -> a -> b
$ (Vector Double -> STFTFrame -> STFTFrame)
-> STFTFrame -> Vector Double -> STFTFrame
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Vector Double -> STFTFrame -> STFTFrame)
 -> STFTFrame -> Vector Double -> STFTFrame)
-> (Vector Double -> STFTFrame -> STFTFrame)
-> STFTFrame
-> Vector Double
-> STFTFrame
forall a b. (a -> b) -> a -> b
$ VocoderParams -> Vector Double -> STFTFrame -> STFTFrame
synthesisBlock VocoderParams
par) f (Vector Double)
ph

-- | Perform frequency domain processing.
processFramesF :: (Applicative f, Monad m) => VocoderParams -> (f Phase, f Phase) -> ConduitT (f STFTFrame) (f STFTFrame) m r -> ConduitT (f Frame) (f Frame) m (r, (f Phase, f Phase))
processFramesF :: VocoderParams
-> (f (Vector Double), f (Vector Double))
-> ConduitT (f STFTFrame) (f STFTFrame) m r
-> ConduitT
     (f (Vector Double))
     (f (Vector Double))
     m
     (r, (f (Vector Double), f (Vector Double)))
processFramesF VocoderParams
par (f (Vector Double)
p1, f (Vector Double)
p2) ConduitT (f STFTFrame) (f STFTFrame) m r
c = (\((f (Vector Double)
p1', r
r), f (Vector Double)
p2') -> (r
r, (f (Vector Double)
p1', f (Vector Double)
p2'))) (((f (Vector Double), r), f (Vector Double))
 -> (r, (f (Vector Double), f (Vector Double))))
-> ConduitT
     (f (Vector Double))
     (f (Vector Double))
     m
     ((f (Vector Double), r), f (Vector Double))
-> ConduitT
     (f (Vector Double))
     (f (Vector Double))
     m
     (r, (f (Vector Double), f (Vector Double)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VocoderParams
-> f (Vector Double)
-> ConduitT (f (Vector Double)) (f STFTFrame) m (f (Vector Double))
forall (f :: * -> *) (m :: * -> *).
(Applicative f, Monad m) =>
VocoderParams
-> f (Vector Double)
-> ConduitT (f (Vector Double)) (f STFTFrame) m (f (Vector Double))
analysisF VocoderParams
par f (Vector Double)
p1 ConduitT (f (Vector Double)) (f STFTFrame) m (f (Vector Double))
-> ConduitT (f STFTFrame) (f STFTFrame) m r
-> ConduitT
     (f (Vector Double)) (f STFTFrame) m (f (Vector Double), r)
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` (VocoderParams -> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (f :: * -> *) (m :: * -> *).
(Applicative f, Monad m) =>
VocoderParams -> ConduitT (f STFTFrame) (f STFTFrame) m ()
volumeFixF VocoderParams
par ConduitT (f STFTFrame) (f STFTFrame) m ()
-> ConduitT (f STFTFrame) (f STFTFrame) m r
-> ConduitT (f STFTFrame) (f STFTFrame) m r
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitT (f STFTFrame) (f STFTFrame) m r
c) ConduitT (f (Vector Double)) (f STFTFrame) m (f (Vector Double), r)
-> ConduitT (f STFTFrame) (f (Vector Double)) m (f (Vector Double))
-> ConduitT
     (f (Vector Double))
     (f (Vector Double))
     m
     ((f (Vector Double), r), f (Vector Double))
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` VocoderParams
-> f (Vector Double)
-> ConduitT (f STFTFrame) (f (Vector Double)) m (f (Vector Double))
forall (f :: * -> *) (m :: * -> *).
(Applicative f, Monad m) =>
VocoderParams
-> f (Vector Double)
-> ConduitT (f STFTFrame) (f (Vector Double)) m (f (Vector Double))
synthesisF VocoderParams
par f (Vector Double)
p2