abstract type AbstractLoadFormulation <: AbstractDeviceFormulation end

abstract type AbstractControllablePowerLoadFormulation <: AbstractLoadFormulation end

struct StaticPowerLoad <: AbstractLoadFormulation end

struct InterruptiblePowerLoad <: AbstractControllablePowerLoadFormulation end

struct DispatchablePowerLoad <: AbstractControllablePowerLoadFormulation end

########################### dispatchable load variables ############################################

function activepower_variables(canonical_model::CanonicalModel,
                               devices::IS.FlattenIteratorWrapper{L}) where {L<:PSY.ElectricLoad}
    add_variable(canonical_model,
                 devices,
                 Symbol("P_$(L)"),
                 false,
                 :nodal_balance_active, -1.0;
                 ub_value = x -> PSY.get_maxactivepower(x),
                 lb_value = x -> 0.0)

    return

end

function reactivepower_variables(canonical_model::CanonicalModel,
                                 devices::IS.FlattenIteratorWrapper{L}) where {L<:PSY.ElectricLoad}
    add_variable(canonical_model,
                 devices,
                 Symbol("Q_$(L)"),
                 false,
                 :nodal_balance_reactive, -1.0;
                 ub_value = x -> PSY.get_maxreactivepower(x),
                 lb_value = x -> 0.0)

    return

end

function commitment_variables(canonical_model::CanonicalModel,
                              devices::IS.FlattenIteratorWrapper{L}) where {L<:PSY.ElectricLoad}

    add_variable(canonical_model,
                 devices,
                 Symbol("ON_$(L)"),
                 true)

    return

end

####################################### Reactive Power Constraints ######################################
"""
Reactive Power Constraints on Loads Assume Constant PowerFactor
"""
function reactivepower_constraints(canonical_model::CanonicalModel,
                                   devices::IS.FlattenIteratorWrapper{L},
                                   device_formulation::Type{D},
                                   system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                       D<:AbstractControllablePowerLoadFormulation,
                                                                       S<:PM.AbstractPowerModel}
    time_steps = model_time_steps(canonical_model)
    key = Symbol("reactive_$(L)")
    canonical_model.constraints[key] = JuMPConstraintArray(undef, (PSY.get_name(d) for d in devices), time_steps)

    for t in time_steps, d in devices
        name = PSY.get_name(d)
        pf = sin(atan((PSY.get_maxreactivepower(d)/PSY.get_maxactivepower(d))))
        canonical_model.constraints[key][PSY.get_name(d), t] = JuMP.@constraint(canonical_model.JuMPmodel,
                        canonical_model.variables[Symbol("Q_$(L)")][name, t] == canonical_model.variables[Symbol("P_$(L)")][name, t]*pf)
    end

    return

end

######################## output constraints without Time Series ###################################
function _get_time_series(devices::IS.FlattenIteratorWrapper{T},
                          time_steps::UnitRange{Int64}) where {T<:PSY.ElectricLoad}

    names = Vector{String}(undef, length(devices))
    series = Vector{Vector{Float64}}(undef, length(devices))

    for (ix, d) in enumerate(devices)
        names[ix] = PSY.get_name(d)
        series[ix] = fill(PSY.get_maxactivepower(d), (time_steps[end]))
    end

    return names, series

end

function activepower_constraints(canonical_model::CanonicalModel,
                                 devices::IS.FlattenIteratorWrapper{L},
                                 device_formulation::Type{DispatchablePowerLoad},
                                 system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                     S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)

    if model_has_parameters(canonical_model)
        device_timeseries_param_ub(canonical_model,
                                   _get_time_series(devices, time_steps),
                                   Symbol("active_$(L)"),
                                   UpdateRef{L}(Symbol("P_$(L)")),
                                   Symbol("P_$(L)"))
    else
        range_data = [(PSY.get_name(d), (min = 0.0, max = PSY.get_maxactivepower(d))) for d in devices]
        device_range(canonical_model,
                    range_data,
                    Symbol("activerange_$(L)"),
                    Symbol("P_$(L)")
                    )
    end

    return

end

"""
This function works only if the the Param_L <= PSY.get_maxactivepower(g)
"""
function activepower_constraints(canonical_model::CanonicalModel,
                                 devices::IS.FlattenIteratorWrapper{L},
                                 device_formulation::Type{InterruptiblePowerLoad},
                                 system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                          S<:PM.AbstractPowerModel}
    time_steps = model_time_steps(canonical_model)

    if model_has_parameters(canonical_model)
        device_timeseries_ub_bigM(canonical_model,
                                 _get_time_series(devices, time_steps),
                                 Symbol("active_$(L)"),
                                 Symbol("P_$(L)"),
                                 UpdateRef{L}(Symbol("P_$(L)")),
                                 Symbol("ON_$(L)"))
    else
        device_timeseries_ub_bin(canonical_model,
                                _get_time_series(devices, time_steps),
                                Symbol("active_$(L)"),
                                Symbol("P_$(L)"),
                                Symbol("ON_$(L)"))
    end

    return

end

######################### output constraints with Time Series ##############################################

function _get_time_series(forecasts::Vector{PSY.Deterministic{L}}) where {L<:PSY.ElectricLoad}

    names = Vector{String}(undef, length(forecasts))
    ratings = Vector{Float64}(undef, length(forecasts))
    series = Vector{Vector{Float64}}(undef, length(forecasts))

    for (ix, f) in enumerate(forecasts)
        component = PSY.get_component(f)
        names[ix] = PSY.get_name(component)
        series[ix] = values(PSY.get_timeseries(f))
        ratings[ix] = PSY.get_maxactivepower(component)
    end

    return names, ratings, series

end

function activepower_constraints(canonical_model::CanonicalModel,
                                 devices::Vector{PSY.Deterministic{L}},
                                 device_formulation::Type{DispatchablePowerLoad},
                                 system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                     S<:PM.AbstractPowerModel}

    if model_has_parameters(canonical_model)
        device_timeseries_param_ub(canonical_model,
                                   _get_time_series(devices),
                                   Symbol("active_$(L)"),
                                   UpdateRef{L}(Symbol("P_$(L)")),
                                   Symbol("P_$(L)"))
    else
        device_timeseries_ub(canonical_model,
                            _get_time_series(devices),
                            Symbol("active_$(L)"),
                            Symbol("P_$(L)"))
    end

    return

end

function activepower_constraints(canonical_model::CanonicalModel,
                                 devices::Vector{PSY.Deterministic{L}},
                                 device_formulation::Type{InterruptiblePowerLoad},
                                 system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                     S<:PM.AbstractPowerModel}

    if model_has_parameters(canonical_model)
        device_timeseries_ub_bigM(canonical_model,
                                 _get_time_series(devices),
                                 Symbol("active_$(L)"),
                                 Symbol("P_$(L)"),
                                 UpdateRef{L}(Symbol("P_$(L)")),
                                 Symbol("ON_$(L)"))
    else
        device_timeseries_ub_bin(canonical_model,
                                _get_time_series(devices),
                                Symbol("active_$(L)"),
                                Symbol("P_$(L)"),
                                Symbol("ON_$(L)"))
    end

    return

end

############################ injection expression with parameters ####################################

########################################### Devices ####################################################

function _nodal_expression_param(canonical_model::CanonicalModel,
                                devices::IS.FlattenIteratorWrapper{L},
                                system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                    S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)
    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(devices))
    ts_data_reactive = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(devices))

    for (ix, d) in enumerate(devices)
        bus_number = PSY.get_number(PSY.get_bus(d))
        name = PSY.get_name(d)
        active_power = PSY.get_activepower(d)
        reactive_power = PSY.get_reactivepower(d)
        time_series_vector = ones(time_steps[end])
        ts_data_active[ix] = (name, bus_number, active_power, time_series_vector, -1.0)
        ts_data_reactive[ix] = (name, bus_number, reactive_power, time_series_vector, -1.0)
    end

    include_parameters(canonical_model,
                  ts_data_active,
                  UpdateRef{L}(Symbol("P_$(L)")),
                  :nodal_balance_active)
    include_parameters(canonical_model,
                   ts_data_reactive,
                   UpdateRef{L}(Symbol("Q_$(L)")),
                   :nodal_balance_reactive)

    return

end

function _nodal_expression_param(canonical_model::CanonicalModel,
                                devices::IS.FlattenIteratorWrapper{L},
                                system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                    S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical_model)
    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(devices))

    for (ix, d) in enumerate(devices)
        bus_number = PSY.get_number(PSY.get_bus(d))
        name = PSY.get_name(d)
        active_power = PSY.get_activepower(d)
        time_series_vector = ones(time_steps[end])
        ts_data_active[ix] = (name, bus_number, active_power, time_series_vector, -1.0)
    end

    include_parameters(canonical_model,
                  ts_data_active,
                  UpdateRef{L}(Symbol("P_$(L)")),
                  :nodal_balance_active)

    return

end

############################################## Time Series ###################################
function _nodal_expression_param(canonical_model::CanonicalModel,
                                forecasts::Vector{PSY.Deterministic{L}},
                                system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                    S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)

    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(forecasts))
    ts_data_reactive = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(forecasts))

    for (ix, f) in enumerate(forecasts)
        device = PSY.get_component(f)
        bus_number = PSY.get_number(PSY.get_bus(device))
        name = PSY.get_name(device)
        active_power = PSY.get_maxactivepower(device)
        reactive_power = PSY.get_maxreactivepower(device)
        time_series_vector = values(PSY.get_timeseries(f))
        ts_data_active[ix] = (name, bus_number, active_power, time_series_vector)
        ts_data_reactive[ix] = (name, bus_number, reactive_power, time_series_vector)
    end

    include_parameters(canonical_model,
                    ts_data_active,
                    UpdateRef{L}(Symbol("P_$(L)")),
                    :nodal_balance_active,
                    -1.0)
    include_parameters(canonical_model,
                    ts_data_reactive,
                    UpdateRef{L}(Symbol("Q_$(L)")),
                    :nodal_balance_reactive,
                    -1.0)

    return

end

function _nodal_expression_param(canonical_model::CanonicalModel,
                                forecasts::Vector{PSY.Deterministic{L}},
                                system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                    S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical_model)
    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(forecasts))

    for (ix, f) in enumerate(forecasts)
        device = PSY.get_component(f)
        bus_number = PSY.get_number(PSY.get_bus(device))
        name = PSY.get_name(device)
        active_power = PSY.get_maxactivepower(device)
        time_series_vector = values(PSY.get_timeseries(f))
        ts_data_active[ix] = (name, bus_number, active_power, time_series_vector)
    end

    include_parameters(canonical_model,
                    ts_data_active,
                    UpdateRef{L}(Symbol("P_$(L)")),
                    :nodal_balance_active,
                    -1.0)

    return

end

############################ injection expression with fixed values ####################################

########################################### Devices ####################################################

function _nodal_expression_fixed(canonical_model::CanonicalModel,
                                devices::IS.FlattenIteratorWrapper{L},
                                system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                    S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)

    for t in time_steps, d in devices
        bus_number = PSY.get_number(PSY.get_bus(d))
        active_power = PSY.get_activepower(d)
        reactive_power = PSY.get_reactivepower(d)
        _add_to_expression!(canonical_model.expressions[:nodal_balance_active],
                            bus_number,
                            t,
                            -1*active_power);
        _add_to_expression!(canonical_model.expressions[:nodal_balance_reactive],
                            bus_number,
                            t,
                            -1*reactive_power);
    end

    return

end

function _nodal_expression_fixed(canonical_model::CanonicalModel,
                                devices::IS.FlattenIteratorWrapper{L},
                                system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                    S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical_model)

    for t in time_steps, d in devices
        bus_number = PSY.get_number(PSY.get_bus(d))
        active_power = PSY.get_activepower(d)
        _add_to_expression!(canonical_model.expressions[:nodal_balance_active],
                            bus_number,
                            t,
                            -1*active_power);
    end

    return

end

############################################## Time Series ###################################

function _nodal_expression_fixed(canonical_model::CanonicalModel,
                                forecasts::Vector{PSY.Deterministic{L}},
                                system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                    S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)

    for f in forecasts
        device = PSY.get_component(f)
        bus_number = PSY.get_number(PSY.get_bus(device))
        active_power = PSY.get_maxactivepower(device)
        reactive_power = PSY.get_maxreactivepower(device)
        time_series_vector = values(PSY.get_timeseries(f))
        for t in time_steps
            bus_number = PSY.get_number(PSY.get_bus(device))
            _add_to_expression!(canonical_model.expressions[:nodal_balance_active],
                                bus_number,
                                t,
                                -1 * time_series_vector[t] * active_power)
            _add_to_expression!(canonical_model.expressions[:nodal_balance_reactive],
                                bus_number,
                                t,
                                -1 * time_series_vector[t] * reactive_power)
        end
    end

    return

end

