module Hledger.Utils.Tree where

-- import Data.Char
import Data.List (foldl')
import qualified Data.Map as M
import Data.Tree
-- import Text.Megaparsec
-- import Text.Printf

import Hledger.Utils.Regex
-- import Hledger.Utils.UTF8IOCompat (error')

-- standard tree helpers

root :: Tree a -> a
root = Tree a -> a
forall a. Tree a -> a
rootLabel
subs :: Tree a -> Forest a
subs = Tree a -> Forest a
forall a. Tree a -> Forest a
subForest
branches :: Tree a -> Forest a
branches = Tree a -> Forest a
forall a. Tree a -> Forest a
subForest

-- | List just the leaf nodes of a tree
leaves :: Tree a -> [a]
leaves :: Tree a -> [a]
leaves (Node v :: a
v []) = [a
v]
leaves (Node _ branches :: Forest a
branches) = (Tree a -> [a]) -> Forest a -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> [a]
forall a. Tree a -> [a]
leaves Forest a
branches

-- | get the sub-tree rooted at the first (left-most, depth-first) occurrence
-- of the specified node value
subtreeat :: Eq a => a -> Tree a -> Maybe (Tree a)
subtreeat :: a -> Tree a -> Maybe (Tree a)
subtreeat v :: a
v t :: Tree a
t
    | Tree a -> a
forall a. Tree a -> a
root Tree a
t a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
v = Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just Tree a
t
    | Bool
otherwise = a -> [Tree a] -> Maybe (Tree a)
forall a. Eq a => a -> [Tree a] -> Maybe (Tree a)
subtreeinforest a
v ([Tree a] -> Maybe (Tree a)) -> [Tree a] -> Maybe (Tree a)
forall a b. (a -> b) -> a -> b
$ Tree a -> [Tree a]
forall a. Tree a -> Forest a
subs Tree a
t

-- | get the sub-tree for the specified node value in the first tree in
-- forest in which it occurs.
subtreeinforest :: Eq a => a -> [Tree a] -> Maybe (Tree a)
subtreeinforest :: a -> [Tree a] -> Maybe (Tree a)
subtreeinforest _ [] = Maybe (Tree a)
forall a. Maybe a
Nothing
subtreeinforest v :: a
v (t :: Tree a
t:ts :: [Tree a]
ts) = case (a -> Tree a -> Maybe (Tree a)
forall a. Eq a => a -> Tree a -> Maybe (Tree a)
subtreeat a
v Tree a
t) of
                             Just t' :: Tree a
t' -> Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just Tree a
t'
                             Nothing -> a -> [Tree a] -> Maybe (Tree a)
forall a. Eq a => a -> [Tree a] -> Maybe (Tree a)
subtreeinforest a
v [Tree a]
ts

-- | remove all nodes past a certain depth
treeprune :: Int -> Tree a -> Tree a
treeprune :: Int -> Tree a -> Tree a
treeprune 0 t :: Tree a
t = a -> Forest a -> Tree a
forall a. a -> Forest a -> Tree a
Node (Tree a -> a
forall a. Tree a -> a
root Tree a
t) []
treeprune d :: Int
d t :: Tree a
t = a -> Forest a -> Tree a
forall a. a -> Forest a -> Tree a
Node (Tree a -> a
forall a. Tree a -> a
root Tree a
t) ((Tree a -> Tree a) -> Forest a -> Forest a
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Tree a -> Tree a
forall a. Int -> Tree a -> Tree a
treeprune (Int -> Tree a -> Tree a) -> Int -> Tree a -> Tree a
forall a b. (a -> b) -> a -> b
$ Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) (Forest a -> Forest a) -> Forest a -> Forest a
forall a b. (a -> b) -> a -> b
$ Tree a -> Forest a
forall a. Tree a -> Forest a
branches Tree a
t)

-- | apply f to all tree nodes
treemap :: (a -> b) -> Tree a -> Tree b
treemap :: (a -> b) -> Tree a -> Tree b
treemap f :: a -> b
f t :: Tree a
t = b -> Forest b -> Tree b
forall a. a -> Forest a -> Tree a
Node (a -> b
f (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
$ Tree a -> a
forall a. Tree a -> a
root Tree a
t) ((Tree a -> Tree b) -> [Tree a] -> Forest b
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> Tree a -> Tree b
forall a b. (a -> b) -> Tree a -> Tree b
treemap a -> b
f) ([Tree a] -> Forest b) -> [Tree a] -> Forest b
forall a b. (a -> b) -> a -> b
$ Tree a -> [Tree a]
forall a. Tree a -> Forest a
branches Tree a
t)

-- | remove all subtrees whose nodes do not fulfill predicate
treefilter :: (a -> Bool) -> Tree a -> Tree a
treefilter :: (a -> Bool) -> Tree a -> Tree a
treefilter f :: a -> Bool
f t :: Tree a
t = a -> Forest a -> Tree a
forall a. a -> Forest a -> Tree a
Node
                 (Tree a -> a
forall a. Tree a -> a
root Tree a
t)
                 ((Tree a -> Tree a) -> Forest a -> Forest a
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Bool) -> Tree a -> Tree a
forall a. (a -> Bool) -> Tree a -> Tree a
treefilter a -> Bool
f) (Forest a -> Forest a) -> Forest a -> Forest a
forall a b. (a -> b) -> a -> b
$ (Tree a -> Bool) -> Forest a -> Forest a
forall a. (a -> Bool) -> [a] -> [a]
filter ((a -> Bool) -> Tree a -> Bool
forall a. (a -> Bool) -> Tree a -> Bool
treeany a -> Bool
f) (Forest a -> Forest a) -> Forest a -> Forest a
forall a b. (a -> b) -> a -> b
$ Tree a -> Forest a
forall a. Tree a -> Forest a
branches Tree a
t)

