using MathProgBase
function MathProgBase.initialize(d::SSMOptim, requested_features::Vector{Symbol})
    for feat in requested_features
        if !(feat in [:Grad, :Jac])
            error("Unsupported feature $feat")
        end
    end
end

MathProgBase.features_available(d::SSMOptim) = [:Grad, :Jac]

function MathProgBase.eval_f(d::SSMOptim, x)
    ssm_obj_fun(d, x)
end

function MathProgBase.eval_g(d::SSMOptim, g, x)
end

function MathProgBase.eval_grad_f(d::SSMOptim, grad_f, x)
    ForwardDiff.gradient!(grad_f, y->ssm_obj_fun(d, y), x, d.d.jcfg, Val{false}())
end

MathProgBase.jac_structure(d::SSMOptim) = Int[],Int[]
# lower triangle only
MathProgBase.hesslag_structure(d::SSMOptim) = Int[]

function MathProgBase.eval_jac_g(d::SSMOptim, J, x)
end

function MathProgBase.eval_hesslag(d::SSMOptim, H, x, σ, μ)
end


using LinearAlgebra
using Parameters

@testset "jj2lsj" begin
    @testset "Transform matrix" begin
        R = jj2lsj(ros"1[s] 2[s-p] k[s-d]"...)
        # Since it is a rotation matrix, its inverse is simply the
        # transpose.
        @test norm(inv(Matrix(R)) - R') < 10eps()
    end
    @testset "Transform individual spin-orbitals" begin
        for o in [o"5s", o"5p", o"5d"]
            for so in spin_orbitals(o)
                ℓ = so.orb.ℓ
                @unpack mℓ,spin = so
                mj = mℓ + (spin ? 1 : -1)//2
                pure = abs(mj) == ℓ + 1//2
                linear_combination = jj2lsj(so)
                @test length(linear_combination) == (pure ? 1 : 2)
                @test norm(last.(linear_combination)) ≈ 1
                @test all(mj .== last.(first.(linear_combination)))
            end
        end
    end
end


using GlobalSearchRegression, DataFrames, Distributions, CSV
# Insert here your working directory path
data = DataFrame(Array{Union{Missing,Float64}}(randn(1000,21)))
headers = [ :y ; [ Symbol("x$i") for i = 1:size(data,2) - 1 ] ]
names!(data, headers )
b=@elapsed gsreg("y x*", data; ttest=true, criteria=[:aic], parallel=16, csv="20x1000x16.csv")
t=DataFrame(A=Float64[])
push!(t,b)
CSV.write("juliatimmings.csv", t; append=true)
exit()


@testset "MPolyGcd" begin
  QA, A = PolynomialRing(Hecke.QQ, "A")
  K, a = number_field(A^3 - 2, "a")
  Kx, (x,) = PolynomialRing(K, ["x"])
  g = gcd(-715//2*a^2*x^3 - 715//2*a^2*x^2 - 365525875//2*a^2*x - 365525875//2*a^2, (-715//2*a^2 + 511225)*x^2 - 365525875//2*a^2 + 261351000625)
  @test g == x^2 + 511225
end


using Test
using SparseArrays
using LinearAlgebra
import sbp
using sbp.Sparse: diag_block_matrix_2x2
using sbp.Grid: grid_xp, grid_xm, grid_2d_x, grid_2d_y

st = sbp.StaggeredAcoustic
sop = sbp.OP2019

nx=10
ny=20

function mapping(nx, ny)
        x1, h = grid_xp(nx)
        y1, h = grid_xp(ny)

        p = grid_2d_x(x1, ny)
        q = grid_2d_y(y1, nx)
        x = p
        y = (1 .+ 0.1 * sin.(p)).*q
        return x, y
end

fx, fy = mapping(nx, ny)
desc = sop.load_description()
builder = n -> sop.build_operators(desc, n)
ops = st.init_operators(nx, ny, builder, fx, fy)

HJp = st.pressure_norm(ops)
HJ = st.velocity_norm(ops)
Gp = st.contravariant_metric_tensor(ops)
Ek = HJ*Gp
Ap = st.divergence(ops)
Av = st.gradient(ops, Gp)
S = st.pressure_sat(ops, Gp)
H = st.energy_norm(HJp, HJ, Gp)
A = st.spatial_discretization(Ap, Av - S)
d = st.pressure_source(ops, 0.5, 0.5, 4, 4)

@testset "Grids" begin
        xp, yp = st.grids("p", nx, ny)
        @test size(xp,1) == (nx + 1) * (ny + 1)
        @test size(yp,1) == (nx + 1) * (ny + 1)
        x1, y1 = st.grids("v1", nx, ny)
        @test size(x1,1) == (nx + 0) * (ny + 1)
        @test size(y1,1) == (nx + 0) * (ny + 1)
        x2, y2 = st.grids("v2", nx, ny)
        @test size(x2,1) == (nx + 1) * (ny + 0)
        @test size(y2,1) == (nx + 1) * (ny + 0)
        xn, yn = st.grids("node", nx, ny)
        @test size(xn,1) == (nx + 0) * (ny + 0)
        @test size(yn,1) == (nx + 0) * (ny + 0)
end

@testset "Kinetic energy tensor" begin
        # Check symmetry
        @test isapprox(Ek, (Ek)')
        # Check positivity
        lam = real(eigvals(Matrix(Ek)))
        @test minimum(lam) > 0

end

@testset "Stability" begin
        # Check that the discretization is energy conservative to rounding error
        lam = real(eigvals(Matrix(A)))
        @test minimum(lam) > -1e-13
        @test maximum(lam) < 1e-13

        E = H*A + (H*A)'
        lam = real(eigvals(Matrix(E)))
        @test minimum(lam) > -1e-13
        @test maximum(lam) < 1e-13

end


using ReservoirComputing, DifferentialEquations, Statistics, Random

u0 = [1.0,0.0,0.0]
tspan = (0.0,1000.0)
datasize = 100000
tsteps = range(tspan[1], tspan[2], length = datasize)

function lorenz(du,u,p,t)
    p = [10.0,28.0,8/3]
    du[1] = p[1]*(u[2]-u[1])
    du[2] = u[1]*(p[2]-u[3]) - u[2]
    du[3] = u[1]*u[2] - p[3]*u[3]
end

function prior_model_data_generator(u0, tspan, tsteps, model = lorenz)
    prob = ODEProblem(lorenz, u0, tspan)
    sol = Array(solve(prob, saveat = tsteps))
    return sol
end

train_len = 10000

ode_prob = ODEProblem(lorenz, u0, tspan)
ode_sol = solve(ode_prob, saveat = tsteps)
ode_data =Array(ode_sol)
input_data = ode_data[:, 1:train_len]
target_data = ode_data[:, 2:train_len+1]

test_data = ode_data[:, train_len+1:end][:, 1:1000]
predict_len = size(test_data, 2)
tspan_train = (tspan[1], ode_sol.t[train_len])

hybrid = Hybrid(prior_model_data_generator, u0, tspan_train, train_len)

Random.seed!(77)
esn = ESN(input_data,
    reservoir = RandSparseReservoir(300),
    variation = hybrid)

output_layer = train(esn, target_data, StandardRidge(0.3))

output = esn(Generative(predict_len), output_layer)

@test mean(abs.(test_data[1:100] .- output[1:100])) ./ mean(abs.(test_data[1:100])) < 0.11


naolaplacian00(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (2*ud + r*ud2)/(2*sqrt(π)*r)
naolaplacian1m1(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(3/π)*sin(θ)*sin(ϕ)*(-2*u + r*(2*ud + r*ud2)))/(2*r^2)
naolaplacian10(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(3/π)*cos(θ)*(-2*u + r*(2*ud + r*ud2)))/(2*r^2)
naolaplacian11(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(3/π)*cos(ϕ)*sin(θ)*(-2*u + r*(2*ud + r*ud2)))/(2*r^2)
naolaplacian2m2(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(15/π)*sin(θ)^2*sin(2*ϕ)*(-6*u + r*(2*ud + r*ud2)))/(4*r^2))
naolaplacian2m1(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(15/π)*sin(2*θ)*sin(ϕ)*(-6*u + r*(2*ud + r*ud2)))/(4*r^2)
naolaplacian20(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(5/π)*(1 + 3*cos(2*θ))*(-6*u + r*(2*ud + r*ud2)))/(8*r^2)
naolaplacian21(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(15/π)*cos(ϕ)*sin(2*θ)*(-6*u + r*(2*ud + r*ud2)))/(4*r^2)
naolaplacian22(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(15/π)*cos(2*ϕ)*sin(θ)^2*(-6*u + r*(2*ud + r*ud2)))/(4*r^2)
naolaplacian3m3(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(35/(2*π))*sin(θ)^3*sin(3*ϕ)*(-12*u + r*(2*ud + r*ud2)))/(4*r^2)
naolaplacian3m2(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(105/π)*cos(θ)*sin(θ)^2*sin(2*ϕ)*(-12*u + r*(2*ud + r*ud2)))/(4*r^2))
naolaplacian3m1(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(21/(2*π))*(3 + 5*cos(2*θ))*sin(θ)*sin(ϕ)*(-12*u + r*(2*ud + r*ud2)))/(8*r^2)
naolaplacian30(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(7/π)*(3*cos(θ) + 5*cos(3*θ))*(-12*u + r*(2*ud + r*ud2)))/(16*r^2)
naolaplacian31(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(21/(2*π))*(3 + 5*cos(2*θ))*cos(ϕ)*sin(θ)*(-12*u + r*(2*ud + r*ud2)))/(8*r^2)
naolaplacian32(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(105/π)*cos(θ)*cos(2*ϕ)*sin(θ)^2*(-12*u + r*(2*ud + r*ud2)))/(4*r^2)
naolaplacian33(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(35/(2*π))*cos(3*ϕ)*sin(θ)^3*(-12*u + r*(2*ud + r*ud2)))/(4*r^2)
naolaplacian4m4(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(35/π)*sin(θ)^4*sin(4*ϕ)*(-20*u + r*(2*ud + r*ud2)))/(16*r^2))
naolaplacian4m3(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/(2*π))*cos(θ)*sin(θ)^3*sin(3*ϕ)*(-20*u + r*(2*ud + r*ud2)))/(4*r^2)
naolaplacian4m2(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(5/π)*(5 + 7*cos(2*θ))*sin(θ)^2*sin(2*ϕ)*(-20*u + r*(2*ud + r*ud2)))/(16*r^2))
naolaplacian4m1(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(5/(2*π))*(2*sin(2*θ) + 7*sin(4*θ))*sin(ϕ)*(-20*u + r*(2*ud + r*ud2)))/(32*r^2)
naolaplacian40(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*(9 + 20*cos(2*θ) + 35*cos(4*θ))*(-20*u + r*(2*ud + r*ud2)))/(128*sqrt(π)*r^2)
naolaplacian41(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(5/(2*π))*cos(ϕ)*(2*sin(2*θ) + 7*sin(4*θ))*(-20*u + r*(2*ud + r*ud2)))/(32*r^2)
naolaplacian42(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(5/π)*(5 + 7*cos(2*θ))*cos(2*ϕ)*sin(θ)^2*(-20*u + r*(2*ud + r*ud2)))/(16*r^2)
naolaplacian43(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/(2*π))*cos(θ)*cos(3*ϕ)*sin(θ)^3*(-20*u + r*(2*ud + r*ud2)))/(4*r^2)
naolaplacian44(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/π)*cos(4*ϕ)*sin(θ)^4*(-20*u + r*(2*ud + r*ud2)))/(16*r^2)
naolaplacian5m5(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(77/(2*π))*sin(θ)^5*sin(5*ϕ)*(-30*u + r*(2*ud + r*ud2)))/(16*r^2)
naolaplacian5m4(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(385/π)*cos(θ)*sin(θ)^4*sin(4*ϕ)*(-30*u + r*(2*ud + r*ud2)))/(16*r^2))
naolaplacian5m3(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(385/(2*π))*(7 + 9*cos(2*θ))*sin(θ)^3*sin(3*ϕ)*(-30*u + r*(2*ud + r*ud2)))/(32*r^2)
naolaplacian5m2(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(1155/π)*(5*cos(θ) + 3*cos(3*θ))*sin(θ)^2*sin(2*ϕ)*(-30*u + r*(2*ud + r*ud2)))/(32*r^2))
naolaplacian5m1(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(165/π)*(15 + 28*cos(2*θ) + 21*cos(4*θ))*sin(θ)*sin(ϕ)*(-30*u + r*(2*ud + r*ud2)))/(128*r^2)
naolaplacian50(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(11/π)*(30*cos(θ) + 35*cos(3*θ) + 63*cos(5*θ))*(-30*u + r*(2*ud + r*ud2)))/(256*r^2)
naolaplacian51(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(165/π)*(15 + 28*cos(2*θ) + 21*cos(4*θ))*cos(ϕ)*sin(θ)*(-30*u + r*(2*ud + r*ud2)))/(128*r^2)
naolaplacian52(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(1155/π)*(5*cos(θ) + 3*cos(3*θ))*cos(2*ϕ)*sin(θ)^2*(-30*u + r*(2*ud + r*ud2)))/(32*r^2)
naolaplacian53(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(385/(2*π))*(7 + 9*cos(2*θ))*cos(3*ϕ)*sin(θ)^3*(-30*u + r*(2*ud + r*ud2)))/(32*r^2)
naolaplacian54(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(385/π)*cos(θ)*cos(4*ϕ)*sin(θ)^4*(-30*u + r*(2*ud + r*ud2)))/(16*r^2)
naolaplacian55(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(77/(2*π))*cos(5*ϕ)*sin(θ)^5*(-30*u + r*(2*ud + r*ud2)))/(16*r^2)
naolaplacian6m6(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(3003/(2*π))*sin(θ)^6*sin(6*ϕ)*(-42*u + r*(2*ud + r*ud2)))/(32*r^2))
naolaplacian6m5(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(1001/(2*π))*cos(θ)*sin(θ)^5*sin(5*ϕ)*(-42*u + r*(2*ud + r*ud2)))/(16*r^2)
naolaplacian6m4(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(91/π)*(9 + 11*cos(2*θ))*sin(θ)^4*sin(4*ϕ)*(-42*u + r*(2*ud + r*ud2)))/(64*r^2))
naolaplacian6m3(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(1365/(2*π))*(21*cos(θ) + 11*cos(3*θ))*sin(θ)^3*sin(3*ϕ)*(-42*u + r*(2*ud + r*ud2)))/(64*r^2)
naolaplacian6m2(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(1365/(2*π))*(35 + 60*cos(2*θ) + 33*cos(4*θ))*sin(θ)^2*sin(2*ϕ)*(-42*u + r*(2*ud + r*ud2)))/(256*r^2))
naolaplacian6m1(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(273/π)*(50*cos(θ) + 45*cos(3*θ) + 33*cos(5*θ))*sin(θ)*sin(ϕ)*(-42*u + r*(2*ud + r*ud2)))/(256*r^2)
naolaplacian60(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(13/π)*(50 + 105*cos(2*θ) + 126*cos(4*θ) + 231*cos(6*θ))*(-42*u + r*(2*ud + r*ud2)))/(1024*r^2)
naolaplacian61(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(273/π)*(50*cos(θ) + 45*cos(3*θ) + 33*cos(5*θ))*cos(ϕ)*sin(θ)*(-42*u + r*(2*ud + r*ud2)))/(256*r^2)
naolaplacian62(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(1365/(2*π))*(35 + 60*cos(2*θ) + 33*cos(4*θ))*cos(2*ϕ)*sin(θ)^2*(-42*u + r*(2*ud + r*ud2)))/(256*r^2)
naolaplacian63(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(1365/(2*π))*(21*cos(θ) + 11*cos(3*θ))*cos(3*ϕ)*sin(θ)^3*(-42*u + r*(2*ud + r*ud2)))/(64*r^2)
naolaplacian64(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(91/π)*(9 + 11*cos(2*θ))*cos(4*ϕ)*sin(θ)^4*(-42*u + r*(2*ud + r*ud2)))/(64*r^2)
naolaplacian65(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(1001/(2*π))*cos(θ)*cos(5*ϕ)*sin(θ)^5*(-42*u + r*(2*ud + r*ud2)))/(16*r^2)
naolaplacian66(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(3003/(2*π))*cos(6*ϕ)*sin(θ)^6*(-42*u + r*(2*ud + r*ud2)))/(32*r^2)
naolaplacian7m7(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(715/π)*sin(θ)^7*sin(7*ϕ)*(-56*u + r*(2*ud + r*ud2)))/(64*r^2)
naolaplacian7m6(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(5005/(2*π))*cos(θ)*sin(θ)^6*sin(6*ϕ)*(-56*u + r*(2*ud + r*ud2)))/(32*r^2))
naolaplacian7m5(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(385/π)*(11 + 13*cos(2*θ))*sin(θ)^5*sin(5*ϕ)*(-56*u + r*(2*ud + r*ud2)))/(128*r^2)
naolaplacian7m4(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(385/π)*(27*cos(θ) + 13*cos(3*θ))*sin(θ)^4*sin(4*ϕ)*(-56*u + r*(2*ud + r*ud2)))/(128*r^2))
naolaplacian7m3(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/π)*(189 + 308*cos(2*θ) + 143*cos(4*θ))*sin(θ)^3*sin(3*ϕ)*(-56*u + r*(2*ud + r*ud2)))/(512*r^2)
naolaplacian7m2(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(35/(2*π))*(350*cos(θ) + 275*cos(3*θ) + 143*cos(5*θ))*sin(θ)^2*sin(2*ϕ)*(-56*u + r*(2*ud + r*ud2)))/(512*r^2))
naolaplacian7m1(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(105/π)*(350 + 675*cos(2*θ) + 594*cos(4*θ) + 429*cos(6*θ))*sin(θ)*sin(ϕ)*(-56*u + r*(2*ud + r*ud2)))/(2048*r^2)
naolaplacian70(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(15/π)*(175*cos(θ) + 189*cos(3*θ) + 231*cos(5*θ) + 429*cos(7*θ))*(-56*u + r*(2*ud + r*ud2)))/(2048*r^2)
naolaplacian71(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(105/π)*(350 + 675*cos(2*θ) + 594*cos(4*θ) + 429*cos(6*θ))*cos(ϕ)*sin(θ)*(-56*u + r*(2*ud + r*ud2)))/(2048*r^2)
naolaplacian72(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/(2*π))*(350*cos(θ) + 275*cos(3*θ) + 143*cos(5*θ))*cos(2*ϕ)*sin(θ)^2*(-56*u + r*(2*ud + r*ud2)))/(512*r^2)
naolaplacian73(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/π)*(189 + 308*cos(2*θ) + 143*cos(4*θ))*cos(3*ϕ)*sin(θ)^3*(-56*u + r*(2*ud + r*ud2)))/(512*r^2)
naolaplacian74(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(385/π)*(27*cos(θ) + 13*cos(3*θ))*cos(4*ϕ)*sin(θ)^4*(-56*u + r*(2*ud + r*ud2)))/(128*r^2)
naolaplacian75(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(385/π)*(11 + 13*cos(2*θ))*cos(5*ϕ)*sin(θ)^5*(-56*u + r*(2*ud + r*ud2)))/(128*r^2)
naolaplacian76(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(5005/(2*π))*cos(θ)*cos(6*ϕ)*sin(θ)^6*(-56*u + r*(2*ud + r*ud2)))/(32*r^2)
naolaplacian77(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(715/π)*cos(7*ϕ)*sin(θ)^7*(-56*u + r*(2*ud + r*ud2)))/(64*r^2)
naolaplacian8m8(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(12155/π)*sin(θ)^8*sin(8*ϕ)*(-72*u + r*(2*ud + r*ud2)))/(256*r^2))
naolaplacian8m7(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(12155/π)*cos(θ)*sin(θ)^7*sin(7*ϕ)*(-72*u + r*(2*ud + r*ud2)))/(64*r^2)
naolaplacian8m6(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(7293/(2*π))*(13 + 15*cos(2*θ))*sin(θ)^6*sin(6*ϕ)*(-72*u + r*(2*ud + r*ud2)))/(128*r^2))
naolaplacian8m5(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(17017/π)*(11*cos(θ) + 5*cos(3*θ))*sin(θ)^5*sin(5*ϕ)*(-72*u + r*(2*ud + r*ud2)))/(256*r^2)
naolaplacian8m4(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(1309/π)*(99 + 156*cos(2*θ) + 65*cos(4*θ))*sin(θ)^4*sin(4*ϕ)*(-72*u + r*(2*ud + r*ud2)))/
      (1024*r^2))
