import LightGraphs
import Polyhedra
import GLPK
import Plots
import Random
import Statistics
import Triangulate

"""
    gen_obstacle(max_side_len)

Creates an obstacle from a given set of points in the unit square.
"""
function gen_obstacle(max_side_len::Real)
    # TODO: Make this function accept a list of points to take the convex hull
    # Faces must be passed in counter-clockwise order
    # Use rational points (convert if necessary). Rid redundant points and compute HRep.
    # Assure obstacles have at least 4 extreme points for pairwise IB to work

    # Julia Notes
    # # "varargs" or slurping up arguments
    # function myAdd(x, y, z...) # the ... here means collect additional arguments into z
    #     if isempty(z)
    #         return x + y
    #     end
    #     return x + myAdd(y, z...)
    #     # note the function calls it self (recursion)
    #     # z is a container; the ... is necessary so that its contents can be passed as arguments to myAdd
    # end

    # This might also be helpful
    # v = Polyhedra.convexhull(map(i -> collect(points[i]), face)...)

    # Remove redundant points with

    point = rand(2)

    v = Polyhedra.convexhull(point, map(x -> x > 1 ? 1 : x, point + max_side_len * rand(2)), map(x -> x > 1 ? 1 : x, point + max_side_len * rand(2)))

    if Polyhedra.npoints(poly) < 4
        # something
    end

    poly = Polyhedra.polyhedron(v, Polyhedra.DefaultLibrary{Rational{Int64}}(GLPK.Optimizer))
    Polyhedra.removevredundancy!(poly)
    Polyhedra.hrep(poly)
    return poly

end

