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

import Language.ZipperAG
import Language.StrategicData

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)

-- monad transformers, to extract the monad from adhocTP
import Control.Monad.Trans.Class
import Control.Monad.Trans.Maybe

import Control.Monad.State.Lazy
import System.Random (randomRIO, initStdGen, setStdGen)


gatherChildren (trav, untrav) z = maybe [] ((navigation:) . gatherChildrenR navigation) $ down' z
 where navigation = (fromJust . down' . trav, untrav . fromJust . up)


gatherChildrenR (toNthChild, fromNthChild) z = maybe [] ((navigation:) . gatherChildrenR navigation) $ right z
 where navigation = (fromJust . right . toNthChild, fromNthChild)



breadthFirst_tdTP f = MkTP $ \z -> bf' f z [(id, id)]
 where bf' f z [] = return z
       bf' f z ((trav, untrav):travs) =
           do
             let thisNode = trav z
             thisNodeTransformed <- applyTP f thisNode
             let newTravs = gatherChildren (trav, untrav) thisNodeTransformed
                 baseTransformed = untrav thisNodeTransformed
             bf' f baseTransformed (travs ++ newTravs)

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

newtype TP m = MkTP (forall a. (Typeable a, StrategicData a) => Zipper a -> m (Zipper a))

unTP (MkTP f) = f

applyTP :: (Typeable a, StrategicData a) => TP m -> Zipper a -> m (Zipper a)
applyTP = unTP

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

{-
New stuff
-}

atRoot :: Monad m => TP m -> TP m
atRoot tp = MkTP (\z -> moveM up (\v -> Just $ v.$ arity z) (applyTP tp z) (applyTP (atRoot tp)) z)

full_uptdTP :: (Monad m) => TP m  -> TP m
full_uptdTP f = allTPup (full_uptdTP f) `seqTP` f

full_upbuTP :: (Monad m) => TP m  -> TP m
full_upbuTP f = f `seqTP` allTPup (full_upbuTP f)

once_uptdTP :: (MonadPlus m) => TP m -> TP m
once_uptdTP f =  oneTPup (once_upbuTP f) `choiceTP` f

once_upbuTP :: (MonadPlus m) => TP m -> TP m
once_upbuTP f = f `choiceTP` oneTPup (once_uptdTP f)


full_tdTPupwards :: forall m a. (Typeable a, Monad m) => Proxy a -> TP m -> TP m
full_tdTPupwards _ tp = MkTP $ \z ->
        let (Just v) = (getHole @(Maybe a) z)
            z' = trans forbid z
            traversed = applyTP (atRoot (full_tdTP tp)) z'
        in fmap (setHole v) traversed



-- counts a node if mutable 
counter_func :: (Typeable a, MonadPlus m) => (a -> m a) -> a -> StateT Int m a
counter_func tr e = do
  t <- lift $ tr e
  modify succ
  return e

-- counts the mutable nodes
counting :: (StrategicData a, Typeable n, MonadPlus m) =>
        Zipper a -> (n -> m n) -> m Int
counting r tr = execStateT (applyTP (full_tdTP step) r) 0
 where step = idTP `adhocTPSeq` counter_func tr

-- applies mutation if node is transformable and we are supposed to mutate now
mutation_func :: (MonadPlus m) => (a -> m a) -> a -> StateT Int m a
mutation_func tr e = do
  --run the transformation to check for mzeros (which will end computations)
  t <- lift $ tr e
  modify pred
  x <- get
  if x==0
    then return t
    else return e

-- applies mutation_func with a random counter deciding when to actually mutate
mutating :: (StrategicData a, Typeable b, MonadPlus m) =>
     Zipper a -> Int -> (b -> m b) -> m (Zipper a)
mutating r index tr = evalStateT (applyTP (full_tdTP step) r) index
 where step = idTP `adhocTPSeq` mutation_func tr

once_RandomTP :: (Typeable n, StrategicData a) => Zipper a -> (n -> Maybe n) -> IO (Zipper a)
once_RandomTP r tr = do
  let Just n = counting r tr
  s <- initStdGen
  setStdGen s
  index <- randomRIO (1, n)
  let Just v = mutating r index tr
  return v

mutations :: (Data a, Typeable n, StrategicData a) =>
             a -> (n -> Maybe n) -> [a]
mutations z tr = applyTU (full_tdTU step) $ toZipper z
    where step = failTU `adhocTUZ` select tr
--           select :: (Typeable n, StrategicData a) =>
--                     (n -> Maybe n) -> n -> Zipper a -> [a]
          select tr node zipper = case tr node of
            Nothing -> []
            Just newNode -> let newZipper = setHole newNode zipper
                            in [fromZipper newZipper]


{-
/New stuff
-}

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 b) -> 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 b) -> TP m
adhocTPZSeq f g = MkTP (\z -> do
             val <- elevateZ g z
             maybe (applyTP f z) (applyTP (tryTP f)) val)
                `choiceTP` f

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 b) -> Zipper a -> m (Maybe (Zipper a))
elevateZ g z = runMaybeT $ transM (MaybeT . maybe (return Nothing) (\b -> maybe (return Nothing) (fmap cast . g b) (cast z)) . cast) z

