# This file provides optimal integer-programming-powered algorithms for
# tree layout. JuMP is not a hard dependency of GraphLayout.jl, but if
# a user wishes to use these methods then they will need JuMP and an
# integer programming solver installed, such as GLPK, Cbc, or one of the
# commercial solvers like Gurobi or CPLEX
# Ideally, only when one of these methods is used will JuMP will be
# loaded. However, one of the phases only has a JuMP-dependent version
# so we need to always load JuMP for now.
########################################################################

using JuMP

"""
    Given a layer assignment, decide a permutation for each layer
    that minimizes edge crossings using integer programming.

    Based on the IP described in
      M. Junger, E. K. Lee, P. Mutzel, and T. Odenthal.
      A polyhedral approach to the multi-layer crossing minimization problem.
      In G. Di Battista, editor, Graph Drawing: 5th International Symposium,
      GD ’97, volume 1353 of Lecture Notes in Computer Science, pages 13–24,
      Rome, Italy, September 1997. Springer-Verlag.

    Arguments:
    adj_list        Directed graph in adjacency list format
    layers          Assignment of vertices
    layer_verts     Dictionary of layer => vertices (initial perm.)

    Returns:
    new_layer_verts An improved dictionary of layer => vertices (opt. perm.)
"""
function _ordering_ip{T}(adj_list::AdjList{T}, layers, layer_verts)
    num_layers = maximum(layers)

    m = Model()

    # Define crossing binary variables
    @variable(m, c[L=1:num_layers,        # for each layer
                 i=layer_verts[L],      # for each vertex in this layer
                 j=adj_list[i],         # and vertex in the next layer
                 k=layer_verts[L],      # for each vertex in this layer
                 l=adj_list[k]], Bin)   # and vertex in the next layer

    # Objective: minimize crossings
    @objective(m, Min, sum(c))

    # Define permutation variables for each layer
    # We'll define for both (i,j) and (j,i), and ensure they consistency
    # by adding constraints. Presolve in the IP solver will simplify
    # the problem for us by removing one of the variables.
    @variable(m, x[L=1:num_layers, i=layer_verts[L], j=layer_verts[L]], Bin)
    for L in 1:num_layers
        for i in layer_verts[L]
            for j in layer_verts[L]
                j <= i && continue  # Don't double-add
                # Ensure x[i,j] and x[j,i] are consistent
                @constraint(m, x[L,i,j] == 1 - x[L,j,i])
                # And ensure that triples are consistent
                for k in layer_verts[L]
                    k <= j && continue
                    @constraint(m, 0 <= x[L,i,j] + x[L,j,k] - x[L,i,k] <= 1)
                end
            end
        end
    end

    # Link permutations to crossings
    for L in 1:num_layers-1
        # For all (i,j)
        for i in layer_verts[L]
        for j in adj_list[i]
            # For all (k,l)
            for k in layer_verts[L]
            k == i && continue  # Can't cross if starting from same vertex!
            for l in adj_list[k]
                @constraint(m, -c[L,i,j,k,l] <= x[L+1,j,l] - x[L,i,k])
                @constraint(m,  c[L,i,j,k,l] >= x[L+1,j,l] - x[L,i,k])
            end
            end
        end
        end
    end

    # Solve the IP
    solve(m)

    # Extract permutation from solution
    x_sol = getvalue(x)
    new_layer_verts = [L => Int[] for L in 1:num_layers]
    for L in 1:num_layers
        old_perm = layer_verts[L]
        # For each vertex, count the number of times it is "in front"
        # of the other vertices. The higher this number, the earlier
        # the vertex appears in the layer.
        scores = zeros(length(old_perm))
        for (p,i) in enumerate(old_perm)
            for j in old_perm
                i == j && continue
                if round(Integer, x_sol[L,i,j]) == 1
                    # i appears before j
                    scores[p] += 1
                end
            end
        end
        new_layer_verts[L] = old_perm[sortperm(scores,rev=true)]
    end

    return new_layer_verts
end

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

"""
    Given a layer assignment and permutation, decide the coordinates for
    each vertex. The objective is to encourage straight edges, especially
    for longer edges. This function uses an integer program to decide the
    coordinates (although it is solved as a linear program), as described in
      Gansner, Emden R., et al.
      A technique for drawing directed graphs.
      Software Engineering, IEEE Transactions on 19.3 (1993): 214-230.

    Arguments:
    adj_list        Directed graph in adjacency list format
    layers          Assignment of vertices
    layer_verts     Dictionary of layer => vertices (final perm.)
    orig_n          Number of original (non-dummy) vertices
    widths          Width of each vertex
    xsep            Minimum seperation between each vertex

    Returns:
    layer_coords    For each layer and vertex, the x-coord
"""
function _coord_ip{T}(adj_list::AdjList{T}, layers, layer_verts, orig_n, widths, xsep)
    num_layers = maximum(layers)

    m = Model()

    # One variable for each vertex
    @variable(m, x[L=1:num_layers, i=layer_verts[L]] >= 0)

    # Constraint: must respect permutation, and spacign constraint
    for L in 1:num_layers
        for i in 1:length(layer_verts[L])-1
            a = layer_verts[L][i]
            b = layer_verts[L][i+1]
            @constraint(m, x[L,b] - x[L,a] >=
                (widths[a] + widths[b])/2 + xsep)
        end
    end

    # Objective: minimize total misalignment
    # Use the weights from the Ganser paper:
    #   1 if both nodes "real"
    #   2 if one of the nodes is "real"
    #   8 if neither node is "real"
    # We use absolute distance in the objective so we'll need
    # auxilary variables for each pair of edges
    obj = AffExpr()
    @variable(m, absdiff[L=1:num_layers-1,
                        i=layer_verts[L], j=adj_list[i]] >= 0)
    for L in 1:num_layers-1
        for i in layer_verts[L]
            for j in adj_list[i]
                @constraint(m, absdiff[L,i,j] >= x[L,i] - x[L+1,j])
                @constraint(m, absdiff[L,i,j] >= x[L+1,j] - x[L,i])
                if i > orig_n && j > orig_n
                    # Both dummy vertices
                    obj += 8*absdiff[L,i,j]
                elseif (i <= orig_n && j >  orig_n) ||
                       (i >  orig_n && j <= orig_n)
                    # Only one dummy vertix
                    obj += 2*absdiff[L,i,j]
                else
                    # Both real
                    obj += absdiff[L,i,j]
                end
            end
        end
    end
    @objective(m, Min, obj)

    # Solve it...
    solve(m)

    # ... and mangle the solution into shape
    x_sol = getvalue(x)
    locs_x = zeros(length(layers))
    for L in 1:num_layers
        for i in layer_verts[L]
            locs_x[i] = x_sol[L,i]
        end
    end
    return locs_x
end


function update_weights(data, e, closest_same, closest_other, weights, m, max_f_vals, min_f_vals)
	for t = 1:size(data, 2)
		weights[t] = weights[t] - (abs(e[t] - closest_same[t])/(max_f_vals[t] - min_f_vals[t] + eps(Float64)))/m +
			(abs(e[t] - closest_other[t])/(max_f_vals[t] - min_f_vals[t] + eps(Float64)))/m
	end
	return weights
end


using Printf

include("polynm.jl")

function my_sin(x)
    return sin(10*x)
end

function integ_my_sin(x)
    return -cos(10*x)/10
end

function d_my_sin(x)
    return 10*cos(10*x)
end

function test_main()
    Npoints = 10
    xa = range(1.0, 1.5, length=Npoints)
    ya = my_sin.(xa)

    m = 0
    for i in 1:Npoints
        x = xa[i]
        y = polynm(m, Npoints, xa, ya, x)
        @printf("%18.10f %18.10f %18.10f\n", x, y, ya[i])
    end

    x_plt = range(1.0, 1.5, length=5*Npoints)
    for x in x_plt
        y = polynm(m, Npoints, xa, ya, x)
        #@printf("%18.10f %18.10f %18.10f\n", x, y, my_sin(x))
        #@printf("%18.10f %18.10f %18.10f\n", x, y, integ_my_sin(x))
    end
end

test_main()


