using SymEngine

"""
Creates the necessary differentiated versions of base kernels required by GLOM.
You must pass it a SymEngine Basic object with the variables already declared with the @vars command.
The created function will look like this

    \$kernel_name(hyperparameters::Vector{<:Real}, δ::Real; dorder::Vector{<:Integer}=zeros(Int64, length(hyperparameters) + 2))

For example, you could define a kernel like so:

    "Radial basis function GP kernel (aka squared exonential, ~gaussian)"
    function se_kernel_base(λ::Number, δ::Number)
        return exp(-δ ^ 2 / (2 * λ ^ 2))
    end

And then calculate the necessary derivative versions like so:

    @vars δ λ
    kernel_coder(se_kernel_base(λ, δ), "se_kernel")

The function is saved in src/kernels/\$kernel_name.jl, so you can use it with a command akin to this:

    include("src/kernels/" * kernel_name * ".jl")

"""

function kernel_coder(
    symbolic_kernel_original::Basic,
    kernel_name::String;
    periodic_var::String="",
    cutoff_var::String="")
    # add_white_noise::Bool=false)

    δ = symbols("δ")
    periodic = periodic_var != ""
    if periodic
        P_sym_str = kernel_name * "_P"
        P_sym = symbols(P_sym_str)
        π_sym = symbols("π_sym")
        @assert occursin(periodic_var, SymEngine.toString(symbolic_kernel_original)) "can't find periodic variable"
        symbolic_kernel_original = subs(symbolic_kernel_original, abs(symbols(periodic_var))=>2*sin(π_sym*abs(δ)/P_sym))
        symbolic_kernel_original = subs(symbolic_kernel_original, symbols(periodic_var)=>2*sin(π_sym*δ/P_sym))
        # length(periodic_vars)==1 ? P_sym_str = kernel_name * "_P" : P_sym_str = kernel_name * "_P$i"
        # for i in 1:length(periodic_vars)
        #     length(periodic_vars)==1 ? P_sym_str = kernel_name * "_P" : P_sym_str = kernel_name * "_P$i"
        #     P_sym = symbols(P_sym_str)
        #     append!(symbols, [P_sym])
        #     append!(symbols_str, [P_sym_str])
        #     symbolic_kernel_original = subs(symbolic_kernel_original, periodic_vars[i]=>2*sin(π*δ/P_sym))
        # end
        # kernel_name *= "_periodic"
    end

    # if add_white_noise
    #     σ_white = symbols("σ_white")
    #     symbolic_kernel_original += σ_white * σ_white
    # end

    # get the symbols of the passed function
    symbs = free_symbols(symbolic_kernel_original)

    δ_inds = findall(x -> x==δ, symbs)
    @assert length(δ_inds) <= 1
    if length(δ_inds)==1
        deleteat!(symbs, δ_inds[1])
    end
    π_inds = findall(x -> x==symbols("π_sym"), symbs)
    if length(π_inds) > 0
        deleteat!(symbs, π_inds[1])
    end
    symbs_str = [string(symb) for symb in symbs]

    hyper_amount = length(symbs)

    # open the file we will write to
    kernel_name *= "_kernel"
    file_loc = "src/kernels/" * kernel_name * ".jl"
    io = open(file_loc, "w")

    # begin to write the function including assertions that the amount of hyperparameters are correct
    write(io, """import GPLinearODEMaker.powers_of_negative_one\n\n\"\"\"
$kernel_name function created by kernel_coder(). Requires $hyper_amount hyperparameters. Likely created using $(kernel_name)_base() as an input.
Use with include(\"src/kernels/$kernel_name.jl\").
hyperparameters == $symbs_str
\"\"\"
function $kernel_name(
    hyperparameters::Vector{<:Real},
    δ::Real,
    dorder::Vector{<:Integer};
    shift_ind::Integer=0)

    include_shift = shift_ind!=0
    @assert length(hyperparameters)==$hyper_amount+Int(include_shift) \"hyperparameters is the wrong length\"
    dorder_len = $(hyper_amount + 2) + Int(include_shift)
    @assert length(dorder)==dorder_len \"dorder is the wrong length\"
    @assert maximum(dorder) < 3 \"No more than two derivatives for each time or hyperparameter can be calculated\"

    dorder2 = dorder[2]
    dorder[2] += dorder[1]

    if include_shift
        dorder[2] += dorder[shift_ind+2]
        δ += hyperparameters[shift_ind]
        hyperparameters_view = view(hyperparameters, 1:$(hyper_amount+1) .!= shift_ind)
        dorder_view = view(dorder, 2:dorder_len)
        dorder_view = view(dorder_view, 1:(dorder_len-1) .!= (shift_ind+1))
    else
        hyperparameters_view = hyperparameters
        dorder_view = view(dorder, 2:dorder_len)
    end\n\n""")

    # map the hyperparameters that will be passed to this function to the symbol names
    for i in 1:(hyper_amount)
        write(io, "    " * symbs_str[i] * " = hyperparameters_view[$i]\n")
    end
    # if add_white_noise
    #     write(io, "    σ_white *= Int(δ == 0)\n")  # accounting for white noise only being on the diagonal
    # end
    write(io, "\n")

    if cutoff_var!=""
        @assert cutoff_var in symbs_str
        write(io, "    if abs(δ) > $cutoff_var\n")
        write(io, "        dorder[2] = dorder2\n")
        write(io, "        return 0\n")
        write(io, "    end\n\n")
    end

    abs_vars = String[]
    for i in append!(["δ"], symbs_str)
        println("abs($i)")
        if occursin("abs($i)", SymEngine.toString(symbolic_kernel_original))
            append!(abs_vars, [i])
            println("found something")
            # write(io, "    dabsδ = sign(δ)  # store derivative of abs()\n")
            write(io, "    dabs_$i = powers_of_negative_one($i < 0)  # store derivative of abs()\n")
            write(io, "    abs_$i = abs($i)\n\n")
        end
    end

    max_δ_derivs = 7  # δf has 7 derivatives (0-(4+2), 4 from 2nd derivatives of GP and 2 from a possible time delay hyperparameter)
    max_hyper_derivs = 3  # other symbols have 3 (0-2)
    # calculate all of the necessary derivations we need for the GLOM model
    # for two symbols (δ and a hyperparameter), dorders is of the form:
    # [6 2; 4 2; 3 2; 2 2; 1 2; 0 2; ... ]
    # where the dorders[n, :]==[dorder of δ, dorder of symbol 2]
    # can be made for any number of symbols

    dorders = zeros(Int64, max_δ_derivs * (max_hyper_derivs ^ hyper_amount), hyper_amount + 1)
    amount_of_dorders = size(dorders,1)
    for i in 1:amount_of_dorders
          quant = amount_of_dorders - i
          dorders[i, 1] = rem(quant, max_δ_derivs)
          quant = div(quant, max_δ_derivs)
          for j in 1:(hyper_amount)
                dorders[i, j+1] = rem(quant, max_hyper_derivs)
                quant = div(quant, max_hyper_derivs)
          end
    end

    # for each differentiated version of the kernel
    for i in 1:amount_of_dorders

        # dorder = convert(Vector{Int64}, dorders[i, :])
        dorder = dorders[i, :]

        # only record another differentiated version of the function if we will actually use it
        # i.e. no instances where differentiations of multiple, non-time symbols are asked for
        if sum(dorder[2:end]) < max_hyper_derivs
            symbolic_kernel = copy(symbolic_kernel_original)

            # performing the differentiations
            symbolic_kernel = diff(symbolic_kernel, δ, dorder[1])
            for j in 1:hyper_amount
                # println(symbs[j], dorder[j+1])
                symbolic_kernel = diff(symbolic_kernel, symbs[j], dorder[j+1])
            end

            for j in abs_vars
                variable, abs_var, dabs_var = symbols("$j abs_$j dabs_$j")
                symbolic_kernel = subs(symbolic_kernel, diff(abs(variable), variable)=>dabs_var)
                for i in 2:6
                    symbolic_kernel = subs(symbolic_kernel, diff(abs(variable), variable, i)=>0)
                end
                for i in 1:2
                    symbolic_kernel = subs(symbolic_kernel, dabs_var^(i*2)=>1)
                end
                symbolic_kernel = subs(symbolic_kernel, abs(variable)=>abs_var)
            end

            symbolic_kernel_str = SymEngine.toString(symbolic_kernel)
            symbolic_kernel_str = replace(symbolic_kernel_str, "π_sym"=>"π")
            symbolic_kernel_str = " " * symbolic_kernel_str

            # some simplifications
            symbolic_kernel_str = replace(symbolic_kernel_str, "sqrt(δ^2)"=>"abs(δ)")
            symbolic_kernel_str = replace(symbolic_kernel_str, " 0.0 - "=>" -")
            symbolic_kernel_str = replace(symbolic_kernel_str, "(0.0 - "=>"(-")
            symbolic_kernel_str = replace(symbolic_kernel_str, "(0.0 + "=>"(")
            symbolic_kernel_str = replace(symbolic_kernel_str, " 0.0 + "=>" ")
            symbolic_kernel_str = replace(symbolic_kernel_str, " 1.0*"=>" ")
            symbolic_kernel_str = replace(symbolic_kernel_str, "-1.0*"=>"-")

            # println(symbolic_kernel_str)

            write(io, "    if dorder_view==" * string(dorder) * "\n")
            write(io, "        func =" * symbolic_kernel_str * "\n    end\n\n")
        end

    end

    write(io, "    dorder[2] = dorder2  # resetting dorder[2]\n")
    write(io, "    return powers_of_negative_one(dorder2) * float(func)  # correcting for amount of t2 derivatives")
    write(io, "\n\nend\n\n\n")
    write(io, "return $kernel_name, $hyper_amount  # the function handle and the number of kernel hyperparameters\n")
    close(io)

    @warn "The order of hyperparameters in the function may be different from the order given when you made them in @vars. Check the created function file (or the print statement below) to see what order you need to actually use."
    println("$kernel_name() created at $file_loc")
    println("hyperparameters == ", symbs_str)
end


export SpikeMixture

# TODO: Add `AbstractSuperposition <: AbstractMeasure`, and make SpikeMixture a
# subtype of this
struct SpikeMixture{T,S} <: AbstractMeasure
    m::T # parent
    w::S # relative weight of parent
    s::S # scale
end
SpikeMixture(m, w) = SpikeMixture(m, w, one(w))

function Base.show(io::IO, μ::SpikeMixture)
    io = IOContext(io, :compact => true)
    print(io, "(", μ.s * μ.w, "*", string(μ.m), " + ", μ.s * (1 - μ.w), "Dirac(0))")
end

function basemeasure(μ::SpikeMixture)
    # Compare formula (1.4) in Joris Bierkens, Sebastiano Grazzi, Frank van der Meulen, Moritz Schauer:
    # Sticky PDMP samplers for sparse and local inference problems. 2020. [https://arxiv.org/abs/2103.08478].
    ki = (1 / μ.w - 1) / density(μ.m, 0)
    SpikeMixture(basemeasure(μ.m), 1 / (1 + ki), μ.s * (1 + ki))
end

function logdensity(μ::SpikeMixture, x)
    return log(μ.w) + logdensity(μ.m, x)
end

function sampletype(μ::SpikeMixture)
    sampletype(μ.m)
end

function Base.rand(rng::AbstractRNG, T::Type, μ::SpikeMixture)
    μ.s != 1 && throw(ArgumentError("Not a probability measure"))
    return (rand(rng, T) < μ.w) * rand(rng, T, μ.m)
end

testvalue(μ::SpikeMixture) = testvalue(μ.m)


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    ModularProduct, modularproduct,
    CartesianProduct, cartesianproduct

# Modular product

struct ModularProductNode <: AbstractNode
    g::Int
    h::Int
end

struct ModularProductEdge <: UndirectedEdge
    hasedge::Bool
end

struct ModularProduct <: OrderedGraph
    neighbormap::Vector{Dict{Int,Int}}
    edges::Vector{Tuple{Int,Int}}
    nodeattrs::Vector{ModularProductNode}
    edgeattrs::Vector{ModularProductEdge}
    cache::Dict{Symbol,Any}
end

"""
    modularproduct(G::OrderedGraph, H::OrderedGraph) -> ModularProduct

Return the modular product of graphs G and H.
"""
function modularproduct(G::OrderedGraph, H::OrderedGraph;
            nodematcher=(g,h)->true,
            edgefilter=(g1,g2,h1,h2)->hasedge(G,g1,g2)==hasedge(H,h1,h2),
            kwargs...)
    product = ModularProduct([], [], [], [], Dict())
    ndict = Dict{Int,Dict{Int,Int}}() # Ref to node indices of the product
    # Modular product nodes
    nattrs = ModularProductNode[]
    for g in 1:nodecount(G)
        ndict[g] = Dict{Int,Int}()
        for h in 1:nodecount(H)
            ndict[g][h] = addnode!(product, ModularProductNode(g, h))
        end
    end
    (nodecount(G) < 2 || nodecount(H) < 2) && return product
    # Modular product edges
    for (g1, g2) in combinations(nodecount(G))
        for (h1, h2) in combinations(nodecount(H))
            edgefilter(g1, g2, h1, h2) || continue
            if nodematcher(g1, h1) && nodematcher(g2, h2)
                addedge!(
                    product, ndict[g1][h1], ndict[g2][h2],
                    ModularProductEdge(hasedge(G, g1, g2))
                )
            end
            if nodematcher(g1, h2) && nodematcher(g2, h1)
                addedge!(
                    product, ndict[g1][h2], ndict[g2][h1],
                    ModularProductEdge(hasedge(G, g1, g2))
                )
            end
        end
    end
    return product
