text:
abstract type NoiseModelFlag end
struct ErdosRenyiNoise  <: NoiseModelFlag end
struct DuplicationNoise <: NoiseModelFlag end

#=------------------------------------------------------------------------------
                        Local File Experiments Routines
------------------------------------------------------------------------------=#

#TODO: update to use kwargs properly
function self_alignment(ssten_files::Array{String,1};method="LambdaTAME")

    Best_alignment_ratio = Array{Float64}(undef,length(ssten_files))

    for f in ssten_files
        A = load(f;type="COOTen")

        #permute tensor to align against
        p = sort(1:A.cubical_dimension,by=i->rand())
        A_permuted = COOTen(A.indices,A.vals,A.cubical_dimension)
        permute_tensor!(A_permuted,p)

        best_score, _, _ = align_tensors(A,A_permuted;method=method)

        println("$f aligned with $best_score")
    end

end

function apply_Metis_permutation!(A::SparseMatrixCSC,k::Int=100)

	n = size(A,1)
	metis_partition = Metis.Metis.partition(A, 100)
	p = sort(1:n, by= i-> metis_partition[i])

	permute!(A,p,p)
end

#TODO: put in kwargs
function pairwise_alignment(dir)

    #align all .ssten files
    ssten_files = sort([f for f in readdir(dir) if occursin(".ssten",f)])

    Best_alignment_ratio = Array{Float64}(undef,length(ssten_files),length(ssten_files))

    for i in 1:length(ssten_files)
        for j in i+1:length(ssten_files)
            Best_alignment_ratio[i,j] = align_from_files(dir*ssten_files[i],dir*ssten_files[j])
            Best_alignment_ratio[i,j] = Best_alignment_ratio[j,i]
        end
    end

    return ssten_files, Best_alignment_ratio
end

#=
function distributed_pairwise_alignment(dir::String;kwargs...)

    #align all .ssten files
    ssten_files = sort([f for f in readdir(dir) if occursin(".ssten",f)])
    distributed_pairwise_alignment(ssten_files,dir;kwargs...)

end
=#

"""------------------------------------------------------------------------------
  This function runs alignemnt experiments in parallel using the @spawn macro of
  Distributed.jl. This function should be run in the src file, as the @everywhere
  macro will only work properly there. Any additional kwargs are passed to
  'align_matrices' or 'align_tensors' depending on the file extensions.

  Inputs
  ------
  * files - (Array{String,1}):
    List of the file names. Files may be either all SMAT or all STEN files.
    'align_tensors' or 'align_matrices' will be called depending on the extension.
  * dirpath - (String):
    The path to where the files exist.
  * method - (String):
    What type of alignment method to run, options are 'LambdaTAME', 'TAME',
    'LowRankTAME', 'EigenAlign', 'LowRankEigenAlign', 'Degree', or 'Random'.
  * profile - (Bool):
    Whether or not to profile the algorithm. Setting this flag to true will
    record different things depending on the algorithm, but will always return
    runtime information.
Output
------
  * 'exp_results' - (type omitted):
    The final array of results for the series of experiments, the size of each
    entry will differ depending on whether or not profile was True. At the very
    least, each entry will contain the the pair of file names aligned, the
    number of triangles matched in the alignemnt, and the maximum number of
    triangles possible to align (the minimum number of triangles in either
    network).

------------------------------------------------------------------------------"""
function distributed_pairwise_ssten_alignment(files::Array{String,1},dirpath::String;
                                        method::AlignmentMethod=ΛTAME_M(),profile=false,kwargs...)

    futures = []
	exp_results = []

    for i in 1:length(files)
        for j in i+1:length(files)
            future = @spawn align_tensors(dirpath*files[i],dirpath*files[j];profile=profile,method=method,kwargs...)
            push!(futures,((i,j),future))
        end
    end

    for ((i,j), future) in futures

        if method === ΛTAME_M() || method === LowRankTAME_M() || method === TAME_M()
            alignmentOutput = fetch(future)
            matched_tris = alignmentOutput.matchScore
            best_matching = alignmentOutput.matching
            A_tens, B_tens = alignmentOutput.motifCounts
            if profile
                profiling_results = alignmentOutput.profile
			end
		end

        data_to_save = Array{Any,1}(undef,0)
        push!(data_to_save,files[i])
        push!(data_to_save,files[j])
        push!(data_to_save,matched_tris)
        push!(data_to_save,best_matching)
        push!(data_to_save,A_tens)
        push!(data_to_save,B_tens)

        if profile
            push!(data_to_save,profiling_results)
        end

        push!(exp_results,data_to_save)
    end

    return exp_results
end

function distributed_pairwise_smat_alignment(files::Array{String,1},dirpath::String;
                                             method::AlignmentMethod=ΛTAME_M(),profile=false,kwargs...)

    futures = []
	exp_results = []

    for i in 1:length(files)
        for j in i+1:length(files)
            future = @spawn align_matrices(dirpath*files[i],dirpath*files[j];profile,method,kwargs...)
            push!(futures,((i,j),future))
        end
    end

    for ((i,j), future) in futures

        if method === ΛTAME_M() || method === LowRankTAME_M()
            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()
                alignmentOutput = fetch(future)
            else
                alignmentOutput, postProcessingOutput = fetch(future)
            end
            matched_tris = alignmentOutput.matchScore
            max_tris = min(alignmentOutput.motifCounts...)
            best_matching = alignmentOutput.matching
            A_tens, B_tens = alignmentOutput.motifCounts
            if profile
                profiling_results = alignmentOutput.profile
			end
        elseif method === TAME_M()
            alignmentOutput = fetch(future)
            matched_tris = alignmentOutput.matchScore
            max_tris = min(alignmentOutput.motifCounts...)
            best_matching = alignmentOutput.matching
            A_tens, B_tens = alignmentOutput.motifCounts
            if profile
                profiling_results = alignmentOutput.profile
			end
        elseif method === ΛTAME_MultiMotif_M() || method === LowRankTAME_MultiMotif_M() || method === TAME_MultiMotif_M()
            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()
                A_motifDistribution, B_motifDistribution,alignmentOutput = fetch(future)
            else
                A_motifDistribution, B_motifDistribution,alignmentOutput, postProcessingOutput = fetch(future)
            end

            if profile
                profiling_results = alignmentOutput.profile
			end
        elseif method === LowRankEigenAlign_M()

            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()
                alignmentOutput = fetch(future)
            else
                alignmentOutput,postProcessingOutput = fetch(future)
            end

            matched_tris = alignmentOutput.matchScore
            A_tens,B_tens = alignmentOutput.motifCounts
            max_tris = min(alignmentOutput.motifCounts...)
            best_matching = alignmentOutput.matching

            if profile
                profiling_results = alignmentOutput.profile
            end

        elseif method === EigenAlign_M() || method === Degree_M() || method === Random_M()
            A_tens, B_tens, matched_tris, best_matching, profiling_results = fetch(future)
            max_tris = min(A_tens, B_tens)
            profile = true
		end

        data_to_save = Array{Any,1}(undef,0)
        push!(data_to_save,files[i])
        push!(data_to_save,files[j])

        if method === ΛTAME_MultiMotif_M() || method === LowRankTAME_MultiMotif_M() || method === TAME_MultiMotif_M()
            push!(data_to_save,alignmentOutput.matchScore)
            push!(data_to_save,alignmentOutput.motifCounts[1])
            push!(data_to_save,alignmentOutput.motifCounts[2])
            push!(data_to_save,A_motifDistribution)
            push!(data_to_save,B_motifDistribution)
        else
            push!(data_to_save, matched_tris)
            push!(data_to_save, A_tens)
            push!(data_to_save, B_tens)
        end

        if profile
            push!(data_to_save,profiling_results)
        end

        if haskey(kwargs,:postProcessing)
            if typeof(kwargs[:postProcessing]) <: KlauAlgo
                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_tris_matched)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.setup_rt)
                    push!(data_to_save,postProcessingOutput.klau_rt)
                end
                push!(data_to_save,postProcessingOutput.L_sparsity)
                push!(data_to_save,postProcessingOutput.f_status)
            elseif typeof(kwargs[:postProcessing]) <: SuccessiveKlauAlgo

                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_tris_matched)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.profiling)
                end
            elseif typeof(kwargs[:postProcessing]) <: LocalSearch

                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.localsearch_edges_matched)
                push!(data_to_save,postProcessingOutput.localsearch_tris_matched)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.profiling)
                    push!(data_to_save,postProcessingOutput.full_runtime)
                end

            end
        end

        push!(exp_results,data_to_save)
    end

    return exp_results
end

function distributed_pairwise_alignment(files::Array{String,1},args...;kwargs...)
     if all( [f[end-5:end] == ".ssten" for f in files]) #tensors
        subkwargs = Dict([(k,v) for (k,v) in kwargs if k != :postProcessing]) #ssten can't support post processing
        return distributed_pairwise_ssten_alignment(files,args...;subkwargs...)
     elseif all( [f[end-4:end] == ".smat" for f in files]) #matrices
         return distributed_pairwise_smat_alignment(files,args...;kwargs...)
     else
         throw(ArgumentError("all files must be the same file type, either all '.ssten' or '.smat'."))
     end
 end

function self_alignment(dir::String;kwargs...)

    #align all .ssten files
    ssten_files = sort([dir*"/"*f for f in readdir(dir) if occursin(".ssten",f)])
    return self_alignment(ssten_files;kwargs...)

end

#=------------------------------------------------------------------------------
                      Generated Graph Experiments Routines
