#Unscented Kalman Inversion: specific structures and function definitions

"""
    Unscented{FT<:AbstractFloat, IT<:Int} <: Process

An unscented Kalman Inversion process.
"""
mutable struct Unscented{FT <: AbstractFloat, IT <: Int} <: Process
    "a vector of arrays of size `N_parameters` containing the mean of the parameters (in each `uki` iteration a new array of mean is added)"
    u_mean::Vector{Vector{FT}}
    "a vector of arrays of size (`N_parameters x N_parameters`) containing the covariance of the parameters (in each `uki` iteration a new array of `cov` is added)"
    uu_cov::Vector{Matrix{FT}}
    "a vector of arrays of size `N_y` containing the predicted observation (in each `uki` iteration a new array of predicted observation is added)"
    obs_pred::Vector{Vector{FT}}
    "weights in UKI"
    c_weights::Vector{FT}
    mean_weights::Vector{FT}
    cov_weights::Vector{FT}
    "covariance of the artificial evolution error"
    Σ_ω::Matrix{FT}
    "covariance of the artificial observation error"
    Σ_ν_scale::FT
    "regularization parameter"
    α_reg::FT
    "regularization vector"
    r::Vector{FT}
    "update frequency"
    update_freq::IT
    "current iteration number"
    iter::IT
end

# outer constructors
function EnsembleKalmanProcess(
    obs_mean::Vector{FT},
    obs_noise_cov::Matrix{FT},
    process::Unscented{FT, IT};
    Δt = FT(1),
) where {FT <: AbstractFloat, IT <: Int}

    #initial parameters stored as columns
    init_params = [DataContainer(update_ensemble_prediction!(process), data_are_columns = true)]
    # Number of parameter
    N_u = length(process.u_mean[1])
    # ensemble size
    N_ens = 2N_u + 1 #stored with data as columns
    #store for model evaluations
    g = []
    # error store
    err = FT[]
    # timestep store
    Δt = Array([Δt])

    EnsembleKalmanProcess{FT, IT, Unscented}(init_params, obs_mean, obs_noise_cov, N_ens, g, err, Δt, process)
end

"""
    Unscented(
        u0_mean::Vector{FT},
        uu0_cov::Matrix{FT},
        α_reg::FT = 1.0;
        update_freq::IT = 1,
        modified_unscented_transform::Bool = true,
        prior_mean::Union{Vector{FT}, Nothing} = nothing,
    ) where {FT <: AbstractFloat, IT <: Int}

Construct an Unscented Inversion EnsembleKalmanProcess.

Arguments
=========

  - `u0_mean`: Mean at initialization.
  - `uu0_cov`: Covariance at initialization.
  - `α_reg`: Hyperparameter controlling regularization toward the prior mean (0 < `α_reg` ≤ 1),
  default should be 1, without regulariazion.
  - `update_freq`: Set to 0 when the inverse problem is not identifiable,
  namely the inverse problem has multiple solutions, the covariance matrix
  will represent only the sensitivity of the parameters, instead of
  posterior covariance information; set to 1 (or anything > 0) when
  the inverse problem is identifiable, and the covariance matrix will
  converge to a good approximation of the posterior covariance with an
  uninformative prior.
  - `modified_unscented_transform`: Modification of the UKI quadrature given
    in Huang et al (2021).
  - `prior_mean`: Prior mean used for regularization.
"""
function Unscented(
    u0_mean::Vector{FT},
    uu0_cov::Matrix{FT};
    α_reg::FT = 1.0,
    update_freq::IT = 1,
    modified_unscented_transform::Bool = true,
    prior_mean::Union{Vector{FT}, Nothing} = nothing,
) where {FT <: AbstractFloat, IT <: Int}

    N_u = size(u0_mean, 1)
    # ensemble size
    N_ens = 2 * N_u + 1

    c_weights = zeros(FT, N_u)
    mean_weights = zeros(FT, N_ens)
    cov_weights = zeros(FT, N_ens)

    # set parameters λ, α
    α = min(sqrt(4 / N_u), 1.0)
    λ = α^2 * N_u - N_u

    c_weights[1:N_u] .= sqrt(N_u + λ)
    mean_weights[1] = λ / (N_u + λ)
    mean_weights[2:N_ens] .= 1 / (2 * (N_u + λ))
    cov_weights[1] = λ / (N_u + λ) + 1 - α^2 + 2.0
    cov_weights[2:N_ens] .= 1 / (2 * (N_u + λ))

    if modified_unscented_transform
        mean_weights[1] = 1.0
        mean_weights[2:N_ens] .= 0.0
    end

    u_mean = Vector{FT}[]  # array of Vector{FT}'s
    push!(u_mean, u0_mean) # insert parameters at end of array (in this case just 1st entry)
    uu_cov = Matrix{FT}[]  # array of Matrix{FT}'s
    push!(uu_cov, uu0_cov) # insert parameters at end of array (in this case just 1st entry)

    obs_pred = Vector{FT}[]  # array of Vector{FT}'s

    Σ_ω = (2 - α_reg^2) * uu0_cov
    Σ_ν_scale = 2.0

    r = isnothing(prior_mean) ? u0_mean : prior_mean
    iter = 0

    Unscented(
        u_mean,
        uu_cov,
        obs_pred,
        c_weights,
        mean_weights,
        cov_weights,
        Σ_ω,
        Σ_ν_scale,
        α_reg,
        r,
        update_freq,
        iter,
    )
end

"""
    construct_sigma_ensemble(
        process::Unscented,
        x_mean::Array{FT},
        x_cov::Matrix{FT},
    ) where {FT <: AbstractFloat, IT <: Int}

Construct the sigma ensemble based on the mean `x_mean` and covariance `x_cov`.
"""
function construct_sigma_ensemble(
    process::Unscented,
    x_mean::Array{FT},
    x_cov::Matrix{FT},
) where {FT <: AbstractFloat, IT <: Int}

    N_x = size(x_mean, 1)

    c_weights = process.c_weights

    chol_xx_cov = cholesky(Hermitian(x_cov)).L

    x = zeros(FT, N_x, 2 * N_x + 1)
    x[:, 1] = x_mean
    for i in 1:N_x
        x[:, i + 1] = x_mean + c_weights[i] * chol_xx_cov[:, i]
        x[:, i + 1 + N_x] = x_mean - c_weights[i] * chol_xx_cov[:, i]
    end

    return x
end

"""
construct_mean `x_mean` from ensemble `x`.
"""
function construct_mean(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::Matrix{FT},
) where {FT <: AbstractFloat, IT <: Int}
    N_x, N_ens = size(x)

    @assert(uki.N_ens == N_ens)

    x_mean = zeros(FT, N_x)

    mean_weights = uki.process.mean_weights

    for i in 1:N_ens
        x_mean += mean_weights[i] * x[:, i]
    end

    return x_mean
end

"""
construct_cov `xx_cov` from ensemble `x` and mean `x_mean`.
"""
function construct_cov(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::Matrix{FT},
    x_mean::Array{FT},
) where {FT <: AbstractFloat, IT <: Int}
    N_ens, N_x = uki.N_ens, size(x_mean, 1)

    cov_weights = uki.process.cov_weights

    xx_cov = zeros(FT, N_x, N_x)

    for i in 1:N_ens
        xx_cov .+= cov_weights[i] * (x[:, i] - x_mean) * (x[:, i] - x_mean)'
    end

    return xx_cov
end

"""
construct_cov `xy_cov` from ensemble x and mean `x_mean`, ensemble `obs_mean` and mean `y_mean`.
"""
function construct_cov(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::Matrix{FT},
    x_mean::Array{FT},
    obs_mean::Matrix{FT},
    y_mean::Array{FT},
) where {FT <: AbstractFloat, IT <: Int, P <: Process}
    N_ens, N_x, N_y = uki.N_ens, size(x_mean, 1), size(y_mean, 1)

    cov_weights = uki.process.cov_weights

    xy_cov = zeros(FT, N_x, N_y)

    for i in 1:N_ens
        xy_cov .+= cov_weights[i] * (x[:, i] - x_mean) * (obs_mean[:, i] - y_mean)'
    end

    return xy_cov
end

"""
uki prediction step : generate sigma points
"""
function update_ensemble_prediction!(process::Unscented) where {FT <: AbstractFloat, IT <: Int}

    process.iter += 1
    # update evolution covariance matrix
    if process.update_freq > 0 && process.iter % process.update_freq == 0
        process.Σ_ω = (2 - process.α_reg^2) * process.uu_cov[end]
    end

    u_mean = process.u_mean[end]
    uu_cov = process.uu_cov[end]

    α_reg = process.α_reg
    r = process.r
    Σ_ω = process.Σ_ω

    N_u = length(process.u_mean[1])
    ############# Prediction step:

    u_p_mean = α_reg * u_mean + (1 - α_reg) * r
    uu_p_cov = α_reg^2 * uu_cov + Σ_ω

    ############ Generate sigma points
    u_p = construct_sigma_ensemble(process, u_p_mean, uu_p_cov)
    return u_p
end

"""
uki analysis step
g is the predicted observations  `Ny x N_ens` matrix
"""
function update_ensemble_analysis!(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    u_p::Matrix{FT},
    g::Matrix{FT},
) where {FT <: AbstractFloat, IT <: Int}

    obs_mean = uki.obs_mean
    Σ_ν = uki.process.Σ_ν_scale * uki.obs_noise_cov

    ############# Prediction step:

    u_p_mean = construct_mean(uki, u_p)
    uu_p_cov = construct_cov(uki, u_p, u_p_mean)

    ###########  Analysis step

    g_mean = construct_mean(uki, g)
    gg_cov = construct_cov(uki, g, g_mean) + Σ_ν
    ug_cov = construct_cov(uki, u_p, u_p_mean, g, g_mean)

    tmp = ug_cov / gg_cov

    u_mean = u_p_mean + tmp * (obs_mean - g_mean)
    uu_cov = uu_p_cov - tmp * ug_cov'

    ########### Save results
    push!(uki.process.obs_pred, g_mean) # N_ens x N_data
    push!(uki.process.u_mean, u_mean) # N_ens x N_params
    push!(uki.process.uu_cov, uu_cov) # N_ens x N_data

    push!(uki.g, DataContainer(g, data_are_columns = true))

    compute_error!(uki)
end

function update_ensemble!(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    g_in::Matrix{FT},
) where {FT <: AbstractFloat, IT <: Int}
    #catch works when g_in non-square
    if !(size(g_in)[2] == uki.N_ens)
        throw(DimensionMismatch("ensemble size in EnsembleKalmanProcess and g_in do not match, try transposing or check ensemble size"))
    end

    u_p_old = get_u_final(uki)

    #perform analysis on the model runs
    update_ensemble_analysis!(uki, u_p_old, g_in)
    #perform new prediction output to model parameters u_p
    u_p = update_ensemble_prediction!(uki.process)

    push!(uki.u, DataContainer(u_p, data_are_columns = true))

    return u_p
end

function get_u_mean_final(uki::EnsembleKalmanProcess{FT, IT, Unscented}) where {FT <: AbstractFloat, IT <: Int}
    return uki.process.u_mean[end]
end

function compute_error!(uki::EnsembleKalmanProcess{FT, IT, Unscented}) where {FT <: AbstractFloat, IT <: Int}
    mean_g = uki.process.obs_pred[end]
    diff = uki.obs_mean - mean_g
    X = uki.obs_noise_cov \ diff # diff: column vector
    newerr = dot(diff, X)
    push!(uki.err, newerr)
end

function get_error(uki::EnsembleKalmanProcess{FT, IT, Unscented}) where {FT <: AbstractFloat, IT <: Int}
    return uki.err
end

function Gaussian_2d(
    u_mean::Vector{FT},
    uu_cov::Matrix{FT},
    Nx::IT,
    Ny::IT;
    xx = nothing,
    yy = nothing,
) where {FT <: AbstractFloat, IT <: Int}
    # 2d Gaussian plot
    u_range = [min(5 * sqrt(uu_cov[1, 1]), 5); min(5 * sqrt(uu_cov[2, 2]), 5)]

    if xx === nothing
        xx = Array(LinRange(u_mean[1] - u_range[1], u_mean[1] + u_range[1], Nx))
    end
    if yy == nothing
        yy = Array(LinRange(u_mean[2] - u_range[2], u_mean[2] + u_range[2], Ny))
    end
    X, Y = repeat(xx, 1, Ny), Array(repeat(yy, 1, Nx)')
    Z = zeros(FT, Nx, Ny)

    det_uu_cov = det(uu_cov)

    for ix in 1:Nx
        for iy in 1:Ny
            Δxy = [xx[ix] - u_mean[1]; yy[iy] - u_mean[2]]
            Z[ix, iy] = exp(-0.5 * (Δxy' / uu_cov * Δxy)) / (2 * pi * sqrt(det_uu_cov))
        end
    end

    return xx, yy, Z
end


function tensor_create_linear(T::Union{Type{Tensor{order, dim}}, Type{SymmetricTensor{order, dim}}}, f) where {order, dim}
    return Expr(:tuple, [f(i) for i=1:n_components(T)]...)
end

function tensor_create(::Type{Tensor{order, dim}}, f) where {order, dim}
    if order == 1
        ex = Expr(:tuple, [f(i) for i=1:dim]...)
    elseif order == 2
        ex = Expr(:tuple, [f(i,j) for i=1:dim, j=1:dim]...)
    elseif order == 4
        ex = Expr(:tuple, [f(i,j,k,l) for i=1:dim, j=1:dim, k = 1:dim, l = 1:dim]...)
    end
    return ex
end

function tensor_create(::Type{SymmetricTensor{order, dim}}, f) where {order, dim}
    ex = Expr(:tuple)
    if order == 2
        for j in 1:dim, i in j:dim
            push!(ex.args, f(i, j))
        end
    elseif order == 4
        for l in 1:dim, k in l:dim, j in 1:dim, i in j:dim
            push!(ex.args, f(i, j, k, l))
        end
    end
    return ex
end

# reduced two expressions by summing the products
# madd = true uses muladd instructions which is faster
# in some cases, like in double contraction
function reducer(ex1i, ex2i, madd=false)
    ex1, ex2 = remove_duplicates(ex1i, ex2i)
    N = length(ex1)
    expr = :($(ex1[1]) * $(ex2[1]))

    for i in 2:N
        expr = madd ? :(muladd($(ex1[i]), $(ex2[i]), $expr)) :
                      :($(expr) + $(ex1[i]) * $(ex2[i]))
    end
    return expr
end

function remove_duplicates(ex1in, ex2in)
    ex1out, ex2out = Expr[], Expr[]
    exout = Expr[]
    factors = ones(Int, length(ex1in))

    for (ex1ine, ex2ine) in zip(ex1in, ex2in)
        prod = :($ex1ine * $ex2ine)
        i = findfirst(isequal(prod), exout) # check if this product exist in the output
        if i == nothing # this product does not exist yet
            push!(ex1out, ex1ine)
            push!(ex2out, ex2ine)
            push!(exout, prod)
        else # found a duplicate
            factors[i] += 1
        end
    end
    for i in 1:length(ex1out)
        factors[i] != 1 && (ex1out[i] = :($(factors[i]) * $(ex1out[i])))
    end
    return ex1out, ex2out
end

# check symmetry and return
remove_duplicates(::Type{<:Tensor}, ex) = ex # do nothing if return type is a Tensor
function remove_duplicates(::Type{SymmetricTensor{order, dim}}, ex) where {order, dim}
    ex.args = ex.args[SYMMETRIC_INDICES[order][dim]]
    return ex
end

# return types
# double contraction
function dcontract end
@pure getreturntype(::typeof(dcontract), ::Type{<:FourthOrderTensor{dim}}, ::Type{<:FourthOrderTensor{dim}}) where {dim} = Tensor{4, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:SymmetricTensor{4, dim}}, ::Type{<:SymmetricTensor{4, dim}}) where {dim} = SymmetricTensor{4, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:Tensor{4, dim}}, ::Type{<:SecondOrderTensor{dim}}) where {dim} = Tensor{2, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:SymmetricTensor{4, dim}}, ::Type{<:SecondOrderTensor{dim}}) where {dim} = SymmetricTensor{2, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:SecondOrderTensor{dim}}, ::Type{<:Tensor{4, dim}}) where {dim} = Tensor{2, dim}
@pure getreturntype(::typeof(dcontract), ::Type{<:SecondOrderTensor{dim}}, ::Type{<:SymmetricTensor{4, dim}}) where {dim} = SymmetricTensor{2, dim}

