{-# LANGUAGE RebindableSyntax #-}
module Fourier where

import qualified Math.FFT as FFT

import qualified Synthesizer.Generic.Analysis as Ana

import qualified Data.StorableVector as SV
import qualified Data.StorableVector.CArray as SVCArr

import qualified Data.Array.CArray as CArray
import qualified Data.Array.IArray as IArray
import Data.Array.CArray (CArray)
import Data.Array.IArray ((//), (!))
import Data.Tuple.HT (mapFst)

import Foreign.Storable (Storable)

import Control.Monad (liftM2)

import qualified Data.Complex as Complex

import qualified Algebra.IntegralDomain as Integral
import qualified Algebra.Additive as Additive
import qualified Algebra.Ring as Ring
import NumericPrelude
import Prelude ()


-- see synthesizer-core:Synthesizer.Basic.NumberTheory
{- |
It's not awfully efficient, but ok for our uses.
-}
ceilingPower :: (Integral.C a, Ord a) => a -> a -> a
ceilingPower base n = base ^ fromIntegral (ceilingLog base n)

ceilingLog :: (Integral.C a, Ord a) => a -> a -> Int
ceilingLog base =
   length . takeWhile (>0) . iterate (flip div base) . subtract 1


{- |
It must hold @transformSize > chunkSize@.
-}
layoutBlocks :: Int -> Int -> SV.Vector Float -> CArray (Int, Int) Float
layoutBlocks transformSize chunkSize xs =
   let numChunks = Integral.divUp (SV.length xs) chunkSize
   in  CArray.listArray ((0,0), (numChunks-1, transformSize-1)) (repeat 0)
       //
       zip (liftM2 (,) [0..] (take chunkSize [0..])) (SV.unpack xs)

mul1d2d ::
   (Ring.C a, Storable a) =>
   CArray Int a -> CArray (Int,Int) a -> CArray (Int,Int) a
mul1d2d xs ys =
   IArray.array (IArray.bounds ys) $
   map (\((i,j),y) -> ((i,j), xs!j * y)) $
   IArray.assocs ys

permuteClip ::
   (IArray.Ix i, IArray.Ix j, Additive.C a, Storable a) =>
   (i -> j) -> (j,j) -> CArray i a -> CArray j a
permuteClip f bnds xs =
   IArray.accumArray (+) zero bnds $
   filter (IArray.inRange bnds . fst) $
   map (mapFst f) $ IArray.assocs xs

sizes :: Int -> (Int, Int)
sizes len =
   let transformSize = ceilingPower 2 $ 2 * len
       chunkSize = transformSize - len + 1
   in  (transformSize, chunkSize)

padRight :: Int -> SV.Vector Float -> SV.Vector Float
padRight size xs =
   SV.append xs $ SV.replicate (size - SV.length xs) 0

convolve :: SV.Vector Float -> SV.Vector Float -> SV.Vector Float
convolve xs ys =
   let (transformSize, chunkSize) = sizes $ SV.length xs
       dims = [1]
   in  SVCArr.from $
       permuteClip
          (\(i,j) -> i*chunkSize+j)
          (0, SV.length xs + SV.length ys - 2) $
       FFT.dftCRN dims $
       mul1d2d
          (FFT.dftRC $ SVCArr.to $ padRight transformSize xs)
          (FFT.dftRCN dims $ layoutBlocks transformSize chunkSize ys)

correlate :: SV.Vector Float -> SV.Vector Float -> SV.Vector Float
correlate xs ys =
   let (transformSize, chunkSize) = sizes $ SV.length xs
       dims = [1]
   in  SVCArr.from $
       permuteClip
          (\(i,j) -> i*chunkSize+j - SV.length xs + 1)
          (0, SV.length ys - 1) $
       FFT.dftCRN dims $
       mul1d2d
          (FFT.dftRC $ SVCArr.to $ padRight transformSize $ SV.reverse xs)
          (FFT.dftRCN dims $ layoutBlocks transformSize chunkSize ys)


pillow :: Int -> CArray Int Float
pillow n =
   CArray.listArray (0, n-1) $
   map (\i -> sin (pi * fromIntegral i / fromIntegral n)) [(0::Int) .. ]

{- |
The array is not padded.
Instead we choose the longest array that only contains valid data.
-}
layoutOverlapBlocks :: Int -> Int -> SV.Vector Float -> CArray (Int, Int) Float
layoutOverlapBlocks shift chunkSize xs =
   let lastChunk = Integral.div (SV.length xs - chunkSize) shift
       bnds = ((0,0), (lastChunk, chunkSize-1))
   in  CArray.array bnds $ flip map (IArray.range bnds) $
       \(i,j) -> ((i,j), SV.index xs $ i*shift + j)

absoluteBlockSpectra :: Int -> Int -> SV.Vector Float -> CArray (Int, Int) Float
absoluteBlockSpectra shift chunkSize xs =
   CArray.amap Complex.magnitude $
   FFT.dftRCN [1] $ mul1d2d (pillow chunkSize) $
   layoutOverlapBlocks shift chunkSize xs

slice :: CArray (Int, Int) Float -> [SV.Vector Float]
slice arr =
   let ((firstChunk, firstElem), (lastChunk, lastElem)) = IArray.bounds arr
   in  map (\i ->
          SV.sample (IArray.rangeSize (firstElem, lastElem)) $
             \j -> arr ! (i, firstElem+j)) $
       IArray.range (firstChunk, lastChunk)

spectralFlatness :: SV.Vector Float -> Float
spectralFlatness xs =
   2 ** Ana.average (SV.map (logBase 2) xs) / Ana.average xs