------------------------------------------------------------------------------=#
"""------------------------------------------------------------------------------
  This function runs randomly generated experiments in parallel using the @spawn
  macro of Distributed.jl. This function should be run in the src file, as the
  @everywhere macro will only work properly there. Any additional kwargs are
  passed to 'random_graph_exp'.

  Inputs
  ------
  * 'trial_count' - (Int):
    The number of trials to run for each of the experiments conducted.
  * 'seed_exps' - (Bool):
    Indicates whether or not to seed the experiments. Seeds are generated by
    seeding Random.seed!(0), and then generating an array of UInts using the
    number of 'p_remove' and 'n_sizes' variables, and the 'trial_count'. Changing
    these may lead to different seeds.
  * 'n_sizes' - (Array{Int,1}):
    The different sizes of graphs to generate experiments for.
  * 'p_remove' - (Array{Float,1}):
    The different probability of edge removals to test in the ER noise models.
  * profile - (Bool):
    Whether or not to profile the algorithm. Setting this flag to true will
    record different things depending on the algorithm, but will always return
    runtime information.
Output
------
  * 'exp_results' - (type omitted):
    The final array of results for the series of experiments, the size of each
    entry will differ depending on whether or not profile was True. At the very
    least, the each entry will contain the seed used, the 'p_remove' value, the
    number of vertices in the graph, the ground truth accuracy of the alignment,
    the degree weighted accuracy of the alignment, the number of matched
    triangles, the number of triangles each graph, and the maximum number of
    triangles possible to align. Anything additional returned by the functions
    with the profile flag on, is appended at the end.
------------------------------------------------------------------------------"""
function distributed_random_trials(trial_count::Int,noise_model::ErdosRenyiNoise,seed_exps::Bool=false
                                   ;method::AlignmentMethod=ΛTAME_M(),graph::RandomGraphType=ErdosRenyi(),
                                   n_sizes = [100, 500, 1000, 2000,5000],p_remove = [.01,.05],
                                   profile=false,kwargs...)

	if seed_exps
		Random.seed!(0)
		seeds = rand(UInt64, length(p_remove),length(n_sizes), trial_count)
	end

	seed = nothing
	results = []
    p_index = 1

    futures = []

    for p in p_remove

        n_index = 1

        for n in n_sizes

            for trial = 1:trial_count

                if seed_exps
                    seed = seeds[p_index,n_index,trial]
                end
                future = @spawn random_graph_exp(n,p,graph; profile=profile,seed=seed,
                                                 method=method,noise_model=noise_model,
                                                 kwargs...)
                push!(futures,(seed,p,n,future))

            end

            n_index += 1
        end
        p_index += 1
    end

    #TODO: fix parsing of returned functions
    for (seed,p,n,future) in futures

        if method === ΛTAME_M() || method === LowRankTAME_M()

            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()

                d_A, d_B, perm, alignmentOutput = fetch(future)
            else
                d_A, d_B, perm, (alignmentOutput, postProcessingOutput) = fetch(future)
            end

            matched_tris = alignmentOutput.matchScore
            max_tris = min(alignmentOutput.motifCounts...)
            A_tris, B_tris = alignmentOutput.motifCounts
            best_matching = alignmentOutput.matching
            if profile
                profiling_results = alignmentOutput.profile
            end
        elseif method === TAME_M()
            d_A, d_B, perm, alignmentOutput = fetch(future)
            matched_tris = alignmentOutput.matchScore
            A_tris, B_tris = alignmentOutput.motifCounts
            max_tris = min(alignmentOutput.motifCounts...)
            best_matching = alignmentOutput.matching
            if profile
                profiling_results = alignmentOutput.profile
            end
        elseif method === ΛTAME_MultiMotif_M() || method === LowRankTAME_MultiMotif_M() || method === TAME_MultiMotif_M()
            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()
                d_A, d_B, perm, (A_motifDistribution, B_motifDistribution, alignmentOutput) = fetch(future)
            else
                d_A, d_B, perm, (A_motifDistribution, B_motifDistribution, alignmentOutput, postProcessingOutput) = fetch(future)
            end
            if profile
                profiling_results = alignmentOutput.profile
            end
            A_motifCounts, B_motifCounts = alignmentOutput.motifCounts
            best_matching_score = alignmentOutput.matchScore
            best_matching = alignmentOutput.matching
            A_tris = -1
            B_tris = -1
        elseif method === LowRankEigenAlign_M()

            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()
                d_A, d_B, perm, alignmentOutput = fetch(future)
            else
                d_A, d_B, perm, (alignmentOutput,postProcessingOutput) = fetch(future)
            end

            matched_tris = alignmentOutput.matchScore
            A_tris, B_tris = alignmentOutput.motifCounts
            max_tris = min(alignmentOutput.motifCounts...)
            best_matching = alignmentOutput.matching

            if profile
                profiling_results = alignmentOutput.profile
            end

        elseif method === EigenAlign_M() || method == Degree_M() || method === Random_M() || method === LowRankEigenAlignOnlyEdges_M()
            d_A, d_B, perm, (A_tris, B_tris, matched_tris, best_matching, _) = fetch(future)
            max_tris = minimum((A_tris,B_tris))
        end

        if length(alignmentOutput.matching) > 0
            accuracy = sum([1 for (i,j) in alignmentOutput.matching if get(perm,j,-1) == i])/n

            D_A = sum(d_A)
            D_B = sum(d_B)
            degree_weighted_accuracy = sum([(get(perm,j,-1) == i) ? ((d_A[i] + d_B[j])/(D_A+D_B)) : 0.0 for (i,j) in best_matching])
        else
            accuracy = -1.0
            degree_weighted_accuracy = -1.0
        end

        data_to_save = Array{Any,1}(undef,0)

        push!(data_to_save,seed)
        push!(data_to_save,p)
        push!(data_to_save,n)
        push!(data_to_save,accuracy)
        push!(data_to_save, degree_weighted_accuracy)
        if method === ΛTAME_MultiMotif_M() || method === LowRankTAME_MultiMotif_M() || method === TAME_MultiMotif_M()
            push!(data_to_save,best_matching_score)
            push!(data_to_save,A_motifCounts)
            push!(data_to_save,B_motifCounts)
            push!(data_to_save,A_motifDistribution)
            push!(data_to_save,B_motifDistribution)
        else
            push!(data_to_save, matched_tris)
            push!(data_to_save, A_tris)
            push!(data_to_save, B_tris)
            push!(data_to_save, max_tris)
        end

        if profile
            push!(data_to_save,profiling_results)
        end

        if haskey(kwargs,:postProcessing)
            if typeof(kwargs[:postProcessing]) <: KlauAlgo
                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_tris_matched)
                push!(data_to_save,sum([1 for (i,j) in postProcessingOutput.matching if get(perm,j,-1) == i])/n)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.setup_rt)
                    push!(data_to_save,postProcessingOutput.klau_rt)
                end
                push!(data_to_save,postProcessingOutput.L_sparsity)
                push!(data_to_save,postProcessingOutput.f_status)
            elseif typeof(kwargs[:postProcessing]) <: SuccessiveKlauAlgo

                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_tris_matched)
                push!(data_to_save,sum([1 for (i,j) in postProcessingOutput.matching if get(perm,j,-1) == i])/n)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.profiling)
                end
            elseif typeof(kwargs[:postProcessing]) <: LocalSearch

                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.localsearch_edges_matched)
                push!(data_to_save,postProcessingOutput.localsearch_tris_matched)
                push!(data_to_save,sum([1 for (i,j) in postProcessingOutput.matching if get(perm,j,-1) == i])/n)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.profiling)
                    push!(data_to_save,postProcessingOutput.full_runtime)
                end

            end
        end

        push!(results,data_to_save)

    end

    return results

end

function distributed_random_trials(trial_count::Int,noise_model::DuplicationNoise,seed_exps::Bool=false;
                                    method::AlignmentMethod=ΛTAME_M(),graph::RandomGraphType=ErdosRenyi(),
                                    n_sizes = [100, 500, 1000, 2000,5000],
                                    edge_inclusion_p = [.5,1.0],
                                    step_percentage=[.1],profile=false,kwargs...)

    if seed_exps
        Random.seed!(0)
        seeds = rand(UInt64, length(step_percentage),length(edge_inclusion_p),length(n_sizes), trial_count)
    end

    seed = nothing
    results = []

    step_index = 1
    p_index = 1

    futures = []

    for sp in step_percentage

        p_index = 1
        for p in edge_inclusion_p

            n_index = 1

            for n in n_sizes

                for trial = 1:trial_count

                    if seed_exps
                        seed = seeds[step_index,p_index,n_index,trial]
                    end
                    future = @spawn random_graph_exp(n,p,graph; profile=profile,noise_model=noise_model,step_percent=sp,
                                                     seed=seed,method=method, kwargs...)
                    push!(futures,(seed,p,n,sp,future))
                end
                n_index += 1
            end
            p_index += 1
        end
        step_index += 1
    end

    #TODO: fix parsing of returned functions
    for (seed,p,n,sp,future) in futures
        if method === ΛTAME_M() || method === LowRankTAME_M()

            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()
                perm, dup_vertices, alignmentOutput = fetch(future)
            else
                perm, dup_vertices, (alignmentOutput, postProcessingOutput) = fetch(future)
            end

            matched_tris = alignmentOutput.matchScore
            A_tris,B_tris = alignmentOutput.motifCounts
            max_tris = min(alignmentOutput.motifCounts...)
            best_matching = alignmentOutput.matching
            if profile
                profiling_results = alignmentOutput.profile
            end

        elseif method === TAME_M()
            perm, dup_vertices, alignmentOutput = fetch(future)
            A_tris,B_tris = alignmentOutput.motifCounts
            matched_tris = alignmentOutput.matchScore
            max_tris = min(alignmentOutput.motifCounts...)
            best_matching = alignmentOutput.matching
            if profile
                profiling_results = alignmentOutput.profile
            end
        elseif method === ΛTAME_MultiMotif_M() || method === LowRankTAME_MultiMotif_M() || method === TAME_MultiMotif_M()
            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()
                perm, dup_vertices, (A_motifDistribution, B_motifDistribution, alignmentOutput) = fetch(future)
            else
                perm, dup_vertices, (A_motifDistribution, B_motifDistribution, alignmentOutput, postProcessingOutput) = fetch(future)
            end
            if profile
                profiling_results = alignmentOutput.profile
            end

            best_matching = alignmentOutput.matching
            A_tris = -1
            B_tris = -1
        elseif method === LowRankEigenAlign_M()

            if !haskey(kwargs,:postProcessing) || kwargs[:postProcessing] === noPostProcessing()
                perm, dup_vertices, alignmentOutput = fetch(future)
            else
                perm, dup_vertices, (alignmentOutput,postProcessingOutput) = fetch(future)
            end

            matched_tris = alignmentOutput.matchScore
            A_tris, B_tris = alignmentOutput.motifCounts
            max_tris = min(alignmentOutput.motifCounts...)
            best_matching = alignmentOutput.matching

            if profile
                profiling_results = alignmentOutput.profile
            end

        elseif method === EigenAlign_M() || method == Degree_M() || method === Random_M() || method === LowRankEigenAlignOnlyEdges_M()
            perm, dup_vertices, (A_tris, B_tris, matched_tris, best_matching, _) = fetch(future)
            max_tris = minimum((A_tris, B_tris))
        end

        if length(alignmentOutput.matching) > 0
            accuracy = sum([1 for (i,j) in alignmentOutput.matching if get(perm,i,-1) == j])/n
            #align_matrices is called with B,A since B > A

            dup_tolerant_perm = replaced_duplications_with_originals(perm,n, dup_vertices)
            dup_vertex_tolerant_accuracy = sum([1 for (j,i) in alignmentOutput.matching if get(dup_tolerant_perm,i,-1) == j])/n
        else
            accuracy = -1.0
            dup_vertex_tolerant_accuracy = -1.0
        end

        data_to_save = Array{Any,1}(undef,0)
        push!(data_to_save,seed)
        push!(data_to_save,p)
        push!(data_to_save,n)
        push!(data_to_save,sp)
        push!(data_to_save,accuracy)
        push!(data_to_save, dup_vertex_tolerant_accuracy)

        if method === ΛTAME_MultiMotif_M() || method === LowRankTAME_MultiMotif_M() || method === TAME_MultiMotif_M()
            push!(data_to_save,alignmentOutput.matchScore)
            push!(data_to_save,alignmentOutput.motifCounts[1])
            push!(data_to_save,alignmentOutput.motifCounts[2])
            push!(data_to_save,A_motifDistribution)
            push!(data_to_save,B_motifDistribution)
        else
            push!(data_to_save, matched_tris)
            push!(data_to_save, A_tris)
            push!(data_to_save, B_tris)
            push!(data_to_save, max_tris)
        end

        if profile
            push!(data_to_save,profiling_results)
        end
        if haskey(kwargs,:postProcessing)
            if typeof(kwargs[:postProcessing]) <: KlauAlgo
                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_tris_matched)
                push!(data_to_save,sum([1 for (j,i) in postProcessingOutput.matching if get(perm,j,-1) == i])/n)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.setup_rt)
                    push!(data_to_save,postProcessingOutput.klau_rt)
                end
                push!(data_to_save,postProcessingOutput.L_sparsity)
                push!(data_to_save,postProcessingOutput.f_status)
            elseif typeof(kwargs[:postProcessing]) <: SuccessiveKlauAlgo

                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_edges_matched)
                push!(data_to_save,postProcessingOutput.klau_tris_matched)
                push!(data_to_save,sum([1 for (j,i) in postProcessingOutput.matching if get(perm,j,-1) == i])/n)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.profiling)
                end
            elseif typeof(kwargs[:postProcessing]) <: LocalSearch

                push!(data_to_save,postProcessingOutput.original_edges_matched)
                push!(data_to_save,postProcessingOutput.localsearch_edges_matched)
                push!(data_to_save,postProcessingOutput.localsearch_tris_matched)
                push!(data_to_save,sum([1 for (j,i) in postProcessingOutput.matching if get(perm,j,-1) == i])/n)
                push!(data_to_save,postProcessingOutput.matching)
                if profile
                    push!(data_to_save,postProcessingOutput.profiling)
                    push!(data_to_save,postProcessingOutput.full_runtime)
                end
            end
        end

        push!(results,data_to_save)
    end
    return results
