module LMM_module

using Statistics
using LinearAlgebra
using Distributions
using Optim
using RCall

#####################
### Sub-functions ###
#####################

"""
# _______________________________________
# Polyalgorithmic matrix inverse function

`inverse(A::Array{Float64,2})`

# Inputs
1. *A* [Array{Float64,2}]: square or rectangular matrix

# Output
1. Inverse or pseudo inverse of A

# Examples
```
S = convert(Array{Float64,2}, reshape(collect(1:9), 3, 3))
LMM_module.inverse(S)
R = convert(Array{Float64,2}, reshape(collect(1:10), 5, 2))
LMM_module.inverse(R)
```
"""
function inverse(A::Array{Float64,2})
	out = try
			LinearAlgebra.inv(A)
		catch
			LinearAlgebra.pinv(A)
		end
	return(out)
end

"""
# ________________________________________________________________
# Linear mixed model parameter optimization and effects estimation

`LMM_optim(parameters::Array{Float64,1}; X::Array{Float64,2}, y::Array{Float64,1}, Z::Array{Float64,2}, METHOD_VAR_EST::String=["ML", "REML"][1], OPTIMIZE::Bool=true, VARFIXEF::Bool=false)`

The mixed linear model is defined as y = Xb + Zu + e, where:
	- X [n,p] is the centered matrix of allele frequencies
	- Z [n,n] is the square symmetric matrix of relatedness
	- y [n,1] is the centered vector of phenotypic values
	- no intercept is explicitly fitted but implicitly set at the mean phenotypic value as a consequence of centering y
	- u ~ N(0, σ²uI)
	- e ~ N(0, σ²eI)
	- y ~ N(0, V); V = (Z (σ²uI) Z') + (σ²eI)
	- variance component (σ²e, σ²u) are estimated via maximum likelihood (ML) or restricted maximum likelihood (REML)
	- fixed effects (b) are estimated by solving: (X' * inverse(V)) * inverse(X * X' * inverse(V)) * y
	- random effects (u) are estimated by solving: (σ²uI) * Z' * inverse(V) * (y - (X*b))

# Inputs
1. parameters [Array{Float64,1}]: the error variance (σ2e) and random effects variance (σ2u)
2. *X* [Array{Float64,2}; xij ∈ [0.00, 1.00]]: matrix of allele frequencies (Pool-seq data) set as the fixed effects
3. *y* [Array{Float64,1}]: vector of phenotypes
4. *Z* [Array{Any,2}]: matrix of relatedness (Fst or XX'/p where p is the number of columns in X) set as the random effects
5. *METHOD_VAR_EST* [String]: method of variance parameter estimation choose "ML" for maximum likelihood or "REML" for restricted maximum likelihood (default="ML")
6. *OPTIMIZE* [Bool]: logical value to perform optimization to find estimate the variance parameters if TRUE or to estimate the fixed and random effects if FALSE (default=true)
7. *VARFIXEF* [Bool]: logical value to estimate the variances of the fixed effects which can be computationally and memory intensive for large datasets (default=false)

# Outputs
1. *If OPTIMIZE==true*: Negative log-likelihood of the parameters (σ2e and σ2u) given the dataset (X, y, and Z)
2. *If OPTIMIZE==false*: Fixed effects (b_hat), random effects (u_hat), and the optional fixed effects variances (σ2b_hat)

# Examples
```
using Distributions
using LinearAlgebra
using Optim
n=5; l=200; Ve=1.00; Vu=1.00;
X = reshape(abs.(rand(n*l)), n, l)
Z = reshape(abs.(rand(n^2)), n, n)
b = zeros(l); b[[5,10,100]] = [1.00, 1.00, 1.00]
u = rand(Distributions.Normal(0.00, Vu), n)
e = rand(Distributions.Normal(0.00, Ve), n)
y = (X*b) .+ (Z*u) .+ e
par = Optim.optimize(par -> LMM_module.LMM_optim(par, X=X, y=y, Z=Z), [0.0,0.0], [Inf,Inf], [1.00,1.00]).minimizer
b_hat, u_hat, Vb_hat = LMM_module.LMM_optim(par, X=X, y=y, Z=Z, OPTIMIZE=false)
```
"""
function LMM_optim(parameters::Array{Float64,1}; X::Array{Float64,2}, y::Array{Float64,1}, Z::Array{Float64,2}, METHOD_VAR_EST::String=["ML", "REML"][1], OPTIMIZE::Bool=true, VARFIXEF::Bool=false)
	### y = Xb + Zu + e; where Xb=fixed and Za=random
	σ2e = parameters[1] # variance of the error effects (assuming homoscedasticity)
	σ2u = parameters[2] # variance of the other random effects (assuming homoscedasticity)
	n = length(y)		# number of individual samples
	l = size(X, 2)		# number of fixed effects
	q = size(Z, 2)		# number of random effects excluding error effect
	### Variance components
	R = σ2e * I					# error variance-covariance matrix
	D = σ2u * I					# random effects variance-covariance matrix
	V = (Z * D * Z') + R		# y variance-covariance matrix
	inv_V = inverse(V)
	### Mixed model equations
	###### Fixed effects:
	b_hat = (X' * inv_V) * inverse(X * X' * inv_V) * y
	###### Random effects:
	u_hat = D * Z' * inv_V * (y - (X*b_hat))
	### Calculate the log-likelihood of the parameters (variances of the error and random effects: σ2e and σ2u) given the data
	### Or estimate the fixed and random effects as well as the fixed effects variance given the optimum parameters
	if OPTIMIZE
		################
		### OPTIMIZE ###
		################
		### Calculate the -log-likelihood given the error variance (σ2e) and random effects variance (σ2u) for maximum likelihood optimization to find the best fit σ2e and σ2u
		OUT = []
		if METHOD_VAR_EST == "ML"
			mu_y = (X * b_hat)
			n = length(y)
			push!(OUT, 0.5 * ( log(abs(det(V))) + ((y - mu_y)' * inv_V * (y - mu_y)) + (n*log(2*pi)) )) ### The negative log-likelihood function y given σ2e and σ2u
		elseif METHOD_VAR_EST == "REML"
			### NOTE: Z MUST BE SQUARE!
			M = try
					inverse(LinearAlgebra.cholesky(σ2u*Z + R)) ### Cholesky decomposition
				catch
					inverse(LinearAlgebra.lu(σ2u*Z + R).L) ### LU decomposition
				end
			y_new = M' * y
			intercept_new = sum(M', dims=2) #M' * ones(n)
			V_new = M' * V * M
			n = length(y_new)
			push!(OUT, 0.5 * ( log(abs(det(V_new))) .+ ((y_new - intercept_new)' * inverse(V_new) * (y_new - intercept_new)) .+ (n*log(2*pi)) )[1,1]) ### negative log-likelihood of σ2e and σ2u given X, y, and Z
		else
			println(string("Sorry. ", METHOD_VAR_EST, " is not a valid method of estimating the variances of the random effects effects. Please pick ML or REML."))
			exit()
		end
	else
		################
		### ESTIMATE ###
		################
		### Return the fixed effects (b_hat), random effects (u_hat), and optional fixed effects variance (σ2b_hat) estimates
		OUT = []
		if VARFIXEF
			### Computationally intensive!!!
			Vb_hat = inverse(X' * inverse(V) * X)
			σ2b_hat = diag(Vb_hat)
		else
			σ2b_hat = nothing
		end
		# push!(OUT, (fixef=b_hat, ranef=u_hat, var_fixef=σ2b_hat))
		push!(OUT, [b_hat, u_hat, σ2b_hat])
	end
	return(OUT[1])
end

#####################
### main function ###
#####################
"""
# ______________________
# Linear Mixed Modelling

`LMM(;X::Array{Float64,2}, y::Array{Float64,1}, Z::Array{Float64,2}, METHOD_VAR_EST::String=["ML", "REML"][1])`

The mixed linear model is defined as y = Xb + Zu + e, where:
	- no intercept is fitted
	- X is centered
	- y is centered which means the intercept is automatically the mean of y
	- Z is a square symmetric matrix of relatedness
	- u ~ N(0, σ2uI)
	- e ~ N(0, σ2eI)
	- y ~ (Z (σ2uI) Z') + (σ2eI)

# Inputs
1. *X* [Array{Float64,2}; xij ∈ [0.00, 1.00]]: matrix of allele frequencies (Pool-seq data) set as the fixed effects
2. *y* [Array{Float64,1}]: vector of phenotypes
3. *Z* [Array{Any,2}]: matrix of relatedness (Fst or XX'/p where p is the number of columns in X) set as the random effects
4. *METHOD_VAR_EST* [String]: method of variance parameter estimation choose "ML" for maximum likelihood or "REML" for restricted maximum likelihood (default="ML")

# Outputs
1. Intercept (b0) the mean of the phenotype data
2. Fixed effects estimates of X (b_hat)
3. Random effects estimates of Z (u_hat)

# Examples
```
using Distributions
using LinearAlgebra
n=5; l=200; Ve=1.00; Vu=1.00;
X = reshape(abs.(rand(n*l)), n, l)
Z = reshape(abs.(rand(n^2)), n, n)
b = zeros(l); b[[5,10,100]] = [1.00, 1.00, 1.00]
u = rand(Distributions.Normal(0.00, Vu), n)
e = rand(Distributions.Normal(0.00, Ve), n)
y = (X*b) .+ (Z*u) .+ e
b0, b_hat, u_hat = LMM_module.LMM(X=X, y=y, Z=Z)
b0, b_hat, u_hat = LMM_module.LMM(X=X, y=y, Z=Z, METHOD_VAR_EST="REML")
```
"""
function LMM(;X::Array{Float64,2}, y::Array{Float64,1}, Z::Array{Float64,2}, METHOD_VAR_EST::String=["ML", "REML"][1])
	### MODEL: y = Xb + Zu + e, where: u~N(0,Vu) and e~(0,Ve)
	### Generate copies of the input vector and matrices so that we don't modify the original ones
	X = copy(X); y = copy(y); Z = copy(Z);
	### Center y vector, X matrix and Z matrix (This way we don't need to estimate the intercept! It will be at zero!)
	uy = mean(y); uX = mean(X); uZ = mean(Z);
	sy = std(y); sX = std(X); sZ = std(Z);
	y[:] = (y .- uy)
	X[:,:] = (X .- uX)
	Z[:,:] = (Z .- uZ)
	### Maximum likelihood estimation of the error variance (Ve) and random covariate variance (Vu)
	# lower_limit = [0.0, 0.0]
	# upper_limit = [Inf, Inf]
	lower_limit = [1.0e-50, 1.0e-50]
	upper_limit = [1.0e+50, 1.0e+50]
	initial_VeVu = [1.0, 1.0]
	# optim_out = Optim.optimize(VeVu -> LMM_optim(VeVu, X=X, y=y, Z=Z, METHOD_VAR_EST=METHOD_VAR_EST, METHOD_FIX_EST=METHOD_FIX_EST, alfa=alfa, lambda=lambda, OPTIMIZE=true), lower_limit, upper_limit, initial_VeVu)
	### Optimization in Julia 1.0 is too darn slow! Going with R then! Importing Julia functions and variables into R.
	@rput inverse;
	@rput LMM_optim;
	@rput X;
	@rput y;
	@rput Z;
	@rput METHOD_VAR_EST;
	@rput lower_limit;
	@rput upper_limit;
	@rput initial_VeVu;
	R"optim_out = optim(par=initial_VeVu, fn=LMM_optim, X=X, y=y, Z=Z, METHOD_VAR_EST=METHOD_VAR_EST, OPTIMIZE=TRUE, method='L-BFGS-B', lower=lower_limit, upper=upper_limit)"
	@rget optim_out;
	Ve, Vu = optim_out[:par]
	### Fixed effects (b), random effects (u), and fixed effects variance (Vb) estimation (NOTE: No need to back-transform these estimated effects because we simply centered the explanatory variables meaning they're still on the same space)
	b, u, Vb = LMM_optim([Ve, Vu], X=X, y=y, Z=Z, METHOD_VAR_EST=METHOD_VAR_EST, OPTIMIZE=false)
	### Output
	return (b0=uy, b=b, u=u, Ve=Ve, Vu=Vu)
end

end #end of LMM module


using CircGeometry, Test

between_buffer = 5.0
vf = 0.4
n_bodies = 50

# test circle outline
material = MaterialParameters(vf,n_bodies)
outline = OutlineCircle(1.0,Point(0.0,0.0))
ps = generate_porous_structure(outline,material,between_buffer;log=false)
vf_exp = compute_volume_fraction(ps,outline)

CircGeometry.copyArraysToCenters!(ps)

# test rectangle outline
outline = OutlineRectangle(Point(-2.0,0.0),Point(1.0,2.0))
@test CircGeometry.compute_outline_area(outline) == 6.0
ps = generate_porous_structure(outline,material,between_buffer;log=false)
vf_exp = compute_volume_fraction(ps,outline)

# test polygon outline
outline = OutlinePolygon([Point(-1.0,0.0),Point(1.0,0.0),
            Point(1.0,2.0),Point(0.0,1.0),Point(-1.0,2.0)])
@test CircGeometry.compute_outline_area(outline) == 3.0
ps = generate_porous_structure(outline,material,between_buffer;log=false)
vf_exp = compute_volume_fraction(ps,outline)


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMQuad.jl/blob/master/LICENSE

### Gauss quadrature rules for tetrahedrons

""" Gauss-Legendre quadrature, 1 point rule on tetrahedron. """
function get_quadrature_points(::Type{Val{:GLTET1}})
    weights = (1.0/6.0, )
    points = ((1.0/4.0, 1.0/4.0, 1.0/4.0), )
    return zip(weights, points)
end

function get_order(::Type{Val{:GLTET1}})
    return 1
end

""" Gauss-Legendre quadrature, 4 point rule on tetrahedron. """
function get_quadrature_points(::Type{Val{:GLTET4}})
    a = (5.0+3.0*sqrt(5.0))/20.0
    b = (5.0-sqrt(5.0))/20.0
    w = 1.0/24.0
    weights = (w, w, w, w)
    points = ((a, b, b), (b, a, b), (b, b, a), (b, b, b))
    return zip(weights, points)
end

function get_order(::Type{Val{:GLTET4}})
    return 2
end

""" Gauss-Legendre quadrature, 5 point rule on tetrahedron. """
function get_quadrature_points(::Type{Val{:GLTET5}})
    a = 1.0/4.0
    b = 1.0/6.0
    c = 1.0/2.0
    weights = (-2.0/15.0, 3.0/40.0, 3.0/40.0, 3.0/40.0, 3.0/40.0)
    points = ((a, a, a), (b, b, b), (b, b, c), (b, c, b), (c, b, b))
    return zip(weights, points)
end

function get_order(::Type{Val{:GLTET5}})
    return 3
end

""" Gauss-Legendre quadrature, 15 point rule on tetrahedron. """
function get_quadrature_points(::Type{Val{:GLTET15}})
    a = 1.0/4.0
    b1 = 1.0/34.0*(7.0 + sqrt(15.0))
    b2 = 1.0/34.0*(7.0 - sqrt(15.0))
    c1 = 1.0/34.0*(13.0 - 3.0*sqrt(15.0))
    c2 = 1.0/34.0*(13.0 + 3.0*sqrt(15.0))
    d = 1.0/20.0*(5.0 - sqrt(15.0))
    f = 1.0/20.0*(5.0 + sqrt(15.0))
    w1 = 8.0/405.0
    w2 = (2665.0 - 14.0*sqrt(15.0))/226800.0
    w3 = (2665.0 + 14.0*sqrt(15.0))/226800.0
    w4 = 5.0/567.0
    weights = (w1, w2, w2, w2, w2, w3, w3, w3, w3, w4, w4, w4, w4, w4, w4)
    points = ((a, a, a), (b1, b1, b1), (b1, b1, c1), (b1, c1, b1), (c1, b1, b1),
              (b2, b2, b2), (b2, b2, c2), (b2, c2, b2), (c2, b2, b2), (d, d, f),
              (d, f, d), (f, d, d), (d, f, f), (f, d, f), (f, f, d))
    return zip(weights, points)
end

function get_order(::Type{Val{:GLTET15}})
    return 4
end


## sample random effects given latent class membership, observed data and parameters
function sample_rand(g::Int, data::Data, param::Param)
    @unpack beta, mu, S, L, sigma, b = param
    @unpack Y, times, n, Xl, Z = data

    N, q = size(b)
    z = Matrix{Float64}(undef, N, q)
    invSigma = inv_cov_SL(S[g,:], L[g,:,:])

    tt = 1:n[1]
    Zi = Z[1,tt,:]
    ll_mean = Zi'*(Y[1,tt] - Xl[1,tt,:]*beta[g,:])
    ll_inv_cov = Zi'*Zi / sigma^2
    post_cov = inv(Symmetric(invSigma + ll_inv_cov))
    post_mean = post_cov*(invSigma*mu[g,:] + ll_mean / sigma^2)

    rand(MvNormal(post_mean, post_cov))
end

f1(t, eta=-1.5) = (1 .+ t).^eta .- 1
f2(t, nu=0.0) = t.^(nu+1) ./ (t.+1).^nu

## sample longitudinal trjactory given latent class membership, random effects, covariates and parameters
function sample_lmm(time::AbstractVector{T}, g::Int, b::Vector{T}, data::Data, param::Param) where T <: Real
    @unpack beta, sigma = param
    @unpack Xl, n = data

    X = Xl[1,1,:]
    y = zeros(T, length(time))

    for j in 1:length(time)
        Z = [1, f1(time[j]), f2(time[j])]
        Xl = vcat(X[1:3], X[4:6]*f1(time[j]), X[7:9]*f2(time[j]))
        mean_y = dot(Xl, beta[g,:]) + dot(Z, b)
        y[j] = mean_y ##rand(Normal(mean_y, sigma))
    end
    return y
end

surv_g(time, shape, scale, delta, Xe) = exp.(-exp(dot(Xe, delta))*weibull_cumhaz(time, shape, scale))

## prediction of conditional survival function and longitudinal trajectory conditional on new_data
## new_data is interpreted as data from a new subject if out_of_sample=true or as additional data from an existing
## subject in data if out_of_sample=false (for within sample prediction new data can be same as old data)
## for new subjects latent class indicator and random effects are sampled
## for existing subjects the existing posterior samples of latent class indicator and random effects are used
function predict(time::AbstractVector{T}, new_data::Data, x::Tf; out_of_sample=true) where {T <: Real, Tf <: Fit}
    param = get_samples(x)

    if length(new_data.T) > 1
        @warn "Only using first subject in new_data"
    end

    t0 = new_data.T[1]
    @assert minimum(time) >= t0

    obs_t = new_data.times[1, 1:new_data.n[1]]
    obs_y = new_data.Y[1,1:new_data.n[1]]

    if t0 > maximum(obs_t)
        time = vcat(t0, sort(time))
    end

    M = length(param)
    b = Matrix{T}(undef, M, size(param[1].b,2))
    y = Matrix{T}(undef, M, length(t))
    S = Matrix{T}(undef, M, length(t))
    g = Vector{Int}(undef, M)

    for m in 1:M
        theta = param[m]

        if out_of_sample
            g[m] = sample_cind(theta, new_data)[1] ##sample_class(new_data, theta)
            b[m,:] = sample_rand(g[m], new_data, theta)
        else
            g[m] = theta.cind[1]
            b[m,:] = theta.b[1,:]
        end

        y[m,:] = sample_lmm(time, g[m], b[m,:], new_data, theta)
        S0 = surv_g(new_data.T[1], theta.bl_shape[g[m]], theta.bl_scale[g[m]], theta.delta[g[m],:], new_data.Xe[1,:])
        S[m,:] = min.(surv_g(time, theta.bl_shape[g[m]], theta.bl_scale[g[m]], theta.delta[g[m],:], new_data.Xe[1,:]) ./ S0, 1.0)
    end

    summary_predict(obs_t, obs_y, time, y, S)
end

function summary_predict(obs_t, obs_y, t, y, S)
    S0 = repeat([1.0], inner=length(obs_t))

    S_ci_lo = vcat(S0, [quantile(S[:,i], 0.025) for i in 1:length(t)])
    S_ci_hi = vcat(S0, [quantile(S[:,i], 0.975) for i in 1:length(t)])

    y_ci_lo = vcat(obs_y, [quantile(y[:,i], 0.025) for i in 1:length(t)])
    y_ci_hi = vcat(obs_y, [quantile(y[:,i], 0.975) for i in 1:length(t)])

    DataFrame(time=vcat(obs_t, t), type=vcat(repeat(["observed"], inner=length(obs_t)), repeat(["predicted"], inner=length(t))),
              y=vcat(obs_y, mean(y, dims=1)[1,:]), sd_y=vcat(zeros(Float64, length(obs_t)), sqrt.(var(y, dims=1))[1,:]), ci_lo_y=y_ci_lo, ci_hi_y=y_ci_hi,
              S=vcat(S0, mean(S, dims=1)[1,:]), sd_S=vcat(zeros(Float64, length(obs_t)), sqrt.(var(S, dims=1)[1,:])), ci_lo_S=S_ci_lo, ci_hi_S=S_ci_hi)
end

function plot_predict(df::DataFrame; use_R=true)
    if use_R
        R"library(ggplot2)"
        @rput df
        R"ggplot(df, aes(x=df$time, y=df$y, ymin=df$ci_lo_y, ymax=df$ci_hi_y)) + geom_line() + geom_ribbon()"
        R"ggplot(df, aes(x=df$time, y=df$S, ymin=df$ci_lo_S, ymax=df$ci_hi_S)) + geom_line() + geom_ribbon()"
    else
        df_obs = df[df.type .== "observed",:]
        df_pred = df[df.type .== "predicted",:]

        p1 = plot(layer(x=df_pred.time, y=df_pred.y, ymin=df_pred.ci_lo_y, ymax=df_pred.ci_hi_y, Geom.line, Geom.ribbon),
                  layer(x=df_obs.time, y=df_obs.y, Geom.point), Guide.xlabel("Time"), Guide.ylabel("log(PSA)"),
                  Coord.cartesian(xmin=0, xmax=15, ymin=0, ymax=10))
        p2 = plot(x=df_pred.time, y=df_pred.S, ymin=df_pred.ci_lo_S, ymax=df_pred.ci_hi_S, Geom.line, Geom.ribbon, Guide.xlabel("Time"), Guide.ylabel("Survival Probability"), Coord.cartesian(xmin=0, xmax=15, ymin=0, ymax=1))

        hstack(p1, p2)
    end
end


module Born

import GeoPhyInv.Acquisition
import GeoPhyInv.Models
import GeoPhyInv.Data
using SpecialFunctions
using FFTW
using DSP

mutable struct Param
end

"""
k = wavenumber = 2pif/v0
"""
function G0_homo_acou(x::T, z::T, k::T, ρ0::T) where {T<:AbstractFloat}
	sqrt(x*x + z*z) == 0.0 ? error("distance cannot be zero") : nothing
	k == 0.0 ? error("wavenumber cannot be zero") : nothing
	G0 = -0.25 * ρ0 * im * complex.(besselj0(k*sqrt(x*x + z*z)), -bessely0(k*sqrt(x*x + z*z)))
	return G0
end

function mod(;
	     vp0::Float64=2000.0,
	     ρ0::Float64=2000.0,
	     model_pert::Models.Seismic=nothing,
             born_flag::Bool=false,
	     tgridmod::StepRangeLen=nothing,
	     tgrid::StepRangeLen = tgridmod,
	     acqgeom::Acquisition.Geom=nothing,
	     acqsrc::Acquisition.Src=nothing,
	     src_flag::Int64=2,
		  )
	(getfield(acqgeom,:nss) != getfield(acqsrc,:nss))  ? error("different supersources") : nothing
	(getfield(acqgeom,:ns) != getfield(acqsrc,:ns))  ? error("different sources") : nothing

	nt = (length(tgridmod) == length(acqsrc.tgrid)) ? length(tgrid) : error("acqsrc tgrid")
	np2 = nextpow(2, 2*nt);

	if(born_flag)
		mesh_x=model_pert.mgrid[2]
		mesh_z=model_pert.mgrid[1]
		nz=length(mesh_z)
		nx=length(mesh_x)
		δx = step(mesh_x)
		δz = step(mesh_z)

		δmodtt = Models.Seismic_get(model_pert, :KI) - (vp0 * vp0 * ρ0)^(-1)
		δmodrr = Models.Seismic_get(model_pert, :ρI) - (ρ0)^(-1)
	end

	fnpow2grid = DSP.fftfreq(np2,inv(step(tgridmod)));

	nss = acqgeom.nss
	nr = acqgeom.nr
	ns = acqgeom.ns
	data = Data.TD(
	      [zeros(nt,nr[iss]) for iss=1:nss, ifield=1:1],
	      [:P],
	      tgridmod,acqgeom)

	for ifield = 1:length(data.fields), iss = 1:nss
		sx = acqgeom.sx[iss][:]
		rx = acqgeom.rx[iss][:]
		sz = acqgeom.sz[iss][:]
		rz = acqgeom.rz[iss][:]

		for ir = 1:nr[iss]

		dtemp=zeros(nt)
		dpow2all=complex.(zeros(np2), zeros(np2));
		wpow2=complex.(zeros(np2), zeros(np2));

		for is=1:acqgeom.ns[iss]
			# zero pad wavelet
			for it in 1:nt
				wpow2[it]=complex(acqsrc.wav[iss,ifield][it,is])
			end
			FFTW.fft!(wpow2) # source wavelet in the frequency domain

			x = sx[is] - rx[ir]
			z = sz[is] - rz[ir]
			# analytical expression for every frequency, except zero
			dpow2 = complex.(zeros(np2), zeros(np2));
			dpow2[1] = complex.(0.0, 0.0)
			for iω in 2:np2
				ω = 2. * pi * abs(fnpow2grid[iω])
				k = ω / vp0

				if(born_flag)
					term = complex(0., 0.)
					for ix=1:nx
						@simd for iz=1:nz
							if(δmodtt[iz,ix] ≠ 0.0)
								term += (G0_homo_acou(sx[is]-mesh_x[ix], sz[is]-mesh_z[iz], k, ρ0)[1]
									 * G0_homo_acou(rx[ir]-mesh_x[ix], rz[ir]-mesh_z[iz], k, ρ0)[1] .* ω .* ω .* δmodtt[iz,ix]) * δx * δz
								# factor due to integration
							end
						end
					end

				else
					if(src_flag == 2)
						term = G0_homo_acou(x, z, k, ρ0);
					elseif(src_flag == 1)
						dpow2[iω] = G0_homo_acou(x, z, k, ρ0) * im * abs(fnpow2grid[iω])
					else
						error("invalid src_flag")
					end
				end
				if(fnpow2grid[iω] > 0)
					dpow2[iω] = term;
				elseif(fnpow2grid[iω] < 0)
					dpow2[iω] = conj(term);
				end

			end
			# convolution and stack over simultaneous sources
			dpow2all += dpow2 .* wpow2;
		end

		# back to time domain
		FFTW.ifft!(dpow2all)

		# truncate
		for it in 1:nt
			data.d[iss,ifield][it,ir] = real(dpow2all[it])
		end
		end
	end
	return data
end

end # module


import Projections, Test, Distributions, LinearAlgebra

include("tests.jl")

divergences = [Projections.Burg(),
               Projections.Hellinger(),
               Projections.ChiSquare(),
               Projections.ModifiedChiSquare(),
               Projections.KullbackLeibler()]

norms = [Projections.Linf(),
         Projections.Lone(),
         Projections.Ltwo()]

constraints = vcat(divergences, norms)

n  = 100
εs = vcat(1e-5:0.01:0.1, 0.2, 0.5, 1)

Test.@testset "All tests" begin

    Test.@testset "Generating functions" begin
        Test.@testset "$(typeof(d))" for d in divergences
            test_generate(d)
        end
    end

    Test.@testset "DRO" begin
        Test.@testset "ε = $ε" for ε in εs
            Test.@testset "Constraint = $(Projections.name(d))" for d in constraints
                q   = LinearAlgebra.normalize(rand(n), 1)
                c1  = rand(n)
                c2  = max.(min.(rand(n), 0.8), 0.2)

                Test.@testset "standard" begin
                    m = Projections.DRO(d, q, c1, ε)
                    test_model(m)
                end
                Test.@testset "sorted" begin
                    m = Projections.DRO(d, sort(q; rev = true), sort(c1), ε)
                    test_model(m)
                end
                Test.@testset "bounds" begin
                    m = Projections.DRO(d, q, c2, ε)
                    test_model(m)
                end
            end
        end
    end

    Test.@testset "Simplex" begin
        q  = rand(n)
        lb = rand(n)./n
        ub = 1 .+ rand(n)

        m = Projections.Simplex(q, lb, ub)
        test_model(m)
    end
end


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

using Base.Test
using ModifiedStokesSolver
import AnalyticDomains

@testset "BlockCorrections" begin
    # Params
    alpha = 10.0
    # Discretize
    numpanels = 20
    panelorder = 16
    curve = AnalyticDomains.starfish(n_arms = 3, amplitude=0.2)
    dcurve = discretize(curve, numpanels, panelorder)
    N = dcurve.numpoints

    # Assemble complete double layer matrix
    Dmat = doublelayer_matrix(dcurve, alpha)

    # Assemble block correction matrix
    BCmat = ModifiedStokesSolver.correction_block_matrix(dcurve, alpha)

    tovec(x) = reshape(x, :)
    fromvec(x) = reshape(x, 2, :)
    matvec(x) = tovec(doublelayer_self(dcurve, fromvec(x), alpha, zerolimit=false)) + BCmat*x

    x = rand(2*N)
    v1 = Dmat*x
    v2 = matvec(x)
    relerr = norm(v1-v2, Inf) / norm(v1, Inf)

    @test relerr < 1e-14
end


# Random array with two phases
rand_array = rand(Float32, (50, 50, 50))
rand_array = map(x -> (x<0.3) ? 0 : 1, rand_array)

@testcase "Check s²(a, x) = l²(a, x) = 0 for all x where a is an array without phase 2." begin
    for p in (false, true)
        for func in (Directional.s2, Directional.l2)
            corr = mean(func(rand_array, 2; periodic = p, directions = known_directions))
            @test all(x -> x == 0, corr)
        end
    end
end

@testcase "Check that corr(a, len1, phase) = corr(a, len2, phase)[1:len1] for len2>len1" begin
    for phase in (0, 1)
        for func in (Directional.s2,
                     Directional.l2,
                     Directional.c2,
                     Directional.surfsurf,
                     Directional.surfvoid)
            corr1 = mean(func(rand_array, phase; len = 30))
            corr2 = mean(func(rand_array, phase; len = 35))
            @test corr1 == corr2[1:30]
        end
    end
end

# Probability of randomly choosing 0 or 1
p = count(x -> x == 0, rand_array) / length(rand_array) # P{random dot ∈ phase 0}
q = 1 - p                                               # P{random dot ∈ phase 1}
prob = [p, q]

@testcase "sᶦ(a,0) = lᶦ(a,0) = P{randomly choosing i}" begin
    for p in (false, true)
        for phase in 0:1
            for func in (Directional.s2, Directional.l2)
                f = mean(func(rand_array, phase; len = 1, directions = known_directions))[1]
                @test relerr(f, prob[phase+1]) < 0.02
            end
        end
    end
end

@testcase "sᶦ(a,1) = lᶦ(a,1)" begin
    for p in (false, true)
        for phase in 0:1
            s = mean(Directional.s2(rand_array, phase;
                                    len = 2, directions = known_directions))[2]
            l = mean(Directional.l2(rand_array, phase;
                                    len = 2, directions = known_directions))[2]
            @test s ≈ l
        end
    end
end

@testcase "Check that l2 is non-increasing" begin
    for p in (false, true)
        for phase in 0:1
            l = mean(Directional.l2(rand_array, phase; directions = known_directions))
            @test minimum(map(-, l, l[2:end])) >= 0
        end
    end
end

@testcase "Check that sⁱ(x) = P{randomly independently choosing i twice} (x > 0)" begin
    for p in (false, true)
        for phase in 0:1
            s = mean(Directional.s2(rand_array, phase;
                                    directions = known_directions, periodic = p))
            @test relerr(maximum(s[2:end]), prob[phase+1]^2) < 0.04
        end
    end
end

@testcase "Check pore size and chord length sum" begin
    for phase in 0:1
        @test sum(pore_size(rand_array, phase).weights) ≈ 1
        @test sum(Directional.chord_length(rand_array, phase)[1].weights) ≈ 1
    end
end


#=

    This is a port of the SPPL model poisson_mean_inference.pynb

    Here is the PDF of Y = 0..20 (exact probabilities) from the SPPL model:
    (0, 0.19572329209779615),
    (1, 0.24541518485806976),
    (2, 0.15633363576277434),
    (3, 0.07275764007128807),
    (4, 0.03749784880330323),
    (5, 0.03225414602603766),
    (6, 0.036272539178041456),
    (7, 0.03919655164760693),
    (8, 0.03782130413444809),
    (9, 0.032547494609049904),
    (10, 0.025221805385485633),
    (11, 0.017769927037907036),
    (12, 0.011477227802063514),
    (13, 0.006844292955361095),
    (14, 0.00379364785941744),
    (15, 0.0019704640279916506),
    (16, 0.000975431521715664),
    (17, 0.0004844975064651556),
    (18, 0.0002799391781190319),
    (19, 0.00023526933546341676)
    (20, 0.0002864142016352986)

    This Turing.jl model:

    Summary Statistics
    parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
        Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

            X    0.0256    1.0030     0.0100    0.0124   6401.6166    0.9999     2199.1125
            Y    4.2629    7.6756     0.0768    0.0975   6534.2074    1.0000     2244.6607

    Quantiles
    parameters      2.5%     25.0%     50.0%     75.0%     97.5%
        Symbol   Float64   Float64   Float64   Float64   Float64

            X   -1.9514   -0.6385    0.0195    0.7112    2.0227
            Y    0.0000    1.0000    2.0000    4.0000   26.0000

    Distributions of variable Y:
    1.00000 =>    2448  (0.244800)
    0.00000 =>    1827  (0.182700)
    2.00000 =>    1776  (0.177600)
    3.00000 =>     987  (0.098700)
    4.00000 =>     614  (0.061400)
    5.00000 =>     382  (0.038200)
    6.00000 =>     287  (0.028700)
    7.00000 =>     237  (0.023700)
    8.00000 =>     192  (0.019200)
    9.00000 =>     145  (0.014500)
    11.00000 =>     113  (0.011300)
    10.00000 =>     106  (0.010600)
    13.00000 =>      94  (0.009400)
    12.00000 =>      80  (0.008000)
    14.00000 =>      75  (0.007500)
    15.00000 =>      59  (0.005900)
    16.00000 =>      55  (0.005500)
    20.00000 =>      55  (0.005500)
    17.00000 =>      42  (0.004200)
    18.00000 =>      36  (0.003600)
    24.00000 =>      30  (0.003000)
    21.00000 =>      27  (0.002700)
    23.00000 =>      25  (0.002500)
    19.00000 =>      23  (0.002300)
    26.00000 =>      21  (0.002100)
    27.00000 =>      19  (0.001900)
    25.00000 =>      15  (0.001500)
    22.00000 =>      15  (0.001500)
    30.00000 =>      13  (0.001300)
    34.00000 =>      13  (0.001300)
    36.00000 =>      12  (0.001200)
    32.00000 =>      11  (0.001100)
    28.00000 =>      11  (0.001100)
    29.00000 =>      11  (0.001100)
    46.00000 =>      10  (0.001000)
    52.00000 =>       9  (0.000900)
    31.00000 =>       9  (0.000900)
    35.00000 =>       8  (0.000800)
    39.00000 =>       8  (0.000800)
    33.00000 =>       8  (0.000800)
    40.00000 =>       8  (0.000800)
    38.00000 =>       7  (0.000700)
    37.00000 =>       7  (0.000700)
    42.00000 =>       6  (0.000600)
    57.00000 =>       5  (0.000500)
    64.00000 =>       4  (0.000400)
    49.00000 =>       4  (0.000400)
    47.00000 =>       4  (0.000400)
    60.00000 =>       3  (0.000300)
    45.00000 =>       3  (0.000300)
    41.00000 =>       3  (0.000300)
    61.00000 =>       3  (0.000300)
    53.00000 =>       3  (0.000300)
    48.00000 =>       3  (0.000300)
    56.00000 =>       2  (0.000200)
    55.00000 =>       2  (0.000200)
    110.00000 =>       2  (0.000200)
    67.00000 =>       2  (0.000200)
    73.00000 =>       2  (0.000200)
    43.00000 =>       2  (0.000200)
    50.00000 =>       2  (0.000200)
    51.00000 =>       2  (0.000200)
    139.00000 =>       1  (0.000100)
    63.00000 =>       1  (0.000100)
    54.00000 =>       1  (0.000100)
    58.00000 =>       1  (0.000100)
    69.00000 =>       1  (0.000100)
    95.00000 =>       1  (0.000100)
    66.00000 =>       1  (0.000100)
    93.00000 =>       1  (0.000100)
    44.00000 =>       1  (0.000100)
    70.00000 =>       1  (0.000100)
    80.00000 =>       1  (0.000100)
    65.00000 =>       1  (0.000100)
    97.00000 =>       1  (0.000100)

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

@model function poisson_mean_inference(n=20)
    X ~ truncated(Normal(),-4, 4)
    Y ~ Poisson(1+2*abs(X^3))
    return Y
end

n=20
model = poisson_mean_inference(n)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, HMC(0.1,6), 1_000)
# chns = sample(model, NUTS(), 1_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns,:Y)


#
# Tensor the matrix V of order n up to order N,
# where V acts on qubits: j, j + 1,... k; j < k
#
# and |ψ₀> = |0>₁ ⊗ |0>₂ ⊗ ... ⊗ |0>(N)
#
# V(k) -> I₁ ⊗ ... ⊗ I(k - 1) ⊗ V(k) ⊗ I(k + n - 1) ⊗ ... ⊗ I(N)
#

module Tensor

using LinearAlgebra

export tensor, ⊗

const C32 = Complex{Float32}

function tensor(V::Matrix{C32}, indx::Tuple{Vararg{Int}}, N::Int)
    if length(indx) == 2
        j, k = indx
        if j > k
            # Qₖ ⊗ Qⱼ -> Qⱼ ⊗ Q(j - 1)
            Ṽ = tensor(V, j, N)
            R, L = σ(j, k, N)
        else
            # Qⱼ ⊗ Qₖ -> Q(k - 1) ⊗ Qₖ
            Ṽ = tensor(V, k, N)
            R, L = σ(k - 1, j, N)
        end
        return L * Ṽ * R
    else
        i, j, k = indx
        if i > j && j > k
            Ṽ = tensor(V, i, N)
            R₁, L₁ = σ(i, k, N)
            R₂, L₂ = σ(i - 1, j - 1, N)
        elseif i < j && j > k
            if i < k
                Ṽ = tensor(V, j, N)
                R₁, L₁ = σ(j, k, N)
                R₂, L₂ = σ(j - 2, i, N)
            else
                Ṽ = tensor(V, j, N)
                R₁, L₁ = σ(j, k, N)
                R₂, L₂ = σ(j - 2, i - 1, N)
            end
        elseif i > j && j < k
            if i < k
                Ṽ = tensor(V, k, N)
                R₁, L₁ = σ(k - 1, j, N)
                R₂, L₂ = σ(j - 2, i - 1, N)
            else
                Ṽ = tensor(V, i, N)
                R₁, L₁ = σ(i, k, N)
                R₂, L₂ = σ(i - 1, j, N)
            end
        else
            Ṽ = tensor(V, k, N)
            R₁, L₁ = σ(k - 1, j, N)
            R₂, L₂ = σ(k - 2, i, N)
        end
        return L₁ * L₂ * Ṽ * R₂ * R₁
    end
end

function tensor(V::Matrix{C32}, k::Int, N::Int)
    n = Int(log(2, size(V, 1)))
    L = diagm(ones(C32, 2^(k - n)))
    R = diagm(ones(C32, 2^(N - k)))
    L ⊗ V ⊗ R
end

⊗(x, y) = kron(x, y)

const SWAP = [1 0 0 0;
              0 0 1 0;
              0 1 0 0;
              0 0 0 1]

# Q(i) -> Q(i + 1)
τ(i, N) = tensor(C32.(SWAP), i + 1, N)

# k < j => Q(k) -> Q(j)
σ(j, k, N) = begin
    if k < j
        τs = [τ(k + j - i - 1, N) for i = k:j-1];
        return (*(τs...), *(reverse(τs)...))
    else
        return (C32(1), C32(1))
    end
end

end


abstract type AbstractGroup end
abstract type AbstractAbelianGroup <: AbstractGroup end

abstract type AbstractCharge end
struct Trivial <: AbstractCharge end
abstract type AbstractNonTrivialCharge end

struct Z2Charge <: AbstractCharge
    charge :: Bool
end
"U1Charge for U(1) abelian group"
struct U1Charge <: AbstractCharge
    charge :: Int
end
const U1 = U1Charge

@inline Base.inv(c::U1Charge) = U1Charge(-c.charge)
@inline Base.:*(a::Int, c::U1Charge) = U1Charge(a*c.charge)
@inline Base.:+(c1::U1Charge, c2::U1Charge) = U1Charge(c1.charge + c2.charge)
@inline Base.:+(c::U1Charge, a::Int) = U1Charge(c.charge + a)
@inline Base.:+(c::U1Charge) = c
@inline Base.:-(c1::U1Charge, c2::U1Charge) = U1Charge(c1.charge - c2.charge)
@inline Base.:-(c::U1Charge, a::Int) = U1Charge(c.charge - a)
@inline Base.:-(c::U1Charge) = inv(c)
@inline Base.div(c::U1Charge, a::Int) = div(c.charge, a)
@inline Base.zero(::Type{U1Charge}) = U1Charge(0)
@inline Base.isodd(c::U1Charge) = isodd(c.charge)
@inline Base.isless(c1::U1Charge, c2::U1Charge) = isless(c1.charge, c2.charge)
@inline Base.isequal(c1::U1Charge, c2::U1Charge) = isequal(c1.charge, c2.charge)
@inline Base.isequal(c1::U1Charge, c2::Int) = isequal(c1.charge, c2)
@inline Base.length(::U1Charge) = 1
@inline Base.iterate(c::U1Charge) = (c, nothing)
convert(::Type{U1Charge}, s::Int) = U1Charge(s)
convert(::Type{Int}, c::U1Charge) = c.charge

function Base.show(io::IO, c::AbstractCharge)
    print(io, c.charge)
    nothing
end


#=

  Euler #23 in Julia.

  Problem 23:
  """
  A perfect number is a number for which the sum of its proper divisors
  is exactly equal to the number. For example, the sum of the proper divisors
  of 28 would be 1 + 2 + 4 + 7 + 14 = 28, which means that 28 is a perfect number.

  A number n is called deficient if the sum of its proper divisors is less than
  n and it is called abundant if this sum exceeds n.

  As 12 is the smallest abundant number, 1 + 2 + 3 + 4 + 6 = 16, the smallest number
  that can be written as the sum of two abundant numbers is 24. By mathematical
  analysis, it can be shown that all integers greater than 28123 can be written
  as the sum of two abundant numbers. However, this upper limit cannot be reduced
  any further by analysis even though it is known that the greatest number that
  cannot be expressed as the sum of two abundant numbers is less than this limit.

  Find the sum of all the positive integers which cannot be written as the sum of
  two abundant numbers.
  """

  This Julia program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

include("Euler.jl")

#
# Note: This is a solution ported from my JavaScript program via Picat code
# which in turn is from some C++(?) code,
# 0.00790154s
function euler23a()
    limit = 20161;
    arr = ones(limit)
    for i in 2:limit
        for j in i*2:i:limit
            arr[j] = arr[j] + i
        end
    end

    abundant = [i for i in 12:limit if arr[i] > i]
    for a in abundant
        for b in abundant
            if b > a || a + b >= limit
                break;
            else
                arr[a + b] = 0
            end
        end
    end

    return sum([i for i in 1:limit if arr[i] != 0])
end

run_euler(euler23a)


f = x -> x^3
f_dash = x -> 3 * x^2
partition = collect(-1.0:1.0:1.0)

@testset "sampling tests with x^3 for error tolerance" begin
    PR.silence()
    p = deepcopy(partition)
    for i = 1:10
        lb = -1.0
        ub = 1.0
        λ = rand()
        x_val = λ * lb + (1 - λ) * ub
        m = Model(cbc_optimizer)
        @variable(m, -1.0 <= x <= 1.0)
        @variable(m, y)
        construct_univariate_relaxation!(m, f, x, y, p, true, error_tolerance = 1e-2)
        @constraint(m, x == x_val)

        # solve for min y
        @objective(m, Min, y)
        optimize!(m)
        y_min = objective_value(m)

        # solve for max y
        @objective(m, Max, y)
        optimize!(m)
        y_max = objective_value(m)
        @test abs(y_max - y_min) <= 1e-2
    end
end

@testset "sampling tests with x^3 for binary variable budget" begin
    num_vars = 10
    tol = 1e-5
    for l in [0.1, 0.25, 0.5, 1.0]
        base_partition = collect(-1.0:l:1.0)
        num_base = length(base_partition) - 1
        m = Model(cbc_optimizer)
        @variable(m, -1.0 <= x <= 1.0)
        @variable(m, y)
        formulation_info = construct_univariate_relaxation!(
            m,
            f,
            x,
            y,
            base_partition,
            true,
            error_tolerance = tol,
            num_additional_partitions = num_vars,
        )
        @test length(formulation_info.variables[:z]) == num_base + num_vars
    end
end

@testset "sampling tests with x^3 for length tolerance" begin
    err_tol = 1e-7
    len_tol = 0.1
    p = deepcopy(partition)
    m = Model(cbc_optimizer)
    @variable(m, -1.0 <= x <= 1.0)
    @variable(m, y)
    construct_univariate_relaxation!(
        m,
        f,
        x,
        y,
        p,
        true,
        error_tolerance = err_tol,
        length_tolerance = len_tol,
    )
    for i = 1:length(p)-1
        @test p[i+1] - p[i] >= len_tol
    end
end

@testset "sampling tests with x^3 for derivative tolerance" begin
    err_tol = 1e-7
    der_tol = 0.1
    p = deepcopy(partition)
    m = Model(cbc_optimizer)
    @variable(m, -1.0 <= x <= 1.0)
    @variable(m, y)
    construct_univariate_relaxation!(
        m,
        f,
        x,
        y,
        p,
        true,
        error_tolerance = err_tol,
        derivative_tolerance = der_tol,
    )

    for i = 1:length(p)-1
        d = f_dash(p[i])
        d_next = f_dash(p[i+1])
        @test abs(d_next - d) >= der_tol
    end
end


mutable struct ILUTPreconditioner{Tv, Ti} <: AbstractPreconditioner{Tv,Ti}
    A::ExtendableSparseMatrix{Tv,Ti}
    fact::IncompleteLU.ILUFactorization{Tv,Ti}
    droptol::Float64
    function ILUTPreconditioner{Tv,Ti}(;droptol=1.0e-3) where {Tv,Ti}
        p=new()
        p.droptol=droptol
        p
    end
end

"""
```
ILUTPreconditioner(;droptol=1.0e-3,valuetype=Float64, indextype=Int64)
ILUTPreconditioner(matrix; droptol=1.0e-3)
```

Create the [`ILUTPreconditioner`](@ref) wrapping the one
from [IncompleteLU.jl](https://github.com/haampie/IncompleteLU.jl)
For using this, you need to issue `using IncompleteLU`.
"""
ILUTPreconditioner(;valuetype::Type=Float64, indextype::Type=Int64, kwargs...)=ILUTPreconditioner{valuetype,indextype}(;kwargs...)

@eval begin
    @makefrommatrix ILUTPreconditioner
end

function update!(precon::ILUTPreconditioner)
    A=precon.A
    @inbounds flush!(A)
    precon.fact=IncompleteLU.ilu(A.cscmatrix,τ=precon.droptol)
end


using GeometricEquations
using Requires

import ..Diagnostics: compute_invariant_error

export lotka_volterra_2d_ode,  lotka_volterra_2d_dae,
       lotka_volterra_2d_pode, lotka_volterra_2d_pdae,
       lotka_volterra_2d_iode, lotka_volterra_2d_idae,
       lotka_volterra_2d_lode, lotka_volterra_2d_ldae,
       lotka_volterra_2d_hode, lotka_volterra_2d_hdae,
       lotka_volterra_2d_dg,   lotka_volterra_2d_slrk,
       lotka_volterra_2d_idae_spark

export lotka_volterra_2d_ode_poincare_invariant_1st,
       lotka_volterra_2d_iode_poincare_invariant_1st

export compute_energy_error

const Δt = 0.1
const nt = 1000

const parameters = (a₁=-1.0, a₂=-1.0, b₁=1.0, b₂=2.0)
const reference_solution = [0.34520671417508214, 1.1140647339845946]

const q₀ = [2.0, 1.0]
const v₀=[v₁(0, q₀, parameters), v₂(0, q₀, parameters)]

function f_loop(s)
   rx = 0.2
   ry = 0.3
   x0 = 1.0
   y0 = 1.0

   xs = x0 + rx*cos(2π*s)
   ys = y0 + ry*sin(2π*s)

   qs = [xs, ys]

   return qs
end

function f_loop(i, n)
   f_loop(i/n)
end

function initial_conditions_loop(n)
   q₀ = zeros(2, n)

   for i in axes(q₀,2)
       q₀[:,i] .= f_loop(i, n)
   end

   return q₀
end

compute_energy_error(t,q,params) = compute_invariant_error(t,q, (t,q) -> hamiltonian(t,q,params))

"Creates an ODE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_ode(q₀=q₀; params=parameters)
    ODE(lotka_volterra_2d_v, q₀; parameters=params, invariants=(h=hamiltonian,))
end

"Creates a Hamiltonian ODE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_hode(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀); params=parameters)
    HODE(lotka_volterra_2d_v, lotka_volterra_2d_f, hamiltonian_pode, q₀, p₀;
         parameters=params)
end

"Creates an implicit ODE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_iode(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀); params=parameters)
    IODE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
         lotka_volterra_2d_g, q₀, p₀;
         parameters=params, invariants=(h=hamiltonian_iode,), v̄=lotka_volterra_2d_v)
end

"Creates a partitioned ODE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_pode(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀); params=parameters)
    PODE(lotka_volterra_2d_v, lotka_volterra_2d_f, q₀, p₀;
         parameters=params, invariants=(h=hamiltonian_pode,))
end

"Creates a variational ODE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_lode(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀); params=parameters)
    LODE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
         lotka_volterra_2d_g, lagrangian, lotka_volterra_2d_ω, q₀, p₀;
         parameters=params, invariants=(h=hamiltonian_iode,), v̄=lotka_volterra_2d_v)
end

"Creates a DAE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_dae(q₀=vcat(q₀,v₀), λ₀=zero(q₀); params=parameters)
    DAE(lotka_volterra_2d_v_dae, lotka_volterra_2d_u_dae, lotka_volterra_2d_ϕ_dae, q₀, λ₀;
        parameters=params, invariants=(h=hamiltonian,))
end

"Creates a Hamiltonian DAE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_hdae(q₀=q₀, p₀=ϑ(0, q₀), λ₀=zero(q₀); params=parameters)
    HDAE(lotka_volterra_2d_v, lotka_volterra_2d_f,
         lotka_volterra_2d_u, lotka_volterra_2d_g, lotka_volterra_2d_ϕ,
         lotka_volterra_2d_u̅, lotka_volterra_2d_g̅, lotka_volterra_2d_ψ,
         hamiltonian_pode, q₀, p₀, λ₀; parameters=params)
end

"Creates an implicit DAE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_idae(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀), λ₀=zero(q₀); params=parameters)
    IDAE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
         lotka_volterra_2d_u, lotka_volterra_2d_g, lotka_volterra_2d_ϕ,
         q₀, p₀, λ₀; parameters=params, invariants=(h=hamiltonian_iode,),
         v̄=lotka_volterra_2d_v)
end

"Creates an implicit DAE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_idae_spark(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀), λ₀=zero(q₀); params=parameters)
    IDAE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f_ham,
         lotka_volterra_2d_u, lotka_volterra_2d_g, lotka_volterra_2d_ϕ,
         q₀, p₀, λ₀; parameters=params, invariants=(h=hamiltonian_iode,),
         v̄=lotka_volterra_2d_v, f̄=lotka_volterra_2d_f)
end

"Creates a partitioned DAE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_pdae(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀), λ₀=zero(q₀); params=parameters)
    PDAE(lotka_volterra_2d_v_ham, lotka_volterra_2d_f_ham,
         lotka_volterra_2d_u, lotka_volterra_2d_g, lotka_volterra_2d_ϕ,
         q₀, p₀, λ₀; parameters=params, invariants=(h=hamiltonian_pode,),
         v̄=lotka_volterra_2d_v, f̄=lotka_volterra_2d_f)
end

"Creates a variational DAE object for the Lotka-Volterra 2D model."
function lotka_volterra_2d_ldae(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀), λ₀=zero(q₀); params=parameters)
    LDAE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f_ham,
         lotka_volterra_2d_u, lotka_volterra_2d_g, lotka_volterra_2d_ϕ,
         lotka_volterra_2d_u̅, lotka_volterra_2d_g̅, lotka_volterra_2d_ψ_lode,
         lagrangian, lotka_volterra_2d_ω,
         q₀, p₀, λ₀; parameters=params, invariants=(h=hamiltonian_iode,),
         v̄=lotka_volterra_2d_v, f̄=lotka_volterra_2d_f)
end

"Creates a variational DAE object for the Lotka-Volterra 2D model for use with SLRK integrators."
function lotka_volterra_2d_slrk(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀), λ₀=zero(q₀); params=parameters)
    LDAE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
            lotka_volterra_2d_u, lotka_volterra_2d_g, lotka_volterra_2d_ϕ,
            lotka_volterra_2d_u̅, lotka_volterra_2d_g̅, lotka_volterra_2d_ψ,
            lagrangian, lotka_volterra_2d_ω,
            q₀, p₀, λ₀; parameters=params, invariants=(h=hamiltonian_iode,),
            v̄=lotka_volterra_2d_v, f̄=lotka_volterra_2d_f)
