struct GridIndex{T, Index <: AbstractVector{Int}} <: AbstractIndex
    x0::T
    y0::T
    spacing::T
    n_x::Int
    n_y::Int
    start_indices::Vector{Int}
    index::Index
end

Base.summary(i::GridIndex) = "GridIndex ($(i.n_x)×$(i.n_y) cells of spacing $(i.spacing))"

Base.@propagate_inbounds Base.getindex(g::GridIndex, i, j) = g[i+g.n_x*(j-1)]
Base.@propagate_inbounds Base.getindex(g::GridIndex, i) = (g.start_indices[i]):(g.start_indices[i+1]-1)

function AcceleratedArrays.accelerate(points::AbstractVector{<:StaticVector{3, T}}, ::Type{GridIndex}; spacing = one(T)) where {T}
    grid = GridIndex(points; spacing = spacing)
    return AcceleratedArray(points, grid)
end

function AcceleratedArrays.accelerate!(points::AbstractVector{<:StaticVector{3, T}}, ::Type{GridIndex}; spacing = one(T), perm = Ref{Vector{Int}}()) where {T}
    grid = GridIndex(points; spacing = spacing)
    perm[] = grid.index
    permute!(points, grid.index)
    index = keys(points)
    return AcceleratedArray(points, GridIndex(grid.x0, grid.y0, grid.spacing, grid.n_x, grid.n_y, grid.start_indices, index))
end

"""
    GridIndex(points; spacing = 1)

An acceleration index for 3D points using a 2D grid. Inside each grid cell, the points are
ordered in the third dimension. Suitable for "two-and-a-half dimensional data" such as point
clouds on the surface of the Earth, and where the query radius is close to the grid
`spacing`.

Note that `GridIndex`s are typically created by the `accelerate` and `accelerate!`
functions (which also accept the `spacing` keyword argument. In the latter case, the points
are re-ordered to increase query speed.
"""
function GridIndex(points::AbstractVector{<:StaticVector{3, T}}; spacing = one(T)) where {T}
    grid_spacing = convert(T, spacing)
    n_points = length(points)
    if n_points < 1
        return GridIndex(zero(T), zero(T), grid_spacing, 1, 1, [1, 1], Int[])
    end
    if length(points[1]) != 3
        error("Points must be of length 3")
    end

    # Initialize some grid parameters
    (xmin, xmax, ymin, ymax) = bounds(points)
    xmin = prevfloat(xmin)
    ymin = prevfloat(ymin)

    n_x = Int(cld(xmax - xmin, grid_spacing))
    n_y = Int(cld(ymax - ymin, grid_spacing))

    # Pass over the points to construct the grid
    cells = Matrix{Vector{Int}}(undef, n_x, n_y)
    for i = 1:(n_x*n_y)
        @inbounds cells[i] = Vector{Int}()
    end

    @inbounds for i ∈ 1:n_points
        p = points[i]
        cell_x = Int(cld(p[1] - xmin, grid_spacing))
        cell_y = Int(cld(p[2] - ymin, grid_spacing))
        push!(cells[cell_x, cell_y], i)
    end

    # Create the permutation and sort the cells
    permutation = Vector{Int}(undef, n_points)
    start_indices = Vector{Int}(undef, n_x::Int*n_y::Int+1)
    k = 1
    for iy ∈ 1:n_y
        for ix ∈ 1:n_x
            @inbounds start_indices[ix + n_x*(iy-1)] = k
            @inbounds cell = cells[ix,iy]
            sort!(cell, Base.Sort.QuickSortAlg(), Base.Order.By(i -> (Base.@_inline_meta; @inbounds return points[i][3])))
            k_end = k + length(cell) - 1
            @inbounds permutation[k:k_end] = cell
            k = k_end + 1
        end
    end
    @inbounds start_indices[end] = n_points+1

    grid = GridIndex(xmin, ymin, grid_spacing, n_x, n_y, start_indices, permutation)
end

function bounds(points::AbstractVector)
    PointType = eltype(points)
    T = eltype(PointType)
    xmin = typemax(T)
    xmax = typemin(T)
    ymin = typemax(T)
    ymax = typemin(T)

    @inbounds for p ∈ points
        @inbounds x = p[1]
        if x < xmin
            xmin = x
        end
        if x > xmax
            xmax = x
        end

        @inbounds y = p[2]
        if y < ymin
            ymin = y
        end
        if y > ymax
            ymax = y
        end
    end

    return (xmin, xmax, ymin, ymax)
end

function cells_bbox(grid::GridIndex, region::BoundingBox)
    # Find the rough rectangle of intersecting cells
    # Axis-aligned bounding box:
    xmin = region.xmin
    xmax = region.xmax
    ymin = region.ymin
    ymax = region.ymax

    # Make a square grid
    gx_min = Int(cld(xmin - grid.x0, grid.spacing))
    if gx_min < 1
        gx_min = 1
    end
    gx_max = Int(cld(xmax - grid.x0, grid.spacing))
    if gx_max > grid.n_x
        gx_max = grid.n_x
    end

    gy_min = Int(cld(ymin - grid.y0, grid.spacing))
    if gy_min < 1
        gy_min = 1
    end
    gy_max = Int(cld(ymax - grid.y0, grid.spacing))
    if gy_max > grid.n_y
        gy_max = grid.n_y
    end

    return (gx_min, gx_max, gy_min, gy_max)
end

# Accelerations

function Base.findall(pred::Base.Fix2{typeof(in), <:AbstractRegion}, points::AcceleratedArray{<:Any, <:Any, <:Any, <:GridIndex})
    indices = Vector{Int}()
    findall!(indices, pred, points)
    return indices
end

function findall!(indices::Vector{<:Integer}, pred::Base.Fix2{typeof(in), <:AbstractRegion}, points::AcceleratedArray{<:Any, <:Any, <:Any, <:GridIndex})
    # First find the relevant grid cells
    grid = points.index
    bbox = boundingbox(pred.x)
    (gx_min, gx_max, gy_min, gy_max) = cells_bbox(grid, bbox)

    # Now we'll iterate over these cells and fill the indices
    @inbounds for gy ∈ gy_min:gy_max
        for gx ∈ gx_min:gx_max
            for i ∈ grid[gx, gy]
                # Could use min/max height information here, as in old code
                j = grid.index[i]
                p = points[j]

                if p[3] > bbox.zmax
                    break
                end

                if pred(p)
                    push!(indices, j)
                end
            end
        end
    end

    return indices
end

function Base.count(pred::Base.Fix2{typeof(in), <:AbstractRegion}, points::AcceleratedArray{<:Any, <:Any, <:Any, <:GridIndex})
    # First find the relevant grid cells
    grid = points.index
    (gx_min, gx_max, gy_min, gy_max) = cells_bbox(grid, boundingbox(pred.x))

    out = 0

    # Now we'll iterate over these cells and fill the indices
    @inbounds for gy ∈ gy_min:gy_max
        for gx ∈ gx_min:gx_max
            for i ∈ grid[gx, gy]
                # Could use min/max height information here, as in commented code below
                j = grid.index[i]
                p = points[j]

                out += pred(p)
            end
        end
    end

    return out
end

function Base.filter(pred::Base.Fix2{typeof(in), <:AbstractRegion}, points::AcceleratedArray{<:Any, <:Any, <:Any, <:GridIndex})
    # First find the relevant grid cells
    grid = points.index
    (gx_min, gx_max, gy_min, gy_max) = cells_bbox(grid, boundingbox(pred.x))

    out = empty(points.parent)

    # Now we'll iterate over these cells and fill the indices
    @inbounds for gy ∈ gy_min:gy_max
        for gx ∈ gx_min:gx_max
            for i ∈ grid[gx, gy]
                # Could use min/max height information here, as in commented code below
                j = grid.index[i]
                p = points[j]

                if pred(p)
                    push!(out, p)
                end
            end
        end
    end

    return out
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    RotationConvention

Convention used to specify rotations in 2D and 3D spaces.
"""
abstract type RotationConvention end

"""
    axesseq(convention) -> Symbol

Sequence of three axes by which the rotations are made, e.g. `:ZXZ`.
"""
function axesseq(::Type{<:RotationConvention}) end

"""
    orientation(convention) -> Tuple{Symbol,Symbol,Symbol}

Informs for each of the three rotation if it is clockwise (`:CW`) or
counter-clockwise (`:CCW`). Adopts right-hand rule: orientation defined
looking towards the negative direction of the axis.
"""
function orientation(::Type{<:RotationConvention}) end

"""
    angleunits(convention) -> Symbol

Informs the angle units of the rotations in radians (`:RAD`) or in degrees (`:DEG`).
"""
function angleunits(::Type{<:RotationConvention}) end

"""
    mainaxis(convention) -> Symbol

Informs if the main axis of the rotation is `:X` or `:Y`.
"""
function mainaxis(::Type{<:RotationConvention}) end

"""
    isextrinsic(convention) -> Bool

Informs if the rotation is extrinsic or intrinsic.
"""
function isextrinsic(::Type{<:RotationConvention}) end

"""
    rotmat(angles, convention)

Return the direction cosine matrix of rotation for the given `angles`
according to the rotation `convention`.
"""
function rotmat(angles, convention::Type{<:RotationConvention})
  nangles = length(angles)
  @assert nangles ∈ [1,3] "invalid number of angles"

  Dim = nangles == 1 ? 2 : 3

  # convert to radian if necessary
  θs = angleunits(convention) == :DEG ? deg2rad.(angles) : angles

  # invert sign if necessary
  _0 = zero(eltype(θs))
  Dim == 2 && (θs = [θs[1], _0, _0])
  intr = (orientation(convention) .== :CW) .& !isextrinsic(convention)
  extr = (orientation(convention) .== :CCW) .& isextrinsic(convention)
  inds = collect(intr .| extr)
  θs[inds] *= -1

  # rotation matrix
  R = angle_to_dcm(θs..., axesseq(convention))[SOneTo(Dim),SOneTo(Dim)]

  isextrinsic(convention) ? R' : R
end

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

include("conventions/TaitBryanExtr.jl")
include("conventions/TaitBryanIntr.jl")
include("conventions/EulerExtr.jl")
include("conventions/EulerIntr.jl")
include("conventions/GSLIB.jl")
include("conventions/Leapfrog.jl")
include("conventions/Datamine.jl")


"""
# RootSolvers

Module containing functions for solving roots of non-linear
equations. See [`find_zero`](@ref).

## Example

```
using CLIMA.Utilities.RootSolvers

x_root, converged = find_zero(x -> x^2 - 100^2, 0.0, 1000.0, SecantMethod())
```

"""
module RootSolvers

export find_zero, SecantMethod, RegulaFalsiMethod, NewtonsMethod

import ForwardDiff

abstract type RootSolvingMethod end
Base.broadcastable(method::RootSolvingMethod) = Ref(method)

struct SecantMethod <: RootSolvingMethod end
struct RegulaFalsiMethod <: RootSolvingMethod end
struct NewtonsMethod <: RootSolvingMethod end

# TODO: CuArrays.jl has trouble with isapprox on 1.1
# we use simple checks for now, will switch to relative checks later.

"""
    x, converged = find_zero(f, x0[, x1], method;
                             xatol=0, xrtol=sqrt(eps(eltype(x0))),
                             yatol=sqrt(eps(eltype(x0))), maxiters=10_000)

Finds the nearest root of `f` to `x0` and `x1`. Returns a the value of the root `x` such
that `f(x) ≈ 0`, and a Boolean value `converged` indicating convergence.

