"""
```
forecast(m, system, z0; cond_type = :none,
    enforce_zlb = false, shocks = Matrix{S}(undef, 0,0))

forecast(m, altpolicy, z0, states, obs, pseudo, shocks)

forecast(system, z0, shocks; enforce_zlb = false)

forecast(m, system, z0, shocks; cond_type = :none, enforce_zlb = false)
```

The first method produces a forecast, given a state space system,
initial state, and shocks, using information about the desired forecast
contained in `m`. It enforces the ZLB by using monetary policy shocks.

The second method is similar but differs in two ways. First,
it produces forecasts specifically when an alternative policy
is used. Second, it enforces the ZLB by treating it as a temporary
alternative policy.

The third and fourth methods are internal functions used by the first two methods.

### Inputs

- `system::System{S}`: state-space system matrices
- `z0::Vector{S}`: state vector in the final historical period
- `shocks::Matrix{S}`: `nshocks` x `nperiods` matrix of shocks to use when
  forecasting. Note that in the first method, `nperiods` doesn't necessarily
  have to equal `forecast_horizons(m)`; it will be truncated or padded with
  zeros appropriately

**Method 1 and 2 only:**

- `m::AbstractDSGEModel`

**Method 2 only:**

- `altpolicy::Symbol`: Which alternative policy is being used
- `obs::Matrix{S <: Real}`: matrix of forecasted observables

where `S<:AbstractFloat`.

### Keyword Arguments

- `cond_type::Symbol`: one of `:none`, `:semi`, or `:full`, used to determine
  how many periods to forecast ahead. If `cond_type in [:semi, :full]`, the
  forecast horizon is reduced by the number of periods of conditional
  data. Defaults to `:none`.
- `enforce_zlb::Bool`: whether to enforce the zero lower bound. Defaults to
  `false`.
- `shocks::Matrix{S}`: matrix of size `nshocks` x `shock_horizon` of shock
  innovations under which to forecast. If `shock_horizon > horizon`, the extra
  periods of shocks will be ignored; if `shock_horizon < horizon`, zeros will be
  filled in for the shocks hitting the remaining forecasted periods.
- `draw_shocks::Bool`: if `isempty(shocks)`, indicates whether to draw shocks
  according to:

  1. If `forecast_tdist_shocks(m)`, draw `horizons` many shocks from a
     `Distributions.TDist(forecast_tdist_df_val(m))`
  2. Otherwise, draw `horizons` many shocks from a
     `DegenerateMvNormal(zeros(nshocks), sqrt(system[:QQ]))`

  or to set `shocks` to a `nshocks` x `horizon` matrix of zeros. Defaults to
  `false`. If `shocks` is provided as a keyword argument, this flag has no
  effect.

**Method 2 only:**

- `set_zlb_regime_vals::Function`: user-provided function that adds additional regimes to
    regime-switching parameters if not enough regimes exist to impose the ZLB
    as a temporary alternative policy. Defaults to `identity`, and nothing will happen
    if this is the case.

- `tol::{<: Real}`: Tolerance for the smallest permissible value for the nominal interest rate.
    Defaults to -1e-14.

### Outputs

- `states::Matrix{S}`: matrix of size `nstates` x `horizon` of forecasted states
- `obs::Matrix{S}`: matrix of size `nobs` x `horizon` of forecasted observables
- `pseudo::Matrix{S}`: matrix of size `npseudo` x `horizon` of forecasted
  pseudo-observables
- `shocks::Matrix{S}`: matrix of size `nshocks` x `horizon` of shock innovations
"""
function forecast(m::AbstractDSGEModel, system::Union{RegimeSwitchingSystem{S}, System{S}},
                  z0::Vector{S}; cond_type::Symbol = :none, enforce_zlb::Bool = false,
                  shocks::AbstractMatrix{S} = Matrix{S}(undef, 0, 0), draw_shocks::Bool = false) where {S<:AbstractFloat}

    # Numbers of things
    nshocks = n_shocks_exogenous(m)
    horizon = forecast_horizons(m; cond_type = cond_type)

    if isempty(shocks)
        # Populate shocks matrix
        if draw_shocks
            # If Regime Switching System, need to get proper QQ in each forecast period.
            if isa(system, RegimeSwitchingSystem)
                shocks = zeros(nshocks, horizon)
                μ = zeros(nshocks)
                # Forecast regime indices (starting with 1 in 1st fcast period, going up to horizon)
                regime_inds = get_fcast_regime_inds(m, horizon, cond_type)
                # The regime in which forecast starts (indexing from presample_start as reg=1)
                sys_ind = cond_type == :none ? get_setting(m, :reg_forecast_start) :
                max(get_setting(m, :reg_forecast_start), get_setting(m, :reg_post_conditional_end))
                for ts in regime_inds
                    σ = sqrt.(system[sys_ind, :QQ])
                    dist = if forecast_tdist_shocks(m)
                        # Use t-distributed shocks
                        ν = forecast_tdist_df_val(m)
                        DegenerateDiagMvTDist(μ, σ, ν)
                    else
                        # Use normally distributed shocks
                        DegenerateMvNormal(μ, σ)
                    end
                    shocks[:, ts] = rand(dist, length(ts))
                    sys_ind += 1
                end
            else
                μ = zeros(S, nshocks)
                σ = sqrt.(system[:QQ])
                dist = if forecast_tdist_shocks(m)
                    # Use t-distributed shocks
                    ν = forecast_tdist_df_val(m)
                    DegenerateDiagMvTDist(μ, σ, ν)
                else
                    # Use normally distributed shocks
                    DegenerateMvNormal(μ, σ)
                end
                shocks = rand(dist, horizon)
            end

            # Forecast without anticipated shocks
            if n_mon_anticipated_shocks(m) > 0
                ind_ant1 = m.exogenous_shocks[Symbol(get_setting(m, :monetary_policy_shock), :l1)]
                ind_antn = m.exogenous_shocks[Symbol(get_setting(m, :monetary_policy_shock), "l$(n_mon_anticipated_shocks(m))")]
                ant_shock_inds = ind_ant1:ind_antn
                shocks[ant_shock_inds, :] .= 0
            end
        else
            shocks = zeros(S, nshocks, horizon)
        end
    else
        # Adjust size of shocks matrix, padding with zeros or cutting off
        # periods of shocks if necessary
        shock_horizon = size(shocks, 2)
        if shock_horizon <= horizon
            shocks0 = zeros(nshocks, horizon - shock_horizon)
            shocks = hcat(shocks, shocks0)
        else
            shocks = shocks[:, 1:horizon]
        end
    end

    if haskey(get_settings(m), :n_periods_no_shocks)
        if get_setting(m, :n_periods_no_shocks) > 0
            shocks[:, 1:get_setting(m, :n_periods_no_shocks)] = zeros(size(shocks, 1), get_setting(m, :n_periods_no_shocks))
        end
    end

    # Populate shocks matrix under alternative policy, if
    # user has specified a function to do so
    #=    alt_policy = alternative_policy(m)

    # TODO: check that these handling of the altpolicy cases work properly
    if (alt_policy.solve != identity &&
    alt_policy.forecast_init != identity && (!haskey(m.settings, :initialize_pgap_ygap) || get_setting(m, :initialize_pgap_ygap)))
    shocks, z0 = alt_policy.forecast_init(m, shocks, z0, cond_type = cond_type)
    end=#

    #=
    # This code block has been moved to forecast_one_draw in drivers.jl
    # b/c we want to actually alter the smoothed states so that pgap or ygap
    # starts from the user-specified gap value and that this start is actually recorded.
    # TODO: check that these handling of the altpolicy cases work properly
    # Added separately to cover case where you're not using altpolicy but are using gensys2
    if haskey(m.settings, :pgap_type) && haskey(get_settings(m), :pgap_value)
    if get_setting(m, :pgap_type) == :ngdp
    _, z0 = ngdp_forecast_init(m, shocks, z0, cond_type = cond_type)
    end
    end

    if haskey(m.settings, :ygap_type) && haskey(get_settings(m), :ygap_value)
    if get_setting(m, :ygap_type) == :smooth_ait_gdp
    _, z0 = smooth_ait_gdp_forecast_init(m, shocks, z0, cond_type = cond_type)
    end
    end
    =#

    # Get variables necessary to enforce the zero lower bound in the forecast
    ind_r = m.observables[get_setting(m, :nominal_rate_observable)]
    ind_r_sh = m.exogenous_shocks[get_setting(m, :monetary_policy_shock)]
    zlb_value = forecast_zlb_value(m)

    if isa(system, RegimeSwitchingSystem)
        forecast(m, system, z0, shocks; cond_type = cond_type, enforce_zlb = enforce_zlb,
                 ind_r = ind_r, ind_r_sh = ind_r_sh, zlb_value = zlb_value)
    else
        forecast(system, z0, shocks; enforce_zlb = enforce_zlb,
                 ind_r = ind_r, ind_r_sh = ind_r_sh, zlb_value = zlb_value)
    end
end

function forecast(system::System{S}, z0::Vector{S},
                  shocks::Matrix{S}; enforce_zlb::Bool = false, ind_r::Int = -1,
                  ind_r_sh::Int = -1, zlb_value::S = 0.1/4) where {S<:AbstractFloat}

    # Unpack system
    T, R, C = system[:TTT], system[:RRR], system[:CCC]
    Q, Z, D = system[:QQ], system[:ZZ], system[:DD]
    Z_pseudo, D_pseudo = system[:ZZ_pseudo], system[:DD_pseudo]

    # Setup
    nshocks = size(R, 2)
    nstates = size(T, 2)
    nobs    = size(Z, 1)
    npseudo = size(Z_pseudo, 1)
    horizon = size(shocks, 2)

    # Define our iteration function
    function iterate(z_t1, ϵ_t)
        z_t = C + T*z_t1 + R*ϵ_t

        # Change monetary policy shock to account for 0.13 interest rate bound
        if enforce_zlb
            interest_rate_forecast = getindex(D + Z*z_t, ind_r)
            if interest_rate_forecast < zlb_value
                # Solve for interest rate shock causing interest rate forecast to be exactly ZLB
                ϵ_t[ind_r_sh] = 0.
                z_t = C + T*z_t1 + R*ϵ_t
                ϵ_t[ind_r_sh] = getindex((zlb_value - D[ind_r] - Z[ind_r, :]'*z_t) / (Z[ind_r, :]' * R[:, ind_r_sh]), 1)

                # Forecast again with new shocks
                z_t = C + T*z_t1 + R*ϵ_t

                # Confirm procedure worked
                interest_rate_forecast = getindex(D + Z*z_t, ind_r)
                @assert interest_rate_forecast >= zlb_value - 0.01 "interest_rate_forecast = $interest_rate_forecast must be >= zlb_value - 0.01 = $(zlb_value - 0.01)"
            end
        end
        return z_t, ϵ_t
    end

    # Iterate state space forward
    states = zeros(S, nstates, horizon)
    states[:, 1], shocks[:, 1] = iterate(z0, shocks[:, 1])
    for t in 2:horizon
        states[:, t], shocks[:, t] = iterate(states[:, t-1], shocks[:, t])
    end

    # Apply measurement and pseudo-measurement equations
    obs    = D .+ Z*states
    pseudo = D_pseudo .+ Z_pseudo * states

    # Return forecasts
    return states, obs, pseudo, shocks
end

