# Tests formulations in prob/scopf
@testset "test scopf" begin

scopf_dc_cuts_soft_woc_objective = [14642.16, 26982.17]
@testset "scopf cuts dc soft 2, without cuts - $(i)" for (i,network) in enumerate(networks)
    network = deepcopy(network)
    network["gen_flow_cuts"] = []
    network["branch_flow_cuts"] = []

    result = run_scopf_cuts_soft(network, DCPPowerModel, lp_solver)

    @test isapprox(result["termination_status"], OPTIMAL)
    @test isapprox(result["objective"], scopf_dc_cuts_soft_woc_objective[i]; atol = 1e0)
end

@testset "scopf cuts dc soft 2, infeasible" begin
    network = deepcopy(network_infeasible)
    network["gen_flow_cuts"] = []
    network["branch_flow_cuts"] = []

    result = run_scopf_cuts_soft(network, DCPPowerModel, lp_solver)
    @test isapprox(result["termination_status"], INFEASIBLE)
end

scopf_dc_cuts_soft_wc_objective = [14642.16, 37403.76]
@testset "scopf cuts dc soft 2, with cuts - $(i)" for (i,network) in enumerate(networks)
    network = deepcopy(network)
    network["gen_flow_cuts"] = []
    network["branch_flow_cuts"] = []

    result = run_opf_cheap(network, DCPPowerModel, lp_solver)
    @test isapprox(result["termination_status"], OPTIMAL)

    update_active_power_data!(network, result["solution"])

    cuts = check_contingencies_branch_power(network, total_cut_limit=2, gen_flow_cuts=[], branch_flow_cuts=[])

    #println(length(cuts.gen_cuts) + length(cuts.branch_cuts))
    #cuts_found = sum(length(c.gen_cuts)+length(c.branch_cuts) for c in cuts)
    append!(network["gen_flow_cuts"], cuts.gen_cuts)
    append!(network["branch_flow_cuts"], cuts.branch_cuts)

    result = run_scopf_cuts_soft(network, DCPPowerModel, lp_solver)

    @test isapprox(result["termination_status"], OPTIMAL)
    @test isapprox(result["objective"], scopf_dc_cuts_soft_wc_objective[i]; atol = 1e0)
end

scopf_ac_cuts_soft_wc_objective = [14676.95, 37915.33]
@testset "scopf cuts ac soft 2, with cuts - $(i)" for (i,network) in enumerate(networks)
    network = deepcopy(network)
    network["gen_flow_cuts"] = []
    network["branch_flow_cuts"] = []

    result = run_opf_cheap(network, ACPPowerModel, nlp_solver)
    @test isapprox(result["termination_status"], LOCALLY_SOLVED)

    update_active_power_data!(network, result["solution"])

    cuts = check_contingencies_branch_power(network, total_cut_limit=2, gen_flow_cuts=[], branch_flow_cuts=[])

    #println(length(cuts.gen_cuts) + length(cuts.branch_cuts))
    #cuts_found = sum(length(c.gen_cuts)+length(c.branch_cuts) for c in cuts)
    append!(network["gen_flow_cuts"], cuts.gen_cuts)
    append!(network["branch_flow_cuts"], cuts.branch_cuts)

    result = run_scopf_cuts_soft(network, ACPPowerModel, nlp_solver)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], scopf_ac_cuts_soft_wc_objective[i]; atol = 1e0)
end

scopf_dc_cuts_soft_woc_objective = [14642.16, 26982.17]
@testset "scopf cuts dc soft 2, without cuts - $(i)" for (i,network) in enumerate(networks)
    network = deepcopy(network)
    network["gen_flow_cuts"] = []
    network["branch_flow_cuts"] = []

    result = run_scopf_cuts_soft_bpv(network, DCPPowerModel, lp_solver)

    @test isapprox(result["termination_status"], OPTIMAL)
    @test isapprox(result["objective"], scopf_dc_cuts_soft_woc_objective[i]; atol = 1e0)
end

@testset "scopf cuts dc soft 2, infeasible" begin
    network = deepcopy(network_infeasible)
    network["gen_flow_cuts"] = []
    network["branch_flow_cuts"] = []

    result = run_scopf_cuts_soft_bpv(network, DCPPowerModel, lp_solver)
    @test isapprox(result["termination_status"], INFEASIBLE)
end

scopf_dc_cuts_soft_wc_objective = [14642.16, 30737.94]
@testset "scopf cuts dc soft 2, with cuts - $(i)" for (i,network) in enumerate(networks)
    network = deepcopy(network)
    network["gen_flow_cuts"] = []
    network["branch_flow_cuts"] = []

    result = run_opf_cheap(network, DCPPowerModel, lp_solver)
    @test isapprox(result["termination_status"], OPTIMAL)

    update_active_power_data!(network, result["solution"])

    cuts = check_contingencies_branch_power_bpv(network, total_cut_limit=2, gen_flow_cuts=[], branch_flow_cuts=[])

    #println(length(cuts.gen_cuts) + length(cuts.branch_cuts))
    #cuts_found = sum(length(c.gen_cuts)+length(c.branch_cuts) for c in cuts)
    append!(network["gen_flow_cuts"], cuts.gen_cuts)
    append!(network["branch_flow_cuts"], cuts.branch_cuts)

    result = run_scopf_cuts_soft_bpv(network, DCPPowerModel, lp_solver)

    @test isapprox(result["termination_status"], OPTIMAL)
    @test isapprox(result["objective"], scopf_dc_cuts_soft_wc_objective[i]; atol = 1e0)
end

scopf_ac_cuts_soft_wc_objective = [14676.95, 28346.36]
@testset "scopf cuts ac soft 2, with cuts - $(i)" for (i,network) in enumerate(networks)
    network = deepcopy(network)
    network["gen_flow_cuts"] = []
    network["branch_flow_cuts"] = []

    result = run_opf_cheap(network, ACPPowerModel, nlp_solver)
    @test isapprox(result["termination_status"], LOCALLY_SOLVED)

    update_active_power_data!(network, result["solution"])

    cuts = check_contingencies_branch_power_bpv(network, total_cut_limit=2, gen_flow_cuts=[], branch_flow_cuts=[])

    #println(length(cuts.gen_cuts) + length(cuts.branch_cuts))
    #cuts_found = sum(length(c.gen_cuts)+length(c.branch_cuts) for c in cuts)
    append!(network["gen_flow_cuts"], cuts.gen_cuts)
    append!(network["branch_flow_cuts"], cuts.branch_cuts)

    result = run_scopf_cuts_soft_bpv(network, ACPPowerModel, nlp_solver)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], scopf_ac_cuts_soft_wc_objective[i]; atol = 1e0)
end

end # close test group


lb_param_ub = [
0.95    0.99   	1; 	    # beta
0.0	    2      	5;		# gam
0    	0.9   	1;	    # rho1
0       0.02 	0.1;    # sigma1
0    	0.7     1;      # rho2
0	    0.01  	0.1;    # sigma2
6/24    8/24	9/24	# nss
]


# HARMONIC APPROXIMATION THEORY

# Collision rate per atom
function equilibrium_collrate(field::HarmonicField, species, Np, T, t)
    ω_x, ω_y, ω_z = field.ωx(t), field.ωy(t), field.ωz(t)
    n̄ = Np * ω_x * ω_y * ω_z * (species.m / (4π * kB * T))^1.5
    v̄ = sqrt(8 * kB * T / (π * species.m))
    Γₐ = sqrt(2) * n̄ * species.σ * v̄
    return Γₐ
end

# Speed distribution
function equilibrium_speeds(m, T, max_speed = 0.01)
    N₀ = 4π * (m / (2π * kB * T))^1.5 # Harmonic normalisation constant
    f(v) = N₀ * v^2 * exp(-0.5 * m * v^2 / (kB * T)) # Ideal distribution
    speed_domain = range(0, stop=1.2 * max_speed, length=300)
    speed_density = f.(speed_domain)
    return speed_domain, speed_density
end

# Harmonic theory
function harmonic_theory(species, ωx, ωy, ωz, T₀, N₀, trap_depth, duration, τ_bg)
    dt = 0.000001

    ωx, ωy, ωz = time_parametrize(ωx, ωy, ωz)
    trap_depth = time_parametrize(trap_depth)

    iter_count = ceil(Int64, duration / dt)
    N_series = zeros(Float64, iter_count)
    T_series = zeros(Float64, iter_count)
    Γ_series = zeros(Float64, iter_count)
    t_series = zeros(Float64, iter_count)

    m = species.m
    a = species.aₛ
    σ = species.σ

    t = 0
    N = N₀
    T = T₀
    i = 1
    while i <= iter_count
        n₀ = N * ωx(t) * ωy(t) * ωz(t) * (m / (2 * kB * T))^3/2
        #Γcoeff = 8*sqrt(2)*a^2 * m / (π * kB) * ωx(t) * ωy(t) * ωz(t)
        #Γ =  Γcoeff * N / T # Peak collision rate (per atom?)
        Γ = n₀ * σ * sqrt(16 * kB * T / (π * m))
        εₜ = trap_depth(t)
        η = εₜ / (kB * T) # Truncation parameter
        Ṅₑ = - N * Γ * η/sqrt(2) * exp(-η) # Evaporation rate
        Ṅₗ = - N / τ_bg # Background loss rate
        Ṅ = Ṅₑ + Ṅₗ # Total removal rate
        ε̄ = 3 * kB * T # Average energy
        γ = (εₜ / ε̄ - 1) * (Ṅₑ / Ṅ) # Evaporation efficiency
        Ṫ = T * γ * (Ṅ / N) # Rate of temperature change

        dN = Ṅ * dt
        dT = Ṫ * dt

        N_series[i] = N
        T_series[i] = T
        Γ_series[i] = Γ / (2 * sqrt(2)) # Convert to AVG collrate
        t_series[i] = t

        N += dN
        T += dT
        t += dt

        i += 1
    end

    return t_series, N_series, T_series, Γ_series
end

function purdue_theory(species, ωx, ωy, ωz, T₀, N₀, trap_depth, duration, τ_bg = Inf, K = 0)
    dt = 0.000001

    ωx, ωy, ωz = time_parametrize(ωx, ωy, ωz)
    trap_depth = time_parametrize(trap_depth)

    iter_count = ceil(Int64, duration / dt)
    N_series = zeros(Float64, iter_count)
    T_series = zeros(Float64, iter_count)
    Γ_series = zeros(Float64, iter_count)
    t_series = zeros(Float64, iter_count)

    m = species.m
    σ = species.σ

    E₀ = 3 * N₀ * kB * T₀

    Γ_bg = 1 / τ_bg

    t = 0
    N = N₀
    E = E₀
    T = T₀

    i = 1
    while i <= iter_count
        n₀ = N * ωx(t) * ωy(t) * ωz(t) * (m / (2π * kB * T))^1.5 # Peak density
        n̄ = n₀ / sqrt(8) # Mean density
        vᵣ = sqrt(16 * kB * T / (π * m)) # Mean relative velocity

        εₜ = trap_depth(t) # Trap depth
        η = εₜ / (kB * T) # Truncation parameter
        Γ_el = n̄ * σ * vᵣ; # Elastic collision rate per atom
        Γ_ev = (η-4)*exp(-η) * Γ_el # Evaporation rate per atom
        Γ_3B = K * n₀^2 / (3 * sqrt(3)) # Three-body collision rate per atom

        Ṅ = - (Γ_ev + Γ_bg + Γ_3B) * N
        Ė_ev = - N * Γ_ev * (η + (η-5)/(η-4)) * kB * T
        Ė_bg = - Γ_bg * E
        Ė_3B = - Γ_3B * (2/3) * E

        Ė = Ė_ev + Ė_bg + Ė_3B

        # Note: ignoring effects of adiabatic change in trap shape

        dN = Ṅ * dt
        dE = Ė * dt

        N_series[i] = N
        T_series[i] = T
        Γ_series[i] = Γ_el
        t_series[i] = t

        N += dN
        E += dE
        T = E / (3 * N * kB)
        t += dt

        i += 1
    end

    return t_series, N_series, T_series, Γ_series
end


using RosDataProcess
using PyPlot

dx = 0.01
# Test simple time series
a = TimeSeries(sin, -2pi:dx:6pi)
b = TimeSeries(cos, -3pi:dx:2pi)
c = TimeSeries(atan, -5pi:dx:4pi)
series = [a, b, c]
interpolated = intersect_interpolate(series)

# Optional plots to illustrate the above examples
#figure(); foreach(plot, series)
#figure(); plot(interpolated)


module TestJacobian
include("preamble.jl")

using DiffEqBase: ODEProblem
using StaticArrays: SVector
using SymEngine: diff, subs, symbols

using Bifurcations.Continuations: residual, residual_jacobian, get_prob_cache
using Bifurcations.Codim2: DiffEqCodim2Problem, NormalizingAS
using Bifurcations.Examples.PredatorPrey

∙(a, b) = sum(a .* b)
sjac(f, x) = [diff(f[i], x[j]) for i in 1:length(f), j in 1:length(x)]

nx = length(PredatorPrey.u0)
np = length(PredatorPrey.p)
f = generalized_f(PredatorPrey)
pe = [symbols("p_$i") for i in 1:np]
xe = [symbols("x_$i") for i in 1:nx]
ve = [symbols("v_$i") for i in 1:nx]
all_symbols = vcat(pe, xe, ve)
fe = f(xe, pe, 0)
fe = Vector(fe)  # just for safety
dfdx = sjac(fe, xe)

