"""
    mpifit_hsvgp(get_data, n_parts, n_dims, bounds_low, bounds_high;
    ni_low     = 5,
    ni_high    = 50,
    batch_size = 20,
    n_iters    = 200,
    n_subcycle_ELBO_local  = 20,
    n_subcycle_ELBO_global = 20)

Fit HSVGP to data for distributed computation

    get_data:                function to define how a local partition should read in its data
    ni_low:                  number of local inducing points
    ni_high:                 number of global inducing points
    batch_size:              size of batches for stochastic optimizer
    n_iters:                 number of optimization steps
    n_subcycle_ELBO_local:   subcycles for local  ELBO before MPI communication
    n_subcycle_ELBO_global:  subcycles for global ELBO before MPI communication

WARNING: Currently writes out several jld2 files containing local and global GPss
"""
function mpifit_hsvgp(get_data, n_parts, n_dims, bounds_low, bounds_high;
    ni_low     = 5,               # Number of input dimensions
    ni_high    = 50,              # Number of local inducing points
    batch_size = 20,              # Number of points used in updating local ELBO
    n_iters    = 200,             # Total updates in whole cycle
    n_subcycle_ELBO_local  = 20,  # subcycles for local ELBO before comm.
    n_subcycle_ELBO_global = 20)  # subcycles for global ELBO before comm.

    #TODO:
    #     Need to clean up saving out of local and global objects at the end.
    #     It's messy now and will get worse with many MPI processes.
    #
    #     Would be nice to write a function for prediction to identify which
    #     partition a new value of X would be in, then direct the prediction to
    #     that local svgp
    #
    #     Need to generalize the interface here. Should have a way to pass in
    #     the structure of data to be worked with. Maybe pass in a get_data
    #     function that returns the x,y data for each partition.
    #
    #     Would like to clean up some of those arguments also. I dont think
    #     we should need to have to pass all that stuff in, but not sure the
    #     best way to give global bounds information without simply specifying
    #     it for cases where all the data cannot be stored on one node.
    #

    MPI.Init()
    comm = MPI.COMM_WORLD
    myrank = MPI.Comm_rank(comm)
    mysize = MPI.Comm_size(comm)

    root = 0 # global node

    part_split = [[ii] for ii in 1:(mysize-1)]

    if n_parts > mysize - 1
        for ii = mysize:n_parts
            append!(part_split[1 + ((ii-1) % (mysize - 1))], ii)
        end
    end

    ######################## DATA SPLIT ########################################
    # This part need not be done separately. Data will already be stored on nodes

    # Root process is the global node and rest of the processes are local processes
    # datarow goes from 2 to 48601
    # numrowscsv = 48600
    # numrowseachrank = Int(numrowscsv/(mysize-1))

    if myrank != root

        data_array = [get_data(part) for part in part_split[myrank]]
        len        = length(part_split[myrank])

    end # rank != root

    # Initializing "trace" array to store ELBO values
    if myrank == root
        trace_h       = zeros(n_iters)          # To collect the global ELBO value
        store_trace_l = zeros(n_iters,n_parts) # Store trace_l from all proesses
    end
    if myrank != root
        trace_l = [zeros(n_iters) for ii in 1:len] # To collect the local ELBO value
    end

    # Now that each local node has seen the data, it will send xi_low, cov_low,
    # mean_low to global node.
    # Once the global node gets data from local nodes, it will evaluate xi_high,
    # cov_high, mean_high and theta_high
    # Following that, global node will calculate prior_mean and prior_cov and
    # send it back to local nodes.

    ###################### INITIALIZATION ######################################

    if myrank != root
        # Step 1
        # Build local SVGPs ###################
        local_svgps = [SVGP_obj(data_array[ii][1], float(data_array[ii][2]), ni_low) for ii in 1:len]

        for ii in 1:len
            # Step 2
            # Sending mean_factor, xi_low, cov_low, mean_low and N to global #######
            send_inducing_points_to_global(local_svgps[ii].params.inducing_locs, local_svgps[ii].params.inducing_mean, root, comm)
        end

        # Step 3
        # Initialize local optimizers
        lopt_lrho = [ADAM(1. * 0.0075) for ii in 1:len]     #     opt  = ADAM(0.0075, (0.5, 0.8))
        lopt_lkap = [ADAM(1. * 0.0075) for ii in 1:len]      #     opt  = ADAM(0.0075, (0.5, 0.8))
        lopt_lsig = [ADAM(1. * 0.0075) for ii in 1:len]      #     opt  = ADAM(0.0075, (0.5, 0.8))
        loptm     = [ADAM(1. * 0.01)   for ii in 1:len]        #     optm = ADAM(0.02, (0.5, 0.8))
        loptS     = [ADAM(1. * 0.002)  for ii in 1:len]       #     optS = ADAM(0.002, (0.5, 0.8))
        loptx     = [ADAM(1. * 0.01)   for ii in 1:len]        #     optx = ADAM(0.01, (0.5, 0.8))

    end # myrank != root

    if myrank == root

        # Step 3
            # Initialize global svgp and optimizers
            # Root process will receive inducing point data
            # Root process will update global parameters
            # Root process will calculate prior_mean prior_cov for local nodes
            # Root process will send prior_mean prior_cov back to local processes

        # Initialize global optimizers
        opt_lrho = ADAM(1. * 0.0075)     #     opt  = ADAM(0.0075, (0.5, 0.8))
        opt_lkap = ADAM(1. * 0.0075)     #     opt  = ADAM(0.0075, (0.5, 0.8))
        opt_lsig = ADAM(1. * 0.0075)     #     opt  = ADAM(0.0075, (0.5, 0.8))
        optm     = ADAM(1. * 0.01)       #     optm = ADAM(0.02, (0.5, 0.8))
        optS     = ADAM(1. * 0.002)      #     optS = ADAM(0.002, (0.5, 0.8))
        optx     = ADAM(1. * 0.01)       #     optx = ADAM(0.01, (0.5, 0.8))

        # Initialize global svgp
        # TODO:
        #     Need bounds on input space and some idea of scale for output
        #         to set initial GP variance parameter and
        #         Latin Hypercube for inducing point locations
        #
        #     Assuming we will work in the case that we can't send all inducing point
        #         data to global. Instead should probably get bounds as arguments to original
        #         function or otherwise.
        fake_data_to_init = SVGP_data(
            vcat(bounds_low', bounds_high'),
            randn(4)*10 # sampling 4 points with something like the variance I was to initialize at
        )
        global_svgp = SVGP_obj(ni_high, SVGP_data("global"), SVGP_params(fake_data_to_init, ni_high), "gaussian")

        for part = (root+1):n_parts
            recv_id = 1 + (part - 1) % (mysize - 1)
            # TODO: Would prefer to not do this sequentially by rank number but as communications arrive
            #

            # println("In Process(1): ",recv_id,"\n")

            # Receiving xi_low and yi_low #
            xi_low, yi_low = receive_inducing_points_in_global(ni_low, n_dims, recv_id, comm) ### TODO: IMPORTANT! Need to communicate which partition is being recieved

            # Update global parameters from local data #
            grads  = gradient(gp -> svgp_elbo_global(xi_low, yi_low, gp, global_svgp), global_svgp.params)[1][]

            Flux.Optimise.update!(opt_lrho, global_svgp.params.log_rho,       -grads.log_rho)
            Flux.Optimise.update!(opt_lkap, global_svgp.params.log_kappa,     -grads.log_kappa)
            Flux.Optimise.update!(opt_lsig, global_svgp.params.log_sigma,     -grads.log_sigma)
            Flux.Optimise.update!(optm,     global_svgp.params.inducing_mean, -grads.inducing_mean)
            Flux.Optimise.update!(optS,     global_svgp.params.inducing_L,    -LowerTriangular(grads.inducing_L) )
            Flux.Optimise.update!(optx,     global_svgp.params.inducing_locs, -grads.inducing_locs)
            # println("global ELBO optimized")

            # Sending prior_mean, prior_cov to local process ###################
            send_prior_to_local( pred_vgp(xi_low, global_svgp), recv_id, comm) ### TODO: IMPORTANT! Need to communicate which partition is being recieved
            # println("prior_mean, prior_cov sent to local")

        end # for loop

    end # myrank == root

    ###################### END OF INITIALIZATION ###############################

    ##################### OUTER LOOP ###########################################

    for iter= 1:n_iters
        if mod(iter, 10) == 0
            print(iter, " ")
        end

        if myrank != root

            for ii in 1:len ## TODO: Clean this up with Kelin to process non-sequentially
                            ## TODO: ENSURE THE MPI PROPERLY ASSOCIATED DATA COMMUNICATED TO CORRECT PARTITION
                # Step 4
                # Each local process will receive prior_mean and prior_cov

                # Receiving prior_mean, prior_cov to local process #################
                prior_mean, prior_sd = receive_prior_in_local(
                                                        ni_low,
                                                        root,
                                                        comm)
                # println("prior_mean, prior_cov received in local")

                # Step 5
                # Update local svgp using prior_mean prior_cov from global
                for sub_iter = 1:n_subcycle_ELBO_local
                    inds   = rand(1:size(local_svgps[ii].data.y)[1], batch_size)

                    grads  = gradient(gp -> svgp_elbo_local(
                            local_svgps[ii].data.x[inds,:],
                            local_svgps[ii].data.y[inds],
                            gp,
                            local_svgps[ii],
                            (prior_mean, prior_sd)
                            ), local_svgps[ii].params)[1][]

                    Flux.Optimise.update!(lopt_lrho[ii], local_svgps[ii].params.log_rho,       -grads.log_rho)
                    Flux.Optimise.update!(lopt_lkap[ii], local_svgps[ii].params.log_kappa,     -grads.log_kappa)
                    Flux.Optimise.update!(lopt_lsig[ii], local_svgps[ii].params.log_sigma,     -grads.log_sigma)
                    Flux.Optimise.update!(loptm[ii],     local_svgps[ii].params.inducing_mean, -grads.inducing_mean)
                    Flux.Optimise.update!(loptS[ii],     local_svgps[ii].params.inducing_L,    -LowerTriangular(grads.inducing_L) )
                    Flux.Optimise.update!(loptx[ii],     local_svgps[ii].params.inducing_locs, -grads.inducing_locs)
                    if sub_iter == n_subcycle_ELBO_local
                        # Collecting local ELBO in an array ###########################
                        trace_l[ii][iter] = svgp_elbo_local(
                            local_svgps[ii].data.x[inds,:],
                            local_svgps[ii].data.y[inds],
                            local_svgps[ii].params,
                            local_svgps[ii],
                            (prior_mean, prior_sd)
                        )
                    end
                end
                # println("local ELBO optimized")

                # Step 6
                # Sending inducing point data to global
                send_inducing_points_to_global(local_svgps[ii].params.inducing_locs, local_svgps[ii].params.inducing_mean, root, comm)
            end

        end # myrank != root

        if myrank == root
            step_elbo = 0.
            for part = (root+1):n_parts
                status  = MPI.Probe(MPI.MPI_ANY_SOURCE, MPI.MPI_ANY_TAG, comm)
                recv_id = status.source

                xi_low  = zeros(ni_low, n_dims)
                rmsg    = MPI.Recv!(xi_low, recv_id, 0, comm)
                yi_low  = zeros(ni_low)
                rmsg    = MPI.Recv!(yi_low, recv_id, 1, comm)s

                # Update global parameters from local data #
                for sub_iter = 1:n_subcycle_ELBO_global
                    grads  = gradient(gp -> svgp_elbo_global(xi_low, yi_low, gp, global_svgp), global_svgp.params)[1][]

                    Flux.Optimise.update!(opt_lrho, global_svgp.params.log_rho,       -grads.log_rho)
                    Flux.Optimise.update!(opt_lkap, global_svgp.params.log_kappa,     -grads.log_kappa)
                    Flux.Optimise.update!(opt_lsig, global_svgp.params.log_sigma,     -grads.log_sigma)
                    Flux.Optimise.update!(optm,     global_svgp.params.inducing_mean, -grads.inducing_mean)
                    Flux.Optimise.update!(optS,     global_svgp.params.inducing_L,    -LowerTriangular(grads.inducing_L) )
                    Flux.Optimise.update!(optx,     global_svgp.params.inducing_locs, -grads.inducing_locs)
                end
                # println("global ELBO optimized")

                # Update step_elbo
                step_elbo += svgp_elbo(xi_low, yi_low, global_svgp.params, global_svgp)

                # Sending prior_mean, prior_cov to local process ###################
                send_prior_to_local( pred_vgp(xi_low, global_svgp), recv_id, comm) # MAKE SURE COMMUNICATION INDICATES WHICH PARTITION IN A PROCESS THIS COMM IS FOR
                # println("prior_mean, prior_cov sent to local")

            end # for loop

            trace_h[iter] = step_elbo
        end # myrank == root

        MPI.Barrier(comm)
    end # outeriter

    ############################################################################
    # Returning trace_h and trace_l to main function
    if myrank != root

        for ii in 1:len
            send_trace(trace_l[ii], root ,comm)
        end

    end #myrank != root

    if myrank == root

        for part = (root+1):n_parts
            recv_id = 1 + (part - 1) % (mysize - 1)
            trace_l = receive_trace(n_iters, recv_id, comm)
            store_trace_l[:,part] = trace_l
        end # recv_id

        return_trace = hcat(trace_h, store_trace_l)
        # println(return_trace)

        # Write trace to a csv file
        df_return_trace = DataFrame(return_trace)
        CSV.write("./return_trace.csv", df_return_trace)

    end # myrank == root
    ############################################################################

    ############################################################################
    # Save svgp_params
    if myrank != root

        # send_svgp_params(local_svgp.params, root ,comm)
        for ii in 1:len
            save(string("local_fitted_parameters",part_split[myrank][ii],"_",mysize-1, ".jld2"), string("local_svgp",part_split[myrank][ii]), local_svgps[ii])
        end

    end #myrank != root

    if myrank == root

        # local_svgp_params = [receive_svgp_params(recv_id, comm) for recv_id = 1:mysize-1]

        save(string("global_fitted_svgp_",mysize-1,".jld2"), "global_svgp", global_svgp)

    end # myrank == root
    ############################################################################

    MPI.Finalize()