naolaplacian8m3(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(19635/π)*(126*cos(θ) + 91*cos(3*θ) + 39*cos(5*θ))*sin(θ)^3*sin(3*ϕ)*(-72*u + r*(2*ud + r*ud2)))/(1024*r^2)
naolaplacian8m2(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(595/(2*π))*(210 + 385*cos(2*θ) + 286*cos(4*θ) + 143*cos(6*θ))*sin(θ)^2*sin(2*ϕ)*(-72*u + r*(2*ud + r*ud2)))/(2048*r^2))
naolaplacian8m1(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(17/π)*(70*sin(2*θ) + 154*sin(4*θ) + 286*sin(6*θ) + 715*sin(8*θ))*sin(ϕ)*(-72*u + r*(2*ud + r*ud2)))/(8192*r^2)
naolaplacian80(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(17/π)*(1225 + 2520*cos(2*θ) + 2772*cos(4*θ) + 3432*cos(6*θ) + 6435*cos(8*θ))*(-72*u + r*(2*ud + r*ud2)))/(32768*r^2)
naolaplacian81(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(17/π)*cos(ϕ)*(70*sin(2*θ) + 154*sin(4*θ) + 286*sin(6*θ) + 715*sin(8*θ))*(-72*u + r*(2*ud + r*ud2)))/(8192*r^2)
naolaplacian82(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(595/(2*π))*(210 + 385*cos(2*θ) + 286*cos(4*θ) + 143*cos(6*θ))*cos(2*ϕ)*sin(θ)^2*(-72*u + r*(2*ud + r*ud2)))/(2048*r^2)
naolaplacian83(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(19635/π)*(126*cos(θ) + 91*cos(3*θ) + 39*cos(5*θ))*cos(3*ϕ)*sin(θ)^3*(-72*u + r*(2*ud + r*ud2)))/(1024*r^2)
naolaplacian84(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(1309/π)*(99 + 156*cos(2*θ) + 65*cos(4*θ))*cos(4*ϕ)*sin(θ)^4*(-72*u + r*(2*ud + r*ud2)))/(1024*r^2)
naolaplacian85(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(17017/π)*(11*cos(θ) + 5*cos(3*θ))*cos(5*ϕ)*sin(θ)^5*(-72*u + r*(2*ud + r*ud2)))/(256*r^2)
naolaplacian86(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(7293/(2*π))*(13 + 15*cos(2*θ))*cos(6*ϕ)*sin(θ)^6*(-72*u + r*(2*ud + r*ud2)))/(128*r^2)
naolaplacian87(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(12155/π)*cos(θ)*cos(7*ϕ)*sin(θ)^7*(-72*u + r*(2*ud + r*ud2)))/(64*r^2)
naolaplacian88(u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(12155/π)*cos(8*ϕ)*sin(θ)^8*(-72*u + r*(2*ud + r*ud2)))/(256*r^2)

# naolaplacian(l::Int, m::Int, u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = naolaplaciantab[l,m](u,ud,ud2,r,θ,ϕ)

naolaplacian(l::Int, m::Int, u::Float64, ud::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) =
    naolaplaciantab[i_lm(l,m)](u,ud,ud2,r,θ,ϕ)::Float64

naolaplaciantab = [
    naolaplacian00,
    naolaplacian1m1,
    naolaplacian10,
    naolaplacian11,
    naolaplacian2m2,
    naolaplacian2m1,
    naolaplacian20,
    naolaplacian21,
    naolaplacian22,
    naolaplacian3m3,
    naolaplacian3m2,
    naolaplacian3m1,
    naolaplacian30,
    naolaplacian31,
    naolaplacian32,
    naolaplacian33,
    naolaplacian4m4,
    naolaplacian4m3,
    naolaplacian4m2,
    naolaplacian4m1,
    naolaplacian40,
    naolaplacian41,
    naolaplacian42,
    naolaplacian43,
    naolaplacian44,
    naolaplacian5m5,
    naolaplacian5m4,
    naolaplacian5m3,
    naolaplacian5m2,
    naolaplacian5m1,
    naolaplacian50,
    naolaplacian51,
    naolaplacian52,
    naolaplacian53,
    naolaplacian54,
    naolaplacian55,
    naolaplacian6m6,
    naolaplacian6m5,
    naolaplacian6m4,
    naolaplacian6m3,
    naolaplacian6m2,
    naolaplacian6m1,
    naolaplacian60,
    naolaplacian61,
    naolaplacian62,
    naolaplacian63,
    naolaplacian64,
    naolaplacian65,
    naolaplacian66,
    naolaplacian7m7,
    naolaplacian7m6,
    naolaplacian7m5,
    naolaplacian7m4,
    naolaplacian7m3,
    naolaplacian7m2,
    naolaplacian7m1,
    naolaplacian70,
    naolaplacian71,
    naolaplacian72,
    naolaplacian73,
    naolaplacian74,
    naolaplacian75,
    naolaplacian76,
    naolaplacian77,
    naolaplacian8m8,
    naolaplacian8m7,
    naolaplacian8m6,
    naolaplacian8m5,
    naolaplacian8m4,
    naolaplacian8m3,
    naolaplacian8m2,
    naolaplacian8m1,
    naolaplacian80,
    naolaplacian81,
    naolaplacian82,
    naolaplacian83,
    naolaplacian84,
    naolaplacian85,
    naolaplacian86,
    naolaplacian87,
    naolaplacian88
]


"""
    constants
"""

const GSL_SQRT_DBL_EPSILON = 1.4901161193847656e-08
const GSL_LOG_DBL_MIN = -7.0839641853226408e02

const synch_c0 = π / sqrt(3.0)

const c01 = 0.2257913526447274323630976
const cond1 = 2 * sqrt(2.0) * GSL_SQRT_DBL_EPSILON
const cond3 = -8.0 * GSL_LOG_DBL_MIN / 7.0

const synchrotron1_data = [ 30.364682982501076273,
                            17.079395277408394574,
                            4.560132133545072889,
                            0.549281246730419979,
                            0.372976075069301172e-01,
                            0.161362430201041242e-02,
                            0.481916772120371e-04,
                            0.10512425288938e-05,
                            0.174638504670e-07,
                            0.22815486544e-09,
                            0.240443082e-11,
                            0.2086588e-13,
                            0.15167e-15
                            ]

const synchrotron1a_data = [ 2.1329305161355000985,
                             0.741352864954200240e-01,
                             0.86968099909964198e-02,
                             0.11703826248775692e-02,
                             0.1645105798619192e-03,
                             0.240201021420640e-04,
                             0.35827756389389e-05,
                             0.5447747626984e-06,
                             0.838802856196e-07,
                             0.13069882684e-07,
                             0.2053099071e-08,
                             0.325187537e-09,
                             0.517914041e-10,
                             0.83002988e-11,
                             0.13352728e-11,
                             0.2159150e-12,
                             0.349967e-13,
                             0.56994e-14,
                             0.9291e-15,
                             0.152e-15,
                             0.249e-16,
                             0.41e-17,
                             0.7e-18
                            ]

const synchrotron2_data = [  0.4490721623532660844,
                             0.898353677994187218e-01,
                             0.81044573772151290e-02,
                             0.4261716991089162e-03,
                             0.147609631270746e-04,
                             0.3628633615300e-06,
                             0.66634807498e-08,
                             0.949077166e-10,
                             0.1079125e-11,
                             0.10022e-13,
                             0.77e-16,
                             0.5e-18
                            ]

@inline @fastmath function cheb_eval(coeff::Vector{<:Real}, order::Integer, x::Real)

    d = 0.0
    dd = 0.0

    x2 = 2x

    @inbounds for j = order:-1:2
        temp = d
        d = x2 * d - dd + coeff[j]
        dd = temp
    end

    d = x * d - dd + 0.5 * coeff[1]

    return d
end

"""
    F(x::Real)

Computes the first synchrotron function at a given frequency ratio ``x = \\frac{\\nu}{\\nu_0}``.

``F(x) = x \\int_x^\\infty K_{\\frac{5}{3}}(t) dt``

"""
@inline function F(x::Real)

    if x < cond1

        z = x^(1/3)
        cf = 1 - 8.43812762813205e-01 * z * z
        return 2.14952824153447863671 * z * cf

    elseif x <= 4.0

        px = ∛(x)
        px11 = px*px*px*px*px*px*px*px*px*px*px
        t = x^2 / 8.0 - 1.0
        result_c1 = cheb_eval( synchrotron1_data, 13, t )

        result_c2 = cheb_eval( synchrotron2_data, 12, t )

        return px * result_c1 - px11 * result_c2 - synch_c0 * x

    elseif x < cond3

        t = (12.0 - x) / (x + 4.0)

        result_c1 = cheb_eval( synchrotron1a_data, 23, t )

        return sqrt(x) * result_c1 * exp(c01 - x)

    else
        return 0.0
    end
