"""
```
shock_decompositions(m, system, histshocks)

shock_decompositions(system, forecast_horizons, histshocks, start_index,
    end_index)

shock_decompositions(m, system, histshocks, start_date, end_date)

shock_decompositions(m, system, forecast_horizons, histshocks, start_index,
    end_index, regime_inds, cond_type)

shock_decompositions(m, system, old_system,
                              histshocks, old_histshocks, start_date, end_date,
                              cond_type; full_shock_decomp)

shock_decompositions(m, system, old_system,
                              forecast_horizons, histshocks, old_histshocks,
                              start_index, end_index,
                              regime_inds, cond_type; full_shock_decomp)
```

### Inputs

- `system::System{S}` or `RegimeSwitchingSystem{S}`: state-space system matrices
- `forecast_horizons::Int`: number of periods ahead to forecast
- `histshocks::Matrix{S}`: matrix of size `nshocks` x `hist_periods` of
  historical smoothed shocks
- `start_index::Int`: first index from which to return computed shock
  decompositions
- `end_index::Int`: last index for which to return computed shock decompositions
- `start_date::Date`: initial date for deterministic trends
- `end_date::Date`: final date for deterministic trends
- `regime_inds::Vector{UnitRange}`: indices of the data corresponding to each regime.
- `old_system::RegimeSwitchingSystem{S}`: state-space system matrices for the old model
- `old_histshocks::Matrix{S}`: matrix of size `nshocks` x `hist_periods` of
  historical smoothed shocks using the old system
- `full_shock_decomp::Bool`: If true for old_system case, return difference in shockdecs
    between all shocks for the new and old systems. Else, return old system's shocks
    with 1 cumulative AIT shock added on.

where `S<:AbstractFloat`.

### Outputs

- `states::Array{S, 3}`: matrix of size `nstates` x `nperiods` x `nshocks` of
  state shock decompositions
- `obs::Array{S, 3}`: matrix of size `nobs` x `nperiods` x `nshocks` of
  observable shock decompositions
- `pseudo::Array{S, 3}`: matrix of size `npseudo` x `nperiods` x `nshocks` of
  pseudo-observable shock decompositions

where `nperiods = `end_index - start_index + 1`.
"""
function shock_decompositions(m::AbstractDSGEModel,
    system::System{S}, histshocks::Matrix{S}) where {S<:AbstractFloat}

    horizon   = forecast_horizons(m)
    start_ind = index_shockdec_start(m)
    end_ind   = index_shockdec_end(m)

    shock_decompositions(system, horizon, histshocks, start_ind, end_ind)
end

function shock_decompositions(system::System{S},
    forecast_horizons::Int, histshocks::Matrix{S},
    start_index::Int, end_index::Int) where {S<:AbstractFloat}

    # Setup
    nshocks      = size(system[:RRR], 2)
    nstates      = size(system[:TTT], 2)
    nobs         = size(system[:ZZ], 1)
    npseudo      = size(system[:ZZ_pseudo], 1)
    histperiods  = size(histshocks, 2)
    allperiods   = histperiods + forecast_horizons

    states = zeros(S, nstates, allperiods, nshocks)
    obs    = zeros(S, nobs, allperiods, nshocks)
    pseudo = zeros(S, npseudo, allperiods, nshocks)

    # Check dates
    if forecast_horizons <= 0 || start_index < 1 || end_index > allperiods
        throw(DomainError("At least one of forecast_horizons, the index for the start date of the shock decomposition, and the index for the end date of the shock decomposition is outside the permissible range."))
    end

    # Set constant system matrices to 0
    system = zero_system_constants(system)

    z0 = zeros(S, nstates)
    for i = 1:nshocks
        # Isolate single shock
        shocks = zeros(S, nshocks, allperiods)
        shocks[i, 1:histperiods] = histshocks[i, :]

        # Iterate state space forward
        states[:, :, i], obs[:, :, i], pseudo[:, :, i], _ = forecast(system, z0, shocks)
    end

    # Return shock decompositions in appropriate range
    if start_index == 1 && end_index == allperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range, :], obs[:, range, :], pseudo[:, range, :]
    end
end

function shock_decompositions(m::AbstractDSGEModel{S},
                              system::RegimeSwitchingSystem{S}, histshocks::Matrix{S},
                              start_date::Dates.Date = date_presample_start(m),
                              end_date::Dates.Date = prev_quarter(date_forecast_start(m)),
                              cond_type::Symbol = :none) where {S<:AbstractFloat}
    # Set up dates
    horizon     = forecast_horizons(m; cond_type = cond_type)
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Set up regime indices. Note that we do not need to account
    # for ZLB split b/c `histshocks` should have zeros for anticipated shocks
    # in the pre-ZLB periods.
    regime_inds = regime_indices(m, start_date, end_date)
    if regime_inds[1][1] < 1 # remove periods occuring before desired start date
        regime_inds[1] = 1:regime_inds[1][end]
    end

    shock_decompositions(m, system, horizon, histshocks, start_index, end_index, regime_inds, cond_type)
end

function shock_decompositions(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S},
                              forecast_horizons::Int, histshocks::Matrix{S},
                              start_index::Int, end_index::Int,
                              regime_inds::Vector{UnitRange{Int}},
                              cond_type::Symbol) where {S<:AbstractFloat}

    # Setup
    nshocks     = size(system[1, :RRR], 2)
    nstates     = size(system[1, :TTT], 2)
    nobs        = size(system[1, :ZZ], 1)
    npseudo     = size(system[1, :ZZ_pseudo], 1)
    histperiods = size(histshocks, 2)
    allperiods  = histperiods + forecast_horizons

    states = zeros(S, nstates, allperiods, nshocks)
    obs    = zeros(S, nobs,    allperiods, nshocks)
    pseudo = zeros(S, npseudo, allperiods, nshocks)
    shocks = zeros(S, nshocks, histperiods)

    # Check dates
    if forecast_horizons <= 0 || start_index < 1 || end_index > allperiods
        throw(DomainError("At least one of forecast_horizons, the index for the start date of the shock decomposition, and the index for the end date of the shock decomposition is outside the permissible range."))
    end

    # Set constant system matrices to 0 and start at stationary steady state
    system = zero_system_constants(system)
    z0     = zeros(S, nstates)

    fcast_shocks = zeros(S, nshocks, forecast_horizons) # these are always zero
    for i = 1:nshocks
        # Isolate single shock
        shocks[i, :] = histshocks[i, :]

        # Looping through each historical regime
        for (reg_num, reg_ind) in enumerate(regime_inds)
            if reg_ind[end] > histperiods
                break
            end
            init_state = (reg_num == 1) ? zeros(S, nstates) : states[:, reg_ind[1] - 1, i] # update initial state

            states[:, reg_ind, i], obs[:, reg_ind, i], pseudo[:, reg_ind, i], _ =
                forecast(system[reg_num], init_state, shocks[:, reg_ind])
        end

        # Run forecast using regime-switching forecast
        if allperiods > histperiods
            fcast_inds = (histperiods + 1):allperiods
            states[:, fcast_inds, i], obs[:, fcast_inds, i], pseudo[:, fcast_inds, i], _ =
                forecast(m, system, states[:, histperiods, i], fcast_shocks, cond_type = cond_type)
        end

        # zero out shocks b/c want effects of single shocks
        shocks[i, :] .= 0.
    end

    # Return shock decompositions in appropriate range
    if start_index == 1 && end_index == allperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range, :], obs[:, range, :], pseudo[:, range, :]
    end
end

# Get Shock Decompositions with an old system and a bar that gives
## effect of using new system compared to old.
function shock_decompositions(m::AbstractDSGEModel{S},
                              system::RegimeSwitchingSystem{S}, old_system::RegimeSwitchingSystem{S},
                              histshocks::Matrix{S}, old_histshocks::Matrix{S},
                              start_date::Dates.Date = date_presample_start(m),
                              end_date::Dates.Date = prev_quarter(date_forecast_start(m)),
                              cond_type::Symbol = :none; full_shock_decomp = true) where {S<:AbstractFloat}
    # Set up dates
    horizon     = forecast_horizons(m; cond_type = cond_type)
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Set up regime indices. Note that we do not need to account
    # for ZLB split b/c `histshocks` should have zeros for anticipated shocks
    # in the pre-ZLB periods.
    regime_inds = regime_indices(m, start_date, end_date)
    if regime_inds[1][1] < 1 # remove periods occuring before desired start date
        regime_inds[1] = 1:regime_inds[1][end]
    end

    shock_decompositions(m, system, old_system, horizon, histshocks, old_histshocks, start_index, end_index,
                         regime_inds, cond_type, full_shock_decomp = full_shock_decomp)
end

function shock_decompositions(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S},
                              old_system::RegimeSwitchingSystem{S},
                              forecast_horizons::Int,
                              histshocks::Matrix{S}, old_histshocks::Matrix{S},
                              start_index::Int, end_index::Int,
                              regime_inds::Vector{UnitRange{Int}},
                              cond_type::Symbol; full_shock_decomp = true) where {S<:AbstractFloat}

    # Setup
    nshocks     = size(system[1, :RRR], 2)
    nstates     = size(system[1, :TTT], 2)
    nobs        = size(system[1, :ZZ], 1)
    npseudo     = size(system[1, :ZZ_pseudo], 1)
    histperiods = size(histshocks, 2)
    allperiods  = histperiods + forecast_horizons

    states = zeros(S, nstates, allperiods, nshocks)
    obs    = zeros(S, nobs,    allperiods, nshocks)
    pseudo = zeros(S, npseudo, allperiods, nshocks)
    shocks = zeros(S, nshocks, histperiods)
    old_shocks = zeros(S, nshocks, histperiods)

    # Old System Values
    old_states = zeros(S, nstates, allperiods, nshocks)
    old_obs    = zeros(S, nobs,    allperiods, nshocks)
    old_pseudo = zeros(S, npseudo, allperiods, nshocks)

    # Check dates
    if forecast_horizons <= 0 || start_index < 1 || end_index > allperiods
        throw(DomainError("At least one of forecast_horizons, the index for the start date of the shock decomposition, and the index for the end date of the shock decomposition is outside the permissible range."))
    end

    # Set constant system matrices to 0 and start at stationary steady state
    system = zero_system_constants(system)
    old_system = zero_system_constants(old_system)
    z0     = zeros(S, nstates)

    fcast_shocks = zeros(S, nshocks, forecast_horizons) # these are always zero
    for i = 1:nshocks
        # Isolate single shock
        shocks[i, :] = histshocks[i, :]
        old_shocks[i, :] = old_histshocks[i, :]

        # Old System Values
        # init_state_old = zeros(S, nstates)
        # old_states[:,1:histperiods,i], old_obs[:,1:histperiods,i], old_pseudo[:,1:histperiods,i], _ =
        #     forecast(old_system, init_state_old, shocks[:, 1:histperiods])

        counts = 1
        # Looping through each historical regime
        for (reg_num, reg_ind) in enumerate(regime_inds)
            if reg_ind[end] > histperiods
                break
            end
            init_state = (reg_num == 1) ? zeros(S, nstates) : states[:, reg_ind[1] - 1, i] # update initial state

            states[:, reg_ind, i], obs[:, reg_ind, i], pseudo[:, reg_ind, i], _ =
                forecast(system[reg_num], init_state, shocks[:, reg_ind])

            # Old System Values
            #=if reg_ind[end] < old_system
                old_system2 = system[reg_num]
                counts = reg_num
            else
                old_system2 = system[1]#system[counts]
            end=#
            old_system2 = old_system
            init_state_old = (reg_num == 1) ? zeros(S, nstates) : old_states[:, reg_ind[1] - 1, i]
            old_states[:, reg_ind, i], old_obs[:, reg_ind, i], old_pseudo[:, reg_ind, i], _ =
                forecast(old_system2[reg_num], init_state_old, old_shocks[:,reg_ind])
        end

        # Run forecast using regime-switching forecast
        if allperiods > histperiods
            fcast_inds = (histperiods + 1):allperiods
            states[:, fcast_inds, i], obs[:, fcast_inds, i], pseudo[:, fcast_inds, i], _ =
                forecast(m, system, states[:, histperiods, i], fcast_shocks, cond_type = cond_type)

            # Old System Values
            old_system2 = old_system#system[1]#system[counts]
            old_states[:, fcast_inds, i], old_obs[:, fcast_inds, i], old_pseudo[:, fcast_inds, i], _ =
                forecast(m, old_system2, old_states[:, histperiods, i], fcast_shocks, cond_type = cond_type)
                #forecast(old_system2, old_states[:, histperiods, i], fcast_shocks)
        end

        # zero out shocks b/c want effects of single shocks
        shocks[i, :] .= 0.
        old_shocks[i, :] .= 0.
    end

    # Add new system as shock to old system
    #=old_states = cat(old_states, zeros(size(old_states[:,:,1])), dims = 3)
    old_obs = cat(old_obs, zeros(size(old_obs[:,:,1])), dims = 3)
    old_pseudo = cat(old_pseudo, zeros(size(old_pseudo[:,:,1])), dims = 3)=#

    for i in 1:size(states,1)
        for j in 1:size(states,2)
            old_states[i,j,end] = sum([states[i,j,k] for k in 1:(size(states,3)-1)]) - sum([old_states[i,j,k] for k in 1:(size(states,3)-1)])
            full_shock_decomp && (old_states[i,j,1:end-1] .= states[i,j,1:end-1] .- old_states[i,j,1:end-1])
        end
    end

    for i in 1:size(obs,1)
        for j in 1:size(obs,2)
            old_obs[i,j,end] = sum([obs[i,j,k] for k in 1:(size(obs,3)-1)]) - sum([old_obs[i,j,k] for k in 1:(size(obs,3)-1)])
            full_shock_decomp && (old_obs[i,j,1:end-1] .= obs[i,j,1:end-1] .- old_obs[i,j,1:end-1])
        end
    end

    for i in 1:size(pseudo,1)
        for j in 1:size(pseudo,2)
            old_pseudo[i,j,end] = sum([pseudo[i,j,k] for k in 1:(size(pseudo,3)-1)]) - sum([old_pseudo[i,j,k] for k in 1:(size(pseudo,3)-1)])
            full_shock_decomp && (old_pseudo[i,j,1:end-1] .= pseudo[i,j,1:end-1] .- old_pseudo[i,j,1:end-1])
        end
    end

    # Return shock decompositions in appropriate range
    if start_index == 1 && end_index == allperiods
        return old_states, old_obs, old_pseudo
    else
        range = start_index:end_index
        return old_states[:, range, :], old_obs[:, range, :], old_pseudo[:, range, :]
    end
