"""
`init_subspec!(m::HetDSGE)`

Initializes a model subspecification by overwriting parameters from
the original model object with new parameter objects. This function is
called from within the model constructor.
"""
function init_subspec!(m::HetDSGE)
    if subspec(m) == "ss0"
        return
    # All but g_shock
    elseif subspec(m) == "ss1"
        return ss1!(m)
    # All but b_shock
    elseif subspec(m) == "ss2"
        return ss2!(m)
    # All but μ shock
    elseif subspec(m) == "ss3"
        return ss3!(m)
    # All but z shock
    elseif subspec(m) == "ss4"
        return ss4!(m)
    # All but λ_f shock
    elseif subspec(m) == "ss5"
        return ss5!(m)
    # All but λ_w shock
    elseif subspec(m) == "ss6"
        return ss6!(m)
    # All but rm_shock
    elseif subspec(m) == "ss7"
        return ss7!(m)
    # All but steady state parameters
    elseif subspec(m) == "ss8"
        return ss8!(m)
    else
        error("This subspec should be a 0")
    end
end

"""
```
ss1!(m::HetDSGE)
```

Initializes subspec 1 of `HetDSGE`.
This shuts down all shocks except the government spending process (g shock).
"""
function ss1!(m::HetDSGE)
    m <= parameter(:spp, 4., (-15., 15.), (-15., 15.), Untransformed(),
                   Normal(4., 1.5), fixed = true,
                   description = "S'': The second derivative of households' cost of adjusting investment.",
                   tex_label = "S''")
    m <= parameter(:ρR , 0.75, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.75, 0.10), fixed = true,
                   description = "ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho_R")
    m <= parameter(:ψπ , 1.5, (1e-5, 10.), (1e-5, 10.0), ModelConstructors.Exponential(),
                   Normal(1.5, 0.25), fixed = true,
                   description = "ψ1: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1")

    m <= parameter(:ρ_G, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g")

    m <= parameter(:ρ_B, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_b: AR(1) coefficient in the intertemporal preference shifter process.",
                   tex_label = "\\rho_B")
    m <= parameter(:ρ_μ, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label = "\\rho_{\\mu}")
    m <= parameter(:ρ_z, 0.0,(1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_z: AR(1) coefficient in the technology process.",
                   tex_label = "\\rho_z")
    m <= parameter(:ρ_lamf, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_f: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_f}")
    m <= parameter(:ρ_lamw, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}")
    m <= parameter(:ρ_mon, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_rm: AR(1) coefficient in the monetary policy shock process.",
                   tex_label = "\\rho_{r^m}")

    m <= parameter(:σ_b, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_b: The standard deviation of the intertemporal preference shifter process.",
                   tex_label = "\\sigma_{b}")
    m <= parameter(:σ_μ, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_μ: The standard deviation of the exogenous marginal efficiency of investment shock process.",
                   tex_label = "\\sigma_{\\mu}")
    m <= parameter(:σ_z, 0.0,(1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_z: The standard deviation of the process describing the stationary component of productivity.",
                   tex_label = "\\sigma_z")
    m <= parameter(:σ_λ_f, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   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.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_λ_w",
                   tex_label = "\\sigma_{\\lambda_w}")
    m <= parameter(:σ_rm, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_r_m: The standard deviation of the monetary policy shock.", tex_label = "\\sigma_{r^m}")

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

    m <= parameter(:Lmean, -45.9364, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed=true,
                   description="Lmean: Mean level of hours.",
                   tex_label="\\bar{L}")

    m <= parameter(:κ_p, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                 fixed = true, description = "κ_p : The slope of the Price Phillips curve",
                 tex_label = "\\kappa")

    m <= parameter(:κ_w, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = true, description = "κ_w: The slope of the Wage Phillips curve",
                   tex_label = "\\kappa")

    m <= parameter(:ψy , 0.5, (-0.5, 0.5), (-0.5, 0.5), Untransformed(),
                   Normal(0.12, 0.05), fixed = true,
                   description = "ψy: Weight on output growth in monetary policy rule")

    m <= parameter(:ϕh, 2., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   Normal(2, 0.75), fixed = true, description = "inverse frisch elasticity",
                   tex_label = "\\phi_h")
end

"""
```
ss2!(m::HetDSGE)
```

Initializes subspec 2 of `HetDSGE`.
This shuts down all shocks except the preference shifter process (b shock).
"""
function ss2!(m::HetDSGE)
    m <= parameter(:spp, 4., (-15., 15.), (-15., 15.), Untransformed(),
                   Normal(4., 1.5), fixed = true,
                   description = "S'': The second derivative of households' cost of adjusting investment.",
                   tex_label = "S''")
    m <= parameter(:ρR , 0.75, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.75, 0.10), fixed = true,
                   description = "ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho_R")
    m <= parameter(:ψπ , 1.5, (1e-5, 10.), (1e-5, 10.0), ModelConstructors.Exponential(),
                   Normal(1.5, 0.25), fixed = true,
                   description = "ψ1: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1")

    m <= parameter(:ρ_G, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g")

    m <= parameter(:ρ_B, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_b: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_b")

    m <= parameter(:ρ_μ, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label = "\\rho_{\\mu}")
    m <= parameter(:ρ_z, 0.0,(1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_z: AR(1) coefficient in the technology process.",
                   tex_label = "\\rho_z")
    m <= parameter(:ρ_lamf, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_f: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_f}")
    m <= parameter(:ρ_lamw, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}")
    m <= parameter(:ρ_mon, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_rm: AR(1) coefficient in the monetary policy shock process.",
                   tex_label = "\\rho_{r^m}")

    m <= parameter(:σ_g, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_g: The standard deviation of the government spending process.",
                   tex_label = "\\sigma_{g}")
    m <= parameter(:σ_μ, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_μ: The standard deviation of the exogenous marginal efficiency of investment shock process.",
                   tex_label = "\\sigma_{\\mu}")
    m <= parameter(:σ_z, 0.0,(1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_z: The standard deviation of the process describing the stationary component of productivity.",
                   tex_label = "\\sigma_z")
    m <= parameter(:σ_λ_f, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   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.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_λ_w",
                   tex_label = "\\sigma_{\\lambda_w}")
    m <= parameter(:σ_rm, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_r_m: The standard deviation of the monetary policy shock.", tex_label = "\\sigma_{r^m}")

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

    m <= parameter(:Lmean, -45.9364, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed=true,
                   description="Lmean: Mean level of hours.",
                   tex_label="\\bar{L}")

    m <= parameter(:κ_p, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                 fixed = true, description = "κ_p : The slope of the Price Phillips curve",
                 tex_label = "\\kappa")

    m <= parameter(:κ_w, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = true, description = "κ_w: The slope of the Wage Phillips curve",
                   tex_label = "\\kappa")

    m <= parameter(:ψy , 0.5, (-0.5, 0.5), (-0.5, 0.5), Untransformed(),
                   Normal(0.12, 0.05), fixed = true,
                   description = "ψy: Weight on output growth in monetary policy rule")

    m <= parameter(:ϕh, 2., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   Normal(2, 0.75), fixed = true, description = "inverse frisch elasticity",
                   tex_label = "\\phi_h")
end

"""
```
ss3!(m::HetDSGE)
```

Initializes subspec 3 of `HetDSGE`.
This shuts down all shocks except the capital adjustment cost process (μ shock).
"""
function ss3!(m::HetDSGE)
    m <= parameter(:spp, 4., (-15., 15.), (-15., 15.), Untransformed(),
                   Normal(4., 1.5), fixed = true,
                   description = "S'': The second derivative of households' cost of adjusting investment.",
                   tex_label = "S''")
    m <= parameter(:ρR , 0.75, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.75, 0.10), fixed = true,
                   description = "ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho_R")
    m <= parameter(:ψπ , 1.5, (1e-5, 10.), (1e-5, 10.0), ModelConstructors.Exponential(),
                   Normal(1.5, 0.25), fixed = true,
                   description = "ψ1: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1")

    m <= parameter(:ρ_G, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g")
    m <= parameter(:ρ_B, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_b: AR(1) coefficient in the intertemporal preference shifter process.",
                   tex_label = "\\rho_B")

    m <= parameter(:ρ_μ, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_μ: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_\\mu")

    m <= parameter(:ρ_z, 0.0,(1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_z: AR(1) coefficient in the technology process.",
                   tex_label = "\\rho_z")
    m <= parameter(:ρ_lamf, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_f: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_f}")
    m <= parameter(:ρ_lamw, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}")
    m <= parameter(:ρ_mon, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_rm: AR(1) coefficient in the monetary policy shock process.",
                   tex_label = "\\rho_{r^m}")

    m <= parameter(:σ_g, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_g: The standard deviation of the government spending process.",
                   tex_label = "\\sigma_{g}")
    m <= parameter(:σ_b, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_b: The standard deviation of the intertemporal preference shifter process.",
                   tex_label = "\\sigma_{b}")
    m <= parameter(:σ_z, 0.0,(1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_z: The standard deviation of the process describing the stationary component of productivity.",
                   tex_label = "\\sigma_z")
    m <= parameter(:σ_λ_f, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   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.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_λ_w",
                   tex_label = "\\sigma_{\\lambda_w}")
    m <= parameter(:σ_rm, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_r_m: The standard deviation of the monetary policy shock.", tex_label = "\\sigma_{r^m}")

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

    m <= parameter(:Lmean, -45.9364, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed=true,
                   description="Lmean: Mean level of hours.",
                   tex_label="\\bar{L}")

    m <= parameter(:κ_p, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                 fixed = true, description = "κ_p : The slope of the Price Phillips curve",
                 tex_label = "\\kappa")

    m <= parameter(:κ_w, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = true, description = "κ_w: The slope of the Wage Phillips curve",
                   tex_label = "\\kappa")

    m <= parameter(:ψy , 0.5, (-0.5, 0.5), (-0.5, 0.5), Untransformed(),
                   Normal(0.12, 0.05), fixed = true,
                   description = "ψy: Weight on output growth in monetary policy rule")

    m <= parameter(:ϕh, 2., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   Normal(2, 0.75), fixed = true, description = "inverse frisch elasticity",
                   tex_label = "\\phi_h")
end

"""
```
ss4(m::HetDSGE)
```

Initializes subspec 4 of `HetDSGE`.
This shuts down all shocks except the technology process (z shock).
"""
function ss4!(m::HetDSGE)
    m <= parameter(:spp, 4., (-15., 15.), (-15., 15.), Untransformed(),
                   Normal(4., 1.5), fixed = true,
                   description = "S'': The second derivative of households' cost of adjusting investment.",
                   tex_label = "S''")
    m <= parameter(:ρR , 0.75, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.75, 0.10), fixed = true,
                   description = "ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho_R")
    m <= parameter(:ψπ , 1.5, (1e-5, 10.), (1e-5, 10.0), ModelConstructors.Exponential(),
                   Normal(1.5, 0.25), fixed = true,
                   description = "ψ1: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1")

    m <= parameter(:ρ_G, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g")
    m <= parameter(:ρ_B, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_b: AR(1) coefficient in the intertemporal preference shifter process.",
                   tex_label = "\\rho_B")
    m <= parameter(:ρ_μ, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label = "\\rho_{\\mu}")

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

    m <= parameter(:ρ_lamf, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_f: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_f}")
    m <= parameter(:ρ_lamw, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}")
    m <= parameter(:ρ_mon, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_rm: AR(1) coefficient in the monetary policy shock process.",
                   tex_label = "\\rho_{r^m}")

    m <= parameter(:σ_g, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_g: The standard deviation of the government spending process.",
                   tex_label = "\\sigma_{g}")
    m <= parameter(:σ_b, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_b: The standard deviation of the intertemporal preference shifter process.",
                   tex_label = "\\sigma_{b}")
    m <= parameter(:σ_μ, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_μ: The standard deviation of the exogenous marginal efficiency of investment shock process.",
                   tex_label = "\\sigma_{\\mu}")
    m <= parameter(:σ_λ_f, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   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.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_λ_w",
                   tex_label = "\\sigma_{\\lambda_w}")
    m <= parameter(:σ_rm, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_r_m: The standard deviation of the monetary policy shock.", tex_label = "\\sigma_{r^m}")

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

    m <= parameter(:Lmean, -45.9364, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed=true,
                   description="Lmean: Mean level of hours.",
                   tex_label="\\bar{L}")

    m <= parameter(:κ_p, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                 fixed = true, description = "κ_p : The slope of the Price Phillips curve",
                 tex_label = "\\kappa")

    m <= parameter(:κ_w, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = true, description = "κ_w: The slope of the Wage Phillips curve",
                   tex_label = "\\kappa")

    m <= parameter(:ψy , 0.5, (-0.5, 0.5), (-0.5, 0.5), Untransformed(),
                   Normal(0.12, 0.05), fixed = true,
                   description = "ψy: Weight on output growth in monetary policy rule")

    m <= parameter(:ϕh, 2., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   Normal(2, 0.75), fixed = true, description = "inverse frisch elasticity",
                   tex_label = "\\phi_h")

end

"""
```
ss5!(m::HetDSGE)
```

Initializes subspec 5 of `HetDSGE`.
This shuts down all shocks except the price mark-up shock process (λ_f shock).
"""
function ss5!(m::HetDSGE)
    m <= parameter(:spp, 4., (-15., 15.), (-15., 15.), Untransformed(),
                   Normal(4., 1.5), fixed = true,
                   description = "S'': The second derivative of households' cost of adjusting investment.",
                   tex_label = "S''")
    m <= parameter(:ρR , 0.75, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.75, 0.10), fixed = true,
                   description = "ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho_R")
    m <= parameter(:ψπ , 1.5, (1e-5, 10.), (1e-5, 10.0), ModelConstructors.Exponential(),
                   Normal(1.5, 0.25), fixed = true,
                   description = "ψ1: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1")

    m <= parameter(:ρ_G, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g")
    m <= parameter(:ρ_B, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_b: AR(1) coefficient in the intertemporal preference shifter process.",
                   tex_label = "\\rho_B")
    m <= parameter(:ρ_μ, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label = "\\rho_{\\mu}")
    m <= parameter(:ρ_z, 0.0,(1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_z: AR(1) coefficient in the technology process.",
                   tex_label = "\\rho_z")

    m <= parameter(:ρ_lamf, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_lamf: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_{\\lambda_f}")

    m <= parameter(:ρ_lamw, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}")
    m <= parameter(:ρ_mon, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_rm: AR(1) coefficient in the monetary policy shock process.",
                   tex_label = "\\rho_{r^m}")

    m <= parameter(:σ_g, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_g: The standard deviation of the government spending process.",
                   tex_label = "\\sigma_{g}")
    m <= parameter(:σ_b, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_b: The standard deviation of the intertemporal preference shifter process.",
                   tex_label = "\\sigma_{b}")
    m <= parameter(:σ_μ, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_μ: The standard deviation of the exogenous marginal efficiency of investment shock process.",
                   tex_label = "\\sigma_{\\mu}")
    m <= parameter(:σ_z, 0.0,(1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_z: The standard deviation of the process describing the stationary component of productivity.",
                   tex_label = "\\sigma_z")
    m <= parameter(:σ_λ_w, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_λ_w",
                   tex_label = "\\sigma_{\\lambda_w}")
    m <= parameter(:σ_rm, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_r_m: The standard deviation of the monetary policy shock.", tex_label = "\\sigma_{r^m}")

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

    m <= parameter(:Lmean, -45.9364, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed=true,
                   description="Lmean: Mean level of hours.",
                   tex_label="\\bar{L}")

    m <= parameter(:κ_p, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                 fixed = true, description = "κ_p : The slope of the Price Phillips curve",
                 tex_label = "\\kappa")

    m <= parameter(:κ_w, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = true, description = "κ_w: The slope of the Wage Phillips curve",
                   tex_label = "\\kappa")

    m <= parameter(:ψy , 0.5, (-0.5, 0.5), (-0.5, 0.5), Untransformed(),
                   Normal(0.12, 0.05), fixed = true,
                   description = "ψy: Weight on output growth in monetary policy rule")

    m <= parameter(:ϕh, 2., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   Normal(2, 0.75), fixed = true, description = "inverse frisch elasticity",
                   tex_label = "\\phi_h")

end

"""
```
ss6!(m::HetDSGE)
```

Initializes subspec 6 of `HetDSGE`.
This shuts down all shocks except the wage mark-up shock process (λ_w shock).
"""
function ss6!(m::HetDSGE)
    m <= parameter(:spp, 4., (-15., 15.), (-15., 15.), Untransformed(),
                   Normal(4., 1.5), fixed = true,
                   description = "S'': The second derivative of households' cost of adjusting investment.",
                   tex_label = "S''")
    m <= parameter(:ρR , 0.75, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.75, 0.10), fixed = true,
                   description = "ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho_R")
    m <= parameter(:ψπ , 1.5, (1e-5, 10.), (1e-5, 10.0), ModelConstructors.Exponential(),
                   Normal(1.5, 0.25), fixed = true,
                   description = "ψ1: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1")

    m <= parameter(:ρ_G, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g")
    m <= parameter(:ρ_B, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_b: AR(1) coefficient in the intertemporal preference shifter process.",
                   tex_label = "\\rho_B")
    m <= parameter(:ρ_μ, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label = "\\rho_{\\mu}")
    m <= parameter(:ρ_z, 0.0,(1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_z: AR(1) coefficient in the technology process.",
                   tex_label = "\\rho_z")
    m <= parameter(:ρ_lamf, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_f: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_f}")
    m <= parameter(:ρ_lamw, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_w: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}")

    m <= parameter(:ρ_mon, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_rm: AR(1) coefficient in the monetary policy shock process.",
                   tex_label = "\\rho_{r^m}")

    m <= parameter(:σ_g, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_g: The standard deviation of the government spending process.",
                   tex_label = "\\sigma_{g}")
    m <= parameter(:σ_b, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_b: The standard deviation of the intertemporal preference shifter process.",
                   tex_label = "\\sigma_{b}")
    m <= parameter(:σ_μ, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_μ: The standard deviation of the exogenous marginal efficiency of investment shock process.",
                   tex_label = "\\sigma_{\\mu}")
    m <= parameter(:σ_z, 0.0,(1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_z: The standard deviation of the process describing the stationary component of productivity.",
                   tex_label = "\\sigma_z")
    m <= parameter(:σ_λ_f, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   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(:σ_rm, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_r_m: The standard deviation of the monetary policy shock.", tex_label = "\\sigma_{r^m}")

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

    m <= parameter(:Lmean, -45.9364, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed=true,
                   description="Lmean: Mean level of hours.",
                   tex_label="\\bar{L}")

    m <= parameter(:κ_p, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                 fixed = true, description = "κ_p : The slope of the Price Phillips curve",
                 tex_label = "\\kappa")

    m <= parameter(:κ_w, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = true, description = "κ_w: The slope of the Wage Phillips curve",
                   tex_label = "\\kappa")

    m <= parameter(:ψy , 0.5, (-0.5, 0.5), (-0.5, 0.5), Untransformed(),
                   Normal(0.12, 0.05), fixed = true,
                   description = "ψy: Weight on output growth in monetary policy rule")

    m <= parameter(:ϕh, 2., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   Normal(2, 0.75), fixed = true, description = "inverse frisch elasticity",
                   tex_label = "\\phi_h")

end

"""
```
ss7!(m::HetDSGE)
```

Initializes subspec 7 of `HetDSGE`.
This shuts down all shocks except the monetary policy shock process (rm shock).
"""
function ss7!(m::HetDSGE)
    m <= parameter(:spp, 4., (-15., 15.), (-15., 15.), Untransformed(),
                   Normal(4., 1.5), fixed = true,
                   description = "S'': The second derivative of households' cost of adjusting investment.",
                   tex_label = "S''")
    m <= parameter(:ρR , 0.75, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(),
                   BetaAlt(0.75, 0.10), fixed = true,
                   description = "ρ: The degree of inertia in the monetary policy rule.",
                   tex_label="\\rho_R")
    m <= parameter(:ψπ , 1.5, (1e-5, 10.), (1e-5, 10.0), ModelConstructors.Exponential(),
                   Normal(1.5, 0.25), fixed = true,
                   description = "ψ1: Weight on inflation gap in monetary policy rule.",
                   tex_label = "\\psi_1")

    m <= parameter(:ρ_G, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_g: AR(1) coefficient in the government spending process.",
                   tex_label = "\\rho_g")
    m <= parameter(:ρ_B, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_b: AR(1) coefficient in the intertemporal preference shifter process.",
                   tex_label = "\\rho_B")
    m <= parameter(:ρ_μ, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_μ: AR(1) coefficient in capital adjustment cost process.",
                   tex_label = "\\rho_{\\mu}")
    m <= parameter(:ρ_z, 0.0,(1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_z: AR(1) coefficient in the technology process.",
                   tex_label = "\\rho_z")
    m <= parameter(:ρ_lamf, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_f: AR(1) coefficient in the price mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_f}")
    m <= parameter(:ρ_lamw, 0.0, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_λ_w: AR(1) coefficient in the wage mark-up shock process.",
                   tex_label = "\\rho_{\\lambda_w}")

    m <= parameter(:ρ_mon, 0.5, (1e-5, 1 - 1e-5), (1e-5, 1-1e-5), SquareRoot(),
                   BetaAlt(0.5, 0.2), fixed = true,
                   description = "ρ_rm: AR(1) coefficient in the monetary policy shock process.",
                   tex_label = "\\rho_{r^m}")

    m <= parameter(:σ_g, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_g: The standard deviation of the government spending process.",
                   tex_label = "\\sigma_{g}")
    m <= parameter(:σ_b, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_b: The standard deviation of the intertemporal preference shifter process.",
                   tex_label = "\\sigma_{b}")
    m <= parameter(:σ_μ, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_μ: The standard deviation of the exogenous marginal efficiency of investment shock process.",
                   tex_label = "\\sigma_{\\mu}")
    m <= parameter(:σ_z, 0.0,(1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_z: The standard deviation of the process describing the stationary component of productivity.",
                   tex_label = "\\sigma_z")
    m <= parameter(:σ_λ_f, 0.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   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.0, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(),
                   RootInverseGamma(2, 0.10), fixed = true,
                   description = "σ_λ_w",
                   tex_label = "\\sigma_{\\lambda_w}")

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

    m <= parameter(:Lmean, -45.9364, (-1000., 1000.), (-1e3, 1e3), Untransformed(),
                   Normal(-45., 5.), fixed=true,
                   description="Lmean: Mean level of hours.",
                   tex_label="\\bar{L}")
    m <= parameter(:κ_p, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                 fixed = true, description = "κ_p : The slope of the Price Phillips curve",
                 tex_label = "\\kappa")

    m <= parameter(:κ_w, 0.5, (1e-5, 5.), (1e-5, 5.), SquareRoot(), GammaAlt(0.5, 0.3),
                   fixed = true, description = "κ_w: The slope of the Wage Phillips curve",
                   tex_label = "\\kappa")

    m <= parameter(:ψy , 0.5, (-0.5, 0.5), (-0.5, 0.5), Untransformed(),
                   Normal(0.12, 0.05), fixed = true,
                   description = "ψy: Weight on output growth in monetary policy rule")

    m <= parameter(:ϕh, 2., (1e-5, 10.), (1e-5, 10.), ModelConstructors.Exponential(),
                   Normal(2, 0.75), fixed = true, description = "inverse frisch elasticity",
                   tex_label = "\\phi_h")
end

"""
```
ss8!(m::HetDSGE)
```

Initializes subspec 8 of `HetDSGE`.
Estimates all non-steadystate parameters.
Right now it's empty becaus we do this by default
"""
function ss8!(m::HetDSGE)

end


@testset "Vax Float F" begin
    f4_vax = [  0x00004080,
                0x0000C080,
                0x00004160,
                0x0000C160,
                0x0FD04149,
                0x0FD0C149,
                0xBDC27DF0,
                0xBDC2FDF0,
                0x1CEA0308,
                0x1CEA8308,
                0x0652409E,
                0x0652C09E ]

    f4_ieee = Array{Float32}([  1.000000,
                               -1.000000,
                                3.500000,
                               -3.500000,
                                3.141590,
                               -3.141590,
                                9.9999999E+36,
                               -9.9999999E+36,
                                9.9999999E-38,
                               -9.9999999E-38,
                                1.23456789,
                               -1.23456789 ])

    @testset "Conversion..." begin
        for (vax, ieee) in zip(f4_vax, f4_ieee)
            @test VaxFloatF(vax) == VaxFloatF(ieee)
            @test convert(Float32, VaxFloatF(vax)) == ieee
        end
    end

    @testset "Promotion..." begin
        for t in [subtypes(VaxInt); Int8; Int16; Int32; Int64; Int128; Float16; Float32; VaxFloatF]
            @test isa(one(t)*VaxFloatF(1), Float32)
        end

        @test isa(one(Float64)*VaxFloatF(1), Float64)
        @test isa(one(BigInt)*VaxFloatF(1), BigFloat)
        @test isa(one(BigFloat)*VaxFloatF(1), BigFloat)

        un = one(VaxFloatF)
        @test promote(un, un, un) == (1.0, 1.0, 1.0)
        @test promote(un, un, un, un) == (1.0, 1.0, 1.0, 1.0)
    end

    @testset "Number definitions" begin
        @test floatmax(VaxFloatF) == typemax(VaxFloatF)
        @test floatmin(VaxFloatF) == typemin(VaxFloatF)

        @test zero(VaxFloatF) == 0
        @test one(VaxFloatF) == 1
    end

    @testset "Edge cases" begin
        # Reserved Vax floating point operand
        @test_throws InexactError convert(Float32, VaxFloatF(UInt32(0x8000)))

        # Inf and NaN should error too
        @test_throws InexactError VaxFloatF(Inf32)
        @test_throws InexactError VaxFloatF(-Inf32)
        @test_throws InexactError VaxFloatF(NaN32)

        # Both IEEE zeros should be converted to Vax true zero
        @test VaxFloatF(-0.0f0) === VaxFloatF(0.0f0) === zero(VaxFloatF)

        # Dirty zero
        @test convert(Float32, VaxFloatF(UInt32(0x40))) === zero(Float32)

        # Numbers smaller than floatmin(VaxFloatF) should underflow
        @test VaxFloatF(prevfloat(convert(Float32, floatmin(VaxFloatF)))) === zero(VaxFloatF)
        @test VaxFloatF(convert(Float32, floatmin(VaxFloatF))) === floatmin(VaxFloatF)

        # Numbers larger than floatmax(VaxFloatF) should error
        @test_throws InexactError VaxFloatF(nextfloat(convert(Float32, floatmax(VaxFloatF))))
        @test VaxFloatF(convert(Float32, floatmax(VaxFloatF))) === floatmax(VaxFloatF)
    end

    @testset "IO" begin
        io = IOBuffer(reinterpret(UInt8, ones(VaxFloatF, 4)))
        @test read(io, VaxFloatF) === one(VaxFloatF)
        @test read!(io, Vector{VaxFloatF}(undef, 3)) == ones(VaxFloatF, 3)
    end
end


using DataArrays
using Base: @get!, promote_eltype
using Base.Broadcast: bitcache_chunks, bitcache_size, dumpbitcache
using Compat: promote_eltype_op

if isdefined(Base, :OneTo)
    if VERSION < v"0.6.0-dev.1121"
        _broadcast_shape(x...) = Base.to_shape(Base.Broadcast.broadcast_shape(x...))
    else
        _broadcast_shape(x...) = Base.to_shape(Base.Broadcast.broadcast_indices(x...))
    end
else
    const _broadcast_shape = Base.Broadcast.broadcast_shape
end

# Check that all arguments are broadcast compatible with shape
# Differs from Base in that we check for exact matches
function check_broadcast_shape(shape::Dims, As::(@compat Union{AbstractArray,Number})...)
    samesize = true
    for A in As
        if ndims(A) > length(shape)
            throw(DimensionMismatch("cannot broadcast array to have fewer dimensions"))
        end
        for k in 1:length(shape)
            n, nA = shape[k], size(A, k)
            samesize &= (n == nA)
            if n != nA != 1
                throw(DimensionMismatch("array could not be broadcast to match destination"))
            end
        end
    end
    samesize
end

# Get ref for value for a PooledDataArray, adding to the pool if
# necessary
_unsafe_pdaref!(Bpool, Brefdict::Dict, val::NAtype) = 0
function _unsafe_pdaref!(Bpool, Brefdict::Dict, val)
    @get! Brefdict val begin
        push!(Bpool, val)
        length(Bpool)
    end
end

# Generate a branch for each possible combination of NA/not NA. This
# gives good performance at the cost of 2^narrays branches.
function gen_na_conds(f, nd, arrtype, outtype, daidx=find([arrtype...] .!= AbstractArray), pos=1, isna=())
    if pos > length(daidx)
        args = Any[Symbol("v_$(k)") for k = 1:length(arrtype)]
        for i = 1:length(daidx)
            if isna[i]
                args[daidx[i]] = NA
            end
        end

        # Needs to be gensymmed so that the compiler won't box it
        val = gensym("val")
        quote
            $val = $(Expr(:call, f, args...))
            $(if outtype == DataArray
                :(@inbounds unsafe_dasetindex!(Bdata, Bc, $val, ind))
            elseif outtype == PooledDataArray
                :(@inbounds (@nref $nd Brefs i) = _unsafe_pdaref!(Bpool, Brefdict, $val))
            end)
        end
    else
        k = daidx[pos]
        quote
            if $(Symbol("isna_$(k)"))
                $(gen_na_conds(f, nd, arrtype, outtype, daidx, pos+1, tuple(isna..., true)))
            else
                $(if arrtype[k] == DataArray
                    :(@inbounds $(Symbol("v_$(k)")) = $(Symbol("data_$(k)"))[$(Symbol("state_$(k)_0"))])
                else
                    :(@inbounds $(Symbol("v_$(k)")) = $(Symbol("pool_$(k)"))[$(Symbol("r_$(k)"))])
                end)
                $(gen_na_conds(f, nd, arrtype, outtype, daidx, pos+1, tuple(isna..., false)))
            end
        end
    end
end

# Broadcast implementation for DataArrays
#
# TODO: Fall back on faster implementation for same-sized inputs when
# it is safe to do so.
function gen_broadcast_dataarray(nd::Int, arrtype::@compat(Tuple{Vararg{DataType}}), outtype, f::Function)
    F = Expr(:quote, f)
    narrays = length(arrtype)
    As = [Symbol("A_$(i)") for i = 1:narrays]
    dataarrays = find([arrtype...] .== DataArray)
    abstractdataarrays = find([arrtype...] .!= AbstractArray)
    have_fastpath = outtype == DataArray && all(x->!(x <: PooledDataArray), arrtype)

    @eval let
        function _F_(B::$(outtype), $(As...))
            @assert ndims(B) == $nd

            # Set up input DataArray/PooledDataArrays
            $(Expr(:block, [
                arrtype[k] == DataArray ? quote
                    $(Symbol("na_$(k)")) = $(Symbol("A_$(k)")).na.chunks
                    $(Symbol("data_$(k)")) = $(Symbol("A_$(k)")).data
                    $(Symbol("state_$(k)_0")) = $(Symbol("state_$(k)_$(nd)")) = 1
                    @nexprs $nd d->($(Symbol("skip_$(k)_d")) = size($(Symbol("data_$(k)")), d) == 1)
                end : arrtype[k] == PooledDataArray ? quote
                    $(Symbol("refs_$(k)")) = $(Symbol("A_$(k)")).refs
                    $(Symbol("pool_$(k)")) = $(Symbol("A_$(k)")).pool
                end : nothing
            for k = 1:narrays]...))

            # Set up output DataArray/PooledDataArray
            $(if outtype == DataArray
                quote
                    Bdata = B.data
                    # Copy in case aliased
                    # TODO: check for aliasing?
                    Bna = falses(size(Bdata))
                    Bc = Bna.chunks
                    ind = 1
                end
            elseif outtype == PooledDataArray
                quote
                    Bpool = B.pool = similar(B.pool, 0)
                    Brefs = B.refs
                    Brefdict = Dict{eltype(Bpool),eltype(Brefs)}()
                end
            end)

            @nloops($nd, i, $(outtype == DataArray ? (:Bdata) : (:Brefs)),
                # pre
                d->($(Expr(:block, [
                    arrtype[k] == DataArray ? quote
                        $(Symbol("state_$(k)_")){d-1} = $(Symbol("state_$(k)_d"));
                        $(Symbol("j_$(k)_d")) = $(Symbol("skip_$(k)_d")) ? 1 : i_d
                    end : quote
                        $(Symbol("j_$(k)_d")) = size($(Symbol("A_$(k)")), d) == 1 ? 1 : i_d
                    end
                for k = 1:narrays]...))),

                # post
                d->($(Expr(:block, [quote
                    $(Symbol("skip_$(k)_d")) || ($(Symbol("state_$(k)_d")) = $(Symbol("state_$(k)_0")))
                end for k in dataarrays]...))),

                # body
                begin
                    # Advance iterators for DataArray and determine NA status
                    $(Expr(:block, [
                        arrtype[k] == DataArray ? quote
                            @inbounds $(Symbol("isna_$(k)")) = Base.unsafe_bitgetindex($(Symbol("na_$(k)")), $(Symbol("state_$(k)_0")))
                        end : arrtype[k] == PooledDataArray ? quote
                            @inbounds $(Symbol("r_$(k)")) = @nref $nd $(Symbol("refs_$(k)")) d->$(Symbol("j_$(k)_d"))
                            $(Symbol("isna_$(k)")) = $(Symbol("r_$(k)")) == 0
                        end : nothing
                    for k = 1:narrays]...))

                    # Extract values for ordinary AbstractArrays
                    $(Expr(:block, [
                        :(@inbounds $(Symbol("v_$(k)")) = @nref $nd $(Symbol("A_$(k)")) d->$(Symbol("j_$(k)_d")))
                    for k = find([arrtype...] .== AbstractArray)]...))

                    # Compute and store return value
                    $(gen_na_conds(F, nd, arrtype, outtype))

                    # Increment state
                    $(Expr(:block, [:($(Symbol("state_$(k)_0")) += 1) for k in dataarrays]...))
                    $(if outtype == DataArray
                        :(ind += 1)
                    end)
                end)

            $(if outtype == DataArray
                :(B.na = Bna)
            end)
        end
        _F_
    end
end

datype(A_1::PooledDataArray, As...) = tuple(PooledDataArray, datype(As...)...)
datype(A_1::DataArray, As...) = tuple(DataArray, datype(As...)...)
datype(A_1, As...) = tuple(AbstractArray, datype(As...)...)
datype() = ()

datype_int(A_1::PooledDataArray, As...) = (@compat(UInt64(2)) | (datype_int(As...) << 2))
datype_int(A_1::DataArray, As...) = (@compat(UInt64(1)) | (datype_int(As...) << 2))
datype_int(A_1, As...) = (datype_int(As...) << 2)
datype_int() = @compat UInt64(0)

# The following four methods are to avoid ambiguity warnings on 0.4
Base.map!(f::Base.Callable, B::DataArray) =
    invoke(map!, Tuple{Base.Callable, AbstractArray}, f, B)
Base.map!(f::Base.Callable, B::PooledDataArray) =
    invoke(map!, Tuple{Base.Callable, AbstractArray}, f, B)
Base.broadcast!(f::Base.Function, B::DataArray) =
    invoke(map!, Tuple{Base.Callable, AbstractArray}, f, B)
Base.broadcast!(f::Base.Function, B::PooledDataArray) =
    invoke(map!, Tuple{Base.Callable, AbstractArray}, f, B)

for bsig in (DataArray, PooledDataArray), asig in (Union{Array,BitArray,Number},DataArray, PooledDataArray,)
    @eval let cache = Dict{Function,Dict{UInt64,Dict{Int,Function}}}()
        function Base.map!(f::Base.Callable, B::$bsig, As::$asig...)
            nd = ndims(B)
            length(As) <= 8 || throw(ArgumentError("too many arguments"))
            samesize = check_broadcast_shape(size(B), As...)
            samesize || throw(DimensionMismatch("Argument dimensions must match"))
            arrtype = datype_int(As...)

            cache_f    = @get! cache      f        Dict{UInt64,Dict{Int,Function}}()
            cache_f_na = @get! cache_f    arrtype  Dict{Int,Function}()
            func       = @get! cache_f_na nd       gen_broadcast_dataarray(nd, datype(As...), $bsig, f)

            func(B, As...)
            B
        end

        function Base.broadcast!(f::Function, B::$bsig, As::$asig...)
            nd = ndims(B)
            length(As) <= 8 || throw(ArgumentError("too many arguments"))
            samesize = check_broadcast_shape(size(B), As...)
            arrtype = datype_int(As...)

            cache_f    = @get! cache      f        Dict{UInt64,Dict{Int,Function}}()
            cache_f_na = @get! cache_f    arrtype  Dict{Int,Function}()
            func       = @get! cache_f_na nd       gen_broadcast_dataarray(nd, datype(As...), $bsig, f)

            # println(code_typed(func, typeof(tuple(B, As...))))
            func(B, As...)
            B
        end
    end
end

databroadcast(f::Function, As...) = broadcast!(f, DataArray(promote_eltype(As...), _broadcast_shape(As...)), As...)
pdabroadcast(f::Function, As...) = broadcast!(f, PooledDataArray(promote_eltype(As...), _broadcast_shape(As...)), As...)

function exreplace!(ex::Expr, search, rep)
    for i = 1:length(ex.args)
        if ex.args[i] == search
            splice!(ex.args, i, rep)
            break
        else
            exreplace!(ex.args[i], search, rep)
        end
    end
    ex
end
exreplace!(ex, search, rep) = ex

macro da_broadcast_vararg(func)
    if (func.head != :function && func.head != :(=)) ||
       func.args[1].head != :call || !isa(func.args[1].args[end], Expr) ||
       func.args[1].args[end].head != :...
        throw(ArgumentError("@da_broadcast_vararg may only be applied to vararg functions"))
    end

    va = func.args[1].args[end]
    defs = Any[]
    for n = 1:4, aa = 0:n-1
        def = deepcopy(func)
        rep = Any[Symbol("A_$(i)") for i = 1:n]
        push!(rep, va)
        exreplace!(def.args[2], va, rep)
        rep = Vector{Any}(n+1)
        for i = 1:aa
            rep[i] = Expr(:(::), Symbol("A_$i"), AbstractArray)
        end
        for i = aa+1:n
            rep[i] = Expr(:(::), Symbol("A_$i"), (@compat Union{DataArray, PooledDataArray}))
        end
        rep[end] = Expr(:..., Expr(:(::), va.args[1], AbstractArray))
        exreplace!(def.args[1], va, rep)
        push!(defs, def)
    end
    esc(Expr(:block, defs...))
end

macro da_broadcast_binary(func)
    if (func.head != :function && func.head != :(=)) ||
       func.args[1].head != :call ||
       length(func.args[1].args) != 3
        throw(ArgumentError("@da_broadcast_binary may only be applied to two-argument functions"))
    end
    (ff, A, B) = func.args[1].args
    f = esc(ff)
    body = func.args[2]
    quote
        ($f)($A::(@compat Union{DataArray, PooledDataArray}), $B::(@compat Union{DataArray, PooledDataArray})) = $(body)
        ($f)($A::(@compat Union{DataArray, PooledDataArray}), $B::AbstractArray) = $(body)
        ($f)($A::AbstractArray, $B::(@compat Union{DataArray, PooledDataArray})) = $(body)
    end
end

# Broadcasting DataArrays returns a DataArray
@da_broadcast_vararg Base.broadcast(f::Function, As...) = databroadcast(f, As...)

# Definitions for operators,
(.*)(A::BitArray, B::(@compat Union{DataArray{Bool}, PooledDataArray{Bool}})) = databroadcast(*, A, B)
(.*)(A::(@compat Union{DataArray{Bool}, PooledDataArray{Bool}}), B::BitArray) = databroadcast(*, A, B)
@da_broadcast_vararg (.*)(As...) = databroadcast(*, As...)
@da_broadcast_binary (.%)(A, B) = databroadcast(%, A, B)
@da_broadcast_vararg (.+)(As...) = broadcast!(+, DataArray(promote_eltype_op(@functorize(+), As...), _broadcast_shape(As...)), As...)
@da_broadcast_binary (.-)(A, B) =
    broadcast!(-, DataArray(promote_op(@functorize(-), eltype(A), eltype(B)),
                            _broadcast_shape(A,B)), A, B)
@da_broadcast_binary (./)(A, B) =
    broadcast!(/, DataArray(promote_op(@functorize(/), eltype(A), eltype(B)),
                            _broadcast_shape(A, B)), A, B)
@da_broadcast_binary (.\)(A, B) =
    broadcast!(\, DataArray(promote_op(@functorize(\), eltype(A), eltype(B)),
                            _broadcast_shape(A, B)), A, B)
(.^)(A::(@compat Union{DataArray{Bool}, PooledDataArray{Bool}}), B::(@compat Union{DataArray{Bool}, PooledDataArray{Bool}})) = databroadcast(>=, A, B)
(.^)(A::BitArray, B::(@compat Union{DataArray{Bool}, PooledDataArray{Bool}})) = databroadcast(>=, A, B)
(.^)(A::AbstractArray{Bool}, B::(@compat Union{DataArray{Bool}, PooledDataArray{Bool}})) = databroadcast(>=, A, B)
(.^)(A::(@compat Union{DataArray{Bool}, PooledDataArray{Bool}}), B::BitArray) = databroadcast(>=, A, B)
(.^)(A::(@compat Union{DataArray{Bool}, PooledDataArray{Bool}}), B::AbstractArray{Bool}) = databroadcast(>=, A, B)
@da_broadcast_binary (.^)(A, B) =
    broadcast!(^, DataArray(promote_op(@functorize(^), eltype(A), eltype(B)),
                            _broadcast_shape(A, B)), A, B)

# XXX is a PDA the right return type for these?
Base.broadcast(f::Function, As::PooledDataArray...) = pdabroadcast(f, As...)
(.*)(As::PooledDataArray...) = pdabroadcast(*, As...)
(.%)(A::PooledDataArray, B::PooledDataArray) = pdabroadcast(%, A, B)
(.+)(As::PooledDataArray...) =
    broadcast!(+, PooledDataArray(promote_eltype_op(@functorize(+), As...), _broadcast_shape(As...)), As...)
(.-)(A::PooledDataArray, B::PooledDataArray) =
    broadcast!(-, PooledDataArray(promote_op(@functorize(-), eltype(A), eltype(B)),
                                  _broadcast_shape(A,B)), A, B)
(./)(A::PooledDataArray, B::PooledDataArray) =
    broadcast!(/, PooledDataArray(promote_op(@functorize(/), eltype(A), eltype(B)),
                                  _broadcast_shape(A, B)), A, B)
(.\)(A::PooledDataArray, B::PooledDataArray) =
    broadcast!(\, PooledDataArray(promote_op(@functorize(\), eltype(A), eltype(B)),
                                  _broadcast_shape(A, B)), A, B)
(.^)(A::PooledDataArray{Bool}, B::PooledDataArray{Bool}) = databroadcast(>=, A, B)
(.^)(A::PooledDataArray, B::PooledDataArray) =
    broadcast!(^, PooledDataArray(promote_op(@functorize(^), eltype(A), eltype(B)),
                                  _broadcast_shape(A, B)), A, B)

for (sf, vf) in zip(scalar_comparison_operators, array_comparison_operators)
    @eval begin
        # ambiguity
        $(vf)(A::(@compat Union{PooledDataArray{Bool},DataArray{Bool}}), B::(@compat Union{PooledDataArray{Bool},DataArray{Bool}})) =
            broadcast!($sf, DataArray(Bool, _broadcast_shape(A, B)), A, B)
        $(vf)(A::(@compat Union{PooledDataArray{Bool},DataArray{Bool}}), B::AbstractArray{Bool}) =
            broadcast!($sf, DataArray(Bool, _broadcast_shape(A, B)), A, B)
        $(vf)(A::AbstractArray{Bool}, B::(@compat Union{PooledDataArray{Bool},DataArray{Bool}})) =
            broadcast!($sf, DataArray(Bool, _broadcast_shape(A, B)), A, B)

        @da_broadcast_binary $(vf)(A, B) = broadcast!($sf, DataArray(Bool, _broadcast_shape(A, B)), A, B)
    end
end


# TODO(kykim): Try a lock-free approach using Atomics.
mutable struct PISActionNode{S,A}
    id::Int
    a_label::A
    n::Int
    q::Float64
    transitions::Vector{Tuple{S,Float64}}
    unique_transitions::Set{S}
    n_a_children::Int
    c_cdf_est::RunningCDFEstimator  # Conditional CDF.
    a_lock::ReentrantLock
end
PISActionNode(id::Int, a::A, n::Int, q::Float64, transitions::Vector{Tuple{S,Float64}}) where {S,A} =
    PISActionNode{S,A}(id, a, n, q, transitions, Set{S}(), 0, RunningCDFEstimator([0.0], [1e-5]), ReentrantLock())

@inline n(n::PISActionNode) = n.n
@inline q(n::PISActionNode) = n.q
@inline n_a_children(n::PISActionNode) = n.n_a_children

mutable struct PISStateNode{S,A}
    id::Int
    s_label::S
    total_n::Int
    children::Vector{A}
    s_lock::ReentrantLock
    # Action nodes currently being evaluated. Used for applying virtual loss.
    a_selected::Set{A}
end
PISStateNode(id::Int, s::S, total_n::Int, children::Vector{A}) where {S,A} =
    PISStateNode{S,A}(id, s, total_n, children, ReentrantLock(), Set{A}())

@inline total_n(n::PISStateNode) = n.total_n
@inline children(n::PISStateNode) = n.children
@inline n_children(n::PISStateNode) = length(children(n))
@inline isroot(n::PISStateNode) = (n.id == 1)

mutable struct PISTree{S,A}
    root::Union{Nothing,S}

    state_nodes::Dict{S, PISStateNode}
    state_action_nodes::Dict{Tuple{S,A}, PISActionNode}

    cdf_est::RunningCDFEstimator
    costs::Vector{Float64}
    weights::Vector{Float64}

    state_nodes_lock::ReentrantLock
    state_action_nodes_lock::ReentrantLock
    costs_weights_lock::ReentrantLock

    # Used to assign unique ids to nodes.
    _s_id_counter::Threads.Atomic{Int}
    _a_id_counter::Threads.Atomic{Int}

    function PISTree{S,A}(root::Union{Nothing,S}=nothing) where {S,A}
        return new(root,
                   Dict{S, PISStateNode}(),
                   Dict{Tuple{S,A}, PISActionNode}(),

                   RunningCDFEstimator([0.0], [1e-7]),
                   sizehint!(Float64[], 10_000),
                   sizehint!(Float64[], 10_000),

                   ReentrantLock(),
                   ReentrantLock(),
                   ReentrantLock(),

                   Threads.Atomic{Int}(1),
                   Threads.Atomic{Int}(1))
    end
end

Base.isempty(tree::PISTree) = isempty(tree.state_nodes)

function insert_state_node!(tree::PISTree{S,A}, s::S) where {S,A}
    snode = nothing
    Base.@lock tree.state_nodes_lock begin
        haskey(tree.state_nodes, s) && return tree.state_nodes[s]
        id = Threads.atomic_add!(tree._s_id_counter, 1)
        snode = PISStateNode(id, s, 0, Vector{A}())
        tree.state_nodes[s] = snode
    end
    return snode
end

function sanode_key(s_label::S, a_label::A) where {S,A}
    return tuple(s_label, a_label)
end

function insert_action_node!(tree::PISTree{S,A}, snode::PISStateNode{S,A},
                             a::Union{S,A}, n0::Int, q0::Float64) where {S,A}
    sanode = nothing
    Base.@lock tree.state_action_nodes_lock begin
        state_action_key = sanode_key(snode.s_label, a)
        haskey(tree.state_action_nodes, state_action_key) && return tree.state_action_nodes[state_action_key]
        id = Threads.atomic_add!(tree._a_id_counter, 1)
        sanode = PISActionNode(id, a, n0, q0, Vector{Tuple{S,Float64}}())
        tree.state_action_nodes[state_action_key] = sanode
    end
    Base.@lock snode.s_lock begin
        push!(snode.children, a)
        snode.total_n += n0
    end
    return sanode
end


export CbcOptimizer

using MathOptInterface
using Cbc.CbcCInterface

const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const CbcCI = CbcCInterface

mutable struct CbcOptimizer <: MOI.AbstractOptimizer
    inner::CbcModel
    CbcOptimizer() = new(CbcModel()) # Initializes with an empty model
end

struct CbcModelFormat
    num_rows::Int
    num_cols::Int
    row_idx::Vector{Int}
    col_idx::Vector{Int}
    values::Vector{Float64}
    col_lb::Vector{Float64}
    col_ub::Vector{Float64}
    obj::Vector{Float64}
    row_lb::Vector{Float64}
    row_ub::Vector{Float64}
    function CbcModelFormat(num_rows::Int, num_cols::Int)
        obj = fill(0.0, num_cols)
        row_idx = Int[]
        col_idx = Int[]
        values = Float64[]
        col_lb = fill(-Inf, num_cols)
        col_ub = fill(Inf, num_cols)
        row_lb = fill(-Inf, num_rows)
        row_ub = fill(Inf, num_rows)
        constraint_matrix = Tuple{Int,Int,Float64}[]
        new(num_rows, num_cols, row_idx, col_idx, values, col_lb, col_ub, obj, row_lb, row_ub)
    end
end

function load_constraint(ci::MOI.ConstraintIndex, cbc_model_format::CbcModelFormat, mapping::MOIU.IndexMap,
    f::MOI.SingleVariable, s::MOI.EqualTo)
    cbc_model_format.col_lb[mapping.varmap[f.variable].value] = s.value
    cbc_model_format.col_ub[mapping.varmap[f.variable].value] = s.value
end

function load_constraint(ci::MOI.ConstraintIndex, cbc_model_format::CbcModelFormat, mapping::MOIU.IndexMap,
    f::MOI.SingleVariable, s::MOI.LessThan)
    cbc_model_format.col_ub[mapping.varmap[f.variable].value] = s.upper
end

function load_constraint(ci::MOI.ConstraintIndex, cbc_model_format::CbcModelFormat, mapping::MOIU.IndexMap,
    f::MOI.SingleVariable, s::MOI.GreaterThan)
    cbc_model_format.col_lb[mapping.varmap[f.variable].value] = s.lower
end

function load_constraint(ci::MOI.ConstraintIndex, cbc_model_format::CbcModelFormat, mapping::MOIU.IndexMap,
    f::MOI.SingleVariable, s::MOI.Interval)
    cbc_model_format.col_lb[mapping.varmap[f.variable].value] = s.lower
    cbc_model_format.col_ub[mapping.varmap[f.variable].value] = s.upper
end

function push_terms(row_idx::Vector{Int}, col_idx::Vector{Int}, values::Vector{Float64},
    ci::MOI.ConstraintIndex{F,S}, terms::Vector{MOI.ScalarAffineTerm{Float64}},
    mapping::MOIU.IndexMap) where {F,S}
    for term in terms
        push!(row_idx, mapping.conmap[ci].value)
        push!(col_idx, mapping.varmap[term.variable_index].value)
        push!(values, term.coefficient)
    end
end

function load_constraint(ci::MOI.ConstraintIndex, cbc_model_format::CbcModelFormat,
    mapping::MOIU.IndexMap, f::MOI.ScalarAffineFunction, s::MOI.EqualTo)
    push_terms(cbc_model_format.row_idx, cbc_model_format.col_idx,
               cbc_model_format.values, ci, f.terms, mapping)
    cbc_model_format.row_lb[mapping.conmap[ci].value] = s.value - f.constant
    cbc_model_format.row_ub[mapping.conmap[ci].value] = s.value - f.constant
end

function load_constraint(ci::MOI.ConstraintIndex, cbc_model_format::CbcModelFormat,
    mapping::MOIU.IndexMap, f::MOI.ScalarAffineFunction, s::MOI.GreaterThan)
    push_terms(cbc_model_format.row_idx, cbc_model_format.col_idx,
               cbc_model_format.values, ci, f.terms, mapping)
    cbc_model_format.row_lb[mapping.conmap[ci].value] = s.lower - f.constant
end

function load_constraint(ci::MOI.ConstraintIndex, cbc_model_format::CbcModelFormat, mapping::MOIU.IndexMap,
    f::MOI.ScalarAffineFunction, s::MOI.LessThan)
    push_terms(cbc_model_format.row_idx, cbc_model_format.col_idx,
               cbc_model_format.values, ci, f.terms, mapping)
    cbc_model_format.row_ub[mapping.conmap[ci].value] = s.upper - f.constant
end

function load_constraint(ci::MOI.ConstraintIndex, cbc_model_format::CbcModelFormat, mapping::MOIU.IndexMap,
    f::MOI.ScalarAffineFunction, s::MOI.Interval)
    push_terms(cbc_model_format.row_idx, cbc_model_format.col_idx,
               cbc_model_format.values, ci, f.terms, mapping)
    cbc_model_format.row_ub[mapping.conmap[ci].value] = s.upper - f.constant
    cbc_model_format.row_lb[mapping.conmap[ci].value] = s.lower - f.constant
end

function load_obj(cbc_model_format::CbcModelFormat, mapping::MOIU.IndexMap,
    f::MOI.ScalarAffineFunction)
    # We need to increment values of objective function with += to handle cases like $x_1 + x_2 + x_1$
    # This is safe becasue objective function is initialized with zeros in the constructor
    for term in f.terms
        cbc_model_format.obj[mapping.varmap[term.variable_index].value] += term.coefficient
    end
end

function copy_constraints!(cbc_model_format::CbcModelFormat, user_optimizer::MOI.ModelLike,
    mapping::MOIU.IndexMap)
    for (F,S) in MOI.get(user_optimizer, MOI.ListOfConstraints())
        if !(S <: Union{MOI.ZeroOne, MOI.Integer})
            for ci in MOI.get(user_optimizer, MOI.ListOfConstraintIndices{F,S}())
                f = MOI.get(user_optimizer, MOI.ConstraintFunction(), ci)
                s = MOI.get(user_optimizer,  MOI.ConstraintSet(), ci)
                load_constraint(ci, cbc_model_format, mapping, f, s)
            end
        end
    end
end

function update_bounds_for_binary_vars!(col_lb::Vector{Float64}, col_ub::Vector{Float64}, zero_one_indices::Vector{Int})
    for idx in zero_one_indices
        if col_lb[idx] < 0.0
            col_lb[idx] = 0.0
        end
        if col_ub[idx] > 1.0
            col_ub[idx] = 1.0
        end
    end
end

function update_zeroone_indices(user_optimizer::MOI.ModelLike, mapping::MOIU.IndexMap,
    ci::Vector{MOI.ConstraintIndex{F,S}}, zero_one_indices::Vector{Int}) where {F, S}
    for i in 1:length(ci)
        f = MOI.get(user_optimizer, MOI.ConstraintFunction(), ci[i])
        push!(zero_one_indices, mapping.varmap[f.variable].value)
    end
end

function update_integer_indices(user_optimizer::MOI.ModelLike, mapping::MOIU.IndexMap,
    ci::Vector{MOI.ConstraintIndex{F,S}}, integer_indices::Vector{Int}) where {F, S}
    for i in 1:length(ci)
        f = MOI.get(user_optimizer, MOI.ConstraintFunction(), ci[i])
        push!(integer_indices, mapping.varmap[f.variable].value)
    end
end

"""
    function copy!(cbc_optimizer, user_optimizer; copynames=false)

Receive a cbc_optimizer which contains the pointer to the cbc C object and instantiate the object cbc_model_format::CbcModelFormat based on user_optimizer::AbstractModel (also provided by the user).
Function loadProblem of CbcCInterface requires all information stored in cbc_model_format.
"""
function MOI.copy!(cbc_optimizer::CbcOptimizer,
    user_optimizer::MOI.ModelLike; copynames=false)

    mapping = MOIU.IndexMap()

    num_cols = MOI.get(user_optimizer, MOI.NumberOfVariables())
    var_index = MOI.get(user_optimizer, MOI.ListOfVariableIndices())
    for i in 1:num_cols
        mapping.varmap[var_index[i]] = MOI.VariableIndex(i)
    end

    zero_one_indices = Int[]
    integer_indices = Int[]
    list_of_constraints = MOI.get(user_optimizer, MOI.ListOfConstraints())
    num_rows = 0
    for (F,S) in list_of_constraints
        if !(MOI.supportsconstraint(cbc_optimizer, F, S))
            return MOI.CopyResult(MOI.CopyUnsupportedConstraint,
            "Cbc MOI Interface does not support constraints of type " * (F,S) * ".", nothing)
        end

        ci = MOI.get(user_optimizer, MOI.ListOfConstraintIndices{F,S}())

        if F == MOI.SingleVariable
            if S == MOI.ZeroOne
                update_zeroone_indices(user_optimizer, mapping, ci, zero_one_indices)
            elseif S == MOI.Integer
                update_integer_indices(user_optimizer, mapping, ci, integer_indices)
            end
        else
            ## Update conmap for (F,S) for F != MOI.SingleVariable
            ## Single variables are treated by bounds in Cbc, so no
            ## need to add a row
            for i in 1:length(ci)
                mapping.conmap[ci[i]] = MOI.ConstraintIndex{F,S}(num_rows + i)
            end
            num_rows += MOI.get(user_optimizer, MOI.NumberOfConstraints{F,S}())
        end
    end

    cbc_model_format = CbcModelFormat(num_rows, num_cols)

    copy_constraints!(cbc_model_format, user_optimizer, mapping)
    update_bounds_for_binary_vars!(cbc_model_format.col_lb, cbc_model_format.col_ub, zero_one_indices)

    ## Copy objective function
    objF = MOI.get(user_optimizer, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    load_obj(cbc_model_format, mapping, objF)
    sense = MOI.get(user_optimizer, MOI.ObjectiveSense())
    MOI.set!(cbc_optimizer, MOI.ObjectiveSense(), sense)

    ## Load the problem to Cbc
    CbcCI.loadProblem(cbc_optimizer.inner, sparse(cbc_model_format.row_idx, cbc_model_format.col_idx,
                                                  cbc_model_format.values, cbc_model_format.num_rows,
                                                  cbc_model_format.num_cols),
                      cbc_model_format.col_lb, cbc_model_format.col_ub, cbc_model_format.obj,
                      cbc_model_format.row_lb, cbc_model_format.row_ub)

    empty!(cbc_model_format.row_idx)
    empty!(cbc_model_format.col_idx)
    empty!(cbc_model_format.values)

    ## Set integer variables
    for idx in vcat(integer_indices, zero_one_indices)
        CbcCI.setInteger(cbc_optimizer.inner, idx-1)
    end

    return MOI.CopyResult(MOI.CopySuccess, "Model was copied succefully.", MOIU.IndexMap(mapping.varmap, mapping.conmap))
end

function MOI.optimize!(cbc_optimizer::CbcOptimizer)
    # Call solve function
    CbcCI.solve(cbc_optimizer.inner)
end

## canadd, canset, canget functions

function MOI.canaddvariable(cbc_optimizer::CbcOptimizer)
    return false
end

## supports constraints

MOI.supportsconstraint(::CbcOptimizer, ::Type{<:Union{MOI.ScalarAffineFunction{Float64}, MOI.SingleVariable}},
::Type{<:Union{MOI.EqualTo{Float64}, MOI.Interval{Float64}, MOI.LessThan{Float64},
MOI.GreaterThan{Float64}, MOI.ZeroOne, MOI.Integer}}) = true

MOI.supports(cbc_optimizer::CbcOptimizer, object::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}) = true

## Set functions

function MOI.write(cbc_optimizer::CbcOptimizer, filename::String)
    if !endswith("filename", "mps")
        error("CbcOptimizer only supports writing .mps files")
    else
        writeMps(cbc_optimizer.inner, filename)
    end
end

# empty!
function MOI.empty!(cbc_optimizer::CbcOptimizer)
    cbc_optimizer.inner = CbcModel()
end

function MOI.set!(cbc_optimizer::CbcOptimizer, object::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    if sense == MOI.MaxSense
        CbcCI.setObjSense(cbc_optimizer.inner, -1)
    else ## Other senses are set as minimization (cbc default)
        CbcCI.setObjSense(cbc_optimizer.inner, 1)
    end
end

## Get functions

function MOI.canget(cbc_optimizer::CbcOptimizer, object::MOI.PrimalStatus)
    if object.N != 1
        return false
    end
    return MOI.get(cbc_optimizer, MOI.ResultCount()) == 1
end

MOI.canget(cbc_optimizer::CbcOptimizer, object::Union{MOI.NodeCount, MOI.ResultCount,
MOI.TerminationStatus, MOI.ObjectiveSense, MOI.ObjectiveValue, MOI.ObjectiveBound, MOI.NumberOfVariables}) = true

MOI.canget(cbc_optimizer::CbcOptimizer, object::MOI.VariablePrimal, indexTypeOrObject::Type{MOI.VariableIndex}) = true

function MOI.isempty(cbc_optimizer::CbcOptimizer)
    return (CbcCI.getNumCols(cbc_optimizer.inner) == 0 && CbcCI.getNumRows(cbc_optimizer.inner) == 0)
end

MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.NumberOfVariables) = getNumCols(cbc_optimizer.inner)

MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.ObjectiveBound) = CbcCI.getBestPossibleObjValue(cbc_optimizer.inner)

MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.NodeCount) = CbcCI.getNodeCount(cbc_optimizer.inner)

function MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.ObjectiveValue)
    return CbcCI.getObjValue(cbc_optimizer.inner)
end

function MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.VariablePrimal, ref::MOI.VariableIndex)
    variablePrimals = CbcCI.getColSolution(cbc_optimizer.inner)
    return variablePrimals[ref.value]
end

function MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.VariablePrimal, ref::Vector{MOI.VariableIndex})
    variablePrimals = CbcCI.getColSolution(cbc_optimizer.inner)
    return [variablePrimals[vi.value] for vi in ref]
end

function MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.ResultCount)
    if (isProvenInfeasible(cbc_optimizer.inner) || isContinuousUnbounded(cbc_optimizer.inner)
        || isAbandoned(cbc_optimizer.inner) || CbcCI.getObjValue(cbc_optimizer.inner) >= 1e300)
        return 0
    end
    return 1
end

function MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.ObjectiveSense)
    CbcCI.getObjSense(cbc_optimizer.inner) == 1 && return MOI.MinSense
    CbcCI.getObjSense(cbc_optimizer.inner) == -1 && return MOI.MaxSense
end

function MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.TerminationStatus)

    if isProvenInfeasible(cbc_optimizer.inner)
        return MOI.InfeasibleNoResult
    elseif isContinuousUnbounded(cbc_optimizer.inner)
        return MOI.InfeasibleOrUnbounded
    elseif isNodeLimitReached(cbc_optimizer.inner)
        return MOI.NodeLimit
    elseif isSecondsLimitReached(cbc_optimizer.inner)
        return MOI.TimeLimit
    elseif isSolutionLimitReached(cbc_optimizer.inner)
        return MOI.SolutionLimit
    elseif (isProvenOptimal(cbc_optimizer.inner) || isInitialSolveProvenOptimal(cbc_optimizer.inner)
        || MOI.get(cbc_optimizer, MOI.ResultCount()) == 1)
        return MOI.Success
    elseif isAbandoned(cbc_optimizer.inner)
        return MOI.Interrupted
    else
        error("Internal error: Unrecognized solution status")
    end

end

function MOI.get(cbc_optimizer::CbcOptimizer, object::MOI.PrimalStatus)
    if isProvenOptimal(cbc_optimizer.inner) || isInitialSolveProvenOptimal(cbc_optimizer.inner)
        return MOI.FeasiblePoint
    elseif isProvenInfeasible(cbc_optimizer.inner)
        return MOI.InfeasiblePoint
    end
end


using Oceananigans.TimeSteppers: update_state!

import Oceananigans.Fields: set!

using Oceananigans.Utils: @apply_regionally, apply_regionally!

"""
    set!(model; kwargs...)

Set velocity and tracer fields of `model`. The keyword arguments
`kwargs...` take the form `name=data`, where `name` refers to one of the
fields of `model.velocities` or `model.tracers`, and the `data` may be an array,
a function with arguments `(x, y, z)`, or any data type for which a
`set!(ϕ::AbstractField, data)` function exists.

Example
=======
```julia
model = HydrostaticFreeSurfaceModel(grid=RectilinearGrid(size=(32, 32, 32), length=(1, 1, 1))

# Set u to a parabolic function of z, v to random numbers damped
# at top and bottom, and T to some silly array of half zeros,
# half random numbers.

u₀(x, y, z) = z/model.grid.Lz * (1 + z/model.grid.Lz)
v₀(x, y, z) = 1e-3 * rand() * u₀(x, y, z)

T₀ = rand(size(model.grid)...)
T₀[T₀ .< 0.5] .= 0

set!(model, u=u₀, v=v₀, T=T₀)
```
"""
function set!(model::HydrostaticFreeSurfaceModel; kwargs...)
    for (fldname, value) in kwargs
        if fldname ∈ propertynames(model.velocities)
            ϕ = getproperty(model.velocities, fldname)
        elseif fldname ∈ propertynames(model.tracers)
            ϕ = getproperty(model.tracers, fldname)
        elseif fldname ∈ propertynames(model.free_surface)
            ϕ = getproperty(model.free_surface, fldname)
        else
            throw(ArgumentError("name $fldname not found in model.velocities, model.tracers, or model.free_surface"))
        end

        @apply_regionally set!(ϕ, value)
    end

    update_state!(model)

    return nothing
end


module Scotch

# Define before execution, for instance in your /~.juliarc.jl add the following line
# const SCOTCH_LIBSCOTCH = "/path/to/scotch/lib/libscotch.dylib"
const libscotch = Main.SCOTCH_LIBSCOTCH

mutable struct Graph
end

mutable struct Strat
end

mutable struct Ordering
end

mutable struct File # libc's File
end

function version()
    ver = Cint[0]
    rel = Cint[0]
    pat = Cint[0]
    ccall((:SCOTCH_version, libscotch), Void,
            (Ptr{Cint}, Ptr{Cint}, Ptr{Cint}),
            pointer(ver), pointer(rel), pointer(pat))
    return (ver[1], rel[1], pat[1])
end

function graphAlloc()
    graph = ccall((:SCOTCH_graphAlloc, libscotch), Ptr{Graph}, ())
    if graph == C_NULL
        error("Error in graphAlloc")
    end
    return graph
end

function graphInit(graph::Ptr{Graph})
    err = ccall((:SCOTCH_graphInit, libscotch), Cint,
                (Ptr{Graph},),
                graph)
    err == 0 || error("Error in graphInit")
    return
end

function graphLoad(graph::Ptr{Graph}, filename, baseval::Integer, flagval::Integer)
    m = "r"
    file = ccall((:fopen, "libc"), Ptr{File},
             (Ptr{Char}, Ptr{Char}),
             pointer(filename), pointer(m))
    file != C_NULL || error("Error in graphInit, fopen")
    err = ccall((:SCOTCH_graphLoad, libscotch), Cint,
                (Ptr{Graph}, Ptr{File}, Cint, Cint),
                graph, file, Cint(baseval), Cint(flagval))
    err == 0 || error("Error in graphInit")
    return
end

function graphBuild(graph::Ptr{Graph}, A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti <: Integer}
    A.n == A.m || error("Matrix should be square")
    # Remove self loops
    vertnbr = A.n
    verttab = Ptr{Cint}(Libc.malloc(sizeof(Cint) * (vertnbr+1)))
    edgetab = Ptr{Cint}(Libc.malloc(sizeof(Cint) * length(A.rowval)))
    unsafe_store!(verttab, 1)
    k = 1
    edgenbr = 0
    for i = 1:vertnbr
        for j = A.colptr[i]:(A.colptr[i+1]-1)
            if A.rowval[j] != i
                unsafe_store!(edgetab, A.rowval[j], k)
                k += 1
                edgenbr += 1
            end
        end
        unsafe_store!(verttab, k, i+1)
    end
    # FIXME: resize edgetab to k
    err = ccall((:SCOTCH_graphBuild, libscotch), Cint,
          (Ptr{Graph}, Cint, Cint,    Ptr{Cint},  Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Cint,    Ptr{Cint},  Ptr{Cint}),
          graph,       1,    vertnbr, verttab,    C_NULL,    C_NULL,    C_NULL,    edgenbr, edgetab,    C_NULL)
    err == 0 || error("Error in graphBuild")
    return
end

function graphCheck(graph::Ptr{Graph})
    err = ccall((:SCOTCH_graphCheck, libscotch), Cint,
                (Ptr{Graph},),
                graph)
    err == 0 || error("Error in graphCheck")
end

function graphSize(graph::Ptr{Graph})
    nv = Cint[0]
    ne = Cint[0]
    ccall((:SCOTCH_graphSize, libscotch), Void,
               (Ptr{Graph}, Ptr{Cint}, Ptr{Cint}),
               graph, pointer(nv), pointer(ne))
    return (nv[1], ne[1])
end

function graphFree(graph::Ptr{Graph})
    ccall((:SCOTCH_graphFree, libscotch), Void,
          (Ptr{Graph},),
          graph)
end

function graphExit(graph::Ptr{Graph})
    ccall((:SCOTCH_graphExit, libscotch), Void,
          (Ptr{Graph},),
          graph)
end

function graphStat(graph::Ptr{Graph})
    A1 = Cint[0,0,0,0,0,0,0,0]
    A2 = Cdouble[0,0,0,0,0,0]
    @show A1, A2
    ccall((:SCOTCH_graphStat, libscotch), Void,
          (Ptr{Graph},
           Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble},
           Ptr{Cint}, Ptr{Cint},            Ptr{Cdouble}, Ptr{Cdouble},
           Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble}),
          graph,
          pointer(A1, 1), pointer(A1, 2), pointer(A1, 3), pointer(A2, 1), pointer(A2, 2),
          pointer(A1, 4), pointer(A1, 5),                 pointer(A2, 3), pointer(A2, 4),
          pointer(A1, 6), pointer(A1, 7), pointer(A1, 8), pointer(A2, 5), pointer(A2, 6))
    @show A1, A2
    return (A1[1], A1[2], A1[3], A2[1], A2[2],
            A1[4], A1[5],        A2[3], A2[4],
            A1[6], A1[7], A1[8], A2[5], A2[6])
end

# Strategy

function stratAlloc()
    strat = ccall((:SCOTCH_stratAlloc, libscotch), Ptr{Strat}, ())
    strat != C_NULL || error("Error in stratAlloc")
    return strat
end

function stratInit(strat::Ptr{Strat})
    err = ccall((:SCOTCH_stratInit, libscotch), Cint,
                (Ptr{Strat},),
                strat)
    err == 0 || error("Error in stratInit")
    return
end

function stratGraphOrder(strat::Ptr{Strat}, st::String)
    err = ccall((:SCOTCH_stratGraphOrder, libscotch), Cint,
                (Ptr{Strat}, Ptr{Char}),
                strat, pointer(st))
    err == 0 || error("Error in stratGraphOrder")
    return
end

function stratExit(strat::Ptr{Strat})
    ccall((:SCOTCH_stratExit, libscotch), Void,
          (Ptr{Strat},),
          strat)
end

# Ordering

function graphOrder(graph::Ptr{Graph}, strat::Ptr{Strat})
    (nv, ne) = graphSize(graph)
    permtab = Array{Cint, 1}(nv)
    peritab = Array{Cint, 1}(nv)
    cblkptr = Array{Cint, 1}(1)
    rangtab = Array{Cint, 1}(nv + 1)
    treetab = Array{Cint, 1}(nv)
    err = ccall((:SCOTCH_graphOrder, libscotch), Cint,
            (Ptr{Graph}, Ptr{Strat}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}),
            graph,       strat,      pointer(permtab), pointer(peritab), pointer(cblkptr), pointer(rangtab), pointer(treetab))
    err == 0 || error("Error in graphOrder")
    return (permtab, peritab, cblkptr, rangtab, treetab)
end

function orderAlloc()
    optr = ccall((:SCOTCH_orderAlloc, libscotch), Ptr{Ordering}, ())
    optr != C_NULL || error("Error in orderAlloc")
    return optr
end

function graphOrderInit(graph::Ptr{Graph}, ordering::Ptr{Ordering}, permtab::Array{Cint,1}, peritab::Array{Cint,1}, cblkptr::Array{Cint,1}, rangtab::Array{Cint,1}, treetab::Array{Cint,1})
    err = ccall((:SCOTCH_graphOrderInit, libscotch), Cint,
                (Ptr{Graph}, Ptr{Ordering}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}),
                graph, ordering, pointer(permtab), pointer(peritab), pointer(cblkptr), pointer(rangtab), pointer(treetab))
    err == 0 || error("Error in graphOrderInit")
    return
end

function graphOrderComputeList(graph::Ptr{Graph}, ordering::Ptr{Ordering}, listtab::Array{Cint,1}, strat::Ptr{Strat})
    err = ccall((:SCOTCH_graphOrderComputeList, libscotch), Cint,
                (Ptr{Graph}, Ptr{Ordering}, Cint, Ptr{Cint}, Ptr{Strat}),
                graph, ordering, Cint(length(listtab)), pointer(listtab), strat)
    err == 0 || error("Error in graphOrderComputeList")
    return
end

function graphOrderCompute(graph::Ptr{Graph}, ordering::Ptr{Ordering}, strat::Ptr{Strat})
    err = ccall((:SCOTCH_graphOrderCompute, libscotch), Cint,
                (Ptr{Graph}, Ptr{Ordering}, Ptr{Strat}),
                graph, ordering, strat)
    err == 0 || error("Error in graphOrderCompute")
    return
end

function graphOrderCheck(graph::Ptr{Graph}, ordering::Ptr{Ordering})
    err = ccall((:SCOTCH_graphOrderCheck, libscotch), Cint,
                (Ptr{Graph}, Ptr{Ordering}),
                graph, ordering)
    err == 0 || error("Error in graphOrderCheck")
end

function graphOrderExit(graph::Ptr{Graph}, ordering::Ptr{Ordering})
    ccall((:SCOTCH_graphOrderExit, libscotch), Void,
                (Ptr{Graph}, Ptr{Ordering}),
                graph, ordering)
end

end # module Scotch


using OrderedCollections

# Initialize model object
m = AnSchorfheide(testing = true)
m <= Setting(:forecast_horizons, 12)

estroot = joinpath(dirname(@__FILE__), "..", "reference")
overrides = forecast_input_file_overrides(m)
overrides[:mode] = joinpath(estroot, "optimize.h5")

# Zero out measurement error (or else filtering shocks won't work exactly)
for para in [:e_y, :e_π, :e_R]
    m[para].value = 0
end

# Simulate dummy scenarios
def = Scenario(:defscen, "Test Default Scenario", [:obs_gdp, :obs_cpi], Symbol[], "REF")
alt = Scenario(:altscen, "Test Alternative Scenario", [:obs_gdp, :obs_cpi], [:g_sh, :rm_sh], "REF")

forecast_scenario(m, def, verbose = :none)
forecast_scenario(m, alt, verbose = :none)

expect = JLD2.jldopen(get_scenario_input_file(m, alt), "r") do file
    read(file, "arr")
end

JLD2.jldopen(get_scenario_output_files(m, alt, [:forecastobs])[:forecastobs], "r") do file
    observable_indices = read(file, "observable_indices")
    global inds = map(var -> observable_indices[var], def.target_names)
    global actual = read(file, "arr")[:, inds, :]
end

@testset "Test dummy scenarios" begin
    @test expect ≈ actual
end

# Simulate always switching
probs_enter = vcat([1], zeros(11))
probs_exit  = zeros(12)
allalt = SwitchingScenario(:allalt, alt, def, probs_enter, probs_exit)
simulate_switching(m, allalt, verbose = :none)

original_draws = JLD2.jldopen(get_scenario_output_files(m, alt, [:forecastobs])[:forecastobs]) do file
    read(file, "arr")
end
actual = JLD2.jldopen(get_scenario_output_files(m, allalt, [:forecastobs])[:forecastobs]) do file
    read(file, "arr")
end
@testset "Test switching scenarios" begin
    @test original_draws ≈ actual
end

# Simulate never switching
probs_enter = zeros(12)
probs_exit  = zeros(12)
alldef = SwitchingScenario(:alldef, alt, def, probs_enter, probs_exit)
simulate_switching(m, alldef, verbose = :none)

default_draws = JLD2.jldopen(get_scenario_output_files(m, def, [:forecastobs])[:forecastobs]) do file
    read(file, "arr")
end

actual = JLD2.jldopen(get_scenario_output_files(m, alldef, [:forecastobs])[:forecastobs]) do file
    read(file, "arr")
end

@testset "Test never switching scenarios" begin
    for i = 1:10 # Check that the ith draw of actual matches some draw j from default_draws
        matches_one = false
        for j = 1:10
            (actual[i, :, :] ≈ default_draws[j, :, :]) && (matches_one = true)
        end
        @test matches_one
    end
end

# Simulate switching in at t = 5 and out at t = 9
probs_enter = vcat(zeros(4), [1], zeros(7))
probs_exit  = vcat(zeros(8), [1], zeros(3))
somealt = SwitchingScenario(:somealt, alt, def, probs_enter, probs_exit)
simulate_switching(m, somealt, verbose = :none)

@testset "Test switching scenarios, where the switch happens at particular points in time" begin
    global actual = JLD2.jldopen(get_scenario_output_files(m, somealt, [:forecastobs])[:forecastobs]) do file
        read(file, "arr")
    end
    for i = 1:10
        @test actual[i, :, 5:8] ≈ original_draws[i, :, 1:4]
        matches_one = false
        for j = 1:10
            (actual[i, :, [1:4;9:12]] ≈ default_draws[j, :, [1:4;9:12]]) && (matches_one = true)
        end
        @test matches_one
    end
end

# Transform single scenarios
scenario_means_bands(m, alt, verbose = :none)
scenario_means_bands(m, allalt, verbose = :none)
global actual = dropdims(mean(original_draws, dims = 1), dims=1)
@testset "Test single scenarios" begin
    mb1 = read_scenario_mb(m, alt, :forecastutobs)
    mb2 = read_scenario_mb(m, allalt, :forecastutobs)
    for mb in [mb1, mb2]
        for (var, ind) in mb.metadata[:indices]
            @test mb.means[!, var] ≈ actual[ind, :]
        end
    end

    mb1 = read_scenario_mb(m, alt, :forecastobs)
    mb2 = read_scenario_mb(m, allalt, :forecastobs)
    for mb in [mb1, mb2]
        @test mb.means[!, :obs_gdp] ≈ quartertoannual(actual[1, :])
        @test mb.means[!, :obs_cpi] ≈ quartertoannual(actual[2, :])
        @test mb.means[!, :obs_nominalrate] ≈ actual[3, :]
    end

    mb1 = read_scenario_mb(m, alt, :forecast4qobs)
    mb2 = read_scenario_mb(m, allalt, :forecast4qobs)
    for mb in [mb1, mb2]
        @test mb.means[!, :obs_gdp] ≈ loggrowthtopct_4q_approx(actual[1, :], zeros(3))
        @test mb.means[!, :obs_cpi] ≈ loggrowthtopct_4q_approx(actual[2, :], zeros(3))
        @test mb.means[!, :obs_nominalrate] ≈ actual[3, :]
    end
end
# Aggregate scenarios equally
aggall = ScenarioAggregate(:aggall, "Test Scenario Aggregate", AbstractScenario[def, alt],
                           [0.5, 0.5], 20, false, "REF")
scenario_means_bands(m, aggall, verbose = :none)
global actual = dropdims(mean(cat(default_draws, original_draws, dims = 1), dims = 1), dims=1)

@testset "Test scenario aggregation" begin
    mb = read_scenario_mb(m, aggall, :forecastutobs)
    for (var, ind) in mb.metadata[:indices]
        @test mb.means[!, var] ≈ actual[ind, :]
    end
    # Aggregate scenarios to only have alternative scenaro
    aggalt = ScenarioAggregate(:aggalt, "Test Scenario Aggregate", AbstractScenario[def, alt],
                               [0.0, 1.0], 20, false, "REF")
    scenario_means_bands(m, aggalt, verbose = :none)
    global actual = dropdims(mean(repeat(original_draws, outer = [2, 1, 1]), dims = 1), dims=1)

    mb = read_scenario_mb(m, aggalt, :forecastutobs)
    for (var, ind) in mb.metadata[:indices]
        @test mb.means[!, var] ≈ actual[ind, :]
    end
end

# Aggregate scenarios, drawing with replacement
aggrep1 = ScenarioAggregate(:aggrep1, "Test Scenario Aggregate", AbstractScenario[def, alt],
                            [0.5, 0.5], 20, true, "REF")
aggrep2 = ScenarioAggregate(:aggrep2, "Test Scenario Aggregate", AbstractScenario[def, alt],
                            [0.0, 1.0], 20, true, "REF")

# Simulate dummy scenario with shock scaling
scale_scenario = Scenario(:altscen, "Test Shock Scaling Scenario", [:obs_gdp, :obs_cpi], [:g_sh, :rm_sh], "REF",
                 shock_scaling = 2.0)
forecast_scenario(m, scale_scenario, verbose = :none)

expect = JLD2.jldopen(get_scenario_input_file(m, scale_scenario), "r") do file
    2.0 * read(file, "arr")
end

JLD2.jldopen(get_scenario_output_files(m, alt, [:forecastobs])[:forecastobs], "r") do file
    observable_indices = read(file, "observable_indices")
    global inds = map(var -> observable_indices[var], def.target_names)
    global actual = read(file, "arr")[:, inds, :]
end
@testset "Test scenario aggregation, drawing with replacement" begin
    @test expect ≈ actual
end

# Delete outputted files
#=
raw_fn_root  = "../reference/output_data/an_schorfheide/ss0/scenarios/raw/"
work_fn_root = "../reference/output_data/an_schorfheide/ss0/scenarios/work/"
rm(joinpath(dirname(@__FILE__), raw_fn_root, "forecastobs_scen=altscen_svin=REF_test.jld2"))
rm(joinpath(dirname(@__FILE__), raw_fn_root, "forecastpseudo_scen=altscen_svin=REF_test.jld2"))
rm(joinpath(dirname(@__FILE__), work_fn_root, "mbforecast4qobs_scen=altscen_svin=REF_test.jld2"))
rm(joinpath(dirname(@__FILE__), work_fn_root, "mbforecast4qpseudo_scen=altscen_svin=REF_test.jld2"))
rm(joinpath(dirname(@__FILE__), work_fn_root, "mbforecastobs_scen=altscen_svin=REF_test.jld2"))
rm(joinpath(dirname(@__FILE__), work_fn_root, "mbforecastpseudo_scen=altscen_svin=REF_test.jld2"))
rm(joinpath(dirname(@__FILE__), work_fn_root, "mbforecastutobs_scen=altscen_svin=REF_test.jld2"))
rm(joinpath(dirname(@__FILE__), work_fn_root, "mbforecastutpseudo_scen=altscen_svin=REF_test.jld2"))
=#
nothing