"""
	vcongruence(
        V::Lar.Points; epsilon=1e-6
    )::Tuple{Lar.Points, Array{Array{Int,1},1}}

Evaluates the Vertex Congruence for 3D-points ``V ∈ ℳ(3,n)``.

The function determines the points of ``V`` closer than ``epsilon`` and
builds a new Vertex Set made of the representative of each point cluster.

The method returns:
 - the new Vertex Set
 - a map that, for every new vertex, identifies the old vertices it is made of
"""
function vcongruence(V::Lar.Points; epsilon=1e-6)
    vclasses, visited = [], []
    V = convert(Array{Float64,2}, V)
    kdtree = NearestNeighbors.KDTree(V)
    for vidx = 1 : size(V, 2)  if !(vidx in visited)
        nearvs = NearestNeighbors.inrange(kdtree, V[:,vidx], epsilon)
        push!(vclasses, nearvs)
        append!(visited, nearvs)
    end  end
    W = hcat([sum(V[:,class], dims=2)/length(class) for class in vclasses]...)
    return W, vclasses
end

"""
    cellCongruence(
        Delta::Lar.ChainOp,
        inclasses::Array{Array{Int,1},1};
        dim::Int = 0
    )::Tuple{Lar.Cells, Array{Array{Int,1},1}}

Evaluates the Cell Congruence for Cochain ``Delta`` with classes ``inclasses``.

The function determines the new Cochain (as ``Lar.Cells``) built from ``cop``
where the lower order cells are merged according to ``inclasses`` map.

The parameter ``dim``, if specified, represent the order of the cells
(that is the least number of lower order cells a current order cell is made of).

The method returns:
- the new Cochain Operator as ``Lar.Cells``
- a map that, for every new cell, identifies the old cells it is made of
"""

function cellcongruence(Delta, inclasses; dim = 0)
    cellarray = Lar.cop2lar(Delta)
    newcell = Vector(undef, size(Delta,2))
    for (k, class) in enumerate(inclasses)  for e in class
            newcell[e] = k
    end end
    cells = [sort(map(x -> newcell[x], face)) for face in cellarray]
    okcells = [cell for cell in cells if length(Set(cell)) > dim] # non-empty cells
    classes = DefaultOrderedDict{Vector, Vector}([])
    for (k,face) in enumerate(okcells)
        classes[face] == [] ? classes[face] = [k] : append!(classes[face], [k])
    end
    cells = collect(keys(classes))
    outclasses = collect(values(classes))
    return cells, outclasses
end

"""
    chaincongruence(
        W::Lar.Points,
        Delta_0::Lar.ChainOp,
        Delta_1::Lar.ChainOp
        epsilon = 1e-6
    )::Tuple{Lar.Points, Lar.Cells, Lar.Cells}

Cell Congruence Enabling Algorithm for 3D Complexes.

The algorithm performs the identify vertices within ``epsilon`` and coherently
update the Chain Operators.
"""
function chaincongruence(W, Delta_0, Delta_1; epsilon = 1e-6)
    V, vclasses = vcongruence(W, epsilon = epsilon)
    EV, eclasses = cellcongruence(Delta_0, vclasses, dim=1)
    FE, _ = cellcongruence(Delta_1, eclasses, dim=2)
    return V, EV, FE
end

"""
	signedCellCongruence(
		cop::Lar.ChainOp,
		lo_cls::Array{Array{Int,1},1},
		lo_sign::Array{Array{Int8,1},1};
		imp = false,
		d = 0
	)
Evaluates the Cell Congruence for a Cochain ``cop`` with classes ``lo_cls``.
The function determines the new Cochain Operator built from ``cop`` where
the lower order cells are merged according to ``lo_cls`` map.
``lo_sign`` specifies whether a cell must be considered in reverse order.
If optional paramenter ``imp`` is set to ``true`` then FP imprecisions
are taken into account in the sense that lower order cells may have collided.
The parameter ``d`` represent then the order of the cell (that also is the
least number of lower order cells a current order cell is made of).
The method returns:
 - the new Cochain Operator
 - a map that, for every new cell, identifies the old cells it is made of
 - a map that, for every new cell, specify if old cells have changed ordering.
"""
function signedCellCongruence(cop, lo_cls, lo_sign; imp = false, dim = 0)::Tuple{
        Lar.ChainOp,
		Array{Array{Int,1},1},
		Array{Array{Int8,1},1}
    }
#function cellCongruence(
#		cop::Lar.ChainOp,
#		lo_cls::Array{Array{Int,1},1},
#		lo_sign::Array{Array{Int8,1},1};
#		imp = false,
#		dim = 0
#	)::Tuple{
#		Lar.ChainOp,
#		Array{Array{Int,1},1},
#		Array{Array{Int8,1},1}
#	}

	# Collide columns
	copCols = []
	for i = 1 : length(lo_cls)
		col = sum([
			cop[:, lo_cls[i][j]] .* lo_sign[i][j]
			for j = 1 : length(lo_cls[i])
		])
#		if imp
#			#TODO remove zeros
#			if length(col.nzind) > dim  push!(copCols, col)  end
#		end
		push!(copCols, col)
	end

	# Retrieve Matrix Form and extract rows
	cop = hcat(copCols...)
	rows = [cop[row, :] for row = 1 : cop.m]

	# Adjustinfg signs such that the first value always is `-1`
	sign = ones(Int8, cop.m)
	for row = 1 : cop.m  if rows[row].nzval[1] > 0
		rows[row] = -rows[row]
		sign[row] = -1
	end  end

	# Sort rows in order to collide them quickly
	rows_ord = sortperm(rows)
	nrows = unique(rows[rows_ord])
	ho_cls = [Array{Int,1}() for i = 1 : length(nrows)]
	nidx = 1

	# Collide cells with same structure
	for cidx = 1 : cop.m
		if rows[rows_ord[cidx]] != nrows[nidx]  nidx = nidx + 1;  end
		push!(ho_cls[nidx], rows_ord[cidx])
	end

	# Reshaping Sign and Cochain
	ho_sign = [[sign[el] for el in cell] for cell in ho_cls]
#	cop = convert(Lar.ChainOp, hcat(nrows...)')

	return hcat(nrows...)', ho_cls, ho_sign
end


using Random
using InvertedIndices
using TensorToolbox

function rapid_tensor_from_theta(theta)
    tensor_shape = size(theta)
    d = ndims(theta)

    rs = []
    for k = 1:d
        theta_k = theta[ ntuple( j->j == k ? (:) : 1, d)... ]
        r = ones( tensor_shape[k] )
        r[1] = exp( theta_k[1] )
        for i = 2:tensor_shape[k]
            r[i] = r[i-1] * exp( theta_k[i] )
        end
        push!(rs, r)
    end

    tensor = ttt( vec(rs[1]), vec(rs[2]) )
    for k = 3:d
        tensor = ttt(tensor, vec(rs[k]) )
    end

    return tensor
end

function rapid_eta_from_tensor(T)
    tensor_shape = size(T)
    eta = zeros( tensor_shape )
    d = ndims(T)
    one_to_N = [1:d;]
    for i = 1:d
        tensor_shape_i = ones(Int, d)
        tensor_shape_i[i] = tensor_shape[i]
        tensor_shape_i_tp = tuple( tensor_shape_i... )
        sums = vec(sum(T, dims=one_to_N[Not(i)]))
        for idxes in reverse(Tuple(CartesianIndices( tensor_shape_i_tp )))
            eta[idxes] = sums[ idxes[i] ]
            if idxes[i] != tensor_shape[i]
                tmp = zeros(Int, d)
                tmp[i] = 1
                eta[idxes] += eta[ idxes + CartesianIndex( tuple(tmp...) ) ]
            end
        end
    end

    return eta
end

function rapid_theta_from_tensor(T)
    tensor_shape = size(T)
    theta = zeros( tensor_shape )
    d = ndims(T)
    for i = 1:d
        tensor_shape_i = ones(Int, d)
        tensor_shape_i[i] = tensor_shape[i]
        tensor_shape_i_tp = tuple( tensor_shape_i... )
        for idxes in CartesianIndices( tensor_shape_i_tp )
            theta[ idxes ] = log( T[ idxes ] )
            if idxes[i] > 1
                tmp = zeros(Int, d)
                tmp[i] = 1
                theta[ idxes ] -= log( T[ idxes - CartesianIndex( tuple(tmp...))] )
            end
        end
    end
    return theta
end

function get_uniform_eta( tensor_shape )
    eta = ones(tensor_shape)
    for idxes in CartesianIndices( tensor_shape )
        tmp = 1.0
        for i = 1:length(idxes)
            tmp *= (tensor_shape[i]-idxes[i]+1)
        end
        eta[ idxes ] = tmp
    end
    eta ./= prod(tensor_shape)
    return eta
end

function get_uniform_theta( tensor_shape )
    theta = zeros(tensor_shape)
    theta[1] = -1.0 * log(prod( tensor_shape ))
    return theta
end


using Test

# A recursive descent parser for a string representation of boolean expressions

mutable struct ExprString
    data::String
    pos::Int
    ExprString(data, pos) = pos != 1 ? error("start must be 1") : new(data, 1)
end

ExprString(s::String) = ExprString(s, 1)

hasmore(e::ExprString) = e.pos <= length(e.data)
isparen(c::Char) = c == '(' || c == ')'
isbinop(c::Char) = c == '&' || c == '|'
isunop(c::Char) = c == '!'
isalphanum(c::Char) = isletter(c) || isnumeric(c)

function skip_when(f, s, start)::Int
    len = length(s)
    if start <= len
        for i in start:len
            if !f(s[i])
                return i
            end
        end
    end
    return len + 1
end

function next_pos(s::String, pos::Int)::Int
    c = s[pos]
    if isparen(c) || isunop(c) || isbinop(c)
        return skip_when(isspace, s, pos + 1)
    elseif isletter(c) #We're inside a variable name, skip alphanumeric
        return skip_when(isalphanum, s, pos + 1)
    else
        throw(ArgumentError("Expected alphanum, parens, \"&\", \"|\", or \"!\", found \"$c\""))
    end
end

function peek(e::ExprString)::String
    c = e.data[e.pos]
    if isparen(c) || isbinop(c)
        return string(c)
    end
    endPos = next_pos(e.data, e.pos)
    strip(e.data[e.pos:endPos-1])
end

function consume(e::ExprString, s::String)
    nextStartPos = next_pos(e.data, e.pos)
    expected = strip(e.data[e.pos:nextStartPos-1])
    if expected == s
        e.pos = nextStartPos
    else
        throw(ArgumentError("Expected $expected, saw $s"))
    end
end

function parse_expr(e::ExprString)::BoolExpr
    term = parse_term(e)
    if hasmore(e)
        tok_string = peek(e)
        if tok_string == "&"
            consume(e, "&")
            return And([term, parse_expr(e)])
        elseif tok_string == "|"
            consume(e, "|")
            return Or([term, parse_expr(e)])
        end
    end
    return term
end

function parse_term(e::ExprString)::BoolExpr
    tok_string = peek(e)
    if tok_string == "!"
        consume(e, "!")
        result = parse_term(e)
        return Not(result)
    elseif tok_string == "("
        consume(e, "(")
        result = parse_expr(e)
        consume(e, ")")
        return result
    elseif true # isAlpha(tok_string)
        consume(e, tok_string)
        return Var(tok_string)
    else
        throw(ArgumentError("Syntax error: expected variable name or \"(\", found \"$tok_string\")"))
    end
end

Base.show(io::IO, exprString::ExprString) = print(io, "ExprString(data=$e.data, pos=$e.pos)")
Base.show(io::IO, boolExpr::BoolExpr) = BoolExprs.show(io, boolExpr)
mkexpr(s) = parse_expr(ExprString(s))

"""
"""
#######

Base.:(==)(x::BoolExpr, y::BoolExpr)::Bool = equiv(x, y)

@testset "Make expressions from strings" begin
    x1 = mkexpr("x1")
    @test x1 == Var("x1")
    x2 = mkexpr("x2")
    @test x2 == Var("x2")
    x3 = mkexpr("x3")
    @test x3 == Var("x3")

    not_x3 = mkexpr("!x3")
    @test not_x3 == Not(Var("x3"))

    and_x1_x2 = mkexpr("x1&x2")
    @test and_x1_x2 == And([x1, x2])

    and_x2_x3 = And([x2, x3])
    @test and_x2_x3 == And([x2, x3])

    or_x1_x2 = mkexpr("x1|x2")
    @test or_x1_x2 == Or([x1, x2])

    or_x2_x3 = mkexpr("x2|x3")
    @test or_x2_x3 == Or([x2, x3])
end
#######

@testset "DNF" begin
    s = "(x1&x2) & (x2&x3)"
    dnf_expr = dnf(mkexpr(s))
    @test dnf_expr == Or([And([Var("x1"), Var("x2"), Var("x2"), Var("x3")])])
    println("dnf($s) = $dnf_expr")
    s = "(x1&x2) | (x2&x3)"
    dnf_expr = dnf(mkexpr(s))
    @test dnf_expr == Or([And([Var("x1"), Var("x2")]), And([Var("x2"), Var("x3")])])
    println("dnf($s) = $dnf_expr")
    s = "(x1|x2) & (x2|x3)"
    dnf_expr = dnf(mkexpr(s))
    @test dnf_expr == Or([And([Var("x1"), Var("x2")]), And([Var("x1"), Var("x3")]), And([Var("x2"), Var("x2")]), And([Var("x2"), Var("x3")])])
    println("dnf($s) = $dnf_expr")
end

@testset "CNF" begin
    s = "(x1&x2) & (x2&x3)"
    cnf_expr = cnf(mkexpr(s))
    @test cnf_expr == And([Var("x1"), Var("x2"), Var("x2"), Var("x3")])
    println("cnf($s) = $cnf_expr")
    s = "(x1&x2) | (x2&x3)"
    cnf_expr = cnf(mkexpr(s))
    @test cnf_expr == And([Or([Var("x1"), Var("x2")]), Or([Var("x1"), Var("x3")]), Or([Var("x2"), Var("x2")]), Or([Var("x2"), Var("x3")])])
    println("cnf($s) = $cnf_expr")
    s = "(x1|x2) & (x2|x3)"
    cnf_expr = cnf(mkexpr(s))
     @test cnf_expr == And([Or([Var("x1"), Var("x2")]), Or([Var("x2"), Var("x3")])])
   println("cnf($s) = $cnf_expr")
