import Base.rand
import Base.Filesystem
import LinearAlgebra.det
import LinearAlgebra.dot
import LinearAlgebra.norm
import Plots
import JSON
using Logging
using LoggingExtras
import Glob.glob

#------------------------------------------------------------------------------
# Datastructures and functions related to delaunay triangulation
# - Edge: structure used to describe an edge
# - Triangule: structure used to describe a triangle
# - create_nodes: create a set of random nodes
# - in_circumcircle: find if a node is in the circumcircle of a triangle
# - is_clockwise: find if a set of 3 nodes is ordered clockwise
# - create_supertriangle: create a triangle containing all the nodes
# - require_flip: tell if two neighbor triangles need to be flipped because
#   they don't follow the criterion on opposite angles
# - compute_angle: compute the angle (in radians) between two edges
# - get_tri_node_ids: get the ids of the nodes of a triangle
# - get_triangulation_nodes: get the ids of the nodes used in a triangulation
# - flip_triangulation!: flip the triangles of the triangulation when needed
# - save_step: save the details of a step of the triangulation to a file
# - plot_step: create a visualization for a step of the triangulation
# - inside_triangle: find if a node is located inside a triangle
# - triangulate: create the triangulation of a set of nodes
#------------------------------------------------------------------------------

struct Edge
    # Note: possible improvement using static arrays to nodes and in_triangles
    #       in_triangles::StaticArrays.SVector{2,Int}

    # indices of the nodes
    nodes::Array{Int, 1}
    # indices of the triangles containing this edge
    in_triangles::Array{Int, 1}
end

struct Triangle
    # Note: possible improvement using static arrays to nodes and in_triangles
    #       edges::StaticArrays.SVector{3,Int}

    # indices of the edges of the triangle
    edges::Array{Int, 1}
end

function create_nodes(number::Int, x_bounds, y_bounds)
    # number:   number of nodes created
    # x_bounds: X bounds of nodes - list of 2 reals
    # y_bounds: Y bounds of nodes - list of 2 reals
    nodes = [[rand(Float64), rand(Float64)] for i in 1:number]
    # apply X and Y bounds
    for i=1:number
        nodes[i][1] = nodes[i][1]*(x_bounds[2] - x_bounds[1]) + x_bounds[1]
        nodes[i][2] = nodes[i][2]*(y_bounds[2] - y_bounds[1]) + y_bounds[1]
    end
    return nodes
end

function in_circumcircle(triangle_ids, id, nodes)
    # Tell if node with if "id" is in the circumcircle of triangle
    # made of the nodes of "triangle_ids"
    # - triangle_ids : id of the nodes of the triangle - list of 3 int
    # - id : id of the node investigated - int
    # - nodes : coordinates of the nodes - list of [x, y] arrays of floats
    if (length(triangle_ids) != 3)
        error("Argument triangle_ids must be a list of 3 int")
    end
    # vertices of the triangle
    A = nodes[triangle_ids[1]]
    B = nodes[triangle_ids[2]]
    C = nodes[triangle_ids[3]]
    P = nodes[id]
    # switch vertices if the triangle is clockwise
    if (is_clockwise(A, B, C))
        B, C = C, B
    end
    # create the matrix (see https://fr.wikipedia.org/wiki/Triangulation_de_Delaunay)
    mat = zeros(Float64, 3, 3)
    mat[1, 1] = A[1] - P[1]
    mat[2, 1] = B[1] - P[1]
    mat[3, 1] = C[1] - P[1]
    mat[1, 2] = A[2] - P[2]
    mat[2, 2] = B[2] - P[2]
    mat[3, 2] = C[2] - P[2]
    mat[1, 3] = A[1]^2 - P[1]^2 + A[2]^2 - P[2]^2
    mat[2, 3] = B[1]^2 - P[1]^2 + B[2]^2 - P[2]^2
    mat[3, 3] = C[1]^2 - P[1]^2 + C[2]^2 - P[2]^2
    return (det(mat) > 0)
end

function is_clockwise(A, B, C)
    # A, B & C: list of 2 reals representing 2D points
    AB = B - A
    AC = C - A
    determinant = AB[1]*AC[2] - AB[2]*AC[1]
    return (determinant < 0)
end

function create_supertriangle(nodes)
    # Create a triangle containing all the nodes
    # - nodes : coordinates of the nodes - list of [x, y] arrays of floats

    # bounds of the nodes
    nodes_x = [a_node[1] for a_node in nodes]
    nodes_y = [a_node[2] for a_node in nodes]
    x_bounds = [minimum(nodes_x), maximum(nodes_x)]
    y_bounds = [minimum(nodes_y), maximum(nodes_y)]

    # take a margin (relative to the X & Y bounds)
    m = 0.05
    min_length = 1e-4
    lx = x_bounds[2] - x_bounds[1] + min_length
    x_bounds = [x_bounds[1]-m*lx, x_bounds[2]+m*lx]
    ly = y_bounds[2] - y_bounds[1] + min_length
    y_bounds = [y_bounds[1]-m*ly, y_bounds[2]+m*ly]

    # get the circle containing the nodes
    right_top = [x_bounds[2], y_bounds[2]]
    center = [0.5*(x_bounds[1]+x_bounds[2]), 0.5*(y_bounds[1]+y_bounds[2])]
    radius = sqrt( (right_top[1]-center[1])^2 + (right_top[2]-center[2])^2 )

    # build an equilateral triangle with edges adjacent to the circle
    # base's center = lowest point of the circle on the Y axis
    base_mid = [center[1], center[2]-radius]
    # radius projection on an edge other than the base
    edge_mid = [center[1]+radius*cos(pi/6), center[2]+radius*sin(pi/6)]
    edge_l = 4*abs(edge_mid[1] - base_mid[1])
    tri_lower_right = [base_mid[1]+edge_l/2, base_mid[2]]
    tri_lower_left = [base_mid[1]-edge_l/2, base_mid[2]]
    tri_top = [base_mid[1], base_mid[2]+edge_l*sin(pi/3)]

    return [tri_lower_right, tri_top, tri_lower_left]
end

function require_flip(tri1_id, tri2_id, triangles, edges, nodes)
    # Evaluate if the triangles must be flipped, that is to say if the opposite angles don't
    # match the delaunay criterion (sum <= 180°)
    # - tri1_id : id of the first triangle
    # - tri2_id : id of the second triangle
    # - triangles : array(n) containing all the triangles
    # - edges : array(n) containing all the edges
    # - nodes : coordinates of the all nodes - list of [x, y] arrays of floats

    triangles_pair = [triangles[tri1_id], triangles[tri2_id]]

    # common_edge
    common_edge = 0
    for triangle1_edge in triangles_pair[1].edges
        if triangle1_edge in triangles_pair[2].edges
            common_edge = triangle1_edge
            break
        end
    end
    if common_edge == 0
        error("Common edge not found")
    end

    # angle opposite to the common edge
    angles = [NaN, NaN]
    for i in [1, 2]
        angle_edges = [edges[e] for e in triangles_pair[i].edges if e != common_edge]
        angle_node = intersect(angle_edges[1].nodes, angle_edges[2].nodes)[1]

        # ensure that the first point of both edges is their common point
        if (angle_edges[1].nodes[1] != angle_node)
            angle_edges[1] = Edge(reverse(angle_edges[1].nodes), [])
        end
        if (angle_edges[2].nodes[1] != angle_node)
            angle_edges[2] = Edge(reverse(angle_edges[2].nodes), [])
        end

        # compute angle
        angles[i] = compute_angle(angle_edges[1], angle_edges[2], nodes)
    end
    return (abs(angles[1]) + abs(angles[2]) > pi)
end

function compute_angle(edge1, edge2, nodes)
    # - edge1 : first edge (type Edge)
    # - edge2 : second edge (type Edge)
    # - nodes : coordinates of the all nodes - list of [x, y] arrays of floats
    p1 = nodes[edge1.nodes[1]]
    p2 = nodes[edge1.nodes[2]]
    v1 = [p2[1]-p1[1], p2[2]-p1[2]]
    v1 = v1/norm(v1)

    p1 = nodes[edge2.nodes[1]]
    p2 = nodes[edge2.nodes[2]]
    v2 = [p2[1]-p1[1], p2[2]-p1[2]]
    v2 = v2/norm(v2)

    return acos(dot(v1, v2))
end

function get_tri_node_ids(triangle_id, triangles, edges)
    # - triangle_id : id of the triangle
    # - triangles : array(n) containing all the triangles
    # - edges : array(n) containing all the edges
    triangle_edges = [edges[edge_id] for edge_id in triangles[triangle_id].edges]
    all_nodes = [e.nodes for e in triangle_edges]
    return unique(Iterators.flatten(all_nodes))
end

function get_triangulation_nodes(edges)
    # - edges : array(n) containing all the edges
    node_ids = unique(Iterators.flatten([e.nodes for e in values(edges)]))
    return node_ids
end

function flip_triangulation!(triangles, edges, nodes, do_save_steps, step_file_pattern; init_flip_step=0)
    # Flip the triangles of the triangulation that don't passe delaunay's criterion
    # This function modifies the triangles and edges input arguments
    # - triangles : array(n) containing all the triangles
    # - edges : array(n) containing all the edges
    # - nodes : coordinates of the all nodes - list of [x, y] arrays of floats

    flip_step = init_flip_step
    nb_edges_flipped = 0
    for edge_id in keys(edges)
        common_edge = edges[edge_id]
        common_nodes = common_edge.nodes
        # find if neighbor triangles separated by this edge must be flipped
        neighbors_id = common_edge.in_triangles
        if 0 in neighbors_id
            # edge is on the border of the triangulation: no flip
            continue
        elseif require_flip(neighbors_id[1], neighbors_id[2], triangles, edges, nodes)
            flip_step = flip_step + 1
            nb_edges_flipped = nb_edges_flipped + 1
            # do flip
            # 1. analyse the neighbor triangles: opposite nodes and id of other edges
            tri1_nodes = get_tri_node_ids(neighbors_id[1], triangles, edges)
            tri2_nodes = get_tri_node_ids(neighbors_id[2], triangles, edges)
            @debug "triangle 1 : $(neighbor_tri_ids[1]) has nodes $(tri1_nodes)"
            @debug "triangle 2 : $(neighbor_tri_ids[2]) has nodes $(tri2_nodes)"
            opposite_nodes = [setdiff(tri1_nodes, common_nodes)[1], setdiff(tri2_nodes, common_nodes)[1]]
            #top_nodes = setdiff(union(triangles[neighbors_id[1]].nodes, triangles[neighbors_id[1]].nodes)
            #       common_edge.nodes)
            @debug "opposite_nodes : $(opposite_nodes)"
            if length(unique(opposite_nodes)) != 2
                @error "Flip: same opposite nodes"
                error("Flip: same opposite nodes")
            end
            other_edges_ids = setdiff(triangles[neighbors_id[1]].edges, edge_id)
            append!(other_edges_ids, setdiff(triangles[neighbors_id[2]].edges, edge_id))
            # TODO: rewrite with union
            # other_edges_ids = setdiff(union(triangles[neighbors_id[1]].edges, triangles[neighbors_id[2]].edges), common_edge)

            # 2. reconstruct the common edge and the neighbor triangles
            #    -> the common edge is now built with the opposite nodes
            #    -> the nodes of the former common edge are now the new opposite nodes

            # 2.1 common edge
            common_edge = Edge(opposite_nodes, neighbors_id)
            @debug "Common edge $(common_edge_id) updated: nodes = $(opposite_nodes)"
            edges[edge_id] = common_edge

            # 2.2 triangles
            for i_tri in 1:2
                @debug "Rework triangle $(i_tri)"
                tri_id = neighbors_id[i_tri]
                # use a (former) common node and the new common edge
                top_node = common_nodes[i_tri]
                @debug "- Top node = $(top_node)"
                side_edges_ids = [e_id for e_id in other_edges_ids if top_node in edges[e_id].nodes]
                @debug "- Side edge 1 = $(side_edges_ids[1]): $(edges[side_edges_ids[1]])"
                @debug "- Side edge 2 = $(side_edges_ids[2]): $(edges[side_edges_ids[2]])"
                # - create new triangle (replace existing one with same id)
                triangles[tri_id] = Triangle([edge_id, side_edges_ids[1], side_edges_ids[2]])
                @debug "- Triangle $(tri_id) updated: $(triangles[tri_id].edges)"
                # - update the neighbor information (in_triangle) for the side edges
                for i_edge in side_edges_ids
                    # before the flip: side edge is in
                    # - one of the triangles that will be flipped: an id in in neighbor_tri_ids
                    # - a triangle out of the parent triangle unchanged : outer_tri
                    #   -> outer_tri_id is 0 is this outer triangle does not exist
                    # after the flip: side edge is in
                    # - the triangle that has been flipped: id = tri_id
                    # - the outer triangle with id outer_tri_id
                    # => replace id of triangle-to-be-flipped by id of new triangle
                    in_triangles = copy(edges[i_edge].in_triangles)
                    if (in_triangles[1] in neighbors_id)
                        in_triangles[1] = tri_id
                    elseif (in_triangles[2] in neighbors_id)
                        in_triangles[2] = tri_id
                    else
                        @error "- Edge $(i_edge) is not in flipped triangle $(tri_id)"
                        error("- Edge ", i_edge, " is not in flipped triangle ", tri_id)
                    end
                    edge_nodes = edges[i_edge].nodes
                    #edges[i_edge] = Edge(edge_nodes, [tri_id, outer_tri_id])
                    edges[i_edge] = Edge(edge_nodes, in_triangles)
                    @debug "- Edge $(i_edge) now in triangles $(in_triangles)"
                end
            end

            # save result
            if do_save_steps
                new_edges_ids = union(other_edges_ids, edge_id)
                save_step(step_file_pattern, "flip_$(flip_step)", nodes, 0, edges,
                new_edges_ids, [], [], triangles)
            end
        end
    end
    if nb_edges_flipped > 0
        @info "$(nb_edges_flipped) on $(length(edges)) edges have been flipped"
    end

    # check if there remain edges to flip after this first step
    nb_edges_to_flip = 0
    for edge_id in keys(edges)
        common_edge = edges[edge_id]
        # find if neighbor triangles separated by this edge must be flipped
        neighbors_id = common_edge.in_triangles
        if 0 in neighbors_id
            # edge is on the border of the triangulation: no flip
            continue
        elseif require_flip(neighbors_id[1], neighbors_id[2], triangles, edges, nodes)
            nb_edges_to_flip = nb_edges_to_flip + 1
        end
    end
    if (nb_edges_to_flip > 0)
        @info "Recursive call to flip_triangulation to fix $(nb_edges_to_flip) edges"
        flip_triangulation!(triangles, edges, nodes, do_save_steps, step_file_pattern, init_flip_step=flip_step)
    end