end

#used as a helper function for mapping duplicated vertices back to vertices in original graph
function replaced_duplications_with_originals(perm,n, dup_vertices)
    dup_tolerant_perm = copy(perm)
    dup_vertex_mapping = Dict([(n+i,v) for (i,v) in enumerate(dup_vertices)])

    while maximum(dup_tolerant_perm) > n
        dup_tolerant_perm = [haskey(dup_vertex_mapping,i) ? dup_vertex_mapping[i] : i for i in dup_tolerant_perm]
    end

    return dup_tolerant_perm

end

"""-----------------------------------------------------------------------------
                        ΛTAME Matching Experiments
-----------------------------------------------------------------------------"""
function ΛTAME_matching_exp(A_file::String, B_file::String, args...)

    @assert A_file[end-5:end] == ".ssten"
    @assert B_file[end-5:end] == ".ssten"

    A = load_SymTensorUnweighted(A_file,'\t') #delimiter may need to be updated
    B = load_SymTensorUnweighted(B_file,'\t')

    return ΛTAME_matching_exp(A,B,args...)

end

function ΛTAME_matching_exp(A_file::String, B_file::String, output_file::String=nothing, args...)

    @assert A_file[end-5:end] == ".ssten"
    @assert B_file[end-5:end] == ".ssten"
    @assert output_file[end-4:end] == ".json"

    A = load_SymTensorUnweighted(A_file,'\t') #delimiter may need to be updated
    B = load_SymTensorUnweighted(B_file,'\t')

    open(output_file,"w") do f
        JSON.print(f,ΛTAME_matching_exp(A,B,args...))
    end

end

