{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Trustworthy #-}
module Data.Fold.L1
  ( L1(..)
  ) where

import Control.Applicative
import Control.Arrow
import Control.Category
import Control.Lens
import Data.Fold.Class
import Data.Fold.Internal
import Data.Functor.Apply
import Data.Pointed
import Data.Profunctor
import Data.Profunctor.Unsafe
import Data.Semigroupoid
import Prelude hiding (id,(.))
import Unsafe.Coerce

-- | A Mealy Machine
data L1 a b = forall c. L1 (c -> b) (c -> a -> c) (a -> c)

instance Scan L1 where
  run1 a (L1 k _ z) = k (z a)
  prefix1 a (L1 k h z) = L1 k h (h (z a))
  postfix1 (L1 k h z) a = L1 (\c -> k (h c a)) h z
  interspersing a (L1 k h z) = L1 k (\x b -> h (h x a) b) z
  {-# INLINE run1 #-}
  {-# INLINE prefix1 #-}
  {-# INLINE postfix1 #-}
  {-# INLINE interspersing #-}

instance Functor (L1 a) where
  fmap f (L1 k h z) = L1 (f.k) h z
  {-# INLINE fmap #-}
  b <$ _ = pure b
  {-# INLINE (<$) #-}

instance Pointed (L1 a) where
  point x = L1 (\() -> x) (\() _ -> ()) (\_ -> ())
  {-# INLINE point #-}

instance Apply (L1 a) where
  (<.>) = (<*>)
  {-# INLINE (<.>) #-}
  (<.) m = \_ -> m
  {-# INLINE (<.) #-}
  _ .> m = m
  {-# INLINE (.>) #-}

instance Applicative (L1 a) where
  pure x = L1 (\() -> x) (\() _ -> ()) (\_ -> ())
  {-# INLINE pure #-}
  L1 kf hf zf <*> L1 ka ha za = L1
    (\(Pair' x y) -> kf x (ka y))
    (\(Pair' x y) a -> Pair' (hf x a) (ha y a))
    (\a -> Pair' (zf a) (za a))
  (<*) m = \ _ -> m
  {-# INLINE (<*) #-}
  _ *> m = m
  {-# INLINE (*>) #-}

instance Monad (L1 a) where
  return x = L1 (\() -> x) (\() _ -> ()) (\_ -> ())
  {-# INLINE return #-}
  m >>= f = L1 (\xs a -> walk xs (f a)) Snoc1 First <*> m where
  {-# INLINE (>>=) #-}
  _ >> n = n
  {-# INLINE (>>) #-}

instance Semigroupoid L1 where
  o = (.)
  {-# INLINE o #-}

instance Category L1 where
  id = arr id
  {-# INLINE id #-}
  L1 k h z . L1 k' h' z' = L1 (\(Pair' b _) -> k b) h'' z'' where
    z'' a = Pair' (z (k' b)) b where b = z' a
    h'' (Pair' c d) a = Pair' (h c (k' d')) d' where d' = h' d a
  {-# INLINE (.) #-}

instance Arrow L1 where
  arr h = L1 h (\_ a -> a) id
  {-# INLINE arr #-}
  first (L1 k h z) = L1 (first k) h' (first z) where
    h' (a,b) (c,_) = (h a c, b)
  {-# INLINE first #-}
  second (L1 k h z) = L1 (second k) h' (second z) where
    h' (a,b) (_,c) = (a, h b c)
  {-# INLINE second #-}
  L1 k h z *** L1 k' h' z' = L1 (k *** k') h'' (z *** z') where
    h'' (a,b) (c,d) = (h a c, h' b d)
  {-# INLINE (***) #-}
  L1 k h z &&& L1 k' h' z' = L1 (k *** k') h'' (z &&& z') where
    h'' (c,d) a = (h c a, h' d a)
  {-# INLINE (&&&) #-}

instance Profunctor L1 where
  dimap f g (L1 k h z) = L1 (g.k) (\a -> h a . f) (z.f)
  {-# INLINE dimap #-}
  lmap f (L1 k h z) = L1 (k) (\a -> h a . f) (z.f)
  {-# INLINE lmap #-}
  rmap g (L1 k h z) = L1 (g.k) h z
  {-# INLINE rmap #-}
  ( #. ) _ = unsafeCoerce
  {-# INLINE (#.) #-}
  x .# _ = unsafeCoerce x
  {-# INLINE (.#) #-}

instance Strong L1 where
  first' = first
  {-# INLINE first' #-}
  second' = second
  {-# INLINE second' #-}

instance Choice L1 where
  left' (L1 k h z) = L1 (_Left %~ k) step (_Left %~ z) where
    step (Left x) (Left y) = Left (h x y)
    step (Right c) _ = Right c
    step _ (Right c) = Right c
  {-# INLINE left' #-}

  right' (L1 k h z) = L1 (_Right %~ k) step (_Right %~ z) where
    step (Right x) (Right y) = Right (h x y)
    step (Left c) _ = Left c
    step _ (Left c) = Left c
  {-# INLINE right' #-}

instance ArrowChoice L1 where
  left (L1 k h z) = L1 (_Left %~ k) step (_Left %~ z) where
    step (Left x) (Left y) = Left (h x y)
    step (Right c) _ = Right c
    step _ (Right c) = Right c
  {-# INLINE left #-}

  right (L1 k h z) = L1 (_Right %~ k) step (_Right %~ z) where
    step (Right x) (Right y) = Right (h x y)
    step (Left c) _ = Left c
    step _ (Left c) = Left c
  {-# INLINE right #-}

data SnocList1 a = Snoc1 (SnocList1 a) a | First a

walk :: SnocList1 a -> L1 a b -> b
walk xs0 (L1 k h z) = k (go xs0) where
  go (First a) = z a
  go (Snoc1 as a) = h (go as) a
{-# INLINE walk #-}