module CKP

  include("MDInput.jl")
  include("Atoms.jl")
  include("MDTraj.jl")
  include("Forces.jl")

  include("save.jl")
  include("plot_tools.jl")

  include("setindex2D.jl")
  include("update_status.jl")

  include("resetForces.jl")
  include("reduceForces.jl")

  include("compute_uf_partials_QUAD.jl")
  include("uf_QUAD.jl")

  include("compute_uf_partials_LJ.jl")
  include("uf_LJ.jl")

  include("image.jl")
  include("initial.jl")
  include("kinetic.jl")
  include("normal.jl")

  include("update_positions.jl")
  include("update_velocities.jl")
  include("md.jl")

  include("animate.jl")

  include("SIRDInput.jl")
  include("SIRDTraj.jl")
  include("sird.jl")
  include("plotsird.jl")

  include("fiteval.jl")
  include("fit.jl")
  include("plotfit.jl")

  include("plotU.jl")

end


# Removal of LinearKernel
Base.@deprecate LinearKernel(a::Real=1, c::Real=1) PolynomialKernel(a, c, 1)

# Renaming of GammaRational to GammaRationalQuadratic
Base.@deprecate GammaRationalKernel GammaRationalQuadraticKernel

# Removal of PeriodicKernel
struct SineSquared <: PreMetric end
@inline base_aggregate(::SineSquared, s::T, x::T, y::T) where {T} = s + sin(x-y)^2
@inline isstationary(::SineSquared) = true

struct PeriodicKernel{T<:AbstractFloat} <: MercerKernel{T}
    α::T
    function PeriodicKernel{T}(α::Real) where {T<:AbstractFloat}
        Base.depwarn("PeriodicKernel will be removed in the next major release", :PeriodicKernel)
        @check_args(PeriodicKernel, α, α > zero(α), "α > 0")
        new{T}(α)
    end
end
PeriodicKernel(α::T₁ = 1.0) where {T₁<:Real} = PeriodicKernel{promote_float(T₁)}(α)

@inline basefunction(::PeriodicKernel) = SineSquared()

@inline function kappa(κ::PeriodicKernel{T}, z::T) where {T}
    return exp(-κ.α*z)
end


using DelimitedFiles:readdlm
using Random:shuffle!
using Flux
using Plots
using CUDA

using DrugRepositioning

CUDA.allowscalar(false)

"原始数据存储路径"
data_dir=joinpath(@__DIR__,"data","origin")

"对比方法测试结果"
compare_dir=joinpath(@__DIR__,"data","all_result")
ms=compared_scores(compare_dir)

"高玲的方法的测试结果"
gao_dir=joinpath(@__DIR__,"data","all_result","gao")
P=readdlm(joinpath(gao_dir,"2","0mean5_P.txt"))[:]
FPR=readdlm(joinpath(gao_dir,"2","0mean5_FPR.txt"))[:]
TPR=readdlm(joinpath(gao_dir,"2","0mean5_TPR.txt"))[:]
R=TPR
pushfirst!(ms,MetricsScore("GPred_2",R,P,FPR,TPR,[],[]))
P=readdlm(joinpath(gao_dir,"1","0mean5_P.txt"))[:]
FPR=readdlm(joinpath(gao_dir,"1","0mean5_FPR.txt"))[:]
TPR=readdlm(joinpath(gao_dir,"1","0mean5_TPR.txt"))[:]
R=TPR
pushfirst!(ms,MetricsScore("GPred_1",R,P,FPR,TPR,[],[]))

"读取CSV"
function read_csv(file_path;d_type::DataType=Float64,i_type::DataType=String,h_dtype::DataType=String)
    f=readdlm(file_path,',')
    header=f[1,2:end]|>Array{String}
    indexer=f[2:end,1]|>Array{String}
    data=f[2:end,2:end]|>Array{d_type}
    indexer,header,data
end

"药物疾病关联"
function drug_disease()
    return read_csv(joinpath(data_dir,"drug_disease_mat.csv");d_type=Bool)
end

"疾病相似性"
function disease_similarity()
    return read_csv(joinpath(data_dir,"disease_similarity_mat.csv");d_type=Float32)
end

"药物靶标域"
drug_domain()=read_csv(joinpath(data_dir,"drug_target_domain_mat.csv");d_type=Bool)
"药物靶标GO"
drug_go()=read_csv(joinpath(data_dir,"drug_target_go_mat.csv");d_type=Bool)
"药物公共化学子结构"
drug_pubchem()=read_csv(joinpath(data_dir,"drug_pubchem_mat.csv");d_type=Bool)

drugs,diseases,RD=drug_disease()
_,_,D1=disease_similarity()
_,_,RP1=drug_domain()
_,_,RP2=drug_go()
_,_,RP3=drug_pubchem()
_,_,RP4=drug_disease()

props=Dict{Symbol,Array{Matrix}}(
    :drug=>[RP1',RP2',RP3',RP4'],
    :disease=>[]
)

dataset=Dataset(drugs,diseases,RD,props)
add_similarity!(dataset,:disease,D1)
model=CombineModel(dataset;feature_left=128,feature_right=512)|>gpu
# trainer=Trainer(model,Flux.Losses.logitcrossentropy,dataset;epochs=80)
# training(trainer,Flux.Optimise.ADAM(3e-4))
res=ValidateResult("result")
df=analyze_by_drugs(res)
R,P,FPR,TPR=pr_roc(res)
my_score=MetricsScore("MSRD",R,P,FPR,TPR,[],[])
pushfirst!(ms,my_score)
pr_roc_curve(ms)|>display
plot(res)


abstract type FitBackend end
abstract type OptimFitBackend <: FitBackend end

backend_type(::Val{:optim}) = OptimFitBackend

fit(args...; backend::Symbol = :optim, kwargs...) = fit( backend_type(Val(backend)), args...; kwargs...)
fit(::Type{OptimFitBackend}, args...; kwargs...) = opt_fit(args...; kwargs...)


function print_results(fileflag, seqpa, seqpo, J, H, L, ctype, lambda_o, lambda_e, μext, μint, seqtest, seqins, init, fin, start, fname, satmf, beta, T0, el_time)

	enmf = compute_potts_en(J, H, seqpa, L,ctype)
	enhmm = compute_potts_en(J,H, seqtest, L ,ctype)
	costmf = compute_cost_function(J,H,seqpo,L,ctype,lambda_o,lambda_e,μext,μint)
	costhmm = compute_cost_function(J,H,seqins,L,ctype,lambda_o,lambda_e,μext,μint)
	distmf, mgmf, lgmf, dmmf = hammingdist(seqpa, seqtest)

	@printf(fileflag, ">%s/%d-%d\t", fname, init + parse(Int64, start) - 1, fin + parse(Int64, start) - 1)
	@printf(fileflag, "sat: %s beta: %.2f muext: %.1f muint: %.1f T0: %s Dist: %d %d %d %d Time: %.2f\n", satmf, beta, μext, μint, T0, distmf,mgmf,lgmf,dmmf,el_time)
	@printf(fileflag, "test     %s %.3f %.3f\n", seqtest, enhmm, costhmm)
	@printf(fileflag, "dcalign  %s %.3f %.3f\n", seqpa, enmf, costmf)
	@printf(fileflag, "ins test     %s\n", seqins)
	@printf(fileflag, "ins dcalign  %s\n", seqpo)
	flush(fileflag)
end

function align_all(
    q::Int,
    L::Int,
    filename_par::String,
    filename_full::String,
    inspen_file::String,
    μext::Real,
    μint::Real;
    delta::Int = 20,
    typel::Symbol = :bm,
    filename_align::String = "",
    filename_ins::String = "",
    β::Real = 1.0,
    filename_out::String = "out.fasta",
    filename_flag::String = "flag.dat",
    seedrng::Int = 0,
    mindec::Int = 50,
    nprint::Int = 100,
    maxiter::Int = 1000,
    epsconv::Real = 1e-5,
    verbose::Bool = false,
)

    if q == 5
        #ctype = Symbol("nbase")
        ctype = :nbase
        delta = 1000
    else
        #ctype = Symbol("amino")
        ctype = :amino
    end

    if typel == :bm
        gap = 0
    else
        gap = q
    end
    println("### Reading parameters from file ###")
    J, H = read_parameters(filename_par, q, L, gap = gap, typel = typel)
    file_out = open(filename_out, "a")
    println("### Reading sequences ###")
    typedic = true
    if filename_align == "" || filename_ins == ""
        println("No delta")
        println("File full: ", filename_full)
        al = enveloptoalign(filename_full, ctype = ctype)
        names = collect(keys(al))
        fake = '-'^L
    else
        println("### Cutting the full length sequence using delta =", delta)
        file_flag = open(filename_flag, "a")
        typedic = false
        println("File full: ", filename_full)
        println("File aligned seqs: ", filename_align)
        println("File ins seqs: ", filename_ins)
        al = enveloptoalign(
            filename_full,
            filename_align,
            filename_ins,
            delta = delta,
            ctype = ctype,
        )
    end

    M = length(al)

    Lambda_all = readdlm(inspen_file)

    lambda_o = Lambda_all[:, 1]
    lambda_e = Lambda_all[:, 2]
    damp = 0.0
    println("Tot number of sequences to be aligned ", M)

    tenp = ceil(Int64, M / 10)
    println("### Starting loop ###")
    for k = 1:M
        if !typedic
            (name, hit) = split(al[k][1], "/")
            (start, finish) = split(hit, "-")
            seqtest = al[k][3]
            seqins = al[k][4]
            fullseq = al[k][2]
            seq = Seq(seqtest, fullseq, ctype)
        else
            fname = names[k]
            fullseq = al[fname]
            if occursin("/", fname)
                (name, hit) = split(fname, "/")
                (start, finish) = split(hit, "-")
            else
                name = fname
                start = "1"
                finish = @sprintf("%s", L)
            end
            seq = Seq(fake, fullseq, ctype)
        end
        if parse(Int64, start) > parse(Int64, finish) # it often happens in RNA seq
            aux = finish
            finish = start
            start = aux
        end

        N = length(fullseq)
        T0 = true
        stime = time()
        res = palign(
            seq,
            β .* J,
            β .* H,
            β .* lambda_o,
            β .* lambda_e,
            ctype = ctype,
            mindec = mindec,
            nprint = nprint,
            maxiter = maxiter,
            T0 = T0,
            μint = β * μint,
            seed = seedrng,
            damp = damp,
            μext = β * μext,
            epsconv = epsconv,
            verbose = verbose,
        )
        el_time = time() -stime
        P = res[3].pbf.P
        out = decodeposterior(P, seq.strseq)
        seqpa0 = out.pa
        seqpo0 = out.po
        sat = check_assignment(P, false, N)
        if sat == false
            (seqpa0, P) = decimate_post(res[3], T0)
            out = decodeposterior(P, seq.strseq)
            seqpo0 = out.po
        end
        b0 = out.start
        e0 = out.finish
        enT0 = compute_potts_en(J, H, seqpa0, L, ctype)
        costT0 = compute_cost_function(J,H,seqpo0,L,ctype,lambda_o,lambda_e,μext,μint)

        if !typedic
            print_results(
                file_flag,
                seqpa0,
                seqpo0,
                β .* J,
                β .* H,
                L,
                ctype,
                β * lambda_o,
                β * lambda_e,
                β * μext,
                β * μint,
                seqtest,
                seqins,
                b0,
                e0,
                start,
                name,
                sat,
                β,
                T0,
                el_time,
            )
        end

        T0 = false
        stime = time()
        res = palign(
            seq,
            β .* J,
            β .* H,
            β .* lambda_o,
            β .* lambda_e,
            ctype = ctype,
            mindec = mindec,
            nprint = nprint,
            maxiter = maxiter,
            T0 = T0,
            μint = β * μint,
            seed = seedrng,
            damp = damp,
            μext = β * μext,
            epsconv = epsconv,
            verbose = verbose,
        )
        el_time = time() - stime
        P = res[3].pbf.P
        out = decodeposterior(P, seq.strseq)
        seqpa = out.pa
        seqpo = out.po
        sat = check_assignment(P, false, N)
        if sat == false
            (seqpa, P) = decimate_post(res[3], T0)
            out = decodeposterior(P, seq.strseq)
            seqpo = out.po
        end
        enT = compute_potts_en(J, H, seqpa, L, ctype)
        costT = compute_cost_function(J,H,seqpo,L,ctype,lambda_o,lambda_e,μext,μint)
        b = out.start
        e = out.finish

        #if enT < enT0
        if costT < costT0
            @printf(
                file_out,
                ">%s/%i-%i\n",
                name,
                b + parse(Int64, start) - 1,
                e + parse(Int64, start) - 1
            )
            @printf(file_out, "%s\n", seqpa)
        else
            @printf(
                file_out,
                ">%s/%i-%i\n",
                name,
                b0 + parse(Int64, start) - 1,
                e0 + parse(Int64, start) - 1
            )
            @printf(file_out, "%s\n", seqpa0)
        end
        flush(file_out)

        if !typedic
            print_results(
                file_flag,
                seqpa,
                seqpo,
                β .* J,
                β .* H,
                L,
                ctype,
                β * lambda_o,
                β * lambda_e,
                β * μext,
                β * μint,
                seqtest,
                seqins,
                b,
                e,
                start,
                name,
                sat,
                β,
                T0,
                el_time,
            )
        end

        if mod(k, tenp) == 0
            println("Done...", k, "/", M)
        end

    end
    println("### End ###")