end     # mpifit_HSVGP


# Timings for the skeleton BATSRUS for porting to GPU with OpenACC.
# 4 subplots, each of block size 8^3, 16^3, 32^3, and 64^3.
# Each subplot contains a bar plot, with nBlock as x axis, timing as y axis,
# and different columns represent nWork.
# Column have two colors, blue represent CPU, and orange represent GPU.
#
# Hongyang Zhou, hyzhou@umich.edu 04/22/2020

using PyPlot, CSV, DataFrames

#
df = DataFrame(CSV.File("results/timing_gpu.txt"))
t1 = df[df.nCell.==8,:]
t2 = df[df.nCell.==16,:]
t3 = df[df.nCell.==32,:]
t4 = df[df.nCell.==64,:]

t1_cpu = t1[t1.type.=="CPU",:]

# set width of bar
const barWidth = 0.3

# Set position of bar on X axis
r1 = collect(0:size(t1,1)-1)
r2 = [x + barWidth for x in r1]
r3 = [x + barWidth for x in r2]

fig, ax = subplots(2,2,figsize=(14,6))
#axis("tight")

# Make the plot
#a1 = ax[1,1].bar(r1, t1[:,:time], width=barWidth, edgecolor="white", label="1 thread * 256 MPI")
#a2 = ax[1,2].bar(r2, t2[:,:time], width=barWidth, edgecolor="white", label="8 threads * 32 MPI")
#a3 = ax[2,1].bar(r3, t3[:,:time], width=barWidth, edgecolor="white", label="32 threads * 8 MPI")
#a3 = ax[2,2].bar(r3, t4[:,:time], width=barWidth, edgecolor="white", label="32 threads * 8 MPI")
#=
# Add xticks on the middle of the group bars
xlabel("module", fontweight=:bold, fontsize=18)
xticks([r + barWidth for r in 0:length(t1)-1],
   ["Advance", "FaceValue", "FaceFlux", "Update", "UpdateCheck", "MessagePass"],
   fontsize=18)
plt.minorticks_on()
ax.tick_params(axis=:x, which=:minor, bottom=false)
ax.yaxis.grid(true, linestyle="--")
#ax.yaxis.grid(true, which="minor")
ylabel("time [s]", fontsize=18)
ax.set_title("Timings per time step for 256 cores, explicit 2nd Order Linde + mc3 limiter",fontsize=18)

# Create legend & Show graphic
legend(fontsize=18)

function autolabel(rects, t_percent)
   """Attach a text label above each bar in *rects*, displaying its height."""
   for (i, rect) in enumerate(rects)
      i == 1 && continue
      height = rect.get_height()
      percent = t_percent[i]
      ax.annotate(string(percent),
         xy=(rect.get_x() + rect.get_width() / 2, height),
         xytext=(0, 3),  # 3 points vertical offset
         textcoords="offset points",
         ha=:center, va=:bottom)
   end
end

autolabel(a1, t1_percent)
autolabel(a2, t2_percent)
autolabel(a3, t3_percent)
=#


using JLD
using PyPlot
include("utilities.jl")
include("self_confidence.jl")
pygui(false)

searchdir(path,key) = filter(x->contains(x,key), readdir(path))

folder = "gcloud_data"

for file in searchdir(folder,"jld")
    println("Processing: $file")
    fname = split("$(folder)/$file",".")[1]
    fsize = 12

    data = JLD.load("$fname.jld")
    max_steps = data["max_steps"]
    u_vals = data["u_vals"]
    its_axis = data["its_axis"]
    d_axis = data["d_axis"]
    DA = data["DA"]
    R = data["R"]
    rewards = data["rewards"]
    IA = data["IA"]
    ST = data["ST"]

    fig,ax_ary = PyPlot.subplots(1,1,sharex=true)
    fig[:set_size_inches](8.0,4.0)

    (x_vals,x_ticks,x_lbls,x_rng) = return_indices(DA[:])

    #  vline_frac = (mdp.road_net.gprops[:net_stats].diam - x_rng[1])/(x_rng[2]-x_rng[1])
    #  vline_tick = minimum(x_ticks)+vline_frac*(maximum(x_ticks)-minimum(x_ticks))

    # points = # of points where the Kernel Density Estimator is estimated
    # bw_method = bandwidth of the KDE
    trusted_solver_num = 9
    #  ax_ary[:violinplot](R[1:end .!= trusted_solver_num,:]',x_ticks[1:end .!= trusted_solver_num],widths=0.5,points=250,showmeans=false,showmedians=true,showextrema=true)
    tbp = ax_ary[:boxplot](R[trusted_solver_num,:]',positions=[x_ticks[trusted_solver_num]],widths=0.25,notch=false,showfliers=false,autorange=true,patch_artist=true)
    cbp = ax_ary[:boxplot](R[1:end .!= trusted_solver_num,:]',positions=x_ticks[1:end .!= trusted_solver_num],widths=0.25,notch=false,showfliers=false,autorange=true)
    #  for x in x_ticks
        #  ax_ary[:scatter](float(x)*ones(length(R[x,:])),R[x,:],alpha=0.1)
    #  end
    #  trusted_violin = ax_ary[:violinplot](R[trusted_solver_num,:]',[x_ticks[trusted_solver_num]],widths=0.5,points=250,showmeans=false,showmedians=true,showextrema=true)
    for element in ["boxes","whiskers","means","caps"]
        PyPlot.setp(cbp[element],color=:red)
    end
    for element in ["boxes","whiskers","means","caps"]
        PyPlot.setp(tbp[element],color=:blue)
    end
    for patch in tbp["boxes"]
        patch[:set_facecolor](:blue)
        patch[:set_alpha](0.35)
    end

    for i in x_ticks
        #  display(R[i,:])
        SQ = X3(R[i,:],R[trusted_solver_num,:],global_rwd_range=[minimum(R),maximum(R)])
        println("SQ: $SQ")
        #  @info "X3 at $i: $SQ"
        ax_ary[:annotate]("SQ: $(@sprintf("%0.2f",SQ))",xy=(i,mean(R[i,:])),size=8)
    end

    ax_ary[:set_ylabel]("Reward",size=fsize)
    ax_ary[:set_xticks](x_ticks)
    ax_ary[:xaxis][:set_ticklabels](x_lbls)
    ax_ary[:set_xlabel]("MCTS Depth",size=fsize)
    ax_ary[:set_title]("Reward vs. MCTS Depth",size=fsize)

    PyPlot.savefig("$fname.pdf",dpi=300,transparent=true)
end


## Dense matrix interface

dense_matrix_rows(s::CDenseMatrix) = ccall((:dense_matrix_rows, libsymengine), Int, (Ptr{Cvoid}, ), s.ptr)
dense_matrix_cols(s::CDenseMatrix) = ccall((:dense_matrix_cols, libsymengine), Int, (Ptr{Cvoid}, ), s.ptr)
function dense_matrix_rows_cols(mat::CDenseMatrix, r::UInt, c::UInt)
    ccall((:dense_matrix_rows_cols, libsymengine), Int, (Ptr{Cvoid}, UInt, UInt), s.ptr, r, c)
end

function dense_matrix_get_basic(s::CDenseMatrix, r::Int, c::Int)
    result = Basic()
    ccall((:dense_matrix_get_basic, libsymengine), Nothing, (Ref{Basic}, Ptr{Cvoid}, UInt, UInt), result, s.ptr, UInt(r), UInt(c))
    result
end

function dense_matrix_set_basic(s::CDenseMatrix, val, r::Int, c::Int)
    value = Basic(val)
    ccall((:dense_matrix_set_basic, libsymengine), Nothing, (Ptr{Cvoid}, UInt, UInt, Ref{Basic}), s.ptr, UInt(r), UInt(c), value)
    value
end

## Basic operations det, inv, transpose
function dense_matrix_det(s::CDenseMatrix)
    result = Basic()
    ccall((:dense_matrix_det, libsymengine), Nothing, (Ref{Basic}, Ptr{Cvoid}), result, s.ptr)
    result
end
function dense_matrix_inv(s::CDenseMatrix)
    result = CDenseMatrix()
    ccall((:dense_matrix_inv, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}), result.ptr, s.ptr)
    result
end

function dense_matrix_transpose(s::CDenseMatrix)
    result = CDenseMatrix()
    ccall((:dense_matrix_transpose, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}), result.ptr, s.ptr)
    result
end

function dense_matrix_submatrix(mat::CDenseMatrix, r1::Int, c1::Int, r2::Int, c2::Int, r::Int, c::Int)
    s = CDenseMatrix()
    ccall((:dense_matrix_inv, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, UInt, UInt, UInt, UInt, UInt, UInt),
          s.ptr, mat.ptr, UInt(r1), UInt(c1), UInt(r2), UInt(c2), UInt(r), UInt(c))
    s
end

## some matrix arithmetic methods
## These are unncecessary, as +,-,*,^ are inhertied by the AbstractArray Interface. Those return Array{Basic,2}. These
## return CDenseMatrix objects.
function dense_matrix_add_matrix(a::CDenseMatrix, b::CDenseMatrix)
    result = CDenseMatrix()
    ccall((:dense_matrix_add_matrix, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), result.ptr, a.ptr, b.ptr)
    result
end

function dense_matrix_mul_matrix(a::CDenseMatrix, b::CDenseMatrix)
    result = CDenseMatrix()
    ccall((:dense_matrix_mul_matrix, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), result.ptr, a.ptr, b.ptr)
    result
end

function dense_matrix_add_scalar(a::CDenseMatrix, b::Basic)
    result = CDenseMatrix()
    ccall((:dense_matrix_add_scalar, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), result.ptr, a.ptr, b.ptr)
    result
end

function dense_matrix_mul_scalar(a::CDenseMatrix, b::Basic)
    result = CDenseMatrix()
    ccall((:dense_matrix_mul_scalar, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), result.ptr, a.ptr, b.ptr)
    result