end

function save_step(file_pattern, file_suffix, nodes, last_node, edges, last_edges,
                    supertriangle_node_ids, supertriangle_edge_ids, triangles)
    # Create the name of the step from the pattern by inserting a suffix
    # between the pattern and the file extension
    (base, ext) = splitext(file_pattern)
    file = base*file_suffix*ext
    output = Dict("nodes" => nodes,
                  "last_node" => last_node,
                  "edges" => edges,
                  "last_edges" => last_edges,
                  "supertriangle_node_ids" => supertriangle_node_ids,
                  "supertriangle_edge_ids" => supertriangle_edge_ids,
                  "triangles" => triangles,
                  )
    open(file, "w") do io
        write(io, JSON.json(output))
    end
end

function plot_step(step_file, show_non_delaunay_triangles)
    # plot the nodes, edges and triangles for a step
    # warning: an existing file with the same name will be overwritten
    io = open(step_file, "r")
    string_content = read(io, String)
    close(io)
    json_content = JSON.parse(string_content)
    nodes = json_content["nodes"]
    last_node = json_content["last_node"]
    edges = Dict{Int, Edge}()
    for item in collect(json_content["edges"])
        key = parse(Int, item[1])
        edges[key] = Edge(item[2]["nodes"], item[2]["in_triangles"])
    end
    last_edges = json_content["last_edges"]
    supertriangle_edge_ids = json_content["supertriangle_edge_ids"]
    supertriangle_node_ids = json_content["supertriangle_node_ids"]
    triangles = Dict{Int, Triangle}()
    for item in collect(json_content["triangles"])
        key = parse(Int, item[1])
        triangles[key] = Triangle(item[2]["edges"])
    end

    # create a node to init figure - it will be reploted after
    Plots.scatter([nodes[1][1]], [nodes[1][2]], aspect_ratio=:equal, label="", color=:red)

    # non-delaunay triangles
    if show_non_delaunay_triangles
        for (tri_id, tri) in triangles
            tri_node_ids = get_tri_node_ids(tri_id, triangles, edges)
            # find if any node used in the triangulation is in this triangle's circumcircle
            for lookup_node_id in get_triangulation_nodes(edges)
                if !(lookup_node_id in tri_node_ids)
                    if in_circumcircle(tri_node_ids, lookup_node_id, nodes)
                        # the triangle is non-delaunay -> plot it
                        tri_points = [Tuple(nodes[i]) for i in tri_node_ids]
                        #append!(shapes, Plots.Shape(tri_points))
                        Plots.plot!([Plots.Shape(tri_points)], color=:red, label="")
                        # end nodes lookup
                        break
                    end
                end
            end
        end
    end

    # edges
    for e in edges
        if e[1] in last_edges
            # this edge will be plotted with the last edges - skip it
            continue
        elseif e[1] in supertriangle_edge_ids
            # this edge will be plotted with the supertriangle edges - skip it
            continue
        end
        edge_x = [nodes[n][1] for n in e[2].nodes]
        edge_y = [nodes[n][2] for n in e[2].nodes]
        Plots.plot!(edge_x, edge_y, label="", color=:black)
    end

    # supertriangle's edges (if any)
    for edge_id in supertriangle_edge_ids
        if edge_id in last_edges
            # this edge will be plotted with the last edges - skip it
            continue
        end
        edge_nodes = edges[edge_id].nodes
        edge_x = [nodes[n][1] for n in edge_nodes]
        edge_y = [nodes[n][2] for n in edge_nodes]
        Plots.plot!(edge_x, edge_y, label="", color=:chocolate3)
    end

    # last edges (if any)
    for edge_id in last_edges
        edge_nodes = edges[edge_id].nodes
        edge_x = [nodes[n][1] for n in edge_nodes]
        edge_y = [nodes[n][2] for n in edge_nodes]
        Plots.plot!(edge_x, edge_y, label="", color=:green4)
    end

    # nodes
    nodes_x = [n[1] for n in nodes]
    nodes_y = [n[2] for n in nodes]
    Plots.scatter!(nodes_x, nodes_y, label="", color=:deepskyblue)

    # supertriangle's nodes
    nodes_x = [nodes[i][1] for i in supertriangle_node_ids]
    nodes_y = [nodes[i][2] for i in supertriangle_node_ids]
    Plots.scatter!(nodes_x, nodes_y, label="", color=:burlywood1)

    # last node inserted
    if last_node > 0
        Plots.scatter!([nodes[last_node][1]], [nodes[last_node][2]], label="", color=:red)
    end

    # title
    figure_title = basename(splitext(step_file)[1])
    Plots.title!(figure_title)

    # save fig
    figure_file = splitext(step_file)[1]*".png"
    #rm(figure_file, force=true)
    Plots.savefig(figure_file)
end

function inside_triangle(triangle_ids, node_id, nodes)
    # Evaluate if a node is inside a triangle using barycentric coordinates method
    # https://en.wikipedia.org/wiki/Barycentric_coordinate_system
    #   See #Edge_approach for the calculation of the barycentric coordinates
    #   and #Determining_location_with_respect_to_a_triangle for the inside/outside criterion
    # - triangle_ids : id of the nodes of the triangle - list of 3 int
    # - node_id : id of the node investigated - int
    # - nodes : coordinates of the nodes - list of nodes coordinates [x, y]
    (x, y) = nodes[node_id]
    (x1, y1) = nodes[triangle_ids[1]]
    (x2, y2) = nodes[triangle_ids[2]]
    (x3, y3) = nodes[triangle_ids[3]]
    denominator = (y2 - y3)*(x1 - x3) + (x3 - x2)*(y1 - y3)
    lambda1 = ((y2 - y3)*(x - x3) + (x3 - x2)*(y - y3))/denominator
    lambda2 = ((y3 - y1)*(x - x3) + (x1 - x3)*(y - y3))/denominator
    lambda3 = 1 - lambda1 - lambda2
    inside = (0 < lambda1 < 1) && (0 < lambda2 < 1) && (0 < lambda3 < 1)
    return inside
end

function triangulate(input_nodes, must_save_steps=false, step_file_pattern="./steps/step_.txt")
    # - input_nodes : coordinates of the all nodes - list of [x, y] arrays of floats
    # - must_save_steps : save each step to a file (named from pattern)
    # - step_file_pattern : pattern of the step files - ex: "./steps/step_.txt"

    # nodes: work on a shallow copy of input nodes (we will add then remove the 3 supernodes)
    nodes = copy(input_nodes)
    number_of_nodes = length(nodes)
    if number_of_nodes < 3
        # TODO: return empty set instead of printing an error
        @error "The number of nodes must be 3 or more"
        error("The number of nodes must be 3 or more")
    end

    # remove existing step files
    if (must_save_steps)
        @info "Removing existing step files"
        step_file_pattern_basename = splitext(basename(step_file_pattern))[1]
        for step_file in readdir(dirname(step_file_pattern), join=true)
            if startswith(basename(step_file), step_file_pattern_basename)
                rm(step_file)
                @debug "Removed: "*step_file
            end
        end
    end

    # create dicts to store the triangles (1 triangle is described by the indices of its 3 nodes)
    # notice: dicts are used instead of arrays to avoid renumbering elements each time an element
    #         is removed
    triangles = Dict{Int, Triangle}()
    current_triangle_index = 0
    edges = Dict{Int, Edge}()
    current_edge_index = 0
    # create the super triangle
    # -> its id is 1
    # -> its nodes have ids n+1, n+2 and n+3
    # -> its edges are numbered 1, 2 and 3
    supertriangle_nodes = create_supertriangle(nodes)
    supertriangle_node_ids = [number_of_nodes+1, number_of_nodes+2, number_of_nodes+3]
    append!(nodes, supertriangle_nodes)

    edge1 = Edge([supertriangle_node_ids[1], supertriangle_node_ids[2]], [1, 0])
    edges[1] = edge1
    edge2 = Edge([supertriangle_node_ids[2], supertriangle_node_ids[3]], [1, 0])
    edges[2] = edge2
    edge3 = Edge([supertriangle_node_ids[3], supertriangle_node_ids[1]], [1, 0])
    edges[3] = edge3
    current_edge_index = 3

    supertriangle_edge_ids = [1, 2, 3]
    supertriangle = Triangle(supertriangle_edge_ids)
    triangles[1] = supertriangle
    current_triangle_index = 1

    # iterate over the nodes to insert
    for node_id=1:number_of_nodes
        @info "Insert node $(node_id)"

        # find the parent triangle : triangle which circumcircle contains this node
        # --------------------------
        parent_tri_id = 0
        for tri_id in keys(triangles)
            # get nodes of the triangle
            tri_nodes = get_tri_node_ids(tri_id, triangles, edges)
            if inside_triangle(tri_nodes, node_id, nodes)
                parent_tri_id = tri_id
                break
            end
            # Warning: what to do if the node is ON an edge of the triangle ?
        end
        if (parent_tri_id == 0)
            @error "parent triangle not found for node $(node_id)"
            error("parent triangle not found for node ", node_id)
        end
        # DEBUG: remove
        @debug "Id of the parent the triangle: $(parent_tri_id)"

        # update the triangulation: create the new triangles from the parent triangle and the inserted node
        # -------------------------
        # - edges of the parent triangle
        parent_edges = triangles[parent_tri_id].edges
        parent_nodes = get_tri_node_ids(parent_tri_id, triangles, edges)

        # - create the new edges (nodes of the parent tri + new node)
        new_edges = [Edge([parent_nodes[i], node_id], [0, 0]) for i in 1:3]
        new_edges_ids = [0, 0, 0]
        for i_new_edge in 1:3
            if length(unique(new_edges[i_new_edge].nodes)) != 2
                @error "triangulation: error edge with same node twice $(new_edges[i_new_edge])"
                error("triangulation: error edge with same node twice ", new_edges[i_new_edge])
            end
            current_edge_index = current_edge_index + 1
            new_edges_ids[i_new_edge] = current_edge_index
            edges[current_edge_index] = new_edges[i_new_edge]
        end

        # - create the new triangles with a bound edge and the two matching new edges
        new_triangles_indices = []
        for parent_edge_id in parent_edges
            @debug "Create new triangle from parent edge $(parent_edge_id) : $(edges[parent_edge_id])"
            #
            ref_edge = edges[parent_edge_id]
            ref_nodes = ref_edge.nodes
            # find the edges of the new triangle
            tri_edges_ids = [parent_edge_id]
            for i_edge=1:3
                if (new_edges[i_edge].nodes[1] in ref_nodes) || (new_edges[i_edge].nodes[2] in ref_nodes)
                    append!(tri_edges_ids, new_edges_ids[i_edge])
                end
            end
            if length(tri_edges_ids) != 3
                @error "Newly created triangle has not 3 edges"
                error("Newly created triangle has not 3 edges")
            end

            # create the triangle and add it to the table of triangles
            current_triangle_index = current_triangle_index + 1
            new_triangle_id = current_triangle_index
            append!(new_triangles_indices, new_triangle_id)
            triangles[current_triangle_index] = Triangle(tri_edges_ids)
            @debug "Create triangle #$(new_triangle_id) with edges $(tri_edges_ids)"

            # - update neighbor info of the edge from the parent triangle
            ref_edge_triangles = ref_edge.in_triangles
            replace!(ref_edge_triangles, parent_tri_id => new_triangle_id)
            edges[parent_edge_id] = Edge(ref_nodes, ref_edge_triangles)
        end

        # - update the new edges: neighbor triangles
        for i_new_edge in new_edges_ids

            # find the new triangle containing this edge
            tri_neighbors = []
            for tri_id in new_triangles_indices
                if i_new_edge in triangles[tri_id].edges
                    append!(tri_neighbors, tri_id)
                end
            end
            # update the edge
            @debug "Update edge #$(i_new_edge) : in_triangles = $(tri_neighbors)"
            edge_nodes = edges[i_new_edge].nodes
            edges[i_new_edge] = Edge(edge_nodes, tri_neighbors)
        end

        # - remove the parent triangle
        @debug "Delete parent triangle with id $(parent_tri_id)"
        pop!(triangles, parent_tri_id)

        # save step
        if must_save_steps
            save_step(step_file_pattern, string(node_id), nodes, node_id, edges,
                new_edges_ids, supertriangle_node_ids, supertriangle_edge_ids, triangles)
        end

        # flip the new triangles, if needed
        # ---------------------------------
        # we iterate over the edges that made the former parent triangle
        # -> for each, we need to see if the triangles on each side respect the angle criterion

        flip_step_file_pattern = splitext(step_file_pattern)[1]*string(node_id)*"_"*splitext(step_file_pattern)[2]
        flip_triangulation!(triangles, edges, nodes, must_save_steps, flip_step_file_pattern)
    end

    #--------------------------------------------------------------------------
    # Remove the data related to the (now deleted) supertriangle: nodes, edges
    #--------------------------------------------------------------------------

    @info "Cleanup"

    # find the edges containg a supernode
    edges_to_delete_id = []
    triangles_to_delete_id = []
    for (edge_id, edge) in edges
        if intersect(edge.nodes, supertriangle_node_ids) != []
            @debug "Edge $(edge_id) contains a node from the supertriangle -> to delete"
            # mark the edge for deleting
            append!(edges_to_delete_id, edge_id)
            # mark the triangle(s) made with this edge for deletion
            for tri_id in edge.in_triangles
                if tri_id != 0
                    append!(triangles_to_delete_id, tri_id)
                end
            end
        end
    end
    triangles_to_delete_id = unique(triangles_to_delete_id)

    # delete the edges
    for edge_id in edges_to_delete_id
        pop!(edges, edge_id)
    end

    # update information for the edges in triangles about to be deleted
    for (edge_id, edge) in edges
        in_deleted_tri_ids = intersect(edge.in_triangles, triangles_to_delete_id)
        # 3 cases for the current edge:
        # - in 2 triangles marked for deletion -> it has *already* been deleted -> nothing to do
        # - in 1 triangle marked for deletion -> replace its id by 0
        # - in 0 triangles marked for deleting -> nothing to do
        if length(in_deleted_tri_ids) == 1
            in_triangles_updated = replace(edge.in_triangles, in_deleted_tri_ids[1] => 0)
            edges[edge_id] = Edge(edge.nodes, in_triangles_updated)
        end
    end

    # delete the triangles
    for tri_id in triangles_to_delete_id
        @debug "Delete triangle $(tri_id)"
        pop!(triangles, tri_id)
    end

    # delete the nodes of the supertriangle: 3 last elements of "nodes"
    nodes = nodes[1:number_of_nodes]

    # write the final state
    if must_save_steps
        save_step(step_file_pattern, "final", nodes, 0, edges, [], [], [], triangles)
    end

    return edges, triangles