-- | is predicate true in any node of tree ?
treeany :: (a -> Bool) -> Tree a -> Bool
treeany :: (a -> Bool) -> Tree a -> Bool
treeany f :: a -> Bool
f t :: Tree a
t = a -> Bool
f (Tree a -> a
forall a. Tree a -> a
root Tree a
t) Bool -> Bool -> Bool
|| (Tree a -> Bool) -> [Tree a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((a -> Bool) -> Tree a -> Bool
forall a. (a -> Bool) -> Tree a -> Bool
treeany a -> Bool
f) (Tree a -> [Tree a]
forall a. Tree a -> Forest a
branches Tree a
t)

-- treedrop -- remove the leaves which do fulfill predicate.
-- treedropall -- do this repeatedly.

-- | show a compact ascii representation of a tree
showtree :: Show a => Tree a -> String
showtree :: Tree a -> String
showtree = [String] -> String
unlines ([String] -> String) -> (Tree a -> [String]) -> Tree a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
filter (String -> String -> Bool
regexMatches "[^ \\|]") ([String] -> [String])
-> (Tree a -> [String]) -> Tree a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines (String -> [String]) -> (Tree a -> String) -> Tree a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree String -> String
drawTree (Tree String -> String)
-> (Tree a -> Tree String) -> Tree a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> String) -> Tree a -> Tree String
forall a b. (a -> b) -> Tree a -> Tree b
treemap a -> String
forall a. Show a => a -> String
show

-- | show a compact ascii representation of a forest
showforest :: Show a => Forest a -> String
showforest :: Forest a -> String
showforest = (Tree a -> String) -> Forest a -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> String
forall a. Show a => Tree a -> String
showtree