end


"""

        composite_and_radialaverage_withmask(azimuthal_average,composite_addition,buf,radiusbins,surface_pressure,variable_of_interest)
    """

function composite_and_radialaverage_in_chunks!(azimuthal_average,composite_addition,buf1,buf2,radiusbins,variable_name,pre_iterator_time,nchunks,file_data,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = true)
    totalcyclonecount = 0
    variables_diag = ["convec_heating_anomaly","rad_heating_anomaly","buoyancy_anomaly"]
    variables_orig_3d = ["U","V","QV","TABS","QRAD","PP","W"]
    variables_orig_2d = ["PSFC","USFC","VSFC","PW","Prec","LHF","SHF"]
    chunk_size_2d = 2*length(pre_iterator_time)÷nchunks
    chunk_size_3d = length(pre_iterator_time)÷nchunks

    nd = ndims = if (variable_name ∈ variables_orig_3d)
            nd = ndims = 4
        elseif (variable_name ∈ variables_diag)
            nd = ndims = 4
        elseif (variable_name ∈ variables_orig_2d)
            nd = ndims = 3
    end
    for chunk in 1:nchunks
        iterator_time_2d = (1200 + (chunk - 1)*chunk_size_2d + 1 ): 2 : (1200 + chunk*chunk_size_2d)
        iterator_time_data =  ((chunk - 1)*chunk_size_3d + 1 ) : (chunk*chunk_size_3d)
        iterator_time_3d = 600 .+ iterator_time_data
        if (variable_name ∈ variables_orig_3d)
            variable_of_interest = Dataset(file_data) do ds
                variable(ds,variable_name)[:,:,:,iterator_time_3d]     :: Array{Float32,4}
            end
        elseif (variable_name ∈ variables_diag)
            variable_of_interest = gather_diagnostic_3d_in_chunks(diagnostics_data_dir,exp_name,variable_name,chunk,chunk_size_3d,nchunks)
        elseif (variable_name ∈ variables_orig_2d)
            variable_of_interest = Dataset(file_data) do ds
                variable(ds,variable_name)[:,:,iterator_time_2d]     :: Array{Float32,3}
            end
        end

        #centers_labels_and_cyclones = jldopen(file_centers) do file
        #    read(file,"centers_data")
        #end

        #@info nchunks size(variable_of_interest) #size(centers_labels_and_cyclones)

        for timeindex in 1:size(variable_of_interest,nd)
            centers_and_labels,cyclones = centers_labels_and_cyclones[timeindex]
            if !isnothing(centers_and_labels[1])
                count = add_allcyclones!(composite_addition,buf1,buf2,selectdim(variable_of_interest,nd,timeindex),cyclones,centers_and_labels;maskcyclones)
                totalcyclonecount += count
            end
        end
    end

    if !iszero(totalcyclonecount)
        composite_addition ./= totalcyclonecount
        if nd == 4
            for idx in 1:(length(radiusbins) - 1)
                azimuthal_average[idx,:] .= averageallindistance((radiusbins[idx],radiusbins[idx+1]),composite_addition,(256,256), 2000 )
            end
        elseif nd == 3
            for idx in 1:(length(radiusbins) - 1)
                azimuthal_average[idx] = averageallindistance((radiusbins[idx],radiusbins[idx+1]),composite_addition,(256,256), 2000 )
            end
        end

    end
end

function gather_diagnostic_3d_in_chunks(data_dir,exp_name,desired_var,chunk,chunksize,nchunks)
    diagnostic = zeros(Float32,512,512,80,chunksize)
    howmanyfiles = 12÷nchunks
    @info "how many files" howmanyfiles
    for iter in 1:howmanyfiles
        filenumber = ((chunk - 1)*howmanyfiles + iter)

        @info "this is iter" iter
        @info "file number and chunk:" filenumber chunk
        file_name = string(exp_name,"_nosmoothing",filenumber+12,".jld")
        @info "I am opening the file $file_name and reading the variable $desired_var"
        file_path = joinpath(data_dir,file_name)
        file = jldopen(file_path, "r")
        index_initial = (iter - 1)*(50)+1
        index_end = iter*50
        @info "indices to write:" index_initial,index_end
        var_from_file = read(file, desired_var)
        diagnostic[:,:,:,index_initial:index_end] .=  var_from_file[:,:,:,:]
        close(file)
    end
    return diagnostic
end

"""
        getcomposites(file2d,file3d,diagnostics_dir,exp_,initial_timeindex,final_timeindex,partnumber,outputInterval)

"""
function get_composites_in_chunks(file2d,file3d,exp_name)

    @info "Starting composite routine"
    diagnostics_data_dir = "/global/cscratch1/sd/aramreye/for_postprocessing/ApeBudgetOutputs_nosmoothing/"
    composites_data_dir = "/global/cscratch1/sd/aramreye/for_postprocessing/CompositeOutputs_nosmoothing/"
    azimuthal_averages_data_dir = "/global/cscratch1/sd/aramreye/for_postprocessing/AzimuthualAverageOutputs_nosmoothing/"
    centers_data_file = string(composites_data_dir,"centers_detection_",exp_name,".jld")

    #variables_diag = ["convec_heating_anomaly","rad_heating_anomaly","buoyancy_anomaly"]
    #variables_orig_3d = ["U","V","QV","TABS","QRAD","PP","W"]
    #variables_orig_2d = ["PSFC","USFC","VSFC","PW","Prec","LHF","SHF"]

    initial_timeindex = 601
    #    initial_timeindex = 1195
    final_timeindex = 1200
    #final_timeindex = 650

    pre_iterator_time_2d    = initial_timeindex*2:2:final_timeindex*2
    pre_iterator_time = pre_iterator_time_3d   = initial_timeindex:final_timeindex

    radiusbins = 1000:2000:512000;
    ds3d                = Dataset(file3d)
    ds2d                = Dataset(file2d)
    z                   = variable(ds3d,"z")[:]                       :: Array{Float32,1}
    t                   = variable(ds3d,"time")[pre_iterator_time_3d]     :: Array{Float32,1}
    surface_pressure    = variable(ds2d,"PSFC")[:,:,pre_iterator_time_2d] :: Array{Float32,3}
    close(ds2d)
    close(ds3d)
    pressure_anomaly    = surface_pressure .- mean(surface_pressure,dims=(1,2))

    @info size(surface_pressure)

    ##Create buffers to reuse on allocating functions

    size_var_of_interest_2d = (512,512)
    size_var_of_interest_3d = (size_var_of_interest_2d...,length(z))
    ## Create 3D buffers
    composite_addition_3d   = zeros(Float32,size_var_of_interest_3d)
    buf_3d_2                = similar(composite_addition_3d)
    buf_3d_1                = similar(composite_addition_3d)
    azimuthal_average_3d    = zeros(Float32,length(radiusbins) - 1,length(z))
    ## Create 2D buffers
    composite_addition_2d   = zeros(Float32,size_var_of_interest_2d)
    buf_2d_1                = similar(composite_addition_2d)
    buf_2d_2                = similar(composite_addition_2d)
    azimuthal_average_2d    = zeros(Float32,length(radiusbins) - 1)
    @info size_var_of_interest_2d size_var_of_interest_3d

    centers_labels_and_cyclones = @views [
        detect_cyclones!(buf_2d_1, buf_2d_2, pressure_anomaly[:,:,timeindex],-9.0,2000) for
        timeindex in 1:size(surface_pressure,3)]; pressure_anomaly = []; surface_pressure = [];

    composite_filename_withmask = string(composites_data_dir,exp_name        ,"_withmask.jld")
    composite_filename_nomask   = string(composites_data_dir,exp_name        ,"_nomask.jld")
    azimuthal_filename_withmask = string(azimuthal_averages_data_dir,exp_name,"_withmask.jld")
    azimuthal_filename_nomask   = string(azimuthal_averages_data_dir,exp_name,"_nomask.jld")

    ### *** add allcyclones ***
    @info "Creating composite for surface pressure "
    flush(stdout)
    composite_addition_2d .= 0.0
    buf_2d_1              .= 0.0
    buf_2d_2              .= 0.0
    azimuthal_average_2d  .= 0.0

    nchunks=6

    #create files:
    ds =  jldopen(composite_filename_withmask, "w") do file end
    ds =  jldopen(composite_filename_nomask, "w") do file end
    ds =  jldopen(azimuthal_filename_withmask, "w") do file end
    ds =  jldopen(azimuthal_filename_nomask, "w") do file end

    for current_variable in variables_orig_3d

         composite_addition_3d .= 0.0
         buf_3d_1              .= 0.0
         buf_3d_2              .= 0.0
         azimuthal_average_3d  .= 0.0

    #     ## Create products with mask
         @info "Creating composite for: " current_variable
         flush(stdout)

        composite_and_radialaverage_in_chunks!(azimuthal_average_3d,composite_addition_3d,buf_3d_1,buf_3d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file3d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = true)

         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_3d)
         end
         jldopen(azimuthal_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_3d)
         end
         GC.gc()
         GC.gc()
    #     ## Create products without mask
         @info "Creating composite without mask for: " current_variable
         flush(stdout)
         composite_addition_3d .= 0.0
         buf_3d_1              .= 0.0
         buf_3d_2              .= 0.0
         azimuthal_average_3d  .= 0.0
        composite_and_radialaverage_in_chunks!(azimuthal_average_3d,composite_addition_3d,buf_3d_1,buf_3d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file3d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = false)
         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_3d)
         end
         jldopen(azimuthal_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_3d)
         end
         GC.gc()
    end

    for current_variable in variables_orig_2d
         composite_addition_2d .= 0.0
         buf_2d_1              .= 0.0
         buf_2d_2              .= 0.0
         azimuthal_average_2d  .= 0.0

         ## Create products with mask
         @info "Creating composite for: " current_variable
         flush(stdout)

        composite_and_radialaverage_in_chunks!(azimuthal_average_2d,composite_addition_2d,buf_2d_1,buf_2d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file2d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = true)
         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_2d)
         end
         jldopen(azimuthal_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_2d)
         end
         GC.gc()
         GC.gc()
         ## Create products without mask
         @info "Creating composite without mask for: " current_variable
         flush(stdout)
         composite_addition_2d .= 0.0
         buf_2d_1              .= 0.0
         buf_2d_2              .= 0.0
         azimuthal_average_2d  .= 0.0

         composite_and_radialaverage_in_chunks!(azimuthal_average_2d,composite_addition_2d,buf_2d_1,buf_2d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file2d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = false)
         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_2d)
         end
         jldopen(azimuthal_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_2d)
         end
         GC.gc()
    end

    for current_variable in variables_diag

         composite_addition_3d .= 0.0
         buf_3d_1              .= 0.0
         buf_3d_2              .= 0.0
         azimuthal_average_3d  .= 0.0
         ## Create products with mask
         @info "Creating composite for: " current_variable
         flush(stdout)

        #try
            composite_and_radialaverage_in_chunks!(azimuthal_average_3d,composite_addition_3d,buf_3d_1,buf_3d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file3d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = true)
        #catch
        #    nothing
        #end
         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_3d)
         end
         jldopen(azimuthal_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_3d)
         end
         GC.gc()
         GC.gc()
         ## Create products without mask
         @info "Creating composite without mask for: " current_variable
         flush(stdout)
         composite_addition_3d .= 0.0
         buf_3d_1              .= 0.0
         buf_3d_2              .= 0.0
          azimuthal_average_3d  .= 0.0
        #try
            composite_and_radialaverage_in_chunks!(azimuthal_average_3d,composite_addition_3d,buf_3d_1,buf_3d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file3d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = false)
        #catch
        #    nothing
        #end

         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_3d)
         end
         jldopen(azimuthal_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_3d)
         end
         GC.gc()
         GC.gc()
    end

    return nothing

