{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleContexts #-}
module Synthesizer.LLVM.Parameter (
   T,
   ($#),
   get,
   value,
   with,

   Tuple(..),
   withTuple,
   withTuple1,
   withTuple2,

   -- * for implementation of new processes
   word32,
   ) where

import qualified LLVM.Extra.Class as Class
import qualified LLVM.Extra.Memory as Memory
import Foreign.Storable.Tuple ()
import Foreign.Storable (Storable, )

import qualified Algebra.Transcendental as Trans
import qualified Algebra.Algebraic as Algebraic
import qualified Algebra.Field as Field
import qualified Algebra.Ring as Ring
import qualified Algebra.Additive as Additive

import qualified Control.Category as Cat
import qualified Control.Arrow as Arr
import qualified Control.Applicative as App
import Control.Applicative (pure, liftA2, (<$>), )

import qualified Data.Tuple.HT as TupleHT
import Data.Tuple.HT (mapFst, )
import Data.Word (Word32, )

import NumericPrelude.Numeric
import Prelude (Functor, Monad, fmap, error, (.), ($), const, id, )
import qualified Prelude as P


{- |
This data type is for parameters of parameterized signal generators and causal processes.
It is better than using plain functions of type @p -> a@
since it allows for numeric instances
and we can make explicit,
whether a parameter is constant.

We recommend to use parameters for atomic types.
Although a parameter of type @T p (a,b)@ is possible,
it means that the whole parameter is variable
if only one of the pair elements is variable.
This way you may miss optimizations.
-}
data T p a =
   Constant a |
   Variable (p -> a)


get :: T p a -> (p -> a)
get (Constant a) = const a
get (Variable f) = f


{- |
The call @value param v@ requires
that @v@ represents the same value as @valueTupleOf (get param p)@ for some @p@.
However @v@ might be the result of a load operation
and @param@ might be a constant.
In this case it is more efficient to use @valueTupleOf (get param undefined)@
since the constant is translated to an LLVM constant
that allows for certain optimizations.

This is the main function for taking advantage of a constant parameter
in low-level implementations.
For simplicity we do not omit constant parameters in the parameter struct
since this would mean to construct types at runtime and might become ugly.
Instead we just check using 'value' at the according places in LLVM code
whether a parameter is constant
and ignore the parameter from the struct in this case.
In many cases there will be no speed benefit
because the parameter will be loaded to a register anyway.
It can only lead to speed-up if subsequent optimizations
can precompute constant expressions.
Another example is 'drop' where a loop with constant loop count can be generated.
For small loop counts and simple loop bodies the loop might get unrolled.
-}
value ::
   (Class.MakeValueTuple tuple, Class.ValueTuple tuple ~ value) =>
   T p tuple -> value -> value
value (Constant a) _ = Class.valueTupleOf a
value (Variable _) v = v


{- |
This function provides specialised variants of 'get' and 'value',
that use the unit type for constants
and thus save space in parameter structures.
-}
with ::
   (Storable tuple, Class.MakeValueTuple tuple,
    Class.ValueTuple tuple ~ value, Memory.C value) =>
   T p tuple ->
   (forall parameters.
    (Storable parameters,
     Class.MakeValueTuple parameters,
     Memory.C (Class.ValueTuple parameters)) =>
    (p -> parameters) ->
    (Class.ValueTuple parameters -> value) ->
    a) ->
   a
with (Constant a) f = f (const ()) (\() -> Class.valueTupleOf a)
with (Variable v) f = f v id


word32 :: T p Int -> T p Word32
word32 = fmap fromIntegral


infixl 0 $#

($#) :: (T p a -> b) -> (a -> b)
($#) f a = f (pure a)


{- |
@.@ can be used for fetching a parameter from a super-parameter.
-}
instance Cat.Category T where
   id = Variable id
   Constant f . _ = Constant f
   Variable f . Constant a = Constant (f a)
   Variable f . Variable g = Variable (f . g)

{- |
@arr@ is useful for lifting parameter selectors to our parameter type
without relying on the constructor.
-}
instance Arr.Arrow T where
   arr = Variable
   first f = Variable (mapFst (get f))



{- |
Useful for splitting @T p (a,b)@ into @T p a@ and @T p b@
using @fmap fst@ and @fmap snd@.
-}
instance Functor (T p) where
   fmap f (Constant a) = Constant (f a)
   fmap f (Variable g) = Variable (f . g)

{- |
Useful for combining @T p a@ and @T p b@ to @T p (a,b)@
using @liftA2 (,)@.
However, we do not recommend to do so
because the result parameter can only be constant
if both operands are constant.
-}
instance App.Applicative (T p) where
   pure a = Constant a
   Constant f <*> Constant a = Constant (f a)
   f <*> a = Variable (\p -> get f p (get a p))

instance Monad (T p) where
   return = pure
   Constant x >>= f = f x
   Variable x >>= f =
      Variable (\p -> get (f (x p)) p)


instance Additive.C a => Additive.C (T p a) where
   zero = pure zero
   negate = fmap negate
   (+) = liftA2 (+)
   (-) = liftA2 (-)

instance Ring.C a => Ring.C (T p a) where
   one = pure one
   (*) = liftA2 (*)
   x^n = fmap (^n) x
   fromInteger = pure . fromInteger

instance Field.C a => Field.C (T p a) where
   (/) = liftA2 (/)
   recip = fmap recip
   fromRational' = pure . fromRational'

instance Algebraic.C a => Algebraic.C (T p a) where
   x ^/ r = fmap (^/ r) x
   sqrt = fmap sqrt
   root n = fmap (Algebraic.root n)

instance Trans.C a => Trans.C (T p a) where
   pi      = pure   pi
   exp     = fmap   exp
   log     = fmap   log
   logBase = liftA2 logBase
   (**)    = liftA2 (**)
   sin     = fmap   sin
   tan     = fmap   tan
   cos     = fmap   cos
   asin    = fmap   asin
   atan    = fmap   atan
   acos    = fmap   acos
   sinh    = fmap   sinh
   tanh    = fmap   tanh
   cosh    = fmap   cosh
   asinh   = fmap   asinh
   atanh   = fmap   atanh
   acosh   = fmap   acosh


{-
Instances for Haskell98 numeric type classes
that are useful when working together with other libraries on fixed types.
-}
instance P.Eq a => P.Eq (T p a) where
   (==) = error "Synthesizer.LLVM.Parameter: Num instance requires Eq but we cannot define that"

instance P.Show a => P.Show (T p a) where
   show _ = "Synthesizer.LLVM.Parameter"

instance P.Num a => P.Num (T p a) where
   (+) = liftA2 (P.+)
   (-) = liftA2 (P.-)
   (*) = liftA2 (P.*)
   negate = fmap P.negate
   abs = fmap P.abs
   signum = fmap P.signum
   fromInteger = pure . P.fromInteger

instance P.Fractional a => P.Fractional (T p a) where
   (/) = liftA2 (P./)
   fromRational = pure . P.fromRational

instance P.Floating a => P.Floating (T p a) where
   pi = pure P.pi
   exp = fmap P.exp
   sqrt = fmap P.sqrt
   log = fmap P.log
   (**) = liftA2 (P.**)
   logBase = liftA2 P.logBase
   sin = fmap P.sin
   tan = fmap P.tan
   cos = fmap P.cos
   asin = fmap P.asin
   atan = fmap P.atan
   acos = fmap P.acos
   sinh = fmap P.sinh
   tanh = fmap P.tanh
   cosh = fmap P.cosh
   asinh = fmap P.asinh
   atanh = fmap P.atanh
   acosh = fmap P.acosh



class Tuple tuple where
   type Composed tuple :: *
   type Source tuple :: *
   decompose :: T (Source tuple) (Composed tuple) -> tuple

instance Tuple (T p a) where
   type Composed (T p a) = a
   type Source (T p a) = p
   decompose = id

instance (Tuple a, Tuple b, Source a ~ Source b) => Tuple (a,b) where
   type Composed (a,b) = (Composed a, Composed b)
   type Source (a,b) = Source a
   decompose p = (decompose $ P.fst <$> p, decompose $ P.snd <$> p)

instance
   (Tuple a, Tuple b, Tuple c, Source a ~ Source b, Source b ~ Source c) =>
      Tuple (a,b,c) where
   type Composed (a,b,c) = (Composed a, Composed b, Composed c)
   type Source (a,b,c) = Source a
   decompose p =
      (decompose $ TupleHT.fst3 <$> p,
       decompose $ TupleHT.snd3 <$> p,
       decompose $ TupleHT.thd3 <$> p)

{- |
Provide all elements of a nested tuple as separate parameters.

If you do not use one of the tuple elements,
you will get a type error like
@Couldn't match type `Param.Composed t0' with `Int'@.
The problem is that the type checker cannot infer
that an element is a @Parameter.T@ if it remains unused.
-}
withTuple ::
   (Tuple tuple, Source tuple ~ p, Composed tuple ~ p) =>
   (tuple -> f p) -> f p
withTuple f = idFromFunctor $ f . decompose

idFromFunctor :: (T p p -> f p) -> f p
idFromFunctor f = f Cat.id

withTuple1 ::
   (Tuple tuple, Source tuple ~ p, Composed tuple ~ p) =>
   (tuple -> f p a) -> f p a
withTuple1 f = idFromFunctor1 $ f . decompose

idFromFunctor1 :: (T p p -> f p a) -> f p a
idFromFunctor1 f = f Cat.id

withTuple2 ::
   (Tuple tuple, Source tuple ~ p, Composed tuple ~ p) =>
   (tuple -> f p a b) -> f p a b
withTuple2 f = idFromFunctor2 $ f . decompose

idFromFunctor2 :: (T p p -> f p a b) -> f p a b
idFromFunctor2 f = f Cat.id