@testset "quadratic problems" begin
    @testset "qp1_test" begin
        #=
            min x^2 + xy + y^2 + yz + z^2
        s.t.
            x + 2y + 3z >= 4 (a)
            x +  y      >= 1 (b)
            x,y \in R

        dual
            max 4a +  b - w1^2 - w1 w2 - w2^2 - w2 w3 - w3^2
        s.t.
            - a -  b + 2w1 + 1w2        = 0
            -2a -  b + 1w1 + 2w2 + 1w3  = 0
            -3a            + 1w2 + 2w3  = 0
            a >= 0
            b >= 0
        =#
        primal_model = qp1_test()
        dual_model, primal_dual_map = dual_model_and_map(primal_model)

        @test MOI.get(dual_model, MOI.NumberOfVariables()) == 2 + 3
        list_of_cons =  MOI.get(dual_model, MOI.ListOfConstraints())
        @test Set(list_of_cons) == Set([
            (SVF, MOI.GreaterThan{Float64})
            (SAF{Float64}, MOI.EqualTo{Float64})
        ])
        @test MOI.get(dual_model, MOI.NumberOfConstraints{SVF, MOI.GreaterThan{Float64}}()) == 2
        @test MOI.get(dual_model, MOI.NumberOfConstraints{SAF{Float64}, MOI.EqualTo{Float64}}()) == 3
        obj_type = MOI.get(dual_model, MOI.ObjectiveFunctionType())
        @test obj_type == SQF{Float64}
        obj = MOI.get(dual_model, MOI.ObjectiveFunction{obj_type}())
        @test MOI.constant(obj) == 0.0
        @test MOI.coefficient.(obj.affine_terms) == [1.0; 4.0]
        @test MOI.coefficient.(obj.quadratic_terms) == [-2.0; -1.0; -2.0; -1.0; -2.0]

        eq_con1_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(1))
        eq_con1_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(1))
        @test MOI.coefficient.(eq_con1_fun.terms) == [1.0; 1.0; -2.0; -1.0]
        @test MOI.constant.(eq_con1_fun) == 0.0
        @test MOI.constant(eq_con1_set) == 0.0
        eq_con2_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(2))
        eq_con2_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(2))
        @test MOI.coefficient.(eq_con2_fun.terms) == [2.0; 1.0; -1.0; -2.0; -1.0]
        @test MOI.constant.(eq_con2_fun) == 0.0
        @test MOI.constant(eq_con2_set) == 0.0
        eq_con3_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(3))
        eq_con3_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(3))
        @test MOI.coefficient.(eq_con3_fun.terms) == [3.0; -1.0; -2.0]
        @test MOI.constant.(eq_con3_fun) == 0.0
        @test MOI.constant(eq_con3_set) == 0.0

        primal_con_dual_var = primal_dual_map.primal_con_dual_var
        @test primal_con_dual_var[CI{SAF{Float64}, MOI.GreaterThan{Float64}}(1)] == [VI(1)]
        @test primal_con_dual_var[CI{SAF{Float64}, MOI.GreaterThan{Float64}}(2)] == [VI(2)]

        primal_var_dual_con = primal_dual_map.primal_var_dual_con
        @test primal_var_dual_con[VI(1)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(1)
        @test primal_var_dual_con[VI(2)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(2)
        @test primal_var_dual_con[VI(3)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(3)

        primal_var_dual_quad_slack = primal_dual_map.primal_var_dual_quad_slack
        @test primal_var_dual_quad_slack[VI(1)] == VI(2+1)
        @test primal_var_dual_quad_slack[VI(2)] == VI(2+2)
        @test primal_var_dual_quad_slack[VI(3)] == VI(2+3)
    end
    @testset "qp2_test" begin
        #=
            min 2 x^2 + y^2 + xy + x + y + 1
        s.t.
            x + y = 1 (a)
            x    >= 0 (b)
               y >= 0 (c)

        dual
            max a + 1 - 2 w1^2 - w2^2 - w1 w2
        s.t.
            - a - b     + 4w1 + 1w2 + 1 = 0
            - a     - c + 1w1 + 2w2 + 1 = 0
            a \in R
            b >= 0
            c >= 0
        =#
        primal_model = qp2_test()
        dual_model, primal_dual_map = dual_model_and_map(primal_model)

        @test MOI.get(dual_model, MOI.NumberOfVariables()) == 3 + 2
        list_of_cons =  MOI.get(dual_model, MOI.ListOfConstraints())
        @test Set(list_of_cons) == Set([
            (SVF, MOI.GreaterThan{Float64})
            (SAF{Float64}, MOI.EqualTo{Float64})
        ])
        @test MOI.get(dual_model, MOI.NumberOfConstraints{SVF, MOI.GreaterThan{Float64}}()) == 2
        @test MOI.get(dual_model, MOI.NumberOfConstraints{SAF{Float64}, MOI.EqualTo{Float64}}()) == 2
        obj_type = MOI.get(dual_model, MOI.ObjectiveFunctionType())
        @test obj_type == SQF{Float64}
        obj = MOI.get(dual_model, MOI.ObjectiveFunction{obj_type}())
        @test MOI.constant(obj) == 1.0
        @test MOI.coefficient.(obj.affine_terms) == [1.0]
        @test MOI.coefficient.(obj.quadratic_terms) == [-4.0; -1.0; -2.0]

        eq_con1_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(1))
        eq_con1_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(1))
        @test MOI.coefficient.(eq_con1_fun.terms) == [1.0; 1.0; -4.0; -1.0]
        @test MOI.constant.(eq_con1_fun) == 0.0
        @test MOI.constant(eq_con1_set) == 1.0
        eq_con2_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(2))
        eq_con2_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(2))
        @test MOI.coefficient.(eq_con2_fun.terms) == [1.0; 1.0; -1.0; -2.0]
        @test MOI.constant.(eq_con2_fun) == 0.0
        @test MOI.constant(eq_con2_set) == 1.0

        primal_con_dual_var = primal_dual_map.primal_con_dual_var
        @test primal_con_dual_var[CI{SAF{Float64}, MOI.EqualTo{Float64}}(1)] == [VI(1)]
        @test primal_con_dual_var[CI{SVF, MOI.GreaterThan{Float64}}(1)] == [VI(2)]
        @test primal_con_dual_var[CI{SVF, MOI.GreaterThan{Float64}}(2)] == [VI(3)]

        primal_var_dual_con = primal_dual_map.primal_var_dual_con
        @test primal_var_dual_con[VI(1)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(1)
        @test primal_var_dual_con[VI(2)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(2)

        primal_var_dual_quad_slack = primal_dual_map.primal_var_dual_quad_slack
        @test primal_var_dual_quad_slack[VI(1)] == VI(3+1)
        @test primal_var_dual_quad_slack[VI(2)] == VI(3+2)
    end
end


using Surrogates, XGBoost

#1D
obj_1D = x -> 3*x+1
x = [1.0,2.0,3.0,4.0,5.0]
y = obj_1D.(x)
a = 0.0
b = 10.0
num_round = 2
my_forest_1D = RandomForestSurrogate(x,y,a,b,num_round)
val = my_forest_1D(3.5)
add_point!(my_forest_1D,6.0,19.0)
add_point!(my_forest_1D,[7.0,8.0],obj_1D.([7.0,8.0]))

#ND
lb = [0.0,0.0,0.0]
ub = [10.0,10.0,10.0]
x = sample(5,lb,ub,SobolSample())
obj_ND = x -> x[1] * x[2]^2 * x[3]
y = obj_ND.(x)
my_forest_ND = RandomForestSurrogate(x,y,lb,ub,num_round)
val = my_forest_ND((1.0,1.0,1.0))
add_point!(my_forest_ND,(1.0,1.0,1.0),1.0)
add_point!(my_forest_ND,[(1.2,1.2,1.0),(1.5,1.5,1.0)],[1.728,3.375])


using Statistics
using BlackBoxOptim
using NLopt
using Distributions

## Optimization
# Generate random index
indfast=rand(1:Param.nsim, Param.nfast)
indfast[1]=1 # Keeping the first element fixed.

# schain=SharedArray{Float64,2}(mchain[indfast,:])
mchain_op=mchain[indfast,:]

## Objective function

function objfun(gamma0...; Par::Par=Param, mchain=mchain_op)
    @unpack nsim, nfast = Par
    gamma = gamma0
    N = size(mchain[1,:])
    val=0
    logunif=log.(rand(nfast))
    geta=mchain[1,:]
    hM=zeros(N)
    for r=1:nfast
        gtry=mchain[r,:]
        for j=1:size(gtry,2)
            val=+gtry[:,j].*gamma[j]-geta[:,j].*gamma[j]
        end
        geta=logunif[r] < val[1] ? gtry : geta
        hM=+geta./nfast
    end
    mhM=mean(hM)
    var=Statistics.var(hM)

    obj=(mhM^2)/var

    return obj
end

function objfun_bbo(gamma0; Par::Par=Param, mchain=mchain_op)
    @unpack nsim, nfast = Par
    gamma = gamma0
    N = size(mchain[1,:])
    val=0
    logunif=log.(rand(nfast))
    geta=mchain[1,:]
    hM=zeros(N)
    for r=1:nfast
        gtry=mchain[r,:]
        for j=1:size(gtry,2)
            val=+gtry[:,j].*gamma[j]-geta[:,j].*gamma[j]
        end
        geta=logunif[r] < val[1] ? gtry : geta
        hM=+geta./nfast
    end
    mhM=mean(hM)
    var=Statistics.var(hM)

    obj=(mhM^2)/var

    return obj
end
## testing obj fun

objfun(collect(1:6))

objfun_bbo(repeat([200],6))
objfun_bbo(collect(1:6))
#
fop=nothing
## Solvin BBO
res = bboptimize(x->objfun_bbo(x); Sea = (-10e300,10e300), NumDimensions = 6)

## Getting results

minbb = best_fitness(res)
γ0 = best_candidate(res)
TS_bb = minbb*6
p_val_bb = ccdf(Chisq(4),TS_bb)
## Optimizing

fop=nothing
fop=Model(NLopt.Optimizer)
set_optimizer_attribute(fop, "algorithm", :LN_BOBYQA)
set_optimizer_attribute(fop, "xtol_rel", 1e-8)
@variable(fop, gamma0[1:6])

register(fop, :objfun, 6, objfun; autodiff=true )
@NLobjective(fop, Min, objfun(gamma0...))

for i=1:6
    set_start_value(gamma0[i],γ0[i])
end

## Solving in NLopt

JuMP.optimize!(fop)

## Getting Results

termination_status(fop)
primal_status(fop)

value.(gamma0)
obj_v=objective_value(fop)
# n=size(q,2)
TS=6*obj_v
p_val = ccdf(Chisq(4),TS)

## Testing function

function RP_Cournot_test(;Par::Par=Param, mchain=mchain_op)
    N = size(mchain,2)
    # First optimization BlackBox
    res = bboptimize(x->objfun_bbo(x); Sea = (-10e300,10e300), NumDimensions = 6)
    ## Getting results
    minbb = best_fitness(res)
    γ0 = best_candidate(res)
    # TS_bb = minbb*N
    # p_val_bb = ccdf(Chisq(N),TS_bb)

    ## Second optimization NLopt
    fop=nothing
    fop=Model(NLopt.Optimizer)
    set_optimizer_attribute(fop, "algorithm", :LN_BOBYQA)
    set_optimizer_attribute(fop, "xtol_rel", 1e-8)
    @variable(fop, gamma0[1:N])

    register(fop, :objfun, N, objfun; autodiff=true )
    @NLobjective(fop, Min, objfun(gamma0...))

    for i=1:N
        set_start_value(gamma0[i],γ0[i])
    end

    ## Solving in NLopt
    JuMP.optimize!(fop)

    ## Getting Results
    termination_status(fop)
    primal_status(fop)

    γ1=value.(gamma0)
    obj_v=objective_value(fop)
    TS=N*obj_v
    p_val = ccdf(Chisq(N),TS)

    ## Second optimization for refinement
    for i=1:N
        set_start_value(gamma0[i],γ1[i])
    end
    ## Solving in NLopt
    JuMP.optimize!(fop)
    obj_v=objective_value(fop)
    TS=N*obj_v
    p_val = ccdf(Chisq(N),TS)

    fop=nothing

    return TS, p_val
end


# # HMM with continuous observations

using Distributions
using HMMBase
using Plots

a = [0.6, 0.4]
A = [0.7 0.3; 0.4 0.6]
B = [MvNormal([0.0,5.0],[1.0,1.0]), MvNormal([5.0,10.0],[1.0,1.0])]
hmm = HMM(a, A, B)
#-

z, y = rand(hmm, 250, seq = true)
z_viterbi = viterbi(hmm, y);
#-

plot(y, linetype=:steppre, label=["Observations (1)", "Observations (2)"], size=(600,200))
#-

plot(z, linetype=:steppre, label="True hidden state", size=(600,200))
#-

plot(z_viterbi, linetype=:steppre, label="Viterbi decoded hidden state", size=(600,200))


const neghalflog2pi = -.5log(2pi)

#######################################
### loglik == Natural Parameter space
#######################################

"""
    loglik( g::GaussNP, x )

Exact log likelihood of a Gaussian expressed in the Natural Parameter space.
By default x is assumed to be of size D x N but if axis=1 then it will consider
it is given as N x D (where N=number of points, D=dimension).
"""
function loglik( g::GaussNP, x::AbstractArray{Float};
                 axis::Int=2 )::Union{Float,Vector{Float}}
    precmu   = g.theta1
    sqrtprec = chol(-g.theta2)
    tmp      = sqrtprec'\precmu
    fun(x)   = ( dot(x,g.theta2*x) + 2dot(x,precmu) )/2
    result   = sum(neghalflog2pi + log.(diag(sqrtprec))) - norm(tmp)^2/2
    result  += funarray(fun, x, length(precmu), axis)
end

"""
    loglik( g::DGaussNP, x )

Exact log likelihood of a Diagonal Gaussian expressed in the Natural Parameter
space. By default x is assumed to be of size D x N but if axis=1 then it will
consider it is given as N x D (where N=number of points, D=dimension).
"""
function loglik( g::DGaussNP, x::AbstractArray{Float};
                 axis::Int=2 )::Union{Float,Vector{Float}}
    precmu   = g.theta1
    sqrtprec = sqrt.(-g.theta2)
    tmp      = precmu./sqrtprec
    fun(x)   = ( dot(x, g.theta2 .* x) + 2dot(x, precmu) )/2
    result   = sum(neghalflog2pi + log.(sqrtprec)) - norm(tmp)^2/2
    result  += funarray(fun, x, length(precmu), axis)
end

#####################################
### loglik == Mean Parameter space
#####################################

"""
    loglik( g::GaussMP, x )

Exact log likelihood of a Gaussian expressed in the Mean Parameter space.
By default x is assumed to be of size D x N but if axis=1 then it will consider
it is given as N x D (where N=number of points, D=dimension).
"""
function loglik( g::GaussMP, x::AbstractArray{Float};
                 axis::Int=2 )::Union{Float,Vector{Float}}
    sqrtcov = chol(cov(g))
    fun(x)  = -norm(sqrtcov'\(x-g.mu1))^2/2
    result  = sum(neghalflog2pi - log.(diag(sqrtcov)))
    result += funarray(fun, x, length(g.mu1), axis)
end

"""
    loglik( g::DGaussMP, x )

Exact log likelihood of a Diagonal Gaussian expressed in the Mean Parameter
space. By default x is assumed to be of size D x N but if axis=1 then it will
consider it is given as N x D (where N=number of points, D=dimension).
"""
function loglik( g::DGaussMP, x::AbstractArray{Float};
                 axis::Int=2 )::Union{Float,Vector{Float}}
    stds    = std(g)
    fun(x)  = -norm((x-g.mu1)./stds)^2/2
    result  = sum(neghalflog2pi - log.(stds))
    result += funarray(fun, x, length(g.mu1), axis)
end

#######################################
### uloglik == Natural Parameter space
#######################################

"""
    uloglik( g::GaussNP, x )

Loglikelihood of a GaussNP without the constants (numerically more stable as
it does not require computing determinants and inversions). This is useful
in a sampler where the normalising constant can be neglected.
It corresponds to -0.5(<x, Px> -2<x,Pmu>).
By default x is assumed to be of size D x N but if axis=1 then it will consider
it is given as N x D (where N=number of points, D=dimension).
"""
function uloglik( g::GaussNP, x::AbstractArray{Float};
                  axis::Int=2 )::Union{Float,Vector{Float}}
    fun(x) = dot(x, g.theta1+g.theta2*x/2)
    funarray(fun, x, length(g), axis)
end

"""
    uloglik( g::DGaussNP, x )

Loglikelihood of a DGaussNP without the constants (numerically more stable as
it does not require computing determinants and inversions). This is useful
in a sampler where the normalising constant can be neglected.
It corresponds to -0.5(<x, Px> -2<x,Pmu>).
By default x is assumed to be of size D x N but if axis=1 then it will consider
it is given as N x D (where N=number of points, D=dimension).
"""
function uloglik( g::DGaussNP, x::AbstractArray{Float};
                  axis::Int=2 )::Union{Float,Vector{Float}}
    fun(x) = dot(x, g.theta1+g.theta2.*x/2)
    funarray(fun, x, length(g), axis)
end

#####################################
### uloglik == Mean Parameter space
#####################################

"""
    uloglik( g::GaussMP, x )

Loglikelihood of a GaussMP without the constants (numerically more stable as
it does not require computing determinants and inversions). This is useful
in a sampler where the normalising constant can be neglected.
It corresponds to -0.5(<x, Px> -2<x,Pmu>).
"""
function uloglik( g::GaussMP, x::AbstractArray{Float};
                  axis::Int=2 )::Union{Float,Vector{Float}}
    S = cov(g)
    fun(x) = dot(x, S\(g.mu1-x/2))
    funarray(fun, x, length(g), axis)
end

"""
    uloglik( g::DGaussMP, x )

Loglikelihood of a DGaussMP without the constants (numerically more stable as
it does not require computing determinants and inversions). This is useful
in a sampler where the normalising constant can be neglected.
It corresponds to -0.5(<x, Px> -2<x,Pmu>).
"""
function uloglik( g::DGaussMP, x::AbstractArray{Float};
                  axis::Int=2 )::Union{Float,Vector{Float}}
    S = cov(g)
    fun(x) = dot(x, (g.mu1-x/2)./S)
    funarray(fun, x, length(g), axis)
end

################
### gradloglik
################

# NOTE at the moment only handles point inputs

gradloglik(g::GaussNP,  x::Vector{Float}) = g.theta1 + g.theta2*x
gradloglik(g::DGaussNP, x::Vector{Float}) = g.theta1 + g.theta2 .* x
gradloglik(g::GaussMP,  x::Vector{Float}) = cov(g)\(g.mu1 - x)
gradloglik(g::DGaussMP, x::Vector{Float}) = (g.mu1 - x)./cov(g)


@testset "LP Solver" begin
    @testset "Issue 83 TimeLimit" begin
        glpk_optimizer =
            optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => GLPK.GLP_MSG_OFF)
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => glpk_optimizer,
            "logging" => [],
            "time_limit" => 0.01,
        ))

        # Variables
        @variable(model, inclusion[h = 1:3], Bin)
        @variable(model, 0 <= allocations[h = 1:3, a = 1:3] <= 1, Int)
        @variable(model, 0 <= days[h = 1:3, a = 1:3] <= 5, Int)

        # Constraints
        @constraint(
            model,
            must_include[h = 1:3],
            sum(allocations[h, a] for a in 1:3) <= inclusion[h]
        )
        # at least n
        @constraint(model, min_hospitals, sum(inclusion[h] for h in 1:3) >= 3)
        # every h must be allocated at most one a
        @constraint(model, must_visit[h = 1:3], sum(allocations[h, a] for a in 1:3) <= 1)
        # every allocated h must have fewer than 5 days of visits per week
        @constraint(
            model,
            max_visits[h = 1:3],
            sum(days[h, a] for a in 1:3) <= 5 * inclusion[h]
        )

        @objective(model, Max, sum(days[h, a] * 5 for h in 1:3, a in 1:3))
        optimize!(model)
        @test JuMP.termination_status(model) == MOI.TIME_LIMIT ||
              JuMP.termination_status(model) == MOI.OPTIMAL
        if JuMP.termination_status(model) == MOI.OPTIMAL
            @test JuMP.solve_time(model) <= 0.1
        end
    end

    @testset "Issue 83 max_bt_steps" begin
        glpk_optimizer =
            optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => GLPK.GLP_MSG_OFF)
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => glpk_optimizer,
            "logging" => [],
            "max_bt_steps" => 2,
        ))

        # Variables
        @variable(model, inclusion[h = 1:3], Bin)
        @variable(model, 0 <= allocations[h = 1:3, a = 1:3] <= 1, Int)
        @variable(model, 0 <= days[h = 1:3, a = 1:3] <= 5, Int)

        # Constraints
        @constraint(
            model,
            must_include[h = 1:3],
            sum(allocations[h, a] for a in 1:3) <= inclusion[h]
        )
        # at least n
        @constraint(model, min_hospitals, sum(inclusion[h] for h in 1:3) >= 3)
        # every h must be allocated at most one a
        @constraint(model, must_visit[h = 1:3], sum(allocations[h, a] for a in 1:3) <= 1)
        # every allocated h must have fewer than 5 days of visits per week
        @constraint(
            model,
            max_visits[h = 1:3],
            sum(days[h, a] for a in 1:3) <= 5 * inclusion[h]
        )

        @objective(model, Max, sum(days[h, a] * 5 for h in 1:3, a in 1:3))
        optimize!(model)
        if JuMP.termination_status(model) != MOI.OTHER_LIMIT
            @test JuMP.objective_value(model) ≈ 75
        end
    end

    @testset "Issue 83" begin
        glpk_optimizer =
            optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => GLPK.GLP_MSG_OFF)
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => glpk_optimizer,
            "logging" => [],
        ))

        # Variables
        @variable(model, inclusion[h = 1:3], Bin)
        @variable(model, 0 <= allocations[h = 1:3, a = 1:3] <= 1, Int)
        @variable(model, 0 <= days[h = 1:3, a = 1:3] <= 5, Int)

        # Constraints
        @constraint(
            model,
            must_include[h = 1:3],
            sum(allocations[h, a] for a in 1:3) <= inclusion[h]
        )
        # at least n
        @constraint(model, min_hospitals, sum(inclusion[h] for h in 1:3) >= 3)
        # every h must be allocated at most one a
        @constraint(model, must_visit[h = 1:3], sum(allocations[h, a] for a in 1:3) <= 1)
        # every allocated h must have fewer than 5 days of visits per week
        @constraint(
            model,
            max_visits[h = 1:3],
            sum(days[h, a] for a in 1:3) <= 5 * inclusion[h]
        )

        @objective(model, Max, sum(days[h, a] * 5 for h in 1:3, a in 1:3))
        optimize!(model)
        @test JuMP.objective_value(model) ≈ 75

        com = CS.get_inner_model(model)
    end

    @testset "Combine lp with table constraint" begin
        cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => cbc_optimizer,
            "logging" => [],
        ))

        # Variables
        @variable(model, 1 <= x[1:10] <= 15, Int)

        # Constraints
        @constraint(model, sum(x[1:5]) >= 10)
        @constraint(model, sum(x[6:10]) < 16)

        table_left = [
            1 2 3 4 5 # sum 16
            1 2 3 4 6 #     17
            4 5 5 3 4 #     21
        ]

        table_right = [
            1 2 3 1 1 # sum 8
            1 3 2 2 1 #     9
            1 1 1 1 4 #     8
        ]

        @constraint(model, x[1:5] in CS.TableSet(table_left))
        @constraint(model, x[6:10] in CS.TableSet(table_right))

        @objective(model, Max, sum(x))
        optimize!(model)

        @test JuMP.objective_value(model) ≈ 30
        @test sum(JuMP.value.(x[6:10])) ≈ 9
        @test sum(JuMP.value.(x[1:5])) ≈ 21
    end

    @testset "LP with bigger table constraint" begin
        cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => cbc_optimizer,
            "logging" => [],
            "keep_logs" => true,
        ))

        # Variables
        @variable(model, 1 <= x[1:10] <= 15, Int)

        # Constraints
        @constraint(model, sum(x[1:4]) > 9)
        @constraint(model, sum(x[5:10]) <= 15)

        table_left = Array{Int,2}(undef, (0, 5))
        for r in permutations(1:10, 5)
            s = sum(r)
            if 10 < s < 20 || 28 < s < 47
                table_left = vcat(table_left, r')
            end
        end

        table_right = [
            1 2 3 1 1 # sum 8
            1 3 3 2 1 #    10
            1 1 3 2 1 #     8
            1 1 1 2 4 #     9
            4 5 5 3 4 #     21
        ]

        @constraint(model, x[1:5] in CS.TableSet(table_left))
        @constraint(model, x[6:10] in CS.TableSet(table_right))

        @objective(model, Max, sum(x))
        optimize!(model)

        @test JuMP.termination_status(model) == MOI.OPTIMAL
        @test JuMP.objective_value(model) ≈ 49
        @test sum(JuMP.value.(x[1:4])) >= 10
        @test sum(JuMP.value.(x[5:10])) <= 15
        @test allunique(JuMP.value.(x[1:5]))
    end

    @testset "Combine lp with all different + not equal" begin
        cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => cbc_optimizer,
            "logging" => [],
            "traverse_strategy" => :DBFS,
        ))

        # Variables
        @variable(model, 1 <= x[1:10] <= 15, Int)

        # Constraints
        @constraint(model, sum(x[1:5]) >= 10)
        @constraint(model, sum(x[6:10]) <= 15)
        # disallow x[1] = 11 and x[2] == 12
        @constraint(model, 1.5 * x[1] + 2 * x[2] != 40.5)
        @constraint(model, x in CS.AllDifferentSet())

        @objective(model, Max, sum(x))
        optimize!(model)
        # possible solution 11+12+13+14+15  + 1+2+3+4+5
        # only works fast if the all different bound works
        @test JuMP.objective_value(model) ≈ 80
        @test sum(JuMP.value.(x[6:10])) ≈ 15
        @test !(1.5 * JuMP.value(x[1]) + 2 * JuMP.value(x[2]) ≈ 40.5)
    end

    @testset "Combine lp with all different + not equal + DFS" begin
        cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => cbc_optimizer,
            "logging" => [],
            "traverse_strategy" => :DFS,
            "branch_split" => :Biggest,
            "keep_logs" => true,
        ))

        # Variables
        @variable(model, 1 <= x[1:10] <= 15, Int)

        # Constraints
        @constraint(model, sum(x[1:5]) >= 10)
        @constraint(model, sum(x[6:10]) <= 15)
        # disallow x[1] = 11 and x[2] == 12
        @constraint(model, 1.5 * x[1] + 2 * x[2] != 40.5)
        @constraint(model, x in CS.AllDifferentSet())

        @objective(model, Max, sum(x))
        optimize!(model)
        # possible solution 11+12+13+14+15  + 1+2+3+4+5
        # only works fast if the all different bound works
        @test JuMP.objective_value(model) ≈ 80
        @test sum(JuMP.value.(x[6:10])) ≈ 15
        @test !(1.5 * JuMP.value(x[1]) + 2 * JuMP.value(x[2]) ≈ 40.5)

        com = CS.get_inner_model(model)
        @test is_solved(com)
        @test general_tree_test(com)
    end

    @testset "Combine lp with all different + not equal + DFS + Split Half" begin
        cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
        model = Model(optimizer_with_attributes(
            CS.Optimizer,
            "lp_optimizer" => cbc_optimizer,
            # "logging" => [],
            "traverse_strategy" => :DFS,
            "branch_split" => :InHalf,
        ))

        # Variables
        @variable(model, 1 <= x[1:10] <= 15, Int)

        # Constraints
        @constraint(model, sum(x[1:5]) >= 10)
        @constraint(model, sum(x[6:10]) <= 15)
        # disallow x[1] = 11 and x[2] == 12
        @constraint(model, 1.5 * x[1] + 2 * x[2] != 40.5)
        @constraint(model, x in CS.AllDifferentSet())

        @objective(model, Max, sum(x))
        optimize!(model)
        # possible solution 11+12+13+14+15  + 1+2+3+4+5
        # only works fast if the all different bound works
        @test JuMP.objective_value(model) ≈ 80
        @test sum(JuMP.value.(x[6:10])) ≈ 15
        @test !(1.5 * JuMP.value(x[1]) + 2 * JuMP.value(x[2]) ≈ 40.5)
    end