`method` can be one of:
- `SecantMethod()`: [Secant method](https://en.wikipedia.org/wiki/Secant_method)
- `RegulaFalsiMethod()`: [Regula Falsi method](https://en.wikipedia.org/wiki/False_position_method#The_regula_falsi_(false_position)_method).
- `NewtonsMethod()`: [Newton's method](https://en.wikipedia.org/wiki/Newton%27s_method)
  - The `x1` argument is omitted for Newton's method.

The keyword arguments:
- `xatol` is the absolute tolerance of the input.
- `maxiters` is the maximum number of iterations.
"""
function find_zero end

function find_zero(f::F, x0::T, x1::T, ::SecantMethod;
                   xatol=T(1e-3),
                   maxiters=10_000) where {F, T<:AbstractFloat}
  y0 = f(x0)
  y1 = f(x1)
  for i in 1:maxiters
    Δx = x1 - x0
    Δy = y1 - y0
    x0, y0 = x1, y1
    x1 -= y1 * Δx / Δy
    y1 = f(x1)
    if abs(x0-x1) <= xatol
      return x1, true
    end
  end
  return x1, false
end

function find_zero(f::F, x0::T, x1::T, ::RegulaFalsiMethod;
                   xatol=T(1e-3),
                   maxiters=10_000) where {F, T<:AbstractFloat}
  y0 = f(x0)
  y1 = f(x1)
  @assert y0 * y1 < 0
  lastside = 0
  local x
  for i in 1:maxiters
    x = (x0 * y1 - x1 * y0)/ (y1 - y0)
    y = f(x)
    if y * y0 < 0
      if abs(x-x1) <= xatol
        return x, true
      end
      x1, y1 = x, y
      if lastside == +1
        y0 /= 2
      end
      lastside = +1
    else
      if abs(x0-x) <= xatol
        return x, true
      end
      x0, y0 = x, y
      if lastside == -1
        y1 /= 2
      end
      lastside = -1
    end
  end
  return x, false
end

"""
    value_deriv(f, x)

Compute the value and derivative `f(x)` using ForwardDiff.jl.
"""
function value_deriv(f, x::T) where {T}
    tag = typeof(ForwardDiff.Tag(f, T))
    y = f(ForwardDiff.Dual{tag}(x,one(x)))
    ForwardDiff.value(tag, y), ForwardDiff.partials(tag, y, 1)
end

function find_zero(f::F, x0::T, ::NewtonsMethod;
                   xatol=1e-3,
                   maxiters=10_000) where {F, T<:AbstractFloat}
  for i in 1:maxiters
    y,y′ = value_deriv(f, x0)
    x1 = x0 - y/y′
    if abs(x0-x1) <= xatol
      return x1, true
    end
    x0 = x1
  end
  return x0, false
end

end


using StrainRecon
using Base.Test

@testset "StrainRecon" begin
	println("=== testing StrainRecon.jl code ===")
	println("")
	include("testFindw.jl")
	include("testFindM.jl")
	include("testSolveMAP.jl")
	include("testSolveConvexMAP.jl")
	include("testReducedMisfit.jl")
	include("testRandomw.jl")
	include("testRandomM.jl")
	include("testBCD.jl")
	include("testIntegratePosterior.jl")
	include("testASQP.jl")
	include("testIPLP.jl")
	include("testEntropyM.jl")
	include("testSrDistance.jl")
	include("testIntmat2binmat.jl")
end


#jl #! format: off
# Using PowerSystems to calculate network matrices

# **Originally Contributed by**: Clayton Barrows

# ## Introduction

# PowerSystems.jl supports the calculation of several different matrix representations of
# power system networks. This example demonstrates how to use PowerSystems.jl to calculate:
#  - Y bus
#  - Power transfer distribution factor (PTDF)
#  - Line outage distribution  factor (LODF)

# ### Dependencies
# Let's use a dataset from the [tabular data parsing example](https://nbviewer.jupyter.org/github/NREL-SIIP/SIIPExamples.jl/blob/master/notebook/2_PowerSystems_examples/parse_matpower.ipynb)
using SIIPExamples

pkgpath = pkgdir(SIIPExamples)
include(joinpath(pkgpath, "test", "2_PowerSystems_examples", "02_parse_matpower.jl"))

# ### Ybus
ybus = Ybus(sys)

# ### PTDF
ptdf = PTDF(sys)

# ### LODF
lodf = LODF(sys)

# ### Indexing
# Note that the axes of these matrices that correspond to buses are indexed by bus number
# (::Int64) while the branch axes are indexed by branch name (::String). You can access
# specific elements of the matrices as follows:

ptdf["bus3-bus4-i_6", 3]

# Additionally, PowerSystems provides accessors to the network matrices that take `Componets`
# as arguments so that you can pass references to the components themselves rather than the
# name or number. For example:
buses = collect(get_components(Bus, sys))
ybus[buses[1], buses[2]]

# If you would instead like to index by bus name, something like the following should work:
busname2num = get_components(Bus, sys) |> (c -> Dict(zip(get_name.(c), get_number.(c))))
ptdf["bus3-bus4-i_6", busname2num["bus3"]]


using RDatasets
using DataFrames
random_dataframe = DataFrame(A = randn(5), B = randn(5), C = randn(5))
random_modelframe = ModelFrame(A ~ B + C, random_dataframe)
random_modelmatrix = ModelMatrix(ModelFrame(A ~ B + C, random_dataframe))


# This file reproduces just enough functionality of BEAST to allow
# for meaningful testing of non-trivial user cases
using StaticArrays
using CompScienceMeshes

struct RTRefSpace{T<:Real} end

function (f::RTRefSpace{T})(x) where {T}
    u, v = parametric(x)
    j = jacobian(x)

    tu = tangents(x,1)
    tv = tangents(x,2)

    d = 2/j

    return SVector((
        ((tu*(u-1) + tv*v    ) / j, d),
        ((tu*u     + tv*(v-1)) / j, d),
        ((tu*u     + tv*v    ) / j, d)
    ))
end

numfunctions(::RTRefSpace) = 3


using JuMP
using Ipopt
using CPLEX
include("io.jl")

function create_and_solve_det_milp(data::Dict{AbstractString,Any}, ref::Dict{Symbol,Any}, attacks::Vector{Attack}; k = 2)

    ϵ = 1e-5
    m = Model()
    # setsolver(m, GurobiSolver(LogToConsole=0))
    # setsolver(m, GurobiSolver())
    setsolver(m, CplexSolver(CPX_PARAM_PREIND=0))

    notgen_bus_indexes = Int[] # includes plain buses, buses with loads and condensors
    gen_bus_indexes = Int[] # includes only generators and not condensors

    w_vertex = Dict{Int,Float64}()
    for i in keys(ref[:bus])
        buses = ref[:bus]
        if buses[i]["bus_type"] == 2 || buses[i]["bus_type"] == 3
            total_pmin = sum([abs(ref[:gen][g]["pmin"]) for g in ref[:bus_gens][i]])
            total_pmax = sum([abs(ref[:gen][g]["pmax"]) for g in ref[:bus_gens][i]])
            is_condensor = (total_pmin == 0 && total_pmax == 0) ? true : false
            if is_condensor
                w_vertex[i] = abs(buses[i]["pd"])
                push!(notgen_bus_indexes, i)
                continue
            else
                w_vertex[i] = 0.0
                push!(gen_bus_indexes, i)
                continue
            end
        else
            w_vertex[i] = abs(buses[i]["pd"])
            push!(notgen_bus_indexes, i)
        end
    end

    @variable(m, x[i in keys(ref[:branch])], Bin) # x[e] = 1 means the edge is removed
    @variable(m, y[v in keys(ref[:bus])], Bin) # y[v] = 1 means the vertex is isolated
    @variable(m, dummy_edge[v in gen_bus_indexes] == 0) # dummy edges are added between a source and a gen_bus

    w_edge = Dict(i => 1 for i in keys(ref[:branch]))
    @constraint(m, gen_off[v in gen_bus_indexes], y[v] == 0)
    @constraint(m, sum(w_edge[i]*x[i] for i in keys(ref[:branch])) == k)
    @objective(m, Max, sum(w_vertex[i]*y[i] for i in keys(ref[:bus])))

    index_map = IndexMap()
    bus_indexes = collect(keys(ref[:bus]))
    for i in 1:length(bus_indexes)
        index_map.bus_to_vertex[bus_indexes[i]] = i
        index_map.vertex_to_bus[i] = bus_indexes[i]
    end
    index_map.bus_to_vertex[0] = length(keys(ref[:bus])) + 1
    index_map.vertex_to_bus[length(keys(ref[:bus]))+1] = 0

    function lazycuts(cb)
        x_vals = Dict(i => getvalue(x[i]) for i in keys(ref[:branch]))
        y_vals = Dict(v => getvalue(y[v]) for v in notgen_bus_indexes)
        y_vals = filter((v, val) -> val >= (1-ϵ), y_vals)
        g = simple_graph(length(keys(index_map.vertex_to_bus)), is_directed=false)
        dists = Float64[]
        # create support graph edges
        for (key, value) in x_vals
            f_bus = ref[:branch][key]["f_bus"]
            t_bus = ref[:branch][key]["t_bus"]
            i = index_map.bus_to_vertex[f_bus]
            j = index_map.bus_to_vertex[t_bus]
            add_edge!(g, i, j)
            if !haskey(index_map.parallel_edges, (i, j))
                index_map.parallel_edges[(i, j)] = Dict{AbstractString, Any}("count" => 1, "index" => [key])
                index_map.parallel_edges[(j, i)] = Dict{AbstractString, Any}("count" => 1, "index" => [key])
            else
                index_map.parallel_edges[(i, j)]["count"] += 1
                push!(index_map.parallel_edges[(i, j)]["index"], key)
                index_map.parallel_edges[(j, i)]["count"] += 1
                push!(index_map.parallel_edges[(j, i)]["index"], key)
            end
            push!(dists, value)
        end
        # create support graph dummy edges
        for gen in gen_bus_indexes
            f_bus = 0
            t_bus = gen
            i = index_map.bus_to_vertex[f_bus]
            j = index_map.bus_to_vertex[t_bus]
            add_edge!(g, i, j)
            if !haskey(index_map.parallel_edges, (i, j))
                index_map.parallel_edges[(i, j)] = Dict{AbstractString, Any}("count" => 1, "index" => [0])
                index_map.parallel_edges[(j, i)] = Dict{AbstractString, Any}("count" => 1, "index" => [0])
            end
            push!(dists, 0.0)
        end
        r = dijkstra_shortest_paths(g, dists, index_map.bus_to_vertex[0])
        if length(keys(y_vals)) != 0
            paths = enumerate_indices(r.parent_indices, [index_map.bus_to_vertex[i] for i in keys(y_vals)])
            for path in paths
                target_bus = index_map.vertex_to_bus[path[length(path)]]
                @assert path[1] == length(keys(ref[:bus]))+1
                @assert target_bus in notgen_bus_indexes
                @assert target_bus in keys(y_vals)
                path_indexes = Int[]
                for i in 2:length(path)-1
                    path_index = 0.0
                    index = index_map.parallel_edges[(path[i], path[i+1])]
                    if index["count"] > 1
                        path_index = (x_vals[index["index"][1]] > 1-ϵ) ? index["index"][2] : index["index"][1]
                    else
                        path_index = index["index"][1]
                    end
                    push!(path_indexes, path_index)
                end

                if sum([x_vals[i] for i in path_indexes]) < y_vals[target_bus]
                    @lazyconstraint(cb, sum(x[i] for i in path_indexes) >= y[target_bus])
                end
            end
        end
    end

    addlazycallback(m, lazycuts)
    status = solve(m)
    x_vals = Dict(i => getvalue(x[i]) for i in keys(ref[:branch]))
    x_vals = filter((i, val) -> val > (1-ϵ), x_vals)
    y_vals = Dict(i => getvalue(y[i]) for i in notgen_bus_indexes)
    y_vals = filter((i, val) -> val > (1-ϵ), y_vals)
    attack = Attack()
    for bus_id in collect(keys(y_vals))
        push!(attack.bus_ids, bus_id)
    end
    for branch_id in collect(keys(x_vals))
        push!(attack.branch_ids, branch_id)
    end
    attack.total_load = getobjectivevalue(m)
    return attack
end

function create_and_solve_stoch_milp(data::Dict{AbstractString,Any}, ref::Dict{Symbol,Any}, attacks::Vector{Attack}, log_p; k = 2)

    ϵ = 1e-5
    m = Model()
    setsolver(m, CplexSolver(CPX_PARAM_PREIND=0))

    notgen_bus_indexes = Int[] # includes plain buses, buses with loads and condensors
    gen_bus_indexes = Int[] # includes only generators and not condensors

    w_vertex = Dict{Int,Float64}()
    for i in keys(ref[:bus])
        buses = ref[:bus]
        if buses[i]["bus_type"] == 2 || buses[i]["bus_type"] == 3
            total_pmin = sum([abs(ref[:gen][g]["pmin"]) for g in ref[:bus_gens][i]])
            total_pmax = sum([abs(ref[:gen][g]["pmax"]) for g in ref[:bus_gens][i]])
            is_condensor = (total_pmin == 0 && total_pmax == 0) ? true : false
            if is_condensor
                w_vertex[i] = abs(buses[i]["pd"])
                push!(notgen_bus_indexes, i)
                continue
            else
                w_vertex[i] = 0.0
                push!(gen_bus_indexes, i)
                continue
            end
        else
            w_vertex[i] = abs(buses[i]["pd"])
            push!(notgen_bus_indexes, i)
        end
    end

    @variable(m, x[i in keys(ref[:branch])], Bin) # x[e] = 1 means the edge is removed
    @variable(m, y[v in keys(ref[:bus])], Bin) # y[v] = 1 means the vertex is isolated
    @variable(m, dummy_edge[v in gen_bus_indexes] == 0) # dummy edges are added between a source and a gen_bus
    @variable(m, p <= 0)
    @variable(m, z <= 1e6)
    @variable(m, 1e-6 <= eta <= 1e6)

    w_edge = Dict(i => 1 for i in keys(ref[:branch]))
    @constraint(m, gen_off[v in gen_bus_indexes], y[v] == 0)
    @constraint(m, sum(w_edge[i]*x[i] for i in keys(ref[:branch])) == k)
    @constraint(m, p == sum(x[i]*log_p[i] for i in keys(ref[:branch])))
    @constraint(m, eta == sum(w_vertex[i]*y[i] for i in keys(ref[:bus])))
    @objective(m, Max, z)

    index_map = IndexMap()
    bus_indexes = collect(keys(ref[:bus]))
    for i in 1:length(bus_indexes)
        index_map.bus_to_vertex[bus_indexes[i]] = i
        index_map.vertex_to_bus[i] = bus_indexes[i]
    end
    index_map.bus_to_vertex[0] = length(bus_indexes) + 1
    index_map.vertex_to_bus[length(bus_indexes)+1] = 0

    function lazycuts(cb)
        eta_val = getvalue(eta)
        p_val = getvalue(p)
        if eta_val < ϵ
            eta_val = 0.01
        end
        if eta_val > 0.0
            inv_eta_val = 1/eta_val
            @lazyconstraint(cb, z <= p + log(eta_val) + inv_eta_val*(eta - eta_val))
        end

        x_vals = Dict(i => getvalue(x[i]) for i in keys(ref[:branch]))
        y_vals = Dict(v => getvalue(y[v]) for v in notgen_bus_indexes)
        y_vals = filter((v, val) -> val >= (1-ϵ), y_vals)
        g = simple_graph(length(keys(index_map.vertex_to_bus)), is_directed=false)
        dists = Float64[]
        # create support graph edges
        for (key, value) in x_vals
            f_bus = ref[:branch][key]["f_bus"]
            t_bus = ref[:branch][key]["t_bus"]
            i = index_map.bus_to_vertex[f_bus]
            j = index_map.bus_to_vertex[t_bus]
            add_edge!(g, i, j)
            if !haskey(index_map.parallel_edges, (i, j))
                index_map.parallel_edges[(i, j)] = Dict{AbstractString, Any}("count" => 1, "index" => [key])
                index_map.parallel_edges[(j, i)] = Dict{AbstractString, Any}("count" => 1, "index" => [key])
            else
                index_map.parallel_edges[(i, j)]["count"] += 1
                push!(index_map.parallel_edges[(i, j)]["index"], key)
                index_map.parallel_edges[(j, i)]["count"] += 1
                push!(index_map.parallel_edges[(j, i)]["index"], key)
            end
            push!(dists, value)
        end
        # create support graph dummy edges
        for gen in gen_bus_indexes
            f_bus = 0
            t_bus = gen
            i = index_map.bus_to_vertex[f_bus]
            j = index_map.bus_to_vertex[t_bus]
            add_edge!(g, i, j)
            if !haskey(index_map.parallel_edges, (i, j))
                index_map.parallel_edges[(i, j)] = Dict{AbstractString, Any}("count" => 1, "index" => [0])
                index_map.parallel_edges[(j, i)] = Dict{AbstractString, Any}("count" => 1, "index" => [0])
            end
            push!(dists, 0.0)
        end
        r = dijkstra_shortest_paths(g, dists, index_map.bus_to_vertex[0])
        if length(keys(y_vals)) != 0
            paths = enumerate_indices(r.parent_indices, [index_map.bus_to_vertex[i] for i in keys(y_vals)])
            for path in paths
                target_bus = index_map.vertex_to_bus[path[length(path)]]
                @assert path[1] == length(keys(ref[:bus]))+1
                @assert target_bus in notgen_bus_indexes
                @assert target_bus in keys(y_vals)
                path_indexes = Int[]
                for i in 2:length(path)-1
                    path_index = 0.0
                    index = index_map.parallel_edges[(path[i], path[i+1])]
                    if index["count"] > 1
                        path_index = (x_vals[index["index"][1]] > 1-ϵ) ? index["index"][2] : index["index"][1]
                    else
                        path_index = index["index"][1]
                    end
                    push!(path_indexes, path_index)
                end

                if sum([x_vals[i] for i in path_indexes]) < y_vals[target_bus]
                    @lazyconstraint(cb, sum(x[i] for i in path_indexes) >= y[target_bus])
                end
            end
        end
    end

    addlazycallback(m, lazycuts)
    status = solve(m)
    x_vals = Dict(i => getvalue(x[i]) for i in keys(ref[:branch]))
    x_vals = filter((i, val) -> val > (1-ϵ), x_vals)
    y_vals = Dict(i => getvalue(y[i]) for i in notgen_bus_indexes)
    y_vals = filter((i, val) -> val > (1-ϵ), y_vals)
    p_val = getvalue(p)
    actual_prob = exp(p_val)
    log_eta = getobjectivevalue(m) - p_val
    eta_val = exp(log_eta)
    exp_load_shed = actual_prob*eta_val
    attack = Attack()
    for bus_id in collect(keys(y_vals))
        push!(attack.bus_ids, bus_id)
    end
    for branch_id in collect(keys(x_vals))
        push!(attack.branch_ids, branch_id)
    end
    attack.total_load = exp_load_shed
    return attack
end

type IndexMap
    bus_to_vertex::Dict{Int,Int}
    vertex_to_bus::Dict{Int,Int}
    parallel_edges::Dict{Any,Any}
    function IndexMap()
        imap = new()
        imap.bus_to_vertex = Dict{Int,Int}()
        imap.vertex_to_bus = Dict{Int,Int}()
        imap.parallel_edges = Dict{Any,Any}()
        return imap
    end
end


@testset "PLL discriminator" begin

    correlator_minus60off = EarlyPromptLateCorrelator(
        SVector(
            -0.5 + sqrt(3) / 2im,
            -1 + sqrt(3) * 1im,
            -0.5 + sqrt(3) / 2im
        )
    )
    correlator_0off = EarlyPromptLateCorrelator(
        SVector(
            0.5 + 0.0im,
            1 + 0.0im,
            0.5 + 0.0im
        )
    )
    correlator_plus60off = EarlyPromptLateCorrelator(
        SVector(
            0.5 + sqrt(3) / 2im,
            1 + sqrt(3) * 1im,
            0.5 + sqrt(3) / 2im
        )
    )
    gpsl1 = GPSL1()
    correlator_sample_shifts = SVector(-1, 0, 1)
    @test @inferred(Tracking.pll_disc(gpsl1, correlator_minus60off, correlator_sample_shifts)) == -π / 3  #-60°
    @test @inferred(Tracking.pll_disc(gpsl1, correlator_0off, correlator_sample_shifts)) == 0
    @test @inferred(Tracking.pll_disc(gpsl1, correlator_plus60off, correlator_sample_shifts)) == π / 3  #+60°
end

@testset "DLL discriminator" begin

    sample_shifts = SVector(-2, 0, 2)
    index_shift = 1
    delta = 0.25
    @test @inferred(get_early_late_sample_spacing(sample_shifts, index_shift)) == 4

    gpsl1 = GPSL1()
    very_early_correlator = EarlyPromptLateCorrelator(SVector(1.0 + 0.0im, 0.5 + 0.0im, 0.0 + 0.0im))
    early_correlator = EarlyPromptLateCorrelator(SVector(0.75 + 0.0im, 0.75 + 0.0im, 0.25 + 0.0im))
    prompt_correlator = EarlyPromptLateCorrelator(SVector(0.5 + 0.0im, 1 + 0.0im, 0.5 + 0.0im))
    late_correlator = EarlyPromptLateCorrelator(SVector(0.25 + 0.0im, 0.75 + 0.0im, 0.75 + 0.0im))
    very_late_correlator = EarlyPromptLateCorrelator(SVector(0.0 + 0.0im, 0.5 + 0.0im, 1.0 + 0.0im))

    @test @inferred(
        Tracking.dll_disc(gpsl1, very_early_correlator, sample_shifts, index_shift, delta)
    ) == -0.5
    @test @inferred(
        Tracking.dll_disc(gpsl1, early_correlator, sample_shifts, index_shift, delta)
    ) == -0.25
    @test @inferred(
        Tracking.dll_disc(gpsl1, prompt_correlator, sample_shifts, index_shift, delta)
    ) == 0
    @test @inferred(
        Tracking.dll_disc(gpsl1, late_correlator, sample_shifts, index_shift, delta)
    ) == 0.25
    @test @inferred(
        Tracking.dll_disc(gpsl1, very_late_correlator, sample_shifts, index_shift, delta)
    ) == 0.5
end


immutable FidMtx
  MVxxBVxx   :: Array{Float64,1}
  MVxxBTxx   :: SparseMatrixCSC{Float64,Int64}
  MVxzBVxz   :: Array{Float64,1}
  MVxzBTxz   :: SparseMatrixCSC{Float64,Int64}
  MVzxBVzx   :: Array{Float64,1}
  MVzxBTxz   :: SparseMatrixCSC{Float64,Int64}
  MVzzBVzz   :: Array{Float64,1}
  MVzzBTzz   :: SparseMatrixCSC{Float64,Int64}
  MTxxxBVx   :: SparseMatrixCSC{Float64,Int64}
  MTxxxBTxxx :: Array{Float64,1}
  MTxxzBVz   :: SparseMatrixCSC{Float64,Int64}
  MTxxzBTxxz :: Array{Float64,1}
  MTzzxBVx   :: SparseMatrixCSC{Float64,Int64}
  MTzzxBTzzx :: Array{Float64,1}
  MTzzzBVz   :: SparseMatrixCSC{Float64,Int64}
  MTzzzBTzzz :: Array{Float64,1}
  MTxzxBVz   :: SparseMatrixCSC{Float64,Int64}
  MTxzxBTxzx :: Array{Float64,1}
  MTxzzBVx   :: SparseMatrixCSC{Float64,Int64}
  MTxzzBTxzz :: Array{Float64,1}
end

function vel2lm(vp::Array{Float64,2}, vs::Array{Float64,2}, rho::Array{Float64,2})
    (m, n) = size(vp)
    lambda = zeros(typeof(vp[1]), m, n)
    mu     = zeros(typeof(vp[1]), m, n)
    for j = 1: n
        for i = 1: m
            mu[i,j] = rho[i,j] * vs[i,j]^2
            lambda[i,j] = rho[i,j] * vp[i,j]^2 - 2*mu[i,j]
        end
    end
    return lambda, mu
end

function lm2vel(lambda::Array{Float64,2}, mu::Array{Float64,2}, rho::Array{Float64,2})
    (m, n) = size(lambda)
    vp = Array(typeof(lambda[1]), m, n)
    vs = Array(typeof(lambda[1]), m, n)
    for j = 1: n
        for i = 1: m
            vp[i,j] = sqrt((lambda[i,j]+2*mu[i,j])/rho[i,j])
            vs[i,j] = sqrt(mu[i,j]/rho[i,j])
        end
    end
    return vp, vs
end

function DispStable!(vmax::Float64, vmin::Float64, f0::Float64, dz::Float64, dx::Float64, dt::Float64)
    h = vmin / 5 / f0
    d = minimum([dz dx])
    tt = 6*d / (7*sqrt(3)*vmax)
    if dx >= h || dz >= h || dt >= tt
       println("maximum acceptable grid size: $h")
       println("maximum acceptable time step: $tt")
       error("unstable or frequency dispersion")
    end
    return nothing
end

function MVxxBVxx(dVxx::SparseMatrixCSC{Float64,Int64}, rho::Array{Float64,2}, dt::Float64)
    (m,n) = size(dVxx)
    Mdiag = zeros(m*n)
    for ix = 1 : n-1
        for iz = 1 : m
            a3 = (rho[iz,ix+1]+rho[iz,ix])/2/dt + dVxx[iz,ix]/2
            tmp= (rho[iz,ix+1]+rho[iz,ix])/2/(dt*a3) - dVxx[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    for iz = 1 : m
        a3 = rho[iz,n]/dt + dVxx[iz,n]/2
        tmp= rho[iz,n]/(dt*a3) - dVxx[iz,n]/(2*a3)
        Mdiag[(n-1)*m+iz] = tmp
    end
    M = Mdiag
    return M
end

function MVxxBTxx(dVxx::SparseMatrixCSC{Float64,Int64}, rho::Array{Float64,2}, dt::Float64, dx::Float64)
    a1 =  9/8
    a2 = -1/24
    (m, n) = size(dVxx)
    A3 = zeros(m*n)
    for ix = 1:n-1
        for iz = 1:m
            A3[(ix-1)*m+iz] = 1 / ((rho[iz,ix+1]+rho[iz,ix])/2/dt + dVxx[iz,ix]/2)
        end
    end
    for iz = 1 : m
        A3[(n-1)*m+iz] =  1 / (rho[iz,n]/dt + dVxx[iz,n]/2)
    end
    A = spzeros(n, n)
    A[1, 1] = -a1/dx; A[1, 2] = a1/dx; A[1, 3] = a2/dx;
    for ix = 2 : n-2
        A[ix, ix-1] = -a2/dx
        A[ix, ix  ] = -a1/dx
        A[ix, ix+1] =  a1/dx
        A[ix, ix+2] =  a2/dx
    end
    A[n-1, n-2] = -a2/dx;
    A[n-1, n-1] = -a1/dx; A[n-1, n  ] = a1/dx;
    A[n  , n-1] = -a2/dx; A[n  , n  ] = -a1/dx;
    Im = speye(m)
    M = kron(A, Im)
    M = spdiagm(A3) * M
    return M
end

function MVxzBVxz(dVxz::SparseMatrixCSC{Float64,Int64}, rho::Array{Float64,2}, dt::Float64)
    (m,n) = size(dVxz)
    Mdiag = zeros(m*n)
    for ix = 1 : n-1
        for iz = 1 : m
            a3 = (rho[iz,ix+1]+rho[iz,ix])/2/dt + dVxz[iz,ix]/2
            tmp= (rho[iz,ix+1]+rho[iz,ix])/2/(dt*a3) - dVxz[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    for iz = 1 : m
        a3 = rho[iz,n]/dt + dVxz[iz,n]/2
        tmp= rho[iz,n]/(dt*a3) - dVxz[iz,n]/(2*a3)
        Mdiag[(n-1)*m+iz] = tmp
    end
    M = Mdiag
    return M
end

function MVxzBTxz(dVxz::SparseMatrixCSC{Float64, Int64}, rho::Array{Float64,2}, dt::Float64, dz::Float64)
    a1 =  9/8
    a2 = -1/24
    (m,n) = size(dVxz)
    A3 = zeros(m*n)
    for ix = 1:n-1
        for iz = 1:m
            A3[(ix-1)*m+iz] = 1 / ((rho[iz,ix+1]+rho[iz,ix])/2/dt + dVxz[iz,ix]/2)
        end
    end
    for iz = 1 : m
        A3[(n-1)*m+iz] =  1 / (rho[iz,n]/dt + dVxz[iz,n]/2)
    end
    A = spzeros(m, m)
    A[1, 1] = a1/dz; A[1, 2] = a2/dz;
    A[2, 1] =-a1/dz; A[2, 2] = a1/dz; A[2, 3] = a2/dz;
    for iz = 3 : m-1
        A[iz, iz-2] = -a2/dz
        A[iz, iz-1] = -a1/dz
        A[iz, iz  ] =  a1/dz
        A[iz, iz+1] =  a2/dz
    end
    A[m, m-2] = -a2/dz;
    A[m, m-1] = -a1/dz; A[m, m] = a1/dz;
    In = speye(n)
    M = kron(In, A)
    M = spdiagm(A3) * M
    return M
end

function MVzxBVzx(dVzx::SparseMatrixCSC{Float64,Int64}, rho::Array{Float64,2}, dt::Float64)
    (m,n) = size(dVzx)
    Mdiag = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m-1
            a3 = (rho[iz+1,ix]+rho[iz,ix])/2/dt + dVzx[iz,ix]/2
            tmp= (rho[iz+1,ix]+rho[iz,ix])/2/(dt*a3) - dVzx[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    for ix = 1 : n
        a3 = rho[m,ix]/dt + dVzx[m,ix]/2
        tmp= rho[m,ix]/(dt*a3) - dVzx[m,ix]/(2*a3)
        Mdiag[(ix-1)*m+m] = tmp
    end
    M = Mdiag
    return M
end

function MVzxBTxz(dVzx::SparseMatrixCSC{Float64, Int64}, rho::Array{Float64,2}, dt::Float64, dx::Float64)
    a1 = 9/8
    a2 = -1/24
    (m,n) = size(dVzx)
    A3 = zeros(m*n)
    for ix = 1:n
        for iz = 1:m-1
            A3[(ix-1)*m+iz] = 1 / ((rho[iz+1,ix]+rho[iz,ix])/2/dt+dVzx[iz,ix]/2)
        end
    end
    for ix = 1 : n
        A3[(ix-1)*m+m] =  1 / (rho[m,ix]/dt + dVzx[m,ix]/2)
    end
    A = spzeros(n, n)
    A[1, 1] = a1/dx; A[1, 2] = a2/dx;
    A[2, 1] =-a1/dx; A[2, 2] = a1/dx;
    A[2, 3] = a2/dx;
    for ix = 3:n-1
        A[ix, ix-2] = -a2/dx
        A[ix, ix-1] = -a1/dx
        A[ix, ix  ] =  a1/dx
        A[ix, ix+1] =  a2/dx
    end
    A[n, n-2] =-a2/dx;
    A[n, n-1] =-a1/dx; A[n, n] = a1/dx;
    Im = speye(m)
    M = kron(A, Im)
    M = spdiagm(A3) * M
    return M
end

function MVzzBVzz(dVzz::SparseMatrixCSC{Float64,Int64}, rho::Array{Float64,2}, dt::Float64)
    (m,n) = size(dVzz)
    Mdiag = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m-1
            a3 = (rho[iz+1,ix]+rho[iz,ix])/2/dt + dVzz[iz,ix]/2
            tmp= (rho[iz+1,ix]+rho[iz,ix])/2/(dt*a3) - dVzz[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    for ix = 1 : n
        a3 = rho[m,ix]/dt + dVzz[m,ix]/2
        tmp= rho[m,ix]/(dt*a3) - dVzz[m,ix]/(2*a3)
        Mdiag[(ix-1)*m+m] = tmp
    end
    M = Mdiag
    return M
end

function MVzzBTzz(dVzz::SparseMatrixCSC{Float64,Int64}, rho::Array{Float64,2}, dt::Float64, dz::Float64)
    a1 = 9/8
    a2 = -1/24
    (m,n) = size(dVzz)
    A3 = zeros(m*n)
    for ix = 1:n
        for iz = 1:m-1
            A3[(ix-1)*m+iz] = 1 / ((rho[iz+1,ix]+rho[iz,ix])/2/dt + dVzz[iz,ix]/2)
        end
    end
    for ix = 1 : n
        A3[(ix-1)*m+m] =  1 / (rho[m,ix]/dt + dVzz[m,ix]/2)
    end
    A = spzeros(m, m)
    A[1, 1] =-a1/dz; A[1, 2] = a1/dz; A[1, 3] = a2/dz;
    for iz = 2 : m-2
        A[iz, iz-1] = -a2/dz
        A[iz, iz  ] = -a1/dz
        A[iz, iz+1] =  a1/dz
        A[iz, iz+2] =  a2/dz
    end
    A[m-1, m-2] = -a2/dz;
    A[m-1, m-1] = -a1/dz; A[m-1, m] = a1/dz;
    A[m  , m-1] = -a2/dz; A[m  , m] = -a1/dz;
    In = speye(n)
    M = kron(In, A)
    M = spdiagm(A3) * M
    return M
end

function MTxxxBTxxx(dTxxx::SparseMatrixCSC{Float64,Int64}, dt::Float64)
    (m,n) = size(dTxxx)
    Mdiag = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a3 = 1/dt + dTxxx[iz,ix]/2
            tmp= 1/(dt*a3) - dTxxx[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    M = Mdiag
    return M
end

function MTxxxBVx(dTxxx::SparseMatrixCSC{Float64,Int64}, lambda::Array{Float64,2}, mu::Array{Float64,2}, dt::Float64, dx::Float64)
    a1 = 9/8
    a2 = -1/24
    (m,n) = size(dTxxx)
    A3 = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            A3[(ix-1)*m+iz] = (lambda[iz,ix]+2*mu[iz,ix]) / (1/dt+dTxxx[iz,ix]/2)
        end
    end
    A = spzeros(n, n)
    A[1, 1] = a1/dx; A[1, 2] = a2/dx;
    A[2, 1] =-a1/dx; A[2, 2] = a1/dx; A[2, 3] = a2/dx;
    for ix = 3:n-1
        A[ix, ix-2] = -a2/dx
        A[ix, ix-1] = -a1/dx
        A[ix, ix  ] =  a1/dx
        A[ix, ix+1] =  a2/dx
    end
    A[n, n-2] = -a2/dx;
    A[n, n-1] = -a1/dx; A[n, n] = a1/dx;
    Im = speye(m)
    M = kron(A, Im)
    M = spdiagm(A3) * M
    return M
end

function MTxxzBTxxz(dTxxz::SparseMatrixCSC{Float64,Int64}, dt::Float64)
    (m,n) = size(dTxxz)
    Mdiag = zeros(m*n)
    for ix = 1:n
        for iz = 1:m
            a3 = 1/dt + dTxxz[iz,ix]/2
            tmp= 1/(dt*a3) - dTxxz[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    M = Mdiag
    return M
end

function MTxxzBVz(dTxxz::SparseMatrixCSC{Float64,Int64}, lambda::Array{Float64,2}, dt::Float64, dz::Float64)
    a1 = 9/8
    a2 = -1/24
    (m,n) = size(dTxxz)
    A3 = zeros(m*n)
    for ix = 1:n
        for iz = 1:m
            A3[(ix-1)*m+iz] = lambda[iz,ix] / (1/dt+dTxxz[iz,ix]/2)
        end
    end
    A = spzeros(m, m)
    A[1, 1] = a1/dz; A[1, 2] = a2/dz;
    A[2, 1] =-a1/dz; A[2, 2] = a1/dz; A[2, 3] = a2/dz;
    for iz = 3:m-1
        A[iz, iz-2] = -a2/dz
        A[iz, iz-1] = -a1/dz
        A[iz, iz  ] =  a1/dz
        A[iz, iz+1] =  a2/dz
    end
    A[m, m-2] = -a2/dz;
    A[m, m-1] = -a1/dz; A[m, m] = a1/dz;
    In = speye(n)
    M = kron(In, A)
    M = spdiagm(A3) * M
    return M
end

function MTzzxBTzzx(dTzzx::SparseMatrixCSC{Float64,Int64}, dt::Float64)
    (m,n) = size(dTzzx)
    Mdiag = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a3 = 1/dt + dTzzx[iz,ix]/2
            tmp= 1/(dt*a3) - dTzzx[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    M = Mdiag
    return M
end

function MTzzxBVx(dTzzx::SparseMatrixCSC{Float64,Int64}, lambda::Array{Float64,2}, dt::Float64, dx::Float64)
    a1 = 9/8
    a2 =-1/24
    (m, n) = size(dTzzx)
    A3 = zeros(m*n)
    for ix = 1:n
        for iz = 1:m
            A3[(ix-1)*m+iz] = lambda[iz,ix] / (1/dt+dTzzx[iz,ix]/2)
        end
    end
    A = spzeros(n, n)
    A[1, 1] = a1/dx; A[1, 2] = a2/dx;
    A[2, 1] =-a1/dx; A[2, 2] = a1/dx; A[2, 3] = a2/dx;
    for ix = 3:n-1
        A[ix, ix-2] = -a2/dx
        A[ix, ix-1] = -a1/dx
        A[ix, ix  ] =  a1/dx
        A[ix, ix+1] =  a2/dx
    end
    A[n, n-2] = -a2/dx;
    A[n, n-1] = -a1/dx; A[n, n] = a1/dx;
    Im = speye(m)
    M = kron(A, Im)
    M = spdiagm(A3) * M
    return M
end

function MTzzzBTzzz(dTzzz::SparseMatrixCSC{Float64,Int64}, dt::Float64)
    (m,n) = size(dTzzz)
    Mdiag = zeros(m*n)
    for ix = 1:n
        for iz = 1:m
            a3 = 1/dt + dTzzz[iz,ix]/2
            tmp= 1/(dt*a3) - dTzzz[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    M = Mdiag
    return M
end

function MTzzzBVz(dTzzz::SparseMatrixCSC{Float64,Int64}, lambda::Array{Float64,2}, mu::Array{Float64,2}, dt::Float64, dz::Float64)
    a1 = 9/8
    a2 = -1/24
    (m,n) = size(dTzzz)
    A3 = zeros(m*n)
    for ix = 1:n
        for iz = 1:m
            A3[(ix-1)*m+iz] = (lambda[iz,ix]+2*mu[iz,ix]) / (1/dt+dTzzz[iz,ix]/2)
        end
    end
    A = spzeros(m, m)
    A[1, 1] = a1/dz; A[1, 2] = a2/dz;
    A[2, 1] =-a1/dz; A[2, 2] = a1/dz; A[2, 3] = a2/dz;
    for iz = 3 : m-1
        A[iz, iz-2] = -a2/dz
        A[iz, iz-1] = -a1/dz
        A[iz, iz  ] =  a1/dz
        A[iz, iz+1] =  a2/dz
    end
    A[m, m-2] = -a2/dz;
    A[m, m-1] = -a1/dz; A[m, m] = a1/dz;
    In = speye(n)
    M = kron(In, A)
    M = spdiagm(A3) * M
    return M
end

function MTxzxBTxzx(dTxzx::SparseMatrixCSC{Float64,Int64}, dt::Float64)
    (m,n) = size(dTxzx)
    Mdiag = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a3 = 1/dt + dTxzx[iz,ix]/2
            tmp= 1/(dt*a3) - dTxzx[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    M = Mdiag
    return M
end

function MTxzxBVz(dTxzx::SparseMatrixCSC{Float64,Int64}, mu::Array{Float64,2}, dt::Float64, dx::Float64)
    a1 = 9/8
    a2 =-1/24
    (m, n) = size(dTxzx)
    A3 = zeros(m*n)
    for ix = 1 : n-1
        for iz = 1 : m-1
            A3[(ix-1)*m+iz] = (mu[iz,ix]+mu[iz+1,ix]+mu[iz,ix+1]+mu[iz+1,ix+1])/4 / (1/dt+dTxzx[iz,ix]/2)
        end
    end
    for iz = 1 : m-1
        A3[(n-1)*m+iz] = (mu[iz,n]+mu[iz+1,n])/2 / (1/dt+dTxzx[iz,n]/2)
    end
    for ix = 1 : n-1
        A3[(ix-1)*m+m] = (mu[m,ix]+mu[m,ix+1])/2 / (1/dt+dTxzx[m,ix]/2)
    end
    A3[end] = mu[m,n] / (1/dt+dTxzx[m,n]/2)
    A = spzeros(n, n)
    A[1,1] =-a1/dx; A[1,2] = a1/dx; A[1,3] = a2/dx;
    for ix = 2 : n-2
        A[ix, ix-1] = -a2/dx
        A[ix, ix  ] = -a1/dx
        A[ix, ix+1] =  a1/dx
        A[ix, ix+2] =  a2/dx
    end
    A[n-1,n-2] =-a2/dx;
    A[n-1,n-1] =-a1/dx; A[n-1,n] = a1/dx;
    A[n  ,n-1] =-a2/dx; A[n  ,n] =-a1/dx;
    Im = speye(m)
    M = kron(A, Im)
    M = spdiagm(A3) * M
    return M
end

function MTxzzBTxzz(dTxzz::SparseMatrixCSC{Float64,Int64}, dt::Float64)
    (m, n) = size(dTxzz)
    Mdiag  = zeros(m*n)
    for ix = 1:n
        for iz = 1:m
            a3 = 1/dt + dTxzz[iz,ix]/2
            tmp= 1/(dt*a3) - dTxzz[iz,ix]/(2*a3)
            Mdiag[(ix-1)*m+iz] = tmp
        end
    end
    M = Mdiag
    return M
end

function MTxzzBVx(dTxzz::SparseMatrixCSC{Float64,Int64}, mu::Array{Float64,2}, dt::Float64, dz::Float64)
    a1 = 9/8
    a2 =-1/24
    (m,n) = size(dTxzz)
    A3 = zeros(m*n)
    for ix = 1 : n-1
        for iz = 1 : m-1
            A3[(ix-1)*m+iz] = (mu[iz,ix]+mu[iz+1,ix]+mu[iz,ix+1]+mu[iz+1,ix+1])/4 / (1/dt+dTxzz[iz,ix]/2)
        end
    end
    for iz = 1 : m-1
        A3[(n-1)*m+iz] = (mu[iz,n]+mu[iz+1,n])/2 / (1/dt+dTxzz[iz,n]/2)
    end
    for ix = 1 : n-1
        A3[(ix-1)*m+m] = (mu[m,ix]+mu[m,ix+1])/2 / (1/dt+dTxzz[m,ix]/2)
    end
    A3[end] = mu[m,n] / (1/dt+dTxzz[m,n]/2)
    A = spzeros(m, m)
    A[1,1] =-a1/dz; A[1,2] = a1/dz; A[1,3] = a2/dz;
    for iz = 2:m-2
        A[iz, iz-1] = -a2/dz
        A[iz, iz  ] = -a1/dz
        A[iz, iz+1] =  a1/dz
        A[iz, iz+2] =  a2/dz
    end
    A[m-1,m-2] =-a2/dz;
    A[m-1,m-1] =-a1/dz; A[m-1,m] = a1/dz;
    A[m  ,m-1] =-a2/dz; A[m  ,m] =-a1/dz;
    In = speye(n)
    M = kron(In, A)
    M = spdiagm(A3) * M
    return M
end

function CreateFidMtx(nz::Int64, nx::Int64, ext::Int64, iflag::Int64, vp::Array{Float64,2}, vs::Array{Float64,2}, rho::Array{Float64,2}, dz::Float64, dx::Float64, dt::Float64, f0::Float64)
    vmax = maximum(vp); vmin = minimum(vs);
    DispStable!(vmax, vmin, f0, dz, dx, dt)
    (lambda, mu) = vel2lm(vp, vs, rho)
    (lambda, mu, rho) = modelPadding(lambda, mu, rho, ext, iflag)
    dpCoef = DampBound(nz, nx, vmax, ext, rho, iflag, dz)
    MVxxBVxx1 = MVxxBVxx(dpCoef.dVxx, rho, dt)
    MVxxBTxx1 = MVxxBTxx(dpCoef.dVxx, rho, dt, dx)
    MVxzBVxz1 = MVxzBVxz(dpCoef.dVxz, rho, dt)
    MVxzBTxz1 = MVxzBTxz(dpCoef.dVxz, rho, dt, dz)
    MVzxBVzx1 = MVzxBVzx(dpCoef.dVzx, rho, dt)
    MVzxBTxz1 = MVzxBTxz(dpCoef.dVzx, rho, dt, dx)
    MVzzBVzz1 = MVzzBVzz(dpCoef.dVzz, rho, dt)
    MVzzBTzz1 = MVzzBTzz(dpCoef.dVzz, rho, dt, dz)
    MTxxxBVx1   =   MTxxxBVx(dpCoef.dTxxx, lambda, mu, dt, dx)
    MTxxxBTxxx1 = MTxxxBTxxx(dpCoef.dTxxx, dt)
    MTxxzBVz1   =   MTxxzBVz(dpCoef.dTxxz, lambda, dt, dz)
    MTxxzBTxxz1 = MTxxzBTxxz(dpCoef.dTxxz, dt)
    MTzzxBVx1   =   MTzzxBVx(dpCoef.dTzzx, lambda, dt, dx)
    MTzzxBTzzx1 = MTzzxBTzzx(dpCoef.dTzzx, dt)
    MTzzzBVz1   =   MTzzzBVz(dpCoef.dTzzz, lambda, mu, dt, dz)
    MTzzzBTzzz1 = MTzzzBTzzz(dpCoef.dTzzz, dt)
    MTxzxBVz1   =   MTxzxBVz(dpCoef.dTxzx, mu, dt, dx)
    MTxzxBTxzx1 = MTxzxBTxzx(dpCoef.dTxzx, dt)
    MTxzzBVx1   =   MTxzzBVx(dpCoef.dTxzz, mu, dt, dz)
    MTxzzBTxzz1 = MTxzzBTxzz(dpCoef.dTxzz, dt)
    fidMtx = FidMtx(MVxxBVxx1, MVxxBTxx1,   MVxzBVxz1, MVxzBTxz1,
                    MVzxBVzx1, MVzxBTxz1,   MVzzBVzz1, MVzzBTzz1,
                    MTxxxBVx1, MTxxxBTxxx1, MTxxzBVz1, MTxxzBTxxz1,
                    MTzzxBVx1, MTzzxBTzzx1, MTzzzBVz1, MTzzzBTzzz1,
                    MTxzxBVz1, MTxzxBTxzx1, MTxzzBVx1, MTxzzBTxzz1)
    return fidMtx
end


using Images, CausticsEngineering

image = Images.load("./examples/cat_posing.jpg") # Check current working directory with pwd()
mesh, image_caustics = engineer_caustics(image);
nothing;

image_caustics'


using Plots, SNN

SNN.IF_curve(SNN.IF, 0:0.1:20) |> display
SNN.IF_curve(SNN.IZ, 0:0.1:20) |> display
SNN.IF_curve(SNN.HH, 1e-5*(0:0.01:4)) |> display


type XGBoost <: DecisionTree
    n_estimators::Int64
    learning_rate::Float64
    min_samples_split::Int64
    min_impurity::Float64
    max_depth::Int64
    init_estimate::Union{String, Array}
    loss::LogisticLoss
    trees::Vector{XGBoostRegressionTree}
end

function XGBoost(;
                 n_estimators::Int64 = 2,
                 learning_rate::Float64 = 0.001,
                 min_samples_split::Int64 = 2,
                 min_impurity::Float64 = 0.00001,
                 max_depth::Int64 = 2,
                 init_estimate = "nothing",
                 loss::LogisticLoss = LogisticLoss())
    trees = []
    single_boost_tree = XGBoostRegressionTree(min_samples_split = min_samples_split,
            min_gain = min_impurity, max_depth = max_depth, loss = loss)
    for i in 1:n_estimators
        push!(trees, single_boost_tree)
    end
    return XGBoost(n_estimators, learning_rate,
                   min_samples_split, min_impurity,
                   max_depth, init_estimate, loss, trees)
end

function train!(model::XGBoost, X::Matrix, y::Vector)
    y = one_hot(y)
    n_sample, n_feature = size(X)
    model.init_estimate = mean(y, 1)
    y_pred = repmat(model.init_estimate, size(y,1),1)
    for i in 1:model.n_estimators
        y_and_pred = hcat(y, y_pred)
        print("$(y_and_pred)")
        train!(model.trees[i], X, y_and_pred)
        update_pred = predict(model.trees[i],X)
        y_pred += model.learning_rate * update_pred
        progress = 1/model.n_estimators * i * 100
        #println("Progress: $(progress) %")
    end
end

function predict(model::XGBoost,
                 x::Matrix)
    n_sample = size(x)
    flag = 1
    y_pred = 0
    for tree in model.trees
        update_pred = predict(tree,x)
        update = model.learning_rate * update_pred
        if flag == 1
            y_pred = update
            flag = 0
        else
            y_pred += update
        end
    end
    res = zeros(size(y_pred,1))
    for i in 1:size(y_pred,1)
        exp_pred = exp(y_pred[i,:])
        y_pred[i,:] = exp_pred / sum(exp_pred)
        print("$(y_pred[i,:])")
        res[i] = indmax(y_pred[i,:])
    end
    return res-1
end

function test_xgboost()
    X_train, X_test, y_train, y_test = make_iris()
    model = XGBoost()
    train!(model,X_train, y_train)
    predictions = predict(model,X_test)
    print("$(predictions)")
    print("classification accuracy", accuracy(y_test, predictions))
end


module Poisson_examples
using FinEtools
using FinEtoolsHeatDiff

function Poisson_FE_H20_example()
    println("""

    Heat conduction example described by Amuthan A. Ramabathiran
    http://www.codeproject.com/Articles/579983/Finite-Element-programming-in-Julia:
    Unit cube, with known temperature distribution along the boundary,
    and uniform heat generation rate inside.  Mesh of regular quadratic HEXAHEDRA,
    in a grid of 30 x 30 x 30 edges (100079 degrees of freedom).
    Version: 06/03/2017
    """
    )
    t0 = time()

    A = 1.0 # dimension of the domain (length of the side of the square)
    thermal_conductivity = [i==j ? one(FFlt) : zero(FFlt) for i=1:3, j=1:3]; # conductivity matrix
    Q = -6.0; # internal heat generation rate
    function getsource!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        forceout[1] = Q; #heat source
    end
    tempf(x) = (1.0 .+ x[:,1].^2 + 2.0 .* x[:,2].^2);#the exact distribution of temperature
    N = 30;# number of subdivisions along the sides of the square domain

    println("Mesh generation")
    @time fens,fes = H20block(A, A, A, N, N, N)

    geom = NodalField(fens.xyz)
    Temp = NodalField(zeros(size(fens.xyz,1),1))

    println("Searching nodes  for BC")
    Tolerance = 1.0/N/100.0
    @time l1 = selectnode(fens; box=[0. 0. 0. A 0. A], inflate = Tolerance)
    @time l2 = selectnode(fens; box=[A A 0. A 0. A], inflate = Tolerance)
    @time l3 = selectnode(fens; box=[0. A 0. 0. 0. A], inflate = Tolerance)
    @time l4 = selectnode(fens; box=[0. A A A 0. A], inflate = Tolerance)
    @time l5 = selectnode(fens; box=[0. A 0. A 0. 0.], inflate = Tolerance)
    @time l6 = selectnode(fens; box=[0. A 0. A A A], inflate = Tolerance)
    List = vcat(l1, l2, l3, l4, l5, l6)
    @time setebc!(Temp, List, true, 1, tempf(geom.values[List,:])[:])
    @time applyebc!(Temp)
    @time numberdofs!(Temp)

    println( "Number of free degrees of freedom: $(Temp.nfreedofs)")
    t1 = time()

    material = MatHeatDiff(thermal_conductivity)

    femm = FEMMHeatDiff(IntegDomain(fes, GaussRule(3, 3)), material)

    println("Conductivity")
    @time K = conductivity(femm, geom, Temp)
    println("Nonzero EBC")
    @time F2 = nzebcloadsconductivity(femm, geom, Temp);
    println("Internal heat generation")
    # fi = ForceIntensity(FFlt, getsource!);# alternative  specification
    fi = ForceIntensity(FFlt[Q]);
    @time F1 = distribloads(femm, geom, Temp, fi, 3);

    println("Solution of the system")
    @time U = K\(F1+F2)
    scattersysvec!(Temp,U[:])

    println("Total time elapsed = $(time() - t0) [s]")
    println("Solution time elapsed = $(time() - t1) [s]")

    Error= 0.0
    for k=1:size(fens.xyz,1)
        Error = Error + abs.(Temp.values[k,1] - tempf(reshape(fens.xyz[k,:], (1,3)))[1])
    end
    println("Error =$Error")

    # File =  "a.vtk"
    # MeshExportModule.vtkexportmesh (File, fes.conn, [geom.values Temp.values], MeshExportModule.T3; scalars=Temp.values, scalars_name ="Temperature")

    true

end # Poisson_FE_H20_example

function Poisson_FE_T10_example()
    println("""

    Heat conduction example described by Amuthan A. Ramabathiran
    http://www.codeproject.com/Articles/579983/Finite-Element-programming-in-Julia:
    Unit cube, with known temperature distribution along the boundary,
    and uniform heat generation rate inside.  Mesh of regular quadratic HEXAHEDRA,
    in a grid of 30 x 30 x 30 edges (100079 degrees of freedom).
    Version: 06/03/2017
    """
    )
    t0 = time()

    A = 1.0 # dimension of the domain (length of the side of the square)
    thermal_conductivity = [i==j ? one(FFlt) : zero(FFlt) for i=1:3, j=1:3]; # conductivity matrix
    Q = -6.0; # internal heat generation rate
    function getsource!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        forceout[1] = Q; #heat source
    end
    tempf(x) = (1.0 .+ x[:,1].^2 + 2.0 .* x[:,2].^2);#the exact distribution of temperature
    N = 30;# number of subdivisions along the sides of the square domain

    println("Mesh generation")
    @time fens,fes = H20block(A, A, A, N, N, N)

    geom = NodalField(fens.xyz)
    Temp = NodalField(zeros(size(fens.xyz,1),1))

    println("Searching nodes  for BC")
    Tolerance = 1.0/N/100.0
    @time l1 = selectnode(fens; box=[0. 0. 0. A 0. A], inflate = Tolerance)
    @time l2 = selectnode(fens; box=[A A 0. A 0. A], inflate = Tolerance)
    @time l3 = selectnode(fens; box=[0. A 0. 0. 0. A], inflate = Tolerance)
    @time l4 = selectnode(fens; box=[0. A A A 0. A], inflate = Tolerance)
    @time l5 = selectnode(fens; box=[0. A 0. A 0. 0.], inflate = Tolerance)
    @time l6 = selectnode(fens; box=[0. A 0. A A A], inflate = Tolerance)
    List = vcat(l1, l2, l3, l4, l5, l6)
    @time setebc!(Temp, List, true, 1, tempf(geom.values[List,:])[:])
    @time applyebc!(Temp)
    @time numberdofs!(Temp)

    println( "Number of free degrees of freedom: $(Temp.nfreedofs)")
    t1 = time()

    material = MatHeatDiff(thermal_conductivity)

    femm = FEMMHeatDiff(IntegDomain(fes, GaussRule(3, 3)), material)

    println("Conductivity")
    @time K = conductivity(femm, geom, Temp)
    println("Nonzero EBC")
    @time F2 = nzebcloadsconductivity(femm, geom, Temp);
    println("Internal heat generation")
    # fi = ForceIntensity(FFlt, getsource!);# alternative  specification
    fi = ForceIntensity(FFlt[Q]);
    @time F1 = distribloads(femm, geom, Temp, fi, 3);

    println("Solution of the system")
    @time U = K\(F1+F2)
    scattersysvec!(Temp,U[:])

    println("Total time elapsed = $(time() - t0) [s]")
    println("Solution time elapsed = $(time() - t1) [s]")

    Error= 0.0
    for k=1:size(fens.xyz,1)
        Error = Error + abs.(Temp.values[k,1] - tempf(reshape(fens.xyz[k,:], (1,3)))[1])
    end
    println("Error =$Error")

    # File =  "a.vtk"
    # MeshExportModule.vtkexportmesh (File, fes.conn, [geom.values Temp.values], MeshExportModule.T3; scalars=Temp.values, scalars_name ="Temperature")

    true

end # Poisson_FE_T10_example

function Poisson_FE_T4_example()
    t0 = time()

    A = 1.0 # dimension of the domain (length of the side of the square)
    thermal_conductivity = [i==j ? one(FFlt) : zero(FFlt) for i=1:3, j=1:3]; # conductivity matrix
    Q = -6.0; # internal heat generation rate
    function getsource!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        forceout[1] = Q; #heat source
    end
    tempf(x) = (1.0 .+ x[:,1].^2 + 2.0 .* x[:,2].^2);#the exact distribution of temperature
    N = 30;# number of subdivisions along the sides of the square domain

    println("Mesh generation")
    @time fens,fes = T4block(A, A, A, N, N, N)

    println("""
    Heat conduction example described by Amuthan A. Ramabathiran
    http://www.codeproject.com/Articles/579983/Finite-Element-programming-in-Julia:
    Unit cube, with known temperature distribution along the boundary,
    and uniform heat generation rate inside.  Mesh of regular quadratic TETRAHEDRA,
    in a grid of $(N) x $(N) x $(N) edges ($(count(fens)) degrees of freedom).
    Version: 07/03/2017
    """
    )

    geom = NodalField(fens.xyz)
    Temp = NodalField(zeros(size(fens.xyz,1),1))

    println("Searching nodes  for BC")
    Tolerance = 1.0/N/100.0
    @time l1 = selectnode(fens; box=[0. 0. 0. A 0. A], inflate = Tolerance)
    @time l2 = selectnode(fens; box=[A A 0. A 0. A], inflate = Tolerance)
    @time l3 = selectnode(fens; box=[0. A 0. 0. 0. A], inflate = Tolerance)
    @time l4 = selectnode(fens; box=[0. A A A 0. A], inflate = Tolerance)
    @time l5 = selectnode(fens; box=[0. A 0. A 0. 0.], inflate = Tolerance)
    @time l6 = selectnode(fens; box=[0. A 0. A A A], inflate = Tolerance)
    List = vcat(l1, l2, l3, l4, l5, l6)
    @time setebc!(Temp, List, true, 1, tempf(geom.values[List,:])[:])
    @time applyebc!(Temp)
    @time numberdofs!(Temp)

    println( "Number of free degrees of freedom: $(Temp.nfreedofs)")
    t1 = time()

    material = MatHeatDiff(thermal_conductivity)

    femm = FEMMHeatDiff(IntegDomain(fes, TetRule(1)), material)

    println("Conductivity")
    @time K = conductivity(femm, geom, Temp)
    println("Nonzero EBC")
    @time F2 = nzebcloadsconductivity(femm, geom, Temp);
    println("Internal heat generation")
    # fi = ForceIntensity(FFlt, getsource!);# alternative  specification
    fi = ForceIntensity(FFlt[Q]);
    @time F1 = distribloads(femm, geom, Temp, fi, 3);

    println("Solution of the system")
    @time U = K\(F1+F2)
    scattersysvec!(Temp,U[:])

    println("Total time elapsed = $(time() - t0) [s]")
    println("Solution time elapsed = $(time() - t1) [s]")

    Error= 0.0
    for k=1:size(fens.xyz,1)
        Error = Error + abs.(Temp.values[k,1] - tempf(reshape(fens.xyz[k,:], (1,3)))[1])
    end
    println("Error =$Error")

    # File =  "a.vtk"
    # MeshExportModule.vtkexportmesh (File, fes.conn, [geom.values Temp.values], MeshExportModule.T3; scalars=Temp.values, scalars_name ="Temperature")

    true

end # Poisson_FE_T4_example

function allrun()
    println("#####################################################")
    println("# Poisson_FE_H20_example ")
    Poisson_FE_H20_example()
    println("#####################################################")
    println("# Poisson_FE_T10_example ")
    Poisson_FE_T10_example()
    println("#####################################################")
    println("# Poisson_FE_T4_example ")
    Poisson_FE_T4_example()
end # function allrun

end # module Poisson_examples


#Written By Johnathan Bizzano
export MultiVarPolynomial

struct MultiVarPolynomial{T, N} <: AbstractMathFun{T, N}
    summation::RegressableSummation{T, N}

    function MultiVarPolynomial{T, N}(powers::AbstractArray) where {T, N}
        new{T, N}(Functions.RegressableSummation{T, N}(GeneratePowers(powers, T, N)))
    end

    function GeneratePowers(iter::AbstractArray, T::Type, Var_Len)
        iter = adjoint(iter)
        arr = TermsArray{T, Var_Len}(undef, size(iter, 1))
        i = 1
        vars = [Functions.VarFunction{T}(v) for v in 1:Var_Len]
        for row in eachrow(iter)
            powCollection = MFCollection(undef, Var_Len)
            for col in 1:Var_Len
                powCollection[col] = Functions.PowFunction{T, Var_Len}(vars[col], row[col])
            end
            arr[i] = Functions.ProductFunction{T, Var_Len}(powCollection)
            i += 1
        end
        arr
    end

    (f::MultiVarPolynomial{T, N})(args::AbstractArray{T}; Data_Type::Type=T) where {T, N} = f.summation(args, Data_Type = Data_Type)
    Functions.train!(f::MultiVarPolynomial{T, N}, data::AbstractDataSet; Data_Type::Type=T) where {T, N} = Functions.train!(f.summation, data, Data_Type = Data_Type)
    Base.sizeof(x::MultiVarPolynomial) = sizeof(x.summation)
    Base.string(x::MultiVarPolynomial) = string(x.summation)
    Functions.simplify(x::MultiVarPolynomial) = Functions.simplify(x.summation)
    Functions.trainer(f::MultiVarPolynomial{T, N}; precision=10) where {T, N} = RegressableSummationTrainer{T, N}(f.summation, precision = precision)
end


function rot_elem(N::Int, v::Vector{Int})
    elem = 0.0

    for i in 1:N
        l = v[2i-1]
        elem += l*(l+1)
    end

    elem
end

function sparse_rot(block::Block{N}) where {N}
    rows = Int[]
    elems = Float64[]

    for row in 1:block.size
        v = block.vectors[row]
        elem = rot_elem(N, v)

        push!(rows, row)
        push!(elems, elem)
    end

    sparse(rows, rows, elems, block.size, block.size)
end

"""
    sparse_rot(basis::SingleBlockBasis)

Compute a sparse representation of the dimensionless rotational energy matrix.
"""
sparse_rot(basis::SingleBlockBasis) = sparse_rot(basis.block)

"""
    diag_rot(basis::SingleBlockBasis)

Compute a diagonal representation of the dimensionless rotational energy
matrix.
"""
diag_rot(basis::SingleBlockBasis) = diag(sparse_rot(basis.block))

"""
    diag_rot(basis::MultiBlockBasis)

Compute a diagonal representation of the dimensionless rotational energy
matrix.
"""
function diag_rot(basis::MultiBlockBasis)
    result = Dict{BlockLabel,Vector{Float64}}()

    for (label, block) in basis.blocks
        result[label] = diag(sparse_rot(block))
    end

    result
end

function C(l::Int, lp::Int, m::Int)
    if lp == l-1
        sqrt((l-m)*(l+m)/((2l-1)*(2l+1)))
    elseif lp == l+1
        sqrt((lp-m)*(lp+m)/((2lp-1)*(2lp+1)))
    end
end

function Sp(l::Int, lp::Int, m::Int)
    if lp == l-1
        -sqrt((l-m-1)*(l-m)/((2l-1)*(2l+1)))
    elseif lp == l+1
        sqrt((lp+m)*(lp+m+1)/((2lp-1)*(2lp+1)))
    end
end

function Sm(l::Int, lp::Int, m::Int)
    if lp == l-1
        sqrt((l+m-1)*(l+m)/((2l-1)*(2l+1)))
    elseif lp == l+1
        -sqrt((lp-m)*(lp-m+1)/((2lp-1)*(2lp+1)))
    end
end

function pot_elem(l1::Int, m1::Int, l2::Int, m2::Int, l1p::Int, m1p::Int, l2p::Int, m2p::Int)
    if m1p == m1-1 && m2p == m2+1
        Sm(l1, l1p, m1) * Sp(l2, l2p, m2)
    elseif m1p == m1+1 && m2p == m2-1
        Sp(l1, l1p, m1) * Sm(l2, l2p, m2)
    elseif m1p == m1 && m2p == m2
        -4.0 * C(l1, l1p, m1) * C(l2, l2p, m2)
    end
end

function pot_rows(f::Function, N::Int, l_max::Int, l_total_max::Int, v::Vector{Int}, w::Vector{Int})
    l_total = 0
    for i in 1:N
        l_total += v[2i-1]
    end

    # All rotor pairs.
    for i in 1:(N-1)
        for j in (i+1):N
            # All potentially allowed transitions.
            for delta_li in [-1, 1]
                for delta_lj in [-1, 1]
                    for delta_m in -1:1
                        w .= v
                        w[2i-1] += delta_li
                        w[2i] += delta_m
                        w[2j-1] += delta_lj
                        w[2j] -= delta_m

                        # Bounds on l.
                        0 <= w[2i-1] <= l_max || continue
                        0 <= w[2j-1] <= l_max || continue
                        l_total + delta_li + delta_lj <= l_total_max || continue
                        # Bounds on m.
                        -w[2i-1] <= w[2i] <= w[2i-1] || continue
                        -w[2j-1] <= w[2j] <= w[2j-1] || continue

                        unscaled = pot_elem(v[2i-1], v[2i], v[2j-1], v[2j], w[2i-1], w[2i], w[2j-1], w[2j])
                        elem = 0.5 * unscaled / (j-i)^3

                        f(elem)
                    end
                end
            end
        end
    end

    nothing
end

function sparse_pot(basis::AbstractBasis{N}, block::Block{N}) where {N}
    rows = Int[]
    cols = Int[]
    elems = Float64[]

    w = Array{Int}(undef, 2N)

    for col in 1:block.size
        v = block.vectors[col]

        pot_rows(N, basis.l_max, basis.l_total_max, v, w) do elem
            row = block.lookup[w]

            push!(rows, row)
            push!(cols, col)
            push!(elems, elem)
        end
    end

    sparse(rows, cols, elems, block.size, block.size)
end

"""
    sparse_pot(basis::SingleBlockBasis)

Compute a sparse representation of the dimensionless potential energy matrix
for dipole-dipole interactions between all rotor pairs.
"""
sparse_pot(basis::SingleBlockBasis) = sparse_pot(basis, basis.block)

"""
    dense_pot(basis::SingleBlockBasis)

Compute a dense representation of the dimensionless potential energy matrix for
dipole-dipole interactions between all rotor pairs.
"""
dense_pot(basis::SingleBlockBasis) = Matrix(sparse_pot(basis, basis.block))

"""
    dense_pot(basis::MultiBlockBasis)

Compute a dense representation of the dimensionless potential energy matrix for
dipole-dipole interactions between all rotor pairs.
"""
function dense_pot(basis::MultiBlockBasis)
    result = Dict{BlockLabel,Matrix{Float64}}()

    for (label, block) in basis.blocks
        result[label] = Matrix(sparse_pot(basis, block))
    end

    result
end

"""
    scale_dense_pot(pot::Matrix{Float64}, k::Float64)

Scale the matrix elements of `pot` in place by a factor of `k`.
"""
function scale_dense_pot(pot::Matrix{Float64}, k::Float64)
    pot .*= k
end

"""
    scale_dense_pot(pot::Dict{BlockLabel,Matrix{Float64}}, k::Float64)

Scale the matrix elements of `pot` in place by a factor of `k`.
"""
function scale_dense_pot(pot::Dict{BlockLabel,Matrix{Float64}}, k::Float64)
    for block in values(pot)
        block .*= k
    end
end

"""
    dense_x(basis::SingleBlockBasis)

Compute a dense representation of the single-rotor x operator.
"""
function dense_x(basis::SingleBlockBasis)
    l_max = basis.l_max

    result = zeros(Float64, (l_max+1)^2, (l_max+1)^2)

    for l in 0:l_max
        for m in (-l):l
            col = l*(l+1) + m + 1

            for lp in 0:l_max
                for mp in (-lp):lp
                    row = lp*(lp+1) + mp + 1

                    if lp == l+1 && mp == m+1
                        result[row, col] -= 0.5 * sqrt((l+m+1)*(l+m+2)/((2l+1)*(2l+3)))
                    elseif lp == l-1 && mp == m+1
                        result[row, col] += 0.5 * sqrt((l-m-1)*(l-m)/((2l-1)*(2l+1)))
                    elseif lp == l+1 && mp == m-1
                        result[row, col] += 0.5 * sqrt((l-m+1)*(l-m+2)/((2l+1)*(2l+3)))
                    elseif lp == l-1 && mp == m-1
                        result[row, col] -= 0.5 * sqrt((l+m-1)*(l+m)/((2l-1)*(2l+1)))
                    end
                end
            end
        end
    end

    result
end

"""
    dense_iy(basis::SingleBlockBasis)

Compute a dense representation of the single-rotor iy operator.
"""
function dense_iy(basis::SingleBlockBasis)
    l_max = basis.l_max

    result = zeros(Float64, (l_max+1)^2, (l_max+1)^2)

    for l in 0:l_max
        for m in (-l):l
            col = l*(l+1) + m + 1

            for lp in 0:l_max
                for mp in (-lp):lp
                    row = lp*(lp+1) + mp + 1

                    if lp == l+1 && mp == m+1
                        result[row, col] -= 0.5 * sqrt((l+m+1)*(l+m+2)/((2l+1)*(2l+3)))
                    elseif lp == l-1 && mp == m+1
                        result[row, col] += 0.5 * sqrt((l-m-1)*(l-m)/((2l-1)*(2l+1)))
                    elseif lp == l+1 && mp == m-1
                        result[row, col] -= 0.5 * sqrt((l-m+1)*(l-m+2)/((2l+1)*(2l+3)))
                    elseif lp == l-1 && mp == m-1
                        result[row, col] += 0.5 * sqrt((l+m-1)*(l+m)/((2l-1)*(2l+1)))
                    end
                end
            end
        end
    end

    result
end

"""
    dense_z(basis::SingleBlockBasis)

Compute a dense representation of the single-rotor z operator.
"""
function dense_z(basis::SingleBlockBasis)
    l_max = basis.l_max

    result = zeros(Float64, (l_max+1)^2, (l_max+1)^2)

    for l in 0:l_max
        for m in (-l):l
            col = l*(l+1) + m + 1

            for lp in 0:l_max
                for mp in (-lp):lp
                    row = lp*(lp+1) + mp + 1

                    if lp == l+1 && mp == m
                        result[row, col] += sqrt((l-m+1)*(l+m+1)/((2l+1)*(2l+3)))
                    elseif lp == l-1 && mp == m
                        result[row, col] += sqrt((l-m)*(l+m)/((2l-1)*(2l+1)))
                    end
                end
            end
        end
    end

    result
end


# Carries out a Kepler step for bodies i & j
function keplerij!(m::Array{Float64,1},x::Array{Float64,2},v::Array{Float64,2},i::Int64,j::Int64,h::Float64,jac_ij::Array{Float64,2})
# The state vector has: 1 time; 2-4 position; 5-7 velocity; 8 r0; 9 dr0dt; 10 beta; 11 s; 12 ds
# Initial state:
state0 = zeros(Float64,12)
# Final state (after a step):
state = zeros(Float64,12)
delx = zeros(Float64,NDIM)
delv = zeros(Float64,NDIM)
# jac_ij should be the Jacobian for going from (x_{0,i},v_{0,i},m_i) &  (x_{0,j},v_{0,j},m_j)
# to  (x_i,v_i,m_i) &  (x_j,v_j,m_j), a 14x14 matrix for the 3-dimensional case.
# Fill with zeros for now:
fill!(jac_ij,0.0)
for k=1:NDIM
  state0[1+k     ] = x[k,i] - x[k,j]
  state0[1+k+NDIM] = v[k,i] - v[k,j]
end
gm = GNEWT*(m[i]+m[j])
# The following jacobian is just computed for the Keplerian coordinates (i.e. doesn't include
# center-of-mass motion, or scale to motion of bodies about their common center of mass):
jac_kepler = zeros(Float64,7,7)
kepler_step!(gm, h, state0, state, jac_kepler)
for k=1:NDIM
  delx[k] = state[1+k] - state0[1+k]
  delv[k] = state[1+NDIM+k] - state0[1+NDIM+k]
end
# Compute COM coords:
mijinv =1.0/(m[i] + m[j])
xcm = zeros(Float64,NDIM)
vcm = zeros(Float64,NDIM)
mi = m[i]*mijinv # Normalize the masses
mj = m[j]*mijinv
for k=1:NDIM
  xcm[k] = mi*x[k,i] + mj*x[k,j]
  vcm[k] = mi*v[k,i] + mj*v[k,j]
end
# Compute the Jacobian:
jac_ij[ 7, 7] = 1.0  # the masses don't change with time!
jac_ij[14,14] = 1.0
for k=1:NDIM
   jac_ij[   k,   k] +=   mi
   jac_ij[   k, 3+k] += h*mi
   jac_ij[   k, 7+k] +=   mj
   jac_ij[   k,10+k] += h*mj
   jac_ij[ 3+k, 3+k] +=   mi
   jac_ij[ 3+k,10+k] +=   mj
   jac_ij[ 7+k,   k] +=   mi
   jac_ij[ 7+k, 3+k] += h*mi
   jac_ij[ 7+k, 7+k] +=   mj
   jac_ij[ 7+k,10+k] += h*mj
   jac_ij[10+k, 3+k] +=   mi
   jac_ij[10+k,10+k] +=   mj
   for l=1:NDIM
# Compute derivatives of \delta x_i with respect to initial conditions:
     jac_ij[   k,   l] += mj*jac_kepler[  k,  l]
     jac_ij[   k, 3+l] += mj*jac_kepler[  k,3+l]
     jac_ij[   k, 7+l] -= mj*jac_kepler[  k,  l]
     jac_ij[   k,10+l] -= mj*jac_kepler[  k,3+l]
# Compute derivatives of \delta v_i with respect to initial conditions:
     jac_ij[ 3+k,   l] += mj*jac_kepler[3+k,  l]
     jac_ij[ 3+k, 3+l] += mj*jac_kepler[3+k,3+l]
     jac_ij[ 3+k, 7+l] -= mj*jac_kepler[3+k,  l]
     jac_ij[ 3+k,10+l] -= mj*jac_kepler[3+k,3+l]
# Compute derivatives of \delta x_j with respect to initial conditions:
     jac_ij[ 7+k,   l] -= mi*jac_kepler[  k,  l]
     jac_ij[ 7+k, 3+l] -= mi*jac_kepler[  k,3+l]
     jac_ij[ 7+k, 7+l] += mi*jac_kepler[  k,  l]
     jac_ij[ 7+k,10+l] += mi*jac_kepler[  k,3+l]
# Compute derivatives of \delta v_j with respect to initial conditions:
     jac_ij[10+k,   l] -= mi*jac_kepler[3+k,  l]
     jac_ij[10+k, 3+l] -= mi*jac_kepler[3+k,3+l]
     jac_ij[10+k, 7+l] += mi*jac_kepler[3+k,  l]
     jac_ij[10+k,10+l] += mi*jac_kepler[3+k,3+l]
   end
# Compute derivatives of \delta x_i with respect to the masses:
   jac_ij[   k, 7] += (x[k,i]+h*v[k,i]-xcm[k]-mj*state[1+k])*mijinv + GNEWT*mj*jac_kepler[  k,7]
   jac_ij[   k,14] += (x[k,j]+h*v[k,j]-xcm[k]+mi*state[1+k])*mijinv + GNEWT*mj*jac_kepler[  k,7]
# Compute derivatives of \delta v_i with respect to the masses:
   jac_ij[ 3+k, 7] += (v[k,i]-vcm[k]-mj*state[4+k])*mijinv + GNEWT*mj*jac_kepler[3+k,7]
   jac_ij[ 3+k,14] += (v[k,j]-vcm[k]+mi*state[4+k])*mijinv + GNEWT*mj*jac_kepler[3+k,7]
# Compute derivatives of \delta x_j with respect to the masses:
   jac_ij[ 7+k, 7] += (x[k,i]+h*v[k,i]-xcm[k]-mj*state[1+k])*mijinv - GNEWT*mi*jac_kepler[  k,7]
   jac_ij[ 7+k,14] += (x[k,j]+h*v[k,j]-xcm[k]+mi*state[1+k])*mijinv - GNEWT*mi*jac_kepler[  k,7]
# Compute derivatives of \delta v_j with respect to the masses:
   jac_ij[10+k, 7] += (v[k,i]-vcm[k]-mj*state[4+k])*mijinv - GNEWT*mi*jac_kepler[3+k,7]
   jac_ij[10+k,14] += (v[k,j]-vcm[k]+mi*state[4+k])*mijinv - GNEWT*mi*jac_kepler[3+k,7]
end
# Advance center of mass & individual Keplerian motions:
centerm!(m,mijinv,x,v,vcm,delx,delv,i,j,h)
return
end


using Flux
import KitML

# dataset
cd(@__DIR__)
dataset = KitML.load_data("data_mn.csv"; mode = :csv, dlm = ",")

baseDimension = 4 #length of u

for name in names(dataset)
    println(":", name, ", ")
end

L = size(dataset, 1)

# surrogate 1: moments -> entropy
begin
    X = zeros(baseDimension, size(dataset, 1))
    Y = zeros(1, size(dataset, 1))
    for j in axes(X, 2)
        X[1, j] = dataset.u_0[j]
        X[2, j] = dataset.u_1[j]
        X[3, j] = dataset.u_2[j]
        X[4, j] = dataset.u_3[j]
        Y[1, j] = dataset.h[j]
    end
    data = Flux.Data.DataLoader(X, Y, batchsize = 256, shuffle = true)

    #Base.show(nn)
    #nn = Chain(
    #    Dense(4, 16, relu),
    #    Dense(16, 16, relu),
    #    Dense(16, 16, relu),
    #    Dense(16, 1),
    #)
    nn = KitML.create_neural_closure(baseDimension, 1, 8, 3)

end
KitML.sci_train!(nn, data, ADAM(), 1000)
KitML.save_model(nn)
println("End ICNN Training")

# surrogate 2: ResNet
nn2 = Chain(
    Dense(4, 16, relu),
    KitML.Shortcut(Chain(Dense(16, 16, relu), Dense(16, 16, relu))),
    Dense(16, 1),
)
KitML.sci_train!(nn2, data, ADAM(), 2)

"""Notice the training of mapping u to Lagrangian multiplier λ seems not working well"""


"""
Creates lossy converter model between AC and DC grid
```
pconv_ac[i] + pconv_dc[i] == a + bI + cI^2
```
"""
function constraint_converter_losses(pm::AbstractACPModel, n::Int, cnd::Int, i::Int, a, b, c, plmax)
    pconv_ac = PowerModels.var(pm, n, cnd, :pconv_ac, i)
    pconv_dc = PowerModels.var(pm, n, cnd, :pconv_dc, i)
    iconv = PowerModels.var(pm, n, cnd, :iconv_ac, i)

    PowerModels.con(pm, n, cnd, :conv_loss)[i] = @NLconstraint(pm.model, pconv_ac + pconv_dc == a + b*iconv + c*iconv^2)
end
"""
Links converter power & current
```
pconv_ac[i]^2 + pconv_dc[i]^2 == vmc[i]^2 * iconv_ac[i]^2
```
"""
function constraint_converter_current(pm::AbstractACPModel, n::Int, cnd::Int, i::Int, Umax, Imax)
    vmc = PowerModels.var(pm, n, cnd, :vmc, i)
    pconv_ac = PowerModels.var(pm, n, cnd, :pconv_ac, i)
    qconv_ac = PowerModels.var(pm, n, cnd, :qconv_ac, i)
    iconv = PowerModels.var(pm, n, cnd, :iconv_ac, i)

    PowerModels.con(pm, n, cnd, :conv_i)[i] = @NLconstraint(pm.model, pconv_ac^2 + qconv_ac^2 == vmc^2 * iconv^2)
end
"""
Converter transformer constraints
```
p_tf_fr ==  g/(tm^2)*vm_fr^2 + -g/(tm)*vm_fr*vm_to * cos(va_fr-va_to) + -b/(tm)*vm_fr*vm_to*sin(va_fr-va_to)
q_tf_fr == -b/(tm^2)*vm_fr^2 +  b/(tm)*vm_fr*vm_to * cos(va_fr-va_to) + -g/(tm)*vm_fr*vm_to*sin(va_fr-va_to)
p_tf_to ==  g*vm_to^2 + -g/(tm)*vm_to*vm_fr  *    cos(va_to - va_fr)     + -b/(tm)*vm_to*vm_fr    *sin(va_to - va_fr)
q_tf_to == -b*vm_to^2 +  b/(tm)*vm_to*vm_fr  *    cos(va_to - va_fr)     + -g/(tm)*vm_to*vm_fr    *sin(va_to - va_fr)
```
"""
function constraint_conv_transformer(pm::AbstractACPModel, n::Int, cnd::Int, i::Int, rtf, xtf, acbus, tm, transformer)
    ptf_fr = PowerModels.var(pm, n, cnd, :pconv_tf_fr, i)
    qtf_fr = PowerModels.var(pm, n, cnd, :qconv_tf_fr, i)
    ptf_to = PowerModels.var(pm, n, cnd, :pconv_tf_to, i)
    qtf_to = PowerModels.var(pm, n, cnd, :qconv_tf_to, i)

    vm = PowerModels.var(pm, n, cnd, :vm, acbus)
    va = PowerModels.var(pm, n, cnd, :va, acbus)
    vmf = PowerModels.var(pm, n, cnd, :vmf, i)
    vaf = PowerModels.var(pm, n, cnd, :vaf, i)

    ztf = rtf + im*xtf
    if transformer
        ytf = 1/(rtf + im*xtf)
        gtf = real(ytf)
        btf = imag(ytf)
        gtf_sh = 0
        c1, c2, c3, c4 = ac_power_flow_constraints(pm.model, gtf, btf, gtf_sh, vm, vmf, va, vaf, ptf_fr, ptf_to, qtf_fr, qtf_to, tm)

        PowerModels.con(pm, n, cnd, :conv_tf_p_fr)[i] = c1
        PowerModels.con(pm, n, cnd, :conv_tf_q_fr)[i] = c2

        PowerModels.con(pm, n, cnd, :conv_tf_p_to)[i] = c3
        PowerModels.con(pm, n, cnd, :conv_tf_q_to)[i] = c4
    else
        PowerModels.con(pm, n, cnd, :conv_tf_p_fr)[i] = @constraint(pm.model, ptf_fr + ptf_to == 0)
        PowerModels.con(pm, n, cnd, :conv_tf_q_fr)[i] = @constraint(pm.model, qtf_fr + qtf_to == 0)
        @constraint(pm.model, va == vaf)
        @constraint(pm.model, vm == vmf)
    end
end
"constraints for a voltage magnitude transformer + series impedance"
function ac_power_flow_constraints(model, g, b, gsh_fr, vm_fr, vm_to, va_fr, va_to, p_fr, p_to, q_fr, q_to, tm)
    c1 = @NLconstraint(model, p_fr ==  g/(tm^2)*vm_fr^2 + -g/(tm)*vm_fr*vm_to * cos(va_fr-va_to) + -b/(tm)*vm_fr*vm_to*sin(va_fr-va_to))
    c2 = @NLconstraint(model, q_fr == -b/(tm^2)*vm_fr^2 +  b/(tm)*vm_fr*vm_to * cos(va_fr-va_to) + -g/(tm)*vm_fr*vm_to*sin(va_fr-va_to))
    c3 = @NLconstraint(model, p_to ==  g*vm_to^2 + -g/(tm)*vm_to*vm_fr  *    cos(va_to - va_fr)     + -b/(tm)*vm_to*vm_fr    *sin(va_to - va_fr))
    c4 = @NLconstraint(model, q_to == -b*vm_to^2 +  b/(tm)*vm_to*vm_fr  *    cos(va_to - va_fr)     + -g/(tm)*vm_to*vm_fr    *sin(va_to - va_fr))
    return c1, c2, c3, c4
end
"""
Converter reactor constraints
```
-pconv_ac == gc*vmc^2 + -gc*vmc*vmf*cos(vac-vaf) + -bc*vmc*vmf*sin(vac-vaf)
-qconv_ac ==-bc*vmc^2 +  bc*vmc*vmf*cos(vac-vaf) + -gc*vmc*vmf*sin(vac-vaf)
p_pr_fr ==  gc *vmf^2 + -gc *vmf*vmc*cos(vaf - vac) + -bc *vmf*vmc*sin(vaf - vac)
q_pr_fr == -bc *vmf^2 +  bc *vmf*vmc*cos(vaf - vac) + -gc *vmf*vmc*sin(vaf - vac)
```
"""
function constraint_conv_reactor(pm::AbstractACPModel, n::Int, cnd::Int, i::Int, rc, xc, reactor)
    pconv_ac = PowerModels.var(pm, n, cnd, :pconv_ac, i)
    qconv_ac = PowerModels.var(pm, n, cnd, :qconv_ac, i)
    ppr_fr = PowerModels.var(pm, n, cnd, :pconv_pr_fr, i)
    qpr_fr = PowerModels.var(pm, n, cnd, :qconv_pr_fr, i)

    vmf = PowerModels.var(pm, n, cnd, :vmf, i)
    vaf = PowerModels.var(pm, n, cnd, :vaf, i)
    vmc = PowerModels.var(pm, n, cnd, :vmc, i)
    vac = PowerModels.var(pm, n, cnd, :vac, i)

    zc = rc + im*xc
    if reactor
        yc = 1/(zc)
        gc = real(yc)
        bc = imag(yc)
        PowerModels.con(pm, n, cnd, :conv_pr_p)[i] = @NLconstraint(pm.model, -pconv_ac == gc*vmc^2 + -gc*vmc*vmf*cos(vac-vaf) + -bc*vmc*vmf*sin(vac-vaf))
        PowerModels.con(pm, n, cnd, :conv_pr_q)[i] = @NLconstraint(pm.model, -qconv_ac ==-bc*vmc^2 +  bc*vmc*vmf*cos(vac-vaf) + -gc*vmc*vmf*sin(vac-vaf))
        @NLconstraint(pm.model, ppr_fr ==  gc *vmf^2 + -gc *vmf*vmc*cos(vaf - vac) + -bc *vmf*vmc*sin(vaf - vac))
        @NLconstraint(pm.model, qpr_fr == -bc *vmf^2 +  bc *vmf*vmc*cos(vaf - vac) + -gc *vmf*vmc*sin(vaf - vac))
    else
        ppr_to = -pconv_ac
        qpr_to = -qconv_ac
        PowerModels.con(pm, n, cnd, :conv_pr_p)[i] = @constraint(pm.model, ppr_fr + ppr_to == 0)
        PowerModels.con(pm, n, cnd, :conv_pr_q)[i] = @constraint(pm.model, qpr_fr + qpr_to == 0)
        @constraint(pm.model, vac == vaf)
        @constraint(pm.model, vmc == vmf)

    end
end
"""
Converter filter constraints
```
ppr_fr + ptf_to == 0
qpr_fr + qtf_to +  (-bv) * filter *vmf^2 == 0
```
"""
function constraint_conv_filter(pm::AbstractACPModel, n::Int, cnd::Int, i::Int, bv, filter)
    ppr_fr = PowerModels.var(pm, n, cnd, :pconv_pr_fr, i)
    qpr_fr = PowerModels.var(pm, n, cnd, :qconv_pr_fr, i)
    ptf_to = PowerModels.var(pm, n, cnd, :pconv_tf_to, i)
    qtf_to = PowerModels.var(pm, n, cnd, :qconv_tf_to, i)

    vmf = PowerModels.var(pm, n, cnd, :vmf, i)

    PowerModels.con(pm, n, cnd, :conv_kcl_p)[i] = @constraint(pm.model,   ppr_fr + ptf_to == 0 )
    PowerModels.con(pm, n, cnd, :conv_kcl_q)[i] = @NLconstraint(pm.model, qpr_fr + qtf_to +  (-bv) * filter *vmf^2 == 0)
end
"""
LCC firing angle constraints
```
pconv_ac == cos(phi) * Srated
qconv_ac == sin(phi) * Srated
```
"""
function constraint_conv_firing_angle(pm::AbstractACPModel, n::Int, cnd::Int, i::Int, S, P1, Q1, P2, Q2)
    p = PowerModels.var(pm, n, cnd, :pconv_ac, i)
    q = PowerModels.var(pm, n, cnd, :qconv_ac, i)
    phi = PowerModels.var(pm, n, cnd, :phiconv, i)

    PowerModels.con(pm, n, cnd, :conv_cosphi)[i] = @NLconstraint(pm.model,   p == cos(phi) * S)
    PowerModels.con(pm, n, cnd, :conv_sinphi)[i] = @NLconstraint(pm.model,   q == sin(phi) * S)
end

function constraint_dc_droop_control(pm::AbstractACPModel, n::Int, cnd::Int, i::Int, busdc_i, vref_dc, pref_dc, k_droop)
    pconv_dc = PowerModels.var(pm, n, cnd, :pconv_dc, i)
    vdc = PowerModels.var(pm, n, cnd, :vdcm, busdc_i)

    PowerModels.con(pm, n, cnd, :conv_dc_droop)[i] = @constraint(pm.model, pconv_dc == pref_dc - sign(pref_dc) * 1 / k_droop * (vdc - vref_dc))
end


using Libm
using BenchmarkTools

xerf1 = linspace(-0.84375,0.84375,10000)
xerf2 = linspace(-2e-28,2e-28,10000)
xerf3 = linspace(-0.84375,0.84375,10000)
xerf4 = linspace(0.84375, 6,10000)
xerf5 = linspace(-6,-0.84375,10000);
xerf = union(xerf1,xerf2,xerf3,xerf4,xerf5);

xerfc1 = linspace(-0.84375,0.84375,10000)
xerfc2 = linspace(-2e-56,2e-56,10000)
xerfc3 = linspace(-0.25,0.25,10000)
xerfc4 = linspace(0.84375, 28,10000)
xerfc5 = linspace(-28,-0.84375,10000);
xercf = union(xerfc1,xerfc2,xerfc3,xerfc4,xerfc5)

t1 = @benchmark Libm.erf.($xerf)
t2 = @benchmark Base.erf.($xerf)
t3 = @benchmark Libm.erfc.($xercf)
t4 = @benchmark Base.erfc.($xercf)

println("***** Benchmark erf *****")
println()
println(ratio(mean(t1),mean(t2)))
println()
println("Details\n")
println(t1,"\n")
println(t2,"\n")

println("***** Benchmark erfc *****")
println()
println(ratio(mean(t3),mean(t4)))
println()
println("Details\n")
println(t3,"\n")
println(t4,"\n")


################################################################################
### For generic methods
################################################################################

# Heuristic for selecting the level with the (presumably) most reliable
# standard error estimate:
# Take the highest lvl with at least 32 bins.
# (Chose 32 based on https://doi.org/10.1119/1.3247985)
function _reliable_level(B::ErrorPropagator{T,N})::Int64 where {T,N}
    isempty(B) && (return 1)                # results in NaN in std_error
    i = findlast(x -> x >= 32, B.count)
    return something(i, 1)
end

_eachlevel(B::ErrorPropagator) = 1:findlast(x -> x > 1, B.count)

################################################################################
### Statistics
################################################################################

"""
    mean(B::ErrorPropagator, i[, lvl = 1])

Calculates the mean for the i-th argument of the error propagator at an optional
binning level.
"""
function mean(B::ErrorPropagator, i::Integer, lvl = 1)
    return B.sums1D[lvl][i] / B.count[lvl]
end

"""
    var(B::ErrorPropagator, i[, lvl])

Calculates the variance for the i-th argument of the error propagator at a given
binning level.
"""
function var(B::ErrorPropagator{<: Real}, i::Integer, lvl = _reliable_level(B))
    n = B.count[lvl]
    X = B.sums1D[lvl][i]
    X2 = B.sums2D[lvl][i, i]

    # lvl = 1 <=> original values
    # correct variance:
    # (∑ xᵢ^2) / (N-1) - (∑ xᵢ)(∑ xᵢ) / (N(N-1))
    return X2 / (n - 1) - X^2 / (n*(n - 1))
end

function var(B::ErrorPropagator{<: Complex}, i::Integer, lvl = _reliable_level(B))
    n = B.count[lvl]
    X = B.sums1D[lvl][i]
    X2 = B.sums2D[lvl][i, i]

    return (real(X2) + imag(X2)) / (n - 1) - (real(X)^2 + imag(X)^2) / (n*(n - 1))
end

function var(B::ErrorPropagator{<: AbstractArray{<: Real}}, i::Integer, lvl = _reliable_level(B))
    n = B.count[lvl]
    X = B.sums1D[lvl][i]
    X2 = B.sums2D[lvl][i, i]

    return @. X2 / (n - 1) - X^2 / (n*(n - 1))
end

function var(B::ErrorPropagator{<: AbstractArray{<: Complex}}, i::Integer, lvl = _reliable_level(B))
    n = B.count[lvl]
    X = B.sums1D[lvl][i]
    X2 = B.sums2D[lvl][i, i]

    return @. (real(X2) + imag(X2)) / (n - 1) - (real(X)^2 + imag(X)^2) / (n*(n - 1))
end

"""
    varN(B::ErrorPropagator, i[, lvl])

Calculates the variance/N for the i-th argument of the error propagator at a
given binning level.
"""
function varN(B::ErrorPropagator, i::Integer, lvl = _reliable_level(B))
    return var(B, i, lvl) / B.count[lvl]
end

"""
    tau(B::ErrorPropagator, i[, lvl])

Calculates the autocorrelation time tau for the i-th argument of the error
propagator at a given binning level.
"""
function tau(B::ErrorPropagator{<: Number}, i::Integer, lvl = _reliable_level(B))
    return 0.5 * (varN(B, i, lvl) / varN(B, i, 1) - 1)
end
function tau(B::ErrorPropagator{<: AbstractArray}, i::Integer, lvl = _reliable_level(B))
    return 0.5 * (varN(B, i, lvl) ./ varN(B, i, 1) .- 1)
end
autocorrelation(B::ErrorPropagator, i::Integer, lvl = _reliable_level(B)) = tau(B, i, lvl)
autocorrelation_time(B::ErrorPropagator, i::Integer, lvl = _reliable_level(B)) = tau(B, i, lvl)

"""
    std_error(B::ErrorPropagator, i[, lvl])

Calculates the standard error of the mean for the i-th argument of the error
propagator at a given binning level.
"""
function std_error(B::ErrorPropagator{<: Number}, i::Integer, lvl=_reliable_level(B))
    return sqrt(varN(B, i, lvl))
end
function std_error(B::ErrorPropagator{<: AbstractArray}, i::Integer, lvl=_reliable_level(B))
    return sqrt.(varN(B, i, lvl))
end

################################################################################
### Additional methods
################################################################################

for name in [:mean, :varN, :var, :tau, :std_error, :autocorrelation, :autocorrelation_time]
    # generates functions fs(B[, lvl]) = [f(B, 1, lvl), .., f(B, N_args, lvl)]
    @eval begin
        function $(Symbol(name, :s))(B::ErrorPropagator, lvl=_reliable_level(B))
            return [$name(B, i, lvl) for i in eachindex(B.sums1D[1])]
        end
    end

    # all_fs(ep) = [[f(ep, 1, lvl), .., f(ep, N_args, lvl)] for lvl in eachlvl]
    @eval begin
        function $(Symbol(:all_, name, :s))(B::ErrorPropagator)
            return [$(Symbol(name, :s))(B, lvl) for lvl in _eachlevel(B)]
        end
    end
end

# Docs
@doc """
    varNs(B::ErrorPropagator[, lvl])

Calculates the variance/N for each argument of the error propagator at a given
binning level.
""" varNs
@doc """
    vars(B::ErrorPropagator[, lvl])

Calculates the variance for each argument of the error propagator at a given
binning level.
""" vars
@doc """
    means(B::ErrorPropagator[, lvl])

Calculates the mean for each argument of the error propagator at a given
binning level.
""" means
@doc """
    std_errors(B::ErrorPropagator[, lvl])

Calculates the standard error of the mean for each argument of the error
propagator at a given binning level.
""" std_errors

"""
    covmat(B::ErrorPropagator[, lvl])

Returns the covariance matrix for a given level of the error propgator.
"""
function covmat(
        B::ErrorPropagator{T, N}, lvl = _reliable_level(B)
    ) where {N , T <: Number}

    invN = 1.0 / B.count[lvl]
    invN1 = 1.0 / (B.count[lvl] - 1)
    return [
        (
            B.sums2D[lvl][i, j] -
            B.sums1D[lvl][i] * conj(B.sums1D[lvl][j]) * invN
        ) * invN1
        for i in eachindex(B.sums1D[lvl]), j in eachindex(B.sums1D[lvl])
    ]
end
function covmat(
        B::ErrorPropagator{T, N}, lvl = _reliable_level(B)
    ) where {N , T <: AbstractArray}

    invN = 1.0 / B.count[lvl]
    invN1 = 1.0 / (B.count[lvl] - 1)
    return [
        (
            B.sums2D[lvl][i, j] -
            B.sums1D[lvl][i] .* conj(B.sums1D[lvl][j]) * invN
        ) * invN1
        for i in eachindex(B.sums1D[lvl]), j in eachindex(B.sums1D[lvl])
    ]
end

"""
    var(B::ErrorPropagator, gradient[, lvl])

Gives the first-order variance estimate of a function `f` acting on the
arguments of the error propagator. `gradient` is either the gradient of `f` (a
function) or a vector `∇f(means(B))`. To get an estimate mean value of `f`,
`mean(B, f)` can be used.
"""
function var(
        B::ErrorPropagator{T, N},
        gradient::Vector,
        lvl = _reliable_level(B)
    ) where {T <: Real, N}

    result = 0.0
    invN = 1.0 / B.count[lvl]
    # invNN1 = 1.0 / (B.count[lvl] * (B.count[lvl] - 1))
    invN1 = 1.0 / (B.count[lvl] - 1)
    for i in eachindex(B.sums1D[lvl])
        for j in eachindex(B.sums1D[lvl])
            result += gradient[i] * gradient[j] * (
                B.sums2D[lvl][i, j] -
                B.sums1D[lvl][i] * B.sums1D[lvl][j] * invN
            ) * invN1
        end
    end

    return result
end

function var(
        B::ErrorPropagator{T, N},
        gradient::Vector,
        lvl = _reliable_level(B)
    ) where {T <: Complex, N}

    result = 0.0 * 0.0im
    invN = 1.0 / B.count[lvl]
    # invNN1 = 1.0 / (B.count[lvl] * (B.count[lvl] - 1))
    invN1 = 1.0 / (B.count[lvl] - 1)
    for i in eachindex(B.sums1D[lvl])
        for j in eachindex(B.sums1D[lvl])
            result += gradient[i] * conj(gradient[j]) * (
                B.sums2D[lvl][i, j] -
                B.sums1D[lvl][i] * conj(B.sums1D[lvl][j]) * invN
            ) * invN1
        end
    end

    return abs(result)
end

# Wrappers
function var(B::ErrorPropagator, gradient::Function, lvl = _reliable_level(B))
    grad = gradient(means(B))
    if typeof(grad) <: AbstractArray
        return var(B, grad, lvl)
    else
        return var(B, [grad], lvl)
    end
end

"""
    varN(B::ErrorPropagator, gradient[, lvl])

    Gives the first-order variance/N estimate of a function `f` acting on the
    arguments of the error propagator. `gradient` is either the gradient of `f` (a
    function) or a vector `∇f(means(B))`. To get an estimate mean value of `f`,
    `mean(B, f)` can be used.
"""
function varN(B::ErrorPropagator, gradient::Function, lvl = _reliable_level(B))
    return var(B, gradient, lvl) / B.count[lvl]
end

"""
    std_error(B::ErrorPropagator, gradient[, lvl])

Gives the first-order standard error estimate of a function `f` acting on the
arguments of the error propagator. `gradient` is either the gradient of `f` (a
function) or a vector `∇f(means(B))`. To get an estimate mean value of `f`,
`mean(B, f)` can be used.
"""
function std_error(B::ErrorPropagator, gradient::Function, lvl = _reliable_level(B))
    return sqrt(var(B, gradient, lvl) / B.count[lvl])
end

"""
    mean(B, f[, lvl=1])

Returns an estimate for the mean value of `f`, where `f` is a function acting
on the sample pushed the error porpagator. `f` must be of the form `f(v)`, where
`v = [mean_arg1, mean_arg2, ..., mean_argN]` is a vector containing the averages
of each argument pushed to the error propagator.
"""
function mean(B::ErrorPropagator, f::Function, lvl = 1)
    return f(means(B, lvl))
end


Double64(x::Float64) = Double64((x, 0.0))

for T in (:Float32, :Float16, :Int32, :Int16, :Int8)
    @eval Double64(x::$T) = ((Float64(x), 0.0))
end

for T in (:Int64, :Int128)
  @eval function Double64(x::$T)
    if abs(x) <= maxintfloat(Float64)
        hi = Float64(x)
        lo = 0.0
    else
        bf = Float128(x)
        hi = Float64(x)
        lo = Float64(bf - hi)
    end
    return Double64((hi, lo))
  end
end

function Double64(x::BigFloat)
    hi = Float64(x)
    lo = Float64(x - hi)
    return Double64((hi, lo))
end

function Double64(x::Float128)
    hi = Float64(x)
    lo = Float64(x - hi)
    return Double64((hi, lo))
end

function Double64(x::BigInt)
    if abs(x) <= typemax(Int128)
        Double64(Int128(x))
    else
         Double64(BigFloat(x))
    end
end

Double64(x::Real, y::Real) = Double64(Float64(x), Float64(y))

Double64(x::ComplexD64) = real(x)
Double64(x::ComplexF64) = Double64(real(x))

ComplexD64(x::Float64) = ComplexD64(Double64(x))
ComplexD64(x::Float64, y::Float64) = ComplexD64( Double64(x), Double64(y) )
ComplexD64(x::ComplexF64) = ComplexD64(x.re, x.im)

# inverse constructors
Base.Float64(x::Double64) = Hi(x)
Base.Float64(x::ComplexD64) = real(Hi(x))
Base.ComplexF64(x::Double64) = ComplexF64(Hi(x), 0.0)
Base.ComplexF64(x::ComplexD64) = ComplexF64(Hi(x)...)
Base.BigFloat(x::Double64) = BigFloat(Hi(x)) + BigFloat(Lo(x))
Base.BigFloat(x::ComplexD64) = BigFloat(real(Hi(x))) + BigFloat(real(Lo(x)))
Base.BigInt(x::Double64) = BigInt(BigFloat(x))
Base.BigInt(x::ComplexD64) = BigInt(BigFloat(x))
Base.Int128(x::Double64) = Int128(BigInt(x))
Base.Int128(x::ComplexD64) = Int128(BigInt(x))

function Base.Int64(x::Double64)
    Int64(Hi(x)) + Int64(Lo(x))
end

for T in (:Float32, :Float16, :Int32, :Int16, :Int8)
    @eval Base.$T(x::Double64) = $T(Hi(x))
end

Float128(x::Double64) = Float128(Hi(x)) + Float128(Lo(x))
Double64(x::Float128) = Double64((Float64(x), Float64(x-Float64(x))))
ComplexD64(x::ComplexF128) = (Double64(real(x)), Double64(imag(x)))
ComplexF128(x::ComplexD64) = ComplexF128(real(x), imag(x))


struct K80rel <: K80
  κ::Float64
  function K80rel(κ::Float64;
                  safe::Bool=true)
    if safe
      if κ <= 0.
        error("K80 parameter κ must be positive")
      end
    end
    new(κ)
  end
end

function Base.show(io::IO, object::K80rel)
  print(io, "\r\e[0m\e[1mK\e[0mimura 19\e[1m80\e[0m model (relative rate form)
κ = $(object.κ)")
end

@inline function Q(mod::K80rel)
  κ = mod.κ

  Q₁ = 0.25 * κ
  Q₂ = 0.25
  Q₃ = -(2 * Q₂ + Q₁)

  return Qmatrix(Q₃, Q₂, Q₁, Q₂,
                 Q₂, Q₃, Q₂, Q₁,
                 Q₁, Q₂, Q₃, Q₂,
                 Q₂, Q₁, Q₂, Q₃)
end

@inline function P(mod::K80rel, t::Float64)
  if t < 0.0
    error("t must be positive")
  end
  κ = mod.κ

  e₁ = 0.25 * exp(-t)
  e₂ = 0.5 * exp(-0.5 * t * (κ + 1))

  P₁ = 0.25 + e₁ + e₂
  P₂ = 0.25 + e₁ - e₂
  P₃ = 0.25 - e₁

  return Pmatrix(P₁, P₃, P₂, P₃,
                 P₃, P₁, P₃, P₂,
                 P₂, P₃, P₁, P₃,
                 P₃, P₂, P₃, P₁)
end


include("../utils.jl")

@testset "FunctionCurve manifold" begin
    M = FunctionCurveSpace(Euclidean(2))
    @testset "FunctionCurve Basics" begin
        @test manifold_dimension(M) === Inf
    end

    x = t -> SA[t, t - 1]
    y = t -> SA[t^2, sin(t)]
    z = t -> SA[2 - 2, cos(t)]

    test_manifold(
        M,
        [x, y, z],
        test_forward_diff=false,
        test_reverse_diff=false,
        test_vector_spaces=true,
        test_project_tangent=false,
        test_musical_isomorphisms=false,
        test_default_vector_transport=false,
        test_tangent_vector_broadcasting=false,
        test_representation_size=false,
        is_mutating=false,
        expected_dimension_type=Float64,
    )
end


"""
Adapted from sandreza/Learning/sandbox/gaussian_process.jl
https://github.com/sandreza/Learning/blob/master/sandbox/gaussian_process.jl
Changed handling of kernel functions; changed some variable names;
added log marginal likelihood function.
"""

using LinearAlgebra

"""
GP_multiple
# Description
- data structure for GPR computations where each gridpoint in the prediction has a different predictor
# Data Structure and Description
    GPs, Array of GP objects
    kernel, Kernel object
    x_train
"""
struct GP_multiple
    GPs::Array{GP}
    kernel::Kernel
    x_train
end

function stencil_range(D, stencil_size, i)
    k = Int(floor(stencil_size/2))
    if i-k < 1
        return 1:stencil_size
    elseif i-k+stencil_size-1 > D
        return D-stencil_size+1:D
    else
        start = i-k
        start:start+stencil_size-1
    end
end

stencil(stencil_range, data) = [x[stencil_range] for x in data]

function model(𝒟::ProfileData; kernel::Kernel = Kernel(), stencil_size=)
    # create instance of GP using data from ProfileData object
    stencil_ranges = [stencil_range(D,stencil_size,i) for i=1:D]
    GPs = [model(stencil(𝒟.x_train,r), stencil(𝒟.y_train,r), kernel, 𝒟.zavg[r]) for range in stencil_ranges]
    return GP_multiple(GPs, kernel, 𝒟.x_train);
end

"""
model_output(x, 𝒢::GP_multiple)
# Description
- Given state x, GP 𝒢, returns the mean GP prediction
# Arguments
- `x`: single scaled state
- `𝒢`: GP_multiple object with which to make the prediction
# Return
- `y`: scaled prediction
"""
function model_output(x, 𝒢::GP_multiple)
    return [model_output(x[i],𝒢.GPs[i]) for i in 1:length(x)]
end

"""
uncertainty(x, 𝒢::GP)
# Description
- Given state x and GP 𝒢, output the variance at a point
# Arguments
- `x`: state
# Return
- `var`: variance
"""
function uncertainty(x, 𝒢::GP)
    tmpv = zeros(size(𝒢.x_train)[1])
    for i in eachindex(𝒢.x_train)
        tmpv[i] = 𝒢.kernel(x, 𝒢.x_train[i])
    end
    # no ldiv for suitesparse
    tmpv2 = 𝒢.CK \ tmpv
    var = k(x, x) .- tmpv'*tmpv2  # var(f*) = k(x*,x*) - tmpv'*tmpv2
    return var
end

"""
compute_kernel_matrix(kernel, x)
# Description
- Computes the kernel matrix for GPR
# Arguments
- `k` : (Kernel) kernel function k(a,b).
- `x` : (array of predictors). x[1] is a vector
# Return
- `sK`: (symmetric matrix). A symmetric matrix with entries sK[i,j] = k(x[i], x[j]). This is only meaningful if k(x,y) = k(y,x) (it should)
"""
function compute_kernel_matrix(k, x)

    K = [k(x[i], x[j]) for i in eachindex(x), j in eachindex(x)]

    if typeof(K[1,1]) <: Number
        sK = Symmetric(K)
    else
        sK = K
    end
    return sK
end

"""
mean_log_marginal_loss(y_train, 𝒢::GP; add_constant=false)
# Description
Computes log marginal loss for each element in the output and averages the results.
Assumes noise-free observations.

log(p(y|X)) = -(1/2) * (y'*α + 2*sum(Diagonal(CK)) + n*log(2*pi))
where n is the number of training points and

# Arguments
- `y_train`: (Array). training outputs (prediction), must have the same number as x_train
- `𝒢`: (GP).
# Keyword Arguments
- `add_constant`: (bool). whether to give the exact value of the loss or leave out an added constant for efficiency.

"""
function mean_log_marginal_loss(y_train, 𝒢::GP; add_constant=false)
    n = length(𝒢.x_train)
    D = length(𝒢.x_train[1])

    ys = hcat(y_train...)' # n x D

    if add_constant
        c = sum([log(𝒢.CK[i,i]) for i in 1:n]) + 0.5*n*log(2*pi)
        total_loss=0.0
        for i in 1:D
            total_loss -= 0.5*ys[:,i]'*𝒢.α[:,i] + c
        end
    else
        total_loss=0.0
        for i in 1:D
            total_loss -= 0.5*ys[:,i]'*𝒢.α[:,i]
        end
    end

    return total_loss / D
end