end


verbose = true
if verbose && !isdefined(Main,:RvSpectML)  println("# Loading RvSpecML")    end
 using RvSpectML

# USER: You must create a data_paths.jl file in one of the default_paths_to_search listed below. It need only contain one line:
# solar_data_path = "/home/eford/Data/SolarSpectra/NEID_solar/"
target_subdir = "20190918"   # USER: Replace with directory of your choice
 fits_target_str = "Solar"
 output_dir = "examples/output"
 default_paths_to_search = [pwd(),"examples",joinpath(pkgdir(RvSpectML),"examples"), "/gpfs/group/ebf11/default/ebf11/neid_solar"]
 # NOTE: make_manifest does not update its paths_to_search when default_paths_to_search is defined here, so if you change the line above, you must also include "paths_to_search=default_paths_to_search" in the make_manifest() function call below
 pipeline_plan = PipelinePlan()

RvSpectML.Pipeline.reset_all_needs!(pipeline_plan)
 if need_to(pipeline_plan,:read_spectra)
   if verbose println("# Finding what data files are avaliable.")  end
   df_files = make_manifest(target_subdir, NEID )

   if verbose println("# Reading in customized parameters from param.jl.")  end
   eval(code_to_include_param_jl())

   if verbose println("# Reading in ", size(df_files_use,1), " FITS files.")  end
   @time all_spectra = map(NEID.read_solar_data,eachrow(df_files_use))
   dont_need_to!(pipeline_plan,:read_spectra)

   if verbose println("# Applying wavelength corrections.")  end
   NEID.read_drift_corrections!(joinpath(ancilary_solar_data_path,"SolarRV20190918_JD_SciRV_CalRV.txt"), df_files_use)
   NEID.read_barycentric_corrections!(joinpath(ancilary_solar_data_path,"SolarTelescope2019-09-18_inclGravRedshiftAirMassAltitude.csv"), df_files_use)
   NEID.read_differential_extinctions!(joinpath(ancilary_solar_data_path,"20190918_diff_ex_full_fixed.txt"), df_files_use)
   apply_doppler_boost!(all_spectra,df_files_use)
   all_spectra
 end


#
# Protein - TMAO (compare new and old implementations)
#

using ComplexMixtures
using PDBTools
using Plots
using DelimitedFiles
ENV["GKSwstype"] = "nul"

# Here we use the PDBTools package to read the pdb file (from http://github.com/m3g/PDBTools)
atoms = PDBTools.readPDB("../../test/data/NAMD/structure.pdb")

# The solute is a single protein molecule (infinte dilution case). In this case,
# use the option nmols=1
protein = PDBTools.select(atoms,"protein")
solute = ComplexMixtures.Selection(protein,nmols=1)

# The solvent is TMAO, which has 14 atoms. Use the natomspermol to indicate how many
# atoms each molecule has, such that there is no ambiguity on how to split the coordinates
# of the selection into individual molecules.
tmao = PDBTools.select(atoms,"resname TMAO")
solvent = ComplexMixtures.Selection(tmao,natomspermol=14)

# Initialize trajectroy data structure and open input stream
trajectory = ComplexMixtures.Trajectory("../../test/data/NAMD/trajectory.dcd",solute,solvent)

# Input options for the calculation
options = ComplexMixtures.Options(binstep=0.2)

# Run ComplexMixtures calculation, and get the results in the R structure
R = ComplexMixtures.mddf(trajectory,options)

# Save data for future loeading
ComplexMixtures.save(R,"example1.json")

# Save data in human-readable format
ComplexMixtures.write(R,"example1.dat",solute,solvent)

old = readdlm("./old_reference.dat",comments=true,comment_char='#')

plot(layout=(6,1))

sp=1
plot!(ylabel="mddf",subplot=sp)
plot!(old[:,1],old[:,2],subplot=sp,label="old")
plot!(R.d,R.mddf,subplot=sp,label="new - mddf")
plot!(R.d,R.rdf,subplot=sp,label="new - rdf")
plot!(legend=:topright,subplot=sp)

sp=2
plot!(ylabel="KB",subplot=sp)
plot!(old[:,1],old[:,3],subplot=sp,label="old")
plot!(R.d,R.kb,subplot=sp,label="new - mddf")
plot!(R.d,R.kb_rdf,subplot=sp,label="new - rdf")
plot!(legend=:topright,subplot=sp)

sp=3
plot!(ylabel="Count",subplot=sp)
plot!(old[:,1],old[:,4],subplot=sp,label="old")
scatter!(R.d,R.md_count,subplot=sp,label="new")
plot!(old[:,1],old[:,5],subplot=sp,label="old - rand")
scatter!(R.d,R.md_count_random,subplot=sp,label="new -rand")

sp=4
plot!(ylabel="Shell vol", subplot=sp)
plot!(old[:,1],old[:,8],subplot=sp,label="old")
plot!(R.d,R.volume.shell,subplot=sp,label="new")

sp=5
plot!(ylabel="Sum MD", subplot=sp)
plot!(old[:,1],old[:,6],subplot=sp,label="old - md")
scatter!(R.d,R.sum_md_count,subplot=sp,label="new - md")

sp=6
plot!(ylabel="Sum RAND", subplot=sp)
plot!(old[:,1],old[:,7],subplot=sp,label="old - rand")
scatter!(R.d,R.sum_md_count_random,subplot=sp,label="new - rand")
scatter!(R.d,R.sum_rdf_count,subplot=sp,label="new - rdf")
plot!(legend=:topleft,subplot=sp)

#sp=7
#plot!(ylabel="Count",subplot=sp)
#z = zeros(R.nbins)
#@. z = old[:,5]/R.md_count_random
#plot!(R.d,z,subplot=sp,label="old/rand")
#println(z[R.nbins])

plot!(size=(800,1300))
savefig("./example1.pdf")


#%%
using Revise
using GdMn2O5
const Gd = GdMn2O5
using LinearAlgebra
using LaTeXStrings
using ProgressMeter
using FileIO
pyplot()

l  = Ledger(physical_model())
const start = 3.5
const stop  = 5.0
const len   = 120
const Hr    = H_sweep_range(start, stop, len)
const hsweep = generate_Hsweep_states(l, 10, Hr)
const N = 252
const margin = 2.1 * Plots.Measures.mm * 2
const BLACK = RGB(0, 0, 0)
const DARK_GREEN = RGB(0, 0.5, 0)
const DARK_BLUE = RGB(0, 0, 0.5)
const LIGHT_GREEN = RGB(0.1, 1, 0.1)
const LIGHT_BLUE = RGB(0.1, 0.1, 1.0)
const RED = RGB(1, 0, 0)
const PATH_COLOR = RGB(240 / 255, 92 / 255, 7 / 255)
const PATH_COLOR = RGB(1.0, 1.0, 1.0)
function plot_3angle_Pb(l; a1 = 9.80, a2 = 10, a3 = 10.35)
    optimize(l, Spin, L)
    N = 400
    Hr = H_sweep_range(0, 10, N)[1:end - N + 1]
    h_sweep_state0 = generate_Hsweep_states(l, a1, Hr)

    h_sweep_state11 = generate_Hsweep_states(l, a2, Hr)
    h_sweep_state20 = generate_Hsweep_states(l, a3, Hr)
    colors = [fill(BLACK, N-1);fill(DARK_GREEN, N-1); fill(DARK_BLUE, N-1); fill(LIGHT_GREEN, N-1)]
    p3 = plot(Hr[1:2N + 1], Pb(h_sweep_state0[1:2N + 1]),  color = [fill(BLACK, N - 1);fill(LIGHT_GREEN, N + 1)])
    p4 = plot(Hr, Pb(h_sweep_state11), color = colors);
    p5 = plot(Hr[1:2N + 1], Pb(h_sweep_state20[1:2N + 1]), color = [fill(BLACK, N -1);fill(LIGHT_GREEN, N + 1)]);
    return p3, p4, p5
end

function calculate_Etots(hsweep...)
    p = Progress(length(hsweep))
    L1_ϕs = range(0, 2π, length = N)
    L2_ϕs = range(0, 2π, length = N)
    E_tots = [zeros(N, N) for i = 1:length(hsweep)]
    for ih = 1:length(hsweep)
        for e in Gd.@entities_in(hsweep[ih][Spin])
            hsweep[ih][e] = Optimize()
        end
        thread_ls = [deepcopy(hsweep[ih]) for t = 1:Threads.nthreads()]
        Threads.@threads for i1 = 1:N
            for (i2, ϕ2) in enumerate(L2_ϕs)
                ϕ1 = L1_ϕs[i1]
                tl = thread_ls[Threads.threadid()]
                tl[L].data[1] = GdMn2O5.L(ϕ = ϕ1)
                tl[L].data[2] = GdMn2O5.L(ϕ = ϕ2)
                res = optimize!(tl, Spin, Pb)
                E_tots[ih][i1,i2] = calc_etot(tl).e
            end
        end
        next!(p)
    end
    return E_tots
end
function plot_heatmap(h_ledger, hsweep, L1_traj_shift, L2_traj_shift, offset = 1.0, r=0.65)
    ϕ1_r = range(0, 2π, length = N)
    ϕ2_r = range(0, 2π, length = N)
    E_tots =calculate_Etots(h_ledger)[1]
    E_tots = log.(E_tots.-minimum(E_tots).+offset)
    mi = minimum(E_tots)
    ma = maximum(E_tots)
    E_tots = map(x-> (ma - x)/(ma - mi) < r ? NaN : x, E_tots)
    cs = :rainbow
    p = heatmap(ϕ1_r, ϕ2_r, E_tots, colorbar_title = "E", color = cs, aspect_ratio = 1, legend=false)
    plot!(p, map(x->x[L][2].ϕ + L2_traj_shift, hsweep), map(x->x[L][1].ϕ + L1_traj_shift, hsweep), color = PATH_COLOR)
    scatter!(p, [h_ledger[L][2].ϕ + L2_traj_shift], [h_ledger[L][1].ϕ + L1_traj_shift], color = PATH_COLOR, markersize = 4)
    return p
end
#%%
# %%
### ENERGY BARRIERS
# %%
## NEB
Hr_neb = H_sweep_range(3.5, 5, 120)
hsweep_neb = generate_Hsweep_states(l, 10, Hr_neb)
E_barriers, Hs = energy_barriers(hsweep_neb, find_hysteresis(hsweep_neb)...,  10, 10)

start_H, stop_H = find_hysteresis(hsweep_neb)
hsweep_energy = generate_Hsweep_states(l, 10, H_sweep_range(start, stop, 100))
fontsize = 35
#%%
p1 = plot_energy_barriers(E_barriers[1:end - 1], Hs[1:end - 1], 0.03;
                          title = "",
                          xguidefontsize = fontsize,
                          xtickfontsize = fontsize,
                          yguidefontsize = fontsize,
                          ytickfontsize = fontsize,
                          linewidth = 4,
                          colorbar_title = "",
                          colorbar = true,
                          ylabel = "E",
                          bottom_margin = 4margin)
plot(p1,
     plot(H_magnitudes(hsweep_energy), [Etot(hsweep_energy, :E_dipole) Etot(hsweep_energy, :E_Gd_L) Etot(hsweep_energy, :E_Gd) Etot(hsweep_energy, :E_L)],
          label = [L"dipole" L"Gd \leftrightarrow L" L"Gd \leftrightarrow H" L"L \leftrightarrow L + L \leftrightarrow H + L \leftrightarrow easy"],
          yaxis = "E (eV)",
          xaxis = "H (T)",
          legend = :left,
          legendfontsize = 20,
          linewidth = 4,
          xguidefontsize = fontsize,
          yguidefontsize = fontsize,
          xtickfontsize = fontsize,
          ytickfontsize = fontsize),
          layout = (2, 1),
          framestyle = :box,
          size = (800, 1600))
# %%
# savefig(papersdir("fig/neb_energy_contributions.svg"))

# savefig(papersdir("fig/neb.svg"))
#%%
## Heatmaps
#%%
l = Ledger(physical_model())
hsweep1 = generate_Hsweep_states(l, 9.85, Hr)
hsweep2 = generate_Hsweep_states(l, 10, Hr)
hsweep3 = generate_Hsweep_states(l, 10.35, Hr)

