##### List of functions to analyze data

function extend_axis(df::AbstractDataFrame, xlabel, ylabel, xaxis, val)
    aux = DataFrame()
    aux[xlabel] = xaxis
    extended = join(aux, df, on = xlabel, kind = :left)
    sort!(extended, cols = [xlabel])
    return convert(Array, extended[ylabel], val)
end

function extend_axis(xsmall, ysmall, xaxis, val)
    df = DataFrame(x = xsmall, y = ysmall)
    return extend_axis(df, :x, :y, xaxis, val)
end

"""
    `_locreg(df, xaxis::LinSpace, x, y; kwargs...)`

Apply loess regression, training the regressor with `x` and `y` and
predicting `xaxis`
"""
function _locreg(df, xaxis::LinSpace, x, y; kwargs...)
    predicted = fill(NaN,length(xaxis))
    within = minimum(df[x]).<= xaxis .<= maximum(df[x])
    if any(within)
        model = Loess.loess(convert(Vector{Float64},df[x]),convert(Vector{Float64},df[y]); kwargs...)
        predicted[within] = Loess.predict(model,xaxis[within])
    end
    return predicted
end

"""
    `_locreg(df, xaxis, x, y; kwargs...)`

In the discrete case, the function computes the conditional expectation of `y` for
a given value of `x`
"""
function _locreg(df, xaxis, x,  y)
  ymean = by(df, x) do dd
      DataFrame(m = mean(dd[y]))
  end
  return extend_axis(ymean, x, :m, xaxis, NaN)
end

"""
    `_density(df,xaxis::LinSpace, x; kwargs...)`

Kernel density of `x`, computed along `xaxis`
"""
_density(df,xaxis::LinSpace, x; kwargs...) = pdf(KernelDensity.kde(df[x]; kwargs...),xaxis)

"""
    `_density(df, xaxis, x)`

Normalized histogram of `x` (which is discrete: every value is its own bin)
"""
function _density(df,xaxis, x)
    xhist = by(df, x) do dd
        DataFrame(length = size(dd,1)/size(df,1))
    end
    return extend_axis(xhist, x, :length, xaxis, 0.)
end

"""
    `_cumulative(df, xaxis, x) = ecdf(df[x])(xaxis)`

Cumulative density function of `x`, computed along `xaxis`
"""
_cumulative(df, xaxis, x) = ecdf(df[x])(xaxis)

#### Method to compute and plot grouped error plots using the above functions

type GroupedError{S, T<:AbstractString}
    x::Vector{Vector{S}}
    y::Vector{Vector{Float64}}
    err::Vector{Vector{Float64}}
    group::Vector{T}
    axis_type::Symbol
    show_error::Bool
end

get_axis(column) = sort!(union(column))
get_axis(column, npoints::Int64) = linspace(minimum(column),maximum(column),npoints)

function get_axis(column, axis_type::Symbol)
    if axis_type == :discrete
        return get_axis(column)
    elseif axis_type == :continuous
        return get_axis(column, 100)
    else
        error("Unexpected axis_type: only :discrete and :continuous allowed!")
    end
end

# f is the function used to analyze dataset: define it as nan when it is not defined,
# the input is: dataframe used, points chosen on the x axis, x (and maybe y) column labels
# the output is the y value for the given xvalues

get_symbol(s::Symbol) = s
get_symbol(s) = s[1]

function new_symbol(s, l::AbstractArray{Symbol})
    ns = s
    while ns in l
        ns = Symbol(ns,:_)
    end
    return ns
end

new_symbol(s, df::AbstractDataFrame) = new_symbol(s, names(df))

"""
    get_groupederror(trend,variation, f, splitdata::GroupedDataFrame, xvalues, args...; kwargs...)

Apply function `f` to `splitdata`, then compute summary statistics
`trend` and `variation` of those values. A shared x axis `xvalues` is needed: use
`LinSpace` for continuous x variable and a normal vector for the discrete case. Remaining arguments
are label of x axis variable and extra arguments for function `f`. `kwargs...` are passed
to `f`
"""
function get_groupederror(trend,variation, f, splitdata::GroupedDataFrame, xvalues::AbstractArray, args...; kwargs...)
    v = Array(Float64, length(xvalues), length(splitdata));
    for i in 1:length(splitdata)
        v[:,i] = f(splitdata[i],xvalues, args...; kwargs...)
    end
    mean_across_pop = Array(Float64, length(xvalues));
    sem_across_pop = Array(Float64, length(xvalues));
    for j in 1:length(xvalues)
        notnan = !isnan(v[j,:])
        mean_across_pop[j] = trend(v[j,notnan])
        sem_across_pop[j] = variation(v[j,notnan])
    end
    valid = !isnan(mean_across_pop) & !isnan(sem_across_pop)
    return xvalues[valid], mean_across_pop[valid], sem_across_pop[valid]
end

