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

-- | This module implements a tree of 'Data.Map.Map'.
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
import           Data.Monoid (Monoid(..))
import           Data.Traversable (Traversable(..))
import           Data.Typeable (Typeable)

-- * The 'Path' type

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

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

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

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

-- * The 'TreeMap' type

type TreeMap k x = Data.Map.Map k (Node k x)
data Ord k => Node k x
 =   Node
 { node_size        :: Int -- ^ The number of non-'Nothing' 'node_content's reachable from this 'Node'.
 , node_content     :: Maybe x -- ^ Some content, 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_content = Nothing
		 , node_size = 0
		 , node_descendants = mempty
		 }
	mappend
	 Node{node_content=x0, node_descendants=m0}
	 Node{node_content=x1, node_descendants=m1} =
		let m = union const m0 m1 in
		let x = x0 `mappend` x1 in
		Node
		 { node_content = 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_content=x, node_descendants=m, node_size} =
		Node
		 { node_content = fmap f x
		 , node_descendants = Hcompta.Lib.TreeMap.map f m
		 , node_size
		 }

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

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

-- * Contructors

empty :: TreeMap k x
empty = Data.Map.empty

singleton :: Ord k => Path k -> x -> TreeMap k x
singleton ks x = insert const ks x Data.Map.empty

leaf :: Ord k => x -> Node k x
leaf x =
	Node
	 { node_content     = Just x
	 , node_descendants = Data.Map.empty
	 , node_size        = 1
	 }

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

-- | Return a 'TreeMap' associating the given 'Path' to the given content,
-- merging content 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 the same 'Path's as both given 'TreeMap's,
-- merging contents (in respective order) when a 'Path' leads
-- to a non-'Nothing' 'node_content' in both given 'TreeMap's.
union :: Ord k => (x -> x -> x) -> TreeMap k x -> TreeMap k x -> TreeMap k x
union merge =
	Data.Map.unionWith
	 (\Node{node_content=x0, node_descendants=m0}
	   Node{node_content=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_content = x
		 , node_descendants = m
		 , node_size = size m + maybe 0 (const 1) x
		 })

-- | 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

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

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

-- * Extractors

-- | Return the number of non-'Nothing' 'node_content's in the given 'TreeMap'.
size :: Ord k => TreeMap k x -> Int
size = Data.Map.foldr ((+) . node_size) 0

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

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

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

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