"""
Schnyder's Algorithm for straight-line planar embeddings

A module for computing the (x,y) coordinates for a straight-line planar
embedding of any connected planar graph with at least three vertices.  Uses
Walter Schnyder's Algorithm from [Sch1990]_.

AUTHORS:

- Jonathan Bober, Emily Kirkman (2008-02-09) --  initial version
"""
# ****************************************************************************
#      Copyright (C) 2008 Jonathan Bober and Emily Kirkman
#
# Distributed  under  the  terms  of  the  GNU  General  Public  License (GPL)
#                         https://www.gnu.org/licenses/
# ****************************************************************************
from __future__ import absolute_import

from sage.sets.set import Set
from .all import DiGraph

def _triangulate(g, comb_emb):
    """
    Helper function to schnyder method for computing coordinates in the plane to
    plot a planar graph with no edge crossings.

    Given a connected graph g with at least 3 vertices and a planar combinatorial
    embedding comb_emb of g, modify g in place to form a graph whose faces are
    all triangles, and return the set of newly created edges. Also ``comb_emb``
    is updated in place.

    The simple way to triangulate a face is to just pick a vertex and draw
    an edge from that vertex to every other vertex in the face. Think that this
    might ultimately result in graphs that don't look very nice when we draw them
    so we have decided on a different strategy.  After handling special cases,
    we add an edge to connect every third vertex on each face.  If this edge is
    a repeat, we keep the first edge of the face and retry the process at the next
    edge in the face list.  (By removing the special cases we guarantee that this
    method will work on one of these attempts.)

    INPUT:

    - g -- the graph to triangulate
    - ``comb_emb`` -- a planar combinatorial embedding of g

    OUTPUT:

    A list of edges that are added to the graph (in place)

    EXAMPLES::

        sage: from sage.graphs.schnyder import _triangulate
        sage: g = Graph(graphs.CycleGraph(4))
        sage: g.is_planar(set_embedding=True)
        True
        sage: _triangulate(g, g._embedding)
        [(2, 0), (1, 3)]

        sage: g = graphs.PathGraph(3)
        sage: g.is_planar(set_embedding=True)
        True
        sage: new_edges = _triangulate(g, g._embedding)
        sage: [sorted(e) for e in new_edges]
        [[0, 2]]
    """
    # first make sure that the graph has at least 3 vertices, and that it is connected
    if g.order() < 3:
        raise ValueError("A Graph with less than 3 vertices doesn't have any triangulation.")
    if not g.is_connected():
        raise NotImplementedError("_triangulate() only knows how to handle connected graphs.")

    # At this point we know that the graph is connected, has at least 3
    # vertices. This is where the real work starts.

    faces = g.faces(comb_emb)
    # We start by finding all of the faces of this embedding.

    edges_added = []   # The list of edges that we add to the graph.
    # This will be returned at the end.

    for face in faces:
        new_face = []
        if len(face) < 3:
            raise RuntimeError('Triangulate method created face %s with < 3 edges.' % face)
        if len(face) == 3:
            continue  # This face is already triangulated
        elif len(face) == 4:  # In this special case just add diagonal edge to square
            u, v, w, x = (e[0] for e in face)
            if w == u or g.has_edge(w,u):
                u, v, w, x = v, w, x, u
            new_face = (w, u)
            comb_emb[w].insert(comb_emb[w].index(x), u)
            comb_emb[u].insert(comb_emb[u].index(v), w)
            g.add_edge(new_face)
            edges_added.append(new_face)
        else:
            N = len(face)
            i = 0
            while i < N - 1:
                new_edge = (face[i + 1][1], face[i][0])  # new_edge is from third vertex in face to first
                if g.has_edge(new_edge) or new_edge[0] == new_edge[1]:  # check for repeats
                    new_face.append(face[i])  # if repeated, keep first edge in face instead
                    if i == N - 2:   # if we are two from the end, found a triangle already
                        break
                    i += 1
                    continue

                g.add_edge(new_edge)
                edges_added.append(new_edge)
                comb_emb[new_edge[0]].insert(comb_emb[new_edge[0]].index((face + new_face)[i + 2][1]), new_edge[1])
                comb_emb[new_edge[1]].insert(comb_emb[new_edge[1]].index(face[i][1]), new_edge[0])
                new_face.append((new_edge[1], new_edge[0]))
                i += 2
            if i != N:
                new_face.append(face[-1])
            faces.append(new_face)

    return edges_added

def _normal_label(g, comb_emb, external_face):
    r"""
    Helper function to schnyder method for computing coordinates in
    the plane to plot a planar graph with no edge crossings.

    Constructs a normal labelling of a triangular graph g, given the
    planar combinatorial embedding of g and a designated external
    face.  Returns labels dictionary.  The normal label is constructed
    by first contracting the graph down to its external face, then
    expanding the graph back to the original while simultaneously
    adding angle labels.

    INPUT:

    - g -- the graph to find the normal labeling of (g must be triangulated)
    - ``comb_emb`` -- a planar combinatorial embedding of g
    - ``external_face`` -- the list of three edges in the external face of g

    OUTPUT:

    x -- tuple with entries

        x[0] = dict of dicts of normal labeling for each vertex of g and each
        adjacent neighbors u,v (u < v) of vertex:

        { vertex : { (u,v): angel_label } }

        x[1] = (v1,v2,v3) tuple of the three vertices of the external face.

    EXAMPLES::

        sage: from sage.graphs.schnyder import _triangulate, _normal_label, _realizer
        sage: g = Graph(graphs.CycleGraph(7))
        sage: g.is_planar(set_embedding=True)
        True
        sage: faces = g.faces(g._embedding)
        sage: _triangulate(g, g._embedding)
        [(2, 0), (4, 2), (6, 4), (1, 3), (6, 1), (3, 5), (4, 0), (6, 3)]
        sage: tn = _normal_label(g, g._embedding, faces[0])
        sage: _realizer(g, tn)
        ({0: [<sage.graphs.schnyder.TreeNode object at ...>]},
         (1, 0, 2))
    """
    contracted = []
    contractible = []

    labels = {}

    # For now we will not take the order of the outer face into account.
    # We will correct this in the end of this function.
    external_vertices = sorted([external_face[0][0],
                                external_face[1][0],
                                external_face[2][0]])
    v1, v2, v3 = external_vertices
    v1_neighbors = Set(g.neighbors(v1))

    neighbor_count = {}
    for v in g.vertices():
        neighbor_count[v] = 0
    for v in g.neighbors(v1):
        neighbor_count[v] = len(v1_neighbors.intersection(Set(g.neighbors(v))))

    for v in v1_neighbors:
        if v in [v1, v2, v3]:
            continue
        if neighbor_count[v] == 2:
            contractible.append(v)

    # contraction phase:

    while g.order() > 3:
        try:
            v = contractible.pop()
        except Exception:
            raise RuntimeError('Contractible list is empty but graph still has %d vertices.  (Expected 3.)' % g.order())

            break
        # going to contract v
        v_neighbors = Set(g.neighbors(v))
        contracted.append((v, v_neighbors,
                           v_neighbors - v1_neighbors - Set([v1])))
        g.delete_vertex(v)
        v1_neighbors -= Set([v])
        for w in v_neighbors - v1_neighbors - Set([v1]):
            # adding edge (v1, w)
            g.add_edge((v1, w))
        if g.order() == 3:
            break
        v1_neighbors += v_neighbors - Set([v1])
        contractible = []
        for w in g.neighbors(v1):
            if (len(v1_neighbors.intersection(Set(g.neighbors(w)))) == 2
                    and w not in [v1, v2, v3]):
                contractible.append(w)

    # expansion phase:

    v1, v2, v3 = g.vertices()  # always in sorted order

    labels[v1] = {(v2, v3): 1}
    labels[v2] = {(v1, v3): 2}
    labels[v3] = {(v1, v2): 3}

    while len(contracted):
        v, new_neighbors, neighbors_to_delete = contracted.pop()
        # going to add back vertex v
        labels[v] = {}

        for w in neighbors_to_delete:
            g.delete_edge((v1, w))

        if len(neighbors_to_delete) == 0:
            # we are adding v into the face new_neighbors
            w1, w2, w3 = sorted(new_neighbors)

            labels[v] = {(w1, w2): labels[w3].pop((w1, w2)),
                         (w2, w3): labels[w1].pop((w2, w3)),
                         (w1, w3): labels[w2].pop((w1, w3))}
            labels[w1][tuple(sorted((w2, v)))] = labels[v][(w2, w3)]
            labels[w1][tuple(sorted((w3, v)))] = labels[v][(w2, w3)]

            labels[w2][tuple(sorted((w1, v)))] = labels[v][(w1, w3)]
            labels[w2][tuple(sorted((w3, v)))] = labels[v][(w1, w3)]

            labels[w3][tuple(sorted((w1, v)))] = labels[v][(w1, w2)]
            labels[w3][tuple(sorted((w2, v)))] = labels[v][(w1, w2)]
        else:
            new_neighbors_set = Set(new_neighbors)
            angles_out_of_v1 = set()
            vertices_in_order = []
            l = []
            for angle in labels[v1].keys():
                if len(Set(angle).intersection(new_neighbors_set)) == 2:
                    angles_out_of_v1.add(angle)
                    l = l + list(angle)
            # find a unique element in l
            l.sort()
            i = 0
            while i < len(l):
                if l[i] == l[i + 1]:
                    i += 2
                else:
                    break

            angle_set = Set(angles_out_of_v1)

            vertices_in_order.append(l[i])
            while len(angles_out_of_v1) > 0:
                for angle in angles_out_of_v1:
                    if vertices_in_order[-1] in angle:
                        break
                if angle[0] == vertices_in_order[-1]:
                    vertices_in_order.append(angle[1])
                else:
                    vertices_in_order.append(angle[0])
                angles_out_of_v1.remove(angle)

            w = vertices_in_order

            # is w[0] a 2 or a 3?
            top_label = labels[w[0]][tuple(sorted((v1, w[1])))]
            if top_label == 3:
                bottom_label = 2
            else:
                bottom_label = 3
            i = 0
            while i < len(w) - 1:
                labels[v][tuple(sorted((w[i], w[i + 1])))] = 1
                labels[w[i]][tuple(sorted((w[i + 1], v)))] = top_label
                labels[w[i + 1]][tuple(sorted((w[i], v)))] = bottom_label
                i += 1

            labels[v][tuple(sorted((v1, w[0])))] = bottom_label
            labels[v][tuple(sorted((v1, w[-1])))] = top_label

            labels[w[0]][tuple(sorted((v1, v)))] = top_label
            labels[w[-1]][tuple(sorted((v1, v)))] = bottom_label
            labels[v1][tuple(sorted((w[0], v)))] = 1
            labels[v1][tuple(sorted((w[-1], v)))] = 1

            # delete all the extra labels

            for angle in angle_set:
                labels[v1].pop(angle)

            labels[w[0]].pop(tuple(sorted((v1, w[1]))))
            labels[w[-1]].pop(tuple(sorted((v1, w[-2]))))

            i = 1
            while i < len(w) - 1:
                labels[w[i]].pop(tuple(sorted((v1, w[i + 1]))))
                labels[w[i]].pop(tuple(sorted((v1, w[i - 1]))))
                i += 1

        for w in new_neighbors:
            g.add_edge((v, w))

    # Up to this point we did not take the order of the external face into
    # account. Since the combinatorial embedding of a triangulation is unique up
    # to the choice of the outer face and reflection, this might lead to a
    # reflection of the Schnyder drawing resulting from this labeling which is
    # not conformal with comb_emb any longer. Therefore, we might have to swap
    # the labels 1 and 2.
    if (v1, v2) in external_face:
        for u in labels:
            for v, w in labels[u]:
                if labels[u][v, w] == 1:
                    labels[u][v, w] = 2
                elif labels[u][v, w] == 2:
                    labels[u][v, w] = 1
        v1, v2 = v2, v1

    return labels, (v1, v2, v3)