end

#------------------------------------------------------------------------------
# Functions related to the CLI application
# - parse_args: parses the options from a user's input string
# - print_help: displays the help of the CLI
# - get_arg_value: get the value of a couple of options (long/short) from
#   parse_args's output
# - get_main_parameters: get the CLI parameters from the user's input string
# - main: perform the operation requested by the user
#------------------------------------------------------------------------------

function parse_args(args_string::String)
    # parse the input arguments
    output = Dict()
    output["non_options"] = String[]
    option_regex = r"^-[a-zA-Z]"
    long_option_regex = r"^--[a-zA-Z]"
    # note: elements of the result of split are of type SubString{String}
    #       -> convert to String to ease future handling
    args = [string(s) for s in split(args_string)]
    index = 1
    while index <= length(args)
        currentArg = string(args[index])

        if currentArg == "-"
            error("Invalid argument '-'")
        end

        if occursin(option_regex, currentArg)
            # options found -> extract each one are store it in the dict
            number_of_options = length(currentArg) - 1
            if number_of_options == 1
                # single option
                # -> the next argument can hold a value, or be another option
                option = string(currentArg[2])
                if index < length(args)
                    nextArg = args[index+1]
                    if occursin(option_regex, nextArg)
                        # option: skip it
                        value = nothing
                    else
                        # value: store it
                        value = nextArg
                        # increment index because arg[index+1] as been dealt with
                        index = index + 1
                    end
                else
                    # no following argument: empty value
                    value = nothing
                end
                # store option
                output[option] = value
            else
                # several options stacked -> they cannot have values
                for option in currentArg[2:end]
                    output[string(option)] = nothing
                end
            end
        elseif occursin(long_option_regex, currentArg)
            # long option
            if occursin("=", currentArg)
                # the long option is passed with a value
                option, value = split(currentArg[3:end], "=")
                value = string(value)   # convert from SubString{String} to String to ease handling
                if value == ""
                    error("No value for option $(option)")
                end
            else
                option = currentArg[3:end]
                value = nothing
            end
            output[option] = value
        else
            # value not related to an option
            push!(output["non_options"], currentArg)
        end

        # go to the next index in the arguments array
        index = index + 1
    end

    return output
end

function print_help()
    println("Help: Delaunay triangulation of a set of nodes")
    println("-----")
    println("")
    println("Syntax for options: -s short_value --long-option=long_value")
    println("")
    println("Options:")
    println("-c               Create a random set of nodes. The value is the number of nodes")
    println("                 (default: 100")
    println("--create-nodes   Same as -c")
    println("-f               Name of the text file containing the nodes (1 node per line, ")
    println("                 x et y coordinates separated by a comma)")
    println("                 [!] NOT YET IMPLEMENTED")
    println("--from-file      Same as -f")
    println("-h               Display the help")
    println("--help           Same as -h")
    println("-l               When used with -p (or --plot), do not generate a triangulation ")
    println("                 and plot from a result file. The value is the name of the ")
    println("                 result file. Give a regex to plot several files.")
    println("--load-result    Same as -l")
    println("-o               Name of the output file.")
    println("--output         Same as -o")
    println("-p               Plot the triangulation. The possible values are:")
    println("                   'final'   (default) only plot the final triangulation")
    println("                   'all'     plot the final triangulation and each step")
    println("--plot           Same as -p")
    println("-s               Save the steps. The file names are built from the output: ")
    println("                 <output>_step<stepId>_flip<flipId>.<output_ext>")
    println("--save-steps     Same as -s")

end

function get_arg_value(options, args_dict)
    # Get the value of an option which can have several names
    # - options: list of names corresponding to the same option.
    #            eg: ["c", "create"] to get short option "-c" and long option "--create"
    # - args_dict: dictionary containing the options and their values (nothing if no value)
    # Return
    #  - false if none of the options is in the arguments dict
    #  - the value (string or nothing) of the option (or of all options is they match)
    #  - error if several options are in the dict with different values
    options_in_args = [n for n in options if n in keys(args_dict)]
    if options_in_args == []
        # None of the options is in the arguments dict -> use the default value
        return false
    end

    if length(options_in_args) == 1
        return args_dict[options_in_args[1]]
    else
        # Redundant options: check that they have the same values
        if length(unique([args_dict[o] for o in options_in_args])) == 1
            return args_dict[options_in_args[1]]
        else
            # error
            error("Options passed with different value: $(options_in_args)")
        end
    end
end

function get_main_parameters(args_string::String)

    # initialize the parameters' dict with default values
    params = Dict(
        "create_nodes" => false,
        "create_nodes_number" => 100,
        "load_node_file" => false,
        "load_result_files" => false,
        "output" =>  "./triangulation.json",
        "plot_results" => false,
        "save_steps" => false,
        "help" => false
    )

    # parse arguments
    args = parse_args(args_string)

    # handle invalid options: signal, set help and return
    all_options = ["c", "create-nodes", "f", "from-file", "h", "help", "l", "load-result",
                     "o", "output", "p", "plot", "s", "save-steps", "non_options"]
    has_invalid_option = false
    for key in keys(args)
        if ! (key in all_options)
            pop!(args, key)
            if length(key) == 1
                option_name = "-$(key)"
            else
                option_name = "--$(key)"
            end
            println("Invalid option $(option_name)")
            has_invalid_option = true
        end
    end
    if has_invalid_option === true
        params["help"] = true
        return params
    end

    # if there is no option, set help and return
    if length(keys(args)) == 1 && "non_options" in keys(args)
        params["help"] = true
        return params
    end

    # get the parameters from input arguments & default value

    # * create nodes: -c and --create_nodes
    arg_create_nodes = get_arg_value(["c", "create-nodes"], args)
    if arg_create_nodes != false
        # either "c" or "create-nodes" is passed
        params["create_nodes"] = true
        if arg_create_nodes !== nothing
            # passed with a value: convert it to int
            params["create_nodes_number"] = parse(Int, arg_create_nodes)
        end
    end

    # * get nodes from a file: -f and --from-file
    arg_nodes_file = get_arg_value(["f", "from-file"], args)
    if arg_nodes_file === nothing
        error("Options -f and --from-file must be used with a value")
    elseif arg_nodes_file != false
        if Filesystem.isfile(arg_nodes_file)
            params["load_node_file"] = arg_nodes_file
        else
            error("Node file passed to option -f or --from-file does not exist")
        end
    end

    # * display help
    arg_help = get_arg_value(["h", "help"], args)
    if arg_help != false
        params["help"] = true
    end

    # * get the results to plot from a file: -l and --load-result
    arg_load_result = get_arg_value(["l", "load-result"], args)
    if arg_load_result === nothing
        error("Options -l and --load-result must be used with a value")
    elseif arg_load_result != false
        # check the existence of files matching the pattern
        if length(glob(arg_load_result)) > 0
            params["load_result_files"] = arg_load_result
        else
            error("No file match the pattern given for option -l or --load-result")
        end
    end

    # * output file: -o and --output
    arg_output = get_arg_value(["o", "output"], args)
    if typeof(arg_output) === String
        # use the value given
        params["output"] = arg_output
    end

    # * plot results: -p and --plot
    arg_plot_results = get_arg_value(["p", "plot"], args)
    if arg_plot_results === nothing
        params["plot_results"] = "final"
    elseif typeof(arg_plot_results) == String
        if arg_plot_results in ["final", "all"]
            # use the value given
            params["plot_results"] = arg_plot_results
        else
            error("Valid values for -p and --plot are \"final\" and \"all\"")
        end
    end

    # * save the triangulation steps: -s and --save
    arg_save_steps = get_arg_value(["s", "save-steps"], args)
    if arg_save_steps !== false
        params["save_steps"] = true
    end

    return params
end

function main(args)
    # get parameters from the user's arguments
    args_string = join(args, " ")
    params = Dict()     # need to create the variable at this level, before the try/catch
    try
        params = get_main_parameters(args_string)
    catch ErrorException
        # something is invalid with the input parameters
        # -> print help and exit
        print_help()
        return
    end

    # apply a scenario, depending on the inputs
    if params["help"] === true
        # scenario 1: display help
        print_help()
    elseif params["load_result_files"] !== false
        # scenario 2: plot the results of a previous triangulation from result files
        step_file_pattern = params["load_result_files"]
        for filename in glob(step_file_pattern)
            plot_step(filename, true)
        end
    else
        # scenario 3: calculate a triangulation
        # - nodes
        nodes = []
        if params["load_node_file"] !== false
            # load the nodes from a file
            error("Load the nodes from a file: not implemented")
        elseif params["create_nodes"] === true
            # create random nodes
            nodes = create_nodes(params["create_nodes_number"], [-1., 1.], [-1., 1.])
        else
            # no policy for the nodes: notify error and exit
            println("The nodes must be either created or loaded from a file")
            println("  - created: -c or --create-nodes | optional value: number of nodes")
            println("  - loaded: -f or --from-file | value: name of the file")
            return
        end
        # - triangulation
        output = params["output"]
        if params["save_steps"] === true
            do_save_steps = true
            # create the pattern for the step files
            step_file_pattern = splitext(output)[1]*"_step_"*splitext(output)[2]
        else
            do_save_steps = false
            step_file_pattern = ""
        end
        (edges, triangles) = triangulate(nodes, do_save_steps, step_file_pattern)
        save_step(output, "", nodes, 0, edges, [], [], [], triangles)
        # - plot
        if params["plot_results"] !== false
            if params["plot_results"] == "final"
                plot_step(params["output"], true)
            elseif params["plot_results"] == "all"
                output_glob_pattern = splitext(output)[1]*"*"*splitext(output)[2]
                for filename in glob(output_glob_pattern)
                    plot_step(filename, true)
                end
            end
        end
    end
end
;

# CLI call from a terminal: call the main function with CLI args
if abspath(PROGRAM_FILE) == @__FILE__
    main(ARGS)
end
;


X = collect(Float64, reshape(1:40, 10, 4))
x = rand(10) * 10

D = DataFrame(A=rand(10), B=1:10, C=string.(11:20; base=16))
D_NA = allowmissing!(copy(D))
D_NA[1, :A] = missing