function forecast(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S}, z0::Vector{S},
                  shocks::Matrix{S}; cond_type::Symbol = :none, enforce_zlb::Bool = false, ind_r::Int = -1,
                  ind_r_sh::Int = -1, zlb_value::S = 0.1/4) where {S<:AbstractFloat}

    # Determine how many regimes occur in the forecast, depending
    # on whether we need to subtract conditional forecast regimes or not
    n_fcast_reg = get_setting(m, :n_fcast_regimes)
    if cond_type != :none
        n_fcast_reg -= get_setting(m, :reg_post_conditional_end) - # Remove number of regimes switched since starting
        get_setting(m, :reg_forecast_start)                    # the conditional forecast
        if n_fcast_reg == 0 # Then no new regimes after conditional forecast ends
            n_fcast_reg = 1 # So we use the last conditional forecast regime
        end
    end

    Ts = Vector{Matrix{Float64}}(undef, n_fcast_reg)
    Rs = Vector{Matrix{Float64}}(undef, n_fcast_reg)
    Cs = Vector{Vector{Float64}}(undef, n_fcast_reg)
    Qs = Vector{Matrix{Float64}}(undef, n_fcast_reg)
    Zs = Vector{Matrix{Float64}}(undef, n_fcast_reg)
    Ds = Vector{Vector{Float64}}(undef, n_fcast_reg)
    Z_pseudos = Vector{Matrix{Float64}}(undef, n_fcast_reg)
    D_pseudos = Vector{Vector{Float64}}(undef, n_fcast_reg)

    # Determine in which regime the forecast starts, after accounting for
    # conditional forecast regimes, if applicable
    reg_fcast_cond_start = (cond_type == :none) ? get_setting(m, :reg_forecast_start) :
    max(get_setting(m, :reg_forecast_start), get_setting(m, :reg_post_conditional_end))

    # n_regimes is incorrectly being computed, as is n_fcast_reg
    # Unpack system
    for (ss_ind, sys_ind) in enumerate(reg_fcast_cond_start:get_setting(m, :n_regimes))
        Ts[ss_ind], Rs[ss_ind], Cs[ss_ind] = system[sys_ind, :TTT], system[sys_ind, :RRR], system[sys_ind, :CCC]
        Qs[ss_ind], Zs[ss_ind], Ds[ss_ind] = system[sys_ind, :QQ], system[sys_ind, :ZZ], system[sys_ind, :DD]
        Z_pseudos[ss_ind], D_pseudos[ss_ind] = system[sys_ind, :ZZ_pseudo], system[sys_ind, :DD_pseudo]
    end

    # Setup
    nshocks = size(Rs[1], 2)
    nstates = size(Ts[1], 2)
    nobs    = size(Zs[1], 1)
    npseudo = size(Z_pseudos[1], 1)
    horizon = size(shocks, 2)

    # Define our iteration function
    check_has_unant_mp_sh = (haskey(get_settings(m), :gensys2) ? get_setting(m, :gensys2) : false) &&
    (!isempty(intersect([:zlb_rule, :zero_rate], [x.alternative_policy.key for x in values(get_setting(m, :regime_eqcond_info))])))
    function iterate(z_t1, ϵ_t, T, R, C, Q, Z, D)
        z_t = C + T*z_t1 + R*ϵ_t

        # Change monetary policy shock to account for 0.13 interest rate bound
        if enforce_zlb
            interest_rate_forecast = getindex(D + Z*z_t, ind_r)
            if interest_rate_forecast < zlb_value
                continue_enforce = check_has_unant_mp_sh ? abs.(Z[ind_r, :]' * R[:, ind_r_sh]) > 1e-4 : true

                if continue_enforce
                    # Solve for interest rate shock causing interest rate forecast to be exactly ZLB
                    ϵ_t[ind_r_sh] = 0. # get forecast when MP shock
                    z_t = C + T*z_t1 + R*ϵ_t # is zeroed out
                    z_t_old = C + T*z_t1 + R*ϵ_t
                    ϵ_t[ind_r_sh] = getindex((zlb_value - D[ind_r] - Z[ind_r, :]'*z_t) / (Z[ind_r, :]' * R[:, ind_r_sh]), 1)

                    # Forecast again with new shocks
                    z_t = C + T*z_t1 + R*ϵ_t

                    # Confirm procedure worked
                    interest_rate_forecast = getindex(D + Z*z_t, ind_r)
                    if isnan(interest_rate_forecast)
                        ϵ_t[ind_r_sh] = 0. # get forecast when MP shock
                        z_t = C + T*z_t1 + R*ϵ_t # is zeroed out
                        z_t_old = C + T*z_t1 + R*ϵ_t
                        ϵ_t[ind_r_sh] = getindex((zlb_value - D[ind_r] - Z[ind_r, :]'*z_t) / (Z[ind_r, :]' * R[:, ind_r_sh]), 1)
                    end

                    # Subtract a small number to deal with numerical imprecision
                    @assert interest_rate_forecast >= zlb_value - 0.01 "interest_rate_forecast = $interest_rate_forecast must be >= zlb_value - 0.01 = $(zlb_value - 0.01)."
                end
            end
        end
        return z_t, ϵ_t
    end

    # These are the fcast regime indices
    regime_inds = get_fcast_regime_inds(m, horizon, cond_type)

    # Iterate state space forward
    states = zeros(nstates, horizon)
    obs = zeros(nobs, horizon)
    pseudo = zeros(npseudo, horizon)

    states[:, 1], shocks[:, 1] = iterate(z0, shocks[:, 1], Ts[1], Rs[1], Cs[1], Qs[1], Zs[1], Ds[1])
    obs[:, 1] = Ds[1] .+ Zs[1] * states[:, 1]
    pseudo[:, 1] = D_pseudos[1] .+ Z_pseudos[1] * states[:, 1]

    # If there's multiple regimes in forecast period, go through each set of indices. Otherwise, just take the first set
    if length(regime_inds) > 1
        for i in 1:length(regime_inds)
            ts = regime_inds[i]
            if length(ts) > 0 # Check in case the regime is empty
                if ts[1] == 1 # We already did this period in the initialization
                    ts = 2:ts[end]
                end
            end

            for t in ts
                states[:, t], shocks[:, t] = iterate(states[:, t - 1], shocks[:, t], Ts[i], Rs[i], Cs[i], Qs[i], Zs[i], Ds[i])
                obs[:, t] = Ds[i] .+ Zs[i] * states[:, t]
                pseudo[:, t] = D_pseudos[i] .+ Z_pseudos[i] * states[:, t]
            end
        end
    else
        for t in 2:horizon
            states[:, t], shocks[:, t] = iterate(states[:, t - 1], shocks[:, t], Ts[1], Rs[1], Cs[1], Qs[1], Zs[1], Ds[1])
            obs[:, t] = Ds[1] .+ Zs[1] * states[:, t]
            pseudo[:, t] = D_pseudos[1] .+ Z_pseudos[1] * states[:, t]
        end
    end

    # Return forecasts
    return states, obs, pseudo, shocks
end

# Automatic enforcing of the ZLB as a temporary alternative policy
function forecast(m::AbstractDSGEModel, z0::Vector{S}, states::AbstractMatrix{S},
                  obs::AbstractMatrix{S}, pseudo::AbstractMatrix{S}, shocks::AbstractMatrix{S};
                  cond_type::Symbol = :none, set_zlb_regime_vals::Function = identity,
                  set_info_sets_altpolicy::Function = auto_temp_altpolicy_info_set,
                  update_regime_eqcond_info!::Function =
                  (a, b, c, d) -> default_update_regime_eqcond_info!(a, b, c, d, alternative_policy(m)),
                  tol::S = -1e-5, rerun_smoother::Bool = false,
                  df = nothing, draw_states::Bool = false,
                  nan_failures::Bool = false,
                  histstates::AbstractMatrix{S} = Matrix{S}(undef, 0, 0),
                  histshocks::AbstractMatrix{S} = Matrix{S}(undef, 0, 0),
                  histpseudo::AbstractMatrix{S} = Matrix{S}(undef, 0, 0),
                  initial_states::AbstractVector{S} = Vector{S}(undef, 0)) where {S <: Real}
    # Grab "original" settings" so they can be restored later
    altpol = alternative_policy(m)
    is_regime_switch = haskey(get_settings(m), :regime_switching) ? get_setting(m, :regime_switching) : false
    is_replace_eqcond = haskey(get_settings(m), :replace_eqcond) ? get_setting(m, :replace_eqcond) : false
    is_gensys2 = haskey(get_settings(m), :gensys2) ? get_setting(m, :gensys2) : false
    info_set = haskey(get_settings(m), :tvis_information_set) ? get_setting(m, :tvis_information_set) : UnitRange{Int64}[]
    original_info_set = deepcopy(info_set)
    eqcond_dict = haskey(get_settings(m), :regime_eqcond_info) ? get_setting(m, :regime_eqcond_info) : Dict{Int, EqcondEntry}()
    original_eqcond_dict = deepcopy(eqcond_dict)
    orig_regimes      = haskey(get_settings(m), :n_regimes) ? get_setting(m, :n_regimes) : 1
    regime_dates = haskey(get_settings(m), :regime_dates) ? get_setting(m, :regime_dates) : Dict{Int, Date}()
    orig_regime_dates = deepcopy(regime_dates)
    orig_temp_zlb     = haskey(get_settings(m), :temporary_altpolicy_length) ? get_setting(m, :temporary_altpolicy_length) : nothing
    orig_identical_eqcond_regimes = haskey(get_settings(m), :identical_eqcond_regimes) ? deepcopy(get_setting(m, :identical_eqcond_regimes)) : nothing
    orig_perf_cred_identical_transitions = haskey(get_settings(m), :perfect_credibility_identical_transitions) ? deepcopy(get_setting(m, :perfect_credibility_identical_transitions)) : nothing
    orig_model2para_regime = haskey(get_settings(m), :model2para_regime) ? deepcopy(get_setting(m, :model2para_regime)) : nothing

    # Grab some information about the forecast
    n_hist_regimes = haskey(get_settings(m), :n_hist_regimes) ? get_setting(m, :n_hist_regimes) : 1
    has_reg_dates = haskey(get_settings(m), :regime_dates) # calculate dates of first & last regimes we need to add for the alt policy
    if has_reg_dates && length(get_setting(m, :regime_dates)) > 1
        start_altpol_reg  = get_setting(m, :reg_forecast_start)
        start_altpol_date = get_setting(m, :regime_dates)[start_altpol_reg]
        end_altpol_date   = date_forecast_end(m)
    else
        start_altpol_reg  = n_hist_regimes + 1
        start_altpol_date = date_forecast_start(m)
        end_altpol_date   = iterate_quarters(start_altpol_date, forecast_horizons(m; cond_type = cond_type) - 1)
    end
    altpol_reg_qtrrange   = quarter_range(start_altpol_date, end_altpol_date)
    n_altpol_reg_qtrrange = length(altpol_reg_qtrrange)

    # Determine the number of regimes before the start of the forecast (used to adjust
    # indexing)

    pre_fcast_regimes = n_hist_regimes
    #=
    if cond_type != :none
        if is_regime_switch
            pre_fcast_regimes += get_setting(m, :reg_post_conditional_end) - get_setting(m, :reg_forecast_start)
        else
            pre_fcast_regimes += subtract_quarters(get_setting(m, :date_conditional_end),
                                                   get_setting(m, :date_forecast_start)) + 1
        end
    end
    =#

    # Determine which quarters in the forecast have sub-zlb nominal rates
    has_neg_rates = view(obs, get_observables(m)[:obs_nominalrate], :) .<= forecast_zlb_value(m)

    first_endo_zlb     = findfirst(has_neg_rates)
    # Information set - start of awareness of ZLB
    first_aware = if haskey(get_settings(m), :tvis_information_set)
        findfirst([maximum(info_set[i]) == get_setting(m, :n_regimes) for i in keys(info_set)])
    elseif haskey(get_settings(m), :regime_eqcond_info)
        findfirst([info_set[i].key == :zlb_rule for i in keys(info_set)])
    else
        first_endo_zlb
    end

    ## 1. Determine if we need to do anything (are there any further negative nominal rates)
    # If not, return the forecast as is
    if isnothing(first_endo_zlb)
        if rerun_smoother
            return states, obs, pseudo, histstates, histshocks, histpseudo, initial_states
        else
            return states, obs, pseudo
        end
        # If there are negative nominal rates in the forecast, we continue to computing the endogenous
        # zlb forecast
    else
        min_zlb = haskey(m.settings, :min_temporary_altpolicy_length) ? get_setting(m, :min_temporary_altpolicy_length) : 0
        first_endo_zlb = max(first_endo_zlb, min_zlb + 1)
        last_endo_zlb = findfirst(.!has_neg_rates[first_endo_zlb:end])
        first_endo_zlb += pre_fcast_regimes

        last_endo_zlb = !isnothing(last_endo_zlb) ? last_endo_zlb-1 + first_endo_zlb : size(obs, 2) -3

        max_zlb_regimes = haskey(get_settings(m), :max_temporary_altpolicy_length) ? get_setting(m, :max_temporary_altpolicy_length) : size(obs, 2) - 3
        last_endo_zlb = min(last_endo_zlb, max_zlb_regimes)

        # set up the information sets TODO: add checkfor whether or not we even need to update the tvis_info_set
        if isnothing(first_aware)
            first_aware = get_setting(m, :reg_forecast_start)
        end
        set_info_sets_altpolicy(m, get_setting(m, :n_regimes), first_aware)

        endozlb_forecast::Function = (zlb_start, zlb_end; unant_enforce_zlb = false) -> forecast_endozlb_helper(m, zlb_start, zlb_end, z0, states, shocks,
            orig_regimes, eqcond_dict, regime_dates, info_set;
            cond_type = cond_type, unant_enforce_zlb = unant_enforce_zlb,
            set_zlb_regime_vals = set_zlb_regime_vals,
            set_info_sets_altpolicy = set_info_sets_altpolicy,
            update_regime_eqcond_info! = update_regime_eqcond_info!,
            tol = tol, rerun_smoother = rerun_smoother, df = df,
            draw_states = draw_states, histstates = histstates,
            histshocks = histshocks, histpseudo = histpseudo,
            initial_states = initial_states)
        ## 2. If we don't liftoff post minimum zlb, then enforce the rest of the contiguous zlb
        ##    intervals using two rule (zlb_rule alternative policy)
        if !has_neg_rates[min_zlb+1]
            endo_success = all(obs[get_observables(m)[:obs_nominalrate], :] .> get_setting(m, :zlb_rule_value) / 4. + tol)
            if !endo_success
                system = compute_system(m; tvis = true)
                states, obs, pseudo = forecast(m, system, z0; cond_type = cond_type, shocks = shocks,
                                               enforce_zlb = true)
            end
            if rerun_smoother
                return states, obs, pseudo, histstates, histshocks, histpseudo, initial_states
            else
                return states, obs, pseudo
            end
        else
            # Ensure that the Eqcond Dict is correctly cleared for endo zlb
            temp_altpol_length = 0
            if haskey(m.settings, :regime_eqcond_info)
                for (reg, v) in get_setting(m, :regime_eqcond_info)
                    if v.alternative_policy.key == :zlb_rule || v.alternative_policy.key == :zero_rate
                        if reg >= first_endo_zlb
                            @warn "Regime $reg of regime_eqcond_info used zero rate--to avoid gensys errors in computing the endogenous zlb, this regime is now being set to use $(altpol.key)."
                            v.alternative_policy = altpol
                        else
                            temp_altpol_length += 1
                        end
                    end
                end
            end
            m <= Setting(:temporary_altpolicy_length, temp_altpol_length)

            # Run extended zlb forecast
            states, obs, pseudo, histstates, histshocks, histpseudo, initial_states =
            endozlb_forecast(first_endo_zlb, last_endo_zlb+1, unant_enforce_zlb = false)
            ## 3. Check if this forecast lifts off after the end of the two-rule zlb.
            liftoff = obs[get_observables(m)[:obs_nominalrate], last_endo_zlb-pre_fcast_regimes+1] .> forecast_zlb_value(m)+tol

            ## If it does, binary search between min_zlb and last_endo_zlb to find the earliest liftoff
            if liftoff
                high = last_endo_zlb
                low  = first_endo_zlb
                lookback = haskey(m.settings, :endogenous_zlb_lookback) ? get_setting(m, :endogenous_zlb_lookback)-1 : 2
                iter = max(last_endo_zlb - lookback, first_endo_zlb)

                while true
                    # Run extended zlb forecast
                    states, obs, pseudo, histstates, histshocks, histpseudo, initial_states =
                    endozlb_forecast(first_endo_zlb, iter+1, unant_enforce_zlb = false)
                    ## 3. Check if this forecast lifts off after the end of the two-rule zlb.
                    liftoff = obs[get_observables(m)[:obs_nominalrate], iter-pre_fcast_regimes+1] .> get_setting(m, :zlb_rule_value) / 4. + tol
                    if !liftoff
                        low = iter + 1
                        if iter == last_endo_zlb
                            # We have hit the max number of regimes for this leg of the ZLB, so we've
                            # run Fixed ZLB with unanticipated shocks to enforce ZLB
                            break
                        elseif low-1 == high
                            # We should never reach this point with endo_success false or the
                            # forecast unenforced with unant shocks, but just in case...
                            # (if we do end up here, the forecast will be NaNed)
                            break
                        elseif iter < last_endo_zlb && iter >= last_endo_zlb - lookback
                            # If our initial guess did work but lookback didn't, check ahead
                            # stepwise
                            iter = iter + 1
                        else
                            # Continue Binary Search
                            iter = floor(Int64, (low + high) / 2)
                        end
                    else
                        # If ZLB works (we lift off after the first consecutive stretch)
                        high = iter # Set high to iter so we search for ZLB regimes less than or equal to iter
                        if low == high || iter == low
                            break
                        elseif iter == last_endo_zlb
                            iter = last_endo_zlb - lookback
                        elseif iter == last_endo_zlb - lookback
                            # If first_guess - 2 works, let's try the first possible ZLB regime value
                            iter = first_endo_zlb
                        elseif iter == first_endo_zlb
                            # Minimal number of ZLB works, so we break here
                            # and use iter = first_zlb_regime
                            break
                        elseif iter < last_endo_zlb  && iter > last_endo_zlb - lookback
                            # if we're in this block, last_zlb_reg worked but first_guess -
                            # lookback didn't, so we're proceeding stepwise forward and pick the
                            # first regime that works
                            break
                        else
                            # Continue Binary Search
                            iter = floor(Int64, (low + high) / 2)
                        end
                    end
                end

                last_endo_zlb = iter

                ## If it does not, iterate forward until it does.
            else
                while !liftoff && last_endo_zlb < max_zlb_regimes
                    last_endo_zlb += 1
                    # Run extended zlb forecast
                    states, obs, pseudo, histstates, histshocks, histpseudo, initial_states =
                    endozlb_forecast(first_endo_zlb, last_endo_zlb+1, unant_enforce_zlb = false)
                    ## 3. Check if this forecast lifts off after the end of the two-rule zlb.
                    liftoff = obs[get_observables(m)[:obs_nominalrate], last_endo_zlb-pre_fcast_regimes+1] .> get_setting(m, :zlb_rule_value) / 4. + tol
                end
            end
        end
        ## 6. If we still have further noncontiguous periods of negative rates, run a forecast enforcing zlb over those remaining periods
        ##    using unanticipated monetary policy shocks.
        endo_success = all(obs[get_observables(m)[:obs_nominalrate], :] .> get_setting(m, :zlb_rule_value) / 4. + tol)
        if !endo_success
            states, obs, pseudo, histstates, histshocks, histpseudo, initial_states =
            endozlb_forecast(first_endo_zlb, last_endo_zlb+1, unant_enforce_zlb = true)
        end
        # Restore original settings
        m <= Setting(:regime_switching, is_regime_switch)
        m <= Setting(:regime_dates, orig_regime_dates)
        if !isempty(orig_regime_dates)
            setup_regime_switching_inds!(m; cond_type = cond_type)
        end
        m <= Setting(:replace_eqcond,   is_replace_eqcond)
        m <= Setting(:gensys2,          is_gensys2)
        if !is_replace_eqcond
            delete!(get_settings(m), :regime_eqcond_info)
        end
        if isempty(original_eqcond_dict)
            delete!(get_settings(m), :regime_eqcond_info)
        else
            m <= Setting(:regime_eqcond_info, original_eqcond_dict)
        end
        if haskey(get_settings(m), :tvis_information_set)
            m <= Setting(:tvis_information_set, original_info_set)
        end
        if isnothing(orig_temp_zlb)
            delete!(get_settings(m), :temporary_altpolicy_length)
        else
            m <= Setting(:temporary_altpolicy_length, orig_temp_zlb)
        end
        if isnothing(orig_identical_eqcond_regimes)
            delete!(get_settings(m), :identical_eqcond_regimes)
        else
            m <= Setting(:identical_eqcond_regimes, orig_identical_eqcond_regimes)
        end
        if isnothing(orig_perf_cred_identical_transitions)
            delete!(get_settings(m), :perfect_credibility_identical_transitions)
        else
            m <= Setting(:perfect_credibility_identical_transitions, orig_perf_cred_identical_transitions)
        end
        if isnothing(orig_model2para_regime)
            delete!(get_settings(m), :model2para_regime)
        else
            m <= Setting(:model2para_regime, orig_model2para_regime)
        end
        if rerun_smoother
            return states, obs, pseudo, histstates, histshocks, histpseudo, initial_states
        else
            return states, obs, pseudo
        end
    end
end

function forecast_endozlb_helper(m::AbstractDSGEModel, first_endo_zlb::Int64, liftoff_reg::Int64,
                                 z0::Vector{S}, states::AbstractMatrix{S}, shocks::AbstractMatrix{S},
                                 orig_regimes::Int64, original_eqcond_dict::Dict{Int64, EqcondEntry},
                                 orig_regime_dates::Dict{Int, Date}, original_info_set::Union{UnitRange{Int64}, Array{UnitRange{Int64}}};
                                 unant_enforce_zlb::Bool = false, cond_type::Symbol = :none, set_zlb_regime_vals::Function = identity,
                                 set_info_sets_altpolicy::Function = auto_temp_altpolicy_info_set,
                                 update_regime_eqcond_info!::Function =
                                 (a, b, c, d) -> default_update_regime_eqcond_info!(a, b, c, d, alternative_policy(m)),
                                 tol::S = -1e-5, rerun_smoother::Bool = false,
                                 df = nothing, draw_states::Bool = false,
                                 histstates::AbstractMatrix{S} = Matrix{S}(undef, 0, 0),
                                 histshocks::AbstractMatrix{S} = Matrix{S}(undef, 0, 0),
                                 histpseudo::AbstractMatrix{S} = Matrix{S}(undef, 0, 0),
                                 initial_states::AbstractVector{S} = Vector{S}(undef, 0)) where {S <: Real}

    altpol = alternative_policy(m)
    is_regime_switch = haskey(get_settings(m), :regime_switching) ? get_setting(m, :regime_switching) : false
    is_replace_eqcond = haskey(get_settings(m), :replace_eqcond) ? get_setting(m, :replace_eqcond) : false
    is_gensys2 = haskey(get_settings(m), :gensys2) ? get_setting(m, :gensys2) : false
    # Grab some information about the forecast
    n_hist_regimes = haskey(get_settings(m), :n_hist_regimes) ? get_setting(m, :n_hist_regimes) : 1
    has_reg_dates = haskey(get_settings(m), :regime_dates) # calculate dates of first & last regimes we need to add for the alt policy
    if has_reg_dates && length(get_setting(m, :regime_dates)) > 1
        start_altpol_reg  = get_setting(m, :reg_forecast_start)
        start_altpol_date = get_setting(m, :regime_dates)[start_altpol_reg]
        end_altpol_date   = date_forecast_end(m)
    else
        start_altpol_reg  = n_hist_regimes + 1
        start_altpol_date = date_forecast_start(m)
        end_altpol_date   = iterate_quarters(start_altpol_date, forecast_horizons(m; cond_type = cond_type) - 1)
    end
    altpol_reg_qtrrange   = quarter_range(start_altpol_date, end_altpol_date)
    n_altpol_reg_qtrrange = length(altpol_reg_qtrrange)
    orig_temp_zlb     = haskey(get_settings(m), :temporary_altpolicy_length) ? get_setting(m, :temporary_altpolicy_length) : nothing

    # Set up regime dates
    altpol_regime_dates = Dict{Int, Date}(1 => date_presample_start(m))
    if is_regime_switch # Add historical regimes
        for regind in 2:n_hist_regimes
            altpol_regime_dates[regind] = orig_regime_dates[regind]
        end
    end

    altpol_reg_range = start_altpol_reg:liftoff_reg
    for (regind, date) in zip(altpol_reg_range, altpol_reg_qtrrange)
        altpol_regime_dates[regind] = date
    end

    # in the case where we're shortening the zlb, this is necessary to stop the
    # setup_regime_switching_inds! call from truncating on a zlb reg and throwing
    # a gensys error
    if haskey(get_settings(m), :regime_eqcond_info) ? haskey(get_setting(m, :regime_eqcond_info), liftoff_reg) : false
        get_setting(m, :regime_eqcond_info)[liftoff_reg].alternative_policy = alternative_policy(m)
    end

    if haskey(get_settings(m), :cred_vary_until) ? (isempty(altpol_reg_range) ? true : get_setting(m, :cred_vary_until) >= maximum(altpol_reg_range)) : false
        m <= Setting(:regime_switching, true)
    else
        # NOTE: the following will truncate regime_eqcond_info after liftoff regime unless cred_vary_until
        # is set. At present, this DOES NOT HANDLE cases where there are more policy changes after the
        # liftoff regime. If this is the case, use cred_vary_until up to the end of those changes.
        if haskey(get_settings(m), :regime_eqcond_info)
            for i in keys(get_setting(m, :regime_eqcond_info))
                if i > liftoff_reg
                    delete!(get_setting(m, :regime_eqcond_info), i)
                end
            end
        end
        m <= Setting(:regime_dates, altpol_regime_dates)
        m <= Setting(:regime_switching, true)
        setup_regime_switching_inds!(m; cond_type = cond_type)
    end

    # Set up replace_eqcond entries
    m <= Setting(:replace_eqcond, true)
    m <= Setting(:gensys2, true)

    # In general, the update_regime_eqcond_info! function should not change the
    # EqcondEntry during historical/conditional horizon regimes, but any other
    # regimes in the forecast horizon should be/can be set.
    #
    # From above, we also assume that no ZLBs occur in the eqcond_dict that
    # is passed to update_regime_eqcond_info! and assume that
    # if there is a particular sequence of alternative policies the user
    # wants to occur after the ZLB, then the user will have incorporated
    # code that implements that sequence in their own update_regime_eqcond_info!
    # function. In the default DSGE policy, the regimes after the ZLB ends
    # are updated only if there is time-varying credibility
    # (specified by the Setting :cred_vary_until).
    update_regime_eqcond_info!(m, original_eqcond_dict, first_endo_zlb, liftoff_reg)
    min_zlb = haskey(m.settings, :min_temporary_altpolicy_length) ? get_setting(m, :min_temporary_altpolicy_length) : 0
    min_zlb += haskey(m.settings, :historical_temporary_altpolicy_length) ? get_setting(m, :historical_temporary_altpolicy_length) : 0
    m <= Setting(:temporary_altpolicy_length, min_zlb + liftoff_reg-first_endo_zlb)

    # Set up parameters if there are switching parameter values.
    #
    # User needs to provide a function which takes in the model object `m`
    # and the total number of regimes (after adding the required temporary regimes),
    # and sets up regime-switching parameters for these new additional regimes.
    if set_zlb_regime_vals != identity
        set_zlb_regime_vals(m, liftoff_reg)
    end
    # if we have a mapping between model regimes and parameter regimes, make sure to
    # extend that mapping if we've added regimes in the process of running the endo zlb
    if haskey(m.settings, :model2para_regime)
        for para in m.parameters
            model2para_dict = get_setting(m, :model2para_regime)
            if haskey(model2para_dict, para.key)
                for i in first_endo_zlb:liftoff_reg
                    if !haskey(model2para_dict[para.key], i)
                        model2para_dic[para.key][i] = model2para_dict[para.key][i-1]
                    end
                end
            end
        end
    end

    # set up the information sets TODO: add checkfor whether or not we even need to update the tvis_info_set
    first_aware = if haskey(get_settings(m), :tvis_information_set)
        findfirst([maximum(original_info_set[i]) == get_setting(m, :n_regimes) for i in keys(original_info_set)])
    elseif haskey(get_settings(m), :regime_eqcond_info)
        findfirst([original_info_set[i].key == :zlb_rule for i in keys(original_info_set)])
    else
        nothing
    end
    if isnothing(first_aware)
        first_aware = first_endo_zlb
    end
    set_info_sets_altpolicy(m, get_setting(m, :n_regimes), first_aware)

    # TODO: maybe delete tvis, since tvis should always be on given the fact we always run set_info_sets
    # tvis = haskey(get_settings(m), :tvis_information_set) && !isempty(get_setting(m, :tvis_information_set))

    # Adjust the identical perfect credibility transitions for the new zlb
    # Note: the following assumes that there are no more policy changes after the switch to flexible ait at
    # the end of the zlb, and that all zlb regimes enforced with eqcond are consecutive.
    # If either of these are no longer true, this will break!
    if haskey(m.settings, :perfect_credibility_identical_transitions)
        for reg in first_endo_zlb:get_setting(m, :n_regimes)
            delete!(get_setting(m, :perfect_credibility_identical_transitions), reg)
        end
        for reg in liftoff_reg+2:get_setting(m, :n_regimes)
            get_setting(m, :perfect_credibility_identical_transitions)[reg] = liftoff_reg+1
        end
    end
    if haskey(m.settings, :identical_eqcond_regimes)
        for reg in first_endo_zlb:get_setting(m, :n_regimes)
            delete!(get_setting(m, :identical_eqcond_regimes), reg)
        end
        first_zlb_reg = minimum(collect(keys(get_setting(m, :regime_eqcond_info))))
        for reg in first_zlb_reg+1:liftoff_reg-1
            get_setting(m, :identical_eqcond_regimes)[reg] = first_zlb_reg
        end
        for reg in liftoff_reg+2:get_setting(m, :n_regimes)
            get_setting(m, :identical_eqcond_regimes)[reg] = liftoff_reg+1
        end
    end

    # Recompute to account for new regimes
    system = compute_system(m; tvis = true)
    if rerun_smoother # if state space system changes, then the smoothed states will also change generally
        histstates, histshocks, histpseudo, initial_states =
        smooth(m, df, system; cond_type = cond_type, draw_states = draw_states)
        z0 = histstates[:, end]
    end

    # Forecast! Note that since forecast(m, system, z0) allocates new matrices for states, obs, and pseudo,
    # this step merely changes to which matrices the local variables states, obs, and pseudo point
    # w/in this function's closure. Thus, the matrices states, obs, and pseudo which are first
    # passed into this function will not be over-written.
    states, obs, pseudo = forecast(m, system, z0; cond_type = cond_type, shocks = shocks,
                                   enforce_zlb = unant_enforce_zlb)

    # Delete extra regimes added to implement the temporary alternative policy, or else updating the parameters
    # in forecast_one will not work.
    if set_zlb_regime_vals != identity
        for p in m.parameters
            if haskey(p.regimes, :value)
                if length(p.regimes[:value]) > orig_regimes
                    for i in (orig_regimes + 1):length(p.regimes[:value])
                        delete!(p.regimes[:value], i)
                    end
                end
            end
        end
    end

    return states, obs, pseudo, histstates, histshocks, histpseudo, initial_states
end

function get_fcast_regime_inds(m::AbstractDSGEModel, horizon::Int, cond_type::Symbol;
                               start_index::Int = 0)

    fcast_post_cond_date = (cond_type == :none) ? date_forecast_start(m) : # If cond forecast, then we want to start forecasting
    max(date_forecast_start(m), iterate_quarters(date_conditional_end(m), 1)) # from the period after the conditional end period
    # last_date = iterate_quarters(fcast_post_cond_date, -1) # minus 1 to get one date before the forecast start
    last_date = fcast_post_cond_date
    last_ind = 1

    # Construct vector of time periods for each regime in the forecast periods after the conditional forecast
    start_reg   = (cond_type == :none) ? get_setting(m, :reg_forecast_start) :
    max(get_setting(m, :reg_forecast_start), get_setting(m, :reg_post_conditional_end))
    regime_inds = Vector{UnitRange{Int}}(undef, get_setting(m, :n_regimes) - start_reg + 1)
    for (rgi, i) in enumerate(start_reg:(get_setting(m, :n_regimes) - 1)) # Last regime handled separately
        qtr_diff = subtract_quarters(get_setting(m, :regime_dates)[i + 1], last_date)
        regime_inds[rgi] = (start_index + last_ind):(start_index + last_ind + qtr_diff - 1)
        last_ind += qtr_diff
        last_date = get_setting(m, :regime_dates)[i + 1]
    end
    regime_inds[end] = (start_index + last_ind):(start_index + horizon) # Covers case where final hist regime is also first (and only) forecast regime.

    return regime_inds
end

function auto_temp_altpolicy_info_set(m::AbstractDSGEModel, n_regimes::Int, first_zlb_regime::Int)
    tvis_infoset = Vector{UnitRange{Int}}(undef, n_regimes)
    for r in 1:n_regimes
        if r < first_zlb_regime
            tvis_infoset[r] = r:r
        else
            tvis_infoset[r] = r:n_regimes
        end
    end
    m <= Setting(:tvis_information_set, tvis_infoset)
end

"""
```
default_update_regime_eqcond_info!(m::AbstractDSGEModel, eqcond_dict::AbstractDict{Int64, EqcondEntry},
                                   zlb_start_regime::Int64, liftoff_regime::Int64, liftoff_policy::AltPolicy)
```
is the default method for updating the setting `regime_eqcond_info` during the automatic
endogenous ZLB enforcement as a temporary policy.

The input `eqcond_dict` should be a dictionary whose historical/conditional horizon regimes
should already be set (and won't be affected by this function.
The regimes which will be altered are those for which the temporary ZLB will apply
(according to `zlb_start_regime` and `liftoff_regime`), and if there is
time-varying credibility, forecast regimes past the ZLB.

In general, other implementations of the `update_regime_eqcond_info!` function should not change the
EqcondEntry during historical/conditional horizon regimes, but any other
regimes in the forecast horizon should be/can be set.
"""
function default_update_regime_eqcond_info!(m::AbstractDSGEModel, eqcond_dict::AbstractDict{Int64, EqcondEntry},
                                            zlb_start_regime::Int64, liftoff_regime::Int64, liftoff_policy::AltPolicy)

    # TODO: add a check for making sure weights are properly set (e.g. enough regimes with weights to avoid error)
    # Populate eqcond_dict for regimes in which the ZLB should apply and update the weights as appropriate
    for regind in zlb_start_regime:(liftoff_regime - 1)
        if haskey(eqcond_dict, regind)
            # If regime index exists already, just directly change the
            # alternative_policy (which leaves the credibility weights unchanged)
            eqcond_dict[regind].alternative_policy = zlb_rule() # Temp ZLB rule in this regimes
        else
            # If the regime index doesn't already exist, then add an EqcondEntry
            # for a ZLB that is perfectly credible
            weights = zeros(haskey(get_settings(m), :alternative_policies) ? length(get_setting(m, :alternative_policies)) + 1 : 1)
            weights[1] = 1.
            eqcond_dict[regind] = EqcondEntry(zlb_rule(), weights)
        end
    end

    # Set the liftoff regime's EqcondEntry, using same logic as above
    if haskey(eqcond_dict, liftoff_regime)
        eqcond_dict[liftoff_regime].alternative_policy = liftoff_policy
    else
        weights = zeros(haskey(m.settings, :alternative_policies) ? length(get_setting(m, :alternative_policies)) + 1 : 1)
        weights[1] = 1.
        eqcond_dict[liftoff_regime] = EqcondEntry(liftoff_policy, weights)
    end

    # NOTE: the following block assumes there is only one temporary altpol
    # Update eqcond_dict to use the lift-off policy for any regimes between (inclusive)
    # the liftoff regime and the first regime for which credibility stops varying.
    if haskey(get_settings(m), :cred_vary_until) && get_setting(m, :cred_vary_until) >= liftoff_regime
        for z in liftoff_regime:(get_setting(m, :cred_vary_until))
            eqcond_dict[z].alternative_policy = liftoff_policy
        end
    end

    m <= Setting(:regime_eqcond_info, eqcond_dict)
    m <= Setting(:n_regimes, maximum(collect(keys(eqcond_dict))))

    return m
end


using IICGP
using Dates

rootdir = joinpath(homedir(), "Documents/git/ICGP-results/")
# rootdir = dirname(@__DIR__)
resdir = joinpath(rootdir, "results/")
#fetch_backup(rootdir, clean=true)

# Filters parameters
min_date = DateTime(2021, 11, 30, 14)
max_date = DateTime(2021, 11, 30, 15)
#savedir_index = 1
labels = Vector{String}()
colors = Vector{Symbol}()
rom_names = ["boxing"]
reducers = ["pooling"]
dotime = false
dosave = true
baselines = true

objectives_names = ["Atari score", "Atari score (≠ seed)"]
pareto_gen = [1, 9000, 17000, 25000]
pareto_xlim = (0, 1)
pareto_ylim = (0, 1)

colors = [:skyblue3]
labels = ["NSGA2"]

for rom_name in rom_names
    exp_dirs, games = get_exp_dir(resdir, min_date=min_date, max_date=max_date,
        games=[rom_name], reducers=reducers)
    process_nsga2_results(exp_dirs, games, objectives_names, colors,
        labels, pareto_gen=pareto_gen, pareto_xlim=pareto_xlim,
        pareto_ylim=pareto_ylim)
end


# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

function rhs!(du, u, t,
              mesh::ParallelP4estMesh{3}, equations,
              initial_condition, boundary_conditions, source_terms,
              dg::DG, cache)
  # Start to receive MPI data
  @trixi_timeit timer() "start MPI receive" start_mpi_receive!(cache.mpi_cache)

  # Prolong solution to MPI interfaces
  @trixi_timeit timer() "prolong2mpiinterfaces" prolong2mpiinterfaces!(
    cache, u, mesh, equations, dg.surface_integral, dg)

  # Start to send MPI data
  @trixi_timeit timer() "start MPI send" start_mpi_send!(
    cache.mpi_cache, mesh, equations, dg, cache)

  # Reset du
  @trixi_timeit timer() "reset ∂u/∂t" reset_du!(du, dg, cache)

  # Calculate volume integral
  @trixi_timeit timer() "volume integral" calc_volume_integral!(
    du, u, mesh,
    have_nonconservative_terms(equations), equations,
    dg.volume_integral, dg, cache)

  # Prolong solution to interfaces
  @trixi_timeit timer() "prolong2interfaces" prolong2interfaces!(
    cache, u, mesh, equations, dg.surface_integral, dg)

  # Calculate interface fluxes
  @trixi_timeit timer() "interface flux" calc_interface_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg.surface_integral, dg, cache)

  # Prolong solution to boundaries
  @trixi_timeit timer() "prolong2boundaries" prolong2boundaries!(
    cache, u, mesh, equations, dg.surface_integral, dg)

  # Calculate boundary fluxes
  @trixi_timeit timer() "boundary flux" calc_boundary_flux!(
    cache, t, boundary_conditions, mesh, equations, dg.surface_integral, dg)

  # Prolong solution to mortars
  @trixi_timeit timer() "prolong2mortars" prolong2mortars!(
    cache, u, mesh, equations, dg.mortar, dg.surface_integral, dg)

  # Calculate mortar fluxes
  @trixi_timeit timer() "mortar flux" calc_mortar_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg.mortar, dg.surface_integral, dg, cache)

  # Finish to receive MPI data
  @trixi_timeit timer() "finish MPI receive" finish_mpi_receive!(
    cache.mpi_cache, mesh, equations, dg, cache)

  # Calculate MPI interface fluxes
  @trixi_timeit timer() "MPI interface flux" calc_mpi_interface_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg.surface_integral, dg, cache)

  # Calculate surface integrals
  @trixi_timeit timer() "surface integral" calc_surface_integral!(
    du, u, mesh, equations, dg.surface_integral, dg, cache)

  # Apply Jacobian from mapping to reference element
  @trixi_timeit timer() "Jacobian" apply_jacobian!(
    du, mesh, equations, dg, cache)

  # Calculate source terms
  @trixi_timeit timer() "source terms" calc_sources!(
    du, u, t, source_terms, equations, dg, cache)

  # Finish to send MPI data
  @trixi_timeit timer() "finish MPI send" finish_mpi_send!(cache.mpi_cache)

  return nothing
