"""
SmetsWouters{T} <: AbstractModel{T}

The SmetsWouters type defines the structure of the FRBNY DSGE
model. We can then concisely pass around a Model object to the
remaining steps of the model (solve, estimate, and forecast).

### Fields

#### Parameters and Steady-States
* `parameters::Vector{AbstractParameter}`: Vector of all time-invariant model parameters.

* `steady_state::Vector`: Model steady-state values, computed as a
function of elements of `parameters`.

* `keys::Dict{Symbol,Int}`: Maps human-readable names for all model
parameters and steady-states to their indices in `parameters` and
`steady_state`.

#### Inputs to Measurement and Equilibrium Condition Equations

The following fields are dictionaries that map human-readible names to
row and column indices in the matrix representations of of the
measurement equation and equilibrium conditions.

* `endogenous_states::Dict{Symbol,Int}`: Maps each state to a column
in the measurement and equilibrium condition matrices.

* `exogenous_shocks::Dict{Symbol,Int}`: Maps each shock to a column in
the measurement and equilibrium condition matrices.

* `expected_shocks::Dict{Symbol,Int}`: Maps each expected shock to a
column in the measurement and equilibrium condition matrices.

* `equilibrium_conditions::Dict{Symbol,Int}`: Maps each equlibrium
condition to a row in the model's equilibrium condition matrices.

* `endogenous_states_augmented::Dict{Symbol,Int}`: Maps lagged states
to their columns in the measurement and equilibrium condition
equations. These are added after Gensys solves the model.

* `observables::Dict{Symbol,Int}`: Maps each observable to a row in
the model's measurement equation matrices.

#### Model Specifications and Settings

* `spec::AbstractString`: The model specification identifier, "m990",
cached here for filepath computation.

* `subspec::AbstractString`: The model subspecification number,
indicating that some parameters from the original model spec ("ss0")
are initialized differently. Cached here for filepath computation.

* `settings::Dict{Symbol,Setting}`: Settings/flags that affect
computation without changing the economic or mathematical setup of
the model.

* `test_settings::Dict{Symbol,Setting}`: Settings/flags for testing mode

#### Other Fields

* `rng::MersenneTwister`: Random number generator. By default, it is
seeded to ensure replicability in algorithms that involve randomness
(such as Metropolis-Hastings).

* `testing::Bool`: Indicates whether the model is in testing mode. If
`true`, settings from `m.test_settings` are used in place of those in
`m.settings`.

"""
type SmetsWouters{T} <: AbstractModel{T}
    parameters::ParameterVector{T}                  # vector of all time-invariant model parameters
    steady_state::ParameterVector{T}                # model steady-state values
    keys::Dict{Symbol,Int}                          # human-readable names for all the model
                                                    # parameters and steady-n_states

    endogenous_states::Dict{Symbol,Int}             # these fields used to create matrices in the
    exogenous_shocks::Dict{Symbol,Int}              # measurement and equilibrium condition equations.
    expected_shocks::Dict{Symbol,Int}               #
    equilibrium_conditions::Dict{Symbol,Int}        #
    endogenous_states_augmented::Dict{Symbol,Int}   #
    observables::Dict{Symbol,Int}                   #

    spec::ASCIIString                               # Model specification number
    subspec::ASCIIString                            # Model subspecification
    settings::Dict{Symbol,Setting}                  # Settings/flags for computation
    test_settings::Dict{Symbol,Setting}             # Settings/flags for testing mode
    rng::MersenneTwister                            # Random number generator
    testing::Bool                                   # Whether we are in testing mode or not
end

description(m::SmetsWouters) = "Smets-Wouters Model"

"""
`init_model_indices!(m::SmetsWouters)`

Arguments:
`m:: SmetsWouters`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and equilibrium conditions.
"""
function init_model_indices!(m::SmetsWouters)
    # Endogenous states
    endogenous_states = [[
        :y_t, :c_t, :i_t, :qk_t, :k_t, :kbar_t, :u_t, :rk_t, :mc_t,
        :π_t, :μ_ω_t, :w_t, :L_t, :R_t, :g_t, :b_t, :μ_t, :z_t,
        :λ_f_t, :λ_f_t1, :λ_w_t, :λ_w_t1, :rm_t, :Ec_t, :Eqk_t, :Ei_t,
        :Eπ_t, :EL_t, :Erk_t, :Ew_t, :y_f_t, :c_f_t,
        :i_f_t, :qk_f_t, :k_f_t, :kbar_f_t, :u_f_t, :rk_f_t, :w_f_t,
        :L_f_t, :r_f_t, :Ec_f_t, :Eqk_f_t, :Ei_f_t, :EL_f_t, :Erk_f_t, :ztil_t];
        [symbol("rm_tl$i") for i = 1:n_anticipated_shocks(m)]]

    # Exogenous shocks
    exogenous_shocks = [[
        :g_sh, :b_sh, :μ_sh, :z_sh, :λ_f_sh, :λ_w_sh, :rm_sh];
        [symbol("rm_shl$i") for i = 1:n_anticipated_shocks(m)]]

    # Expectations shocks
    expected_shocks = [
        :Ec_sh, :Eqk_sh, :Ei_sh, :Eπ_sh, :EL_sh, :Erk_sh, :Ew_sh, :Ec_f_sh,
        :Eqk_f_sh, :Ei_f_sh, :EL_f_sh, :Erk_f_sh]

    # Equilibrium conditions
    equilibrium_conditions = [[
        :eq_euler, :eq_inv, :eq_capval, :eq_output, :eq_caputl, :eq_capsrv, :eq_capev,
        :eq_mkupp, :eq_phlps, :eq_caprnt, :eq_msub, :eq_wage, :eq_mp, :eq_res, :eq_g, :eq_b, :eq_μ, :eq_z,
        :eq_λ_f, :eq_λ_w, :eq_rm, :eq_λ_f1, :eq_λ_w1, :eq_Ec,
        :eq_Eqk, :eq_Ei, :eq_Eπ, :eq_EL, :eq_Erk, :eq_Ew, :eq_euler_f, :eq_inv_f,
        :eq_capval_f, :eq_output_f, :eq_caputl_f, :eq_capsrv_f, :eq_capev_f, :eq_mkupp_f, :eq_caprnt_f, :eq_msub_f,
        :eq_res_f, :eq_Ec_f, :eq_Eqk_f, :eq_Ei_f, :eq_EL_f, :eq_Erk_f, :eq_ztil];
        [symbol("eq_rml$i") for i=1:n_anticipated_shocks(m)]]

    # Additional states added after solving model
    # Lagged states and observables measurement error
    endogenous_states_augmented = [
        :y_t1, :c_t1, :i_t1, :w_t1, :π_t1, :L_t1, :Et_π_t]

    # Measurement equation observables
    observables = [[
        :obs_gdp,              # quarterly output growth
        :obs_hours,            # aggregate hours growth
        :obs_wages,            # real wage growth
        :obs_gdpdeflator,      # inflation (GDP deflator)
        :obs_nominalrate,      # nominal interest rate
        :obs_consumption,      # consumption growth
        :obs_investment];       # investment growth
        [symbol("obs_nominalrate$i") for i=1:n_anticipated_shocks(m)]] # compounded nominal rates

    for (i,k) in enumerate(endogenous_states);            m.endogenous_states[k]            = i end
    for (i,k) in enumerate(exogenous_shocks);             m.exogenous_shocks[k]             = i end
    for (i,k) in enumerate(expected_shocks);              m.expected_shocks[k]              = i end
    for (i,k) in enumerate(equilibrium_conditions);       m.equilibrium_conditions[k]       = i end
    for (i,k) in enumerate(endogenous_states);            m.endogenous_states[k]            = i end
    for (i,k) in enumerate(endogenous_states_augmented); m.endogenous_states_augmented[k] = i+length(endogenous_states) end
    for (i,k) in enumerate(observables);                  m.observables[k]                  = i end
end

