module SignalProcessing where

import qualified Rate
import Parameters (Freq, freq, )

import qualified Synthesizer.Plain.Filter.Recursive.Universal as UniFilt
import qualified Synthesizer.Plain.Filter.Recursive.FirstOrder as Filt1
import qualified Synthesizer.Plain.Filter.Recursive as FiltRec
import qualified Synthesizer.Causal.Process as Causal
import qualified Synthesizer.Generic.Analysis as Ana
import qualified Synthesizer.Generic.Cut as Cut
import qualified Synthesizer.Generic.Signal as SigG
import qualified Synthesizer.State.Signal as SigS

import qualified Data.StorableVector.Lazy as SVL
import qualified Data.StorableVector as SV
import Foreign.Storable (Storable, )

import qualified Control.Category as Cat
import Control.Arrow (Arrow, (^<<), (<<<), (&&&))
import Control.DeepSeq (NFData, ($!!))

import qualified Data.NonEmpty as NonEmpty
import qualified Data.Foldable as Fold
import qualified Data.List.Match as Match
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import Data.Traversable (Traversable, mapAccumL)
import Data.Foldable (Foldable, )
import Data.Maybe.HT (toMaybe, )
import Data.Tuple.HT (swap, )

import qualified Algebra.RealRing as Real
import qualified Algebra.Field as Field
import qualified Algebra.Additive as Additive
import NumericPrelude.Numeric
import NumericPrelude.Base
import Prelude ()