ue = vcat(xe, ve, pe[1], pe[3])
He = vcat(fe,
          dfdx * ve,
          ve ∙ ve - 1)
dHdu = sjac(He, ue)

make_prob(;
        p = (0.0, 3, 5, 3),
        u0 = SVector(0.0, 0.0),
        tspan = (0.0, 30.0),
        ode = ODEProblem{!(u0 isa SVector)}(PredatorPrey.f, u0, tspan, p),
        kwargs...) =
    DiffEqCodim2Problem(
        ode,
        (@lens _[1]),
        (@lens _[3]),
        ([-Inf, -Inf], [Inf, Inf]);
        augmented_system = NormalizingAS(),
        kwargs...)

function test_predator_prey(
        xr, vr, p1, p3;
        u0 = SVector(0.0, 0.0),
        prob = make_prob(u0 = u0),
        prob_cache = get_prob_cache(prob))

    pr = collect(Float64, prob.de_prob.p)
    pr[1] = p1
    pr[3] = p3

    ur = vcat(xr, vr, pr[1], pr[3])
    if prob.de_prob.u0 isa SVector
        ur = SVector(ur...)
    else
        utype = typeof(prob.de_prob.u0)
        if !(ur isa utype)
            ur = utype(ur)
        end
    end

    all_values = vcat(pr, xr, vr) :: Vector{Float64}
    evaluated = (e) -> Float64(subs(e, zip(all_symbols, all_values)...))

    H_desired = evaluated.(He)
    J_desired = evaluated.(dHdu)

    H1_actual = residual(ur, prob_cache)
    @test H1_actual ≈ H_desired

    H2_actual, J_actual = residual_jacobian(ur, prob_cache)
    @test H2_actual ≈ H_desired
    @test J_actual ≈ J_desired
end

rng = MersenneTwister(0)
@testset begin
    for args in [
            ([0, 0], [0, 1], 0.6, 5),
            ([0, 0], [-1, 0], 0.384806, 7.79613),
            ],
        u0 in [SVector(0.0, 0.0), [0.0, 0.0]]

        test_predator_prey(args...; u0=u0)
    end
    for _ in 1:100,
        u0 in [SVector(0.0, 0.0), [0.0, 0.0]]

        p1 = randn(rng)
        p3 = randn(rng)
        xr = randn(rng, nx)
        vr = normalize(randn(rng, nx))
        test_predator_prey(xr, vr, p1, p3; u0=u0)
    end
end

end  # module


# BFGS solver for DMD problem
mutable struct DMD_BFGS_Options{T<:AbstractFloat}
	itm::Integer
	tol::T
	ptf::Integer
	# pre-allocate variables
	ar  ::Vector{T}
	nar ::Vector{T}
	gar ::Vector{T}
	ngar::Vector{T}
	par ::Vector{T}
	sar ::Vector{T}
	yar ::Vector{T}
	Har ::Matrix{T}
end

# constructor for BFGS options
function DMD_BFGS_Options(σ, params; itm=100, tol=1e-5, ptf=10)
	T = typeof(real(params.X[1]));
	tol = T(tol);
	k = params.k;

	ar   = zeros(T, 2*k);
	nar  = zeros(T, 2*k);
	gar  = zeros(T, 2*k);
	ngar = zeros(T, 2*k);
	par  = zeros(T, 2*k);
	sar  = zeros(T, 2*k);
	yar  = zeros(T, 2*k);
	Har  = diagm(0 => fill(σ, 2*k));

	return DMD_BFGS_Options(itm, tol, ptf, ar, nar,
		gar, ngar, par, sar, yar, Har)
end

# BFGS solver
function solveDMD_withBFGS(params, opts)
	# load all the variables
	T    = typeof(real(params.X[1]));
	itm  = opts.itm;
	tol  = opts.tol;
	ptf  = opts.ptf;

	ar   = params.ar;
	nar  = opts.nar;
	gar  = opts.gar;
	ngar = opts.ngar;
	par  = opts.par;
	sar  = opts.sar;
	yar  = opts.yar;
	Har  = opts.Har;

	# first gradient step
	nobj = aBFunc(params);
	aBGrad(ngar, params);
	err = norm(ngar, Inf);
	noi = 0;

	obj_his = zeros(T, itm);
	err_his = zeros(T, itm);

	while err ≥ tol
		# calculate the direction
		BLAS.gemv!('N', T(-1.0), Har, ngar, T(0.0), par);
		# decend line search
		η = 1.0;
		copyto!(nar, ar); BLAS.axpy!(η, par, ar);
		obj = aBFunc(params);
		while obj ≥ nobj
			η *= 0.5;
			η < 1e-10 && break;
			copyto!(ar, nar); BLAS.axpy!(η, par, ar);
			obj = aBFunc(params);
		end

		# update gradient
		aBGrad(gar, params);
		# update differences
		sar .= ar - nar; copyto!(nar, ar);
		yar .= gar - ngar; copyto!(ngar, gar);

		# update Hessian approximation
		ρ = T(1.0)/dot(yar, sar);
		μ = sum(abs2, sar);
		ρ = min(T(10.0), ρ*μ)/μ;
		if ρ > 0.0
			BLAS.gemv!('N', T(1.0), Har, yar, T(0.0), par);
			β = dot(yar, par)*ρ^2 + ρ;
			BLAS.ger!(-ρ, sar, par, Har);
			BLAS.ger!(-ρ, par, sar, Har);
			BLAS.ger!( β, sar, sar, Har);
		end

		# update information
		nobj = obj;
		err = norm(par, Inf);
		noi = noi + 1;

		obj_his[noi] = obj;
		err_his[noi] = err;

		# print information
		noi % ptf == 0 && @printf("iter %4d, obj %1.2e, err %1.2e\n",
			noi, obj, err);
		noi ≥ itm && break;
	end

	return obj_his[1:noi], err_his[1:noi]
end


#*************************************************************************#
#
# FeedforwardNetwork Type
#
#*************************************************************************#

# Constructor I
"""
    FeedforwardNetwork(ninputs, activfunc)

Initialize `ParametricModel` representing feedforward neural network.
`ninputs` is the network number of inputs and `activfunc`
is a vector of `AparametricModel` containg activation
functions for each one of the layers.
"""
function FeedforwardNetwork{T<:AparametricModel}(ninputs::Int,
                                                 activfunc::Vector{T})
    nlayers = length(activfunc)
    mdls = Vector{StaticModel}(3*nlayers)
    node_noutputs = ninputs
    k = 1
    for i = 1:nlayers
        node_ninputs = node_noutputs
        node_noutputs = activfunc[i].ninputs
        mdls[k] = Linear(node_ninputs, node_noutputs)
        k += 1
        mdls[k] = Bias(node_noutputs)
        k += 1
        mdls[k] = activfunc[i]
        k += 1
    end

    SeriesModel(mdls)
end

# Constructor II
"""
    FeedforwardNetwork(ninputs, noutputs, nhidden)

Initialize `ParametricModel` representing feedforward neural network.
`ninputs` and `outputs` are the network number of inputs and outputs
and `nhidden` is a vector containg the number of nodes in each hidden
layer.
"""
function FeedforwardNetwork(ninputs::Int,
                            noutputs::Int,
                            nhidden::Vector{Int})
    nlayers = length(nhidden)+1

    activfunc = Vector{ElementwiseOperator}(nlayers)
    for k = 1:nlayers
        activfunc[k] = (k==nlayers) ? Identity(noutputs) : HyperbolicTangent(nhidden[k])
    end

    FeedforwardNetwork(ninputs, activfunc)
end


module Dyestuff2Test

println()
@info "Dyestuff2 example"

using GLMCopula, RDatasets, Test

# Dataframe with columns: Batch (Categorical), Yield (Int32)
dyestuff2 = dataset("lme4", "Dyestuff2")
groups = unique(dyestuff2[!, :Batch])
n, p, m = length(groups), 1, 1
gcs = Vector{GaussianCopulaVCObs{Float64}}(undef, n)
for (i, grp) in enumerate(groups)
    gidx = dyestuff2[!, :Batch] .== grp
    ni = count(gidx)
    y = Float64.(dyestuff2[gidx, :Yield])
    X = ones(ni, 1)
    V = [ones(ni, ni)]
    gcs[i] = GaussianCopulaVCObs(y, X, V)
end
gcm = GaussianCopulaVCModel(gcs);

# initialize β and τ from least square solution
@info "Initial point:"
init_β!(gcm)
@show gcm.β
# update σ2 from β and τ using the MM algorithm
fill!(gcm.Σ, 1)
update_Σ!(gcm)
@show gcm.τ
@show gcm.Σ;
@test loglikelihood!(gcm, true, false) ≈ -81.43652408
@show gcm.∇β
@show gcm.∇τ
@show gcm.∇Σ

# fit model using NLP on profile loglikelihood
@info "MLE:"
@time GLMCopula.fit!(gcm, IpoptSolver(print_level=0))
@show gcm.β
@show gcm.τ
@show gcm.Σ
@test loglikelihood!(gcm, true, false) ≈  -81.43651914
@show gcm.∇β
@show gcm.∇τ
@show gcm.∇Σ

end


#' ## Генерация априорной информации

using NonLinearReactionAdvectionDiffusionWithFrontData;
using NonLinearReactionAdvectionDiffusionWithFrontData: phidetermination, Φ;
using NonLinearReactionAdvectionDiffusionWithFrontData: f1, f2;

u_l(t) = -8 #+ cos(2*π * t);
u_r(t) =  4 #+ (1 + sin(2*π * t));
qf(x) = sin(3 * π * x);         # Коэффициент линейного усиления, который в обратной
                                # задаче необходимо определить, но при генерации априорной
                                # информации мы задаем некоторый коэффициент, который,
                                # собственно, после имея априорную информацию и будем определять.
ε = 0.2;                        # Малый параметр при старшей производной
a, b = 0, 1;                    # Область по X
t₀, T = 0, 0.36;                # Область по T
N, M = 100, 80;                 # Кол-во разбиений по X, T
h = (b-a)/N;                    # шаг по X
τ = (T-t₀)/M;                   # шаг по T
Xₙ = [a  + n*h for n in 0:N];   # Сетка по Х
Tₘ = [t₀ + m*τ for m in 0:M];   # Сетка по Т
qₙ =     qf.(Xₙ);               # Сеточные значения коэффициента линейного усиления
ulₘ=    u_l.(Tₘ);               # Сеточные значения левого  ГУ
urₘ=    u_r.(Tₘ);               # Сеточные значения правого ГУ
u₀ = u_init.(Xₙ);               # Начальные условия
nothing #hide

u, XX, TP   = solve(u₀, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ);                              # Нам вернут решение, сетку, f_1^s
#
ϕl          = phidetermination(qₙ, ulₘ, Xₙ, N, Tₘ, M);                               # Левый вырожденный корень
ϕr          = phidetermination(qₙ, urₘ,Xₙ, N, Tₘ, M, reverseX = true);               # Нужно подать инвертированную сетку
ϕr          = reverse(ϕr, dims=1);                                                   # А после — инвертировать решение по X
ϕ           = Φ(ϕl, ϕr, N, M);                                                       # Серединный корень
f1_data     = f1(ϕ, u, Xₙ, N, M);                                                    # Положение переходного слоя
f2_data     = f2(f1_data, u, Xₙ, N, M);                                              # Значение функции на переходном слое
nothing #hide
#########################################################################################

q₀ = [ 0.0 * x for x in qₙ];
q₀ = q_guess;

S = 500;
β = 0.01;
qs, Js, Qs = minimize(q₀, u₀, ulₘ, urₘ, Xₙ, N, Tₘ, M, ε, f1_data, f2_data, S = S, β = β, w = 0.00001)

frames = [1:20; 21:div(S,50):S];    # Первые двадцать без пропусков
frames = collect(1:div(S, 50):S);   # С пропусками, чтобы всего было 50 кадров
make_minimzation_gif(Js, Qs, qₙ, Xₙ, name = "Minimization_uniform.gif", frames_to_write = frames, β = β)


using LinearAlgebra
BLAS.set_num_threads(1)

using Zygote, Flux, BenchmarkTools, Test, DataFrames, GLM, Statistics, StaticArrays
using Plots

to_smatrix(x) = x
function to_smatrix(x::Matrix)
    return SMatrix{size(x)...}(x)
end

function build_model(num_layers::Int, size::Int)
    layers = Any[]
    for layer_idx in 1:num_layers
        push!(layers, LSTM(size, size))
    end
    return mapleaves(to_smatrix, mapleaves(Flux.data, Chain(layers...)))
end

@info "Reporting (num_layers, seq_len, feature_size, batch_size, time)"
function sweep_batchsizes(batch_sizes, num_layers = 1, seq_len = 4, feature_size = 4)
    timings = Float64[]
    for batch_size in batch_sizes
        model = build_model(num_layers, feature_size)
        x = SMatrix{feature_size,batch_size}(randn(feature_size, batch_size))
        test(m, x) = gradient(m, x) do m, x
            x_t = x
            for idx in 1:seq_len
                x_t = m(x_t)
            end
            return sum(x_t)
        end
        push!(timings, minimum(@benchmark $test($model, $x)).time)
        @info num_layers, seq_len, feature_size, batch_size, timings[end]
    end
    return timings
end

