module Data.TreeSeq.Strict.Zipper where

import Control.Arrow (Kleisli(..))
import Control.Category (Category(..), (>>>))
import Control.Applicative (Applicative(..), Alternative(..))
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 a = NonEmpty (Cursor a)

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

-- | Return a 'Zipper' starting at the left-most 'Tree' of the given 'Trees'.
zippers :: Trees a -> [Zipper 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 a -> Cursor a
root = NonEmpty.head . List.last . runAxis axis_ancestor_or_self

-- | Like 'root', but concatenate the 'Cursor' into a 'Trees'.
roots :: Zipper a -> Trees a
roots z = ps <> (s <| fs)
	where Cursor ps s fs = root z

-- | Return the keys within the 'Tree' nodes
-- leading to the current 'Cursor' of the given 'Zipper'.
zipath :: Zipper a -> [a]
zipath z =
	List.reverse $
	unTree . cursor_self
	 <$> NonEmpty.toList z

-- | Return the 'Tree's selected by the given 'Axis' from the given 'Zipper'.
select :: Axis a -> Zipper a -> [Tree 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 a -> (Zipper a -> Maybe b) -> Zipper a -> [b]
filter axis f z = f `mapMaybe` runAxis axis z

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

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

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

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

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

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

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

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

-- ** Axis @repeat@
-- | Collect all 'Zipper's along a given axis,
--   including the first 'Zipper'.
axis_repeat :: AxisAlt Maybe a -> Axis 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 a -> Axis a
axis_repeat_without_self f = Kleisli $ \z -> maybe [] (runAxis $ axis_repeat f) (runAxisAlt f z)

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

axis_filter_current :: Axis a -> (Tree a -> Bool) -> Axis 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 a -> Axis a
axis_first axis = Kleisli $ List.take 1 . runAxis axis

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

-- ** Axis @at@
axis_at :: Alternative f => Axis a -> Int -> AxisAlt f 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 a
axis_self = Kleisli pure

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

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

axis_child_lookup :: (a -> Bool) -> Axis a
axis_child_lookup f = Kleisli $ \z@(Cursor _ps t _fs :| _) ->
	let ns = subTrees t in
	(<$> Seq.findIndicesL (f . unTree) 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

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

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

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

axis_ancestor_or_self :: Axis a
axis_ancestor_or_self = axis_repeat axis_parent

axis_root :: Alternative f => AxisAlt f a
axis_root = Kleisli $ pure . List.last . runAxis axis_ancestor_or_self

-- ** Axis @descendant@
axis_descendant_or_self :: Axis 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 a
axis_descendant_or_self_reverse = Kleisli go
	where go z = z : List.concatMap go (List.reverse $ runAxis axis_child z)

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

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

axis_preceding_sibling_nearest :: Alternative f => AxisAlt f 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 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 a
axis_preceding =
	axis_ancestor_or_self >>>
	axis_preceding_sibling >>>
	axis_descendant_or_self_reverse

-- ** Axis @following@
axis_following_sibling :: Axis a
axis_following_sibling = axis_repeat_without_self axis_following_sibling_nearest

axis_following_sibling_nearest :: Alternative f => AxisAlt f 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 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 a
axis_following =
	axis_ancestor_or_self >>>
	axis_following_sibling >>>
	axis_descendant_or_self

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

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