end

# Cartesian product

struct CartesianProductNode <: AbstractNode
    g::Int
    h::Int
end

struct CartesianProduct <: OrderedGraph
    neighbormap::Vector{Dict{Int,Int}}
    edges::Vector{Tuple{Int,Int}}
    nodeattrs::Vector{CartesianProductNode}
    cache::Dict{Symbol,Any}
end

"""
    cartesianproduct(G::OrderedGraph, H::OrderedGraph) -> CartesianProduct

Return the cartesian product of graphs G and H.
"""
function cartesianproduct(G::OrderedGraph, H::OrderedGraph)
    product = CartesianProduct([], [], [], Dict())
    ndict = Dict{Int,Dict{Int,Int}}() # Ref to node indices of the product
    #Cartesian product nodes
    nattrs = CartesianProductNode[]
    for g in 1:nodecount(G)
        ndict[g] = Dict{Int,Int}()
        for h in 1:nodecount(H)
            ndict[g][h] = addnode!(product, CartesianProductNode(g, h))
        end
    end
    (nodecount(G) < 2 || nodecount(H) < 2) && return product
    # Cartesian product edges
    for i in 1:nodecount(G)
        for (u, v) in edgesiter(H)
            addedge!(product, ndict[i][u], ndict[i][v])
        end
    end
    for i in 1:nodecount(H)
        for (u, v) in edgesiter(G)
            addedge!(product, ndict[u][i], ndict[v][i])
        end
    end
    return product
end


using Memoize

function testlcs()
	s1::String = "Please write comments if you find anything incorrect"
	s2::String = "or you want to share more information about the topic discussed above"
	@time a = lcs(s1, s2)
	@show a
end

function lcs(s1::String, s2::String)
	@memoize function rec(n1::Int, n2::Int)
		(n1 <= 0 || n2 <= 0) && return 0
		return max(	rec(n1 - 1, n2),
					rec(n1, n2 - 1),
					rec(n1 - 1, n2 - 1) +
					(s1[n1] == s2[n2] ? 1 : 0))
	end
	return rec(length(s1), length(s2))
end


module TaylorTests

using ForwardDiff
using Morbit
using Test

using LinearAlgebra: norm

const RFD = Morbit.RFD
const Trees = RFD.Trees

ensure_vec( x ) = vec(x)
ensure_vec( x :: Number ) = [x,]

const test_funcs = [
	x -> [x[1],],
	x -> [π + ℯ * sum(x.^2),],
	x -> x.^2,
	x -> [ exp(sum(x)), x[1] + x[end]^3 ]
]

const test_stamps64 = [
	RFD.CFDStamp(1,2),
	RFD.CFDStamp(1,4),
	RFD.CFDStamp(1,6),
	RFD.FFDStamp(1,1),
	RFD.FFDStamp(1,2),
	RFD.FFDStamp(1,3),
	RFD.BFDStamp(1,1),
	RFD.BFDStamp(1,2),
	RFD.BFDStamp(1,3),
]

h32 = RFD.stepsize(RFD.stepsize(Float32))
const test_stamps32 = [
	RFD.CFDStamp(1,2, h32),
	RFD.CFDStamp(1,4, h32),
	RFD.CFDStamp(1,6, h32),
	RFD.FFDStamp(1,1, h32),
	RFD.FFDStamp(1,2, h32),
	RFD.FFDStamp(1,3, h32),
	RFD.BFDStamp(1,1, h32),
	RFD.BFDStamp(1,2, h32),
	RFD.BFDStamp(1,3, h32),
]

h16 = RFD.stepsize(RFD.stepsize(Float16))
const test_stamps16 = [
	RFD.CFDStamp(1,2, h16),
	RFD.CFDStamp(1,4, h16),
	RFD.CFDStamp(1,6, h16),
	RFD.FFDStamp(1,1, h16),
	RFD.FFDStamp(1,2, h16),
	RFD.FFDStamp(1,3, h16),
	RFD.BFDStamp(1,1, h16),
	RFD.BFDStamp(1,2, h16),
	RFD.BFDStamp(1,3, h16),
]

@testset "Recursive Finite Difference Trees" begin
	for ET in [Float16, Float32, Float64]
		stamps = ET == Float16 ? test_stamps16 : ET == Float32 ? test_stamps32 : test_stamps64

		for (j,func) in enumerate(test_funcs)
			n_in = rand(1:5)
			x0 = ensure_vec(rand(ET, n_in))
			fx0 = func(x0)
			for stamp in stamps
				for order in [1,2]
					# first, use convenience function to fill tree
					dw_1 = RFD.DiffWrapper(; x0, fx0, stamp, order)
					RFD.prepare_tree!(dw_1, func)

					# secondly, use the Morbit routine
					dw_2 = RFD.DiffWrapper(; x0, fx0, stamp, order)
					RFD.substitute_leaves!(dw_2)
					all_sites = RFD.collect_leave_sites( dw_2 )
					all_vals = func.(all_sites)
					RFD.set_leave_values!(dw_2, all_vals)

					if order > 1
						for i = 1 : length(fx0)
							Hi = RFD.hessian(dw_1; output_index = i)
							@test Hi ≈ RFD.hessian( dw_2; output_index = i)
							# I don't test against automatic differencing because
							# finite diff hessians become somewhat instable without careful tuning
							# of stepsize …
							# H_ad = ForwardDiff.hessian( x -> func(x)[i], x0 )
							# ET == Float16 || @test norm( Hi .- H_ad ) <= .1
						end
					end

					J1 = RFD.jacobian(dw_1)

					for i = 1 : length(fx0)
						Gi = RFD.gradient( dw_1; output_index = i )
						@test Gi ≈ RFD.gradient( dw_2; output_index = i)
						@test Gi ≈ J1[i,:]
					end

					# does it roughly compare to a jacobian determined by automatic differencing
					J_ad = ForwardDiff.jacobian( func, x0 )
					ET == Float16 || @test norm( J1 .- J_ad ) <= 0.1

				end
			end
		end
	end

end#testset

@testset "Accurate Linear Models, Unconstrained " begin
	mop = Morbit.MixedMOP(2)

	# The gradients for a linear objective will be exact and a linear model
	# should then equal the linear objective globally
	Morbit.add_objective!( mop, x -> sum(x), Morbit.TaylorConfig(;degree=1) )

	x0 = [π, -ℯ]

	smop, iter_data, data_base, sc, ac = Morbit.initialize_data(mop, x0, Float32[] )
	Morbit.update_surrogates!( sc, mop, iter_data, data_base, ac )

	mod = sc.surrogates[1].model
	objf = Morbit.list_of_objectives(smop)[1]

	@test Morbit.eval_models( mod, x0 ) ≈ Morbit.eval_objf(objf,x0)
	@test Morbit.get_gradient( mod, x0, 1 ) ≈ [1, 1]
	@test Morbit.get_jacobian(sc, x0) ≈ [ 1 1 ]

end#testset

@testset "Accurate Linear Models, Constrained " begin
	mop = Morbit.MixedMOP( fill(-10.0, 2), fill(10.0, 2) )

	Morbit.add_objective!( mop, x -> sum(x), Morbit.TaylorConfig(;degree=1) )

	x0 = [π, -ℯ]

	smop, iter_data, data_base, sc, ac = Morbit.initialize_data(mop, x0, Float32[] )
	#Morbit.update_surrogates!( sc, mop, iter_data, data_base, ac )

	mod = sc.surrogates[1].model
	objf = Morbit.list_of_objectives(smop)[1]

	x̂0 = Morbit.scale( x0, smop)

	@test Morbit.eval_models( mod, x̂0 ) ≈ Morbit.eval_objf(objf, x0 )
	@test Morbit.get_gradient( mod, x̂0, 1 ) ≈ ForwardDiff.gradient( ξ -> Morbit.eval_objf(objf, Morbit.unscale(ξ, smop) )[end], x̂0 )
	@test Morbit.get_jacobian( mod, x̂0 ) ≈ ForwardDiff.jacobian( ξ -> Morbit.eval_all_objectives(smop, ξ), x̂0 )

end#testset

@testset "Quadratic Models, Unconstrained " begin
	mop = Morbit.MixedMOP(2)

	Morbit.add_objective!( mop, x -> sum(x.^2), Morbit.TaylorConfig(;degree=2) )

	x0 = [π, -ℯ]

	smop, iter_data, data_base, sc, ac = Morbit.initialize_data(mop, x0, Float32[] )
	#Morbit.update_surrogates!( sc, mop, iter_data, data_base, ac )

	mod = sc.surrogates[1].model
	objf = Morbit.list_of_objectives(smop)[1]

	@test Morbit.eval_models( mod, x0 ) ≈ Morbit.eval_objf(objf,x0)
	@test Morbit.get_gradient( mod, x0, 1 ) ≈ 2 .* x0
	@test Morbit.get_jacobian(sc, x0) ≈ 2 .* x0'
end#testset

@testset "Two Parabolas, Taylor Models, 3D2D" begin
	mop = MixedMOP(3)

	add_objective!( mop, x -> sum( ( x .- 1 ).^2 ), Morbit.TaylorApproximateConfig(;degree=2,mode=:fdm) )
	add_objective!( mop, x -> sum( ( x .+ 1 ).^2 ), Morbit.TaylorApproximateConfig(;degree=2,mode=:autodiff) )

	x_fin, f_fin, _ = optimize( mop, [-π, ℯ, 0])

	@test isapprox(x_fin[1],x_fin[2], atol = 1e-3) && isapprox(x_fin[2], x_fin[3]; atol = 1e-3)

	mop = MixedMOP(3)

	add_objective!( mop, x -> sum( ( x .- 1 ).^2 ),
		TaylorCallbackConfig(; degree = 1, gradients = [ x -> 2 .* ( x .- 1 ), ] ))
	add_objective!( mop, x -> sum( ( x .+ 1 ).^2 ),
		TaylorCallbackConfig(; degree=2, gradients = [ x -> 2 .* ( x .+ 1 ), ],
		hessians = [ x -> [ 2.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 2.0] ] ) )

	x_fin, f_fin, _ = optimize( mop, [-π, ℯ, 0])

	@test isapprox(x_fin[1],x_fin[2], atol = 1e-3) && isapprox(x_fin[2], x_fin[3]; atol = 1e-3)

end#testset

@testset "Two Parabolas, Derivative Accuracy" begin
	# deg 1 autodiff
	mop = Morbit.MixedMOP( 2 )

	Morbit.add_objective!( mop, x -> sum(x), Morbit.TaylorApproximateConfig(;degree=1, mode=:autodiff) )

	x0 = [π, -ℯ]

	smop, iter_data, data_base, sc, ac = Morbit.initialize_data(mop, x0, Float32[] )
#	Morbit.update_surrogates!( sc, mop, iter_data, data_base, ac )

	mod = sc.surrogates[1].model
	objf = Morbit.list_of_objectives(smop)[1]

	x̂0 = Morbit.scale( x0, smop)

	@test Morbit.eval_models( mod, x0 ) ≈ Morbit.eval_objf(objf,x0)
	@test Morbit.get_gradient( mod, x0, 1 ) ≈ [1, 1]
	@test Morbit.get_jacobian(sc, x0) ≈ [ 1 1 ]

	#===================================================================#
	# deg 2 autodiff
	mop = Morbit.MixedMOP( 2 )

	Morbit.add_objective!( mop, x -> sum(x.^2), Morbit.TaylorApproximateConfig(;degree=2, mode=:autodiff) )

	x0 = [π, -ℯ]

	smop, iter_data, data_base, sc, ac = Morbit.initialize_data(mop, x0, Float32[] )
	# Morbit.update_surrogates!( sc, mop, iter_data, data_base, ac )

	mod = sc.surrogates[1].model
	objf = Morbit.list_of_objectives(smop)[1]

	x̂0 = Morbit.scale( x0, smop)

	@test Morbit.eval_models( mod, x0 ) ≈ Morbit.eval_objf(objf,x0)
	@test Morbit.get_gradient( mod, x0, 1 ) ≈ 2 .* x0
	@test Morbit.get_jacobian(sc, x0) ≈ 2 .* x0'
	@test mod.H[1] ≈ [ 2.0 0.0; 0.0 2.0 ]	# hessian of Taylor Model

	#===================================================================#
	# deg 1 finite diff
	mop = Morbit.MixedMOP( 2 )

	Morbit.add_objective!( mop, x -> sum(x), Morbit.TaylorApproximateConfig(;degree=1, mode=:fdm) )

	x0 = [π, -ℯ]

	smop, iter_data, data_base, sc, ac = Morbit.initialize_data(mop, x0, Float32[] )
	#Morbit.update_surrogates!( sc, mop, iter_data, data_base, ac )

	mod = sc.surrogates[1].model
	objf = Morbit.list_of_objectives(smop)[1]

	x̂0 = Morbit.scale( x0, smop)

	@test Morbit.eval_models( mod, x0 ) ≈ Morbit.eval_objf(objf,x0)
	@test Morbit.get_gradient( mod, x0, 1 ) ≈ [1, 1]
	@test Morbit.get_jacobian(sc, x0) ≈ [ 1 1 ]

	#===================================================================#
	# deg 2 autodiff
	mop = Morbit.MixedMOP( 2 )

	Morbit.add_objective!( mop, x -> sum(x.^2), Morbit.TaylorApproximateConfig(;degree=2, mode=:fdm) )

	x0 = [π, -ℯ]

	smop, iter_data, data_base, sc, ac = Morbit.initialize_data(mop, x0, Float32[] )