end


"""
    mutation:classic(A::AbstractMatrix, lower_limits::AbstractArray, upper_limits::AbstractArray, F::Real)
Given a matrix `A` with `lower_limits` and `upper_limits` arrays, crates a new mutated
matrix with a constant sacale factor `F`.

# Parameters
**`A::AbstractMatrix`**:

Input matrix

## **`lower_limits::AbstractArray`**:

Array that contains all the minimun values
that each individue could obtain.

## **`upper_limits::AbstractArray`**:

Array that contains all the maximum values
that each individue could obtain.

## **`F::Real`**:

The scale factor, F ∈ (0,1+), is a positive real number that controls
the rate at which the population evolves. While there is no upper limit on F,
effective values are seldom greater than 1.0.
(*Price, K., Storn, R. M., & Lampinen, J. A. (2006).
Differential evolution: a practical approach to global optimization.
Springer Science & Business Media.*)

# Examples
```julia-repl
A = 5 .*rand(5,3)
5×3 Matrix{Float64}:
 1.18384  0.587246  4.46021
 4.58842  3.52444   1.39847
 1.08242  4.8591    4.9314
 3.54709  3.17776   4.42353
 4.62508  4.84452   4.59266
```
`A` is a matrix with minimun and maximum possible values equal to 0 and 5 respectively.
```julia-repl
julia> mutation_classic(A,[0,0,0],[5,5,5],1.3)
5×3 Matrix{Float64}:
 4.53165   4.40841   2.46325
 1.55285   1.47686   4.64034
 3.72591   0.302261  1.29996
 1.2315    2.30335   0.0760031
 0.797949  1.3387    0.738242
```
"""
function mutation_classic(A::AbstractMatrix, lower_limits::AbstractArray, upper_limits::AbstractArray, F::Real)

    mutated = similar(A)
    population_size = size(A,1)
    number_of_parameters = size(A,2)

    @simd for individue in 1:population_size

        #Obtaining 3 differents individues without taking into account the current individue

        rnd_individues = rand_exclusive(deleteat!(collect(1:population_size),individue),3)

        @simd for parameter in 1:number_of_parameters
            @inbounds mutated[individue,parameter] = A[rnd_individues[1],parameter] + F*(A[rnd_individues[2],parameter]-A[rnd_individues[3],parameter])
            if (mutated[individue,parameter] < lower_limits[parameter]) | (mutated[individue,parameter] > upper_limits[parameter])
                @inbounds mutated[individue,parameter] = (upper_limits[parameter] - lower_limits[parameter])*rand() + lower_limits[parameter]
            end
        end
    end

    return mutated

end

"""
    mutation_jither(A::AbstractMatrix, lower_limits::AbstractArray, upper_limits::AbstractArray, F::Real)
Similar to **`mutation_classic`**. Given a population matrix `A`, perform the mutation with a random scale factor
between (0,1]. Each scale factor is diferent for each parameter.
"""
function mutation_jither(A::AbstractMatrix, lower_limits::AbstractArray, upper_limits::AbstractArray, F::Real)

    mutated = similar(A)
    population_size = size(A,1)
    number_of_parameters = size(A,2)

    @simd for individue in 1:population_size

        #Obtaining 3 differents individues without taking into account the current individue

        rnd_individues = rand_exclusive(deleteat!(collect(1:population_size),individue),3)

        @simd for parameter in 1:number_of_parameters
            @inbounds mutated[individue,parameter] = A[rnd_individues[1],parameter] + rand()*F*(A[rnd_individues[2],parameter]-A[rnd_individues[3],parameter])
            if (mutated[individue,parameter] < lower_limits[parameter]) | (mutated[individue,parameter] > upper_limits[parameter])
                @inbounds mutated[individue,parameter] = (upper_limits[parameter] - lower_limits[parameter])*rand() + lower_limits[parameter]
            end
        end
    end

    return mutated
end

"""
    mutation_dither(A::AbstractMatrix, lower_limits::AbstractArray, upper_limits::AbstractArray, F::Real)
Similar to **`mutation_classic`**. Given a population matrix `A`, perforn the mutation with a random scale factor
in the interval (0,F). Each scale factor is diferent for each individue but the same for all parameters.
"""
function mutation_dither(A::AbstractMatrix, lower_limits::AbstractArray, upper_limits::AbstractArray, F::Real)

    mutated = similar(A)
    population_size = size(A,1)
    number_of_parameters = size(A,2)

    @simd for individue in 1:population_size

        #Obtaining 3 differents individues without taking into account the current individue

        rnd_individues = rand_exclusive(deleteat!(collect(1:population_size),individue),3)
        rnd = rand()
        @simd for parameter in 1:number_of_parameters
            @inbounds mutated[individue,parameter] = A[rnd_individues[1],parameter] + rnd*F*(A[rnd_individues[2],parameter]-A[rnd_individues[3],parameter])
            if (mutated[individue,parameter] < lower_limits[parameter]) | (mutated[individue,parameter] > upper_limits[parameter])
                @inbounds mutated[individue,parameter] = (upper_limits[parameter] - lower_limits[parameter])*rand() + lower_limits[parameter]
            end
        end
    end

    return mutated
end

"""
    mutation_JADE_no_archive(A::AbstractMatrix, lower_limits::AbstractArray, upper_limits::AbstractArray, F::Real)
Named DE/current-to-pbest/1

"""
function mutation_current2pbest(xi::AbstractArray, x_best::AbstractArray, x_r1::AbstractArray, x_r2::AbstractArray, F::Real)
    v = similar(xi)
    @inbounds for parameter in 1:length(xi)
        v[parameter] = xi[parameter] + F*(x_best[parameter] - xi[parameter]) + F*(x_r1[parameter] - x_r2[parameter])
    end

    return v
end


include("../../__lib__/math/ode/src/ode_module.jl")
include("../../__lib__/math/ode/ode_test_functions/ode_test_functions.jl")
include("../../__lib__/math/common/numder/src/noisy_discrete_diff.jl")

include("../src/sindy_module.jl")

# create from the Xi matrix a diffeq
# solve the ode on the same interval and settings
# check the error in the time domain

module stest
using ..sindy
using ..noisydiff
using ..odesample
using ..ode
using PyPlot
PyPlot.pygui(true)

enable_plot_diffs  = true
enable_plot_odesol = false

function exitation(t)
    return 3 * sin(10 * t)
end

# --------------------------------------------------------------------------- #
# problem definition and parameters
# --------------------------------------------------------------------------- #

omega = 4
D = 0.05
m = 1.0
c = m * omega^2
d = 2 * D * omega # 2*0.05*3 = 3

params = odesample.HarmonicOscillatorOptions(;
    m = m,
    d = d,
    c = c,
    f = odesample.std_zero)  # exitation
ode_fcn(der_q, t, q) = odesample.harmonic_oscillator!(der_q, t, q, params)

# --------------------------------------------------------------------------- #
# ode solving
# --------------------------------------------------------------------------- #

time_interval = ode.TimeInterval(35.0)
init_cond     = [0.0, 1.0]

ode_problem = ode.OdeProblem(
    ode_fcn,
    time_interval,
    init_cond)

# ode_solver = ode.BogaczkiShampine23()
# ode_solver = ode.DormandPrince45()
ode_solver = ode.Tsitouras45()

basic_options = ode.BasicOdeOptions(
    ode_problem;
    abs_tol = 1e-15,
    rel_tol = 1e-15,
    max_iter = 10_000_000
)

step_size_controller = ode.ClassicalStepSizeControl()

stepper_options = ode.StepperOptions(;
    max_step_size = 5.0)

@time (ode_res, step_stat) = ode.solve(
    ode_problem,
    ode_solver;
    basic_options        = basic_options,
    step_size_controller = step_size_controller,
    stepper_options      = stepper_options
    )

ode_time = ode.time(ode_res)
ode_vars = ode.variables(ode_res)

# --------------------------------------------------------------------------- #
# Sindy
# --------------------------------------------------------------------------- #

window_size = 3
t_mids = ode_time[window_size+1:end]

dx1 = noisydiff.diff_sliding_least_squares_d1(ode_time, ode_vars[1,:], w = window_size)
dx2 = noisydiff.diff_sliding_least_squares_d1(ode_time, ode_vars[2,:], w = window_size)

dermat  = [dx1 dx2]
datamat = ode_vars[:, (window_size+1):end]'
u       = params.f.( ode_time )
u       = u[(window_size+1):end]
datamat = [datamat u]

poly_order = 2
(libmat, names) = sindy.create_poly_lib(datamat, polyorder = poly_order)

sparsifing_opt = sindy.SparsifingOptions(
    threshold = 0.1,
    max_iter = 10)

Xi = sindy.sparsify_dynamics(dermat, libmat, sparsifing_opt)

# dxdt = v
# dvdt = 1/m * (f(t) - c * x - d * v )

# dxdt = v
# dvdt = -c/m * x - d/m * v

Xi_theo = [
            0 0;
            0 -c/m;
            1 -d/m]

# --------------------------------------------------------------------------- #
# Visualize
# --------------------------------------------------------------------------- #

if enable_plot_odesol
    PyPlot.figure()
    PyPlot.subplot(2,1,1)
    PyPlot.grid()
    PyPlot.plot( ode_res.t, ode_res.q[:, 1] )

    PyPlot.subplot(2,1,2)
    PyPlot.grid()
    PyPlot.plot( ode_res.t, ode_res.q[:, 2] )
    PyPlot.plot(t_mids, dx1)
end

if enable_plot_diffs
    PyPlot.figure()

    PyPlot.subplot(2,1,1)
    PyPlot.grid()
    PyPlot.plot(t_mids, dx1)

    PyPlot.subplot(2,1,2)
    PyPlot.grid()
    PyPlot.plot(t_mids, dx2)
end

end


# %% Entropy comparison
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

he = Systems.henon()
N, q = 10000, 2
X = trajectory(he, N, Ttr = 10)
ε = 10 .^ (-5.0:0.1:1)
H = [genentropy(X, e; q) for e in ε]
C = correlationsum(X, ε)

fig, axs = subplots(2, 1; figsize = (figx/2, figy), sharex = true)

axs[1].plot(log.(ε), -H, COLORS[1])

x = -5:0
D = 1.22

axs[2].plot(log.(ε), log.(C), COLORS[1])

for (j, (ax, x)) in enumerate(zip(axs, (-5:0, -10:0)))
    Dl = D .* x
    Dl = Dl .- 2
    ax.axvspan(x[1], x[end], color = COLORS[4], alpha = 0.25)
    ax.grid("on")
    ax.plot(x, Dl, c = COLORS[3])
    ax.text(x[2] + 1.5 + 0.5*(j-1), Dl[2], "\$D_2\$ = $D", color = COLORS[3], size = 30)
end

xtick = -11:3:1
xtickl = string.(xtick)
xtickl[3] = ""
axs[2].set_xticks(xtick)
axs[2].set_xticklabels(xtickl)
axs[2].set_xlabel("\$\\log(\\varepsilon)\$"; labelpad = -20)
# axs[1].set_yticks()
axs[1].set_ylabel("\$-H_$q\$")
axs[2].set_ylabel("\$\\log(C)\$")
axs[1].set_yticks(0:-3:-10)
axs[2].set_yticks(0:-4:-12)
fig.tight_layout(pad=0.3)
wsave(plotsdir("5", "dimension"), fig)


"""
## fromSkyline

Helper function to convert a Skyline vector to a full matrix.

### Type
```julia
fromSkyline(skyline::Vector{Float64}, kdiag::Vector{Int})
```

### Arguments
```julia
* skyline::Vector{Float64}     : 1D Line(nod, nodof)
* kdiag::Vector{Int}         : 2D Triangle(nod, nodof)
```
"""
function fromSkyline(skyline::Vector{Float64}, kdiag::Vector{Int})
	neq = size(kdiag, 1)
	km = zeros(neq, neq)
	km[1, 1] = skyline[kdiag[1]]
	for i in 2:neq
		km[i, i] = skyline[kdiag[i]]
		for j in (kdiag[i-1] + 1):kdiag[i]
			km[i, i - (kdiag[i]-j)] = skyline[j]
			km[i - (kdiag[i]-j), i] = skyline[j]
    end
  end
	km
end


"""Tunes `param` until `expr` evaluates to zero within δ.
`expr` must be monotonically increasing in terms of `param`."""
macro binary_opt(expr, param, min, max, δ)
    :(m = $min;M = $max;z=zero($δ);
    while(true)
        $(esc(param)) = (m+M)/2
        println("$((m+M)/2)")
        e = $(esc(expr))
        if norm(e) < $δ
            break
        elseif e < z
            m = $(esc(param))
        else
            M = $(esc(param))
        end
    end)
end

"""Returns the location and values of local maxima in `sig`."""
function local_maxima(sig)
    l = length(sig)
    res = Tuple{Int, T}[]
    if l < 3
        return res
    end
    @inbounds for i = 2:l-1
        if sig[i-1] < sig[i] && sig[i+1] < sig[i]
            append!(res, (i, sig[i]))
        end
    end
    res
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Julia 1.3.1
#     language: julia
#     name: julia-1.3
# ---

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# # `MeshArrays.jl` defines the `MeshArray` type
#
# Each `MeshArray` contains an array of elementary arrays that (1) are connected at their edges and (2) collectively form a global grid. Overall grid specifications are contained within `gcmgrid` instances, which merely define array sizes and how e.g. grid variables are represented in memory. Importantly, it is only when e.g. grid variables are read from file that sizable memory is allocated.

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Initialize Framework
#
# 1. import `MeshArrays` and plotting tools
# 2. choose e.g. a standard `MITgcm` grid
# 3. download the grid if needed
#

# + {"cell_style": "center", "slideshow": {"slide_type": "subslide"}}
using MeshArrays, Plots

pth="../inputs/GRID_LLC90/"
γ=GridSpec("LatLonCap",pth)