function SmetsWouters(subspec::AbstractString="ss0")

    # Model-specific specifications
    spec               = split(basename(@__FILE__),'.')[1]
    subspec            = subspec
    settings           = Dict{Symbol,Setting}()
    test_settings      = Dict{Symbol,Setting}()
    rng                = MersenneTwister()        # Random Number Generator
    testing            = false

    # initialize empty model
    m = SmetsWouters{Float64}(
            # model parameters and steady state values
            Vector{AbstractParameter{Float64}}(), Vector{Float64}(), Dict{Symbol,Int}(),

            # model indices
            Dict{Symbol,Int}(), Dict{Symbol,Int}(), Dict{Symbol,Int}(), Dict{Symbol,Int}(), Dict{Symbol,Int}(), Dict{Symbol,Int}(),

            spec,
            subspec,
            settings,
            test_settings,
            rng,
            testing)

    # Set settings
    settings_smets_wouters!(m)
    default_test_settings!(m)

    # Initialize parameters
    m <= parameter(:α,      0.24, (1e-5, 0.999), (1e-5, 0.999),   SquareRoot(),     Normal(0.30, 0.05),         fixed=false,
                   description="α: Capital elasticity in the intermediate goods sector's Cobb-Douglas production function.",
                   tex_label="\\alpha")

    m <= parameter(:ζ_p,   0.7813, (1e-5, 0.999), (1e-5, 0.999),   SquareRoot(),     BetaAlt(0.5, 0.1),          fixed=false,
                   description="ζ_p: The Calvo parameter. In every period, intermediate goods producers optimize prices with probability (1-ζ_p). With probability ζ_p, prices are adjusted according to a weighted average of the previous period's inflation (π_t1) and steady-state inflation (π_star).",
                   tex_label="\\zeta_p")

    m <= parameter(:ι_p,   0.3291, (1e-5, 0.999), (1e-5, 0.999),   SquareRoot(),     BetaAlt(0.5, 0.15),         fixed=false,
                   description="ι_p: The weight attributed to last period's inflation in price indexation. (1-ι_p) is the weight attributed to steady-state inflation.",

                   tex_label="\\iota_p")

    m <= parameter(:δ,      0.025,  fixed=true,
                   description="δ: The capital depreciation rate.", tex_label="\\delta" )

    m <= parameter(:Upsilon,  1.000,  (0., 10.),     (1e-5, 0.),      DSGE.Exponential(),    GammaAlt(1., 0.5),          fixed=true,
                   description="Υ: The trend evolution of the price of investment goods relative to consumption goods. Set equal to 1.",
                   tex_label="\\mathcal{\\Upsilon}")

    m <= parameter(:Φ,   1.4672, (1., 10.),     (1.00, 10.00),   DSGE.Exponential(),    Normal(1.25, 0.12),         fixed=false,
                   description="Φ: Fixed costs.",
                   tex_label="\\Phi")

    m <= parameter(:S′′,       6.3325, (-15., 15.),   (-15., 15.),     Untransformed(),  Normal(4., 1.5),            fixed=false,
                   description="S'': The second derivative of households' cost of adjusting investment.", tex_label="S\\prime\\prime")

    m <= parameter(:h,        0.7205, (1e-5, 0.999), (1e-5, 0.999),   SquareRoot(),     BetaAlt(0.7, 0.1),          fixed=false,
                   description="h: Consumption habit persistence.", tex_label="h")

    m <= parameter(:ppsi,     0.2648, (1e-5, 0.999), (1e-5, 0.999),   SquareRoot(),     BetaAlt(0.5, 0.15),         fixed=false,
                   description="ppsi: Utilization costs.", tex_label="ppsi")

    m <= parameter(:ν_l,     2.8401, (1e-5, 10.),   (1e-5, 10.),     DSGE.Exponential(),    Normal(2, 0.75),            fixed=false,
                   description="ν_l: The coefficient of relative risk aversion on the labor term of households' utility function.", tex_label="\\nu_l")

    m <= parameter(:ζ_w,   0.7937, (1e-5, 0.999), (1e-5, 0.999),   SquareRoot(),     BetaAlt(0.5, 0.1),          fixed=false,
                   description="ζ_w: (1-ζ_w) is the probability with which households can freely choose wages in each period. With probability ζ_w, wages increase at a geometrically weighted average of the steady state rate of wage increases and last period's productivity times last period's inflation.",
                   tex_label="\\zeta_w")

    m <= parameter(:ι_w,   0.4425, (1e-5, 0.999), (1e-5, 0.999),   SquareRoot(),     BetaAlt(0.5, 0.15),         fixed=false,
                   description="ι_w: No description available.",
                   tex_label="\\iota_w")

    m <= parameter(:λ_w,      1.5000,                                                                               fixed=true,
                   description="λ_w: The wage markup, which affects the elasticity of substitution between differentiated labor services.",
                   tex_label="\\lambda_w")

    m <= parameter(:β, 0.7420, (1e-5, 10.),   (1e-5, 10.),     DSGE.Exponential(),    GammaAlt(0.25, 0.1),        fixed=false,  scaling = x -> 1/(1 + x/100),
                   description="β: Discount rate.",
                   tex_label="\\beta ")

    m <= parameter(:ψ1,  1.7985, (1e-5, 10.),   (1e-5, 10.00),   DSGE.Exponential(),    Normal(1.5, 0.25),          fixed=false,
                   description="ψ₁: Weight on inflation gap in monetary policy rule.",
                   tex_label="\\psi_1")

    m <= parameter(:ψ2,  0.0893, (-0.5, 0.5),   (-0.5, 0.5),     Untransformed(),  Normal(0.12, 0.05),         fixed=false,
                   description="ψ₂: Weight on output gap in monetary policy rule.",
                   tex_label="\\psi_2")

    m <= parameter(:ψ3, 0.2239, (-0.5, 0.5),   (-0.5, 0.5),     Untransformed(),  Normal(0.12, 0.05),         fixed=false,
                   description="ψ₃: Weight on rate of change of output gap in the monetary policy rule.",
                   tex_label="\\psi_3")

    m <= parameter(:π_star,   0.7000, (1e-5, 10.),   (1e-5, 10.),     DSGE.Exponential(),    GammaAlt(0.62, 0.1),        fixed=false,  scaling = x -> 1 + x/100,
                   description="π_star: The steady-state rate of inflation.",
                   tex_label="\\pi_*")

    m <= parameter(:σ_c, 1.2312, (1e-5, 10.),   (1e-5, 10.),     DSGE.Exponential(),    Normal(1.5, 0.37),          fixed=false,
                   description="σ_c: No description available.",
                   tex_label="\\sigma_{c}")

    m <= parameter(:ρ,      .8258, (1e-5, 0.999), (1e-5, 0.999),   SquareRoot(),     BetaAlt(0.75, 0.10),        fixed=false,
                   description="ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho")

    m <= parameter(:ϵ_p,     10.000,                                                                               fixed=true,
                   description="ϵ_p: No description available.",
                   tex_label="\\varepsilon_{p}")

    m <= parameter(:ϵ_w,     10.000,                                                                               fixed=true,
                   description="ϵ_w: No description available.",
                   tex_label="\\varepsilon_{w}")

    # exogenous processes - level
    m <= parameter(:γ,      0.3982, (-5.0, 5.0),     (-5., 5.),     Untransformed(), Normal(0.4, 0.1),            fixed=false, scaling = x -> x/100,
                   description="γ: The log of the steady-state growth rate of technology.",
                   tex_label="\\gamma")

    m <= parameter(:Lmean,  875., (-1000., 1000.), (-1e3, 1e3),   Untransformed(), Normal(-45, 5),   fixed=false,
                   description="Lmean: No description available.",
                   tex_label="Lmean")

    m <= parameter(:g_star,    0.1800,                                                                               fixed=true,
                   description="g_star: No description available.",
                   tex_label="g_*")

    # exogenous processes - autocorrelation
    m <= parameter(:ρ_g,      0.9930, (1e-5, 0.999),   (1e-5, 0.999), SquareRoot(),    BetaAlt(0.5, 0.2),           fixed=false,
                   description="ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label="\\rho_g")

    m <= parameter(:ρ_b,      0.2703, (1e-5, 0.999),   (1e-5, 0.999), SquareRoot(),    BetaAlt(0.5, 0.2),           fixed=false,
                   description="ρ_b: AR(1) coefficient in the intertemporal preference shifter process.",
                   tex_label="\\rho_b")

    m <= parameter(:ρ_μ,     0.5724, (1e-5, 0.999),   (1e-5, 0.999), SquareRoot(),    BetaAlt(0.5, 0.2),           fixed=false,
                   description="ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label="\\rho_{\\mu}")

    m <= parameter(:ρ_z,      0.9676, (1e-5, 0.999),   (1e-5, 0.999), SquareRoot(),    BetaAlt(0.5, 0.2),           fixed=false,
                   description="ρ_z: AR(1) coefficient in the technology process.",
                   tex_label="\\rho_z")

    m <= parameter(:ρ_λ_f,    0.8692, (1e-5, 0.999),   (1e-5, 0.999), SquareRoot(),    BetaAlt(0.5, 0.2),           fixed=false,
                   description="ρ_λ_f: AR(1) coefficient in the price mark-up shock process.",
                   tex_label="\\rho_{\\lambda_f}")

    m <= parameter(:ρ_λ_w,    0.9546, (1e-5, 0.999),   (1e-5, 0.999), SquareRoot(),    BetaAlt(0.5, 0.2),           fixed=false,
                   description="ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.", # CHECK THIS
                   tex_label="\\rho_{\\lambda_f}")

    #monetary policy shock - see eqcond
    m <= parameter(:ρ_rm,     0.3000, (1e-5, 0.999),   (1e-5, 0.999), SquareRoot(),    BetaAlt(0.5, 0.2),           fixed=false,
                   description="ρ_rm: AR(1) coefficient in the monetary policy shock process.", # CHECK THIS
                   tex_label="\\rho_{rm}")

    # exogenous processes - standard deviation
    m <= parameter(:σ_g,      0.6090, (1e-8, 5.),      (1e-8, 5.),    DSGE.Exponential(),   RootInverseGamma(2., 0.10),  fixed=false,
                   description="σ_g: The standard deviation of the government spending process.",
                   tex_label="\\sigma_{g}")

    m <= parameter(:σ_b,      0.1818, (1e-8, 5.),      (1e-8, 5.),    DSGE.Exponential(),   RootInverseGamma(2., 0.10),  fixed=false,
                   description="σ_b: No description available.",
                   tex_label="\\sigma_{b}")

    m <= parameter(:σ_μ,     0.4601, (1e-8, 5.),      (1e-8, 5.),    DSGE.Exponential(),   RootInverseGamma(2., 0.10),  fixed=false,
                   description="σ_μ: The standard deviation of the exogenous marginal efficiency of investment shock process.",
                   tex_label="\\sigma_{\\mu}")

    m <= parameter(:σ_z,      0.4618, (1e-8, 5.),      (1e-8, 5.),    DSGE.Exponential(),   RootInverseGamma(2., 0.10),  fixed=false,
                   description="σ_z: No description available.",
                   tex_label="\\sigma_{z}")

    m <= parameter(:σ_λ_f,    0.1455, (1e-8, 5.),      (1e-8, 5.),    DSGE.Exponential(),   RootInverseGamma(2., 0.10),  fixed=false,
                   description="σ_λ_f: The mean of the process that generates the price elasticity of the composite good.  Specifically, the elasticity is (1+λ_{f,t})/(λ_{f_t}).",
                   tex_label="\\sigma_{\\lambda_f}")

    m <= parameter(:σ_λ_w,    0.2089, (1e-8, 5.),      (1e-8, 5.),    DSGE.Exponential(),   RootInverseGamma(2., 0.10),  fixed=false,
                   description="σ_λ_w: No description available.",
                   tex_label="\\sigma_{\\lambda_w}")

    m <= parameter(:σ_rm,     0.2397, (1e-8, 5.),      (1e-8, 5.),    DSGE.Exponential(),   RootInverseGamma(2., 0.10),  fixed=false,
                   description="σ_rm: No description available.",
                   tex_label="\\sigma_{rm}")

    m <= parameter(:η_gz,       0.0500, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.50, 0.20), fixed=false,
                   description="η_gz: Correlate g and z shocks.",
                   tex_label="\\eta_{gz}")

    m <= parameter(:η_λ_f,      0.7652, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.50, 0.20),         fixed=false,
                   description="η_λ_f: No description available.",
                   tex_label="\\eta_{\\lambda_f}")

    m <= parameter(:η_λ_w,      0.8936, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.50, 0.20),         fixed=false,
                   description="η_λ_w: AR(2) coefficient on wage markup shock process.",
                   tex_label="\\eta_{\\lambda_w}")

    # steady states
    m <= SteadyStateParameter(:zstar,  NaN, description="steady-state growth rate of productivity", tex_label="\\z_*")
    m <= SteadyStateParameter(:rstar,   NaN, description="steady-state something something", tex_label="\\r_*")
    m <= SteadyStateParameter(:Rstarn,  NaN, description="steady-state something something", tex_label="\\R_*_n")
    m <= SteadyStateParameter(:rkstar,  NaN, description="steady-state something something", tex_label="\\BLAH")
    m <= SteadyStateParameter(:wstar,   NaN, description="steady-state something something", tex_label="\\w_*")
    m <= SteadyStateParameter(:Lstar,   NaN, description="steady-state something something", tex_label="\\L_*")
    m <= SteadyStateParameter(:kstar,   NaN, description="Effective capital that households rent to firms in the steady state.", tex_label="\\k_*")
    m <= SteadyStateParameter(:kbarstar, NaN, description="Total capital owned by households in the steady state.", tex_label="\\bar{k}_*")
    m <= SteadyStateParameter(:istar,  NaN, description="Detrended steady-state investment", tex_label="\\i_*")
    m <= SteadyStateParameter(:ystar,  NaN, description="steady-state something something", tex_label="\\y_*")
    m <= SteadyStateParameter(:cstar,  NaN, description="steady-state something something", tex_label="\\c_*")
    m <= SteadyStateParameter(:wl_c,   NaN, description="steady-state something something", tex_label="\\wl_c")

    init_model_indices!(m)
    init_subspec!(m)
    steadystate!(m)
    return m
end

"""
```
steadystate!(m::SmetsWouters)
```

Calculates the model's steady-state values. `steadystate!(m)` must be called whenever the parameters of `m` are updated.
"""
function steadystate!(m::SmetsWouters)
    m[:zstar]    = log(1+m[:γ]) + m[:α]/(1-m[:α])*log(m[:Upsilon])
    m[:rstar]    = exp(m[:σ_c]*m[:zstar]) / m[:β]
    m[:Rstarn]   = 100*(m[:rstar]*m[:π_star] - 1)
    m[:rkstar]   = m[:rstar]*m[:Upsilon] - (1-m[:δ])
    m[:wstar]    = (m[:α]^m[:α] * (1-m[:α])^(1-m[:α]) * m[:rkstar]^(-m[:α]) / m[:Φ])^(1/(1-m[:α]))
    m[:Lstar]    = 1.
    m[:kstar]    = (m[:α]/(1-m[:α])) * m[:wstar] * m[:Lstar] / m[:rkstar]
    m[:kbarstar] = m[:kstar] * (1+m[:γ]) * m[:Upsilon]^(1 / (1-m[:α]))
    m[:istar]    = m[:kbarstar] * (1-((1-m[:δ])/((1+m[:γ]) * m[:Upsilon]^(1/(1-m[:α])))))
    m[:ystar]    = (m[:kstar]^m[:α]) * (m[:Lstar]^(1-m[:α])) / m[:Φ]
    m[:cstar]    = (1-m[:g_star])*m[:ystar] - m[:istar]
    m[:wl_c]     = (m[:wstar]*m[:Lstar])/(m[:cstar]*m[:λ_w])

    return m
end

function settings_smets_wouters!(m::SmetsWouters)

    default_settings!(m)

    # Anticipated shocks
    m <= Setting(:n_anticipated_shocks, 0)
    m <= Setting(:n_anticipated_shocks_padding, 20)
    m <= Setting(:n_anticipated_lags, 26)

    # Estimation
    m <= Setting(:reoptimize, true)
    m <= Setting(:recalculate_hessian, true)

    # Data vintage
    m <= Setting(:data_vintage, "150827")

    m.settings
end


using LinearAlgebra

export stamp_cutter #
export gen_pix_mask #
export condCovEst_wdiag #
export build_cov! #

"""
    stamp_cutter(cx,cy,residimIn,star_im,maskim;Np=33) -> data_in, stars_in, kmasked2d

Cuts out local stamps around each star of the various input images to be used for
per star statistics calculations.

# Arguments:
- `cx`: center coorindate x of the stamp
- `cy`: center coorindate y of the stamp
- `residimIn`: residual image with infilling from which covariance was estimated
- `star_im`: input image of model of stars only. abs(mod_im-sky_im)
- `maskim`: input image of upstream masked pixels

# Keywords:
- `Np`: size of covariance matrix footprint around each star (default 33)

# Outputs:
- `data_in`: local stamp of the (non-infilled) residual image
- `stars_in`: local stamp of model of stars only
- `kmasked2d`: local stamp of upstream masked pixels
"""
function stamp_cutter(cx,cy,residimIn,star_im,maskim;Np=33)
    radNp = (Np-1)÷2

    cov_stamp = cx-radNp:cx+radNp,cy-radNp:cy+radNp
    @views data_in = residimIn[cov_stamp[1],cov_stamp[2]]
    @views stars_in = star_im[cov_stamp[1],cov_stamp[2]]
    @views kmasked2d = maskim[cov_stamp[1],cov_stamp[2]];
    return data_in, stars_in, kmasked2d
end