# otimes
function otimes end
@pure getreturntype(::typeof(otimes), ::Type{<:Tensor{1, dim}}, ::Type{<:Tensor{1, dim}}) where {dim} = Tensor{2, dim}
@pure getreturntype(::typeof(otimes), ::Type{<:SecondOrderTensor{dim}}, ::Type{<:SecondOrderTensor{dim}}) where {dim} = Tensor{4, dim}
@pure getreturntype(::typeof(otimes), ::Type{<:SymmetricTensor{2, dim}}, ::Type{<:SymmetricTensor{2, dim}}) where {dim} = SymmetricTensor{4, dim}


@testset "base" begin
    @testset "Trig" begin
        @testset "Basics" for x = (Float64(π)-0.01, Complex(π, π/2))
            test_scalar(sec, x)
            test_scalar(csc, x)
            test_scalar(cot, x)
            test_scalar(sinpi, x)
            test_scalar(cospi, x)
        end
        @testset "Hyperbolic" for x = (Float64(π)-0.01, Complex(π-0.01, π/2))
            test_scalar(sech, x)
            test_scalar(csch, x)
            test_scalar(coth, x)
        end
        @testset "Degrees" begin
            x = 45.0
            test_scalar(sind, x)
            test_scalar(cosd, x)
            test_scalar(tand, x)
            test_scalar(secd, x)
            test_scalar(cscd, x)
            test_scalar(cotd, x)
        end
        @testset "Inverses" for x = (0.5, Complex(0.5, 0.25))
            test_scalar(asec, 1/x)
            test_scalar(acsc, 1/x)
            test_scalar(acot, 1/x)
        end
        @testset "Inverse hyperbolic" for x = (0.5, Complex(0.5, 0.25),  Complex(-2.1 -3.1im))
            test_scalar(asinh, x)
            test_scalar(acosh, x + 1)  # +1 accounts for domain for real
            test_scalar(atanh, x)
            test_scalar(asech, x)
            test_scalar(acsch, x)
            test_scalar(acoth, x + 1)
        end

        @testset "Inverse degrees" for x = (0.5, Complex(0.5, 0.25))
            test_scalar(asind, x)
            test_scalar(acosd, x)
            test_scalar(atand, x)
            test_scalar(asecd, 1/x)
            test_scalar(acscd, 1/x)
            test_scalar(acotd, 1/x)
        end

        @testset "sinc" for x = (0.0, 0.434, Complex(0.434, 0.25))
            test_scalar(sinc, x)
        end
    end  # Trig

    @testset "Angles" begin
        for x in (-0.1, 6.4, 0.5 + 0.25im)
            test_scalar(deg2rad, x)
            test_scalar(rad2deg, x)
        end
    end

    @testset "Unary complex functions" begin
        for x in (-4.1, 6.4, 0.0, 0.0 + 0.0im, 0.5 + 0.25im)
            test_scalar(real, x)
            test_scalar(imag, x)
            test_scalar(hypot, x)
            test_scalar(adjoint, x)
        end
    end

    @testset "Complex" begin
        test_scalar(Complex, randn())
        test_scalar(Complex, randn(ComplexF64))

        test_frule(Complex, randn(), randn())
        test_rrule(Complex, randn(), randn())
    end

    @testset "*(x, y) (scalar)" begin
        # This is pretty important so testing it fairly heavily
        test_points = (0.0, -2.1, 3.2, 3.7+2.12im, 14.2-7.1im)
        @testset "($x) * ($y)" for
            x in test_points, y in test_points

            # ensure all complex if any complex for FiniteDifferences
            x, y = Base.promote(x, y)

            test_frule(*, x, y)
            test_rrule(*, x, y)
        end
    end

    @testset "ldexp" begin
        for n in (0,1,20)
            test_frule(ldexp, 10rand(), n)
            test_rrule(ldexp, 10rand(), n)
        end
    end

    @testset "\\(x::$T, y::$T) (scalar)" for T in (Float64, ComplexF64)
        test_frule(*, randn(T), randn(T))
        test_rrule(*, randn(T), randn(T))
    end

    @testset "mod" begin
        test_frule(mod, 10rand(), rand())
        test_rrule(mod, 10rand(), rand())
    end

    @testset "identity" for T in (Float64, ComplexF64)
        test_frule(identity, randn(T))
        test_frule(identity, randn(T, 4))
        test_frule(
            identity,
            #Tuple(randn(T, 3)) ⊢ Composite{Tuple{T, T, T}}(randn(T, 3)...)
            Tuple(randn(T, 3))
        )

        test_rrule(identity, randn(T))
        test_rrule(identity, randn(T, 4))
        test_rrule(
            identity,
            Tuple(randn(T, 3)) ⊢ Composite{Tuple{T, T, T}}(randn(T, 3)...);
            output_tangent = Composite{Tuple{T, T, T}}(randn(T, 3)...)
        )
    end

    @testset "Constants" for x in (-0.1, 6.4, 1.0+0.5im, -10.0+0im, 0.0+200im)
        test_scalar(one, x)
        test_scalar(zero, x)
    end

    @testset "muladd(x::$T, y::$T, z::$T)" for T in (Float64, ComplexF64)
        test_frule(muladd, 10randn(), randn(), randn())
        test_rrule(muladd, 10randn(), randn(), randn())
    end

    @testset "fma" begin
        test_frule(fma, 10randn(), randn(), randn())
        test_rrule(fma, 10randn(), randn(), randn())
    end

    @testset "clamp"  begin
        # to left
        test_frule(clamp, 1., 2., 3.)
        test_rrule(clamp, 1., 2., 3.)

        # in the middle
        test_frule(clamp, 2.5, 2., 3.)
        test_rrule(clamp, 2.5, 2., 3.)

        # to right
        test_frule(clamp, 4., 2., 3.)
        test_rrule(clamp, 4., 2., 3.)
    end

    @testset "rounding" begin
        for x in (-0.6, -0.2, 0.1, 0.6)
            # thanks to RoundNearest
            if 0 > x % 1 > 0.5 || -1 < x % 1 <= 0.5
                test_scalar(round, x; fdm=backward_fdm(5,1))
            else
                test_scalar(round, x; fdm=forward_fdm(5,1))
            end
            test_scalar(floor, x; fdm=backward_fdm(5, 1))
            test_scalar(ceil, x; fdm=forward_fdm(5, 1))
        end
    end

    @testset "type" begin
        @test frule((NO_FIELDS, DoesNotExist(), DoesNotExist()), typejoin, Array{Float32,4}, Array{Float32,3}) !== nothing
        @test rrule(typejoin, Array{Float32,4}, Array{Float32,3}) !== nothing
    end

    @testset "Logging" begin
        @test frule((NO_FIELDS, DoesNotExist(), DoesNotExist()), Base.depwarn, "message", :f) !== nothing
        @test rrule(Base.depwarn, "message", :f) !== nothing
    end
end


function lyapunovs_convergence(integ, N, dt::Real, Ttr::Real = 0.0)

    T = stateeltype(integ)
    t0 = integ.t
    if Ttr > 0
        while integ.t < t0 + Ttr
            step!(integ, dt)
            qrdec = LinearAlgebra.qr(get_deviations(integ))
            set_deviations!(integ, qrdec.Q)
        end
    end
    k = size(get_deviations(integ))[2]
    t0 = integ.t; t = zeros(T, N); t[1] = t0
    λs = [zeros(T, k) for i in 1:N];

    for i in 2:N
        step!(integ, dt)
        qrdec = LinearAlgebra.qr(get_deviations(integ))
        for j in 1:k
            @inbounds λs[i][j] = λs[i-1][j] + log(abs(qrdec.R[j,j]))
        end
        t[i] = integ.t
        set_deviations!(integ, qrdec.Q)
    end
    popfirst!(λs); popfirst!(t)
    for j in eachindex(t)
        λs[j] ./= (t[j] - t0)
    end
    return λs, t
end


@testset "Edit distance" begin

    gtri = random_regular_graph(3, 2)
    gquad = random_regular_graph(4, 2)
    gpent = random_regular_graph(5, 2)

    @testset "edit_distance $triangle, $quadrangle, $pentagon" for triangle in testgraphs(gtri), quadrangle in testgraphs(gquad), pentagon in testgraphs(gpent)
        d, λ = @inferred(edit_distance(triangle, quadrangle, subst_cost=MinkowskiCost(1:3, 1:4)))
        @test d == 1.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (0, 4)]

        d, λ = @inferred(edit_distance(quadrangle, triangle, subst_cost=MinkowskiCost(1:4, 1:3)))
        @test d == 1.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (4, 0)]

        d, λ = @inferred(edit_distance(triangle, pentagon, subst_cost=MinkowskiCost(1:3, 1:5)))
        @test d == 2.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (0, 4), (0, 5)]

        d, λ = @inferred(edit_distance(pentagon, triangle, subst_cost=MinkowskiCost(1:5, 1:3)))
        @test d == 2.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (4, 0), (5, 0)]
    end

    @testset "Minkowski cost / bounded Minkowski" begin
        cost = @inferred(MinkowskiCost(1:3, 1:3))
        bcost = @inferred(BoundedMinkowskiCost(1:3, 1:3))
        for i = 1:3
          @test cost(i, i) == 0.
          @test bcost(i, i) == 2 / 3
        end
    end

    g1c = complete_graph(4)
    g2c = complete_graph(4)
    rem_edge!(g2c, 1, 2)
    @testset "edit_distance $g1, $g2" for g1 in testgraphs(g1c), g2 in testgraphs(g2c)
        d, λ = @inferred(edit_distance(g1, g2))
        @test d == 2.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (4, 4)]
    end
end


using MinkReduction
using Test

@testset "MinkReduction.jl" begin
    @test m = DeviousMat(10); a = m[1,:]; b = m[2,:]; c = m[3,:]; d,e,f = minkReduce(a,b,c); hcat(d,e,f) == [0.0 -1.0 0.0; 0.0 0.0 -1.0; 1.0 0.0 0.0]

end


using Tmatrix
using Zygote

rx, rz = 1e-6, 1e-6
n_max = 1
k1_r = 1e7;
k1_i = 0.0;
k2_r = 1.5e7;
k2_i = 1e3;

k1 = Complex(k1_r, k1_i)
k2 = Complex(k2_r, k2_i)

n_θ_points = 10
θ_array = collect(LinRange(1e-6, pi, n_θ_points))
r_array = rx * ones(size(θ_array))
ϕ_array = zeros(size(θ_array))
rotationally_symmetric = true
symmetric_about_plane_perpendicular_z = false
BigFloat_precision = nothing

# testing out T-matrix calculation and automatic differentiation using <<T_matrix_SeparateRealImag_arbitrary_mesh>>
@time T_sph = calculate_Tmatrix_for_spheroid_SeparateRealImag(
    rx,
    rz,
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i;
    n_θ_points = n_θ_points,
    rotationally_symmetric = true,
    symmetric_about_plane_perpendicular_z = false,
);
println("T-matrix using spheroid <<calculate_Tmatrix_for_spheroid_SeparateRealImag>>")
display(T_sph)

@time T = Tmatrix.T_matrix_SeparateRealImag_arbitrary_mesh(
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    rotationally_symmetric,
    symmetric_about_plane_perpendicular_z,
    BigFloat_precision,
);
println("T-matrix using an arbitrary mesh <<T_matrix_SeparateRealImag_arbitrary_mesh>>")
display(T)

@time ∂T = Zygote.jacobian(
    Tmatrix.T_matrix_SeparateRealImag_arbitrary_mesh,
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    rotationally_symmetric,
    symmetric_about_plane_perpendicular_z,
    BigFloat_precision,
)
println("Jacobian for <<T_matrix_SeparateRealImag_arbitrary_mesh>> calcuated by Zygote:")
display(∂T)

# testing gradient descent optimization for <<sin>>
function target_to_be_minimized(t)
    return sin(t)
end

function ∂target_to_be_minimized(t)
    return Zygote.gradient(target_to_be_minimized, t)[1]
end

"""
learning_rate = 0.01
loss_here = 1e6
t_0 = 0.3
∂loss = 0
while loss_here > 0.01
    t_0 -= learning_rate * ∂loss
    loss_here = target_to_be_minimized(t_0)
    ∂loss = ∂target_to_be_minimized(t_0)
    println("loss_here = loss_here, ∂loss = ∂loss")
end
"""

# testing gradient descent optimization for <<T_matrix_SeparateRealImag_arbitrary_mesh>>
# TODO: use an optimization package instead, so that it has better stop criteria and tricks for faster convergence
rx, rz = 1e-6, 1e-6
n_max = 2
k1_r = 1e7;
k1_i = 0.0;
k2_r = 1.5e7;
k2_i = 1e3;

n_θ_points = 20
θ_array = collect(LinRange(1e-6, pi, n_θ_points))
r_array = rx * ones(size(θ_array))
ϕ_array = zeros(size(θ_array))
rotationally_symmetric = true
symmetric_about_plane_perpendicular_z = false
BigFloat_precision = nothing