#	Morbit.update_surrogates!( sc, mop, iter_data, data_base, ac )

	mod = sc.surrogates[1].model
	objf = Morbit.list_of_objectives(smop)[1]

	x̂0 = Morbit.scale( x0, smop)

	@test Morbit.eval_models( mod, x0 ) ≈ Morbit.eval_objf(objf,x0)
	@test Morbit.get_gradient( mod, x0, 1 ) ≈ 2 .* x0
	@test Morbit.get_jacobian(sc, x0) ≈ 2 .* x0'
	@test norm(mod.H[1] .- [ 2.0 0.0; 0.0 2.0 ]) <= 0.1	# hessian of Taylor Model

end

end#module

using .TaylorTests


using GeometricIntegrators.Equations

ϑ₁(t, q) = q[4] + A₁(t,q)
ϑ₂(t, q) = q[5] + A₂(t,q)
ϑ₃(t, q) = q[6] + A₃(t,q)

function ϑ(t, q, p)
    p[1] = ϑ₁(t,q)
    p[2] = ϑ₂(t,q)
    p[3] = ϑ₃(t,q)
    p[4] = zero(eltype(q))
    p[5] = zero(eltype(q))
    p[6] = zero(eltype(q))
    nothing
end

# ϕ₀(x) = E₀*sin(2π*x)
# ϕ(t,q) = ϕ₀(q[3])

# E₁(t,q) = zero(eltype(q))
# E₂(t,q) = zero(eltype(q))
# E₃(t,q) = - 2π*E₀*cos(2π*q[3])

hamiltonian(t,q) = 0.5 * (q[4]^2 + q[5]^2 + q[6]^2) + ϕ(t, q)

angular_momentum(t,q) = q[1] * ϑ₂(t,q) - q[2] * ϑ₁(t,q)

function charged_particle_3d_pᵢ(qᵢ, tᵢ=0)
    pᵢ = zero(qᵢ)

    if ndims(qᵢ) == 1
        ϑ(tᵢ, qᵢ, pᵢ)
    else
        for i in 1:size(qᵢ,2)
            @views ϑ(tᵢ, qᵢ[:,i], pᵢ[:,i])
        end
    end
    pᵢ
end

charged_particle_3d_iode_ϑ(t, q, v, p) = ϑ(t, q, p)

function charged_particle_3d_iode_f(t, q, v, f)
    f[1] = dA₁dx₁(t,q) * v[1] + dA₂dx₁(t,q) * v[2] + dA₃dx₁(t,q) * v[3] + E₁(t,q)
    f[2] = dA₁dx₂(t,q) * v[1] + dA₂dx₂(t,q) * v[2] + dA₃dx₂(t,q) * v[3] + E₂(t,q)
    f[3] = dA₁dx₃(t,q) * v[1] + dA₂dx₃(t,q) * v[2] + dA₃dx₃(t,q) * v[3] + E₂(t,q)
    f[4] = v[1] - q[4]
    f[5] = v[2] - q[5]
    f[6] = v[3] - q[6]
    nothing
end

function charged_particle_3d_iode_g(t, q, λ, g)
    g[1] = dA₁dx₁(t,q) * λ[1] + dA₂dx₁(t,q) * λ[2] + dA₃dx₁(t,q) * λ[3]
    g[2] = dA₁dx₂(t,q) * λ[1] + dA₂dx₂(t,q) * λ[2] + dA₃dx₂(t,q) * λ[3]
    g[3] = dA₁dx₃(t,q) * λ[1] + dA₂dx₃(t,q) * λ[2] + dA₃dx₃(t,q) * λ[3]
    g[4] = λ[1]
    g[5] = λ[2]
    g[6] = λ[3]
    nothing
end

function charged_particle_3d_iode_v(t, q, v)
    v[1] = q[4]
    v[2] = q[5]
    v[3] = q[6]
    v[4] = E₁(t,q) + q[5] * B₃(t,q) - q[6] * B₂(t,q)
    v[5] = E₂(t,q) + q[6] * B₁(t,q) - q[4] * B₂(t,q)
    v[6] = E₃(t,q) + q[4] * B₂(t,q) - q[5] * B₁(t,q)
    nothing
end

function charged_particle_3d_iode(q₀=q₀)
    IODE(charged_particle_3d_iode_ϑ, charged_particle_3d_iode_f,
         charged_particle_3d_iode_g, q₀, charged_particle_3d_pᵢ(q₀);
         h=hamiltonian, v=charged_particle_3d_iode_v)
end


tol = 1e-13

pmax = 3

for ii = 1:pmax

  (ploc,) = luteos.genlocal3D( ii )

  (ϕ,∇ϕ) = luteos.basisFuncTetLag( Val{ii}, ploc[:,2], ploc[:,3], ploc[:,4] )
  println( "      3D")

  sz  = size( ϕ )
  for jj = 1:sz[1], kk = 1:sz[2]
    if jj == kk
      @test abs(ϕ[jj,kk] - 1.0) < tol
    else
      @test abs(ϕ[jj,kk]) < tol
    end
  end

  # check evaluation for random point
  xloc   = [0.11698  0.2678  0.05263;
            0.06905  0.1376  0.08579]
  (ϕ,∇ϕ) = luteos.basisFuncTetLag( Val{ii}, xloc[:,1], xloc[:,2], xloc[:,3] )

  for jj in 1:2, dd in 1:3
    xloctest = ϕ[:,jj]' * ploc[:,dd+1]
    @test abs( xloctest - xloc[jj,dd] ) < tol

  end

end


mutable struct Segment
   P0::SVector{3,Float64}    # Point 0 (support vector of line segment)
   P1::SVector{3,Float64}    # Point 1
   function Segment(P0::SVector{3,Float64},P1::SVector{3,Float64})
     new(P0,P1)
   end
end

mutable struct Plane
   V0::SVector{3,Float64}    # Point 0 of plane (support vector)
   V1::SVector{3,Float64}    # Point 1
   V2::SVector{3,Float64}    # Point 1
   n::SVector{3,Float64}     # normal of plane
   function Plane(V0::SVector{3,Float64}, V1::SVector{3,Float64}, V2::SVector{3,Float64})
      n = cross(V1-V0, V2-V0)
      new(V0,V1,V2,n)
   end
end

function intersect3DSegmentPlane(seg::Segment, plane::Plane, neps)
   u = seg.P1 - seg.P0
   w = seg.P0 - plane.V0

   dividend = -dot(plane.n, w)
   divisor = dot(plane.n, u)

   # checks if segment and plane are parallel
   if abs(divisor) < neps  # segment is parallel to plane
      if (abs(dividend) < neps)  # segemnt is parallel to plane AND lies in plane
         error("Segment is parallel to plane and lies in plane")
         return (2, nothing)
      else
         error("Segment is parallel to plane")
         return (0, nothing)
      end
   end

   # computes intersection point of plane and segment
   SI = dividend/divisor
   I = seg.P0 + SI*u
   return (1, I)
end

function sameSideTriangle(p1, p2, a, b)::Bool
   cp1 = cross(b-a, p1-a)
   cp2 = cross(b-a, p2-a)
   return (dot(cp1, cp2) >= 0.0)
end
pointInTriangle(p,a,b,c) = (sameSideTriangle(p,a,b,c) && sameSideTriangle(p,b,a,c) && sameSideTriangle(p,c,a,b))

function doesRayIntersectPortal(r1,r2,r3,point,neps)
   plane = Plane(r1,r2,r3)
   segment = Segment(point, SVector(0.0, 0.0, 0.0))

   (value, intersectionPoint) = intersect3DSegmentPlane(segment, plane, neps)
   if value == 1
      if !pointInTriangle(intersectionPoint, r1, r2, r3)
         error("Ray = ", point ," does not intersect portal (r1,r2,r3) = ", r1, r2, r3)
         return false
      end
   end

   return true
end

function analyzeFinalPortal(r1, r2, r3, r4, neps)
   plane = Plane(r1,r2,r3)
   segment = Segment(r4, SVector(0.0, 0.0, 0.0))
   (value, intersectionPoint) = intersect3DSegmentPlane(segment, plane, neps)
   if value == 1
      if !pointInTriangle(intersectionPoint, r1, r2, r3)
         println("Ray r4 = ", r4 ," is not intersecting last portal (r1,r2,r3) = ", r1, r2, r3)
         return false
      end
   end

   areaPlane = norm(plane.n)*1/2
   println("areaPlane = ", areaPlane)
    # det([r2-r1, r3-r1, r4-r1])*1/6
   volumeTetraeder1 = abs(dot(r2-r1,cross(r3-r1,r4-r1)))*1/6
   #volumeTetraeder2 = abs(det(hcat(r2-r1,r3-r1,r4-r1))) *1/6
   println("volumeTetraeder = ", volumeTetraeder1)
   #println("volumeTetraeder2 = ", volumeTetraeder2)
   return true
end

function sameSideTetrahedron(v1,v2,v3,v4,p)
   normal = cross(v2-v1,v3-v1)
   dotV4  = dot(normal, v4-v1)
   dotP   = dot(normal, p-v1)
   return ( sign(dotV4) == sign(dotP) )
end
pointInTetrahedron(v1,v2,v3,v4,p) = (sameSideTetrahedron(v1,v2,v3,v4,p) && sameSideTetrahedron(v2,v3,v4,v1,p) && sameSideTetrahedron(v3,v4,v1,v2,p) && sameSideTetrahedron(v4,v1,v2,v3,p) )


## First order perturbation approach
## Local solution for log-linear model around a non-stochastic steady state
## Blanchard and Kahn (1980)
## Based on Exercises for Practical DSGE Modelling
module DsgeSolves
using Distributions

type DSGE
  b::Int64
  f::Int64
  s::Int64
  A::Array
  B::Array
  C::Array
  v::Distribution
end

function DSGE(b,f,s)
  A = zeros(b+f,b+f)
  B = zeros(b+f,b+f)
  C = zeros(b+f,s)
  v = Normal(0,1)
  return DSGE(b,f,s,A,B,C,v)
end