end

"""
```
deterministic_trends(m, system, z0)

deterministic_trends(system, z0, nperiods, start_index, end_index)

deterministic_trends(m, system, z0, start_date, end_date)

deterministic_trends(m, system, z0, nperiods, start_index, end_index,
    regime_inds, regimes, cond_type)

deterministic_trends(m, system, old_system, z0, start_date, end_date)

deterministic_trends(m, system, old_system, z0, nperiods, start_index, end_index,
    regime_inds, regimes, cond_type)
```

Compute deterministic trend values of states, observables, and
pseudo-observables, given a model object and system matrices. The deterministic
trend for a single draw is simply the series that would be obtained by iterating
the state-space system forward, beginning from a state vector `z0` in the last
presample period.

### Inputs

- `m::AbstractDSGEModel`: model object
- `system::System{S}` or `RegimeSwitchingSystem`: state-space system matrices
- `z0`::Vector{S}: initial state vector
- `start_date::Date`: initial date for deterministic trends
- `end_date::Date`: final date for deterministic trends
- `regime_inds::Vector{UnitRange}`: indices of the data corresponding to each regime.
- `regimes::UnitRange`: which regimes are involved in the date range for which
    we want to compute the deterministic trends
- `old_system::`RegimeSwitchingSystem`: state-space system matrices for old model

where `S<:AbstractFloat`.

### Outputs

- `states::Matrix{S}`: matrix of size `nstates` x `nperiods` of state
  steady-state values
- `obs::Matrix{S}`: matrix of size `nobs` x `nperiods` of observable
  steady-state values
- `pseudo::Matrix{S}`: matrix of size `npseudo` x `nperiods` of
  pseudo-observable steady-state values

where `nperiods` is the number of quarters between
`date_shockdec_start(m)` and `date_shockdec_end(m)`, inclusive.
"""
function deterministic_trends(m::AbstractDSGEModel{S},
                              system::System{S}, z0::Vector{S}) where {S<:AbstractFloat}

    # Dates: We compute the deterministic trend starting from the
    # first historical period.  However, since it is only used to
    # compute shock decompositions, we truncate and only store
    # results for periods corresponding to the shockdec period.
    nperiods    = subtract_quarters(date_forecast_end(m), date_mainsample_start(m)) + 1
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    deterministic_trends(system, z0, nperiods, start_index, end_index)
end

function deterministic_trends(system::System{S}, z0::Vector{S}, nperiods::Int,
                              start_index::Int, end_index::Int) where {S<:AbstractFloat}

    # Set constant system matrices to 0
    system = zero_system_constants(system)

    # Construct matrix of 0 shocks for entire history and forecast horizon
    nshocks  = size(system[:RRR], 2)
    shocks   = zeros(S, nshocks, nperiods)

    # Use forecast to iterate state-space system forward without shocks or the ZLB procedure
    states, obs, pseudo, _ = forecast(system, z0, shocks)

    # Return deterministic trends in appropriate range
    if start_index == 1 && end_index == nperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range], obs[:, range], pseudo[:, range]
    end
end

function deterministic_trends(m::AbstractDSGEModel{S},
                              system::RegimeSwitchingSystem{S}, z0::Vector{S},
                              start_date::Dates.Date = date_presample_start(m),
                              end_date::Dates.Date = prev_quarter(date_forecast_start(m)),
                              cond_type::Symbol = :none) where {S<:AbstractFloat}

    # Dates: We compute the deterministic trend starting from the
    # first historical period.  However, since it is only used to
    # compute shock decompositions, we truncate and only store
    # results for periods corresponding to the shockdec period.
    nperiods    = subtract_quarters(date_forecast_end(m), date_mainsample_start(m)) + 1
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Now intersect these periods with regime indices
    regime_inds = regime_indices(m, start_date, end_date) # do not need to account for ZLB split b/c shocks don't matter for trends
    if regime_inds[1][1] < 1
        regime_inds[1] = 1:regime_inds[1][end]
    end
    if regime_inds[end][end] > end_index
        regime_inds[end] = regime_inds[end][1]:end_index # if the end index is in the middle of a regime or is past the regime's end
    end
    if length(regime_inds[end]) == 0
        pop!(regime_inds)
    end

    return deterministic_trends(m, system, z0, nperiods, start_index, end_index, regime_inds, cond_type)
end

function deterministic_trends(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S}, z0::Vector{S}, nperiods::Int,
                              start_index::Int, end_index::Int, regime_inds::Vector{UnitRange{Int}},
                              cond_type::Symbol) where {S<:AbstractFloat}

    # Set constant system matrices to 0
    system = zero_system_constants(system)

    # Initialize storage matrices
    nshocks = size(system[1, :RRR], 2) # this is used in the forecasting loops to create shock matrices
    nstates = size(system[1, :TTT], 1)
    nobs    = size(system[1, :ZZ], 1)
    npseudo = size(system[1, :ZZ_pseudo], 1)
    states = Matrix{S}(undef, nstates, nperiods)
    obs    = Matrix{S}(undef, nobs,    nperiods)
    pseudo = Matrix{S}(undef, npseudo, nperiods)

    # Use forecast to iterate state-space system forward
    # without shocks or the ZLB procedure during history
    for (reg_num, reg_ind) in enumerate(regime_inds)
        init_state = (reg_num == 1) ? z0 : states[:, reg_ind[1] - 1] # update initial state

        # Construct matrix of 0 shocks for this regime
        shocks = zeros(S, nshocks, length(reg_ind))
        states[:, reg_ind], obs[:, reg_ind], pseudo[:, reg_ind], _ = forecast(system[reg_num], init_state, shocks)
    end

    # Use regime-switching forecast to iterate state-space system forward
    # during the forecast horizon
    if nperiods > regime_inds[end][end]
        fcast_shocks = zeros(S, nshocks, nperiods - regime_inds[end][end])
        fcast_inds = (regime_inds[end][end] + 1):nperiods

        states[:, fcast_inds], obs[:, fcast_inds], pseudo[:, fcast_inds], _ =
            forecast(m, system, states[:, regime_inds[end][end]], fcast_shocks;
                     cond_type = cond_type)
    end

    if start_index == 1 && end_index == nperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range], obs[:, range], pseudo[:, range]
    end
end

# Get Deterministic Trends with an old system
function deterministic_trends(m::AbstractDSGEModel{S},
                              system::RegimeSwitchingSystem{S}, old_system::RegimeSwitchingSystem{S},
                              z0::Vector{S},
                              start_date::Dates.Date = date_presample_start(m),
                              end_date::Dates.Date = prev_quarter(date_forecast_start(m)),
                              cond_type::Symbol = :none) where {S<:AbstractFloat}

    # Dates: We compute the deterministic trend starting from the
    # first historical period.  However, since it is only used to
    # compute shock decompositions, we truncate and only store
    # results for periods corresponding to the shockdec period.
    nperiods    = subtract_quarters(date_forecast_end(m), date_mainsample_start(m)) + 1
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Now intersect these periods with regime indices
    regime_inds = regime_indices(m, start_date, end_date) # do not need to account for ZLB split b/c shocks don't matter for trends
    if regime_inds[1][1] < 1
        regime_inds[1] = 1:regime_inds[1][end]
    end
    if regime_inds[end][end] > end_index
        regime_inds[end] = regime_inds[end][1]:end_index # if the end index is in the middle of a regime or is past the regime's end
    end
    if length(regime_inds[end]) == 0
        pop!(regime_inds)
    end

    return deterministic_trends(m, system, old_system, z0, nperiods, start_index, end_index, regime_inds, cond_type)
end

function deterministic_trends(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S}, old_system::RegimeSwitchingSystem{S},
                              z0::Vector{S}, nperiods::Int,
                              start_index::Int, end_index::Int, regime_inds::Vector{UnitRange{Int}},
                              cond_type::Symbol) where {S<:AbstractFloat}

    # Set constant system matrices to 0
    system = zero_system_constants(system)
    old_system = zero_system_constants(old_system)

    # Initialize storage matrices
    nshocks = size(system[1, :RRR], 2) # this is used in the forecasting loops to create shock matrices
    nstates = size(system[1, :TTT], 1)
    nobs    = size(system[1, :ZZ], 1)
    npseudo = size(system[1, :ZZ_pseudo], 1)
    states = Matrix{S}(undef, nstates, nperiods)
    obs    = Matrix{S}(undef, nobs,    nperiods)
    pseudo = Matrix{S}(undef, npseudo, nperiods)

    # Old System
    states_old = Matrix{S}(undef, nstates, nperiods)
    obs_old    = Matrix{S}(undef, nobs,    nperiods)
    pseudo_old = Matrix{S}(undef, npseudo, nperiods)

    # Use forecast to iterate state-space system forward
    # without shocks or the ZLB procedure during history
    for (reg_num, reg_ind) in enumerate(regime_inds)
        init_state = (reg_num == 1) ? z0 : states[:, reg_ind[1] - 1] # update initial state

        # Construct matrix of 0 shocks for this regime
        shocks = zeros(S, nshocks, length(reg_ind))
        states[:, reg_ind], obs[:, reg_ind], pseudo[:, reg_ind], _ = forecast(system[reg_num], init_state, shocks)

        # Do for old system too
        old_system2 = old_system[reg_num]#reg_ind[end] >= old_system ? system[1] : system[reg_num]
        init_state_old = (reg_num == 1) ? z0 : states_old[:, reg_ind[1] - 1] # update initial state
        states_old[:, reg_ind], obs_old[:, reg_ind], pseudo_old[:, reg_ind], _ = forecast(old_system2, init_state, shocks)
    end

    # Use regime-switching forecast to iterate state-space system forward
    # during the forecast horizon
    if nperiods > regime_inds[end][end]
        fcast_shocks = zeros(S, nshocks, nperiods - regime_inds[end][end])
        fcast_inds = (regime_inds[end][end] + 1):nperiods

        states[:, fcast_inds], obs[:, fcast_inds], pseudo[:, fcast_inds], _ =
            forecast(m, system, states[:, regime_inds[end][end]], fcast_shocks;
                     cond_type = cond_type)

        # Old System
        ##TODO: (Same in shock_decompositions) - currently assuming that we use
        ## Taylor in all forecast regimes, but generalize based on regime_inds[end][end]
        ## and old_system values
        states_old[:, fcast_inds], obs_old[:, fcast_inds], pseudo_old[:, fcast_inds], _ =
            forecast(m, old_system, states_old[:, regime_inds[end][end]], fcast_shocks, cond_type = cond_type)
        #forecast(old_system, states_old[:, regime_inds[end][end]], fcast_shocks)
    end

    # Old System Update
    states = states .- states_old
    obs = obs .- obs_old
    pseudo = pseudo .- pseudo_old

    if start_index == 1 && end_index == nperiods
        return states, obs, pseudo
    else
        range = start_index:end_index
        return states[:, range], obs[:, range], pseudo[:, range]
    end