batch_sizes = 1:8
layer_sizes = 1:3
timings = Dict()
for num_layers in layer_sizes,
    seq_len in (4,),
    feature_size in (4,)

    timings[(num_layers, seq_len, feature_size, batch_sizes)] = sweep_batchsizes(batch_sizes, num_layers, seq_len, feature_size)
end

# These gotten via count_ops/count_lstm_ops.jl
num_ops = Dict(
    1 => 255,
    2 => 491,
    3 => 727,
    4 => 963,
)

overhead_estimates = Dict()
workload_estimates = Dict()
for k in keys(timings)
    num_layers, seq_len, feature_size, batch_sizes = k
    data = DataFrame(X=batch_sizes, Y=timings[k])
    ols = lm(@formula(Y ~ X), data)
    overhead_estimates[k] = coeftable(ols).cols[1][1]/num_ops[num_layers]
    workload_estimates[k] = coeftable(ols).cols[1][2]
end

gr()
ENV["GKSwstype"]="100"
p = plot()
for key in sort(collect(keys(timings)))
    num_layers, seq_len, feature_size, batch_sizes = key
    plot!(p, batch_sizes, timings[key]; title="LSTM Runtime", ylabel="Absolute Runtime (ns)", xlabel="Batch Size", label="$(num_layers) layers")
    extrapolate(x) = overhead_estimates[key]*num_ops[num_layers] + workload_estimates[key]*x
    x_points = [0, batch_sizes...]
    plot!(p, x_points, extrapolate.(x_points); style=:dashdot, label="$(num_layers) layers extrapolation")
end
savefig(p, joinpath(@__DIR__, "lstm_runtime.pdf"))


#=
The tests here merely check that the interface functions work.
Testing of the implementation is done in selinv.jl.
=#

using Test
using Random
using SparseArrays
using LinearAlgebra
using SparseFactorizations

@testset "selinv_api" begin
    @testset "selinv" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64, 10,10,0.1)
            A = I + A + adj(A)
            F = simplicial_ldlt(A; conj=conj)
            B = selinv(F)
            Bref = inv(Matrix(A)) .* ldlt_structure(A)
            @test B ≈ Bref
        end
    end
end


@testset "$(rpad("QuadratureRule",80))" begin

    o = 2
    b = [0.5, 0.5]
    c = [0.0, 1.0]

    quad = QuadratureRule(o,c,b)

    @test order(quad) == o
    @test nodes(quad) == c
    @test weights(quad) == b
    @test nnodes(quad) == length(b)
    @test eachindex(quad) == 1:2
    @test eltype(quad) == Float64

    @test quad(x -> 1)   == 1
    @test quad(x -> x)   == 1//2
    @test quad(x -> x^2) == 1//2

    quad2 = QuadratureRule{Float64,2}(o, c, b)
    quadϵ = QuadratureRule{Float64,2}(o, c .+ 8eps() * rand(2), b .+ 8eps() * rand(2))

    @test hash(quad) == hash(quad2)
    @test hash(quad) != hash(quadϵ)

    @test  isequal(quad, quad2)
    @test !isequal(quad, quadϵ)

    @test isapprox(quad, quad2)
    @test isapprox(quad, quadϵ)

    @test quad ≈   quad2
    @test quad ==  quad2
    @test quad === quad2

    @test quad ≈   quadϵ
    @test quad !=  quadϵ
    @test quad !== quadϵ

end


###############################################################################
"""
Provides utilities for removal of radial distortion, by predicting the true u,v
coordinates of a pixel from the distorted coordinates, and the intrinsic
calibration parameters of the camera.

Also provides utilities to simulate radial distoriton.
"""
module RemoveImageDistortion
###############################################################################

###############################################################################
import ExtrinsicCameraCalibration: IntrinsicCalibrationParameters
# using Statistics
###############################################################################

###############################################################################
export RemoveRadialDistortion, SimulateRadialDistortion
###############################################################################

###############################################################################
function RemoveRadialDistortion(u_d::Array, v_d::Array, ι::IntrinsicCalibrationParameters)
    return RemoveRadialDistortion(u_d::Array, v_d::Array, ι.radial_distortion_poly_degrees,
                                  ι.radial_distortion_poly_coeffs, ι.u₀, ι.v₀)
end
function RemoveRadialDistortion(u_d::Array, v_d::Array, Degrees, Coeffs, u₀, v₀)
    # Convert the *observed* pixel to polar coordinates:
    ρ_d = abs.((u_d .- u₀) .+ (v_d .- v₀)im)
    Θ_d = angle.((u_d .- u₀) .+ (v_d .- v₀)im)
    # Calculate the Δρ correction needed at each point:
    Δρ = Polyvalc(ρ_d, Coeffs, Degrees)
    Δu = cos.(Θ_d) .* Δρ
    Δv = sin.(Θ_d) .* Δρ
    # Apply the correction to ρ and convert back to pixel coordinates:
    u_p = u_d - Δu
    v_p = v_d - Δv
    return u_p,v_p
end
###############################################################################

###############################################################################
function SimulateRadialDistortion(uv_p::Array, ι::IntrinsicCalibrationParameters)
    return SimulateRadialDistortion(u_p, v_p, ι)
end
function SimulateRadialDistortion(u_p::Array, v_p::Array, ι::IntrinsicCalibrationParameters)
    return SimulateRadialDistortion(u_p, v_p,
                                    ι.radial_distortion_poly_degrees,
                                    ι.radial_distortion_poly_coeffs,
                                    ι.u₀, ι.v₀)
end
function SimulateRadialDistortion(u_p::Array, v_p::Array, Degrees, Coeffs, u₀, v₀)
    ρ_p = abs.((u_p .- u₀) .+ (v_p .- v₀)im)
    Θ_p = angle.((u_p .- u₀) .+ (v_p .- v₀)im)
    # Calculate the Δρ correction needed at each point:
    Δρ = Polyvalc(ρ_p, Coeffs, Degrees)
    Δu = cos.(Θ_p) .* Δρ
    Δv = sin.(Θ_p) .* Δρ
    # Apply the correction to ρ and convert back to pixel coordinates:
    u_d = u_p + Δu
    v_d = v_p + Δv
    return u_d,v_d
end
###############################################################################

###############################################################################
# Polyfitc is similar to polyfit, but allows choosing powers (5,3,1, etc.)
# (based on Matlab file exchange Polyfitc calculations):
function Polyfitc(X,Y,Degrees)
    lsqM = ones(length(X), length(Degrees))
    for n = 1:length(Degrees)
        lsqM[:, n] = X.^Degrees[n]
    end
    Coeffs = (lsqM \ Y)'
    return Coeffs[:]
end
# ----------
function Polyvalc(X, Coeffs, Degrees)
    val(x,Coeffs,Degrees) = sum( [ Coeffs[i]*x^(Degrees[i]) for i in 1:length(Coeffs) ] )
    Y = [ val(x,Coeffs,Degrees) for x in X ]
end
###############################################################################

###############################################################################
###############################################################################
###############################################################################
end                             # end of module
###############################################################################
###############################################################################
###############################################################################


using Pkg; Pkg.activate("GraphFusedElasticnet.jl/")
include("GraphFusedElasticNet.jl/src/ARS.jl")
using Plots

μ, σ = 1.0, 2.0
f(x) = exp(-0.5(x - μ)^2 / σ^2) / sqrt(2pi * σ^2)
support = (-10.0, 2.0)

# Build the sampler and simulate 10,000 samples
sampler = RejectionSampler(f, support, (-1.0, 1.0), max_segments = 3)
@time sim = run_sampler!(sampler, 1000);

x = range(-10.0, 10.0, length=100)
envelop = [eval_envelop(sampler.envelop, xi) for xi in x]
target = [f(xi) for xi in x]

histogram(sim, normalize = true, label = "Histogram")
plot!(x, [target envelop], width = 2, label = ["Normal(μ, σ)" "Envelop"])


## Config (changes allowed)
latlen = 220          # lattice length
iters = 220           # iterations
start_rule = 113      # starting rule
order2_rule = ">|+"   # the rule of changing the rule
           #      - no changes (regular auomata)
           # >    - shift one to the right
           # >>   - shift two to the right
           # <    - shift one to the left
           # <<   - shift two to the left
           # +    - next rule (increases Wolfram's number)
           # -    - previous rule (decreases Wolfram's number)
           # |    - flip the rule by bits (mirror)
           # !    - perform "not" operation on each bit
           # +>>  - combination (increase rule and move two rules to the right). + operation cometh first
init = "r"        # initial distribution of cells
           # .  - dot
           # r  - random
           # i  - inverse
           # :  - alternating black and white cells
           # .i - combination (inverse dot)
           # a  - get from init_arr
init_arr = []     # initial distribution of cells (leave blank if init is not "a").
                  # watch out, init_arr needs to be of length latlen

## Some interesting configurations
#=
starting with dot:
# 60 >>
# 60 >>|<<
# 60 >|<
# 60 <|>
# 60 >+<
# 101
# 101 >
# 101 >>
# 101 >|
# 101 |
# 105 >>
# 105 >|>
# 105 >|<
random start:
# 40 +
# 40 >+< any even?
# 50 >|<
# 60 >|<
# 105 >>
# 105 >>+
# 110 +<
# 110 <
# 110 +
# 110 -
# 110 --
# 110 >+++<
=#

## Config processing (changes not recommended)
# rule is a vector of integers (1s and 0s)
rule = [Int(i)-48 for i::Char in Base.bin(Unsigned(start_rule), 8, false)]

# a function to change the rule number by one
function next_rule!(rulevec::Vector{Int}, step::Int)::Nothing
    num::Int = parse(Int, join(rulevec), base=2)
    num += step
    if num > 255
        num -= 255
    elseif num < 0
        num = 256 + num
    end
    b::String = Base.bin(Unsigned(num), 8, false)
    for i in 1:length(rulevec)
        rulevec[i] = Int(b[i])-48
    end
    nothing
end

# a function to change the rule with a bitwise "not"
function not!(rulevec::Vector{Int})::Nothing
    for i in 1:length(rulevec)
        rulevec[i] = Int(rulevec[i] == 0)
    end
    nothing
end

# a mapping from characters to actions (order2 rules)
char2rule = Dict([
    '+' => r -> next_rule!(r, 1),
    '-' => r -> next_rule!(r, -1),
    '>' => r -> circshift!(r, copy(r), 1),
    '<' => r -> circshift!(r, copy(r), -1),
    '|' => r -> reverse!(r),
    '!' => r -> not!(r)
])

## Import and setup for visualization
using PyCall
import PyPlot; const plt = PyPlot
rcParams = PyDict(plt.matplotlib."rcParams")
rcParams["image.cmap"] = "binary"

## Data init (first line)
data = [] # comment this line to append the new auomaton graph to the previous one
if occursin("r", init)
    push!(data, [rand((0, 1)) for i::Int in 1:latlen])
elseif occursin(".", init)
    push!(data, [((i==round(latlen/2)) ? 1 : 0) for i::Int in 1:latlen])
elseif occursin("a", init)
    push!(data, init_arr)
elseif occursin(":", init)
    push!(data, [i%2 for i::Int in 1:latlen])
end
if occursin("i", init)
    data = [[abs(i-1) for i::Int in data[1]]]
end

## The magic
past = length(data)-1
for iter::Int in 1:iters
    println(parse(Int, join(rule), base=2), rule)
    new_line::Vector{Int} = []
    for i::Int in 1:length(data[past+iter])
        neigh::String = ""
        for j::Int in i-1:i+1
            aj::Unsigned = j # actual index to look at
            if (j <= 0)
                aj = length(data[past+iter])-1+j
            elseif (j > length(data[past+iter]))
                aj = j - length(data[past+iter])
            end
            neigh *= string(data[past+iter][aj])
        end
        push!(new_line, rule[8-parse(Int, neigh, base=2)])
    end
    push!(data, new_line)
    # change the rule according to the order2_rule
    for c::Char in order2_rule
        char2rule[c](rule)
    end
end

## Visualization
fig, ax = plt.subplots(figsize=(80, 80))
ax.matshow(data)
ax.axis(false)
display(plt.gcf()) # for Juno Plots window
plt.show()         # for terminal usage


using PyPlot
using JLD
function plot_sens()
	X = load("../data/obj_erg_avg/cos4y_s2.jld")
	s2_arr = X["s2"]
	J_arr = X["J"]
	fig, ax = subplots(1,1)
	ax.plot(s2_arr, J_arr, ".", ms=10.0)
    ax.xaxis.set_tick_params(labelsize=32)
    ax.yaxis.set_tick_params(labelsize=32)
    ax.set_xlabel(L"$s_2$",fontsize=32)
    ax.set_ylabel(L"$\langle J\rangle$",fontsize=32)
	ax.grid(true)

	X = load("../data/sens/dJds_s2.jld")

	dJds = X["dJds"]
	s2 = X["s2"]

	X = load("../data/obj_erg_avg/cos4y_s2_sens.jld")
	J = X["J"]

	eps = 2.e-2
	n = size(dJds)[1]
	J_pts = reshape([J .- eps*dJds J .+ eps*dJds], n,
					2)'
	s_pts = reshape([s2 .- eps s2 .+ eps], n,
					2)'

	ax.plot(s_pts[:,1:n-1], J_pts[:,1:n-1], "k",lw=6.0)
	ax.plot(s_pts[:,n], J_pts[:,n], "k",lw=6.0,label=L"$d_s\langle J\rangle$ from S3")

	ax.legend(fontsize=32)

	return dJds

end


