{-# LANGUAGE
    DoAndIfThenElse
  , FlexibleContexts
  , ScopedTypeVariables
  , UnicodeSyntax
  #-}
-- | FIXME
module Codec.Audio.WavPack.Decorrelation
    ( DecorrPass(..)
    , decorrStereoPass
    )
    where
import Control.Monad.ST
import Data.Bits
import Data.Int
import Data.STRef
import qualified Data.Vector.Generic.Mutable as MV
import Prelude.Unicode

{-
maxTerm ∷ Num a ⇒ a
maxTerm = 8
-}

-- | FIXME
data DecorrPass v s
    = DecorrPass {
      -- | The decorrelation term: @(term /= 0) && ((-3 <= term <= 8)
      --   || (term == 17) || (term <= 18))@
        dpTerm     ∷ !Int8
      -- | The decorrelation delta: @0 <= delta <= 8@
      , dpDelta    ∷ !Int8
      -- | The decorrelation weight for channel A: @-1024 <= weight <=
      --   1024@
      , dpWeightA  ∷ !(STRef s Int16)
      -- | The decorrelation weight for channel B: @-1024 <= weight <=
      --   1024@
      , dpWeightB  ∷ !(STRef s Int16)
      -- | The decorrelation samples for channel A.
      , dpSamplesA ∷ !(v s Int32)
      -- | The decorrelation samples for channel B.
      , dpSamplesB ∷ !(v s Int32)
      }

-- |General function to perform stereo decorrelation pass on specified
-- buffer (although since this is the reverse function it might
-- technically be called /correlation/ instead). This version handles
-- all sample resolutions and weight deltas. The dpSamplesX data is
-- *not* returned normalized for term values 1-8, so it should be
-- normalized if it is going to be used to call this function again.
decorrStereoPass ∷ ∀v s. MV.MVector v Int32
                 ⇒ DecorrPass v s
                 → v s Int32
                 → ST s ()
{-# INLINEABLE decorrStereoPass #-}
decorrStereoPass dp buffer
    | MV.length buffer `rem` 2 ≢ 0
        = fail "decorrStereoPass: assertion failed: MV.length buffer `rem` 2 ≡ 0"
    | dpTerm dp ≡ 17
        = flip mapM_ [0 .. (MV.length buffer `div` 2) - 1] $ \n →
              do a0 ← MV.unsafeRead (dpSamplesA dp) 0
                 a1 ← MV.unsafeRead (dpSamplesA dp) 1
                 MV.unsafeWrite (dpSamplesA dp) 1 a0
                 tmpA ← MV.unsafeRead buffer (n ⋅ 2)
                 weiA ← readSTRef (dpWeightA dp)
                 let samA  = 2 ⋅ a0 - a1
                     samA' = applyWeight weiA samA + tmpA
                 MV.unsafeWrite (dpSamplesA dp) 0 samA'
                 MV.unsafeWrite buffer (n ⋅ 2) samA'
                 writeSTRef (dpWeightA dp) $ updateWeight weiA (dpDelta dp) tmpA samA

                 b0 ← MV.unsafeRead (dpSamplesB dp) 0
                 b1 ← MV.unsafeRead (dpSamplesB dp) 1
                 MV.unsafeWrite (dpSamplesB dp) 1 b0
                 tmpB ← MV.unsafeRead buffer (n ⋅ 2 + 1)
                 weiB ← readSTRef (dpWeightB dp)
                 let samB  = 2 ⋅ b0 - b1
                     samB' = applyWeight weiB samB + tmpB
                 MV.unsafeWrite (dpSamplesB dp) 0 samB'
                 MV.unsafeWrite buffer (n ⋅ 2 + 1) samB'
                 writeSTRef (dpWeightB dp) $ updateWeight weiB (dpDelta dp) tmpB samB
    | otherwise
        = fail "FIXME"

applyWeight ∷ Int16 → Int32 → Int32
{-# INLINE applyWeight #-}
applyWeight weight sample
    | sample `shiftR` 0xFFFF ≡ 0
        = (fromIntegral weight ⋅ sample + 512) `shiftR` 10
    | otherwise
        = ( (((sample .&. 0xFFFF) ⋅ fromIntegral weight) `shiftR` 9) +
            (((sample .&. complement 0xFFFF) `shiftR` 9) ⋅ fromIntegral weight) +
            1
          ) `shiftR` 1

updateWeight ∷ Int16 → Int8 → Int32 → Int32 → Int16
{-# INLINE updateWeight #-}
updateWeight weight δ source result
    | source ≢ 0 ∧ result ≢ 0
        = let s = fromIntegral $ (source `xor` result) `shiftR` 31
          in
            (fromIntegral δ `xor` s) + (weight - s)
    | otherwise
        = weight