{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
-- | Provides an interface to take a raw PCMBuffer
-- and get a FreqBuffer back. The provided Buffer
-- will have a (fftWindow)/2+1 size.
module Sound.Analysis.Spectrum where


import qualified Data.Vector.Unboxed as UB
import qualified Numeric.FFT.Vector.Invertible as FFTW
import qualified Data.List.Split as DL

import Data.Complex
import Numeric.FFT.Vector.Plan

import Sound.Data.PCM
import Sound.Data.Stream
import Sound.Data.Freq


--Top level frequency getter
--Takes a PCMData and returns a FreqData
{-# INLINE freqStream #-}
freqStream :: PCMData -> FreqData
freqStream (PCMData i s) = FreqData i freqStream'
    where   freqStream' = fmap (freqs fftPlan) s
            fftSize = fftWindow i
            fftPlan = FFTW.plan FFTW.dftR2C fftSize

--Takes Raw PCM data and creates a Frequency Frame
{-# INLINE freqs #-}
--freqs :: Int -> Int -> Int -> PCM -> FreqFrame
--freqs :: -> Int -> Int -> PCM -> FreqFrame
freqs   :: Plan Double (Complex Double)
        -> PCM
        -> FreqFrame
freqs fftPlan b = dBDFT--The frequency will now be deduced
                                    --by the sampleRate/fftSize when it's
                                    --needed as the data is in the Stream
    where   monoV = monoize b   --Make Mono
            normV = hanningWindow monoV --Hanning window up on to it

            dft = FFTW.execute fftPlan normV
            magDFT = magFromDFT dft
            dBDFT = dbFromMagDFT magDFT          --capture dB from dft

{-# INLINE magFromDFT #-}
magFromDFT  :: (RealFloat a, UB.Unbox a) =>
               UB.Vector (Complex a)
            -> UB.Vector a
magFromDFT dft = mag
    where   dftSq = UB.map (^2) 
            --Square both sides of r2c dft
            dftReal = dftSq $ UB.map realPart dft
            dftImag = dftSq $ UB.map imagPart dft
            --NO NEED
            --to sqrt changing log10 to ((log x)/log 10)* 10
            mag  =  UB.zipWith (+) dftReal dftImag

--Translate magnitude into decibels
{-# INLINE dbFromMagDFT #-}
dbFromMagDFT :: UB.Vector Double -> UB.Vector Double
dbFromMagDFT dft = db
    where db = UB.map log10 dft
            
lx :: Double
lx = log 10

log10 :: Double -> Double
log10 x = (log x/lx)*10

--Apply a hanning window to a vector
{-# INLINE hanningWindow #-}
hanningWindow   :: PCM
                -> PCM
hanningWindow v = UB.imap hanning v
    where   hanning :: Int -> Double -> Double
            hanning i e = 0.5*(1.0 - (cos $ ( pi2* fromIntegral i)/(totalSamples-1.0)))*e
            pi2 = 2.0*pi
            totalSamples = fromIntegral $ UB.length v

--turn a stereo signal into a mono one
{-# INLINE monoize #-}
monoize ::     PCM
            -> PCM
monoize v = UB.fromList l'
    where   l = UB.toList v
            l'= avgSum $ pairup l

            pairup  = DL.chunksOf 2
            avgSum  = map (\x -> sum x/2.0) 

