text:
"""
    mpifit_psvgp(get_data, n_parts, n_dims, bounds_low, bounds_high;
    ni     = 5,
    batch_size = 20,
    n_iters    = 200,
    frac_local = 0.75)

Fit PSVGP to data for distributed computation

    get_data:                function to define how a local partition should read in its data
    ni:                      number of local inducing points
    batch_size:              size of batches for stochastic optimizer
    n_iters:                 number of optimization steps
    frac_local:              what fraction of steps should a cell query neighbors

WARNING: Currently writes out several jld2 files containing local GPs
"""
function mpifit_psvgp(get_data, n_parts, n_dims, bounds_low, bounds_high;
        ni         = 5,               # Number of input dimensions
        batch_size = 20,              # Number of points used in updating local ELBO
        n_iters    = 2000,
        frac_local = 0.75
        )

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

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

    # Fixing this for Gaussian case for now.
    #     Need to generalize this for the future but that will come in part of a large generalization of the code
    n_latent = 1

    # Since Julia is 1-indexed and not 0, making variable for array indexing
    rank_ind = myrank + 1

    grad_clip = 1.e-3

    part_split = [[ii] for ii in 1:mysize]
    part_nbors = [[[ii - 1, ii + 1, ii + 20, ii - 20]] for ii in 1:mysize]

    for ii in 1:mysize
        if mod(ii,20) == 1
            filter!(aa -> aa != (ii - 1), part_nbors[ii][1])
        elseif mod(ii,20) == 0
            filter!(aa -> aa != (ii + 1), part_nbors[ii][1])
        end
        if ii <= 20
            filter!(aa -> aa != (ii - 20), part_nbors[ii][1])
        elseif ii + 20 > 400
            filter!(aa -> aa != (ii + 20), part_nbors[ii][1])
        end
    end

    if n_parts > mysize
        for ii = (mysize+1):n_parts
            nbor_inds = [ii - 1, ii + 1, ii + 20, ii - 20]
            if mod(ii,20) == 1
                filter!(aa -> aa != (ii - 1), nbor_inds)
            elseif mod(ii,20) == 0
                filter!(aa -> aa != (ii + 1), nbor_inds)
            end
            if ii <= 20
                filter!(aa -> aa != (ii - 20), nbor_inds)
            elseif ii + 20 > 400
                filter!(aa -> aa != (ii + 20), nbor_inds)
            end
            append!(part_split[1 + ((ii-1) % mysize)], ii)
            append!(part_nbors[1 + ((ii-1) % mysize)], [nbor_inds])
        end
    end

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

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

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

    # Initializing "trace" array to store ELBO values
    trace_l = [zeros(n_iters) for ii in 1:len] # To collect the local ELBO value

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

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

    @time begin
        local_svgps = [Inference_obj(data_array[ii][1], float(data_array[ii][2]), ni) for ii in 1:len]

        # Step 3
        # Initialize local optimizers
        lopt_cmean = [[Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075)) for jj in 1:n_latent] for ii in 1:len]
        lopt_lrho  = [[Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075)) for jj in 1:n_latent] for ii in 1:len]
        lopt_lkap  = [[Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075)) for jj in 1:n_latent] for ii in 1:len]
        lopt_lsig  = [[Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075)) for jj in 1:n_latent] for ii in 1:len]
        loptm      = [[Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.01)  ) for jj in 1:n_latent] for ii in 1:len]
        loptS      = [[Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.002) ) for jj in 1:n_latent] for ii in 1:len]
        loptx      = [[Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.01)  ) for jj in 1:n_latent] for ii in 1:len]

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

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

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

            for ii in 1:len
                # Step 4
                # Check to see if a neighbor needs data from node
                is_request, recv_id, tag_ind = check_for_query(comm)
                while is_request
                    dummy   = [0]
                    # Receive message to clear out
                    rmsg    = MPI.Recv!(dummy, recv_id, tag_ind, comm) #TODO: Do this better. Shouldn't need to send dummy data.

                    inds    = rand(1:length(local_svgps[tag_ind].data.y), batch_size)
                    send_x  = local_svgps[tag_ind].data.x[inds,:]
                    send_y  = local_svgps[tag_ind].data.y[inds]
                    smsg    = MPI.Send(send_x, recv_id, length(part_split[recv_id+1])+1, comm)
                    smsg    = MPI.Send(send_y, recv_id, length(part_split[recv_id+1])+2, comm)
                    is_request, recv_id, tag_ind = check_for_query(comm)
                end

                # Step 5
                # Select either local data or neighbor data
                select_local = rand(1)[1] < frac_local
                if select_local
                    inds    = rand(1:size(local_svgps[ii].data.y)[1], batch_size)
                    batch_x = local_svgps[ii].data.x[inds,:]
                    batch_y = local_svgps[ii].data.y[inds]
                else
                    nbor = rand(part_nbors[rank_ind][ii])

                    if nbor in part_split[rank_ind]
                        part_ind = filter(aa -> part_split[rank_ind][aa] == nbor, 1:size(part_split[rank_ind])[1])[1]

                        inds     = rand(1:size(local_svgps[part_ind].data.y)[1], batch_size)
                        batch_x  = local_svgps[part_ind].data.x[inds,:]
                        batch_y  = local_svgps[part_ind].data.y[inds]
                    else
                        nbor_ind = filter(aa -> nbor in part_split[aa], 1:size(part_split)[1])[1]
                        part_ind = filter(aa -> part_split[nbor_ind][aa] == nbor, 1:size(part_split[nbor_ind])[1])[1]
                        # Query neighbor for data
                        smsg     = MPI.Send([0], nbor_ind - 1, part_ind, comm)
                        got_data = false
                        while !got_data
                            is_request, recv_id, tag_ind = check_for_query(comm)
                            if is_request
                                if tag_ind < len+1 # if < len+1 then tag is a partition label
                                    dummy   = [0]
                                    # Receive message to clear out
                                    rmsg    = MPI.Recv!(dummy, recv_id, tag_ind, comm) #TODO: Do this better. Shouldn't need to send dummy data.
                                    inds    = rand(1:size(local_svgps[tag_ind].data.y)[1], batch_size)
                                    send_x  = local_svgps[tag_ind].data.x[inds,:]
                                    send_y  = local_svgps[tag_ind].data.y[inds]
                                    smsg    = MPI.Send(send_x, recv_id, length(part_split[recv_id+1])+1, comm)
                                    smsg    = MPI.Send(send_y, recv_id, length(part_split[recv_id+1])+2, comm)
                                end
                                if tag_ind > len
                                    batch_x  = zeros(batch_size, n_dims)
                                    rmsg     = MPI.Recv!(batch_x, recv_id, len+1, comm)

                                    batch_y  = zeros(batch_size)
                                    rmsg     = MPI.Recv!(batch_y, recv_id, len+2, comm)
                                    got_data = true
                                end # tag_ind > len
                            end # is_request
                        end # !got_data
                    end # nbor in part_split[rank_ind]
                end # select_local

                grads  = gradient(gps -> inference_elbo(
                        local_svgps[ii].data.x[inds,:],
                        local_svgps[ii].data.y[inds],
                        local_svgps[ii].data.n,
                        gps
                        ), local_svgps[ii])[1]

                for jj in 1:n_latent
                    Flux.Optimise.update!(lopt_cmean[ii][jj], local_svgps[ii].params[jj].const_mean,    -grads.params[jj][].const_mean)
                    Flux.Optimise.update!(lopt_lrho[ii][jj],  local_svgps[ii].params[jj].log_rho,       -grads.params[jj][].log_rho)
                    Flux.Optimise.update!(lopt_lkap[ii][jj],  local_svgps[ii].params[jj].log_kappa,     -grads.params[jj][].log_kappa)
                    Flux.Optimise.update!(lopt_lsig[ii][jj],  local_svgps[ii].params[jj].log_sigma,     -grads.params[jj][].log_sigma)
                    Flux.Optimise.update!(loptm[ii][jj],      local_svgps[ii].params[jj].inducing_mean, -grads.params[jj][].inducing_mean)
                    Flux.Optimise.update!(loptS[ii][jj],      local_svgps[ii].params[jj].inducing_L,    -LowerTriangular(grads.params[jj][].inducing_L) )
                    Flux.Optimise.update!(loptx[ii][jj],      local_svgps[ii].params[jj].inducing_locs, -grads.params[jj][].inducing_locs)
                end
                # end
                # Collecting local ELBO in an array ###########################
                trace_l[ii][iter] = inference_elbo(
                    batch_x,
                    batch_y,
                    local_svgps[ii].data.n,
                    local_svgps[ii]
                )
                # println("local ELBO optimized")
            end

        end # outeriter

        barrier_req = MPI.Ibarrier(comm)

        all_done = false
        while !all_done
            is_request, recv_id, tag_ind = check_for_query(comm)
            if is_request
                if tag_ind < len+1 # if < len+1 then tag is a partition label
                    dummy   = [0]
                    # Receive message to clear out
                    rmsg    = MPI.Recv!(dummy, recv_id, tag_ind, comm) #TODO: Do this better. Shouldn't need to send dummy data.
                    inds    = rand(1:size(local_svgps[tag_ind].data.y)[1], batch_size)
                    send_x  = local_svgps[tag_ind].data.x[inds,:]
                    send_y  = local_svgps[tag_ind].data.y[inds]
                    smsg    = MPI.Send(send_x, recv_id, length(part_split[recv_id+1])+1, comm)
                    smsg    = MPI.Send(send_y, recv_id, length(part_split[recv_id+1])+2, comm)
                end
                if tag_ind > len
                    println("Shouldnt get a message sending data! Something is wrong!")
                end # tag_ind > len
            end # is_request
            all_done, barrier_status = MPI.Test!(barrier_req)
        end # !all_done
    end

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

    store_trace = zeros(n_iters, n_parts)

    # Collect all traces to write out at once
    if myrank != 0
        for ii in 1:len
            send_trace(trace_l[ii], 0 ,comm)
        end

    end #myrank != 0

    if myrank == 0
        for ii in 1:len
            store_trace[:, part_split[1][ii] ] = trace_l[ii]
        end

        for proc = 2:mysize
            for ii in 1:size(part_split[proc])[1]
                trace_l = receive_trace(n_iters, proc - 1, comm)
                store_trace[:, part_split[proc][ii] ] = trace_l
            end
        end

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

    end # myrank == 0
    ############################################################################

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

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

    MPI.Finalize()

end     # mpifit_HSVGP

max_stars_repo_name: ayanbiswas/PRISM

max_issues_repo_name: ayanbiswas/PRISM

max_forks_repo_name: ayanbiswas/PRISM


text:
"""
    SerialExecution

Functor object for using serial execution in a lshaped algorithm. Create by supplying a [`Serial`](@ref) object through `execution` in the `LShapedSolver` factory function and then pass to a `StochasticPrograms.jl` model.

"""
struct SerialExecution{F <: AbstractFeasibility,
                       T <: AbstractFloat,
                       A <: AbstractVector,
                       S <: LQSolver} <: AbstractExecution
    subproblems::Vector{SubProblem{F,T,A,S}}
    subobjectives::A
    model_objectives::A

    function SerialExecution(::Type{F}, ::Type{T}, ::Type{A}, ::Type{S}) where {F <: AbstractFeasibility, T <: AbstractFloat, A <: AbstractVector, S <: LQSolver}
        return new{F,T,A,S}(Vector{SubProblem{F,T,A,S}}(), A(), A())
    end
end

function nthetas(lshaped::AbstractLShapedSolver, ::SerialExecution)
    return nthetas(lshaped.nscenarios, lshaped.aggregation)
end

function initialize_subproblems!(execution::SerialExecution{F,T,A,S},
                                 scenarioproblems::ScenarioProblems,
                                 x::AbstractVector,
                                 subsolver::MPB.AbstractMathProgSolver) where {F <: AbstractFeasibility,
                                                                               T <: AbstractFloat,
                                                                               A <: AbstractVector,
                                                                               S <: LQSolver}
    for i = 1:StochasticPrograms.nscenarios(scenarioproblems)
        m = subproblem(scenarioproblems, i)
        y₀ = convert(A, rand(m.numCols))
        push!(execution.subproblems, SubProblem(m,
                                                parentmodel(scenarioproblems),
                                                i,
                                                probability(scenario(scenarioproblems,i)),
                                                copy(x),
                                                y₀,
                                                subsolver,
                                                F))
    end
    return nothing
end