p1, p2, p3 = plot_3angle_Pb(l)
p2
p4 = plot_heatmap(hsweep1[62], hsweep1, -π, 2π, 0.005, 0.65)
p5 = plot_heatmap(hsweep2[177], hsweep2, π, 2π, 0.005, 0.45)
p6 = plot_heatmap(hsweep3[62], hsweep3, 0, 2π, 0.005, 0.65)
#%%
fontsize = 18
margin = margin
pyplot()
p1 = plot(p1, ylabel = L"P_b \mathrm{~(}\mu \mathrm{C/cm}^2\mathrm{)}", title = L"\phi_H = 9.85^{\circ}")
p2 = plot(p2, ylabel = nothing, yticks = (-0.2:0.1:0.2, ["" for x in -0.2:0.1:0.2]), title = L"\phi_H = 10^{\circ}")
p3 = plot(p3, ylabel = nothing, yticks = (-0.2:0.1:0.2, ["" for x in -0.2:0.1:0.2]), title = L"\phi_H = 10.35^{\circ}", right_margin=2margin)

ptop = plot(p1,p2,p3, layout=(1,3), size=(1200,400),dpi=100, xlabel = "H (T)",ylims = [-0.2,0.2],legend = false,linewidth = 4, left_margin=2margin,               xtickfontsize=fontsize,
               ytickfontsize=fontsize,
               xguidefontsize=fontsize,
               yguidefontsize=fontsize,
               titlefontsize=fontsize,
               box=true)
# savefig(ptop, plotsdir("ptop.svg"))
p4 = plot(p4, ylabel = L"ϕ_{L_1}", yticks = ([0, π, 2π], ["", L"π", L"2π"]))
p5 = plot(p5, ylabel=nothing, legend=false, yticks = ([0, π, 2π], ["", "", ""]))
p6 = plot(p6,  ylabel=nothing, legend=false, yticks = ([0, π, 2π], ["", "", ""]))
pbottom = plot(p4,p5,p6, layout=(1,3), size=(1200,400),dpi=100, left_margin=2margin, right_margin=0.0margin, ylims = [0,2π],
               xticks = ([0,π, 2π], ["0","π", "2π"]),
               xlims = [0,2π],
               linewidth = 2,
               markersize = 10,
               xlabel = L"ϕ_{L_2}",
               titlefontsize=fontsize,
               xtickfontsize=fontsize,
               ytickfontsize=fontsize,
               xguidefontsize=fontsize,
               yguidefontsize=fontsize,grid=false,
               box=true)
# savefig(pbottom, plotsdir("pbottom.svg"))

for p in ("ptop", "pbottom")
    img = load(plotsdir("$p.png"))

    t = similar(img)
    for i in 1:size(img, 1), j in 1:size(img, 2)
        x = img[i,j]
        # if i > div(size(img, 1), 2)
        t[i, j] = x != RGBA{ColorSchemes.N0f8}(PATH_COLOR.r, PATH_COLOR.g, PATH_COLOR.b) ? RGBA(x.r, x.g, x.b, 1 - x.r) : PATH_COLOR
        # else
            # t[i,j] = x == RGBA{ColorSchemes.N0f8}(1, 1, 1, 1) ? RGBA{ColorSchemes.N0f8}(1, 1, 1, 0.0) : x
        # end
    end
    v1 = findfirst(i->any(x->x.r != 1 || x.b != 1 || x.g != 1, t[i, :]), 1:size(t, 1))
    v2 = findlast(i->any(x->x.r != 1 || x.b != 1 || x.g != 1, t[i, :]), 1:size(t, 1))
    h1 = findfirst(i->any(x->x.r != 1 || x.b != 1 || x.g != 1, t[:, i]), 1:size(t, 2))
    h2 = findlast(i->any(x->x.r != 1 || x.b != 1 || x.g != 1, t[:, i]), 1:size(t, 2))
    img_cropped = t[v1:v2, h1:h2]
    plt = plot(img_cropped, box=:none)
    savefig(papersdir("fig/$p.svg"))
end
v1 = findfirst(i->any(x->x != RGBA{Float32}(1.0, 1.0, 1.0, 0.0), t[i, :]), 1:size(t, 1))
v2 = findlast(i->any(x->x != RGBA{Float32}(1.0, 1.0, 1.0, 0.0), t[i, :]), 1:size(t, 1))
h1 = findfirst(i->any(x->x != RGBA{Float32}(1.0, 1.0, 1.0, 0.0), t[:, i]), 1:size(t, 2))
h2 = findlast(i->any(x->x != RGBA{Float32}(1.0, 1.0, 1.0, 0.0), t[:, i]), 1:size(t, 2))
img_cropped = t[v1:v2, h1:h2]
# save(papersdir("fig/pb_heatmaps.svg"), img_cropped)
#%%


import Revise
import GLRenderer
import PoseComposition
import Rotations
import FileIO

R = Rotations
P = PoseComposition
GL = GLRenderer

obj_path = joinpath(@__DIR__, "035_power_drill/textured_simple.obj")
texture_path = joinpath(@__DIR__, "035_power_drill/texture_map.png")

camera_intrinsics = GL.CameraIntrinsics(
    640, 480,
    1000.0, 1000.0,
    320.0, 240.0,
    0.01, 5.0
)
camera_intrinsics = GL.scale_down_camera(camera_intrinsics, 4)

renderer = GL.setup_renderer(camera_intrinsics, GL.DepthMode())

mesh_data = GL.get_mesh_data_from_obj_file(obj_path)
GL.load_object!(renderer, mesh_data)

depth_image = GL.gl_render(renderer, [1], [P.Pose([0.0, 0.0, 1.0], R.RotXYZ(0.1, 0.4, 0.9))], P.IDENTITY_POSE)

times = [
    let
        r = rand() * 2*π
        t = @elapsed GL.gl_render(renderer, [1], [P.Pose([0.0, 0.0, 1.0], R.RotXYZ(0.1, 0.4, r))], P.IDENTITY_POSE)
        t
    end
    for _ in 1:1000
]
avg_time = sum(times)/length(times)
@show avg_time

FileIO.save("imgs/speed_test_img.png", GL.view_depth_image(depth_image))


using Fermi.Orbitals
function update_amp(T2::Array{T, 4}, newT1::Array{T,2}, newT2::Array{T,4}, foo::Array{T,2}, fov::Array{T,2}, fvv::Array{T,2}, ints::IntegralHelper, alg::A) where { T <: AbstractFloat,
                                                                                A <: CCAlgorithm }

    fill!(newT1, 0.0)
    fill!(newT2, 0.0)

    # Get new amplitudes
    update_T1(T2,newT1,foo,fov,fvv,ints,alg)
    update_T2(T2,newT2,foo,fov,fvv,ints,alg)
