{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module Synthesizer.LLVM.Generator.Render where

import qualified Synthesizer.LLVM.Causal.Parameterized as Parameterized
import qualified Synthesizer.LLVM.Generator.Source as Source
import qualified Synthesizer.LLVM.Storable.ChunkIterator as ChunkIt
import qualified Synthesizer.LLVM.Storable.LazySizeIterator as SizeIt
import qualified Synthesizer.LLVM.EventIterator as EventIt
import Synthesizer.LLVM.Generator.Private (T(Cons))

import qualified Synthesizer.LLVM.Frame.Stereo as Stereo
import qualified Synthesizer.LLVM.Storable.Vector as SVU
import qualified Synthesizer.LLVM.ForeignPtr as ForeignPtr
import qualified Synthesizer.LLVM.ConstantPiece as Const

import qualified Synthesizer.PiecewiseConstant.Signal as PC
import qualified Synthesizer.Causal.Class as CausalClass

import qualified LLVM.DSL.Execution as Exec
import qualified LLVM.DSL.Expression as Expr
import LLVM.DSL.Expression (Exp(Exp))

import qualified LLVM.Extra.Multi.Value.Storable as Storable
import qualified LLVM.Extra.Multi.Value.Marshal as Marshal
import qualified LLVM.Extra.Multi.Value as MultiValue
import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.MaybeContinuation as MaybeCont
import qualified LLVM.Extra.Maybe as Maybe
import qualified LLVM.Extra.Control as C
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Tuple as Tuple

import qualified LLVM.Core as LLVM

import qualified Type.Data.Num.Decimal as TypeNum

import qualified Data.StorableVector.Lazy as SVL
import qualified Data.StorableVector.Base as SVB
import qualified Data.StorableVector as SV

import qualified Data.EventList.Relative.BodyTime as EventList
import qualified Numeric.NonNegative.Wrapper as NonNeg
import qualified Numeric.NonNegative.Chunky as NonNegChunky

import Control.Monad (join)
import Control.Applicative (liftA3)

import Foreign.ForeignPtr (touchForeignPtr)
import Foreign.Ptr (Ptr)

import Data.Foldable (traverse_)
import Data.Tuple.Strict (mapPair, mapTriple)
import Data.Word (Word, Word8, Word32)

import qualified System.Unsafe as Unsafe


foreign import ccall safe "dynamic" derefFillPtr ::
   Exec.Importer (LLVM.Ptr param -> Word -> Ptr struct -> IO Word)


compile ::
   (Storable.C a, MultiValue.T a ~ value,
    Marshal.C param, Marshal.Struct param ~ paramStruct) =>
   (Exp param -> T value) ->
   IO (LLVM.Ptr paramStruct -> Word -> Ptr a -> IO Word)
compile sig =
   Exec.compile "signal" $
   Exec.createFunction derefFillPtr "fill" $ \paramPtr size bPtr ->
   case sig (Exp (Memory.load paramPtr)) of
      Cons next start stop -> do
         (global,s) <- start
         local <- LLVM.alloca
         (pos,_) <- Storable.arrayLoopMaybeCont size bPtr s $ \ ptri s0 -> do
            (y,s1) <- next global local s0
            MaybeCont.lift $ Storable.store y ptri
            return s1
         stop global
         return pos

runAux ::
   (Marshal.C p, Storable.C a, MultiValue.T a ~ value) =>
   (Exp p -> T value) -> IO (IO () -> Int -> p -> IO (SV.Vector a))
runAux sig = do
   fill <- compile sig
   return $ \final len param ->
      Marshal.with param $ \paramPtr ->
      SVB.createAndTrim len $ \ptr -> do
         n <- fill paramPtr (fromIntegral len) ptr
         final
         return $ fromIntegral n

run_ ::
   (Marshal.C p, Storable.C a, MultiValue.T a ~ value) =>
   (Exp p -> T value) -> IO (Int -> p -> IO (SV.Vector a))
run_ = fmap ($ return ()) . runAux


foreign import ccall safe "dynamic" derefStartPtr ::
   Exec.Importer (LLVM.Ptr param -> IO (LLVM.Ptr globalState))

foreign import ccall safe "dynamic" derefStopPtr ::
   Exec.Importer (LLVM.Ptr globalState -> IO ())

foreign import ccall safe "dynamic" derefChunkPtr ::
   Exec.Importer (LLVM.Ptr globalState -> Word -> Ptr a -> IO Word)


type MemoryPtr a = LLVM.Ptr (Memory.Struct a)

type WithGlobalState param = LLVM.Struct (param, ())

type Pair a b = LLVM.Struct (a,(b,()))
type Triple a b c = LLVM.Struct (a,(b,(c,())))

tripleStruct ::
   (LLVM.IsSized a, LLVM.IsSized b, LLVM.IsSized c) =>
   LLVM.Value a -> LLVM.Value b -> LLVM.Value c ->
   LLVM.CodeGenFunction r (LLVM.Value (Triple a b c))
tripleStruct a b c = do
   s0 <- LLVM.insertvalue Tuple.undef a TypeNum.d0
   s1 <- LLVM.insertvalue s0 b TypeNum.d1
   LLVM.insertvalue s1 c TypeNum.d2

{- |
This is a pretty ugly hack, but its seems to be the least ugly one.
We need to solve the following problem:
We have a function of type @Exp param -> T value@.
This means that all methods in @T value@ depend on @Exp param@.
We need to choose one piece of LLVM code in @Exp param@
that generates appropriate code for all methods in @T value@.
If we access a function parameter via @Memory.load paramPtr@
this means that all methods must end up in the same LLVM function
in order to access this parameter.
Thus I have to put all functionality in one LLVM function
and then the three functions in 'compileChunky'
jump into the handler function with a 'Word8' code
specifying the actual sub-routine.
We need to squeeze all possible inputs and outputs
through one function interface.

However, since the handler is marked as internal
the optimizer inlines it in the three functions from 'compileChunky'
and eliminates dead code.
This way, we end up with the code that we would have written otherwise.

The alternative would be to construct @T value@ multiple times.
Due to existential quantification we cannot prove
that the pointer types of different methods match,
so we need to cast pointers.
However, with the current approach we also have to do that.
-}
compileHandler ::
   (Marshal.C param, Marshal.Struct param ~ paramStruct,
    Storable.C a, MultiValue.T a ~ value) =>
   (Exp param -> T value) ->
   LLVM.CodeGenModule
      (LLVM.Function
         (Word8 -> LLVM.Ptr paramStruct -> Word -> Ptr a ->
          IO (Pair (LLVM.Ptr (WithGlobalState paramStruct)) Word)))
compileHandler sig =
   LLVM.createNamedFunction LLVM.InternalLinkage "handlesignal" $
   \phase paramPtr loopLen bufferPtr ->
   case sig $ Exp (Memory.load paramPtr) of
      Cons next start stop -> do
         paramGlobalStatePtr <- LLVM.bitcast paramPtr

         let create = do
               newParamGlobalStatePtr <- LLVM.malloc
               (global,state) <- start
               flip LLVM.store newParamGlobalStatePtr =<<
                  join
                     (liftA3 tripleStruct
                        (LLVM.load paramPtr)
                        (Memory.compose global)
                        (Memory.compose state))
               newOpaqueParamGlobalStatePtr <-
                  LLVM.bitcast
                     (newParamGlobalStatePtr `asTypeOf` paramGlobalStatePtr)
               LLVM.insertvalue Tuple.undef
                  newOpaqueParamGlobalStatePtr TypeNum.d0

         let delete = do
               globalPtr <-
                  LLVM.getElementPtr0 paramGlobalStatePtr (TypeNum.d1, ())
               stop =<< Memory.load globalPtr
               LLVM.free paramGlobalStatePtr
               return Tuple.undef

         let fill = do
               globalPtr <-
                  LLVM.getElementPtr0 paramGlobalStatePtr (TypeNum.d1, ())
               statePtr <-
                  LLVM.getElementPtr0 paramGlobalStatePtr (TypeNum.d2, ())
               global <- Memory.load globalPtr
               sInit <- Memory.load statePtr
               local <- LLVM.alloca
               (pos,sExit) <-
                  Storable.arrayLoopMaybeCont loopLen bufferPtr sInit $
                     \ ptr s0 -> do
                  (y,s1) <- next global local s0
                  MaybeCont.lift $ Storable.store y ptr
                  return s1
               Memory.store (Maybe.fromJust sExit) statePtr
               LLVM.insertvalue Tuple.undef pos TypeNum.d1

         doCreate <- A.cmp LLVM.CmpEQ (LLVM.valueOf 0) phase
         doDelete <- A.cmp LLVM.CmpEQ (LLVM.valueOf 1) phase
         C.ret =<<
            (C.ifThenElse doCreate create $
             C.ifThenElse doDelete delete fill)

compileChunky ::
   (LLVM.IsSized paramStruct, LLVM.Value (LLVM.Ptr paramStruct) ~ pPtr,
    Memory.C state, Memory.Struct state ~ stateStruct,
    Memory.C global, Memory.Struct global ~ globalStruct,
    Triple paramStruct globalStruct stateStruct ~ triple,
    LLVM.IsSized local,
    Storable.C a, MultiValue.T a ~ value) =>
   (forall r z. (Tuple.Phi z) =>
    pPtr -> global -> LLVM.Value (LLVM.Ptr local) ->
    () -> state -> MaybeCont.T r z (value, state)) ->
   (forall r. pPtr -> LLVM.CodeGenFunction r (global, state)) ->
   (forall r. pPtr -> global -> LLVM.CodeGenFunction r ()) ->
   IO (LLVM.Ptr paramStruct -> IO (LLVM.Ptr triple),
       Exec.Finalizer triple,
       LLVM.Ptr triple -> Word -> Ptr a -> IO Word)
compileChunky next start stop =
   Exec.compile "signal-chunky" $
   liftA3 (,,)
      (Exec.createFunction derefStartPtr "startsignal" $
         \paramPtr -> do
            paramGlobalStatePtr <- LLVM.malloc
            (global,state) <- start paramPtr
            flip LLVM.store paramGlobalStatePtr =<<
               join
                  (liftA3 tripleStruct
                     (LLVM.load paramPtr)
                     (Memory.compose global)
                     (Memory.compose state))
            return paramGlobalStatePtr)
      (Exec.createFinalizer derefStopPtr "stopsignal" $
         \paramGlobalStatePtr -> do
            paramPtr <-
               LLVM.getElementPtr0 paramGlobalStatePtr (TypeNum.d0, ())
            stop paramPtr =<<
               Memory.load =<<
               LLVM.getElementPtr0 paramGlobalStatePtr (TypeNum.d1, ())
            LLVM.free paramGlobalStatePtr)
      (Exec.createFunction derefChunkPtr "fillsignal" $
         \paramGlobalStatePtr loopLen ptr -> do
            paramPtr <-
               LLVM.getElementPtr0 paramGlobalStatePtr (TypeNum.d0, ())
            global <-
               Memory.load =<<
               LLVM.getElementPtr0 paramGlobalStatePtr (TypeNum.d1, ())
            statePtr <-
               LLVM.getElementPtr0 paramGlobalStatePtr (TypeNum.d2, ())
            sInit <- Memory.load statePtr
            local <- LLVM.alloca
            (pos,sExit) <-
               Storable.arrayLoopMaybeCont loopLen ptr sInit $
                  \ ptri s0 -> do
               (y,s1) <- next paramPtr global local () s0
               MaybeCont.lift $ Storable.store y ptri
               return s1
            Memory.store (Maybe.fromJust sExit) statePtr
            return pos)


runChunkyAux ::
   (Storable.C a, MultiValue.T a ~ value, Marshal.C p) =>
   (Exp p -> T value) -> IO (IO () -> SVL.ChunkSize -> p -> IO (SVL.Vector a))
runChunkyAux sig = do
   paramd <-
      Parameterized.fromProcessPtr "Signal.run" (CausalClass.fromSignal . sig)
   case paramd of
      Parameterized.Cons next start stop -> do
         (startFunc,stopFunc,fill) <- compileChunky next start stop
         return $ \final (SVL.ChunkSize size) p -> do
            statePtr <- ForeignPtr.newParamMV stopFunc startFunc p

            let go =
                  Unsafe.interleaveIO $ do
                     v <-
                        ForeignPtr.with statePtr $ \sptr ->
                        SVB.createAndTrim size $
                        fmap (fromIntegral :: Word -> Int) .
                        fill sptr (fromIntegral size)
                     (if SV.length v > 0
                        then fmap (v:)
                        else id) $
                        (if SV.length v < size
                           then final >> return []
                           else go)
            fmap SVL.fromChunks go

runChunky ::
   (Storable.C a, MultiValue.T a ~ value, Marshal.C p) =>
   (Exp p -> T value) -> IO (SVL.ChunkSize -> p -> IO (SVL.Vector a))
runChunky = fmap ($ return ()) . runChunkyAux


runChunkyOnVector ::
   (Storable.C a, MultiValue.T a ~ al) =>
   (Storable.C b, MultiValue.T b ~ bl) =>
   (T al -> T bl) ->
   IO (SVL.ChunkSize -> SV.Vector a -> IO (SVL.Vector b))
runChunkyOnVector sig = do
   f <- runChunkyAux (sig . Source.storableVector)
   return $ \chunkSize av -> do
      let (fp,ptr,l) = SVU.unsafeToPointers av
      f (touchForeignPtr fp) chunkSize (Source.consStorableVector ptr l)


class Run f where
   type DSL f
   type Shape f
   build ::
      (Marshal.C p) =>
      (Exp p -> DSL f) -> IO (IO (p, IO ()) -> Shape f -> f)

instance (Storable.C a) => Run (SVL.Vector a) where
   type DSL (SVL.Vector a) = T (MultiValue.T a)
   type Shape (SVL.Vector a) = SVL.ChunkSize
   build =
      fmap (\f create shape -> Unsafe.performIO $ buildIOGen f create shape) .
      runChunkyAux

instance (Storable.C a) => Run (SV.Vector a) where
   type DSL (SV.Vector a) = T (MultiValue.T a)
   type Shape (SV.Vector a) = Int
   build =
      fmap (\f create shape -> Unsafe.performIO $ buildIOGen f create shape) .
      runAux

instance (RunIO a) => Run (IO a) where
   type DSL (IO a) = T (DSL_IO a)
   type Shape (IO a) = ShapeIO a
   build = buildIO

instance (RunArg a, Run f) => Run (a -> f) where
   type DSL (a -> f) = DSLArg a -> DSL f
   type Shape (a -> f) = Shape f
   build sig =
      case buildArg of
         BuildArg pass createA -> do
            f <- build (Expr.uncurry $ \p -> sig p . pass)
            return $ \createP shape av ->
               f (do (p,finalP) <- createP
                     (pa,finalA) <- createA av
                     return ((p,pa), finalA >> finalP))
                  shape


class RunIO a where
   type DSL_IO a
   type ShapeIO a
   buildIO ::
      (Marshal.C p) =>
      (Exp p -> T (DSL_IO a)) -> IO (IO (p, IO ()) -> ShapeIO a -> IO a)

instance (Storable.C a) => RunIO (SVL.Vector a) where
   type DSL_IO (SVL.Vector a) = MultiValue.T a
   type ShapeIO (SVL.Vector a) = SVL.ChunkSize
   buildIO = fmap buildIOGen . runChunkyAux

instance (Storable.C a) => RunIO (SV.Vector a) where
   type DSL_IO (SV.Vector a) = MultiValue.T a
   type ShapeIO (SV.Vector a) = Int
   buildIO = fmap buildIOGen . runAux

buildIOGen ::
   (Monad m) => (final -> shape -> p -> m a) -> m (p, final) -> shape -> m a
buildIOGen f create shape = do (p,final) <- create; f final shape p


data BuildArg a =
   forall al. Marshal.C al =>
   BuildArg (Exp al -> DSLArg a) (a -> IO (al, IO ()))

class RunArg a where
   type DSLArg a
   buildArg :: BuildArg a

instance RunArg () where
   type DSLArg () = ()
   buildArg = BuildArg (\ _unit -> ()) (\() -> return ((), return ()))

instance (RunArg a, RunArg b) => RunArg (a,b) where
   type DSLArg (a,b) = (DSLArg a, DSLArg b)
   buildArg =
      case (buildArg,buildArg) of
         (BuildArg passA createA, BuildArg passB createB) ->
            BuildArg
               (mapPair (passA,passB) . Expr.unzip)
               (\(a,b) -> do
                  (pa,finalA) <- createA a
                  (pb,finalB) <- createB b
                  return ((pa,pb), finalB>>finalA))

instance (RunArg a, RunArg b, RunArg c) => RunArg (a,b,c) where
   type DSLArg (a,b,c) = (DSLArg a, DSLArg b, DSLArg c)
   buildArg =
      case (buildArg,buildArg,buildArg) of
         (BuildArg passA createA, BuildArg passB createB,
          BuildArg passC createC) ->
            BuildArg
               (mapTriple (passA,passB,passC) . Expr.unzip3)
               (\(a,b,c) -> do
                  (pa,finalA) <- createA a
                  (pb,finalB) <- createB b
                  (pc,finalC) <- createC c
                  return ((pa,pb,pc), finalC>>finalB>>finalA))

primitiveArg :: (Marshal.C a, DSLArg a ~ Exp a) => BuildArg a
primitiveArg = BuildArg id (\a -> return (a, return ()))

instance RunArg Float where
   type DSLArg Float = Exp Float
   buildArg = primitiveArg

instance RunArg Int where
   type DSLArg Int = Exp Int
   buildArg = primitiveArg

instance RunArg Word where
   type DSLArg Word = Exp Word
   buildArg = primitiveArg

instance RunArg Word32 where
   type DSLArg Word32 = Exp Word32
   buildArg = primitiveArg

instance (RunArg a) => RunArg (Stereo.T a) where
   type DSLArg (Stereo.T a) = Stereo.T (DSLArg a)
   buildArg =
      case buildArg of
         BuildArg pass create ->
            BuildArg
               (fmap pass . Stereo.unExpression)
               (\s -> do
                  pf <- traverse create s
                  return (fst<$>pf, traverse_ snd pf))

instance
   (TypeNum.Natural n, Marshal.C a, LLVM.IsSized (Marshal.Struct a),
    TypeNum.Natural (n TypeNum.:*: LLVM.SizeOf (Marshal.Struct a))) =>
      RunArg (MultiValue.Array n a) where
   type DSLArg (MultiValue.Array n a) = Exp (MultiValue.Array n a)
   buildArg = primitiveArg

instance (Storable.C a) => RunArg (SV.Vector a) where
   type DSLArg (SV.Vector a) = T (MultiValue.T a)
   buildArg =
      BuildArg
         Source.storableVector
         (\av -> do
            let (fp,ptr,l) = SVU.unsafeToPointers av
            return (Source.consStorableVector ptr l, touchForeignPtr fp))

newtype Buffer a = Buffer (SV.Vector a)

buffer :: SV.Vector a -> Buffer a
buffer = Buffer

instance (Storable.C a) => RunArg (Buffer a) where
   type DSLArg (Buffer a) = Exp (Source.StorableVector a)
   buildArg =
      BuildArg id
         (\(Buffer av) -> do
            let (fp,ptr,l) = SVU.unsafeToPointers av
            return (Source.consStorableVector ptr l, touchForeignPtr fp))

newDisposeArg ::
   (Marshal.C handle) =>
   (a -> IO handle) -> (handle -> IO ()) ->
   (Exp handle -> DSLArg a) -> BuildArg a
newDisposeArg new dispose fetch =
   BuildArg fetch
      (\x -> do
         it <- new x
         return (it, dispose it))

instance (Storable.C a) => RunArg (SVL.Vector a) where
   type DSLArg (SVL.Vector a) = T (MultiValue.T a)
   buildArg =
      newDisposeArg ChunkIt.new ChunkIt.dispose Source.storableVectorLazy

class TimeInteger int where
   subdivideLong :: EventList.T (NonNeg.T int) a -> EventList.T NonNeg.Int a

instance TimeInteger Int where
   subdivideLong = id

instance TimeInteger Integer where
   subdivideLong = PC.subdivideLongStrict

instance
   (time ~ NonNeg.T int, TimeInteger int, Marshal.C a) =>
      RunArg (EventList.T time a) where
   type DSLArg (EventList.T time a) = T (Const.T (MultiValue.T a))
   buildArg =
      newDisposeArg
         (EventIt.new . subdivideLong) EventIt.dispose Source.eventList

instance (a ~ SVL.ChunkSize) => RunArg (NonNegChunky.T a) where
   type DSLArg (NonNegChunky.T a) = T (Const.T ())
   buildArg =
      newDisposeArg SizeIt.new SizeIt.dispose Source.lazySize

{-
do f <- run (\n -> takeWhile (<*n) (iterate (1+) 0) <> takeWhile (<*n) (iterate (2+) 0)); f SVL.defaultChunkSize (12::Float) :: IO (SVL.Vector Float)
do f <- Sig.run (\n -> Sig.takeWhile (Expr.<*n) (Sig.iterate (1+) 0) <> Sig.takeWhile (Expr.<*n) (Sig.iterate (2+) 0)); f SVL.defaultChunkSize (12::Float) :: IO (SVL.Vector Float)
-}
run :: (Run f) => DSL f -> IO (Shape f -> f)
run sig = do
   act <- build (const sig)
   return $ act (return ((), return ()))