"""
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


import logging
import traceback
import os
from typing import Dict
from math import ceil
from itertools import repeat, takewhile

import multiprocessing
import numpy as np

from playground.main import get_args
from playground import ColorizedLogger, Configuration, timeit

# Create loggers with different colors to use in each problem
main_logger = ColorizedLogger('Main', 'yellow')
p1_logger = ColorizedLogger('Problem1', 'blue')
p2_logger = ColorizedLogger('Problem2', 'green')
p3_logger = ColorizedLogger('Problem3', 'magenta')
extra_ch_logger = ColorizedLogger('ExtraMain', 'yellow')
extra_sub_ch_logger = ColorizedLogger('ExtraSub', 'cyan')

# Global Vars (For the Extra Challenges)
# lock: multiprocessing.Lock
# multi_list: List = []

def my_pid(x: int) -> None:
    """ Problem 1 function to be called using pool.map

    Parameters:
        x: the id of the worker
    """

    pid = os.getpid()
    p1_logger.info(f"Hi, I’m worker {x} (with {pid})")

def problem1(conf: Dict) -> None:
    """ Problem 1 solution

    Parameters:
         conf: The config loaded from the yml file
            Example:
                properties:
                  pool_size: 4
                  chunk_size: 1
                  x_min: 0
                  x_max: 9
                conf_type: required
    """

    p1_logger.info("Starting Problem 1..")
    conf_props = conf['properties']
    # Generate iterable from `x_min` to `x_max`
    xs = range(conf_props["x_min"], conf_props["x_max"] + 1)
    p1_logger.info(f"Pool Size: {conf_props['pool_size']}: Will call `my_pid` for x in {tuple(xs)}")
    # Call my_pid() using pool.map()
    with multiprocessing.Pool(processes=conf_props['pool_size']) as pool:
        # https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool.map
        pool.map(func=my_pid,
                 iterable=xs,
                 chunksize=conf_props['chunk_size'])

# Call the timeit wrapper from main.py and pass it a custom string to print
# It already supports string formatting for func_name`, `args`, and `duration`
# To reference the first positional argument of the function I wrap, I can use {0}
@timeit(custom_print='N={0}: Calculation of pi took: {duration:2.5f} sec(s) total')
def py_pi(N: int, real_pi: float) -> None:
    """ Problem 2 function to be called using pool.starmap

    Parameters:
        N: Number of terms to be used to calculate pi
        real_pi: Correct value of pi
    """

    # Construct the equation
    first_term = (4 / N)
    list_to_be_summed = [1 / (1 + ((i - 0.5) / N) ** 2) for i in range(1, N + 1)]
    second_term = sum(list_to_be_summed)
    calced_pi = first_term * second_term
    # Calculate the absolute difference between the calculated pi and real value of pi
    pi_diff = abs(real_pi - calced_pi)
    p2_logger.info(f"N={N}: Pi({N}) = {calced_pi} (Real is {real_pi}, difference is {pi_diff})")

def problem2(conf: Dict) -> None:
    """ Problem 2 solution

    Parameters:
         conf: The config loaded from the yml file
            Example:
                properties:
                  pool_size: 4
                  chunk_size: 1
                  num_terms_min: 10
                  num_terms_step: 5
                  num_terms_max: 3906250
                conf_type: required

    """

    p2_logger.info("Starting Problem 2..")
    real_pi = np.pi
    conf_props = conf['properties']
    # Generate iterable with number of terms to be used
    # Lambda function that multiplies each element by 5 in the power of the element's index
    # e.g. [(10, 0), (10, 1), (10, 2)] -> [10*5^0, 10*5^1, 10*5^2]
    # the first ind_and_num would be (10, 0)
    multiplied_series = lambda ind_and_num: (ind_and_num[1] *
                                             (conf_props["num_terms_step"] ** (ind_and_num[0])))
    # Map an infinite enumerated iterable of (index, 10) to the lambda function
    # Stop when the series exceeds 3906250
    num_terms_range = takewhile(lambda series_el: series_el <= conf_props["num_terms_max"],
                                map(multiplied_series, enumerate(repeat(conf_props["num_terms_min"]))))
    num_terms_range = tuple(num_terms_range)  # I only do this because I want to print them
    # Create the iterable of arguments to be passed to py_pi using pool.starmap()
    # repeat() propagates the same `real_pi` value as many times as necessary
    # to zip it with num_terms_range
    args = zip(num_terms_range, repeat(real_pi))
    p2_logger.info(f"Will call `py_pi` for N in {tuple(num_terms_range)}")
    # Call py_pi() using pool.starmap() (starmap accepts iterable with multiple arguments)
    with multiprocessing.Pool(processes=conf_props['pool_size']) as pool:
        # https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool.map
        pool.starmap(func=py_pi,
                     iterable=args,
                     chunksize=conf_props['chunk_size'])

def py_pi_better(N: int, i_start: int, i_stop: int) -> float:
    """ Problem 3 function to be called using pool.starmap

    Parameters:
        N: Number of terms to be used to calculate pi
        i_start: The starting index of the sum used in pi's approximation
        i_stop: The ending index of the sum used in pi's approximation
    Returns:
        partial_calced_pi: The part of pi's approximation calculated
    """

    # Construct the equation
    first_term = (4 / N)
    list_to_be_summed = [1 / (1 + ((i - 0.5) / N) ** 2) for i in range(i_start, i_stop + 1)]
    second_term = sum(list_to_be_summed)
    cacled_pi = first_term * second_term
    return cacled_pi

def problem3(conf: Dict) -> None:
    """ Problem 3 solution

    Parameters:
         conf: The config loaded from the yml file
            Example:
                properties:
                  pool_size: 4
                  chunk_size: 1
                  num_terms:
                    - 100
                    - 500
                    - 1000
                    - 2000
                    - 10000
                    - 50000
                conf_type: required
    """

    p3_logger.info("Starting Problem 3..")
    conf_props = conf['properties']
    p3_logger.info(f"Will call `py_pi_better` for N in {tuple(conf_props['num_terms'])}")
    # Run the pi calculation once for each number of terms requested in the yml
    for num_term in conf_props["num_terms"]:
        # Split the work of `num_terms` into `pool_size` number of parts
        step = ceil(num_term / conf_props["pool_size"])
        i_start = range(1, num_term, step)
        i_stop = map(lambda el: el + step - 1 if el + step - 1 <= num_term else num_term, i_start)
        # Zip N with the i_start and i_stop iterables. Propagate the same N value using repeat
        args = zip(repeat(num_term), i_start, i_stop)
        # Call py_pi_better() using pool.starmap() (starmap accepts iterable with multiple arguments)
        with multiprocessing.Pool(processes=conf_props['pool_size']) as pool:
            # timeit can be used as a context manager too. Pass it a custom string and count the
            # total time to calculate pi
            custom_string = f'N={num_term}: Parallel calculation of pi took: ' + \
                            '{duration:2.5f} sec(s) total'
            with timeit(custom_print=custom_string):
                # https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool.map
                pi_chunks = pool.starmap(func=py_pi_better,
                                         iterable=args,
                                         chunksize=conf_props['chunk_size'])
                calced_pi = sum(pi_chunks)
        real_pi = np.pi
        pi_diff = abs(real_pi - calced_pi)
        p3_logger.info(f"N={num_term}: Pi({num_term}) = {calced_pi} (Real is {real_pi}, "
                       f"difference is {pi_diff})")

def extra_1(conf_props: Dict):
    """ Extra Challenge 1 solution

    Parameters:
         conf_props: The config loaded from the yml file
            Example:
                pool_sizes:
                  - 2
                  - 4
                  - 8
                  - 16
                  - 32
                chunk_size: 1
                num_term: 8000000
    """

    num_term = int(conf_props["num_term"])
    extra_ch_logger.info("Will call `py_pi_better` for pool_size in "
                         f"{tuple(conf_props['pool_sizes'])}")
    for pool_size in conf_props["pool_sizes"]:
        pool_size = int(pool_size)
        # Slightly modified code from problem 3
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Calling workers for N={num_term}")
        step = ceil(num_term / pool_size)
        i_start = range(1, num_term, step)
        i_stop = map(lambda el: el + step - 1 if el + step - 1 <= num_term else num_term, i_start)
        args = zip(repeat(num_term), i_start, i_stop)
        with multiprocessing.Pool(processes=pool_size) as pool:
            custom_string = f'Pool Size={pool_size}: Calculation of pi for N={num_term} took: ' + \
                            '{duration:2.5f} sec(s) total'
            with timeit(custom_print=custom_string):
                pi_chunks = pool.starmap(func=py_pi_better,
                                         iterable=args,
                                         chunksize=conf_props['chunk_size'])
                calced_pi = sum(pi_chunks)
        real_pi = np.pi
        pi_diff = abs(real_pi - calced_pi)
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Pi({num_term}) = {calced_pi} "
                                 f"(Real is {real_pi}, difference is {pi_diff})")

def py_pi_better_with_queue(N: int, i_start: int, i_stop: int,
                            m_queue: multiprocessing.Queue) -> None:
    """ Extra Challenge 2 function to be called using pool.call_async

    Parameters:
        N: Number of terms to be used to calculate pi
        i_start: The starting index of the sum used in pi's approximation
        i_stop: The ending index of the sum used in pi's approximation
        m_queue: A queue in which to append the results
    """

    # Construct the equation
    first_term = (4 / N)
    list_to_be_summed = [1 / (1 + ((i - 0.5) / N) ** 2) for i in range(i_start, i_stop + 1)]
    second_term = sum(list_to_be_summed)
    cacled_pi = first_term * second_term
    m_queue.put(cacled_pi)

def extra_2(conf_props: Dict):
    """ Extra Challenge 2 solution

    Parameters:
         conf_props: The config loaded from the yml file
            Example:
                pool_sizes:
                  - 2
                  - 4
                  - 8
                  - 16
                  - 32
                num_term: 8000000
    """

    num_term = int(conf_props["num_term"])

    extra_ch_logger.info("Will call `py_pi_better_with_queue` for pool_size in "
                         f"{tuple(conf_props['pool_sizes'])}")
    for pool_size in conf_props["pool_sizes"]:
        pool_size = int(pool_size)
        # Create a multiprocessing manager and a Queue
        multi_manager = multiprocessing.Manager()
        multi_queue = multi_manager.Queue()
        # Slightly modified code from problem 3
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Calling Async workers for N={num_term}")
        step = ceil(num_term / pool_size)
        i_start = range(1, num_term, step)
        i_stop = map(lambda el: el + step - 1 if el + step - 1 <= num_term else num_term, i_start)
        args = zip(repeat(num_term), i_start, i_stop, repeat(multi_queue))
        # Setup manually a Pool
        custom_string = f'Pool Size={pool_size}: Calculation of pi for N={num_term} and  ' + \
                        'took: {duration:2.5f} sec(s) total'
        with timeit(custom_print=custom_string):
            pool = multiprocessing.Pool(processes=pool_size)
            pool.starmap_async(func=py_pi_better_with_queue, iterable=args)
            pi_chunks = []
            while len(pi_chunks) < pool_size:
                pi_chunks.append(multi_queue.get())
            calced_pi = sum(pi_chunks)
        real_pi = np.pi
        pi_diff = abs(real_pi - calced_pi)
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Pi({num_term}) = {calced_pi}"
                                 f"(Real is {real_pi}, difference is {pi_diff})")
        pool.close()
        pool.join()

def init(_lock, _multi_list):
    global lock
    global multi_list
    lock = _lock
    multi_list = _multi_list

def py_pi_better_with_list(N: int, i_start: int, i_stop: int) -> None:
    """ Extra Challenge 2 function to be called using pool.call_async

    Parameters:
        N: Number of terms to be used to calculate pi
        i_start: The starting index of the sum used in pi's approximation
        i_stop: The ending index of the sum used in pi's approximation
    """

    # Construct the equation
    first_term = (4 / N)
    list_to_be_summed = [1 / (1 + ((i - 0.5) / N) ** 2) for i in range(i_start, i_stop + 1)]
    second_term = sum(list_to_be_summed)
    cacled_pi = first_term * second_term
    lock.acquire()
    multi_list.append(cacled_pi)
    lock.release()

def extra_3(conf_props: Dict):
    """ Extra Challenge 3 solution

    Parameters:
         conf_props: The config loaded from the yml file
    """

    num_term = int(conf_props["num_term"])

    extra_ch_logger.info("Will call `py_pi_better_with_list` for pool_size in "
                         f"{tuple(conf_props['pool_sizes'])}")
    for pool_size in conf_props["pool_sizes"]:
        pool_size = int(pool_size)
        # Create a multiprocessing manager and a Queue
        manager = multiprocessing.Manager()
        multi_list = manager.list()
        _lock = multiprocessing.Lock()
        # Slightly modified code from problem 3
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Calling Async workers for N={num_term}")
        step = ceil(num_term / pool_size)
        i_start = range(1, num_term, step)
        i_stop = map(lambda el: el + step - 1 if el + step - 1 <= num_term else num_term, i_start)
        args = zip(repeat(num_term), i_start, i_stop)
        # Setup manually a Pool
        custom_string = f'Pool Size={pool_size}: Calculation of pi for N={num_term} and  ' + \
                        'took: {duration:2.5f} sec(s) total'
        with timeit(custom_print=custom_string):
            pool = multiprocessing.Pool(initializer=init, initargs=(_lock, multi_list),
                                        processes=pool_size)
            pool.starmap_async(func=py_pi_better_with_queue, iterable=args)
            pi_chunks = []
            while len(multi_list) < pool_size:
                pi_chunks = multi_list
            calced_pi = sum(pi_chunks)
        real_pi = np.pi
        pi_diff = abs(real_pi - calced_pi)
        extra_sub_ch_logger.info(f"Pool Size={pool_size}: Pi({num_term}) = {calced_pi}"
                                 f"(Real is {real_pi}, difference is {pi_diff})")
        pool.close()
        pool.join()

def extra_challenges(conf: Dict) -> None:
    """ Extra Challenges solution

    Parameters:
         conf: The config loaded from the yml file
            Example: Examples of each sub-config in each problem's function
    """

    extra_ch_logger.info("Starting the Extra Challenges..")
    # Run all the extra challenge that have the property enabled: True
    if conf['properties']["ch1"]["enabled"]:
        extra_ch_logger.info("1: Test performance for different pool sizes.")
        extra_1(conf['properties']["ch1"])
    if conf['properties']["ch2"]["enabled"]:
        extra_ch_logger.info("2: Collect results from workers using multiprocessing list "
                             "(using starmap_async).")
        extra_2(conf['properties']["ch2"])
    if conf['properties']["ch3"]["enabled"]:
        extra_ch_logger.info("3: Collect results from workers using a pythonic list "
                             "(using starmap_async and Lock()).")
        extra_3(conf['properties']["ch3"])
    else:
        extra_ch_logger.info("3: Collect results from workers using a pythonic list "
                             "(using starmap_async and Lock()).")
        extra_ch_logger.info("Last progress report:")
        extra_ch_logger.info("For this and the 4th problem, I managed to create a Lock() and share it")
        extra_ch_logger.info("with the workers by creating an initializer function for the Pool().")
        extra_ch_logger.info("The problem was that when I did the same with the multiprocessing "
                             "List() or Array(), ")
        extra_ch_logger.info("I wasn't able to access them from the function that spawned them.")

@timeit()
def main():
    """ This is the main function of assignment.py

    Example:
        python assignment1/assignment.py \
            -c ../confs/assignment1.yml \
            -l ../logs/assignment.log
    """

    # Initialize
    args = get_args()
    ColorizedLogger.setup_logger(args.log, args.debug)
    main_logger.info("Starting Assignment 1")
    # Load the configuration
    conf = Configuration(config_src=args.config_file)
    # Start the problems defined in the configuration
    main_logger.info(f"{' Required Problems ':-^{100}}")
    check_required = lambda conf_type, tag: ((conf_type == 'required' or tag != 'required_only')
                                             and conf_type != 'disabled')
    # For each problem present in the config file, call the appropriate function
    if 'problem1' in conf.config_keys:
        for bench_conf in conf.get_config(config_name='problem1'):
            if check_required(bench_conf['type'], conf.tag):
                problem1(bench_conf)
    if 'problem2' in conf.config_keys:
        for bench_conf in conf.get_config(config_name='problem2'):
            if check_required(bench_conf['type'], conf.tag):
                problem2(bench_conf)
    if 'problem3' in conf.config_keys:
        for bench_conf in conf.get_config(config_name='problem3'):
            if check_required(bench_conf['type'], conf.tag):
                problem3(bench_conf)
    # Run the extra challenges if the tag of the conf is not set as "required_only"
    main_logger.info(f"{' Optional Problems ':-^{100}}")
    if 'extra_challenges' in conf.config_keys:
        for bench_conf in conf.get_config(config_name='extra_challenges'):
            if check_required(bench_conf['type'], conf.tag):
                extra_challenges(bench_conf)
    main_logger.info("Assignment 1 Finished")

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        logging.error(str(e) + '\n' + str(traceback.format_exc()))
        raise e


"""