end


include("dependencies_for_runtests.jl")

using Oceananigans.Solvers: solve!
using Statistics

function identity_operator!(b, x)
    parent(b) .= parent(x)
    return nothing
end

function run_identity_operator_test(grid)
    b = CenterField(grid)
    solver = PreconditionedConjugateGradientSolver(identity_operator!, template_field = b, reltol=0, abstol=10*sqrt(eps(eltype(grid))))
    initial_guess = solution = similar(b)
    set!(initial_guess, (x, y, z) -> rand())

    solve!(initial_guess, solver, b)

    @test norm(solution) .< solver.abstol
end

function run_poisson_equation_test(grid)
    arch = architecture(grid)
    # Solve ∇²ϕ = r
    ϕ_truth = CenterField(grid)

    # Initialize zero-mean "truth" solution with random numbers
    set!(ϕ_truth, (x, y, z) -> rand())
    parent(ϕ_truth) .-= mean(ϕ_truth)
    fill_halo_regions!(ϕ_truth)

    # Calculate Laplacian of "truth"
    ∇²ϕ = r = CenterField(grid)
    compute_∇²!(∇²ϕ, ϕ_truth, arch, grid)

    solver = PreconditionedConjugateGradientSolver(compute_∇²!, template_field=ϕ_truth, reltol=eps(eltype(grid)))

    # Solve Poisson equation
    ϕ_solution = CenterField(grid)
    solve!(ϕ_solution, solver, r, arch, grid)
    fill_halo_regions!(ϕ_solution)

    # Diagnose Laplacian of solution
    ∇²ϕ_solution = CenterField(grid)
    compute_∇²!(∇²ϕ_solution, ϕ_solution, arch, grid)

    # Test
    extrema_tolerance = 1e-12
    std_tolerance = 1e-13

    CUDA.@allowscalar begin
        @test minimum(abs, interior(∇²ϕ_solution) .- interior(∇²ϕ)) < extrema_tolerance
        @test maximum(abs, interior(∇²ϕ_solution) .- interior(∇²ϕ)) < extrema_tolerance
        @test          std(interior(∇²ϕ_solution) .- interior(∇²ϕ)) < std_tolerance

        @test   minimum(abs, interior(ϕ_solution) .- interior(ϕ_truth)) < extrema_tolerance
        @test   maximum(abs, interior(ϕ_solution) .- interior(ϕ_truth)) < extrema_tolerance
        @test            std(interior(ϕ_solution) .- interior(ϕ_truth)) < std_tolerance
    end

    return nothing
end

@testset "PreconditionedConjugateGradientSolver" begin
    for arch in archs
        @info "Testing PreconditionedConjugateGradientSolver [$(typeof(arch))]..."
        grid = RectilinearGrid(arch, size=(4, 8, 4), extent=(1, 3, 1))
        run_identity_operator_test(grid)
        run_poisson_equation_test(grid)
    end
end


function flatten(weights::Params)::Array{Float32, 1}
    cat([ndims(w) > 1 ? reshape(w, length(w)) : w for w in weights]..., dims=1)
end

function unflatten_set(directions::Array{Float32, 1}, weights::Params)::Params
    unflattened = deepcopy(weights)
    idx = 1
    for uf in unflattened
        uf[:] = reshape(directions[idx:idx + length(uf) - 1], size(uf))
        idx += length(uf)
    end
    @assert idx - 1 == length(directions) "Failed to correctly unflatten, please report this issue."
    unflattened
