{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NoStrict #-}

-- Using laziness to build cyclic lists for the output

module Nes.APU.State.Filter.Fir (FirFilter (..), lowPassFirFilter) where

import Control.Monad.IO.Class
import qualified Data.Vector.Unboxed.Mutable as V
import Nes.APU.State.Filter.Class
import Nes.APU.State.Filter.Constants

-- | Finite impulse response (FIR) filter
data FirFilter = MkFirF
    { FirFilter -> IOVector Float
kernel :: !(V.IOVector Float)
    , FirFilter -> IOVector Float
inputs :: !(V.IOVector Float)
    , FirFilter -> Int
inputIndex :: {-# UNPACK #-} !Int
    }

instance (MonadIO m) => Filter m FirFilter where
    output :: FirFilter -> m Float
output MkFirF{Int
IOVector Float
kernel :: FirFilter -> IOVector Float
inputs :: FirFilter -> IOVector Float
inputIndex :: FirFilter -> Int
kernel :: IOVector Float
inputs :: IOVector Float
inputIndex :: Int
..} =
        IO Float -> m Float
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Float -> m Float) -> IO Float -> m Float
forall a b. (a -> b) -> a -> b
$
            (Float -> Int -> Float -> IO Float)
-> Float -> MVector (PrimState IO) Float -> IO Float
forall (m :: * -> *) a b.
(PrimMonad m, Unbox a) =>
(b -> Int -> a -> m b) -> b -> MVector (PrimState m) a -> m b
V.ifoldM'
                ( \Float
acc Int
idx Float
n -> do
                    Float
n' <- MVector (PrimState IO) Float -> Int -> IO Float
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
V.read IOVector Float
MVector (PrimState IO) Float
inputs ((Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inputIndex) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` IOVector Float -> Int
forall a s. Unbox a => MVector s a -> Int
V.length IOVector Float
inputs)
                    Float -> IO Float
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Float -> IO Float) -> Float -> IO Float
forall a b. (a -> b) -> a -> b
$ Float
acc Float -> Float -> Float
forall a. Num a => a -> a -> a
+ (Float
n Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
n')
                )
                Float
0
                IOVector Float
MVector (PrimState IO) Float
kernel
    consume :: Float -> FirFilter -> m FirFilter
consume Float
sample FirFilter
f = IO FirFilter -> m FirFilter
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO FirFilter -> m FirFilter) -> IO FirFilter -> m FirFilter
forall a b. (a -> b) -> a -> b
$ do
        MVector (PrimState IO) Float -> Int -> Float -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
V.write (FirFilter -> IOVector Float
inputs FirFilter
f) (FirFilter -> Int
inputIndex FirFilter
f) Float
sample
        FirFilter -> IO FirFilter
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (FirFilter -> IO FirFilter) -> FirFilter -> IO FirFilter
forall a b. (a -> b) -> a -> b
$
            FirFilter
f
                { inputIndex = newInputIndex
                }
      where
        newInputIndex :: Int
newInputIndex = (FirFilter -> Int
inputIndex FirFilter
f Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` IOVector Float -> Int
forall a s. Unbox a => MVector s a -> Int
V.length (FirFilter -> IOVector Float
inputs FirFilter
f)

lowPassFirFilter :: SampleRate -> Cutoff -> Int -> IO FirFilter
lowPassFirFilter :: Float -> Float -> Int -> IO FirFilter
lowPassFirFilter Float
sampleRate Float
cutoff Int
windowSize = do
    let inputIndex :: Int
inputIndex = Int
0
    IOVector Float
kernel <- Float -> Float -> Int -> IO (IOVector Float)
windowedSincKernel Float
sampleRate Float
cutoff (Int
windowSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    IOVector Float
inputs <- Int -> Float -> IO (MVector (PrimState IO) Float)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
V.replicate (Int
windowSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Float
0
    FirFilter -> IO FirFilter
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return MkFirF{Int
IOVector Float
kernel :: IOVector Float
inputs :: IOVector Float
inputIndex :: Int
inputIndex :: Int
kernel :: IOVector Float
inputs :: IOVector Float
..}

windowedSincKernel :: SampleRate -> Cutoff -> Int -> IO (V.IOVector Float)
windowedSincKernel :: Float -> Float -> Int -> IO (IOVector Float)
windowedSincKernel Float
sampleRate Float
cutoff Int
windowSize = do
    let fc :: Float
fc = Float
cutoff Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
sampleRate
    IOVector Float
kernelV <- Int -> (Int -> Float) -> IO (MVector (PrimState IO) Float)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> (Int -> a) -> m (MVector (PrimState m) a)
V.generate Int
windowSize ((Int -> Float) -> IO (MVector (PrimState IO) Float))
-> (Int -> Float) -> IO (MVector (PrimState IO) Float)
forall a b. (a -> b) -> a -> b
$ \Int
i -> (Int -> Float -> Int -> Float
sinc Int
i Float
fc Int
windowSize) Float -> Float -> Float
forall a. Num a => a -> a -> a
* Int -> Int -> Float
blackmanWindow Int
i Int
windowSize
    IOVector Float -> IO ()
normalise IOVector Float
kernelV
    IOVector Float -> IO (IOVector Float)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return IOVector Float
kernelV
  where
    blackmanWindow :: Int -> Int -> Float
    blackmanWindow :: Int -> Int -> Float
blackmanWindow Int
idx Int
winSize =
        let
            fIdx :: Float
fIdx = Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx
            fWinSize :: Float
fWinSize = Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
winSize
            tau :: Float
tau = Float
2 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
forall a. Floating a => a
pi
         in
            Float
0.42
                Float -> Float -> Float
forall a. Num a => a -> a -> a
- Float
0.5
                    Float -> Float -> Float
forall a. Num a => a -> a -> a
* ((Float -> Float
forall a. Floating a => a -> a
cos ((Float
tau Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
fIdx) Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
fWinSize)) Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
0.08 Float -> Float -> Float
forall a. Num a => a -> a -> a
* (Float -> Float
forall a. Floating a => a -> a
cos ((Float
2 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
tau Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
fIdx Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
fWinSize))))
    sinc :: Int -> Float -> Int -> Float
    sinc :: Int -> Float -> Int -> Float
sinc Int
idx Float
fc Int
winSize =
        let
            fIdx :: Float
fIdx = Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx
            fWinSize :: Float
fWinSize = Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
winSize
            shiftedIndex :: Float
shiftedIndex = Float
fIdx Float -> Float -> Float
forall a. Num a => a -> a -> a
- (Float
fWinSize Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
2)
            tau :: Float
tau = Float
2 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
forall a. Floating a => a
pi
         in
            if Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
windowSize Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
                then Float
tau Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
fc
                else (Float -> Float
mySin (Float
tau Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
fc Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
shiftedIndex)) Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
shiftedIndex
    normalise :: V.IOVector Float -> IO ()
    normalise :: IOVector Float -> IO ()
normalise IOVector Float
vec = do
        Float
vecSum <- (Float -> Float -> Float)
-> Float -> MVector (PrimState IO) Float -> IO Float
forall (m :: * -> *) a b.
(PrimMonad m, Unbox a) =>
(b -> a -> b) -> b -> MVector (PrimState m) a -> m b
V.foldl' Float -> Float -> Float
forall a. Num a => a -> a -> a
(+) Float
0 IOVector Float
MVector (PrimState IO) Float
vec
        (Int -> Float -> IO ()) -> MVector (PrimState IO) Float -> IO ()
forall (m :: * -> *) a b.
(PrimMonad m, Unbox a) =>
(Int -> a -> m b) -> MVector (PrimState m) a -> m ()
V.imapM_ (\Int
i Float
a -> MVector (PrimState IO) Float -> Int -> Float -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
V.write IOVector Float
MVector (PrimState IO) Float
vec Int
i (Float
a Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
vecSum)) IOVector Float
MVector (PrimState IO) Float
vec

-- | Faster implementation of the sin function,
--
-- Stolen from https://www.youtube.com/watch?v=72dI7dB3ZvQ
mySin :: Float -> Float
mySin :: Float -> Float
mySin Float
t =
    let
        j0 :: Float
j0 = Float
t Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
0.15915
        j1 :: Float
j1 = Float
j0 Float -> Float -> Float
forall a. Num a => a -> a -> a
- Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Float -> Int
forall b. Integral b => Float -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Float
j0 :: Int)
     in
        Float
20.785 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
j1 Float -> Float -> Float
forall a. Num a => a -> a -> a
* (Float
j1 Float -> Float -> Float
forall a. Num a => a -> a -> a
- Float
0.5) Float -> Float -> Float
forall a. Num a => a -> a -> a
* (Float
j1 Float -> Float -> Float
forall a. Num a => a -> a -> a
- Float
1)