end

function prolong2mpiinterfaces!(cache, u,
                                mesh::ParallelP4estMesh{3},
                                equations, surface_integral, dg::DG)
  @unpack mpi_interfaces = cache
  index_range = eachnode(dg)

  @threaded for interface in eachmpiinterface(dg, cache)
    # Copy solution data from the local element using "delayed indexing" with
    # a start value and a step size to get the correct face and orientation.
    # Note that in the current implementation, the interface will be
    # "aligned at the primary element", i.e., the index of the primary side
    # will always run forwards.
    local_side = mpi_interfaces.local_sides[interface]
    local_element = mpi_interfaces.local_element_ids[interface]
    local_indices = mpi_interfaces.node_indices[interface]

    i_element_start, i_element_step_i, i_element_step_j = index_to_start_step_3d(local_indices[1], index_range)
    j_element_start, j_element_step_i, j_element_step_j = index_to_start_step_3d(local_indices[2], index_range)
    k_element_start, k_element_step_i, k_element_step_j = index_to_start_step_3d(local_indices[3], index_range)

    i_element = i_element_start
    j_element = j_element_start
    k_element = k_element_start
    for j in eachnode(dg)
      for i in eachnode(dg)
        for v in eachvariable(equations)
          mpi_interfaces.u[local_side, v, i, j, interface] = u[v, i_element, j_element, k_element, local_element]
        end
        i_element += i_element_step_i
        j_element += j_element_step_i
        k_element += k_element_step_i
      end
      i_element += i_element_step_j
      j_element += j_element_step_j
      k_element += k_element_step_j
    end
  end

  return nothing
