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

This module defines some useful frequency-domain filters as conduits.
It includes convenience wrappers for filters defined in the vocoder package.
-}
{-# LANGUAGE RankNTypes #-}
module Vocoder.Conduit.Filter(
      Filter,
      runFilter,
      idFilter,
      composeFilters,
      realtimeFilter,
      amplitudeFilter,
      linearAmplitudeFilter,
      amplify,
      lowpassBrickwall,
      highpassBrickwall,
      bandpassBrickwall,
      bandstopBrickwall,
      lowpassButterworth,
      highpassButterworth,
      bandpassButterworth,
      bandstopButterworth,
      pitchShiftInterpolate,
      convolutionFilter,
      envelopeFilter,
      randomPhaseFilter,
      playSpeed
    ) where

import Vocoder
import qualified Vocoder.Filter as F
import Data.Conduit
import Control.Monad.IO.Class
import qualified Data.Vector.Storable as V
import qualified Data.Conduit.Combinators as DCC

-- | Conduit frequency-domain filter type. A conduit filter extends 
--   basic frequency-domain filters by using a conduit instead of a
--   pure function. This enables time transformation filters.
newtype Filter m = Filter { Filter m
-> forall (f :: * -> *).
   Traversable f =>
   FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ()
runFilter :: forall f. Traversable f => F.FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m () }

-- | Identity filter
idFilter :: Monad m => Filter m
idFilter :: Filter m
idFilter = (forall (f :: * -> *).
 Traversable f =>
 FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
forall (m :: * -> *).
(forall (f :: * -> *).
 Traversable f =>
 FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
Filter ((forall (f :: * -> *).
  Traversable f =>
  FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
 -> Filter m)
-> (forall (f :: * -> *).
    Traversable f =>
    FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
_ -> (f STFTFrame -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (m :: * -> *) i o r.
Monad m =>
(i -> ConduitT i o m r) -> ConduitT i o m ()
awaitForever f STFTFrame -> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield

-- | Sequential filter composition.
composeFilters :: Monad m => Filter m -> Filter m -> Filter m
composeFilters :: Filter m -> Filter m -> Filter m
composeFilters (Filter forall (f :: * -> *).
Traversable f =>
FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ()
f1) (Filter forall (f :: * -> *).
Traversable f =>
FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ()
f2) = (forall (f :: * -> *).
 Traversable f =>
 FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
forall (m :: * -> *).
(forall (f :: * -> *).
 Traversable f =>
 FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
Filter ((forall (f :: * -> *).
  Traversable f =>
  FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
 -> Filter m)
-> (forall (f :: * -> *).
    Traversable f =>
    FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
step -> FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (f :: * -> *).
Traversable f =>
FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ()
f1 FreqStep
step ConduitT (f STFTFrame) (f STFTFrame) m ()
-> ConduitT (f STFTFrame) (f STFTFrame) m ()
-> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (f :: * -> *).
Traversable f =>
FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ()
f2 FreqStep
step

-- | Use a basic frequency-domain filter as a conduit filter.
realtimeFilter :: Monad m => F.Filter m -> Filter m
realtimeFilter :: Filter m -> Filter m
realtimeFilter Filter m
f = (forall (f :: * -> *).
 Traversable f =>
 FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
forall (m :: * -> *).
(forall (f :: * -> *).
 Traversable f =>
 FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
Filter (\FreqStep
step -> (f STFTFrame -> m (f STFTFrame))
-> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> ConduitT a b m ()
DCC.mapM ((f STFTFrame -> m (f STFTFrame))
 -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> (f STFTFrame -> m (f STFTFrame))
-> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall a b. (a -> b) -> a -> b
$ (STFTFrame -> m STFTFrame) -> f STFTFrame -> m (f STFTFrame)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((STFTFrame -> m STFTFrame) -> f STFTFrame -> m (f STFTFrame))
-> (STFTFrame -> m STFTFrame) -> f STFTFrame -> m (f STFTFrame)
forall a b. (a -> b) -> a -> b
$ Filter m
f FreqStep
step)

-- | Creates a conduit filter which transforms only amplitudes, leaving
--   phase increments unchanged.
amplitudeFilter :: Monad m => (F.FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter :: (FreqStep -> Moduli -> Moduli) -> Filter m
amplitudeFilter = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m)
-> ((FreqStep -> Moduli -> Moduli) -> Filter m)
-> (FreqStep -> Moduli -> Moduli)
-> Filter m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FreqStep -> Moduli -> Moduli) -> Filter m
forall (m :: * -> *).
Monad m =>
(FreqStep -> Moduli -> Moduli) -> Filter m
F.amplitudeFilter

-- | Creates a filter which scales amplitudes depending on frequency.
linearAmplitudeFilter :: Monad m => (Double -> Double) -> Filter m
linearAmplitudeFilter :: (FreqStep -> FreqStep) -> Filter m
linearAmplitudeFilter = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m)
-> ((FreqStep -> FreqStep) -> Filter m)
-> (FreqStep -> FreqStep)
-> Filter m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FreqStep -> FreqStep) -> Filter m
forall (m :: * -> *). Monad m => (FreqStep -> FreqStep) -> Filter m
F.linearAmplitudeFilter

-- | Creates an "amplifier" which scales all frequencies.
amplify :: Monad m => Double -> Filter m
amplify :: FreqStep -> Filter m
amplify = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m)
-> (FreqStep -> Filter m) -> FreqStep -> Filter m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FreqStep -> Filter m
forall (m :: * -> *). Monad m => FreqStep -> Filter m
F.amplify

-- | Creates a brickwall lowpass filter.
lowpassBrickwall :: Monad m => Double -> Filter m
lowpassBrickwall :: FreqStep -> Filter m
lowpassBrickwall FreqStep
t = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> Filter m
forall (m :: * -> *). Monad m => FreqStep -> Filter m
F.lowpassBrickwall FreqStep
t

-- | Creates a brickwall highpass filter.
highpassBrickwall :: Monad m => Double -> Filter m
highpassBrickwall :: FreqStep -> Filter m
highpassBrickwall FreqStep
t = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> Filter m
forall (m :: * -> *). Monad m => FreqStep -> Filter m
F.highpassBrickwall FreqStep
t

-- | Creates a brickwall bandpass filter.
bandpassBrickwall :: Monad m => Double -> Double -> Filter m
bandpassBrickwall :: FreqStep -> FreqStep -> Filter m
bandpassBrickwall FreqStep
t FreqStep
u = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> FreqStep -> Filter m
forall (m :: * -> *). Monad m => FreqStep -> FreqStep -> Filter m
F.bandpassBrickwall FreqStep
t FreqStep
u

-- | Creates a brickwall bandstop filter.
bandstopBrickwall :: Monad m => Double -> Double -> Filter m
bandstopBrickwall :: FreqStep -> FreqStep -> Filter m
bandstopBrickwall FreqStep
t FreqStep
u = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> FreqStep -> Filter m
forall (m :: * -> *). Monad m => FreqStep -> FreqStep -> Filter m
F.bandstopBrickwall FreqStep
t FreqStep
u

-- | Creates an n-th degree Butterworth-style lowpass filter.
lowpassButterworth :: Monad m => Double -> Double -> Filter m
lowpassButterworth :: FreqStep -> FreqStep -> Filter m
lowpassButterworth FreqStep
n FreqStep
t = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> FreqStep -> Filter m
forall (m :: * -> *). Monad m => FreqStep -> FreqStep -> Filter m
F.lowpassButterworth FreqStep
n FreqStep
t

-- | Creates an n-th degree Butterworth-style highpass filter.
highpassButterworth :: Monad m => Double -> Double -> Filter m
highpassButterworth :: FreqStep -> FreqStep -> Filter m
highpassButterworth FreqStep
n FreqStep
t = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> FreqStep -> Filter m
forall (m :: * -> *). Monad m => FreqStep -> FreqStep -> Filter m
F.highpassButterworth FreqStep
n FreqStep
t

-- | Creates an n-th degree Butterworth-style bandpass filter.
bandpassButterworth :: Monad m => Double -> Double -> Double -> Filter m
bandpassButterworth :: FreqStep -> FreqStep -> FreqStep -> Filter m
bandpassButterworth FreqStep
n FreqStep
t FreqStep
u = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> FreqStep -> FreqStep -> Filter m
forall (m :: * -> *).
Monad m =>
FreqStep -> FreqStep -> FreqStep -> Filter m
F.bandpassButterworth FreqStep
n FreqStep
t FreqStep
u

-- | Creates an n-th degree Butterworth-style bandstop filter.
bandstopButterworth :: Monad m => Double -> Double -> Double -> Filter m
bandstopButterworth :: FreqStep -> FreqStep -> FreqStep -> Filter m
bandstopButterworth FreqStep
n FreqStep
t FreqStep
u = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> FreqStep -> FreqStep -> Filter m
forall (m :: * -> *).
Monad m =>
FreqStep -> FreqStep -> FreqStep -> Filter m
F.bandstopButterworth FreqStep
n FreqStep
t FreqStep
u

-- | Creates an interpolative pitch-shifting filter.
pitchShiftInterpolate :: Monad m => Double -> Filter m
pitchShiftInterpolate :: FreqStep -> Filter m
pitchShiftInterpolate FreqStep
n = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ FreqStep -> Filter m
forall (m :: * -> *). Monad m => FreqStep -> Filter m
F.pitchShiftInterpolate FreqStep
n

-- | Creates a filter which convolves the spectrum using a kernel.
convolutionFilter :: Monad m => V.Vector Double -> Filter m
convolutionFilter :: Moduli -> Filter m
convolutionFilter Moduli
ker = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ Moduli -> Filter m
forall (m :: * -> *). Monad m => Moduli -> Filter m
F.convolutionFilter Moduli
ker

-- | Creates a filter which replaces the amplitudes with their envelope.
envelopeFilter :: Monad m => Length -> Filter m
envelopeFilter :: Length -> Filter m
envelopeFilter Length
ksize = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ Length -> Filter m
forall (m :: * -> *). Monad m => Length -> Filter m
F.envelopeFilter Length
ksize

-- | Sets the phase increments so that the bins have horizontal consistency.
--   This erases the phase information, introducing "phasiness".
randomPhaseFilter :: MonadIO m => Filter m
randomPhaseFilter :: Filter m
randomPhaseFilter = Filter m -> Filter m
forall (m :: * -> *). Monad m => Filter m -> Filter m
realtimeFilter (Filter m -> Filter m) -> Filter m -> Filter m
forall a b. (a -> b) -> a -> b
$ Filter m
forall (m :: * -> *). MonadIO m => Filter m
F.randomPhaseFilter

-- | Changes play speed by replicating or dropping frames.
playSpeed :: Monad m => Rational -> Filter m
playSpeed :: Rational -> Filter m
playSpeed Rational
coeff = (forall (f :: * -> *).
 Traversable f =>
 FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
forall (m :: * -> *).
(forall (f :: * -> *).
 Traversable f =>
 FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
Filter ((forall (f :: * -> *).
  Traversable f =>
  FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
 -> Filter m)
-> (forall (f :: * -> *).
    Traversable f =>
    FreqStep -> ConduitT (f STFTFrame) (f STFTFrame) m ())
-> Filter m
forall a b. (a -> b) -> a -> b
$ \FreqStep
_ -> [f STFTFrame]
-> Rational -> ConduitT (f STFTFrame) (f STFTFrame) m ()
forall (m :: * -> *) o.
Monad m =>
[o] -> Rational -> ConduitT o o m ()
f [] Rational
0
    where
    f :: [o] -> Rational -> ConduitT o o m ()
f [o]
l Rational
c
        | Rational
c Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
< Rational
1 = do
            Maybe o
next <- ConduitT o o m (Maybe o)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await
            case Maybe o
next of
                Maybe o
Nothing -> (o -> ConduitT o o m ()) -> [o] -> ConduitT o o m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ o -> ConduitT o o m ()
forall i o (m :: * -> *). i -> ConduitT i o m ()
leftover ([o] -> ConduitT o o m ()) -> [o] -> ConduitT o o m ()
forall a b. (a -> b) -> a -> b
$ [o] -> [o]
forall a. [a] -> [a]
reverse [o]
l
                Just o
i -> [o] -> Rational -> ConduitT o o m ()
f (o
io -> [o] -> [o]
forall a. a -> [a] -> [a]
:[o]
l) (Rational
c Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
+ Rational
coeff)
        | Bool
otherwise = [o] -> Rational -> ConduitT o o m ()
g [o]
l Rational
c
    g :: [o] -> Rational -> ConduitT o o m ()
g [o]
l Rational
c
        | Rational
c Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
>= Rational
1 = do
            o -> ConduitT o o m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield (o -> ConduitT o o m ()) -> o -> ConduitT o o m ()
forall a b. (a -> b) -> a -> b
$ [o]
l [o] -> Length -> o
forall a. [a] -> Length -> a
!! Length
0
            [o] -> Rational -> ConduitT o o m ()
g [o]
l (Rational
c Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
- Rational
1)
        | Bool
otherwise = [o] -> Rational -> ConduitT o o m ()
f [] Rational
c