function create_subclusters_labels!(labels::AbstractArray{Int64,2},points::AbstractArray{Float64,2},cluster_params::splittable_cluster_params)
    if size(labels,1) == 0
        return
    end
    lr_arr = zeros(length(labels), 2)
    log_likelihood!(lr_arr[:,1],points,cluster_params.cluster_params_l.distribution)
    log_likelihood!(lr_arr[:,2],points,cluster_params.cluster_params_r.distribution)
    lr_arr[:,1] .+= log(cluster_params.lr_weights[1])
    lr_arr[:,2] .+= log(cluster_params.lr_weights[2])
    sample_log_cat_array!(labels,lr_arr)
end

function sample_labels!(labels::AbstractArray{Int64,2},points::AbstractArray{Float64,2},clusters_samples::Dict,final::Bool)
    parr = zeros(length(labels), length(clusters_samples))
    for (k,v) in clusters_samples
        log_likelihood!((@view parr[:,k]),points,v)
    end
    sample_log_cat_array!(labels,parr,final)
end

function create_splittable_from_params(params::cluster_parameters, α::Float64)
    params_l = deepcopy(params)
    params_l.distribution = sample_distribution(params.posterior_hyperparams)
    params_r = deepcopy(params)
    params_r.distribution = sample_distribution(params.posterior_hyperparams)
    #params_l = cluster_parameters(params.distribution_hyper_params, sample_distribution(params.posterior_hyperparams), params.suff_stats, params.posterior_hyperparams)
    #params_r = cluster_parameters(params.distribution_hyper_params, sample_distribution(params.posterior_hyperparams), params.suff_stats, params.posterior_hyperparams)
    lr_weights = rand(Dirichlet([α / 2, α / 2]))
    return splittable_cluster_params(params,params_l,params_r,lr_weights, false, ones(20).*-Inf)
end

function merge_clusters_to_splittable(cpl::cluster_parameters,cpr::cluster_parameters, α::Float64)
    suff_stats = aggregate_suff_stats(cpl.suff_statistics,cpr.suff_statistics)
    posterior_hyperparams = calc_posterior(cpl.hyperparams, suff_stats)
    lr_weights = rand(Dirichlet([cpl.suff_statistics.N + (α / 2), cpr.suff_statistics.N + (α / 2)]))
    cp = cluster_parameters(cpl.hyperparams, cpl.distribution, suff_stats, posterior_hyperparams)
    return splittable_cluster_params(cp,cpl,cpr,lr_weights, false, ones(20).*-Inf)
end

function should_split!(should_split::AbstractArray{Float64,1}, cluster_params::splittable_cluster_params, α::Float64, final::Bool)

    cpl = cluster_params.cluster_params_l
    cpr = cluster_params.cluster_params_r
    cp = cluster_params.cluster_params
    if final || cpl.suff_statistics.N == 0 ||cpr.suff_statistics.N == 0
        should_split .= 0
        return
    end
    log_likihood_l = log_marginal_likelihood(cpl.hyperparams, cpl. posterior_hyperparams, cpl.suff_statistics)
    log_likihood_r = log_marginal_likelihood(cpr.hyperparams, cpr. posterior_hyperparams, cpr.suff_statistics)
    log_likihood = log_marginal_likelihood(cp.hyperparams, cp. posterior_hyperparams, cp.suff_statistics)
    log_HR = log(α) + logabsgamma(cpl.suff_statistics.N)[1] + log_likihood_l + logabsgamma(cpr.suff_statistics.N)[1] + log_likihood_r-(logabsgamma(cp.suff_statistics.N)[1] + log_likihood)
    if log_HR > log(rand())
        should_split .= 1
    end
end

function should_split!(cluster_params::splittable_cluster_params, α::Float64, final::Bool)

    cpl = cluster_params.cluster_params_l
    cpr = cluster_params.cluster_params_r
    cp = cluster_params.cluster_params
    if final || cpl.suff_statistics.N == 0 ||cpr.suff_statistics.N == 0
        should_split .= 0
        return
    end
    log_likihood_l = log_marginal_likelihood(cpl.hyperparams, cpl. posterior_hyperparams, cpl.suff_statistics)
    log_likihood_r = log_marginal_likelihood(cpr.hyperparams, cpr. posterior_hyperparams, cpr.suff_statistics)
    log_likihood = log_marginal_likelihood(cp.hyperparams, cp. posterior_hyperparams, cp.suff_statistics)
    log_HR = log(α) + logabsgamma(cpl.suff_statistics.N)[1] + log_likihood_l + logabsgamma(cpr.suff_statistics.N)[1] + log_likihood_r-(logabsgamma(cp.suff_statistics.N)[1] + log_likihood)
    # println(log_HR)
end

function update_splittable_cluster_params!(splittable_cluser::splittable_cluster_params,
        points::AbstractArray{Float64,2},
        sub_labels::AbstractArray{Int64,1},
        is_global::Bool,
        pts_to_groups = -1)
    cpl = splittable_cluser.cluster_params_l
    cpr = splittable_cluser.cluster_params_r
    cp = splittable_cluser.cluster_params
    if is_global
        cp.suff_statistics = create_sufficient_statistics(cp.hyperparams,cp.posterior_hyperparams, points, pts_to_groups)
        pts_gl = @view pts_to_groups[(@view (sub_labels .<= 2)[:])]
        pts_gr = @view pts_to_groups[(@view (sub_labels .> 2)[:])]
        cpl.suff_statistics = create_sufficient_statistics(cpl.hyperparams,cpl.posterior_hyperparams, (@view points[:,(@view (sub_labels .<= 2)[:])]),pts_gl)
        cpr.suff_statistics = create_sufficient_statistics(cpr.hyperparams,cpr.posterior_hyperparams, (@view points[:,(@view (sub_labels .> 2)[:])]),pts_gr)
    else
        cp.suff_statistics = create_sufficient_statistics(cp.hyperparams,cp.posterior_hyperparams, points)
        cpl.suff_statistics = create_sufficient_statistics(cpl.hyperparams,cpl.posterior_hyperparams, @view points[:,(@view (sub_labels .% 2 .== 1)[:])])
        cpr.suff_statistics = create_sufficient_statistics(cpr.hyperparams,cpr.posterior_hyperparams, @view points[:,(@view (sub_labels .% 2 .== 0)[:])])
    end
    cp.posterior_hyperparams = calc_posterior(cp.hyperparams, cp.suff_statistics)
    cpl.posterior_hyperparams = calc_posterior(cpl.hyperparams, cpl.suff_statistics)
    cpr.posterior_hyperparams = calc_posterior(cpr.hyperparams, cpr.suff_statistics)
end

function update_splittable_cluster_params(splittable_cluser::splittable_cluster_params,
        points::AbstractArray{Float64,2},
        sub_labels::AbstractArray{Int64,1},
        is_global::Bool,
        pts_to_groups = -1)
    cpl = splittable_cluser.cluster_params_l
    cpr = splittable_cluser.cluster_params_r
    cp = splittable_cluser.cluster_params
    # println(sum((@view (sub_labels .> 2)[:])))
    if is_global
        cp.suff_statistics = create_sufficient_statistics(cp.hyperparams,cp.posterior_hyperparams, points, pts_to_groups)
        pts_gl = @view pts_to_groups[(@view (sub_labels .<= 2)[:])]
        pts_gr = @view pts_to_groups[(@view (sub_labels .> 2)[:])]
        cpl.suff_statistics = create_sufficient_statistics(cpl.hyperparams, cpl.posterior_hyperparams, (@view points[:,(@view (sub_labels .<= 2)[:])]),pts_gl)
        cpr.suff_statistics = create_sufficient_statistics(cpr.hyperparams, cpr.posterior_hyperparams,(@view points[:,(@view (sub_labels .> 2)[:])]),pts_gr)
    else
        cp.suff_statistics = create_sufficient_statistics(cp.hyperparams,cp.posterior_hyperparams,points)
        cpl.suff_statistics = create_sufficient_statistics(cpl.hyperparams, cpl.posterior_hyperparams,@view points[:,(@view (sub_labels .% 2 .== 1)[:])])
        cpr.suff_statistics = create_sufficient_statistics(cpr.hyperparams, cpr.posterior_hyperparams,@view points[:,(@view (sub_labels .% 2 .== 0)[:])])
    end
    begin
        cp.posterior_hyperparams = calc_posterior(cp.hyperparams, cp.suff_statistics)
        cpl.posterior_hyperparams = calc_posterior(cpl.hyperparams, cpl.suff_statistics)
        cpr.posterior_hyperparams = calc_posterior(cpr.hyperparams, cpr.suff_statistics)
    end
    return splittable_cluser
end

function sample_cluster_params!(params::splittable_cluster_params, α::Float64, is_zero_dim = false)
    points_count = Vector{Float64}()
    params.cluster_params.distribution = sample_distribution(params.cluster_params.posterior_hyperparams)
    params.cluster_params_l.distribution = sample_distribution(params.cluster_params_l.posterior_hyperparams)
    params.cluster_params_r.distribution = sample_distribution(params.cluster_params_r.posterior_hyperparams)
    push!(points_count, params.cluster_params_l.suff_statistics.N)
    push!(points_count, params.cluster_params_r.suff_statistics.N)
    points_count .+= α/2
    params.lr_weights = rand(Dirichlet(points_count))
    log_likihood_l = log_marginal_likelihood(params.cluster_params_l.hyperparams,params.cluster_params_l.posterior_hyperparams, params.cluster_params_l.suff_statistics)
    log_likihood_r = log_marginal_likelihood(params.cluster_params_r.hyperparams,params.cluster_params_r.posterior_hyperparams, params.cluster_params_r.suff_statistics)

    # params.logsublikelihood_hist[1:4] = params.logsublikelihood_hist[2:5]
    # params.logsublikelihood_hist[5] = log_likihood_l + log_likihood_r
    # logsublikelihood_now = 0.0
    # for i=1:5
    #     logsublikelihood_now += params.logsublikelihood_hist[i] *0.20
    # end
    # if logsublikelihood_now != -Inf && logsublikelihood_now - params.logsublikelihood_hist[5] < 1e-2 # propogate abs change to other versions?
    #     params.splittable = true
    # end
    # println(split_delays)
    # logsublikelihood_now = (log_likihood_l+log_likihood_r) /  (params.cluster_params_l.suff_statistics.N+ params.cluster_params_r.suff_statistics.N)
    # if logsublikelihood_now - params.logsublikelihood_hist[1] < 0
    #     params.splittable = true
    # end
    # params.logsublikelihood_hist[1] = logsublikelihood_now
    #
    # if split_delays == false
    #     params.splittable = true
    # end
    params.logsublikelihood_hist[1:burnout_period-1] = params.logsublikelihood_hist[2:burnout_period]
    params.logsublikelihood_hist[burnout_period] = log_likihood_l + log_likihood_r
    logsublikelihood_now = 0.0

    for i=1:burnout_period
        logsublikelihood_now += params.logsublikelihood_hist[i] *(1/(burnout_period-0.1))
    end

    if logsublikelihood_now != -Inf && logsublikelihood_now - params.logsublikelihood_hist[burnout_period] < 1e-2 # propogate abs change to other versions?
        # println(params.logsublikelihood_hist)
        params.splittable = true
    end
    if is_zero_dim == true
        params.splittable = true
    end
    return params.cluster_params.suff_statistics.N
end

function sample_cluster_params!(params::splittable_cluster_params, α::Float64, counts::AbstractArray{Int64,1})
    points_count = [params.cluster_params_l.suff_statistics.N, params.cluster_params_r.suff_statistics.N]
    params.cluster_params.distribution = sample_distribution(params.cluster_params.posterior_hyperparams)
    params.cluster_params_l.distribution = sample_distribution(params.cluster_params_l.posterior_hyperparams)
    params.cluster_params_r.distribution = sample_distribution(params.cluster_params_r.posterior_hyperparams)
    points_count .+= α/2
    # println(points_count)
    params.lr_weights = rand(Dirichlet(points_count))
    log_likihood_l = log_marginal_likelihood(params.cluster_params_l.hyperparams,params.cluster_params_l.posterior_hyperparams, params.cluster_params_l.suff_statistics)
    log_likihood_r = log_marginal_likelihood(params.cluster_params_r.hyperparams,params.cluster_params_r.posterior_hyperparams, params.cluster_params_r.suff_statistics)

    # params.logsublikelihood_hist[1:4] = params.logsublikelihood_hist[2:5]
    # params.logsublikelihood_hist[5] = log_likihood_l + log_likihood_r
    # logsublikelihood_now = 0.0
    # for i=1:5
    #     logsublikelihood_now += params.logsublikelihood_hist[i] *0.20
    # end
    # if logsublikelihood_now != -Inf && logsublikelihood_now - params.logsublikelihood_hist[5] < 1e-2 # propogate abs change to other versions?
    #     params.splittable = true
    # end
    params.logsublikelihood_hist[1:burnout_period-1] = params.logsublikelihood_hist[2:burnout_period]
    params.logsublikelihood_hist[burnout_period] = log_likihood_l + log_likihood_r
    logsublikelihood_now = 0.0
    for i=1:burnout_period
        logsublikelihood_now += params.logsublikelihood_hist[i] *(1/(burnout_period-0.1))
    end

    if logsublikelihood_now != -Inf && logsublikelihood_now - params.logsublikelihood_hist[burnout_period] < 1e-2 # propogate abs change to other versions?
        # println(params.logsublikelihood_hist)
        params.splittable = true
    end
    return params.cluster_params.suff_statistics.N
end


using SpikingNN
using Plots

# SRM0 params
η₀ = 5.0
τᵣ = 1.0

# Input spike train params
rate = 0.01
T = 20
∂t = 0.01
n = convert(Int, ceil(T / ∂t))

srm = Neuron(QueuedSynapse(Synapse.Delta()), SRM0(η₀, τᵣ), Threshold.Poisson(5.0, 0.5, 0.1))
input = ConstantRate(rate)
spikes = excite!(srm, input, n)

# callback to record voltages
voltages = Float64[]
record = function ()
    push!(voltages, getvoltage(srm))
end

# simulate
@time output = simulate!(srm, n; dt = ∂t, cb = record, dense = true)

# plot raster plot
raster_plot = rasterplot(∂t .* spikes, ∂t .* output, label = ["Input", "Output"], xlabel = "Time (sec)",
                title = "Raster Plot (δ response)")
xlims!(0, T)

# plot dense voltage recording
plot(∂t .* collect(1:n), voltages,
    title = "SRM Membrane Potential with Varying Presynaptic Responses", xlabel = "Time (sec)", ylabel = "Potential (V)", label = "δ response")

# resimulate using presynaptic response
srm = Neuron(QueuedSynapse(Synapse.Alpha()), SRM0(η₀, τᵣ), Threshold.Poisson(5.0, 0.5, 0.1))
voltages = Float64[]
excite!(srm, spikes)
@time simulate!(srm, n; dt = ∂t, cb = record, dense = true)

# plot voltages with response function
voltage_plot = plot!(∂t .* collect(1:n), voltages, label = "α response")
xlims!(0, T)

plot(raster_plot, voltage_plot, layout = grid(2, 1), xticks = 0:T)


using EVQueues, Plots

lambda=120.0;
mu=1.0;
gamma=1.0;
C=60.0;

Tfinal=1000.0;