end

## Factorizations ##################################################
function dense_matrix_LU(mat::CDenseMatrix)
    ## need square?
    L = CDenseMatrix()
    U = CDenseMatrix()
    ccall((:dense_matrix_LU, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), L.ptr, U.ptr,  mat.ptr)
    (L, U)
end

## LDL decomposition
function dense_matrix_LDL(mat::CDenseMatrix)
    L = CDenseMatrix()
    D = CDenseMatrix()
    ccall((:dense_matrix_LDL, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), L.ptr, D.ptr,  mat.ptr)
    (L, D)
end

##  Fraction free LU factorization
function dense_matrix_FFLU(mat::CDenseMatrix)
    LU = CDenseMatrix()
    ccall((:dense_matrix_FFLU, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}), LU.ptr, mat.ptr)
    LU
end

## Fraction free LDU factorization
function dense_matrix_FFLDU(mat::CDenseMatrix)
    L = CDenseMatrix()
    D = CDenseMatrix()
    U = CDenseMatrix()

    ccall((:dense_matrix_FFLDU, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), L.ptr, D.ptr, U.ptr, mat.ptr)
    (L, D, U)
end

function dense_matrix_LU_solve(A::CDenseMatrix, b::CDenseMatrix)
    x = CDenseMatrix()
    ccall((:dense_matrix_LU_solve, libsymengine), Nothing, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), x.ptr, A.ptr, b.ptr)
    x
end

function dense_matrix_zeros(::Type{CDenseMatrix}, r::Int, c::Int)
    result = CDenseMatrix()
    ccall((:dense_matrix_zeros, libsymengine), Nothing, (Ptr{Cvoid}, UInt, UInt), result.ptr, UInt(r), UInt(c))
    result
end

function dense_matrix_ones(r::Int, c::Int)
    result = CDenseMatrix()
    ccall((:dense_matrix_ones, libsymengine), Nothing, (Ptr{Cvoid}, UInt, UInt), result.ptr, UInt(r), UInt(c))
    result
end

## dense_matrix_diag XXX don't have CVecBasic constructor

function dense_matrix_eye(N::Int, M::Int, k::Int=0)
    s = CDenseMatrix()
    ccall((:dense_matrix_eye, libsymengine), Nothing, (Ptr{Cvoid}, UInt, UInt, Int), s.ptr, UInt(N), UInt(M), k)
    s
end

dense_matrix_eq(a::CDenseMatrix, b::CDenseMatrix) = ccall((:dense_matrix_eq, libsymengine), Int, (Ptr{Cvoid},Ptr{Cvoid}), a.ptr, b.ptr)

## Plug into Julia's interface ##################################################

import Base: ==
==(a::CDenseMatrix, b::CDenseMatrix)  = dense_matrix_eq(a, b) == 1

## Abstract Array Interface
Base.size(s::CDenseMatrix) = (dense_matrix_rows(s), dense_matrix_cols(s))
Base.getindex(s::CDenseMatrix, r::Int, c::Int) = dense_matrix_get_basic(s, r-1, c-1)
Base.setindex!(s::CDenseMatrix, val, r::Int, c::Int) = dense_matrix_set_basic(s, val, r-1, c-1)

## special matrices
Base.zeros(::Type{CDenseMatrix}, r::Int, c::Int) = dense_matrix_zeros(s, r-1, c-1)
Base.ones(::Type{CDenseMatrix}, r::Int, c::Int) = dense_matrix_ones(r-1, c-1)

## basic functions
LinearAlgebra.det(s::CDenseMatrix) = dense_matrix_det(s)
Base.inv(s::CDenseMatrix) = dense_matrix_inv(s)
Base.transpose(s::CDenseMatrix) = dense_matrix_transpose(s)

LinearAlgebra.factorize(M::CDenseMatrix) = factorize(convert(Matrix, M))

"""
LU decomposition for CDenseMatrix, dense matrices of symbolic values

Also: lufact(a, val{:false}) for non-pivoting lu factorization
"""
function LinearAlgebra.lu(a::CDenseMatrix)
    l, u = dense_matrix_LU(a)
    convert(Matrix, l), convert(Matrix, u), Matrix{Basic}(LinearAlgebra.I, size(l)[1], size(l)[1])
end

# solve using LU_solve
import Base: \
\(A::CDenseMatrix, b::CDenseMatrix) = dense_matrix_LU_solve(A, b)
\(A::CDenseMatrix, b::Vector) = A \ convert(CDenseMatrix, [convert(Basic,u) for u in b])


function run_opf(
      parent_folder,
      folder,
      objective,
      gen_areas,
      area_transfer,
      gen_zones,
      zone_transfer,
      results_dict,
      fix_f_override,
      direct_pq,
      master_subnet,
      suffix,
      print_results;
      override_param::Dict=Dict(),
      start_vals=Dict("sn"=>Dict()),
      no_converter_loss::Bool=false,
      regularize_f::Float64=0.0,
      ipopt_max_iter::Int64=10000,
      output_to_files=output_to_files
   )
   (result, res_summary, solution_pm, binding_cnstr) = multifrequency_opf(
      "$parent_folder/$(folder)", objective,
      gen_areas=gen_areas, area_interface=area_transfer,
      gen_zones=gen_zones, zone_interface=zone_transfer,
      print_results=print_results,
      fix_f_override=fix_f_override, direct_pq=direct_pq,
      master_subnet=1, suffix=suffix, override_param=override_param,
      start_vals=start_vals, no_converter_loss=no_converter_loss,
      output_to_files=output_to_files,
      regularize_f=regularize_f,
      ipopt_max_iter=ipopt_max_iter
      )
   results_dict_out = deepcopy(results_dict)
   for (label, value) in result
      # println("label: $label")
      if !(label in keys(results_dict_out))
         results_dict_out[label] = Dict{String, Any}()
      end
      # Comment '|| true' to prevent saving outputs when not solved to optimum
      if res_summary.status[1] in [LOCALLY_SOLVED, ALMOST_LOCALLY_SOLVED] || true
         # results_dict_out[label][rstrip(folder,'/')*suffix] = value
         results_dict_out[label][rstrip(folder,'/')] = value
      end
   end
   return (result, res_summary, solution_pm, results_dict_out)
end


mutable struct LinQuadSOS
    kind::Symbol
    weights::Vector{Float64}
    indices::Vector{Int}
end

mutable struct MockLinQuadModel # <: LinQuadOptInterface.LinQuadOptimizer

    sense::Symbol

    A::Matrix{Float64}
    b::Vector{Float64}
    c::Vector{Float64}

    #Q::SparseVector{Matrix{Float64}}

    range::Vector{Float64} # constraint UB for range

    lb::Vector{Float64}
    ub::Vector{Float64}

    vartype::Vector{Cchar} # :Bin, :Int, :Con
    contype::Vector{Cchar} # :LEQ :GEQ, :EQ, :RANGE

    # sos::LinQuadSOS{LinQuadSOS}

    termination_status::MOI.TerminationStatusCode
    primal_status::MOI.ResultStatusCode
    dual_status::MOI.ResultStatusCode

    variable_primal_solution::Vector{Float64}
    variable_dual_solution::Vector{Float64}

    constraint_primal_solution::Vector{Float64}
    constraint_dual_solution::Vector{Float64}
    function MockLinQuadModel(env::Void,str::String="defaultname")
        m = new()

        m.sense = :minimize

        m.A = zeros(0,0)
        m.b = zeros(0)
        m.c = zeros(0)

        m.range = zeros(0)

        m.lb = zeros(0)
        m.ub = zeros(0)

        m.vartype = Cchar[]
        m.contype = Cchar[]

        m.termination_status = MOI.Success
        m.primal_status = MOI.FeasiblePoint
        m.dual_status = MOI.FeasiblePoint

        m.variable_primal_solution = zeros(0)
        m.variable_dual_solution = zeros(0)

        m.constraint_primal_solution = zeros(0)
        m.constraint_dual_solution = zeros(0)

        return m
    end
end
num_vars(inner::MockLinQuadModel) = length(inner.c)
num_cons(inner::MockLinQuadModel) = length(inner.rhs)

function set_variable_primal_solution!(inner::MockLinQuadModel,input::Vector)
    m.variable_primal_solution = input
end
function set_variable_dual_solution!(inner::MockLinQuadModel,input::Vector)
    m.variable_dual_solution = input
end
function set_constraint_primal_solution!(inner::MockLinQuadModel,input::Vector)
    m.constraint_primal_solution = input
end
function set_constraint_dual_solution!(inner::MockLinQuadModel,input::Vector)
    m.constraint_dual_solution = input
end
function set_termination_status!(inner::MockLinQuadModel,input)
    m.termination_status = input
end
function set_primal_status!(inner::MockLinQuadModel,input)
    m.primal_status = input
end
function set_dual_status!(inner::MockLinQuadModel,input)
    m.dual_status = input
end

const LQOI = LinQuadOptInterface
const MOI  = LQOI.MOI

const SUPPORTED_OBJECTIVES = [
    LQOI.Linear,
    #LQOI.Quad
]

const SUPPORTED_CONSTRAINTS = [
    (LQOI.Linear, LQOI.EQ),
    (LQOI.Linear, LQOI.LE),
    (LQOI.Linear, LQOI.GE),
    (Linear, IV),
    #(LQOI.Quad, LQOI.EQ),
    #(LQOI.Quad, LQOI.LE),
    #(LQOI.Quad, LQOI.GE),
    (LQOI.SinVar, LQOI.EQ),
    (LQOI.SinVar, LQOI.LE),
    (LQOI.SinVar, LQOI.GE),
    (LQOI.SinVar, LQOI.IV),
    (LQOI.SinVar, MOI.ZeroOne),
    (LQOI.SinVar, MOI.Integer),
    #(LQOI.VecVar, LQOI.SOS1),
    #(LQOI.VecVar, LQOI.SOS2),
    (LQOI.VecVar, MOI.Nonnegatives),
    (LQOI.VecVar, MOI.Nonpositives),
    (LQOI.VecVar, MOI.Zeros),
    (LQOI.VecLin, MOI.Nonnegatives),
    (LQOI.VecLin, MOI.Nonpositives),
    (LQOI.VecLin, MOI.Zeros)
]

mutable struct MockLinQuadOptimizer <: LQOI.LinQuadOptimizer
    LQOI.@LinQuadOptimizerBase
    params::Dict{String,Any}
    MockLinQuadOptimizer(::Void) = new()
end

LQOI.LinearQuadraticModel(::Type{MockLinQuadOptimizer},env::Void) = MockLinQuadModel(env,"defaultname")

function MockLinQuadOptimizer(;kwargs...)

    m = MockLinQuadOptimizer(nothing)
    m.params = Dict{String,Any}()
    MOI.empty!(m)
    for (name,value) in kwargs
        m.params[string(name)] = value
        setparam!(m.inner, string(name), value)
    end
    return m
end

function MOI.empty!(m::MockLinQuadOptimizer)
    MOI.empty!(m,nothing)
    for (name,value) in m.params
        setparam!(m.inner, name, value)
    end
end

LQOI.supported_constraints(s::MockLinQuadOptimizer) = SUPPORTED_CONSTRAINTS
LQOI.supported_objectives(s::MockLinQuadOptimizer) = SUPPORTED_OBJECTIVES

cintvec(v::Vector) = convert(Vector{Int32}, v)

LQOI.backend_type(m::MockLinQuadOptimizer, ::MOI.EqualTo{Float64})     = Cchar('E')#Cchar('=')
LQOI.backend_type(m::MockLinQuadOptimizer, ::MOI.LessThan{Float64})    = Cchar('L')#Cchar('<')
LQOI.backend_type(m::MockLinQuadOptimizer, ::MOI.GreaterThan{Float64}) = Cchar('G')#Cchar('>')
LQOI.backend_type(m::MockLinQuadOptimizer, ::MOI.Zeros)                = Cchar('E')#Cchar('=')
LQOI.backend_type(m::MockLinQuadOptimizer, ::MOI.Nonpositives)         = Cchar('L')#Cchar('<')
LQOI.backend_type(m::MockLinQuadOptimizer, ::MOI.Nonnegatives)         = Cchar('G')#Cchar('>')