function objective_function_normal_mesh_density(r_array)
    # println("I am <<objective_function>>, and the type of r_array is $(typeof(r_array))")
    # println("r_array = $r_array")
    ϕ_array = zeros(size(θ_array))
    T = Tmatrix.T_matrix_SeparateRealImag_arbitrary_mesh(
        n_max,
        k1_r,
        k1_i,
        k2_r,
        k2_i,
        (collect(r_array)),
        θ_array,
        ϕ_array,
        rotationally_symmetric,
        symmetric_about_plane_perpendicular_z,
        BigFloat_precision,
    )
    # return abs(T[1]) # if we need to minimize the first element of T-matrix

    # return sum(abs.(T)) # if we need to minimize the sum of all real and imag

    # return Tmatrix.get_OrentationAv_scattering_CrossSection_from_Tmatrix(T, Complex(k1_r, k1_i)) # if we need to minimize scattering cross section

    return 1 * Tmatrix.get_OrentationAv_emissivity_from_Tmatrix(
        T,
        Complex(k1_r, k1_i),
        Tmatrix.calculate_surface_area_of_axisymmetric_particle(r_array, θ_array),
    )
end

function objective_function_double_mesh_density(r_array_input)
    r_θ_array = double_mesh_density(r_array_input, θ_array)
    r_array_double_density = r_θ_array[:, 1]
    θ_array_double_density = r_θ_array[:, 2]
    ϕ_array_double_density = zeros(size(θ_array_double_density))

    # println("I am <<objective_function>>, and the type of r_array_double_density is $(typeof(r_array_double_density))")
    println("r_array_input = $r_array_input")
    println("r_array_double_density = $r_array_double_density")

    T = Tmatrix.T_matrix_SeparateRealImag_arbitrary_mesh(
        n_max,
        k1_r,
        k1_i,
        k2_r,
        k2_i,
        (collect(r_array_double_density)),
        θ_array_double_density,
        ϕ_array_double_density,
        rotationally_symmetric,
        symmetric_about_plane_perpendicular_z,
        BigFloat_precision,
    )
    # return abs(T[1]) # if we need to minimize the first element of T-matrix

    # return sum(abs.(T)) # if we need to minimize the sum of all real and imag

    # return Tmatrix.get_OrentationAv_scattering_CrossSection_from_Tmatrix(T, Complex(k1_r, k1_i)) # if we need to minimize scattering cross section

    return 1 * Tmatrix.get_OrentationAv_emissivity_from_Tmatrix(
        T,
        Complex(k1_r, k1_i),
        Tmatrix.calculate_surface_area_of_axisymmetric_particle(
            r_array_double_density,
            θ_array_double_density,
        ),
    )
end

# objective_function = objective_function_normal_mesh_density
objective_function = objective_function_double_mesh_density

function ∂objective_function(r_array)
    # println("I am <<∂objective_function>>, and the type of r_array is $(typeof(r_array))")
    # println("r_array = $r_array")
    return Zygote.gradient(objective_function, (collect(r_array)))[1]
end

using Plots
n_max = 2
n_θ_points = 20
θ_array = collect(LinRange(1e-6, pi, n_θ_points))
r_array_initial, _ = Tmatrix.ellipsoid(0.1e-7, 0.12e-7, θ_array)
# r_array_initial = 0.1e-7 * ones(size(θ_array))
r_array = r_array_initial
ϕ_array = zeros(size(θ_array))
xyz_initial = vcat(
    Tmatrix.convert_coordinates_Sph2Cart.(
        r_array_initial,
        θ_array,
        zeros(size(r_array_initial)),
    )...,
)
rotationally_symmetric = true
symmetric_about_plane_perpendicular_z = false
BigFloat_precision = nothing

learning_rate = 0.5e-18
# learning_rate = 1e-2
global loss_here = -1e6
global ∂loss = 0
println("Starting optimization for T-matrix")
global n_iteration = 0
loss_array = []
loss_old = 9999999999999
while (n_iteration < 50) # (loss_here < 0)
    n_iteration += 1
    r_array = r_array .- learning_rate .* ∂loss

    loss_here = objective_function(r_array)
    ∂loss = ∂objective_function(r_array)
    append!(loss_array, loss_here)

    println()
    println("iteration #$n_iteration: loss_here = $loss_here, ∂loss = $∂loss")
    println("r_array = $r_array")

    xyz = vcat(
        Tmatrix.convert_coordinates_Sph2Cart.(r_array, θ_array, zeros(size(r_array)))...,
    )
    p1 = plot(
        xyz_initial[:, 1],
        xyz_initial[:, 3],
        aspect_ratio = :equal,
        label = "initial particle",
    )
    p1 = plot!(
        xyz[:, 1],
        xyz[:, 3],
        aspect_ratio = :equal,
        label = "particle after $n_iteration iterations",
        title = "scattering cross section = $loss_here m",
    )
    p2 = plot(
        1:length(loss_array),
        loss_array,
        xlabel = "iteration #",
        ylabel = "scattering cross section (m)",
    )
    fig = plot(p1, p2, layout = (1, 2), size = (1200, 800))
    savefig(
        fig,
        "cache/iteration_particle_plots/maximizing_emissivity_07_12_2021__/particle_geom_iteration_$n_iteration.png",
    )
    # if (loss_here - loss_old) < 1e-6; break; end
    loss_old = loss_here
end


# model.jl

"""
    abstract type AbstractModel end

Implements a statistical model. Has so far only SDEModel as subtype.
"""
abstract type AbstractModel end

"""
    type SDEModel{T<:Union{Float64, <:Function}} <: AbstractModel
    σ::T
    beginvalue::Float64
    end

Implements an SDE model with begin value beginvalue and σ.

# Warning
It is assumed that
for every b under consideration, the laws of dX_t=b(X_t)dt+σ(X_t)dW_t are
equivalent.

# Example

```julia
SDEModel(1.0, 0.0)
# Implements the model dX_t=b(X_t)dt+σ(X_t)dW_t, with X_0=0.0.
#
SDEModel(identity, 0.0)
# Implements the model dX_t=b(X_t)+X_tdW_t, with X_0=0.0.
```
"""
type SDEModel{T<:Union{Float64, <:Function}} <: AbstractModel
    σ::T
    beginvalue::Float64
end


export iauRx
"""
Rotate an r-matrix about the x-axis.

This function is part of the International Astronomical Union's
SOFA (Standards Of Fundamental Astronomy) software collection.

Status:  vector/matrix support function.

Given:
   phi    double          angle (radians)

Given and returned:
   r      double[3][3]    r-matrix, rotated

Notes:

   1. Calling this function with positive phi incorporates in the
   supplied r-matrix r an additional rotation, about the x-axis,
   anticlockwise as seen looking towards the origin from positive x.

   2. The additional rotation can be represented by this matrix:

      (  1        0            0      )
      (                               )
      (  0   + cos(phi)   + sin(phi)  )
      (                               )
      (  0   - sin(phi)   + cos(phi)  )

This revision:  2013 June 18

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauRx(double phi, double r[3][3])
function iauRx(phi::Real, r::AbstractMatrix{<:Real})

   # Transpose matrix upfront
   r = Array{Float64, 2}(r') # Transpose input up front

   ccall((:iauRx, libsofa_c), Cvoid,
           (Cdouble, Ptr{Cdouble}),
           phi, pointer(r))

   return convert(Array{Float64, 2}, r')
end


@symbol_func function K_nT(cur_reactor::AbstractReactor)
  cur_K = 0.1581

  cur_K *= ( 1 + cur_reactor.f_D )

  cur_K *= ( 1 + cur_reactor.nu_n )

  cur_K *= ( 1 + cur_reactor.nu_T )

  cur_K /= ( 1 + cur_reactor.nu_n + cur_reactor.nu_T )

  cur_K
end


function pack(X, U)
    z = zeros(num_var - p)
    for t = 1:T
        z[x_idx[t]] = X[t]
        t == T && continue
        z[u_idx[t]] = U[t]
    end
    return z
end

function unpack(z)
    x = [z[x_idx[t]] for t = 1:T]
    u = [z[u_idx[t]] for t = 1:T-1]
    return x, u
end

function primal_bounds(x1, xT)
    zl = -Inf * ones(num_var)
    zu = Inf * ones(num_var)

	zl[x_idx[1]] = x1
	zu[x_idx[1]] = x1

	# zl[x_idx[T]] = xT
	# zu[x_idx[T]] = xT

    return zl, zu
end

function constraint_bounds()
    cl = zeros(num_con)
    cu = zeros(num_con)
    return cl, cu
end

function eval_objective(prob::MOIGeneralProblem, z)
	J = 0.0
	for t = 1:T
		if t < T
			x = z[x_idx[t]]
			u = z[u_idx[t]]
			J += objective(x, u, t)
		else
			x = z[x_idx[t]]
			J += objective(x, nothing, t)
		end
	end
	# parameter regularization
	θ = z[p_idx]
	J += 100.0 * (θ - θ0_policy)' * (θ - θ0_policy)
	return J
end

function eval_objective_gradient!(∇J, z, prob::MOIGeneralProblem)
    ∇J .= 0.0
	for t = 1:T
		if t < T
			x = z[x_idx[t]]
			u = z[u_idx[t]]
			obj_x(a) = objective(a, u, t)
			obj_u(a) = objective(x, a, t)
			∇J[x_idx[t]] = ForwardDiff.gradient(obj_x, x)
			∇J[u_idx[t]] = ForwardDiff.gradient(obj_u, u)
		else
			x = z[x_idx[t]]
			obj_xT(a) = objective(a, nothing, t)
			∇J[x_idx[t]] = ForwardDiff.gradient(obj_xT, x)
		end
	end
	θ = z[p_idx]
	∇J[p_idx] = 2.0 * 100.0 * (θ - θ0_policy)

    return nothing
end

function eval_constraint!(c, z, prob::MOIGeneralProblem)
	# dynamics
    for t = 1:T-1
		x = z[x_idx[t]]
		u = z[u_idx[t]]
		y = z[x_idx[t+1]]
		c[n_idx[t]] = dynamics(y, x, u, t)
	end

	# policy
	θ = z[p_idx]
	for t = 1:T-1
		x = z[x_idx[t]]
		u = z[u_idx[t]]
		c[num_dyn .+ m_idx[t]] = u - policy(x, θ)[1]
	end
    return nothing
end

function eval_constraint_jacobian!(∇c, z, prob::MOIGeneralProblem)
	shift = 0

    for t = 1:T-1
        x = view(z, x_idx[t])
        u = view(z, u_idx[t])
        y = view(z, x_idx[t + 1])

        dyn_x(a) = dynamics(y, a, u, t)
        dyn_u(a) = dynamics(y, x, a, t)
        dyn_y(a) = dynamics(a, x, u, t)

        r_idx = n_idx[t]

        s = n * n
        ∇c[shift .+ (1:s)] = vec(ForwardDiff.jacobian(dyn_x, x))
        shift += s

        s = n * m
        ∇c[shift .+ (1:s)] = vec(ForwardDiff.jacobian(dyn_u, u))
        shift += s

        s = n * n
        ∇c[shift .+ (1:s)] = vec(ForwardDiff.jacobian(dyn_y, y))
        shift += s
    end

	# policy
	θ = z[p_idx]
	for t = 1:T-1
		x = z[x_idx[t]]
		u = z[u_idx[t]]

		px, pθ = policy_gradients(x, θ)

		r_idx = num_dyn .+ m_idx[t]

		s = m * m
		∇c[shift .+ (1:s)] = vec(Diagonal(ones(m)))
		shift += s

		s = m * n
		∇c[shift .+ (1:s)] = -vec(px)
		shift += s

		s = m * p
		∇c[shift .+ (1:s)] = -vec(pθ)
		shift += s
	end

    return nothing
end

function sparsity_jacobian(prob::MOIGeneralProblem)

	row = []
	col = []

	for t = 1:T-1
		r_idx = n_idx[t]
		row_col!(row, col, r_idx, x_idx[t])
		row_col!(row, col, r_idx, u_idx[t])
		row_col!(row, col, r_idx, x_idx[t + 1])
	end

	for t = 1:T-1
		r_idx = num_dyn .+ (m_idx[t])
		row_col!(row, col, r_idx, u_idx[t])
		row_col!(row, col, r_idx, x_idx[t])
		row_col!(row, col, r_idx, p_idx)
	end

	return collect(zip(row, col))
end

## Trajectory Optimization

# time
T = 11
h = 1.0

# dynamics
n = 2
m = 1

A = SMatrix{n,n}([1.0 1.0; 0.0 1.0])
B = SMatrix{n,m}([0.0; 1.0])

function dynamics(y, x, u, t)
	y - (A * x + B * u)
end

# objective
function objective(x, u, t)
	J = 0.0
	if t == T
		Q = Diagonal(100.0 * ones(n))
		J += x' * Q  * x
	else
		Q = Diagonal(ones(n))
		R = Diagonal(0.1 * ones(m))
		J += x' * Q  * x
		J += transpose(u) * R * u
	end
	return J
end

## policy

T_mpc = 5
n = 2
m = 1

x1 = [1.0; 0.0]
xT = [0.0; 0.0]

x_init = x1
_Q = sqrt.(ones(n))
_R = sqrt.(1.0 * ones(m))
_QT = sqrt.(1.0 * ones(n))
x_final = xT

nθ_mpc = n + n + m + n + n + n * n + n * m

num_x_mpc = n * T_mpc
num_u_mpc = m * (T_mpc - 1)
num_y_mpc = n * T_mpc
nz_mpc = num_x_mpc + num_u_mpc + num_y_mpc

x_idx_mpc = [collect((t - 1) * n .+ (1:n)) for t = 1:T_mpc]
u_idx_mpc = [collect(num_x_mpc + (t - 1) * m .+ (1:m)) for t = 1:T_mpc-1]
y_idx_mpc = [collect(num_x_mpc + num_u_mpc + (t - 1) * n .+ (1:n)) for t = 1:T_mpc]

x_mpc = linear_interpolation(x_init, x_final, T_mpc)
u_mpc = [0.01 * randn(m) for t = 1:T_mpc-1]
y_mpc = [zeros(n) for t = 1:T_mpc]

z0 = [vcat(x_mpc...); vcat(u_mpc...); vcat(y_mpc...)]
θ0 = [x_init; _Q; _R; _QT; x_final; vec(A); vec(B)]

function lagrangian(z, θ)
	x_init = θ[1:n]
	_Q = θ[n .+ (1:n)]
	Qt = Diagonal(_Q.^2.0)
	_R = θ[n + n .+ (1:m)]
	Rt = Diagonal(_R.^2.0)
	_QT = θ[n + n + m .+ (1:n)]
	QT = Diagonal(_QT.^2.0)
	x_final = θ[n + n + m + n .+ (1:n)]
	Aθ = reshape(θ[n + n + m + n + n .+ (1:n * n)], n, n)
	Bθ = reshape(θ[n + n + m + n + n + n * n .+ (1:n * m)], n, m)

	L = 0.0

	x1 = view(z, x_idx_mpc[1])
	y1 = view(z, y_idx_mpc[1])

	L += transpose(y1) * (x_init - x1)

	for t = 1:T_mpc-1
		xt = view(z, x_idx_mpc[t])
		ut = view(z, u_idx_mpc[t])
		yt⁺ = view(z, y_idx_mpc[t+1])
		xt⁺ = view(z, x_idx_mpc[t+1])

		L += transpose(xt - x_final) * Qt * (xt - x_final)
		L += (transpose(ut) * Rt * ut)

		L += transpose(yt⁺) *  (xt⁺ - (Aθ * xt + Bθ * ut))
	end

	xT = view(z, x_idx_mpc[T_mpc])
	# yT⁺ = view(z, y_idx_mpc[T+1])

	L += transpose(xT - x_final) * QT * (xT - x_final)
	# L += transpose(yT⁺) * (x_final - xT)

	return L
end

lagrangian(z0, θ0)

@variables z_sym[1:nz_mpc]
@variables θ_sym[1:nθ_mpc]
@variables κ_sym[1:1]

L = lagrangian(z_sym, θ_sym);
L = simplify.(L);

dL = Symbolics.gradient(L, z_sym)
ddL = Symbolics.sparsehessian(L, z_sym)

L_grad = eval(Symbolics.build_function(dL, z_sym, θ_sym, κ_sym)[1])
L_hess = eval(Symbolics.build_function(ddL, z_sym, θ_sym)[1])

L_grad! = eval(Symbolics.build_function(dL, z_sym, θ_sym, κ_sym)[2])
L_hess! = eval(Symbolics.build_function(ddL, z_sym, θ_sym)[2])

L_grad(z0, θ0, 1.0)
L_hess(z0, θ0)

dLθ = Symbolics.jacobian(dL, θ_sym)

dL_jacθ = eval(Symbolics.build_function(dLθ, z_sym, θ_sym)[1])
dL_jacθ! = eval(Symbolics.build_function(dLθ, z_sym, θ_sym)[2])

dL_jacθ(z0, θ0)

# options
opts = InteriorPointOptions(
	κ_init = 1.0,
	κ_tol = 1.0,
	r_tol = 1.0e-8,
	diff_sol = true)

# solver
ip = interior_point(z0, θ0,
	r! = L_grad!, rz! = L_hess!,
	rz = similar(ddL, Float64),
	rθ! = dL_jacθ!,
	rθ = similar(dLθ, Float64),
	opts = opts)

# solve
status = interior_point!(ip)

using Test, Plots
# test
@test status
@test norm(ip.z[x_idx_mpc[1]] - x_init) < 1.0e-8

norm(ip.z[x_idx_mpc[T_mpc]] - x_final)

x_traj_mpc = [ip.z[x_idx_mpc[t]] for t = 1:T_mpc]
u_traj_mpc = [ip.z[u_idx_mpc[t]] for t = 1:T_mpc-1]

plot(hcat(x_traj_mpc...)')
plot(hcat(u_traj_mpc...)')

function policy(x, θ)
	x_mpc = linear_interpolation(x, x_final, T_mpc)
	u_mpc = [0.01 * randn(m) for t = 1:T_mpc-1]
	y_mpc = [zeros(n) for t = 1:T_mpc]

	# initialized solver primals
	ip.z .= copy([vcat(x_mpc...); vcat(u_mpc...); vcat(y_mpc...)])

	# initialized solver parameters
	_Q = θ[1:n]
	_R = θ[n .+ (1:m)]
	_QT = θ[n + m .+ (1:n)]
	_A = θ[n + m + n .+ (1:n * n)]
	_B = θ[n + m + n + n * n .+ (1:n * m)]

	ip.θ .= copy([x; _Q; _R; _QT; x_final; _A; _B])

	status = interior_point!(ip)

	!status && (@warn "solver failure")

	return ip.z[u_idx_mpc[1]], ip
	# return reshape(θ, m, n) * x, nothing
end

function policy_gradients(x, θ)
	_, ip = policy(x, θ)

	ux = ip.δz[u_idx_mpc[1], 1:n]
	uθ = ip.δz[u_idx_mpc[1], n .+ (1:(n + m + n + n * n + n * m))]

	return ux, uθ

	# px(a) = policy(a, θ)[1]
	# pθ(a) = policy(x, a)[1]
	# return ForwardDiff.jacobian(px, x), ForwardDiff.jacobian(pθ, θ)
end

# policy(x1, 0.01 * randn(p))
# policy_gradients(x1, 0.01 * randn(p))
θ0_policy = [_Q; _R; _QT; vec(A); vec(B)]
policy(x1, θ0_policy)[1]
policy_gradients(x1, θ0_policy)

# p = m * n
p = n + m + n + n * n + n * m

# trajectory indices
x_idx = [(t - 1) * (n + m) .+ (1:n) for t = 1:T]
u_idx = [(t - 1) * (n + m) + n .+ (1:m) for t = 1:T-1]
n_idx = [(t - 1) * n .+ (1:n) for t = 1:T-1]
m_idx = [(t - 1) * m .+ (1:m) for t = 1:T-1]
p_idx = n * T + m * (T - 1) .+ (1:p)

# MOI setup
num_x = n * T
num_u = m * (T - 1)
num_var = num_x + num_u + p

num_dyn = n * (T - 1)
num_policy = m * (T - 1)
num_con = num_dyn + num_policy

# initialization
x1 = [1.0; 0.0]
xT = [0.0; 0.0]

# U = [0.1 * randn(m) for t = 1:T-1]

# rollout
x_hist = [copy(x1)]
u_hist = []
Js = 0.0
for t = 1:T-1
	push!(u_hist, policy(x_hist[end], θ0_policy)[1])
	# push!(u_hist, policy(x_hist[end], θ_sol)[1])

	Js += objective(x_hist[end], u_hist[end], t)
	push!(x_hist, A * x_hist[end] + B * u_hist[end])
end
Js += objective(x_hist[end], nothing, T)
@show Js
@show norm(x_hist[end] - xT)

z = [pack(x_hist, u_hist); θ0_policy]
# z = [pack(X, U); 0.01 * randn(p)]# _Q; _R; _QT]

prob = MOIGeneralProblem(num_var, num_con, primal_bounds(x1, xT), constraint_bounds())
prob_moi = moi_problem(prob)

z_sol, info = solve(prob_moi, z,
		tol = 1.0e-2,
		c_tol = 1.0e-2,
		max_iter = 1000,
		nlp = :ipopt)

x_sol, u_sol = unpack(z_sol)
θ_sol = z_sol[p_idx]

norm(θ_sol - θ0_policy)
# M = reshape(θ_sol, m, n)

using Plots

plot(hcat(x_sol...)')
plot!(hcat(x_hist...)')
plot(hcat(u_sol..., u_sol[end])', linetype = :steppost)
plot!(hcat(u_hist..., u_hist[end])', linetype = :steppost)

## simulate policy
x1_sim = copy(x1) + 0.0 * randn(n)

x_hist = [x1_sim]
u_hist = []
Js = 0.0
for t = 1:T-1
	# push!(u_hist, policy(x_hist[end], θ0_policy)[1])
	push!(u_hist, policy(x_hist[end], θ_sol)[1])

	Js += objective(x_hist[end], u_hist[end], t)
	push!(x_hist, A * x_hist[end] + B * u_hist[end])
end
Js += objective(x_hist[end], nothing, T)
@show Js
plot(hcat(x_hist...)')
plot(hcat(u_hist..., u_hist[end])', linetype = :steppost)

@show norm(x_hist[end] - xT)
Q_sol = Diagonal(θ_sol[1:n].^2.0)
R_sol = Diagonal(θ_sol[n .+ (1:m)].^2.0)
QT_sol = Diagonal(θ_sol[n + m .+ (1:n)].^2.0)
reshape(θ_sol[n + m + n .+ (1:n * n)], n, n)
reshape(θ_sol[n + m + n + n * n .+ (1:n * m)], n, m)

eigen(reshape(θ_sol[n + m + n .+ (1:n * n)], n, n))


"""

    lossy_fit(xs,ys;Rdc=1e0,Rs=1e-6,L0=1e-9,epsilon=1e-3,C0=1e-12,show_trace=false,maxIter=1000)

