{-# LANGUAGE Arrows #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

{- |
Utilities to create 'ClSF's.
The fundamental effect that 'ClSF's have is
reading the time information of the clock.
It can be used for many purposes, for example digital signal processing.
-}
module FRP.Rhine.ClSF.Util where

-- base
import Control.Arrow
import Control.Category (Category)
import Control.Category qualified (id)

-- containers
import Data.Sequence

-- transformers
import Control.Monad.Trans.Reader (ask, asks)

-- dunai
import Data.Automaton.Trans.Reader (readerS)

-- simple-affine-space
import Data.VectorSpace

-- time-domain
import Data.TimeDomain

-- rhine
import FRP.Rhine.ClSF.Core
import FRP.Rhine.ClSF.Except
import FRP.Rhine.Clock

-- * Read time information

-- | Read the environment variable, i.e. the 'TimeInfo'.
timeInfo :: (Monad m) => ClSF m cl a (TimeInfo cl)
timeInfo = constM ask

{- | Utility to apply functions to the current 'TimeInfo',
such as record selectors:
@
printAbsoluteTime :: ClSF IO cl () ()
printAbsoluteTime = timeInfoOf absolute >>> arrMCl print
@
-}
timeInfoOf :: (Monad m) => (TimeInfo cl -> b) -> ClSF m cl a b
timeInfoOf f = constM $ asks f

-- | Continuously return the time difference since the last tick.
sinceLastS :: (Monad m) => ClSF m cl a (Diff (Time cl))
sinceLastS = timeInfoOf sinceLast

-- | Continuously return the time difference since clock initialisation.
sinceInitS :: (Monad m) => ClSF m cl a (Diff (Time cl))
sinceInitS = timeInfoOf sinceInit

-- | Continuously return the absolute time.
absoluteS :: (Monad m) => ClSF m cl a (Time cl)
absoluteS = timeInfoOf absolute

-- | Continuously return the tag of the current tick.
tagS :: (Monad m) => ClSF m cl a (Tag cl)
tagS = timeInfoOf tag

{- |
Calculate the time passed since this 'ClSF' was instantiated,
i.e. since the first tick on which this 'ClSF' was run.

This is _not_ the same as 'sinceInitS',
which measures the time since clock initialisation.

For example, the following gives a sawtooth signal:

@
sawtooth = safely $ do
  try $ sinceStart >>> proc time -> do
    throwOn () -< time > 1
    returnA    -< time
  safe sawtooth
@

If you replace 'sinceStart' by 'sinceInitS',
it will usually hang after one second,
since it doesn't reset after restarting the sawtooth.

Even in the absence of conditional activation of 'ClSF's,
there is a difference:
For a clock that doesn't tick at its initialisation time,
'sinceStart' and 'sinceInitS' will have a constant offset of the duration between initialisation time and first tick.
-}
sinceStart :: (Monad m, TimeDomain time) => BehaviourF m time a (Diff time)
sinceStart =
  absoluteS >>> proc time -> do
    startTime <- keepFirst -< time
    returnA -< time `diffTime` startTime

-- * Useful aliases

-- TODO Is it cleverer to generalise to Arrow?

{- | Alias for 'Control.Category.>>>' (sequential composition)
with higher operator precedence, designed to work with the other operators, e.g.:

> clsf1 >-> clsf2 @@ clA |@| clsf3 >-> clsf4 @@ clB

The type signature specialises e.g. to

> (>->) :: Monad m => ClSF m cl a b -> ClSF m cl b c -> ClSF m cl a c
-}
infixr 6 >->

(>->) ::
  (Category cat) =>
  cat a b ->
  cat b c ->
  cat a c
(>->) = (>>>)

-- | Alias for 'Control.Category.<<<'.
infixl 6 <-<

(<-<) ::
  (Category cat) =>
  cat b c ->
  cat a b ->
  cat a c
(<-<) = (<<<)

{- | Output a constant value.
Specialises e.g. to this type signature:

> arr_ :: Monad m => b -> ClSF m cl a b
-}
arr_ :: (Arrow a) => b -> a c b
arr_ = arr . const

-- | The identity synchronous stream function.
clId :: (Monad m) => ClSF m cl a a
clId = Control.Category.id

-- * Basic signal processing components

-- ** Integration and differentiation

{- | The output of @integralFrom v0@ is the numerical Euler integral
   of the input, with initial offset @v0@.
-}
integralFrom ::
  ( Monad m
  , VectorSpace v s
  , s ~ Diff td
  ) =>
  v ->
  BehaviorF m td v v
integralFrom v0 = proc v -> do
  _sinceLast <- timeInfoOf sinceLast -< ()
  sumFrom v0 -< _sinceLast *^ v

-- | Euler integration, with zero initial offset.
integral ::
  ( Monad m
  , VectorSpace v s
  , s ~ Diff td
  ) =>
  BehaviorF m td v v
integral = integralFrom zeroVector

{- | The output of @derivativeFrom v0@ is the numerical derivative of the input,
   with a Newton difference quotient.
   The input is initialised with @v0@.
-}
derivativeFrom ::
  ( Monad m
  , VectorSpace v s
  , s ~ Diff td
  ) =>
  v ->
  BehaviorF m td v v
derivativeFrom v0 = proc v -> do
  vLast <- delay v0 -< v
  TimeInfo {..} <- timeInfo -< ()
  returnA -< (v ^-^ vLast) ^/ sinceLast

-- | Numerical derivative with input initialised to zero.
derivative ::
  ( Monad m
  , VectorSpace v s
  , s ~ Diff td
  ) =>
  BehaviorF m td v v
derivative = derivativeFrom zeroVector

{- | Like 'derivativeFrom', but uses three samples to compute the derivative.
   Consequently, it is delayed by one sample.
-}
threePointDerivativeFrom ::
  ( Monad m
  , VectorSpace v s
  , s ~ Diff td
  , Num s
  ) =>
  -- | The initial position
  v ->
  BehaviorF m td v v
threePointDerivativeFrom v0 = proc v -> do
  dv <- derivativeFrom v0 -< v
  dv' <- delay zeroVector -< dv
  returnA -< (dv ^+^ dv') ^/ 2

{- | Like 'threePointDerivativeFrom',
   but with the initial position initialised to 'zeroVector'.
-}
threePointDerivative ::
  ( Monad m
  , VectorSpace v s
  , s ~ Diff td
  , Num s
  ) =>
  BehaviorF m td v v
threePointDerivative = threePointDerivativeFrom zeroVector

-- ** Averaging and filters

{- | A weighted moving average signal function.
   The output is the average of the first input,
   weighted by the second input
   (which is assumed to be always between 0 and 1).
   The weight is applied to the average of the last tick,
   so a weight of 1 simply repeats the past value unchanged,
   whereas a weight of 0 outputs the current value.
-}
weightedAverageFrom ::
  ( Monad m
  , VectorSpace v s
  , s ~ Diff td
  , Num s
  ) =>
  -- | The initial position
  v ->
  BehaviorF m td (v, s) v
weightedAverageFrom v0 = feedback v0 $ proc ((v, weight), vAvg) -> do
  let
    vAvg' = weight *^ vAvg ^+^ (1 - weight) *^ v
  returnA -< (vAvg', vAvg')

{- | An exponential moving average, or low pass.
   It will average out, or filter,
   all features below a given time constant @t@.
   (Equivalently, it filters out frequencies above @1 / (2 * pi * t)@.)
-}
averageFrom ::
  ( Monad m
  , VectorSpace v s
  , Floating s
  , s ~ Diff td
  ) =>
  -- | The initial position
  v ->
  -- | The time scale on which the signal is averaged
  Diff td ->
  BehaviorF m td v v
averageFrom v0 t = proc v -> do
  TimeInfo {..} <- timeInfo -< ()
  let
    weight = exp $ -(sinceLast / t)
  weightedAverageFrom v0 -< (v, weight)

-- | An average, or low pass, initialised to zero.
average ::
  ( Monad m
  , VectorSpace v s
  , Floating s
  , s ~ Diff td
  ) =>
  -- | The time scale on which the signal is averaged
  Diff td ->
  BehaviourF m td v v
average = averageFrom zeroVector

{- | A linearised version of 'averageFrom'.
   It is more efficient, but only accurate
   if the supplied time scale is much bigger
   than the average time difference between two ticks.
-}
averageLinFrom ::
  ( Monad m
  , VectorSpace v s
  , Floating s
  , s ~ Diff td
  ) =>
  -- | The initial position
  v ->
  -- | The time scale on which the signal is averaged
  Diff td ->
  BehaviourF m td v v
averageLinFrom v0 t = proc v -> do
  TimeInfo {..} <- timeInfo -< ()
  let
    weight = t / (sinceLast + t)
  weightedAverageFrom v0 -< (v, weight)

-- | Linearised version of 'average'.
averageLin ::
  ( Monad m
  , VectorSpace v s
  , Floating s
  , s ~ Diff td
  ) =>
  -- | The time scale on which the signal is averaged
  Diff td ->
  BehaviourF m td v v
averageLin = averageLinFrom zeroVector

-- *** First-order filters

-- | Alias for 'average'.
lowPass ::
  ( Monad m
  , VectorSpace v s
  , Floating s
  , s ~ Diff td
  ) =>
  Diff td ->
  BehaviourF m td v v
lowPass = average

-- | Filters out frequencies below @1 / (2 * pi * t)@.
highPass ::
  ( Monad m
  , VectorSpace v s
  , Floating s
  , Eq s
  , s ~ Diff td
  ) =>
  -- | The time constant @t@
  Diff td ->
  BehaviourF m td v v
highPass t = clId ^-^ lowPass t

-- | Filters out frequencies other than @1 / (2 * pi * t)@.
bandPass ::
  ( Monad m
  , VectorSpace v s
  , Floating s
  , Eq s
  , s ~ Diff td
  ) =>
  -- | The time constant @t@
  Diff td ->
  BehaviourF m td v v
bandPass t = lowPass t >>> highPass t

-- | Filters out the frequency @1 / (2 * pi * t)@.
bandStop ::
  ( Monad m
  , VectorSpace v s
  , Floating s
  , Eq s
  , s ~ Diff td
  ) =>
  -- | The time constant @t@
  Diff td ->
  BehaviourF m td v v
bandStop t = clId ^-^ bandPass t

-- * Delays

-- | Remembers and indefinitely outputs ("holds") the first input value.
keepFirst :: (Monad m) => ClSF m cl a a
keepFirst = safely $ do
  a <- try throwS
  safe $ arr $ const a

{- | Remembers all input values that arrived within a given time window.
   New values are appended left.
-}
historySince ::
  (Monad m, Ord (Diff (Time cl)), TimeDomain (Time cl)) =>
  -- | The size of the time window
  Diff (Time cl) ->
  ClSF m cl a (Seq (TimeInfo cl, a))
historySince dTime = readerS $ accumulateWith appendValue empty
  where
    appendValue (ti, a) tias = takeWhileL (recentlySince ti) $ (ti, a) <| tias
    recentlySince ti (ti', _) = diffTime (absolute ti) (absolute ti') < dTime

{- | Delay a signal by certain time span,
   initialising with the first input.
-}
delayBy ::
  (Monad m, Ord (Diff td), TimeDomain td) =>
  -- | The time span to delay the signal
  Diff td ->
  BehaviorF m td a a
delayBy dTime = historySince dTime >>> arr (viewr >>> safeHead) >>> lastS undefined >>> arr snd
  where
    safeHead EmptyR = Nothing
    safeHead (_ :> a) = Just a

-- * Timers

{- | Throws an exception after the specified time difference,
   outputting the time passed since the 'timer' was instantiated.
-}
timer ::
  ( Monad m
  , TimeDomain td
  , Ord (Diff td)
  ) =>
  Diff td ->
  BehaviorF (ExceptT () m) td a (Diff td)
timer diff = proc _ -> do
  time <- sinceStart -< ()
  _ <- throwOn () -< time > diff
  returnA -< time

-- | Like 'timer_', but doesn't output the remaining time at all.
timer_ ::
  ( Monad m
  , TimeDomain td
  , Ord (Diff td)
  ) =>
  Diff td ->
  BehaviorF (ExceptT () m) td a ()
timer_ diff = timer diff >>> arr (const ())

-- | Like 'timer', but divides the remaining time by the total time.
scaledTimer ::
  ( Monad m
  , TimeDomain td
  , Fractional (Diff td)
  , Ord (Diff td)
  ) =>
  Diff td ->
  BehaviorF (ExceptT () m) td a (Diff td)
scaledTimer diff = timer diff >>> arr (/ diff)