{-# INLINE zerothMoment #-}
zerothMoment :: Causal.T Float Float
zerothMoment =
   Causal.consInit 0

{-# INLINE firstMoment #-}
firstMoment :: Causal.T Float Float
firstMoment =
   (\(x0,_x1,x2) -> (x2-x0)/2) ^<< lag2

{-# INLINE secondMoment #-}
secondMoment :: Causal.T Float Float
secondMoment =
   (\(x0,x1,x2) -> x2-2*x1+x0) ^<< lag2

{-# INLINE lag2 #-}
lag2 :: (Additive.C a) => Causal.T a (a,a,a)
lag2 = lag2Init zero

{-# INLINE lag2Init #-}
lag2Init :: a -> Causal.T a (a,a,a)
lag2Init x =
   (\((x0,x1),x2) -> (x0,x1,x2))
   ^<<
   (Causal.consInit x &&& Cat.id <<< Causal.consInit x) &&& Cat.id


{-# INLINE bandpass #-}
bandpass :: (Rate.C rate) => rate -> Float -> Freq -> Causal.T Float Float
bandpass rate q f =
   UniFilt.bandpass
   ^<<
   UniFilt.causal
   <<<
   Causal.feedConstFst (UniFilt.parameter (FiltRec.Pole q (freq rate f)))

{-# INLINE highpass #-}
highpass :: (Rate.C rate) => rate -> Float -> Freq -> Causal.T Float Float
highpass rate q f =
   UniFilt.highpass
   ^<<
   UniFilt.causal
   <<<
   Causal.feedConstFst (UniFilt.parameter (FiltRec.Pole q (freq rate f)))

{-# INLINE lowpass #-}
lowpass :: (Rate.C rate) => rate -> Float -> Freq -> Causal.T Float Float
lowpass rate q f =
   UniFilt.lowpass
   ^<<
   UniFilt.causal
   <<<
   Causal.feedConstFst (UniFilt.parameter (FiltRec.Pole q (freq rate f)))


twoPasses ::
   (Storable a) =>
   (SVL.Vector a -> SVL.Vector a) -> SVL.Vector a -> SVL.Vector a
twoPasses f  =  SVL.reverse . f . SVL.reverse . f

lowpassOnePass ::
   (Rate.C rate) => rate -> Freq -> SVL.Vector Float -> SVL.Vector Float
lowpassOnePass rate f sig =
   let fr = freq rate f
   in  Causal.apply
          (Filt1.lowpass_
           ^<<
           Filt1.causalInit
              (Ana.average $ Cut.take (ceiling (1/fr)) $ SigG.toState sig)
           <<<
           Causal.feedConstFst (Filt1.parameter fr))
          sig

lowpassTwoPass ::
   (Rate.C rate) => rate -> Freq -> SVL.Vector Float -> SVL.Vector Float
lowpassTwoPass rate f = twoPasses (lowpassOnePass rate f)



-- ToDo: move to synthesizer-core:Causal.Filter
{-# INLINE differentiate #-}
differentiate :: Causal.T Float Float
differentiate =
   Cat.id - Causal.consInit zero

{-# INLINE differentiateMin3 #-}
differentiateMin3 :: Causal.T Float Float
differentiateMin3 = differentiateMin3Init zero

{-# INLINE differentiateMin3Init #-}
differentiateMin3Init :: Float -> Causal.T Float Float
differentiateMin3Init x =
   Cat.id
   -
   ((\(x0,x1,x2) -> x0 `min` x1 `min` x2) ^<< lag2Init x
      <<< Causal.consInit x)

{- |
This one produces narrow pikes for the click beginnings.
However this turns out to make the click recognition worse.
-}
differentiateMin3Narrow :: Causal.T Float Float
differentiateMin3Narrow =
   (\(x, (x0,x1,x2)) ->
       let xmin = x0 `min` x1 `min` x2
           xmax = x0 `max` x1 `max` x2
       in  if x>xmax then x-xmin else x-x2) ^<<
      Cat.id &&& (lag2 <<< Causal.consInit zero)


downSampleMax :: Int -> SVL.Vector Float -> SVL.Vector Float
downSampleMax k =
   SigG.fromState SigG.defaultLazySize .
   fmap SVL.maximum . Cut.sliceVertical k


downSampleChunkSizes ::
   (Field.C t, Real.C t) =>
   t -> SigS.T Int
downSampleChunkSizes sizeFrac =
   SigS.unfoldR
      (\sizeRem -> Just $ splitFraction $ sizeRem + sizeFrac)
      zero


{-
ToDo: move to synthesizer-core
could be generalized to Causal arrows
sometimes reversed parameter seems to be more appropriate
-}
chop :: (Traversable f, Cut.Transform sig) => sig -> f Int -> f sig
chop xs  =  snd . mapAccumL (\xsr d -> swap $ Cut.splitAt d xsr) xs


chopFrac :: Double -> SVL.Vector Float -> SigS.T (SVL.Vector Float)
chopFrac sizeFrac xs =
   SigG.crochetL (\n xi -> toMaybe (not $ SVL.null xi) $ SVL.splitAt n xi) xs $
   downSampleChunkSizes sizeFrac

{- |
It must be @sizeFrac >= 1@.
-}
downSampleMaxFrac :: Double -> SVL.Vector Float -> SVL.Vector Float
downSampleMaxFrac sizeFrac =
   SigG.fromState SigG.defaultLazySize . fmap SVL.maximum .
   chopFrac sizeFrac

downSampleMaxAbsFrac :: Double -> SVL.Vector Float -> SVL.Vector Float
downSampleMaxAbsFrac sizeFrac =
   SigG.fromState SigG.defaultLazySize . fmap Ana.volumeMaximum .
   chopFrac sizeFrac

downSampleAvgFrac :: Double -> SVL.Vector Float -> SVL.Vector Float
downSampleAvgFrac sizeFrac =
   SigG.fromState SigG.defaultLazySize . fmap Ana.average .
   chopFrac sizeFrac



takeSlices :: (Cut.Transform sig) => Int -> sig -> [sig] -> [sig]
takeSlices blockSize xs =
   Match.take (SigS.toList $ Cut.sliceVertical blockSize xs)

-- ToDo: move to synthesizer-core
sliceOverlappingAbs, sliceOverlappingDiff, sliceOverlappingRel,
   sliceOverlapping :: (Cut.Transform sig) => Int -> (Int,Int) -> sig -> [sig]

sliceOverlappingAbs blockSize (pre,suf) xs =
   takeSlices blockSize xs $
   map
      (\t ->
         let start = max 0 $ t-pre
             stop = t+blockSize+suf
         in  Cut.take (stop-start) $ Cut.drop start xs) $
   iterate (blockSize+) 0

sliceOverlappingDiff blockSize (pre,suf) xs =
   let offsets = iterate (blockSize+) (-pre)
   in  takeSlices blockSize xs $
       zipWith (\offset -> Cut.take (blockSize+pre+suf+min 0 offset)) offsets $
       snd $ List.mapAccumL (\xsi k -> (Cut.drop k xsi, xsi)) xs $
       ListHT.mapAdjacent subtract $ map (max 0) offsets

sliceOverlappingRel blockSize (pre,suf) xs =
   let offsets = iterate (min 0 . (blockSize+)) (-pre)
   in  takeSlices blockSize xs $
       zipWith (\offset -> Cut.take (blockSize+pre+suf+offset)) offsets $ snd $
       List.mapAccumL
          (\xsi offset -> (Cut.drop (max 0 $ offset+blockSize) xsi, xsi))
          xs offsets

sliceOverlapping = sliceOverlappingRel

propSliceOverlapping :: NonEmpty.T [] Int -> ([Int], [Int]) -> String -> Bool
propSliceOverlapping blockSizeList (preList,sufList) xs =
   let blockSize = length $ NonEmpty.flatten blockSizeList
       pre = length preList
       suf = length sufList
   in  ListHT.allEqual $
          sliceOverlappingAbs blockSize (pre,suf) xs :
          sliceOverlappingDiff blockSize (pre,suf) xs :
          sliceOverlappingRel blockSize (pre,suf) xs :
          []


centroidVariance3 ::
   Field.C a => (a, a) -> (a, a) -> (a, a) -> (a, a)
centroidVariance3 (f0,x0) (f1,x1) (f2,x2) =
   let s = x0+x1+x2
       mean y0 y1 y2 = (y0*x0 + y1*x1 + y2*x2) / s
       center = mean f0 f1 f2
   in  (center, mean ((f0-center)^2) ((f1-center)^2) ((f2-center)^2))


svlConcat :: SVL.Vector Float -> SV.Vector Float
svlConcat = SV.concat . SVL.chunks


foldLength :: (Foldable f) => f a -> Int
foldLength = length . Fold.toList


{-
ToDo: move to utility-ht?
could be the basis of a Synthesizer.State.argmax function
it is now included in 'semigroups' but with swapped order
better maxKey :: (Ord b) => (a -> b) -> a -> a -> a ?
-}
argMax :: (Ord b) => (a,b) -> (a,b) -> (a,b)
argMax x0 x1 =
   if snd x0 >= snd x1 then x0 else x1

argMin :: (Ord b) => (a,b) -> (a,b) -> (a,b)
argMin x0 x1 =
   if snd x0 <= snd x1 then x0 else x1

-- ToDo: move to synthesizer-core
argMaximum ::
   (NFData a, Ord a, Additive.C a, SigG.Read sig a) => sig a -> (Int, a)
argMaximum =
   SigS.foldL (\x y -> (argMax $!! x) y) (0,zero) .
   SigS.zip (SigS.iterate (1+) 0) . SigG.toState



{-# INLINE fanout3 #-}
fanout3 ::
   (Arrow arrow) =>
   arrow a b -> arrow a c -> arrow a d -> arrow a (b,c,d)
fanout3 arrb arrc arrd =
   (\(b,(c,d)) -> (b,c,d)) ^<< arrb &&& arrc &&& arrd



chirpingPauseDur :: SVL.Vector Float -> Int
chirpingPauseDur xs =
   let x = Ana.average xs
   in  SVL.length $ SVL.takeWhile (<=x) $ SVL.reverse xs

chirpingMainDur :: SVL.Vector Float -> Int
chirpingMainDur xs =
   SVL.length xs - chirpingPauseDur xs