@testset "Array" begin
    XX = deepcopy(X)
    mu = center!(XX, obsdim=1)
    @test sum(abs.(mean(XX; dims=1))) == 0
    @test all(std(XX; dims=1) .== std(X; dims=1))
    @test all(mu .== vec(mean(X; dims=1)))

    XX = deepcopy(X)
    mu = center!(XX, ObsDim.First())
    @test sum(abs.(mean(XX; dims=1))) == 0
    @test all(std(XX; dims=1) .== std(X; dims=1))
    @test all(mu .== vec(mean(X; dims=1)))

    XX = deepcopy(X)
    mu = center!(XX, ObsDim.Last())
    @test sum(abs.(mean(XX; dims=2))) == 0
    @test all(std(XX; dims=2) .== std(X; dims=2))
    @test all(mu .== vec(mean(X; dims=2)))

    XX = deepcopy(X)
    mu = center!(XX)
    @test sum(abs.(mean(XX; dims=2))) == 0
    @test all(std(XX; dims=2) .== std(X; dims=2))
    @test all(mu .== vec(mean(X; dims=2)))

    XX = deepcopy(X)
    mu = vec(mean(X; dims=1))
    center!(XX, mu, obsdim=1)
    @test sum(abs.(mean(XX; dims=1))) == 0
    @test all(std(XX; dims=1) .== std(X; dims=1))

    XX = deepcopy(X)
    mu = vec(mean(X; dims=1))
    center!(XX, mu, ObsDim.First())
    @test sum(abs.(mean(XX; dims=1))) == 0
    @test all(std(XX; dims=1) .== std(X; dims=1))

    XX = deepcopy(X)
    mu = vec(mean(XX; dims=2))
    center!(XX, mu, obsdim=2)
    @test sum(abs.(mean(XX; dims=2))) == 0
    @test all(std(XX; dims=2) .== std(X; dims=2))

    XX = deepcopy(X)
    mu = vec(mean(XX; dims=2))
    center!(XX, mu, ObsDim.Last())
    @test sum(abs.(mean(XX; dims=2))) == 0
    @test all(std(XX; dims=2) .== std(X; dims=2))

    XX = deepcopy(X)
    mu = vec(mean(X[:,[1,3]]; dims=1))
    center!(XX, mu, obsdim=1, operate_on=[1, 3])
    @test sum(abs.(mean(XX[:,[1,3]]; dims=1))) == 0
    @test all(XX[:,2] .== X[:,2])
    @test all(std(XX; dims=1) .== std(X; dims=1))

    XX = deepcopy(X)
    mu = vec(mean(X[[1,3],:]; dims=2))
    center!(XX, mu, obsdim=2, operate_on=[1, 3])
    @test sum(abs.(mean(XX[[1,3],:]; dims=2))) == 0
    @test all(XX[2,:] .== X[2,:])
    @test all(std(XX; dims=2) .== std(X; dims=2))
    println()

    xx = deepcopy(x)
    center!(xx)
    @test mean(xx) <= 10e-10

    xx = deepcopy(x)
    mu = mean(xx)
    center!(xx, mu)
    @test mean(xx) <= 10e-10

    xx = deepcopy(x)
    mu = ones(size(xx))
    center!(xx, mu)
    @test mean(xx) - mean(x) ≈ -1

    xx = deepcopy(x)
    mu = ones(size(xx))
    center!(xx, mu)
    @test mean(xx) - mean(x) ≈ -1
end

@testset "DataFrame" begin
    # Center DataFrame
    DD = deepcopy(D)
    center!(DD)
    @test abs.(mean(DD[:A])) <= 10e-10
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D)
    center!(DD, operate_on=[:B])
    @test all(DD[:A] .== D[:A])
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D)
    mu = center!(DD, operate_on=[:A, :B])
    @test abs.(mean(DD[:A])) <= 10e-10
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])
    @test all(mu .== [mean(D[:A]), mean(D[:B])])

    DD = deepcopy(D_NA)
    mu = center!(DD, [:A, :B, :C])
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[2:end, :A] .== D[2:end, :A])
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D_NA)
    mu = [0.0, mean(DD[:B]), 0.0]
    mu = center!(DD, mu, [:A, :B, :C])
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[2:end, :A] .== D[2:end, :A])
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D_NA)
    mu = 0.0
    mu = center!(DD, mu, :A)
    @test all(DD[2:end, :A] .== D[2:end, :A])
    @test all(DD[:B] .== D[:B])
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D)
    mu =  [mean(D[:A]), mean(D[:B])]
    @test all(center!(DD, mu, operate_on=[:A, :B]) .== mu)
    @test abs.(mean(DD[:A])) <= 10e-10
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D_NA)
    center!(DD)
    @test all(DD[2:end, :A] .== D[2:end, :A])
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])
    @test ismissing(DD[1, :A])
end


##############################################################################
#
# In this test problem we consider an active-set method.
#
# Note that there is no optimization of the evaluations here.
#
# Note the use of a structure for the algorithmic parameters which is
# forwarded to all the 3 steps. If a parameter is not mentioned, then the default
# entry in the algorithm will be taken.
#
#############################################################################

#include("penalty.jl")

##############################################################################
#
# First, we create a subtype of AbstractNLPModel to represent the unconstrained
# subproblem we "solve" at each iteration of the activeset.
#
import NLPModels: obj, grad, hess, hprod

mutable struct ActifNLP <: AbstractNLPModel
  nlp::AbstractNLPModel
  x0::Vector #reference vector
  I::Vector #set of active indices
  Ic::Vector #set of inactive indices
  meta::AbstractNLPModelMeta
  counters::Counters
end

function obj(anlp::ActifNLP, x::Vector)
  t = anlp.x0
  t[anlp.Ic] = x
  return obj(anlp.nlp, t)
end
function grad(anlp::ActifNLP, x::Vector)
  t = anlp.x0
  t[anlp.Ic] = x
  return grad(anlp.nlp, t)[anlp.Ic]
end
function hess(anlp::ActifNLP, x::Vector)
  t = anlp.x0
  t[anlp.Ic] = x
  return hess(anlp.nlp, t)[anlp.Ic, anlp.Ic]
end
function hprod(anlp::ActifNLP, x::Vector, v::Vector, y::Vector)
  return hess(anlp, x) * v
end

##############################################################################
#
# Active-set algorithm for bound constraints optimization
# fill_in! used instead of update! (works but usually more costly in evaluations)
# subproblems are solved via Newton method
#
#############################################################################

function activeset(
  stp::NLPStopping;
  active::Float64 = stp.meta.tol_check(stp.meta.atol, stp.meta.rtol, stp.meta.optimality0),
  prms = nothing,
)
  xt = stp.current_state.x
  n = length(xt)
  all = findall(xt .== xt)

  if maximum(vcat(max.(xt - stp.pb.meta.uvar, 0.0), max.(-xt + stp.pb.meta.lvar, 0.0))) > 0.0
    #OK = true; stp.meta.fail_sub_pb = true
    #xt is not feasible
    xt = max.(min.(stp.current_state.x, stp.pb.meta.uvar), stp.pb.meta.lvar)
  end

  fill_in!(stp, xt)
  OK = start!(stp)

  Il = findall(abs.(-xt + stp.pb.meta.lvar) .<= active)
  Iu = findall(abs.(xt - stp.pb.meta.uvar) .<= active)
  I = union(Il, Iu)
  Ic = setdiff(all, I)
  nI = max(0, length(xt) - length(Il) - length(Iu)) #lvar_i != uvar_i
  @show xt, I
  while !OK

    #prepare the subproblem stopping:
    subpb = ActifNLP(nlp, xt, I, Ic, NLPModelMeta(nI), Counters())
    #the subproblem stops if he solved the unconstrained nlp or iterate is infeasible
    feas(x, y) =
      maximum(vcat(max.(y.x - stp.pb.meta.uvar[Ic], 0.0), max.(-y.x + stp.pb.meta.lvar[Ic], 0.0)))
    check_func(x, y) = feas(x, y) > 0.0 ? 0.0 : unconstrained_check(x, y)
    substp = NLPStopping(subpb, NLPAtX(xt[Ic]), main_stp = stp, optimality_check = check_func)

    #we solve the unconstrained subproblem:
    global_newton(substp, prms)
    @show status(substp, list = true)

    if feas(substp.pb, substp.current_state) > 0.0 #new iterate is infeasible
      #then we need to project
      xt[Ic] = max.(min.(substp.current_state.x, stp.pb.meta.uvar[Ic]), stp.pb.meta.lvar[Ic])
      #we keep track of the new active indices
      Inew = setdiff(
        union(
          findall(abs.(-xt + stp.pb.meta.lvar) .<= active),
          findall(abs.(x0 - stp.pb.meta.uvar) .<= active),
        ),
        I,
      )
    else
      Inew = []
    end

    fill_in!(stp, xt) #the lazy update

    OK = update_and_stop!(stp, evals = stp.pb.counters)

    if !OK #we use a relaxation rule based on an approx. of Lagrange multipliers
      Irmv = findall(stp.current_state.mu .< 0.0)
      I = union(setdiff(I, Irmv), Inew)
      Ic = setdiff(all, I)
    end
    @show xt, I
  end #end of main loop

  return stp
end


using Interpolations
using Test, LinearAlgebra

@testset "Scaling" begin
    # Model linear interpolation of y = -3 + .5x by interpolating y=x
    # and then scaling to the new x range

    itp = interpolate(1:1.0:10, BSpline(Linear()))

    sitp = @inferred(scale(itp, -3:.5:1.5))
    @test typeof(sitp) <: Interpolations.ScaledInterpolation
    @test parent(sitp) === itp

    for (x,y) in zip(-3:.05:1.5, 1:.1:10)
        @test sitp(x) ≈ y
    end

    # Verify that it works in >1D, with different types of ranges

    gauss(phi, mu, sigma) = exp(-(phi-mu)^2 / (2sigma)^2)
    testfunction(x,y) = gauss(x, 0.5, 4) * gauss(y, -.5, 2)

    xs = -5:.5:5
    ys = -4:.2:4
    zs = Float64[testfunction(x,y) for x in xs, y in ys]

    itp2 = interpolate(zs, BSpline(Quadratic(Flat(OnGrid()))))
    sitp2 = @inferred scale(itp2, xs, ys)

    for x in xs, y in ys
        @test testfunction(x,y) ≈ sitp2(x,y)
    end

    # Test gradients of scaled grids
    xs = -pi:.1:pi
    ys = map(sin, xs)
    itp = interpolate(ys, BSpline(Linear()))
    sitp = @inferred scale(itp, xs)

    for x in -pi:.1:pi
        g = @inferred(Interpolations.gradient(sitp, x))[1]
        @test ≈(cos(x),g,atol=0.05)
    end

    # Verify that return types are reasonable
    @inferred(sitp2(-3.4, 1.2))
    @inferred(sitp2(-3, 1))
    @inferred(sitp2(-3.4, 1))

    sitp32 = @inferred scale(interpolate(Float32[testfunction(x,y) for x in -5:.5:5, y in -4:.2:4], BSpline(Quadratic(Flat(OnGrid())))), -5f0:.5f0:5f0, -4f0:.2f0:4f0)
    @test typeof(@inferred(sitp32(-3.4f0, 1.2f0))) == Float32

    # Iteration
    itp = interpolate(rand(3,3,3), BSpline(Quadratic(Flat(OnCell()))))
    knots = map(d->1:10:21, 1:3)
    sitp = @inferred scale(itp, knots...)

    iter = @inferred(eachvalue(sitp))

    function foo!(dest, sitp)
        i = 0
        for s in eachvalue(sitp)
            dest[i+=1] = s
        end
        dest
    end
    function bar!(dest, sitp)
        for I in CartesianIndices(size(dest))
            dest[I] = sitp(I)
        end
        dest
    end
    rfoo = Array{Float64}(undef, Interpolations.ssize(sitp))
    rbar = similar(rfoo)
    foo!(rfoo, sitp)
    bar!(rbar, sitp)
    @test rfoo ≈ rbar
end