function create_matrices(m)
  A = inv(m.A)*m.B
  B = inv(m.A)*m.C
  (λ,P) = real.(eig(A))
  λ = real(λ)
  t = sortrows([λ P'])
  Λ = diagm(t[:,1])
  P = t[:,2:(m.b+m.f+1)]'
  pstar = inv(P)
  R = pstar*B
 return A, B, Λ, P, pstar, R
end

function create_partition(m,pstar,Λ,R)
  Λ1 = Λ[1:m.b,1:m.b]
  Λ2 = Λ[(m.b+1):(m.b+m.f),(m.b+1):(m.b+m.f)]

  P11 = pstar[1:m.b,1:m.b]
  P12 = pstar[1:m.b,(m.b+1):(m.b+m.f)]
  P21 = pstar[(m.b+1):(m.b+m.f),1:m.b]
  P22 = pstar[(m.b+1):(m.b+m.f),(m.b+1):(m.b+m.f)]

  R1 = R[1:m.s,1:m.s]
  R2 = R[(m.b+1:m.s),1:m.s]
  return Λ1, Λ2, P11, P12, P21, P22, R1, R2
end

function simulating_dsge(m,pstar,Λ,R,Σ,rep)
  (Λ1, Λ2, P11, P12, P21, P22, R1, R2) = create_partition(m,pstar,Λ,R)
  w = zeros(m.b,rep)
  y = zeros(m.f,rep)
    for i in 2:rep
      w[:,i] = real(inv(P11-P12*inv(P22)*P21)*Λ1*(P11-P12*inv(P22)*P21))*w[:,i-1]+real((P11-P12*inv(P22)*P21)*R1)*Σ*rand(m.v,m.s)
      y[:,i] = -real(inv(P22)*P21)*w[:,i]
    end
return w, y
end

function irf_dsge(m,pstar,Λ,R,shk,rep)
  (Λ1, Λ2, P11, P12, P21, P22, R1, R2) = create_partition(m,pstar,Λ,R)
  shock = zeros(m.s,rep)
  shock[shk,2] = 1
  w = zeros(m.b,rep)
  y = zeros(m.f,rep)
    for i in 2:rep
      w[:,i] = real(inv(P11-P12*inv(P22)*P21)*Λ1*(P11-P12*inv(P22)*P21))*w[:,i-1]+real((P11-P12*inv(P22)*P21)*R1)*shock[:,i]
      y[:,i] = -real(inv(P22)*P21)*w[:,i]
    end
return w, y
end

function fevd(m,pstar,Λ,R,rep) ##restart from this function using a for loop
                               ## with the function irf_dsge
  (Λ1, Λ2, P11, P12, P21, P22, R1, R2) = create_partition(m,pstar,Λ,R)
  irf_w = zeros(m.b,rep,m.b*m.s)
  irf_y = zeros(m.f,rep,m.f*m.s)
  for j in 1:m.s
    shock = zeros(m.s,rep)
    shock[j,2] = 1
    for i in 2:rep
      w[:,i] = real(inv(P11-P12*inv(P22)*P21)*Λ1*(P11-P12*inv(P22)*P21))*w[:,i-1]+real((P11-P12*inv(P22)*P21)*R1)*shock[:,i]
      y[:,i] = -real(inv(P22)*P21)*w[:,i]
    end
    irf_w[:,:,j] = w
    irf_y[:,:,j] = y
  end
  sum(irf_w[1,:,1].^2)Σ[j,j]/(sum(irf_w[1,:,1].^2) + sum(irf_w[1,:,2].^2))
    irf_w[1,i,1]
    den =

return irf_w, irf_y
end

export DSGE, create_matrices, simulating_dsge, irf_dsge

end # end module


module XfromProjections

using SparseArrays
using LinearAlgebra
using PyCall
using Suppressor
using LinearOperators
using Logging

# analytic
include("analytic/filter_proj.jl")
export filter_proj, bp_slices
# include("analytic/gridrec.jl")
# export recon2d_gridrec

# iterative
include("iterative/util_convexopt.jl")
include("iterative/tv_primaldual.jl")
include("iterative/tv_primaldual_flow.jl")
include("iterative/ctv_primaldual.jl")
include("iterative/sirt.jl")
export recon2d_sirt!, recon2d_stack_sirt!, _compute_sum_rows_cols
export recon2d_tv_primaldual!, recon2d_stack_tv_primaldual!, recon2d_ctv_primaldual!, recon2d_tv_primaldual_flow
export get_flows, compute_warping_operator

end


using BAT
using Distributions
using IntervalSets
using TypedTables
using NamedArrays
using LinearAlgebra
using PrettyTables

likelihood = params -> begin
    r1 = logpdf.(Normal(-10.0, 1.2), params.a)
    r2 = logpdf.(Normal(-5.0, 0.8), params.c[1])
    r3 = logpdf.(Normal(12.0, 1.4), params.c[2])

    return LogDVal(r1+r2+r3)
end

prior = BAT.NamedTupleDist(
    a = Normal(-3, 4.5),
    b = -10..10,
    c = [-20.0..20.0, -10..10]
)

posterior = PosteriorDensity(likelihood, prior);

samples, chains = bat_sample(posterior, MCMCSampling(sampler = MetropolisHastings(), nsteps = 10^5));
#samples = bat_sample(posterior, SobolSampler(nsamples = 10^5)).result;

sd = SampledDensity(posterior, samples, generator=BAT.MCMCSampleGenerator(chains))
display(sd)

# # write summary to txt-file
# io = open("sampled_density.txt", "w")
# show(io, "text/plain", sd)
# close(io)
#
# # ### write output to html-file
# io = open("sampled_density.html", "w")
# show(io, "text/html", sd)
# close(io)


function Pixelberechner(x0,x1,y0,y1,nx,ny)
n = zeros(nx*ny,1);
Pixelx = zeros(nx*ny,1);
Pixely = zeros(nx*ny,1);

deltax = (x1-x0)/(2*nx); # Verschiebung, um zur x-Mitte des Pixels zu gelangen
deltay = (y1-y0)/(2*ny); # Verschiebung, um zur y-Mitte des Pixels zu gelangen
counter = 1;
for i = 1:ny # damit spaltenweise Pixel nummeriert werden
    for k = 1:ny
        n[counter,1] = counter;
        Pixelx[counter] = x0+i-deltax;
        Pixely[counter] = y1-k+deltay;
        counter = counter+1;
    end
end
return n, Pixelx, Pixely;
end


# source:
# https://github.com/JuliaGeometry/GeometryTypes.jl/issues/71#issuecomment-302175458

"""
A `Cone3` is a 3D cylinder defined by its origin point, its extremity
and a radius. `origin`, `extremity` and `r`, must be specified.
"""
struct Cone3{T} <: GeometryPrimitive{3,T}
  origin::Point{3,T}
  extremity::Point{3,T}
  r::T
end

origin(c::Cone3) = c.origin
extremity(c::Cone3) = c.extremity
radius(c::Cone3) = c.r
height(c::Cone3) = norm(c.extremity-c.origin)
direction(c::Cone3) = (c.extremity .- c.origin)./height(c)
function rotation(c::Cone3)
  u = [direction(c)...];
  v = abs(u[1])>0 || abs(u[2])>0 ? [u[2],-u[1],0] : [0,-u[3],u[2]]
  v ./= norm(v)
  w = [u[2]*v[3]-u[3]*v[2],-u[1]*v[3]+u[3]*v[1],u[1]*v[2]-u[2]*v[1]]
  return hcat(v,w,u)
end

isdecomposable(::Type{T}, ::Type{C}) where {T<:Point, C<:Cone3} = true
isdecomposable(::Type{T}, ::Type{C}) where {T<:Face, C<:Cone3} = true

function decompose1(PT::Type{Point{3,T}},c::Cone3{T},facets=12) where {T}
  isodd(facets) ? facets = 2*div(facets,2) : nothing
  facets<8 ? facets = 8 : nothing
  nbv = Int(facets/2)
  M = rotation(c)
  h = height(c)
  position = 1
  vertices = Array{PT}(undef, 2*nbv)
  for j = 1:nbv
    phi                = T((2*pi*(j-1))/nbv)
    vertices[position] = PT(M*[c.r*cos(phi);c.r*sin(phi);0])+PT(c.origin)
    vertices[position+1] = PT([0;0;h])+PT(c.origin)
    position += 2
  end
  vertices
end

function decompose1(::Type{FT},c::Cone3{T},facets=12) where FT<:Face where T
  isodd(facets) ? facets = 2*div(facets,2) : nothing
  facets<8 ? facets = 8 : nothing
  nbv = Int(facets/2)
  indexes = Array{Face{3,Int}}(undef, facets)
  index = 1
  for j = 1:(nbv-1)
    indexes[index] = (index,index+1,index+2)
    indexes[index+1] = (index+2,index+1,index+3)
    index += 2
  end
  indexes[index] = (index,index+1,1)
  indexes[index+1] = (1,index+1,2)
  indexes
end

function ConeMesh(orig, extr, rad, nfacets)
  obj = Cone3(orig, extr, rad)
  vert = decompose1(Point3f0, obj, nfacets)
  fac = decompose1(GLTriangle, obj, nfacets)
  fovMesh = GLNormalMesh(vertices=vert, faces=fac)
end


using ParamLevelSet
using jInv.Mesh
using LinearAlgebra

println("Testing the rotation itself (extended RBFs):")

n = [64,64,64];
Mesh = getRegularMesh([0.0;5.0;0.0;5.0;0.0;5.0],n);
midmesh = [(Mesh.domain[2]+Mesh.domain[1])/2 ; (Mesh.domain[4]+Mesh.domain[3])/2; (Mesh.domain[5]+Mesh.domain[6])/2];
lenmesh = [(Mesh.domain[2]-Mesh.domain[1]) ; (Mesh.domain[4]-Mesh.domain[3]); (Mesh.domain[6]-Mesh.domain[5])];
numParamOfRBF = 10;
nRBF = 5
m = zeros(nRBF*numParamOfRBF);
for k=1:nRBF
	offset = (k-1)*numParamOfRBF + 1;
	m[offset] = randn();
	#A = randn(3,3); A = 0.3*A'*A + 0.5*eye(3);
	A = diagm(rand(3).+1.0);
	# L = chol(A)';
	# m[(offset+1):(offset+6)] = L[L.!=0.0];
	m[(offset+1):(offset+6)] = A[idx_trilA()];
	m[(offset+7):(offset+9)] = midmesh .+ 0.05*lenmesh.*randn(size(lenmesh));
end
u, = ParamLevelSetModelFunc(Mesh,m;computeJacobian=0,numParamOfRBF=numParamOfRBF);

theta_phi = deg2rad.([37.0 24.0 ; 27.0 14.0]);
b = [1.0 5.0 7.0 ; -2.3 5.2 3.5]*sum((Mesh.h))/3.0;

Xc = convert(Array{Float32,2},getCellCenteredGrid(Mesh));

N = prod(size(u));
ur1 = zeros(N ,2);
ur2 = copy(ur1);

u = reshape(u,tuple(n...));

@elapsed ur1[:,1] = reshape(rotateAndMove3D(u,theta_phi[1,:],b[1,:]./Mesh.h)[1],N,1);
@elapsed ur1[:,2] = reshape(rotateAndMove3D(u,theta_phi[2,:],b[2,:]./Mesh.h)[1],N,1);;

mr2,Jrot = rotateAndMoveRBFsimple(m,Mesh,theta_phi,b;computeJacobian = 0,numParamOfRBF=numParamOfRBF);
@elapsed ur2[:,1] = ParamLevelSetModelFunc(Mesh,mr2[:,1];computeJacobian=0,Xc = Xc,numParamOfRBF=numParamOfRBF)[1];
@elapsed ur2[:,2] = ParamLevelSetModelFunc(Mesh,mr2[:,2];computeJacobian=0,Xc = Xc,numParamOfRBF=numParamOfRBF)[1];

err = norm(ur1-ur2)./(2*norm(u));
if err < 2e-1
	println("RBF Rotation test succeeded: ",err);
else
	error("RBF Rotation test not succeeded: ",err);
end

println("Testing the sensitivity of the rotation (simple version - fixed theta/phi and b):");
println("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
mr2,Jrot = rotateAndMoveRBFsimple(m,Mesh,theta_phi,b;computeJacobian = 1,numParamOfRBF=numParamOfRBF);
m0 = m;
dm = 0.01*randn(length(m));

# for k=1:nRBF
	# offset = (k-1)*numParamOfRBF + 1;
	# dm[offset] = 0.01*randn();
	# dm[(offset+1):(offset+6)] = 0.01*randn(6);
	# dm[(offset+7):(offset+9)] = 0.01*randn(3);
# end

hh = 1.0;
for k=0:5
	hhh = (0.5^k)*hh;
	mt = rotateAndMoveRBFsimple(m + hhh*dm,Mesh,theta_phi,b;computeJacobian = 0,numParamOfRBF=numParamOfRBF)[1];
	mtt= mr2 + reshape(Jrot*hhh*dm,size(mr2));
	println("norm(mt-mr2): ",norm(mt[:]-mr2[:]),", norm(mt - mr2 - Jrot*dm): ",norm(mt[:] - mr2[:] - Jrot*hhh*dm));
end

println("Testing the sensitivity of the rotation (with varying theta/phi and b):");
println("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
nRBF = div(length(m),numParamOfRBF);
nMoves = size(theta_phi,1);
m_wrapped = wrapRBFparamAndRotationsTranslations(m,theta_phi,b);
m_t,theta_phi_t,b_t = splitRBFparamAndRotationsTranslations(m_wrapped,nRBF,nMoves,numParamOfRBF);

if norm(m_t - m) > 1e-5 || norm(theta_phi_t - theta_phi)>1e-5 || norm(b_t - b)>1e-5
	error("wrap/split RBFparamAndRotationsTranslations failed");
end

m0 = wrapRBFparamAndRotationsTranslations(m,theta_phi,b);
dm = 0.01*randn(length(m0));

dm,dtheta_phi,db = splitRBFparamAndRotationsTranslations(m_wrapped,nRBF,nMoves,numParamOfRBF);
dm = wrapRBFparamAndRotationsTranslations(dm,dtheta_phi,db);

mr2,Jrot = rotateAndMoveRBF(m,Mesh,theta_phi,b;computeJacobian = 1,numParamOfRBF=numParamOfRBF);

hh = 1.0;
for k=0:10
	hhh = (0.5^k)*hh;
	(mt,theta_phi_t,b_t) = splitRBFparamAndRotationsTranslations(m0+hhh*dm,nRBF,nMoves,numParamOfRBF);
	mt2 = rotateAndMoveRBF(mt,Mesh,theta_phi_t,b_t;computeJacobian = 0,numParamOfRBF=numParamOfRBF)[1];
	println("norm(mt2-mr2): ",norm(mt2[:]-mr2[:]),", norm(mt2 - mr2 - Jrot*dm): ",norm(mt2[:] - mr2[:] - Jrot*hhh*dm));
end


# Процессы для которых требуется указать только число точек

poisson(n; dim=2) = rand(dim, n)

function circle_uniform(n)
    ws = range(0., length=n, step=2π / n)
    [f(w) for f in (cos, sin), w in ws]
end

function circle_poisson(n)
    ws = rand(n) * 2π
    [f(w) for f in (cos, sin), w in ws]
end

normal(n; dim=2) = rand(Normal(), dim, n)


"""Model of algorithms and competition"""

module init

    export init_game

    using Parameters, Statistics, Combinatorics, NLsolve, StructTypes

    # Model properties
    @with_kw mutable struct model
        """Default Properties"""
        n::Int8 = 2                             # Number of firms
        alpha::Float32 = 0.15                   # Learning parameter
        beta::Float32 = 4e-6                    # Learning parameter
        delta::Float32 = 0.95                   # Discount factor
        mu::Float32 = 0.25                      # Product differentiation
        c::Int8 = 1                             # Marginal cost
        a::Int8 = 2                             # Value of the product
        a0::Int8 = 0                            # Value of the outside option
        k::Int8 = 15                            # Dimension of the price grid
        tstable::Int32 = 1e5                    # Number of iterations needed for stability
        tmax::Int32 = 1e7                       # Maximum number of iterations

        """Derived Properties"""
        sdim::Array{Int8,1} = zeros(1)          # Dimension of the state
        s0::Array{Int8,1} = zeros(1)            # Initial state
        A::Array{Float32,1} = zeros(1)          # Action space
        p_minmax::Array{Float32,2} = zeros(1,1) # Minimum and maximum prices
        PI::Array{Float32,3} = zeros(1,1,1)     # Profit matrix
        Q::Array{Float32,4} = zeros(1,1,1,1)    # Q-function of the firms

    end

    # Save struct type
    StructTypes.StructType(::Type{model}) = StructTypes.Mutable()

    function demand(game::model, p::Array{Float32,1})::Array{Float32,1}
        """Compute Demand"""
        e = exp.((game.a .- p) ./ game.mu);
        d = e ./ (sum(e) + exp(game.a0 / game.mu));
        return d
    end

    function foc(game::model, p::Array{Float32,1})::Array{Float32,1}
        """Compute first order condition"""
        d = demand(game, p);
        zero = 1 .- (p .- game.c) .* (1 .- d) ./ game.mu;
        return zero
    end

    function foc_monopoly(game::model, p::Array{Float32,1})::Array{Float32,1}
        """Compute first order condition of a monopolist"""
        d = demand(game, p);
        d1 = reverse(d, dims=1);
        p1 = reverse(p, dims=1);
        zero = 1 .- (p .- game.c) .* (1 .- d) ./ game.mu .+ (p1 .- game.c) .* d1 ./ game.mu;
        return zero
    end

    function compute_p_competitive_monopoly(game::model)::Array{Float32,2}
        """Computes competitive and monopoly prices"""
        p0 = Float32.(ones(game.n) * 3);
        p_competitive = nlsolve((p -> foc(game, p)), p0).zero;
        p_monopoly = nlsolve((p -> foc_monopoly(game, p)), p0).zero;
        return [p_competitive p_monopoly]
    end

    function init_actions(game::model)::Array{Float32,1}
        """Get action space of the firms"""
        a = range(min(game.p_minmax[:,1]...), max(game.p_minmax[:,2]...), length=game.k-2);
        delta = a[2] - a[1];
        A = collect(a[1]-delta:delta:a[end]+2*delta)
        return A
    end

    function init_state(game::model)::Tuple{Array{Int8,1},Array{Int8,1}}
        """Get state dimension and initial state"""
        sdim = [game.k, game.k];
    	s0 = Int8.([1,1,1]);
        return sdim, s0
    end

    function compute_profits(game::model, p::Array{Float32,1})::Array{Float32,1}
        """Compute payoffs"""
        d = demand(game, p);
        profits = (p .- game.c) .* d;
        return profits
    end

    function init_PI(game::model)::Array{Float32,3}
        """Initialize Profits (k^n x n)"""
        PI = zeros(game.k, game.k, game.n);
        for a1=1:game.k
            for a2=1:game.k
                p = [game.A[a1], game.A[a2]];
                PI[a1, a2, :] = compute_profits(game, p);
            end
        end
        return PI
    end

    function init_Q(game::model)::Array{Float32,4}
        """Initialize Q function (n x #s x k)"""
        Q = zeros(game.n, game.k, game.k, game.k);
        for n=1:game.n
            profits = reshape(mean(game.PI[:, :, n], dims=3-n), (1, 1, 1, game.k));
            Q[n, :, :, :] = repeat(profits ./ (1 - game.delta), outer=[1, game.k, game.k, 1]);
        end
        return Q
    end

    function init_game()::model
        """This function initializes the game"""
        game = model();
        game.p_minmax = compute_p_competitive_monopoly(game);
        game.A = init_actions(game);
        game.sdim, game.s0 = init_state(game);
        game.PI = init_PI(game);
        game.Q = init_Q(game);
        return game
    end

end


using COSMO
using Distributions
using Plots
using PortfolioOpt
using PortfolioOpt.TestUtils:
    backtest_po, get_test_data, mean_variance,
    percentchange, timestamp, rename!

############ Read Prices #############
prices = get_test_data()
numD, numA = size(prices) # A: Assets    D: Days

############# Calculating returns #####################
returns_series = percentchange(prices)

# risk free asset
# rf =  fill(3.2e-4, numD)

############# backtest Parameters #####################
DEFAULT_SOLVER = optimizer_with_attributes(
    COSMO.Optimizer, "verbose" => false, "max_iter" => 900000
)

start_date = timestamp(returns_series)[100]

## Strategies
# Look back days
k_back = 60

############# backtest with limit return strategies #####################

wealth_markowitz_limit_R, strategy_returns =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        R = 0.0015 / current_wealth
        formulation = MeanVariance(;
            predicted_mean = r̄,
            predicted_covariance = Σ,
        )

        # model
        model = po_min_variance_limit_return(formulation, R; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_markowitz_limit_R, :markowitz_limit_R);

wealth_markowitz_infeasible_limit_R, strategy_returns =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        R = 0.005 / current_wealth
        formulation = MeanVariance(;
            predicted_mean = r̄,
            predicted_covariance = Σ,
        )

        # model
        model = po_min_variance_limit_return(formulation, R; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_markowitz_infeasible_limit_R, :markowitz_infeasible_limit_R);

wealth_soyster_limit_R, strategy_returns =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        R = 0.0015 / current_wealth
        formulation = RobustBertsimas(;
            predicted_mean = r̄,
            predicted_covariance = Σ,
            uncertainty_delta = std(returns[(end - k_back):end, :]; dims=1)'[:,1] / 5,
            bertsimas_budget = Float64(numA),
        )

        # model
        model = po_min_variance_limit_return(formulation, R; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_soyster_limit_R, :soyster_limit_R);

wealth_bertsimas_limit_R, strategy_returns =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        R = 0.0015 / current_wealth
        formulation = RobustBertsimas(;
            predicted_mean = r̄,
            predicted_covariance = Σ,
            uncertainty_delta = std(returns[(end - k_back):end, :]; dims=1)'[:,1] / 5,
            bertsimas_budget = 3.0,
        )

        # model
        model = po_min_variance_limit_return(formulation, R; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_bertsimas_limit_R, :bertsimas_limit_R);

wealth_bental_limit_R, strategy_returns =
backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        R = 0.0015 / current_wealth
        formulation = RobustBenTal(;
            predicted_mean = r̄,
            predicted_covariance = Σ,
            uncertainty_delta = 0.025
        )

        # model
        model = po_min_variance_limit_return(formulation, R; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_bental_limit_R, :bental_limit_R);

############# backtest with limit variance strategies #####################

wealth_markowitz_limit_var, returns_markowitz_limit_var =
backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ_s, r̄_s = mean_variance(returns[(end - 14):end, :])
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        max_risk = 0.8
        formulation = MeanVariance(;
            predicted_mean = r̄_s,
            predicted_covariance = Σ,
        )

        # model
        model = po_max_return_limit_variance(formulation, max_risk; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_markowitz_limit_var, :markowitz_limit_var);

wealth_soyster_limit_var, returns_soyster_limit_var =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ_s, r̄_s = mean_variance(returns[(end - 14):end, :])
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        max_risk = 0.8
        formulation = RobustBertsimas(;
            predicted_mean = r̄_s,
            predicted_covariance = Σ,
            uncertainty_delta = std(returns[(end - k_back):end, :]; dims=1)'[:,1] / 3,
            bertsimas_budget = Float64(numA),
        )

        # model
        model = po_max_return_limit_variance(formulation, max_risk; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_soyster_limit_var, :soyster_limit_var);

wealth_bertsimas_2_limit_var, returns_bertsimas_2_limit_var =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ_s, r̄_s = mean_variance(returns[(end - 14):end, :])
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        max_risk = 0.8
        formulation = RobustBertsimas(;
            predicted_mean = r̄_s,
            predicted_covariance = Σ,
            uncertainty_delta = std(returns[(end - k_back):end, :]; dims=1)'[:,1] / 3,
            bertsimas_budget = 2.0,
        )

        # model
        model = po_max_return_limit_variance(formulation, max_risk; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_bertsimas_2_limit_var, :bertsimas_2_limit_var);

wealth_bertsimas_4_limit_var, returns_bertsimas_4_limit_var =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ_s, r̄_s = mean_variance(returns[(end - 14):end, :])
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        max_risk = 0.8
        formulation = RobustBertsimas(;
            predicted_mean = r̄_s,
            predicted_covariance = Σ,
            uncertainty_delta = std(returns[(end - k_back):end, :]; dims=1)'[:,1] / 3,
            bertsimas_budget = 4.0,
        )

        # model
        model = po_max_return_limit_variance(formulation, max_risk; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_bertsimas_4_limit_var, :bertsimas_4_limit_var);

wealth_bental_limit_var, returns_bental_limit_var =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ_s, r̄_s = mean_variance(returns[(end - 14):end, :])
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

        # Parameters
        max_risk = 0.8
        formulation = RobustBenTal(;
            predicted_mean = r̄_s,
            predicted_covariance = Σ,
            uncertainty_delta = 0.028
        )

        # model
        model = po_max_return_limit_variance(formulation, max_risk; rf = rf)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_bental_limit_var, :bental_limit_var);

############# backtest with Data-Driven RO strategies #####################
wealth_betina, returns_betina =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        # Parameters
        j_robust = 25
        R = -0.015 / current_wealth
        num_train = 20
        kst = 14
        klt = 30
        kmom = 20
        Q = 2

        formulation = RobustBetina(;
            sampled_returns = returns[(end - j_robust):end, :]
        )

        # Predict return
        r_bar_t = zeros(numA)
        for i in 1:numA
            r_bar_t[i] = mixed_signals_predict_return(
                returns[:, i], num_train, kst, klt, kmom, Q
            )
        end

        # model
        model = po_max_predicted_return_limit_return(formulation, R; predicted_mean = r_bar_t)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_betina, :betina);

############# backtest with DRO strategies #####################

wealth_delague, returns_delague_limit_var =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ_s, r̄_s = mean_variance(returns[(end - 14):end, :])
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])
        # Parameters
        # a =  [0.020646446303211906; 0.01456965970660176; 0.011268181574938444; 0.009189943838602721; 0.007760228688512961; 0.006716054241470283; 0.005919840492702679; 0.005292563196132596; 0.004785576102825234; 0.004367285293323039; 0.004367285293323039]
        # b = [12351.417663268794; 11949.933298064285; 11797.83904955914; 11743.66244416854; 11734.986173090612; 11749.533056761096; 11776.549758689594; 11810.379744602318; 11847.862045043921; 11887.152715141867; 11887.152715141867]

        formulation = RobustDelague(;
            predicted_mean = r̄_s,
            predicted_covariance = Σ,
            γ1 = 0.0065,
            γ2 = 5.5,
            utility_coeficients = [1.0],
            utility_intercepts = [0.0],
        )

        model = po_max_utility_return(formulation)
        x = compute_solution(model, DEFAULT_SOLVER)
        return x * current_wealth
    end;