http="https://github.com/gaelforget/GRID_LLC90"
!isdir(pth) ? run(`git clone $http $pth`) : nothing;

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Read Example
#
# A `MeshArray` instance, on the chosen grid, can be obtained from a file path (argument #1). Format conversion occur inside the `read` function based on a propotype argument (#2). `read` / `write` calls then convert back and forth between `MeshArray` and `Array` formats.

# + {"slideshow": {"slide_type": "subslide"}}
D=γ.read(γ.path*"Depth.data",MeshArray(γ,Float64))
tmp1=write(D); tmp2=read(tmp1,D)
show(D)

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Subdomain Arrays
#
# The heatmap function as specialized in `../examples/Plots.jl` operates on each `inner-array` sequentially, one after the other.

# + {"slideshow": {"slide_type": "subslide"}}
p=dirname(pathof(MeshArrays))
include(joinpath(p,"../examples/Plots.jl"))
heatmap(D,title="Ocean Depth",clims=(0.,6000.))

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## A `MeshArray` Behaves Like A `Array`
#
# Here are a few examples that would be coded similarly in both cases

# +
size(D)
eltype(D)
view(D,:)

D .* 1.0
D .* D
1000*D
D*1000

D[findall(D .> 300.)] .= NaN
D[findall(D .< 1.)] .= NaN

D[1]=0.0 .+ D[1]
tmp=cos.(D);
# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## Let's switch grid now

# +
pth="../inputs/GRID_CS32/"
γ=GridSpec("LatLonCap",pth)

http="https://github.com/gaelforget/GRID_CS32"
!isdir(pth) ? run(`git clone $http $pth`) : nothing

γ=GridSpec("CubeSphere",pth)
D=γ.read(γ.path*"Depth.data",MeshArray(γ,Float32))
show(D)

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## The `exchange` Function
#
# It adds neighboring points at face edges to slightly extend the computational domain as often needed e.g. to compute partial derivatives.
# -

Dexch=exchange(D,4)
show(Dexch)

# + {"slideshow": {"slide_type": "subslide"}}
P=heatmap(D.f[6],title="Ocean Depth (D, Face 6)",lims=(-4,36))
Pexch=heatmap(Dexch.f[6],title="...(Dexch, Face 6)",lims=(0,40))
plot(P,Pexch)

# + {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# # Diffusion-Based Smoothing
#
# The unit testing of `MeshArrays.jl` uses the `smooth()` function. Starting from a random noise field, the smoothing efficiency is predictable and can be set via a smoothing scale parameter [(see Weaver and Courtier, 2001)](https://doi.org/10.1002/qj.49712757518).
#
# This example also illustrates the generality of the `MeshArrays` approach, where the same code in `demo2` is readily applicable to any `PeriodicDomain`, `PeriodicChanel`, `CubeSphere`, or `LatLonCap` grid. Here the chosen grid maps onto the `6` faces of a cube with `16*16` points per face, with distances, areas, etc all set to `1.0`.

# + {"slideshow": {"slide_type": "subslide"}}
p=dirname(pathof(MeshArrays))
include(joinpath(p,"../examples/Demos.jl"))
Γ=GridOfOnes("CubeSphere",6,16)
Δ=demo2(Γ);

# + {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# The initial noise field is `D[1]` while the smoothed one is `D[2]`. After `smooth()` has been applied via `demo2()`, the noise field is visibly smoother and more muted.

# + {"cell_style": "split"}
heatmap(Δ[1],title="initial noise",clims=(-0.5,0.5))

# + {"cell_style": "split"}
heatmap(Δ[2],title="smoothed noise",clims=(-0.5,0.5))

# + {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# The computational cost of `smooth()` predictably  increases with the decorrelation scale. For more about how this works, please refer to **Weaver and Courtier, 2001** _Correlation modelling on the sphere using a generalized diffusion equation_ https://doi.org/10.1002/qj.49712757518
# -

Rini=Δ[1]
DXCsm=Δ[3]
DYCsm=Δ[4]
@time Rend=smooth(Rini,DXCsm,DYCsm,Γ);
@time Rend=smooth(Rini,2DXCsm,2DYCsm,Γ);


module Discrete

export ConstSequence, Sequence, DFun;
export fdiff, bdiff, cdiff;
export euler, runge_kutta;
export trapezoid, simpsons;

import FastAnonymous;

immutable ConstSequence
  i_start::Int;
  values::Vector{Float64};

  Sequence(i_start::Int, values::Vector{Float64}) = new(i_start, values);
end

function call(s::ConstSequence, n::Int)
  if n < s.i_start || n > s.i_start + length(s.values)
    error("Out of domain on which sequence is defined");
  end

  return s.values[s.i_start + n];
end

function call(s::ConstSequence, nreal::Real)
  const n = convert(Int, nreal);
  if n < s.i_start || n > s.i_start + length(s.values)
    error("Out of domain on which sequence is defined");
  end

  return s.values[s.i_start + n];
end

typealias DFun     Union{Function, FastAnonymous.Fun};
typealias Sequence Union{DFun, ConstSequence};

#! Compute the forward difference of a sequence
#!
#! \param     seq     Sequence
#! \param     n       Index at which to take compute difference
#! \param     order   Order of difference
#! \return    Forward difference
function fdiff(seq::Sequence, n::Int; order::Int=1)
  @assert(order > 0, "unable to compute $order order forward difference");

  sum = 0.0;
  for i=0:order
    sum += (((order-i)%2==0) ?
              binomial(order, i) * seq(n + i) :
             -binomial(order, i) * seq(n + i));
  end

  return sum;
end

#! Compute the forward difference of a sequence
#!
#! \param     seq     Sequence
#! \param     n_range Range of index for which to compute differences
#! \param     order   Order of difference
#! \return    Forward differences
function fdiff(seq::Sequence, n_range::UnitRange{Int}; order::Int=1)
  return pmap(n -> fdiff(seq, n, order=order), n_range);
end

#! Compute the forward difference of a sequence
#!
#! \param     seq     Sequence
#! \param     n_range Range of index for which to compute differences
#! \param     order   Order of difference
#! \return    Forward differences
function fdiff(seq::Sequence, n_range::UnitRange{Float64}; order::Int=1)
  return pmap(n -> fdiff(seq, n, order=order),
              convert(UnitRange{Int}, n_range));
end

#! Compute the backward difference of a sequence
#!
#! \param     seq     Sequence
#! \param     n       Index at which to take compute difference
#! \param     order   Order of difference
#! \return    Backward difference
function bdiff(seq::Sequence, n::Int; order::Int=1)
  @assert(order > 0, "unable to compute $order order backward difference");

  sum = 0.0;
  for i=0:order
    sum += ((i % 2 == 0) ?
              binomial(order, i) * seq(n - i) :
             -binomial(order, i) * seq(n - i));
  end

  return sum;
end

#! Compute the backward difference of a sequence
#!
#! \param     seq     Sequence
#! \param     n_range Range of index for which to compute differences
#! \param     order   Order of difference
#! \return    Backward differences
function bdiff(seq::Sequence, n_range::UnitRange{Int}; order::Int=1)
  return pmap(n -> bdiff(seq, n, order=order), n_range);
end

#! Compute the backward difference of a sequence
#!
#! \param     seq     Sequence
#! \param     n_range Range of index for which to compute differences
#! \param     order   Order of difference
#! \return    Backward differences
function bdiff(seq::Sequence, n_range::UnitRange{Float64}; order::Int=1)
  return pmap(n -> bdiff(seq, n, order=order),
              convert(UnitRange{Int}, n_range));
end

#! Calculate forward difference on a continuous function
#!
#! \param   f     Continuous function
#! \param   x     Value to calculate difference at
#! \param   h     Step size
#! \param   order Order of difference
#! \retrun  Difference
function fdiff(f::DFun, x::Real, h::Real; order=1)
  sum = 0.0;
  for i=0:order
    sum += ((i % 2 == 0) ? binomial(order, i) * f(x + (n - i) * h) :
                          -binomial(order, i) * f(x + (n - i) * h));
  end
  return sum;
end

#! Calculate backward difference on a continuous function
#!
#! \param   b     Continuous function
#! \param   x     Value to calculate difference at
#! \param   h     Step size
#! \param   order Order of difference
#! \retrun  Difference
function bdiff(f::DFun, x::Real, h::Real; order=1)
  sum = 0.0;
  for i=0:order
    sum += ((i % 2 == 0) ? binomial(order, i) * f(x - i * h) :
                          -binomial(order, i) * f(x - i * h));
  end
  return sum;
end

#! Calculate central difference on a continuous function
#!
#! \param   b     Continuous function
#! \param   x     Value to calculate difference at
#! \param   h     Step size
#! \param   order Order of difference
#! \retrun  Difference
function cdiff(f::DFun, x::Real, h::Real; order=1)
  sum = 0.0;
  for i=0:order
    sum += ((i % 2 == 0) ? binomial(order, i) * f(x + (n/2 - i) * h) :
                          -binomial(order, i) * f(x + (n/2 - i) * h));
  end
  return sum;
end

#! Approximate the solution to an ODE using Euler's method
#!
#! \param   df      Differential equation, df = df(x,t)
#! \param   x0      Initial value of x
#! \param   t0      Initial value of t
#! \param   tf      Final value of t
#! \param   nsteps  Number of timesteps
#! \param   Approximate x(tf)
function euler(df::DFun, x0::Real, t0::Real, tf::Real, nsteps::Int)
  const h   =   (tf - t0) / nsteps;
  x         =   x0;
  t         =   t0;

  for i=1:nsteps
    x += h * df(x, t);
    t += h;
  end

  return x;
end

#! Approximate the solution to an ODE using Runge-Kutta method
#!
#! \param   df      Differential equation, df = df(x,t)
#! \param   x0      Initial value of x
#! \param   t0      Initial value of t
#! \param   tf      Final value of t
#! \param   nsteps  Number of timesteps
#! \param   Approximate x(tf)
function runge_kutta(df::DFun, x0::Real, t0::Real, tf::Real, nsteps::Int)
  const h   =   (tf - t0) / nsteps;
  x         =   x0;
  t         =   t0;

  for i=1:nsteps
    k1    =   h * df(x, t);
    k2    =   h * df(x + 0.5*k1, t + 0.5*h);
    k3    =   h * df(x + 0.5*k2, t + 0.5*h);
    k4    =   h * df(x + k3, t + h);

    x    +=   1.0/6.0 * (k1 + 2*k2 + 2*k3 + k4);
  end

  return x;
end

#! Integrate using trapezoidal rule
#!
#! \param   f     Function to integrate
#! \param   x_0   Starting x
#! \param   x_1   Ending x
#! \param   n     Number of discrete points (must be even)
#! \return        Approximate integral
function trapezoid(f::DFun, x_0::Real, x_1::Real, n::Int)
  const dx  =   (x_1 - x_0) / n;
  sum       =   dx * (f(x_1)+f(x_0)) / 2.0;

  for (i, x) = enumerate(linspace(x_0+dx, x_1-dx, n-1));
    sum += dx * f(x);
  end

  return sum;
end

#! Integrate using simpsons rule
#!
#! \param   f     Function to integrate
#! \param   x_0   Starting x
#! \param   x_1   Ending x
#! \param   n     Number of discrete points (must be even)
#! \return        Approximate integral
function simpsons(f::DFun, x_0::Real, x_1::Real, n::Int)
  @assert(n % 2 == 0, "Number of integration points must be even for "         *
                      "Simpson's rule. $n is not an even number.");

  const dx  =   (x_1 - x_0) / n;
  sum       =   dx * (f(x_1)+f(x_0)) / 3.0;

  for (i, x) = enumerate(linspace(x_0+dx, x_1-dx, n-1));
    sum += dx * f(x) * ((i % 2 == 0) ? 2.0/3.0 : 4.0/3.0);
  end

  return sum;
end

end # end module


using Statistics
using StatsBase
using Tullio
# ------------------------------------------------------------------------------------------------ #
#                             Functions for normalising feature vectors                            #
# ------------------------------------------------------------------------------------------------ #

# ------------------------------------- General normalisation ------------------------------------ #
tanh(x::Number, centre::Number=0, scale::Number=1) = Base.tanh(x - centre)*scale; export tanh
standardise(x::Number, centre::Number=0, scale::Number=1) = (x - centre)/(scale); export standardise
logistic(x::Number, centre::Number=0, scale::Number=1) = 1/(1 + exp(-(x-centre)/scale)); export logistic
centre(x::Number, centre::Number=0, scale::Number=1) = x - centre; export centre

function normalise(x::Union{AbstractVector, Number}, f::Function=standardise, μ::Number=mean(x), σ::Number=std(x))
    y = f.(x, μ, σ)
end
function normalise(X::AbstractArray, f::Function=standardise, dim::Int=2, args...) # Rows by default
    Y = mapslices(x -> normalise(x, f, args...), X, dims=dim)
end
function normalise(X::AbstractArray, 𝛍::AbstractArray, 𝛔::AbstractArray, f::Function=standardise, dim::Int=2)
    # Need a generalised mapslices
    if dim == 1
        @tullio Y[i, j] := normalise(X[i, j], f, 𝛍[j], 𝛔[j])
    elseif dim == 2
        @tullio Y[i, j] := normalise(X[i, j], f, 𝛍[i], 𝛔[i])
    end
end
function normalise(X::AbstractFeatureArray, 𝛍::AbstractFeatureArray, 𝛔::AbstractFeatureArray, f::Function=standardise, dim::Int=2)
    X, 𝛍 = intersectFeatures(X, 𝛍)
    X, 𝛔 = intersectFeatures(X, 𝛔)
    X, 𝛍 = intersectFeatures(X, 𝛍) # In case sigma is different from mu in features
    FeatureArray(normalise(X, vec(𝛍), vec(𝛔), f, dim), getnames(X))
end
export normalise

unitInterval(x::AbstractVector) = normalise(x, standardise, min(x...), abs(-(extrema(x)...)))
unitInterval(X::AbstractMatrix, dim::Int=2) = mapslices(unitInterval, X, dims=dim)
export unitInterval

standardise(x::AbstractVector, args...) = normalise(x, standardise, args...)
standardise(X::AbstractMatrix, dim::Int=2) = mapslices(standardise, X, dims=dim)
export standardise