end
function BCCD{Ta}(refwfn::RHF, ints::IntegralHelper, newT1::Array{Tb, 2}, newT2::Array{Tc,4}, alg::A; ecT1 = Array{Float64}(undef,0,0), ecT2 = Array{Float64}(undef,0,0,0,0)) where { Ta <: AbstractFloat,
                                                                                Tb <: AbstractFloat,
                                                                                Tc <: AbstractFloat,
                                       A <: CCAlgorithm }

    # Print intro
    Fermi.CoupledCluster.print_header()
    Fermi.CoupledCluster.print_bcc_alg(alg)
    newT1 = convert(Array{Ta},newT1)
    newT2 = convert(Array{Ta},newT2)
    ecT1 = convert(Array{Ta},ecT1)
    ecT2 = convert(Array{Ta},ecT2)

    @output repeat("-",80)*"\n"
    @output "\tBasis: {}\n" ints.bname["primary"]
    @output "\tComputing and Transforming Integrals..."
    tint = @elapsed Fermi.CoupledCluster.compute_integrals(ints,alg)
    @output " done in {} s\n" tint
    for key in keys(ints.cache)
        ints[key] = convert(Array{Ta},ints[key])
    end

    # Process Fock matrix, important for non HF cases
    foo = convert(Array{Ta},ints["FOO"] - Diagonal(ints["FOO"]))
    fvv = convert(Array{Ta},ints["FVV"] - Diagonal(ints["FVV"]))
    fov = convert(Array{Ta},ints["FOV"])

    d = Ta[i - a for i = diag(ints["FOO"]), a = diag(ints["FVV"])]
    D = Ta[i + j - a - b for i = diag(ints["FOO"]), j = diag(ints["FOO"]), a = diag(ints["FVV"]), b = diag(ints["FVV"])]

    # Compute Guess Energy
    oovv = convert(Array{Ta},Fermi.CoupledCluster.compute_oovv(ints,alg))
    Ecc = update_energy(newT1, newT2, fov, ints, alg)
    Eguess = Ecc+refwfn.energy

    #@output "Initial Amplitudes Guess: MP2\n"
    @output "\tGuess Energy:   {:15.10f}\n\n" Ecc
    @output "\tGuess Total Energy:   {:15.10f}\n\n" Ecc+refwfn.energy

    # Start CC iterations

    cc_max_iter = Fermi.CurrentOptions["cc_max_iter"]
    cc_e_conv = Fermi.CurrentOptions["cc_e_conv"]
    cc_max_rms = Fermi.CurrentOptions["cc_max_rms"]
    precision_override = Fermi.CurrentOptions["precision_override"]
    bcc_tol = Fermi.CurrentOptions["bcc_tol"]

    preconv_T1 = Fermi.CurrentOptions["preconv_T1"]
    dp = Fermi.CurrentOptions["cc_damp_ratio"]
    do_diis = Fermi.CurrentOptions["cc_diis"]
    ndiis = Fermi.CurrentOptions["cc_ndiis"]
    diis_start = Fermi.CurrentOptions["diis_start"]
    diis_relax = Fermi.CurrentOptions["cc_diis_relax"]
    prec_selector = Dict{String,Any}(
                                     "half" => Float16,
                                     "single" => Float32,
                                     "double" => Float64
                                    )
    diis_prec = prec_selector[Fermi.CurrentOptions["diis_prec"]]
    do_diis ? DM_T1 = Fermi.DIIS.DIISManager{Ta,diis_prec}(size=ndiis) : nothing
    do_diis ? DM_T2 = Fermi.DIIS.DIISManager{Ta,diis_prec}(size=ndiis) : nothing
    nocc = size(newT2,1)
    nvir = size(newT2,3)

    drop_occ = Fermi.CurrentOptions["drop_occ"]
    drop_vir = Fermi.CurrentOptions["drop_vir"]
    @output "\tDropped Occupied Orbitals →  {:3.0f}\n" Int(Fermi.CurrentOptions["drop_occ"])
    @output "\tDropped Virtual Orbitals  →  {:3.0f}\n\n" Int(Fermi.CurrentOptions["drop_vir"])
    @output "\n"*repeat("-",80)*"\n"
    @output "\tOptions:\n"
    @output "\tPrecision? {}\n" Ta
    @output "\tDIIS? {}\n" do_diis
    @output "\tDIIS Vectors? {}\n" (do_diis ? ndiis : 0)
    @output "\tDIIS Precision? {}\n" diis_prec
    @output "\tPreconverge T1 amplitudes? {}\n" preconv_T1
    @output "\tDamping percentage ? {}\n" dp
    @output "\t\tcc_max_iter →  {:3.0d}\n" Int(cc_max_iter)
    @output "\t\tcc_e_conv   →  {:2.0e}\n" cc_e_conv
    @output "\t\tcc_max_rms  →  {:2.0e}\n" cc_max_rms
    if (cc_e_conv < eps(Ta)) && !(precision_override)
        @output "\t⚠ WARNING⚠   cc_e_conv set to less than ϵ ({}) for precision {}\n" eps(Ta) Ta
        @output "\t             CCSD is unlikely to converge to your standards.\n"
        @output "\t             OVERRIDING set convergence criterion.\n"
        @output "\t             Use @set precision_override true to perform the computation as entered.\n"
        cc_e_conv = eps(Ta)
    end
    if (cc_max_rms < eps(Ta)) && !(precision_override)
        @output "\t⚠ WARNING⚠   cc_max_rms set to less than ϵ ({}) for precision {}\n" eps(Ta) Ta
        @output "\t             CCSD is unlikely to converge to your standards.\n"
        @output "\t             OVERRIDING set convergence criterion.\n"
        @output "\t             Use @set precision_override true to perform the computation as entered.\n"
        cc_max_rms = eps(Ta)
    end

    r1 = 1
    r2 = 1
    dE = 1
    rms = 1
    ite = 1
    T1 = deepcopy(newT1)
    T2 = deepcopy(newT2)

    #dfsz,nocc,nvir = size(Bov)

    @output repeat("-",80)*"\n"
    @output "    Starting CC Iterations\n\n"
    preconv_T1 ? T1_time = 0 : nothing
    dummyT2 = zeros(Ta,size(newT2))
    if preconv_T1
        @output "Preconverging T1 amplitudes\n"
        @output "Taking one T2 step\n"
        @output "{:10s}    {: 15s}    {: 12s}    {:12s}    {:10s}\n" "Iteration" "CC Energy" "ΔE" "Max RMS (T1)" "Time (s)"
        t = @elapsed begin
            update_amp(T2, newT1, newT2, foo, fov, fvv, ints, alg)

            #apply external correction
            apply_ec(newT1,dummyT2,ecT1,ecT2)

            # Apply resolvent
            newT1 ./= d
            newT2 ./= D

            # Compute residues
            r1 = sqrt(sum((newT1 .- T1).^2))/length(T1)
            r2 = sqrt(sum((newT2 .- T2).^2))/length(T2)

            if do_diis
                e1 = (newT1 - T1)
                e2 = (newT2 - T2)
                push!(DM_T1,newT1,e1)
                push!(DM_T2,newT2,e2)
            end

            newT1 .= (1-dp)*newT1 .+ dp*T1
            newT2 .= (1-dp)*newT2 .+ dp*T2
        end
        T1_time += t

        rms = max(r1,r2)
        oldE = Ecc
        Ecc = update_energy(newT1, newT2, fov, ints, alg)
        dE = Ecc - oldE
        @output "    {:<5}    {:<15.10f}    {:<12.10f}    {:<12.10f}    {:<10.5f}\n" "pre" Ecc dE rms t

        while abs(dE) > cc_e_conv || rms > cc_max_rms
            if ite > cc_max_iter
                @output "\n⚠  CC Equations did not converge in {:1.0d} iterations.\n" cc_max_iter
                break
            end
            t = @elapsed begin
                T1 .= newT1
                T2 .= newT2
                update_T1(T2,newT1,foo,fov,fvv,ints,alg)
                apply_ec(newT1,dummyT2,ecT1,ecT2)
                newT1 ./= d
                if do_diis
                    e1 = newT1 - T1
                    push!(DM_T1,newT1,e1)
                    newT1 = Fermi.DIIS.extrapolate(DM_T1)
                end

                # Compute residues
                r1 = sqrt(sum((newT1 .- T1).^2))/length(T1)

                newT1 .= (1-dp)*newT1 .+ dp*T1

                rms = r1
                oldE = Ecc
                Ecc = update_energy(newT1, newT2, fov, ints, alg)
                dE = Ecc - oldE
            end
            T1_time += t
            dE > 0 ? sign = "+" : sign = "-"
            @output "    {:<5.0d}    {:<15.10f}    {}{:>12.10f}    {:<12.10f}    {:<10.5f}\n" ite Ecc sign abs(dE) rms t
            ite += 1
        end
        @output "\nT1 pre-convergence took {}s\n" T1_time
    end

    ints.orbs["BCC"] = deepcopy(ints.orbs["FU"])
    Fermi.Orbitals.activate!(ints.orbs,"BCC")
    ints.orbs.frozencore = drop_occ
    ints.orbs.frozenvir = drop_vir
    C = ints.orbs["[FU]"]
    #display(C)

    dE = 1
    rms = 1
    ite = 1

    do_diis ? DM_T1 = Fermi.DIIS.DIISManager{Ta,diis_prec}(size=ndiis) : nothing
    do_diis ? DM_T2 = Fermi.DIIS.DIISManager{Ta,diis_prec}(size=ndiis) : nothing
    if preconv_T1
        @output "Including T2 update\n"
    end

    main_time = 0
    @output "{:10s}    {: 15s}    {: 12s}    {:12s}    {:10s} {:3s}\n" "Iteration" "CC Energy" "ΔE" "Max RMS" "Time (s)" "BCC"
    relax = 1
    function myexp(A::Array,N=6)
        sum = zeros(size(A))
        for n=0:N
            sum += A^n/factorial(n)
        end
        sum
    end

    mT1 = 0
    while (abs(dE) > cc_e_conv || rms > cc_max_rms) || mT1 > bcc_tol
        if ite > cc_max_iter
            @output "\n⚠  CC Equations did not converge in {:1.0d} iterations.\n" cc_max_iter
            break
        end
        t = @elapsed begin

            bcc = false
            T1 .= newT1
            T2 .= newT2
            d = Ta[i - a for i = diag(ints["FOO"]), a = diag(ints["FVV"])]
            D = Ta[i + j - a - b for i = diag(ints["FOO"]), j = diag(ints["FOO"]), a = diag(ints["FVV"]), b = diag(ints["FVV"])]

            update_T2(T2,newT2,foo,fov,fvv,ints,alg)
            update_T1(T2,newT1,foo,fov,fvv,ints,alg)
            newT2 ./= D
            newT1 ./= d
            oldE = Ecc
            Ecc = update_energy(newT1, newT2, fov, ints, alg)

            oldC = ints.orbs["[FU]"]

            #update_amp(T2, newT1, newT2, foo, fov, fvv, ints, alg)
            apply_ec(newT1,newT2,ecT1,ecT2)
            # Apply resolvent
            if do_diis
                relax -= 1
                e1 = (newT1 - T1)
                e2 = (newT2 - T2)
                push!(DM_T1,newT1,e1)
                push!(DM_T2,newT2,e2)
                if relax == 0
                    newT2 = Fermi.DIIS.extrapolate(DM_T2)
                    newT1 = Fermi.DIIS.extrapolate(DM_T1)
                    relax = diis_relax
                end
            end

            newT1 .= (1-dp)*newT1 .+ dp*T1
            newT2 .= (1-dp)*newT2 .+ dp*T2

            #println(maximum(abs.(T1)))

            # Compute residues
            r1 = sqrt(sum((newT1 .- T1).^2))/length(T1)
            r2 = sqrt(sum((newT2 .- T2).^2))/length(T2)

            mT1 = maximum(abs.(newT1))

            #if there's no significant rotation, lets just skip the integral transform
            if mT1 >= bcc_tol #&& max(r1,r2) < 1E-10#&& ite >= 3
                do_diis ? DM_T1 = Fermi.DIIS.DIISManager{Ta,diis_prec}(size=ndiis) : nothing
                do_diis ? DM_T2 = Fermi.DIIS.DIISManager{Ta,diis_prec}(size=ndiis) : nothing
                bcc = true
                X = Array([ zeros(nocc,nocc) zeros(size(T1)); newT1' zeros(nvir,nvir)])
                U = exp(X - X')
                Fermi.Orbitals.rotate!(ints.orbs,1.0*U,fc=drop_occ,fv=drop_vir)
                delete_integrals(ints,alg)
                Fermi.Orbitals.matchphase(ints.orbs,oldC)
                foo = convert(Array{Ta},ints["FOO"] - Diagonal(ints["FOO"]))
                fvv = convert(Array{Ta},ints["FVV"] - Diagonal(ints["FVV"]))
                fov = convert(Array{Ta},ints["FOV"])
            end

        end
        rms = max(r1,r2)
        dE = Ecc - oldE
        main_time += t
        dE > 0 ? sign = "+" : sign = "-"
        @output "    {:<5.0d}    {:<15.10f}    {}{:>12.10f}    {:<12.10f}    {:<10.5f} {:3s}\n" ite Ecc sign abs(dE) rms t bcc
        ite += 1
    end
    @output "\nMain BCC iterations done in {}s\n" main_time

    # Converged?
    if abs(dE) < cc_e_conv && rms < cc_max_rms
        @output "\n 🍾 Equations Converged!\n"
    end
    @output "\n⇒ Final BCCD Energy:     {:15.10f}\n" Ecc+refwfn.energy
    @output "\n  Final T1 norm:         {:15.10f}\n" sqrt(sum(newT1.^2))
    @output repeat("-",80)*"\n"

    return RCCSD{Ta}(Eguess, Ecc+refwfn.energy, Fermi.MemTensor(newT1), Fermi.MemTensor(newT2))
end

function apply_ec(T1,T2,ecT1,ecT2)
    s1,s2,s3,s4 = size(ecT2)
    r1 = 1:s1
    r2 = 1:s2
    r3 = 1:s3
    r4 = 1:s4
    T1[r1,r3] .+= ecT1
    T2[r1,r2,r3,r4] .+= ecT2
end


#=
This is an implementation of a priority_queue data structure, customized for modern cpu architectures

Special Thanks to Gandalf and other people in the Hoj discord server for helping me out.

The purpose of this implementation:
    priority_queue, despite being one crucial data structure in many applications
    (pathfinding, best-first-search, etc) has not been optimized for modern cpu architectures.
    Commonly used heaps (pairing heap, binary heap) has poor data locality.
    Moreover, this makes heavy use of inlining, making it fast.

This is a minimum viable version.

#Decrease_key is possible, but would require altering the elements/heaps to be aware of their locations.
There are several ways of doing this.
1) Making a handle class which contains the information and the data, then specialize the dispatch in case the heap stores that handle.
2) Making a pointer back to the parent heap.
3) Storing the information of the position somewhere.
4) Make a lookup table storing position for key positions.
2) is likely the most efficient option, but not gonna do it here.

There are also many ways to "defrag", such as by "transposing", merging, etc... but not doing it here either.
Even merging likely not implemented.
=#

struct summary_key{sub_heap_type,Dtype}
    key::sub_heap_type
    data::Dtype
    summary_key(a::sub_heap_type, b::Dtype) where{sub_heap_type, Dtype} = new{sub_heap_type,Dtype}(a,b)
end

#=
syms = [:<, :(==) , :>, :<=, :>=]
for i in syms
    @eval Base.$i(a::summary_key, b::summary_key) = $i(a.data, b.data)
end
No longer needed, the cmp function will now be generated during type compute.
=#
mutable struct Base_layer{Dtype, Base_size, cmp}
    data::Vector{Dtype}
    size::UInt8
    function Base_layer{Dtype, Base_size, cmp}() where {Dtype,Base_size, cmp}
        x = new{Dtype,Base_size, cmp}()
        x.data = Vector{Dtype}(undef,Base_size)
        x.size = 0
        x
    end
end

mutable struct Heap_layer{sub_layer, summary_layer}
    #TODO, fix summary to make it point to appropriate vector.
    summary::summary_layer
    #data::Vector{sub_layer} Fixing the issue.
    #size::UInt
    function Heap_layer{sub_layer, summary_layer}() where {sub_layer, summary_layer}
        x = new{sub_layer, summary_layer}()
        x.summary = summary_layer()
        #x.data = Vector{sub_layer}(undef, max_size(x.summary))
        #x.size = 0
        x
    end
end

function Heap_Heap_type(Dtype::Type, x::Integer, cmp::Function, Base_size::Integer = UInt8(32))
    #TODO: Fix summary key thing.
    if x==0
        return Base_layer{Dtype, Base_size, cmp}
    else
        @inline cmp_key(a,b) = cmp(a.data, b.data)
        sub_heap_type = Heap_Heap_type(Dtype, x-1, cmp, Base_size)
        return Heap_layer{sub_heap_type, Heap_Heap_type(summary_key{sub_heap_type,Dtype}, x-1, cmp_key, Base_size)}
    end
end

@inline function max_size(::Base_layer{Dtype, Base_size, cmp})::UInt where {Dtype, Base_size, cmp}
    return Base_size
end

@inline function max_size(A::Heap_layer{sub_layer, summary_layer})::UInt where {sub_layer, summary_layer}
    return max_size(A.summary)^2
end

@inline function need_splitting(X::Base_layer{Dtype, Base_size, cmp}) where {Dtype, Base_size, cmp}
    return X.size >= Base_size #Maybe replace this with == for improved performance? Considering that X.size only increases by one each time.
end
@inline function push!(X::Base_layer{Dtype, Base_size, cmp}, content::Dtype) where {Dtype, Base_size, cmp}
    #This does not need splitting.
    #Binary search the array to insert.
    #println("Working in progress.")
    #Wait, can I also unroll this too? #Later.

    if X.size == 0
        X.size = 1
        X.data[1] = content
        return
    elseif !cmp(content,X.data[X.size])
        X.size += 1
        X.data[X.size] = content
        return

    end

    mindex::UInt8 = 1
    maxdex::UInt8 = X.size
    while mindex < maxdex
        middex = (mindex+maxdex)>>1
        if cmp(content, X.data[middex])
            maxdex = middex
        else
            mindex = middex+1
        end
    end
    #May replace "end" with "X.size" (careful consideration needed)
    X.size += 1
    X.data[mindex+1:end] =  X.data[mindex:end-1]
    X.data[mindex] = content
    return