end

"""
    ℱ(x::Real)

Computes the first synchrotron function at a given frequency ratio ``x = \\frac{\\nu}{\\nu_0}``.

``F(x) = x \\int_x^\\infty K_{\\frac{5}{3}}(t) dt``

Wrapper for [`F`](@ref).
"""
ℱ(x::Real) = F(x)


#----------------------------------------------------------------------------------------------------------------------------------------
# Implements solution to the diffusion equation for the fluence in the steady-state and time-domain in a semi-infinite medium [1][2][3].
#
# [1] Alwin Kienle and Michael S. Patterson, "Improved solutions of the steady-state and the time-resolved diffusion equations
#     for reflectance from a semi-infinite turbid medium," J. Opt. Soc. Am. A 14, 246-254 (1997)
# [2] Martelli, Fabrizio, et al. Light propagation through biological tissue and other diffusive media: theory, solutions and software.
#     Vol. 10. No. 3.824746. Bellingham: SPIE press, 2010.
#----------------------------------------------------------------------------------------------------------------------------------------

""" Structure containing inputs for simulating the fluence under the diffusion approximation in the semi-infinite space."""
struct DAsemiinf_params{T <: AbstractFloat} <: DiffusionParameters
    ρ::T                                 # distance away from isotropic point source (cm)
    μa::T                                # absorption coefficient (cm⁻¹)
    μsp::T                               # reduced scattering coefficient (cm⁻¹)
    n_med::T                             # medium's index of refraction
    n_ext::T                             # external medium's index of refraction
    ω::T                                 # modulation frequency (1/ns)
    z::T                                 # z depth of detector within medium (cm)
    t::Union{T, AbstractVector{T}}       # time vector (ns)

    # do not need to provide these arguments (calculated from previous inputs)
    D::T                                 # Diffusion coefficient
    ν::T                                 # Speed of light in medium (cm/ns)
    z0::T                                # isotroptic source depth
    zb::T                                # Extrapolation length
    function DAsemiinf_params{T}(ρ::T, μa::T, μsp::T, n_med::T, n_ext::T, ω::T, z::T, t::Union{T, AbstractVector{T}}) where {T <: AbstractFloat}
        @assert ρ >= zero(T) "ρ must greater than or equal to 0"
        @assert μa >= zero(T) "μa must greater than or equal to 0"
        @assert μsp > zero(T) "μsp must greater than 0"
        @assert n_med > zero(T) "n_med must be positive"
        @assert n_ext > zero(T) "n_ext must be positive"
        @assert all(t .> zero(T)) "t must be positive"
        D = D_coeff(μsp)
        return new{T}(ρ, μa, μsp, n_med, n_ext, ω, z, t, D, ν_coeff(n_med), z0_coeff(μsp), zb_coeff(A_coeff(n_med / n_ext), D))
    end
end

"""
    Generator function for DAsemiinf_params structure.

Provides default parameters to use in the semi-infinite space fluence calculation in either the CW, FD, or TD.

# Keyword Arguments
- `ρ`: source-detector separation (cm⁻¹)
- `μa`: absorption coefficient (cm⁻¹)
- `μsp`: reduced scattering coefficient (cm⁻¹)
- `n_med`: medium's index of refraction
- `n_ext`: external medium's index of refraction (air or detector)
- `ω`: modulation frequency (1/ns)
- `z`: the z-depth orthogonal from the boundary (cm)
- `t`: the time vector (ns).

# Examples
```
julia> data = DAsemiinf_params() # return default parameters
julia> data = DAsemiinf_params(ρ = 1.5) # return ρ = 1.5 with the rest of the parameters given by defaults
julia> fluence_DA_semiinf_CW(data) # can then be used with corresponding functions
```
"""
function DAsemiinf_params(;
    ρ::T = 1.0,
    μa::T = 0.1,
    μsp::T = 10.0,
    n_med::T = 1.0,
    n_ext::T = 1.0,
    ω::T = 0.0,
    z::T = 0.0,
    t::Union{T, AbstractVector{T}} = 1.0
) where {T<:AbstractFloat}
    return DAsemiinf_params{T}(ρ, μa, μsp, n_med, n_ext, ω, z, t)
end

#------------------------------
# Steady-State Fluence
#------------------------------

"""
    fluence_DA_semiinf_CW(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0)

Compute the steady-state fluence in a semi-infinite geometry.

# Arguments
- `ρ`: the source detector separation (cm⁻¹)
- `μa`: absorption coefficient (cm⁻¹)
- `μsp`: reduced scattering coefficient (cm⁻¹)

# Keyword arguments
- `n_ext`: the boundary's index of refraction (air or detector)
- `n_med`: the sample medium's index of refraction
- `z`: the z-depth orthogonal from the boundary (cm)

# Examples
```
julia> fluence_DA_semiinf_CW(1.0, 0.1, 10.0)
julia> fluence_DA_semiinf_CW(1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, z = 0.0)
```
"""
function fluence_DA_semiinf_CW(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0)
    params = DiffusionKernelParams(μsp, n_med, n_ext)
    μeff = sqrt(3 * μa * μsp)
    return _kernel_fluence_DA_semiinf_CW(μeff, ρ, z, params.z0, params.zb, params.D)
end
"""
    fluence_DA_semiinf_CW(data::DiffusionParameters)

Wrapper to `fluence_DA_semiinf_CW(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0...)`

# Examples
```
julia> data = DAsemiinf_params(ρ = 1.0) # use structure to generate inputs
julia> fluence_DA_semiinf_CW(data) # then call the function
```
"""
function fluence_DA_semiinf_CW(data::DiffusionParameters)
    μeff = sqrt(3 * data.μa * data.μsp)
    return _kernel_fluence_DA_semiinf_CW(μeff, data.ρ, data.z, data.z0, data.zb, data.D)
end

function _kernel_fluence_DA_semiinf_CW(μeff, ρ, z, z0, zb, D)
    ϕ = exp(-μeff * sqrt(ρ^2 + (z - z0)^2)) / (sqrt(ρ^2 + (z - z0)^2))
    ϕ -= exp(-μeff * sqrt(ρ^2 + (z + 2 * zb + z0)^2)) / (sqrt(ρ^2 + (z + 2 * zb + z0)^2))
    return ϕ / (4 * π * D)
end

#------------------------------
# Time-Domain Fluence
#------------------------------

"""
    fluence_DA_semiinf_TD(t, ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0)

Compute the time-domain fluence in a semi-infinite medium.

# Arguments
- `t`: the time vector (ns).
- `ρ`: the source detector separation (cm⁻¹)
- `μa`: absorption coefficient (cm⁻¹)
- `μsp`: reduced scattering coefficient (cm⁻¹)

# Optional arguments
- `n_ext`: the boundary's index of refraction (air or detector)
- `n_med`: the sample medium's index of refraction
- `z`: the z-depth orthogonal from the boundary (cm)

# Examples
```
julia> fluence_DA_semiinf_TD(0.1:0.1:2.0, 1.0, 0.1, 10.0)
julia> fluence_DA_semiinf_TD(0.2:0.2:1.0, 1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, z = 0.0)
```
"""
function fluence_DA_semiinf_TD(t, ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0)
    params = DiffusionKernelParams(μsp, n_med, n_ext)
    return map(t -> _kernel_fluence_DA_semiinf_TD(params.D, params.ν, t, μa, z, params.z0, ρ, params.zb), t)
end

"""
    fluence_DA_semiinf_TD(t, data::DiffusionParameters)

Wrapper to `fluence_DA_semiinf_TD(t, ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0)`

# Examples
```
julia> data = DAsemiinf_params() # use structure to generate inputs
julia> fluence_DA_semiinf_TD(data) # then call the function
```
"""
function fluence_DA_semiinf_TD(data::DiffusionParameters)
    return map(t -> _kernel_fluence_DA_semiinf_TD(data.D, data.ν, t, data.μa, data.z, data.z0, data.ρ, data.zb), data.t)
end

@inline function _kernel_fluence_DA_semiinf_TD(D, ν, t, μa, z, z0, ρ, zb)
    tmp1 = 4 * D * ν * t
    ϕ = ν * exp(-μa * ν * t) / (π * tmp1)^(3/2)
    ϕ *= (exp(-((z - z0)^2 + ρ^2) / tmp1) - exp(-((z + z0 + 2*zb)^2 + ρ^2) / tmp1))
    return ϕ
end

#------------------------------
# Frequency-Domain Fluence
#------------------------------
"""
fluence_DA_semiinf_FD(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0, ω = 1.0)

Compute the frequency domain fluence in a semi-infinite geometry.

# Arguments
- `ρ`: the source detector separation (cm⁻¹)
- `μa`: absorption coefficient (cm⁻¹)
- `μsp`: reduced scattering coefficient (cm⁻¹)

# Keyword arguments
- `n_ext`: the boundary's index of refraction (air or detector)
- `n_med`: the sample medium's index of refraction
- `z`: the z-depth orthogonal from the boundary (cm)
- `ω`: the modulation frequency (1/ns)

# Examples
```
julia> fluence_DA_semiinf_FD(1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, z = 0.0, ω = 1.0)
```
"""
function fluence_DA_semiinf_FD(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0, ω = 1.0)
    params = DiffusionKernelParams(μsp, n_med, n_ext)
    μa_complex = μa + ω * im / params.ν
    μeff = sqrt(3 * μa_complex * data.μsp)
    return _kernel_fluence_DA_semiinf_CW(μeff, ρ, z, params.z0, params.zb, params.D)
end
"""
    fluence_DA_semiinf_FD(data::DiffusionParameters)

Wrapper to `fluence_DA_semiinf_FD(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0, ω = 1.0)`

# Examples
```
julia> data = DAsemiinf_params(ω = 1.0) # use structure to generate inputs
julia> fluence_DA_semiinf_FD(data) # then call the function
```
"""
function fluence_DA_semiinf_FD(data::DiffusionParameters)
    μa_complex = data.μa + data.ω * im / data.ν
    μeff = sqrt(3 * μa_complex * data.μsp)
    return _kernel_fluence_DA_semiinf_CW(μeff, data.ρ, data.z, data.z0, data.zb, data.D)
end

#------------------------------
# Time-Domain Flux
#------------------------------
"""
    flux_DA_semiinf_TD(t, ρ, μa, μsp; n_ext = 1.0, n_med = 1.0)

Compute the time-domain flux (D*∂ϕ(t)/∂z @ z = 0) from a semi-infinite medium from Eqn. 36 Contini 97 (m = 0).

# Arguments
- `t`: the time vector (ns).
- `ρ`: the source detector separation (cm⁻¹)
- `μa`: absorption coefficient (cm⁻¹)
- `μsp`: reduced scattering coefficient (cm⁻¹)

# Keyword arguments
- `n_ext`: the boundary's index of refraction (air or detector)
- `n_med`: the sample medium's index of refraction

# Examples
```
julia> flux_DA_semiinf_TD(0.1:0.1:2.0, 1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0)
```
"""
function flux_DA_semiinf_TD(t, ρ, μa, μsp; n_ext = 1.0, n_med = 1.0)
    params = DiffusionKernelParams(μsp, n_med, n_ext)
    z3m = - params.z0
    z4m = 2 * params.zb + params.z0

    return map(t -> _kernel_flux_DA_semiinf_TD(params.D, params.ν, t, μa, z3m, z4m, ρ), t)
 end

 """
    flux_DA_semiinf_TD(t, data::DiffusionParameters)

Wrapper to `flux_DA_semiinf_TD(t, ρ, μa, μsp; n_ext = 1.0, n_med = 1.0, z = 0.0)`

# Examples
```
julia> data = DAsemiinf_params() # use structure to generate inputs
julia> flux_DA_semiinf_TD(data) # then call the function
```
"""
function flux_DA_semiinf_TD(data::DiffusionParameters)
    z3m = - data.z0
    z4m = 2 * data.zb + data.z0

    return map(t -> _kernel_flux_DA_semiinf_TD(data.D, data.ν, t, data.μa, z3m, z4m, data.ρ), data.t)
end

@inline function _kernel_flux_DA_semiinf_TD(D, ν, t, μa, z3m, z4m, ρ)
    tmp1 = 4 * D * ν * t
    Rt = -(ρ^2 / (tmp1))
    Rt -= μa * ν * t
    Rt = -exp(Rt) / (2 * (4 * π * D * ν)^(3/2) * sqrt(t * t * t * t * t))
    Rt *= (z3m * exp(-(z3m^2 / (tmp1))) - z4m * exp(-(z4m^2 / (tmp1))))

    return Rt
end

#------------------------------
# Steady-State Flux
#------------------------------
"""
    flux_DA_semiinf_CW(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0)

Compute the steady-state flux (D*∂ϕ(ρ)/∂z @ z = 0) from a semi-infinite medium.

# Arguments
- `ρ`: the source detector separation (cm⁻¹)
- `μa`: absorption coefficient (cm⁻¹)
- `μsp`: reduced scattering coefficient (cm⁻¹)
- `ndet`: the boundary's index of refraction (air or detector)
- `nmed`: the sample medium's index of refraction

# Examples
```
julia> flux_DA_semiinf_CW(1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0)
```
"""
function flux_DA_semiinf_CW(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0)
    params = DiffusionKernelParams(μsp, n_med, n_ext)
    return params.D * ForwardDiff.derivative(z -> fluence_DA_semiinf_CW(ρ, μa, μsp, n_med = n_med, n_ext = n_ext, z = z), 0.0)
end

"""
    flux_DA_semiinf_CW(data::DiffusionParameters)

Wrapper to `flux_DA_semiinf_CW(ρ, μa, μsp; n_ext = 1.0, n_med = 1.0)`

# Examples
```
julia> data = DAsemiinf_params(ρ = 1.0) # use structure to generate inputs
julia> flux_DA_semiinf_CW(data) # then call the function
```
"""
function flux_DA_semiinf_CW(data)
    return flux_DA_semiinf_CW(data.ρ, data.μa, data.μsp, n_ext = data.n_ext, n_med = data.n_med)
end


######################################################################
# ergodic_manager/r2.jl
#
# handles stuff needed for ergodicity
######################################################################

export ErgodicManagerR2