sigmoidNormalise(x::AbstractVector, args...) = normalise(x, logistic, args...)
sigmoidNormalise(X::AbstractMatrix, dim::Int=2) = mapslices(sigmoidNormalise, X, dims=dim)
export sigmoidNormalise

function robustNormalise(x::Union{AbstractVector, Number}, f::Function=standardise, μ::Number=median(x), σ::Number=iqr(x))
    y = f.(x, μ, σ/1.35)
end
function robustNormalise(X::AbstractArray, f::Function=standardise, dim::Int=2, args...) # Rows by default
    Y = mapslices(x -> robustNormalise(x, f, args...), X, dims=dim)
end
function robustNormalise(X::AbstractArray, 𝛍::AbstractArray, 𝛔::AbstractArray, f::Function=standardise, dim::Int=2)
    if dim == 1
        @tullio Y[i, j] := robustNormalise(X[i, j], f, 𝛍[j], 𝛔[j])
    elseif dim == 2
        @tullio Y[i, j] := robustNormalise(X[i, j], f, 𝛍[i], 𝛔[i])
    end
end
# function robustNormalise(X::AbstractFeatureArray, 𝛍::AbstractFeatureArray, 𝛔::AbstractFeatureArray, f::Function=standardise, dim::Int=2)
#     X, 𝛍 = intersectFeatures(X, 𝛍)
#     X, 𝛔 = intersectFeatures(X, 𝛔)
#     X, 𝛍 = intersectFeatures(X, 𝛍) # In case sigma is different from mu in features
#     robustNormalise(X, vec(𝛍), vec(𝛔), f, dim)
# end
export robustNormalise

robustStandardise(x::AbstractVector, args...) = robustNormalise(x, standardise, args...)
robustStandardise(X::AbstractMatrix, dim::Int=2) = mapslices(standardise, X, dims=dim)
export robustStandardise

robustSigmoidNormalise(x::AbstractVector, args...) = robustNormalise(x, logistic, args...)
robustSigmoidNormalise(X::AbstractMatrix, dim::Int=2) = mapslices(robustSigmoidNormalise, X, dims=dim)
export robustSigmoidNormalise

# ------------------------------------ Scale to L1 unit vector ----------------------------------- #
unitL1(x::AbstractVector) = x./sum(x)
unitL1(X::AbstractMatrix) = mapslices(unitL1, X, dims=2)
export unitL1

# ------------------------------------ Scale to unit interval ------------------------------------ #
# function unitInterval(x::AbstractVector)
#     y = x .- min(x...)
#     y = y./max(y...)
# end
# unitInterval(X::AbstractMatrix) = mapslices(unitInterval, X, dims=2)
# export unitInterval

# ------------------------------------------------------------------------------------------------ #
#                                  Functions for filtering arrays                                  #
# ------------------------------------------------------------------------------------------------ #
function constantrows(F::AbstractArray; tol=1e-10)
    idxs = Array(StatsBase.std(F, dims=2) .< tol)
end
export constantrows

function _noconstantrows(F::AbstractArray; tol=1e-10)
    idxs = constantrows(F; tol=tol)
    if any(idxs)
        @warn "$(+(idxs...)) constant rows are being removed"
        return F -> F[collect(.!idxs)[:], :]
    end
    return F -> F
end
noconstantrows(F::AbstractArray; tol=1e-10) = _noconstantrows(F; tol)(F)
export noconstantrows

function nanrows(F::AbstractArray)
    idxs = any(isnan.(Array(F)), dims=2)
end
export nanrows

function infrows(F::AbstractArray)
    idxs = any(isinf.(Array(F)), dims=2)
end
export nanrows

function _nonanrows(F::AbstractArray)
    idxs = nanrows(F)
    if any(idxs)
        @warn "$(+(idxs...)) NaN rows will be removed" # Maybe more detail?
        return F -> F[collect(.!idxs)[:], :]
    end
    return F -> F
end
nonanrows(F::AbstractArray) = _nonanrows(F)(F)
export nonanrows

function _noinfrows(F::AbstractArray)
    idxs = infrows(F)
    if any(idxs)
        @warn "$(+(idxs...)) Inf rows will be removed"
        return F -> F[collect(.!idxs)[:], :]
    end
    return F -> F
end
noinfrows(F::AbstractArray) = _noinfrows(F)(F)
export noinfrows

function nonans(F::AbstractVector)
    idxs = isnan.(Array(F))
    if any(idxs)
        @warn "$(+(idxs...)) NaN elements are being removed"
    end
    return F[.!idxs]
end
export nonans

function nomissing(F::AbstractVector)
    idxs = ismissing.(Array(F))
    # if any(idxs)
    #     @warn "$(+(idxs...)) missing elements are being removed"
    # end
    return F[.!idxs]
end
export nomissing

function noinf(F::AbstractVector)
    idxs = isinf.(Array(F))
    return F[.!idxs]
end
export noinf


#*******************************************************/
#* Copyright(c) 2018 by Artelys                        */
#* This source code is subject to the terms of the     */
#* MIT Expat License (see LICENSE.md)                  */
#*******************************************************/
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# This file contains routines to implement problemDef.h for
# a simple nonlinear least-squares problem.
#
# min  ( x0*1.309^x1 - 2.138 )^2 +( x0*1.471^x1 - 3.421 )^2 +( x0*1.49^x1 - 3.597 )^2
#        +( x0*1.565^x1 - 4.34 )^2 +( x0*1.611^x1 - 4.882 )^2 +( x0*1.68^x1-5.66 )^2
#
# The standard start point(1.0, 5.0) usually converges to the standard
# minimum at(0.76886, 3.86041), with final objective = 0.00216.
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

using KNITRO
using Test

function example_lsq2(; verbose=true)
    #*------------------------------------------------------------------*
    #*     FUNCTION callbackEvalR                                       *
    #*------------------------------------------------------------------*
    # The signature of this function matches KNITRO.KN_eval_callback in knitro.h.
    # Only "rsd" is set in the KNITRO.KN_eval_result structure.
    function callbackEvalR(kc, cb, evalRequest, evalResult, userParams)
        if evalRequest.evalRequestCode != KNITRO.KN_RC_EVALR
            println(
                "*** callbackEvalR incorrectly called with eval type ",
                evalRequest.evalRequestCode,
            )
            return -1
        end
        x = evalRequest.x

        # Clamp x1 to 1000 so as to avoid overflow
        if x[2] > 1000.0
            x[2] = 1000.0
        end

        # Evaluate nonlinear residual components
        evalResult.rsd[1] = x[1] * 1.309^x[2]
        evalResult.rsd[2] = x[1] * 1.471^x[2]
        evalResult.rsd[3] = x[1] * 1.49^x[2]
        evalResult.rsd[4] = x[1] * 1.565^x[2]
        evalResult.rsd[5] = x[1] * 1.611^x[2]
        evalResult.rsd[6] = x[1] * 1.68^x[2]

        return 0
    end

    #*------------------------------------------------------------------*
    #*     FUNCTION callbackEvalRJ                                      *
    #*------------------------------------------------------------------*
    # The signature of this function matches KNITRO.KN_eval_callback in knitro.h.
    # Only "rsdJac" is set in the KNITRO.KN_eval_result structure.
    function callbackEvalRJ(kc, cb, evalRequest, evalResult, userParams)
        if evalRequest.evalRequestCode != KNITRO.KN_RC_EVALRJ
            println(
                "*** callbackEvalRJ incorrectly called with eval type %d" %
                evalRequest.evalRequestCode,
            )
            return -1
        end
        x = evalRequest.x

        # Clamp x1 to 1000 so as to avoid overflow
        if x[2] > 1000.0
            x[2] = 1000.0
        end

        # Evaluate non-zero residual Jacobian elements(row major order).
        evalResult.rsdJac[1] = 1.309^x[2]
        evalResult.rsdJac[2] = x[1] * log(1.309) * 1.309^x[2]
        evalResult.rsdJac[3] = 1.471^x[2]
        evalResult.rsdJac[4] = x[1] * log(1.471) * 1.471^x[2]
        evalResult.rsdJac[5] = 1.49^x[2]
        evalResult.rsdJac[6] = x[1] * log(1.49) * 1.49^x[2]
        evalResult.rsdJac[7] = 1.565^x[2]
        evalResult.rsdJac[8] = x[1] * log(1.565) * 1.565^x[2]
        evalResult.rsdJac[9] = 1.611^x[2]
        evalResult.rsdJac[10] = x[1] * log(1.611) * 1.611^x[2]
        evalResult.rsdJac[11] = 1.68^x[2]
        evalResult.rsdJac[12] = x[1] * log(1.68) * 1.68^x[2]

        return 0
    end

    #*------------------------------------------------------------------*
    #*     main                                                         *
    #*------------------------------------------------------------------*

    # Create a new Knitro solver instance.
    kc = KNITRO.KN_new()

    # Add the variables/parameters.
    # Note: Any unset lower bounds are assumed to be
    # unbounded below and any unset upper bounds are
    # assumed to be unbounded above.
    n = 2 # # of variables/parameters
    KNITRO.KN_add_vars(kc, n)

    # In order to prevent the possiblity of numerical
    # overflow from very large numbers, we set a
    # reasonable upper bound on variable x[1] and set the
    # "honorbnds" option for this variable to enforce
    # that all trial x[1] values satisfy this bound.
    KNITRO.KN_set_var_upbnd(kc, 1, 100.0)
    KNITRO.KN_set_var_honorbnd(kc, 1, KNITRO.KN_HONORBNDS_ALWAYS)

    # Add the residuals.
    m = 6 # # of residuals
    KNITRO.KN_add_rsds(kc, m)

    # Set the array of constants in the residuals
    KNITRO.KN_add_rsd_constants_all(kc, [-2.138, -3.421, -3.597, -4.34, -4.882, -5.66])

    # Add a callback function "callbackEvalR" to evaluate the nonlinear
    # residual components.  Note that the constant terms are added
    # separately above, and will not be included in the callback.
    cb = KNITRO.KN_add_lsq_eval_callback(kc, callbackEvalR)

    # Also add a callback function "callbackEvalRJ" to evaluate the
    # Jacobian of the residuals.  If not provided, Knitro will approximate
    # the residual Jacobian using finite-differencing.  However, we recommend
    # providing callbacks to evaluate the exact Jacobian whenever
    # possible as this can drastically improve the performance of Knitro.
    # We specify the residual Jacobian in "dense" row major form for simplicity.
    # However for models with many sparse residuals, it is important to specify
    # the non-zero sparsity structure of the residual Jacobian for efficiency
    #(this is true even when using finite-difference gradients).
    KNITRO.KN_set_cb_rsd_jac(kc, cb, KNITRO.KN_DENSE_ROWMAJOR, callbackEvalRJ)

    kn_outlev = verbose ? KNITRO.KN_OUTLEV_ALL : KNITRO.KN_OUTLEV_NONE
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_OUTLEV, kn_outlev)

    # Solve the problem.
    #
    # Return status codes are defined in "knitro.h" and described
    # in the Knitro manual.

    nRC = KNITRO.KN_solve(kc)

    # An example of obtaining solution information.
    # Return status codes are defined in "knitro.h" and described
    # in the Knitro manual.
    nStatus, obj, x, lambda_ = KNITRO.KN_get_solution(kc)

    if nStatus != 0 && verbose
        println("Knitro successful. The optimal solution is:")
        for i in 1:n
            println("x[$i]= ", x[i])
        end
    end

    # Delete the knitro solver instance.
    KNITRO.KN_free(kc)

    @testset "Example LSQ2" begin
        @test nStatus == 0
        @test obj ≈ 0.00216 atol = 1e-4
        @test x ≈ [0.76886, 3.86041] atol = 1e-4
    end
end

example_lsq2(; verbose=isdefined(Main, :KN_VERBOSE) ? KN_VERBOSE : true)


"""
Takes a time series and produces a matrix containing all the possible lags.

Those points which do not have lags of that magnitude are set to Null.
"""
function lagmatrix{T}(x::AbstractVector{T})
    n = length(x)
    X = fill(Nullable{T}(),n,n)
    for j in 1:n, i in 1:n
        X[i,j] = j-i>=0 ? x[j-i+1] : Nullable()
        end
    X
end

"""
Create a lagmatrix for any given Quantity
"""
function makelagmatrix(q::Quantity)
    lagmatrix(quantity(q))
end

function preparedata(p::Stage,d::Discharge,range,M)
    H = makelagmatrix(p)
    Q = quantity(d)
    Hm = validate(H,range,M)
    Qm = validate(Q,range,M)
    Hm,Qm
end

"""
Given a matrix formed by `lagmatrix`, a range of valid values desired and a number of lags,
return those data points which are valid data
"""
function validate{T}(X::Matrix{Nullable{T}},range,n::Int)
    X = X[:,range]
    i = 1
    while any(map(isnull,X[1:n,i]))
        i+=1
    end
    _,m = size(X)
    Y = zeros(T,n,m-i+1)
    for k in 1:n, j in i:m
        Y[k,j-i+1] = get(X[k,j])
    end
    Y
end

"""
The same idea, but give all the valid response values from a vector of discharges

Try not to pass in a range which is invalid for the vector
"""
function validate(x::AbstractVector,range,n::Int)
    if range.start < n
        return x[n:range.stop]
    else
        return x[range]
    end
end

include("statsutils.jl")


# computational utilities

"""
    θ(x)

Heaviside step function. Returns one if the argument is positive, zero otherwise.
"""
function θ(x::T) where T
    if x > zero(T)
        return one(T)
    end
    return zero(T)
end

"""
    xθ(x)

Heaviside step times x. Returns x if the argument is positive, zero otherwise.
"""
function xθ(x::T) where T
    if x > zero(T)
        return x
    end
    return zero(T)
end

