{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{- |
A wrapper around @(->)@ or @Causal.Process@
that adds amplitude handling to the Arrow paradigm.
This wrapper unifies "Synthesizer.Dimensional.Map"
and "Synthesizer.Dimensional.Causal.Process".
-}
module Synthesizer.Dimensional.Arrow where

import qualified Synthesizer.Dimensional.Signal.Private as SigA
import qualified Synthesizer.Dimensional.Amplitude.Flat as Flat
import qualified Synthesizer.Dimensional.Amplitude as Amp
import qualified Synthesizer.Dimensional.Rate as Rate

import qualified Synthesizer.Causal.Arrow as CausalArrow

import qualified Control.Arrow as Arrow
import qualified Control.Category as Category
import Control.Arrow (Arrow, ArrowLoop, )
import Control.Category (Category, )

import Control.Applicative (Applicative, 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 NumericPrelude (one)
import Prelude hiding (map, id, fst, snd, )



{- |
Note that @amp@ can also be a pair of amplitudes
or a more complicated ensemble of amplitudes.
-}
newtype T amp0 amp1 arrow =
   Cons (amp0 -> (arrow, amp1))


{-
It is tempting to declare a rate parameter for the process type,
instead of putting the rate phantom into the arrow.
However, Map would then be defined as

> type Map amp0 amp1 yv0 yv1 = T (forall rate. rate) amp0 amp1 (yv0->yv1)@

which is at least ugly. Even more, in module Rate we would need

> class Applicable process signal | signal -> process
> instance Applicable (Phantom s) (Phantom s)
> instance Applicable (forall process. process) (Actual rate)

and this is not possible, at all.

With the current approach we can have
both generic apply functions and generic arrow combinators.
-}

class CausalArrow.C arrow => Applicable arrow rate

instance Applicable (->) rate




infixl 9 `apply`

{-# INLINE apply #-}
apply ::
   (SigG2.Transform sig yv0 yv1, Applicable arrow rate) =>
   T amp0 amp1 (arrow yv0 yv1) ->
   SigA.T rate amp0 (sig yv0) ->
   SigA.T rate amp1 (sig yv1)
apply (Cons f) (SigA.Cons rate xAmp samples) =
   let (arrow, yAmp) = f xAmp
   in  SigA.Cons rate yAmp (CausalArrow.apply arrow samples)

{-# INLINE applyFlat #-}
applyFlat ::
   (Flat.C yv0 amp0,
    SigG2.Transform sig yv0 yv1, Applicable arrow rate) =>
   T (Amp.Flat yv0) amp1 (arrow yv0 yv1) ->
   SigA.T rate amp0 (sig yv0) ->
   SigA.T rate amp1 (sig yv1)
applyFlat f =
   apply (canonicalizeFlat >>> f)

{-# INLINE canonicalizeFlat #-}
canonicalizeFlat ::
   (Flat.C y flat, Arrow arrow) =>
   T flat (Amp.Flat y) (arrow y y)
canonicalizeFlat =
   Cons $ \ amp -> (Arrow.arr (Flat.amplifySample amp), Amp.Flat)


{-# INLINE applyConst #-}
applyConst ::
   (Amp.C amp1, Ring.C y0, CausalArrow.C arrow) =>
   T (Amp.Numeric amp0) amp1 (arrow y0 yv1) ->
   amp0 ->
   SigA.T (Rate.Phantom s) amp1 (Sig.T yv1)
applyConst (Cons f) x =
   let (arrow, yAmp) = f (Amp.Numeric x)
   in  SigA.Cons Rate.Phantom yAmp
          (CausalArrow.apply arrow (Sig.repeat one))


infixl 0 $/:, $/-

{-# INLINE ($/:) #-}
($/:) ::
   (Applicative f, SigG2.Transform sig yv0 yv1,
    Applicable arrow rate) =>
   f (T amp0 amp1 (arrow yv0 yv1)) ->
   f (SigA.T rate amp0 (sig yv0)) ->
   f (SigA.T rate amp1 (sig yv1))
($/:) = liftA2 apply

{-# INLINE ($/-) #-}
($/-) ::
   (Amp.C amp1, Functor f, Ring.C y0, CausalArrow.C arrow) =>
   f (T (Amp.Numeric amp0) amp1 (arrow y0 yv1)) ->
   amp0 ->
   f (SigA.T (Rate.Phantom s) amp1 (Sig.T yv1))
($/-) p x = fmap (flip applyConst x) p



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



{-# INLINE id #-}
id ::
   (Category arrow) =>
   T amp amp (arrow yv yv)
id =
   Cons (\amp -> (Category.id, amp))


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

(>>>) = compose

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


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

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

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

(***) = split

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

(&&&) = fanout




-- * map functions

{- |
This function can be abused to bring the amplitudes out of order.
So be careful!
-}
independentMap ::
   (Arrow arrow) =>
   (amp0 -> amp1) -> (yv0 -> yv1) ->
   T amp0 amp1 (arrow yv0 yv1)
independentMap f g =
   Cons (\amp -> (Arrow.arr g, f amp))

double ::
   (Arrow arrow) =>
   T amp (amp, amp) (arrow y (y, y))
double =
   let aux = \x -> (x, x)
   in  independentMap aux aux

{-# INLINE forceDimensionalAmplitude #-}
forceDimensionalAmplitude ::
   (Dim.C v, Field.C y, Module.C y yv, Arrow arrow) =>
   DN.T v y ->
   T (Amp.Dimensional v y) (Amp.Dimensional v y) (arrow yv yv)
forceDimensionalAmplitude ampOut =
   Cons $ \(Amp.Numeric ampIn) ->
      (Arrow.arr (DN.divToScalar ampIn ampOut *>),
       Amp.Numeric ampOut)



{- |
I will call the connection from input to output amplitudes of type @amp@,
the looping channel.
It is essential, that the looping channel decouples output from input amplitude.
You can achieve this by inserting one of the @forceAmplitude@ functions
somewhere in the looping channel.
-}
{-# INLINE loop #-}
loop ::
   (ArrowLoop arrow) =>
   T (restAmpIn, amp) (restAmpOut, amp)
     (arrow (restSampIn, yv) (restSampOut, yv)) ->
   T restAmpIn restAmpOut (arrow restSampIn restSampOut)
loop (Cons f) =
   Cons $ \restAmpIn ->
      let (arrow, (restAmpOut, amp)) = f (restAmpIn, amp)
      in  (Arrow.loop arrow, restAmpOut)


{-# INLINE loopVolume #-}
loopVolume ::
   (Field.C y, Module.C y yv, Dim.C v,
    ArrowLoop arrow) =>
   DN.T v y ->
   T (restAmpIn, Amp.Dimensional v y)
     (restAmpOut, Amp.Dimensional v y)
     (arrow (restSampIn, yv) (restSampOut, yv)) ->
   T restAmpIn restAmpOut
     (arrow restSampIn restSampOut)
loopVolume ampIn f =
   loop (f >>> second (forceDimensionalAmplitude ampIn))


{-# INLINE loop2Volume #-}
loop2Volume ::
   (Field.C y0, Module.C y0 yv0, Dim.C v0,
    Field.C y1, Module.C y1 yv1, Dim.C v1,
    ArrowLoop arrow) =>
   (DN.T v0 y0, DN.T v1 y1) ->
   T (restAmpIn,  (Amp.Numeric (DN.T v0 y0), Amp.Numeric (DN.T v1 y1)))
     (restAmpOut, (Amp.Numeric (DN.T v0 y0), Amp.Numeric (DN.T v1 y1)))
     (arrow (restSampIn,  (yv0,yv1))
            (restSampOut, (yv0,yv1))) ->
   T restAmpIn restAmpOut
     (arrow restSampIn restSampOut)
loop2Volume (ampIn0,ampIn1) f =
   loop (f >>> second
      (forceDimensionalAmplitude ampIn0 ***
       forceDimensionalAmplitude ampIn1))