"""
    gen_pix_mask(kmasked2d,psfmodel,circmask,x_star,y_star,flux_star;Np=33,thr=20) -> psft, kstar, kpsf2d, kcond0, kcond, kpred, dnt

Assigns pixels in the local subimage around a star to either be "good", "hidden", or "ignored"
based on user settings and the flux of the star. Reads in masked pixels from the quality
flags on pixels coming from the community pipeline `kmasked2d`, a PSF model for the star,
and a precomputed circular mask `circmask` to exclude pixels at a large radius from the
stellar center since they have little impact on the regression of hidden pixels. The pixels
assigned as "hidden" and to be interpolated are determined by a `thr` on the pixel values
for `flux_star` times the PSF model. We use a parametric PSFs that varies with position and
query the PSF at the stellar position for each star.

# Arguments:
- `kmasked2d`: Bool mask from upstream pixel quality flags to assign pixels as "ignored"
- `psfmodel`: parametric PSF model that can be queried at different positions
- `circmask`: static Bool mask assigning pixels beyond some radius of the stellar center as "ignored"
- `x_star`: x-coordinate of the star (used only for flexible PSF model query)
- `y_star`: y-coordinate of the star (used only for flexible PSF model query)
- `flux_star`: flux of star in ADU to determine how large a region to make "hidden"

# Keywords:
- `Np`: size of local covariance matrix in pixels (default 33)
- `thr`: threshold for psf-based masking of the residuals (larger more "hidden")

# Outputs:
- `psft`: static array (image) of the stellar PSF
- `kstar`: Boolean indexes the NOT "good" pixels
- `kpsf2d`: Boolean indexes the "hidden" pixels
- `kcond0`: initial number of "good" pixels
- `kcond`: final number of "good" pixels after fallbacks
- `kpred`: the number of pixels "hidden"
- `dnt`: quality flag bits on the solution
"""
function gen_pix_mask(kmasked2d,psfmodel,circmask,x_star,y_star,flux_star;Np=33,thr=20)

    psft = psfmodel(x_star,y_star,Np)

    if flux_star < 1e4  #these are the pixels we want the cov of
        kpsf2d = (psft .> thr/1e4)
    else
        kpsf2d = (psft .> thr/flux_star)
    end

    kstar = (kmasked2d .| kpsf2d .| circmask)
    kcond0 = Np^2-count(kstar)

    dnt = 0
    if kcond0 < 4*Np
        dnt += 1
        kstar = (kmasked2d .| kpsf2d)
    end

    kcond = Np^2-count(kstar)

    if kcond < 4*Np #this is about a 10% cut, and is the sum of bndry
        dnt += 2
        kstar[1,:] .= 0
        kstar[end,:] .= 0
        kstar[:,1] .= 0
        kstar[:,end] .= 0

        kpsf2d[1,:] .= 0
        kpsf2d[end,:] .= 0
        kpsf2d[:,1] .= 0
        kpsf2d[:,end] .= 0
    end

    psfr = minimum(psft)/maximum(psft)
    if psfr < 0
        dnt += 8
        psft .= abs.(psft)
    end
    if psfr < -1e-3
        dnt += 16
    end
    if psfr < -1e-1
        dnt += 32
    end

    kpred = count(kpsf2d)
    kcond = Np^2-count(kstar)

    return psft, kstar[:], kpsf2d, kcond0, kcond, kpred, dnt
end

