"""CS 61A HW 06
Name:
Login:
TA:
Section:
"""


from itree import *


#### Q1
def depth(t, v):
    """Returns the depth of value v in ITree t if v is contained in t.  If v is
    not in t, return None.

    >>> test_tree = 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))))))))
    >>> depth(test_tree, 1)
    0
    >>> depth(test_tree, 42) # Returns None
    >>> depth(test_tree, 6)
    2
    >>> depth(test_tree, 15)
    3
    """
    "*** YOUR CODE HERE ***"


#### Q2
def prune_leaves(t, vals):
    """Return a modified copy of t with all leaves that have a datum that
    appears in vals removed.  Return None if the entire tree is pruned away.

    >>> prune_leaves(make_itree(2), (2,)) # Returns None
    >>> test_tree1 = make_itree(1, (make_itree(2), make_itree(3)))
    >>> test_tree2 = make_itree(1,
    ...                         (make_itree(2,
    ...                                     (make_itree(3),
    ...                                      make_itree(2))),
    ...                          make_itree(3)))
    >>> print(itree_str(test_tree1))
    1
    |
    +-2
    |
    +-3
    >>> print(itree_str(prune_leaves(test_tree1, (2, 3))))
    1
    >>> print(itree_str(prune_leaves(test_tree1, (2,))))
    1
    |
    +-3
    >>> print(itree_str(test_tree2))
    1
    |
    +-2
    | |
    | +-3
    | |
    | +-2
    |
    +-3
    >>> print(itree_str(prune_leaves(test_tree2, (2, 3))))
    1
    |
    +-2
    >>> print(itree_str(prune_leaves(test_tree2, (2,))))
    1
    |
    +-2
    | |
    | +-3
    |
    +-3
    """
    "*** YOUR CODE HERE ***"


def sprout_leaves(t, vals):
    """Sprout new leaves containing the data in vals at each leaf in the
    original tree t and return the resulting tree.

    >>> print(itree_str(sprout_leaves(make_itree(1), (1, 2, 3))))
    1
    |
    +-1
    |
    +-2
    |
    +-3
    >>> test_tree1 = make_itree(1, (make_itree(2), make_itree(3)))
    >>> print(itree_str(test_tree1))
    1
    |
    +-2
    |
    +-3
    >>> print(itree_str(sprout_leaves(test_tree1, (4, 5, 6))))
    1
    |
    +-2
    | |
    | +-4
    | |
    | +-5
    | |
    | +-6
    |
    +-3
      |
      +-4
      |
      +-5
      |
      +-6
    >>> test_tree2 = make_itree(1,
    ...                         (make_itree(2,
    ...                                     (make_itree(3),
    ...                                      make_itree(2))),
    ...                          make_itree(3)))
    >>> print(itree_str(test_tree2))
    1
    |
    +-2
    | |
    | +-3
    | |
    | +-2
    |
    +-3
    >>> print(itree_str(sprout_leaves(test_tree2, (4, 5, 6))))
    1
    |
    +-2
    | |
    | +-3
    | | |
    | | +-4
    | | |
    | | +-5
    | | |
    | | +-6
    | |
    | +-2
    |   |
    |   +-4
    |   |
    |   +-5
    |   |
    |   +-6
    |
    +-3
      |
      +-4
      |
      +-5
      |
      +-6
    """
    "*** YOUR CODE HERE ***"


#### Q4
def same_shape(t1, t2):
    """Return True if t1 is indentical in shape to t2.

    >>> test_tree1 = make_itree(1, (make_itree(2), make_itree(3)))
    >>> test_tree2 = make_itree(4, (make_itree(5), make_itree(6)))
    >>> test_tree3 = make_itree(1,
    ...                         (make_itree(2,
    ...                                     (make_itree(3),)),))
    >>> test_tree4 = make_itree(4,
    ...                         (make_itree(5,
    ...                                     (make_itree(6),)),))
    >>> same_shape(test_tree1, test_tree2)
    True
    >>> same_shape(test_tree3, test_tree4)
    True
    >>> same_shape(test_tree2, test_tree4)
    False
    """
    "*** YOUR CODE HERE ***"

#### Q5
def make_mobile(left_branch, right_branch):
    """Construct a mobile with the given left and right branches."""
    return (left_branch, right_branch)

def mobile_left(m):
    """Return the left branch of the mobile m."""
    return m[0]

def mobile_right(m):
    """Return the right branch of the mobile m."""
    return m[1]

def make_branch(length, item):
    """Construct a branch for the mobile with a given length and an item
    hanging off that branch (either another mobile or a weight).
    """
    return (length, item)

def branch_len(b):
    """Return the length of the branch b."""
    return b[0]

def branch_item(b):
    """Return the item hanging from the branch b."""
    return b[1]

def is_weight(x):
    """Return True if x is not a mobile but a weight (a number)."""
    # Not the best function, but it'll do.
    return type(x) is not tuple

def total_weight(x):
    """Return the total weight of this mobile (either a full mobile or just a
    weight).

    >>> total_weight(4)
    4
    >>> total_weight(make_mobile(make_branch(4, 5),
    ...                          make_branch(4,
    ...                                      make_mobile(make_branch(3, 6),
    ...                                                  make_branch(2, 3)))))
    14
    """
    "*** YOUR CODE HERE ***"

def is_balanced(m):
    """Return true if the given mobile is balanced.

    >>> is_balanced(5) # weights are balanced
    True
    >>> is_balanced(make_mobile(make_branch(4, 3),
    ...                         make_branch(3,
    ...                                     make_mobile(make_branch(2, 2),
    ...                                                 make_branch(2, 2)))))
    True
    >>> is_balanced(make_mobile(make_branch(4, 3),
    ...                         make_branch(3, 5)))
    False
    >>> is_balanced(make_mobile(make_branch(4, 3),
    ...                         make_branch(3,
    ...                                     make_mobile(make_branch(2, 3),
    ...                                                 make_branch(3, 2)))))
    False
    """
    "*** YOUR CODE HERE ***"