{-# LANGUAGE RebindableSyntax #-}
module Measurement where

import qualified LabelChain
import qualified Fourier
import qualified Class
import qualified Named
import qualified SpectralDistribution as SD
import qualified Arithmetic as Arith
import qualified Parameters as Params
import qualified Rate
import qualified Signal
import qualified SignalProcessingMethods as SPMethods
import qualified SignalProcessing as SP
import SignalProcessingMethods (Triple, )
import Parameters (Freq(Freq), )

import qualified Synthesizer.Generic.Analysis as AnaG
import qualified Synthesizer.Basic.Binary as Bin

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

import qualified Control.Applicative.HT as AppHT
import qualified Control.Functor.HT as FuncHT
import Control.DeepSeq (NFData, rnf, )
import Control.Applicative (Applicative, pure, (<*>), )

import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.NonEmpty as NonEmpty
import Data.Tuple.HT (mapPair, )
import Data.Maybe (mapMaybe, )

import qualified Algebra.Field as Field
import NumericPrelude.Numeric
import NumericPrelude.Base hiding (id)


-- | LabelChain.mergePhases drops a trailing isolated chunk
mergeClickPhases :: [[a]] -> [[a]]
mergeClickPhases =
   let go (x0:x1:xs) = (x0++x1) : go xs
       go _ = []
   in  go

halfLife :: (Ord a, Field.C a) => NonEmpty.T [] a -> Int
halfLife xs =
   let xm = NonEmpty.maximum xs
   in  length $ takeWhile (>= xm / 2) $
       dropWhile (<xm) $ NonEmpty.flatten xs


countEmphasized :: Params.T -> [Float] -> Int
countEmphasized params clickAmplitudes =
   let progression = iterate (one+) zero
       (c0,c1) = Arith.linearRegression $ zip progression clickAmplitudes
   in  length $ takeWhile (zero<) $ zipWith (-) clickAmplitudes $
       map (\k -> Params.emphasisExcess params * (c0+c1*k)) progression


type ClassFeatures = Class.Sound (Int, Int, Int) Int Int (Int, Int, Int)

type ChunkFeatures = ((Int, Int, Int), Int)

chunkFeatures ::
   Params.T -> SVL.Vector Float -> SVL.Vector Float -> ChunkFeatures
chunkFeatures params volume featSig =
   let clicks =
         mergeClickPhases $
         LabelChain.chopMonotony
            (mapPair (LabelChain.spanWeakRising, LabelChain.spanWeakFalling) $
             Params.weakCounterSlopeSizes params) $
         SVL.unpack featSig
   in  ((length clicks,
         sum $ map halfLife $ mapMaybe NonEmpty.fetch clicks,
         countEmphasized params $
         map NonEmpty.maximum $ mapMaybe (NonEmpty.fetch . SVL.unpack) $
         flip SP.chop (map length clicks) $
         SVL.zipWith (*) volume featSig),
        SP.chirpingMainDur featSig)


bandFreq0, bandFreq1, bandFreq2 :: Float
bandFreq0 = 1000
bandFreq1 = 2500
bandFreq2 = 4000

bandFreqs :: Triple Freq
bandFreqs = (Freq bandFreq0, Freq bandFreq1, Freq bandFreq2)


data SpectralParameters a =
   SpectralParameters {
      spectralFlatness, spectralMaximum :: a,
      spectralBandParams :: (a,a),
      spectralDistribution :: (SD.T a)
   } deriving Show

instance (NFData a) => NFData (SpectralParameters a) where
   rnf (SpectralParameters specFlat specMax bands distr) =
      rnf (specFlat, specMax, bands, distr)

instance Functor SpectralParameters where
   fmap = Trav.fmapDefault

instance Fold.Foldable SpectralParameters where
   foldMap = Trav.foldMapDefault

instance Trav.Traversable SpectralParameters where
   traverse f (SpectralParameters specFlat specMax bands distr) =
      pure SpectralParameters
         <*> f specFlat <*> f specMax
         <*> AppHT.mapPair (f,f) bands
         <*> Trav.traverse f distr

instance Applicative SpectralParameters where
   pure x = SpectralParameters x x (x,x) (pure x)
   SpectralParameters fSpecFlat fSpecMax fbands fdistr <*>
      SpectralParameters specFlat specMax bands distr =
         SpectralParameters
            (fSpecFlat specFlat) (fSpecMax specMax)
            (mapPair fbands bands) (fdistr <*> distr)


spectralParameters ::
   (Float, Float) -> ((Float, Float), SD.T Float) ->
   SpectralParameters Float
spectralParameters (flat, maxf) (bp, distr) =
   SpectralParameters flat maxf bp distr

spectrogramParameters :: [SV.Vector Float] -> (Float, Float)
spectrogramParameters specs =
   (let blockFlats = map Fourier.spectralFlatness specs
    in  if null blockFlats then 1 else AnaG.average blockFlats,
    let amax block =
           fromIntegral (fst (SP.argMaximum block))
           /
           fromIntegral (2 * (SV.length block-1))
        blockMaxs = map amax specs
    in  if null blockMaxs then 0 else AnaG.average blockMaxs)


classFromChunkFeatures ::
   ChunkFeatures -> Class.Sound rasping chirping ticking growling ->
   ClassFeatures
classFromChunkFeatures (clickMeasure@(numClicks, _, _), chirpMain) cls =
   case cls of
      Class.Rasping _ -> Class.Rasping clickMeasure
      Class.Chirping _ -> Class.Chirping chirpMain
      Class.Ticking _ -> Class.Ticking numClicks
      Class.Growling _ -> Class.Growling clickMeasure
      Class.Other str -> Class.Other str

measureSignal ::
   SPMethods.T -> Params.T ->
   Signal.SoxLabelled (Class.Sound rasping chirping ticking growling) ->
   ([ChunkFeatures],
    (Signal.LabelChain Rate.Measure (SpectralParameters Float, ClassFeatures),
     ([Float], Signal.T Rate.Measure [Named.Signal])))
measureSignal sigProc params labelled =
   let (sig, classified) = FuncHT.unzip labelled
       intervalSizes = Fold.toList . fmap fst . LabelChain.intervalSizes
       measRate = Params.measureSampleRate params
       classifiedHighRate = Signal.body classified
       classifiedMeasRate =
          Signal.body $ Signal.labelResample measRate classified
       dehummed = SPMethods.dehum sigProc sig
       (volume, relEnv) = SPMethods.envelopeLowRate sigProc measRate dehummed
       chunkFeats =
          case intervalSizes classifiedMeasRate of
             chunkSizes ->
                zipWith (chunkFeatures params)
                   (SP.chop volume chunkSizes)
                   (SP.chop relEnv chunkSizes)
       fourierStep = Params.fourierBlockStep params
       fourierSize = Params.fourierBlockSize params
       chunkSizesBlockRate =
          intervalSizes $
          LabelChain.mapTime (max 0) $
          LabelChain.mapTime
             (\n -> div (n + div (fourierStep-fourierSize) 2) fourierStep) $
          classifiedHighRate
       spectroParams =
          map spectrogramParameters $
          flip SP.chop chunkSizesBlockRate $
          Fourier.slice $
          Fourier.absoluteBlockSpectra fourierStep fourierSize $
          SP.svlConcat $ SVL.map Bin.toCanonical $ Signal.body sig
       spectralDists =
          zipWith spectralParameters spectroParams $
          SPMethods.bandParameters sigProc bandFreqs dehummed $
          intervalSizes classifiedHighRate
   in  (chunkFeats,
        (Signal.Cons measRate $
         LabelChain.zipWithList (,) spectralDists $
         LabelChain.zipWithList classFromChunkFeatures
            chunkFeats classifiedMeasRate,
         ([1, 0.4],
          Signal.Cons measRate
            [Named.Cons "volume" volume, Named.Cons "envelope" relEnv])))