-- | An efficient-to-build tree suggested by Cale Gibbard, probably
-- better than accountNameTreeFrom.
newtype FastTree a = T (M.Map a (FastTree a))
  deriving (Int -> FastTree a -> ShowS
[FastTree a] -> ShowS
FastTree a -> String
(Int -> FastTree a -> ShowS)
-> (FastTree a -> String)
-> ([FastTree a] -> ShowS)
-> Show (FastTree a)
forall a. Show a => Int -> FastTree a -> ShowS
forall a. Show a => [FastTree a] -> ShowS
forall a. Show a => FastTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [FastTree a] -> ShowS
$cshowList :: forall a. Show a => [FastTree a] -> ShowS
show :: FastTree a -> String
$cshow :: forall a. Show a => FastTree a -> String
showsPrec :: Int -> FastTree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> FastTree a -> ShowS
Show, FastTree a -> FastTree a -> Bool
(FastTree a -> FastTree a -> Bool)
-> (FastTree a -> FastTree a -> Bool) -> Eq (FastTree a)
forall a. Eq a => FastTree a -> FastTree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: FastTree a -> FastTree a -> Bool
$c/= :: forall a. Eq a => FastTree a -> FastTree a -> Bool
== :: FastTree a -> FastTree a -> Bool
$c== :: forall a. Eq a => FastTree a -> FastTree a -> Bool
Eq, Eq (FastTree a)
Eq (FastTree a) =>
(FastTree a -> FastTree a -> Ordering)
-> (FastTree a -> FastTree a -> Bool)
-> (FastTree a -> FastTree a -> Bool)
-> (FastTree a -> FastTree a -> Bool)
-> (FastTree a -> FastTree a -> Bool)
-> (FastTree a -> FastTree a -> FastTree a)
-> (FastTree a -> FastTree a -> FastTree a)
-> Ord (FastTree a)
FastTree a -> FastTree a -> Bool
FastTree a -> FastTree a -> Ordering
FastTree a -> FastTree a -> FastTree a
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (FastTree a)
forall a. Ord a => FastTree a -> FastTree a -> Bool
forall a. Ord a => FastTree a -> FastTree a -> Ordering
forall a. Ord a => FastTree a -> FastTree a -> FastTree a
min :: FastTree a -> FastTree a -> FastTree a
$cmin :: forall a. Ord a => FastTree a -> FastTree a -> FastTree a
max :: FastTree a -> FastTree a -> FastTree a
$cmax :: forall a. Ord a => FastTree a -> FastTree a -> FastTree a
>= :: FastTree a -> FastTree a -> Bool
$c>= :: forall a. Ord a => FastTree a -> FastTree a -> Bool
> :: FastTree a -> FastTree a -> Bool
$c> :: forall a. Ord a => FastTree a -> FastTree a -> Bool
<= :: FastTree a -> FastTree a -> Bool
$c<= :: forall a. Ord a => FastTree a -> FastTree a -> Bool
< :: FastTree a -> FastTree a -> Bool
$c< :: forall a. Ord a => FastTree a -> FastTree a -> Bool
compare :: FastTree a -> FastTree a -> Ordering
$ccompare :: forall a. Ord a => FastTree a -> FastTree a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (FastTree a)
Ord)

emptyTree :: FastTree a
emptyTree = Map a (FastTree a) -> FastTree a
forall a. Map a (FastTree a) -> FastTree a
T Map a (FastTree a)
forall k a. Map k a
M.empty

mergeTrees :: (Ord a) => FastTree a -> FastTree a -> FastTree a
mergeTrees :: FastTree a -> FastTree a -> FastTree a
mergeTrees (T m :: Map a (FastTree a)
m) (T m' :: Map a (FastTree a)
m') = Map a (FastTree a) -> FastTree a
forall a. Map a (FastTree a) -> FastTree a
T ((FastTree a -> FastTree a -> FastTree a)
-> Map a (FastTree a) -> Map a (FastTree a) -> Map a (FastTree a)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith FastTree a -> FastTree a -> FastTree a
forall a. Ord a => FastTree a -> FastTree a -> FastTree a
mergeTrees Map a (FastTree a)
m Map a (FastTree a)
m')

treeFromPath :: [a] -> FastTree a
treeFromPath :: [a] -> FastTree a
treeFromPath []     = Map a (FastTree a) -> FastTree a
forall a. Map a (FastTree a) -> FastTree a
T Map a (FastTree a)
forall k a. Map k a
M.empty
treeFromPath (x :: a
x:xs :: [a]
xs) = Map a (FastTree a) -> FastTree a
forall a. Map a (FastTree a) -> FastTree a
T (a -> FastTree a -> Map a (FastTree a)
forall k a. k -> a -> Map k a
M.singleton a
x ([a] -> FastTree a
forall a. [a] -> FastTree a
treeFromPath [a]
xs))

treeFromPaths :: (Ord a) => [[a]] -> FastTree a
treeFromPaths :: [[a]] -> FastTree a
treeFromPaths = (FastTree a -> FastTree a -> FastTree a)
-> FastTree a -> [FastTree a] -> FastTree a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' FastTree a -> FastTree a -> FastTree a
forall a. Ord a => FastTree a -> FastTree a -> FastTree a
mergeTrees FastTree a
forall a. FastTree a
emptyTree ([FastTree a] -> FastTree a)
-> ([[a]] -> [FastTree a]) -> [[a]] -> FastTree a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> FastTree a) -> [[a]] -> [FastTree a]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> FastTree a
forall a. [a] -> FastTree a
treeFromPath