{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
module Synthesizer.LLVM.Simple.Value (
   T, decons,
   twoPi,
   lift0, lift1, lift2,
   constantValue, constant,
   Flatten(flatten, unfold), flattenTraversable, unfoldFunctor,
   ) where

import qualified LLVM.Extra.ScalarOrVector as SoV
import qualified LLVM.Extra.Arithmetic as A

import LLVM.Core
   (CodeGenFunction,
    Value, valueOf, CmpRet,
    IsArithmetic, IsConst, IsType, IsFloating, )
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 as Frame
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 System.IO.Unsafe (unsafePerformIO, )

import NumericPrelude.Numeric
import NumericPrelude.Base


{-
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 (Value a)}

decons :: T a -> (forall r. LLVM.CodeGenFunction r (Value 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 (Value a)) -> T a
consUnique code0 =
   unsafePerformIO $
   fmap (consKey code0) Vault.newKey

consKey :: (forall r. Compute r (Value a)) -> Vault.Key (Value 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 (IsArithmetic a, IsType a) => Additive.C (T a) where
   zero = constantValue (LLVM.value LLVM.zero)
   (+) = lift2 A.add
   (-) = lift2 A.sub
   negate = lift1 A.neg

instance (IsArithmetic a, SoV.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 (SoV.PseudoModule a v, SoV.IntegerConstant a) =>
      Module.C (T a) (T v) where
   (*>) = lift2 SoV.scale

instance (IsArithmetic a, SoV.IntegerConstant a) => Enum (T a) where
   succ x = x + one
   pred x = x - one
   fromEnum _ = error "CodeGenFunction Value: fromEnum"
   toEnum = 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 (IsFloating a, SoV.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 (IsFloating a, SoV.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 (IsFloating a, SoV.RationalConstant a, Trans.C a) => Trans.C (T a) where
   pi = constant 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 ::
   (IsFloating a, IsConst a, Trans.C a) =>
   T a
twoPi = constant (2*pi)
{-
twoPi ::
   (Cmp a b, P.Floating a, IsConst a, IsFloating a) =>
   Compute r a
twoPi = P.fromInteger 2 P.* P.pi
-}

instance (SoV.Real a, SoV.IntegerConstant a, CmpRet a Bool) =>
      Absolute.C (T a) where
   abs = lift1 A.abs
   signum = lift1 Frame.signum

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

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

instance
   (Algebraic.C a, NormedEuc.Sqr (T a) (T v), SoV.RationalConstant a, IsFloating a) =>
      NormedEuc.C (T a) (T v) where
   norm = NormedEuc.defltNorm

{-
instance (Ring.C a, IsArithmetic a, IsConst a, CmpRet a Bool) =>
      NormedMax.C (T a) (T a) where
   norm = lift1 A.abs
-}


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

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

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


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

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


class Flatten value register | value -> register where
   flatten :: value -> CodeGenFunction r register
   unfold :: register -> value

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

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


instance (Flatten ah al, Flatten bh bl) =>
      Flatten (ah,bh) (al,bl) where
   flatten (a,b) =
      liftM2 (,) (flatten a) (flatten b)
   unfold (a,b) =
      (unfold a, unfold b)

instance (Flatten ah al, Flatten bh bl, Flatten ch cl) =>
      Flatten (ah,bh,ch) (al,bl,cl) where
   flatten (a,b,c) =
      liftM3 (,,) (flatten a) (flatten b) (flatten c)
   unfold (a,b,c) =
      (unfold a, unfold b, unfold c)

instance Flatten v r =>
      Flatten (Stereo.T v) (Stereo.T r) where
   flatten = flattenTraversable
   unfold = unfoldFunctor

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

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


instance Flatten (T a) (Value a) where
   flatten = decons
   unfold  = constantValue
instance Flatten () () where
   flatten = return
   unfold  = id