lossy_fit accepts two positional arguments

xs::Float{64,1} - Array of sampled frequencies
ys::Float{64,1} - Array of frequency response values

and 7 keyword arguments

Rdc::Float64
Rs::Float64
L0::Float64
epsilon::Float64
C0::Float64
show_trace::Bool
maxIter::Int32

where the first 5 keyword arguments are the sizes (in SI units) of the guesses
of each of the fitted parameter values and the last two control the behavior of
LsqFit.jl.

lossy_fit computes a fit of the parameters Rdc (DC resistive loss), Rs (skin effect loss),
L0 (series inductance), loss tangent, C0 (shunt capacitance) provided the
frequency-domain response of a device under test assuming the model described in
[[RINEHART2017]].
"""
function lossy_fit(xs,ys;Rdc=1e0,Rs=1e-6,L0=1e-9,epsilon=1e-3,C0=1e-12,show_trace=false,maxIter=1000)
    z(p) = f -> tline_z_fd(p...)(f)
    # re-define the model to conform to LsqFit expectations
    model(f,p) = vcat(real(z(p).(f)), imag(z(p).(f)))

    # ForwardDiff.jacobian takes a function that only depends on the parameters to
    # be solved for.
    function model_jac(w,p)
        return ForwardDiff.jacobian(pj -> model(w,pj), p)
    end
    # guess
    p0 = [Rdc, Rs, L0, epsilon, C0]
    fit = LsqFit.curve_fit(model, model_jac, xs, ys, p0,
                           show_trace=show_trace,
                           lower=zeros(5)+eps(),
                           maxIter=maxIter,
                        )
    fit_errors = LsqFit.estimate_errors(fit, .95)
    return Dict("vals" => fit.param, "errors" => fit_errors)
end

"""
    impedance(ZR,ZL,ZG,ZC)