"""
    wavevectors3D([T=Float64], dims, box_size=(2π, 2π, 2π))

# Arguments:
- `T` : output element type
- `dims`: array or tuple containing size
- `box_size`: side lengths of box, normalizes the wavenumber

# Returns:
- `Tuple{Vector{T},Vector{T},Vector{T}}`: kx, ky, kz wavevectors

# Examples
```julia-repl
julia> CosmoMMF.wavevectors3D((6,2,2))
([0.0, 1.0, 2.0, -3.0, -2.0, -1.0], [0.0, -1.0], [0.0, -1.0])
```
"""
function wavevectors3D(T::Type{<:Real}, dims, box_size=(2π, 2π, 2π))
    sample_rate = T.(2π .* dims ./ box_size)
    kx = fftfreq(dims[1], sample_rate[1]) .* (2π / dims[1])
    ky = fftfreq(dims[2], sample_rate[2]) .* (2π / dims[2])
    kz = fftfreq(dims[3], sample_rate[3]) .* (2π / dims[3])
    return kx, ky, kz
end

# default T is Float64
wavevectors3D(dims, box_size=(2π, 2π, 2π)) = wavevectors3D(Float64, dims, box_size)

function test_print()
    print("This test function works! (totally)")
end

function wall(n::Int64)
    """
    Builds a wall out of a nxnxn array. All values in the wall
    are = 1. Other points = 0.
    """
    array = zeros((n,n,n))
    array[:,:,:] .= 0.1
    index = n ÷ 2
    array[index,:,:] .= 1
    return array
end

function cylinder(n::Int64,r::Float64)
    """
    Returns a nxnxn array with a cylinder in the center with
    radius r. All points within the radius are = 1, the rest
    are 0.
    """
    array  = zeros(Float64,(n,n,n))
    array[:,:,:] .= 0.1
    for x::Int64 in 1:n
        for y::Int64 in 1:n
            for z::Int64 in 1:n
                if ((x - (n)/2)^2 + (y - (n)/2)^2) < r^2
                    array[x,y,z] = 1
                end
            end
        end
    end
    return array
end

function sphere(n::Int64,r::Float64)
    """
    Returns a nxnxn array with a sphere in the center with
    radius r. All points within the radius are = 1, the rest
    are 0.
    """
    array = zeros(Float64,(n,n,n))
    array[:,:,:] .= 0.1
    for x::Int64 in 1:n
        for y::Int64 in 1:n
            for z::Int64 in 1:n
                if ((x - (n)/2)^2 + (y - (n)/2)^2 + (z - (n)/2)^2) < r^2
                    array[x,y,z] = 1
                end
            end
        end
    end
    return array
end

function save_max_sigs(output_directory, id_string, sig_array::AbstractArray{T,4}) where T
    """
    A simple function to save the signatures to an output directory in the form of a jld2 file.
    """
    @save output_directory * id_string * ".jld2" sig_array
end

function save_hessian_component(output_directory, Rs, id_string, hessian_component::AbstractArray{T,3}) where T
    """
    A helper function to save a hessian component.

    Args:
        output_directory [string]: where should the hessian component be saved to?
        id_string [string]: component number in matrix
                            -> [[1 2 3],
                                [  4 5],
                                [    6]]
        subvolume_divisions [Int]: number of subvolumes to divide original hessian into on each side. If = 2, then we make 2^3 subvolumes. (Default = 2)
        hessian_component [AbstractArray{T,3}]: one of the components
    """
    dims = size(hessian_component)

    div_x, div_y, div_z = dims .÷ 2

    #chop hessian into 8ths
    hessian_component_000 = hessian_component[1:div_x, 1:div_y, 1:div_z]
    hessian_component_001 = hessian_component[1:div_x, 1:div_y, div_z+1:dims[3]]
    hessian_component_010 = hessian_component[1:div_x, div_y+1:dims[2], 1:div_z]
    hessian_component_100 = hessian_component[div_x+1:dims[1], 1:div_y, 1:div_z]
    hessian_component_011 = hessian_component[1:div_x, div_y+1:dims[2], div_z+1:dims[3]]
    hessian_component_110 = hessian_component[div_x+1:dims[1], div_y+1:dims[2], 1:div_z]
    hessian_component_101 = hessian_component[div_x+1:dims[1], 1:div_y, div_z+1:dims[3]]
    hessian_component_111 = hessian_component[div_x+1:dims[1], div_y+1:dims[2], div_z+1:dims[3]]

    @save output_directory * "hessian_component_000" * "_" * string(Rs) * id_string * ".jld2" hessian_component_000
    @save output_directory * "hessian_component_001" * "_" * string(Rs) * id_string * ".jld2" hessian_component_001
    @save output_directory * "hessian_component_010" * "_" * string(Rs) * id_string * ".jld2" hessian_component_010
    @save output_directory * "hessian_component_100" * "_" * string(Rs) * id_string * ".jld2" hessian_component_100
    @save output_directory * "hessian_component_011" * "_" * string(Rs) * id_string * ".jld2" hessian_component_011
    @save output_directory * "hessian_component_110" * "_" * string(Rs) * id_string * ".jld2" hessian_component_110
    @save output_directory * "hessian_component_101" * "_" * string(Rs) * id_string * ".jld2" hessian_component_101
    @save output_directory * "hessian_component_111" * "_" * string(Rs) * id_string * ".jld2" hessian_component_111

end

function compile_sub_hessians(f_Rn::AbstractArray{T,3}, Rs, subvolume, input_directory) where T

    id_strings = ["_1", "_2", "_3", "_4", "_5", "_6"]
    dims = size(f_Rn) .÷ 2
    hessian_sub = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))

#     for (i, id_string) in enumerate(id_strings)
#         @load input_directory * "hessian_component_000" * id_string * ".jld2" hessian_component_000
#         hessian_sub[:,:,:,i] .= hessian_component_000
#         #figure out how to delete files?
#     end

#     @save input_directory * "hessian_component_000" * "_full" * ".jld2" real(hessian_sub)
#     hessian_sub = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))

#     for (i, id_string) in enumerate(id_strings)
#         @load input_directory * "hessian_component_001" * id_string * ".jld2" hessian_component_001
#         hessian_sub[:,:,:,i] .= hessian_component_001
#         #figure out how to delete files?
#     end

#     @save input_directory * "hessian_component_001" * "_full" * ".jld2" real(hessian_sub)
#     hessian_sub = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))

#     for (i, id_string) in enumerate(id_strings)
#         @load input_directory * "hessian_component_010" * id_string * ".jld2" hessian_component_010
#         hessian_sub[:,:,:,i] .= hessian_component_010
#         #figure out how to delete files?
#     end

#     @save input_directory * "hessian_component_010" * "_full" * ".jld2" real(hessian_sub)
#     hessian_sub = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))

#     for (i, id_string) in enumerate(id_strings)
#         @load input_directory * "hessian_component_100" * id_string * ".jld2" hessian_component_100
#         hessian_sub[:,:,:,i] .= hessian_component_100
#         #figure out how to delete files?
#     end

#     @save input_directory * "hessian_component_100" * "_full" * ".jld2" real(hessian_sub)
#     hessian_sub = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))

#     for (i, id_string) in enumerate(id_strings)
#         @load input_directory * "hessian_component_011" * id_string * ".jld2" hessian_component_011
#         hessian_sub[:,:,:,i] .= hessian_component_011
#         #figure out how to delete files?
#     end

#     @save input_directory * "hessian_component_011" * "_full" * ".jld2" real(hessian_sub)
#     hessian_sub = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))

#     for (i, id_string) in enumerate(id_strings)
#         @load input_directory * "hessian_component_110" * id_string * ".jld2" hessian_component_110
#         hessian_sub[:,:,:,i] .= hessian_component_110
#         #figure out how to delete files?
#     end

#     @save input_directory * "hessian_component_110" * "_full" * ".jld2" real(hessian_sub)
#     hessian_sub = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))

#     for (i, id_string) in enumerate(id_strings)
#         @load input_directory * "hessian_component_101" * id_string * ".jld2" hessian_component_101
#         hessian_sub[:,:,:,i] .= hessian_component_101
#         #figure out how to delete files?
#     end

#     @save input_directory * "hessian_component_101" * "_full" * ".jld2" real(hessian_sub)
#     hessian_sub = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))

    for (i, id_string) in enumerate(id_strings)
        if subvolume == "000"
            @load input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2" hessian_component_000
        hessian_sub[:,:,:,i] .= hessian_component_000

        #delete old slice we don't need anymore
        file = input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2"
        run(`rm $file`)

        elseif subvolume == "001"
            @load input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2" hessian_component_001
        hessian_sub[:,:,:,i] .= hessian_component_001

        #delete old slice we don't need anymore
        file = input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2"
        run(`rm $file`)

        elseif subvolume == "010"
            @load input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2" hessian_component_010
        hessian_sub[:,:,:,i] .= hessian_component_010

        #delete old slice we don't need anymore
        file = input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2"
        run(`rm $file`)

        elseif subvolume == "100"
            @load input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2" hessian_component_100
        hessian_sub[:,:,:,i] .= hessian_component_100

        #delete old slice we don't need anymore
        file = input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2"
        run(`rm $file`)

        elseif subvolume == "011"
            @load input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2" hessian_component_011
        hessian_sub[:,:,:,i] .= hessian_component_011

        #delete old slice we don't need anymore
        file = input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2"
        run(`rm $file`)

        elseif subvolume == "110"
            @load input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2" hessian_component_110
        hessian_sub[:,:,:,i] .= hessian_component_110

        #delete old slice we don't need anymore
        file = input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2"
        run(`rm $file`)

        elseif subvolume == "101"
            @load input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2" hessian_component_101
        hessian_sub[:,:,:,i] .= hessian_component_101

        #delete old slice we don't need anymore
        file = input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2"
        run(`rm $file`)

        elseif subvolume == "111"
            @load input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2" hessian_component_111
        hessian_sub[:,:,:,i] .= hessian_component_111

        #delete old slice we don't need anymore
        file = input_directory * "hessian_component_" * subvolume * "_" * string(Rs) * id_string * ".jld2"
        run(`rm $file`)

        else
            print("Give a correct subvolume string. Example: '000' ")
        end

    end

    hessian_sub = real(hessian_sub)
    @save input_directory * "hessian_component_" * subvolume * "_full" * ".jld2" hessian_sub

end


n = 20
Z = rand(1:5, n)
ϑ = dropdims(ones(1,n) + rand(1,n), dims = 1)
μ = mean(ϑ)
kmax = 4

ω(x, α) = (x[1]+1)^(-α[x[2]]*x[2])

α = repeat([2.0], kmax)

Ω = sumOfExteriorDegreesIntensityFunction(ω, kmax)
H = sampleSBM(Z, ϑ, Ω;α=α, kmax=kmax, kmin = 1)
D = H.D
all_partitions = collect(Iterators.flatten([(partitions(i,j)) for i =1:kmax for j=1:i]))

p = [3, 1, 1]

s1 = evalSumNaive(p, Z, D)
s2 = evalSumNaive2(p, Z, D)

## d
@test s1 == s2

sumsNaive = evalSumsNaive(Z,D,kmax)
sumsPV    = evalSumsPV(Z,D,kmax)
sums      = evalSums(Z,D,kmax)[3]

@test all([sumsNaive[p] == sumsPV[p] for p in all_partitions])
@test all([sumsNaive[p] == sums[p] for p in all_partitions])

# voldiff

ℓ = maximum(Z) # number of total groups

# let's move all the nodes in group 4 to group 5
to_move = findall(==(4), Z)
t = 5
Z_ = copy(Z)

## first we'll step through, moving each node one-by-one
V, μ, M = evalSums(Z, D, kmax; constants=false, bigInt=false)

for i in to_move
    global V, μ, M
    ΔV, Δμ, ΔM = increments(V, μ, M, i, t, D, Z_)
    V, μ, M = HyperModularity.addIncrements(V, μ, M, ΔV, Δμ, ΔM)
    # carry out the change in membership
    Z_[i] = t
end
step_wise = M

# next we'll compute the complete set of increments at once
V, μ, M = evalSums(Z, D, kmax; constants=false, bigInt=false);
ΔV, Δμ, ΔM = increments(V, μ, M, to_move, t, D, Z)
V, μ, M = HyperModularity.addIncrements(V, μ, M, ΔV, Δμ, ΔM)
batch = M
#
@test all([step_wise[p] == batch[p] for p in keys(M)])


#
# Methods for orbit states and state vectors.
#

keplerian(state::CartesianStateVector, μ) = keplerian(position(state), velocity(state), μ)
keplerian(orbit::CartesianR2BPOrbit) = keplerian(state(orbit), massparameter(system(orbit)))
keplerian(orbit::KeplerianR2BPOrbit) = orbit

cartesian(orbit::CartesianR2BPOrbit) = orbit
cartesian(orbit::KeplerianR2BPOrbit) = cartesian(state(orbit), massparameter(system(orbit)))
cartesian(state::KeplerianState, μ) = cartesian(
    eccentricity(state),
    semimajor_axis(state),
    inclination(state),
    RAAN(state),
    argument_of_periapsis(state),
    true_anomoly(state),
    μ
)

Base.position(state::CartesianStateVector) = OrbitalStates.get_r(state) * lengthunit(state)
Base.position(orbit::KeplerianR2BPOrbit) = position(CartesianState(cartesian(state(orbit), massparameter(state(orbit)))))
Base.position(orbit::Orbit) = position(state(orbit))

distance(state::CartesianStateVector) = norm(position(state))
distance(orbit::Orbit) = norm(position(orbit))

velocity(state::CartesianStateVector) = OrbitalStates.get_v(state) * velocityunit(state)
velocity(orbit::Orbit) = velocity(state(orbit))

speed(state::CartesianStateVector) = norm(velocity(state))
speed(orbit::KeplerianR2BPOrbit) = speed(semi_parameter(semimajor_axis(state(orbit)), eccentricity(state(orbit))), semimajor_axis(state(orbit)), massparameter(system(orbit)))
speed(orbit::Orbit) = norm(velocity(state(orbit)))

eccentricity(state::KeplerianState) = OrbitalStates.get_e(state)
eccentricity(orbit::KeplerianR2BPOrbit) = eccentricity(state(orbit))
eccentricity(orbit::CartesianR2BPOrbit) = eccentricity(position(state(orbit)), velocity(state(orbit)), massparameter(system(orbit)))

