{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
module Synthesizer.LLVM.Plug.Input where

import qualified Synthesizer.Zip as Zip

import qualified Synthesizer.LLVM.ConstantPiece as Const

import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Class as Class
import qualified LLVM.Extra.Control as C

import qualified LLVM.Core as LLVM
import LLVM.Extra.Class (MakeValueTuple, ValueTuple, )

import qualified Types.Data.Num as TypeNum

import Control.Applicative (liftA2, )

import qualified Data.Map as Map

import Data.Tuple.HT (mapFst, mapPair, swap, )

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 (Storable, pokeElemOff, )

import Data.Word (Word32, )


{-
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 paramTuple.
      (Storable paramTuple,
       MakeValueTuple paramTuple,
       Memory.C (ValueTuple paramTuple),
       Memory.C state) =>
   Cons
      (forall r.
       ValueTuple paramTuple ->
       state -> LLVM.CodeGenFunction r (b, state))
          -- compute next value
      (forall r.
       ValueTuple paramTuple ->
       LLVM.CodeGenFunction r state)
          -- initial state
      (a -> IO (ioContext, paramTuple))
          {- 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 a, MakeValueTuple a, Memory.C (Class.ValueTuple a)) =>
      Default (SV.Vector a) where
   type Element (SV.Vector a) = Class.ValueTuple 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 a, MakeValueTuple a, Memory.C (ValueTuple a)) =>
      Default (EventListBT.T NonNegW.Int a) where
   type Element (EventListBT.T NonNegW.Int a) = ValueTuple 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) (sa0,sb0) -> do
      (a,sa1) <- nextA parameterA sa0
      (b,sb1) <- nextB parameterB sb0
      return ((a,b), (sa1,sb1)))
   (\(parameterA, parameterB) ->
      liftA2 (,)
         (startA parameterA)
         (startB parameterB))
   (\(Zip.Cons a b) -> do
      (ca,paramA) <- createA a
      (cb,paramB) <- createB b
      return ((ca,cb), (paramA, paramB)))
   (\(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 a, MakeValueTuple a, ValueTuple a ~ value, Memory.C value) =>
   T (SV.Vector a) value
storableVector =
   Cons
      (\ _ p ->
         liftA2 (,)
            (Memory.load p)
            (A.advanceArrayElementPtr 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 a, MakeValueTuple a, ValueTuple a ~ value,
    Memory.C value) =>
   T (EventListBT.T NonNegW.Int a) value
piecewiseConstant =
   case rmap (uncurry Zip.Cons .
              mapPair (SV.pack . map ((fromIntegral :: Int -> Word32) .
                                      NonNegW.toNumber),
                       SV.pack) .
              swap . unzip . EventListBT.toPairList) $
        fmap (uncurry Const.Cons) $
        split storableVector storableVector of
      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 Class.zeroTuple,
                      next param s))
            length2 <- A.dec length1
            return (y1, (Const.Cons length2 y1, s1)))
         (\param ->
            fmap ((,) (Const.Cons Class.zeroTuple Class.undefTuple)) $
            start param)
         create delete


{- |
Return an Array and not a pointer to an array,
in order to forbid writing to the array.
-}
controllerSet ::
   (TypeNum.NaturalT n,
    Storable a, MakeValueTuple a, ValueTuple a ~ LLVM.Value a,
    Memory.FirstClass a, LLVM.IsSized a, LLVM.IsSized (Memory.Stored a)) =>
   n -> T (PCS.T Int a) (LLVM.Value (LLVM.Array n a))
controllerSet n =
   case storableVector of
      Cons next start create delete -> Cons
         (\((arrPtr, _), param) state0 -> do
            (length2, s2) <-
               C.whileLoopShared state0
                  (\(len0, s0) ->
                     (A.cmp LLVM.CmpEQ len0 Class.zeroTuple,
                      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
                        (SV.pack .
                         map (\((i,a),len) ->
                            (fromIntegral len :: Word32,
                             (fromIntegral i :: Word32, a))) .
                         EventListBT.toPairList $
                         bt)

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

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

                  return (context, ((sampleArray (\i -> maybe Class.undefTuple Class.valueTupleOf $ 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.
-}