returns the impedance function to be fit.
"""
function impedance(ZR,ZL,ZG,ZC)
    frac(p) = f -> (ZR(p[1],p[2])(f)+ZL(p[3])(f)) / (ZG(p[4],p[5])(f) + ZC(p[5])(f))
    function Z(p)
        return function(f)
            f==0?sqrt(p[1]):abs(frac(p)(f))*cis(angle(frac(p)(f)))
        end
    end
    return Z
end


struct Order{N} end
Order(N) = Order{N}()

function get_process(::Order{1}, T)
    1
end

function get_process(::Order{2}, T)
    @SMatrix [
        1 T
        0 1
    ]
end

function get_process(::Order{3}, T)
    @SMatrix [
        1 T T^2 / 2
        0 1 T
        0 0 1
    ]
end

function get_process_covariance(::Order{1}, T)
    @SMatrix [T] # SMatrix more efficient than scalar for cholesky
end

function get_process_covariance(::Order{2}, T)
    @SMatrix [
        T^3/3 T^2/2
        T^2/2 T
    ]
end

function get_process_covariance(::Order{3}, T)
    @SMatrix [
        T^5/20 T^4/8 T^3/6
        T^4/8 T^3/3 T^2/2
        T^3/6 T^2/2 T
    ]
end

function soft_bound(value, noise, upper_bound, lower_bound, relative_soft_bounding = 0.2)
    value += value[1] > upper_bound * (1 - sign(upper_bound) * relative_soft_bounding) && value[end] > 0 && noise[end] > 0 ?
        SVector(noise[1], -noise[end]) : noise
    value += value[1] < lower_bound * (1 + sign(lower_bound) * relative_soft_bounding) && value[end] < 0 && noise[end] < 0 ?
        SVector(noise[1], -noise[end]) : noise
    value
end


using PolyChaos, Test, StaticArrays

deg = 5
op = genLaguerreOrthoPoly(deg,1.23)

n = 10
X = [ exp(3 + rand()), exp.(3 .+ rand(n)) ]

α, β, αs, βs = op.α, op.β, SVector(op.α...), SVector(op.β...)
combinations = Iterators.product([α, αs], [β, βs])

@testset "evaluate univariate basis at point(s)" begin
    for x in X
        [ @test evaluate(collect(0:deg),x,op) == evaluate(collect(0:deg), x, a, b)  for (a, b) in combinations ]
        [ @test evaluate(SVector(0:deg...),x,op) == evaluate(SVector(0:deg...), x, a, b)  for (a, b) in combinations ]
        [ @test evaluate(0:deg,x,op) == evaluate(0:deg, x, a, b)  for (a, b) in combinations ]
        [ @test evaluate(0:2:deg,x,op) == evaluate(0:2:deg, x, a, b)  for (a, b) in combinations ]
        for d in 0:deg
            if typeof(x) <: Real
                [ @test evaluate(d,x,a,b) == evaluate(d,x,op) == evaluate(d,[x],a,b) for (a, b) in combinations ]
            else
                [ @test evaluate(d,x,a,b) == evaluate(d,x,op) for (a, b) in combinations ]
            end
        end

        res = hcat(map(i -> evaluate(i,x,op.α,op.β),collect(0:op.deg))...)
        if typeof(x) <: Real
            [ @test evaluate(x,op) == evaluate([x], op) == evaluate(collect(0:op.deg),[x],op) == evaluate(collect(0:op.deg),[x],a,b) == res for (a, b) in  combinations ]
            [ @test evaluate(x,op) == evaluate([x], op) == evaluate(0:op.deg,[x],op) == evaluate(0:op.deg,[x],a,b) == res for (a, b) in  combinations ]
        else
            [ @test evaluate(x, op) == evaluate(collect(0:op.deg),x,op) == evaluate(collect(0:op.deg),x,a,b) == res for (a, b) in combinations ]
            [ @test evaluate(x, op) == evaluate(0:op.deg,x,op) == evaluate(0:op.deg,x,a,b) == res for (a, b) in combinations ]
        end
    end
end

ops = [  genLaguerreOrthoPoly(deg,1.23),
        genHermiteOrthoPoly(deg,3.),
        LegendreOrthoPoly(deg),
        JacobiOrthoPoly(deg,4.3,10.)
     ]

mop = MultiOrthoPoly(ops,deg)

nop = length(ops)
X = [ exp.(3 .+ rand(nop)) , exp.(3 .+ rand(n,nop)) ]

α, β = coeffs(mop)
αs, βs = SVector(α...), SVector(β...)
combinations = Iterators.product([α, αs], [β, βs])

@testset "evaluate multivariate basis at point(s)" begin
    for x in X
        [ @test evaluate(mop.ind,x,a,b) == evaluate(mop.ind,x,mop) == evaluate(x,mop) for (a, b) in combinations ]
        if typeof(x) <: AbstractVector{<:Real}
            [ @test evaluate(mop.ind,x,a,b) == evaluate(mop.ind,x,mop) == evaluate(x,mop) == evaluate(mop.ind,reshape(x,1,length(x)),coeffs(mop)...) for (a, b) in combinations]
            [ @test evaluate(x,mop) == evaluate(mop.ind,reshape(x,1,length(x)),mop) == evaluate(mop.ind,reshape(x,1,length(x)),a,b) for (a, b) in combinations ]
        else
            [ @test evaluate(x,mop) == evaluate(mop.ind,x,mop) == evaluate(mop.ind,x,a,b) for (a, b) in combinations ]
        end
        if VERSION >= VersionNumber("1.1.1")
            for d in eachrow(mop.ind[1:3,:])
                [ @test evaluate(collect(d),x,a,b) == evaluate(collect(d),x,mop) for (a, b) in combinations ]
                [ @test evaluate(d,x,a,b) == evaluate(collect(d),x,mop) for (a, b) in combinations ]
            end
        end
    end
end


use_cuda = false

using Mocha

srand(12345678)

############################################################
# Prepare Random Data
############################################################
N = 5    # works with arbitrary minibatch size as long as
         # N == batch_size in MemoryDataLayer so it cycles through
         # and gets the same data during forward()
M = 10
P = 4

X = rand(M, N)
W = rand(M, P)
B = rand(P, 1)

Y = (W'*X .+ B)
Y = Y + 0.01*randn(size(Y))

############################################################
# Define network
############################################################
backend = CPUBackend()
init(backend)

data_layer = MemoryDataLayer(batch_size=N, data=Array[X,Y])

w1 = InnerProductLayer(neuron=Neurons.Sigmoid(), name="ip1",output_dim=20, tops=[:a], bottoms=[:data])
w2 = InnerProductLayer(neuron=Neurons.Identity(), name="ip2",output_dim=4, tops=[:b], bottoms=[:a])
loss_layer = SquareLossLayer(name="loss", bottoms=[:b, :label] )

net = Net("TEST", backend, [w1,w2, loss_layer, data_layer])

# epsilon:     milage may vary 1e-4 - 1e-8
# digit:       compare this many digits to check for 'identity'
# visualize:   prints out correct and failed positions of gradients
# return value: false if any gradients fails, true if all passes
test_gradients(net, epsilon=1e-8, digit=6, visual=true )

shutdown(backend)


using SparseArrays
using GalerkinSparseGrids

# The exact coefficients for f(x_1, ..., x_D) = 1 in the modal basis
function get_one_modal(D::Int, k::Int, n::Int)
    len = get_size(Val(D), k, n)
    return [i==1 ? 1.0 : 0.0 for i in 1:len]
end

# The exact coefficients of the function f(x_1, ..., x_D) = x_i
function get_xi_modal(D::Int, i::Int, k::Int, n::Int)
    len = get_size(Val(1), k, n)
    coeffs1_1D::Array{Float64, 1} = [i==1 ? 1 : 0 for i in 1:len] # Coeffs for 1
    coeffs2_1D::Array{Float64, 1} = [i==2 ? 1/(sqrt(3)) : 0 for i in 1:len]
    coeff_array = [d==i ? coeffs2_1D : coeffs1_1D for d in 1:D]
    return tensor_construct(D, k, n, coeff_array)
end

function get_xi2_modal(D::Int, i::Int, k::Int, n::Int)
    len = get_size(Val(1), k, n)
    # Coefficients for 1:
    coeffs1_1D::Array{Float64, 1} = [i==1 ? 1 : 0 for i in 1:len]
    # Coefficients for (v-1/2)^2:
    coeffs2_1D::Array{Float64, 1} = [i==3 ? 2/(3*sqrt(5)) : 0 for i in 1:len]
    coeffs2_1D += 1/3 * coeffs1_1D
    coeff_array = [d==i ? coeffs2_1D : coeffs1_1D for d in 1:D]
    return tensor_construct(D, k, n, coeff_array)
end

# The exact coefficients for f(x) = |x|^2 in the modal basis
function get_r2_modal(D::Int, k::Int, n::Int)
    coeff_array = [get_xi2_modal(D, i, k, n) for i in 1:D]

    return sum(coeff_array)
end

# The exact coefficients for f(x_1, ..., x_D) = x_i = x_i in the point basis
# ^* insofar as the transformation matrices are exact
function get_xi_point(D::Int, i::Int, k::Int, n::Int,
                        m2n::SparseMatrixCSC{Float64, Int},
                        n2p::SparseMatrixCSC{Float64, Int})
    return n2p * (m2n * get_xi_modal(D, i, k, n))
end

# The exact coefficients for f(x_1, ..., x_D) = |x|^2 in the point basis
# ^* insofar as the transformation matrices are exact
function get_r2_point(D::Int, k::Int, n::Int,
                        m2n::SparseMatrixCSC{Float64, Int},
                        n2p::SparseMatrixCSC{Float64, Int})
    v = n2p * (m2n * get_r2_modal(D, k, n))
    v[v .< 0] .= 0
    return v
end


using MCObservables

function test_VectorObservable(MCS, SObs, VObs)
  inform("Start testing $VObs")

  n = 2

  s = zeros(SObs, n)
  v = VObs()

  for i in 1:MCS
    x = zeros(n)
    for j in 1:n
      x[j] = randn()
      s[j] << x[j]
    end
    v << x
  end

  vjk = jackknife(v)

  for i in 1:n
    @test mean(s[i]) ==  mean(v)[i]
    @test stderror(s[i]) == stderror(v)[i]
    sjk = jackknife(s[i])
    @test_approx_eq_eps mean(sjk) mean(vjk)[i] 1e-10
    @test_approx_eq_eps stderror(sjk) stderror(vjk)[i] 1e-10
  end

  inform("Success testing $VObs")
end

result_simpleObservable = test_VectorObservable(8192, SimpleObservable, SimpleVectorObservable)
result_simpleObservable = test_VectorObservable(8192, BinningObservable, BinningVectorObservable)


using SpecialFunctions

# sort by increasing absolute value of the real part. Hack to get poles ordered
# for testing
csort(x) = sort(x, lt = (x,y) -> abs(real(x)) < abs(real(y)))

# OK, so these tolerances are not as good as those shown in the AAA paper
function test_aaa_spiral()
    zz = range(-0.5, complex(0.5, 0.15pi), length=100)
    yy = tan.(pi*zz/2)
    f = aaa(zz, yy)
    pol, res, zer = prz(f)
    spol = csort(pol)

    # somehow, running in the Pkg test environment, I'm getting a
    # spurious pole at -1.46 -46.51im but not when I run locally.  How
    # can that be?  As a hack, I filter it out here, but should investigate
    # further
    filter!( x -> abs(imag(x)) < 0.01, spol)

    p1 = isapprox(abs(spol[1]),  1.0, atol=1e-8)
    p2 = isapprox(abs(spol[2]),  1.0, atol=1e-8)
    p3 = isapprox(abs(spol[3]),  3.0, atol=5e-3)
    p4 = isapprox(abs(spol[4]),  3.0, atol=5e-3)
    return all((p1, p2, p3, p4))
end

# we test the values of the first 4 located poles to the accuracy noted in
# the AAA paper.  It is very likely that these values will make for brittle
# testing.
function test_aaa_gamma_poles()
    xx = range(-1.5, 1.5, length=100)
    yy = gamma.(xx)
    f = aaa(xx, yy)
    pol, res, zer = prz(f)
    spol = csort(pol)
    p1 = isapprox(spol[1],  0.0, atol=1e-15)
    p2 = isapprox(spol[2], -1.0, atol=1e-15)
    p3 = isapprox(spol[3], -2.0, atol=2e-7)
    p4 = isapprox(spol[4], -3.0, atol=3e-3)
    return p1 & p2 & p3 & p4
end

function test_aaa_exp(tol=1e-8)
    Z = range(-1, 1, length=1000)
    F = exp.(Z)
    r = aaa(Z, F)
    t1 = norm(F - r(Z), Inf) < tol
    t2 = isnan(r(NaN))                        # check that r(NaN) = NaN
    t3 = !isinf(r(Inf))                       # r(inf) = sum(w.*f)/sum(w)
    return all((t1, t2, t3))
end

function test_aaa_length(tol=1e-8)
    Z = range(-1, 1, length=1000)
    F = exp.(Z)
    r = aaa(Z, F)
    m1 = length(r.x)
    r = aaa(Z, F, mmax=m1-1)
    t1 = length(r.x) == m1 - 1

    r = aaa(Z, F, tol=1e-3)
    t2 = length(r.x) < m1
    return all((t1, t2))
end

function test_aaa_tan(tol=1e-8)
    Z = range(-1, 1, length=1000)
    F = z -> tan(pi*z)
    r = aaa(Z, F)
    pol, res, zer  = prz(r)
    t1 = norm(F.(Z) - r(Z), Inf) < 10*tol
    t2 = minimum(abs.(zer)) < tol
    t3 = minimum(abs.(pol .- 0.5)) < tol
    t4 = minimum(abs.(res)) > 1e-13     # Test for spurious poles.
    return all((t1, t2, t3, t4))
end

function test_aaa_gamma(tol=1e-8)
    r = aaa([-0.9:0.05:1.0;], gamma)
    return abs(r(1.5) - gamma(1.5)) < 1e-3
end

function test_aaa_ranlog(tol=1e-8)
    Random.seed!(1137)
    Z = randn(10000) + 3im*randn(10000)
    f = z -> log(5 - z) / (1 + z^2)
    r = aaa(Z, f.(Z))
    return abs(r(0) - f(0)) < tol
end

function test_aaa_infinp(tol=1e-8)
    Z = complex.(range(-1.0, 1.0, length=101), zeros(101))
    r = aaa(Z, gamma.(Z))
    return abs(r(0.63) - gamma(0.63)) < 1e-3
end

function test_aaa_naninp(tol=1e-8)
    X = range(0, 20, length=100)
    F = sin.(X) ./ X
    r = aaa(X, F)
    return abs(r(2) - sin(2)/2) < 1e-3
end

function test_aaa_residuals(tol=1e-8)
    X = range(-1.337, 2, length=537)
    r = aaa(X, exp.(X) ./ X)
    pol, res, zer = prz(r)
    ii = findfirst(abs.(pol) .< 1e-8)
    t1 = abs(res[ii] - 1.0) < 1e-10

    r = aaa(X, (1+1im) .* gamma.(X))
    pol, res, zer = prz(r)
    ii = findfirst(abs.(pol .- (-1)) .< 1e-8)
    t2 = abs(res[ii] + (1+1im)) < 1e-10

    return all((t1, t2))
end

# The following two tests do not pass

function test_aaa_case2(tol=1e-8)
    # Case |Z| = 2: needs special treatment.
    Z = [0, 1]
    F = [1, 2]
    r = aaa(Z, F)
    t1 = norm(F .- r(Z), Inf) < tol
    t2 = r(Inf) == -Inf

    return all((t1, t2))
end

function test_aaa_scale_invar(tol=1e-8)
    # Check for exact scale-invariance
    Z = range(0.3, 1.5, length=100)
    F = exp.(Z) ./ (1+1im)
    r1 = aaa(Z, F)
    r2 = aaa(Z, BigFloat(2)^311*F)
    r3 = aaa(Z, BigFloat(2)^-311*F)
    t1 = r1(0.2im) == BigFloat(2)^-311*r2(0.2im)
    t2 = r1(1.4) == BigFloat(2)^311*r3(1.4)

    return all((t1, t2))
end

# Now a couple of tests of 'normal' functions on regular grids.

function do_aaa_func(func=abs, inc=0.01)
    x = [-1.0:inc:1.0;]
    y = func.(x)
    f = aaa(x, y)
    xt = [(-1.0 + inc):inc:(1.0 - inc);]
    yt = f.(xt)
    yexact = func.(xt)
    err = maximum(abs.(yt - yexact))
end

function test_aaa_abs_x(tol=1e-10)
    err = do_aaa_func()
    if err > tol
        println("AAA test failed  error: ", err, "   tol: ", tol)
        return false
    end
    return true
end

function runge(x)
    return 1.0 /(1 + x^2)
end

function test_aaa_runge(tol=1e-10)
    err = do_aaa_func(runge)
    if err > tol
        println("AAA test failed  error: ", err, "   tol: ", tol)
        return false
    end
    return true
end

#=
cs = [
2.4685030026219157  4.328148466145844e-5
2.2877690812473563  4.208926791402662e-5
2.165905087675962  4.1378862441314513e-5
2.0717086214734195  4.090525788259946e-5
1.9941763744769134  4.057575180437743e-5
1.9278704490832148  4.0343910121344306e-5
1.8697766805850384  4.018267056175463e-5
1.817980188860383  4.007483985348351e-5
1.7711964027423774  4.0008750298368046e-5
1.7285267566201805  3.9975982351327924e-5
1.6892386141964113  3.997041506117895e-5
1.6528756041444737  3.998735197995217e-5
1.6190201684981171  4.002313429694444e-5
1.5873476689145403  4.007488693018958e-5
1.5576041134013925  4.0140274206484625e-5
1.5295736436900844  4.021739144781512e-5
1.5030595092484886  4.030468761641359e-5
1.477921932766034  4.0400775660381905e-5
1.454088259456376  4.050457505125146e-5
1.431321568193522  4.061523389017265e-5
1.4096202820729862  4.073182521865663e-5
1.388866749525395  4.0853635453992726e-5
1.3690044324914987  4.0980262953198336e-5
1.3499576165246097  4.111096457002233e-5
1.3316503622294995  4.1245355727745365e-5
1.3140767513577072  4.138302834451439e-5
1.2971408925810646  4.1523626817923604e-5
1.2808179835881397  4.166684722541318e-5
1.265100234243082  4.1812370690389976e-5
1.2498687696967827  4.195998980734356e-5
1.2351709265347175  4.210948376459288e-5
1.2209602643145323  4.226070045734215e-5
1.207188872283871  4.241309483664375e-5
1.1938622378301857  4.256708644232688e-5
1.1809227477508621  4.2721984153255756e-5
1.168369539371024  4.2878130564666566e-5
1.1561858817429516  4.30350041771981e-5
1.1443463892653818  4.319269655519794e-5
1.1328417764637249  4.3351234144270344e-5
1.1216548642512816  4.351024305591165e-5]
=#


using SuperTB, PyPlot, ProgressMeter, LinearAlgebra

# set up a model (amorphous system with AII-class hopping)
hopfun(ri,rj,per) = hopAgarwala(ri, rj, per,
                                1.0/6,  # R
                                -.5,    # M
                                0.25,   # t₂
                                0.5,    # λ
                                1.0/64) # a
tb = TBModel(hopfun)
tb.periodicity = trues(2) # periodic boundary conditions
tb.basisdim    = 2

# run over this parameter range
sitessqrt  = 24
sites      = sitessqrt^2
Nensembles = 5
getbottatindex = sites .+ (-(2*sitessqrt+1):(2*sitessqrt+1))

# preallocation
par  = zeros(sites*tb.basisdim, Nensembles)      # participation ratios
εeig = zeros(sites*tb.basisdim, Nensembles)      # energy eigenvalues
β    = zeros(size(getbottatindex,1), Nensembles) # bott indices

close("all")
# compute and visualize during...
@showprogress 1 "Calculating ensembles..." for i in 1:Nensembles
    tb.pos = rand(sites,2)
    # solve the system
    H = Hermitian(Matrix(setHopping!(tb)))
    F = eigen(H)

    # compute participation ratios & plot
    par[:,i]  = parratio(F.vectors)
	εeig[:,i] = F.values
	figure()
	plot(εeig[:,i], par[:,i], ".")

    # compute bott indices & plot
    β[:,i] = bott(F.vectors,  tb,getbottatindex)
	plot(εeig[getbottatindex,i], β[:,i]./10, ".") # scale to fit in same plot...
end

#PyPlot.svg(true)
#savefig("sdos_ensembleaveraged")


using Test, StatsModels
using DataFrames
using StatsBase

using Unfold
include("test_utilities.jl")

data, evts = loadtestdata("test_case_3a") #
data_r = reshape(data, (1, :))

data_e, times = Unfold.epoch(data = data_r, tbl = evts, τ = (0.0, 1), sfreq = 20) # cut the data into epochs
basisfunction = firbasis(τ = (0.0, 1), sfreq = 20, name = "basisA")

# f  = @formula 0~ 1 * beta_0 + contA*beta_1
# fit => beta_0 = -1.5, beta_1 = 2
# predict => (contA=0.5) y_hat = 1 * -1.5 + 0.5 * 2 = -0.5

f = @formula 0 ~ 1 # 1
m_mul = fit(UnfoldModel, f, evts, data_e, times)
m_tul = fit(UnfoldModel, f, evts, data_r, basisfunction)
m_mul_results = coeftable(m_mul)
m_tul_results = coeftable(m_tul)

conditionA = [0, 1.0]
continuousA = [-1.0, 0, 1.0]

tmp = reshape(
    [[x, y] for x in conditionA, y in continuousA],
    length(conditionA) * length(continuousA),
)
evts_grid = DataFrame(collect(hcat(tmp...)'), ["conditionA", "continuousA"])

yhat_tul = predict(m_tul, evts_grid)
yhat_mul = predict(m_mul, evts_grid, times)

@test unique(yhat_mul.times) == times

@test size(m_tul_results)[1] * size(evts_grid)[1] == size(yhat_tul)[1]
@test size(m_mul_results)[1] * size(evts_grid)[1] == size(yhat_mul)[1]

@test yhat_mul[end-3, :yhat] ≈ mean(data[data.!=0])
@test yhat_tul[end-3, :yhat] ≈ mean(data[data.!=0])

@test yhat_mul.times[1] == 0.0

f = @formula 0 ~ 1 + conditionA + continuousA# 1
m_tul = fit(UnfoldModel, f, evts, data_r, basisfunction)
m_mul = fit(UnfoldModel, f, evts, data_e, times)

m_mul_results = coeftable(m_mul)
m_tul_results = coeftable(m_tul)

yhat_tul = predict(m_tul, evts_grid)
yhat_mul = predict(m_mul, evts_grid, times)
#@test predict(m_mul,evts).yhat ≈ yhat.yhat

@test isapprox(sum(coef(m_tul)[[5, 25, 45]]), yhat_tul.yhat[end-3], atol = 0.001)

@test isapprox(
    yhat_tul[yhat_tul.times.==0.5, :yhat],
    yhat_mul[yhat_mul.times.==0.5, :yhat],
    atol = 0.001,
)
@test isapprox(yhat_tul[yhat_tul.times.==0.5, :yhat], [-2, 1, 2, 5, 6, 9.0], atol = 0.0001)

if 1 == 0
    #for visualization
    using AlgebraOfGraphics
    yhat_mul.conditionA = categorical(yhat_mul.conditionA)
    m = mapping(:times, :yhat, color = :continuousA, linestyle = :conditionA)
    df = yhat_mul
    AlgebraOfGraphics.data(df) * visual(Lines) * m |> draw
end


export expected_policy_values, nash_conv

function expected_policy_values(π::AbstractPolicy, env::AbstractEnv)
    if get_terminal(env)
        [get_reward(env, p) for p in get_players(env) if p != get_chance_player(env)]
    elseif get_current_player(env) == get_chance_player(env)
        vals = [0. for p in get_players(env) if p != get_chance_player(env)]
        for a::ActionProbPair in get_legal_actions(env)
            vals .+= a.prob .* expected_policy_values(π, child(env, a))
        end
        vals
    else
        vals = [0. for p in get_players(env) if p != get_chance_player(env)]
        actions = get_actions(env)
        probs = get_prob(π, env)
        @assert length(actions) == length(probs)

        for (a, p) in zip(actions, probs)
            if p > 0 #= ignore illegal action =#
                vals .+= p .* expected_policy_values(π, child(env, a))
            end
        end
        vals
    end
end

function nash_conv(π, env;is_reduce=true,kw...)
    e = copy(env)
    RLBase.reset!(e)

    σ′ = [
        best_response_value(BestResponsePolicy(π, e, i;kw...), e)
        for i in get_players(e)
        if i != get_chance_player(e)
    ]

    σ = expected_policy_values(π, e)
    if is_reduce
        mapreduce(-, +, σ′, σ)
    else
        map(-, σ′, σ)
    end
end


@testset "read" begin
    dir = @__DIR__
    file = joinpath(dir,  "..", "example", "data", "Malitson.csv")

    d0 = SellmeierFit.read(file)
    d1 = SellmeierFit.read(file, unit_prefix=1.0)
    @test (d0.λ == micro .* d1.λ) && d0.ε==d1.ε

    d2 = SellmeierFit.read(file, unit_prefix=1.0, s_col=2, n_col=1)
    @test (d2.λ .^2 == sort(d1.ε)) && (sort(d2.ε) == d1.λ .^ 2)

    d3 = SellmeierFit.read(file, s_var=FREQUENCY)
    @test (d3.λ == reverse(c₀ ./ d0.λ)) && (d3.ε == reverse(d0.ε))

    d4 = SellmeierFit.read(file, s_var=ENERGY)
    @test (d4.λ == reverse(ℎc₀ ./ d0.λ)) && (d4.ε == reverse(d0.ε))

    url = "https://refractiveindex.info/tmp/data/main/SiO2/Malitson.csv"
    try
        d5 = SellmeierFit.read(url)
        @test d5 == d0
    catch e
        @test e isa Downloads.RequestError
    end
end  # @testset "read"


using DataFrames, InteractiveFixedEffectModels, CSV, LinearAlgebra, Test

precision = 2e-1
df = DataFrame(CSV.File(joinpath(dirname(pathof(InteractiveFixedEffectModels)), "../dataset/Cigar.csv")))

for method in [:dogleg, :levenberg_marquardt, :gauss_seidel]
	println(method)
	model = @formula Sales ~ Price + ife(State, Year, 1)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ [328.1653237715761, -1.0415042260420706] .- 1)  < precision
	@test norm(abs(result.augmentdf[1, :factors1]) / 0.228 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 851.514 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -9.7870 - 1) < precision
	@test result.r2_within > 0.0

	show(result)
	model = @formula Sales ~ Price + ife(State, Year, 2)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ [163.01350, -0.40610] .- 1) < precision
	@test norm(abs(result.augmentdf[1, :factors1]) / 0.227 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 184.1897 - 1) < precision
	@test norm(result.augmentdf[1, :residuals]  / -1.774 - 1) < precision
	@test result.r2_within > 0.0

	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(State)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ -0.425389 .- 1) < precision
	@test norm(abs(result.augmentdf[1, :factors1])  / 0.2474 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 123.460 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -5.222 - 1) < precision
	@test result.r2_within > 0.0

	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(Year)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ -0.3744296120563005 .- 1) < precision
	@test norm(abs(result.augmentdf[1, :factors1])  / 0.1918 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 102.336 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -2.2153 - 1) < precision
	@test result.r2_within > 0.0

	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(State) + fe(Year)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ -0.524157 .- 1) < precision
	@test norm(abs(result.augmentdf[1, :factors1]) / 0.256 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 60.0481 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -5.614 - 1) < precision
	@test result.r2_within > 0.0

	# subset
	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(State)
	result = regife(df, model, method = method, subset = df.State .<= 30, save = true)
	@test size(result.augmentdf, 1) == size(df, 1)
	@test norm(abs(result.augmentdf[1, :factors1]) / 0.25965 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 107.832 - 1) < precision
	@test norm(abs(result.augmentdf[1, :factors2])  / 0.2655  - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings2]) / 15.551 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -3.8624  - 1) < precision
	@test norm(result.augmentdf[1, :fe_State] / 131.6162 - 1) < precision
end

# check high dimentional fixed effects are part of factor models
df.State2 = categorical(df.State)
@test_throws ErrorException regife(df, @formula(Sales ~ Price + ife(State, Year, 2) + fe(State2)))

# local minima
using DataFrames, InteractiveFixedEffectModels, Test
precision = 2e-1

for method in [:levenberg_marquardt, :dogleg]
	println(method)

	local df = DataFrame(CSV.File(joinpath(dirname(pathof(InteractiveFixedEffectModels)), "../dataset/Cigar.csv")))
	df.State = categorical(df.State)
	df.Year = categorical(df.Year)
	model = @formula Sales ~ Price + ife(State, Year, 1) + fe(State)
	result = regife(df, model, weights = :Pop, method = method, save = true)
	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(State)
	result = regife(df, model, weights = :Pop, method = method, save = true)

	local df = DataFrame(CSV.File(joinpath(dirname(pathof(InteractiveFixedEffectModels)), "../dataset/EmplUK.csv")))
	df.id1 = df.Firm
	df.id2 = df.Year
	df.pid1 = categorical(df.id1)
	df.pid2 = categorical(df.id2)
	df.y = df.Wage
	df.x1 = df.Emp
	df.w = df.Output
	model = @formula y ~ x1 + ife(pid1, pid2, 2)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ [4.53965, -0.0160858] .- 1) < precision
	model = @formula y ~ x1 + ife(pid1, pid2, 2)
	result = regife(df, model, weights = :w, method = method, save = true)
	@test norm(result.coef ./ [3.47551,-0.017366] .- 1) < precision
	model = @formula y ~ x1 + ife(pid1, pid2, 1)
	result = regife(df, model, weights = :w, method = method, save = true)
	@test norm(result.coef ./ [ -2.62105, -0.0470005] .- 1) < precision
end

# Check formula can be constructed programatically
using StatsModels
result1 = regife(df, @formula(Sales ~ NDI + fe(State) + ife(State, Year, 2)))
result2 = regife(df, Term(:Sales) ~ Term(:NDI) + fe(Term(:State)) + ife(Term(:State), Term(:Year), 2))
@test result1.coef ≈ result2.coef


N = 10
Ngroups = N
groupsInds = distributeAinVecN(Vector(1:Ngroups),N)
tmpDic = countmap(groupsInds)
#number of nodes per group
NgroupsMemb = sortrows([[i for i in keys(tmpDic)] [  tmpDic[i] for i in keys(tmpDic) ] ])

deltaN = 3
groupsPar =  ones(Ngroups)
degGroups =  DegSeq2graphDegSeq(fooErgmBin1, Array{Real,1}(linspace(1+deltaN,N - round(0.1*N)-deltaN,N)))
#tmpdegs = [8, 4, 7, 8, 5, 5, 6, 5, 6, 0]
#sort([7, 5, 8, 9, 8, 5, 4, 2, 8, 4],rev = true)# DegSeq2graphDegSeq(fooErgmBin1, sort([9, 8, 4, 7, 3, 6, 4, 7, 4, 2],rev = true))
Mod = ErgmBin1(degGroups,groupsPar,groupsInds)
groupsPar, nIterm,estMod = estimate(Mod)
firstOrderCond(Mod;deg = tmpdegs,par = groupsPar, groupsInds = groupsInds )
expMatrix(Mod,groupsPar[groupsInds])

## test sampling
groupsPar[10] -= 10
Nsample = 1000
sam = sampl(estMod,Nsample)
meanMatSam = squeeze(mean(sam,3),3)
degsSam = squeeze(sum(meanMatSam,2),2)
degGroupsSam = zeros(Ngroups) ; for i =1:N degGroupsSam[groupsInds[i]] += degsSam[i] end
relErrSam = maximum(abs.(( degGroupsSam - degGroups )./degGroups ))
## Test estimate umbiasedness
estParSam = zeros(N,Nsample)
@time for n=1:Nsample
    degs_n = expValStatsFromMat(fooErgmBin1,sam[:,:,n] )
    estParSam[:,n], ~,~ = estimate(Mod; deg= degs_n)
end

##
using Plots, AverageShiftedHistograms, Distributions, StatPlots
plotly()
listPlots=""
    for n=1:N
        err = ((estParSam .- groupsPar))[n,:]
        rng_space = -3:.1:3
        n==10 ? (rng_space = -300:.1:300) : ()
        tmpAsh = ash(err;rng = rng_space)
        pA1 = plot(tmpAsh)
        expr = "PW$(n) = plot!(Normal(0,1),linewidth = 5) "
        eval(parse(expr))
        listPlots *= "PW$(n) , "
    end

exprLast = "plot(" * listPlots  *" titlefont = font(12) ,titleloc=:center,size=(1350,600),legend = :none )"
 eval(parse(exprLast))

##

#


using Test
using DFTK: energy_ewald
using LinearAlgebra

@testset "Hydrogen atom" begin
    lattice = 16 * Diagonal(ones(3))
    charges = [1]
    positions = [[0,0,0]]

    ref = -0.088665545  # TODO source?
    γ_E = energy_ewald(lattice, charges, positions)
    @test abs(γ_E - ref) < 1e-8
end

@testset "Silicon diamond structure" begin
    lattice = [0.0  5.131570667152971 5.131570667152971;
               5.131570667152971 0.0 5.131570667152971;
               5.131570667152971 5.131570667152971  0.0]
    positions = [ones(3)/8, -ones(3)/8]
    charges = [14, 14]

    ref = -102.8741963352893
    γ_E = energy_ewald(lattice, charges, positions)
    @test abs(γ_E - ref) < 1e-8
end

@testset "Boron molecule" begin
    lattice = 16 * Diagonal(ones(3))
    positions = [[0,0,0], [0.14763485355139283, 0, 0]]
    charges = [5, 5]

    ref = 1.790634595  # TODO source?
    γ_E = energy_ewald(lattice, charges, positions)
    @test abs(γ_E - ref) < 1e-7
end

@testset "Hydrogen molecule" begin
    lattice = 16 * Diagonal(ones(3))
    positions = [
        [0.45312500031210007, 1/2, 1/2],
        [0.5468749996028622, 1/2, 1/2],
    ]
    charges = [1, 1]

    ref = 0.31316999  # TODO source?
    γ_E = energy_ewald(lattice, charges, positions)
    @test abs(γ_E - ref) < 1e-7
end

@testset "Forces" begin
    lattice = [0.0  5.131570667152971 5.131570667152971;
               5.131570667152971 0.0 5.131570667152971;
               5.131570667152971 5.131570667152971  0.0]
    # perturb positions away from equilibrium to get nonzero force
    positions = [ones(3)/8+rand(3)/20, -ones(3)/8]
    charges = [14, 14]

    forces = zeros(Vec3{Float64}, 2)
    γ1 = energy_ewald(lattice, charges, positions, forces=forces)

    # Compare forces to finite differences
    disp = [rand(3)/20, rand(3)/20]
    ε = 1e-8
    γ2 = energy_ewald(lattice, charges, positions .+ ε .* disp)
    @test (γ2-γ1)/ε ≈ -dot(disp, forces) atol=abs(γ1*1e-6)
end


using TensorOperations,SUNRepresentations;

let

d = SUNRepresentations.CGC(Irrep(2,1,0),Irrep(2,1,0),Irrep(2,1,0));
@tensor temp[-1 -2;-3 -4] := d[1 2 -1 -2]*conj(d[1 2 -3 -4])

end


export TranslationOperation

export apply_operation
export domaintype
export isidentity, istranslation, ispoint

"""
    TranslationOperation{S<:Real}