def _realizer(g, x, example=False):
    """
    Given a triangulated graph g and a normal labeling constructs the
    realizer and returns a dictionary of three trees determined by the
    realizer, each spanning all interior vertices and rooted at one of
    the three external vertices.

    A realizer is a directed graph with edge labels that span all interior
    vertices from each external vertex.  It is determined by giving direction
    to the edges that have the same angle label on both sides at a vertex.
    (Thus the direction actually points to the parent in the tree.)  The
    edge label is set as whatever the matching angle label is.  Then from
    any interior vertex, following the directed edges by label will
    give a path to each of the three external vertices.

    INPUT:

    - g -- the graph to compute the realizer of
    - x -- tuple with entries

        x[0] = dict of dicts representing a normal labeling of g.  For
        each vertex of g and each adjacent neighbors u,v (u < v) of
        vertex:  { vertex : { (u,v): angle_label } }

        x[1] = (v1, v2, v3) tuple of the three external vertices (also
        the roots of each tree)

    OUTPUT:

    - x -- tuple with entries

        x[0] = dict of lists of TreeNodes:

        { root_vertex : [ list of all TreeNodes under root_vertex ] }

        x[1] = (v1,v2,v3) tuple of the three external vertices (also the
        roots of each tree)

    EXAMPLES::

        sage: from sage.graphs.schnyder import _triangulate, _normal_label, _realizer
        sage: g = Graph(graphs.CycleGraph(7))
        sage: g.is_planar(set_embedding=True)
        True
        sage: faces = g.faces(g._embedding)
        sage: _triangulate(g, g._embedding)
        [(2, 0), (4, 2), (6, 4), (1, 3), (6, 1), (3, 5), (4, 0), (6, 3)]
        sage: tn = _normal_label(g, g._embedding, faces[0])
        sage: _realizer(g, tn)
        ({0: [<sage.graphs.schnyder.TreeNode object at ...>]},
         (1, 0, 2))

    """
    normal_labeling, (v1, v2, v3) = x
    realizer = DiGraph()

    tree_nodes = {}
    for v in g:
        tree_nodes[v] = [TreeNode(label=v, children=[]),
                         TreeNode(label=v, children=[]),
                         TreeNode(label=v, children=[])]

    for v in g:
        ones = []
        twos = []
        threes = []
        l = [ones, twos, threes]
        for angle, value in normal_labeling[v].items():
            l[value - 1] += list(angle)

        ones.sort()
        twos.sort()
        threes.sort()

        i = 0
        while i < len(ones) - 1:
            if ones[i] == ones[i + 1]:
                realizer.add_edge((ones[i], v), label=1)
                tree_nodes[v][0].append_child(tree_nodes[ones[i]][0])
                i += 1
            i += 1
        i = 0
        while i < len(twos) - 1:
            if twos[i] == twos[i + 1]:
                realizer.add_edge((twos[i], v), label=2)
                tree_nodes[v][1].append_child(tree_nodes[twos[i]][1])
                i += 1
            i += 1
        i = 0
        while i < len(threes) - 1:
            if threes[i] == threes[i + 1]:
                realizer.add_edge((threes[i], v), label=3)
                tree_nodes[v][2].append_child(tree_nodes[threes[i]][2])
                i += 1
            i += 1

    _compute_coordinates(realizer, (tree_nodes, (v1, v2, v3)))

    if example:
        realizer.show(talk=True, edge_labels=True)

    return tree_nodes, (v1, v2, v3)

def _compute_coordinates(g, x):
    r"""
    Given a triangulated graph g with a dict of trees given by the
    realizer and tuple of the external vertices, we compute the
    coordinates of a planar geometric embedding in the grid.

    The coordinates will be set to the ``_pos`` attribute of g.

    INPUT:

    - g -- the graph to compute the coordinates of
    - x -- tuple with entries

        x[0] = dict of tree nodes for the three trees with each external
        vertex as root:

        { root_vertex : [ list of all TreeNodes under root_vertex ] }

        x[1] = (v1, v2, v3) tuple of the three external vertices (also
        the roots of each tree)

    EXAMPLES::

        sage: from sage.graphs.schnyder import _triangulate, _normal_label, _realizer, _compute_coordinates
        sage: g = Graph(graphs.CycleGraph(7))
        sage: g.is_planar(set_embedding=True)
        True
        sage: faces = g.faces(g._embedding)
        sage: _triangulate(g, g._embedding)
        [(2, 0), (4, 2), (6, 4), (1, 3), (6, 1), (3, 5), (4, 0), (6, 3)]
        sage: tn = _normal_label(g, g._embedding, faces[0])
        sage: r = _realizer(g, tn)
        sage: _compute_coordinates(g,r)
        sage: g.get_pos()
        {0: [0, 5], 1: [5, 1], 2: [1, 0], 3: [4, 1], 4: [1, 1], 5: [2, 2], 6: [1, 4]}
    """

    tree_nodes, (v1, v2, v3) = x
    # find the roots of each tree:
    t1, t2, t3 = tree_nodes[v1][0], tree_nodes[v2][1], tree_nodes[v3][2]

    # Compute the number of descendants and depth of each node in
    # each tree.
    t1.compute_number_of_descendants()
    t2.compute_number_of_descendants()
    t3.compute_number_of_descendants()

    t1.compute_depth_of_self_and_children()
    t2.compute_depth_of_self_and_children()
    t3.compute_depth_of_self_and_children()

    coordinates = {}  # the dict to pass to g.set_pos()

    # Setting coordinates for external vertices
    coordinates[t1.label] = [g.order() - 2, 1]
    coordinates[t2.label] = [0, g.order() - 2]
    coordinates[t3.label] = [1, 0]

    for v in g.vertices():
        if v not in [t1.label, t2.label, t3.label]:
            # Computing coordinates for v
            r = list((0, 0, 0))

            for i in [0, 1, 2]:
                # Computing size of region i:

                # Tracing up tree (i + 1) % 3
                p = tree_nodes[v][(i + 1) % 3]
                while p is not None:
                    q = tree_nodes[p.label][i].number_of_descendants
                    # Adding number of descendants from Tree i nodes with
                    # labels on path up tree (i + 1) % 3
                    r[i] += q
                    p = p.parent

                # Tracing up tree (i - 1) % 3
                p = tree_nodes[v][(i - 1) % 3]
                while p is not None:
                    q = tree_nodes[p.label][i].number_of_descendants
                    # Adding number of descendants from Tree i nodes with
                    # labels on path up tree (i - 1) % 3
                    r[i] += q
                    p = p.parent

                q = tree_nodes[v][i].number_of_descendants
                # Subtracting
                r[i] -= q

                # Subtracting
                q = tree_nodes[v][(i - 1) % 3].depth
                r[i] -= q

            if sum(r) != g.order() - 1:
                raise RuntimeError("Computing coordinates failed: vertex %s's coordinates sum to %s.  Expected %s" % (v, sum(r), g.order() - 1))

            coordinates[v] = r[:-1]

    g.set_pos(coordinates)  # Setting _pos attribute to store coordinates

class TreeNode(object):
    """
    A class to represent each node in the trees used by ``_realizer`` and
    ``_compute_coordinates`` when finding a planar geometric embedding in
    the grid.

    Each tree node is doubly linked to its parent and children.

    INPUT:

    - ``parent`` -- the parent TreeNode of ``self``
    - ``children`` -- a list of TreeNode children of ``self``
    - ``label`` -- the associated realizer vertex label

    EXAMPLES::

        sage: from sage.graphs.schnyder import TreeNode
        sage: tn = TreeNode(label=5)
        sage: tn2 = TreeNode(label=2,parent=tn)
        sage: tn3 = TreeNode(label=3)
        sage: tn.append_child(tn3)
        sage: tn.compute_number_of_descendants()
        2
        sage: tn.number_of_descendants
        2
        sage: tn3.number_of_descendants
        1
        sage: tn.compute_depth_of_self_and_children()
        sage: tn3.depth
        2
    """
    def __init__(self, parent=None, children=None, label=None):
        """
        INPUT:

        - ``parent`` -- the parent TreeNode of ``self``
        - ``children`` -- a list of TreeNode children of ``self``
        - ``label`` -- the associated realizer vertex label

        EXAMPLES::

            sage: from sage.graphs.schnyder import TreeNode
            sage: tn = TreeNode(label=5)
            sage: tn2 = TreeNode(label=2,parent=tn)
            sage: tn3 = TreeNode(label=3)
            sage: tn.append_child(tn3)
            sage: tn.compute_number_of_descendants()
            2
            sage: tn.number_of_descendants
            2
            sage: tn3.number_of_descendants
            1
            sage: tn.compute_depth_of_self_and_children()
            sage: tn3.depth
            2
        """
        if children is None:
            children = []
        self.parent = parent
        self.children = children
        self.label = label
        self.number_of_descendants = 1

    def compute_number_of_descendants(self):
        """
        Computes the number of descendants of self and all descendants.

        For each TreeNode, sets result as attribute self.number_of_descendants

        EXAMPLES::

            sage: from sage.graphs.schnyder import TreeNode
            sage: tn = TreeNode(label=5)
            sage: tn2 = TreeNode(label=2,parent=tn)
            sage: tn3 = TreeNode(label=3)
            sage: tn.append_child(tn3)
            sage: tn.compute_number_of_descendants()
            2
            sage: tn.number_of_descendants
            2
            sage: tn3.number_of_descendants
            1
            sage: tn.compute_depth_of_self_and_children()
            sage: tn3.depth
            2

        """
        n = 1
        for child in self.children:
            n += child.compute_number_of_descendants()
        self.number_of_descendants = n
        return n

    def compute_depth_of_self_and_children(self):
        """
        Computes the depth of self and all descendants.

        For each TreeNode, sets result as attribute self.depth

        EXAMPLES::

            sage: from sage.graphs.schnyder import TreeNode
            sage: tn = TreeNode(label=5)
            sage: tn2 = TreeNode(label=2,parent=tn)
            sage: tn3 = TreeNode(label=3)
            sage: tn.append_child(tn3)
            sage: tn.compute_number_of_descendants()
            2
            sage: tn.number_of_descendants
            2
            sage: tn3.number_of_descendants
            1
            sage: tn.compute_depth_of_self_and_children()
            sage: tn3.depth
            2
        """
        if self.parent is None:
            self.depth = 1
        else:
            self.depth = self.parent.depth + 1
        for child in self.children:
            child.compute_depth_of_self_and_children()

    def append_child(self, child):
        """
        Add a child to list of children.

        EXAMPLES::

            sage: from sage.graphs.schnyder import TreeNode
            sage: tn = TreeNode(label=5)
            sage: tn2 = TreeNode(label=2,parent=tn)
            sage: tn3 = TreeNode(label=3)
            sage: tn.append_child(tn3)
            sage: tn.compute_number_of_descendants()
            2
            sage: tn.number_of_descendants
            2
            sage: tn3.number_of_descendants
            1
            sage: tn.compute_depth_of_self_and_children()
            sage: tn3.depth
            2
        """
        if child in self.children:
            return
        self.children.append(child)
        child.parent = self

