{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
module Synthesizer.LLVM.Simple.Value (
   T, decons,
   twoPi, square, sqrt,
   max, min, limit,
   lift0, lift1, lift2,
   unlift0, unlift1, unlift2,
   constantValue, constant,

   Flatten(flatten, unfold), Registers,
   flattenTraversable, unfoldFunctor,
   flattenFunction,
   ) where

import qualified LLVM.Extra.Arithmetic as A

import LLVM.Core (CodeGenFunction, )
import qualified LLVM.Core as LLVM

import qualified Synthesizer.Basic.Phase as Phase

import qualified Data.Vault as Vault
import qualified Control.Monad.Trans.Class as MT
import qualified Control.Monad.Trans.State as MS
import Control.Monad (liftM2, liftM3, )

import qualified Synthesizer.LLVM.Frame.Stereo as Stereo

-- import qualified Algebra.NormedSpace.Maximum   as NormedMax
import qualified Algebra.NormedSpace.Euclidean as NormedEuc
import qualified Algebra.NormedSpace.Sum       as NormedSum

import qualified Algebra.Transcendental as Trans
import qualified Algebra.Algebraic as Algebraic
import qualified Algebra.RealRing as RealRing
import qualified Algebra.Absolute as Absolute
import qualified Algebra.Module as Module
import qualified Algebra.Field as Field
import qualified Algebra.Ring as Ring
import qualified Algebra.Additive as Additive

import qualified Number.Complex as Complex

import qualified Data.Traversable as Trav

import qualified System.Unsafe as Unsafe

import qualified Prelude as P
import NumericPrelude.Numeric hiding (pi, sqrt, )
import NumericPrelude.Base hiding (min, max, )


{-
The @r@ type parameter must be hidden and forall-quantified
because otherwise we would need an impossible type
where we have to quantify for @r@ and @t@ in different scopes
while having a class constraint that involves both of them.

> osci ::
>    (RealRing.C (Value.T r t),
>     IsFirstClass t, IsSized t size, IsFloating t,
>     IsPrimitive t, IsConst t) =>
>    (forall r. Wave.T (Value.T r t) (Value.T r y)) ->
>    t -> t -> T (Value y)

-}
newtype T a = Cons {code :: forall r. Compute r a}

decons :: T a -> (forall r. LLVM.CodeGenFunction r a)
decons value =
   MS.evalStateT (code value) Vault.empty


type Compute r a =
   MS.StateT Vault.Vault (LLVM.CodeGenFunction r) a

consUnique :: (forall r. Compute r a) -> T a
consUnique code0 =
   Unsafe.performIO $
   fmap (consKey code0) Vault.newKey

consKey :: (forall r. Compute r a) -> Vault.Key a -> T a
consKey code0 key =
   Cons (do
      ma <- MS.gets (Vault.lookup key)
      case ma of
         Just a -> return a
         Nothing -> do
            a <- code0
            MS.modify (Vault.insert key a)
            return a)

{- |
We do not require a numeric prelude superclass,
thus also LLVM only types like vectors are instances.
-}
instance (A.Additive a) => Additive.C (T a) where
   zero = constantValue A.zero
   (+) = lift2 A.add
   (-) = lift2 A.sub
   negate = lift1 A.neg

instance (A.PseudoRing a, A.IntegerConstant a) =>
      Ring.C (T a) where
   one = constantValue (A.fromInteger' 1)
   (*) = lift2 A.mul
   fromInteger = constantValue . A.fromInteger'

{-
This instance is enough for Module here.
The difference to Module instances on Haskell tuples is,
that LLVM vectors cannot be nested.
-}
instance (A.PseudoModule a v, A.IntegerConstant a) =>
      Module.C (T a) (T v) where
   (*>) = lift2 A.scale

instance (A.Additive a, A.IntegerConstant a) => Enum (T a) where
   succ x = x + constantValue A.one
   pred x = x - constantValue A.one
   fromEnum _ = error "CodeGenFunction Value: fromEnum"
   toEnum = constantValue . A.fromInteger' . fromIntegral

{-
instance (IsArithmetic a, Cmp a b, Num a, IsConst a) => Real (T a) where
   toRational _ = error "CodeGenFunction Value: toRational"

instance (Cmp a b, Num a, IsConst a, IsInteger a) => Integral (T a) where
   quot = lift2 idiv
   rem  = lift2 irem
   quotRem x y = (quot x y, rem x y)
   toInteger _ = error "CodeGenFunction Value: toInteger"
-}

instance (A.Field a, A.RationalConstant a) => Field.C (T a) where
   (/) = lift2 A.fdiv
   fromRational' =
      constantValue . A.fromRational' . fromRational'

{-
instance (Cmp a b, Fractional a, IsConst a, IsFloating a) => RealFrac (T a) where
   properFraction _ = error "CodeGenFunction Value: properFraction"
-}

instance (A.Transcendental a, A.RationalConstant a) => Algebraic.C (T a) where
   sqrt = lift1 A.sqrt
   root n x = lift2 A.pow x (1 / fromInteger n)
   x^/r = lift2 A.pow x (fromRational' r)

instance (A.Transcendental a, A.RationalConstant a) => Trans.C (T a) where
   pi = lift0 A.pi
   sin = lift1 A.sin
   cos = lift1 A.cos
   (**) = lift2 A.pow
   exp = lift1 A.exp
   log = lift1 A.log

   asin _ = error "LLVM missing intrinsic: asin"
   acos _ = error "LLVM missing intrinsic: acos"
   atan _ = error "LLVM missing intrinsic: atan"


instance
   (A.PseudoRing a, A.Real a, A.IntegerConstant a) =>
      P.Num (T a) where
   fromInteger = constantValue . A.fromInteger'
   (+) = lift2 A.add
   (-) = lift2 A.sub
   (*) = lift2 A.mul
   negate = lift1 A.neg
   abs = lift1 A.abs
   signum = lift1 A.signum

instance
   (A.Field a, A.Real a, A.RationalConstant a) =>
      P.Fractional (T a) where
   fromRational =
      constantValue . A.fromRational' . P.fromRational
   (/) = lift2 A.fdiv

instance
   (A.Transcendental a, A.Real a, A.RationalConstant a) =>
      P.Floating (T a) where
   pi = lift0 A.pi
   sin = lift1 A.sin
   cos = lift1 A.cos
   (**) = lift2 A.pow
   exp = lift1 A.exp
   log = lift1 A.log

   asin _ = error "LLVM missing intrinsic: asin"
   acos _ = error "LLVM missing intrinsic: acos"
   atan _ = error "LLVM missing intrinsic: atan"

   sinh x  = (exp x - exp (-x)) / 2
   cosh x  = (exp x + exp (-x)) / 2
   asinh x = log (x + sqrt (x*x + 1))
   acosh x = log (x + sqrt (x*x - 1))
   atanh x = (log (1 + x) - log (1 - x)) / 2


twoPi ::
   (A.Transcendental a, A.RationalConstant a) =>
   T a
twoPi = 2 * Trans.pi

square ::
   (A.PseudoRing a) =>
   T a -> T a
square = lift1 A.square

{- |
The same as 'Algebraic.sqrt',
but needs only Algebraic constraint, not Transcendental.
-}
sqrt ::
   (A.Algebraic a) =>
   T a -> T a
sqrt = lift1 A.sqrt


min, max :: (A.Real a) => T a -> T a -> T a
min = lift2 A.min
max = lift2 A.max

limit :: (A.Real a) => (T a, T a) -> T a -> T a
limit (l,u) = max l . min u


instance (A.Real a, A.PseudoRing a, A.IntegerConstant a) =>
      Absolute.C (T a) where
   abs = lift1 A.abs
   signum = lift1 A.signum

{-
For useful instances with different scalar and vector type,
we would need a more flexible superclass.
-}
instance (A.Real a, A.IntegerConstant a, A.PseudoModule a a) =>
      NormedSum.C (T a) (T a) where
   norm = lift1 A.abs

instance (A.Real a, A.IntegerConstant a, A.PseudoModule a a) =>
      NormedEuc.Sqr (T a) (T a) where
   normSqr = lift1 A.square

instance
   (NormedEuc.Sqr (T a) (T v),
    A.RationalConstant a, A.Algebraic a) =>
      NormedEuc.C (T a) (T v) where
   norm = lift1 A.sqrt . NormedEuc.normSqr

{-
instance (A.Real a, A.IntegerConstant a, A.PseudoModule a a) =>
      NormedMax.C (T a) (T a) where
   norm = lift1 A.abs
-}


lift0 ::
   (forall r. CodeGenFunction r a) ->
   T a
lift0 f =
   consUnique $ MT.lift $ f

lift1 ::
   (forall r. a -> CodeGenFunction r b) ->
   T a -> T b
lift1 f x =
   consUnique $ MT.lift . f =<< code x

lift2 ::
   (forall r. a -> b -> CodeGenFunction r c) ->
   T a -> T b -> T c
lift2 f x y =
   consUnique $ do
      xv <- code x
      yv <- code y
      MT.lift $ f xv yv


unlift0 ::
   T a ->
   (forall r. CodeGenFunction r a)
unlift0 = decons

unlift1 ::
   (T a -> T b) ->
   (forall r. a -> CodeGenFunction r b)
unlift1 f x =
   decons (f (constantValue x))

unlift2 ::
   (T a -> T b -> T c) ->
   (forall r. a -> b -> CodeGenFunction r c)
unlift2 f x y =
   decons (f (constantValue x) (constantValue y))


constantValue :: a -> T a
constantValue x =
   consUnique (return x)

constant :: (LLVM.IsConst a) => a -> T (LLVM.Value a)
constant = constantValue . LLVM.valueOf


class Flatten value where
   type Registers value :: *
   flatten :: value -> CodeGenFunction r (Registers value)
   unfold :: (Registers value) -> value

flattenTraversable ::
   (Flatten value, Trav.Traversable f) =>
   f value -> CodeGenFunction r (f (Registers value))
flattenTraversable =
   Trav.mapM flatten

unfoldFunctor ::
   (Flatten value, Functor f) =>
   f (Registers value) -> f value
unfoldFunctor =
   fmap unfold


flattenFunction ::
   (Flatten a, Flatten b) =>
   (a -> b) -> (Registers a -> CodeGenFunction r (Registers b))
flattenFunction f =
   flatten . f . unfold

{-
This function is hardly useful,
since most functions are not of type
@(Registers a -> (forall r. CodeGenFunction r (Registers b)))@
but of type
@(forall r. Registers a -> CodeGenFunction r (Registers b))@.
We would also need a method unfoldF.
See ValueUnfoldF for some implementations.

unfoldFunction ::
   (Flatten a, Flatten b) =>
   (Registers a -> (forall r. CodeGenFunction r (Registers b))) -> (a -> b)
unfoldFunction f x =
   unfoldF (f =<< flatten x)
-}


instance (Flatten a, Flatten b) => Flatten (a,b) where
   type Registers (a,b) = (Registers a, Registers b)
   flatten (a,b) =
      liftM2 (,) (flatten a) (flatten b)
   unfold (a,b) =
      (unfold a, unfold b)

instance (Flatten a, Flatten b, Flatten c) => Flatten (a,b,c) where
   type Registers (a,b,c) = (Registers a, Registers b, Registers c)
   flatten (a,b,c) =
      liftM3 (,,) (flatten a) (flatten b) (flatten c)
   unfold (a,b,c) =
      (unfold a, unfold b, unfold c)

instance Flatten a => Flatten (Stereo.T a) where
   type Registers (Stereo.T a) = Stereo.T (Registers a)
   flatten = flattenTraversable
   unfold = unfoldFunctor

instance Flatten a => Flatten (Complex.T a) where
   type Registers (Complex.T a) = Complex.T (Registers a)
   flatten s =
      liftM2 (Complex.+:)
         (flatten $ Complex.real s)
         (flatten $ Complex.imag s)
   unfold = unfoldFunctor

instance (RealRing.C a, Flatten a) => Flatten (Phase.T a) where
   type Registers (Phase.T a) = Registers a
   flatten s =
      flatten $ Phase.toRepresentative s
   unfold s =
      -- could also be unsafeFromRepresentative
      Phase.fromRepresentative $ unfold s


instance Flatten (T a) where
   type Registers (T a) = a
   flatten = decons
   unfold  = constantValue

instance Flatten () where
   type Registers () = ()
   flatten = return
   unfold  = id