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

"""
Tree class

"""

from .tree_tools import (to_nested_parenthesis,
        infix_str_concatenate,
        postfix_concatenate,
        show_tree,
        )
from .coroutine import *
from .str2 import str2
from .operator import OPERATORS

__all__ = ["Tree", "MutableTree"]


class Tree(object):

    """
    Binary tree

    This is the chosen structure to manage calculus in PyMath.

    The class is not mutable to preserve its integrity.

    """

    def __init__(self, node, left_value, right_value):
        """
        Initiate a tree with tuple (node, (left value, right value))

        :example:

        >>> t = Tree("+", 1, 2)
        >>> t.node
        '+'
        >>> t.left_value
        1
        >>> t.right_value
        2
        """
        if node is None or right_value is None:
            raise TypeError("Tree can't have an empty node or an empty right leaf")

        self.node = node

        self.left_value = left_value
        self.right_value = right_value

    @classmethod
    def from_str(cls, expression):
        """ Initiate a tree from an string expression

        :example:

        >>> t = MutableTree.fromStr("2+3*4")
        >>> print(t)
        +
         > 2
         > *
         | > 3
         | > 4
        >>> t = MutableTree.fromStr("(2+3)*4")
        >>> print(t)
        *
         > +
         | > 2
         | > 3
         > 4
        """
        t = MutableTree.from_str(expression)
        return cls.from_mutable_tree(t)

    @classmethod
    def from_nested_parenthesis(cls, nested_parenthesis):
        """
        Initiate recursively a tree with tuple (node, (left value, right value))

        :example:

        >>> nested_par = ("+", (1, 2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.node
        '+'
        >>> t.left_value
        1
        >>> t.right_value
        2
        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.node
        '+'
        >>> type(t.left_value)
        <class 'mapytex.calculus.core.tree.Tree'>
        >>> t.right_value
        2

        """
        try:
            nested_len = len(nested_parenthesis)
            num_len = len(nested_parenthesis[1])
        except TypeError:
            raise ValueError("Nested parenthesis are not composed of lists")

        if nested_len != 2 and \
            num_len != 2:
                raise ValueError("Nested parenthesis don't have right shape")

        node = nested_parenthesis[0]

        try:
            left_value = Tree.from_nested_parenthesis(nested_parenthesis[1][0])
        except ValueError:
            left_value = nested_parenthesis[1][0]
        try:
            right_value = Tree.from_nested_parenthesis(nested_parenthesis[1][1])
        except ValueError:
            right_value = nested_parenthesis[1][1]

        return cls(node, left_value, right_value)

    @classmethod
    def from_mutable_tree(cls, mutable_tree):
        """ Initial a Tree from a MutableTree

        :example:
        >>> t = MutableTree("*", 1, 2)
        >>> print(Tree.from_mutable_tree(t))
        *
         > 1
         > 2
        >>> t1 = MutableTree("*", 1, 2)
        >>> t2 = MutableTree("*", t1, 3)
        >>> print(Tree.from_mutable_tree(t2))
        *
         > *
         | > 1
         | > 2
         > 3
        >>> t = MutableTree("*", 1)
        >>> print(t)
        *
         > 1
         > None
        >>> Tree.from_mutable_tree(t)
        Traceback (most recent call last):
            ...
        TypeError: Tree can't have empty node or leaf

        """
        node = mutable_tree.node
        left_value = mutable_tree.left_value
        right_value = mutable_tree.right_value

        if node is None or \
                left_value is None or \
                right_value is None:
                    raise TypeError("Tree can't have empty node or leaf")

        try:
            l_value = Tree.from_mutable_tree(left_value)
        except AttributeError:
            l_value = left_value

        try:
            r_value = Tree.from_mutable_tree(right_value)
        except AttributeError:
            r_value = right_value

        return cls(node, l_value, r_value)

    def map_on_leaf(self, function):
        """ Map on leafs a function

        :param function: take leaf value returns other value
        :returns: Tree with calculated leaf

        :example:

        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> print(t)
        +
         > *
         | > 3
         | > 4
         > 2
        >>> print(t.map_on_leaf(lambda x:2*x))
        +
         > *
         | > 6
         | > 8
         > 4

        """
        try:
            left_applied = self.left_value.map_on_leaf(function)
        except AttributeError:
            left_applied = function(self.left_value)

        try:
            right_applied = self.right_value.map_on_leaf(function)
        except AttributeError:
            right_applied = function(self.right_value)

        return Tree(self.node, left_applied, right_applied)

    def apply_on_last_level(self, function):
        """ Apply the function on last level of the tree before leaf

        :param function: (op, a, a) -> b
        :returns: b if it is a 1 level Tree, Tree otherwise


        :example:

        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> print(t)
        +
         > *
         | > 3
         | > 4
         > 2
        >>> from .tree_tools import infix_str_concatenate
        >>> tt = t.apply_on_last_level(lambda *x: infix_str_concatenate(*x))
        >>> print(tt)
        +
         > 3 * 4
         > 2
        >>> from .evaluate import compute
        >>> tt = t.apply_on_last_level(compute)
        >>> print(tt)
        +
         > 12
         > 2
        >>> ttt = tt.apply_on_last_level(compute)
        >>> print(ttt)
        14
        """
        left_is_leaf = 0
        right_is_leaf = 0
        try:
            left_applied = self.left_value.apply_on_last_level(function)
        except AttributeError:
            left_applied = self.left_value
            left_is_leaf = 1

        try:
            right_applied = self.right_value.apply_on_last_level(function)
        except AttributeError:
            right_applied = self.right_value
            right_is_leaf = 1

        if left_is_leaf and right_is_leaf:
            return function(self.node, self.left_value, self.right_value)
        else:
            return Tree(self.node, left_applied, right_applied)

    def apply(self, function):
        """ Apply the function on every node of the tree

        :param function: (op, a, a) -> b
        :returns: b

        :example:

        >>> def to_nested(op, left, right):
        ...     return (op, (left, right))
        >>> nested_par = ("+", (1, 2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.apply(to_nested)
        ('+', (1, 2))

        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.apply(to_nested)
        ('+', (('*', (3, 4)), 2))

        >>> from .evaluate import compute
        >>> t.apply(compute)
        14

        """
        try:
            left_value = self.left_value.apply(function)
        except AttributeError:
            left_value = self.left_value

        try:
            right_value = self.right_value.apply(function)
        except AttributeError:
            right_value = self.right_value

        return function(self.node, left_value, right_value)

    def to_nested_parenthesis(self):
        """ Transform the Tree into its nested parenthesis description
        :example:

        >>> nested_par = ("+", (1, 2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.to_nested_parenthesis()
        ('+', (1, 2))
        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.to_nested_parenthesis()
        ('+', (('*', (3, 4)), 2))
        """
        return self.apply(to_nested_parenthesis)

    def to_postfix(self):
        """ Transform the Tree into postfix notation

        :example:

        >>> nested_par = ("+", (1, 2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.to_postfix()
        [1, 2, '+']
        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.to_postfix()
        [3, 4, '*', 2, '+']
        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    ("*", (5, 6)),
        ...     ))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> t.to_postfix()
        [3, 4, '*', 5, 6, '*', '+']
        """
        return self.apply(postfix_concatenate)

    def get_leafs(self, callback = lambda x:x):
        """ Generator which yield all the leaf value of the tree.
        Callback act on every leaf.

        :param callback: function on leaf

        :example:

        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> [l for l in t.get_leafs()]
        [3, 4, 2]
        >>> {type(l) for l in t.get_leafs()}
        {<class 'int'>}
        """
        try:
            yield from self.left_value.get_leafs(callback)
        except AttributeError:
            yield callback(self.left_value)
        try:
            yield from self.right_value.get_leafs(callback)
        except AttributeError:
            yield callback(self.right_value)

    def get_nodes(self, callback = lambda x:x):
        """ Generator which yield all nodes of the tree.
        Callback act on every nodes.

        :param callback: function on node

        :example:

        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> [l for l in t.get_nodes()]
        ['+', '*']
        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    ("*", (3, 4))
        ...    ))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> [l for l in t.get_nodes()]
        ['+', '*', '*']
        """
        yield self.node
        try:
            yield from self.left_value.get_nodes(callback)
        except AttributeError:
            pass
        try:
            yield from self.right_value.get_nodes(callback)
        except AttributeError:
            pass

    def __str__(self):
        """ Overload str method

        :example:

        >>> nested_par = ("+", (1, 2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> print(t)
        +
         > 1
         > 2
        >>> nested_par = ("+", (
        ...    ("*", (3, 4)),
        ...    2))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> print(t)
        +
         > *
         | > 3
         | > 4
         > 2
        >>> nested_par = ("+", (
        ...    ("*", (1, 2)),
        ...    ("*", (3, 4)),
        ...    ))
        >>> t = Tree.from_nested_parenthesis(nested_par)
        >>> print(t)
        +
         > *
         | > 1
         | > 2
         > *
         | > 3
         | > 4
        """
        return self.apply(show_tree)


