{-# LANGUAGE LambdaCase #-}

module AFRP
  ( Mealy(..)
  , eff
  , Event(..)
  , hold
  , events
  , switch
  , preMapAccum
  , preMapAccumUTCTime
  , mapAccum
  , mapAccumUTCTime
  , changes
  , whenA
  , filterA
  , thenA
  , (>>|)
  , toEvent
  ) where

import Control.Category (Category(..), (>>>))
import Prelude hiding ((.), id)
import Control.Arrow (Arrow(..), ArrowChoice(..), ArrowLoop(..))
import Data.Time (UTCTime)
import Control.Monad.Fix (MonadFix (mfix))
import Data.Either (fromLeft)
import Data.Bool (bool)

newtype Mealy eff a b = Mealy
  { runMealy :: forall m. MonadFix m => (forall x. eff x -> m x) -> UTCTime -> a -> m (b, Mealy eff a b) }

eff :: (a -> eff b) -> Mealy eff a b
eff f = Mealy $ \nt _ x ->
  nt (f x) >>= \b -> pure (b, eff f)

instance Category (Mealy eff) where
  id = Mealy (\_ _ x -> pure (x, id))
  (Mealy f) . (Mealy g) = Mealy $ \nt t a -> do
    (b, g') <- g nt t a
    (c, f') <- f nt t b
    pure (c, f' . g')

instance Arrow (Mealy eff) where
  arr f = Mealy $ \_ _ b -> pure (f b, arr f)
  first (Mealy f) = Mealy $ \nt t (b,d) -> do
    (c, f') <- f nt t b
    pure ((c, d), first f')

instance ArrowChoice (Mealy eff) where
  left (Mealy f) = Mealy $ \nt t -> \case
    Left b -> do
      (c, f') <- f nt t b
      pure (Left c, left f')
    Right d -> pure (Right d, left (Mealy f))

instance ArrowLoop (Mealy eff) where
  loop (Mealy f) = Mealy $ \nt t b -> do
    ((c,_), f') <- mfix $ \((_,d), _) -> f nt t (b,d)
    pure (c, loop f')

instance Functor (Mealy eff a) where
  fmap f (Mealy g) = Mealy $ \nt t a -> do
    (b, g') <- g nt t a
    pure (f b, fmap f g')

instance Applicative (Mealy eff a) where
  pure b = Mealy $ \_ _ _ -> pure (b, pure b)
  Mealy f <*> Mealy x = Mealy $ \nt t a -> do
    (f', fNext) <- f nt t a
    (x', xNext) <- x nt t a
    pure (f' x', fNext <*> xNext)

data Event a
  = Tick
  | Event a
  deriving (Show, Functor, Foldable, Traversable)

hold :: a -> Mealy eff (Event a) a
hold a = Mealy $ \_ _ -> \case
  Tick -> pure (a, hold a)
  Event a' -> pure (a', hold a')

events :: Mealy eff (Event a) (Either () a)
events = arr $ \case
  Tick -> Left ()
  Event a -> Right a

switch :: Mealy eff a (b, Event c) -> (c -> Mealy eff a b) -> Mealy eff a b
switch (Mealy f) s = Mealy $ \nt t a -> do
  ((b, ev), f') <- f nt t a
  case ev of
          Tick -> pure (b, switch f' s)
          Event x -> runMealy (s x) nt t a

preMapAccum :: (x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
preMapAccum f x extract = go x
  where
    go b = Mealy $ \_ _ a ->
      let next = f b a
      in pure (extract b, go next)

preMapAccumUTCTime :: (UTCTime -> x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
preMapAccumUTCTime f x extract = go x
  where
    go b = Mealy $ \_ t a ->
      let next = f t b a
      in pure (extract b, go next)

mapAccum :: (x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
mapAccum f x extract = go x
  where
    go b = Mealy $ \_ _ a ->
      let next = f b a
      in pure (extract next, go next)

mapAccumUTCTime :: (UTCTime -> x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
mapAccumUTCTime f x extract = go x
  where
    go b = Mealy $ \_ t a ->
      let next = f t b a
      in pure (extract next, go next)

changes :: Eq a => Mealy eff a (Event a)
changes = mapAccum go Nothing (maybe Tick snd)
  where
    go :: Eq a => Maybe (a, Event a) -> a -> Maybe (a, Event a)
    -- The first observed value is not a change I think
    go Nothing x = Just (x, Tick)
    go (Just (y, _)) x | x == y = Just (x, Tick)
                       | otherwise = Just (x, Event x)

whenA :: (a -> Bool) -> Mealy eff a () -> Mealy eff a ()
whenA predicate auto = arr (\a -> if predicate a then Left a else Right ()) >>> left auto >>> arr (fromLeft ())

filterA :: (a -> Bool) -> Mealy eff a (Either () a)
filterA f = arr $ \a -> bool (Left ()) (Right a) (f a)

thenA :: (ArrowChoice cat, Arrow cat) => cat a (Either b1 c) -> cat c (Either b1 b2) -> cat a (Either b1 b2)
thenA f g = f >>> arr Left ||| g

(>>|) :: (ArrowChoice cat, Arrow cat) => cat a (Either b1 c) -> cat c (Either b1 b2) -> cat a (Either b1 b2)
(>>|) = thenA

infixl 1 >>|

toEvent :: Mealy eff (Either () a) (Event a)
toEvent = arr (either (const Tick) Event)