end

function calc_mpi_interface_flux!(surface_flux_values,
                                  mesh::ParallelP4estMesh{3},
                                  nonconservative_terms,
                                  equations, surface_integral, dg::DG, cache)
  @unpack local_element_ids, node_indices, local_sides = cache.mpi_interfaces
  @unpack contravariant_vectors = cache.elements
  index_range = eachnode(dg)

  @threaded for interface in eachmpiinterface(dg, cache)
    # Get element and side index information on the local element
    local_element = local_element_ids[interface]
    local_indices = node_indices[interface]
    local_direction = indices2direction(local_indices)
    local_side = local_sides[interface]

    # Create the local i,j,k indexing on the local element used to pull normal direction information
    i_element_start, i_element_step_i, i_element_step_j = index_to_start_step_3d(local_indices[1], index_range)
    j_element_start, j_element_step_i, j_element_step_j = index_to_start_step_3d(local_indices[2], index_range)
    k_element_start, k_element_step_i, k_element_step_j = index_to_start_step_3d(local_indices[3], index_range)

    i_element = i_element_start
    j_element = j_element_start
    k_element = k_element_start

    # Initiate the node indices to be used in the surface for loop,
    # the surface flux storage must be indexed in alignment with the local element indexing
    local_surface_indices = surface_indices(local_indices)
    i_surface_start, i_surface_step_i, i_surface_step_j = index_to_start_step_3d(local_surface_indices[1], index_range)
    j_surface_start, j_surface_step_i, j_surface_step_j = index_to_start_step_3d(local_surface_indices[2], index_range)
    i_surface = i_surface_start
    j_surface = j_surface_start

    for j in eachnode(dg)
      for i in eachnode(dg)
        # Get the normal direction on the local element
        # Contravariant vectors at interfaces in negative coordinate direction
        # are pointing inwards. This is handled by `get_normal_direction`.
        normal_direction = get_normal_direction(local_direction, contravariant_vectors,
                                                i_element, j_element, k_element,
                                                local_element)

        calc_mpi_interface_flux!(surface_flux_values, mesh, nonconservative_terms, equations,
                                 surface_integral, dg, cache,
                                 interface, normal_direction,
                                 i, j, local_side,
                                 i_surface, j_surface, local_direction, local_element)

        # Increment local element indices to pull the normal direction
        i_element += i_element_step_i
        j_element += j_element_step_i
        k_element += k_element_step_i
        # Increment the surface node indices along the local element
        i_surface += i_surface_step_i
        j_surface += j_surface_step_i
      end
      # Increment local element indices to pull the normal direction
      i_element += i_element_step_j
      j_element += j_element_step_j
      k_element += k_element_step_j
      # Increment the surface node indices along the local element
      i_surface += i_surface_step_j
      j_surface += j_surface_step_j
    end
  end

  return nothing