# TODO - improve single type
function LQOI.change_variable_bounds!(instance::MockLinQuadOptimizer, colvec, valvec, sensevec)
    # colvec = colvec+1

    lb_len = count(x->x==Cchar('L'), sensevec)
    LB_val = Array{Float64}(0)
    sizehint!(LB_val, lb_len)
    LB_col = Array{Cint}(0)
    sizehint!(LB_col, lb_len)

    ub_len = count(x->x==Cchar('G'), sensevec)
    UB_val = Array{Float64}(0)
    sizehint!(UB_val, ub_len)
    UB_col = Array{Cint}(0)
    sizehint!(UB_col, ub_len)

    for i in eachindex(valvec)
        if sensevec[i] == Cchar('L')
            push!(LB_col, colvec[i])
            push!(LB_val, valvec[i])
        elseif sensevec[i] == Cchar('G')
            push!(UB_col, colvec[i])
            push!(UB_val, valvec[i])
        end
    end

    if lb_len > 0
        for i in eachindex(LB_col)
            instance.inner.lb[LB_col[i]] = LB_val[i]
        end
    end

    if ub_len > 0
        for i in eachindex(UB_col)
            instance.inner.ub[UB_col[i]] = UB_val[i]
        end
    end

    nothing
end

function LQOI.get_variable_lowerbound(instance::MockLinQuadOptimizer, col)
    instance.inner.lb[col]
end

function LQOI.get_variable_upperbound(instance::MockLinQuadOptimizer, col)
    instance.inner.ub[col]
end

function LQOI.get_number_linear_constraints(instance::MockLinQuadOptimizer)
    size(instance.inner.A)[1]
end