function initialize_subproblems!(execution::SerialExecution{F,T,A,S},
                                 scenarioproblems::DScenarioProblems,
                                 x::AbstractVector,
                                 subsolver::MPB.AbstractMathProgSolver) where {F <: AbstractFeasibility,
                                                                               T <: AbstractFloat,
                                                                               A <: AbstractVector,
                                                                               S <: LQSolver}
    for i = 1:StochasticPrograms.nscenarios(scenarioproblems)
        m = subproblem(scenarioproblems, i)
        y₀ = convert(A, rand(m.numCols))
        push!(execution.subproblems,SubProblem(m,
                                               i,
                                               probability(scenario(scenarioproblems,i)),
                                               copy(x),
                                               y₀,
                                               masterterms(scenarioproblems,i),
                                               subsolver,
                                               F))
    end
    return nothing
end

function finish_initilization!(lshaped::AbstractLShapedSolver, execution::SerialExecution)
    append!(execution.subobjectives, fill(1e10, nthetas(lshaped)))
    append!(execution.model_objectives, fill(-1e10, nthetas(lshaped)))
    return nothing
end

function resolve_subproblems!(lshaped::AbstractLShapedSolver, execution::SerialExecution{F,T}) where {F <: AbstractFeasibility, T <: AbstractFloat}
    # Update subproblems
    update_subproblems!(execution.subproblems, lshaped.x)
    # Assume no cuts are added
    added = false
    # Solve subproblems
    for subproblem ∈ execution.subproblems
        cut::SparseHyperPlane{T} = subproblem()
        added |= aggregate_cut!(lshaped, lshaped.aggregation, cut)
    end
    added |= flush!(lshaped, lshaped.aggregation)
    # Return current objective value and cut_added flag
    return current_objective_value(lshaped), added
end

function calculate_objective_value(lshaped::AbstractLShapedSolver, execution::SerialExecution)
    return get_obj(lshaped)⋅decision(lshaped) + sum([subproblem.π*subproblem(decision(lshaped)) for subproblem in execution.subproblems])
end

function fill_submodels!(lshaped::AbstractLShapedSolver, scenarioproblems::ScenarioProblems, execution::SerialExecution)
    for (i, submodel) in enumerate(scenarioproblems.problems)
        execution.subproblems[i](decision(lshaped))
        fill_submodel!(submodel, execution.subproblems[i])
    end
    return nothing
end

function fill_submodels!(lshaped::AbstractLShapedSolver, scenarioproblems::DScenarioProblems, execution::SerialExecution)
    j = 0
    @sync begin
        for w in workers()
            n = remotecall_fetch((sp)->length(fetch(sp).problems), w, scenarioproblems[w-1])
            for i in 1:n
                k = i+j
                execution.subproblems[k](decision(lshaped))
                @async remotecall_fetch((sp,i,x,μ,λ,C) -> fill_submodel!(fetch(sp).problems[i],x,μ,λ,C),
                                        w,
                                        scenarioproblems[w-1],
                                        i,
                                        get_solution(execution.subproblems[k])...)
            end
            j += n
        end
    end
    return nothing
end

# API
# ------------------------------------------------------------
function (execution::Serial)(::Integer, ::Type{F}, ::Type{T}, ::Type{A}, ::Type{S}) where {F <: AbstractFeasibility, T <: AbstractFloat, A <: AbstractVector, S <: LQSolver}
    return SerialExecution(F, T, A, S)
end

function str(::Serial)
    return ""
end

ext: jl


text:
##############################################################################################
module SingleExperimentDriver
##############################################################################################
##############################################################################################
# imports and exports
using ..FilterExps, ..SmootherExps, ..GenerateTimeSeries, JLD2
export filter_state_exp, filter_param_exp, classic_smoother_state_exp,
       classic_smoother_param_exp, single_iteration_smoother_state_exp,
       single_iteration_smoother_param_exp, iterative_smoother_state_exp,
       iterative_smoother_param_exp, L96_time_series_exp,
       IEEE39bus_time_series_exp, L96_time_series_test, IEEE39bus_time_series_test

##############################################################################################
# Description
##############################################################################################
# The following are used as wrappers for the experiments to work with standard inputs
# for debugging, benchmarking and profiling.  These will work simply with the standard macros,
# while the parallel submission scripts are used for performance on servers
##############################################################################################
##############################################################################################
##############################################################################################
## Generate time series data
##############################################################################################
##############################################################################################

##############################################################################################
## Lorenz-96 model
##############################################################################################
# Lorenz-96(-s) truth twin data generated as a single function call, arguments are
# seed, state_dim, tanl, nanl, spin, diffusion, F = args

function L96_time_series_exp()
    args = (0, 40, 0.05, 50000, 5000, 0.00, 8.0)
    L96_time_series(args)
end

##############################################################################################
## IEEE 39 bus test case model
##############################################################################################
# IEEE 39 bus test case truth twin data generated as a single function call, arguments are
# seed, tanl, nanl, spin, diffusion = args

function IEEE39bus_time_series_exp()
    args = (0, 0.01, 50000, 5000, 0.0)
    IEEE39bus_time_series(args)
end

##############################################################################################
##############################################################################################
## DA experiments begin here
##############################################################################################
##############################################################################################

##############################################################################################
## Time series data
##############################################################################################
# Observation timeseries to load into DA experiments as a truth twin, without
# needing to generate the timeseries with the model twin.

# path to time series directory
path = joinpath(@__DIR__, "../data/time_series/")

# file names
#fname = "IEEE39bus_time_series_seed_0000_diff_0.000_tanl_0.01_nanl_50000_spin_5000_" *
#        "h_0.010.jld2"
#fname = "IEEE39bus_time_series_seed_0000_diff_0.012_tanl_0.01_nanl_50000_spin_5000_" *
#        "h_0.010.jld2"
fname = "L96_time_series_seed_0000_dim_40_diff_0.000_F_08.0_tanl_0.05_nanl_50000_" *
        "spin_5000_h_0.050.jld2"
#fname = "L96_time_series_seed_0000_dim_40_diff_0.000_F_08.0_tanl_0.10_nanl_50000_" *
#        "spin_5000_h_0.050.jld2"
#fname = "L96_time_series_seed_0000_dim_40_diff_0.100_F_08.0_tanl_0.05_nanl_50000_" *
#         "spin_5000_h_0.005.jld2"
#fname = "L96_time_series_seed_0000_dim_40_diff_0.100_F_08.0_tanl_0.10_nanl_50000_" *
#        "spin_5000_h_0.005.jld2"

# load the file name with the path
time_series = path * fname

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

##############################################################################################
# Filters
##############################################################################################
# filter state estimation, arguments are
# time_series, scheme, seed, nanl, obs_un, obs_dim, γ, N_ens, infl = args

function filter_state_exp()
    args = (time_series, "etkf", 0, 2500, 1.0, 40, 1.00, 21, 1.02)
    filter_state(args)
end

##############################################################################################
# filter joint state-parameter estimation, arguments are
# time_series, scheme, seed, nanl, obs_un, obs_dim, γ, param_err, param_wlk, N_ens,
# state_infl, param_infl = args

function filter_param_exp()
    args = (time_series, "etkf", 0, 2500, 1.0, 40, 1.0, 0.10, 0.0010, 21, 1.02, 1.0)
    filter_param(args)
end

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

##############################################################################################
# Classic smoothers
##############################################################################################
# classic EnKS style smoothing for state estimation, arguments are
# time_series, method, seed, nanl, lag, shift, obs_un, obs_dim, γ, N_ens, infl = args

function classic_smoother_state_exp()
    args = (time_series, "etks", 0, 2500, 10, 1, 1.0, 40, 1.0, 21, 1.02)
    classic_state(args)
end

##############################################################################################
# classic EnKS style smoothing for joint state-parameter estimation, arguments are
# time_series, method, seed, nanl, lag, shift, obs_un, obs_dim, γ,
# param_err, param_wlk, N_ens, state_infl, param_infl = args

function classic_smoother_param_exp()
    args = (time_series, "etks", 0, 2500, 10, 1, 1.0, 40, 1.0, 0.03, 0.001, 21, 1.02, 1.0)
    classic_param(args)
end

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

##############################################################################################
# Single iteration smoothers
##############################################################################################
# Single iteration ensemble Kalman smoothing for state estimation, arguments are
# time_series, method, seed, nanl, lag, shift, mda, obs_un, obs_dim, γ, N_ens, infl = args

function single_iteration_smoother_state_exp()
    args = (time_series, "etks", 0, 2500, 10, 1, false, 0.1, 20, 1.0, 21, 1.02)
    single_iteration_state(args)
end

##############################################################################################
# Single iteration ensemble Kalman smoothing for joint state-parameter estimation,
# arguments are
# time_series, method, seed, nanl, lag, shift, mda, obs_un, obs_dim, γ,
# param_err, param_wlk, N_ens, state_infl, param_infl = args

function single_iteration_smoother_param_exp()
    args = (time_series,
            "etks", 0, 2500, 10, 1, false, 0.1, 20, 1.0, 0.03, 0.0010, 21, 1.02, 1.00)
    single_iteration_param(args)
end

##############################################################################################
# Iterative smoothers
##############################################################################################
# IEnKS style iterative smoothing for state estmation, arugments are
# time_series, method, seed, nanl, lag, shift, mda, obs_un, obs_dim, γ, N_ens, infl = args

function iterative_smoother_state_exp()
    args = (time_series,
            "ienks-transform", 0, 2500, 10, 1, false, 0.1, 20, 1.0, 21, 1.02)
    iterative_state(args)
end

##############################################################################################
# IEnKS style iterative smoothing for joint state-parameter estmation, arugments are
# time_series, method, seed, nanl, lag, shift, mda, obs_un, obs_dim, γ,
# param_err, param_wlk, N_ens, state_infl, param_infl = args

function iterative_smoother_param_exp()
    args = (time_series, "ienks-transform",
            0, 2500, 10, 1, true, 1.0, 20, 1.0, 0.03, 0.0001,  21, 1.02, 1.0)
    iterative_param(args)
end

##############################################################################################
# end module

end

max_stars_repo_name: cgrudz/DataAssimilationBenchmarks

max_issues_repo_name: cgrudz/DataAssimilationBenchmarks


text:
#Aim: have map operate on FB

@doc Markdown.doc"""
    compose_mod(x::nmod_poly, y::nmod_poly, z::nmod_poly) -> nmod_poly

  Compute x(y) mod z
"""
function compose_mod(x::nmod_poly, y::nmod_poly, z::nmod_poly)
  check_parent(x,y)
  check_parent(x,z)
  r = parent(x)()
  ccall((:nmod_poly_compose_mod, :libflint), Nothing,
          (Ref{nmod_poly}, Ref{nmod_poly}, Ref{nmod_poly}, Ref{nmod_poly}), r, x, y, z)
  return r
end

function compose_mod(x::gfp_poly, y::gfp_poly, z::gfp_poly)
  check_parent(x,y)
  check_parent(x,z)
  r = parent(x)()
  ccall((:nmod_poly_compose_mod, :libflint), Nothing,
          (Ref{gfp_poly}, Ref{gfp_poly}, Ref{gfp_poly}, Ref{gfp_poly}), r, x, y, z)
  return r
end

@doc Markdown.doc"""
    taylor_shift(x::nmod_poly, c::UInt) -> nmod_poly

  Compute x(t-c)
"""
function taylor_shift(x::nmod_poly, c::UInt)
  r = parent(x)()
  ccall((:nmod_poly_taylor_shift, :libflint), Nothing,
          (Ref{nmod_poly}, Ref{nmod_poly}, UInt), r, x, c)
  return r
end

function induce(FB::Hecke.NfFactorBase, A::Map)
  K = domain(A)
  f = A(gen(K)) # essentially a polynomial in the primitive element

  O = order(FB.ideals[1])
  prm = Array{Tuple{Int, Int}, 1}()

  if f == gen(K)
    return PermutationGroup(length(FB.ideals))()
  end

  for p in FB.fb_int.base
    FP = FB.fb[p]
    if length(FP.lp) < 3 || isindex_divisor(O, p) || p > 2^60
      lp = [x[2] for x = FP.lp]
      for (i, P) in FP.lp
        Q = induce_image(A, P)
        id = findfirst(isequal(Q), lp)
        @assert id !== nothing
        push!(prm, (i, FP.lp[id][1]))
      end
    else
      px = PolynomialRing(GF(Int(p), cached=false), "x", cached=false)[1]
      fpx = px(f)
      gpx = px(K.pol)
      #idea/ reason
      # if p is no index divisor, then a second generator is just
      #   an irreducible factor of gpx (Kummer/ Dedekind)
      # an ideal is divisible by P iff the canonical 2nd generator of the prime ideal
      # divides the 2nd generator of the target (CRT)
      # so
      lp = [gcd(px(K(P[2].gen_two)), gpx) for P = FP.lp]
      # this makes lp canonical (should be doing nothing actually)

      for (i,P) in FP.lp
        hp = px(K(P.gen_two))
        if degree(hp)==1
          im = fpx + coeff(hp, 0)
        else
          im = compose_mod(hp, fpx, gpx)
          # the image, directly mod p...
        end
        im = Hecke.gcd!(im, gpx, im)
        # canonical
        push!(prm, (i, FP.lp[findfirst(isequal(im), lp)][1]))
        # and find it!
      end
    end
  end
  sort!(prm, lt=(a,b) -> a[1] < b[1])
  G = PermutationGroup(length(prm))
  return G([x[2] for x = prm])