end

@inline function length(X::Base_layer)
    return X.size
end

@inline function empty(X::Base_layer)
    return length(X) == 0
end

@inline function pop!(X::Base_layer)
    data = X.data[1]
    #May replace "end" with "X.size"
    X.data[1:end-1] = X.data[2:end] #May need optimiza tion.I don't know how Julia arrays are implemented.
    X.size -= 1
    #X.data[end] = undef #Free the memory.
    return data
    #return the first element.
end

@inline function peek(X::Base_layer)
    return X.data[1]
end

@inline function replace!(X::Base_layer{Dtype, Base_size, cmp}, content::Dtype) where {Dtype, Base_size, cmp}
    #pop one element, then push one element in, like python's heapq.
    Out::Dtype = X.data[1]
    if !cmp(content,X.data[X.size])
        #In last element, needs to be taken care of separately.
        X.data[1:X.size-1] = X.data[2:X.size]
        X.data[X.size] = content
        return Out
    end

    mindex::UInt8 = 2 #The 1st element will be popped out regardless. The element this can come before is the 2nd element onward.
    #No need for special case of size 1: this will be tackled.
    maxdex::UInt8 = X.size
    #The 1st element shouldn't matter because it gets popped anyway?
    while mindex < maxdex
        middex = (mindex+maxdex)>>1
        if cmp(content, X.data[middex])
            maxdex = middex
        else
            mindex = middex+1
        end
    end
    X.data[1:mindex-2] = X.data[2:mindex-1]
    X.data[mindex-1] = content
    #What to do next?
    return Out
end

@inline function split!(X::Base_layer{Dtype, Base_size, cmp}) where {Dtype, Base_size, cmp}
    #create one new base layers.
    #Assumes full X.
    new_one = Base_layer{Dtype, Base_size, cmp}()
    cutoff = div(Base_size, 2)
    new_one.size = Base_size-cutoff
    new_one.data[1:new_one.size] = X.data[cutoff+1:end]
    X.size = cutoff
    return new_one
end

@inline function need_splitting(X::Heap_layer)
    return need_splitting(X.summary)
end

@inline function push!(X::Heap_layer{sub_layer, summary_layer}, content) where {sub_layer, summary_layer}
    #X.size += 1
    if empty(X.summary)
        new_one = sub_layer()
        push!(new_one, content)
        push!(X.summary, summary_key(new_one, content))
        return
        #println("Working in progress.")
    end

    summary_peek = peek(X.summary)
    sub_heap = summary_peek.key
    if need_splitting(sub_heap)
        #Do_something
        new_one = split!(sub_heap)
        push!(sub_heap, content)
        #Get new summary key.

        push!(X.summary, summary_key(new_one, peek(new_one)))
        #println("Working in progress")
    else
        push!(sub_heap, content)
    end
    return
end

@inline function length(X::Heap_layer)
    #Might not even store X, in which case, it is 0.
    #return X.size
end

@inline function empty(X::Heap_layer)
    return empty(X.summary)
end

@inline function pop!(X::Heap_layer)
    #Pop the 1st one, then the summary heap.
    #X.size -= 1
    summary_peek = peek(X.summary)
    sub_heap = summary_peek.key
    out = pop!(sub_heap)
    if empty(sub_heap)
        pop!(X.summary)
        #println("Working in progress")
        #Todo: make sure the data stays in place.
    else
        #println(X.summary)
        replace!(X.summary, summary_key(sub_heap, peek(sub_heap)))
    end
    return out

end

@inline function peek(X::Heap_layer)
    return peek(X.summary).data
end

@inline function replace!(X::Heap_layer, content)
    summary_peek = peek(X.summary)
    sub_heap = summary_peek.key
    out = replace!(sub_heap, content)
    replace!(X.summary, summary_key(sub_heap, peek(sub_heap)))
    return out
    # Replace the pointed heap then replace the summary.
end

@inline function split!(X::Heap_layer{sub_layer, summary_layer}) where {sub_layer, summary_layer}
    new_summary = split!(X.summary)
    new_heap = Heap_layer{sub_layer, summary_layer}()
    new_heap.summary = new_summary
    return new_heap
end

#Trust the compiler.
@inline function debug_print(X)
    print(X, " ")
end
@inline function debug_print(X::summary_key)
    #println("Summary key with key:")
    #print("key = ")
    #debug_print(X.key)
    print("data = ")
    debug_print(X.data)
end
@inline function debug_print(X::Base_layer)
    #println("Base layer:")
    #println(typeof(X))
    println("[")
    for i in 1:X.size
        debug_print(X.data[i])
        print(",")
    end
    println("]")
end

@inline function debug_print(X::Heap_layer)
    #println("Heap layer with summary:")
    print("summary = ")
    debug_print(X.summary)
end

function New_heap(Dtype::Type; layer::Integer=2, cmp::Function = <, base_size::Integer = 64)
    type = Heap_Heap_type(Dtype,layer, cmp, base_size)
    return type()
end


import Base.convert

# How to convert a Dual to a Float
convert{N,T}(Float64, x::Dual{N,T}) = value(x)

"""
    Assemble Jacobian and Residual
"""
function assembleResidualAndJacobian{T}(solver::Solver, sys::System{T})
    startLog(main_perf_log, "assembleResidualAndJacobian()")

    mesh = sys.mesh

    updateGhostedSolution!(solver)

    solution = solver.ghosted_solution

    vars = sys.variables
    n_vars = length(vars)

    # Reset the Residual and Jacobian
    zero!(solver.rhs)
    zero!(solver.mat)

    # Each inner array is of length n_dofs for each var
    var_residuals = Array{Array{T}}(n_vars)
    for i in 1:n_vars
        var_residuals[i] = Array{T}(0)
    end

    # Each inner matrix is i_n_dofs * j_n_dofs
    var_jacobians = Matrix{Matrix{Float64}}((n_vars, n_vars))
    for i in 1:n_vars
        for j in 1:n_vars
            var_jacobians[i,j] = Matrix{Float64}((0,0))
        end
    end

    # Execute the element loop and accumulate Kernel contributions
    for elem in mesh.local_elements
        reinit!(sys, elem, solution)

        # Resize and zero residual vectors and jacobian matrices
        for i_var in sys.variables
            resize!(var_residuals[i_var.id], i_var.n_dofs)
            fill!(var_residuals[i_var.id], 0.)

            for j_var in sys.variables
                # Because for some inexplicable reason Julia doesn't provide a resize!() for Matrices...
                if size(var_jacobians[i_var.id, j_var.id]) != (i_var.n_dofs, j_var.n_dofs)
                    var_jacobians[i_var.id, j_var.id] = Matrix{Float64}((i_var.n_dofs, j_var.n_dofs))
                end

                fill!(var_jacobians[i_var.id, j_var.id], 0.)
            end
        end

        # Get the Residual/Jacobian contributions from all Kernels
        computeResidualAndJacobian!(var_residuals, var_jacobians, vars, sys.kernels)

        # Scatter those entries back out into the Residual and Jacobian
        for i_var in sys.variables
            plusEquals!(solver.rhs, var_residuals[i_var.id], i_var.dofs)

            for j_var in sys.variables
                plusEquals!(solver.mat, var_jacobians[i_var.id, j_var.id], i_var.dofs, j_var.dofs)
            end
        end
    end

    assemble!(solver.mat)
    assemble!(solver.rhs)

    # Now apply BCs

    boundary_info = mesh.boundary_info
    bcs = sys.bcs

    # First: get the set of boundary IDs we need to operate on:
    bids = Set{Int64}()
    for bc in bcs
        union!(bids, bc.bids)
    end

    # Reusable storage for calling residual and jacobian calculations on NodalBCs
    temp_residual = Array{T}(1)
    temp_jacobian = Array{Float64}(n_vars)

    rows_to_zero = []

    entries_in_bc_rows= []

    proc_id = MPI.Comm_rank(MPI.COMM_WORLD)

    # Now go over each nodeset and apply the BCs
    for bid in bids

        # Grab the nodeset for this bid
        node_list = boundary_info.node_list[bid]

        # Iterate over each node and apply the boundary conditions
        for node in node_list
            if node.processor_id == proc_id
                reinit!(sys, node, solution)

                # Apply all of the BCs that should be applied here
                for bc in bcs
                    if bid in bc.bids
                        computeResidualAndJacobian!(temp_residual, temp_jacobian, vars, bc)

                        # First set the residual
                        solver.rhs[bc.u.nodal_dof] = temp_residual[1]

                        # Now - we need to zero out the row in the matrix corresponding to this dof
                        push!(rows_to_zero, bc.u.nodal_dof)

                        # And put this piece in place
                        for v in vars
                            #solver.mat[bc.u.nodal_dof,v.nodal_dof] = temp_jacobian[v.id]
                            push!(entries_in_bc_rows, (bc.u.nodal_dof,v.nodal_dof,temp_jacobian[v.id]))
                        end
                    end
                end
            end
        end
    end

    zeroRows!(solver.mat, rows_to_zero)
    assemble!(solver.mat)

    for entry in entries_in_bc_rows
        solver.mat[entry[1], entry[2]] = entry[3]
    end

    assemble!(solver.mat)
    assemble!(solver.rhs)

    stopLog(main_perf_log, "assembleResidualAndJacobian()")
end


__precompile__(true)

module SLAM

export
    # Types
    EKFSlamState,
    Scene,
    SimData,
    SlamState,
    PFSlamState,
    Vehicle,

    # Functions
    add_features,
    associate,
    compute_association,
    compute_landmark_ellipses,
    ellipse,
    local_to_global,
    get_waypoints,
    inbounds,
    initial_pose,
    laser_lines,
    mpi_to_pi,
    predict,
    predict_observation,
    steer!,
    step_vehicle!,
    update

include("common.jl")
include("ekf.jl")
include("data-association.jl")

end # SLAM


"""
    Calculate firm variables and populate struct with steady state values.
"""
function Save_Firm_Solution(l, l_P, l_G, D_I, D_E, D; par::Pars_v3)
    type = [:BS, :LS, :LU, :BU]

    D_S  = sold_data.(D_I, l_G, type; par)
    D_G  = data_gen.(l_G, type; par)
    prof = firm_profits.(l, l_P, l_G, D_S, D_E, type; par)
    K    = knowledge.(l_P, D, type; par)
    res = Firm_sol()
    @set! res.prof_BS_ss = prof[1] # Labor good production big customer base, sophisticated
    @set! res.prof_LS_ss = prof[2] # Labor good production small customer base, sophisticated
    @set! res.prof_LU_ss = prof[3] # Labor good production small customer base, unsophisticated
    @set! res.prof_BU_ss = prof[4] # Labor good production big customer base, unsophisticated

    @set! res.l_BS_ss = l[1] # Labor good production big customer base, sophisticated
    @set! res.l_LS_ss = l[2] # Labor good production small customer base, sophisticated
    @set! res.l_LU_ss = l[3] # Labor good production small customer base, unsophisticated
    @set! res.l_BU_ss = l[4] # Labor good production big customer base, unsophisticated

    @set! res.l_P_BS_ss = l_P[1] # Labor good production big customer base, sophisticated
    @set! res.l_P_LS_ss = l_P[2] # Labor good production small customer base, sophisticated
    @set! res.l_P_LU_ss = l_P[3] # Labor good production small customer base, unsophisticated
    @set! res.l_P_BU_ss = l_P[4] # Labor good production big customer base, unsophisticated

    @set! res.l_G_BS_ss = l_G[1] # Labor good production big customer base, sophisticated
    @set! res.l_G_LS_ss = l_G[2] # Labor good production small customer base, sophisticated
    @set! res.l_G_LU_ss = l_G[3] # Labor good production small customer base, unsophisticated
    @set! res.l_G_BU_ss = l_G[4] # Labor good production big customer base, unsophisticated

    @set! res.D_S_BS_ss = D_S[1] # data sharing big customer base, sophisticated
    @set! res.D_S_LS_ss = D_S[2] # data sharing small customer base, sophisticated
    @set! res.D_S_LU_ss = D_S[3] # data sharing small customer base, unsophisticated
    @set! res.D_S_BU_ss = D_S[4] # data sharing big customer base, unsophisticated

    @set! res.D_G_BS_ss = D_G[1] # data generated big customer base, sophisticated
    @set! res.D_G_LS_ss = D_G[2] # data generated small customer base, sophisticated
    @set! res.D_G_LU_ss = D_G[3] # data generated small customer base, unsophisticated
    @set! res.D_G_BU_ss = D_G[4] # data generated big customer base, unsophisticated

    @set! res.D_BS_ss = D[1] # data generated big customer base, sophisticated
    @set! res.D_LS_ss = D[2] # data generated small customer base, sophisticated
    @set! res.D_LU_ss = D[3] # data generated small customer base, unsophisticated
    @set! res.D_BU_ss = D[4] # data generated big customer base, unsophisticated

    @set! res.D_E_BS_ss = D_E[1] # data generated big customer base, sophisticated
    @set! res.D_E_LS_ss = D_E[2] # data generated small customer base, sophisticated
    @set! res.D_E_LU_ss = D_E[3] # data generated small customer base, unsophisticated
    @set! res.D_E_BU_ss = D_E[4] # data generated big customer base, unsophisticated

    @set! res.D_I_BS_ss = D_I[1] # data generated big customer base, sophisticated
    @set! res.D_I_LS_ss = D_I[2] # data generated small customer base, sophisticated
    @set! res.D_I_LU_ss = D_I[3] # data generated small customer base, unsophisticated
    @set! res.D_I_BU_ss = D_I[4] # data generated big customer base, unsophisticated

    @set! res.K_BS_ss = K[1] # knowledge employed big customer base, sophisticated
    @set! res.K_LS_ss = K[2] # knowledge employed small customer base, sophisticated
    @set! res.K_LU_ss = K[3] # knowledge employed small customer base, unsophisticated
    @set! res.K_BU_ss = K[4] # knowledge employed big customer base, unsophisticated

    return res