end

function pca_directions(
    weights::Params, checkpoints::Vector{String},
)::Array{Float32, 2}
    directions_matrix = []
    for checkpoint in checkpoints
        @load checkpoint checkpoint_weights
        directions = create_directions(weights, params(checkpoint_weights)) |> flatten
        push!(directions_matrix, directions)
    end
    directions_matrix = cat(directions_matrix..., dims=2)
    model = fit(PCA, directions_matrix, maxoutdim=2)
    projection(model)
end

function project(
    directions::Array{Float32, 1}, pca_projection::Array{Float32, 2},
    method::Symbol = :cos,
)::Array{Float32, 1}
    if method == :lstsq
        return pca_projection \ directions
    end
    normalizer = [norm(pca_projection[:, i]) for i in 1:size(pca_projection, 2)]
    (pca_projection' * directions) ./ normalizer
end

function trajectory(
    weights::Params, pca_projection::Array{Float32, 2},
    checkpoints::Vector{String}, projection::Symbol = :cos,
)::Array{Float32, 2}
    coordinates = zeros(Float32, length(checkpoints), 2)
    @inbounds for (i, checkpoint) in enumerate(checkpoints)
        @load checkpoint checkpoint_weights
        directions = create_directions(weights, params(checkpoint_weights)) |> flatten
        projection = project(directions, pca_projection, :cos)
        coordinates[i, :] = projection
    end
    coordinates
end

"""
```julia
function create_trajectory(
    weights::Params, checkpoints::Vector{String}, projection::Symbol = :cos,
)::Array{Float32, 2}
```

Create trajectory between model checkpoints.
This allows for path visualization that optimizer took
during training, for example.

# Arguments
- `weights::Params`: Target weights of the model. E.g. from last epoch.
- `checkpoints::Vector{String}`: Checkpoints of the model from previous epochs.
They will be loaded in sorted order and should contain weights
under `checkpoint_weights` key.
- `projection::Symbol`: Projection method. Either `:lstsq` or `:cos`.

!!! note
    Checkpoint files should contain weights under `checkpoint_weights` key.
"""
function create_trajectory(
    weights::Params, checkpoints::Vector{String}, projection::Symbol = :cos,
)::Array{Float32, 2}
    directions = pca_directions(weights, checkpoints)
    trajectory(weights, directions, checkpoints, projection)
end


using Pkg,LinearAlgebra, SparseArrays, Random,Distributions;

y=[1,2,3]
function get_Y(y)
    m=length(y)#number of rows
    n=m*(m-1) #number of columns
    res = zeros(m,n) #set the result
    start = 1
    for i=1:length(y)
        row = copy(y)
        row = deleteat!(row, i)
        res[i,start:start+m-2]=row
        start = start+m-1
    end
    return res
end
get_Y(y)

#fake data
y1=[1,2,3]
y2=[4,5,6]
Y1=get_Y(y1)
Y2=get_Y(y2)
Y=[Y1,Y2]
R=[1 0 0;0 2 0; 0 0 3]
w1=[6,7,8]
w2=[6,3,4]
W=[w1,w2]
t_square = 3
lambda0 = 4
get_Y(y1)

function get_para(Y,R,t_square,lambda0)
    m = size(Y[1])[1]#number of rows
    n=m*(m-1) #number of columns
    first = zeros(n,n)
    second=zeros(n,1)
    for i=1:length(Y)
        first += Y[i]'*inv(R)*Y[i]
        second += Y[i]'*inv(R)*W[i]
    end
    first += Diagonal(repeat([t_square],n))
    second += repeat([lambda0*t_square],n)
    mu = vec(inv(first)*second)
    return mu,Symmetric(inv(first))
end

#MCMC
mu,var = get_para(Y,R,t_square,lambda0)
rand(MvNormal(mu,Symmetric(var)))

y = [1,2,3]
function get_my_Y(y)
    m=length(y)#number of rows
    n::Int64=(m*(m-1))/2 #number of columns
    res = zeros(m,n) #set the result
    start = 1
    for i = 1:length(y)-1
        row = y[2:i+1]
        res[i+1,start:start + i - 1] = row
        start = start + i
    end
    return res
end
get_my_Y(y)

#Making fake data
y1=[1,2,3]
y2=[4,5,6]
Y1=get_my_Y(y1)
Y2=get_my_Y(y2)
Y=[Y1,Y2]
R=[1 0 0;0 2 0; 0 0 3]
w1=[6,7,8]
w2=[6,3,4]
W=[w1,w2]
t_square = 3
lambda0 = 4

function get_para(Y,R,t_square,lambda0,W)
    m = size(Y[1])[1]#number of rows
    n::Int64=(m*(m-1))/2 #number of columns
    first = zeros(n,n)
    second=zeros(n,1)
    for i=1:length(Y)
        first += Y[i]'*inv(R)*Y[i]
        second += Y[i]'*inv(R)*W[i]
    end
    first += Diagonal(repeat([t_square],n))
    second += repeat([lambda0*t_square],n)
    mu = vec(inv(first)*second)
    return mu,Symmetric(inv(first))
end

@time mu,var = get_para(Y,R,t_square,lambda0,W)
#rand(MvNormal(mu,Symmetric(var)))

mu

var

res = randn(3)
res = L.U*res + mu
function sampling(mu,var,n)
    res = randn(n)
    L = cholesky(var)
    res = L.L*res + mu
    return res
end
sampling(mu,var,3)

#Making fake data
y1=[1,2,3]
y2=[4,5,6]
Y1=get_my_Y(y1)
Y2=get_my_Y(y2)
Y=[Y1,Y2]
R=[1 0 0;0 2 0; 0 0 3]
w1=[6 7 8]
w2=[6 3 4]
W=[w1 w2]
t_square = 3
lambda0 = 4
number=2

#Y is the data for all individuals

#unlist function will tranfrom y from a list to a big matrix
function unlist(Y)
    res = Y[1]
    for i = 2:length(Y)
        res = vcat(res,Y[i])
    end
    return res
end

unlist(Y)

function get_para(Y,R,W,t_square,lambda0,number)
    m = size(Y[1])[1]#number of rows
    n::Int64=(m*(m-1))/2 #number of columns
    Y = unlist(Y) #Set big Y matrix

    #Define sparse matrix
    I = collect(1:number*m); J=collect(1:number*m);V = repeat(diag(R),number)
    R = sparse(I,J,V)

    #formula calculation
    first = Y'*inv(Matrix(R))*Y + Diagonal(ones(m,m))
    second = Y'*inv(Matrix(R))*transpose(vcat(W)) + ones(m,1)*lambda0*t_square
    mu = vec(inv(first)*second)
    return mu,Symmetric(inv(first))
end

a, b = get_para(Y,R,W,t_square,lambda0,number)

#Y is a data for single persion
y1 = [1,2,3]
function get_sparse_Y(Y)
    m = size(Y)[1] #number of rows
    n::Int64=(m*(m-1))/2  #number of columns
    col_index = collect(1:n)
    row_index = [];value = []
    for i = 2:m
        current = fill(i,i-1) #get a element i with i-1 time
        row_index = vcat(row_index,current)
        current_value = Y[2:i]
        value = vcat(value,current_value)
    end

    #Return the sparse matrix
    row_index = convert(Array{Int64,1}, row_index) #change Any type
    value = convert(Array{Float64,1}, value) #Change Any type

    res = sparse(row_index,col_index,value)
end

get_sparse_Y(y1)

#Fake data
y1=[1,2,3]
y2=[4,5,6]
ys1=get_sparse_Y(y1)
ys2=get_sparse_Y(y2)
YS = [ys1,ys2]
R=[1 0 0;0 2 0; 0 0 3]
w1=[6,7,8]
w2=[6,3,4]
W=[w1,w2]
t_square = 3
lambda0 = 4
m = 3
number = 10000

YS = repeat([get_sparse_Y(rand(3))],10000)
YS = unlist(YS)
#W=unlist(W)
W = rand(30000)

function get_sparse_para(YS,R,W,t_square,lambda0,number,m)

    #Define residual matrix
    value = diag(R).^(-1)
    RS = kron(sparse(I,number,number),sparse(1:m,1:m,value))

    #formula calculation
    YS_RS_product = YS'*RS
    first = YS_RS_product*YS + t_square*sparse(I,m,m)
    second = YS_RS_product*W + ones(m,1)*lambda0*t_square

    #Return the result
    first = inv(Matrix(first))
    mu = vec(first*second)
    var = Symmetric(first)

    #Take sample
    res = rand(MvNormal(mu,var))
    return res
end

#a,b = get_sparse_para(YS,R,W,t_square,lambda0,number,3)#
@time lambda= get_sparse_para(YS,R,W,t_square,lambda0,number,m)

y1 = rand(1000);y2 = rand(1000);y3=rand(1000)
BigY = [y1;y2;y3]
Lambda = [1 0 0; lambda[1] 1 0;lambda[2] lambda[3] 1]

function Modify_Y(Lambda,Y,n_individuals)
    Identy = Matrix{Int32}(I,n_individuals,n_individuals)
    res = kron(Identy,Lambda)*BigY
    return res
end

Modify_Y(Lambda,BigY,1000)


module SecondOrderSeawaterPolynomials

export
    SecondOrderSeawaterPolynomial,
    RoquetSeawaterPolynomial,
    RoquetEquationOfState

using SeawaterPolynomials: AbstractSeawaterPolynomial, BoussinesqEquationOfState

import SeawaterPolynomials: ρ′, thermal_expansion, haline_contraction

"""
    struct SecondOrderSeawaterPolynomial{FT} <: AbstractSeawaterPolynomial

Container of coefficients for a second-order polynomial function of
absolute salinity `Sᴬ`, conservative temperature `Θ`, and geopotential
depth `Z` for seawater density.

The coefficients have the form

```math
R_ᵦᵪᵩ ,
```

where ``β, χ, φ`` denote the order of the term to which the coefficent corresponds:
``β`` is the polynomial order of absolute salinity, ``Sᴬ``, ``χ`` is the polynomial order
of conservative temperature, ``Θ``, and ``φ`` is the order of geopotential height, ``Z``.

For a `SecondOrderSeawaterPolynomial`, ``β + χ + φ < 3``.

The coefficient ``R₁₁₀`` arises in the seawater polynomial as

```
seawater_polynomial(Θ, Sᴬ, Z) = ⋯ + R₁₁₀ * Sᴬ * Θ + ⋯
```
"""
@Base.kwdef struct SecondOrderSeawaterPolynomial{FT} <: AbstractSeawaterPolynomial
    R₁₀₀ :: FT = 0
    R₀₁₀ :: FT = 0
    R₁₀₁ :: FT = 0
    R₀₁₁ :: FT = 0
    R₁₁₀ :: FT = 0
    R₀₂₀ :: FT = 0
    R₂₀₀ :: FT = 0
end

const EOS₂ = BoussinesqEquationOfState{<:SecondOrderSeawaterPolynomial}

@inline ρ′(Θ, Sᴬ, Z, eos::EOS₂) = (  eos.seawater_polynomial.R₁₀₀ * Sᴬ
                                   + eos.seawater_polynomial.R₀₁₀ * Θ
                                   + eos.seawater_polynomial.R₀₂₀ * Θ^2
                                   - eos.seawater_polynomial.R₀₁₁ * Θ * Z
                                   + eos.seawater_polynomial.R₂₀₀ * Sᴬ^2
                                   - eos.seawater_polynomial.R₁₀₁ * Sᴬ * Z
                                   + eos.seawater_polynomial.R₁₁₀ * Sᴬ * Θ )

@inline thermal_expansion(Θ, Sᴬ, Z, eos::EOS₂) = (      eos.seawater_polynomial.R₀₁₀
                                                  + 2 * eos.seawater_polynomial.R₀₂₀ * Θ
                                                  -     eos.seawater_polynomial.R₀₁₁ * Z
                                                  +     eos.seawater_polynomial.R₁₁₀ * Sᴬ )

@inline haline_contraction(Θ, Sᴬ, Z, eos::EOS₂) = (      eos.seawater_polynomial.R₁₀₀
                                                   + 2 * eos.seawater_polynomial.R₂₀₀ * Sᴬ
                                                   -     eos.seawater_polynomial.R₁₀₁ * Z
                                                   +     eos.seawater_polynomial.R₁₁₀ * Θ )

"""
    RoquetSeawaterPolynomial([FT=Float64,] coefficient_set=:SecondOrder)

Returns a `SecondOrderSeawaterPolynomial` with coefficients optimized by

> Roquet et al., "Defining a Simplified yet 'Realistic' Equation of State for Seawater", Journal of Physical Oceanography (2015).

The `coefficient_set` is a symbol or string that selects one of the "sets" of
optimized second order coefficients.

Coefficient sets
================

    - :Linear: a linear equation of state, ρ = ρᵣ + R₁₀₀ * Θ + R₀₁₀ * Sᴬ

    - :Cabbeling: includes quadratic temperature term,
                  ρ = ρᵣ + R₁₀₀ * Θ + R₀₁₀ * Sᴬ + R₀₂₀ * Θ^2

    - :CabbelingThermobaricity: includes 'thermobaricity' term,
                                ρ = ρᵣ + R₁₀₀ * Θ + R₀₁₀ * Sᴬ + R₀₂₀ * Θ^2 - R₀₁₁ * Θ * Z

    - :Freezing: same as :cabbeling_thermobaricity with modified constants to increase
                   accuracy near freezing

    - :SecondOrder: includes quadratic salinity, halibaricity, and thermohaline term,
                    ρ = ρᵣ + R₁₀₀ * Θ + R₀₁₀ * Sᴬ + R₀₂₀ * Θ^2 - R₀₁₁ * T * Z
                           + R₂₀₀ * Sᴬ^2 - R₁₀₁ * Sᴬ * Z + R₁₁₀ * Sᴬ * Θ

The optimized coefficients are reported in Table 3 of Roquet et al., "Defining a Simplified
yet 'Realistic' Equation of State for Seawater", Journal of Physical Oceanography (2015), and
further discussed around equations (12)--(15). The optimization minimizes errors in estimated
horizontal density gradient estiamted from climatological temperature and salinity distributions
between the 5 simplified forms chosen by Roquet et. al and the full-fledged
[TEOS-10](http://www.teos-10.org) equation of state.
"""
RoquetSeawaterPolynomial(FT::DataType, coefficient_set=:SecondOrder) =
    eval(Symbol(coefficient_set, :RoquetSeawaterPolynomial))(FT)

# For Float64 default
RoquetSeawaterPolynomial(coefficient_set=:SecondOrder) =
    eval(Symbol(coefficient_set, :RoquetSeawaterPolynomial))(Float64)

"""
    RoquetEquationOfState([FT=Float64,] coefficient_set=:SecondOrder; reference_density=1024.6)

Returns an `BoussinesqEquationOfState` with a `RoquetSeawaterPolynomial` corresponding to
`coefficient_set` and with `reference density = 1024.6 kg m⁻³`, the average density of seawater
at the surface of the world ocean.

Type

```julia
help?> RoquetSeawaterPolynomial
```

for options for the `coefficient_set`. The optimzed coefficient sets for the
`RoquetSeawaterPolynomial` are listed in Table 3 in

> Roquet et al., "Defining a Simplified yet 'Realistic' Equation of State for Seawater", Journal of Physical Oceanography (2015).
"""
RoquetEquationOfState(FT::DataType, coefficient_set=:SecondOrder; reference_density=1024.6) =
    BoussinesqEquationOfState(RoquetSeawaterPolynomial(FT, coefficient_set), reference_density)

RoquetEquationOfState(coefficient_set=:SecondOrder; reference_density=1024.6) =
    RoquetEquationOfState(Float64, coefficient_set, reference_density=1024.6)

"""
    LinearRoquetSeawaterPolynomial(FT=Float64)

Parameters for a linear equation of state optimized for the 'current' oceanic
temperature and salinity distribution.

For more information, type `help?> RoquetSeawaterPolynomial`.
"""
LinearRoquetSeawaterPolynomial(FT=Float64) =
    SecondOrderSeawaterPolynomial{FT}(
                                      R₀₁₀ = - 1.775e-1,
                                      R₁₀₀ =   7.718e-1,
                                     )

"""
    CabbelingRoquetSeawaterPolynomial(FT=Float64)

Parameters for a minimal equation of state that describes cabbeling,
optimized for the 'current' oceanic temperature and salinity distribution.

For more information, type `help?> RoquetSeawaterPolynomial`.
"""
CabbelingRoquetSeawaterPolynomial(FT=Float64) =
    SecondOrderSeawaterPolynomial{FT}(
                                      R₀₁₀ = - 0.844e-1,
                                      R₁₀₀ =   7.718e-1,
                                      R₀₂₀ = - 4.561e-3,
                                     )

"""
    CabbelingThermobaricityRoquetSeawaterPolynomial(FT=Float64)

Parameters for a minimal equation of state that describes cabbeling and thermobaric
effects on sewater density, optimized for the 'current' oceanic temperature and salinity
distribution.

For more information, type `help?> RoquetSeawaterPolynomial`.
"""
CabbelingThermobaricityRoquetSeawaterPolynomial(FT=Float64) =
    SecondOrderSeawaterPolynomial{FT}(
                                      R₀₁₀ = - 0.651e-1,
                                      R₁₀₀ =   7.718e-1,
                                      R₀₂₀ = - 5.027e-3,
                                      R₀₁₁ = - 2.5681e-5,
                                     )

"""
    FreezingRoquetSeawaterPolynomial(FT=Float64)

Parameters for a minimal equation of state that describes seawater density near its
freezing point, optimized for the 'current' oceanic temperature and salinity
distribution.

For more information, type `help?> RoquetSeawaterPolynomial`.
"""
FreezingRoquetSeawaterPolynomial(FT=Float64) =
    SecondOrderSeawaterPolynomial{FT}(
                                      R₀₁₀ = - 0.491e-1,
                                      R₁₀₀ =   7.718e-1,
                                      R₀₂₀ = - 5.027e-3,
                                      R₀₁₁ = - 2.5681e-5,
                                     )

"""
    SecondOrderRoquetSeawaterPolynomial(FT=Float64)

Parameters for a fully second-order equation of state for seawater,
optimized for the 'current' oceanic temperature and salinity  distribution.

For more information, type `help?> RoquetSeawaterPolynomial`.
"""
SecondOrderRoquetSeawaterPolynomial(FT=Float64) =
    SecondOrderSeawaterPolynomial{FT}(
                                      R₀₁₀ =   0.182e-1,
                                      R₁₀₀ =   8.078e-1,
                                      R₀₂₀ = - 4.937e-3,
                                      R₀₁₁ = - 2.4677e-5,
                                      R₂₀₀ = - 1.115e-4,
                                      R₁₀₁ = - 8.241e-6,
                                      R₁₁₀ = - 2.446e-3
                                     )

end # module


# This file contains various methods for solving the
# Kolmogorov-Forward equation in the steady state
# computation for HANK models:
# 0 = A' * g
# They can be generally classified as
# methods for well-conditioned A matrices and ill-conditioned A matrices.
# All of these methods return the stacked distribution.

function solve_kfe(A::SparseMatrixCSC{ComplexF64, Int64}, partitions::Vector{Float64},
                   dimensions::Vector{Int64})

    if size(partitions, 1) != 1 || size(dimensions, 1) != 1
        error("Need to enter partitions and/or dimensions as a 1 x n array or as a vector")
    end

    AT       = A'
    dim_prod = prod(dimensions) # Need total size of dimensions

    # Fixing one value to ensure matrix is not singular, may require further generalization later
    i_fix    = 1
    b        = zeros(dim_prod)
    b[i_fix] = .1
    for j = 1:dim_prod
        AT[i_fix,j] = 0
    end
    AT[i_fix, i_fix] = 1

    # Solve linear system for distribution
    gg    = AT \ b
    g_sum = gg'*ones(dim_prod, 1) * prod(partitions)
    gg    = gg ./ g_sum
    return vec(gg)
end

# Iterative method for ill-conditioned A
# g0: initial distribution guess provided by user
# integrate_mat: weights distribution points by some multiple of their Lebesgue measure across wealth
function solve_kfe(A::SparseMatrixCSC{ComplexF64, Int64}, g0::Matrix{Float64},
                   weight_mat::SparseMatrixCSC{Float64, Int64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    if size(weight_mat) != size(A)
        error("Dimension of weight matrix is incorrect.")
    end
    dim_size = size(A, 1)
    gg = vec(g0) # stack distribution matrix into vector

    # Solve linear system
    for ikfe = 1:maxit_kfe
        gg_tilde = weight_mat * gg # weight distribution points by their measure across wealth

        # 0 = A'g -> g = A'g + g -> (I - A')g = g, but possible singularity so perturb from zero
        # by adding g/Δ_KFE, where Δ_KFE is large => g/Δ_KFE ≈ 0. Thus,
        # g/Δ_KFE = A'*g + g/Δ_KFE => g = Δ_KFE A'*g + g
        # => (I - Δ_KFE A')g = g; could also use (I + Δ_KFE A')g = g
        gg1_tilde = (speye(ComplexF64, dim_size) - Δ_kfe * A') \ gg_tilde
        gg1_tilde = gg1_tilde ./ sum(gg1_tilde) # normalize
        gg1       = weight_mat \ gg1_tilde # undo integration over wealth

        # Check iteration for convergence
        err_kfe = maximum(abs.(gg1 - gg))
        if err_kfe < tol_kfe
            break
        end
        gg = gg1
    end

    return vec(gg)
end

# Previous function but for complex g0
function solve_kfe(A::SparseMatrixCSC{ComplexF64, Int64}, g0::Matrix{ComplexF64},
                   weight_mat::SparseMatrixCSC{Float64, Int64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    if size(weight_mat) != size(A)
        error("Dimension of weight matrix is incorrect.")
    end
    dim_size = size(A, 1)
    gg = vec(g0) # stack distribution matrix into vector

    # Solve linear system
    for ikfe = 1:maxit_kfe
        gg_tilde = weight_mat * gg # weight distribution points by their measure across wealth

        # 0 = A'g -> g = A'g + g -> (I - A')g = g, but to avoid issues with small numbers
        # we use a large scalar Delta_KFE such that
        # 0 = A' Delta_KFE g -> (I - Delta_KFE A') g = g
        gg1_tilde = (speye(ComplexF64, dim_size) - Δ_kfe * A') \ gg_tilde # Left divide
        gg1_tilde = gg1_tilde ./ sum(gg1_tilde) # normalize to one to get pdf
        gg1 = weight_mat \ gg1_tilde # undo integration over wealth

        # Check iteration for convergence
        err_kfe = maximum(abs.(gg1 - gg))
        if err_kfe < tol_kfe
            break
        end

        gg = gg1
    end

    return vec(gg)
end

# Iterative methods when weight_mat is (1) not sparse or
# (2) a dense matrix that integrates out non-wealth variables (rather than a diagonal one),
function solve_kfe(A::SparseMatrixCSC{ComplexF64, Int64}, g0::Vector{Float64},
                   weight_mat::Matrix{Float64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    if size(weight_mat) == size(A)
        return solve_kfe(A, g0, sparse(weight_mat); maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
    else
        weight_mat = spdiagm(vec(weight_mat), 0)
        return solve_kfe(A, g0, weight_mat; maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
    end
end

# Iterative method when weight_mat is a vector of differences of the wealth grid
# e.g. if x = [0 .05, .1, .15], then weight_mat = [.05, .05, .05]
function solve_kfe(A::SparseMatrixCSC{ComplexF64, Int64}, g0::Vector{Float64},
                   grid::Vector{Float64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    reps = size(A, 1) / length(grid)
    if Float64(reps) != Float64(Int64(reps))
        error("Dimension of A does not match dimensions of the grid.")
    end
    diagonal = zeros(size(A, 1))
    for i = 1:reps
        reps[1 + length(grid) * (i - 1): length(grid) * i] = grid
    end
    weight_mat = spdiagm(diagonal, 0)
    return solve_kfe(A, g0, weight_mat; maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
end

# Allows for complex g0
function solve_kfe(A::SparseMatrixCSC{ComplexF64, Int64}, g0::Vector{ComplexF64},
                   weight_mat::Matrix{Float64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    if size(weight_mat) == size(A)
        return solve_kfe(A, g0, sparse(weight_mat); maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
    else
        weight_mat = spdiagm(vec(weight_mat), 0)
        return solve_kfe(A, g0, weight_mat; maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
    end
end

function solve_kfe(A::SparseMatrixCSC{ComplexF64, Int64}, g0::Vector{ComplexF64},
                   grid::Vector{Float64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    reps = size(A, 1) / length(grid)
    if Float64(reps) != Float64(Int64(reps))
        error("Dimension of A does not match dimensions of the grid.")
    end
    diagonal = zeros(size(A, 1))
    for i = 1:reps
        reps[1 + length(grid) * (i - 1): length(grid) * i] = grid
    end
    weight_mat = spdiagm(diagonal, 0)
    return solve_kfe(A, g0, weight_mat; maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
end

# Repeats previous code when entries of A are floats.
# Exact method for well-conditioned A
function solve_kfe(A::SparseMatrixCSC{Float64, Int64}, partitions::Vector{Float64},
                   dimensions::Vector{Int64})
    AT = A'
    dim_prod = prod(dimensions)
    # Fixing one value to ensure matrix is not singular
    i_fix = 1
    b = zeros(dim_prod, 1)
    b[i_fix] = .1
    for j = 1:dim_prod
        AT[i_fix,j] = 0
    end
    AT[i_fix, i_fix] = 1

    # Solve linear system for distribution
    gg = AT \ b
    g_sum = gg'*ones(dim_prod, 1) * prod(partitions)
    gg = gg ./ g_sum
    return vec(gg)
end
function solve_kfe(A::SparseMatrixCSC{Float64, Int64}, partitions::Matrix{Float64},
                   dimensions::Matrix{Int64})
    if size(partitions, 1) != 1 || size(dimensions, 1) != 1
        error("Need to enter partitions and/or dimensions as a 1 x n array or as a vector")
    end
    AT = A'
    dim_prod = prod(dimensions)
    # Fixing one value to ensure matrix is not singular
    i_fix = 1
    b = zeros(dim_prod, 1)
    b[i_fix] = .1
    for j = 1:dim_prod
        AT[i_fix,j] = 0
    end
    AT[i_fix, i_fix] = 1

    # Solve linear system for distribution
    gg = AT \ b
    g_sum = gg'*ones(dim_prod, 1) * prod(partitions)
    gg = gg ./ g_sum
    return vec(gg)
end

# Iterative method for ill-conditioned A
function solve_kfe(A::SparseMatrixCSC{Float64, Int64}, g0::Vector{Float64},
                   weight_mat::SparseMatrixCSC{Float64, Int64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    if size(weight_mat) != size(A)
        error("Dimension of weight matrix is incorrect.")
    end
    dim_size = size(A, 1)
    gg = vec(g0) # stack distribution matrix into vector

    # Solve linear system
    for ikfe = 1:maxit_kfe
        gg_tilde = weight_mat * gg # weight distribution points by their measure across wealth

        # 0 = A'g -> g = A'g + g -> (I - A')g = g, but possible singularity so, given scalar Delta_KFE,
        # 0 = A' Delta_KFE g -> (I - Delta_KFE A') g = g
        gg1_tilde = (speye(Float64, dim_size) - Δ_kfe * A') \ gg_tilde
        gg1_tilde = gg1_tilde ./ sum(gg1_tilde) # normalize
        gg1 = weight_mat \ gg1_tilde # undo integration over wealth

        # Check iteration for convergence
        err_kfe = maximum(abs.(gg1 - gg))
        if err_kfe < tol_kfe
            break
        end

        gg = gg1
    end

    return vec(gg)
end
# Other iterative methods when weight_mat is (1) not sparse,
# (2) a matrix that integrates out non-wealth variables,
# (3) a vector of differences over the wealth grid
function solve_kfe(A::SparseMatrixCSC{Float64, Int64}, g0::Vector{Float64},
                   weight_mat::Matrix{Float64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    if size(weight_mat) == size(A)
        return solve_kfe(A, g0, sparse(weight_mat); maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
    else
        weight_mat = spdiagm(vec(weight_mat), 0)
        return solve_kfe(A, g0, weight_mat; maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
    end
end
function solve_kfe(A::SparseMatrixCSC{Float64, Int64}, g0::Vector{Float64},
                   grid::Vector{Float64};
                   maxit_kfe::Int64 = 1000, tol_kfe::Float64 = 1e-12,
                   Δ_kfe::Float64 = 1e6)
    reps = size(A, 1) / length(grid)
    if Float64(reps) != Float64(Int64(reps))
        error("Dimension of A does not match dimensions of the grid.")
    end
    diagonal = zeros(size(A, 1))
    for i = 1:reps
        reps[1 + length(grid) * (i - 1): length(grid) * i] = grid
    end
    weight_mat = spdiagm(diagonal, 0)
    return solve_kfe(A, g0, weight_mat; maxit_kfe = maxit_kfe, tol_kfe = tol_kfe, Δ_kfe = Δ_kfe)
end


function f_model(
    f_x::R,
    grad_f_x::AbstractArray{C},
    res::AbstractArray{C},
    gamma::R,
) where {R<:Real,C<:RealOrComplex{R}}
    return f_x - real(dot(grad_f_x, res)) + (0.5 / gamma) * norm(res)^2
end


"""
A collection of methods and submodules useful for computation of quantum sytems.
"""
module QBase

using LinearAlgebra
using Base.Iterators: flatten
using Combinatorics: permutations, combinations
using RandomMatrices: Haar

# Methods extended by QBase.jl
import Base: convert, *, kron, adjoint, show, sqrt
import LinearAlgebra: rank, eigvals

const ATOL = 1e-7

# math utilities
include("math/matrices.jl")
include("math/combinatorics.jl")

# type exports
include("types/probabilities.jl")
include("types/brakets.jl")
include("types/operators.jl")
include("types/states.jl")
include("types/unitaries.jl")
include("types/measurements.jl")

# TODO: Unitary
# TODO: KrausOperator
# TODO: AbstractChannel
# TODO: AbstractSuperOperator
# TODO: AbstractObservable

include("constructors/brakets.jl")
include("constructors/states.jl")
include("constructors/unitaries.jl")
include("constructors/measurements.jl")

# main methods
include("evolve.jl")
include("measure.jl")
include("channels.jl")
include("information.jl")

end


using OrdinaryDiffEq, CUDA, LinearAlgebra, Test
function f(u,p,t)
    A*u
end
function f(du,u,p,t)
    mul!(du,A,u)
end
function jac(J,u,p,t)
    J .= A
end
function jac(u,p,t)
    A
end
function tgrad(du,u,p,t)
    du .= 0
end
function tgrad(u,p,t)
    zero(u)
end
ff = ODEFunction(f,jac=jac,tgrad=tgrad)
CUDA.allowscalar(false)
A = cu(-rand(3,3))
u0 = cu([1.0;0.0;0.0])
tspan = (0.,100.)

prob = ODEProblem(ff,u0,tspan)
sol = solve(prob,Tsit5())
@test solve(prob,Rosenbrock23()).retcode == :Success
solve(prob,Rosenbrock23(autodiff=false));

prob_oop = ODEProblem{false}(ff,u0,tspan)
CUDA.allowscalar(false)
sol = solve(prob_oop,Tsit5())
@test solve(prob_oop,Rosenbrock23()).retcode == :Success
@test solve(prob_oop,Rosenbrock23(autodiff=false)).retcode == :Success

prob_nojac = ODEProblem(f,u0,tspan)
@test solve(prob_nojac,Rosenbrock23()).retcode == :Success
@test solve(prob_nojac,Rosenbrock23(autodiff=false)).retcode == :Success
@test solve(prob_nojac,Rosenbrock23(autodiff=false,diff_type = Val{:central})).retcode == :Success
@test solve(prob_nojac,Rosenbrock23(autodiff=false,diff_type = Val{:complex})).retcode == :Success

#=
prob_nojac_oop = ODEProblem{false}(f,u0,tspan)
DiffEqBase.prob2dtmin(prob_nojac_oop)
@test_broken solve(prob_nojac_oop,Rosenbrock23()).retcode == :Success
@test_broken solve(prob_nojac_oop,Rosenbrock23(autodiff=false)).retcode == :Success
@test_broken solve(prob_nojac_oop,Rosenbrock23(autodiff=false,diff_type = Val{:central})).retcode == :Success
@test_broken solve(prob_nojac_oop,Rosenbrock23(autodiff=false,diff_type = Val{:complex})).retcode == :Success
=#

# Complex Numbers Adaptivity DifferentialEquations.jl#460
f_complex(u,nothing,t) = 1/2 .*u
u0 = cu(rand(32,32).+ 1im*rand(32,32));
prob = ODEProblem(f_complex,u0,(0.0,1.0))
@test_nowarn sol = solve(prob,Tsit5())


# SPDX-License-Identifier: MIT

"""
    loggamma_dot

Wrapper around the log-gamma version to avoid conflicts with CUDA
"""
loggamma_dot(x::AbstractVector) = loggamma.(x)
loggamma_dot(x::CuArray) = lgamma.(x)

"""
    pdmat_logdet(M)

Log-determinant of a positive definite matrix.
"""
function pdmat_logdet(M::AbstractMatrix)
    U = cholesky(M).U
    sum(2 * log.(diag(U)))
end
pdmat_logdet(M::Diagonal) = sum(log.(diag(M)))

"""
    pdmat_inverse(M)

Inverse of a positive definite matrix.
"""
function pdmat_inverse(M::AbstractMatrix)
    U = cholesky(M).U
    I = typeof(M)(LinearAlgebra.I, size(M)...)
    U⁻ = I / U
    U⁻ * U⁻'
end
pdmat_inverse(M::Diagonal) = Diagonal(1 ./ diag(M))

"""
    vec_tril(M)

Returns the vectorized low-triangular part (diagonal not included) of
the matrix.

See also [`inv_vec_tril`](@ref), [`matrix`](@ref).
"""
function vec_tril(M)
    D, _ = size(M)
    vcat([diag(M, -i) for i in 1:D-1]...)
end

"""
    inv_vec_tril(v)

Returns a lower triangular matrix from a vectorized form. The diagonal
of the matrix is set to 0.

See also [`vec_tril`](@ref), [`matrix`](@ref)
"""
function inv_vec_tril(v)
    # Determine the dimension of the matrix from the length of the
    # vector. To do this, we solve the equation:
    #   x²/2 + x/2 - l = 0
    # The dimension is given by s (the positive solution) + 1.
    l = length(v)
    D = Int((-1 + sqrt(1 + 8*l))/2) + 1

    M = similar(v, D, D)
    fill!(M, 0)
    #M = zeros(eltype(v), D, D)
    offset = 0
    for i in 1:D-1
        M[diagind(M, -i)] = v[offset+1:offset+D-i]
        offset += D-i
    end
    M

    #LowerTriangular(M)
end

"""
    matrix(diagM, vec_trilM)

Returns a symmetrix matrix from the diagonal and the "tril" form of a
matrix.

See also [`vec_tril`](@ref), [`inv_vec_tril`](@ref)
"""
function matrix(diagM, vec_trilM)
    T = eltype(diagM)
    trilM = inv_vec_tril(vec_trilM)
    retval = trilM + trilM'
    retval[diagind(retval)] .= diagM
    retval
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# %%
A = 10(1:9) .+ (1:9)'

# %%
i = 4
@views A00 = A[begin:i-1, begin:i-1]

# %%
@views a10 = A[begin:i-1, i:i]

# %%
@views A02 = A[begin:i-1, i+1:end]

# %%
@views A10 = A[i:i, begin:i-1]

# %%
@views a11 = A[i:i, i:i]

# %%
@views a12 = A[i:i, i+1:end]

# %%
@views A20 = A[i+1:end, begin:i-1]

# %%
@views a21 = A[i+1:end, i:i]

# %%
@views A22 = A[i+1:end, i+1:end]

# %%
?"subarray"

# %%
?view

# %%
?@view

# %%
?@views

# %%


#
# 2D Triangulations
# =================
#
# These examples can be loaded into Julia (Revise.jl recommended)
#
# These examples return either a [`SimplexGridBuilder`](@ref) struct which can
# be turned into a grid by calling [`simplexgrid`](@ref)(builder), or
# just an ExtendableGrid with default simplex grid data.
#
# The control flags for [Triangle](https://github.com/JuliaGeometry/Triangulate.jl)
# are created based on default options
# provided by this module which try to ensure "good" grids for
# FEM and FVM computations. These are documented in [`default_options`](@ref)
# Occasional [`options!`](@ref) statements in the examples overwrite these defaults.
#
#
# This test code is released under the license conditions of
# Triangulate.jl
#

using SimplexGridFactory
using ExtendableGrids
using LinearAlgebra
using Triangulate

# ## Domain triangulation
# Here we just describe a domain as a polygon and
# mesh it.
#
function triangulation_of_domain()

    builder=SimplexGridBuilder(Generator=Triangulate)

    p1=point!(builder,0,0)
    p2=point!(builder,1,0)
    p3=point!(builder,1,2)
    p4=point!(builder,0,1)
    p5=point!(builder,-1,2)

    facet!(builder,p1,p2)
    facet!(builder,p2,p3)
    facet!(builder,p3,p4)
    facet!(builder,p4,p5)
    facet!(builder,p5,p1)

    builder
end

#
# ![](triangulation_of_domain.svg)
#

#
# ## Triangulation with size control and boundary markers
# The previous example was a little bit bland.
# We miss:
# - size control for the triangles
# - differently marking of boundary parts
#
function nicer_triangulation_of_domain()

    builder=SimplexGridBuilder(Generator=Triangulate)

    p1=point!(builder,0,0)
    p2=point!(builder,1,0)
    p3=point!(builder,1,2)
    p4=point!(builder,0,1)
    p5=point!(builder,-1,2)

    facetregion!(builder,1)
    facet!(builder,p1,p2)
    facet!(builder,p2,p3)
    facetregion!(builder,2)
    facet!(builder,p3,p4)
    facet!(builder,p4,p5)
    facetregion!(builder,3)
    facet!(builder,p5,p1)

    options!(builder,maxvolume=0.01)

    builder
end
#
# ![](nicer_triangulation_of_domain.svg)
#

# ## Triangulation with subregions
# Here we create different subregions and
# apply the maxvolume constraint to
# the subregions
#
function triangulation_of_domain_with_subregions()

    builder=SimplexGridBuilder(Generator=Triangulate)

    p1=point!(builder,0,0)
    p2=point!(builder,1,0)
    p3=point!(builder,1,2)
    p4=point!(builder,0,1)
    p5=point!(builder,-1,2)

    facetregion!(builder,1)
    facet!(builder,p1,p2)
    facet!(builder,p2,p3)
    facetregion!(builder,2)
    facet!(builder,p3,p4)
    facet!(builder,p4,p5)
    facetregion!(builder,3)
    facet!(builder,p5,p1)

    facetregion!(builder,4)
    facet!(builder,p1,p4)

    cellregion!(builder,2)
    maxvolume!(builder, 0.1)
    regionpoint!(builder, -0.1,0.5)

    cellregion!(builder,3)
    maxvolume!(builder, 0.01)
    regionpoint!(builder, 0.2,0.2)

    builder
end
#
# ![](triangulation_of_domain_with_subregions.svg)
#

# ## Direct specification of input arrays
# Of course we can specify the input for Triangle
# directly. The aim of SimplexBuilder is to avoid
# the tedious and error prone counting connected
# with this approach.
function direct_square(Generator=Triangulate)
    simplexgrid(Generator;
                points=[0 0 ; 0 1 ; 1 1 ; 1 0]',
                bfaces=[1 2 ; 2 3 ; 3 4 ; 4 1 ]',
                bfaceregions=[1, 2, 3, 4],
                regionpoints=[0.5 0.5;]',
                regionnumbers=[1],
                regionvolumes=[0.01])
end
#
# ![](direct_square.svg)
#

# ## Local refinement
#
# We can interface to Triangle's unsuitable mechanism
#
function square_localref()

    builder=SimplexGridBuilder(Generator=Triangulate)
    cellregion!(builder,1)
    maxvolume!(builder,0.01)
    regionpoint!(builder,0.5,0.5)

    p1=point!(builder,0,0)
    p2=point!(builder,1,0)
    p3=point!(builder,1,1)
    p4=point!(builder,0,1)

    facetregion!(builder,1)
    facet!(builder,p1,p2)
    facetregion!(builder,2)
    facet!(builder,p2,p3)
    facetregion!(builder,3)
    facet!(builder,p3,p4)
    facetregion!(builder,4)
    facet!(builder,p4,p1)

    refinement_center=[0.5,0.5]
    function unsuitable(x1,y1,x2,y2,x3,y3, area)
        bary=[(x1+x2+x3)/3,(y2+y2+y3)/3]
        dist=norm(bary-refinement_center)
        if area > 0.01*dist
            return 1
        else
            return 0
        end
    end
    options!(builder, unsuitable=unsuitable)
    builder
end
#
# ![](square_localref.svg)
#

# ## Domain with holes
# We can generate domains with holes.
# This at once shall demonstrate how the chosen
# API approach eases bookeeping of features added to the
# geometry description
#
function swiss_cheese_2d()

    function circlehole!(builder, center, radius; n=20)
        points=[point!(builder, center[1]+radius*sin(t),center[2]+radius*cos(t)) for t in range(0,2π,length=n)]
        for i=1:n-1
            facet!(builder,points[i],points[i+1])
        end
        facet!(builder,points[end],points[1])
        holepoint!(builder,center)
    end

    builder=SimplexGridBuilder(Generator=Triangulate)
    cellregion!(builder,1)
    maxvolume!(builder,0.1)
    regionpoint!(builder,0.1,0.1)

    p1=point!(builder,0,0)
    p2=point!(builder,10,0)
    p3=point!(builder,10,10)
    p4=point!(builder,0,10)

    facetregion!(builder,1)
    facet!(builder,p1,p2)
    facet!(builder,p2,p3)
    facet!(builder,p3,p4)
    facet!(builder,p4,p1)

    holes=[8.0 4.0;
           1.0 2.0;
           8.0 9.0;
           3.0 4.0;
           4.0 6.0;
           7.0 9.0;
           4.0 7.0;
           7.0 5.0;
           2.0 1.0;
           4.0 1.0;
           4.0 8.0;
           2.0 8.0;
           3.0 6.0;
           4.0 9.0;
           9.0 1.0;
           9.0 1.0;
           6.0 9.0;
           8.0 9.0;
           3.0 5.0;
           1.0 4.0]'

    radii=[0.15, 0.15, 0.1, 0.35, 0.2, 0.3, 0.1, 0.4, 0.1, 0.4, 0.4, 0.15, 0.2, 0.2, 0.2, 0.35, 0.15, 0.25, 0.15, 0.25]

    for i=1:length(radii)
        facetregion!(builder,i+1)
        circlehole!(builder,holes[:,i], radii[i])
    end

    builder
end
#
# ![](swiss_cheese_2d.svg)
#


function kmr_markov_matrix_sequential{T<:Real}(n::Integer, p::T, ε::T)
    """
    Generate the markov matrix for the KMR model with *sequential* move

    n: number of players
    p: level of p-dominance for action 2
       = the value of p such that action 2 is the BR for (1-q, q) for any q > p,
         where q (1-q, resp.) is the prob that the opponent plays action 2 (1, resp.)
    ε: mutation probability

    References:
        kmr_markov_matrix_sequential is contributed from https://github.com/oyamad
    """
    P = zeros(T, n+1, n+1)

    P[1, 1], P[1, 2] = 1 - ε/2, ε/2
    @inbounds for i = 1:n-1
        P[i+1, i] = (i/n) * (ε/2 + (1 - ε) *
                             (((i-1)/(n-1) < p) + ((i-1)/(n-1) == p)/2))
        P[i+1, i+2] = ((n-i)/n) * (ε/2 + (1 - ε) *
                                 ((i/(n-1) > p) + (i/(n-1) == p)/2))
        P[i+1, i+1] = 1 - P[i+1, i] - P[i+1, i+2]
    end
    P[end, end-1], P[end, end] = ε/2, 1 - ε/2
    return P
end

function Base.isapprox{T<:Real,S<:Real}(x::Vector{Vector{T}},
                                        y::Vector{Vector{S}})
    n = length(x)
    length(y) == n || return false
    for i in 1:n
        isapprox(x[i], y[i])::Bool || return false
    end
    return true
end

@testset "Testing mc_tools.jl" begin
    # Matrix with two recurrent classes [1, 2] and [4, 5, 6],
    # which have periods 2 and 3, respectively
    Q = [0 1 0 0 0 0
         1 0 0 0 0 0
         1//2 0 0 1//2 0 0
         0 0 0 0 1 0
         0 0 0 0 0 1
         0 0 0 1 0 0]
    Q_stationary_dists = Vector{Rational{Int}}[
        [1//2, 1//2, 0, 0, 0, 0], [0, 0, 0, 1//3, 1//3, 1//3]
    ]
    Q_dict_Rational = Dict(
        "P" => Q,
        "stationary_dists" => Q_stationary_dists,
        "comm_classes" => Vector{Int}[[1, 2], [3], [4, 5, 6]],
        "rec_classes" => Vector{Int}[[1, 2], [4, 5, 6]],
        "is_irreducible" => false,
        "period" => 6,
        "is_aperiodic" => false,
    )
    Q_dict_Float64 = copy(Q_dict_Rational)
    Q_dict_Float64["P"] = convert(Matrix{Float64}, Q)
    Q_dict_Float64["stationary_dists"] =
        convert(Vector{Vector{Float64}}, Q_stationary_dists)

    # examples from
    # Graph-Theoretic Analysis of Finite Markov Chains by J.P. Jarvis & D. R. Shier

    fig1_p = zeros(Rational{Int64}, 5, 5)
    fig1_p[[3, 4, 9, 10, 11, 13, 18, 19, 22, 24]] =
        [1//2, 2//5, 1//10, 1, 1, 1//5, 3//10, 1//5, 1, 3//10]
    fig2_p = zeros(Rational{Int64}, 5, 5)
    fig2_p[[3, 10, 11, 13, 14, 17, 18, 19, 22]] =
        [1//3, 1, 1, 1//3, 1//2, 1//2, 1//3, 1//2, 1//2]

    fig1_dict_Rational = Dict(
        "P" => fig1_p,
        "stationary_dists" => Vector{Rational{Int}}[[0, 1//2, 0, 0, 1//2]],
        "comm_classes" => Vector{Int}[[1, 3, 4], [2, 5]],
        "rec_classes" => Vector{Int}[[2, 5]],
        "is_irreducible" => false,
        "period" => 2,
        "is_aperiodic" => false,
    )
    fig2_dict_Rational = Dict(
        "P" => fig2_p,
        "stationary_dists" => Vector{Rational{Int}}[[1//6, 0, 3//6, 2//6, 0]],
        "comm_classes" => Vector{Int}[[1, 3, 4], [2, 5]],
        "rec_classes" => Vector{Int}[[1, 3, 4]],
        "is_irreducible" => false,
        "period" => 1,
        "is_aperiodic" => true,
    )
    fig1_dict_Float64 = copy(fig1_dict_Rational)
    fig2_dict_Float64 = copy(fig2_dict_Rational)
    for (d_R, d_F) in zip((fig1_dict_Rational, fig2_dict_Rational),
                          (fig1_dict_Float64, fig2_dict_Float64))
        d_F["P"] = convert(Matrix{Float64}, d_R["P"])
        d_F["stationary_dists"] =
            convert(Vector{Vector{Float64}}, d_R["stationary_dists"])
    end

    testcases_Float64 = [
        Q_dict_Float64,
        Dict(
            "P" => [0.4 0.6; 0.2 0.8],
            "stationary_dists" => Vector{Float64}[[0.25, 0.75]],
            "comm_classes" => Vector{Int}[collect(1:2)],
            "rec_classes" => Vector{Int}[collect(1:2)],
            "is_irreducible" => true,
            "period" => 1,
            "is_aperiodic" => true,
            # "cyclic_classes" => Vector[collect(1:2)],
        ),
        Dict(
            "P" => [0. 1.; 1. 0.],
            "stationary_dists" => Vector{Float64}[[0.5, 0.5]],
            "comm_classes" => Vector{Int}[collect(1:2)],
            "rec_classes" => Vector{Int}[collect(1:2)],
            "is_irreducible" => true,
            "period" => 2,
            "is_aperiodic" => false,
            # "cyclic_classes" => Vector[[1], [2]],
        ),
        Dict(
            "P" => eye(2),
            "stationary_dists" => Vector{Float64}[[1, 0], [0, 1]],
            "comm_classes" => Vector{Int}[[1], [2]],
            "rec_classes" => Vector{Int}[[1], [2]],
            "is_irreducible" => false,
            "period" => 1,
            "is_aperiodic" => true,
        ),
        # Reducible mc with a unique recurrent class,
        # where n-1 is a transient state
        Dict(
            "P" => [1. 0.; 1. 0.],
            "stationary_dists" => Vector{Float64}[[1, 0]],
            "comm_classes" => Vector{Int}[[1], [2]],
            "rec_classes" => Vector{Int}[[1]],
            "is_irreducible" => false,
            "period" => 1,
            "is_aperiodic" => true,
        ),
        # these matrices come from RMT4 section 2.2.1
        Dict(
            "P" => [1 0 0; .2 .5 .3; 0 0 1],
            "stationary_dists" => Vector{Float64}[[1, 0, 0], [0, 0, 1]],
            "comm_classes" => Vector{Int}[[1], [2], [3]],
            "rec_classes" => Vector{Int}[[1], [3]],
            "is_irreducible" => false,
            "period" => 1,
            "is_aperiodic" => true,
        ),
        Dict(
            "P" => [.7 .3 0; 0 .5 .5; 0 .9 .1],
            "stationary_dists" => Vector{Float64}[[0, 9/14, 5/14]],
            "comm_classes" => Vector{Int}[[1], [2, 3]],
            "rec_classes" => Vector{Int}[[2, 3]],
            "is_irreducible" => false,
            "period" => 1,
            "is_aperiodic" => true,
        ),
        fig1_dict_Float64,
        fig2_dict_Float64,
    ]

    @testset "test MarkovChain with Float64" begin
        for test_dict in testcases_Float64
            mc = MarkovChain(test_dict["P"])
            stationary_dists = @inferred stationary_distributions(mc)
            @test isapprox(stationary_dists, test_dict["stationary_dists"])
            @test isequal(
                sort(communication_classes(mc), by=(x->x[1])),
                test_dict["comm_classes"]
                )
            @test isequal(
                sort(recurrent_classes(mc), by=(x->x[1])),
                test_dict["rec_classes"]
                )
            @test is_irreducible(mc) == test_dict["is_irreducible"]
            @test period(mc) == test_dict["period"]
            @test is_aperiodic(mc) == test_dict["is_aperiodic"]
        end
    end

    testcases_Rational = [
        Q_dict_Rational,
        Dict(
            "P" => [2//3 1//3; 1//4 3//4],
            "stationary_dists" => Vector{Rational{Int}}[[3//7, 4//7]],
            "comm_classes" => Vector{Int}[collect(1:2)],
            "rec_classes" => Vector{Int}[collect(1:2)],
            "is_irreducible" => true,
            "period" => 1,
            "is_aperiodic" => true,
            # "cyclic_classes" => Vector[collect(1:2)],
        ),
        fig1_dict_Rational,
        fig2_dict_Rational,
    ]

    @testset "test MarkovChain with Rational" begin
        for test_dict in testcases_Rational
            mc = MarkovChain(test_dict["P"])
            stationary_dists = @inferred stationary_distributions(mc)
            @test isequal(stationary_dists, test_dict["stationary_dists"])
            @test isequal(
                sort(communication_classes(mc), by=(x->x[1])),
                test_dict["comm_classes"]
                )
            @test isequal(
                sort(recurrent_classes(mc), by=(x->x[1])),
                test_dict["rec_classes"]
                )
            @test is_irreducible(mc) == test_dict["is_irreducible"]
            @test period(mc) == test_dict["period"]
            @test is_aperiodic(mc) == test_dict["is_aperiodic"]
        end
    end

    testcases_Int = [
        Dict(
            "P" => [0 1; 1 0],
            "stationary_dists" => Vector{Rational{Int}}[[1//2, 1//2]],
            "comm_classes" => Vector{Int}[collect(1:2)],
            "rec_classes" => Vector{Int}[collect(1:2)],
            "is_irreducible" => true,
            "period" => 2,
            "is_aperiodic" => false,
            # "cyclic_classes" => Vector[[1], [2]],
        ),
        Dict(
            "P" => eye(Int, 2),
            "stationary_dists" => Vector{Rational{Int}}[[1, 0], [0, 1]],
            "comm_classes" => Vector{Int}[[1], [2]],
            "rec_classes" => Vector{Int}[[1], [2]],
            "is_irreducible" => false,
            "period" => 1,
            "is_aperiodic" => true,
        ),
        # Reducible mc with a unique recurrent class,
        # where n-1 is a transient state
        Dict(
            "P" => [1 0; 1 0],
            "stationary_dists" => Vector{Rational{Int}}[[1, 0]],
            "comm_classes" => Vector{Int}[[1], [2]],
            "rec_classes" => Vector{Int}[[1]],
            "is_irreducible" => false,
            "period" => 1,
            "is_aperiodic" => true,
        ),
        Dict(
            "P" => [0 1 0; 1 0 0; 0 0 1],
            "stationary_dists" =>
                Vector{Rational{Int}}[[1//2,1//2,0//1], [0//1,0//1,1//1]],
            "comm_classes" => Vector{Int}[[1, 2], [3]],
            "rec_classes" => Vector{Int}[[1, 2], [3]],
            "is_irreducible" => false,
            "period" => 2,
            "is_aperiodic" => false,
        ),
    ]

    @testset "test MarkovChain with Int" begin
        for test_dict in testcases_Int
            mc = MarkovChain(test_dict["P"])
            stationary_dists = @inferred stationary_distributions(mc)
            @test isequal(stationary_dists, test_dict["stationary_dists"])
            @test isequal(
                sort(communication_classes(mc), by=(x->x[1])),
                test_dict["comm_classes"]
                )
            @test isequal(
                sort(recurrent_classes(mc), by=(x->x[1])),
                test_dict["rec_classes"]
                )
            @test is_irreducible(mc) == test_dict["is_irreducible"]
            @test period(mc) == test_dict["period"]
            @test is_aperiodic(mc) == test_dict["is_aperiodic"]
        end
    end

    @testset "test Sparse MarkovChain" begin
        # Test 2 testcases from each of testcases_*
        for test_dict in [testcases[i] for i in 1:2,
                                           testcases in (testcases_Float64,
                                                         testcases_Rational,
                                                         testcases_Int)]
            mc = MarkovChain(sparse(test_dict["P"]))
            stationary_dists = @inferred stationary_distributions(mc)
            @test isapprox(stationary_dists, test_dict["stationary_dists"])
            @test isequal(
                sort(communication_classes(mc), by=(x->x[1])),
                test_dict["comm_classes"]
                )
            @test isequal(
                sort(recurrent_classes(mc), by=(x->x[1])),
                test_dict["rec_classes"]
                )
            @test is_irreducible(mc) == test_dict["is_irreducible"]
            @test period(mc) == test_dict["period"]
            @test is_aperiodic(mc) == test_dict["is_aperiodic"]
        end
    end

    tol = 1e-15
    kmr_matrices = (
        kmr_markov_matrix_sequential(27, 1/3, 1e-2),
        kmr_markov_matrix_sequential(3, 1/3, 1e-14)
    )

    @testset "test gth_solve with KMR matrices" begin
        for P in kmr_matrices
            x = gth_solve(P)

            # Check elements sum to one
            @test isapprox(sum(x), 1; atol=tol)

            # Check elements are nonnegative
            for i in 1:length(x)
                @test x[i] >= -tol
            end

            # Check x is a left eigenvector of P
            @test isapprox(vec(x'*P), x; atol=tol)
        end
    end

    @testset "test gth_solve with generator matrices" begin
        P_dict_Int = Dict(
            "P" => [-3 3; 4 -4],
            "stationary_dist" => [4//7, 3//7],
        )

        P_dict_Float64 = Dict(
            "P" => [-3. 3.; 4. -4.],
            "stationary_dist" => [4/7, 3/7],
        )

        x = gth_solve(P_dict_Int["P"])
        @test isequal(x, P_dict_Int["stationary_dist"])

        x = gth_solve(P_dict_Float64["P"])
        @test isapprox(x, P_dict_Float64["stationary_dist"])
    end

    @testset "test gth_solve with reducible matrix" begin
        P_dict = Dict(
            "P" => [0.4 0 0.6 0; 0 0 0 1; 0.2 0 0.8 0; 0 1 0 0],
            # Stationary dist for the recurrent class {1, 3}
            "stationary_dist" => [0.25, 0, 0.75, 0],
        )

        x = gth_solve(P_dict["P"])
        @test isapprox(x, P_dict["stationary_dist"])
    end

    @testset "test MarkovChain with KMR matrices" begin
        for P in kmr_matrices
            mc = MarkovChain(P)
            stationary_dists = stationary_distributions(mc)
            for x in stationary_dists
                # Check elements sum to one
                @test isapprox(sum(x), 1; atol=tol)

                # Check elements are nonnegative
                for i in 1:length(x)
                    @test x[i] >= -tol
                end

                # Check x is a left eigenvector of P
                @test isapprox(vec(x'*P), x; atol=tol)
            end
        end
    end

    @testset "test MarkovChain throws errors" begin
        # not square
        @test_throws DimensionMismatch MarkovChain(rand(4, 5))

        # state_values to long
        @test_throws DimensionMismatch MarkovChain([0.5 0.5; 0.2 0.8], rand(3))

        # first row doesn't sum to 1
        @test_throws ArgumentError MarkovChain([0.0 0.5; 0.2 0.8])

        # negative element, but sums to 1
        @test_throws ArgumentError MarkovChain([-1 1; 0.2 0.8])
    end

    mc3 = MarkovChain([0.4 0.6; 0.2 0.8])

    @testset "test simulate_indices shape" begin

        for mc in (mc3, MarkovChain(sparse(mc3.p)))
            ts_length = 10
            init = 2
            nums_sims = [3, 1]

            @test size(simulate_indices(mc, ts_length)) == (ts_length,)
            @test size(simulate_indices(mc, ts_length; init=init)) ==
                       (ts_length, )
            for num_sims in nums_sims
                X = Array(Int64, ts_length, num_sims)
                @test size(simulate_indices!(X, mc)) ==
                           (ts_length, num_sims)
                @test size(simulate_indices!(X, mc; init=init)) ==
                           (ts_length, num_sims)
            end
        end
    end  # testset

    @testset "test simulate_indices init array" begin
        for mc in (mc3, MarkovChain(sparse(mc3.p)))
            ts_length = 10
            init = 2
            Y = simulate_indices(mc, ts_length; init=init)
            @test Y[1] == init

            num_sims = 5
            X = Array(Int64, ts_length, num_sims)
            simulate_indices!(X, mc; init=init)
            @test vec(X[1, :]) == init .* ones(Int, num_sims)

            init = [2, 1]
            simulate_indices!(X, mc; init=init)
            @test vec(X[1, :]) == collect(take(cycle(init), num_sims))
        end
    end  # testset

    @testset "Sparse Matrix" begin
        # use fig1_p from above because we already checked the graph theory
        # algorithms for this stochastic matrix.
        p = fig1_p
        p_s = sparse(p)

        @testset "constructors" begin
            mc_s = MarkovChain(p_s)
            @test isa(mc_s, MarkovChain{eltype(p), typeof(p_s)})
        end

        mc_s = MarkovChain(p_s)
        mc = MarkovChain(p)

        @testset "basic correspondence with dense version" begin
            @test n_states(mc) == n_states(mc_s)
            @test maxabs(mc.p - mc_s.p) == 0.0
            @test recurrent_classes(mc) == recurrent_classes(mc_s)
            @test communication_classes(mc) == communication_classes(mc_s)
            @test is_irreducible(mc) == is_irreducible(mc_s)
            @test is_aperiodic(mc) == is_aperiodic(mc_s)
            @test period(mc) == period(mc_s)
            @test isapprox(stationary_distributions(mc_s),
                           stationary_distributions(mc))
        end
    end

    @testset "simulate" begin
        for _mc in (mc3, MarkovChain(sparse(mc3.p)))
            for T in [Float16, Float32, Float64, Int8, Int16, Int32, Int64]
                mc = MarkovChain(_mc.p, rand(T, size(_mc.p, 1)))
                ts_length = 10
                init = 2
                nums_sims = [3, 1]

                @test size(@inferred(simulate(mc, ts_length))) == (ts_length,)
                for num_sims in nums_sims
                    X = Array(T, ts_length, num_sims)
                    @test size(simulate!(X, mc)) ==
                           (ts_length, num_sims)
                    @test size(simulate!(X, mc; init=init)) ==
                           (ts_length, num_sims)
                end
            end  # state_values eltypes
        end

        @testset "test simulate init array" begin
            for mc in (mc3, MarkovChain(sparse(mc3.p)))
                ts_length = 10
                init = 2
                Y = simulate(mc, ts_length; init=init)
                @test Y[1] == mc.state_values[init]

                num_sims = 5
                X = Array(eltype(mc.state_values), ts_length, num_sims)
                simulate!(X, mc; init=init)
                @test vec(X[1, :]) == mc.state_values[init .* ones(Int, num_sims)]

                init = [2, 1]
                simulate!(X, mc; init=init)
                @test vec(X[1, :]) == mc.state_values[collect(take(cycle(init), num_sims))]
            end
        end  # testset
    end

    @testset "simulate iterators" begin
        p = [0.0 1.0 0.0 0.0
             0.0 0.0 1.0 0.0
             0.0 0.0 0.0 1.0
             1.0 0.0 0.0 0.0]
        mc = MarkovChain(p, [10.0, 20.0, 30.0, 40.0])

        mcis = MCIndSimulator(mc, 50, 1)

        want = collect(take(cycle(1:4), 50))
        have = Array(Int, 50)
        for (ix, i) in enumerate(mcis)
            have[ix] = i
        end
        @test have == want

        @test start(mcis) == (1, 0)
        for i in 1:49
            @test !done(mcis, (1, i))
        end
        @test done(mcis, (1, 50))
        @test length(mcis) == 50
        @test eltype(mcis) == Int

        mcs = MCSimulator(mc, 50, 1)
        want = collect(take(cycle([10, 20, 30, 40]), 50))
        have = zeros(Float64, 50)
        for (ix, i) in enumerate(mcs)
            have[ix] = i
        end

        @test have == want

        @test start(mcs) == (1, 0)
        for i in 1:49
            @test !done(mcs, (1, i))
        end
        @test done(mcs, (1, 50))
        @test length(mcs) == 50
        @test eltype(mcs) == Float64

        if isdefined(Base, :iteratorsize)
            @test Base.iteratorsize(mcs) == Base.HasLength()
            @test Base.iteratorsize(mcis) == Base.HasLength()
        end
    end

end  # testset


module TestHessian
    using Base.Test
    import FiniteDiff

    negative_sin(x) = -sin(x)
    negative_cos(x) = -cos(x)

    f1(x) = sin(x[1]) + cos(x[2])
    function h1!(out, x)
        out[1, 1] = negative_sin(x[1])
        out[1, 2] = 0.0
        out[2, 1] = 0.0
        out[2, 2] = negative_cos(x[2])
        return
    end

    f2(x) = 2.3 * sin(x[1]) + 3.1 * cos(x[2])
    function h2!(out, x)
        out[1, 1] = 2.3 * negative_sin(x[1])
        out[1, 2] = 0.0
        out[2, 1] = 0.0
        out[2, 2] = 3.1 * negative_cos(x[2])
        return
    end

    f3(x) = 0.1 * (x[1] - 1.0)^2 - 0.3 * (x[2] - 0.7)^2
    function h3!(out, x)
        out[1, 1] = 0.1 * 2.0
        out[1, 2] = 0.0
        out[2, 1] = 0.0
        out[2, 2] = -0.3 * 2.0
        return
    end

    f4(x) = x[1]^2 * x[2]^2
    function h4!(out, x)
        out[1, 1] = 2.0 * x[2]^2
        out[1, 2] = 4.0 * x[1] * x[2]
        out[2, 1] = 4.0 * x[1] * x[2]
        out[2, 2] = 2.0 * x[1]^2
        return
    end

    function check_error(y_true, y_approximate, err = 1 // 10_000)
        n = length(y_true)
        @test length(y_approximate) == n
        for i in 1:n
            # TODO: Decide how to handle very small inputs.
            if y_true[i] > eps(y_true[i])
                @test abs(y_true[i] - y_approximate[i]) / abs(y_true[i]) < err
            end
        end
    end

    function run_tests(n::Integer)
        funcs = (
            (f1, h1!),
            (f2, h2!),
            (f3, h3!),
            (f4, h4!),
        )

        x = Array(Float64, 2)
        output = Array(Float64, 2, 2)
        buffer = Array(Float64, 2)
        H = Array(Float64, 2, 2)

        n_dims = length(x)

        for (f, f′′!) in funcs
            for itr in 1:n
                for i in 1:n_dims
                    x[i] = itr
                end

                f′′!(H, x)

                FiniteDiff.hessian!(output, f, x, buffer)
                check_error(H, output)

                output2 = FiniteDiff.hessian(f, x)
                check_error(H, output2)

                tmp! = FiniteDiff.hessian(f, mutates = true)
                tmp!(output, x, buffer)
                check_error(H, output)

                tmp = FiniteDiff.hessian(f, mutates = false)
                output2 = tmp(x)
                check_error(H, output2)
            end
        end
    end

    @testset "hessian tests" begin
        run_tests(100)
    end
end


"""
    default_weights(g)

Return the binary undirected adjacency matrix of `g`.
"""
function default_weights(g::G) where {G<:AbstractGraph}
    m = spzeros(nv(g), nv(g))
    for e in edges(g)
        m[src(e), dst(e)] = 1
    end
    return m
end