Mapytex/mapytex/calculus/random/core/random_tree.py

from ...core.tree import MutableTree, Tree
from ...core.MO import moify
from .grammar import extract_letters, eval_words
from .generate import random_generator
from .str2 import rdstr2


class RandomTree(MutableTree):
    """MutableTree that accept {a} syntax for random generation

    :example:
    >>> t = RandomTree()
    >>> type(t)
    <class 'mapytex.calculus.random.core.random_tree.RandomTree'>
    """

    @classmethod
    def from_str(cls, expression):
        """Initiate a random tree from a string that need to be parsed

        :exemple:
        >>> t = RandomTree.from_str("{b}*x+{c}")
        >>> print(t)
        +
         > *
         | > {b}
         | > x
         > {c}
        >>> t = RandomTree.from_str("{a}*({b}*x+{c})")
        >>> print(t)
        *
         > {a}
         > +
         | > *
         | | > {b}
         | | > x
         | > {c}
        """
        str_2_mut_tree = rdstr2(cls.sink)
        return str_2_mut_tree(expression)

    @property
    def random_leaves(self) -> list[str]:
        """Get list of random leaves

        :example:
        >>> from .leaf import RdLeaf
        >>> random_tree = RandomTree("+", RdLeaf("a"), RdLeaf("a*k"))
        >>> random_tree.random_leaves
        ['a', 'a*k']
        >>> random_tree = RandomTree("+", RdLeaf("a"), 2)
        >>> random_tree.random_leaves
        ['a']
        """
        rd_leafs = []
        for leaf in self.get_leafs():
            try:
                leaf.rdleaf
            except AttributeError:
                pass
            else:
                rd_leafs.append(leaf.name)
        return rd_leafs

    @property
    def random_value(self) -> set[str]:
        """Get set of random values to generate

        :example:
        >>> from .leaf import RdLeaf
        >>> random_tree = RandomTree("+", RdLeaf("a"), RdLeaf("a*k"))
        >>> random_tree.random_value == {'a', 'k'}
        True
        >>> random_tree = RandomTree("+", RdLeaf("a"), 2)
        >>> random_tree.random_value
        {'a'}
        """
        return extract_letters(self.random_leaves)

    def eval_random_leaves(self, leaves_value: dict[str, int]):
        """Given random leaves value get the tree

        :example:
        >>> from .leaf import RdLeaf
        >>> rd_t = RandomTree("+", RdLeaf("a"), RdLeaf("a*k"))
        >>> leaves_values = {'a': 2, 'a*k': 6}
        >>> t = rd_t.eval_random_leaves(leaves_values)
        >>> type(t)
        <class 'mapytex.calculus.core.tree.Tree'>
        >>> print(t)
        +
         > 2
         > 6
        """

        def replace(leaf):
            try:
                return leaf.replace(leaves_value)
            except AttributeError:
                return leaf

        return self.map_on_leaf(replace).map_on_leaf(moify)

    def generate(
        self, conditions: list[str] = [], global_config: dict = {}, configs: dict = {}
    ) -> Tree:
        """Generate a random version of self

        :param conditions: list of conditions
        :param config: global configuration for generated values
        :param configs: specific configuration for each generated values

        """
        generated_values = random_generator(
            self.random_value, conditions, global_config, configs
        )
        leaves = eval_words(self.random_leaves, generated_values)
        return self.eval_random_leaves(leaves)