"""
    condCovEst_wdiag(cov_loc,μ,km,kpsf2d,data_in,stars_in,psft;Np=33,export_mean=false,n_draw=0,diag_on=true) -> out

Using a local covariance matrix estimate `cov_loc` and a set of known ("good") pixels `km`
and "hidden" pixels `kpsf2d`, this function computes a prediction for the mean value
of the `kpsf2d` pixels and the covariance matrix of the `kpsf2d` pixels. In terms of
statistics use to adjust the photometry of a star, we are only interested in the
pixels masked as a result of the star (i.e. not a detector defect or cosmic ray nearby).
The residual image `data_in` and a model of the counts above the background coming from the
star `stars_in` for the local patch are also inputs of the function. Correction factors for
the photometric flux and flux uncertainities are outputs as well as a chi2 value for the
"good" pixels. The output list can conditionally include the mean reconstruction and
draws from the distribution of reconstructions.

# Arguments:
- `cov_loc`: local covariance matrix
- `μ`: vector containing mean value for each pixel in the patch
- `km`: unmasked pixels
- `kpsf2d`: pixels masked due to the star of interest
- `data_in`: (non-infilled) residual image in local patch
- `psft`: static array (image) of the stellar PSF

# Keywords:
- `Np`: size of local covariance matrix in pixels (default 33)
- `export_mean`: when true, returns the mean conditional prediction for the "hidden" pixels (default false)
- `n_draw`: when nonzero, returns that number of realizations of the conditional infilling (default 0)
- `diag_on`: flag for adding to the pixelwise uncertainty based on the photoelectron counts of the modeled star (default true)

# Outputs:
- `out[1][1]`: flux uncertainity of the star
- `out[1][2]`: flux uncertainity of the star assuming the covariance matrix were diagonal
- `out[1][3]`: flux correction which must be added to correct the input flux estimate
- `out[1][4]`: flux correction coming from the residuals (fdb_res)
- `out[1][5]`: flux correction coming from the predicted background (fdb_pred)
- `out[1][6]`: chi2 for the "good" pixels under `cov_loc` as a metric on how good our assumptions are
- `out[2]`: local region (image) with "hidden" pixels replaced by the mean conditional estimate (optional output)
- `out[end]`: local region (image) with "hidden" pixels replaced by the draws from the conditional distribution (optional output). Array is flattened to npix x n_draw.
"""
function condCovEst_wdiag(cov_loc,μ,km,kpsf2d,data_in,stars_in,psft;Np=33,export_mean=false,n_draw=0,diag_on=true)
    k = .!km
    kstar = kpsf2d[:]
    if diag_on
        for i=1:Np*Np cov_loc[i,i] += stars_in[i] end
    end
    cov_kk = Symmetric(cov_loc[k,k])
    cov_kkstar = cov_loc[k,kstar];
    cov_kstarkstar = cov_loc[kstar,kstar];
    icov_kkC = cholesky(cov_kk)
    icovkkCcovkkstar = icov_kkC\cov_kkstar
    predcovar = Symmetric(cov_kstarkstar - (cov_kkstar'*icovkkCcovkkstar))
    ipcovC = cholesky(predcovar)

    @views uncond_input = data_in[:]
    @views cond_input = data_in[:].- μ

    kstarpredn = (cond_input[k]'*icovkkCcovkkstar)'
    kstarpred = kstarpredn .+ μ[kstar]
    @views p = psft[kpsf2d][:]
    ipcovCp = ipcovC\p

    #@views std_wdiag = sqrt(abs(sum((pw.^(2)).*diag(predcovar[kpsf1d_kstar,kpsf1d_kstar]))))/sum(p2w)
    @views var_wdb = (p'*ipcovCp)

    var_diag = 0
    for i in 1:size(predcovar)[1]
        var_diag += (p[i]^2)/predcovar[i,i]
    end

    @views resid_mean = (uncond_input[kstar]'*ipcovCp)./var_wdb
    @views pred_mean = (kstarpred'*ipcovCp)./var_wdb

    #if we can afford it, a nice check would be how good of a covariance matrix
    #cov is for the k pixels (which we think are clean)
    chi2 = cond_input[k]'*(icov_kkC\cond_input[k])

    # Currently limited to the Np region. Often useful to have some context with a larger
    # surrounding region... TO DO to implement
    out = []
    push!(out,[sqrt(var_wdb^(-1)) sqrt(var_diag^(-1)) pred_mean-resid_mean resid_mean pred_mean chi2])
    if export_mean
        mean_out = copy(data_in)
        mean_out[kstar] .= kstarpred
        push!(out,mean_out)
    end
    if n_draw != 0
        covsvd = svd(predcovar)
        sqrt_cov = covsvd.V*diagm(sqrt.(covsvd.S))*covsvd.Vt;
        noise = sqrt_cov*randn(size(sqrt_cov)[1],n_draw)

        draw_out = repeat(copy(data_in)[:],outer=[1 n_draw])
        draw_out[kstar,:] .= repeat(kstarpred,outer=[1 n_draw]) .+ noise
        push!(out,draw_out)
    end

    return out
end

"""
    build_cov!(cov::Array{T,2},μ::Array{T,1},cx::Int,cy::Int,bimage::Array{T,2},bism::Array{T,4},Np::Int,widx::Int,widy::Int) where T <:Union{Float32,Float64}

Constructs the local covariance matrix and mean for an image patch of size `Np` x `Np` pixels around a location
of interest (`cx`,`cy`). The construction is just a lookup of pixel values from the stored boxcar-smoothed copies
of the input image times itself shifted in `bism`. Passing the smoothed image `bimage` and the widths of the boxcar
mean `widx` and `widy` is helpful for the mean and normalization. The covariance and mean are updated in place
for speed since this operation may be performed billions of times since we construct a new covariance matrix for
every detection. Math may either be performed `Float32` or `Float64`.

# Arguments:
- `cov::Array{T,2}`: preallocated output array for local covariance matrix
- `μ::Array{T,1}`: preallocated output vector for local mean
- `cx::Int`: x-coordinate of the center of the local region
- `cy::Int`: y-coordinate of the center of the local region
- `bimage::Array{T,2}`: boxcar smoothed unshifted image
- `bism::Array{T,4}`: boxcar-smoothed image products for all shifts
- `Np::Int`: size of local covariance matrix in pixels
- `widx::Int`: width of boxcar window in x which determines size of region used for samples for the local covariance estimate
- `widy::Int`: width of boxcar window in y which determines size of region used for samples for the local covariance estimate
"""
function build_cov!(cov::Array{T,2},μ::Array{T,1},cx::Int,cy::Int,bimage::Array{T,2},bism::Array{T,4},Np::Int,widx::Int,widy::Int) where T <:Union{Float32,Float64}
    Δx = (widx-1)÷2
    Δy = (widy-1)÷2
    halfNp = (Np-1) ÷ 2
    Δr, Δc = cx-(halfNp+1), cy-(halfNp+1)
    # Δr, Δc = cx-(halfNp-1), cy-(halfNp-1)
    for dc=0:Np-1       # column shift loop
        pcr = 1:Np-dc
        for dr=1-Np:Np-1# row loop, incl negatives
            if (dr < 0) & (dc == 0)
                continue
            end
            if dr >= 0
                prr = 1:Np-dr
            end
            if (dr < 0) & (dc > 0)
                prr = 1-dr:Np
            end

            for pc=pcr, pr=prr
                i = ((pc   -1)*Np)+pr
                j = ((pc+dc-1)*Np)+pr+dr
                @inbounds μ1μ2 = bimage[pr+Δr,pc+Δc]*bimage[pr+dr+Δr,pc+dc+Δc]/((widx*widy)^2)
                @inbounds t = bism[pr+Δr,pc+Δc,dr+Np,dc+1]/(widx*widy) - μ1μ2
                @inbounds cov[i,j] = t
                @inbounds cov[j,i] = t
                if i == j
                    @inbounds μ[i] = μ1μ2
                end
            end
        end
    end
    cov .*= (widx*widy)/((widx*widy)-1)
    return
end


module ParametricOptInterface

using MathOptInterface
const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const DD = MOIU.DoubleDicts
const POI = ParametricOptInterface
const PARAMETER_INDEX_THRESHOLD = 1_000_000_000_000_000_000
const SUPPORTED_SETS = [MOI.LessThan{Float64}, MOI.EqualTo{Float64}, MOI.GreaterThan{Float64}]

"""
    Parameter(Float64)

The `Parameter` structure stores the numerical value associated to a given parameter.
# Example:
```julia-repl
julia> ParametricOptInterface.Parameter(5)
ParametricOptInterface.Parameter(5)
```
"""
struct Parameter <: MOI.AbstractScalarSet
    val::Float64
end

"""
    ParametricOptimizer{T, OT <: MOI.ModelLike} <: MOI.AbstractOptimizer

Declares a `ParametricOptimizer`, which allows the handling of parameters in a optimization model.
# Example:
```julia-repl
julia> ParametricOptInterface.ParametricOptimizer(GLPK.Optimizer())
ParametricOptInterface.ParametricOptimizer{Float64,GLPK.Optimizer}
```
"""
mutable struct ParametricOptimizer{T, OT <: MOI.ModelLike} <: MOI.AbstractOptimizer
    optimizer::OT
    parameters::Dict{MOI.VariableIndex, T}
    parameters_name::Dict{MOI.VariableIndex, String}
    updated_parameters::Dict{MOI.VariableIndex, T}
    variables::Dict{MOI.VariableIndex, MOI.VariableIndex}
    last_variable_index_added::Int
    last_parameter_index_added::Int
    affine_constraint_cache::DD.DoubleDict{Vector{MOI.ScalarAffineTerm{Float64}}}
    quadratic_constraint_cache_pv::Dict{MOI.ConstraintIndex, Vector{MOI.ScalarQuadraticTerm{Float64}}}
    quadratic_constraint_cache_pp::Dict{MOI.ConstraintIndex, Vector{MOI.ScalarQuadraticTerm{Float64}}}
    quadratic_constraint_cache_pc::DD.DoubleDict{Vector{MOI.ScalarAffineTerm{Float64}}}
    quadratic_constraint_variables_associated_to_parameters_cache::Dict{MOI.ConstraintIndex, Vector{MOI.ScalarAffineTerm{T}}}
    quadratic_added_cache::Dict{MOI.ConstraintIndex, MOI.ConstraintIndex}
    last_quad_add_added::Int
    affine_objective_cache::Vector{MOI.ScalarAffineTerm{T}}
    quadratic_objective_cache_pv::Vector{MOI.ScalarQuadraticTerm{T}}
    quadratic_objective_cache_pp::Vector{MOI.ScalarQuadraticTerm{T}}
    quadratic_objective_cache_pc::Vector{MOI.ScalarAffineTerm{T}}
    quadratic_objective_variables_associated_to_parameters_cache::Vector{MOI.ScalarAffineTerm{T}}
    multiplicative_parameters::BitSet
    dual_value_of_parameters::Vector{Float64}
    evaluate_duals::Bool
    number_of_parameters_in_model::Int
    function ParametricOptimizer(optimizer::OT; evaluate_duals::Bool=true) where OT
        new{Float64, OT}(
            optimizer,
            Dict{MOI.VariableIndex, Float64}(),
            Dict{MOI.VariableIndex, String}(),
            Dict{MOI.VariableIndex, Float64}(),
            Dict{MOI.VariableIndex, MOI.VariableIndex}(),
            0,
            PARAMETER_INDEX_THRESHOLD,
            DD.DoubleDict{Vector{MOI.ScalarAffineTerm{Float64}}}(),
            Dict{MOI.ConstraintIndex, Vector{MOI.ScalarQuadraticTerm{Float64}}}(),
            Dict{MOI.ConstraintIndex, Vector{MOI.ScalarQuadraticTerm{Float64}}}(),
            DD.DoubleDict{Vector{MOI.ScalarAffineTerm{Float64}}}(),
            Dict{MOI.ConstraintIndex, Vector{MOI.ScalarAffineTerm{Float64}}}(),
            Dict{MOI.ConstraintIndex, MOI.ConstraintIndex}(),
            0,
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            Vector{MOI.ScalarQuadraticTerm{Float64}}(),
            Vector{MOI.ScalarQuadraticTerm{Float64}}(),
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            Vector{MOI.ScalarAffineTerm{Float64}}(),
            BitSet(),
            Vector{Float64}(),
            evaluate_duals,
            0
        )
    end
end

include("utils.jl")
include("duals.jl")
include("update_parameters.jl")

function MOI.is_empty(model::ParametricOptimizer)
    return MOI.is_empty(model.optimizer) &&
    isempty(model.parameters) &&
    isempty(model.parameters_name) &&
    isempty(model.variables) &&
    isempty(model.updated_parameters) &&
    isempty(model.variables) &&
    model.last_variable_index_added == 0 &&
    model.last_parameter_index_added == PARAMETER_INDEX_THRESHOLD &&
    isempty(model.affine_constraint_cache) &&
    isempty(model.quadratic_constraint_cache_pv) &&
    isempty(model.quadratic_constraint_cache_pp) &&
    isempty(model.quadratic_constraint_cache_pc) &&
    isempty(model.quadratic_constraint_variables_associated_to_parameters_cache) &&
    isempty(model.quadratic_added_cache) &&
    model.last_quad_add_added == 0 &&
    isempty(model.affine_objective_cache) &&
    isempty(model.quadratic_objective_cache_pv) &&
    isempty(model.quadratic_objective_cache_pp) &&
    isempty(model.quadratic_objective_cache_pc) &&
    isempty(model.quadratic_objective_variables_associated_to_parameters_cache) &&
    isempty(model.dual_value_of_parameters) &&
    model.number_of_parameters_in_model == 0
end

function MOI.supports_constraint(
    model::ParametricOptimizer,
    F::Union{Type{MOI.SingleVariable},
             Type{MOI.ScalarAffineFunction{T}},
             Type{MOI.VectorOfVariables},
             Type{MOI.VectorAffineFunction{T}}},
    S::Type{<:MOI.AbstractSet}) where T

    return MOI.supports_constraint(model.optimizer, F, S)
end

function MOI.supports_constraint(
    model::ParametricOptimizer,
    ::Type{MOI.ScalarQuadraticFunction{T}},
    S::Type{<:MOI.AbstractSet}) where T
    return MOI.supports_constraint(model.optimizer, MOI.ScalarAffineFunction{T}, S)
end

function MOI.supports(
    model::ParametricOptimizer,
    attr::Union{MOI.ObjectiveSense,
            MOI.ObjectiveFunction{MOI.SingleVariable},
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}},
            MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{T}}
            }) where T
    return MOI.supports(model.optimizer, attr)
end

function MOI.empty!(model::ParametricOptimizer{T}) where T
    MOI.empty!(model.optimizer)
    empty!(model.parameters)
    empty!(model.parameters_name)
    empty!(model.updated_parameters)
    empty!(model.variables)
    model.last_variable_index_added = 0
    model.last_parameter_index_added = PARAMETER_INDEX_THRESHOLD
    empty!(model.affine_constraint_cache)
    empty!(model.quadratic_constraint_cache_pv)
    empty!(model.quadratic_constraint_cache_pp)
    empty!(model.quadratic_constraint_cache_pc)
    empty!(model.quadratic_constraint_variables_associated_to_parameters_cache)
    empty!(model.quadratic_added_cache)
    model.last_quad_add_added = 0
    empty!(model.affine_objective_cache)
    empty!(model.quadratic_objective_cache_pv)
    empty!(model.quadratic_objective_cache_pp)
    empty!(model.quadratic_objective_cache_pc)
    empty!(model.quadratic_objective_variables_associated_to_parameters_cache)
    empty!(model.dual_value_of_parameters)
    model.number_of_parameters_in_model = 0
    return
end

function MOI.set(model::ParametricOptimizer, attr::MOI.VariableName, v::MOI.VariableIndex, name::String)
    if is_parameter_in_model(model, v)
        model.parameters_name[v] = name
    else
        return MOI.set(model.optimizer, attr, v, name)
    end
end

function MOI.get(model::ParametricOptimizer, attr::MOI.VariableName, v::MOI.VariableIndex)
    if is_parameter_in_model(model, v)
        return get(model.parameters_name, v, "")
    else
        return MOI.get(model.optimizer, attr, v)
    end
end

function MOI.supports(model::ParametricOptimizer, attr::MOI.VariableName, tp::Type{MOI.VariableIndex})
    return MOI.supports(model.optimizer, attr, tp)
end

function MOI.set(model::ParametricOptimizer, attr::MOI.ConstraintName, c::MOI.ConstraintIndex, name::String)
    return MOI.set(model.optimizer, attr, c, name)
end

function MOI.get(model::ParametricOptimizer, attr::MOI.ConstraintName, c::MOI.ConstraintIndex)
    return MOI.get(model.optimizer, attr, c)
end

function MOI.get(model::ParametricOptimizer, ::MOI.NumberOfVariables)
    return length(model.parameters) + length(model.variables)
end

function MOI.supports(model::ParametricOptimizer, attr::MOI.ConstraintName, tp::Type{<:MOI.ConstraintIndex})
    return MOI.supports(model.optimizer, attr, tp)
end

# TODO
# This is not correct, you need to put the parameters back into the function
# function MOI.get(model::ParametricOptimizer, attr::MOI.ConstraintFunction, ci::MOI.ConstraintIndex{F, S}) where {F, S}
#     MOI.get(model.optimizer, attr, ci)
# end

function MOI.get(model::ParametricOptimizer, attr::MOI.ConstraintSet, ci::MOI.ConstraintIndex{F, S}) where {F, S}
    MOI.throw_if_not_valid(model, ci)
    return MOI.get(model.optimizer, attr, ci)
end

function MOI.get(model::ParametricOptimizer, attr::MOI.ObjectiveSense)
    return MOI.get(model.optimizer, attr)
end

function MOI.get(model::ParametricOptimizer, attr::MOI.ObjectiveFunctionType)
    if !isempty(model.quadratic_constraint_cache_pv) || !isempty(model.quadratic_constraint_cache_pc)
        return MOI.ScalarQuadraticFunction{Float64}
    end
    return MOI.get(model.optimizer, attr)
end

# TODO
# Same as ConstraintFunction getter. And you also need to convert to F
# function MOI.get(
#     model::ParametricOptimizer,
#     attr::MOI.ObjectiveFunction{F}) where F <: Union{MOI.SingleVariable,MOI.ScalarAffineFunction{T}} where T

#     MOI.get(model.optimizer, attr)
# end

# TODO
# You might have transformed quadratic functions into affine functions so this is incorrect
# function MOI.get(model::ParametricOptimizer, ::MOI.ListOfConstraints)
#     constraints = Set{Tuple{DataType, DataType}}()
#     inner_ctrs = MOI.get(model.optimizer, MOI.ListOfConstraints())
#     for (F, S) in inner_ctrs
#         push!(constraints, (F,S))
#     end

#     collect(constraints)
# end

function MOI.get(
    model::ParametricOptimizer,
    attr::MOI.ListOfConstraintIndices{F, S}
) where {S, F<:Union{
    MOI.VectorOfVariables,
    MOI.SingleVariable,
}}
    return MOI.get(model.optimizer, attr)
end

# TODO
# You might have transformed quadratic functions into affine functions so this is incorrect
# function MOI.get(
#     model::ParametricOptimizer,
#     ::MOI.ListOfConstraintIndices{F, S}
# ) where {S<:MOI.AbstractSet, F<:Union{
#     MOI.ScalarAffineFunction{T},
#     MOI.VectorAffineFunction{T},
# }} where T
#     MOI.get(model.optimizer, MOI.ListOfConstraintIndices{F, S}())
# end

function MOI.supports_add_constrained_variable(
    ::ParametricOptimizer, ::Type{Parameter})
    return true
end

function MOI.supports_add_constrained_variables(
    model::ParametricOptimizer, ::Type{MOI.Reals})
    return MOI.supports_add_constrained_variables(model.optimizer, MOI.Reals)
end

function MOI.add_variable(model::ParametricOptimizer)
    next_variable_index!(model)
    v_p = MOI.VariableIndex(model.last_variable_index_added)
    v = MOI.add_variable(model.optimizer)
    model.variables[v_p] = v
    return v_p
end

"""
    MOI.add_constrained_variable(model::ParametricOptimizer, set::Parameter)

Adds a parameter, that is, a variable parameterized to the value provided in `set`, to the `model` specified.
The `model` must be a `ParametricOptInterface.ParametricOptimizer` model.
Returns the MOI.VariableIndex of the parameterized variable and the MOI.ConstraintIndex associated.

#Example:
```julia-repl
julia> w, cw = MOI.add_constrained_variable(optimizer, POI.Parameter(0))
(MathOptInterface.VariableIndex(1), MathOptInterface.ConstraintIndex{MathOptInterface.SingleVariable,ParametricOptInterface.Parameter}(1))
```
"""
function MOI.add_constrained_variable(model::ParametricOptimizer, set::Parameter)
    next_parameter_index!(model)
    p = MOI.VariableIndex(model.last_parameter_index_added)
    model.parameters[p] = set.val
    cp = MOI.ConstraintIndex{MOI.SingleVariable, Parameter}(model.last_parameter_index_added)
    update_number_of_parameters!(model)
    return p, cp
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.SingleVariable, set::MOI.AbstractScalarSet)
    if is_parameter_in_model(model, f.variable)
        error("Cannot constrain a parameter")
    elseif is_variable_in_model(model, f.variable)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        error("Variable not in the model")
    end
end

function add_constraint_with_parameters_on_function(model::ParametricOptimizer, f::MOI.ScalarAffineFunction{T}, set::MOI.AbstractScalarSet) where T
    vars, params, param_constant = separate_possible_terms_and_calculate_parameter_constant(model, f.terms)
    ci = MOIU.normalize_and_add_constraint(model.optimizer, MOI.ScalarAffineFunction(vars, f.constant + param_constant), set)
    model.affine_constraint_cache[ci] = params
    return ci
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.ScalarAffineFunction{T}, set::MOI.AbstractScalarSet) where T
    if !function_has_parameters(model, f)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        return add_constraint_with_parameters_on_function(model, f, set)
    end
end

function MOI.get(model::ParametricOptimizer, attr::MOI.VariablePrimal, v::MOI.VariableIndex)
    if is_parameter_in_model(model, v)
        return model.parameters[v]
    elseif is_variable_in_model(model, v)
        return MOI.get(model.optimizer, attr, model.variables[v])
    else
        error("Variable not in the model")
    end
end

"""
    MOI.set(model::ParametricOptimizer, ::MOI.ConstraintSet, cp::MOI.ConstraintIndex{MOI.SingleVariable, Parameter}, set::Parameter)

Sets the parameter to a given value, using its `MOI.ConstraintIndex` as reference.

#Example:
```julia-repl
julia> MOI.set(optimizer, MOI.ConstraintSet(), cw, POI.Parameter(2.0))
2.0
```
"""
function MOI.set(model::ParametricOptimizer, ::MOI.ConstraintSet, cp::MOI.ConstraintIndex{MOI.SingleVariable, Parameter}, set::Parameter)
    p = MOI.VariableIndex(cp.value)
    if is_parameter_in_model(model,p)
        return model.updated_parameters[p] = set.val
    else
        error("Parameter not in the model")
    end
end

struct ParameterValue <: MOI.AbstractVariableAttribute end

"""
    MOI.set(model::ParametricOptimizer, ::ParameterValue, vi::MOI.VariableIndex, val::Float64)

Sets the parameter to a given value, using its `MOI.VariableIndex` as reference.

#Example:
```julia-repl
julia> MOI.set(model, ParameterValue(), w, 2.0)
2.0
```
"""
function MOI.set(model::ParametricOptimizer, ::ParameterValue, vi::MOI.VariableIndex, val::Float64)
    if is_parameter_in_model(model, vi)
        return model.updated_parameters[vi] = val
    else
        error("Parameter not in the model")
    end
end
function MOI.set(model::ParametricOptimizer, ::ParameterValue, vi::MOI.VariableIndex, val::Real)
    return MOI.set(model, ParameterValue(), vi, convert(Float64, val))
end

function empty_objective_function_caches!(model:: ParametricOptimizer)
    empty!(model.affine_objective_cache)
    empty!(model.quadratic_objective_cache_pv)
    empty!(model.quadratic_objective_cache_pp)
    empty!(model.quadratic_objective_cache_pc)
    return nothing
end

function MOI.set(model::ParametricOptimizer, attr::MOI.ObjectiveFunction, f::MOI.ScalarAffineFunction{T}) where T

    # clear previously defined objetive function cache
    empty_objective_function_caches!(model)

    if !function_has_parameters(model, f)
        MOI.set(model.optimizer, attr, f)
        return
    else
        vars, params, param_constant = separate_possible_terms_and_calculate_parameter_constant(model, f.terms)
        MOI.set(model.optimizer, attr, MOI.ScalarAffineFunction(vars, f.constant + param_constant))
        model.affine_objective_cache = params
        return
    end
end

function MOI.set(model::ParametricOptimizer, attr::MOI.ObjectiveFunction, v::MOI.SingleVariable)
    if haskey(model.parameters, v)
        error("Cannot use a parameter as objective function alone")
    elseif haskey(model.variables, v)
        return MOI.set(model.optimizer, attr, MOI.SingleVariable(model.variables[v.variable]))
    else
        error("Variable not in the model")
    end
end

function MOI.set(model::ParametricOptimizer, attr::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    return MOI.set(model.optimizer, attr, sense)
end

function MOI.get(model::ParametricOptimizer, attr::T) where {
                                                                T <: Union{
                                                                    MOI.TerminationStatus,
                                                                    MOI.ObjectiveValue,
                                                                    MOI.PrimalStatus
                                                                    }
                                                            }
    return MOI.get(model.optimizer, attr)
end

function MOI.set(model::ParametricOptimizer, ::MOI.Silent, bool::Bool)
    MOI.set(model.optimizer, MOI.Silent(), bool)
end

function MOI.set(model::ParametricOptimizer, attr::MOI.RawParameter, val::Any)
    MOI.set(model.optimizer, attr, val)
end

function MOI.set(model::ParametricOptimizer, attr::String, val::Any)
    MOI.set(model.optimizer, MOI.RawParameter(attr), val)
end

function MOI.get(model::ParametricOptimizer, ::MOI.SolverName)
    return "ParametricOptimizer with " *
           MOI.get(model.optimizer, MOI.SolverName()) *
           " attached"
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.VectorOfVariables, set::MOI.AbstractVectorSet) where T
    if function_has_parameters(model, f)
        error("VectorOfVariables does not allow parameters")
    else
        return MOI.add_constraint(model.optimizer, f, set)
    end
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.VectorAffineFunction{T}, set::MOI.AbstractVectorSet) where T
    if function_has_parameters(model, f)
        error("VectorAffineFunction does not allow parameters")
    else
        return MOI.add_constraint(model.optimizer, f, set)
    end