Reads a file containing columns of data and creates a dictionary according to header information
using the following format:

1. comments start with ``#`` in the **first column** : ``# My Comment``

2. header information starts with ``#!``, where ``#`` is also in the **first column**. It must precede the data.

3. each columns is represented by ``name[dtype, col.nr.]/``,  *dtype* (optional) is the data type and *col.nr*
   (starting at 0) is the column number

4. data are separated by white space **NOT commas!**

5. data types:

      - s : string
      - f : float
      - i : integer

6. blank lines are ignored

 **NOTE**: if data formats are entered they should be specified for all columns

Example data::

   #! p_miss[0]/ siglt[2]/ s01[3]/ alt[4]/
   200. 1.35e-4 -1.e-3    0.1
   220. 2.56e-4 -2.e-4    -0.1
   230. 3.47e-6 -3.e-5    1.1

The header can also contain data type information::

   #! p_miss[f,0]/ siglt[f,2]/ s01[f,3]/ alt[f,4]/

   200. 1.35e-4 -1.e-3    0.1
   220. 2.56e-4 -2.e-4    -0.1
   230. 3.47e-6 -3.e-5    1.1

Example that opens a file and create a dfile object::

   >>> f0 = dfile('sig_LT.dat')

Loop over content::

   >>> for l in f0:
   >>> ...pm = l['p_miss']
   >>> ...sig_lt = l['siglt']*10000.
   >>> ...print pm, sig_lt, l['alt']
   >>> # end for l

