{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# OPTIONS_GHC -fno-warn-tabs #-}
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.Maybe (Maybe(..), maybe)
import Data.Monoid (Monoid(..))
import Data.Sequence (ViewL(..), ViewR(..), (<|), (|>))
import Data.Semigroup (Semigroup(..))
import Data.Typeable (Typeable)
import Prelude (undefined)
import Text.Show (Show(..))
import qualified Data.Sequence as Seq
import qualified Data.List as List

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

keyTree :: Tree k a -> k
keyTree (TreeN k _) = k
keyTree Tree0{}     = undefined

-- * Type 'Zipper'
data Zipper k a
 =   Zipper
 {   zipper_path :: [Zipper_Step k a]
 ,   zipper_curr :: Trees k a
 } deriving (Eq, Show, Typeable)

zipper :: Trees k a -> Zipper k a
zipper = Zipper []

zipper_root :: Zipper k a -> Trees k a
zipper_root = zipper_curr . List.last . zipper_ancestor_or_self

path_of_zipper :: Zipper k x -> [k]
path_of_zipper z =
	keyTree . zipper_step_self <$>
	List.reverse (zipper_path z)

-- * Type 'Zipper_Step'
data Zipper_Step k a
 =   Zipper_Step
 {   zipper_step_prec :: Trees k a
 ,   zipper_step_self :: Tree  k a
 ,   zipper_step_foll :: Trees k a
 } deriving (Eq, Show, Typeable)

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

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

-- ** Axis self
zipper_self :: Zipper k a -> [Tree k a]
zipper_self (Zipper (Zipper_Step _ t _ : _) _) = [t]
zipper_self _ = []

-- ** Axis child
zipper_child :: Zipper k a -> [Zipper k a]
zipper_child z =
	zipper_child_first z >>=
	zipper_collect zipper_foll

zipper_child_lookup ::
 Alternative f =>
 (k -> Bool) -> Zipper k a -> f (Zipper k a)
zipper_child_lookup fk z = safeHead $ zipper_childs_lookup fk z

zipper_childs_lookup ::
 (k -> Bool) -> Zipper k a -> [Zipper k a]
zipper_childs_lookup fk (Zipper path ts) =
	(<$> Seq.findIndicesL (\case TreeN k _ -> fk k; Tree0{} -> False) ts) $ \i ->
		let (ps, ps') = Seq.splitAt i ts in
		case Seq.viewl ps' of
		 EmptyL -> undefined
		 t :< fs ->
			Zipper
			 { zipper_path = Zipper_Step ps t fs : path
			 , zipper_curr = nodesTree t
			 }

zipper_child_first :: Alternative f => Zipper k a -> f (Zipper k a)
zipper_child_first (Zipper path trees) =
	case Seq.viewl trees of
	 EmptyL -> empty
	 t :< ts -> pure $ Zipper
		 { zipper_path = Zipper_Step mempty t ts : path
		 , zipper_curr = nodesTree t
		 }

zipper_child_last :: Alternative f => Zipper k a -> f (Zipper k a)
zipper_child_last (Zipper path trees) =
	case Seq.viewr trees of
	 EmptyR -> empty
	 ts :> t -> pure $ Zipper
		 { zipper_path = Zipper_Step ts t mempty : path
		 , zipper_curr = nodesTree t
		 }

-- ** Axis ancestor
zipper_ancestor :: Zipper k a -> [Zipper k a]
zipper_ancestor = zipper_collect_without_self zipper_parent

zipper_ancestor_or_self :: Zipper k a -> [Zipper k a]
zipper_ancestor_or_self = zipper_collect zipper_parent

-- ** Axis descendant
zipper_descendant_or_self :: Zipper k a -> [Zipper k a]
zipper_descendant_or_self =
	collect_child []
	where
		collect_child acc z =
			z : maybe acc
			 (collect_foll acc)
			 (zipper_child_first z)
		collect_foll  acc z =
			collect_child
			 (maybe acc
				 (collect_foll acc)
				 (zipper_foll z)
			 ) z

zipper_descendant_or_self_reverse :: Zipper k a -> [Zipper k a]
zipper_descendant_or_self_reverse z =
	z : List.concatMap
	 zipper_descendant_or_self_reverse
	 (List.reverse $ zipper_child z)

zipper_descendant :: Zipper k a -> [Zipper k a]
zipper_descendant = List.tail . zipper_descendant_or_self

-- ** Axis preceding
zipper_prec :: Alternative f => Zipper k a -> f (Zipper k a)
zipper_prec (Zipper [] _curr) = empty
zipper_prec (Zipper (Zipper_Step ps c fs : path) _curr) =
	case Seq.viewr ps of
	 EmptyR -> empty
	 ts :> t -> pure Zipper
		 { zipper_path = Zipper_Step ts t (c <| fs) : path
		 , zipper_curr = nodesTree t
		 }

zipper_preceding :: Zipper k a -> [Zipper k a]
zipper_preceding =
	zipper_ancestor_or_self >=>
	zipper_preceding_sibling >=>
	zipper_descendant_or_self_reverse

zipper_preceding_sibling :: Zipper k a -> [Zipper k a]
zipper_preceding_sibling = zipper_collect_without_self zipper_prec

-- ** Axis following
zipper_foll :: Alternative f => Zipper k a -> f (Zipper k a)
zipper_foll (Zipper [] _curr) = empty
zipper_foll (Zipper (Zipper_Step ps c fs:path) _curr) =
	case Seq.viewl fs of
	 EmptyL -> empty
	 t :< ts -> pure $ Zipper
		 { zipper_path = Zipper_Step (ps |> c) t ts : path
		 , zipper_curr = nodesTree t
		 }

zipper_following :: Zipper k a -> [Zipper k a]
zipper_following =
	zipper_ancestor_or_self >=>
	zipper_following_sibling >=>
	zipper_descendant_or_self

zipper_following_sibling :: Zipper k a -> [Zipper k a]
zipper_following_sibling = zipper_collect_without_self zipper_foll

-- ** Axis parent
zipper_parent :: Alternative f => Zipper k a -> f (Zipper k a)
zipper_parent (Zipper [] _) = empty
zipper_parent (Zipper (Zipper_Step ps c fs : path) curr) =
	pure Zipper
	 { zipper_path = path
	 , zipper_curr = (ps |> m) <> fs
	 }
	where
	m = case c of
	 TreeN k _ -> TreeN k curr
	 Tree0{} -> undefined

-- ** Filter
zipper_filter ::
 (Zipper k a -> [Zipper k a]) ->
 (Zipper k a -> Bool) ->
 (Zipper k a -> [Zipper k a])
zipper_filter axis p z = List.filter p (axis z)
infixl 5 `zipper_filter`

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

zipper_null ::
 (Zipper k a -> [Zipper k a]) ->
 Zipper k a -> Bool
zipper_null axis = List.null . axis