function pvcongruence(W)
	err, i, todelete, vclasses = 10^-6, 1, [], []
	verts = convert(Lar.Points, W')
	kdtree = NearestNeighbors.KDTree(verts)
	newverts = zeros(Int, size(verts,2))
	for vi in 1:size(verts,2)
		if !(vi in todelete)
			nearvs = NearestNeighbors.inrange(kdtree, verts[:,vi], err)
			push!(vclasses,nearvs)
			newverts[nearvs] .= i
			nearvs = setdiff(nearvs, vi)
			todelete = union(todelete, nearvs)
			i += 1
		end
	end
	V = zeros(3,length(vclasses))
	for (k,class) in enumerate(vclasses)
			V[:,k] = sum(W[class,:],dims=1)/length(class)
	end
	return V, vclasses
end

V, vclasses = vcongruence(W)

GL.VIEW([
	GL.GLAxis(GL.Point3d(-1,-1,-1),GL.Point3d(1,1,1))
  #GL.GLPol( V,[[1,2,3,4,5,6,7,8]], GL.COLORS[1])
	GL.GLPol( V,FV, GL.COLORS[1])
]);

function EVcongruence2(Delta_0,vclasses)
	Y = SparseVector[]
	for class in vclasses
		y = sparsevec(sum(Delta_0[:,class],dims=2))
		push!(Y,y)
	end
	copEV = hcat(Y...)

	function econgruence(copEV)
		eclasses = sort([(findnz(copEV[k,:])[1],k) for k=1:copEV.m])
		class = [eclasses[1]]
		classes = []
		for h=2:length(eclasses)
			if eclasses[h][1]!=class[end][1]
				push!(classes, class)
				class = []
			end
			push!(class, eclasses[h])
		end
		return classes
	end

	eclasses = [[idx for (edge,idx) in class] for class in econgruence(copEV)]
	lastclass = setdiff(1:Delta_0.m, cat(eclasses))
	push!(eclasses,lastclass)

	newedges = [eclass[1] for eclass in eclasses]
	newEV = hcat([copEV[k,:] for k=1:copEV.m if k in newedges]...)

	return newEV,eclasses
end


# TODO: add more target distributions and make them iteratable
# TODO: Integrate with https://github.com/xukai92/VecTargets.jl to achieve goal noted
#       above.

# Dimension of testing distribution
const D = 5
# Tolerance ratio
const TRATIO = Int == Int64 ? 1 : 2
# Deterministic tolerance
const DETATOL = 1e-3 * D * TRATIO
# Random tolerance
const RNDATOL = 5e-2 * D * TRATIO * 2

# Hand-coded multivariate Gaussian

const gaussian_m = zeros(D)
const gaussian_s = ones(D)

struct Gaussian{Tm, Ts}
    m::Tm
    s::Ts
end

function ℓπ_gaussian(g::AbstractVecOrMat{T}, s) where {T}
    return .-(log(2 * T(pi)) .+ 2 .* log.(s) .+ abs2.(g) ./ s.^2) ./ 2
end

ℓπ_gaussian(m, s, x) = ℓπ_gaussian(m .- x, s)

function ∇ℓπ_gaussianl(m, s, x)
    g = m .- x
    v = ℓπ_gaussian(g, s)
    return v, g
end

function get_ℓπ(g::Gaussian)
    ℓπ(x::AbstractVector) = sum(ℓπ_gaussian(g.m, g.s, x))
    ℓπ(x::AbstractMatrix) = dropdims(sum(ℓπ_gaussian(g.m, g.s, x); dims=1); dims=1)
    return ℓπ
end

function get_∇ℓπ(g::Gaussian)
    function ∇ℓπ(x::AbstractVector)
        val, grad = ∇ℓπ_gaussianl(g.m, g.s, x)
        return sum(val), grad
    end
    function ∇ℓπ(x::AbstractMatrix)
        val, grad = ∇ℓπ_gaussianl(g.m, g.s, x)
        return dropdims(sum(val; dims=1); dims=1), grad
    end
    return ∇ℓπ
end

ℓπ = get_ℓπ(Gaussian(gaussian_m, gaussian_s))
∂ℓπ∂θ = get_∇ℓπ(Gaussian(gaussian_m, gaussian_s))

# For the Turing model
# @model gdemo() = begin
#     s ~ InverseGamma(2, 3)
#     m ~ Normal(0, sqrt(s))
#     1.5 ~ Normal(m, sqrt(s))
#     2.0 ~ Normal(m, sqrt(s))
#     return s, m
# end

using Distributions: logpdf, InverseGamma, Normal
using Bijectors: invlink, logpdf_with_trans

function invlink_gdemo(θ)
    s = invlink(InverseGamma(2, 3), θ[1])
    m = θ[2]
    return [s, m]
end

function ℓπ_gdemo(θ)
    s, m = invlink_gdemo(θ)
    logprior = logpdf_with_trans(InverseGamma(2, 3), s, true) + logpdf(Normal(0, sqrt(s)), m)
    loglikelihood = logpdf(Normal(m, sqrt(s)), 1.5) + logpdf(Normal(m, sqrt(s)), 2.0)
    return logprior + loglikelihood
end

test_show(x) = test_show(s -> length(s) > 0, x)
function test_show(pred, x)
    io = IOBuffer(; append = true)
    show(io, x)
    s = read(io, String)
    @test pred(s)
end


import DoubleFloats: Double64, ComplexDF64

@testset "LinearAlgebra" begin
	@testset "Constructor + basics" begin
		A = rand(ComplexF32, 12, 12)
		WS = HC.MatrixWorkspace(A)
		@test WS isa HC.MatrixWorkspace{Float64}
		@test_throws ArgumentError HC.update!(WS, rand(12, 13))

		WS.factorized[] = true
		B = randn(ComplexF64, 12, 12)
		HC.update!(WS, B)
		@test WS.A == B
		@test WS.A !== B
		@test WS.factorized[] == false

		HC.factorization!(WS, HC.LU_FACT)
		@test WS.fact[] == HC.LU_FACT
		WS.factorized[] = true
		HC.factorization!(WS, HC.QR_FACT)
		@test WS.fact[] == HC.QR_FACT
		@test WS.factorized[] == false

		HC.factorization!(WS, HC.LU_FACT)
		@test (@allocated HC.factorize!(WS)) == 0
		HC.factorization!(WS, HC.QR_FACT)
		@test WS.factorized[] == false
		HC.factorize!(WS) # first factorize resizes the work buffers
		WS.factorized[] = false
		@test (@allocated HC.factorize!(WS)) == 0

		# test indexing etc
		@test WS[2,1] == B[2, 1]
		@test WS[2] == B[2]
		WS[2,3] = 5.0
		@test WS[2,3] == 5.0
		WS[5] = 2.32
		@test WS[5] == 2.32
		@test size(WS) == (12,12)
		C = randn(ComplexF64, 12, 12)
		WS .= C
		@test WS.A == C
	end

	@testset "ldiv" begin
		A = randn(ComplexF64, 13, 13)
		b = randn(ComplexF64, 13)
		x = zeros(ComplexF64, 13)
		WS = HC.MatrixWorkspace(A)
		HC.update!(WS, A)
		ldiv!(x, WS, b)
		WS.factorized[] = false
		@test (@allocated ldiv!(x, WS, b)) == 0
		@test (lu(A) \ b) ≈ x atol=1e-12

		HC.factorization!(WS, HC.QR_FACT)
		x .= 0
		ldiv!(x, WS, b)
		@test (qr(A, Val(true)) \ b) ≈ x atol=1e-12
		WS.factorized[] = false
		@test (@allocated ldiv!(x, WS, b)) == 0

		# overdetermined
		A = randn(ComplexF64, 13, 8)
		x = randn(ComplexF64, 8)
		b = A*x
		WS = HC.MatrixWorkspace(A)
		HC.update!(WS, A)
		@test_throws ArgumentError HC.factorization!(WS, HC.LU_FACT)
		ldiv!(x, WS, b)
		WS.factorized[] = false
		x .= 0
		@test (@allocated ldiv!(x, WS, b)) == 0
		@test (qr(A, Val(true)) \ b) ≈ x atol=1e-12
	end

	@testset "(Mixed precision) residual" begin
		A = randn(ComplexF64, 13, 13)
		x = randn(ComplexF64, 13)
		b = A * x
		û = convert.(ComplexF64, ComplexDF64.(A) * ComplexDF64.(x) - ComplexDF64.(b))

		u = zeros(ComplexF64, 13)
		# fixed precision residual should be of order unit roundoff
		HC.residual!(u, A, x, b)
		@test maximum(abs, ComplexDF64.(u) - ComplexDF64.(û)) < 1e-14
		@test maximum(abs, ComplexDF64.(u) - ComplexDF64.(û)) > 1e-18

		HC.residual!(u, A, x, b, Complex{Double64})
		@test maximum(abs, ComplexDF64.(u) - ComplexDF64.(û)) < 1e-26
	end

	@testset "(Mixed precision) iterative refinement" begin
		A = randn(ComplexF64, 4, 4)
		b = randn(ComplexF64, 4)
		x̂ = zeros(ComplexF64, 4)
		x = copy(x̂)
		WS = HC.MatrixWorkspace(A)
		U, S, VT = svd(A)
		S[end] *= 1e-4
		A2 = U * diagm(0=>S) * VT'
		update!(WS, A2)
		ldiv!(x̂, WS, b)
		δx1 = HC.iterative_refinement_step!(x, WS, x̂, b, HC.InfNorm(), ComplexDF64)
		@test δx1 / norm(x, Inf) > 1e-14
		δx2 = HC.iterative_refinement_step!(WS, x, b, HC.InfNorm(), ComplexDF64)
		@test δx2/ norm(x, Inf) < 1e-16

		δx = HC.iterative_refinement_step!(x, WS, x̂, b)
		@test δx1 / norm(x, Inf) > 1e-14

		# overdetermined
		A = randn(ComplexF64, 4, 3)
		x̂ = zeros(ComplexF64, 3)
		x = copy(x̂)
		WS = HC.MatrixWorkspace(A)
		U, S, VT = svd(A)
		S[end] *= 1e-4
		A2 = U * diagm(0=>S) * VT'
		update!(WS, A2)
		b = A2 * randn(ComplexF64, 3)
		ldiv!(x̂, WS, b)
		δx1 = HC.iterative_refinement_step!(x, WS, x̂, b, HC.InfNorm(), ComplexDF64)
		@test δx1 / norm(x, Inf) > 1e-14
		δx2 = HC.iterative_refinement_step!(WS, x, b, HC.InfNorm(), ComplexDF64)
		@test δx2/ norm(x, Inf) < 1e-16

		δx = HC.iterative_refinement_step!(x, WS, x̂, b)
		@test δx1 / norm(x, Inf) > 1e-14
	end

	@testset "rcond" begin
		A = randn(ComplexF64, 13, 13)
		WS = HC.MatrixWorkspace(A)
		U, S, VT = svd(A)
		S[end] *= 1e-5
		A2 = U * diagm(0=>S) * VT'
		HC.update!(WS, A2)
		HC.factorization!(WS, HC.LU_FACT)
		HC.factorize!(WS)
		@test cond(A2, Inf) ≈ inv(HC.rcond(WS))
		HC.factorization!(WS, HC.QR_FACT)
		@test 0.5 ≤ cond(A2, Inf) / inv(HC.rcond(WS)) ≤ 1.5
		# @test (@allocated rcond(WS)) == 0

		A = randn(ComplexF64, 13, 10)
		WS = HC.MatrixWorkspace(A)
		U, S, VT = svd(A)
		S[end] *= 1e-8
		A2 = U * diagm(0=>S) * VT'
		HC.update!(WS, A2)
		HC.factorize!(WS)
		@test 1 / HC.rcond(WS) ≈ cond(UpperTriangular(WS.qr.R), Inf) rtol=1e-14
		# @test (@allocated rcond(WS)) == 0

		# meddle aroudn with the internals to check dead branch
		WS.fact[] = HC.LU_FACT
		@test isnan(HC.rcond(WS))
	end

	@testset "inf norm est" begin
		A = randn(ComplexF64, 5, 5)
		U, s, Vt = svd(A)
		s[end] *= 1e-6
		B = U*diagm(0=>s)*Vt'

		d = rand(5)*1000
		D = diagm(0=>d)
		g = rand(5) .* 1000
		G = diagm(0=>g)

		@test 0.1 ≤ opnorm(inv(B), Inf) / HC.inf_norm_est(lu(B)) ≤ 10
		@test 0.1 ≤ opnorm(inv(D)*inv(B), Inf) / HC.inf_norm_est(lu(B), nothing, d) ≤ 10
		@test 0.1 ≤ opnorm(inv(B)*G, Inf) / HC.inf_norm_est(lu(B), g, nothing) ≤ 10
		@test 0.1 ≤ opnorm(inv(D)*inv(B)*G, Inf) / HC.inf_norm_est(lu(B), g, d) ≤ 10

		WS = HC.MatrixWorkspace(B)
		@test 0.1 ≤ opnorm(inv(B), Inf) / HC.inf_norm_est(WS) ≤ 10
		@test 0.1 ≤ opnorm(inv(D)*inv(B), Inf) / HC.inf_norm_est(WS, nothing, d) ≤ 10
		@test 0.1 ≤ opnorm(inv(B)*G, Inf) / HC.inf_norm_est(WS, g, nothing) ≤ 10
		@test 0.1 ≤ opnorm(inv(D)*inv(B)*G, Inf) / HC.inf_norm_est(WS, g, d) ≤ 10
	end

	@testset "JacobianMonitor" begin
		A = randn(ComplexF64, 6, 6)
		b = randn(ComplexF64, 6)
		x̂ = similar(b)
		x = similar(b)
		JM = HC.JacobianMonitor(zeros(ComplexF64, 6, 6))
		@test JM isa HC.JacobianMonitor
		HC.jacobian(JM) .= A
		HC.updated!(JM)
		ldiv!(x̂, JM, b)
		ldiv!(x, jacobian(JM), b)
		@test x == x̂
	end
	@testset "Forward error" begin
		A = [1.81451+0.46473im -0.473651+0.813117im 0.130822-0.254704im -0.203277+0.397106im;
			 -2.04468-1.59228im 0.441157-1.39994im -0.0984804+1.42619im -0.818292-0.487898im;
			 0.35501+0.0235574im 1.32618+0.577264im 1.12561+0.43383im 0.579533+1.49484im;
			 0.144381+0.104175im 0.665384-1.43389im 0.22097+0.258016im 0.374103-0.635066im]
		b = [-0.14581-0.868847im, 0.0162454+0.91156im, -0.133769+0.510834im, -0.867022+0.752196im]
		x̂ = similar(b)
		x = similar(b)
		JM = HC.JacobianMonitor(A)
		HC.updated!(JM)
		ldiv!(x̂, JM, b)

		ferr = HC.forward_err!(x, JM, x̂, b, HC.InfNorm())
		@test HC.forward_err(JM) == ferr
		@test ferr < eps()*cond(A)
		@test ferr > 1e-12
		ferrD64 = HC.forward_err!(x, JM, x̂, b, HC.InfNorm(), ComplexDF64)
		@test 0.5 ≤ ferrD64 / ferr ≤ 2

		ferr = HC.forward_err!(JM, x̂, b, HC.InfNorm(), ComplexF64)
		@test ferr < eps()*cond(A)
		@test ferr > 1e-12

		D = diagm(0 => [10.0^(2i) for i in 1:4])
		JM = HC.JacobianMonitor(D*A)
		HC.updated!(JM)
		ldiv!(x̂, JM, D*b)
		conderr = HC.cond(D*A)*eps()
		sferr = HC.strong_forward_err!(x, JM, x̂, D*b, HC.InfNorm())
		@test HC.forward_err(JM) == sferr
		ferr = HC.forward_err!(x, JM, x̂, D*b, HC.InfNorm())
		@test ferr < sferr < conderr*1e-4 # check that error estimate is actually substantially less

		sferr2 = HC.strong_forward_err!(JM, x̂, D*b, HC.InfNorm())
		@test sferr == sferr2
		@test x̂ ≈ x

		DR = diagm(0=>[10.0^(2i-4) for i in 1:4])
		JM = HC.JacobianMonitor(A*DR)
		ldiv!(x̂, JM, b)

		WN = HC.WeightedNorm(HC.InfNorm(), x̂)
		WN .= abs.(x̂)

		wferr = HC.forward_err!(x, JM, x̂, b, WN)
		ferr = HC.forward_err!(x, JM, x̂, b, HC.InfNorm())

		strong_wferr = HC.strong_forward_err!(x, JM, x̂, b, WN)
		strong_ferr = HC.strong_forward_err!(x, JM, x̂, b, HC.InfNorm())
		@test ferr ≤ strong_ferr ≤ eps()*inv(HC.rcond(jacobian(JM)))
		@test wferr ≤ strong_wferr ≤ eps()*inv(HC.rcond(jacobian(JM)))

		# overdetermined
		A = randn(ComplexF64, 3, 2)
		x = randn(ComplexF64, 2)
		b = A * x
		x̂ = similar(x)

		JM = HC.JacobianMonitor(A)
		ldiv!(x̂, JM, b)

		ferr = HC.forward_err!(x, JM, x̂, b, HC.InfNorm())
		strong_ferr = HC.strong_forward_err!(x, JM, x̂, b, HC.InfNorm())
		@test ferr < 1e-15
		@test ferr == strong_ferr # test fallback to ferr for qr
	end

	@testset "cond" begin
		A = randn(ComplexF64, 13, 13)
		D = diagm(0=>[2.0^(i-5) for i in 1:13])
		JM = HC.JacobianMonitor(D*A)
		comp_cond = HC.cond!(JM)
		@test comp_cond === cond(JM)
		B = diagm(0=>inv.(JM.J.row_scaling)) * D * A
		@test abs(cond(JM) - cond(B, Inf)) / cond(JM) < 2

		A = randn(ComplexF64, 13, 12)
		JM = HC.JacobianMonitor(D*A)
		comp_cond = HC.cond!(JM)
		@test comp_cond === cond(JM)
	end

	@testset "Hermite Normal Form" begin
        A = [0 3 1 0; -2 2 -1 2; 1 -1 2 3; -3 3 3 2]
        H, U = HC.hnf(A)
        @test A * U == H

        D = zeros(Int, 1,1)
        D[1,1] = -2
        H, U = HC.hnf(D)
        @test H[1,1] == 2
        @test U[1,1] == -1

		# overflow A
		A = [2 4 4 2 -4 -3 -4 3 0 3;
		     -4 -2 -5 -5 1 0 -3 0 -1 -2;
			 -3 5 -3 1 4 3 -1 -4 -1 0;
			 1 3 3 4 -2 -3 -2 -5 -4 -3;
			 2 2 0 2 -4 4 3 -4 -2 -4;
			 -3 -3 1 1 -4 4 0 4 4 -4;
			 -3 -1 -1 -1 2 -4 -3 -4 4 -4;
			 0 1 -2 4 5 4 3 1 -5 2;
			 -5 -4 -5 3 1 5 0 -3 -3 -1;
			 2 -5 -3 -1 -1 5 -4 5 -3 1]
		@test_throws OverflowError HC.hnf(A)
		@test_throws OverflowError HC.hnf(A, Int128)
		H, U = HC.hnf(A, BigInt)
		@test A * U == H
    end
end


__precompile__()

module EmpiricalRiskMinimization

using Compat

# Losses
export loss
export SquareLoss, AbsoluteLoss, HuberLoss, HingeLoss, LogisticLoss, SigmoidLoss
export MultiLogisticOneHotLoss
include("losses.jl")

# Regularizers
export L1Reg, L2Reg, L1L2Reg, SqrtReg, NonnegReg
include("regularizers.jl")

# Internal imports (private)
include("internal/source.jl")
include("internal/results_and_data.jl") # Results, Data
include("internal/optimizer.jl") # Gradient descent, QR, CVX solvers

# Models
export Model, status, setfeatures, train, trainpath, trainfolds, Ytrain, Ytest, Xtrain, Xtest, Utrain, Utest, Vtrain, Vtest
include("model.jl")

# Validation
export thetaopt, trainloss, testloss, lambda, lambdaopt
export lambdapath, testlosspath, trainlosspath, thetapath
include("validation.jl")

# Prediction
export predict
export predict_y_from_test, predict_y_from_train, predict_v_from_test, predict_v_from_train
export predict_y_from_u, predict_y_from_u, predict_v_from_u, predict_v_from_u
include("prediction.jl")

# Convenience functions
export sigm, matrix
include("convenience.jl")

end # module


using MINPACK
using Test, Printf, LinearAlgebra

#=
NOTE: the contents of this file were taken from NLsolve.jl.

The license for that package appears below.

The NLsolve.jl package is licensed under the MIT License:
Copyright (c) 2017 NLsolve.jl contributors

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

=#

# If the results should be printed to a
const PRINT_FILE = false

struct DifferentiableMultivariateFunction{Tf,Tg}
    f!::Tf
    g!::Tg
end

@testset "minpack" begin

function rosenbrock()
    function f!(fvec::Vector, x::Vector)
        fvec[1] = 1 - x[1]
        fvec[2] = 10(x[2]-x[1]^2)
    end
    function g!(fjac::Matrix, x::Vector)
        fjac[1,1] = -1
        fjac[1,2] = 0
        fjac[2,1] = -20x[1]
        fjac[2,2] = 10
    end
    (DifferentiableMultivariateFunction(f!, g!), [-1.2; 1.], "Rosenbrock")
end

function powell_singular()
    function f!(fvec::Vector, x::Vector)
        fvec[1] = x[1] + 10x[2]
        fvec[2] = sqrt(5)*(x[3] - x[4])
        fvec[3] = (x[2] - 2x[3])^2
        fvec[4] = sqrt(10)*(x[1] - x[4])^2
    end
    function g!(fjac::Matrix, x::Vector)
        fill!(fjac, 0)
        fjac[1,1] = 1
        fjac[1,2] = 10
        fjac[2,3] = sqrt(5)
        fjac[2,4] = -fjac[2,3]
        fjac[3,2] = 2(x[2] - 2x[3])
        fjac[3,3] = -2fjac[3,2]
        fjac[4,1] = 2sqrt(10)*(x[1] - x[4])
        fjac[4,4] = -fjac[4,1]
    end
    (DifferentiableMultivariateFunction(f!, g!), [3.; -1.; 0.; 1.], "Powell singular")
end

function powell_badly_scaled()
    c1 = 1e4
    c2 = 1.0001
    function f!(fvec::Vector, x::Vector)
        fvec[1] = c1*x[1]*x[2] - 1
        fvec[2] = exp(-x[1]) + exp(-x[2]) - c2
    end
    function g!(fjac::Matrix, x::Vector)
        fjac[1,1] = c1*x[2]
        fjac[1,2] = c1*x[1]
        fjac[2,1] = -exp(-x[1])
        fjac[2,2] = -exp(-x[2])
    end
    (DifferentiableMultivariateFunction(f!, g!), [0.; 1.], "Powell badly scaled")
end

function wood()
    c3 = 2e2
    c4 = 2.02e1
    c5 = 1.98e1
    c6 = 1.8e2

    function f!(fvec::Vector, x::Vector)
        temp1 = x[2] - x[1]^2
        temp2 = x[4] - x[3]^2
        fvec[1] = -c3*x[1]*temp1 - (1 - x[1])
        fvec[2] = c3*temp1 + c4*(x[2] - 1) + c5*(x[4] - 1)
        fvec[3] = -c6*x[3]*temp2 - (1 - x[3])
        fvec[4] = c6*temp2 + c4*(x[4] - 1) + c5*(x[2] - 1)
    end

    function g!(fjac::Matrix, x::Vector)
        fill!(fjac, 0)
        temp1 = x[2] - 3x[1]^2
        temp2 = x[4] - 3x[3]^2
        fjac[1,1] = -c3*temp1 + 1
        fjac[1,2] = -c3*x[1]
        fjac[2,1] = -2*c3*x[1]
        fjac[2,2] = c3 + c4
        fjac[2,4] = c5
        fjac[3,3] = -c6*temp2 + 1
        fjac[3,4] = -c6*x[3]
        fjac[4,2] = c5
        fjac[4,3] = -2*c6*x[3]
        fjac[4,4] = c6 + c4
    end
    (DifferentiableMultivariateFunction(f!, g!), [-3.; -1; -3; -1], "Wood")
end

function helical_valley()
    tpi = 8*atan(1)
    c7 = 2.5e-1
    c8 = 5e-1

    function f!(fvec::Vector, x::Vector)
        if x[1] > 0
            temp1 = atan(x[2]/x[1])/tpi
        elseif x[1] < 0
            temp1 = atan(x[2]/x[1])/tpi + c8
        else
            temp1 = c7*sign(x[2])
        end
        temp2 = sqrt(x[1]^2+x[2]^2)
        fvec[1] = 10(x[3] - 10*temp1)
        fvec[2] = 10(temp2 - 1)
        fvec[3] = x[3]
    end

    function g!(fjac::Matrix, x::Vector)
        temp = x[1]^2 + x[2]^2
        temp1 = tpi*temp
        temp2 = sqrt(temp)
        fjac[1,1] = 100*x[2]/temp1
        fjac[1,2] = -100*x[1]/temp1
        fjac[1,3] = 10
        fjac[2,1] = 10*x[1]/temp2
        fjac[2,2] = 10*x[2]/temp2
        fjac[2,3] = 0
        fjac[3,1] = 0
        fjac[3,2] = 0
        fjac[3,3] = 1
    end
    (DifferentiableMultivariateFunction(f!, g!), [-1.; 0; 0], "Helical Valley")
end

function watson(n::Integer)
    c9 = 2.9e1

    function f!(fvec::Vector, x::Vector)
        fill!(fvec, 0)
        for i = 1:29
            ti = i/c9
            sum1 = 0.0
            temp = 1.0
            for j = 2:n
                sum1 += (j-1)*temp*x[j]
                temp *= ti
            end
            sum2 = 0.0
            temp = 1.0
            for j = 1:n
                sum2 += temp*x[j]
                temp *= ti
            end
            temp1 = sum1 - sum2^2 - 1
            temp2 = 2*ti*sum2
            temp = 1/ti
            for k = 1:n
                fvec[k] += temp*(k-1-temp2)*temp1
                temp *= ti
            end
        end
        temp = x[2] - x[1]^2 - 1
        fvec[1] += x[1]*(1-2temp)
        fvec[2] += temp
    end

    function g!(fjac::Matrix, x::Vector)
        fill!(fjac, 0)
        for i = 1:29
            ti = i/c9
            sum1 = 0.0
            temp = 1.0
            for j = 2:n
                sum1 += (j-1)*temp*x[j]
                temp *= ti
            end
            sum2 = 0.0
            temp = 1.0

            for j = 1:n
                sum2 += temp*x[j]
                temp *= ti
            end
            temp1 = 2(sum1 - sum2^2 - 1)
            temp2 = 2*sum2
            temp = ti^2
            tk = 1.0

            for k = 1:n
                tj = tk
                for j = k:n
                    fjac[k,j] += tj*((Float64(k-1)/ti - temp2)*(Float64(j-1)/ti - temp2) - temp1)
                    tj *= ti
                end
                tk *= temp
            end
        end
        fjac[1,1] += 6x[1]^2 - 2x[2] + 3
        fjac[1,2] -= 2x[1]
        fjac[2,2] += 1
        for k = 1:n
            for j = k:n
                fjac[j,k] = fjac[k,j]
            end
        end
    end
    (DifferentiableMultivariateFunction(f!, g!), zeros(n), "Watson$n")
end

function chebyquad(n::Integer)
    tk = 1/n

    function f!(fvec::Vector, x::Vector)
        fill!(fvec, 0)
        for j = 1:n
            temp1 = 1.0
            temp2 = 2x[j]-1
            temp = 2temp2
            for i = 1:n
                fvec[i] += temp2
                ti = temp*temp2 - temp1
                temp1 = temp2
                temp2 = ti
            end
        end
        iev = -1.0
        for k = 1:n
            fvec[k] *= tk
            if iev > 0
                fvec[k] += 1/(k^2-1)
            end
            iev = -iev
        end
    end

    function g!(fjac::Matrix, x::Vector)
        for j = 1:n
            temp1 = 1.
            temp2 = 2x[j] - 1
            temp = 2*temp2
            temp3 = 0.0
            temp4 = 2.0
            for k = 1:n
                fjac[k,j] = tk*temp4
                ti = 4*temp2 + temp*temp4 - temp3
                temp3 = temp4
                temp4 = ti
                ti = temp*temp2 - temp1
                temp1 = temp2
                temp2 = ti
            end
        end
    end
    (DifferentiableMultivariateFunction(f!, g!), collect(1:n)/(n+1), "Chebyquad")
end

function brown_almost_linear(n::Integer)
    function f!(fvec::Vector, x::Vector)
        sum1 = sum(x) - (n+1)
        for k = 1:(n-1)
            fvec[k] = x[k] + sum1
        end
        fvec[n] = prod(x) - 1
    end

    function g!(fjac::Matrix, x::Vector)
        fill!(fjac, 1)
        fjac[diagind(fjac)] .= 2
        prd = prod(x)
        for j = 1:n
            if x[j] == 0.0
                fjac[n,j] = 1.
                for k = 1:n
                    if k != j
                        fjac[n,j] *= x[k]
                    end
                end
            else
                fjac[n,j] = prd/x[j]
            end
        end
    end
    (DifferentiableMultivariateFunction(f!, g!), 0.5*ones(n), "Brown almost-linear$n")
end

function discrete_boundary_value(n::Integer)
    h = 1/(n+1)

    function f!(fvec::Vector, x::Vector)
        for k = 1:n
            temp = (x[k] + k*h + 1)^3
            if k != 1
                temp1 = x[k-1]
            else
                temp1 = 0.0
            end
            if k != n
                temp2 = x[k+1]
            else
                temp2 = 0.0
            end
            fvec[k] = 2x[k] - temp1 - temp2 + temp*h^2/2
        end
    end

    function g!(fjac::Matrix, x::Vector)
        for k = 1:n
            temp = 3*(x[k]+k*h+1)^2
            for j = 1:n
                fjac[k,j] = 0
            end
            fjac[k,k] = 2 + temp*h^2/2
            if k != 1
                fjac[k,k-1] = -1
            end
            if k != n
                fjac[k,k+1] = -1
            end
        end
    end
    initial_x = collect(1:n)*h
    initial_x = initial_x .* (initial_x .- 1)

    (DifferentiableMultivariateFunction(f!, g!), initial_x, "Discrete boundary value")
end

function discrete_integral_equation(n::Integer)
    h = 1/(n+1)

    function f!(fvec::Vector, x::Vector)
        for k = 1:n
            tk = k*h
            sum1 = 0.0
            for j = 1:k
                tj = j*h
                sum1 += tj*(x[j] + tj + 1)^3
            end
            sum2 = 0.0
            kp1 = k+1
            if n >= kp1
                for j = kp1:n
                    tj = j*h
                    sum2 += (1-tj)*(x[j] + tj + 1)^3
                end
            end
            fvec[k] = x[k] + h*((1-tk)*sum1 + tk*sum2)/2
        end
    end

    function g!(fjac::Matrix, x::Vector)
        for k = 1:n
            tk = k*h
            for j = 1:n
                tj = j*h
                fjac[k,j] = h*min(tj*(1-tk), tk*(1-tj))*3(x[j] + tj + 1)^2/2
            end
            fjac[k,k] += 1
        end
    end

    initial_x = collect(1:n)*h
    initial_x = initial_x .* (initial_x .- 1)

    (DifferentiableMultivariateFunction(f!, g!), initial_x, "Discrete integral equation")
end

function trigonometric(n::Integer)
    function f!(fvec::Vector, x::Vector)
        for j = 1:n
            fvec[j] = cos(x[j])
        end
        sum1 = sum(fvec)
        for k = 1:n
            fvec[k] = n+k-sin(x[k]) - sum1 - k*fvec[k]
        end
    end

    function g!(fjac::Matrix, x::Vector)
        for j = 1:n
            temp = sin(x[j])
            for k = 1:n
                fjac[k,j] = temp
            end
            fjac[j,j] = (j+1)*temp - cos(x[j])
        end
    end
    (DifferentiableMultivariateFunction(f!, g!), ones(n)/n, "Trigonometric")
end

function variably_dimensioned(n::Integer)
    function f!(fvec::Vector, x::Vector)
        sum1 = 0.0
        for j = 1:n
            sum1 += j*(x[j]-1)
        end
        temp = sum1*(1+2sum1^2)
        for k = 1:n
            fvec[k] = x[k] - 1 + k*temp
        end
    end

    function g!(fjac::Matrix, x::Vector)
        sum1 = 0.0
        for j = 1:n
            sum1 += j*(x[j]-1)
        end
        temp = 1 + 6sum1^2
        for k = 1:n
            for j = k:n
                fjac[k,j] = k*j*temp
                fjac[j,k] = fjac[k,j]
            end
            fjac[k,k] += 1
        end
    end
    (DifferentiableMultivariateFunction(f!, g!), 1 .- collect(1:n)/n, "Variably dimensioned")
end

function broyden_tridiagonal(n::Integer)
    function f!(fvec::Vector, x::Vector)
        for k = 1:n
            temp = (3-2x[k])*x[k]
            if k != 1
                temp1 = x[k-1]
            else
                temp1 = 0.0
            end
            if k != n
                temp2 = x[k+1]
            else
                temp2 = 0.0
            end
            fvec[k] = temp - temp1 - 2temp2 + 1
        end
    end

    function g!(fjac::Matrix, x::Vector)
        fill!(fjac, 0)
        for k = 1:n
            fjac[k,k] = 3-4x[k]
            if k != 1
                fjac[k,k-1] = -1
            end
            if k != n
                fjac[k,k+1] = -2
            end
        end
    end
    (DifferentiableMultivariateFunction(f!, g!), -ones(n), "Broyden tridiagonal")
end

function broyden_banded(n::Integer)
    ml = 5
    mu = 1

    function f!(fvec::Vector, x::Vector)
        for k = 1:n
            k1 = max(1, k-ml)
            k2 = min(k+mu, n)
            temp = 0.0
            for j = k1:k2
                if j != k
                    temp += x[j]*(1+x[j])
                end
            end
            fvec[k] = x[k]*(2+5x[k]^2) + 1 - temp
        end
    end

    function g!(fjac::Matrix, x::Vector)
        fill!(fjac, 0)
        for k = 1:n
            k1 = max(1, k-ml)
            k2 = min(k+mu, n)
            for j = k1:k2
                if j != k
                    fjac[k,j] = -(1+2x[j])
                end
            end
            fjac[k,k] = 2+15x[k]^2
        end
    end
    (DifferentiableMultivariateFunction(f!, g!), -ones(n), "Broyden banded")
end

alltests = [rosenbrock();
            powell_singular();
            powell_badly_scaled();
            wood();
            helical_valley();
            watson(6);
            watson(9);
            chebyquad(5);
            chebyquad(6);
            chebyquad(7);
            #chebyquad(8);
            chebyquad(9);
            brown_almost_linear(10);
            brown_almost_linear(30);
            brown_almost_linear(40);
            discrete_boundary_value(10);
            discrete_integral_equation(1);
            discrete_integral_equation(10);
            trigonometric(10);
            variably_dimensioned(10);
            broyden_tridiagonal(10);
            broyden_banded(10);
           ]

if PRINT_FILE; f_out = open("minpack_results.dat", "w"); end

@printf("%-45s   %5s   %5s   %5s   %14s     %10s\n", "Function", "Dim", "NFEV",
        "NJEV", "Final inf-norm", "total time")
println("-"^86)

if PRINT_FILE
    @printf(f_out, "%-45s   %5s   %5s   %5s   %14s\n", "Function", "Dim", "NFEV",
            "NJEV", "Final inf-norm")
    println(f_out, "-"^86)
end

for (df, initial, name) in alltests
    for method in (:hybr, :lm)
        if name == "Trigonometric" && method == :hybr
            continue
        end
        if (Sys.WORD_SIZE == 32 &&
            name in ["Brown almost-linear30", "Brown almost-linear40"] &&
            method == :hybr )
            continue
        end
        tot_time = @elapsed r = fsolve(df.f!, df.g!, initial, method=method, tracing=true)
        @printf("%-45s   %5d   %5d   %5d   %14e   %10e\n", name*"-"*string(method), length(initial),
                r.trace.f_calls, r.trace.g_calls, norm(r.f, Inf), tot_time)
        # @show name, method, r.converged
        @test r.converged
        # without gradient
        if name == "Watson9" && method == :lm
            continue
        end
        tot_time = @elapsed r_no_g = fsolve(df.f!, initial, method=method, tracing=true)
        @printf("%-45s   %5d   %5d   %5d   %14e   %10e\n", name*"-"*string(method)*"-no_g",
                length(initial), r_no_g.trace.f_calls, r_no_g.trace.g_calls, norm(r_no_g.f, Inf), tot_time)
        if PRINT_FILE
            @printf(f_out, "%-45s   %5d   %5d   %5d   %14e\n", name*"-"*string(method)*"-no_g", length(initial),
                    r_no_g.trace.f_calls, r_no_g.trace.g_calls, norm(r_no_g.f, Inf))
            @printf(f_out, "%-45s   %5d   %5d   %5d   %14e\n", name*"-"*string(method), length(initial),
                    r.trace.f_calls, r.trace.g_calls, norm(r.f, Inf))
        end
        # @show name, method, r_no_g.converged
        @test r_no_g.converged
    end
end

if PRINT_FILE; close(f_out); end

end


####### Payoffs definition

### Spot
include("payoffs/spot.jl")

### European Payoffs
include("payoffs/forward.jl")
include("payoffs/european_option.jl")
include("payoffs/binary_european_option.jl")

### Barrier Payoffs
include("payoffs/barrier_do_option.jl")
include("payoffs/barrier_di_option.jl")
include("payoffs/barrier_uo_option.jl")
include("payoffs/barrier_ui_option.jl")
include("payoffs/double_barrier_option.jl")

### American Payoffs
include("payoffs/american_option.jl")
include("payoffs/binary_american_option.jl")

### Path Dependent Payoffs
include("payoffs/bermudan_option.jl")

### Asian Payoffs
include("payoffs/asian_fixed_strike_option.jl")
include("payoffs/asian_floating_strike_option.jl")

### Basket Payoffs
include("payoffs/basket/n_european_option.jl")


using PyCall

@pyimport dolo

modelpy = dolo.yaml_import("examples/models/consumption_savings_iid.yaml")
modelpy.__exogenous__ = modelpy.exogenous.processes[1]

@pyimport dolo.algos.time_iteration_helpers as tih
Fpy = tih.Euler(modelpy)

using Dolo

model = Dolo.yaml_import("examples/models/consumption_savings_iid.yaml")

@time Dolo.time_iteration(model; maxit=500)

@time Dolo.improved_time_iteration(model)

F = Dolo.Euler(model)

# set equal dprocess

Dolo.n_inodes(F.dprocess,1)
Dolo.inode(F.dprocess,1,1)
Dolo.iweight(F.dprocess,1,1)

Fpy.dprocess.n_inodes(1)
Fpy.dprocess.iweight(1,0)

jpoints = [Dolo.inode(F.dprocess, 1, j) for j=1:Dolo.n_inodes(F.dprocess, 1)]
jweights = [Dolo.iweight(F.dprocess, 1, j) for j=1:Dolo.n_inodes(F.dprocess, 1)]

ppoints  = [Fpy.dprocess.inode(0,j) for j=0:6]
pweights  = [Fpy.dprocess.iweight(0,j) for j=0:6]

F.dprocess.integration_weights = pweights
F.dprocess.integration_nodes = copy(hcat(ppoints...)')

# check derivatives without complementarities

flatten(x0::Dolo.MSM) = cat([e for e in x0.data]...;dims=1)

# same residual
x0_py = Fpy.x0
respy = Fpy(x0_py, x0_py)
resj = F(F.x0, F.x0; ignore_constraints=false)
maximum(abs, (respy.data)[:] - flatten(resj))

# same jacobian
rpy, Jpy = Fpy.d_A(x0_py, x0_py)
Jj = Dolo.df_A(F, F.x0, F.x0)
diff = [Jj.data[n] - (Jpy.data)[n,:,:] for n=1:length(Jj.data)]
ddiff = cat(diff...; dims=1)
maximum(abs, ddiff)

# same step
delta_py = Jpy.solve(rpy)
delta_jl = Jj\resj
maximum(abs, (delta_py.data)[:] - flatten(delta_jl))

x0j = F.x0
x0p = Fpy.x0

x0j = x1j
x0p = x1p

x1j = x0j - delta_jl
x1p = x0p - delta_py
maximum(abs, (x1p.data)[:] - flatten(x1j))

res1j = F(x1j, x0j; ignore_constraints=false)
res1p = Fpy(x1p, x0p; ignore_constraints=false)

J1j = Dolo.df_A(F, x1j, x0j)
r1p, J1p = Fpy.d_A(x1p, x0p)

Dolo.norm(res1j)
res1p.norm()

maximum(abs, (res1p.data)[:] - flatten(res1j))

diff1 = [J1j.data[n] - (J1p.data)[n,:,:] for n=1:length(J1j.data)]
ddiff1 = cat(diff...; dims=1)
maximum(abs, ddiff1)

#####

time_iteration(model, maxit=100)

dolo.time_iteration(modelpy, maxit=100)

sol = Dolo.improved_time_iteration(model)

tab = Dolo.tabulate(model, sol.dr, :w)

using SimplePlots

plot(tab[:w], tab[:w], ylims=(0,2))
plot!(tab[:w], tab[:c])


function TriMesh_Search( x, eConn, eAdjacency, point )
#=
   TriMesh_Search - Locates an element that contains a given point.
   The function provides only one such element and makes an
   explicit assumption that the triangulation is Delauney.

   Usage:
     e, r, s = TriMesh_Search(x,eConn,eAdjacency,point)

   where:
     x    is an array of the nodal coordinates.
          Dimension: [ nNodes, 2 ]

     eConn  is the element connectivity.  The standard orientation is assumed
            so the first three nodes describe the vertices.
            Dimension: [ nElements, nLocalNoes ]

     eAdjacency  an [nElements, 3] array containing adjacent elements or -1
                 where -1 indicates a boundary (no adjacent element).

     point  coordinates of a point of interest in the mesh.

     e      the index of the element where "point" is located.

     r,s    the finite element coordinate of "point" in the reference element.

   Licensing:
     This code is distributed under the MIT license.

   Author: Jeff Borggaard, Virginia Tech
           part of FEMfunctions.jl
=#

  nElements = size(eConn,1)

  element = 1   # initial guess of the element containing the point
  elementOld = 1
  count = 0
  found = false
  global iso1 = 0.0
  global iso2 = 0.0
  global iso3 = 0.0
  while ~found
    xLocal = x[ eConn[element,1:3], : ]    # view?

#    plot( [ xLocal[1,1] xLocal[2,1] ],[ xLocal[1,2] xLocal[2,2]],...
#          [ xLocal[2,1] xLocal[3,1] ],[ xLocal[2,2] xLocal[3,2]],...
#          [ xLocal[3,1] xLocal[1,1] ],[ xLocal[3,2] xLocal[1,2]] )

    delta = ( xLocal[3,1]-xLocal[2,1] )*( xLocal[1,2]-xLocal[2,2] ) -
            ( xLocal[3,2]-xLocal[2,2] )*( xLocal[1,1]-xLocal[2,1] )

    global iso1  = ( ( xLocal[2,1]-point[1] )*( xLocal[3,2]-point[2] ) -
              ( xLocal[3,1]-point[1] )*( xLocal[2,2]-point[2] ) ) / delta;

    global iso2  = ( ( xLocal[3,1]-point[1] )*( xLocal[1,2]-point[2] ) -
              ( xLocal[1,1]-point[1] )*( xLocal[3,2]-point[2] ) ) / delta;

    global iso3  = ( ( xLocal[1,1]-point[1] )*( xLocal[2,2]-point[2] ) -
              ( xLocal[2,1]-point[1] )*( xLocal[1,2]-point[2] ) ) / delta;

    if iso1>=0 && iso2>=0 && iso3>=0
      found = true

    else
      count = count + 1;
      index = sortperm([iso1, iso2, iso3])

      if eAdjacency[element,index[1]]>0
        element_old = element;
        element = eAdjacency[element,index[1]]
        choice = 1
      elseif eAdjacency[element,index[2]]>0
        element_old = element
        element = eAdjacency[element,index[2]]
        choice = 2
      elseif eAdjacency[element,index[3]]>0
        element_old = element
        element = eAdjacency[element,index[3]]
        choice = 3;
      else
         # iso1
         # iso2
         # iso3
         # eAdjacency[element,:]
         # index
         error("no next element found, Delaunay mesh?")
      end
#       if ( index(choice)==1 )
#         plot( (xLocal(2,1)+xLocal(3,1))/2, (xLocal(2,2)+xLocal(3,2))/2, '*')
#       elseif ( index(choice)==2 )
#         plot( (xLocal(1,1)+xLocal(3,1))/2, (xLocal(1,2)+xLocal(3,2))/2, '*')
#       else
#         plot( (xLocal(1,1)+xLocal(2,1))/2, (xLocal(1,2)+xLocal(2,2))/2, '*')
#       end
       if count>nElements   # if we are somehow in a loop, get out...
         element = nElements
         count   = 0
       end
    end

  end

  return element, iso2, iso3

end


################################################################################
#  Copyright 2021, Tom Van Acker, Arpan Koirala                                #
################################################################################
# StochasticPowerModels.jl                                                     #
# An extention package of PowerModels.jl for Stochastic (Optimal) Power Flow   #
# See http://github.com/timmyfaraday/StochasticPowerModels.jl                  #
################################################################################

# input data
""
function parse_dst(dst, pa, pb, deg)
    dst == "Beta"    && return _PCE.Beta01OrthoPoly(deg, pa, pb; Nrec=5*deg)
    dst == "Normal"  && return _PCE.GaussOrthoPoly(deg; Nrec=5*deg)
    dst == "Uniform" && return _PCE.Uniform01OrthoPoly(deg; Nrec=5*deg)
end

"""
    StochasticPowerModels.build_stochastic_data(data::Dict{String,Any}, deg::Int)

Function to build the multi-network data representative of the polynomial chaos
expansion of a single-network data dictionary.
"""
function build_stochastic_data(data::Dict{String,Any}, deg::Int)
    # add maximum current
    for (nb, branch) in data["branch"]
        f_bus = branch["f_bus"]
        branch["cmax"] = branch["rate_a"] / data["bus"]["$f_bus"]["vmin"]
    end

    # build mop
    opq = [parse_dst(ns[2]["dst"], ns[2]["pa"], ns[2]["pb"], deg) for ns in data["sdata"]]
    mop = _PCE.MultiOrthoPoly(opq, deg)

    # build load matrix
    Nd, Npce = length(data["load"]), mop.dim
    pd, qd = zeros(Nd, Npce), zeros(Nd, Npce)
    for nd in 1:Nd
        # reactive power
        qd[nd,1] = data["load"]["$nd"]["qd"]
        # active power
        nb = data["load"]["$nd"]["load_bus"]
        ni = data["bus"]["$nb"]["dst_id"]
        if ni == 0
            pd[nd,1] = data["load"]["$nd"]["pd"]
        else
            base = data["baseMVA"]
            μ, σ = data["bus"]["$nb"]["μ"] / base, data["bus"]["$nb"]["σ"] / base
            if mop.uni[ni] isa _PCE.GaussOrthoPoly
                pd[nd,[1,ni+1]] = _PCE.convert2affinePCE(μ, σ, mop.uni[ni])
            else
                pd[nd,[1,ni+1]] = _PCE.convert2affinePCE(μ, σ, mop.uni[ni], kind="μσ")
            end
        end
    end

    # replicate the data
    data = _PM.replicate(data, Npce)

    # add the stochastic data
    data["T2"] = _PCE.Tensor(2,mop)
    data["T3"] = _PCE.Tensor(3,mop)
    data["T4"] = _PCE.Tensor(4,mop)
    data["mop"] = mop
    for nw in 1:Npce, nd in 1:Nd
        data["nw"]["$nw"]["load"]["$nd"]["pd"] = pd[nd,nw]
        data["nw"]["$nw"]["load"]["$nd"]["qd"] = qd[nd,nw]
    end

    return data
end

# output data
"""
    StochasticPowerModels.pce_coeff(result, element::String, id::Int, var::String)

Returns all polynomial chaos coefficients associated with the variable `var` of
the `id`th element `element`.
"""
pce_coeff(result, element::String, id::Int, var::String) =
    [nw[2][element]["$id"][var] for nw in sort(collect(result["solution"]["nw"]), by=x->parse(Int,x[1]))]

"""
    StochasticPowerModels.sample(sdata, result, element::String, id::Int, var::String; sample_size::Int=1000)

Return an `sample_size` sample of the variable `var` of the `id`th element
`element`.
"""
sample(result, element::String, id::Int, var::String; sample_size::Int=1000) =
    _PCE.samplePCE(sample_size, pce_coeff(result, element, id, var), result["mop"])

"""
    StochasticPowerModels.density(sdata, result, element::String, id::Int, var::String; sample_size::Int=1000)

Return an kernel density estimate of the variable `var` of the `id`th element
`element`.
"""
density(result, element::String, id::Int, var::String; sample_size::Int=1000) =
    _KDE.kde(sample(result, element, id, var; sample_size=sample_size))

function print_summary(obj::Dict{String,<:Any}; kwargs...)
    if _IM.ismultinetwork(obj)
        for (n,nw) in obj["nw"]
            println("----------------")
            println("PCE index $n")
            _PM.summary(stdout, nw; kwargs...)
        end
    end
end


using Random: shuffle, seed!

seed!(1729)

@testset "Numeric" begin
    @syms a::Integer b c d x::Real y::Number
    @test simplify(x - y) == x + -1*y
    @test simplify(x - sin(y)) == x + -1*sin(y)
    @test simplify(-sin(x)) == -sin(x)
    @test simplify(1 * x * 2) == 2 * x
    @test simplify(1 + x + 2) == 3 + x
    @test simplify(b*b) == b^2 # tests merge_repeats
    @test simplify((a*b)^c) == a^c * b^c

    @test simplify(1x + 2x) == 3x
    @test simplify(3x + 2x) == 5x

    @test simplify(a + b + (x * y) + c + 2 * (x * y) + d)     == (3 * x * y) + a + b + c + d
    @test simplify(a + b + 2 * (x * y) + c + 2 * (x * y) + d) == (4 * x * y) + a + b + c + d

    @test simplify(a * x^y * b * x^d) == (a * b * (x ^ (d + y)))

    @test simplify(a + b + 0*c + d) == a + b + d
    @test simplify(a * b * c^0 * d) == a * b * d
    @test simplify(a * b * 1*c * d) == a * b * c * d
end

@testset "Pythagorean Identities" begin
    @syms a::Integer x::Real y::Number

    @test simplify(cos(x)^2 + 1 + sin(x)^2) == 2
    @test simplify(cos(y)^2 + 1 + sin(y)^2) == 2
    @test simplify(sin(y)^2 + cos(y)^2 + 1) == 2

    @test simplify(1 + y + tan(x)^2) == sec(x)^2 + y
    @test simplify(1 + y + cot(x)^2) == csc(x)^2 + y
end

#@testset "2π Identities" begin
#    @syms a::Integer x::Real y::Number
#
#    @test simplify(cos(x + 2π + a)) == cos(a + x)
#    @test simplify(tan(x + 2π * a)) == tan(x)
#end

@testset "Depth" begin
    @syms x
    R = RuleSet([@rule(sin(~x) => cos(~x)),
                 @rule( ~x + 1 => ~x - 1)])
    @test R(sin(sin(sin(x + 1)))) == cos(cos(cos(x - 1)))
    @test R(sin(sin(sin(x + 1))), depth=2) == cos(cos(sin(x + 1)))
end

@testset "RuleRewriteError" begin
    pred(x) = error("Fail")

    @syms a b

    rs = RuleSet([@rule ~x + ~y::pred => ~x])
    @test_throws SymbolicUtils.RuleRewriteError rs(a+b)
    err = try rs(a+b) catch err; err; end
    @test sprint(io->Base.showerror(io, err)) == "Failed to apply rule ~x + ~(y::pred) => ~x on expression a + b"
end

@testset "timerwrite" begin
    @syms a b c d
    expr1 = foldr((x,y)->rand([*, /])(x,y), rand([a,b,c,d], 100))
    SymbolicUtils.@timerewrite simplify(expr1)
end

@testset "Shuffle Rules" begin
    @syms a::Integer b c d x::Real y::Number
    R1 = RuleSet(SymbolicUtils.SIMPLIFY_RULES)
    R2 = RuleSet(shuffle(R1.rules))
    simplify_shuffle_tester(ex) = (R2 ∘ R1)(ex) == (R1 ∘ R2)(ex)

    for ex ∈ [foldr((x,y)->rand([*, +, -, ^, /])(x,y),
                        rand([a, b, c, d, x, y, 0, 1.0, 2], 5))
              for _ ∈ 1:20]
        @test isequal((R2 ∘ R1)(ex), (R1 ∘ R2)(ex))
    end
end


import POMDPs
using POMDPModelTools

mutable struct HistoryUpdater <: POMDPs.Updater end

# update(up, pomdp, b, a, o)
# function POMDPs.update(up::HistoryUpdater, pomdp::FireWorld, b, a::Array{Int64,1}, o::FireObs)
function POMDPs.update(up::HistoryUpdater, pomdp::FireWorld, b::SparseCat{Array{FireState,1},Array{Float64,1}}, a::Array{Int64,1}, o::FireObs)
    # particle filter without rejection
    states = FireState[]
    probabilities = Array{Float64,1}(undef,0)

    next_states = FireState[]
    weights = Array{Float64,1}(undef,0)

    belief_particles = ParticleCollection(b.vals)
    for i in 1:n_particles(belief_particles)
        s_i = rand(rng, belief_particles)
        sp_gen = rand(rng, transition(pomdp, s_i, a))
        obs_dist = observation(pomdp, a, sp_gen)
        w_i = in_dist_obs(obs_dist, o)
        push!(next_states, sp_gen)
        push!(weights, w_i)
    end
    if sum(weights) == 0 # all next states do not have observation o
        println("All zero. No observation o.")
        weights = ones(length(weights)) * (1/length(weights))
    end
    for_sampling = SparseCat(next_states, weights)

    # resample
    for i in 1:n_particles(belief_particles)
        sp_sampled = sample(next_states, Weights(weights))
        in_sample, idx = in_dist_states(states, sp_sampled)
        if !in_sample # new state
            push!(states, sp_sampled)
            push!(probabilities, pdf(for_sampling, sp_sampled))
        else # state already represented
            probabilities[idx] += probabilities[idx]
        end
    end
    return SparseCat(states, normalize!(probabilities,1))
end

POMDPs.update(up::HistoryUpdater, b::SparseCat{Array{FireState,1},Array{Float64,1}}, a::Array{Int64,1}, o::FireObs) = update(up, pomdp, b, a, o)
POMDPs.initialize_belief(updater::HistoryUpdater, belief::Any) = belief

# returns probability of observing o
function in_dist_obs(obs_dist::SparseCat{Array{FireObs,1},Array{Float64,1}}, o::FireObs)
    in_dist = false
    prob = 0.0
    for obs in obs_dist.vals
        burning = obs.burning
        if burning == o.burning
            in_dist = true
            prob = pdf(obs_dist, obs)
        end
    end
    return prob
end

# returns boolean whether a state is in distribution and its index
function in_dist_states(s_dist::Array{FireState,1}, s::FireState)
    in_dist = false
    idx = 0
    for i in 1:length(s_dist)
        state = s_dist[i]
        if isequal(state, s)
            in_dist = true
            idx = i
        end
    end
    return in_dist, idx
end


abstract type StageCost end

struct StageCosts <: Objective
    cost::Vector{StageCost}
    T::Int
end

g(obj::StageCosts, x, u, t) = 0.0

function objective(obj::StageCosts, x, u)
    T = obj.T
    J = 0.0
    for t = 1:T-1
        J += g(obj, x[t], u[t], t)
    end
    J += g(obj, x[T], nothing, T)
    return J
end

"""
    nonlinear stage cost
"""
struct NonlinearCost <: StageCost
end

NonlinearCosts = Vector{NonlinearCost}

"""
    quadratic stage cost
"""
struct QuadraticCost <: StageCost
    Q
    q
    R
    r
end

QuadtraticCosts = Vector{QuadraticCost}