end

"""
```
trends(system::System{S}) where {S<:AbstractFloat}
trends(system::RegimeSwitchingSystem{S}) where {S<:AbstractFloat}
trends(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S},
                start_date::Dates.Date = date_presample_start(m),
                end_date::Dates.Date = prev_quarter(date_forecast_start(m)),
                cond_type::Symbol) where {S<:AbstractFloat}
```

Compute trend (steady-state) states, observables, and pseudo-observables. The
trend is used for plotting shock decompositions. The first method
applies to non-regime switching systems, the second to regime-switching systems
that do not involve time variation in the `CCC` or `DD` vectors,
and the third to regime-switching systems with time variation in `CCC` or `DD`.
"""
function trends(system::System{S}) where {S<:AbstractFloat}

    state_trend  = system[:CCC]
    obs_trend    = system[:ZZ] * system[:CCC] + system[:DD]
    pseudo_trend = system[:ZZ_pseudo] * system[:CCC] + system[:DD_pseudo]

    return state_trend, obs_trend, pseudo_trend
end

function trends(system::RegimeSwitchingSystem{S}) where {S<:AbstractFloat}

    nreg  = n_regimes(system)
    input = 1:nreg

    state_trends  = Matrix{S}(undef, length(system[1, :CCC]), nreg)
    obs_trends    = Matrix{S}(undef, length(system[1, :DD]), nreg)
    pseudo_trends = Matrix{S}(undef, length(system[1, :DD_pseudo]), nreg)
    for i in input
        state_trends[:, i] .= 0.
        obs_trends[:, i]    = system[i, :DD]
        pseudo_trends[:, i] = system[i, :DD_pseudo]
    end

    return state_trends, obs_trends, pseudo_trends
end

# Handles case of time-varying CCC and DD
function trends(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S},
                start_date::Dates.Date = date_presample_start(m),
                end_date::Dates.Date = prev_quarter(date_forecast_start(m)),
                cond_type::Symbol = :none) where {S<:AbstractFloat}

    # Dates: We compute the trend starting from the
    # first historical period.  However, since it is only used to
    # compute shock decompositions, we truncate and only store
    # results for periods corresponding to the shockdec period.
    nperiods    = subtract_quarters(date_forecast_end(m), date_mainsample_start(m)) + 1
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Now calculate the regime indices in history and forecast
    hist_regime_inds = regime_indices(m, start_date, end_date) # do not need to account for ZLB split b/c shocks don't matter for trends
    if hist_regime_inds[1][1] < 1
        hist_regime_inds[1] = 1:hist_regime_inds[1][end]
    end
    if hist_regime_inds[end][end] >= end_index  # if the end index is in the middle of a regime or is past the regime's end
        hist_regime_inds[end] = hist_regime_inds[end][1]:end_index
        fcast_regime_inds = nothing
    else
        fcast_regime_inds = get_fcast_regime_inds(m, forecast_horizons(m; cond_type = cond_type), cond_type,
                                                  start_index = hist_regime_inds[end][end])
        fcast_cutoff = findfirst([regind[end] >= end_index for regind in fcast_regime_inds])
        if isnothing(fcast_cutoff)
            error("The index_shockdec_end(m) occurs past the index of the final forecast date.")
        end
        fcast_regime_inds = fcast_regime_inds[1:fcast_cutoff]
        fcast_regime_inds[end] = fcast_regime_inds[end][1]:end_index
    end
    if length(hist_regime_inds[end]) == 0
        pop!(hist_regime_inds)
    end

    # Initialize output
    state_trends  = Matrix{S}(undef, length(system[1, :CCC]), end_index)
    obs_trends    = Matrix{S}(undef, length(system[1, :DD]), end_index)
    pseudo_trends = Matrix{S}(undef, length(system[1, :DD_pseudo]), end_index)

    # Check if any CCC are nonzero, otherwise we can do a faster version
    first_nonzero_CCC = findfirst([any(.!(system[i, :CCC] .≈ 0.)) for i in 1:n_regimes(system)])
    if isnothing(first_nonzero_CCC)
        state_trends .= 0.
        for (reg, inds) in enumerate(hist_regime_inds)
            for t in inds
                obs_trends[:, t]    = system[reg, :DD]
                pseudo_trends[:, t] = system[reg, :DD_pseudo]
            end
        end
        if !isnothing(fcast_regime_inds)
            regbuffer = cond_type == :none ? get_setting(m, :reg_forecast_start) - 1 :
                get_setting(m, :reg_post_conditional_end) - 1
            for (i, inds) in enumerate(fcast_regime_inds)
                reg = i + regbuffer
                for t in inds
                    obs_trends[:, t]    = system[reg, :DD]
                    pseudo_trends[:, t] = system[reg, :DD_pseudo]
                end
            end
        end
    else
        first_nonzero_CCC_hist, first_nonzero_CCC_fcast = if first_nonzero_CCC > length(hist_regime_inds)
            length(hist_regime_inds) + 1, first_nonzero_CCC - length(hist_regime_inds)
        else
            first_nonzero_CCC, 1
        end

        for (reg, inds) in enumerate(hist_regime_inds[1:(first_nonzero_CCC_hist - 1)])
            state_trends[:, inds]  .= 0.
            obs_trends[:, inds]    .= system[reg, :DD]
            pseudo_trends[:, inds] .= system[reg, :DD_pseudo]
        end

        for (i, inds) in enumerate(hist_regime_inds[first_nonzero_CCC_hist:end])
            reg = i + first_nonzero_CCC_hist - 1
            for t in inds
                state_trends[:, t]  = system[reg, :CCC] + system[reg, :TTT] * state_trends[:, t - 1]
                obs_trends[:, t]    = system[reg, :ZZ] * state_trends[:, t] + system[reg, :DD]
                pseudo_trends[:, t] = system[reg, :ZZ_pseudo] * state_trends[:, t] + system[reg, :DD_pseudo]
            end
        end

        if !isnothing(fcast_regime_inds)
            reg_buffer = cond_type == :none ? get_setting(m, :reg_forecast_start) - 1 :
                get_setting(m, :reg_post_conditional_end) - 1
            for (i, inds) in enumerate(fcast_regime_inds[1:(first_nonzero_CCC_fcast - 1)])
                reg = i + reg_buffer
                for t in inds
                    state_trends[:, inds]  .= 0.
                    obs_trends[:, inds]    .= system[reg, :DD]
                    pseudo_trends[:, inds] .= system[reg, :DD_pseudo]
                end
            end

            for (i, inds) in enumerate(fcast_regime_inds[first_nonzero_CCC_fcast:end])
                reg = i + first_nonzero_CCC_fcast - 1 + reg_buffer
                for t in inds
                    state_trends[:, t]  = system[reg, :CCC] + system[reg, :TTT] * state_trends[:, t - 1]
                    obs_trends[:, t]    = system[reg, :ZZ] * state_trends[:, t] + system[reg, :DD]
                    pseudo_trends[:, t] = system[reg, :ZZ_pseudo] * state_trends[:, t] + system[reg, :DD_pseudo]
                end
            end
        end
    end

    if start_index == 1
        return state_trends, obs_trends, pseudo_trends
    else
        return state_trends[:, start_index:end], obs_trends[:, start_index:end], pseudo_trends[:, start_index:end]
    end
end

# Sequence of Shock Decompositions (Effect of time t shock on obs/pseudo/states at each time for all t)
"""
```
shock_decompositions_sequence(m, system, histshocks, n_back, back_shocks)

shock_decompositions_sequence(system, forecast_horizons, histshocks, start_index,
    end_index, n_back, back_shocks)

shock_decompositions_sequence(m, system, histshocks, start_date, end_date, n_back, back_shocks)

shock_decompositions_sequence(m, system, forecast_horizons, histshocks, start_index,
    end_index, regime_inds, cond_type, n_back, back_shocks)

shock_decompositions_sequence(m, system, old_system,
                              histshocks, old_histshocks, start_date, end_date,
                              cond_type; full_shock_decomp, n_back, back_shocks)

shock_decompositions_sequence(m, system, old_system,
                              forecast_horizons, histshocks, old_histshocks,
                              start_index, end_index,
                              regime_inds, cond_type; full_shock_decomp, n_back, back_shocks)
```

### Inputs

- `system::System{S}` or `RegimeSwitchingSystem{S}`: state-space system matrices
- `forecast_horizons::Int`: number of periods ahead to forecast
- `histshocks::Matrix{S}`: matrix of size `nshocks` x `hist_periods` of
  historical smoothed shocks
- `start_index::Int`: first index from which to return computed shock
  decompositions
- `end_index::Int`: last index for which to return computed shock decompositions
- `start_date::Date`: initial date for deterministic trends
- `end_date::Date`: final date for deterministic trends
- `regime_inds::Vector{UnitRange}`: indices of the data corresponding to each regime.
- `old_system::RegimeSwitchingSystem{S}`: state-space system matrices for the old model
- `old_histshocks::Matrix{S}`: matrix of size `nshocks` x `hist_periods` of
  historical smoothed shocks using the old system
- `full_shock_decomp::Bool`: If true for old_system case, return difference in shockdecs
    between all shocks for the new and old systems. Else, return old system's shocks
    with 1 cumulative AIT shock added on.

where `S<:AbstractFloat`.

### Outputs

- `states::Array{S, 4}`: matrix of size `nhistperiods` x `nstates` x `nperiods` x `nshocks` of
  state shock decompositions
- `obs::Array{S, 4}`: matrix of size `nhistperiods` x `nobs` x `nperiods` x `nshocks` of
  observable shock decompositions
- `pseudo::Array{S, 4}`: matrix of size `nhistperiods` x `npseudo` x `nperiods` x `nshocks` of
  pseudo-observable shock decompositions

where `nperiods = `end_index - start_index + 1`
and `nhistperiods` = `size(histshocks,2)`
"""
function shock_decompositions_sequence(m::AbstractDSGEModel,
    system::System{S}, histshocks::Matrix{S};
    n_back::Int64 = 0, back_shocks::Vector{Symbol} = Symbol[]) where {S<:AbstractFloat}

    horizon   = forecast_horizons(m)
    start_ind = index_shockdec_start(m)
    end_ind   = index_shockdec_end(m)

    shock_decompositions_sequence(system, horizon, histshocks, start_ind, end_ind, n_back = n_back, back_shocks = back_shocks)
end