def minimal_schnyder_wood(graph, root_edge=None, minimal=True, check=True):
    """
    Return the minimal Schnyder wood of a planar rooted triangulation.

    INPUT:

    - graph -- a planar triangulation, given by a graph with an embedding.

    - root_edge -- a pair of vertices (default is from ``-1`` to ``-2``)
      The third boundary vertex is then determined using the orientation and
      will be labelled ``-3``.

    - minimal -- boolean (default ``True``), whether to return a
      minimal or a maximal Schnyder wood.

    - check -- boolean (default ``True``), whether to check if the input
      is a planar triangulation

    OUTPUT:

    A planar graph, with edges oriented and colored. The three outer
    edges of the initial graph are removed. For the three outer vertices the
    list of the neighbors stored in the combinatorial embedding is in the order
    of the incident edges between the two incident (and removed) outer edges,
    and not a cyclic shift of it.

    The algorithm is taken from [Bre2000]_ (section 4.2).

    EXAMPLES::

        sage: from sage.graphs.schnyder import minimal_schnyder_wood
        sage: g = Graph([(0,-1),(0,-2),(0,-3),(-1,-2),(-2,-3),
        ....:  (-3,-1)], format='list_of_edges')
        sage: g.set_embedding({-1:[-2,0,-3],-2:[-3,0,-1],
        ....:  -3:[-1,0,-2],0:[-1,-2,-3]})
        sage: newg = minimal_schnyder_wood(g)
        sage: newg.edges()
        [(0, -3, 'red'), (0, -2, 'blue'), (0, -1, 'green')]
        sage: newg.plot(color_by_label={'red':'red','blue':'blue',
        ....:  'green':'green',None:'black'})
        Graphics object consisting of 8 graphics primitives

    A larger example::

        sage: g = Graph([(0,-1),(0,2),(0,1),(0,-3),(-1,-3),(-1,2),
        ....: (-1,-2),(1,2),(1,-3),(2,-2),(1,-2),(-2,-3)], format='list_of_edges')
        sage: g.set_embedding({-1:[-2,2,0,-3],-2:[-3,1,2,-1],
        ....: -3:[-1,0,1,-2],0:[-1,2,1,-3],1:[-2,-3,0,2],2:[-1,-2,1,0]})
        sage: newg = minimal_schnyder_wood(g)
        sage: sorted(newg.edges(), key=lambda e:(str(e[0]),str(e[1])))
        [(0, -1, 'green'),
         (0, -3, 'red'),
         (0, 2, 'blue'),
         (1, -2, 'blue'),
         (1, -3, 'red'),
         (1, 0, 'green'),
         (2, -1, 'green'),
         (2, -2, 'blue'),
         (2, 1, 'red')]
        sage: newg2 = minimal_schnyder_wood(g, minimal=False)
        sage: sorted(newg2.edges(), key=lambda e:(str(e[0]),str(e[1])))
        [(0, -1, 'green'),
         (0, -3, 'red'),
         (0, 1, 'blue'),
         (1, -2, 'blue'),
         (1, -3, 'red'),
         (1, 2, 'green'),
         (2, -1, 'green'),
         (2, -2, 'blue'),
         (2, 0, 'red')]

    TESTS::

        sage: minimal_schnyder_wood(graphs.RandomTriangulation(5))
        Digraph on 5 vertices
        sage: minimal_schnyder_wood(graphs.CompleteGraph(5))
        Traceback (most recent call last):
        ...
        ValueError: not a planar graph
        sage: minimal_schnyder_wood(graphs.WheelGraph(5))
        Traceback (most recent call last):
        ...
        ValueError: not a triangulation
        sage: minimal_schnyder_wood(graphs.OctahedralGraph(),root_edge=(0,5))
        Traceback (most recent call last):
        ...
        ValueError: not a valid root edge
    """
    if root_edge is None:
        a = -1
        b = -2
    else:
        a, b = root_edge

    if check:
        if not graph.is_planar():
            raise ValueError('not a planar graph')
        if not all(len(u) == 3 for u in graph.faces()):
            raise ValueError('not a triangulation')
        if not(a in graph.neighbors(b)):
            raise ValueError('not a valid root edge')

    new_g = DiGraph()
    emb = graph.get_embedding()

    # finding the third outer vertex c
    emb_b = emb[b]
    idx_a = emb_b.index(a)
    c = emb_b[(idx_a + 1) % len(emb_b)]

    # initialisation
    for i in emb[c]:
        if i != a and i != b:
            new_g.add_edge((i, -3, 'red'))

    path = list(emb[c])
    idxa = path.index(a)
    path = path[idxa:] + path[:idxa]
    neighbors_in_path = {i: len([u for u in graph.neighbors(i) if u in path])
                         for i in graph}
    removable_nodes = [u for u in path if neighbors_in_path[u] == 2 and
                       u != a and u != b]

    # iterated path shortening
    while len(path) > 2:
        if minimal:
            v = removable_nodes[-1]   # node to be removed from path
        else:
            v = removable_nodes[0]   # node to be removed from path
        idx_v = path.index(v)
        left = path[idx_v - 1]
        new_g.add_edge((v, left, 'green'))
        right = path[idx_v + 1]
        new_g.add_edge((v, right, 'blue'))
        neighbors_v = emb[v]
        idx_left = neighbors_v.index(left)
        neighbors_v = neighbors_v[idx_left:] + neighbors_v[:idx_left]
        idx_right = neighbors_v.index(right)
        inside = neighbors_v[1:idx_right]
        new_g.add_edges([(w, v, 'red') for w in inside])
        path = path[:idx_v] + inside + path[idx_v + 1:]
        # updating the table of neighbors_in_path
        for w in inside:
            for x in graph.neighbors(w):
                neighbors_in_path[x] += 1
        for x in graph.neighbors(v):
            neighbors_in_path[x] -= 1
        # updating removable nodes
        removable_nodes = [u for u in path if neighbors_in_path[u] == 2 and
                           u != a and u != b]

    def relabel(w):
        return -3 if w == c else w

    emb = {relabel(v): [relabel(u) for u in emb[v]] for v in graph}
    for u, v, w in (a, b, -3), (b, -3, a), (-3, a, b):
        idx = emb[u].index(v)
        if idx == 0:
            emb[u] = emb[u][1:-1]
        else:
            emb[u] = emb[u][idx+1:] + emb[u][:idx-1]

    new_g.set_embedding(emb)
    return new_g


"""First-order exploration of pandas Series object"""

import numpy as np              #python's array proccesing / linear algebra library
import pandas as  pd            #data processing / stats library
import matplotlib.pyplot as plt #data visualization
import matplotlib.dates as dates
import csv
import datetime
from py_utils import printme 	#home-made formatting utilities

live=True

#Create some data to be used in a time series
data=[10, 20, 30, 40]

#Here's how to create a pandas Series object:
ser=pd.Series(data=data)
printme("Let's have a look:", "index  value", ser)

#Slicing operations are supported - you get an index "for free"
printme('second value only is:', ser[1])
printme('a slice from position 1 to 3:', ser[1:3])

#Like a list, values can be heterogenous objects:
ser[0]='some string'
ser[1]=(1,2,3)
ser[55]=5555555
printme(ser, '\n')

#we can remove and retrieve objects using pop() on the index
value_3=ser.pop(3)
value_5=ser.pop(55)
print('value_3:  {}'.format(value_3))
print('value_5:  {}'.format(value_5))
printme(ser)

#get() is also supported
default=ser.get(666, "nope")
printme(default)

#We can use the keyword 'in' to test indices
look_for= 0
result=look_for in ser
printme('Is {} in the series?  {}'.format(look_for, result))
printme()

#We can find values using the isin() method of Series
print('isin() produces a boolean mask; the argument is a list\n')
look_for=[ (1,2,3) ]
print('looking for', look_for)
printme(ser.isin( look_for ))

look_for=[ (1,2,3), 30, 50]
printme('looking for', look_for)
printme(ser.isin( look_for ))

#Better indices ... pandas work and play well with dates

#Create an index object (DatetimeIndex) for the time series
#   - Anything that looks like a date works. Default is range(len(data))
#   - Optional arguments can localize/normalize to 2400:
#         tz='America/Los_Angeles', normalize=True
#   - Frequencies can be ms, s, m, h, d, a <several flavors of a (annual)>
#       a-dec, a-mar, etc. For offset from specific day: pd.DateOffset(years=2)

#These all produce dates from Jan 1, 2017 - Jan 5, 2017
import datetime
data=[10, 20, 30, 40]
index=pd.date_range(start='1/1/2017',   periods=4, freq='d')
index=pd.date_range(start='Jan 1 2017', periods=4, freq='d')
index=pd.date_range(start=datetime.datetime(2017,1,1), periods=4, freq='d')
index=pd.date_range(start='1 Jan 2017', end='4 Jan 2017')

printme("Here's our index:")
printme(index)

#We can load our data and index into a Series object in one fell swoop
data=[10, 20, 30, 40]
index=pd.date_range(start='1 Jan 2017', end='4 Jan 2017')
ser=pd.Series(data=data, index=index)
printme("... and here's our series:")
printme(ser)

#It's easy top operate on the entire index or data vector:
#   advance index by a day and increment each value
printme("first date:  {}  first value: {}".format(ser.index[0], ser[0]))
ser.index=ser.index+1
ser=ser+100
printme("now, check it out:", ser)

#We can shift it out a couple years if we want
add_years=2
ser.index=ser.index+ pd.DateOffset(years=add_years)
printme("extended {} years:".format(add_years))
printme(ser)

