import Data.Complex

j = 0 :+ 1

exp (pi * j)

(-1.0) :+ 1.2246467991473532e-16

{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedStrings #-}

import Data.Monoid (Sum(Sum), getSum)
import Control.Lens.Indexed (itraverse,imap,TraversableWithIndex)
import Control.Lens
import Data.Functor.Identity (Identity(Identity, runIdentity))
import Graphics.Plotly qualified as PT
import WGL.IHaskell qualified as WGL
import Data.Ratio (numerator,(%))

dftFactor :: (RealFloat a) => Int -> Int -> Int -> Complex a
dftFactor m k n = 
    let m' = fromIntegral m
        n' = fromIntegral n
        k' = fromIntegral k
    in exp ((-2)*pi*j*k'*n'/m')

dftFactor 256 10 3

oneSft m k n x = x * dftFactor m k n

:t itraverse (\n x -> Identity $ oneSft 0 0 n x) ([] :: [Complex Double])

oneSum :: (TraversableWithIndex Int f) => f (Complex Double) -> Int -> Int -> Complex Double
oneSum dataIn m k = getSum . foldMap Sum . runIdentity $ itraverse (\n x -> Identity $ oneSft m k n x) dataIn

-- slow fourrier transform
sft :: [Complex Double] -> [Complex Double]
sft dataIn = let m = length dataIn in imap (\k _x -> oneSum dataIn m k) dataIn

genSinTestData :: Double -> Double -> Double -> Double -> Int -> [Double]
genSinTestData freq amp phase rate n = take n $ [ amp * sin (2*pi*freq*(fromIntegral t / rate) - phase) | t <- [0..]]

testFreq = 1e6
testAmp = 1
testPhase = 0
testSampleRate = 25e6 -- Hz
testSize = 2048

testData1 = genSinTestData testFreq testAmp testPhase testSampleRate testSize
testData2 = genSinTestData (2*testFreq) (testAmp/2) testPhase testSampleRate testSize

testData = zipWith (+) testData1 testData2

sftResult = fftShift $ sft $ map (\r -> r :+ 0) testData

simpleLinePlot xs ys = let points = zip xs ys in PT.plotly "simple" [PT.line (PT.aes & PT.x .~ fst & PT.y .~ snd) points]
simplerLinePlot = simpleLinePlot [0..]

fftShift :: [a] -> [a]
fftShift xs = 
    let nSplit = numerator (length xs % 2)
        posFreqs = take nSplit xs
        negFreqs = drop nSplit xs
    in negFreqs <> posFreqs

testDataPlot = simplerLinePlot testData
testFTPlot =
    let ys = map magnitude sftResult
        xs = [-testSampleRate/2 + i*testSampleRate/testSize | i <- [0..testSize-1]]
    in simpleLinePlot xs ys

WGL.displayPlotly testDataPlot
WGL.displayPlotly testFTPlot

0.7409511253549591 :+ (-0.6715589548470183)

butterfly :: (Num a) => [a] -> [a] -> [a] -> ([a], [a])
butterfly w x0 x1 = let m = zipWith (*) w x1 in (zipWith (+) x0 m, zipWith (-) x0 m)

everyOther [] = []
everyOther [x] = [x]
everyOther (x:_:xs) = x : everyOther xs

fft [x] = [x]
fft samples =
    let nSamples = length samples
        halfNSamples = numerator (nSamples % 2)
        ws = map (dftFactor nSamples 1) [0..halfNSamples-1]
        xEven = everyOther samples
        xOdd = everyOther (drop 1 samples)
        (lo,hi) = butterfly ws (fft xEven) (fft xOdd)
    in lo <> hi

fftResult = fft $ map (:+ 0) testData

plotRawFFT fr =
    let ys = map magnitude $ fftShift fr
        xs = [-(testSampleRate/2) + i*testSampleRate/testSize | i <- [0..testSize-1]]
    in simpleLinePlot xs ys

plotRawFFTdB = plotRawFFT . map (\x -> 10 * logBase 10 x)

fftPlot = plotRawFFTdB fftResult
WGL.displayPlotly fftPlot

-- Should be very small
maximum $ map magnitude $ zipWith (-) (fftShift fftResult) sftResult

2.4377009118881024e-10

hannWindow samples =
    let m = length samples
        halfM = numerator (m % 2)
        hann i = let n = i - halfM in (1 + cos (2 * pi * fromIntegral n / fromIntegral m)) / 2
    in imap (\i x -> x * hann i) samples

realComplex = map (:+ 0)

-- hannWindow $ replicate 128 1

windowedResult = fft $ realComplex $ hannWindow testData

WGL.displayPlotly $ plotRawFFTdB windowedResult

Hann Window¶