end

#= implementation from Butler's Fundamntal Algorithm for Permutation Groups
  Algo 4: Dimino
  Tested for cyclic groups - unfortunately only.
  I still need to generate other input
=#
#function orbit_in_FB(op::Array{Tuple{Map, Generic.Perm}, 1}, a::nf_elem, s::SRow)
function orbit_in_FB(op::Array, a::nf_elem, s::SRow)
  function op_smat(n::SRow, p::Generic.Perm)
    r = [(p[i], v) for (i,v) = n]
    sort!(r, lt = (a,b)->a[1]<b[1])
    return typeof(n)(r)
  end

  Ss = Dict{nf_elem, typeof(s)}()
#  Ss = Dict{typeof(s), nf_elem}()
  Ss[a] = s
  # start with the cyclic group be op[1]
  n = s
  b = op[1][1](a)
  while b != a
    n = op_smat(n, op[1][2])
    Ss[b] = n
    b = op[1][1](b)
  end

  for i=2:length(op)
    bb = op[i][1](a)
    if haskey(Ss, bb)
      continue
    end
    old = collect(Ss)
    for (bb, n) in old # one redundant - step
      Ss[op[i][1](bb)] = op_smat(n, op[i][2])
    end
    while true
      done = true
      for j = 1:length(op)
        bb = op[j][1](b)
        if haskey(Ss, bb)
          continue;
        end
        done = false
        b = bb
        for (bbb, n) in old
          # TODO: What is going on with the b's?
          Ss[op[j][1](bbb)] = op_smat(n, op[j][2])
        end
      end
      if done
        break
      end
    end
  end
  return Ss
end

function generated_subgroup(op::Array) #pairs: permutations and Map
  elt = Array{Any, 1}()
  push!(elt, (x->x, parent(op[1][2])()))
  ord = 1
  g = op[1]
  while g[2] != parent(op[1][2])()
    let c_g = g
      push!(elt, c_g)
#      g = (x->op[1][1](c_g[1](x)), op[1][2]*c_g[2])
      g = (x->op[1][1](c_g[1](x)), c_g[2]*op[1][2])
    end
  end
  ord = length(elt)

  for i=2:length(op)
    c_i = i
    if op[i][2] in [x[2] for x=elt]
      continue
    end
    pord = ord
    push!(elt, op[i])
    for j=2:pord
      c_j = j
#      push!(elt, (x->elt[c_j][1](op[c_i][1](x)), elt[c_j][2]*op[c_i][2]))
      push!(elt, (x->elt[c_j][1](op[c_i][1](x)), op[c_i][2]*elt[c_j][2]))
    end
    ord = length(elt)
    rpos = pord + 1
    while true
      for s in op
        let c_rpos = rpos, c_s = s
#          g = (x->elt[c_rpos][1](c_s[1](x)), elt[c_rpos][2]*c_s[2])
          g = (x->elt[c_rpos][1](c_s[1](x)), c_s[2]*elt[c_rpos][2])
          if g[2] in [x[2] for x=elt]
            continue
          end
        end
        let c_g = g
          push!(elt, c_g)
          for j = 2:pord
            c_j = j
#            push!(elt, (x->elt[c_j][1](c_g[1](x)), elt[c_j][2]*c_g[2]))
            push!(elt, (x->elt[c_j][1](c_g[1](x)), c_g[2]*elt[c_j][2]))
          end
        end
        ord = length(elt)
      end
      rpos += pord
      if rpos > length(elt)
        break
      end
    end
  end
  return elt
end


text:
function opt(block::AbstractString, key::AbstractString)
    conf = ConfParse(ascii(joinpath(HOMEDIR, "config", "user_settings.ini")))
    parse_conf!(conf)

    r = ""
    try
        r = retrieve(conf, ascii(lowercase(block)), ascii(lowercase(key)))
    catch
        r = retrieve(default_settings, ascii(lowercase(block)), ascii(lowercase(key)))
    end
    Meta.parse(r)
end
opt(key::AbstractString) = opt("default", key)

# runtime
function init_opt()

    suffix = Sys.islinux() ? "_linux" : ""
    default_settings = ConfParse(ascii(joinpath(HOMEDIR, "config", "default_settings$(suffix).ini")))
    parse_conf!(default_settings)

    if !isfile(joinpath(HOMEDIR, "config", "user_settings.ini"))
        cp(joinpath(HOMEDIR, "config", "default_settings$(suffix).ini"), joinpath(HOMEDIR, "config", "user_settings.ini"))
    end

end


text:
module model

export generate_vel_mask, write_grid, write_model, write_lambda, write_dust, Grid, size_au
export AbstractParameters, ParametersStandard, ParametersVertical, ParametersVerticalEta, ParametersNuker, convert_vector, convert_dict, registered_params
export lnprior

# The double dot is because we are now inside the model module, and we want to import the
# constants module, which is part of the enclosing DiskJockey package.
using ..constants
using Printf
using QuadGK
# using Dierckx
# using Optim
# using ODE

"Write the wavelength sampling file. Only run on setup"
function write_lambda(lams::AbstractArray, basedir::AbstractString)
    fcam = open(basedir * "camera_wavelength_micron.inp", "w")
    nlam = length(lams)
    @printf(fcam, "%d\n", nlam)
    for lam in lams
        @printf(fcam, "%.9e\n", lam) # [microns]
    end

    close(fcam)
end

"Generate a boolean mask corresponding to the velocities which fall inside the user-specified mask ranges.
arr is an array of [start, end] pairs, like arr = [[1.0, 3.0], [4.0, 5.2]]
vels are the velocities corresponding to the dataset, like vels = [1.2, 2.4, 3.6], etc."
function generate_vel_mask(arr, vels)
    # initialize a boolean array of all trues
    mask = trues(length(vels))

    # iterate over the array of exclude values
    for (left, right) in arr
        # set those velocities which fall inside the range of channels to be false
        mask[(left .< vels) .& (vels .< right)] .= false
    end

    return mask
end

const eqmirror = true # mirror the grid about the z=0 midplane ?
# if we decide to mirror, then ncells = 1/2 of the true value

"Define a grid object which stores all of these variables
This will not change for the duration of the run."
struct Grid
    nr::Int
    ntheta::Int
    nphi::Int
    ncells::Int
    # cell edges
    Rs::Vector{Float64} # [cm]
    Thetas::Vector{Float64}
    Phis::Vector{Float64}
    # cell centers
    rs::Vector{Float64} # [cm]
    thetas::Vector{Float64}
    phis::Vector{Float64}
end

