{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
module Synthesizer.LLVM.Parameterized.SignalPrivate where

import qualified Synthesizer.LLVM.Simple.SignalPrivate as Sig
import qualified Synthesizer.LLVM.Parameter as Param
import qualified LLVM.Extra.MaybeContinuation as MaybeCont
import qualified LLVM.Extra.Either as Either
import qualified LLVM.Extra.Maybe as Maybe
import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Arithmetic as A

import qualified LLVM.Core as LLVM
import LLVM.Extra.Class (MakeValueTuple, ValueTuple, Undefined, )
import LLVM.Core (CodeGenFunction, )
import LLVM.Util.Loop (Phi, )

import Control.Arrow ((&&&), )
import Control.Monad (liftM, liftM2, )
import Control.Applicative (Applicative, pure, (<*>), )

import Foreign.Storable.Tuple ()
import Foreign.Storable (Storable, )
import Foreign.Ptr (Ptr, )

import Data.Monoid (Monoid, mempty, mappend, )
import Data.Semigroup (Semigroup, (<>), )

import qualified Number.Ratio as Ratio
import qualified Algebra.Field as Field
import qualified Algebra.Ring as Ring
import qualified Algebra.Additive as Additive

import NumericPrelude.Base hiding (and, iterate, map, zip, zipWith, )

import qualified Prelude as P


{-
In this attempt we use a Haskell value as parameter supply.
This is okay, since the Haskell value will be converted to internal parameters
and then to LLVM values only once.
We can even have a storable vector as parameter.
However, this way we cannot easily implement
the Vanilla signal using Parameterized.Value as element type.

This separation is nice for maximum efficiency,
but it cannot be utilized by Generic.Signal methods.
Consider an expression like @iterate ((0.5 ** recip halfLife) *) 1@.
How shall we know, that the sub-expression @(0.5 ** recip halfLife)@
needs to be computated only once?
I do not try to do such optimization, instead I let LLVM do it.
However, this means that parameter initialization
will be performed (unnecessarily) at the beginning of every chunk.
For Generic.Signal method instances
we will always set the @(p -> paramTuple)@ to 'id'.

Could we drop parameterized signals at all
and rely entirely on Causal processes?
Unfortunately 'interpolateConstant' does not fit into the Causal process scheme.
(... although it would be causal for stretching factor being at least one.
It would have to maintain the waiting signal as state,
i.e. the state would grow linearly with time.)
Consider a signal algorithm, where the LFO frequency is a parameter.
-}
data T p a =
   forall context state local ioContext parameters.
      (Storable parameters,
       MakeValueTuple parameters,
       Memory.C (ValueTuple parameters),
       Memory.C context,
       Memory.C state) =>
   Cons
      (forall r c.
       (Phi c) =>
       context -> local -> state -> MaybeCont.T r c (a, state))
          -- compute next value
      (forall r.
       CodeGenFunction r local)
          -- allocate temporary variables before a loop
      (forall r.
       ValueTuple parameters ->
       CodeGenFunction r (context, state))
          -- allocate initial state
      (forall r.
       context -> state ->
       CodeGenFunction r ())
          {- cleanup
          You must make sure to call this
          whenever you allocated context and state with the 'start' function.
          You must call it with the latest state returned from the 'next' function.
          -}
      (p -> IO (ioContext, parameters))
          {- initialization from IO monad
          This will be run within Unsafe.performIO,
          so no observable In/Out actions please!
          -}
      (ioContext -> IO ())
          -- finalization from IO monad, also run within Unsafe.performIO


instance Sig.C (T p) where
   simpleAlloca next alloca0 start =
      Cons
         (\() local -> next local)
         alloca0
         (const $ fmap ((,) ()) start)
         (const $ const $ return ())
         (const $ return ((), ()))
         (const $ return ())

   alter f (Cons next0 alloca0 start0 stop0 create delete) =
      case f (Sig.Core (uncurry next0) return id) of
         Sig.Core next1 start1 stop1 ->
            Cons
               (curry next1)
               alloca0
               (withStart start0 start1)
               (\c -> stop0 c . stop1)
               create delete


withStart ::
   Monad m =>
   (startParam -> m (context, state0)) ->
   (state0 -> m state1) ->
   startParam -> m (context, state1)
withStart start act p = do
   (c,s) <- start p
   liftM ((,) c) $ act s

combineStart ::
   Monad m =>
   (paramA -> m (contextA, stateA)) ->
   (paramB -> m (contextB, stateB)) ->
   (paramA, paramB) -> m ((contextA, contextB), (stateA, stateB))
combineStart startA startB (paramA, paramB) =
   liftM2
      (\(ca,sa) (cb,sb) -> ((ca,cb), (sa,sb)))
      (startA paramA)
      (startB paramB)

combineStop ::
   Monad m =>
   (contextA -> stateA -> m ()) ->
   (contextB -> stateB -> m ()) ->
   (contextA, contextB) -> (stateA, stateB) -> m ()
combineStop stopA stopB (ca, cb) (sa, sb) =
   stopA ca sa >> stopB cb sb

combineCreate ::
   Monad m =>
   (p -> m (ioContextA, contextA)) ->
   (p -> m (ioContextB, contextB)) ->
   p -> m ((ioContextA, ioContextB), (contextA, contextB))
combineCreate createIOContextA createIOContextB p = do
   (ca,paramA) <- createIOContextA p
   (cb,paramB) <- createIOContextB p
   return ((ca,cb), (paramA,paramB))

combineDelete ::
   (Monad m) =>
   (ca -> m ()) -> (cb -> m ()) -> (ca, cb) -> m ()
combineDelete deleteIOContextA deleteIOContextB (ca,cb) =
   deleteIOContextA ca >>
   deleteIOContextB cb



simple ::
   (Storable parameters,
    MakeValueTuple parameters,
    Memory.C (ValueTuple parameters),
    Memory.C context,
    Memory.C state) =>
   (forall r c.
    (Phi c) =>
    context -> state -> MaybeCont.T r c (al, state)) ->
   (forall r.
    ValueTuple parameters ->
    CodeGenFunction r (context, state)) ->
   Param.T p parameters -> T p al
simple f start param =
   Param.with param $ \getParam valueParam ->
   Cons
      (\context () state -> f context state)
      (return ())
      (start . valueParam)
      (const $ const $ return ())
      (return . (,) () . getParam)
      (const $ return ())


constant ::
   (Storable a, MakeValueTuple a, ValueTuple a ~ al,
    Memory.C al) =>
   Param.T p a -> T p al
constant =
   simple
      (\pl () -> return (pl, ()))
      (return . flip (,) ())


map ::
   (Storable ph, MakeValueTuple ph, ValueTuple ph ~ pl, Memory.C pl) =>
   (forall r. pl -> a -> CodeGenFunction r b) ->
   Param.T p ph ->
   T p a -> T p b
map f param =
   Sig.map (uncurry f) . zip (constant param)

-- for backwards compatibility
mapSimple ::
   (forall r. a -> CodeGenFunction r b) ->
   T p a -> T p b
mapSimple = Sig.map


zipWith ::
   (Storable ph, MakeValueTuple ph, ValueTuple ph ~ pl, Memory.C pl) =>
   (forall r. pl -> a -> b -> CodeGenFunction r c) ->
   Param.T p ph ->
   T p a -> T p b -> T p c
zipWith f param as bs =
   map (uncurry . f) param $ zip as bs

zip :: T p a -> T p b -> T p (a,b)
zip (Cons nextA allocaA startA stopA createIOContextA deleteIOContextA)
    (Cons nextB allocaB startB stopB createIOContextB deleteIOContextB) =
   Cons
      (\(parameterA, parameterB) (localA, localB) (sa0,sb0) -> do
         (a,sa1) <-
            MaybeCont.onFail (stopB parameterB sb0) $
            nextA parameterA localA sa0
         (b,sb1) <-
            MaybeCont.onFail (stopA parameterA sa1) $
            nextB parameterB localB sb0
         return ((a,b), (sa1,sb1)))
      (liftM2 (,) allocaA allocaB)
      (combineStart startA startB)
      (combineStop stopA stopB)
      (combineCreate createIOContextA createIOContextB)
      (combineDelete deleteIOContextA deleteIOContextB)

{-
maintained for backwards compatibility
It is a specialisation of Sig.zipWith.
However, we cannot define zipWithSimple = Sig.zipWith,
since Sig.zipWith depends on Applicative.liftA2,
which depends on zipWithSimple.
-}
zipWithSimple ::
   (forall r. a -> b -> CodeGenFunction r c) ->
   T p a -> T p b -> T p c
zipWithSimple f as bs =
   mapSimple (uncurry f) $ zip as bs


instance Functor (T p) where
   fmap f = mapSimple (return . f)

{- |
ZipList semantics
-}
instance Applicative (T p) where
   pure x =
      simple
         (\() () -> return (x, ()))
         (\() -> return ((),()))
         (return ())
   (<*>) = zipWithSimple (\f a -> return (f a))


instance (A.Additive a) => Additive.C (T p a) where
   zero = pure A.zero
   negate = mapSimple A.neg
   (+) = zipWithSimple A.add
   (-) = zipWithSimple A.sub

instance (A.PseudoRing a, A.IntegerConstant a) => Ring.C (T p a) where
   one = pure A.one
   fromInteger n = pure (A.fromInteger' n)
   (*) = zipWithSimple A.mul

instance (A.Field a, A.RationalConstant a) => Field.C (T p a) where
   fromRational' x = pure (A.fromRational' $ Ratio.toRational98 x)
   (/) = zipWithSimple A.fdiv


instance (A.PseudoRing a, A.Real a, A.IntegerConstant a) => P.Num (T p a) where
   fromInteger n = pure (A.fromInteger' n)
   negate = mapSimple A.neg
   (+) = zipWithSimple A.add
   (-) = zipWithSimple A.sub
   (*) = zipWithSimple A.mul
   abs = mapSimple A.abs
   signum = mapSimple A.signum

instance (A.Field a, A.Real a, A.RationalConstant a) => P.Fractional (T p a) where
   fromRational x = pure (A.fromRational' x)
   (/) = zipWithSimple A.fdiv


{- |
For restrictions see 'Sig.append'.
-}
append :: (Phi a, Undefined a) => T p a -> T p a -> T p a
append
      (Cons nextA allocaA startA stopA createIOContextA deleteIOContextA)
      (Cons nextB allocaB startB stopB createIOContextB deleteIOContextB) =
   Cons
      (\parameterB (localA, localB) ecs0 -> MaybeCont.fromMaybe $ do
         ecs1 <-
            Either.run ecs0
               (\(ca, sa0) ->
                  MaybeCont.resolve
                     (nextA ca localA sa0)
                     (fmap Either.right $ startB parameterB)
                     (\(a1,sa1) -> return (Either.left (a1, (ca, sa1)))))
               (return . Either.right)

         Either.run ecs1
            (\(a1,cs1) ->
               return (Maybe.just (a1, Either.left cs1)))
            (\(cb,sb0) ->
               MaybeCont.toMaybe $
               fmap (\(b,sb1) -> (b, Either.right (cb,sb1))) $
               nextB cb localB sb0))
      (liftM2 (,) allocaA allocaB)
      (\(parameterA, parameterB) -> do
         cs <- startA parameterA
         return (parameterB, Either.left cs))
      (\ _parameterB s -> Either.run s (uncurry stopA) (uncurry stopB))
      (combineCreate createIOContextA createIOContextB)
      (combineDelete deleteIOContextA deleteIOContextB)

instance (Phi a, Undefined a) => Semigroup (T p a) where
   (<>) = append

instance (Phi a, Undefined a) => Monoid (T p a) where
   mempty = Sig.empty
   mappend = append


iterate ::
   (Storable ph, MakeValueTuple ph, ValueTuple ph ~ pl, Memory.C pl,
    Storable a,  MakeValueTuple a, ValueTuple a ~ al, Memory.C al) =>
   (forall r. pl -> al -> CodeGenFunction r al) ->
   Param.T p ph ->
   Param.T p a -> T p al
iterate f param initial = simple
   (\pl al0 ->
      MaybeCont.lift $ fmap (\al1 -> (al0,al1)) (f pl al0))
   return
   (param &&& initial)

malloc :: (LLVM.IsSized a) => T p (LLVM.Value (Ptr a))
malloc =
   Cons
      (\ptr () () -> return (ptr, ()))
      (return ())
      (const $ fmap (flip (,) ()) $ LLVM.malloc)
      (\ptr () -> LLVM.free ptr)
      (const $ return ((), ()))
      (const $ return ())