{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# OPTIONS_GHC -fno-warn-tabs #-}

module Hcompta.Lib.TreeMap.Zipper where

import           Control.Applicative ((<$>), (<*>), pure)
import           Control.DeepSeq (NFData(..))
import           Control.Monad (Monad(..), (>=>))
import           Data.Bool
import           Data.Eq (Eq)
import           Data.Data (Data)
import           Data.Foldable (Foldable, foldMap)
import           Data.Functor (Functor(..))
import           Data.Ord (Ord(..))
import qualified Data.List as List
import qualified Data.List.NonEmpty
import           Data.List.NonEmpty (NonEmpty(..))
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Maybe (Maybe(..), maybe, maybeToList)
import           Data.Monoid (Monoid(..))
import qualified Data.Strict.Maybe as Strict
import           Data.Traversable (Traversable(..))
import           Data.Typeable (Typeable)
import           Prelude (($), (.), Int, Num(..), Show, const, flip, id, seq)

import qualified Hcompta.Lib.Strict as Strict ()
import           Hcompta.Lib.TreeMap (TreeMap(..))
import qualified Hcompta.Lib.TreeMap as TreeMap

-- * Type 'Zipper'

data Zipper k x
 =   Zipper
 {   zipper_path :: [Zipper_Step k x]
 ,   zipper_curr :: TreeMap k x
 } deriving (Data, Eq, Show, Typeable)

zipper :: TreeMap k x -> Zipper k x
zipper = Zipper []

zipper_root :: Ord k => Zipper k x -> TreeMap k x
zipper_root =
	zipper_curr . List.last .
	zipper_collect zipper_parent

-- * Type 'Zipper_Step'

data Zipper_Step k x
 =   Zipper_Step
 {   zipper_step_prec :: TreeMap k x
 ,   zipper_step_self :: (k, TreeMap.Node k x)
 ,   zipper_step_foll :: TreeMap k x
 } deriving (Data, 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 :: Ord k => Zipper k x -> Maybe (k, TreeMap.Node k x)
zipper_self z =
	case z of
	 Zipper{ zipper_path=
	         Zipper_Step{zipper_step_self}
	         : _ } -> Just zipper_step_self
	 _ -> Nothing

-- ** Axis child

zipper_child :: Ord k => Zipper k x -> [Zipper k x]
zipper_child z =
	maybeToList (zipper_child_first z)
	>>= zipper_collect zipper_foll

zipper_child_at :: Ord k => k -> Zipper k x -> Maybe (Zipper k x)
zipper_child_at k (Zipper path (TreeMap m)) =
	case Map.splitLookup k m of
	 (_, Nothing, _) -> Nothing
	 (ps, Just s, fs) ->
		Just $ Zipper
		 { zipper_path = Zipper_Step (TreeMap ps) (k, s) (TreeMap fs) : path
		 , zipper_curr = TreeMap.node_descendants s
		 }

zipper_child_first :: Ord k => Zipper k x -> Maybe (Zipper k x)
zipper_child_first (Zipper path (TreeMap m)) =
	case Map.minViewWithKey m of
	 Nothing -> Nothing
	 Just ((k', s'), fs') ->
		Just $ Zipper
		 { zipper_path = Zipper_Step TreeMap.empty (k', s') (TreeMap fs') : path
		 , zipper_curr = TreeMap.node_descendants s'
		 }

zipper_child_last :: Ord k => Zipper k x -> Maybe (Zipper k x)
zipper_child_last (Zipper path (TreeMap m)) =
	case Map.maxViewWithKey m of
	 Nothing -> Nothing
	 Just ((k', s'), ps') ->
		Just $ Zipper
		 { zipper_path = Zipper_Step (TreeMap ps') (k', s') TreeMap.empty : path
		 , zipper_curr = TreeMap.node_descendants s'
		 }

-- ** Axis ancestor

zipper_ancestor :: Ord k => Zipper k x -> [Zipper k x]
zipper_ancestor = zipper_collect_without_self zipper_parent

zipper_ancestor_or_self :: Ord k => Zipper k x -> [Zipper k x]
zipper_ancestor_or_self = zipper_collect zipper_parent

-- ** Axis descendant

zipper_descendant_or_self :: Ord k => Zipper k x -> [Zipper k x]
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 :: Ord k => Zipper k x -> [Zipper k x]
zipper_descendant_or_self_reverse z =
	z : List.concatMap
	 zipper_descendant_or_self_reverse
	 (List.reverse $ zipper_child z)

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

zipper_descendant_at :: Ord k => TreeMap.Path k -> Zipper k x -> Maybe (Zipper k x)
zipper_descendant_at (k:|ks) =
	case ks of
	 []     -> zipper_child_at k
	 k':ks' -> zipper_child_at k >=> zipper_descendant_at (k':|ks')

-- ** Axis preceding

zipper_prec :: Ord k => Zipper k x -> Maybe (Zipper k x)
zipper_prec (Zipper path _curr) =
	case path of
	 [] -> Nothing
	 Zipper_Step (TreeMap ps) (k, s) (TreeMap fs):steps ->
		case Map.maxViewWithKey ps of
		 Nothing -> Nothing
		 Just ((k', s'), ps') ->
			Just $ Zipper
			 { zipper_path = Zipper_Step (TreeMap ps')
			                             (k', s')
			                             (TreeMap $ Map.insert k s $ fs)
			                 : steps
			 , zipper_curr = TreeMap.node_descendants s'
			 }

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

zipper_preceding_sibling :: Ord k => Zipper k x -> [Zipper k x]
zipper_preceding_sibling = zipper_collect_without_self zipper_prec

-- ** Axis following

zipper_foll :: Ord k => Zipper k x -> Maybe (Zipper k x)
zipper_foll (Zipper path _curr) =
	case path of
	 [] -> Nothing
	 Zipper_Step (TreeMap ps) (k, s) (TreeMap fs):steps ->
		case Map.minViewWithKey fs of
		 Nothing -> Nothing
		 Just ((k', s'), fs') ->
			Just $ Zipper
			 { zipper_path = Zipper_Step (TreeMap $ Map.insert k s $ ps)
			                             (k', s')
			                             (TreeMap fs')
			                 : steps
			 , zipper_curr = TreeMap.node_descendants s'
			 }

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

zipper_following_sibling :: Ord k => Zipper k x -> [Zipper k x]
zipper_following_sibling = zipper_collect_without_self zipper_foll

-- ** Axis parent

zipper_parent :: Ord k => Zipper k x -> Maybe (Zipper k x)
zipper_parent (Zipper path curr) =
	case path of
	 [] -> Nothing
	 Zipper_Step (TreeMap ps) (k, s) (TreeMap fs):steps ->
		let node = TreeMap.Node
			 { TreeMap.node_value       = TreeMap.node_value s
			 , TreeMap.node_size        = TreeMap.size curr
			 , TreeMap.node_descendants = curr
			 } in
		Just $ Zipper
		 { zipper_path = steps
		 , zipper_curr = TreeMap $ Map.union ps $ Map.insert k node $ fs
		 }