PolyTree.hs

{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Data.PolyTree (
  -- * Types
  Tree(..)
, TreeForest(..)
, Tree'
, TreeForest'
, TreeList
, TreeList'
, Tree1
, Tree1'
  -- * Construction
, makeTree
, makeChild
, makeLeaves
, makeChildren
, singleton
  -- * Traversals
, dfs
, bfs
  -- * Utility
, unfoldTree
, unfoldTreeM
, pruneLeaves
, countNodes
, countLeaves
, levels
  -- * Optics classes
, GetTree(..)
, HasTree(..)
, GetTreeForest(..)
, HasTreeForest(..)
, ReviewTree(..)
, AsTree(..)
, ReviewTreeForest(..)
, AsTreeForest(..)
  -- * Optics
, treeForest'
, treeSubForest
, treeLeaves
, treeForestChildren
  -- * Conversions
, baseTree
) where

#if !MIN_VERSION_base(4,18,0)
import Control.Applicative ( Applicative(liftA2) )
#endif
import Control.Lens hiding ((<.>), levels)
import Data.Bifoldable ( Bifoldable(bifoldMap) )
import Data.Bitraversable ( Bitraversable(..) )
import Data.Functor.Apply ( Apply(liftF2, (<.>)) )
import Data.Functor.Classes
    ( showsBinaryWith,
      showsUnaryWith,
      Eq1(..),
      Eq2(..),
      Ord1(..),
      Ord2(..),
      Show1(liftShowsPrec),
      Show2(..) )
import Data.Functor.Identity ()
import Data.List.NonEmpty ( NonEmpty(..), nonEmpty, toList )
import Data.Monoid ( Sum(..) )
import Data.Semigroup.Bifoldable ( Bifoldable1(bifoldMap1) )
import Data.Semigroup.Bitraversable ( Bitraversable1(bitraverse1) )
import Data.Semigroup.Foldable ( Foldable1(foldMap1) )
import Data.Semigroup.Traversable ()
import qualified Data.Tree as Tree

-- $setup
-- >>> import Control.Lens hiding ((<.>), levels)
-- >>> import Data.Functor.Apply ((<.>))
-- >>> import Data.Functor.Identity
-- >>> import Data.List.NonEmpty (NonEmpty(..))
-- >>> import qualified Data.PolyTree
-- >>> :set -XOverloadedStrings

-- | A forest of trees, represented as a functor containing either leaves (b) or subtrees.
-- The 'TreeForest' type allows for a polymorphic container type 'f', enabling different
-- forest representations (lists, vectors, etc.).
--
-- Examples:
--
-- >>> TreeForest [] :: TreeForest [] String String
-- TreeForest []
--
-- >>> TreeForest [Left "leaf1", Left "leaf2"] :: TreeForest [] String String
-- TreeForest [Left "leaf1",Left "leaf2"]
--
-- >>> TreeForest [Right (Tree "node" (TreeForest []))] :: TreeForest [] String String
-- TreeForest [Right (Tree "node" (TreeForest []))]
newtype TreeForest f a b =
  TreeForest (f (Either b (Tree f a b)))

type TreeForest' f a =
  TreeForest f a a

instance (TreeForest f a b ~ t) =>
  Rewrapped (TreeForest f' a' b') t

instance Wrapped (TreeForest f a b) where
  type Unwrapped (TreeForest f a b) =
    f (Either b (Tree f a b))
  _Wrapped' =
    iso (\(TreeForest x) -> x) TreeForest

instance Eq1 f => Eq2 (TreeForest f) where
  liftEq2 f g (TreeForest x1) (TreeForest x2) =
    liftEq (liftEq2 g (liftEq2 f g)) x1 x2

instance Ord1 f => Ord2 (TreeForest f) where
  liftCompare2 f g (TreeForest x1) (TreeForest x2) =
    liftCompare (liftCompare2 g (liftCompare2 f g)) x1 x2

instance Show1 f => Show2 (TreeForest f) where
  liftShowsPrec2 spA slA spB slB d (TreeForest x) =
    let spT =
          liftShowsPrec2 spA slA spB slB
        slT =
          liftShowList2 spA slA spB slB
    in  showsUnaryWith
          (liftShowsPrec
                (liftShowsPrec2 spB slB spT slT)
                (liftShowList2 spB slB spT slT))
          "TreeForest"
          d
          x

instance (Eq a, Eq1 f) => Eq1 (TreeForest f a) where
  liftEq =
    liftEq2 (==)

instance (Ord a, Ord1 f) => Ord1 (TreeForest f a) where
  liftCompare =
    liftCompare2 compare

instance (Show a, Show1 f) => Show1 (TreeForest f a) where
  liftShowsPrec =
    liftShowsPrec2 showsPrec showList

-- |
--
-- >>> TreeForest [Left "a"] == TreeForest [Left "a"] :: Bool
-- True
--
-- >>> TreeForest [Left "a"] == TreeForest [Left "b"] :: Bool
-- False
instance (Eq a, Eq1 f, Eq b) => Eq (TreeForest f a b) where
  (==) =
    liftEq (==)

-- |
--
-- >>> compare (TreeForest [Left "a"]) (TreeForest [Left "b"]) :: Ordering
-- LT
--
-- >>> compare (TreeForest [Left "z"]) (TreeForest [Left "a"]) :: Ordering
-- GT
instance (Ord a, Ord1 f, Ord b) => Ord (TreeForest f a b) where
  compare =
    liftCompare compare

-- |
--
-- >>> show (TreeForest [Left "a", makeChild "b" []]) :: String
-- "TreeForest [Left \"a\",Right (Tree \"b\" (TreeForest []))]"
instance (Show a, Show1 f, Show b) => Show (TreeForest f a b) where
  showsPrec =
    liftShowsPrec showsPrec shows

-- | Map over both the node labels and leaf values in a forest.
--
-- >>> bimap (+1) (*10) (TreeForest [Left 5, makeChild 3 [Left 7]]) :: TreeForest [] Int Int
-- TreeForest [Left 50,Right (Tree 4 (TreeForest [Left 70]))]
instance Functor f => Bifunctor (TreeForest f) where
  bimap f g (TreeForest x) =
    TreeForest (fmap (bimap g (bimap f g)) x)

-- | Map over the leaf values in a forest, keeping node labels unchanged.
--
-- >>> fmap (*2) (TreeForest [Left 5, Left 10]) :: TreeForest [] Int Int
-- TreeForest [Left 10,Left 20]
instance Functor f => Functor (TreeForest f a) where
  fmap =
    bimap id

-- NOTE: Biapply and Biapplicative instances are NOT possible for TreeForest
-- because the recursive structure makes it impossible to correctly handle the
-- case where we have a tree of functions and a single value (or vice versa).
-- The bifunctor structure over (a, b) doesn't align with the Either-based
-- recursive forest representation.

-- | Apply instance for TreeForest using semigroup on node labels.
--
-- >>> TreeForest [Left (+1), Left (*2)] <.> TreeForest [Left 5] :: TreeForest [] String Int
-- TreeForest [Left 6,Left 10]
instance (Apply f, Semigroup a) => Apply (TreeForest f a) where
  TreeForest x1 <.> TreeForest x2 =
    let combine (Left f) (Left x) =
          Left (f x)
        combine (Left f) (Right tx) =
          Right (fmap f tx)
        combine (Right tf) (Left x) =
          Right (fmap ($ x) tf)
        combine (Right tf) (Right tx) =
          Right (tf <.> tx)
    in  TreeForest (liftF2 combine x1 x2)

-- | Applicative instance for TreeForest using monoid on node labels.
--
-- >>> pure 42 :: TreeForest [] String Int
-- TreeForest [Left 42]
--
-- >>> TreeForest [Left (+1)] <*> TreeForest [Left 5, Left 10] :: TreeForest [] String Int
-- TreeForest [Left 6,Left 11]
instance (Applicative f, Monoid a) => Applicative (TreeForest f a) where
  pure b =
    TreeForest (pure (Left b))
  TreeForest x1 <*> TreeForest x2 =
    let combine (Left f) (Left x) =
          Left (f x)
        combine (Left f) (Right tx) =
          Right (fmap f tx)
        combine (Right tf) (Left x) =
          Right (fmap ($ x) tf)
        combine (Right tf) (Right tx) =
          Right (tf <*> tx)
    in  TreeForest (liftA2 combine x1 x2)

-- | Fold over both node labels and leaf values in a forest.
--
-- >>> import Data.Monoid (Sum(..))
-- >>> bifoldMap Sum Sum (TreeForest [Left 5, makeChild 3 [Left 7]]) :: Sum Int
-- Sum {getSum = 15}
instance Foldable f => Bifoldable (TreeForest f) where
  bifoldMap f g (TreeForest x) =
    foldMap (either g (bifoldMap f g)) x

-- | Bifoldable1 for non-empty forests.
--
-- >>> import Data.Semigroup (Sum(..))
-- >>> bifoldMap1 Sum Sum (TreeForest (Left 5 :| [Left 3])) :: Sum Int
-- Sum {getSum = 8}
instance Foldable1 f => Bifoldable1 (TreeForest f) where
  bifoldMap1 f g (TreeForest x) =
    foldMap1 (either g (bifoldMap1 f g)) x

-- | Fold over the leaf values in a forest.
--
-- >>> foldMap (:[]) (TreeForest [Left "a", Left "b", makeChild 1 [Left "c"]]) :: [String]
-- ["a","b","c"]
instance Foldable f => Foldable (TreeForest f a) where
  foldMap f (TreeForest x) =
    foldMap (either f (foldMap f)) x

-- | Foldable1 for non-empty forests.
--
-- >>> import Data.Semigroup (Sum(..))
-- >>> foldMap1 Sum (TreeForest (Left 5 :| [Left 3])) :: Sum Int
-- Sum {getSum = 8}
instance Foldable1 f => Foldable1 (TreeForest f a) where
  foldMap1 f (TreeForest x) =
    foldMap1 (either f (foldMap1 f)) x

-- | Traverse both node labels and leaf values with effects.
--
-- >>> bitraverse Just Just (TreeForest [Left "x"]) :: Maybe (TreeForest [] String String)
-- Just (TreeForest [Left "x"])
instance Traversable f => Bitraversable (TreeForest f) where
  bitraverse f g (TreeForest x) =
    TreeForest <$> traverse (either (fmap Left . g) (fmap Right . bitraverse f g)) x

-- | Bitraversable1 for non-empty forests.
--
-- >>> import Data.Functor.Identity
-- >>> bitraverse1 Identity Identity (TreeForest (Left "x" :| [])) :: Identity (TreeForest NonEmpty String String)
-- Identity (TreeForest (Left "x" :| []))
instance Traversable1 f => Bitraversable1 (TreeForest f) where
  bitraverse1 f g (TreeForest x) =
    TreeForest <$> traverse1 (either (fmap Left . g) (fmap Right . bitraverse1 f g)) x

-- | Traverse the leaf values with effects.
--
-- >>> traverse Just (TreeForest [Left "x", Left "y"]) :: Maybe (TreeForest [] Int String)
-- Just (TreeForest [Left "x",Left "y"])
instance Traversable f => Traversable (TreeForest f a) where
  traverse f (TreeForest x) =
    TreeForest <$> traverse (either (fmap Left . f) (fmap Right . traverse f)) x

-- | Traversable1 for non-empty forests.
--
-- >>> import Data.Functor.Identity
-- >>> traverse1 Identity (TreeForest (Left "x" :| [])) :: Identity (TreeForest NonEmpty Int String)
-- Identity (TreeForest (Left "x" :| []))
instance Traversable1 f => Traversable1 (TreeForest f a) where
  traverse1 f (TreeForest x) =
    TreeForest <$> traverse1 (either (fmap Left . f) (fmap Right . traverse1 f)) x

-- | Read-only access to a tree forest structure via a 'Getter'.
class GetTreeForest x f a b | x -> f a b where
  -- | Extract the tree forest from a structure.
  getTreeForest ::
    Getter x (TreeForest f a b)

-- |
--
-- >>> view getTreeForest (TreeForest [Left "a", makeChild "b" []]) == TreeForest [Left "a", Right (Tree "b" (TreeForest []))]
-- True
instance GetTreeForest (TreeForest f a b) f a b where
  getTreeForest =
    to id

-- | Read-write access to a tree forest structure via a 'Lens'.
-- This extends 'GetTreeForest' to allow modification.
class GetTreeForest x f a b => HasTreeForest x f a b | x -> f a b where
  -- | Access the tree forest with read-write capability.
  treeForest ::
    Lens' x (TreeForest f a b)

-- |
--
-- >>> view treeForest (TreeForest []) :: TreeForest [] String String
-- TreeForest []
instance HasTreeForest (TreeForest f a b) f a b where
  treeForest =
    id

-- | Construction-only access to a tree forest via a 'Review'.
class ReviewTreeForest x f a b | x -> f a b where
  -- | Construct a value from a tree forest.
  reviewTreeForest ::
    Review x (TreeForest f a b)

-- |
--
-- >>> review reviewTreeForest (TreeForest [Left "x"]) :: TreeForest [] String String
-- TreeForest [Left "x"]
instance ReviewTreeForest (TreeForest f a b) f a b where
  reviewTreeForest =
    unto id

-- | Full prism access to a tree forest structure.
-- This extends 'ReviewTreeForest' to allow both construction and pattern matching.
class ReviewTreeForest x f a b => AsTreeForest x f a b | x -> f a b where
  -- | Access the tree forest as a prism (construct or pattern match).
  _TreeForest ::
    Prism' x (TreeForest f a b)

-- |
--
-- >>> preview _TreeForest (TreeForest [Left "y"]) :: Maybe (TreeForest [] String String)
-- Just (TreeForest [Left "y"])
instance AsTreeForest (TreeForest f a b) f a b where
  _TreeForest =
    id

-- | A polymorphic rose tree with different types for node labels (a) and leaf values (b).
-- The container type 'f' allows for different representations of child forests.
--
-- Examples:
--
-- >>> Tree "root" (TreeForest []) :: TreeList String String
-- Tree "root" (TreeForest [])
--
-- >>> Tree "root" (TreeForest [Left "leaf1", Left "leaf2"]) :: TreeList String String
-- Tree "root" (TreeForest [Left "leaf1",Left "leaf2"])
--
-- >>> Tree "root" (TreeForest [makeChild "child" [Left "leaf"]]) :: TreeList String String
-- Tree "root" (TreeForest [Right (Tree "child" (TreeForest [Left "leaf"]))])
data Tree f a b =
  Tree a (TreeForest f a b)

-- | A tree where node labels and leaf values have the same type.
type Tree' f a =
  Tree f a a

-- | A tree using lists for the forest container.
type TreeList a b =
  Tree [] a b

-- | A list-based tree where nodes and leaves have the same type.
type TreeList' a =
  TreeList a a

-- | A tree with a single child (using Identity).
type Tree1 a b =
  Tree Identity a b

-- | A single-child tree where nodes and leaves have the same type.
type Tree1' a =
  Tree1 a a

instance Eq1 f => Eq2 (Tree f) where
  liftEq2 f g (Tree a t1) (Tree b t2) =
    f a b &&
    liftEq2 f g t1 t2

instance Ord1 f => Ord2 (Tree f) where
  liftCompare2 f g (Tree a t1) (Tree b t2) =
    f a b <>
    liftCompare2 f g t1 t2

instance Show1 f => Show2 (Tree f) where
  liftShowsPrec2 spA slA spB slB d (Tree a t) =
    showsBinaryWith
      spA
      (liftShowsPrec2 spA slA spB slB)
      "Tree"
      d
      a
      t

instance (Eq a, Eq1 f) => Eq1 (Tree f a) where
  liftEq =
    liftEq2 (==)

instance (Ord a, Ord1 f) => Ord1 (Tree f a) where
  liftCompare =
    liftCompare2 compare

instance (Show a, Show1 f) => Show1 (Tree f a) where
  liftShowsPrec =
    liftShowsPrec2 showsPrec showList

-- |
--
-- >>> Tree "a" (TreeForest []) == Tree "a" (TreeForest []) :: Bool
-- True
--
-- >>> Tree "a" (TreeForest [Left 1]) == Tree "a" (TreeForest [Left 2]) :: Bool
-- False
instance (Eq a, Eq1 f, Eq b) => Eq (Tree f a b) where
  (==) =
    liftEq (==)

-- |
--
-- >>> compare (Tree "a" (TreeForest [])) (Tree "b" (TreeForest [])) :: Ordering
-- LT
instance (Ord a, Ord1 f, Ord b) => Ord (Tree f a b) where
  compare =
    liftCompare compare

-- |
--
-- >>> show (Tree "root" (TreeForest [Left "x"])) :: String
-- "Tree \"root\" (TreeForest [Left \"x\"])"
instance (Show a, Show1 f, Show b) => Show (Tree f a b) where
  showsPrec =
    liftShowsPrec showsPrec shows

-- | Map over both node labels and leaf values.
--
-- >>> bimap (+1) (*10) (Tree 5 (TreeForest [Left 3])) :: TreeList Int Int
-- Tree 6 (TreeForest [Left 30])
instance Functor f => Bifunctor (Tree f) where
  bimap f g (Tree a t) =
    Tree (f a) (bimap f g t)

-- | Map over leaf values, keeping node labels unchanged.
--
-- >>> fmap (*2) (Tree "root" (TreeForest [Left 5, Left 10])) :: TreeList String Int
-- Tree "root" (TreeForest [Left 10,Left 20])
instance Functor f => Functor (Tree f a) where
  fmap =
    bimap id

-- NOTE: Biapply and Biapplicative instances are NOT possible for Tree
-- because the underlying TreeForest cannot support these instances due to
-- the recursive Either-based structure. See the note on TreeForest above.

-- | Apply with semigroup combination of node labels.
--
-- >>> Tree "a" (TreeForest [Left (+1)]) <.> Tree "b" (TreeForest [Left 5]) :: TreeList String Int
-- Tree "ab" (TreeForest [Left 6])
instance (Apply f, Semigroup a) => Apply (Tree f a) where
  Tree a1 t1 <.> Tree a2 t2 =
    Tree (a1 <> a2) (t1 <.> t2)

-- |
--
-- >>> Tree "a" (TreeForest []) <*> Tree "b" (TreeForest []) :: TreeList String String
-- Tree "ab" (TreeForest [])
--
-- >>> Tree "a" (TreeForest [Left Prelude.reverse]) <*> Tree "b" (TreeForest [Left "xyz"]) :: TreeList String String
-- Tree "ab" (TreeForest [Left "zyx"])
--
-- >>> Tree "a" (TreeForest [Left Prelude.reverse]) <*> Tree "b" (TreeForest [Left "xyz", makeChild "c" [Left "pqr"], makeChild "d" [Left "mno"]]) :: TreeList String String
-- Tree "ab" (TreeForest [Left "zyx",Right (Tree "c" (TreeForest [Left "rqp"])),Right (Tree "d" (TreeForest [Left "onm"]))])
instance (Applicative f, Monoid a) => Applicative (Tree f a) where
  pure =
    Tree mempty . pure
  Tree a1 t1 <*> Tree a2 t2 =
    Tree (a1 <> a2) (t1 <*> t2)

-- | Fold over both node labels and leaf values.
--
-- >>> import Data.Monoid (Sum(..))
-- >>> bifoldMap Sum Sum (Tree 1 (TreeForest [Left 5, makeChild 3 []])) :: Sum Int
-- Sum {getSum = 9}
instance Foldable f => Bifoldable (Tree f) where
  bifoldMap f g (Tree a t) =
    f a <> bifoldMap f g t

-- | Bifoldable1 for trees with non-empty forests.
--
-- >>> import Data.Semigroup (Sum(..))
-- >>> bifoldMap1 Sum Sum (Tree 1 (TreeForest (Left 5 :| []))) :: Sum Int
-- Sum {getSum = 6}
instance Foldable1 f => Bifoldable1 (Tree f) where
  bifoldMap1 f g (Tree a t) =
    f a <> bifoldMap1 f g t

-- | Fold over leaf values only (ignores node labels).
--
-- >>> foldMap (:[]) (Tree "root" (TreeForest [Left "a", Left "b"])) :: [String]
-- ["a","b"]
instance Foldable f => Foldable (Tree f a) where
  foldMap f (Tree _ t) =
    foldMap f t

-- | Foldable1 for trees with non-empty forests.
--
-- >>> import Data.Semigroup (Sum(..))
-- >>> foldMap1 Sum (Tree "root" (TreeForest (Left 5 :| [Left 10]))) :: Sum Int
-- Sum {getSum = 15}
instance Foldable1 f => Foldable1 (Tree f a) where
  foldMap1 f (Tree _ t) =
    foldMap1 f t

-- | Traverse both node labels and leaf values with effects.
--
-- >>> bitraverse Just Just (Tree "a" (TreeForest [Left "b"])) :: Maybe (TreeList String String)
-- Just (Tree "a" (TreeForest [Left "b"]))
instance Traversable f => Bitraversable (Tree f) where
  bitraverse f g (Tree a t) =
    Tree <$> f a <*> bitraverse f g t

-- | Bitraversable1 for trees with non-empty forests.
--
-- >>> import Data.Functor.Identity
-- >>> bitraverse1 Identity Identity (Tree "a" (TreeForest (Left "b" :| []))) :: Identity (Tree NonEmpty String String)
-- Identity (Tree "a" (TreeForest (Left "b" :| [])))
instance Traversable1 f => Bitraversable1 (Tree f) where
  bitraverse1 f g (Tree a t) =
    Tree <$> f a <.> bitraverse1 f g t

-- | Traverse leaf values with effects.
--
-- >>> traverse Just (Tree "root" (TreeForest [Left "x"])) :: Maybe (TreeList String String)
-- Just (Tree "root" (TreeForest [Left "x"]))
instance Traversable f => Traversable (Tree f a) where
  traverse f (Tree a t) =
    Tree a <$> traverse f t

-- | Traversable1 for trees with non-empty forests.
--
-- >>> import Data.Functor.Identity
-- >>> traverse1 Identity (Tree "root" (TreeForest (Left "x" :| []))) :: Identity (Tree NonEmpty String String)
-- Identity (Tree "root" (TreeForest (Left "x" :| [])))
instance Traversable1 f => Traversable1 (Tree f a) where
  traverse1 f (Tree a t) =
    Tree a <$> traverse1 f t

-- | Plated instance for recursive tree traversal.
--
-- >>> transform (\(Tree a f) -> Tree (a <> "!") f) (Tree "a" (TreeForest [makeChild "b" []])) :: TreeList String String
-- Tree "a!" (TreeForest [Right (Tree "b!" (TreeForest []))])
instance Traversable f => Plated (Tree f a b) where
  plate f (Tree a (TreeForest x)) =
    Tree a . TreeForest <$> traverse
      ( either
          (pure . Left)
          (fmap Right . f)
      ) x

-- | Lens into a tree's forest.
--
-- >>> view treeForest' (Tree "a" (TreeForest [Left "x"])) :: TreeForest [] String String
-- TreeForest [Left "x"]
--
-- >>> set treeForest' (TreeForest [Left "y"]) (Tree "a" (TreeForest [Left "x"])) :: TreeList String String
-- Tree "a" (TreeForest [Left "y"])
treeForest' ::
  Lens
    (Tree f a b)
    (Tree f' a b')
    (TreeForest f a b)
    (TreeForest f' a b')
treeForest' f (Tree a t) =
  fmap (Tree a) (f t)

-- | Traversal over the immediate children and leaves of a tree.
--
-- >>> toListOf treeSubForest (Tree "root" (TreeForest [Left "x", makeChild "c" []])) :: [Either String (TreeList String String)]
-- [Left "x",Right (Tree "c" (TreeForest []))]
treeSubForest ::
  Traversable f =>
  Traversal
    (Tree f a b)
    (Tree f a b')
    (Either b (Tree f a b))
    (Either b' (Tree f a b'))
treeSubForest =
  treeForest' . _Wrapped . traverse

-- | Traversal over all immediate leaves in a tree (not recursive).
--
-- >>> toListOf treeLeaves (Tree "root" (TreeForest [Left "x", Left "y"])) :: [String]
-- ["x","y"]
--
-- >>> over treeLeaves (*2) (Tree "root" (TreeForest [Left 5, Left 10])) :: TreeList String Int
-- Tree "root" (TreeForest [Left 10,Left 20])
treeLeaves ::
  Traversable f =>
  Traversal'
    (Tree f a b)
    b
treeLeaves =
  treeSubForest . _Left

-- | Traversal over the immediate child trees.
--
-- >>> lengthOf treeForestChildren (Tree "root" (TreeForest [Left "x", makeChild "c1" [], makeChild "c2" []]))
-- 2
treeForestChildren ::
  Traversable f =>
  Traversal'
    (Tree f a b)
    (Tree f a b)
treeForestChildren =
  treeSubForest . _Right

-- | Read-only access to a tree structure via a 'Getter'.
-- This is the most general class, allowing types to expose a tree view
-- without necessarily allowing modification.
class GetTree x f a b | x -> f a b where
  -- | Extract the tree from a structure.
  getTree ::
    Getter x (Tree f a b)
  {-# INLINE getTreeLabel #-}
  -- | Extract just the tree's label.
  getTreeLabel ::
    Getter x a
  getTreeLabel =
    getTree . getTreeLabel

-- |
--
-- >>> view getTreeLabel (Tree "a" (TreeForest []))
-- "a"
--
-- >>> view getTree (Tree "a" (TreeForest [])) == Tree "a" (TreeForest [])
-- True
instance GetTree (Tree f a b) f a b where
  getTree =
    to id
  {-# INLINE getTreeLabel #-}
  getTreeLabel =
    to (\(Tree a _) -> a)

-- | Read-write access to a tree structure via a 'Lens'.
-- This extends 'GetTree' to allow modification.
class GetTree x f a b => HasTree x f a b | x -> f a b where
  -- | Access the tree with read-write capability.
  tree ::
    Lens' x (Tree f a b)
  {-# INLINE treeLabel #-}
  -- | Access the tree's label with read-write capability.
  treeLabel ::
    Lens' x a
  treeLabel =
    tree . treeLabel

-- |
--
-- >>> view treeLabel (Tree "a" (TreeForest []))
-- "a"
--
-- >>> set treeLabel "z" (Tree "a" (TreeForest [])) :: TreeList String String
-- Tree "z" (TreeForest [])
--
-- >>> view treeForest (Tree "a" (TreeForest []))
-- TreeForest []
--
-- >>> view treeForest (Tree "b" (TreeForest [Left "xyz", makeChild "c" [Left "pqr"], makeChild "d" [Left "mno"]]))
-- TreeForest [Left "xyz",Right (Tree "c" (TreeForest [Left "pqr"])),Right (Tree "d" (TreeForest [Left "mno"]))]
instance HasTree (Tree f a b) f a b where
  tree =
    id
  {-# INLINE treeLabel #-}
  treeLabel f (Tree a t) =
    fmap (`Tree` t) (f a)

-- |
--
-- >>> view getTreeForest (Tree "root" (TreeForest [Left "a"])) :: TreeForest [] String String
-- TreeForest [Left "a"]
instance GetTreeForest (Tree f a b) f a b where
  getTreeForest =
    to (\(Tree _ forest) -> forest)

-- |
--
-- >>> view treeForest (Tree "root" (TreeForest [Left "a"])) :: TreeForest [] String String
-- TreeForest [Left "a"]
--
-- >>> set treeForest (TreeForest [Left "b"]) (Tree "root" (TreeForest [Left "a"])) :: TreeList String String
-- Tree "root" (TreeForest [Left "b"])
instance HasTreeForest (Tree f a b) f a b where
  treeForest =
    treeForest'

-- | Construction-only access to a tree via a 'Review'.
-- This allows building a type from a tree.
class ReviewTree x f a b | x -> f a b where
  -- | Construct a value from a tree.
  reviewTree ::
    Review x (Tree f a b)

-- |
--
-- >>> review reviewTree (Tree "hello" (TreeForest [])) :: TreeList String String
-- Tree "hello" (TreeForest [])
instance ReviewTree (Tree f a b) f a b where
  reviewTree =
    unto id

-- | Full prism access to a tree structure.
-- This extends 'ReviewTree' to allow both construction and pattern matching.
class ReviewTree x f a b => AsTree x f a b | x -> f a b where
  -- | Access the tree as a prism (construct or pattern match).
  _Tree ::
    Prism' x (Tree f a b)

-- |
--
-- >>> preview _Tree (Tree "test" (TreeForest [])) :: Maybe (TreeList String String)
-- Just (Tree "test" (TreeForest []))
instance AsTree (Tree f a b) f a b where
  _Tree =
    id

-- |
--
-- >>> dfs (Tree 1 (TreeForest []))
-- Left 1 :| []
--
-- >>> dfs (Tree 1 (TreeForest [Left 2]))
-- Left 1 :| [Right 2]
--
-- >>> dfs (Tree 1 (TreeForest [Left 2, makeChild 3 []]))
-- Left 1 :| [Right 2,Left 3]
--
-- >>> dfs (Tree 1 (TreeForest [Left 2, makeChild 3 [], Left 4]))
-- Left 1 :| [Right 2,Left 3,Right 4]
--
-- >>> dfs (Tree 1 (TreeForest [Left 2, makeChild 3 [Left 5], Left 4]))
-- Left 1 :| [Right 2,Left 3,Right 5,Right 4]
--
-- >>> dfs (Tree 1 (TreeForest [Left 2, makeChild 3 [Left 5], Left 4, makeChild 6 []]))
-- Left 1 :| [Right 2,Left 3,Right 5,Right 4,Left 6]
dfs ::
  Foldable f =>
  Tree f a b ->
  NonEmpty (Either a b)
dfs (Tree a t) =
  Left a :| foldMap (either (\b -> [Right b]) (toList . dfs)) (view _Wrapped t)

-- |
--
-- >>> bfs (Tree 1 (TreeForest []))
-- Left 1 :| []
--
-- >>> bfs (Tree 1 (TreeForest [Left 2]))
-- Left 1 :| [Right 2]
--
-- >>> bfs (Tree 1 (TreeForest [Left 2, makeChild 3 []]))
-- Left 1 :| [Right 2,Left 3]
--
-- >>> bfs (Tree 1 (TreeForest [Left 2, makeChild 3 [], Left 4]))
-- Left 1 :| [Right 2,Right 4,Left 3]
--
-- >>> bfs (Tree 1 (TreeForest [Left 2, makeChild 3 [Left 5], Left 4]))
-- Left 1 :| [Right 2,Right 4,Left 3,Right 5]
--
-- >>> bfs (Tree 1 (TreeForest [Left 2, makeChild 3 [Left 5], Left 4, makeChild 6 []]))
-- Left 1 :| [Right 2,Right 4,Left 3,Right 5,Left 6]
bfs ::
  Foldable f =>
  Tree f a b
  -> NonEmpty (Either a b)
bfs root =
  let go (Tree a t :| rest) =
        let (leaves, c) =
              foldMap (either (\b -> ([Right b], [])) (\tr -> ([], [tr]))) (view _Wrapped t)
        in  case nonEmpty (rest <> c) of
              Nothing -> Left a :| leaves
              Just q  -> Left a :| (leaves <> toList (go q))
  in  go (root :| [])

-- |
--
-- >>> makeTree 1 []
-- Tree 1 (TreeForest [])
--
-- >>> makeTree 1 [Left "a"]
-- Tree 1 (TreeForest [Left "a"])
--
-- >>> makeTree 1 [Left "a", Right (makeTree 2 [])]
-- Tree 1 (TreeForest [Left "a",Right (Tree 2 (TreeForest []))])
makeTree ::
  a
  -> f (Either b (Tree f a b))
  -> Tree f a b
makeTree a t =
  Tree a (TreeForest t)

-- |
--
-- >>> makeChild 1 []
-- Right (Tree 1 (TreeForest []))
--
-- >>> makeChild 1 [Left "a"]
-- Right (Tree 1 (TreeForest [Left "a"]))
--
-- >>> makeChild 1 [Left "a", makeChild 2 []]
-- Right (Tree 1 (TreeForest [Left "a",Right (Tree 2 (TreeForest []))]))
makeChild ::
  a
  -> f (Either b (Tree f a b))
  -> Either x (Tree f a b)
makeChild a t =
  Right (makeTree a t)

-- |
--
-- >>> makeLeaves 1 []
-- Tree 1 (TreeForest [])
--
-- >>> makeLeaves 1 [makeLeaves 2 []]
-- Tree 1 (TreeForest [Left (Tree 2 (TreeForest []))])
--
-- >>> makeLeaves 1 [makeLeaves 2 [makeChild 3 []]]
-- Tree 1 (TreeForest [Left (Tree 2 (TreeForest [Left (Right (Tree 3 (TreeForest [])))]))])
makeLeaves ::
  Functor f =>
  a
  -> f b
  -> Tree f a b
makeLeaves a bs =
  Tree a (TreeForest (Left <$> bs))

-- |
--
-- >>> makeChildren 1 []
-- Tree 1 (TreeForest [])
--
-- >>> makeChildren 1 [makeChildren 2 []]
-- Tree 1 (TreeForest [Right (Tree 2 (TreeForest []))])
--
-- >>> makeChildren 1 [makeChildren 2 [], makeLeaves 3 []]
-- Tree 1 (TreeForest [Right (Tree 2 (TreeForest [])),Right (Tree 3 (TreeForest []))])
makeChildren ::
  Functor f =>
  a
  -> f (Tree f a b)
  -> Tree f a b
makeChildren a cs =
  Tree a (TreeForest (Right <$> cs))

-- |
--
-- >>> singleton "root" :: TreeList String String
-- Tree "root" (TreeForest [])
singleton ::
  Monoid (f (Either b (Tree f a b))) =>
  a
  -> Tree f a b
singleton a =
  Tree a (TreeForest mempty)

-- |
--
-- >>> unfoldTree (\n -> (n, if n < 3 then [Left (n * 10), Right (n + 1)] else [])) 1 :: TreeList Int Int
-- Tree 1 (TreeForest [Left 10,Right (Tree 2 (TreeForest [Left 20,Right (Tree 3 (TreeForest []))]))])
unfoldTree ::
  Functor f =>
  (a -> (a, f (Either b a)))
  -> a
  -> Tree f a b
unfoldTree f seed =
  let (label, forest) = f seed
  in  Tree label (TreeForest (fmap (fmap (unfoldTree f)) forest))

-- |
--
-- >>> unfoldTreeM (\n -> pure (n, if n < 3 then [Left (n * 10), Right (n + 1)] else [])) 1 :: Maybe (TreeList Int Int)
-- Just (Tree 1 (TreeForest [Left 10,Right (Tree 2 (TreeForest [Left 20,Right (Tree 3 (TreeForest []))]))]))
unfoldTreeM ::
  (Monad m, Traversable f) =>
  (a -> m (a, f (Either b a)))
  -> a
  -> m (Tree f a b)
unfoldTreeM f seed = do
  (label, forest) <- f seed
  forest' <- traverse (either (pure . Left) (fmap Right . unfoldTreeM f)) forest
  pure (Tree label (TreeForest forest'))

-- |
--
-- >>> pruneLeaves (> 2) (Tree 1 (TreeForest [Left 1, Left 3, makeChild 2 [Left 2, Left 4]])) :: TreeList Int Int
-- Tree 1 (TreeForest [Left 3,Right (Tree 2 (TreeForest [Left 4]))])
pruneLeaves ::
  (Traversable f, Applicative f, Monoid (f (Either b (Tree f a b)))) =>
  (b -> Bool)
  -> Tree f a b
  -> Tree f a b
pruneLeaves p (Tree a (TreeForest forest)) =
  Tree a (TreeForest (foldMap go forest))
  where
    go (Left b) = if p b then pure (Left b) else mempty
    go (Right t) = pure (Right (pruneLeaves p t))

-- |
--
-- >>> countNodes (Tree 1 (TreeForest [Left 2, makeChild 3 [Left 4], makeChild 5 []])) :: Int
-- 3
countNodes ::
  Foldable f =>
  Tree f a b
  -> Int
countNodes (Tree _ (TreeForest forest)) =
  1 + getSum (foldMap (either (const (Sum 0)) (Sum . countNodes)) forest)

-- |
--
-- >>> countLeaves (Tree 1 (TreeForest [Left 2, makeChild 3 [Left 4], makeChild 5 []])) :: Int
-- 2
countLeaves ::
  Foldable f =>
  Tree f a b
  -> Int
countLeaves (Tree _ (TreeForest forest)) =
  getSum (foldMap (either (const (Sum 1)) (Sum . countLeaves)) forest)

-- | Group tree elements by their depth level.
--
-- >>> Data.PolyTree.levels (Tree 1 (TreeForest [Left 2, makeChild 3 [], Left 7])) :: [[Either Int Int]]
-- [[Left 1],[Right 2],[Right 7],[Left 3]]
levels ::
  Foldable f =>
  Tree f a b
  -> [[Either a b]]
levels t =
  case bfs t of
    root :| rest -> groupByLevel [root] rest
  where
    groupByLevel acc [] = [acc]
    groupByLevel acc xs =
      let (level, remaining) = splitAt (length acc) xs
      in  if null level
            then [acc]
            else acc : groupByLevel level remaining

-- |
--
-- >>> view baseTree (Tree 1 (TreeForest []))
-- Node {rootLabel = 1, subForest = []}
--
-- >>> view baseTree (Tree 1 (TreeForest [Left 2, makeChild 3 [Left 5], Left 4, makeChild 6 []]))
-- Node {rootLabel = 1, subForest = [Node {rootLabel = 2, subForest = []},Node {rootLabel = 3, subForest = [Node {rootLabel = 5, subForest = []}]},Node {rootLabel = 4, subForest = []},Node {rootLabel = 6, subForest = []}]}
--
-- >>> review baseTree (Tree.Node 1 [])
-- Tree 1 (TreeForest [])
--
-- >>> review baseTree (Tree.Node 1 [Tree.Node 2 [],Tree.Node 3 [Tree.Node 5 []],Tree.Node 4 [],Tree.Node 6 []])
-- Tree 1 (TreeForest [Left 2,Right (Tree 3 (TreeForest [Left 5])),Left 4,Left 6])
baseTree ::
  Iso' (TreeList' a) (Tree.Tree a)
baseTree =
  iso
    (
      let go (Tree a t) =
            Tree.Node a (fmap (either pure go) (view _Wrapped t))
      in  go)
    (
      let perNode (Tree.Node a []) =
            Left a
          perNode tr@(Tree.Node _ (_:_)) =
            Right tr
          go (Tree.Node a t) =
            Tree a (TreeForest (fmap (fmap go . perNode) t))
      in  go)