#pandas has lots of built-in 'intelligent dates'
from pandas.core.datetools import BYearEnd, Easter, QuarterEnd, BMonthEnd
new_year=pd.datetime(2017, 1, 1)
easter= pd.date_range(start=new_year,periods=3, freq=Easter())
printme('Happy Easter!', easter)
qend=pd.date_range(start=new_year, periods=4, freq=QuarterEnd())
printme("End of quarters:", qend)

#Plotting can be super-easy with matplotlib (but it's ugly)
plt.plot(ser)
if live: plt.show()

#Tease out 'figure' and 'axes' objects to apply formatting
fig, ax = plt.subplots()

#add the Series object (we could add x and y values separately)
ax.plot(ser)

#to apply axes formatting, create and apply a formatting object
x_format = dates.DateFormatter('%Y-%m-%d')
ax.xaxis.set_major_formatter(x_format)

#this automatically 'slants' the text so it's legible
fig.autofmt_xdate()

#this 'pads out' the chart with a little extra margin
plt.tight_layout()

#we can optionally add a title
plt.title("The most boring plot ever - bigly")

#this renders the plot
plt.show()
a=1


# testDriver.py
"""Volume II Lab 7: Nearest Neighbor Search. Test Driver.

< IN DEVELOPMENT >

"""

import inspect
import numpy as np
from scipy.spatial import KDTree
from solutions import metric, postal_problem

def test(student_module):
    """Test script. You must import the students file as a module.

     5 points for problem 1
     5 points for problem 2
    10 points for problem 3
    10 points for problem 4
    20 points for problem 5
    10 points for problem 6

    Inputs:
        student_module: the imported module for the student's file.

    Returns:
        score (int): the student's score, out of 80.
        feedback (str): a printout of test results for the student.
    """
    tester = _testDriver()
    tester.test_all(student_module)
    return tester.score, tester.feedback

class _testDriver(object):
    """Class for testing a student's work. See test.__doc__ for more info."""
    def __init__(self):
        self.feedback = ""

    # Main routine -----------------------------------------------------------
    def test_all(self, student_module, total=50):
        """Grade the provided module on each problem and compile feedback."""
        # Reset feedback and score.
        self.feedback = ""
        self.score = 0

        def test_one(problem, label, value):
            """Test a single problem, checking for errors."""
            try:
                self.feedback += "\n\n{} ({} points):".format(label, value)
                points = problem(student_module)
                self.score += points
                self.feedback += "\nScore += {}".format(points)
            except BaseException as e:
                self.feedback += "\n{}: {}".format(self._errType(e), e)

        # Grade each problem.
        test_one(self.problem1, "Problem 1", 5)
        test_one(self.problem2, "Problem 2", 5)
        test_one(self.problem5, "Problems 4 and 5", 30)
        test_one(self.problem6, "Problem 6", 10)

        # Report final score.
        percentage = (100. * self.score) / total
        self.feedback += "\n\nTotal score: {}/{} = {}%".format(
                                    self.score, total, round(percentage, 2))
        if   percentage >=  98: self.feedback += "\n\nExcellent!"
        elif percentage >=  90: self.feedback += "\n\nGreat job!"

        # Add comments (optionally).
        print(self.feedback)
        comments = str(raw_input("Comments: "))
        if len(comments) > 0:
            self.feedback += '\n\n\nComments:\n\t{}'.format(comments)

    # Helper Functions --------------------------------------------------------
    @staticmethod
    def _errType(error):
        """Get just the name of the exception 'error' in string format."""
        return str(type(error).__name__)

    def _eqTest(self, correct, student, message):
        """Test to see if 'correct' and 'student' are equal.
        Report the given 'message' if they are not.
        """
        if np.allclose(correct, student, atol=1e-04):
            return 1
        else:
            self.feedback += "\n{}".format(message)
            self.feedback += "\n\tCorrect response: {}".format(correct)
            self.feedback += "\n\tStudent response: {}".format(student)
            return 0

    def _grade(self, points, message=None):
        """Manually grade a problem worth 'points'. Return the score.
        If full points are not earned, get feedback on the problem.
        """
        credit = -1
        while credit > points or credit < 0:
            try:
                credit = int(input("\nScore out of {}: ".format(points)))
            except:
                credit = -1
        if credit != points:
            # Add comments (optionally),
            comments = raw_input("Comments: ")
            if len(comments) > 0:
                self.feedback += "\n{}".format(comments)
            # Or add a predetermined error message.
            elif message is not None:
                self.feedback += "\n{}".format(message)
        return credit

    def neighbor(self, m, k, func):
        """Do a single nearest neighbor search trial for mxk data,
        solved with the function 'func'.
        """
        data = np.random.random((m, k))
        target = np.random.random(k)
        tree = KDTree(data)
        dist, index = tree.query(target)
        point = tree.data[index]
        spoint, sdist = func(data, target) # func solves the problem
        p1 = self._eqTest(point, spoint,
            "\n\t"+func.__name__+"() failed: incorrect nearest neighbor")
        p2 = self._eqTest(dist, sdist,
            "\n\t"+func.__name__+"() failed: incorrect minimum distance")
        return p1 + p2

    @staticmethod
    def get_code(func):
        rawcode = inspect.getsource(func).splitlines()[len(
                                            func.__doc__.splitlines())+1:]
        for line in rawcode: print line

    # Problems ----------------------------------------------------------------
    def problem1(self, s):
        """Test metric(). 5 Points."""

        # Test with good inputs (4 points)
        x = np.array([1, 2])
        y = np.array([2, 2])
        points = self._eqTest(metric(x,y), s.metric(x,y),
                                            "\n\tmetric() failed.")

        x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
        y = np.array([2, 6, 4, 8, 0, 2, 4, 7, 5, 11])
        points += self._eqTest(metric(x,y), s.metric(x,y),
                                            "\n\tmetric() failed.")

        x = (np.random.random(100)-.5)*200
        y = (np.random.random(100)-.5)*200
        points += self._eqTest(metric(x,y), s.metric(x,y),
                                        "\n\tmetric() failed.")*2

        # Test with bad inputs (1 point)
        x = np.array([1, 2])
        y = np.array([1, 2, 3])
        try:
            s.metric(x, y)
            self.feedback += "\n\tmetric() failed to raise a "
            self.feedback += "ValueError for vectors of different lengths"
        except:
            points += 1

        return points

    def problem2(self, s):
        """Test exhaustive_search(). 5 points."""

        points  = self.neighbor(100, 10, s.exhaustive_search)
        points += self.neighbor(10, 100, s.exhaustive_search)
        points += 1

        _testDriver.get_code(s.exhaustive_search)
        print "\n(Check that scipy.spatial.KDTree is not used)"
        points *= self._grade(1)

        return points

    def problem3(self, s):
        """Test the KDTNode class. 10 points."""

        points = 0

        # Test KDTNode.__init__ (can only hold np.ndarrays; 2 points)
        try:
            s.KDTNode("This is not a numpy array")
            self.feedback += "\n\tKDTNode(x) failed to raise a TypeError "
            self.feedback += "for x not a numpy array (np.ndarray)"
        except:
            points += 2

        # Test KDTNode.__sub__ (euclidean distance; 2 points)
        x = np.random.random(10); y = np.random.random(10)
        A =   KDTNode(x); B =   KDTNode(y)
        C = s.KDTNode(x); D = s.KDTNode(y)
        points += 2*self._eqTest(A-B, C-D, "\n\tKDTNode.__sub__ failed")

        # Test KDTNode.__eq__ (1 Point)
        D = s.KDTNode(1.5*x)
        if not (C == D):
            points += 1
        else:
            self.feedback += "\n\tKDTNode.__eq__ failed on nonequal"

        # Test KDTNode.__lt__ and KDTNode.__gt__ (5 points)
        x = s.KDTNode(np.array([3,1,0,5], dtype=np.int)); x.axis = 0
        y = s.KDTNode(np.array([1,2,4,3], dtype=np.int)); y.axis = 1
        if x < y: points += 1
        else: self.feedback += "\n\tKDTNode.__lt__ failed"
        if y < x: points += 1
        else: self.feedback += "\n\tKDTNode.__lt__ failed"

        x.axis = 2; y.axis = 3
        if x > y: points += 1
        else: self.feedback += "\n\tKDTNode.__gt__ failed"
        if y > x: points += 2
        else: self.feedback += "\n\tKDTNode.__gt__ failed"

        return points

    def problem5(self, s):
        """Test nearest_neighbor(). 30 points."""
        points = 0

        points  = self.neighbor( 10,  10, s.nearest_neighbor)*3
        points += self.neighbor(100,  10, s.nearest_neighbor)*3
        points += self.neighbor( 10, 100, s.nearest_neighbor)*3
        points += self.neighbor(100, 100, s.nearest_neighbor)*3
        points += self.neighbor(100, 100, s.nearest_neighbor)*3

        _testDriver.get_code(s.nearest_neighbor)
        print "\n(Check that scipy.spatial.KDTree is not used)"
        points *= self._grade(1)

        return points

    def problem6(self, s):
        """Test postal_problem(). 10 points."""

        print("Correct responses:")
        postal_problem(grading=True)
        print("\nStudent responses:")
        x = s.postal_problem()
        if x is not None:
            print x

        return self._grade(10)

if __name__ == '__main__':
    import solutions as sol
    score, feedback = test(sol)

# =============================== END OF FILE =============================== #