end

# Inlined version of the interface flux computation for conservation laws
@inline function calc_mpi_interface_flux!(surface_flux_values,
                                          mesh::P4estMesh{3},
                                          nonconservative_terms::Val{false}, equations,
                                          surface_integral, dg::DG, cache,
                                          interface_index, normal_direction,
                                          interface_i_node_index, interface_j_node_index, local_side,
                                          surface_i_node_index, surface_j_node_index,
                                          local_direction_index, local_element_index)
  @unpack u = cache.mpi_interfaces
  @unpack surface_flux = surface_integral

  u_ll, u_rr = get_surface_node_vars(u, equations, dg,
                                     interface_i_node_index, interface_j_node_index, interface_index)

  if local_side == 1
    flux_ = surface_flux(u_ll, u_rr, normal_direction, equations)
  else # local_side == 2
    flux_ = -surface_flux(u_ll, u_rr, -normal_direction, equations)
  end

  for v in eachvariable(equations)
    surface_flux_values[v, surface_i_node_index, surface_j_node_index,
                        local_direction_index, local_element_index] = flux_[v]
  end
end

end # muladd


# This file is a part of Julia. License is MIT: https://julialang.org/license

## UMFPACK

## Type of solve
const UMFPACK_A     =  0     # Ax=b
const UMFPACK_At    =  1     # adjoint(A)x=b
const UMFPACK_Aat   =  2     # transpose(A)x=b
const UMFPACK_Pt_L  =  3     # adjoint(P)Lx=b
const UMFPACK_L     =  4     # Lx=b
const UMFPACK_Lt_P  =  5     # adjoint(L)Px=b
const UMFPACK_Lat_P =  6     # transpose(L)Px=b
const UMFPACK_Lt    =  7     # adjoint(L)x=b
const UMFPACK_Lat   =  8     # transpose(L)x=b
const UMFPACK_U_Qt  =  9     # U*adjoint(Q)x=b
const UMFPACK_U     =  10    # Ux=b
const UMFPACK_Q_Ut  =  11    # Q*adjoint(U)x=b
const UMFPACK_Q_Uat =  12    # Q*transpose(U)x=b
const UMFPACK_Ut    =  13    # adjoint(U)x=b
const UMFPACK_Uat   =  14    # transpose(U)x=b

## Sizes of Control and Info arrays for returning information from solver
const UMFPACK_INFO = 90
const UMFPACK_CONTROL = 20
const UMFPACK_PRL = 1

## Status codes
const UMFPACK_OK = 0
const UMFPACK_WARNING_singular_matrix       = 1
const UMFPACK_WARNING_determinant_underflow = 2
const UMFPACK_WARNING_determinant_overflow  = 3
const UMFPACK_ERROR_out_of_memory           = -1
const UMFPACK_ERROR_invalid_Numeric_object  = -3
const UMFPACK_ERROR_invalid_Symbolic_object = -4
const UMFPACK_ERROR_argument_missing        = -5
const UMFPACK_ERROR_n_nonpositive           = -6
const UMFPACK_ERROR_invalid_matrix          = -8
const UMFPACK_ERROR_different_pattern       = -11
const UMFPACK_ERROR_invalid_system          = -13
const UMFPACK_ERROR_invalid_permutation     = -15
const UMFPACK_ERROR_internal_error          = -911
const UMFPACK_ERROR_file_IO                 = -17
const UMFPACK_ERROR_ordering_failed         = -18


# Carries out a fit to TRAPPIST-1 data from December 2019
# Optimizes, runs HMC.  No prior on longitude difference.
# Eccentricity vectors used without the 1/e prior for now.
# Student's-t distribution for likelihood.
# Uses full data, a larger timestep.

use_student = true

# Carry out a numerical check of the gradient once (NOTE: this is slow!)
test_gradient = false
if test_gradient
  include("../compute_grad_num.jl")