function _nodal_expression_fixed(canonical_model::CanonicalModel,
                                forecasts::Vector{PSY.Deterministic{L}},
                                system_formulation::Type{S}) where {L<:PSY.ElectricLoad,
                                                                    S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical_model)

    for f in forecasts
        device = PSY.get_component(f)
        bus_number = PSY.get_number(PSY.get_bus(device))
        active_power = PSY.get_maxactivepower(device)
        time_series_vector = values(PSY.get_timeseries(f))
        for t in time_steps
            bus_number = PSY.get_number(PSY.get_bus(device))
            _add_to_expression!(canonical_model.expressions[:nodal_balance_active],
                                bus_number,
                                t,
                                -1 * time_series_vector[t] * active_power)
        end
    end

    return

end

##################################### FormulationControllable Load Cost ######################################

function cost_function(canonical_model::CanonicalModel,
                       devices::IS.FlattenIteratorWrapper{L},
                       device_formulation::Type{DispatchablePowerLoad},
                       system_formulation::Type{S}) where {L<:PSY.ControllableLoad,
                                                           S<:PM.AbstractPowerModel}

    add_to_cost(canonical_model,
                devices,
                Symbol("P_$(L)"),
                :variable,
                -1.0)

    return

end

function cost_function(canonical_model::CanonicalModel,
                       devices::IS.FlattenIteratorWrapper{L},
                       device_formulation::Type{InterruptiblePowerLoad},
                       system_formulation::Type{S}) where {L<:PSY.ControllableLoad,
                                                           S<:PM.AbstractPowerModel}

    add_to_cost(canonical_model,
                devices,
                Symbol("ON_$(L)"),
                :fixed,
                -1.0)

    return

end


abstract type AbstractPlotList{T<:Tuple} end

plottype(::Type{<:AbstractPlotList{T}}) where {T} = T.parameters

"""
`Plotlist(plots...)`

Experimental feature. Create an object that can encode multiple series.
"""
struct PlotList{T<:Tuple} <: AbstractPlotList{T}
    plots::Tuple
    function PlotList(plots...)
        T = Tuple{unique(plottype.(plots))...}
        new{T}(plots)
    end
end

Base.getindex(m::PlotList, i) = getindex(m.plots, i)
Base.length(m::PlotList) = length(m.plots)
Base.iterate(m::PlotList, args...) = iterate(m.plots, args...)

@recipe(MultiplePlot) do scene
    default_theme(scene)
end

function default_theme(scene, ::Type{<:Combined{multipleplot, Tuple{P}}}) where {P<:AbstractPlotList}
    merge((default_theme(scene, pt) for pt in plottype(P))...)
end
# Allow MultiplePlot to prevail on user input: the plot type of each series will be defined in convert_arguments
plottype(::Type{<: Combined{Any}}, A::Type{<:MultiplePlot}, argvalues...) = A
plottype(::Type{<: Combined{T}}, A::Type{<:MultiplePlot}, argvalues...) where T = A

function convert_arguments(P::PlotFunc, m::PlotList)
    function convert_series(plot::PlotSpec)
        ptype = plottype(P, plottype(plot))
        to_plotspec(ptype, convert_arguments(ptype, plot.args...); plot.kwargs...)
    end
    pl = PlotList(convert_series.(m.plots)...)
    PlotSpec{MultiplePlot}(pl)
end

combine(val1, val2; palette = nothing) = val2

function combine!(theme1::Theme, theme2::Theme)
    palette = get(theme1, :palette, current_default_theme()[:palette])
    for (key, val) in theme2
        tv = convert(Node, get(theme1, key, nothing))
        pv = convert(Node, get(palette, key, nothing))
        theme1[key] = lift((t, p, v) -> combine(t, v, palette = p), tv, pv, val)
    end
    theme1
end
combine(theme1::Theme, theme2) = combine!(copy(theme1), theme2)

# This allows plotting an arbitrary combination of series form one argument
# The recipe framework can be constructed using this as a building block and computing
# PlotList with convert_arguments
function plot!(p::Combined{multipleplot, <:Tuple{PlotList}})
    mp = to_value(p[1]) # TODO how to preserve interactivity here, as number of series may change?
    theme = Theme(p)
    for s in mp.plots
        attr = combine(theme, Theme(; s.kwargs...))
        plot!(p, plottype(s), attr, s.args...)
    end
end

function default_theme(scene, ::Type{<:Combined{S, T}}) where {S<:Tuple, T}
    merge((default_theme(scene, Combined{pt}) for pt in S.parameters)...)
end

function plot!(p::Combined{S, <:Tuple{Vararg{Any, N}}}) where {S <: Tuple, N}
    for pt in S.parameters
        plot!(p, Combined{pt}, Theme(p), p[1:N]...)
    end
end

function Base.:*(::Type{<:Combined{S}}, ::Type{<:Combined{T}}) where {S, T}
    params1 = S isa Type{<:Tuple} ? S.parameters : [S]
    params2 = T isa Type{<:Tuple} ? T.parameters : [T]
    params = union(params1, params2)
    Combined{Tuple{params...}}
end


#=

This file specifies how to pretty-print Scattering module types

=#

function Base.show(io::IO, ::MIME"text/plain", x::Aerosol)

    println(io, "\n------------------")
    println(io, "Aerosol Struct")
    println(io, "------------------")
    println(io, "Size distribution: $(x.size_distribution)")
    println(io, "Real refractive Index nᵣ: $(x.nᵣ)")
    println(io, "Imaginary refractive Index nᵣ: $(x.nᵢ)")
    # start,stop = quantile(x.size_distribution,[0.0025,0.9975])
    # x_axis = collect(range(start,stop,20))
    # y = pdf.(x.size_distribution, x_axis)
    # print(io, UnicodePlots.lineplot(x_axis, y,title="Aerosol Size Distribution", xlabel="Radius (μm)" , ylabel="Frequency"))
end

# Overload the show method for ScatteringMatrix
function Base.show(io::IO, x::ScatteringMatrix)

    println(io, "  f₁₁| $(length(x.f₁₁))-element vector")
    println(io, "  f₁₂| $(length(x.f₁₂))-element vector")
    println(io, "  f₂₂| $(length(x.f₂₂))-element vector")
    println(io, "  f₃₃| $(length(x.f₃₃))-element vector")
    println(io, "  f₃₄| $(length(x.f₃₄))-element vector")
    print(io, "  f₄₄| $(length(x.f₄₄))-element vector")
end


"""
Calculate value, gradient, and hessian of the variational ELBO.
"""
module DeterministicVI

using Base.Threads: threadid, nthreads

import ..Config
using ..BivariateNormals: BivariateNormalDerivatives, BvnComponent,
                          GalaxySigmaDerivs, get_bvn_cov, eval_bvn_pdf!,
                          get_bvn_derivs!, transform_bvn_derivs!
using ..Model
using ..Model: SkyPatch, BvnBundle, populate_fsm!
import ..Celeste: Const, @aliasscope, @unroll_loop
import ..Infer
using ..SensitiveFloats
import ..Log
using ..Transform
import DataFrames
import Optim
import ForwardDiff.Dual
using StaticArrays
import Base: convert

export ElboArgs, generic_init_source, catalog_init_source, init_sources,
       VariationalParams, elbo, ElboIntermediateVariables

function init_thread_pool!(pool::Vector, create)
    if length(pool) != Base.Threads.nthreads()
        empty!(pool)
        for i = 1:Base.Threads.nthreads()
            push!(pool, create())
        end
    end
end

"""
Return a default-initialized VariationalParams instance.
"""
function generic_init_source(init_pos::Vector{Float64})
    ret = Vector{Float64}(length(CanonicalParams))
    ret[ids.is_star] = 0.5
    ret[ids.pos] = init_pos
    ret[ids.flux_loc] = log(2.0)
    ret[ids.flux_scale] = 1e-3
    ret[ids.gal_fracdev] = 0.5
    ret[ids.gal_ab] = 0.5
    ret[ids.gal_angle] = 0.0
    ret[ids.gal_scale] = 1.0
    ret[ids.k] = 1.0 / size(ids.k, 1)
    ret[ids.color_mean] = 0.0
    ret[ids.color_var] =  1e-2
    ret
end

"""
Return VariationalParams instance initialized form a catalog entry
"""
function catalog_init_source(ce::CatalogEntry; max_gal_scale=Inf)
    # TODO: sync this up with the transform bounds
    ret = generic_init_source(ce.pos)

    # TODO: don't do this thresholding for background sources,
    # just for sources that are being optimized
    ret[ids.is_star[1]] = ce.is_star ? 0.8: 0.2
    ret[ids.is_star[2]] = ce.is_star ? 0.2: 0.8

    ret[ids.flux_loc[1]] = log(max(0.1, ce.star_fluxes[3]))
    ret[ids.flux_loc[2]] = log(max(0.1, ce.gal_fluxes[3]))

    function get_color(color_var, color_mean)
        color_var > 0 && color_mean > 0 ? min(max(log(color_var / color_mean), -9.), 9.) :
            color_var > 0 && color_mean <= 0 ? 3.0 :
                color_var <= 0 && color_mean > 0 ? -3.0 : 0.0
    end

    function get_colors(raw_fluxes)
        [get_color(raw_fluxes[c+1], raw_fluxes[c]) for c in 1:4]
    end

    ret[ids.color_mean[:, 1]] = get_colors(ce.star_fluxes)
    ret[ids.color_mean[:, 2]] = get_colors(ce.gal_fluxes)

    ret[ids.gal_fracdev] = min(max(ce.gal_frac_dev, 0.015), 0.985)

    ret[ids.gal_ab] = ce.is_star ? .8 : min(max(ce.gal_ab, 0.015), 0.985)
    ret[ids.gal_angle] = ce.gal_angle
    ret[ids.gal_scale] = ce.is_star ? 0.2 : min(max_gal_scale, max(ce.gal_scale, 0.2))

    ret
end

function init_sources(target_sources::Vector{Int}, catalog::Vector{CatalogEntry})
    ret = Vector{Vector{Float64}}(length(catalog))
    for s in 1:length(catalog)
        ret[s] = catalog_init_source(catalog[s])
    end
    for s in target_sources
        ret[s][:] = generic_init_source(catalog[s].pos)
    end
    ret
end

include("deterministic_vi/elbo_args.jl")
include("deterministic_vi/elbo_kl.jl")
include("deterministic_vi/source_brightness.jl")
include("deterministic_vi/elbo_objective.jl")
include("deterministic_vi/ConstraintTransforms.jl")
include("deterministic_vi/ElboMaximize.jl")

"""
Infers one light source. This routine is intended to be called in parallel,
once per target light source.

Arguments:
    images: a collection of astronomical images
    neighbors: the other light sources near `entry`
    entry: the source to infer
"""
function infer_source(config::Config,
                      images::Vector{Image},
                      neighbors::Vector{CatalogEntry},
                      entry::CatalogEntry)
    if length(neighbors) > 100
        msg = string("object at RA, Dec = $(entry.pos) has an excessive",
                     "number ($(length(neighbors))) of neighbors")
        Log.warn(msg)
    end

    # It's a bit inefficient to call the next 5 lines every time we optimize_f.
    # But, as long as runtime is dominated by the call to maximize!, that
    # isn't a big deal.
    cat_local = vcat([entry], neighbors)
    vp = init_sources([1], cat_local)
    patches = Infer.get_sky_patches(images, cat_local)
    Infer.load_active_pixels!(config, images, patches)

    ea = ElboArgs(images, patches, [1])
    f_evals, max_f, max_x, nm_result = ElboMaximize.maximize!(ea, vp)
    return vp[1]
end

# legacy wrapper
function infer_source(images::Vector{Image},
                      neighbors::Vector{CatalogEntry},
                      entry::CatalogEntry)
    infer_source(Config(), images, neighbors, entry)
end

end


function construct_device!(canonical::CanonicalModel, sys::PSY.System,
                           model::DeviceModel{St, D},
                           ::Type{S};
                           kwargs...) where {St<:PSY.Storage,
                                             D<:AbstractStorageFormulation,
                                             S<:PM.AbstractPowerModel}

    devices = PSY.get_components(St, sys)

    if validate_available_devices(devices, St)
        return
    end

    parameters = get(kwargs, :parameters, true)

    #Variables
    active_power_variables(canonical, devices)

    reactive_power_variables(canonical, devices)

    energy_storage_variables(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, D)

    #Constraints
    active_power_constraints(canonical, devices, D, S)

    reactive_power_constraints(canonical, devices, D, S)

    energy_capacity_constraints(canonical, devices, D, S)

    feedforward!(canonical, St, model.feedforward)

    # Energy Balanace limits
    energy_balance_constraint(canonical, devices, D, S)

    return

end

function construct_device!(canonical::CanonicalModel, sys::PSY.System,
                           model::DeviceModel{St, D},
                           ::Type{S};
                           kwargs...) where {St<:PSY.Storage,
                                             D<:AbstractStorageFormulation,
                                             S<:PM.AbstractActivePowerModel}

    devices = PSY.get_components(St, sys)

    if validate_available_devices(devices, St)
        return
    end

    parameters = get(kwargs, :parameters, true)

    #Variables
    active_power_variables(canonical, devices);

    energy_storage_variables(canonical, devices);

    #Initial Conditions

    initial_conditions!(canonical, devices, D)

    #Constraints
    active_power_constraints(canonical, devices, D, S)

    energy_capacity_constraints(canonical, devices, D, S)

    feedforward!(canonical, St, model.feedforward)

    # Energy Balanace limits
    energy_balance_constraint(canonical, devices, D, S)

    return

