module Nes.APU.State.Filter.Chain (FilterChain (..), newFilterChain) where

import Control.Monad
import qualified Data.Vector.Mutable as V
import Nes.APU.State.Filter.Class
import Nes.APU.State.Filter.Constants
import Nes.APU.State.Filter.Fir
import Nes.APU.State.Filter.Iir
import Nes.APU.State.Filter.Sampled
import Prelude hiding (filter)

data FilterChain = MkFC
    { FilterChain -> IOVector SampledFilter
filters :: !(V.IOVector SampledFilter)
    , FilterChain -> Sample
dt :: {-# UNPACK #-} !Float
    }

newFilterChain :: OutputRate -> IO FilterChain
newFilterChain :: Sample -> IO FilterChain
newFilterChain Sample
outputRate = do
    [SampledFilter]
filtersList <- do
        FirFilter
_firFilter <- Sample -> Sample -> Int -> IO FirFilter
lowPassFirFilter Sample
intermediateSampleRate (Sample
outputRate Sample -> Sample -> Sample
forall a. Num a => a -> a -> a
* Sample
0.45) Int
160
        [SampledFilter] -> IO [SampledFilter]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return
            [ Either IirFilter FirFilter -> Sample -> SampledFilter
newSampledFilter (IirFilter -> Either IirFilter FirFilter
forall a b. a -> Either a b
Left IirFilter
identityIirFilter) Sample
1.0
            , Either IirFilter FirFilter -> Sample -> SampledFilter
newSampledFilter (IirFilter -> Either IirFilter FirFilter
forall a b. a -> Either a b
Left (IirFilter -> Either IirFilter FirFilter)
-> IirFilter -> Either IirFilter FirFilter
forall a b. (a -> b) -> a -> b
$ Sample -> Sample -> IirFilter
lowPassIirFilter Sample
clockRate Sample
intermediateCutoff) Sample
clockRate
            , Either IirFilter FirFilter -> Sample -> SampledFilter
newSampledFilter (IirFilter -> Either IirFilter FirFilter
forall a b. a -> Either a b
Left (IirFilter -> Either IirFilter FirFilter)
-> IirFilter -> Either IirFilter FirFilter
forall a b. (a -> b) -> a -> b
$ Sample -> Sample -> IirFilter
highPassIirFilter Sample
intermediateSampleRate Sample
90) Sample
intermediateSampleRate
            , Either IirFilter FirFilter -> Sample -> SampledFilter
newSampledFilter (IirFilter -> Either IirFilter FirFilter
forall a b. a -> Either a b
Left (IirFilter -> Either IirFilter FirFilter)
-> IirFilter -> Either IirFilter FirFilter
forall a b. (a -> b) -> a -> b
$ Sample -> Sample -> IirFilter
highPassIirFilter Sample
intermediateSampleRate Sample
440) Sample
intermediateSampleRate
            , Either IirFilter FirFilter -> Sample -> SampledFilter
newSampledFilter (IirFilter -> Either IirFilter FirFilter
forall a b. a -> Either a b
Left (IirFilter -> Either IirFilter FirFilter)
-> IirFilter -> Either IirFilter FirFilter
forall a b. (a -> b) -> a -> b
$ Sample -> Sample -> IirFilter
lowPassIirFilter Sample
intermediateSampleRate Sample
14000) Sample
intermediateSampleRate
            -- , newSampledFilter (Right firFilter) intermediateSampleRate
            ]
    IOVector SampledFilter
filters <- Int -> IO (MVector (PrimState IO) SampledFilter)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
V.new (Int -> IO (MVector (PrimState IO) SampledFilter))
-> Int -> IO (MVector (PrimState IO) SampledFilter)
forall a b. (a -> b) -> a -> b
$ [SampledFilter] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SampledFilter]
filtersList
    [(Int, SampledFilter)] -> ((Int, SampledFilter) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [SampledFilter] -> [(Int, SampledFilter)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [SampledFilter]
filtersList) (((Int, SampledFilter) -> IO ()) -> IO ())
-> ((Int, SampledFilter) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> SampledFilter -> IO ()) -> (Int, SampledFilter) -> IO ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (MVector (PrimState IO) SampledFilter
-> Int -> SampledFilter -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
V.write IOVector SampledFilter
MVector (PrimState IO) SampledFilter
filters)
    FilterChain -> IO FilterChain
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return MkFC{Sample
IOVector SampledFilter
filters :: IOVector SampledFilter
dt :: Sample
filters :: IOVector SampledFilter
dt :: Sample
..}
  where
    clockRate :: Sample
clockRate = Sample
21477272 Sample -> Sample -> Sample
forall a. Fractional a => a -> a -> a
/ Sample
12
    intermediateSampleRate :: Sample
intermediateSampleRate = Sample
outputRate Sample -> Sample -> Sample
forall a. Num a => a -> a -> a
* Sample
2 Sample -> Sample -> Sample
forall a. Num a => a -> a -> a
+ (Sample
forall a. Floating a => a
pi Sample -> Sample -> Sample
forall a. Fractional a => a -> a -> a
/ Sample
32)
    intermediateCutoff :: Sample
intermediateCutoff = Sample
outputRate Sample -> Sample -> Sample
forall a. Num a => a -> a -> a
* Sample
0.4
    dt :: Sample
dt = Sample
1 Sample -> Sample -> Sample
forall a. Fractional a => a -> a -> a
/ Sample
clockRate

instance Filter IO FilterChain where
    consume :: Sample -> FilterChain -> IO FilterChain
consume = Sample -> FilterChain -> IO FilterChain
filterChainConsumeSample
    output :: FilterChain -> IO Sample
output = FilterChain -> IO Sample
filterChainOutput

filterChainConsumeSample :: Sample -> FilterChain -> IO FilterChain
filterChainConsumeSample :: Sample -> FilterChain -> IO FilterChain
filterChainConsumeSample Sample
sample FilterChain
fc = do
    MVector (PrimState IO) SampledFilter
-> (SampledFilter -> IO SampledFilter) -> Int -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
V.modifyM (FilterChain -> IOVector SampledFilter
filters FilterChain
fc) (Sample -> SampledFilter -> IO SampledFilter
forall (m :: * -> *) a. Filter m a => Sample -> a -> m a
consume Sample
sample) Int
0
    SampledFilter
firstFilter <- MVector (PrimState IO) SampledFilter -> Int -> IO SampledFilter
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.read (FilterChain -> IOVector SampledFilter
filters FilterChain
fc) Int
0
    SampledFilter
_ <-
        (SampledFilter -> Int -> SampledFilter -> IO SampledFilter)
-> SampledFilter
-> MVector (PrimState IO) SampledFilter
-> IO SampledFilter
forall (m :: * -> *) b a.
PrimMonad m =>
(b -> Int -> a -> m b) -> b -> MVector (PrimState m) a -> m b
V.ifoldM'
            ( \SampledFilter
prev Int
currIdx SampledFilter
curr ->
                if Int
currIdx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                    then SampledFilter -> IO SampledFilter
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return SampledFilter
curr
                    else do
                        !SampledFilter
newCurr <- SampledFilter -> Sample -> SampledFilter -> IO SampledFilter
filterChainConsumeIteration SampledFilter
prev (FilterChain -> Sample
dt FilterChain
fc) SampledFilter
curr
                        MVector (PrimState IO) SampledFilter
-> Int -> SampledFilter -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
V.write (FilterChain -> IOVector SampledFilter
filters FilterChain
fc) Int
currIdx SampledFilter
newCurr
                        SampledFilter -> IO SampledFilter
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return SampledFilter
newCurr
            )
            SampledFilter
firstFilter
            (FilterChain -> IOVector SampledFilter
filters FilterChain
fc)
    FilterChain -> IO FilterChain
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return FilterChain
fc

filterChainConsumeIteration :: SampledFilter -> Float -> SampledFilter -> IO SampledFilter
filterChainConsumeIteration :: SampledFilter -> Sample -> SampledFilter -> IO SampledFilter
filterChainConsumeIteration SampledFilter
prev Sample
dt SampledFilter
current =
    if SampledFilter -> Sample
periodCounter SampledFilter
current Sample -> Sample -> Bool
forall a. Ord a => a -> a -> Bool
>= SampledFilter -> Sample
samplePeriod SampledFilter
current
        then do
            let
                newPeriodCounter :: Sample
newPeriodCounter = SampledFilter -> Sample
periodCounter SampledFilter
current Sample -> Sample -> Sample
forall a. Num a => a -> a -> a
- SampledFilter -> Sample
samplePeriod SampledFilter
current
            Sample
previousOutput <- Either IirFilter FirFilter -> IO Sample
forall (m :: * -> *) a. Filter m a => a -> m Sample
output (Either IirFilter FirFilter -> IO Sample)
-> Either IirFilter FirFilter -> IO Sample
forall a b. (a -> b) -> a -> b
$ SampledFilter -> Either IirFilter FirFilter
filter SampledFilter
prev
            SampledFilter
newCurrent <- Sample -> SampledFilter -> IO SampledFilter
forall (m :: * -> *) a. Filter m a => Sample -> a -> m a
consume Sample
previousOutput (SampledFilter -> IO SampledFilter)
-> SampledFilter -> IO SampledFilter
forall a b. (a -> b) -> a -> b
$ SampledFilter
current{periodCounter = newPeriodCounter}
            SampledFilter -> Sample -> SampledFilter -> IO SampledFilter
filterChainConsumeIteration
                SampledFilter
prev
                Sample
dt
                SampledFilter
newCurrent
        else
            let newPeriodCounter :: Sample
newPeriodCounter = SampledFilter -> Sample
periodCounter SampledFilter
current Sample -> Sample -> Sample
forall a. Num a => a -> a -> a
+ Sample
dt
             in SampledFilter -> IO SampledFilter
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (SampledFilter -> IO SampledFilter)
-> SampledFilter -> IO SampledFilter
forall a b. (a -> b) -> a -> b
$ SampledFilter
current{periodCounter = newPeriodCounter}

{-# INLINE filterChainOutput #-}
filterChainOutput :: FilterChain -> IO Sample
filterChainOutput :: FilterChain -> IO Sample
filterChainOutput FilterChain
fc = case IOVector SampledFilter -> Int
forall s a. MVector s a -> Int
V.length (IOVector SampledFilter -> Int) -> IOVector SampledFilter -> Int
forall a b. (a -> b) -> a -> b
$ FilterChain -> IOVector SampledFilter
filters FilterChain
fc of
    Int
0 -> Sample -> IO Sample
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Sample
0
    Int
l -> (IirFilter -> IO Sample)
-> (FirFilter -> IO Sample)
-> Either IirFilter FirFilter
-> IO Sample
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall (m :: * -> *) a. Filter m a => a -> m Sample
output @IO) (forall (m :: * -> *) a. Filter m a => a -> m Sample
output @IO) (Either IirFilter FirFilter -> IO Sample)
-> (SampledFilter -> Either IirFilter FirFilter)
-> SampledFilter
-> IO Sample
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SampledFilter -> Either IirFilter FirFilter
filter (SampledFilter -> IO Sample) -> IO SampledFilter -> IO Sample
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState IO) SampledFilter -> Int -> IO SampledFilter
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.read (FilterChain -> IOVector SampledFilter
filters FilterChain
fc) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)