end

"Creates an implicit ODE object for the Lotka-Volterra 2D model for use with DG integrators."
function lotka_volterra_2d_dg(q₀=q₀, p₀=lotka_volterra_2d_pᵢ(q₀); params=parameters)
    IODE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f, lotka_volterra_2d_g,
         q₀, p₀; parameters=params, invariants=(h=hamiltonian_iode,), v̄=lotka_volterra_2d_v)
end

function lotka_volterra_2d_loop_ode(n)
   lotka_volterra_2d_ode(initial_conditions_loop(n))
end

function lotka_volterra_2d_loop_iode(n)
   lotka_volterra_2d_iode(initial_conditions_loop(n))
end

function __init__()
    @require PoincareInvariants = "26663084-47d3-540f-bd97-40ca743aafa4" begin

        function lotka_volterra_2d_ode_poincare_invariant_1st(Δt, nloop, ntime, nsave, DT=Float64)
            PoincareInvariant1st(lotka_volterra_2d_ode, f_loop, ϑ, Δt, 2, nloop, ntime, nsave, DT)
        end

        function lotka_volterra_2d_iode_poincare_invariant_1st(Δt, nloop, ntime, nsave, DT=Float64)
            PoincareInvariant1st(lotka_volterra_2d_iode, f_loop, ϑ, Δt, 2, nloop, ntime, nsave, DT)
        end

    end