rename!(wealth_delague, :delague);

############# extra strategies #####################

wealth_sharpe, returns_sharpe =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        returns = values(past_returns)
        Σ, r̄ = mean_variance(returns[(end - k_back):end, :])
        # Parameters
        x = max_sharpe(Σ, r̄, rf)
        return x * current_wealth
    end;
rename!(wealth_sharpe, :sharpe);

wealth_equal_weights, returns_equal =
    backtest_po(returns_series; start_date=start_date) do past_returns, current_wealth, rf
        # Prep
        numD, numA = size(past_returns)
        # Parameters
        x = equal_weights(numA)
        return x * current_wealth
    end;
rename!(wealth_equal_weights, :equal_weights);

############# plot results  #####################
## Wealth
plt = plot(
    wealth_sharpe;
    title="Culmulative Wealth",
    xlabel="Time",
    ylabel="Wealth",
    legend=:outertopright,
);
plot!(plt, wealth_equal_weights);
# plot!(plt, wealth_markowitz_limit_R);
# plot!(plt, wealth_soyster_limit_R);
# plot!(plt, wealth_bertsimas_limit_R);
# plot!(plt, wealth_bental_limit_R);
# plot!(plt, wealth_markowitz_infeasible);
plot!(plt, wealth_markowitz_limit_var);
plot!(plt, wealth_soyster_limit_var);
plot!(plt, wealth_bertsimas_2_limit_var);
plot!(plt, wealth_bertsimas_4_limit_var);
plot!(plt, wealth_bental_limit_var);
plot!(plt, wealth_delague);
plot!(plt, wealth_betina);
plt

## Stats
plt = scatter(
    [std(returns_sharpe)],
    [mean(returns_sharpe)];
    label="Sharpe",
    title="Mean Vs Std",
    xlabel="σ",
    ylabel="E[r]",
    legend=:outertopright,
);
scatter!(plt, [std(returns_equal)], [mean(returns_equal)]; label="Equal");
scatter!(
    plt,
    [std(returns_markowitz_limit_var)],
    [mean(returns_markowitz_limit_var)];
    label="Markowitz",
);
scatter!(
    plt,
    [std(returns_soyster_limit_var)],
    [mean(returns_soyster_limit_var)];
    label="Soyster",
);
scatter!(
    plt,
    [std(returns_bertsimas_2_limit_var)],
    [mean(returns_bertsimas_2_limit_var)];
    label="Bertsimas 2",
);
scatter!(
    plt,
    [std(returns_bertsimas_4_limit_var)],
    [mean(returns_bertsimas_4_limit_var)];
    label="Bertsimas 4",
);
scatter!(
    plt, [std(returns_bental_limit_var)], [mean(returns_bental_limit_var)]; label="Bental"
);
scatter!(
    plt,
    [std(returns_delague_limit_var)],
    [mean(returns_delague_limit_var)];
    label="Delague",
);
scatter!(plt, [std(returns_betina)], [mean(returns_betina)]; label="Betina");
plt


using JuMP
using Gurobi

const output_flag = 0

