{-# OPTIONS -Wall #-}
{-# LANGUAGE GADTs, KindSignatures, TypeOperators, GeneralizedNewtypeDeriving
  , ScopedTypeVariables, Arrows #-}
----------------------------------------------------------------------
-- |
-- Module      :  Data.Bot.Edit
-- Copyright   :  (c) Conal Elliott 2008
-- License     :  BSD3
-- 
-- Maintainer  :  conal@conal.net
-- Stability   :  experimental
-- 
-- Edit-bots
----------------------------------------------------------------------

module Data.Bot.Edit where

import Prelude hiding (id,(.))
import Control.Category
import Control.Arrow

import Data.Maybe (fromMaybe)

-- import Control.Arrow.Transformer
import Control.Arrow.Transformer.Automaton

import Data.Stream hiding ()
import qualified Data.Stream as S

-- | An edit/change to a value.  Represents @a -> a@ but as data, to allow
-- analysis.
data Edit :: * -> * where
  Konst   :: a -> Edit a                     -- const
  Idy     :: Edit a                          -- id
  (:***)  :: Edit a -> Edit b -> Edit (a,b)  -- ***
  (:>>>)  :: Edit a -> Edit a -> Edit a      -- (>>>)

-- Note: why use *** instead of First?  Because the smart constructor
-- rewrite rules seem to come out more simply.

-- | Interpret an edit as an arrow value (e.g., function).
edit :: Arrow (~>) => Edit a -> (a ~> a)
edit (Konst a)      = arr (const a)
edit Idy            = arr id
edit (eda :*** edb) = edit eda *** edit edb
edit (ed :>>> ed')  = edit ed  >>> edit ed'

-- TODO: import Data.Pair.  Replace unzip* with unpair and zip* with
-- pair.  Would add a TypeCompose dependency.

-- | Zip up two edits into a pair edit
zipE :: Edit a -> Edit b -> Edit (a,b)
zipE = (***:)

-- Try splitting one edit @ed@ into two independent edits @eda@ and @edb@
-- such that @ed == eda *** edb@.
unzipE :: Edit (a,b) -> Maybe (Edit a, Edit b)

unzipE (Konst (a,b))  = Just (Konst a, Konst b)
unzipE Idy            = Just (Idy, Idy)
unzipE (eda :*** edb) = Just (eda, edb)
unzipE (_   :>>> _  ) = Nothing

-- | Smart constructor for '(:***)'.  Applies some rewrites.
(***:) :: Edit a -> Edit b -> Edit (a,b)

Idy     ***: Idy     = Idy
Konst a ***: Konst b = Konst (a,b)
ed      ***: ed'     = ed :*** ed'

-- | Smart constructor for '(:>>>)'.  Applies some rewrites.
(>>>:) :: Edit a -> Edit a -> Edit a

Idy        >>>: f            = f
f          >>>: Idy          = f
_          >>>: k@(Konst _)  = k
(f :*** g) >>>: (f' :*** g') = (f >>>: f') ***: (g >>>: g')
ed         >>>: ed'          = ed :>>> ed'


-- | A tweak is an edit and a new value.  This new value can be computed
-- from the edit and a previous value.  It's stored here (lazily) so it
-- can be shared.
data Tweak a = Tweak (Edit a) a

-- | Simple 'Konst' tweaker.  A catch-all, but not helpful for performance.
tweak :: a -> Tweak a
tweak a = Tweak (Konst a) a

-- Zip up two tweaks into a pair tweak.
zipT :: Tweak a -> Tweak b -> Tweak (a,b)
zipT (Tweak ea a) (Tweak eb b) = Tweak (zipE ea eb) (a,b)

-- | Split a pair tweak in two.
unzipT :: Tweak (a,b) -> (Tweak a, Tweak b)
unzipT (Tweak ed (a,b)) = (Tweak eda a, Tweak edb b)
 where
   (eda,edb) = fromMaybe (Konst a, Konst b) (unzipE ed)

-- | Stream of tweaks
type Tweaks a = Stream (Tweak a)

-- Zip up two tweak streams into a pair-valued tweak stream.
zipTs :: Tweaks a -> Tweaks c -> Tweaks (a,c)
zipTs = (fmap.fmap.fmap) (uncurry zipT) S.zip

-- zipTs tas tcs = fmap (uncurry zipT) $ S.zip tas tcs

-- | Split a pair-valued tweak stream in two.
unzipTs :: Tweaks (a,c) -> (Tweaks a, Tweaks c)
unzipTs = S.unzip . fmap unzipT


-- | Reactive (discretely varying) value.  An initial value and a stream
-- of tweaks.  
-- 
-- Although 'Reactive' is a 'Functor', I don't recommend its use.  'fmap'
-- cannot do much optimization without knowing something about the
-- function being mapped.  Maybe better not to have the 'Functor'
-- instance.
data Reactive a = R a (Tweaks a)

-- Maybe punt Reactive and rely on the first stream element to have a Konst

instance Functor Reactive where
  fmap f (R a ts) = R b (fmapR' f b ts) where b = f a

-- Carry value along to re-use for identity edits.
fmapR' :: (a -> b) -> b -> Stream (Tweak a) -> Stream (Tweak b)
fmapR' f b (Cons (Tweak Idy _) ts) = Cons (Tweak Idy b) (fmapR' f b ts)
fmapR' f _ (Cons (Tweak _ a) ts) =
  Cons (tweak b') (fmapR' f b' ts) where b' = f a

-- | Zip up two reactive values into a reactive pair.
zipR :: Reactive a -> Reactive c -> Reactive (a,c)
zipR (R a tas) (R c tcs) = R (a,c) (zipTs tas tcs)

-- | Split a reactive pair into a pair of reactive values.
unzipR :: Reactive (a,c) -> (Reactive a, Reactive c)
unzipR (R (a,c) ts) = (R a tsa, R c tsc)
 where
   (tsa,tsc) = unzipTs ts


-- | Arrow between two reactive values.
newtype ArrowReactive (~>) a b = AR (Reactive a ~> Reactive b)

instance Category (~>) => Category (ArrowReactive (~>)) where
  id          = AR id
  AR g . AR f = AR (g . f)

instance Arrow (~>) => Arrow (ArrowReactive (~>)) where
  arr f         = AR (arr (fmap f))
  first (AR  f) = AR (arr unzipR >>> first f >>> arr (uncurry zipR))


-- | 'Automaton'-style tweak arrow.  Could be split into follow and lead.
data AutR (~>) a b = AutR (a ~> (b, Automaton (~>) (Tweak a) (Tweak b)))

instance Arrow (~>) => Category (AutR (~>)) where
  id = arr id
  AutR g . AutR f = AutR $ proc a -> do (b,ab) <- f -< a
                                        (c,bc) <- g -< b
                                        returnA -< (c, bc . ab)

instance Arrow (~>) => Arrow (AutR (~>)) where
  arr f = AutR (arr (\ a -> let b = f a in (b, arrAut f (f a))))
  first (AutR f) = AutR $
    proc (a,c) ->
      do (b,ab) <- f -< a
         returnA -< ((b,c), arr unzipT >>> first ab >>> arr (uncurry zipT))

arrAut :: Arrow (~>) => (a -> b) -> b -> Automaton (~>) (Tweak a) (Tweak b)
arrAut f b = aut
 where
   aut = Automaton (arr g)
   g (Tweak Idy _) = (Tweak Idy b, aut)
   g (Tweak _   a) = (tweak b', arrAut f b') where b' = f a