{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
module Synthesizer.LLVM.Plug.Input (
   T(..),
   Default(..),
   rmap,
   split,
   fanout,
   lazySize,
   ignore,
   storableVector,
   piecewiseConstant,
   controllerSet,
   ) where

import qualified Synthesizer.Zip as Zip

import qualified Synthesizer.LLVM.ConstantPiece as Const

import qualified LLVM.Extra.Storable as Storable
import qualified LLVM.Extra.Marshal as Marshal
import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Tuple as Tuple
import qualified LLVM.Extra.Control as C

import qualified LLVM.Core as LLVM

import qualified Type.Data.Num.Decimal as TypeNum
import Type.Base.Proxy (Proxy)

import qualified Synthesizer.MIDI.PiecewiseConstant.ControllerSet as PCS
import qualified Synthesizer.Generic.Signal as SigG
import qualified Data.EventList.Relative.BodyTime as EventListBT
import qualified Data.EventList.Relative.MixedTime as EventListMT
import qualified Data.EventList.Relative.TimeTime as EventListTT

import qualified Numeric.NonNegative.Wrapper as NonNegW

import qualified Synthesizer.LLVM.Storable.Vector as SVU
import qualified Data.StorableVector as SV

import qualified Foreign.Marshal.Array as Array
import qualified Foreign.Marshal.Alloc as Alloc
import qualified Foreign.ForeignPtr as FPtr
import Foreign.Storable (pokeElemOff)

import Control.Applicative (liftA2)

import qualified Data.Map as Map

import Data.Tuple.Strict (mapFst, mapPair, swap, zipPair)
import Data.Word (Word)


{-
This datatype does not provide an early exit option, e.g. by Maybe.T,
since we warrant that the driver function will always
read only as much data as is available.
To this end you must provide a @length@ function
via an instance of 'Synthesizer.Generic.Cut.Read'.
-}
data T a b =
   forall state ioContext parameters.
      (Marshal.C parameters, Memory.C state) =>
   Cons
      (forall r.
       Tuple.ValueOf parameters ->
       state -> LLVM.CodeGenFunction r (b, state))
          -- compute next value
      (forall r.
       Tuple.ValueOf parameters ->
       LLVM.CodeGenFunction r state)
          -- initial state
      (a -> IO (ioContext, parameters))
          {- initialization from IO monad
          This is called once input chunk.
          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 Functor (T a) where
   fmap f (Cons next start create delete) =
      Cons (\p s -> fmap (mapFst f) $ next p s) start create delete


class Default a where
   type Element a :: *
   deflt :: T a (Element a)

instance (Default a, Default b) => Default (Zip.T a b) where
   type Element (Zip.T a b) = (Element a, Element b)
   deflt = split deflt deflt

instance Default SigG.LazySize where
   type Element SigG.LazySize = ()
   deflt = lazySize

instance (Storable.C a) => Default (SV.Vector a) where
   type Element (SV.Vector a) = Tuple.ValueOf a
   deflt = storableVector


{-
This is intentionally restricted to NonNegW.Int aka StrictTimeShort,
since chunks must fit into memory.
If you have good reasons to allow other types,
see the versioning history for an according hack.
-}
instance
   (Storable.C a, Memory.C (Tuple.ValueOf a)) =>
      Default (EventListBT.T NonNegW.Int a) where
   type Element (EventListBT.T NonNegW.Int a) = Tuple.ValueOf a
   deflt = piecewiseConstant



rmap :: (a -> b) -> T b c -> T a c
rmap f (Cons next start create delete) =
   Cons next start (create . f) delete


split :: T a c -> T b d -> T (Zip.T a b) (c,d)
split (Cons nextA startA createA deleteA)
      (Cons nextB startB createB deleteB) = Cons
   (\(parameterA, parameterB) (sa,sb) ->
      liftA2 zipPair
         (nextA parameterA sa)
         (nextB parameterB sb))
   (\(parameterA, parameterB) ->
      liftA2 (,)
         (startA parameterA)
         (startB parameterB))
   (\(Zip.Cons a b) ->
      liftA2 zipPair
         (createA a)
         (createB b))
   (\(ca,cb) ->
      deleteA ca >>
      deleteB cb)

fanout :: T a b -> T a c -> T a (b,c)
fanout f g = rmap (\a -> Zip.Cons a a) $ split f g

lazySize :: T SigG.LazySize ()
lazySize = ignore

ignore :: T a ()
ignore =
   Cons
      (\ _ _ -> return ((), ()))
      return
      (\ _a -> return ((), ()))
      (const $ return ())

storableVector ::
   (Storable.C a, Tuple.ValueOf a ~ value) => T (SV.Vector a) value
storableVector =
   Cons
      (\ _ p -> liftA2 (,) (Storable.load p) (Storable.incrementPtr p))
      return
      (\vec ->
         let (fp,ptr,_l) = SVU.unsafeToPointers vec
         in  return (fp,ptr))
      -- keep the foreign ptr alive
      FPtr.touchForeignPtr

{-
I would like to re-use code from ConstantPiece here.
Unfortunately, it is based on the LLVM-Maybe-Monad,
but here we do not accept early exit.
-}
piecewiseConstant ::
   (Storable.C a, Tuple.ValueOf a ~ value, Memory.C value) =>
   T (EventListBT.T NonNegW.Int a) value
piecewiseConstant =
   expandConstantPieces $
   rmap (uncurry Zip.Cons .
         mapPair
            (SV.pack .
               map ((fromIntegral :: Int -> Word) . NonNegW.toNumber),
             SV.pack) .
         swap . unzip . EventListBT.toPairList) $
   fmap (uncurry Const.Cons) $
   split storableVector storableVector

expandConstantPieces ::
   (Memory.C value) => T events (Const.T value) -> T events value
expandConstantPieces (Cons next start create delete) = Cons
   (\param state0 -> do
      (Const.Cons length1 y1, s1) <-
         C.whileLoopShared state0
            (\(Const.Cons len _y, s) ->
               (A.cmp LLVM.CmpEQ len Tuple.zero,
                next param s))
      length2 <- A.dec length1
      return (y1, (Const.Cons length2 y1, s1)))
   (\param ->
      fmap ((,) (Const.Cons Tuple.zero Tuple.undef)) $
      start param)
   create delete


{- |
Return an Array and not a pointer to an array,
in order to forbid writing to the array.
-}
controllerSet ::
   (TypeNum.Natural n, Storable.C a,
    LLVM.Storable a, Tuple.ValueOf a ~ LLVM.Value a, LLVM.IsSized a) =>
   Proxy n -> T (PCS.T Int a) (LLVM.Value (LLVM.Array n a))
controllerSet pn =
   controllerSetFromSV pn $
   split storableVector $ split storableVector storableVector

controllerSetFromSV ::
   (TypeNum.Natural n,
    LLVM.Storable a, Tuple.ValueOf a ~ LLVM.Value a, LLVM.IsSized a) =>
   Proxy n ->
   T (Zip.T (SV.Vector Word) (Zip.T (SV.Vector Word) (SV.Vector a)))
     (LLVM.Value Word, (LLVM.Value Word, LLVM.Value a)) ->
   T (PCS.T Int a) (LLVM.Value (LLVM.Array n a))
controllerSetFromSV pn (Cons next start create delete) = Cons
   (\((arrPtr, _), param) state0 -> do
      (length2, s2) <-
         C.whileLoopShared state0
            (\(len0, s0) ->
               (A.cmp LLVM.CmpEQ len0 Tuple.zero,
                do ((len1, (i,a)), s1) <- next param s0
                   LLVM.store a =<< LLVM.getElementPtr arrPtr (i, ())
                   return (len1, s1)))
      length3 <- A.dec length2
      arr <- LLVM.load =<< LLVM.bitcast arrPtr
      return (arr, (length3, s2)))
   (\((_, initialTime), param) -> do
      state <- start param
      return (initialTime, state))

   (\pcs ->
      EventListMT.switchTimeL
         (\initialTime bt -> do
            (context, param) <-
               create
                  (uncurry Zip.Cons .
                   mapPair
                     (SV.pack,
                      uncurry Zip.Cons . mapPair (SV.pack, SV.pack). unzip) .
                   unzip .
                   map (\((i,a),len) ->
                      (fromIntegral len :: Word,
                       (fromIntegral i :: Word, a))) .
                   EventListBT.toPairList $
                   bt)

            -- FIXME: handle memory exhaustion
            let n = TypeNum.integralFromProxy pn
            arr <- Array.mallocArray n
            flip mapM_ (Map.toList $ PCS.initial pcs) $ \(i,a) ->
               if i >= n
                 then error "Plug.Input.controllerSet: array too small"
                 else pokeElemOff arr i a

            return
               ((arr, context),
                ((LLVM.fromPtr arr, fromIntegral initialTime :: Word), param)))
            {-
            It would be more elegant,
            if we could pass Arrays around just like Vectors.

            return (context, ((sampleArray (\i -> maybe Tuple.undef Tuple.valueOf $ Map.lookup i (PCS.initial pcs)), time), param)))
            -}
         (EventListTT.flatten (PCS.stream pcs)))
   (\(arr, context) ->
      Alloc.free arr >> delete context)

{-
We might provide a plug that maps from a sequence of time-stamped controller events
to a stream of (Array Controller Value).
This way, we could select controllers more easily from within an causal arrow.
The disadvantage is, that MIDI controller numbers are then hard-wired into the arrow.
Instead we could use a stream of (Array Index Value)
and a global mapping (Array Controller (Maybe Index)).
This way would both save memory and make the controller numbers exchangeable.
We also have to cope with initialization of values
and have to assert that the exponential function
is computed only once per constant piece in controllerExponential.
-}