end

function add_constraint_with_parameters_on_function(
                            model::ParametricOptimizer,
                            f::MOI.ScalarQuadraticFunction{T},
                            set::S) where {T, S <: MOI.AbstractScalarSet}

    if function_quadratic_terms_has_parameters(model, f.quadratic_terms)
        (quad_vars,
         quad_aff_vars,
         quad_params,
         aff_terms,
         variables_associated_to_parameters,
         quad_param_constant) = separate_possible_terms_and_calculate_parameter_constant(model, f.quadratic_terms)
    else
        quad_vars = f.quadratic_terms
        quad_aff_vars = MOI.ScalarQuadraticTerm{T}[]
        quad_params = MOI.ScalarQuadraticTerm{T}[]
        aff_terms = MOI.ScalarAffineTerm{T}[]
        quad_param_constant = zero(T)
        variables_associated_to_parameters = MOI.VariableIndex[]
    end

    if function_affine_terms_has_parameters(model, f.affine_terms)
        (aff_vars,
         aff_params,
         terms_with_variables_associated_to_parameters,
         aff_param_constant) = separate_possible_terms_and_calculate_parameter_constant(
                                                                                model,
                                                                                f.affine_terms,
                                                                                variables_associated_to_parameters
                                                                            )
    else
        aff_vars = f.affine_terms
        aff_params = MOI.ScalarAffineTerm{T}[]
        terms_with_variables_associated_to_parameters = MOI.ScalarAffineTerm{T}[]
        aff_param_constant = zero(T)
    end

    aff_terms = vcat(aff_terms, aff_vars)
    const_term = f.constant + aff_param_constant + quad_param_constant
    quad_terms = quad_vars

    f_quad = if !isempty(quad_vars)
        MOI.ScalarQuadraticFunction(
                    aff_terms,
                    quad_terms,
                    const_term
                )
    else
        MOI.ScalarAffineFunction(
                    aff_terms,
                    const_term
                )
    end

    model.last_quad_add_added += 1
    ci = MOIU.normalize_and_add_constraint(model.optimizer, f_quad, set)
    # This part is used to remember that ci came from a quadratic function
    # It is particularly useful because sometimes
    new_ci = MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{T}, S}(model.last_quad_add_added)
    model.quadratic_added_cache[new_ci] = ci

    fill_quadratic_constraint_caches!(
        model,
        new_ci,
        quad_aff_vars,
        quad_params,
        aff_params,
        terms_with_variables_associated_to_parameters
    )

    return ci
end

function MOI.add_constraint(model::ParametricOptimizer, f::MOI.ScalarQuadraticFunction{T}, set::MOI.AbstractScalarSet) where T
    if !function_has_parameters(model, f)
        return MOI.add_constraint(model.optimizer, f, set)
    else
        return add_constraint_with_parameters_on_function(model, f, set)
    end
end

function MOI.set(model::ParametricOptimizer, attr::MOI.ObjectiveFunction, f::MOI.ScalarQuadraticFunction{T}) where T

    # clear previously defined objetive function cache
    empty_objective_function_caches!(model)

    if !function_has_parameters(model, f)
        MOI.set(model.optimizer, attr, f)
        return
    else
        if function_quadratic_terms_has_parameters(model, f.quadratic_terms)
            (quad_vars,
             quad_aff_vars,
             quad_params,
             aff_terms,
             variables_associated_to_parameters,
             quad_param_constant) = separate_possible_terms_and_calculate_parameter_constant(model, f.quadratic_terms)
        else
            quad_vars = f.quadratic_terms
            quad_aff_vars = MOI.ScalarQuadraticTerm{T}[]
            quad_params = MOI.ScalarQuadraticTerm{T}[]
            aff_terms = MOI.ScalarAffineTerm{T}[]
            quad_param_constant = zero(T)
            variables_associated_to_parameters = MOI.VariableIndex[]
        end

        if function_affine_terms_has_parameters(model, f.affine_terms)
            (aff_vars,
             aff_params,
             terms_with_variables_associated_to_parameters,
             aff_param_constant) = separate_possible_terms_and_calculate_parameter_constant(
                                                                                model,
                                                                                f.affine_terms,
                                                                                variables_associated_to_parameters
                                                                                )
        else
            aff_vars = f.affine_terms
            aff_params = MOI.ScalarAffineTerm{T}[]
            terms_with_variables_associated_to_parameters = MOI.ScalarAffineTerm{T}[]
            aff_param_constant = zero(T)
        end

        aff_terms = vcat(aff_terms, aff_vars)
        const_term = f.constant + aff_param_constant + quad_param_constant
        quad_terms = quad_vars

        if !isempty(quad_vars)
            MOI.set(
                model.optimizer,
                attr,
                MOI.ScalarQuadraticFunction(aff_terms, quad_terms, const_term)
            )
        else
            MOI.set(
                model.optimizer,
                MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
                MOI.ScalarAffineFunction(aff_terms, const_term)
            )
        end

        if !isempty(quad_terms)
            f_quad = MOI.ScalarQuadraticFunction(
                        aff_terms,
                        quad_terms,
                        const_term
                    )

            MOI.set(model.optimizer, attr, f_quad)
        else
           f_quad = MOI.ScalarAffineFunction(
                        aff_terms,
                        const_term
                    )

            MOI.set(model.optimizer, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(), f_quad)
        end

        model.quadratic_objective_cache_pv = quad_aff_vars
        model.quadratic_objective_cache_pp = quad_params
        model.quadratic_objective_cache_pc = aff_params
        model.quadratic_objective_variables_associated_to_parameters_cache = terms_with_variables_associated_to_parameters

        return
    end
end

function MOI.optimize!(model::ParametricOptimizer)
    if !isempty(model.updated_parameters)
        update_parameters!(model)
    end

    MOI.optimize!(model.optimizer)

    if model.evaluate_duals
        calculate_dual_of_parameters(model)
    end
end

end # module


"""
    NeuronIndices

Treat vertices as having parameters which represents neurons when formulating the size change problem.

This means that individual indices will be aligned so that the function to the largest extent possible is the same after resizing.
"""
struct NeuronIndices end

Δsizetype(vs::AbstractVector{<:AbstractVertex}) = partialsort!(Δsizetype.(vs), 1; by=Δsizetypeprio, rev=true)
Δsizetype(v::AbstractVertex, seen=Set()) = v in seen ? nothing : Δsizetype(trait(v), v, push!(seen, v))
Δsizetype(t::DecoratingTrait, v, seen) = Δsizetype(base(t), v, seen)
Δsizetype(t::MutationTrait, v, seen) = Δsizetype(t, base(v), seen)
Δsizetype(::MutationTrait, ::InputVertex, seen) = nothing
Δsizetype(::MutationTrait, v::CompVertex, seen) = Δsizetype(v.computation)
Δsizetype(f) = NeuronIndices()

# TODO: Label elemwise and concatenation to allow them to return some low prio Δsizetype? Now things like BatchNorm will not have an opinion on its Δsizetype
# Perhaps better idea: pass down trait to Δsizetype(f) and default to low prio Δsizetype for SizeTransparent? Requires that one always looks at all vertices, but that seems to be the requirement anyways
Δsizetype(::SizeTransparent, v, seen) = partialsort!(vcat(map(vi -> Δsizetype(vi, seen), inputs(v)), map(vo -> Δsizetype(vo, seen), outputs(v))), 1; by=Δsizetypeprio, rev=true)

Δsizetypeprio(::Nothing) = 0
Δsizetypeprio(::ScalarSize) = 1
Δsizetypeprio(::NeuronIndices) = 2

"""
    defaultutility(v::AbstractVertex)

Default function used to calculate utility of output neurons.

Implement either `defaultutility(f)` or `defaultutility(t, f)` where `f` is the computation performed
by `CompVertex` and `t` is [`trait(v)`](@ref) to set the default for `f`.

# Examples

```jldoctest
julia> using NaiveNASlib, Statistics

julia> struct Affine{T}
        W::Matrix{T}
       end;

julia> NaiveNASlib.defaultutility(l::Affine) = mean(abs, l.W; dims=2);

julia> NaiveNASlib.defaultutility(Affine(ones(2,3)))
2×1 Matrix{Float64}:
 1.0
 1.0
```
"""
defaultutility(v::AbstractVertex) = defaultutility(trait(v), v)
defaultutility(t, v::AbstractVertex) = defaultutility(t, base(v))
defaultutility(t, ::InputVertex) = 1
defaultutility(t, v::CompVertex) = defaultutility(t, v.computation)
defaultutility(t, f) = resolve_utility(t, defaultutility(f))

struct NoDefaultUtility end

defaultutility(f) = NoDefaultUtility()

resolve_utility(t::DecoratingTrait, u::NoDefaultUtility) = resolve_utility(base(t),u)
resolve_utility(::MutationTrait, ::NoDefaultUtility) = 1
resolve_utility(::SizeTransparent, ::NoDefaultUtility) = 0 # rationale is that utility of other vertices will be tied to this anyways
resolve_utility(t, val) = val

Δsize!(::Nothing, s::AbstractΔSizeStrategy, vs::AbstractVector{<:AbstractVertex}) = false

# Main entry point for computing and applying size changes with NeuronIndices

Δsize!(case::NeuronIndices, s::AbstractΔSizeStrategy, vs::AbstractVector{<:AbstractVertex}) = Δsize!(defaultutility, case,s , vs)
function Δsize!(utilityfun, case::NeuronIndices, s::AbstractΔSizeStrategy, vs::AbstractVector{<:AbstractVertex})
    success, ins, outs = solve_outputs_selection(s, vs, utilityfun)
    if success
        applyΔsize!(case, s, ins, outs, vs)
    end
    return success
end

