{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE NamedFieldPuns #-}

-- | This module implements a 'TreeMap',
-- which is like a 'Map'
-- but whose key is now a 'NonEmpty' list of 'Map' keys (a 'Path')
-- enabling the possibility to gather mapped values
-- by 'Path' prefixes (inside a 'Node').
module Hcompta.Lib.TreeMap where

import           Control.Applicative ((<$>), (<*>), pure)
import           Data.Data (Data)
import           Data.Foldable (Foldable(..))
import qualified Data.List
import qualified Data.List.NonEmpty
import           Data.List.NonEmpty (NonEmpty(..))
import qualified Data.Map.Strict as Data.Map
import           Data.Map.Strict (Map)
import           Data.Monoid (Monoid(..))
import           Data.Traversable (Traversable(..))
import           Data.Typeable (Typeable)
import           Prelude hiding (filter, null, reverse)

-- * Type 'TreeMap'

newtype TreeMap k x
 =      TreeMap (Map k (Node k x))
 deriving (Data, Eq, Read, Show, Typeable)

instance (Ord k, Monoid v) => Monoid (TreeMap k v) where
	mempty = empty
	mappend = union const
	-- mconcat = Data.List.foldr mappend mempty
instance Ord k => Functor (TreeMap k) where
	fmap f (TreeMap m) = TreeMap $ fmap (fmap f) m
instance Ord k => Foldable (TreeMap k) where
	foldMap f (TreeMap m) = foldMap (foldMap f) m
instance Ord k => Traversable (TreeMap k) where
	traverse f (TreeMap m) = TreeMap <$> traverse (traverse f) m

-- * Type 'Path'

-- | A 'Path' is a non-empty list of 'Map' keys.
type Path k = NonEmpty k

path :: k -> [k] -> Path k
path = (:|)

list :: Path k -> [k]
list = Data.List.NonEmpty.toList

reverse :: Path k -> Path k
reverse = Data.List.NonEmpty.reverse

-- * Type 'Node'
data Ord k
 => Node k x
 =  Node
 { node_size        :: Int -- ^ The number of non-'Nothing' 'node_value's reachable from this 'Node'.
 , node_value       :: Maybe x -- ^ Some value, or 'Nothing' if this 'Node' is intermediary.
 , node_descendants :: TreeMap k x -- ^ Descendants 'Node's.
 } deriving (Data, Eq, Read, Show, Typeable)

instance (Ord k, Monoid v) => Monoid (Node k v) where
	mempty =
		Node
		 { node_value = Nothing
		 , node_size = 0
		 , node_descendants = TreeMap mempty
		 }
	mappend
	 Node{node_value=x0, node_descendants=m0}
	 Node{node_value=x1, node_descendants=m1} =
		let m = union const m0 m1 in
		let x = x0 `mappend` x1 in
		Node
		 { node_value = x
		 , node_size = size m + maybe 0 (const 1) x
		 , node_descendants = union const m0 m1
		 }
	-- mconcat = Data.List.foldr mappend mempty

instance Ord k => Functor (Node k) where
	fmap f Node{node_value=x, node_descendants=m, node_size} =
		Node
		 { node_value = fmap f x
		 , node_descendants = Hcompta.Lib.TreeMap.map f m
		 , node_size
		 }

instance Ord k => Foldable (Node k) where
	foldMap f Node{node_value=Nothing, node_descendants=TreeMap m} =
		foldMap (foldMap f) m
	foldMap f Node{node_value=Just x, node_descendants=TreeMap m} =
		f x `mappend` foldMap (foldMap f) m

instance Ord k => Traversable (Node k) where
	traverse f Node{node_value=Nothing, node_descendants=TreeMap m, node_size} =
		Node node_size <$> pure Nothing <*> (TreeMap <$> traverse (traverse f) m)
	traverse f Node{node_value=Just x, node_descendants=TreeMap m, node_size} =
		Node node_size <$> (Just <$> f x) <*> (TreeMap <$> traverse (traverse f) m)

-- * Construct

-- | Return the empty 'TreeMap'.
empty :: TreeMap k x
empty = TreeMap Data.Map.empty

-- | Return a 'TreeMap' only mapping the given 'Path' to the given value.
singleton :: Ord k => Path k -> x -> TreeMap k x
singleton ks x = insert const ks x empty

-- | Return a 'Node' only containing the given value.
leaf :: Ord k => x -> Node k x
leaf x =
	Node
	 { node_value = Just x
	 , node_descendants = empty
	 , node_size = 1
	 }

-- | Return the given 'TreeMap' associating the given 'Path' with the given value,
-- merging values if the given 'TreeMap' already associates the given 'Path'
-- with a non-'Nothing' 'node_value'.
insert :: Ord k => (x -> x -> x) -> Path k -> x -> TreeMap k x -> TreeMap k x
insert merge (k:|[]) x (TreeMap m) =
	TreeMap $
	Data.Map.insertWith
	 (\_ Node{node_value=x1, node_descendants=m1, node_size=s1} ->
		Node
		 { node_value = maybe (Just x) (Just . merge x) x1
		 , node_descendants = m1
		 , node_size = maybe (s1 + 1) (const s1) x1
		 })
	 k (leaf x) m
insert merge (k:|k':ks) x (TreeMap m) =
	TreeMap $
	Data.Map.insertWith
	 (\_ Node{node_value=x1, node_descendants=m1} ->
		let m' = insert merge (path k' ks) x m1 in
		Node{node_value=x1, node_descendants=m', node_size=size m' + maybe 0 (const 1) x1})
	 k
	 (Node
		 { node_value = Nothing
		 , node_descendants = insert merge (path k' ks) x empty
		 , node_size = 1
		 })
	 m

-- | Return a 'TreeMap' associating for each tuple of the given list
-- the 'Path' to the value,
-- merging values of identical 'Path's (in respective order).
from_List :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
from_List merge = Data.List.foldl (\acc (p, x) -> insert merge p x acc) empty

-- | Return a 'TreeMap' associating for each key and value of the given 'Map'
-- the 'Path' to the value,
-- merging values of identical 'Path's (in respective order).
from_Map :: Ord k => (x -> x -> x) -> Map (Path k) x -> TreeMap k x
from_Map merge = Data.Map.foldlWithKey (\acc p x -> insert merge p x acc) empty

-- * Size

-- | Return the 'Map' in the given 'TreeMap'.
nodes :: TreeMap k x -> Map k (Node k x)
nodes (TreeMap m) = m

-- | Return 'True' iif. the given 'TreeMap' is 'empty'.
null :: TreeMap k x -> Bool
null (TreeMap m) = Data.Map.null m

-- | Return the number of non-'Nothing' 'node_value's in the given 'TreeMap'.
--
--   * Complexity: O(r) where r is the size of the root 'Map'.
size :: Ord k => TreeMap k x -> Int
size = Data.Map.foldr ((+) . node_size) 0 . nodes

-- * Find

-- | Return the value (if any) associated with the given 'Path'.
find :: Ord k => Path k -> TreeMap k x -> Maybe x
find (k:|[]) (TreeMap m) = maybe Nothing node_value $ Data.Map.lookup k m
find (k:|k':ks) (TreeMap m) =
	maybe Nothing (find (path k' ks) . node_descendants) $
	Data.Map.lookup k m

-- * Union

-- | Return a 'TreeMap' associating the same 'Path's as both given 'TreeMap's,
-- merging values (in respective order) when a 'Path' leads
-- to a non-'Nothing' 'node_value' in both given 'TreeMap's.
union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
union merge (TreeMap tm0) (TreeMap tm1) =
	TreeMap $
	Data.Map.unionWith
	 (\Node{node_value=x0, node_descendants=m0}
	   Node{node_value=x1, node_descendants=m1} ->
		let m = union merge m0 m1 in
		let x = maybe x1 (\x0' -> maybe (Just x0') (Just . merge x0') x1) x0 in
		Node
		 { node_value = x
		 , node_descendants = m
		 , node_size = size m + maybe 0 (const 1) x
		 })
	 tm0 tm1

-- | Return the 'union' of the given 'TreeMap's.
--
-- NOTE: use 'Data.List.foldl'' to reduce demand on the control-stack.
unions :: Ord k => (x -> x -> x) -> [TreeMap k x] -> TreeMap k x
unions merge ts = Data.List.foldl' (union merge) empty ts

-- foldl' :: (a -> b -> a) -> a -> [b] -> a
-- foldl' f = go
-- 	where
-- 		go z []     = z
-- 		go z (x:xs) = z `seq` go (f z x) xs

-- * Map

-- | Return the given 'TreeMap' with each non-'Nothing' 'node_value'
-- mapped by the given function.
map :: Ord k => (x -> y) -> TreeMap k x -> TreeMap k y
map f =
	TreeMap .
	Data.Map.map
	 (\n@Node{node_value=x, node_descendants=m} ->
		n{ node_value=maybe Nothing (Just . f) x
		 , node_descendants=Hcompta.Lib.TreeMap.map f m
		 }) .
	nodes

-- | Return the given 'TreeMap' with each 'node_value'
-- mapped by the given function supplied with
-- the already mapped 'node_descendants' of the current 'Node'.
map_by_depth_first :: Ord k => (TreeMap k y -> Maybe x -> y) -> TreeMap k x -> TreeMap k y
map_by_depth_first f =
	TreeMap .
	Data.Map.map
	 (\n@Node{node_value, node_descendants} ->
		let m = map_by_depth_first f node_descendants in
		let x = f m node_value in
		n{ node_value = Just x
		 , node_descendants = m
		 , node_size = size m + 1
		 }) .
	nodes

-- * Alter

alterl_path :: Ord k => (Maybe x -> Maybe x) -> Path k -> TreeMap k x -> TreeMap k x
alterl_path fct =
	go fct . list
	where
		go :: Ord k
		 => (Maybe x -> Maybe x) -> [k]
		 -> TreeMap k x -> TreeMap k x
		go _f [] m = m
		go f (k:p) (TreeMap m) =
			TreeMap $
			Data.Map.alter
			 (\c ->
				let (cv, cm) =
					case c of
					 Just Node{node_value=v, node_descendants=d} -> (v, d)
					 Nothing -> (Nothing, empty) in
				let fx = f cv in
				let gm = go f p cm in
				case (fx, size gm) of
				 (Nothing, 0) -> Nothing
				 (_, s) -> Just
					Node
					 { node_value = fx
					 , node_descendants = gm
					 , node_size = s + 1
					 }
			 ) k m

-- * Fold

-- | Return the given accumulator folded by the given function
-- applied on non-'Nothing' 'node_value's
-- from left to right through the given 'TreeMap'.
foldl_with_Path :: Ord k => (a -> Path k -> x -> a) -> a -> TreeMap k x -> a
foldl_with_Path =
	foldp []
	where
		foldp :: Ord k
		 => [k] -> (a -> Path k -> x -> a)
		 -> a -> TreeMap k x -> a
		foldp p fct a (TreeMap m) =
			Data.Map.foldlWithKey
			 (\acc k Node{node_value, node_descendants} ->
				let acc' = maybe acc (fct acc (reverse $ path k p)) node_value in
				foldp (k:p) fct acc' node_descendants) a m

-- | Return the given accumulator folded by the given function
-- applied on non-'Nothing' 'Node's and 'node_value's
-- from left to right through the given 'TreeMap'.
foldl_with_Path_and_Node :: Ord k => (a -> Node k x -> Path k -> x -> a) -> a -> TreeMap k x -> a
foldl_with_Path_and_Node =
	foldp []
	where
		foldp :: Ord k
		 => [k] -> (a -> Node k x -> Path k -> x -> a)
		 -> a -> TreeMap k x -> a
		foldp p fct a (TreeMap m) =
			Data.Map.foldlWithKey
			 (\acc k n@Node{node_value, node_descendants} ->
				let acc' = maybe acc (fct acc n (reverse $ path k p)) node_value in
				foldp (k:p) fct acc' node_descendants) a m

-- | Return the given accumulator folded by the given function
-- applied on non-'Nothing' 'node_value's
-- from right to left through the given 'TreeMap'.
foldr_with_Path :: Ord k => (Path k -> x -> a -> a) -> a -> TreeMap k x -> a
foldr_with_Path =
	foldp []
	where
		foldp :: Ord k
		 => [k] -> (Path k -> x -> a -> a)
		 -> a -> TreeMap k x -> a
		foldp p fct a (TreeMap m) =
			Data.Map.foldrWithKey
			 (\k Node{node_value, node_descendants} acc ->
				let acc' = foldp (k:p) fct acc node_descendants in
				maybe acc' (\x -> fct (reverse $ path k p) x acc') node_value) a m

-- | Return the given accumulator folded by the given function
-- applied on non-'Nothing' 'Node's and 'node_value's
-- from right to left through the given 'TreeMap'.
foldr_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> a -> a) -> a -> TreeMap k x -> a
foldr_with_Path_and_Node =
	foldp []
	where
		foldp :: Ord k
		 => [k] -> (Node k x -> Path k -> x -> a -> a)
		 -> a -> TreeMap k x -> a
		foldp p fct a (TreeMap m) =
			Data.Map.foldrWithKey
			 (\k n@Node{node_value, node_descendants} acc ->
				let acc' = foldp (k:p) fct acc node_descendants in
				maybe acc' (\x -> fct n (reverse $ path k p) x acc') node_value) a m

-- | Return the given accumulator folded by the given function
-- applied on non-'Nothing' 'node_value's
-- from left to right along the given 'Path'.
foldl_path :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
foldl_path fct =
	go fct [] . list
	where
		go :: Ord k
		 => (Path k -> x -> a -> a) -> [k] -> [k]
		 -> TreeMap k x -> a -> a
		go _f _ [] _t a = a
		go f p (k:n) (TreeMap t) a =
			case Data.Map.lookup k t of
			 Nothing -> a
			 Just Node{node_value=v, node_descendants=d} ->
				case v of
				 Nothing -> go f (k:p) n d a
				 Just x  -> go f (k:p) n d (f (reverse $ path k p) x a)

-- | Return the given accumulator folded by the given function
-- applied on non-'Nothing' 'node_value's
-- from right to left along the given 'Path'.
foldr_path :: Ord k => (Path k -> x -> a -> a) -> Path k -> TreeMap k x -> a -> a
foldr_path fct =
	go fct [] . list
	where
		go :: Ord k
		 => (Path k -> x -> a -> a) -> [k] -> [k]
		 -> TreeMap k x -> a -> a
		go _f _ [] _t a = a
		go f p (k:n) (TreeMap t) a =
			case Data.Map.lookup k t of
			 Nothing -> a
			 Just Node{node_value=v, node_descendants=d} ->
				case v of
				 Nothing -> go f (k:p) n d a
				 Just x  -> f (reverse $ path k p) x $ go f (k:p) n d a

-- * Flatten

-- | Return a 'Map' associating each 'Path'
-- leading to a non-'Nothing' 'node_value' in the given 'TreeMap',
-- with its value mapped by the given function.
flatten :: Ord k => (x -> y) -> TreeMap k x -> Map (Path k) y
flatten =
	flat_map []
	where
		flat_map :: Ord k
		 => [k] -> (x -> y)
		 -> TreeMap k x
		 -> Map (Path k) y
		flat_map p f (TreeMap m) =
			Data.Map.unions $
			(
			Data.Map.mapKeysMonotonic (reverse . flip path p) $
			Data.Map.mapMaybe (\Node{node_value=x} -> f <$> x) m
			) :
			Data.Map.foldrWithKey
			 (\k -> (:) . flat_map (k:p) f . node_descendants)
			 [] m

-- * Filter

-- | Return the given 'TreeMap'
--   keeping only its non-'Nothing' 'node_value's
--   passing the given predicate.
filter :: Ord k => (x -> Bool) -> TreeMap k x -> TreeMap k x
filter f =
	map_Maybe_with_Path
	 (\_p x -> if f x then Just x else Nothing)

-- | Like 'filter' but with also the current 'Path' given to the predicate.
filter_with_Path :: Ord k => (Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
filter_with_Path f =
	map_Maybe_with_Path
	 (\p x -> if f p x then Just x else Nothing)

-- | Like 'filter_with_Path' but with also the current 'Node' given to the predicate.
filter_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Bool) -> TreeMap k x -> TreeMap k x
filter_with_Path_and_Node f =
	map_Maybe_with_Path_and_Node
	 (\n p x -> if f n p x then Just x else Nothing)

-- | Return the given 'TreeMap'
--   mapping its non-'Nothing' 'node_value's
--   and keeping only the non-'Nothing' results.
map_Maybe :: Ord k => (x -> Maybe y) -> TreeMap k x -> TreeMap k y
map_Maybe f = map_Maybe_with_Path (const f)

-- | Like 'map_Maybe' but with also the current 'Path' given to the predicate.
map_Maybe_with_Path :: Ord k => (Path k -> x -> Maybe y) -> TreeMap k x -> TreeMap k y
map_Maybe_with_Path f = map_Maybe_with_Path_and_Node (const f)

-- | Like 'map_Maybe_with_Path' but with also the current 'Node' given to the predicate.
map_Maybe_with_Path_and_Node :: Ord k => (Node k x -> Path k -> x -> Maybe y) -> TreeMap k x -> TreeMap k y
map_Maybe_with_Path_and_Node =
	go []
	where
		go :: Ord k
		 => [k] -> (Node k x -> Path k -> x -> Maybe y)
		 -> TreeMap k x
		 -> TreeMap k y
		go p test (TreeMap m) =
			TreeMap $
			Data.Map.mapMaybeWithKey
			 (\k node@Node{node_value=v, node_descendants=ns} ->
				let node_descendants = go (k:p) test ns in
				let node_size = size node_descendants in
				case v of
				 Just x ->
					let node_value = test node (reverse $ path k p) x in
					case node_value of
					 Nothing | null node_descendants -> Nothing
					 Nothing -> Just Node{node_value, node_descendants, node_size=1 + node_size}
					 Just _  -> Just Node{node_value, node_descendants, node_size}
				 _ ->
					if null node_descendants
					then Nothing
					else Just Node{node_value=Nothing, node_descendants, node_size}
			 ) m