end

function construct_device!(canonical::CanonicalModel, sys::PSY.System,
                                        model::DeviceModel{St, BookKeepingwReservation},
                                        ::Type{S};
                                        kwargs...) where {St<:PSY.Storage,
                                                          S<:PM.AbstractPowerModel}

    devices = PSY.get_components(St, sys)

    if validate_available_devices(devices, St)
        return
    end

    parameters = get(kwargs, :parameters, true)

    #Variables
    active_power_variables(canonical, devices)

    reactive_power_variables(canonical, devices)

    energy_storage_variables(canonical, devices)

    storage_reservation_variables(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, model.formulation)

    #Constraints
    active_power_constraints(canonical, devices, model.formulation, S)

    reactive_power_constraints(canonical, devices, model.formulation, S)

    energy_capacity_constraints(canonical, devices, model.formulation, S)

    feedforward!(canonical, St, model.feedforward)

    # Energy Balanace limits
    energy_balance_constraint(canonical, devices, model.formulation, S)

    return

end

function construct_device!(canonical::CanonicalModel, sys::PSY.System,
                           model::DeviceModel{St, BookKeepingwReservation},
                           ::Type{S};
                           kwargs...) where {St<:PSY.Storage,
                                             S<:PM.AbstractActivePowerModel}

    devices = PSY.get_components(St, sys)

    if validate_available_devices(devices, St)
        return
    end

    parameters = get(kwargs, :parameters, true)

    #Variables
    active_power_variables(canonical, devices)

    energy_storage_variables(canonical, devices)

    storage_reservation_variables(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, model.formulation)

    #Constraints
    active_power_constraints(canonical, devices, model.formulation, S)

    energy_capacity_constraints(canonical, devices, model.formulation, S)

    feedforward!(canonical, St, model.feedforward)

    # Energy Balanace limits
    energy_balance_constraint(canonical, devices, model.formulation, S)

    return

end


# runtests.jl

using Test: @test, @testset, @test_throws, @inferred, detect_ambiguities
using Sound # sound, soundsc, record, pick_output
using SampledSignals: SampleBuf
using PortAudio: devices, PortAudioDevice

# include("aqua.jl") # fails in Interval, Unitful ...

@show devices()

if isempty(devices())
    @warn "No devices so no tests on this OS."
else

@testset "output" begin
    dev = @inferred Sound.get_default_output_device()
    @test dev isa PortAudioDevice

    index = findfirst(==(dev), devices())
    io_out = IOBuffer()
    io_in = IOBuffer("$index")
    tmp = @inferred Sound.pick_output( ; io_in, io_out)
    @test tmp === dev

    io_in = IOBuffer("\n")
    tmp = @inferred Sound.pick_output( ; io_in, io_out)
    @test tmp === dev

    io_in = IOBuffer("$index")
    sound(:pick, randn(500), 1000; io_in, io_out)
    sound(index, randn(500), 1000)
    @test_throws String sound(0, randn(500), 1000)
end

@testset "Sound" begin
    S = 8192 # sampling rate in Hz
    x = 0.8 * cos.(2π*(1:800)*440/S)
    y = 0.7 * sin.(2π*(1:800)*660/S)
    sound(x, S) # specify sampling rate
    @test_throws ErrorException sound(y) # unknown sampling rate
    sound([x y], S) # stereo
    soundsc(0.1 * reshape(x,1,:), S) # scale and row vector

    sb = SampleBuf(x, S)
    sound(sb)
    soundsc(sb)

    @test_throws String sound(ones(5,3), S) # array size

    @test isempty(detect_ambiguities(Sound))
end

@testset "record" begin
    dev = @inferred Sound.get_default_input_device()
    @test dev isa PortAudioDevice

    data, S = record(0.001)
    @test S isa Real
    @test data isa Vector
end

end # !isempty(devices())

include("hann.jl")
include("findpeaks.jl")
include("phase_vocoder.jl")


function analyze()

  measure, measureinfo = loadmeasure()
  @unpack prefix, analysispath, outputpath, resultid = measureinfo

  try
    PARAM[:refreshverify] && verifymeasure(; measure, measureinfo)
    PARAM[:refreshleadlag] && analyzeleadlag(; measure, measureinfo)

  catch err

    mv(outputpath, "$(analysispath)\\$prefix\\ERR_$resultid")
    #show(stdout, MIME"text/plain"(), stacktrace(catch_backtrace()))
    @error "Verificaiton failed" exception=(err, catch_backtrace())
    throw("Verification failed")
  end

end