function ΛTAME_matching_exp(A::Union{ThirdOrderSymTensor,SymTensorUnweighted{S}},
                            B::Union{ThirdOrderSymTensor,SymTensorUnweighted{S}},
                            max_iter::Int,alphas::Array{Float64,1}=[.5,1.0],
                            betas::Array{Float64,1}=[0.0,1.0,10.0,100.0]) where {S <:Motif}

    tol = 1e-16 # we want to track how triangles change over iterations
    A_motifs= size(A.indices,2)
    B_motifs= size(B.indices,2)

    results = []
    for α in alphas
        for β in betas

            matching_counts = []
            U,V = ΛTAME(A, B, β,max_iter,tol,α)

            d = size(U,2)

            for i=1:d
                motifs_matched,_,_ = TAME_score(A,B,U[:,1:i]*V[:,1:i]')
                push!(matching_counts,motifs_matched)
            end

            push!(results,(α,β,matching_counts))
        end
    end

    A_motifs, B_motifs, results
end

"""-----------------------------------------------------------------------------
                              SSHOPM Experiments
-----------------------------------------------------------------------------"""
function distributed_SSHOPM_exps(tensor_files::Array{String,1},output_path::String,seed::Bool;kwargs...)

    #TODO: test
    seeds = Array{UInt,2}(undef,1,1)

    if seed
        seeds = rand(UInt,length(tensor_files),length(tensor_files))
    end

    futures = []
    for i =1:length(tensor_files)
        for j=i+1:length(tensor_files)

            future = @spawn SSHOPM_exp(tensor_files[i],tensor_files[j],output_path,seeds[i,j])
            push!(futures,(tensor_files[i],tensor_files[j],future))
        end
    end

    #TODO: join using a sigpoll equivilant
    for (ten_A, ten_B, future) in futures
        fetch(future)
        println("joined experiment $ten_A + $ten_B")
    end
end

function distributed_SSHOPM_exps(tensor_files::Array{Tuple{String,String},1},output_path::String,seed::Bool;kwargs...)

    #TODO: test
    seeds = Array{UInt,2}(undef,1,1)

    if seed
        seeds = rand(UInt,length(tensor_files),length(tensor_files))
    end

    futures = []
    for (tensor_A_file, tensor_B_file) in tensor_files

        future = @spawn SSHOPM_exp(tensor_A_file, tensor_B_file, output_path,seeds[i,j])
        push!(futures,(tensor_A_file, tensor_B_file,future))

    end

    #TODO: join using a sigpoll equivilant
    for (ten_A, ten_B, future) in futures
        fetch(future)
        println("joined experiment $ten_A + $ten_B")
    end
end

function SSHOPM_exp(tensor_A_file::String, tensor_B_file::String,output_path::String,seed::UInt)

    @assert tensor_A_file[end-5:end] == ".ssten"
    @assert tensor_B_file[end-5:end] == ".ssten"
    @assert output_path[end] == '/'

    A = load_ThirdOrderSymTensor(tensor_A_file)
    B = load_ThirdOrderSymTensor(tensor_B_file)

    A_root = split(split(tensor_A_file,"/")[end],".ssten")[1]
    B_root = split(split(tensor_B_file,"/")[end],".ssten")[1]

    tol = 1e-16
    max_iter = 30
    β = 0.0
    samples = 200

    #delimeter heirarchy:
    #    - > : > +
    exp_filename = output_path*"align:$(A_root)+$(B_root)-beta:$β-max_iter:$max_iter-samples:$samples-seed:$seed-tol:$tol-results.json"

    relative_λ_diff, extremal_idx, eig_vals, extremal_vecs = SSHOPM_exp(A,B,samples,tol,max_iter,β,seed)

    #TODO: add in json save
    open(exp_filename,"w") do f
        JSON.print(f,[relative_λ_diff, extremal_idx, eig_vals, extremal_vecs])
    end

end

function SSHOPM_exp(A::ThirdOrderSymTensor, B::ThirdOrderSymTensor, samples::Int, tol::Float64,max_iter::Int,β::Float64,seed::UInt=nothing)

    if seed !== nothing
        Random.seed!(seed)
    end

    #A_V,A_Λ = SSHOPM_sample(A,samples,β,max_iter,tol)
    #B_V,B_Λ = SSHOPM_sample(B,samples,β,max_iter,tol)
    #AB_V,AB_Λ = SSHOPM_sample(A,B,samples,β,max_iter,tol)

    A_argmax_vec,A_argmax_val,A_Λ = SSHOPM_sample(A,samples,β,max_iter,tol)
    B_argmax_vec,B_argmax_val,B_Λ = SSHOPM_sample(B,samples,β,max_iter,tol)
    AB_argmax_vec,AB_argmax_val,AB_Λ = SSHOPM_sample(A,B,samples,β,max_iter,tol)

    A_i  = argmax([abs(x) for x in A_Λ])
    B_i  = argmax([abs(x) for x in B_Λ])
    AB_i = argmax([abs(x) for x in AB_Λ])

    relative_λ_diff = abs(A_argmax_val*B_argmax_val - AB_argmax_val)/abs(AB_argmax_val)

    println(relative_λ_diff)

    #r = rank(AB_V[:,:,AB_i])
   #U,S,Vt = svd(AB_V[:,:,AB_i])

    # not assuming anything about the rank of argmax(AB_Λ),
    # error will trigger if not rank 1.
    #u_Avec_inner_prod = dot(U[:,1:r],A_V[:,A_i])
    #v_Bvec_inner_prod = dot(Vt[:,1:r],B_V[:,B_i])

    #  -- subspace angle  --  #
    #u_Av_subspaceAngle = acos(dot(U[:,1:r],A_V[:,A_i]))
    #v_Bv_subspaceAngle = acos(dot(Vt[:,1:r],B_V[:,B_i]))

    #  -- relative norm code  --  #
    #first_nnz_sign = x-> sign(x[findfirst(x .!= 0.0)])
    #relative_u_norm = norm( U[:,1:r]*Diagonal([first_nnz_sign( U[:,j]) for j in 1:r]) - first_nnz_sign(A_V[:,A_i])*A_V[:,A_i])/norm(A_V[:,A_i])
    #relative_v_norm = norm(Vt[:,1:r]*Diagonal([first_nnz_sign(Vt[:,j]) for j in 1:r]) - first_nnz_sign(B_V[:,B_i])*B_V[:,B_i])/norm(B_V[:,B_i])

    # -- group the variables to return from the experiment -- #
    #comparisons = (relative_λ_diff, u_Avec_inner_prod, v_Bvec_inner_prod)
    extremal_idx = (A_i, B_i, AB_i)
    eig_vals = (A_Λ, B_Λ, AB_Λ)
    extremal_vecs = (A_argmax_vec, B_argmax_vec, AB_argmax_vec)

    return relative_λ_diff, extremal_idx, eig_vals, extremal_vecs
end

function test_SSHOPM_exp_file(filename)

    relative_λ_diff, extremal_idx, eig_vals, extremal_vecs = open(filename,"r") do f
        JSON.parse(f)
    end

    subspace_angle = dot(extremal_vecs[1]*extremal_vecs[2]',hcat(extremal_vecs[3]...))
    λ_A  = eig_vals[1][extremal_idx[1]]
    λ_B  = eig_vals[2][extremal_idx[2]]
    λ_AB = eig_vals[3][extremal_idx[3]]

    println("A.n = $(length(extremal_vecs[1]))  B.n=$(length(extremal_vecs[2]))")
    println("λ_A: $λ_A  -- λ_B: $λ_B  -- λ_AB: $λ_AB")
    println("λ_Aλ_B - λ_AB = $(λ_A*λ_B - λ_AB)")
    println("<(v_A*v_B, v_AB) = $subspace_angle")
end

struct HyperKron <: RandomGraphType end
"""-----------------------------------------------------------------------------
  Runs an instance of a graph alignment problem using a random graph model. Once
  a graph is generated, we use the function 'ER_noise_model' to generate a
  second network to align against it. Once the second network is generated it is
  randomly permuted. All additional kwargs are passed to the 'align_matrices'
  function.

  #TODO: update docs

  Inputs
  ------
  * n - (Int):
    The node of nodes to use in each network.
  * 'p_remove' -(Float):
    The probability of removing an edge in the ER noise model. The probability
    of adding in an edge is computed using a function of this parameter at the
    density of the generated network.
  * 'graph_type' - (String):
    The type of graph used to generate the network. Options include 'ER' for an
    Erdos Renyi network, 'RandomGeometric' for a Random Geometric network from
    a 2D unit square, and 'HyperKron' which comes from the Triangle Generalized
    Preferential Attachment model.
  * seed - (Any):
    A seed used to replicate the results. Our experiments typically use UInts.
  * 'use_metis' - (Bool):
    Indicates whether or not to use metis clusters to permute the adjacency
    matrix. For large networks this can improve the runtime by improving the
    access patterns.
  * degreedist - (UnivariateDistribution):
    A random variable generator from Distributions.jl which is used to generate
    the degrees for random nodes. We primarily use
    Distributions.LogNormal(log(5,1)).
  * 'p_edges' - (func Int - Float):
    An anonymous function used for computing the probability of adding edges.
    Primarily used in the ER graphs for generating the probability of adding
    in edges from the number of nodes in the network. an example expected input
    would be 'n -> 2*log(n)/n'.

    Outputs
    -------
    * 'd_A','d_B' - (Array{Float,1}):
      The degrees of all the nodes in graphs A and B. The ith entry of each
      being the degree of node i in the respective networks.
    * perm - (Array{Int,1}):
      The permutation used to shuffle the nodes in graph B. Used for computing
      the precision of the methods.
    * All additional return variables are from 'align_matrices' function call.
-----------------------------------------------------------------------------"""
function random_graph_exp(n::Int, perturbation_p::Float64,graph::RandomGraphType;
                          noise_model::NoiseModelFlag,seed=nothing,step_percent=.1,
                          use_metis=false,degreedist=nothing,p_edges=nothing,kwargs...)

    if seed !== nothing
		Random.seed!(seed)
	end

    if typeof(graph) === ErdosRenyi
        if degreedist === nothing
            if p_edges === nothing
                p = 2*log(n)/n
            else
                p = p_edges(n)
            end
            A = erdos_renyi(n,p)
        else
            A = erdos_renyi(n;degreedist)
            p = nnz(A)/n^2
        end

	elseif typeof(graph) === RandomGeometric

		if degreedist === nothing
            k = 10
            if p_edges === nothing
                p = k/n
            end
			A = random_geometric_graph(n,k)
		else
			d = 2
			A = spatial_network(n, d;degreedist= degreedist)
			p = nnz(A)/n^2
		end

    elseif typeof(graph) == HyperKron
        if p_edges === nothing
            p = .4#2*log(n)/n
        end
		r = .4

		A = sparse(gpa_graph(n,p,r,5))

	else
		error("invalid graph type: $graph_type\n must be either 'ER','RandomGeometric' or 'HyperKron'")
	end

    if typeof(noise_model) === ErdosRenyiNoise
        p_add = (p*perturbation_p)/(1-p)
        #B = ER_noise_model(A,0.0,perturbation_p)#,p_add)
        B = ER_noise_model(A,perturbation_p,p_add)
    elseif  typeof(noise_model) === DuplicationNoise
        steps = Int(ceil(step_percent*size(A,1))) #adds in 10% new nodes by default
        B,dup_vertices = duplication_perturbation_noise_model(A,steps, perturbation_p)
    end

    perm = shuffle(1:size(B,1))
    B = B[perm,perm]

    #return A, B , perm

	if use_metis
		apply_Metis_permutation!(A,100)
		apply_Metis_permutation!(B,100)
	end

    d_A = A*ones(size(A,1))
    d_B = B*ones(size(B,1))

    if typeof(noise_model) === DuplicationNoise
        #  B is always larger than A
        perm,dup_vertices,align_matrices(B,A;kwargs...)
    else
        return d_A,d_B,perm,align_matrices(A,B;kwargs...)
    end
end

function random_graph_exp(file::String,perturbation_p::Float64;noise_model::NoiseModelFlag,seed=nothing,step_percent=.1,
                          use_metis=false,degreedist=nothing,p_edges=nothing,kwargs...)

    if seed !== nothing
        println("seeding with $seed")
		Random.seed!(seed)
	end

    A = MatrixNetworks.readSMAT(file)
    n = size(A,1)

    #=
	elseif typeof(graph) === RandomGeometric

		if degreedist === nothing
            k = 10
            if p_edges === nothing
                p = k/n
            end
			A = random_geometric_graph(n,k)
		else
			d = 2
			A = spatial_network(n, d;degreedist= degreedist)
			p = nnz(A)/n^2
		end

    elseif typeof(graph) == HyperKron
        if p_edges === nothing
            p = .4#2*log(n)/n
        end
		r = .4

		A = sparse(gpa_graph(n,p,r,5))

	else
		error("invalid graph type: $graph_type\n must be either 'ER','RandomGeometric' or 'HyperKron'")
	end
    =#

    if typeof(noise_model) === ErdosRenyiNoise
        p = nnz(A)/size(A,1)^2
        p_add = (p*perturbation_p)/(1-p)
        #B = ER_noise_model(A,0.0,perturbation_p)#,p_add)
        B = ER_noise_model(A,perturbation_p,p_add)
    elseif  typeof(noise_model) === DuplicationNoise
        steps = Int(ceil(step_percent*size(A,1))) #adds in 10% new nodes by default
        B,dup_vertices = duplication_perturbation_noise_model(A,steps, perturbation_p)
    end

    perm = shuffle(1:size(B,1))
    B = B[perm,perm]

    #return A, B , perm

	if use_metis
		apply_Metis_permutation!(A,100)
		apply_Metis_permutation!(B,100)
	end

    d_A = A*ones(size(A,1))
    d_B = B*ones(size(B,1))

    if typeof(noise_model) === DuplicationNoise
        #  B is always larger than A
        return perm,dup_vertices,align_matrices(B,A;kwargs...),A,B
    else
        return d_A,d_B,perm,align_matrices(A,B;kwargs...)
    end
end

function ER_noise_model(A,p_remove::F,p_add::F) where {F <: AbstractFloat}
    n = size(A,1)
    B = copy(A)

    is,js,_ = findnz(erdos_renyi(n,p_remove))
    for (i,j) in zip(is,js)
        if A[i,j] > 0
            B[i,j] = 0
            B[j,i] = 0
        end
    end

    is,js,_ = findnz(erdos_renyi(n,p_add))
    for (i,j) in zip(is,js)
        if A[i,j] == 0.0
            B[i,j] = 1
            B[j,i] = 1
        end
    end

    dropzeros!(B)
    return B
end

function duplication_perturbation_noise_model(A::SparseMatrixCSC{T,Int},steps::Integer, new_edge_p::Float64,seed::Union{UInt,Nothing}=nothing) where T

    is_undirected(A)                      || throw(ArgumentError("A must be undirected."))
    (new_edge_p >= 0 && new_edge_p <= 1)  || throw(ArgumentError("new_edge_p must be a probability."))
    steps >= 0                            || throw(ArgumentError("Must take a non-negative number of steps."))
    # let it steps equal 0 for testing purposes

    if seed !== nothing
        Random.seed!(seed)
    end

    duplicated_vertices = Array{Int,1}(undef, steps)

    n,_ = size(A) # n will be updated

    #store A as an edge list so it's fast to sample
    A_edge_list = Array{Array{Tuple{Int,T},1},1}(undef,n+steps)
    for i = 1:n
        A_edge_list[i] = collect(zip(findnz(A[i,:])...))
    end
    for i = n+1:n+steps
        A_edge_list[i] = Array{Tuple{Int,T},1}(undef,0)
    end

    for step in 1:steps

        dup_vertex = rand(1:n)
        duplicated_vertices[step] = dup_vertex
        for (neighbor,weight) in A_edge_list[dup_vertex]
            if rand() < new_edge_p
                push!(A_edge_list[n+1],(neighbor,weight))
                push!(A_edge_list[neighbor],(n+1,weight))
            end
        end
        n += 1
    end

    #convert edge list back into a MatrixNetwork
    total_edges = 0
    for i=1:n
        total_edges += length(A_edge_list[i])
    end

    Is = Array{Int,1}(undef,total_edges)
    Js = Array{Int,1}(undef,total_edges)
    Vs = Array{T,1}(undef,total_edges)

    edge_idx = 1
    for i=1:n
        for (n_j,weight) in A_edge_list[i]
            Is[edge_idx] = i
            Js[edge_idx] = n_j
            Vs[edge_idx] = weight
            edge_idx += 1
        end
    end

    return sparse(Js,Is,Vs,n,n), duplicated_vertices
end

max_stars_repo_name: chuckcol/LambdaTAME

max_issues_repo_name: chuckcol/LambdaTAME

max_forks_repo_name: chuckcol/LambdaTAME


text:
"""
Calculate the contributions of a single source for a single pixel to
the sensitive floats E_G_s and var_G_s, which are cleared and updated in place.

Args:
    - ea: Model parameters
    - vp: the variational parameters
    - elbo_vars: Elbo intermediate values, with updated fs1m and fs0m.
    - sb: Source brightnesse
    - s: The source, in 1:ea.S
    - b: The band

Returns:
    Updates E_G_s, E_G2_s, and var_G_s in place with the brightness
    for this source at this pixel.
"""
function calculate_G_s!(vp::VariationalParams{T},
                        elbo_vars::ElboIntermediateVariables{T},
                        sb::SourceBrightness{T},
                        b::Int,
                        s::Int,
                        is_active_source::Bool) where {T<:Number}
    E_G_s = elbo_vars.E_G_s
    E_G2_s = elbo_vars.E_G2_s
    var_G_s = elbo_vars.var_G_s

    @assert E_G_s.local_P == var_G_s.local_P == length(CanonicalParams)
    @assert E_G_s.local_S == var_G_s.local_S == 1
    @assert elbo_vars.fs0m.local_P == length(StarPosParams)
    @assert elbo_vars.fs1m.local_P == length(GalaxyPosParams)

    # we'd like to get rid of these calls to `fill!`
    if is_active_source
        SensitiveFloats.zero!(E_G_s)
        SensitiveFloats.zero!(E_G2_s)
        SensitiveFloats.zero!(var_G_s)
    else
        E_G_s.v[] = 0.0
        E_G2_s.v[] = 0.0
        var_G_s.v[] = 0.0
    end

    @inbounds for i in 1:NUM_SOURCE_TYPES # Celestial object types (e.g. stars and galaxies)
        fsm_i = (i == 1) ? elbo_vars.fs0m : elbo_vars.fs1m
        a_i = vp[s][ids.is_star[i]]
        sb_E_l_a_b_i = sb.E_l_a[b, i]
        sb_E_ll_a_b_i = sb.E_ll_a[b, i]

        fsm_i_v = fsm_i.v[]
        sb_E_l_a_b_i_v = sb_E_l_a_b_i.v[]
        sb_E_ll_a_b_i_v = sb_E_ll_a_b_i.v[]

        fsm_i_d = fsm_i.d
        sb_E_l_a_b_i_d = sb_E_l_a_b_i.d
        sb_E_ll_a_b_i_d = sb_E_ll_a_b_i.d
        E_G_s_d = E_G_s.d
        E_G2_s_d = E_G2_s.d

        E_G_s_h = E_G_s.h
        E_G2_s_h = E_G2_s.h

        lf = sb_E_l_a_b_i_v * fsm_i_v
        llff = sb_E_ll_a_b_i_v * fsm_i_v^2

        E_G_s.v[] += a_i * lf
        E_G2_s.v[] += a_i * llff

        ############ Only gradient and hessian code below ##############
        (is_active_source && elbo_vars.elbo.has_gradient) || continue

        E_G_s_d[ids.is_star[i], 1] += lf
        E_G2_s_d[ids.is_star[i], 1] += llff

        p0_shape = shape_standard_alignment[i]
        p0_bright = brightness_standard_alignment[i]
        u_ind = i == 1 ? star_ids.pos : gal_ids.pos

        # Derivatives with respect to the spatial parameters
        tmp1 = sb_E_l_a_b_i_v * a_i
        tmp2 = sb_E_ll_a_b_i_v * 2 * fsm_i_v * a_i
        for p0_shape_ind in 1:length(p0_shape)
            E_G_s_d[p0_shape[p0_shape_ind], 1] += tmp1 * fsm_i_d[p0_shape_ind, 1]
            E_G2_s_d[p0_shape[p0_shape_ind], 1] += tmp2 * fsm_i_d[p0_shape_ind, 1]
        end

        # Derivatives with respect to the brightness parameters.
        for p0_bright_ind in 1:length(p0_bright)
            E_G_s_d[p0_bright[p0_bright_ind], 1] =
                a_i * fsm_i_v * sb_E_l_a_b_i_d[p0_bright_ind, 1]
            E_G2_s_d[p0_bright[p0_bright_ind], 1] =
                a_i * (fsm_i_v^2) * sb_E_ll_a_b_i_d[p0_bright_ind, 1]
        end

        ############ only hessian code below ##############
        elbo_vars.elbo.has_hessian || continue

        E_G_s_hsub = elbo_vars.E_G_s_hsub_vec[i]
        E_G2_s_hsub = elbo_vars.E_G2_s_hsub_vec[i]

        # The (a, a) block of the hessian is zero.

        # The (bright, bright) block:
        for p0_ind1 in 1:length(p0_bright), p0_ind2 in 1:length(p0_bright)
            E_G_s_h[p0_bright[p0_ind1], p0_bright[p0_ind2]] =
                a_i * sb_E_l_a_b_i.h[p0_ind1, p0_ind2] * fsm_i_v
            E_G2_s_h[p0_bright[p0_ind1], p0_bright[p0_ind2]] =
                (fsm_i_v^2) * a_i * sb_E_ll_a_b_i.h[p0_ind1, p0_ind2]
        end

        # The (shape, shape) block:
        p1, p2 = size(E_G_s_hsub.shape_shape)
        for ind1 = 1:p1, ind2 = 1:p2
            E_G_s_hsub.shape_shape[ind1, ind2] =
                a_i * sb_E_l_a_b_i_v * fsm_i.h[ind1, ind2]
            E_G2_s_hsub.shape_shape[ind1, ind2] =
                2 * a_i * sb_E_ll_a_b_i_v * (
                    fsm_i_v * fsm_i.h[ind1, ind2] +
                    fsm_i_d[ind1, 1] * fsm_i_d[ind2, 1])
        end

        # The u_u submatrix of this assignment will be overwritten after
        # the loop.
        for p0_ind1 in 1:length(p0_shape), p0_ind2 in 1:length(p0_shape)
            E_G_s_h[p0_shape[p0_ind1], p0_shape[p0_ind2]] =
                a_i * sb_E_l_a_b_i_v * fsm_i.h[p0_ind1, p0_ind2]
            E_G2_s_h[p0_shape[p0_ind1], p0_shape[p0_ind2]] =
                E_G2_s_hsub.shape_shape[p0_ind1, p0_ind2]
        end

        # Since the u_u submatrix is not disjoint between different i, accumulate
        # it separate and add it at the end.
        for u_ind1 = 1:2, u_ind2 = 1:2
            E_G_s_hsub.u_u[u_ind1, u_ind2] =
                E_G_s_hsub.shape_shape[u_ind[u_ind1], u_ind[u_ind2]]
            E_G2_s_hsub.u_u[u_ind1, u_ind2] =
                E_G2_s_hsub.shape_shape[u_ind[u_ind1], u_ind[u_ind2]]
        end

        # All other terms are disjoint between different i and don't involve
        # addition, so we can just assign their values (which is efficient in
        # native julia).

        # The (a, bright) blocks:
        for p0_ind in 1:length(p0_bright)
            E_G_s_h[p0_bright[p0_ind], ids.is_star[i]] =
                fsm_i_v * sb_E_l_a_b_i_d[p0_ind, 1]
            E_G2_s_h[p0_bright[p0_ind], ids.is_star[i]] =
                (fsm_i_v ^ 2) * sb_E_ll_a_b_i_d[p0_ind, 1]
            E_G_s_h[ids.is_star[i], p0_bright[p0_ind]] =
                E_G_s_h[p0_bright[p0_ind], ids.is_star[i]]
            E_G2_s_h[ids.is_star[i], p0_bright[p0_ind]] =
                E_G2_s_h[p0_bright[p0_ind], ids.is_star[i]]
        end

        # The (a, shape) blocks.
        for p0_ind in 1:length(p0_shape)
            E_G_s_h[p0_shape[p0_ind], ids.is_star[i]] =
                sb_E_l_a_b_i_v * fsm_i_d[p0_ind, 1]
            E_G2_s_h[p0_shape[p0_ind], ids.is_star[i]] =
                sb_E_ll_a_b_i_v * 2 * fsm_i_v * fsm_i_d[p0_ind, 1]
            E_G_s_h[ids.is_star[i], p0_shape[p0_ind]] =
                E_G_s_h[p0_shape[p0_ind], ids.is_star[i]]
            E_G2_s_h[ids.is_star[i], p0_shape[p0_ind]] =
                E_G2_s_h[p0_shape[p0_ind], ids.is_star[i]]
        end

        for ind_b in 1:length(p0_bright), ind_s in 1:length(p0_shape)
            E_G_s_h[p0_bright[ind_b], p0_shape[ind_s]] =
                a_i * sb_E_l_a_b_i_d[ind_b, 1] * fsm_i_d[ind_s, 1]
            E_G2_s_h[p0_bright[ind_b], p0_shape[ind_s]] =
                2 * a_i * sb_E_ll_a_b_i_d[ind_b, 1] * fsm_i_v * fsm_i_d[ind_s]

            E_G_s_h[p0_shape[ind_s], p0_bright[ind_b]] =
                E_G_s_h[p0_bright[ind_b], p0_shape[ind_s]]
            E_G2_s_h[p0_shape[ind_s], p0_bright[ind_b]] =
                E_G2_s_h[p0_bright[ind_b], p0_shape[ind_s]]
        end
    end # i loop

    if elbo_vars.elbo.has_hessian
        # Accumulate the pos Hessian. pos is the only parameter that is shared between
        # different values of i.

        # This is
        # for i = 1:NUM_SOURCE_TYPES
        #     E_G_u_u_hess += elbo_vars.E_G_s_hsub_vec[i].u_u
        #     E_G2_u_u_hess += elbo_vars.E_G2_s_hsub_vec[i].u_u
        # end
        # For each value in 1:NUM_SOURCE_TYPES, written this way for speed.
        for u_ind1 = 1:2, u_ind2 = 1:2
            elbo_vars.E_G_s.h[ids.pos[u_ind1], ids.pos[u_ind2]] =
                elbo_vars.E_G_s_hsub_vec[1].u_u[u_ind1, u_ind2] +
                elbo_vars.E_G_s_hsub_vec[2].u_u[u_ind1, u_ind2]

            elbo_vars.E_G2_s.h[ids.pos[u_ind1], ids.pos[u_ind2]] =
                elbo_vars.E_G2_s_hsub_vec[1].u_u[u_ind1, u_ind2] +
                elbo_vars.E_G2_s_hsub_vec[2].u_u[u_ind1, u_ind2]
        end
    end

    ####################################################
    #### the code below loads var_G_s ##################

    var_G_s.v[] = E_G2_s.v[] - (E_G_s.v[] ^ 2)

    ############### only gradient and hessian code below
    (is_active_source && elbo_vars.elbo.has_gradient) || return

    var_G_s_d = elbo_vars.var_G_s.d
    E_G2_s_d = elbo_vars.E_G2_s.d
    E_G_s_d = elbo_vars.E_G_s.d

    @assert length(var_G_s_d) == length(E_G2_s_d) == length(E_G_s_d)

    @inbounds for ind1 = 1:length(var_G_s_d)
        var_G_s_d[ind1] = E_G2_s_d[ind1] - 2 * E_G_s.v[] * E_G_s_d[ind1]
    end

    ########## only hessian code below
    elbo_vars.elbo.has_hessian || return

    var_G_s_h = elbo_vars.var_G_s.h
    E_G2_s_h = elbo_vars.E_G2_s.h
    E_G_s_h = elbo_vars.E_G_s.h
    p1, p2 = size(var_G_s_h)
    @inbounds for ind2 = 1:p2, ind1 = 1:ind2
        var_G_s_h[ind1, ind2] =
            E_G2_s_h[ind1, ind2] - 2 * (
                E_G_s.v[] * E_G_s_h[ind1, ind2] +
                E_G_s_d[ind1, 1] * E_G_s_d[ind2, 1])
        var_G_s_h[ind2, ind1] = var_G_s_h[ind1, ind2]
    end
end

"""
Add the contributions from a single source at a single pixel to the
sensitive floast E_G and var_G, which are updated in place.
"""
function accumulate_source_pixel_brightness!(
        ea::ElboArgs,
        vp::VariationalParams{T},
        elbo_vars::ElboIntermediateVariables{T},
        sb::SourceBrightness{T},
        b::Int, s::Int,
        is_active_source::Bool) where {T<:Number}

    calculate_G_s!(vp, elbo_vars, sb, b, s, is_active_source)

    if is_active_source
        sa = findfirst(ea.active_sources, s)
        add_sources_sf!(elbo_vars.E_G, elbo_vars.E_G_s, sa)
        add_sources_sf!(elbo_vars.var_G, elbo_vars.var_G_s, sa)
    else
        # If the sources is inactive, simply accumulate the values.
        elbo_vars.E_G.v[] += elbo_vars.E_G_s.v[]
        elbo_vars.var_G.v[] += elbo_vars.var_G_s.v[]
    end
end

"""
Add the lower bound to the log term to the elbo for a single pixel.

Args:
     - elbo_vars: Intermediate variables
     - x_nbm: The photon count at this pixel
     - iota: The optical sensitivity

 Returns:
    Updates elbo_vars.elbo in place by adding the lower bound to the log
    term.
"""
function add_elbo_log_term!(elbo_vars::ElboIntermediateVariables{T},
                            E_G::SensitiveFloat{T},
                            var_G::SensitiveFloat{T},
                            elbo::SensitiveFloat{T},
                            x_nbm::AbstractFloat,
                            iota::AbstractFloat) where {T<:Number}
    # See notes for a derivation. The log term is
    # log E[G] - Var(G) / (2 * E[G] ^2 )

    @inbounds begin
        E_G_v = E_G.v[]
        var_G_v = var_G.v[]

        # The gradients and Hessians are written as a f(x, y) = f(E_G2, E_G)
        log_term_value = log(E_G_v) - var_G_v / (2.0 * E_G_v ^ 2)

        # Add x_nbm * (log term * log(iota)) to the elbo.
        # If not calculating derivatives, add the values directly.
        elbo.v[] += x_nbm * (log(iota) + log_term_value)

        if elbo_vars.elbo.has_gradient
            elbo_vars.combine_grad[1] = -0.5 / (E_G_v ^ 2)
            elbo_vars.combine_grad[2] = 1 / E_G_v + var_G_v / (E_G_v ^ 3)

            if elbo_vars.elbo.has_hessian
                elbo_vars.combine_hess[1, 1] = 0.0
                elbo_vars.combine_hess[1, 2] = elbo_vars.combine_hess[2, 1] = 1 / E_G_v^3
                elbo_vars.combine_hess[2, 2] =
                    -(1 / E_G_v ^ 2 + 3 * var_G_v / (E_G_v ^ 4))
            end

            # Calculate the log term.
            combine_sfs!(
                var_G, E_G, elbo_vars.elbo_log_term,
                log_term_value, elbo_vars.combine_grad, elbo_vars.combine_hess)

            # Add to the ELBO.
            elbo_d = elbo.d
            elbo_vars_elbo_log_term_d = elbo_vars.elbo_log_term.d

            for ind in 1:length(elbo_d)
                elbo_d[ind] += x_nbm * elbo_vars_elbo_log_term_d[ind]
            end

            if elbo_vars.elbo.has_hessian
                elbo_h = elbo.h
                elbo_vars_elbo_log_term_h = elbo_vars.elbo_log_term.h
                for ind in 1:length(elbo_h)
                    elbo_h[ind] += x_nbm * elbo_vars_elbo_log_term_h[ind]
                end
            end
        end
    end
end

function add_pixel_term!(
        ea::ElboArgs,
        vp::VariationalParams{T},
        n::Int, h::Int, w::Int,
        bvn_bundle::BvnBundle{T},
        sbs::Vector{SourceBrightness{T}},
        elbo_vars::ElboIntermediateVariables = ElboIntermediateVariables(T, ea.Sa)) where {T<:Number}
    img = ea.images[n]

    SensitiveFloats.zero!(elbo_vars.E_G)
    SensitiveFloats.zero!(elbo_vars.var_G)

    for s in 1:ea.S
        p = ea.patches[s,n]

        h2 = h - p.bitmap_offset[1]
        w2 = w - p.bitmap_offset[2]

        H2, W2 = size(p.active_pixel_bitmap)
        if 1 <= h2 <= H2 && 1 <= w2 < W2 && p.active_pixel_bitmap[h2, w2]
            is_active_source = s in ea.active_sources

            # this if/else block is for reporting purposes only
            if is_active_source
                elbo_vars.active_pixel_counter[] += 1
            else
                elbo_vars.inactive_pixel_counter[] += 1
            end

            Model.star_light_density!(elbo_vars.fs0m, p, h, w, vp[s][ids.pos], is_active_source)
            Model.populate_gal_fsm!(elbo_vars.fs1m,
                                    bvn_bundle.bvn_derivs,
                                    s, h, w,
                                    is_active_source,
                                    p.wcs_jacobian,
                                    bvn_bundle.gal_mcs)

            accumulate_source_pixel_brightness!(ea, vp, elbo_vars,
                sbs[s], ea.images[n].b, s, is_active_source)
        end
    end

    # There are no derivatives with respect to epsilon, so can safely add
    # to the value.
    elbo_vars.E_G.v[] += img.sky[h, w]

    # Add the terms to the elbo given the brightness.
    add_elbo_log_term!(elbo_vars,
                       elbo_vars.E_G,
                       elbo_vars.var_G,
                       elbo_vars.elbo,
                       img.pixels[h,w],
                       img.nelec_per_nmgy[h])
    add_scaled_sfs!(elbo_vars.elbo,
                    elbo_vars.E_G,
                    -img.nelec_per_nmgy[h])

    # Subtract the log factorial term. This is not a function of the
    # parameters so the derivatives don't need to be updated. Note that
    # even though this does not affect the ELBO's maximum, it affects
    # the optimization convergence criterion, so I will leave it in for now.
    elbo_vars.elbo.v[] -= lgamma(img.pixels[h,w] + 1.0)
end

"""
Return the expected log likelihood for all bands in a section
of the sky.
Returns: A sensitive float with the log likelihood.
"""
function elbo_likelihood(ea::ElboArgs,
                         vp::VariationalParams{T},
                         elbo_vars::ElboIntermediateVariables = ElboIntermediateVariables(T, ea.Sa),
                         bvn_bundle::BvnBundle{T} = BvnBundle{T}(ea.psf_K, ea.S)) where {T}
    zero!(elbo_vars)
    Model.zero!(bvn_bundle)

    # this call loops over light sources (but not images)
    sbs = load_source_brightnesses(ea, vp)

    for n in 1:ea.N
        img = ea.images[n]

        # could preallocate these---outside of elbo_likehood even to use for
        # all ~50 evalulations of the likelihood
        # This convolves the PSF with the star/galaxy model, returning a
        # mixture of bivariate normals.
        star_mcs, gal_mcs = Model.load_bvn_mixtures!(
                                    bvn_bundle.star_mcs,
                                    bvn_bundle.gal_mcs,
                                    ea.S, ea.patches,
                                    vp, ea.active_sources,
                                    ea.psf_K, n,
                                    elbo_vars.elbo.has_gradient,
                                    elbo_vars.elbo.has_hessian)

        # if there's only one active source, we know each pixel we visit
        # hasn't been visited before, so no need to allocate memory.
        # currently length(ea.active_sources) > 1 only in unit tests, never
        # when invoked from `bin`.
        already_visited = length(ea.active_sources) == 1 ?
                              falses(0, 0) :
                              falses(size(img.pixels))

        # iterate over the pixels by iterating over the patches, and visiting
        # all the pixels in the patch that are active and haven't already been
        # visited
        for s in ea.active_sources
            p = ea.patches[s, n]
            H2, W2 = size(p.active_pixel_bitmap)
            for w2 in 1:W2, h2 in 1:H2
                # (h2, w2) index the local patch, while (h, w) index the image
                h = p.bitmap_offset[1] + h2
                w = p.bitmap_offset[2] + w2

                if !p.active_pixel_bitmap[h2, w2]
                    continue
                end

                # if there's only one source to visit, we know this pixel is new
                if length(ea.active_sources) != 1
                    if already_visited[h,w]
                        continue
                    end
                    already_visited[h,w] = true
                end

                # Some pixels that are NaN in the original image may be active
                # for the convolution code.
                if isnan(img.pixels[h, w])
                    continue
                end

                # if we're here it's a unique active pixel.
                # Note that although we are iterating over pixels within a
                # single patch, add_pixel_term /also/ iterates over patches to
                # find all patches that overlap with this pixel.
                add_pixel_term!(ea, vp, n, h, w, bvn_bundle, sbs, elbo_vars)
            end
        end
    end

    assert_all_finite(elbo_vars.elbo)
    elbo_vars.elbo
end

"""
Calculates and returns the ELBO and its derivatives for all the bands
of an image.
Returns: A sensitive float containing the ELBO for the image.
"""
function elbo(ea::ElboArgs,
              vp::VariationalParams{T},
              elbo_vars::ElboIntermediateVariables =
                  ElboIntermediateVariables(T, ea.Sa, true,  T<:AbstractFloat),
              bvn_bundle::BvnBundle = BvnBundle{T}(ea.psf_K, ea.S)) where {T}
    @assert(all(all(isfinite, vs) for vs in vp), "vp contains NaNs or Infs")
    result = elbo_likelihood(ea, vp, elbo_vars, bvn_bundle)
    ea.include_kl && KLDivergence.subtract_kl_all_sources!(ea, vp, result)
    assert_all_finite(elbo_vars.elbo)
    return result
end

ext: jl


text:
import Base:isless
using LightGraphs: Graph, neighbors, degree, nv, ne
using Combinatorics: permutations

export are_isomorphic, enumerate_steiner, label_fst, enumerate_fst

"find unlabeled nodes with at most one unlabeled neighbor"
function find_next_nodes_to_label(graph::Graph, labels::Vector{Int})
    n = nv(graph)
    @assert n == length(labels)
    nolabel = (1:n)[labels .== 0]
    filter(v -> sum(labels[neighbors(graph, v)] .== 0) ≤ 1, nolabel)
end

"combine neighbor labels into tentative labels"
function tentative_labels(graph::Graph, labels::Vector{Int}, nodes::Vector{Int})
    @assert all(labels[nodes] .== 0)
    tent = Tuple{Vector{Int},Int}[]
    for u in nodes
        nblabels = filter(l -> l ≠ 0, labels[neighbors(graph, u)])
        sort!(nblabels)
        push!(tent, (nblabels, u))
    end
    sort!(tent)
    tent
end

"""
Are two given (undirected) trees isomorphic?

Reference: algorithm 2 from Fampa et al.: A specialized branch-and-bound
algorithm for the Euclidean Steiner tree problem in n-space.
"""
function are_isomorphic(tree1::Topology, tree2::Topology)
    @assert is_tree(tree1) && is_tree(tree2)
    g1 = Graph(digraph_from_topology(tree1))
    g2 = Graph(digraph_from_topology(tree2))

    if nv(g1) != nv(g2) || ne(g1) != ne(g2)
        return false
    end

    label1, label2 = fill(0, nv(g1)), fill(0, nv(g2))
    label = 1

    # first label all the leaf nodes
    leaves1 = [n for n=1:nv(g1) if degree(g1, n) == 1]
    leaves2 = [n for n=1:nv(g2) if degree(g2, n) == 1]
    if length(leaves1) ≠ length(leaves2)
        return false
    end
    label1[leaves1] .= label
    label2[leaves2] .= label
    label += 1

    # add more labels until difference or isomorphism is found
    while true
        # find next nodes to give labels to
        next1 = find_next_nodes_to_label(g1, label1)
        next2 = find_next_nodes_to_label(g2, label2)
        if length(next1) ≠ length(next2)
            return false
        elseif length(next1) == 0
            return true
        end

        # try labels from multiset of neighbors' labels
        tent1 = tentative_labels(g1, label1, next1)
        tent2 = tentative_labels(g2, label2, next2)
        if any(t1[1] ≠ t2[1] for (t1, t2) in zip(tent1, tent2))
            return false
        end

        # actually assign the label to smallest of tentatives
        smallest = tent1[1][1]
        label1[[t[2] for t in tent1 if t[1] == smallest]] .= label
        label2[[t[2] for t in tent2 if t[1] == smallest]] .= label
        label += 1
    end
end

"""
Enumerate all representative trees.

This considers only the trees interconnecting the Steiner nodes for full Steiner
trees of n terminals.

Following algorithm 1 in Fampa et al.: A specialized branch-and-bound algorithm
for the Euclidean Steiner tree problem in n-space.
"""
function enumerate_steiner(n::Int)
    nodes = [Node(i,i) for i=1:n-2]
    nclass = fill(0, n - 2) # by number of steiner nodes
    repr = Dict()

    # start with single node
    nclass[1] = 1
    repr[(1, 1)] = Topology(nodes[1:1], Arc[])

    # add more steiner nodes
    for j in 2:(n - 2)
        # look at trees of one size smaller
        for k in 1:nclass[j-1]
            t_base = repr[(j-1, k)]
            g_base = Graph(digraph_from_topology(t_base))

            # and all potential neighbors
            for i in 1:(j-1)
                # maximum degree
                if degree(g_base, i) == 3
                    continue
                end

                # connect new node here then for new tree
                t_new = Topology(nodes[1:j], vcat(t_base.arcs, [Arc(i,j)]))

                # avoid isomorphies
                if any(are_isomorphic(t_new, repr[(j,l)]) for l in 1:nclass[j])
                    continue
                end

                # add to collection
                nclass[j] += 1
                repr[(j, nclass[j])] = t_new
            end
        end
    end
    nclass, repr
end

# needed below
isless(x::T, y::Tuple{T,T}) where {T<:Int} = true
isless(x::Tuple{T,T}, y::T) where {T<:Int} = false

"""
Computes a labeling of a Full Steiner Tree.

This can be used to check for isomorphic trees (having same label).

Following algorithm 4 in Fampa et al.: A specialized branch-and-bound algorithm
for the Euclidean Steiner tree problem in n-space.
"""
function label_fst(fst::Topology)
    n = length(fst.nodes)
    g = Graph(digraph_from_topology(fst))
    terms = [v for v in 1:n if degree(g, v) == 1]
    stein = [v for v in 1:n if degree(g, v) == 3]
    @assert length(terms) + length(stein) == n

    label = Any[-1 for i = 1:n]
    label[terms] = terms

    # rooted binary tree with lexicographical ordering
    root = minimum(terms)
    current_level, next_level = [root], []
    bfs_edges = []
    visited = fill(false, n)
    # BFS
    while length(current_level) > 0
        visited[current_level] .= true
        empty!(next_level)
        for v in current_level
            for w in neighbors(g, v)
                if !visited[w]
                    push!(next_level, w)
                    push!(bfs_edges, (v,w))
                end
            end
        end
        current_level = next_level[:] # copy!
    end

    # work up from the leaves
    reverse!(bfs_edges)
    for depth in 1:2:length(bfs_edges)-1
        t1, h1 = bfs_edges[depth]
        t2, h2 = bfs_edges[depth + 1]
        @assert t1 == t2
        @assert label[t1] == -1
        l1, l2 = label[h1], label[h2]
        label[t1] = l1 < l2 ? (l1, l2) : (l2, l1)
    end

    # all labeled
    @assert all(label .!= -1)

    # label of root's adjacent Steiner node is label for tree
    t, h = bfs_edges[end]
    @assert t == root
    label[h]
end

"""
Enumerate all representative full steiner trees.

Following algorithm 3 in Fampa et al.: A specialized branch-and-bound algorithm
for the Euclidean Steiner tree problem in n-space.
"""
function enumerate_fst(terminals)
    nterms = length(terminals)
    @assert nterms ≥ 3
    nstein = nterms - 2

    stein = [Node(-1,i) for i=1:nstein]
    nodes = vcat(stein, terminals)

    # get all classes for trees of Steiner nodes
    ncls, repr = enumerate_steiner(nterms)
    nclasses = ncls[nstein] # only "full" trees

    # topologies indexed by label
    trees = Dict()

    # early exit for special case of 3 terminals
    if nterms == 3
        tree = Topology(nodes, [Arc(1,2), Arc(1,3), Arc(1,4)])
        trees[label_fst(tree)] = tree
        return trees
    end

    # iterate over all "inner" trees on Steiner nodes
    for c in nclasses
        inner = repr[nstein, c]
        graph = Graph(digraph_from_topology(inner))
        deg1 = [s for s in 1:nstein if degree(graph, s) == 1]
        deg2 = [s for s in 1:nstein if degree(graph, s) == 2]
        deg3 = [s for s in 1:nstein if degree(graph, s) == 3]
        @assert sum(map(length, [deg1, deg2, deg3])) == nstein
        @assert 2 * length(deg1) + length(deg2) == nterms
        K = length(deg1)

        # iterate over permutations of terminals
        termidx = collect((nstein+1):(nstein+nterms))
        for perm in permutations(termidx)
            # TODO: skip those permutations that yield the same tree,
            # i.e. where two terminals are swapped that are connected to
            # the same degree-1 Steiner node

            # connect terminals to Steiner nodes in order
            arcs = inner.arcs[:]
            for k in 1:K
                push!(arcs, Arc(perm[2k - 1], deg1[k]))
                push!(arcs, Arc(perm[2k],     deg1[k]))
            end
            for l in 1:length(deg2)
                push!(arcs, Arc(perm[2K + l], deg2[l]))
            end
            tree = Topology(nodes, arcs)

            # check if new tree is isomorphic to old one
            label = label_fst(tree)
            if haskey(trees, label)
                # TODO: select lexicographically smaller tree
                continue
            end

            trees[label] = tree
        end
    end
    trees
end


text:
### A Pluto.jl notebook ###
# v0.17.2

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 2f0ce4ef-4c45-4175-a9a7-df70cb268629
begin
    CURRENT_DIR = @__DIR__
    ROOT_DIR = if basename(CURRENT_DIR) == "pluto"
        dirname(dirname(CURRENT_DIR))
    else
        CURRENT_DIR
    end

    using Pkg
    Pkg.activate(ROOT_DIR; io = devnull)

    using Logging
    disable_logging(Logging.Warn)

    println('\n', " "^4, "> Loading the packages...")

    using LaTeXStrings
    using Plots
    using PlutoUI

    # Use the GR backend for plots
    gr()

    # Change some of the default parameters for plots
    default(fontfamily = "Computer Modern", dpi = 300, legend = nothing)

    # Define the paths to output directories
    PLOTS_DIR = joinpath(ROOT_DIR, "plots")
    NOTEBOOKS_PLOTS_DIR = joinpath(PLOTS_DIR, "notebooks")
    M2_DIR = joinpath(NOTEBOOKS_PLOTS_DIR, "m2")

    # Make sure the needed directories exist
    mkpath(M2_DIR)

    # An alias for a floating point type
    F = Float64

    md"This is a Pluto notebook to test the parameters of the second model."
end

# ╔═╡ 52342993-f23a-4a50-b3e0-e9e4b032f1cf
md"""
Masses are in units of ``M_0 = 2.325 \times 10^7 M_\odot``, other parameters are in ``\text{kpc}``.

Initial values of the parameters for the bulge, the thin disk and the thick disk are taken from Pouliasis et al. (2017, model I), initial values of the parameters for the halo are taken from Eilers (2018), ranges are taken from Granados et al. (2021).
"""

# ╔═╡ 1fa86bf7-e1dc-4cd1-9ff0-73c14225ef2c
begin
    # Bulge
    M_B_S = @bind M_B NumberField(
        43.0:0.001:2580.0,
        default = 460.0,
    )
    A_B_S = @bind A_B NumberField(0.0:0.001:0.05, default = 0.0)
    B_B_S = @bind B_B NumberField(0.0:0.001:2.0, default = 0.3)

    # Thin disk
    M_THIN_D_S = @bind M_THIN_D NumberField(
        1500.0:0.001:12903.0,
        default = 1700.0,
    )
    A_THIN_D_S = @bind A_THIN_D NumberField(1.0:0.001:20.0, default = 5.3)
    B_THIN_D_S = @bind B_THIN_D NumberField(0.1:0.001:16.0, default = 0.25)

	# Thick disk
	M_THICK_D_S = @bind M_THICK_D NumberField(
        1500.0:0.001:12903.0,
        default = 1700.0,
    )
    A_THICK_D_S = @bind A_THICK_D NumberField(1.0:0.001:20.0, default = 2.6)
    B_THICK_D_S = @bind B_THICK_D NumberField(0.1:0.001:16.0, default = 0.8)

    # Halo
    M_H_S = @bind M_H NumberField(
        43.0:0.001:12903.0,
        default = 4 * π * (1.06 * 1e7) * (14.8)^3 / (2.325 * 1e7),
    )
    A_H_S = @bind A_H NumberField(0.1:0.001:30.0, default = 14.8)

    md"""
    ``M_\text{b}``: $(M_B_S) ``a_\text{b}``: $(A_B_S) ``b_\text{b}``: $(B_B_S) \
    ``M_{\text{thin}}``: $(M_THIN_D_S) ``M_{\text{thick}}``: $(M_THICK_D_S) \
    ``a_{\text{thin}}``: $(A_THIN_D_S) ``a_{\text{thick}}``: $(A_THICK_D_S) \
    ``b_{\text{thin}}``: $(B_THIN_D_S) ``b_{\text{thick}}``: $(B_THICK_D_S) \
    ``M_\text{h}``: $(M_H_S) ``a_\text{h}``: $(A_H_S)
    """
end

# ╔═╡ a5d896ff-570f-49a9-adf0-19f11cb899e0
begin
    "Calculate the value of the R derivative of the Miyamoto & Nagai potential
    ``[ 100 \\, \\text{km} \\, \\text{s}^{-2} ]``"
    function miyamoto_nagai_phi_dr(r, z, m, a, b)::F
        return m * r / (r^2 + (a + √(z^2 + b^2))^2)^(1.5)
    end

    "Calculate the value of the R derivative of the truncated spherical potential
    ``[ 100 \\, \\text{km} \\, \\text{s}^{-2} ]``"
    function truncated_spherical_phi_dr(r, z, m, a)::F
        sq_sum = r^2 + z^2
        return m * r / sq_sum^(1.5) -
		       (r ≥ 100 ? 0 : m * r / sq_sum^(0.49) /
			   a^(2.02) / (sq_sum^(0.51) / a^(1.02) + 1))
    end

    "Calculate the value of the R derivative of the Navarro-Frenk-White potential
    ``[ 100 \\, \\text{km} \\, \\text{s}^{-2} ]``"
    function navarro_frenk_white_phi_dr(r, z, m, a)::F
        sq_sum = r^2 + z^2
        k_1 = √(sq_sum) / a + 1.0
        return m * r * log(k_1) / sq_sum^(1.5) - m * r / a / sq_sum / k_1
    end

    "Calculate the value of circular velocity according to
    ``v_c^2(r) = R \\left. \\frac{\\partial \\Phi(R, Z)}{\\partial R} \\right|_{Z=0}``
    ``[ \\text{km}^2 \\, \\text{s}^{-2} ]``"
    circ_vel(r) = √(r * phi_dr(r, 0) * 100)

    "Calculate the value of ``\\partial \\Phi(R, Z) / \\partial R``
    ``[ 100 \\, \\text{km} \\, \\text{s}^{-2} ]``"
    function phi_dr(r, z)::F
        return miyamoto_nagai_phi_dr(r, z, M_B, A_B, B_B) +
               miyamoto_nagai_phi_dr(r, z, M_THIN_D, A_THIN_D, B_THIN_D) +
		       miyamoto_nagai_phi_dr(r, z, M_THICK_D, A_THICK_D, B_THICK_D) +
		       navarro_frenk_white_phi_dr(r, z, M_H, A_H)
    end

    println(" "^4, "> Plotting the rotation curve...")

    # Plot the rotation curve
    p1 = plot(
        circ_vel, 0.2, 25;
        title = "Rotation curve",
        xlabel = L"R \;\, [\mathrm{kpc}]",
        ylabel = L"v_c \; [\mathrm{km \, s^{-2}}]",
        xminorticks = 8,
        yminorticks = 8
    )

    # Save the figure
    savefig(p1, joinpath(M2_DIR, "Rotation curve.pdf"))
    savefig(p1, joinpath(M2_DIR, "Rotation curve.png"))

    p1
end

# ╔═╡ d4ba276c-a578-4a64-9f54-c33c81225a0b
begin
    # Plot the rotation curve
    p2 = plot(
        circ_vel, 0.2, 200;
        title = "Rotation curve",
        xlabel = L"R \;\, [\mathrm{kpc}]",
        ylabel = L"v_c \; [\mathrm{km \, s^{-1}}]",
        xticks = [1e-1, 1, 1e1, 1e2],
        xlims = (0.1, 390),
        xscale = :log10,
        xminorticks = 9,
        yminorticks = 8
    )

    # Save the figure
    savefig(p2, joinpath(M2_DIR, "Rotation curve (log).pdf"))
    savefig(p2, joinpath(M2_DIR, "Rotation curve (log).png"))

    println()

    p2
end

# ╔═╡ bb727eda-9d7c-48e2-bac6-23e90aa6aab1
md"""
From Eilers (2018):

$(PlutoUI.LocalResource(joinpath(ROOT_DIR, "assets", "Rotation curve of the second model (Eilers, 2018).png")))
"""

# ╔═╡ Cell order:
# ╟─2f0ce4ef-4c45-4175-a9a7-df70cb268629
# ╟─52342993-f23a-4a50-b3e0-e9e4b032f1cf
# ╠═1fa86bf7-e1dc-4cd1-9ff0-73c14225ef2c
# ╟─a5d896ff-570f-49a9-adf0-19f11cb899e0
# ╟─d4ba276c-a578-4a64-9f54-c33c81225a0b
# ╟─bb727eda-9d7c-48e2-bac6-23e90aa6aab1


text:
using SampleJuliaPackage
using Test

@testset "SampleJuliaPackage.jl" begin
    # Write your own tests here.
    @test test0([1,2,3], [4,5,6]) == 5
end


text:
module AdventOfCode

include("utils.jl")
# include("day1.jl")
include("day2.jl")
include("day3.jl")
include("day4.jl")
include("day5.jl")
include("day6.jl")
include("day7.jl")
include("day8.jl")
include("day9.jl")
include("day10.jl")
include("day11.jl")
include("day12.jl")
include("day13.jl")
include("day14.jl")
include("day15.jl")
include("day17.jl")

end


text:
#=
Created on Saturday 22 Feb 2020
Last update: -

@author: Michiel Stock
michielfmstock@gmail.com

Example plots convergence unconstrained convex optimization.
=#

using STMO
using Plots, LaTeXStrings
using LinearAlgebra
using STMO.TestFuns
using STMO.Unconstrained

# starting point
x0q = [9.0, 2.0];
x0nq = [-1.0, 0.75];

# GRADIENT DESCENT
trackerqgd = PathTrack(x0q);
xstarq = Unconstrained.gradient_descent(fquadr, copy(x0q), grad_fquadr, tracker=trackerqgd);

trackernqgd = PathTrack(x0nq)
xstarnq = gradient_descent(fnonquadr, copy(x0nq), grad_fnonquadr,
                tracker=trackernqgd);

# COORDINATE DESCENT

trackerqcd = PathTrack(x0q);
xstarq = Unconstrained.coordinate_descent(fquadr, copy(x0q), grad_fquadr,
                tracker=trackerqcd);

trackernqcd = PathTrack(x0nq);
xstarq = Unconstrained.coordinate_descent(fnonquadr, copy(x0nq), grad_fnonquadr,
            tracker=trackernqcd);

# NEWTON

trackerqnm = PathTrack(x0q)
xstarq = Unconstrained.newtons_method(fquadr, copy(x0q), grad_fquadr, hess_fquadr, tracker=trackerqnm)

trackernqnm = PathTrack(x0nq)
xstarq = Unconstrained.newtons_method(fnonquadr, copy(x0nq), grad_fnonquadr, hess_fnonquadr, tracker=trackernqnm)

# PLOTTING

pq = contour(-10:0.1:10, -5:0.1:5, (x1, x2) -> fquadr((x1, x2)), xlabel="\$ x_1 \$",
                ylabel="\$ x_2 \$", title="quadratic", fill=false)
path!(trackerqgd, label="GD")
path!(trackerqcd, label="CD", color=myblue)
path!(trackerqnm, label="Newton", color=mygreen)

savefig(pq, "figures/path_quadr.png")

pnq = contour(-2:0.1:2, -1:0.1:1, (x1, x2) -> fnonquadr((x1, x2)), xlabel="\$ x_1 \$",
                ylabel="\$ x_2 \$", title="non-quadratic", fill=false)
path!(trackernqgd, label="GD")
path!(trackernqcd, label="CD", color=myblue)
path!(trackernqnm, label="Newton", color=mygreen)
savefig(pnq, "figures/path_nonquadr.png")

# convergence

fnqexact = newtons_method(fnonquadr, copy(x0nq), grad_fnonquadr, hess_fnonquadr, ϵ=1e-12) |> fnonquadr
nonquadrerr = x -> abs(fnonquadr(x) .- fnqexact);

p1 = plotobj(fquadr, trackerqgd, label="GD", yscale=:log10, ylabel="\$ f(x) - f(x^\\star)\$", title="quadratic", ylim=(1e-15, 10));
plotobj!(fquadr, trackerqcd, color=myblue, label="CD");
plotobj!(fquadr, trackerqnm, color=mygreen, label="Newton");
p2 = plotobj(nonquadrerr, trackernqgd, label="GD", yscale=:log10, ylabel="\$ f(x) - f(x^\\star)\$", title="non-quadratic", ylim=(1e-15, 10));
plotobj!(nonquadrerr, trackernqcd, color=myblue, label="CD");
plotobj!(fquadr, trackernqnm, color=mygreen, label="Newton");
pconv = plot(p1, p2);
savefig(pconv, "figures/unconstrained_conv.png")


text:
using Oceananigans.Grids: unpack_grid
using Oceananigans.BoundaryConditions: PBC, NFBC

const PeriodicBC = PBC
const NoFluxBC = NFBC

function generate_discrete_eigenvalues(grid, pressure_bcs)
    kx² = λi(grid, pressure_bcs.x.left)
    ky² = λj(grid, pressure_bcs.y.left)
    kz² = λk(grid, pressure_bcs.z.left)
    return kx², ky², kz²
end

"""
    λi(grid, ::PeriodicBC)

Return an Nx×1×1 array of eigenvalues satisfying the discrete form of Poisson's
equation with periodic boundary conditions in the x-dimension on `grid`.
"""
function λi(grid, ::PeriodicBC)
    Nx, Ny, Nz, Lx, Ly, Lz = unpack_grid(grid)
    is = reshape(1:Nx, Nx, 1, 1)
    return @. (2sin((is - 1) * π / Nx) / (Lx / Nx))^2
end

"""
    λi(grid, ::NoFluxBC)

Return an Nx×1×1 array of eigenvalues satisfying the discrete form of Poisson's
equation with staggered Neumann boundary conditions in the x-dimension on `grid`.
"""
function λi(grid, ::NoFluxBC)
    Nx, Ny, Nz, Lx, Ly, Lz = unpack_grid(grid)
    is = reshape(1:Nx, Nx, 1, 1)
    return @. (2sin((is - 1) * π / 2Nx) / (Lx / Nx))^2
end

"""
    λj(grid, ::PeriodicBC)

Return an 1×Ny×1 array of eigenvalues satisfying the discrete form of Poisson's
equation with periodic boundary conditions in the y-dimension on `grid`.
"""
function λj(grid, ::PeriodicBC)
    Nx, Ny, Nz, Lx, Ly, Lz = unpack_grid(grid)
    js = reshape(1:Ny, 1, Ny, 1)
    return @. (2sin((js - 1) * π / Ny) / (Ly / Ny))^2
end

"""
    λj(grid, ::NoFluxBC)

Return an 1×Ny×1 array of eigenvalues satisfying the discrete form of Poisson's
equation with staggered Neumann boundary conditions in the y-dimension on `grid`.
"""
function λj(grid, ::NoFluxBC)
    Nx, Ny, Nz, Lx, Ly, Lz = unpack_grid(grid)
    js = reshape(1:Ny, 1, Ny, 1)
    return @. (2sin((js - 1) * π / 2Ny) / (Ly / Ny))^2
end

"""
    λk(grid, ::PeriodicBC)

Return an 1×1×Nz array of eigenvalues satisfying the discrete form of Poisson's
equation with periodic boundary conditions in the z-dimension on `grid`.
"""
function λk(grid, ::PeriodicBC)
    Nx, Ny, Nz, Lx, Ly, Lz = unpack_grid(grid)
    ks = reshape(1:Nz, 1, 1, Nz)
    return @. (2sin((ks - 1) * π / 2Nz) / (Lz / Nz))^2
end

"""
    λk(grid, ::NoFluxBC)

Return an 1×1×Nz array of eigenvalues satisfying the discrete form of Poisson's
equation with staggered Neumann boundary conditions in the z-dimension on `grid`.
"""
function λk(grid, ::NoFluxBC)
    Nx, Ny, Nz, Lx, Ly, Lz = unpack_grid(grid)
    ks = reshape(1:Nz, 1, 1, Nz)
    return @. (2sin((ks - 1) * π / 2Nz) / (Lz / Nz))^2
end


text:
import MacroTools
using MacroTools: postwalk, @capture

# symbols that should not be interpreted as variable names - must be a better way to do this
const RESERVED_SYMBOLS = Set([:+, :*, :/, :-, :^, :(==), :(!=), :<, :>, :(<=), :(>=), :≤, :≥])

iscoef(coef::Symbol) = startswith(String(coef), r"[βα]")
getcoef(coef::QuoteNode) = getcoef(coef.value)
getcoef(coef::Symbol) = iscoef(coef) ? coef : error("unable to parse coef $coef")

#=
define a utility function like so:
    @utility begin
        outcome ~ :asc_1 + :b_car * car
        :asc_1 = 2
        :b_car = 3f
    end

Symbols are treated as coefficients. Starting values for coefficients are defined with equal signs. Fixed values should be postfixed with f.
Coefficients that are not defined will be treated as estimated with a starting value of 0
=#
macro utility(ex::Expr)
    # inner macros are not parsed first, so expand all of them (e.g. @β)
    ex = macroexpand(Main, ex)

    # first, extract coefficient definitions
    # turn coefs into a vector, reference them by integer values
    coefnames = Vector{Symbol}()
    starting_values = Vector{Float64}()
    coef_indices = Dict{Symbol, Int64}()
    fixed_coefs = Dict{Symbol, Number}()
    postwalk(ex) do subex
        if @capture(subex, (coefnode_ = (starting_val_, fixed)))
            coef = getcoef(coefnode)
            !(starting_val isa Number) && error("Fixed starting value must be a number, for coef $coef")
            # fixed coefficient
            (haskey(coef_indices, coef) || haskey(fixed_coefs, coef)) && error("Coef $coef defined multiple times")
            fixed_coefs[coef] = starting_val
        elseif @capture(subex, coefnode_ = starting_val_)
            coef = getcoef(coefnode)
            !(starting_val isa Number) && error("Starting value must be a number, for coef $coef")
            # non-fixed coefficient
            (haskey(coef_indices, coef) || haskey(fixed_coefs, coef)) && error("Coef $coef defined multiple times")
            push!(coefnames, coef)
            push!(starting_values, convert(Float64, starting_val))
            coef_indices[coef] = length(coefnames)
        end

        return subex  # make sure we don't mangle the expression while we're iterating over it
    end

    # now find all utility functions, and number alternatives
    util_funcs = Vector{Expr}()
    alt_numbers = Dict{Any, Int64}()
    columns = Set{Any}()
    postwalk(ex) do subex
        if @capture(subex, lhs_ ~ rhs_)
            # turn the right hand side into a function of the params
            parsed_rhs = postwalk(rhs) do x
                if x isa Symbol
                    if iscoef(x)
                        # convert :coef to coefs - either references into the params array for
                        # non-fixed coefs, or literal values for fixed coefs
                        if haskey(fixed_coefs, x)
                            # interpolate in literal value
                            return :($(fixed_coefs[x]))
                        elseif haskey(coef_indices, x)
                            # interpolate in reference into params array
                            return :(params[$(coef_indices[x])])
                        else
                            # coef not specifically defined, create it implicitly with starting value 0
                            push!(coefnames, x)
                            push!(starting_values, zero(Float64))
                            coef_indices[x] = length(coefnames)
                            return :(params[$(length(coefnames))])
                        end
                    elseif Base.isoperator(x)
                        # TODO would be useful to also allow things like trig functions here. Maybe we need
                        # something like I(...) like they have in R.
                        return x
                    else
                        # convert bare values to data values
                        push!(columns, x)
                        return :(row.$x)
                    end
                else
                    return x
                end
            end

            # turn alternatives into numbers for processing speed
            push!(util_funcs, :(function (params::Vector{T}, row) where T <: Number
                # ensure the return value is always a T, even when the function is e.g. a constant 0
                convert(T, @inbounds($parsed_rhs))
            end))
            alt_numbers[lhs] = length(util_funcs)
        end

        return subex
    end

    # TODO some kind of error checking that there aren't other expressions that people meant to be
    # utility function or coefficient definitions, but didn't parse correctly

    return quote
        (
            coefnames = $coefnames,
            starting_values = $starting_values,
            fixed_coefs = $fixed_coefs,
            utility_functions = [$(util_funcs...)],
            alt_numbers = $alt_numbers,
            columnnames=$columns
        )
    end
end

"""
Create betas for all of the names in the vector
"""
macro β(outcomenode, variablenode, includealpha=false)
    outcomes = eval(outcomenode)
    vars = eval(variablenode)

    res = Vector{Expr}()
    for outcome in outcomes
        full, short = if outcome isa Pair
            outcome[1], outcome[2]
        else
            outcome, outcome
        end

        expressions = map(vars) do varn
            β = (Symbol("β$(short)_$varn"))
            x = Symbol(varn)
            :($β * $x)
        end

        if includealpha
            expressions = convert(Vector{Union{Expr, Symbol}}, expressions)
            pushfirst!(expressions, Symbol("α$short"))
        end

        rhs = if length(expressions) == 0
            quote end
        elseif length(expressions) == 1
            expressions[1]
        else
            Expr(:call, :+, expressions...)
        end

        push!(res, :($full ~ $rhs))
    end

    return esc(quote $(res...) end)
end

macro dummy_code(prefix, column, valuenode)
    values = eval(valuenode)
    expressions = map(values) do val
        β = (Symbol("$(prefix)_$(column)_$(val)"))
        :($β * ($column == $val))
    end

    if length(expressions) == 0
        quote end
    elseif length(expressions) == 1
        esc(expressions[1])
    else
        esc(Expr(:call, :+, expressions...))
    end
end


text:
using HidApi
using Test

@testset "HidApi.jl" begin

    # just a smoke test, as CI has not a known USB target to test
    @test enumerate_devices() isa Vector{HidDevice}
    @test nothing === find_device(0x0000,0x0000)
    @test shutdown() === nothing
end