end

function weight_of_type(type::Symbol; par::Pars_v3)
    if type == :LU # little, unsophisticated
        (1 - par.w_B) * (1 - par.w_S)
    elseif type == :LS # little, sophisticated
        (1 - par.w_B) * par.w_S
    elseif type == :BU # big, unsophisticated
        par.w_B * (1 - par.w_S)
    elseif type == :BS # big, sophisticated
        par.w_B * par.w_S
    else
        error("Firm type needs to be either LU, LS, BU or BS")
    end
end

"""
    Data Generation Productivity / Customer Base depending on type.
"""
function A_G_of_type(type::Symbol; par::Pars_v3)
    if type == :LU # little, unsophisticated
        par.A_G_S
    elseif type == :LS # little, sophisticated
        par.A_G_S
    elseif type == :BU # big, unsophisticated
        par.A_G_B
    elseif type == :BS # big, sophisticated
        par.A_G_B
    else
        error("Firm type needs to be either LU, LS, BU or BS")
    end
end

"""
    Data Processing Productivity depending on type.
"""
function A_P_of_type(type::Symbol; par::Pars_v3)
    if type == :LU # little, unsophisticated
        par.A_P_U
    elseif type == :LS # little, sophisticated
        par.A_P_S
    elseif type == :BU # big, unsophisticated
        par.A_P_U
    elseif type == :BS # big, sophisticated
        par.A_P_S
    else
        error("Firm type needs to be either LU, LS, BU or BS")
    end
end

"""
    Squares the minimizer and returns the solution values
"""
function get_solution(sol; par::Pars_v3, type::Symbol)
    if length(sol.minimizer) == 4
        l_P = sol.minimizer[1]^2
        l_G = sol.minimizer[2]^2
        D_E = sol.minimizer[3]^2
        D_I = sol.minimizer[4]^2
        D = bundle(D_I, D_E; par, type)
        l = labor_demand(l_P, D; par, type)
    elseif length(sol.minimizer) == 3
        l_P = sol.minimizer[1]^2
        l_G = sol.minimizer[2]^2
        D_E = sol.minimizer[3]^2
        D_I = data_gen(l_G; par, type)
        D = bundle(D_I, D_E; par, type)
        l = labor_demand(l_P, D; par, type)
    else
        error("I should never be here")
    end

    return l_P, l_G, l, D_E, D_I, D
end

function get_solution_nl(sol; par::Pars_v3, type::Symbol)
    if length(sol.zero) == 4
        l_P = sol.zero[1]^2
        l_G = sol.zero[2]^2
        D_E = sol.zero[3]^2
        D_I = sol.zero[4]^2
        D = bundle(D_I, D_E; par, type)
        l = labor_demand(l_P, D; par, type)
    elseif length(sol.zero) == 3
        l_P = sol.zero[1]^2
        l_G = sol.zero[2]^2
        D_E = sol.zero[3]^2
        D_I = data_gen(l_G; par, type)
        D = bundle(D_I, D_E; par, type)
        l = labor_demand(l_P, D; par, type)
    else
        error("I should never be here")
    end

    return l_P, l_G, l, D_E, D_I, D
end

function get_solution_nl2(sol; par::Pars_v3, type::Symbol)
    if length(sol.zero) == 4
        l_P = sol.zero[1]^2
        l_G = sol.zero[2]^2
        D_E = sol.zero[3]^2
        D_I = sol.zero[4]^2
        D = bundle(D_I, D_E; par, type)
        l = labor_demand(l_P, D; par, type)
        return l_P, l_G, l, D_E, D_I, D
    elseif length(sol.zero) == 3
        l_P = sol.zero[1]^2
        l_G = sol.zero[2]^2
        D_E = sol.zero[3]^2
        D_I = data_gen(l_G; par, type)
        D = bundle(D_I, D_E; par, type)
        l = labor_demand(l_P, D; par, type)
        return l_P, l_G, l, D_E, D_I, D
    elseif length(sol.zero) == 8
        l_P = sol.zero[1]^2
        l_G = sol.zero[2]^2
        D_E = sol.zero[3]^2
        D_I = sol.zero[4]^2
        l_P2 = sol.zero[5]^2
        l_G2 = sol.zero[6]^2
        D_E2 = sol.zero[7]^2
        D_I2 = sol.zero[8]^2
        D = bundle(D_I, D_E; par, type)
        D2 = bundle(D_I2, D_E2; par, type)
        l = labor_demand(l_P, D; par, type, sec_period=:no)
        l2 = labor_demand(l_P2, D2; par, type, sec_period=:yes)
        return l_P, l_G, l, D_E, D_I, D, l_P2, l_G2, l2, D_E2, D_I2, D2
    elseif length(sol.zero) == 6 #! Need to take also into account an intermediate case in which firms switch from data sharing to not sharing... but unlikely to arise in our setting
        l_P = sol.zero[1]^2
        l_G = sol.zero[2]^2
        D_E = sol.zero[3]^2
        l_P2 = sol.zero[4]^2
        l_G2 = sol.zero[5]^2
        D_E2 = sol.zero[6]^2
        D_I = data_gen(l_G; par, type)
        D_I2 = data_gen(l_G2; par, type)
        D = bundle(D_I, D_E; par, type)
        D2 = bundle(D_I2, D_E2; par, type)
        l = labor_demand(l_P, D; par, type, sec_period = :no)
        l2 = labor_demand(l_P2, D2; par, type, sec_period = :yes)
    else
        error("I should never be here")
    end
end

function set_a_Y(compl_prod, σ, ζ)
    if compl_prod == :yes
        α_Y = (σ*ζ - σ + 1) / (σ*ζ) # auxiliary paramter
    elseif compl_prod == :no
        α_Y = 0
    else
        error("Field compl_prod needs to be either :no or :yes")
    end
    return α_Y
end

"""
    Return the correct gamma
"""
function gamma_of_type(par::Pars_v3; type::Symbol)
    if type in [:BS, :LS]
        par.γ_S
    elseif type in [:BU, :LU]
        par.γ_U
    else
        error("type needs to be either :BS, :LS, :BU or :LU")
    end
end

function K_last_of_type(par::Pars_v3; type::Symbol)
    if type == :LU # little, unsophisticated
        par.K_last_LU
    elseif type == :LS # little, sophisticated
        par.K_last_LS
    elseif type == :BU # big, unsophisticated
        par.K_last_BU
    elseif type == :BS # big, sophisticated
        par.K_last_BS
    else
        error("Firm type needs to be either LU, LS, BU or BS")
    end
end

function D_last_of_type(par::Pars_v3; type::Symbol)
    if type == :LU # little, unsophisticated
        par.D_last_LU
    elseif type == :LS # little, sophisticated
        par.D_last_LS
    elseif type == :BU # big, unsophisticated
        par.D_last_BU
    elseif type == :BS # big, sophisticated
        par.D_last_BS
    else
        error("Firm type needs to be either LU, LS, BU or BS")
    end
end


# this is a script to train a basic VAE on the protobuffer pdf data
using GenerativeModels
using GMExtensions
using ConditionalDists
using PdfFraud
using JLD2
using FileIO
using Flux
using CuArrays
using Dates
using ValueHistories
using ProgressMeter
using ArgParse
using MLDataPattern
using DrWatson
quickactivate(dirname(dirname(pathof(PdfFraud)))) # this is so that DrWatson saves all the info

# this is for running the script from REPL - default values
if length(ARGS) == 0
	push!(ARGS, "2")
	push!(ARGS, "2")
	push!(ARGS, "4")
	push!(ARGS, "--test")
end

# parse arguments
s = ArgParseSettings()
@add_arg_table s begin
    "ldimsize"
    	required = true
    	arg_type = Int
    	help = "size of latent dimension"
   "channels"
    	required = true
    	arg_type = Int
    	help = "a list of channel numbers"
    	nargs = '+'
    "--scaling"
    	default = [2]
    	help = "a scalar or a vector of scaling factors"
    	arg_type = Int
    	nargs = '+'
    "--kernelsize"
    	default = [5]
    	help = "a scalar or a vector of kernelsizes"
    	arg_type = Int
    	nargs = '+'
    "--eta"
    	default = 0.001f0
    	arg_type = Float32
    	help = "learning rate"
    "--nepochs"
    	default = 10
    	arg_type = Int
    	help = "number of training epochs"
    "--batchsize"
    	default = 128
    	arg_type = Int
    	help = "batch size"
    "--beta"
    	default = 1.0f0
    	arg_type = Float32
    	help = "elbo β"
    "--savepath"
    	default = ""
    	arg_type = String
    	help = "where should the model be saved"
    "--save-frequency"
    	default = 1
    	arg_type = Int
    	help = "how often (in a number of epochs) should a model be saved"
    "--test"
    	action = :store_true
    	help = "test run"
end
parsed_args = parse_args(ARGS, s)

# extract settings from the parsed arguments
vectorize(x) = (length(x) == 1) ? repeat(x, nconv) : x
ldim = parsed_args["ldimsize"]
channels = parsed_args["channels"]
nconv = length(channels)
scalings = vectorize(parsed_args["scaling"])
kernelsizes = vectorize(parsed_args["kernelsize"])
η = parsed_args["eta"]
nepochs = parsed_args["nepochs"]
batchsize = parsed_args["batchsize"]
β = parsed_args["beta"]
savepath = parsed_args["savepath"]
save_freq = parsed_args["save-frequency"]
test = parsed_args["test"]
if test
	batchsize = 2
	savepath = "test"
end

# prepare the savepath
savepath = (length(savepath) > 0 && savepath[1] == "/") ? savepath :
	joinpath(PdfFraud.experiment_path, savepath, "models")
mkpath(savepath)

# get data
X, labels, sha256, page_nums = PdfFraud.load_pdf_data()
X = X |> gpu;

# construct the convolutional VAE model
(h, w, c) = size(X)[1:3]
densedims = []
vecdim = reduce(*, (h,w,c))
T = Float32

encoder = GMExtensions.conv_encoder((h, w, c), ldim*2, kernelsizes, channels, scalings;
	densedims = densedims)
enc_dist = CMeanVarGaussian{DiagVar}(encoder)

decoder = GMExtensions.conv_decoder((h, w, c), ldim, reverse(kernelsizes),
	reverse(channels), reverse(scalings); densedims = densedims, vec_output = true,
	vec_output_dim = vecdim + 1)
dec_dist = CMeanVarGaussian{ScalarVar}(decoder)

model = VAE(ldim, enc_dist, dec_dist) |> gpu

# create stuff for training
loss(x) = PdfFraud.elbo(model, x, β=β)
opt = ADAM(η)
if !test
	data = [(gpu(Array(x)),) for x in RandomBatches(X, batchsize, ceil(Int, size(X,4)/batchsize))];
	validation_data = gpu(Array(collect(RandomBatches(X, batchsize, 1))[1]));
else
	data = [(X[:,:,:,1:batchsize],) for _ in 1:100];
	validation_data = data[1][1];
end
@info "Training with $(length(data)) batches of $(batchsize) samples."
h = MVHistory()
tstart = "$(now())"
iepoch = 0