sim = ev_edf(lambda,mu,gamma,Tfinal,C,snapshots=[Tfinal])
compute_statistics!(sim)

#TIme plot of occupation and power
p=plot(sim.timetrace)
display(p)

#CDF of departure attained workloads
p=servicecdf(sim.EVs)
display(p)

#State space of the last snapshot
snap = sim.snapshots[end];

p=stateplot(snap.charging)
display(p)

#fairness measure over time

t=collect(0:10:Tfinal);
h=20.0;
J=compute_fairness(sim,t,h)

p = plot(   xlabel = "Time",
            ylabel = "J",
            title = "Fairness measure",
            ylim = (0,1)
            )

plot!(p,t,J,lw=2,legend=:none)

display(p)


function C3(t, fft_u, fft_v,
	    A1, A2,
	    matr, conjmatr,
	    sigma, llambda,
	    Ktaubis, epsilon, Ntaumm)

    h1u, h1v = ftau(t, fft_u, fft_v, A1, A2,
		    matr, conjmatr, sigma, llambda)

    h1u_fft = fft(h1u, dims=1)
    h1v_fft = fft(h1v, dims=1)

    dtu0u .= h1u_fft[1, :] / Ntaumm
    dtu0v .= h1v_fft[1, :] / Ntaumm
    h1u_fft[1, :] .= 0
    h1v_fft[1, :] .= 0
    h1u = epsilon * ifft(h1u_fft / (1im * Ktaubis), dims=1)
    h1v = epsilon * ifft(h1v_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t, fft_u, fft_v,
			      dtu0u, dtu0v,
			      A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t, fft_u, fft_v,
			      A1, A2,
			      matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu = champu1 + champu2
    champv = champv1 + champv2

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    dttu0u = champu_fft[1, :] / Ntaumm
    dttu0v = champv_fft[1, :] / Ntaumm

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth1u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dth1v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    champu, champv = ftau(t, fft_u .+ h1u, fft_v .+ h1v,
			  A1, A2,
			  matr, conjmatr,
			  sigma, llambda)

    h2u = champu - dth1u
    h2v = champv - dth1v

    h2u_fft = fft(h2u, dims=1)
    h2v_fft = fft(h2v, dims=1)

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    dtu1u = champu_fft[1, :] / Ntaumm
    dtu1v = champv_fft[1, :] / Ntaumm

    h2u_fft[1, :] .= 0
    h2v_fft[1, :] .= 0

    h2u = epsilon * ifft(h2u_fft ./ (1im * Ktaubis), dims=1)
    h2v = epsilon * ifft(h2v_fft ./ (1im * Ktaubis), dims=1)

    champu1, champv1 = d2uftau(t, fft_u, fft_v,
			       dtu0u, dtu0v, dtu0u, dtu0v,
			       A1, A2, matr, conjmatr,
			       sigma, llambda, Ntaumm)

    champu2, champv2 = dtduftau(t, fft_u, fft_v,
				dtu0u, dtu0v,
				A1, A2, matr, conjmatr,
				sigma, llambda, Ntaumm)

    champu3, champv3 = duftau(t, fft_u, fft_v,
			      dttu0u, dttu0v,
			      A1, A2, matr, conjmatr,
			      sigma, llambda, Ntaumm)

    champu4, champv4 = d2tftau(t, fft_u, fft_v,
			       A1, A2, matr, conjmatr,
			       sigma, llambda, Ntaumm)

    champu = champu1 .+ 2 * champu2 .+ champu3 .+ champu4
    champv = champv1 .+ 2 * champv2 .+ champv3 .+ champv4

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dtth1u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dtth1v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t, fft_u .+ h1u, fft_v .+ h1v,
                              dtu1u .+ dth1u, dtu1v .+ dth1v,
			      A1, A2, matr,
                              conjmatr, sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t, fft_u .+ h1u, fft_v .+ h1v,
			      A1, A2, matr,
                              conjmatr, sigma, llambda, Ntaumm)

    champu = champu1 + champu2 - dtth1u
    champv = champv1 + champv2 - dtth1v

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth2u = epsilon * ifft(champu_fft ./ (1im * Ktaubis), dims=1)
    dth2v = epsilon * ifft(champv_fft ./ (1im * Ktaubis), dims=1)

    h3u, h3v = ftau(t, fft_u .+ h2u, fft_v .+ h2v,
		    A1, A2, matr, conjmatr,
                    sigma, llambda)

    h3u = h3u - dth2u
    h3v = h3v - dth2v
    h3u_fft = fft(h3u, dims=1)
    h3v_fft = fft(h3v, dims=1)
    h3u_fft[1, :] .= 0
    h3v_fft[1, :] .= 0
    h3u = epsilon * ifft(h3u_fft / (1im * Ktaubis), dims=1)
    h3v = epsilon * ifft(h3v_fft / (1im * Ktaubis), dims=1)

    C3u = fft_u .+ h3u
    C3v = fft_v .+ h3v

    C3u, C3v

end


j = 30
maxn = 70
ω = 1.0
ω₀ = 1.0
δ = 1.0

λ = 1.0
λᵪ= sqrt(ω * ω₀) / (1.0 + δ)

G = 2.0 * λ / (sqrt(2.0 * j) * ω)
logG = log(G)

function Index(n, m)
    return n * (2.0 * j + 1) + m + j + 1
end

function LogFactorial(n)
    if n <= 1
        return 0
    else
        return log(n) + LogFactorial(n - 1)
    end
end

function Minus(n)
    if n % 2 == 0
        return 1
    else
        return -1
    end
end

function fu(n1, n2)
    coef = 0.5 * (LogFactorial(n1) + LogFactorial(n2) - G*G)

    sum = 0
    for r = 0:min(n1, n2)
        sum += Minus(r) * exp(logG * (n1 + n2 - 2.0 * r) - LogFactorial(n1 - r) - LogFactorial(n2 - r) - LogFactorial(r))
    end

    return sum
end

function Fuf(np, mp, n, m)
    if m + 1 == mp
        return -0.5 * ω₀ * sqrt(j * (j + 1) - m * (m + 1)) * Minus(n) * fu(np + 1, n + 1)
    elseif m - 1 == mp
        return -0.5 * ω₀ * sqrt(j * (j + 1) - m * (m - 1)) * Minus(np) * fu(np + 1, n + 1)
    end

    return 0.0
end


module QuasiHamiltonianRTN
using Expokit # expmv
using ExpmV
using SparseArrays
using LinearAlgebra

include("SpecialUnitary.jl")
include("Utilities.jl")

include("Lattice1D.jl")
include("StarGraph.jl")

export quasiHamiltonian, evolution, bloch_vector, density_operator, localized_state

"""
    Vnoise(n, γ)

Returns the generator ``V`` for the transition probabilities
between the various noise configurations for `n` fluctuators,
with switching rate ``\\gamma``.

It is defined as

``\\mathbf{V} = \\sum_{i=1}^N \\mathbf{V}_\\gamma^{(i)}``

with ``\\mathbf{V}_\\gamma^{(i)} = \\mathbf{I}_2^{\\otimes i-1} \\otimes \\mathbf{V}_\\gamma \\otimes \\mathbf{I}_2^{\\otimes N-i}``.
"""
function Vnoise(n, γ)
    V = sparse([-γ γ; γ -γ])

    function Vi(i)
        foldl(kron, [j == i ? V : sparse(1.0I, 2, 2) for j = 1:n])
    end

    sum([Vi(i) for i = 1:n])
end

"""
    quasiHamiltonian(Hamiltonian, γ)

Returns the quasi-Hamiltonian ``H_q`` for the Hamiltonian specified by the
array of matrices `H`, representing all the possible configurations of the
Hamiltonian, with switching rate γ.

The quasi-Hamiltonian is defined as in Eq. (20) of Joynt et al.
This function returns ``-i H_q``, in order to deal with real matrices instead of
imaginary ones.

## Arguments

* `Hamiltonian` is an array of matrices, each of them representing the Hamiltonian
with a different configuration of the noise.
* `γ` is the switching rate of the noise.
"""
function quasiHamiltonian(Hamiltonian, γ)::SparseMatrixCSC{Float64,Int64}
    n = size(Hamiltonian[1], 1)
    Nc = length(Hamiltonian)
    RTN_number = Int(log2(Nc))
    λ::Array{SparseArrays.SparseMatrixCSC{Complex{Float64},Int64},1} = sun_generators(n)
    f::Array{SparseArrays.SparseMatrixCSC{Float64,Int64},3} = structure_constants(λ)

    res::Array{SparseArrays.SparseMatrixCSC{Float64,Int64},1} = map(h -> igen(f,λ,h), Hamiltonian)
    Hq = blockdiag(res...) + kron(-Vnoise(RTN_number,γ), sparse(1.0I, n^2-1, n^2-1))

    return Hq
end

"""
    evolution(Hq, r0, t; [expmpkg=:ExpmV])

Evaluates the state of the system with quasi-Hamiltonian matrix `Hq`, starting
from an initial Bloch vector `r0`, at time `t`.

If `t` is a `StepRangeLen` object of length ```N_t``, return a ``N_r \\times N_t``
array containing the Bloch vector of the system at each time step as columns.
"""
function evolution(Hq::SparseMatrixCSC, r0::Vector, t; expmpkg=:ExpmV)
    @assert expmpkg ∈ (:ExpmV, :Expokit) "expmpkg must be either ExpmV or Expokit"

    n = size(Hq, 1)
    # n = Nc * Nq
    Nq = length(r0)
    Nc = Int(n / Nq)

    rtnstate = ones(Nc)/sqrt(Nc)

    v0 = kron(rtnstate, r0)
    y = kron(rtnstate, sparse(1.0I, Nq, Nq))'

    if expmpkg == :ExpmV
        res = ExpmV.expmv(-t, Hq, v0)

    elseif expmpkg == :Expokit
        res = Array{Float64}(undef, n, length(t))
        for (i, ti) in enumerate(t)
            res[:, i] = Expokit.expmv(-ti, Hq, v0)
        end
    end

    return y * res
end

function evolution(Hamiltonian::Array{SparseMatrixCSC{Float64,Int64},1},
            r0::Vector, t; γ=1, expmpkg=:ExpmV)
    n = size(Hamiltonian[1], 1)
    Nc = length(Hamiltonian)
    Nq = length(r0)

    if n^2-1 != Nq
        throw(ArgumentError("Hamiltonian and r0 have incompatible dimensions"))
    end

    Hq = quasiHamiltonian(Hamiltonian, γ)

    return evolution(Hq, r0, t; expmpkg=expmpkg)
end

end # module


function realization(ngrid,H,Br,m,y,bx,by,pd,q)

# computes the velocity and acceleration on a lat,long grid
# at a series of snapshots.
#
# input
#
# ngrid -  number of grid points in colatitude
# H    - layer thickness (km)
# Np   - dimensionless buoyancy frequency (N/Omega)
# Br   - radial magnetic field  (mT)
# m    - angular order m
# y    - meridional coordinate
# bx,by - x and y components of magnetic perturbation
# pd    - period  (years)
# q     - quality factor
#
# output
#
# tvec      - array to times to evaluate velocity
# lat,long  -  array of latitude and longitude
# vy,vx     -  theta and phi components of velocity at CMB
# ay,ax    -  theta and phi components of acceleration

# array of times
dt = 0.5;
tvec = (0.0 : dt : abs(pd));

# latitude and longitude grid
#ngrid = 100;
dlat = pi / ngrid;
dlong = pi / ngrid;
lat = (-pi/2+dlat/2 : dlat : pi/2-dlat/2);
long = (dlong/2 : dlong : 2 * pi-dlong/2);
colat = (dlat/2 : dlat : pi-dlat/2);  # reverse order

# allocate memory for velocity and acceleration
vx = Array{Float64,3}(undef,length(tvec),ngrid,2*ngrid);
vy = Array{Float64,3}(undef,length(tvec),ngrid,2*ngrid);
ax = Array{Float64,3}(undef,length(tvec),ngrid,2*ngrid);
ay = Array{Float64,3}(undef,length(tvec),ngrid,2*ngrid);

# loop over times
for k = 1 : length(tvec)

   tstep = tvec[k];
   vxt,vyt,axt,ayt = velocity(H,Br,m,y,lat,long,bx,by,pd,q,tstep);
   vx[k,:,:]=vxt;
   vy[k,:,:]=vyt;
   ax[k,:,:]=axt;
   ay[k,:,:]=ayt;

end

# collect into arrays
t = collect(tvec);
clt = collect(colat);
lg = collect(long);

return t,clt,lg,vx,vy,ax,ay

end


"""
    draw_gif_result(time, x, y, ψ, shape, file_path, [, interval, fps]) -> gif

Draw the gif animation from simulation result.
"""
function draw_gif_result(time, x, y, ψ, shape, file_path; interval = 10, fps = 100)
    anim = @animate for i = 1:interval:length(time)
        plot(
            x,
            y,
            label = "",
            xlabel = "x [m]",
            ylabel = "y [m]",
            linestyle = :dot,
            aspect_ratio = :equal,
        )
        ship = square(x[i], y[i], shape, ψ[i])
        plot!(Shape(ship[1], ship[2]), label = "")
        scatter!(
            [x[i]],
            [y[i]],
            seriestype = :scatter,
            title = "time = $(time[i])",
            label = "",
        )
    end
    gif(anim, file_path, fps = fps)
end

function rotate_pos(pos, angle)
    rotate_matrix = [cos(angle) -sin(angle); sin(angle) cos(angle)]
    rotate_matrix * pos
end

function square(center_x, center_y, shape, angle)
    square_xy = [
        rotate_pos([shape[1], shape[2]] / 2, angle) + [center_x, center_y],
        rotate_pos([-shape[1], shape[2]] / 2, angle) + [center_x, center_y],
        rotate_pos([-shape[1], -shape[2]] / 2, angle) + [center_x, center_y],
        rotate_pos([shape[1], -shape[2]] / 2, angle) + [center_x, center_y],
    ]
    xy = hcat(square_xy...)
    [xy[1, :], xy[2, :]]
end


using QuantumCumulants
using Test
using OrdinaryDiffEq, ModelingToolkit

@testset "higher-order" begin

# Parameters
tspan = range(0.0,10.0,length=101)
ω = range(-2pi,2pi,length=201)

# Define parameters
@cnumbers Δ g γ κ ν

# Define hilbert space
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom,(:g,:e))
h = hf ⊗ ha

# Define the fundamental operators
a = Destroy(h,:a)
s = Transition(h,:σ,:g,:e)

# Hamiltonian
H = Δ*a'*a + g*(a'*s + a*s')