semimajor_axis(state::KeplerianState) = OrbitalStates.get_a(state) * lengthunit(state)
semimajor_axis(orbit::CartesianR2BPOrbit) = semimajor_axis(distance(state(orbit)), speed(state(orbit)), massparameter(system(orbit)))
semimajor_axis(orbit::KeplerianR2BPOrbit) = semimajor_axis(state(orbit))

inclination(state::KeplerianState) = OrbitalStates.get_i(state) * angularunit(state)
inclination(orbit::KeplerianR2BPOrbit) = inclination(state(orbit))

RAAN(state::KeplerianState) = OrbitalStates.get_Ω(state) * angularunit(state)
RAAN(orbit::KeplerianR2BPOrbit) = RAAN(state(orbit))

argument_of_periapsis(state::KeplerianState) = OrbitalStates.get_ω(state) * angularunit(state)
argument_of_periapsis(orbit::KeplerianR2BPOrbit) = argument_of_periapsis(state(orbit))

true_anomoly(state::KeplerianState) = OrbitalStates.get_ν(state) * angularunit(state)
true_anomoly(orbit::KeplerianR2BPOrbit) = true_anomoly(state(orbit))
true_anomoly(orbit::CartesianR2BPOrbit) = true_anomoly(distance(orbit), specific_angular_momentum(orbit), eccentricity(orbit), massparameter(system(orbit)))

conic(orbit::R2BPOrbit) = conic(eccentricity(orbit))

specific_angular_momentum_vector(orbit::CartesianR2BPOrbit) = specific_angular_momentum_vector(position(state(orbit)), velocity(state(orbit)))

specific_angular_momentum(orbit::CartesianR2BPOrbit) = specific_angular_momentum(position(state(orbit)), velocity(state(orbit)))

specific_energy(orbit::CartesianR2BPOrbit) = specific_energy(distance(state(orbit)), speed(state(orbit)), massparameter(system(orbit)))
specific_energy(orbit::KeplerianR2BPOrbit) = specific_energy(semimajor_axis(orbit), massparameter(system(orbit)))

C3(orbit::R2BPOrbit) = C3(distance(orbit), speed(orbit), massparameter(system(orbit)))

v_infinity(orbit::R2BPOrbit) = v_infinity(distance(orbit), speed(orbit), massparameter(system(orbit)))

eccentricity_vector(orbit::CartesianR2BPOrbit) = eccentricity_vector(position(state(orbit)), velocity(state(orbit)), massparameter(system(orbit)))

semi_parameter(orbit::R2BPOrbit) = semi_parameter(semimajor_axis(orbit), eccentricity(orbit))

periapsis_radius(orbit::R2BPOrbit) = periapsis_radius(semimajor_axis(orbit), eccentricity(orbit))

apoapsis_radius(orbit::R2BPOrbit) = apoapsis_radius(semimajor_axis(orbit), eccentricity(orbit))

periapsis_velocity(orbit::R2BPOrbit) = speed(periapsis_radius(orbit), semimajor_axis(orbit), massparameter(system(orbit)))

apoapsis_velocity(orbit::R2BPOrbit) = speed(apoapsis_radius(orbit), semimajor_axis(orbit), massparameter(system(orbit)))

period(orbit::R2BPOrbit) = conic(orbit) ∈ (Circular, Elliptical) ? period(semimajor_axis(orbit), massparameter(system(orbit))) : conic(orbit) == Parabolic ? Inf * timeunit(orbit) : NaN * timeunit(orbit)

mean_motion(orbit::R2BPOrbit) = mean_motion(semimajor_axis(orbit), massparameter(system(orbit)))

function mean_motion_vector(orbit::R2BPOrbit)
#   î = SVector{3, Float64}([1, 0, 0])
#   ĵ = SVector{3, Float64}([0, 1, 0])
    k̂ = SVector{3}(0, 0, 1)
    return k̂ × specific_angular_momentum_vector(orbit)
end

time_since_periapsis(orbit::R2BPOrbit) = time_since_periapsis(mean_motion(orbit), eccentricity(orbit), eccentric_anomoly(orbit))

specific_potential_energy(orbit::CartesianR2BPOrbit) = specific_potential_energy(distance(orbit.state), massparameter(system(orbit)))


# simple program to test the new k-means (not ready yet)

using Base.Test
using Clustering

srand(34568)

m = 3
n = 1000
k = 10

x = rand(m, n)

# non-weighted
r = kmeans(x, k; maxiter=50)
@test isa(r, KmeansResult{Float64})
@test size(r.centers) == (m, k)
@test length(r.assignments) == n
@test all(r.assignments .>= 1) && all(r.assignments .<= k)
@test length(r.costs) == n
@test length(r.counts) == k
@test sum(r.counts) == n
@test r.cweights == @compat(map(Float64, r.counts))
@test_approx_eq sum(r.costs) r.totalcost

# non-weighted (float32)
r = kmeans(@compat(map(Float32, x)), k; maxiter=50)
@test isa(r, KmeansResult{Float32})
@test size(r.centers) == (m, k)
@test length(r.assignments) == n
@test all(r.assignments .>= 1) && all(r.assignments .<= k)
@test length(r.costs) == n
@test length(r.counts) == k
@test sum(r.counts) == n
@test r.cweights == @compat(map(Float64, r.counts))
@test_approx_eq sum(r.costs) r.totalcost

# weighted
w = rand(n)
r = kmeans(x, k; maxiter=50, weights=w)
@test isa(r, KmeansResult{Float64})
@test size(r.centers) == (m, k)
@test length(r.assignments) == n
@test all(r.assignments .>= 1) && all(r.assignments .<= k)
@test length(r.costs) == n
@test length(r.counts) == k
@test sum(r.counts) == n

cw = zeros(k)
for i = 1:n
    cw[r.assignments[i]] += w[i]
end
@test_approx_eq r.cweights cw

@test_approx_eq dot(r.costs, w) r.totalcost


using Test
using CompScienceMeshes
using BEAST

# G0 is the interior of the metal (no simulation required here)
# G1 is the unbounded exterior
# G2 is the bounded interior

fn = joinpath(@__DIR__, "assets", "thick_cladding.msh")
G01 = CompScienceMeshes.read_gmsh_mesh(fn, physical="Gamma01")
G02 = CompScienceMeshes.read_gmsh_mesh(fn, physical="Gamma02")
G12 = CompScienceMeshes.read_gmsh_mesh(fn, physical="Gamma12")

G = CompScienceMeshes.read_gmsh_mesh(fn)

@test numcells(G01) + numcells(G02) + numcells(G12) == numcells(G)

# Build region boundaries with normals pointing inward:
G1 = weld(-G01, -G12)
@test CompScienceMeshes.isoriented(G1)

G2 = weld(-G02, G12)
@test CompScienceMeshes.isoriented(G2)

isclosed(m) = (numcells(boundary(m)) == 0)

@test isclosed(G1)
@test isclosed(G2)

# Test the validity of the embedding matrices
E = skeleton(G,1)
E1 = skeleton(G1,1)
A1 = CompScienceMeshes.embedding(E1,E)

nd = BEAST.nedelec(G,E)
@test numfunctions(nd) == numcells(E)
nd1 = BEAST.nedelec(G1,E1)
@test numfunctions(nd1) == numcells(E1)
Z_direct = assemble(BEAST.Identity(), nd1, nd)
Z_expand = assemble(BEAST.Identity(), nd1, nd1) * A1
@test norm(Z_direct - Z_expand, Inf) ≤ √(eps(1.0))


SetID = Int64

"不确定有穷自动机"
struct FiniteAutomata
    𝐊::StateSet             # 状态集
    𝚺::SymbolTabel          # 字母表
    f::Transition{StateSet} # 状态映射
    𝐒::StateSet             # 开始状态集合
    𝐙::StateSet             # 接受状态集合
end

"可达状态集合"
function move(fa::FiniteAutomata, 𝐈::StateSet, α::Symbol)
    states = StateSet()
    for s in 𝐈
        union!(states, fa.f[s][α])
    end
    states
end

"ϵ-closure ϵ-闭包"
function ϵ_colsure(fa::FiniteAutomata, 𝐈::StateSet)
    states = StateSet(𝐈)
    queue = collect(State, 𝐈)
    while !isempty(queue)
        s = popfirst!(queue)
        for next in get(fa.f[s], :ϵ, StateSet())
            if next ∉ states
                push!(states, next)
                push!(queue, next)
            end
        end
    end
    states
end

"是否为确定的有穷自动机"
function isDeterministic(fa::FiniteAutomata)
    for (_, t) in fa.f
        for (_, s) in t
            "若可达状态数大于`1`,则该状态机为不确定的"
            length(s) > 1 && return false
        end
    end
    true
end

"确定的有穷自动机"
struct DeterministicFiniteAutomata
    𝐊::StateSet             # 状态集合
    𝚺::SymbolTabel          # 字母表
    f::Transition{State}    # 状态转移
    𝐒::State                # 开始状态
    𝐙::StateSet             # 接受状态集合
end

function createSym!(tb::Dict{Symbol,StateSet}, counter::Ref{Int64}, sts::StateSet)
    sym = Symbol("St_$(counter)")
    tb[sym] = sts
    counter += 1
    return sym
end

"子集法产生新的状态集"
function subset(fa::FiniteAutomata)
    counter = 1
    table = Dict{State,StateSet}()
    f = Transition{State}()
    𝐂₀ = ϵ_colsure(fa, fa.𝐊)
    s₀ = createSym!(table, Ref(counter), 𝐂₀)
    𝐂 = Tuple{StateSet,State}[(𝐂₀, s₀),]
    𝐓 = Set{StateSet}()
        while !isempty(𝐂)
        (T, s) = popfirst!(𝐂)
        push!(𝐓, T)
        for α in fa.𝚺
            𝐔 = ϵ_colsure(fa, move(fa, T, α))
            𝐔 ∈ 𝐓 && continue
            sym = createSym!(table, Ref(counter), 𝐔)
            !haskey(f, s) && f[s] = Dict{Symbol,State}()
            f[s][α] = sym
            push!(𝐂, (𝐔, sym))
        end
    end
    table, f
end

"Transform NFA to DFA"
function determine(fa::FiniteAutomata)
    𝐊, f = subset(fa)
    S = Symbol("St_1")
    𝐙 = filter(keys(𝐊)) do s
        any(∈(fa.𝐙), 𝐊[s])
    end
    DeterministicFiniteAutomata(keys(𝐊), fa.𝚺, f, S, 𝐙)
end

function remove_useless!(dfa::DeterministicFiniteAutomata)
    queue = State[dfa.𝐒,]
    𝐊 = StateSet()
    push!(𝐊, dfa.𝐒)
    while !isempty(queue)
        s = popfirst!(queue)
        for α in dfa.𝚺
            !haskey(dfa.f[s], α) && continue
            ns = dfa.f[s][α]
            ns ∈ 𝐊 && continue
            push!(queue, ns)
            push!(𝐊, ns)
        end
    end
    empty!(dfa.𝐊)
    union!(dfa.𝐊, 𝐊)
end

function state2stid(𝐏::Vector{StateSet})
    set_tabel = Dict{State,SetID}()
    for (i, st) in 𝐏
        for s ∈ st
            set_tabel[s] = i
        end
    end
    set_tabel
end

function segmentation!(𝐏::Vector{StateSet}, f::Transition{State}, 𝚺::SymbolTabel)
    set_tabel = state2stid(𝐏)
    i = 1
    while i <= length(𝐏)
        stateset = 𝐏[i]
        length(stateset) == 1 && i += 1 && continue
        cnt = 0
        for α ∈ 𝚺
            new_set = Dict{SetID,StateSet}()
            for state in stateset
                next_state = get(f[state], α, state)
                state_id = set_tabel[next_state]
                !haskey(new_set, state_id) && new_set[state_id] = StateSet()
                push!(new_set[state_id], state)
            end
            if length(new_set) > 1
                for set in values(new_set)
                    splice!(𝐏, i, set)
                    foreach(x -> set_tabel[x] = length(𝐏), set)
                end
                break
            end
            cnt += 1
        end
        cnt == length(𝚺) && i += 1
    end
    𝐏
end

function reset!(f::Transition{State}, src::State, dst::State)
    for k in keys(f)
        if k == src
            delete!(f, k)
        else
            for e in keys(f[k])
                v = f[k][e]
                v == src && f[k][e] = dst
            end
        end
    end
    f
end

"DFA化简（最小化）"
function simplify!(dfa::DeterministicFiniteAutomata)
    remove_useless!(dfa)
    # 分割法
    𝐏 = StateSet[setdiff(dfa.𝐊, dfa.𝐙),dfa.𝐙]
    segmentation!(𝐏, dfa.f, dfa.𝚺)
    for set in 𝐏
        state = pop!(set)
        while !isempty(set)
            removed = pop!(set)
            reset!(dfa.f, removed, state)
        end
    end
    dfa
end


module Bautin

using DiffEqBase: ODEProblem
using Parameters: @with_kw, @unpack
using StaticArrays: SVector
using Setfield: @lens

using ...Bifurcations: BifurcationProblem
using ...Codim2: DiffEqCodim2Problem

@with_kw struct BautinParam{B1, B2, S}
    β₁::B1 = -1.0
    β₂::B2 = -1.0
    s::S = -1
end

function f(u::SVector, p, t)
    @unpack β₁, β₂, s = p
    x, y = u
    #=
    ρ² = x^2 + y^2
    ρ⁴ = ρ²^2
    =#
    ρ² = x * x + y * y
    ρ⁴ = ρ² * ρ²
    a = β₂ * ρ² + s * ρ⁴
    dx = β₁ * x -      y + a * x
    dy =      x + β₁ * y + a * y
    return SVector(dx, dy)
end

# TODO: make complex Bautin normal form work
#=
function f(u::SVector, p, t)
    dz = f(u[1] + u[2] * im, p, t)
    return SVector(real(dz), imag(dz))
end

function f(z::Complex, p, t)
    @unpack β₁, β₂, s = p
    ρ² = abs2(z)
    ρ⁴ = ρ²^2
    return (β₁ + 1im) * z + β₂ * z * ρ² + s * z * ρ⁴
end
=#

f(u::A, p, t) where {A <: AbstractVector} = A(f(SVector{2}(u), p, t))