Represents translation symmetry operation

# Fields
* `displacement`: displacement vector

# Examples
```jldoctest
julia> using TightBindingLattice

julia> TranslationOperation([1, 2])
TranslationOperation{Int64}([1, 2])

julia> TranslationOperation([0.5, 0.0, 0.5])
TranslationOperation{Float64}([0.5, 0.0, 0.5])
```
"""
struct TranslationOperation{S<:Real} <:AbstractSpaceSymmetryOperation{S}
    displacement::Vector{S}

    @doc """
        TranslationOperation{S}(displacement::AbstractVector{<:Real}) where {S}

    Construct a translation operation of `displacement` on coordinate space of type `S`.
    """
    function TranslationOperation{S}(displacement::AbstractVector{<:Real}) where {S}
        new{S}(displacement)
    end

    @doc """
        TranslationOperation(displacement::AbstractVector{S}) where {S<:Real}

    Construct a translation operation of `displacement` on coordinate space of type `S`.
    """
    function TranslationOperation(displacement::AbstractVector{S}) where {S<:Real}
        new{S}(displacement)
    end
end

function Base.convert(::Type{TranslationOperation{S}}, obj::IdentityOperation{S}) where S
    dim = dimension(obj)
    return TranslationOperation{S}(zeros(S, dim))
end

function Base.convert(::Type{TranslationOperation{S}}, displacement::AbstractVector{S}) where S
    return TranslationOperation{S}(displacement)
end

function Base.promote_rule(::Type{TranslationOperation{S}}, ::Type{IdentityOperation{S}}) where S
    return TranslationOperation{S}
end

## properties
"""
    isidentity(t::TranslationOperation)