function shock_decompositions_sequence(system::System{S},
    forecast_horizons::Int, histshocks::Matrix{S},
    start_index::Int, end_index::Int;
    n_back::Int64 = 0, back_shocks::Vector{Symbol} = Symbol[]) where {S<:AbstractFloat}

    # Setup
    nshocks      = size(system[:RRR], 2)
    nstates      = size(system[:TTT], 2)
    nobs         = size(system[:ZZ], 1)
    npseudo      = size(system[:ZZ_pseudo], 1)
    histperiods  = size(histshocks, 2)
    allperiods   = histperiods + forecast_horizons

    dim3_len = min(end_index-start_index+1, allperiods)
    states = zeros(S, histperiods, nstates, dim3_len, nshocks)
    obs    = zeros(S, histperiods, nobs, dim3_len, nshocks)
    pseudo = zeros(S, histperiods, npseudo, dim3_len, nshocks)

    for t in 1:histperiods
        histshocks_t = copy(histshocks)
        histshocks_t[:,vcat(1:t-1,t+1:histperiods)] .= 0.0

        states[t,:,:,:], obs[t,:,:,:], pseudo[t,:,:,:] = shock_decompositions(system, forecast_horizons,
                                                                                histshocks_t, start_index, end_index)
    end

    # Get the matrices back to 3 dimensions by having bars for contemporaneous shocks and shocks each of n_back time periods back.
    ## n_back is given by user and shocks for which we don't just have contemporaneous shocks is given by user.
    states3 = zeros(size(states,2), size(states,3), size(states,4) + (n_back+1)*length(back_shocks))
    obs3 = zeros(size(obs,2), size(obs,3), size(obs,4) + (n_back+1)*length(back_shocks))
    pseudo3 = zeros(size(pseudo,2), size(pseudo,3), size(pseudo,4) + (n_back+1)*length(back_shocks))

    ## Get indices for vector of shocks
    shock_inds = [m.exogenous_shocks[shock_ind] for shock_ind in back_shocks]
    shock_len = length(back_shocks)

    ## Actually do the conversion
    for i in 1:size(states,3)
        st_ind = min(start_index+i-1, size(states,1))
        states3[:,i,1:size(states,4)] = states[st_ind,:,i,:]
        obs3[:,i,1:size(states,4)] = obs[st_ind,:,i,:]
        pseudo3[:,i,1:size(states,4)] = pseudo[st_ind,:,i,:]

        for k in 1:n_back
            st_ind2 = min(start_index+i-1-k, size(states,1)-k)
            for j in 1:length(back_shocks)
                states3[:,i,size(states,4)+(k-1)*shock_len+j] = states[st_ind2,:,i,shock_inds[j]]
                obs3[:,i,size(obs,4)+(k-1)*shock_len+j] = obs[st_ind2,:,i,shock_inds[j]]
                pseudo3[:,i,size(pseudo,4)+(k-1)*shock_len+j] = pseudo[st_ind2,:,i,shock_inds[j]]

                # Add this "lagged shock" to the list of exogenous shocks
                if !haskey(m.exogenous_shocks, Symbol(back_shocks[j], :_, k))
                    m.exogenous_shocks[Symbol(back_shocks[j], :_, k)] = length(m.exogenous_shocks)+1
                end
            end
        end
        ## Add in a bar that sums up shocks before n_back periods back.
        st_ind3 = min(start_index+i-2-n_back, size(states,1)-n_back-1)
        for j in 1:length(back_shocks)
            states3[:,i,size(states,4)+n_back*shock_len+j] = sum(states[1:st_ind3,:,i,shock_inds[j]],dims=1)
            obs3[:,i,size(obs,4)+n_back*shock_len+j] = sum(obs[1:st_ind3,:,i,shock_inds[j]],dims=1)
            pseudo3[:,i,size(pseudo,4)+n_back*shock_len+j] = sum(pseudo[1:st_ind3,:,i,shock_inds[j]],dims=1)

            # Add this "lagged shock" to the list of exogenous shocks
            if !haskey(m.exogenous_shocks, Symbol(back_shocks[j], :_History))
                m.exogenous_shocks[Symbol(back_shocks[j], :_History)] = length(m.exogenous_shocks)+1
            end
        end
    end

    return states3, obs3, pseudo3
end

function shock_decompositions_sequence(m::AbstractDSGEModel{S},
                              system::RegimeSwitchingSystem{S}, histshocks::Matrix{S},
                              start_date::Dates.Date = date_presample_start(m),
                              end_date::Dates.Date = prev_quarter(date_forecast_start(m)),
                              cond_type::Symbol = :none;
                              n_back::Int64 = 0, back_shocks::Vector{Symbol} = Symbol[]) where {S<:AbstractFloat}
    # Set up dates
    horizon     = forecast_horizons(m; cond_type = cond_type)
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Set up regime indices. Note that we do not need to account
    # for ZLB split b/c `histshocks` should have zeros for anticipated shocks
    # in the pre-ZLB periods.
    regime_inds = regime_indices(m, start_date, end_date)
    if regime_inds[1][1] < 1 # remove periods occuring before desired start date
        regime_inds[1] = 1:regime_inds[1][end]
    end

    shock_decompositions_sequence(m, system, horizon, histshocks, start_index, end_index,
                                  regime_inds, cond_type, n_back = n_back, back_shocks = back_shocks)
end

function shock_decompositions_sequence(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S},
                              forecast_horizons::Int, histshocks::Matrix{S},
                              start_index::Int, end_index::Int,
                              regime_inds::Vector{UnitRange{Int}},
                              cond_type::Symbol;
                              n_back::Int64 = 0, back_shocks::Vector{Symbol} = Symbol[]) where {S<:AbstractFloat}

    # Setup
    nshocks     = size(system[1, :RRR], 2)
    nstates     = size(system[1, :TTT], 2)
    nobs        = size(system[1, :ZZ], 1)
    npseudo     = size(system[1, :ZZ_pseudo], 1)
    histperiods = size(histshocks, 2)
    allperiods  = histperiods + forecast_horizons

    states = zeros(S, nstates, allperiods, nshocks)
    obs    = zeros(S, nobs,    allperiods, nshocks)
    pseudo = zeros(S, npseudo, allperiods, nshocks)
    shocks = zeros(S, nshocks, histperiods)

    dim3_len = min(end_index-start_index+1, allperiods)
    states = zeros(S, histperiods, nstates, dim3_len, nshocks)
    obs    = zeros(S, histperiods, nobs, dim3_len, nshocks)
    pseudo = zeros(S, histperiods, npseudo, dim3_len, nshocks)

    for t in 1:histperiods
        histshocks_t = copy(histshocks)
        histshocks_t[:,vcat(1:t-1,t+1:histperiods)] .= 0.0

        states[t,:,:,:], obs[t,:,:,:], pseudo[t,:,:,:] = shock_decompositions(m, system, forecast_horizons,
                                                                              histshocks_t, start_index, end_index,
                                                                              regime_inds, cond_type)
    end

    # Get the matrices back to 3 dimensions by having bars for contemporaneous shocks and shocks each of n_back time periods back.
    ## n_back is given by user and shocks for which we don't just have contemporaneous shocks is given by user.
    states3 = zeros(size(states,2), size(states,3), size(states,4) + (n_back+1)*length(back_shocks))
    obs3 = zeros(size(obs,2), size(obs,3), size(obs,4) + (n_back+1)*length(back_shocks))
    pseudo3 = zeros(size(pseudo,2), size(pseudo,3), size(pseudo,4) + (n_back+1)*length(back_shocks))

    ## Get indices for vector of shocks
    shock_inds = [m.exogenous_shocks[shock_ind] for shock_ind in back_shocks]
    shock_len = length(back_shocks)

    ## Actually do the conversion
    for i in 1:size(states,3)
        st_ind = min(start_index+i-1, size(states,1))
        states3[:,i,1:size(states,4)] = states[st_ind,:,i,:]
        obs3[:,i,1:size(states,4)] = obs[st_ind,:,i,:]
        pseudo3[:,i,1:size(states,4)] = pseudo[st_ind,:,i,:]

        for k in 1:n_back
            st_ind2 = min(start_index+i-1-k, size(states,1)-k)
            for j in 1:length(back_shocks)
                states3[:,i,size(states,4)+(k-1)*shock_len+j] = states[st_ind2,:,i,shock_inds[j]]
                obs3[:,i,size(obs,4)+(k-1)*shock_len+j] = obs[st_ind2,:,i,shock_inds[j]]
                pseudo3[:,i,size(pseudo,4)+(k-1)*shock_len+j] = pseudo[st_ind2,:,i,shock_inds[j]]

                # Add this "lagged shock" to the list of exogenous shocks
                if !haskey(m.exogenous_shocks, Symbol(back_shocks[j], :_, k))
                    m.exogenous_shocks[Symbol(back_shocks[j], :_, k)] = length(m.exogenous_shocks)+1
                end
            end
        end
        ## Add in a bar that sums up shocks before n_back periods back.
        st_ind3 = min(start_index+i-2-n_back, size(states,1)-n_back-1)
        for j in 1:length(back_shocks)
            states3[:,i,size(states,4)+n_back*shock_len+j] = sum(states[1:st_ind3,:,i,shock_inds[j]],dims=1)
            obs3[:,i,size(obs,4)+n_back*shock_len+j] = sum(obs[1:st_ind3,:,i,shock_inds[j]],dims=1)
            pseudo3[:,i,size(pseudo,4)+n_back*shock_len+j] = sum(pseudo[1:st_ind3,:,i,shock_inds[j]],dims=1)

            # Add this "lagged shock" to the list of exogenous shocks
            if !haskey(m.exogenous_shocks, Symbol(back_shocks[j], :_History))
                m.exogenous_shocks[Symbol(back_shocks[j], :_History)] = length(m.exogenous_shocks)+1
            end
        end
    end

    return states3, obs3, pseudo3
end

# Get Shock Decompositions with an old system and a bar that gives
## effect of using new system compared to old.
function shock_decompositions_sequence(m::AbstractDSGEModel{S},
                              system::RegimeSwitchingSystem{S}, old_system::RegimeSwitchingSystem{S},
                              histshocks::Matrix{S}, old_histshocks::Matrix{S},
                              start_date::Dates.Date = date_presample_start(m),
                              end_date::Dates.Date = prev_quarter(date_forecast_start(m)),
                              cond_type::Symbol = :none; full_shock_decomp = true,
                              n_back::Int64 = 0, back_shocks::Vector{Symbol} = Symbol[]) where {S<:AbstractFloat}
    # Set up dates
    horizon     = forecast_horizons(m; cond_type = cond_type)
    start_index = index_shockdec_start(m)
    end_index   = index_shockdec_end(m)

    # Set up regime indices. Note that we do not need to account
    # for ZLB split b/c `histshocks` should have zeros for anticipated shocks
    # in the pre-ZLB periods.
    regime_inds = regime_indices(m, start_date, end_date)
    if regime_inds[1][1] < 1 # remove periods occuring before desired start date
        regime_inds[1] = 1:regime_inds[1][end]
    end

    shock_decompositions_sequence(m, system, old_system, horizon, histshocks, old_histshocks, start_index, end_index,
                                  regime_inds, cond_type, full_shock_decomp = full_shock_decomp,
                                  n_back = n_back, back_shocks = back_shocks)
end

