{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
module Synthesizer.LLVM.Filter.ComplexFirstOrder (
   Parameter, parameter,
   causal, causalP,
   ) where

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

import qualified Synthesizer.LLVM.Frame.Stereo as Stereo
import qualified Synthesizer.LLVM.Complex as Complex

import qualified LLVM.Extra.ScalarOrVector as SoV
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Class as Class
import LLVM.Extra.Class (Undefined, undefTuple, )

import qualified LLVM.Core as LLVM
import LLVM.Core
   (Value, Struct, value, valueOf,
    IsArithmetic, IsFloating, IsSized,
    CodeGenFunction, )
import LLVM.Util.Loop (Phi, phis, addPhis, )

import Data.TypeLevel.Num (d0, d1, d2, )

import qualified Control.Applicative as App
import qualified Data.Foldable as Fold
import qualified Data.Traversable as Trav
import Control.Applicative (liftA2, liftA3, (<*>), )

import qualified Algebra.Transcendental as Trans

import NumericPrelude.Numeric
import NumericPrelude.Base


data Parameter a =
   Parameter a (Complex.T a)

instance Functor Parameter where
   {-# INLINE fmap #-}
   fmap f (Parameter k c) =
      Parameter (f k) (fmap f c)

instance App.Applicative Parameter where
   {-# INLINE pure #-}
   pure x = Parameter x (x Complex.+: x)
   {-# INLINE (<*>) #-}
   Parameter fk fc <*> Parameter pk pc =
      Parameter (fk pk) $
         (Complex.real fc $ Complex.real pc)
         Complex.+:
         (Complex.imag fc $ Complex.imag pc)

instance Fold.Foldable Parameter where
   {-# INLINE foldMap #-}
   foldMap = Trav.foldMapDefault

instance Trav.Traversable Parameter where
   {-# INLINE sequenceA #-}
   sequenceA (Parameter k c) =
      liftA2 Parameter k $
      liftA2 (Complex.+:) (Complex.real c) (Complex.imag c)


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

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


parameterMemory ::
   (Memory.C l s, IsSized s ss) =>
   Memory.Record r (Struct (s, (s, (s, ())))) (Parameter l)
parameterMemory =
   liftA3 (\amp kr ki -> Parameter amp (kr Complex.+: ki))
      (Memory.element (\(Parameter  amp _) -> amp) d0)
      (Memory.element (\(Parameter _amp k) -> Complex.real k) d1)
      (Memory.element (\(Parameter _amp k) -> Complex.imag k) d2)

instance (Memory.C l s, IsSized s ss) =>
      Memory.C (Parameter l) (Struct (s, (s, (s, ())))) where
   load = Memory.loadRecord parameterMemory
   store = Memory.storeRecord parameterMemory
   decompose = Memory.decomposeRecord parameterMemory
   compose = Memory.composeRecord parameterMemory

instance (Value.Flatten ah al) =>
      Value.Flatten (Parameter ah) (Parameter al) where
   flatten = Value.flattenTraversable
   unfold =  Value.unfoldFunctor


parameter, _parameter ::
   (Trans.C a, IsFloating a, SoV.RationalConstant a) =>
   Value a -> Value a -> CodeGenFunction r (Parameter (Value a))
parameter reson freq =
   let amp = recip $ Value.unfold reson
   in  Value.flatten $ Parameter amp $
       Complex.scale (1-amp) $ Complex.cis $
       Value.unfold freq * Value.twoPi

_parameter reson freq = do
   amp <- A.fdiv (valueOf 1) reson
   k   <- A.sub  (valueOf 1) amp
   w  <- A.mul freq =<< Value.decons Value.twoPi
   kr <- A.mul k =<< A.cos w
   ki <- A.mul k =<< A.sin w
   return (Parameter amp (kr Complex.+: ki))


{-
Synthesizer.Plain.Filter.Recursive.FirstOrderComplex.step
cannot be used directly, because Filt1C has complex amplitude
-}
next, _next ::
   (IsArithmetic a, SoV.IntegerConstant a) =>
   (Parameter (Value a), Stereo.T (Value a)) ->
   Complex.T (Value a) ->
   CodeGenFunction r (Stereo.T (Value a), Complex.T (Value a))
next inp state =
   let stereoFromComplex ::
          Complex.T (Value a) -> Complex.T (Value.T a) ->
          Stereo.T (Value.T a)
       stereoFromComplex _ c =
          Stereo.cons (Complex.real c) (Complex.imag c)
       (Parameter amp k, x) = Value.unfold inp
       xc = Stereo.left x  Complex.+:  Stereo.right x
       y = Complex.scale amp xc + k * Value.unfold state
   in  Value.flatten (stereoFromComplex state y, y)

_next (Parameter amp k, x) s = do
   let kr = Complex.real k
       ki = Complex.imag k
       sr = Complex.real s
       si = Complex.imag s
   yr <- Value.decons $
      Value.lift0 (A.mul (Stereo.left x) amp) +
      Value.lift0 (A.mul kr sr) - Value.lift0 (A.mul ki si)
   yi <- Value.decons $
      Value.lift0 (A.mul (Stereo.right x) amp) +
      Value.lift0 (A.mul kr si) + Value.lift0 (A.mul ki sr)
   return (Stereo.cons yr yi, yr Complex.+: yi)


start ::
   (LLVM.IsType a, SoV.IntegerConstant a) =>
   CodeGenFunction r (Complex.T (Value a))
start =
   return (value LLVM.zero Complex.+: value LLVM.zero)

causal ::
   (IsSized a sa, SoV.IntegerConstant a,
    Memory.FirstClass a am, IsSized am amsize,
    IsFloating a) =>
   Causal.T
      (Parameter (Value a), Stereo.T (Value a))
      (Stereo.T (Value a))
causal =
   Causal.mapAccum next start

causalP ::
   (IsSized a sa, SoV.IntegerConstant a,
    Memory.FirstClass a am, IsSized am amsize,
    IsFloating a) =>
   CausalP.T p
      (Parameter (Value a), Stereo.T (Value a))
      (Stereo.T (Value a))
causalP =
   CausalP.mapAccumSimple next start