function LQOI.add_linear_constraints!(instance::MockLinQuadOptimizer, rowvec, colvec, coefvec, sensevec, rhsvec)

    rows = length(rhsvec)
    cols = size(instance.inner.A)[2]
    push!(rowvec,length(colvec)+1)

    # An = full(sparse(rowvec,colvec,coefvec,rows,cols))
    # @show cols,rows,rowvec,colvec,coefvec
    An = full(SparseMatrixCSC(cols,rows,rowvec,colvec,coefvec)')

    instance.inner.A = vcat(instance.inner.A,An)

    append!(instance.inner.b,rhsvec)
    append!(instance.inner.contype,sensevec)

    append!(instance.inner.range,[Inf for i in 1:rows])
    nothing
end

function LQOI.add_ranged_constraints!(instance::MockLinQuadOptimizer, rowvec, colvec, coefvec, rhsvec, ubvec)
    rows = length(rhsvec)
    cols = size(instance.inner.A)[2]
    push!(rowvec,length(colvec)+1)

    # An = full(sparse(rowvec,colvec,coefvec,rows,cols))
    # @show cols,rows,rowvec,colvec,coefvec
    An = full(SparseMatrixCSC(cols,rows,rowvec,colvec,coefvec)')

    instance.inner.A = vcat(instance.inner.A,An)

    append!(instance.inner.b,rhsvec)
    append!(instance.inner.contype,['R' for i in 1:rows])

    append!(instance.inner.range,ubvec)

    nothing
end

function modify_ranged_constraints!(instance::MockLinQuadOptimizer, rows, rhsvec, ubvec)
    for i in rows
        instance.inner.b[i] = rhsvec[i]
        instance.inner.range[i] = ubvec[i]
    end
    nothing
end

function LQOI.get_rhs(instance::MockLinQuadOptimizer, row)
    instance.inner.b[row]
end

function LQOI.get_linear_constraint(instance::MockLinQuadOptimizer, idx)
    vals = instance.inner.A[idx,:]

    outvals = Float64[]
    outinds = Int[]
    for i in eachindex(vals)
        if abs(vals[i]) > 0.0
            push!(outvals,vals[i])
            push!(outinds,i-1)
        end
    end
    return outinds, outvals
end

# TODO SPLIT THIS ONE
function LQOI.change_coefficient!(instance::MockLinQuadOptimizer, row, col, coef)
    # if row == 0
    #     set_dblattrlist!(instance.inner, "Obj", Cint[col], Float64[coef])
    # elseif col == 0
    #     set_dblattrlist!(instance.inner, "RHS", Cint[row], Float64[coef])
    # else
    #     chg_coeffs!(instance.inner, row, col, coef)
    #     #TODO fix this function in gurobi
    # end
end

function LQOI.delete_linear_constraints!(instance::MockLinQuadOptimizer, rowbeg, rowend)
    if rowbeg > 1 && rowend < length(instance.inner.b)
        instance.inner.A = vcat(instance.inner.A[1:rowbeg-1,:],instance.inner.A[rowend+1:end,:])
        instance.inner.b = vcat(instance.inner.b[1:rowbeg-1],instance.inner.b[rowend+1:end])
        instance.inner.contype = vcat(instance.inner.contype[1:rowbeg-1],instance.inner.contype[rowend+1:end])
        instance.inner.range = vcat(instance.inner.range[1:rowbeg-1],instance.inner.range[rowend+1:end])
    elseif rowbeg > 1
        instance.inner.A = instance.inner.A[1:rowbeg-1,:]
        instance.inner.b = instance.inner.b[1:rowbeg-1]
        instance.inner.contype = instance.inner.contype[1:rowbeg-1]
        instance.inner.range = instance.inner.range[1:rowbeg-1]
    elseif rowend < length(instance.inner.b)
        instance.inner.A = instance.inner.A[rowend+1:end,:]
        instance.inner.b = instance.inner.b[rowend+1:end]
        instance.inner.contype = instance.inner.contype[rowend+1:end]
        instance.inner.range = instance.inner.range[rowend+1:end]
    else
        instance.inner.A = zeros(0,length(instance.inner.c))
        instance.inner.b = Vector{Float64}[]
        instance.inner.contype = Cchar[]
        instance.inner.range = Vector{Float64}[]
    end
end

# TODO fix types
function variabletype(::MockLinQuadOptimizer, typeval)
    if typeval == 'B'
        return 'B'
    elseif typeval == 'I'
        return 'I'
    else
        return 'C'
    end
end
function LQOI.change_variable_types!(instance::MockLinQuadOptimizer, colvec, typevec)
    for i in eachindex(colvec)
        instance.inner.vartype[colvec[i]] = variabletype(instance, typevec[i])
    end
end

function sensetype(::MockLinQuadOptimizer, typeval)
    if typeval == 'G'
        return 'G'
    elseif typeval == 'E'
        return 'E'
    elseif typeval == 'L'
        return 'L'
    else
        return 'R'
    end
end
function LQOI.change_linear_constraint_sense!(instance::MockLinQuadOptimizer, rowvec, sensevec)
    for i in eachindex(rowvec)
        instance.inner.vartype[rowvec[i]] = sensetype(instance, sensevec[i])
    end
end

#=
LQOI.add_sos_constraint!(instance::MockLinQuadOptimizer, colvec, valvec, typ) = (add_sos!(instance.inner, typ, colvec, valvec);update_model!(instance.inner))

LQOI.delete_sos!(instance::MockLinQuadOptimizer, idx1, idx2) = (del_sos!(instance.inner, cintvec(collect(idx1:idx2)));update_model!(instance.inner))

# TODO improve getting processes
function LQOI.get_sos_constraint(instance::MockLinQuadOptimizer, idx)
    A, types = get_sos_matrix(instance.inner)
    line = A[idx,:] #sparse vec
    cols = line.nzind
    vals = line.nzval
    typ = types[idx] == Cint(1) ? :SOS1 : :SOS2
    return cols, vals, typ
end

LQOI.get_number_quadratic_constraints(instance::MockLinQuadOptimizer) = num_qconstrs(instance.inner)

#   NOTE:
# LQOI assumes 0.5 x' Q x, but Gurobi requires x' Q x so we multiply V by 0.5
LQOI.add_quadratic_constraint!(instance::MockLinQuadOptimizer, cols,coefs,rhs,sense, I,J,V) = add_qconstr!(instance.inner, cols, coefs, I, J, 0.5 * V, sense, rhs)

# LQOI.change_range_value!(instance::MockLinQuadOptimizer, rows, vals) = chg_rhsrange!(instance.inner, cintvec(rows), -vals)

# function LQOI.set_quadratic_objective!(instance::MockLinQuadOptimizer, I, J, V)
#     delq!(instance.inner)
#     for i in eachindex(V)
#         if I[i] == J[i]
#             V[i] /= 2
#         end
#     end
#     add_qpterms!(instance.inner, I, J, V)
#     return nothing
# end
=#

function LQOI.set_linear_objective!(instance::MockLinQuadOptimizer, colvec, coefvec)
    nvars = num_vars(instance.inner)
    obj = zeros(Float64, nvars)

    for i in eachindex(colvec)
        obj[colvec[i]] = coefvec[i]
    end

    instance.inner.c = obj
    nothing
end

function LQOI.change_objective_sense!(instance::MockLinQuadOptimizer, symbol)
    if symbol == :min
        instance.inner.sense = :minimize
    else
        instance.inner.sense = :maximize
    end
end

function LQOI.get_linear_objective!(instance::MockLinQuadOptimizer, x)
    for i in eachindex(x)
        x[i] = instance.inner.c[i]
    end
end

function LQOI.get_objectivesense(instance::MockLinQuadOptimizer)
    s = instance.inner.sense
    if s == :maximize
        return MOI.MaxSense
    else
        return MOI.MinSense
    end
end

LQOI.get_number_variables(instance::MockLinQuadOptimizer) = length(instance.inner.c)

function LQOI.add_variables!(instance::MockLinQuadOptimizer, int)
    append!(instance.inner.ub,Inf*ones(int))
    append!(instance.inner.lb,-Inf*ones(int))
    append!(instance.inner.c,0.0*ones(int))
    append!(instance.inner.vartype,['C' for i in 1:int])
    instance.inner.A = hcat(instance.inner.A,zeros(length(instance.inner.b),int))
end

function LQOI.delete_variables!(instance::MockLinQuadOptimizer, colbeg, colend)
    if colbeg > 1 && colend < length(instance.inner.lb)
        instance.inner.A = hcat(instance.inner.A[:,1:colbeg-1],instance.inner.A[:,colend+1:end])
        instance.inner.lb = vcat(instance.inner.lb[1:colbeg-1],instance.inner.lb[colend+1:end])
        instance.inner.ub = vcat(instance.inner.ub[1:colbeg-1],instance.inner.ub[colend+1:end])
        instance.inner.c = vcat(instance.inner.c[1:colbeg-1],instance.inner.c[colend+1:end])
        instance.inner.vartype = vcat(instance.inner.vartype[1:colbeg-1],instance.inner.vartype[colend+1:end])
    elseif colbeg > 1
        instance.inner.A = instance.inner.A[:,1:colbeg-1]
        instance.inner.lb = instance.inner.lb[1:colbeg-1]
        instance.inner.ub = instance.inner.ub[1:colbeg-1]
        instance.inner.c = instance.inner.c[1:colbeg-1]
        instance.inner.vartype = instance.inner.vartype[1:colbeg-1]
    elseif colend < length(instance.inner.lb)
        instance.inner.A = instance.inner.A[:,colend+1:end]
        instance.inner.lb = instance.inner.lb[colend+1:end]
        instance.inner.ub = instance.inner.ub[colend+1:end]
        instance.inner.c = instance.inner.c[colend+1:end]
        instance.inner.vartype = instance.inner.vartype[colend+1:end]
    else
        instance.inner.A = zeros(length(instance.inner.b),0)
        instance.inner.lb = Vector{Float64}[]
        instance.inner.ub = Vector{Float64}[]
        instance.inner.cß = Vector{Float64}[]
        instance.inner.vartype = Cchar[]
    end
end

# function LQOI.add_mip_starts!(instance::MockLinQuadOptimizer, colvec::Vector{Int}, valvec::Vector)
#     x = zeros(num_vars(instance.inner))
#     for (col, val) in zip(colvec, valvec)
#         x[col] = val
#     end
#     loadbasis(instance.inner, x)
# end

LQOI.solve_mip_problem!(instance::MockLinQuadOptimizer) = nothing

LQOI.solve_quadratic_problem!(instance::MockLinQuadOptimizer) = nothing

LQOI.solve_linear_problem!(instance::MockLinQuadOptimizer) = nothing

function LQOI.get_termination_status(instance::MockLinQuadOptimizer)
    return MOI.Success
end

function LQOI.get_primal_status(instance::MockLinQuadOptimizer)
    return MOI.FeasiblePoint
end

function LQOI.get_dual_status(instance::MockLinQuadOptimizer)
    return MOI.FeasiblePoint
end

function LQOI.get_variable_primal_solution!(instance::MockLinQuadOptimizer, place)
    for i in eachindex(place)
        place[i] = instance.inner.variable_primal_solution[i]
    end
    nothing
end

function LQOI.get_linear_primal_solution!(instance::MockLinQuadOptimizer, place)
    # Fix QUADRATICS
    for i in eachindex(place)
        place[i] = instance.inner.constraint_primal_solution[i]
    end
    nothing
end

#=
function LQOI.get_quadratic_primal_solution!(instance::MockLinQuadOptimizer, place)
    get_dblattrarray!(place, instance.inner, "QCSlack", 1)
    rhs = get_dblattrarray(instance.inner, "QCRHS", 1, num_qconstrs(instance.inner))
    for i in eachindex(place)
        place[i] = -place[i]+rhs[i]
    end
    nothing
end
=#

function LQOI.get_variable_dual_solution!(instance::MockLinQuadOptimizer, place)
    for i in eachindex(place)
        place[i] = instance.inner.variable_dual_solution[i]
    end
    nothing
end

function LQOI.get_linear_dual_solution!(instance::MockLinQuadOptimizer, place)
    # Fix QUADRATICS
    for i in eachindex(place)
        place[i] = instance.inner.constraint_dual_solution[i]
    end
    nothing
end

#LQOI.get_quadratic_dual_solution!(instance::MockLinQuadOptimizer, place) = get_dblattrarray!(place, instance.inner, "QCPi", 1)

LQOI.get_objective_value(instance::MockLinQuadOptimizer) = get_objval(instance.inner)

#=
LQOI.get_objective_bound(instance::MockLinQuadOptimizer) = get_objval(instance.inner)

function LQOI.get_relative_mip_gap(instance::MockLinQuadOptimizer)
    L = get_objval(instance.inner)
    U = get_objbound(instance.inner)
    return abs(U-L)/U
end

LQOI.get_iteration_count(instance::MockLinQuadOptimizer)  = get_iter_count(instance.inner)

LQOI.get_barrier_iterations(instance::MockLinQuadOptimizer) = get_barrier_iter_count(instance.inner)

LQOI.get_node_count(instance::MockLinQuadOptimizer) = get_node_count(instance.inner)
=#
#=
LQOI.get_farkas_dual!(instance::MockLinQuadOptimizer, place) = get_dblattrarray!(place, instance.inner, "FarkasDual", 1)

function hasdualray(instance::MockLinQuadOptimizer)
    try
        get_dblattrarray(instance.inner, "FarkasDual", 1, num_constrs(instance.inner))
        return true
    catch
        return false
    end
end

LQOI.get_unbounded_ray!(instance::MockLinQuadOptimizer, place) = get_dblattrarray!(place, instance.inner, "UnbdRay", 1)

function hasprimalray(instance::MockLinQuadOptimizer)
    try
        get_dblattrarray(instance.inner, "UnbdRay", 1, num_vars(instance.inner))
        return true
    catch
        return false
    end
end
=#
MOI.free!(m::MockLinQuadOptimizer) = nothing


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

AbstractSkipList API definition, typedefs, and
shared code for child types.

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

#===========================
Typedefs
===========================#

"""
    SkipList{T,M} <: AbstractSkipList{T,M}

A non-concurrent skip list. `T` is the type of the element that can be stored in the
`SkipList`, while `M` is the list's mode:

- `M == :List`: the same key can be stored in a skip list multiple times.
- `M == :Set`: only one copy of each key is allowed in the skip list. `SkipListSet{T}` is
  an alias for `SkipList{T,:Set}`.

# Examples

```jldoctest; setup = :(using SkipLists)
julia> list = SkipList{Int64}();

julia> insert!(list, 1); insert!(list, 2); insert!(list, 3);

julia> collect(list)
3-element Array{Int64,1}:
 1
 2
 3
```

"""
mutable struct SkipList{T,M} <: AbstractSkipList{T,M}
    # height_p: the parameter used by the geometric distribution for which the height
    # of a newly-created node is sampled. See random_height for more information.
    height_p::Float64

    # The maximum height of a node in the skip list
    max_height::Int64

    # The maximum number of elements that can be stored in a single SkipList node
    node_capacity::Int64

    # predecessor / successor buffers used to store the results found by the find_node
    # function
    predecessor_buffer::Vector{Node{T,M}}
    successor_buffer::Vector{Node{T,M}}

    # left / right sentinel nodes marking the left and right ends of the skip list
    left_sentinel::Node{T,M}
    right_sentinel::Node{T,M}

    # The height of the skip list. This should be an upper bound on the height of all
    # nodes within the skip list
    height::Int64

    # The number of elements in the skip list
    length::Int64
end

"""
    SkipListSet{T}

A non-concurrent skip list that only allows one copy of each key. `SkipListSet{T}` is an alias for
`SkipList{T,:Set}`.
"""
SkipListSet{T} = SkipList{T,:Set}

#===========================
Shared constructors
===========================#

for list_type in (:SkipList,)
    @eval $list_type{T}(args...; kws...) where T = $list_type{T,:List}(args...; kws...)
end

#===========================
Shared AbstractSkipList API
===========================#

function Base.show(io::IO, list::AbstractSkipList)
    write(io, string(list))
end

function Base.string(list::AbstractSkipList{T,M}) where {T,M}
    type_str = if M == :Set
        "SkipListSet{$T}"
    else
        "SkipList{$T}"
    end

    "$type_str(length = $(length(list)), height = $(height(list)))"
end

"""
    max_height(list::AbstractSkipList)

Return the maximum height that a node can attain in the skip list.
"""
max_height(list::AbstractSkipList) = list.max_height

"""
    height_p(list::AbstractSkipList)

Return the `height_p` parameter used by the skip list to sample a random height
for a newly generated node.
"""
height_p(list::AbstractSkipList) = list.height_p

Base.eltype(::AbstractSkipList{T}) where T = T


"""
    setup_problem(N, Npairs, pair_order)

Create the list of objects necessary for the PAPA reconstruction.
"""
function setup_problem(N::Int,Npairs::Int,pair_order::Array{Int,2})

    # Dimensions and parameter numbers
    dim = 2^N
    Choi_dim = dim^2
    dim2q = 16
    Nel = convert(Int,dim2q*(dim2q+1)/2)

    # Two-qubit bases
    Bases = zeros(Npairs,dim2q,dim,dim)
    B2 = zeros(1,16,4,4)

    bases_create!(Bases,N,Npairs,pair_order)
    Bases = convert(Array{ComplexF64,4},Bases)

    p_temp = zeros(Int,1,2)
    p_temp[1,:] = [1 2]
    # p_temp = convert(Array{Int,2},p_temp)
    bases_create!(B2,2,1,p_temp)
    B2 = convert(Array{ComplexF64,4},B2)

    # Maximally entangled states
    mrho = zeros(dim^2,dim^2)
    mrho2 = zeros(16,16)
    max_ent_state!(dim;mrho=mrho)
    max_ent_state!(4,mrho=mrho2)

    # Identities
    IN = iden(dim)
    I4 = iden(4)

    # Choi and chi matrices
    Choi = zeros(ComplexF64,Npairs,Choi_dim,Choi_dim)
    Choi2q = zeros(ComplexF64,dim2q,dim2q)
    chitemp = zeros(ComplexF64,dim2q,dim2q)
    rhotemp = zeros(ComplexF64,dim,dim)
    tempE = zeros(ComplexF64,16)
    tempF = zeros(ComplexF64,Nel)

    # Counters and qubit labels
    start = 0
    stop = 0
    qstemp = ones(Int,N)

    return (Bases,B2,N,Npairs,pair_order,mrho,mrho2,IN,I4,Choi_dim,dim2q,Nel,Choi,Choi2q,chitemp,rhotemp,tempE,tempF,start,stop,qstemp)
end

"""
Functionalized version of the full process solver
"""
function process_solver(chilist::Vector{Float64},sigmaVec::Vector{Float64},setup_tuple::Tuple,CP_penalty::Float64)
    return process_solver_full(chilist,sigmaVec,CP_penalty,setup_tuple...)
end

function process_solver!(F::Vector{Float64},chilist::Vector{Float64},sigmaVec::Vector{Float64},setup_tuple::Tuple,CP_penalty::Float64)
    F[:] = process_solver_full(chilist,sigmaVec,CP_penalty,setup_tuple...)
    nothing
end

"""
Full process solver that implements the PAPA_BEAR bootstrapping
"""
function process_solver_full(chilist::Vector{Float64},sigmaVec::Vector{Float64},CP_penalty::Float64,Bases::Array{ComplexF64,4},B2::Array{ComplexF64,4},N::Int,Npairs::Int,pair_order::Array{Int,2},mrho::Array{Float64,2},mrho2::Array{Float64,2},IN::Array{ComplexF64,2},I4::Array{ComplexF64,2},Choi_dim::Int,dim2q::Int,Nel::Int,Choi::Array{ComplexF64,3},Choi2q::Array{ComplexF64,2},chitemp::Array{ComplexF64,2},rhotemp::Array{ComplexF64,2},tempE::Array{ComplexF64,1},tempF::Array{ComplexF64,1},start::Int,stop::Int,qstemp::Vector{Int})::Vector{Float64}

    F = []

    fill!(Choi,0.)

    for nn = 1:1:Npairs
        fill!(chitemp,0.)
        start = 1 + (nn-1)*Nel
        stop = nn*Nel;
        chitemp[triu(trues(dim2q,dim2q))] = chilist[start:1:stop] + 1im*chilist[(start+Npairs*Nel):1:(stop+Npairs*Nel)]
        chitemp[:,:] = chitemp+chitemp'
        chitemp[:,:] = chitemp - diagm(0 => diag(chitemp))/2

        fill!(Choi2q,0.)
        for aa = 1:1:dim2q
            for bb = 1:1:dim2q
                # Choi state for the 2-qubit process on all N qubits
                Choi[nn,:,:] = Choi[nn,:,:] + chitemp[aa,bb]*kron(IN,Bases[nn,aa,:,:])*mrho*kron(IN,Bases[nn,bb,:,:]')

                # Choi state for the 2-qubit process on the 2-qubits only (for constraints use)
                Choi2q[:,:] = Choi2q[:,:] + chitemp[aa,bb]*kron(I4,B2[1,aa,:,:])*mrho2*kron(I4,B2[1,bb,:,:]')
            end
        end

        tempE = eigvals(Choi2q)

        F = vcat(F,[real(4*partial_trace(Choi2q,[0,1],[4,4],2)[:]-I4[:]);imag(4*partial_trace(Choi2q,[0,1],[4,4],2)[:]-I4[:]);CP_penalty*abs(sum(tempE[tempE.<0]))])

    end

    for pp = 1:1:Npairs
        fill!(chitemp,0.)  #ignore the name, used as a temp array

        fill!(qstemp,1)
        qstemp[pair_order[pp,1]] = 0
        qstemp[pair_order[pp,2]] = 0

        # Calculate the reduced process
        for dd = 1:1:dim2q
            fill!(rhotemp,0.) #ignore the name, used as a temp array
            full_process(rhotemp,Bases[pp,dd,:,:],Choi,IN)
            chitemp[:,:] = chitemp[:,:] + kron(B2[1,dd,:,:],partial_trace(rhotemp,qstemp,2*ones(Int,N),N));
        end
        chitemp[:,:] = chitemp[:,:]/2^N

        # Compare to input data
        start = 1 + (pp-1)*Nel
        stop = pp*Nel
        tempF[:] = chitemp[triu(trues(dim2q,dim2q))]

        F = vcat(F,[real(tempF)-sigmaVec[start:1:stop];imag(tempF)-sigmaVec[(start+Npairs*Nel):1:(stop+Npairs*Nel)]])
    end

    return abs.(F)

end


abstract type Preset end

export Glider, BHeptomino, GosperGun

struct Glider <: Preset
    data::Array{Int,2}
    size::Tuple{Int,Int}
end
Glider() = Glider([
    0 1 0
    0 0 1
    1 1 1
], (3,3))

struct BHeptomino <: Preset
    data::Array{Int,2}
    size::Tuple{Int,Int}
end
BHeptomino() = BHeptomino([
    1 0 1 1
    1 1 1 0
    0 1 0 0
], (3,4))

struct GosperGun <: Preset
    data::Array{Int,2}
    size::Tuple{Int,Int}
end
GosperGun() = GosperGun([
    0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0
    0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 0 0 0 0 0 0 0 0 0 0 0
    0 0 0 0 0 0 0 0 0 0 0 0 1 1 0 0 0 0 0 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 1 1
    0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 1 0 0 0 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 1 1
    1 1 0 0 0 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0
    1 1 0 0 0 0 0 0 0 0 1 0 0 0 1 0 1 1 0 0 0 0 1 0 1 0 0 0 0 0 0 0 0 0 0 0
    0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0
    0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
    0 0 0 0 0 0 0 0 0 0 0 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
], (9, 36)
)


function init_op_list(length)
    operator_list::Vector{NTuple{2,Int}} = [(0, 0) for _ in 1:length]
    return operator_list
end

abstract type AbstractQMCState{D,N} end

abstract type AbstractGroundState{D,N} <: AbstractQMCState{D,N} end
abstract type AbstractThermalState{D,N} <: AbstractQMCState{D,N} end

struct BinaryGroundState{N} <: AbstractGroundState{2,N}
    left_config::BitArray{N}
    right_config::BitArray{N}
    propagated_config::BitArray{N}

    operator_list::Vector{NTuple{2,Int}}

    linked_list::Vector{Int}
    leg_types::BitVector
    associates::Vector{NTuple{3,Int}}

    first::Vector{Int}
end

function BinaryGroundState(H::Hamiltonian{2,N}, M::Int) where N
    operator_list = init_op_list(2*M)

    len = 2*nspins(H) + 4*length(operator_list)
    linked_list = zeros(Int, len)
    leg_types = falses(len)
    associates = [(0, 0, 0) for _ in 1:len]

    first = zeros(Int, nspins(H))

    BinaryGroundState{N}(zero(H), zero(H), zero(H),
                         operator_list,
                         linked_list, leg_types, associates,
                         first)
end

function BinaryGroundState(left_config::BitArray{N}, right_config::BitArray{N}, operator_list::Vector{NTuple{2,Int}}) where N
    @assert left_config !== right_config "left_config and right_config can't be the same array!"

    len = 2*length(left_config) + 4*length(operator_list)
    linked_list = zeros(Int, len)
    leg_types = falses(len)
    associates = [(0, 0, 0) for _ in 1:len]

    first = zeros(Int, length(left_config))

    BinaryGroundState{N}(left_config, right_config, copy(left_config),
                         operator_list,
                         linked_list, leg_types, associates,
                         first)
end

struct BinaryThermalState{N} <: AbstractThermalState{2,N}
    left_config::BitArray{N}
    right_config::BitArray{N}
    propagated_config::BitArray{N}

    operator_list::Vector{NTuple{2,Int}}

    linked_list::Vector{Int}
    leg_types::BitVector
    associates::Vector{NTuple{3,Int}}

    first::Vector{Int}
    last::Vector{Int}
end

function BinaryThermalState(H::Hamiltonian{2,N}, cutoff::Int) where N
    operator_list = init_op_list(cutoff)

    len = 4*length(operator_list)
    linked_list = zeros(Int, len)
    leg_types = falses(len)
    associates = [(0, 0, 0) for _ in 1:len]

    first = zeros(Int, nspins(H))
    last = zeros(Int, nspins(H))

    BinaryThermalState{N}(zero(H), zero(H), zero(H),
                          operator_list,
                          linked_list, leg_types, associates,
                          first, last)
end

function BinaryThermalState(left_config::BitArray{N}, right_config::BitArray{N}, operator_list::Vector{NTuple{2,Int}}) where N
    @assert left_config !== right_config "left_config and right_config can't be the same array!"

    len = 4*length(operator_list)
    linked_list = zeros(Int, len)
    leg_types = falses(len)
    associates = [(0, 0, 0) for _ in 1:len]

    first = zeros(Int, length(left_config))
    last = copy(first)

    BinaryThermalState{N}(left_config, right_config, copy(left_config),
                          operator_list,
                          linked_list, leg_types, associates,
                          first, last)
end

const BinaryQMCState{N} = Union{BinaryGroundState{N}, BinaryThermalState{N}}


module GeneticAlgorithms

    # -------
    using RandomExtensions
    using Distributed

    include("Types.jl")
    using .Types

    import Dates

    import Base, Base.show, Base.isless

    export  runga,
            isless,
            freeze,
            defrost,
            generation_num,
            population

    # -------

    #isless(lhs::Entity, rhs::Entity) = lhs.fitness < rhs.fitness

    fitness!(ent, fitness_score) = ent.fitness = fitness_score

    distance(ind1 , ind2 ) = abs(ind1.fitness - ind2.fitness)

    global _g_model

    # -------

    function freeze(model::GAmodel, entity::EntityData)
        push!(model.freezer, entity)
        println("Freezing: ", entity)
    end

    function freeze(model::GAmodel, entity)
        entitydata = EntityData(entity, model.params.currentGeneration)
        freeze(model, entitydata)
    end

    freeze(entity) = freeze(_g_model, entity)

    function defrost(model::GAmodel, generation::Int)
        filter(model.freezer) do entitydata
            entitydata.generation == generation
        end
    end

    defrost(generation::Int) = defrost(_g_model, generation)

    generation_num(model::GAmodel = _g_model) = model.params.currentGeneration

    population(model::GAmodel = _g_model) = model.population

    function show_simulation(model :: GAmodel, ent)
        printed = model.specific_fitness.simulate_entity(ent, model.instructionsSet)
        @show ent.fitness
        #@show model.specific_fitness.fitness(ent)
        println(printed)
        return "fitness : $(ent.fitness) \n $printed"
    end

    function runga(mdl::Module, fit_mdl :: Module )
        # get the parameters of the specific fitness
        model = GAmodel(fit_mdl.getParams())
        model.specific_fitness = fit_mdl

        # save the ga (mutation and crossover units function)
        model.ga = mdl
        # save once the dictionary of instructions Set
        model.instructionsSet = mdl.getInstructionsSet()
        #run the ga
        runga(model; resume = false)
    end

    #= function runga(mdl::Module)  #deprecated
        model = GAmodel()
        model.instructionsSet = mdl.getInstructionsSet()
        #model.initial_pop_size = initial_pop_size
        model.ga = mdl
        # TODO, this function need a module for the model.specific_fitness !!
        runga(model; resume = false)
    end=#

    function runga(model::GAmodel; resume = false)
        stop = false

        #set _expandAmount & _expandRate, TODO I think they are not placed very well in case we want specifi paramters for each fitness
        _expandAmount = 2
        _expandRate = 4000

        if (!resume)
            #  initialize params.
            model.params.totalFitness = 0;
            model.params.targetFitness = model.specific_fitness.getTargetFitness();
            model.params.targetFitnessCount = 0;
            model.params.currentGeneration = 0;
            #stop = false;

            reset_model(model)
            create_initial_population(model)
            evaluate_population(model)
        end

        if model.params.historyPath != nothing
            #to save params in the history file
            write(model.params.historyPath, "params : " * string(model.params) * " \n")
        end

        while (model.params.currentGeneration < model.params.generations && !stop)
            #CreatenextGeneration();
            crossover_population(model, [])
            evaluate_population(model)

            lastIdx = model.params.populationSize #length(model.population)
            _fitness = model.population[lastIdx].fitness

            if model.params.currentGeneration % 100 == 0
                _log = "$(Dates.now()) , "
                _log *= "Gen : $(model.params.currentGeneration) , "
                l_1 = lastIdx -1
                _log *= "BEST: $_fitness , "
                _log *= "SECOND: $(model.population[l_1].fitness) \n"
                print(_log)
                if model.params.currentGeneration % 1000 ==0

                    _log *= show_simulation(model, model.population[lastIdx]) * "\n"
                end

                if model.params.historyPath != nothing
                    write(model.params.historyPath, _log)
                end

            end

            #println("hehe $_fitness  $(model.params.targetFitness)")
            if (model.params.targetFitness > 0 && _fitness >= model.params.targetFitness)
                #println("/////////////////////// $(model.params.targetFitnessCount)")
                model.params.targetFitnessCount = model.params.targetFitnessCount +1
                if (model.params.targetFitnessCount > 100) # default : set to 1000
                    break;
                end
            else
                # in case we expand and lose the best one
                model.params.targetFitnessCount = 0;
            end

            #=if (OnGenerationFunction != null)
                OnGenerationFunction(this);
            end=#
            if (_expandAmount > 0 && model.params.currentGeneration > 0 && model.params.currentGeneration % _expandRate == 0 && model.params.genomeSize < model.params.maxGenomeSize )
                model.params.genomeSize +=  _expandAmount;
                #_bestStatus.Fitness = 0; # Update display of best program, since genome has changed and we have a better/worse new best fitness.
            end

            model.params.currentGeneration = model.params.currentGeneration + 1
        end

        if model.params.historyPath != nothing
            best = model.population[model.params.populationSize] #[length(model.population)]
            _log = show_simulation(model, best)
            _log *= "\n Generation : $(model.params.currentGeneration) \n"
            _log *= "fitness $(best.fitness) \n"
            _log *= "dna : $(best.dna) \n"
            write(model.params.historyPath, _log)
            close(model.params.historyPath)
        end
        model
    end

    # -------

    function reset_model(model::GAmodel)
        global _g_model = model

        model.params.currentGeneration = 1
        empty!(model.population)
    end

    function create_initial_population(model::GAmodel)
        for i = 1:model.params.populationSize
            entity = model.ga.create_entity(i, model.params.genomeSize)
            push!(model.population, entity)
        end
    end

    function evaluate_population(model::GAmodel)
        #pmap(model.ga.entityToBfInstructions!, model.population)
        #scores = [ model.specific_fitness.fitness(ent, model.instructionsSet) for ent in model.population ]

        # for each entity, translate dna to bf instructions, then perform the specific fitness,
        # it saves its fitness & bonus in itself and return the fitness, see the fitness function
        scores = pmap(
            ent -> (model.ga.entityToBfInstructions!(ent); model.specific_fitness.fitness(ent, model.instructionsSet))
            , model.population)

        model.params.totalFitness = sum(scores)

        #pmap(fitness!, model.population, scores) # TODO decomment for normal use

        # the isless for the entities is defined in bfga
        sort!(model.population; rev = false)
        model.scores = sort!(scores; rev = false)
    end

    function crossover_population(model::GAmodel, groupings)
        thisGeneration = (population(model))
        _length = 1

        model.population = Any[]

        model.params.currentGeneration += 1
        g = model.ga.create_entity(model.params.currentGeneration, model.params.genomeSize)
        g2 = model.ga.create_entity(model.params.currentGeneration, model.params.genomeSize)
        #println("Generation n' $(model.params.currentGeneration) ")
        if model.params.elitism
            l = model.params.populationSize #length(thisGeneration)
            g = thisGeneration[l]
            l_1 = (model.params.populationSize) -1
            g2 = thisGeneration[l-1]
            #push!(model.population, g)
            #push!(model.pop_data, EntityData(g, model.params.currentGeneration))
            #push!(model.population, g2)
            #push!(model.pop_data, EntityData(g2, model.params.currentGeneration))

            _length += 2;
        end
        #println("_______________________________________________________________________________")
        #@show model.scores

        for i in _length:2:model.params.populationSize
            pidx1 = rouletteSelection(model)
            pidx2 = rouletteSelection(model)
            #println("$pidx1    $pidx2") # for debugging
            child1, child2 = nothing, nothing
            parent1 = thisGeneration[pidx1]
            parent2 = thisGeneration[pidx2]

            if (rand() < model.params.crossoverRate)
                child1, child2 = model.ga.crossover(Any[parent1, parent2])
            else
                child1, child2 = parent1, parent2
                parent1.bonus = 0
                parent2.bonus = 0
            end

            push!(model.population, child1)
            push!(model.population, child2)
        end

        if model.params.elitism
            push!(model.population, g)
            #push!(model.pop_data, EntityData(g, model.params.currentGeneration))
            push!(model.population, g2)
            #push!(model.pop_data, EntityData(g2, model.params.currentGeneration))
        end

        model.params.populationSize = length(model.population)
        #pmap(model.ga.clearCode!, model.population)

        # Expand genomes.
        if model.population[1].m_length != model.params.genomeSize
            #println("expannndd")
            newGenomeSize = model.params.genomeSize
            for m in model.population
                if m.m_length != newGenomeSize
                    model.ga.expand(m, newGenomeSize)
                end
            end
        end
    end

    function mutate_population(model::GAmodel) # deprecated, the mutation is performed in the crossover
        pmap(model.ga.mutate, model.population)
        #=for entity in model.population
            model.ga.mutate(entity)
        end=#
    end

    # We can add strangers in the population
    function add_stranger(model :: GAmodel, num)
        #println("-_-")
        for i in 1:num
            entity = model.ga.create_entity(model.params.currentGeneration, model.population[1].m_length)
            push!(model.population, entity)
            #@async push!(model.pop_data, EntityData(entity, model.params.currentGeneration))
        end
    end

    function rouletteSelection(model :: GAmodel)
        #idx = trunc(Int, rand()*length(model.scores))
        #_idx = idx ==0 ? 1 : idx
        #rand(1:trunc(Int, _idx))

        n = model.params.populationSize #length(model.scores)
        randomFitness = rand() * (model.scores[n] == 0 ? 1 : model.scores[n])

        idx = -1
        first = 1
        last = model.params.populationSize
        mid = trunc(Int , (last - first)/2)

        #  ArrayList's BinarySearch is for exact values only
        #  so do this by hand.
        while (idx == -1 && first <= last)
            if (randomFitness < model.scores[mid])
                last = mid;
            elseif (randomFitness > model.scores[mid])
                first = mid;
            end
            mid = trunc(Int, (first + last)/2)
            #  lies between i and i+1
            if ((last - first) == 1)
                idx = last;
            end
        end

        return idx
    end
end


"""
$(DocStringExtensions.TYPEDSIGNATURES)

A short description of the function

```jldoctest
julia> 1 + 1
2
```
"""
function determine_edge_probabilities(graph, strand)
    kmers = graph.gprops[:kmers]
    outgoing_edge_probabilities = SparseArrays.spzeros(length(kmers), length(kmers))

    for (kmer_index, kmer) in enumerate(kmers)
        if !strand
            kmer = BioSequences.reverse_complement(kmer)
        end

        downstream_neighbor_indices = Int[]
        for neighbor in BioSequences.neighbors(kmer)
            index = get_kmer_index(kmers, BioSequences.canonical(neighbor))
            # kmer must be in our dataset and there must be a connecting edge
            if !isnothing(index) && Graphs.has_edge(graph, ordered_edge(kmer_index, index))
                push!(downstream_neighbor_indices, index)
            end
        end
        sort!(unique!(downstream_neighbor_indices))

        downstream_edge_weights = Int[
            length(get(graph.edge_evidence, ordered_edge(kmer_index, neighbor_index), EdgeEvidence[])) for neighbor_index in downstream_neighbor_indices
        ]

        non_zero_indices = downstream_edge_weights .> 0
        downstream_neighbor_indices = downstream_neighbor_indices[non_zero_indices]
        downstream_edge_weights = downstream_edge_weights[non_zero_indices]

        downstream_edge_likelihoods = downstream_edge_weights ./ sum(downstream_edge_weights)

        for (neighbor_index, likelihood) in zip(downstream_neighbor_indices, downstream_edge_likelihoods)
            outgoing_edge_probabilities[kmer_index, neighbor_index] = likelihood
        end
    end
    return outgoing_edge_probabilities
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

A short description of the function

```jldoctest
julia> 1 + 1
2
```
"""
function determine_edge_probabilities(graph)
    outgoing_edge_probabilities = determine_edge_probabilities(graph, true)
    incoming_edge_probabilities = determine_edge_probabilities(graph, false)
    return outgoing_edge_probabilities, incoming_edge_probabilities
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

Get dna (db = "nuccore") or protein (db = "protein") sequences from NCBI
or get fasta directly from FTP site

```jldoctest
julia> 1 + 1
2
```
"""
function oriented_unbranching_walk(kmer_graph, vertex, orientation)
    walk = []
    viable_neighbors = find_unbranched_neighbors(kmer_graph, vertex, orientation)
    while length(viable_neighbors) == 1
#         @show "found a viable neighbor!!"
        viable_neighbor = first(viable_neighbors)
        edge = Graphs.Edge(vertex, viable_neighbor)
        push!(walk, edge)
        vertex = edge.dst
        viable_neighbors = Set{Int}()
        destination_orientations = [o.destination_orientation for o in kmer_graph.eprops[edge][:orientations]]
        for destination_orientation in destination_orientations
            union!(viable_neighbors, find_unbranched_neighbors(kmer_graph, vertex, destination_orientation))
        end
    end
    return walk
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

Get dna (db = "nuccore") or protein (db = "protein") sequences from NCBI
or get fasta directly from FTP site

```jldoctest
julia> 1 + 1
2
```
"""
function find_unbranched_neighbors(kmer_graph, vertex, orientation)
    downstream_vertices = find_downstream_vertices(kmer_graph, vertex, orientation)
#     backtrack_vertices
#     @show downstream_vertices
    if length(downstream_vertices) == 1
        downstream_vertex = first(downstream_vertices)
#         @show downstream_vertex
        edge = Graphs.Edge(vertex, downstream_vertex)
#         @show edge
        destination_orientations = [o.destination_orientation for o in kmer_graph.eprops[edge][:orientations]]
#         @show destination_orientations
        for destination_orientation in destination_orientations
            backtrack_vertices = find_downstream_vertices(kmer_graph, downstream_vertex, !destination_orientation)
#             @show backtrack_vertices
            # if the only backtrack is the vertex we're on, then we can simplify
            if backtrack_vertices == Set([vertex])
                return downstream_vertices
            end
        end
    end
    return Int[]
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

Get dna (db = "nuccore") or protein (db = "protein") sequences from NCBI
or get fasta directly from FTP site

```jldoctest
julia> 1 + 1
2
```
"""
function find_downstream_vertices(kmer_graph, vertex, orientation)
    viable_neighbors = Set{Int}()
    for neighbor in Graphs.neighbors(kmer_graph, vertex)
        not_same_vertex = vertex != neighbor
        candidate_edge = Graphs.Edge(vertex, neighbor)
        # palindromes can have multiple viable orientations
        # check each viable orientation individually
        edge_src_orientations = [e.source_orientation for e in kmer_graph.eprops[candidate_edge][:orientations]]
        for edge_src_orientation in edge_src_orientations
#             edge_src_orientation = kmer_graph.eprops[candidate_edge][:orientations].source_orientation
            viable_orientation = edge_src_orientation == orientation
            if not_same_vertex && viable_orientation
                push!(viable_neighbors, neighbor)
            end
        end
    end
    return viable_neighbors
end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function reverse_oriented_path(oriented_path)
#     reversed_path = copy(oriented_path)
#     for (index, state) in enumerate(oriented_path)
#         reversed_path[index] = Mycelia.OrientedKmer(index = state.index, orientation = !state.orientation)
#     end
#     return reverse!(reversed_path)
# end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function maximum_likelihood_walk(graph, connected_component)
#     max_count = maximum(graph.counts[connected_component])
#     max_count_indices = findall(count -> count == max_count, graph.counts[connected_component])
#     initial_node_index = rand(max_count_indices)
#     initial_node = connected_component[initial_node_index]
#     outgoing_edge_probabilities, incoming_edge_probabilities = Mycelia.determine_edge_probabilities(graph)
#     forward_walk = maximum_likelihood_walk(graph, [Mycelia.OrientedKmer(index = initial_node, orientation = true)], outgoing_edge_probabilities, incoming_edge_probabilities)
#     reverse_walk = maximum_likelihood_walk(graph, [Mycelia.OrientedKmer(index = initial_node, orientation = false)], outgoing_edge_probabilities, incoming_edge_probabilities)
#     reversed_reverse_walk = reverse!(
#         [
#             Mycelia.OrientedKmer(index = oriented_kmer.index, orientation = oriented_kmer.orientation)
#             for oriented_kmer in reverse_walk[2:end]
#         ]
#         )
#     full_path = [reversed_reverse_walk..., forward_walk...]
# end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function maximum_likelihood_walk(graph, path::Vector{Mycelia.OrientedKmer}, outgoing_edge_probabilities, incoming_edge_probabilities)
#     done = false
#     while !done
#         maximum_path_likelihood = 0.0
#         maximum_likelihood_path = Vector{Mycelia.OrientedKmer}()
#         for neighbor in Graphs.neighbors(graph.graph, last(path).index)
#             this_path = [last(path).index, neighbor]
#             this_oriented_path, this_path_likelihood =
#                 Mycelia.assess_path(this_path,
#                     graph.kmers,
#                     graph.counts,
#                     last(path).orientation,
#                     outgoing_edge_probabilities,
#                     incoming_edge_probabilities)
#             if this_path_likelihood > maximum_path_likelihood
#                 maximum_path_likelihood = this_path_likelihood
#                 maximum_likelihood_path = this_oriented_path
#             end
#         end
#         if isempty(maximum_likelihood_path) && (maximum_path_likelihood == 0.0)
#             done = true
#         else
#             append!(path, maximum_likelihood_path[2:end])
#         end
#     end
#     return path
# end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function take_a_walk(graph, connected_component)
#     max_count = maximum(graph.counts[connected_component])
#     max_count_indices = findall(count -> count == max_count, graph.counts[connected_component])
#     initial_node_index = rand(max_count_indices)
#     initial_node = connected_component[initial_node_index]
#     outgoing_edge_probabilities, incoming_edge_probabilities = Mycelia.determine_edge_probabilities(graph)

#     # walk forwards from the initial starting node
#     forward_walk = take_a_walk(graph, [Mycelia.OrientedKmer(index = initial_node, orientation = true)], outgoing_edge_probabilities, incoming_edge_probabilities)

#     # walk backwards from the initial starting node
#     reverse_walk = take_a_walk(graph, [Mycelia.OrientedKmer(index = initial_node, orientation = false)], outgoing_edge_probabilities, incoming_edge_probabilities)

#     # we need to reverse everything to re-orient against the forward walk
#     reverse_walk = reverse_oriented_path(reverse_walk)

#     # also need to drop the last node, which is equivalent to the first node of the
#     @assert last(reverse_walk) == first(forward_walk)
#     full_path = [reverse_walk[1:end-1]..., forward_walk...]
# #     @show full_path
# end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function take_a_walk(graph, path::Vector{Mycelia.OrientedKmer}, outgoing_edge_probabilities, incoming_edge_probabilities)
#     done = false
#     while !done
#         maximum_path_likelihood = 0.0
#         maximum_likelihood_path = Vector{Mycelia.OrientedKmer}()
#         for neighbor in Graphs.neighbors(graph.graph, last(path).index)
#             this_path = [last(path).index, neighbor]
#             this_oriented_path, this_path_likelihood =
#                 Mycelia.assess_path(this_path,
#                     graph.kmers,
#                     graph.counts,
#                     last(path).orientation,
#                     outgoing_edge_probabilities,
#                     incoming_edge_probabilities)
#             if this_path_likelihood > maximum_path_likelihood
#                 maximum_path_likelihood = this_path_likelihood
#                 maximum_likelihood_path = this_oriented_path
#             end
#         end
#         if isempty(maximum_likelihood_path) && (maximum_path_likelihood == 0.0)
#             done = true
#         else
#             append!(path, maximum_likelihood_path[2:end])
#         end
#     end
#     return path
# end

# use default cost 1 for breadth first seach
# use default cost 0 for depth first search
function dijkstra_step!(graph, distances, arrival_paths, queue; default_cost = 1)
    current_kmer, cost = DataStructures.dequeue_pair!(queue)
    current_orientation = BioSequences.iscanonical(current_kmer)
    current_canonical_kmer = BioSequences.canonical(current_kmer)
    current_index = graph[current_canonical_kmer, :kmer]

    present_neighbors = Vector{typeof(current_kmer)}()
    for neighbor in BioSequences.neighbors(current_kmer)
        try
            graph[BioSequences.canonical(neighbor), :kmer]
            push!(present_neighbors, neighbor)
        catch
            continue
        end
    end

    true_neighbors =
        filter(neighbor ->
            Graphs.has_edge(
                graph,
                graph[BioSequences.canonical(current_kmer), :kmer],
                graph[BioSequences.canonical(neighbor), :kmer]),
            present_neighbors)

    if !isempty(true_neighbors)
        canonical_true_neighbors = BioSequences.canonical.(true_neighbors)
        true_neighbor_is_canonical = BioSequences.iscanonical.(true_neighbors)
        neighbor_indices = map(canonical_neighbor -> graph[canonical_neighbor, :kmer], canonical_true_neighbors)
        neighbor_weights = map(neighbor_i -> MetaGraphs.get_prop(graph, current_index, neighbor_i, :weight), neighbor_indices)

        # inverse probability!!
        neighbor_costs = 1 .- (neighbor_weights ./ sum(neighbor_weights))
        neighbor_costs .+= default_cost
        # afterwards, 100% likelihood edges cost 1
        # 50% likelihood costs 1.5
        # 0% likelihood is impossible and not included

        for (neighbor, cost) in zip(true_neighbors, neighbor_costs)
            candidate_path = vcat(arrival_paths[current_kmer], current_kmer)
            candidate_path_cost = distances[current_kmer] + cost
            if distances[neighbor] > candidate_path_cost
                arrival_paths[neighbor] = candidate_path
                DataStructures.enqueue!(queue, neighbor => cost)
                distances[neighbor] = candidate_path_cost
            end
        end
    end
    return distances, arrival_paths, queue
end

function evaluate_hits(hits, forward_distances, reverse_distances, forward_arrival_paths, reverse_arrival_paths)

    lowest_cost = Inf
    optimal_path = Int[]

    for hit in hits
        reverse_complement_hit = BioSequences.reverse_complement(hit)
        total_cost = forward_distances[hit] + reverse_distances[reverse_complement_hit]
        if total_cost < lowest_cost
            forward_path = vcat(forward_arrival_paths[hit], hit)
            reverse_path = vcat(reverse_arrival_paths[reverse_complement_hit], reverse_complement_hit)
            reverse_path = reverse!(BioSequences.reverse_complement.(reverse_path))
            full_path = vcat(forward_path, reverse_path[2:end])

            lowest_cost = total_cost
            optimal_path = full_path
        end
    end

    return lowest_cost, optimal_path
end

function bidirectional_dijkstra(graph, a::T, b::T) where T <: BioSequences.AbstractMer
    forward_distances = DataStructures.DefaultDict{T, Float64}(Inf)
    forward_distances[a] = 0

    forward_arrival_paths = Dict{T, Vector{T}}()
    forward_arrival_paths[a] = Int[]

    forward_queue = DataStructures.PriorityQueue{T, Float64}()
    DataStructures.enqueue!(forward_queue, a, 0)

    reverse_b = BioSequences.reverse_complement(b)
    reverse_distances = DataStructures.DefaultDict{T, Float64}(Inf)
    reverse_distances[reverse_b] = 0

    reverse_arrival_paths = Dict{T, Vector{T}}()
    reverse_arrival_paths[reverse_b] = Int[]

    reverse_queue = DataStructures.PriorityQueue{T, Float64}()
    DataStructures.enqueue!(reverse_queue, reverse_b, 0)

    hits = Set{T}()

    while isempty(hits)
        dijkstra_step!(graph, forward_distances, forward_arrival_paths, forward_queue)
        dijkstra_step!(graph, reverse_distances, reverse_arrival_paths, reverse_queue)
        hits = intersect(
                    keys(forward_arrival_paths),
                    BioSequences.reverse_complement.(keys(reverse_arrival_paths)))
    end
    lowest_cost, optimal_path = evaluate_hits(hits, forward_distances, reverse_distances, forward_arrival_paths, reverse_arrival_paths)
    return lowest_cost, optimal_path
end

# # simple point to point
# function dijkstra(graph, a::T, b::T) where {T <: BioSequences.AbstractMer}
#     distances = DataStructures.DefaultDict{T, Float64}(Inf)
#     distances[a] = 0

#     arrival_paths = Dict{T, Vector{T}}()
#     arrival_paths[a] = Int[]

#     queue = DataStructures.PriorityQueue{T, Float64}()
#     DataStructures.enqueue!(queue, a, 0)

#     current_kmer, cost = a, Inf
#     while current_kmer != b
#         distances, arrival_paths, queue = dijkstra_step!(graph, distances, arrival_paths, queue)
#     end
#     shortest_path = vcat(arrival_paths[b], b)
#     return shortest_path, distances[b]
# end

function dijkstra(graph, a::T, targets::Set{T}; search_strategy::Union{Symbol, Missing}=missing) where {T <: BioSequences.AbstractMer}
    if ismissing(search_strategy)
        # local breadth first search to find the single target
        # a-b shortest path
        if length(targets) == 1
            search_strategy = :BFS
        # a - no targets - go find the ends
        # a - series of any target
        else
            search_strategy = :DFS
        end
    end
    if search_strategy == :BFS
        default_cost = 1
    elseif search_strategy == :DFS
        default_cost = 0
    else
        error("please specify either :BFS (breadth-first search) or :DFS (depth-first search)")
    end

    distances = DataStructures.DefaultDict{T, Float64}(Inf)
    distances[a] = 0

    arrival_paths = Dict{T, Vector{T}}()
    arrival_paths[a] = Int[]

    queue = DataStructures.PriorityQueue{T, Float64}()
    DataStructures.enqueue!(queue, a, 0)

    while !isempty(queue) && !(first(DataStructures.peek(queue)) in targets)
        distances, arrival_paths, queue =
            Mycelia.dijkstra_step!(graph, distances, arrival_paths, queue, default_cost = default_cost)
    end
    if !isempty(queue)
        discovered_destination = DataStructures.dequeue!(queue)
        @assert discovered_destination in targets
        shortest_path = vcat(arrival_paths[discovered_destination], discovered_destination)
        return shortest_path, distances[discovered_destination]
    else
        longest_path = Int[]
        distance = Inf
        for (destination, arrival_path) in arrival_paths
            this_distance = distances[destination]
            this_path = vcat(arrival_path, destination)
            if (length(this_path) > length(longest_path)) && (this_distance <= distance)
                longest_path = this_path
                distance = this_distance
            end

        end
        return longest_path, distance
    end
end

function update_remaining_targets(current_walk::AbstractVector{T}, remaining_targets::AbstractSet{T}) where T <: BioSequences.AbstractMer
    # assess whether targets have been hit in the canonical space
    remaining_targets = setdiff(BioSequences.canonical.(remaining_targets), BioSequences.canonical.(current_walk))
    # blow back out into forward and reverse_complement space
    remaining_targets = Set{T}(vcat(remaining_targets, BioSequences.reverse_complement.(remaining_targets)))
    return remaining_targets
end

function assess_downstream_weight(graph, kmer)
    # here we look to see if walking forward or backward from the initial node gets us to heavier weight options
    score = 0
    for neighbor in BioSequences.neighbors(kmer)
        try
            score += MetaGraphs.get_prop(graph, graph[BioSequences.canonical(neighbor), :kmer], :count)
        catch
            continue
        end
    end
    return score
end

# vertices should either be entire graph (by default) or a connected component
# if people want to work on just the connected component, let them induce a subgraph
function find_graph_core(graph; seed=rand(Int))

    Random.seed!(seed)

    T = typeof(MetaGraphs.get_prop(graph, 1, :kmer))

    # targets = [MetaGraphs.get_prop(graph, v, :kmer) for v in Graphs.vertices(graph)]
    # sortperm

    # targets = [
    #         MetaGraphs.get_prop(graph, v_index, :count) for v_index in
    #         sortperm([MetaGraphs.get_prop(graph, v, :count) for v in Graphs.vertices(graph)], rev=true)[1:Int(ceil(Graphs.nv(graph)/10))]]

    targets = [MetaGraphs.get_prop(graph, v, :kmer) for v in Graphs.vertices(graph) if MetaGraphs.get_prop(graph, v, :count) > 1]
            # sortperm([MetaGraphs.get_prop(graph, v, :count) for v in Graphs.vertices(graph)], rev=true)[1:Int(ceil(Graphs.nv(graph)/10))]]

    @show length(targets)
    starting_kmer = first(targets)
    max_degree = 0
    for node in targets
        node_degree = Graphs.degree(graph, graph[node, :kmer])
        if node_degree > max_degree
            max_degree = node_degree
            starting_kmer = node
        end
    end

    current_walk = [starting_kmer]
    prior_walk_length = length(current_walk)
    remaining_targets = update_remaining_targets(current_walk, Set(targets))
    done = isempty(remaining_targets)

    while !done
        # here we look to see if walking forward or backward from the current ends gets us to heavier weight options
        # we want to prioritize walks toward higher coverage nodes
        forward_score = assess_downstream_weight(graph, last(current_walk))
        reverse_score = assess_downstream_weight(graph, BioSequences.reverse_complement(first(current_walk)))
        if reverse_score > forward_score
            current_walk = reverse(BioSequences.reverse_complement.(current_walk))
        end

        forward_source = last(current_walk)
        forward_walk, forward_distance = Mycelia.dijkstra(graph, forward_source, remaining_targets, search_strategy=:DFS)
        current_walk = vcat(current_walk, forward_walk[2:end])
        remaining_targets = update_remaining_targets(current_walk, remaining_targets)
        if isempty(remaining_targets)
            done = true
        else
            reverse_source = BioSequences.reverse_complement(first(current_walk))
            reverse_walk, reverse_distance = Mycelia.dijkstra(graph, reverse_source, remaining_targets, search_strategy=:DFS)
            current_walk = vcat(reverse(BioSequences.reverse_complement.(reverse_walk))[1:end-1], current_walk)
            remaining_targets = update_remaining_targets(current_walk, remaining_targets)
        end
        @show length(current_walk)
        failed_this_expansion = length(current_walk) == prior_walk_length
        prior_walk_length = length(current_walk)
        if isempty(remaining_targets)
            done = true
        elseif failed_this_expansion
            done = true
        end
    end
    return current_walk
end


module CodeGenerator
import Base:show

# "文法"类型定义
Terminal = Symbol
Nonterminal = Symbol
Token = Union{Terminal,Nonterminal}
Statement = Vector{Token}
Production = Dict{Token,Set{Statement}}

# 状态机类型定义
State = Symbol
StateSet = Set{State}
SymbolTabel = Set{Symbol}
Transition{T} = Dict{State,Dict{Symbol,T}}

include("utils.jl")
include("grammar.jl")
include("generator.jl")
include("NFA.jl")

end # module