"""
    gen_obstacle_four(point_per_dim)

Creates an obstacle with four extreme points from a grid of points in the unit square.
"""
function gen_obstacle_four(points_per_dim::Int)
    points_per_dim = max(points_per_dim, 6)  # Need at least 6 interior points
    nums = LinRange(0//1, 1//1, points_per_dim+2)[2:end-1]  # Do not want obstacles touching the unit square boundaries

    # Select the bottom left point at random and select the other three points with respect to it
    x_index = rand(2:(length(nums)-4))
    y_index = rand(2:(length(nums)-4))
    bottom_left = [ nums[x_index] ; nums[y_index] ]
    bottom_right = [ nums[x_index + 3 + rand(-1:1)] ; nums[y_index + rand(-1:1)] ]
    top_left = [ nums[x_index + rand(-1:1)] ; nums[y_index + 3 + rand(-1:1)] ]
    top_right = [ nums[x_index + 3 + rand(-1:1)] ; nums[y_index + 3 + rand(-1:1)] ]

    v = Polyhedra.convexhull(top_left, top_right, bottom_left, bottom_right)
    poly = Polyhedra.polyhedron(v, Polyhedra.DefaultLibrary{Rational{Int64}}(GLPK.Optimizer))
    Polyhedra.removevredundancy!(poly)
    # Loop below assures the polyhedron has 4 extreme points (3 can occur under certain combinations) for pairwise IB purposes
    while Polyhedra.npoints(poly) < 4
        top_right = [ nums[x_index + 3 + rand(-1:1)] ; nums[y_index + 3 + rand(-1:1)] ]
        v = Polyhedra.convexhull(top_left, top_right, bottom_left, bottom_right)
        poly = Polyhedra.polyhedron(v, Polyhedra.DefaultLibrary{Rational{Int64}}(GLPK.Optimizer))
        Polyhedra.removevredundancy!(poly)
    end

    Polyhedra.hrep(poly)
    return poly
end

"""
    gen_field(num_obstacles; custom, seed)

Creates an obstacle course, an array of obstacles in the unit square. If
obstacles overlap, the convex hull of their combined extreme points will be taken
instead.
"""
function gen_field(num_obstacles::Int; custom::Bool=false, points_it::Set{Vector{T}}=Set{Vector{Rational}}(), seed::Int=1) where {T}
    if custom
        obstacles = []
        for points in points_it
            push!(obstacles, gen_obstacle(points))
        end
    else
        Random.seed!(seed)
        points_per_dim = 20
        obstacles = map(_n -> gen_obstacle_four(points_per_dim), 1:num_obstacles)
    end

    return remove_overlaps(obstacles)
end

"""
    remove_overlaps(obstacles)

Given an obstacle course, check for overlapping obstacles and replace offending
obstacles with the convex hull of their combined extreme points.
"""
function remove_overlaps(obstacles)
    if !has_overlaps(obstacles)
        return obstacles
    end

    res = []
    fused = []

    for i = 1:length(obstacles)
        if i in fused
            continue
        end

        accumulator = obstacles[i]
        for j = 1:length(obstacles)
            if i != j && !(j in fused)
                other = obstacles[j]

                overlap = Polyhedra.intersect(accumulator, other)
                if !isempty(overlap)
                    accumulator = Polyhedra.convexhull(accumulator, other)
                    Polyhedra.removevredundancy!(accumulator)
                    Polyhedra.removehredundancy!(accumulator)
                    push!(fused, j)
                end
            end
        end

        push!(res, accumulator)
    end

    return remove_overlaps(res)
end

"""
    has_overlaps(obstacles)

Returns true if any obstacles in the set are overlapping.
"""
function has_overlaps(obstacles)::Bool
    for i = 1:length(obstacles)
        for j = 1:length(obstacles)
            if i != j
                if !isempty(Polyhedra.intersect(obstacles[i], obstacles[j]))
                    return true
                end
            end
        end
    end

    return false
end

"""
    find_intersections(obstacles)

Analyzes the halfspaces that make up each obstacle and unit square border, and finds
where they intersect. Returns an array of all intersections, as well as an array
mapping each halfspace to the points of intersection that lie on it, sorted. Points
in the interior of obstacles are stored at the end of the array.
"""
function find_intersections(obstacles)
    res = []

    halfspaces = @pipe map(obstacle -> Polyhedra.hrep(obstacle).halfspaces, obstacles) |> Iterators.flatten(_) |> collect(_)
    unique!(halfspaces)

    boundaries = [
        Polyhedra.HalfSpace([-1//1, 0//1], 0//1),
        Polyhedra.HalfSpace([1//1, 0//1], 1//1),
        Polyhedra.HalfSpace([0//1, -1//1], 0//1),
        Polyhedra.HalfSpace([0//1, 1//1], 1//1)
    ]
    filter!(b -> !(b in halfspaces), boundaries)

    lines = [halfspaces; boundaries]

    # Save indices of horizontal and vertical halfspaces to keep track of which points need
    # the same y- or x- coordinates (necessary to sort correctly)
    vert_ind = []
    hor_ind = []
    for i = 1:length(lines)
        if abs(lines[i].a'[1]) == 1//1 && abs(lines[i].a'[2]) == 0//1 # set tolerance?
            push!(vert_ind, i)
        end
        if abs(lines[i].a'[1]) == 0//1 && abs(lines[i].a'[2]) == 1//1
            push!(hor_ind, i)
        end
    end

    points = fill([], length(lines))    # For storing the points on each halfspace
    for i = 1:length(lines)
        points[i] = []
    end

    # Function for determining if a point should be considered a duplicate of another. If so, the index
    # of the existing one is returned
    function point_duplicates(point, res)
        tol = Rational(1.e-13)
        for index in 1:length(res)
            if abs(point.first - res[index].first) < tol && abs(point.second - res[index].second) < tol
                return true, index
            end
        end
        return false, -1
    end

    tol = Rational(1.e-13)
    x_vert = Dict([]) # keys are the indices in "lines"
    y_hor = Dict([])
    for i = 1:length(lines)
        for j = (i + 1):length(lines)
            A = vcat(lines[i].a', lines[j].a')
            b = [lines[i].β; lines[j].β]
            try
                x = A \ b

                # Assures boundary points do not have floating point errors
                if abs(x[1] - 0//1) < tol
                    x[1] = 0//1
                elseif abs(x[1] - 1//1) < tol
                    x[1] = 1//1
                end

                if abs(x[2] - 0//1) < tol
                    x[2] = 0//1
                elseif abs(x[2] - 1//1) < tol
                    x[2] = 1//1
                end

                if (x[1] >= 0//1 && x[1] <= 1//1 && x[2] >= 0//1 && x[2] <= 1//1)
                    # Assures points on vertical or horizontal halfspaces have the same x- or y- coordinates, resp.
                    if i in vert_ind
                        if !(i in keys(x_vert))
                            x_vert[i] = x[1]
                        else
                            x[1] = x_vert[i]
                        end
                    elseif i in hor_ind
                        if !(i in keys(y_hor))
                            y_hor[i] = x[2]
                        else
                            x[2] = y_hor[i]
                        end
                    end
                    if j in vert_ind
                        if !(j in keys(x_vert))
                            x_vert[j] = x[1]
                        else
                            x[1] = x_vert[j]
                        end
                    elseif j in hor_ind
                        if !(j in keys(y_hor))
                            y_hor[j] = x[2]
                        else
                            x[2] = y_hor[j]
                        end
                    end

                    point = x[1] => x[2]

                    duplicate, index = point_duplicates(point, res)
                    if !(duplicate)
                        # If point is not a duplicate, push it into res and the halfspaces it is in
                        push!(res, point)
                        push!(points[i], point)
                        push!(points[j], point)
                    else
                        # If point is a duplicate, push it into the halfspaces it's in if not added yet
                        if !(res[index] in points[i])
                            push!(points[i], res[index])
                        end
                        if !(res[index] in points[j])
                            push!(points[j], res[index])
                        end
                    end
                end
            catch e
                # No intersections found
            end
        end
    end

    # Find vertices in the interior of an obstacle
    points_vec = res
    inside_indices = filter(i -> begin
        inside = false
        for ob in obstacles
            if Polyhedra.ininterior([points_vec[i].first ; points_vec[i].second], ob)
                inside = true
                # println("Inside Point: $(points_vec[i].first), $(points_vec[i].second)")
                break
            end
        end
        return inside
    end, 1:length(points_vec))

    # Move inside vertices to the end of the points vector, so that they may be removed from the points
    # vector and the lightgraph graph without altering the vertex labels and edges remaining.
    # Removal does not occur here.
    for i = 1:(length(inside_indices))
        temp = points_vec[inside_indices[i]]
        points_vec[inside_indices[i]] = points_vec[end - (i-1)]
        points_vec[end - (i-1)] = temp
    end

    sorted_points = map(list -> sort(list), points)

    # Third return value is for knowing how many inside points there are
    return (points_vec, sorted_points, length(inside_indices))
end

"""
    find_points(obs)
Compute the points corresponding to obstacle extreme points and the four corners
of the unit square.
"""
function find_points(obs)
    obstacle_faces = []   # Store obstacle face vertices (should be ordered)
    points = []   # Store points as Vector of Pairs
    vertex_num = 0
    for ob in obs
        vertices_in_ob = []
        iter = Polyhedra.points(ob)
        # Store x and y coordinates in points
        for point in iter
            push!(points, point[1] => point[2])
            vertex_num += 1
            push!(vertices_in_ob, vertex_num)
        end

        push!(obstacle_faces, vertices_in_ob)
    end

    points = [points; 0//1 => 0//1 ; 0//1 => 1//1; 1//1 => 1//1; 1//1 => 0//1]
    # unique!(points)   # Code above relies on the points being unique

    return points
end

"""
    construct_graph(obs)

Given a list of obstacles with no overlaps, finds the intersections of the lines
that make up each obstacle's halfspaces and constructs a graph with the
intersections as nodes and the halfspaces as edges. Returns a tuple containing
the obstacles, a list of the location of the nodes' coordinates, the graph, a set of
every face of obstacle space, and a set of every face of free space. (Each face is
a clockwise-ordered vector of vertices).
"""
function construct_graph(obs)
    points, mapped, inside_quant = ClutteredEnvPathOpt.find_intersections(obs)

    # Create map from point to its neighbors
    neighbors = ClutteredEnvPathOpt._find_neighbors(points, mapped)
    # Create lightgraph
    graph = ClutteredEnvPathOpt._gen_graph(neighbors)

    # TODO: Tolerances issues due to approximating an irrational with a rational
    # angles[i,j] is the angle point j makes with point i
    angles = fill(Rational(Inf), length(points), length(points))
    for i in 1:length(points)
        for j in 1:length(points)
            angles[i, j] = Rational(atan(points[j].second - points[i].second, points[j].first - points[i].first))
            if angles[i, j] < 0//1
                angles[i, j] += 2 * pi # Rational(2*3.141592653589793)
            end
        end
    end

    function greatest_angle_neighbor(source, last_angle, visited)
        #tol = Rational(1.e-15)
        unvisited = filter(vertex -> !(vertex in visited), neighbors[source])
        neighbor_angles = map(neighbor -> angles[source, neighbor], unvisited)
        # For sorting angles
        for i in 1:length(neighbor_angles)
            if neighbor_angles[i] >= Rational(3.141592653589793)  # Use pi?
                neighbor_angles[i] -= Rational(2*3.141592653589793)
            end
        end
        zipped = @pipe zip(unvisited, neighbor_angles) |> collect(_)
        # zipped is an array of tuples of the form: (unvisited node, angle it makes with source)

        # Two cases to assure we are constructing the most compact face (ie, no face with subfaces)
        if last_angle >= 0//1 && last_angle < Rational(3.141592653589793)
            greater = filter(tup -> (tup[2] > last_angle) && (tup[2] < last_angle + Rational(3.141592653589793)), zipped)
        else
            greater = filter(tup -> (tup[2] > last_angle - Rational(2*3.141592653589793)) && (tup[2] < last_angle - Rational(3.141592653589793)), zipped)
        end

        if isempty(greater)
            return (-1, -1) # bad state
        else
            sort!(greater, by=(tup -> tup[2]))
            destination = greater[end][1]
        end

        return (destination, angles[source, destination])
    end

    faces = []

    for i in 1:length(points)
        start = i

        for j in 1:length(neighbors[start])
            current = neighbors[start][j]
            last_angle = angles[start, current]

            face = [start, current]

            while start != current

                if (length(face) > 2 && (start in neighbors[current]))
                    push!(faces, copy(face))
                    break
                end

                current, last_angle = greatest_angle_neighbor(current, last_angle, face)

                if current == -1
                    break
                end

                push!(face, current)
            end
        end
    end

    # Find irredundant faces
    face_sets = []
    unique_faces = filter(face -> begin
        face_set = Set(face)
        included = false

        if !(face_set in face_sets)
            push!(face_sets, face_set)
            included = true
        end

        return included
    end, faces)

    # Find free faces (by determining which do not overlap obstacles)
    free_faces = filter(face -> begin
        face_set = Set(face)   # face is already a set, luckily taking a Set of a Set doesn't do anything
        # TODO: It's not even used! Comment it and test, then remove

        face_set_v = Polyhedra.convexhull(map(i -> collect(points[i]), face)...)
        face_set_poly = Polyhedra.polyhedron(face_set_v, Polyhedra.DefaultLibrary{Rational{Int64}}(GLPK.Optimizer))
        face_set_overlaps_obs_faces = false
        included = false

        # If the intersection of face and an obstacle is the face itself, than that face is a subface of an obstacle
        tol =  Rational(1.e-10)
        for n in 1:length(obs)  # need to loop in this manner for 'continue' to work properly
            indicator = 0
            intersect_poly = Polyhedra.intersect(face_set_poly, obs[n])
            Polyhedra.vrep(intersect_poly)  # compute V-representation

            if isempty(Polyhedra.points(intersect_poly)) == true
                continue  # intersection is empty
            end

            ext_p_intersect_poly = sort(collect(intersect_poly.vrep.points.points))
            ext_p_face_set_poly = sort(collect(face_set_poly.vrep.points.points))
            if length(ext_p_intersect_poly) != length(ext_p_face_set_poly)
                continue  # intersection is not the face
            end

            for i in 1:length(ext_p_face_set_poly)
                for j in 1:length(ext_p_face_set_poly[i])  # should be 1:2 (x and y coordinates)
                    if abs(ext_p_intersect_poly[i][j] - ext_p_face_set_poly[i][j]) > tol
                        indicator += 1  # the extreme points do not coincide at the x or y coordinate
                        break
                    end
                end
                if indicator != 0
                    break
                end
            end
            if indicator == 0
                face_set_overlaps_obs_faces = true
                break
            end
        end

        if !(face_set_overlaps_obs_faces)
            included = true
        end

        return included
    end, unique_faces)

    # Swap obstacle (sub)faces with whole faces, then remove inside nodes and inside edges
    # CHAD SUBSET APPROACH

    # Remove inside vertices
    if inside_quant > 0
        length_p = length(points)
        for i = 1:inside_quant
            LightGraphs.rem_vertex!(graph, length_p - (i-1))
            pop!(points)
        end
    end

    # Determine the vertices that make up each obstacle's extreme points
    obs_euclid_points = map(ob -> Polyhedra.points(ob), obs)
    obstacle_faces = map(ob_points -> begin
        tol = Rational(1e-13)
        face_points = []
        for ob_point in ob_points
            vertex_index = findfirst(p -> abs(ob_point[1] - p.first) < tol && abs(ob_point[2] - p.second) < tol, points)
            if vertex_index !== nothing
                push!(face_points, vertex_index)
            end
        end
        return face_points
    end, obs_euclid_points)

    # Find all vertices associated with an obstacle face and remove inside edges

    boundary_edges = fill([], length(obs))  # store boundary edges of each obstacle
    for i = 1:length(obs)
        boundary_edges[i] = []
    end

    edges_to_remove = []  # edges crossing through an obstacle are to be removed
    for edge in LightGraphs.edges(graph)
        on_boundary = true
        for i = 1:length(obs)
            if Polyhedra.in([points[edge.src].first; points[edge.src].second], obs[i]) && Polyhedra.in([points[edge.dst].first; points[edge.dst].second], obs[i])
                # Edge is contained in obstacle i
                # println("Found Edge: $(edge.src) => $(edge.dst) with i = $i")
                push!(obstacle_faces[i], edge.src)
                push!(obstacle_faces[i], edge.dst)
                unique!(obstacle_faces[i])
                if Polyhedra.ininterior([(points[edge.src].first + points[edge.dst].first)//2 ; (points[edge.src].second + points[edge.dst].second)//2], obs[i])
                    # Edge crosses through obstacle i
                    #println("Removing edge $(edge.src) => $(edge.dst)")
                    push!(edges_to_remove, (edge.src, edge.dst))
                    on_boundary = false
                end
                if on_boundary   # this can be change to an if-else with the above
                    # println("Adding edge $(edge.src) => $(edge.dst) with i = $i")
                    push!(boundary_edges[i], edge)
                end
                break
            end
        end
    end

    for (a,b) in edges_to_remove
        # println("Removing edge $a => $b")
        LightGraphs.rem_edge!(graph, a, b)
    end

    # for i in 1:length(boundary_edges)
    #     println("Obstacle $i")
    #     for edge in boundary_edges[i]
    #         println("$edge")
    #     end
    # end

    # Now need to order points on obstacle faces counter-clockwise (they're then returned clockwise)

    # Start with left-most point (if multiple, use bottom-most)
    starting_points = []
    for ob_face in obstacle_faces
        x_coor = map(vertex -> points[vertex].first, ob_face)
        min_x = min(x_coor...)
        candidates = filter(vertex -> points[vertex].first == min_x, ob_face)
        if length(candidates) > 1
            y_coor = map(vertex -> points[vertex].second, candidates)
            min_y = min(y_coor...)
            candidates = filter(vertex -> points[vertex].second == min_y, candidates)
        end
        push!(starting_points, candidates[1])
    end

    # println("Starting Points:")
    # for sp in starting_points
    #     println("$sp")
    # end

    obstacle_faces_ordered = []
    for i = 1:length(starting_points)
        sp = starting_points[i]
        sp_adj = []  #  adjacent nodes (should be exactly 2)
        for edge in boundary_edges[i]
            if sp == edge.src
                push!(sp_adj, edge.dst)
            elseif sp == edge.dst
                push!(sp_adj, edge.src)
            end
        end
        sp_n1 = sp_adj[1]
        sp_n2 = sp_adj[2]

        angle_n1 = angles[sp, sp_n1]
        if angle_n1 >= Rational(3.141592653589793)  # for ordering
            angle_n1 -= Rational(2*3.141592653589793)
        end
        angle_n2 = angles[sp, sp_n2]
        if angle_n2 >= Rational(3.141592653589793)
            angle_n2 -= Rational(2*3.141592653589793)
        end

        # Determine counter-clockwise order of the three points
        n1 = sp
        if angle_n1 > angle_n2
            face = [sp_n1; sp; sp_n2]
            n0 = sp_n1
            n2 = sp_n2
        else
            face = [sp_n2; sp; sp_n1]
            n0 = sp_n2
            n2 = sp_n1
        end

        # From current node n2, determine the next node in the face by checking which other node n2
        # is adjacent to (other than n1). Repeat until the next node is n0.
        while n2 != n0
            next_found = false
            for edge in boundary_edges[i]
                if edge.src == n2 && edge.dst != n1
                    n1 = copy(n2)
                    n2 = edge.dst
                    next_found = true
                    break
                elseif edge.dst == n2 && edge.src != n1
                    n1 = copy(n2)
                    n2 = edge.src
                    next_found = true
                    break
                end
            end
            if !next_found
                println("Infinite while loop. Face may not be complete.")
                break
            end
            if n2 != n0
                push!(face, n2)
            end
        end

        push!(obstacle_faces_ordered, face)
    end

    return (obs, points, graph, Set(map(face -> reverse(face), obstacle_faces_ordered)), Set(map(f -> reverse(f), free_faces)))
end

"""
    construct_graph_delaunay(obs)

Generate the constrained Delaunay triangulation partition of the free space.
Returns a tuple containing the obstacles, a list of the location of the nodes' coordinates,
the graph, a set of every face of obstacle space, and a set of every face of free space.
(Each face is a clockwise-ordered vector of vertices).
"""
function construct_graph_delaunay(obs)

    S = []   # Store segments
    obstacle_faces = []   # Store obstacle face vertices (should be ordered)
    points = []   # Store points as Vector of Pairs
    vertex_num = 0
    for ob in obs
        vertices_in_ob = []
        iter = Polyhedra.points(ob)
        # Store x and y coordinates in points
        for point in iter
            push!(points, point[1] => point[2])
            vertex_num += 1
            push!(vertices_in_ob, vertex_num)
        end
        # Push edges of obstacle boundary (assumes they're in order)
        for i = 1:(length(vertices_in_ob)-1)
            push!(S, [vertices_in_ob[i] ; vertices_in_ob[i+1]])
        end
        push!(S, [vertices_in_ob[end] ; vertices_in_ob[1]])

        push!(obstacle_faces, vertices_in_ob)
    end

    points = [points; 0//1 => 0//1 ; 0//1 => 1//1; 1//1 => 1//1; 1//1 => 0//1]
    # unique!(points)   # Code above relies on the points being unique

    # Add boundary segments
    num_points = length(points)
    for i = (num_points-3):(num_points-1)
        push!(S, [i ; i+1])
    end
    push!(S, [num_points ; num_points-3])

    # Initialize triangulate structure
    triin = Triangulate.TriangulateIO()

    # Store points in matrix
    C = zeros(2,length(points))
    for j = 1:length(points)
        C[:,j] = [points[j].first ; points[j].second]
    end
    triin.pointlist = Matrix{Cdouble}(C)

    # Store segments in matrix
    segments = zeros(Int,2,length(S))
    for j=1:length(S)
        segments[:,j] = [S[j][1]; S[j][2]]
    end
    triin.segmentlist = Matrix{Cint}(segments)

    # Add hole list
    hole_center = zeros(2,length(obstacle_faces))
    for (j,face) in enumerate(obstacle_faces)
        x_locations = map(i -> points[i].first, face)
        y_locations = map(i -> points[i].second, face)
        avg_x = Statistics.mean(x_locations)
        avg_y = Statistics.mean(y_locations)
        hole_center[:,j] = [avg_x; avg_y]
    end
    triin.holelist = Matrix{Cdouble}(hole_center)

    # Computed CDT
    (triout, _) = Triangulate.triangulate("pQ", triin)

    triangles = Matrix{Int64}(triout.trianglelist)   # matrix with 3 rows
    # edges = Matrix{Int64}(triout.edgelist)   # matrix with 2 rows

    # Create graph
    graph = LightGraphs.SimpleGraph(length(points))
    # Add edges
    for j = 1:size(triangles,2)
        LightGraphs.add_edge!(graph, triangles[1,j] => triangles[2,j])
        LightGraphs.add_edge!(graph, triangles[1,j] => triangles[3,j])
        LightGraphs.add_edge!(graph, triangles[2,j] => triangles[3,j])
    end
    # Store free faces as vectors (utilizing holes, all triangles are free faces)
    free_faces = []
    for j = 1:size(triangles,2)
        push!(free_faces, triangles[:,j])
    end

    # TODO: Face merging

    return (obs, points, graph, Set(map(face -> face, obstacle_faces)), Set(map(f -> reverse(f), free_faces)))
end

"""
    _find_neighbors(points, mapped)

Given a list of points and an association of halfspace lines to points that lie
on the line, return an array where res[i] <- the neighbors of node i.
"""
function _find_neighbors(points::Vector{Any}, mapped::Vector{Vector{Any}})::Vector{Vector{Any}}
    # This function should only be used in construct_graph(), before "inside" points are removed
    neighbors = Dict()

    for point in points
        neighbors[point] = []

        for colinear in mapped
            if in(point, colinear)
                index = findfirst(p -> p == point, colinear)
                # Do not push the same point more than once
                if index - 1 > 0 && !(colinear[index - 1] in neighbors[point])
                    push!(neighbors[point], colinear[index - 1])
                end

                if index + 1 <= length(colinear) && !(colinear[index + 1] in neighbors[point])
                    push!(neighbors[point], colinear[index + 1])
                end
            end
        end
    end

    # Res will contain the neighbors by their node label
    res = fill([], length(points))
    for i in 1:length(res)
        res[i] = map(point -> findfirst(p -> p == point, points), neighbors[points[i]])
    end

    return res
end

"""
    _gen_graph(neighbors)

Given an array where res[i] <- the neighbors of node i, return the
corresponding lightgraph.
"""
function _gen_graph(neighbors)
    res = LightGraphs.SimpleGraph(length(neighbors))

    for i in 1:length(neighbors)
        for j in 1:length(neighbors[i])
            LightGraphs.add_edge!(res, i => neighbors[i][j])
        end
    end

    return res
end

"""
    plot_faces(faces, points; plot_name, col, new_plot, individually)

Given a set of vectors corresponding to faces of a planar graph, plots the faces in the unit square.
"""
function plot_faces(faces::Set{Vector{T}}, points::Vector{Any}; plot_name::String="Faces", col::String="green3", new_plot::Bool=true, individually::Bool=false) where {T}
    if new_plot && !individually
        Plots.plot()
    end

    for (j,face) in enumerate(faces)
        v = Polyhedra.convexhull(map(i -> collect(points[i]), face)...)
        x_locations = map(i -> points[i].first, face)
        y_locations = map(i -> points[i].second, face)

        # avg_x = Statistics.mean(x_locations)
        # avg_y = Statistics.mean(y_locations)

        polygon = Polyhedra.polyhedron(
            v,
            Polyhedra.DefaultLibrary{Rational{Int64}}(GLPK.Optimizer)
        )

        if individually
            # Plot each face on a separate plot with node labels
            Plots.plot(polygon, color=col, alpha=0.9)
            Plots.plot!(x_locations, y_locations, series_annotations=([Plots.text(string(x), :center, 8, "courier") for x in face]))
            display(Plots.plot!(xlims=(-0.05,1.05), ylims=(-0.05,1.05), title="Face $j: $face"))
        else
            # Accumulate the faces on the same plot
            Plots.plot!(polygon, color=col, alpha=0.9)
            # Plots.plot!([avg_x], [avg_y], series_annotations=([Plots.text("$j", :center, 8, "courier")]))
            # display(Plots.plot!(polygon, title="Free Face # $j", xlims=(-0.05,1.05), ylims=(-0.05,1.05)))
        end
    end

    if !individually
        display(Plots.plot!(title=plot_name, xlims=(-0.05,1.05), ylims=(-0.05,1.05)))
    end
end

"""
    plot_edges(lg, points; plot_name, col, new_plot)

Given a LabeledGraph, plots its nodes and edges in the unit square.
"""
function plot_edges(lg::LabeledGraph{T}, points::Vector{Any}; plot_name::String="Edges", col::String="colorful", new_plot::Bool=true, vertices::Dict{T,T}=Dict{T,T}()) where {T}
    if new_plot
        Plots.plot()
    end

    # Need to map nodes to the point they refer to in "points"
    rev = ClutteredEnvPathOpt._reverse_labels(lg.labels)
    for edge in LightGraphs.edges(lg.graph)
        if col == "colorful"
            Plots.plot!([points[rev[edge.src]].first, points[rev[edge.dst]].first], [points[rev[edge.src]].second, points[rev[edge.dst]].second],linewidth=2)
            # display(Plots.plot!(title="Edge ($(rev[edge.src]), $(rev[edge.dst]))"))
        else
            Plots.plot!([points[rev[edge.src]].first, points[rev[edge.dst]].first], [points[rev[edge.src]].second, points[rev[edge.dst]].second], color=col,linewidth=2)
            # display(Plots.plot!(title="Edge ($(rev[edge.src]), $(rev[edge.dst]))"))
        end
    end

    # ClutteredEnvPathOpt.plot_intersections(obstacles, vertices=vertices)
    ClutteredEnvPathOpt.plot_points(points, vertices=vertices)

    display(Plots.plot!(title=plot_name, xlims=(-0.05,1.05), ylims=(-0.05,1.05), legend=false))
end

"""
    plot_field(field)

Plots the obstacles to the existing active plot.
"""
function plot_field(field)
    for i = 1:length(field)
        Plots.plot!(field[i], xlims = (-0.05,1.05), ylim = (-0.05, 1.05))
    end
end

"""
    plot_lines(field)

Plots the lines that make up the obstacles' halfspaces to the existing active
plot.
"""
function plot_lines(field)
    halfspaces = @pipe map(obstacle -> Polyhedra.hrep(obstacle).halfspaces, field) |> Iterators.flatten(_) |> collect(_)
    unique!(halfspaces)

    for h in halfspaces
        if abs(h.a[2]) != 0//1
            f = x -> (h.β - h.a[1] * x) // h.a[2]

            x = LinRange(0//1, 1//1, 11)
            y = map(f, x)
        else  # vertical lines
            x = fill(abs(h.β // h.a[1]), 11)
            y = LinRange(0//1, 1//1, 11)
        end

        Plots.plot!(x, y)
    end
end

"""
    plot_borders()

Plots the lines that make up the unit square's borders to the existing active
plot.
"""
function plot_borders()
    halfspaces = [
        Polyhedra.HalfSpace([-1//1, 0//1], 0//1),
        Polyhedra.HalfSpace([1//1, 0//1], 1//1),
        Polyhedra.HalfSpace([0//1, -1//1], 0//1),
        Polyhedra.HalfSpace([0//1, 1//1], 1//1)
    ]

    for h in halfspaces
        if abs(h.a[2]) != 0//1
            f = x -> (h.β - h.a[1] * x) // h.a[2]

            x = LinRange(0//1, 1//1, 11)
            y = map(f, x)
        else  # vertical lines
            x = fill(abs(h.β // h.a[1]), 11)
            y = LinRange(0//1, 1//1, 11)
        end

        Plots.plot!(x, y)
    end
end

"""
    plot_intersections(field)

Plots and labels the points where the lines that make up the obstacles'
halfspaces intersect to the existing active plot.
"""
function plot_intersections(field; vertices::Dict{T,T}=Dict{T,T}()) where {T}
    intersections, _, inside_quant = find_intersections(field)

    # Remove points inside obstacles
    for _ = 1:inside_quant
        pop!(intersections)
    end

    if !isempty(vertices)
        points_filtered = []
        for v in keys(vertices)
            push!(points_filtered, intersections[v])
        end
        x = map(point -> point.first, points_filtered)
        y = map(point -> point.second, points_filtered)
        Plots.scatter!(x,y, color="red", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in keys(vertices)]))
    else
        x = map(point -> point.first, intersections)
        y = map(point -> point.second, intersections)
        Plots.scatter!(x,y, color="red", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in 1:length(points)]))
    end

    # TODO: Delete this once certain it works
    # if !isempty(vertices)
    #     intersections = @pipe filter(v -> v in keys(vertices), 1:length(intersections)) |> intersections[_]
    # end

    # x = map(point -> point[1], intersections)
    # y = map(point -> point[2], intersections)

    # # Plots.scatter!(x,y, color="red3", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in 1:length(x)]))
    # Plots.scatter!(x,y, color="red", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in 1:length(x)]))
end

"""
    plot_points(points; vertices)

Plots and labels points given.
"""
function plot_points(points::Vector{Any}; vertices::Dict{Int,Int}=Dict{Int,Int}())

    Plots.plot!(legend=false)

    if !isempty(vertices)
        points_filtered = []
        for v in keys(vertices)
            push!(points_filtered, points[v])
        end
        x = map(point -> point.first, points_filtered)
        y = map(point -> point.second, points_filtered)
        Plots.scatter!(x,y, color="red", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in keys(vertices)]))
    else
        x = map(point -> point.first, points)
        y = map(point -> point.second, points)
        Plots.scatter!(x,y, color="red", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in 1:length(points)]))
    end

end

"""
    plot_new(n, name; custom, seed, save_image)

Generates a new obstacle course with n obstacles. If custom = true, obstacles will be
created from an array containing the points for each obstacle. If save_image = true,
the course will be plotted and saved to an image.
"""
function plot_new(n::Int, name::String; custom::Bool=false, seed::Int=1, save_image::Bool=false, partition::String="CDT")
    if custom
        # obs = gen_field(num_obstacles::Int; custom::Bool=false, points_it::Set{Vector{T}}=Set{Vector{Rational}}(), seed::Int=1) where {T}
    else
        obs = gen_field(n, seed = seed)
    end

    if save_image
        Plots.plot()
        plot_field(obs)
        points = ClutteredEnvPathOpt.find_points(obs)
        ClutteredEnvPathOpt.plot_points(points)
        #plot_lines(obs)
        #plot_borders()
        #plot_intersections(obs)
        Plots.png(name)
        display(Plots.plot!(title="Field"))
    end

    if partition == "CDT"
        return construct_graph_delaunay(obs)
    else
        return construct_graph(obs)
    end
end


#   Unit examples of the FourierAnalysis Package for julia language
#
#   MIT License
#   Copyright (c) 2019-2022,
#   Marco Congedo, CNRS, Grenobe, France:
#   https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#   This example shows how to compute univariate measures in
#   the time-frequency domain.

#   ~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~  #
#                                                                             #
#   ~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~  #

using FourierAnalysis, FFTW, LinearAlgebra, Statistics, Plots

# add module for reading the two EEG text files to be used ater
push!(LOAD_PATH, @__DIR__)
using IOtxt

# Get EEG file names with complete path (they have extension .txt)
S=getFilesInDir(@__DIR__; ext=(".txt",))

# read the two EEG data files and put them in a Matrix object
X1=readEEG(S[1])
X2=readEEG(S[2])

#############################################################################
# Bivariate measures
# see Congedo, 2018: https://hal.archives-ouvertes.fr/hal-01868538/document

Pz=15
Fz=5
𝐱₁=[X1[:, Pz], X2[:, Pz]] # get the two times-series at electrode Pz
𝐱₂=[X1[:, Fz], X2[:, Fz]] # get the two times-series at electrode Fz
sr, wl, bandwidth=128, 512, 2
𝐘₁=TFanalyticsignal(𝐱₁, sr, wl, bandwidth; fmax=32, nonlinear=false)
𝐘₂=TFanalyticsignal(𝐱₂, sr, wl, bandwidth; fmax=32, nonlinear=false)
𝐀₁=TFamplitude(𝐘₁)
𝐀₂=TFamplitude(𝐘₂)

# Comodulation Eq. 0.11 (Com)
# compute the Com averaging in a TF region from a TFAnalyticSignalVector object
# NB If you compute it from Analytic Signal or Amplitude objects, those
# must be linear (note the nonlinear=false above)
Com=comodulation(𝐘₁, 𝐘₂, (8, 12), (1, 512); mode=mean)
# compute the Com averaging in a TF region from a TFAmplitudeVector object
Com=comodulation(𝐀₁, 𝐀₂, (8, 12), (1, 512); mode=mean)
# compute the Com averaging in a TF region directly from data
# In this care you don't have to worry about linearity of the analytic signal
Com=comodulation(𝐱₁, 𝐱₂, sr, wl, (8, 12), (1, 512), bandwidth; mode=mean)
# You can compute comodulation from smoothed amplitude:
Com=comodulation(𝐱₁, 𝐱₂, sr, wl, (8, 12), (1, 512), bandwidth;
                 mode=mean,
                 fsmoothing=blackmanSmoother,
                 tsmoothing=blackmanSmoother)

# you can go faster pre-computing a FFTW plan.
# This is aslo useful when you have to call the comodulation function several times
plan=Planner(plan_patient, 5, wl, Float64, true)
Com=comodulation(𝐱₁, 𝐱₂, sr, wl, (8, 12), (1, 512), bandwidth; mode=mean, planner=plan)

# compute the Com in a TF region from a TFAnalyticSignalVector object
Com=comodulation(𝐘₁, 𝐘₂, (8, 12), (1, 512); mode=extract)
# compute the Com in a TF region from a TFAmplitudeVector object
Com=comodulation(𝐀₁, 𝐀₂, (8, 12), (1, 512); mode=extract)
# compute the Com in a TF region directly from data
Com=comodulation(𝐱₁, 𝐱₂, sr, wl, (8, 12), (1, 512), bandwidth; mode=extract)

# All these operations can be done also for coherence measures, for example
# avoid conflict with package DSP, which also exports 'coherence'
Coh=FourierAnalysis.coherence(𝐘₁, 𝐘₂, (8, 12), (1, 512); mode=mean)
Coh=FourierAnalysis.coherence(𝐘₁, 𝐘₂, (8, 12), (1, 512); mode=extract)

# Compute all 5 coherence types
Coh=FourierAnalysis.coherence(𝐘₁, 𝐘₂, (8, 12), (1, 512); mode=extract, allkinds=true)

# phase coherence (phase-locking value)
𝐘₁=TFanalyticsignal(𝐱₁, sr, wl, bandwidth; fmax=32, nonlinear=true)
𝐘₂=TFanalyticsignal(𝐱₂, sr, wl, bandwidth; fmax=32, nonlinear=true)
Coh=FourierAnalysis.coherence(𝐘₁, 𝐘₂, (8, 12), (1, 512); mode=mean, nonlinear=true)

# or directly from data (no need to compute non-linear analytic signal in this case)
Coh=FourierAnalysis.coherence(𝐱₁, 𝐱₂, sr, wl, (8, 12), (1, 512), bandwidth; mode=mean, nonlinear=true)

# and also for non-linear meausures
# compute non-linear analyticSignal
𝐘₁=TFanalyticsignal(𝐱₁, sr, wl, bandwidth; fmax=32, nonlinear=true)
𝐘₂=TFanalyticsignal(𝐱₂, sr, wl, bandwidth; fmax=32, nonlinear=true)
# although you are allowed to compute the amplitude of non-linear
# analytic sygnal this way, it does not make sense as the amplitude
# is 1.0 everywhere, as you can check
𝐀₁=TFamplitude(𝐘₁)
𝐀₂=TFamplitude(𝐘₂)

##################################################################


######### ARM Ch03: kid score example  ###########

using StanSample, Test

kid = "
data {
  int<lower=0> N;
  vector[N] kid_score;
  vector[N] mom_hs;
  vector[N] mom_iq;
}
parameters {
  vector[3] beta;
  real<lower=0> sigma;
}
model {
  kid_score ~ normal(beta[1] + beta[2] * mom_hs + beta[3] * mom_iq, sigma);
}
"

kiddata = Dict("N" => 434,
  "kid_score" => [65, 98, 85, 83, 115, 98, 69, 106, 102, 95, 91, 58, 84, 78, 102,
  110, 102, 99, 105, 101, 102, 115, 100, 87, 99, 96, 72, 78, 77, 98, 69, 130, 109,
  106, 92, 100, 107, 86, 90, 110, 107, 113, 65, 102, 103, 111, 42, 100, 67, 92,
  100, 110, 56, 107, 97, 56, 95, 78, 76, 86, 79, 81, 79, 79, 56, 52, 63, 80, 87,
  88, 92, 100, 94, 117, 102, 107, 99, 73, 56, 78, 94, 110, 109, 86, 92, 91, 123,
  102, 105, 114, 96, 66, 104, 108, 84, 83, 83, 92, 109, 95, 93, 114, 106, 87, 65,
  95, 61, 73, 112, 113, 49, 105, 122, 96, 97, 94, 117, 136, 85, 116, 106, 99, 94,
  89, 119, 112, 104, 92, 86, 69, 45, 57, 94, 104, 89, 144, 52, 102, 106, 98, 97,
  94, 111, 100, 105, 90, 98, 121, 106, 121, 102, 64, 99, 81, 69, 84, 104, 104,
  107, 88, 67, 103, 94, 109, 94, 98, 102, 104, 114, 87, 102, 77, 109, 94, 93, 86,
  97, 97, 88, 103, 87, 87, 90, 65, 111, 109, 87, 58, 87, 113, 64, 78, 97, 95, 75,
  91, 99, 108, 95, 100, 85, 97, 108, 90, 100, 82, 94, 95, 119, 98, 100, 112, 136,
  122, 126, 116, 98, 94, 93, 90, 70, 110, 104, 83, 99, 81, 104, 109, 113, 95, 74,
  81, 89, 93, 102, 95, 85, 97, 92, 78, 104, 120, 83, 105, 68, 104, 80, 120, 94,
  81, 101, 61, 68, 110, 89, 98, 113, 50, 57, 86, 83, 106, 106, 104, 78, 99, 91,
  40, 42, 69, 84, 58, 42, 72, 80, 58, 52, 101, 63, 73, 68, 60, 69, 73, 75, 20, 56,
  49, 71, 46, 54, 54, 44, 74, 58, 46, 76, 43, 60, 58, 89, 43, 94, 88, 79, 87, 46,
  95, 92, 42, 62, 52, 101, 97, 85, 98, 94, 90, 72, 92, 75, 83, 64, 101, 82, 77,
  101, 50, 90, 103, 96, 50, 47, 73, 62, 77, 64, 52, 61, 86, 41, 83, 64, 83, 116,
  100, 42, 74, 76, 92, 98, 96, 67, 84, 111, 41, 68, 107, 82, 89, 83, 73, 74, 94,
  58, 76, 61, 38, 100, 84, 99, 86, 94, 90, 50, 112, 58, 87, 76, 68, 110, 88, 87,
  54, 49, 56, 79, 82, 80, 60, 102, 87, 42, 119, 84, 86, 113, 72, 104, 94, 78, 80,
  67, 104, 96, 65, 64, 95, 56, 75, 91, 106, 76, 90, 108, 86, 85, 104, 87, 41, 106,
  76, 100, 89, 42, 102, 104, 59, 93, 94, 76, 50, 88, 70],
  "mom_hs" => [1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1,
  1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1,
  1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1,
  1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 1,
  1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1,
  1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
  1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1,
  0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 1, 0,
  1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1,
  1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1,
  1, 1, 0, 0, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0,
  1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1,
  1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0,
  0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
  1, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1,
  0, 0, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1,
  0, 0, 0, 1, 0, 1, 1],
  "mom_iq" => [121.117528602603, 89.3618817100663, 115.443164881725,
  99.4496394360723, 92.7457099982118, 107.901837758501, 138.893106071162,
  125.145119475328, 81.6195261789843, 95.0730686206496, 88.5769977185567,
  94.8597081943671, 88.9628008509596, 114.114297012333, 100.534071915245,
  120.419145591086, 114.426876891447, 111.592357580831, 133.849227208159,
  97.2648010634673, 110.09680614075, 126.72399416984, 97.9115903092628,
  99.9257251603133, 97.5950080521511, 121.748013460479, 98.7480786989934,
  97.9152543291671, 80.3585564632336, 114.307860633927, 109.138316302971,
  101.817179733939, 117.965104313227, 108.633496913048, 96.528619145464,
  92.8713754661642, 95.8981342875886, 107.015457730577, 87.1970117452858,
  89.3618817100663, 102.530987846246, 130.166860235388, 83.4141025980336,
  125.75346623026, 85.8103765615827, 126.520072628695, 79.203048104857,
  113.165907175696, 110.33138786508, 99.4092237400546, 102.425526488872,
  124.900437749856, 94.8597081943671, 92.3685645501674, 101.817179733939,
  79.8335329627323, 96.226139267492, 82.3554723942338, 76.5756473159814,
  110.013482886319, 111.592357580831, 113.657407368227, 101.164556773121,
  101.341094009698, 101.698163476419, 98.8191545781969, 93.4981963041352,
  87.2938898998626, 87.316028002806, 86.2094574206895, 89.8379674343074,
  121.870014908527, 127.544378629405, 109.991344783376, 90.4463141892395,
  88.8949026973177, 113.046890918176, 78.0131542683097, 95.3896508777613,
  83.6164421099846, 86.2462806103944, 95.3896508777613, 111.252314499127,
  127.014435946011, 89.951950871924, 99.2731021994947, 81.1655785576865,
  110.52128746677, 136.493846917085, 123.862011656634, 134.602392343459,
  91.0989371500582, 103.269371773029, 126.287072933559, 127.544378629405,
  92.4738903034255, 84.7940885713045, 100.534071915245, 108.003012055173,
  78.8071449713118, 121.117528602603, 103.056011346747, 100.437193760669,
  88.5548596156134, 85.3055571716597, 119.448617651507, 90.5517755466136,
  96.3340383364003, 127.644920803886, 119.856558886853, 112.018920897562,
  92.5143059994432, 96.023799755541, 91.4168421288192, 103.07814944969,
  100.243630139074, 89.8158293313641, 115.057361749322, 96.650620593512,
  105.577950778248, 97.7216907075729, 131.835771186485, 101.817179733939,
  93.4981963041352, 104.026539286327, 103.708634307566, 107.368863177393,
  92.5511291891481, 94.4425837627742, 93.1447908573186, 83.4141025980336,
  71.037405135663, 127.66705890683, 115.687846607198, 90.9762568726337,
  132.865336903467, 72.5022963925581, 100.710609151823, 132.688799666889,
  98.5457391870425, 115.665708504254, 117.062799760565, 136.493846917085,
  106.548478717828, 134.126306619218, 104.657024144202, 99.3785635568688,
  100.437193760669, 89.5253875551931, 120.92294779354, 113.143769072753,
  126.414746875437, 122.737163481443, 115.375266728083, 91.6067417305091,
  88.6603209729875, 107.178963575703, 104.969604023316, 103.8630334412,
  110.09680614075, 88.4543174411322, 117.032139577379, 125.656588075683,
  122.601041940883, 121.049630448962, 93.1447908573186, 99.3785635568688,
  123.892807443936, 113.043226898272, 115.771169861628, 114.938345491802,
  93.6208765815597, 121.748013460479, 90.3494360346627, 83.4104385781293,
  99.4864626257772, 113.910375471188, 97.5950080521511, 89.8158293313641,
  119.584739192066, 97.3816476258686, 84.8774118257353, 102.328648334295,
  99.2731021994947, 99.0750497482455, 96.9277000045708, 108.313250636032,
  100.534071915245, 99.4864626257772, 113.987614356911, 92.7702011694918,
  88.0298361151122, 106.111557481547, 103.232548583324, 99.295240302438,
  106.742042339422, 99.1725600250135, 102.447664591815, 92.8677114462598,
  116.826136045524, 85.4245734291798, 115.565166329773, 113.174564860055,
  106.861058596942, 95.8302361339468, 77.1092540192799, 89.2908058308629,
  112.412742040396, 110.643967744195, 109.537397162078, 107.901837758501,
  104.969604023316, 107.682765778157, 109.466321282875, 133.532644951047,
  128.80901236506, 101.270018130495, 87.316028002806, 118.090769781179,
  120.255639745959, 90.3457720147583, 79.9046088419358, 136.577170171515,
  102.425526488872, 132.688799666889, 96.2972151466954, 118.209786038699,
  118.090769781179, 101.504599854825, 94.3346846938659, 84.67140829388,
  86.7688663993614, 129.245933601341, 127.781042344446, 106.861058596942,
  105.754488014826, 117.032139577379, 119.35173949693, 102.765569570576,
  108.205351567124, 126.414746875437, 101.90050298837, 99.1562556370934,
  91.8127452623644, 84.1414656104858, 84.3548260367683, 79.1197248504261,
  91.4168421288192, 90.4463141892395, 79.0007085929061, 109.263981770924,
  76.7521845525589, 97.5950080521511, 128.80901236506, 113.657407368227,
  114.206686337254, 112.513284214877, 75.3368157031739, 81.5226480244075,
  102.638886915154, 101.186694876064, 112.222842438706, 123.408064035336,
  86.5665268874105, 100.116947483653, 102.328648334295, 87.9465128606813,
  80.2580142887524, 80.2580142887524, 91.0767990471149, 85.5789725628141,
  113.165907175696, 112.05933659358, 101.270018130495, 81.8328866052668,
  81.0945026784831, 81.3290844028133, 120.609045192776, 102.447664591815,
  90.9762568726337, 103.708634307566, 102.328648334295, 121.870014908527,
  88.6603209729875, 90.0291897576466, 83.5331188555537, 91.9206443312728,
  87.1933477253815, 80.4640178206077, 84.3180028470634, 77.7397388771553,
  82.272149139803, 77.3826694104343, 89.715287156883, 79.9046088419358,
  82.1531328822829, 80.2616783086568, 77.3826694104343, 82.7836177401583,
  89.2074825764321, 89.088466318912, 88.9872920222396, 118.923594151005,
  83.8510238343148, 84.044587455909, 105.052927277747, 107.999348035269,
  97.3816476258686, 103.791957561996, 97.5950080521511, 88.6603209729875,
  107.246861729345, 89.8379674343074, 108.906912304203, 100.63953327262,
  118.840270896574, 95.5123311551858, 84.7719504683612, 100.747432341528,
  87.0034481236916, 77.8587551346754, 85.3018931517554, 92.7702011694918,
  74.8607299789328, 115.248584072661, 107.372527197298, 102.202982866342,
  97.9115903092628, 96.8566241253673, 111.465674925409, 109.466321282875,
  82.2500110368597, 92.7702011694918, 115.568830349677, 84.044587455909,
  116.199315207553, 99.295240302438, 87.1933477253815, 76.7117688565413,
  80.6985995449379, 82.4633714631422, 85.1798917037074, 93.0737149781151,
  88.5769977185567, 103.056011346747, 86.8767654682698, 97.9152543291671,
  120.92294779354, 91.7294220079335, 87.5072503261452, 90.4463141892395,
  84.1636037134291, 96.6542846134164, 95.703553478525, 88.7682200418959,
  92.1397163116164, 108.003012055173, 80.3893522505354, 99.8030448828889,
  90.0291897576466, 93.9387815603207, 103.708634307566, 77.1092540192799,
  117.55716307788, 79.8335329627323, 83.6164421099846, 88.731396852191,
  82.7836177401583, 86.8767654682698, 88.731396852191, 92.2408906082888,
  90.5517755466136, 84.1414656104858, 93.0737149781151, 85.8348677328627,
  98.8559777679018, 95.5123311551858, 101.067678618544, 85.9323780096308,
  82.7836177401583, 98.6647554445626, 82.1494688623785, 85.6157957525191,
  79.4376298291872, 75.3368157031739, 80.8921631665322, 115.771169861628,
  97.1591040033394, 80.8921631665322, 79.505527982829, 98.0121324837439,
  108.866496608185, 80.6985995449379, 100.747432341528, 91.0989371500582,
  82.9859572521092, 89.715287156883, 131.010705519546, 124.514634617453,
  91.6104057504134, 92.9903917236843, 91.8838211415678, 88.6909811561733,
  126.093509311965, 84.9853108946437, 82.9859572521092, 101.795041630996,
  100.243630139074, 110.013482886319, 88.5769977185567, 106.742042339422,
  95.2921406009932, 108.313250636032, 110.33138786508, 116.829800065428,
  86.6855431449306, 96.8566241253673, 90.2482617379903, 89.2908058308629,
  74.2302451210575, 91.8838211415678, 96.4607209918222, 97.4037857288119,
  131.533291308512, 78.2445582670783, 127.675716591188, 124.514634617453,
  80.4640178206077, 74.8607299789328, 84.8774118257353, 92.9903917236843,
  94.8597081943671, 96.8566241253673, 91.2533362836924],
  "mom_hs_new" => 1,
  "mom_iq_new" => 100)

sm = SampleModel("kid", kid);

rc = stan_sample(sm, data=kiddata)

if success(rc)
  samples = read_samples(sm)

  # Fetch the cmdstan summary.

  df = read_summary(sm)
  @test df[df.parameters .== Symbol("beta[1]"), :mean][1] ≈ 25.5 rtol=0.1
  @test df[df.parameters .== Symbol("beta[2]"), :mean][1] ≈ 6.01 rtol=0.1
  @test df[df.parameters .== Symbol("beta[3]"), :mean][1] ≈ 0.565 rtol=0.1
  @test df[df.parameters .== Symbol("sigma"), :mean][1] ≈ 18.2 rtol=0.1
end


type SigmoidLayer <: Layer
    inputActivation::Vol
    outputActivation::Vol

    function SigmoidLayer(dimensions)
        new(Vol(dimensions, 0.0), Vol(dimensions, 0.0))
    end
end

function forward(self::SigmoidLayer, inputActivation::Vol)
    @assert inputSize(self) == outputSize(self)
    self.inputActivation = inputActivation
    self.outputActivation.w = logistic(inputActivation.w)
    return self.outputActivation
end

function backward(self::SigmoidLayer)
    @assert inputSize(self) == outputSize(self)
    self.inputActivation.dw = self.outputActivation.w .* (1.0 .- self.outputActivation.w) .* self.outputActivation.dw
end

gradientStates(self::SigmoidLayer) = []
weights(self::SigmoidLayer) = []

type RectifiedLinearLayer
    inputActivation::Vol
    outputActivation::Vol

    function RectifiedLinearLayer(dimensions)
        new(Vol(dimensions, 0.0), Vol(dimensions, 0.0))
    end
end

function forward(self::RectifiedLinearLayer, inputActivation::Vol)
    @assert size(inputActivation) == size(self.inputActivation)
    self.inputActivation = inputActivation
    self.outputActivation = max(self.inputActivation, 0)
    return self.outputActivation
end

function backward(self::RectifiedLinearLayer)
    @assert inputSize(self) == outputSize(self)
    fill!(self.inputActivation.dw, 0.0)
    self.inputActivation.dw = (0 .< self.outputActivation.w) .* self.outputActivation.dw
end

gradientStates(self::RectifiedLinearLayer) = []


const order = 5
const a2 = Float64[1/5]
const a3 = Float64[3/40, 9/40]
const a4 = Float64[44/45, -56/15, 32/9]
const a5 = Float64[19372/6561, -25360/2187, 64448/6561, -212/729]
const a6 = Float64[9017/3168, -355/33, 46732/5247, 49/176, -5103/18656]
const a7 = Float64[35/384, 0., 500/1113, 125/192, -2187/6784, 11/84]
const a = {Float64[] a2 a3 a4 a5 a6 a7}
const bs = Float64[5179/57600, 0., 7571/16695, 393/640, -92097/339200, 187/2100, 1/40]
const c = Float64[0., 1/5, 3/10, 4/5, 8/9, 1., 1.]

function substep{T}(x::Vector{T}, x0::Vector{T}, h::Float64, coeffs::Vector{Float64}, k::Vector{Vector{T}})
    @inbounds for m=1:length(x0)
        dx::T = 0.
        @inbounds for i=1:length(coeffs)
            dx += coeffs[i]::Float64*k[i][m]
        end
        x[m] = x0[m] + h*dx
    end
    return nothing
end

function step{T}(F::Function, t::Float64, h::Float64,
                x0::Vector{T}, xp::Vector{T}, xs::Vector{T}, k::Vector{Vector{T}})
    for i=2:length(c)
        substep(xp, x0, h, a[i], k)
        F(t + h*c[i], xp, k[i])
    end
    substep(xs, x0, h, bs, k)
    return nothing
end

function allocate_memory{T}(x::Vector{T})
    xp = zeros(T, length(x))
    xs = zeros(T, length(x))
    k = Vector{Vector{T}}[]
    for i=1:7
        push!(k, zeros(T, length(x))
    return xp, xs, k
end

function error_estimate(xp, xs, abstol, reltol)
    err::Float64 = 0.
    for i=1:length(xp)
        sc_i = abstol + reltol*max(abs(xp[i]), abs(xs[i]))
        err += abs2(xp[i]-xs[i])/sc_i^2
    return sqrt(err/length(xp))
end

function initial_stepsize(F, x, k, tmp1, tmp2)
    d0 = norm(x,2)
    d1 = norm(k[1],2)
    h0 = ((d0<1e-5 | d1<1e-5) ? 1e-6 : 0.01*d0/d1)
    substep(tmp1, x, h0, [1.], k)
    F(t+h0, tmp1, tmp2)
    for i=1:length(tmp2)
        tmp2[i] -= k[1][i]
    end
    d2 = norm(tmp2, 2)
    if max(d1,d2)<1e-15
        h1 = max(1e-6, h0*1e-3)
    else
        h1 = (0.01/max(d1,d2))^(1./order)
    return min(100*h0, h1)
end

function ode{T}(F, tspan::Vector{Float64}, x0::Vector{T};
                    reltol = 1.0e-5,
                    abstol = 1.0e-8,
                    h0 = 0,
                    hmin = (tspan[end]-tspan[1])/1e9,
                    hmax = (tspan[end]-tspan[1]),
                    fout=(t,x)->nothing,)
    t, tend = tspan[1], tspan[end]

    fout(t, x)
    x = 1*x0
    xp, xs, k = allocate_memory(x0)

    F(t,x,k[1])
    h = (h0==0 ? initial_stepsize(F, x, k, k[2], k[3]) : h0)
    was_rejected = false
    while t < tfinal
        step(F, t, h, x, xp, xs, k)
        err = error_estimate(xp, xs, abstol, reltol)
        facmin = (was_rejected ? 1. : 5.)
        hnew = h*min(facmin, max(0.2, 0.9*(1./err)^(1./order)))
        hnew = min(hmax, hnew)
        if hnew<hmin
            error("Stepsize below hmin.")
        end
        if err>1
            was_rejected = true
            continue
        end
        was_rejected = false
        xp, x = x, xp
        fout(t, x)
    end
    return nothing
end


@testset "BBC & SID exp" begin
    α = big(1.1)

    @testset "BBC exp" begin
        x = 1.1
        p0 = 1 # Order of approx for 0 matrix multiplies
        (graph, cref) = graph_bbc_exp(0)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p0 + 1

        x = 1.1
        p1 = 2 # Order of approx for 1 matrix multiplies
        (graph, cref) = graph_bbc_exp(1)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p1 + 1

        x = 1.1
        p2 = 4 # Order of approx for 2 matrix multiplies
        (graph, cref) = graph_bbc_exp(2, T = BigFloat)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p2 + 1

        x = 1.1
        p3 = 8 # Order of approx for 3 matrix multiplies
        (graph, cref) = graph_bbc_exp(3)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p3 + 1

        x = 2.0
        p4bbc = 12 # Order of approx for 4 matrix multiplies
        (graph, cref) = graph_bbc_exp(4)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p4bbc + 1

        x = 2.3
        p5bbc = 18 # Order of approx for 5 matrix multiplies
        (graph, cref) = graph_bbc_exp(5)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p5bbc + 1
    end

    @testset "SID exp" begin
        x = 2.0
        p0sid = 1 # Order of approx for 0 matrix multiplies
        (graph, cref) = graph_sid_exp(0)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p0sid + 1

        x = 2.0
        p1sid = 2 # Order of approx for 1 matrix multiplies
        (graph, cref) = graph_sid_exp(1)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p1sid + 1

        x = 2.0
        p2sid = 4 # Order of approx for 2 matrix multiplies
        (graph, cref) = graph_sid_exp(2)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p2sid + 1

        x = 2.0
        p3sid = 6 # Order of approx for 3 matrix multiplies
        (graph, cref) = graph_sid_exp(3)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p3sid + 1

        x = 2.0
        p4sid = 15 # Order of approx for 4 matrix multiplies
        (graph, cref) = graph_sid_exp(4)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p4sid + 1

        x = 3.2
        p5sid = 21 # Order of approx for 5 matrix multiplies
        (graph, cref) = graph_sid_exp(5)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p5sid + 1

        x = 3.5
        p6sid = 24 # Order of approx for 6 matrix multiplies
        (graph, cref) = graph_sid_exp(6)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p6sid + 1

        x = 5.2
        p7sid = 30 # Order of approx for 7 matrix multiplies
        (graph, cref) = graph_sid_exp(7)
        err1 = eval_graph(big(graph), big(x)) - exp(big(x))
        err2 = eval_graph(big(graph), big(x) / α) - exp(big(x) / α)
        @test -log(abs(err2 / err1)) / log(α) > p7sid + 1

        # Test scaling and squaring
        x = 6.5
        (graph, cref) = graph_sid_exp(8)
        @test abs(eval_graph(big(graph), big(x)) - exp(big(x))) < 5e2 * eps()
    end
end


mutable struct Nnet
    y::Array{Int64,1}
    X::Array{Float64,2}
    N::Int64
    D::Int64
    hidden_layers::Int64
    nodes::Int64
    weights::Int64
    λ::Array{Float64,2}
    μ::Array{Float64,1}
    σ::Array{Float64,1}
    prior::Distributions.Normal{Float64}
    iter::Int64
end

"Construct Neural Network"
function Nnet(y::Array{Int64, 1},
              X::Array{Float64, 2};
              hidden_layers::Int64=1,
              nodes::Int64=2,
              prior::Distributions.Normal{Float64}=Normal())

    N = length(y)
    D = size(X)[2]

    # still need to account for more than one hidden layer
    weights = D * nodes + nodes + # input to layer 1
              # layer n to layer n + 1
              (hidden_layers - 1) * ((nodes + 1) * nodes)  +
              nodes + 1 # final layer to f(X)

    # initialize variational parameters
    λ = rand(Normal(), weights, 2)
    truncate!(λ, weights)

    μ = λ[:, 1]
    σ = softplus.(λ[:, 2])

    Nnet(y, X, N, D,
         hidden_layers, nodes, weights,
         λ, μ, σ, prior, 0)
end

# estimate of evidence lower bound
function ℒ(nn::Nnet, μ, σ)
    # extract elements from Nnet
    y = nn.y
    X = nn.X
    prior = nn.prior
    hidden_layer = nn.hidden_layers
    nodes = nn.nodes

    # draw values from q
    ϵ = rand(Normal(), length(μ))
    z = μ + ϵ .* σ

    # initialize vectors for calculating ELBO
    # note that the type is important for AD
    log_lik = zeros(eltype(μ), 1)
    log_prior = zeros(eltype(μ), length(μ))
    entropy = zeros(eltype(μ), length(μ))
    f = zeros(eltype(μ), length(μ))

    # get number of covariates
    N, D = size(X)

    # store layer calculations
    layer = zeros(eltype(μ), nodes)
    layer_p1 = zeros(eltype(μ), nodes)
    out = zeros(eltype(μ), N)

    for n in 1:N
        out[n] = fwd_prop(X[n, :],
                          nodes, hidden_layer,
                          z,
                          layer, layer_p1)

        log_lik += logpdf(Bernoulli(out[n]), y[n])
    end

    # construct ELBO estimate
    for i in 1:length(μ)
        # first calculate log prior and entropy
        log_prior[i] = logpdf(nn.prior, z[i])
        entropy[i] = logpdf(Normal(μ[i], σ[i]), z[i])
        f[i] = log_lik[1] + log_prior[i] - entropy[i]
    end

    return f
end

function fit(nn::Nnet,
             opt::GradDescent.Optimizer;
             tol=0.005,
             maxiter=1000,
             msg=100,
             elboiter=10)
    N = length(nn.y)
    D = nn.weights

    mean_δ = 1.0
    i = 1

    while mean_δ > tol && i < maxiter
        nn.μ = nn.λ[:, 1]
        nn.σ = softplus.(nn.λ[:, 2])

        g = ForwardDiff.jacobian(ϕ -> ℒ(nn, ϕ[1:D], ϕ[(D+1):2D]),
                                 vcat(nn.μ, nn.σ))
        ∇ℒ = hcat(diag(g[1:D, 1:D]),
                  diag(g[1:D, (D+1):2D]))

        δ = update(opt, ∇ℒ)
        nn.λ += δ
        truncate!(nn.λ, D)

        mean_δ = mean(δ .^ 2)

        if (i % msg == 0)
            println(i, " ", mean_δ)
        end

        i += 1
    end

    # calculate elbo
    nn.μ = nn.λ[:, 1]
    nn.σ = softplus.(nn.λ[:, 2])
    elbo = 0

    for i in 1:elboiter
        elbo += mean(ℒ(nn, nn.μ, nn.σ))
    end

    elbo = elbo / elboiter

    return elbo
end

function predict(nn::Nnet,
                 X::Array{Float64, 2})
    # extract elements from Nnet
    hidden_layer = nn.hidden_layers
    nodes = nn.nodes

    # get number of covariates
    N, D = size(X)

    # store layer calculations
    layer = zeros(nodes)
    layer_p1 = zeros(nodes)
    out = zeros(N)

    # for right now, fix z at mean
    z = nn.μ

    for n in 1:N
        out[n] = fwd_prop(X[n, :],
                          nodes, hidden_layer,
                          z,
                          layer, layer_p1)
    end

    return out
end


# # Event counting with Negative Binomial

# We load all the necessary packages
using AbstractGPs
using ApproximateGPs
using AugmentedGPLikelihoods
using Distributions
using LinearAlgebra
# Plotting libraries
using Plots

# We create some random data (sorted for plotting reasons)
N = 100
x = range(-10, 10; length=N)
kernel = with_lengthscale(SqExponentialKernel(), 2.0)
gp = GP(kernel)
lik = NegativeBinomialLikelihood(NBParamFailure(15))
lf = LatentGP(gp, lik, 1e-6)
f, y = rand(lf(x));
# We plot the sampled data
plt = scatter(x, y; label="Data")
plot!(plt, x, f; color=:red, label="Latent GP")
# ## CAVI Updates
# We write our CAVI algorithmm
function u_posterior(fz, m, S)
    return posterior(SparseVariationalApproximation(Centered(), fz, MvNormal(m, S)))
end

function cavi!(fz::AbstractGPs.FiniteGP, x, y, m, S, qΩ; niter=10)
    K = ApproximateGPs._chol_cov(fz)
    for _ in 1:niter
        post_u = u_posterior(fz, m, S)
        post_fs = marginals(post_u(x))
        aux_posterior!(qΩ, lik, y, post_fs)
        S .= inv(Symmetric(inv(K) + Diagonal(only(expected_auglik_precision(lik, qΩ, y)))))
        m .= S * (only(expected_auglik_potential(lik, qΩ, y)) + K \ mean(fz))
    end
    return m, S
end;
# Now we just initialize the variational parameters
m = zeros(N)
S = Matrix{Float64}(I(N))
qΩ = init_aux_posterior(lik, N)
fz = gp(x, 1e-8);
# And visualize the current posterior
x_te = -10:0.01:10
plot!(
    plt, x_te, u_posterior(fz, m, S); color=:blue, alpha=0.3, label="Initial VI Posterior"
)
# We run CAVI for 3-4 iterations
cavi!(fz, x, y, m, S, qΩ; niter=4);
# And visualize the obtained variational posterior
plot!(
    plt,
    x_te,
    u_posterior(fz, m, S);
    color=:darkgreen,
    alpha=0.3,
    label="Final VI Posterior",
)
# ## ELBO
# How can one compute the Augmented ELBO?
# Again AugmentedGPLikelihoods provides helper functions
# to not have to compute everything yourself
function aug_elbo(lik, u_post, x, y)
    qf = marginals(u_post(x))
    qΩ = aux_posterior(lik, y, qf)
    return expected_logtilt(lik, qΩ, y, qf) - aux_kldivergence(lik, qΩ, y) -
           kldivergence(u_post.approx.q, u_post.approx.fz)     # approx.fz is the prior and approx.q is the posterior
end

aug_elbo(lik, u_posterior(fz, m, S), x, y)
# ## Gibbs Sampling
# We create our Gibbs sampling algorithm (we could do something fancier with
# AbstractMCMC)
function gibbs_sample(fz, f, Ω; nsamples=200)
    K = ApproximateGPs._chol_cov(fz)
    Σ = zeros(length(f), length(f))
    μ = zeros(length(f))
    return map(1:nsamples) do _
        aux_sample!(Ω, lik, y, f)
        Σ .= inv(Symmetric(inv(K) + Diagonal(only(auglik_precision(lik, Ω, y)))))
        μ .= Σ * (only(auglik_potential(lik, Ω, y)) + K \ mean(fz))
        rand!(MvNormal(μ, Σ), f)
        return copy(f)
    end
end;
# We initialize our random variables
f = randn(N)
Ω = init_aux_variables(lik, N);
# Run the sampling for default number of iterations (200)
fs = gibbs_sample(fz, f, Ω);
# And visualize the samples overlapped to the variational posterior
# that we found earlier.
for f in fs
    plot!(plt, x, f; color=:black, alpha=0.07, label="")
end
plt


"""
    OpticalFlowAlgorithm

An optical flow algorithm with given parameters.
"""
abstract type OpticalFlowAlgorithm end

"""
    Farneback(Args...)

A method for dense optical flow estimation developed by Gunnar Farneback. It
computes the optical flow for all the points in the frame using the polynomial
representation of the images. The idea of polynomial expansion is to approximate
the neighbourhood of a point in a 2D function with a polynomial. Displacement
fields are estimated from the polynomial coefficients depending on how the
polynomial transforms under translation.

# Options
Various options for the fields of this type are described in more detail
below.

## Choices for `iterations`

Number of iterations the displacement estimation algorithm is run at each point.
If left unspecified a default value of seven iterations is assumed.

## Choices for `estimation_window`

Determines the neighbourhood size over which information will be intergrated
when determining the displacement of a pixel. The total size equals
`2*estimation_window + 1`. If left unspecified a default value of eleven is
assumed.

## Choices for `σ_estimation_window`

Standard deviation of a Gaussian weighting filter used to weigh the contribution
of a pixel's neighbourhood when determining the displacement of a pixel.
If left unspecified a default value of nine is assumed.

## Choices for `expansion_window`

Determines the size of the pixel neighbourhood used to find polynomial expansion
for each pixel; larger values mean that the image will be approximated with
smoother surfaces, yielding more robust algorithm and more blurred motion field.
The total size equals `2*expansion_window + 1`. If left unspecified a default
value of six is assumed.

## Choices for `σ_expansion_window`

Standard deviation of the Gaussian that is used to smooth the image for the purpose
of approximating it with a polynomial expansion. If left unspecified a default
value of five is assumed.

# References

1. Farnebäck G. (2003) Two-Frame Motion Estimation Based on Polynomial Expansion. In: Bigun J., Gustavsson T. (eds) Image Analysis. SCIA 2003. Lecture Notes in Computer Science, vol 2749. Springer, Berlin, Heidelberg
2. Farnebäck, G.: Polynomial Expansion for Orientation and Motion Estimation. PhD thesis, Linköping University, Sweden, SE-581 83 Linköping, Sweden (2002) Dissertation No 790, ISBN 91-7373-475-6.

"""
struct Farneback{F <: Float64, I <: Int} <: OpticalFlowAlgorithm
    iterations::I
    estimation_window::I
    σ_estimation_window::F
    expansion_window::I
    σ_expansion_window::F
end

Farneback(iterations::Int; estimation_window::Int = 11, σ_estimation_window::Real = 9.0, expansion_window::Int =  6, σ_expansion_window::Real = 5.0) = Farneback(iterations, estimation_window, σ_estimation_window, expansion_window, σ_expansion_window)

"""
    LucasKanade(Args...)

A differential method for optical flow estimation developed by Bruce D. Lucas
and Takeo Kanade. It assumes that the flow is essentially constant in a local
neighbourhood of the pixel under consideration, and solves the basic optical flow
equations for all the pixels in that neighbourhood by the least squares criterion.

# Options
Various options for the fields of this type are described in more detail
below.

## Choices for `iterations`

The termination criteria of the iterative search algorithm, that is, the number of
iterations.

## Choices for `window_size`

Size of the search window at each pyramid level; the total size of the
window used is 2*window_size + 1. If left unspecified, a default value of
eleven is assumed.

## Choices for `pyramid_levels`

0-based maximal pyramid level number; if set to 0, pyramids are not used
(single level), if set to 1, two levels are used, and so on. IF left unspecified,
a default value of four is assumed.

## Choices for `eigenvalue_threshold`

The algorithm calculates the minimum eigenvalue of a (2 x 2) normal matrix of
optical flow equations, divided by number of pixels in a window; if this value
is less than `eigenvalue_threshold`, then a corresponding feature is filtered
out and its flow is not processed (default value is 1e-4).

## `ϵ` termination criteria

Minimum required change in displacement at which the iterative algorithm continues to work.
Default value is `1e-2`.

## References

1. B. D. Lucas, & Kanade. "An Interative Image Registration Technique with an Application to Stereo Vision," DARPA Image Understanding Workshop, pp 121-130, 1981.
2. J.-Y. Bouguet, “Pyramidal implementation of the afﬁne lucas-kanade feature tracker description of the algorithm,” Intel Corporation, vol. 5,no. 1-10, p. 4, 2001.
"""
struct LucasKanade <: OpticalFlowAlgorithm
    iterations::Int
    window_size::Int
    pyramid_levels::Int
    eigenvalue_threshold::Float64
    ϵ::Float64
end

LucasKanade(
    iterations::Integer = 20; window_size::Integer = 11, pyramid_levels::Integer = 4,
    eigenvalue_threshold::Real = 1e-4, ϵ::Real = 1e-2,
) = LucasKanade(iterations, window_size,  pyramid_levels, eigenvalue_threshold, ϵ)

"""
    flow = optical_flow(source, target, Farneback(Args...))
    flow = optical_flow(source, target, displacement, Farneback(Args...))

Returns the dense optical flow from the `source` to the `target` image using the `Farneback` algorithm.

# Details

The `source` and `target` images must be `Gray` types.

The `displacement` argument allows you to specify an initial
guess for the optical flow and must be of type `Array{SVector{2, Float64}, 2}`.
The elements of `displacement` should represent the flow required to map the
`(row, column)` of each pixel in the `source` image to the `target` image.

"""
function optical_flow(source::AbstractArray{T, 2}, target::AbstractArray{T,2}, algorithm::Farneback) where T <: Gray
    # Sanity checks.
    @assert size.(axes(source)) == size.(axes(target)) "Images must have the same size"
    optflow(source, target, algorithm)
end

function optical_flow(source::AbstractArray{T, 2}, target::AbstractArray{T,2}, displacement::Array{SVector{2, Float64}, 2}, algorithm::Farneback) where T <: Gray
    # Sanity checks.
    @assert size.(axes(source)) == size.(axes(target)) "Images must have the same size"
    @assert size.(axes(source)) == size.(axes(displacement)) "Optical flow field must match image size"
    optflow!(source, target, copy(displacement), algorithm)
end

"""
    flow, indicator  = optical_flow(source, target, points, LucasKanade(Args...))
    flow, indicator  = optical_flow(source, target, points, displacement, LucasKanade(Args...))

Returns the optical flow from the `source` to the `target` image for the specified `points` using the `LucasKanade` algorithm.

# Details

The `source` and `target` images must be `Gray` types.

The `points` argument is of type `Array{SVector{2, Float64}, 2}` and represents
a set of keypoints in the `source` image for which the optical flow is to be
computed. The coordinates  of a `point` represent the  `(row, column)` of a
pixel in the image. The `displacement` argument allows you to specify an initial
guess for the optical flow.

The function returns `flow` of type `Array{SVector{2, Float64}, 2}` which
matches the length of `points` and represents the displacement needed to map
a point in the `source` image to the `target` image.

The `indicator` is a vector of boolean values (one for each point in `points`)
which signals whether a particular point was successfully tracked or not.

In order to use the `flow` to index the corresponding point in the `target` image
you first need to round it to the nearest integer, and check that it falls within
the bounds of the `target` image dimensions.

"""
function optical_flow(source::AbstractArray{T, 2}, target::AbstractArray{T,2}, points::Array{SVector{2, Float64}, 1}, algorithm::LucasKanade) where T <: Gray
    # Sanity checks.
    @assert size.(axes(source)) == size.(axes(target)) "Images must have the same size"

    optflow(source, target, points, algorithm)
end

function optical_flow(source::AbstractArray{T, 2}, target::AbstractArray{T,2}, points::Array{SVector{2, Float64}, 1},  displacement::Array{SVector{2, Float64}, 1}, algorithm::LucasKanade) where T <: Gray
    # sanity checks
    @assert size.(axes(source)) == size.(axes(target)) "Images must have the same size"
    @assert size.(axes(points)) == size.(axes(displacement)) "Vector of points and vector of displacement must have the same size"
    optflow!(source, target, points, copy(displacement), algorithm)
end

#------------------
# IMPLEMENTATIONS
#------------------

include("lucas_kanade.jl")
include("farneback.jl")


###########
# Generic methods for obtaining Jacobi matrices with one of the a, b and c parameters lowered by 1.
# passing symbols :a, :b or :c into lowindex determines which parameter is lowered
######
function initialα_gen(P::SemiclassicalJacobi, lowindex::Symbol)
    t = P.t; a = P.a; b = P.b; c = P.c;
    A,B,_ = recurrencecoefficients(P)
    # this makes use of Q_1 == A[1]x + B[1]
    if lowindex == :a
        return -(A[1]*sum(SemiclassicalJacobiWeight(t, a, b, c)) + B[1]*sum(SemiclassicalJacobiWeight(t, a-1, b, c)))/(sum(SemiclassicalJacobiWeight(t, a-1, b, c)))
    elseif lowindex == :b
        return -(A[1]*sum(SemiclassicalJacobiWeight(t, a+1, b-1, c)) + B[1]*sum(SemiclassicalJacobiWeight(t, a, b-1, c)))/(sum(SemiclassicalJacobiWeight(t, a, b-1, c)))
    else # lowindex == :c
        return -(A[1]*sum(SemiclassicalJacobiWeight(t, a+1, b, c-1)) + B[1]*sum(SemiclassicalJacobiWeight(t, a, b, c-1)))/(sum(SemiclassicalJacobiWeight(t, a, b, c-1)))
    end
end

mutable struct LoweringCoefficients{T} <: AbstractCachedVector{T}
    P::SemiclassicalJacobi{T}
    lowind::Symbol # options for lowering are :a, :b and :c
    data::Vector{T}
    datasize::Tuple{Int}
    LoweringCoefficients{T}(P::SemiclassicalJacobi{T}, lowindex::Symbol) where T = new{T}(P, lowindex, [initialα_gen(P, lowindex)], (1,))
end
LoweringCoefficients(P::SemiclassicalJacobi{T}, lowindex::Symbol) where T = LoweringCoefficients{T}(P, lowindex)
size(::LoweringCoefficients) = (ℵ₀,)

function getindex(α::LoweringCoefficients, I::UnitRange{Int})
    resizedata!(α, maximum(I))
    return getindex(α.data,I)
end

function getindex(α::LoweringCoefficients, I::Int)
    resizedata!(α, I)
    return getindex(α.data,I)
end

function resizedata!(α::LoweringCoefficients, nm::Int)
    νμ = length(α.data)
    if nm > νμ
        olddata = copy(α.data)
        nmax = maximum(nm)
        α.data = similar(olddata,nmax)
        α.data[1:νμ] = olddata[1:νμ]
        inds = νμ:nmax
        cache_filldata!(α, inds)
        α.datasize = (nmax,)
    end
    return α
end
cache_filldata!(α::LoweringCoefficients, inds::UnitRange{Int}) =  αgenfillerbackwards!(α.data, 100, α.P, α.lowind, inds)

# fill in coefficients via Miller recurrence
function αgenfillerbackwards!(α::Vector{T}, addscale::Int, P::SemiclassicalJacobi{T}, lowindex::Symbol, inds::UnitRange{Int}) where T
    maxI::Int = maximum(inds)
    minI::Int = minimum(inds)
    oldval = α[minI]
    n = maxI
    k = one(T)
    k2 = zero(T)
    if lowindex == :a
        _,B,C = recurrencecoefficients(P)
        while abs(k2-k) > 1e-10
            k2, k = k, one(T)
            n += addscale
            @inbounds for j = n:-1:maxI
                k = -C[j+1]/(k+B[j+1])
            end
        end
        α[maxI] = k
        @inbounds for j = maxI:-1:minI+1
            α[j-1] = -C[j]/(α[j]+B[j])
        end
    elseif lowindex == :b
        A,B,C = recurrencecoefficients(P)
        while abs(k2-k) > 1e-10
            k2, k = k, one(T)
            n += addscale
            @inbounds for j = n:-1:maxI
                k = -C[j+1]/(k+A[j+1]+B[j+1])
            end
        end
        α[maxI] = k
        @inbounds for j = maxI:-1:minI+1
            α[j-1] = -C[j]/(α[j]+A[j]+B[j])
        end
    else # lowindex == :c
        A,B,C = recurrencecoefficients(P)
        while abs(k2-k) > 1e-10
            k2, k = k, one(T)
            n += addscale
            @inbounds for j = n:-1:maxI
                k = -C[j+1]/(k+A[j+1]*P.t+B[j+1])
            end
        end
        α[maxI] = k
        @inbounds for j = maxI:-1:minI+1
            α[j-1] = -C[j]/(α[j]+A[j]*P.t+B[j])
        end
    end
    α[minI:maxI] = ((oldval)/α[inds[1]]).*α[minI:maxI]
    α
end

function symlowered_jacobimatrix(Q::SemiclassicalJacobi, lowindex::Symbol)
    bands = SemiclassicalLoweredJacobiBands(Q,lowindex)
    return SymTridiagonal(bands[1,:],bands[2,:])
end

mutable struct SemiclassicalLoweredJacobiBands{T} <: AbstractCachedMatrix{T}
    P::SemiclassicalJacobi{T}
    data::Array{T} # to avoid redundant re-computations, the bands are stored as a (∞,2)-array
    αcfs::LoweringCoefficients{T}
    datasize::Tuple{Int,Int}
end
size(r::SemiclassicalLoweredJacobiBands) = (2,ℵ₀)

function SemiclassicalLoweredJacobiBands(P::SemiclassicalJacobi{T}, lowindex::Symbol) where T
    cachedα = LoweringCoefficients(P,lowindex)
    a, b, c, t, αcfs = P.a, P.b, P.c, P.t, cachedα[1]
    C, A = P.X[2,1], P.X[1,1]
    if lowindex == :a
        SemiclassicalLoweredJacobiBands{T}(P, [A-C*αcfs, C/(sqrt(sum(SemiclassicalJacobiWeight(t,a-1,b,c)))/sqrt((((αcfs^2-2*αcfs*A/C+A^2/C^2)*sum(SemiclassicalJacobiWeight(t, a-1, b, c))+(2*αcfs/C-2*A/C^2)*sum(SemiclassicalJacobiWeight(t, a, b, c))))+(1/C^2)*(sum(SemiclassicalJacobiWeight(t, a+1, b, c)))))], cachedα, (2,1))
    elseif lowindex == :b
        SemiclassicalLoweredJacobiBands{T}(P, [A-C*αcfs, C/(sqrt(sum(SemiclassicalJacobiWeight(t,a,b-1,c)))/sqrt((((αcfs^2-2*αcfs*A/C+A^2/C^2)*sum(SemiclassicalJacobiWeight(t, a, b-1, c))+(2*αcfs/C-2*A/C^2)*sum(SemiclassicalJacobiWeight(t, a+1, b-1, c))))+(1/C^2)*(sum(SemiclassicalJacobiWeight(t, a+2, b-1, c)))))], cachedα, (2,1))
    elseif lowindex == :c
        SemiclassicalLoweredJacobiBands{T}(P, [A-C*αcfs, C/(sqrt(sum(SemiclassicalJacobiWeight(t,a,b,c-1)))/sqrt((((αcfs^2-2*αcfs*A/C+A^2/C^2)*sum(SemiclassicalJacobiWeight(t, a, b, c-1))+(2*αcfs/C-2*A/C^2)*sum(SemiclassicalJacobiWeight(t, a+1, b, c-1))))+(1/C^2)*(sum(SemiclassicalJacobiWeight(t, a+2, b, c-1)))))], cachedα, (2,1))
    end
end

function getindex(r::SemiclassicalLoweredJacobiBands, I::Vararg{Int,2})
    r.αcfs[maximum(I)+1] # expanding the cache of coefficients here prevents redundant recomputation later
    resizedata!(r,maximum(I))
    getindex(r.data,I[1],I[2])
end
function getindex(r::SemiclassicalLoweredJacobiBands, I::Int, J::UnitRange{Int})
    r.αcfs[maximum(J)+1] # expanding the cache of coefficients here prevents redundant recomputation later
    resizedata!(r,maximum(J))
    getindex(r.data,I,J)
end
function resizedata!(r::SemiclassicalLoweredJacobiBands, nm::Int)
    νμ = length(r.data[1,:])
    if nm > νμ
        olddata = copy(r.data)
        r.data = similar(olddata,2,nm)
        r.data[1:2,1:νμ] = olddata
        inds = νμ:nm
        cache_filldata!(r, inds)
        r.datasize = (2,nm)
    end
    return r
end
cache_filldata!(r::SemiclassicalLoweredJacobiBands, inds::UnitRange{Int}) =  loweringjacobibandfill!(r.data, r.αcfs, r.P, inds)

function loweringjacobibandfill!(data::Array, αcfs::LoweringCoefficients, P::SemiclassicalJacobi, inds::UnitRange{Int})
    lowI = minimum(inds):maximum(inds)-1
    shiftI = minimum(inds)+1:maximum(inds)
    C,A = subdiagonaldata(P.X), diagonaldata(P.X)
    data[1,shiftI] = -C[shiftI].*αcfs[shiftI]+C[lowI].*αcfs[lowI]+A[shiftI]
    data[2,shiftI] = sqrt.(C[lowI].*C[shiftI].*αcfs[shiftI]./αcfs[lowI])
    data
end

###########
# Methods for the special case of computing SemiclassicalJacobi(t,0,0,-1)
######
# As this can be built from SemiclassicalJacobi(t,0,0,0) which is just shifted and normalized Legendre(), we have more explicit methods at our disposal due to explicitly known coefficients for the Legendre bases.

###
#   α coefficients implementation, cached, direct and via recurrences
###
# inital value n=0 for α_{n,n-1}(t) coefficients
initialα(t) = t-2/(log1p(t)-log(t-1))

# compute n-th coefficient from direct evaluation formula
αdirect(n, t) = 2*n*_₂F₁general2((1+n)/2,(n+2)/2,(2*n+3)/2,1/t^2)/(t*2*(1+2*n)*_₂F₁general2(n/2,(n+1)/2,(2*n+1)/2,1/t^2))
# this version takes a pre-existing vector v and fills in the missing data guided by indices in inds using explicit formula
function αdirect!(α, t, inds)
    @inbounds for n in inds
        α[n] = 2*n*_₂F₁general2((1+n)/2,(2+n)/2,(2*n+3)/2,1/t^2)/(t*2*(1+2*n)*_₂F₁general2(n/2,(n+1)/2,(2*n+1)/2,1/t^2))
    end
end

# takes a previously computed vector of α_{n,n-1}(t) that has been increased in size and fills in the missing data guided by indices in inds via forward recurrence
function αcoefficients!(α, t, inds)
    @inbounds for n in inds
       α[n] = (t*(2*n-1)/n-(n-1)/(n*α[n-1]))
    end
end

# takes a previously computed vector of α_{n,n-1}(t) that has been increased in size and fills in the missing data guided by indices in inds, using a final condition computation and backward recurrence
function backαcoeff!(α, t, inds)
    α[end] = αdirect(length(α),t)
    @inbounds for n in reverse(inds)[1:end-1]
       α[n-1] = (n-1)/(t*(2*n-1)-n*α[n])
    end
end

# cached implementation using stable back recurrence to fill data
mutable struct neg1c_αcfs{T} <: AbstractCachedVector{T}
    t::T
    data::Vector{T}
    datasize::Tuple{Int}
    neg1c_αcfs{T}(t::T) where T = new{T}(t, [initialα(2*t-1),αdirect(2,2*t-1)], (2,))
end
neg1c_αcfs(t::T) where T = neg1c_αcfs{T}(t)
size(α::neg1c_αcfs) = (ℵ₀,)
cache_filldata!(α::neg1c_αcfs, inds) = backαcoeff!(α.data, 2*α.t-1, inds)

function resizedata!(α::neg1c_αcfs, nm)
    νμ = length(α.data)
    if nm > νμ
        olddata = copy(α.data)
        α.data = similar(olddata,maximum(nm))
        α.data[1:νμ] = olddata[1:νμ]
        inds = Array(νμ:maximum(nm))
        cache_filldata!(α, inds)
        α.datasize = (nm,)
    end
    α
end

###
#   cached implementation of normalization constants
###
mutable struct neg1c_normconstant{T} <: AbstractCachedVector{T}
    t::T
    data::Vector{T}
    datasize::Tuple{Int}
    neg1c_normconstant{T}(t::T) where T = new{T}(t, neg1c_normconstinitial(t, 10), (10,))
end
neg1c_normconstant(t::T) where T = neg1c_normconstant{T}(t)
size(B::neg1c_normconstant) = (ℵ₀,)
cache_filldata!(B::neg1c_normconstant, inds) = neg1c_normconstextension!(B.data, inds, B.t)

function resizedata!(B::neg1c_normconstant, nm)
    νμ = length(B.data)
    if nm > νμ
        olddata = copy(B.data)
        B.data = similar(olddata,maximum(nm))
        B.data[1:νμ] = olddata[1:νμ]
        inds = Array(νμ-1:maximum(nm))
        cache_filldata!(B, inds)
        B.datasize = (nm,)
    end
    B
end

function neg1c_normconstinitial(t::T, N::Integer) where T
    # generate α coefficients for OPs via recurrence
    α = zeros(T,N-1)
    α[1] = initialα(2*t-1)
    backαcoeff!(α,2*t-1,(2:N-1))
    # normalization constants
    B = [one(T)]
    append!(B,sqrt(2*acoth(2*t-1))*sqrt.((1:N-1)./(2 .*α)))
    return B
end

function neg1c_normconstextension!(B::Vector, inds, t::T) where T
    n = maximum(inds)
    m = minimum(inds)-1
    # generate missing α coefficients
    α = zeros(T, n)
    backαcoeff!(α,2*t-1,(m:n))
    # normalization constants
    norm = sqrt(2*acoth(2*t-1))*sqrt.((m:n)./(2 .*α[m:n]))
    B[m+1:n] = norm[1:end-1]
    B
end

###
#   bidiagonal operator which converts from OPs wrt (t-x)^-1 to shifted Legendre
###
function neg1c_tolegendre(t::T) where T
    norm = neg1c_normconstant(t)
    α = neg1c_αcfs(t)
    supD = norm.*Vcat(zeros(T,1),-α)
    D = norm
    return Bidiagonal(D,supD[2:end],:U)
end

###
#   explicit polynomial evaluations
###
# Evaluate the n-th normalized OP wrt 1/(t-x) at point x
function evalQn(n::Integer, x, t::T) where T
    # this version recomputes α based on t
    t <= 1 && throw(ArgumentError("t must be greater than 1."))
    n == 0 && return one(T)
    α = αdirect(n,2*t-1)
    return sqrt(2*acoth(2*t-1))*(-α*jacobip(n-1,0,0,2*x-1)+jacobip(n,0,0,2*x-1))*sqrt(n/(2*α))
end
# Evaluate the n-th non-normalized OP wrt 1/(t-x) at point x
function evalϕn(n::Integer, x, t::T) where T
    # this version recomputes α based on t
    t <= 1 && throw(ArgumentError("t must be greater than 1."))
    n == 0 && return one(T)
    return -αdirect(n,2*t-1)*jacobip(n-1,0,0,2*x-1)+jacobip(n,0,0,2*x-1)
end


int main () {
  printInt(fact(7)) ;
  printInt(factr(7)) ;
  return 0 ;
} ;

// iteracyjnie
int fact (int n) {
  int i,r ;
  i = 1 ;
  r = 1 ;
  while (i < n+1) {
    r = r * i ;
    i++ ;
  } ;
  return r ;
} ;

// rekurencyjnie
int factr (int n) {
  if (n < 2)
    return 1 ;
  else
    return (n * factr(n-1)) ;
} ;


"""
    maximum_error(mw::MWState)

Compute maximum error of synthetic data on the query set. Result is not
differentially private.
"""
function maximum_error(mw::MWState)
    norm(evaluate(mw.queries, mw.synthetic) - mw.real_answers, Inf)
end

"""
    mean_squared_error(mw::MWState)

Compute mean_squared_error of synthetic data on the query set. Result is not
differentially private.
"""
function mean_squared_error(mw::MWState)

    errors = evaluate(mw.queries, mw.synthetic) - mw.real_answers
    (norm(errors)^2)/length(errors)
end

function total_squared_error(mw::MWState)

    errors = evaluate(mw.queries, mw.synthetic) - mw.real_answers
    (norm(errors)^2)/length(errors)
end

"""
   kl_divergence(a, b)

Compute KL-divergence between two normalized histograms a and b.
"""
function kl_divergence(a::Array{Float64, 1}, b::Array{Float64, 1})
    r = 0.0
    for i = 1 : length(a)
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        if ai > 0 && bi > 0
            r += ai * log(2, ai / bi)
        end
    end
    r
end

function kl_divergence_error(mw::MWState)

    #error = kl_divergence(mw.synthetic.weights, mw.real.weights)
    error = kl_divergence(mw.real.weights, mw.synthetic.weights)
    error
end


struct FaceQuadratures
    quads::Any
    facetoquad::Any
    ncells::Any
    nfaces::Any
    function FaceQuadratures(quads, facetoquad)
        nphase, nfaces, ncells = size(facetoquad)
        @assert nfaces == 4
        @assert nphase == 2
        @assert all(facetoquad .>= 0)
        @assert all(facetoquad .<= length(quads))
        new(quads, facetoquad, ncells, nfaces)
    end
end

function FaceQuadratures(cutmesh, levelset, numqp)

    ncells = number_of_cells(cutmesh)

    quads = face_quadratures(numqp)
    quad1d = ReferenceQuadratureRule(numqp)
    facetoquad = zeros(Int, 2, 4, ncells)

    for cellid = 1:ncells
        s = cell_sign(cutmesh,cellid)
        if s == +1
            facetoquad[1, :, cellid] .= 1:4
        elseif s == -1
            facetoquad[2, :, cellid] .= 1:4
        elseif s == 0
            nodeids = nodal_connectivity(background_mesh(cutmesh),cellid)
            load_coefficients!(levelset,cellid)

            pquad = face_quadratures(interpolater(levelset), +1, quad1d)
            idxstart = length(quads) + 1
            append!(quads, pquad)
            facetoquad[1, :, cellid] .= idxstart:(idxstart+3)

            nquad = face_quadratures(interpolater(levelset), -1, quad1d)
            idxstart = length(quads) + 1
            append!(quads, nquad)
            facetoquad[2, :, cellid] .= idxstart:(idxstart+3)
        else
            error("Expected cellsign ∈ {-1,0,+1}, got cellsign = $s")
        end
    end
    return FaceQuadratures(quads, facetoquad)
end

function Base.getindex(facequads::FaceQuadratures, s, faceid, cellid)
    phaseid = cell_sign_to_row(s)
    return facequads.quads[facequads.facetoquad[phaseid, faceid, cellid]]
end

function uniform_face_quadratures(facequads::FaceQuadratures)
    return facequads.quads[1:facequads.nfaces]
end

function number_of_faces_per_cell(facequads::FaceQuadratures)
    return facequads.nfaces
end

function Base.show(io::IO, facequads::FaceQuadratures)
    ncells = facequads.ncells
    nuniquefacequads = length(facequads.quads)
    str = "FaceQuadratures\n\tNum. Cells: $ncells\n\tNum. Unique Quadratures: $nuniquefacequads"
    print(io, str)
end

function face_quadratures(numqp)
    quad1d = tensor_product_quadrature(1, numqp)
    facequads = [extend_to_face(quad1d, faceid) for faceid = 1:4]
    return facequads
end

function face_quadrature(faceid, levelset, signcondition, quad1d)
    dir, coordval = reference_face(faceid)
    quad = QuadratureRule(ImplicitDomainQuadrature.one_dimensional_quadrature(
        [x -> levelset(extend(x, dir, coordval))],
        [signcondition],
        -1.0,
        +1.0,
        quad1d,
    ))
    return extend_to_face(quad, faceid)
end

function face_quadratures(levelset, signcondition, quad1d)
    quads = [face_quadrature(faceid, levelset, signcondition, quad1d) for faceid = 1:4]
    return quads
end

function extend_to_face(quad::QuadratureRule, faceid)
    extp = extend_to_face(quad.points, faceid)
    return QuadratureRule(extp, quad.weights)
end

function extend_to_face(points, faceid)
    dir, coordval = reference_face(faceid)
    @assert dir == 1 || dir == 2
    flipdir = dir == 1 ? 2 : 1
    return extend([coordval], flipdir, points)
end

function reference_face(faceid)
    if faceid == 1
        return (2, -1.0)
    elseif faceid == 2
        return (1, +1.0)
    elseif faceid == 3
        return (2, +1.0)
    elseif faceid == 4
        return (1, -1.0)
    else
        error("Expected faceid ∈ {1,2,3,4}, got faceid = $faceid")
    end
end

function update_face_quadrature!(facequads::FaceQuadratures, s, faceid, cellid, quad)
    row = cell_sign_to_row(s)
    idx = facequads.facetoquad[row, faceid, cellid]
    @assert idx > 4 "Attempting to modify the uniform cell face quadrature rules"
    facequads.quads[idx] = quad
end