function problem_setup(m, K, seed; total_resource=500.0, u = 1, v = 10)
    rng_gt = MersenneTwister(seed + 200)

    b = ones(K) * m * total_resource / K
    # Assume cost is normally distributed around a mean_ik for each good and supplier
    # pre compute costs
    C_mean = u + (v - u) * rand(rng_gt, m, K)
    C = randn(rng_gt, m, n, K)
    for j = 1:n
        C[:, j, :] = C_mean + C[:, j, :] * 0.2
    end
    # r_c_mean = u + (v - u) * rand(rng_gt, m)
    r_c_mean = mean(C_mean, 2) + randn(rng_gt) * 0.2
    temp = r_c_mean
    r_c_mean = zeros(m)
    for i = 1:m
        r_c_mean[i] = temp[i]
    end

    println("Problem:")
    println("b = $b")
    println("Mean Cost Per Item, Supplier")
    for k = 1:K
        println("Supplier $k, mean item cost = $(C_mean[:, k])")
    end
    println("Mean Bid Pricing = $r_c_mean")

    rng = MersenneTwister(seed)
    bid_generator = function () return rand(rng, 0:1, m) end
    c_generator = function (j) return C[:, j, :] end
    pi_generator = function (a_k) return transpose(r_c_mean) * a_k + randn(rng) * 0.2 end

    return b, bid_generator, c_generator, pi_generator, rng
end

function offline_lp(env, n, m, K, b, bid_generator::Function, c_generator::Function, pi_generator::Function)
    # collect all data first
    q = zeros(n + 1, K) # resources used
    A = zeros(m, n)
    C = zeros(m, n, K)
    π = zeros(n)
    value = zeros(n + 1)

    for j = 1:n
        A[:, j] = bid_generator()
        C[:, j, :] = c_generator(j)
        π[j] = pi_generator(A[:, j])
    end

    _, prices, X, Y, _ = solve_primal(env, n, m, K, b, A, C, π)

    for j = 1:n
        value[j + 1] = value[j] + X[j] * π[j] - sum(Y[:, j, :] .* C[:, j, :])
        for k = 1:K
            q[j + 1, k] = q[j, k] + sum(Y[:, j, k])
        end
    end

    return prices, X, value, q

end

function online_lp(env, n, h_vector, m, K, b, bid_generator::Function, c_generator::Function, pi_generator::Function;
    greedy=false, print_every=100)
    q = zeros(n + 1, K) # resources used
    A = zeros(m, n)
    C = zeros(m, n, K)
    π = zeros(n)
    prices = zeros(K, n)
    X = zeros(n)
    Y = zeros(m, n, K)
    value = zeros(n + 1)

    # initialize price so that first k bids are not fulfilled
    h = shift!(h_vector)
    h1 = h
    for j = 1:n
        # draw bid
        A[:, j] = bid_generator()
        C[:, j, :] = c_generator(j)
        π[j] = pi_generator(A[:, j])

        # solve ip to get y
        if j >= h1
            feasible, y, obj = ip_sub(env, m, K, b - q[j, :], A[:, j], C[:, j, :], prices[:, j])
            if feasible == true && (π[j] > obj || greedy == true)
                X[j] = 1.0
                Y[:, j, :] = y
            else
                X[j] = 0.0
                Y[:, j, :] = 0
            end
        else
            X[j] = 0.0
            Y[:, j, :] = 0
        end

        for k = 1:K
            q[j + 1, k] = q[j, k] + sum(Y[:, j, k])
        end
        value[j + 1] = value[j] + X[j] * π[j] - sum(Y[:, j, :] .* C[:, j, :])

        # learning step
        if j == h && j < n
            if !isempty(h_vector) h = shift!(h_vector) end
            _, prices[:, j + 1], _, _, _ = solve_primal(env, j, m, K, (b - q[j, :]) * (j / n), A[:, 1:j], C[:, 1:j, :], π[1:j])
        elseif j < n
            prices[:, j + 1] = prices[:, j]
        end
    end
    return prices, X, value, q

end

function solve_primal(env, n, m, K, b, A, C, π)
    model = Model(solver=GurobiSolver(env, OutputFlag = 0))
    @variable(model, 0 <= x[1:n] <= 1)
    @variable(model, y[1:m, 1:n, 1:K] >= 0)

    resource_constraint = []
    for i = 1:m
        for j = 1:n
            push!(resource_constraint, @constraint(model, sum(y[i, j, :]) == A[i, j] * x[j]))
        end
    end

    production_constraint = []
    for k = 1:K
        push!(production_constraint, @constraint(model, sum(y[:, :, k]) <= b[k]))
    end

    @objective(model, Max, sum(x .* π) - sum(C .* y))

    status = solve(model)

    if status == :Optimal
        return getdual(resource_constraint), getdual(production_constraint), getvalue(x), getvalue(y), getobjectivevalue(model)
    else
        print(model)
        error("No production prices found. Status = $status")
    end
end

function solve_dual(env, n, m, K, b, A, C, π)
    model = Model(solver = GurobiSolver(env, OutputFlag=0))
    @variable(model, λ[1:m, 1:n])
    @variable(model, p[1:K] >= 0)
    @variable(model, μ[1:n] >= 0)

    c1 = []
    for j = 1:n
        push!(c1, @constraint(model, μ[j] + sum(A[:, j] .* λ[:, j]) >= π[j]))
    end

    c2 = []
    for i = 1:m
        for j = 1:n
            for k = 1:K
                push!(c2, @constraint(model, λ[i, j] - p[k] <= C[i, j, k]))
            end
        end
    end

    @objective(model, Min, sum(b .* p) + sum(μ))

    status = solve(model)

    if status == :Optimal
        return getdual(c1), getdual(c2), getvalue(λ), getvalue(p), getvalue(μ), getobjectivevalue(model)
    else
        print(model)
        error("No optimal solution to dual found. Status = $status")
    end
end

function ip_sub(env, m, K, b, A, C, p)
    model = Model(solver=GurobiSolver(env, OutputFlag=output_flag))
    @variable(model, λ[1:m])
    @variable(model, 0 <= y[1:m, 1:K] <= 1)

    for i = 1:m
        @constraint(model, λ[i] == sum(y[i, :] .* (p + C[i, :])))
    end

    for i = 1:m
        @constraint(model, sum(y[i, :]) == A[i])
    end

    for k = 1:K
        @constraint(model, sum(y[:, k]) <= b[k])
    end

    @objective(model, Min, sum(λ .* A))

    status = solve(model, suppress_warnings=true)

    if status == :Optimal
        return true, getvalue(y), getobjectivevalue(model)
    elseif status == :Infeasible
        return false, zeros(m, K), Inf
    else
        print(model)
        error("Subproblem not optimal or infeasible. Status = $status")
    end
end

function simulation_q6(n, m, K; seed=1234)
    b, bg, cg, pig, rng = problem_setup(m, K, seed)

    env = Gurobi.Env()

    prices = zeros(K, n, 5)
    X = zeros(n, 5)
    value = zeros(n + 1, 5)
    q = zeros(n + 1, K, 5)

    # offline lp
    srand(rng, seed)
    println("Computing offline LP")
    prices_opt, X[:, 1], value[:, 1], q[:, :, 1] = offline_lp(env, n, m, K, b, bg, cg, pig)
    for k = 1:K
        prices[k, :, 1] = prices_opt[k]
    end

    # online lp with k_vector
    h_vector = [50]
    srand(rng, seed)
    println("Computing online LP, h = 50")
    prices[:, :, 2], X[:, 2], value[:, 2], q[:, :, 2]  = online_lp(env, n, h_vector, m, K, b, bg, cg, pig)

    h_vector = [100]
    srand(rng, seed)
    println("Computing online LP, h = 100")
    prices[:, :, 3], X[:, 3], value[:, 3], q[:, :, 3]  = online_lp(env, n, h_vector, m, K, b, bg, cg, pig)

    h_vector = [200]
    srand(rng, seed)
    println("Computing online LP, h = 200")
    prices[:, :, 4], X[:, 4], value[:, 4], q[:, :, 4]  = online_lp(env, n, h_vector, m, K, b, bg, cg, pig)

    # greedy
    h_vector = [50]
    srand(rng, seed)
    println("Computing online greedy LP, h = 50")
    prices[:, :, 5], X[:, 5], value[:, 5], q[:, :, 5]  = online_lp(env, n, h_vector, m, K, b, bg, cg, pig, greedy=true)

    return b, prices, X, value, q
end


#=
Written by Jackson Thea for CPSC 490 Fall 2015

Provides
  akpw! and akpw (a wrapper for akpw!)

AKPW is an algorithm published by Alon, Karp, Peleg and West in 1995 to produce low stretch spanning trees. This
is my implementation of the algorithm. For more information on how to run AKPW, see the doc string above akpw!'s
function declaration.
=#

function normalizeEdgeWeights(mat::SparseMatrixCSC, kind, nEdges)

  if (kind == :max)
    for eInd in 1:nEdges
      mat.nzval[eInd] \= 1.0
    end #for
  end #if

  minEdgeWeight = minimum(mat.nzval)

  for eInd in 1:nEdges
    mat.nzval[eInd] /= minEdgeWeight
  end #for

  return minEdgeWeight
end #normalizeEdgeWeights

function divideEdgesIntoClasses(mat::SparseMatrixCSC, nEdges, rows, columns, edgeWeights)
  nVertices = mat.n
  x = exp(sqrt(log(nVertices) * log(log(nVertices))))
  y = x * log(nVertices) #placeholder for testing
  edgeClasses = zeros(Int64, nEdges)

  for eInd in 1:nEdges
    edgeClasses[eInd] = floor(log(y, edgeWeights[eInd])) + 1
  end #for

  return edgeClasses

end #divideEdgesIntoClasses

function unsortedArray(n)
  permutation = zeros(Int64, n)
  for i in 1:n
    permutation[i] = i
  end

  for i in 1:n
    j = rand(1:i)
    permutation[i] = permutation[j]
    permutation[j] = i
  end #for

  return permutation
end #unsortedArray

#could speed this up by creating linked list for clusterCount (cuz it'll usually be mostly zeros)
#this makes partitionList potentially unsorted. Will that mess shit up..?
# should be greatest total weight! not just the greatest number of neighbors
function reshuffleClusters(
  mat,
  partitionList,
  vertexToCluster,
  starts,
  vertexToClusterLocation,
  finalRoundClusterVertices,
  rows,
  columns,
  edgeWeights
)
  nVertices = mat.n
  visited = zeros(Bool, nVertices)
  nClusters = length(partitionList)

  for vInd in 1:nVertices
    validVertex = true
    if starts[vertexToCluster[vInd]] == vInd || !finalRoundClusterVertices[vInd]
      continue
    end #if

    clusterNeighborCount = Dict{Int64, Int64}()

    newCluster = vertexToCluster[vInd]
    newClusterCount = 0

    for eInd in mat.colptr[vInd]:(mat.colptr[vInd+1]-1) #eInd is the edge
      otherV = mat.rowval[eInd]

      # clusterNeighborCount[vertexToCluster[otherV]] += edgeWeights[eInd]
      if haskey(clusterNeighborCount, vertexToCluster[otherV])
        clusterNeighborCount[vertexToCluster[otherV]] += 1
      else
        clusterNeighborCount[vertexToCluster[otherV]] = 1
      end #if

      if clusterNeighborCount[vertexToCluster[otherV]] > newClusterCount
        newClusterCount = clusterNeighborCount[vertexToCluster[otherV]]
        newCluster = vertexToCluster[otherV]
      end

      # if any of neighbors
        # 1: in same cluster
        # 2: have no other neighbors in that cluster
        #   don’t shuffle vertex
      if vertexToCluster[otherV] == vertexToCluster[vInd]
        hasOtherNeighborsInOwnCluster = false
        for otherEInd in mat.colptr[otherV]:(mat.colptr[otherV+1]-1)
          otherOtherV = mat.rowval[otherEInd]
          if otherOtherV != vInd && vertexToCluster[otherOtherV] == vertexToCluster[vInd]
            hasOtherNeighborsInOwnCluster = true
          end #if
        end #for
        if !hasOtherNeighborsInOwnCluster
          validVertex = false
        end
      end #if

    end #for

    if !validVertex
      continue
    end

    if newCluster != vertexToCluster[vInd]

      #remove vInd from it's old clusterList
      partitionList[vertexToCluster[vInd]][vertexToClusterLocation[vInd]] = -1

      #push to new clusterList
      push!(partitionList[newCluster], vInd)

      # change in vertexToCluster
      vertexToCluster[vInd] = newCluster
    end #if
  end #for

  for i in 1:length(partitionList)
    sort!(partitionList[i])
    for j in 1:length(partitionList[i])
      if partitionList[i][j] != -1
        vertexToClusterLocation[partitionList[i][j]] = j
      end
    end #for
  end #for

  return partitionList, vertexToCluster, starts, vertexToClusterLocation
end #reshuffleClusters

# takes a matrix from which partition will be constructed (referred to as the parent matrix)
# a partition is a double array of nodes. each sub-array c (cluster) has n vertices (the n of the parent mat)
# and either a 1 or 0 indicating whether or not v is in cluster c
# this function also returns an array of edge indices (of the parent mat) with the edges in the tree added

