{-# LANGUAGE FlexibleContexts #-}
module Synthesizer.Dimensional.Causal.Process (
   module Synthesizer.Dimensional.Causal.Process,
   Flat(Flat),
   ) where

import qualified Synthesizer.Dimensional.Arrow as ArrowD
import qualified Synthesizer.Dimensional.Map as Map

import qualified Synthesizer.Dimensional.RatePhantom as RP
import qualified Synthesizer.Dimensional.Amplitude.Signal as SigA
import qualified Synthesizer.Dimensional.Abstraction.HomogeneousGen as Hom
import qualified Synthesizer.Dimensional.Amplitude as Amplitude
import qualified Synthesizer.Dimensional.Abstraction.Flat as Flat

import Synthesizer.Dimensional.Amplitude (Flat(Flat))

import qualified Synthesizer.Causal.Process as Causal

import Control.Applicative (Applicative, liftA, liftA2, )

import qualified Synthesizer.State.Signal as Sig
import qualified Synthesizer.Generic.Signal2 as SigG2

import qualified Algebra.Module as Module
import qualified Algebra.Field  as Field
import qualified Algebra.Ring   as Ring
import Algebra.Module ((*>))

import qualified Number.DimensionTerm        as DN
import qualified Algebra.DimensionTerm       as Dim

import qualified Control.Arrow as Arrow

import Data.Tuple.HT as TupleHT (mapSnd, )

import NumericPrelude (one)
import Prelude hiding (map, id, fst, snd, )



{-
TODO:
This differs from Rate.Process and Amplitude.Signal in the following way:
Here we expect, that @amp@ are types that contain physical units,
whereas Rate.Process.T has separate type variables for unit and values.
Thus Rate.Process.T is limited to DimensionalTerm numbers.
We need the additional flexibility here
because @amp@ can also be a pair of amplitudes
or a more complicated ensemble of amplitudes.

Should the 's' parameter be provided by a RatePhantom?
There are causal processes, namely @map@s,
which do not depend on the sample rate.
For these it would make sense to omit the 's'.
On the other hand what other wrappers could be useful?
RateWrapper around T is not sensible,
since it provides the sample rate as value,
not as an input parameter.
Note, that RatePhantom has the signal element type as parameter.
This would accidentally match here, but is it sensible?
-}
newtype T s amp0 amp1 yv0 yv1 =
   Cons (amp0 -> (amp1, Causal.T yv0 yv1))

instance ArrowD.C (T s) where
   map = map
   (>>>) = (>>>)
   first = first
   second = second
   (***) = (***)
   (&&&) = (&&&)


{-# INLINE apply #-}
apply ::
   (Hom.C amp0 Sig.T signal0, Hom.C amp1 Sig.T signal1) =>
   T s amp0 amp1 yv0 yv1 ->
   RP.T s signal0 yv0 -> RP.T s signal1 yv1
apply (Cons f) x =
   let (xAmp, samples) = Hom.unwrap x
       (yAmp, causal) = f xAmp
   in  Hom.wrap (yAmp, Causal.apply causal samples)

{-# INLINE applyGeneric #-}
applyGeneric ::
   (Hom.C amp0 storage signal0, Hom.C amp1 storage signal1,
    SigG2.Transform storage yv0 yv1) =>
   T s amp0 amp1 yv0 yv1 ->
   RP.T s signal0 yv0 -> RP.T s signal1 yv1
applyGeneric (Cons f) x =
   let (xAmp, samples) = Hom.unwrap x
       (yAmp, causal) = f xAmp
   in  Hom.wrap (yAmp, Causal.applyGeneric causal samples)


{-# INLINE applyConst #-}
applyConst :: (Dim.C v0, Dim.C v1, Ring.C y0) =>
   T s (DN.T v0 y0) (DN.T v1 y1) y0 yv1 ->
   DN.T v0 y0 -> SigA.R s v1 y1 yv1
applyConst (Cons f) x =
   let (yAmp, causal) = f x
   in  SigA.fromSamples yAmp (Causal.applyConst causal one)


infixl 0 $/:, $/-

{-# INLINE ($/:) #-}
($/:) :: (Dim.C v0, Dim.C v1, Applicative f) =>
   f (T s (DN.T v0 y0) (DN.T v1 y1) yv0 yv1) ->
   f (SigA.R s v0 y0 yv0) -> f (SigA.R s v1 y1 yv1)
($/:) = liftA2 apply

{-# INLINE ($/-) #-}
($/-) :: (Dim.C v0, Dim.C v1, Applicative f, Ring.C y0) =>
   f (T s (DN.T v0 y0) (DN.T v1 y1) y0 yv1) ->
   DN.T v0 y0 -> f (SigA.R s v1 y1 yv1)
($/-) p x = liftA (flip applyConst x) p


infixl 9 `apply`, `applyFst`, `applyFlat`, `applyFlatFst`

{-# INLINE applyFst #-}
applyFst, applyFst' :: (Dim.C v) =>
   T s (DN.T v y, restAmpIn) restAmpOut (yv, restSampIn) restSampOut ->
   SigA.R s v y yv ->
   T s restAmpIn restAmpOut restSampIn restSampOut
applyFst c x = c <<< feedFst x

applyFst' (Cons f) x =
   Cons $ \yAmp ->
      let (zAmp, causal) = f (SigA.amplitude x, yAmp)
      in  (zAmp, Causal.applyFst causal (SigA.samples x))


{-# INLINE feedFst #-}
feedFst :: (Dim.C v) =>
   SigA.R s v y yv ->
   T s restAmp (DN.T v y, restAmp) restSamp (yv, restSamp)
feedFst x =
   Cons $ \yAmp ->
      ((SigA.amplitude x, yAmp), Causal.feedFst (SigA.samples x))


{-# INLINE applyFlat #-}
applyFlat :: (Dim.C v1, Flat.C sig yv0) =>
   T s Flat (DN.T v1 y1) yv0 yv1 ->
   RP.T s sig yv0 -> SigA.R s v1 y1 yv1
applyFlat (Cons f) x =
   let (yAmp, causal) = f Flat
   in  SigA.fromSamples yAmp (Causal.apply causal (Flat.toSamples x))


{-# INLINE applyFlatFst #-}
applyFlatFst, applyFlatFst' :: (Flat.C sig yv) =>
   T s (Flat, restAmpIn) restAmpOut (yv, restSampIn) restSampOut ->
   RP.T s sig yv ->
   T s restAmpIn restAmpOut restSampIn restSampOut
applyFlatFst f x =
   f <<< feedFlatFst x

applyFlatFst' (Cons f) x =
   Cons $ \yAmp ->
      let (zAmp, causal) = f (Flat, yAmp)
      in  (zAmp, Causal.applyFst causal (Flat.toSamples x))

{-# INLINE feedFlatFst #-}
feedFlatFst :: (Flat.C sig yv) =>
   RP.T s sig yv ->
   T s restAmp (Flat, restAmp) restSamp (yv, restSamp)
feedFlatFst x =
   Cons $ \yAmp ->
      ((Flat, yAmp), Causal.feedFst (Flat.toSamples x))



{-# INLINE map #-}
map ::
   Map.T amp0 amp1 yv0 yv1 ->
   T s amp0 amp1 yv0 yv1
map (Map.Cons f) =
   Cons $ mapSnd Causal.map . f


{- |
We restrict the amplitude types to those of class 'Amplitude'.
Otherwise 'mapAmplitude' could be abused
for bringing amplitudes and respective sample values out of sync.
For mapping amplitudes that are nested in some pairs,
use it in combination with 'first' and 'second'.
-}
{-# INLINE mapAmplitude #-}
mapAmplitude ::
   (Amplitude.C amp0, Amplitude.C amp1) =>
   (amp0 -> amp1) ->
   T s amp0 amp1 yv yv
mapAmplitude f =
   Cons $ \ xAmp -> (f xAmp, Causal.id)

{-# INLINE mapAmplitudeSameType #-}
mapAmplitudeSameType ::
   (amp -> amp) ->
   T s amp amp yv yv
mapAmplitudeSameType f =
   Cons $ \ xAmp -> (f xAmp, Causal.id)

{- |
Lift a low-level homogeneous process to a dimensional one.

Note that the @amp@ type variable is unrestricted.
This way we show, that the amplitude is not touched,
which also means that the underlying low-level process must be homogeneous.
-}
{-# INLINE homogeneous #-}
homogeneous ::
   Causal.T yv0 yv1 ->
   T s amp amp yv0 yv1
homogeneous c =
   Cons $ \ xAmp -> (xAmp, c)


infixr 3 ***
infixr 3 &&&
infixr 1 >>>, ^>>, >>^
infixr 1 <<<, ^<<, <<^


{-# INLINE compose #-}
{-# INLINE (>>>) #-}
compose, (>>>) ::
   T s amp0 amp1 yv0 yv1 ->
   T s amp1 amp2 yv1 yv2 ->
   T s amp0 amp2 yv0 yv2
compose (Cons f) (Cons g) =
   Cons $ \ xAmp ->
      let (yAmp, causalXY) = f xAmp
          (zAmp, causalYZ) = g yAmp
      in  (zAmp, Causal.compose causalXY causalYZ)

(>>>) = compose

{-# INLINE (<<<) #-}
(<<<) ::
   T s amp1 amp2 yv1 yv2 ->
   T s amp0 amp1 yv0 yv1 ->
   T s amp0 amp2 yv0 yv2
(<<<) = flip (>>>)


{-# INLINE first #-}
first ::
   T s amp0 amp1 yv0 yv1 ->
   T s (amp0, amp) (amp1, amp) (yv0, yv) (yv1, yv)
first (Cons f) =
   Cons $ \ (xAmp, amp) ->
      let (yAmp, causal) = f xAmp
      in  ((yAmp, amp), Causal.first causal)

{-# INLINE second #-}
second ::
   T s amp0 amp1 yv0 yv1 ->
   T s (amp, amp0) (amp, amp1) (yv, yv0) (yv, yv1)
second (Cons f) =
   Cons $ \ (amp, xAmp) ->
      let (yAmp, causal) = f xAmp
      in  ((amp, yAmp), Causal.second causal)

{-# INLINE split #-}
{-# INLINE (***) #-}
split, (***) ::
   T s amp0 amp1 yv0 yv1 ->
   T s amp2 amp3 yv2 yv3 ->
   T s (amp0, amp2) (amp1, amp3) (yv0, yv2) (yv1, yv3)
split f g =
   compose (first f) (second g)

(***) = split

{-# INLINE fanout #-}
{-# INLINE (&&&) #-}
fanout, (&&&) ::
   T s amp amp0 yv yv0 ->
   T s amp amp1 yv yv1 ->
   T s amp (amp0, amp1) yv (yv0, yv1)
fanout f g =
   compose (map Map.double) (split f g)

(&&&) = fanout


-- * map functions

{-# INLINE (^>>) #-}
-- | Precomposition with a pure function.
(^>>) ::
   Map.T amp0 amp1 yv0 yv1 ->
   T s amp1 amp2 yv1 yv2 ->
   T s amp0 amp2 yv0 yv2
f ^>> a = map f >>> a

{-# INLINE (>>^) #-}
-- | Postcomposition with a pure function.
(>>^) ::
   T s amp0 amp1 yv0 yv1 ->
   Map.T amp1 amp2 yv1 yv2 ->
   T s amp0 amp2 yv0 yv2
a >>^ f = a >>> map f

{-# INLINE (<<^) #-}
-- | Precomposition with a pure function (right-to-left variant).
(<<^) ::
   T s amp1 amp2 yv1 yv2 ->
   Map.T amp0 amp1 yv0 yv1 ->
   T s amp0 amp2 yv0 yv2
a <<^ f = a <<< map f

{-# INLINE (^<<) #-}
-- | Postcomposition with a pure function (right-to-left variant).
(^<<) ::
   Map.T amp1 amp2 yv1 yv2 ->
   T s amp0 amp1 yv0 yv1 ->
   T s amp0 amp2 yv0 yv2
f ^<< a = map f <<< a



{-# INLINE loop #-}
-- loop :: a (b, d) (c, d) -> a b c
loop ::
   (Field.C y, Module.C y yv, Dim.C v) =>
   DN.T v y ->
   T s (restAmpIn, DN.T v y) (restAmpOut, DN.T v y) (restSampIn, yv) (restSampOut, yv) ->
   T s restAmpIn restAmpOut restSampIn restSampOut
loop ampIn (Cons f) =
   Cons $ \restAmpIn ->
      let ((restAmpOut, ampOut), causal) = f (restAmpIn, ampIn)
      in  (restAmpOut,
           Causal.loop (causal Arrow.>>^
              mapSnd (DN.divToScalar ampOut ampIn *>)))

{-# INLINE loop2 #-}
-- loop2 :: a (b, (d,e)) (c, (d,e)) -> a b c
loop2 (amp0,amp1) p =
   loop amp0 $
   loop amp1 $
   (Map.balanceRight ^>> p >>^ Map.balanceLeft)

loop2, loop2' ::
   (Field.C y0, Module.C y0 yv0, Dim.C v0,
    Field.C y1, Module.C y1 yv1, Dim.C v1) =>
   (DN.T v0 y0, DN.T v1 y1) ->
   T s
     (restAmpIn,  (DN.T v0 y0, DN.T v1 y1))
     (restAmpOut, (DN.T v0 y0, DN.T v1 y1))
     (restSampIn,  (yv0,yv1))
     (restSampOut, (yv0,yv1)) ->
   T s restAmpIn restAmpOut restSampIn restSampOut
loop2' ampIn@(ampIn0,ampIn1) (Cons f) =
   Cons $ \restAmpIn ->
      let ((restAmpOut, (ampOut0,ampOut1)), causal) = f (restAmpIn, ampIn)
      in  (restAmpOut,
           Causal.loop (causal Arrow.>>^
              Arrow.second ((DN.divToScalar ampOut0 ampIn0 *>) Arrow.***
                            (DN.divToScalar ampOut1 ampIn1 *>))))



{-# INLINE id #-}
id ::
   T s amp amp yv yv
id =
   homogeneous Causal.id