#///////////////////////////////////////
#// File Name: ileqg_test.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2020/11/03
#// Description: Test code for src/ileqg.jl
#///////////////////////////////////////

using LinearAlgebra
using Test

@testset "iLEQG Test" begin
    f(x, u) = x + u;
    c(k, x, u) = k;
    h(x) = 1.0;
    W(k) = Matrix(1.0I, 2, 2);
    N = 10;

    prob = FiniteHorizonRiskSensitiveOptimalControlProblem(f, c, h, W, N)

    u_array = [ones(2) for ii = 1 : prob.N];
    x_array = simulate_dynamics(prob, zeros(2), u_array);

    @test x_array[1] == zeros(2)
    @test all([x_array[ii + 1] == f(x_array[ii], u_array[ii]) for ii = 1 : prob.N - 1])

    L_array = [ones(2, 2) for ii = 1:prob.N];
    x_array_new, u_array_new = simulate_dynamics(prob, x_array, u_array, L_array)
    @test all([u_array_new[ii] == u_array[ii] for ii = 1:prob.N])
    @test all([x_array_new[ii] == x_array[ii] for ii = 1:prob.N])

    # integrate_cost test
    cost = integrate_cost(prob, x_array, u_array)
    @test cost ≈ sum([c(ii - 1, x_array[ii], u_array[ii]) for ii = 1 : prob.N]) + h(x_array[end])

    # initialize test
    solver = ILEQGSolver(prob);
    initialize!(solver, prob, zeros(2), u_array, 0.0);
    @test solver.l_array == u_array
    @test solver.L_array == [zeros(2, 2) for ii = 1 : prob.N]
    @test solver.x_array == x_array
    @test solver.μ == 0.0
    @test solver.Δ == solver.Δ_0
    @test solver.d_current == Inf
    @test solver.iter_current == 0;
    @test solver.ϵ_history == Tuple{Float64, Float64}[]
    dp_result_init = solve_approximate_dp(
                    approximate_model(prob, u_array, x_array),
                    [zeros(2, 2) for ii = 1 : N], θ=0.0, μ=0.0);
    @test solver.value_current ≈ dp_result_init.s_array[1];

    # approximate_model test
    prob.c(k, x, u) = 0.5*x'*x + 1.0*u'*u + x'*u
    prob.h(x) = 0.5*x'*x
    approx_result = approximate_model(prob, u_array, x_array);

    @test all([isapprox(approx_result.q_array[ii], 0.5*(2*(ii - 1)^2) + 1.0*2 + 2*(ii - 1)) for ii = 1 : length(u_array)])
    @test isapprox(approx_result.q_array[end], h(x_array[end]))

    @test all([isapprox(approx_result.q_vec_array[ii], x_array[ii] + ones(2)) for ii = 1:length(u_array)])
    @test isapprox(approx_result.q_vec_array[end], x_array[end])
    @test all([isapprox(approx_result.Q_array[ii], Matrix(1.0I, 2, 2)) for ii = 1:length(x_array)])
    @test all([isapprox(approx_result.r_array[ii], x_array[ii] + 2.0*ones(2)) for ii = 1:length(u_array)])
    @test all([isapprox(approx_result.R_array[ii], 2.0*Matrix(1.0I, 2, 2)) for ii = 1:length(u_array)])
    @test all([isapprox(approx_result.P_array[ii], Matrix(1.0I, 2, 2)) for ii = 1:length(u_array)])
    @test all([approx_result.W_array[ii] == Matrix(1.0I, 2, 2) for ii = 1:length(u_array)])

    prob.c(k, x, u) = 0.5*x'*x + 1.0*u'*u;
    prob.h(x) = 0.5*x'*x
    approx_result = approximate_model(prob, u_array, x_array);
    # solve_approximate_dp test
    dp_result, dl_array_new = solve_approximate_dp!(solver, approx_result, false, θ=0.0);
    @test length(dp_result.s_array) == length(x_array)
    @test length(dp_result.s_vec_array) == length(x_array)
    @test all([length(s_vec) == 2 for s_vec in dp_result.s_vec_array])
    @test length(dp_result.S_array) == length(x_array)
    @test all([size(S) == (2, 2) for S in dp_result.S_array])
    @test all([issymmetric(S) for S in dp_result.S_array])
    @test all([isposdef(S) for S in dp_result.S_array])
    @test length(dp_result.g_array) == length(u_array)
    @test all([length(g) == 2 for g in dp_result.g_array])
    @test length(dp_result.G_array) == length(u_array)
    @test all([size(G) == (2, 2) for G in dp_result.G_array])
    @test length(dp_result.H_array) == length(u_array)
    @test all([size(H) == (2, 2) for H in dp_result.H_array])
    # # gains should match LQR solution.
    begin
        L_array_lqr = Vector{Matrix{Float64}}(undef, N);
        S_array_lqr = Vector{Matrix{Float64}}(undef, N + 1);
        S_array_lqr[N + 1] = approx_result.Q_array[N + 1];
        for ii = Iterators.reverse(1 : N)
            Q = approx_result.Q_array[ii];
            R = approx_result.R_array[ii];
            A = approx_result.A_array[ii];
            B = approx_result.B_array[ii];
            S_array_lqr[ii] = Q + (A')*S_array_lqr[ii + 1]*A -
                              (A')*S_array_lqr[ii + 1]*B/(R + (B')*S_array_lqr[ii + 1]*B)*(B')*S_array_lqr[ii + 1]*A;
        end
        for ii = 1 : N
            R = approx_result.R_array[ii];
            A = approx_result.A_array[ii];
            B = approx_result.B_array[ii];
            L_array_lqr[ii] = -(R + (B')*S_array_lqr[ii + 1]*B)\B*S_array_lqr[ii + 1]*A
        end
        @test all([all(L_array_lqr[ii] .≈ solver.L_array[ii]) for ii = 1 : N])
    end
    # # Nominal control offsets should be 0 (LQR is linear feedback)
    @test all(isapprox.(norm.(u_array .+ dl_array_new .- solver.L_array.*x_array[1:end-1]), 0.0, atol=1e-8))

    dp_result_2, dl_array_new_2 = solve_approximate_dp!(solver, approx_result, false, θ=1e-8);
    @test length(dp_result_2.s_array) == length(x_array)
    @test length(dp_result_2.s_vec_array) == length(x_array)
    @test all([length(s_vec) == 2 for s_vec in dp_result_2.s_vec_array])
    @test length(dp_result_2.S_array) == length(x_array)
    @test all([size(S) == (2, 2) for S in dp_result_2.S_array])
    @test all([issymmetric(S) for S in dp_result_2.S_array])
    @test all([isposdef(S) for S in dp_result_2.S_array])
    @test length(dp_result_2.g_array) == length(u_array)
    @test all([length(g) == 2 for g in dp_result_2.g_array])
    @test length(dp_result_2.G_array) == length(u_array)
    @test all([size(G) == (2, 2) for G in dp_result_2.G_array])
    @test length(dp_result_2.H_array) == length(u_array)
    @test all([size(H) == (2, 2) for H in dp_result_2.H_array])
    @test isapprox(dp_result.s_array[1], dp_result_2.s_array[1], rtol=1e-5)
    @test all([isapprox(dl_array_new[ii], dl_array_new_2[ii]) for ii = 1 : N])

    solve_approximate_dp!(solver, approx_result, false, θ=0.0);

    dp_result_3 = solve_approximate_dp(approx_result, solver.L_array, dl_array_new, θ=0.0, μ=0.0);
    @test dp_result_3.s_array == dp_result.s_array

    # line search test (for linear system)
    line_search!(solver, prob, dl_array_new, 0.0, false)
    @test solver.value_current .≈ dp_result.s_array[1] # line search should find the same optimal solution

    # increase_μ_and_Δ test
    solver = ILEQGSolver(prob);
    initialize!(solver, prob, zeros(2), u_array, 0.0)
    increase_μ_and_Δ!(solver);
    @test solver.Δ == 4.0
    @test solver.μ == 1e-6

    # decrease_μ_and_Δ test
    solver = ILEQGSolver(prob);
    initialize!(solver, prob, zeros(2), u_array, 0.0)
    decrease_μ_and_Δ!(solver);
    @test solver.Δ == 0.5
    @test solver.μ == 0.0

    # Testing with nonlinear model
    f(x, u) = x.^1.3 + u.^1.5;
    c(k, x, u) = sum(x.^2.5 + u.^2.5);
    h(x) = 1.0;
    W(k) = Matrix(0.01I, 2, 2);
    N = 10;
    θ = 0.5;

    u_array = [0.1*ones(2) for ii = 1 : prob.N];

    prob = FiniteHorizonRiskSensitiveOptimalControlProblem(f, c, h, W, N)

    solver = ILEQGSolver(prob);
    initialize!(solver, prob, zeros(2), u_array, θ)
    approx_result = approximate_model(prob, solver.l_array, solver.x_array,
                                      solver.A_array, solver.B_array);
    dp_result, dl_array = solve_approximate_dp!(solver, approx_result, false, θ=θ);
    value_new = line_search!(solver, prob, dl_array, θ, false)
    @test length(solver.ϵ_history) == 1
    @test solver.ϵ_history[1][1] == 1.0;
    @test solver.ϵ_history[1][2] < 0.0

    x_array, l_array, L_array, value, ϵ_history =
    solve!(solver, prob, zeros(2), u_array, θ=0.0, verbose=false);
    @test all([all(isapprox.(x_array[ii], zeros(2), atol=1e-4)) for ii = 1 : N + 1])
end


abstract type InitializationStrategyGamma <: InitializationStrategy end

struct FixedGammaStrategy <: InitializationStrategyGamma
    kernel
end

calculate_gamma(model, strategy::FixedGammaStrategy) = MLKernels.getvalue(strategy.kernel.alpha)
function calculate_gamma(model::SubSVDD, strategy::FixedGammaStrategy, subspace_idx)
    MLKernels.getvalue(strategy.kernel[subspace_idx].alpha)
end

"""
Original publication:
Silverman, Bernard W. Density estimation for statistics and data analysis. Routledge, 2018.
"""
struct RuleOfThumbSilverman <: InitializationStrategyGamma end

function rule_of_thumb_silverman(data::Array{T,2}) where T <: Real
    return (size(data, 2) * (size(data, 1) + 2) / 4.0)^(-1.0 / (size(data,1) + 4.0))
end

calculate_gamma(model, strategy::RuleOfThumbSilverman) = rule_of_thumb_silverman(model.data)

function calculate_gamma(model::SubSVDD, strategy::RuleOfThumbSilverman, subspace_idx)
    return rule_of_thumb_silverman(model.data[model.subspaces[subspace_idx], :])
end

"""
Original publication:
Scott, David W. Multivariate density estimation: theory, practice, and visualization. John Wiley & Sons, 2015.
"""
struct RuleOfThumbScott <: InitializationStrategyGamma end

function rule_of_scott(data::Array{T,2}) where T <: Real
    return size(data, 2)^(-1.0/(size(data, 1) + 4))
end

calculate_gamma(model, strategy::RuleOfThumbScott) = rule_of_scott(model.data)

function calculate_gamma(model::SubSVDD, strategy::RuleOfThumbScott, subspace_idx)
    return rule_of_scott(model.data[model.subspaces[subspace_idx], :])
end

"""
Generate binary data to tune a one class classifier according to the following paper:
Wang, S. et al. 2018. Hyperparameter selection of one-class support vector machine by self-adaptive data
shifting. Pattern Recognition. 74, 2018.

Changes compared to the paper:
1. We only use the strategy here to fit gamma and not a combination of gamma and C.
2. We also fit on the real data. This is because the algorithm sometimes generates no pseudo inliers.
3. We test on pseudo-inliers, pseudo-outliers and real data. We assume that the real data are all inliers when testing.
4. We use the F1 score instead of an averaged FNR and FPR.
"""
struct WangGammaStrategy <: InitializationStrategyGamma
    solver
    gamma_search_range
    C
    scoring_function
end

WangGammaStrategy(solver) = WangGammaStrategy(solver, 1.0)
WangGammaStrategy(solver, C) = WangGammaStrategy(solver, 10.0.^range(-2, stop=2, length=50), C, f1_scoring)
WangGammaStrategy(solver, gamma_search_range, C) = WangGammaStrategy(solver, gamma_search_range, C, f1_scoring)

function generate_binary_data_for_tuning(data, k=nothing, threshold=0.1)
    if k === nothing
        k = round(Int, ceil(5 * log10(size(data, 2))))
    end
    tree = KDTree(data)
    edge_idx = Int[]
    norm_vec = []
    data_target = []
    l_ns = 0

    for i in 1:size(data, 2)
        idx, dist = knn(tree, data[:, i], k + 1, true)
        v_ij = mapslices(normalize, data[:, i] .- data[:, idx[2:end]], dims=1)
        n_i = sum(v_ij, dims=2)
        θ_ij = sum(v_ij .* n_i, dims=1)
        l_i = 1 / k * sum(θ_ij .>= 0)
        if l_i >= 1 - threshold
            # add new edge
            push!(edge_idx, i)
            push!(norm_vec, n_i)
            l_ns += 1 / k * sum(dist)
        end

        # generate pseudo inlier (shift direction is ∇p(x) = -n)
        n_i = -normalize(vec(n_i))
        Λ_i_positive = sum(n_i .* (data[:, idx[2:end]] .- data[:, i]), dims=1)
        if length(Λ_i_positive[Λ_i_positive .> 0]) > 0
            # shift along positive direction of data density gradient
            x_ij_min_positive = minimum(Λ_i_positive[Λ_i_positive .> 0])
            push!(data_target, data[:, i] + x_ij_min_positive * n_i)
        end
    end

    # compute negative shift amount
    l_ns *= 1 / length(edge_idx)
    # generate pseudo outliers by shifting along negative direction
    data_outliers = data[:, edge_idx] + mapslices(normalize, hcat(norm_vec...), dims=1) * l_ns

    data_target = hcat(data_target...)

    return data_target, data_outliers
end

function f1_scoring(predictions, ground_truth)
    return MLBase.f1score(MLBase.roc(ground_truth .== :outlier, predictions .== :outlier))
end

function calculate_gamma(model::SubSVDD, strategy::WangGammaStrategy, subspace_idx)
    info(LOGGER, "[Gamma Search] Using VanillaSVDD to search for gamma in subspace $subspace_idx")
    calculate_gamma(VanillaSVDD(model.data[model.subspaces[subspace_idx], :]), strategy)
end

function calculate_gamma(model::SubSVDD, strategy::WangGammaStrategy)
    info(LOGGER, "[Gamma Search] Using VanillaSVDD to estimate a global gamma.")
    calculate_gamma(VanillaSVDD(model.data), strategy)
end
function calculate_gamma(model, strategy::WangGammaStrategy)
    m = deepcopy(model)
    data_target, data_outliers = generate_binary_data_for_tuning(m.data)
    if isempty(data_target)
        data_predictions = hcat(m.data, data_outliers)
        ground_truth = vcat(fill(:inlier, size(m.data, 2)),
                            fill(:outlier, size(data_outliers, 2)))
    else
        data_predictions = hcat(m.data, data_target, data_outliers)
        ground_truth = vcat(fill(:inlier, size(m.data, 2) + size(data_target, 2)),
                            fill(:outlier, size(data_outliers, 2)))
    end

    debug(LOGGER, "[Gamma Search] Searching for parameter gamma.")
    best_gamma = 1.0
    best_score = -Inf
    for gamma in strategy.gamma_search_range
        debug(LOGGER, "[Gamma Search] Testing gamma = $gamma")
        init_strategy = FixedParameterInitialization(GaussianKernel(gamma), strategy.C)
        initialize!(m, init_strategy)
        set_adjust_K!(m, true)
        try
            fit!(m, strategy.solver)
        catch e
            debug(LOGGER, "[Gamma Search] Fitting failed for gamma $gamma.")
            debug(LOGGER, e)
            continue
        end
        predictions = classify.(predict(m, data_predictions))
        score = strategy.scoring_function(ground_truth, predictions)
        if score > best_score
            debug(LOGGER, "[Gamma Search] New best found with gamma = $gamma and score = $score.")
            best_gamma = gamma
            best_score = score
        end
    end
    info(LOGGER, "[Gamma Search] Finished.")
    return best_gamma
end


"""
Test function of a parabola
"""
function parabola(t, x; active::Bool=false)
    if active == true
        return [3*t^2]
    else
        return [0.0]
    end
end

"""
Simple test function for spherical earth dynamics
"""
function point_earth(epc::Epoch, x::Array{<:Real, 1})
    # Restrict inputs to position only. Considered in body frame
    r  = x[1:3]
    v  = x[4:6]

    # Acceleration
    a = -GM_EARTH * r/norm(r)^3

    return vcat(v, a)
end

# Test basic integration of parabola
let
    rk4 = RK4(parabola, active=false)
    @test istep(rk4, 0.0, 1.0, [0.0])[1] == 0.0

    rk4 = RK4(parabola, active=true)
    @test istep(rk4, 0.0, 1.0, [0.0])[1] == 1.0

    t = 0.0
    for i in 1:10
        t = istep(rk4, 0, i, [0.0])[1]
    end

    @test t == 1000
end

# Test stability of integration of conservative motion (point-mass orbit)
let
    rk4 = RK4(point_earth)

    epc0 = Epoch(2019, 1, 1)
    oe0  = [R_EARTH + 500e3, 0.01, 90.0, 0, 0, 0]
    eci0 = sOSCtoCART(oe0, use_degrees=true)

    epcf = epc0 + orbit_period(oe0[1])

    # Manually propagate orbit
    epc = deepcopy(epc0)
    eci = copy(eci0)
    while epc < epcf
        dt  = min(1.0, epcf-epc) # Set step-size
        eci = istep(rk4, epc, dt, eci)
        epc += dt
    end

    # Test that initial state equals intial propagated state
    @test isapprox(eci[1, 1], eci0[1], atol=1e-5)
    @test isapprox(eci[2, 1], eci0[2], atol=1e-5)
    @test isapprox(eci[3, 1], eci0[3], atol=1e-5)
    @test isapprox(eci[4, 1], eci0[4], atol=1e-5)
    @test isapprox(eci[5, 1], eci0[5], atol=1e-5)
    @test isapprox(eci[6, 1], eci0[6], atol=1e-5)

    oe = sCARTtoOSC(eci[:, 1], use_degrees=true)
    @test isapprox(oe[1], oe0[1], atol=1e-5)
    @test isapprox(oe[2], oe0[2], atol=1e-5)
    @test isapprox(oe[3], oe0[3], atol=1e-5)
    @test isapprox(oe[4], oe0[4], atol=1e-5)
    @test isapprox(oe[5], oe0[5], atol=1e-5)
    @test isapprox(oe[6], oe0[6], atol=1e-5)
end

let
    # Initialize integrator
    rk4 = RK4(point_earth)

    epc = Epoch(2019, 1, 1)
    eci = sOSCtoCART([R_EARTH + 500e3, 0.01, 90.0, 0, 0, 0], use_degrees=true)
    phi = diagm(0 => ones(Float64, length(eci)))

    # No step variational equations STM should be identity
    xu, phiu = istep(rk4, epc, 0.0, eci, phi)

    for i in 1:length(eci)
        @test phiu[i,i] == 1.0
    end

    # Small step STM should be nearly identity but non-zero
    xu, phiu = istep(rk4, epc, 1.0, eci, phi)

    for i in 1:length(eci)
        @test phiu[i, i] != 0.0
        @test isapprox(phiu[i, i], 1.0, atol=1.0e-5)
    end

    # Actually performing a perturbed step should equal the STM output
    pert = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0]

    # Perturbed position from integrating pertturbed initial conditions
    ecip   = istep(rk4, epc, 1.0, eci + pert)

    # Pertrubed position from variational equation approximation
    ecistm = istep(rk4, epc, 1.0, eci) + phiu*pert

    # Assert equal
    for i in 1:6
        @test isapprox(ecip[i], ecistm[i], atol=1.0e-9)
    end
end


wavenumber_low_volfrac(ω::Number, medium::Medium, specie::Specie; kws...) =
  wavenumber_low_volfrac(ω, medium, [specie]; kws...)

wavenumber_low_volfrac(ωs::AbstractVector{T}, medium::Medium{T}, species::Vector{Specie{T}}; kws...) where T<:Number =
  [wavenumber_low_volfrac(ω, medium, species; kws...) for ω in ωs]

wavenumber_low_volfrac(ωs::AbstractVector{T}, medium::Medium{T}, specie::Specie{T}; kws...) where T<:Number = [wavenumber_low_volfrac(ω, medium, [specie]; kws...) for ω in ωs]

"Explicit formula for one effective wavenumber based on a low particle volume fraction expasion."
function wavenumber_low_volfrac(ω::T, medium::Medium{T}, species::Vector{Specie{T}}; tol::T =1e-6,
    hankel_order::Int = maximum_hankel_order(ω, medium, species; tol=tol),
    radius_multiplier::T = 1.005, verbose::Bool = true) where T <: Number

  volume_fraction = sum(pi*sp.r^2.0*sp.num_density for sp in species)
  if volume_fraction >= 0.4 && verbose
    @warn("the volume fraction $(round(100*volume_fraction))% is too high, expect a relative error of approximately $(round(100*volume_fraction^3.0))%")
  end
  num_density = sum(sp.num_density for sp in species)
  # Add incident wavenumber
  kT2 = (ω/medium.c)^2.0
  # Add far-field contribution
  kT2 += - 4.0im*num_density*far_field_pattern(ω, medium, species; hankel_order=hankel_order)(0.0)
  # Add pair-field contribution
  kT2 += - 4.0im*num_density^(2.0)*pair_field_pattern(ω, medium, species;
      radius_multiplier=radius_multiplier, hankel_order=hankel_order)(0.0)

  return (imag(sqrt(kT2)) > zero(T)) ? sqrt(kT2) : -sqrt(kT2)
end

reflection_coefficient_low_volfrac(ωs::AbstractVector{T}, medium::Medium{T}, species::Vector{Specie{T}}; kws... ) where T<:Number =
    [reflection_coefficient_low_volfrac(ω, medium, species; kws... ) for ω in ωs]

"An explicit formula for the refleciton coefficient based on a low particle volume fraction."
function reflection_coefficient_low_volfrac(ω::T, medium::Medium{T}, species::Vector{Specie{T}};
        θin::T = zero(T), kws... ) where T<:Number

    θ_ref = T(π) - T(2)*θin
    fo = far_field_pattern(ω, medium, species; kws...)
    dfo = diff_far_field_pattern(ω, medium, species; kws...)
    foo = pair_field_pattern(ω, medium, species; kws...)

    k = ω/medium.c
    α = k*cos(θin)
    R1 = im*fo(θ_ref)
    R2 = 2.0*fo(zero(T))/(α^2.0)
    R2 = im*foo(θ_ref) + R2*(sin(θin)*cos(θin)*dfo(θ_ref) - fo(θ_ref))

    num_density = sum(s.num_density for s in species)
    R = (R1 + num_density*R2)*num_density/(α^2.0)
    return R
end

function wavenumber_very_low_volfrac(ω::Number, medium::Medium, species::Array{Specie{T}};
    tol=1e-6, radius_multiplier = 1.005, verbose = false)  where T<:Number

  volume_fraction = sum(pi*sp.r^2.0*sp.num_density for sp in species)
  if volume_fraction >= 0.4
    @warn("the volume fraction $(round(100*volume_fraction))% is too high, expect a relative error of approximately $(round(100*volume_fraction^3.0))%")
  end
  kT2 = (ω/medium.c)^2.0
  next_order = 4.0im*sum(sp.num_density*Zn(ω,sp,medium,0) for sp in species)
  hankel_order=1

  # sum more hankel orders until the relative error < tol
  while abs(next_order/kT2) > tol
    kT2 += next_order
    next_order = 4.0im*sum(sp.num_density*Zn(ω,sp,medium,m) for sp in species, m in (-hankel_order,hankel_order))
    hankel_order +=1
  end
  kT2 += next_order
  hankel_order +=1
  if verbose println("max Hankel order = $hankel_order") end

  return sqrt(kT2)
end


module TestFillArrays

using Test
using SIMDIndexing

include("../examples/vdot.jl")
using LinearAlgebra
using FillArrays

arr1 = randn(2^20)
arr2 = Ones{Float64}(length(arr1))

@test vdot(arr1, arr2) ≈ dot(arr1, arr2)

end  # module


#-----------------------------------------------------------------------
#       BEGIN   Gordon.jl
#-----------------------------------------------------------------------
#-----------------------------------------------------------------------
#   DYAD
#-----------------------------------------------------------------------
"""
    uugen( u::RealVec, X::RealVec )

The function `uugen` constructs a dyad from the vector u.

"""
function uugen( u::RealVec, X::RealVec )
    return transpose(u).*u
end     #---------------------------------------------------------------

#-----------------------------------------------------------------------
#   COSINE OF THETA
#-----------------------------------------------------------------------
"""
    cθ( X::RealVec )

The function `cθ` computes the cosine of θ (as defined in spherical
polar coordinates, using the physicist convention).

"""
function cθ( X::RealVec )
    return X[4]/norm(X[2:4])
end     #---------------------------------------------------------------

#-----------------------------------------------------------------------
#   SINE OF PHI
#-----------------------------------------------------------------------
"""
    sϕ( X::RealVec )

The function `sϕ` computes the sine of ϕ (as defined in spherical
polar coordinates, using the physicist convention).

"""
function sϕ( X::RealVec )
    return X[3]/norm(X[3:4])
end     #---------------------------------------------------------------

#-----------------------------------------------------------------------
#   THE GORDON METRIC
#-----------------------------------------------------------------------
"""
    gGordon( X::RealVec , n::Function=nIR , gfunc::Function=giso ,
             U::RealVec=Float64[-1;0;0;0] )

The function `gGordon` computes the Gordon metric. The argument `n` can
be replaced with a user-supplied index of refraction function, the
argument `gfunc` can be replaced with a user-supplied background metric,
and the argument `U` can be replaced with a user-supplied fluid
four-velocity function.

"""
function gGordon( X::RealVec , n::Function=nIR ,
                  gfunc::Function=giso , U::RealVec=Float64[-1;0;0;0] )

    tpfl=typeof(X[1])

    gs = gfunc( X )

    Unsq = abs(transpose(U)*gs*U)

    UU = uugen( U , X )

    return gs + ( 1-1/(n(X)^2) )*UU/Unsq
end     #---------------------------------------------------------------

#-----------------------------------------------------------------------
#       END     Gordon.jl
#-----------------------------------------------------------------------


Sphere(x...) = HyperSphere(x...)

Circle(x...) = HyperSphere(x...)

widths(c::HyperSphere{N, T}) where {N, T} = Vec{N, T}(radius(c)*2)
radius(c::HyperSphere) = c.r
origin(c::HyperSphere) = c.center
minimum(c::HyperSphere{N, T}) where {N, T} = Vec{N, T}(origin(c)) - Vec{N, T}(radius(c))
maximum(c::HyperSphere{N, T}) where {N, T} = Vec{N, T}(origin(c)) + Vec{N, T}(radius(c))


function perform_thresholding(f, M, t)
    """
        Only 3 types of thresholding currently implemented
    """
    if t == "largest"
        a = sort(abs(f)[:])[end:-1:1] #sort a 1D copy of F in descending order
        T = a[M]
        y = f.*(abs(f) .> T)
    elseif t == "soft"
        s = abs(f) - M
        s = (s + abs(s))./2
        y = (f./abs(f)).*s
    elseif t == "hard"
        y = f.*(abs(f) .> M)
    end
    return y
end


############################################################################################
# Constants
"RNG for sampling based solutions"
const _rng = Random.GLOBAL_RNG   # shorthand
Random.seed!(_rng, 1)

"Tolerance for stochastic solutions"
const _TOL = 1.0e-6

"Number of samples"
N = 10^3

############################################################################################
"Normalize Vector inplace and return normalizing constant"
@inline function normalize!(vec::AbstractVector)
    norm = sum(vec)
    vec ./= norm
    return norm
end
"Faster max function for Viterbi algorithm in HMM. Not exported."
function vec_maximum(vec::AbstractVector)
    m = vec[1]
    @inbounds for i = Base.OneTo(length(vec))
        if vec[i] > m
            m = vec[i]
        end
    end
    return m
end

############################################################################################
"Compute Stationary distribution for given transition matrix"
function get_stationary!(Transition::AbstractMatrix{T}) where T<:Real
    # From: https://github.com/QuantEcon/QuantEcon.jl/blob/f454d4dfbaf52f550ddd52eff52471e4b8fddb9d/src/markov/mc_tools.jl
    # Grassmann-Taksar-Heyman (GTH) algorithm (Grassmann, Taksar, and Heyman 1985)
    n = size(Transition, 1)
    x = zeros(T, n)

    @inbounds for k in 1:n-1
        scale = sum(Transition[k, k+1:n])
        if scale <= zero(T)
            # There is one (and only one) recurrent class contained in
            # {1, ..., k};
            # compute the solution associated with that recurrent class.
            n = k
            break
        end
        Transition[k+1:n, k] /= scale

        for j in k+1:n, i in k+1:n
            Transition[i, j] += Transition[i, k] * Transition[k, j]
        end
    end

    # backsubstitution
    x[n] = 1
    @inbounds for k in n-1:-1:1, i in k+1:n
        x[k] += x[i] * Transition[i, k]
    end

    # normalisation
    x /= sum(x)

    return x
end
get_stationary(Transition::AbstractMatrix{T}) where {T<:Real} = get_stationary!(copy(Transition))

############################################################################################
"Check if ancestors are all in correct order - Not that first column can have different order as ancestors not resampled at 0."
function check_ancestors(ancestors::AbstractMatrix{I}) where {I<:Integer}
    for Ncol in 2:size(ancestors, 2)
        for Nrow in Base.OneTo(size(ancestors, 1))
            if ancestors[Nrow, Ncol] != Nrow
                return false
            end
        end
    end
    return true
end

############################################################################################
include("hmm.jl")
include("hmmMV.jl")
include("hsmm.jl")
include("HigherOrderMarkov.jl")

############################################################################################
resamplemethods = [Systematic(), Residual(), Stratified(), BaytesFilters.Multinomial()]
referencemethods = [Conditional(), Ancestral(), Marginal()]
objectives = [deepcopy(markov_objective), deepcopy(markov_MV_objective), deepcopy(semimarkov_objective)]
generating = [UpdateFalse(), UpdateTrue()]


import Base: show, string, typejoin

abstract type AbstractLogicTerm end
abstract type SententialTerm{T<:AbstractLogicTerm} <: AbstractLogicTerm end
abstract type JunctionTerm{T<:AbstractLogicTerm} <: SententialTerm{T} end
abstract type QuantifierTerm{T<:AbstractLogicTerm} <: SententialTerm{T} end

# we make Variable <: SententialTerm to enable FOOL extensions
struct Variable <: SententialTerm{AbstractLogicTerm}
  name::String
end
Variable(name::AbstractString) = Variable(name)
Variable(v::Variable) = Variable(v.name)
Variable(name::AbstractString, t::Symbol) = Variable(name)
Variable(v::Variable, t::Symbol) = Variable(v.name)
string(v::Variable) = v.name

struct Functor <: AbstractLogicTerm
  name::String
end
string(f::Functor) = f.name

show(io::IO, q::AbstractLogicTerm) = print(io, string(q))

struct FunctionTerm <: AbstractLogicTerm
  name::Functor
  args::Vector{Union{FunctionTerm, Variable}}
end
FunctionTerm(n::AbstractString, args) = FunctionTerm(Functor(string(n)), args)
ConstantTerm(n::Functor) = FunctionTerm(n, Variable[])
ConstantTerm(n) = FunctionTerm(Functor(string(n)), Variable[])

function listString(xs::AbstractVector, head::AbstractString,
                    sep::AbstractString, tail::AbstractString)
  foldl((x,y) -> x*sep*string(y), xs[2:end]; init=head*string(xs[1]))*tail
end

struct PredicateTerm{T<:AbstractLogicTerm} <: SententialTerm{T}
  name::Functor
  args::Vector{T}
end
PredicateTerm(n::String, args) = PredicateTerm(Functor(n), args)

equalTerm(t1::Union{PredicateTerm,FunctionTerm,Variable},
          t2::Union{PredicateTerm,FunctionTerm,Variable}) =
            PredicateTerm(Functor("="), [t1,t2])

string(f::Union{FunctionTerm,PredicateTerm}) =
  isempty(f.args) ? string(f.name) : listString(f.args, string(f.name)*"(", ",", ")")

struct NegationTerm{T<:AbstractLogicTerm} <: SententialTerm{T}
  scope::T
end
string(n::NegationTerm) = "¬(" * string(n.scope) * ")"

const LiteralTerm = Union{PredicateTerm{T}, NegationTerm{PredicateTerm{T}}} where {T<:Union{Variable,FunctionTerm}}

Base.typejoin(::Type{PredicateTerm},::Type{NegationTerm{PredicateTerm}}) = LiteralTerm

struct AndTerm{T<:AbstractLogicTerm} <: JunctionTerm{T}
  juncts::Vector{T}
end
AndTerm(s::T, ss...) where {T<:SententialTerm} =
  AndTerm{foldr((si,T)->typejoin(T,typeof(si)), ss; init=T)}([s; collect(ss)])
string(a::AndTerm) = listString(a.juncts, "(", ")∧(", ")")

struct OrTerm{T<:AbstractLogicTerm} <: JunctionTerm{T}
  juncts::Vector{T}
end
OrTerm(s::T, ss...) where {T<:SententialTerm} =
  OrTerm{foldr((si,T)->typejoin(T,typeof(si)), ss; init=T)}([s; collect(ss)])
string(a::OrTerm) = listString(a.juncts, "(", ")∨(", ")")

struct EQuantifierTerm{T<:AbstractLogicTerm} <: QuantifierTerm{T}
  variables::Vector{Variable}
  scope::T
end
function EQuantifierTerm(args...)
  x = collect(args);  EQuantifierTerm{typeof(x[end])}(x[1:end-1], x[end])
end
string(q::EQuantifierTerm) = "∃" * listString(q.variables, "", ",", "") * " " * string(q.scope)

struct AQuantifierTerm{T<:AbstractLogicTerm} <: QuantifierTerm{T}
  variables::Vector{Variable}
  scope::T
end
function AQuantifierTerm(args...)
  x = collect(args);  AQuantifierTerm{typeof(x[end])}(x[1:end-1], x[end])
end
string(q::AQuantifierTerm) = "∀" * listString(q.variables, "", ",", "") * " " * string(q.scope)

for T in (Variable, FunctionTerm, PredicateTerm, NegationTerm, AndTerm, OrTerm,
          AQuantifierTerm,EQuantifierTerm)
  @eval rootType(_::$T) = $T
end

pairedQuantifierType(_::AndTerm) = AQuantifierTerm
pairedQuantifierType(_::OrTerm) = EQuantifierTerm

"""
    freeVar(a::AbstractLogicTerm)

Return a Vector{Variable} of the free variables in a.
"""
freeVar(v::Variable) = [v]
freeVar(fp::Union{FunctionTerm,PredicateTerm}) =
  foldl((u,x) -> union!(u, freeVar(x)), fp.args; init=Variable[])
freeVar(n::NegationTerm) = freeVar(n.scope)
freeVar(j::JunctionTerm) = foldl((u,x) -> union!(u, freeVar(x)), j.juncts; init=Variable[])
freeVar(q::QuantifierTerm) = setdiff(freeVar(q.scope), q.variables)


"""

```
radial_distribution(
  sys::System{T},
  trajfile::String;
  dmax::Float64=10.,
  nbins::Int=100
) where T
```

Computes the g(r). If `T=Point3D` the system will be read as three-dimensional.

"""
function radial_distribution(
  sys::System{T},
  trajfile::String;
  dmax::Float64=10.,
  nbins::Int=100
) where T

  function shell_volume(r,Δr,T::DataType)
    if T == Point2D
      return π*((r+Δr)^2 - r^2)
    elseif T == Point3D
      return (4π/3)*((r+Δr)^3 - r^3)
    else
      error("T must be Point2D or Point3D")
    end
  end

  # Initialize histogram
  @unpack n, sides = sys
  x = zeros(T,n)
  density = n / prod(sides)

  Δr = dmax/nbins
  r = [ (i-1/2)*Δr for i in 1:nbins ]
  g = zeros(nbins)

  # Reading coordinates, computing distances, and adding to histogram
  nframes = 0
  open(trajfile) do file
    while !(eof(file))
      readline(file) # number of atoms, ignore
      readline(file) # title, ignore
      nframes += 1
      # read atomic coordinates
      for i in 1:n
        atom_data = split(readline(file))
        x[i] = T(ntuple(i -> parse(Float64,atom_data[i+1]), dim(T)))
      end
      # Compute distances and add to histogram
      for i in 1:n-1
        for j in i+1:n
          dx = image(x[j] - x[i],sides)
          idist = ceil(Int,nbins*(norm(dx)/dmax))
          if idist <= nbins
            g[idist] += 1.
          end
        end
      end
    end
  end
  println("Number of frames read = ", nframes)

  # Normalizing by the number of frames and number of atoms
  # Normalizing by the number of atoms expected in an area or volume
  for i in 1:nbins
    g[i] = g[i] / (nframes*((n-1)/2))
    g[i] = g[i] / (density*shell_volume(r[i],Δr,T))
  end

  return r, g
end


# Utils

export 	getF,
		linearXMapping

"""
    getF(date::Date)

Return F factor to convert PJUR into Risk-weigthed asset (RWA), using
the relationship ``RWA = (1/F) * PJUR``.
F is the minimum Basel Index required at date `date::Date`.

"""
function getF(date::Date)

	if date <= Date(2015, 12, 31)
		F = 0.11
	elseif date <= Date(2016, 12, 31)
		F = 0.09875
	elseif date <= Date(2017, 12, 31)
		F = 0.0925
	elseif date <= Date(2018, 12, 31)
		F = 0.08625
	else
		F = 0.08
	end #if

	return(F)

end #function

"""
	linearXMapping(maturities::Array{Int64,1}, exposures::Array{Float64,1}, standardizedVertices::Array{Int64,1})

Map exposures into standardized vertices (SVs), using linear method.
Suppose we have an exposure of ``X`` BRL that matures in ``T`` business days,
and ``n`` SVs ``t_1, \ldots, \t_n`` in each we must allocate our exposure.
If ``T`` is shorter than any SV, we allocate a fraction of the exposure in the first SV:

`` SV(1) = \frac{T}{t_0} X ``

If ``T`` is longer than any SV, we allocate a multiple of the exposure in the last SV:

`` SV(n) = \frac{T}{t_n} X ``

If ``T`` lies between two vertices (``i`` and ``i+1``), we allocate a fraction of the exposure in the two adjacent vertices,
inversely proportional to the distance between ``T`` and the vertex:

`` SV(i) = \frac{ T - t_{i+1} }{ t_{i+1} - t_i } X ``

`` SV(i+1) = \frac{ t_i - T }{ t_{i+1} - t_i } X ``.

# Parameters
* `maturities::Array{Int64,1}`: an array of maturities in business days. It is not required to pass unique or ordered maturities.
* `exposures::Array{Float64,1}`: an array of exposures (present value) to each maturity.
* `standardizedVertices::Array{Int64,1}`: an array of ordered standardized vertices in business days,
such as package constants `VERTICES_RWAJUR1`, `VERTICES_RWAJUR2`, `VERTICES_RWAJUR3` and `VERTICES_RWAJUR4`.

"""
function linearXMapping(maturities::Array{Int64,1}, exposures::Array{Float64,1}, standardizedVertices::Array{Int64,1})
	if size(maturities, 1) != size(exposures, 1)
		error("maturities and exposures must have the same length")
	end #if
	if !issorted(standardizedVertices)
		error("stantardized vertices must be sorted")
	end

	N = size(maturities, 1)
	NS = size(standardizedVertices, 1)
	out = zeros(NS)
	for i = 1:N

		pos = sum(standardizedVertices .<= maturities[i])
		if pos == 0 # before the first vertex
			out[1] += (maturities[i]/standardizedVertices[1]) * exposures[i]
		elseif pos == size(standardizedVertices, 1) # after the last vertex
			out[NS] += (maturities[i]/standardizedVertices[NS]) * exposures[i]
		else # between two vertices
			alpha = (standardizedVertices[pos+1] - maturities[i])/(standardizedVertices[pos+1] - standardizedVertices[pos])
			out[pos] += alpha * exposures[i]
			out[pos+1] += (1-alpha) * exposures[i]
		end

	end #for
	return(out)
end #function

"""
	linearXMapping(date::Date, maturities::Array{Date, 1}, exposures::Array{Float64, 1}, standardizedVertices::Array{Int64, 1})

Map exposures into standardized vertices (SVs), using linear method and calculating maturities in business days using
Brazilian holiday calendar.
See `linearXMapping(maturities::Array{Int64,1}, exposures::Array{Float64,1}, standardizedVertices::Array{Int64,1})` for more info.

# Parameters
* `date::Date`: date of evaluation
* `maturities::Array{Int64,1}`: an array of maturities (dates). It is not required to pass unique or ordered maturities.
* `exposures::Array{Float64,1}`: an array of exposures (present value) to each maturity.
* `standardizedVertices::Array{Int64,1}`: an array of ordered standardized vertices in business days,
such as package constants `VERTICES_RWAJUR1`, `VERTICES_RWAJUR2`, `VERTICES_RWAJUR3` and `VERTICES_RWAJUR4`.

"""
function linearXMapping(date::Date, maturities::Array{Date, 1}, exposures::Array{Float64, 1}, standardizedVertices::Array{Int64, 1})
	linearXMapping(map(x -> x.value,
				     bdays(:Brazil, date, maturities)),
				 exposures,
				 standardizedVertices)
end #function


using MultivariateOrthogonalPolynomials, Test

include("test_modalinterlace.jl")
include("test_disk.jl")
include("test_rectdisk.jl")
include("test_triangle.jl")
# include("test_dirichlettriangle.jl")


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Desription
# ==============================================================================
#
#   Tests related to conversion from Euler angles to direction cosine matrices.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# File: ./src/conversions/angle_to_dcm.jl
# =======================================

# Functions: angle_to_dcm
# -----------------------

@testset "Euler angles => DCM (Float64)" begin
    T = Float64

    for rot_seq in valid_rot_seqs
        # Sample Euler angles.
        ea = EulerAngles(_rand_ang(T), _rand_ang(T), _rand_ang(T), rot_seq)

        # Convert to DCM.
        D = angle_to_dcm(ea)
        @test eltype(D) === T

        # Create the DCM by composing the rotations.
        rot_seq_str = string(rot_seq)
        D₁ = create_rotation_matrix(ea.a1, Symbol(rot_seq_str[1]))
        D₂ = create_rotation_matrix(ea.a2, Symbol(rot_seq_str[2]))
        D₃ = create_rotation_matrix(ea.a3, Symbol(rot_seq_str[3]))
        De = D₃ * D₂ * D₁

        # Compare.
        @test D ≈ De
    end
end

@testset "Euler angles => DCM (Float32)" begin
    T = Float32

    for rot_seq in valid_rot_seqs
        # Sample Euler angles.
        ea = EulerAngles(_rand_ang(T), _rand_ang(T), _rand_ang(T), rot_seq)

        # Convert to DCM.
        D = angle_to_dcm(ea)
        @test eltype(D) === T

        # Create the DCM by composing the rotations.
        rot_seq_str = string(rot_seq)
        D₁ = create_rotation_matrix(ea.a1, Symbol(rot_seq_str[1]))
        D₂ = create_rotation_matrix(ea.a2, Symbol(rot_seq_str[2]))
        D₃ = create_rotation_matrix(ea.a3, Symbol(rot_seq_str[3]))
        De = D₃ * D₂ * D₁

        # Compare.
        @test D ≈ De
    end
end

@testset "Euler angles => DCM (Errors)" begin
    @test_throws ArgumentError angle_to_dcm(1, 2, 3, :ZZY)
end

# Functions: smallangle_to_dcm
# ----------------------------

@testset "Small Euler angles => DCM (Float64)" begin
    T = Float64

    # Sample three small angles.
    θx = _rand_ang(T) * T(0.001)
    θy = _rand_ang(T) * T(0.001)
    θz = _rand_ang(T) * T(0.001)

    # Create the rotation matrix without normalization.
    D = smallangle_to_rot(θx, θy, θz, normalize = false)
    @test eltype(D) === T

    # Expected result.
    De = DCM(
          1, +θz, -θy,
        -θz,   1, +θx,
        +θy, -θx,   1
    )'

    @test De ≈ D

    # Create the rotation matrix with normalization.
    D = smallangle_to_rot(θx, θy, θz, normalize = true)
    @test eltype(D) === T

    @test orthonormalize(De) ≈ D
end

@testset "Small Euler angles => DCM (Float32)" begin
    T = Float32

    # Sample three small angles.
    θx = _rand_ang(T) * T(0.001)
    θy = _rand_ang(T) * T(0.001)
    θz = _rand_ang(T) * T(0.001)

    # Create the rotation matrix without normalization.
    D = smallangle_to_rot(θx, θy, θz, normalize = false)
    @test eltype(D) === T

    # Expected result.
    De = DCM(
          1, +θz, -θy,
        -θz,   1, +θx,
        +θy, -θx,   1
    )'

    @test De ≈ D

    # Create the rotation matrix with normalization.
    D = smallangle_to_rot(θx, θy, θz, normalize = true)
    @test eltype(D) === T

    @test orthonormalize(De) ≈ D
end


module ArnoldCatMap
export arnold_cat_map
using ..ExampleBase: LEDemo, DiscreteExample

@inline function phase_dynamics!(u_next, u, M, t)
    A_mul_B!(u_next, M, u)
    u_next .= mod.(u_next, 1)
end

@inline @views function tangent_dynamics!(u_next, u, M, t)
    phase_dynamics!(u_next[:, 1], u[:, 1], M, t)
    for i in 2:size(u)[2]
        @inbounds A_mul_B!(u_next[:, i], M, u[:, i])
    end
end

"""
Return a [`LEDemo`](@ref) for the Arnold's cat map

* <https://en.wikipedia.org/wiki/Arnold%27s_cat_map>
"""
function arnold_cat_map(;
        u0::AbstractVector = [0.1, 0.1],
        M::AbstractMatrix = [2 1; 1 1],
        t_renorm=10,
        t_attr=10000,
        atol=0, rtol=1e-4,
        terminator_options = [:atol => atol, :rtol => rtol],
        kwargs...)
    @assert (length(u0), length(u0)) == size(M)
    LEDemo(DiscreteExample(
        "Arnold's cat map",
        phase_dynamics!, u0, t_renorm, M,
        tangent_dynamics!,
        t_attr,
        log.(sort!(abs.(eigvals(M)), rev=true)), # known_exponents
        atol, rtol,
        terminator_options,
    ); kwargs...)
end

end


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ b9fc511c-6007-11eb-0c6b-1f6871a40710
using Pkg, DrWatson

# ╔═╡ 20d377b2-6008-11eb-364a-617b6934ecb2
begin
	@quickactivate "StatsModelComparisons"
	using StatsModelComparisons
	using StanSample, StatsFuns, StatsPlots
	using DataFrames, CSV, RDatasets
end

# ╔═╡ af6b0b20-6008-11eb-2fa1-2f61145ab7db
md"
!!! note

	This script assumes that RDatasets.jl is loaded. RDatasets.jl is not included in the dependencies of PSIS.jl (as it would require R to be installed)."

# ╔═╡ 20d43cda-6008-11eb-09f0-53489a26110d
df = RDatasets.dataset("datasets", "cars");

# ╔═╡ 20e22570-6008-11eb-1565-f13541b41861
cars_stan = "
data {
    int N;
    vector[N] speed;
    vector[N] dist;
}
parameters {
    real a;
    real b;
    real sigma;
}
transformed parameters{
    vector[N] mu;
    mu = a + b * speed;
}
model {
    a ~ normal(0, 100);
    b ~ normal(0, 10);
    sigma ~ exponential(1);
    dist ~ normal(mu, sigma)    ;
}
generated quantities {
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = normal_lpdf(dist[i] | mu[i], sigma);
}
";

# ╔═╡ 20e2ebb6-6008-11eb-34f5-61f1ec3d024c
begin
	cars_stan_model = SampleModel("cars.model", cars_stan)
	data = (N = size(df, 1), speed = df.Speed, dist = df.Dist)
	rc = stan_sample(cars_stan_model; data)

	if success(rc)
		cars_df = read_samples(cars_stan_model; output_format=:dataframe)
		isdefined(Main, :StatisticalRethinking) &&
            PRECIS(cars_df[:, [:a, :b, :sigma]])
	end
end

# ╔═╡ 5fc59200-6008-11eb-3e06-1d0bcdf11d7d
if success(rc)
	nt_cars = read_samples(cars_stan_model);
	log_lik = nt_cars.log_lik'
end;

# ╔═╡ 20ed768a-6008-11eb-13f4-458ca1a29592
begin
    ns, n = size(log_lik)
	lppd = reshape(logsumexp(log_lik .- log(ns); dims=1), n)
	pwaic = [var(log_lik[:, i]) for i in 1:n]
	-2(sum(lppd) - sum(pwaic))
end

# ╔═╡ efe960f0-600a-11eb-1df4-5be83899715a
begin
	waic_vec = -2(lppd - pwaic)
	sqrt(n*var(waic_vec))
end

# ╔═╡ 20f5a31e-6008-11eb-2ce1-075893273872
waic(log_lik)

# ╔═╡ 6b67e38e-6009-11eb-3d9f-afc517b7d9fb
begin
	loo, loos, pk = psisloo(log_lik)
	loo
end

# ╔═╡ 82232dfe-6009-11eb-3b32-dbc487c6e4e7
pk_qualify(pk)

# ╔═╡ 8ddf22b0-6009-11eb-08da-5198ba046628
pk_plot(pk)

# ╔═╡ Cell order:
# ╟─af6b0b20-6008-11eb-2fa1-2f61145ab7db
# ╠═b9fc511c-6007-11eb-0c6b-1f6871a40710
# ╠═20d377b2-6008-11eb-364a-617b6934ecb2
# ╠═20d43cda-6008-11eb-09f0-53489a26110d
# ╠═20e22570-6008-11eb-1565-f13541b41861
# ╠═20e2ebb6-6008-11eb-34f5-61f1ec3d024c
# ╠═5fc59200-6008-11eb-3e06-1d0bcdf11d7d
# ╠═20ed768a-6008-11eb-13f4-458ca1a29592
# ╠═efe960f0-600a-11eb-1df4-5be83899715a
# ╠═20f5a31e-6008-11eb-2ce1-075893273872
# ╠═6b67e38e-6009-11eb-3d9f-afc517b7d9fb
# ╠═82232dfe-6009-11eb-3b32-dbc487c6e4e7
# ╠═8ddf22b0-6009-11eb-08da-5198ba046628


function sinkhornTransport(a,b,K,U,lambda,stoppingCriterion="marginalDifference",p_norm=Inf,tolerance=.5e-2,maxIter= 5000,VERBOSE=0)
    # Compute the N Sinkhorn divergence for all the pairs
    # D = [d(a_1,b_1),...,d(a_N,b_N)]
    #--------------------------------------------------
    # Inputs:
    #--------------------------------------------------
    # a : is either
    #   - a d1 x 1 column vector in the probability simpolex (non-negative summing to 1). This is the [1-vs-N mode]
    #   - a d1 x N matrix, where each column vector is in the probability simplex.This is the [N x 1-vs-1 mode]
    # b : is a d2 x N matrix of N vectors in the probability simplex
    #
    # K is a d1 x d2 matrix, equal to exp(-lambda M), where M is the d1 x d2 matrix of pairwise distances between
    # bins described in a and bins in the b_1,...b_N histograms.
    # In the most simple case d_1=d_2 and M is simply a distance matrix (zero on the diagonal and such that m_ij < m_ik + m_kj
    #
    # U = K.*M is a d1 x d2 matrix, pre-stored to speed up the computation of the distances.
    #--------------------------------------------------
    # Optional Inputs:
    #--------------------------------------------------
    # Outputs:
    #--------------------------------------------------
    # D : vector of N dual-sinkhorn divergences, or upper bounds to the EMD.
    #
    # L : vector of N lower bounds to the original OT problem, a.k.a EMD. This is computed by using the dual variables of
    # the smoothed problem, which, when modified adequately, are feasible for the original (non-smoothed) OT dual problem
    #
    # u : d1 x N matrix of left scalings
    # v : d2 x N matrix of right scalings
    # The smoothed optimal transport between (a_i,b_i) can be recovered as
    # T_i = diag(u(:,i)) * K * diag(v(:,i));
    #
    # or, equivalently and substantially faster:
    # TO-DO: the julia syntax is not working now
    # T_i = broadcast(*,v(:,i)',(broadcast(*,u(:,i),K)))
    # Check this for details: https://docs.julialang.org/en/v0.6.1/manual/arrays/
    #
    # Relevant paper:
    # M. Cuturi,
    # Sinkhorn Distances : Lightspeed Computation of Optimal Transport,
    # Advances in Neural Information Processing Systems (NIPS) 26, 2013
    #
    # Optional Input Check:
    #--------------------------------------------------
    # First make sure that a and b are two dimensional array
    # TO-DO:
    #
    ## Checking the type of computation: 1-vs-N points or many pairs
    a_r1, a_c = size(a)
    b_r, b_c = size(b)

    if a_c == 1
        ONE_VS_N = true # We are computing [D(a,b_1), ... , D(a,b_N)]
    elseif a_c == b_c
        ONE_VS_N = false # We are computing [D(a_1,b_1), ... , D(a_N,b_N)]
    else
        error("The first parameter a is either a column vector in the probability simplex, or N column vectors in the probability simplex where N is size(b,2)")
    end
    ## Checking dimensionality:
    if b_c > b_r
        BIGN = true
    else
        BIGN = false
    end
    ## Small changes in the 1-vs-N case to go a bit faster.
    if ONE_VS_N # if computing 1-vs-N make sure all components of a are >0. Otherwise we can get rid of some   lines of K to go faster.
        I = find(a) # Changed from Matlab
        someZeroValues = false
        if length(I)<a_r1 # need to update some vectors and matrices if a does not have full support
            someZeroValues=true
            K = K[I,:]
            U = U[I,:]
            a = a[I]
        end
        ainvK= broadcast(/,K,a) # precomputation of this matrix saves a d1 x N Schur product at each iteration.
    end
    # display(ainvK[1:5,1:5])

    a_r = length(a)
    ## Fixed point counter
    compt=0;
    ## Initialization of
    u =  (1/a_r).*ones(a_r,b_c)./a_r #broadcast(/,ones(a_r,b_c),a_r)#
    v = zeros(b_r,b_c)
    if cmp(stoppingCriterion, "distanceRelativeDecrease")==0
        Dold = ones(1,b_c); # initialization of vector of distances.
    end
    while compt < maxIter
        if ONE_VS_N # 1-vs-N mode
            if BIGN
                u = 1./(ainvK*broadcast(/,b,(K'*u)))
            else
                u=1./(ainvK*broadcast(/,b,(u'*K)'))
            end
        else
            if BIGN
                u=broadcast(/,a,(K*broadcast(/,b,(u'*K)')))
            else
                u=broadcast(/,a,(K*broadcast(/,b,(K'*u))))
            end
        end
        compt=compt+1
        # check the stopping criterion every 20 fixed point iterations
        # or, if that's the case, before the final iteration to store the most
        # recent value for the matrix of right scaling factors v.
        if mod(compt,20)==1 || compt==maxIter
        # split computations to recover right and left scalings.
            if BIGN
                v=broadcast(/,b,(K'*u)) # main iteration of Sinkhorn's algorithm
            else
                v=broadcast(/,b,(u'*K)')
            end
            if ONE_VS_N # 1-vs-N mode
                u=1./(ainvK*v)
            else
                u=broadcast(/,a,(K*v))
            end
            # check stopping criterion
            if cmp(stoppingCriterion, "distanceRelativeDecrease")==0
                D=sum(broadcast(*,u,(U*v)))
                Criterion=norm(broadcast(/,D,Dold)-1,p_norm)
                if Criterion<tolerance || isnan(Criterion)
                    break
                end
                 Dold=D
            elseif cmp(stoppingCriterion,"marginalDifference")==0
                Criterion=norm(sum(abs.(broadcast(*,v,(K'*u))-b)),p_norm)
                if Criterion<tolerance || isnan(Criterion)
                    break
                end
            else
                error("Stopping Criterion not recognized")
            end# norm of all || . ||_1 differences between the marginal of the current solution with the actual marginals.
            compt=compt+1
            if VERBOSE>0
                display("Iteration : $compt, Criterion: $Criterion")
            end
            # if any(isnan(Criterion)), # stop all computation if a computation of one of the pairs goes wrong.
            #     error("NaN values have appeared during the fixed point iteration. This problem appears because of insufficient machine precision when processing computations with a regularization value of lambda that is too high. Try again with a reduced regularization parameter lambda or with a thresholded metric matrix M.");
            # end
        end
    end
    if cmp(stoppingCriterion,"marginalDifference")==0 # if we have been watching marginal differences, we need to compute the vector of distances.
        D=sum(u.*(U*v))
    end

    alpha = log.(u)
    beta = log.(v)
    [beta[t]==-Inf ? beta[t]=0 : beta[t]=beta[t] for t in eachindex(beta)]
    if ONE_VS_N
        L= (a'* alpha + sum(b.*beta))/lambda
    else
        [alpha[t]==-Inf ? alpha[t]=0 : alpha[t]=alpha[t] for t in eachindex(alpha)]
        L= (sum(broadcast(*,a,alpha)) + sum(broadcast(*,b,beta)))/lambda
    end
    # display(L)
    if ONE_VS_N && someZeroValues
        uu=u
        u=zeros(a_r1,b_c)
        u[I,:]=uu
    end
    return D, L, u, v
end


####
# FB Oracle tests
####

facts("suppFcnTest FB") do
	mfs     = [.1, .2]
	mbs     = [-.05, -.25]
	sigfs   = [1., 2.]
	sigbs   = [2., .5]
	w       = FBOracle(mfs, mbs, sigfs, sigbs, .1)

	zstar, ustar = suppFcn([1, 1], w, :Min)
	@fact zstar -->  roughly(-4.724022297688992)
	@fact ustar[1] --> roughly(-4.213785691942581)
	@fact ustar[2] --> roughly(-0.5102366057464114)

	zstar, ustar = suppFcn([1, 1], w, :Max)
	@fact zstar --> roughly(5.098525912188081)
	@fact ustar[1] --> roughly(1.0597051824376162)
	@fact ustar[2] --> roughly(4.038820729750465)

	zstar, ustar = suppFcn([-1,  1], w, :Min)
	@fact zstar --> roughly(-2.7492629560940403)
	@fact ustar[1] --> roughly(2.0194103648752324)
	@fact ustar[2] --> roughly(-0.7298525912188081)

	zstar, ustar = suppFcn([-1,  1], w, :Max)
	@fact zstar --> roughly(6.319708517540586)
	@fact ustar[1] --> roughly(-3.0848542587702927)
	@fact ustar[2] --> roughly(3.234854258770293)
end

facts("portTest for FB") do
	srand(8675309); data = randn(500, 2)
	w = FBOracle(data, .1, .2)
	portTest(w, -1.791121672173153, [0.5840372921578592, 0.41596270784214084],TOL=1e-2)
end


import DPFEHM
import GaussianRandomFields
import Optim
import PyPlot
import Random
import Zygote

Random.seed!(0)

function getobsnodes(coords, obslocs)
	obsnodes = Array{Int}(undef, length(obslocs))
	for i = 1:length(obslocs)
		obsnodes[i] = findmin(map(j->sum((obslocs[i] .- coords[:, j]) .^ 2), 1:size(coords, 2)))[2]
	end
	return obsnodes
end

mins = [0, 0]#meters
maxs = [100, 100]#meters
allns = [[101, 101], [201, 201], [401, 401], [801, 801]]
allns = [[allns[end]]; allns]
num_eigenvectors = 200
x_true = randn(num_eigenvectors)
x0 = zeros(num_eigenvectors)
sqrtnumobs = 16
obslocs_x = range(mins[1], maxs[1]; length=sqrtnumobs + 2)[2:end - 1]
obslocs_y = range(mins[2], maxs[2]; length=sqrtnumobs + 2)[2:end - 1]
obslocs = collect(Iterators.product(obslocs_x, obslocs_y))[:]
observations = Array{Float64}(undef, length(obslocs))
for (i, ns) in enumerate(allns)
	Random.seed!(0)
	global x_true
	global x0
	global observations
	coords, neighbors, areasoverlengths, _ = DPFEHM.regulargrid2d(mins, maxs, ns, 1.0)
	Qs = zeros(size(coords, 2))
	boundaryhead(x, y) = 5 * (x - maxs[1]) / (mins[1] - maxs[1])
	dirichletnodes = Int[]
	dirichleths = zeros(size(coords, 2))
	for i = 1:size(coords, 2)
		if coords[1, i] == mins[1] || coords[1, i] == maxs[1] || coords[2, i] == mins[2] || coords[2, i] == maxs[2]
			push!(dirichletnodes, i)
			dirichleths[i] = boundaryhead(coords[1:2, i]...)
		end
	end

	lambda = 50.0#meters -- correlation length of log-conductivity
	sigma = 1.0#standard deviation of log-conductivity
	mu = 0.0#mean of log-permeability can be arbitrary because this is steady-state and there are no fluid sources

	cov = GaussianRandomFields.CovarianceFunction(2, GaussianRandomFields.Matern(lambda, 1; σ=sigma))
	x_pts = range(mins[1], maxs[1]; length=ns[1])
	y_pts = range(mins[2], maxs[2]; length=ns[2])
	@time grf = GaussianRandomFields.GaussianRandomField(cov, GaussianRandomFields.KarhunenLoeve(num_eigenvectors), x_pts, y_pts)
	@time logKs = GaussianRandomFields.sample(grf)
	parameterization = copy(grf.data.eigenfunc)
	sigmas = copy(grf.data.eigenval)

	if i == 1#only plot this the first time around
		#plot a realization
		fig, ax = PyPlot.subplots()
		ax.imshow(logKs)
		ax.title.set_text("Random Conductivity Field")
		display(fig)
		println()
		PyPlot.close(fig)

		#plot the first few eigenvectors and some random ones
		fig, axs = PyPlot.subplots(2, 4, figsize=(16, 8))
		rp = Random.randperm(size(parameterization, 2))
		for (i, ax) in enumerate(axs)
			if i <= length(axs) / 2
				ax.imshow(reshape(parameterization[:, i], ns...))
				ax.title.set_text("Eigenvector $i")
			else
				ax.imshow(reshape(parameterization[:, rp[i]], ns...))
				ax.title.set_text("Eigenvector $(rp[i])")
			end
		end
		display(fig)
		println()
		PyPlot.close(fig)
	end

	#logKs2Ks_neighbors(Ks) = map(neighbor->exp(0.5 * (Ks[neighbor[1]] + Ks[neighbor[2]])), neighbors)
	#logKs2Ks_neighbors(Ks) = exp.([0.5 * (Ks[i1] + Ks[i2]) for (i1, i2) in neighbors])
	logKs2Ks_neighbors(Ks) = exp.(0.5 * (Ks[map(p->p[1], neighbors)] .+ Ks[map(p->p[2], neighbors)]))#Zygote differentiates this efficiently but the definitions above are ineffecient with Zygote
	function solveforh(logKs, dirichleths)
		@assert length(logKs) == length(Qs)
		if maximum(logKs) - minimum(logKs) > 25
			return fill(NaN, length(Qs))#this is needed to prevent the solver from blowing up if the line search takes us somewhere crazy
		else
			Ks_neighbors = logKs2Ks_neighbors(logKs)
			return DPFEHM.groundwater_steadystate(Ks_neighbors, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs)
		end
	end
	solveforheigs(x) = solveforh(x2logKs(x), dirichleths)
	x2logKs(x) = reshape(parameterization * (sigmas .* x), ns...)
	logKs_true = x2logKs(x_true)
	h_true = solveforheigs(x_true)

	numobs = 250
	obsnodes = getobsnodes(coords, obslocs)
	obssigma = 1e-3
	observations .= h_true[obsnodes]#set up the observations
	f(logKs) = sum((solveforh(logKs, dirichleths)[obsnodes] - observations) .^ 2 ./ obssigma ^ 2)
	h = solveforh(logKs, dirichleths)
	h = solveforh(logKs, dirichleths)
	print("forward solve time")
	@time h = solveforh(logKs_true, dirichleths)
	zg = Zygote.gradient(f, logKs)[1]
	zg = Zygote.gradient(f, logKs)[1]#work out the precompilation
	print("gradient time")
	@time zg = Zygote.gradient(f, logKs)[1]

	#now reformulate the function in terms of the cofficients of the eigenvectors with some regularization in there
	feigs(x) = f(x2logKs(x)) + sum(x .^ 2)
	print("eigs forward solve time")
	@time feigs(x_true)
	zgeigs = Zygote.gradient(feigs, x_true)[1]
	print("eigs gradient time")
	@time zgeigs = Zygote.gradient(feigs, x_true)[1]

	if i == 1#first time around plot a bunch of stuff and set up the observations
		#plot the solution, the difference between the solution and the solution for a uniform medium, and the logKs
		fig, axs = PyPlot.subplots(1, 4, figsize=(32, 8))
		ims = axs[1].imshow(reshape(h, ns...))
		axs[1].title.set_text("Head for Random Conductivity")
		fig.colorbar(ims, ax=axs[1])
		ims = axs[2].imshow(reshape(h .- map(i->boundaryhead(coords[:, i]...), 1:size(coords, 2)), ns...))
		axs[2].title.set_text("Head Fluctuation for Random Conductivity")
		fig.colorbar(ims, ax=axs[2])
		ims = axs[3].imshow(logKs)
		axs[3].title.set_text("Random Conductivity Field")
		fig.colorbar(ims, ax=axs[3])
		ims = axs[4].imshow(zg)
		axs[4].title.set_text("Gradient of Loss w.r.t. Conductivity")
		fig.colorbar(ims, ax=axs[4])
		fig.tight_layout()
		display(fig)
		println()
		PyPlot.close(fig)

		#plot the solution, the difference between the solution and the solution for a uniform medium, and the logKs
		fig, axs = PyPlot.subplots(1, 4, figsize=(32, 8))
		ims = axs[1].imshow(reshape(h_true, ns...))
		axs[1].title.set_text("Head for Random Conductivity")
		fig.colorbar(ims, ax=axs[1])
		ims = axs[2].imshow(reshape(h_true .- map(i->boundaryhead(coords[:, i]...), 1:size(coords, 2)), ns...))
		axs[2].title.set_text("Head Fluctuation for Random Conductivity")
		fig.colorbar(ims, ax=axs[2])
		ims = axs[3].imshow(x2logKs(x_true))
		axs[3].title.set_text("True Conductivity Field")
		fig.colorbar(ims, ax=axs[3])
		ims = axs[4].plot(zgeigs)
		axs[4].title.set_text("Gradient of Loss w.r.t. Eigenvector Coefficients")
		fig.tight_layout()
		display(fig)
		println()
		PyPlot.close(fig)

	else#otherwise, do the optimization
		options = Optim.Options(iterations=200, extended_trace=false, store_trace=true, show_trace=ifelse(i == length(allns), true, false), x_tol=1e-6, time_limit=60 * 60 * 2)
		print("Optimization time")
		@time opt = Optim.optimize(feigs, x->Zygote.gradient(feigs, x)[1], x0, Optim.LBFGS(), options; inplace=false)
		#@time opt = Optim.optimize(feigs, x->Zygote.gradient(feigs, x)[1], x0, Optim.GradientDescent(), options; inplace=false)
		x_est = opt.minimizer
		logKs_est = x2logKs(x_est)

		#plot results from the inverse analysis
		fig, axs = PyPlot.subplots(2, 3, figsize=(12, 8))
		ims = axs[1].imshow(logKs_true)
		axs[1].title.set_text("True Conductivity")
		fig.colorbar(ims, ax=axs[1])
		ims = axs[2].imshow(logKs_est)
		axs[2].title.set_text("Estimated Conductivity")
		fig.colorbar(ims, ax=axs[2])
		ims = axs[3].imshow(reshape(solveforh(logKs_true, dirichleths), ns...))
		axs[3].title.set_text("True Head")
		#ims = axs[3].imshow(x2logKs(x0))
		#axs[3].title.set_text("Previous Conductivity Estimate")
		fig.colorbar(ims, ax=axs[3])
		#ims = axs[4].imshow(logKs_est - x2logKs(x0))
		#axs[4].title.set_text("Conductivity Update")
		ims = axs[4].imshow(reshape(solveforh(logKs_est, dirichleths), ns...))
		axs[4].title.set_text("Estimated Head")
		fig.colorbar(ims, ax=axs[4])
		axs[5].plot([0, 5], [0, 5], "k", alpha=0.5)
		axs[5].plot(observations, solveforh(logKs_est, dirichleths)[obsnodes], "k.")
		axs[5].set_xlabel("Observed Head")
		axs[5].set_ylabel("Predicted Head")
		axs[6].semilogy(map(i->opt.trace[i].value, 1:length(opt.trace)))
		axs[6].set_xlabel("Iteration")
		axs[6].set_ylabel("Loss")
		fig.tight_layout()
		fig.savefig("figs/level_$(i - 1).pdf")
		display(fig)
		println()
		PyPlot.close(fig)

		x0 = x_est#update x0 for the next optimization
	end
end


using QuasiGeostrophy
using LinearAlgebra, Test, FFTW, BenchmarkTools, Plots

# boiler plate definitions
Ωxy = S¹(0,2π) × S¹(0,2π)
Nx = 2^5; Ny = 2^5;
fourier_grid = FourierGrid((Nx, Ny), Ωxy)
x, y = fourier_grid.grid
kx, ky = fourier_grid.wavenumbers
ψ1 = @. 0 * sin(x) - sin(y) + 0im
ψ2 = @. 0 * sin(x) - sin(y) + 0im
ψ̂1 = copy(ψ1)
ψ̂2 = copy(ψ2)
q̂1 = copy(ψ2)
q̂2 = copy(ψ2)
q1 = copy(ψ2)
q2 = copy(ψ2)
FFTW.set_num_threads(Threads.nthreads())
P = plan_fft(ψ1) # typeof(iP) <: AbstractFFTs.ScaledPlan
iP = plan_ifft(ψ1)
mul!(ψ̂1, P, ψ1)

## Define Parameters
const λ  = 1e-2
const U  = 1e-2
const ν  = 1e-6
const κ  = 1e-0
const Δt = 1e-1
const L  = 1.0
const Q  = 1.0

## Define forcing
forcing = @. 0 * sin(x) + Q * sin(y/L) + 0im
f̂ = copy(forcing)
mul!(f̂, P, forcing)

## Define Operators
∂x = FourierOperator(im .* kx)
∂y = FourierOperator(im .* ky)
Δ = ∂x^2 + ∂y^2
Δ⁵ = Δ^5
G = inv( (Δ + -1 /(2*λ^2) )^2 + -(1/(2*λ^2))^2 )

function get_pv!(q̂1, q̂2, ψ̂1, ψ̂2)
    tmp1 = Δ(ψ̂1) + 1 /(2*λ^2) * (ψ̂2 - ψ̂1)
    q̂1 .= tmp1
    tmp2 = Δ(ψ̂2) + 1 /(2*λ^2) * (ψ̂1 - ψ̂2)
    q̂2 .= tmp2
    return nothing
end

get_pv!(q̂1, q̂2, ψ̂1, ψ̂2)

function get_stream_function!(q1, q2, ψ1, ψ2)
    tmp1 = G(Δ(q1) + 1 /(2*λ^2) * (-q2 - q1))
    ψ1 .= tmp1
    tmp2 = G(Δ(q2) + 1 /(2*λ^2) * (-q1 - q2))
    ψ2 .= tmp2
    return nothing
end

function closure_J(∂x, ∂y, P, iP)
    function J(a,b)
        term_1 = filter_convolve(∂x(a) , ∂y(b), P, iP)
        term_2 = filter_convolve(∂x(b) , ∂y(a), P, iP)
        return  term_1 - term_2
    end
end

J = closure_J(∂x, ∂y, P, iP)

function closure_rhs(∂x, ∂y, Δ, Δ⁵, J, λ, ν, U, κ, forcing)
    function rhs(state)
        q1 = state[1]
        q2 = state[2]
        ψ1 = copy(q1)
        ψ2 = copy(q2)
        get_stream_function!(q1, q2, ψ1, ψ2)
        drag = - 2 .* κ .* (Δ(ψ2))
        q1_rhs = -J(ψ1, q1) - (U/λ^2) .* ∂x(ψ1) - U .* ∂x(q1) + forcing
        q2_rhs = -J(ψ2, q2) + (U/λ^2) .* ∂x(ψ2) + U .* ∂x(q2) - forcing + drag
        return [q1_rhs, q2_rhs]
    end
end

rhs = closure_rhs(∂x, ∂y, Δ, Δ⁵, J, λ, ν, U, κ, f̂)

function closure_lhs(∂x, ∂y, Δ, Δ⁵, J, λ, ν, U, Δt)
    function lhs()
        q1_lhs = 1 + Δt * ν * Δ^4
        return [q1_lhs, q1_lhs]
    end
end

lhs = closure_lhs(∂x, ∂y, Δ, Δ⁵, J, λ, ν, U, Δt)
operators = lhs()
lhs_operator = [inv(operator) for operator in operators]

function imex_step!(state, lhs_operator, rhs, Δt)
    rhs_state = rhs(state)
    current_state = copy(state)
    for i in eachindex(state)
        tmp = lhs_operator[i](current_state[i] .+ Δt .* rhs_state[i])
        @. state[i] =  tmp
    end
end
# @btime imex_step!(state, lhs_operator, rhs, Δt);

##
q1 = @. 1 * sin(x) - cos(y) + 0im
q2 = @. 1 * sin(x) + sin(y) + 0im
fq1 = P * q1
fq2 = P * q2
τ = (q1 - q2) ./ 2
τ̂ = (fq1 - fq2) ./ 2
tmp = P * (400 * sin.(x) .+ 0 * sin.(y))
#fq1 .= good_start1
#fq2 .= good_start2
state = [fq1, fq2]
##
gr(size = (400,400))
for i in 1:2000
    imex_step!(state, lhs_operator, rhs, Δt)
    if (i%100) == 0
        @. τ̂ = (fq1 - fq2) / 2
        mul!(τ, iP, τ̂)
        local p1 = contourf(x[:], y[:], real.(τ)', linewidth = 0, color = :thermometer)
        local p2 = plot(sum(real.(τ), dims = 1)[:] ./ Nx, x[:] )
        display(plot(p1,p2))
        sleep(0.01)
    end
end
#good_start1 = copy(fq1)
#good_start2 = copy(fq2)
##

@. τ̂ = (fq1 - fq2) / 2
mul!(τ, iP, τ̂)

##

# Check PV and Stream Function Calculation
tmp = @. 1 * sin(2*x) + Q * sin(y/L * 3) + 0im
tmp2 = @. 1 * sin(x) - Q * sin(y/L) + 0im
ψ̂1 = P * tmp
ψ̂2 = P * tmp2
q1 = copy(q̂1)
q2 = copy(q̂2)
ψ1 = copy(ψ̂1)
ψ2 = copy(ψ̂2)

get_pv!(q̂1, q̂2, ψ̂1, ψ̂2)
get_stream_function!(q̂1, q̂2, ψ̂1, ψ̂2)
norm(ψ1 - ψ̂1)
norm(ψ2 - ψ̂2)
##
i = 1
@. state[i] *= lhs_operator[i].op


using LinearAlgebra

mutable struct Broyden{R<:Real,C<:Union{R,Complex{R}},T<:AbstractArray{C}}
    H
    theta_bar::R
    function Broyden{R,C,T}(x::T, H, theta_bar) where {R,C,T}
        new(H, theta_bar)
    end
end

Broyden(
    x::T;
    H = I,
    theta_bar = R(0.2),
) where {R<:Real,C<:Union{R,Complex{R}},T<:AbstractArray{C}} =
    Broyden{R,C,T}(x, H, theta_bar)

_sign(x::R) where {R} = x == 0 ? R(1) : sign(x)

function update!(L::Broyden{R,C,T}, s, y) where {R,C,T}
    Hy = L.H * y
    sH = s' * L.H
    delta = dot(Hy, s) / norm(s)^2
    theta = if abs(delta) >= L.theta_bar
        R(1)
    else
        (1 - _sign(delta) * L.theta_bar) / (1 - delta)
    end
    L.H += (s - Hy) / dot(s, (1 / theta - 1) * s + Hy) * sH
end

import Base: *

function (*)(L::Broyden, v)
    w = similar(v)
    mul!(w, L, v)
end

import LinearAlgebra: mul!

mul!(d::T, L::Broyden, v::T) where {T} = mul!(d, L.H, v)


module FindProc

using JuMP
using Cbc

include("Format.jl")
using .Format

export findproc

B = []

# Determines procedure with minimum piece size alpha -- Work in progress
# Author: Jason Liu
function findproc(m::Int64, s::Int64, alpha::Rational{Int64}; output::Int64 = 2)
    # Prepare variables for solving
    c = numerator(alpha)
    d = denominator(alpha)
    b = lcm(s, d)

    a = Int64(b * c / d)
    global B = [x for x = a:(b-a)]

    M = []
    S = []
    procedures = []

    # Display header
    alphaF = formatFrac(alpha)
    if output > 0
        printHeader(center("FIND PROCEDURE"))
        printfT(
            "Overview",
            "Find procedures for dividing $m muffins among $s students where $alphaF is the smallest piece size",
            "",
            "Let the common denominator be $b",
        )
    end

    # Exit early if not solvable via Cbc
    if m <= 0 || s <= 0 || alpha <= 0 || alpha > 1
        if output > 0
            printf("Find Procedure does not apply", line = true)
            printEnd()
        end
        return [b, Nothing]
    elseif m % s == 0 && toFrac(alpha) == 1
        b = 1
        give = repeat([1], Int64(m / s))
        if output > 0
            printfT(
                "Procedure",
                "Divide $m muffins {1}",
                "Give $s students {$(join(give, ","))}",
            )
            printEnd()
        end
        return [b, [[[[m, [1]]], [[s, give]]]]]
    end

    try
        # Determine vector sets
        L = b - 2a + 1
        V = Int64(ceil(2m / s))
        T = Int64(b * m / s)
        M = vectorize(unionF(f(L, b, 2), f(L, b, 3)))
        S = vectorize(unionF(f(L, T, V), f(L, T, V - 1)))

        # Initialize model
        model = Model(Cbc.Optimizer)
        set_optimizer_attribute(model, "logLevel", 0)
        set_optimizer_attribute(model, "allowableGap", 1e-5)
        set_optimizer_attribute(model, "cuts", "off")

        @variable(model, 0 <= x[i = 1:length(M)], Int)
        @variable(model, 0 <= y[i = 1:length(S)], Int)

        for k = 1:L
            @constraint(model, sum(x .* getindex.(M, k)) == sum(y .* getindex.(S, k)))
        end

        @constraint(model, sum(x) == m)
        @constraint(model, sum(y) == s)

        # Solve for a natural number solutions
        for i = 1:min(length(M), 3)
            @objective(model, Min, x[i])
            optimize!(model)
            append!(procedures, [[round.(value.(x)), round.(value.(y))]])

            @objective(model, Max, x[i])
            optimize!(model)
            append!(procedures, [[round.(value.(x)), round.(value.(y))]])
        end

        procedures = unique(procedures)
    catch
        Nothing
    end

    # Exit if no procedures
    if length(procedures) == 0
        if output > 0
            printf("No procedures for muffins($m, $s, $alphaF)", line = true)
            printEnd()
        end
        return [b, Nothing]
    end

    # Output each procedure
    M = [unpack(M[i]) for i = 1:length(M)]
    S = [unpack(S[i]) for i = 1:length(S)]
    formatProc = []
    for k = 1:length(procedures)
        x = procedures[k][1]
        y = procedures[k][2]

        divide = [
            "Divide $(Int64(x[i])) muffins {$(join(M[i], ","))}"
            for i = 1:length(M) if x[i] > 0
        ]
        give = [
            "Give $(Int64(y[i])) students {$(join(S[i], ","))}"
            for i = 1:length(S) if y[i] > 0
        ]
        append!(
            formatProc,
            [[
                [[Int64(x[i]), M[i]] for i = 1:length(M)],
                [[Int64(y[i]), S[i]] for i = 1:length(S)],
            ]],
        )

        if output > 1
            printfT("Procedure $k", divide..., give...)
        elseif output > 0
            printfT("Procedure", divide..., give...)
            printEnd()
            output = 0
        end
    end
    output > 0 && printEnd()

    [b, formatProc]
end

# Helper function for findproc -- "vectorizes" arrays in set S based on B
function vectorize(S)
    S == Nothing ? Nothing :
    sort([[count(i -> (i == x), s) for x in B] for s in S], rev = true)
end

# Helper function for findproc -- "unpacks" vectors into strings based on B
function unpack(V)
    vcat(map(i -> repeat([B[i]], V[i]), 1:length(V))...)
end

# Helper function for findproc -- recursively determines k-size subsets of B[1:i] that sum to T
function f(i, T, k, memo = Dict())
    # Base cases
    if i >= 1 && T == 0 && k == 0
        Set()
    elseif T == 0 && k >= 1
        Nothing
    elseif i == 0 && T >= 1
        Nothing
    elseif T >= 1 && k == 0
        Nothing
    elseif T <= -1
        Nothing
    elseif i == 1
        T == k * B[1] ? Set([repeat([B[1]], k)]) : Nothing
    elseif i >= 1 && k == 1
        index = findall(x -> x == T, B[1:i])
        length(index) > 0 ? Set([[B[index[1]]]]) : Nothing
        # Recursive case with memoization
    else
        if haskey(memo, (i, T, k))
            get(memo, (i, T, k), 0)
        else
            set = unionF(f(i - 1, T, k, memo), mapCat(f(i, T - B[i], k - 1, memo), [B[i]]))
            get!(memo, (i, T, k), set)
        end
    end
end

# Helper function for f -- set unions with Nothing
function unionF(x, y)
    if x == Nothing && y == Nothing
        Nothing
    elseif x == Nothing
        y
    elseif y == Nothing
        x
    else
        union(x, y)
    end
end

# Helper function for f -- concatenate e with each element in set S
function mapCat(S, e)
    S == Nothing ? Nothing : Set([vcat(s, e) for s in S])
end

end


# These functions are only conditionally loaded with Clustering.jl
# Code adapted from @cormullion's [ColorSchemeTools](https://github.com/JuliaGraphics/ColorSchemeTools.jl).
function get_colorscheme(
    img,
    ncolors;
    maxiter=Clustering._kmeans_default_maxiter,
    tol=Clustering._kmeans_default_tol,
)
    # Cluster in Lab color space
    data = reshape(channelview(Lab.(img)), 3, :)
    R = Clustering.kmeans(data, ncolors; maxiter=maxiter, tol=tol)

    # Make color scheme out of cluster centers
    cs = Lab{Float64}[]
    for i in 1:3:length(R.centers)
        push!(cs, Lab(R.centers[i], R.centers[i + 1], R.centers[i + 2]))
    end
    return cs
end

function _colordither(
    ::Type{T},
    img,
    alg,
    ncolors::Int;
    maxiter=Clustering._kmeans_default_maxiter,
    tol=Clustering._kmeans_default_tol,
    kwargs...,
) where {T}
    cs = get_colorscheme(img, ncolors; maxiter=maxiter, tol=tol)
    return _colordither(T, img, alg, cs; kwargs...)
end

"""
    dither!([out,] img, alg::AbstractDither, ncolors; maxiter, tol, kwargs...)

Dither image `img` using algorithm `alg`.
A color palette with `ncolors` is computed by Clustering.jl's K-means clustering.
The amount of `maxiter` and tolerance `tol` default to those exported by Clustering.jl.
"""
dither!(img, alg::AbstractDither, ncolors::Int; kwargs...)

"""
    dither([T::Type,] img, alg::AbstractDither, ncolors; maxiter, tol, kwargs...)

Dither image `img` using algorithm `alg`.
A color palette with `ncolors` is computed by Clustering.jl's K-means clustering.
The amount of `maxiter` and tolerance `tol` default to those exported by Clustering.jl.
"""
dither(::Type, img, alg::AbstractDither, ncolors::Int; kwargs...)


#=
Response surfaces implementantion, following:
"A Taxonomy of Global Optimization Methods Based on Response Surfaces"
by DONALD R. JONES
=#
mutable struct RadialBasis{F,X,Y,B,C} <: AbstractSurrogate
    phi::F
    dim_poly::Int
    x::X
    y::Y
    bounds::B
    coeff::C
end

"""
Calculates current estimate of array value 'val' with respect to RadialBasis object.
"""
function (rad::RadialBasis)(val)
    n = length(rad.x)
    d = length(rad.x[1])
    q = rad.dim_poly
    central_point = zeros(eltype(rad.x[1]), d)
    sum = zero(eltype(rad.x[1]))
    @inbounds for i = 1:d
        central_point[i] = (rad.bounds[i][1]+rad.bounds[i][2])/2
        sum += (rad.bounds[i][2]-rad.bounds[i][1])/2
    end
    half_diameter_domain = sum/d
    approx = zero(eltype(rad.x[1]))
    for i = 1:n
        approx = approx + rad.coeff[i]*rad.phi(val .- rad.x[i])
    end
    for i = n+1:n+q
        approx = approx + rad.coeff[i]*centralized_monomial(val,n+1-i,half_diameter_domain,central_point)
    end
    return approx
end

"""
Calculates current estimate of value 'val' with respect to RadialBasis object.
"""
function (rad::RadialBasis)(val::Number)
    approx = zero(eltype(rad.x[1]))
    n = length(rad.x)
    q = rad.dim_poly
    Chebyshev(x,k) = cos(k*acos(-1 + 2/(rad.bounds[2]-rad.bounds[1])*(x.-rad.bounds[1])))
    for i = 1:n
        approx = approx + rad.coeff[i]*rad.phi(val .- rad.x[i])
    end
    for i = n+1:n+q
        approx = approx + rad.coeff[i]*Chebyshev(val,n+1-i)
    end
    return approx
end

#=
linear_basis_function = Basis(z->norm(z), 1)
cubic_basis_function = Basis(z->norm(z)^3, 2)
thinplate_basis_function = Basis(z->norm(z)^2*log(norm(z)),2)
function multiquadric_basis_function(lambda)
    return Basis(z->sqrt(norm(z)^2 + lambda^2),1)
end
=#

"""
    RadialBasis(x,y,a::Number,b::Number,phi::Function,q::Int)

Constructor for RadialBasis surrogate
- (x,y): sampled points
- '(a,b)': interval of interest
-'phi': radial basis of choice
-'q': number of polynomial elements
"""
function RadialBasis(x,y,a::Number,b::Number,phi::Function,q::Int)
    coeff = _calc_coeffs(x,y,a,b,phi,q)
    RadialBasis(phi,q,x,y,(a,b),coeff)
end

function _calc_coeffs(x,y,a,b,phi,q)
    Chebyshev(x,k) = cos(k*acos(-1 + 2/(b-a)*(x-a)))
    n = length(x)
    size = n+q
    D = zeros(eltype(x[1]), size, size)
    d = zeros(eltype(x[1]),size)

    @inbounds for i = 1:n
        d[i] =  y[i]
        for j = 1:n
            D[i,j] = phi(x[i] .- x[j])
        end
        if i < n + 1
            for k = n+1:size
                D[i,k] = Chebyshev(x[i],k)
            end
        end
    end
    Sym = Symmetric(D,:U)
    coeff = Sym\d
end

"""
RadialBasis(x,y,bounds,phi::Function,q::Int)

Constructor for RadialBasis surrogate

- (x,y): sampled points
- bounds: region of interest of the form [[a,b],[c,d],...,[w,z]]
- phi: radial basis of choice
- q: number of polynomial elements
"""
function RadialBasis(x,y,bounds,phi::Function,q::Int)
    coeff = _calc_coeffs(x,y,bounds,phi,q)
    RadialBasis(phi,q,x,y,bounds,coeff)
end

function _calc_coeffs(x,y,bounds,phi,q)
    n = length(x)
    d = length(x[1])
    central_point = zeros(eltype(x[1]), d)
    sum = zero(eltype(x[1]))
    @inbounds for i = 1:d
        central_point[i] = (bounds[i][1]+bounds[i][2])/2
        sum += (bounds[i][2]-bounds[i][1])/2
    end
    half_diameter_domain = sum/d

    size = n+q
    D = zeros(eltype(x[1]), size, size)
    d = zeros(eltype(x[1]),size)

    @inbounds for i = 1:n
        d[i] =  y[i]
        for j = 1:n
            D[i,j] = phi(x[i] .- x[j])
        end
        if i < n + 1
            for k = n+1:size
                D[i,k] = centralized_monomial(x[i],k,half_diameter_domain,central_point)
            end
        end
    end
    Sym = Symmetric(D,:U)
    coeff = Sym\d
end

"""
    centralized_monomial(vect,alpha,half_diameter_domain,central_point)

Returns the value at point vect[] of the alpha degree monomial centralized.

#Arguments:
-'vect': vector of points i.e [x,y,...,w]
-'alpha': degree
-'half_diameter_domain': half diameter of the domain
-'central_point': central point in the domain
"""
function centralized_monomial(vect,alpha,half_diameter_domain,central_point)
    mul = 1
    @inbounds for i = 1:length(vect)
        mul *= vect[i]
    end
    return ((mul-norm(central_point))/(half_diameter_domain))^alpha
end

"""
    add_point!(rad::RadialBasis,new_x,new_y)

Add new samples x and y and updates the coefficients. Return the new object radial.
"""
function add_point!(rad::RadialBasis,new_x,new_y)
    if (length(new_x) == 1 && length(new_x[1]) == 1) || ( length(new_x) > 1 && length(new_x[1]) == 1 && length(rad.bounds[1])>1)
        push!(rad.x,new_x)
        push!(rad.y,new_y)
        if length(rad.bounds[1]) == 1
            rad.coeff = _calc_coeffs(rad.x,rad.y,rad.bounds[1],rad.bounds[2],rad.phi,rad.dim_poly)
        else
            rad.coeff = _calc_coeffs(rad.x,rad.y,rad.bounds,rad.phi,rad.dim_poly)
        end
    else
        append!(rad.x,new_x)
        append!(rad.y,new_y)
        if length(rad.bounds[1]) == 1
            rad.coeff = _calc_coeffs(rad.x,rad.y,rad.bounds[1],rad.bounds[2],rad.phi,rad.dim_poly)
        else
            rad.coeff = _calc_coeffs(rad.x,rad.y,rad.bounds,rad.phi,rad.dim_poly)
        end
    end
    nothing
end


include("../ds.jl")

struct compact_mnm_hyper <: distribution_hyper_params
    α::AbstractArray{Float64,1}
end

mutable struct compact_mnm_sufficient_statistics <: sufficient_statistics
    N::Float64
    points_sum::AbstractArray{Int64,1}
end

function calc_posterior(prior:: compact_mnm_hyper, suff_statistics::compact_mnm_sufficient_statistics)
    if suff_statistics.N == 0
        return prior
    end
    return compact_mnm_hyper(prior.α + suff_statistics.points_sum)
end

function sample_distribution(hyperparams::compact_mnm_hyper)
    cat_dist = rand(Dirichlet(hyperparams.α))
    return compact_mnm_dist(log.(cat_dist))
end

function create_sufficient_statistics(hyper::compact_mnm_hyper,posterior::compact_mnm_hyper,points::AbstractArray{Float64,2}, pts_to_group = 0)
    if length(points) == 0
        return compact_mnm_sufficient_statistics(size(points,2),zeros(Int64,length(hyper.α)))
    end
    pt_count = counts(Int.(points),length(hyper.α))
    return compact_mnm_sufficient_statistics(size(points,2),pt_count)
end

function log_multivariate_gamma(x::Number, D::Number)
    res::Float64 = D*(D-1)/4*log(pi)
    for d = 1:D
        res += logabsgamma(x+(1-d)/2)[1]
    end
    return res
end

function log_marginal_likelihood(hyper::compact_mnm_hyper, posterior_hyper::compact_mnm_hyper, suff_stats::compact_mnm_sufficient_statistics)
    D = length(suff_stats.points_sum)
    logpi = log(pi)
    val = logabsgamma(sum(hyper.α))[1] -logabsgamma(sum(posterior_hyper.α))[1] + sum((x-> logabsgamma(x)[1]).(posterior_hyper.α) - (x-> logabsgamma(x)[1]).(hyper.α))
    return val
end

function aggregate_suff_stats(suff_l::compact_mnm_sufficient_statistics, suff_r::compact_mnm_sufficient_statistics)
    return compact_mnm_sufficient_statistics(suff_l.N+suff_r.N, suff_l.points_sum + suff_r.points_sum)
end


# Tested and developed in Windows with Julia 1.5.1, October 2020.
# 2-D Topoplogy Optimization in Julia
# IMPORT PACKAGES
using Printf
using SparseArrays
using LinearAlgebra
using Plots: heatmap
# FUNCTIONS
function optimize!(F::Array{Float64,1}, U::Array{Float64,1},
                  freedofs::Array{Int64,1},
                  penal::Float64, E0::Float64, Emin::Float64, ft::Int64,
                  nelx::Int64, nely::Int64,
                  iK::Array{Int64,1},
                  jK::Array{Int64,1},
                  edofMat::Array{Int64,2},
                  KE::Array{Float64,2},
                  H::SparseMatrixCSC{Float64,Int64},
                  Hs::Array{Float64,2}, volfrac::Float64)::Array{Float64,2}
    x = repeat([volfrac],nely,nelx)
    xPhys = copy(x)
    xnew = copy(x)
    loop = 0
    change = 1.0
    sK = zeros(64*nelx*nely)
    K = spzeros(2*(nely+1)*(nelx+1),2*(nely+1)*(nelx+1))
    dc = zeros(nely, nelx)
    dv = zeros(nely, nelx)
    ce = zeros(nely, nelx)
    while change > 0.01
      loop = loop +1
      # FE-ANALYSIS
      sK = reshape(vec(KE)*(Emin.+vec(xPhys)'.^penal*(E0-Emin)),64*nelx*nely)
      K = sparse(iK,jK,sK)
      U[freedofs] .= Symmetric(K[freedofs,freedofs])\F
      # OBJECTIVE FUNCTION AND SENSITIVITY ANALYSIS
      ce = reshape(sum((U[edofMat]*KE).*U[edofMat], dims=2),nely,nelx)
      c = sum((Emin.+xPhys.^penal.*(E0-Emin)).*ce)
      dc = -penal*(E0-Emin).*xPhys.^(penal-1).*ce
      dv = ones(nely,nelx)
      # FILTERING/MODIFICATION OF SENSITIVITIES
      if ft == 1 # sensitivity filtering
        dc = reshape(H*(vec(x).*vec(dc))./Hs./max.(1e-3,vec(x)),nely,nelx)
      elseif ft == 2 # density filtering
        dc = reshape(H*(vec(dc)./Hs),nely,nelx)
        dv = reshape(H*(vec(dv)./Hs),nely,nelx)
      end
      # OPTIMALITY CRITERIA UPDATE OF DESIGN VARIABLES AND PHYSICAL DENSITIES
      xPhys, xnew = optimalityCriteria!(x, xPhys, xnew, dc, dv, H, Hs, nelx, nely, ft, volfrac)
      change = maximum(abs.(vec(xnew) - vec(x)))
      x = copy(xnew)
      s = @sprintf "It: %i  Obj.: %.2f  Ch: %.3f" loop c change;
      println(s)
    end
    return xPhys
end
function optimalityCriteria!(x::Array{Float64,2}, xPhys::Array{Float64,2}, xnew::Array{Float64,2},
                            dc::Array{Float64,2}, dv::Array{Float64,2}, H::SparseMatrixCSC{Float64,Int64},
                            Hs::Array{Float64,2}, nelx::Int64, nely::Int64, ft::Int64,
                            volfrac::Float64)::Tuple{Array{Float64,2},Array{Float64,2}}
    l1 = 0.0; l2 = 1.0e9; move = 0.2
    while (l2-l1)/(l1+l2) > 1.0e-3
      lmid = 0.5*(l2+l1)
      xnew = max.(0.0,max.(x.-move,min.(1.0,min.(x.+move,x.*sqrt.(-dc./dv./lmid)))))
      if ft == 1
        xPhys = copy(xnew)
      elseif ft == 2
        xPhys = reshape((H*vec(xnew))./Hs,nely,nelx)
      end
      if (sum(xPhys) > volfrac*nelx*nely) l1 = lmid else l2 = lmid end
    end
    return xPhys, xnew
end
function prepFilter(nelx::Int64, nely::Int64, rmin::Int64)::Tuple{SparseMatrixCSC{Float64,Int64}, Array{Float64,2}}
  iH = ones(nelx*nely*(2*(ceil(rmin)-1)+1)^2)
  jH = ones(size(iH))
  sH = zeros(size(iH))
  k = 0
  for i1 in 1:nelx
    for j1 in 1:nely
      e1 = (i1-1)*nely+j1
      for i2 in max(i1-(ceil(rmin)-1),1):min(i1+(ceil(rmin)-1),nelx)
        for j2 in max(j1-(ceil(rmin)-1),1):min(j1+(ceil(rmin)-1),nely)
          e2 = (i2-1)*nely+j2
          k = k+1
          iH[k] = e1
          jH[k] = e2
          sH[k] = max(0.0,rmin-sqrt((i1-i2)^2+(j1-j2)^2))
        end
      end
    end
  end
  iH = Int.(iH)
  jH = Int.(jH)
  H = sparse(iH,jH,sH)
  Hs = sum(H, dims=2)
  return (H,Hs)
end
################################################################################
# MAIN CODE
function TopOptOC(nelx = 150::Int64, nely = 50::Int64,
                  rmin = 6::Int64, volfrac = 0.5::Float64,
                  penal = 3.0::Float64, ft = 1::Int64)
  # TOPOPT SETTINGS
  # nelx, elements in x
  # nely, elements in y
  # rmin, filter radius
  # volfrac, volume fraction
  # penal, SIMP penalty
  # ft, filtering option: 1 = sensitivity filter, 2 density filter
  nele = nelx * nely
  # MATERIAL PROPERTIES
  E0 = 1.0
  Emin = 1.0e-9
  nu = 0.3
  # ELEMENT STIFFNESS MATRIX
  A11 = Float64[12 3 -6 -3; 3 12 3 0; -6 3 12 -3; -3 0 -3 12.0]
  A12 = Float64[-6 -3 0 3; -3 -6 -3 -6; 0 -3 -6 3; 3 -6 3 -6.0]
  B11 = Float64[-4 3 -2 9; 3 -4 -9 4; -2 -9 -4 -3; 9 4 -3 -4.0]
  B12 = Float64[2 -3 4 -9; -3 2 9 -2; 4 9 2 3; -9 -2 3 2.0]
  KE = 1/(1-nu^2)/24 * ([A11 A12; A12' A11] + nu*[B11 B12; B12' B11])
  # CONNECTIVITY
  nodenrs = reshape(1:(1+nelx)*(1+nely), 1+nely, 1+nelx)
  edofVec = reshape(2*nodenrs[1:end-1,1:end-1]+ones(nely,nelx), nelx*nely, 1)
  edofMat = Int.(repeat(edofVec, 1, 8) + repeat([0 1 2*nely .+ [2 3 0 1] -2 -1], nelx*nely, 1))
  iK = reshape(kron(edofMat, ones(8,1))', 64*nelx*nely)
  jK = reshape(kron(edofMat, ones(1,8))', 64*nelx*nely)
  iK = convert(Array{Int64}, iK)
  jK = convert(Array{Int64}, jK)
  # DEFINE LOADS AND SUPPORTS (HALF MBB-BEAM)
  F = zeros(2*(nely+1)*(nelx+1))
  F[2,1] = -1.0
  U = zeros(2*(nely+1)*(nelx+1))
  fixeddofs = union(collect(1:2:2*(nely+1)), [2*(nelx+1)*(nely+1)])
  alldofs = collect(1:2*(nely+1)*(nelx+1))
  freedofs = setdiff(alldofs,fixeddofs)
  # PREPARE FILTER
  H,Hs = prepFilter(nelx,nely,rmin)
  # INITIALIZE ITERATION
  x = Array{Float64}(undef, nely, nelx)
  x = optimize!(F[freedofs], U, freedofs, penal, E0, Emin, ft, nelx, nely, iK, jK, edofMat, KE, H, Hs, volfrac)
  # PLOT FINAL DESIGN
  heatmap(1.0.-x[end:-1:1,:],  yaxis=false, xaxis=false, legend = :none,
  color = :greys, grid=false, border=nothing, aspect_ratio=:equal)
end
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# This code was written by Nicolò Pollini,                                %
# Department of Wind Energy,                                              %
# Technical University of Denmark.                                        %
#                                                                         %
# Contact: nicolo@alumni.technion.ac.il                                   %
#                                                                         %
# Code repository: https://github.com/pollinico/topopt_jl                 %
#                                                                         %
# Disclaimer:                                                             %
# The author reserves all rights but does not guarantee that the code is  %
# free from errors. Furthermore, the author shall not be liable in any    %
# event caused by the use of the program.                                 %
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


# This file is a part of TextSearch.jl

#####
using CategoricalArrays
export EntropyWeighting

"""
    EntropyWeighting(; smooth=0.0, lowerweight=0.0, weights=:balance)

Entropy weighting uses the empirical entropy of the vocabulary along classes to produce a notion of importance for each token
"""
struct EntropyWeighting <: GlobalWeighting end

function entropy(dist)
    e = 0.0
    ipop = 1/sum(dist)

    for x in dist
        p = x * ipop

        if p > 0.0
            e -= p * log2(p)
        end
    end

    e
end

"""
    VectorModel(ent::EntropyWeighting, lw::LocalWeighting, corpus::BOW, labels;
        mindocs::Integer=1,
        smooth::Float64=0.0,
        weights=:balance,
        lowerweight=0.0
    )

Creates a vector model using the input corpus.
"""
function VectorModel(ent::EntropyWeighting, lw::LocalWeighting, corpus::AbstractVector{BOW}, labels=nothing;
        mindocs=1,
        smooth::Float64=0.0,
        weights=:balance
    )
    labels === nothing && error("EntropyWeighting requires labels as a categorical array to work")
    nclasses = length(levels(labels))
    D = Dict{UInt64,Vector{Float64}}()
    V = Vocabulary()

    for i in eachindex(corpus)
        vec = corpus[i]
        code = labels.refs[i]
        #M = sum(x for x in values(vec))
        for (t, occs) in vec
            s = get(V, t, UnknownTokenStats)
            V[t] = TokenStats(s.occs + occs, s.ndocs + 1, 0f0)
            dist = get(D, t, nothing)
            if dist === nothing
                D[t] = dist = fill(smooth, nclasses)
            end

            dist[code] += 1  # occs/M # log2(1 + occs)
        end
    end

    weights = weights isa String ? Symbol(weights) : weights

    if weights === :balance
        weights = sum(Vector{Float64}, values(D))
        weights .= sum(weights) ./ weights
    elseif weights === :none
        weights = ones(Float64, nclasses)
    end

    tokens, maxfreq = _create_vocabulary_with_entropy(V, D, weights, nclasses, mindocs)

    VectorModel(ent, lw, tokens, maxfreq, length(tokens), length(corpus))
end

function _create_vocabulary_with_entropy(V, D, weights, nclasses, mindocs)
    tokens = Vocabulary()
    tokenID = 0
    maxfreq = 0
    maxent = log2(nclasses)

    for (t, s) in V
        s.ndocs < mindocs && continue
        dist = D[t]
        dist .= dist .* weights
        e = 1.0 - entropy(dist) / maxent
        tokens[t] = TokenStats(s.occs, s.ndocs, e)
        maxfreq = max(maxfreq, s.occs)
    end

    tokens, maxfreq
end

global_weighting(::VectorModel{EntropyWeighting}, s::TokenStats) = s.weight


#=
-------------------------------------------------------------------------------------------------------------
This file runs the 1D Shallow Water Equations using a unified CG/DG method using the AGGP storage
introduced in F.X. Giraldo's Introduction to Element-based Galerkin Methods using
Tensor-Product Bases: Analysis, Algorithms, and Applications.
Note that it only requires the storage of a global mass matrix, inside of TI_LSRK,
the CREATE_RHS routine builds the RHS without using global matrices.

The strategy followed is very similar to that described in Alg. 7.6 except that we do not build the
Flux matrix explicitly. This is all done in CREATE_RHS.jl.

Written by F.X. Giraldo on July 6, 2021.
Department of Applied Mathematics
Naval Postgraduate School
Monterey, CA 93943

For time-integration, we use LSRK45 (4th order, 5-stage).
The linearized shallow water equations run indefinitely using either CG or DG. For the nonlinear
equations, we would need some form of stabilization.

The interpolation points used are the following:
ipoints=1: Lobatto
ipoints=2: Legendre
ipoints=3: Chebyshev
ipoints=4: Equi-spaced

The integration points used are the following:
qpoints=1: Lobatto
qpoints=2: Legendre
-------------------------------------------------------------------------------------------------------------
=#

using Plots
using LinearAlgebra

#Include all files used here
include("QuadraturePoints.jl")
include("element_matrices.jl")
include("create_grid.jl")
include("global_matrices.jl")
include("exact_solution.jl")
include("compute_norms.jl")
include("ti_LSRK.jl")

#Some Constants
DFloat = Float64
machine_zero=eps(DFloat)

#Main function
function main()

    #-----------------------------Only Change these Input parameters---------------------------------#
    Ne=20
    N=4
    Q=N
    ipoints=1
    qpoints=1
    #space_method="CG" #time_final=0.4
    space_method="DG" #time_final=0.4
    case=1 #1=exponential, 2=square wave
    Δ_NL=0
    time_final=DFloat(1.0)
    plot_movie=true
    plot_solution=false
    #-----------------------------Only Change these Input parameters---------------------------------#

    Np=N+1
    Nq=Q+1
    Npoin_dg=Ne*Np
    Npoin_cg=Ne*N + 1
    dt=DFloat(0.0001)
    Courant_max=0.5
    ntime = ceil(Int64, time_final / dt)
    @show(N,Q,Ne,Npoin_cg,Npoin_dg,dt,ntime,case)

    #Select Interpolation Points
    if ipoints == 1
        (ξ,ω) = QuadraturePoints.lobatto_gauss(Np,DFloat) #Lobatto
    elseif ipoints == 2
        (ξ,ω) = QuadraturePoints.legendre_gauss(Np,DFloat) #Legendre
    elseif ipoints == 3
        (ξ,ω) = QuadraturePoints.chebyshev_gauss(Np,DFloat) #Chebyshev
    elseif ipoints == 4
        (ξ,ω) = QuadraturePoints.equispaced_points(Np,DFloat) #Equi-spaced
    end

    #Select Quadrature Points
    if qpoints == 1
        (ξq,ωq) = QuadraturePoints.lobatto_gauss(Nq,DFloat) #Lobatto
    elseif qpoints == 2
        (ξq,ωq) = QuadraturePoints.legendre_gauss(Nq,DFloat) #Legendre
    end

    #Construct Lagrange Polynomials
    (ψ,dψ) = QuadraturePoints.lagrange_basis(Np,Nq,ξ,ξq,DFloat)

    #Construct Local Mass and Differentiation Matrices
    (Me,De,Dwe,Fe) = element_matrices(ψ,dψ,Np,Nq,ωq,DFloat)

    #Construct Grid and (i,e)->I map
    (coord_cg,coord_dg,intma_cg,intma_dg,periodicity_cg,periodicity_dg) = create_grid(Np,Npoin_cg,Npoin_dg,Ne,ξ,DFloat)

    #CG or DG?
    if (space_method == "CG")
        Npoin=Npoin_cg
        coord=copy(coord_cg)
        intma=copy(intma_cg)
        periodicity=copy(periodicity_cg)
    elseif (space_method == "DG")
        Npoin=Npoin_dg
        coord=copy(coord_dg)
        intma=copy(intma_dg)
        periodicity=copy(periodicity_dg)
    end
    @show(space_method,Npoin,Ne,N,Q)

    #Construct Global Mass and Differentiation Matrices (only need Mass)
    (M,~,~,~) = global_matrices(Me,De,Dwe,Fe,intma,coord,Npoin,Ne,Np,periodicity,DFloat)

    #Compute Initial Condition
    time=DFloat(0.0)
    (qe,qb,gravity) = exact_solution(coord,Npoin,time,case,DFloat)
    q0=copy(qe)
    dx_min=coord[2]-coord[1]

    #Recompute DT, NTIME
    umax=DFloat(0.0)
    for i=1:Npoin
        h=q0[i,1]+qb[i]
        u=q0[i,2]/h
        umax=max(umax, abs(u) + sqrt(0.5*h^2) )
    end
    dt=Courant_max*dx_min/umax
    courant=umax*dt/dx_min
    ntime = ceil(Int64, time_final / dt)
    println("Courant = ",courant," dt = ",dt," ntime = ",ntime)

    #Time Integration
    #-----------------------------Students Add their Functions inside ti_LSRK--------------------------#
    (q0,time) = ti_LSRK!(q0,qb,coord,M,Dwe,intma,periodicity,time,ntime,dt,gravity,Δ_NL,space_method,plot_movie,case,DFloat)
    #-----------------------------Students Add their Functions inside ti_LSRK--------------------------#

    #Print Time and Extrema
    println(" ntime = ",ntime," time = ",time," qmax = ",maximum(q0)," qmin = ",minimum(q0))

    #Compute L2 Norm
    (qe,qb,gravity) = exact_solution(coord,Npoin,time,case,DFloat)
    norms = compute_norms(q0,qe,DFloat)
    println(" norm.h = ",norms[1]," norm.U = ",norms[2])

    #Plot Final Solution
     if plot_solution
        p1 = plot(coord,qe[:,1],linecolor = :blue,label = "Analytic Solution",legend = :bottomright,legendfont=font(5))
        plot!(p1,coord,q0[:,1],linestyle = :dashdot,linecolor = :red,label = "$space_method Solution")
        title!(p1,"Height (hs)")
        xlabel!(p1,"x")
        ylabel!(p1,"hs(x,t)")
        plot!(p1,xlims = (first(coord),last(coord)+.1),ylims = (-1,0.5))

        p2 = plot(coord,qe[:,2],linecolor = :blue,label = "Analytic Solution",legend = :bottomright,legendfont=font(5))
        plot!(p2,coord,q0[:,2],linestyle = :dashdot,linecolor = :red,label = "$space_method Solution")
        title!(p2,"Momentum (U)")
        xlabel!(p2,"x")
        ylabel!(p2,"U(x,t)")
        plot!(p2,xlims = (first(coord),last(coord)+.1),ylims = (-1,0.5))

        figure_1 = plot(p1,p2,layout = (1,2))
        display(figure_1)
    end
    println("**Simulation Finished**") #output
end
#Main

#----------------------------------#
# Run the main function
#----------------------------------#
main()


using Omega
using Test

function testcond()
  x = normal(0.0, 1.0)
  x_ = cond(x, x < 1.0)
  x_samples = rand(x_, 10000, alg = RejectionSample)
  x_samples = filter(s -> s != nothing, x_samples)
  @test maximum(x_samples) < 1.0
  x__ = cond(x_, x > -1.0)
  x__samples = rand(x__, 10000, alg = RejectionSample)
  x__samples = filter(s -> s != nothing, x__samples)
  @test minimum(x__samples) > -1.0
end

testcond()

function testcond2()
  x = cond(poisson(2.0), iseven)
  xsamples = rand(x, 100, alg = RejectionSample)
  @test all(iseven, xsamples)
end

testcond2()


# ---
# title: 789. Escape The Ghosts
# id: problem789
# author: Indigo
# date: 2021-01-29
# difficulty: Medium
# categories: Math
# link: <https://leetcode.com/problems/escape-the-ghosts/description/>
# hidden: true
# ---
#
# You are playing a simplified PAC-MAN game on an infinite 2-D grid. You start
# at the point `[0, 0]`, and you are given a destination point `target =
# [xtarget, ytarget]`, which you are trying to get to. There are several ghosts
# on the map with their starting positions given as an array `ghosts`, where
# `ghosts[i] = [xi, yi]` represents the starting position of the `ith` ghost.
# All inputs are **integral coordinates**.
#
# Each turn, you and all the ghosts may independently choose to either **move 1
# unit** in any of the four cardinal directions: north, east, south, or west or
# **stay still**. All actions happen **simultaneously**.
#
# You escape if and only if you can reach the target **before** any ghost
# reaches you. If you reach any square (including the target) at the **same
# time** as a ghost, it **does not** count as an escape.
#
# Return `true` _if it is possible to escape, otherwise return_`false` _._
#
#
#
# **Example 1:**
#
#
#
#     Input: ghosts = [[1,0],[0,3]], target = [0,1]
#     Output: true
#     Explanation: You can reach the destination (0, 1) after 1 turn, while the ghosts located at (1, 0) and (0, 3) cannot catch up with you.
#
#
# **Example 2:**
#
#
#
#     Input: ghosts = [[1,0]], target = [2,0]
#     Output: false
#     Explanation: You need to reach the destination (2, 0), but the ghost at (1, 0) lies between you and the destination.
#
#
# **Example 3:**
#
#
#
#     Input: ghosts = [[2,0]], target = [1,0]
#     Output: false
#     Explanation: The ghost can reach the target at the same time as you.
#
#
# **Example 4:**
#
#
#
#     Input: ghosts = [[5,0],[-10,-2],[0,-5],[-2,-2],[-7,1]], target = [7,7]
#     Output: false
#
#
# **Example 5:**
#
#
#
#     Input: ghosts = [[-1,0],[0,1],[-1,0],[0,1],[-1,0]], target = [0,0]
#     Output: true
#
#
#
#
# **Constraints:**
#
#   * `1 <= ghosts.length <= 100`
#   * `ghosts[i].length == 2`
#   * `-104 <= xi, yi <= 104`
#   * There can be **multiple ghosts** in the same location.
#   * `target.length == 2`
#   * `-104 <= xtarget, ytarget <= 104`
#
#
## @lc code=start
using LeetCode

function escape_ghosts(ghosts::Vector{Vector{Int}}, target::Vector{Int})
    dist1 = sum(abs, target)
    dist2 = minimum(sum(abs, g .- target) for g in ghosts)
    return dist1 < dist2
end
## @lc code=end


@testset "SS-HOPM" begin

Random.seed!(1)
eigvec = randn(20)
eigval = norm(eigvec)
eigvec ./= eigval
T = eigval * eigvec .* eigvec' .* reshape(eigvec, 1, 1, 20)

@testset "Dense representation" begin
    @time (lbd, x) = sshopm(T, 1, verbose=true)
    @test norm(eigvec .- x) < 1e-5
    @test norm(eigval .- lbd) < 1e-5
end

@testset "Sparse representation" begin
    spT = SparseArray(T)

    @time (lbd, x) = sshopm(spT, 1, verbose=true)
    @test norm(eigvec .- x) < 1e-5
    @test norm(eigval .- lbd) < 1e-5
end

end