# only shuffle vertices at the last level of BFS cluster Tree
function partitionMatrix(
  mat::SparseMatrixCSC,
  bigIteration,
  edgeClasses,
  bigEdgeMapReversed,
  bigMatNVertices,
  rows,
  columns,
  edgeWeights,
  randomClusters,
  shuffleClusters,
  exponentialX
)

  nVertices = mat.n
  partitionList = Array{Int64, 1}[]
  nClusters = 0
  visited = zeros(Bool, nVertices)
  nAddedToClusters = 0
  vertexToCluster = zeros(Int64, nVertices)
  vertexToClusterLocation = zeros(Int64, nVertices) #returns the index where the vertex is located inside clusterList
  starts = Int64[]
  newVertices = zeros(Int64, nVertices)
  nNewVertices = 0
  newVerticesArray = zeros(Bool, nVertices)

  if randomClusters
    unsortedNodes = unsortedArray(nVertices)
  end

  finalRoundClusterVertices = zeros(Bool, nVertices)

  x::Float64

  if nVertices > 2
    if exponentialX
      x = exp(sqrt(log(nVertices) * log(log(nVertices))))
    else
      x = log(nVertices+1)/log(2)
    end
  else
    x = 5.0
  end

  lowestUnvisitedUnsortedIndex = 1

  while (nAddedToClusters != nVertices)
    nClusters += 1
    push!(partitionList, Int64[])

    # pick a random starting node between 1 and nVertices each time
    if randomClusters
      while visited[unsortedNodes[lowestUnvisitedUnsortedIndex]]
        lowestUnvisitedUnsortedIndex += 1
      end #while
      start = unsortedNodes[lowestUnvisitedUnsortedIndex]
    else
      # OR (if not random), just select next unvisited node
      while visited[lowestUnvisitedUnsortedIndex]
        lowestUnvisitedUnsortedIndex += 1
      end #while
      start = lowestUnvisitedUnsortedIndex
    end

    push!(starts, start)
    visited[start] = true

    nNodesInCluster = 1
    nAddedToClusters += 1

    push!(partitionList[nClusters], start)
    vertexToCluster[start] = nClusters
    vertexToClusterLocation[start] = nNodesInCluster

    lastnNodesInCluster = 0

    calculatedVolume = 0.0
    calculatedBoundary = 0.0

    for eInd in mat.colptr[start]:(mat.colptr[start+1]-1)
      calculatedVolume += 1.0
      calculatedBoundary += 1.0
    end #for

    #building a single cluster
    while !(nAddedToClusters == nVertices)

      if (calculatedVolume / calculatedBoundary >= x) || lastnNodesInCluster == nNodesInCluster
        break
      end #if

      lastnNodesInCluster = nNodesInCluster

      for i in 1:nNewVertices
        newVerticesArray[newVertices[i]] = 0
      end #for

      nNewVertices = 0
      for v in partitionList[nClusters] #for each vertex in the current cluster
        for eInd in mat.colptr[v]:(mat.colptr[v+1]-1) #eInd is the edge

          if edgeClasses[bigEdgeMapReversed[eInd]] <= bigIteration
            otherV = mat.rowval[eInd] #this gives us the other vertex
            #doing this step before for boundary... maybe can combine?
            if !visited[otherV]
              if !newVerticesArray[otherV]
                nNewVertices += 1
                newVertices[nNewVertices] = otherV
              end #if
              newVerticesArray[otherV] = true
            end #if
          end #if

        end #for
      end #for

      for nthNewVertex in 1:nNewVertices
        newV = newVertices[nthNewVertex]
        push!(partitionList[nClusters], newV)
        vertexToCluster[newV] = nClusters
        nAddedToClusters += 1
        nNodesInCluster += 1
        vertexToClusterLocation[newV] = nNodesInCluster
        visited[newV] = true

        #updating boundary and volume
        for eInd in mat.colptr[newV]:(mat.colptr[newV+1]-1)
          otherV = mat.rowval[eInd]
          if vertexToCluster[otherV] == vertexToCluster[newV]
            calculatedBoundary -= 1.0

            #LEAVE THIS IN IF Vol = sum of degrees of nodes. Take out of Vol = # edges
            calculatedVolume += 1.0
          else
            calculatedVolume += 1.0
            calculatedBoundary += 1.0
          end #if
        end #for

      end #for

    end #while

    for nthNewVertex in 1:nNewVertices
      finalRoundClusterVertices[newVertices[nthNewVertex]] = true
    end #for

  end #while

  if shuffleClusters
    return reshuffleClusters(mat, partitionList, vertexToCluster, starts, vertexToClusterLocation, finalRoundClusterVertices, rows, columns, edgeWeights)
  end

  return partitionList, vertexToCluster, starts, vertexToClusterLocation
end #partitionMatrix

# for a single cluster, updates vertexToCluster to include another cluster for each extra component
# inside the original cluster
function updateVertexToCluster{Tv,Ti}(
  mat::SparseMatrixCSC{Tv,Ti},
  nVertices,
  vertexToCluster,
  c,
  nClusters,
  partitionList
)

  nVerticesInCluster = length(partitionList[c])
  visited = zeros(Bool, nVertices)
  compList = Array{Int64, 1}[]
  order = Array(Int64, nVertices)

  numComps = 0

  for clusterInd in 1:nVerticesInCluster
    vInd = partitionList[c][clusterInd]
    if !visited[vInd]
      numComps += 1
      visited[vInd] = true
      if numComps > 1
        vertexToCluster[vInd] = nClusters + numComps-1
      end

      if mat.colptr[vInd+1] > mat.colptr[vInd]
        ptr = 1
        orderLen = 2
        order[ptr] = vInd

        while ptr < orderLen
          curNode = order[ptr]

          for ind in mat.colptr[curNode]:(mat.colptr[curNode+1]-1)
            nbr = mat.rowval[ind]
            if vertexToCluster[nbr] == c && !visited[nbr]
              visited[nbr] = true
              if numComps > 1
                vertexToCluster[nbr] = nClusters + numComps-1
              end
              order[orderLen] = nbr
              orderLen += 1
            end # if
          end # for
          ptr += 1
        end # while
      end # if
    end
  end #for

  if numComps > 1
    nClusters += numComps-1
  end

  return nClusters
  # return maximum(vertexToCluster)

end #clusterComponents

function partitionListFromMap(nVertices, nClusters, vertexToCluster)
  vertexToClusterLocation = zeros(Int64, nVertices)
  starts = zeros(Int64, nClusters)
  partitionList = Array{Int64, 1}[]

  for i in 1:nClusters
    push!(partitionList, Int64[])
  end

  for i in 1:nVertices
    thisCluster = vertexToCluster[i]
    push!(partitionList[thisCluster], i)
    nVerticesInCluster = length(partitionList[thisCluster])
    vertexToClusterLocation[i] = nVerticesInCluster
    if nVerticesInCluster == 1
      starts[thisCluster] = i
    end #if
  end

  return vertexToClusterLocation, starts, partitionList
end #partitionListFromMap

function metisPartition(mat, nClusters)
  nVertices = mat.n
  g = Graph(mat)

  if nVertices == 2 || nClusters >= nVertices || nClusters <= 2
    vertexToCluster = ones(Int64, nVertices)
    nClusters = 1
  else
    vertexToCluster = partGraphRecursive(g, nClusters)[2]
    while minimum(vertexToCluster) > 1
      vertexToCluster = vertexToCluster.- 1
      nClusters -= 1
    end
  end #if

  vertexToClusterLocation, starts, partitionList = partitionListFromMap(nVertices, nClusters, vertexToCluster)
  for c in 1:nClusters
    nClusters = updateVertexToCluster(mat, nVertices, vertexToCluster, c, nClusters, partitionList) #what if cluster comps is a list of lists?
  end

  vertexToClusterLocation, starts, partitionList = partitionListFromMap(nVertices, nClusters, vertexToCluster)

  return partitionList, vertexToCluster, starts, vertexToClusterLocation, nClusters
end #metisPartition

#definitions for a potentially faster lightgraphs heap

# immutable DijkstraHeapEntry{T}
#   vertex::Int
#   dist::T
# end

# isless(e1::DijkstraHeapEntry, e2::DijkstraHeapEntry) = e1.dist < e2.dist

function shortestPathsForCluster(mat, clusterList, vertexToCluster, start, vertexToClusterLocation)
  nVerticesInMat = mat.n
  nVerticesInCluster = length(clusterList)

  visited = zeros(Bool, nVerticesInCluster)

  nh = intHeap(nVerticesInCluster)
  dists = nh.keys

  pArray = zeros(Int64, nVerticesInCluster)
  pArrayNonZeros = Int64[]

  # sizehint!(nh, nVerticesInCluster)

  intHeapAdd!(nh, vertexToClusterLocation[start], 0.0)

  pArray[vertexToClusterLocation[start]] = vertexToClusterLocation[start]

  while nh.nitems > 0
    vIndInCluster::Int64 = intHeapPop!(nh)
    visited[vIndInCluster] = true

    dv = dists[vIndInCluster]
    vInd = clusterList[vIndInCluster]

    if vInd == -1
      continue
    end

    for eInd in mat.colptr[vInd]:(mat.colptr[vInd+1]-1)
      otherVInd = mat.rowval[eInd]
      otherVIndInCluster = vertexToClusterLocation[otherVInd]

      if vertexToCluster[otherVInd] == vertexToCluster[vInd] && !visited[otherVIndInCluster]
        newdist = dv + mat.nzval[eInd]
        if newdist < dists[otherVIndInCluster]
          dists[otherVIndInCluster] = newdist
          intHeapAdd!(nh, otherVIndInCluster, newdist)
          pArray[otherVIndInCluster] = vIndInCluster
          push!(pArrayNonZeros, otherVIndInCluster)
        end # if
      end # if
    end # for

  end # while

  #trying the lightgraphs heap

  # dists = fill(Inf, nVerticesInCluster)
  # visited = zeros(Bool, nVerticesInCluster)
  # pArray = zeros(Int64, nVerticesInCluster)
  # H = Vector{Float64}()
  # dists[vertexToClusterLocation[start]] = 0.0

  # sizehint!(H, nVerticesInCluster)

  # heappush!(H, DijkstraHeapEntry(vertexToClusterLocation[start], dists[vertexToClusterLocation[start]]))

  # while !isempty(H)
  #   hentry = heappop!(H)
  #   vIndInCluster = hentry.vertex
  #   vInd = clusterList[vIndInCluster]
  #   dv = dists[vIndInCluster]

  #   if vInd == -1
  #     continue
  #   end
  #   # vIndInCluster
  #   for eInd in mat.colptr[vInd]:(mat.colptr[vInd+1]-1)
  #     otherVInd = mat.rowval[eInd]
  #     otherVIndInCluster = vertexToClusterLocation[otherVInd]

  #     if !visited[otherVIndInCluster]
  #       dists[otherVIndInCluster] = dv
  #       pArray[otherVIndInCluster] = vIndInCluster
  #       visited[otherVIndInCluster] = true
  #       heappush!(H, DijkstraHeapEntry{T}(otherVIndInCluster, alt))
  #     else
  #       if dv < dists[otherVIndInCluster]
  #         dists[otherVIndInCluster] = dv
  #         pArray[otherVIndInCluster] = vIndInCluster
  #         heappush!(H, DijkstraHeapEntry{T}(otherVIndInCluster, dv))
  #       end
  #       if dv == dists[otherVIndInCluster]
  #           pathcounts[otherVIndInCluster] += pathcounts[vIndInCluster]
  #       end #if
  #     end #if/else
  #   end #for
  # end #while

  newTreeInds = Int64[]

  for vInd in pArrayNonZeros
    bigMatVInd = clusterList[vInd]
    for eInd in mat.colptr[bigMatVInd]:(mat.colptr[bigMatVInd+1]-1) #all edges connecting to vertex j
      otherV = mat.rowval[eInd]
      clusterIndOtherV = vertexToClusterLocation[otherV]
      if clusterIndOtherV == pArray[vInd] && vertexToCluster[otherV] == vertexToCluster[bigMatVInd]
        push!(newTreeInds, eInd)
      end #if
    end #for
  end #for

  # helpful for debugging a bug that causes a non-connected tree

  # checking connectedness
  # for vInd in 1:length(visited)
  #   if !visited[vInd] && clusterList[vInd] != -1
  #     println("for cluster: ", vertexToCluster[clusterList[vInd]], " vertex: ", clusterList[vInd], " is not connected")
  #     println("vertex: ", clusterList[vInd], " is connected to: ")
  #     for eInd in mat.colptr[clusterList[vInd]]:(mat.colptr[clusterList[vInd]+1]-1)
  #       otherVInd = mat.rowval[eInd]
  #       println("\t otherV: ", otherVInd, " in cluster: ", vertexToCluster[otherVInd])
  #     end
  #   end #if
  # end #for

  return newTreeInds
end # shortestPaths

function collapsePartition(mat::SparseMatrixCSC, partitionList, map, nEdges, rows, columns, edgeWeights)

  nClusters = length(partitionList)
  newRows = Int64[]
  newColumns = Int64[]
  newEdgeWeights = Float64[]
  for i in 1:nEdges
    if map[rows[i]] != map[columns[i]]
      push!(newRows, map[rows[i]])
      push!(newColumns, map[columns[i]])
      push!(newEdgeWeights, edgeWeights[i])
    end #if
  end #for

  newMat = sparse(newRows, newColumns, newEdgeWeights, nClusters, nClusters, min)

  return newMat
end #collapsePartition

function getBigEdgeMapReversed(
  mat,
  newMat,
  bigMapD,
  nEdges,
  rows,
  columns,
  edgeWeights,
  newNEdges,
  newRows,
  newColumns,
  newEdgeWeights
)

  bigEdgeMapReversed = zeros(Int64, newNEdges)

  for eInd in 1:nEdges
    c = columns[eInd]
    r = rows[eInd]

    # this makes this function MUCH MUCH faster -> usually O(E) rather than O(E^2)... sort of
    if bigMapD[c] == bigMapD[r]
      continue
    end

    for e in newMat.colptr[bigMapD[c]]:(newMat.colptr[bigMapD[c]+1]-1) #eInd is the edge
      if newMat.rowval[e] == bigMapD[r] && (bigEdgeMapReversed[e] == 0 || edgeWeights[eInd] < edgeWeights[bigEdgeMapReversed[e]])
        bigEdgeMapReversed[e] = eInd
      end #if
    end #for
  end #for

  return bigEdgeMapReversed
end #getBigEdgeMapReversed