Return true if the translation operation is an identity, i.e. `iszero(t.displacement)`
"""
isidentity(arg::TranslationOperation) = iszero(arg.displacement)

"""
    istranslation(arg::TranslationOperation)

Always return `true`, since `arg` is already a translation operation.
"""
istranslation(arg::TranslationOperation) = true

"""
    ispoint(arg::TranslationOperation)

Return `true` if `arg` is a point operation. The only way this can be true is when `arg` is
an identity operation.
"""
ispoint(arg::TranslationOperation) = iszero(arg.displacement)

"""
    dimension(arg::TranslationOperation)

Spatial dimension of the translation operation
"""
dimension(arg::TranslationOperation) = length(arg.displacement)

Base.hash(arg::TranslationOperation{S}) where S = hash(arg.displacement, hash(TranslationOperation{S}))

## operators
function Base.:(==)(lhs::TranslationOperation{S}, rhs::TranslationOperation{S}) where S
    return lhs.displacement == rhs.displacement
end
function Base.:(==)(top::TranslationOperation{S}, iden::IdentityOperation{S}) where S
    return iszero(top.displacement)
end
function Base.:(==)(iden::IdentityOperation{S}, top::TranslationOperation{S}) where S
    return iszero(top.displacement)
end

function Base.:(*)(lhs::TranslationOperation{S}, rhs::TranslationOperation{S}) where S
    return TranslationOperation{S}(lhs.displacement .+ rhs.displacement)
end

function Base.:(^)(lhs::TranslationOperation{S}, rhs::Real) where S
    return TranslationOperation{S}(lhs.displacement .* rhs)
end

function Base.inv(arg::TranslationOperation{S}) where S
    return TranslationOperation{S}(-arg.displacement)
end

## apply
"""
    apply_operation(symop::TranslationOperation{S}, coordinate::AbstractArray{S}) where S

Return the translated coordinate.
"""
function apply_operation(symop::TranslationOperation{S},
                         coord::AbstractArray{S}) where {S<:Real}
    return coord .+ symop.displacement
end

function (symop::TranslationOperation{S})(coord::AbstractVector{S}) where S
    return coord .+ symop.displacement
end


using BenchmarkTools
using Random
using OMEinsum

Random.seed!(0)

const SUITE = BenchmarkGroup()

SUITE["manyinds"] = BenchmarkGroup()
suite = SUITE["manyinds"]

for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args = ["abcdefghijklmnop,flnqrcipstujvgamdwxyz->bcdeghkmnopqrstuvwxyz"]
    dims = [2]
    for str in args, dim in dims
        code = ein"abcdefghijklmnop,flnqrcipstujvgamdwxyz->bcdeghkmnopqrstuvwxyz"
        arr1 = rand(T, map(i->2, OMEinsum.getixs(code)[1])...)
        arr2 = rand(T, map(i->2, OMEinsum.getixs(code)[2])...)
        suite[string(T)]["tiny"] = @benchmarkable code($arr1,$arr2)
    end
end

# Matrix multiplication
SUITE["matmul"] = BenchmarkGroup()
suite = SUITE["matmul"]
for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args = ["tiny"   => rand(T,2,2),
    "small"  => rand(T,10,10),
    "medium" => rand(T,10^2,10^2),
    "large"  => rand(T, 10^3,10^3)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ij,jk -> ik"($m,$m)
    end
end

# Matrix batch-multiplication
SUITE["batchmul"] = BenchmarkGroup()
suite = SUITE["batchmul"]

for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args = ["tiny"   => rand(T,2,2,3),
    "small"  => rand(T,10,10,3),
    "medium" => rand(T,10^2,10^2,3),
    "large"  => rand(T, 10^3,10^3,3)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ijk,jlk -> ilk"($m,$m)
    end
end

#inner - reduction to scalar
SUITE["dot"] = BenchmarkGroup()
suite = SUITE["dot"]
for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args = ["tiny"   => rand(T,fill(2,3)...)
    "small"  => rand(T,fill(10,3)...)
    "medium" => rand(T,fill(20,3)...)
    "large"  => rand(T,fill(50,3)...)
    "huge"   => rand(T,fill(10^2,3)...)]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ijk,ijk -> "($m, $m)
    end
end

#trace
SUITE["trace"] = BenchmarkGroup()
suite = SUITE["trace"]
for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args = ["tiny"   => rand(T,fill(2,2)...)
    "small"  => rand(T,fill(10,2)...)
    "medium" => rand(T,fill(10^2,2)...)
    "large"  => rand(T,fill(10^3,2)...)]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ii -> "($m)
    end
end

#partial trace
SUITE["ptrace"] = BenchmarkGroup()
suite = SUITE["ptrace"]
for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args = ["tiny"   => rand(T,fill(2,3)...)
    "small"  => rand(T,fill(10,3)...)
    "medium" => rand(T,fill(20,3)...)
    "large"  => rand(T,fill(50,3)...)
    "huge"   => rand(T,fill(10^2,3)...)]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"iij -> j"($m)
    end
end

#diagonal
SUITE["diag"] = BenchmarkGroup()
suite = SUITE["diag"]

for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args =  [
    "tiny"   => rand(T,fill(2,3)...)
    "small"  => rand(T,fill(10,3)...)
    "medium" => rand(T,fill(20,3)...)
    "large" => rand(T,fill(30,3)...)
    "huge" => rand(T,fill(100,3)...)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ijj -> ij"($m)
    end
end

#permutation
SUITE["perm"] = BenchmarkGroup()
suite = SUITE["perm"]
for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args =  [
    "tiny"   => rand(T,fill(1,4)...)
    "small"  => rand(T,fill(2,4)...)
    "medium" => rand(T,fill(10,4)...)
    "large"  => rand(T,fill(30,4)...)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ijkl -> ljki"($m)
    end
end

# tensor contraction
SUITE["tcontract"] = BenchmarkGroup()
suite = SUITE["tcontract"]

for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args =  [
    "tiny"   => rand(T,fill(1,3)...)
    "small"  => rand(T,fill(2,3)...)
    "medium" => rand(T,fill(10,3)...)
    "large"  => rand(T,fill(30,3)...)
    ]

    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ijk, jlk -> il"($m,$m)
    end
end

# star contraction
SUITE["star"] = BenchmarkGroup()
suite = SUITE["star"]

for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args =  [
    "tiny"   => rand(T,fill(2,2)...)
    "small"  => rand(T,fill(10,2)...)
    "medium" => rand(T,fill(30,2)...)
    "large"  => rand(T,fill(100,2)...)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ij,ik,il -> jkl"($m,$m,$m)
    end
end

#star and contract
SUITE["starandcontract"] = BenchmarkGroup()
suite = SUITE["starandcontract"]

for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args =  [
    "tiny"   => rand(T,fill(2,2)...)
    "small"  => rand(T,fill(10,2)...)
    "medium" => rand(T,fill(30,2)...)
    "large"  => rand(T,fill(100,2)...)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ij,ik,ik -> j"($m,$m,$m)
    end
end

# index-sum
SUITE["indexsum"] = BenchmarkGroup()
suite = SUITE["indexsum"]
for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args =  [
    "tiny"   => rand(T,fill(2,3)...)
    "small"  => rand(T,fill(10,3)...)
    "medium" => rand(T,fill(30,3)...)
    "large"  => rand(T,fill(100,3)...)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ijk -> ik"($m)
    end
end
# Hadamard
SUITE["hadamard"] = BenchmarkGroup()
suite = SUITE["hadamard"]
for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args =  [
    "tiny"   => rand(T,fill(2,3)...)
    "small"  => rand(T,fill(10,3)...)
    "medium" => rand(T,fill(30,3)...)
    "large"  => rand(T,fill(100,3)...)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ijk,ijk -> ijk"($m,$m)
    end
end

# Outer
SUITE["outer"] = BenchmarkGroup()
suite = SUITE["outer"]

for T in (Float32, Float64, ComplexF32, ComplexF64)
    suite[string(T)] = BenchmarkGroup()
    args =  [
    "tiny"   => rand(T,fill(2,2)...)
    "small"  => rand(T,fill(10,2)...)
    "medium" => rand(T,fill(50,2)...)
    "large"  => rand(T,fill(100,2)...)
    ]
    for (k,m) in args
        suite[string(T)][k] = @benchmarkable ein"ij,kl -> ijkl"($m,$m)
    end
end


include("heat3d_model.jl")

# computes out <- exp(A * NSTEPS * dt) * b
function _expv(A, b, NSTEPS, dt; hermitian=false, m=min(30, size(A, 1)), tol=1e-7)

    # initialization of the krylov subspace
    TA, Tb = eltype(A), eltype(b)
    T = promote_type(TA, Tb)
    Ks = KrylovSubspace{T, real(T)}(length(b), m)
    arnoldi!(Ks, A, b; m=m, ishermitian=hermitian, tol=tol)

    out = similar(b)
    expv!(out, NSTEPS*dt, Ks)

    return out
end

function _phiv(A, b, NSTEPS, dt; hermitian=false, m=min(30, size(A, 1)), tol=1e-7)

    # initialization of the krylov subspace
    TA, Tb = eltype(A), eltype(b)
    T = promote_type(TA, Tb)
    Ks = KrylovSubspace{T, real(T)}(length(b), m)
    arnoldi!(Ks, A, b; m=m, ishermitian=hermitian, tol=tol)

    out = Matrix{Float64}(undef, size(A, 1), 3)
    phiv!(out, NSTEPS*dt, Ks, 2)

    return view(out, :, 3) .* dt^2
end

function _support_function_Phi_dir(ℓ, A, X0, δ; m=30, tol=1e-7, hermitian=false)
    At = transpose(A)
    vec = _expv(At, ℓ, 1, δ, m=m, tol=tol, hermitian=hermitian)
    ρ(vec, X0)
end

function _support_function_Eplus(ℓ::AbstractVector{N}, A::AbstractMatrix{N}, X0::Hyperrectangle, δ::Float64; cutoff=eps(N), m=min(30, size(A, 1)), tol=1e-7, hermitian=false) where {N}
    c = center(X0)
    r = radius_hyperrectangle(X0)
    A2 = A^2
    rin = abs.(A2 * c) + abs.(A2) * r
    rin = Vector(rin)
    Aabs = abs.(A)
    Pv = _phiv(Aabs, rin, 1, δ; m=m, tol=tol, hermitian=hermitian)
    return dot(abs.(ℓ), Pv)
end

# Non-Krylov version
function forward_non_krylov(instance=heat01; δ=0.02)
    A, Aᵀδ, X0, ℓ = instance(δ=δ)
    prob = @ivp(x' = A*x, x(0) ∈ X0)
    D = discretize(normalize(prob), δ, Forward())
    Ω0 = initial_state(D)
    res = ρ(ℓ, Ω0)
end

# Krylov version of Forward:
# ρ(ℓ, CH(X0, Φ*X0 ⊕ E⨥)) = ρ(ℓ, CH(X0, Φ*X0 ⊕ E⨥)) = max(ρ(ℓ, X0), ρ(Φ^T * ℓ, X0) + ρ(ℓ, E⨥))
function forward_krylov(instance=heat01; δ=0.02, m=30, tol=1e-7, hermitian=false)
    A, Aᵀδ, X0, ℓ = instance(δ=δ)
    aux1 = ρ(ℓ, X0)
    aux2 = _support_function_Phi_dir(ℓ, A, X0, δ, m=m, tol=tol, hermitian=hermitian)
    aux3 = _support_function_Eplus(ℓ, A, X0, δ, m=m, tol=tol, hermitian=hermitian)
    res = max(aux1, aux2 + aux3)
end

# Let Ω0 = CH(X0, Φ*X0 ⊕ E⨥), then we are interested in ρ(ℓ, Φ^k Ω0) for k = 0, 1, ..., NSTEPS-1
function forward_krylov_propagation(instance=heat01; NSTEPS=1, δ=0.02, m=30, tol=1e-7, hermitian=false)
    A, Aᵀδ, X0, ℓ = instance(δ=δ)
    At = copy(transpose(A))

    # initialization of the krylov subspace
    TA, Tb = eltype(At), eltype(ℓ)
    T = promote_type(TA, Tb)
    Ks = KrylovSubspace{T, real(T)}(length(ℓ), m)
    arnoldi!(Ks, At, ℓ; m=m, ishermitian=hermitian, tol=tol)

    ℓ′ = copy(ℓ)
    out = Vector{Float64}()
    for k in 0:NSTEPS-1
        if k > 0
            # ℓ′ = _expv(At, ℓ, k, δ; hermitian=false, m=min(30, size(A, 1)), tol=1e-7)

            # more efficient: use always the same precomputed Krylov subspace
            expv!(ℓ′, k*δ, Ks)
        end
        aux1 = ρ(ℓ′, X0)
        aux2 = _support_function_Phi_dir(ℓ′, A, X0, δ, m=m, tol=tol, hermitian=hermitian)
        aux3 = _support_function_Eplus(ℓ′, A, X0, δ, m=m, tol=tol, hermitian=hermitian)
        res = max(aux1, aux2 + aux3)
        push!(out, res)
    end
    out
end

function run_benchmarks()
    filename = "experiment_krylov_table_1.dat"
    open(filename, "w") do f
        forward_non_krylov(heat01) # warm-up runs
        forward_krylov(heat01)

        # ------------------------------------------------------------------------------
        write(f, "Case Heat01: 5x5x5 grid\n")
        t = @elapsed begin
                res = forward_non_krylov(heat01)
            end
        t = round(t, digits=3)
        write(f, "    Forward non-krylov: runtime: $t, result: $res\n")

        t = @elapsed begin
                res = forward_krylov(heat01, m=30, tol=1e-10, hermitian=true)
            end
        t = round(t, digits=3)
        write(f, "    Forward krylov: runtime: $t, result: $res\n")

        # ------------------------------------------------------------------------------
        write(f, "Case Heat02: 10x10x10 grid\n")
        t = @elapsed begin
                res = forward_non_krylov(heat02)
            end
        t = round(t, digits=3)
        write(f, "    Forward non-krylov: runtime: $t, result: $res\n")

        t = @elapsed begin
                res = forward_krylov(heat02, m=100, tol=1e-10, hermitian=true)
            end
        t = round(t, digits=3)
        write(f,"    Forward krylov: runtime: $t, result: $res\n")

        # ------------------------------------------------------------------------------
        write(f, "Case Heat03: 20x20x20 grid\n")
        write(f, "    Forward non-krylov: runtime: N/A, result: N/A\n")

        t = @elapsed begin
                res = forward_krylov(heat03, m=100, tol=1e-10, hermitian=true)
            end
        t = round(t, digits=3)
        write(f, "    Forward krylov: runtime: $t, result: $res\n")

        # ------------------------------------------------------------------------------
        write(f, "Case Heat04: 50x50x50 grid\n")
        write(f, "    Forward non-krylov: runtime: N/A, result: N/A\n")

        t = @elapsed begin
                res = forward_krylov(heat04, m=100, tol=1e-10, hermitian=true)
            end
        t = round(t, digits=3)
        write(f, "    Forward krylov: runtime: $t, result: $res\n")
    end
end

println("Running Heat3D model ...")
run_benchmarks()

nothing


using OmegaCore
using Distributions
using OmegaDistributions
using Test

@testset "Pointwie" begin
  x = 1 ~ Normal(0, 1)
  y = abs.(x)
  y_, x_ = randsample((y, x))
  @test y_ == abs(x_)
  @test isinferred(randsample, x)
end

@testset "Pointwise class" begin
  μ = @~ Normal(0, 1)
  ϵs = ~ Normal(0, 1)
  ys = μ .+ ϵs
  y1 = 1 ~ ys
  y2 = 2 ~ ys
  randsample((y1, y2))
end

@testset "Pointwise func" begin
  issin = @~ Bernoulli(0.5)
  f = ifelse.(issin, sin, sqrt)

  μ = @~ Normal(0, 1)
  ϵs = ~ Normal(0, 1)
  ys = μ .+ ϵs
  y1 = 1 ~ ys
  y2 = 2 ~ ys
  randsample((y1, y2))
end

function f()
  y = @~ Normal(0, 1)
  z = y .+ y
  randsample(z)
  ω = defω()
  z(ω)
end


using LinearAlgebra

struct PDS
    ∇h
    proxF
    proxG
    L
    Lᵀ

    iterations::Integer
    τ::Real
    σ::Real
    λ::Real

    absTol::Real
end

function PDS(∇h; proxF=(t,a)->t, proxG=(t,a)->t, L=identity, adjL=identity, τ=1.0, σ=1.0, λ=1.0, iterations=nothing, absTol=1e2*eps())
    PDS(∇h, proxF, proxG, L, adjL, iterations, τ, σ, λ, absTol)
end

# FB-based primal-dual splitting algorithm
# N. Komodakis and J. Pesquet, "Playing with Duality," IEEE Signal Processing Magazine, vol. 32, no. 6, pp. 31--54, 2015.
# Algorithm 3
# argmin_{x} f(x) + g(Lx) + h(x)
# f() and g(): proxiable function
# h(x): smooth convex function having a Lipschitzian gradient
# ∇loss : gradient of loss function f(x)
# L: linear operator
# Lᵀ: adjoint of L
function (pds::PDS)(x0, v0=pds.L(x0); iterations=pds.iterations, verboseFunction::Function=(args...)->nothing)
    xₖ = x0
    v = v0
    cproxG = fmconj(pds.proxG)
    for k = 1:iterations
        xₖ₋₁ = xₖ
        p = pds.proxF(xₖ - pds.τ*(pds.∇h(xₖ) + pds.Lᵀ(v)), pds.τ)
        q = cproxG(v + pds.L(2.0*p - xₖ₋₁), pds.σ^-1)

        # (x, v) <- (x, v) + λ((p,q) - (x,v)) for　λ == 1
        xₖ, v = (x, v) .+ λ .* (p - x, q - v)
        # xₖ, v = p, q

        errx = norm(xₖ - xₖ₋₁)^2 /2
        verboseFunction(nItr, xₖ, errx)
        # TODO: break condition
    end
    (xₖ, ())
end


"""
    BurgersPeriodicSemidiscretisation