end
include("../CGS.jl")
include("log_students_prob.jl")
include("../extract_planet.jl")
include("../nlog_prior.jl")
include("../loglinspace.jl")
include("../regress.jl")
# Use NbodyGradient code:
include("../../src/NbodyGradient/src/ttv.jl")

using Main.CGS
#using PyPlot
using ForwardDiff
using DelimitedFiles
using Printf
using SpecialFunctions
using LinearAlgebra
using JLD2

# Uses e\cos{\omega} & e\sin{\omega} as parameters for Markov chain.

# Read in elements_grid from prior optimization that showed a jump
# in the minimum:

function run_hmc(fname::String,felements::String,foutput::String,nstep::Int64,epsilon0::Float64,nleap0::Int64,showplots::Bool)
# Number of bodies:
n = 8
nplanet = n-1

# Prior on masses:
sigm2 = 1e-8^2

# Number of grid points for each mass - this will give ~0.12-sigma spacing:
ngrid = 1
# Mass grid in standard deviations:
sig_grid = zeros(ngrid);
if ngrid > 1
  sig_grid .= linearspace(-3.0,3.0,ngrid)
end
mass_grid = copy(sig_grid)
# Use Simon's values as a guide; convert these to mass-ratio:
sgrimm = [0.082,0.073,0.022,0.027,0.026,0.027,0.034]*MEARTH/(0.09*MSUN)

t0 = 7257.93115525
h  = 0.06
tmax = 800.0

# Read in initial conditions from the first optimization:
#elements_minimum = readdlm("elements_dompi.txt",',')
#elements_minimum = readdlm("elements_noprior_mixture_rescale.txt",',')
#elements_minimum = readdlm("elements_noprior_students.txt",',')
elements_minimum = readdlm(felements,',')
elements_minimum = elements_minimum[1:n,:]
elements = copy(elements_minimum)
elements_trial = copy(elements)

# Make an array, tt,  to hold transit times:
# First, though, make sure it is large enough:
ntt = zeros(Int64,n);
for i=2:n
  ntt[i] = ceil(Int64,tmax/elements[i,2])+100
end
#println("ntt: ",ntt)
tt1 = zeros(n,maximum(ntt));
tt2 = zeros(n,maximum(ntt));
dtdq0 = zeros(n,maximum(ntt),7,n);
dtdelements = zeros(n,maximum(ntt),7,n);
# Save a counter for the actual number of transit times of each planet:
count1 = zeros(Int64,n);
# Set the "size" of the star (only checks that close to transit):
rstar = 1e11
# Call the ttv function:
@time dq=ttv_elements!(n,t0,h,tmax,elements,tt1,count1,0.0,0,0,rstar)
@time ttv_elements!(n,t0,h,tmax,elements,tt1,count1,dtdq0,rstar)

# Now call with half the timestep:
count2 = zeros(Int64,n)

# Set the initial outlier parameters: ndof =  number of degrees of freedom (\nu);
#    lnV1 = log ratio by which variance increases

#ndof = 3.27; lnV1 =  -0.36

#ndof = 3.36016234786184600
#lnV1 = -0.34445676197382336

# Convert to new parameterization:
#lgndof = log(ndof)
#V1exp2nuinv = exp(lnV1+0.5/ndof)

#lgndof = 0.962757625597931
#V1exp2nuinv = 0.6089636752275194

# Set to a (nearly) Chi-square distribution:
lgndof_prior = 12.0
V1exp2nuinv_prior = 1.0
lgndof = copy(lgndof_prior)
V1exp2nuinv = copy(V1exp2nuinv_prior)
weightprior = 1e16

# Now, try to improve the fit to the data:
# First, find the indices that match for each planet:
iplanet = zeros(Int64,0)
indx = zeros(Int64,0)
tobs_tot = zeros(Float64,0)
sobs_tot = zeros(Float64,0)
ntrans = 0
chi_ttv = 0.0
chi_ttv_error = 0.0
chi_ttv += ntrans*log_norm_students_t([lgndof,V1exp2nuinv])
data = readdlm(fname,',')
for ip=1:nplanet
  eobs,tobs,sobs,nobs = extract_planet(data,ip)
  ittv = 1
  for j=1:nobs
    ittv = 1; dt = Inf
    for jttv=1:count1[ip+1]
      if abs(tt1[ip+1,jttv]-tobs[j]) < dt
        ittv = jttv
        dt = abs(tt1[ip+1,jttv]-tobs[j])
      end
    end
    push!(iplanet,ip+1)
    push!(indx,ittv)
    push!(tobs_tot,tobs[j])
    push!(sobs_tot,sobs[j])
    chi_ttv += lnprob([tobs[j],tt1[ip+1,ittv],sobs[j],lgndof,V1exp2nuinv])
    ntrans +=1
  end
end
# Now add normalization:
println("Total number of transits: ",ntrans," Chi-square: ",chi_ttv," error: ",chi_ttv_error)

# Now, optimize the model:

# First, vary just the period & the t0 values for all planets:
param_vary = [1,2,3,4,5]
nvary = length(param_vary)
dx = [1e-10,1e-5,1e-5,1e-5,1e-5]
#nparam = length(param_vary) * (n-1) # Exclude the star
if use_student
  nparam = length(param_vary) * (n-1) + 2 # Exclude the star, include outlier parameters
else
  nparam = length(param_vary) * (n-1)  # Exclude the star, exclude outlier/mixture parameters
end

# Now, try to implement a  Levenburg-Marquardt algorithm:
dchi = Inf
# Tolerance with which to converge to:
tol  = 3e-5
iter = 0
itmax = 100
hessian = zeros(Float64,nparam,nparam)
#lambda = [1e3,316.,100.,31.6,10.0,3.16,1.0,3.16e-1,1e-1,3.16e-2,1e-2,3.16e-3,1e-3,3.16e-4,1e-4,3.16e-5,1e-5,3.16e-6,1e-6]
#lambda = [1e3,316.,100.,31.6,10.0,3.16,1.0,3.16e-1,1e-1,1e-2,3.16e-3,1e-3,3.16e-4,1e-4]
#lambda = [1e4,1e3,100.,10.0,1.0,0.1,.01,.001,1e-4,1e-5,1e-6]
#lambda = [1e3,1e2,10.,1.0,0.1,.01,.001,1e-4,1e-5,1e-6,1e-7,0.]
#nlambda = 12
lambda = [1e3,1e2,10.,1.0,0.1,.01,.001,1e-4,0.]
nlambda = 9
#lambda = [1.0]
#nlambda = 1
#lambda = 5e-4
# Gradient contains derivatives with respect to mass, period, t0, e*cos(omega), e*sin(omega) for each planet:
gradf = zeros(Float64,nparam)
# Numerical gradient check:
gradf_num = zeros(Float64,nparam)

