"""The itree module contains functions implementing the immutable tree (ITree)
data abstractions for CS61A at UC Berkeley.

ITrees are (immutable) trees, a common class of data structures in computer
science which organizes data into a hierarchical structure.

One might notice that none of the doctests for either ADT actually "show" an
ITree. This is because this ADT is meant to be treated as a black box: the
tests do not assume any specific implementation for ITrees.
"""


# General Immutable Trees Abstraction


def make_itree(datum, children=()):
    """Create an ITree with datum and children.

    >>> print(itree_str(make_itree(1)))
    1
    >>> print(itree_str(make_itree(1,
    ...                            (make_itree(2),))))
    1
    |
    +-2
    >>> print(itree_str(make_itree(1,
    ...                            (make_itree(2,
    ...                                        (make_itree(3),
    ...                                         make_itree(4))),
    ...                             make_itree(5,
    ...                                        (make_itree(6),
    ...                                         make_itree(7)))))))
    1
    |
    +-2
    | |
    | +-3
    | |
    | +-4
    |
    +-5
      |
      +-6
      |
      +-7
    """
    return (datum, children)


def itree_datum(t):
    """Return the datum of ITree t.

    >>> test_tree = make_itree(4, (make_itree(11),
    ...                            make_itree(23)))
    >>> itree_datum(test_tree)
    4
    """
    return t[0]


def itree_children(t):
    """Return the children of ITree t.
    >>> test_tree = make_itree(4, (make_itree(11),
    ...                            make_itree(23)))
    >>> itree_datum(itree_children(test_tree)[0])
    11
    """
    return t[1]


# General Immutable Trees Utility Functions


def itree_items(t):
    """Returns a tuple of all items in the ITree t.
    
    >>> test_tree = make_itree(1,
    ...                        (make_itree(2,
    ...                                    (make_itree(3),
    ...                                     make_itree(4))),
    ...                         make_itree(5,
    ...                                    (make_itree(6),
    ...                                     make_itree(7)))))
    >>> print(itree_str(test_tree))
    1
    |
    +-2
    | |
    | +-3
    | |
    | +-4
    |
    +-5
      |
      +-6
      |
      +-7
    >>> itree_items(test_tree)
    (1, 2, 3, 4, 5, 6, 7)
    """
    from functools import reduce
    from operator import add
    return (itree_datum(t),) + \
           reduce(add, map(itree_items, itree_children(t)), ())

           
def itree_size(t):
    """Returns the number of items in an ITree t.
    
    >>> test_tree = make_itree(1,
    ...                        (make_itree(2,
    ...                                    (make_itree(3),
    ...                                     make_itree(4))),
    ...                         make_itree(5,
    ...                                    (make_itree(6),
    ...                                     make_itree(7)))))
    >>> itree_size(test_tree)
    7
    """
    return len(itree_items(t))


def is_leaf(t):
    """Returns True if the ITree t is a leaf (or has no children).

    >>> is_leaf(make_itree(2))
    True
    >>> is_leaf(make_itree(1, (make_itree(2), make_itree(3))))
    False
    """
    return len(itree_children(t)) == 0
    

def itree_map(func, t):
    """Returns a new ITree with the results of applying func to each of the
    items in t.

    >>> test_tree = make_itree(1,
    ...                        (make_itree(2,
    ...                                    (make_itree(3),
    ...                                     make_itree(4))),
    ...                         make_itree(5,
    ...                                    (make_itree(6),
    ...                                     make_itree(7)))))
    >>> print(itree_str(test_tree))
    1
    |
    +-2
    | |
    | +-3
    | |
    | +-4
    |
    +-5
      |
      +-6
      |
      +-7
    >>> print(itree_str(itree_map(lambda x: x * x, test_tree)))
    1
    |
    +-4
    | |
    | +-9
    | |
    | +-16
    |
    +-25
      |
      +-36
      |
      +-49
    """
    return make_itree(func(itree_datum(t)),
                      tuple(map(lambda child: itree_map(func, child),
                                itree_children(t))))


def itree_str(t, depth=0, more_children=False):
    """Returns a string representation of the ITree t.

    Arguments:
    t               --  the tree to be printed
    depth           --  the depth (in case it's a subtree)
    more_children   --  do we need to put a character at the front of each line
                        to connect more children after this tree?

    >>> print(itree_str(make_itree(1,
    ...                            (make_itree(2),
    ...                             make_itree(3, (make_itree(4),)),
    ...                             make_itree(5, (make_itree(6),
    ...                                            make_itree(7)))))))
    1
    |
    +-2
    |
    +-3
    | |
    | +-4
    |
    +-5
      |
      +-6
      |
      +-7
    >>> print(itree_str(make_itree(1,
    ...                      (make_itree(2,
    ...                            (make_itree(3,
    ...                                  (make_itree(4),
    ...                                   make_itree(5))),
    ...                             make_itree(6,
    ...                                  (make_itree(7),
    ...                                   make_itree(8))))),
    ...                       make_itree(9,
    ...                            (make_itree(10,
    ...                                  (make_itree(11),
    ...                                   make_itree(12))),
    ...                             make_itree(13,
    ...                                  (make_itree(14),
    ...                                   make_itree(15)))))))))
    1
    |
    +-2
    | |
    | +-3
    | | |
    | | +-4
    | | |
    | | +-5
    | |
    | +-6
    |   |
    |   +-7
    |   |
    |   +-8
    |
    +-9
      |
      +-10
      | |
      | +-11
      | |
      | +-12
      |
      +-13
        |
        +-14
        |
        +-15
    """
    result = repr(itree_datum(t))
    last_child = len(itree_children(t)) - 1
    for child_n in range(len(itree_children(t))):
        result += "\n|\n+-"
        result += itree_str(itree_children(t)[child_n],
                            depth + 1,
                            child_n != last_child)

    # Add front padding if we're part of a bigger tree we're converting to a
    # string.  Perhaps not the most elegant code.
    if depth > 0:
        if more_children:
            leader = "| "
        else:
            leader = "  "
        result = result.replace("\n", "\n" + leader)
    return result