function Δsize!(utilityfun, case::NeuronIndices, s::TruncateInIndsToValid, vs::AbstractVector{<:AbstractVertex})
    success, ins, outs = solve_outputs_selection(s.strategy, vs, utilityfun)
    if success
        for (vv, ininds) in ins
            for innr in eachindex(ininds)
                ininds[innr] = aligntomax(nin(vv)[innr], ininds[innr])
            end
            newins, newouts = align_outs_to_ins(vv, ins[vv], outs[vv])
            ins[vv] = newins
            outs[vv] = newouts
        end
        applyΔsize!(case, s, ins, outs, vs)
    end
    return success
end

aligntomax(maxval, ::Missing) = missing
function aligntomax(maxval, arr)
    arr = copy(arr)
    maxind = argmax(arr)
    while arr[maxind] > maxval
        newval = maxval
        while newval in arr && newval > -1
            newval -= 1
        end
        arr[maxind] = newval == 0 ? -1 : newval
        maxind = argmax(arr)
    end
    return arr
end

align_outs_to_ins(v, ins, outs) = align_outs_to_ins(trait(v), v, ins, outs)
align_outs_to_ins(t::DecoratingTrait, v, ins, outs) = align_outs_to_ins(base(t), v, ins, outs)
align_outs_to_ins(t, v, ins, outs) = ins,outs
function align_outs_to_ins(::SizeInvariant, v, ins::AbstractVector, outs)
    isempty(ins) && return ins, outs
    inds = ins[1]
    newins = repeat([inds], length(ins))
    newouts = ismissing(outs) ? outs : inds
    return newins, newouts
end

# Main entry point for applying size changes with NeuronIndices after we have compted them
# Peel off all DecoratingJuMPΔSizeStrategies until we find either an AbstractAfterΔSizeStrategy (which is needed in after_Δnin etc)
# or we hit the bottom (in the form of an AbstractΔSizeStrategy).
function applyΔsize!(case::NeuronIndices, s::DecoratingJuMPΔSizeStrategy, ins::AbstractDict, outs::AbstractDict, vs::AbstractVector{<:AbstractVertex})
    applyΔsize!(case, base(s), ins, outs, vs)
end
function applyΔsize!(case::NeuronIndices, s::AbstractAfterΔSizeStrategy, ins::AbstractDict, outs::AbstractDict, vs::AbstractVector{<:AbstractVertex})
    _applyΔsize!(case, s, ins, outs, vs)
end
function applyΔsize!(case::NeuronIndices, s::AbstractΔSizeStrategy, ins::AbstractDict, outs::AbstractDict, vs::AbstractVector{<:AbstractVertex})
    _applyΔsize!(case, s, ins, outs, vs)
end

function _applyΔsize!(case::NeuronIndices, s, ins::AbstractDict, outs::AbstractDict, vs::AbstractVector{<:AbstractVertex})

    Δnins = map(vs) do vi
        vins = ins[vi]
        ismissing(vins) && return false
        mapreduce(|, zip(vins, nin(vi)); init=false) do (vin, ninorg)
            ismissing(vin) && return false
            return length(vin) != ninorg
        end
    end
    Δnouts = [length(outs[vi]) != nout(vi) for vi in vs]

    for vi in vs
        applyΔsize!(case, s, vi, ins[vi], outs[vi])
    end

    for (i, vi) in enumerate(vs)
        after_Δnin(s, vi, ins[vi], Δnins[i])
        after_Δnout(s, vi, outs[vi], Δnouts[i])
    end
end

# Main entry point for applying size changes to a single vertex with NeuronIndices

function applyΔsize!(case::NeuronIndices, s::WithKwargs, v::AbstractVertex, ins::AbstractVector, outs::AbstractVector; kws...)
    applyΔsize!(case, base(s), v, ins, outs; kws..., s.kwargs...)
end
function applyΔsize!(case::NeuronIndices, s::DecoratingJuMPΔSizeStrategy, v::AbstractVertex, ins::AbstractVector, outs::AbstractVector; kws...)
     applyΔsize!(case, base(s), v, ins, outs; kws...)
end
function applyΔsize!(case::NeuronIndices, s::AbstractΔSizeStrategy, v::AbstractVertex, ins::AbstractVector, outs::AbstractVector; kws...)
    applyΔsize!(case, v, ins, outs; kws...)
end

function applyΔsize!(::NeuronIndices, v::InputSizeVertex, ins::AbstractVector, outs::AbstractVector; kwargs...)
    if !all(isempty, skipmissing(ins))
        throw(ArgumentError("Try to change input neurons of InputVertex $(name(v)) to $ins"))
    end
    if outs != 1:nout(v)
        throw(ArgumentError("Try to change output neurons of InputVertex $(name(v)) to $outs"))
    end
end

function applyΔsize!(case::NeuronIndices, v::AbstractVertex, ins::AbstractVector, outs::AbstractVector; kwargs...)
    applyΔsize!(case, base(v), ins, outs; kwargs...)
end
function applyΔsize!(case::NeuronIndices, v::CompVertex, ins::AbstractVector, outs::AbstractVector; kwargs...)
    # I don't love how we go back to Δsize! from here, but I don't want to export applyΔsize! and I can't think of another name right now...
    Δsize!(case, v.computation, ins, outs; kwargs...)
end
Δsize!(::NeuronIndices, f, ins::AbstractVector, outs::AbstractVector; kwargs...) = Δsize!(f, ins, outs; kwargs...)
"""
    Δsize!(f::F, ins::AbstractVector, outs::AbstractVector; kwargs...)

Apply the changes to `f` so that input neurons in `ins` and output neurons in `outs` are selected and/or inserted.

Argument `outs` is a vector of indices to select/insert while `ins` has one vector of indices per input vertex.

Shall be implemented for any type `F` which holds parameters for which the shape shall be modified by NaiveNASlib.

Tip: the function [`parselect`](@ref) can be used to change parameter arrays according to `ins` and `outs`.

Tip: `kwargs` can be passed using [`WithKwargs`](@ref).
"""
function Δsize!(f, ins::AbstractVector, outs::AbstractVector; kwargs...) end

"""
    parselect(pars::AbstractArray{T,N}, elements_per_dim...; newfun = (T, dim, size...) -> 0) where {T, N}

Return a new array of same type as `pars` which has a subset of the elements of `pars` as well as potentially
new elements.

Which elements to select/insert is determined by `elements_per_dim` which is a `Pair{Int, Vector{Int}}` mapping
dimension (first memeber) to which elements to select/insert in that dimension (second memeber).

For a each `dim=>elems` pair, the following holds: `selectdim(output, dim, i) == selectdim(pars, dim, elems[i])`
if `elems[i]` is positive and `selectdim(output, dim, i) .== newfun(T, dim, size)[j]` if `elems[i]` is the `j:th`
negative value and `size` is `sum(elems .< 0)`.

# Examples
```jldoctest
julia> using NaiveNASlib, NaiveNASlib.Extend

julia> pars = reshape(1:3*5, 3,5)
3×5 reshape(::UnitRange{Int64}, 3, 5) with eltype Int64:
 1  4  7  10  13
 2  5  8  11  14
 3  6  9  12  15

julia> NaiveNASlib.parselect(pars, 1 => [-1, 1,3,-1,2], 2=>[3, -1, 2, 1]; newfun = (T, d, s...) -> fill(-T(d), s))
5×4 Matrix{Int64}:
 -1  -2  -1  -1
  7  -2   4   1
  9  -2   6   3
 -1  -2  -1  -1
  8  -2   5   2
```
"""
function parselect(pars::AbstractArray{T,N}, elements_per_dim...; newfun = (T, dim, size...) -> 0) where {T, N}
    psize = collect(size(pars))
    assign = repeat(Any[Colon()], N)
    access = repeat(Any[Colon()], N)

    elements_per_dim = filter(!ismissing ∘ last, elements_per_dim)

    for (dim, elements) in elements_per_dim
        psize[dim] = length(elements)
    end
    newpars = similar(pars, psize...)

    for (dim, elements) in elements_per_dim
        indskeep = filter(ind -> ind > 0, elements)
        indsmap = elements .> 0
        newmap = .!indsmap

        assign[dim] = findall(indsmap)
        access[dim] = indskeep
        tsize = copy(psize)
        tsize[dim] = sum(newmap)
        selectdim(newpars, dim, newmap) .= newfun(T, dim, tsize...)
    end

    newpars[assign...] = pars[access...]
    return newpars
end
parselect(::Missing, args...;kwargs...) = missing

function solve_outputs_selection(s::LogΔSizeExec, vertices::AbstractVector{<:AbstractVertex}, utilityfun)
    @logmsg s.level s.msgfun(vertices[1])
    return solve_outputs_selection(base(s), vertices, utilityfun)
end

solve_outputs_selection(::ΔSizeFailNoOp, vs::AbstractVector{<:AbstractVertex}, ::Any) = false, Dict(v => [1:n for n in nin(v)] for v in vs), Dict(v => 1:nout(v) for v in vs)
solve_outputs_selection(s::ThrowΔSizeFailError, vertices::AbstractVector{<:AbstractVertex}, ::Any) = throw(ΔSizeFailError(s.msgfun(vertices)))
solve_outputs_selection(s::WithUtilityFun, vs::AbstractVector{<:AbstractVertex}, ::Any) = solve_outputs_selection(s.strategy, vs, s.utilityfun)
"""
    solve_outputs_selection(s::AbstractΔSizeStrategy, vertices::AbstractVector{<:AbstractVertex}, utilityfun)

Returns a tuple `(success, nindict, noutdict)` where `nindict[vi]` are new input neuron indices and `noutdict[vi]` are new output neuron indices for each vertex `vi` in `vertices`.

The function generally tries to maximize `sum(utilityfun(vi) .* selected[vi]) ∀ vi in vertices` where `selected[vi]` is all elements in `noutdict[vi]` larger than 0 (negative values
in `noutdict` indicates a new output shall be inserted at that position). This however is up to the implementation of the `AbstractΔSizeStrategy s`.

Since selection of outputs is not guaranteed to work in all cases, a flag `success` is also returned. If `success` is `false` then applying the new indices may (and probably will) fail.
"""
function solve_outputs_selection(s::AbstractJuMPΔSizeStrategy, vs::AbstractVector{<:AbstractVertex}, orgutilityfun)
    utilityfun = remaputilityfun(orgutilityfun, s, vs)
    model = selectmodel(NeuronIndices(), s, vs, utilityfun)
    case = NeuronIndices()

    # TODO: Cache and rescale utilityfun?

    # The binary variables `outselectvars` tells us which existing output indices to select
    # The integer variables `outinsertvars` tells us where in the result we shall insert -1 where -1 means "create a new output (e.g. a neuron)
    # Thus, the result will consist of all selected indices with possibly interlaced -1s
    outselectvars, outinsertvars, noutdict = createselectvars!(case, s, vs, (;model, utilityfun))

    objexpr = sizeobjective!(case, s, vs, (;model, outselectvars, outinsertvars, noutdict, utilityfun))
    for v in vs
        data = (;model, outselectvars, outinsertvars, noutdict, objexpr, utilityfun)
        vertexconstraints!(case, v, s, data)
        objexpr = selectobjective!(case, s, v, data)
    end

    @objective(model, Max, objexpr)

    JuMP.optimize!(model)

    !accept(case, s, model) && return   let fbstrat = fallback(s)
                                            solve_outputs_selection(fbstrat, add_participants!(fbstrat, copy(vs)), orgutilityfun)
                                        end

    return true, extract_ininds_and_outinds(s, outselectvars, outinsertvars)...
end

remaputilityfun(utilityfun, s::DecoratingJuMPΔSizeStrategy, vs, args...) = remaputilityfun(utilityfun, base(s), vs, args...)
function remaputilityfun(utilityfun, ::AbstractJuMPΔSizeStrategy, vs, floorval=1e-2, ceilval=1e5)
    cache = Dict{eltype(vs), Any}(v => utilityfun(v) for v in vs)

    # For any problem which one thinks can be solved by writing code one will always generate new problems
    # Learn to keep the number of moving parts in a codebase to a minimum or spend an eternity writing
    # code like this
    !any(vs) do v
        issizemutable(v) || return false
        any(>(0), cache[v])
    end && return v -> cache[v]

    # We would prefer the range of the utility values to stay well within typical solver tolerances
    # Not sure how to do this correctly. Some papers point towards scaling the cost/constraint matrix,
    # but it does not seem to be availble from JuMP

    # What this ad-hoc-y algorithm tries to do is to check if the magnitudes are within bounds [floorval, ceilval]
    # If they are not, it tries to adjust them to be so by uniform scaling.
    # The scaling is limited so that it never pushes something out of bounds.
    # If magnitudes outside both bounds exist (i.e both too large and too small values exist) it does not do anything.
    # TODO: Use quantiles to not let outliers prevent scaling? Some optimization criterion for maximizing the fraction of values within bounds?

    # Don't let 0s affect the result or else all those SizeTransparent vertices is going to push us to max scaling all the time.
    minval, maxval = extrema(abs, Iterators.filter(!=(0), Iterators.flatten(val for (v, val) in cache if issizemutable(v))))
    lb = max(minval, floorval)
    ub = min(maxval, ceilval)
    maxscaledown = lb / floorval
    scaledown = min(maxscaledown, maxval / ub)
    maxscaleup = ceilval / ub
    scaleup = min(maxscaleup, lb / minval)

    scal = scaledown > scaleup ? 1 / scaledown : scaleup

    for (v, val) in cache
        if issizemutable(v)
            cache[v] = val .* scal
        end
    end
    return v -> cache[v]