# Save gradient to compare with prior step when new chi-square is rejected:
# Set up a grids for looping over masses of each planet (there are nplanet=n-1 of these):
chi_grid = zeros(Float64,nplanet,ngrid)
prior_grid = zeros(Float64,nplanet,ngrid)
elements_grid = zeros(Float64,nplanet,n,7,ngrid)
lgndof_grid = zeros(Float64,nplanet,ngrid)
V1exp2nuinv_grid = zeros(Float64,nplanet,ngrid)
tt_grid = zeros(nplanet,n,maximum(ntt),ngrid)
#jgrid=[11,10,9,8,7,6,5,4,3,2,1,12,13,14,15,16,17,18,19,20,21]
#jgrid=[16,15,14,13,12,11,10,9,8,7,6,5,4,3,2,1,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31]
jgrid=[1]
# Loop over the mass-ratios of each planet:
#for imass = 1:nplanet
# We're interested in the optimum for now, so just "iterate" over one planet:
for imass = 1:1
# Loop over the grid points for each planet:
for igrid in jgrid
  start = time()
  if igrid == jgrid[1]
    elements = copy(elements_minimum)
    # First, optimize fit, so no constraint upon mass:
    sigm2 = Inf
    if imass == 1
      compute_grad_num_q = true
    else
      compute_grad_num_q = false
    end
  elseif igrid == jgrid[1]+1
    elements = copy(elements_grid[imass,:,:,jgrid[1]])
    compute_grad_num_q = false
  end
  # Compute chi-square
  dq0 = ttv_elements!(n,t0,h,tmax,elements,tt2,count2,0.0,0,0,rstar)
  # Add in prior:
  prior_ttv = (elements[1+imass,1]-mass_grid[igrid])^2/sigm2
  lgndof_best = lgndof; V1exp2nuinv_best = V1exp2nuinv
  prior_ttv += (lgndof-lgndof_prior)^2*weightprior + (V1exp2nuinv - V1exp2nuinv_prior)^2*weightprior
  chi_ttv = prior_ttv
  for j=1:ntrans
    chi_ttv += lnprob([tobs_tot[j],tt2[iplanet[j],indx[j]],sobs_tot[j],lgndof_best,V1exp2nuinv_best])
  end
  dchi = Inf; iter = 0
  while abs(dchi) > tol && iter < itmax
    tstart = time()
    dtdelements = ttv_elements!(n,t0,h,tmax,elements,tt1,count1,dtdq0,rstar)
    # Compute the "chi-square", gradient & Hessian:
    fill!(hessian,0.0)
    fill!(gradf,0.0)
    i = 0; chi_ttv = 0.0
  # Iterate over planets & varying parameters:
    for ip=1:nplanet, ivary=1:nvary
      i +=1
      # If varying mass, this is at end of dtdq0 parameter list:
      if ivary == 1; iv = 7; else; iv= ivary-1; end
      # Gradient of "chi-square" with respect to each varying parameter:
      for it=1:ntrans
        # Set up parameter vector for computing likelihood of this data point:
        p = [tobs_tot[it],tt1[iplanet[it],indx[it]],sobs_tot[it],lgndof,V1exp2nuinv]
        if i == 1
          # Compute the chi-square (but only once):
          chi_ttv += lnprob(p)
        end
        # Compute gradient of -2*ln(p_i) with respect to these parameters:
        grad_tmp = grad_point(p)
        # Take derivative of -2*log likelihood with respect to model at this point, and then compute derivative of model with
        # respect to model parameters:
        gradf[i] += grad_tmp[2]*dtdelements[iplanet[it],indx[it],iv,ip+1]
        if use_student
          if i == 1
          # Derivative wrt ndof:
            gradf[nparam-1] += grad_tmp[4]
          # Derivative wrt V1exp2nuinv:
            gradf[nparam  ] += grad_tmp[5]
          end
        end
  #      gradf[i] -= dtdelements[iplanet[it],indx[it],iv,ip+1]*(tobs_tot[it]-tt1[iplanet[it],indx[it]])/sobs_tot[it]^2
      end
      # Now compute the Hessian:
      j = 0
      for jp=1:nplanet, jvary=1:nvary
        j +=1
        if jvary == 1; jv = 7; else; jv= jvary-1; end
        for it=1:ntrans
          # Set up parameter vector for computing likelihood of this data point:
          p = [tobs_tot[it],tt1[iplanet[it],indx[it]],sobs_tot[it],lgndof,V1exp2nuinv]
          # Compute hessian of -2ln(p_i) with respect to these parameters:
          hess_tmp = hess_point(p)
          hessian[i,j] += hess_tmp[2,2]*dtdelements[iplanet[it],indx[it],iv,ip+1]*dtdelements[iplanet[it],indx[it],jv,jp+1]
          if use_student
            if j == 1
              hessian[i,nparam-1] += hess_tmp[2,4]*dtdelements[iplanet[it],indx[it],iv,ip+1]
              hessian[i,nparam  ] += hess_tmp[2,5]*dtdelements[iplanet[it],indx[it],iv,ip+1]
            end
            if i == 1
              hessian[nparam-1,j] += hess_tmp[4,2]*dtdelements[iplanet[it],indx[it],jv,jp+1]
              hessian[nparam  ,j] += hess_tmp[5,2]*dtdelements[iplanet[it],indx[it],jv,jp+1]
            end
            if i == 1 && j == 1
              hessian[nparam-1,nparam-1] += hess_tmp[4,4]
              hessian[nparam-1,nparam  ] += hess_tmp[4,5]
              hessian[nparam  ,nparam-1] += hess_tmp[5,4]
              hessian[nparam  ,nparam  ] += hess_tmp[5,5]
            end
          end
  #        hessian[i,j] += dtdelements[iplanet[it],indx[it],iv,ip+1]*dtdelements[iplanet[it],indx[it],jv,jp+1]/sobs_tot[it]^2
        end
      end
    end
    # Add in grid constraint upon mass:
    chi_ttv += (elements[imass+1,1]-mass_grid[igrid])^2/sigm2
    chi_ttv += (lgndof-lgndof_prior)^2*weightprior + (V1exp2nuinv - V1exp2nuinv_prior)^2*weightprior
    # Add in normalization:
    chi_ttv += ntrans*log_norm_students_t([lgndof,V1exp2nuinv])
    # Add in prior gradient:
    gradf[(imass-1)*5+1] += 2.0*(elements[imass+1,1]-mass_grid[igrid])/sigm2
    # Add in gradient with respect to (lgndof,V1exp2nuinv):
    gradf[nparam-1:nparam] .+= ntrans*grad_norm([lgndof,V1exp2nuinv])
    gradf[nparam-1] += 2*(lgndof-lgndof_prior)*weightprior
    gradf[nparam] += 2*(V1exp2nuinv - V1exp2nuinv_prior)*weightprior
    # Compare with numerical gradient:
    if compute_grad_num_q && iter == 0 && test_gradient == true
       gradf_num = convert(Array{Float64,1},
        compute_grad_num(n,h,t0,tmax,elements,tt2,count2,rstar,tobs_tot,sobs_tot,lgndof,V1exp2nuinv,indx,nparam,nplanet,nvary))
      println("grad: ",gradf," gradf_num: ",gradf_num," diff: ",gradf-gradf_num)
    end
    # Add in Hessian with respect to prior:
    hessian[(imass-1)*5+1,(imass-1)*5+1] += 2.0/sigm2
    # Add in hessian of normalization:
    hessian[nparam-1:nparam,nparam-1:nparam] .+= ntrans*hess_norm([lgndof,V1exp2nuinv])
    hessian[nparam-1,nparam-1] += 2*weightprior
    hessian[nparam,nparam] += 2*weightprior
    # Now, carry out Newton step (can no longer use Levenberg-Marquardt!):
    chi_best = chi_ttv; elements_best = copy(elements); lgndof_best = copy(lgndof); V1exp2nuinv_best=copy(V1exp2nuinv);  prior_best = copy(prior_ttv)
    elements_trial = copy(elements); lgndof_trial = copy(lgndof); V1exp2nuinv_trial = copy(V1exp2nuinv); prior_trial = copy(prior_ttv)
    for ilambda = 1:nlambda
    # Update the parameters:
      mat = copy(hessian)
      for j=1:nparam
        mat[j,j] += lambda[ilambda]*hessian[j,j]
      end
      # Newton step:
      deltax = -\(mat,gradf)
      # deltax may take some parameters out of bounds - in which case
      # we'll shrink it until inbounds:
      out_of_bounds = true; factor = 1.0;  ecc2_trial = zeros(n-1)
      elements_trial = copy(elements); lgndof_trial = copy(lgndof); V1exp2nuinv_trial = copy(V1exp2nuinv)
      while out_of_bounds && factor > 1e-3
        elements_trial = copy(elements); lgndof_trial = copy(lgndof); V1exp2nuinv_trial = copy(V1exp2nuinv)
      # Loop over planets & varied parameters:
        i = 0
        for ip=2:n
          for ivary=1:nvary
            i += 1
            elements_trial[ip,param_vary[ivary]] += factor*deltax[i]
          end
          ecc2_trial[ip-1] = elements_trial[ip,4]^2+elements_trial[ip,5]^2
        end
        if use_student
          lgndof_trial+= factor*deltax[nparam-1]
          V1exp2nuinv_trial += factor*deltax[nparam]
        end
        if maximum(ecc2_trial) < 0.2 && log(abs(V1exp2nuinv_trial)) > -10.0 && exp(lgndof_trial) > 0.5 && exp(lgndof_trial) < 1e8 && V1exp2nuinv_trial > 0.0
          out_of_bounds = false
        else
          # Since we're out of bounds, shrink deltax:
          factor *= 0.9
        end
      end
    # Compute chi-square
      if maximum(ecc2_trial) < 0.2 && log(abs(V1exp2nuinv_trial)) > -10.0 && exp(lgndof_trial) > 1.0 && exp(lgndof_trial) < 1e8 && V1exp2nuinv_trial > 0.0
        chi_trial = Inf
        while chi_trial > chi_ttv && factor > 1e-2
          # Now, do a line search:
          dq0 = ttv_elements!(n,t0,h,tmax,elements_trial,tt2,count2,0.0,0,0,rstar)
          chi_trial = 0.0
          for j=1:ntrans
            chi_trial += lnprob([tobs_tot[j],tt2[iplanet[j],indx[j]],sobs_tot[j],lgndof_trial,V1exp2nuinv_trial])
          end
          # Add in normalization:
          chi_trial += ntrans*log_norm_students_t([lgndof_trial,V1exp2nuinv_trial])
          # Add in prior:
          prior_trial = (elements_trial[imass+1,1]-mass_grid[igrid])^2/sigm2
          prior_trial += (lgndof_trial-lgndof_prior)^2*weightprior + (V1exp2nuinv_trial - V1exp2nuinv_prior)^2*weightprior
          chi_trial += prior_trial
          if chi_trial > chi_ttv
            factor *= 0.5
            elements_trial = copy(elements); lgndof_trial = copy(lgndof); V1exp2nuinv_trial = copy(V1exp2nuinv)
      # Loop over planets & varied parameters:
            i = 0
            for ip=2:n, ivary=1:nvary
              i += 1
              elements_trial[ip,param_vary[ivary]] += factor*deltax[i]
            end
            if use_student
              lgndof_trial += factor*deltax[nparam-1]
              V1exp2nuinv_trial += factor*deltax[nparam]
            end
          end
        end
      else
        chi_trial = Inf
      end
      if chi_trial < chi_ttv && chi_trial < chi_best
        println("New chi2: ",chi_trial," ",chi_trial," prior_trial: ",prior_trial,
           " target_mass: ",mass_grid[imass]," elapsed time: ",(time()-start)/60," min")
      # A better chi-square is found - update elements, transit times & chi-square:
        dchi = chi_ttv - chi_trial
        elements_best .= elements_trial
        if use_student
          lgndof_best = copy(lgndof_trial)
          V1exp2nuinv_best = copy(V1exp2nuinv_trial)
        end
        tt1 .= tt2
        chi_best = copy(chi_trial)
        prior_best = copy(prior_trial)
      else
        println("Chi-square worse. ",chi_trial," best: ",chi_best," chi_ttv: ",chi_ttv)
      end
    end
    # Don't include prior:
    if chi_best < chi_ttv
      telapse = time()-tstart
      println("New chi2: ",chi_best," prior: ",prior_best," target m: ",mass_grid[igrid]," elapsed time: ",telapse/60," min")
      chi_grid[imass,igrid] = chi_best - prior_best
      if use_student
        lgndof_grid[imass,igrid] = lgndof_best
        V1exp2nuinv_grid[imass,igrid] = V1exp2nuinv_best
      end
      elements_grid[imass,:,:,igrid] .= elements_best
      tt_grid[imass,:,:,igrid] .= tt1
      prior_grid[imass,igrid] = prior_best
      chi_ttv = copy(chi_best); prior_ttv=copy(prior_best); elements .= elements_best
      if use_student
        lgndof = copy(lgndof_best); V1exp2nuinv = copy(V1exp2nuinv_best)
      end
      iter += 1
    else
#      println("Can't improve chi-square. elapsed time: ",toq()/60.," min")
      # Save the initial parameters:
      chi_grid[imass,igrid] = chi_best - prior_best
      if use_student
        lgndof_grid[imass,igrid] = lgndof_best
        V1exp2nuinv_grid[imass,igrid] = V1exp2nuinv_best
      end
      elements_grid[imass,:,:,igrid] .= elements_best
      tt_grid[imass,:,:,igrid] .= tt1
      prior_grid[imass,igrid] = prior_best
      chi_ttv = copy(chi_best); prior_ttv=copy(prior_best); elements .= elements_best
      if use_student
        lgndof = copy(lgndof_best); V1exp2nuinv = copy(V1exp2nuinv_best)
      end
      iter = itmax
    end
  end
  println("Planet ",imass," Finished grid point: ",igrid," ",elements," ",lgndof," ",V1exp2nuinv,
          " ",chi_ttv-prior_ttv," ",prior_ttv," Elapsed time: ",time()-start)
  # For first grid point, we optimized chi-square.  Now, we want to add grid constraints.
  if igrid == jgrid[1]
    # Now, create mass grid, spanning +-3-sigma of Simon Grimm's uncertainty:
    mass_grid = sig_grid.*sgrimm[imass]
    # Offset grid about optimum:
    mass_grid .+= elements_grid[imass,imass+1,1,jgrid[1]]
    # Now, cause grid to be reinstated:
    sigm2 = 1e-8^2
    itmax = 10
  end
end
end

cov_save = inv(0.5*hessian)

hess_save = 0.5*hessian

@save "T1_grimm_students.jld2" elements hess_save cov_save

planet = ["b","c","d","e","f","g","h"]
sigm = zeros(nplanet)
mass = zeros(nplanet)
ecc = zeros(nplanet)
pomega = zeros(nplanet)
for i=1:nplanet
  mass[i] = elements[i+1,1]*.09*MSUN/MEARTH
  sigm[i] = sqrt(abs(cov_save[5*(i-1)+1,5*(i-1)+1]))*.09*MSUN/MEARTH
  ecc[i] = sqrt(elements[i+1,4]^2+elements[i+1,5]^2)
  pomega[i] = atan(elements[i+1,5],elements[i+1,4])*180/pi
  println(planet[i]," ",@sprintf("%6.4f",mass[i]),"+-",@sprintf("%6.4f",sigm[i])," ",
   @sprintf("%6.4f",elements[i+1,4]),"+-",@sprintf("%6.4f",sqrt(abs(cov_save[5*(i-1)+4,5*(i-1)+4]))),
   " ",@sprintf("%6.4f",elements[i+1,5]),"+-",@sprintf("%6.4f",sqrt(abs(cov_save[5*(i-1)+5,5*(i-1)+5]))),
   " ",@sprintf("%6.4f",ecc[i])," ",@sprintf("%7.2f",pomega[i]))
end

# Print deviations of each transit:
chi_planet = zeros(Float64,n-1)
nobs_planet = zeros(Int64,n-1)
dev = zeros(Float64,ntrans)
elements = elements_grid[1,:,:,jgrid[1]]
#        dq0 = ttv_elements!(n,t0,h,tmax,elements,tt2,count2,0.0,0,0,rstar)
dtdelements = ttv_elements!(n,t0,h,tmax,elements,tt1,count1,dtdq0,rstar)
for j=1:ntrans
  dev[j] = (tobs_tot[j]-tt1[iplanet[j],indx[j]])/sobs_tot[j]
  chi_planet[iplanet[j]-1] += dev[j]^2
  nobs_planet[iplanet[j]-1] += 1
  if abs(dev[j]) > 3
    println(j," ",iplanet[j]," ",indx[j]," ",tt1[iplanet[j],indx[j]]," ",tobs_tot[j]," ",sobs_tot[j]," ",dev[j]," ",dev[j]^2)
  end
end

