{-# LANGUAGE Rank2Types #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
{-# HLINT ignore "Use camelCase" #-}
module Language.Memo.Ztrategic where

import Language.ZipperAG
import Data.Generics.Zipper hiding (left, right, up, down')
import Data.Generics.Aliases
import Data.Maybe
import Data.Data
import Control.Monad -- (join, mplus, MonadPlus)

import Control.Monad.Trans.Maybe -- for MaybeT in elevate func

import Language.StrategicData
import Language.Memo.AGMemo

----------
----
--- TP
----
----------

newtype TP m = MkTP (forall d mm. (StrategicData (d mm), Memoizable d mm) => Zipper (d mm) -> m (Zipper (d mm)))
unTP (MkTP f) = f

applyTP_clean :: (Monad m, StrategicData (d mm), Memoizable d mm) => TP m -> Zipper (d mm) -> m (Zipper (d mm))
applyTP_clean f = fmap cleanMemoTable . unTP f

applyTP :: (StrategicData (d mm), Memoizable d mm) => TP m -> Zipper (d mm) -> m (Zipper (d mm))
applyTP = unTP

elevate :: (Monad m, Typeable b) => (b -> m b) -> Zipper a -> m (Maybe (Zipper a))
elevate g z = runMaybeT $ transM (MaybeT . maybe (return Nothing) (fmap cast . g) . cast) z


elevateZ :: (Monad m, Typeable a, Typeable b, Typeable c) => (b -> Zipper c -> m (Zipper c)) -> Zipper a -> m (Maybe (Zipper a))
elevateZ f z = maybe (return Nothing) (\b -> maybe (return Nothing) (fmap cast . f b) (cast z)) (getHole z)

full_tdTP :: (Monad m) => TP m  -> TP m
full_tdTP f = f `seqTP` allTPdown (full_tdTP f) `seqTP` allTPright (full_tdTP f)

full_buTP :: (Monad m) => TP m -> TP m
full_buTP f = allTPright (full_buTP f) `seqTP` allTPdown (full_buTP f) `seqTP` f

once_tdTP :: (MonadPlus m) => TP m -> TP m
once_tdTP f = f `choiceTP` oneTPdown (once_tdTP f) `choiceTP` oneTPright (once_tdTP f)

once_buTP :: (MonadPlus m) => TP m -> TP m
once_buTP f = oneTPright (once_buTP f) `choiceTP` oneTPdown (once_buTP f) `choiceTP` f

--Experimental
stop_tdTP :: (MonadPlus m) => TP m -> TP m
stop_tdTP f = f `choiceTP` (allTPdown (stop_tdTP f) `seqTP` allTPright (stop_tdTP f))

--Experimental
stop_buTP :: (MonadPlus m) => TP m -> TP m
stop_buTP f = (allTPdown (stop_tdTP f) `seqTP` allTPright (stop_tdTP f)) `choiceTP` f

adhocTP :: (Monad m, Typeable b) => TP m -> (b -> m b) -> TP m
adhocTP f g = MkTP $ \z -> do
          val <- elevate g z
          maybe (applyTP f z) return val

adhocTPSeq :: (MonadPlus m, Typeable b) => TP m -> (b -> m b) -> TP m
adhocTPSeq f g = MkTP (\z -> do
             val <- elevate g z
             maybe (applyTP f z) (applyTP (tryTP f)) val)
                `choiceTP` f

adhocTPZ :: (Monad m, Typeable a, Typeable b) => TP m -> (b -> Zipper a -> m (Zipper a)) -> TP m
adhocTPZ f g = MkTP $ \z -> do
          val <- elevateZ g z
          maybe (applyTP f z) return val

adhocTPZSeq :: (MonadPlus m, Typeable a, Typeable b) => TP m -> (b -> Zipper a -> m (Zipper a)) -> TP m
adhocTPZSeq f g = MkTP (\z -> do
             val <- elevateZ g z
             maybe (applyTP f z) (applyTP (tryTP f)) val)
                `choiceTP` f

--Identity function
idTP :: Monad m => TP m
idTP = MkTP return


--Failing function
failTP :: MonadPlus m => TP m
failTP = MkTP (const mzero)

allTPright :: (Monad m) => TP m  -> TP m
allTPright f = MkTP $ \z -> moveM right left (return z) (applyTP f) z

allTPdown :: (Monad m) => TP m  -> TP m
allTPdown f = MkTP $ \z -> moveM down' up (return z) (applyTP f) z

oneTPright :: (MonadPlus m) => TP m -> TP m
oneTPright f = MkTP $ moveM right left mzero (applyTP f)

oneTPdown :: (MonadPlus m) => TP m -> TP m
oneTPdown f = MkTP $ moveM down' up mzero (applyTP f)

--Sequential composition, ignores failure
seqTP :: Monad m => TP m -> TP m -> TP m
seqTP f g = MkTP (unTP f `mseq` unTP g)
f `mseq` g =  f >=> g

--Sequential composition, chooses rightmost only if possible
choiceTP :: MonadPlus m => TP m -> TP m -> TP m
choiceTP f g = MkTP (unTP f `mchoice` unTP g)
f `mchoice` g   =  \x -> f x `mplus` g x


--Apply a function, fail the composition if it fails
monoTP :: (MonadPlus m, Typeable b) => (b -> m b) -> TP m
monoTP = adhocTP failTP

--Apply a function with access to the zipper, fail the composition if it fails
monoTPZ :: (MonadPlus m, Typeable a, Typeable b) => (b -> Zipper a -> m (Zipper a)) -> TP m
monoTPZ = adhocTPZ failTP

--Try to apply a zipper function, and apply identity if it fails
tryTP :: MonadPlus m => TP m -> TP m
tryTP s = s `choiceTP` idTP


repeatTP :: MonadPlus m => TP m -> TP m
repeatTP s =  tryTP (s `seqTP` repeatTP s)

innermost :: (MonadPlus m) => TP m -> TP m
innermost s = allTPright (innermost s)
      `seqTP` allTPdown (innermost s)
      `seqTP` tryTP (s `seqTP` innermost s)

innermost' :: (MonadPlus m) => TP m -> TP m
innermost' s  = repeatTP (once_buTP s)

outermost  :: (MonadPlus m) => TP m -> TP m
outermost s = repeatTP (once_tdTP s)

----------
----
--- TU
----
----------


newtype TU m r = MkTU (forall d mm . (StrategicData (d mm), Memoizable d mm) => Zipper (d mm) -> (m r, Zipper (d mm)))
unTU (MkTU f) = f

applyTU_clean :: (Memoizable d mm, StrategicData (d mm)) => TU m r -> Zipper (d mm) -> (m r, Zipper (d mm))
applyTU_clean f z = case unTU f z of
                (r, z1) -> (r, cleanMemoTable z1)

applyTU :: (Memoizable d mm, StrategicData (d mm)) => TU m r -> Zipper (d mm) -> (m r, Zipper (d mm))
applyTU = unTU

foldr1TU :: (Memoizable d mm, StrategicData (d mm), Monoid (m r), Foldable m, MonadPlus m) =>
    TU m r ->
    Zipper (d mm) -> (r -> r -> r) -> r
foldr1TU f z red = foldr1 red $ fst (applyTU (full_tdTU f) z)

foldl1TU :: (Memoizable d mm, StrategicData (d mm), Monoid (m r), Foldable m, MonadPlus m) =>
    TU m r ->
    Zipper (d mm) -> (r -> r -> r) -> r
foldl1TU f z red = foldl1 red $ fst (applyTU (full_tdTU f) z)

foldrTU :: (Memoizable d mm, StrategicData (d mm), Monoid (m r), Foldable m, MonadPlus m) =>
    TU m r ->
    Zipper (d mm) -> (r -> s -> s) -> s -> s
foldrTU f z red i = foldr red i $ fst (applyTU (full_tdTU f) z)

foldlTU :: (Memoizable d mm, StrategicData (d mm), Monoid (m r), Foldable m, MonadPlus m) =>
    TU m r ->
    Zipper (d mm) -> (s -> r -> s) -> s -> s
foldlTU f z red i = foldl red i $ fst (applyTU (full_tdTU f) z)

full_tdTU :: (MonadPlus m, Monoid (m r)) => TU m r -> TU m r
full_tdTU f = allTUright (full_buTU f) `seqTU` allTUdown (full_buTU f) `seqTU` f

full_buTU :: (MonadPlus m, Monoid (m r)) => TU m r -> TU m r
full_buTU f = f `seqTU` allTUdown (full_tdTU f) `seqTU` allTUright (full_tdTU f)

once_tdTU :: (MonadPlus m, Monoid (m r)) => TU m r -> TU m r
once_tdTU f = f `choiceTU` allTUdown (once_tdTU f)  `choiceTU` allTUright (once_tdTU f)

once_buTU :: (MonadPlus m, Monoid (m r)) => TU m r -> TU m r
once_buTU f = allTUright (once_buTU f) `choiceTU` allTUdown (once_buTU f) `choiceTU` f

stop_tdTU :: (MonadPlus m, Monoid (m r)) => TU m r -> TU m r
stop_tdTU f = (f `choiceTU` allTUdown (stop_tdTU f)) `seqTU` allTUright (stop_tdTU f)

stop_buTU :: (MonadPlus m, Monoid (m r)) => TU m r -> TU m r
stop_buTU f = allTUright (stop_buTU f) `seqTU` (allTUdown (stop_buTU f) `choiceTU` f)


allTUdown :: (Monoid (m r)) => TU m r -> TU m r
allTUdown f = MkTU $ \z -> moveQ down' (mempty, z) (back . applyTU f) z
 where back (a,b) = (a, fromJust $ up b)


allTUright :: (Monoid (m r)) => TU m r -> TU m r
allTUright f = MkTU $ \z -> moveQ right (mempty, z) (back . applyTU f) z
 where back (a,b) = (a, fromJust $ left b)


adhocTU :: (Monad m, Typeable a) => TU m r -> (a -> m r) -> TU m r
adhocTU f g = MkTU $ \z -> maybe (applyTU f z) (\c -> (c, z)) (getHole z >>= return . g)

adhocTUZ :: (Monad m, Typeable a, Typeable d, Typeable mm) => TU m r -> (a -> Zipper (d mm) -> (m r, Zipper (d mm))) -> TU m r
adhocTUZ f g = MkTU $ \z -> fromMaybe (applyTU f z) (reduce g z)

reduce :: (Typeable a, Typeable d, Typeable d', Typeable mm', Typeable mm) => (a -> Zipper (d' mm') -> (m r, Zipper (d' mm'))) -> Zipper (d mm) -> Maybe (m r, Zipper (d mm))
reduce f z = getHole z >>= \v -> cast z >>= (\(r, m) -> cast m >>= \mc -> return (r, mc)) . f v

seqTU :: (Monoid (m r)) => TU m r -> TU m r -> TU m r
seqTU x y = MkTU $ \z ->
  let (yr, z') = applyTU y z
      (xr, z'') = applyTU x z'
  in (xr `mappend` yr, z'')

choiceTU :: (MonadPlus m) => TU m r -> TU m r -> TU m r
choiceTU x y = MkTU $ \z ->
  let (yr, z') = applyTU y z
      (xr, z'') = applyTU x z'
  in (xr `mplus` yr, z'')


failTU :: (MonadPlus m) => TU m r
failTU = MkTU $ \z -> (mzero, z)

constTU :: Monad m => r -> TU m r
constTU v = MkTU $ \z -> (return v, z)

--Apply a function, fail the composition if it fails
monoTU :: (MonadPlus m, Typeable a) => (a -> m r) -> TU m r
monoTU = adhocTU failTU

--Apply a function with access to the zipper, fail the composition if it fails
monoTUZ :: (MonadPlus m, Typeable a, Typeable d, Typeable mm) => (a -> Zipper (d mm) -> (m r, Zipper (d mm))) -> TU m r
monoTUZ = adhocTUZ failTU