"""
    get_groupederror(trend,variation, f, df::AbstractDataFrame, xvalues::AbstractArray, ce, args...; kwargs...)

Get `GropedDataFrame` from `df` according to `ce`. `ce = (:across, col_name)` will split
across column `col_name`, whereas `ce = (:bootstrap, n_samples)` will generate `n_samples`
fake datasets distributed like the real dataset (nonparametric bootstrapping).
Then compute `get_groupederror` of the `GroupedDataFrame`.
"""
function get_groupederror(trend,variation, f, df::AbstractDataFrame, xvalues::AbstractArray, ce, args...; kwargs...)

    if ce == :none
        mean_across_pop = f(df,xvalues, args...; kwargs...)
        sem_across_pop = zeros(length(xvalues));
        valid = ~isnan(mean_across_pop)
        return xvalues[valid], mean_across_pop[valid], sem_across_pop[valid]
    elseif ce[1] == :across
        # get mean value and sem of function of interest
        splitdata = groupby(df, ce[2])
        return get_groupederror(trend,variation, f, splitdata, xvalues, args...; kwargs...)
    elseif ce[1] == :bootstrap
        n_samples = ce[2]
        indexes = Array(Int64,0)
        split_var = Array(Int64,0)
        for i = 1:n_samples
            append!(indexes, rand(1:size(df,1),size(df,1)))
            append!(split_var, fill(i,size(df,1)))
        end
        split_col = new_symbol(:split_col, df)
        bootstrap_data = df[indexes,:]
        bootstrap_data[split_col] = split_var
        ends = collect(size(df,1)*(1:n_samples))
        starts = ends - size(df,1) + 1
        splitdata = GroupedDataFrame(bootstrap_data,[split_col],collect(1:ends[end]), starts, ends)
        return get_groupederror(trend,variation, f, splitdata, xvalues, args...; kwargs...)
    else
        error("compute_error can only be equal to :none, :across,
        (:across, col_name), :bootstrap or (:bootstrap, n_samples)")
    end
end

"""
    get_groupederror(trend,variation, f, df::AbstractDataFrame, axis_type, ce, args...; kwargs...)

Choose shared axis according to `axis_type` (`:continuous` or `:discrete`) then
compute `get_groupederror`.
"""
function get_groupederror(trend,variation, f, df::AbstractDataFrame, axis_type, ce,  args...; kwargs...)
    # define points on x axis
    xvalues = get_axis(df[args[1]], axis_type)
    return get_groupederror(trend,variation, f, df::AbstractDataFrame, xvalues, ce, args...; kwargs...)
end

"""

    groupapply(f::Function, df, args...;
                axis_type = :auto, compute_error = :none, group = [],
                summarize = (get_symbol(compute_error) == :bootstrap) ? (mean, std) : (mean, sem),
                kwargs...)

Split `df` by `group`. Then apply `get_groupederror` to get a population summary of the grouped data.
Output is a `GroupedError` with error computed according to the keyword `compute_error`.
It can be plotted using `plot(g::GroupedError)`
Seriestype can be specified to be `:path`, `:scatter` or `:bar`
"""
function groupapply(f::Function, df, args...;
                    axis_type = :auto, compute_error = :none, group = [],
                    summarize = (get_symbol(compute_error) == :bootstrap) ? (mean, std) : (mean, sem),
                    compute_axis = :separate,
                    kwargs...)
    if !(axis_type in [:discrete, :continuous])
        axis_type = (typeof(df[args[1]])<:PooledDataArray) ? :discrete : :continuous
    end
    if (axis_type == :continuous) & !(eltype(df[args[1]])<:Real)
        warn("Changing to discrete axis, x values are not real numbers!")
        axis_type = :discrete
    end
    mutated_xtype = (axis_type == :continuous) ? Float64 : eltype(df[args[1]])

    # Add default for :across and :bootstrap
    if compute_error == :across
        row_name = new_symbol(:rows, df)
        df[row_name] = 1:size(df,1)
        ce = (:across, row_name)
    elseif compute_error == :bootstrap
        ce = (:bootstrap, 1000)
    else
        ce = compute_error
    end

    g = GroupedError(
                    Array(Vector{mutated_xtype},0),
                    Array(Vector{Float64},0),
                    Array(Vector{Float64},0),
                    Array(AbstractString,0),
                    axis_type,
                    ce != :none
                    )
    if group == []
        xvalues,yvalues,shade = get_groupederror(summarize..., f, df, axis_type, ce, args...; kwargs...)
        push!(g.x, xvalues)
        push!(g.y, yvalues)
        push!(g.err, shade)
        push!(g.group, "")
    else
        #group_array = isa(group, AbstractArray) ? group : [group]
        by(df,group) do dd
            label = isa(group, AbstractArray) ?
                    string(["$(dd[1,column]) " for column in group]...) : string(dd[1,group])
            xvalues,yvalues,shade = get_groupederror(summarize...,f, dd, axis_type, ce, args...; kwargs...)
            push!(g.x, xvalues)
            push!(g.y, yvalues)
            push!(g.err, shade)
            push!(g.group, label)
            return
        end
    end
    if compute_error == :across; delete!(df, row_name); end

    return g
end

builtin_funcs = Dict(zip([:locreg, :density, :cumulative], [_locreg, _density, _cumulative]))

"""
    groupapply(s::Symbol, df, args...; kwargs...)

`s` can be `:locreg`, `:density` or `:cumulative`, in which case the corresponding built in
analysis function is used. `s` can also be a symbol of a column of `df`, in which case the call
is equivalent to `groupapply(:locreg, df, args[1], s; kwargs...)`
"""
function groupapply(s::Symbol, df, args...; kwargs...)
    if s in keys(builtin_funcs)
        analysisfunction = builtin_funcs[s]
        return groupapply(analysisfunction, df, args...; kwargs...)
    else
        return groupapply(_locreg, df, args[1], s; kwargs...)
    end
end

"""
    groupapply(df::AbstractDataFrame, x, y; kwargs...)

Equivalent to `groupapply(:locreg, df::AbstractDataFrame, x, y; kwargs...)`
"""

groupapply(df::AbstractDataFrame, x, y; kwargs...) = groupapply(_locreg, df, x, y; kwargs...)

@recipe function f(g::GroupedError)
    if !(:seriestype in keys(d)) || (d[:seriestype] == :path)
        for i = 1:length(g.group)
            @series begin
                seriestype := :shadederror
                x := cycle(g.x,i)
                y := cycle(g.y, i)
                shade := cycle(g.err,i)
                label --> cycle(g.group,i)
                ()
            end
        end
    elseif d[:seriestype] == :scatter
        for i = 1:length(g.group)
            @series begin
                seriestype := :scatter
                x := cycle(g.x,i)
                y := cycle(g.y, i)
                if g.show_error
                    err := cycle(g.err,i)
                end
                label --> cycle(g.group,i)
                ()
            end
        end
    elseif d[:seriestype] == :bar
        if g.axis_type == :continuous
            warn("Bar plot with continuous x axis doesn't make sense!")
        end
        xaxis = sort!(union((g.x)...))
        ys = extend_axis.(g.x, g.y, [xaxis], [NaN])
        y = hcat(ys...)
        if g.show_error
            errs = extend_axis.(g.x, g.err, [xaxis], [NaN])
            err := hcat(errs...)
        end
        label --> hcat(g.group...)
        StatPlots.GroupedBar((xaxis,y))
    end
end

@shorthands GroupedError


"""

Structure that contains a two-dimensional point, subtype of FieldVector from StaticArrays.

"""
struct Point2D <: FieldVector{2,Float64}
  x::Float64
  y::Float64
end
norm2(p::Point2D) = (p.x^2 + p.y^2)
norm(p::Point2D) = sqrt(norm2(p))

"""

Structure that contains a three-dimensional point, subtype of FieldVector from StaticArrays.

"""
struct Point3D <: FieldVector{3,Float64}
  x::Float64
  y::Float64
  z::Float64
end
norm2(p::Point3D) = (p.x^2 + p.y^2 + p.z^2)
norm(p::Point3D) = sqrt(norm2(p))

"""

```
dim(T::DataType)
```

Returns the dimension associated to a `Point2D` (`2`) or `Point3D` (`3`).

"""
function dim(T::DataType)
  T == Point2D && return 2
  T == Point3D && return 3
end


export difference

# alias for set difference
import Base: \

const IA = IntervalArithmetic

"""
    \\(X::LazySet, Y::LazySet)

Convenience alias for set difference.

### Input

- `X` -- a set
- `Y` -- another set

### Output

The set difference between `X` and `Y`.

### Notes

If `X` and `Y` are intervals, `X \\ Y` is used in some libraries to denote
the left division, as the example below shows. However, it should not be
confused with the *set difference*. For example,

```jldoctest
julia> X = Interval(0, 2); Y = Interval(1, 4);

julia> X \\ Y   # computing the set difference
Interval{Float64, IntervalArithmetic.Interval{Float64}}([0, 1])

julia> X.dat \\ Y.dat  # computing the left division
[0.5, ∞]
```
"""
\(X::LazySet, Y::LazySet) = difference(X, Y)

# =================================
# Set difference between intervals
# =================================

"""
    difference(I1::IN, I2::IN) where {N, IN<:Interval{N}}

Return the set difference between the given intervals.

The set difference is defined as:

```math
    I₁ \\setminus I₂ = \\{x: x ∈ I₁ \\text{ and } x ∉ I₂ \\}
```

The backslash symbol, `\\`, can be used as an alias.

### Input

- `I1` -- first interval
- `I2` -- second interval

### Output

Depending on the position of the intervals, the output is one of the following:

- An `EmptySet`.
- An `Interval`.
- A `UnionSet` of two `Interval` sets.

### Algorithm

Let ``I₁ = [a, b]`` and ``I₂ = [c, d]`` be intervals. Their set difference is
``I₁ \\setminus I₂ = \\{x: x ∈ I₁ \\text{ and } x ∉ I₂ \\}`` and depending on their
position three different results may occur:

- If ``I₁`` and ``I₂`` do not overlap, i.e. if their intersection is empty, then the
  set difference is just ``I₁``.
- Otherwise, let `I₁₂ = I₁ ∩ I₂` and assume that it is not empty, then either
  ``I₁₂`` splits `I₁` into one interval or into two intervals. The latter case
  happens when the inclusion is strict on both ends of ``I₂``.

To check for strict inclusion, we assume that the inclusion is strict and then
check if the resulting intervals that cover `I₁` (one to its left and one to its
right, let them be `Ileft` and `Iright`), obtained by intersection with `I₂`,
are flat or not. Three cases may arise:

- If both `Ileft` and `Iright` are flat then it means that `I₁ = I₂`, then the
  set difference is the empty set.
- If only `Ileft` is flat, then the remaining interval not covered by `I₂` is
  `Iright`. In a similar manner, if only `Iright` is flat, then `Ileft` is returned.
- Finally, if none of the intervals is flat, then `I₂` is strictly contained in
  `I₁` and the set union of `Ileft` and `Iright` is returned.
"""
function difference(I1::IN, I2::IN) where {N, IN<:Interval{N}}
    I12 = intersection(I1, I2)
    if isempty(I12)
        return I1
    else
        Ileft = Interval(min(I1), min(I12))
        Iright = Interval(max(I12), max(I1))

        flat_left = isflat(Ileft)
        flat_right = isflat(Iright)

        if flat_left && flat_right
            return EmptySet{N}(1)
        elseif flat_left && !flat_right
            return Iright
        elseif !flat_left && flat_right
            return Ileft
        else
            return UnionSet(Ileft, Iright)
        end
    end
end

# ============================================
# Set difference between hyperrectangular sets
# ============================================

"""
    difference(X::AbstractHyperrectangle{N}, Y::AbstractHyperrectangle{N}) where {N}

Return the set difference between the given hyperrectangular sets.

### Input

- `X` -- first hyperrectangular set
- `Y` -- second hyperrectangular set

The set difference is defined as:

```math
    X \\setminus Y = \\{x: x ∈ X \\text{ and } x ∉ Y \\}
```

### Output

A `UnionSetArray` consisting of the union of hyperrectangles. Note that this
union is in general not convex.

### Algorithm

This function calls the implementation in `IntervalArithmetic.setdiff`.

### Notes

The backslash symbol, `\\`, can be used as an alias.
"""
function difference(X::AbstractHyperrectangle{N}, Y::AbstractHyperrectangle{N}) where {N}
    Xib = convert(IA.IntervalBox, X)
    Yib = convert(IA.IntervalBox, Y)
    return UnionSetArray(convert.(Hyperrectangle, IA.setdiff(Xib, Yib)))
end


# Problem 1
println("Problem 1:")
reduce(+, 0, filter(x -> x % 3 == 0 || x % 5 == 0, 1:999))

# Problem 2
# It seams that the problem is wrong.
# The Fib Seq starts with 1, 1, 2, 3, ...
println("Problem 2:")
a, b, s = 1, 1, 0
while a < 4000000
    if iseven(a)
        s += a
    end
    a, b = b, a+b
end
println(s)

# Problem 3
println("Problem 3:")
function  next_prime(primes)
    x = maximum(primes) + 1
    while true
        found_prime = true
        for prime in primes
            if x % prime == 0
                found_prime = false
                break
            end
        end

        if found_prime
            break
        end

        x = x + 1
    end
    return x
end

function p3()
    n = 600851475143
    primes = IntSet([2])
    prime = 2

    while n > 1.0
        if n % prime == 0
            n = n / prime
        else
            prime = next_prime(primes)
            union!(primes, [prime])
        end
    end
    println(prime)
end
@time p3()


"""
  Constructs and returns the energy system model.
"""

function buildmodel(input)

    println("\nBuilding model...")

    @unpack REGION, PLANT, PLANTFACT, HOUR, numregions, CO2Cap, load, maxcap, assum, discountrate, cf, hydro_inflow = input

    m = Model(Gurobi.Optimizer)

    @variables m begin
        Electricity[r in REGION, p in PLANT, h in HOUR]       >= 0        # MW
        InstalledCapacity[r in REGION, p in PLANT]            >= 0        # MW
        ReservoirContent[h in HOUR]                           >= 0        # TWh
        Emission                                              >= 0        # Mton

        # # Exercise 2
        BatteryContent[r in REGION, h in HOUR]                >= 0        # MWh / MW
        BatteryInput[r in REGION, h in HOUR]                  >= 0        # MW
        #
        # # Exercise 3
        TransmissionCapacity[r in REGION, R in REGION]        >= 0        # MW
        Transmission[r in REGION, R in REGION, h in HOUR]     >= 0        # MW, r: transmission in, R: transmission out
    end #variables
    #print(Capacity)

    #Variable bounds
    for r in REGION, p in PLANT
        set_upper_bound(InstalledCapacity[r, p], maxcap[r, p]) # MW
    end

    function a(d,p)
        1-(1/(1+d)^assum[:Lifetime,p])
    end

    # test expression TODO: check all expressions
    @expression(m, Runningcost[r in REGION], sum(Electricity[r,p,h].*assum[:RunCost,p] for p in PLANT, h in HOUR)) # €
    @expression(m, Investmentcost[r in REGION],
        sum(InstalledCapacity[r,p].*(assum[:InvestmentCost,p].*(10^3)).*discountrate./a(discountrate,p) for p in PLANT)) # €
    @expression(m, Fuelcost[r in REGION], sum(Electricity[r,p,h].*assum[:FuelCost,p]./assum[:Efficiency,p] for h in HOUR, p in PLANT)) #€
    @expression(m, Systemcost[r in REGION], Runningcost[r] + Investmentcost[r] + Fuelcost[r]) # €

    @expression(m, GeneratedElectricity[r in REGION, h in HOUR], sum(Electricity[r,p,h] for p in [:Wind, :PV, :Gas, :Hydro]))#, :Nuclear])) # MW
    # @expression(m, GeneratedElectricityNet[r in REGION, h in HOUR], sum(Electricity[r,p,h] for p in PLANT)-BatteryInput[r,h])
    @expression(m, GeneratedElectricityNet[r in REGION, h in HOUR],
        sum(Electricity[r,p,h] for p in PLANT)-BatteryInput[r,h]-sum(Transmission[R, r, h] for R in REGION)./assum[:Efficiency,:Transmission]) # transmission input in Electricity[r, :Transmission,h]

    @expression(m, CapacityPerHour[r in REGION, p in PLANT, h in HOUR], sum(InstalledCapacity[r, p].*cf[r, p, h])) # MW

    @constraints m begin
        Generation[r in REGION, p in PLANT, h in HOUR], # name of constraint
            Electricity[r, p, h] <= CapacityPerHour[r, p, h]

        Syscost[r in REGION], # name of constraint
            Systemcost[r] >= 0 # sum of all annualized costs

        CapacityConstraints[r in REGION, p in PLANT],
            InstalledCapacity[r, p] <= maxcap[r,p]

        #GenElec[r in REGION, h in HOUR],
        #   GeneratedElectricity[r, h] >= load[r, h]
        GenElec[r in REGION, h in HOUR],
            GeneratedElectricityNet[r, h] >= load[r, h]

        ReservoirBalance[h in HOUR[1:end-1]],
            ReservoirContent[h+1] <= ReservoirContent[h] + hydro_inflow[h] - Electricity[:SE, :Hydro, h]./(10^6)

        ReservoirEndOfYear,
            ReservoirContent[1] == ReservoirContent[end]

        ReservoirLimit[h in HOUR],
            ReservoirContent[h] <= 33 # TWh

        # #Exercise 2:
        MinEmission,
            Emission >= sum((Electricity[r, :Gas, h]./assum[:Efficiency, :Gas]).*assum[:EmissionFactor, :Gas] for r in REGION, h in HOUR)./(10^6)

        MaxEmission,
            Emission <= CO2Cap #Mton CO2

        BatteryBalance[r in REGION, h in HOUR[1:end-1]],
            BatteryContent[r, h+1] <= BatteryContent[r, h] - Electricity[r, :Batteries, h]./assum[:Efficiency, :Batteries] + BatteryInput[r, h] #TODO: h-1 correct?

        BatteryInputMax[r in REGION, h in HOUR],
            BatteryInput[r, h] <= sum(Electricity[r,p,h] for p in PLANT) # can get input from batteries, but would be inefficient

        BatteryMax[r in REGION, h in HOUR],
            BatteryContent[r, h] <= InstalledCapacity[r, :Batteries]

        # #Exercise 3
        TransmissionCapacityEquality[r in REGION, R in REGION],
            TransmissionCapacity[r, R] == TransmissionCapacity[R, r]

        TransmissionCapacityMax[r in REGION, R in REGION],
            TransmissionCapacity[r, R] <= maxcap[r, :Transmission]

            # Kolla så att blev rätt
        TransmissionInstalledCapacity[r in REGION],
            sum(TransmissionCapacity[r, R] for R in REGION)./2 == InstalledCapacity[r, :Transmission]

        TransmissionMax[r in REGION, R in REGION, h in HOUR],
            Transmission[r, R, h] <= TransmissionCapacity[r, R]

        TransmissionProduction[r in REGION, h in HOUR], # transmission in
            sum(Transmission[r, R, h] for R in REGION) == Electricity[r, :Transmission, h]

    end #constraints

    @objective m Min begin
        sum(Systemcost[r] for r in REGION)
    end # objective

    return (; m, InstalledCapacity, Electricity, Emission, GeneratedElectricity, TransmissionCapacity)

end # buildmodel

function calculate_generation(input)

end # calculate_generation


const K = 2
const N = 10
const α_neal = 10.0

@model function gmm_joint(z, x, ::Type{T}=Float64) where T
    if z === missing
        z = Vector{Int}(undef, N)
    end

    if x === missing
        x = Vector{T}(undef, N)
    end

    # Cluster centers.
    μ ~ filldist(Normal(), K)

    # Cluster association prior.
    w ~ Dirichlet(K, 1.0)

    # Cluster assignments.
    z ~ filldist(DiscreteNonParametric(1:K, w), N)

    # Observations.
    for n = 1:N
        x[n] ~ Normal(μ[z[n]], 1.0)
    end

    return z, x
end

@model function hmm_joint(s, x, ::Type{T}=Float64) where {T<:Real}
    if s === missing
        s = Vector{Int}(undef, N)
    end

    if x === missing
        x = Vector{T}(undef, N)
    end

    # Emission matrix.
    m = Vector{T}(undef, K)

    # Transition matrix.
    T = Vector{Vector{T}}(undef, K)

    # Assign distributions to each element
    # of the transition matrix and the
    # emission matrix.
    for i = 1:K
        T[i] ~ Dirichlet(K, 1.0)
        m[i] ~ Normal(i, 0.5)
    end

    # Observe each point of the input.
    s[1] ~ Categorical(K)
    x[1] ~ Normal(m[s[1]], 0.1)

    for i = 2:N
        s[i] ~ Categorical(T[s[i-1]])
        x[i] ~ Normal(m[s[i]], 0.1)
    end

    return s, x
end

function stickbreak(v)
    K = length(v) + 1
    cumprod_one_minus_v = cumprod(1 .- v)

    return map(1:K) do k
        if k == 1
            v[1]
        elseif k == K
            cumprod_one_minus_v[K - 1]
        else
            v[k] * cumprod_one_minus_v[k - 1]
        end
    end
end

# from https://luiarthur.github.io/TuringBnpBenchmarks/dpsbgmm
@model function imm_stick_joint(z, y, ::Type{T}=Float64) where {T}
    if z === missing
        z = Vector{Int}(undef, N)
    end

    if x === missing
        x = Vector{T}(undef, N)
    end

    crm = DirichletProcess(α_neal)

    v ~ filldist(StickBreakingProcess(crm), K - 1)
    w = stickbreak(v)

    # Cluster assignments
    for n = 1:N
        z[n] ~ Categorical(w)
    end

    # Cluster centers
    L = identity(K)
    μ ~ filldist(Normal(), L)

    # Observations
    for n = 1:N
        y[n] ~ Normal(μ[z[n]], 1.0)
    end
end

# test function
g(θ, x) = cat(θ, x; dims=1)

function geweke_gmm()
    rand_θ_given(x) = rand(StaticConditional(gmm_joint(missing, x), :z))

    test = perform(GewekeTest(2_000), gmm_joint, rand_θ_given, g)
    # plot(test, gmm_joint; size=(300, 300), title="GMM Geweke statistic")
    s = compute_statistic!(test, g)
    @test s.pvalue < 0.05
end


using Random
using JSON

mutable struct Network
    layers::Vector{Int64}
    neurons::Vector{Vector{Float64}}
    weights::Vector{Matrix{Float64}}
    biases::Vector{Vector{Float64}}
    z::Vector{Vector{Float64}}
    batch::Int64
    epochs::Int64
    η::Float64
    act
    act_p
    C
    ∇C_a

    function Network(layers,batch,epochs,η,act,act_p,C,∇C_a)
        new(layers::Vector{Int64},
        [[0.0 for j=1:layers[i]] for i=1:length(layers)],
        [randn(Float64,layers[i+1],layers[i]) for i=1:(length(layers)-1)],
        [rand!(MersenneTwister(1234), zeros(layers[i])) for i=1:(length(layers)-1)],
        [[0.0 for j=1:layers[i]] for i=1:length(layers)],
        batch::Int64,epochs::Int64,η::Float64,act,act_p,C,∇C_a)
    end
end

function feedforward!(net::Network,input::Vector{Float64})::Vector{Float64}
    net.layers[1] = input

    # a^(l+1) = σ(w^l*a^l + b^1)
    for i=1:(length(net.layers)-1)
        net.z[i] = net.weights[i]*net.neurons[i] + net.biases[i]
        net.neurons[i+1] = net.z[i].net.act()
    end

    return net.neurons[end]
end

function backprop(net::Network,input::Vector{Float64},label::Vector{Float64})::Vector{Vector{Float64}}
    feedforward!(net,input)

    # δ^L = ∇C_a ⊙ σ'(z^L)
    output_error = net.∇C_a(net.neurons[end],label).*net.act_p(net.z[end])
    errors = [output_error]

    for i=1:(length(net.layers)-1)
        # δ^l=((w^l+1)T δ^l+1) ⊙ σ′(z^l)
        next_error = (net.weights[length(net.layers)-i+1]*errors[1]).*net.act_p(net.z[length(net.layers)-i])
        prepend!(errors,next_error)
    end

    return errors, net.neurons
end

function sgd!(net::Network,train::Vector{Tuple{Vector{Float64},Vector{Float64}}})::Nothing
    for e=1:net.epochs
        train = shuffle(train)
        batches = [train[(1+(i*net.batch)):((i+1)*net.batch)] for i=0:(Int(length(train)/net.batch)-1)]

        for b in batches
            errors = []
            activations = []

            for x in b
                push!(errors,backprop(net,x[1],x[2])[1])
                push!(activations,backprop(net,x[1],x[2])[2])
            end

            # change w & b using errors = []
            for w1=1:length(net.weights)
                net.weights[wl] = net.weights[wl].-(net.η/net.batch).*sum([errors[e][wl]*transpose(activations[e][wl-1]) for e=1:length(errors)])

            end
            for bl=length(net.biases)
                net.biases[bl] = net.biases[bl].-(net.η/net.batch).*sum([errors[e][bl] for e=1:length(errors)])

            end
        end
    end

    return nothing
end

function evaluate(net::Network,test::Vector{Tuple{Vector{Float64},Vector{Float64}}},rule::Function)::Float64
    c = sum([1 for t in test if rule(feedforward(net,t[0])) == rule(t[1])])
    return "$(c) / $(length(test)) = $(100*(c/length(test)))%"
end


function manialn(X::AbstractMatrix{T},Y::AbstractMatrix{T}) where T<:Real
    μ,dim=0.9,30
    n1,n2=size(X,1),size(Y,1)
    W₁,W₂=X.+1,Y.+1
    ℐ=Matrix(I,n1,n2)
    μ = μ*(sum(W₁)+sum(W₂))/(2*sum(ℐ))
    𝕎 = [W₁ μ*ℐ; μ*ℐ' W₂]
    L=diagm(vec(sum(abs.(𝕎),dims=1))).-𝕎
    # λ,V =KrylovKit.eigsolve(L,35,:SR,krylovdim=40)
    # V=hcat(V)
    λ,V = eigen(L)
    i=real(λ).>=1e-8
    V=real(V[:,i])
    dim=min(dim,size(V,2))
    V=V[:,1:dim]
    aln0=V[1:n1,:]
    aln1=V[n1+1:end,:]
    d = norm.((aln0.-aln1)[i,:] for i = 1:n1)
    # _,idx=findmax(dd)
    return d
end


"Construction of the Kravchuk matrix using three-term recurrence relation"

export kmatrix

begin
    k₁(x, k, j) = ((2*x)/√((2 + 2*j - k)*(k - 1)))
    k₂(k, j) = √(((3 + 2*j - k)*(k - 2))/((2 + 2*j - k)*(k - 1)))
end

function kmatrix(j)
    vk = zeros(Float64, (Integer(2*j + 1), Integer(2*j + 1)))
    Threads.@threads for x in -j:j
        vk[Integer(x + j + 1), 1] = 2.0^(-j) * √binomial(Integer(2*j), Integer(j + x))
        vk[Integer(x + j + 1), 2] = ((2*x)/√(2*j)) * vk[Integer(x + j + 1), 1]
        for k in 2:Integer(2*j+1)
            if k == 2
                vk[Integer(x + j + 1), k] = k₁(x, k, j) * vk[Integer(x + j + 1), k - 1]
            else
                vk[Integer(x + j + 1), k] = k₁(x, k, j) * vk[Integer(x + j + 1), k - 1] - k₂(k, j) * vk[Integer(x + j + 1), k - 2]
            end
        end
    end
    return vk
end


######################################################################
# lqr.jl
#
# Solves LQR problems
######################################################################

export LQR
export apply_LQR_gains

function LQR(A::MF, B::MF, Q::MF, R::MF, N::Int)
	P = Array(Matrix{Float64}, N+1)
	G = Array(Matrix{Float64}, N)
	K = Array(Matrix{Float64}, N)
	P[N+1] = Q

	for n = (N-1):-1:0
		G[n+1] = R + (B' * P[n+1+1] * B)
		K[n+1] = inv(G[n+1]) * B' * P[n+1+1] * A
		P[n+1] = Q + (A' * P[n+1+1] * A) - (K[n+1]' * G[n+1] * K[n+1])
	end

	return K
end

function apply_LQR_gains(A::MF, B::MF, K::Vector{MF}, x0::Vector{Float64})
	x = [x0]
	u = Array(Vector{Float64}, 0)

	N = length(K)

	for n = 1:N
		push!(u, -K[n]*x[n])
		push!(x, A*x[n] + B*u[n])
	end
	return x, u
end


struct Sphere <: SDF
    Radius::Float64
    ShaderID::Shd_Id
end

function to_code(Sdf::Sphere)
    temp = gensym()
    a = :((norm($temp)-$(Sdf.Radius)),$(Sdf.ShaderID))
    SingFunc([a],temp)
end

struct Ball <: SDF
    Radius::Float64
    ShaderID::Shd_Id

        Metric::Float64 ## Infinity allowed for Max norm
end

function to_code(Sdf::Ball)
    temp = gensym()
    a = :((norm($temp,$(Sdf.Metric))-$(Sdf.Radius)),$(Sdf.ShaderID))
    SingFunc([a],temp)
end

struct Plane <: SDF
    Normal::space
    ShaderID::Shd_Id
end

function to_code(Sdf::Plane)
    temp = gensym()
    a = Expr(:tuple,Expr(:call,:dot,Sdf.Normal,temp),Sdf.ShaderID)
    SingFunc([a],temp)
end


c = 0
for n = 3:1000
  global c
  v = Rational{BigInt}(1, 2)
  for i = 1:n-1
    v += 2
    v = 1/v
  end
  v += 1
  if length(string(numerator(v))) > length(string(denominator(v)))
    c += 1
  end
end
println(c)


@testset "solve" begin

    @testset "total degree (simple)" begin
        @var x y
        affine_square = System([
            2.3 * x^2 + 1.2 * y^2 + 3x - 2y + 3,
            2.3 * x^2 + 1.2 * y^2 + 5x + 2y - 5,
        ])
        @test count(is_success, track.(total_degree(affine_square)...)) == 2

        @var x y z
        proj_square = System([
            2.3 * x^2 + 1.2 * y^2 + 3x * z - 2y * z + 3 * z^2,
            2.3 * x^2 + 1.2 * y^2 + 5x * z + 2y * z - 5 * z^2,
        ])
        @test count(is_success, track.(total_degree(proj_square)...)) == 4

        @var x y
        affine_overdetermined = System([
            (x^2 + y^2 + x * y - 3) * (x + 3),
            (x^2 + y^2 + x * y - 3) * (y - x + 2),
            2x + 5y - 3,
        ])
        @test count(is_success, track.(total_degree(affine_overdetermined)...)) == 2

        @var x y
        affine_overdetermined_reordering = System([
            (x^2 + y^2 + x * y - 3) * (x + 3),
            2x + 5y - 3,
            (x^2 + y^2 + x * y - 3) * (y^2 - x + 2),
        ])
        tracker, starts = total_degree(affine_overdetermined_reordering)
        @test length(starts) == 4 * 3
        @test count(is_success, track.(tracker, starts)) == 2

        @var x y z
        proj_overdetermined = System([
            (x^2 + y^2 + x * y - 3 * z^2) * (x + 3z),
            (x^2 + y^2 + x * y - 3 * z^2) * (y - x + 2z),
            2x + 5y - 3z,
        ])
        @test count(is_success, track.(total_degree(proj_overdetermined)...)) == 2

        @var x y
        proj_overdetermined_reordering = System([
            (x^2 + y^2 + x * y - 3 * z^2) * (x + 3z),
            2x + 5y - 3z,
            (x^2 + y^2 + x * y - 3 * z^2) * (y^2 - x * z + 2 * z^2),
        ])
        tracker, starts = total_degree(proj_overdetermined_reordering)
        @test length(starts) == 4 * 3
        @test count(is_success, track.(tracker, starts)) == 2

        @var x y
        affine_underdetermined = System([2.3 * x^2 + 1.2 * y^2 + 3x - 2y + 3])
        @test_throws HC2.FiniteException total_degree(affine_underdetermined)

        @var x y z
        proj_underdetermined = System([2.3 * x^2 + 1.2 * y^2 + 3x * z])
        @test_throws HC2.FiniteException total_degree(proj_underdetermined)
    end

    @testset "total degree (variable groups)" begin
        @var x y v w
        affine_sqr = System([x * y - 2, x^2 - 4], variable_groups = [[x], [y]])
        tracker, starts = total_degree(affine_sqr)
        @test length(collect(starts)) == 2
        @test count(is_success, track.(tracker, starts)) == 2

        @var x y v w
        proj_sqr =
            System([x * y - 2v * w, x^2 - 4 * v^2], variable_groups = [[x, v], [y, w]])
        tracker, starts = total_degree(proj_sqr)
        @test length(collect(starts)) == 2
        @test count(is_success, track.(tracker, starts)) == 2

        @var x y v w
        affine_overdetermined = System(
            [(x^2 - 4) * (x * y - 2), x * y - 2, x^2 - 4],
            variable_groups = [[x], [y]],
        )
        tracker, starts = total_degree(affine_overdetermined)
        @test count(is_success, track.(tracker, starts)) == 2
        @var x y v w
        proj_overdetermined = System(
            [(x^2 - 4 * v^2) * (x * y - v * w), x * y - v * w, x^2 - v^2],
            variable_groups = [[x, v], [y, w]],
        )
        tracker, starts = total_degree(proj_overdetermined)
        @test count(is_success, track.(tracker, starts)) == 2
    end

    @testset "polyhedral" begin
        @var x y
        affine_square = System([
            2.3 * x^2 + 1.2 * y^2 + 3x - 2y + 3,
            2.3 * x^2 + 1.2 * y^2 + 5x + 2y - 5,
        ])
        @test count(is_success, track.(polyhedral(affine_square)...)) == 2

        @var x y z
        proj_square = System([
            2.3 * x^2 + 1.2 * y^2 + 3x * z - 2y * z + 3 * z^2,
            2.3 * x^2 + 1.2 * y^2 + 5x * z + 2y * z - 5 * z^2,
        ])
        @test count(is_success, track.(polyhedral(proj_square)...)) == 4

        @var x y
        affine_overdetermined = System([
            (x^2 + y^2 + x * y - 3) * (x + 3),
            (x^2 + y^2 + x * y - 3) * (y - x + 2),
            2x + 5y - 3,
        ])
        @test count(is_success, track.(polyhedral(affine_overdetermined)...)) == 2

        @var x y z
        proj_overdetermined = System([
            (x^2 + y^2 + x * y - 3 * z^2) * (x + 3z),
            (x^2 + y^2 + x * y - 3 * z^2) * (y - x + 2z),
            2x + 5y - 3z,
        ])
        @test count(is_success, track.(polyhedral(proj_overdetermined)...)) == 2

        @var x y
        affine_underdetermined = System([2.3 * x^2 + 1.2 * y^2 + 3x - 2y + 3])
        @test_throws HC2.FiniteException polyhedral(affine_underdetermined)

        @var x y z
        proj_underdetermined = System([2.3 * x^2 + 1.2 * y^2 + 3x * z])
        @test_throws HC2.FiniteException polyhedral(proj_underdetermined)
    end

    @testset "overdetermined" begin
        @testset "3 by 5 minors" begin
            res = track.(total_degree(minors())...)
            @test count(is_success, res) == 80
        end
    end

    @testset "Result" begin
        d = 2
        @var x y a[1:6]
        F = System(
            [
                (a[1] * x^d + a[2] * y) * (a[3] * x + a[4] * y) + 1,
                (a[1] * x^d + a[2] * y) * (a[5] * x + a[6] * y) + 1,
            ];
            parameters = a,
        )
        res = solve(F; target_parameters = [0.257, -0.139, -1.73, -0.199, 1.79, -1.32])

        @test startswith(sprint(show, res), "Result with 3 solutions")
        @test seed(res) isa UInt32
        test_treeviews(res)

        seeded_res = solve(
            F;
            target_parameters = [0.257, -0.139, -1.73, -0.199, 1.79, -1.32],
            seed = seed(res),
        )
        @test seed(seeded_res) == seed(res)
        test_treeviews(res)

        @test length(path_results(res)) == ntracked(res) == 7
        @test length(results(res)) == nresults(res) == 3
        @test length(solutions(res)) == 3
        @test length(findall(is_success, res)) == 3
        @test real_solutions(res) isa Vector{Vector{Float64}}
        @test length(real_solutions(res)) == nreal(res) == 1
        @test length(nonsingular(res)) == nnonsingular(res) == 3
        @test isempty(singular(res))
        @test nsingular(res) == 0
        @test length(at_infinity(res)) == nat_infinity(res) == 4
        @test isempty(failed(res))
        @test nfailed(res) == 0
        @test nexcess_solutions(res) == 0
        @test !isempty(sprint(show, statistics(res)))

        # singular result
        @var x y
        g = System([(29 / 16) * x^3 - 2 * x * y, x^2 - y])
        res = solve(g; start_system = :total_degree)
        @test startswith(sprint(show, res), "Result with 3 solutions")
        @test seed(res) isa UInt32
        test_treeviews(res)
        @test !isempty(sprint(show, statistics(res)))
        @test !isempty(sprint(show, res))
    end

    @testset "composition" begin
        @var a b c x y z u v
        e = System([u + 1, v - 2])
        f = System([a * b - 2, a * c - 1])
        g = System([x + y, y + 3, x + 2])

        res = solve(e ∘ f ∘ g; start_system = :total_degree)
        @test nsolutions(res) == 2

        res = solve(e ∘ f ∘ g; start_system = :polyhedral)
        @test nsolutions(res) == 2
    end

    @testset "paths to track" begin
        @var x y
        f = System([2y + 3 * y^2 - x * y^3, x + 4 * x^2 - 2 * x^3 * y])

        @test paths_to_track(f; start_system = :total_degree) == 16
        @test paths_to_track(f; start_system = :polyhedral) == 8
        @test paths_to_track(f; start_system = :polyhedral, only_non_zero = true) == 3
        @test paths_to_track(f) == 8
        @test_deprecated bezout_number(f) == 16
        @test mixed_volume(f) == 3
    end

    @testset "solve (parameter homotopy)" begin
        # affine
        @var x a y b
        F = System([x^2 - a, x * y - a + b], [x, y], [a, b])
        s = [1, 1]
        res = solve(F, [s]; start_parameters = [1, 0], target_parameters = [2, 4])
        @test nsolutions(res) == 1
        res = solve(
            ModelKitSystem(F),
            [s];
            start_parameters = [1, 0],
            target_parameters = [2, 4],
            threading = false,
        )
        @test nsolutions(res) == 1

        @var x a y b
        F = System([x^2 - a], [x, y], [a, b])
        s = [1, 1]
        @test_throws FiniteException(1) solve(
            F,
            [s];
            start_parameters = [1, 0],
            target_parameters = [2, 4],
        )

        # proj
        @var x a y b z
        F_proj = System([x^2 - a * z^2, x * y + (b - a) * z^2], [x, y, z], [a, b])
        s = [1, 1, 1]
        res = solve(F_proj, [s]; start_parameters = [1, 0], target_parameters = [2, 4])
        @test nsolutions(res) == 1
        res = solve(ModelKitSystem(F_proj), [s]; p₁ = [1, 0], p₀ = [2, 4])
        @test nsolutions(res) == 1

        F_proj_err = System([x * y + (b - a) * z^2], [x, y, z], [a, b])
        @test_throws FiniteException solve(F_proj_err, [s]; p₁ = [1, 0], p₀ = [2, 4])

        # multi-proj
        @var x y v w a b
        F_multi_proj = System(
            [x * y - a * v * w, x^2 - b * v^2],
            parameters = [a, b],
            variable_groups = [[x, v], [y, w]],
        )
        S = [
            [
                -1.808683149843597 + 0.2761582523875564im,
                -0.9043415749217985 + 0.1380791261937782im,
                -0.0422893850686111 - 0.7152002569359284im,
                -0.0422893850686111 - 0.7152002569359283im,
            ],
            [
                -0.36370464807054353 + 0.6777414371333245im,
                0.18185232403527177 - 0.33887071856666223im,
                -0.3348980281838583 - 0.7759382220656511im,
                0.3348980281838583 + 0.7759382220656511im,
            ],
        ]
        res = solve(F_multi_proj, S; start_parameters = [2, 4], target_parameters = [3, 5])
        @test nsolutions(res) == 2
        res = solve(ModelKitSystem(F_multi_proj), S; p₁ = [2, 4], p₀ = [3, 5])
        @test nsolutions(res) == 2

        F_multi_proj_err = System(
            [x * y - a * v * w],
            parameters = [a, b],
            variable_groups = [[x, v], [y, w]],
        )
        @test_throws FiniteException(1) solve(F_multi_proj_err, S; p₁ = [2, 4], p₀ = [3, 5])
    end

    @testset "solve (Homotopy)" begin
        @var x a y b
        F = System([x^2 - a, x * y - a + b], [x, y], [a, b])
        s = [1, 1]
        H = ParameterHomotopy(F, [1, 0], [2, 4])
        res = solve(H, [s])
        @test nsolutions(res) == 1
    end

    @testset "solve (start target)" begin
        @var x a y b
        f = System([x^2 - a, x * y - a + b], parameters = [a, b])
        s = [1, 1]
        res = solve(f, f, [s]; start_parameters = [1, 0], target_parameters = [2, 4])
        @test nsolutions(res) == 1

        G = FixedParameterSystem(f, [1, 0])
        F = FixedParameterSystem(f, [2, 4])
        res = solve(G, F, [s];)
        @test nsolutions(res) == 1
    end

    @testset "solve (Vector{Expression})" begin
        @var x a y b
        F = [x^2 - a, x * y - a + b]
        s = [1, 1]
        res = solve(
            F,
            [s];
            parameters = [a, b],
            start_parameters = [1, 0],
            target_parameters = [2, 4],
        )
        @test nsolutions(res) == 1
    end

    @testset "solve (DynamicPolynomials)" begin
        @polyvar x y
        # define the polynomials
        f₁ = (x^4 + y^4 - 1) * (x^2 + y^2 - 2) + x^5 * y
        f₂ = x^2 + 2x * y^2 - 2 * y^2 - 1 / 2
        result = solve([f₁, f₂])
        @test nsolutions(result) == 18

        @polyvar x a y b
        F = [x^2 - a, x * y - a + b]
        s = [1, 1]
        res = solve(
            F,
            [s];
            parameters = [a, b],
            start_parameters = [1, 0],
            target_parameters = [2, 4],
        )
        @test nsolutions(res) == 1
    end

    @testset "change parameters" begin
        @var x a y b
        F = System([x^2 - a, x * y - a + b]; parameters = [a, b])
        s = [1.0, 1.0 + 0im]
        S = solver(F, generic_parameters = [2.2, 3.2])
        start_parameters!(S, [1, 0])
        target_parameters!(S, [2, 4])
        @test is_success(track(S, s))
    end

    @testset "solve (threading)" begin
        res = solve(cyclic(7), threading = true, show_progress = false)
        @test nsolutions(res) == 924
    end

    @testset "stop early callback" begin
        @var x
        first_result = nothing
        results = solve(
            [(x - 3) * (x + 6) * (x + 2)];
            stop_early_cb = r -> begin
                first_result = r
                true
            end,
            threading = false,
            show_progress = false,
            start_system = :total_degree,
        )
        @test length(results) == 1
        @test first(results) === first_result

        nresults = 0
        result = solve(
            [(x - 3) * (x + 6) * (x + 2)],
            stop_early_cb = r -> (nresults += 1) == 2,
            start_system = :total_degree,
            show_progress = false,
            threading = false,
        )
        @test length(result) == 2

        # threading
        @var x y z
        first_result = nothing
        # this has 5^3 = 125 solutions, so we should definitely stop early if we
        # have less than 64 threads
        results = solve(
            [
                (x - 3) * (x + 6) * (x + 2) * (x - 2) * (x + 2.5),
                (y + 2) * (y - 2) * (y + 3) * (y + 5) * (y - 1),
                (z + 2) * (z - 2) * (z + 3) * (z + 5) * (z - 2.1),
            ];
            stop_early_cb = r -> begin
                first_result = r
                true
            end,
            threading = true,
            show_progress = false,
            start_system = :total_degree,
        )
        @test length(results) < 125
    end

    @testset "Many parameters solver" begin
        # Setup
        @var x y
        f = x^2 + y^2 - 1

        @var a b c
        l = a * x + b * y + c
        F = [f, l]

        # Compute start solutions S₀ for given start parameters p₀
        p₀ = randn(ComplexF64, 3)
        S₀ = solutions(solve(subs(F, [a, b, c] => p₀)))
        # The parameters we are intersted in
        params = [rand(3) for i = 1:100]

        result1 = solve(
            F,
            S₀,
            ;
            start_parameters = p₀,
            target_parameters = params,
            parameters = [a, b, c],
            threading = true,
        )
        @test typeof(result1) == Vector{Tuple{Result,Vector{Float64}}}
        result1 = solve(
            F,
            S₀,
            ;
            start_parameters = p₀,
            target_parameters = params,
            parameters = [a, b, c],
            show_progress = false,
            threading = false,
        )
        @test typeof(result1) == Vector{Tuple{Result,Vector{Float64}}}

        # Only keep real solutions
        result2 = solve(
            F,
            S₀,
            ;
            start_parameters = p₀,
            target_parameters = params,
            parameters = [a, b, c],
            transform_result = (r, p) -> real_solutions(r),
            threading = true,
        )
        @test typeof(result2) == Vector{Vector{Vector{Float64}}}
        @test !isempty(result2)

        # Now instead of an Array{Array{Array{Float64,1},1},1} we want to have an
        # Array{Array{Float64,1},1}
        result3 = solve(
            F,
            S₀,
            ;
            start_parameters = p₀,
            target_parameters = params,
            parameters = [a, b, c],
            transform_result = (r, p) -> real_solutions(r),
            flatten = true,
            threading = false,
        )
        @test typeof(result3) == Vector{Vector{Float64}}
        @test !isempty(result3)

        # The passed `params` do not directly need to be the target parameters.
        # Instead they can be some more concrete informations (e.g. an index)
        # and we can them by using the `transform_parameters` method
        result4 = solve(
            F,
            S₀,
            ;
            start_parameters = p₀,
            target_parameters = 1:100,
            parameters = [a, b, c],
            transform_result = (r, p) -> (real_solutions(r), p),
            transform_parameters = _ -> rand(3),
        )
        @test typeof(result4) == Vector{Tuple{Vector{Vector{Float64}},Int64}}
    end
end


function observables(R, xroot, lroot, ϵroot, κ, norm, coef, energ, j, mass)

    @info "Extracting observables..."

    sinϕ = @. mass/√(energ^2+mass^2)
    cosϕ = @. √(1-sinϕ^2)
    vnjl² = @. (1-energ/√(energ^2+mass^2))/2

    # Average mass computed with vnjl² and sinϕ
    Mvavg = sum(@. mass*j*vnjl²)/sum(@. j*vnjl²)
    Msavg = sum(@. mass*j*sinϕ) / sum(@. j*sinϕ)

    # Average mass computed as an integral
    Mcoef(r) = sum(@. 1/2*(2*j+1)/(4*π)*sinϕ*norm^2*(sphericalbessel(xroot/R*r, lroot)^2+sphericalbessel(xroot/R*r, lroot+ϵroot)^2))
    M(r) = 4*Nc*Nf*G*ħ^3*Mcoef(r)
    Mint(r) = 4*π*r^2*M(r)
    Mavg = quadgk(Mint, 0, R)[1]
    Mavg = Mavg / (4*π*R^3/3)

    # Chiral condensate extracted from the radial mass
    # For dubious reasons, for certain values it returns a complex value with very small imaginary part...
    ψ̅ψ(r) = - complex(1/(2*G) * M(r))^(1/3)
    # ψ̅ψ(r) = - (1/(2*G) * M(r))^(1/3)
    ψ̅ψint(r) = 4*π*r^2*ψ̅ψ(r)
    ψ̅ψavg = quadgk(ψ̅ψint, 0, R)[1]
    ψ̅ψavg = ψ̅ψavg / (4*π*R^3/3)

    # Trying to brute force get rid of the imaginary part...
    if imag(ψ̅ψavg) != 0
        ψ̅ψavg = NaN
    else
        ψ̅ψavg = real(ψ̅ψavg)
    end

    # Radial dependence of mass
    Mr = r -> M(r)
    # Radial dependence of chiral condensate
    qq̅ = r -> ψ̅ψ(r)

    # Energy density of the chiral condensate
    ε₁ = - sum(@. j*energ*cosϕ)
    ε₂ = - sum([(coef[k₁, k₂]*sinϕ[k₁]*sinϕ[k₂]) for k₁ in 1:κ for k₂ in 1:κ])
    ε₂ = ε₂ *2*Nc*Nf*G*ħ^3
    ε = (ε₁ + ε₂) / (4*π*R^3/3)

    @info "Done!"
    @debug "Average mass: " Mavg
    @debug "Average value of chiral condensate: " ψ̅ψavg
    @debug "Energy density of a chiral pair: " ε

    return Mvavg, Msavg, Mr, Mavg, qq̅, ψ̅ψavg, ε

end


#    Copyright (c) 2021 J.A. Duffek
#    Copyright (c) 2000 D.M. Spink
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, see <http://www.gnu.org/licenses/>.

"""
    nrbrevolve(crv::NURBS1D{I,F},vec::Vector{F},
               pnt::Vector{F}=[0.0;0.0;0.0],theta::F=2.0*π
              )::NURBS2D{I,F} where {I<:Integer,F<:AbstractFloat}

Construct a NURBS surface by revolving the profile NURBS curve around an axis
defined by a point and vector.

# Arguments:
- `crv`: NURBS curve to revolve, see [`nrbmak`](@ref).
- `vec`: Vector defining the direction of the rotation axis.
- `pnt`: Coordinate of the point used to define the axis of rotation.
- `theta`: Angle to revolve the curve, default `2*π`.

# Output:
- `srf`: Constructed surface.

# Examples:
```julia
julia> srf = nrbrevolve(crv,vec)
julia> srf = nrbrevolve(crv,vec,pnt)
julia> srf = nrbrevolve(crv,vec,pnt,ang)
```

Construct a sphere by rotating a semicircle around a x-axis.
```julia
julia> crv = nrbcirc(1.0,[0.0;0.0;0.0],0.0,pi);
julia> srf = nrbrevolve(crv,[1.0;0.0;0.0],[0.0;0.0;0.0]);
julia> nrbplot(srf,[20;20]);
```

# The algorithm:

1. vectrans the point to the origin `[0.0;0.0;0.0]`
2. rotate the vector into alignment with the z-axis
3. for each control point along the curve
   1. determine the radius and angle of control point to the z-axis
   2. construct a circular arc in the x-y plane with this radius and start
      angle and sweep angle theta
   3. combine the arc and profile, coefs and weights.
   4. next control point
4. rotate and vectrans the surface back into position by reversing 1 and 2.
"""
function nrbrevolve(crv::NURBS1D{I,F},vec::Vector{F},
                    pnt::Vector{F}=[0.0;0.0;0.0],theta::F=2.0*π
                   )::NURBS2D{I,F} where {I<:Integer,F<:AbstractFloat}
# TODO: optional arguments in julia are always at the end i switched vec and
#       pnt argument > examine if that did change the behaviour in this toolbox

if length(vec) != 3 || length(pnt) !=3
  throw(ArgumentError("All point and vector coordinates must be 3D!"));
end # if

# Translate and rotate the curve into alignment with the z-axis
T  = vectrans(-pnt);
angx = vecangle(vec[1],vec[3]);
RY = vecroty(-angx);
vectmp = RY * [vecnorm_toolbox(vec);1.0];
angy = vecangle(vectmp[2],vectmp[3]);
RX = vecrotx(angy);
curve = nrbtform(crv,RX*RY*T);

# Construct an arc
arc = nrbcirc(1.0,[0.0;0.0;0.0],0.0,theta);

# Construct the surface
coefs = zeros(F,(4,arc.number,curve.number));
angle = vecangle(curve.coefs[2:2,:],curve.coefs[1:1,:]);
radius = vecmag(curve.coefs[1:2,:]);
for i in 1:curve.number
  coefs[:,:,i] = vecrotz(angle[i]) *
                 vectrans([0.0;0.0;curve.coefs[3,i]]) *
                 vecscale([radius[i];radius[i]]) * arc.coefs;
  coefs[4,:,i] = coefs[4,:,i] * curve.coefs[4,i];
end # for i
surf = nrbmak(coefs,[arc.knots,curve.knots]);

# Rotate and vectrans the surface back into position
T = vectrans(pnt);
RX = vecrotx(-angy);
RY = vecroty(angx);
return nrbtform(surf,T*RY*RX);
end # nrbrevolve


push!(LOAD_PATH, string(pwd(),"/src"))

using Revise
using SpecialFunctions
using ModifiedStokesSolver
using Base.Test

@testset "NewPowerSeriesTests" begin

    z = 0.99
    #(T1Sp, T1Lp, T2Sp, T2Lp, T3Sp, T3Lp) =
    Tp = ModifiedStokesSolver.Ti_split_pow(z)
    #(T1Sd, T1Ld, T2Sd, T2Ld, T3Sd, T3Ld) =
    Td = ModifiedStokesSolver.Ti_split_direct(z)
    Ep = Tp .- Td
    @test maximum(abs.(Ep./Td)) < 9e-14

    z = 1.8
    #(T1Sp, T1Lp, T2Sp, T2Lp, T3Sp, T3Lp) =
    Tp = ModifiedStokesSolver.Ti_split_pow(z,true)
    #(T1Sd, T1Ld, T2Sd, T2Ld, T3Sd, T3Ld) =
    Td = ModifiedStokesSolver.Ti_split_direct(z,true)
    Ep = Tp .- Td
    @test maximum(abs.(Ep./Td)) < 9e-14

end

# @testset "PowerSeries" begin

# # Series evaluator
# function evalseries(fn, z, nstart)
#     y = 0.0
#     dy = 1.0
#     n = nstart
#     while abs(dy)>=eps(y)
#         dy = fn(z, n)*z^(2*n)
#         y += dy
#         n += 1
#     end
#     return y
# end
# # Definitions for modified Bessel functions
# i0(z,n) = 1/(4^n*factorial(n)^2)
# i1(z, n) = 1/(2*4^n*factorial(n)*factorial(n+1))
# k0(z,n) = ((digamma(n+1) + log(2) - log(z)) / (4^n * factorial(n)^2))
# k1(z, n) = (2*log(z) - 2*log(2) - digamma(n+1) - digamma(n+2)) /
#     (4^(n+1)*factorial(n)*factorial(n+1))
# # Splits
# k0S(z,n) = ((digamma(n+1) + log(2)) / (4^n * factorial(n)^2))
# k0L(z,n) = (-1) / (4^n * factorial(n)^2)
# k1S(z, n) = ( - 2*log(2) - digamma(n+1) - digamma(n+2)) /
#     (4^(n+1)*factorial(n)*factorial(n+1))
# k1L(z, n) = (2) / (4^(n+1)*factorial(n)*factorial(n+1))

# # Direct power series for split K0, K1
# # Returns: (K0S, K0L, K1S, K1L, K1C)
# function besselpowsplit(z)
#     # Outputs
#     K0S = 0
#     K0L = 0
#     K1S = 0
#     K1L = 0
#     K1C = 1
#     # Recursion variables
#     n = 0
#     z2 = z*z
#     z2n = 1; # z^(2*n)
#     digamma_np1 = -eulergamma # digamma(n+1)
#     digamma_np2 = -eulergamma + 1
#     ifact_n = 1 # 1/factorial(n)
#     ifact_np1 = 1
#     i4n = 1 # 1/4^n
#     i4np1 = 1/4
#     log2 = log(2)
#     # Loop
#     maxit = 20
#     while n<maxit
#         k0n_denom = i4n*ifact_n*ifact_n
#         k1n_denom = i4np1*ifact_n*ifact_np1
#         k0Ln = (-1)*k0n_denom
#         k0Sn = (digamma_np1 + log2)*k0n_denom
#         k1Ln = (2)*k1n_denom
#         k1Sn = (-2*log2-digamma_np1-digamma_np2)*k1n_denom
#         K0S += k0Sn*z2n
#         K0L += k0Ln*z2n
#         K1S += k1Sn*z2n*z
#         K1L += k1Ln*z2n*z
#         # Update recursions
#         n += 1
#         z2n *= z2;
#         inp1 = 1/(n+1)
#         (digamma_np1, digamma_np2) = (digamma_np2, digamma_np2+inp1)
#         (ifact_n, ifact_np1) = (ifact_np1, ifact_np1*inp1)
#         (i4n, i4np1) = (i4np1, i4np1*0.25)
#     end
#     return (K0S, K0L, K1S, K1L, K1C)
# end

# # Tests
# z = 1/10

# ref = besseli(0, z)
# I0test = evalseries(i0, z, 0)
# @test abs((ref-I0test)/ref) < 1e-14
# ref = besseli(1, z)
# I1test = z*evalseries(i1, z, 0)
# @test abs((ref-I1test)/ref) < 1e-14
# ref = besselk(0, z)
# K0test = evalseries(k0, z, 0)
# K0split = evalseries(k0S, z, 0) + log(z)*evalseries(k0L, z, 0)
# @test abs((ref-K0test)/ref) < 1e-14
# @test abs((ref-K0split)/ref) < 1e-14
# ref = besselk(1, z)
# K1test = 1/z + z*evalseries(k1, z, 0)
# K1split = 1/z + z*evalseries(k1S, z, 0) + z*log(z)*evalseries(k1L, z, 0)
# @test abs((ref-K1test)/ref) < 1e-14
# @test abs((ref-K1split)/ref) < 1e-14
# # Test direct split power series
# ref = (besselk(0, z), besselk(1, z))
# (K0S, K0L, K1S, K1L, K1C) = besselpowsplit(z)
# split = (log(z)*K0L+K0S, K1S+log(z)*K1L+K1C/z)
# @test maximum(@. abs((ref-split)/ref) ) < 1e-14

# # Evaluate T, high precision reference from Mathematica
# for n=1:2
#     if n==1
#         z = 0.001
#         T1exact = 0.00028941085628129279453492498775927
#         T2exact = -318.30984639510755467178629422316
#         T3exact = -0.00059871605412101857968611091116825
#     end
#     if n==2
#         z = 0.99
#         T1exact = 0.023658849757364503594228422585464
#         T2exact = -0.28952030153253490875314823065265
#         T3exact = -0.063320117528153811164290571635020
#     end
#     # z = 1
#     # (T1exact, T2exact, T3exact) =
#     #     (0.023620976507380996090769874829016,
#     #      -0.28607692796680641340443170447384,
#     #      -0.063358432123254121248207660793628)

#     # z = 2
#     # (T1exact, T2exact, T3exact) =
#     #     (0.016930005787967333888517613314308,
#     #      -0.11224071607388210664375507421808,
#     #      -0.057317125084941282339599571205834)

#     println("z = ", z)
#     println()

#     K0 = besselk(0, z)
#     K1 = besselk(1, z)
#     K0S = K0 + besseli(0, z)*log(z)
#     K1S = K1 - besseli(1, z)*log(z) - 1/z
#     I0 = besseli(0,z)
#     I1 = besseli(1,z)

#     ( T1Spow, T1Lpow,
#       T2Spow, T2Lpow, T2Cpow,
#       T3Spow, T3Lpow) = ModifiedStokesSolver.Ti_split_pow(z)

#     ## T1
#     T1 = -( z^2*(2*K0 + z*K1) + 4*(z*K1 - 1) ) / (2*pi*z*z^2)
#     t1(z, n) = 2*k0(z, n) + 4*k1(z, n) + k1(z, n-1)
#     t1S(z, n) = 2*k0S(z, n) + 4*k1S(z, n) + k1S(z, n-1)
#     t1L(z, n) = 2*k0L(z, n) + 4*k1L(z, n) + k1L(z, n-1)
#     T1series = -evalseries(t1,z,1)/(2*pi*z)
#     T1Sseries = -evalseries(t1S,z,1)/(2*pi*z)
#     T1Lseries = -evalseries(t1L,z,1)/(2*pi*z)
#     T1splitseries = (T1Sseries + T1Lseries*log(z))
#     T1pow = (T1Spow + T1Lpow*log(z))

#     #quit()

#     T1S = -(2*z*K0S + (z^2+4)*K1S + z) / (2*pi*z^2)
#     T1L = (2*z*I0 - (z^2+4)*I1) / (2*pi*z^2)
#     T1split = (T1S + T1L*log(z))

#     println("== T1")
#     println("exact  = ", T1exact)
#     println("series = ", T1series, "  relerr = ", (T1series-T1exact)/T1exact)
#     T1spliterr = (T1splitseries-T1exact)/T1exact
#     println("spl+ser= ", T1splitseries, "  relerr = ", T1spliterr)
#     #@test abs(T1spliterr)<1e-14
#     println("regular= ", T1, "  relerr = ", (T1-T1exact)/T1exact)
#     println("reg+spl= ", T1split, "  relerr = ", (T1split-T1exact)/T1exact)
#     #@test abs( (T1series-T1exact)/T1exact )<1e-14
#     T1powerr = (T1pow-T1exact)/T1exact
#     println("final  = ", T1pow, "  relerr = ", T1powerr)
#     #@test abs(T1powerr)<1e-14
#     println()

#     ## T2
#     T2 = (4*z^2*K0 + (z^2+8)*z*K1 - 8) / (pi*z^3)
#     t2(z, n) = 4*k0(z, n) + 8*k1(z, n) + k1(z, n-1)
#     T2series = -1/(pi*z) + evalseries(t2,z,1)/(pi*z)
#     println("== T2")
#     println("exact  = ", T2exact)
#     println("series = ", T2series, "  relerr = ", (T2series-T2exact)/T2exact)
#     println("regular= ", T2, "  relerr = ", (T2-T2exact)/T2exact)
#     @test abs( (T2series-T2exact)/T2exact )<1e-14
#     T2pow = (T2Spow + T2Lpow*log(z) + T2Cpow/z)
#     T2powerr = (T2pow-T2exact)/T2exact
#     println("final  = ", T2pow, "  relerr = ", T2powerr)
#     @test abs(T2powerr)<1e-14
#     println()

#     ## T3
#     T3 = (z*K1 - 1) / (2*pi*z)
#     t3(z, n) = k1(z, n)
#     T3series = z*evalseries(t3,z,0)/(2*pi)
#     println("== T3")
#     println("exact  = ", T3exact)
#     println("series = ", T3series, "  relerr = ", (T3series-T3exact)/T3exact)
#     println("regular= ", T3, "  relerr = ", (T3-T3exact)/T3exact)
#     @test abs( (T3series-T3exact)/T3exact )<1e-14
#     T3pow = (T3Spow + T3Lpow*log(z))
#     T3powerr = (T3pow-T3exact)/T3exact
#     println("final  = ", T3pow, "  relerr = ", T3powerr)
#     @test abs(T3powerr)<1e-14

#     println()
# end

#end # End testset


"""
    Order0()

The `Order0` method is engineered to be a more robust, though possibly
slower, alternative to the other derivative-free root-finding
methods. The implementation roughly follows the algorithm described in
*Personal Calculator Has Key to Solve Any Equation ``f(x) = 0``*, the
SOLVE button from the
[HP-34C](http://www.hpl.hp.com/hpjournal/pdfs/IssuePDFs/1979-12.pdf).
The basic idea is to use a secant step. If along the way a bracket is
found, switch to a bracketing algorithm, using `AlefeldPotraShi`.  If the secant
step fails to decrease the function value, a quadratic step is used up
to ``4`` times.

This is not really ``0``-order: the secant method has order
``1.6...`` [Wikipedia](https://en.wikipedia.org/wiki/Secant_method#Comparison_with_other_root-finding_methods)
and the the bracketing method has order
``1.6180...`` [Wikipedia](http://www.ams.org/journals/mcom/1993-61-204/S0025-5718-1993-1192965-2/S0025-5718-1993-1192965-2.pdf)
so for reasonable starting points and functions, this algorithm should be
superlinear, and relatively robust to non-reasonable starting points.

"""
struct Order0 <: AbstractSecant end

# special case Order0 to be hybrid
function init(
    𝑭𝑿::ZeroProblem,
    M::Order0,
    p′=nothing;
    p=nothing,
    verbose::Bool=false,
    tracks=NullTracks(),
    kwargs...,
)
    p = p′ === nothing ? p : p′
    init(𝑭𝑿, Secant(), AlefeldPotraShi(); p=p, verbose=verbose, tracks=tracks, kwargs...)
end

## When passing 2 methods, any parameters must be passed as a named argument through
## the keyword p
function init(
    𝑭𝑿::ZeroProblem,
    M::AbstractUnivariateZeroMethod,
    N::AbstractBracketing;
    p=nothing,
    verbose::Bool=false,
    tracks=NullTracks(),
    kwargs...,
)
    F = Callable_Function(M, 𝑭𝑿.F, p)
    state = init_state(M, F, 𝑭𝑿.x₀)
    options = init_options(M, state; kwargs...)
    l = Tracks(verbose, tracks, state)
    incfn(l, initial_fncalls(M))
    ZeroProblemIterator(M, N, F, state, options, l)
end

# Robust version using some tricks: idea from algorithm described in
# [The SOLVE button from the
# HP-34]C(http://www.hpl.hp.com/hpjournal/pdfs/IssuePDFs/1979-12.pdf).
# * use bracketing method if one identifed
# * limit steps so as not too far or too near the previous one
# * if not decreasing, use a quad step upto 4 times to bounce out of trap, if possible
# First uses M, then N if bracket is identified
function solve!(𝐙::ZeroProblemIterator{𝐌,𝐍}; verbose=false) where {𝐌,𝐍<:AbstractBracketing}
    M, N, F, state, options, l = 𝐙.M, 𝐙.N, 𝐙.F, 𝐙.state, 𝐙.options, 𝐙.logger

    incfn(l, 2)
    log_step(l, M, state; init=true)

    quad_ctr = 0
    flag = :not_converged
    ctr = 0
    α = NaN * state.xn1
    while true
        ctr += 1
        flag, converged = assess_convergence(M, state, options)

        converged && break
        ctr >= options.maxevals && break

        state0 = state
        state0, stopped = update_state(M, F, state0, options) # state0 is proposed step

        ## did we find a zero or a bracketing interval?
        if iszero(state0.fxn1)
            state = state0
            break
        elseif sign(state0.fxn0) * sign(state0.fxn1) < 0
            log_step(l, M, state0)
            !isa(l, NullTracks) && log_message(
                l,
                "Used bracketing method $N on  [$(min(state0.xn0,state0.xn1)),$(max(state0.xn0,state0.xn1))]",
            )

            Fₙ = Callable_Function(N, F)
            stateₙ = init_state(N, state0, Fₙ) # save function calls by using state0 values
            optionsₙ = init_options(N, stateₙ)
            α = solve!(init(N, Fₙ, stateₙ, optionsₙ, l))
            break
        end

        ## did we move too far?
        adj = false
        r, a, b = state0.xn1, state.xn0, state.xn1
        Δr = abs(r - b)
        Δx = abs(b - a)
        ts, TB = 1e-3, 1e2 # too small, too big
        if Δr >= TB * Δx
            adj = true
            r = b + sign(r - b) * TB * Δx  ## too big
        elseif Δr <= ts * Δx
            adj = true
            r = b + sign(r - b) * ts * Δx
        end

        @set! state0.xn1 = r
        @set! state0.fxn1 = first(F(r))
        incfn(l)

        # a sign change after shortening?
        if sign(state.fxn1) * sign(state0.fxn1) < 0
            log_step(l, M, state)
            a, b = state.xn1, state0.xn1
            fa, fb = state.fxn1, state0.fxn1
            !isa(l, NullTracks) && log_message(
                l,
                "Used bracketing method $N on  [$a,$b]",
            )

            Fₙ = Callable_Function(N, F)
            stateₙ = init_state(N, Fₙ, a, b, fa, fb)
            optionsₙ = init_options(N, stateₙ)
            α = solve!(init(N, Fₙ, stateₙ, optionsₙ, l))
            break
        end

        ## did we improve?
        if adj || abs(state0.fxn1) < abs(state.fxn1)
            if isnan(state0.xn1) ||
               isnan(state0.fxn1) ||
               isinf(state0.xn1) ||
               isinf(state0.fxn1)
                break
            end
            state = state0
            log_step(l, M, state)
            quad_ctr = 0
            continue
        end

        ## try quad_vertex, unless that has gotten old
        if quad_ctr > 4
            state = state0
            break
        else
            quad_ctr += 1
            r = quad_vertex(
                state0.xn1,
                state0.fxn1,
                state.xn1,
                state.fxn1,
                state.xn0,
                state.fxn0,
            )

            if isnan(r) || isinf(r)
                state = state0
            else
                fr = F(r)
                incfn(l)

                @set! state0.xn1 = r
                @set! state0.fxn1 = fr
                state = state0
            end
        end
        log_step(l, M, state)
    end

    val, stopped = assess_convergence(M, state, options)
    log_convergence(l, val)
    log_method(l, M)
    log_nmethod(l, N)
    log_last(l, α)
    verbose && display(l)

    isnan(α) ? decide_convergence(M, F, state, options, flag) : α
end

function find_zero(
    fs,
    x0,
    M::AbstractUnivariateZeroMethod,
    N::AbstractBracketing;
    verbose=false,
    kwargs...,
)
    𝐏 = ZeroProblem(fs, x0)
    solve!(init(𝐏, M, N; verbose=verbose, kwargs...), verbose=verbose)
end

# Switch to bracketing method
# deprecate soon, not used
function run_bisection(N::AbstractBracketing, f, ab, state)
    steps, fnevals = state.steps, state.fnevals
    f = Callable_Function(N, f)
    init_state!(state, N, f; clear=true)
    find_zero(N, f, state, init_options(N, state))
    a, b = _extrema(ab)
    u, v = a > b ? (b, a) : (a, b)
    state.message *= "Bracketing used over ($u, $v), those steps not shown. "
    return nothing
end


# Orbit Space Covariance

function get_global_covariance(c1::EPRCoordinate, ot1::Symbol,
                               c2::EPRCoordinate, ot2::Symbol,
                               p::AbstractVector, norms::AbstractVector)

    x = S3(c1.energy, c1.pitch, c1.r)
    y = S3(c2.energy, c2.pitch, c2.r)
    mu = (x .- y)./norms

    GO = Set((:trapped,:potato,:stagnation,:co_passing))
    same_ot = ((ot1 in GO) && (ot2 in GO)) || (ot1 == ot2)

    same_orb = c1 == c2

    np = length(p)
    if np == 2
        k = same_ot*(p[1]^2)*exp(-0.5*dot(mu,mu)/(p[2]^2)) + same_orb*1e-6
    elseif np == 3
        k = same_ot*(p[1]^2)*exp(-0.5*dot(mu,mu)/(p[2]^2)) + same_orb*p[3]^2
    elseif np == 4
        Σ_rho_inv = Diagonal(p[2:4].^(-2))
        k = same_ot*(p[1]^2)*exp(-0.5*dot(mu,Σ_rho_inv*mu)) + same_orb*1e-6
    elseif np == 5
        Σ_rho_inv = Diagonal(p[2:4].^(-2))
        k = same_ot*(p[1]^2)*exp(-0.5*dot(mu,Σ_rho_inv*mu)) + same_orb*p[5]^2
    else
        error("Unsupported number of parameters")
    end

    return k
end

function get_global_covariance(o1::Orbit, o2::Orbit, p::Vector; norms=S3(1.0,1.0,1.0))
    c1 = o1.coordinate
    c2 = o2.coordinate
    ot1 = o1.class
    ot2 = o2.class
    return get_global_covariance(c1,ot1,c2,ot2,p,norms)
end

function get_global_covariance_matrix(orbs, p::Vector; norms=S3(1.0,1.0,1.0))
    n = length(orbs)
    Σ = zeros(n,n)
    for i=1:n
        oi = orbs[i]
        for j = i:n
            oj = orbs[j]
            k = get_global_covariance(oi.coordinate, oi.class,
                                      oj.coordinate, oj.class,
                                      p, norms)
            Σ[i,j] = k
            Σ[j,i] = k
        end
    end
    return Σ
end

function get_global_covariance_matrix(orbs1::Vector, orbs2::Vector, p::Vector; norms=S3(1.0,1.0,1.0))
    n1 = length(orbs1)
    n2 = length(orbs2)
    Σ = zeros(n1,n2)
    for i=1:n1
        oi = orbs1[i]
        for j = 1:n2
            oj = orbs2[j]
            k = get_global_covariance(oi.coordinate, oi.class,
                                      oj.coordinate, oj.class,
                                      p, norms)
            Σ[i,j] = k
        end
    end
    return Σ
end

function epr_cov(energy, pitch, r, orbit_type, p; independent=false, norm = ones(3))
    norb = length(energy)
    np = length(p)

    GO = Set((:trapped,:potato,:stagnation,:co_passing))
    Sigma_F = zeros(norb,norb)
    for i=1:norb
        x = [energy[i], pitch[i], r[i]]
        ox = orbit_type[i]
        for j=i:norb
            y = [energy[j], pitch[j], r[j]]
            oy = orbit_type[j]
            if independent
                same_ot = ox == oy
            else
                same_ot = ((ox in GO) && (oy in GO)) || (ox == oy)
            end
            mu = (x .- y)./norm
            if np == 2
                k = same_ot*(p[1]^2)*exp(-0.5*dot(mu,mu)/(p[2]^2)) + (i==j)*1e-6
            elseif np == 3
                k = same_ot*(p[1]^2)*exp(-0.5*dot(mu,mu)/(p[2]^2)) + (i==j)*p[3]^2
            elseif np == 4
                Σ_rho_inv = Diagonal(p[2:4].^(-2))
                k = same_ot*(p[1]^2)*exp(-0.5*dot(mu,Σ_rho_inv*mu)) + (i==j)*1e-6
            elseif np == 5
                Σ_rho_inv = Diagonal(p[2:4].^(-2))
                k = same_ot*(p[1]^2)*exp(-0.5*dot(mu,Σ_rho_inv*mu)) + (i==j)*p[5]^2
            else
                error("Unsupported number of parameters")
            end

	    Sigma_F[i,j] = k
            Sigma_F[j,i] = k
        end
    end

    return Sigma_F
end

# EPRZ Covariance
struct RepeatedBlockDiagonal{T,N,A<:AbstractArray} <: AbstractArray{T,N}
    data::A
    n::Int
end

Base.size(A::RepeatedBlockDiagonal) = size(A.data).*A.n

function Base.getindex(A::RepeatedBlockDiagonal,I::Int...)
    i, j = I
    ni, nj = size(A.data)
    ii = mod(i-1, ni) + 1
    jj = mod(j-1, nj) + 1

    return div(i-1,ni) == div(j-1,nj) ? A.data[ii,jj] : zero(eltype(A.data))
end

Base.getindex(A::RepeatedBlockDiagonal,I::Int) = A[CartesianIndices(A)[I]]

function RepeatedBlockDiagonal(A::Matrix,n::Int)
    RepeatedBlockDiagonal{eltype(A),2,typeof(A)}(A,n)
end

Base.:(*)(A::RepeatedBlockDiagonal,a::T) where {T<:Number} = RepeatedBlockDiagonal(a*A.data,A.n)
Base.:(*)(a::T, A::RepeatedBlockDiagonal) where {T<:Number} = RepeatedBlockDiagonal(a*A.data,A.n)
Base.:(/)(A::RepeatedBlockDiagonal,a::T) where {T<:Number} = RepeatedBlockDiagonal(A.data/a,A.n)

function Base.:(*)(A::RepeatedBlockDiagonal , a::Union{SparseVector{S,T},SparseMatrixCSC{S,T}}) where {S,T}
    nr,nc = size(A.data)
    colptr = vcat(1:nr:(nr*nc),fill(nr*nc+1,nr*A.n + 1 - nc))
    colptr2 = Array{Int}(undef,length(colptr))
    rowval = repeat(1:nr,nc)
    rowval2 = Array{Int}(undef,length(rowval))
    K = SparseMatrixCSC(size(A)..., colptr, rowval, vec(A.data))

    p = Progress(A.n)
    channel = RemoteChannel(()->Channel{Bool}(A.n), 1)
    B = fetch(@sync begin
        @async while take!(channel)
            ProgressMeter.next!(p)
        end
        @async begin
            B = @distributed (+) for i=1:A.n
                circshift!(colptr2,colptr,(i-1)*nc)
                rowval2 .= rowval .+ (i-1)*nr
                K = SparseMatrixCSC(K.m,K.n,colptr2,rowval2, K.nzval)
                Ka = K*a
                put!(channel,true)
                Ka
            end
            put!(channel, false)
            B
        end
    end)
    return B
end

function Base.:(*)(a::Union{SparseVector{S,T},SparseMatrixCSC{S,T}}, A::RepeatedBlockDiagonal) where {S,T}
    nr,nc = size(A.data)
    colptr = vcat(1:nr:(nr*nc),fill(nr*nc+1,nr*A.n + 1 - nc))
    colptr2 = Array{Int}(undef,length(colptr))
    rowval = repeat(1:nr,nc)
    rowval2 = Array{Int}(undef,length(rowval))
    K = SparseMatrixCSC(size(A)..., colptr, rowval, vec(A.data))

    p = Progress(A.n)
    channel = RemoteChannel(()->Channel{Bool}(A.n), 1)
    B = fetch(@sync begin
        @async while take!(channel)
            ProgressMeter.next!(p)
        end
        @async begin
            B = @distributed (+) for i=1:A.n
                circshift!(colptr2,colptr,(i-1)*nc)
                rowval2 .= rowval .+ (i-1)*nr
                K = SparseMatrixCSC(K.m,K.n,colptr2,rowval2, K.nzval)
                aK = a*K
                put!(channel,true)
                aK
            end
            put!(channel, false)
            B
        end
    end)
    return B
end

function transform_eprz_cov(C,R)
    nr,nc = size(C.data)
    colptr = vcat(1:nr:(nr*nc),fill(nr*nc+1,nr*C.n + 1 - nc))
    colptr2 = Array{Int}(undef,length(colptr))
    rowval = repeat(1:nr,nc)
    rowval2 = Array{Int}(undef,length(rowval))
    K = SparseMatrixCSC(size(C)..., colptr, rowval, vec(C.data))
    S = @distributed (+) for i=1:C.n
        circshift!(colptr2,colptr,(i-1)*nc)
        rowval2 .= rowval .+ (i-1)*nr
        K = SparseMatrixCSC(K.m,K.n,colptr2,rowval2, K.nzval)
        Array((R*(R*K)'))
    end

    return S
end

function transform_eprz_cov2(C,R)
    nr,nc = size(C.data)
    Rnr = size(R)[1]
    Cs = sparse(C.data)

    S = @distributed (+) for i=1:nr:(C.n*nr)
        Ri = R[:,i:(i+nr-1)]
        if nnz(Ri) != 0
            Si = Ri*(Ri*Cs)'
        else
            Si = spzeros(Rnr,Rnr)
        end
        Si
    end

    return Array(S)
end

function transform_eprz_cov3(C,R)
    nr,nc = size(C.data)
    Rnr = size(R)[1]
    Cs = sparse(C.data)

    Ris = SparseMatrixCSC{Float64,Int64}[]
    for i=1:nr:(C.n*nr)
        Ri = R[:,i:(i+nr-1)]
        nnz(Ri) == 0 && continue
        push!(Ris,Ri)
    end

    Si = [zeros(Rnr,Rnr) for i=1:Threads.nthreads()]
    Threads.@threads for Ri in Ris
        id = Threads.threadid()
        Si[id] .= Si[id] .+ Ri*(Ri*Cs)'
    end

    return sum(Si)
end

function LinearAlgebra.inv(A::RepeatedBlockDiagonal)
    RepeatedBlockDiagonal(inv(A.data),A.n)
end

function ep_cov(energy, pitch, sigE, sigp)

    nenergy = length(energy)
    npitch = length(pitch)
    nep = nenergy*npitch

    Σ_ep = zeros(typeof(sigE),nep,nep)
    Σ_p_inv = SMatrix{2,2}(inv(Diagonal([sigE,sigp].^2)))
    subs = CartesianIndices((nenergy,npitch))
    @inbounds for i=1:nep
        ie,ip = Tuple(subs[i])
        x = SVector{2}(energy[ie],pitch[ip])
        for j=i:nep
            je,jp = Tuple(subs[j])
            y = SVector{2}(energy[je],pitch[jp])
            d = x .- y
            l = d'*Σ_p_inv*d
            Σ_ep[i,j] = exp(-0.5*l)
            Σ_ep[j,i] = Σ_ep[i,j]
        end
    end

    return Σ_ep
end

function ep_cov(energy, pitch, p::Vector)
    if length(p) == 2
        return ep_cov(energy, pitch, p[1], p[2])
    elseif length(p) == 3
        return (p[1]^2)*ep_cov(energy, pitch, p[2], p[3])
    else
        error("Unsupported number of parameters")
    end
end

function eprz_cov(energy, pitch, r, z, p::Vector)
    nr = length(r)
    nz = length(z)
    if length(p) == 2
        Σ_ep = ep_cov(energy, pitch, p[1], p[2])
    elseif length(p) == 3
        Σ_ep = (p[1]^2)*ep_cov(energy, pitch, p[2], p[3])
    else
        error("Unsupported number of parameters")
    end
    return RepeatedBlockDiagonal(Σ_ep,nr*nz)
end

function eprz_kernel(x,y,Σ_inv)
    theta_x = acos(clamp(x[2],-1.0,1.0))
    theta_y = acos(clamp(y[2],-1.0,1.0))
    d = S4(x[1] - y[1], # dE
           (theta_x - theta_y)/(2pi), # dtheta/2pi
           x[3] - y[3], # dr
           x[4] - y[4]) # dz
    #d = x .- y
    l = d'*Σ_inv*d
    return exp(-0.5*l)
end

function make_jacobian_spline(J::Vector{T}, t) where T<:Number
    if length(t) == 0
        Js = scale(interpolate(zeros(2),BSpline(Linear())),range(0.0,1.0,length=2))
    else
        Js = scale(interpolate(J, BSpline(Linear())), t)
    end
    return Js
end

make_jacobian_spline(J::T, t) where T<:AbstractInterpolation = J

function get_covariance(o1::OrbitSpline, J1::T, o2::OrbitSpline, J2::T, Σ_inv::S44, atol) where {T<:AbstractInterpolation}
    (length(o1) == 0 || length(o2) == 0) && return 0.0
    K, kerr = hcubature(x -> J1(x[1])*J2(x[2])*eprz_kernel(o1(x[1]), o2(x[2]), Σ_inv),
                        (0.0,0.0), (1.0,1.0), atol=atol)
    return abs(K)
end

function get_covariance(o1::OrbitSpline, J1::T, o2::OrbitSpline, J2::T, sigma::AbstractVector, atol) where {T<:AbstractInterpolation}
    Σ_inv = S44(inv(Diagonal(sigma.^2)))
    return get_covariance(o1, J1, o2, J2, Σ_inv, atol)
end

function get_covariance(M::AxisymmetricEquilibrium, o, sigma;
                        J = Float64[], atol = 1e-3, kwargs...)

    n = length(o)
    t = range(0.0, 1.0, length = n)
    oi = OrbitSpline(o, t)
    if isempty(J)
        J1i = make_jacobian_spline(get_jacobian(M, o; kwargs...),t)
    else
        J1i = make_jacobian_spline(J, t)
    end

    return get_covariance(oi, J1i, oi, J1i, sigma, atol)
end

function get_covariance(M::AxisymmetricEquilibrium, o1, o2, sigma;
                        J1 = Float64[], J2 = Float64[], atol=1e-3, kwargs...)

    n1 = length(o1)
    n2 = length(o2)
    if n1 == 0 || n2 == 0
        return 0.0
    end

    t1 = range(0.0, 1.0, length=n1)
    o1i = OrbitSpline(o1, t1)

    t2 = range(0.0, 1.0, length=n2)
    o2i = OrbitSpline(o2, t2)

    if isempty(J1)
        J1i = make_jacobian_spline(get_jacobian(M, o1; kwargs...), t1)
    else
        J1i = make_jacobian_spline(J1, t1)
    end

    if isempty(J2)
        J2i = make_jacobian_spline(get_jacobian(M, o2; kwargs...), t2)
    else
        J2i = make_jacobian_spline(J2, t2)
    end

    return get_covariance(o1i, J1i, o2i, J2i, sigma, atol)
end

function get_covariance(M::AxisymmetricEquilibrium, c1::T, c2::T, sigma;
                         k1::Float64 = 0.0, k2::Float64 = 0.0,
                         J1 = Float64[], J2 = Float64[],
                        atol = 1e-3, kwargs...) where T <: Union{AbstractOrbitCoordinate,AbstractParticle}
    o1 = get_orbit(M, c1; kwargs...)
    o2 = get_orbit(M, c2; kwargs...)
    return get_covariance(M, o1, o2, sigma; J1 = J1, k1 = k1, J2 = J2, k2 = k2, atol = atol, kwargs...)
end

function get_correlation(M::AxisymmetricEquilibrium, o1, o2, sigma;
                         k1::Float64 = 0.0, k2::Float64 = 0.0,
                         J1 = Float64[], J2 = Float64[],
                         atol = 1e-3, kwargs...)

    n1 = length(o1)
    n2 = length(o2)
    if n1 == 0 || n2 == 0
        return 0.0
    end

    t1 = range(0.0, 1.0, length=n1)
    o1i = OrbitSpline(o1, t1)

    t2 = range(0.0, 1.0, length=n2)
    o2i = OrbitSpline(o2, t2)

    if isempty(J1)
        J1i = make_jacobian_spline(get_jacobian(M, o1; kwargs...), t1)
    else
        J1i = make_jacobian_spline(J1, t1)
    end

    if isempty(J2)
        J2i = make_jacobian_spline(get_jacobian(M, o2; kwargs...), t2)
    else
        J2i = make_jacobian_spline(J2, t2)
    end

    if k1 == 0.0
        k11 = get_covariance(o1i, J1i, o1i, J1i, sigma, atol)
    else
        k11 = k1
    end

    if k2 == 0.0
        k22 = get_covariance(o2i, J2i, o2i, J2i, sigma, atol)
    else
        k22 = k2
    end

    if k11 == 0.0 || k22 == 0.0
        return k11*k22
    end

    k12 = get_covariance(o1i, J1i, o2i, J2i, sigma, atol)

    return k12/sqrt(k11*k22)
end

#1-arg Threaded Covariance
function compute_covariance_matrix(orbs, Js, Σ_inv, atol)
    n = length(orbs)
    N = Threads.nthreads()
    Σ = zeros(n,n)
    Threads.@threads for id=1:N
        @inbounds for i = id:N:n
            oi = orbs[i]
            Ji = Js[i]
            for j=i:n
                oj = orbs[j]
                Jj = Js[j]
                s = get_covariance(oi, Ji, oj, Jj, Σ_inv, atol)
                Σ[i,j] = s
                Σ[j,i] = s
            end
        end
    end

    return Σ
end

#1-arg Threaded Sparse Covariance
function compute_covariance_matrix(orbs, Js, Σ_inv, atol, minval)
    n = length(orbs)
    N = Threads.nthreads()
    Σs = [spzeros(n,n) for i=1:N]
    Threads.@threads for id=1:N
        @inbounds for i=id:N:n
            oi = orbs[i]
            Ji = Js[i]
            Σs[id][i,i] = 1.0
            for j=(i+1):n
                oj = orbs[j]
                Jj = Js[j]
                s = get_covariance(oi, Ji, oj, Jj, Σ_inv, atol)
                if s > minval
                    Σs[id][i,j] = s
                    Σs[id][j,i] = s
                end
            end
        end
    end
    Σ = sum(Σs)

    return Σ
end

#1-arg Distributed Covariance
function compute_covariance_matrix(orbs, Js, Σ_inv, atol, pool::AbstractWorkerPool; batch_size = 0)
    n = length(orbs)
    indices = [(i,j) for i=1:n, j=1:n if i >= j]

    if batch_size == 0
        batch_size = round(Int, length(indices)/(5*nprocs()))
    end

    c = pmap(x -> get_covariance(orbs[x[1]],Js[x[1]],orbs[x[2]],Js[x[2]], Σ_inv, atol),
             pool, indices, on_error=x->0.0, retry_delays=zeros(3), batch_size=batch_size)

    Σ = zeros(n,n)
    for (ii,I) in enumerate(indices)
        i,j = I
        Σ[i,j] = c[ii]
        Σ[j,i] = c[ii]
    end

    return Σ
end

#1-arg Distributed Sparse Covariance
function compute_covariance_matrix(orbs, Js, Σ_inv, atol, minval, pool::AbstractWorkerPool; batch_size=0)
    Σ = compute_covariance_matrix(orbs, Js, Σ_inv, atol, pool, batch_size=batch_size)
    Σ[Σ .< minval] .= zero(eltype(Σ))

    return sparse(Σ)
end

#1-arg Covariance
function get_covariance_matrix(M::AxisymmetricEquilibrium, orbits::Vector, sigma::AbstractVector;
                               Js::Vector{Vector{Float64}} = Vector{Float64}[],
                               sparse::Bool = false, atol::Float64 = 1e-3,
                               minval::Float64 = atol, distributed = false,
                               batch_size = 0, kwargs...)
    n = length(orbits)
    ns = length.(orbits)
    ts = [range(0.0, 1.0, length=nn) for nn in ns]
    orbs = [OrbitSpline(o) for o in orbits]

    if isempty(Js)
        J = Array{Vector{Float64}}(undef, n)
        @inbounds Threads.@threads for i=1:n
            oi = orbits[i]::Orbit
            Ji = get_jacobian(M, oi; kwargs...)
            J[i] = Ji
        end
    else
        J = Js
    end
    Jis = [make_jacobian_spline(jj,tt) for (jj, tt) in zip(J,ts)]

    Σ_inv = S44(inv(Diagonal(sigma.^2)))

    if distributed
        pool = CachingPool(workers())
        if sparse
            Σ = compute_covariance_matrix(orbs, Jis, Σ_inv, atol, minval, pool; batch_size=batch_size)
        else
            Σ = compute_covariance_matrix(orbs, Jis, Σ_inv, atol, pool; batch_size=batch_size)
        end
        clear!(pool)
        return Σ
    else
        if sparse
            return compute_covariance_matrix(orbs, Jis, Σ_inv, atol, minval)
        end
        return compute_covariance_matrix(orbs, Jis, Σ_inv, atol)
    end
end

#1-arg Correlation
function get_correlation_matrix(cov)
    nr,nc = size(cov)
    nr != nc && error("Covariance matrix must be square")
    K = sqrt.(diag(cov))
    return cov./(K*K')
end

function get_correlation_matrix(M::AxisymmetricEquilibrium, orbits::Vector, sigma::AbstractVector;
                                Js::Vector{Vector{Float64}} = Vector{Float64}[],
                                sparse::Bool = false, atol::Float64 = 1e-3,
                                minval::Float64 = atol, distributed = false,
                                batch_size = 0, kwargs...)

    Σ = get_covariance_matrix(M, orbits, sigma; Js=Js, sparse=sparse,
                              atol=atol, minval=minval, distributed=distributed,
                              batch_size=batch_size, kwargs...)

    return get_correlation_matrix(Σ)
end

#2-arg Threaded Covariance
function compute_covariance_matrix(orbs1, Js1, orbs2, Js2, Σ_inv, atol)
    n1 = length(orbs1)
    n2 = length(orbs2)
    Σ = zeros(n1,n2)
    N = Threads.nthreads()
    Threads.@threads for id=1:N
        @inbounds for j=id:N:n2
            oj = orbs2[j]
            Jj = Js2[j]
            for i=1:n1
                oi = orbs1[i]
                Ji = Js1[i]
                Σ[i,j] = get_covariance(oi, Ji, oj, Jj, Σ_inv, atol)
            end
        end
    end
    return Σ
end

#2-arg Threaded Sparse Covariance
function compute_covariance_matrix(orbs1, Js1, orbs2, Js2, Σ_inv, atol, minval)
    n1 = length(orbs1)
    n2 = length(orbs2)
    N = Threads.nthreads()
    Σs = [spzeros(n1,n2) for i=1:N]
    Threads.@threads for id=1:N
        @inbounds for j=id:N:n2
            oj = orbs2[j]
            Jj = Js2[j]
            for i=1:n1
                oi = orbs1[i]
                Ji = Js1[i]
                s = get_covariance(oi, Ji, oj, Jj, Σ_inv, atol)
                if s > minval
                    Σs[id][i,j] = s
                end
            end
        end
    end
    Σ = sum(Σs)
    return Σ
end

#2-arg Distributed Covariance
function compute_covariance_matrix(orbs1, Js1, orbs2, Js2, Σ_inv, atol, pool::AbstractWorkerPool; batch_size=0)
    n1 = length(orbs1)
    n2 = length(orbs2)
    indices = vec([(i,j) for i=1:n1, j=1:n2])

    if batch_size == 0
        batch_size = round(Int, length(indices)/(5*nprocs()))
    end

    Σ = pmap(x -> get_covariance(orbs1[x[1]],Js1[x[1]],orbs2[x[2]],Js2[x[2]], Σ_inv, atol),
             pool, indices, on_error=x->0.0, retry_delays=zeros(3), batch_size=batch_size)

    return reshape(Σ,(n1,n2))
end

#2-arg Distributed Sparse Covariance
function compute_covariance_matrix(orbs1, Js1, orbs2, Js2, Σ_inv, atol, minval, pool::AbstractWorkerPool; batch_size=0)
    Σ = compute_covariance_matrix(orbs1, Js1, orbs2, Js2, Σ_inv, atol, pool; batch_size=batch_size)
    Σ[Σ .< minval] .= zero(eltype(Σ))
    return sparse(Σ)
end

#2-arg Covariance
function get_covariance_matrix(M::AxisymmetricEquilibrium, orbits_1::Vector, orbits_2::Vector, sigma::AbstractVector;
                                Js_1::Vector{Vector{Float64}} = Vector{Float64}[],
                                Js_2::Vector{Vector{Float64}} = Vector{Float64}[],
                                sparse::Bool = false, atol::Float64 = 1e-3,
                                minval::Float64 = atol, distributed=false,
                                batch_size=0, kwargs...)

    Σ_inv = S44(inv(Diagonal(sigma.^2)))

    n1 = length(orbits_1)
    ns1 = length.(orbits_1)
    ts1 = [range(0.0, 1.0, length=nn) for nn in ns1]
    orbs1 = [OrbitSpline(o) for o in orbits_1]

    if isempty(Js_1)
        J1 = Array{Vector{Float64}}(undef, n1)
        @inbounds Threads.@threads for i=1:n1
            oi = orbits_1[i]::Orbit
            Ji = get_jacobian(M, oi; kwargs...)
            J1[i] = Ji
        end
    else
        J1 = Js_1
    end
    Jis1 = [make_jacobian_spline(jj,tt) for (jj, tt) in zip(J1,ts1)]

    n2 = length(orbits_2)
    ns2 = length.(orbits_2)
    ts2 = [range(0.0, 1.0, length=nn) for nn in ns2]
    orbs2 = [OrbitSpline(o) for o in orbits_2]

    if isempty(Js_2)
        J2 = Array{Vector{Float64}}(undef, n2)
        @inbounds Threads.@threads for i=1:n2
            oi = orbits_2[i]::Orbit
            Ji = get_jacobian(M, oi; kwargs...)
            J2[i] = Ji
        end
    else
        J2 = Js_2
    end
    Jis2 = [make_jacobian_spline(jj,tt) for (jj, tt) in zip(J2,ts2)]

    if distributed
        pool = CachingPool(workers())
        if sparse
            Σ = compute_covariance_matrix(orbs1, Jis1, orbs2, Jis2, Σ_inv, atol, minval, pool; batch_size=batch_size)
        else
            Σ = compute_covariance_matrix(orbs1, Jis1, orbs2, Jis2, Σ_inv, atol, pool; batch_size=batch_size)
        end
        clear!(pool)
        return Σ
    else
        if sparse
            return compute_covariance_matrix(orbs1, Jis1, orbs2, Jis2, Σ_inv, atol, minval)
        end
        return compute_covariance_matrix(orbs1, Jis1, orbs2, Jis2, Σ_inv, atol)
    end
end

#2-arg Threaded Correlation
function compute_correlation_matrix(orbs1, Js1, Ks1, orbs2, Js2, Ks2, Σ_inv, atol)
    n1 = length(orbs1)
    n2 = length(orbs2)
    Σ = zeros(n1,n2)
    N = Threads.nthreads()
    Threads.@threads for id=1:N
        @inbounds for j=id:N:n2
            oj = orbs2[j]
            Jj = Js2[j]
            kj = Ks2[j]
            for i=1:n1
                oi = orbs1[i]
                Ji = Js1[i]
                ki = Ks1[i]
                Σ[i,j] = get_covariance(oi, Ji, oj, Jj, Σ_inv, atol)/sqrt(ki*kj)
            end
        end
    end
    return Σ
end

#2-arg Threaded Sparse Correlation
function compute_correlation_matrix(orbs1, Js1, Ks1, orbs2, Js2, Ks2, Σ_inv, atol, minval)
    n1 = length(orbs1)
    n2 = length(orbs2)
    N = Threads.nthreads()
    Σs = [spzeros(n1,n2) for i=1:N]
    Threads.@threads for id=1:N
        @inbounds for j=id:N:n2
            oj = orbs2[j]
            Jj = Js2[j]
            kj = Ks2[j]
            for i=1:n1
                oi = orbs1[i]
                Ji = Js1[i]
                ki = Ks1[i]
                s = get_covariance(oi, Ji, oj, Jj, Σ_inv, atol)/sqrt(ki*kj)
                if s > minval
                    Σs[id][i,j] = s
                end
            end
        end
    end
    Σ = sum(Σs)
    return Σ
end

#2-arg Distributed Correlation
function compute_correlation_matrix(orbs1, Js1, Ks1, orbs2, Js2, Ks2, Σ_inv, atol, pool::AbstractWorkerPool; batch_size=0)
    n1 = length(orbs1)
    n2 = length(orbs2)
    indices = vec([(i,j) for i=1:n1, j=1:n2])

    if batch_size == 0
        batch_size = round(Int, length(indices)/(5*nprocs()))
    end

    Σ = pmap(x -> get_covariance(orbs1[x[1]],Js1[x[1]],orbs2[x[2]],Js2[x[2]], Σ_inv, atol)/sqrt(Ks1[x[1]]*Ks2[x[2]]),
             pool, indices, on_error=x->0.0, retry_delays=zeros(3), batch_size=batch_size)

    return reshape(Σ,(n1,n2))
end

#2-arg Distributed Sparse Correlation
function compute_correlation_matrix(orbs1, Js1, Ks1, orbs2, Js2, Ks2, Σ_inv, atol, minval, pool::AbstractWorkerPool; batch_size=0)
    Σ = compute_correlation_matrix(orbs1, Js1, Ks1, orbs2, Js2, Ks2, Σ_inv, atol, pool; batch_size=batch_size)
    Σ[Σ .< minval] .= zero(eltype(Σ))
    return sparse(Σ)
end

function get_correlation_matrix(M::AxisymmetricEquilibrium, orbits_1::Vector, orbits_2::Vector, sigma::AbstractVector;
                                ks_1::Vector{Float64} = Float64[],
                                Js_1::Vector{Vector{Float64}} = Vector{Float64}[],
                                ks_2::Vector{Float64} = Float64[],
                                Js_2::Vector{Vector{Float64}} = Vector{Float64}[],
                                sparse::Bool = false, atol::Float64 = 1e-3,
                                minval::Float64 = atol, distributed=false,
                                batch_size = 0, kwargs...)

    Σ_inv = S44(inv(Diagonal(sigma.^2)))

    n1 = length(orbits_1)
    ns1 = length.(orbits_1)
    ts1 = [range(0.0, 1.0, length=nn) for nn in ns1]
    orbs1 = [OrbitSpline(o) for o in orbits_1]

    if isempty(Js_1)
        J1 = Array{Vector{Float64}}(undef, n1)
        @inbounds Threads.@threads for i=1:n1
            oi = orbits_1[i]::Orbit
            Ji = get_jacobian(M, oi; kwargs...)
            J1[i] = Ji
        end
    else
        J1 = Js_1
    end
    Jis1 = [make_jacobian_spline(jj,tt) for (jj, tt) in zip(J1,ts1)]

    if isempty(ks_1)
        K1 = zeros(n1)
        @inbounds Threads.@threads for i=1:n1
            oi = orbs1[i]
            Ji = Jis1[i]
            K1[i] = get_covariance(oi, Ji, oi, Ji, Σ_inv, atol)
        end
    else
        K1 = ks_1
    end

    n2 = length(orbits_2)
    ns2 = length.(orbits_2)
    ts2 = [range(0.0, 1.0, length=nn) for nn in ns2]
    orbs2 = [OrbitSpline(o) for o in orbits_2]

    if isempty(Js_2)
        J2 = Array{Vector{Float64}}(undef, n2)
        @inbounds Threads.@threads for i=1:n2
            oi = orbits_2[i]::Orbit
            Ji = get_jacobian(M, oi; kwargs...)
            J2[i] = Ji
        end
    else
        J2 = Js_2
    end
    Jis2 = [make_jacobian_spline(jj,tt) for (jj, tt) in zip(J2,ts2)]

    if isempty(ks_2)
        K2 = zeros(n2)
        @inbounds Threads.@threads for i=1:n2
            oi = orbs2[i]
            Ji = Jis2[i]
            K2[i] = get_covariance(oi, Ji, oi, Ji, Σ_inv, atol)
        end
    else
        K2 = ks_2
    end

    if distributed
        pool = CachingPool(workers())
        if sparse
            Σ = compute_correlation_matrix(orbs1, Jis1, K1, orbs2, Jis2, K2, Σ_inv, atol, minval, pool, batch_size=batch_size)
        else
            Σ = compute_correlation_matrix(orbs1, Jis1, K1, orbs2, Jis2, K2, Σ_inv, atol, pool, batch_size=batch_size)
        end
        clear!(pool)
        return Σ
    else
        if sparse
            return compute_correlation_matrix(orbs1, Jis1, K1, orbs2, Jis2, K2, Σ_inv, atol, minval)
        end
        return compute_correlation_matrix(orbs1, Jis1, K1, orbs2, Jis2, K2, Σ_inv, atol)
    end

end


export dcs_exhaustive
using Devectorize

function dcs_exhaustive{T}(G::Array{SparseMatrixCSC{T, Int64}, 1})
  m = length(G)
  n = size(G[1], 1)

  prev_state = zeros(Bool, n)
  prev_nodes = zeros(Int64, n)
  prev_num_nodes = 0

  binary_state = zeros(Bool, n)

  curr_state = zeros(Bool, n)
  curr_nodes = zeros(Int64, n)
  curr_num_nodes = 0

  opt_idx = 0
  opt_deg = 0.0

  curr_ne = zeros(m, 1)
  curr_deg = zeros(m, 1)

  for j=[1:(2^n-1);]
    int_to_binary_bool!(binary_state, j, n)
    binary_to_gray!(curr_state, binary_state)
    curr_num_nodes = find_set_bits!(curr_nodes, curr_state)

    # Find which node has been added or deleted
    changed_idx = 1
    while (curr_state[changed_idx] == prev_state[changed_idx])
      changed_idx = changed_idx + 1
    end

    # Has the node been added or deleted?
    if curr_state[changed_idx]
      m_sign = 1
    else
      m_sign = -1
    end
    @inbounds begin
      for k=[1:m;]
        delta_deg = 0
        if prev_num_nodes > 0
          for ll=[1:prev_num_nodes;]
            other_node = prev_nodes[ll]
            delta_deg = delta_deg + G[k][other_node, changed_idx] + G[k][changed_idx, other_node]
          end
        end
        curr_ne[k] = curr_ne[k] + m_sign  * delta_deg
        curr_deg[k] = curr_ne[k] / curr_num_nodes
        #      println("G[$k] j: $j: $(curr_ne[k]), $(curr_deg[k])")
      end
    end
    @devec curr_min_deg = minimum(curr_deg)
    if curr_min_deg > opt_deg
      opt_deg = curr_min_deg
      opt_idx = j
    end

    for k=[1:n;]
      prev_state[k] = curr_state[k]
      prev_nodes[k] = curr_nodes[k]
    end
    prev_num_nodes = curr_num_nodes
  end
  int_to_binary_bool!(binary_state, opt_idx, n)
  binary_to_gray!(curr_state, binary_state)
  curr_num_nodes = find_set_bits!(curr_nodes, curr_state)
  opt_S = curr_nodes[1:curr_num_nodes]
  # println("S: $opt_S min_deg: $opt_deg")
  return opt_S , opt_deg
end


using DynamicalSystems, Neighborhood

function correlationsum_tree(X, ε; q = 2, metric = Euclidean(), w = 0)
    tree = KDTree(X, metric)
    qf = q == 2 ? identity : x -> float(x)^(q-1)
    _correlationsum_new(tree, X, ε, w, qf)
end

function _correlationsum_tree(tree, X, ε, w, qf)
    theiler = Neighborhood.Theiler(w)
    searchtype = Neighborhood.WithinRange(maximum(ε))
    N = length(X)
    Cq = similar(ε, Float64)
    @inbounds for i in 1:N
        idxs = Neighborhood.isearch(tree, X[i], searchtype, theiler(i))
        isempty(idxs) && continue
        dist = [Neighborhood.evaluate(tree.metric, X[i], X[j]) for j in idxs]
        sort!(dist)
        for k in 1:length(ε)-1
            j = searchsortedfirst(dist, ε[k])
            Cq[k] += qf(j-1)
        end
        Cq[end] += qf(length(dist))
    end
    return Cq ./ (N*qf(N-w-1))
end


using SDEs, Base.Test
include("setup.jl")
include("test_utils.jl")

Ncheck = 1000
tolerance = 2/sqrt(Ncheck) #CLT for comparing processes with different random number generators
id_tolerance = 1e-7 #For comparing things with same random number generators
θ=0.5
η=0.5

println(prob, ", dt:", t[2]-t[1], ", θ:",θ, ", η:",η)

EMPdev = check(prob, x0, t, SDEs.EMproto(), 1:Ncheck)
println("proto EM  :", EMPdev)

EMdev = check(prob, x0, t, EM(), 1:Ncheck)
println("EM        :", EMdev)

DeqEMdev = check(prob, x0, t, DEQEM(), 1:Ncheck)
println("Diffeq    :", DeqEMdev)

PCEPdev = check(prob, x0, t, SDEs.PCEproto(), 1:Ncheck; θ=θ, η=η)
println("proto PCE :", PCEPdev)

PCEdev = check(prob, x0, t, PCE(), 1:Ncheck; θ=θ, η=η)
println("PCE       :", PCEdev)

@test EMdev ≈ DeqEMdev atol=EMdev*tolerance
@test EMdev ≈ EMPdev atol=EMdev*id_tolerance
@test PCEPdev ≈ PCEdev atol=PCEdev*id_tolerance
@test PCEdev < EMdev