{-# LANGUAGE RebindableSyntax #-}
module SignalProcessingSpecific where

import qualified Signal
import qualified Named
import qualified Rate
import Parameters (Freq(Freq), formatFreq, )

import qualified SpectralDistribution as SD
import qualified SignalProcessingMethods as Methods
import qualified SignalProcessing as SP
import SignalProcessingMethods (Triple, )
import SignalProcessing
          (bandpass, highpass, lowpassTwoPass,
           downSampleMaxFrac, downSampleMaxAbsFrac, fanout3, )

import qualified Synthesizer.Generic.Signal as SigG
import qualified Synthesizer.Causal.Process as Causal
import qualified Synthesizer.Basic.Binary as Bin

import qualified Data.StorableVector.Lazy as SVL
import qualified Data.StorableVector as SV

import qualified Data.List as List
import Control.Arrow ((&&&), (^<<), (<<^), )
import Data.Tuple.HT (mapSnd, fst3, snd3, thd3, )

import NumericPrelude.Numeric
import NumericPrelude.Base


{-# INLINE dehum #-}
dehum :: Rate.Sample -> Causal.T Float Float
dehum rate = highpass rate 1 (Freq 800)


filterBand :: Float -> Freq -> Signal.Sox -> Named.Signal
filterBand q f (Signal.Cons rate sig) =
   Named.Cons ("band " ++ formatFreq f) $
   Causal.apply (abs ^<< bandpass rate q f <<^ Bin.toCanonical) sig

{-
We have checked with a chirp that the bands slightly overlap.
-}
filterBands ::
   Signal.Sox -> (Named.Signal, (Named.Signal, Named.Signal, Named.Signal))
filterBands sig =
   ((filterBand 2 (Freq 1600) sig) {Named.name = "dehummed"},
    (filterBand 10 (Freq 1200) sig,
     filterBand 10 (Freq 2000) sig,
     filterBand 10 (Freq 4000) sig))

bandEnvelopes ::
   Signal.Sox -> (Named.Signal, (Named.Signal, Named.Signal, Named.Signal))
bandEnvelopes sig@(Signal.Cons rate _) =
   let (broadband, (band12, band20, band40)) = filterBands sig
       volume = lowpassTwoPass rate (Freq 20) $ Named.body broadband
       envelope xs = SVL.zipWith (/) (lowpassTwoPass rate (Freq 200) xs) volume
   in  (Named.Cons "volume" volume,
        (fmap envelope band12, fmap envelope band20, fmap envelope band40))


bandEnvelopesLowRate ::
   Rate.Feature -> Signal.Sox ->
   (Named.Signal, (Named.Signal, Named.Signal, Named.Signal))
bandEnvelopesLowRate featRate sig =
   let (broadband, (band12, band20, band40)) = filterBands sig
       -- ToDo: would be simpler, if broadband contains the sample rate
       k = Rate.ratio (Signal.sampleRate sig) featRate
       volume =
          lowpassTwoPass featRate (Freq 20) $
          downSampleMaxFrac k $ Named.body broadband
       envelope xs = SVL.zipWith (/) (downSampleMaxFrac k xs) volume
   in  (Named.Cons "volume" volume,
        (fmap envelope band12, fmap envelope band20, fmap envelope band40))


{-# INLINE bandsDerivatives #-}
bandsDerivatives ::
   Triple Freq -> Signal.Sampled Float ->
   SVL.Vector (Triple Float, Triple Float)
bandsDerivatives bandFreqs (Signal.Cons rate sig) =
   Causal.apply
      (let band f = bandpass rate 10 (f bandFreqs)
       in  fanout3 (band fst3) (band snd3) (band thd3)
       &&&
       fanout3 SP.zerothMoment SP.firstMoment SP.secondMoment)
      sig


spectralDistribution1Slow, spectralDistribution2Slow ::
   SVL.Vector (Triple Float) -> SD.T Float
spectralDistribution1Slow chunk =
   let partSum sel = SigG.sum $ SigG.map (abs.sel) $ SigG.toState chunk
   in  SD.spectralDistribution1
          (partSum fst3) (partSum snd3) (partSum thd3)

spectralDistribution2Slow chunk =
   let partSum sel = SigG.sum $ SigG.map ((^2).sel) $ SigG.toState chunk
   in  SD.spectralDistribution2
          (partSum fst3) (partSum snd3) (partSum thd3)


sumSV :: SV.Vector Float -> Float
sumSV = SV.foldl' (+) 0

_sumSVL :: SVL.Vector Float -> Float
_sumSVL = sum . map sumSV . SVL.chunks


addSumSV :: Float -> SV.Vector Float -> Float
addSumSV = SV.foldl' (+)

{- |
Consistently sum with left associativity.
This is consistent with the LLVM implementation.
-}
sumSVL :: SVL.Vector Float -> Float
sumSVL = List.foldl' addSumSV 0 . SVL.chunks


spectralDistribution1, spectralDistribution2 ::
   SVL.Vector (Triple Float) -> SD.T Float
spectralDistribution1 chunk =
   SD.spectralDistribution1
      (sumSVL $ SVL.map (abs.fst3) chunk)
      (sumSVL $ SVL.map (abs.snd3) chunk)
      (sumSVL $ SVL.map (abs.thd3) chunk)

spectralDistribution2 chunk =
   SD.spectralDistribution2
      (sumSVL $ SVL.map ((^2).fst3) chunk)
      (sumSVL $ SVL.map ((^2).snd3) chunk)
      (sumSVL $ SVL.map ((^2).thd3) chunk)


spectralBandDistr :: Triple Freq -> SVL.Vector (Triple Float) -> (Float, Float)
spectralBandDistr (Freq bandFreq0, Freq bandFreq1, Freq bandFreq2) chunk =
   mapSnd sqrt $
   SP.centroidVariance3
      (bandFreq0, sumSVL $ SVL.map (abs.fst3) chunk)
      (bandFreq1, sumSVL $ SVL.map (abs.snd3) chunk)
      (bandFreq2, sumSVL $ SVL.map (abs.thd3) chunk)

bandParameters ::
   Triple Freq -> SVL.Vector (Triple Float, Triple Float) ->
   ((Float, Float), SD.T Float)
bandParameters bandFreqs chunk =
   (spectralBandDistr bandFreqs (SVL.map fst chunk),
    spectralDistribution1 (SVL.map snd chunk))


methods :: Methods.T
methods =
   Methods.Cons {
      Methods.dehum =
         \(Signal.Cons rate xs) ->
            Signal.Cons rate $ Causal.apply (dehum rate <<^ Bin.toCanonical) xs,

      Methods.rumble =
         \(Signal.Cons rate xs) ->
            Signal.Cons rate $
            Causal.apply (SP.lowpass rate 5 (Freq 220) <<^ Bin.toCanonical) xs,

      Methods.downSampleAbs =
         \featRate (Signal.Cons rate xs) ->
            downSampleMaxAbsFrac (Rate.unpack rate / featRate) xs,

      Methods.bandpassDownSample =
         \featRate f (Signal.Cons rate xs) ->
            downSampleMaxAbsFrac (Rate.ratio rate featRate) $
            Causal.apply (bandpass rate 10 f <<^ Bin.toCanonical) xs,

      Methods.bandParameters =
         \bandFreqs sig sizes ->
            map (bandParameters bandFreqs) $
            SP.chop (bandsDerivatives bandFreqs sig) sizes
   }