@testset "ReachNumber" begin
  heuristic(state, goal) = abs(goal - state)

  @testset "ReachNumberGoalDefaultHeuristic" begin
    start = 0
    goal = 10

    getneighbours(state) = [state - 1, state + 1]

    res = astar(getneighbours, start, goal)
    @test res.status == :success
    @test res.path == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    @test res.cost == 10
  end

  @testset "ReachNumberGoal" begin
    start = 0
    goal = 10

    getneighbours(state) = [state - 1, state + 1]

    res = astar(getneighbours, start, goal; heuristic)
    @test res.status == :success
    @test res.path == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    @test res.cost == 10
  end

  @testset "ReachNumberGoal Negative" begin
    start = 10
    goal = 0

    getneighbours(state) = [state - 1, state + 1]

    res = astar(getneighbours, start, goal; heuristic)
    @test res.status == :success
    @test res.path == [10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
    @test res.cost == 10
  end

  @testset "Timeout" begin
    start = 1
    goal = 3

    getneighbours(state) = [state + 2, state + 1]

    res = astar(getneighbours, start, goal; heuristic, timeout = -1)# timeout immediately
    @test res.status == :timeout
    @test res.path == [1]
    @test res.cost == 0
  end

  @testset "Multiple paths" begin
    start = 1
    goal = 4

    getneighbours(state) = [state + 2, state + 1, state - 1]

    res = astar(getneighbours, start, goal; heuristic)
    @test res.status == :success
    @test res.path == [1, 3, 4]
    @test res.cost == 2
  end

  @testset "Cost argument" begin
    start = 0
    goal = 10

    getneighbours(state) = [state - 1, state + 1]
    cost(a, b) = a % 2 == 0 ? abs(b - a) : 2 * abs(b - a)

    res = astar(getneighbours, start, goal; heuristic, cost)
    @test res.status == :success
    @test res.path == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    @test res.cost == 15
  end

  @testset "Multiple paths, a double step costs 3, exercise the new path" begin
    start = 1
    goal = 4

    getneighbours(state) = [state + 2, state + 1, state - 1]
    cost(a, b) = b - a == 2 ? 3 : 1

    res = astar(getneighbours, start, goal; heuristic, cost)
    @test res.status == :success
    @test res.path == [1, 2, 3, 4]
    @test res.cost == 3
  end
end


"""
Struct for Standard American Option

		amOption=AmericanOption(T::num1,K::num2,isCall::Bool=true) where {num1 <: Number,num2 <: Number}

Where:\n
		T	=	Time to maturity of the Option.
		K	=	Strike Price of the Option.
		isCall  = true for CALL, false for PUT.
"""
mutable struct AmericanOption{num1 <: Number, num2 <: Number, numtype <: Number} <: AmericanPayoff{numtype}
    T::num1
    K::num2
    isCall::Bool
    function AmericanOption(T::num1, K::num2, isCall::Bool = true) where {num1 <: Number, num2 <: Number}
        if T <= 0.0
            error("Time to Maturity must be positive")
        elseif K <= 0.0
            error("Strike Price must be positive")
        else
            zero_typed = zero(num1) + zero(num2)
            return new{num1, num2, typeof(zero_typed)}(T, K, isCall)
        end
    end
end

export AmericanOption;

function payout(Sti::numtype_, amPayoff::AmericanOption) where {numtype_ <: Number}
    iscall = amPayoff.isCall ? 1 : -1
    return ((Sti - amPayoff.K) * iscall > 0.0) ? (Sti - amPayoff.K) * iscall : zero(numtype_)
end


"""
    AutoCorrelationFunction{T}

A representation of the autocorrelation function.

It supports accessing the coeficients starting from lag 0 by
- `coef(acf::AutoCorrelationFunction)`: implements StatsBase.coef
- `coef(acf::AutoCorrelationFunction, lag::Integer)`: coefficient for lag

Wrapping the vector of coefficients into its own type helps avoiding
method ambiguities.

# Examples
```jldoctest am; output = false
using StatsBase: coef
acf = AutoCorrelationFunction([1,0.4,0.1])
coef(acf) == [1,0.4,0.1]
coef(acf,1) == 0.4
# output
true
```
"""
struct AutoCorrelationFunction{T}
    coef::T
end
# implements StatsBase coef
AutoCorrelationFunction(coef::AbstractVector{<:Number}) =
    AutoCorrelationFunction{typeof(coef)}(coef)
coef(acf::AutoCorrelationFunction) = acf.coef
coef(acf::AutoCorrelationFunction, lag::Integer) = acf.coef[lag+1]

cormatrix_for_acf(n::Int,acf::AutoCorrelationFunction) =
    cormatrix_for_acf(n, coef(acf))

function cormatrix_for_acf(n::Int,acf::AbstractVector)
    nacf::Int = length(acf)
    corrM = BandedMatrix{Float64}(undef, (n,n), (nacf-1,nacf-1))
    corrM[band(0)] .= acf[1]
    for i in 1:(nacf-1)
      corrM[band(i)] .= corrM[band(-i)] .= acf[i+1]
    end
    corrM
end


function plot_convergence(exact, y0, tmax, Na, Nb, propagator::MagnusPropagator, subs, label)
    Ns,errors,times = Magnus.convergence(exact, y0, tmax, Na, Nb, propagator)
    subplot(subs[1])
    plot(Ns, errors, ".-", label = label)
    subplot(subs[2])
    plot(errors, 1e-3times, ".-", label = label)
end


########################################################
function realresampleddataforpca()
########################################################

    A = h5read("Data/AGNs.hdf5","simulations/Mrk509/AGNGALresamp")

    x = Matrix(A[1:3, 2, :]')  # select filters out of available six

    Σ =  Matrix(A[1:3, 3, :]')

    return x, Σ

end

########################################################
function realresampleddataforblr()
########################################################

    A = h5read("Data/AGNs.hdf5","simulations/Mrk509/AGNGALresamp")

    x = Matrix(A[1:6, 2, :]')  # select filters out of available six

    srt = sortperm(vec(A[1,1,:])) # sort according to time,
                                  # not necessary, just convenient for plotting

    idx  = [ones(Int, 110)   ; ones(Int, 110)*2 ; ones(Int, 110)*3 ; ones(Int, 110)*4 ; ones(Int, 110)*5 ; ones(Int, 110)*6]
    flux = [vec(x[srt,1])    ; vec(x[srt,2])    ; vec(x[srt,3])    ; vec(x[srt,4])    ; vec(x[srt,5])    ; vec(x[srt,6])]
    Σ    = [vec(A[1, 3, srt]); vec(A[2, 3, srt]); vec(A[3, 3, srt]); vec(A[4, 3, srt]); vec(A[5, 3, srt]); vec(A[6, 3, srt])];

    t    = A[1,1,srt] .- minimum(A[1,1,srt]) # no need to subtract,
                                             # just convenient that times starts at 0

    idx1 = findall(idx .== 1)
    idx2 = findall(idx .== 2)
    idx3 = findall(idx .== 3)
    idx4 = findall(idx .== 4)
    idx5 = findall(idx .== 5)
    idx6 = findall(idx .== 6)

    figure(1)
    cla()
    title("plotting three first LSST filters")
    plot3D(flux[idx1], flux[idx2], flux[idx3], "bo", alpha=0.4)

    figure(2)
    cla()
    plot(t, flux[idx1], ".", label = "1 - u")
    plot(t, flux[idx2], ".", label = "2 - g")
    plot(t, flux[idx3], ".", label = "3 - r")
    plot(t, flux[idx4], ".", label = "4 - i")
    plot(t, flux[idx5], ".", label = "5 - y")
    plot(t, flux[idx6], ".", label = "6 - z")

    legend()

    # return [t;t;t; t;t;t], flux, Σ, idx
    return [t;t;t], flux[1:330], Σ[1:330], idx[1:330]

end


module FESOM
using UUIDs
using Random
using Statistics
using LinearAlgebra

export fesom

function quality_of_fit(isf_m::Array{Float64,1},isf::Array{Float64,1},isf_error::Array{Float64,1})
    mean(broadcast((x,y,z) -> ((x - y)/z)^2,isf,isf_m,isf_error))
end

function quality_of_fit!(isf_m::Array{Float64,1},isf::Array{Float64,1},isf_error::Array{Float64,1})
    mean(broadcast!((x,y,z) -> ((x - y)/z)^2,isf_m,isf,isf_m,isf_error))
end

function set_isf_term(imaginary_time::Array{Float64,1},frequency_bins::Array{Float64,1},beta::Float64)
    b = beta;
    isf_term = Array{Float64,2}(undef,size(imaginary_time,1),size(frequency_bins,1));
    for i in 1:size(imaginary_time,1)
        t = imaginary_time[i];
        for j in 1:size(frequency_bins,1)
            f = frequency_bins[j];
            isf_term[i,j] = (exp(-t*f) + exp(-(b - t)*f));
            #isf_term[i,j] = (exp(-t*f) + exp(-(b - t)*f))/(1 + exp(-b*f));
        end
    end
    isf_term
end

function set_isf_trapezoidal!( isf_m::Array{Float64,1}, isf_m2::Array{Float64,1},
                               isf_term::Array{Float64,2}, isf_term2::Array{Float64,2},
                               dsf::Array{Float64,1}
                               )
    mul!(isf_m,isf_term,dsf);
    mul!(isf_m2,isf_term2,dsf);
    isf_m .+= isf_m2;
    nothing
end

function fesom( dsf::Array{Float64,1},isf::Array{Float64,1},isf_error::Array{Float64,1},
                frequency::Array{Float64,1},imaginary_time::Array{Float64,1};
                temperature::Float64 = 1.2,
                number_of_iterations::Int64 = 10000,
                stop_minimum_fitness::Float64 = 1.0,
                seed::Int64 = 1,
                track_stats::Bool = false,
                number_of_blas_threads::Int64 = 0)
    rng = MersenneTwister(seed);
    fesom(rng,dsf,isf,isf_error,frequency,imaginary_time,
          temperature = temperature,
          number_of_iterations = number_of_iterations,
          stop_minimum_fitness = stop_minimum_fitness,
          number_of_blas_threads = number_of_blas_threads)
end

function fesom( rng::MersenneTwister,
                dsf::Array{Float64,1},isf::Array{Float64,1},isf_error::Array{Float64,1},
                frequency::Array{Float64,1},imaginary_time::Array{Float64,1};
                temperature::Float64 = 1.2,
                number_of_iterations::Int64=10000,
                stop_minimum_fitness::Float64 = 1.0,
                track_stats::Bool = false,
                number_of_blas_threads::Int64 = 0)
    #_bts = ccall((:openblas_get_num_threads64_, Base.libblas_name), Cint, ())
    #println("Number of BLAS threads: $(_bts)");
    if number_of_blas_threads > 0
        LinearAlgebra.BLAS.set_num_threads(number_of_blas_threads)
    end

    beta = 1/temperature;
    moment0 = isf[1];
    random_vector = Array{Float64,1}(undef,size(dsf,1));
    dsf_new = Array{Float64,1}(undef,size(dsf,1));

    isf_term = set_isf_term(imaginary_time,frequency,beta);
    isf_term2 = copy(isf_term);

    df = frequency[2:end] .- frequency[1:size(frequency,1) - 1];
    dfrequency1 = zeros(size(frequency,1));
    dfrequency2 = zeros(size(frequency,1));
    for i in 1:(size(frequency,1) - 1)
        dfrequency1[i] = df[i]/2
        dfrequency2[i+1] = df[i]/2
    end
    isf_term .*= dfrequency1';
    isf_term2 .*= dfrequency2';

    isf_m = Array{Float64,1}(undef,size(isf,1));
    isf_m2 = Array{Float64,1}(undef,size(isf,1));

    set_isf_trapezoidal!(isf_m, isf_m2, isf_term, isf_term2, dsf );

    χ2 = quality_of_fit!(isf_m,isf,isf_error);
    χ2_new = 0.0;
    normalization = 1.0;
    normalization_term1 = (1 .+ exp.(-beta .* frequency)) .* dfrequency1;
    normalization_term2 = (1 .+ exp.(-beta .* frequency)) .* dfrequency2;
    normalization_term = normalization_term1 .+ normalization_term2;

    if track_stats
        minimum_fitness = zeros(number_of_iterations);
    end
    nsteps = 0;
    for i in 1:number_of_iterations
        randn!(rng,random_vector);
        broadcast!((x,y) -> abs(x*(1 + y)),dsf_new, dsf, random_vector);

        #normalization: store intermediate results in random_vector
        normalization = dot(dsf_new,normalization_term);
        broadcast!(*,dsf_new, dsf_new, moment0/normalization);

        set_isf_trapezoidal!(isf_m, isf_m2, isf_term, isf_term2, dsf_new );

        χ2_new = quality_of_fit!(isf_m,isf,isf_error);
        if χ2_new < χ2
            dsf_new, dsf = dsf, dsf_new;
            χ2_new, χ2 = χ2, χ2_new;
        end

        nsteps +=1 ;

        if track_stats
            minimum_fitness[i] = χ2;
        end

        if χ2 < stop_minimum_fitness
            break
        end
    end
    println(χ2)

    if track_stats
        return rng,dsf,χ2,minimum_fitness[1:nsteps]
    else
        return rng,dsf,χ2,zeros(2)
    end

end
end


"""
Functions for applying empirical Bayes to a set of test statistics.
""";

using Discretizers
using Distributions
using Interpolations
using DataFrames
using GLM

"""
    discretize_test_statistics(test_statistics::AbstractArray{<:Real}, n)

Bin the values in test_statistics into n uniform width bins.
Returns the midpoints of the bins, the counts in each bin and the bin width.
"""
function discretize_test_statistics(test_statistics::AbstractArray{<:Real}, n)
    try
        n = Int(n) # Bad form, but allows values like 10.0 to be silently used as 10
    catch e
        error("n must be an Int or convertible to an Int")
    end

    # Calculate midpoints
    min_ts, max_ts = extrema(test_statistics)
    bin_width = (max_ts - min_ts) / n
    midpoints = [min_ts + i*bin_width - bin_width/2 for i in 1:n]

    # Get counts for each bin
    bin_edges = binedges(DiscretizeUniformWidth(n), test_statistics)
    lindisc = LinearDiscretizer(bin_edges)
    counts = get_discretization_counts(lindisc, test_statistics)

    return midpoints, counts, bin_width
end

"""
    keep_bottom_proportion_of_hist(midpoints, counts, proportion_to_keep, verbose)

Remove the highest proportion of values from a histogram. Returns the new
midpoints and counts for the truncated histogram.

# Arguments
- `midpoints` : list of midpoints of the histogram bins.
- `counts` : list of counts for the histogram bins.
- `proportion_to_keep` : proportion of lowest values in the histogram to keep.
"""
function keep_bottom_proportion_of_hist(midpoints, counts, proportion_to_keep, verbose)
    @assert 0.0 < proportion_to_keep <= 1.0

    # Number of values to keep
    num_vals_total = sum(counts)
    to_keep = Int(round(proportion_to_keep * num_vals_total))

    # Iterate through counts to find which bins to keep
    # Keep updating total until it is more than we want to keep
    cur_bin = 1
    cur_total = 0
    while (cur_total < to_keep) && (cur_bin <= length(counts))
        cur_total += counts[cur_bin]
        cur_bin += 1
    end

    # Only keep bins whose entire counts are within the kept interval
    final_bin = cur_total == to_keep ? cur_bin-1 : cur_bin-2
    counts = counts[1:final_bin]
    midpoints = midpoints[1:final_bin]

    if verbose
        kept_proportion = sum(counts) / num_vals_total
        println("Kept $kept_proportion of test statistics")
    end

    return midpoints, counts
end

"""
    estimate_fit_parameters(num_bins, bin_width, midpoints, counts, distr::Type{Gamma})

Estimates the poisson parameters required to fit a mode matched null distribution
to the input histogram. Solves the regression problem outlined in equations (6)
and (8) in Schwarzman 2009 (https://projecteuclid.org/euclid.aoas/1231424213).
Returns the parameters C, eta1, eta2
"""
function estimate_fit_parameters(num_bins, bin_width, midpoints, counts, distr::Type{Gamma})
    # Get the log of the counts
    log_counts = log.(counts)

    # Filter out non finite
    log_counts = [isfinite(x) ? x : eps() for x in log_counts]

    # Minus difference h to create regression target
    h = log(num_bins * bin_width)
    y = log_counts .- h

    # Construct and solve the regression problem
    data = DataFrame(Y = y, X1 = midpoints, X2 = log.(midpoints))
    C, eta1, eta2 = coef(glm(@formula(Y ~ X1 + X2), data, Normal(), minstepfac = eps()))

    return C, eta1, eta2
end
"""
    estimate_fit_parameters(num_bins, bin_width, midpoints, counts, distr::Type{Normal})

Estimates the parameters required to fit a mode matched null distribution to the
input histogram. Solves the regression problem outlined in equation (4.5) and (4.6)
in Efron 2007 (https://projecteuclid.org/euclid.aos/1188405614).
Returns the parameters beta0, beta1, beta2
"""
function estimate_fit_parameters(num_bins, bin_width, midpoints, counts, distr::Type{Normal})
    # Get the log of the counts
    log_counts = log.(counts)

    # Filter out non finite
    log_counts = [isfinite(x) ? x : eps() for x in log_counts]

    # Construct and solve the regression problem
    data = DataFrame(Y = log_counts, X1 = midpoints, X2 = midpoints.^2)
    beta0, beta1, beta2 = coef(glm(@formula(Y ~ X1 + X2), data, Normal(), minstepfac = eps()))

    return beta0, beta1, beta2
end

"""
    get_distr_parameters(C, eta1, eta2, distr::Type{Gamma})

Gets the parameters which describe a gamma distribution based on the poisson
parameters input. Converts the poisson parameter inputs into a chi squared
distribution, then converts that into a gamma distribution.
"""
function get_distr_parameters(C, eta1, eta2, distr::Type{Gamma})

    # Convert eta to chi squared parameters
    a = -(1 / (2*eta1))
    v = 2 * (eta2 + 1)

    # Convert chi squared parameters to gamma shape and scale parameters
    k = v/2
    theta = 2*a

    # TODO: Decide what to do about p0
    # Find p0 by solving the following:
    # log(p0) = C + log(gammafunc(eta2 + 1) / (-eta1)^(eta2 + 1))
    # p0 = exp(C + log(gamma(eta2 + 1) / (-eta1)^(eta2 + 1)))
    p0 = 0

    params_valid = k > zero(k) && theta > zero(theta)
    if !params_valid
        error_string = "Error: resulting gamma parameters are invalid. Got k = ", k, " and Θ = ", theta, ".
        Input data might not be gamma distributed, or number of bins needs to change.
        Exiting"
        error(error_string)
    end

    return k, theta, p0
end
"""
    get_distr_parameters(beta0, beta1, beta2, distr::Type{Normal})

Gets the parameters which describe a normal distribution based on the
parameters input. Converts the quadratic parameter inputs into a normal
distribution.
"""
function get_distr_parameters(beta0, beta1, beta2, distr::Type{Normal})

    sigma = (-2 * beta2)^-(1/2)

    mu = beta1 * sigma^2

    # TODO: Decide what to do about p0
    p0 = 0

    return mu, sigma, p0
end

"""
    fit_null_distribution(midpoints, counts, num_bins, bin_width, proportion_to_keep, distr, verbose=true)

Fit a distribution to the discretized inputs. Uses mode-matching, since
the input test statistics may have been truncated.
"""
function fit_null_distribution(midpoints, counts, num_bins, bin_width, proportion_to_keep, distr; verbose = true)

    # Remove highest test statistics, that dont fit the zero assumption
    midpoints, counts = keep_bottom_proportion_of_hist(midpoints, counts, proportion_to_keep, verbose)

    # Estimate parameters using regression
    param0, param1, param2 = estimate_fit_parameters(num_bins, bin_width, midpoints, counts, distr)

    # TODO: Decide what to do about p0
    # Calculate distribution parameters
    k, theta, p0 = get_distr_parameters(param0, param1, param2, distr)

    # Create distribution
    f0 = distr(k, theta)

    return f0
end

"""
    fit_mixture_distribution(midpoints, counts, bin_width)

Fit a cubic spline distribution to the discretized inputs.
"""
function fit_mixture_distribution(midpoints, counts, bin_width)
    # Add some zero bins on the beginning to prevent curve going negative
    number_of_zero_bins = 10 #3
    zero_bins = zeros(number_of_zero_bins)
    pushfirst!(counts, zero_bins...)
    for i in 1 : number_of_zero_bins
        pushfirst!(midpoints, midpoints[1] - bin_width)
    end

    # Normalize counts
    num_statsitics = sum(counts)
    counts = counts / (num_statsitics * bin_width)

    # Fit spline to histogram points
    interpolation = interpolate(counts, BSpline(Cubic(Line(OnCell()))))
    r = range(midpoints[1], stop=midpoints[end], length=length(midpoints))
    scaled_interpolation = Interpolations.scale(interpolation, r)
    scaled_extrapolation = extrapolate(scaled_interpolation, Line())

    fhat(x) = max(scaled_extrapolation(x), 0)
    fhat(x::AbstractArray) = [fhat(i) for i in x]

    return fhat
end

"""
    calculate_posterior(test_statistics, priors, null_distr, mixture_pdf, tail, w0=0.0)

Calculate the empirical Bayes posterior using the priors, null distribution and
mixture distribution.

# Arguments
- `test_statistics` : list of test statistics.
- `priors` : one dimensional list of prior values that correspond with the
   test_statistics, such that `priors[i]` is the prior for `test_statistics[i]`.
- `null_distr` : the null distribution of test statistics
- `mixture_pdf` : function representing the mixture distribution of test statistics
- `tail` : Whether the test is two-tailed (:two) or one-tailed (:lower or :upper)
- `w0` : Default constant for the prior calculation
- `null_value = -Inf` : what to set the unwanted tail values to
"""
function calculate_posterior(test_statistics, priors, null_distr, mixture_pdf, tail; w0=0.0, null_value = -Inf)
    @assert length(test_statistics) == size(priors, 1)

    null_pdf(x) = pdf(null_distr, x)

    # Prior function
    prior_fn(x) = exp(w0) / ( exp(w0) + exp(x) )

    num_test_statistics = length(test_statistics)

    # Posterior array
    posterior = Array{Float64}(undef, num_test_statistics)

    for i in 1:num_test_statistics
        ts = test_statistics[i]
        prior_val = prior_fn(priors[i])
        null_val = null_pdf(ts)
        mix_val = mixture_pdf(ts)

        # if mixture distr equals 0, then just return a 0 posterior
        if mix_val == zero(mix_val)
            posterior[i] = 0.0
            continue
        end

        fdr = null_val / mix_val
        p1 = 1 - prior_val * fdr
        posterior[i] = p1
    end

    # Get rid of the lower tail for an upper-tailed test, or vice-versa
    if tail == :upper || tail == :lower
        null_mode = mode(null_distr)
        compare = tail == :upper ? (<=) : (>=)
        zero_indices = findall(t -> compare(t, null_mode), test_statistics)
        posterior[zero_indices] .= null_value
    end

    return posterior
end

"""
    calculate_posterior(test_statistics, null_distr, mixture_pdf, tail, w0=0.0)

Calculate the empirical Bayes posterior using no priors, null distribution and
mixture distribution.

# Arguments
- `test_statistics` : list of test statistics.
- `null_distr` : the null distribution of test statistics
- `mixture_pdf` : function representing the mixture distribution of test statistics
- `tail` : Whether the test is two-tailed (:two) or one-tailed (:lower or :upper)
- `w0` : Default constant for the prior calculation
- `null_value = -Inf` : what to set the unwanted tail values to
"""
function calculate_posterior(test_statistics, null_distr, mixture_pdf, tail; w0=0.0, null_value = -Inf)
    priors = [0 for _ in test_statistics]
    return calculate_posterior(test_statistics, priors, null_distr, mixture_pdf, tail, w0=w0)
end

"""
    empirical_bayes(test_statistics, priors, num_bins, distr, proportion_to_keep=1.0, tail=:two, w0=0.0)

Calculate the empirical Bayes posteriors of the input statistics using the priors.

# Arguments
- `test_statistics` : list of test statistics.
- `priors` : one dimensional list of prior values that correspond with the
   test_statistics, such that `priors[i]` is the prior for `test_statistics[i]`.
- `num_bins::Integer` : number of uniform width bins to discretize into.
- `distr` : form of the null distribution to be fitted.
- `proportion_to_keep=1.0` : Proportion of lowest test statistics to
   keep when calculating null distribution.
- `tail=:two` : Whether the test is two-tailed (:two) or one-tailed (:lower or :upper)
- `w0` : Default constant for the prior calculation
- `null_value = -Inf` : what to set the unwanted tail values to
"""
function empirical_bayes(test_statistics, priors, num_bins, distr; proportion_to_keep=1.0, tail=:two, w0=0.0, null_value = -Inf)
     midpoints, counts, bin_width = discretize_test_statistics(test_statistics, num_bins)
     null_distr = fit_null_distribution(midpoints, counts, num_bins, bin_width, proportion_to_keep, distr)
     mixture_pdf = fit_mixture_distribution(midpoints, counts, bin_width)
     posteriors = calculate_posterior(test_statistics, priors, null_distr, mixture_pdf, tail, w0=w0, null_value = null_value)
     return posteriors
end
function empirical_bayes(test_statistics, priors, num_bins, distr::Symbol; proportion_to_keep=1.0, tail=:two, w0=0.0, null_value = -Inf)
    return empirical_bayes(test_statistics, priors, num_bins, get_distr(distr); proportion_to_keep=proportion_to_keep, tail=tail, w0=w0, null_value = null_value)
end

"""
    empirical_bayes(test_statistics, num_bins, distr, proportion_to_keep=1.0, tail=:two, w0=0.0)

Calculate the empirical Bayes posteriors of the input statistics with null priors.

# Arguments
- `test_statistics` : list of test statistics.
- `num_bins::Integer` : number of uniform width bins to discretize into.
- `distr` : form of the null distribution to be fitted.
- `proportion_to_keep=1.0` : Proportion of lowest test statistics to
   keep when calculating null distribution.
- `tail=:two` : Whether the test is two-tailed (:two) or one-tailed (:lower or :upper)
- `w0` : Default constant for the prior calculation
- `null_value = -Inf` : what to set the unwanted tail values to
"""
function empirical_bayes(test_statistics, num_bins, distr; proportion_to_keep=1.0, tail=:two, w0=0.0, null_value = -Inf)
    priors = [0 for _ in test_statistics]
    return empirical_bayes(test_statistics, priors, num_bins, distr, proportion_to_keep=proportion_to_keep, tail=tail, w0=w0, null_value = null_value)
end
function empirical_bayes(test_statistics, num_bins, distr::Symbol; proportion_to_keep=1.0, tail=:two, w0=0.0, null_value = -Inf)
    return empirical_bayes(test_statistics, num_bins, get_distr(distr); proportion_to_keep=proportion_to_keep, tail=tail, w0=w0, null_value = null_value)
end

# Utility function for getting a distribution from a symbol, so that the caller needn't
# depend on Distributions.jl
function get_distr(d::Symbol)
    distrs = Dict(
        :Gamma => Gamma,
        :Normal => Normal
    )
    return distrs[d]
end


using Krylov, Stopping

include("random_coordinate_descent_method.jl")
include("stop_random_coordinate_descent_method.jl")
include("stop_random_coordinate_descent_method_LS.jl")

n = 5000;
A = rand(n, n)
sol = rand(n)
b = A * sol

#@time stp = StopRandomizedCD(A,b, max_iter = 1000)

x0 = zeros(size(A, 2))
pb = LinearSystem(A, b)
s0 = GenericState(x0, similar(b))
mcnt = Stopping._init_max_counters_linear_operators()
#
@time meta = StoppingMeta(
  max_cntrs = mcnt,
  atol = 1e-7,
  rtol = 1e-15,
  max_iter = 99,
  retol = false,
  tol_check = (atol, rtol, opt0) -> (atol + rtol * opt0),
  optimality_check = (pb, state) -> state.res,
)
s0 = GenericState(x0, similar(b))
@time stp = LAStopping(pb, meta, cheap_stop_remote_control(), s0)
@time StopRandomizedCD(stp)

@time x, OK = RandomizedCD(A, b, max_iter = 100)

###############################################################################
#
# We compare the stopping_randomized_CD with the same version with an
# artificial substopping created and called at each iteration.
# It appears that the expansive part is to create the SubStopping, and in
# particular create the StoppingMeta.
#
@time meta = StoppingMeta(
  max_cntrs = mcnt,
  atol = 1e-7,
  rtol = 1e-15,
  max_iter = 99,
  retol = false,
  tol_check = (atol, rtol, opt0) -> (atol + rtol * opt0),
  optimality_check = (pb, state) -> state.res,
)
s0 = GenericState(x0, similar(b))
@time stp2 = LAStopping(pb, meta, cheap_stop_remote_control(), s0)
@time StopRandomizedCD_LS(stp2)

;


## Implements bitonic merge networks that can merge two sorted SIMD vectors.
## The generated function supports any possible type and size.

using SIMD

"""
    bitonic_merge(input_a::Vec{N,T}, input_b::Vec{N,T}) where {N,T}

Merges two `SIMD.Vec` objects of the same type and size using a bitonic sort network. The inputs are assumed to be sorted. Returns a pair of vectors with the first and second halves of the merged sequence.
"""
@generated function bitonic_merge(input_a::Vec{N,T}, input_b::Vec{N,T}) where {N,T}

    pat = Val{ntuple(x->N-x, N)}

    ex = [
        Expr(:meta, :inline),
        :(la_0 = input_a),
        :(lb_0 = shufflevector(input_b, $pat)),
        :(L_0 = min(la_0, lb_0)),
        :(H_0 = max(la_0, lb_0))
    ]

    p = mylog(N)
    for n in 1:p
        la = Symbol("la_", n)
        lb = Symbol("lb_", n)
        Lp = Symbol("L_", n-1)
        Hp = Symbol("H_", n-1)
        L = Symbol("L_", n)
        H = Symbol("H_", n)

        ih = inverse_shuffle(bitonic_step(2^(n), 2^(p-n+1)))
        sh = bitonic_step(2^(n+1), 2^(p-n))
        pat_a = Val{tuple((ih[sh[1:N]].-1)...)}
        pat_b = Val{tuple((ih[sh[(N+1):end]].-1)...)}

        append!(ex, [
            :($la = shufflevector($Lp, $Hp, $pat_a)),
            :($lb = shufflevector($Lp, $Hp, $pat_b)),
            :($L = min($la, $lb)),
            :($H = max($la, $lb))
        ])
    end

    la = Symbol("la_", p+1)
    lb = Symbol("lb_", p+1)
    Lp = Symbol("L_", p)
    Hp = Symbol("H_", p)

    ih = inverse_shuffle(bitonic_step(2N, 1))
    pat_a = Val{tuple((ih[1:N].-1)...)}
    pat_b = Val{tuple((ih[(N+1):end].-1)...)}

    append!(ex, [
        :($la = shufflevector($Lp, $Hp, $pat_a)),
        :($lb = shufflevector($Lp, $Hp, $pat_b)),
        :(($la, $lb))
    ])

    quote $(ex...) end
end

function inverse_shuffle(tt)
    out = copy(tt)
    for n in 1:length(tt)
        out[tt[n]] = n
    end
    out
end

function bitonic_step(G, Glen)
    N = G * Glen
    gg = reshape(1:N, Glen, G)
    t_lh = [gg[:,1:2:end][:]; gg[:,2:2:end][:]]
end

"""
    merge_vecs(input...)

This brave function merges 4, 8 vectors or even more. Relies on the great ability from SIMD.jl and LLVM to handle
large vectors. Performance not yet tested."""
@generated function merge_vecs(input::Vararg{Vec{N,T}, L})::Vec{N*L,T} where {L,N,T}

    ex = [Expr(:meta, :inline)]

    for l in 1:L
        m = Symbol("m_0_", l)
        push!(ex, :($m = input[$l]))
    end

    p = mylog(L)
    for l in 1:p
        for b in 1:2^(p-l)
            m = Symbol("m_", l, "_", b)
            a1 = Symbol("m_", l-1, "_", b*2-1)
            a2 = Symbol("m_", l-1, "_", b*2)
            push!(ex, :($m = bitonic_merge_concat($a1, $a2)))
        end
    end

    m = Symbol("m_", p, "_", 1)
    push!(ex, :(return $m))

    quote $(ex...) end
end

@inline function bitonic_merge_concat(input_a::Vec{N,T}, input_b::Vec{N,T})::Vec{2*N,T} where {N,T}
    m_a, m_b = bitonic_merge(input_a, input_b)
    concat(m_a, m_b)
end

"""
    bitonic_merge_interleaved(output_a, output_b, input_a, input_b, Val(N))

Merges K pairs of vectors using a bitonic sort network, with interleaved execution. The inputs are assumed to be sorted. Inputs and outputs are acessed directly in memory.
"""
@generated function bitonic_merge_interleaved(
    output::AbstractArray{T,3},
    input::Array{Ptr{T},1},
    ::Val{V},
    J,
    ::Val{K}
) where {T,V,K}

    allex = []
    for indy in 1:K
        iSymbol(x...) = Symbol("indy_", indy, "_", x...)
        pat = Val{ntuple(x->V-x, V)}

        n = 0
        la = iSymbol("la_", n)
        lb = iSymbol("lb_", n)
        L = iSymbol("L_", n)
        H = iSymbol("H_", n)
        ex = [
            :($la = vloada(Vec{V,T}, pointer(@view output[:,J,$indy]))),
            :($lb = shufflevector(vloada(Vec{V,T}, input[$indy]), $pat)),
            :($L = min($la, $lb)),
            :($H = max($la, $lb))
        ]

        p = mylog(V)
        for n in 1:p
            la = iSymbol("la_", n)
            lb = iSymbol("lb_", n)
            Lp = iSymbol("L_", n-1)
            Hp = iSymbol("H_", n-1)
            L = iSymbol("L_", n)
            H = iSymbol("H_", n)

            ih = inverse_shuffle(bitonic_step(2^(n), 2^(p-n+1)))
            sh = bitonic_step(2^(n+1), 2^(p-n))
            pat_a = Val{tuple((ih[sh[1:V]].-1)...)}
            pat_b = Val{tuple((ih[sh[(V+1):end]].-1)...)}

            append!(ex, [
                :($la = shufflevector($Lp, $Hp, $pat_a)),
                :($lb = shufflevector($Lp, $Hp, $pat_b)),
                :($L = min($la, $lb)),
                :($H = max($la, $lb))
            ])
        end

        la = iSymbol("la_", p+1)
        lb = iSymbol("lb_", p+1)
        Lp = iSymbol("L_", p)
        Hp = iSymbol("H_", p)

        ih = inverse_shuffle(bitonic_step(2V, 1))
        pat_a = Val{tuple((ih[1:V].-1)...)}
        pat_b = Val{tuple((ih[(V+1):end].-1)...)}

        append!(ex, [
            :(vstorea(shufflevector($Lp, $Hp, $pat_a), pointer(@view output[:,J,$indy]))),
            :(vstorea(shufflevector($Lp, $Hp, $pat_b), pointer(@view output[:,J+1,$indy])))
        ])
        push!(allex, ex)
    end

    inteleaved = [x for xx in zip(allex...) for x in xx]

    final_ex = [
        Expr(:meta, :inline),
        inteleaved...
    ]

    quote $(final_ex...) end
end


#####################################
#     Basic interface functions     #
#####################################

"""
    cost,i1,i2 = dtw(seq1, seq2, [dist=SqEuclidean])

Find a set of indices (`i1`,`i2`) that align two time series (`seq1`,`seq2`) by
dynamic time warping. Also returns the distance (after warping) according to
the SemiMetric `dist`, which defaults to squared Euclidean distance (see
Distances.jl). If `seq1` and `seq2` are matrices, each column is considered
an observation.
"""
function dtw(
        seq1::Sequence,
        seq2::Sequence,
        dist::SemiMetric=SqEuclidean()
    )

    D = dtw_cost_matrix(seq1, seq2, dist)
    return trackback(D)
end

"""
    cost,i1,i2 = dtw(seq1,seq2,i2min,i2max,[dist=SqEuclidean])

Do DTW to align `seq1` and `seq2` confined to a window. Vectors `i2min` and
`i2max` specify (inclusive) lower and upper bounds for `seq2` for each index in
`seq1`. Thus, `i2min` and `i2max` are required to be the same length as `seq1`.
"""
function dtw(
        seq1::Sequence,
        seq2::Sequence,
        i2min::AbstractVector,
        i2max::AbstractVector,
        dist::SemiMetric = SqEuclidean()
    )

    D = dtw_cost_matrix(seq1, seq2, i2min, i2max, dist)
    return trackback(D)
end

##############################
#  Cost matrix computations  #
##############################
function dtw_cost_matrix(
        seq1::Sequence{N,T},
        seq2::Sequence{N,T},
        dist::SemiMetric=SqEuclidean()
    ) where {N,T}
    # Build the cost matrix
    m = length(seq2)
    n = length(seq1)
    D = zeros(T, m, n)

    # Initialize first column and first row
    D[1,1] = evaluate(dist, seq1[1], seq2[1])
    for r=2:m
        D[r,1] = D[r-1,1] + evaluate(dist, seq1[1], seq2[r])
    end
    for c=2:n
        D[1,c] = D[1,c-1] + evaluate(dist, seq1[c], seq2[1])
    end

    # Complete the cost matrix
    for c=2:n
        for r=2:m
            best_neighbor_cost = min(D[r-1,c], D[r-1,c-1], D[r,c-1])
            D[r,c] = best_neighbor_cost + evaluate(dist, seq1[c], seq2[r])
        end
    end

    return D
end

function dtw_cost_matrix(
        seq1::Sequence{N,T},
        seq2::Sequence{N,T},
        i2min::AbstractVector{U},
        i2max::AbstractVector{U},
        dist::SemiMetric = SqEuclidean()
    ) where {N,T,U<:Integer}
    m = length(seq2) # of rows in cost matrix
    n = length(seq1) # of columns in cost matrix
    n == length(i2min) || throw(ArgumentError("i2min does not match length of seq1."))
    n == length(i2max) || throw(ArgumentError("i2max does not match length of seq1."))
    1 == i2min[1] || throw(ArgumentError("i2min must start at 1."))
    m == i2max[end] || throw(ArgumentError("i2max must end at length(seq2)."))

    # Build the (n x m) cost matrix into a WindowedMatrix, because it's ragged.
    # That type gives efficient storage with convenient [r,c] indexing and returns
    # Inf when accessed outside the window.
    D = WindowedMatrix(i2min, i2max, Inf)

    # First column first
    D[1,1] = evaluate(dist, seq1[1], seq2[1])
    for r=2:i2max[1]
        D[r,1] = D[r-1,1]  + evaluate(dist, seq1[1], seq2[r])
    end

    # Complete the cost matrix from columns 2 to m.
    for c=2:n
        for r=i2min[c]:i2max[c]
            best_neighbor_cost = min(D[r-1,c], D[r-1,c-1], D[r,c-1])
            D[r,c] = best_neighbor_cost + evaluate(dist, seq1[c], seq2[r])
        end
    end

    return D
end

########################################
#  Find Best Path through Cost Matrix  #
########################################

"""
    cost,cols,rows = trackback(D::Matrix)

Given the cost matrix `D`, computes the optimal track from end to beginning.
Returns `cols` and `rows` which are vectors respectively holding the track.
"""
function trackback(D::AbstractMatrix{T}) where {T<:Number}

    # initialize trackback throught rows/columns
    r,c = size(D)
    rows,cols = Int[r],Int[c]

    # estimate that we'll need N⋅logN elements
    N = max(r,c)
    sz = 2*N
    sizehint!(rows,sz)
    sizehint!(cols,sz)

    # do trackback
    while r > 1 && c > 1
        tb,r,c = indmin3(D[r-1,c-1], D[r-1,c], D[r,c-1], r, c)
        push!(rows,r)
        push!(cols,c)
    end
    # Possibly either r>1 or c>1 at this point (but not both).
    # Add the unfinished part of the track to reach [1,1]
    for r=r-1:-1:1
        push!(rows,r)
        push!(cols,1)
    end
    for c=c-1:-1:1
        push!(rows,1)
        push!(cols,c)
    end
    return D[end,end], reverse(cols), reverse(rows)
end

###############################################
# Wrapper functions for abstract array inputs #
###############################################

function dtw(
        seq1::AbstractArray,
        seq2::AbstractArray,
        dist::SemiMetric=SqEuclidean()
    )
    dtw(Sequence(seq1),Sequence(seq2),dist)
end

function dtw(
        seq1::AbstractArray,
        seq2::AbstractArray,
        i2min::AbstractVector,
        i2max::AbstractVector,
        dist::SemiMetric = SqEuclidean()
    )
    dtw(Sequence(seq1),Sequence(seq2),i2min,i2max,dist)
end


plan_flows_fft(a::Array, args...; kwargs...) = plan_fft(a, args...; kwargs...)
plan_flows_rfft(a::Array, args...; kwargs...) = plan_rfft(a, args...; kwargs...)

"""
A placeholder grid object for `0D` problems (in other words, systems of ODEs).
"""
struct ZeroDGrid{T, A} <: AbstractGrid{T, A} end

"""
A one-dimensional `grid` object.
"""
struct OneDGrid{T<:AbstractFloat, Tk, Tx, Tfft, Trfft} <: AbstractGrid{T, Tk}
        nx :: Int
        nk :: Int
       nkr :: Int

        dx :: T
        Lx :: T

         x :: Tx

         k :: Tk
        kr :: Tk
    invksq :: Tk
   invkrsq :: Tk

   fftplan :: Tfft
  rfftplan :: Trfft

  # Range objects that access the aliased part of the wavenumber range
    kalias :: UnitRange{Int}
   kralias :: UnitRange{Int}
end

"""
    OneDGrid(nx, Lx; x0=-Lx/2, nthreads=Sys.CPU_THREADS, effort=FFTW.MEASURE,
                      T=Float64, dealias=1/3, ArrayType=Array)

Constructs a OneDGrid object with size `Lx`, resolution `nx`, and leftmost position `x0`.
FFT plans are generated for `nthreads` CPUs using FFTW flag `effort`. The float type is `T`
and the array types is `ArrayType`. The `dealias` keyword determines the highest  wavenubers
that are being zero-ed out by `dealias()` function; 1/3 is the nominal value for quadratic
nonlinearities.
"""
function OneDGrid(nx, Lx; x0=-Lx/2, nthreads=Sys.CPU_THREADS, effort=FFTW.MEASURE,
                  T=Float64, dealias=1/3, ArrayType=Array)

  dx = Lx/nx

   nk = nx
  nkr = Int(nx/2 + 1)

  # Physical grid
  x = range(T(x0), step=T(dx), length=nx)

  # Wavenubmer grid
   k = ArrayType{T}( fftfreq(nx, 2π/Lx*nx))
  kr = ArrayType{T}(rfftfreq(nx, 2π/Lx*nx))

   invksq = @. 1 / k^2
  invkrsq = @. 1 / kr^2
  CUDA.@allowscalar  invksq[1] = 0
  CUDA.@allowscalar invkrsq[1] = 0

  FFTW.set_num_threads(nthreads)
   fftplan = plan_flows_fft(ArrayType{Complex{T}, 1}(undef, nx), flags=effort)
  rfftplan = plan_flows_rfft(ArrayType{T, 1}(undef, nx), flags=effort)

  kalias, kralias = getaliasedwavenumbers(nk, nkr, dealias)

  Tx = typeof(x)
  Tk = typeof(k)
  Tfft = typeof(fftplan)
  Trfft = typeof(rfftplan)

  return OneDGrid{T, Tk, Tx, Tfft, Trfft}(nx, nk, nkr, dx, Lx, x, k, kr,
                                      invksq, invkrsq, fftplan, rfftplan, kalias, kralias)
end

"A two-dimensional `grid` object."
struct TwoDGrid{T<:AbstractFloat, Tk, Tx, Tfft, Trfft} <: AbstractGrid{T, Tk}
        nx :: Int
        ny :: Int
        nk :: Int
        nl :: Int
       nkr :: Int

        dx :: T
        dy :: T
        Lx :: T
        Ly :: T

         x :: Tx
         y :: Tx

         k :: Tk
         l :: Tk
        kr :: Tk
       Ksq :: Tk
    invKsq :: Tk
      Krsq :: Tk
   invKrsq :: Tk

   fftplan :: Tfft
  rfftplan :: Trfft

  # Range objects that access the aliased part of the wavenumber range
    kalias :: UnitRange{Int}
   kralias :: UnitRange{Int}
    lalias :: UnitRange{Int}
end

"""
    TwoDGrid(nx, Lx, ny=nx, Ly=Lx; x0=-Lx/2, y0=-Ly/2, nthreads=Sys.CPU_THREADS,
                      effort=FFTW.MEASURE, T=Float64, dealias=1/3, ArrayType=Array)

Constructs a TwoDGrid object with size `Lx`, `Ly`, resolution `nx`, `ny`, and leftmost
positions `x0`, `y0`. FFT plans are generated for `nthreads` CPUs using
FFTW flag `effort`. The float type is `T` and the array types is `ArrayType`.
The `dealias` keyword determines the highest wavenubers that are being zero-ed
out by `dealias()` function; 1/3 is the nominal value for quadratic nonlinearities.
"""
function TwoDGrid(nx, Lx, ny=nx, Ly=Lx; x0=-Lx/2, y0=-Ly/2, nthreads=Sys.CPU_THREADS,
                  effort=FFTW.MEASURE, T=Float64, dealias=1/3, ArrayType=Array)

  dx = Lx/nx
  dy = Ly/ny

   nk = nx
   nl = ny
  nkr = Int(nx/2 + 1)

  # Physical grid
  x = range(T(x0), step=T(dx), length=nx)
  y = range(T(y0), step=T(dy), length=ny)

  # Wavenubmer grid
   k = ArrayType{T}(reshape( fftfreq(nx, 2π/Lx*nx), (nk, 1)))
   l = ArrayType{T}(reshape( fftfreq(ny, 2π/Ly*ny), (1, nl)))
  kr = ArrayType{T}(reshape(rfftfreq(nx, 2π/Lx*nx), (nkr, 1)))

     Ksq = @. k^2 + l^2
  invKsq = @. 1 / Ksq
  CUDA.@allowscalar invKsq[1, 1] = 0

     Krsq = @. kr^2 + l^2
  invKrsq = @. 1 / Krsq
  CUDA.@allowscalar invKrsq[1, 1] = 0

  # FFT plans
  FFTW.set_num_threads(nthreads)
  fftplan = plan_flows_fft(ArrayType{Complex{T}, 2}(undef, nx, ny), flags=effort)
  rfftplan = plan_flows_rfft(ArrayType{T, 2}(undef, nx, ny), flags=effort)

  # Index endpoints for aliasfrac i, j wavenumbers
  kalias, kralias = getaliasedwavenumbers(nk, nkr, dealias)
  lalias, _ = getaliasedwavenumbers(nl, nl, dealias)

  Tx = typeof(x)
  Tk = typeof(k)
  Tfft = typeof(fftplan)
  Trfft = typeof(rfftplan)

  return TwoDGrid{T, Tk, Tx, Tfft, Trfft}(nx, ny, nk, nl, nkr, dx, dy, Lx, Ly, x, y, k, l, kr, Ksq, invKsq, Krsq, invKrsq,
                                      fftplan, rfftplan, kalias, kralias, lalias)
end

"A three-dimensional `grid` object."
struct ThreeDGrid{T<:AbstractFloat, Tk, Tx, Tfft, Trfft} <: AbstractGrid{T, Tk}
        nx :: Int
        ny :: Int
        nz :: Int
        nk :: Int
        nl :: Int
        nm :: Int
       nkr :: Int

        dx :: T
        dy :: T
        dz :: T
        Lx :: T
        Ly :: T
        Lz :: T

         x :: Tx
         y :: Tx
         z :: Tx

         k :: Tk
         l :: Tk
         m :: Tk
        kr :: Tk
       Ksq :: Tk
    invKsq :: Tk
      Krsq :: Tk
   invKrsq :: Tk

   fftplan :: Tfft
  rfftplan :: Trfft

  # Range objects that access the aliased part of the wavenumber range
    kalias :: UnitRange{Int}
   kralias :: UnitRange{Int}
    lalias :: UnitRange{Int}
    malias :: UnitRange{Int}
end

"""
    ThreeDGrid(nx, Lx, ny=nx, Ly=Lx, nz=nx, Lz=Lx; x0=-Lx/2, y0=-Ly/2, z0=-Lz/2,
    nthreads=Sys.CPU_THREADS, effort=FFTW.MEASURE, T=Float64, dealias=1/3, ArrayType=Array)

 Constructs a TwoDGrid object with size `Lx`, `Ly`, `Lz`, resolution `nx`, `ny`,
 `nz` and leftmost positions `x0`, `y0`, `z0`. FFT plans are generated for `nthreads`
 CPUs using FFTW flag `effort`. The float type is `T` and the array types is `ArrayType`.
 The `dealias` keyword determines the highest wavenubers that are being zero-ed
 out by `dealias()` function; 1/3 is the nominal value for quadratic nonlinearities.
"""
function ThreeDGrid(nx, Lx, ny=nx, Ly=Lx, nz=nx, Lz=Lx; x0=-Lx/2, y0=-Ly/2, z0=-Lz/2,
                  nthreads=Sys.CPU_THREADS, effort=FFTW.MEASURE, T=Float64, dealias=1/3, ArrayType=Array)

  dx = Lx/nx
  dy = Ly/ny
  dz = Lz/nz

  nk = nx
  nl = ny
  nm = nz
  nkr = Int(nx/2 + 1)

  # Physical grid
  x = range(T(x0), step=T(dx), length=nx)
  y = range(T(y0), step=T(dy), length=ny)
  z = range(T(z0), step=T(dz), length=nz)

  # Wavenubmer grid
   k = ArrayType{T}(reshape( fftfreq(nx, 2π/Lx*nx), (nk, 1, 1)))
   l = ArrayType{T}(reshape( fftfreq(ny, 2π/Ly*ny), (1, nl, 1)))
   m = ArrayType{T}(reshape( fftfreq(nz, 2π/Lz*nz), (1, 1, nm)))
  kr = ArrayType{T}(reshape(rfftfreq(nx, 2π/Lx*nx), (nkr, 1, 1)))

     Ksq = @. k^2 + l^2 + m^2
  invKsq = @. 1 / Ksq
  CUDA.@allowscalar invKsq[1, 1, 1] = 0

     Krsq = @. kr^2 + l^2 + m^2
  invKrsq = @. 1 / Krsq
  CUDA.@allowscalar invKrsq[1, 1, 1] = 0

  # FFT plans
  FFTW.set_num_threads(nthreads)
  fftplan = plan_flows_fft(ArrayType{Complex{T}, 3}(undef, nx, ny, nz), flags=effort)
  rfftplan = plan_flows_rfft(ArrayType{T, 3}(undef, nx, ny, nz), flags=effort)

  # Index endpoints for aliasfrac i, j wavenumbers
  kalias, kralias = getaliasedwavenumbers(nk, nkr, dealias)
  lalias, malias = getaliasedwavenumbers(nl, nm, dealias)

  Tx = typeof(x)
  Tk = typeof(k)
  Tfft = typeof(fftplan)
  Trfft = typeof(rfftplan)

  return ThreeDGrid{T, Tk, Tx, Tfft, Trfft}(nx, ny, nz, nk, nl, nm, nkr, dx, dy, dz, Lx, Ly, Lz,
                                        x, y, z, k, l, m, kr, Ksq, invKsq, Krsq, invKrsq, fftplan, rfftplan,
                                        kalias, kralias, lalias, malias)
end

Base.eltype(grid::OneDGrid) = eltype(grid.x)
Base.eltype(grid::TwoDGrid) = eltype(grid.x)
Base.eltype(grid::ThreeDGrid) = eltype(grid.x)

OneDGrid(dev::CPU, args...; kwargs...) = OneDGrid(args...; ArrayType=Array, kwargs...)
TwoDGrid(dev::CPU, args...; kwargs...) = TwoDGrid(args...; ArrayType=Array, kwargs...)
ThreeDGrid(dev::CPU, args...; kwargs...) = ThreeDGrid(args...; ArrayType=Array, kwargs...)

"""
    gridpoints(grid::TwoDGrid)
    gridpoints(grid::ThreeDGrid)

Returns the collocation points of the `grid` in 2D or 3D arrays `X, Y` (and `Z`).
"""
function gridpoints(grid::TwoDGrid{T, A}) where {T, A}
  X = [ grid.x[i₁] for i₁=1:grid.nx, i₂=1:grid.ny ]
  Y = [ grid.y[i₂] for i₁=1:grid.nx, i₂=1:grid.ny ]

  return A(X), A(Y)
end

function gridpoints(grid::ThreeDGrid{T, A}) where {T, A}
  X = [ grid.x[i₁] for i₁=1:grid.nx, i₂=1:grid.ny, i₃=1:grid.nz ]
  Y = [ grid.y[i₂] for i₁=1:grid.nx, i₂=1:grid.ny, i₃=1:grid.nz ]
  Z = [ grid.z[i₃] for i₁=1:grid.nx, i₂=1:grid.ny, i₃=1:grid.nz ]

  return A(X), A(Y), A(Z)
end

"""
    getaliasedwavenumbers(nk, nkr, aliasfraction)

Returns the top `aliasfraction` highest wavenumbers.
"""
function getaliasedwavenumbers(nk, nkr, aliasfraction)
  # Index endpoints for aliased i, j wavenumbers
  # 1/3 aliasfraction => upper 1/6 of +/- wavenumbers (1/3 total) are set to 0 after performing fft.
  # 1/2 aliasfraction => upper 1/4 of +/- wavenumbers (1/2 total) are set to 0 after performing fft.
  L = (1 - aliasfraction)/2 # (1 - 1/3) / 2 + 1 = 1/3.
  R = (1 + aliasfraction)/2 # (1 + 1/3) / 2 - 1 = 2/3.
  iL = floor(Int, L * nk) + 1
  iR =  ceil(Int, R * nk)

  aliasfraction < 1 || error("`aliasfraction` must be less than 1") # aliasfraction=1 is not sensible.
   kalias = (aliasfraction > 0) ? (iL:iR) : Int(nx/2 + 1)
  kralias = (aliasfraction > 0) ? (iL:nkr) : nkr

  return kalias, kralias
end

"""
    dealias!(a, g, kalias)

Dealias array `a` on the grid `g` with aliased x-wavenumbers `kalias`.
"""
function dealias!(a, g::OneDGrid)
  kalias = size(a, 1) == g.nkr ? g.kralias : g.kalias
  dealias!(a, g, kalias)
  return nothing
end

function dealias!(a, g::OneDGrid, kalias)
  @views @. a[kalias, :] = 0

  return nothing
end

function dealias!(a, g::TwoDGrid)
  kalias = size(a, 1) == g.nkr ? g.kralias : g.kalias
  dealias!(a, g, kalias)

  return nothing
end

function dealias!(a, g::TwoDGrid, kalias)
  @views @. a[kalias, g.lalias, :] = 0

  return nothing
end

function dealias!(a, g::ThreeDGrid)
  kalias = size(a, 1) == g.nkr ? g.kralias : g.kalias
  dealias!(a, g, kalias)

  return nothing
end

function dealias!(a, g::ThreeDGrid, kalias)
  @views @. a[kalias, g.lalias, g.malias, :] = 0

  return nothing
end

"""
    makefilter(K; order=4, innerK=0.65, outerK=1)

Returns a filter acting on the non-dimensional wavenumber K that decays exponentially
for K>innerK, thus removing high-wavenumber content from a spectrum it is multiplied
with. The decay rate is determined by order and outerK determines the outer wavenumber
at which the filter is smaller than Float64 machine precision.
"""
function makefilter(K::Array; order=4, innerK=0.65, outerK=1)
  TK = typeof(K)
  K = Array(K)
  decay = 15*log(10) / (outerK-innerK)^order # decay rate for filtering function
  filt = @. exp( -decay*(K-innerK)^order )
  filt[real.(K) .< innerK] .= 1

  return TK(filt)
end

function makefilter(g::OneDGrid; realvars=true, kwargs...)
  K = realvars ? g.kr*g.dx/π : @.(abs(g.k*g.dx/π))

  return makefilter(K; kwargs...)
end

function makefilter(g::TwoDGrid; realvars=true, kwargs...)
  K = realvars ?
      @.(sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2)) : @.(sqrt((g.k*g.dx/π)^2 + (g.l*g.dy/π)^2))

  return makefilter(K; kwargs...)
end

function makefilter(g::ThreeDGrid; realvars=true, kwargs...)
  K = realvars ?
      @.(sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2 + (g.m*g.dz/π)^2)) : @.(sqrt((g.k*g.dx/π)^2 + (g.l*g.dy/π)^2 + (g.m*g.dz/π)^2))

  return makefilter(K; kwargs...)
end

makefilter(g, T, sz; kwargs...) = ones(T, sz) .* makefilter(g; realvars=sz[1]==g.nkr, kwargs...)
makefilter(eq; kwargs...) = makefilter(eq.grid, fltype(eq.T), eq.dims; kwargs...)

"""
    griddevice(grid)

Returns the device on which the `grid` lives on.
"""
griddevice(grid::AbstractGrid{T, A}) where {T, A} = A<:Array ? "CPU" : "GPU"

show(io::IO, g::OneDGrid{T}) where T =
     print(io, "OneDimensionalGrid\n",
               "  ├─────────── Device: ", griddevice(g), '\n',
               "  ├──────── FloatType: $T", '\n',
               "  ├────────── size Lx: ", g.Lx, '\n',
               "  ├──── resolution nx: ", g.nx, '\n',
               "  ├── grid spacing dx: ", g.dx, '\n',
               "  └─────────── domain: x ∈ [$(g.x[1]), $(g.x[end])]")

show(io::IO, g::TwoDGrid{T}) where T =
     print(io, "TwoDimensionalGrid\n",
               "  ├───────────────── Device: ", griddevice(g), '\n',
               "  ├────────────── FloatType: $T", '\n',
               "  ├────────── size (Lx, Ly): ", (g.Lx, g.Ly), '\n',
               "  ├──── resolution (nx, ny): ", (g.nx, g.ny), '\n',
               "  ├── grid spacing (dx, dy): ", (g.dx, g.dy), '\n',
               "  └───────────────── domain: x ∈ [$(g.x[1]), $(g.x[end])]", '\n',
               "                             y ∈ [$(g.y[1]), $(g.y[end])]")

show(io::IO, g::ThreeDGrid{T}) where T =
     print(io, "ThreeDimensionalGrid\n",
               "  ├───────────────────── Device: ", griddevice(g), '\n',
               "  ├────────────────── FloatType: $T", '\n',
               "  ├────────── size (Lx, Ly, Lz): ", (g.Lx, g.Ly, g.Ly), '\n',
               "  ├──── resolution (nx, ny, nz): ", (g.nx, g.ny, g.nz), '\n',
               "  ├── grid spacing (dx, dy, dz): ", (g.dx, g.dy, g.dz), '\n',
               "  └────────────────────  domain: x ∈ [$(g.x[1]), $(g.x[end])]", '\n',
               "                                 y ∈ [$(g.y[1]), $(g.y[end])]", '\n',
               "                                 z ∈ [$(g.z[1]), $(g.z[end])]")


# Return log density for population of states
function eval_population( theta::Array{Float64,2}, calc_log_pdf::Function )
   popsize = size(theta,2)
   logpdf = Array{Float64}( popsize )
   for i in 1:popsize
      logpdf[i] = calc_log_pdf(theta[:,i])
   end
   return logpdf
end

## Differential Evolution Sampler
#    based on ter Braak (2006) and Nelson et al. (2014)

# Generate ensemble of trial states according to DEMCMC
function generate_trial_states_demcmc( theta::Array{Float64,2}; gamma_o::Float64 = 2.38/sqrt(2*size(theta,1)), epsilon = 0.01  )
   num_param = size(theta,1)
   num_pop = size(theta,2)
   theta_trial = similar(theta)

   for i in 1:num_pop
      # Choose head (k) and tail (j) for perturbation vector
  j = rand(1:(num_pop-1))
  if( j >= i ) j = j+1 end
  k = rand(1:(num_pop-2))
  if( k >= i ) k = k+1 end
  if( k >= j ) k = k+1 end

  # Choose scale factor
  scale = (1.0 + epsilon*randn()) * gamma_o

  theta_trial[:,i] = theta[:,i] + scale * (theta[:,k] - theta[:,j])
   end
   return theta_trial
end

# Evolve population of states (theta_init) with target density
function run_demcmc( theta_init::Array{Float64,2}, calc_log_pdf::Function; num_gen::Integer = 100, epsilon::Float64 = 0.01 )
   @assert(num_gen >=1 )
   num_param = size(theta_init,1)
   num_pop = size(theta_init,2)
   @assert(num_pop > num_param)

   # Allocate arrays before starting loop
   pop = Array{Float64}(num_param, num_pop, num_gen)
   poplogpdf = Array{Float64}(num_pop, num_gen )
   accepts_chain = zeros(Int64,num_pop )
   rejects_chain = zeros(Int64,num_pop )
   accepts_generation = zeros(Int64,num_gen )
   rejects_generation = zeros(Int64,num_gen )

   # Make a working copy of the current population of parameter values
   theta = copy(theta_init)
   logpdf = eval_population( theta, calc_log_pdf )

   for g in 1:num_gen # first(pop.range):last(pop.range)  # loop over generations
      # Generate population of trial sets of model parameters
      gamma_o = (mod(g,10)!=0) ? 2.38/sqrt(2*num_param) : 1.0 # every 10th generation try full-size steps
      theta_trial = generate_trial_states_demcmc( theta, gamma_o=gamma_o, epsilon=epsilon )

      # Evaluate model for each set of trial parameters
      logpdf_trial = eval_population( theta_trial, calc_log_pdf )

      # For each member of population
      for i in 1:num_pop
        log_pdf_ratio = logpdf_trial[i] - logpdf[i]
        if( (log_pdf_ratio>0) || (log_pdf_ratio>log(rand())) )     # Decide whether to Accept
            theta[:,i] = theta_trial[:,i]
            logpdf[i] = logpdf_trial[i]
            accepts_chain[i] = accepts_chain[i]+1
            accepts_generation[g] = accepts_generation[g]+1
        else
            rejects_chain[i] = rejects_chain[i]+1
            rejects_generation[g] = rejects_generation[g]+1
        end
      end

      # Log results
      if(true)  # (in(g,pop.range))
        pop[:,:,g] = copy(theta)
        poplogpdf[:,g] = copy(logpdf)
      end
   end
   return Dict("theta_last"=>theta, "logpdf_last"=>logpdf, "theta_all"=>pop, "logpdf_all"=>poplogpdf,
           "accepts_chain"=>accepts_chain, "rejects_chain"=>rejects_chain, "accepts_generation"=>accepts_generation, "rejects_generation"=>rejects_generation )
end


"""
    struct Affine
        w
        b
        σ
    end

Affine layer of neural network

"""
struct Affine{T1<:AbstractArray,T2}
    w::T1
    b::T2
    σ::Function
end

function Affine(
    in::Integer,
    out::Integer,
    σ = identity::Function;
    fw = randn::Function,
    fb = zeros::Function,
    isBias = true::Bool,
    precision = Float32,
)
    if isBias
        return Affine(fw(precision, out, in), fb(precision, out), σ)
    else
        return Affine(fw(precision, out, in), Flux.Zeros(precision, out), σ)
    end
end

Flux.@functor Affine # works with Flux.params

(L::Affine)(x::AbstractVector) = L.w * x .+ L.b

"""
    struct FastAffine{I,F,F2} <: DiffEqFlux.FastLayer
        out::I
        in::I
        σ::F
        initial_params::F2
    end
Equivalent FastDense layer with controllable type
"""
struct FastAffine{I,F,F2} <: DiffEqFlux.FastLayer
    out::I
    in::I
    σ::F
    initial_params::F2
end

function FastAffine(
    in::Integer,
    out::Integer,
    σ = identity;
    fw = randn,
    fb = Flux.zeros,
    precision = Float32,
)
    initial_params() = vcat(vec(fw(precision, out, in)), fb(precision, out))
    return FastAffine(out, in, σ, initial_params)
end

(f::FastAffine)(x, p) =
    f.σ.(reshape(p[1:(f.out*f.in)], f.out, f.in) * x .+ p[(f.out*f.in+1):end])

DiffEqFlux.paramlength(f::FastAffine) = f.out * (f.in + 1)
DiffEqFlux.initial_params(f::FastAffine) = f.initial_params()


using DifferentialEquations, ModelingToolkit
using HoJMTK.Electrical
using Plots

# Based on https://mtk.sciml.ai/dev/tutorials/acausal_components/
R = 1.0
C = 1.0
V = 1.0
@named resistor = Resistor(R=R)
@named capacitor = Capacitor(C=C)
@named source = DCVoltageSource(V=V)
@named ground = Ground()

rc_eqs = [
          connect(source.p, resistor.p)
          connect(resistor.n, capacitor.p)
          connect(capacitor.n, source.n, ground.g)
         ]

@variables t
@named _rc_model = ODESystem(rc_eqs, t)
@named rc_model = compose(_rc_model, [resistor, capacitor, source, ground])
sys = structural_simplify(rc_model)
u0 = [
      capacitor.v => 0.0
     ]
prob = ODAEProblem(sys, u0, (0, 10.0))
sol = solve(prob, Tsit5())

plot(sol)


function solvenjl(κ, coef, energ, j)

    @info "Solving the NJL equations..."

    constant = 4*Nc*Nf*G*ħ^3
    function f!(F, x)
        for k in 1:κ
            F[k] = constant*sum([coef[k, kprime]*(2*j[kprime]+1)*x[kprime]/√(x[kprime]^2+energ[kprime]^2) for kprime in 1:κ]) - x[k]
        end
    end

    # The quark masses are initialized as having 350 [MeV]
    initial_x = 350.0*ones(κ)
    result = nlsolve(f!, initial_x)
    mass = result.zero

    @info "Done!"
    @debug "Masses: " mass

    return mass

end