{-# LANGUAGE
    NoMonomorphismRestriction,
    Rank2Types,
    ExistentialQuantification #-}

import Prelude hiding ((.))
import Control.Applicative
import Control.Category
import Control.Monad
import Control.Arrow

-- A representation for sequences (signals) that can be represented as
-- iterated functions + initial value.

-- ( Note: type aliases are used for the "proto types" to expose the
-- structure of the computation without unnecessary
-- wrapping/unwrapping.  The idea is to wrap these in existentially
-- quantified data definitions so they can be type class instances.
-- State is the 2nd element in the tuple to make some functions look
-- simpler. )

type Sig s a = (s, (s -> (s, a)))

-- Monad, starting from fmap and join.
-- 
-- Join flattens two monadic layers into one.  This boils down to
-- rewriting the state update equations into a single one.

_join :: (Sig s1 (Sig s2 a)) -> Sig (s1,s2) a
_join (i1, u1) = ((i1, i2), u12) where
  -- i1, i2   : inital state value
  -- u1, u2   : state update function
  -- s1, s2   : input state
  -- s1', s2' : output state
  (_, (i2, _)) = u1 i1 -- (**) 
  u12 (s1, s2) = ((s1', s2'), a) where
    (s1', (_, u2)) = u1 s1
    (s2', a)       = u2 s2

-- (**) Dig up the inital value which is buried deeply.  I don't know
-- how to get at this value using some clever recursive / lazy binding
-- trick so I just use the knowledge that it is constant and
-- independent of the input.  This allows the initial value to be used
-- as input to unlock the goodies that way.


-- The map is a bit special in that it doesn't need to change the type
-- of the state.  This means that the Functor instance doesn't need an
-- existential wrapper.

_fmap :: (a -> b) -> Sig s a -> Sig s b
_fmap f (s0, u) = (s0, u') where
  u' s = (s', f a) where
    (s', a) = u s
    

-- With return we're done.
_return :: a -> Sig () a
_return x = ((), \() -> ((), x))


-- The rest is wrapping and composing.

_bind :: (Sig s1 a) -> (a -> Sig s2 b) -> (Sig (s1,s2) b)
_bind ma f = _join ((_fmap f) ma)

-- The state type (s1,(s2,())) is an artefact of defing _ap using
-- _return in the definition of _ap.  It could just have well been
-- defined manually as (s1,s2).
_ap :: Sig s1 (a -> b) -> Sig s2 a -> Sig (s1, (s2, ())) b
_ap f a = f `_bind` \f' ->
          a `_bind` \a' ->
          _return $ f' a' 

_run (init, tick) = f init where
  f s = (v:vs) where
    (s', v) = tick s
    vs = f s'


-- Define an existential wrapper to hide the state types.
data Signal a = forall s. Signal (Sig s a)

run (Signal a) = _run a

-- Functor and Applicative instances are straightforward.
instance Functor Signal where
  fmap f (Signal a) = Signal $ _fmap f a
instance Applicative Signal where
  pure  = Signal . _return
  (<*>) (Signal f) (Signal a) = Signal $ _ap f a


-- However, I can't get Monad to work.  The "unpacking" that's done in
-- _bind or _join seem to somehow be structured in such a way that its
-- impossible (or just not straightforward?) to type it correctly.
  
{-
instance Monad Signal where
  return = Signal . _return
  (>>=) (Signal ma) f = Signal $ _bind ma f' where 
    f' a = case (f a) of
      Signal mb -> mb
-}

-- So let's skip that for now and define some more abstractions.
  
-- Operators are defined as Kleisli arrows in the pseudo-monad (I'm
-- going to just call it monad from now on.)


-- Signal operators are Kleisli arrows of the monad.
type Opr s i o = i -> Sig s o
_kleisli :: Opr s1 a b -> Opr s2 b c -> Opr (s1,s2) a c
_kleisli f g = \x -> (f x) `_bind` g

-- This can be abstracted in another existential type.
data Operator i o = forall s. Operator (Opr s i o)

instance Category Operator where
  id = Operator _return
  (.) (Operator f) (Operator g) = Operator $ g `_kleisli` f


-- It doesn't look like defining the Kleisli instance works without
-- having a Monad instance first.  I wonder if it's possible to
-- construct an Arrow instance directly without running into typing
-- restrictions as for Monad.
  
_arr :: (a -> b) -> Opr () a b  
_arr f a =  _return $ f a

{-

-- Manual definition for _first.  Note this has a simpler type than
-- the composed version which has a spurious () due to the use of
-- _return, but since we're hiding types anyway that doesn't make such
-- a difference.

_first :: Opr s a b -> Opr s (a, z) (b, z)
_first opr = opr' where
  opr' (a,z) = (s0, u') where
    (s0, u) = opr a   -- Use input a to get at Sig
    u' s = (s', (b, z)) where
      (s', b) = u s   -- Use state to get at output and next state.
-}

_first :: Opr s a b -> Opr (s, ()) (a, z) (b, z)
_first arr (a, z) =
  arr a `_bind` \b ->
  _return (b, z)


instance Arrow Operator where
  arr   = Operator . _arr
  first (Operator a) = Operator $ _first a


-- Convert a convenient state space model sigature to a Kleisli arrow.
_operator :: s -> ((s,i) -> (s,o)) -> Opr s i o
_operator s0 f = f' where
  f' i = (s0, u) where
    u s = (s', o) where
      (s',o) = f (s,i)
          
-- Similar, but wrapped as an Operator instance.
operator i u  = Operator $ _operator i u
              
-- Example: single integrator written in terms of update function, and
-- a double integrator constructed as a composition.
intUpdate (s,i) = (s',o) where 
  s' = i + s
  o  = s'

-- Instantiate it with naked types... 
_int  = _operator 0 intUpdate
_int2 = _int `_kleisli` _int    
    
-- and wrapped with class functionality attached.
int  = Operator _int
int2 = int . int


-- Applying operators to signals is just _bind.
_ones = _return 1

_disp = (take 10) . _run

_t0 = _disp _ones
_t1 = _disp (_ones `_bind` _int)
_t2 = _disp (_ones `_bind` _int `_bind` _int)
_t3 = _disp (_ones `_bind` (_int `_kleisli` _int))  -- same as _t2

-- But we don't have a true bind for Signal so let's work around it by
-- defining application between Operator and Signal.

(.$) (Operator o) (Signal s) = Signal $ s `_bind` o

disp = (take 10) . run
ones = Signal $ _return 1
t0 = disp ones
t1 = disp $ int .$ ones
t2 = disp $ int .$ (int .$ ones)
t3 = disp $ (int . int) .$ ones

compM :: (Functor m, Monad m) => 
         (a -> m b) -> (b -> m c) -> (a -> m c)
compM f g = f .> (fmap g) .> join
  where (.>) = flip (.)