module Data.TreeSeq.Strict.Zipper where

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

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

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

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

zipper :: Tree k a -> Zipper k a
zipper t = Node mempty t mempty :| []

zippers :: Trees k a -> [Zipper k a]
zippers ts = ns >>= axis_collect axis_following_first
	where ns =
		case Seq.viewl ts of
		 EmptyL -> empty
		 l :< ls -> pure $ Node mempty l ls :| []

zipper_root :: Zipper k a -> Tree k a
zipper_root =
	zip_self .
	NonEmpty.head .
	List.last .
	axis_ancestor_or_self

path :: Zipper k x -> [k]
path ns =
	List.reverse $
	NonEmpty.toList ns >>= \n ->
		case zip_self n of
		 TreeN k _ -> [k]
		 Tree0{}   -> []

current :: Zipper k a -> Tree k a
current (Node _ t _ :| _) = t

at :: Alternative f =>
      Axis k a -> Int ->
      (Zipper k a -> f (Zipper k a))
at axis i n =
	case List.drop i (axis n) of
	 []  -> empty
	 a:_ -> pure a
infixl 5 `at`

null :: Axis k a -> Zipper k a -> Bool
null axis = List.null . axis

-- ** Type 'Node'
data Node k a
 =   Node
 {   zip_prec :: Trees k a
 ,   zip_self :: Tree  k a
 ,   zip_foll :: Trees k a
 } deriving (Eq, Show, Typeable)

-- * Type 'Axis'
type Axis k a = 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 = Alternative f => Zipper k a -> f (Zipper k a)

-- | Collect all 'Zipper's along a given axis,
--   including the first 'Zipper'.
axis_collect :: (n -> Maybe n) -> n -> [n]
axis_collect f n = n : maybe [] (axis_collect f) (f n)

-- | Collect all 'Zipper's along a given axis,
--   excluding the first 'Zipper'.
axis_collect_without_self :: (n -> Maybe n) -> n -> [n]
axis_collect_without_self f n = maybe [] (axis_collect f) (f n)

-- ** Axis self
axis_self :: Applicative f => Zipper k a -> f (Tree k a)
axis_self (Node _ t _ :| _) = pure t

-- ** Axis child
axis_child :: Axis k a
axis_child n =
	axis_child_first n >>=
	axis_collect axis_following_first

axis_child_lookup_first :: (k -> Bool) -> AxisAlt f k a
axis_child_lookup_first fk n = safeHead $ axis_child_lookup fk n

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

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

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

-- ** Axis ancestor
axis_ancestor :: Axis k a
axis_ancestor = axis_collect_without_self axis_parent

axis_ancestor_or_self :: Axis k a
axis_ancestor_or_self = axis_collect axis_parent

-- ** Axis descendant
axis_descendant_or_self :: Axis k a
axis_descendant_or_self =
	collect_child []
	where
	collect_child acc n =
		n : maybe acc
		 (collect_following_first acc)
		 (axis_child_first n)
	collect_following_first acc n =
		collect_child
		 (maybe acc
			 (collect_following_first acc)
			 (axis_following_first n)
		 ) n

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

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

-- ** Axis preceding
axis_preceding_first :: AxisAlt f k a
axis_preceding_first (Node ps t fs :| ns) =
	case Seq.viewr ps of
	 EmptyR -> empty
	 rs :> r -> pure $ Node rs r (t <| fs) :| ns

axis_preceding_sibling :: Axis k a
axis_preceding_sibling = axis_collect_without_self axis_preceding_first

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

-- ** Axis following
axis_following_first :: AxisAlt f k a
axis_following_first (Node ps t fs :| ns) =
	case Seq.viewl fs of
	 EmptyL -> empty
	 l :< ls -> pure $ Node (ps |> t) l ls :| ns

axis_following_sibling :: Axis k a
axis_following_sibling = axis_collect_without_self axis_following_first

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

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

-- ** Filter
axis_filter :: Axis k a -> (Zipper k a -> Bool) -> Axis k a
axis_filter axis p n = List.filter p (axis n)
infixl 5 `axis_filter`