function denormalizeEdgeWeights(oldMinEdgeWeight, mat, kind, nEdges)
  for eInd in 1:nEdges
    mat.nzval[eInd] *= oldMinEdgeWeight
  end #for

  if (kind == :max)
    for eInd in 1:nEdges
      mat.nzval[eInd] \= 1.0
    end #for
  end #if
end #denormalizeEdgeWeights

function sparseMatrixFromTreeIndices(mat, treeInds)
  nVertices = mat.n

  rows, columns, edgeWeights = findnz(mat)

  nRows = Int64[]
  nColumns = Int64[]
  nEdgeWeights = Float64[]

  for i in 1:length(treeInds)
    if treeInds[i]
      push!(nRows, rows[i])
      push!(nColumns, columns[i])
      push!(nEdgeWeights, edgeWeights[i])
    end #if
  end #for

  tree = sparse(nRows,nColumns,nEdgeWeights,nVertices,nVertices, min)
  tree = tree + tree'
  return tree

end #sparseMatrixFromTreeIndices

"""
Constructs a low stretch tree using the Alon, Karp, Peleg, West algorithm. This version (akpw! instead of akpw)
modifies the graph slightly changing the edges weights, then changing them back, which may lead to floating
point imprecisions. akpw! is faster (about 10-20%), but akpw doesn't have float imprecisions.

The function has a few options:

kind: default is :max, which regards each edge weight as the inverse of its length (just like kruskal). If this is
  set to anything else (e.g. :min), it will regard edge weight as length

randomClusters: default is false. This means the partition function searches for the beginning of the next cluster
  in node order, rather than randomly choosing nodes. If this is set to false, it will
  randomly choose the next node. This slows down akpw, but may produce better stretch.

metisClustering: default is false. If this is set to false, the graph will be partitioned
  each time by metis, rather than by the akpw partitioning method.

shuffleClusters: default is true. This preserves the "reshuffleClusters" method after each each graph is
  partitioned into clusters. If set to false, the function will skip this step. May be faster
  but have worse stretch.

exponentialX: default is true, where the funciton exp(sqrt(log(nVertices) * log(log(nVertices)))) is used for X.
  If set fo false, the function log(nVertices+1)/log(2) will be used for X instead.

EXAMPLE:

[2, 1]  =  0.631273
[3, 1]  =  0.40103
[1, 2]  =  0.631273
[4, 2]  =  0.147018
[1, 3]  =  0.40103
[4, 3]  =  0.772661
[2, 4]  =  0.147018
[3, 4]  =  0.772661

      |
      |
      V

[2, 1]  =  0.631273
[3, 1]  =  0.40103
[1, 2]  =  0.631273
[1, 3]  =  0.40103
[4, 3]  =  0.772661
[3, 4]  =  0.772661
"""
function akpw!(
  mat::SparseMatrixCSC;
  kind=:max,
  randomClusters= false,
  metisClustering= false,
  shuffleClusters= true,
  exponentialX= true
)

  # useful for debugging!
  # println("Starting up AKPW...")
  # info("Starting up AKPW...")

  nVertices = mat.n
  nEdges = nnz(mat)

  oldMinEdgeWeight = normalizeEdgeWeights(mat, kind, nEdges)
  rows, columns, edgeWeights = findnz(mat)

  newMat = mat
  newNEdges = nEdges
  newRows, newColumns, newEdgeWeights = rows, columns, edgeWeights

  bigMapD = zeros(Int64, nVertices) #assuming map does C->D, bigMapD does A->D.
  for i in 1:nVertices #the first iteration maps to iself
    bigMapD[i] = i
  end #for

  bigEdgeMapReversed = zeros(Int64, nEdges) #goes C->A with edges (picks min of the many possibilities)
  for i in 1:nEdges #the first iteration maps to iself
    bigEdgeMapReversed[i] = i
  end #for

  treeInds = zeros(Bool, nEdges)
  edgeClasses = divideEdgesIntoClasses(mat, nEdges, rows, columns, edgeWeights)

  bigIteration = 1
  while (true)
    if !metisClustering
      partitionList, vertexToCluster, starts, vertexToClusterLocation = partitionMatrix(newMat, bigIteration, edgeClasses, bigEdgeMapReversed, nVertices, newRows, newColumns, newEdgeWeights, randomClusters, shuffleClusters, exponentialX)
      nClusters = length(partitionList)
    else
      nClusters = Int64(floor((newMat.n)/900)+2)
      # nClusters = Int64(floor(log((newMat.n)+1)/log(2)))
      partitionList, vertexToCluster, starts, vertexToClusterLocation, nClusters = metisPartition(newMat, nClusters)
    end

    # useful for debugging!
    # println("nClusters ", nClusters)

    for i in 1:length(bigMapD)
      bigMapD[i] = vertexToCluster[bigMapD[i]]
    end #for

    #this, for each cluster, creates its own sparse matrix and runs shortest paths on that, then returns the
    # original mapped edges corresponding to that shortest paths tree. Then it adds those indices to treeInds
    for cInd in 1:nClusters
      newTreeInds = shortestPathsForCluster(newMat, partitionList[cInd], vertexToCluster, starts[cInd], vertexToClusterLocation)

      for i in newTreeInds
        treeInds[bigEdgeMapReversed[i]] = true
      end #for
    end #for

    if (nClusters == 1)
      break
    end #if

    newMat = collapsePartition(newMat, partitionList, vertexToCluster, newNEdges, newRows, newColumns, newEdgeWeights)
    newNEdges = nnz(newMat)
    newRows, newColumns, newEdgeWeights = findnz(newMat)

    bigEdgeMapReversed = getBigEdgeMapReversed(mat, newMat, bigMapD, nEdges, rows, columns, edgeWeights, newNEdges, newRows, newColumns, newEdgeWeights)

    bigIteration += 1
  end #while

  denormalizeEdgeWeights(oldMinEdgeWeight, mat, kind, nEdges)

  finalTree = sparseMatrixFromTreeIndices(mat, treeInds)

  return finalTree
end #akpw!

"""
This is a wrapper for akpw!. It's slower, but won't modify the original graph. See akpw! documentation for more
details.
"""
function akpw(
  origMat::SparseMatrixCSC;
  kind=:max,
  randomClusters=false,
  metisClustering=false,
  shuffleClusters=true,
  exponentialX=true
)

  rows, columns, edgeWeights = findnz(origMat)
  mat = sparse(rows, columns, edgeWeights)

  return akpw!(mat, kind=kind, randomClusters=randomClusters, metisClustering=metisClustering, shuffleClusters=shuffleClusters, exponentialX=exponentialX)
end #akpw


"""
     add_singleton(x::AbstractArray, ::Val{dim}) where {dim}

Add an additional dimension `dim` of size 1 to array `x`.
"""
function add_singleton(x::AbstractArray, ::Val{dim}) where {dim}
    shape = ntuple(max(ndims(x) + 1, dim)) do i
        return i < dim ? size(x, i) : (i > dim ? size(x, i - 1) : 1)
    end
    return reshape(x, shape)
end

"""
    dot_matwise(x::AbstractArray, y::AbstractArray)

Compute the inner product of all matrices in `x` and `y`.

At least one of `x` and `y` has to be a matrix.
"""
dot_matwise(x::AbstractMatrix, y::AbstractMatrix) = dot(x, y)
function dot_matwise(x::AbstractArray, y::AbstractMatrix)
    xmat = reshape(x, size(x, 1) * size(x, 2), :)
    return reshape(reshape(y, 1, :) * xmat, size(x)[3:end])
end
dot_matwise(x::AbstractMatrix, y::AbstractArray) = dot_matwise(y, x)

"""
    checksize(μ::AbstractVecOrMat, ν::AbstractVecOrMat, C::AbstractMatrix)

Check that dimensions of source and target marginals `μ` and `ν` are consistent with cost
matrix `C`.
"""
function checksize(μ::AbstractVecOrMat, ν::AbstractVecOrMat, C::AbstractMatrix)
    size(C) == (size(μ, 1), size(ν, 1)) || throw(
        DimensionMismatch("cost matrix `C` must be of size `(size(μ, 1), size(ν, 1))`")
    )
    return nothing
end

"""
    checksize2(x::AbstractVecOrMat, y::AbstractVecOrMat)

Check if arrays `x` and `y` are compatible, then return a tuple of its broadcasted second
dimension.
"""
checksize2(::AbstractVector, ::AbstractVector) = ()
function checksize2(μ::AbstractVecOrMat, ν::AbstractVecOrMat)
    size_μ_2 = size(μ, 2)
    size_ν_2 = size(ν, 2)
    if size_μ_2 > 1 && size_ν_2 > 1 && size_μ_2 != size_ν_2
        throw(DimensionMismatch("size of source and target marginals is not compatible"))
    end
    return (max(size_μ_2, size_ν_2),)
end

"""
     isallapprox(x::AbstractVector)

Check that all entries of `x` are approximately equal
"""
function isallapprox(x::AbstractVecOrMat)
    return all(y -> isapprox(y, x[1]), x[2:end])
end

"""
     checkbalanced(μ::AbstractVecOrMat, ν::AbstractVecOrMat)

Check that source and target marginals `μ` and `ν` are balanced.
"""
function checkbalanced(μ::AbstractVector, ν::AbstractVector)
    sum(μ) ≈ sum(ν) || throw(ArgumentError("source and target marginals are not balanced"))
    return nothing
end
function checkbalanced(x::AbstractVecOrMat, y::AbstractVecOrMat)
    all(isapprox.(sum(x; dims=1), sum(y; dims=1))) ||
        throw(ArgumentError("source and target marginals are not balanced"))
    return nothing
end
function checkbalanced(x::AbstractMatrix)
    isallapprox(sum(x; dims=1)) ||
        throw(ArgumentError("source and target marginals are not balanced"))
    return nothing
end

"""
    A_batched_mul_B!(c::AbstractVector, A::AbstractMatrix, b::AbstractVector)

Compute the matrix-vector product `Ab` and write the result to `c`.
"""
A_batched_mul_B!(c::AbstractVector, A::AbstractMatrix, b::AbstractVector) = mul!(c, A, b)

"""
    A_batched_mul_B!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix)

Compute the matrix-matrix product `AB` and write the result to `C`.
"""
A_batched_mul_B!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix) = mul!(C, A, B)

"""
    A_batched_mul_B!(C::AbstractMatrix, A::AbstractArray{<:Any,3}, B::AbstractMatrix)

Compute the matrix-vector products `A[:, :, i] B[:, i]` and write them to `C[:, i]`.
"""
function A_batched_mul_B!(C::AbstractMatrix, A::AbstractArray{<:Any,3}, B::AbstractMatrix)
    return NNlib.batched_mul!(add_singleton(C, Val(2)), A, add_singleton(B, Val(2)))
end

"""
    At_batched_mul_B!(c::AbstractVector, A::AbstractMatrix, b::AbstractVector)

Compute the matrix-vector product `transpose(A)b` and write the result to `c`.
"""
function At_batched_mul_B!(c::AbstractVector, A::AbstractMatrix, b::AbstractVector)
    return mul!(c, transpose(A), b)
end

"""
    At_batched_mul_B!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix)

Compute the matrix-matrix product `transpose(A)B` and write the result to `C`.
"""
function At_batched_mul_B!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix)
    return mul!(C, transpose(A), B)
end

"""
    At_batched_mul_B!(C::AbstractMatrix, A::AbstractArray{<:Any,3}, B::AbstractMatrix)

Compute the matrix-vector products `transpose(A[:, :, i]) B[:, i]` and write them to
`C[:, i]`.
"""
function At_batched_mul_B!(C::AbstractMatrix, A::AbstractArray{<:Any,3}, B::AbstractMatrix)
    return NNlib.batched_mul!(add_singleton(C, Val(1)), add_singleton(B, Val(1)), A)
end

struct FiniteDiscreteMeasure{X<:AbstractVector,P<:AbstractVector}
    support::X
    p::P

    function FiniteDiscreteMeasure{X,P}(support::X, p::P) where {X,P}
        length(support) == length(p) || error("length of `support` and `p` must be equal")
        isprobvec(p) || error("`p` must be a probability vector")
        return new{X,P}(support, p)
    end
end

"""
    discretemeasure(
        support::AbstractVector,
        probs::AbstractVector{<:Real}=fill(inv(length(support)), length(support))
    )

Construct a finite discrete probability measure with `support` and corresponding
`probabilities`. If the probability vector argument is not passed, then
equal probability is assigned to each entry in the support.

# Examples
```julia
using KernelFunctions
# rows correspond to samples
μ = discretemeasure(RowVecs(rand(7,3)), normalize!(rand(10),1))

# columns correspond to samples, each with equal probability
ν = discretemeasure(ColVecs(rand(3,12)))
```

!!! note
    If `support` is a 1D vector, the constructed measure will be sorted,
    e.g. for `mu = discretemeasure([3, 1, 2],[0.5, 0.2, 0.3])`, then
    `mu.support` will be `[1, 2, 3]` and `mu.p` will be `[0.2, 0.3, 0.5]`.
"""
function discretemeasure(
    support::AbstractVector{<:Real},
    probs::AbstractVector{<:Real}=fill(inv(length(support)), length(support)),
)
    return DiscreteNonParametric(support, probs)
end
function discretemeasure(
    support::AbstractVector,
    probs::AbstractVector{<:Real}=fill(inv(length(support)), length(support)),
)
    return FiniteDiscreteMeasure{typeof(support),typeof(probs)}(support, probs)
end

Distributions.support(d::FiniteDiscreteMeasure) = d.support
Distributions.probs(d::FiniteDiscreteMeasure) = d.p