hinv = 0.5*(inv(hess_save)+inv(hess_save)')
cholhinv = cholesky(hinv)

x = zeros(Float64,nparam)
# Insert elements into parameter vector:
for i=2:n, j=1:5
  x[(i-2)*5+j]= elements[i,j]
end
# Select values of mixture model at optimum parameters:
lgndof = lgndof_grid[1,jgrid[1]]
V1exp2nuinv = V1exp2nuinv_grid[1,jgrid[1]]
# Add in mixture model parameters to parameter vector:
if use_student
  # Convert to log(ndof) and V1*exp(1/(2ndof)):
  x[nparam-1] = lgndof
  x[nparam  ] = V1exp2nuinv
end

# Compute the negative log likelihood:
function f0(x)
# computes negative log likelihood:
for i=2:n, j=1:5
  elements_trial[i,j] = x[(i-2)*5+j]
end
dq0 = ttv_elements!(n,t0,h,tmax,elements_trial,tt2,count2,0.0,0,0,rstar)
# Ignoring priors:
chi_trial = 0.0
lgndof = x[nparam-1]; V1exp2nuinv = x[nparam]
for j=1:ntrans
  chi_trial += lnprob([tobs_tot[j],tt2[iplanet[j],indx[j]],sobs_tot[j],lgndof,V1exp2nuinv])
end
# Add in normalization:
chi_trial += ntrans*log_norm_students_t([lgndof,V1exp2nuinv])
# Convert from -2log(likelihood) to -log(likelihood)
chi_trial *= 0.5
# Add in prior:
chi_trial += nlog_prior(x)
return chi_trial
end

# Check that this routine works:
chi_opt = chi_grid[1,jgrid[1]]
println("f0(x): ",f0(x)," 0.5 chi_opt: ",0.5*chi_opt)

x_opt = x
writedlm("optimum_values.txt",x_opt,chi_opt)

#read(stdin,Char)

nrand = 10
dxrand = zeros(nparam,nrand)
nll_exact = zeros(nrand)
nll_approx = zeros(nrand)
for i=1:nrand
#for i=1:nparam
# Draw random parameters:
  dev = randn(nparam)/sqrt(nparam)
#  dev = zeros(nparam); dev[i]=1.0
  dx = cholhinv.L*dev
  x_rand = x + dx
  dxrand[:,i] = dx
  nll_exact[i]  = f0(x_rand);
  nll_approx[i] =  0.5*chi_opt + 0.5*dx'*hess_save*dx; #  Figure out the correct quadratic approximation
  # Compare:
  println(i," nll_exact: ",nll_exact[i]-0.5*chi_opt," nll_approx: ",nll_approx[i]-0.5*chi_opt," Difference: ",nll_approx[i]-nll_exact[i])
#  read(stdin,Char)
end

#--------------------------------------------------------------
# Now, run a Hamiltonian Markov chain:

# Create a parameter vector to hold orbital elements & mass ratios of planets:
npar = (n-1)*5
if use_student
  # Add two variables for Student's t model:
  npar += 2
end

# Recompute chi-square:
chi_ttv = 0.0
dq0 = ttv_elements!(n,t0,h,tmax,elements,tt2,count2,0.0,0,0,rstar)
chi_ttv = 0.0
for j=1:ntrans
  chi_ttv += lnprob([tobs_tot[j],tt2[iplanet[j],indx[j]],sobs_tot[j],lgndof,V1exp2nuinv])
end
# Add in normalization:
chi_ttv += ntrans*log_norm_students_t([lgndof,V1exp2nuinv])

# Check that prior is zero:
println("-Log prior: ",nlog_prior(x))

# The optimum chi-square:
chi_opt = chi_grid[1,jgrid[1]]

# Now invert to compute momentum "inverse mass" matrix:
#invmass = inv(hess_save)
invmass = 0.5*(inv(hess_save)+inv(hess_save)')

# "Time" step is 0.05:
#epsilon0 = 0.20

# Start out with a modest number of steps:
#nstep = 1
#nstep = 500
state = zeros(2*npar+2,nstep)

# Cholesky decompose hessian for computation of momentum:
#hessian = .5*(hessian + hessian')
hessian = .5*(hess_save + hess_save')
#cholh = transpose(chol(hessian))
#cholh = permutedims(chol(hessian))
cholh = permutedims(cholesky(hessian).U)

# Leapfrog (Neal 2012):
#nleap0 = 1000
#nleap0 = 10
if showplots
  substep = zeros(npar,nleap0)
end
nacc = 0
grad = zeros(Float64,npar)
hamnew = 0.0
#tic(); elapsed = 0.0
tstart = time(); elapsed = 0.0

function compute_gradient!(grad,x)
# Check that variables are defined:
#println("n: ",n," ntrans: ",ntrans," t0: ",t0," h: ",h," tmax: ",tmax)
# Begin gradient computation.
# Put the x parameters into the elements matrix:
for i=2:n
  for j=1:5
    elements_trial[i,j] = x[(i-2)*5+j]
  end
end
dtdelements = ttv_elements!(n,t0,h,tmax,elements_trial,tt1,count1,dtdq0,rstar)
# Compute the gradient & Hessian:
i = 0; fill!(grad,0.0)
lgndof = x[nparam-1]; V1exp2nuinv = x[nparam]
# Iterate over planets & varying parameters:
for ip=1:n-1, ivary=1:nvary
  i +=1
  # If varying mass, this is at end of dtdq0 parameter list:
  if ivary == 1; iv = 7; else; iv= ivary-1; end
  # Gradient of chi-square with respect to each varying parameter:
  for it=1:ntrans
    p = [tobs_tot[it],tt1[iplanet[it],indx[it]],sobs_tot[it],lgndof,V1exp2nuinv]
    # Take derivative of -2*log likelihood with respect to model at this point, and then compute derivative of model with
    # respect to model parameters:
    grad_tmp = grad_point(p)
    dtdxi = dtdelements[iplanet[it],indx[it],iv,ip+1]
    grad[i] += 0.5*grad_tmp[2]*dtdxi
    if use_student
      if i == 1
        # Derivative wrt lgndof:
        grad[nparam-1] += 0.5*grad_tmp[4]
        # Derivative wrt V1exp2nuinv:
        grad[nparam  ] += 0.5*grad_tmp[5]
      end
    end
  end
end
# Add in normalization:
grad[nparam-1:nparam] .+= 0.5*ntrans*grad_norm([lgndof,V1exp2nuinv])
# Add in prior:
grad .+= grad_nlog_prior(x)
# End gradient computation
return
end

function check_bounds(x)
# Check that we're within bounds:
check = true
for i=2:n
  # Check if eccentricity >= 1.0
  if x[(i-2)*5+4]^2+x[(i-1)*5]^2 >= 1.0
    check = false
  end
  # Check if masses are non-positive:
  if x[(i-2)*5+1] <= 0.0
    check = false
  end
end
# Check if ndof is in bounds:
if exp(x[nparam-1]) < 0.5 || exp(x[nparam-1]) > 100.0
  check = false
end
return check
end

# Create array to save data on Markov chain:
stats = zeros(Float64,5,nstep)  # save epsilon, nleap, alpha, accept/reject
fofx = f0(x)
fofxnew = 0.0

# Loop over steps in Hamiltonian Markov Chain:
for istep = 1:nstep
# Draw momentum from the kinetic energy distribution:
  p = cholh*randn(npar)
  ham = 0.5*p'*invmass*p + fofx
  pnew = p
  xnew = x
  nleap = round(Int64,nleap0*(0.8+0.2*rand()))
  if showplots
    # Save substeps to make plots of progress:
    substep = zeros(npar,nleap+1)
    substep[:,1]=xnew
  end
  # Compute gradient with respect to parameters:
  compute_gradient!(grad,xnew)
  # Choose a random value of epsilon:
#  epsilon = abs(randn())*epsilon0
# Choose between 0.03-0.04:
  epsilon = (0.75+0.25*rand())*epsilon0
  # Initial step of leapfrog integration:
  pnew -= 0.5*epsilon*grad
  # Loop over leapfrog steps:
  ileap = 1
  while ileap <= nleap
    # parameter update:
    xnew += epsilon*invmass*pnew
#    println("xnew: ",xnew)
    # check that we're within bounds for eccentricity:
    if check_bounds(xnew)
      if showplots
        substep[:,ileap+1]=xnew
      end
    # momentum update:
      compute_gradient!(grad,xnew)
      pnew -= epsilon*grad
      ileap += 1
    else
      println("Parameters out of bounds: ",xnew)
      ileap = nleap+1
    end
  end
  # Final momentum update:
  pnew += 0.5*epsilon*grad
  # Now, compute hamiltonian again:
  if check_bounds(xnew)
    fofxnew = f0(xnew)
    hamnew = 0.5*pnew'*invmass*pnew+ fofxnew
  # Compute whether to accept or reject:
    alpha = exp(ham-hamnew)
  else
    alpha = 0.0
    hamnew = 0.0
  end
  uni = rand()
  if uni < alpha
    x .= xnew
    p .= pnew
    fofx = copy(fofxnew)
    println(istep," accepted, 1: ",ham," 2: ",hamnew," eps: ",epsilon," L: ",nleap," ",fofx)
    ham = copy(hamnew)
    nacc += 1
    # plot masses of inner two planets:
    if showplots
      plot(substep[1,:],substep[6,:])
    end
    stats[:,istep] = [epsilon;float(nleap);alpha;uni;1.0]
  else
    println(istep," rejected, 1: ",ham," 2: ",hamnew," eps: ",epsilon," L: ",nleap," ",fofx)
    stats[:,istep] = [epsilon;float(nleap);alpha;uni;0.0]
  end
  state[:,istep] = [x;-p;ham;fofx]
  if mod(istep,10) == 0
    elapsed += time()-tstart
    println(istep," elapsed: ",elapsed/3600.," facc: ",nacc/istep," ",(nstep-istep)/istep*elapsed/3600.," hr remaining")
    tstart = time()
    # That's it!  Now save every variable to an HDF5 JLD2 file (checkpoint):
    @save foutput n cov_save hess_save fname felements nparam t0 h tmax ntrans iplanet indx tobs_tot sobs_tot data count1 state hessian cholh nleap0 nacc nstep stats elements chi_opt x_opt
  end
end
println("Fraction accepted: ",nacc/nstep)
#--------------------------------------------------------------

# End of run_hmc
return
end

# Example of running this routine:

# File containing the transit times:
#fname = "T1_timings_20191203.txt"

#felements = "elements_noprior_students.txt"

# Name an output file:
#foutput = "T1_run_hmc_student_ecc_lgndof_V1exp2nuinv_001.jld2"

# Call the likelihood profile function:
#run_hmc(fname,felements,foutput,nstep,epsilon0,nleap0,showplots)
#run_hmc(fname,felements,foutput,100,0.2,10,true)


using Cxx
using CxxStd: WrappedCppPrimArray

cxx"""
#include <tf/transform_datatypes.h>
"""

const tfVector3    = Union{cxxt"tf::Vector3",    cxxt"tf::Vector3&"}
const tfQuaternion = Union{cxxt"tf::Quaternion", cxxt"tf::Quaternion&"}
const tfMatrix3x3  = Union{cxxt"tf::Matrix3x3",  cxxt"tf::Matrix3x3&"}

function Base.unsafe_wrap(::Type{DenseArray}, v::T) where {T<:tfVector3}
    WrappedCppPrimArray(icxx"&$v[0];",3)
end

function Base.unsafe_wrap(::Type{DenseArray}, q::T) where {T<:tfQuaternion}
    WrappedCppPrimArray(icxx"&$q[0];",4)
end

function Base.convert(::Type{CT}, v::T) where {CT<:AbstractArray,T<:Union{tfVector3, tfQuaternion}}
    convert(CT, unsafe_wrap(DenseArray, v))
end

function Base.convert(::Type{CT}, m::T) where {CT<:AbstractArray, T<:tfMatrix3x3}
    [DenseArray(icxx"$m[0];"), DenseArray(icxx"$m[1];"), DenseArray(icxx"$m[2];")]
end


function dicttest()
  com = LanguageDynamics.ZipfTravellerPopulation{LanguageDynamics.MomentumSelector}(7, 2.5)

  for i = 1:200
    LanguageDynamics.inject!(LanguageDynamics.MomentumSelector(0.01, 0.02, 1.0, 4, 0.1), com)
  end

  #for i = 1:100
  #  eject!(com)
  #end

  while com.nactions < 0.5*200*100#_000
    LanguageDynamics.travel!(com)
    LanguageDynamics.rendezvous!(com)
  end
end

function arrtest()
  com = PoolPopulation{EmptyInteractor}()

  for i = 1:500
    inject!(EmptyInteractor(), com)
  end

  #for i = 1:100
  #  eject!(com)
  #end

  for i = 1:1000
    rendezvous!(com)
  end
end