function f(du, u, p, t)
    du .= f(SVector{2}(u), p, t)
    nothing
end

make_prob(
        p = BautinParam();
        u0 = SVector(0.0, 0.0),
        tspan = (0.0, 30.0),
        ode = ODEProblem{!(u0 isa SVector)}(f, u0, tspan, p),
        param_axis = (@lens _.β₁),
        t_domain = (-2.0, 2.0),
        kwargs...) =
    BifurcationProblem(ode, param_axis, t_domain;
                       kwargs...)

prob = make_prob()
ode = prob.p.de_prob
u0 = ode.u0
param_axis = prob.p.param_axis

function make_codim2_prob(
        p = ode.p;
        u0 = SVector(0.0, 0.0),
        tspan = (0.0, 30.0),
        ode = ODEProblem{!(u0 isa SVector)}(f, u0, tspan, p),
        param_axis1 = (@lens _.β₁),
        param_axis2 = (@lens _.β₂),
        t_domain = ([-2.0, -2.0], [2.0, 2.0]),
        kwargs...)
    return DiffEqCodim2Problem(
        ode,
        param_axis1,
        param_axis2,
        t_domain;
        kwargs...)
end

end  # module


immutable Gamma <: ContinuousUnivariateDistribution
    α::Float64
    β::Float64

    function Gamma(α::Real, β::Real)
        α > zero(α) && β > zero(β) || error("Gamma: both shape and scale must be positive")
        new(float64(α), float64(β))
    end

    Gamma(α::Real) = Gamma(α, 1.0)
    Gamma() = new(1.0, 1.0)
end

@_jl_dist_2p Gamma gamma

@distr_support Gamma 0.0 Inf

#### Parameters

shape(d::Gamma) = d.α
scale(d::Gamma) = d.β
rate(d::Gamma) = 1.0 / d.β

params(d::Gamma) = (d.α, d.β)

#### Statistics

mean(d::Gamma) = d.α * d.β

var(d::Gamma) = d.α * d.β^2

skewness(d::Gamma) = 2.0 / sqrt(d.α)

kurtosis(d::Gamma) = 6.0 / d.α

function mode(d::Gamma)
    (α, β) = params(d)
    α >= 1.0 ? β * (α - 1.0) : error("Gamma has no mode when shape < 1.0")
end

function entropy(d::Gamma)
    (α, β) = params(d)
    α + lgamma(α) + (1.0 - α) * digamma(α) + log(β)
end

mgf(d::Gamma, t::Real) = (1.0 - t * d.β)^(-d.α)

cf(d::Gamma, t::Real) = (1.0 - im * t * d.β)^(-d.α)

#### Evaluation

gradlogpdf(d::Gamma, x::Float64) =
    insupport(Gamma, x) ? (d.α - 1.0) / x - 1.0 / d.β : 0.0

#### Fit model

immutable GammaStats <: SufficientStats
    sx::Float64      # (weighted) sum of x
    slogx::Float64   # (weighted) sum of log(x)
    tw::Float64      # total sample weight

    GammaStats(sx::Real, slogx::Real, tw::Real) = new(float64(sx), float64(slogx), float64(tw))
end

function suffstats(::Type{Gamma}, x::Array)
    sx = 0.
    slogx = 0.
    for xi = x
        sx += xi
        slogx += log(xi)
    end
    GammaStats(sx, slogx, length(x))
end

function suffstats(::Type{Gamma}, x::Array, w::Array{Float64})
    n = length(x)
    if length(w) != n
        throw(ArgumentError("Inconsistent argument dimensions."))
    end

    sx = 0.
    slogx = 0.
    tw = 0.
    for i = 1:n
        @inbounds xi = x[i]
        @inbounds wi = w[i]
        sx += wi * xi
        slogx += wi * log(xi)
        tw += wi
    end
    GammaStats(sx, slogx, tw)
end

function gamma_mle_update(logmx::Float64, mlogx::Float64, a::Float64)
    ia = 1.0 / a
    z = ia + (mlogx - logmx + log(a) - digamma(a)) / (abs2(a) * (ia - trigamma(a)))
    1.0 / z
end

function fit_mle(::Type{Gamma}, ss::GammaStats;
    alpha0::Float64=NaN, maxiter::Int=1000, tol::Float64=1.0e-16)

    mx = ss.sx / ss.tw
    logmx = log(mx)
    mlogx = ss.slogx / ss.tw

    a::Float64 = isnan(alpha0) ? 0.5 / (logmx - mlogx) : alpha0
    converged = false

    t = 0
    while !converged && t < maxiter
        t += 1
        a_old = a
        a = gamma_mle_update(logmx, mlogx, a)
        converged = abs(a - a_old) <= tol
    end

    Gamma(a, mx / a)
end


# Aiyagari model
# Endogenous Grid Points with IID Income
# Greg Kaplan 2017
# Translated by Tom Sweeney Jan 2021

using Random, Interpolations, NLsolve, Plots
include("discrete_normal.jl")
include("lininterp1.jl")

# PARAMETERS

## preferences
risk_aver = 2
beta = 0.95

## production
deprec = 0.10
capshare = 0.4

## income risk: discretized N(mu,sigma^2)
mu_y = 1
sd_y = 0.2
ny = 5

## asset grids
na = 40
amax = 50
borrow_lim = 0
agrid_par = 0.5 # 1 for linear, 0 for L-shaped

## computation
max_iter = 1000
tol_iter = 1.0e-6
Nsim = 50000
Tsim = 500

maxiter_KL = 70
tol_KL = 1.0e-5
step_KL = 0.005
rguess = 1/beta-1-0.001 # a bit lower than inverse of discount rate
KLratioguess = ((rguess + deprec)/capshare)^(1/(capshare-1))

# OPTIONS
Display = 1
MakePlots = 1

## which function to interpolation
InterpCon = 0
InterpEMUC = 1

## tolerance for non-linear solver
TolX=1.0e-6

# UTILITY FUNCTION

if risk_aver==1
    u(c) = log.(c)
else
    u(c) = (c.^(1-risk_aver).-1)./(1-risk_aver)
end

u1(c) = c.^(-risk_aver)
u1inv(u) = u.^(-1 ./risk_aver)

# SET UP GRIDS

## assets
agrid = range(0,1,length=na)
agrid = agrid.^(1 ./ agrid_par)
agrid = borrow_lim .+ (amax.-borrow_lim).*agrid

## income: disretize normal distribution
width = nlsolve(x -> discrete_normal(ny,mu_y,sd_y,x...)[1],[2.0]).zero
temp, ygrid, ydist = discrete_normal(ny,mu_y,sd_y,width...)
ycumdist = cumsum(ydist)

# DRAW RANDOM NUMBERS
Random.seed!(2021)
yrand = rand(Nsim,Tsim)

# SIMULATE LABOR EFFICIENCY REALIZATIONS
if Display>=1
    println("Simulating labor efficiency realizations in advance")
end
yindsim = zeros(Int,Nsim,Tsim)

for it = 1:Tsim

    # income realization: note we vectorize simulations at once because
    # of matlab, in other languages we would loop over individuals
    yindsim[yrand[:,it].<=ycumdist[1],it] .= 1
    for iy = 2:ny
        yindsim[(yrand[:,it].>ycumdist[iy-1]) .& (yrand[:,it].<=ycumdist[iy]),it] .= iy;
    end
end

ysim = ygrid[yindsim]

# ITERATE OVER KL RATIO
KLratio = KLratioguess

iterKL = 0
KLdiff = 1