"Hold the ray-tracing grid."
function Grid(nr::Int, ntheta::Int, r_in::Real, r_out::Real) #, eqmirror::Bool=true)
    # Specify a 2D axisymmetric *separable* grid in spherical coordinates:
    # {r, theta, phi}, where theta is angle from zenith, phi is azimuth

    # Number of cells in each dimension
    nphi = 1 # axisymmetric disk
    ncells = nr * ntheta * nphi
    r_in = convert(Float64, r_in) * AU # [cm] Inner extent of disk
    r_out = convert(Float64, r_out) * AU # [cm] Outer extent of disk

    #Define the cell *walls*
    # Rs = logspace(log10(r_in), log10(r_out), nr+1) # [cm] logarithmically spaced
    Rs = 10 .^ range(log10(r_in), stop=log10(r_out), length=nr+1) # [cm] logarithmically spaced

    eqmirror = true

    if eqmirror
        ped = 0.1
        #Thetas = LinRange(0, pi/2., ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        # Thetas = pi/2. - (logspace(log10(ped), log10(pi/2. + ped), ntheta+1) - ped)[end:-1:1]
        Thetas = pi/2.0 .- (10 .^ range(log10(ped), stop=log10(pi/2.0 + ped), length=ntheta+1) .- ped)[end:-1:1]
        #Logarithmically spaced closer near the z=0
    else
        Thetas = LinRange(0, pi, ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        # Equally spaced in theta.
    end

    Phis = Float64[0.0, 0.0] # [rad] cell walls for inactive coordinate

    #Define the cell centers as the average between walls
    rs = 0.5 * (Rs[1:end-1] + Rs[2:end]) # [cm]
    thetas = 0.5 * (Thetas[1:end-1] + Thetas[2:end])
    phis = Float64[0.0]

    return Grid(nr, ntheta, nphi, ncells, Rs, Thetas, Phis, rs, thetas, phis)

end

"Hold the ray-tracing grid."
function Grid(nr::Int, ntheta::Int, nphi::Int, r_in::Real, r_out::Real)
    # Specify a 2D axisymmetric *separable* grid in spherical coordinates:
    # {r, theta, phi}, where theta is angle from zenith, phi is azimuth

    # Number of cells in each dimension
    ncells = nr * ntheta * nphi
    r_in = convert(Float64, r_in) * AU # [cm] Inner extent of disk
    r_out = convert(Float64, r_out) * AU # [cm] Outer extent of disk

    #Define the cell *walls*
    # Rs = logspace(log10(r_in), log10(r_out), nr+1) # [cm] logarithmically spaced
    Rs = 10 .^ range(log10(r_in), stop=log10(r_out), length=nr+1) # [cm] logarithmically spaced

    eqmirror = true

    if eqmirror
        ped = 0.1
        #Thetas = LinRange(0, pi/2., ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        # Thetas = pi/2. - (logspace(log10(ped), log10(pi/2. + ped), ntheta+1) - ped)[end:-1:1]
        Thetas = pi/2.0 .- (10 .^ range(log10(ped), stop=log10(pi/2. + ped), length=ntheta+1) .- ped)[end:-1:1]
        #Logarithmically spaced closer near the z=0
    else
        Thetas = LinRange(0, pi, ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        # Equally spaced in theta.
    end

    Phis = LinRange(0, 2pi, nphi+1) # [rad] cell walls for inactive coordinate

    #Define the cell centers as the average between walls
    rs = 0.5 * (Rs[1:end-1] + Rs[2:end]) # [cm]
    thetas = 0.5 * (Thetas[1:end-1] + Thetas[2:end])
    phis = 0.5 * (Phis[1:end-1] + Phis[2:end])

    return Grid(nr, ntheta, nphi, ncells, Rs, Thetas, Phis, rs, thetas, phis)

end

"Create a grid object using a logarithmic then linear then logarithmic radial spacing"
function Grid(r_in::Real, r_linstart::Real, r_linend::Real, r_out::Real, n_in::Int, n_mid::Int, n_out::Int, ntheta::Int, eqmirror::Bool=true)
    # Number of cells in each dimension
    nphi = 1 # axisymmetric disk
    nr = n_in + n_mid + n_out
    ncells = nr * ntheta * nphi
    r_in = convert(Float64, r_in) * AU # [cm] Inner extent of disk
    r_linstart = convert(Float64, r_linstart) * AU # [cm] Start of linear section
    r_linend = convert(Float64, r_linend) * AU # [cm] End of linear section
    r_out = convert(Float64, r_out) * AU # [cm] Outer extent of disk

    #Define the cell *walls*
    # logarithmically spaced inner grid
    # Rs_in = logspace(log10(r_in), log10(r_linstart), n_in + 1) # [cm]
    Rs_in = 10 .^ range(log10(r_in), stop=log10(r_linstart), length=n_in + 1) # [cm]

    # linearly spaced middle grid
    Rs_mid = LinRange(r_linstart, r_linend, n_mid + 1) # [cm]

    # logarithmically spaced outer grid
    # Rs_out = logspace(log10(r_linend), log10(r_out), n_out + 1) # [cm]
    Rs_out = 10 .^ range(log10(r_linend), stop=log10(r_out), length=n_out + 1) # [cm]

    Rs = cat(1, Rs_in[1:end-1], Rs_mid, Rs_out[2:end])

    if eqmirror
        ped = 0.1
        #Thetas = LinRange(0, pi/2., ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        # Thetas = pi/2. - (logspace(log10(ped), log10(pi/2. + ped), ntheta+1) - ped)[end:-1:1]
        Thetas = pi/2.0 .- (10 .^ range(log10(ped), stop=log10(pi/2. + ped), length=ntheta+1) - ped)[end:-1:1]
        #Spaced closer near the z=0
    else
        Thetas = LinRange(0, pi, ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
    end

    Phis = Float64[0.0, 0.0] # [rad] cell walls for inactive coordinate

    #Define the cell centers as the average between walls
    rs = 0.5 * (Rs[1:end-1] + Rs[2:end]) # [cm]
    thetas = 0.5 * (Thetas[1:end-1] + Thetas[2:end])
    phis = Float64[0.0]

    return Grid(nr, ntheta, nphi, ncells, Rs, Thetas, Phis, rs, thetas, phis)
end

"Read from a dictionary, then choose how to make the grid based upon the arguments."
function Grid(d::Dict)
    if "r_linstart" in keys(d)
        # We're going for a log-linear-log grid
        names = ["r_in", "r_linstart", "r_linend", "r_out", "n_in", "n_mid", "n_out", "ntheta"]
    elseif d["nphi"] > 1
        names = ["nr", "ntheta", "nphi", "r_in", "r_out"]
    else
        # We're just going for a log spaced grid
        names = ["nr", "ntheta", "r_in", "r_out"]
    end
    vec = [d[name] for name in names]
    # Unroll these into an actual parameter
    return Grid(vec...)
end

"This function only needs to be run once, upon setup."
function write_grid(basedir::AbstractString, grid::Grid)
    #amr_grid.inp
    f = open(basedir * "amr_grid.inp", "w")

    #Write the header
    @printf(f, "%d\n", 1) #iformat
    @printf(f, "%d\n", 0) #regular grid (no AMR or Oct-tree)
    @printf(f, "%d\n", 100) #spherical coordiantes
    @printf(f, "%d\n", 0) #gridinfo (none needed for now)
    #incl_r incl_phi incl_z #use this axis?
    if length(grid.Phis) > 2
        @printf(f, "%d %d %d \n", 1, 1, 1) # 2D axisymmetric
    else
        @printf(f, "%d %d %d \n", 1, 1, 0) # 2D axisymmetric
    end

    #n_r    n_phi   n_z #number of cells in this dimension
    @printf(f, "%d %d %d \n", grid.nr, grid.ntheta, grid.nphi)

    for R in grid.Rs
        @printf(f, "%.9e\n", R)
    end

    for Theta in grid.Thetas
        @printf(f, "%.9e\n", Theta)
    end

    for Phi in grid.Phis
        @printf(f, "%.9e\n", Phi)
    end

    close(f)
end

# Define an abstract Parameters type, then subset of parameters that can be used to dispatch specifics

abstract type AbstractParameters end

"Parameters for the standard model."
mutable struct ParametersStandard <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface density gradient exponent
    log_M_gas::Float64 # [log10 M_sun] total disk mass
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
    Sigma_c::Float64 # [g/cm^2] surface density at r_c
end

# constructor for ParametersStandard takes all but Sigma_c
function ParametersStandard(M_star::Float64, r_c::Float64, T_10::Float64, q::Float64, gamma::Float64, log_M_gas::Float64, ksi::Float64, dpc::Float64,incl::Float64, PA::Float64, vel::Float64, mu_RA::Float64, mu_DEC::Float64)

  # calculate the normalization constant
  Sigma_c::Float64 = 10^log_M_gas * M_sun * (2 - gamma) / (2 * pi * (r_c * AU)^2)

  # initialize the actual type
  return ParametersStandard(M_star, r_c, T_10, q, gamma, log_M_gas, ksi, dpc,incl, PA, vel, mu_RA, mu_DEC, Sigma_c)
end

"Parameters for the vertical model."
mutable struct ParametersVertical <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    T_10m::Float64 # [K] temperature at 10 AU, midplane
    q_m::Float64 # midplane temperature gradient exponent
    T_10a::Float64 # [K] temperature at 10 AU, atmosphere
    q_a::Float64 # atmosphere temperature gradient exponent
    T_freeze::Float64 # [K] temperature below which to reduce CO abundance
    X_freeze::Float64 # [ratio] amount to reduce CO abundance
    sigma_s::Float64 # Photodissociation boundary in units of A_V.
    gamma::Float64 # surface temperature gradient exponent
    h::Float64 # Number of scale heights that z_q is at, typically fixed to 4
    delta::Float64 # Shape exponent, currently fixed to 2
    log_M_gas::Float64 # [log10 M_sun] total disk mass
    ksi::Float64 # [cm s^{-1}] micsroturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
    Sigma_c::Float64 # [g/cm^2] surface density at r_c
end

"Parameters for the vertical model with variable slope."
mutable struct ParametersVerticalEta <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    T_10m::Float64 # [K] temperature at 10 AU, midplane
    q_m::Float64 # midplane temperature gradient exponent
    T_freeze::Float64 # [K] temperature below which to reduce CO abundance
    X_freeze::Float64 # [ratio] amount to reduce CO abundance
    sigma_s::Float64 # Photodissociation boundary in units of A_V.
    gamma::Float64 # surface density gradient exponent
    h::Float64 # Number of scale heights that z_q is at, typically fixed to 4
    eta::Float64 # Exponent for zq profile
    delta::Float64 # Shape exponent, currently fixed to 2
    log_M_gas::Float64 # [log10 M_sun] total disk mass
    ksi::Float64 # [cm s^{-1}] micsroturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
    Sigma_c::Float64 # [g/cm^2] surface density at r_c
end

"Parameters for the NUKER model."
mutable struct ParametersNuker <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface density gradient within r_c (negative values yield holes)
    log_alpha::Float64 #  log10 sharpness of transition (2 = smooth, 16 = sharp)
    beta::Float64 # gradient power law outside r_c (~7)
    log_M_gas::Float64 # [log10 M_sun] total disk mass
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
    Sigma_c::Float64 # [g/cm^2] surface density at r_c
end

# constructor for ParametersStandard takes all but Sigma_c
function ParametersNuker(M_star::Real, r_c::Real, T_10::Real, q::Real, gamma::Real, log_alpha::Real, beta::Real, log_M_gas::Real, ksi::Real, dpc::Real,incl::Real, PA::Real, vel::Real, mu_RA::Real, mu_DEC::Real)

    if r_c <= 0.0
        throw(ModelException("r_c cannot be negative, $r_c"))
    end

    M_gas = 10^log_M_gas * M_sun
    alpha = 10^log_alpha

    function integrand(r)
        return (r/r_c)^(-gamma) * (1 + (r/r_c)^alpha)^((gamma - beta)/alpha) * r
    end

    # calculate the normalization constant via an integral
    # break it into two pieces bounded at r_c, since this might help accuracy
    res, err = quadgk(integrand, 1e-3, r_c, 1e4) # [AU^2]
    # convert area to cm^2
    res_cm = res * AU^2 # [cm^2]
    Sigma_c = M_gas / (2 * pi * res_cm) # [g/cm^2]

     # calculate the normalization constant
    # we need to do an integral
    # M_gas = 2 * pi * Sigma_c * integrate((r/r_c)^(-gamma) * (1 + (r/r_c)^alpha)^((gamma - beta)/alpha) * r * dr)

    # initialize the actual type
    return ParametersNuker(M_star, r_c, T_10, q, gamma, log_alpha, beta, log_M_gas, ksi, dpc,incl, PA, vel, mu_RA, mu_DEC, Sigma_c)
end

"A dictionary of parameter lists for conversion."
registered_params = Dict([("standard", ["M_star", "r_c", "T_10", "q", "gamma", "log_M_gas", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]),
("vertical", ["M_star", "r_c", "T_10m", "q_m", "T_10a", "q_a", "T_freeze", "X_freeze", "sigma_s", "gamma", "h", "delta", "log_M_gas", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]),
("verticalEta", ["M_star", "r_c", "T_10m", "q_m", "T_freeze", "X_freeze", "sigma_s", "gamma", "h", "eta", "delta", "log_M_gas", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]),
("nuker", ["M_star", "r_c", "T_10", "q", "gamma", "log_alpha", "beta", "log_M_gas", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"])])

registered_types = Dict([("standard", ParametersStandard), ("vertical", ParametersVertical), ("verticalEta", ParametersVerticalEta), ("nuker", ParametersNuker)])

"Unroll a vector of parameter values into a parameter type."
function convert_vector(p::Vector{Float64}, model::AbstractString, fix_params::Vector; args...)
    args = Dict{Symbol}{Float64}(args)

    # The goal is to assemble a length-nparam vector that can be unrolled into a parameter type
    # e.g., ParametersStandard(M_star, r_c, T_10, q, gamma, log_M_gas, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC)

    # Select the registerd parameters corresponding to this model
    # These are the names listed in this file (model.jl)
    reg_params = registered_params[model]

    # fix_params is a list of strings from a config file that list the names of parameters to be fixed to the config
    # value. These names correspond to the registered names.

    # fit_params are the ones in reg_params that are not in (∉) fix_params
    fit_params = filter(x->∉(x,fix_params), reg_params)

    nparams = length(reg_params)

    # Make an empty vector of this same length
    par_vec = Array{Float64}(undef, nparams)

    # This requires assigning p to fit_params
    # Find the indexes that correspond to fit_params
    # par_indexes = findin(reg_params, fit_params)
    par_indexes = findall((in)(fit_params), reg_params)
    # Stuff p directly into these
    par_vec[par_indexes] = p

    # Then reading fix_params from args
    # First, create an array of fixed values analogous to p
    p_fixed = Float64[args[Symbol(par)] for par in fix_params]
    par_indexes = findall((in)(fix_params), reg_params)
    par_vec[par_indexes] = p_fixed

    # Then assembling these in the same orignial order as registered_params, into the parameter
    # type corresponding to the model.
    parameters = registered_types[model](par_vec...)

    return parameters

end

"Used to turn a dictionary of parameter values (from config.yaml) directly into a parameter type. Generally used for synthesis and plotting command line scripts."
function convert_dict(p::Dict, model::AbstractString)
    # Select the registerd parameters corresponding to this model
    reg_params = registered_params[model]
    nparams = length(reg_params)

    # Using this order of parameters, unpack the dictionary p into a vector
    # reg_params reads Sigma_c, not logSigma_c; alpha, not logalpha
    par_vec = Float64[p[par_name] for par_name in reg_params]

    return registered_types[model](par_vec...)

end

"The common sense priors that apply to all parameter values"
function lnprior_base(pars::AbstractParameters)
    # Create a giant short-circuit or loop to test for sensical parameter values.
    if pars.M_star <= 0.0 || pars.ksi <= 0.0 || pars.T_10 <= 0.0 || pars.r_c <= 0.0  || pars.T_10 > 1500. || pars.q < 0. || pars.q > 1.0 || pars.incl < 0. || pars.incl > 180. || pars.PA < -180. || pars.PA > 520.
        # println("M_star ", pars.M_star)
        # println("r_c ", pars.r_c)
        # println("T_10 ", pars.T_10)
        # println("q ", pars.q)
        # println("incl ", pars.incl)
        # println("PA ", pars.PA)
        # println("Exiting in lnprior base")
        throw(ModelException("Parameters outside of prior range $pars"))
    end

    # If we've passed all the hard-cut offs by this point, return the geometrical inclination prior.
    # natural log, not log10
    return log(0.5 * sind(pars.incl))

end

function lnprior(pars::ParametersStandard, grid::Grid)

    lnp = lnprior_base(pars)

    r_out = grid.Rs[end]/AU # [AU]
    # A somewhat arbitrary cutoff regarding the gridsize to prevent the disk from being too large
    # to fit on the model grid.
    if (3 * pars.r_c) > r_out
        throw(ModelException("Model radius ($(pars.r_c)) too large for grid size ($r_out)"))
    else
        return lnp
    end

end

function lnprior(pars::ParametersVertical, grid::Grid)
    # Create a giant short-circuit or loop to test for sensical parameter values.
    if pars.M_star <= 0.0 || pars.ksi <= 0.0 || pars.T_10a <= 0.0 || pars.T_10m <= 0.0 || pars.r_c <= 0.0  || pars.T_10a > 1500.0 || pars.q_m < 0.0 || pars.q_a < 0.0 || pars.q_m > 1.0 || pars.q_a > 1.0 || pars.incl < 0. || pars.incl > 180.0 || pars.PA < -180.0 || pars.PA > 520.0 || pars.X_freeze > 1.0 || pars.sigma_s < 0.0
        throw(ModelException("Parameters outside of prior range $pars"))
    end

    # Check to see that the temperatures make sense
    if pars.T_10m > pars.T_10a
        throw(ModelException("Atmosphere ($(pars.T_10a)) is cooler than midplane ($(pars.T_10m)."))
    end

    # If we've passed all the hard-cut offs by this point, return the geometrical inclination prior.
    # natural log, not log10
    lnp = log(0.5 * sind(pars.incl))

    r_out = grid.Rs[end]/AU # [AU]
    # A somewhat arbitrary cutoff regarding the gridsize to prevent the disk from being too large
    # to fit on the model grid.
    if (3 * pars.r_c) > r_out
        throw(ModelException("Model radius ($(pars.r_c)) too large for grid size ($r_out)."))
    else
        return lnp
    end

end

function lnprior(pars::ParametersVerticalEta, grid::Grid)
    # Create a giant short-circuit or loop to test for sensical parameter values.
    if pars.M_star <= 0.0 || pars.ksi <= 0.0 || pars.T_10m <= 60.0 || pars.r_c <= 0.0  || pars.q_m < 0.0 || pars.q_m > 1.0 || pars.incl < 0.0 || pars.incl > 180.0 || pars.PA < -180.0 || pars.PA > 520.0 || pars.X_freeze > 1.0 || pars.sigma_s < 0.0 || pars.eta < 0.2 || pars.eta > 6.0  || pars.h < 0.5 || pars.h > 6.0 || pars.delta < 0.5 || pars.delta > 6.0 || pars.gamma < 0.5 || pars.gamma > 3.0
        throw(ModelException("Parameters outside of prior range $pars"))
    end

    # If we've passed all the hard-cut offs by this point, return the geometrical inclination prior.
    # natural log, not log10
    lnp = log(0.5 * sind(pars.incl))

    r_out = grid.Rs[end]/AU # [AU]
    # A somewhat arbitrary cutoff regarding the gridsize to prevent the disk from being too large
    # to fit on the model grid.
    if (3 * pars.r_c) > r_out
        throw(ModelException("Model radius ($(pars.r_c)) too large for grid size ($r_out)."))
    else
        return lnp
    end

end

function lnprior(pars::ParametersNuker, grid::Grid)

    # println("In lnprior(Nuker)")
    lnp = lnprior_base(pars)

    if (pars.log_alpha < 0.0 ) || (pars.log_alpha > 2.0) || (pars.beta < 2) || (pars.beta > 10)
        # println("alpha: ", pars.alpha)
        # println("beta: ", pars.beta)
        # println("Alpha or Beta outside of prior range.")
        throw(ModelException("log_alpha ($(pars.log_alpha)) or beta ($(pars.beta)) outside of prior range."))
    end

    # Add on the prior on gamma (log, not log10)
    lnp_gamma = log(1/(1 + exp(-5 * (pars.gamma + 3))) - 1/(1 + exp(-15 * (pars.gamma - 2))))
    lnp += lnp_gamma

    r_out = grid.Rs[end]/AU # [AU]
    # A somewhat arbitrary cutoff regarding the gridsize to prevent the disk from being too large
    # to fit on the model grid.
    if (3 * pars.r_c) > r_out
        # println("Model radius too large for grid size.")
        throw(ModelException("Model radius ($(pars.r_c)) too large for grid size ($r_out)."))
    else
        # println("Returning lnprior as ", lnp)
        return lnp
    end

end

# Determine the physical size of the image. Size_arcsec is the full width of the image,
# and so sizeau is the full size of the image as well (RADMC3D conventions).
# Because there seems to be a small but constant shift in the size of the image
function size_au(size_arcsec::Real, dpc::Real, grid::Grid)
    outer_model = (1.1 * grid.rs[end] / AU) # [AU]
    sizeau_desired = size_arcsec * dpc # [AU]

    if sizeau_desired < outer_model
        # We don't want this error to be caught by the lnprob routine, since this should
        # always halt execution of the script.
        throw(ErrorException("Image size ($size_au AU) must be larger than 110% the model grid size ($outer_model AU)."))
    end

    # because there seems to be a slight offset between what sizeau is specified and the actual
    # size of the image, we will also specifiy sizeau_command, which is the value to give to RADMC
    sizeau_command = sizeau_desired/(1.0 + RADMC_SIZEAU_SHIFT)

    return (sizeau_desired, sizeau_command) # [AU]

end

# Assume all inputs to these functions are in CGS units and in *cylindrical* coordinates.
# Parametric type T allows passing individual Float64 or Vectors.
# # Alternate functions accept pars passed around, where pars is in M_star, AU, etc...
function velocity(r::T, M_star::Float64) where {T}
    sqrt.(G * M_star ./ r)
end
velocity(r::T, pars::AbstractParameters) where {T} = velocity(r, pars.M_star * M_sun)

# For the vertical temperature gradient
function velocity(r::Float64, z::Float64, M_star::Float64)
    sqrt.(G * M_star / (r^2 + z^2)^(3.0/2)) * r
end
velocity(r::Float64, z::Float64, pars::ParametersVertical) = velocity(r, z, pars.M_star * M_sun)
velocity(r::Float64, z::Float64, pars::ParametersVerticalEta) = velocity(r, z, pars.M_star * M_sun)

function temperature(r::T, T_10::Float64, q::Float64) where {T}
    T_10 * (r ./ (10.0 * AU)).^(-q)
end
temperature(r::T, pars::AbstractParameters) where {T} = temperature(r, pars.T_10, pars.q)

function Hp(r::T, M_star::Float64, T_10::Float64, q::Float64) where {T}
    temp = temperature(r, T_10, q)
    sqrt.(kB * temp .* r.^3.0/(mu_gas * m_H * G * M_star))
end
Hp(r::T,  pars::AbstractParameters) where {T} = Hp(r, pars.M_star * M_sun, pars.T_10, pars.q)
# Scale height computed from midplane temperature for vertical temperature gradient model
Hp(r::T,  pars::Union{ParametersVertical,ParametersVerticalEta}) where {T} = Hp(r, pars.M_star * M_sun, pars.T_10m, pars.q_m)

# For the vertical temperature gradient model
T_mid(r::T, pars::ParametersVertical) where {T} = temperature(r, pars.T_10m, pars.q_m)
T_atm(r::T, pars::ParametersVertical) where {T} = temperature(r, pars.T_10a, pars.q_a)

T_mid(r::T, pars::ParametersVerticalEta) where {T} = temperature(r, pars.T_10m, pars.q_m)

function T_atm(r::T, pars::ParametersVerticalEta) where {T}
    Tm = T_mid(r, pars)

    # the atmosphere temperature is the maximum of 500 or twice the midplane
    # in order to get the structure close to the star correct
    Ta = maximum([500, 2 * Tm])
    return Ta
end

# Atmosphere height computed as multiple of scale height (computed at midplane)
function z_q(r::T, pars::ParametersVertical) where {T}
    return pars.h * Hp(r, pars)
end

function z_q(r::T, pars::ParametersVerticalEta) where {T}

    # Calculate the scale height at this radius
    H = Hp(r, pars) # [cm]
    r0 = 10 * AU # [cm]

    if r < r0
      return pars.h * H
    else
      return pars.h * H * (r / r0)^pars.eta
    end
end

function temperature(r::Float64, z::Float64, pars::ParametersVertical)
    zq = z_q(r, pars)
    Ta = T_atm(r, pars)
    Tm = T_mid(r, pars)
    if Tm > Ta
      return Ta
    end

    if z >= zq
        return Ta
    else
        # return Ta + (Tm - Ta) * cos(pi * z/ (2 * zq))^pars.delta # Dartois
        return Tm + (Ta - Tm) * sin(pi * z/ (2 * zq))^(2 * pars.delta) # Williams and Best 14
    end
end

function temperature(r::Float64, z::Float64, pars::ParametersVerticalEta)
    zq = z_q(r, pars)
    Tm = T_mid(r, pars)
    Ta = T_atm(r, pars)

    H = Hp(r, pars) # scale height

    # If we are less than one scale height from the midplane, just return the midplane temperature
    if z < H
      return Tm
    end

    if Tm > Ta
      return Ta
    end

    if z >= zq
        return Ta
    else
        return Tm + (Ta - Tm) * sin(pi * (z - H)/ (2 * zq))^(2 * pars.delta)
    end
end

"Calculate the gas surface density"
function Sigma(r::Float64, pars::Union{ParametersStandard, ParametersVertical, ParametersVerticalEta})
    r_c = pars.r_c * AU

    gamma = pars.gamma
    Sigma_c = pars.Sigma_c

    S = Sigma_c * (r/r_c)^(-gamma) * exp(-(r/r_c)^(2 - gamma))

    return S
end

"
    Sigma(r::Float64, pars::ParametersNuker)

Calculate the gas surface density using the Nuker profile."
function Sigma(r::Float64, pars::ParametersNuker)
    r_c = pars.r_c * AU

    gamma = pars.gamma
    Sigma_c = pars.Sigma_c

    alpha = 10^pars.log_alpha
    beta = pars.beta

    S = Sigma_c * (r/r_c)^(-gamma) * (1 + (r/r_c)^alpha)^((gamma - beta)/alpha)

    return S
end

# Delivers a gas density in g/cm^3
function rho_gas(r::Float64, z::Float64, pars::AbstractParameters)
    H = Hp(r, pars)
    S = Sigma(r, pars)

    # Calculate the density
    rho = S/(sqrt(2.0 * pi) * H) * exp(-0.5 * (z/H)^2)

    return rho
end

# Determines the midplane density assuming vertically isothermal
function rho_gas_mid(r::Float64, pars::AbstractParameters)
    H = Hp(r, pars)
    S = Sigma(r, pars)

    # Calculate the density at the midplane
    rho = S/(sqrt(2.0 * pi) * H)

    return rho

end

# The temperature change, dT/dz
function dT(r::Float64, z::Float64, pars::ParametersVertical)
    zq = z_q(r, pars)
    Ta = T_atm(r, pars)

    if z >= zq
        return 0.0
    else
        Tm = T_mid(r, pars)
        # return - pi * pars.delta / (2 * zq) * (Tm - Ta) * (cos(pi * z/(2 * zq)))^(pars.delta - 1) * sin(pi * z / (2 * zq))
        return pars.delta * pi/zq * (Ta - Tm) * cos(pi * z/(2 * zq)) * (sin(pi * z/(2 * zq)))^(2 * pars.delta - 1)
    end
end

# The temperature change, dT/dz
function dT(r::Float64, z::Float64, pars::ParametersVerticalEta)

    zq = z_q(r, pars)
    Ta = T_atm(r, pars)

    H = Hp(r, pars)

    if z >= zq
        return 0.0
    elseif z < H
      return 0.0
    else
        Tm = T_mid(r, pars)

        return pars.delta * pi/zq * (Ta - Tm) * cos(pi * (z - H)/(2 * zq)) * (sin(pi * (z - H)/(2 * zq)))^(2 * pars.delta - 1)
    end
end

# The integrand
function dlnrho(r::Float64, z::Float64, pars::Union{ParametersVertical,ParametersVerticalEta})
    T = temperature(r, z, pars)
    dTemp = dT(r, z, pars)
    return -(dTemp/T + (mu_gas * m_H)/(kB * T) * (G * pars.M_star * M_sun * z)/(r^2 + z^2)^1.5)
end

# The integrand for a vertically isothermal testcase
function dlnrho(r::Float64, z::Float64, pars::ParametersStandard)
    return -(mu_gas * m_H)/(kB * T) * (G * pars.M_star * M_sun * z)/r^3
end

function z_top(r::Float64, pars::Union{ParametersVertical, ParametersVerticalEta})
    zq = z_q(r, pars) # The height of the disk "atmosphere"

    return zq * 5
end

# function rho_gas(r::Float64, z::Float64, pars::Union{ParametersVertical, ParametersVerticalEta})
#
#     # Calculate the "top" of the atmosphere,
#     # a height where we are sure we can enforce that density = 0
#     ztop = z_top(r, pars)
#
#     # If we are querying a height above this, just return 0 now
#     if z > ztop
#         return constants.rho_gas_zero
#     end
#
#     # Calculate what the maximum surface density would be if we integrated all the way to  the midplane
#     sigma = Sigma(r, pars) / 2
#
#     # Calculate the photodissociation height
#     # Threshold column density
#     thresh_H2 = pars.sigma_s * Av_sigmaH / 2 # [n_H2/cm^2]
#     thresh = thresh_H2 * (mu_gas * amu / X_H2) # [g/cm^2] of gas
#
#     # If there is not even enough column density at the midplane in order to exceed the threshold
#     # to protect against photodissociation of CO, just return 0 now
#     # if sigma < thresh
#     #     return constants.rho_gas_zero
#     # end
#
#     # Create an array of heights that goes from the midplane (0, 0.001, ..., ztop)
#     zs = cat(1, [0], logspace(log10(0.001 * AU), log10(ztop), 64 - 1))
#
#     # Define the ODE at this radius r
#     function f(z, y)
#         # calculate y'
#         dlnrho(r, z, pars) * y
#     end
#
#     # Choose a starting guess corresponding to the midplane density for an isothermal model.
#     # This helps us get in the ballpark and avoid a lot of the numerical errors that result
#     start = rho_gas_mid(r, pars)
#
#     # Now solve the ODE on the grid of postulated values
#     # For the inner disk, we require a larger value of abstol
#     # We need to experiment with what is best
#     abstol=1e-22
#     z_out, y_out = ode45(f, start, zs; abstol=1e-22)
#
#     # where y_out is less than the minimum accuracy we can trust, just set it to the minimum value
#     y_out[y_out .< abstol] = abstol
#
#     # Integrate the output from the ODE solver to find the total (un-normalized) surface density
#     spl = Spline1D(z_out, y_out)
#     tot = integrate(spl, z_out[1], z_out[end])
#
#     # Apply a correction factor to the un-normalized densities
#     cor = sigma / tot
#     rhos = y_out .* cor
#
#     # This is the function that will be minimized
#     function g(zvar::Float64)
#         return abs(thresh - cor * integrate(spl, zvar, z_out[end]))
#     end
#
#     z_phot = optimize(g, z_out[1], z_out[end]).minimum
#
#     # Evaluate rho at this point
#     rho = cor * evaluate(spl, z)
#
#     # If we are above this height, return the gas density reduced by a factor of 100
#     if z > z_phot
#         rho = 1e-2 * rho
#     end
#
#     if rho < 0.0
#       println("r ", r/AU, " z ", z/AU, " z_top ", ztop/AU, " z_phot ", z_phot/AU, " y_out ", y_out)
#       return (start, z_out, y_out)
#       throw(ModelException("Rho less than 0.0"))
#     end
#
#     return rho
# end

# Now, replace these functions to simply multiply rho_gas by X_12CO/m_12CO, or X_13CO/m_13CO, etc.
n_12CO(r::Float64, z::Float64, pars::AbstractParameters) = number_densities["12CO"] * rho_gas(r, z, pars)

n_13CO(r::Float64, z::Float64, pars::AbstractParameters) = number_densities["13CO"] * rho_gas(r, z, pars)

n_C18O(r::Float64, z::Float64, pars::AbstractParameters) = number_densities["C18O"] * rho_gas(r, z, pars)

# It is realistic to include freezout of CO onto dust grains.
# This is the amount by which the number density of the CO is reduced (X_freeze) relative to
# the nominal value.
function X_freeze(temp::Float64, pars::AbstractParameters)
    # If it's cooler than the freezout temperature, reduce the number density by the given factor
    if temp <= pars.T_freeze
        return pars.X_freeze
    # Otherwise, just keep it as is.
    else
        return 1.0
    end
end

function rho_dust(r::Float64, z::Float64, pars::AbstractParameters)

    # Use the rho_gas function to get the total density in [g/cm^3]
    mGas = rho_gas(r, z, pars)

    # Convert from mGas to mDust using Gas/Dust ratio of 100
    mDust = mGas * 0.01 # [g]

    return mDust
end

# Ksi is microturbulent broadining width in units of km/s. Output of this function
# is in cm/s for RADMC (RADMC manual, eqn 7.12)
function microturbulence(ksi::Float64)
    return ksi * 1.e5 # convert from km/s to cm/s
end

microturbulence(pars::AbstractParameters) = microturbulence(pars.ksi)

function write_model(pars::AbstractParameters, basedir::AbstractString, grid::Grid, species::AbstractString)

    funcs = Dict([("12CO", n_12CO), ("13CO", n_13CO), ("C18O", n_C18O)])
    n_CO = funcs[species]

    # numberdens_co.inp
    fdens = open(basedir * "numberdens_" * molnames[species] * ".inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)

    # gas_velocity.inp
    fvel = open(basedir * "gas_velocity.inp", "w")
    @printf(fvel, "%d\n", 1) #iformat
    @printf(fvel, "%d\n", grid.ncells)

    # gas_temperature.inp
    ftemp = open(basedir * "gas_temperature.inp", "w")
    @printf(ftemp, "%d\n", 1) #iformat
    @printf(ftemp, "%d\n", grid.ncells)

    # microturbulence.inp
    fmicro = open(basedir * "microturbulence.inp", "w")
    @printf(fmicro, "%d\n", 1) #iformat
    @printf(fmicro, "%d\n", grid.ncells)

    # Now, we will need to write the three other files as a function of grid position.
    # Therefore we will do *one* loop over these indices, calculate the required value,
    # and write it to the appropriate file.

    #Looping over the cell centers
    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", n_CO(r_cyl, z, pars))
                @printf(fvel, "0 0 %.9e\n", velocity(r_cyl, pars))
                @printf(ftemp, "%.9e\n", temperature(r_cyl, pars))
                @printf(fmicro, "%.9e\n", microturbulence(pars))
            end
        end
    end

    close(fdens)
    close(fvel)
    close(ftemp)
    close(fmicro)

end

function P_x(var)
    mat =  Float64[[1, 0, 0]  [0, cos(var), -sin(var)] [0, sin(var), cos(var)]]
    return mat
end

function P_z(var)
    mat =   Float64[[cos(var), -sin(var), 0] [sin(var), cos(var),   0] [0,        0,          1]]
    return mat
end

# function to transform spherical to cartesian
function P_project(theta, phi)
    mat = Float64[ [sin(theta) * cos(phi), sin(theta) * sin(phi), cos(theta)] [cos(theta) * cos(phi), cos(theta) * sin(phi), -sin(theta)] [-sin(phi), cos(phi), 0]]
    return mat
end

# # Get the components of the vector in RADMC coordinates
# function velocity_inner(pars, r, theta, phi)
#
#     Pproj = P_project(theta, phi)
#     # define rotation matrices
#     Px = P_x(pars.incl_inner * pi/180)
#     Pz = P_z(pars.Omega * pi/180)
#
#
#     # get primed coordinates
#     cart, cart_prime, sphere_prime = convert_position(pars, r, theta, phi)
#
#     # query the velocity field of the inner disk, in the primed frame,
#     # which will be projected along the hat{\phi^\prime} direction.
#     # for now, just call this v_phi_prime
#     r_prime, theta_prime, phi_prime = sphere_prime
#     r_cyl_prime = r_prime * sin(theta_prime)
#     z_cyl_prime = r_prime * cos(theta_prime)
#     vel_sphere_prime = Float64[0.0, 0.0, velocity(r_cyl_prime, z_cyl_prime, pars)]
#
#     # project it to the primed cartesian unit vectors
#     vel_cart_prime = Pproj * vel_sphere_prime
#
#     # now rotate this from the primed cartesian frame to the RADMC-3D cartesian frame
#     vel_cart = Pz * Px * vel_cart_prime
#
#     # now project it on to the spherical unit vectors in the RADMC-3D frame
#     vel_sphere = Pproj.' * vel_cart
#
#     # assert that the norm of all of these vectors is always zero
#     @assert isapprox(norm(vel_cart_prime), norm(vel_sphere_prime)) "Cartesian norm doesn't match "
#     @assert isapprox(norm(vel_cart), norm(vel_sphere_prime)) "Rotated cartesion norm doesn't match"
#     @assert isapprox(norm(vel_sphere), norm(vel_sphere_prime)) "Spherical vel doesn't match"
#
#     return vel_sphere
#
# end

function write_model(pars::Union{ParametersVertical, ParametersVerticalEta}, basedir::AbstractString, grid::Grid, species::AbstractString)

    function n_CO(r_cyl, z)
        return number_densities[species] * rho_gas(r_cyl, z, pars)
    end

    # numberdens_co.inp
    fdens = open(basedir * "numberdens_" * molnames[species] * ".inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)

    # gas_velocity.inp
    fvel = open(basedir * "gas_velocity.inp", "w")
    @printf(fvel, "%d\n", 1) #iformat
    @printf(fvel, "%d\n", grid.ncells)

    # gas_temperature.inp
    ftemp = open(basedir * "gas_temperature.inp", "w")
    @printf(ftemp, "%d\n", 1) #iformat
    @printf(ftemp, "%d\n", grid.ncells)

    # microturbulence.inp
    fmicro = open(basedir * "microturbulence.inp", "w")
    @printf(fmicro, "%d\n", 1) #iformat
    @printf(fmicro, "%d\n", grid.ncells)

    # Now, we will need to write the three other files as a function of grid position.
    # Therefore we will do *one* loop over these indices, calculate the required value,
    # and write it to the appropriate file.

    #Looping over the cell centers
    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                temp = temperature(r_cyl, z, pars)

                @printf(fdens, "%.9e\n", X_freeze(temp, pars) * n_CO(r_cyl, z))
                @printf(fvel, "0 0 %.9e\n", velocity(r_cyl, z, pars))
                @printf(ftemp, "%.9e\n", temp)
                @printf(fmicro, "%.9e\n", microturbulence(pars))
            end
        end
    end

    close(fdens)
    close(fvel)
    close(ftemp)
    close(fmicro)

end

function write_dust(pars::AbstractParameters, basedir::AbstractString, grid::Grid)
    fdens = open(basedir * "dust_density.inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)
    @printf(fdens, "%d\n", 1) # number of dust species

    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", rho_dust(r_cyl, z, pars))
            end
        end
    end

    close(fdens)
end

end


text:
"""
Computes random samples of functions from aligned data using Gaussian model

    gauss_model(fn, timet, qn, gam; n=1, sort_samples=false)
    :param fn: aligned functions (M,N)
    :param timet: vector (M) describing time
    :param qn: aligned srvfs (M,N)
    :param gam: warping functions (M,N)
    :param n: number of samples
    :param sort_samples: sort samples

    Returns Dict containing
    :return fs: random aligned functions
    :return gams: random warping functions
    :return ft: random functions
"""
function gauss_model(fn, timet, qn, gam; n=1, sort_samples=false)
    # Parameters
    eps1 = eps(Float64);
    binsize = mean(diff(timet));
    M = length(timet);

    # compute mean and covariance in q-domain
    mq_new = mean(qn,2);
    mididx = round(Integer, M/2);
    m_new = sign(fn[mididx,:]) .* sqrt(abs(fn[mididx,:]));
    mqn = vec([mq_new; mean(m_new)]);
    qn2 = vcat(qn, m_new');
    C = Base.covm(qn2, mean(qn2,2), 2);

    q_s = mvnrand(mqn, C, n);

    # compute the correspondence to the orignal function domain
    fs = zeros(M, n);
    for k in 1:n
        fs[:, k] = cumtrapzmid(timet, q_s[1:M,k].*abs(q_s[1:M,k]),
                               sign(q_s[M+1,k])*(q_s[M+1,k]^2));
    end

    # random warping generation
    rgam = random_gamma(gam, n);
    gams = zeros(M, n);
    for k in 1:n
        gams[:,k] = invert_gamma(rgam[:, k]);
    end

    if sort_samples
        mx = maximum(fs,1);
        seq1 = sortperm(vec(mx));

        # compute the psi-function
        psi = zeros(size(rgam));
        for ii in 1:size(rgam, 2)
            fy = gradient(rgam[:, ii], binsize);
            psi[:,ii] = fy ./ sqrt(abs(fy) + eps1);
        end
        ip = zeros(n);
        len = zeros(n);

        for i in 1:n
            tmp = ones(1,M);
            psi_tmp = (psi[:,i]/M);
            ip[i] = (tmp * psi_tmp)[1];
            len[i] = acos(ip[i]);
        end

        seq2 = sortperm(len);

        # combine x-variability and y-variability
        ft = zeros(M, n);
        time0 = [0:(M-1)]./(M-1);
        for k in 1:n
            ft[:,k] = approx(time0, fs[:,seq1[k]], gams[:,seq2[k]]);
            tmp = isnan(ft[:, k]);
            while any(tmp)
                rgam2 = random_gamma(gam, 1);
                gam_tmp = invert_gamma(rgam2);
                ft[:, k] = approx(time0, fs[:,seq1[k]], gam_tmp);
            end
        end
    else
        # combine x-variability and y-variability
        ft = zeros(M, n);
        time0 = collect(0:(M-1))./(M-1);
        for k in 1:n
            ft[:,k] = approx(time0, fs[:,k], gams[:,k]);
            tmp = isnan(ft[:, k]);
            while any(tmp)
                rgam2 = random_gamma(gam, 1);
                gam_tmp = invert_gamma(rgam2);
                ft[:, k] = approx(time0, fs[:,k], gam_tmp);
            end
        end
    end

    out = Dict("fs" => fs, "gams" => rgam, "ft" => ft);
    return out

end


text: Dict{String,Any}("million" => Dict{String,Any}("value" => "1000000","type" => "float"),"minustenth" => Dict{String,Any}("value" => "-0.1","type" => "float"),"beast" => Dict{String,Any}("value" => "666","type" => "float"))

max_stars_repo_name: SmiteWindows/julia

max_issues_repo_name: SmiteWindows/julia

max_forks_repo_name: SmiteWindows/julia


text:
# Axis-aligned multivariate normal (i.e., independent entries, i.e., diagonal covariance matrix) using conjugate prior.
module MVNaaC

module MVNaaCmodel # submodule for component family definitions
export Theta, Data, log_marginal, new_theta, Theta_clear!, Theta_adjoin!, Theta_remove!,
       Hyperparameters, construct_hyperparameters, update_hyperparameters!

const Data = Array{Float64,1}

type Theta
    n::Int64                 # number of data points assigned to this cluster
    sum_x::Array{Float64,1}  # sum of the data points x assigned to this cluster
    sum_xx::Array{Float64,1} # sum of x.*x for the data points assigned to this cluster
    Theta(d) = (p=new(); p.n=0; p.sum_x=zeros(d); p.sum_xx=zeros(d); p)
end

new_theta(H) = Theta(H.d)

Theta_clear!(p) = (p.sum_x[:] = 0.; p.sum_xx[:] = 0.; p.n = 0)

Theta_adjoin!(p,x) = (for i=1:length(x); p.sum_x[i] += x[i]; p.sum_xx[i] += x[i]*x[i]; end; p.n += 1)
Theta_remove!(p,x) = (for i=1:length(x); p.sum_x[i] -= x[i]; p.sum_xx[i] -= x[i]*x[i]; end; p.n -= 1)

# In each dimension independently,
# X_1,...,X_n ~ Normal(mu,1/lambda) with Normal(mu|m,1/(c*lambda))Gamma(lambda|a,b) prior on mean=mu, precision=lambda.
function log_marginal(p,H)
    n = p.n
    LB = 0.0
    # For each dimension
    for i = 1:H.d
        ## update to b
        t1 = H.b + 0.5 * p.sum_xx[i] * H.alpha_wt ###
        t2 = 0.5 * p.sum_x[i]*p.sum_x[i] / n * H.alpha_wt
        t3 = 0.5 * H.c * n * H.alpha_wt * (p.sum_x[i]/n - H.m)^2 / (H.c + n * H.alpha_wt)
        #LB += log(H.b + 0.5 * p.sum_xx[i] - 0.5 * p.sum_x[i]*p.sum_x[i] / n + \
        #          0.5 * H.c * n * (p.sum_x[i]/n - H.m)^2/(H.c + n))
        LB += log(t1 - t2 + t3)
    end

    # update to a
    aa = H.a + 0.5 * n * H.alpha_wt ##
    # log term w/ update to c
    cc = 0.5 * log(H.c + n * H.alpha_wt) ##

    term = H.d * (H.constant - 0.5*n*log(2*pi)*H.alpha_wt - cc + H.log_Ga[p.n]) - aa * LB
    #return H.d * (H.constant - 0.5*n*log(2*pi) - 0.5*log(H.c+n) + H.log_Ga[n]) - (H.a+0.5*n) * LB
    return term
end

function log_marginal(x,p,H)
    Theta_adjoin!(p,x)
    result = log_marginal(p,H)
    Theta_remove!(p,x)
    return result
end

type Hyperparameters
    d::Int64    # dimension
    m::Float64  # prior mean of mu's
    c::Float64  # prior precision multiplier for mu's
    a::Float64  # prior shape of lambda's
    b::Float64  # prior rate of lambda's
    constant::Float64
    log_Ga::Array{Float64,1}
    alpha_wt::Float64 # alpha weight (not technically a hyperparameter)
end

function construct_hyperparameters(options)
    x = options.x
    n = length(x)
    d = length(x[1])
    mu = mean(x)
    v = mean([xi.*xi for xi in x]) - mu.*mu  # sample variance
    @assert(all(abs.(mu) .< 1e-10) && all(abs.(v - 1.0) .< 1e-10), "Data must be normalized to zero mean, unit variance.")
    m = 0.0
    c = 1.0
    a = 1.0
    b = 1.0

    # specify the weight
    alpha_eval = eval(parse(options.alpha_wt_str))
    alpha_wt_fn(n) = Base.invokelatest(alpha_eval, n)
    alpha_wt = alpha_wt_fn(options.n)

    log_Ga = lgamma.(a+0.5*(1:n+1) * alpha_wt)

    # normalizing constant of the log Normal-gamma pdf involving hyperparams
    constant = 0.5*log(c) + a*log(b) - lgamma(a)

    return Hyperparameters(d,m,c,a,b,constant,log_Ga, alpha_wt)
end

function update_hyperparameters!(H,theta,list,t,x,z)
    #error("update_hyperparameters is not yet implemented.")
end

end # module MVNaaCmodel
using .MVNaaCmodel

# Include generic code
include("generic.jl")

# Include core sampler code
include("coreConjugate.jl")

end # module MVNaaC


text:
"""
    rm_annotations(x)

Remove type annotation of given expression.
"""
function rm_annotations(x)
    x isa Expr || return x
    if x.head == :(::)
        if length(x.args) == 1 # anonymous
            return
        else
            return x.args[1]
        end
    elseif x.head in [:(=), :kw] # default values
        return rm_annotations(x.args[1])
    else
        return x
    end
end

function annotations(x)
    x isa Expr || return x
    if x.head == :(::)
        x.args[end]
    elseif x.head in [:(=), :kw]
        return annotations(x.args[1])
    else
        return x
    end
end

function splatting_variables(variables, free)
    Expr(:(=), Expr(:tuple, variables...), free)
end

export gate_count

function gate_count(spec)
    ci, _ = code_yao(Semantic.main, spec; optimize = true)
    return gate_count(ci)
end

function gate_count(ci::CodeInfo)
    count = Dict{Symbol,Any}()
    for stmt in ci.code
        is_quantum_statement(stmt) || continue
        qt = quantum_stmt_type(stmt)

        if qt === :gate || qt === :ctrl
            gc = get!(count, qt, IdDict{Any,Int}())
            if stmt.head === :invoke
                gate, gt = obtain_gate_stmt(stmt.args[3], ci)
            elseif stmt.head === :call
                gate, gt = obtain_gate_stmt(stmt.args[2], ci)
            end

            if gate isa IntrinsicRoutine && isempty(gate.variables)
                gc[gate] = get(gc, gate, 0) + 1
            else
                gc[parent(gate)] = get(gc, parent(gate), 0) + 1
            end
        else
            count[qt] = get(count, qt, 0) + 1
        end
    end
    return count
end

get_type(x::SSAValue, ci::CodeInfo) = widenconst(ci.ssavaluetypes[x.id])
get_type(x::SSAValue, ir::IRCode) = widenconst(ir.stmts.type[x.id])
get_type(x::SSAValue, ir::YaoIR) = get_type(x, ir.ir)

eachstmt(ci::CodeInfo) = enumerate(ci.code)
eachstmt(ir::IRCode) = enumerate(ir.stmts.inst)
eachstmt(ir::YaoIR) = eachstmt(ir.ir)

function count_qubits(ir)
    min_loc, max_loc = 0, 0
    for (v, e) in eachstmt(ir)
        minmax = find_minmax_locations(e, ir)

        # non-constant location
        if minmax === false
            return
        end

        if !isnothing(minmax)
            min_loc = min(minmax[1], min_loc)
            max_loc = max(minmax[2], max_loc)
        end
    end
    return max_loc - min_loc
end

function find_minmax_locations(@nospecialize(e), ir)
    if e isa Locations
        return minimum(e.storage), maximum(e.storage)
    elseif e isa Expr
        min_loc, max_loc = 0, 0
        for each in e.args
            minmax = find_minmax_locations(each, ir)
            if isnothing(minmax)
                continue
            elseif minmax === false
                return false
            else
                min_loc = min(minmax[1], min_loc)
                max_loc = max(minmax[2], max_loc)
            end
        end
        return min_loc, max_loc
        @show println(ir)
    elseif e isa SSAValue && get_type(e, ir) <: AbstractLocations
        return false
    else
        return
    end
end


text:
using NPZ
using JuMP
using Gurobi
using LinearAlgebra

function facto(m::Int64)
    factori = 1
    for j = 1:m
        factori = factori * j
    end
    return factori
end

function comb(n::Int64, m::Int64)
    combi = facto(n) / (facto(m) * facto(n-m))
    return convert(Int, combi)
end

function setUpFitting(deg::Int64, c::Float64)
	normCoeffs = Array{Int}(undef, deg+1)
	for i=1:deg + 1
		normCoeffs[i] = comb(deg, i-1) * c^(deg-i+1)
	end
	return normCoeffs
end

function read_np_data()
	A = npzread("tmp_jl/A.npz")["arr_0"]
	B = npzread("tmp_jl/B.npz")["arr_0"]
	C = npzread("tmp_jl/C.npz")["arr_0"]
	h = npzread("tmp_jl/h.npz")["arr_0"]
	flow_k = npzread("tmp_jl/flow_k.npz")["arr_0"]
	data = npzread("tmp_jl/data.npz")["arr_0"]
	dxdb_k = npzread("tmp_jl/dxdb.npz")["arr_0"]
	dxdg_k = npzread("tmp_jl/dxdg.npz")["arr_0"]
	g_k = npzread("tmp_jl/gk.npz")["arr_0"]
	beta_k = npzread("tmp_jl/betak.npz")["arr_0"]
	pars = npzread("tmp_jl/par.npz")["arr_0"]
	return A, B, C, h, flow_k, data, dxdg_k, dxdb_k, g_k, beta_k, pars
end

function jointBilevel()

	# Read optimization data
	A, B, C, h, flow_k, data, dxdg_k, dxdb_k, g_k, beta_k, pars = read_np_data()
	trust_region_g = pars[1,1]
	trust_region_beta = pars[2,1]
	scaling = pars[3,1]
	c = pars[4,1]
	lambda_1 = pars[5,1]
	numNodes = Int(pars[6,1])
	numLinks = Int(pars[7,1])
	numOD = Int(pars[8,1])
	numPoly = size(B,2)
	K = 1
	# create variables
	m  =  Model(with_optimizer(Gurobi.Optimizer, Presolve=0, OutputFlag=0))
	@variable(m, y[1:numOD,1:numNodes])
	@variable(m, g[1:numOD]>=0)
	@variable(m, epsilon[1:K]>=0)
	@variable(m, beta[1:numPoly])# >=0 )

	y_ = [y[i,j] for i = 1:numOD for j = 1:numNodes]

	# build IP-1 constraint
	constr = A*y_ + B*beta + C*epsilon.*scaling + h
	@constraint(m, constr .<= 0);
	@constraint(m, beta[1] == 1);

	# Dual variables constraint
	@variable(m, v[1:size(A,1)]>=0)
	@variable(m, ksi>=0)
	@constraint(m, A'*v .==0);

	# Primal Dual gap constraint
	# Creating H1 and H2
	H1 = Matrix(I, K, K);
	H2 = Matrix(Diagonal((1 ./setUpFitting(numPoly-1, c))));

	iH1 =  inv(H1)

	iH2 = inv(H2)

	M1 = (1/4)*C*iH1*C'  ;
	M1 = (M1 + M1')/2;
	M2 =  (1/4)*B*iH2*B'  ;
	M2 = (M2 + M2')/2 +1e-3*Matrix(I, size(M2,1), size(M2,1));

	qd = dot(beta, H2*beta)+ dot(v, M1*v) + dot(v, M2*v)
	ln = dot(-h,v)
	nlconstr = qd+ln
	@constraint(m, nlconstr <= ksi)

	# Creating Objective function
	DeltaF = zeros(numOD+numPoly+length(lambda_1) , 1);
	for i = 1 :numLinks
		DeltaF += [2*(flow_k[i] - data[i]).*dxdb_k[i,:] ;
					2*(flow_k[i] - data[i]).*dxdg_k[i,:] ;
					lambda_1];
	end

	# Constraints on trust region
	k=0
	for g_ka in g_k
	 	k=k+1
		@constraint(m, g[k] <= g_ka + trust_region_g)
		@constraint(m, g[k] >= g_ka - trust_region_g)
	end

	@constraint(m, beta_k - trust_region_beta .<= beta[1:numPoly] )
	@constraint(m, beta[1:numPoly] .<= beta_k + trust_region_beta)

	obj = DeltaF'*[(beta-beta_k) ; (g -g_k) ; ksi];
	#obj = DeltaF'*[(beta-beta_k) ; (g -g_k) ; (0)];
	@objective(m, Min, obj[1])

	optimize!(m)

	beta_k_1 = [value(beta[i]) for i =1:length(beta)];
	g_k_1 = [value(g[i]) for i =1:length(g)];
	ksi_k_1 = [value(ksi)];
	obj_k_1 = objective_value(m)

	npzwrite("tmp_jl/beta_k_1.npz", beta_k_1)
	npzwrite("tmp_jl/g_k_1.npz", g_k_1)
	npzwrite("tmp_jl/ksi_k_1.npz", ksi_k_1)
end

jointBilevel()

max_stars_repo_name: salomonw/tnet

max_issues_repo_name: salomonw/tnet

max_forks_repo_name: salomonw/tnet


text:
function parse_ast(code::AbstractString)
    start = 1
    args = []
    while true
        ex, next = Meta.parse(code, start)
        next === start && break
        start = next
        push!(args, ex)
    end
    return Expr(:block, args...)
end

function match_macro(prefix, name, ex)
    x = Symbol(name)
    return ex in (x, :($prefix.$x))
end

preprocess_ast(x) = x
function preprocess_ast(ex::Expr)
    #! format: off
    if isexpr(ex, :macrocall) && (
        ex.args[1] == GlobalRef(Core, Symbol("@doc")) ||
        match_macro(:Base, "@inline", ex.args[1]) ||
        match_macro(:Base, "@propagate_inbounds", ex.args[1]) ||
        match_macro(:Base, "@pure", ex.args[1]) ||
        match_macro(:Base, "@nospecialize", ex.args[1]) ||
        match_macro(:Base, "@static", ex.args[1]) ||
        match_macro(:Test, "@testset", ex.args[1])
    )
        return preprocess_ast(ex.args[end])
    elseif isexpr(ex, :macrocall) && match_macro(:Test, "@test", ex.args[1])
        return preprocess_ast(ex.args[3])
    elseif (
        isexpr(ex, :macrocall) &&
        match_macro(:Base, "@deprecate", ex.args[1]) &&
        length(ex.args) >= 4
    )
        # Handle: @deprecate old new
        return :($(ex.args[3]) = $(ex.args[4]))
    end
    #! format: on
    return Expr(ex.head, preprocess_ast.(ex.args)...)
end

vsolve(ex::Expr) = JuliaVariables.solve(ex)
vsolve(code::AbstractString) = vsolve(preprocess_ast(parse_ast(code)))

resolvedvar(x::Symbol) = x
function resolvedvar(scopedvar::JuliaVariables.ScopedVar)
    scope = scopedvar.scope
    sym = scopedvar.sym
    return scope[sym]
end

assymbol(x) = resolvedvar(x)::Symbol

isglobalvar(_) = false
# isglobalvar(::JuliaVariables.GlobalVar) = true
isglobalvar(scopedvar::JuliaVariables.ScopedVar) =
    resolvedvar(scopedvar) isa JuliaVariables.GlobalVar

function maybe_push_lvar!(vars, x)
    if isglobalvar(x)
        push!(vars.left, assymbol(x))
    end
    return vars
end

function maybe_push_rvar!(vars, x)
    if isglobalvar(x)
        push!(vars.right, assymbol(x))
    end
    return vars
end

global_variables!(vars, _) = vars
function global_variables!(vars, x::JuliaVariables.ScopedFunc)
    if isexpr(x.func, :function) || isexpr(x.func, :(=))
        sig = x.func.args[1]
        if isexpr(sig, :where)
            sig = sig.args[1]
        end
        if isexpr(sig, :call)
            if isexpr(sig.args[1], :.)
                # Handle: M.f() = ...
                maybe_push_rvar!(vars, sig.args[1].args[1])
            else
                maybe_push_lvar!(vars, sig.args[1])
            end
        end
    end
    global_variables!(vars, x.func.args[2])
    return vars
end

# function global_variables!(vars, x::JuliaVariables.GlobalVar)
#     push!(vars.right, assymbol(x))
#     return vars
# end

global_variables!(vars, x::JuliaVariables.ScopedVar) =
    maybe_push_rvar!(vars, x)

function global_variables!(vars, ex::Expr)
    if isexpr(ex, :(=))
        lvar = ex.args[1]
        if isexpr(lvar, :curly)
            lvar = lvar.args[1]
        end
        maybe_push_lvar!(vars, lvar)
        return global_variables!(vars, ex.args[2])
    elseif (isexpr(ex, :import) || isexpr(ex, :using))
        if isexpr(ex.args[1], :(:))
            # Add `x`, `y`, ... from `using M: x, y, ...`:
            args = [isexpr(x, :(.)) ? x.args[1] : x for x in ex.args[1].args[2:end]]
        else
            # Add `M` from `using M`:
            args = Any[ex.args[1].args[1]]
        end
        globals = Iterators.filter(isglobalvar, args)
        # globals = collect(globals)
        # @show globals
        mapfoldl(assymbol, push!, globals, init = vars.left)
        return vars
    elseif isexpr(ex, :struct)
        if isexpr(ex.args[2], :(<:))
            typeex = ex.args[2].args[1]
            global_variables!(vars, ex.args[2].args[2])
        else
            typeex = ex.args[2]
        end
        if isexpr(typeex, :curly)
            maybe_push_lvar!(vars, typeex.args[1])
        else
            maybe_push_lvar!(vars, typeex)
        end
        # It seems JuliaVariables detects fields as global variables.
        # Workaround it by not recursing inside.
        return vars
    elseif isexpr(ex, :abstract)
        lvar = ex.args[1]
        if isexpr(lvar, :<:)
            global_variables!(vars, lvar.args[2])
            lvar = lvar.args[1]
        end
        if isexpr(lvar, :curly)
            lvar = lvar.args[1]
        end
        maybe_push_lvar!(vars, lvar)
        return vars
    elseif isexpr(ex, :macro)
        # Macro arguments are treated as globals
        # https://github.com/thautwarm/JuliaVariables.jl/issues/10
        # So, let's not recurse into the expression manually.
        if isexpr(ex.args[1], :call)
            lvar = ex.args[1].args[1]
            if isglobalvar(lvar)
                s = assymbol(lvar)
                if !startswith(string(s), "@")
                    s = Symbol('@', s)
                end
                push!(vars.left, s)
            end
        end
        return vars
    elseif isexpr(ex, :macrocall)
        if ex.args[1] isa Symbol
            push!(vars.right, ex.args[1])
        end
    elseif isexpr(ex, :module)
        maybe_push_lvar!(vars, ex.args[2])
        # TODO: support sub-module
        # return vars
    elseif isexpr(ex, :.)
        return vars
    end
    return foldl(global_variables!, ex.args, init = vars)
end

global_variables(ex::Expr) = global_variables!((left = Set([]), right = Set([])), ex)
global_variables(code::AbstractString) = global_variables(vsolve(code))

merge_global_variables!(::Nothing, y) = y
function merge_global_variables!(x, y)
    union!(x.left, y.left)
    union!(x.right, y.right)
    return x
end

# TODO: handle submodules
module_variables(dir::AbstractString) = module_variables(sort!(jlfiles(dir)))
function module_variables(paths::AbstractVector{<:AbstractString})
    return mapfoldl(merge_global_variables!, paths; init = nothing) do p
        # @debug "Analyzing $p"
        try
            global_variables(read(p, String))
        catch err
            @error(
                "global_variables(read($(repr(p)), String))",
                exception = (err, catch_backtrace()),
            )
            global_variables("")
        end
    end
end

function module_variables(pkg::Base.PkgId)
    path = Base.locate_package(pkg)
    return module_variables(dirname(path))
end

function depsvars_from_tomlpath(tomlpath, imported_modules = ())
    deps = get(TOML.parsefile(tomlpath), "deps", Dict{String,Any}())
    if !isempty(imported_modules)
        imported_modules = unique(String.(imported_modules))
        notfound = setdiff(imported_modules, keys(deps))
        if !isempty(notfound)
            @warn "Packages not found in `$tomlpath`:\n$(join(notfound, "\n"))"
        end
        deps = filter(deps) do (pkgname, _)
            pkgname in imported_modules
        end
    end
    return depsvars_from_deps(deps)
end

function depsvars_from_deps(deps)
    depsvars = []
    for (pkgname, pkguuid) in deps
        pkg = Base.PkgId(Base.UUID(pkguuid), pkgname)
        @debug "Analyzing $pkg..."
        push!(depsvars, (pkg, module_variables(pkg)))
    end
    sort!(depsvars, by = ((pkg, _),) -> pkg.name)
    # TODO: make it `export`-aware
    return depsvars
end

__imported_modules(_) = Symbol[]
function __imported_modules(ex::Expr)::Vector{Symbol}
    @match ex begin
        :(using $name:$(_...)) || :(import $name:$(_...)) => return [name]
        :(using $name) || :(import $name) => [name]

        # Handle: using A.B: C, D
        :(using $(names...):$(_...)) || :(import $(names...):$(_...)) => [names[1]]
        Expr(:using, Expr(:(:), x, _...)) || Expr(:import, Expr(:(:), x, _...)) => begin
            @match x begin
                Expr(:., name, _...) => [name]
                Expr(:., :., _...) => []  # using .A: B
                _ => error("Unexpected `using`/`import` argument `$x` in:\n", ex)
            end
        end

        # Handle: using A.B.C.D
        :(using $(names...)) || :(import $(names...)) => [names[1]]

        # Handle: using A, B.C, D
        Expr(:using, args...) || Expr(:import, args...) => begin
            map(args) do x
                @match x begin
                    Expr(:., name, _...) => name
                    _ => error("Unexpected `using`/`import` argument `$x` in:\n", ex)
                end
            end
        end

        # Recurse:
        _ => mapreduce(__imported_modules, vcat, ex.args; init = Symbol[])
    end
end

_imported_modules(code::AbstractString) = __imported_modules(parse_ast(code))

imported_modules(ex::Expr) = __imported_modules(ex)
imported_modules(path::AbstractString) = _imported_modules(read(path, String))
imported_modules(paths::AbstractVector{<:AbstractString}) =
    mapreduce(imported_modules, vcat, paths; init = Symbol[])


text:
using Pkg
Pkg.clone("https://github.com/sisl/MultivariateTimeSeries.jl")


text:
function aggregate(::Type{StatisticalSummary}, arg0::Iterable)
    return jcall(StatisticalSummary, "aggregate", StatisticalSummary, (Iterable,), arg0)
end

function aggregate(::Type{StatisticalSummary}, arg0::Vector{StatisticalSummary})
    return jcall(StatisticalSummary, "aggregate", StatisticalSummary, (Vector{StatisticalSummary},), arg0)
end

function get_max(obj::StatisticalSummary)
    return jcall(obj, "getMax", jdouble, ())
end

function get_mean(obj::StatisticalSummary)
    return jcall(obj, "getMean", jdouble, ())
end

function get_min(obj::StatisticalSummary)
    return jcall(obj, "getMin", jdouble, ())
end

function get_n(obj::StatisticalSummary)
    return jcall(obj, "getN", jlong, ())
end

function get_standard_deviation(obj::StatisticalSummary)
    return jcall(obj, "getStandardDeviation", jdouble, ())
end

function get_sum(obj::StatisticalSummary)
    return jcall(obj, "getSum", jdouble, ())
end

function get_variance(obj::StatisticalSummary)
    return jcall(obj, "getVariance", jdouble, ())
end