# Collapse operators
J = [a,s,s']
rates = [κ,γ,ν]

# Custom filter function -- include only phase-invaraint terms
ϕ(x) = 0
ϕ(::Destroy) = -1
ϕ(::Create) = 1
function ϕ(t::Transition)
    if (t.i==:e && t.j==:g)
        1
    elseif (t.i==:g && t.j==:e)
        -1
    else
        0
    end
end
ϕ(avg::Average) = ϕ(avg.arguments[1])
function ϕ(t::QuantumCumulants.QMul)
    p = 0
    for arg in t.args_nc
        p += ϕ(arg)
    end
    return p
end
phase_invariant(x) = iszero(ϕ(x))

# Derive equations
he_n = meanfield(a'*a,H,J;rates=rates)

## Fourth order
he4 = complete(he_n;order=4,filter_func=phase_invariant)

# Numerical solution
ps = (Δ, g, γ, κ, ν)
p0 = (0.5, 1.5, 0.25, 1, 4)
sys4 = ODESystem(he4)
u0 = zeros(ComplexF64, length(he4))
prob4 = ODEProblem(sys4,u0,(0.0,tspan[end]),ps .=> p0)
sol4 = solve(prob4,RK4())

n4 = real.(getindex.(sol4.(tspan), 1))

# Correlation function and spectrum
c4 = CorrelationFunction(a', a, he4; steady_state=true, filter_func=phase_invariant)
S4 = Spectrum(c4,ps)
s4 = S4(ω,sol4.u[end],p0)

## Sixth order
he6 = complete(he_n;order=6,filter_func=phase_invariant)

# Numerical solution
ps = (Δ, g, γ, κ, ν)
p0 = (0.5, 1.5, 0.25, 1, 4)
sys6 = ODESystem(he6)
u0 = zeros(ComplexF64, length(he6))
prob6 = ODEProblem(sys6,u0,(0.0,tspan[end]),ps .=> p0)
sol6 = solve(prob6,RK4())

n6 = real.(getindex.(sol6.(tspan), 1))

# Correlation function and spectrum
c6 = CorrelationFunction(a', a, he6; steady_state=true, filter_func=phase_invariant)
S6 = Spectrum(c6,ps)
s6 = S6(ω,sol6.u[end],p0)

## Test
@test maximum(abs.(n6 .- n4)) < 0.05
@test maximum(abs.(s6 .- s4)) < 0.2

end # testset


export ARC_generic_ls

"""
A one dimensionnal ARC algorithm designed for Line Search purposes.

In this generic version we assume a second order polynomial enriched with a cubic
regularization term is used to approximate our h function.
We use c(d) = h(θ) + h'(θ)⋆d + 0.5 * h''(θ) ⋆ d² + Δ/3 |d|³.
Also h''(θ) is sometimes approximated by a secant (Sec) or an improved secant (SecA).
That approximation
is decided with hyper parameter "direction"

Some specific parameters for theses algorithms:
  - Δ :: The cubic regularization parameter
Some parameters are used for TR and ARC methods:
  -eps1 and eps2 :: treshold to determine if q is a bad or a good++ approximation
    of h
  -red and aug :: Controls the speed at which we reduce or augment the size of
    the trust region depending on the previously mentionned treshold.

All the parameters mentionned above can be found in the LS_Function_Meta
"""
function ARC_generic_ls(h         :: LineModel, stop_ls   :: LS_Stopping;
                        f_meta = LS_Function_Meta(),
                        φ_dφ      :: Function = (x, y) -> phi_dphi(x, y),
                        verboseLS :: Bool = false, kwargs...)

    state = stop_ls.current_state

    update!(state, x = 1.0, ht = obj(h, 1.0), gt = grad(h, 1.0), tmps = time())

    # Check if 1.0 is an admissible step size.
    OK = start!(stop_ls)

    # If 1.0 isn't an admissible step size we start our line search algorithm
    iter = 0
    # If 1.0 is a good step size we don't want to compute  unnecessary function
    # evaluations
    H = NaN; φt = NaN; dφt = NaN                   # only for the verbose
    if OK == false
      φt, dφt = φ_dφ(h, state)

      tprec = NaN; φtestTR = NaN; dφtestTR = NaN; φtprec = NaN; dφtprec = NaN;

      # H will denote the approximation to φ'' hereafter
      if f_meta.dir == :Nwt
        H = hess(h, state.x)
      elseif f_meta.dir == :Sec || f_meta.dir == :SecA
        H = 1.0
      end
    end # if OK == false

    verboseLS && @printf("   iter   t       φt        dφt        Δ")
    verboseLS && @printf("   Successful        d        ratio\n")
    verboseLS && @printf("%4d %9.2e  %9.2e  %9.2e %9.2e\n",
                         iter, state.x, φt, dφt, f_meta.Δ);

    while !OK # admissible: respecte armijo et wolfe,
                                # tired: nb d'itérations

        d = ARC_step_computation(H, dφt, f_meta.Δ; kwargs...)

        # We see where the direction d gets us
        candidate_state = copy(state)
        update!(candidate_state, x = state.x + d)
        φtestTR, dφtestTR = φ_dφ(h, candidate_state; kwargs...)

        # depending on the approximation of the second derivate we need different
        # information
        if f_meta.dir == :Sec || f_meta.dir == :SecA
          tprec = state.x
          dφtprec = dφt
        end
        if f_meta.dir == :SecA φtprec = φt end

        # test d'arrêt sur dφ
        (pred, ared, ratio) =
               pred_ared_computation(dφt, φt, H, d, φtestTR, dφtestTR)

        if (state.x + d < 0.0) || ratio < f_meta.eps1  # Unsuccessful
          f_meta.Δ = f_meta.red * f_meta.Δ
          iter += 1
          verboseLS &&
            @printf("%4d %9.2e %9.2e  %9.2e  %9.2e %9.2e %7.2e\n",
                    iter, state.x, φt, dφt, f_meta.Δ, 0, d);
        else             # Successful
            stop_ls.current_state = copy(candidate_state)
            state = stop_ls.current_state
            φt, dφt = φtestTR, dφtestTR
            H = update_H(f_meta.dir, h, state.x, tprec, φt, dφt, φtprec, dφtprec)

          if ratio > f_meta.eps2
            f_meta.Δ = f_meta.aug * f_meta.Δ
          end
          iter += 1
          verboseLS &&
            @printf("%4d %9.2e %9.2e  %9.2e  %9.2e  %9.2e %7.2e \n",
                    iter, state.x, φt, dφt, f_meta.Δ, 1, d);
        end;
        OK = stop!(stop_ls)
    end;

    return (state, stop_ls.meta.optimal)

end


using DrWatson
@quickactivate "PECUZAL_Julia"

using DynamicalSystems
using DelayEmbeddings
using DelimitedFiles
using PyPlot

include("../../src/pecuzal_method.jl")
include("../../src/data_analysis_functions.jl")

## This script was for testing reasons only and had no impact of any results
# shown in the paper

## We analyze the reconstruction from standard time delay embedding, the
# MDOP embedding algorithm and the Garcia & Almeida method and compare it to
# our propsed method: The pecuzal method.
# For evaluation of the different reconstructions we consider Uzal's L-statistic,
# the fraction of recurrence rates from the JRP of the reconstruction and the
# reference and only the reference. As a third criterion we look at the
# recurrence time entropy, since it is related to the Kolmogorov-entropy as a
# dynamical invariant. We also look at two other RQA-measures and, moreover, the
# generalized mutual false nearest neighbors.

# Here we look at a chaotic system: The Roessler attractor in the funnel regime

## Integrate system and determine theiler window
using Random
using DifferentialEquations
Random.seed!(234)
# set time interval for integration
N = 15000 # number of samples
transients = 2000
dt_tra = 0.01
tspan = (dt_tra*N) + (dt_tra*transients)

function roessler!(du,u,p,t)
    x,y,z = u
    a,b,c = p
    du[1] = -(y+z) # dx
    du[2] = x + (a*y) # dy
    du[3] = b+z*(x-c) # dz
end

function prob_func_roe(prob,i,repeat)
  @. prob.u0 = [2*rand(); 2*rand(); 2*rand()]
  prob
end

# parameter values
a = .2925
b = .1
c = 8.5
p = [a, b, c]

u0 = [2*rand(); 2*rand(); 2*rand()]
roe = ODEProblem(roessler!,u0,tspan,p)
tr = solve(roe; saveat = dt_tra)
tr = tr.u

roe_ensemble = EnsembleProblem(roe; prob_func = prob_func_roe)
sim = solve(roe_ensemble; saveat = dt_tra, trajectories = 1000)

pygui(true)

figure()
plot3D(tr[:,1], tr[:,2], tr[:,3])

a = .36
b = .01
c = 14
u0 = [rand(); 5+rand(); 50+rand()]
roe = Systems.roessler(u0 ; a=a, b=b, c=c)

# set time interval for integration
N = 10000 # number of samples
dt_tra = 0.01
t = 0:dt_tra:(dt_tra*N)

tr = trajectory(roe, (N*dt_tra); dt = dt_tra, Ttr = (2000*dt_tra))
tr = regularize(tr)

w1 = estimate_delay(tr[:,1], "mi_min")
w2 = estimate_delay(tr[:,2], "mi_min")

w = maximum(hcat(w1,w2))
taus = 0:100

## Perform reconstructions and compute according L-value

# L-value of reference
L_ref = uzal_cost(tr, Tw = (4*w), w = w, samplesize=1)

#Standard TDE
@time Y_tde, _ = standard_embedding_cao(tr[:,2])
L_tde = uzal_cost(Y_tde, Tw = (4*w), w = w, samplesize=1)

#MDOP embedding
println("Computation time MDOP:")
@time begin
    tw1 = mdop_maximum_delay(tr[:,1])
    tw2 = mdop_maximum_delay(tr[:,2])
    lm1 = DelayEmbeddings.findlocalminima(tw1[2])
    if lm1[1] ≤ 2
        lmm1 = try lm1[2]
        catch
            lm1[1]
        end
    else
        lmm1 = lm1[1]
    end
    lm2 = DelayEmbeddings.findlocalminima(tw2[2])
    if lm2[1] ≤ 2
        lmm2 = try lm2[2]
        catch
            lm2[1]
        end
    else
        lmm2 = lm2[1]
    end
    tau_max = maximum(hcat(lmm1, lmm2))
    taus_mdop = 0:tau_max

    Y_mdop, τ_vals_mdop, ts_vals_mdop, FNNs_mdop , βs_mdop = mdop_embedding(tr[:,1:2];
                                                        τs = taus_mdop , w = w)
    L_mdop = uzal_cost(Y_mdop, Tw = (4*w), w = w, samplesize=1)
end

#Garcia & Almeida
println("Computation time Garcia & Almeida method:")
@time Y_GA, τ_vals_GA, ts_vals_GA, FNNs_GA , ns_GA = garcia_almeida_embedding(tr[:,1:2];
                                                    τs = taus , w = w, T = w)
L_GA = uzal_cost(Y_GA, Tw = (4*w), w = w, samplesize=1)

#Pecuzal
println("Computation time pecuzal method:")
@time Y_pec, τ_vals_pec, ts_vals_pec, Ls_pec , εs_pec = pecuzal_embedding_update(tr[:,1:2];
                                                            τs = taus , w = w)
L_pec = minimum(Ls_pec)

@time Y_pec, τ_vals_pec, ts_vals_pec, Ls_pec , εs_pec = pecuzal_embedding_update(tr[:,1:2];
                                                            τs = taus , w = w)

## compute other evaluation statistics

@time mfnn1, mfnn2, mfnn3, mfnn4, f1, f2, f3, f4, RQA_ref, RQA1, RQA2,
                                            RQA3, RQA4, R_ref, R1, R2, R3, R4 =
        perform_recurrence_analysis(tr, Dataset(Y_tde), Dataset(Y_mdop),
                    Dataset(Y_GA), Dataset(Y_pec); ε = 0.08, w = w, kNN = 10)

NN = 1000
figure()
subplot(231)
RR_ref = grayscale(R_ref[1:NN,1:NN])
imshow(RR_ref, cmap = "binary_r", extent = (1, size(RR_ref)[1], 1, size(RR_ref)[2]))
title("reference")
subplot(232)
RR1 = grayscale(R1[1:NN,1:NN])
imshow(RR1, cmap = "binary_r", extent = (1, size(RR_ref)[1], 1, size(RR_ref)[2]))
title("TDE")
subplot(233)
RR2 = grayscale(R2[1:NN,1:NN])
imshow(RR2, cmap = "binary_r", extent = (1, size(RR_ref)[1], 1, size(RR_ref)[2]))
title("MDOP")
subplot(234)
RR3 = grayscale(R3[1:NN,1:NN])
imshow(RR3, cmap = "binary_r", extent = (1, size(RR_ref)[1], 1, size(RR_ref)[2]))
title("GA")
subplot(235)
RR4 = grayscale(R4[1:NN,1:NN])
imshow(RR4, cmap = "binary_r", extent = (1, size(RR_ref)[1], 1, size(RR_ref)[2]))
title("PECUZAL")

figure()
subplot(131)
RR_ref = grayscale(R_ref[1:NN,1:NN])
imshow(RR_ref, cmap = "binary_r", extent = (1, size(RR_ref)[1], 1, size(RR_ref)[2]))
title("Reference")

subplot(132)
RR4 = grayscale(R4[1:NN,1:NN])
imshow(RR4, cmap = "binary_r", extent = (1, size(RR_ref)[1], 1, size(RR_ref)[2]))
title("Reconstruction")

subplot(133)
JRP = grayscale(elementwise_product(R4[1:NN,1:NN],R_ref[1:NN,1:NN]))
imshow(JRP, cmap = "binary_r", extent = (1, size(RR_ref)[1], 1, size(RR_ref)[2]))
title("Joint Recurrence Plot (JRP)")

figure()
plot3D(tr[:,1], tr[:,2], tr[:,3])
title("reference")
figure()
plot3D(Y_tde[:,1], Y_tde[:,2], Y_tde[:,3])
title("TDE")
figure()
plot3D(Y_GA[:,1], Y_GA[:,2], Y_GA[:,3])
title("GA")
figure()
plot3D(Y_mdop[:,1], Y_mdop[:,2], Y_mdop[:,3])
title("MDOP")
figure()
plot3D(Y_pec[:,1], Y_pec[:,2], Y_pec[:,3])
title("PECUZAL")

display("mfnn_tde: $mfnn1")
display("mfnn_mdop: $mfnn2")
display("mfnn_GA: $mfnn3")
display("mfnn_pec: $mfnn4")

display("*****")

display("f_tde: $f1")
display("f_mdop: $f2")
display("f_GA: $f3")
display("f_pec: $f4")

display("*****")

display("DET_ref: $(abs(RQA_ref.DET - RQA_ref.DET)/RQA_ref.DET)")
display("DET_tde: $(abs(RQA_ref.DET - RQA1.DET)/RQA_ref.DET)")
display("DET_mdop: $(abs(RQA_ref.DET - RQA2.DET)/RQA_ref.DET)")
display("DET_GA: $(abs(RQA_ref.DET - RQA3.DET)/RQA_ref.DET)")
display("DET_pec: $(abs(RQA_ref.DET - RQA4.DET)/RQA_ref.DET)")

display("*****")

display("RTE_ref: $(abs(RQA_ref.RTE - RQA_ref.RTE)/RQA_ref.RTE)")
display("RTE_tde: $(abs(RQA_ref.RTE - RQA1.RTE)/RQA_ref.RTE)")
display("RTE_mdop: $(abs(RQA_ref.RTE - RQA2.RTE)/RQA_ref.RTE)")
display("RTE_GA: $(abs(RQA_ref.RTE - RQA3.RTE)/RQA_ref.RTE)")
display("RTE_pec: $(abs(RQA_ref.RTE - RQA4.RTE)/RQA_ref.RTE)")

display("*****")

display("Lmax_ref: $(abs(RQA_ref.Lmax - RQA_ref.Lmax)/RQA_ref.Lmax)")
display("Lmax_tde: $(abs(RQA_ref.Lmax - RQA1.Lmax)/RQA_ref.Lmax)")
display("Lmax_mdop: $(abs(RQA_ref.Lmax - RQA2.Lmax)/RQA_ref.Lmax)")
display("Lmax_GA: $(abs(RQA_ref.Lmax - RQA3.Lmax)/RQA_ref.Lmax)")
display("Lmax_pec: $(abs(RQA_ref.Lmax - RQA4.Lmax)/RQA_ref.Lmax)")

display("*****")

display("ENTR_ref: $(abs(RQA_ref.ENTR - RQA_ref.ENTR)/RQA_ref.ENTR)")
display("ENTR_tde: $(abs(RQA_ref.ENTR - RQA1.ENTR)/RQA_ref.ENTR))")
display("ENTR_mdop: $(abs(RQA_ref.ENTR - RQA2.ENTR)/RQA_ref.ENTR))")
display("ENTR_GA: $(abs(RQA_ref.ENTR - RQA3.ENTR)/RQA_ref.ENTR))")
display("ENTR_pec: $(abs(RQA_ref.ENTR - RQA4.ENTR)/RQA_ref.ENTR))")

display("*****")

display("L_ref: $(abs(RQA_ref.L - RQA_ref.L)/RQA_ref.L)")
display("L_tde: $(abs(RQA_ref.L - RQA1.L)/RQA_ref.L)")
display("L_mdop: $(abs(RQA_ref.L - RQA2.L)/RQA_ref.L)")
display("L_GA: $(abs(RQA_ref.L - RQA3.L)/RQA_ref.L)")
display("L_pec: $(abs(RQA_ref.L - RQA4.L)/RQA_ref.L)")

display("*****")

display("MRT_ref: $(abs(RQA_ref.MRT - RQA_ref.MRT)/RQA_ref.MRT)")
display("MRT_tde: $(abs(RQA_ref.MRT - RQA1.MRT)RQA_ref.MRT)")
display("MRT_mdop: $(abs(RQA_ref.MRT - RQA2.MRT)RQA_ref.MRT)")
display("MRT_GA: $(abs(RQA_ref.MRT - RQA3.MRT)RQA_ref.MRT)")
display("MRT_pec: $(abs(RQA_ref.MRT - RQA4.MRT)RQA_ref.MRT)")

display("*****")

display("LAM_ref: $(abs(RQA_ref.LAM - RQA_ref.LAM)/RQA_ref.LAM)")
display("LAM_tde: $(abs(RQA_ref.LAM - RQA1.LAM)RQA_ref.LAM)")
display("LAM_mdop: $(abs(RQA_ref.LAM - RQA2.LAM)RQA_ref.LAM)")
display("LAM_GA: $(abs(RQA_ref.LAM - RQA3.LAM)RQA_ref.LAM)")
display("LAM_pec: $(abs(RQA_ref.LAM - RQA4.LAM)RQA_ref.LAM)")

display("*****")

display("RTE-MRT_ref: $(abs(RQA_ref.RTE/RQA_ref.MRT - RQA_ref.RTE/RQA_ref.MRT)/(RQA_ref.RTE/RQA_ref.MRT))")
display("RTE-MRT_tde: $(abs(RQA_ref.RTE/RQA_ref.MRT - RQA1.RTE/RQA1.MRT)/(RQA_ref.RTE/RQA_ref.MRT))")
display("RTE-MRT_mdop: $(abs(RQA_ref.RTE/RQA_ref.MRT - RQA2.RTE/RQA2.MRT)/(RQA_ref.RTE/RQA_ref.MRT))")
display("RTE-MRT_GA: $(abs(RQA_ref.RTE/RQA_ref.MRT - RQA3.RTE/RQA3.MRT)/(RQA_ref.RTE/RQA_ref.MRT))")
display("RTE-MRT_pec: $(abs(RQA_ref.RTE/RQA_ref.MRT - RQA4.RTE/RQA4.MRT)/(RQA_ref.RTE/RQA_ref.MRT))")

display("*****")


"""
Module UniformHexMesh

Includes mesh related utility functions
"""

module UniformHexMesh

using SparseArrays
using CommonUtils # for meshgrid

export uniform_hex_mesh
export hex_face_vertices

"""
uniform_hex_mesh(Kx::Int,Ky::Int)

Matlab uniform hexahedral mesh.

# Examples
```jldoctest
```
"""

function uniform_hex_mesh(Nx,Ny,Nz)
        Nxp = Nx+1
        Nyp = Ny+1
        Nzp = Nz+1
        Nv = convert(Int,Nxp*Nyp*Nzp)
        K  = convert(Int,Nx*Ny*Nz)

        x1D = LinRange(-1,1,Nxp)
        y1D = LinRange(-1,1,Nyp)
        z1D = LinRange(-1,1,Nzp)

        # hack to work with old vertex ordering
        function mymeshgrid(x1D,y1D,z1D)
                Np = length(x1D)*length(y1D)*length(z1D)
                x = zeros(Np)
                y = zeros(Np)
                z = zeros(Np)
                sk = 1
                for k = 1:length(z1D)
                        for j = 1:length(y1D)
                                for i = 1:length(x1D)
                                        x[sk] = x1D[i]
                                        y[sk] = y1D[j]
                                        z[sk] = z1D[k]
                                        sk += 1
                                end
                        end
                end
                return x,y,z
        end
        x, y, z = mymeshgrid(x1D,y1D,z1D)

        EToV = zeros(Int,K,8)
        k = 1
        for e = 1:K
                em = e-1
                k = div(em,(Nx*Ny))
                j = div(em - k*Nx*Ny,Nx)
                i = em % Nx

                EToV[e,1] = i     + Nxp*j     + Nxp*Nyp*k
                EToV[e,2] = (i+1) + Nxp*j     + Nxp*Nyp*k
                EToV[e,3] = i     + Nxp*(j+1) + Nxp*Nyp*k
                EToV[e,4] = (i+1) + Nxp*(j+1) + Nxp*Nyp*k
                EToV[e,5] = i     + Nxp*j     + Nxp*Nyp*(k+1)
                EToV[e,6] = (i+1) + Nxp*j     + Nxp*Nyp*(k+1)
                EToV[e,7] = i     + Nxp*(j+1) + Nxp*Nyp*(k+1)
                EToV[e,8] = (i+1) + Nxp*(j+1) + Nxp*Nyp*(k+1)
        end
        EToV = @. EToV + 1 # re-index to 1 index

        VX = x[:];
        VY = y[:];
        VZ = z[:];
        return VX[:],VY[:],VZ[:],EToV
end

function uniform_hex_mesh(Kx)
        return uniform_hex_mesh(Kx,Kx,Kx)
end

function hex_face_vertices()
        x1D = LinRange(-1,1,2)
        r, s, t = meshgrid(x1D,x1D,x1D)
        fv1 = map(x->x[1], findall(@. abs(r+1) < 1e-10))
        fv2 = map(x->x[1], findall(@. abs(r-1) < 1e-10))
        fv3 = map(x->x[1], findall(@. abs(s+1) < 1e-10))
        fv4 = map(x->x[1], findall(@. abs(s-1) < 1e-10))
        fv5 = map(x->x[1], findall(@. abs(t+1) < 1e-10))
        fv6 = map(x->x[1], findall(@. abs(t-1) < 1e-10))
        return fv1,fv2,fv3,fv4,fv5,fv6
end

end


#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

using FMI
using Flux
using DifferentialEquations: Tsit5

import Random
Random.seed!(1234);

FMUPath = joinpath(dirname(@__FILE__), "..", "model", "SpringPendulumExtForce1D.fmu")

t_start = 0.0
t_step = 0.01
t_stop = 1.0 # 5.0
tData = t_start:t_step:t_stop

# generate traing data
myFMU = fmiLoad(FMUPath)
fmiInstantiate!(myFMU; loggingOn=false)
fmiSetupExperiment(myFMU, t_start, t_stop)
fmiSetReal(myFMU, "mass_s0", 1.3)
fmiEnterInitializationMode(myFMU)
fmiExitInitializationMode(myFMU)
success, realSimData = fmiSimulateCS(myFMU, t_start, t_stop; recordValues=["mass.a"], setup=false, reset=false, saveat=tData)

# sine(t) as external force
function extForce(t)
    return [sin(t)]
end
accData = collect(data[1] for data in realSimData.saveval)

# loss function for training
function losssum()
    solution = problem(extForce, t_step)

    accNet = collect(data[1] for data in solution)

    Flux.Losses.mse(accNet, accData)
end

# callback function for training
global iterCB = 0
global lastLoss = 0.0
function callb()
    global iterCB += 1
    global lastLoss

    if iterCB == 1
        lastLoss = losssum()
    end

    if iterCB % 50 == 0
        loss = losssum()
        @info "Loss: $loss"
        @test loss < lastLoss
        lastLoss = loss
    end
end

# NeuralFMU setup
numInputs = length(myFMU.modelDescription.inputValueReferences)
numOutputs = length(myFMU.modelDescription.outputValueReferences)

net = Chain(inputs -> fmiInputDoStepCSOutput(myFMU, t_step, inputs),
            Dense(numOutputs, 16, tanh),
            Dense(16, 16, tanh),
            Dense(16, numOutputs))

problem = CS_NeuralFMU(myFMU, net, (t_start, t_stop); saveat=tData)
@test problem != nothing

solutionBefore = problem(extForce, t_step)
vals = collect(data[1] for data in solutionBefore)

# train it ...
p_net = Flux.params(problem)

optim = ADAM()
Flux.train!(losssum, p_net, Iterators.repeated((), 100), optim; cb=callb)

# check results
solutionAfter = problem(extForce, t_step)
vals = collect(data[1] for data in solutionAfter)

fmiUnload(myFMU)


import PotentialUQ
using LinearAlgebra: svd, svdvals, cond, pinv, I, diagm, norm
using Statistics
using Distributions
using TransformVariables
using Random
using Turing

num_configs = 40
r = Vector{PotentialUQ.Potentials.Configuration}(undef, num_configs)
file_path = "../examples/GaN/DATA"
for j = 1:num_configs
    c = PotentialUQ.Potentials.load_lammps(file_path * "/" * string(j) * "/DATA"; atom_names = [:Ga, :N],
                    radii = [0.5, 0.5], weights = [1.0, 0.5], boundary_type = ["p", "p", "p"], units = "metal")
    r[j] = c
end

num_configs_test = 20
rtest = Vector{PotentialUQ.Potentials.Configuration}(undef, num_configs_test)
file_path = "../examples/GaN/DATA"
for j = 1:num_configs_test
    c = PotentialUQ.Potentials.load_lammps(file_path * "/" * string(40+j) * "/DATA"; atom_names = [:Ga, :N],
            radii = [0.5, 0.5], weights = [1.0, 0.5], boundary_type = ["p", "p", "p"], units = "metal")
    rtest[j] = c
end

# Set up true potential
ϵ_Ga_Ga = 0.643
σ_Ga_Ga = 2.390
ϵ_N_N   = 1.474
σ_N_N   = 1.981
A_Ga_N  = 608.54
ρ_Ga_N  = 0.435
q_Ga    = 3.0
q_N     = -3.0
ϵ0      = 55.26349406

lj_Ga_Ga = PotentialUQ.Potentials.LennardJones(ϵ_Ga_Ga,σ_Ga_Ga)
lj_N_N   = PotentialUQ.Potentials.LennardJones(ϵ_N_N, σ_N_N)
bm_Ga_N  = PotentialUQ.Potentials.BornMayer(A_Ga_N,ρ_Ga_N)
c_Ga_N   = PotentialUQ.Potentials.Coulomb(q_Ga, q_N, ϵ0)
gan = PotentialUQ.Potentials.GaN(lj_Ga_Ga, lj_N_N, bm_Ga_N, c_Ga_N)
gantest = PotentialUQ.Potentials.GaN(lj_Ga_Ga, lj_N_N, bm_Ga_N, c_Ga_N)

println(gan)
pe = PotentialUQ.Potentials.potential_energy(r, gan)
ptest = PotentialUQ.Potentials.potential_energy(rtest, gantest)
println("GaN Energy ", pe)
fo = PotentialUQ.Potentials.force(r, gan)
fotest = PotentialUQ.Potentials.force(rtest, gantest)
println("GaN Forces ")
show(stdout, "text/plain", fo[1][1:3])
println(" ")
v = PotentialUQ.Potentials.virial(r, gan)
vtest = PotentialUQ.Potentials.virial(rtest, gantest)

println("GaN virial ", v)
v_tensor = PotentialUQ.Potentials.virial_stress(r, gan)
println("GaN Virial Tensor ", v_tensor)
println(" ")
# Form right hand side
b = vcat(pe, reduce(vcat, reduce(vcat, fo)), reduce(vcat, v_tensor))
# b = [pe; vec(v_tensor)]
println("Right hand side shape: ", size(b))
show(stdout, "text/plain", b[1:10])
println(" ")

# Form SNAP A matrix
twojmax = 6
rcutfac = 4.0
snap = PotentialUQ.Potentials.SNAP(rcutfac, twojmax, r[1])
A = PotentialUQ.Potentials.get_snap(r, snap)
println("Shape of A: ", size(A))
show(stdout, "text/plain", A[1:20, 1:9])
println(" ")

# Fit
snap.β = A \ b
println("Fitted β = ")
show(stdout, "text/plain", snap.β)
println(" ")
# Test
pe_snap = PotentialUQ.Potentials.potential_energy(rtest, snap)
println("SNAP potential_energy = ", pe_snap)
println("Relative Error = ", mean(abs.(ptest - pe_snap)./abs.(ptest)))
f_snap = PotentialUQ.Potentials.force(rtest, snap)
show(stdout, "text/plain", vcat(fotest...)[1:5])
println(" ")
show(stdout, "text/plain", vcat(f_snap...)[1:5])
println(" ")
ff = vcat( (fotest - f_snap)... )
fff = vcat( fotest... )
println("Force Relative Error = ", mean( norm.(ff)./norm.(fff)  ) )
v_snap = PotentialUQ.Potentials.virial(rtest, snap)
println("SNAP Virial = ", v_snap)
println("Virial Error = ", mean(abs.(vtest - v_snap)./abs.(vtest)))

# Regularize
println("Regularizing")
AA = A' * A
println(" ")
snap.β = pinv(AA, 1e-6) * (A'*b)
println("Regularized Fitted β = ")
show(stdout, "text/plain", snap.β)
println(" ")
pe_snap = PotentialUQ.Potentials.potential_energy(rtest, snap)
println("Regularized SNAP potential_energy = ", pe_snap)
println("Regularized Relative Error = ", mean(abs.(ptest - pe_snap)./abs.(ptest)))
f_snap = PotentialUQ.Potentials.force(rtest, snap)
ff = vcat( (fotest - f_snap)... )
fff = vcat( fotest... )
println("Force Relative Error = ", mean( norm.(ff)./norm.(fff)  ) )
v_snap = PotentialUQ.Potentials.virial(rtest, snap)
println("Regularized SNAP Virial = ", v_snap)
println("Regularized Virial Error = ", mean(abs.(vtest - v_snap)./abs.(vtest)))

############### UQ ######################
# Unknowns
n, m = size(A)

β = pinv(A' * A, 1e-6)*(A'*b)
nt = (β = β, )
# println("v = ", nt)
t = PotentialUQ.namedtp_to_vec(nt)
x = TransformVariables.inverse(t, nt)

# # Prior
F = svd(A)
Σ = diagm(F.S)
scale_A = vcat(200.0 .+ 0.0 .* pe, 1.0 .+ 0.0 .*reduce(vcat, reduce(vcat,fo)), 1.0 .+ 0.0 .*reduce(vcat, v_tensor))
Q =  F.U' * ( diagm(scale_A) ) * (F.U)
Q = 1e-2*I(m) + 0.5*Q + 0.5*Q'
println("Q = ")
show(stdout, "text/plain", Q)
println(" ")
println(" y = ")
y =  F.U' * b
show(stdout, "text/plain", y)
println(" ")
sampler(rng::Random.AbstractRNG, d::ContinuousMultivariateDistribution) = β + 1e-1*randn(length(d))
logpdf(x::Vector) = loglikelihood(MvNormal(β, 1e-1*I(m)), x)
snap_prior = PotentialUQ.SNAP_Prior_Distribution(A, b, x, t, sampler, logpdf)
println("logpdf = ", Distributions.logpdf(snap_prior, x))
snap_likelihood(x::Vector) = loglikelihood( MvNormal( Σ * F.Vt*x, Q ),  y )
println("Likelihood = ", snap_likelihood(x))
dsnap = PotentialUQ.PotentialDistribution(snap, nt, snap_prior, snap_likelihood)

@time samp = PotentialUQ.Sample(dsnap; num_samples = 1000, num_adapts = 500, verbose = true)

snap.β = dsnap.x.β
println("MAP Fitted β = ")
show(stdout, "text/plain", snap.β)
println(" ")
pe_snap = PotentialUQ.Potentials.potential_energy(rtest, snap)
println("MAP SNAP potential_energy = ", pe_snap)
println("MAP Relative Error = ", mean(abs.(ptest - pe_snap)./abs.(ptest)))
f_snap = PotentialUQ.Potentials.force(rtest, snap)
ff = vcat( (fotest - f_snap)... )
fff = vcat( fotest... )
println("Force Relative Error = ", mean( norm.(ff)./norm.(fff)  ) )
v_snap = PotentialUQ.Potentials.virial(rtest, snap)
println("MAP SNAP Virial = ", v_snap)
println("MAP Virial Error = ", mean(abs.(vtest - v_snap)./abs.(vtest)))

# using Plots

# plot(pinv(AA, 1e-6) * (A'*b))
# plot!(vec(m), ribbon = vec(s), fillalpha = 0.5)
# savefig("snap_map_samples.png")


function isvalid(sim::Chains)
    if !all(gelmandiag(sim).value[:,1,1] .< 1.01)
        return false
    end
    for i in 1:size(sim.value)[2]
        data = mean(sim.value[:,i,:], dims = [1])[:]
        # mu = mean(data)
        sig = std(data)
        if sig > 0.05
            return false
        end
    end
    return true
end

function getstats(value::AbstractArray;standardize::Bool = false)
    # get statistics or natural parameter of given MCMC chains
    # in this function, return the first four moments
    m1 = mean(value, dims = [1,3])[:]
    m2 = var(value, dims = [1,3])[:]
    if standardize
        m3 = map(1:size(value)[2]) do i
            return skewness(value[:,i,:][:])
        end
        m3 = m3[:]
        m4 = map(1:size(value)[2]) do i
            return kurtosis(value[:,i,:][:])
        end
        m4 = m4[:]
    else
        m3 = mean(value.^3, dims = [1,3])[:]
        m4 = mean(value.^4, dims = [1,3])[:]
    end
    res = [m1 m2 m3 m4]
    @assert size(res) == (size(value)[2], 4) "size incorrect"
    return res
end


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

# TYPE

mutable struct OLS <: MLE

    sample::Microdata
    β::Vector{Float64}
    V::Matrix{Float64}

    OLS() = new()
end

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

# CONSTRUCTOR

function OLS(MD::Microdata)
    obj        = OLS()
    obj.sample = MD
    return obj
end

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

# ESTIMATION

function _fit!(obj::OLS, ::UnitWeights)
    y     = getvector(obj, :response)
    x     = getmatrix(obj, :control)
    obj.β = x \ y
end

function _fit!(obj::OLS, w::AbstractWeights)
    y     = getvector(obj, :response)
    x     = getmatrix(obj, :control)
    v     = Diagonal(w) * x
    obj.β =  (v' * x) \ (v' * y)
end

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

# SCORE (MOMENT CONDITIONS)

score(obj::OLS) = Diagonal(residuals(obj)) * getmatrix(obj, :control)

# EXPECTED JACOBIAN OF SCORE × NUMBER OF OBSERVATIONS

jacobian(obj::OLS, ::UnitWeights)      = - crossprod(getmatrix(obj, :control))
jacobian(obj::OLS, w::AbstractWeights) = - crossprod(getmatrix(obj, :control), w)

# HOMOSCEDASTIC VARIANCE MATRIX

function _vcov!(obj::OLS, corr::Homoscedastic, w::UnitWeights)
    σ²    = sum(abs2, residuals(obj)) / dof_residual(obj)
    obj.V = lmul!(-σ², inv(jacobian(obj, w)))
end

function _vcov!(obj::OLS, corr::Homoscedastic, w::AbstractWeights)
    σ²    = sum(abs2.(residuals(obj)), w) / dof_residual(obj)
    obj.V = lmul!(-σ², inv(jacobian(obj, w)))
end

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

# LINEAR PREDICTOR

function predict(obj::OLS, MD::Microdata)
    (getnames(obj, :control) != getnames(MD, :control)) && throw("missing variables")
    getmatrix(MD, :control) * obj.β
end

# FITTED VALUES

fitted(obj::OLS, MD::Microdata) = predict(obj, MD)

# DERIVATIVE OF FITTED VALUES

jacobexp(obj::OLS) = copy(getmatrix(obj, :control))

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

# UTILITIES

coefnames(obj::OLS) = getnames(obj, :control)
mtitle(obj::OLS)    = "OLS"

# COEFFICIENT OF DETERMINATION

adjr2(obj::OLS) = 1.0 - (1.0 - r2(obj)) * (nobs(obj) - 1) / dof_residual(obj)
r2(obj::OLS)    = _r2(obj, getweights(obj))

function _r2(obj::OLS, ::UnitWeights)
    y   = response(obj)
    ŷ   = fitted(obj)
    rss = sum(abs2, y .- ŷ)
    tss = sum(abs2, y .- mean(y))
    return 1.0 - rss / tss
end

function _r2(obj::OLS, w::AbstractWeights)
    y   = response(obj)
    μ   = mean(y, w)
    ŷ   = fitted(obj)
    rss = sum(abs2.(y .- ŷ), w)
    tss = sum(abs2.(y .- μ), w)
    return 1.0 - rss / tss
end

# LIKELIHOOD FUNCTION

function _loglikelihood(obj::OLS, w::AbstractWeights)
    n  = sum(w)
    r  = residuals(obj)
    r .= abs2.(r)
    σ² = sum(r, w) / n
    return - (log(σ²) + log2π + 1.0) * n / 2.0
end

# LIKELIHOOD FUNCTION UNDER NULL MODEL

function _nullloglikelihood(obj::OLS, w::AbstractWeights)
    n  = sum(w)
    r  = copy(getvector(obj, :response))
    μ  = mean(r, w)
    r .= abs2.(r .- μ)
    σ² = sum(r, w) / n
    return - (log(σ²) + log2π + 1.0) * n / 2.0
end

# DEVIANCE

function _deviance(obj::OLS, w::AbstractWeights)
    r  = residuals(obj)
    r .= abs2.(r)
    return sum(r, w)
end

# DEVIANCE UNDER NULL MODEL

function _nulldeviance(obj::OLS, w::AbstractWeights)
    r  = copy(getvector(obj, :response))
    μ  = mean(r, w)
    r .= abs2.(r .- μ)
    return sum(r)
end


__precompile__
import GSL

struct Basis
    # number of moments
    N::Int64;
    Nq::Int64;
    states::Int64;

    # precomputed Legendre Polynomials at quadrature points
    PhiQuad::Array{Float64,2};
    PhiQuadCons::Array{Float64,2};
    PhiQuadFull::Array{Float64,2};
    PhiQuadW::Array{Float64,2};
    PhiQuadWCons::Array{Float64,2};
    PhiQuadWFull::Array{Float64,2};

    function Basis(quadrature::Quadrature,settings::Settings)
        N = settings.N;
        Nq = quadrature.Nq;
        NCons = settings.NCons;

        # precompute Legendre basis functions at quad points
        PhiQuad = zeros(quadrature.Nq^2,N^2);
        PhiQuadFull = zeros(quadrature.Nq^2,N^2);
        PhiQuadW = zeros(N^2,quadrature.Nq^2);
        PhiQuadWFull = zeros(N^2,quadrature.Nq^2);
        for i = 1:N
            for j = 1:N
                for k = 1:Nq
                    for q = 1:Nq
                        PhiQuad[(q-1)*Nq+k,(j-1)*N+i] = Phi.(i-1,quadrature.xi[k])*Phi.(j-1,quadrature.xi[q]);
                        PhiQuadW[(j-1)*N+i,(q-1)*Nq+k] = Phi.(i-1,quadrature.xi[k])*Phi.(j-1,quadrature.xi[q])*quadrature.w[k]*quadrature.w[q];
                    end
                end
            end
        end
        PhiQuadFull .= PhiQuad;
        PhiQuadWFull .= PhiQuadW;

        # save conserved quantities
        PhiQuadCons = zeros(quadrature.Nq^2,NCons);
        PhiQuadWCons = zeros(NCons,quadrature.Nq^2);
        for i = 1:NCons
            for k = 1:Nq
                for q = 1:Nq
                    i1 = settings.iCons[i,1];
                    i2 = settings.iCons[i,2];
                    # set conserved quantities to zero
                    PhiQuad[(q-1)*Nq+k,(i2-1)*N+i1] = 0.0;
                    PhiQuadW[(i2-1)*N+i1,(q-1)*Nq+k] = 0.0;
                    # save conserved basis
                    PhiQuadCons[(q-1)*Nq+k,i] = Phi.(i1-1,quadrature.xi[k])*Phi.(i2-1,quadrature.xi[q]);
                    PhiQuadWCons[i,(q-1)*Nq+k] = PhiQuadCons[(q-1)*Nq+k,i]*quadrature.w[k]*quadrature.w[q];
                end
            end
        end

        new(N,Nq,1,PhiQuad,PhiQuadCons,PhiQuadFull,PhiQuadW,PhiQuadWCons,PhiQuadWFull);
    end
end

# Legendre Polynomials on [-1,1]
function Phi(n::Int64,xi::Float64) ## can I use this with input vectors xVal?
    return sqrt(2.0*n+1.0)*GSL.sf_legendre_Pl.(n,xi);
end

# evaluate polynomial with moments u at spatial position x
function Eval(obj::Basis,u::Array{Float64,1},xi)
    y=zeros(length(xi))
    for i = 1:obj.N
        y = y+u[i]*Phi.(i-1,xi);
    end
    return y;
end

# evalueates states at quadrature points. returns solution[quadpoints,states]
function Eval(obj::Basis,u::Array{Float64,2},xi)
    y=zeros(length(xi),obj.states)
    for s = 1:obj.states
        for i = 1:obj.N
            y[:,s] = y[:,s]+u[i,s]*Phi.(i-1,xi);
        end
    end
    return y;
end

function EvalAtQuad(obj::Basis,u::Array{Float64,1})
    return obj.PhiQuadFull*u;
end

function EvalAtQuadCons(obj::Basis,u::Array{Float64,2})
    return obj.PhiQuadCons*u;
end

function EvalAtQuadDLR(obj::Basis,u::Array{Float64,2})
    return obj.PhiQuad*u;
end

# evalueates states at quadrature points. returns solution[quadpoints,states]
function EvalAtQuad(obj::Basis,u::Array{Float64,2})
    return obj.PhiQuadFull*u;
end

# returns N moments of a function evaluated at the quadrature points and stored in uQ
function ComputeMomentsDLR(obj::Basis,uQ::Array{Float64,1})
    return obj.PhiQuadW*uQ;
end

# returns N moments of a function evaluated at the quadrature points and stored in uQ
function ComputeMoments(obj::Basis,uQ::Array{Float64,1})
    return obj.PhiQuadWFull*uQ;
end

# returns N moments of a function evaluated at the quadrature points and stored in uQ
function ComputeMomentsCons(obj::Basis,uQ::Array{Float64,1})
    return obj.PhiQuadWCons*uQ;
end

# returns N moments of all states evaluated at the quadrature points and stored in uQ
# uQ in [Nq,states], returns [N,states]
function ComputeMoments(obj::Basis,uQ::Array{Float64,2})
    return obj.PhiQuadWFull*uQ;
end

function ComputeMomentsDLR(obj::Basis,uQ::Array{Float64,2})
    println(size(obj.PhiQuadWFull*uQ))
    return obj.PhiQuadW*uQ;
end

# evaluate polynomial with moments u at spatial position x
function Eval(obj::Basis,u::Array{Float64,1},xi,eta)
    y=zeros(length(xi))
    for i = 1:obj.N
        for j = 1:obj.N
            y = y.+u[(j-1)*obj.N+i]*Phi.(i-1,xi).*Phi.(j-1,eta);
        end
    end
    return y;
end


#! /usr/bin/julia

# Rosetta Code, Bitmap/Read a PPM file

using Color, Images, FixedPointNumbers

const M_RGB_Y = reshape(Color.M_RGB_XYZ[2,:], 3)

function rgb2gray(img::Image)
    g = red(img)*M_RGB_Y[1] + green(img)*M_RGB_Y[2] + blue(img)*M_RGB_Y[3]
    g = clamp(g, 0.0, 1.0)
    return grayim(g)
end

ima = imread("bitmap_read_ppm_in.ppm")
imb = convert(Image{Gray{Ufixed8}}, ima)
imwrite(imb, "bitmap_read_ppm_out.png")


using LinearAlgebra

export Scale, factor

"""
    Scale{S <: Union{Number, Val}, N, BT <: AbstractBlock{N}} <: TagBlock{BT, N}

`Scale` a block with scalar. it can be either a `Number` or a compile time `Val`.

# Example

```jldoctest
julia> 2 * X
[scale: 2] X gate

julia> im * Z
[+im] Z gate

julia> -im * Z
[-im] Z gate

julia> -Z
[-] Z gate
```
"""
struct Scale{S <: Union{Number, Val}, N, BT <: AbstractBlock{N}} <: TagBlock{BT, N}
    alpha::S
    content::BT
end

content(x::Scale) = x.content
factor(x::Scale{<:Number}) = x.alpha
factor(x::Scale{Val{X}}) where X = X

Base.copy(x::Scale) = Scale(x.alpha, copy(x.content))
Base.adjoint(x::Scale{<:Number}) = Scale(adjoint(x.alpha), adjoint(content(x)))
Base.adjoint(x::Scale{Val{X}}) where X = Scale(Val(adjoint(X)), adjoint(content(x)))

Base.:(==)(x::Scale, y::Scale) = (factor(x) == factor(y)) && (content(x) == content(y))

chsubblocks(x::Scale, blk::AbstractBlock) = Scale(x.alpha, blk)
cache_key(x::Scale) = (factor(x), cache_key(content(x)))

mat(::Type{T}, x::Scale) where T = T(x.alpha) * mat(T, content(x))
mat(::Type{T}, x::Scale{Val{S}}) where {T, S} = T(S) * mat(T, content(x))

function apply!(r::ArrayReg{B}, x::Scale{S, N}) where {S, B, N}
    apply!(r, content(x))
    r.state .*= factor(x)
    return r
end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 33bd93ec-77d9-11eb-2da8-87271f8b30c7
begin
	using Pkg
	Pkg.activate("MLJ_env", shared=true)

	# using MLJ
	using CSV
	using DataFrames
	using PlutoUI
	using Test
	using Printf
	# using Plots
	using LinearAlgebra
end

# ╔═╡ 6d69b148-77d9-11eb-04f5-034603897c31
begin
	using JSON

	important_words = open("../../ML_UW_Spec/C03/data/important_words.json", "r") do f
    	JSON.parse(f)
	end;

	length(important_words)
end

# ╔═╡ 0f6bb770-77df-11eb-0f1f-4dddbd7e2498
include("./C03w01_text_proc.jl");

# ╔═╡ baa8ca1c-77d8-11eb-2e26-3b2445acbfa1
md"""
## Implementing logistic regression from scratch

The goal of this notebook is to implement your own logistic regression classifier. We will:

  - Extract features from Amazon product reviews.
  - Convert an DataFrame into a Julia array.
  - Implement the link function for logistic regression.
  - Write a function to compute the derivative of the log likelihood function with respect to a single coefficient.
  - Implement gradient ascent.
  - Given a set of coefficients, predict sentiments.
  - Compute classification accuracy for the logistic regression model.

"""

# ╔═╡ 3ad4770e-77d9-11eb-3a93-addf3e427025
md"""
### Load review dataset

For this assignment, we will use a subset of the Amazon product review dataset. The subset was chosen to contain similar numbers of positive and negative reviews, as the original dataset consisted primarily of positive reviews.
"""

# ╔═╡ 6e264178-77d9-11eb-33c1-19462dde66e2
begin
	products = train = CSV.File("../../ML_UW_Spec/C03/data/amazon_baby_subset.csv";
		header=true) |> DataFrame;

	first(products, 3)
end

# ╔═╡ 75f2e820-77de-11eb-19eb-a5a5977b5257
size(products)

# ╔═╡ 36f2c616-77de-11eb-03b4-3b92a532949a
## how many missing values for the column :review?
sum(ismissing.(products.review))

# ╔═╡ 6d4ef592-77d9-11eb-3844-1bea7dc71094
## in-place replace(ment) of missing values for review  with empty string
replace!(products.review, missing =>"");

# ╔═╡ d7b9c812-77de-11eb-0ebb-6bb75de9725b
## no more missing values for column :review
@test sum(ismissing.(products.review)) == 0

# ╔═╡ 6e09ca34-77d9-11eb-1a99-6578f9141004
md"""
One column of this dataset is 'sentiment', corresponding to the class label with +1 indicating a review with positive sentiment and -1 indicating one with negative sentiment.
"""

# ╔═╡ 6ded8978-77d9-11eb-3d94-e7507a85197e
products.sentiment[1:20]

# ╔═╡ 6dd12c56-77d9-11eb-0350-1dfc08d17dd6
md"""
Let us quickly explore more of this dataset. The 'name' column indicates the name of the product. Here we list the first 10 products in the dataset. We then count the number of positive and negative reviews.
"""

# ╔═╡ 6db87918-77d9-11eb-2082-edecdf8b7576
first(products, 10).name

# ╔═╡ 6d9e0568-77d9-11eb-1369-bb6990e1e4a7
# num. of positive reviews, num. of negative reviews
(num_pos=sum(products.sentiment .== 1), num_neg=sum(products.sentiment .== -1))

# ╔═╡ 6d316f86-77d9-11eb-107b-1bf5708655d5
md"""
Now, we will perform 2 simple data transformations:

  1. Remove punctuation using Python's built-in string functionality.
  2. Compute word counts (only for important_words)

"""

# ╔═╡ 734ed4d4-77df-11eb-3dab-e5c10b4cda62
md"""
###### Step 1.
"""

# ╔═╡ 0f528d0e-77df-11eb-1415-2fb0e9d6792d
# products[!, :review_clean] = remove_punctuation.(products[:, :review]);
products[!, :review] = remove_punctuation.(products[:, :review]);

# ╔═╡ 0f362fec-77df-11eb-02a6-210a58c2839b
names(products)

# ╔═╡ 0f1c6b16-77df-11eb-2197-93c4a82cc135
md"""
###### Step 2.

For each word in `important_words`, we compute a count for the number of times the word occurs in the review. We will store this count in a separate column (one for each word). The result of this feature processing is a single column for each word in `important_words` which keeps a count of the number of times the respective word occurs in the review text.

"""

# ╔═╡ 0057a50e-77e0-11eb-29d6-85c2dac67743
function wc(text::String, word::String)
 	# length(filter(w -> w == word, split(text)))
	length(findall(word, text))
end

# ╔═╡ 51a94e8e-77e3-11eb-04c8-d1389b1f5856
function addcols!(df, words)
	for word ∈ words
		df[!, word] = wc.(df.review_clean, word)
	end

	#for word ∈ words
	#	transform!(df, :review_clean => r -> wc.(r, word) => word)
	#end
end

# ╔═╡ 0f03e474-77df-11eb-2791-f165fd011737
@time addcols!(products, important_words)
# 135.474563 seconds (2.43 G allocations: 181.104 GiB, 9.77% gc time)
# transform!  5.6s
# df[!, wrod] 5.5s

# ╔═╡ 0ee68550-77df-11eb-0a6e-c5bf5e3fd9ac
length(names(products))

# ╔═╡ 0ecc0054-77df-11eb-21a4-3f677548ecf9
products[!, :perfect], length(products[!, :perfect])

# ╔═╡ 0eb50d9a-77df-11eb-0b7d-557ab3ea2c81
md"""
Now, write some code to compute the number of product reviews that contain the word perfect.
"""

# ╔═╡ 0e9884fe-77df-11eb-3199-13ab2fbdcbc8
sum(products.perfect .>= 1)

# ╔═╡ 0e815eaa-77df-11eb-3422-515926dab50d
md"""
**Quiz Question**. How many reviews contain the word perfect?
  - cf. above cell.
"""

# ╔═╡ 0e4f6222-77df-11eb-3c39-239382844c12
md"""
#### Using Julia array
"""

# ╔═╡ 0e383a04-77df-11eb-2ec9-c5aeffc4b131
function get_data(df::DF, features, output) where {DF <: Union{DataFrame, SubDataFrame}}
	s_features = [Symbol(f) for f ∈ features]
	df[:, :intercept] .= 1.0
	s_features = [:intercept, s_features...]
	X_matrix = convert(Matrix, select(df, s_features))
	y = df[!, output]

	(X_matrix, y)
end

# ╔═╡ 0e1d5f86-77df-11eb-12c9-75b0b884ad25
@time feature_matrix, sentiment = get_data(products, important_words, :sentiment);

# ╔═╡ 0e03942c-77df-11eb-1bde-4f7b6739a3ee
size(feature_matrix)

# ╔═╡ 0deb769c-77df-11eb-2b53-7b73905353cd
md"""
**Quiz Question**: How many features are there in the feature_matrix?
"""

# ╔═╡ 0dd0bc6c-77df-11eb-2df1-09ca3d5aaafc
size(feature_matrix)[2]

# ╔═╡ fd5a8f7c-77e6-11eb-34dd-11ba6b2d7409
md"""
**Quiz Question**: Assuming that the intercept is present, how does the number of features in feature_matrix relate to the number of features in the logistic regression model?
Let x = [number of features in feature_matrix] and y = [number of features in logistic regression model].

  - [ ] y = x - 1
  - [x] y = x
  - [ ] y = x + 1
  - [ ] None of the above

"""

# ╔═╡ fd419c9c-77e6-11eb-2925-0d832a3fe2ca
sentiment

# ╔═╡ fd23d180-77e6-11eb-2091-af1e84d8fbcf
md"""
####  Estimating conditional probability with link function

Recall from lecture that the link function is given by:

$$P(y_i = +1 | \mathbf{x}_i,\mathbf{w}) = \frac{1}{1 + \exp(-\mathbf{w}^T h(\mathbf{x}_i))},$$

where the feature vector $h(\mathbf{x}_i)$ represents the word counts of `important_words` in the review  $\mathbf{x}_i$.

Implement the `predict_probability` function that implements the link function:
"""

# ╔═╡ fd120ad4-77e6-11eb-15ea-bb7d7ea35587
function predict_probability(feature_matrix::Matrix, coeffs)
	1. ./ (1. .+ exp.(-(feature_matrix * coeffs)))
end

# ╔═╡ fcf1dd2e-77e6-11eb-0619-e7fd2a622877
md"""
**Aside**. How the link function works with matrix algebra

Since the word counts are stored as columns in **feature_matrix**, each $i$-th row of the matrix corresponds to the feature vector $h(\mathbf{x}_i)$:

```math
[feature\_matrix] =
\left[
	\begin{array}{c}
		h(\mathbf{x}_1)^T \\
		h(\mathbf{x}_2)^T \\
		\vdots \\
		h(\mathbf{x}_N)^T
	\end{array}
\right] = \left[
	\begin{array}{cccc}
	h_0(\mathbf{x}_1) & h_1(\mathbf{x}_1) & \cdots & h_D(\mathbf{x}_1) \\
	h_0(\mathbf{x}_2) & h_1(\mathbf{x}_2) & \cdots & h_D(\mathbf{x}_2) \\
	\vdots & \vdots & \ddots & \vdots \\
	h_0(\mathbf{x}_N) & h_1(\mathbf{x}_N) & \cdots & h_D(\mathbf{x}_N)
\end{array}
\right]
```

By the rules of matrix multiplication, the score vector containing elements $\mathbf{w}^T h(\mathbf{x}_i)$ is obtained by multiplying **feature_matrix** and the coefficient vector $\mathbf{w}$.

```math
[score] =
[feature\_matrix]\mathbf{w} =
\left[
\begin{array}{c}
h(\mathbf{x}_1)^T \\
h(\mathbf{x}_2)^T \\
\vdots \\
h(\mathbf{x}_N)^T
\end{array}
\right]
\mathbf{w}
= \left[
\begin{array}{c}
h(\mathbf{x}_1)^T\mathbf{w} \\
h(\mathbf{x}_2)^T\mathbf{w} \\
\vdots \\
h(\mathbf{x}_N)^T\mathbf{w}
\end{array}
\right]
= \left[
\begin{array}{c}
\mathbf{w}^T h(\mathbf{x}_1) \\
\mathbf{w}^T h(\mathbf{x}_2) \\
\vdots \\
\mathbf{w}^T h(\mathbf{x}_N)
\end{array}
\right]
```
"""

# ╔═╡ bb677796-77e7-11eb-2baf-ffb0ee41fff6
md"""
**Checkpoint**

Just to make sure we are on the right track, we have provided a few examples. If our predict_probability function is implemented correctly, then the following should pass

"""

# ╔═╡ bb4f23d0-77e7-11eb-1f2d-efd11a1d5f66
begin
	d_feature_matrix = [1. 2. 3.; 1. -1. -1]
	d_coeffs = [1.; 3.; -1]

	# c_scores = d_feature_matrix .* d_coeffs
	c_scores = sum(d_feature_matrix .* d_coeffs', dims=2)
	c_preds  = [1. ./ (1. .+ exp.(-c_scores[1])); 1. / (1. .+ exp.(-c_scores[2]))]

	ϵ = 1e-7
	@test all(t -> abs(t[1] - t[2]) ≤ ϵ,
		zip(predict_probability(d_feature_matrix, d_coeffs), c_preds))
end

# ╔═╡ 9744c78c-77ed-11eb-3bf6-fdc3d8cf5281
md"""
#### Compute derivative of log likelihood with respect to a single coefficient

Recall from lecture:

```math
\frac{\partial\ell}{\partial w_j} = \sum_{i=1}^N h_j(\mathbf{x}_i)\left(\mathbf{1}[y_i = +1] - P(y_i = +1 | \mathbf{x}_i, \mathbf{w})\right)
```

We will now write a function that computes the derivative of log likelihood with respect to a single coefficient $w_j$. The function accepts two arguments:
  - `errors` vector containing $\mathbf{1}[y_i = +1] - P(y_i = +1 | \mathbf{x}_i, \mathbf{w})$ for all $i$.
  - `feature` vector containing $h_j(\mathbf{x}_i)$  for all $i$.

"""

# ╔═╡ 2cddfe30-77ee-11eb-389c-130f0af51f70
feature_derivative(errors::Matrix, feature::Vector) = dot(errors, feature)

# ╔═╡ 2cbeab98-77ee-11eb-2108-7596888619c2
md"""
In the main lecture, our focus was on the likelihood.  In the advanced optional video, however, we introduced a transformation of this likelihood *called the log likelihood* that simplifies the derivation of the gradient and is more numerically stable.  Due to its numerical stability, we will use the log likelihood instead of the likelihood to assess the algorithm.

The log likelihood is computed using the following formula (see the advanced optional video if you are curious about the derivation of this equation):

$$\ell\ell(\mathbf{w}) = \sum_{i=1}^N \Big( (\mathbf{1}[y_i = +1] - 1)\mathbf{w}^T h(\mathbf{x}_i) - \ln\left(1 + \exp(-\mathbf{w}^T h(\mathbf{x}_i))\right) \Big)$$

"""

# ╔═╡ 2ca74a52-77ee-11eb-0f50-fd3c15702be1
function log_likelihood(feature_matrix::Matrix, sentiment, coeffs)
	indic = sentiment .== 1
	scores = feature_matrix * coeffs
	log_exp = log.(1. .+ exp.(-scores))

	mask = isinf.(log_exp)
	log_exp[mask] = -scores[mask]

	sum((indic .- 1) .* scores .- log_exp)
end

# ╔═╡ 2c8d9e7c-77ee-11eb-3897-c5bd68dbea53
md"""
**Checkpoint**

Just to make sure we are on the same page, the following code block should pass.
"""

# ╔═╡ 889c2692-77ef-11eb-057c-47b467b71f21
begin
	# re-using d_feature_matrix, d_coeffs above
	d_sentiment = [-1; 1]

	c_indic  = [-1 == +1; 1 == +1] # [Int(-1 == +1); Int(1 == +1)]
	# c_scores (as above)

	c_1st_term  = [(c_indic[1] - 1) * c_scores[1]; (c_indic[2] - 1) * c_scores[2]]
	c_2nd_term = [log(1. + exp(-c_scores[1])); log(1. + exp(-c_scores[2]))]

	c_ll = sum([c_1st_term[1] .- c_2nd_term[1]; c_1st_term[2] .- c_2nd_term[2]])

	@test abs(log_likelihood(d_feature_matrix, d_sentiment, d_coeffs) - c_ll) ≤ ϵ
end

# ╔═╡ 972e0c72-77ed-11eb-0ae7-c127694faf89
md"""
#### Taking gradient steps

Now we are ready to implement our own logistic regression. All we have to do is to write a gradient ascent function that takes gradient steps towards the optimum.
"""

# ╔═╡ 970fba24-77ed-11eb-1a3e-e769d6049078
function logistic_regression(feature_matrix::Matrix, sentiment, init_coeffs;
		η=1e-7, max_iter=200)
	coeffs = copy(init_coeffs)
	# println("coeffs: $(size(coeffs))")

	for ix ∈ 1:max_iter
		preds = predict_probability(feature_matrix, coeffs)
		indic = sentiment .== 1
		errors = indic .- preds
		# println(size(preds), size(indic), size(errors))

		for jx ∈ 1:length(coeffs)
			# println("==> $(size(feature_matrix[:, jx]))")
			∂c = feature_derivative(errors, feature_matrix[:, jx])
			# println("coeffs: $(size(coeffs[jx])), ∂c: $(size(∂c))")
			coeffs[jx] += η * ∂c
		end

		if ix ≤ 15 || (ix ≤ 100 && ix % 10 == 0) || (ix ≤ 1000 && ix % 100 == 0) ||
			(ix ≤ 10000 && ix % 1000 == 0) || ix % 10000 == 0
			lp = log_likelihood(feature_matrix, sentiment, coeffs)
			@printf("iter: %d / %d: log likelihood of observed labels: %.8f\n",
				ceil(Int, log10(max_iter)), ix, lp)
		end
 	end

	coeffs
end

# ╔═╡ 96f4d234-77ed-11eb-0e9d-8b4483b9d020
@time coeffs = logistic_regression(feature_matrix, sentiment,
	zeros(Float64, (194, 1)); max_iter=301);

# ╔═╡ 656b8c04-780c-11eb-16ef-2d0ab940e7b7
size(feature_matrix)

# ╔═╡ a01911d6-7808-11eb-34e1-9dde4224fca8
md"""
**Quiz Question:** As each iteration of gradient ascent passes, does the log likelihood increase or decrease?
  - increase
"""

# ╔═╡ a000fdf0-7808-11eb-32af-fba57700380e
md"""
### Predicting sentiments

Recall from lecture that class predictions for a data point $\mathbf{x}$ can be computed from the coefficients $\mathbf{w}$ using the following formula:

```math
\hat{y}_i =
\left\{
\begin{array}{ll}
      +1 & \mathbf{x}_i^T\mathbf{w} > 0 \\
      -1 & \mathbf{x}_i^T\mathbf{w} \leq 0 \\
\end{array}
\right.
```

Now, we will write some code to compute class predictions. We will do this in two steps:
  1. First compute the **scores** using **feature_matrix** and **coefficients** using a dot product.
  2.  Using the formula above, compute the class predictions from the scores.

Step 1 can be implemented as follows:
"""

# ╔═╡ 9fe93efc-7808-11eb-3755-11a93bd3374a
function calc_preds(feature_matrix::Matrix, coeffs)
	scores = feature_matrix * coeffs
	ifelse.(scores .> 0., 1, -1)
end

# ╔═╡ 9fd0d0b0-7808-11eb-2f66-c5815f2f4929
preds = calc_preds(feature_matrix, coeffs);

# ╔═╡ 9fb85b86-7808-11eb-1af6-0f648740b822
md"""
**Quiz Question:** How many reviews were predicted to have positive sentiment?
"""

# ╔═╡ 380fef64-780c-11eb-3676-7d2fa5d2f726
begin
	n_pos, n_neg = length(preds[preds .== 1]), length(preds[preds .== -1])
	@assert n_pos + n_neg == length(preds)

	(pred_pos=n_pos, pred_neg=n_neg)
end

# ╔═╡ 6bd80650-780f-11eb-20f9-8f99728a72f3
md"""
### Measuring accuracy

We will now measure the classification accuracy of the model. Recall from the lecture that the classification accuracy can be computed as follows:

```math
\mbox{accuracy} = \frac{{\mbox{\# correctly classified data points}}}{\mbox{\# total data points}}
```
"""

# ╔═╡ c7ff6cc0-780f-11eb-3ac5-ffd1fd251698
t = (num_mistakes=sum(sentiment .- preds .!= 0), total_data_points=size(products)[1])

# ╔═╡ 7e0f62a4-7810-11eb-34d9-85a2f2de4296
md"""
**Quiz Question**: What is the accuracy of the model on predictions made above? (round to 2 digits of accuracy)
"""

# ╔═╡ c7e77f86-780f-11eb-0511-5394adb6b0dd
begin
	accuracy = (t.total_data_points - t.num_mistakes) / t.total_data_points;
	round(accuracy, digits=2)
end

# ╔═╡ 8b92d55a-7810-11eb-2495-0112a869bcc9
md"""
### Which words contribute most to positive & negative sentiments?

Recall that in Module 2 assignment, we were able to compute the "**most positive words**". These are words that correspond most strongly with positive reviews. In order to do this, we will first do the following:

  - Treat each coefficient as a tuple, i.e. (**word**, **coefficient_value**).

  - Sort all the (**word**, **coefficient_value**) tuples by **coefficient_value** in descending order.
"""

# ╔═╡ 8c6e485a-7811-11eb-3fe7-19876cc85ee4
md"""
##### Ten "most positive" words

Now, we compute the 10 words that have the most positive coefficient values. These words are associated with positive sentiment.

"""

# ╔═╡ c7cacf90-780f-11eb-21ac-37de331bfd0d
function top_n(;words=important_words, coeffs=coeffs[2:end], key=:positive)
  	sort(
    	[(w, c) for (w, c) ∈ zip(words, coeffs)],
    	by=t -> t[2],
    	rev= key == :positive
	)
end

# ╔═╡ 28b5ed9a-7811-11eb-2e7a-e9ee7b3c7db2
md"""
**Quiz Question**: Which word is **NOT** present in the top 10 "most positive" words?
"""

# ╔═╡ 42107bfc-7811-11eb-35f7-f52a103f7be6
map(t -> t[1], top_n()[1:10])

# ╔═╡ 7321412c-7811-11eb-14eb-f529581256f1
md"""
  - **cheap** is NOT in previous list.
"""

# ╔═╡ 86669250-7811-11eb-1938-5dedabd25ad7
md"""
##### Ten "most negative" words

Next, we repeat this exercise on the 10 most negative words.  That is, we compute the 10 words that have the most negative coefficient values. These words are associated with negative sentiment.
"""

# ╔═╡ 864d674e-7811-11eb-0071-85e6c429d503
map(t -> t[1], top_n(;key=:negative)[1:10])

# ╔═╡ 862e1b1e-7811-11eb-1212-25bce5527b8b
md"""
  - **need** is NOT in previous list.
"""

# ╔═╡ Cell order:
# ╟─baa8ca1c-77d8-11eb-2e26-3b2445acbfa1
# ╠═33bd93ec-77d9-11eb-2da8-87271f8b30c7
# ╟─3ad4770e-77d9-11eb-3a93-addf3e427025
# ╠═6e264178-77d9-11eb-33c1-19462dde66e2
# ╠═75f2e820-77de-11eb-19eb-a5a5977b5257
# ╠═36f2c616-77de-11eb-03b4-3b92a532949a
# ╠═6d4ef592-77d9-11eb-3844-1bea7dc71094
# ╠═d7b9c812-77de-11eb-0ebb-6bb75de9725b
# ╟─6e09ca34-77d9-11eb-1a99-6578f9141004
# ╠═6ded8978-77d9-11eb-3d94-e7507a85197e
# ╟─6dd12c56-77d9-11eb-0350-1dfc08d17dd6
# ╠═6db87918-77d9-11eb-2082-edecdf8b7576
# ╠═6d9e0568-77d9-11eb-1369-bb6990e1e4a7
# ╠═6d69b148-77d9-11eb-04f5-034603897c31
# ╟─6d316f86-77d9-11eb-107b-1bf5708655d5
# ╠═0f6bb770-77df-11eb-0f1f-4dddbd7e2498
# ╟─734ed4d4-77df-11eb-3dab-e5c10b4cda62
# ╠═0f528d0e-77df-11eb-1415-2fb0e9d6792d
# ╠═0f362fec-77df-11eb-02a6-210a58c2839b
# ╟─0f1c6b16-77df-11eb-2197-93c4a82cc135
# ╠═0057a50e-77e0-11eb-29d6-85c2dac67743
# ╠═51a94e8e-77e3-11eb-04c8-d1389b1f5856
# ╠═0f03e474-77df-11eb-2791-f165fd011737
# ╠═0ee68550-77df-11eb-0a6e-c5bf5e3fd9ac
# ╠═0ecc0054-77df-11eb-21a4-3f677548ecf9
# ╟─0eb50d9a-77df-11eb-0b7d-557ab3ea2c81
# ╠═0e9884fe-77df-11eb-3199-13ab2fbdcbc8
# ╟─0e815eaa-77df-11eb-3422-515926dab50d
# ╟─0e4f6222-77df-11eb-3c39-239382844c12
# ╠═0e383a04-77df-11eb-2ec9-c5aeffc4b131
# ╠═0e1d5f86-77df-11eb-12c9-75b0b884ad25
# ╠═0e03942c-77df-11eb-1bde-4f7b6739a3ee
# ╟─0deb769c-77df-11eb-2b53-7b73905353cd
# ╠═0dd0bc6c-77df-11eb-2df1-09ca3d5aaafc
# ╟─fd5a8f7c-77e6-11eb-34dd-11ba6b2d7409
# ╠═fd419c9c-77e6-11eb-2925-0d832a3fe2ca
# ╟─fd23d180-77e6-11eb-2091-af1e84d8fbcf
# ╠═fd120ad4-77e6-11eb-15ea-bb7d7ea35587
# ╟─fcf1dd2e-77e6-11eb-0619-e7fd2a622877
# ╟─bb677796-77e7-11eb-2baf-ffb0ee41fff6
# ╠═bb4f23d0-77e7-11eb-1f2d-efd11a1d5f66
# ╟─9744c78c-77ed-11eb-3bf6-fdc3d8cf5281
# ╠═2cddfe30-77ee-11eb-389c-130f0af51f70
# ╟─2cbeab98-77ee-11eb-2108-7596888619c2
# ╠═2ca74a52-77ee-11eb-0f50-fd3c15702be1
# ╟─2c8d9e7c-77ee-11eb-3897-c5bd68dbea53
# ╠═889c2692-77ef-11eb-057c-47b467b71f21
# ╟─972e0c72-77ed-11eb-0ae7-c127694faf89
# ╠═970fba24-77ed-11eb-1a3e-e769d6049078
# ╠═96f4d234-77ed-11eb-0e9d-8b4483b9d020
# ╠═656b8c04-780c-11eb-16ef-2d0ab940e7b7
# ╟─a01911d6-7808-11eb-34e1-9dde4224fca8
# ╟─a000fdf0-7808-11eb-32af-fba57700380e
# ╠═9fe93efc-7808-11eb-3755-11a93bd3374a
# ╠═9fd0d0b0-7808-11eb-2f66-c5815f2f4929
# ╟─9fb85b86-7808-11eb-1af6-0f648740b822
# ╠═380fef64-780c-11eb-3676-7d2fa5d2f726
# ╟─6bd80650-780f-11eb-20f9-8f99728a72f3
# ╠═c7ff6cc0-780f-11eb-3ac5-ffd1fd251698
# ╟─7e0f62a4-7810-11eb-34d9-85a2f2de4296
# ╠═c7e77f86-780f-11eb-0511-5394adb6b0dd
# ╟─8b92d55a-7810-11eb-2495-0112a869bcc9
# ╟─8c6e485a-7811-11eb-3fe7-19876cc85ee4
# ╠═c7cacf90-780f-11eb-21ac-37de331bfd0d
# ╟─28b5ed9a-7811-11eb-2e7a-e9ee7b3c7db2
# ╠═42107bfc-7811-11eb-35f7-f52a103f7be6
# ╟─7321412c-7811-11eb-14eb-f529581256f1
# ╟─86669250-7811-11eb-1938-5dedabd25ad7
# ╠═864d674e-7811-11eb-0071-85e6c429d503
# ╟─862e1b1e-7811-11eb-1212-25bce5527b8b


se_kernel, se_n_hyper = include("se_kernel.jl")
scale_kernel, scale_n_hyper = include("scale_kernel.jl")

"""
Custom function created by combining kernels from kernel_coder().
hyperparameters == ["λ, shift, σ1, σ2"]
"""
function custom_kernel(
    hyperparameters::Vector{<:Real},
    δ::Real,
    dorder::Vector{<:Integer};
    outputs::Vector{<:Integer}=[1,1])

    @assert length(hyperparameters)==4 "hyperparameters is the wrong length"
    dorder_len = 6
    @assert length(dorder)==dorder_len "dorder is the wrong length"
    @assert maximum(dorder) < 3 "No more than two derivatives for each time or hyperparameter can be calculated"
    @assert minimum(dorder) >= 0
    @assert maximum(outputs) < 3
    @assert minimum(outputs) > 0

    λ = hyperparameters[1]
    shift = hyperparameters[2]
    σ1 = hyperparameters[3]
    σ2 = hyperparameters[4]

    # dλ = dorder[3]
    # dshift = dorder[4]
    dσ1 = dorder[5]
    dσ2 = dorder[6]

    if outputs==[1,1]
        return se_kernel([λ], δ, dorder[1:3]) + scale_kernel([σ1], δ, append!(dorder[1:2],[dσ2])) * Int(δ == 0)
    end

    if outputs in [[1,2],[2,1]]
        return se_kernel([λ, -shift], δ, dorder[1:4]; shift_ind=2)
    end

    if outputs==[2,2]
        return se_kernel([λ], δ, dorder[1:3]) + scale_kernel([σ2], δ, append!(dorder[1:2],[dσ2])) * Int(δ == 0)
    end

end

return custom_kernel, 4  # the function handle and the number of kernel hyperparameters


"""
    rayleigh(λs, nH_I, nHe_II, nH2)

Absorption coefficient from Rayleigh scattering by neutral H, He, and H2.  Formulations for H and He
are via [Colgan+ 2016](https://ui.adsabs.harvard.edu/abs/2016ApJ...817..116C/abstract).  Formulation
for H2 from
[Dalgarno and Williams 1962](https://ui.adsabs.harvard.edu/abs/1962ApJ...136..690D/abstract).

The Dalgarno and Williams H2 is applicable redward of 1300 Å.  Since Rayleigh scattering breaks down
when the particle size to wavelength ratio gets large, we that all frequencies passed to this
function be equivalent to 1300 Å or greater.

The formulations for H is adapted from
[Lee 2005](https://ui.adsabs.harvard.edu/abs/1962ApJ...136..690D/abstract), which states that it is
applicable redward of Lyman alpha. See Colgan 2016 for details on He.
"""
function rayleigh(νs::AbstractVector{<:Real}, nH_I, nHe_I, nH2)
    @assert c_cgs/maximum(νs) > 1.3e-5
    σth = 6.65246e-25 #Thompson scattering cross section [cm^2]

    #(ħω/ 2E_H)^2 in Colgan+ 2016.  The photon energy over 2Ryd
    E_2Ryd_2 = @. (hplanck_eV * νs / (2 * Rydberg_eV ))^2
    E_2Ryd_4 = E_2Ryd_2.^2
    E_2Ryd_6 = E_2Ryd_2.*E_2Ryd_4
    E_2Ryd_8 = E_2Ryd_4.^2

    #Colgan+ 2016 equation 6
    σH_σth = @. 20.24E_2Ryd_4 + 239.2E_2Ryd_6 + 2256E_2Ryd_8
    #Colgan+ 2016 equation 7
    σHe_σth = @. 1.913E_2Ryd_4 + 4.52E_2Ryd_6 + 7.90E_2Ryd_8

    αH_HE = @. (nH_I*σH_σth + nHe_I*σHe_σth) * σth

    #Dalgarno & Williams 1962 equation 3 (which assumes λ in Å)
    invλ2 = (νs ./ (1e8 * c_cgs)) .^2
    invλ4 = invλ2 .^ 2
    invλ6 = invλ2 .* invλ4
    invλ8 = invλ4 .^ 2
    αH2 = @. (8.14e-13*invλ4 + 1.28e-6*invλ6 + 1.61*invλ8) * nH2

    αH_HE .+ αH2
end

"""
    electron_scattering(nₑ)

Compute the linear absorption coefficient, α, from scattering off of free electrons. This has no
wavelength dependence. It assumes isotropic scattering.  (See, e.g. Gray p 160.)

# Arguments
- `nₑ::F`: number density of free electrons (in cgs)
"""
function electron_scattering(nₑ::F) where {F<:Real}
    8π/3 * (electron_charge_cgs^2/(electron_mass_cgs * c_cgs^2))^2 * nₑ
end


module ideal

include("general.jl")

using Unitful
using NLsolve

∑(x)=sum(x)
P(T,x,CA,i) = x[i]*general.P_Antoine(CA,T,i)
Pₜₒₜ(T,x,CA,n) = ∑(P(T,x,CA,i) for i=1:n)
y_calc(T,x,CA,i,n) = P(T,x,CA,i)/Pₜₒₜ(T,x,CA,n)

"""
    pxy(n,T,CA,xx;uni=u"Torr")

**Los campos de entrada son:**
- `n :: Int` Es el numero de componentes
- `T :: Float` Es la temperatura del sistema
- `CA :: n×3 Array{Float,2}` Constantes de Antoine para los `n` componentes
- `xx :: StepRange ó Vector` Son los puntos en el liquido donde se buscara el equilibrio con el vapor
- `uni` Son las unidades en las que se desea el resultado, por defecto esta en Torr

**Salida:**

Regresa las fracciones de liquido `x`,
las fracciones del vapor `y` y la presión `P` en Torr (`m` es la longitud de `xx`).

- `x :: n×m Array{Float64,2}` liquido
- `y :: n×m Array{Float64,2}` vapor
- `P :: m-element Array{Float64,1}` presión
"""
function pxy(n::Int,T,CA::Array,xx;uni=u"Torr")
    T = T |> u"K"
    xx = general.nc(xx...)
    m = size(xx,2)
    y = Array{Float64}(undef,n,m)
    P = Vector{Float64}(undef,m)*uni
    for i=1:m
        x = xx[:,i]
        for j=1:n
            y[j,i] = y_calc(T,x,CA,j,n)
        end
        P[i] = Pₜₒₜ(T,x,CA,n) |> uni
    end

    return xx,y,P
end

"""
    txy(n,Pobj,CA,xx;uni=u"K")

**Los campos de entrada son:**
- `n :: Int` Es el numero de componentes
- `Pobj :: Float` Es la presión del sistema
- `CA :: n×3 Array{Float,2}` Constantes de Antoine para los `n` componentes
- `xx :: StepRange ó Vector` Son los puntos en el liquido donde se buscara el equilibrio con el vapor
- `uni` Son las unidades en las que se desea el resultado, por defecto esta en grados kelvin

**Salida:**

Regresa las fracciones de liquido `x`,
las fracciones del vapor `y` y la temperatura `T` en grados kelvin (`m` es la longitud de `xx`).

- `x :: n×m Array{Float64,2}` liquido
- `y :: n×m Array{Float64,2}` vapor
- `T :: m-element Array{Float64,1}` temperatura
"""
function txy(n::Int,Pobj,CA::Array,xx;uni=u"K")
    Pobj = Pobj |> u"Torr"
    xx = general.nc(xx...)
    m = size(xx,2)
    y = Array{Float64}(undef,n,m)
    T = Vector{Float64}(undef,m)*u"K"
    Tᵢ = 300.0
    #temperatura inicial
    for i=1:n
        if xx[i,1]==1
            T[1] = general.T_Antoine(CA,Pobj,i)
            Tᵢ  = ustrip(u"K",T[1])
        end
    end
    for i = 1:m
        x = xx[:,i]
        P2ₜₒₜ(T) = Pₜₒₜ(T,x,CA,n) - ustrip(u"Torr",Pobj)
        #println(i)
        Pₛₒₗ = nlsolve(n_ary(P2ₜₒₜ),[Tᵢ])
        Tᵢ = Pₛₒₗ.zero[1]
        T[i] = Tᵢ*u"K"
        for j = 1:n
            y[j,i] = y_calc(T[i],x,CA,j,n)
        end
    end

    return xx,y,T .|> uni
end

end  # module


# Utilities for Gibbs sampling with geometric inter-arrivals

function update_geometric_interarrival_param!(p::Vector{Float64},PP::Vector{Int},T::Vector{Int},n::Int,params::Vector{Float64})
    """
    - `p`: current geometric parameter for the interarrival time distribution (will be updated in place)
    - `PP`: partition/degree vector
    - `T`: arrival times vector
    - `n`: number of observations
    - `params`: parameters of prior Beta distribution
    """
    K = size(T,1)
    update_geometric_interarrival_param!(p,K,n,params)

end

function update_geometric_interarrival_param!(p::Vector{Float64},K::Int,n::Int,params::Vector{Float64})
    """
    - `p`: current geometric parameter for the interarrival time distribution (will be updated in place)
    - `K`: number of blocks in partition (=# of arrivals)
    - `n`: number of observations
    - `params`: parameters of prior Beta distribution
    """
    # sample conjugate p
    p[1] = rand(Beta(K-1+params[1],n-K+params[2]))

end


struct Cost{T}
    evaluate::Any
    gradient::Any
    hessian::Any
    num_gradient::Int
    num_hessian::Int
    sparsity::Vector{Vector{Int}}
    evaluate_cache::Vector{T}
    gradient_cache::Vector{T}
    hessian_cache::Vector{T}
end

function Cost(f::Function, num_state::Int, num_action::Int;
    num_parameter::Int=0,
    evaluate_hessian=false)

    #TODO: option to load/save methods
    @variables x[1:num_state], u[1:num_action], w[1:num_parameter]
    evaluate = f(x, u, w)
    gradient = Symbolics.gradient(evaluate, [x; u])
    num_gradient = num_state + num_action
    evaluate_func = eval(Symbolics.build_function([evaluate], x, u, w)[2])
    gradient_func = eval(Symbolics.build_function(gradient, x, u, w)[2])
    if evaluate_hessian
        hessian = Symbolics.sparsehessian(evaluate, [x; u])
        hessian_func = eval(Symbolics.build_function(hessian.nzval, x, u, w)[2])
        sparsity = [findnz(hessian)[1:2]...]
        num_hessian = length(hessian.nzval)
    else
        hessian_func = Expr(:null)
        sparsity = [Int[]]
        num_hessian = 0
    end

    return Cost(
        evaluate_func,
        gradient_func,
        hessian_func,
        num_gradient,
        num_hessian,
        sparsity,
        zeros(1),
        zeros(num_gradient),
        zeros(num_hessian))
end

Objective{T} = Vector{Cost{T}} where T

function cost(objective::Objective, states, actions, parameters)
    J = 0.0
    for (t, cost) in enumerate(objective)
        cost.evaluate(cost.evaluate_cache, states[t], actions[t], parameters[t])
        J += cost.evaluate_cache[1]
    end
    return J
end

function gradient!(gradient, indices, objective::Objective, states, actions, parameters)
    for (t, cost) in enumerate(objective)
        cost.gradient(cost.gradient_cache, states[t], actions[t], parameters[t])
        @views gradient[indices[t]] .+= cost.gradient_cache
        # fill!(cost.gradient_cache, 0.0) # TODO: confirm this is necessary
    end
end

function hessian!(hessian, indices, objective::Objective, states, actions, parameters, scaling)
    for (t, cost) in enumerate(objective)
        cost.hessian(cost.hessian_cache, states[t], actions[t], parameters[t])
        cost.hessian_cache .*= scaling
        @views hessian[indices[t]] .+= cost.hessian_cache
        # fill!(cost.hessian_cache, 0.0) # TODO: confirm this is necessary
    end
end

function sparsity_hessian(objective::Objective, num_state::Vector{Int}, num_action::Vector{Int})
    row = Int[]
    col = Int[]
    for (t, cost) in enumerate(objective)
        if !isempty(cost.sparsity[1])
            shift = (t > 1 ? (sum(num_state[1:t-1]) + sum(num_action[1:t-1])) : 0)
            push!(row, (cost.sparsity[1] .+ shift)...)
            push!(col, (cost.sparsity[2] .+ shift)...)
        end
    end
    return collect(zip(row, col))
end

function hessian_indices(objective::Objective, key::Vector{Tuple{Int,Int}}, num_state::Vector{Int}, num_action::Vector{Int})
    indices = Vector{Int}[]
    for (t, cost) in enumerate(objective)
        if !isempty(cost.sparsity[1])
            row = Int[]
            col = Int[]
            shift = (t > 1 ? (sum(num_state[1:t-1]) + sum(num_action[1:t-1])) : 0)
            push!(row, (cost.sparsity[1] .+ shift)...)
            push!(col, (cost.sparsity[2] .+ shift)...)
            rc = collect(zip(row, col))
            push!(indices, [findfirst(x -> x == i, key) for i in rc])
        else
            push!(indices, Int[])
        end
    end
    return indices
end

# num_gradient(objective::Objective{T}) where T = sum([obj.num_gradient  for obj in objective])
# num_hessian(objective::Objective{T}) where T  = sum([obj.num_hessian   for obj in objective])


## using KronLinInv

###############################################

function test2D( )

    ## 2D problem, so set nx = 1
    nx = 1
    ny = 20
    nz = 30
    nxobs = 1
    nyobs = 18
    nzobs = 24
    sigmaobs  = [1.0, 0.1, 0.1]
    corlenobs = [0.0, 1.4, 1.4]
    sigmam    = [1.0, 0.8, 0.8]
    corlenm   = [0.0, 2.5, 2.5]

    # run test
    result = test_mean_cov(nx,ny,nz,nxobs,nyobs,nzobs,
                           sigmaobs,corlenobs,sigmam,corlenm)

    return result
end

#########################################

function test3D( )

    ## 2D problem, so set nx = 1
    nx = 8
    ny = 9
    nz = 12
    nxobs = 6
    nyobs = 7
    nzobs = 11
    sigmaobs  = [0.1, 0.1, 0.1]
    corlenobs = [1.4, 1.4, 1.4]
    sigmam    = [0.8, 0.8, 0.8]
    corlenm   = [2.5, 2.5, 2.5]

    # run test
    result = test_mean_cov(nx,ny,nz,nxobs,nyobs,nzobs,
                           sigmaobs,corlenobs,sigmam,corlenm)

    return result
end

########################################

function test_mean_cov(nx::Integer,ny::Integer,nz::Integer,
                       nxobs::Integer,nyobs::Integer,nzobs::Integer,
                       sigmaobs::Array{Float64,1},corlenobs::Array{Float64,1},
                       sigmam::Array{Float64,1},corlenm::Array{Float64,1})

    ##############################
    ## Setup the problem
    ##############################

    ## Covariance matrices
    # covariance on observed data
    Cd1,Cd2,Cd3 = mkCovSTAT(sigmaobs,nxobs,nyobs,nzobs,corlenobs,"gaussian")
    # covariance on model parameters
    Cm1,Cm2,Cm3 = mkCovSTAT(sigmam,nx,ny,nz,corlenm,"gaussian")

    ## Create the covariance matrix structure
    Covs = CovMats(Cd1,Cd2,Cd3,Cm1,Cm2,Cm3)

    ## Forward model operator
    G1 = rand(nxobs,nx)
    G2 = rand(nyobs,ny)
    G3 = rand(nzobs,nz)

    ## Create the forward operator structure
    Gfwd = FwdOps(G1,G2,G3)

    ##############################
    ## Setup the synthetic test
    ##############################

    ## Create a reference model
    refmod = rand(nx*ny*nz)

    ## Create a reference model
    mprior = copy(refmod) #0.5 .* ones(nx*ny*nz)

    ## Create some "observed" data
    ##   (without noise because it's just a test of the algo)
    dobs = kron(G1,kron(G2,G3)) * refmod

    ##############################
    ## Solve the inverse problem
    ##############################

    ## Calculate the required factors
    klifac = calcfactors(Gfwd,Covs)

    ## Calculate the posterior mean model
    postm = posteriormean(klifac,Gfwd,mprior,dobs)

    # serial version
    postm2 = KronLinInv.posteriormean_serial(klifac,Gfwd,mprior,dobs)

    ## Calculate the posterior covariance
    npts = nx*ny*nz
    astart, aend = 1,div(npts,3)
    bstart, bend = 1,div(npts,3)
    postC = blockpostcov(klifac,astart,aend,bstart,bend)

    ## serial version
    postC2 = KronLinInv.blockpostcov_serial(klifac,astart,aend,bstart,bend)

    ##############################
    ## Check results
    ##############################

    if (postm ≈ refmod) && (postm2 ≈ refmod)
        return true
    end
    return false
end

###########################################################

function mkCovSTAT(sigma::Array{Float64,1},nx::Integer,ny::Integer,nz::Integer,
                   corrlength::Array{Float64,1},kind::String)
    #
    #   Stationaly covariance model
    #
    #println("Creating covariance matrix...")

    function cgaussian(dist,corrlength)
        if maximum(dist)==0.0
            return 1.0
        else
            @assert(corrlength>0.0)
            return exp.(-(dist./corrlength).^2)
        end
    end
    function cexponential(dist,corrlength)
        if maximum(dist)==0.0
            return 1.0
        else
            @assert(corrlength>0.0)
            return exp.(-(dist./corrlength))
        end
    end

    npts = nx*ny*nz
    x = [float(i) for i=1:nx]
    y = [float(i) for i=1:ny]
    z = [float(i) for i=1:nz]
    covmat_x = zeros(nx,nx)
    covmat_y = zeros(ny,ny)
    covmat_z = zeros(nz,nz)

    if kind=="gaussian"
        calccovfun = cgaussian
    elseif kind=="exponential"
        calccovfun = cexponential
    else
        println("Error, no or wrong cov 'kind' specified")
        exit()
    end

    for i=1:nx
        covmat_x[i,:] .= sigma[1]^2 .*
            calccovfun(sqrt.((x.-x[i]).^2),corrlength[1])
    end
    for i=1:ny
        covmat_y[i,:] .= sigma[2]^2 .*
            calccovfun(sqrt.(((y.-y[i])).^2),corrlength[2])
    end
    for i=1:nz
        covmat_z[i,:] .= sigma[3]^2 .*
            calccovfun(sqrt.(((z.-z[i])).^2),corrlength[3])
    end

    return covmat_x,covmat_y,covmat_z
 end

##########################################################