#performs file io and initial setup for a new measure estimation
function loadmeasure(;
    burnin::DatePeriod = PARAM[:measureburnin],
    measurefilename::String = PARAM[:analysismeasurefilename],
    analysispath = PARAM[:analysispath],
    returntype::String = PARAM[:crspdatalabel],
    workingpath::String = PARAM[:workingpath],
    hfrname::String = PARAM[:hfrname],
    mfname::String = PARAM[:mfname],
    plname::String = PARAM[:plname],)

  datefrequency = Dict("w"=>:week, "m"=>:month)[PARAM[:crspfrequencysuffix]]

  prefix = (returntype === measurefilename[1:length(returntype)]
    ) ? measurefilename : "$(returntype)_$measurefilename"

  #create a directory to hold anlsyis results
  resultid = "rid$(Dates.format(now(),"yymmdd-HHMM"))"

  #load the measure file
  local measure::DataFrame
  inputpath = "$analysispath\\$prefix\\inputs"

  #each results set gets its own output path- do the below only once
  if !ispath("$analysispath\\$prefix")
    mkdir("$analysispath\\$prefix")
    mkdir("$analysispath\\$prefix\\inputs")
    mv("$analysispath\\$measurefilename.csv",  "$inputpath\\$measurefilename.csv")
    mv("$analysispath\\$measurefilename.txt",  "$inputpath\\$measurefilename.txt")
    cp("$workingpath\\$(returntype)_$(hfrname).csv", "$inputpath\\$(returntype)_$(hfrname).csv")
    cp("$workingpath\\$(returntype)_$(mfname).csv", "$inputpath\\$(returntype)_$(mfname).csv")
    #NOTE: momentum info is disabled for now- need to re-run data if going to try it
    if isfile("$workingpath\\$(returntype)_$(plname).csv")
      @warn "Momentum file detected! Note you will need to disable the error barrier in order
        to actually run the momentum tests"
      cp("$workingpath\\$(returntype)_$(plname).csv", "$inputpath\\$(returntype)_$(plname).csv")
    end
  else
    isfile("$analysispath\\$measurefilename.csv") && throw("Duplicate measure file
      $measurefilename found in analysis folder outside of input/results folder $prefix.
      This should not happen- measurefiles should correspond 1:1 with the input/results folder,
      unless the input/results folder does not exist (at which point it will be created
      immediately).")
  end

  #load the input data
  measure = CSV.File("$inputpath\\$measurefilename.csv") |> DataFrame

  #create a directory to hold anlsyis results
  resultid = "rid$(Dates.format(now(),"yymmdd-HHMM"))"
  outputpath = "$analysispath\\$prefix\\$resultid"

  @assert !ispath("$analysispath\\$prefix\\$resultid")
  mkdir(outputpath)
  cp(PARAMETER_FILE, "$outputpath\\$(resultid)_Parameters.xlsx")

  #now pre-process the columns
  #sums accross assets
  if datefrequency === :week && PARAM[:analysisaggregateweeklycols]
    assetmat = measure[!, findcols(measure, :asset)] |> Matrix{MFloat64}
    oldcols = [findcols(measure, :growth); findcols(measure, :asset);
      propertynames(measure[!, r"^V_"])]
    aggassets = sum(assetmat, dims=2) |> vec
    measure[!, Symbol(:A_, weeklyaggregatename)] = aggassets
    measure[!, Symbol(:G_, weeklyaggregatename)] = [
      missing; aggassets[2:end] ./ aggassets[1:(end-1)]; ]

    #also need to some accross volume
    volumemat = measure[!, propertynames(measure[!, r"^V_"])] |> Matrix{MFloat64}
    aggvolume = sum(volumemat, dims=2) |> vec
    measure[!, Symbol(:V_, weeklyaggregatename)] = aggvolume
    select!(measure, Not(oldcols))
  end

  #pre-process the measure data
  measure = measure[measure.date .≥ (minimum(measure.date) + burnin), :]
  derivemeasurecols!(measure; datefrequency)

  #setup neweywest SEs and moving block bootstrap
  #below was previously ceil(max(nrow(measure),nrow(funds))^0.25) |> Int
  neweylags::Int = nrow(measure)^0.25 |> ceil |> Int
  nwΣ = neweywestΣfunc(neweylags)

  bsaggfunc(v::AbstractVector) = (;se=std(skipmissing(v)), p=nonparametricp(v))
  #bsaggfunc(v) = (;se=std(skipmissing(v)), p=nonparametricp(v))

  bsmeasurecolsregex::Regex = Regex(
    PARAM[:bsmeasurecolsregex] === nothing ? "" : PARAM[:bsmeasurecolsregex])
  bsbenchmarkcolsregex::Regex = Regex(
    PARAM[:bsbenchmarkcolsregex] === nothing ? "" : PARAM[:bsbenchmarkcolsregex])
  bssamples = PARAM[:bssamples]

  #note- the below was previously #ceil(max(nrow(measure),nrow(funds))^(1/3)) |> Int
  bswidth = nrow(measure)^(1/3) |> ceil |> Int

  #we can reuse some of the bootstrap aggregation code for the null simulations
  nullar1aggfunc(v::AbstractVector,statname,stattestvals) = (;
    se=std(skipmissing(v)),
    p=nonparametricnullp(v, stattestvals[statname]))
  nullar1aggfunc(M::AbstractMatrix,statname,stattestvals) = (;
    se=std(M, dims=1) |> vec,
    p=nonparametricnullp(M, stattestvals[statname]))
  nullar1simulations = bssamples

  #lightly package the io data for use as needed
  measureinfo = (; resultid, inputpath, outputpath, returntype,
    analysispath, datefrequency, prefix,
    neweylags, nwΣ,
    bsaggfunc, bsmeasurecolsregex, bsbenchmarkcolsregex, bssamples, bswidth,
    nullar1aggfunc, nullar1simulations)

  return (measure, measureinfo)
end

#helpful for grabbing the column name sans prefix
stripcolprefix(s) = replace(s, r"^l?[A-Z0-9]+_"=>"") |> Symbol
stripcolprefix(s::Symbol) = stripcolprefix(string(s))

#compute the strategy pressure measure- AUM×turnover/aggvol
function derivemeasurecols!(measure;
    Fassets::Vector{Symbol}=findcols(measure, :asset, includelp=false),
    datefrequency::Symbol,
    lpfilterperiods::Vector{<:Union{Day, Nothing}} = PARAM[:analysislpfilterperiods],
    Ftotalvolume::Symbol = Symbol(:desc_, PARAM[:analysistotalvolume]),
    lladditionalreturncolsidx = PARAM[:lladditionalreturncolsidx],
    lladditionalnormalizers = PARAM[:lladditionalnormalizers],
    llmarketcol = PARAM[:llmarketcol])

    #other helper methods
    simpledif(v::AbstractVector) = [missing; v[2:end] .- v[1:(end-1)]]
    _logdif(v) = [missing; log.(v[2:end]) .- log.(v[1:(end-1)])]
    logdif(v) = _logdif(abs.(v)) .* sign.(v)

    #generate log growth columns
    for gcol ∈ findcols(measure, :growth)
      measure[!, gcol] = [missing; measure[2:end, gcol]]
      measure[!, Symbol(:l,gcol)] = log.(abs.(measure[!, gcol])) .* sign.(measure[!, gcol])
      measure[!, Symbol(:g,string(gcol)[2:end])] = measure[!, gcol] .- 1

      rcol = Symbol(:R_, stripcolprefix(gcol))
      acol = Symbol(:A_, stripcolprefix(gcol))
      agcol = Symbol(:ag, string(gcol)[2:end])
      ngcol = Symbol(:ng, string(gcol)[2:end])
      cggcol = Symbol(:cgg, string(gcol)[2:end])
      measure[!, agcol] = measure[!, gcol] .- measure[!, rcol] .- 1.0
      measure[!, ngcol] = measure[!, gcol] .- measure[!, llmarketcol] .- 1.0
      measure[!, cggcol] = measure[!, rcol] |> deepcopy

      #active asset growth
      measure[!, Symbol(:aA, string(gcol)[2:end])] = [missing; measure[1:(end-1), acol] .* measure[2:end, agcol]]
      measure[!, Symbol(:nA, string(gcol)[2:end])] = [missing; measure[1:(end-1), acol] .* measure[2:end, ngcol]]
    end

    #Fassetcols = assetcols(measure)
    #make growth cols
    measure = lpfilter(measure,focalcols = [
      findcols(measure, :loggrowth); findcols(measure, :asset); :mctotal])
    colstodif = findcols(measure, :loggrowth)
    transform!(measure,  colstodif .=>
      (v -> [missing; v[2:end] .- v[1:(end-1)]])
      .=> (s->Symbol(:G, s)).(colstodif))

    Fvols = propertynames(measure[!, r"^V_"])
    Fpurchases = propertynames(measure[!, r"^Vs_"])
    Fpressures = (s->Symbol(:P_, stripcolprefix(s))).(Fassets)
    Fparticipations = (s->Symbol(:p_, stripcolprefix(s))).(Fassets)
    Flgpressures = (s->Symbol(:lG, s)).(Fpressures)
    totalvolume = measure[!, Ftotalvolume]
    totalcapitalization = measure[!, :mctotal]

    Fpurchaseparticipations = (s->Symbol(:pp_, stripcolprefix(s))).(Fassets)
    Fpurchasefractions = (s->Symbol(:VsA_, stripcolprefix(s))).(Fassets)
    Fpurchasecaps = (s->Symbol(:VsMC_, stripcolprefix(s))).(Fassets)

    #compute the presure characteristic
    transform!(measure, Fvols .=>
      ((vol)-> vol ./ totalvolume) .=> Fparticipations)
    transform!(measure, Fpurchases .=>
      ((pur)-> pur ./ totalvolume) .=> Fpurchaseparticipations)
    transform!(measure, zip(Fassets, Fpurchases) .|> collect .=>
      ((asset,pur)-> pur ./ asset) .=> Fpurchasefractions)
    transform!(measure,  Fpurchases .=>
      ((pur)-> pur ./ totalcapitalization) .=> Fpurchasecaps)
    transform!(measure, Fvols .=>
      ((vol)-> log.(vol)) .=> ((s->Symbol(:l,s)).(Fvols)))
    transform!(measure, zip(Fassets, Fvols) .|> collect .=>
      ((asset, vol)->asset .* vol ./ totalvolume) .=> Fpressures)
    transform!(measure, Fpressures .=> logdif .=> Flgpressures)
    lpfilter(measure, focalcols=Flgpressures)
    Flgpressures = [Flgpressures; lpcols.(Flgpressures)...;]

    #compute a version just of the volume-pressure
    Fvpressures = (s->Symbol(:VP_, stripcolprefix(s))).(Fassets)
    Flgvpressures = (s->Symbol(:lGVP_, stripcolprefix(s))).(Fassets)
    transform!(measure, Fvols .=> (v -> logdif(v ./ totalvolume)) .=> Flgvpressures)
    lpfilter(measure, focalcols=Flgvpressures)
    Flgvpressures = [Flgvpressures; lpcols.(Flgvpressures)...]

    #compute a measure of dollar return
    Frets = (s->Symbol(:R_,s)).(stripcolprefix.(Fassets))
    Fdolrets = (s->Symbol(:X,s)).(Frets)
    transform!(measure, zip(Fassets, Frets) .|> collect .=> (
      (a,r)-> a .* r) .=> Fdolrets)

    #somewhat different calculation methedology here for the non-logged version
    #in order to avoid various operations on ratios
    lpfilter(measure, focalcols=Fvols;) #[Fvols; Ftotalvolume;])
    Fvpressures = [Fvpressures; lpcols.(Fvpressures)...;]
    Fgvpressures = (s->Symbol(:G, s)).(Fvpressures)
    Fvolsforvpressures = [Fvols; lpcols.(Fvols)...;]
    Ftotalvolumes = [Ftotalvolume for i ∈ 1:length(Fvolsforvpressures)]
    transform!(measure, zip(Fvolsforvpressures, Ftotalvolumes) .|> collect .=>
      ((v,tv) -> simpledif(v) ./ tv) .=> Fgvpressures)

    #generate the appropriate column labels
    Fmcassetcols = [Fassets; lpcols.(Fassets)...;] #all asset cols
    Fmccols = findcols(measure, :mc)
    Fmccols = [Fmccols; (Fmccol->[Fmccol for i in 1:length(lpcols(Fmccol))]).(Fmccols)...;]
    Famccols = Fassets .|>  stripcolprefix .|> (s)->Symbol(:M_, s)
    Famccols = [Famccols; lpcols.(Famccols)...;]
    Fgamccols = Famccols .|> s-> Symbol(:G, s)

    #these will be used for return predictability
        #WARNING: quick and dirty test
    #measure.mctotal = [missing; measure.mctotal[1:(end-1)]]
    lpfilter(measure, focalcols=[:mctotal]) #need to smooth the mctotal
    Fmctotcols = [:mctotal for i ∈ 1:length(Fassets)]
    Fmctotcols = [Fmctotcols; (Fmctotcol->
      [Fmctotcol for i in 1:length(lpcols(Fmctotcol))]).(Fmctotcols)...;]
    Famctotcols = Fassets .|>  stripcolprefix .|> (s)->Symbol(:MT_, s)
    Famctotcols = [Famctotcols; lpcols.(Famctotcols)...;]
    Flamctotcols = Famctotcols .|> s-> Symbol(:l, s)
    Fgamctotcols = Famctotcols .|> s-> Symbol(:G, s)

    Faamcactiveassetcols = Fassets .|> (s)->Symbol(:a, s)
    Faamcactiveassetcols = [Faamcactiveassetcols; lpcols.(Faamcactiveassetcols)...;]
    Fagamctotcols = Famctotcols .|> s-> Symbol(:aG, s)

    Fnamcactiveassetcols = Fassets .|> (s)->Symbol(:n, s)
    Fnamcactiveassetcols = [Fnamcactiveassetcols; lpcols.(Fnamcactiveassetcols)...;]
    Fngamctotcols = Famctotcols .|> s-> Symbol(:nG, s)

    #NOTE we take the ratio of the average rather than the average of the ratio,
    #while also following the Welch 2020 RoC approach

    transform!(measure,
      zip(Fmcassetcols, Fmccols) .|> collect .=>
        ((asset, mc)->(asset ./ mc)) .=> Famccols,
      zip(Fmcassetcols, Fmccols) .|> collect .=>
        ((asset, mc)->(simpledif(asset) ./ mc)) .=> Fgamccols,

      zip(Fmcassetcols, Fmctotcols) .|> collect .=>
        ((asset, mctot)->(asset ./ mctot)) .=> Famctotcols,
      zip(Fmcassetcols, Fmctotcols) .|> collect .=>
        ((asset, mctot)->log.(abs.(asset ./ mctot)) .* sign.(asset)) .=> Flamctotcols,
      zip(Fmcassetcols, Fmctotcols) .|> collect .=>
        ((asset, mctot)->(simpledif(asset) ./ mctot)) .=> Fgamctotcols,

      zip(Faamcactiveassetcols, Fmctotcols) .|> collect .=>
        ((activeasset, mctot)->(activeasset ./ mctot)) .=> Fagamctotcols,
      zip(Fnamcactiveassetcols, Fmctotcols) .|> collect .=>
        ((netasset, mctot)->(netasset ./ mctot)) .=> Fngamctotcols,
        )

    #the below algorithm normalizes for an arbitrary denominator
    for (prefix, Fdenom) ∈ zip(keys(lladditionalnormalizers), values(lladditionalnormalizers))
      @assert Fdenom ∈ propertynames(measure)
      lpfilter(measure, focalcols=[Fdenom]) #need to smooth the denom

      #form the denominator names
      Fdenoms = [Fdenom for i ∈ 1:length(Fassets)]
      Fdenoms = [Fdenoms; (Fdenomcol->
        [Fdenomcol for i in 1:length(lpcols(Fdenomcol))]).(Fdenoms)...;]

      #form the names for the new fields
      Fadenoms = Fassets .|>  stripcolprefix .|> (s)->Symbol(prefix, s)
      Fadenoms = [Fadenoms; lpcols.(Fadenoms)...;]
      Fladenoms = Fadenoms .|> s-> Symbol(:l, s)
      Fgadenoms = Fadenoms .|> s-> Symbol(:G, s)
      Fagadenoms = Fadenoms .|> s-> Symbol(:aG, s)
      Faladenoms = Fadenoms .|> s-> Symbol(:al, s)

      #normalize as necesary
      transform!(measure,
        zip(Fmcassetcols, Fdenoms) .|> collect .=>
          ((asset, denom)->(asset ./ denom)) .=> Fadenoms,
        zip(Fmcassetcols, Fdenoms) .|> collect .=>
          ((asset, denom)->log.(abs.(asset ./ denom)) .* sign.(asset)) .=> Fladenoms,
        zip(Fmcassetcols, Fdenoms) .|> collect .=>
          ((asset, denom)->(simpledif(asset) ./ denom)) .=> Fgadenoms,
        zip(Faamcactiveassetcols, Fdenoms) .|> collect .=>
          ((activeasset, denom)->(activeasset ./ denom)) .=> Fagadenoms,)

    end

    #lgvpressures doesn't include lpcols, but Fgvpressures does
    allcols = [Flgpressures; Fgvpressures; Flgvpressures]
    doubledifcols = allcols .|> s->Symbol(:G, s)
    transform!(measure, allcols .=> simpledif .=> doubledifcols)

    measure |> CSV.write(
      "$(PARAM[:testpath])\\$(PARAM[:analysismeasurefilename])_derive.csv")

    return measure
    #Fpressures =
end

#helper functions for acquiring data from the measure file
function findcols(df, coltype::Symbol; includelp=true, args...)::Vector{Symbol}

  simplepatterns = Dict(
    :loggrowth => r"^lG_",
    :growth => r"^G_",
    :growthM1 => r"^g_",
    :activegrowth => r"^ag_",
    :netgrowth => r"^ng_",
    :capitalgainsgrowth => r"^cgg_",
    :activeassetalloc => r"^aA_",
    :z => r"^Z_",
    :volume => r"^V_",
    :lvolume => r"^lV_",
    :purchases => r"^Vs_",
    :purchasefraction => r"^VsA_",
    :purchasecapitalization => r"^VsMC_",
    :return => r"^R_",
    :ff => r"^F_",
    :logpressure => r"^lP_",
    :growthpressure => r"^GP_",
    :loggrowthpressure => r"^lGP_",
    :participation => r"^p_",
    :purchaseparticipation => r"^pp_",
    :sumloggrowthpressure => r"^lGSP_",
    :growthvolumepressure => r"^GVP_",
    :loggrowthvolumepressure => r"^lGVP_",
    :assetmarketcap => r"^M_",
    :assetmarketcapgrowth => r"^GM_",

    :assetmf => r"^AN_",
    :lassetmf => r"^lAN_",
    :assetmfgrowth => r"^GAN_",
    :activelassetmf => r"^alAN_",
    :activeassetmfgrowth => r"^aGAN_",

    :assetmfequity => r"^EN_",
    :lassetmfequity => r"^lEN_",
    :assetmfequitygrowth => r"^GEN_",
    :activeassetmfequitygrowth => r"^aGEN_",

    :assetmfequityonly => r"^EON_",
    :lassetmfequityonly => r"^lEON_",
    :assetmfequityonlygrowth => r"^GEON_",
    :activeassetmfequityonlygrowth => r"^aGEON_",

    :assetmfdomesticeo => r"^DEON_",
    :lassetmfdomesticeo => r"^lDEON_",
    :assetmfdomesticeogrowth => r"^GDEON_",
    :activeassetmfdomesticeogrowth => r"^aGDEON_",

    :assetmarketcaptotal => r"^MT_",
    :lassetmarketcaptotal => r"^lMT_",
    :assetmarketcaptotalgrowth => r"^GMT_",
    :activeassetmarketcaptotalgrowth => r"^aGMT_",
    :netassetmarketcaptotalgrowth => r"^nGMT_",
    )

  local foundcols
  if coltype ∈ keys(simplepatterns)
    foundcols = propertynames(df[!, simplepatterns[coltype]])
  else
    foundcols = findcols(df, Val{coltype}(); args...)
  end

  if !includelp #sometimes we may want to exclude the filtered cols
    foundcols = setdiff(foundcols, propertynames(df[!, r"_LP[0-9]+d$"]))
  end
  return foundcols
end

findcols(measure::AbstractDataFrame, ::Val{:asset}; includelp=true) = setdiff(
  propertynames(measure[!, r"^A_"]), includelp ? [] : propertynames(measure[!, r"_LP[0-9]+d$"]))
findcols(measure::AbstractDataFrame, ::Val{:mc};
  mctype=:gross) = propertynames(measure[!, Regex("^mc$(mctype)_")])

#generic fallback
findcols(measure::AbstractDataFrame, reg::Regex; includelp=true) = setdiff(
  propertynames(measure[!, reg]), includelp ? [] : propertynames(measure[!, r"_LP[0-9]+d$"]))
findcols(measure::AbstractDataFrame, reg::String; args...) = findcols(
  measure, Regex(reg), args...)

lpcols(F::Symbol;
  lpfilterperiods::Vector{<:Union{Day, Nothing}} = PARAM[:analysislpfilterperiods]) = (
    (p->Symbol(F, lpsuffix(p))).(setdiff(lpfilterperiods, [nothing])))

differencecol(s::Symbol) = Symbol(:G, s)


module graphops

using PyCall
include("utils.jl")

export
    # graph_transformer
    optimize,
    simplify,
    fuse_einsums,
    # graph_dedup
    dedup,
    # expr_generator
    rewrite_einsum_expr,
    # graph_inv_optimizer
    optimize_inverse,
    prune_inv_node,
    # optimal_tree
    generate_optimal_tree,
    split_einsum,
    generate_optimal_tree_w_constraint,
    # graph_als_optimizer
    generate_sequential_optimal_tree

const eingraph = PyNULL()
const gops = PyNULL()

function __init__()
    copy!(eingraph, pyimport_conda("autohoot.einsum_graph", "autohoot"))
    copy!(gops, pyimport_conda("autohoot.graph_ops", "autohoot"))
end

@func_loader eingraph.expr_generator rewrite_einsum_expr
@func_loader gops.graph_transformer optimize simplify fuse_einsums
@func_loader gops.graph_inv_optimizer optimize_inverse prune_inv_node
@func_loader gops.optimal_tree generate_optimal_tree split_einsum generate_optimal_tree_w_constraint
@func_loader gops.graph_als_optimizer generate_sequential_optimal_tree
@func_loader gops.graph_dedup dedup

end


module ChainRules

using Cassette
using LinearAlgebra
using Base.Broadcast: materialize, materialize!, broadcasted, Broadcasted, broadcastable

import NaNMath, SpecialFunctions, LinearAlgebra, LinearAlgebra.BLAS

export AbstractRule, Rule, frule, rrule

include("differentials.jl")
include("rules.jl")
include("rules/base.jl")
include("rules/array.jl")
include("rules/broadcast.jl")
include("rules/linalg/dense.jl")
include("rules/linalg/diagonal.jl")
include("rules/linalg/symmetric.jl")
include("rules/linalg/factorization.jl")
include("rules/blas.jl")
include("rules/nanmath.jl")
include("rules/specialfunctions.jl")

end # module


module MusicModels

using Flux
using CUDA

device = Flux.cpu

if CUDA.has_cuda()
    device = Flux.gpu
    println("GPU is enabled.")
    CUDA.allowscalar(false)
end

include("pretrain.jl")
include("datasets.jl")
include("MusicTransformer/MusicTransformer.jl")
include("PerformanceRNN/PerformanceRNN.jl")

using .PerformanceRNN
using .MusicTransformer

export generate

function call_registers()
    register_configs(pretrained_configs)
    register_datasets(dataset_configs)
end

@init call_registers()

end


using Caesar, IncrementalInference, RoME, HDF5, JLD
using CloudGraphs, Neo4j

# # Uncomment out for command line operation
# cloudGraph, addrdict = standardcloudgraphsetup(nparticles=true)
# session = addrdict["session"]

# interactive operation
# session = "SESSROX"
# Nparticles = 100
# include(joinpath(dirname(@__FILE__),"blandauthremote.jl"))

# get all the data

d = jldopen("test/data/VicPrkBasic.jld", "r") do file
  read(file, "dBasic")
end

f = jldopen("test/data/VicPrkBasic.jld", "r") do file
  read(file, "fBasic")
end
# MM = jldopen("test/data/VicPrkBasic.jld", "r") do file
#   read(file, "MMBasic")
# end
include("test/data/Vic120MM.jl")
MMr = MMwithRedirects(MM);

isamdict = jldopen("test/data/fgWithISAMrefERR01.jld", "r") do file
  read(file, "isamdict")
end
fgu = jldopen("test/data/fgWithISAMrefERR01.jld", "r") do file
  read(file, "fgu")
end

# do inference

T=100
fg = emptyFactorGraph();
fg.cg = cloudGraph

idx = appendFactorGraph!(fg, d, f, toT=T, lcmode=:mmodal, MM=MMr);

# tree = prepBatchTree!(fg, ordering=:qr,drawpdf=true);
# inferOverTree!(fg,tree, N=100);
#
# tree = wipeBuildNewTree!(fg)
# removeGenericMarginals!(conn)
# inferOverTree!(fg, tree)


module Lale

using PyCall

# load abstract super-types and utils
using AMLPipelineBase
using AMLPipelineBase.AbsTypes
export fit, fit!, transform, transform!, fit_transform!

using AMLPipelineBase
using AMLPipelineBase: AbsTypes, Utils, BaselineModels, Pipelines
using AMLPipelineBase: BaseFilters, FeatureSelectors, DecisionTreeLearners
using AMLPipelineBase: EnsembleMethods, CrossValidators
using AMLPipelineBase: NARemovers

export Machine, Learner, Transformer, Workflow, Computer
export holdout, kfold, score, infer_eltype, nested_dict_to_tuples,
       nested_dict_set!, nested_dict_merge, create_transformer,
       mergedict, getiris, getprofb,
       skipmean,skipmedian,skipstd,
       aggregatorclskipmissing,
       find_catnum_columns,
       train_test_split

export Baseline, Identity
export Imputer,OneHotEncoder,Wrapper
export PrunedTree,RandomForest,Adaboost
export VoteEnsemble, StackEnsemble, BestLearner
export FeatureSelector, CatFeatureSelector, NumFeatureSelector, CatNumDiscriminator
export crossvalidate
export NARemover
export @pipeline, @pipelinex
export Pipeline, ComboPipeline

import AMLPipelineBase.AbsTypes: fit!, transform!

# ------
module LaleAbsTypes
   using ..AbsTypes
   using DataFrames: DataFrame
   using PyCall
   using PyCall: PyIterator, inspect

   export LaleOperator, fit, transform, predict, lalepropertynames
   export pretty_print

   abstract type LaleOperator <: Learner end

   #fit(o::Machine, x::DataFrame, y::Vector=Vector()) = fit!(o,x,y)
   #transform(o::Machine, x::DataFrame) = transform!(o,x)
   predict(o::Machine, x::DataFrame) = transform!(o,x)
   pretty_print(o::LaleOperator,args...) = nothing

   lalepropertynames(o::PyObject) = ispynull(o) ? Symbol[] : map(x->Symbol(first(x)), PyIterator{PyObject}(pycall(inspect."getmembers", PyObject, o)))
end

using .LaleAbsTypes

include("lalelibop.jl")
using .LaleLibOps
export LalePipeOptimizer, lalepipeoptimizers
export LalePipe, visualize, pretty_print
export >>, +, |, |>, &

include("laleop.jl")
using .LaleOps
export LaleOp, skops, autogenops, lalelibops, lalesummary

export laleoperator
export fit!, transform!, fit, transform, predict

function laleoperator(name::String,type::String="sklearn"; args...)
   try
      obj = LaleOp(name,type;args...)
      return obj
   catch ArgumentError
      sk = keys(LaleOps.sk_dict)
      ag = keys(LaleOps.ag_dict)
      ll = keys(LaleOps.ll_dict)
      mb = keys(LaleOps.mb_dict)
      println("Please choose among these pipeline elements:")
      println()
      println("sklearn: ", [sk...])
      println()
      println("autogen: ", [ag...])
      println()
      println("lale: ",    [ll...])
      println()
      println("imblearn: ",[mb...])
      println()
      throw(ArgumentError("$name does not exist"))
   end
end

end


"""
    ParallelCondition

Supertype for all types imposing conditions of parallel.
"""
abstract type ParallelCondition end

@fieldequal ParallelCondition

kwarg(v::ParallelCondition) = :c => v

"""
    Unconditional <: ParallelCondition

Assume some notion of parallel holds without conditions.
"""
struct Unconditional <: ParallelCondition end

show(io::IO, ::Unconditional) =
    get(io, :compact, false) ? print(io, "U") : print(io, "Unconditional")

"""
    unconditional()

Alias for [`Unconditional()`](@ref).
"""
unconditional() = Unconditional()

"""
    CovariateConditional <: ParallelCondition

Supertype for all types assuming some notion of parallel holds
after conditioning on covariates.
"""
abstract type CovariateConditional <: ParallelCondition end

"""
    ParallelStrength

Supertype for all types specifying the strength of parallel.
"""
abstract type ParallelStrength end

kwarg(v::ParallelStrength) = :s => v

"""
    Exact <: ParallelStrength

Assume some notion of parallel holds exactly.
"""
struct Exact <: ParallelStrength end

show(io::IO, ::Exact) =
    get(io, :compact, false) ? print(io, "P") : print(io, "Parallel")

"""
    exact()

Alias for [`Exact()`](@ref).
"""
exact() = Exact()

"""
    Approximate <: ParallelStrength

Supertype for all types assuming some notion of parallel holds approximately.
"""
abstract type Approximate <: ParallelStrength end

"""
    AbstractParallel{C<:ParallelCondition, S<:ParallelStrength}

Supertype for all parallel types.
"""
abstract type AbstractParallel{C<:ParallelCondition, S<:ParallelStrength} end

@fieldequal AbstractParallel

"""
    UnspecifiedParallel{C,S} <: AbstractParallel{C,S}

A parallel trends assumption (PTA) without explicitly specified
relations across treatment groups.
See also [`unspecifiedpr`](@ref).

With this parallel type,
operations for complying with a PTA are suppressed.
This is useful, for example,
when the user-provided regressors and sample restrictions
need to be accepted without any PTA-specific alteration.

# Fields
- `c::C`: an instance of [`ParallelCondition`](@ref).
- `s::S`: an instance of [`ParallelStrength`](@ref).
"""
struct UnspecifiedParallel{C,S} <: AbstractParallel{C,S}
    c::C
    s::S
    function UnspecifiedParallel(c::ParallelCondition=Unconditional(),
            s::ParallelStrength=Exact())
        return new{typeof(c),typeof(s)}(c, s)
    end
end

"""
    unspecifiedpr(c::ParallelCondition=Unconditional(), s::ParallelStrength=Exact())

Construct an [`UnspecifiedParallel`](@ref) with fields set by the arguments.
This is an alias of the inner constructor of [`UnspecifiedParallel`](@ref).
"""
unspecifiedpr(c::ParallelCondition=Unconditional(), s::ParallelStrength=Exact()) =
    UnspecifiedParallel(c, s)

show(io::IO, pr::UnspecifiedParallel) =
    print(IOContext(io, :compact=>true), "Unspecified{", pr.c, ",", pr.s, "}")

function show(io::IO, ::MIME"text/plain", pr::UnspecifiedParallel)
    print(io, pr.s, " among unspecified treatment groups")
    pr.c isa Unconditional || print(io, ":\n  ", pr.c)
end

"""
    TrendParallel{C,S} <: AbstractParallel{C,S}

Supertype for all parallel types that
assume a parallel trends assumption holds over all the relevant time periods.
"""
abstract type TrendParallel{C,S} <: AbstractParallel{C,S} end

"""
    TrendOrUnspecifiedPR{C,S}

Union type of [`TrendParallel{C,S}`](@ref) and [`UnspecifiedParallel{C,S}`](@ref).
"""
const TrendOrUnspecifiedPR{C,S} = Union{TrendParallel{C,S}, UnspecifiedParallel{C,S}}

"""
    istreated(pr::TrendParallel, x)

Test whether `x` represents the treatment time
for a group of units that are not treated.
See also [`istreated!`](@ref).
"""
function istreated end

"""
    istreated!(out::AbstractVector{Bool}, pr::TrendParallel, x::AbstractArray)

For each element in `x`,
test whether it represents the treatment time
for a group of units that are not treated and save the result in `out`.
See also [`istreated`](@ref).
"""
function istreated! end

"""
    NeverTreatedParallel{C,S} <: TrendParallel{C,S}

Assume a parallel trends assumption holds between any group
that received the treatment during the sample periods
and a group that did not receive any treatment in any sample period.
See also [`nevertreated`](@ref).

# Fields
- `e::Tuple{Vararg{ValidTimeType}}`: group indices for units that did not receive any treatment.
- `c::C`: an instance of [`ParallelCondition`](@ref).
- `s::S`: an instance of [`ParallelStrength`](@ref).
"""
struct NeverTreatedParallel{C,S} <: TrendParallel{C,S}
    e::Tuple{Vararg{ValidTimeType}}
    c::C
    s::S
    function NeverTreatedParallel(e, c::ParallelCondition, s::ParallelStrength)
        e = applicable(iterate, e) ? (unique!(sort!([e...]))...,) : (e,)
        isempty(e) && error("field `e` cannot be empty")
        return new{typeof(c),typeof(s)}(e, c, s)
    end
end

istreated(pr::NeverTreatedParallel, x) = !(x in pr.e)

function istreated!(out::AbstractVector{Bool}, pr::NeverTreatedParallel,
        x::AbstractArray{<:Union{ValidTimeType, Missing}})
    e = Set(pr.e)
    out .= .!(x .∈ Ref(e))
end

function istreated!(out::AbstractVector{Bool}, pr::NeverTreatedParallel,
        x::ScaledArray{<:Union{ValidTimeType, Missing}})
    refs = refarray(x)
    invpool = invrefpool(x)
    e = Set(invpool[c] for c in pr.e if haskey(invpool, c))
    out .= .!(refs .∈ Ref(e))
end

show(io::IO, pr::NeverTreatedParallel) =
    print(IOContext(io, :compact=>true), "NeverTreated{", pr.c, ",", pr.s, "}",
        length(pr.e)==1 ? string("(", pr.e[1], ")") : pr.e)

function show(io::IO, ::MIME"text/plain", pr::NeverTreatedParallel)
    println(io, pr.s, " trends with any never-treated group:")
    print(io, "  Never-treated groups: ", join(string.(pr.e), ", "))
    pr.c isa Unconditional || print(io, "\n  ", pr.c)
end

"""
    nevertreated(e, c::ParallelCondition, s::ParallelStrength)
    nevertreated(e; c=Unconditional(), s=Exact())

Construct a [`NeverTreatedParallel`](@ref) with fields set by the arguments.
By default, `c` is set as [`Unconditional()`](@ref)
and `s` is set as [`Exact()`](@ref).
When working with `@formula`,
a wrapper method of `nevertreated` calls this method.

# Examples
```jldoctest; setup = :(using DiffinDiffsBase)
julia> nevertreated(-1)
Parallel trends with any never-treated group:
  Never-treated groups: -1

julia> typeof(nevertreated(-1))
NeverTreatedParallel{Unconditional,Exact}

julia> nevertreated([-1, 0])
Parallel trends with any never-treated group:
  Never-treated groups: -1, 0

julia> nevertreated([-1, 0]) == nevertreated(-1:0) == nevertreated(Set([-1, 0]))
true
```
"""
nevertreated(e, c::ParallelCondition, s::ParallelStrength) =
    NeverTreatedParallel(e, c, s)
nevertreated(e; c::ParallelCondition=Unconditional(), s::ParallelStrength=Exact()) =
    NeverTreatedParallel(e, c, s)

"""
    nevertreated(ts::AbstractTerm...)

A wrapper method of `nevertreated` for working with `@formula`.
"""
@unpack nevertreated

"""
    NotYetTreatedParallel{C,S} <: TrendParallel{C,S}

Assume a parallel trends assumption holds between any group
that received the treatment relatively early
and any group that received the treatment relatively late (or never receved).
See also [`notyettreated`](@ref).

# Fields
- `e::Tuple{Vararg{ValidTimeType}}`: group indices for units that received the treatment relatively late.
- `ecut::Tuple{Vararg{ValidTimeType}}`: user-specified period(s) when units in a group in `e` started to receive treatment or show anticipation effects.
- `c::C`: an instance of [`ParallelCondition`](@ref).
- `s::S`: an instance of [`ParallelStrength`](@ref).

!!! note
    `ecut` could be different from `minimum(e)` if
    - never-treated groups are included and use indices with smaller values;
    - the sample has a rotating panel structure with periods overlapping with some others.
"""
struct NotYetTreatedParallel{C,S} <: TrendParallel{C,S}
    e::Tuple{Vararg{ValidTimeType}}
    ecut::Tuple{Vararg{ValidTimeType}}
    c::C
    s::S
    function NotYetTreatedParallel(e, ecut, c::ParallelCondition, s::ParallelStrength)
        e = applicable(iterate, e) ? (unique!(sort!([e...]))...,) : (e,)
        isempty(e) && throw(ArgumentError("field e cannot be empty"))
        ecut = applicable(iterate, ecut) ? (unique!(sort!([ecut...]))...,) : (ecut,)
        isempty(ecut) && throw(ArgumentError("field ecut cannot be empty"))
        eltype(e) == eltype(ecut) ||
            throw(ArgumentError("e and ecut must have the same element type"))
        return new{typeof(c), typeof(s)}(e, ecut, c, s)
    end
end

istreated(pr::NotYetTreatedParallel, x) = !(x in pr.e)

function istreated!(out::AbstractVector{Bool}, pr::NotYetTreatedParallel,
        x::AbstractArray{<:Union{ValidTimeType, Missing}})
    e = Set(pr.e)
    out .= .!(x .∈ Ref(e))
end

function istreated!(out::AbstractVector{Bool}, pr::NotYetTreatedParallel,
        x::ScaledArray{<:Union{ValidTimeType, Missing}})
    refs = refarray(x)
    invpool = invrefpool(x)
    e = Set(invpool[c] for c in pr.e if haskey(invpool, c))
    out .= .!(refs .∈ Ref(e))
end

show(io::IO, pr::NotYetTreatedParallel) =
    print(IOContext(io, :compact=>true), "NotYetTreated{", pr.c, ",", pr.s, "}",
        length(pr.e)==1 ? string("(", pr.e[1], ")") : pr.e)

function show(io::IO, ::MIME"text/plain", pr::NotYetTreatedParallel)
    println(io, pr.s, " trends with any not-yet-treated group:")
    println(io, "  Not-yet-treated groups: ", join(string.(pr.e), ", "))
    print(io, "  Treated since: ", join(string.(pr.ecut), ", "))
    pr.c isa Unconditional || print(io, "\n  ", pr.c)
end

"""
    notyettreated(e, ecut, c::ParallelCondition, s::ParallelStrength)
    notyettreated(e, ecut=e; c=Unconditional(), s=Exact())

Construct a [`NotYetTreatedParallel`](@ref) with
fields set by the arguments.
By default, `c` is set as [`Unconditional()`](@ref)
and `s` is set as [`Exact()`](@ref).
When working with `@formula`,
a wrapper method of `notyettreated` calls this method.

# Examples
```jldoctest; setup = :(using DiffinDiffsBase)
julia> notyettreated(5)
Parallel trends with any not-yet-treated group:
  Not-yet-treated groups: 5
  Treated since: 5

julia> typeof(notyettreated(5))
NotYetTreatedParallel{Unconditional,Exact}

julia> notyettreated([-1, 5, 6], 5)
Parallel trends with any not-yet-treated group:
  Not-yet-treated groups: -1, 5, 6
  Treated since: 5

julia> notyettreated([4, 5, 6], [4, 5, 6])
Parallel trends with any not-yet-treated group:
  Not-yet-treated groups: 4, 5, 6
  Treated since: 4, 5, 6
```
"""
notyettreated(e, ecut, c::ParallelCondition, s::ParallelStrength) =
    NotYetTreatedParallel(e, ecut, c, s)
notyettreated(e, ecut=e; c::ParallelCondition=Unconditional(), s::ParallelStrength=Exact()) =
        NotYetTreatedParallel(e, ecut, c, s)

"""
    notyettreated(ts::AbstractTerm...)

A wrapper method of `notyettreated` for working with `@formula`.
"""
@unpack notyettreated

termvars(c::ParallelCondition) =
    error("StatsModels.termvars is not defined for $(typeof(c))")
termvars(::Unconditional) = Symbol[]
termvars(s::ParallelStrength) =
    error("StatsModels.termvars is not defined for $(typeof(s))")
termvars(::Exact) = Symbol[]
termvars(pr::AbstractParallel) =
    error("StatsModels.termvars is not defined for $(typeof(pr))")
termvars(pr::UnspecifiedParallel) = union(termvars(pr.c), termvars(pr.s))
termvars(pr::NeverTreatedParallel) = union(termvars(pr.c), termvars(pr.s))
termvars(pr::NotYetTreatedParallel) = union(termvars(pr.c), termvars(pr.s))


# g=5.0
# ENV["JULIA_CUDA_MEMORY_LIMIT"] = convert(Int, round(g * 1024 * 1024 * 1024))

include("exp.jl")

function main_ngraph()
    "Test how many graphs are needed to learn the model."
    # this requires the data generating spec to be different

    # TODO these are a lot of configs
    for d in [10, 20],
        ng in [100, 200, 500, 1000, 2000, 5000, 10000]

        specs = []
        for gtype in [:ER, :SF],
            k in [1]
            push!(specs, DataSpec(d=d, k=k, gtype=gtype,
                    noise=:Gaussian, mat=:CH3,
                    # when ng=100, total graph is 200, testing is 200*0.2=40, not enough for a batch of 64
                    ng=ng, bsize=32,
                    N=1))
        end
        specs = Array{DataSpec}(specs)

        @info "training .." d ng
        expID = exp_train(specs, eq2_model_fn,
                          # TODO I'll need to increase the training steps here
                          # CAUTION feed in the gtype in the model prefix
                          prefix="ngraph-ng=$ng-d=$d", train_steps=1e5,
                          test_throttle = 10,
                          merge=true)
    end
end

function main_ersf124()
    # test on different number of nodes
    for d in [10,20],
        (prefix, model_fn,nsteps) in [
            ("EQ2", eq2_model_fn, 1e5),
        ]

        specs = []
        for gtype in [:ER, :SF, :ER2, :SF2, :ER4, :SF4],
            k in [1]
            push!(specs, DataSpec(d=d, k=k, gtype=gtype,
                    noise=:Gaussian, mat=:CH3))
        end
        specs = Array{DataSpec}(specs)

        @info "training .." prefix d
        expID = exp_train(specs, model_fn,
                          # TODO I'll need to increase the training steps here
                          # CAUTION feed in the gtype in the model prefix
                          prefix="$prefix-ERSF124-d=$d", train_steps=nsteps,
                          test_throttle = 10,
                          merge=true)
    end
end

function main_data()
    # generate data for baseline usage
    for d in [10, 20, 50, 100, 150, 200, 300, 400]
        @info "Generating for" d ".."
        specs = []
        # this is combined specs. What was the filename?
        # will generate separate data files, and combine when loading
        for gtype in [:ER, :SF],
            k in [1]
            push!(specs, DataSpec(d=d, k=k, gtype=gtype,
                    noise=:Gaussian, mat=:CH3))
        end
        specs = Array{DataSpec}(specs)
        ds, test_ds = spec2ds(specs)
    end
end

function main_large()
    # large d, large ng and N, longer training time
    for (d, ng, N, bsize) in [
            (200, 1000, 1, 2),
            (150, 3000, 3, 4),
            # DEBUG test ng=10000
            (100, 10000, 3, 8),
            # DEBUG testing whether d=200 is trainable
            # d=200 is very slow
            # d=400
            (400, 1000, 1, 1)
        ],
        (prefix, model_fn,nsteps) in [
            # previous try didn't actually set ng, N, bsize, and it was marked
            # with old- prefix
            ("EQ2", eq2_model_fn, 5e4)
        ]

        specs = []
        for gtype in [:ER, :SF],
            k in [1]
            push!(specs, DataSpec(d=d, k=k, gtype=gtype,
                    noise=:Gaussian, mat=:CH3,
                    ng=ng, N=N, bsize=bsize))
        end
        specs = Array{DataSpec}(specs)

        # print more frequently for CNN and FC to get more data to print
        test_throttle = if prefix == "EQ2" 10 else 1 end

        @info "training .." prefix d
        expID = exp_train(specs, model_fn,
                          # TODO I'll need to increase the training steps here
                          # CAUTION feed in the gtype in the model prefix
                          prefix="$prefix-ERSF-k1-d=$d-ng=$ng-N=$N", train_steps=nsteps,
                          test_throttle = test_throttle,
                          merge=true)
    end
end

# Train on ER and test on SF and vice versa
function main_ersf()
    for d in [10,20,50, 100],
        (prefix, model_fn,nsteps) in [
            ("EQ2", eq2_model_fn, 3e4),
        ],
        gtype in [:ER, :SF]

        specs = DataSpec(d=d, k=1, gtype=gtype,
                        noise=:Gaussian, mat=:CH3)

        # print more frequently for CNN and FC to get more data to print
        test_throttle = if prefix == "EQ2" 10 else 1 end

        @info "training $prefix-$gtype-k1-d=$d .."
        expID = exp_train(specs, model_fn,
                          # TODO I'll need to increase the training steps here
                          # CAUTION feed in the gtype in the model prefix
                          prefix="$prefix-$gtype-k1-d=$d", train_steps=nsteps,
                          test_throttle = test_throttle,
                          merge=true)
        # Do the test here
        begin
            for gtype in [:ER, :SF]
                test_spec = DataSpec(d=d, k=1, gtype=gtype,
                                    noise=:Gaussian, mat=:CH3)
                exp_test(expID, test_spec, "TEST-gtype=$gtype")
            end
        end
    end
end

function main_ch3()
    for d in [10,20,50,100],
        (prefix, model_fn,nsteps) in [
            ("EQ2", eq2_model_fn, 5e4),
            ("FC", ()->fc_model(d=d, ch=2, z=1024, nlayer=6), 3e4),
            # ("FCreg", ()->fc_model(d=d, ch=2, z=1024, nlayer=6, reg=true), 3e4),
            ("CNN", ()->cnn_model(2), 3e4),
            # ("CNN2", ()->cnn_model(2, 128, (5,5), (2,2)), 3e4)
        ]

        specs = []
        for gtype in [:ER, :SF],
            k in [1]
            push!(specs, DataSpec(d=d, k=k, gtype=gtype,
                    noise=:Gaussian, mat=:CH3))
        end
        specs = Array{DataSpec}(specs)

        # print more frequently for CNN and FC to get more data to print
        test_throttle = if prefix == "EQ2" 10 else 1 end

        @info "training .." prefix d
        expID = exp_train(specs, model_fn,
                          # TODO I'll need to increase the training steps here
                          # CAUTION feed in the gtype in the model prefix
                          prefix="$prefix-ERSF-k1-d=$d", train_steps=nsteps,
                          test_throttle = test_throttle,
                          merge=true)
    end
end

function main_mat()
    # test the different matrix, including: CH3, COV, COR, ensemble on different k
    for d in [20,50],
        (mat, model_fn) in [
            (:CH3, eq2_model_fn),
            (:COV, eq_model_fn),
            (:COR, eq_model_fn)]

        specs = []
        for gtype in [:ER, :SF],
            k in [1,2,4,8]
            push!(specs, DataSpec(d=d, k=k, gtype=gtype,
                    noise=:Gaussian, mat=mat))
        end
        specs = Array{DataSpec}(specs)

        @info "training .." d mat
        expID = exp_train(specs, model_fn,
                          # TODO I'll need to increase the training steps here
                          # CAUTION feed in the gtype in the model prefix
                          prefix="EQ-d=$d-mat=$mat", train_steps=3e4,
                          merge=true)
    end
end

function main_ensemble_d()
    # ensemble on different d
    specs = []
    for gtype in [:ER, :SF],
        d in [20,30,40] # TODO use more aggressive d, e.g. 10, 30, 50

        push!(specs, DataSpec(d=d, k=1, gtype=gtype,
                noise=:Gaussian, mat=:CH3))
    end
    specs = Array{DataSpec}(specs)

    @info "training .."
    expID = exp_train(specs, eq2_model_fn,
                      # TODO I'll need to increase the training steps here
                      # CAUTION feed in the gtype in the model prefix
                      prefix="EQ2-CH3-d=[20,30,40]", train_steps=3e4,
                      merge=false)
end

function main_ensemble_d_cnn()
    # ensemble on different d
    specs = []
    for gtype in [:ER, :SF],
        d in [10,15,20] # TODO use more aggressive d, e.g. 10, 30, 50

        push!(specs, DataSpec(d=d, k=1, gtype=gtype,
                noise=:Gaussian, mat=:CH3))
    end
    specs = Array{DataSpec}(specs)

    @info "training .."
    expID = exp_train(specs, ()->cnn_model(2),
                      # TODO I'll need to increase the training steps here
                      # CAUTION feed in the gtype in the model prefix
                      prefix="CNN-CH3-d=[10,15,20]", train_steps=3e4,
                      test_throttle=1,
                      merge=false)
end


module test_4

using AdventOfCode2018.Day4
using Test
using Dates

test_input = """[1518-11-01 00:00] Guard #10 begins shift
[1518-11-01 00:05] falls asleep
[1518-11-01 00:25] wakes up
[1518-11-01 00:30] falls asleep
[1518-11-01 00:55] wakes up
[1518-11-01 23:58] Guard #99 begins shift
[1518-11-02 00:40] falls asleep
[1518-11-02 00:50] wakes up
[1518-11-03 00:05] Guard #10 begins shift
[1518-11-03 00:24] falls asleep
[1518-11-03 00:29] wakes up
[1518-11-04 00:02] Guard #99 begins shift
[1518-11-04 00:36] falls asleep
[1518-11-04 00:46] wakes up
[1518-11-05 00:03] Guard #99 begins shift
[1518-11-05 00:45] falls asleep
[1518-11-05 00:55] wakes up
"""

parsed_input = [
	(DateTime(1518, 11, 01, 00, 00), 10, false, false),
	(DateTime(1518, 11, 01, 00, 05), -1, true, false),
	(DateTime(1518, 11, 01, 00, 25), -1, false, true),
	(DateTime(1518, 11, 01, 00, 30), -1, true, false),
	(DateTime(1518, 11, 01, 00, 55), -1, false, true),
	(DateTime(1518, 11, 01, 23, 58), 99, false, false),
	(DateTime(1518, 11, 02, 00, 40), -1, true, false),
	(DateTime(1518, 11, 02, 00, 50), -1, false, true),
	(DateTime(1518, 11, 03, 00, 05), 10, false, false),
	(DateTime(1518, 11, 03, 00, 24), -1, true, false),
	(DateTime(1518, 11, 03, 00, 29), -1, false, true),
	(DateTime(1518, 11, 04, 00, 02), 99, false, false),
	(DateTime(1518, 11, 04, 00, 36), -1, true, false),
	(DateTime(1518, 11, 04, 00, 46), -1, false, true),
	(DateTime(1518, 11, 05, 00, 03), 99, false, false),
	(DateTime(1518, 11, 05, 00, 45), -1, true, false),
	(DateTime(1518, 11, 05, 00, 55), -1, false, true),
]

@testset "day 3, part 1" begin

	parsed = parse_input(test_input)
	@test parsed == parsed_input

	map = build_sleep_map(parsed)
	@test part_1(map) == 240
	@test part_2(map) == 4455

end

end


# low-level interface, meaning works with vectors and matrices, not data frames and estimation objects

"Structure to store test results"
struct BootTestResult{T}
  stat::T; stattype::String
  p::T; padj::T
  reps::Int64; repsfeas::Int64
  nbootclust::Int64
  dof::Int64; dof_r::T
  plot::Union{Nothing, NamedTuple{(:X, :p), Tuple{Tuple{Vararg{Vector{T}, N} where N},Vector{T}}}}
  peak::Union{Nothing, NamedTuple{(:X, :p), Tuple{Vector{T}, T}}}
  ci::Union{Nothing, Matrix{T}}
  dist::Matrix{T}
  b::Vector{T}
  V::Matrix{T}
  auxweights::Union{Nothing,Matrix{T}}
  # M::StrBootTest
end

"Return test statistic"
teststat(o::BootTestResult) = o.stat

"Return numerator of test statistic"
statnumer(o::BootTestResult) = o.b

"Return denominator of test statistic"
statvar(o::BootTestResult) = o.V

"""Return type of test statistic: "t", "z", "F", or "χ²" """
stattype(o::BootTestResult) = o.stattype

"Return p value"
p(o::BootTestResult) = o.p

"Return p value after multiple-hypothesis adjustment, if any"
padj(o::BootTestResult) = o.padj

"Return requested number of replications"
reps(o::BootTestResult) = o.reps

"Return actual number of replications, subject to enumeration of Rademacher draws"
repsfeas(o::BootTestResult) = o.repsfeas

"Return number of bootstrapping clusters in test"
nbootclust(o::BootTestResult) = o.nbootclust

"Return degrees of freedom of test"
dof(o::BootTestResult) = o.dof

"Return residual degrees of freedom of test"
dof_r(o::BootTestResult) = o.dof_r

"""
Return data for confidence plot of test.
Return value is a 2-tuple with named entries `X` and `p` holding
the confidence sampling locations and p values respectively. `X` is in turn
a 1- or 2-tuple of vectors of sampling coordinates for each
dimension of the tested hypothesis.
"""
plotpoints(o::BootTestResult) = o.plot

"Return parameter value with peak p value in test"
peak(o::BootTestResult) = o.peak

"Return confidence interval matrix from test, one row per disjoint piece"
ci(o::BootTestResult) = o.ci

"Return bootstrap distribution of statistic or statistic numerator in bootstrap test"
dist(o::BootTestResult) = o.dist

"Return auxilliary weight matrix for wild bootstrap"
auxweights(o::BootTestResult) = o.auxweights

strint(x) = iszero(mod(x,1)) ? "$(Int64(x))" : "$x"
function Base.show(io::IO, o::BootTestResult{T}) where T
	s = stattype(o) * ( iszero(dof_r(o)) ? isone(dof(o)) ? " " : "(" * strint(dof(o)) * ")" :                                                     # z, χ²
	                                       isone(dof(o)) ? "(" * strint(dof_r(o)) * ")" : "(" * strint(dof(o)) * ", " * strint(dof_r(o)) * ")" )  # t, F
	Printf.@printf(io, "%s = %6.4f\n", s, teststat(o))
	Printf.@printf(io, "p%s = %6.4f\n", repeat(" ", length(s)-1), p(o))
	isdefined(o, :ci) && !isnothing(o.ci) && length(o.ci)>0 && print(io, "CI" * repeat(" ", length(s)-2) * " = $(round.(ci(o); sigdigits=4))\n")
end

# single entry point with arguments already converted to standardized types, to allow a smaller set of precompile() calls(?)
function __wildboottest(
	R::Matrix{T},
	r::Vector{T};
	resp::VecOrMat{T},
	predexog::Matrix{T},
	predendog::Matrix{T},
	inst::Matrix{T},
	R1::Matrix{T},
	r1::Vector{T},
	clustid::Matrix{Int64},
	nbootclustvar::Int64,
	nerrclustvar::Int64,
	issorted::Bool,
	hetrobust::Bool,
	nfe::Int64,
	feid::VecOrMat{Int64},
	fedfadj::Int64,
	obswt::VecOrMat{T},
	fweights::Bool,
	maxmatsize::Float16,
	ptype::Symbol,
	bootstrapc::Bool,
	liml::Bool,
	fuller::T,
	kappa::T,
	arubin::Bool,
	small::Bool,
	clusteradj::Bool,
	clustermin::Bool,
	scorebs::Bool,
	reps::Int64,
	imposenull::Bool,
	auxwttype::Symbol,
	rng::AbstractRNG,
	level::T,
	rtol::T,
	madjtype::Symbol,
	nH0::Int16,
	ml::Bool,
	scores::Matrix{T},
	beta::Vector{T},
	A::Symmetric{T,Matrix{T}},
	gridmin::VecOrMat{T},
	gridmax::VecOrMat{T},
	gridpoints::VecOrMat{T},
	diststat::Symbol,
	getci::Bool,
	getplot::Bool,
	getauxweights::Bool) where T

	M = StrBootTest{T}(R, r, R1, r1, resp, predexog, predendog, inst, obswt, fweights, liml, fuller, kappa, arubin,
	                   reps, auxwttype, rng, maxmatsize, ptype, imposenull, scorebs, !bootstrapc, clustid, nbootclustvar, nerrclustvar, issorted, hetrobust, small, clusteradj, clustermin,
	                   nfe, feid, fedfadj, level, rtol, madjtype, nH0, ml, beta, A, scores, getplot,
	                   gridmin, gridmax, gridpoints)

	if getplot || (level<1 && getci)
		plot!(M)
		plot = getplot & isdefined(M, :plotX) ? (X=Tuple(M.plotX), p=M.plotY) : nothing
		peak = M.peak
		ci = level<1 & getci ? M.ci : nothing
	else
		ci = plot = peak = nothing
	end

	padj = getp(M)  # trigger main (re)computation

	BootTestResult{T}(getstat(M),
	                  isone(nrows(R)) ? (small ? "t" : "z") : (small ? "F" : "χ²"),
	                  M.p, padj, M.B, M.BFeas, M.N✻, M.dof, M.dof_r, plot, peak, ci,
	                  getdist(M, diststat),
	                  getb(M), getV(M),
	                  getauxweights && reps>0 ? getv(M) : nothing #=, M=#)
end

vecconvert(T::DataType, X) = Vector(isa(X, AbstractArray) ? vec(    eltype(X)==T ? X : T.(X)) : X)
matconvert(T::DataType, X) = Matrix(isa(X, AbstractArray) ? reshape(eltype(X)==T ? X : T.(X), size(X,1), size(X,2)) : X)

function _wildboottest(T::DataType,
					  R::AbstractVecOrMat,
						r::AbstractVecOrMat;
					  resp::AbstractVecOrMat{<:Real},
					  predexog::AbstractVecOrMat{<:Real}=zeros(T,0,0),
					  predendog::AbstractVecOrMat{<:Real}=zeros(T,0,0),
					  inst::AbstractVecOrMat{<:Real}=zeros(T,0,0),
					  R1::AbstractVecOrMat=zeros(T,0,0),
						r1::AbstractVecOrMat=zeros(T,0),
					  hetrobust::Bool=true,
					  clustid::AbstractVecOrMat{<:Integer}=zeros(Int,0,0),  # bootstrap-only clust vars, then boot&err clust vars, then err-only clust vars
					  nbootclustvar::Integer=ncols(clustid),
					  nerrclustvar::Integer=nbootclustvar,
						issorted::Bool=false,
						nfe::Integer=0,
					  feid::AbstractVecOrMat{<:Integer}=Int8[],
					  fedfadj::Integer=length(feid)>0 ? -1 : 0,
					  obswt::AbstractVecOrMat{<:Real}=T[],
					  fweights::Bool=false,
					  maxmatsize::Number=0,
					  ptype::Symbol=:symmetric,
					  bootstrapc::Bool=false,
					  liml::Bool=false,
					  fuller::Number=0,
					  kappa::Number=NaN,
					  arubin::Bool=false,
					  small::Bool=true,
						clusteradj::Bool=small,
						clustermin::Bool=false,
					  scorebs::Bool=false,
					  reps::Integer=999,
					  imposenull::Bool=true,
					  auxwttype::Symbol=:rademacher,
					  rng::AbstractRNG=MersenneTwister(),
					  level::Number=.95,
					  rtol::Number=1e-6,
					  madjtype::Symbol=:none,
					  nH0::Integer=1,
					  ml::Bool=false,
					  scores::AbstractVecOrMat=Matrix{Float32}(undef,0,0),
					  beta::AbstractVecOrMat=T[],
					  A::AbstractMatrix=zeros(T,0,0),
					  gridmin::Union{VecOrMat{S},VecOrMat{Union{S,Missing}}} where S<:Number = T[],
					  gridmax::Union{VecOrMat{S},VecOrMat{Union{S,Missing}}} where S<:Number = T[],
					  gridpoints::Union{VecOrMat{S},VecOrMat{Union{S,Missing}}} where S<:Number = Int64[],
					  diststat::Symbol=:none,
					  getci::Bool=true,
					  getplot::Bool=getci,
					  getauxweights::Bool=false)

	nrows(R)>2 && (getplot = getci = false)

	@assert any(auxwttype .== (:rademacher, :mammen, :webb, :gamma, :normal)) "auxwttype shoud be :rademacher, :mammen, :webb, :gamma, or :normal"
	@assert any(ptype .==(:symmetric, :equaltail, :lower, :upper)) "ptype should be :symmetric, :equaltail, :lower, or :upper"
	@assert any(madjtype .== (:none, :bonferroni, :sidak)) "madjtype should be :none, :bonferroni, or :sidak"
	@assert any(diststat .== (:none, :t, :numer))
	@assert ml || ncols(resp)==1 "resp should have one column"
  @assert (length(predexog)==0 || nrows(predexog)==nrows(resp)) &&
	        (length(predendog)==0 || nrows(predendog)==nrows(resp)) &&
					(length(inst)==0 || nrows(inst)==nrows(resp)) "All data vectors/matrices must have same height"
	@assert ncols(inst) >= ncols(predendog) "Model has fewer instruments than instrumented variables"
  @assert length(feid)==0 || nrows(feid)==nrows(resp) "feid vector must have same height as data matrices"
	@assert ncols(feid)≤1 "feid should have one column"
  @assert length(clustid)==0 || nrows(clustid)==nrows(resp) "clustid must have same height as data matrices"
  @assert nrows(obswt)==0 || nrows(obswt)==nrows(resp) "obswt must have same height as data matrices"
	@assert ncols(obswt)≤1 "obswt must have one column"
  @assert nrows(R)==nrows(r) "R and r must have same height"
  @assert (ncols(R) == (ml ? nrows(beta) : ncols(predexog)+ncols(predendog)) && isone(ncols(r))) "Wrong number of columns in null specification"
  @assert nrows(R1)==nrows(r1) "R₁ and r₁ must have same height"
  @assert length(R1)==0 || ncols(R1)==ncols(predexog)+ncols(predendog) "Wrong number of columns in model constraint specification"
	@assert ncols(r)==1 "r should have one column"
	@assert length(R1)==0 || ncols(r1)==1 "r1 should have one column"
  @assert nbootclustvar ≤ ncols(clustid) "nbootclustvar > width of clustid"
  @assert nerrclustvar ≤ ncols(clustid) "nerrclustvar > width of clustid"
  @assert reps ≥ 0 "reps < 0"
  @assert level ≥ 0. && level≤1. "level must be in the range [0,1]"
  @assert rtol > 0. "rtol ≤ 0"
  @assert nH0 > 0 "nH0 ≤ 0"
	@assert !liml || (ncols(predendog)>0 && ncols(inst)>0) "For liml, non-empty predendog and inst arguments are needed"
	@assert fuller==0 || (ncols(predendog)>0 && ncols(inst)>0) "For Fuller liml, non-empty predendog and inst arguments are needed"
	@assert iszero(ncols(predendog)) || ncols(inst)>0 "predendog provided without inst"
	@assert !arubin || ncols(predendog)>0 "Anderson-Rubin test requested but predendog not provided"

	if getplot || getci
		@assert iszero(length(gridmin   )) || length(gridmin   )==nrows(R) "Length of gridmin doesn't match number of hypotheses being jointly tested"
		@assert iszero(length(gridmax   )) || length(gridmax   )==nrows(R) "Length of gridmax doesn't match number of hypotheses being jointly tested"
		@assert iszero(length(gridpoints)) || length(gridpoints)==nrows(R) "Length of gridpoints doesn't match number of hypotheses being jointly tested"
		@assert iszero(length(gridmin   )) || ncols(gridmin)   ==1 "gridmin should have one column"
		@assert iszero(length(gridmax   )) || ncols(gridmax)   ==1 "gridmax should have one column"
		@assert iszero(length(gridpoints)) || ncols(gridpoints)==1 "gridpoints should have one column"
	end

	_gridmin = Vector{T}(undef, length(gridmin))
	_gridmax = Vector{T}(undef, length(gridmax))
	_gridpoints = Vector{T}(undef, length(gridpoints))
	for i ∈ 1:length(gridmin)  # cumbersome loops because map() and list comprehensions mess up type inference(?!)
		_gridmin[i] = T(ismissing(gridmin[i]) ? NaN : gridmin[i])
	end
	for i ∈ 1:length(gridmax)
		_gridmax[i] = T(ismissing(gridmax[i]) ? NaN : gridmax[i])
	end
	for i ∈ 1:length(gridpoints)
		_gridpoints[i] = T(ismissing(gridpoints[i]) ? NaN : gridpoints[i])
	end

	__wildboottest(
		matconvert(T,R),
		vecconvert(T,r);
		resp=vecconvert(T,resp),
		predexog=matconvert(T,predexog),
		predendog=matconvert(T,predendog),
		inst=matconvert(T,inst),
		R1=matconvert(T,R1),
		r1=vecconvert(T,r1),
		clustid=matconvert(Int64,clustid),
		nbootclustvar=Int64(nbootclustvar),
		nerrclustvar=Int64(nerrclustvar),
		issorted,
		hetrobust,
		nfe,
		feid=vecconvert(Int64,feid),
		fedfadj,
		obswt=vecconvert(T,obswt),
		fweights,
		maxmatsize=Float16(maxmatsize),
		ptype,
		bootstrapc,
		liml,
		fuller=T(fuller),
		kappa=T(kappa),
		arubin,
		small,
		clusteradj,
		clustermin,
		scorebs,
		reps=Int64(reps),
		imposenull,
		auxwttype,
		rng,
		level=T(level),
		rtol=T(rtol),
		madjtype,
		nH0=Int16(nH0),
		ml,
		scores=matconvert(T,scores),
		beta=vecconvert(T,beta),
		A=Symmetric(matconvert(T,A)),
		gridmin=_gridmin,
		gridmax=_gridmax,
		gridpoints=_gridpoints,
		diststat,
		getci,
		getplot,
		getauxweights)
end

_wildboottest(T::DataType, R, r::Number; kwargs...) = _wildboottest(T, R, [r]; kwargs...)
_wildboottest(T::DataType, R::UniformScaling{Bool}, r; kwargs...) = _wildboottest(T, diagm(fill(T(R.λ),nrows(r))), r; kwargs...)
_wildboottest(R::Union{UniformScaling{Bool},AbstractVecOrMat}, r::Union{Number,AbstractVecOrMat}; kwargs...) = _wildboottest(Float64, R, r; kwargs...)

"""
wildboottest([T::DataType=Float64,] R::AbstractMatrix, r::AbstractVector;
             resp, <optional keyword arguments>) -> WildBootTests.BootTestResult

Function to perform wild-bootstrap-based hypothesis test

# Positional arguments
* `T::DataType`: data type for inputs, results, and computations: Float32 or Float64 (default)
* `R::AbstractMatrix` and `r::AbstractVector`: required matrix and vector expressing the null Rβ=r; see notes below

# Required keyword argument
* `resp::AbstractVector`: response/dependent variable (y or y₁ in Roodman et al. (2019))

# Optional keyword arguments
* `predexog::AbstractVecOrMat`: exogenous predictors, including constant term, if any (X/X₁)
* `predendog::AbstractVecOrMat`: endogenous predictors (Y₂)
* `inst::AbstractVecOrMat`: instruments (X₂)
* `R1::AbstractMatrix` and `r1::AbstractVector`: model constraints; same format as for `R` and `r`
* `clustid::AbstractVecOrMat{<:Integer}`: data vector/matrix of error and bootstrapping cluster identifiers; see Notes
* `nbootclustvar::Integer=1`: number of bootstrap-clustering variables
* `nerrclustvar::Integer=nbootclustvar`: number of error-clustering variables
* `issorted:Bool=false`: time-saving flag: data matrices are already sort by column types 2, then 3, then 1 (see notes)
* `hetrobust::Bool=true`: true unless errors are treated as iid
* `nfe::Integer=0`: number of fixed-effect groups; if 0 yet `feid` is provided, will be computed
* `feid::AbstractVector{<:Integer}`: data vector for one-way fixed effect group identifier
* `fedfadj::Integer=nfe`: degrees of freedom that fixed effects (if any) consume
* `obswt::AbstractVector=[]`: observation weight vector; default is equal weighting
* `fweights::Bool=false`: true for frequency weights
* `maxmatsize::Number`: maximum size of auxilliary weight matrix (v), in gigabytes
* `ptype::Symbol=:symmetric`: p value type (`:symmetric`, `:equaltail`, `:lower`, `:upper`)
* `bootstrapc::Bool=false`: true to request bootstrap-c instead of bootstrap-t
* `liml::Bool=false`: true for LIML or Fuller LIML
* `fuller::Number`: Fuller LIML factor
* `kappa::Number`: fixed κ for _k_-class estimation
* `arubin::Bool=false`: true for Anderson-Rubin test
* `small::Bool=true`: true to multiply test statistics by G/(G-1) × N/(N-k), where G, N, k are number of clusters, observations, and predictors
* `clusteradj::Bool=true`: false to drop G/(G-1) factor
* `clustermin::Bool=false``: for multiway clustering, true to base G/(G-1) factor for all clusterings ]on the smallest G across clusterings
* `scorebs::Bool=false`: true for score bootstrap instead of wild bootstrap
* `reps::Integer=999`: number of bootstrap replications; `reps` = 0 requests classical Rao (or Wald) test if `imposenull` = `true` (or `false`)
* `imposenull::Bool=true`: true to impose null
* `auxwttype::Symbol=:rademacher`: auxilliary weight type (`:rademacher`, `:mammen`, `:webb`, `:normal`, `:gamma`)
* `rng::AbstractRNG=MersenneTwister()`: randon number generator
* `level::Number=.95`: significance level (0-1)
* `rtol::Number=1e-6`: tolerance for ci bound determination
* `madjtype::Symbol=:none`: multiple hypothesis adjustment (`none`, `:bonferroni`, `:sidak`)
* `nH0::Integer=1`: number of hypotheses tested, including one being tested now
* `ml::Bool=false`: true for (nonlinear) ML estimation
* `scores::AbstractVecOrMat`: for ML, pre-computed scores
* `beta::AbstractVector`: for ML, parameter estimates
* `A::AbstractMatrix`: for ML, covariance estimates
* `gridmin`: vector of graph lower bounds; max length 2, `missing`/`NaN` entries ask wildboottest() to choose
* `gridmax`: vector of graph upper bounds; `missing`/`NaN` entries ask wildboottest() to choose
* `gridpoints`: vector of number of sampling points; `missing`/`NaN` entries ask wildboottest() to choose
* `diststat::Symbole=:none`: `:t` to save bootstrap distribution of t/z/F/χ² statistics; `:numer` to save numerators thereof
* `getci::Bool=true`: whether to return confidence interval
* `getplot::Bool=getci`: whether to generate plot data
* `getauxweights::Bool=false`: whether to save auxilliary weight matrix (v)

# Notes
`T`, `ptype`, `auxwttype`, `madjtype`, and `diststat` may also be strings. Examples: `"Float32"` and `"webb"`.

The columns of `R` in the statement of the null should correspond to those of the matrix [`predexog` `predendog`],
where `predendog` is non-empty only in regressions with instruments.

Order the columns of `clustid` this way:
1. Variables only used to define bootstrapping clusters, as in the subcluster bootstrap.
2. Variables used to define both bootstrapping and error clusters.
3. Variables only used to define error clusters.
`nbootclustvar` is then the number of columns of type 1 or 2; `nerrclustvar` is the number of columns of type 2 or 3. Typically `clustid` is a single column of type 2.

`wildboottest()` does not handle missing data values: all data and identifier matrices must
be restricted to the estimation sample.

"""
wildboottest(   R, r; kwargs...) = _wildboottest(                                                  R, r; Dict(a.first => isa(a.second, AbstractString) ? Meta.parse(a.second) : a.second for a ∈ kwargs)...)
wildboottest(T, R, r; kwargs...) = _wildboottest(isa(T, AbstractString) ? eval(Meta.parse(T)) : T, R, r; Dict(a.first => isa(a.second, AbstractString) ? Meta.parse(a.second) : a.second for a ∈ kwargs)...)


addprocs(4)

function test()
  a = randn(1000)
  @parallel (+) for i=1:100000
    fetch(a[randi(end)])
  end
end