# kpixl is a (K+1 x bins) matrix, each entry storing cos(k*pi*x / L)
#  this assumes some discretization
"""
`ErgodicManager(d::Domain, phi::Matrix{Float64}, K::Int)`

`ErgodicManager(example_name::String; K::Int=5, bins::Int=100)`

Valid `example_name` entries are:

* "single gaussian"
* "double gaussian"
"""
mutable struct ErgodicManagerR2 <: ErgodicManager
	domain::Domain				# spatial domain
	K::Int						# number of Fourier coefficients
	phi::Matrix{Float64}		# spatial distribution
	phik::Matrix{Float64}		# distribution's Fourier coefficients

	# constant regardless of phi (depend on k1,k2)
	Lambda::Matrix{Float64}
	hk::Matrix{Float64}

	# to speed up computation
	kpixl::Matrix{Float64}
	kpiyl::Matrix{Float64}

	function ErgodicManagerR2(d::Domain, phi::MF, K::Int=5)
		em = new()
		em.domain = deepcopy(d)
		em.K = K
		em.hk = zeros(K+1,K+1)
		em.phi = deepcopy(phi)
		em.phik = zeros(K+1,K+1)
		em.Lambda = zeros(K+1,K+1)
		em.kpixl = zeros(K+1, d.cells[1])
		em.kpiyl = zeros(K+1, d.cells[2])

		Lambda!(em)
		kpixl!(em)
		hk!(em)
		decompose!(em)
		return em
	end

	function ErgodicManagerR2(d::Domain, K::Int=5)
		em = new()
		em.domain = deepcopy(d)
		em.K = K
		em.hk = zeros(K+1,K+1)
		#em.phi = deepcopy(phi)
		em.phik = zeros(K+1,K+1)
		em.Lambda = zeros(K+1,K+1)
		em.kpixl = zeros(K+1, d.cells[1])
		em.kpiyl = zeros(K+1, d.cells[2])

		Lambda!(em)
		kpixl!(em)
		hk!(em)
		#decompose!(em)
		return em
	end
end

# fills each entry Lambda[k1,k2] in the Lambda matrix
function Lambda!(em::ErgodicManagerR2)
    for k1 = 0:em.K, k2 = 0:em.K
        den = (1.0 + k1*k1 + k2*k2) ^ 1.5
        em.Lambda[k1+1, k2+1] = 1.0 / den
    end
end

# TODO: check if the row/col ordering of this is julia-efficient
# TODO: also, I could probably speed up by iterating over k first?
function kpixl!(em::ErgodicManagerR2)
	Lx = em.domain.lengths[1]
	xmin = x_min(em)
	for xi = 1:x_cells(em)
		# TODO: seems like I just cancel xmin out. Ignore it maybe?
		x = xmin + (xi-0.5)*x_size(em)
		for k = 0:em.K
			em.kpixl[k+1,xi] = cos(k*pi*(x-xmin) / Lx)
		end
	end

	Ly = em.domain.lengths[2]
	ymin = y_min(em)
	for yi = 1:y_cells(em)
		y = ymin + (yi-0.5)*y_size(em)
		for k = 0:em.K
			em.kpiyl[k+1,yi] = cos(k*pi*(y-ymin) / Ly)
		end
	end
end

# generates the hk coefficients for the ergodic manager
# these coefficients only need to be computed once
function hk!(em::ErgodicManagerR2)
    for k1 = 0:em.K, k2 = 0:em.K
        em.hk[k1+1,k2+1] = hk_ij(em, k1, k2)
    end
end

# computes the coefficients for a specific value of k1 and k2
# called by hk!
function hk_ij(em::ErgodicManagerR2, k1::Int, k2::Int)
	val = 0.0
	for xi = 1:x_cells(em)
		cx = em.kpixl[k1+1,xi]
		cx2 = cx * cx
		for yi = 1:y_cells(em)
			cy = em.kpiyl[k2+1,yi]
			val += cx2 * cy * cy * em.domain.cell_size
		end
	end

	return sqrt(val)
end

######################################################################
# Computing Fourier coefficients
######################################################################
# update the Fourier coefficients based on some distribution
# Here, I assume it is discrete, but I should change this...
# TODO: maybe I should do some bounds checking?
function decompose!(em::ErgodicManagerR2, d::Matrix{Float64})
    for k1 = 0:em.K, k2 = 0:em.K
        em.phik[k1+1,k2+1] = phi_ij(em, k1, k2, d)
    end
    em.phi = d
end

# iterate over the state space
function phi_ij(em::ErgodicManagerR2, k1::Int, k2::Int, d::Matrix{Float64})
	val = 0.0
	for xi = 1:x_cells(em)
		cx = em.kpixl[k1+1,xi]
		for yi = 1:y_cells(em)
			cy = em.kpiyl[k2+1,yi]
			val += d[xi,yi] * cx * cy * em.domain.cell_size
		end
	end
	return val / em.hk[k1+1,k2+1]
end

"""
`decompose(em, traj::VVF)`

Decomposes a set of positions into a set of `ck` Fourier coefficients.

"""
function decompose(em::ErgodicManagerR2, traj::VVF)

	# trajectory is really of length N+1
	N = length(traj)-1

	# create matrix to hold trajectory's Fourier coefficients
	ck = zeros(em.K+1, em.K+1)

	# lengths of each dimension
	Lx = em.domain.lengths[1]
	Ly = em.domain.lengths[2]

	# minimum values in each dimension
	xmin = x_min(em)
	ymin = y_min(em)

	for k1 = 0:em.K
		kpiL1 = k1 * pi / Lx
		for k2 = 0:em.K
			kpiL2 = k2 * pi / Ly
			hk = em.hk[k1+1, k2+1]
			fk_sum = 0.0
			# now loop over time
			for n = 0:N
				xn = traj[n + 1]
				c1 = cos(kpiL1 * (xn[1]-xmin))
				c2 = cos(kpiL2 * (xn[2]-ymin))
				fk_sum += c1*c2
			end
			ck[k1+1, k2+1] = fk_sum / (hk * (N+1))
		end
	end
	return ck
end

# reconstructs from Fourier coefficients in ck
function reconstruct(em::ErgodicManagerR2, ck::Matrix{Float64})
	# iterate over all bins
	v = em.domain.cell_size
	vals = zeros(x_cells(em), y_cells(em))

	for xi = 1:x_cells(em)
		x = x_min(em) + (xi-0.5)*x_size(em)
		for yi = 1:y_cells(em)
			y = y_min(em) + (yi-0.5)*y_size(em)
			for k1 = 0:em.K
				cx = em.kpixl[k1+1,xi]
				for k2 = 0:em.K
					cy = em.kpixl[k2+1,yi]
					vals[xi,yi] += ck[k1+1,k2+1]*cx*cy/em.hk[k1+1,k2+1]
				end
			end
		end
	end
	return vals
end


export impute_values

function logistic(x)
    return 1.0 ./ (1.0 .+ exp.(-x))
end

function compute_Z(model::MatFacModel;
                   covariates::Union{Nothing,Matrix}=nothing,
                   factor_set::Union{Nothing,Vector{Int64}}=nothing)

    if factor_set == nothing
        factor_set = collect(1:size(model.X,1))
    end
    X = model.X[factor_set,:]
    Y = model.Y[factor_set,:]

    Z = transpose(X)*Y
    Z .+= model.instance_offset
    Z .+= transpose(model.feature_offset)

    if covariates != nothing
        Z .+= covariates*model.instance_covariate_coeff
    end

    return Z
end

function apply_link_function(Z::Matrix, model::MatFacModel)

    type_vec = typeof.(model.losses)
    logistic_idx = (type_vec .== LogisticLoss)
    poisson_idx = (type_vec .== PoissonLoss)

    Z[:,logistic_idx] .= logistic(Z[:,logistic_idx])
    Z[:,poisson_idx] .= exp.(Z[:,poisson_idx])

    return Z
end

function impute_values(model::MatFacModel;
                       covariates::Union{Nothing,Matrix}=nothing,
                       factor_set::Union{Nothing,Vector{Int64}}=nothing)

    Z = compute_Z(model; covariates=covariates, factor_set=factor_set)
    Z = apply_link_function(Z, model)
    return Z
end


module TestConditional
using Test
using LinearAlgebra
using GaussianDistributions
using GaussianDistributions: ConditionalMean
using ForwardDiff: derivative

const nsigma = 4

