module Data.TreeSeq.Strict.Zipper where

import Control.Arrow (Kleisli(..))
import Control.Category (Category(..), (>>>))
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 . runAxis axis_ancestor_or_self

-- | Like 'root', but concatenate the 'Cursor' into a 'Trees'.
roots :: Zipper k a -> Trees k a
roots z = cursor_preceding_siblings <> (cursor_self <| cursor_following_siblings)
	where Cursor{..} = root z

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

-- | 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 z = cursor_self . NonEmpty.head <$> runAxis axis z

-- | 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 z = f `mapMaybe` runAxis axis z

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

-- | Return the current 'Cursor' of a 'Zipper'.
cursor :: Zipper k a -> Cursor k a
cursor = NonEmpty.head

-- | Set the current 'Cursor' of a 'Zipper'.
setCursor :: Zipper k a -> Cursor k a -> Zipper k a
setCursor (_c :| cs) c = c :| cs

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

-- ** Type 'Axis'
type Axis k a = AxisAlt [] k a

runAxis :: Axis k a -> Zipper k a -> [Zipper k a]
runAxis = runKleisli

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

runAxisAlt :: AxisAlt f k a -> Zipper k a -> f (Zipper k a)
runAxisAlt = runKleisli

-- ** 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 = Kleisli $ \z -> z : maybe [] (runAxis $ axis_repeat f) (runAxisAlt f z)

-- | 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 = Kleisli $ \z -> maybe [] (runAxis $ axis_repeat f) (runAxisAlt f z)

-- ** Axis @filter@
axis_filter :: Axis k a -> (Zipper k a -> Bool) -> Axis k a
axis_filter axis f = Kleisli $ \z -> List.filter f (runAxis axis z)
infixl 5 `axis_filter`

axis_filter_current :: Axis k a -> (Tree k a -> Bool) -> Axis k a
axis_filter_current axis f = Kleisli $ \z -> List.filter (f . current) (runAxis axis z)
infixl 5 `axis_filter_current`

-- ** Axis @first@
axis_first :: Axis k a -> Axis k a
axis_first axis = Kleisli $ List.take 1 . runAxis axis

-- ** Axis @last@
axis_last :: Axis k a -> Axis k a
axis_last axis = Kleisli $ List.take 1 . List.reverse . runAxis axis

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

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

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

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

axis_child_lookup :: (k -> Bool) -> Axis k a
axis_child_lookup fk = Kleisli $ \z@(Cursor _ps t _fs :| _) ->
	let ns = nodesTree t in
	(<$> 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 z
	where
	flt = \case
	 TreeN k _ -> fk k
	 Tree0{}   -> False

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

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

-- ** 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 = Kleisli $ pure . List.last . runAxis axis_ancestor_or_self

-- ** Axis @descendant@
axis_descendant_or_self :: Axis k a
axis_descendant_or_self =
	Kleisli $ collect_child []
	where
	collect_child acc z =
		z : maybe acc
		 (collect_following_first acc)
		 (runAxisAlt axis_child_first z)
	collect_following_first acc z =
		collect_child
		 (maybe acc
			 (collect_following_first acc)
			 (runAxisAlt axis_following_sibling_nearest z)
		 ) z

axis_descendant_or_self_reverse :: Axis k a
axis_descendant_or_self_reverse = Kleisli go
	where go z = z : List.concatMap go (List.reverse $ runAxis axis_child z)

axis_descendant :: Axis k a
axis_descendant = Kleisli $ List.tail . runAxis axis_descendant_or_self

-- ** Axis @preceding@
axis_preceding_sibling :: Axis k a
axis_preceding_sibling = axis_repeat_without_self axis_preceding_sibling_nearest

axis_preceding_sibling_nearest :: Alternative f => AxisAlt f k a
axis_preceding_sibling_nearest = Kleisli $ \(Cursor ps t fs :| cs) ->
	case Seq.viewr ps of
	 EmptyR  -> empty
	 rs :> r -> pure $ Cursor rs r (t <| fs) :| cs

axis_preceding_sibling_farthest :: Alternative f => AxisAlt f k a
axis_preceding_sibling_farthest = Kleisli $ \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_following_sibling :: Axis k a
axis_following_sibling = axis_repeat_without_self axis_following_sibling_nearest

axis_following_sibling_nearest :: Alternative f => AxisAlt f k a
axis_following_sibling_nearest = Kleisli $ \(Cursor ps t fs :| cs) ->
	case Seq.viewl fs of
	 EmptyL  -> empty
	 l :< ls -> pure $ Cursor (ps |> t) l ls :| cs

axis_following_sibling_farthest :: Alternative f => AxisAlt f k a
axis_following_sibling_farthest = Kleisli $ \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 = Kleisli $ \(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