while iterKL<=maxiter_KL && abs(KLdiff)>tol_KL
    iterKL = iterKL + 1

    r = capshare.*KLratio^(capshare-1) - deprec
    R = 1+r
    wage = (1-capshare).* KLratio^capshare

    ## rescale efficiency units of labor so that output = 1
    yscale = (KLratio^(-capshare))./(ygrid'*ydist)

    ## initialize consumption function in first iteration only
    if iterKL==1
        conguess = zeros(na,ny)
        for iy = 1:ny
            conguess[:,iy] = r.*agrid .+ wage.* yscale.*ygrid[iy]
        end
        global con = copy(conguess)
     end

    ## solve for policy functions with EGP

    iter = 0
    cdiff = 1000
    while iter <= max_iter && cdiff>tol_iter
        iter = iter + 1
        global sav = zeros(na,ny)

        conlast = copy(con)

        emuc = u1(conlast)*ydist
        muc1 = beta.*R.*emuc
        con1 = u1inv(muc1)

        ## loop over income
        ass1 = zeros(na,ny)
        for iy = 1:ny

            ass1[:,iy] = (con1 .+ agrid .-wage.* yscale.*ygrid[iy])./R

            ## loop over current period ssets
            for ia  = 1:na
                if agrid[ia]<ass1[1,iy] # borrowing constraint binds
                    sav[ia,iy] = borrow_lim
                else # borrowing constraint does not bind;
                    sav[ia,iy] = lininterp1(ass1[:,iy],agrid,agrid[ia])
                end
            end
            con[:,iy] = R.*agrid .+ wage.*yscale.*ygrid[iy] - sav[:,iy]
        end

        cdiff = maximum(abs.(con - conlast))
        if Display>=2
            println("Iteration no. " * string(iter), " max con fn diff is " * string(cdiff))
        end
    end

    ## simulate: start at assets from last interation
    if iterKL==1
        global asim = zeros(Nsim,Tsim)
    elseif iterKL>1
        # asim[:,1] = Ea.*ones(Nsim,1)
        asim[:,1] = asim[:,Tsim]
    end

    ## create interpolating function
    savinterp = Array{Any}(undef,ny)
    for iy = 1:ny
        savinterp[iy] = interpolate((agrid,), sav[:,iy], Gridded(Linear()))
    end

    ## loop over time periods
    for it = 1:Tsim
        if Display>=2 && mod(it,100)==0
            println("Simulating, time period " * string(it))
        end

        ## asset choice
        if it<Tsim
            for iy = 1:ny
                asim[yindsim[:,it].==iy,it+1] = savinterp[iy](asim[yindsim[:,it].==iy,it])
            end
        end
    end

    ## assign actual labor income values
    labincsim = wage.*yscale.*ysim

    ## mean assets and efficiency units
    global Ea = mean(asim[:,Tsim])
    L = yscale.*mean(ysim[:,Tsim])

    KLrationew = Ea./ L

    KLdiff = KLrationew./KLratio - 1
    if Display>=1
        println("Equm iter " * string(iterKL) * ", r = " * string(r) * ", KL ratio: " * string(KLrationew) * " KL diff: " * string(KLdiff*100) * "%")
    end

    KLratio = (1-step_KL)*KLratio + step_KL*KLrationew
end

# MAKE PLOTS
if MakePlots==1

    ## consumption policy function
    p1 = plot(agrid, [con[:,1] con[:,ny]], xlims=(0,amax), title="Consumption Policy Function", color=[:blue :red], label=["Lowest income state" "Highest income state"])
    display(p1)

    ## savings policy function
    p2 = plot(agrid, [sav[:,1].-agrid[:,1] sav[:,ny].-agrid[:,1]], xlims=(0,amax), title="Savings Policy Function", color=[:blue :red], legend=false)
    plot!(agrid, zeros(na,1), color=:black, lw=0.5)
    display(p2)

    ## nice zoom
    xlimits = (0,1)
    xlimind = trues(na)
    if minimum(agrid) < xlimits[1]
        xlimind = xlimind .& (agrid.>=maximum(agrid[agrid<xlimits[1]]))
    elseif minimum(agrid) > xlimits[2]
        xlimind .= 0
    end
    if maximum(agrid) > xlimits[2]
        xlimind = xlimind .& (agrid.<=minimum(agrid[agrid.>xlimits[2]]))
    elseif maximum(agrid) < xlimits[1]
        xlimind .= 0
    end

    ## consumption policy function: zoomed in
    p3 = plot(agrid[xlimind], [con[xlimind,1] con[xlimind,ny]], xlims=xlimits, title="Consumption: Zoomed", marker=:circle, color=[:blue :red], linewidth=2, legend=false)
    plot!(show=true)
    display(p3)

    ## savings policy function: zoomed in
    p4 = plot(agrid[xlimind], [sav[xlimind,1].-agrid[xlimind] sav[xlimind,ny].-agrid[xlimind]], xlims=xlimits, title="Savings: Zoomed (a'-a)", marker=:circle, color=[:blue :red], linewidth=2, legend=false)
    plot!(agrid, zeros(na,1), color=:black, lw=0.5)
    display(p4)

    ## income distribution
    p5 = histogram(ysim[:,Tsim], bins=[2*ygrid[1]-ygrid[2];ygrid].+(ygrid[2]-ygrid[1])/2, title="Income distribution", color=RGB(0,0.5,0.5), linecolor=:blue, legend=false)
    display(p5)

    ## asset distribution
    p6 = histogram(asim[:,Tsim], nbins=100, title="Asset distribution", color=RGB(.7,.7,.7), linecolor=:black, legend=false)
    display(p6)

    ## convergence check
    p7 = plot(1:Tsim, mean(asim,dims=1)', ylims=(0,2*Ea), title="Mean Asset Convergence", xlabel="Time Period", color=:black, lw=1.5, legend=false)
    display(p7)

    ## asset distribution statistics
    aysim = asim[:,Tsim]./mean(ysim[:,Tsim])
    println("Mean assets: " * string(mean(aysim)))
    println("Fraction borrowing constrained: " * string(sum(aysim.==borrow_lim)/Nsim * 100) * '%')
    println("10th Percentile: " * string(quantile(aysim,.1)))
    println("50th Percentile: " * string(quantile(aysim,.5)))
    println("90th Percentile: " * string(quantile(aysim,.9)))
    println("99th Percentile: " * string(quantile(aysim,.99)))
end


include("../Resistances/heat.jl")

using Plots, Measurements

# Boundary model(s)
T_a = 95.6
T_b = 273.15

@named boundaryT_a =Temperature(T=T_a)
@named boundaryT_b =Temperature(T=T_b)

@variables Q_total(t)
D = Differential(t)

# Sphere model
r_inner = 0.1
th_1 = 0.0025
th_ins = 0.01
th_2 = 0.0025

r_ins_inner = r_inner + th_1
r_ins_outer = r_ins_inner + th_ins
r_outer = r_ins_outer + th_2

@named convectionInner = Convection(surfaceArea=4*3.14159*r_inner^2,alpha=150)
@named linerInner = Sphere(r_inner=r_inner,r_outer=r_ins_inner,lambda=15)
@named contact_1 = Contact(surfaceArea=4*3.14159*r_ins_inner^2,Rc_pp=0.003)
@named insulation = Sphere(r_inner=r_ins_inner,r_outer=r_ins_outer,lambda=0.033)
@named contact_2 = Contact(surfaceArea=4*3.14159*r_ins_outer^2,Rc_pp=0.003)
@named linerOuter = Sphere(r_inner=r_ins_outer,r_outer=r_outer,lambda=15)
@named radiationOuter = Radiation(surfaceArea=4*3.14159*r_outer^2,epsilon=0.7)
@named convectionOuter = Convection(surfaceArea=4*3.14159*r_outer^2,alpha=6.0)

eqs =
[
    connect(convectionInner.port_a,boundaryT_a.port)
    connect(convectionInner.port_b,linerInner.port_a)
    connect(linerInner.port_b,contact_1.port_a)
    connect(contact_1.port_b,insulation.port_a)
    connect(insulation.port_b,contact_2.port_a)
    connect(contact_2.port_b,linerOuter.port_a)
    connect(linerOuter.port_b,radiationOuter.port_a,convectionOuter.port_a)
    connect(convectionOuter.port_b, radiationOuter.port_b,boundaryT_b.port)
    D(Q_total) ~ convectionInner.port_a.Q_flow
]

@named model = ODESystem(eqs, t, [Q_total],[], systems=[boundaryT_a,boundaryT_b,convectionInner,linerInner,contact_1,insulation,contact_2,linerOuter,radiationOuter,convectionOuter])

sys = structural_simplify(model)
u0 = [Q_total=>0.0]
prob = ODAEProblem(sys, u0, (0, 1.0), [T_a, T_b, r_inner, th_1, r_ins_inner, th_ins, r_ins_outer, th_2, r_outer])
sol = solve(prob, Tsit5())
y_min, t_min, y_max, t_max = reportMinMax("Q_total","Thermal Resistances Model",sol)
errorCheck(-69.43455, y_min)
plot(sol)


export newton, linop

struct DualFun{F,T}
    f::F
    J::T
end
DualFun(f::Fun) = DualFun(f,Operator(I,space(f)))

domain(df::DualFun) = domain(df.f)

differentiate(d::DualFun) = DualFun(d.f',Derivative(rangespace(d.J))*d.J)
integrate(d::DualFun) = DualFun(integrate(d.f),Integral(rangespace(d.J))*d.J)
function cumsum(d::DualFun)
    Q=Integral(rangespace(d.J))*d.J
    DualFun(cumsum(d.f),(I-Evaluation(rangespace(Q),leftendpoint))*Q)
end

adjoint(d::DualFun) = differentiate(d)

^(d::DualFun,k::Integer)=DualFun(d.f^k,k*d.f^(k-1)*d.J)

# from DualNumbers
for (funsym, exp) in Calculus.symbolic_derivatives_1arg()
    @eval function $(funsym)(z::DualFun)
        x = z.f
        xp = z.J
        DualFun($(funsym)(x),$exp*xp)
    end
end

for OP in (:+,:-)
    @eval begin
        $OP(a::DualFun,b::Union{Number,Fun}) = DualFun($OP(a.f,b),a.J)
        $OP(a::Union{Number,Fun},b::DualFun) = DualFun($OP(a,b.f),$OP(b.J))
        $OP(a::DualFun,b::DualFun) = DualFun($OP(a.f,b.f),$OP(a.J,b.J))
    end
end
-(a::DualFun)=DualFun(-a.f,-a.J)

*(a::Union{Number,Fun},b::DualFun)=DualFun(a*b.f,a*b.J)
*(a::DualFun,b::Union{Number,Fun})=DualFun(b*a.f,b*a.J)
*(a::DualFun,b::DualFun)=DualFun(a.f*b.f,a.f*b.J+b.f*a.J)

/(a::Union{Number,Fun},b::DualFun)=DualFun(a/b.f,-a/b.f^2*b.J)
/(a::DualFun,b::Union{Number,Fun})=DualFun(a.f/b,a.J/b)
/(a::DualFun,b::DualFun)=DualFun(a.f/b.f,a.J/b.f-a.f/b.f^2*b.J)

(d::DualFun)(x) = DualFun(d.f(x),Evaluation(rangespace(d.J),x)*d.J)
first(d::DualFun) = DualFun(first(d.f),Evaluation(rangespace(d.J),leftendpoint)*d.J)
last(d::DualFun) = DualFun(last(d.f),Evaluation(rangespace(d.J),rightendpoint)*d.J)

jacobian(d::DualFun)=d.J

promote_rule(::Type{DF},::Type{T}) where {DF<:DualFun,T<:Number}=DualFun
convert(::Type{DualFun},b::Number) = DualFun(b,0)

function linop(f::Function,ds::Space)
    if (isgeneric(f) && applicable(f,0)) || (!isgeneric(f)&&arglength(f)==1)
        df=f(DualFun(zeros(ds)))
    elseif (isgeneric(f) && applicable(f,0,0)) || (!isgeneric(f)&&arglength(f)==2)
        df=f(Fun(ds),DualFun(zeros(ds)))
    else
        error("Not implemented")
    end

    if isa(df,Array)
        map(u->u.J,df)
    else
        df.J
    end
end

linop(f::Function,d) = linop(f,Space(d))
linop(f::Function) = linop(f,Chebyshev())  #TODO: UnsetSpace

# full operator should be
# N=u->[B*u-bcs;...]
function newton(N,u0::Fun;maxiterations=15,tolerance=1E-15)
    u=u0
    err=Inf
    for k=1:maxiterations
        DF=N(DualFun(u))
        J=map(jacobian,DF)
        F=map(d->d.f,DF)
        unew=u-J\F
        err=norm(unew-u)
        if err≤10tolerance
            return unew
        else
            u=chop(unew,tolerance)
        end
    end
    @warn "Maximum number of iterations $maxiterations reached, with approximate accuracy of $err."
    return u
end

newton(N,d::Domain;opts...) =
    newton(N,zeros(d);opts...)

newton(N,d;opts...) =
    newton(N,Domain(d);opts...)


using Images, ImageFiltering
#=using Plots
const w = 5
function clean_rows(x::Matrix{Float64})
    n = size(x, 1)
    map(1:n - w) do i
        j = i + w
        x[i, j:end]
    end
end
const kernel = Kernel.gaussian((0,1))
function find_peaks(μ::Vector{Float64})
    img = centered(μ')
    s = vec(parent(imfilter(img, kernel)))
    findlocalmaxima(s, 1, false)
end
function get_parameters(row::Vector{Float64})
    m, M = extrema(row)
    m == M && return (0.0, 0.0, 0)
    Δ = M - m
    p = (row[end] - m)/Δ
    return (Δ, p, findfirst(row .> m + Δ/2))
end
function detect_nodes(file::String)
    a = readcsv(file)
    x = Matrix{Float64}(a[2:end, 2:end])
    lengths = Float64.(a[2:end, 1])
    times = Float64.(a[1, 2:end])
    rows = clean_rows(x)
    μ = mean.(rows)
    j = find_peaks(μ)
    map(j) do i
        Δ, p, ti = get_parameters(rows[i])
        ti += i[1] + w
        t = times[ti]
        l = lengths[i]
        (l, t, Δ, p)
    end
end

file = "/home/yakir/.julia/datadeps/all/9/pos1,3-dr5_pos2,4-cle44-luc-1-5/stage 2/root 1/intensities.csv"
p = detect_nodes(file)=#

#=struct Node{T <: Function}
    length_i::Int
    time_i::Int
    label::T
end=#

struct Node
    length_i::Int
    time_i::Int
    Δ::Float64
    p::Float64
end

const kernel = Kernel.gaussian((0,0.8))
function find_peaks(μ::Vector{Float64})
    img = centered(μ')
    s = vec(parent(imfilter(img, kernel)))
    findlocalmaxima(s, 1, false)
end
function get_parameters(row::Vector{Float64})
    m, M = extrema(row)
    m == M && return (0.0, 0.0, 0)
    Δ = M - m
    p = (row[end] - m)/Δ
    return (M, p, findfirst(row .> m + Δ/2))
end
const w = 15
function detect_nodes(track::Track)
    # track = tracks[2][1]
    n = length(track.intensities)
    m = minimum(minimum(i) for i in track.intensities)
    intensities = fill(m, n, n)
    for i in 1:n, j in 1:i-w
        intensities[j,i] = track.intensities[i][j]
    end
    y = vec(mean(intensities, 2))
    j = find_peaks(y)
    nodes = map(j) do i
        Δ, p, ti = get_parameters(intensities[i, :])
        # Node(i[1], ti, x -> text("($Δ, $p)", 5, RGB(1,1,1)))
        # Node(i[1], ti, x -> text("$Δ-$p", RGBA(0,1,0, Int(x > ti)), 8))
        Node(i[1], ti, Δ, p)
    end
    filter!(node -> node.Δ > 0.2, nodes)
    #=rows = [track.intensities[i][1:end-w] for i in w + 1:n]
    y = [mean(rows[i][min(j, i)] for i in 1:n-w) for j in 1:n-w]

    # μ = mean.(rows)
    j = find_peaks(y)

    map(j) do i
        Δ, p, ti = get_parameters(rows[i])
        li = i[1] + w - 1
        ti += li
        Node(li, ti, x -> text("($Δ, $p)", 5, RGB(1,1,1)))
        # Node(li, ti, x -> text("($Δ, $p)", RGBA(1,1,1, Int(x > li)), 8))
    end=#
end


using StaticArrays

export delta_parallelepiped_sampling!

"""
Calcultate val = int_{parallelepiped} d^3k delta(`etarget` - e(k)) / volume
where `e(k) = e0 + k ⋅ v0`. The parallelepiped region is [-L/2, L/2]^3.
Divide the parallelepiped into six tetrahedra and use tetrahedron integration.
# Sampling
Sample points on the `e(k) = etarget` plane, where the spacing is relative to the edge
lengths. Sample around `nsample_1d` points along each dimension, and at most
`(nsample_1d+1)^2` points are sampled. Sampled points are stored in the array `ksamples`.
"""
@inline function delta_parallelepiped_sampling!(etarget::T, e0, v0, L, nsample_1d, ksamples) where {T}
    if norm(v0) < 1E-10
        # Velocity is too small. Tetrahedron method does not work.
        return zero(T), 0
    end

    L_div_2 = L ./ 2
    e1 = e0 + (-v0[1] * L_div_2[1] - v0[2] * L_div_2[2] - v0[3] * L_div_2[3])
    e2 = e0 + (+v0[1] * L_div_2[1] - v0[2] * L_div_2[2] - v0[3] * L_div_2[3])
    e3 = e0 + (-v0[1] * L_div_2[1] + v0[2] * L_div_2[2] - v0[3] * L_div_2[3])
    e4 = e0 + (+v0[1] * L_div_2[1] + v0[2] * L_div_2[2] - v0[3] * L_div_2[3])
    e5 = e0 + (-v0[1] * L_div_2[1] - v0[2] * L_div_2[2] + v0[3] * L_div_2[3])
    e6 = e0 + (+v0[1] * L_div_2[1] - v0[2] * L_div_2[2] + v0[3] * L_div_2[3])
    e7 = e0 + (-v0[1] * L_div_2[1] + v0[2] * L_div_2[2] + v0[3] * L_div_2[3])
    e8 = e0 + (+v0[1] * L_div_2[1] + v0[2] * L_div_2[2] + v0[3] * L_div_2[3])
    val = zero(T)
    val += delta_tetrahedron(etarget, SVector{4,T}(e1, e2, e4, e5))
    val += delta_tetrahedron(etarget, SVector{4,T}(e1, e3, e4, e5))
    val += delta_tetrahedron(etarget, SVector{4,T}(e2, e4, e6, e5))
    val += delta_tetrahedron(etarget, SVector{4,T}(e3, e4, e7, e5))
    val += delta_tetrahedron(etarget, SVector{4,T}(e4, e7, e8, e5))
    val += delta_tetrahedron(etarget, SVector{4,T}(e4, e6, e8, e5))

    # TODO: If |v0| = 0, skip sampling

    # Find one point on the e(k0) = etarget plane
    e = (e1, e2, e3, e4, e5, e6, e7, e8)
    imin = argmin(e)
    imax = argmax(e)
    emin = e[imin]
    emax = e[imax]
    # Trivial case: etarget outside the energy range
    if etarget < emin || emax < etarget
        return zero(T), 0
    end

    kmin, kmax = KLIST[imin], KLIST[imax]
    ratio = (etarget - emin) / (emax - emin)
    k0 = (kmin + (kmax - kmin) * ratio) .* L_div_2

    # Find two orthogonal unit vectors on the plane
    i = argmin(abs.(v0))
    if i == 1
        u0 = Vec3{T}(1, 0, 0)
    elseif i == 2
        u0 = Vec3{T}(0, 1, 0)
    else
        u0 = Vec3{T}(0, 0, 1)
    end
    u1 = cross(v0, u0)
    u2 = cross(v0, u1)

    # Normalize unit vectors according to the step size
    u1 = u1 / norm(u1 ./ L) / nsample_1d
    u2 = u2 / norm(u2 ./ L) / nsample_1d

    # k vectors on the e(k) = etarget plane is k0 + n1 * u1 + n2 * u2.
    # Sample k vectors on the plane and inside the cube.

    nmax = floor(Int, sqrt(3) * nsample_1d)
    nsamples = 0
    for n1 in -nmax:nmax
        was_inside = false
        k = k0 + n1 * u1 - nmax * u2
        for n2 in -nmax:nmax
            k += u2
            if all(abs.(k) .<= L_div_2)
                was_inside = true
                nsamples += 1
                ksamples[:, nsamples] .= k
            elseif was_inside
                # Was inside the cube, but moved out. The remaining points are all outside.
                break
            end
        end
    end

    return val / 6, nsamples
end