end


###############################################################################
#
#  Twists
#
#
#  (C) 2022 Jeroen Hanselman
#
###############################################################################

export quadratic_twist, quadratic_twists, is_twist, twists

@doc Markdown.doc"""
    quadratic_twist(E::EllCrv{T}}, d::T) -> EllCrv{T}

Compute the quadratic twist of $E$ by $d$.
"""
#Needs more testing over characteristic 2. (Magma's twists don't always work over char 2.
#Adapted from Connell's Handbook of elliptic curves.
function quadratic_twist(E::EllCrv{T}, d) where T<: FieldElem

  a1, a2, a3, a4, a6 = a_invars(E)
  K = base_field(E)
  if characteristic(K) != 2
    return EllipticCurve(K, [a1, a2*d + a1^2*(d-1)//4, a3, a4*d^2 + a1*a3*(d^2-1)//2, a6*d^3 + a3^2*(d^3 -1)//4])
  end

  return EllipticCurve(K, [a1, a2+a1^2*d, a3, a4, a6 + a3^2])

end

@doc Markdown.doc"""
    quadratic_twist(E::EllCrv{FinFieldElem}}) -> EllCrv{FinFieldElem}

Compute the unique quadratic twist of $E$.
"""
function quadratic_twist(E::EllCrv{T}) where T<: FieldElem

  K = base_field(E)
  char = characteristic(K)

 if char == 2
   f, h = hyperelliptic_polynomials(E)
   if iseven(degree(K))
     u = normal_basis(GF(Int(char), 1),K)
   else
     u = one(K)
   end

   return EllipticCurve(f + u*h^2, h)
 end

  a = gen(K)
  return quadratic_twist(E, a)

end

#Test if we can't sometimes get two isomorphic curves
@doc Markdown.doc"""
    quadratic_twists(E::EllCrv{FinFieldElem}}) -> Vector{EllCrv{FinFieldElem}}

Compute all quadratic twists of $E$.
"""
function quadratic_twists(E::EllCrv{T}) where T<: FinFieldElem

  return [E, quadratic_twist(E)]

end

function supersingular_twists2(E::EllCrv{T}) where T<: FinFieldElem
#Adapted from Magma code
  K = base_field(E)
  if isodd(degree(K))
    return EllCrv{T}[EllipticCurve(K, [0, 0, 1, 0, 0]), EllipticCurve(K, [0, 0, 1, 1, 0]), EllipticCurve(K, [0, 0, 1, 1, 1]) ]
  end

  u = gen(K);
  c = u
  while absolute_tr(c) == 0
    c = rand(K)
  end
  #First three curves
  part_1 = EllCrv{T}[EllipticCurve(K, [0, 0, 1, 0, 0]), EllipticCurve(K, [0, 0, 1, c, 0]), EllipticCurve(K, [0, 0, 1, 0, c]) ]
  #The other four
  part_2 = EllCrv{T}[EllipticCurve(K, [0, 0, a, 0, a^2*b]) for (a, b) in [[u, 0], [inv(u), 0], [u, c], [inv(u), c]]]
  return vcat(part_1, part_2)
end

function supersingular_twists3(E::EllCrv{T}) where T<: FinFieldElem
#Adapted from Magma code
    K = base_field(E)
    d = degree(K)

    if mod(d, 3) != 0
        c = one(K)
    else
      c = gen(K)
      while absolute_tr(c) == 0
        c = rand(K)
      end
    end

    if isodd(d)
      return EllCrv{T}[EllipticCurve(K, [1,0]), EllipticCurve(K, [-1, 0]), EllipticCurve(K, [-1, c]), EllipticCurve(K, [-1, -c]) ];
    end
    u = gen(K);
    #First four curves
    part_1 = EllCrv{T}[EllipticCurve(K, [-u^i,0]) for i in (0:3)]
    part_2 = EllCrv{T}[EllipticCurve(K, [-1,c]), EllipticCurve(K, [-u^2,(u^3)*c])]
    return vcat(part_1, part_2)

end

@doc Markdown.doc"""
    twists(E::EllCrv{FinFieldElem}}) -> Vector{EllCrv{FinFieldElem}}

Compute all twists of $E$.
"""
function twists(E::EllCrv{T}) where T<: FinFieldElem
#Adapted from Magma code
   K = base_field(E);
   p = characteristic(K)
   j = j_invariant(E)
   if j != 0 && j != 1728
      return EllCrv{T}[E, quadratic_twist(E)]
   elseif p == 2
      return supersingular_twists2(E)
   elseif p == 3
      return supersingular_twists3(E)
   elseif j == 0
      a = gen(K)
      c4, c6 = c_invars(E)
      c = -c6//864
      n = gcd(6, order(K)-1)
      return EllCrv{T}[ EllipticCurve(K, [0,c*a^i]) for i in (0:n-1) ]
   elseif j == 1728
      a = gen(K)
      c4, c6 = c_invars(E)
      c = -c4//48;
      n = gcd(4, order(K)-1)
      return EllCrv{T}[ EllipticCurve(K, [c*a^i,0]) for i in (0:n-1)]
   end
   #This can never happen:
   error("There is a bug in twists.")
end

@doc Markdown.doc"""
    is_twist(EEllCrv, EllCrv) -> Bool

Return true if $F$ is a twist of $E$
"""
function is_twist(E::EllCrv, F::EllCrv)

  return j_invariant(E) == j_invariant(F)
end


using BandedMatrices, LinearAlgebra, LazyArrays, Test
    import BandedMatrices: MemoryLayout, TriangularLayout, BandedColumnMajor

@testset "Triangular" begin
    @testset "Upper" begin
        A = UpperTriangular(brand(10,10,1,2))
        @test isbanded(A)
        @test MemoryLayout(A) == TriangularLayout{'U','N'}(BandedColumnMajor())
        @test bandwidths(A) == bandwidths(BandedMatrix(A)) == (0,2)
        @test bandwidth(A,2) == 2
        @test BandedMatrix(A) == A
        @test A[2,1] == 0

        x=rand(10)
        @test A*x ≈ BandedMatrix(A)*x ≈ Matrix(A)*x
        @test all(A*x .=== lmul!(A, copy(x)) .=== (similar(x) .= Mul(A,x)) .=== copyto!(similar(x), Mul(A,copy(x))) .===
                    BandedMatrices.tbmv!('U', 'N', 'N', 10, 2, parent(A).data, copy(x)))

        @test A\x ≈ Matrix(A) \ x ≈ BandedMatrix(A) \ x
        @test all((A \ x) .=== ldiv!(A, copy(x)) .=== copyto!(similar(x), Ldiv(A, x)) .=== (similar(x) .= Ldiv(A, x)) .===
                    BandedMatrices.tbsv!('U', 'N', 'N', 10, 2, parent(A).data, copy(x)))

        A = UnitUpperTriangular(brand(10,10,1,2))
        @test isbanded(A)
        @test MemoryLayout(A) == TriangularLayout{'U','U'}(BandedColumnMajor())
        @test bandwidths(A) == bandwidths(BandedMatrix(A)) == (0,2)
        @test bandwidth(A,2) == 2
        @test BandedMatrix(A) == A
        @test A[2,1] == 0
        @test A[2,2] == 1

        x=rand(10)
        @test A*x ≈ Matrix(A)*x
        @test all(A*x .=== lmul!(A, copy(x)) .=== (similar(x) .= Mul(A,x)) .=== copyto!(similar(x), Mul(A,copy(x))) .===
                    BandedMatrices.tbmv!('U', 'N', 'U', 10, 2, parent(A).data, copy(x)))

        @test A\x ≈ Matrix(A) \ x ≈ BandedMatrix(A) \ x
        @test all((A \ x) .=== ldiv!(A, copy(x)) .=== copyto!(similar(x), Ldiv(A, x)) .=== (similar(x) .= Ldiv(A, x)) .===
                    BandedMatrices.tbsv!('U', 'N', 'U', 10, 2, parent(A).data, copy(x)))
    end
    @testset "Lower" begin
        A = LowerTriangular(brand(10,10,1,2))
        @test isbanded(A)
        @test MemoryLayout(A) == TriangularLayout{'L','N'}(BandedColumnMajor())
        @test bandwidths(A) == bandwidths(BandedMatrix(A)) == (1,0)
        @test bandwidth(A,2) == 0
        @test BandedMatrix(A) == A
        @test A[1,2] == 0

        x=rand(10)
        @test A*x ≈ Matrix(A)*x
        @test all(A*x .=== lmul!(A, copy(x)) .=== copyto!(similar(x), Mul(A,copy(x))) .=== (similar(x) .= Mul(A,x)) .===
                    BandedMatrices.tbmv!('L', 'N', 'N', 10, 1, view(parent(A).data, 3:4,:), copy(x)))

        @test A\x ≈ Matrix(A) \ x ≈ BandedMatrix(A) \ x
        @test all((A \ x) .=== ldiv!(A, copy(x)) .=== copyto!(similar(x), Ldiv(A, x)) .=== (similar(x) .= Ldiv(A, x)) .===
                    BandedMatrices.tbsv!('L', 'N', 'N', 10, 1, view(parent(A).data, 3:4,:), copy(x)))

        A = UnitLowerTriangular(brand(10,10,1,2))
        @test isbanded(A)
        @test MemoryLayout(A) == TriangularLayout{'L','U'}(BandedColumnMajor())
        @test bandwidths(A) == bandwidths(BandedMatrix(A)) == (1,0)
        @test bandwidth(A,2) == 0
        @test BandedMatrix(A) == A
        @test A[1,2] == 0
        @test A[2,2] == 1

        x=rand(10)
        @test A*x ≈ Matrix(A)*x
        @test all(A*x .=== lmul!(A, copy(x)) .=== copyto!(similar(x), Mul(A,copy(x))) .=== (similar(x) .= Mul(A,x)) .===
                    BandedMatrices.tbmv!('L', 'N', 'U', 10, 1, view(parent(A).data, 3:4,:), copy(x)))

        @test A\x ≈ Matrix(A) \ x ≈ BandedMatrix(A) \ x
        @test all((A \ x) .=== ldiv!(A, copy(x)) .=== copyto!(similar(x), Ldiv(A, x)) .=== (similar(x) .= Ldiv(A, x)) .===
                    BandedMatrices.tbsv!('L', 'N', 'U', 10, 1, view(parent(A).data, 3:4,:), copy(x)))
    end
end


module Searching
export bfs_parents, bfs_tree, dfs_parents, dfs_tree

using ...CSetDataStructures, ..BasicGraphs

"""
    tree(parents)
Convert a parents array into a directed graph.
"""
function tree(parents::AbstractVector{Int})
    n = T(length(parents))
    t = Graph(n)
    for (v, u) in enumerate(parents)
        if u > 0 && u != v
            add_edge!(t, u, v)
        end
    end
    return t
end

"""
    bfs_parents(g, s[; dir=:out])

Perform a breadth-first search of graph `g` starting from vertex `s`.
Return a vector of parent vertices indexed by vertex. If `dir` is specified,
use the corresponding edge direction (`:in` and `:out` are acceptable values).

### Performance
This implementation is designed to perform well on large graphs. There are
implementations which are marginally faster in practice for smaller graphs,
but the performance improvements using this implementation on large graphs
can be significant.
"""
bfs_parents(g::ACSet, s::Int; dir = :out) =
    (dir == :out) ? _bfs_parents(g, s, outneighbors) : _bfs_parents(g, s, inneighbors)

function _bfs_parents(g::ACSet, source, neighborfn::Function)
    n = nv(g)
    visited = falses(n)
    parents = zeros(Int, nv(g))
    cur_level = Int[]
    sizehint!(cur_level, n)
    next_level = Int[]
    sizehint!(next_level, n)
    @inbounds for s in source
        visited[s] = true
        push!(cur_level, s)
        parents[s] = s
    end
    while !isempty(cur_level)
        @inbounds for v in cur_level
            @inbounds @simd for i in neighborfn(g, v)
                if !visited[i]
                    push!(next_level, i)
                    parents[i] = v
                    visited[i] = true
                end
            end
        end
        empty!(cur_level)
        cur_level, next_level = next_level, cur_level
        sort!(cur_level)
    end
    return parents
end

"""
    bfs_tree(g, s[; dir=:out])
Provide a breadth-first traversal of the graph `g` starting with source vertex `s`,
and return a directed acyclic graph of vertices in the order they were discovered.
If `dir` is specified, use the corresponding edge direction (`:in` and `:out` are
acceptable values).
"""
bfs_tree(g::ACSet, s::Integer; dir = :out) = tree(bfs_parents(g, s; dir = dir))

"""
    dfs_parents(g, s[; dir=:out])

Perform a depth-first search of graph `g` starting from vertex `s`.
Return a vector of parent vertices indexed by vertex. If `dir` is specified,
use the corresponding edge direction (`:in` and `:out` are acceptable values).

### Implementation Notes
This version of DFS is iterative.
"""
dfs_parents(g::ACSet, s::Integer; dir=:out) =
    (dir == :out) ? _dfs_parents(g, s, outneighbors) : _dfs_parents(g, s, inneighbors)

function _dfs_parents(g::ACSet, s::Int, neighborfn::Function)
    parents = zeros(Int, nv(g))
    seen = zeros(Bool, nv(g))
    S = [s]
    seen[s] = true
    parents[s] = s
    while !isempty(S)
        v = S[end]
        u = 0
        for n in neighborfn(g, v)
            if !seen[n]
                u = n
                break
            end
        end
        if u == 0
            pop!(S)
        else
            seen[u] = true
            push!(S, u)
            parents[u] = v
        end
    end
    return parents
end

"""
    dfs_tree(g, s)

Return a directed acyclic graph based on
depth-first traversal of the graph `g` starting with source vertex `s`.
"""
dfs_tree(g::AbstractGraph, s::Integer; dir=:out) = tree(dfs_parents(g, s; dir=dir))

end


# Portions translated from SLICOT-Reference distribution
# Copyright (c) 2002-2020 NICONET e.V.
function run_mb04dz(datfile, io=stdout)
    NIN = 5
    NOUT = 6
    NMAX = 100
    LDA = NMAX
    LDQG = NMAX
    A = Array{ComplexF64,2}(undef, LDA,NMAX)
    QG = Array{ComplexF64,2}(undef, LDQG,NMAX+1)
    SCALE = Array{Float64,1}(undef, NMAX)
    DUMMY = Array{Float64,1}(undef, 1)
    f = open(datfile,"r")
    readline(f)
    vs = split(readline(f))
    N = parse(BlasInt, vs[1])
    JOB =  vs[2][1]
    if ( N<=0 || N>NMAX )
        @error "Illegal N=$N"
    end

    vs = String[]
    _isz,_jsz = (N,N)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       A[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end

    vs = String[]
    _isz,_jsz = (N,N+1)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       QG[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end
    close(f)
    # interp call 1

    ILO, INFO = SLICOT.mb04dz!(JOB, N, A, QG, SCALE)
    @test INFO == 0
    INFO == 0 || return
    println(io, "ILO = $ILO")

# interp output 1
    println(io, "A:")
    _nc = N
    _nr = N
    show(io, "text/plain", A[1:_nr,1:_nc])
    println(io)

# interp output 2
    println(io, "QG:")
    _nc = N+1
    _nr = N
    show(io, "text/plain", QG[1:_nr,1:_nc])
    println(io)

if ( ILO>1 )
        subd = hypot(norm(tril(view(A,2:N,1:ILO-1))),
                    norm(tril(view(QG,1:N,1:ILO-1))))
        println(io, "subdiagonal norm: $subd")
end # if

end # run_mb04dz()


@testset "Inverse" begin
    for z ∈ Zs
        @test z * inv(z) ≈ inv(z) * z ≈ 1v
    end
end


# test the MDP module

include("./script.jl")
import .MDPSimulator

function test1()
    # graph:
    # a <-> b <-> c <-> d <-> e
    # terminals: a (10) and e (1)
    # actions: left or right
    states = ["a", "b", "c", "d", "e"]
    actions = ["L", "R"]
    terminals = ["a", "e"]
    terminal_vals = [10.0, 1.0]
    gamma = 0.2
    graph = MDPSimulator.createMDPGraph(states, actions, terminals, terminal_vals, gamma)
    MDPSimulator.addMDPGraphTransition!(graph, "b", "L", "a", prob=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "b", "R", "c", prob=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "c", "L", "b", prob=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "c", "R", "d", prob=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "d", "L", "c", prob=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "d", "R", "e", prob=1.0)
    MDPSimulator.runMDP!(graph, 100)
    MDPSimulator.displayMDP(graph)
end

function test2()
    # graph:
    # A <-> C
    # A <-> B
    # C <-> B
    # terminals: None
    # actions: C or CC
    states = ["A", "B", "C"]
    actions = ["C", "CC"]
    gamma = 0.5
    terminals::Array{String} = []
    terminal_vals::Array{AbstractFloat} = []
    graph = MDPSimulator.createMDPGraph(states, actions, terminals, terminal_vals, gamma)
    MDPSimulator.addMDPGraphTransition!(graph, "A", "C", "B", prob=1.0, reward=0.0)
    MDPSimulator.addMDPGraphTransition!(graph, "A", "CC", "C", prob=1.0, reward=-2.0)
    MDPSimulator.addMDPGraphTransition!(graph, "B", "C", "A", prob=0.4, reward=-1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "B", "C", "C", prob=0.6, reward=2.0)
    MDPSimulator.addMDPGraphTransition!(graph, "B", "CC", "A", prob=0.6, reward=2.0)
    MDPSimulator.addMDPGraphTransition!(graph, "B", "CC", "C", prob=0.4, reward=-1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "C", "C", "A", prob=0.6, reward=2.0)
    MDPSimulator.addMDPGraphTransition!(graph, "C", "C", "B", prob=0.4, reward=2.0)
    MDPSimulator.addMDPGraphTransition!(graph, "C", "CC", "A", prob=0.4, reward=2.0)
    MDPSimulator.addMDPGraphTransition!(graph, "C", "CC", "B", prob=0.6, reward=0.0)
    MDPSimulator.runMDP!(graph, 100)
    MDPSimulator.displayMDP(graph)
end

function test3()
    # graph:
    # A -> B,C
    # B -> D
    # C -> B,E
    # D -> F
    # E -> D,F
    # terminal: None
    # action: g
    states = ["A", "B", "C", "D", "E", "F"]
    actions = ["g"]
    gamma = 0.5
    terminals::Array{String} = []
    terminal_vals::Array{AbstractFloat} = []
    graph = MDPSimulator.createMDPGraph(states, actions, terminals, terminal_vals, gamma)
    MDPSimulator.addMDPGraphTransition!(graph, "A", "g", "B", prob=0.5, reward=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "A", "g", "C", prob=0.5, reward=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "B", "g", "D", prob=1.0, reward=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "C", "g", "B", prob=0.5, reward=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "C", "g", "E", prob=0.5, reward=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "D", "g", "F", prob=1.0, reward=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "E", "g", "D", prob=0.5, reward=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "E", "g", "F", prob=0.5, reward=1.0)
    MDPSimulator.addMDPGraphTransition!(graph, "F", "g", "F", prob=1.0, reward=0.0)
    MDPSimulator.runMDP!(graph, 100)
    MDPSimulator.displayMDP(graph)
end

println("Test 1")
test1()
println()

println("Test 2")
test2()
println()

println("Test 3")
test3()
println()


using kMeansClustering
using Test
using Random
using Clustering

@testset "kMeansClustering.jl" begin
    Random.seed!(1)
    n_points=1000
    dim=4
    mat = rand(dim,n_points)
    k = 10
    r = mykmeansclustering(mat,k) # organize mat into k clusters
    @test size(r[1]) == (dim,k)
    @test length(r[3]) == n_points
    c=r[4]
    kmeans!(mat,c)
    @test sortslices(r[1],dims=2) == sortslices(c,dims=2)
end


using LoopVectorization

#--------------------------------------------------------------------#

function Generate2DRectMesh(nx,ny,xmin,xmax,ymin,ymax)
    # Generates a 2D rectangular mesh of nx*ny cells
    ncell      = nx*ny
    # Twice resolution mesh
    nx2 = 2*nx + 1
    ny2 = 2*ny + 1
    # 1D axis
    x2  = LinRange(xmin,xmax,nx2)
    y2  = LinRange(ymin,ymax,ny2)
    # 2D mesh: fake ngrid for P-T space ;)
    x2d   = repeat(x2, 1, length(y2))
    y2d   = repeat(y2, 1, length(x2))'
    nodes = zeros(Int64, size(x2d))
    # Number active dofs (midfaces) - Face node numbering
    inum = 0
    for i=1:nx2
        for j=1:ny2
            if (mod(j,2)==1 && mod(i,2)==0) || (mod(j,2)==0 && mod(i,2)==1)
                inum += 1
                nodes[i,j] = inum
            end
        end
    end
    # Face nodes
    xp = zeros(inum)
    yp = zeros(inum)
    tp = zeros(Int64,inum) # node type - 1 boundary
    for i=1:nx2
        for j=1:ny2
            if (mod(j,2)==1 && mod(i,2)==0) || (mod(j,2)==0 && mod(i,2)==1)
                xp[nodes[i,j]] = x2d[i,j]
                yp[nodes[i,j]] = y2d[i,j]
                if (i==1 || i==nx2 || j==1 || j==ny2)
                    tp[nodes[i,j]] = 1
                end
            end
        end
    end
    # # South/North nodes
    # @avx for i=2:2:nx2-1
    #     for j=1:2:ny2
    #         xp[nodes[i,j]] = x2d[i,j]
    #         yp[nodes[i,j]] = y2d[i,j]
    #     end
    # end
    # # West/East nodes
    # @avx for i=1:2:nx2
    #     for j=2:2:ny2-1
    #         xp[nodes[i,j]] = x2d[i,j]
    #         yp[nodes[i,j]] = y2d[i,j]
    #     end
    # end
    # Cell 2 face numbering
    cell = zeros(4, ncell)
    @avx for i=1:nx
        for j=1:ny
            k  = i + (j-1)*nx
            ic = i+1 + (i-1)*1
            jc = j+1 + (j-1)*1
            cell[1,k] =  nodes[ic,jc-1] # South
            cell[2,k] =  nodes[ic+1,jc] # East
            cell[3,k] =  nodes[ic,jc+1] # North
            cell[4,k] =  nodes[ic-1,jc] # West
        end
    end
    return cell,xp,yp,tp
end

#--------------------------------------------------------------------#

# Create sides of mesh
xmin, xmax = 0, 1
ymin, ymax = 0, 1
nx, ny     = 5, 5  # number of elements

@time cell,xp,yp,tp = Generate2DRectMesh(nx,ny,xmin,xmax,ymin,ymax)
@time cell,xp,yp,tp = Generate2DRectMesh(nx,ny,xmin,xmax,ymin,ymax)
@time cell,xp,yp,tp = Generate2DRectMesh(nx,ny,xmin,xmax,ymin,ymax)
@time cell,xp,yp,tp = Generate2DRectMesh(nx,ny,xmin,xmax,ymin,ymax)


@testset "Automorphisms" begin

   G = PermutationGroup(Int8(4))

   @testset "AutSymbol" begin
      @test_throws MethodError Groups.AutSymbol(:a)
      @test_throws MethodError Groups.AutSymbol(:a, 1)
      f = Groups.AutSymbol(:a, 1, Groups.FlipAut(2))
      @test isa(f, Groups.GSymbol)
      @test isa(f, Groups.AutSymbol)
      @test isa(Groups.perm_autsymbol(Int8.([1,2,3,4])), Groups.AutSymbol)
      @test isa(Groups.rmul_autsymbol(1,2), Groups.AutSymbol)
      @test isa(Groups.lmul_autsymbol(3,4), Groups.AutSymbol)
      @test isa(Groups.flip_autsymbol(3), Groups.AutSymbol)
   end

   a,b,c,d = gens(FreeGroup(4))
   D = NTuple{4,FreeGroupElem}([a,b,c,d])

   @testset "flip_autsymbol correctness" begin
      @test Groups.flip_autsymbol(1)(deepcopy(D)) == (a^-1, b,c,d)
      @test Groups.flip_autsymbol(2)(deepcopy(D)) == (a, b^-1,c,d)
      @test Groups.flip_autsymbol(3)(deepcopy(D)) == (a, b,c^-1,d)
      @test Groups.flip_autsymbol(4)(deepcopy(D)) == (a, b,c,d^-1)
      @test inv(Groups.flip_autsymbol(1))(deepcopy(D)) == (a^-1, b,c,d)
      @test inv(Groups.flip_autsymbol(2))(deepcopy(D)) == (a, b^-1,c,d)
      @test inv(Groups.flip_autsymbol(3))(deepcopy(D)) == (a, b,c^-1,d)
      @test inv(Groups.flip_autsymbol(4))(deepcopy(D)) == (a, b,c,d^-1)
   end

   @testset "perm_autsymbol correctness" begin
      σ = Groups.perm_autsymbol([1,2,3,4])
      @test      σ(deepcopy(D)) == deepcopy(D)
      @test inv(σ)(deepcopy(D)) == deepcopy(D)

      σ = Groups.perm_autsymbol([2,3,4,1])
      @test      σ(deepcopy(D)) == (b, c, d, a)
      @test inv(σ)(deepcopy(D)) == (d, a, b, c)

      σ = Groups.perm_autsymbol([2,1,4,3])
      @test      σ(deepcopy(D)) == (b, a, d, c)
      @test inv(σ)(deepcopy(D)) == (b, a, d, c)

      σ = Groups.perm_autsymbol([2,3,1,4])
      @test      σ(deepcopy(D)) == (b, c, a, d)
      @test inv(σ)(deepcopy(D)) == (c, a, b, d)
   end

   @testset "rmul/lmul_autsymbol correctness" begin
      i,j = 1,2
      r = Groups.rmul_autsymbol(i,j)
      l = Groups.lmul_autsymbol(i,j)
      @test      r(deepcopy(D)) == (a*b,   b, c, d)
      @test inv(r)(deepcopy(D)) == (a*b^-1,b, c, d)
      @test      l(deepcopy(D)) == (b*a,   b, c, d)
      @test inv(l)(deepcopy(D)) == (b^-1*a,b, c, d)

      i,j = 3,1
      r = Groups.rmul_autsymbol(i,j)
      l = Groups.lmul_autsymbol(i,j)
      @test      r(deepcopy(D)) == (a, b, c*a,   d)
      @test inv(r)(deepcopy(D)) == (a, b, c*a^-1,d)
      @test      l(deepcopy(D)) == (a, b, a*c,   d)
      @test inv(l)(deepcopy(D)) == (a, b, a^-1*c,d)

      i,j = 4,3
      r = Groups.rmul_autsymbol(i,j)
      l = Groups.lmul_autsymbol(i,j)
      @test      r(deepcopy(D)) == (a, b, c, d*c)
      @test inv(r)(deepcopy(D)) == (a, b, c, d*c^-1)
      @test      l(deepcopy(D)) == (a, b, c, c*d)
      @test inv(l)(deepcopy(D)) == (a, b, c, c^-1*d)

      i,j = 2,4
      r = Groups.rmul_autsymbol(i,j)
      l = Groups.lmul_autsymbol(i,j)
      @test      r(deepcopy(D)) == (a, b*d,   c, d)
      @test inv(r)(deepcopy(D)) == (a, b*d^-1,c, d)
      @test      l(deepcopy(D)) == (a, d*b,   c, d)
      @test inv(l)(deepcopy(D)) == (a, d^-1*b,c, d)
   end

   @testset "AutGroup/Automorphism constructors" begin

      f = Groups.AutSymbol(:a, 1, Groups.FlipAut(1))
      @test isa(Automorphism{3}(f), Groups.GWord)
      @test isa(Automorphism{3}(f), Automorphism)
      @test isa(AutGroup(FreeGroup(3)), AbstractAlgebra.Group)
      @test isa(AutGroup(FreeGroup(1)), Groups.AbstractFPGroup)
      A = AutGroup(FreeGroup(1))
      @test isa(Groups.gens(A), Vector{Automorphism{1}})
      @test length(Groups.gens(A)) == 1
      A = AutGroup(FreeGroup(1), special=true)
      @test length(Groups.gens(A)) == 0
      A = AutGroup(FreeGroup(2))
      @test length(Groups.gens(A)) == 7
      gens = Groups.gens(A)

      @test isa(A(Groups.rmul_autsymbol(1,2)), Automorphism)
      @test A(Groups.rmul_autsymbol(1,2)) in gens

      @test isa(A(Groups.rmul_autsymbol(2,1)), Automorphism)
      @test A(Groups.rmul_autsymbol(2,1)) in gens

      @test isa(A(Groups.lmul_autsymbol(1,2)), Automorphism)
      @test A(Groups.lmul_autsymbol(1,2)) in gens

      @test isa(A(Groups.lmul_autsymbol(2,1)), Automorphism)
      @test A(Groups.lmul_autsymbol(2,1)) in gens

      @test isa(A(Groups.flip_autsymbol(1)), Automorphism)
      @test A(Groups.flip_autsymbol(1)) in gens

      @test isa(A(Groups.flip_autsymbol(2)), Automorphism)
      @test A(Groups.flip_autsymbol(2)) in gens

      @test isa(A(Groups.perm_autsymbol([2,1])), Automorphism)
      @test A(Groups.perm_autsymbol([2,1])) in gens
   end

   A = AutGroup(FreeGroup(4))

   @testset "eltary functions" begin

      f = Groups.perm_autsymbol([2,3,4,1])
      @test (Groups.change_pow(f, 2)).pow == 1
      @test (Groups.change_pow(f, -2)).pow == 1
      @test (inv(f)).pow == 1

      f = Groups.perm_autsymbol([2,1,4,3])
      @test isa(inv(f), Groups.AutSymbol)

      @test_throws MethodError f*f

      @test A(f)^-1 == A(inv(f))
   end

   @testset "reductions/arithmetic" begin
      f = Groups.perm_autsymbol([2,3,4,1])

      f² = Groups.r_multiply(A(f), [f], reduced=false)
      @test Groups.simplifyperms!(f²) == false
      @test f²^2 == A()

      a = A(Groups.rmul_autsymbol(1,2))*Groups.flip_autsymbol(2)
      b = Groups.flip_autsymbol(2)*A(inv(Groups.rmul_autsymbol(1,2)))
      @test a*b == b*a
      @test a^3 * b^3 == A()
      g,h = Groups.gens(A)[[1,8]] # (g, h) = (ϱ₁₂, ϱ₃₂)

      @test Groups.domain(A) == NTuple{4, FreeGroupElem}(gens(A.objectGroup))

      @test (g*h)(Groups.domain(A)) == (h*g)(Groups.domain(A))
      @test (g*h).savedhash == zero(UInt)
      @test (h*g).savedhash == zero(UInt)
      a = g*h
      b = h*g
      @test hash(a) != zero(UInt)
      @test hash(b) == hash(a)
      @test a.savedhash == b.savedhash
      @test length(unique([a,b])) == 1
      @test length(unique([g*h, h*g])) == 1

      # Not so simple arithmetic: applying starting on the left:
      # ϱ₁₂*ϱ₂₁⁻¹*λ₁₂*ε₂ == σ₂₁₃₄

      g = A(Groups.rmul_autsymbol(1,2))
      x1, x2, x3, x4 = Groups.domain(A)
      @test g(Groups.domain(A)) == (x1*x2, x2, x3, x4)
      g = g*inv(A(Groups.rmul_autsymbol(2,1)))
      @test g(Groups.domain(A)) == (x1*x2, x1^-1, x3, x4)
      g = g*A(Groups.lmul_autsymbol(1,2))
      @test g(Groups.domain(A)) == (x2, x1^-1, x3, x4)
      g = g*A(Groups.flip_autsymbol(2))
      @test g(Groups.domain(A)) == (x2, x1, x3, x4)

      @test g(Groups.domain(A)) == A(Groups.perm_autsymbol([2,1,3,4]))(Groups.domain(A))

      @test g == A(Groups.perm_autsymbol([2,1,3,4]))

      g_im = g(Groups.domain(A))
      @test length(g_im[1]) == 5
      @test length(g_im[2]) == 3
      @test length(g_im[3]) == 1
      @test length(g_im[4]) == 1
      @test length.(Groups.reduce!.(g_im)) == (1,1,1,1)

   end

   @testset "specific Aut(F4) tests" begin
      N = 4
      G = AutGroup(FreeGroup(N))
      S = G.gens
      @test isa(S, Vector{Groups.AutSymbol})
      S = [G(s) for s in unique(S)]
      @test isa(S, Vector{Automorphism{N}})
      @test S == gens(G)
      @test length(S) == 51
      S_inv = [S..., [inv(s) for s in S]...]
      @test length(unique(S_inv)) == 75

      G = AutGroup(FreeGroup(N), special=true)
      S = gens(G)
      S_inv = [G(), S..., [inv(s) for s in S]...]
      S_inv = unique(S_inv)
      B_2 = [i*j for (i,j) in Base.product(S_inv, S_inv)]
      @test length(B_2) == 2401
      @test length(unique(B_2)) == 1777
   end

   @testset "linear_repr tests" begin
      N = 3
      G = AutGroup(FreeGroup(N))
      S = unique([gens(G); inv.(gens(G))])
      R = 3

      @test Groups.linear_repr(G()) isa Matrix{Int}
      @test Groups.linear_repr(G()) == Matrix{Int}(I, N, N)

      M = Matrix{Int}(I, N, N)
      M[1,2] = 1
      ϱ₁₂ = G(Groups.rmul_autsymbol(1,2))
      λ₁₂ = G(Groups.rmul_autsymbol(1,2))

      @test Groups.linear_repr(ϱ₁₂) == M
      @test Groups.linear_repr(λ₁₂) == M

      M[1,2] = -1

      @test Groups.linear_repr(ϱ₁₂^-1) == M
      @test Groups.linear_repr(λ₁₂^-1) == M

      @test Groups.linear_repr(ϱ₁₂*λ₁₂^-1) == Matrix{Int}(I, N, N)
      @test Groups.linear_repr(λ₁₂^-1*ϱ₁₂) == Matrix{Int}(I, N, N)

      M = Matrix{Int}(I, N, N)
      M[2,2] = -1
      ε₂ = G(Groups.flip_autsymbol(2))

      @test Groups.linear_repr(ε₂) == M
      @test Groups.linear_repr(ε₂^2) == Matrix{Int}(I, N, N)

      M = [0 1 0; 0 0 1; 1 0 0]

      σ = G(Groups.perm_autsymbol([2,3,1]))
      @test Groups.linear_repr(σ) == M
      @test Groups.linear_repr(σ^3) == Matrix{Int}(I, 3, 3)
      @test Groups.linear_repr(σ)^3 == Matrix{Int}(I, 3, 3)

      function test_homomorphism(S, r)
         for elts in Iterators.product([[g for g in S] for _ in 1:r]...)
            prod(Groups.linear_repr.(elts)) == Groups.linear_repr(prod(elts)) || error("linear representaton test failed at $elts")
         end
         return 0
      end

      @test test_homomorphism(S, R) == 0
   end
end


nmonomers(species::AbstractSpecies) = length(species.mids)

hasmonomer(species::AbstractSpecies, mid::Integer) = mid in species.mids

monomer_fraction(species::AbstractSpecies, mid::Integer) = 0.0

"""
    density!(species)

Update the internal density fields for an `AbstractSpecies`.
"""
density!(species::AbstractSpecies) = nothing

setup!(species::AbstractSpecies, sys::FieldSystem) = nothing

#==============================================================================#

"""
    chain_grid(N_block, Ns)

Calculate the total chain length and block discretization
for a generic polymer chain with variable number of blocks.

The number of contour segments returned for each block needs to be even
for use of Simpson's rule in density/stress calculations.
"""
function chain_grid(N_block::AbstractVector{<:Real}, Ns::Integer)
    N = sum(N_block)
    nblocks = length(N_block)
    ds = N / Ns

    begin_block = Vector{Int}(undef, nblocks+1)
    ds_block = Vector{Float64}(undef, nblocks)

    bgn = 1
    begin_block[1] = bgn

    # Loop over all blocks in the chain
    for iblk = 1:nblocks
        Ns_half = floor(Int, N_block[iblk] / ds / 2.0 + 0.5)

        # Calculate step size for each block
        if Ns_half == 0
            Ns_half = 1
            ds_block[iblk] = N_block[iblk] / 2.0
        else
            ds_block[iblk] = N_block[iblk] / Ns_half / 2.0
        end

        # Calculate location of block start
        bgn = bgn + Ns_half * 2
        begin_block[iblk+1] = bgn
    end

    return N, begin_block, ds_block
end


# Naming convention:
# ------------------
# {type}[S]_{dims}[_d][_n]
# type: C: Continuous, D: Discrete
# scalar: add "S" if any input is scalar
# dims: "nxnuny"
# feedthrough: append "_d" if `D` is present
# names: append "_n" if some inputs/outputs/states are named

@testset "test_statespace" begin
# SCALARS
a_2 = [-5 -3; 2 -9]
CS_111 = ss(-5, 2, 3, [0])
CS_111_d = ss([3], 2, 1, 1)
CS_211 = ss(a_2, [1; 2], [1 0], 0)
CS_221 = ss(a_2, [1 0; 0 2], [1 0], 0)
CS_222 = ss(a_2, [1 0; 0 2], eye_(2), 0)

# CONTINUOUS
a_1 = [-5]
C_111 = ss(a_1, [2], [3], [0])
C_211 = ss(a_2, [1; 2], [1 0], [0])
C_212 = ss(a_2, [1; 2], eye_(2), [0; 0])
C_221 = ss(a_2, [1 0; 0 2], [1 0], [0 0])
C_222 = ss(a_2, [1 0; 0 2], eye_(2), zeros(2,2))
C_222_d = ss(a_2, [1 0; 0 2], eye_(2), eye_(2))
C_022 = ss(4.0*eye_(2))

# DISCRETE
da_1 = [-0.5]
da_2 = [0.2 -0.8; -0.8 0.07]
D_111 = ss(da_1, [2], [3], [0], 0.005)
D_211 = ss(da_2, [1; 2], [1 0], [0], 0.005)
D_221 = ss(da_2, [1 0; 0 2], [1 0], [0 0], 0.005)
D_222 = ss(da_2, [1 0; 0 2], eye_(2), zeros(2,2), 0.005)
D_222_d = ss(da_2, [1 0; 0 2], eye_(2), eye_(2), 0.005)
D_022 = ss(4.0*eye_(2), 0.005)

# Definition of input, output and state names
C_222_d_n = ss(a_2, [1 0; 0 2], eye_(2), eye_(2))

# TESTS
# Contstuct with scalars
@test CS_111 == C_111
@test CS_111_d == ss([3],[2],[1],[1])
@test CS_211 == C_211
@test CS_221 == C_221
@test CS_222 == C_222

# Addition
@test C_111 + C_111 == ss([-5 0; 0 -5],[2; 2],[3 3],[0])
@test C_222 + C_222 == ss([-5 -3 0 0; 2 -9 0 0; 0 0 -5 -3;
        0 0 2 -9],[1 0; 0 2; 1 0; 0 2], [1 0 1 0; 0 1 0 1],[0 0; 0 0])
@test C_222 + 1 == ss([-5 -3; 2 -9],[1 0; 0 2],[1 0; 0 1],[1 1; 1 1])
@test D_111 + D_111 == ss([-0.5 0; 0 -0.5],[2; 2],[3 3],[0], 0.005)

# Subtraction
@test C_111 - C_211 == ss([-5 0 0; 0 -5 -3; 0 2 -9],[2; 1; 2],[3 -1 -0],[0])
@test 1 - C_222 == ss([-5 -3; 2 -9],[1 0; 0 2],[-1 -0; -0 -1],[1 1; 1 1])
@test D_111 - D_211 == ss([-0.5 0 0; 0 0.2 -0.8; 0 -0.8 0.07],[2; 1; 2],
        [3 -1 -0],[0], 0.005)

# Multiplication
@test C_111 * C_221 == ss([-5 2 0; 0 -5 -3; 0 2 -9],
        [0 0; 1 0; 0 2],[3 0 0],[0 0])
@test C_212 * C_111 == ss([-5 -3 3; 2 -9 6; 0 0 -5],
        [0; 0; 2],[1 0 0; 0 1 0],[0; 0])
@test 4*C_222 == ss([-5 -3; 2 -9],[1 0; 0 2],[4 0; 0 4],[0 0; 0 0])
@test D_111 * D_221 == ss([-0.5 2 0; 0 0.2 -0.8; 0 -0.8 0.07],
        [0 0; 1 0; 0 2],[3 0 0],[0 0],0.005)

# Division
@test 1/C_222_d == ss([-6 -3; 2 -11],[1 0; 0 2],[-1 0; -0 -1],[1 -0; 0 1])
@test C_221/C_222_d == ss([-5 -3 -1 0; 2 -9 -0 -2; 0 0 -6 -3;
        0 0 2 -11],[1 0; 0 2; 1 0; 0 2],[1 0 0 0],[0 0])
@test 1/D_222_d == ss([-0.8 -0.8; -0.8 -1.93],[1 0; 0 2],[-1 0; -0 -1],
        [1 -0; 0 1],0.005)

# Indexing
@test size(C_222) == (2, 2)
@test size(C_212) == (2, 1)
@test C_222[1,1] == ss([-5 -3; 2 -9],[1; 0],[1 0],[0])
@test C_222[1:1,1] == ss([-5 -3; 2 -9],[1; 0],[1 0],[0])
@test C_222[1,1:2] == C_221
@test size(C_222[1,[]]) == (1,0)

A = [-1.0 -2.0; 0.0 -1.0]
B = [0.0; -2.0]
C = [1.0 1.0]
D = 1.0
sys = ss(A, B, C, D)

@test sys + 1.0 == ss(A, B, C, D + 1.0)
@test 2.0 + sys == ss(A, B, C, D + 2.0)

@test -sys == ss(A, B, -C, -D)

# Printing
res = ("StateSpace{Float64,Array{Float64,2}}\nA = \n -5.0  -3.0\n  2.0  -9.0\nB = \n 1.0  0.0\n 0.0  2.0\nC = \n 1.0  0.0\n 0.0  1.0\nD = \n 0.0  0.0\n 0.0  0.0\n\nContinuous-time state-space model")
@test sprint(show, C_222) == res
res = ("StateSpace{Float64,Array{Float64,2}}\nA = \n  0.2  -0.8 \n -0.8   0.07\nB = \n 1.0  0.0\n 0.0  2.0\nC = \n 1.0  0.0\n 0.0  1.0\nD = \n 0.0  0.0\n 0.0  0.0\n\nSample Time: 0.005 (seconds)\nDiscrete-time state-space model")
@test sprint(show, D_222) == res
res = ("StateSpace{Float64,Array{Float64,2}}\nD = \n 4.0  0.0\n 0.0  4.0\n\nStatic gain")
@test sprint(show, C_022) == res
res = "StateSpace{Float64,Array{Float64,2}}\nD = \n 4.0  0.0\n 0.0  4.0\n\nSample Time: 0.005 (seconds)\nStatic gain"
@test sprint(show, D_022) == res

# Errors
@test_throws ErrorException C_111 + C_222             # Dimension mismatch
@test_throws ErrorException C_111 - C_222             # Dimension mismatch
@test_throws ErrorException C_111 * C_222             # Dimension mismatch
@test_throws ErrorException D_111 + C_111             # Sampling time mismatch
@test_throws ErrorException D_111 - C_111             # Sampling time mismatch
@test_throws ErrorException D_111 * C_111             # Sampling time mismatch
D_diffTs = ss([1], [2], [3], [4], 0.1)
@test_throws ErrorException D_111 + D_diffTs            # Sampling time mismatch
@test_throws ErrorException D_111 - D_diffTs            # Sampling time mismatch
@test_throws ErrorException D_111 * D_diffTs            # Sampling time mismatch
@test_throws ErrorException 1/C_222                     # Not invertible
@test_throws ErrorException 1/C_212                     # Not invertible
@test_throws ErrorException ss([1 2], [1], [2], [3])      # Not square A
@test_throws ErrorException ss([1], [2 0], [1], [2])      # I/0 dim mismatch
@test_throws ErrorException ss([1], [2], [3 4], [1])      # I/0 dim mismatch
@test_throws ErrorException ss([1], [2], [3], [4], -0.1)  # Negative samping time
@test_throws ErrorException ss(eye_(2), eye_(2), eye_(2), [0]) # Dimension mismatch
end


using Distributions
using Flux.CUDA

"""
	Implements all needed parts for constructing f-AnoGAN model
		(https://www.sciencedirect.com/science/article/pii/S1361841518302640).
"""

struct fAnoGAN
	prior
	generator
	discriminator
	encoder
end

Flux.@functor fAnoGAN

function Flux.trainable(model::fAnoGAN)
	(generator = model.generator, discriminator = model.discriminator, encoder = model.encoder)
end

function izi(m::fAnoGAN, x)
	z = m.encoder(x)
	x̂ = m.generator(z)
	f = m.discriminator[1:end-1]
	fx = f(x)
	fx̂ = f(x̂)
	return x̂, fx, fx̂
end

"""
	function izif_loss(m::fAnoGAN, x, κ)
loss L_{izi_f} = 1/n * ||x-D(x)||^2 + κ / n_g ||f(x) - f(G(x))||^2
"""
function izif_loss(m::fAnoGAN, x, κ)
	x̂, fx, fx̂ = izi(m, x)
	loss = Flux.mse(x,x̂) + κ .* Flux.mse(fx,fx̂)
	return loss
end

"""
	function anomaly_score(m::fAnoGAN, x, κ=1f0)
Anomaly score computed according to formula from fAnoGAN paper.
Version izi_f (image-z-image)
		A(x) = 1/n * ||x-D(x)||^2 + κ / n_g ||f(x) - f(G(x))||^2
"""
function anomaly_score(m::fAnoGAN, x, κ=1f0)
	x̂, fx, fx̂ = izi(m, x)
	xx̂= vec(Flux.sum((x̂ .- x).^2, dims=[1,2,3]))
	fxfx̂ = vec(Flux.sum((fx̂ .- fx).^2, dims=1))
	return xx̂ .+ κ .* fxfx̂
end

"""
	function anomaly_score_gpu(m::fAnoGAN, real_input, κ = 1.0; batch_size=64, to_testmode::Bool=true)
Function which compute anomaly score on GPU.
"""
function anomaly_score_gpu(m::fAnoGAN, real_input, κ = 1f0; batch_size=64, to_testmode::Bool=true)
	real_input = Flux.Data.DataLoader(real_input, batchsize=batch_size)
	(to_testmode == true) ? Flux.testmode!(m) : nothing
	m = m |> gpu
	scores = Array{Float32}([])
	for x in real_input
		x = x |> gpu
		score = anomaly_score(m, x, κ)
		scores = cat(scores, score |> cpu, dims=1)
	end
	(to_testmode == true) ? Flux.testmode!(m, false) : nothing
	return scores
end

# loss functions for Wassersten GAN
wdloss(gen, dis, x, z) = -(Flux.mean(dis(x)) - Flux.mean(dis(gen(z))))
wdloss(model::fAnoGAN, x) = wdloss(model.generator, model.discriminator, x, rand(model.prior, size(x, ndims(x))))
wdloss(model::fAnoGAN, x, z) = wdloss(model.generator, model.discriminator, x, z)

wgloss(gen, dis, z) = - Flux.mean(dis(gen(z)))
wgloss(model::fAnoGAN, z) = wgloss(model.generator, model.discriminator, z)

function fanogan_constructor(
		;idim=(2,2,1),
		zdim::Int=1,
		activation="relu",
		hdim=1024,
		kernelsizes=(1,1),
		channels=(1,1),
		scalings=(1,1),
		init_seed=nothing,
		batchnorm=false,
		kwargs...
	)

	(init_seed !== nothing) ? Random.seed!(init_seed) : nothing

	generator = conv_decoder(idim, zdim, reverse(kernelsizes), reverse(channels),
		reverse(scalings); activation=activation, batchnorm=batchnorm)

	discriminator = Chain(conv_encoder(idim, 1, kernelsizes, channels, scalings,
			activation=activation, batchnorm=batchnorm)...)

	encoder_ =  Chain(conv_encoder(idim, zdim, kernelsizes, channels, scalings,
	activation=activation, batchnorm=batchnorm)..., x->tanh.(x))

	(init_seed !== nothing) ? Random.seed!() : nothing

	model = fAnoGAN(MvNormal(zdim, 1f0), generator, discriminator, encoder_)
end

"""
	function StatsBase.fit!(model::fAnoGAN, data::Tuple, params)

General function for fitting of fAnoGAN.
	model::fAnoGAN
	data::Tupe   		... ((x_train, y_train), (x_valid, y_valid), (x_test, y_test))
	params::NamedTuple  ... all needed parameters for training
		kappa  		 		... scaling coeficent for izi_f loss function
		weight_clip  		... weight clipping parameter (https://arxiv.org/abs/1701.07875)
		lr_gan		 		... learning rate for WGAN
		lr_enc		 		... learning rate for encoder (image-z-iamge_f)
		batch_size::Int 	... batch size
		check_every::Int  	... number of iterations between EarlyStopping evaluation check
		patience::Int		... patience before stoping training
		iters::Int     		... max number of iteration
		mtt::Int			... max train time (if training is too slow)
		n_critic::Int   	... number of discriminator updater per one generator update
		usegpu::Bool    	... to use GPU or not

	Example:
		params = (kappa = 1f0, weight_clip=0.1, lr_gan = 0.001, lr_enc = 0.001, batch_size=128,
			patience=10, check_every=30, iters=10000, mtt=82800, n_critic=1, usegpu=true)

"""
function StatsBase.fit!(model::fAnoGAN, data::Tuple, params::NamedTuple)

	model = params.usegpu ? model |>gpu : model |>cpu
	best_model = deepcopy(model)
	history = MVHistory()
	train_loader, val_loader = prepare_dataloaders(data, params)
	val_batches = length(val_loader)
	"""
		WGAN training
	"""

	opt_G = RMSProp(params.lr_gan)
	opt_D = RMSProp(params.lr_gan)
	opt_E = ADAM(params.lr_enc)
	ps_E = Flux.params(model.encoder)
	ps_D = Flux.params(model.discriminator) #model.gan.discriminator
	ps_G = Flux.params(model.generator) # model.gan.generator
	# no early stopping here <– bahavior of loss function will not allow it
	progress = Progress(length(train_loader))
	for (iter, X) in enumerate(train_loader)
		X = params.usegpu ? gpu(getobs(X)) : getobs(X)
		last_loss_D = 0
		for n = 1:params.n_critic
			z = rand(model.prior, size(X, ndims(X)))
			z = params.usegpu ? z|>gpu : z
			loss_D, back_D = Flux.pullback(ps_D) do
				wdloss(model, X, z)
			end
			grad_D = back_D(1f0)
			Flux.Optimise.update!(opt_D, ps_D, grad_D)
			clip_weights!(ps_D, params.weight_clip)
			last_loss_D = loss_D
		end
		z = rand(model.prior, size(X, ndims(X)))
		z = params.usegpu ? z|>gpu : z
		loss_G, back_G = Flux.pullback(ps_G) do
				wgloss(model, z)
		end
		grad_G = back_G(1f0)
		Flux.Optimise.update!(opt_G, ps_G, grad_G)
		push!(history, :loss_G, iter, loss_G)
		push!(history, :loss_D, iter, last_loss_D)
		next!(progress; showvalues=[
			(:iters, "$(iter)/$(params.iters)"),
			(:loss_G, loss_G),
			(:loss_D, last_loss_D)
			])
		if mod(iter, params.check_every) == 0
			val_loss_G = 0
			val_loss_D = 0
			Flux.testmode!(model)
			for X_val in val_loader
				X_val = params.usegpu ? X_val |> gpu : X_val
				z = rand(model.prior, size(X_val, ndims(X_val)))
				z = params.usegpu ? z|>gpu : z
				val_loss_G += wgloss(model, z)
				val_loss_D += wdloss(model, X_val, z)
			end
			Flux.testmode!(model, false)
			push!(history, :validation_D, iter, val_loss_D/val_batches)
			push!(history, :validation_G, iter, val_loss_G/val_batches)
		end
	end

	"""
		Encoder training
	"""
	# possible memory problems
	CUDA.reclaim() # somehow GPU will run out of memory. This helps to clear cached memory and continue with it

	best_model = deepcopy(model)
	# early stopping
	best_val_loss = Inf
	patience_ = deepcopy(params.patience)
	progress = Progress(length(train_loader))

	start_time = time()
	for (iter, X) in enumerate(train_loader)
		X = params.usegpu ? gpu(getobs(X)) : getobs(X)
		loss_E, back_E = Flux.pullback(ps_E) do
			izif_loss(model, X, params.kappa)
		end
		grad_E = back_E(1f0)
		Flux.Optimise.update!(opt_E, ps_E, grad_E)

		push!(history, :loss_E, iter, loss_E)
		next!(progress; showvalues=[
			(:iters, "$(iter)/$(params.iters)"),
			(:loss_izif, loss_E)
			])

		if mod(iter, params.check_every) == 0
			val_loss_E = 0
			Flux.testmode!(model)
			for X_val in val_loader
				X_val = params.usegpu ? X_val |> gpu : X_val
				val_loss_E += izif_loss(model, X_val, params.kappa)
			end
			Flux.testmode!(model, false)
			push!(history, :validation_E, iter, val_loss_E/val_batches)
			if val_loss_E < best_val_loss
				best_val_loss = val_loss_E
				patience_ = params.patience
				best_model = deepcopy(model)
			else
				patience_ -= 1
				if patience_ == 0
					@info "Stopped training of Autoencoder after $(iter) iterations"
					break
				end
			end
		end
		if time() - start_time > params.mtt # stop early if time is running out
			best_model = deepcopy(model)
			@info "Stopped training after $(i) iterations, $((time() - start_time)/3600) hours."
			break
		end
	end
	return best_model, history, sum(map(p->length(p), Flux.params(model)))
end


#=
ArithematicAnalysis:
- Julia version: 1.4.2
- Author: akhil
- Date: 2020-05-28
=#
module ArithematicAnalysis
    include("Bisection.jl")
    include("GaussianElimination.jl")
    include("StaticEquilibrium.jl")
end


###############################################################################
##### Theta functions
###############################################################################

struct Theta
    nome        ::ComplexF64
    nome_powers ::Array{ComplexF64, 2}
    powers_num  ::Int64
    ϵ           ::Float64
end

function Theta(ω1::ComplexF64,
               ω3::ComplexF64,
               ϵ ::Float64)
    q = nome(ω1, ω3)
    nome_powers = precompute_nome_powers(q, ϵ, ω1+ω3)
    powers_num = size(nome_powers, 2)

    Theta(q, nome_powers, powers_num, ϵ)
end

"""
  nome(ω1::ComplexF64, ω3::ComplexF64)

Return the nome for the giving pair of lattice generators.
"""
function nome(
        ω1::ComplexF64,
        ω3::ComplexF64
        )
    #####
    τ = ω3 / ω1
    q = exp(1im * π * τ)
end

"""
  precompute_nome_powers(; q::ComplexF64, ϵ::Float64, z0::ComplexF64, max_n_terms ::Integer = 10)

Create array of precomputed nome powers for evaluating theta functions with given tolerance.

# RETURN:
- Array{ComplexF64, 2}(2, n)
"""
function precompute_nome_powers(
        # nome
        q           ::ComplexF64,
        # tolerance
        ϵ           ::Float64,
        # point for tolerance checking
        z0          ::ComplexF64,
        # maximum number of terms
        max_n_terms ::Integer = 10
        )
    #####
    n = 0
    while true
        # theta_1(z,q) = 2 sum_{n=0}^infty (-1)^n q^{(n+0.5)^2} sin((2n+1)z)
        dt1 = 2*(2n+1)*cos((2n+1)*z0)*q^((n+0.5)^2)
        # theta_2(z,q) = 2 sum_{n=0}^infty q^{(n+0.5)^2} cos((2n+1)z)
        dt2 = 2*(2n+1)*sin((2n+1)*z0)*q^((n+0.5)^2)
        if (abs(dt1) < ϵ) && (abs(dt1) < ϵ)
            break
        end
        n += 1
        if n > max_n_terms
            error("Series don't converge")
        end
    end
    nome_powers = Array{ComplexF64, 2}(undef, 2, n+1)
    for k in 0 : n
        nome_powers[1, k+1] = q^((k+0.5)^2)
        nome_powers[2, k+1] = nome_powers[1, k+1] * (2k+1)
    end

    return nome_powers
end

@doc """
  theta(th::Theta, z::ComplexF64)

Compute θ_1(z), θ_2(z), θ_1'(z), θ_2'(z) with given precomputed array of nome powers.
"""
function theta(
        th ::Theta,
        z  ::ComplexF64
        )
    #####
    t1  :: ComplexF64 = 0
    t2  :: ComplexF64 = 0
    dt1 :: ComplexF64 = 0
    dt2 :: ComplexF64 = 0
    sign = 1.0

    N = th.powers_num-1
    for n in 0 : N
        z_n = (2n+1)*z
        sin_z = sin(z_n)
        cos_z = cos(z_n)
        # theta_1(z,q) = 2 sum_{n=0}^N (-1)^n q^{(n+0.5)^2} sin((2n+1)z)
        t1 += 2 * sign * th.nome_powers[1, n+1] * sin_z
        # theta_2(z,q) = 2 sum_{n=0}^N q^{(n+0.5)^2} cos((2n+1)z)
        t2 += 2 * th.nome_powers[1, n+1] * cos_z
        # theta_1'(z,q) = 2 sum_{n=0}^N (-1)^n (2n+1) q^{(n+0.5)^2} cos((2n+1)z))
        dt1 += 2 * sign * th.nome_powers[2, n+1] * cos_z
        # theta_2'(z,q) = -2 sum_{n=0}^N (2n+1) q^{(n+0.5)^2} sin((2n+1)z)
        dt2 += -2 * th.nome_powers[2, n+1] * sin_z

        sign = -sign
    end

    return t1, t2, dt1, dt2
end

""""
  theta(θ::Theta; th_k::Integer, d_n::Integer,z::ComplexF64)

Compute ``\\frac{d^n}{dz^n} \\theta_k(z, q)`` with tolerance `ϵ`.
"""
function theta(
        # Theta data struct
        th ::Theta;
        # theta-function number
        th_k ::Integer,
        # derivative number
        d_n  ::Integer,
        # point for computation
        z    ::ComplexF64,
        )
   #####
    q = th.nome
    ϵ = th.ϵ

   if (th_k < 1) || (th_k > 4)
       error("Invalid theta function number")
   end

   rem_d = d_n % 4
   d_sin_sign = (rem_d == 0) || (rem_d == 1) ? 1 : -1
   d_cos_sign = (rem_d == 0) || (rem_d == 3) ? 1 : -1

   res :: ComplexF64 = 0
   if (th_k == 1) || (th_k == 2)
       res = 0
       n = 0
   else
       n = 1
       if d_n == 0
           res = 1
       else
           res = 0
       end
   end

   sign = (-1)^n
   while true
       # theta_1(z,q) = 2 sum_{n=0}^infty (-1)^n q^{(n+0.5)^2} sin((2n+1)z)
       if th_k == 1
           term = 2 * sign * (2n+1)^d_n * q^((n+0.5)^2) * d_sin_sign
           if d_n % 2 == 0
               term *= sin((2n+1)*z)
           else
               term *= cos((2n+1)*z)
           end
       end
       # theta_2(z,q) = 2 sum_{n=0}^infty q^{(n+0.5)^2} cos((2n+1)z)
       if th_k == 2
           term = 2 * (2n+1)^d_n * q^((n+0.5)^2) * d_cos_sign
           if d_n % 2 == 0
               term *= cos((2n+1)*z)
           else
               term *= sin((2n+1)*z)
           end
       end
       # theta_3(z,q) = 1 + 2 sum_{n=1}^infty q^{n^2} cos(2nz)
       if th_k == 3
           term = 2 * (2n)^d_n * q^(n^2) * d_cos_sign
           if d_n % 2 == 0
               term *= cos(2n*z)
           else
               term *= sin(2n*z)
           end
       end
       # theta_4(z,q) = 1 + 2 sum_{n=1}^infty (-1)^n q^{n^2} cos(2nz)
       if th_k == 4
           term = 2 * sign * (2n)^d_n * q^(n^2) * d_cos_sign
           if d_n % 2 == 0
               term *= cos(2n*z)
           else
               term *= sin(2n*z)
           end
       end

       res += term
       n += 1
       sign = -sign

       if abs(term) < ϵ
           break
       end
   end

   return res
end


"""
   interp_chunk_to_grid_gp_brute_force!( flux_out, var_out, chunk_of_spectrum, wavelengths )
Return spectra interpolated onto a grid of points using Gaussian Process interpolation.
# Arguments:
- flux_out: (results stored into this array)
- var_out: (results stored into this array)
- chunk_of_spectrum
- wavelengths: AbstractRange or AbstractArray of locations where chunk is to be interpolated to
# Returns
- flux_out

NOTE:  Using own GP code for now, since include predicting derivatives and can minimize unnecessary dependancies.  We may need to
revisit this if we want improved speed
"""

function interp_chunk_to_shifted_grid_gp_brute_force!( flux_out::AA1, var_out::AA2, chunk::AC, grid::AR, boost_factor::Real ) where { T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}, AC<:AbstractChunkOfSpectrum, AR<:Union{AbstractRange,AbstractArray{T2,1}} }
    @assert size(flux_out) == size(var_out)
    @assert size(flux_out) == size(grid)
    rho = 2*5000/speed_of_light_mps * mean(chunk.λ)
    sigmasq_kernel = 0.25 # Float64(mean(chunk.var))
    println(" rho = ", rho, "  σ²_kernel = ", sigmasq_kernel)
    flux_out .= GPInterpolation.predict_mean(chunk.λ./boost_factor, chunk.flux, grid, sigmasq_obs = chunk.var, kernel = GPs.matern52_sparse_kernel, rho=rho, sigmasq_cor=sigmasq_kernel ) # 	sigmasq_cor=1.0, rho=1
    # TODO: Update var_out to actually use the right GP or at least do something more sensible
    #var_out .= GPInterpolation.predict_mean(chunk.λ, chunk.var, grid, sigmasq_obs = chunk.var, kernel = GPs.matern52_sparse_kernel, rho=5000/speed_of_light_mps) # 	sigmasq_cor=1.0, rho=1
    return flux_out
end

function interp_chunk_to_shifted_grid_gp_brute_force( chunk::AC, grid::AR, boost_factor::Real ) where {  AC<:AbstractChunkOfSpectrum, T2<:Real, AR<:Union{AbstractRange,AbstractArray{T2,1}} }
    flux_out = Array{Float64,1}(undef,length(grid))
    var_out = Array{Float64,1}(undef,length(grid))
    interp_chunk_to_grid_gp_brute_force!(flux_out, var_out, chunk, grid, boost_factor=boost_factor)
    return (flux=flux_out, var=var_out)
end

function interp_chunk_to_grid_gp_brute_force!( flux_out::AA1, var_out::AA2, chunk::AC, grid::AR ) where { T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}, AC<:AbstractChunkOfSpectrum, AR<:Union{AbstractRange,AbstractArray{T2,1}} }
    @assert size(flux_out) == size(var_out)
    @assert size(flux_out) == size(grid)
    rho = 2*5000/speed_of_light_mps * mean(chunk.λ)
    sigmasq_kernel = 0.25 # Float64(mean(chunk.var))
    println(" rho = ", rho, "  σ²_kernel = ", sigmasq_kernel)
    flux_out .= GPInterpolation.predict_mean(chunk.λ, chunk.flux, grid, sigmasq_obs = chunk.var, kernel = GPs.matern52_sparse_kernel, rho=rho, sigmasq_cor=sigmasq_kernel ) # 	sigmasq_cor=1.0, rho=1
    # TODO: Update var_out to actually use the right GP or at least do something more sensible
    #var_out .= GPInterpolation.predict_mean(chunk.λ, chunk.var, grid, sigmasq_obs = chunk.var, kernel = GPs.matern52_sparse_kernel, rho=5000/speed_of_light_mps) # 	sigmasq_cor=1.0, rho=1
    return flux_out
end

function interp_chunk_to_grid_gp_brute_force( chunk::AC, grid::AR ) where {  AC<:AbstractChunkOfSpectrum, T2<:Real, AR<:Union{AbstractRange,AbstractArray{T2,1}} }
    flux_out = Array{Float64,1}(undef,length(grid))
    var_out = Array{Float64,1}(undef,length(grid))
    interp_chunk_to_grid_gp_brute_force!(flux_out, var_out, chunk, grid, )
    return (flux=flux_out, var=var_out)
end


using InteractiveChaos, Makie

ds = Systems.henonheiles()

potential(x, y) = 0.5(x^2 + y^2) + (x^2*y - (y^3)/3)
energy(x,y,px,py) = 0.5(px^2 + py^2) + potential(x,y)
const E = energy(get_state(ds)...)

function complete(y, py, x)
    V = potential(x, y)
    Ky = 0.5*(py^2)
    Ky + V ≥ E && error("Point has more energy!")
    px = sqrt(2(E - V - Ky))
    ic = [x, y, px, py]
    return ic
end

chaotic = get_state(ds)
stable = [0., 0.1, 0.5, 0.]

plane = (1, 0.0)

#does this work?

psos = interactive_poincaresos(ds, plane, (2, 4), complete; markersizes = (-5, -1))