end

createselectvars!(case, s::DecoratingJuMPΔSizeStrategy, vs, data) = createselectvars!(case ,base(s), vs, data)
function createselectvars!(case, s::UpperInsertBound{<:Integer}, vs, data)
    isempty(vs) && return createselectvars!(case, base(s), vs, data)
    createselectvars!(case, UpperInsertBound(vs, s.vbounds; strategy=base(s), fallback=fallback(s)), vs, data)
end
function createselectvars!(case, s::UpperInsertBound, vs, data)
    outselectvars, outinsertvars, noutdict = createselectvars!(case, base(s), vs, data)
    for (v, ub) in s.vbounds
        JuMP.set_upper_bound.(outinsertvars[v], ub)
    end
    return outselectvars, outinsertvars, noutdict
end
function createselectvars!(case, ::AbstractJuMPΔSizeStrategy, vs, data)
    model = data.model
    utilityfun = data.utilityfun
    # We use the lenght of utilityfun as the number of outputs to select as we will use it to add utility to each variable,
    # so their lenghts must not mismatch
    # They might mismatch when there are parameters attached to some size transparent vertex (e.g BatchNorm or even
    # concatenation with some computed activation utility metric) and we are in the process of aligning sizes after
    # adding/removing an edge.
    # Forcing the user to provide a length of utilityfun matching nout(v) in this case is not really fair as it is we who
    # caused the mismatch.
    # TODO: The indicies alignment below will not be ideal in this case. Perhaps it can be mitigated with some strategy
    # providing information on how the edges have changed so that we can better align what hasn't changed.
    selectsizes = map(vs) do v
        outvals = utilityfun(v)
        return outvals isa Number ? nout(v) : length(outvals)
    end

    outselectvars = Dict(v => @variable(model, [1:nvals], Bin) for (v,nvals) in zip(vs, selectsizes))
    outinsertvars = Dict(v => @variable(model, [1:nout(v)], Int, lower_bound=0) for v in vs)

    # This is the same variable name as used when adjusting size only. It allows us to delegate alot of size changing strategies to ScalarSize.
    # Drawback is that it will yield a fair bit of back and forth so maybe it is too clever for its own good
    noutdict = Dict(v => @expression(model, sum(outselectvars[v]) + sum(outinsertvars[v])) for v in vs)

    return outselectvars, outinsertvars, noutdict
end

selectmodel(case::NeuronIndices, s::DecoratingJuMPΔSizeStrategy, vs, utilityfun) = selectmodel(case, base(s), vs, utilityfun)
selectmodel(::NeuronIndices, ::AbstractJuMPΔSizeStrategy, vs, utilityfun) = JuMP.Model(JuMP.optimizer_with_attributes(Cbc.Optimizer, "loglevel"=>0, "seconds" => 20))
function selectmodel(case::NeuronIndices, s::TimeLimitΔSizeStrategy, vs, utilityfun)
    model = selectmodel(case, base(s), vs, utilityfun)
    JuMP.set_time_limit_sec(model, s.limit)
    return model
end

function vertexconstraints!(case::NeuronIndices, ::Immutable, v, s::AbstractJuMPΔSizeStrategy, data)
     JuMP.set_lower_bound.(data.outselectvars[v], 1)
     JuMP.set_upper_bound.(data.outinsertvars[v], 0)
     for (vi, outsize) in zip(inputs(v), nin(v))
        @constraint(data.model, data.noutdict[vi] == outsize)
     end
     # This is needed e.g. to apply (impossible) size constraints incase someone wants to change the input size of an output vertex
     vertexconstraints!(case ,s, v, data)
end

# ScalarSize constraints happen to be good here, but we must ensure that vertexconstraints!(::NeuronIndices,...) gets called instead of vertexconstraints!(::ScalarSize,...)
# For AlignNinToNout we could just have broken out the call to vertexconstraints!, but for AlignNinToNoutVertices we need to call it in the middle of the function
vertexconstraints!(::NeuronIndices, v::AbstractVertex, s::AlignNinToNout, data) = vertexconstraints!(ScalarSize(), v, s, data, NeuronIndices())
vertexconstraints!(::NeuronIndices, v::AbstractVertex, s::AlignNinToNoutVertices, data) = vertexconstraints!(ScalarSize(), v::AbstractVertex, s::AlignNinToNoutVertices, data, NeuronIndices())

function vertexconstraints!(case::NeuronIndices, s::AbstractJuMPΔSizeStrategy, v, data)
    insertconstraints!(case, s, v, data)
    sizeconstraint!(case, s, v, data)
    compconstraint!(case, s, v, (data..., vertex=v))
    inoutconstraint!(case, s, v, data)
end

insertconstraints!(case::NeuronIndices, s::DecoratingJuMPΔSizeStrategy, v, data) = insertconstraints!(case, base(s), v, data)
insertconstraints!(::NeuronIndices, ::AbstractJuMPΔSizeStrategy, v, data) = noinsertgaps!(data.model, data.outselectvars[v], data.outinsertvars[v])

"""
    noinsertgaps!(model, select, insert, maxinsert=length(outsel) * 10)

Add constraints so that `insert` does not create undefined gaps in the result of the neuron selection.

Assume `select` is a set of binary variables where `select[i] = 1` means select the output neuron at position `i`
and `insert[i] = N` means insert `N` new output neurons at the position after `i`.

An example of an undefined gap is if `select = [1, 1, 0]` and `insert = [0, 0, 1]` because this results in the
instruction to use existing output neurons `1 and 2` and then insert a new neuron at position `4`.
In this example position `3` is an undefined gap as one should neither put an existing neuron there nor shall
one insert new neurons. Running this method constrains `model` so that this solution is infeasible.
"""
function noinsertgaps!(model, select, insert, maxinsert=max(length(select) * 10, 200)) # TODO: get maxinsert from strategy instead?
    JuMP.set_upper_bound.(filter(!JuMP.has_upper_bound, insert), maxinsert)

    # This is redundant if we don't allow any inserts and the constraints here add significantly to the solution time
    all(var -> JuMP.upper_bound(var) == 0, insert) && return

    # See  https://discourse.julialang.org/t/help-with-constraints-to-select-and-or-insert-columns-to-a-matrix/63654
    insert_nogap = @variable(model, [1:length(insert)], Bin)

    @constraint(model, sum(insert) <= maxinsert)

    # insert[i] == 0 if insert_nogap[i] == 1
    @constraint(model, [i=1:length(insert)], insert[i] <= 2maxinsert * insert_nogap[i])
    # Monotonicity of insert_nogap, i.e insert_nogap[i] can only be 1 if insert_nogap[i+1] is 1
    @constraint(model, [i=2:length(insert)], insert_nogap[i] <= insert_nogap[i-1])
    # Force insert_nogap to have at least as many ones as the number of not selected neurons
    @constraint(model, length(insert) - sum(select) <= sum(insert_nogap))
end

sizeconstraint!(case::NeuronIndices, s::DecoratingJuMPΔSizeStrategy, v, data) = sizeconstraint!(case, base(s), v, data)
sizeconstraint!(::NeuronIndices, s::AbstractJuMPΔSizeStrategy, v, data) = sizeconstraint!(ScalarSize(), DefaultJuMPΔSizeStrategy(), v, data)
sizeconstraint!(case::NeuronIndices, s::ΔNoutMix, v, data) = sizeconstraint!(case, s.exact, v, data)
sizeconstraint!(::NeuronIndices, s::ΔNout{Exact}, v, data) = sizeconstraint!(ScalarSize(), s, v, data)

inoutconstraint!(case, s::DecoratingJuMPΔSizeStrategy, v, data) = inoutconstraint!(case, base(s), v, data)
inoutconstraint!(case, s, v, data) = inoutconstraint!(case, s, trait(v), v, data)
inoutconstraint!(case::NeuronIndices, s, t::DecoratingTrait, v, data) = inoutconstraint!(case, s, base(t), v, data)

function inoutconstraint!(::NeuronIndices, s, ::MutationTrait, v, data) end
function inoutconstraint!(case::NeuronIndices, s, t::SizeTransparent, v, data)
    onemismatch  = inoutconstraint!(case, s, t, v, data.model, data.outselectvars)
    onemismatch |= inoutconstraint!(case, s, t, v, data.model, data.outinsertvars)

    # We need to ensure that total sizes align if there was a mismatch
    if onemismatch
        ninconstraint!(case, s, v, data)
    end
end

function inoutconstraint!(::NeuronIndices, s, t::SizeStack, v, model, vardict::Dict)
    offs = 1
    var = vardict[v]
    onemismatch = false
    for (i, vi) in enumerate(inputs(v))
        var_i = vardict[vi]
        # Sizes mismatch when vertex/edge was removed (or edge added)
        if length(var_i) == nin(v)[i] && offs+length(var_i)-1 <= length(var)
            @constraint(model, var_i .== var[offs:offs+length(var_i)-1])
        else
            onemismatch = true
        end
        offs += length(var_i)
    end
    # Length var != nout means that we need to enforce the SizeStack constraint that sum(nin.(inputs(v))) == nout(v)
    # as this indicates that there is some parameter array which needs to be aligned with the new size
    return onemismatch || length(var) != nout(v)
end

function inoutconstraint!(::NeuronIndices, s, ::SizeInvariant, v, model, vardict::Dict)
    var = vardict[v]
    onemismatch = false
    for vi in inputs(v)
        # Sizes mismatch when vertex/edge was removed (or edge added)
        var_i = vardict[vi]
        if length(var_i) == length(var)
            @constraint(model, var_i .== var)
        else
            onemismatch = true
        end
    end
    return onemismatch
end

# Minus sign because we maximize objective function in case NeuronIndices
sizeobjective!(case::NeuronIndices, s::AbstractJuMPΔSizeStrategy, vertices, data) = -objective!(case, s, vertices, data)

objective!(case::NeuronIndices, s::ΔNoutMix, vertices, data) = objective!(case, s.relax, vertices, data)
objective!(case::NeuronIndices, s::ΔNout{Relaxed}, vertices, data) = noutrelax!(case, s.Δs, vertices, data)

objective!(case::NeuronIndices, s::DecoratingJuMPΔSizeStrategy, vertices, data) = objective!(case, base(s), vertices, data)
function objective!(::NeuronIndices, s::AbstractJuMPΔSizeStrategy, vertices, data)
    # Or else the problem becomes infeasible
    isempty(vertices) && return @expression(data.model, 0)

    scale = map(v -> maximum(abs, data.utilityfun(v)), vertices)
    noutvars = map(v -> data.noutdict[v], vertices)
    sizetargets = map(v -> sizetarget(data.utilityfun, v), vertices)
    # L1 norm prevents change in vertices which does not need to change.
    # Max norm tries to spread out the change so no single vertex takes most of the change.
    return norm!(SumNorm(0.1 => L1NormLinear(), 0.8 => MaxNormLinear()), data.model, @expression(data.model, objective[i=1:length(noutvars)], scale[i] * (noutvars[i] - sizetargets[i])), sizetargets)
end

function sizetarget(utilityfun, v)
    util = utilityfun(v)
    # Negative scalar utility is probably never useful, but lets make it consistent with an all negative array
    util isa Number && return util < 0 ? 1 : nout(v)
    return max(1, nout(v) - count(<(0), util))
end

function selectobjective!(case::NeuronIndices, s::AbstractJuMPΔSizeStrategy, v, data)
    utility = utilityobjective!(case, s, v, data)
    @objective(data.model, Max, utility)
    # Also makes us prefer to not insert
    insertlast = insertlastobjective!(case, s, v, data)
    return @expression(data.model, data.objexpr + utility + insertlast)
end
utilityobjective!(case::NeuronIndices, s::DecoratingJuMPΔSizeStrategy, v, data) = utilityobjective!(case, base(s), v, data)
utilityobjective!(::NeuronIndices, ::AbstractJuMPΔSizeStrategy, v, data) = @expression(data.model, sum(data.utilityfun(v) .* data.outselectvars[v]))

insertlastobjective!(case::NeuronIndices, s::DecoratingJuMPΔSizeStrategy, v, data) = insertlastobjective!(case::NeuronIndices, base(s), v, data)
function insertlastobjective!(::NeuronIndices, s, v, data)
    insvars = data.outinsertvars[v]

    isempty(insvars) && return @expression(data.model, 0)

    preferend = collect(length(insvars) : -1 : 1)
    scale = minimum(abs, data.utilityfun(v)) / sum(preferend)
    return @expression(data.model, -0.1scale*sum(insvars .* preferend))
end