# train
function train_epoch(model, history, experiment_args)
	# setup callback
	prog = Progress(length(data))
	global iepoch += 1
	filename = joinpath(savepath, tstart*"_epoch-$(iepoch).bson")
	cb = PdfFraud.history_progress_cb(prog, model, validation_data, history, loss)

	# train
	Flux.train!(loss, params(model), data, opt, cb = cb)

	# save model
	if iepoch%save_freq == 0
		PdfFraud.save_model(filename, model, history, experiment_args)
	end
end
Flux.@epochs nepochs train_epoch(model, h, parsed_args)


# Base Quality
# ============
#
# A representation of positions-specific integer quality scores, as in FASTQ.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/Bio.jl/blob/master/LICENSE.md

immutable QualityEncoding
    # valid ASCII character range
    ascii::NTuple{2,UInt8}

    # valid base quality range
    qual::NTuple{2,Int8}

    function QualityEncoding(ascii::Range{Char}, qual::Range{Int})
        if length(ascii) != length(qual)
            throw(ArgumentError("the ranges of ASCII and base quality don't match"))
        end
        return new((first(ascii), last(ascii)), (first(qual), last(qual)))
    end
end

# Return the offset of converting a base quality to an ASCII code (quality + offset).
function ascii_offset(encoding::QualityEncoding)
    return encoding.ascii[1] - encoding.qual[1]
end

"Sanger (Phred+33) quality score encoding"
const SANGER_QUAL_ENCODING     = QualityEncoding('!':'~', 0:93)

"Solexa (Solexa+64) quality score encoding"
const SOLEXA_QUAL_ENCODING     = QualityEncoding(';':'~', -5:62)

"Illumina 1.3 (Phred+64) quality score encoding"
const ILLUMINA13_QUAL_ENCODING = QualityEncoding('@':'~', 0:62)

"Illumina 1.5 (Phred+64) quality score encoding"
const ILLUMINA15_QUAL_ENCODING = QualityEncoding('B':'~', 2:62)

"Illumina 1.8 (Phred+33) quality score encoding"
const ILLUMINA18_QUAL_ENCODING = QualityEncoding('!':'~', 0:93)

# Check quality encoding of `input[start:stop]`.
function check_quality_string(encoding, input, start, stop)
    ascii_lo, ascii_hi = encoding.ascii
    for i in start:stop
        @inbounds a = input[i]
        if !(ascii_lo ≤ a ≤ ascii_hi)
            error("base quality '$(Char(a))' is out of range")
        end
    end
end

# Decode an ASCII string `input[start:stop]` into base quality scores `output`.
function decode_quality_string!(encoding, input, output, start, stop)
    resize!(output, stop - start + 1)
    offset = ascii_offset(encoding)
    for i in 1:(stop - start + 1)
        @inbounds output[i] = input[start + i - 1] - offset
    end
    return output
end

# Encode base quality scores `input` into a ASCII string `output`.
function encode_quality_string!(encoding, input, output, start, stop)
    resize!(output, stop - start + 1)
    offset = ascii_offset(encoding)
    for i in 1:(stop - start + 1)
        @inbounds output[i] = input[start + i - 1] + offset
    end
    return output
end


using RGBDSegmentation
using RGBDSegmentation.JCSA
using RGBDSegmentation.CDNGraph
using RGBDSegmentation.GCF
using ImageSegmentation
using ImageSegmentationEvaluation
using LibPQ
using JSON
using DataStreams
using JLD
using Images
using FileIO
using ArgParse
using YAML
using Statistics
using DataFrames
using Distributed
typedict(x) = Dict(fn=>getfield(x, fn)
                   for fn ∈ fieldnames(typeof(x)))

include("store.jl")

function supervised_metrics(img, segmented_image::ImageSegmentation.SegmentedImage)
    metrics = Dict(
        "FBoundary" => FBoundary(),
        "Precision" => Precision(),
        "FMeasure" => FMeasure(),
        "SegmentationCovering" => SegmentationCovering(),
        "VariationOfInformation" => VariationOfInformation(),
        "RandIndex" => RandIndex(),
        "FMeasureRegions" => FMeasureRegions(),
        "PRObjectsAndParts" => PRObjectsAndParts(),
        "BoundaryDisplacementError" =>  BoundaryDisplacementError()
    )
    result = Dict()
    for metric_name in sort(collect(keys(metrics)))
        @info("Computing $metric_name")
        result[metric_name] = ImageSegmentationEvaluation.evaluate(
            metrics[metric_name],
            segmented_image,
            img.ground_truth)
    end
    return result
end
function unsupervised_metrics(img, segmented_image::ImageSegmentation.SegmentedImage)
    default_params = Dict("ECW" => Dict(
                                   :threshold=>0.5),
                          )
    params = default_params

    RGB_img = colors(img.image)
    segmented_image_RGB = ImageSegmentation.SegmentedImage(
        segmented_image.image_indexmap,
        segmented_image.segment_labels,
        Dict(k => RGB(rm.color...) for (k, rm) in segmented_image.segment_means),
        segmented_image.segment_pixel_count)
    D   = distances(img.image)
    N   = normals(img.image)

    metrics_color = Dict(
        "ECW" => ECW(;params["ECW"]...),
        "Zeboudj" => Zeboudj(radius=5),
        "ValuesEntropy" => ValuesEntropy(),
        "LiuYangF" =>  LiuYangF(),
        "FPrime" => FPrime(),
        #"ErdemMethod" => ErdemMethod(5, 5),
        "Q" => Q())
    metrics_color_distance = Dict(
        "FRCRGBD" => FRCRGBD())

    result = Dict()
    for metric_name in sort(collect(keys(metrics_color)))
        @info("Computing $metric_name")
        result[metric_name] = ImageSegmentationEvaluation.evaluate(
            metrics_color[metric_name],
            RGB_img,
            segmented_image_RGB)
    end

    for metric_name in sort(collect(keys(metrics_color_distance)))
        @info("Computing $metric_name")
        result[metric_name] = ImageSegmentationEvaluation.evaluate(
            metrics_color_distance[metric_name],
            RGB_img,
            D[3, :, :],
            segmented_image)
    end
    return result
end

function dict_to_params(params::Dict)
    return Dict(Symbol(a.first)=>a.second for a in params)
end
function segment_image(dataset, i, algorithms, config)
    K = config["number_of_segments"]
    data_path = config["data_path"]
    image_scale = config["image_scale"]

    image = nothing
    for algo_name in keys(algorithms)
        algorithm_cfg = algorithms[algo_name]
        algorithm = get_or_insert(
            Algorithm(
                name=algo_name,
                params=typedict(algorithm_cfg)),
            conn)
        base_filename = filename(dataset, i)
        file_path = joinpath(data_path, algo_name)
        mkpath(file_path)
        file_path = joinpath(file_path, base_filename)
        segmented_image = SegmentedImage(
            dataset=name(dataset),
            image=i,
            file_path=file_path,
            algorithm=algorithm)
        if exists(segmented_image, conn, exclude=[:metrics])
            @info("Segmentation already stored")
            continue
        end
        if !isfile(string(segmented_image.file_path, ".jld"))
            if image == nothing
                @info "Loading image $i"
                image = dataset[i]
                @info("Resizing image to 50%")
                image = resize(image, image_scale)
            end
            @info("Segmenting image $i with algorithm $algorithm")
            time = @elapsed (
                segments = clusterize(
                    algorithm_cfg,
                    image.image,
                    K))
            segmented_image.metrics["elapsed"] = time
            save(segmented_image,
                 image,
                 segments,
                 conn)
        else
            @info("Inserting already segmented image")
            insert(segmented_image, conn)
        end
    end

end
function segment(config)
    dataset = NYUDataset()
    algorithms = Dict(
        "JCSA"=> (
            JCSA.Config(;
                dict_to_params(config["JCSA"])...)),
        "GCF" => GCF.Config(;
            dict_to_params(config["GCF"])...),
        "CDNGraph" =>  CDNGraph.Config(;
            dict_to_params(config["CDNGraph"])...)
    )
    asyncmap(i->segment_image(dataset, i, algorithms, config),
             1:length(dataset),
             ntasks=4)
end
function evaluate_image(dataset, i, algorithms, config)
    image = nothing
    segmented_image = nothing
    K = config["number_of_segments"]
    data_path = config["data_path"]
    image_scale = config["image_scale"]

    for algo_name in keys(algorithms)
        algorithm_cfg = algorithms[algo_name]
        algorithm = get(
            Algorithm(
                name=algo_name,
                params=typedict(algorithm_cfg)),
            conn)

        base_filename = filename(dataset, i)
        file_path = joinpath(data_path, algo_name)
        mkpath(file_path)
        file_path = joinpath(file_path, base_filename)
        segmented_image_db = get(
            SegmentedImage(
                dataset=name(dataset),
                image=i,
                file_path=file_path,
                algorithm=algorithm),
                conn,
            exclude=[:metrics])
        if (!isempty(segmented_image_db.metrics))
            continue
        end
        try

            if image == nothing
                @info "Loading image $i"
                image = dataset[i]
                @info("Resizing image to 50%")
                image = resize(image, image_scale)
            end
            segmented_image = load(segmented_image_db)
            results = unsupervised_metrics(image, segmented_image)
            merge!(results, supervised_metrics(image, segmented_image))
            segmented_image_db.metrics = results
            save(segmented_image_db, conn)
        catch e
            @warn(e)
        end

    end
    GC.gc()
end
function evaluate(config)
    dataset = NYUDataset()
    algorithms = Dict(
        "JCSA"=> (
            JCSA.Config(;
                dict_to_params(config["JCSA"])...)),
        #"GCF" => GCF.Config(;
        #    dict_to_params(config["GCF"])...),
        "CDNGraph" =>  CDNGraph.Config(;
            dict_to_params(config["CDNGraph"])...)
    )
    #for i=1:length(dataset)
    #    evaluate_image(dataset, i, algorithms, config)
    #end
    asyncmap(i->evaluate_image(dataset, i, algorithms, config),
             1:length(dataset),
             ntasks=4)

end
function parse_commandline()
    s = ArgParseSettings()
    add_arg_table(
        s,
        "--config",
        Dict(:help=>"an option with an argument",
             :required=>true),
        "segment",
        Dict(:help =>"Segment images",
             :action=>:command),
        "evaluate",
        Dict(:help=>"Evaluate segmented image",
             :action=>:command),
        "summarize",
        Dict(:help=>"Summarize metrics",
            :action=>:command)
    )
    return parse_args(ARGS, s)
end
function summarize(config)
    dataset = NYUDataset()
    algorithms = Dict(
        "JCSA"=> (
            JCSA.Config(;
                dict_to_params(config["JCSA"])...)),
        #"GCF" => GCF.Config(;
        #    dict_to_params(config["GCF"])...),
        "CDNGraph" =>  CDNGraph.Config(;
            dict_to_params(config["CDNGraph"])...)
    )
    results = Dict()
    for i=1:length(dataset)
        image = nothing
        segmented_image = nothing
        K = config["number_of_segments"]
        data_path = config["data_path"]
        image_scale = config["image_scale"]
        for algo_name in keys(algorithms)
            algorithm_cfg = algorithms[algo_name]
            algorithm = get(
                Algorithm(
                    name=algo_name,
                    params=typedict(algorithm_cfg)),
                conn)
            base_filename = filename(dataset, i)
            file_path = joinpath(data_path, algo_name)
            mkpath(file_path)
            filepath = joinpath(file_path, base_filename)
            segmented_image_db = get(
                SegmentedImage(
                    dataset=name(dataset),
                    image=i,
                    algorithm=algorithm),
                conn,
                exclude=[:metrics, :file_path])

            if !(haskey(results, algo_name))
                results[algo_name] = Dict()
            end
            for metric in keys(segmented_image_db.metrics)
                if !(haskey(results[algo_name], metric))
                    results[algo_name][metric] = []
                end
                if segmented_image_db.metrics[metric] != nothing
                    push!(results[algo_name][metric], first(segmented_image_db.metrics[metric]))
                end
            end
        end
    end

    algorithm_names = sort(collect(keys(results)))
    for algo_name in algorithm_names
        for metric in sort(collect(keys(results[algo_name])))
            metrics = results[algo_name][metric]
            results[algo_name][metric] = mean(metrics)
        end
    end
    println(results)
end
function main()
    args = parse_commandline()
    config = YAML.load(open(args["config"]))
    if args["%COMMAND%"] == "segment"
        segment(config)
    elseif args["%COMMAND%"] == "evaluate"
        evaluate(config)
    elseif args["%COMMAND%"] == "summarize"
        summarize(config)
    end
end
main()