{-# LANGUAGE BangPatterns #-}

module Main (main) where

import Prelude                      as P
import Control.Monad                as M
import Data.List                    as L
import Data.Vector.Storable         as V
import Data.Vector.Storable.Mutable as MV
import Numeric.LinearAlgebra        as LA
import Numeric.GSL.Fourier          as LA

import Control.Applicative
import Data.Complex
import System.IO
import System.Random
import System.Time
import Text.Printf

lowPassKernel :: Double -> Double -> Int -> Vector Double
lowPassKernel sr fc ksize = raw / V.singleton sum
   where
      n = V.enumFromN 0 ksize
      t = (n - fromIntegral (ksize `div` 2)) / V.singleton sr
      -- sinc function; replace division by zero with limit when t=0
      sinc' = sin (V.singleton (2*pi*fc) * t) / (V.singleton pi * t)
      sinc  = sinc' V.// [(ksize `div` 2, 2 * fc)]
      -- Hamming window function
      kmax = fromIntegral (ksize - 1)
      hamm = 0.54 - 0.46 * cos (V.singleton (2 * pi / kmax) * n)
      -- Normalize the result
      raw = sinc * hamm
      sum = sumElements raw

invertSpectrum :: Vector Double -> Vector Double
invertSpectrum kernel = midVal `seq` (negate kernel V.// [(mid, 1 - midVal)])
   where
      mid = (dim kernel) `div` 2
      midVal = kernel @> mid

highPassKernel :: Double -> Double -> Int -> Vector Double
highPassKernel sr fc ksize =
   invertSpectrum $ lowPassKernel sr fc ksize

bandRejectKernel :: Double -> (Double, Double) -> Int -> Vector Double
bandRejectKernel sr (lfc, hfc) ksize =
   lowPassKernel sr lfc ksize + highPassKernel sr hfc ksize

bandPassKernel :: Double -> (Double, Double) -> Int -> Vector Double
bandPassKernel sr (lfc, hfc) ksize =
   invertSpectrum $ bandRejectKernel sr (lfc, hfc) ksize

-- convolution = integral(kernel(t-tau)*input(tau),tau)
-- t is output vector index (j); products and summation are done with dot product (<.>)
convolve :: Vector Double -> Vector Double -> Vector Double
convolve kernel input = V.generate osize $ \j -> rkernel <.> V.slice j ksize input
   where
      ksize = dim rkernel
      isize = dim input
      osize = isize - ksize
      rkernel = V.reverse kernel

decimate :: Int -> Vector Double -> Vector Double
decimate osize vec =
   V.generate osize $ \j ->
      (sumElements $ V.slice (j * ssize) ssize vec) / fromIntegral ssize
   where
      vsize = dim vec
      ssize = vsize `div` osize

diffClockTimesSec :: ClockTime -> ClockTime -> Double
diffClockTimesSec a b = sec + picosec / 1.0e12
   where
      diff    = diffClockTimes a b
      sec     = fromIntegral $ tdSec diff
      picosec = fromIntegral $ tdPicosec diff

time :: IO a -> IO (a, Double)
time f = do
   start <- getClockTime
   x     <- f
   end   <- x `seq` getClockTime
   return (x, diffClockTimesSec end start)

main :: IO ()
main = do
   let sample_rate = 10000 {-Hz-} :: Double
   let cutoff      =  1000 {-Hz-} :: Double

   let input_size  = 1000000 :: Int
   let kernel_size =     201 :: Int

   seed <- randomIO

   (input,  inputTime)  <- time $ return $ LA.randomVector seed Gaussian input_size
   (kernel, kernelTime) <- time $ return $ lowPassKernel sample_rate cutoff kernel_size
   (result, resultTime) <- time $ return $ convolve kernel input

   --V.mapM_ (printf "%10.6f\n") kernel

   --let fft_result = V.map magnitude $ LA.fft $ V.map (:+0) result
   --V.mapM_ (printf "%10.6f\n") . decimate 500 . V.take (dim fft_result `div` 2) $ fft_result

   V.mapM_ (printf "%10.6f\n") . V.slice 0 50 $ result

   hFlush stdout

   hPutStrLn stderr $ printf "Input  Time: %8.6f seconds" inputTime
   hPutStrLn stderr $ printf "Kernel Time: %8.6f seconds" kernelTime
   hPutStrLn stderr $ printf "Result Time: %8.6f seconds" resultTime