A semidiscretisation of Burgers' equation
    \$\\partial_t u(t,x) + \\partial_x \\frac{u(t,x)^2}{2} = 0\$
with periodic boundary conditions.
"""
struct BurgersPeriodicSemidiscretisation{T,Derivative<:AbstractDerivativeOperator{T},
                                            Dissipation,
                                            SplitForm<:Union{Val{false}, Val{true}}} <: AbstractSemidiscretisation
    derivative::Derivative
    dissipation::Dissipation
    tmp1::Vector{T}
    tmp2::Vector{T}
    split_form::SplitForm

    function BurgersPeriodicSemidiscretisation(derivative::Derivative, dissipation::Dissipation, split_form::SplitForm=Val{false}()) where {T, Derivative<:AbstractDerivativeOperator{T}, Dissipation, SplitForm<:Union{Val{false}, Val{true}}}
        if dissipation != nothing
            @argcheck size(derivative) == size(dissipation) DimensionMismatch
            @argcheck grid(derivative) == grid(dissipation) ArgumentError
        end
        N = size(derivative, 2)
        tmp1 = Array{T}(undef, N)
        tmp2 = Array{T}(undef, N)
        new{T,Derivative,Dissipation,SplitForm}(derivative, dissipation, tmp1, tmp2, split_form)
    end
end

function Base.show(io::IO, semidisc::BurgersPeriodicSemidiscretisation)
    print(io, "Semidiscretisation of Burgers' equation\n")
    print(io, "  \$ \\partial_t u(t,x) + \\partial_x \\frac{u(t,x)^2}{2} = 0 \$ \n")
    print(io, "with periodic boundaries using")
    if semidisc.split_form == Val{true}()
        print(io, " a split form and: \n")
    else
        print(io, " no split form and: \n")
    end
    print(io, semidisc.derivative)
    print(io, semidisc.dissipation)
end

function (disc::BurgersPeriodicSemidiscretisation)(du, u, p, t)
    @unpack tmp1, tmp2, derivative, dissipation, split_form = disc
    @boundscheck begin
        @argcheck length(u) == length(tmp1)
        @argcheck length(u) == length(tmp2)
        @argcheck length(u) == length(du)
    end

    # volume terms
    if typeof(split_form) <: Val{true}
        m1_3 = -one(eltype(u)) / 3

        ## u * D * u
        mul!(tmp1, derivative, u)
        @. du = m1_3 * u * tmp1
        ## D * u^2
        @. tmp2 = u^2
        mul!(tmp1, derivative, tmp2)
        @. du += m1_3 * tmp1
    else
        m1_2 = -one(eltype(u)) / 2
        @. tmp2 = u^2
        mul!(tmp1, derivative, tmp2)
        @. du = m1_2 * tmp1
    end

    # dissipation
    if dissipation != nothing
        mul!(tmp1, dissipation, u)
        @. du += tmp1
    end

    nothing
end

"""
    BurgersNonperiodicSemidiscretisation

A semidiscretisation of Burgers' equation
    \$\\partial_t u(t,x) + \\partial_x \\frac{u(t,x)^2}{2} = 0\$
with boundary conditions `left_bc(t)`, `right_bc(t)`.
"""
struct BurgersNonperiodicSemidiscretisation{T,Derivative<:AbstractDerivativeOperator{T},
                                            Dissipation,
                                            SplitForm<:Union{Val{false}, Val{true}},
                                            LeftBC, RightBC} <: AbstractSemidiscretisation
    derivative::Derivative
    dissipation::Dissipation
    tmp1::Vector{T}
    tmp2::Vector{T}
    split_form::SplitForm
    left_bc::LeftBC
    right_bc::RightBC

    function BurgersNonperiodicSemidiscretisation(derivative::Derivative, dissipation::Dissipation, split_form::SplitForm, left_bc::LeftBC, right_bc::RightBC) where {T, Derivative<:AbstractDerivativeOperator{T}, Dissipation, SplitForm<:Union{Val{false}, Val{true}}, LeftBC, RightBC}
        if dissipation != nothing
            @argcheck size(derivative) == size(dissipation) DimensionMismatch
            @argcheck grid(derivative) == grid(dissipation) ArgumentError
        end
        N = size(derivative, 2)
        tmp1 = Array{T}(undef, N)
        tmp2 = Array{T}(undef, N)
        new{T,Derivative,Dissipation,SplitForm,LeftBC,RightBC}(derivative, dissipation, tmp1, tmp2, split_form, left_bc, right_bc)
    end
end

function Base.show(io::IO, semidisc::BurgersNonperiodicSemidiscretisation)
    print(io, "Semidiscretisation of Burgers' equation\n")
    print(io, "  \$ \\partial_t u(t,x) + \\partial_x \\frac{u(t,x)^2}{2} = 0 \$ \n")
    print(io, "with nonperiodic boundaries using")
    if semidisc.split_form == Val{true}()
        print(io, " a split form and: \n")
    else
        print(io, " no split form and: \n")
    end
    print(io, semidisc.derivative)
    print(io, semidisc.dissipation)
end

function godunov_flux_burgers(uₗ::T, uᵣ::T) where {T<:Real}
    if uₗ < uᵣ
        if uₗ < 0 && 0 < uᵣ
            zero(T)
        else
            min(uₗ^2/2, uᵣ^2/2)
        end
    else
        max(uₗ^2/2, uᵣ^2/2)
    end
end

function (disc::BurgersNonperiodicSemidiscretisation)(du, u, p, t)
    @unpack tmp1, tmp2, derivative, dissipation, split_form, left_bc, right_bc = disc
    @boundscheck begin
        @argcheck length(u) == length(tmp1)
        @argcheck length(u) == length(tmp2)
        @argcheck length(u) == length(du)
    end

    # volume terms
    if typeof(split_form) <: Val{true}
        m1_3 = -one(eltype(u)) / 3

        ## u * D * u
        mul!(tmp1, derivative, u)
        @. du = m1_3 * u * tmp1
        ## D * u^2
        @. tmp2 = u^2
        mul!(tmp1, derivative, tmp2)
        @. du += m1_3 * tmp1
    else
        m1_2 = -one(eltype(u)) / 2
        @. tmp2 = u^2
        mul!(tmp1, derivative, tmp2)
        @. du = m1_2 * tmp1
    end

    # dissipation
    if dissipation != nothing
        mul!(tmp1, dissipation, u)
        @. du += tmp1
    end

    # boundary conditions via Godunov's flux
    @inbounds fnum_left = godunov_flux_burgers(left_bc(t), u[1])
    @inbounds du[1] += (fnum_left - u[1]^2/2) / left_boundary_weight(derivative)
    @inbounds fnum_right = godunov_flux_burgers(u[end], right_bc(t))
    @inbounds du[end] -= (fnum_right - u[end]^2/2) / right_boundary_weight(derivative)

    nothing
end


#Implementation of the EM solver.
struct EMproto <: Solver end #prototyping code to check for correctness
struct EM <: Solver end

function integrate(a!::Function, b!::Function, x0::AbstractVector{T}, t::AbstractVector{<:Real},
    noise::Noise, em::EMproto, callback::Function) where {T<:Number}
    W = W0(noise)
    x = copy(x0)

    n = length(x0)
    f = Vector{T}(n)
    g = Matrix{T}(n, noise.m)
    ΔW = copy(W)

    #save the first data point
    callback(t[1], x, W)

    for i in 1:length(t)-1
        Δt = t[i+1]-t[i]
        randn!(noise, ΔW)
        ΔW = ΔW*sqrt(Δt)

        a!(t[i], x, f)
        b!(t[i], x, g)

        x .= x + f*Δt + g*ΔW
        W .= W + ΔW
        callback(t, x, W)
    end
    return nothing
end

"""
explicitEM function that is optimized
ab! : bias & diffusion function
x0 : Initial state
t : times to use for time evolution
noise : Noise object
callback : a function that runs on t, x, W
"""
function integrate(a!::Function, b!::Function, x0::AbstractVector{T}, t::AbstractVector{<:Real},
    noise::Noise, em::EM, callback::Function) where {T<:Number}
    W = W0(noise)
    x = copy(x0)

    n = length(x0)
    fΔt = Vector{T}(n)
    g = Matrix{T}(n, noise.m)
    ΔW = copy(W)
    gΔW = Vector{T}(n)

    #save the first data point
    callback(t[1], x, W)

    for i in 1:length(t)-1
        Δt = t[i+1]-t[i]
        randn!(noise, ΔW)
        ΔW .*= sqrt(Δt)

        a!(t[i], x, fΔt)
        b!(t[i], x, g)

        fΔt .*= Δt
        A_mul_B!(gΔW, g, ΔW) #This line performs better than the more readable version!

        x .+= fΔt
        x .+= gΔW
        W .+= ΔW
        callback(t, x, W)
    end
    return nothing
end


using Test
using SimpleQuantum

q = Q0
@test measure!(q) == 0
q = Q1
@test measure!(q) == 1

total = 0
for k = 1:1000
    local q = Qubit(3, 4im)
    local v = measure!(q)
    global total += v
    if v == 0
        @test q == Q0
    else
        @test q == Q1
    end
end
println("total = $total which should be around 640")

total = 0
for k = 1:1000
    local q = RandomQubit()
    local v = measure!(q)
    global total += v
end
println("Generated 1000 random qubits and got $total ones")

q = Qubit(3, 4im)
w = q.vec[1]
z = q.vec[2]
@test w == 0.6
@test z == 0.8im