extract_ininds_and_outinds(s::DecoratingJuMPΔSizeStrategy, outselectvars::Dict, outinsertvars::Dict) = extract_ininds_and_outinds(base(s), outselectvars, outinsertvars)
function extract_ininds_and_outinds(s, outselectvars::Dict, outinsertvars::Dict)
    outinds = Dict([v => extract_inds(s, outselectvars[v], outinsertvars[v]) for v in keys(outselectvars)])
    ininds = Dict([v => getall(outinds, inputs(v)) for v in keys(outselectvars)])
    return ininds, outinds
end

function extract_inds(::AbstractJuMPΔSizeStrategy, selectvars::T, insertvars::T) where T <: AbstractVector{JuMP.VariableRef}
    # insertvar is N at indices where a N new output neurons shall be added
    insert = round.(Int, JuMP.value.(insertvars))
    selected = findall(xi -> xi > 0, JuMP.value.(selectvars))
    return join_extracted_inds(selected, insert)
end

function join_extracted_inds(selected, insert)
    result = Vector{Int}(undef, length(selected) + sum(insert))

    # Expected result of algorithm: initialize result = selected, then insert insert[i] -1 elements after each index i in result

    i = 1 # result index
    j = 1 # selected index
    k = 1 # insert index
    while i <= length(result)
        if j <= i && j <= length(selected)
            result[i] = selected[j]
            j += 1
            i += 1
        end
        if k <= length(insert) && insert[k] > 0
            start = i
            stop = i+insert[k]-1
            result[start:stop] .= -1
            i += insert[k]
        end
        k += 1
    end
    return result
end


abstract AbstractMaterial

abstract AbstractMaterialStatus

stiffness() = error("Not implemented")
create_matstat() = error("Not implemented")
stress() = error("Not implemented")
get_kalpha() = error("Not implemented")
get_kalphas() = error("Not implemented")

@lazymod LinearIsotropicMod "linear_isotropic.jl"


"""
	grdproject(cmd0::String="", arg1=nothing, kwargs...)

Project a geographical gridded data set onto a rectangular grid or do the inverse projection.

Full option list at [`grdproject`]($(GMTdoc)grdproject.html)

Parameters
----------

- $(GMT.opt_J)
- **C** | **center** :: [Type => Str | []]      ``Arg = [dx/dy]``

    Let projected coordinates be relative to projection center [Default is relative to lower left corner].
    ($(GMTdoc)grdproject.html#c)
- **D** | **inc** :: [Type => Str | number]     ``Arg = xinc[unit][+e|n][/yinc[unit][+e|n]]``

    Set the grid spacing for the new grid. Append m for arc minute, s for arc second.
    ($(GMTdoc)grdproject.html#d)
- **E** | **dpi** :: [Type => Number]

    Set the resolution for the new grid in dots per inch.
    ($(GMTdoc)grdproject.html#e)
- **F** | **one2one** :: [Type => Str]           ``Arg = [c|i|p|e|f|k|M|n|u]``

    Force 1:1 scaling, i.e., output (or input, see -I) data are in actual projected meters [e].
    ($(GMTdoc)grdproject.html#f)
- **G** | **outgrid** :: [Type => Str]

    Output grid file name. Note that this is optional and to be used only when saving
    the result directly on disk. Otherwise, just use the G = grdproject(....) form.
    ($(GMTdoc)grdproject.html#g)
- **I** | **inverse** :: [Type => Bool]

    Do the Inverse transformation, from rectangular to geographical.
    ($(GMTdoc)grdproject.html#i)
- **M** | **projected_unit** :: [Type => Str]    ``Arg = c|i|p``

    Append c, i, or p to indicate that cm, inch, or point should be the projected measure unit.
    ($(GMTdoc)grdproject.html#m)
- $(GMT.opt_R)
- $(GMT.opt_V)
- $(GMT.opt_n)
- $(GMT.opt_r)
"""
function grdproject(cmd0::String="", arg1=nothing; kwargs...)

	length(kwargs) == 0 && return monolitic("grdproject", cmd0, arg1)

	d = KW(kwargs)
	cmd = parse_common_opts(d, "", [:R :J :V_params :n :r])
	cmd = parse_these_opts(cmd, d, [[:C :center], [:D :inc], [:E :dpi], [:F :one2one],
	                                [:G :outgrid], [:I :inverse], [:M :projected_unit]])

	common_grd(d, cmd0, cmd, "grdproject ", arg1)		# Finish build cmd and run it
end

# ---------------------------------------------------------------------------------------------------
grdproject(arg1, cmd0::String=""; kw...) = grdproject(cmd0, arg1; kw...)


using Plots
using Plots.PlotMeasures
using Colors
using StatsPlots

include("collect_data.jl")

function plot_D()

    mcolor = ARGB(186/255,12/255,47/255, 0.8)
    pcolor = ARGB(39/255, 56/255, 150/255, 0.75)
    ocolor = :darkolivegreen

    mkl = get_mkl()
    oct = get_octavian()
    psi4  = get_psi4()

    nao = [144, 116, 236, 264, 592, 574, 724, 172, 232, 350, 528, 472, 592, 632, 724, 204, 322, 336, 338, 528, 632, 588]
    perm = sortperm(nao)
    nao = nao[perm]

    p1 = groupedbar([oct mkl psi4], c = [ocolor mcolor pcolor], legend=:topright,
                label=["Fermi - Octavian" "Fermi - MKL" "Psi4"], dpi=1200, orientation = :v,
                bar_width=0.8,xticks=([2, 6, 12, 18, 22],["S2","S6", "S12","S18","S22"]), framestyle=:box)

    ylabel!("Time (seconds)")
    xlabel!("Dimer")
    p1
    p2 = scatter(nao, mkl[perm] ./ psi4[perm], label="", yticks=[0.5, 1.0, 1.5], ylims=[0.3,1.8],
                color=:black, ylabel="Relative time\n(Fermi/Psi4)", framestyle=:box, marker=:diamond, markercolor=mcolor,
                xlabel="Number of basis functions", legend=:bottomleft)

    scatter!(p2, nao, oct[perm] ./ psi4[perm], label="", markercolor = ocolor, marker=:diamond)
    hline!(p2, [1], color=:gray, linestyle=:dot, label="")

    plot(p1, p2, layout = grid(2,1, heights=[0.7, 0.3]), leftmargin=0.8cm)
end

function plot_E()

    mcolor = ARGB(186/255,12/255,47/255, 0.8)
    pcolor = ARGB(39/255, 56/255, 150/255, 0.75)
    ocolor = :palegreen4

    mkl   = 1000 .* raw_average(:fermi, "MKL")
    oct   = 1000 .* raw_average(:fermi, "Octavian")
    psi4  = 1000 .* raw_average(:psi4)

    nao = [144, 116, 236, 264, 592, 574, 724, 172, 232, 350, 528, 472, 592, 632, 724, 204, 322, 336, 338, 528, 632, 588]
    perm = sortperm(nao)
    nao = nao[perm]

    p1 = groupedbar([oct mkl psi4], c = [ocolor mcolor pcolor], legend=:topright,
                label=["Fermi - Octavian" "Fermi - MKL" "Psi4"], dpi=1200, orientation = :v,
                yaxis=:log,
                ylim = (10^-0.5,10^6),
                bar_width=0.8,xticks=([2, 6, 12, 18, 22],["S2","S6", "S12","S18","S22"]), framestyle=:box)

    ylabel!("Time (miliseconds)")
    xlabel!("Dimer")
    p1
    p2 = scatter(nao, mkl[perm] ./ psi4[perm], label="MKL", yticks=[0.5, 1.0, 1.5], ylims=[0.3,1.8],
                color=:black, ylabel="Relative time\n(Fermi/Psi4)", framestyle=:box, marker=:diamond, markercolor=mcolor,
                xlabel="Number of basis functions", legend=:bottomleft, legendfontsize=6, legendfonthalign=:left)

    scatter!(p2, nao, oct[perm] ./ psi4[perm], label="Octavian", markercolor = ocolor, marker=:diamond)
    hline!(p2, [1], color=:gray, linestyle=:dot, label="")

    plot(p1, p2, layout = grid(2,1, heights=[0.7, 0.3]), leftmargin=0.8cm, dpi=1200)

    savefig("dfmp2.png")
end


@testset "Insertion" begin
    t = SegmentMap{Int,Int}()
    n = 100000

    @testset "random insertions" begin
        for v in 1:n
            k = randsegment()
            t[k] = v
        end
        @test issorted(collect(keys(t)))
        @test validkeys(t)
        @test validparents(t)
        @test validsiblings(t)
    end

    @testset "ordered insertions" begin
        for v in 1:n
            k = (v, rand(v:maxend))
            t[k] = v
        end
        @test issorted(collect(keys(t)))
        @test validkeys(t)
        @test validparents(t)
        @test validsiblings(t)
    end

    @testset "reverse ordered insertions" begin
        for v in 1:n
            k = (maxend - v, rand((maxend - v):maxend))
            t[k] = v
        end
        @test issorted(collect(keys(t)))
        @test validkeys(t)
        @test validparents(t)
        @test validsiblings(t)
    end

    @testset "bulk insertion" begin
        segments = SegmentValue{Int,Int}[]
        for v in 1:n
            a, b = randsegment()
            push!(segments, SegmentValue{Int,Int}(a, b, v))
        end
        sort!(segments)
        t = SegmentMap{Int,Int}(segments)

        @test collect(keys(t)) == [Segment{Int}(segment.first, segment.last)
                                   for segment in segments]
        @test collect(values(t)) == [segment.value for segment in segments]
        @test issorted(collect(keys(t)))
        @test validkeys(t)
        @test validparents(t)
        @test validsiblings(t)

        # don't break on empty arrays
        t = SegmentMap{Int,Int}(SegmentValue{Int,Int}[])
        @test issorted(collect(keys(t)))
        @test validkeys(t)
        @test validparents(t)
        @test validsiblings(t)

        @test_throws MethodError SegmentTree{Int,Int}(segments, Int[1])
        shuffle!(segments)
        @test_throws UndefVarError SegmentTree{Int,Int}(segments, vals)
    end
end

@testset "Updates" begin
    t = SegmentMap{Int,Int}()
    n = 100000
    ks = [randsegment() for _ in 1:n]
    for (v, k) in enumerate(ks)
        t[k] = v
    end

    shuffle!(ks)
    for (v, k) in enumerate(ks)
        t[k] = v
    end

    @test issorted(collect(keys(t)))
    @test validkeys(t)
    @test validparents(t)
    @test validsiblings(t)
end

@testset "Deletion" begin
    # Set B to smaller number to try to induce more weird merges and borrows and
    # such
    t = SegmentBTree{Int,SegmentValue{Int,Int},8}()
    n = 100000

    # insert n random segments
    segments = [randsegment() for _ in 1:n]
    for (i, segment) in enumerate(segments)
        t[segment] = i
    end
    shuffle!(segments)

    # delete non-existant is a nop
    @test delete!(t, (-1, -1)).n === n

    # delete one
    segment = pop!(segments)
    delete!(t, segment)
    @test validkeys(t)
    @test validparents(t)
    @test validsiblings(t)
    @test issorted(collect(keys(t)))
    @test !haskey(t, segment)
    @test all(segment -> haskey(t, segment), segments)

    # delete a random 50%
    for segment in 1:(n - 1) / 2
        segment = pop!(segments)
        delete!(t, segment)
    end
    @test validkeys(t)
    @test validparents(t)
    @test validsiblings(t)
    @test issorted(collect(keys(t)))
    @test all(segment -> haskey(t, segment), segments)

    # delete the rest
    for segment in segments
        delete!(t, segment)
    end
    @test isempty(t)
    @test validkeys(t)
    @test validparents(t)
    @test validsiblings(t)

    # delete left-to-right
    shuffle!(segments)
    for (i, segment) in enumerate(segments)
        t[segment] = i
    end
    sort!(segments)
    for segment in segments
        delete!(t, segment)
    end
    @test isempty(t)
    @test validkeys(t)
    @test validparents(t)
    @test validsiblings(t)

    # delete right-to-left
    shuffle!(segments)
    for (i, segment) in enumerate(segments)
        t[segment] = i
    end
    reverse!(segments)
    for segment in segments
        delete!(t, segment)
    end
    @test isempty(t)
    @test validkeys(t)
    @test validparents(t)
    @test validsiblings(t)

    # delete from the middle
    shuffle!(segments)
    for (i, segment) in enumerate(segments)
        t[segment] = i
    end
    mid = div(length(segments), 2)
    for i in mid + 1:length(segments)
        delete!(t, segments[i])
    end
    for i in mid:-1:1
    delete!(t, segments[i])
    end
    @test isempty(t)
    @test validkeys(t)
    @test validparents(t)
    @test validsiblings(t)

    # contrive a special case: merge with the left when we have a non-null
    # right
    t = SegmentBTree{Int,SegmentValue{Int,Int},4}()
    for i in 1:24
        t[(i, i)] = i
    end
    @test isa(t.root, InternalNode)
    @test length(t.root.children) === 3
    @test length(t.root.children[2].children) === 2
    delete!(t, (15, 15))
    keys_to_delete = collect(keys(t))
    for key in keys_to_delete
        delete!(t, key)
    end
    @test isempty(t)
    @test validkeys(t)
    @test validparents(t)
    @test validsiblings(t)
end