function shock_decompositions_sequence(m::AbstractDSGEModel, system::RegimeSwitchingSystem{S},
                                       old_system::RegimeSwitchingSystem{S}, forecast_horizons::Int,
                                       histshocks::Matrix{S}, old_histshocks::Matrix{S},
                                       start_index::Int, end_index::Int,
                                       regime_inds::Vector{UnitRange{Int}},
                                       cond_type::Symbol; full_shock_decomp = true,
                                       n_back::Int64 = 0, back_shocks::Vector{Symbol} = Symbol[]) where {S<:AbstractFloat}

    # Setup
    nshocks     = size(system[1, :RRR], 2)
    nstates     = size(system[1, :TTT], 2)
    nobs        = size(system[1, :ZZ], 1)
    npseudo     = size(system[1, :ZZ_pseudo], 1)
    histperiods = size(histshocks, 2)
    allperiods  = histperiods + forecast_horizons

    dim3_len = min(end_index-start_index+1, allperiods)
    states = zeros(S, histperiods, nstates, dim3_len, nshocks)
    obs    = zeros(S, histperiods, nobs, dim3_len, nshocks)
    pseudo = zeros(S, histperiods, npseudo, dim3_len, nshocks)

    for t in 1:histperiods
        histshocks_t = copy(histshocks)
        histshocks_t[:,vcat(1:t-1,t+1:histperiods)] .= 0.0

        old_histshocks_t = copy(old_histshocks)
        old_histshocks_t[:,vcat(1:t-1,t+1:histperiods)] .= 0.0

        states[t,:,:,:], obs[t,:,:,:], pseudo[t,:,:,:] = shock_decompositions(m, system, old_system,
                                                                              forecast_horizons,
                                                                              histshocks_t, old_histshocks_t,
                                                                              start_index, end_index,
                                                                              regime_inds, cond_type,
                                                                              full_shock_decomp = full_shock_decomp)
    end

    # Get the matrices back to 3 dimensions by having bars for contemporaneous shocks and shocks each of n_back time periods back.
    ## n_back is given by user and shocks for which we don't just have contemporaneous shocks is given by user.
    states3 = zeros(size(states,2), size(states,3), size(states,4) + (n_back+1)*length(back_shocks))
    obs3 = zeros(size(obs,2), size(obs,3), size(obs,4) + (n_back+1)*length(back_shocks))
    pseudo3 = zeros(size(pseudo,2), size(pseudo,3), size(pseudo,4) + (n_back+1)*length(back_shocks))

    ## Get indices for vector of shocks
    shock_inds = [m.exogenous_shocks[shock_ind] for shock_ind in back_shocks]
    shock_len = length(back_shocks)

    ## Actually do the conversion
    for i in 1:size(states,3)
        st_ind = min(start_index+i-1, size(states,1))
        states3[:,i,1:size(states,4)] = states[st_ind,:,i,:]
        obs3[:,i,1:size(states,4)] = obs[st_ind,:,i,:]
        pseudo3[:,i,1:size(states,4)] = pseudo[st_ind,:,i,:]

        for k in 1:n_back
            st_ind2 = min(start_index+i-1-k, size(states,1)-k)
            for j in 1:length(back_shocks)
                states3[:,i,size(states,4)+(k-1)*shock_len+j] = states[st_ind2,:,i,shock_inds[j]]
                obs3[:,i,size(obs,4)+(k-1)*shock_len+j] = obs[st_ind2,:,i,shock_inds[j]]
                pseudo3[:,i,size(pseudo,4)+(k-1)*shock_len+j] = pseudo[st_ind2,:,i,shock_inds[j]]

                # Add this "lagged shock" to the list of exogenous shocks
                if !haskey(m.exogenous_shocks, Symbol(back_shocks[j], :_, k))
                    m.exogenous_shocks[Symbol(back_shocks[j], :_, k)] = length(m.exogenous_shocks)+1
                end
            end
        end
        ## Add in a bar that sums up shocks before n_back periods back.
        st_ind3 = min(start_index+i-2-n_back, size(states,1)-n_back-1)
        for j in 1:length(back_shocks)
            states3[:,i,size(states,4)+n_back*shock_len+j] = sum(states[1:st_ind3,:,i,shock_inds[j]],dims=1)
            obs3[:,i,size(obs,4)+n_back*shock_len+j] = sum(obs[1:st_ind3,:,i,shock_inds[j]],dims=1)
            pseudo3[:,i,size(pseudo,4)+n_back*shock_len+j] = sum(pseudo[1:st_ind3,:,i,shock_inds[j]],dims=1)

            # Add this "lagged shock" to the list of exogenous shocks
            if !haskey(m.exogenous_shocks, Symbol(back_shocks[j], :_History))
                m.exogenous_shocks[Symbol(back_shocks[j], :_History)] = length(m.exogenous_shocks)+1
            end
        end
    end

    return states3, obs3, pseudo3
end


#module HARPSNData
#export read_data

import ..EchelleObservation, ..EchelleObservationSet, ..OrdersType
import ..AbstractInstrument
#, ..HARPSNType
#import ..harpsn_flux_hdu,  ..harpsn_all_orders
using FITSIO

"""
read_data(filename; orders)
Example:
   data = read_data(filename_list[1])
"""
function read_data(filename::String; orders::OrdersType = harpsn_all_orders )
   local fits
   try
      fits = FITS(filename)
   catch
      @error("# Can't read " * filename)
   end
   if typeof(orders) <: Integer
      orders = orders:orders
   end
   hdr = read_header(fits[1])
   # Get header info for estimating flux uncertainties
   gain = hdr["TNG DRS CCD CONAD"]
   ron = hdr["TNG DRS CCD SIGDET"]
   # Get barycentric Earth RV correction (should be positive in morning when observer is approaching sun)
   berv = hdr["TNG DRS BERV"]
   berv /= speed_of_light_mps
   flux = read(fits[harpsn_flux_hdu])[:,orders]
   flux *= gain # in ADU
   #flux .*= gain
   var = copy(flux)
   flux[flux .< zero(eltype(flux))] .= zero(eltype(flux))
   var .+= ron^2
   npix,nord = size(flux)
   # Extract pixel->wavelegnth polynomials (thank ACC)
   coeff = zeros(4,nord)
   for ord in 0:(nord-1)
      for poly in 0:3
         n = ord*4+poly
         coeff[1+poly,1+ord] = hdr["TNG DRS CAL TH COEFF LL" * string(n)]
         #push!(coeff,hdr["TNG DRS CAL TH COEFF LL" * string(n)])
         #n += 1
      end
   end
   xpix = 0:(npix-1)
   xarray =  vcat(ones(length(xpix))', xpix', (xpix.^2)', (xpix.^3)' )
   lambda = zeros(npix,length(orders))
   for ord in orders
      lambda[:,ord] .= xarray' * coeff[:,ord]
   end
   metadata = Dict{Symbol,Any}(:Filename=>filename,:Instrument=>:HARPSN,:z_bc=>berv)
#   return EchelleObservation{eltype(lambda),eltype(flux),eltype(var),typeof(berv),HARPSNType}(flux=flux,var=var,lambda=lambda,z_bc=berv,metadata=metadata)
   return EchelleObservation(flux=flux,var=var,lambda=lambda,z_bc=berv,instrument=HARPSNType(), metadata=metadata)
end