--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)

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

allTPup :: (Monad m) => TP m  -> TP m
allTPup f = MkTP $ \z -> moveM up (\v -> Just $ v.$arity z) (return z) (applyTP f) z
-- (Just . flip (.$) (arity z))

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

oneTPup :: (MonadPlus m) => TP m -> TP m
oneTPup f = MkTP $ \z -> moveM up (\v -> Just $ v.$arity z) mzero (applyTP f) z

--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 leftmost 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 b) -> 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)


-- TODO Make sure that this innermost works for all cases?!
-- note that this is significantly faster than innermost'
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 d = MkTU (forall a. (Typeable a, StrategicData a) => Zipper a -> m d)
unTU (MkTU f) = f


applyTU :: (Typeable a, StrategicData a) => TU m d -> Zipper a -> m d
applyTU = unTU


foldr1TU :: (Monoid (m d), Foldable m, StrategicData a) => TU m d -> Zipper a -> (d -> d -> d) -> d
foldr1TU f z red = foldr1 red $ applyTU (full_tdTU f) z

foldl1TU :: (Monoid (m d), Foldable m, StrategicData a) => TU m d -> Zipper a -> (d -> d -> d) -> d
foldl1TU f z red = foldl1 red $ applyTU (full_tdTU f) z

foldrTU :: (Monoid (m d), Foldable m, StrategicData a) => TU m d -> Zipper a -> (d -> c -> c) -> c -> c
foldrTU f z red i = foldr red i $ applyTU (full_tdTU f) z

foldlTU :: (Monoid (m d), Foldable m, StrategicData a) => TU m d -> Zipper a -> (c -> d -> c) -> c -> c
foldlTU f z red i = foldl red i $ applyTU (full_tdTU f) z


full_tdTU :: Monoid (m d) => TU m d -> TU m d
full_tdTU f = f `seqTU` allTUdown (full_tdTU f) `seqTU` allTUright (full_tdTU f)

full_buTU :: Monoid (m d) => TU m d -> TU m d
full_buTU f = allTUright (full_buTU f) `seqTU` allTUdown (full_buTU f) `seqTU` f

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

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

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

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


{-
New!
-}

full_uptdTU :: (MonadPlus m, Monoid (m d)) => TU m d -> TU m d
full_uptdTU f = allTUup (full_uptdTU f) `seqTU` f

full_upbuTU :: (MonadPlus m, Monoid (m d)) => TU m d -> TU m d
full_upbuTU f = f `seqTU` allTUup (full_upbuTU f)

once_uptdTU :: (MonadPlus m, Monoid (m d)) => TU m d -> TU m d
once_uptdTU f =  allTUup (once_uptdTU f) `choiceTU` f

once_upbuTU :: (MonadPlus m, Monoid (m d)) => TU m d -> TU m d
once_upbuTU f = f `choiceTU` allTUup (once_upbuTU f)

allTUup :: Monoid (m d) => TU m d -> TU m d
allTUup f = MkTU $ \z -> moveQ up mempty (applyTU f) z

{-
/New!
-}

allTUdown :: Monoid (m d) => TU m d -> TU m d
allTUdown f = MkTU $ \z -> moveQ down' mempty (applyTU f) z

allTUright :: Monoid (m d) => TU m d -> TU m d
allTUright f = MkTU $ \z -> moveQ right mempty (applyTU f) z

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


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

reduce :: (Typeable a, Typeable c, Typeable e) => (a -> Zipper c -> m d) -> Zipper e -> Maybe (m d)
reduce f z = (\v -> fmap (f v) (cast z)) =<< getHole z


seqTU :: (Monoid (m d)) => TU m d-> TU m d -> TU m d
seqTU x y = MkTU $ \z -> applyTU x z `mappend` applyTU y z

choiceTU :: (MonadPlus m) => TU m d -> TU m d -> TU m d
choiceTU x y = MkTU $ \z -> applyTU x z `mplus` applyTU y z

failTU :: (MonadPlus m) => TU m d
failTU = MkTU (const mzero)

constTU :: Monad m => d -> TU m d
constTU d = MkTU (const . return $ d)

monoTU :: (MonadPlus m, Typeable a) => (a -> m d) -> TU m d
monoTU = adhocTU failTU

monoTUZ :: (MonadPlus m, Typeable a, Typeable e) => (a -> Zipper e -> m d) -> TU m d
monoTUZ = adhocTUZ failTU