{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{- |
Exponential curve with controllable delay.
-}
module Synthesizer.LLVM.Generator.Exponential2 (
   Parameter,
   parameter,
   parameterPlain,
   causalP,

   ParameterPacked,
   parameterPacked,
   parameterPackedPlain,
   causalPackedP,
   ) where

import qualified Synthesizer.LLVM.CausalParameterized.Process as CausalP
import qualified Synthesizer.LLVM.Simple.Value as Value
import qualified Synthesizer.LLVM.Parameter as Param
import qualified Synthesizer.LLVM.SerialVector as Serial

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

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

import qualified Data.TypeLevel.Num as TypeNum

import Foreign.Storable (Storable, )
import qualified Foreign.Storable
-- import qualified Foreign.Storable.Record as Store
import qualified Foreign.Storable.Traversable as Store

import qualified Control.Applicative as App
import qualified Data.Foldable as Fold
import qualified Data.Traversable as Trav
import Control.Applicative (liftA2, (<*>), )
import Control.Arrow (arr, (^<<), (&&&), )
import Control.Monad (liftM2, )

import qualified Algebra.Transcendental as Trans
-- import qualified Algebra.Field as Field
-- import qualified Algebra.Ring as Ring

import NumericPrelude.Numeric
import NumericPrelude.Base


newtype Parameter a = Parameter a
   deriving (Show, Storable)


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

instance App.Applicative Parameter where
   {-# INLINE pure #-}
   pure x = Parameter x
   {-# INLINE (<*>) #-}
   Parameter f <*> Parameter k =
      Parameter (f k)

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

instance Trav.Traversable Parameter where
   {-# INLINE sequenceA #-}
   sequenceA (Parameter k) =
      fmap Parameter k


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

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

instance Class.Zero a => Class.Zero (Parameter a) where
   zeroTuple = Class.zeroTuplePointed

instance
      (Memory.C a s, IsSized s ss) =>
      Memory.C (Parameter a) s where
   load = Memory.loadNewtype Parameter
   store = Memory.storeNewtype (\(Parameter k) -> k)
   decompose = Memory.decomposeNewtype Parameter
   compose = Memory.composeNewtype (\(Parameter k) -> k)


{-
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 ah al) =>
      Class.MakeValueTuple (Parameter ah) (Parameter al) where
   valueTupleOf = Class.valueTupleOfFunctor


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


instance (Vector.ShuffleMatch n v) =>
      Vector.ShuffleMatch n (Parameter v) where
   shuffleMatch = Vector.shuffleMatchTraversable

instance (Vector.Access n a v) =>
      Vector.Access n (Parameter a) (Parameter v) where
   insert  = Vector.insertTraversable
   extract = Vector.extractTraversable


parameter ::
   (Trans.C a, SoV.RationalConstant a, IsFloating a) =>
   Value a ->
   CodeGenFunction r (Parameter (Value a))
parameter halfLife =
   Value.flatten $ parameterPlain $
   Value.constantValue halfLife

parameterPlain ::
   (Trans.C a) =>
   a -> Parameter a
parameterPlain halfLife =
   Parameter $ 0.5 ** recip halfLife


causalP ::
   (Memory.FirstClass a am, IsSized a as, IsSized am ams, SoV.IntegerConstant a,
    IsArithmetic a,
    Storable a, Class.MakeValueTuple a (Value a)) =>
   Param.T p a ->
   CausalP.T p (Parameter (Value a)) (Value a)
causalP initial =
   CausalP.loop initial
      (arr snd &&& CausalP.zipWithSimple (\(Parameter a) -> A.mul a))


data ParameterPacked a =
   ParameterPacked {ppFeedback, ppCurrent :: a}


instance Functor ParameterPacked where
   {-# INLINE fmap #-}
   fmap f p = ParameterPacked
      (f $ ppFeedback p) (f $ ppCurrent p)

instance App.Applicative ParameterPacked where
   {-# INLINE pure #-}
   pure x = ParameterPacked x x
   {-# INLINE (<*>) #-}
   f <*> p = ParameterPacked
      (ppFeedback f $ ppFeedback p)
      (ppCurrent f $ ppCurrent p)

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

instance Trav.Traversable ParameterPacked where
   {-# INLINE sequenceA #-}
   sequenceA p =
      liftA2 ParameterPacked
         (ppFeedback p) (ppCurrent p)


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

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

instance Class.Zero a => Class.Zero (ParameterPacked a) where
   zeroTuple = Class.zeroTuplePointed


{-
storeParameter ::
   Storable a => Store.Dictionary (ParameterPacked a)
storeParameter =
   Store.run $
   liftA2 ParameterPacked
      (Store.element ppFeedback)
      (Store.element ppCurrent)

instance Storable a => Storable (ParameterPacked a) where
   sizeOf    = Store.sizeOf storeParameter
   alignment = Store.alignment storeParameter
   peek      = Store.peek storeParameter
   poke      = Store.poke storeParameter
-}

instance Storable a => Storable (ParameterPacked a) where
   sizeOf    = Store.sizeOf
   alignment = Store.alignment
   peek      = Store.peekApplicative
   poke      = Store.poke


memory ::
   (Memory.C l s, IsSized s ss) =>
   Memory.Record r (LLVM.Struct (s, (s, ()))) (ParameterPacked l)
memory =
   liftA2 ParameterPacked
      (Memory.element ppFeedback TypeNum.d0)
      (Memory.element ppCurrent  TypeNum.d1)

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


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

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

instance (Class.MakeValueTuple ah al) =>
      Class.MakeValueTuple (ParameterPacked ah) (ParameterPacked al) where
   valueTupleOf = Class.valueTupleOfFunctor


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


instance (Vector.ShuffleMatch m v) =>
      Vector.ShuffleMatch m (ParameterPacked v) where
   shuffleMatch = Vector.shuffleMatchTraversable

instance (Vector.Access m a v) =>
      Vector.Access m (ParameterPacked a) (ParameterPacked v) where
   insert  = Vector.insertTraversable
   extract = Vector.extractTraversable



withSize ::
   (n -> m (param (Value (Vector n a)))) ->
   m (param (Value (Vector n a)))
withSize f = f undefined

parameterPacked ::
   (Trans.C a, SoV.RationalConstant a, IsFloating a,
    IsPrimitive a, TypeNum.Pos n) =>
   Value a ->
   CodeGenFunction r (ParameterPacked (Value (Vector n a)))
parameterPacked halfLife = withSize $ \n -> do
   feedback <-
      SoV.replicate =<<
      A.pow (valueOf 0.5) =<<
      A.fdiv (valueOf $ fromIntegral $ TypeNum.toInt n) halfLife
   k <-
      A.pow (valueOf 0.5) =<<
      A.fdiv (valueOf 1) halfLife
   current <-
      Vector.iterate (A.mul k) (valueOf 1)
   return $ ParameterPacked feedback current
{-
   Value.flatten $ parameterPackedPlain $
   Value.constantValue halfLife
-}

withSizePlain ::
   (n -> param (Vector n a)) ->
   param (Vector n a)
withSizePlain f = f undefined

parameterPackedPlain ::
   (Trans.C a,
    TypeNum.Pos n) =>
   a -> ParameterPacked (Vector n a)
parameterPackedPlain halfLife =
   withSizePlain $ \n ->
   ParameterPacked
      (Serial.replicate (0.5 ** (fromIntegral (TypeNum.toInt n) / halfLife)))
      (LLVM.vector $ iterate (0.5 ** recip halfLife *) one)


causalPackedP ::
   (Memory.FirstClass a am, IsSized a as, IsSized am ams, SoV.IntegerConstant a,
    Storable a, Class.MakeValueTuple a (Value a),
    IsArithmetic a, TypeNum.Pos n,
    IsPrimitive a,  TypeNum.Mul n as  vs,  TypeNum.Pos vs,
    IsPrimitive am, TypeNum.Mul n ams vms, TypeNum.Pos vms) =>
   Param.T p a ->
   CausalP.T p (ParameterPacked (Value (Vector n a))) (Value (Vector n a))
causalPackedP initial =
   CausalP.loop
      (Serial.replicate ^<< initial)
      (CausalP.mapSimple $
       \(p, s0) -> liftM2 (,)
          (A.mul (ppCurrent p) s0)
          (A.mul (ppFeedback p) s0))