module Data.TreeSeq.Strict.Zipper where

import Control.Arrow (Kleisli(..))
import Control.Applicative (Applicative(..), Alternative(..))
import Control.Monad (Monad(..), (>=>))
import Data.Bool
import Data.Eq (Eq)
import Data.Function (($), (.))
import Data.Functor ((<$>))
import Data.Int (Int)
import Data.List.NonEmpty (NonEmpty(..))
import Data.Maybe (Maybe(..), maybe, mapMaybe)
import Data.Monoid (Monoid(..))
import Data.Semigroup (Semigroup(..))
import Data.Sequence (ViewL(..), ViewR(..), (<|), (|>))
import Data.Typeable (Typeable)
import Prelude (undefined)
import Text.Show (Show(..))
import qualified Data.List as List
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Sequence as Seq

import Data.TreeSeq.Strict (Trees, Tree(..))

-- * Type 'Zipper'
type Zipper k a = NonEmpty (Cursor k a)

-- | Return a 'Zipper' starting at the given 'Tree'.
zipper :: Tree k a -> Zipper k a
zipper t = Cursor mempty t mempty :| []

-- | Return a 'Zipper' starting at the left-most 'Tree' of the given 'Trees'.
zippers :: Trees k a -> [Zipper k a]
zippers ts =
	case Seq.viewl ts of
	 EmptyL -> empty
	 l :< ls -> pure $ Cursor mempty l ls :| []

-- | Return the 'Cursor' after zipping the given 'Zipper' upto its last parent.
root :: Zipper k a -> Cursor k a
root = NonEmpty.head . List.last . axis_ancestor_or_self

-- | Return the keys within the 'TreeN' nodes
-- leading to the current 'Cursor' of the given 'Zipper'.
path :: Zipper k x -> [k]
path cs =
	List.reverse $
	NonEmpty.toList cs >>= \c ->
		case cursor_self c of
		 TreeN k _ -> [k]
		 Tree0{}   -> []

-- | Return the 'Tree' currently under the 'Cursor'.
current :: Zipper k a -> Tree k a
current (Cursor _ t _ :| _) = t

-- | Return the 'Tree's selected by the given 'Axis' from the given 'Zipper'.
select :: Axis k a -> Zipper k a -> [Tree k a]
select axis c = cursor_self . NonEmpty.head <$> axis c

-- | Return the filtered values selected by the given 'Axis' from the given 'Zipper'.
filter :: Axis k a -> (Zipper k a -> Maybe b) -> Zipper k a -> [b]
filter axis f c = f `mapMaybe` axis c

-- ** Type 'Cursor'
data Cursor k a
 =   Cursor
 {   cursor_precedings :: Trees k a
 ,   cursor_self       :: Tree  k a
 ,   cursor_followings :: Trees k a
 } deriving (Eq, Show, Typeable)

-- * Type 'Axis'
type Axis k a = Zipper k a -> [Zipper k a]

-- ** Type 'KleisliAxis'
type KleisliAxis k a = Kleisli [] (Zipper k a) (Zipper k a)

-- ** Type 'AxisAlt'
-- | Like 'Axis', but generalized with 'Alternative'.
--
-- Useful to return a 'Maybe' instead of a list.
type AxisAlt f k a = Zipper k a -> f (Zipper k a)

-- ** Axis @repeat@
-- | Collect all 'Zipper's along a given axis,
--   including the first 'Zipper'.
axis_repeat :: AxisAlt Maybe k a -> Axis k a
axis_repeat f c = c : maybe [] (axis_repeat f) (f c)

-- | Collect all 'Zipper's along a given axis,
--   excluding the starting 'Zipper'.
axis_repeat_without_self :: AxisAlt Maybe k a -> Axis k a
axis_repeat_without_self f c = maybe [] (axis_repeat f) (f c)

-- ** Axis @filter@
axis_filter :: Axis k a -> (Zipper k a -> Bool) -> Axis k a
axis_filter axis f c = List.filter f (axis c)
infixl 5 `axis_filter`

axis_filter_current :: Axis k a -> (Tree k a -> Bool) -> Axis k a
axis_filter_current axis f c = List.filter (f . current) (axis c)
infixl 5 `axis_filter_current`

-- ** Axis @at@
axis_at :: Alternative f => Axis k a -> Int -> AxisAlt f k a
axis_at axis i c =
	case List.drop i $ axis c of
	 []  -> empty
	 a:_ -> pure a
infixl 5 `axis_at`

-- ** Axis @self@
axis_self :: Applicative f => AxisAlt f k a
axis_self = pure

-- ** Axis @child@
axis_child :: Axis k a
axis_child c =
	axis_child_first c >>=
	axis_repeat axis_following1

