#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2017 lafrite <lafrite@Poivre>
#
# Distributed under terms of the MIT license.

"""
Tools to extract random leafs, random variables, generate random values and
fill new trees

Flow
----

Tree with RdLeaf
|
| Extract rdLeaf
|
List of leafs to generate
|
| extract_rv
|
List random variables to generate
|
| Generate
|
Dictionnary of generated random variables
|
| Compute leafs
|
Dictionnary of computed leafs
|
| Replace
|
Tree with RdLeaf replaced by generated values

:example:

>>> from ..tree import Tree
>>> rd_t = Tree("+", RdLeaf("a"), RdLeaf("a*k"))
>>> print(rd_t)
+
 > {a}
 > {a*k}
>>> extract_rdleaf(rd_t)
['a', 'a*k']


"""

from .leaf import RdLeaf

def extract_rdleaf(tree):
    """ Extract rdLeaf in a Tree

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

def extract_rv(leafs):
    """ Extract the set of random values from the leaf list

    :param leafs: list of leafs
    :return: set of random values

    :example:
    >>> leafs = ["a", "a*k"]
    >>> list(extract_rv(leafs))
    ['a', 'k']
    """
    rd_values = set()
    for leaf in leafs:
        for c in leaf:
            if c.isalpha():
                rd_values.add(c)
    return rd_values

def compute_leafs(leafs, generated_values):
    """ Compute leafs from generated random values

    :param generated_values: Dictionnary of name:generated value
    :param leafs: list of leafs
    :return: Dictionnary of evaluated leafs from generated values

    :example:
    >>> leafs = ["a", "a*k"]
    >>> generated_values = {"a":2, "k":3}
    >>> compute_leafs(leafs, generated_values)
    {'a': 2, 'a*k': 6}
    """
    return {leaf: eval(leaf, generated_values) for leaf in leafs}

def replace_rdleaf(tree, computed_leafs):
    """ Replace RdLeaf by the corresponding computed value

    >>> from ..tree import Tree
    >>> rd_t = Tree("+", RdLeaf("a"), RdLeaf("a*k"))
    >>> computed_leafs = {'a': 2, 'a*k': 6}
    >>> print(replace_rdleaf(rd_t, computed_leafs))
    +
     > 2
     > 6
    """
    def replace(leaf):
        try:
            return leaf.replace(computed_leafs)
        except AttributeError:
            return leaf
    return tree.map_on_leaf(replace)