{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
module Synthesizer.LLVM.Filter.SecondOrderCascade where

import qualified Synthesizer.LLVM.Filter.SecondOrder as Filt2
import qualified Synthesizer.Plain.Filter.Recursive.SecondOrder as Filt2Core

import qualified Synthesizer.LLVM.CausalParameterized.Process as CausalP

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

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

import qualified Data.TypeLevel.Num      as TypeNum
import qualified Data.TypeLevel.Num.Sets as TypeSet

import Data.Word (Word32, )

import qualified Control.Arrow as Arrow
import Control.Arrow ((>>>), (<<<), (&&&), arr, )

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

import NumericPrelude.Numeric
import NumericPrelude.Base


type Parameter n a = LLVM.Array n (Filt2.ParameterStruct a)

newtype ParameterValue n a =
   ParameterValue {parameterValue :: Value (Parameter n a)}
{-
Automatic deriving is not allowed even with GeneralizedNewtypeDeriving
because of IsSized constraint
and it would also be wrong for Functor and friends.
      deriving
         (Phi, Class.Undefined, Class.Zero,
          Functor, App.Applicative, Fold.Foldable, Trav.Traversable)
-}

instance (TypeNum.Nat n, IsSized a s) =>
      Phi (ParameterValue n a) where
   phis bb (ParameterValue r) =
      fmap ParameterValue $ phis bb r
   addPhis bb
        (ParameterValue r)
        (ParameterValue r') =
      addPhis bb r r'

instance (TypeNum.Nat n, IsSized a s) =>
      Class.Undefined (ParameterValue n a) where
   undefTuple = ParameterValue Class.undefTuple

instance (TypeNum.Nat n, IsSized a s) =>
      Class.Zero (ParameterValue n a) where
   zeroTuple = ParameterValue Class.zeroTuple

instance
      (TypeNum.Nat n, Memory.FirstClass a am, IsSized a s, IsSized am ams) =>
      Memory.C (ParameterValue n a) (Parameter n am) where
   load = Memory.loadNewtype ParameterValue
   store = Memory.storeNewtype (\(ParameterValue k) -> k)
   decompose = Memory.decomposeNewtype ParameterValue
   compose = Memory.composeNewtype (\(ParameterValue k) -> k)



withSize ::
   (n -> CausalP.T p (ParameterValue n a, x) y) ->
   CausalP.T p (ParameterValue n a, x) y
withSize f = f undefined

fixSize ::
   n ->
   CausalP.T p (ParameterValue n a, x) y ->
   CausalP.T p (ParameterValue n a, x) y
fixSize _n = id

causalP ::
   (SoV.PseudoModule a v, SoV.IntegerConstant a,
    Memory.FirstClass a am, IsSized a as, IsSized am ams,
    Memory.FirstClass v vm, IsSized v vs, IsSized vm vms,
    IsArithmetic a, TypeSet.Nat n,
    TypeNum.Mul n LLVM.UnknownSize paramSize, TypeSet.Pos paramSize) =>
   CausalP.T p (ParameterValue n a, Value v) (Value v)
causalP =
   withSize $ \n ->
   foldl (\x y -> (arr fst &&& x) >>> y) (arr snd) $
   map
      (\k ->
         Filt2.causalP <<<
         Arrow.first (CausalP.mapSimple
            (\ps -> getStageParameter ps k)))
      (take (TypeNum.toInt n) [0..])

causalPackedP ::
   (IsArithmetic a, SoV.IntegerConstant a,
    Memory.FirstClass a am,
    IsPrimitive a,  IsSized a asize,
    IsPrimitive am, IsSized am amsize,
    TypeNum.Mul d asize vasize, TypeNum.Pos vasize,
    TypeNum.Mul d amsize vmsize, TypeNum.Pos vmsize,
    TypeNum.Pos d,
    TypeNum.Mul n LLVM.UnknownSize paramSize, TypeSet.Pos paramSize,
    TypeSet.Nat n) =>
   CausalP.T p
      (ParameterValue n a, Value (Vector d a)) (Value (Vector d a))
causalPackedP =
   withSize $ \n ->
   foldl (\x y -> (arr fst &&& x) >>> y) (arr snd) $
   map
      (\k ->
         Filt2.causalPackedP <<<
         Arrow.first (CausalP.mapSimple
            (\ps -> getStageParameter ps k)))
      (take (TypeNum.toInt n) [0..])

getStageParameter, getStageParameterMalloc, getStageParameterAlloca ::
   (LLVM.IsFirstClass a, IsSized a as, SoV.IntegerConstant a,
    TypeSet.Nat n, TypeNum.Mul n LLVM.UnknownSize s, TypeSet.Pos s) =>
   ParameterValue n a ->
   Word32 ->
   CodeGenFunction r (Filt2Core.Parameter (Value a))
getStageParameter ps k =
   Filt2.decomposeParameter =<< LLVM.extractvalue (parameterValue ps) k
{-
   Memory.decompose =<<
   flip LLVM.extractvalue k =<<
   Memory.compose ps
-}

{-
Expensive because we need a heap allocation for every sample.
However, we could allocate the memory once in the Causal initialization routine.
-}
getStageParameterMalloc ps k = do
   ptr <- LLVM.malloc
   LLVM.store (parameterValue ps) ptr
   p <-
      Filt2.decomposeParameter
       =<< LLVM.load
       =<< LLVM.getElementPtr0 ptr (valueOf k, ())
   LLVM.free ptr
   return p

{-
With this implementation, LLVM-2.6 generates a stack variable layout
that requires non-aligned access to vector values.
The result is a crash at runtime.
-}
getStageParameterAlloca ps k = do
   ptr <- LLVM.alloca
   LLVM.store (parameterValue ps) ptr
   Filt2.decomposeParameter
    =<< LLVM.load
    =<< LLVM.getElementPtr0 ptr (valueOf k, ())