{-# LANGUAGE
    FlexibleContexts,
    FlexibleInstances,
    UndecidableInstances,
    ExistentialQuantification,
    TypeSynonymInstances,
    Rank2Types,
    NoMonomorphismRestriction #-}

import Prelude hiding ((.), id)
import Control.Category
import Control.Applicative
import Control.Monad
import Control.Arrow
import Data.Functor.Identity

{-

State space models (SSMs) are iterated functions (difference
equations) bundled with their initial conditions.

Note that they resemble the State monad as a base type, but they do
*not* behave as a State monad when composed.  The state is never
observable outside the SSM, while for a State monad it is threaded
from one function to the next.

Composing SSMs is like composing functions, placing them in series
linking the output of one to input of another.  The state
inputs/outputs of both SSMs are bundled in parallel.

Note that the type of an SSM update function is not s -> i -> (s, o)
because partial app doesn't seem so useful.  Using (s,i) -> (s,o) is
easier on the eyes.  The order is (s,i) instead of (i,s) to place
initial state next to input state parameter in the definition of SSM
instances.  This also reflects the usual matrix layout for linear
SSMs.

-}

ssmSer :: ((s1, b)  -> (s1, c)) ->
          ((s2, a)  -> (s2, b)) ->
          ((s1,s2), a) -> ((s1,s2), c)
              
ssmSer f g = fg where
  fg ((s1,s2), a) = ((s1',s2'), c) where
    (s2', b) = g (s2, a)
    (s1', c) = f (s1, b)
    
{-
Alternatively, it's possible to juxtapose two SSM update functions
creating a multi-input/output function.
-}
  
ssmPar :: ((s1, a)  -> (s1, c)) ->
          ((s2, b)  -> (s2, d)) ->
          ((s1,s2), (a,b)) -> ((s1,s2), (c,d))
ssmPar f g = fg where          
  fg ((s1,s2), (a,b)) = ((s1',s2'), (c,d)) where
    (s1',c) = f (s1, a)
    (s2',d) = g (s2, b)
    
{-

Lifting pure functions.  The point of interest is really Functor's
fmap, but it's simpler to define this lifting operation that can be
used in combination with (.) to create fmap.

-}

ssmPure :: (a -> b) -> ((), a) -> ((), b)
ssmPure f (s,a) = (s, f a)

{-

Define a data type to wrap SSMs that support the composition operation
above.  The goal is to make this SSM an instance of Category & Arrow.
However, it's not possible to expose the state type as a parameter in
SSM:

   data SSM s i o = SSM s ((s, i) -> (s, o))

Composing 2 of these will have the type

   SSM s1 b c -> SSM s2 a b -> SSM (s1,s2) a b

which will not fit the expected kind for Category: * -> * -> *.
To hide the "growing" state type we use an existential quantification.

Note that the only thing the state is used for after construction of
an SSM type is to chain it through the state update function.  As a
consequence it can remain completely hidden (i.e. doesn't require an
interface provided by a type class constraint).

-}
  
data SSM i o = forall s. SSM s ((s, i) -> (s, o))


-- SSMs can be composed by composing the update functions and the
-- initial state variables.  This gives rise to Category.  

-- ( Note that in this file the initial state passing is done in the
-- instance declarations; it is not part of the low-level functions. )
instance Category SSM where
  (.) (SSM sf f) (SSM sg g) = SSM (sf,sg) $ f `ssmSer` g
  id  = SSM () $ id
  
-- Adding also parallel composition and lifting of pure functions
-- gives an Arrow instance.  
  
-- ( Note that arr and first are enough, but we already have (***) so
-- use that to implement first. )
instance Arrow SSM where
  arr f = SSM () $ ssmPure f
  (***) (SSM i1 u1) (SSM i2 u2) = SSM (i1,i2) $ ssmPar u1 u2
  first   a = a *** (arr id)
  second  a = (arr id) *** a
  
  


---- Value-oriented interface.
--  
-- Above this line are the point-free interfaces for signal operators.
-- To obtain a value-oriented interface we have to implement Functor,
-- Applicative and Monad.  What I found is that the Arrow above is
-- (equivalent in power?) to the Kleisli arrow of a Monad representing
-- signals.  I wonder if the Monad can be dervied from its Kleisli
-- arrow implementation.
  
-- What about representing the Monad as a Kleisli arrow (Arrow () a)?
-- Then we just need to convert between (Arrow i o) and (i -> Arrow () o).
  
{- From what I can find online
http://lambda-the-ultimate.org/node/2799 it is possible to construct a
monad from an arrow if you can write the type isomorphism:

  iso :: (i -> Kl () o) -> Kl i o

The SSM type in this file has another problem that prevents the
construction of such an iso function so I tried it with a simpler
type:

  data Kl i o = forall s. Kl (i -> s -> (s, o))

I still found it impossible to write down the isomorphism.  Any
function I would think of would fail to typecheck.  The simplest one I
found is this:

  iso f = Kl $ \i -> (\(Kl kl) -> kl ()) (f i)

which can't work because there is no way of guaranteeing that the type
of kl, which is determined by the definition point of f, matches up
with the type we're hiding behind the constructor here.  They are
completely independent entities.  All information about that original
s type was available near the definition point of f, but is not
carried into this function.

The conclusion seems to be that this plainly doesn't work.  The iso
function above does work for the parameteric type (Kl s i o), but in
that case it is not possible to write a Category instance.

For an explanation, see:
http://zwizwa.be/-/compsci/20110820-221428


There is another way using ArrowApply.  However, this seems to run
into the same sort of problem: can't combine types from different
unpackings?


-}

{-
instance ArrowApply SSM where
  app (SSM i1 u1) =
    SSM $ (i1,i2)
-}

---- Tests
  

-- Define the discrete integral & differential with starting value 0.
int  = SSM 0 $ \(s, i) -> let o = i+s in (o,o)
diff = SSM 0 $ \(s, i) -> let o = i-s in (i,o)
z    = SSM 0 $ \(s, i) -> (i, s)

-- Unfold the difference equation given an input sequence.
runSSM :: (SSM i o) -> [i] -> [o]
runSSM (SSM init tick) = f init where
  f _ [] = []
  f s (i:is) = (o:os) where
    (s', o) = tick (s, i)
    os = f s' is

-- Print SSM as truncated impulse response.
dirac = (1:[0,0..])
impres ssm = runSSM ssm dirac
instance (Enum i, Num i, Show o) => Show (SSM i o) where
  show = show . (take 10) . impres