class MutableTree(Tree):

    """
    Mutable Tree.

    It is used to build a new tree before fixing it into a Tree

    """
    def __init__(self,
            node = None,
            left_value = None,
            right_value = None):
        """
        Initiate a tree with potentialy None values

        :Example:

        >>> t = MutableTree()
        >>> t.node, t.left_value, t.right_value
        (None, None, None)
        >>> t = MutableTree('*', 1)
        >>> t.node, t.left_value, t.right_value
        ('*', 1, None)
        """

        self.node = node

        self.left_value = left_value
        self.right_value = right_value

    @classmethod
    def fromStr(cls, expression):
        """ Initiate the MutableTree

        :example:
        >>> t = MutableTree.fromStr("2+3*4")
        >>> print(t)
        +
         > 2
         > *
         | > 3
         | > 4
        >>> t = MutableTree.fromStr("(2+3)*4")
        >>> print(t)
        *
         > +
         | > 2
         | > 3
         > 4
        >>> t = MutableTree.fromStr("4*(-2+3)")
        >>> print(t)
        *
         > 4
         > +
         | > -
         | | > None
         | | > 2
         | > 3
        """
        str2mutTree = str2(cls.sink)
        return str2mutTree(expression)

    @classmethod
    @coroutine
    def sink(cls):
        """ Sink which build a build a mutabletree

        Returns the built tree when STOOOP is thrown

        :example:
        >>> sink = MutableTree.sink()
        >>> for i in ["1", "+", "2", "*", "3"]:
        ...    sink.send(i)
        >>> a = sink.throw(STOOOP)
        >>> print(a)
        +
         > 1
         > *
         | > 2
         | > 3
        >>> sink = MutableTree.sink()
        >>> for i in ["1", "*", "2", "+", "3"]:
        ...    sink.send(i)
        >>> a = sink.throw(STOOOP)
        >>> print(a)
        +
         > *
         | > 1
         | > 2
         > 3
        >>> sink = MutableTree.sink()
        >>> for i in ["-", "1", "+", "2"]:
        ...    sink.send(i)
        >>> a = sink.throw(STOOOP)
        >>> print(a)
        +
         > -
         | > None
         | > 1
         > 2

        """
        ans = cls()
        try:
            while True:
                c = yield
                if c is not None:
                    if c in OPERATORS.keys():
                        try:
                            ans.set_node(c)
                        except ValueError:
                            if OPERATORS[c]["precedence"] > OPERATORS[ans.node]["precedence"]:
                                ans.append_bot(c)
                            else:
                                ans.append_top(c)
                    else:
                        ans.append(c)
        except STOOOP:
            yield ans

    def set_node(self, value):
        """ Set the node value.
        Once it has been changed it can't be changed again.

        :example:
        >>> t = MutableTree()
        >>> t.node
        >>> t.set_node("*")
        >>> t.node
        '*'
        >>> t.set_node("+")
        Traceback (most recent call last):
            ...
        ValueError: The node of the tree is already set
        """
        if self.node is None:
            self.node = value
        else:
            raise ValueError("The node of the tree is already set")

    def set_left_value(self, value):
        """ Set the left value.
        Once it has been changed it can't be changed again.

        :example:
        >>> t = MutableTree()
        >>> t.left_value
        >>> t.set_left_value(1)
        >>> t.left_value
        1
        >>> t.set_left_value(2)
        Traceback (most recent call last):
            ...
        ValueError: The left branch is full, use set_right_value
        """
        if self.left_value is None:
            self.left_value = value
        else:
            try:
                self.left_value.append(value)
            except AttributeError:
                raise ValueError("The left branch is full, use set_right_value")

    def set_right_value(self, value):
        """ Set the right value.
        Once it has been changed it can't be changed again.

        :example:
        >>> t = MutableTree()
        >>> t.right_value
        >>> t.set_right_value(2)
        >>> t.right_value
        2
        >>> t.set_right_value(3)
        Traceback (most recent call last):
            ...
        ValueError: The right branch is full, use append_top or append_bot
        >>> # potentielle source de problèmes??
        >>> t = MutableTree()
        >>> t.set_right_value([1])
        >>> t.right_value
        [1]
        >>> t.set_right_value(3)
        """
        if self.right_value is None:
            self.right_value = value
        else:
            try:
                self.right_value.append(value)
            except AttributeError:
                raise ValueError("The right branch is full, use append_top or append_bot")

    def append(self, value):
        """ Append the value at the bottom of the tree.

        In order to enable operator with arity 1, left value can be set only
        before the node is set, assuming that those operator are placed before
        operand.

        It tries to add it on left branch first.
        If it fails, the value is appened on right branch.

        :example:
        >>> t = MutableTree()
        >>> t.append(1)
        >>> t.set_node("*")
        >>> t.append(2)
        >>> print(t)
        *
         > 1
         > 2
        >>> t.append(3)
        Traceback (most recent call last):
            ...
        ValueError: The right branch is full, use append_top or append_bot

        >>> t1 = MutableTree()
        >>> t1.append(1)
        >>> t1.set_node("*")
        >>> t2 = MutableTree()
        >>> t1.append(t2)
        >>> t1.append(2)
        >>> t2.set_node("+")
        >>> t1.append(3)
        >>> print(t1)
        *
         > 1
         > +
         | > 2
         | > 3

        >>> t1 = MutableTree()
        >>> t1.set_node("-")
        >>> t1.append(1)
        >>> print(t1)
        -
         > None
         > 1
        """
        if self.node is None:
            try:
                self.set_left_value(value)
            except ValueError:
                self.set_right_value(value)
        else:
            self.set_right_value(value)

    def append_top(self, value):
        """ Append node into the tree at the top

        :example:

        >>> t = MutableTree("*", 1, 2)
        >>> t.append_top("+")
        >>> print(t)
        +
         > *
         | > 1
         | > 2
         > None
        """

        self_cp = MutableTree()
        self_cp.set_node(self.node)
        self_cp.set_left_value(self.left_value)
        self_cp.set_right_value(self.right_value)

        self.left_value, self.node, self.right_value = self_cp, value, None

    def append_bot(self, value):
        """ Append node into the tree at the bottom right_value

        :example:

        >>> t = MutableTree("*", 1, 2)
        >>> t.append_bot("+")
        >>> print(t)
        *
         > 1
         > +
         | > 2
         | > None
        """

        rv = self.right_value
        nright_value = MutableTree()
        nright_value.set_node(value)
        nright_value.set_left_value(rv)

        self.right_value = nright_value

# -----------------------------
# Reglages pour 'vim'
# vim:set autoindent expandtab tabstop=4 shiftwidth=4:
# cursor: 16 del