"""
*origin* and *extent* in `~.Axes.imshow`
========================================

:meth:`~.Axes.imshow` allows you to render an image (either a 2D array
which will be color-mapped (based on *norm* and *cmap*) or a 3D RGB(A)
array which will be used as-is) to a rectangular region in data space.
The orientation of the image in the final rendering is controlled by
the *origin* and *extent* kwargs (and attributes on the resulting
`~.AxesImage` instance) and the data limits of the axes.

The *extent* kwarg controls the bounding box in data coordinates that
the image will fill specified as ``(left, right, bottom, top)`` in
**data coordinates**, the *origin* kwarg controls how the image fills
that bounding box, and the orientation in the final rendered image is
also affected by the axes limits.

.. hint:: Most of the code below is used for adding labels and informative
   text to the plots. The described effects of *origin* and *extent* can be
   seen in the plots without the need to follow all code details.

   For a quick understanding, you may want to skip the code details below and
   directly continue with the discussion of the results.
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

def index_to_coordinate(index, extent, origin):
    """Return the pixel center of an index."""
    left, right, bottom, top = extent

    hshift = 0.5 * np.sign(right - left)
    left, right = left + hshift, right - hshift
    vshift = 0.5 * np.sign(top - bottom)
    bottom, top = bottom + vshift, top - vshift

    if origin == 'upper':
        bottom, top = top, bottom

    return {
        "[0, 0]": (left, bottom),
        "[M', 0]": (left, top),
        "[0, N']": (right, bottom),
        "[M', N']": (right, top),
    }[index]

def get_index_label_pos(index, extent, origin, inverted_xindex):
    """
    Return the desired position and horizontal alignment of an index label.
    """
    if extent is None:
        extent = lookup_extent(origin)
    left, right, bottom, top = extent
    x, y = index_to_coordinate(index, extent, origin)

    is_x0 = index[-2:] == "0]"
    halign = 'left' if is_x0 ^ inverted_xindex else 'right'
    hshift = 0.5 * np.sign(left - right)
    x += hshift * (1 if is_x0 else -1)
    return x, y, halign

def get_color(index, data, cmap):
    """Return the data color of an index."""
    val = {
        "[0, 0]": data[0, 0],
        "[0, N']": data[0, -1],
        "[M', 0]": data[-1, 0],
        "[M', N']": data[-1, -1],
    }[index]
    return cmap(val / data.max())

def lookup_extent(origin):
    """Return extent for label positioning when not given explicitly."""
    if origin == 'lower':
        return (-0.5, 6.5, -0.5, 5.5)
    else:
        return (-0.5, 6.5, 5.5, -0.5)

def set_extent_None_text(ax):
    ax.text(3, 2.5, 'equals\nextent=None', size='large',
            ha='center', va='center', color='w')

def plot_imshow_with_labels(ax, data, extent, origin, xlim, ylim):
    """Actually run ``imshow()`` and add extent and index labels."""
    im = ax.imshow(data, origin=origin, extent=extent)

    # extent labels (left, right, bottom, top)
    left, right, bottom, top = im.get_extent()
    if xlim is None or top > bottom:
        upper_string, lower_string = 'top', 'bottom'
    else:
        upper_string, lower_string = 'bottom', 'top'
    if ylim is None or left < right:
        port_string, starboard_string = 'left', 'right'
        inverted_xindex = False
    else:
        port_string, starboard_string = 'right', 'left'
        inverted_xindex = True
    bbox_kwargs = {'fc': 'w', 'alpha': .75, 'boxstyle': "round4"}
    ann_kwargs = {'xycoords': 'axes fraction',
                  'textcoords': 'offset points',
                  'bbox': bbox_kwargs}
    ax.annotate(upper_string, xy=(.5, 1), xytext=(0, -1),
                ha='center', va='top', **ann_kwargs)
    ax.annotate(lower_string, xy=(.5, 0), xytext=(0, 1),
                ha='center', va='bottom', **ann_kwargs)
    ax.annotate(port_string, xy=(0, .5), xytext=(1, 0),
                ha='left', va='center', rotation=90,
                **ann_kwargs)
    ax.annotate(starboard_string, xy=(1, .5), xytext=(-1, 0),
                ha='right', va='center', rotation=-90,
                **ann_kwargs)
    ax.set_title('origin: {origin}'.format(origin=origin))

    # index labels
    for index in ["[0, 0]", "[0, N']", "[M', 0]", "[M', N']"]:
        tx, ty, halign = get_index_label_pos(index, extent, origin,
                                             inverted_xindex)
        facecolor = get_color(index, data, im.get_cmap())
        ax.text(tx, ty, index, color='white', ha=halign, va='center',
                bbox={'boxstyle': 'square', 'facecolor': facecolor})
    if xlim:
        ax.set_xlim(*xlim)
    if ylim:
        ax.set_ylim(*ylim)

def generate_imshow_demo_grid(extents, xlim=None, ylim=None):
    N = len(extents)
    fig = plt.figure(tight_layout=True)
    fig.set_size_inches(6, N * (11.25) / 5)
    gs = GridSpec(N, 5, figure=fig)

    columns = {'label': [fig.add_subplot(gs[j, 0]) for j in range(N)],
               'upper': [fig.add_subplot(gs[j, 1:3]) for j in range(N)],
               'lower': [fig.add_subplot(gs[j, 3:5]) for j in range(N)]}
    x, y = np.ogrid[0:6, 0:7]
    data = x + y

    for origin in ['upper', 'lower']:
        for ax, extent in zip(columns[origin], extents):
            plot_imshow_with_labels(ax, data, extent, origin, xlim, ylim)

    columns['label'][0].set_title('extent=')
    for ax, extent in zip(columns['label'], extents):
        if extent is None:
            text = 'None'
        else:
            left, right, bottom, top = extent
            text = (f'left: {left:0.1f}\nright: {right:0.1f}\n'
                    f'bottom: {bottom:0.1f}\ntop: {top:0.1f}\n')
        ax.text(1., .5, text, transform=ax.transAxes, ha='right', va='center')
        ax.axis('off')
    return columns

###############################################################################
#
# Default extent
# --------------
#
# First, let's have a look at the default ``extent=None``

generate_imshow_demo_grid(extents=[None])

###############################################################################
#
# Generally, for an array of shape (M, N), the first index runs along the
# vertical, the second index runs along the horizontal.
# The pixel centers are at integer positions ranging from 0 to ``N' = N - 1``
# horizontally and from 0 to ``M' = M - 1`` vertically.
# *origin* determines how the data is filled in the bounding box.
#
# For ``origin='lower'``:
#
#    - [0, 0] is at (left, bottom)
#    - [M', 0] is at (left, top)
#    - [0, N'] is at (right, bottom)
#    - [M', N'] is at (right, top)
#
# ``origin='upper'`` reverses the vertical axes direction and filling:
#
#    - [0, 0] is at (left, top)
#    - [M', 0] is at (left, bottom)
#    - [0, N'] is at (right, top)
#    - [M', N'] is at (right, bottom)
#
# In summary, the position of the [0, 0] index as well as the extent are
# influenced by *origin*:
#
# ======  ===============  ==========================================
# origin  [0, 0] position  extent
# ======  ===============  ==========================================
# upper   top left         ``(-0.5, numcols-0.5, numrows-0.5, -0.5)``
# lower   bottom left      ``(-0.5, numcols-0.5, -0.5, numrows-0.5)``
# ======  ===============  ==========================================
#
# The default value of *origin* is set by :rc:`image.origin` which defaults
# to ``'upper'`` to match the matrix indexing conventions in math and
# computer graphics image indexing conventions.
#
#
# Explicit extent
# ---------------
#
# By setting *extent* we define the coordinates of the image area. The
# underlying image data is interpolated/resampled to fill that area.
#
# If the axes is set to autoscale, then the view limits of the axes are set
# to match the *extent* which ensures that the coordinate set by
# ``(left, bottom)`` is at the bottom left of the axes!  However, this
# may invert the axis so they do not increase in the 'natural' direction.
#

extents = [(-0.5, 6.5, -0.5, 5.5),
           (-0.5, 6.5, 5.5, -0.5),
           (6.5, -0.5, -0.5, 5.5),
           (6.5, -0.5, 5.5, -0.5)]

columns = generate_imshow_demo_grid(extents)
set_extent_None_text(columns['upper'][1])
set_extent_None_text(columns['lower'][0])

###############################################################################
#
# Explicit extent and axes limits
# -------------------------------
#
# If we fix the axes limits by explicitly setting `~.axes.Axes.set_xlim` /
# `~.axes.Axes.set_ylim`, we force a certain size and orientation of the axes.
# This can decouple the 'left-right' and 'top-bottom' sense of the image from
# the orientation on the screen.
#
# In the example below we have chosen the limits slightly larger than the
# extent (note the white areas within the Axes).
#
# While we keep the extents as in the examples before, the coordinate (0, 0)
# is now explicitly put at the bottom left and values increase to up and to
# the right (from the viewer's point of view).
# We can see that:
#
# - The coordinate ``(left, bottom)`` anchors the image which then fills the
#   box going towards the ``(right, top)`` point in data space.
# - The first column is always closest to the 'left'.
# - *origin* controls if the first row is closest to 'top' or 'bottom'.
# - The image may be inverted along either direction.
# - The 'left-right' and 'top-bottom' sense of the image may be uncoupled from
#   the orientation on the screen.

generate_imshow_demo_grid(extents=[None] + extents,
                          xlim=(-2, 8), ylim=(-1, 6))

plt.show()


# |------------------------------------------------------------------
# | # Geospatial Data Exercise
# |------------------------------------------------------------------
# |
# | This is an exercise notebook for the fourth lesson of the kaggle course
# | ["Geospatial Analysis"](https://www.kaggle.com/learn/geospatial-analysis)
# | offered by Alexis Cook and Jessica Li. The main goal of the lesson is
# | to get used to __Geocoding__ and __Spatial Join__.

# | ## 1. Introduction
# -------------------------------------------------------
# | Import packages.
# -------------------------------------------------------
import geopandas as gpd
from kaggle_geospatial.kgsp import *
from folium import Choropleth,  Marker, GeoJson
import folium
import plotly.graph_objs as go
from plotly.subplots import make_subplots
import webbrowser
import zipfile
from pathlib import Path
import os
import numpy as np
import pandas as pd
from geopandas.tools import geocode

# | ### Geocoding

# | Geocoding is to convert the names and the addresses of places
# | to the latitudes and the longitude, and vice versa.
# | Here is a quick experiment with 'Marienplatz' in Munich.

result_1 = geocode(["Marienplatz"])

print(type(result_1))
print(result_1.info())
print(result_1['address'])
result_1.head(1)

# | `geocode` returns a GeoDataFrame with Shapely `POINT` object
# | and the address in human-readable form.

# --
result_2 = geocode(["Neuschwanstein"])

print(type(result_2))
print(result_2.info())
print(result_2['address'])
result_2.head(1)

# --
result_3 = geocode(["Augsburger Rathaus"])

print(type(result_3))
print(result_3.info())
print(result_3['address'])
result_3.head(1)

# --
result_4 = geocode(["Rathaus, Altenmuenster"])

print(type(result_4))
print(result_4.info())
print(result_4['address'])
result_4.head(1)

# | Wow!

# | ### Spatial Join
# | See Appendix.
# |
# | ## 2. Task
# |
# | Visualize the distribution of the coffeeshop Starbucks in California.
# | Find out the best county to build the next 'Starbucks Reserve Roastery'
# | (flagship atelier/gallery shops of Starbucks) in California.

# | ## 3. Data
# |
# | 1. Locations of existing (non-Roastery) Starbucks in California.
# | 2. General underlying map.
# | 3. Boundaries of counties in California.
# | 4. Statistics of counties in California, such as population,
# |    area in km2, and median ages.  Among all, a unique information is
# |    the number of __high earners__ (household with annual income over $150,000)
# |    in each county.
# |
# | ## 4. Notebook

# -------------------------------------------------------
# | Set up some directories.

CWD = '/Users/meg/git6/geocode/'
DATA_DIR = '../input/geospatial-learn-course-data/'
KAGGLE_DIR = 'alexisbcook/geospatial-learn-course-data'
GEO_DIR = 'geospatial-learn-course-data'

set_cwd(CWD)
set_data_dir(DATA_DIR, KAGGLE_DIR, GEO_DIR, CWD)
show_whole_dataframe(True)

# -------------------------------------------------------
# | Read the starbucks data.

starbucks = pd.read_csv(DATA_DIR + 'starbucks_locations.csv')
print(starbucks.info())
starbucks.head(3)

# -------------------------------------------------------
# | Ancillary data for the state of California.

CA_data_dir = DATA_DIR+'CA_county_boundaries/CA_county_boundaries/'

CA_counties = gpd.read_file(CA_data_dir+'CA_county_boundaries.shp')
CA_pop = pd.read_csv(DATA_DIR + 'CA_county_population.csv', index_col="GEOID")
CA_high_earners = pd.read_csv(
    DATA_DIR + 'CA_county_high_earners.csv', index_col="GEOID")
CA_median_age = pd.read_csv(
    DATA_DIR + 'CA_county_median_age.csv', index_col="GEOID")

print(CA_counties.info())
print(CA_pop.info())
print(CA_high_earners.info())
print(CA_median_age.info())

print(CA_counties.head(3))
print(CA_pop.head(3))
print(CA_high_earners.head(3))
print(CA_median_age.head(3))

# -------------------------------------------------------
# | 'high earners' here means the number of household with annual
# | income of $150,000 or more.

# -------------------------------------------------------
# | Start with Starbucks data.
# | There are 5 missing 'Longitude' and 'Latitude' in the data.

starbucks.isna().sum()
starbucks.loc[starbucks['Longitude'].isna(), :]
starbucks.loc[starbucks['Latitude'].isna(), :]
missing_locations = starbucks.loc[starbucks['Latitude'].isna(), :]

# -------------------------------------------------------
# | The shops with missing 'Longitude' and 'Latitude' are
# | all in Berkeley. We will use `geocode` to find out
# | the coordinates of these shops from their
# | addresses.

x = pd.concat([geocode(r['Address'])
               for i, r in missing_locations.iterrows()], axis=0)

missing_locations[['Longitude', 'Latitude']] = [
    [p['geometry'].x, p['geometry'].y] for i, p in x.iterrows()]

starbucks = starbucks.combine_first(missing_locations)
starbucks[starbucks['City'] == 'Berkeley']

# -------------------------------------------------------
# | `.combine_first` is to fill `None` with the second
# | DataFrame. Make sure that the indices are aligned [two DataFrames
# | use the same (=consistent) index].

# -------------------------------------------------------
# | Now we have a complete table.

starbucks['Address'].str.contains('CA').mean()

# | All shops are in California.

# -------------------------------------------------------
# | We will start visualization of the locations of the cafes.
# | First, setup the center of the map, tiles, and the zoom factor.

center = [starbucks['Latitude'].mean(), starbucks['Longitude'].mean()]
tiles = 'openstreetmap'
zoom = 8

# -------------------------------------------------------
m_1 = folium.Map(location=center, tiles=tiles, zoom_start=zoom)

dump = [Marker((r['Latitude'], r['Longitude']),
               tooltip=r['Address']).add_to(m_1)
        for i, r in starbucks.iterrows()]

embed_map(m_1, './html/m_1.html')
# --
show_on_browser(m_1, CWD + './html/m_1b.html')

# -------------------------------------------------------
# | Okay, now we start working with the ancillary data.
# | Combine all of them on the index `GEOID`.
# | Make sure that the one with the geometry column
# | is the left most DataFrame.

CA_stats = CA_counties.merge(CA_pop, on='GEOID')
CA_stats = CA_stats.merge(CA_high_earners, on='GEOID')
CA_stats = CA_stats.merge(CA_median_age, on='GEOID')
# CA_stats.set_index('GEOID', inplace=True)
# -------------------------------------------------------
# | Check CRS.
CA_stats.crs

# -------------------------------------------------------
# | It looks like the order of latitude and longitude in `geometry`
# | is not correct. Check it by plotting it.

def style_function(x):
    #    return {'fillColor': 'coral', 'stroke': False}
    return {'fillColor': 'teal', 'stroke': True}

m_2 = folium.Map(location=center, tiles=tiles, zoom_start=zoom)
GeoJson(data=CA_stats.__geo_interface__,
        style_function=style_function).add_to(m_2)
embed_map(m_2, './html/m_2.html')
# --
show_on_browser(m_2, CWD + './html/m_2b.html')

# -------------------------------------------------------
# | Okay somehow GeoJson handles the coordinates correctly.

# -------------------------------------------------------
# | Add population-density, and fraction of high earners.

CA_stats['density'] = CA_stats['population'] / CA_stats['area_sqkm']
CA_stats['fraction_HE'] = CA_stats['high_earners'] / CA_stats['population']

# -------------------------------------------------------
# | Create a couple of choropleths to see
# | the demographic and the economic landscape of California.

tiles = 'Stamen Terrain'
m_3 = folium.Map(location=center, tiles=tiles, zoom_start=zoom)

Choropleth(geo_data=CA_stats.__geo_interface__,
           name='choropleth',
           data=CA_stats,
           columns=['name', 'median_age'],
           key_on='feature.properties.name',
           fill_color='YlGnBu',
           bins=[25, 30, 35, 40, 45, 50, 55, 60],
           legend_name='Median Age of Counties in CA').add_to(m_3)

embed_map(m_3, './html/m_3.html')
# --
show_on_browser(m_3, CWD + './html/m_3b.html')

# -------------------------------------------------------
# | __`Choropleth`__ summary.
# |  - `columns`: two columns of `CA_stats` for statistics to show.
# |  - `key_on` : which one of the two above to use to match with `geo_data`.
# |
# -------------------------------------------------------

# -------------------------------------------------------
# | __What we can see__
# |
# | 1.  In the counties far away from the coast line toward the western hills
# |     of Sierra Nevada, the median ages are high, sometimes over 50.
# |
# | 2.  In and near the two metropolises, Los Angeles and San Francisco,
# |    the median ages are intermediate, 35-45.
# |
# | 3.  In between the cities along the coast and Sierra Nevada, in Great Valley
# |    and Mohave Desert, the median ages are under 35.
# |
# | 4.  We might interpret the picture as a migration of people. Someone who
# |    are raised in rural part of California, come to big cities to work,
# |    and spend their retirement in highland, where the climates are mild.
# |
# -------------------------------------------------------

tiles = 'Stamen Terrain'
m_4 = folium.Map(location=center, tiles=tiles, zoom_start=zoom)

Choropleth(geo_data=CA_stats.__geo_interface__,
           name='choropleth',
           data=CA_stats,
           columns=['name', 'high_earners'],
           key_on='feature.properties.name',
           fill_color='YlGnBu',
           bins=10 ** np.array([1, 2, 3, 4, 5, 6]),
           legend_name='Number of Household with Annual Income > $150k in CA').add_to(m_4)

embed_map(m_4, './html/m_4.html')
# --
show_on_browser(m_4, CWD + './html/m_4b.html')

# -------------------------------------------------------
tiles = 'Stamen Terrain'
m_5 = folium.Map(location=center, tiles=tiles, zoom_start=zoom)

Choropleth(geo_data=CA_stats.__geo_interface__,
           name='choropleth',
           data=CA_stats,
           columns=['name', 'fraction_HE'],
           key_on='feature.properties.name',
           fill_color='YlGnBu',
           bins=[0.0, 0.02, 0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16],
           legend_name='Fraction of Household with Annual Income > $150k in CA').add_to(m_5)

embed_map(m_5, './html/m_5.html')
# --
show_on_browser(m_5, CWD + './html/m_5b.html')

# -------------------------------------------------------
# | 1. Northern part of San Francisco Bay Area has the highest
# | fraction of high earners.
# -------------------------------------------------------
# | Let us look at the number of Starbucks stores in each county.

starbucks = gpd.GeoDataFrame(starbucks,
                             geometry=gpd.points_from_xy(
                                 starbucks['Longitude'], starbucks['Latitude']))

starbucks.crs = {'init': 'epsg:4326'}

number_of_sb = []
for i, c in CA_stats.iterrows():
    number_of_sb.append(sum([c['geometry'].contains(s)
                             for s in starbucks['geometry']]))

CA_stats['number_of_sb'] = number_of_sb

# -------------------------------------------------------
tiles = 'cartodbpositron'
m_6 = folium.Map(location=center, tiles=tiles, zoom_start=zoom)

Choropleth(geo_data=CA_stats.__geo_interface__,
           name='choropleth',
           data=CA_stats,
           columns=['name', 'number_of_sb'],
           key_on='feature.properties.name',
           fill_color='YlGnBu',
           bins=[0, 10, 50, 100, 200, 400, 800],
           legend_name='Number of Starbucks Cafes').add_to(m_6)

embed_map(m_6, './html/m_6.html')
# --
show_on_browser(m_6, CWD + './html/m_6b.html')

# -------------------------------------------------------
CA_stats['sb_per_pop'] = CA_stats['number_of_sb'] / \
    CA_stats['population'] * 10**5

# -------------------------------------------------------
tiles = 'cartodbpositron'
m_7 = folium.Map(location=center, tiles=tiles, zoom_start=zoom)
Choropleth(geo_data=CA_stats.__geo_interface__,
           name='choropleth',
           data=CA_stats,
           columns=['name', 'sb_per_pop'],
           key_on='feature.properties.name',
           fill_color='YlGnBu',
           bins=[0, 2, 4, 6, 8, 10, 12, 14, 16],
           legend_name='Number of Starbucks Cafes per 100,000 People').add_to(m_7)

embed_map(m_7, './html/m_7.html')
# --
show_on_browser(m_7, CWD + './html/m_7b.html')
# -------------------------------------------------------
# |  1. Highest number of Starbucks cafes per capita
# | was Mono County in Sierra Nevada.
# -------------------------------------------------------
CA_stats['sb_per_HE'] = CA_stats['high_earners'] / \
    CA_stats['population']

# -------------------------------------------------------
tiles = 'cartodbpositron'
m_8 = folium.Map(location=center, tiles=tiles, zoom_start=zoom)
Choropleth(geo_data=CA_stats.__geo_interface__,
           name='choropleth',
           data=CA_stats,
           columns=['name', 'sb_per_HE'],
           key_on='feature.properties.name',
           tooltip=folium.features.GeoJsonTooltip(fields=['name']),
           fill_color='YlGnBu',
           bins=[0.0, 0.02, 0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16],
           legend_name='Number of Starbucks Cafes per HE').add_to(m_8)

embed_map(m_8, './html/m_8.html')
# --
show_on_browser(m_8, CWD + './html/m_8b.html')
# -------------------------------------------------------
# | See all of the above in Scatter plot and Bar plots.
# -------------------------------------------------------

trace = go.Scatter(x=CA_stats['median_age'],
                   y=CA_stats['fraction_HE'],
                   mode='markers',
                   marker=dict(color='coral',
                               size=20),
                   #                   textfont=dict(size=32),
                   hovertext=CA_stats['name'],
                   hoverinfo='text',
                   opacity=0.8)

data = [trace]

layout = go.Layout(height=1024, width=1024,
                   font=dict(size=20),
                   xaxis=dict(title=dict(text='Median Age')),
                   yaxis=dict(title=dict(text='Fraction of High Earners')),
                   showlegend=False)

fig = go.Figure(data=data, layout=layout)
# --
embed_plot(fig, './html/p_1.html')
# --
fig.show()
# -------------------------------------------------------
# | 1. There is a cluster at the median age between 36 and 43
# | and the fraction of high earner over 0.06.

# -------------------------------------------------------
CA_HE = CA_stats[(CA_stats['fraction_HE'] >= 0.06)
                 & (CA_stats['median_age'] >= 36)
                 & (CA_stats['median_age'] <= 43)]
len(CA_HE)
CA_HE.head(3)

n_rows = 5
n_cols = 1
fig = make_subplots(rows=n_rows, cols=n_cols,
                    vertical_spacing=0.02,
                    subplot_titles=[
                        #                                    'Population',
                        'Population Density',
                        'Median Age',
                        'Number of Starbucks',
                        'Number of Starbucks per 100,000 People',
                        'Number of Starbucks per High Earner'])

trace_density = go.Bar(y=CA_HE.sort_values('sb_per_HE')['name'],
                       x=CA_HE.sort_values('sb_per_HE')['density'],
                       #                       x=CA_HE.sort_values('sb_per_HE')['population'],
                       xaxis='x', yaxis='y', orientation='h')

trace_med_age = go.Bar(y=CA_HE.sort_values('sb_per_HE')['name'],
                       x=CA_HE.sort_values('sb_per_HE')['median_age'],
                       xaxis='x2', yaxis='y2', orientation='h')

trace_sb = go.Bar(y=CA_HE.sort_values('sb_per_HE')['name'],
                  x=CA_HE.sort_values('sb_per_HE')['number_of_sb'],
                  xaxis='x3', yaxis='y3', orientation='h')

trace_sb_pop = go.Bar(y=CA_HE.sort_values('sb_per_HE')['name'],
                      x=CA_HE.sort_values('sb_per_HE')['sb_per_pop'],
                      xaxis='x4', yaxis='y4', orientation='h')

trace_sb_HE = go.Bar(y=CA_HE.sort_values('sb_per_HE')['name'],
                     x=CA_HE.sort_values('sb_per_HE')['sb_per_HE'],
                     xaxis='x5', yaxis='y5', orientation='h')

data = [trace_density, trace_med_age, trace_sb, trace_sb_pop, trace_sb_HE]

layout = go.Layout(height=640 * 5, width=1024,
                   font=dict(size=20),
                   showlegend=False)

layout = fig.layout.update(layout)
fig = go.Figure(data=data, layout=layout)
# --
embed_plot(fig, './html/p_2.html')
# -
fig.show()

# -------------------------------------------------------
# | ## 5. Conclusion
# | Where to locate the next Starbucks Reserve Roastery in California?
# |
# | 1. The target customer-segment would be young, professional, and early adapter
# |  with the priority in this order.
# |
# | 2. The strategy depends on how we see the number of Starbucks per capita.
# |    The presence of many Starbucks in a county is a sign of success? Or
# |    a sign of saturation?
# |
# | 3. A bit younger median ages among Starbucks-dense counties,
# |  Santa Clara and Alameda stand out.  My pick is __Alameda__.
# |
# |

# -------------------------------------------------------
# | ## 6. Appendix
# | ### Spatial Join
# | 'Spatial join' joins two GeoDataFrames according to their geometrical matching.
# | Suppose that the first GeoDataFrames has `POINT` object in 'geometry' column
# | and the second GeoDataFrames has `POLYGON`. One can join two GeoDataFrame so that
# | the `POINT` is included in `POLYGON`.
# |
# | We assigned the total number of Starbucks in `starbucks` to each county
# | in `CA_stats` above. Let us do the other way around, and add county information
# | to `starbucks`.

starbucks.head(3)
CA_stats.head(3)

x = starbucks.sjoin(CA_stats)
x.head(3)

# -------------------------------------------------------
# | END


# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'
# %%
# Import Libraries
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
plt.style.use('seaborn-pastel')

# %% [markdown]
#  # Lets load the data
#  Okay, we got our libraries loaded, now it is time to use pandas to read
#  the train.csv file that is located in the data folder.
#
#  Pandas read_csv command is able to read and parse csv files into dataframes,
#  you can see the dataframe as a matrix on steroids.
#
#  ## Data Dictionary
#
#
#  | Variable | Definition | Key |
#  | :-: | :-: | :-: |
#  | survival | Survival | 0 = No, 1 = Yes|
#  | pclass | Ticket class | 1 = 1st, 2 = 2nd, 3 = 3rd|
#  | sex | Sex | |
#  | Age | Age in years | |
#  | sibsp | # of siblings / spouses aboard the Titanic | |
#  | parch | # of parents / children aboard the Titanic | |
#  | ticket | Ticket number | |
#  | fare | Passenger fare | |
#  | cabin | Cabin number | |
#  | embarked | Port of Embarkation | C = Cherbourg, Q = Queenstown, S = Southampton |

# %%
dataset = pd.read_csv('data/train.csv')
dataset.head()

# %% [markdown]
#  # Some Terminology
#  In our *dataset* we distinguish between the variables that gives us information that we are going to use
#  in the estimation and we call them *features*, the variable we are trying to estimate is called target.
#
#  In formal machine learning books, you will often see the dataset expressed as a set $D$ composed of
#  $n$ tuples $(\textbf{x},y)$ where $\textbf{x}$ is the feature vector (or predictors) and y is the target variable.
#
#  In this specific problem, **Survived** is the attribute that is our *target* value, all the other attributes from the
#  dataset are *features*. However we don't need (and in this case we should not) use all the avaliable attributes as features for
#  our model.

# %%
dataset.drop(labels=['Name', 'Ticket'],
             axis=1,
             inplace=True)
dataset.head()

# %%
# Separate the data into two subsets to see if there is some difference in distr
# from survived to non survived

survivedSubset = dataset[dataset['Survived'] == 1]
notSurvivedSubset = dataset[dataset['Survived'] == 0]

totalSurvived = len(survivedSubset)
totalNotSurvived = len(notSurvivedSubset)

print(
    f'Total Survived: {totalSurvived}\nTotal Not Survived: {totalNotSurvived}')
print(f'Ratio NS/S: {totalNotSurvived / totalSurvived :.2f}')

# %%
for column in dataset.columns.values:
    hasMissingValues = dataset[column].isnull().values.any()
    if hasMissingValues:
        print(f'Column: {column} has missing values')

# %% [markdown]
# # Some helper functions
# Lets define two auxiliary functions to help us plot some pie charts
# to see how are the how many observations on the dataset are categorized according to
# the attributes ** Sex ** and **Embarked** as well as the ** Survived ** target.

# %%

def drawPieChart(labels,
                 count,
                 title=' '):

    # Pie chart, where the slices will be ordered and plotted counter-clockwise:
    _, ax1 = plt.subplots()
    ax1.set_title(title)
    explode = [0.1 for i in labels]
    ax1.pie(count, labels=labels,
            autopct=lambda perc: f'{perc:.2f}% ({int(perc * sum(count)/100)})',
            shadow=True, explode=explode,
            pctdistance=0.7, startangle=90)
    # Equal aspect ratio ensures that pie is drawn as a circle.
    ax1.axis('equal')

    plt.show()

def getFrequenciesInCategoricalColumn(dataframe, columnName):
    labels = dataframe[columnName].unique()
    # print(f'Antes: {labels}, {list(map(lambda x: x is np.nan,labels))}')
    labels = sorted(labels, key=lambda x: '0' if x is np.nan else x)
    # print(labels)

    def sumatoryFunction(columnContent):
        if columnContent is np.nan:
            return sum(dataframe[columnName].isnull())
        return sum(dataframe[columnName] == columnContent)

    count = list(map(sumatoryFunction,
                     labels))
    return labels, count

# %%
columnName = 'Survived'
labels, count = getFrequenciesInCategoricalColumn(dataset, columnName)
drawPieChart(labels, count, columnName)

# %%
columnName = 'Sex'
labels, count = getFrequenciesInCategoricalColumn(dataset, columnName)
drawPieChart(labels, count, columnName)

# %%
columnName = 'Embarked'
labels, count = getFrequenciesInCategoricalColumn(dataset, columnName)
drawPieChart(labels, count, columnName)

# %%
columnName = 'Survived'
datasetFemale = dataset[dataset['Sex'] == 'female']
datasetMale = dataset[dataset['Sex'] == 'male']

labelsFemale, countFemale = getFrequenciesInCategoricalColumn(
    datasetFemale, columnName)
labelsFemale = list(
    map(lambda x: 'Female Survived' if x else 'Female Not Survived', labelsFemale))

labelsMale, countMale = getFrequenciesInCategoricalColumn(
    datasetMale, columnName)
labelsMale = list(
    map(lambda x: 'Male Survived' if x else 'Male Not Survived', labelsMale))

drawPieChart(list(labelsFemale) + list(labelsMale),
             countFemale + countMale, columnName)

# %%
# Does the gender influences on surviability?

fig, ax = plt.subplots()

width = 0.7
ind = np.arange(2)    # the x locations for the groups

p1 = ax.bar(ind, countMale, width, label='Men')
p2 = ax.bar(ind, countFemale, width,
            bottom=countMale, label='Women')

ax.axhline(0, color='grey', linewidth=0.8)
ax.set_ylabel('Quantity')
ax.set_title('Gender Surviability')
ax.set_xticks(ind)
ax.set_xticklabels(('Not Survived', 'Survived'))
ax.legend()

# Label with label_type 'center' instead of the default 'edge'
ax.bar_label(p1, label_type='center')
ax.bar_label(p2, label_type='center')
ax.bar_label(p2)

plt.show()

# %%
# Lets see the distributions between survived and not survived
# we will use both boxplot and violin plot to see the pros and cons
# in each of them.
# TLDR: Violin shows the distr whereas boxplot only shows quartiles.

labels = ['Survived', 'Not Survived']

fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(18, 8))

# rectangular box plot
bplot1 = ax1.boxplot([survivedSubset['Age'].dropna(axis=0).values,
                      notSurvivedSubset['Age'].dropna(axis=0).values],
                     vert=True,  # vertical box alignment
                     patch_artist=True,  # fill with color
                     labels=labels)  # will be used to label x-ticks
ax1.set_title('Rectangular box plot')

bplot2 = ax2.violinplot([survivedSubset['Age'].dropna(axis=0).values,
                         notSurvivedSubset['Age'].dropna(axis=0).values])
ax2.set_title('Violin plot')

ax2.set_xticks([1, 2])
ax2.set_xticklabels(labels)

# fill with colors
colors = ['pink', 'lightblue']

for patch, color in zip(bplot1['boxes'], colors):
    patch.set_facecolor(color)

# adding horizontal grid lines

ax1.yaxis.grid(True)
ax1.set_xlabel('Survived')
ax1.set_ylabel('Age')

plt.show()

# %%
# Lets look at some histograms to see if something  appears looking at the fare

bins = np.linspace(0, 100, 10)

# plt.hist(survivedSubset['Age'].dropna(axis=0).values,
#          bins, alpha=0.5, label='Survived')
plt.hist([notSurvivedSubset['Age'].dropna(axis=0).values,
          survivedSubset['Age'].dropna(axis=0).values],
         bins,
         alpha=0.5,
         histtype='barstacked',
         stacked=True,
         label=['Not Survived', 'Survived'])
plt.legend(loc='upper right')
plt.ylabel('Nº People')
plt.xlabel('Age')
plt.show()

histogramMin = min(dataset['Fare'].values)
histogramMax = max(dataset['Fare'].values)

bins = np.linspace(histogramMin, histogramMax, 10)
# plt.hist(survivedSubset['Fare'].dropna(axis=0).values,
#          bins, alpha=0.5, label='Survived',
#          histtype='barstacked')
plt.hist([notSurvivedSubset['Fare'].dropna(axis=0).values,
          survivedSubset['Fare'].dropna(axis=0).values],
         bins,
         alpha=0.5,
         label=['Not Survived', 'Survived'],
         histtype='barstacked',
         stacked=True)
plt.legend(loc='upper right')
plt.ylabel('Nº People')
plt.xlabel('Fare')
plt.show()

# %%
# Lets see if there is a difference in surviability given different classes

fig, ax = plt.subplots()

width = 0.7
ind = np.arange(3)    # the x locations for the groups

survivedClass = list(
    map(lambda x: sum(survivedSubset['Pclass'] == x), [1, 2, 3]))
notSurvivedClass = list(
    map(lambda x: sum(notSurvivedSubset['Pclass'] == x), [1, 2, 3]))

p1 = ax.bar(ind, survivedClass, width, label='Survived')
p2 = ax.bar(ind, notSurvivedClass, width,
            bottom=survivedClass, label='Not Survived')

ax.axhline(0, color='grey', linewidth=0.8)
ax.set_ylabel('Quantity')
ax.set_title('pClass Surviability')
ax.set_xticks(ind)
ax.set_xticklabels(('Class 1', 'Class 2', 'Class 3'))
ax.legend()

# Label with label_type 'center' instead of the default 'edge'
# To create the percentage labels whe use list comprehension
# in a zipped list to compute the percentage of survivors
labelsSurvived = [
    round((perc[0]/sum(perc))*100, 2) for perc in zip(survivedClass, notSurvivedClass)]
# the percentage not survived is 100 - percentage survived
labelsNotSurvived = list(map(lambda x: f'{100 - x:.2f}%', labelsSurvived))
# converts to string
labelsSurvived = list(map(lambda x: f'{x}%', labelsSurvived))

ax.bar_label(p1, labels=labelsSurvived, label_type='center')
ax.bar_label(p2, labels=labelsNotSurvived, label_type='center')
ax.bar_label(p2)

plt.show()

# %% [markdown]
# # Feature Engineering + DataVis
# By combining feature engineering to datavis we can see if different age groups have different survival rates.
# We compute the age group by doing an integer division of age by 10, this way we categorize 0-9 years passenger in group 0, 10-19 years in group 1, and so on.
#
# Then we use the **groupby** function in pandas (similar to the *SQL* groupby) and we aggregate values with the np.mean function;

# %%
datasetAges = dataset.dropna(axis=0, subset=['Age'])
datasetAges = datasetAges.assign(AgeGroup=datasetAges['Age'].apply(
    lambda x: x // 10 if x // 10 <= 6 else 6))
# There is a trick here, since Survived is a boolean attribute
# the Survived mean is the same as Rate of Survival (do the math smarty pants!)
ageGroup = datasetAges.groupby(['AgeGroup']).aggregate([np.mean, np.var])

# %%

fig, ax = plt.subplots()
survivedRatio = ageGroup['Survived']['mean'].values
ax.scatter(ageGroup.index.values,
           survivedRatio)
ax.set_xticks(ageGroup.index.values)

ax.set_title('Surviability by Age Group')
ax.set_ylim((0, 1))
ax.set_ylabel('Survival Rate')
ax.set_xlabel('Age Group')

ax.set_xticklabels(['0 - 10', '10 - 20', '20 - 30',
                   '30 - 40', '40 - 50', '50 -60', '60+'])
for i in range(len(survivedRatio)):
    xyAnnotation = list(zip(ageGroup.index.values, survivedRatio))
# The magic number 0.03 serves as a offset so the annotation wont
# be in the same place of the dot in the graph
    xyAnnotationPlace = list(
        map(lambda x: (x[0]+0.03, x[1]+0.03), xyAnnotation))
    ax.annotate(f'{survivedRatio[i]*100:.2f}%',
                xy=xyAnnotation[i],
                xytext=xyAnnotationPlace[i])
# (ageGroup['Survived']['mean'].values)

fig.show()


import torch
import numpy as np
import tensorflow as tf
import tensorflow.keras.backend as K
from collections import Iterable
from typing import Any, Optional

from graphgallery import backend, intx, floatx
from graphgallery.utils.raise_error import assert_kind

__all__ = ['is_iterable',
           'is_list_like',
           'is_scalar_like',
           'is_interger_scalar',
           'infer_type',
           'is_tensor',
           'is_strided_tensor',
           'is_sparse_tensor',
           ]

def is_iterable(obj: Any) -> bool:
    """check whether `x` is an iterable object but not string"""
    return isinstance(obj, Iterable) and not isinstance(obj, str)

def is_list_like(x: Any) -> bool:
    """Check whether `x` is list like, e.g., Tuple or List.

    Parameters:
    ----------
    x: A python object to check.

    Returns:
    ----------
    `True` iff `x` is a list like sequence.
    """
    return isinstance(x, (list, tuple))

def is_scalar_like(x: Any) -> bool:
    """Check whether `x` is a scalar, an array scalar, or a 0-dim array.

    Parameters:
    ----------
    x: A python object to check.

    Returns:
    ----------
    `True` iff `x` is a scalar, an array scalar, or a 0-dim array.
    """
    return np.isscalar(x) or (isinstance(x, np.ndarray) and x.ndim == 0)

def is_interger_scalar(x: Any) -> bool:
    """Check whether `x` is an Integer scalar.

    Parameters:
    ----------
    x: A python object to check.

    Returns:
    ----------
    `True` iff `x` is a Integer scalar (built-in or Numpy integer).
    """
    return isinstance(x, (int, np.int8,
                          np.int16,
                          np.int32,
                          np.int64,
                          np.uint8,
                          np.uint16,
                          np.uint32,
                          np.uint64,
                          ))

def infer_type(x: Any) -> bool:
    """Infer type of the input `x`.

    Parameters:
    ----------
    x: Any python object

    Returns:
    ----------
    dtype: string, the converted type of `x`:
        1. `graphgallery.floatx()` if `x` is floating
        2. `graphgallery.intx()` if `x` is integer
        3. `'bool'` if `x` is bool.

    """
    # For tensor or variable
    if is_th_tensor(x):
        if x.dtype.is_floating_point:
            return floatx()
        elif x.dtype == torch.bool:
            return 'bool'
        elif 'int' in str(x.dtype):
            return intx()
        else:
            raise RuntimeError(f'Invalid input of `{type(x)}`')

    elif is_tf_tensor(x):
        if x.dtype.is_floating:
            return floatx()
        elif x.dtype.is_integer or x.dtype.is_unsigned:
            return intx()
        elif x.dtype.is_bool:
            return 'bool'
        else:
            raise RuntimeError(f'Invalid input of `{type(x)}`')

    if not hasattr(x, 'dtype'):
        x = np.asarray(x)

    if x.dtype.kind in {'f', 'c'}:
        return floatx()
    elif x.dtype.kind in {'i', 'u'}:
        return intx()
    elif x.dtype.kind == 'b':
        return 'bool'
    elif x.dtype.kind == 'O':
        raise RuntimeError(f'Invalid inputs of `{x}`.')
    else:
        raise RuntimeError(f'Invalid input of `{type(x)}`')

def is_sparse_tensor(x: Any, kind: Optional[str] = None) -> bool:
    """Check whether `x` is a sparse Tensor.

    Parameters:
    ----------
    x: A python object to check.

    kind: str, optional.
        "T" for TensorFlow
        "P" for PyTorch
        if not specified, using `backend().kind` instead.

    Returns:
    ----------
    `True` iff `x` is a (tf or torch) sparse-tensor.
    """
    if kind is None:
        kind = backend().kind
    else:
        assert_kind(kind)

    if kind == "T":
        return is_tf_sparse_tensor(x)
    else:
        return is_th_sparse_tensor(x)

def is_strided_tensor(x: Any, kind: Optional[str] = None) -> bool:
    """Check whether `x` is a strided (dense) Tensor.

    Parameters:
    ----------
    x: A python object to check.

    kind: str, optional.
        "T" for TensorFlow
        "P" for PyTorch
        if not specified, using `backend().kind` instead.

    Returns:
    ----------
    `True` iff `x` is a (tf or torch) strided (dense) Tensor.
    """

    if kind is None:
        kind = backend().kind
    else:
        assert_kind(kind)

    if kind == "T":
        return is_tf_strided_tensor(x)
    else:
        return is_th_strided_tensor(x)

def is_tensor(x: Any, kind: Optional[str]=None) -> bool:
    """Check whether `x` is
        tf.Tensor,
        tf.Variable,
        tf.RaggedTensor,
        tf.sparse.SparseTensor,
        torch.Tensor,
        torch.sparse.Tensor.

    Parameters:
    ----------
    x: A python object to check.

    kind: str, optional.
        "T" for TensorFlow
        "P" for PyTorch
        if not specified, using `backend().kind` instead.

    Returns:
    ----------
    `True` iff `x` is a (tf or torch) (sparse-)tensor.
    """
    if kind is None:
        kind = backend().kind
    else:
        assert_kind(kind)

    if kind == "T":
        return is_tf_tensor(x)
    else:
        return is_th_tensor(x)

def is_tf_sparse_tensor(x: Any) -> bool:
    return K.is_sparse(x)

def is_th_sparse_tensor(x: Any) -> bool:
    return is_th_tensor(x) and not is_th_strided_tensor(x)

def is_tf_strided_tensor(x: Any) -> bool:
    return any((isinstance(x, tf.Tensor),
                isinstance(x, tf.Variable),
                isinstance(x, tf.RaggedTensor)))

def is_th_strided_tensor(x: Any) -> bool:
    return is_th_tensor(x) and x.layout == torch.strided

def is_tf_tensor(x: Any) -> bool:
    return is_tf_strided_tensor(x) or is_tf_sparse_tensor(x)

def is_th_tensor(x: Any) -> bool:
    return torch.is_tensor(x)