axis_child_lookup_first :: Alternative f => (k -> Bool) -> AxisAlt f k a
axis_child_lookup_first fk c = listHead $ axis_child_lookup fk c

axis_child_lookup :: (k -> Bool) -> Axis k a
axis_child_lookup fk cs@(Cursor _ps t _fs :| _) =
	(<$> Seq.findIndicesL flt ns) $ \i ->
		let (ps, ps') = Seq.splitAt i ns in
		case Seq.viewl ps' of
		 EmptyL -> undefined
		 l :< ls -> Cursor ps l ls :| NonEmpty.toList cs
	where
	ns = nodesTree t
	flt (TreeN k _) = fk k
	flt Tree0{}     = False

axis_child_first :: Alternative f => AxisAlt f k a
axis_child_first cs@(Cursor _ps t _fs :| _) =
	case Seq.viewl $ nodesTree t of
	 EmptyL -> empty
	 l :< ls -> pure $ Cursor mempty l ls :| NonEmpty.toList cs

axis_child_last :: Alternative f => AxisAlt f k a
axis_child_last cs@(Cursor _ps t _fs :| _) =
	case Seq.viewr $ nodesTree t of
	 EmptyR -> empty
	 rs :> r -> pure $ Cursor rs r mempty :| NonEmpty.toList cs

-- ** Axis @ancestor@
axis_ancestor :: Axis k a
axis_ancestor = axis_repeat_without_self axis_parent

axis_ancestor_or_self :: Axis k a
axis_ancestor_or_self = axis_repeat axis_parent

axis_root :: Alternative f => AxisAlt f k a
axis_root = pure . List.last . axis_ancestor_or_self

-- ** Axis @descendant@
axis_descendant_or_self :: Axis k a
axis_descendant_or_self =
	collect_child []
	where
	collect_child acc c =
		c : maybe acc
		 (collect_following_first acc)
		 (axis_child_first c)
	collect_following_first acc c =
		collect_child
		 (maybe acc
			 (collect_following_first acc)
			 (axis_following1 c)
		 ) c

axis_descendant_or_self_reverse :: Axis k a
axis_descendant_or_self_reverse c =
	c :
	List.concatMap
	 axis_descendant_or_self_reverse
	 (List.reverse $ axis_child c)

axis_descendant :: Axis k a
axis_descendant = List.tail . axis_descendant_or_self

-- ** Axis @preceding@
axis_preceding1 :: Alternative f => AxisAlt f k a
axis_preceding1 (Cursor ps t fs :| cs) =
	case Seq.viewr ps of
	 EmptyR  -> empty
	 rs :> r -> pure $ Cursor rs r (t <| fs) :| cs

axis_preceding_sibling :: Axis k a
axis_preceding_sibling = axis_repeat_without_self axis_preceding1

axis_preceding_sibling_first :: Alternative f => AxisAlt f k a
axis_preceding_sibling_first z@(Cursor ps t fs :| cs) =
	case Seq.viewl (ps |> t) of
	 EmptyL  -> pure z
	 l :< ls -> pure $ Cursor mempty l (ls<>fs) :| cs

axis_preceding :: Axis k a
axis_preceding =
	axis_ancestor_or_self >=>
	axis_preceding_sibling >=>
	axis_descendant_or_self_reverse

-- ** Axis @following@
axis_following1 :: Alternative f => AxisAlt f k a
axis_following1 (Cursor ps t fs :| cs) =
	case Seq.viewl fs of
	 EmptyL  -> empty
	 l :< ls -> pure $ Cursor (ps |> t) l ls :| cs

axis_following_sibling :: Axis k a
axis_following_sibling = axis_repeat_without_self axis_following1

axis_following_sibling_last :: Alternative f => AxisAlt f k a
axis_following_sibling_last z@(Cursor ps t fs :| cs) =
	case Seq.viewr (t <| fs) of
	 EmptyR  -> pure z
	 rs :> r -> pure $ Cursor (ps<>rs) r mempty :| cs

axis_following :: Axis k a
axis_following =
	axis_ancestor_or_self >=>
	axis_following_sibling >=>
	axis_descendant_or_self

-- ** Axis @parent@
axis_parent :: Alternative f => AxisAlt f k a
axis_parent (Cursor ps t fs :| cs) =
	case cs of
	 Cursor ps' (TreeN k _) fs' : cs' ->
		pure $ Cursor ps' (TreeN k $ (ps |> t) <> fs) fs' :| cs'
	 _ -> empty

-- * Utilities
nodesTree :: Tree k a -> Trees k a
nodesTree Tree0{}       = mempty
nodesTree (TreeN _k ts) = ts

listHead :: Alternative f => [a] -> f a
listHead []    = empty
listHead (a:_) = pure a