Variables can also be accessed directly as numpy arrays (if installed):

   >>> pm = f0['p_miss']
   >>> sig_lt = f0['siglt']*10000.

They can also be converted to attributes:

   >>> f0.make_attr()
   >>> print f0.p_miss
   >>> print f0.sig_lt

Variable names are also called keys.

Save the data file as a csv file, this is useful for exporting and formatting for
documents (other than latex)

   >>> f0.write_csv(filename)

---------------------------------------------------------

"""

import sys
import os
import string
import math
import re
import pdb

numpy_ok = True
try:
    import numpy as np
except:
    numpy_ok = False

# function to create a dictionary from 2 lists: key, values
format_dict = {'i':'{:d}', 'f':'{:f}', 's':'{:s}'}
# create a dictionary from  an array of keys and values
#----------------------------------------------------------------------
def find_duplicates(array):
    # array with no duplicates
    na = []
    # counter array
    count = []
    # arrays with duplucates
    da = []
    for a in array:
        try:
            na.index(a)
        except:
            na.append(a)
            count.append(0.)
    # na has no duplicates
    for i,a in enumerate(na):
        for aa in array:
            if aa == a:
                count[i]+= 1
    for i,c in enumerate(count):
        if c>1:
            da.append(na[i])
    return da
#----------------------------------------------------------------------

class dfile:
    """
    Open a file, read and interpret the contents and return a dfile object:

    >>> df = dfile('my_datafile')

    keywords:

        debug = False (default) : print additional iformation

        fast = False (default) : load data in a fast way (using np.loadtxt) this has
                                 some limits for string entries. Best used for large purely
                                 numerical data.

        use_numpy = True (if numpy is installed) : set automatically but can be overridden
                                                   some attributes of datafile are not availables when
                                                   numpy is not present

    """
    def __init__(self, filename, debug = False, new=False, skip = True, use_numpy = numpy_ok, fast = False):
        self.H=re.compile("^#\!") # pattern for header
        self.C=re.compile("^#")   # pattern for comment
        # pattern for splitting header
        # currently:
        # ^#\! : matches line beginning with #!
        # \[[0-9]+\]\/: matches [0]/ ... [8796]/
        # \[ *\w+ *, *\w *\]\/ : matches [   f  , blabla   ]/
        #
        self.S=re.compile("^#\!|\[[0-9]+\]\/|\[ *\w *, *\w+ *\]\/")
        # pattern to find format information in the header
        self.F=re.compile("\w+ *, *\w+")
        # pattern to find a character
        self.Fc=re.compile("[a-zA-Z]")
        # pattern to find a number
        self.Fn=re.compile("[0-9]+ *\]")
        # supported formats: the output format dictionary key must match the supported formats
        self.fmt_characters = ['f','i','s']
        self.dt_formats = {'f':'f8', 'i':'i4','s':'S32'}
        self.output_format= {'f':'%r ','i':'%d ','s':'%s '}
        # empty dictionary for the formats
        self.formats = {}
        self.formats_dt = {}
        self.cols = {}
        # flag for printing debugging information
        self.debug = debug
        # flat to ignore bad line in input
        self.skip = skip
        self.use_numpy = use_numpy
        if not self.use_numpy:
            print('numpy will not be used !')
        if fast and (not use_numpy):
            print ('fast mode not possible as numpy cannot be used!')
            fast = False
        self.fast = fast
        #
        self.header = None
        self.data = []
        self.adata = []
        self.fdata=[]
        self.keys=[]
        self.new = False
        if new:
            new=True
            self.filename = filename
            self.keys.append('indx')
            self.headindex=0
            self.headerline = ('#! ')
            self.adata.append('#! ')
            return
        # pdb.set_trace()
        # open file
        self.filename = filename
        i=open(filename,"r")      # open file
        # remove leading and trailing spaces
        self.adata=[ x.strip() for x in i.readlines()]   # read all data
        if self.debug :
            print("datafile --> data read !")
        if self.debug :
            print("datafile --> spaces removed read !")
        self.remove_blanks()
        if self.debug :
            print("datafile --> blank lines removed read !")
        i.close()
        if (self.find_header() != 0): # find the header
            print("cannot interpret data, header probably missing or wrong")
            return
        if self.debug :
            print("datafile --> create arrays !")
        self.make_array()  # read data
        if self.debug :
            print("datafile --> arrays created !")
    def __getitem__(self,i):
        # allows 'direct' access to the data
        if type(i) is int:
            # it'sm index
            return self.data[i]
        elif type(i) is str:
            # its a lkey return the list of data
            return self.get_data(i)
    def __len__(self):
        return len(self.data)
    def remove_blanks(self):
        # work through the list in reverse order
        indices = list(range( len( self.adata ))) # list of indices
        indices.reverse() # reverse order
        for i in indices: # loop through from the top
            if len(self.adata[i].strip()) == 0:
                del self.adata[i]
    def find_header(self):
        # find header information : look for a line starting with #!
        for l in self.adata:
            if (self.H.match(l) != None):
                self.headindex=self.adata.index(l)
                self.headerline=l
                self.header=re.split(self.S,l)[1:-1] # remove first and last element
                self.keys=[ x.strip() for x in self.header]
                self.keys.append('indx') # store the index into the original
                # handle data format
                # check for duplicate keys
                duplicate_keys = find_duplicates(self.keys)
                if duplicate_keys != []:
                    print("duplicate keys found : ", duplicate_keys)
                    return -1
                # pdb.set_trace()
                fmt = re.findall(self.F,l)
                col_t = re.findall(self.Fn,l)
                col = [cc.replace(']','') for cc in col_t]
                # no format specification, set default format
                if fmt == []:
                    for i,k in enumerate(self.keys[:-1]) : # skip index
                        self.formats[k] = 'f'
                        self.formats_dt[k] = self.dt_formats['f']
                # use format specifications
                else:
                    for i,f in enumerate(fmt):
                        fc = re.findall(self.Fc,f)[0] # use 1st char. for format
                        self.formats[self.keys[i]] = fc
                        self.formats_dt[self.keys[i]] = self.dt_formats[fc]
                # addformat for the index
                self.formats['indx'] ='i'
                self.formats_dt['indx'] = self.dt_formats['i']
                # use format specifications
                for i,cc in enumerate(col):
                    f = self.formats[self.keys[i]]
                    if f == 's':
                        self.cols[self.keys[i]] = ''
                    else:
                        self.cols[self.keys[i]] = col[i]
                self.cols['indx'] ='-1'
                # addformat for the index
                # adata array
                # do some checks
                if (len(self.keys)  != len(self.formats)):
                    print(self.filename,": problem in data formats !")
                    print(self.adata)
                    return -1
                return 0
            #
        print(self.filename,": no header information !")
        print('data dump : ')
        print(self.adata)
        return -1

    def parse_line(self, l):
        if (self.C.match(l) != None) : # not a comment
            return None
        ll=l.split()
        # handle file format problems
        if len(ll) > (len(self.keys)-1) :
            n_data = (len(self.keys)-1)
            print(self.filename, ' error: more data than variable names !')
            print(' ---> check header : ', self.headerline)
            print('current line : ', self.line_n+1, ' ', ll)
            if self.skip:
                print('skipping !')
                return None
            else:
                sys.exit(-1)
        if len(ll) < (len(self.keys)-1) :
            print(self.filename, ' error: too few data values !')
            print(' ---> check header : ', self.headerline)
            print('current line : ', self.line_n+1, ' ', ll)
            if self.skip:
                print('skipping !')
                return None
            else:
                sys.exit(-1)
        f = []
        # now convert to values
        for i,k in enumerate(ll):
            # check if there are format specifications
            f_err = False
            try:
                fmt = self.formats[self.keys[i]]
                fc = fmt
                # check that the format is supported
                fi = self.fmt_characters.index(fmt)
            except:
                f_err = True
                print('unknown format : ', fc, ' set to f')
                fc = 'f'
                self.formats[self.keys[i]] = 'f'
            # now try to convert according to the format
            try:
                if fc == 'f':
                    f.append(float(k)) # create a float
                if fc == 'i':
                    f.append(int(k))   # create an integer
                if fc == 's':
                    f.append(k)        # leave as string
            except:
                f.append(k) # if not possible leave it,
                # but make a comment
                if len(self.formats) == 0 :
                    print('cannot convert : ', k, ' to float, line : ', self.line_n +1)
                else:
                    if (f_err ):
                        print('cannot convert : ', k, ' line : ', self.line_n+1,\
                            ', problem with unknown format : ', fmt)
                    else:
                        print('cannot convert : ', k, ' line : ', self.line_n+1,\
                            ', with format :', fmt)
        f.append(self.adata.index(l))
        # pdb.set_trace()
        return f

    def make_array(self):
        if self.fast:
            # use loadtxt to get the data
            data_types = [self.formats_dt[k] for k in self.keys[:-1]]
            self.data = np.loadtxt(self.filename, dtype = {'names':self.keys[:-1], 'formats':data_types})
        else:
               # create an array of dictionaries
            self.data = []
            n_data = len(self.adata)
            for self.line_n, l in enumerate(self.adata):
                f = self.parse_line(l)
                if f == None:
                    continue
                self.data.append(self.make_dict(self.keys,f) )
                if self.debug:
                    print("processed data : ", self.line_n, " out of ", n_data)

    def make_dict(self,keys,values):
        # create a dictionary from a list of keys and values
        p = []
        for j,key in enumerate(keys): # loop over keys
            try:
                p.append( (key.strip(),values[j]) ) # an array of tuples
                d=dict(p) # create a dictionary
            except:
                print('problem in data : ', key, values, j)
            #
        return d

    def make_attr(self):
        # make attributes out of the keys
        for k in self.keys:
            ka = ''
            try:
                self.__dict__[k]
                # attrtibute exists append _d
                ka = k + '_d'
                print(("{:s} exists using : {:s}".format(k, ka)))
            except:
                ka = k
            # replace periods with underscores
            if self.use_numpy:
                setattr(self, ka.replace('.','_'), np.array(self[k]))
            else:
                setattr(self, ka.replace('.','_'), self[k])

    def scale(self, key, factor):
        """
        multiply all values of key with a factor
        """
        for i, d in enumerate(self.data):
            self.data[i][key] = factor*d[key]
        #

    def sort(self, key, **kwargs):
        """
        sort the data according to the values in key
        """
        index_array = []
        for d in self.data:
            index_array.append((d[key], self.data.index(d))) # store index and values pair
        sorted_array = sorted(index_array,**kwargs) # create a sorted array
        temp = []
        for ip in sorted_array:
            temp.append(self.data[ip[1]])
        self.data = temp
        temp=[]

    def name(self):
        """
        print the filename associate with this instance

        """
        print("Input file name : ", self.filename)
    def show_keys(self):
        """
        print a list of variable names in the dictionary

        """
        print(self.keys)
    def get_keys(self):
        """

        return a list of keys

        """
        return self.keys
    def get_header(self):
        """

        return the header lines

        """
        return self.adata[self.headindex]

    def get_full_header(self):
        """

        return all line up to the header line

        """
        return self.adata[:self.headindex+1]

    def show_data(self,keylist):
        """

        print all the data corresponding to the key list::

           >>> df.show_data('key1:key2:key3')

        """
        akey=keylist.split(":")
        for l in self.data:
            ll=[]
            for i in akey:
                format = self.output_format[ self.formats[i] ]
                ll.append( format%( l[i] ) )
            print(ll)

    def check_data(self, func, data, key, *args):
        """

        function used to check if data fulfill a condition provided
        by the user. The function is assumed to return True of False.

        """
        return func( data, key, *args)

    def select_data(self, sel_func = None, sel_args = None):
        """
        returns an iterator for the data. As in get_data() a selector function and
        its arguments can be supplied:

        conditions can be applied:

        1) sel_func: a user provided function returning True or False

        2) sel_args: a list of arguments used in the sel_func function

        The iterator returned can be used as follows:

        >>> for d in df.select_data():
        >>> ...print d

        or the result can be converted to a list:

        >>> list(df.select_data())

        Using a selector function:

        Example: assume the data contain a variable (key) called 'name'
           you want to select only those data where name contains a certain substring 'sub'

        Define this function:

        >>> def myfind(data, key , what):
        >>> ...where = data[key]
        >>> ...return (str.find(where, what) >= 0)

        now you can select the data using:

        >>> list( df.select_data( myfind, 'name', 'sub') )

        """
        if sel_func == None:
            for i,data in enumerate(self.data):
                yield data
        else:
            for i,data in enumerate(self.data):
                if self.check_data(sel_func, data, *sel_args):
                    yield data
        # that's all

    def select_data_eval(self, eval_str):
        """
        returns an iterator for data selected with an *eval* expression
        stored in the string `eval_str` each dataset item is accessed using the name data::

        >>> df.select_data_eval("data['x'] >= 0.")

        returns only those data items where the value of the x-column in the file
        is larger than or equal to 0.

        """
        for i,data in enumerate(self.data):
            if eval( eval_str):
                yield data
        # that's all

    def eval_data(self, eval_str):
        """
        iterator over all data evaluating an expression contained in eval_str
        """
        for i,data in enumerate(self.data):
            yield eval(eval_str)
        # that's all

    def get_data(self,key, sel_func = None, sel_args = None):
        """

        return all data for `key` subject to the results of a selector function.

        one can define a selector function (`sel_func`) using the arguments stored in `sel_args`.
        This function is evaluated for each data record and only
        those data are returned for which the condition is fulfilled.

        1) sel_func: a user provided function returning True or False

        2) sel_args: a list of arguments used in the sel_func function

        Example::
           assume the data contain a variable (key) called 'name'
           you want to select only those data where name contains a certain substring 'Jo'

        Define this function::

        >>> def myfind(data, key , what):
        >>> ...where = data[key]
        >>> ...return (str.find(where, what) >= 0)

        now you can select the data using::

        >>> df.get_data('name',myfind, ['name','Jo'] )

        This should return a list of names containing the substring 'Jo'

        """
        array=[]
        for l in self.data:
            if sel_func == None:
                array.append(l[key])
            else :
                if self.check_data( sel_func, l, *sel_args):
                    array.append(l[key])
        if self.use_numpy:
            return np.array(array)
        else:
            return array

    def get_data_eval(self,key,eval_str): # return data for a  key
        """

        return all data for the `key` under the condition that the expression in `eval_str`
        is True.

        """
        array=[]
        for l in self.select_data_eval(eval_str):
            array.append(l[key])
        if self.use_numpy:
            return np.array(array)
        else:
            return array

    def get_data_list(self,keylist, sel_func = None, sel_args = None):

        """

        return all the data corresponding to the key list
        as follows::

           >>> a = df.get_data_list('key1:key2:key3')

           >>> a = df.get_data_list('key1:key2:key3', myfind, ['name','J'])

        return those data where the `name` values contain the character `J`

        *a* contains the list of data

        """

        akey=keylist.split(":")
        dd = []
        if sel_func == None:
            for l in self.data:
                ll=[]
                for i in akey:
                    ll.append( l[i] )
                dd.append(ll)
            if self.use_numpy:
                return np.array(dd)
            else:
                return dd
        else:
            for l in self.data:
                if self.check_data(sel_func, l, *sel_args):
                    ll=[]
                    for i in akey:
                        ll.append( l[i] )
                    dd.append(ll)
            if self.use_numpy:
                return np.array(dd)
            else:
                return dd

    def get_data_list_eval(self,keylist, eval_str):
        """

        similar function as select_data_eval but it only returns the
        values for the keys defined in `keylist`.

        """
        akey=keylist.split(":")
        dd = []

        for l in self.select_data_eval(eval_str):
            ll=[]
            for i in akey:
                ll.append( l[i] )
            dd.append(ll)
        if self.use_numpy:
            return np.array(dd)
        else:
            return dd

    def show_all_data(self):
        """

        print all data and keys stored

        """
        for l in self.data:
            print(l)
    def get_all_data(self):
        """

        return a list of all data stored

        """
        return self.data
    def write_complete_header(self,fp):
        """
        write entire header including comments and internal parameters to file with
        handle fp. Example:

        fp = open('mydile.data','w')

        datafile.write_comlete_header(fp)

        """
        for i,l in enumerate(self.adata[0:self.headindex]):
            fp.write(l+'\n')
        # write header line
        h = '#! '
        n = 0
        for k in self.keys[:-1]:
            h = h + k+'['+'%s,%d'%(self.formats[k],n)+']/ '
            n = n +1
        fp.write(h+'\n')
    def update_header(self):
        """
        update header line of this data file including format. This
        makes sure that the header line is in sync with the dictionary keys
        """
        # write header line
        h = '#! '
        n = 0
        for k in self.keys[:-1]:
            h = h + k+'['+'%s,%d'%(self.formats[k],n)+']/ '
            n = n +1
        # update data arrays
        self.adata[self.headindex]=h
        self.headerline = h
    def write_header(self,fp):
        """
        write only header line of this data file into file fp, including format
        """
        # write header line
        h = '#! '
        n = 0
        for k in self.keys[:-1]:
            h = h + k+'['+'%s,%d'%(self.formats[k],n)+']/ '
            n = n +1
        fp.write(h+'\n')
    def write_line(self,fp, i):
        """
        write data line i into file fp
        """
        l = ''
        # check array boundary
        ndata = len(self.data)
        if (i<0) :
            print('index below lower bound, set to 0')
            i = 0
        if (i>=ndata):
            print('index above upper bound, set to : ', ndata-1)
        for k in self.keys[:-1]:
            format = self.output_format[ self.formats[k] ]
            l = l + format%(self.data[i][k])
        fp.write(l+'\n')
    def write_all(self,fp, complete_header = False):
        """
        write all data to a file associated to fp

        if complete_header = True include the complete header including all comments
        """
        if complete_header:
            self.write_complete_header(fp)
        else:
            self.write_header(fp)
        for i in range(len(self.data)):
            self.write_line(fp, i)
        fp.close()
    #   new function to add a key
    def write_selected(self,fp,index_list, complete_header = False):
        """
        write a new datafile with an identical header but enter only
        those data with an index given in index_list. If the index does not
        exist print a message and skip it.

        if complete_header = True include the complete header including all comments

        """
        if complete_header:
            self.write_complete_header(fp)
        else:
            self.write_header(fp)
        for i in index_list:
            try:
                self.write_line(fp, i)
            except:
                print('cannot write index : ', i, ' does it exist ?')
        fp.close()

    def add_key(self,key,format='f'):
        """
        add a new key, this is useful if you want to add new data
        assign new values using a loop like::

        >>> df.add_key('newkey',format='i')

        The new data need to be stored as follows::

        >>> for d in df.data:
        >>>     d['newkey']= some_new_value

        where df is the datafile and 'newkey' the new key. The new data set
        can be saved using ``fp = open('new_file','w')`` and ``df.write_all(fp)``
        """
        # check if key already exists
        try:
            self.keys.index(key)
            print(key, " exists already !")
            return
        except:
            for d in self.data:
                # set initial value to none
                d[key] = None
        # add the new key at the end of the keys list but before indx
        self.keys.insert(len(self.keys)-1,key)
        # add the format information
        try:
            fi = self.fmt_characters.index(format)
        except:
            print('unknown format : ', format, ' use f instead')
            format = 'f'
        self.formats[key] =  format
        self.update_header()
    # end
    def delete_key(self,key):
        """
        remove a key and all values associated with it
        """
        # remove the data for the key
        for d in self.data:
            del d[key]
        # remove the format for this key
        fmt = self.formats.pop(key)
        # now remove the key from the key list
        self.keys.remove(key)
        print('key : ', key , ' and format : ', fmt, ' removed')
        self.update_header()
    # end

    def add_data(self, keys,line):
        """

        add data to the data file

        d.add_data('x:y:z','xval yval zval')

        xval, yval, zval are the values as they would be netered in a data file line

        both arguments are strings !

        IMPORTANT: add the values to all current lines of the data file !

        """

        key_l = keys.split(':')
        # check the key list
        key_index=[]
        for k in self.keys[:-1]:  # skip the indx
            try:
                ki = key_l.index(k)
                key_index.append(ki)
            except:
                print(k, ' does not exist !, nothing added')
                return
        # sort the argumens according to the sequence of the keys
        lsp = line.split()
        fline = [lsp[i] for i in key_index]
        line = ''
        for fl in fline: line += fl+' '
        self.adata.append(line.strip())
        print(line)
        """
        f = self.parse_line(line.strip())
        self.data.append(self.make_dict(self.keys,f) )
        """
        self.make_array()

    def add_header_comment(self, text):
        """
        add a comment line to the header. The # at the beginning is automatically added !

        * for a normal comment start with a space
        * to add a parameter start the comment with a \
        """
        self.adata.insert(self.headindex, '#' + text)
        self.headindex += 1

    def add_parameter(self, name, value):
        """
        add a parameter to the comment section
        """
        cmd = '\ {} = {} '.format(name, value)
        self.add_header_comment(cmd)

    def save(self, file = None):
        """

        Save the current datafile.

        With the keyword: *file* = 'new_file'

        the datafile will be written into the new_file name

        """
        if file == None:
            o = open(self.filename,'w')
        else:
            o = open( file, 'w')
        self.write_all(o, complete_header = True)

    def make_fmt(self):
        fmt = ''
        for k in self.keys[:-1]:
            fmt += format_dict[self.formats[k]] + ', '
        return fmt.strip()[:-1]

    def make_csv(self):
        keys = self.keys[:-1]
        csv = [ ''.join([k + ',' for k in keys]) ]
        for l in self.data:
            fmt = self.make_fmt()
            data = [l[k] for k in keys]
            ll = fmt.format(*data)
            csv.append(ll)
        return csv

    def write_csv(self, f):
        """

        save the current file as a csv file

        f : file name  to be used

        """
        o = open(f,'w')
        csv = self.make_csv()
        for l in csv:
            o.write(l + '\n')
        o.close()

#


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     notebook_metadata_filter: all
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.2'
#       jupytext_version: 1.2.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
#   language_info:
#     codemirror_mode:
#       name: ipython
#       version: 3
#     file_extension: .py
#     mimetype: text/x-python
#     name: python
#     nbconvert_exporter: python
#     pygments_lexer: ipython3
#     version: 3.7.6
# ---

# %% [markdown]
# # ConsPortfolioModel Documentation
#
# [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/econ-ark/DemARK/master?filepath=notebooks%2FConsPortfolioModelDoc.ipynb)

# %% {"code_folding": []}
# Setup stuff

import HARK.ConsumptionSaving.ConsPortfolioModel as cpm
import HARK.ConsumptionSaving.ConsumerParameters as param
import copy
import matplotlib.pyplot as plt
import numpy as np
import warnings
warnings.filterwarnings('ignore')
from HARK.ConsumptionSaving.ConsPortfolioModel import PortfolioSolution

# %% [markdown]
# We implement three different ways to allow portfolio choice.
#            The agent can choose
#               * any portfolio share ('continuous choice')
#               * only a specified set of portfolio shares ('discrete choice')
#                 * With probability 1 (agent always gets to choose)
#                 * With probability 0 < p < 1 (stochastic chance to choose)
#
#            We allow two choices for the description of the
#            distribution of the stochastic variable:
#            1. A generic discrete probability distribution
#               * Nodes and their probabilities are specified
#            2. A true lognormal distribution
#               * The mean return and the standard deviation are specified
#
#            In the discrete portfolio shares case, the user also must
#            input a function that *draws* from the distribution in drawRiskyFunc
#
#            Other assumptions:
#               * distributions are time constant
#               * probability of being allowed to reoptimize is time constant
#                  * If p < 1, you must specify the PortfolioSet discretely
#

# %% {"code_folding": []}
# Set up the model
# Parameters from Mehra and Prescott (1985):
Avg = 1.08 # equity premium
Std = 0.20 # standard deviation of rate-of-return shocks

RiskyDstnFunc = cpm.RiskyDstnFactory(RiskyAvg=Avg, RiskyStd=Std)       # Generates nodes for integration
RiskyDrawFunc = cpm.LogNormalRiskyDstnDraw(RiskyAvg=Avg, RiskyStd=Std) # Function to generate draws from a lognormal distribution

init_portfolio = copy.copy(param.init_idiosyncratic_shocks) # Default parameter values for inf horiz model - including labor income with transitory and permanent shocks
init_portfolio['approxRiskyDstn'] = RiskyDstnFunc
init_portfolio['drawRiskyFunc']   = RiskyDrawFunc
init_portfolio['RiskyCount']      = 2   # Number of points in the approximation; 2 points is minimum
init_portfolio['RiskyShareCount'] = 25  # How many discrete points to allow for portfolio share
init_portfolio['Rfree']           = 1.0 # Riskfree return factor (interest rate is zero)
init_portfolio['CRRA']            = 6.0 # Relative risk aversion

# Uninteresting technical parameters:
init_portfolio['aXtraMax']        = 100
init_portfolio['aXtraCount']      = 50
init_portfolio['BoroCnstArt']     = 0.0 # important for theoretical reasons
# init_portfolio['vFuncBool'] = True # We do not need value function for purposes here

init_portfolio['DiscFac'] = 0.92 # Make them impatient even wrt a riskfree return of 1.08

# Create portfolio choice consumer type
pcct = cpm.PortfolioConsumerType(**init_portfolio)

# %% {"code_folding": []}
# Solve the model under the given parameters

pcct.solve()
aMin = 0   # Minimum ratio of assets to income to plot
aMax = 10  # Maximum ratio of assets to income to plot
aPts = 100 # Number of points to plot

# Campbell-Viceira (2002) approximation to optimal portfolio share in Merton-Samuelson (1969) model
pcct.MertSamCampVicShare = pcct.RiskyShareLimitFunc(RiskyDstnFunc(init_portfolio['RiskyCount']))

# Define grid of points on which to plot
eevalgrid = np.linspace(0,aMax,aPts) # range of values of assets for the plot

# Plot portfolio share by wealth-to-income ratio
plt.plot(eevalgrid, pcct.solution[0].RiskyShareFunc[0][0](eevalgrid))
plt.axhline(pcct.MertSamCampVicShare, c='r') # The Campbell-Viceira approximation
plt.ylim(0,1.05)
plt.text((aMax-aMin)/4,0.45,r'$\uparrow $ limit as  $m \uparrow \infty$',fontsize = 22,fontweight='bold')
plt.show()

# %% {"code_folding": [0]}
# Simulate 20 years of behavior for a set of consumers initially distributed widely
SimPer = 20

pcct.track_vars = ['aNrmNow', 't_age', 'RiskyShareNow']
pcct.T_sim = SimPer
pcct.initializeSim()
pcct.simulate()
pcct.RiskyShareNow_hist

from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_zlim(pcct.MertSamCampVicShare,1.0)
ax.scatter(pcct.aNrmNow_hist, pcct.t_age_hist, pcct.RiskyShareNow_hist)
plt.show()

# The consumers are impatient and so even if they start rich they end up as buffer stock savers
# But with all of their buffer stock savings in the stock market

# %%
# Solve the specialized / simple version for which Campbell-Viceira (2002) is a good approximation
# (as wealth approaches infinity)
# This is the version for which Campbell and Viceira provide an approximate formula
# assuming lognormally distributed shocks

init_lognormportfolio = copy.deepcopy(init_portfolio) # Use same parameter values
init_lognormportfolio['RiskyAvg']   = Avg
init_lognormportfolio['RiskyStd']   = Std
init_lognormportfolio['RiskyCount'] = 11 # Eleven is enough points to do justice to the distribution
lnpcct = cpm.LogNormalPortfolioConsumerType(**init_lognormportfolio)

lnpcct.solve()
lnpcct.MertSamCampVicShare = lnpcct.RiskyShareLimitFunc(RiskyDstnFunc(init_portfolio['RiskyCount']))

fig = plt.figure()
plt.plot(eevalgrid, lnpcct.solution[0].RiskyShareFunc[0][0](eevalgrid))
plt.axhline(lnpcct.MertSamCampVicShare, c='r')
plt.ylim(0,1.05)
plt.text((aMax-aMin)/4,lnpcct.MertSamCampVicShare-0.1,r'$\uparrow $ limit as  $m \uparrow \infty$',fontsize = 22,fontweight='bold')
plt.show()

# %%
# Again simulate a few periods -- similar results (poor buffer stock savers put all their money in the stock market)
lnpcct.track_vars = ['aNrmNow', 't_age', 'RiskyShareNow']
lnpcct.T_sim = SimPer
lnpcct.initializeSim()
lnpcct.simulate()
lnpcct.RiskyShareNow_hist

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(lnpcct.aNrmNow_hist, lnpcct.t_age_hist, lnpcct.RiskyShareNow_hist)
ax.set_zlim(lnpcct.MertSamCampVicShare,1.0)
plt.show()

# %%
# Version where only discrete values of portfolio share of risky assets are allowed

init_portfolio_prb = copy.deepcopy(init_portfolio)

init_portfolio_prb['AdjustPrb'] = 1.0
init_portfolio_prb['PortfolioDomain'] = cpm.DiscreteDomain([0.0, 0.5, 0.6, 1.0])
pcct_prb = cpm.PortfolioConsumerType(**init_portfolio_prb)
pcct_prb.solve()

eevalgrid = np.linspace(0,100,100)
plt.plot(eevalgrid, pcct_prb.solution[0].RiskyShareFunc[0][0](eevalgrid))
plt.axhline(pcct_prb.RiskyShareLimitFunc(RiskyDstnFunc(init_portfolio['RiskyCount'])), c='r')
plt.ylim(0,1.05)
plt.show()

# %%
# Version where you can choose your portfolio share only with some 0 < p < 1

init_portfolio_prb = copy.deepcopy(init_portfolio)

init_portfolio_prb['AdjustPrb'] = 0.5
init_portfolio_prb['PortfolioDomain'] = cpm.DiscreteDomain([0.0, 0.6, 1.0])
pcct_prb = cpm.PortfolioConsumerType(**init_portfolio_prb)

pcct_prb.solve()

plt.plot(eevalgrid, pcct_prb.solution[0].RiskyShareFunc[0][0](eevalgrid))
plt.show()

pcct_prb.track_vars = ['aNrmNow', 't_age', 'RiskyShareNow', 'CantAdjust']
pcct_prb.T_sim = 10
pcct_prb.AgentCount = 30
pcct_prb.initializeSim()
pcct_prb.simulate()
pcct_prb.RiskyShareNow_hist

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(pcct_prb.aNrmNow_hist, pcct_prb.t_age_hist, pcct_prb.RiskyShareNow_hist)
plt.show()


#!/usr/bin/env python
# coding: utf-8

# # Handwritten mathematical sumbols

# ## Data preprocessing

# Hand_written
#
#     --> handwritten_mathematical_model.ipynb
#     --> hasy-data
#     --> hasy-data-labels
#     --> symbols.csv
#     --> hasyv2
#         --> hasy-data
#         --> hasy-data-labels
#         --> symbols.csv
#         --> verification-task
#     --> verification-task
#
# Note:- You can use hasyv2 but for this I have not used that

#Just to check
from IPython.display import Image
Image(url='hasy-data/v2-00010.png')

#importing libraries for preprocessing task
import csv
from PIL import Image as pil_image
import keras.preprocessing.image

#load all images and save there classes

imgs = []
classes = []
with open('hasy-data-labels.csv') as csvFile:
    csvReader = csv.reader(csvFile)
    i = 0
    for row in csvReader:
        if i > 0:
            img = keras.preprocessing.image.img_to_array(pil_image.open(row[0]))
            img /= 255.0
            imgs.append((row[0], row[2], img))
            classes.append(row[2])
        i += 1
print("Total number of images: ", len(imgs)) #print the total number of images

# Randomly split the data into training and test
#
# 80% -> train
#
# 20% -> test

import random
random.shuffle(imgs)
split_idx = int(0.8*len(imgs))
train = imgs[:split_idx]
test = imgs[split_idx:]

#Later take 20% from train for validation

import numpy as np

train_input = np.asarray(list(map(lambda row: row[2], train)))
test_input = np.asarray(list(map(lambda row: row[2], test)))

train_output = np.asarray(list(map(lambda row: row[1], train)))
test_output = np.asarray(list(map(lambda row: row[1], test)))

# learn about oneHotEncoder
# https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.OneHotEncoder.html
#our labels are in the text format therefore using oneHotEncoder to convert them

from sklearn.preprocessing import LabelEncoder, OneHotEncoder

label_encoder = LabelEncoder()
integer_encoded = label_encoder.fit_transform(classes)

onehot_encoder = OneHotEncoder(sparse = False)
integer_encoded = integer_encoded.reshape(len(integer_encoded), 1)
onehot_encoder.fit(integer_encoded)

train_output_int = label_encoder.transform(train_output)
train_output = onehot_encoder.transform(train_output_int.reshape(len(train_output_int), 1))
test_output_int = label_encoder.transform(test_output)
test_ouput = onehot_encoder.transform(test_output_int.reshape(len(test_output_int), 1))

num_classes = len(label_encoder.classes_)

print("Number of classes", num_classes)
print("Shape of the input image: ", np.shape(train_input[0]))

# ## Model

#importing libraries

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import RMSprop, Adam
from tensorflow.keras.layers import (
    Dense,
    Dropout,
    Flatten,
    Conv2D,
    MaxPooling2D
)

# This is very basic CNN architecture with some convolution layer followed by max pooling layer
# which is then followed by flatten and some dense layer.

#model architecture

model = Sequential([
    Conv2D(32, (3,3), activation='relu', input_shape = np.shape(train_input[0])),
    MaxPooling2D(pool_size = (2,2)),
    Conv2D(32, (3,3), activation='relu'),
    MaxPooling2D(pool_size = (2,2)),
    Flatten(),
    Dense(1024, activation = 'tanh'),
    Dropout(0.5),  #to reduce overfitting
    Dense(num_classes, activation='softmax')
])

#compile model
learning_rate = 1e-5 #1e-2 (0.01)

model.compile(
            optimizer = 'adam', #Adam(lr = learning_rate)
            loss = 'categorical_crossentropy',
            metrics = ['accuracy'])

model.summary()

# Training
BATCH_SIZE = 32
EPOCHS = 15 #10
VALIDATION_SPLIT = 0.2

history = model.fit(train_input, train_output,
                   batch_size = BATCH_SIZE,
                   epochs = EPOCHS,
                   verbose = 1,
                   validation_split = VALIDATION_SPLIT)

#Accuracy I achieved was nearly 82% after 10th epoch (Val_accuracy -> 77%)
#tuning hyparameters and changing model architecture will help you gain much more accuracy and reduce overfitting

# Plotting the graphs:
#     1. Training and validation accuracy
#     2. Training and validation loss

import matplotlib.pyplot as plt

acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
loss = history.history['loss']
val_loss = history.history['val_loss']

epochs = range(len(acc))

fig = plt.subplots(figsize=(30, 10))

plt.subplot(1, 2, 1)
plt.plot(acc)
plt.plot(val_acc)
plt.title('Training and Validation Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend(['Training Accuracy', 'Validation Accuracy'])

plt.subplot(1, 2, 2)
plt.plot(loss)
plt.plot(val_loss)
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend(['Training Loss', 'Validation Loss'])

plt.show()

# ## Saving model
model.save('mathSymbolsPredictor.model')
np.save('classes.npy', label_encoder.classes_)

# ## New predictions

import keras.models
#model_ = keras.models.load('mathSymbolsPredictor.model')

label_encoder2 = LabelEncoder()
label_encoder2.classes_ = np.load('classes.npy')

PATH = 'hasy-data/v2-00100.png'

image = keras.preprocessing.image.img_to_array(pil_image.open(PATH))
image /= 255.
predicted = (model.predict(image.reshape(1, 32, 32, 3)))
inverted = label_encoder2.inverse_transform([np.argmax(predicted)])
print(inverted[0], np.max(predicted))
(Image(url=PATH))

# # Convert your model into .tflite
# ## To deploy it on an application

RPS_SAVED_MODEL = "rps_saved_model"
tf.saved_model.save(model, RPS_SAVED_MODEL)

get_ipython().run_cell_magic('bash', '-s $RPS_SAVED_MODEL', 'saved_model_cli show --dir $1 --tag_set serve --signature_def serving_default')

loaded = tf.saved_model.load(RPS_SAVED_MODEL)
print(list(loaded.signatures.keys()))
infer = loaded.signatures["serving_default"]
print(infer.structured_input_signature)
print(infer.structured_outputs)
converter = tf.lite.TFLiteConverter.from_saved_model(RPS_SAVED_MODEL)
converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
tflite_model = converter.convert()

import tensorflow as tf
tflite_model_file = 'handwritten_mathematical_symbol_model_v1.tflite'


import os

import numpy as np
import matplotlib.pyplot as plt

from saf.euler1d.linear.asciireader import ASCIIReader

from helpers import FIGSIZE_SIX_SUBPLOTS as figsize, savefig

# ---
# ## Helper functions

# In[2]:

def znd_read_data():
    output_dir = '_output'
    dirs = os.listdir(output_dir)
    dirs.sort()

    x = []
    rho = []
    u = []
    p = []
    lamda = []

    for d in dirs:
        dirname = os.path.join(output_dir, d)

        r = ASCIIReader(dirname)

        znd_data = r.get_znd_data()
        x.append(znd_data['x'])
        rho.append((d, znd_data['rho']))
        u.append((d, znd_data['u_lab']))
        p.append((d, znd_data['p']))
        lamda.append((d, znd_data['lamda']))

    return x, rho, u, p, lamda

def _znd_plot_quantity(x, quantity, axis, k):
    line_styles = ['-', '--', ':', '-o', '-*']
    assert k >= 0
    assert k < len(line_styles)
    if k != 3 and k != 4:
        axis.plot(x, quantity, line_styles[k])
    else:
        ls = line_styles[k]
        s1 = ls[0]
        s2 = ls[1]
        R = 10
        line, = axis.plot(x, quantity, s1)
        color = line.get_color()
        axis.plot(x[::R], quantity[::R], s2, color=color, markeredgecolor=color)

def znd_plot_data(x, rho, u, p, lamda):
    # Number of rows and columns in figures
    m, n = 3, 2
    E_ACT = 30.0
    X_LIM = -8

    fig, axes = plt.subplots(nrows=m, ncols=n, figsize=figsize)

    assert len(x) == len(rho)
    assert len(x) == len(u)
    assert len(x) == len(p)
    assert len(x) == len(lamda)

    # Density
    ax = axes[0, 0]
    for k, __ in enumerate(rho):
        cur_x = x[k]
        cur_rho = rho[k][1]
        cur_rho = cur_rho/cur_rho[-1]

        _znd_plot_quantity(cur_x, cur_rho, ax, k)
    ax.set_ylabel(r'$\bar{\rho}/\bar{\rho}_{\mathrm{s}}$')
    ax.set_xlabel(r'$x$')
    ax.set_xlim((X_LIM, 0))
    ax.set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
    ax.set_ylim((0, 1.0))

    # Progress variable
    ax = axes[0, 1]
    for k, __ in enumerate(lamda):
        cur_x = x[k]
        cur_lamda = lamda[k][1]

        _znd_plot_quantity(cur_x, cur_lamda, ax, k)
    ax.set_ylabel(r'$\bar{\lambda}$')
    ax.set_xlabel(r'$x$')
    ax.set_xlim((X_LIM, 0))
    ax.set_ylim((0, 1.0))

    # Pressure
    ax = axes[1, 0]
    for k, __ in enumerate(p):
        cur_x = x[k]
        cur_p = p[k][1]
        cur_p = cur_p/cur_p[-1]

        _znd_plot_quantity(cur_x, cur_p, ax, k)
    ax.set_ylabel(r'$\bar{p}/\bar{p}_{\mathrm{s}}$')
    ax.set_xlabel(r'$x$')
    ax.set_xlim((X_LIM, 0))
    ax.set_ylim((0.5, 1.0))

    # Velocity
    ax = axes[1, 1]
    for k, __ in enumerate(u):
        cur_x = x[k]
        cur_label = u[k][0]
        cur_u = u[k][1]
        cur_u = cur_u/cur_u[-1]

        _znd_plot_quantity(cur_x, cur_u, ax, k)
    ax.set_ylabel(r'$\bar{u}/\bar{u}_{\mathrm{s}}$')
    ax.set_xlabel(r'$x$')
    ax.set_xlim((X_LIM, 0))
    ax.set_ylim((0.5, 1.0))

    # Temperature
    ax = axes[2, 0]
    for k, __ in enumerate(p):
        cur_x = x[k]
        cur_p = p[k][1]
        cur_p = cur_p / cur_p[-1]
        cur_rho = rho[k][1]
        cur_rho = cur_rho / cur_rho[-1]
        cur_T = cur_p / cur_rho

        _znd_plot_quantity(cur_x, cur_T, ax, k)
    ax.set_ylabel(r'$\bar{T}/\bar{T}_{\mathrm{s}}$')
    ax.set_xlabel(r'$x$')
    ax.set_xlim((X_LIM, 0))
    ax.set_ylim((0.95, 2.8))

    # Reaction rate
    ax = axes[2, 1]
    for k, __ in enumerate(p):
        cur_x = x[k]
        cur_p = p[k][1]
        cur_rho = rho[k][1]
        cur_lamda = lamda[k][1]
        r = (1 - cur_lamda) * np.exp(-E_ACT*cur_rho/(cur_p))
        r = r / np.max(r)

        _znd_plot_quantity(cur_x, r, ax, k)
    ax.set_ylabel(r'$\bar{\omega}/\bar{\omega}_{\mathrm{max}}$')
    ax.set_xlabel(r'$x$')
    ax.set_xlim((X_LIM, 0))
    ax.set_ylim((0, 1.0))

    fig.tight_layout(pad=0.1, h_pad=0.55)
    savefig('znd-solutions.pdf')

# ---
# ## Reading data

# In[3]:

x, rho, u, p, lamda = znd_read_data()

# ---
# ## Plotting ZND solutions

# The figure below contains profiles of different quantities of ZND solution.
# Solid lines is for $Q=0.5$, dashed line is for $Q=1.0$, dotted line is for $Q=10$, solid line with circle marker is for $Q=50$, solid line with star marker is for $Q=100$.

# In[4]:

znd_plot_data(x, rho, u, p, lamda)