{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Synthesizer.LLVM.Filter.SecondOrder (
   Parameter(Parameter),
   Filt2.c0, Filt2.c1, Filt2.c2, Filt2.d1, Filt2.d2,
   bandpassParameter,
   ParameterStruct, composeParameter, decomposeParameter, -- for cascade
   causal, causalPacked,
   causalP, causalPackedP,
   ) where

import qualified Synthesizer.Plain.Filter.Recursive.SecondOrder as Filt2
import Synthesizer.Plain.Filter.Recursive.SecondOrder (Parameter(Parameter), )

import qualified Synthesizer.Plain.Modifier as Modifier

import qualified Synthesizer.LLVM.CausalParameterized.Process as CausalP
import qualified Synthesizer.LLVM.Causal.Process as Causal
import qualified Synthesizer.LLVM.Frame.SerialVector as Serial
import qualified Synthesizer.LLVM.Simple.Value as Value

import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Class as Class
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Monad as M
import LLVM.Extra.Class (Undefined, undefTuple, )

import qualified LLVM.Core as LLVM
import LLVM.Core (CodeGenFunction, valueOf, )
import LLVM.Util.Loop (Phi, phis, addPhis, )

import qualified Types.Data.Num as TypeNum
import Types.Data.Num (d0, d1, d2, d3, d4, )

import Control.Arrow (arr, (<<<), (&&&), )
import Control.Monad (liftM2, foldM, )
import Control.Applicative (pure, (<*>), )
import Synthesizer.ApplicativeUtility (liftA4, liftA5, )

import NumericPrelude.Numeric
import NumericPrelude.Base


instance (Phi a) => Phi (Parameter a) where
   phis = Class.phisTraversable
   addPhis = Class.addPhisFoldable

instance Undefined a => Undefined (Parameter a) where
   undefTuple = Class.undefTuplePointed

{-
instance LLVM.ValueTuple a => LLVM.ValueTuple (Parameter a) where
   buildTuple f = Class.buildTupleTraversable (LLVM.buildTuple f)

instance LLVM.IsTuple a => LLVM.IsTuple (Parameter a) where
   tupleDesc = Class.tupleDescFoldable
-}

instance (Class.MakeValueTuple a) => Class.MakeValueTuple (Parameter a) where
   type ValueTuple (Parameter a) = Parameter (Class.ValueTuple a)
   valueTupleOf = Class.valueTupleOfFunctor


type ParameterStruct a = LLVM.Struct (a, (a, (a, (a, (a, ())))))

parameterMemory ::
   (Memory.C a) =>
   Memory.Record r (ParameterStruct (Memory.Struct a)) (Parameter a)
parameterMemory =
   liftA5 Parameter
      (Memory.element Filt2.c0 d0)
      (Memory.element Filt2.c1 d1)
      (Memory.element Filt2.c2 d2)
      (Memory.element Filt2.d1 d3)
      (Memory.element Filt2.d2 d4)

decomposeParameter ::
   LLVM.Value (ParameterStruct a) ->
   CodeGenFunction r (Filt2.Parameter (LLVM.Value a))
decomposeParameter param =
   pure Filt2.Parameter
      <*> LLVM.extractvalue param TypeNum.d0
      <*> LLVM.extractvalue param TypeNum.d1
      <*> LLVM.extractvalue param TypeNum.d2
      <*> LLVM.extractvalue param TypeNum.d3
      <*> LLVM.extractvalue param TypeNum.d4

composeParameter ::
   (LLVM.IsSized a) =>
   Filt2.Parameter (LLVM.Value a) ->
   CodeGenFunction r (LLVM.Value (ParameterStruct a))
composeParameter (Filt2.Parameter c0_ c1_ c2_ d1_ d2_) =
   (\param -> LLVM.insertvalue param c0_ TypeNum.d0) =<<
   (\param -> LLVM.insertvalue param c1_ TypeNum.d1) =<<
   (\param -> LLVM.insertvalue param c2_ TypeNum.d2) =<<
   (\param -> LLVM.insertvalue param d1_ TypeNum.d3) =<<
   (\param -> LLVM.insertvalue param d2_ TypeNum.d4) =<<
   return (LLVM.value LLVM.undef)

instance (Memory.C a) => Memory.C (Parameter a) where
   type Struct (Parameter a) = ParameterStruct (Memory.Struct a)
   load = Memory.loadRecord parameterMemory
   store = Memory.storeRecord parameterMemory
   decompose = Memory.decomposeRecord parameterMemory
   compose = Memory.composeRecord parameterMemory


instance (Value.Flatten a) => Value.Flatten (Parameter a) where
   type Registers (Parameter a) = Parameter (Value.Registers a)
   flattenCode = Value.flattenCodeTraversable
   unfoldCode = Value.unfoldCodeTraversable



instance (Phi a) => Phi (Filt2.State a) where
   phis = Class.phisTraversable
   addPhis = Class.addPhisFoldable

instance Undefined a => Undefined (Filt2.State a) where
   undefTuple = Class.undefTuplePointed


type StateStruct a = LLVM.Struct (a, (a, (a, (a, (a, ())))))

stateMemory ::
   (Memory.C a) =>
   Memory.Record r (StateStruct (Memory.Struct a)) (Filt2.State a)
stateMemory =
   liftA4 Filt2.State
      (Memory.element Filt2.u1 d0)
      (Memory.element Filt2.u2 d1)
      (Memory.element Filt2.y1 d2)
      (Memory.element Filt2.y2 d3)


instance (Memory.C a) => Memory.C (Filt2.State a) where
   type Struct (Filt2.State a) = StateStruct (Memory.Struct a)
   load = Memory.loadRecord stateMemory
   store = Memory.storeRecord stateMemory
   decompose = Memory.decomposeRecord stateMemory
   compose = Memory.composeRecord stateMemory

instance (Value.Flatten a) => Value.Flatten (Filt2.State a) where
   type Registers (Filt2.State a) = Filt2.State (Value.Registers a)
   flattenCode = Value.flattenCodeTraversable
   unfoldCode = Value.unfoldCodeTraversable


{-# DEPRECATED bandpassParameter "only for testing, use Universal or Moog filter for production code" #-}
bandpassParameter ::
   (A.Transcendental a, A.RationalConstant a) =>
   a -> a ->
   CodeGenFunction r (Parameter a)
bandpassParameter reson cutoff = do
   rreson <- A.fdiv A.one reson
   k <- A.sub A.one rreson
   k2 <- A.neg =<< A.mul k k
   kcos <-
      A.mul (A.fromInteger' 2) =<< A.mul k =<<
      A.cos =<< A.mul cutoff =<<
      Value.decons Value.twoPi
   return $
      Filt2.Parameter
         rreson A.zero A.zero
         kcos k2

modifier ::
   (a ~ A.Scalar v, A.PseudoModule v, A.IntegerConstant a) =>
   Modifier.Simple
      (Filt2.State (Value.T v))
      (Parameter (Value.T a))
      (Value.T v) (Value.T v)
modifier =
   Filt2.modifier

causal ::
   (Causal.C process,
    a ~ A.Scalar v, A.PseudoModule v, A.IntegerConstant a, Memory.C v) =>
   process (Parameter a, v) v
causal =
   Causal.fromModifier modifier

{-# DEPRECATED causalP "use causal instead" #-}
causalP ::
   (a ~ A.Scalar v, A.PseudoModule v, A.IntegerConstant a, Memory.C v) =>
   CausalP.T p (Parameter a, v) v
causalP =
   CausalP.fromModifier modifier


{-# DEPRECATED causalPackedP "use causalPacked instead" #-}
causalPackedP ::
   (Serial.C v, Serial.Element v ~ a,
    Memory.C v, Memory.C a, A.IntegerConstant v, A.IntegerConstant a,
    A.PseudoRing v, A.PseudoRing a) =>
   CausalP.T p (Parameter a, v) v
causalPackedP = causalPacked

{- |
Vector size must be at least D2.
-}
causalPacked,
  causalRecursivePacked ::
   (Causal.C process,
    Serial.C v, Serial.Element v ~ a,
    Memory.C v, Memory.C a, A.IntegerConstant v, A.IntegerConstant a,
    A.PseudoRing v, A.PseudoRing a) =>
   process (Parameter a, v) v
causalPacked =
   causalRecursivePacked <<<
   (arr fst &&& causalNonRecursivePacked)

_causalRecursivePackedAlt,
  causalNonRecursivePacked ::
   (Causal.C process,
    Serial.C v, Serial.Element v ~ a,
    Memory.C a, A.IntegerConstant v, A.IntegerConstant a,
    A.PseudoRing v, A.PseudoRing a) =>
   process (Parameter a, v) v
causalNonRecursivePacked =
   Causal.mapAccum
      (\(p, v0) (x1,x2) -> do
         (u1n,v1) <- Serial.shiftUp x1 v0
         (u2n,v2) <- Serial.shiftUp x2 v1
         w0 <- A.mul v0 =<< Serial.upsample (Filt2.c0 p)
         w1 <- A.mul v1 =<< Serial.upsample (Filt2.c1 p)
         w2 <- A.mul v2 =<< Serial.upsample (Filt2.c2 p)
         y  <- A.add w0 =<< A.add w1 w2
         return (y, (u1n,u2n)))
      (return (A.zero, A.zero))

{-
A filter of second order can be considered
as the convolution of two filters of first order.

[1,r]*[1,0,r^2] = [1,r,r^2,r^3]
[1,r,r^2,r^3] * [1,s,s^2,s^3]
 = [1,r]*[1,s]*[1,0,r^2]*[1,0,s^2]
     with
       a=r+s
       b=r*s
 = [1,a,b]*[1,0,r^2]*[1,0,s^2]
 = [1,a,b]*[1,0,a^2-2*b,0,b^2]

[1,0,0,0,r^4]*[1,0,0,0,s^4]
 = [1,0,0,0,(a^2-2*b)^2-2*b^2,0,0,0,b^4]
 = [1,0,0,0,a^4-4*a^2*b+2*b^2,0,0,0,b^4]
-}

{-
x = [x0, x1, x2, x3]

filter2 (a,-b) (y1,y2) x
  = [x0 + a*y1 - b*y2,
     x1 + a*x0 + (a^2-b)*y1 - a*b*y2,
     x2 + a*x1 + (a^2-b)*x0 + (a^3-2*a*b)*y1 + (-a^2*b+b^2)*y2,
     x3 + a*x2 + (a^2-b)*x1 + (a^3-2*a*b)*x0 + (a^4-3*a^2*b+b^2)*y1 + (-a^3*b+2*a*b^2)*y2]

(f0x = insert 0 (k*y1) x)
f1x = f0x + a * f0x->1 + b * f0x->2
f2x = f1x + (a^2-2*b) * f1x->2 + b^2 * f1x->4
-}
causalRecursivePacked =
   Causal.mapAccum
      (\(p, x0) y1v -> do
         let size = Serial.size x0

         d1v  <- Serial.upsample (Filt2.d1 p)
         d2v  <- Serial.upsample (Filt2.d2 p)
         d2vn <- A.neg d2v

         y1  <- Serial.extract (valueOf $ fromIntegral size - 1) y1v
         xk1 <-
            Serial.modify (valueOf 0)
               (\u0 -> A.add u0 =<< A.mul (Filt2.d1 p) y1) =<<
            A.add x0 =<< A.mul d2v =<<
            Serial.shiftDownMultiZero (size - 2) y1v

         -- let xk2 = xk1
         xk2 <-
            fmap fst $
            foldM
               (\(y,(a,b)) d ->
                  liftM2 (,)
                     (A.add y =<<
                      M.liftR2 A.add
                         {-
                         Possibility for optimization:
                         In the last step the second operand is a zero vector
                         (LLVM already optimizes this away)
                         and the first operand could be merged
                         with the second operand of the previous step.
                         -}
                         (Serial.shiftUpMultiZero d =<< A.mul y a)
                         (Serial.shiftUpMultiZero (2*d) =<< A.mul y b)) $
                  liftM2 (,)
                     (M.liftR2 A.sub
                         (A.mul a a)
                         (A.mul b (A.fromInteger' 2)))
                     (A.mul b b))
               (xk1,(d1v,d2vn))
               (takeWhile (< size) $ iterate (2*) 1)

         return (xk2, xk2))
      (return A.zero)

_causalRecursivePackedAlt =
   Causal.mapAccum
      (\(p, x0) (x1,x2) -> do
         let size = Serial.size x0
         -- let xk1 = x0
         xk1 <-
            Serial.modify (valueOf 0)
               (\u0 ->
                  A.add u0 =<<
                  M.liftR2 A.add (A.mul (Filt2.d2 p) x2) (A.mul (Filt2.d1 p) x1)) =<<
            Serial.modify (valueOf 1)
               (\u1 -> A.add u1 =<< A.mul (Filt2.d2 p) x1)
            x0

         -- let xk2 = xk1
         d1v <- Serial.upsample (Filt2.d1 p)
         d2v <- Serial.upsample =<< A.neg (Filt2.d2 p)
         xk2 <-
            fmap fst $
            foldM
               (\(y,(a,b)) d ->
                  liftM2 (,)
                     (A.add y =<<
                      M.liftR2 A.add
                         (Serial.shiftUpMultiZero d =<< A.mul y a)
                         (Serial.shiftUpMultiZero (2*d) =<< A.mul y b)) $
                  liftM2 (,)
                     (M.liftR2 A.sub
                         (A.mul a a)
                         (A.mul b (A.fromInteger' 2)))
                     (A.mul b b))
               (xk1,(d1v,d2v))
               (takeWhile (< size) $ iterate (2*) 1)

         y0 <- Serial.extract (valueOf $ fromIntegral size - 1) xk2
         y1 <- Serial.extract (valueOf $ fromIntegral size - 2) xk2
         return (xk2, (y0,y1)))
      (return (A.zero, A.zero))

{-
A filter of second order can also be represented
by a filter of first order with 2x2-matrix coefficients.

filter1 ((d1,d2), (1,0)) (y1,y2) [(x0,0), (x1,0), (x2,0), (x3,0)]

/d1i d2i\ . /d1j d2j\ = /d1i*d1j + d2i  d1i*d2j\
\ 1   0 /   \ 1   0 /   \    d1j            d2j/


With this representation we can also implement filters
with time-variant filter parameters
using time-variant first-order filter.
-}