@testset "conditional" begin
    @testset "multivariate" begin
        d = 3
        μ = randn(d)
        Σ = randn(d, d)
        Σ = Σ'Σ
        G = Gaussian(μ, Σ)

        x = randn()
        tol = 1e-13

        C = G | (1, x)
        @test C == conditional(G, 1, x)
        @test mean(C[1]) ≈ x
        @test cov(C[1]) < tol # collapse of uncertainty
        @test all(vc < vg for (vc, vg) in zip(var(C), var(G))) # uniform reduction in uncertainty

        x = randn(2)
        C = G | (1:2, x)
        @test mean(C[1:2]) ≈ x
        @test sum(std(C[1:2])) < tol # collapse of uncertainty

        # play around with posterior
        x = randn()
        σ = .1
        C = G | (1, x, σ^2)
        @test abs(mean(C[1])-x) < abs(mean(G[1])-x) # posterior mean is closer to data than prior mean
        @test cov(C[1]) < cov(G[1])  # posterior uncertainty is smaller than prior
        @test sum(var(C)) < sum(var(G))

        # test conversion of vector of noise variances to diagonal covariance
        n = 2
        x = randn(n)
        σ² = σ^2 * exp.(randn(n))
        C = G | (1:n, x, σ²)
        E = G | (1:n, x, Diagonal(σ²))
        @test C isa Gaussian
        @test mean(C) ≈ mean(E)
        @test cov(C) ≈ cov(E)
    end

    @testset "bayesian linear regression" begin
        # i.e. conditioning on linear operator data
        n, m = 4, 8
        A = randn(n, m)
        σx = 2
        x = σx * randn(m)
        b = A*x

        x = Gaussian(σx^2*I(m))

        # noiseless conditioning
        P = x | (A, b)
        PC = Matrix(cov(P))
        λ = sort!(eigvals(PC))
        @test all(<(1e-12), λ[1:n])
        @test all(>(1e-12), λ[n+1:end])
        @test A*mean(P) ≈ b

        # test iterative conditioning
        Q = x | (A[1:n-1, :], b[1:n-1])
        Q = Q | (A[n, :], b[n])
        @test mean(Q) ≈ mean(P)
        @test cov(Q) ≈ cov(P)

        # with noise
        σ = 1e-6
        y = Gaussian(b, σ^2*I(n))
        P = x | (A, y)

        # n observations only uniquely determine n components
        PC = Hermitian(Matrix(cov(P)))
        λ = sort!(eigvals(PC))
        @test all(<(σ), λ[1:n])
        @test all(>(σ), λ[n+1:end])

        # if we observe a complete set, we can reconstruct the solution
        @test norm(A*mean(P)-mean(y)) < 3σ

        # test iterative conditioning
        Q = x | (A[1:n-1, :], y[1:n-1])
        Q = Q | (A[n, :], y[n])
        @test mean(Q) ≈ mean(P)
        @test cov(Q) ≈ cov(P)
    end

    @testset "process mean" begin
        # create dataset
        G = 4
        x = sort(randn(G))
        f(x) = sin(2π*x)

        l = .1
        k(x, y) = exp(-(x-y)^2/2l^2)
        G = Gaussian(k)
        K = k.(x, x')
        cm = ConditionalMean(G, x, f.(x), K)

        @test marginal(cm, x) ≈ f.(x)

        # test differentibility
        using ForwardDiff: derivative
        g(x) = derivative(cm, x)
    end

    @testset "process" begin
        # create dataset
        n = 16
        x = sort(randn(n))
        f(x) = sin(2π*x)
        y = f.(x)

        # setup Gaussian process
        l = .1
        k(x, y) = exp(-sum(abs2, x-y)/2l^2)
        G = Gaussian(k)
        K = k.(x, x')
        C = G | (x, y) # noiseless conditioning
        @test mean(C, x) ≈ f.(x)
        tol = 10eps(Float64)

        v = var(C, x)
        @test all(v .< tol) # collapse of uncertainty

        # noisy conditioning
        n = 128
        x = sort(randn(n))
        y = f.(x)
        σ = .01
        @. y += σ*randn() # adding noise
        C = G | (x, y, σ^2)
        d = mean(C, x)-f.(x)
        @test all(abs.(d) .< nsigma*σ)
        v = var(C, x)
        @test all(v .< σ^2) # collapse of uncertainty

        # test differentibility
        g(x) = derivative(mean(C), x)
        @test isapprox(g(0), 2π, rtol = .1)

        # TODO: iterative conditioning
        # xs = randn()
        # ys = f(xs)
        # C = C | (xs, ys, σ^2)

        # d = mean(C)(xs) - f(xs)
        # @test abs(d) < nsigma*σ
        # tol = σ
        # v = std(C)(xs)
        # @test v < tol

        @testset "single datapoint conditioning" begin
            x = randn()
            y = f(x)
            C = G | (x, y)
            @test mean(C)(x) ≈ y
            @test std(C)(x) < 1e-12

            # with noise
            y += σ*randn()
            C = G | (x, y, σ^2)
            @test isapprox(mean(C)(x), y, atol = nsigma*σ)
            @test std(C)(x) < nsigma*σ
            @test std(C)(x) > σ/2 # uncertainty does not collapse completely

            # single point conditioning with 2d input
            x = randn(2)
            y = randn()
            C = G | (x, y)
            @test mean(C)(x) ≈ y
            @test std(C)(x) < 1e-12

            # with noise
            y += σ*randn()
            C = G | (x, y, σ^2)
            @test isapprox(mean(C)(x), y, atol = nsigma*σ)
            @test std(C)(x) < nsigma*σ
        end
    end
end

end


using LinearAlgebra

function diff_svd(A,dA) #for a single parameter
    A_svd = svd(A)
    U = A_svd.U
    V = A_svd.V
    S = A_svd.S

    dS = U'*dA*V
    s = size(dS,1)

    Up = triu(dS,1) #extract the upper triangular part of the dS without the diagonal
    Lo = tril(dS,-1) #extract the lower triangular part of the dS matrix without diagonal

    L2_bar = zeros(s,s)

    for i = 2:s,j=1:(i-1)
        L2_bar[i,j] = (Up[j,i]*S[j] + Lo[i,j]*S[i])/(S[i]^2- S[j]^2)
    end

    dV = V*(L2_bar' - L2_bar)
    return diag(dS),dV
end

function jacob_svd(A,dA)
    p = size(dA,2)
    if isempty(size(A))
        n = 1
        m = 1
    else
        n,m = size(A)
    end

    dV = zeros(n*m,p)
    dS = zeros(max(n,m),p)

    for k in 1:p
        dA_unvec = dA[:,k]
        dA_unvec = reshape(dA_unvec,n,m)
        dS_temp,dV_temp = diff_svd(A,dA_unvec)
        dS[:,k] = dS_temp
        dV[:,k] = vec(dV_temp)
    end

    return dS,dV

end

mutable struct svd_kalman_diff
    # ALL V MATRICES FROM THE SVD ARE STORE TRANSPOSED I.E. Vt
    A # state transition
    G #Observable transition

    #shock of observable variance
    S_Q
    V_Q

    #variance of the shock of the state
    S_R
    V_R

    state_mean

    #variance of the state

    S_var
    V_var

    dA
    dG

    dS_Q
    dV_Q

    dS_R
    dV_R

    dS_var
    dV_var

    diff_state

    D_re #this will always be the square root of what we need
    e_bar

end

function svd_kalman_diff(A,G,Q,R,dA,dG,dR,dQ)

    m,n = size(G)
    st_mean = zeros(n)
    st_var = zeros(n,n)

    dS_Q, dV_Q = jacob_svd(Q,dQ)
    dS_R, dV_R = jacob_svd(R,dR)

    diff_state = zeros(n^2,size(dA,2))

    svd_Q = svd(Q)
    svd_R = svd(R)

    S_R = diagm(sqrt.(svd_R.S))
    S_Q = diagm(sqrt.(svd_Q.S)[1:rank(Q)])

    V_R = svd_R.Vt
    V_Q = svd_Q.Vt

    dS_R = inv(kron(I(m),S_R) + kron(S_R,I(m)))*dS_R
    dS_Q = inv(kron(I(rank(Q)),S_Q) + kron(S_Q,I(rank(Q))))*dS_Q[1:rank(Q),:]

    D_re = zeros(m,m)
    e_bar = zeros(m)

    return svd_kalman_diff(A,G,S_Q,V_Q,S_R,V_R,st_mean,st_var,st_var,dA,dG,dS_Q,dV_Q,dS_R,dV_R,dS_Q,dV_Q,D_re,e_r)
end

function start!(svd_kalman_diff::svd_kalman_diff,x0,v0,dv0)
    m,n = size(svd_kalman_diff.G)

    svd_kalman_diff.state_mean = x0

    svd_v0 = svd(v0)
    S_var = diagm(sqrt.(svd_v0.S))
    svd_kalman_diff.S_var = S_var
    svd_kalman_diff.V_var = svd_v0.Vt

    svd_kalman_diff.diff_state = zeros(size(x0,1))

    dS_var,dV_var = jacob_svd(v0,dv0)

    dS_var = inv(kron(I(n),S_var) + kron(S_var,I(n)))*dS_var

    svd_kalman_diff.dS_var = dS_var
    svd_kalman_diff.dV_var = dV_var

    return nothing
end

function forecast!(svd_kalman_diff::svd_kalman_diff)
    m,n = size(svd_kalman_diff.G)
    Q = svd_kalman_diff.S_Q*svd_kalman_diff.V_Q
    A = svd_kalman_diff.A
    old_vari = svd_kalman_diff.S_var*svd_kalman_diff.V_var
    old_state = svd_kalman_diff.state_mean

    dSvar = svd_kalman_diff.diff_Svar
    dQvar = svd_kalman_diff.diff_Vvar

    new_state = A*old_state

    pre_var = [old_vari*A';Q]

    svd_var = svd(pre_var)

    svd_kalman_diff.state_mean = new_state
    svd_kalman_diff.S_var = diagm(svd_var.S)
    svd_kalman_diff.V_var = svd_var.Vt

    b1 = kron(A*svd_kalman_diff.V_var,I(n))*svd_kalman_diff.dS_var + kron(A,svd_kalman_diff.S_var)*svd_kalman_diff.dV_var + kron(I(n),old_vari)*commutation_matrix(n,n)*svd_kalman_diff.dA
    b2 = kron(svd_kalman_diff.V_Q,I(n))*svd_kalman_diff.dS_Q + kron(I,svd_kalman_diff.S_Q)*svd_kalman_diff.dV_Q
    dVar = [b1;b2]

    svd_kalman_diff.Svar,svd_kalman_diff.Vvar = jacob_svd(diagm(svd_var.S)*svd_var.Vt,dVar)

    svd_kalman_diff.diff_state = kron(new_state',I(n))*svd_kalman_diff.dA + A*svd_kalman_diff.diff_state

    return nothing
end

function update_state!(svd_kalman_diff::svd_kalman_diff,y)
    R = svd_kalman_diff.S_R*svd_kalman_diff.V_R
    old_var = svd_kalman_diff.S_var*svd_kalman_diff.V_var
    G = svd_kalman_diff.G
    old_state = svd_kalman_diff.state_mean
    d_var = kron(svd_kalman_diff.V_var',I(n))*svd_kalman_diff.dS_var + kron(I(n),svd_kalman_diff.S_var)*svd_kalman_diff.dV_var
    m,n = size(G)

    re_mat = [R;old_var*G]
    svd_re = svd(re_mat)

    K_bar = (old_var'*old_var)*G'*svd_re.Vt
    er = *(y - G*old_state)
    e_bar = svd_re.Vt*er

    svd_kalman_diff.e_bar = e_bar
    svd_kalman_diff.D_re = diagm(svd_re.S)

    new_state = old_state + K_bar*diagm(svd_re.S.^(-2))*e_bar

    K = K_bar*diagm(svd_re.S.^(-2))*svd_re.Vt

    pre_var = [old_var*(I - K*G)';R*K']

    svd_var = svd(pre_var)

    svd_kalman_diff.S_var = diagm(svd_var.S)
    svd_kalman_diff.V_var = svd_var.Vt

    b1 = kron(svd_kalman_diff.V_R,I(n))*svd_kalman_diff.dS_R + kron(I(n),svd_kalman_diff.S_R)*commutation_matrix(n,n)*svd_kalman_diff.dV_R

    b2 = kron(G*svd_kalman_diff.V_var,I(n))*svd_kalman_diff.dS_var + kron(G,svd_kalman_diff.S_var)*commutation_matrix(n,n)*svd_kalman_diff.dV_var + kron(I,svd_kalman_diff.S*svd_kalman_diff.V)*commutation_matrix(m,n)*svd_kalman_diff.dG

    diff_re = [b1;b2]
    diff_svd_re = jacob_svd(re_mat,diff_re)

    diff_k_bar = kron(svd_re.Vt*G,I(n))*d_var + kron(svd_re.Vt,old_var)*svd_kalman_diff.dG + kron(I(n),old_var*G')*diff_svd_re.dV

    diff_k = kron(svd_re.V*diagm(svd_re.S .^(-2)),I(n))*diff_k_bar - kron(svd_re.V,K_bar)*(kron(diagm(svd_re.S .^(-2)),diagm(svd_re.S)) + kron(diagm(svd_re.S),diagm(svd_re.S .^(-2))))*diff_svd_re.dS + kron(I(n),K_bar*diagm(svd_re.S .^(-2)))*diff_svd_re.dV

    e_diff = -(kron(old_state',I(m))*svd_kalman_diff.dG + G*svd_kalman_diff.diff_state)
    e_bar_diff = kron(er',I(m))*diff_svd_re.dV + svd_re.Vt*e_diff

    svd_kalman_diff.diff_state = svd_kalman_diff.diff_state + kron(e_bar'*diagm(svd_re.S .^(-2)),I(n))*diff_k_bar - kron(e_bar',K_bar)*(kron(diagm(svd_re.S .^(-2)),diagm(svd_re.S)) + kron(diagm(svd_re.S),diagm(svd_re.S .^(-2))))*diff_svd_re.dS + K_bar*diagm(svd_re.S .^(-2))*e_bar_diff

    b0 = I(n) - K*G
    b1 = kron(b0*svd_kalman_diff.V_var,I(n))*svd_kalman_diff.dS_var + kron(b0,svd_kalman_svd.S_var)*svd_kalman_diff.dV_var - kron(K,svd_kalman_diff.S_var*svd_kalman_diff.V_var)*commutation_matrix(m,n)*svd_kalman_diff.dG - kron(I(n),svd_kalman_diff.S_var*svd_kalman_diff.V_var*G')*commutation_matrix(n,m)*dK
    b2 = kron(K*svd_kalman_diff.V_R',I(m))*svd_kalman_diff.dS_R + kron(K,svd_kalman_diff.S_R)*svd_kalman_res.dQ_R + kron(I(n),svd_kalman_diff.D_R*svd_kalman_diff.Q_R)*commutation_matrix(n,m)*dK

    diff_var = [b1;b2]

    var_pre = [old_var*b0';svd_kalman_diff.S_R*svd_kalman_diff.V_R*K']

    svd_var = svd(var_pre)

    svd_kalman_diff.S_var = diagm(svd_var.S)
    svd_kalman_diff.V_var = svd_var.Vt

    svd_kalman_diff.dS_var,svd_kalman_diff.dV_var = svd_jacob(var_pre,diff_var)

    return nothing

end

function update_and_forecast!(svd_kalman_diff::svd_kalman_diff,y)
    update_state!(svd_kalman_diff,y)
    forecast!(svd_kalman_diff)
    return nothing
end


"""
    MLEstimator

Simple maximum likelihood estimator for response rate `p`.
"""
struct MLEstimator{TD<:Design} <: Estimator

  design::TD

  MLEstimator{TD}(design::TD) where {TD<:Design} = new(design)

end # MLEStimator

"""
    MLEstimator(design::TD) where {TD<:Design}

Create maximum likelihood estimator for response rate under given design.
"""
MLEstimator(design::TD) where {TD<:Design} = MLEstimator{TD}(design)

Base.show(io::IO, estimator::MLEstimator) = print("MLEstimator")

function estimate(estimator::MLEstimator, x1::T, x2::T) where {T<:Integer}

    checkx1x2(x1, x2, design(estimator))
    return (x1 + x2)/samplesize(design(estimator), x1)

end


@symbol_func function G_PU(cur_reactor::AbstractReactor)
  cur_left_part = K_RP(cur_reactor)

  cur_left_part *= cur_reactor.tau_FT

  cur_right_part = K_LP(cur_reactor)

  cur_right_part *= safe_symbol(cur_reactor, :R_0) ^ 2

  cur_right_part *= cur_reactor.T_bar ^ (3/2)

  cur_G = cur_left_part

  cur_G += cur_right_part

  cur_G
end


using Plots
using ColorSchemes
using Printf

pyplot()

function meshgrid(x, y)
    x = reshape(x, 1, length(x))
    y = reshape(y, length(y), 1)
    X = repeat(x, length(y), 1)
    Y = repeat(y, 1, length(x))
    return X, Y
end

function plot_growth_slice(filenc, filejson, fileplot)

    # read json file
    json_string = JSON3.read(read(filejson, String))

    Umax = json_string.jet.Umax
    L    = json_string.jet.L
    fz   = json_string.physics.fz

    # read NetCDF file
    ωs, modes, ks, ms, iEigs, ys = load_spectrum(filenc)

    print("\n--> Plotting the growth rates for a fixed k versus m in ", fileplot, "\n")

    growth = zeros(size(ωs))
    for (ik, k) in enumerate(ks)
        for (im, m) in enumerate(ms)
            for iEig in iEigs
                growth[iEig, im, ik] = max(0, imag(ωs/fz)[iEig, im, ik])
            end
        end
    end

    # scaling wavenumber for plotting
    m = L/(2*π)*ms
    plt = plot()
    for iEig in iEigs
        plot!(plt, m, growth[iEig, :, 1],
                label = @sprintf("%d", iEig),
            linewidth = 2,
            xlim = (m[1], m[end]),
               xlabel = "L m/(2*π)",
                title = "Growth rates for k = 0",
               legend = :topright)
    end
    savefig(plt, fileplot)
end

function plot_growth(filenc, filejson, fileplot)

    # read json file
    json_string = JSON3.read(read(filejson, String))

    Umax = json_string.jet.Umax
    L    = json_string.jet.L
    fz   = json_string.physics.fz

    # read NetCDF file
    ωs, modes, ks, ms, iEigs, ys = load_spectrum(filenc)

    print("\n--> Plotting the growth rates versus (k, m) in ", fileplot, "\n")

    growth = zeros(size(ωs)[2:3])
    for (ik, k) in enumerate(ks)
        for (im, m) in enumerate(ms)
            growth[im, ik] = max(0, imag(ωs/fz)[1, im, ik])
        end
    end

    # scaling wavenumber for plotting
    k = L/(2*π)*ks
    m = L/(2*π)*ms

    kwargs = (
        xlabel = "k",
        ylabel = "m",
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :haline,
      colorbar = true,
          xlim = (k[1], k[end]),
          ylim = (m[1], m[end])
    )

    plt = contour(k, m, growth, title="Growth rates for 1st mode"; kwargs...)
    savefig(plt, fileplot)
end

function plot_modes_1D(filenc, filejson, fileplot, Neigs)

        # read json file
        json_string = JSON3.read(read(filejson, String))

        Umax = json_string.jet.Umax
        L    = json_string.jet.L
        fz   = json_string.physics.fz
        Ny   = json_string.grid.Ny
        Ly   = json_string.grid.Ly

        # read NetCDF file
        ωs, modes, ks, ms, iEigs, ys = load_spectrum(filenc)

        print("\n--> Plotting 1D modes in the file ", fileplot, "\n")

        # scaling wavenumber for plotting
        y = ys[1:Ny+1]
        k = L/(2*π)*ks
        m = L/(2*π)*ms

        ik = 1    # most unstable mode

        for iEig in collect(1:Neigs)
            tmp              = findmax(imag(ωs/fz)[iEig, :, ik])
            max_growth, indm = tmp[1], tmp[2]

            @printf("\n iEig = %2d  k = %10.6f  m = %10.6f  growth = %10.6f",
                    iEig, float(k[ik]), float(m[indm]), max_growth)

            vvec          = zeros(ComplexF64, Ny+1)
            uvec          = modes[1:Ny+1, iEig, indm, ik]
            vvec[2:end-1] = modes[Ny+2:2*Ny, iEig, indm, ik]
            bvec          = modes[2*Ny+1:3*Ny+1, iEig, indm, ik]

            u_plt = plot(y/ L, real(uvec), line=:solid, color=:blue, linewidth=2, label="Re(u)")
            plot!(u_plt, y/ L, imag(uvec), line=:solid, color=:red,  linewidth=2, label="Im(u)")
            v_plt = plot(y/ L, real(vvec), line=:solid, color=:blue, linewidth=2, label="Re(v)")
            plot!(v_plt, y/ L, imag(uvec), line=:solid, color=:red,  linewidth=2, label="Im(v)")
            b_plt = plot(y/ L, real(bvec), line=:solid, color=:blue, linewidth=2, label="Re(b)")
            plot!(b_plt, y/ L, imag(bvec), line=:solid, color=:red,  linewidth=2, label="Im(b)")
            plt = plot(u_plt, v_plt, b_plt, layout=(3,1), size=(800, 600))
            file = string(fileplot, iEig, ".png")
            savefig(plt, file)
        end
end

function plot_modes_2D(filenc, filejson, fileplot, Neigs)

    # read json file
    json_string = JSON3.read(read(filejson, String))

    Umax = json_string.jet.Umax
    L    = json_string.jet.L
    fz   = json_string.physics.fz
    Ny   = json_string.grid.Ny
    Ly   = json_string.grid.Ly

    # read NetCDF file
    ωs, modes, ks, ms, iEigs, ys = load_spectrum(filenc)

    print("\n--> Plotting 2D modes in the file ", fileplot, "\n")

    # scaling wavenumber for plotting
    y = ys[1:Ny+1]
    k = L/(2*π)*ks
    m = L/(2*π)*ms

    ik = 1    # most unstable mode
    iEig = 1
    indm = findmax(imag(ωs/fz)[ik, :, iEig])[2]
    Nz = Ny
    Lz = 2*π/ms[indm]
    dz = Lz/Nz
    z = collect(0:dz:Lz)
    Y,Z = meshgrid(y, z)

    kwargs = (
        xlabel = "y",
        ylabel = "z",
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
#          xlim = (k[1], k[end]),
#          ylim = (m[1], m[end])
    )

    for iEig in collect(1:Neigs)
        tmp              = findmax(imag(ωs/fz)[iEig, :, ik])
        max_growth, indm = tmp[1], tmp[2]

        @printf("\n iEig = %2d  k = %10.6f  m = %10.6f  growth = %10.6f",
                iEig, float(k[ik]), float(m[indm]), max_growth)

        # 1D structures
        uvec            = zeros(ComplexF64, 1, Ny+1)
        vvec            = zeros(ComplexF64, 1, Ny+1)
        bvec            = zeros(ComplexF64, 1, Ny+1)
        uvec[1,:]       = modes[1:Ny+1, iEig, indm, ik]
        vvec[1,2:end-1] = modes[Ny+2:2*Ny, iEig, indm, ik]
        bvec[1,:]       = modes[2*Ny+1:3*Ny+1, iEig, indm, ik]

        # 2D structures
        U = real(repeat(uvec, Nz+1, 1) .* exp.(1im * ms[indm] * Z))
        V = real(repeat(vvec, Nz+1, 1) .* exp.(1im * ms[indm] * Z))
        B = real(repeat(bvec, Nz+1, 1) .* exp.(1im * ms[indm] * Z))

        u_plt = contour(Y/1e3, Z/1e3, U, title="u"; kwargs...)
        v_plt = contour(Y/1e3, Z/1e3, V, title="v"; kwargs...)
        b_plt = contour(Y/1e3, Z/1e3, B, title="b"; kwargs...)
        plt = plot(u_plt, v_plt, b_plt, layout=(1,3), size=(1600, 400))
        file = string(fileplot, iEig, ".png")
        savefig(plt, file)
    end
end

#=
# choose y label to be either y or ϕ, depending on the geometry
function plot_basic_state(grid, background, geometry, file)

    if geometry == Cartesian()
        y = grid.y/1e3
        xlabel = "y (km)"
    elseif geometry == Spherical()
        y = rad2deg.(grid.ϕ)
        xlabel = "ϕ (°)"
    end

    solution = background[geometry]

     p1 = plot(
         y,
         solution.u,
         linewidth = 3,
         label = "u",
         title = "Velocity",
         ylabel = "(m/s)",
         xlims = (y[end], y[1])
         )

     p2 = plot(
         y,
         solution.η,
         linewidth = 3,
         label = "η",
         xlabel = xlabel,
         ylabel = "(m)",
         title = "Free-Surface",
         xlims = (y[end], y[1])
         )

     plt1 = plot(p1, p2, layout = (2,1))
     savefig(plt1, file)
end

#ticks = collect(0:0.2:1)
#ticklabels = [ @sprintf("%5.1f",x) for x in ticks ]
#plot(x,y)
#plot!(xticks=(ticks,ticklabels))

function  make_tick_labels(fs, Nfs)
    fs_max = maximum([real(fs); imag(fs)])
    fs_min = minimum([real(fs); imag(fs)])
    dfs = (fs_max - fs_min)/Nfs

    ticks = collect(fs_min:dfs:fs_max)
    ticklabels = [ @sprintf("%f",f) for f in ticks]

    return ticks, ticklabels
end

function plot_1D_streamfunction(k_index, k, phys, σmodes, geometry, mode_number, file)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        xlabel = "y (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        xlabel = "ϕ (°)"
    end

    N = phys[geometry].grid.N
    Lx = 2*pi/k;
    Nx = N;
     x = LinRange(0, Lx, Nx+1)
  X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];
  ηvec = σmodes[(k,geometry)][2*N+1:3*N+1, mode_number];

    p1 = plot(
      y,
      real(uvec),
      linewidth = 3,
      label = "real",
      title = "u",
      yticks=make_tick_labels(uvec, 3),
      xlims = (y[end], y[1])
      )
  plot!(
      p1,
      y,
      linewidth = 3,
      imag(uvec),
      label = "imag")
  p2 = plot(
      y,
      real(vvec),
      linewidth = 3,
      label = "real",
      title = "v",
      yticks=make_tick_labels(vvec, 3),
       xlims = (y[end], y[1])
      )
  plot!(
      p2,
      y,
      linewidth = 3,
      imag(vvec),
      label = "imag")
  p3 = plot(
      y,
      real(ηvec),
      linewidth = 3,
      label = "real",
      title = "η",
      yticks=make_tick_labels(ηvec, 3),
      xlabel = xlabel,
      xlims = (y[end], y[1])
      )
  plot!(
      p3,
      y,
      linewidth = 3,
      imag(ηvec),
      label = "imag")
  plt = plot(p1, p2, p3, layout = (3,1))
  savefig(plt, file)
  #display(plt)

end

function plot_2D_streamfunction(k_index, k, phys, σmodes, geometry, mode_number, file)

    N = phys[geometry].grid.N
    Nx = N;
    Lx = 2*pi/k;
    x = LinRange(0, Lx, Nx+1)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        ylabel = "y (km)"
        x = x./1e3
        kscaled = k*1e3
        xlabel = "x (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        ylabel = "ϕ (°)"
        x = rad2deg.(x)
        kscaled = k * π / 180
        xlabel = "λ (°)"
     end

  X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];
  ηvec = σmodes[(k,geometry)][2*N+1:3*N+1, mode_number];

    u =     repeat(real(uvec),1,Nx+1).*cos.(kscaled*X) -    repeat(imag(uvec),1,Nx+1).*sin.(kscaled*X);
    v = -k.*repeat(imag(vvec),1,Nx+1).*cos.(kscaled*X) - k.*repeat(real(vvec),1,Nx+1).*sin.(kscaled*X);
    η =     repeat(real(ηvec),1,Nx+1).*cos.(kscaled*X) -    repeat(imag(ηvec),1,Nx+1).*sin.(kscaled*X);

    kwargs = (
        xlabel = "x (km)",
        ylabel = ylabel,
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
          xlim = (x[1],   x[end]),
          ylim = (y[end], y[1])
    )

    #print("x ticks = ", make_tick_labels(x, 2), "\n")

    u_plt = contour(x, y[end:-1:1], u[:, end:-1:1], title="u"; kwargs...)
    v_plt = contour(x, y[end:-1:1], v[:, end:-1:1], title="v"; kwargs...)
    η_plt = contour(x, y[end:-1:1], η[:, end:-1:1], title="η"; kwargs...)

    plt1 = plot(u_plt, v_plt, η_plt, layout = (1,3))
    savefig(plt1, file)

end

function plot_2D_vorticity(phys, k_index, k, σmodes, geometry, mode_number, file)

    N = phys[geometry].grid.N
    Lx = 2*pi/k;
    Nx = N;
    x = LinRange(0, Lx, Nx+1)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        ylabel = "y (km)"
        Dy = phys[geometry].grid.D
        x = x./1e3
        kscaled = k*1e3
        xlabel = "x (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        ylabel = "ϕ (°)"
        Dy = phys[geometry].grid.D/phys[geometry].grid.a
        x = rad2deg.(x)
        kscaled = k * π / 180
        xlabel = "λ (°)"
    end

    X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];

    ζvec = im * k * vvec - Dy * uvec
    ζ    = repeat(real(ζvec),1,Nx+1).*cos.(kscaled*X) - repeat(imag(ζvec),1,Nx+1).*sin.(kscaled*X);

    kwargs = (
        xlabel = xlabel,
        ylabel = ylabel,
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
          xlim = (  x[1], x[end]),
         ylim = (y[end],  y[1])
    )

    plt = contour(x, y[end:-1:1], ζ[:, end:-1:1], title="ζ"; kwargs...)

    savefig(plt, file)

end
=#

#string.(fieldnames(typeof(gridC)))

# Plot vertical slices as well?E


import DataStructures: DefaultDict, OrderedSet
import LinearAlgebra: cross, norm

function match(segments::Vector{LineSegment})::Vector{Polygon}
    cache = DefaultDict{Vertex, Vector{Vertex}}(Vector{Vertex})
    for segment in segments
        push!(cache[segment.a], segment.b)
        push!(cache[segment.b], segment.a)
    end

    polygons = Vector{Polygon}()

    while !isempty(cache)
        polygon = OrderedSet{Vertex}()
        start = first(cache).first
        push!(polygon, start)
        current = pop!(cache[start])
        delete!(cache, start)

        while haskey(cache, current)
            push!(polygon, current)
            a = pop!(cache[current])
            b = pop!(cache[current])
            @assert isempty(cache[current])
            delete!(cache, current)
            current = in(a, polygon) ? b : a
            push!(polygon, a)
            push!(polygon, b)
        end

        push!(polygons, collect(polygon))
    end

    return polygons
end

function simplify(polygons::Vector{Polygon})
    map(simplify, polygons)
end

function simplify(polygon::Polygon)
    [b for ((a, b), (c, d)) in edges(collect(edges(polygon))) if !iscollinear(b - a, d - c)]
end

function iscollinear(a::Vertex, b::Vertex)
    isapprox(norm(cross(a, b)), 0, atol=1e-2)
end

function correct(polygons::Vector{Polygon})
    corrected = Vector{Polygon}()
    boxes = map(boundingbox, polygons)
    for (i, box) in enumerate(boxes)
        num_enclosing = sum(map(other -> issubset(box, other) ? 1 : 0, boxes))
        if iseven(num_enclosing)
            push!(corrected, isclockwise(polygons[i]) ? reverse(polygons[i]) : polygons[i])
        else
            push!(corrected, !isclockwise(polygons[i]) ? reverse(polygons[i]) : polygons[i])
        end
    end
    corrected
end


# @testset "fit moments" begin
# end;

# @testset "fit to quantilepoint and mode - ignore mode" begin
# end;

# @testset "fit two quantiles same" begin
# end;

@testset "fit two quantiles" begin
    qpl = @qp_m(0.5)
    qpu = @qp_uu(5)
    d = fit(Gamma, qpl, qpu);
    @test median(d) ≈ 0.5
    @test quantile(d, 0.975) ≈ 5
    #Plots.plot(d);
end;

i_tmp = () -> begin
    dw = fit(Weibull, @qp_m(0.5), @qp_uu(5));
    dg = fit(Gamma, @qp_m(0.5), @qp_uu(5));
    Plots.plot(dw)
    Plots.plot!(dg, xlim=(0,0.2))
    pdf(dw, 0)
    pdf(dg, 0)
    pdf(dw, 0.01)
    pdf(dg, 0.01)
end


# Advent of Code, Day 20
#
# Transform an image by successively replacing pixels with values
# looked up from a translation table, the index being the value
# of the 9 cells surrounding each pixel, converted from binary to
# decimal. Much complicated by the fact that the input data has a
# 1 in the first position of the translation table, meaning that
# empty areas are filled with 1s, which muck up the pixel count.
#
# AK, 20/12/2021

# Read data into list of rows
f = "sample.txt"
f = "input.txt"
d = split(strip(read(f, String)), "\n")

# First row is "algorithm", convert # to 1 and . to 0
code = map(x -> x == '#' ? 1 : 0, [c for c in d[1]])
#println(code)

# Remaining rows are the image, convert to matrix
rows = map(r -> [c == '#' ? 1 : 0 for c in r], d[3:end])
image = Matrix(transpose(reshape(vcat(rows...), (:, length(rows)))))

# Transform an image by iteratively applying translation function
function transform(img, iterations)

    for iter in 1:iterations

        # If the code starts with a 1, it causes the whole image to
        # flip on each second iteration. To undo this, we neeed to
        # fill the border with 0 on the first iteration, 1 on the
        # second iteration, and so on.
        if code[1] == 1 && iseven(iter)
            bfill = 1
        else
            bfill = 0
        end

        # Enhance the image and display number of on bits
        img = enhance(img, bfill)
        #display(img)
        println("Iteration ", iter, ": bits set = ", sum(img))

    end
end

# Enhance the image, by converting 3x3 areas around a pixel
# to 9-digit binary number, converting this to decimal, and
# then looking up replacement bit in the "algorithm" vector
function enhance(img, padWith)

    # Make a copy of the image with a 2-thick border
    # around it, with the value of the padding alternating
    # to get around the flashing image problem when there
    # is a 1 at the beginning of the translation code.
    padded = addPadding(img, 2, padWith)

    # Make an empty array for the output, one bigger
    # on each side then the original image
    nrows, ncols = size(img)
    out = zeros(Int, nrows+2, ncols+2)

    # Go though the original image, and replace pixels
    # based on algorithm lookup
    for r in 1:nrows+2
        for c in 1:ncols+2
            n = numAt(padded, r+1, c+1)
            out[r,c] = code[Int(n+1)]
        end
    end

    #Return the enhanced  image
    return out

end

# Return the "number" centered at the given coordinates, i.e.,
# take the 3x3 cells around that number, string them along
# row-wise, and convert that binary representation to a decimal
# number
function numAt(img, r, c)

    # Get the 3x3 bits around the given coordinates
    # and turn them into a vector
    bits = reshape(transpose(img[r-1:r+1,c-1:c+1]), 1, :)

    # Convert this binary sequence to a decimal number
    #return bin2dec(bits)
    n = 0
    for b in bits
        n = n * 2 + b
    end
    return n

end

# Add a given number of layer of padding around a matrix,
# using the provide fill value (1/0)
function addPadding(img, n, padWith)

    # Add padding on top and bottom
    nr, nc = size(img)
    v = padWith == 1 ? ones(Int, n, nc) : zeros(Int, n, nc)
    img = vcat(v, img, v)

    # Add padding on sides
    nr, nc = size(img)
    v = padWith == 1 ? ones(Int, nr, n) : zeros(Int, nr, n)
    img = hcat(v, img, v)

    return img

end

# Perform the two transformations (Part 1 has 2 iterations, Part 2 has 50).
# For sample data, should be 35 and 3351, respectively
println("PART 1:")
transform(image, 2)

println("\nPART 2:")
transform(image, 50)


struct WendlandC8{T} <: SPHKernel
    n_neighbours::Int64
    norm_2D::T
    norm_3D::T
end

"""
    WendlandC8(T::DataType=Float64, n_neighbours::Integer=395)

Set up a `WendlandC8` kernel for a given DataType `T`.
"""
WendlandC8(T::DataType=Float64, n_neighbours::Integer=395) = WendlandC8{T}(n_neighbours, 8.0/(3π), 357.0/(64π))

"""
    kernel_value_2D(kernel::WendlandC8{T}, u::Real, h_inv::Real) where T

Evaluate WendlandC6 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_value_2D(kernel::WendlandC8{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_2D * h_inv^2
        t1 = 1 - u
        t9 = t1*t1*t1*t1*t1*t1*t1*t1*t1  # (1.0 - u)^9
        u2 = u*u
        return ( t1 * t9 * (5 + 50u + 210u2 + 450u2 * u + 429u2 * u2)) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    kernel_deriv_2D(kernel::WendlandC8{T}, u::Real, h_inv::Real) where T

Evaluate the derivative of the WendlandC6 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_deriv_2D(kernel::WendlandC8{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_2D * h_inv^3
        t1 = 1.0 - u
        t9 = t1*t1*t1*t1*t1*t1*t1*t1*t1  # (1.0 - u)^9
        u2 = u*u
        return -26t9 * u * (231u2 * u + 159u2 + 45u + 5 ) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    bias_correction_2D(kernel::WendlandC8, density::Real, m::Real, h_inv::Real)

Corrects the density estimate for the kernel bias. See Dehnen&Aly 2012, eq. 18+19.
"""
function bias_correction_2D(kernel::WendlandC8{T}, density::Real, m::Real, h_inv::Real) where T
    return density |> T
end

"""
    kernel_value_3D(kernel::WendlandC8{T}, u::Real, h_inv::Real) where T

Evaluate WendlandC8 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_value_3D(kernel::WendlandC8{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_2D * h_inv^3
        t1 = (1 - u)
        t9 = t1*t1*t1*t1*t1*t1*t1*t1*t1  # (1.0 - u)^9
        u2 = u*u
        return ( t1 * t9 * (5 + 50u + 210u2 + 450u2 * u + 429u2 * u2)) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    kernel_deriv_3D(kernel::WendlandC8, u::Real, h_inv::Real)

Evaluate the derivative of the WendlandC8 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_deriv_3D(kernel::WendlandC8{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_2D * h_inv^2 * h_inv^2
        t1 = 1 - u
        t9 = t1*t1*t1*t1*t1*t1*t1*t1*t1  # (1.0 - u)^9
        u2 = u*u
        return -26t9 * u * (231u2 * u + 159u2 + 45u + 5 ) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    bias_correction_3D(kernel::WendlandC8{T}, density::Real, m::Real, h_inv::Real) where T

Corrects the density estimate for the kernel bias. See Dehnen&Aly 2012, eq. 18+19.
"""
function bias_correction_3D(kernel::WendlandC8{T}, density::Real, m::Real, h_inv::Real) where T
    return density |> T
end


function randpoly(R::Union{ FqNmodPolyRing, FqPolyRing }, dmax::Int)
  r = R()
  F = base_ring(R)
  d = rand(0:dmax)
  for i = 0:d
    setcoeff!(r, i, rand(F))
  end
  return r
end

@testset "Extensions of finite fields" begin
  Fp = GF(2)
  Fpx, x = Fp["x"]
  f = x^3 + x + 1
  Fq = FqNmodFiniteField(f, :$)
  a = gen(Fq)
  Fqy, y = Fq["y"]
  g = y^3 + a*y + a^2
  G, FqytoG = Hecke.field_extension(g)
  @test characteristic(G) == 2
  @test degree(G) == 9
  GtoFqy = pseudo_inv(FqytoG)

  # Check if FqytoG does something reasonable with constant polynomials
  @test iszero(FqytoG(Fqy(1))^8 - FqytoG(Fqy(1)))
  @test iszero(FqytoG(Fqy(a))^8 - FqytoG(Fqy(a)))
  @test iszero(FqytoG(Fqy(a^2))^8 - FqytoG(Fqy(a^2)))

  # Check if FqytoG is linear
  for i = 1:5
    h1 = randpoly(Fqy, 5)
    h2 = randpoly(Fqy, 5)
    @test mod(h1, g) == GtoFqy(FqytoG(h1))
    @test mod(h2, g) == GtoFqy(FqytoG(h2))
    @test FqytoG(h1 + h2) == FqytoG(h1) + FqytoG(h2)
    @test FqytoG(h1*h2) == FqytoG(h1)*FqytoG(h2)
  end
  @test iszero(FqytoG(Fqy()))

  # Now everything with FqFiniteFields
  Fp = GF(fmpz(2))
  Fpx, x = Fp["x"]
  f = x^3 + x + 1
  Fq = FqFiniteField(f, :$)
  a = gen(Fq)
  Fqy, y = Fq["y"]
  g = y^3 + a*y + a^2
  G, FqytoG = Hecke.field_extension(g)
  @test characteristic(G) == 2
  @test degree(G) == 9
  GtoFqy = pseudo_inv(FqytoG)

  # Check if FqytoG does something reasonable with constant polynomials
  @test iszero(FqytoG(Fqy(1))^8 - FqytoG(Fqy(1)))
  @test iszero(FqytoG(Fqy(a))^8 - FqytoG(Fqy(a)))
  @test iszero(FqytoG(Fqy(a^2))^8 - FqytoG(Fqy(a^2)))

  # Check if FqytoG is linear
  for i = 1:5
    h1 = randpoly(Fqy, 5)
    h2 = randpoly(Fqy, 5)
    @test mod(h1, g) == GtoFqy(FqytoG(h1))
    @test mod(h2, g) == GtoFqy(FqytoG(h2))
    @test FqytoG(h1 + h2) == FqytoG(h1) + FqytoG(h2)
    @test FqytoG(h1*h2) == FqytoG(h1)*FqytoG(h2)
    @test iszero(FqytoG(Fqy()))
  end
end


using ApproxOperator

P = 1000.0
Ẽ = 3e6
# ν̃ = 0.3
ν̃ = 0.49999
E = Ẽ/(1-ν̃^2)
ν = ν̃/(1-ν̃)
L = 48.0
D = 12.0
I = D^3/12
EI = E*I

u = (x,y,z) -> (
    -P*y/6/EI*((6*L-3x)*x + (2+ν)*(y^2-D^2/4)),
     P/6/EI*(3*ν*y^2*(L-x) + (4+5*ν)*D^2*x/4 + (3*L-x)*x^2),
    -P/EI*(L-x)*y,
    -P/6/EI*((6*L-3*x)*x + (2+ν)*(3*y^2-D^2/4)),
     P/6/EI*((6*L-3*x)*x - 3*ν*y^2 + (4+5*ν)*D^2/4),
     P/EI*(L-x)*y*ν
)
t = (x,y,z) -> (0.0,P/2/I*(D^2/4-y^2))
b = (x,y,z) -> (0.0,0.0)
n = (x,y,z) -> (1.0,0.0,1.0)

aps = import_msh("./msh/cantilever.msh")
nₚ = length(aps["Domain"][1].nodes)

ops = [
    # 1
    PlaneStress_Ω(b,E,ν),
    # 2
    PlaneStrain_Ωᵛ(Ẽ,ν̃),
    # 3
    PlaneStrain_Ωᵈ(b,Ẽ,ν̃),
    # 4
    PlaneStress_Γᵗ(t),
    # 5
    PlaneStress_Γᵍ_penalty(u,n,E*1e7),
    # 6
    HₑError_PlaneStress(u,E,ν),
    # 7
    VTKExport(
                filename="vtk/cantilever_locking.vtk",
                topic="incompressible test for cantilever beam problem",
                pointdata=Dict{String,Symbol}("disp"=>:u_2D),
                celldata=Dict{String,Symbol}("stress"=>:σ_PlaneStrain),
                parameters=Dict{String,Float64}("E"=>Ẽ,"ν"=>ν̃)
            )
]

## locking
k = zeros(2nₚ,2nₚ)
f = zeros(2nₚ)

ops[1](aps["Domain"],k,f)
ops[4](aps["Traction"],f)
ops[5](aps["EssentialBC"],k,f)

d = k\f

Hₑ_Locking,L₂_Locking = ops[6](aps["Domain"],d)

ops[7](aps["Domain"],d)

## locking-free
k = zeros(2nₚ,2nₚ)
f = zeros(2nₚ)

ops[3](aps["Domain"],k,f)
set_integration_rule!(aps["Domain"],:QuadGI1)
ops[2](aps["Domain"],k,f)
ops[4](aps["Traction"],f)
ops[5](aps["EssentialBC"],k,f)

d = k\f

set_integration_rule!(aps["Domain"],:QuadGI2)
Hₑ_Free,L₂_Free = ops[6](aps["Domain"],d)

ops[7].filename = "vtk/cantilever_free.vtk"
ops[7](aps["Domain"],d)


#=
let E be a symmetric matrix sparsity pattern:
(1) find sparse PSD matrix with given sparsity pattern, "nearest" to A
    max_X tr(A, X) :
    tr(X) = 1
    X in intersect(S_+^n, S^n(E))
(2) find sparse PSD-completable matrix with given sparsity pattern, "nearest" to A
    max_X tr(A, X) :
    tr(X) = 1
    X in proj_E(S_+^n)

adapted from
"Decomposition methods for sparse matrix nearness problems"
(2015) by Sun & Vandenberghe
=#

using SparseArrays

struct NearestPSDJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    side::Int
    use_completable::Bool # solve problem (2) above, else solve problem (1)
    use_chordal_sparsity::Bool # use a chordal sparsity pattern
    use_sparsepsd::Bool # use sparse PSD cone formulation, else dense PSD cone
    use_cholmod_impl::Bool # use sparsePSD CHOLMOD implementation, else dense
end

function build(inst::NearestPSDJuMP{T}) where {T <: Float64}
    side = inst.side
    sparsity = min(3.0 / side, 1.0) # sparsity factor (before chordal extension)

    # generate random symmetric A (indefinite) with sparsity pattern E (nonchordal)
    A = tril!(sprandn(side, side, sparsity)) + Diagonal(randn(side))
    if inst.use_chordal_sparsity
        # compute a (heuristic) chordal extension of A using CHOLMOD functions
        copyto!(A, I)
        A = sparse(cholesky(Symmetric(A, :L)).L)
        (row_idxs, col_idxs, _) = findnz(A)
        A_vals = randn(length(row_idxs))
    else
        (row_idxs, col_idxs, A_vals) = findnz(A)
    end
    diag_idxs = findall(row_idxs .== col_idxs)

    model = JuMP.Model()

    if inst.use_sparsepsd || !inst.use_completable
        JuMP.@variable(model, X[1:length(row_idxs)])
        JuMP.@objective(model, Max, 2 * dot(A_vals, X) -
            sum(A_vals[k] * X[k] for k in diag_idxs)) # tr(A, X)
        JuMP.@constraint(model, sum(X[diag_idxs]) == 1) # tr(X) == 1

        if inst.use_sparsepsd
            rt2 = sqrt(2)
            X_scal = [X[k] * (row_idxs[k] == col_idxs[k] ? 1.0 : rt2)
                for k in eachindex(X)]
            impl = (inst.use_cholmod_impl ? Cones.PSDSparseCholmod :
                Cones.PSDSparseDense)
            cone = Hypatia.PosSemidefTriSparseCone{impl, T, T}
            JuMP.@constraint(model, X_scal in cone(side, row_idxs,
                col_idxs, inst.use_completable))
        else
            X_sparse = sparse(row_idxs, col_idxs, X)
            JuMP.@SDconstraint(model, Symmetric(Matrix(X_sparse), :L) >= 0)
        end
    else
        @assert inst.use_completable
        JuMP.@variable(model, X[1:side, 1:side], PSD)
        # tr(A, X):
        JuMP.@objective(model, Max, 2 * sum(X[row_idxs[k], col_idxs[k]] * A_vals[k]
            for k in eachindex(row_idxs)) - dot(A_vals[diag_idxs], diag(X)))
        JuMP.@constraint(model, tr(X) == 1) # tr(X) == 1
    end

    return model
end


type LikelihoodModel{X, Y}
  xs :: Vector{X}
  ys :: Vector{Y}
  zs :: Vector{Y}
  datas :: Vector{Y}
  Ldx :: Matrix
  Lzx :: Matrix
end

# constructors

LikelihoodModel(xs, ys, zs, datas, measerr) =
  LikelihoodModel(xs, ys, zs, datas, likelihoodmat(datas, ys, measerr), likelihoodmat(zs, ys, measerr))

function LikelihoodModel(; xs=[], phi=Void, ys=[], datas=[], measerr=Void, prior=Void, ndata=0, zs=[], zmult=0, sfact=0, smoothkernel=Void)

  if ys == [] && phi != Void
    ys = phi.(xs)
  end

  if datas == [] && ndata > 0
    datas = phi.(rand(prior, ndata)) + rand(measerr, ndata)
  end

  if zs == [] && zmult > 0
    zs = repmat(ys, zmult) + rand(measerr, length(ys)*zmult)
  end

  if sfact > 0
    datas = smoothdata(smoothkernel, sfact, datas)
    measerr = convolute(measerr, smoothkernel)
  end

  m = LikelihoodModel(xs, ys, zs, datas, measerr)
end

# likelihood computation

likelihoodmat(xs, ys, d::Distribution) = map(a->pdf(d, a), (x-y for x in xs, y in ys))

# model wrappers

npmle(m)     = mple(m, 0)
refprior(m)  = mple(m, 1)

logl(m::LikelihoodModel, w) = logl(w, m.Ldx)
dlogl(m::LikelihoodModel, w) = dlogl(w, m.Ldx)

hz(m::LikelihoodModel, w)  =  hz(w, m.Lzx)
dhz(m::LikelihoodModel, w) = dhz(w, m.Lzx)

function mple_obj(m::LikelihoodModel, reg)
  if reg == 0
    w -> logl(m, w)
  elseif reg == 1
    w -> hz(m, w)
  else
    w -> reg*hz(m,w) + (1-reg) * logl(m,w)
  end
end

function dmple_obj(m::LikelihoodModel, reg)
  if reg == 0
    w -> dlogl(m, w)
  elseif reg == 1
    w -> dhz(m,w)
  else
    w -> reg*dhz(m, w) + (1-reg) * dlogl(m, w)
  end
end


module test_ctriangulate
using Triangulate

function test()
    nodes=Matrix{Cdouble}([1.0 0.0 ; 0.0 1.0 ; -1.0 0.0 ; 0.0 -1.0]')
    faces=Matrix{Cint}([1 2 ; 2 3 ; 3 4 ; 4 1 ]')
    faceregions=Matrix{Cint}([1 2 3 4]')
    regionpoints=Matrix{Cdouble}([0.5 0.5 1 0.01;]')
    triin=Triangulate.CTriangulateIO()
    triout=Triangulate.CTriangulateIO()
    vorout=Triangulate.CTriangulateIO()

    triin.numberofpoints=Cint(size(nodes,2))
    triin.pointlist=pointer(nodes)
    triin.numberofsegments=size(faces,2)
    triin.segmentlist=pointer(faces)
    triin.segmentmarkerlist=pointer(faceregions)
    triin.numberofregions=size(regionpoints,2)
    triin.regionlist=pointer(regionpoints)

    triangulate("paAqQ",triin,triout,vorout)

    triout.numberofpoints==177 && triout.numberoftriangles==319 && triout.numberofsegments==33
end
end


using ITensors

d = 1
i = Index([QN(0, 2) => d, QN(1, 2) => d], "i")

# Parity conserving ITensors
# By default they have flux 0
@show A = randomITensor(i', dag(i))
println()
@show B = randomITensor(i', dag(i))
println()
@show C = randomITensor(QN(1, 2), i', dag(i))
println()

# Add them
@show A + B
println()

# Can't add QN ITensors with different flux
#A + C

# Contract them
@show A' * B
println()
@show A' * C
println()

# Combine the indices to turn into a vector
comb = combiner(i', dag(i))
@show A * comb
println()


module GPs

# https://gregorygundersen.com/blog/2019/09/12/practical-gp-regression/
# NOTE: If A * x = b. Then A \ b solves for x.
# NOTE: cholesky(A) \ b = L' \ (L \ b) = x, where L = cholesky(A).L
# NOTE: A \ b and cholesky(A) \ b are respectively implemented via variants of
# `solve` and `cholesky solve` in other languages.

using Distributions
using Distances

include(joinpath(@__DIR__, "Kernels.jl"))
using .Kernels

import LinearAlgebra
import LinearAlgebra: cholesky, Symmetric

export GP, posterior, Kernels

struct GP
  mu
  kernel
  y
  X
  sigma
  C
  d
  a
end

function GP(y, X, kernel; mu=0, sigma=0)
  Ddata = pairwise(metric(kernel), X, dims=1)
  Kdata = kernel(Ddata)
  C = cholesky(Symmetric(Kdata))
  d = y .- mu
  a = C \ d  # = L' \ (L \ d), where L = cholesky(Kdata).L
  return GP(mu, kernel, y, X, sigma, C, d, a)
end

function posterior(gp, Xnew)
  Dnew = pairwise(metric(gp.kernel), Xnew, dims=1)
  D_new_data = pairwise(metric(gp.kernel), Xnew, gp.X, dims=1)

  Knew = gp.kernel(Dnew)
  K_new_data = gp.kernel(D_new_data, jitter=false)

  mu = gp.mu .+ K_new_data * gp.a
  S = Knew - K_new_data * (gp.C \ K_new_data')

  return MvNormal(mu, Symmetric(S) + LinearAlgebra.I * gp.sigma)
end

quantiles(X, q; dims) = mapslices(v -> quantile(v, q), X, dims=dims)

function ci(preds, alpha=0.05, dims=2)
    q_upper = 1 - alpha / 2
    q_lower = alpha / 2
    lower = vec(quantiles(preds, q_lower, dims=dims))
    upper = vec(quantiles(preds, q_upper, dims=dims))
    return lower, upper
end

end


module MathOptSetDistances

import MathOptInterface
const MOI = MathOptInterface
using BlockDiagonals: BlockDiagonal
import FillArrays
using LinearAlgebra
import ChainRulesCore
const CRC = ChainRulesCore
import FiniteDifferences

export distance_to_set, projection_on_set, projection_gradient_on_set

include("utils.jl")
include("distances.jl")
include("distance_sets.jl")
include("projections.jl")
include("chainrules.jl")

end # module


using OrdinaryDiffEq
using Trixi

###############################################################################
# semidiscretization of the hyperbolic diffusion equations

equations = HyperbolicDiffusionEquations3D()

function initial_condition_poisson_periodic(x, t, equations::HyperbolicDiffusionEquations3D)
  # elliptic equation: -νΔϕ = f
  # depending on initial constant state, c, for phi this converges to the solution ϕ + c
  if iszero(t)
    phi = 0.0
    q1  = 0.0
    q2  = 0.0
    q3  = 0.0
  else
    phi =          sin(2 * pi * x[1]) * sin(2 * pi * x[2]) * sin(2 * pi * x[3])
    q1  = 2 * pi * cos(2 * pi * x[1]) * sin(2 * pi * x[2]) * sin(2 * pi * x[3])
    q2  = 2 * pi * sin(2 * pi * x[1]) * cos(2 * pi * x[2]) * sin(2 * pi * x[3])
    q3  = 2 * pi * sin(2 * pi * x[1]) * sin(2 * pi * x[2]) * cos(2 * pi * x[3])
  end
  return SVector(phi, q1, q2, q3)
end
initial_condition = initial_condition_poisson_periodic

@inline function source_terms_poisson_periodic(u, x, t, equations::HyperbolicDiffusionEquations3D)
  # elliptic equation: -νΔϕ = f
  # analytical solution: phi = sin(2πx)*sin(2πy) and f = -8νπ^2 sin(2πx)*sin(2πy)
  @unpack inv_Tr = equations
  C = -12 * equations.nu * pi^2

  x1, x2, x3 = x
  tmp1 = sinpi(2 * x1)
  tmp2 = sinpi(2 * x2)
  tmp3 = sinpi(2 * x3)
  du1 = -C*tmp1*tmp2*tmp3
  du2 = -inv_Tr * u[2]
  du3 = -inv_Tr * u[3]
  du4 = -inv_Tr * u[4]

  return SVector(du1, du2, du3, du4)
end

solver = DGSEM(polydeg=3, surface_flux=flux_lax_friedrichs)

coordinates_min = (0.0, 0.0, 0.0)
coordinates_max = (1.0, 1.0, 1.0)
mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=3,
                n_cells_max=30_000)

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver,
                                    source_terms=source_terms_poisson_periodic)

###############################################################################
# ODE solvers, callbacks etc.

tspan = (0.0, 2.0)
ode = semidiscretize(semi, tspan);

summary_callback = SummaryCallback()

resid_tol = 5.0e-12
steady_state_callback = SteadyStateCallback(abstol=resid_tol, reltol=0.0)

analysis_interval = 200
analysis_callback = AnalysisCallback(semi, interval=analysis_interval,
                                     extra_analysis_integrals=(entropy, energy_total))

alive_callback = AliveCallback(analysis_interval=analysis_interval)

save_solution = SaveSolutionCallback(interval=100,
                                     save_initial_solution=true,
                                     save_final_solution=true,
                                     solution_variables=cons2prim)

stepsize_callback = StepsizeCallback(cfl=2.4)

callbacks = CallbackSet(summary_callback, steady_state_callback,
                        analysis_callback, alive_callback,
                        save_solution,
                        stepsize_callback)

###############################################################################
# run the simulation

sol = Trixi.solve(ode, Trixi.HypDiffN3Erk3Sstar52(),
                  dt=1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
                  save_everystep=false, callback=callbacks);
summary_callback() # print the timer summary


using Test,TightBinding
using Plots
using LinearAlgebra

@testset "Real to k" begin
    @testset "1D" begin
        @time energies = TightBinding.test_1D()
        sumt = 26.21913117088203
        @test round(sum(abs.(energies))*100) == round(sumt*100)
    end
    @testset "2Dsquare" begin
        sumt = 52.438262341764066
        @test round(sum(abs.(TightBinding.test_2Dsquare()))*100) == round(sumt*100)
    end
    @testset "Graphene" begin
        @time hamk = TightBinding.test_2DGraphene()
        sumt = 5.975026051875501
        @test round(sum(abs.(hamk))*100) == round(sumt*100)
#        @time test_Graphene()
    end
    @testset "Iron pnictides" begin
        @time hamk =  TightBinding.test_pnictides()
        sumt = 2.680985878019015
        @test round(sum(abs.(hamk))*100) == round(sumt*100)

        @time hamk = TightBinding.test_pnictides_5orbitals()
        sumt = 3.0141792528265308
        @test round(sum(abs.(hamk))*100) == round(sumt*100)

        #@time test_Fe2()
        #@time test_Fe5()
    end

end

@testset "k to real" begin
    @testset "Topological insulator" begin
        @time mat_e = TightBinding.test_surface()
        sumt = 17011.010972794524
        @test round(sum(abs.(mat_e))*100) == round(sumt*100)
    end
end


time_start = time()

using PowerModels
using Nonconvex
Nonconvex.@load Ipopt

pkg_load_time = time() - time_start

time_start = time()

file_name = "../data/pglib_opf_case5_pjm.m"
#file_name = "../data/pglib_opf_case118_ieee.m"

data = PowerModels.parse_file(file_name)
PowerModels.standardize_cost_terms!(data, order=2)
PowerModels.calc_thermal_limits!(data)
ref = PowerModels.build_ref(data)[:it][:pm][:nw][0]

data_load_time = time() - time_start

time_start = time()

model = Nonconvex.Model()

bus_pd = Dict(i => 0.0 for (i,bus) in ref[:bus])
bus_qd = Dict(i => 0.0 for (i,bus) in ref[:bus])

bus_gs = Dict(i => 0.0 for (i,bus) in ref[:bus])
bus_bs = Dict(i => 0.0 for (i,bus) in ref[:bus])

for (i,bus) in ref[:bus]
    if length(ref[:bus_loads][i]) > 0
        bus_pd[i] = sum(ref[:load][l]["pd"] for l in ref[:bus_loads][i])
        bus_qd[i] = sum(ref[:load][l]["qd"] for l in ref[:bus_loads][i])
    end

    if length(ref[:bus_shunts][i]) > 0
        bus_gs[i] = sum(ref[:shunt][s]["gs"] for s in ref[:bus_shunts][i])
        bus_bs[i] = sum(ref[:shunt][s]["bs"] for s in ref[:bus_shunts][i])
    end
end

br_g = Dict(i => 0.0 for (i,branch) in ref[:branch])
br_b = Dict(i => 0.0 for (i,branch) in ref[:branch])

br_tr = Dict(i => 0.0 for (i,branch) in ref[:branch])
br_ti = Dict(i => 0.0 for (i,branch) in ref[:branch])
br_tm = Dict(i => 0.0 for (i,branch) in ref[:branch])

br_g_fr = Dict(i => 0.0 for (i,branch) in ref[:branch])
br_b_fr = Dict(i => 0.0 for (i,branch) in ref[:branch])
br_g_to = Dict(i => 0.0 for (i,branch) in ref[:branch])
br_b_to = Dict(i => 0.0 for (i,branch) in ref[:branch])

for (i,branch) in ref[:branch]
    g, b = PowerModels.calc_branch_y(branch)
    tr, ti = PowerModels.calc_branch_t(branch)

    br_g[i] = g
    br_b[i] = b

    br_tr[i] = tr
    br_ti[i] = ti
    br_tm[i] = tr^2 + ti^2

    br_g_fr[i] = branch["g_fr"]
    br_b_fr[i] = branch["b_fr"]
    br_g_to[i] = branch["g_to"]
    br_b_to[i] = branch["b_to"]
end

var_lookup = Dict{String,Int}()

var_idx = 1
for (i,bus) in ref[:bus]
    addvar!(model, -Inf, Inf) #va
    var_lookup["va_$(i)"] = var_idx
    global var_idx += 1

    addvar!(model, bus["vmin"], bus["vmax"], init=1.0) #vm
    var_lookup["vm_$(i)"] = var_idx
    global var_idx += 1
end

for (i,gen) in ref[:gen]
    addvar!(model, gen["pmin"], gen["pmax"]) #pg
    var_lookup["pg_$(i)"] = var_idx
    global var_idx += 1

    addvar!(model, gen["qmin"], gen["qmax"]) #qg
    var_lookup["qg_$(i)"] = var_idx
    global var_idx += 1
end

for (l,i,j) in ref[:arcs]
    branch = ref[:branch][l]

    addvar!(model, -branch["rate_a"], branch["rate_a"]) #p
    var_lookup["p_$(l)_$(i)_$(j)"] = var_idx
    global var_idx += 1

    addvar!(model, -branch["rate_a"], branch["rate_a"]) #q
    var_lookup["q_$(l)_$(i)_$(j)"] = var_idx
    global var_idx += 1
end

@assert var_idx == length(var_lookup)+1

function opf_objective(x::Vector)
    cost = 0.0
    for (i,gen) in ref[:gen]
        pg = x[var_lookup["pg_$(i)"]]
        cost += gen["cost"][1]*pg^2 + gen["cost"][2]*pg + gen["cost"][3]
    end
    return cost
end

set_objective!(model, opf_objective)

#println(model)

model_build_time = time() - time_start

time_start = time()

alg = IpoptAlg()
options = IpoptOptions(print_level = 0)
x0 = NonconvexCore.getinit(model)
#r = optimize(model, alg, x0, options = options)

# Nonconvex v1.0.2
# ERROR: LoadError: MethodError: no method matching getindex(::Dict{Any, Any})
# something to do with :Zygote it seems
r = optimize(model, alg, x0)

solve_time = time() - time_start

println("")
println("\033[1mSummary\033[0m")
println("   case......: $(file_name)")
#println("   cost......: $(round(Int, cost))")
println("   pkg time..: $(pkg_load_time)")
println("   data time.: $(data_load_time)")
println("   build time: $(model_build_time)")
println("   solve time: $(solve_time)")
println("")


__precompile__()

module Projections

export project, project!

abstract type ConvexSet end
abstract type NonconvexSet end

"""
    project!(s, v, y)

Overwrite `v` by projection of point `y` to the box set `b`.
"""
function project!(); end

"""
    project(s, y)

Project point `y` to the box set `b`.
"""
function project(); end

include("box.jl")
include("ball.jl")
include("simplex.jl")
include("affine.jl")
include("consensus.jl")
include("finite.jl")
include("halfspace.jl")
include("l1ball.jl")
include("l_infinity.jl")
include("low-rank-matrix.jl")
include("positive-semidefinite.jl")
include("second-order-cone.jl")
include("stiefel-manifold.jl")
end


function collatz(n::Integer)
    n <= 0 && throw(DomainError())
    iseven(n) ? div(n, 2) : 3n + 1
end

function collatz_steps(n::Integer)
    i = 0
    while n != 1
        n = collatz(n)
        i += 1
    end

    i
end