"""
   read_data!(data_out, obsid, filename; orders)
Writes to data_out[:,:,i]'s flux, var and lambda.
Creates metadata_list including :Filename and :z_bc (redshift of observer relative to solar system barcenter)

Example:
   data = read_data!(data_out, obsid, filename, orders=10:15)
"""
function read_data!(data::EchelleObservationSet{T1,T2,T3,#= T4, =# HARPSNType}, i::Int, filename::String;
      orders::OrdersType = harpsn_all_orders
      ) where {T1<:Number, T2<:Number, T3<:Number, #=T4<:Number,  HARPSNType <:AbstractInstrument =#}
   @assert 1 <= i <= size(data.flux,3)
   local fits
   try
      fits = FITS(filename)
   catch
      @error("# Can't read " * filename)
   end
   if typeof(orders) <: Integer
      orders = orders:orders
   end
   obs = read_data(filename, orders=orders)
   data.flux[:,:,i] .= obs.flux
   data.var[:,:,i] .= obs.var
   data.lambda[:,:,i] .= obs.lambda
   data.metadata_list[i] = Dict{Symbol,Any}(:Filename=>filename,:z_bc=>obs.metadata[:z_bc])
   return data
end

"""
   read_data(filenames; orders, z_bsrv)
Returns struct with flux, var, lambda, and metadata_list including :Filename and :z_bc
If pass z_bsrv, then adds that to the barrycentric correction stored in metadata_list.

Example:
   data = read_data(filename_list, orders=10:15)
"""
function read_data(filename_list::AbstractArray{String,1};
      orders::OrdersType = harpsn_all_orders,
      z_bsrv::AbstractArray = zeros(length(filename_list)) )
   @assert length(filename_list) >= 1
   data_first = read_data(filename_list[1], orders=orders)
   data = EchelleObservationSet{Float64,Float64,Float64,HARPSNType}(size(data_first.flux), length(filename_list), metadata=Dict(:Instrument=>:HARPSN) )
   data.lambda[:,:,1] = data_first.lambda
   data.flux[:,:,1]   = data_first.flux
   data.var[:,:,1]    = data_first.var
   for (i,filename) in enumerate(filename_list)
          read_data!(data, i, filename, orders=orders)
          data.metadata_list[i][:z_bc] += z_bsrv[i]
   end
   return data
end

default_harpsn_blaze_file = joinpath(homedir(),"Downloads","HARPN.2015-07-31T17-02-35.072_blaze_A.fits")
function read_blaze_list(filename::String = default_harpsn_blaze_file,
      pixel_buffer_lo::Integer = default_pixel_buffer_lo,
      pixel_buffer_hi::Integer = default_pixel_buffer_hi,
      log_lambda_lo::AbstractArray{Float64,1} = zeros(length(harpsn_all_orders)),
      log_lambda_hi::AbstractArray{Float64,1} = zeros(length(harpsn_all_orders)) )
   fits = FITS(filename)
   blaze_data = read(fits[1])
   pixels_per_order, num_orders = size(blaze_data)
   pixel_range = (1+default_pixel_buffer_lo):(pixels_per_order-default_pixel_buffer_lo)
   blaze_list = Array{BlazeMap1D{Int64,Float64,Float32},1}(undef,num_orders)
   @assert length(log_lambda_lo) == length(log_lambda_hi) == num_orders
   for i in 1:num_orders
      blaze_list[i] = BlazeMap1D(i,log_lambda_lo[i],log_lambda_hi[i],pixel_range, view(blaze_data,:,i) )
   end

   return BlazeMap2D(blaze_list)
end
export read_blaze_list

function read_blaze_list(data::EchelleObservationSet, filename::String = default_harpsn_blaze_file,
      pixel_buffer_lo::Integer = default_pixel_buffer_lo,
      pixel_buffer_hi::Integer = default_pixel_buffer_hi
      )
   pixels_per_order, num_orders, num_obs = size(data.lambda)
   pixel_range = (1+default_pixel_buffer_lo):(pixels_per_order-default_pixel_buffer_lo)
   log_lambda_lo = map(i->log(minimum(data.lambda[1,i,:])),1:num_orders)
   log_lambda_hi = map(i->log(maximum(data.lambda[end,i,:])),1:num_orders)
   read_blaze_list(filename,pixel_buffer_lo,pixel_buffer_hi, log_lambda_lo, log_lambda_hi )
 end

#end # module


module Manipulator

using LibUMP
import ..Gamepad.GamepadState

const UMP_STEP = [250, 500, 1000, 2000, 4000, 8000]
const UMP_SPEED = [5, 10, 25, 50, 100, 400]

function __init__()
    UMP.set_refresh_time_limit(LibUMP.DEF_HANDLE[], Cint(10))
    UMP.set_timeout(LibUMP.DEF_HANDLE[], Cint(10))
    UMP.set_slow_speed_mode(LibUMP.DEF_HANDLE[], Cint(1), Cint(1))
end

function step!(gamepad::GamepadState)
    X, Y = gamepad.R.x, gamepad.R.y
    LB, RB = gamepad.button.LB, gamepad.button.RB
    ystep = abs(Y) >= 0.4 ? flipsign(-UMP_STEP[speed_step], Y) : 0
    xstep = abs(X) >= 0.4 ? flipsign(-UMP_STEP[speed_step], X) : 0
    zstep = LB ? (RB ? 0 : -UMP_STEP[speed_step]) : (RB ? UMP_STEP[speed_step] : 0)
    if (xstep !== 0 || ystep !== 0 || zstep !== 0)
        move_3D_by(xstep, ystep, zstep, UMP_SPEED[speed_step])
    end
end

end # end


using heuristicopt

function main()
    population = reshape(rand(1:10, 100), (10, 10))
    crossoverRate = 0.5
    mutationRate = 0.1
    iterations = 100
    gaparam = GA.GAParameter(
                         crossoverRate,
                         mutationRate,
                         iterations
                         )

    select = GATool.selectrandom
    crossover = GATool.crossoverOne
    mutate = GATool.mutaterandom
    gafunc = GA.GAFunction(select, crossover, mutate)
    model = GA.GAModel(
        gaparam,
        gafunc,
        objfunc
    )
    predictPopulation, bestFitness = GA.predict(model, population)
    print("bestFitness", bestFitness)
end

function objfunc(individual)
    0
end

main()


struct SimulationInputs
    base_power::Float64
    base_frequency::Float64
    injectors_data::Vector{<:PSY.StaticInjection}
    static_injection_data::Vector{<:PSY.StaticInjection}
    dynamic_branches::Vector{PSY.DynamicBranch}
    counts::Base.ImmutableDict{Symbol, Int}
    Ybus::SparseArrays.SparseMatrixCSC{Complex{Float64}, Int}
    dyn_lines::Bool
    voltage_buses_ix::Vector{Int}
    current_buses_ix::Vector{Int}
    global_index::Dict{String, Dict{Symbol, Int}}
    total_shunts::SparseArrays.SparseMatrixCSC{Complex{Float64}, Int}
    global_vars::Dict{Symbol, Real}
    lookup::Dict{Int, Int}
    DAE_vector::Vector{Bool}
    mass_matrix::SparseArrays.SparseMatrixCSC{Float64, Int}
    aux_arrays::Dict{Int, Vector}
    tspan::NTuple{2, Float64}
end

function SimulationInputs(sys::PSY.System, tspan::NTuple{2, Float64} = (0.0, 0.0))
    injector_data = PSY.get_components(
        PSY.StaticInjection,
        sys,
        x -> PSY.get_dynamic_injector(x) !== nothing && PSY.get_available(x),
    )

    if isempty(injector_data)
        error("System doesn't contain any DynamicInjection devices")
    end

    Ybus, lookup = _get_Ybus(sys)
    n_buses = length(PSY.get_bus_numbers(sys))

    dynamic_branches =
        collect(PSY.get_components(PSY.DynamicBranch, sys, x -> PSY.get_available(x)))

    branch_state_counts = 2 * length(dynamic_branches)
    injector_state_count = sum(PSY.get_n_states.(PSY.get_dynamic_injector.(injector_data)))
    var_count = injector_state_count + 2 * n_buses + branch_state_counts

    mass_matrix = SparseArrays.sparse(LinearAlgebra.I, var_count, var_count)
    mass_matrix[1:(2 * n_buses), 1:(2 * n_buses)] .= 0.0

    static_injection_data = PSY.get_components(
        PSY.StaticInjection,
        sys,
        x -> PSY.get_dynamic_injector(x) === nothing && PSY.get_available(x),
    )

    counts = Base.ImmutableDict(
        :injection_n_states => injector_state_count,
        :branches_n_states => branch_state_counts,
        :first_dyn_injection_pointer => 2 * n_buses + branch_state_counts + 1,
        :first_dyn_branch_point => isempty(dynamic_branches) ? 0 : 2 * n_buses + 1,
        :total_variables => var_count,
        :bus_count => n_buses,
    )

    return SimulationInputs(
        PSY.get_base_power(sys),
        PSY.get_frequency(sys),
        collect(injector_data),
        collect(static_injection_data),
        dynamic_branches,
        counts,
        Ybus,
        !isempty(dynamic_branches),
        Vector{Int}(),
        collect(1:n_buses),
        MAPPING_DICT(),
        SparseArrays.spzeros(Complex{Float64}, n_buses, n_buses),
        GLOBAL_VARS_IX(),
        lookup,
        _init_DAE_vector!(var_count, n_buses),
        mass_matrix,
        Dict{Int, Vector{Real}}(),
        tspan,
    )
end

get_base_power(inputs::SimulationInputs) = inputs.base_power
get_base_frequency(inputs::SimulationInputs) = inputs.base_frequency
get_injectors_data(inputs::SimulationInputs) = inputs.injectors_data
get_dynamic_branches(inputs::SimulationInputs) = inputs.dynamic_branches
get_static_injections_data(inputs::SimulationInputs) = inputs.static_injection_data
get_counts(inputs::SimulationInputs) = inputs.counts
get_voltage_buses_ix(inputs::SimulationInputs) = inputs.voltage_buses_ix
get_current_buses_ix(inputs::SimulationInputs) = inputs.current_buses_ix
get_global_index(inputs::SimulationInputs) = inputs.global_index
get_Ybus(inputs::SimulationInputs) = inputs.Ybus
get_total_shunts(inputs::SimulationInputs) = inputs.total_shunts
get_global_vars(inputs::SimulationInputs) = inputs.global_vars
get_lookup(inputs::SimulationInputs) = inputs.lookup
get_DAE_vector(inputs::SimulationInputs) = inputs.DAE_vector
get_mass_matrix(inputs::SimulationInputs) = inputs.mass_matrix
get_aux_arrays(inputs::SimulationInputs) = inputs.aux_arrays
get_tspan(inputs::SimulationInputs) = inputs.tspan
has_dyn_lines(inputs::SimulationInputs) = inputs.dyn_lines

get_injection_pointer(inputs::SimulationInputs) =
    get_counts(inputs)[:first_dyn_injection_pointer]
get_branches_pointer(inputs::SimulationInputs) = get_counts(inputs)[:first_dyn_branch_point]
get_n_injection_states(inputs::SimulationInputs) = get_counts(inputs)[:injection_n_states]
get_n_branches_states(inputs::SimulationInputs) = get_counts(inputs)[:branches_n_states]
get_variable_count(inputs::SimulationInputs) = get_counts(inputs)[:total_variables]
get_device_index(inputs::SimulationInputs, device::D) where {D <: PSY.DynamicInjection} =
    get_global_index(inputs)[device.name]
get_bus_count(inputs::SimulationInputs) = get_counts(inputs)[:bus_count]
get_ω_sys(inputs::SimulationInputs) = get_global_vars(inputs)[:ω_sys]

function change_vector_type!(inputs::SimulationInputs, ::Type{T}) where {T <: Real}
    for d in get_injectors_data(inputs)
        attach_inner_vars!(PSY.get_dynamic_injector(d), T)
    end
    return
end

function simulation_pre_step!(inputs::SimulationInputs, ::Type{T}) where {T <: Real}
    add_aux_arrays!(inputs, T)
    change_vector_type!(inputs, T)
    return
end

function add_aux_arrays!(inputs::SimulationInputs, ::Type{T}) where {T <: Real}
    @debug "Auxiliary Arrays created with Type $(T)"
    bus_count = get_bus_count(inputs)
    get_aux_arrays(inputs)[1] = collect(zeros(T, bus_count))                       #I_injections_r
    get_aux_arrays(inputs)[2] = collect(zeros(T, bus_count))                       #I_injections_i
    get_aux_arrays(inputs)[3] = collect(zeros(T, get_n_injection_states(inputs)))  #injection_ode
    get_aux_arrays(inputs)[4] = collect(zeros(T, get_n_branches_states(inputs)))   #branches_ode
    get_aux_arrays(inputs)[5] = collect(zeros(Complex{T}, bus_count))              #I_bus
    get_aux_arrays(inputs)[6] = collect(zeros(T, 2 * bus_count))                   #I_balance
    return
end

function _get_Ybus(sys::PSY.System)
    n_buses = length(PSY.get_components(PSY.Bus, sys))
    dyn_lines = PSY.get_components(PSY.DynamicBranch, sys)
    if !isempty(PSY.get_components(PSY.ACBranch, sys))
        Ybus_ = PSY.Ybus(sys)
        Ybus = Ybus_[:, :]
        lookup = Ybus_.lookup[1]
        for br in dyn_lines
            ybus_update!(Ybus, br, lookup, -1.0)
        end
    else
        Ybus = SparseArrays.SparseMatrixCSC{Complex{Float64}, Int}(zeros(n_buses, n_buses))
        lookup = Dict{Int.Int}()
    end
    return Ybus, lookup
end

function _init_DAE_vector!(n_vars::Int, n_buses::Int)
    DAE_vector = trues(n_vars)
    DAE_vector[1:(n_buses * 2)] .= false
    return DAE_vector
end

function _dynamic_lines_inputs!(
    inputs::SimulationInputs,
    state_space_ix::Vector{Int},
    n_buses::Int,
)
    static_bus_var_count = 2 * n_buses
    branches_n_states = 0
    if inputs.dyn_lines
        global_state_index = inputs.global_index
        for br in get_dynamic_branches(inputs)
            index_dynamic_lines!(inputs, br, n_buses)
            add_states_to_global!(global_state_index, state_space_ix, br)
            branches_n_states += PSY.get_n_states(br)
        end

        for (ix, val) in enumerate(inputs.DAE_vector[1:n_buses])
            if val
                #This V_ix should be V_number.
                global_state_index["V_$(ix)"] = Dict(:R => ix, :I => ix + n_buses)
                static_bus_var_count -= 2
                push!(inputs.voltage_buses_ix, ix)
                @assert static_bus_var_count >= 0
            end
        end
        unique!(inputs.voltage_buses_ix)
        setdiff!(inputs.current_buses_ix, inputs.voltage_buses_ix)
    else
        @debug("System doesn't contain Dynamic Branches")
    end
    return branches_n_states, static_bus_var_count
end

function _static_injection_inputs!(inputs::SimulationInputs, ::Vector{Int}, sys::PSY.System)
    for s in PSY.get_components(PSY.StaticInjection, sys)
        index_static_injection(inputs, s)
    end
    return
end

function _dynamic_injection_inputs!(inputs::SimulationInputs, state_space_ix::Vector{Int})
    dynamic_injection_states = 0
    for d in get_injectors_data(inputs)
        @debug PSY.get_name(d)
        _attach_control_refs!(d)
        dynamic_injector = PSY.get_dynamic_injector(d)
        dynamic_injection_states += PSY.get_n_states(dynamic_injector)
        index_dynamic_injection(inputs, dynamic_injector, state_space_ix)
    end
    return dynamic_injection_states
end

function _mass_matrix_inputs!(inputs::SimulationInputs)
    mass_matrix = get_mass_matrix(inputs)
    injection_range = PSID.get_injection_pointer(inputs):PSID.get_variable_count(inputs)
    mass_matrix_injectors = view(mass_matrix, injection_range, injection_range)
    for d in PSID.get_injectors_data(inputs)
        dynamic_injector = PSY.get_dynamic_injector(d)
        device_mass_matrix_entries!(mass_matrix_injectors, dynamic_injector)
    end

    if has_dyn_lines(inputs)
        sys_f = get_base_frequency(inputs)
        shunts = get_total_shunts(inputs)
        n_buses = get_bus_count(inputs)
        for i in 1:n_buses
            val = imag(shunts[i, i])
            if val > 0
                mass_matrix[i, i] =
                    mass_matrix[i + n_buses, i + n_buses] = val * (1 / (2.0 * π * sys_f))
            end
        end
    end
end

# Default implementation for both models. This implementation is to future proof if there is
# a divergence between the required build methods
function _build!(inputs::SimulationInputs, sys::PSY.System)
    n_buses = get_bus_count(inputs)
    state_space_ix = Int[n_buses * 2]
    branches_n_states, static_bus_var_count =
        _dynamic_lines_inputs!(inputs, state_space_ix, n_buses)

    _static_injection_inputs!(inputs, state_space_ix, sys)
    injection_n_states = _dynamic_injection_inputs!(inputs, state_space_ix)

    set_frequency_reference!(inputs, sys)
    IS.@assert_op get_ω_sys(inputs) != -1

    _mass_matrix_inputs!(inputs)

    IS.@assert_op injection_n_states == state_space_ix[1] - branches_n_states - n_buses * 2
    IS.@assert_op n_buses * 2 - static_bus_var_count >= 0
    IS.@assert_op length(inputs.DAE_vector) == state_space_ix[1]

    @debug inputs.counts
    return
end

"""
SimulationInputs build function for MassMatrixModels
"""
function build!(inputs::SimulationInputs, ::Type{MassMatrixModel}, sys::PSY.System)
    return _build!(inputs, sys)
end

"""
SimulationInputs build function for ImplicitModels
"""
function build!(inputs::SimulationInputs, ::Type{ImplicitModel}, sys::PSY.System)
    return _build!(inputs, sys)
end


using YOLO, BenchmarkTools

#First time only (downloads 5011 images & labels!)
#YOLO.download_dataset("voc2007")

settings = YOLO.pretrained.v2_tiny_voc.load(minibatch_size=1) #Process one image at a time
model = YOLO.v2_tiny.load(settings)
YOLO.loadWeights!(model, settings)

voc = YOLO.datasets.VOC.populate()
vocloaded = YOLO.load(voc, settings, indexes = [rand(1:5011)]) #Load a single random image

bt = @benchmark begin
  res = model(vocloaded.imstack_mat);
  predictions = YOLO.postprocess(res, settings, conf_thresh = 0.3, iou_thresh = 0.3)
end

display(bt)


__precompile__()

module Nabla

    using SpecialFunctions
    using LinearAlgebra
    using Statistics

    # Some aliases used repeatedly throughout the package.
    export ∇Scalar, ∇Array, SymOrExpr, ∇ArrayOrScalar
    const ∇Scalar = Number
    const ∇Array = AbstractArray{<:∇Scalar}
    const ∇AbstractVector = AbstractVector{<:∇Scalar}
    const ∇AbstractMatrix = AbstractMatrix{<:∇Scalar}
    const ∇ArrayOrScalar = Union{AbstractArray{<:∇Scalar}, ∇Scalar}
    const SymOrExpr = Union{Symbol, Expr}

    # ones/zeros(::AbstractArray) is deprecated in 0.7 and removed in 1.0, but it's a
    # pretty useful method, so we'll define our own for internal use
    for f in (:ones, :zeros)
        like = Symbol(f, "like")
        @eval begin
            $(like)(a::AbstractArray) = $(f)(eltype(a), size(a))
            $(like)(n::Integer) = $(f)(n)
        end
    end

    # Meta-programming utilities specific to Nabla.
    include("code_transformation/util.jl")
    include("code_transformation/differentiable.jl")

    # Functionality for constructing computational graphs.
    include("core.jl")

    # Functionality for defining new sensitivities.
    include("sensitivity.jl")

    # Finite differencing functionality - only used in tests. Would be good to move this
    # into a separate module at some point.
    include("finite_differencing.jl")

    # Sensitivities for the basics.
    include("sensitivities/indexing.jl")
    include("sensitivities/scalar.jl")
    include("sensitivities/array.jl")

    # Sensitivities for functionals.
    include("sensitivities/functional/functional.jl")
    include("sensitivities/functional/reduce.jl")
    include("sensitivities/functional/reducedim.jl")

    # Linear algebra optimisations.
    include("sensitivities/linalg/generic.jl")
    include("sensitivities/linalg/symmetric.jl")
    include("sensitivities/linalg/strided.jl")
    include("sensitivities/linalg/blas.jl")
    include("sensitivities/linalg/diagonal.jl")
    include("sensitivities/linalg/triangular.jl")
    include("sensitivities/linalg/factorization/cholesky.jl")
    include("sensitivities/linalg/factorization/svd.jl")

    # Checkpointing
    include("checkpointing.jl")

end # module Nabla


# parameters

# wrapper of C types defined in include/bayesopt/parameters.h

@enum learning_type L_FIXED L_EMPIRICAL L_DISCRETE L_MCMC L_ERROR=-1
@enum score_type SC_MTL SC_ML SC_MAP SC_LOOCV SC_ERROR=-1

immutable kernel_parameters
    name::Cstring
    hp_mean::NTuple{128,Cdouble}
    hp_std::NTuple{128,Cdouble}
    n_hp::Csize_t
end

immutable mean_parameters
    name::Cstring
    coef_mean::NTuple{128,Cdouble}
    coef_std::NTuple{128,Cdouble}
    n_coef::Csize_t
end

type bopt_params
    n_iterations::Csize_t
    n_inner_iterations::Csize_t
    n_init_samples::Csize_t
    n_iter_relearn::Csize_t
    init_method::Csize_t
    random_seed::Cint
    verbose_level::Cint
    log_filename::Cstring
    load_save_flag::Csize_t
    load_filename::Cstring
    save_filename::Cstring
    surr_name::Cstring
    sigma_s::Cdouble
    noise::Cdouble
    alpha::Cdouble
    beta::Cdouble
    sc_type::score_type
    l_type::learning_type
    l_all::Cint
    epsilon::Cdouble
    force_jump::Csize_t
    kernel::kernel_parameters
    mean::mean_parameters
    crit_name::Cstring
    crit_params::NTuple{128,Cdouble}
    n_crit_params::Csize_t
end

function Base.show(io::IO, params::bopt_params)
    print(io, "BayesOpt parameters:")
    width = maximum([length(string(f)) for f in fieldnames(bopt_params)])
    for f in fieldnames(bopt_params)
        val = params.(f)
        if f == :sc_type
            val = score2str(val)
        elseif f == :l_type
            val = learn2str(val)
        elseif f in (:kernel, :mean)
            val = val.name
        end
        if isa(val, Number)
            val = string(val)
        elseif isa(val, Cstring)
            val = bytestring(val)
        else
            # skip
            continue
        end
        println(io)
        print(io, "  ", lpad(f, width), " = ", val)
    end
end

function learn2str(name::learning_type)
    return ccall((:learn2str, libbayesopt), Cstring, (learning_type,), name)
end

function score2str(name::score_type)
    return ccall((:score2str, libbayesopt), Cstring, (score_type,), name)
end

# setters
for f in [
      :set_kernel,
      :set_mean,
      :set_criteria,
      :set_surrogate,
      :set_log_file,
      :set_load_file,
      :set_save_file,
      :set_learning,
      :set_score]
    @eval function $f(params::bopt_params, name::String)
        ccall(
            ($(string(f)), libbayesopt),
            Void,
            (Ref{bopt_params}, Cstring),
            params, name
        )
    end
end

function initialize_parameters_to_default()
    ccall(
        (:initialize_parameters_to_default, libbayesopt),
        bopt_params, ())
end


function EllipsoidTessellator(arg0::OneAxisEllipsoid, arg1::TileAiming, arg2::jint)
    return EllipsoidTessellator((OneAxisEllipsoid, TileAiming, jint), arg0, arg1, arg2)
end

function build_simple_zone(::Type{EllipsoidTessellator}, arg0::jdouble, arg1::Vector{GeodeticPoint})
    return jcall(EllipsoidTessellator, "buildSimpleZone", SphericalPolygonsSet, (jdouble, Vector{GeodeticPoint}), arg0, arg1)
end

function build_simple_zone(::Type{EllipsoidTessellator}, arg0::jdouble, arg1::Vector{Vector{jdouble}})
    return jcall(EllipsoidTessellator, "buildSimpleZone", SphericalPolygonsSet, (jdouble, Vector{Vector{jdouble}}), arg0, arg1)
end

function sample(obj::EllipsoidTessellator, arg0::SphericalPolygonsSet, arg1::jdouble, arg2::jdouble)
    return jcall(obj, "sample", List, (SphericalPolygonsSet, jdouble, jdouble), arg0, arg1, arg2)
end

function tessellate(obj::EllipsoidTessellator, arg0::SphericalPolygonsSet, arg1::jdouble, arg2::jdouble, arg3::jdouble, arg4::jdouble, arg5::jboolean, arg6::jboolean)
    return jcall(obj, "tessellate", List, (SphericalPolygonsSet, jdouble, jdouble, jdouble, jdouble, jboolean, jboolean), arg0, arg1, arg2, arg3, arg4, arg5, arg6)
end


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

export get_orders, get_nl

import ACE: standardevaluator, graphevaluator
using SparseArrays: SparseMatrixCSC, sparse
using LinearAlgebra: mul!

"""
`struct RPIBasis`
"""
struct RPIBasis{T, BOP, NZ, TIN} <: IPBasis
   pibasis::PIBasis{BOP, NZ, TIN}
   A2Bmaps::NTuple{NZ, SparseMatrixCSC{T, Int}}
   Bz0inds::NTuple{NZ, UnitRange{Int}}
end

Base.length(basis::RPIBasis, iz0::Integer) = size(basis.A2Bmaps[iz0], 1)

Base.length(basis::RPIBasis) = sum(length(basis, iz0)
                                    for iz0 = 1:numz(basis.pibasis))

fltype(::RPIBasis{T}) where {T}  = T

zlist(basis::RPIBasis) = zlist(basis.pibasis)

cutoff(basis::RPIBasis) = cutoff(basis.pibasis)

standardevaluator(basis::RPIBasis) =
      RPIBasis( standardevaluator(basis.pibasis),
                basis.A2Bmaps, basis.Bz0inds )

graphevaluator(basis::RPIBasis) =
      RPIBasis( graphevaluator(basis.pibasis),
                basis.A2Bmaps, basis.Bz0inds )

_basisfcnidx(basis::RPIBasis, iz0::Integer, ib::Integer) = basis.Bz0inds[iz0][ib]

# ------------------------------------------------------------------------
#    Utilities / Basis analysis code
# ------------------------------------------------------------------------

function get_orders(basis::RPIBasis)
   # get the correlation orders of the PIbasis ...
   Ns_pi = zeros(Int, length(basis.pibasis))
   for iz0 = 1:numz(basis)
      inner = basis.pibasis.inner[iz0]
      Ns_pi[inner.AAindices] .= inner.orders
   end
   # ... and use them to construct the correlation orders of the RPI basis
   Ns = zeros(Int, length(basis))
   for iz0 = 1:numz(basis)
      # loop over basis functions belonging to centre-species iz0
      for ib = 1:length(basis, iz0)
         # find one of the indices of the PI basis that the current
         # RPI basis function belongs to; they are all equivalent in
         # terms of correlation-order, so only one of them matters.
         iPI = findfirst(basis.A2Bmaps[iz0][ib,:] .!= 0)
         Ns[_basisfcnidx(basis, iz0, ib)] = Ns_pi[iPI]
      end
   end
   return Ns
end

function get_nl(basis::RPIBasis)
   # get the correlation orders of the PIbasis ...
   nnll_pi = Vector{Vector{Any}}(undef, length(basis.pibasis))
   for iz0 = 1:numz(basis)
      inner = basis.pibasis.inner[iz0]
      for (key, val) in inner.b2iAA
         b = [ (l = b.l, n = b.n) for b in key.oneps ]
         nnll_pi[inner.AAindices[val]] = b
      end
   end
   # ... and use them to construct the correlation orders of the RPI basis
   nnll = Vector{Vector{Any}}(undef, length(basis))
   for iz0 = 1:numz(basis)
      # loop over basis functions belonging to centre-species iz0
      for ib = 1:length(basis, iz0)
         # find one of the indices of the PI basis that the current
         # RPI basis function belongs to; they are all equivalent in
         # terms of correlation-order, so only one of them matters.
         iPI = findfirst(basis.A2Bmaps[iz0][ib,:] .!= 0)
         nnll[_basisfcnidx(basis, iz0, ib)] = nnll_pi[iPI]
      end
   end
   return nnll
end

# ------------------------------------------------------------------------
#    FIO code
# ------------------------------------------------------------------------

==(B1::RPIBasis, B2::RPIBasis) = (B1.pibasis == B2.pibasis)

write_dict(basis::RPIBasis) = Dict(
      "__id__" => "ACE_RPIBasis",
      "__v__"  => "v0_8_2",
      "pibasis" => write_dict(basis.pibasis),
      "A2Bmaps" => write_dict.(basis.A2Bmaps),
      "Bz0inds" => [ [ur.start, ur.stop] for ur in basis.Bz0inds ]
   )

# v0.8.2 onwards
function read_dict(::Val{:ACE_RPIBasis}, ::Val{:v0_8_2}, D::Dict)
   pibasis = read_dict(D["pibasis"])
   A2Bmaps = tuple( read_dict.(D["A2Bmaps"])... )
   Bz0inds = tuple( [ ur[1]:ur[2] for ur in D["Bz0inds"] ]... )
   return RPIBasis(pibasis, A2Bmaps, Bz0inds)
end

# old version
read_dict(::Val{:ACE_RPIBasis}, D::Dict) =
   RPIBasis(read_dict(D["pibasis"]))

read_dict(::Val{:SHIPs_RPIBasis}, D::Dict) =
   read_dict(Val{:ACE_RPIBasis}(), D)

# ------------------------------------------------------------------------
#    Basis construction code
# ------------------------------------------------------------------------

RPIBasis(basis1p::OneParticleBasis, N::Integer,
         D::AbstractDegree, maxdeg::Real, constants=false) =
   RPIBasis(PIBasis(basis1p, N, D, maxdeg; filter = RPIFilter(constants)))

function RPIBasis(pibasis::PIBasis)
   basis1p = pibasis.basis1p

   # construct the cg matrices
   rotc = Rot3DCoeffs()
   A2Bmaps = ntuple(iz0 -> _rpi_A2B_matrix(rotc, pibasis, iz0), numz(pibasis))

   # construct the indices within the B vector to which the A2Bmaps map.
   Bz0inds = UnitRange{Int}[]
   idx0 = 0
   for i = 1:length(A2Bmaps)
      len = size(A2Bmaps[i], 1)
      push!(Bz0inds, (idx0+1):(idx0+len))
      idx0 += len
   end

   return RPIBasis(pibasis, A2Bmaps, tuple(Bz0inds...))
end

# TODO NOW: graphevaluator, standardevaluator

struct RPIFilter
   constants::Bool
end

(f::RPIFilter)(pib::PIBasisFcn{0}) = f.constants
(f::RPIFilter)(pib::PIBasisFcn{1}) = (pib.oneps[1].l == 0)
(f::RPIFilter)(pib::PIBasisFcn) = (
      iseven( sum(b.l for b in pib.oneps) ) &&
      (sum(b.m for b in pib.oneps) == 0) )

# _rpi_filter(pib::PIBasisFcn{0}) = false
# _rpi_filter(pib::PIBasisFcn{1}) = (pib.oneps[1].l == 0)
# _rpi_filter(pib::PIBasisFcn) = (
#       iseven( sum(b.l for b in pib.oneps) ) &&
#       (sum(b.m for b in pib.oneps) == 0) )

function _rpi_A2B_matrix(rotc::Rot3DCoeffs,
                         pibasis::PIBasis,
                         iz0)
   # allocate triplet format
   Irow, Jcol, vals = Int[], Int[], real(fltype(pibasis.basis1p))[]
   # count the number of PI basis functions = number of rows
   idxB = 0
   # loop through all (zz, kk, ll) tuples; each specifies 1 to several B
   for i = 1:length(pibasis.inner[iz0])
      # get the specification of the ith basis function
      pib = get_basis_spec(pibasis, iz0, i)
      # skip it unless all m are zero, because we want to consider each
      # (nn, ll) block only once.
      if !all(b.m == 0 for b in pib.oneps)
         continue
      end
      # get the rotation-coefficients for this basis group
      # the bs are the basis functions corresponding to the columns
      U, bcols = _rpi_coupling_coeffs(pibasis, rotc, pib)
      # loop over the rows of Ull -> each specifies a basis function
      for irow = 1:size(U, 1)
         idxB += 1
         # loop over the columns of U / over brows
         for (icol, bcol) in enumerate(bcols)
            # this is a subtle step: bcol and bcol_ordered are equivalent
            # permutation-invariant basis functions. This means we will
            # add the same PI basis function several times, but in the call to
            # `sparse` the values will just be added.
            bcol_ordered = ACE._get_ordered(pibasis, bcol)
            idxAA = pibasis.inner[iz0].b2iAA[bcol_ordered]
            push!(Irow, idxB)
            push!(Jcol, idxAA)
            push!(vals, U[irow, icol])
         end
      end
   end
   # create CSC: [   triplet    ]  nrows   ncols
   return sparse(Irow, Jcol, vals, idxB, length(pibasis.inner[iz0]))
end

# U, bcols = rpi_coupling_coeffs(rotc, pib)

"""
this is essentially a wrapper function around Rotations3D.rpi_basis,
and is just meant to translate between different representations
"""
function _rpi_coupling_coeffs(pibasis, rotc::Rot3DCoeffs, pib::PIBasisFcn{N}
                              ) where {N}
   # convert to zz, ll, nn tuples
   zz, nn, ll, _ = _b2znlms(pib)
   # construct the RPI coupling coefficients
   U, Ms = Rotations3D.rpi_basis(rotc, zz, nn, ll)
   # convert the Ms into basis functions
   rpibs = [ _znlms2b(zz, nn, ll, mm, pib.z0) for mm in Ms ]
   return U, rpibs
end

_rpi_coupling_coeffs(pibasis, rotc::Rot3DCoeffs, pib::PIBasisFcn{0}) =
      [ 1.0 ], []

_b2znlms(pib::PIBasisFcn{N}) where {N} = (
   SVector(ntuple(n -> pib.oneps[n].z, N)...),
   SVector(ntuple(n -> pib.oneps[n].n, N)...),
   SVector(ntuple(n -> pib.oneps[n].l, N)...),
   SVector(ntuple(n -> pib.oneps[n].m, N)...) )

_znlms2b(zz, nn, ll, mm = zero(ll), z0 = AtomicNumber(0)) =
   PIBasisFcn( z0, ntuple(i -> PSH1pBasisFcn(nn[i], ll[i], mm[i], zz[i]),
                          length(zz)) )

function combine(basis::RPIBasis, coeffs)
   picoeffs = ntuple(iz0 -> (coeffs[basis.Bz0inds[iz0]]' * basis.A2Bmaps[iz0])[:],
                     numz(basis.pibasis))
   return PIPotential(basis.pibasis, picoeffs)
end

"""
`scaling(basis::RPIBasis, p[; a2b = abs2, fin = ...])` : same as usual but adds a keyword
argument `use` :
* `a2b = identity` : the old buggy scaling
* `a2b = abs` : the bugfix which keeps the behaviour same/similar
* `a2b = abs2` : a different kind of scaling which I believe is the right one
but behaves a bit differently especially for high body-orders.

The `fin` keyword argument is the "finishing" map. E.g. if we take `a2b = abs2`,
then we should apply a sqrt at the end. This is done automatically here.
"""
function scaling(basis::RPIBasis, p; a2b = abs2,
                 fin = (a2b == abs2 ? sqrt : identity) )
   wwpi = scaling(basis.pibasis, p)
   wwrpi = zeros(Float64, length(basis))
   for iz0 = 1:numz(basis)
      wwpi_iz0 = wwpi[basis.pibasis.inner[iz0].AAindices]
      wwrpi[basis.Bz0inds[iz0]] = a2b.(basis.A2Bmaps[iz0]) * a2b.(wwpi_iz0)
   end
   return fin.(wwrpi)
end

# ------------------------------------------------------------------------
#    Evaluation code
# ------------------------------------------------------------------------

alloc_temp(basis::RPIBasis, args...) =
   ( AA = site_alloc_B(basis.pibasis, args...),
     AAr = real(site_alloc_B(basis.pibasis, args...)),
     tmp_pibasis = alloc_temp(basis.pibasis, args...)
   )

function evaluate!(B, tmp, basis::RPIBasis, Rs, Zs, z0)
   iz0 = z2i(basis, z0)
   AA = site_evaluate!(tmp.AA, tmp.tmp_pibasis, basis.pibasis, Rs, Zs, z0)
   # TODO: this could be done better maybe by adding the real function into
   #       site_evaluate!, or by writing a real version of it...
   AAr = @view tmp.AAr[1:length(AA)]
   AAr[:] .= real.(AA)
   Bview = @view B[basis.Bz0inds[iz0]]
   mul!(Bview, basis.A2Bmaps[iz0], AAr)
   return B
end

# ------- gradient

alloc_temp_d(basis::RPIBasis, Rs::AbstractVector, args...) =
   alloc_temp_d(basis, length(Rs))

alloc_temp_d(basis::RPIBasis, nmax::Integer) =
    (
    AA = site_alloc_B(basis.pibasis),
    AAr = real(site_alloc_B(basis.pibasis)),
    tmp_pibasis = alloc_temp(basis.pibasis, nmax),
    dAA = site_alloc_dB(basis.pibasis, nmax),
    tmpd_pibasis = alloc_temp_d(basis.pibasis, nmax),
    )

# TODO: evaluate also B??? the interface seems to command it.
function evaluate_d!(B, dB, tmpd, basis::RPIBasis, Rs, Zs, z0)
   iz0 = z2i(basis, z0)
   AA, dAA = tmpd.AA, tmpd.dAA
   site_evaluate_d!(AA, dAA, tmpd.tmpd_pibasis, basis.pibasis, Rs, Zs, z0)
   len = length(basis.pibasis.inner[iz0])
   for i = 1:len
      AA[i] = real(AA[i])
   end
   mul!((@view B[basis.Bz0inds[iz0]]), basis.A2Bmaps[iz0], (@view AA[1:len]))
   for j = 1:length(Rs)
      # ∂∏A / ∂𝐫ⱼ
      dAAj = @view dAA[1:len, j]
      for i = 1:length(dAAj)
         # TODO: we should incorporate the `real` operation into the PIBasis?
         @inbounds dAAj[i] = real.(dAAj[i])
      end
      # copy into B
      dBview = @view dB[basis.Bz0inds[iz0], j]
      mul!(dBview, basis.A2Bmaps[iz0], dAAj)
   end
   return dB
end


@testset "SymmetricAllDifferent2AllDifferentInverse: $(fct_type), dimension $(dim), $(T)" for fct_type in ["vector of variables", "vector affine function"], dim in [2, 3], T in [Int, Float64]
    mock = MOIU.MockOptimizer(AllDifferentInverseModel{T}())
    model = COIB.SymmetricAllDifferent2AllDifferentInverse{T}(mock)

    if T == Int
        @test MOI.supports_constraint(model, MOI.VariableIndex, MOI.Integer)
    end
    @test MOI.supports_constraint(
        model,
        MOI.VectorAffineFunction{T},
        CP.AllDifferent,
    )
    @test MOI.supports_constraint(
        model,
        MOI.VectorAffineFunction{T},
        CP.Inverse,
    )
    @test MOIB.supports_bridging_constraint(
        model,
        MOI.VectorAffineFunction{T},
        CP.SymmetricAllDifferent,
    )

    if T == Int
        x, _ = MOI.add_constrained_variables(model, [MOI.Integer() for _ in 1:dim])
    elseif T == Float64
        x = MOI.add_variables(model, dim)
    end

    fct = if fct_type == "vector of variables"
        MOI.VectorOfVariables(x)
    elseif fct_type == "vector affine function"
        MOIU.vectorize(x)
    else
        @assert false
    end
    c = MOI.add_constraint(model, fct, CP.SymmetricAllDifferent(dim))

    for i in 1:dim
        @test MOI.is_valid(model, x[i])
    end
    @test MOI.is_valid(model, c)

    bridge = first(MOIBC.bridges(model))[2]

    @testset "Bridge properties" begin
        @test MOIBC.concrete_bridge_type(typeof(bridge), MOI.VariableIndex, CP.SymmetricAllDifferent) == typeof(bridge)
        @test MOIB.added_constrained_variable_types(typeof(bridge)) == Tuple{Type}[]
        @test MOIB.added_constraint_types(typeof(bridge)) == [
            (MOI.VectorAffineFunction{T}, CP.AllDifferent),
            (MOI.VectorAffineFunction{T}, CP.Inverse),
        ]

        @test MOI.get(bridge, MOI.NumberOfVariables()) == 0
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.VectorAffineFunction{T}, CP.AllDifferent}()) == 1
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.VectorAffineFunction{T}, CP.Inverse}()) == 1

        @test MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T}, CP.AllDifferent}()) == [bridge.con_all_diff]
        @test MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T}, CP.Inverse}()) == [bridge.con_inverse]
    end

    @testset "All different" begin
        @test MOI.is_valid(model, bridge.con_all_diff)
        f = MOI.get(model, MOI.ConstraintFunction(), bridge.con_all_diff)
        @test length(f.terms) == dim
        @test MOI.get(model, MOI.ConstraintSet(), bridge.con_all_diff) == CP.AllDifferent(dim)

        for i in 1:dim
            t = f.terms[i]
            @test t.output_index == i
            @test t.scalar_term.coefficient === one(T)
            @test t.scalar_term.variable == x[i]
        end
    end

    @testset "Inverse" begin
        @test MOI.is_valid(model, bridge.con_inverse)
        f = MOI.get(model, MOI.ConstraintFunction(), bridge.con_inverse)
        @test length(f.terms) == 2 * dim
        @test MOI.get(model, MOI.ConstraintSet(), bridge.con_inverse) == CP.Inverse(dim)

        for i in 1:dim
            t = f.terms[i]
            @test t.output_index == i
            @test t.scalar_term.coefficient === one(T)
            @test t.scalar_term.variable == x[i]

            t = f.terms[dim + i]
            @test t.output_index == dim + i
            @test t.scalar_term.coefficient === one(T)
            @test t.scalar_term.variable == x[i]
        end
    end
end