"""
`Transition{T<:AbstractFloat}`

The transition equation of the state-space model takes the form

    `s_t = TTT*s_{t-1} + RRR*ϵ_t + CCC`

The `Transition` type stores the coefficient `Matrix{T}`s (`TTT`, `RRR`) and constant `Vector{T} CCC`.
"""
mutable struct Transition{T<:AbstractFloat}
    TTT::Matrix{T}
    RRR::Matrix{T}
    CCC::Vector{T}
end
function Transition(TTT::Matrix{T}, RRR::Matrix{T}) where T<:AbstractFloat
    CCC = zeros(eltype(TTT), size(TTT, 1))
    Transition{T}(TTT, RRR, CCC)
end
function Transition(TTT::Matrix{T}, RRR::Matrix{T}, CCC::Matrix{T}) where T<:AbstractFloat
    Transition{T}(TTT, RRR, collect(CCC))
end
function Base.getindex(eq::Transition, d::Symbol)
    if d in (:TTT, :RRR, :CCC)
        return getfield(eq, d)
    else
        throw(KeyError(d))
    end
end

"""
`Measurement{T<:AbstractFloat}`

The measurement equation of the state-space model takes the form

   `y_t = ZZ*s_t + DD + u_t`

where the error `u_t` is the measurement error, which is uncorrelated with the
shocks in the transition equation `ϵ_t`.

### Fields

If `Ns` is the number of states `s_t`, `Ny` is the number of
observables `y_t`, and `Ne` is the number of shocks `ϵ_t`:

- `ZZ`: the `Ny` x `Ns`  measurement matrix
- `DD`: the `Ny` x 1 constant vector
- `QQ`: the `Ne` x `Ne` covariance matrix for the shocks `ϵ_t`
- `EE`: the `Ny` x `Ny` covariance matrix for the measurement error `η_t`
"""
mutable struct Measurement{T<:AbstractFloat}
    ZZ::Matrix{T}
    DD::Vector{T}
    QQ::Matrix{T}
    EE::Matrix{T}
end

function Base.getindex(M::Measurement, d::Symbol)
    if d in (:ZZ, :DD, :QQ, :EE)
        return getfield(M, d)
    else
        throw(KeyError(d))
    end
end

function measurement(m::AbstractDSGEModel, trans::Transition; shocks::Bool=true)
    TTT = trans[:TTT]
    RRR = trans[:RRR]
    CCC = trans[:CCC]
    measurement(m, TTT, RRR, CCC; shocks=shocks)
end

"""
```
PseudoMeasurement{T<:AbstractFloat}
```

The pseudo-measurement equation of the state-space model takes the form

   `x_t = ZZ_pseudo*s_t + DD_pseudo`

### Fields

Let `Ns` be the number of states `s_t` and `Nx` be the number of
pseudo-observables `x_t`:

- `ZZ_pseudo`: the `Nx` x `Ns` pseudo-measurement matrix
- `DD_pseudo`: the `Nx` x 1 constant vector
"""
mutable struct PseudoMeasurement{T<:AbstractFloat}
    ZZ_pseudo::Matrix{T}
    DD_pseudo::Vector{T}
end

function Base.getindex(M::PseudoMeasurement, d::Symbol)
    if d in (:ZZ_pseudo, :DD_pseudo)
        return getfield(M, d)
    else
        throw(KeyError(d))
    end
end

"""
`System{T<:AbstractFloat}`

A mutable struct containing the transition and measurement equations for a
state-space model. The matrices may be directly indexed: `sys[:TTT]`
returns `sys.transition.TTT`, `sys[:ZZ]` returns `sys.measurement.ZZ`, etc.
"""
mutable struct System{T<:AbstractFloat}
    transition::Transition{T}
    measurement::Measurement{T}
    pseudo_measurement::PseudoMeasurement{T}
end

function System(transition::Transition{T}, measurement::Measurement{T}) where T<:AbstractFloat
    # Initialize empty pseudo-measurement equation
    _n_states = size(transition.TTT, 1)
    _n_pseudo = 0
    ZZ_pseudo = zeros(_n_pseudo, _n_states)
    DD_pseudo = zeros(_n_pseudo)
    pseudo_measurement = PseudoMeasurement(ZZ_pseudo, DD_pseudo)

    return System(transition, measurement, pseudo_measurement)
end

function Base.getindex(system::System, d::Symbol)
    if d in (:transition, :measurement, :pseudo_measurement)
        return getfield(system, d)
    elseif d in fieldnames(typeof(system.transition))
        return getfield(system.transition, d)
    elseif d in fieldnames(typeof(system.measurement))
        return getfield(system.measurement, d)
    elseif d in fieldnames(typeof(system.pseudo_measurement))
        return getfield(system.pseudo_measurement, d)
    else
        throw(KeyError(d))
    end
end

function Base.copy(system::System)
    trans = Transition(system[:TTT], system[:RRR], system[:CCC])
    meas  = Measurement(system[:ZZ], system[:DD], system[:QQ], system[:EE])
    pseudo_meas = PseudoMeasurement(system[:ZZ_pseudo], system[:DD_pseudo])
    return System(trans, meas, pseudo_meas)
end

"""
```
compute_system(m; apply_altpolicy = false)
```

Given the current model parameters, compute the state-space system
corresponding to model `m`. Returns a `System` object.
"""
function compute_system(m::AbstractDSGEModel{T}; apply_altpolicy = false,
                        verbose::Symbol = :high) where T<:AbstractFloat

    solution_method = get_setting(m, :solution_method)

    # Solve model
    if solution_method == :gensys

        TTT, RRR, CCC = solve(m; apply_altpolicy = apply_altpolicy, verbose = verbose)
        transition_equation = Transition(TTT, RRR, CCC)

        # Solve measurement equation
        measurement_equation = measurement(m, TTT, RRR, CCC)

    elseif solution_method == :klein
        # Unpacking the method from solve to hang on to TTT_jump
        if m.spec == "het_dsge"
            TTT_jump, TTT_state, eu = klein(m)
        else
            TTT_jump, TTT_state, eu = klein(m)
        end
        if eu==-1
            throw(KleinError())
        end

        TTT, RRR = klein_transition_matrices(m, TTT_state, TTT_jump)
        CCC = zeros(n_model_states(m))

        if m.spec == "real_bond_mkup"
            GDPeqn = construct_GDPeqn(m, TTT_jump)
            TTT, RRR, CCC = augment_states(m, TTT, TTT_jump, RRR, CCC, GDPeqn)
            # Measurement (needs the additional TTT_jump argument)
            measurement_equation = measurement(m, TTT, TTT_jump, RRR, CCC, GDPeqn)
        elseif m.spec == "het_dsge" || m.spec == "rep_dsge"
            TTT, RRR, CCC = augment_states(m, TTT, RRR, CCC)
            measurement_equation = measurement(m, TTT, RRR, CCC)
        else
            TTT, RRR, CCC        = augment_states(m, TTT, RRR, CCC)
            measurement_equation = measurement(m, TTT, RRR, CCC)
        end

        transition_equation = Transition(TTT, RRR, CCC)

    else
        throw("solution_method provided does not exist.")
    end

    type_tuple = (typeof(m), Matrix{T}, Matrix{T}, Vector{T})
    if hasmethod(pseudo_measurement, type_tuple)
        # Solve pseudo-measurement equation
        pseudo_measurement_equation = pseudo_measurement(m, TTT, RRR, CCC)
        return System(transition_equation, measurement_equation, pseudo_measurement_equation)
    else
        return System(transition_equation, measurement_equation)
    end
end

"""
```
compute_system(m; apply_altpolicy = false,
               regime_switching = false, n_regimes = 2,
               check_system = false, get_system = false,
               get_population_moments = false, use_intercept = false,
               verbose = :high)
compute_system(m, data; apply_altpolicy = false,
               regime_switching = false, n_regimes = 2,
               check_system = false, get_system = false,
               get_population_moments = false,
               verbose = :high)
```
Given the current model parameters, compute the DSGE-VAR or DSGE-VECM system
corresponding to model `m`. If a matrix `data` is also passed, then
the VAR is estimated on `data` using the DSGE `m` as a prior
with weight λ.

### Keyword Arguments
* `check_system::Bool`: see `?compute_system` that takes the input `m::AbstractDSGEModel`
    and `system::System`.
* `get_system::Bool`: see Outputs
* `get_population_moments::Bool`: see Outputs
* `use_intercept::Bool`: use an intercept term when computing the OLS estimate of the VAR system.

### Outputs
* If `get_system = true`:
    Returns the updated `system` whose measurement matrices `ZZ`, `DD`, and `QQ` correspond
    to the VAR or VECM specified by `m`. If `m` is an `AbstractDSGEVECMModel`,
    then the `system` and the vector implied by additional cointegrating relationships
    are returned as a 2-element tuple.
* If `get_population_moments = true`:
    Returns the limit cross product matrices that describe the DSGE implied
    population moments between the observables and their lags. If `data` is
    also passed as an input, then the sample population moments are also returned.
* Otherwise:
    Returns `β` and `Σ`, the coefficients and observables covariance matrix of the VAR or VECM.
    If `data` is passed in, then `β` and `Σ` are estimated from the data using `m`
    as a prior with weight λ. Otherwise, `β` and `Σ` comprise the VECM approximation
    of the DSGE `m`.
"""
function compute_system(m::AbstractDSGEVARModel{T}; apply_altpolicy::Bool = false,
                        regime_switching::Bool = false, regime::Int = 1, n_regimes::Int = 2,
                        check_system::Bool = false, get_system::Bool = false,
                        get_population_moments::Bool = false, use_intercept::Bool = false,
                        verbose::Symbol = :high) where {T<:Real}
    dsge = get_dsge(m)
    if regime_switching
        regime_system = compute_system(dsge; apply_altpolicy = apply_altpolicy,
                                       regime_switching = regime_switching, n_regimes = n_regimes,
                                       verbose = verbose)
        system = System(regime_system, regime)
    else
        system = compute_system(dsge; verbose = verbose)

    end
    # Update the system to reflect the desired observables using information
    # from the underlying DSGE object
    system = compute_system(dsge, system; observables = collect(keys(get_observables(m))),
                            shocks = collect(keys(get_shocks(m))), check_system = check_system)

    if get_system
        return system
    else
        EE, MM = measurement_error(m)

        return var_approx_state_space(system[:TTT], system[:RRR], system[:QQ],
                                      system[:DD], system[:ZZ], EE, MM, n_lags(m);
                                      get_population_moments = get_population_moments,
                                      use_intercept = use_intercept)
    end
end

function compute_system(m::AbstractDSGEVARModel{T}, data::Matrix{T};
                        apply_altpolicy::Bool = false,
                        regime_switching::Bool = false, regime::Int = 1, n_regimes::Int = 2,
                        check_system::Bool = false, get_system::Bool = false,
                        get_population_moments::Bool = false,
                        verbose::Symbol = :high) where {T<:Real}

    if get_λ(m) == Inf
        # Then we just want the VAR approximation of the DSGE
        return compute_system(m; apply_altpolicy = apply_altpolicy,
                              regime_switching = regime_switching, regime = regime, n_regimes = n_regimes,
                              check_system = check_system, get_system = get_system,
                              get_population_moments = get_population_moments, use_intercept = true,
                              verbose = verbose)
    else
        # Create a system using the method for DSGE-VARs and λ = ∞
        system = compute_system(m; apply_altpolicy = apply_altpolicy,
                                regime_switching = regime_switching, regime = regime,
                                n_regimes = n_regimes, check_system = check_system,
                                get_system = true, use_intercept = true,
                                verbose = verbose)

        if get_system
            return system
        else
            EE, MM = measurement_error(m)

            lags = n_lags(m)
            YYYY, XXYY, XXXX =
                compute_var_population_moments(data, lags; use_intercept = true)
            out = var_approx_state_space(system[:TTT], system[:RRR], system[:QQ],
                                         system[:DD], system[:ZZ], EE, MM, n_lags(m);
                                         get_population_moments = true,
                                         use_intercept = true)

            if get_population_moments
                return out..., YYYY, XXYY, XXXX
            else
                # Compute prior-weighted population moments
                λ = get_λ(m)
                YYYYC = YYYY + λ .* out[1]
                XXYYC = XXYY + λ .* out[2]
                XXXXC = XXXX + λ .* out[3]

                # Draw stationary VAR system
                n_periods = size(data, 2) - lags
                β, Σ =  draw_stationary_VAR(YYYYC, XXYYC, XXXXC,
                                            convert(Int, n_periods + λ * n_periods),
                                            size(data, 1), lags)

                return β, Σ
            end
        end
    end
end

# Same functions as above but for AbstractDSGEVECMModel types
function compute_system(m::AbstractDSGEVECMModel{T}; apply_altpolicy::Bool = false,
                        regime_switching::Bool = false, regime::Int = 1, n_regimes::Int = 2,
                        check_system::Bool = false, get_system::Bool = false,
                        get_population_moments::Bool = false, use_intercept::Bool = false,
                        verbose::Symbol = :high) where {T<:Real}
    dsge = get_dsge(m)
    if regime_switching
        regime_system = compute_system(dsge; apply_altpolicy = apply_altpolicy,
                                       regime_switching = regime_switching, n_regimes = n_regimes,
                                       verbose = verbose)
        system = System(regime_system, regime)
    else
        system = compute_system(dsge; verbose = verbose)

    end
    # Use wrapper compute_system for AbstractDSGEVECMModel types
    # (as opposed to compute_system(m::AbstractDSGEModel, system::System; ...))
    system, DD_coint_add = compute_system(m, system; observables = collect(keys(get_observables(m))),
                                          cointegrating = collect(keys(get_cointegrating(m))),
                                          cointegrating_add = collect(keys(get_cointegrating_add(m))),
                                          shocks = collect(keys(get_shocks(m))), check_system = check_system,
                                          get_DD_coint_add = true)

    if get_system
        return system, DD_coint_add
    else
        EE, MM = measurement_error(m)

        return vecm_approx_state_space(system[:TTT], system[:RRR], system[:QQ],
                                       system[:DD], system[:ZZ], EE, MM, n_observables(m),
                                       n_lags(m), n_cointegrating(m), n_cointegrating_add(m),
                                       DD_coint_add;
                                       get_population_moments = get_population_moments,
                                       use_intercept = use_intercept)
    end
end

function compute_system(m::AbstractDSGEVECMModel{T}, data::Matrix{T};
                        apply_altpolicy::Bool = false,
                        regime_switching::Bool = false, regime::Int = 1, n_regimes::Int = 2,
                        check_system::Bool = false, get_system::Bool = false,
                        get_population_moments::Bool = false,
                        verbose::Symbol = :high) where {T<:Real}

    if get_λ(m) == Inf
        # Then we just want the VECM approximation of the DSGE
        # with no additional cointegration
        return compute_system(m; apply_altpolicy = apply_altpolicy,
                              regime_switching = regime_switching, regime = regime, n_regimes = n_regimes,
                              check_system = check_system, get_system = get_system,
                              get_population_moments = get_population_moments, use_intercept = true,
                              verbose = verbose)
    else
        # Create a system using the method for DSGE-VECMs and λ = ∞
        system, DD_coint_add = compute_system(m; apply_altpolicy = apply_altpolicy,
                                regime_switching = regime_switching, regime = regime,
                                n_regimes = n_regimes, check_system = check_system,
                                get_system = true, use_intercept = true,
                                verbose = verbose)

        if get_system
            return system
        else
            EE, MM = measurement_error(m)

            lags = n_lags(m)
            YYYY, XXYY, XXXX =
                compute_var_population_moments(data, lags; use_intercept = true)
            out = vecm_approx_state_space(system[:TTT], system[:RRR], system[:QQ],
                                         system[:DD], system[:ZZ], EE, MM, size(data, 1),
                                          n_lags(m), n_cointegrating(m),
                                          n_cointegrating_add(m), DD_coint_add;
                                          get_population_moments = true,
                                          use_intercept = true)

            if get_population_moments
                return out..., YYYY, XXYY, XXXX
            else
                # Compute prior-weighted population moments
                λ = get_λ(m)
                YYYYC = YYYY + λ .* out[1]
                XXYYC = XXYY + λ .* out[2]
                XXXXC = XXXX + λ .* out[3]

                # Draw VECM system
                n_periods = size(data, 2) - lags
                β, Σ =  draw_VECM(YYYYC, XXYYC, XXXXC,
                                  convert(Int, n_periods + λ * n_periods),
                                  size(data, 1), lags, n_cointegrating(m))

                return β, Σ
            end
        end
    end
end

"""
```
compute_system(m::PoolModel{T})
```

Given the current model parameters, compute the state-space system
corresponding to the PoolModel model `m`.

Outputs

```
Φ: state transition function
Ψ: likelihood function, given weights on underlying models (the states) and predictive densities
F_ϵ: structural shock distribution
F_u: likelihood function measurement error distribution
F_λ: initial distribution of λ for state transition function
"""
function compute_system(m::PoolModel{T}; verbose::Symbol = :high,
                        regime_switching::Bool = false,
                        n_regimes::Int = 1) where T<:AbstractFloat
    Φ, F_ϵ, F_λ = transition(m)
    Ψ, F_u = measurement(m)
    return Φ, Ψ, F_ϵ, F_u, F_λ
end

"""
```
compute_system(m::AbstractDSGEModel, system::System;
        observables::Vector{Symbol} = collect(keys(m.observables)),
        pseudo_observables::Vector{Symbol} = collect(keys(m.pseudo_observables)),
        states::Vector{Symbol} = vcat(collect(keys(m.endogenou_states)),
                                      collect(keys(m.endogenous_states_augmented)))
        shocks::Vector{Symbol} = collect(keys(m.exogenous_shocks)),
        zero_DD = false, zero_DD_pseudo = false)
compute_system(m::AbstractDSGEVECMModel, system::System;
        observables::Vector{Symbol} = collect(keys(m.observables)),
        pseudo_observables::Vector{Symbol} = collect(keys(m.pseudo_observables)),
        cointegrating::Vector{Symbol} = collect(keys(m.cointegrating)),
        states::Vector{Symbol} = vcat(collect(keys(m.endogenou_states)),
                                      collect(keys(m.endogenous_states_augmented)))
        shocks::Vector{Symbol} = collect(keys(m.exogenous_shocks)),
        zero_DD = false, zero_DD_pseudo = false)
```
computes the corresponding transition and measurement
equations specified by the keywords (e.g. `states`, `pseudo_observables`)
using the existing `ZZ`, `ZZ_pseudo`, and `TTT` matrices in `system`.

Note that this function does not update the EE matrix, which is
set to all zeros. To incorporate measurement errors, the user
must specify the EE matrix after applying compute_system.

### Keywords
* `observables`: variables that should be
    entered into the new `ZZ` and `DD` matrices.
    The `observables` can be both Observables and PseudoObservables,
    but they must be an element of the system already.
* `pseudo_observables`: variables that should be
    entered into the new `ZZ_pseudo` and `DD_pseudo` matrices.
    The `observables` can be both Observables and PseudoObservables,
    but they must be an element of the system already.
* `cointegrating`: variables that should be
    entered into the new `ZZ` and `DD` matrices as cointegrating relationships.
    The `observables` can be both Observables and PseudoObservables,
    but they must be an element of the system already.
* `states`: variables that should be
    entered into the new `TTT` and `RRR` matrices as states.
    They must be existing states.
* `shocks`: variables that should be
    entered into the new `RRR` and `QQ` matrices as shocks.
    They must be existing exogenous shocks.
"""
function compute_system(m::AbstractDSGEModel{S}, system::System;
                        observables::Vector{Symbol} = collect(keys(m.observables)),
                        pseudo_observables::Vector{Symbol} =
                        collect(keys(m.pseudo_observables)),
                        states::Vector{Symbol} =
                        vcat(collect(keys(m.endogenous_states)),
                             collect(keys(m.endogenous_states_augmented))),
                        shocks::Vector{Symbol} = collect(keys(m.exogenous_shocks)),
                        zero_DD::Bool = false, zero_DD_pseudo::Bool = false,
                        check_system::Bool = true) where {S<:Real}

    # Set up indices
    oid  = m.observables # observables indices dictionary
    pid  = m.pseudo_observables # pseudo observables indices dictionary
    sid  = m.endogenous_states
    said = m.endogenous_states_augmented # states augmented

    # Find shocks to keep
    shock_inds = map(k -> m.exogenous_shocks[k], shocks)
    Qout = system[:QQ][shock_inds, shock_inds]

    # Find states to keep
    if !issubset(states, vcat(collect(keys(sid)), collect(keys(said))))
        false_states = setdiff(states, vcat(collect(keys(sid)), collect(keys(said))))
        error("The following states in keyword `states` do not exist in `system`: " *
              join(string.(false_states), ", "))
    elseif !isempty(setdiff(vcat(collect(keys(sid)), collect(keys(said))), states))
        which_states = Vector{Int}(undef, length(states))
        for i = 1:length(states)
            which_states[i] = haskey(sid, states[i]) ? sid[states[i]] : said[states[i]]
        end
        Tout = system[:TTT][which_states, which_states]
        Rout = system[:RRR][which_states, shock_inds]
        Cout = system[:CCC][which_states]
    else
        which_states = 1:n_states_augmented(m)
        Tout = copy(system[:TTT])
        Rout = system[:RRR][:, shock_inds]
        Cout = copy(system[:CCC])
    end

    # Compute new ZZ and DD matrices if different observables than current system
    if !isempty(symdiff(observables, collect(keys(oid))))
        Zout = zeros(S, length(observables), size(Tout, 1))
        Dout = zeros(S, length(observables))
        for (i, obs) in enumerate(observables)
            Zout[i, :], Dout[i] = if haskey(oid, obs)
                system[:ZZ][oid[obs], which_states], zero_DD ? zero(S) : system[:DD][oid[obs]]
            elseif haskey(pid, obs)
                system[:ZZ_pseudo][pid[obs], which_states], zero_DD ? zero(S) : system[:DD_pseudo][pid[obs]]
            else
                error("Observable/PseudoObservable $obs cannot be found in the DSGE model $m")
            end
        end
    else
        Zout = copy(system[:ZZ])[:, which_states]
        Dout = zero_DD ? zeros(S, size(Zout, 1)) : copy(system[:DD])
    end

    Eout = zeros(S, length(observables), length(observables)) # measurement errors are set to zero

    # Compute new ZZ_pseudo, DD_pseudo if different pseudo_observables than current system
    if !isempty(symdiff(pseudo_observables, collect(keys(pid))))
        Zpseudoout = zeros(S, length(pseudo_observables), size(Tout, 1))
        Dpseudoout = zeros(S, length(pseudo_observables))
        for (i, pseudoobs) in enumerate(pseudo_observables)
            Zpseudoout[i, :], Dpseudoout[i] = if haskey(oid, pseudoobs)
                system[:ZZ][oid[pseudoobs], which_states], zero_DD_pseudo ?
                    zero(S) : system[:DD][oid[pseudoobs]]
            elseif haskey(pid, pseudoobs)
                system[:ZZ_pseudo][pid[pseudoobs], which_states], zero_DD_pseudo ?
                    zero(S) : system[:DD_pseudo][pid[pseudoobs]]
            else
                error("Observable/PseudoObservable $pseudoobs cannot be found in the DSGE model $m")
            end
        end
    else
        Zpseudoout = copy(system[:ZZ_pseudo])[:, which_states]
        Dpseudoout = zero_DD_pseudo ? zeros(S, size(Zpseudoout, 1)) : copy(system[:DD_pseudo])
    end

    if check_system
        @assert size(Zout, 2) == size(Tout, 1) "Dimension 2 of new ZZ ($(size(Zout,2))) does not match dimension of states ($(size(Tout,1)))."
        @assert size(Qout, 1) == size(Rout, 2) "Dimension 2 of new RRR ($(size(Zout,2))) does not match dimension of shocks ($(size(Qout,1)))."
    end

    # Construct new system
    return System(Transition(Tout, Rout, Cout),
                  Measurement(Zout, Dout, Qout, Eout),
                  PseudoMeasurement(Zpseudoout, Dpseudoout))
end

function compute_system(m::AbstractDSGEVECMModel{S}, system::System;
                        observables::Vector{Symbol} = collect(keys(get_observables(get_dsge(m)))),
                        cointegrating::Vector{Symbol} = Vector{Symbol}(undef, 0),
                        cointegrating_add::Vector{Symbol} = Vector{Symbol}(undef, 0),
                        pseudo_observables::Vector{Symbol} =
                        collect(keys(get_dsge(m).pseudo_observables)),
                        states::Vector{Symbol} =
                        vcat(collect(keys(get_dsge(m).endogenous_states)),
                             collect(keys(get_dsge(m).endogenous_states_augmented))),
                        shocks::Vector{Symbol} = collect(keys(get_dsge(m).exogenous_shocks)),
                        zero_DD::Bool = false, zero_DD_pseudo::Bool = false,
                        get_DD_coint_add::Bool = false,
                        check_system::Bool = true) where {S<:Real}
    # Cointegrating relationships should exist as observables/pseudo_observables already
    # in the underlying DSGE. We assume cointegrating relationships come after normal observables.
    # Default behavior is to recreate the underlying DSGE's state space representation, however.
    sys = compute_system(get_dsge(m), system; observables = vcat(observables, cointegrating),
                         pseudo_observables = pseudo_observables,
                         states = states, shocks = shocks, zero_DD = zero_DD,
                         zero_DD_pseudo = zero_DD_pseudo, check_system = check_system)
    if get_DD_coint_add
        mtype = typeof(m)
        DD_coint_add = if hasmethod(compute_DD_coint_add, (mtype, Vector{Symbol}))
            compute_DD_coint_add(m, cointegrating_add)
        elseif hasmethod(compute_DD_coint_add, (mtype, ))
            compute_DD_coint_add(m)
        else
            compute_DD_coint_add(m, sys, cointegrating_add)
        end
        return sys, DD_coint_add
    else
        return sys
    end
end

"""
```
function compute_DD_coint_add(m::AbstractDSGEVECMModel{S}, system::System,
                              cointegrating_add::Vector{Symbol}) where {S <: Real}
```
computes `DD_coint_add` for a `DSGEVECM` model. This vector
holds additional cointegrating relationships that do not require
changes to the `ZZ` matrix.

### Note
We recommend overloading this function if there are
cointegrating relationships which a user does not want
to add to the underlying DSGE. The function `compute_system`
first checks for a method `compute_DD_coint_add` that takes
a Type tuple of `(model_type, Vector{Symbol})` and then `(model_type, )`
before calling this method.

This function is generally intended to be internal. As an example of
other such functions, `eqcond` must be user-defined but
is primarily used internally and not directly called by the user in a script.
"""
function compute_DD_coint_add(m::AbstractDSGEVECMModel{S}, system::System,
                              cointegrating_add::Vector{Symbol}) where {S <: Real}
    if !isempty(cointegrating_add)
        oid = get_observables(get_dsge(m))
        pid = get_pseudo_observables(get_dsge(m))
        Dout = zeros(S, length(cointegrating_add))
        for (i, obs) in enumerate(cointegrating_add)
            Dout[i] = if haskey(oid, obs)
                system[:DD][oid[obs]]
            elseif haskey(pid, obs)
                system[:DD_pseudo][pid[obs]]
            else
                error("Observable/PseudoObservable $obs cannot be found in the DSGE model $m")
            end
        end
        return Dout
    else
        @warn "No additional cointegrating relationships specified. Returning an empty vector."
        return Vector{S}(undef, 0)
    end
end

"""
```
compute_system_function(system::System{S}) where S<:AbstractFloat
```

### Inputs

- `system::System`: The output of compute_system(m), i.e. the matrix outputs from solving a given model, m.

### Outputs

- `Φ::Function`: transition equation
- `Ψ::Function`: measurement equation
- `F_ϵ::Distributions.MvNormal`: shock distribution
- `F_u::Distributions.MvNormal`: measurement error distribution
"""
function compute_system_function(system::System{S}) where S<:AbstractFloat
    # Unpack system
    TTT    = system[:TTT]
    RRR    = system[:RRR]
    CCC    = system[:CCC]
    QQ     = system[:QQ]
    ZZ     = system[:ZZ]
    DD     = system[:DD]
    EE     = system[:EE]

    # Define transition and measurement functions
    @inline Φ(s_t1::Vector{S}, ϵ_t::Vector{S}) = TTT*s_t1 + RRR*ϵ_t + CCC
    @inline Ψ(s_t::Vector{S}) = ZZ*s_t + DD

    # Define shock and measurement error distributions
    nshocks = size(QQ, 1)
    nobs    = size(EE, 1)
    F_ϵ = Distributions.MvNormal(zeros(nshocks), QQ)
    F_u = Distributions.MvNormal(zeros(nobs),    EE)

    return Φ, Ψ, F_ϵ, F_u
end

function zero_system_constants(system::System{S}) where S<:AbstractFloat
    system = copy(system)

    system.transition.CCC = zeros(size(system[:CCC]))
    system.measurement.DD = zeros(size(system[:DD]))
    system.pseudo_measurement.DD_pseudo = zeros(size(system[:DD_pseudo]))

    return system
end

"""
```
var_approx_state_space(TTT, RRR, QQQ, DD, ZZ, EE, MM, p; get_population_moments = false,
    use_intercept = false) where {S<:Real}
```
computes the VAR(p) approximation of the linear state space system

```
sₜ = TTT * sₜ₋₁ + RRR * ϵₜ,
yₜ = ZZ * sₜ + DD + uₜ,
```
where the disturbances are assumed to follow
```
ϵₜ ∼ 𝒩 (0, QQ),
uₜ = ηₜ + MM * ϵₜ,
ηₜ ∼ 𝒩 (0, EE).
```
The `MM` matrix implies
```
cov(ϵₜ, uₜ) = QQ * MM'.
```

### Outputs
If `get_population_moments = false`:
* `β`: VAR(p) coefficients
* `Σ`: innovations variance-covariance matrix for the VAR(p) representation
```
yₜ = Xₜβ + μₜ
```
where `Xₜ` appropriately stacks the constants and `p` lags of `yₜ`, and `μₜ ∼ 𝒩 (0, Σ)`.

If `get_population_moments = true`: we return the limit cross product matrices.
* `yyyyd`: 𝔼[y,y]
* `XXXXd`: 𝔼[y,X(lag rr)]
* `XXyyd`: 𝔼[X(lag rr),X(lag ll)]

Using these matrices, the VAR(p) representation is given by
```
β = XXXXd \\ XXyyd
Σ = yyyyd - XXyyd' * β
```

The keyword `use_intercept` specifies whether or not to use an
intercept term in the VAR approximation.
"""
function var_approx_state_space(TTT::AbstractMatrix{S}, RRR::AbstractMatrix{S},
                                QQ::AbstractMatrix{S}, DD::AbstractVector{S},
                                ZZ::AbstractMatrix{S}, EE::AbstractMatrix{S},
                                MM::AbstractMatrix{S}, p::Int;
                                get_population_moments::Bool = false,
                                use_intercept::Bool = false) where {S<:Real}

    n_obs = size(ZZ, 1)

    HH = EE + MM * QQ * MM'
    VV = QQ * MM'

    ## Compute p autocovariances

    ## Initialize Autocovariances
    GAMM0 = zeros(S, n_obs ^ 2, p + 1)
    GA0 =  solve_discrete_lyapunov(TTT, RRR * QQ * RRR')
    Gl   = ZZ * GA0 * ZZ' + ZZ * RRR * VV + (ZZ * RRR * VV)' + HH
    GAMM0[:, 1] = vec(Gl)
    TTl = copy(TTT)
    GA0ZZ = GA0 * ZZ'
    RRRVV = RRR * VV
    for l = 1:p
        Gl = ZZ * TTl * (GA0ZZ + RRRVV) # ZZ * (TTl * GA0Z) * ZZ' + ZZ * (TTl * RRR * VV)
        GAMM0[:, l+1] = vec(Gl)
        TTl = TTl * TTT
    end

    ## Create limit cross product matrices
    yyyyd = zeros(S, n_obs, n_obs)
    if use_intercept
        XXXXd = zeros(S, 1 + p * n_obs, 1 + p * n_obs)
        yyXXd = zeros(S, n_obs, 1 + p * n_obs)

        XXXXd[1, 1] = 1.
        XXXXd[1, 2:1 + p * n_obs] = repeat(DD', 1, p) # same as kron(ones(1, p), DD')
        XXXXd[2:1 + p * n_obs, 1] = repeat(DD, p)     # same as kron(ones(p), DD)
        yyXXd[:, 1] = DD
    else
        XXXXd = zeros(S, p * n_obs, p * n_obs)
        yyXXd = zeros(S, n_obs, p * n_obs)
    end
    DDDD = DD * DD'
    yyyyd = reshape(GAMM0[:, 1], n_obs, n_obs) + DDDD

    shift = use_intercept ? 1 : 0 # for constructing XXXXd, to select the right indices
    for rr = 1:p
        # 𝔼[yy,x(lag rr)]
        yyXXd[:, n_obs * (rr - 1) + 1 + shift:n_obs * rr + shift] =
            reshape(GAMM0[:, rr + 1], n_obs, n_obs) + DDDD

        # 𝔼[x(lag rr),x(lag ll)]
        for ll = rr:p
            yyyydrrll = reshape(GAMM0[:, ll - rr + 1], n_obs, n_obs) + DDDD
            XXXXd[n_obs * (rr - 1) + 1 + shift:n_obs * rr + shift,
                  n_obs * (ll - 1) + 1 + shift:n_obs * ll + shift] = yyyydrrll
            XXXXd[n_obs * (ll - 1) + 1 + shift:n_obs * ll + shift,
                  n_obs * (rr - 1) + 1 + shift:n_obs * rr + shift] = yyyydrrll'
        end
    end

    XXyyd = convert(Matrix{S}, yyXXd')

    if get_population_moments
        return yyyyd, XXyyd, XXXXd
    else
        β = \(XXXXd, XXyyd)
        Σ = yyyyd - XXyyd' * β
        return β, Σ
    end
end

"""
```
vecm_approx_state_space(TTT, RRR, QQQ, DD, ZZ, EE, MM, n_obs, p, n_coint,
    n_coint, n_coint_add, DD_coint_add; get_population_moments = false,
    use_intercept = false) where {S<:Real}
```
computes the VECM(p) approximation of the linear state space system

```
sₜ = TTT * sₜ₋₁ + RRR * ϵₜ,
yₜ = ZZ * sₜ + DD + uₜ,
```
where the disturbances are assumed to follow
```
ϵₜ ∼ 𝒩 (0, QQ),
uₜ = ηₜ + MM * ϵₜ,
ηₜ ∼ 𝒩 (0, EE).
```
The `MM` matrix implies
```
cov(ϵₜ, uₜ) = QQ * MM'.
```

### Outputs
If `get_population_moments = false`:
* `β`: VECM(p) coefficients. The first `n_coint + n_coint_add`
    coefficients for each observable comprise the error correction terms,
    while the following `1 + p * n_obs` terms are the VAR terms.
* `Σ`: innovations variance-covariance matrix for the VECM(p) representation
```
Δyₜ = eₜβₑ + Xₜβᵥ + μₜ
```
where `βₑ` are the coefficients for the error correction terms;
`eₜ` are the error correction terms specifying the cointegrating relationships;
`βᵥ` are the coefficients for the VAR terms; `Xₜ` appropriately stacks
the constants and `p` lags of `Δyₜ`; and `μₜ ∼ 𝒩 (0, Σ)`.

Note that the error correction terms satisfy the mapping
`eₜ' = C * yₜ₋₁`, where `C` is a matrix.

If `get_population_moments = true`: we return the limit cross product matrices.
* `yyyyd`: 𝔼[y,y]
* `XXXXd`: 𝔼[y,X(lag rr)]
* `XXyyd`: 𝔼[X(lag rr),X(lag ll)]

Note that in the rows of `XXyyd` and the rows and columns of `XXXXd`,
the cointegrating relationships are stacked above the constants and
lagged `Δyₜ`.

Using these matrices, the VAR(p) representation is given by
```
β = XXXXd \\ XXyyd
Σ = yyyyd - XXyyd' * β,
```
where `β` has dimensions `n_obs × (n_coint + n_coint_add + 1 + p * n_obs)`,
and `Σ` has dimensions `n_obs × n_obs`.

The keyword `use_intercept` specifies whether or not to use an
intercept term in the VECM approximation.
"""
function vecm_approx_state_space(TTT::AbstractMatrix{S}, RRR::AbstractMatrix{S},
                                 QQ::AbstractMatrix{S}, DD::AbstractVector{S},
                                 ZZ::AbstractMatrix{S}, EE::AbstractMatrix{S},
                                 MM::AbstractMatrix{S}, n_obs::Int, p::Int,
                                 n_coint::Int, n_coint_add::Int = 0,
                                 DD_coint_add::AbstractVector{S} = Vector{S}(undef, 0);
                                 get_population_moments::Bool = false,
                                 use_intercept::Bool = false,
                                 test_GA0::AbstractMatrix{S} =
                                 Matrix{S}(undef, 0, 0)) where {S <: Real}

    # n_obs is number of observables, n_coint is number of cointegrating relationships
    n_coint_all = n_coint + n_coint_add

    # Create variance-covariance matrices w/measurement error included
    HH = EE + MM * QQ * MM'
    VV = QQ * MM'

    ## Compute p autocovariances

    ## Initialize Autocovariances
    GAMM0 = zeros(S, (n_obs + n_coint) ^ 2, p + 1)
    GA0 = isempty(test_GA0) ? solve_discrete_lyapunov(TTT, RRR * QQ * RRR') : test_GA0 # Matlab code uses a different numerical procedure -> some difference in this matrix
    Gl   = ZZ * GA0 * ZZ' + ZZ * RRR * VV + (ZZ * RRR * VV)' + HH
    GAMM0[:, 1] = vec(Gl)
    TTl = copy(TTT)
    GA0ZZ = GA0 * ZZ'
    RRRVV = RRR * VV
    for l = 1:p
        Gl = ZZ * TTl * (GA0ZZ + RRRVV) # ZZ * (TTl * GA0) * ZZ' + ZZ * (TTl * RRR * VV)
        GAMM0[:, l+1] = vec(Gl)
        TTl = TTl * TTT
    end

    ## Create limit cross product matrices
    DDDD = DD * DD'
    yyyyd_coint0 = reshape(GAMM0[:, 1], n_obs + n_coint, n_obs + n_coint) + DDDD
    yyyyd_coint1 = reshape(GAMM0[:, 2], n_obs + n_coint, n_obs + n_coint) + DDDD
    yyyyd = yyyyd_coint0[1:n_obs, 1:n_obs]

    # n_coint_add are treated as the first set of variables in XX
    # n_coint    are treated as the second set of variables in XX
    # composition: n_coint_add - n_coint - constant - lags
    if use_intercept
        XXXXd = zeros(S, 1 + p * n_obs + n_coint_all, 1 + p * n_obs + n_coint_all)
        yyXXd = zeros(S, n_obs, 1 + p * n_obs + n_coint_all)

        # 𝔼[x(n_coint), x(n_coint)]
        XXXXd[n_coint_add + 1:n_coint_all, n_coint_add + 1:n_coint_all] = yyyyd_coint0[n_obs + 1:n_obs + n_coint, n_obs + 1:n_obs + n_coint]

        # 𝔼[x(const), x(const)]
        XXXXd[n_coint_all + 1, n_coint_all + 1] = 1.

        # 𝔼[x(n_coint), x(const)]
        XXXXd[n_coint_add + 1:n_coint_all, n_coint_all + 1] = DD[n_obs + 1:n_obs + n_coint]
        XXXXd[n_coint_all + 1, n_coint_add + 1:n_coint_all] = DD[n_obs + 1:n_obs + n_coint]

        # 𝔼[x(const), x(lags)]
        XXXXd[n_coint_all + 1, n_coint_all + 2:n_coint_all + 1 + p * n_obs] = repeat(DD[1:n_obs]', 1, p)
        XXXXd[n_coint_all + 2:n_coint_all + 1 + p * n_obs, n_coint_all + 1] =
            XXXXd[n_coint_all + 1, n_coint_all + 2:n_coint_all + 1 + p * n_obs]' # same as kron(ones(p), DD[1:n_obs]) but avoids calculation

        # 𝔼[yy, x(n_coint)]
        yyXXd[:, n_coint_add + 1:n_coint_all] = yyyyd_coint1[1:n_obs, n_obs + 1:n_obs + n_coint]

        # 𝔼[yy, x(const)]
        yyXXd[:, n_coint_all + 1] = DD[1:n_obs]
    else
        XXXXd = zeros(S, p * n_obs + n_coint_all, p * n_obs + n_coint_all)
        yyXXd = zeros(S, n_obs, p * n_obs + n_coint_all)

        # 𝔼[x(n_coint), x(n_coint)]
        XXXXd[n_coint_add + 1:n_coint_all, n_coint_add + 1:n_coint_all] =
            yyyyd_coint0[n_obs + 1:n_obs + n_coint, n_obs + 1:n_obs + n_coint]

        # 𝔼[yy, x(n_coint)]
        yyXXd[:, n_coint_add + 1:n_coint_all] = yyyyd_coint1[1:n_obs, n_obs + 1:n_obs + n_coint]
    end

    if n_coint_add > 0
        DD_coint_add_div2 = DD_coint_add ./ 2
        if use_intercept
            # 𝔼[yy, x(n_coint_add)]
            yyXXd[:, 1:n_coint_add] = DD[1:n_obs] * DD_coint_add_div2

            # 𝔼[x(n_coint_add), x(n_coint_add)]
            XXXXd[1:n_coint_add, 1:n_coint_add] = DD_coint_add * DD_coint_add' ./ 3

            # 𝔼[x(n_coint_add), x(n_coint)]
            XXXXd[1:n_coint_add, 1 + n_coint_add:n_coint_all] =
                DD_coint_add_div2 * DD[n_obs + 1: n_obs + n_coint]'
            XXXXd[1 + n_coint_add:n_coint_all, 1:n_coint_add] =
                XXXXd[1:n_coint_add, 1 + n_coint_add:n_coint_all]' # transpose of the previous line

            # 𝔼[x(n_coint_add), x(const)]
            XXXXd[1:n_coint_add, n_coint_all + 1] = DD_coint_add_div2
            XXXXd[n_coint_all + 1, 1:n_coint_add] = DD_coint_add_div2'
        else
            # 𝔼[yy, x(n_coint_add)]
            yyXXd[:, 1:n_coint_add] = DD[1:n_obs] * DD_coint_add_div2

            # 𝔼[x(n_coint_add), x(n_coint_add)]
            XXXXd[1:n_coint_add, 1:n_coint_add] = DD_coint_add * DD_coint_add' ./ 3

            # 𝔼[x(n_coint_add), x(n_coint)]
            XXXXd[1:n_coint_add, 1 + n_coint_add:n_coint_all] =
                DD_coint_add_div2 * DD[n_obs + 1: n_obs + n_coint]'
            XXXXd[1 + n_coint_add:n_coint_all, 1:n_coint_add] =
                XXXXd[1:n_coint_add, 1 + n_coint_add:n_coint_all]' # transpose of the previous line
        end
    end

    shift = use_intercept ? 1 : 0 # for constructing XXXXd, to select the right indices
    for rr = 1:p
        # 𝔼[yy, x(lag rr)]
        yyyyd_cointrr = reshape(GAMM0[:, rr + 1], n_obs + n_coint, n_obs + n_coint) + DDDD
        yyXXd[:, n_coint_all + 1 + n_obs * (rr - 1) + shift:n_coint_all + n_obs * rr + shift] =
            yyyyd_cointrr[1:n_obs, 1:n_obs]

        if n_coint_add > 0
            # 𝔼[x(n_coint_add), x(lag rr)]
            XXXXd[1:n_coint_add, n_coint_all + 1 + n_obs * (rr - 1) + shift:n_coint_all + n_obs * rr + shift] =
                DD_coint_add_div2 * DD[1:n_obs]'
            XXXXd[n_coint_all + 1 + n_obs * (rr - 1) + shift:n_coint_all + n_obs * rr + shift, 1:n_coint_add] =
                XXXXd[1:n_coint_add, n_coint_all + 1 + n_obs * (rr - 1) + shift:n_coint_all + n_obs * rr + shift]'
        end

        # 𝔼[x(n_coint), x(lag rr)]
        yyyyd_cointrr1 = reshape(GAMM0[:, rr], n_obs + n_coint, n_obs + n_coint) + DDDD
        XXXXd[n_coint_add + 1:n_coint_all, n_coint_all + 1 + n_obs * (rr - 1) + shift:n_coint_all + n_obs * rr + shift] =
            yyyyd_cointrr1[n_obs + 1:n_obs + n_coint, 1:n_obs]
        XXXXd[n_coint_all + 1 + n_obs * (rr - 1) + shift:n_coint_all + n_obs * rr + shift, n_coint_add + 1:n_coint_all] =
            yyyyd_cointrr1[n_obs + 1:n_obs + n_coint, 1:n_obs]'

        # 𝔼[x(lag rr), x(lag ll)]
        for ll = rr:p
            yyyyd_cointrrll = reshape(GAMM0[:, ll - rr + 1], n_obs + n_coint, n_obs + n_coint) + DDDD
            XXXXd[n_coint_all + 1 + n_obs * (rr - 1) + shift:n_coint_all + n_obs * rr + shift,
                  n_coint_all + 1 + n_obs * (ll - 1) + shift:n_coint_all + n_obs * ll + shift] = yyyyd_cointrrll[1:n_obs, 1:n_obs]
            XXXXd[n_coint_all + 1 + n_obs * (ll - 1) + shift:n_coint_all + n_obs * ll + shift,
                  n_coint_all + 1 + n_obs * (rr - 1) + shift:n_coint_all + n_obs * rr + shift] = yyyyd_cointrrll[1:n_obs, 1:n_obs]'
        end
    end

    XXyyd = convert(Matrix{S}, yyXXd')

    if get_population_moments
        return yyyyd, XXyyd, XXXXd
    else
        β = XXXXd \ XXyyd
        Σ = yyyyd - XXyyd' * β
        return β, Σ
    end
end


module DEACpopulation
using Random;
const functionNames = ["randPop",
                       "flatPop",
                       "gaussPop",
                       "gaussPop2",
                       "gaussOne",
                       "gaussTwo",
                       "gaussThree"]
#Generate population functions
function randPop(rng::MersenneTwister,genomeSize::Int64,populationSize::Int64,omega::Array{Float64,1})
    rand(rng,genomeSize,populationSize)
end

function flatPop(rng::MersenneTwister,genomeSize::Int64,populationSize::Int64,omega::Array{Float64,1})
    rand(rng) .* ones(genomeSize,populationSize)
end

function gaussPop(rng::MersenneTwister,genomeSize::Int64,populationSize::Int64,omega::Array{Float64,1})
    p = (rand(rng,populationSize));
    mu = (rand(rng,populationSize)*maximum(omega));
    alpha = (rand(rng,populationSize)*maximum(omega)/4);
    (p'.*exp.(-((omega .- mu').^2)./(alpha'.^2)/2)./alpha'/sqrt(2*pi))
end

function gaussPop2(rng::MersenneTwister,genomeSize::Int64,populationSize::Int64,omega::Array{Float64,1})
    p = ones(populationSize);
    mu = Array{Float64}(undef,populationSize);
    for i = 0:populationSize-1
        mu[i+1]=omega[(i%genomeSize)+1];
    end
    alpha = (rand(rng::MersenneTwister,populationSize));
    (p'.*exp.(-((omega .- mu').^2)./(alpha'.^2)/2)./alpha'/sqrt(2*pi));
end

function gaussOne(rng::MersenneTwister,genomeSize::Int64,populationSize::Int64,omega::Array{Float64,1})
    gaussPop(rng,genomeSize,populationSize,omega)
end

function gaussTwo(rng::MersenneTwister,genomeSize::Int64,populationSize::Int64,omega::Array{Float64,1})
    gaussPop(rng,genomeSize,populationSize,omega) + gaussPop(rng,genomeSize,populationSize,omega)
end

function gaussThree(rng::MersenneTwister,genomeSize::Int64,populationSize::Int64,omega::Array{Float64,1})
    gaussPop(rng,genomeSize,populationSize,omega) + gaussPop(rng,genomeSize,populationSize,omega) + gaussPop(rng,genomeSize,populationSize,omega)
end

end


input = joinpath(@__DIR__, "input")
lines = readlines(input)

m = hcat(map(line -> map(n -> parse(Int, n), collect(line)), lines)...)

function neighbors(m, idx)
    x, y = Tuple(idx)
    return filter(c -> checkbounds(Bool, m, c), [
        CartesianIndex(x + 1, y),
        CartesianIndex(x - 1, y),
        CartesianIndex(x, y + 1),
        CartesianIndex(x, y - 1),
        CartesianIndex(x + 1, y + 1),
        CartesianIndex(x - 1, y - 1),
        CartesianIndex(x - 1, y + 1),
        CartesianIndex(x + 1, y - 1)
    ])
end

function step(m, limit)
    counter = 0
    for i in 1:limit
        m = m .+ 1
        octopii = Set()
        flashed = Set()
        for idx in CartesianIndices(m)
            push!(octopii, idx)
            while !isempty(octopii)
                o = pop!(octopii)
                if m[o] > 9 && !(o in flashed)
                    push!(flashed, o)
                    ns = neighbors(m, o)
                    for n in ns
                        m[n] = m[n] + 1
                        push!(octopii, n)
                    end
                end
            end
        end
        for f in flashed
            m[f] = 0
        end
        if length(flashed) == 100
            return (i, counter)
        end
        counter += length(flashed)
    end
    return (limit, counter)
end

(_, p1) = step(m, 100)
(p2, _) = step(m, 1000)

println("-----------------------------------------------------------------------")
println("dumbo octopus -- part one :: $p1")
println("dumbo octopus -- part two :: $p2")
println("-----------------------------------------------------------------------")

@assert(p1 == 1632)
@assert(p2 == 303)


using Random, LinearAlgebra, SpecialFunctions, PyPlot, KernelDensity, PyCall
stats = pyimport("scipy.stats")

"Generate (m, n) matrix uniformly distributed over Stiefel manifold."
function rando(m::Integer, n::Integer)
    M = randn((m, n))
    Q = Matrix(qr(M).Q)
end

"Generate matrix with exponentially-distributed singular values and singular vectors ~ O(n)."
function randA(n::Integer, scale::Real)
    s = randexp(n)*scale
    V = rando(n, n)
    U = rando(n, n)
    A = U*Diagonal(s)*transpose(V)
    return A, s, U, V
end

"Generate random symmetric matrix with exponentially-distributed eigenvalues and eigenvectors ~ O(n)."
function randS(n::Integer, scale::Real)
    s = randexp(n)*scale
    U = rando(n, n)
    A = U*Diagonal(s)*transpose(U)
    return A, s, U
end

"Generate solution vector from x0 + V N(0, I) + Y G N(0, I)."
function randx(x0::Array{<:Real, 1}, V::Array{<:Real, 2}, Y::Array{<:Real, 2}, G::Array{<:Real, 2})
    v = randn(size(V)[2])
    y = randn(size(Y)[2])
    x = x0 + V*v + Y*G*y
end

"Generate solution vector from x0 + N(0, I)."
function randx(x0::Array{<:Real, 1})
    x = x0 + randn(size(x0))
end

"Generate solution vector from x0 + L N(0, I)."
function randx(x0::Array{<:Real, 1}, L::Array{<:Real, 2})
    x = x0 + L*randn(size(L)[2])
end

"Naive implementation of projection method."
function projection(x0::Array{<:Real, 1}, A::Array{<:Real, 2}, b::Array{<:Real, 1}, V::Array{<:Real, 2}, W::Array{<:Real, 2})
    x = x0 + V*inv(transpose(W)*A*V)*transpose(W)*(b - A*x0)
end

"Krylov subspace K_{n}(A, v)."
function Krylov(A::Array{<:Real, 2}, v::Array{<:Real, 1}, n::Integer)
    V = Array{Float64, 2}(undef, size(v)[1], n)
    V[:, 1] = v/norm(v)
    for i=2:n
        V[:, i] = A*V[:, i-1]
        for j=1:i-1
            V[:, i] = V[:, i] - dot(V[:, i], V[:, j])*V[:, j]
        end
        V[:, i] = V[:, i]/norm(V[:, i])
    end
    return V
end

"PDF of chi-squared distribution."
function chi2PDF(x::Real, n::Integer)
    return x^(n/2-1)*exp(-x/2)/(2^(n/2)*gamma(n/2))
end

"PDF of F-distribution."
function FPDF(x::Real, d1::Integer, d2::Integer)
    return (d1/d2)^(d1/2)*x^(d1/2-1)*(1+d1*x/d2)^(-(d1+d2)/2)/beta(d1/2, d2/2)
end

"Compute orthonormal basis for Null(A), where A is full-rank MxN matrix with N>M."
function NullSpace(A::Array{<:Real, 2})
    svdA = svd(A, full=true)
    return svdA.V[:, (size(A)[1]+1):size(A)[2]]
end

"Approximation of int dx|F1-F2| with central Riemann sum."
function Integrate(F1::Array{<:Real, 1}, F2::Array{<:Real, 1}, h::Real)
    F = abs.(F1 - F2)
    s = h*(sum(F[2:(size(F)[1]-1)]) + (F[1] + F[size(F)[1]])/2)
end

"Symmetric matices, separate estimation of the scale for each run, mean of inverse gamma."
function SymSamplingIn(N_rounds::Integer, scale::Real, m::Integer,  N::Integer=100, N_samples::Integer=5000; randomization="A")
    A, s, U = randS(N, scale)
    L = U*Diagonal(s.^(-1/2))
    inv_A = U*Diagonal(s.^(-1))*transpose(U)
    Statistics_inv = zeros(N_samples)
    Statistics_proj = zeros(N_samples)
    if randomization == "A"
        draw = (n=N, l=L) -> randx(zeros(n), l)
    else
        draw = (n=N) -> randx(zeros(n))
    end

    for i=1:N_samples
        x = draw()
        b = A*x
        V = Krylov(A, b, m)
        y = projection(zeros(N), A, b, V, V)
        sigma = inv_A - V*inv(transpose(V)*A*V)*transpose(V)
        statistics = transpose(x - y)*pinv(sigma)*(x - y)
        Statistics_inv[i] = statistics

        alpha = 0
        for i=1:N_rounds
            z = draw()
            b = A*z
            V = Krylov(A, b, m)
            P = I - V*inv(transpose(V)*A*V)*transpose(V)*A
            sample = P*z
            alpha = alpha + norm(sample)^2
        end
        alpha = alpha/(N_rounds*(N-m)+2)

        statistics = (norm(x - y)^2)/alpha
        Statistics_proj[i] = statistics
    end
    x = LinRange(0, 500, 5000)
    kernel_inv = stats.gaussian_kde(Statistics_inv.*(Statistics_inv.<500).*(Statistics_inv.>0))
    kernel_proj = stats.gaussian_kde(Statistics_proj.*(Statistics_proj.<500))
    stat_inv = kernel_inv(x)
    stat_proj = kernel_proj(x)
    return stat_inv, stat_proj
end

"Symmetric matices, the single estimation of the scale for all run, mean of inverse gamma."
function SymSamplingOut(N_rounds::Integer, scale::Real, m::Integer,  N::Integer=100, N_samples::Integer=5000; randomization="A")
    A, s, U = randS(N, scale)
    L = U*Diagonal(s.^(-1/2))
    inv_A = U*Diagonal(s.^(-1))*transpose(U)
    Statistics_inv = zeros(N_samples)
    Statistics_proj = zeros(N_samples)
    if randomization == "A"
        draw = (n=N, l=L) -> randx(zeros(n), l)
    else
        draw = (n=N) -> randx(zeros(n))
    end

    alpha = 0
    for i=1:N_rounds
        z = draw()
        b = A*z
        V = Krylov(A, b, m)
        P = I - V*inv(transpose(V)*A*V)*transpose(V)*A
        sample = P*z
        alpha = alpha + norm(sample)^2
    end
    alpha = alpha/(N_rounds*(N-m)+2)

    for i=1:N_samples
        x = draw()
        b = A*x
        V = Krylov(A, b, m)
        y = projection(zeros(N), A, b, V, V)
        sigma = inv_A - V*inv(transpose(V)*A*V)*transpose(V)
        statistics = transpose(x - y)*pinv(sigma)*(x - y)
        Statistics_inv[i] = statistics

        statistics = (norm(x - y)^2)/alpha
        Statistics_proj[i] = statistics
    end
    x = LinRange(0, 500, 5000)
    kernel_inv = stats.gaussian_kde(Statistics_inv.*(Statistics_inv.<500).*(Statistics_inv.>0))
    kernel_proj = stats.gaussian_kde(Statistics_proj.*(Statistics_proj.<500))
    stat_inv = kernel_inv(x)
    stat_proj = kernel_proj(x)
    return stat_inv, stat_proj
end

"Symmetric matices, separate estimation of the scale for each run, multivariate Student."
function HSymSamplingIn(N_rounds::Integer, scale::Real, m::Integer,  N::Integer=100, N_samples::Integer=5000; randomization="A")
    A, s, U = randS(N, scale)
    L = U*Diagonal(s.^(-1/2))
    inv_A = U*Diagonal(s.^(-1))*transpose(U)
    Statistics_proj = zeros(N_samples)
    if randomization == "A"
        draw = (n=N, l=L) -> randx(zeros(n), l)
    else
        draw = (n=N) -> randx(zeros(n))
    end

    for i=1:N_samples
        alpha = 0
        for i=1:N_rounds
            z = draw()
            b = A*z
            V = Krylov(A, b, m)
            P = I - V*inv(transpose(V)*A*V)*transpose(V)*A
            sample = P*z
            alpha = alpha + norm(sample)^2
        end
        alpha = alpha/N_rounds

        x = draw()
        b = A*x
        V = Krylov(A, b, m)
        y = projection(zeros(N), A, b, V, V)
        statistics = (norm(x - y)^2)/alpha
        Statistics_proj[i] = statistics
    end

    x = LinRange(0, 10, 5000)
    kernel_proj = stats.gaussian_kde(Statistics_proj.*(Statistics_proj.<10))
    stat_proj = kernel_proj(x)
    return stat_proj
end

"Symmetric matices, the single estimation of the scale for all run, multivariate Student."
function HSymSamplingOut(N_rounds::Integer, scale::Real, m::Integer,  N::Integer=100, N_samples::Integer=5000; randomization="A")
    A, s, U = randS(N, scale)
    L = U*Diagonal(s.^(-1/2))
    inv_A = U*Diagonal(s.^(-1))*transpose(U)
    Statistics_proj = zeros(N_samples)
    if randomization == "A"
        draw = (n=N, l=L) -> randx(zeros(n), l)
    else
        draw = (n=N) -> randx(zeros(n))
    end

    alpha = 0
    for i=1:N_rounds
        z = draw()
        b = A*z
        V = Krylov(A, b, m)
        P = I - V*inv(transpose(V)*A*V)*transpose(V)*A
        sample = P*z
        alpha = alpha + norm(sample)^2
    end
    alpha = alpha/N_rounds

    for i=1:N_samples
        x = draw()
        b = A*x
        V = Krylov(A, b, m)
        y = projection(zeros(N), A, b, V, V)
        statistics = (norm(x - y)^2)/alpha
        Statistics_proj[i] = statistics
    end

    x = LinRange(0, 10, 5000)
    kernel_proj = stats.gaussian_kde(Statistics_proj.*(Statistics_proj.<10))
    stat_proj = kernel_proj(x)
    return stat_proj
end


using DynamicalSystems
using Plots
plotly()

# Simulate a sine wave using DynamicalSystems, compare accuracy and timing to builtin sine and use this to construct Process type containing (only) relevant metadata

X0 = [0.0]
T = 2000.0
μ = 1.0 # Hz

function dynamicalSine(X, μ, t)
    dX = μ.*cos(μ.*t) # dx/dt
    return SVector{1}(dX)
end
function dynamicalSineJ(X, μ, t)
    J = @SMatrix [-μ^2*sin(μ.*t)]
end

# (x, μ, t) -> sin(μ.*t)
ds = ContinuousDynamicalSystem(dynamicalSine, X0, μ, dynamicalSineJ) # Easy system defined anonymously
data = trajectory(ds, T; dt=0.01)

plot(data[:, 1], seriestype=:path)
plot!(title="Sine")


using Distributed
# add processes to workspace
addprocs(2)

@everywhere using ModelingToolkit, OrdinaryDiffEq

# create the Lorenz system
@everywhere @parameters t σ ρ β
@everywhere @variables x(t) y(t) z(t)
@everywhere D = Differential(t)

@everywhere eqs = [D(x) ~ σ * (y - x),
    D(y) ~ x * (ρ - z) - y,
    D(z) ~ x * y - β * z]

@everywhere @named de = ODESystem(eqs)
@everywhere ode_func = ODEFunction(de, [x, y, z], [σ, ρ, β])

@everywhere u0 = [19.0, 20.0, 50.0]
@everywhere params = [16.0, 45.92, 4]

@everywhere ode_prob = ODEProblem(ode_func, u0, (0.0, 10.0), params)

@everywhere begin
    using OrdinaryDiffEq
    using ModelingToolkit

    function solve_lorenz(ode_problem)
        print(solve(ode_problem, Tsit5()))
    end
end

solve_lorenz(ode_prob)

future = @spawn solve_lorenz(ode_prob)
@test_broken fetch(future)


using ControlSystems, RobustAndOptimalControl
plotly(size=(1850, 900), show=false)
##
w = 2pi*exp10.(LinRange(-2, 2, 400))
function flexible_servo_loop(; kp=1, kv=1, ki=1, Tf=0.1, Jm=1, Jl=10, c=1, k=10000)

    A = [
        0.0 1 0 0
        -k/Jm -c/Jm k/Jm c/Jm
        0 0 0 1
        k/Jl c/Jl -k/Jl -c/Jl
    ]
    B = [0, 1/Jm, 0, 0]
    C = I(4)
    D11 = D12 = D21 = D22 = 0
    nx,nu,ny = 4,1,1
    P = named_ss(ss(A,B,C,0), u=:u, x=[:qm, :qdm, :qa, :qda], y = [:qm, :qdm, :qa, :qda])

    K = let
        K = cascade_controller(; kp, kv, ki, Tf) # TODO: extend to take in several measurements instead of differentiating
        K = named_ss(K, u=:e, y = :Cu)
        # extended_controller(K)
    end

    sumE = sumblock("e = r - y")

    u1 = [:u,  :y, :e]
    y1 = [:Cu, :qm, :e]
    w1 = [:r]
    G = connect([P, K, sumE]; y1, u1, w1)
    #F = ss([tf(0, [1, 0]); tf(1)])

    (; P, K, G)
end

function cascade_controller(; kp=1, kv=1, ki=1, Tf=0.1)
    ss(tf([kv, kp*kv + ki, kp*ki], [Tf, 1, 0]))
end

#
P,K,G = flexible_servo_loop(kp=0.1, kv=250, ki=15, Tf=0.001)

plot(
    # plot(step(c2d(G.sys, 0.001), 3)),
    plot(step(G, 3)),
    bodeplot(P, w, hz=true, plotphase=false),
    bodeplot(G, w, hz=true, plotphase=false),
    layout=(1,3),
    title="",
)
display(current())


function barzanLinSolve(responseVars, features)

	model = zeros(size(features,2)+1, size(responseVars,2));
	for i=1:size(responseVars,2)
		model[:,i] = llsq([features ones(size(features,1),1)],responseVars[:,i]);
	end

	return model

end


using MDCDL
using LinearAlgebra
using Images, TestImages
using Plots

using Base.Filesystem
using Dates

function main()
    ########## Configurations #########
    # choose NSOLT type: (Rnsolt | Cnsolt)
    Nsolt = Rnsolt
    # TP := eltype(Nsolt) (<:AbstractFloat)
    TP = Float64
    # decimation factor: (<:NTuple{D,Int} where D is #dims)
    df = (2,2)
    # polyphase order: (<:NTuple{D,Int} where D)
    ord = (2,2)
    # number of channels: (<:Union{Integer,Tuple{Int,Int}} for Rnsolt)
    #                     (<:Integer for Cnsolt)
    nch = 8
    # number of tree level (<: Integer)
    level = 3

    # size of minibatches (<:NTuple{D,Int})
    szx = (16,16)
    # number of minibatches (<:Integer)
    nSubData = 16

    # path of log files (do nothing if isa(logdir, Nothing))
    # logdir = joinpath(@__DIR__, "results", Dates.format(now(), "yyyymmdd_HH_MM_SS_sss"))
    logdir = nothing
    # save minibatches?
    do_save_trainingset = false
    do_export_atoms = false

    # options for sparse coding
    sparsity = 0.25
    sparsecoder = SparseCoders.IHT
    sparsecoder_options = (
        iterations = 400,
        nonzeros = trunc(Int, sparsity * prod(szx)), filter_domain=:convolution,)
    # options for dictionary update
    optimizer = Optimizers.Steepest
    optimizer_options = (
        rate = 1e-4,
        iterations=10,)

    # general options of dictionary learning
    options = ( epochs  = 1000,
                verbose = :standard, # :none, :standard, :specified, :loquacious
                sparsecoder = sparsecoder,
                sparsecoder_options = sparsecoder_options,
                optimizer = optimizer,
                optimizer_options = optimizer_options,
                logdir = logdir,)
    ####################################
    logdir != nothing && !isdir(logdir) && mkpath(logdir)

    # original image
    orgImg = testimage("cameraman")

    # generate minibatches of orgImage
    trainingSet = map(1:nSubData) do nsd
        pos = rand.(UnitRange.(0,size(orgImg) .- szx))
        orgImg[UnitRange.(1 .+ pos, szx .+ pos)...]
    end

    if logdir != nothing && do_save_trainingset
        datadir = joinpath(logdir, "data")
        !isdir(datadir) && mkpath(datadir)
        map(idx->save(joinpath(datadir, "$idx.png"), trainingSet[idx]), 1:nSubData)
    end

    # create NSOLT instance
    nsolt = Nsolt(TP, df, ord, nch)
    # set random orthonormal matrices to the initial matrices.
    MDCDL.rand!(nsolt, isPropMat = false, isPropAng = false, isSymmetry = false)
    istype1(nsolt) && MDCDL.vm1constraint!(nsolt)

    msnsolt = Multiscale([ deepcopy(nsolt) for l in 1:level ]...)
    # rand!(msnsolt.filterbanks[level], isSymmetry = false, isPropAng=false, isPropMat = false)

    # dictionary learning
    MDCDL.train!(msnsolt, trainingSet; options...)

    if logdir != nothing && do_export_atoms
        for idx = 1:length(msnsolt)
            png(plot(msnsolt.filterbanks[idx]), joinpath(logdir, string("nsolt_",idx,".png")))
        end
    end

    return msnsolt
end

main()


@model function model_parametric_cases(
    num_impute,        # [Int] num. of days for which to impute infections
    num_total_days,    # [Int] days of observed data + num. of days to forecast
    cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
    deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
    π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    epidemic_start,    # [AbstractVector{<:Int}]
    population,        # [AbstractVector{<:Real}]
    serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
	num_iar_steps,         # [Int] number of weeks (21)
    iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
    lockdown,         # [Int] number of weeks (21)
    num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
	weekday,
	holiday,
	covariates,
	covariates_start,
	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    hospit,
	seroprev_mean,
    seroprev_std,
    seroprev_idx,
	π4,
    predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
    ::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
	invlink = log,#KLogistic(V(3))
	link    = exp,#KLogit(V(k))
) where {TV, V}
    ############ 0.) some definitions
	num_obs        = length(cases)
    num_Rt_steps   = covariates_start-lockdown-1
	num_sero       = length(seroprev_mean)
	num_covariates = size(covariates, 2)
    cases_start    = num_obs-num_case_obs
	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
    num_time_steps = predict ? num_total_days : num_obs

    ############# 2.) time varying reproduction number
	R0        ~ truncated(Normal(3., 1.), 1., 5.)
	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
	σ_rt      ~ truncated(Normal(0.05, .03), 0, .3)
	latent_Rt ~ RandomWalk(num_Rt_steps, σ_rt, invlink(R1))
	α         ~ filldist(Exponential(0.2), num_covariates)

    Rt                                = TV(undef, num_time_steps)
    Rt[1:lockdown]                   .= R0
    Rt[lockdown+1:covariates_start-1] = link.(latent_Rt)
    Rt[covariates_start:num_obs]      = link.(
        latent_Rt[end] .+ covariates * α
    )
    if predict #if prediction
		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
	end

    ############ 3.) infection dynamics
	## 3.1) initialize serial interval
	num_si   = 15
	si       = serial_intervals
	# μ_si   ~ truncated(Normal(5.06, 0.33), 0, Inf)
	# σ_si   = 2.11
	# si     = pdf.( Ref(GammaMeanStd(μ_si, σ_si)), 1:num_si )
	# si    /= sum(si)

	## 3.2) seed epidemic
	τ ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
	T = typeof(τ)
    y ~ truncated(Exponential(τ),T(0),T(1000))

	newly_infected         = TV(undef, num_time_steps)
	cumulative_infected    = TV(undef, num_time_steps)
	newly_infected[1]      = y
	cumulative_infected[1] = zero(V)
	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
	end

    ## 3.3) run disease dynamics
	for t = (num_impute + 1):num_time_steps
		# Update cumulatively infected
		cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		# Adjusts for portion of population that is susceptible
		St = max(population - cumulative_infected[t], 0) / population
		# effective number of infectious individuals
		effectively_infectious = sum(newly_infected[τ] * si[t - τ] for τ = (t - 1):-1:max(t-num_si,1))
		# number of new infections (unobserved)
		newly_infected[t] = St * Rt[t] * effectively_infectious
	end

	########### 4.) derive observables
	## 4.1) initialize delay distributions
	num_i2c = 40
	num_i2h = 40
	num_i2d = 60

	μ_i2c ~ truncated(Normal(4.5, 1.), 3, 7)
    μ_i2h ~ truncated(Normal(12., 3.), 9, 16)
    μ_i2d ~ truncated(Normal(20., 4.), 10, 30)

	ϕ_i2c = 5.41
	ϕ_i2h = 5.41
	ϕ_i2d = 14.26

	i2c = pdf.( Ref(NegativeBinomial2(μ_i2c, ϕ_i2c)), 1:num_i2c )
	i2h = pdf.( Ref(NegativeBinomial2(μ_i2h, ϕ_i2h)), 1:num_i2h )
	i2d = pdf.( Ref(NegativeBinomial2(μ_i2d, ϕ_i2d)), 1:num_i2d )

	i2c /= sum(i2c)
	i2h /= sum(i2h)
	i2d /= sum(i2d)

	### day of week effect
	weekdayeffect_simplex ~ Dirichlet([7,7,7,7,7,3,1])
	num_days = 7
	weekdayeffect = TV(undef, num_days)
	for i in 1:num_days
		weekdayeffect[i] = weekdayeffect_simplex[i] * num_days
	end

	### holiday effect
	holidayeffect ~ Beta(1,1)

    ### 4.2) derive observables from infections
	ifr       ~ truncated(Normal(6/1000, 3/1000), 1/1000, 10/1000)
	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
    iar       ~ Beta(1,10)

    expected_daily_cases     = TV(undef, num_time_steps)
    expected_daily_hospit    = TV(undef, num_time_steps)
    expected_daily_deaths    = TV(undef, num_time_steps)
	expected_seropositive    = TV(undef, num_time_steps)

	expected_daily_cases[1]  = 1e-15 * newly_infected[1]
    expected_daily_hospit[1] = 1e-15 * newly_infected[1]
    expected_daily_deaths[1] = 1e-15 * newly_infected[1]
	for t = 2:num_time_steps
		periodiceffect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
		expected_daily_cases[t]  = iar * sum(newly_infected[τ] * i2c[t - τ] for τ = (t - 1):-1:max(t-num_i2c,1)) * periodiceffect
        expected_daily_hospit[t] = ihr * sum(newly_infected[τ] * i2h[t - τ] for τ = (t - 1):-1:max(t-num_i2h,1))
        expected_daily_deaths[t] = ifr * sum(newly_infected[τ] * i2d[t - τ] for τ = (t - 1):-1:max(t-num_i2d,1))
	end

	serodelay = 16 # 14d self deferral + 16d seroconversion (see Erikstrup et al. https://doi.org/10.1093/cid/ciaa849)
	expected_seropositive[1:serodelay]    .= zero(V)
	expected_seropositive[serodelay+1:end] = cumulative_infected[1:end-serodelay]

    ########### 4.) compare model to observations
	## 4.1) observe cases
	ϕ_c  ~ truncated(Normal(25, 10), 0, Inf)
	ts_c = cases_start:num_obs
	μs_c = expected_daily_cases[ts_c]
	!all( 0 .< μs_c .< population ) && ( @error "expected_daily_cases < 0"; Turing.@addlogprob! -Inf; return)
	cases[ts_c] ~ arraydist(NegativeBinomial2.(μs_c, ϕ_c))

    ## 4.3) observe deaths
	ϕ_d   ~ truncated(Normal(20, 10), 0, Inf)
    ts_d  = epidemic_start:cases_start-1
	μs_d  = expected_daily_deaths[ts_d]
	!all( 0 .< μs_d .< population ) && (@error "expected_daily_deaths"; Turing.@addlogprob! -Inf; return)
	deaths[ts_d] ~ arraydist(NegativeBinomial2.(μs_d, ϕ_d))

    ## 4.3) seroprevalence study
    select_seroprev = 2:4
	seroprev_σ ~ arraydist( InverseGamma2.(seroprev_std[ select_seroprev ], Ref(.1)) )
    seroprev_μ = expected_seropositive[ seroprev_idx[ select_seroprev ] ]
	seroprev_mean[select_seroprev] ~ MvNormal(seroprev_μ, seroprev_σ) # select only early observations

    return (
		expected_daily_cases  = expected_daily_cases,
        expected_daily_deaths = expected_daily_deaths,
        expected_daily_hospit = expected_daily_hospit,
        Rt = Rt,
		expected_seropositive = expected_seropositive,
		iar = iar,
    )
end

## ============================================================================
#                           parametetric model with hospit
# =============================================================================
@model function model_parametric_hospit(
    num_impute,        # [Int] num. of days for which to impute infections
    num_total_days,    # [Int] days of observed data + num. of days to forecast
    cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
    deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
    π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    epidemic_start,    # [AbstractVector{<:Int}]
    population,        # [AbstractVector{<:Real}]
    serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
	num_iar_steps,         # [Int] number of weeks (21)
    iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
    lockdown,         # [Int] number of weeks (21)
    num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
	weekday,
	holiday,
	covariates,
	covariates_start,
	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    hospit,
	seroprev_mean,
    seroprev_std,
    seroprev_idx,
	π4,
    predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
    ::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
	invlink = log,#KLogistic(V(3))
	link    = exp,#KLogit(V(k))
) where {TV, V}
    ############ 0.) some definitions
	num_obs        = length(cases)
    num_Rt_steps   = covariates_start-lockdown-1
	num_sero       = length(seroprev_mean)
	num_covariates = size(covariates, 2)
    cases_start    = num_obs-num_case_obs
	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
    num_time_steps = predict ? num_total_days : num_obs

    ############# 2.) time varying reproduction number
	R0        ~ truncated(Normal(3., 1.), 1., 5.)
	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
	σ_rt      ~ truncated(Normal(0.05, .03), 0, .3)
	latent_Rt ~ RandomWalk(num_Rt_steps, σ_rt, invlink(R1))
	α         ~ filldist(Exponential(0.2), num_covariates)

    Rt                                = TV(undef, num_time_steps)
    Rt[1:lockdown]                   .= R0
    Rt[lockdown+1:covariates_start-1] = link.(latent_Rt)
    Rt[covariates_start:num_obs]      = link.(
        latent_Rt[end] .+ covariates * α
    )
    if predict #if prediction
		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
	end

    ############ 3.) infection dynamics
	## 3.1) initialize serial interval
	num_si   = 15
	si       = serial_intervals
	# μ_si   ~ truncated(Normal(5.06, 0.33), 0, Inf)
	# σ_si   = 2.11
	# si     = pdf.( Ref(GammaMeanStd(μ_si, σ_si)), 1:num_si )
	# si    /= sum(si)

	## 3.2) seed epidemic
	τ ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
	T = typeof(τ)
    y ~ truncated(Exponential(τ),T(0),T(1000))

	newly_infected         = TV(undef, num_time_steps)
	cumulative_infected    = TV(undef, num_time_steps)
	newly_infected[1]      = y
	cumulative_infected[1] = zero(V)
	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
	end

    ## 3.3) run disease dynamics
	for t = (num_impute + 1):num_time_steps
		# Update cumulatively infected
		cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		# Adjusts for portion of population that is susceptible
		St = max(population - cumulative_infected[t], 0) / population
		# effective number of infectious individuals
		effectively_infectious = sum(newly_infected[τ] * si[t - τ] for τ = (t - 1):-1:max(t-num_si,1))
		# number of new infections (unobserved)
		newly_infected[t] = St * Rt[t] * effectively_infectious
	end

	########### 4.) derive observables
	## 4.1) initialize delay distributions
	num_i2c = 40
	num_i2h = 40
	num_i2d = 60

	μ_i2c ~ truncated(Normal(4.5, 1.), 3, 7)
    μ_i2h ~ truncated(Normal(12., 3.), 9, 16)
    μ_i2d ~ truncated(Normal(20., 4.), 10, 30)

	ϕ_i2c = 5.41
	ϕ_i2h = 5.41
	ϕ_i2d = 14.26

	i2c = pdf.( Ref(NegativeBinomial2(μ_i2c, ϕ_i2c)), 1:num_i2c )
	i2h = pdf.( Ref(NegativeBinomial2(μ_i2h, ϕ_i2h)), 1:num_i2h )
	i2d = pdf.( Ref(NegativeBinomial2(μ_i2d, ϕ_i2d)), 1:num_i2d )

	i2c /= sum(i2c)
	i2h /= sum(i2h)
	i2d /= sum(i2d)

	### day of week effect
	weekdayeffect_simplex ~ Dirichlet([7,7,7,7,7,3,1])
	num_days = 7
	weekdayeffect = TV(undef, num_days)
	for i in 1:num_days
		weekdayeffect[i] = weekdayeffect_simplex[i] * num_days
	end

	### holiday effect
	holidayeffect ~ Beta(1,1)

    ### 4.2) derive observables from infections
	ifr       ~ truncated(Normal(6/1000, 3/1000), 1/1000, 10/1000)
	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
    iar       ~ Beta(1,10)

    expected_daily_cases     = TV(undef, num_time_steps)
    expected_daily_hospit    = TV(undef, num_time_steps)
    expected_daily_deaths    = TV(undef, num_time_steps)
	expected_seropositive    = TV(undef, num_time_steps)

	expected_daily_cases[1]  = 1e-15 * newly_infected[1]
    expected_daily_hospit[1] = 1e-15 * newly_infected[1]
    expected_daily_deaths[1] = 1e-15 * newly_infected[1]
	for t = 2:num_time_steps
		periodiceffect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
		expected_daily_cases[t]  = iar * sum(newly_infected[τ] * i2c[t - τ] for τ = (t - 1):-1:max(t-num_i2c,1)) * periodiceffect
        expected_daily_hospit[t] = ihr * sum(newly_infected[τ] * i2h[t - τ] for τ = (t - 1):-1:max(t-num_i2h,1))
        expected_daily_deaths[t] = ifr * sum(newly_infected[τ] * i2d[t - τ] for τ = (t - 1):-1:max(t-num_i2d,1))
	end

	serodelay = 16 # 14d self deferral + 16d seroconversion (see Erikstrup et al. https://doi.org/10.1093/cid/ciaa849)
	expected_seropositive[1:serodelay]    .= zero(V)
	expected_seropositive[serodelay+1:end] = cumulative_infected[1:end-serodelay]

    ########### 4.) compare model to observations
	# 4.2) observe hospitalizations
	ϕ_h   ~ truncated(Normal(25, 10), 0, Inf)
	ts_h  = epidemic_start:num_obs
	μs_h  = expected_daily_hospit[ts_h]
	!all( 0 .< μs_h .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
	hospit[ts_h] ~ arraydist(NegativeBinomial2.(μs_h, ϕ_h))

    ## 4.3) seroprevalence study
    select_seroprev = 2:4
	seroprev_σ ~ arraydist( InverseGamma2.(seroprev_std[ select_seroprev ], Ref(.1)) )
    seroprev_μ = expected_seropositive[ seroprev_idx[ select_seroprev ] ]
	seroprev_mean[select_seroprev] ~ MvNormal(seroprev_μ, seroprev_σ) # select only early observations

    return (
		expected_daily_cases  = expected_daily_cases,
        expected_daily_deaths = expected_daily_deaths,
        expected_daily_hospit = expected_daily_hospit,
        Rt = Rt,
		expected_seropositive = expected_seropositive,
		iar = iar,
    )
end

## ==========================================================================
#                       parametric model with deaths
# ===========================================================================
@model function model_parametric_deaths(
    num_impute,        # [Int] num. of days for which to impute infections
    num_total_days,    # [Int] days of observed data + num. of days to forecast
    cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
    deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
    π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    epidemic_start,    # [AbstractVector{<:Int}]
    population,        # [AbstractVector{<:Real}]
    serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
	num_iar_steps,         # [Int] number of weeks (21)
    iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
    lockdown,         # [Int] number of weeks (21)
    num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
	weekday,
	holiday,
	covariates,
	covariates_start,
	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    hospit,
	seroprev_mean,
    seroprev_std,
    seroprev_idx,
	π4,
    predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
    ::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
	invlink = log,#KLogistic(V(3))
	link    = exp,#KLogit(V(k))
) where {TV, V}
    ############ 0.) some definitions
	num_obs        = length(cases)
    num_Rt_steps   = covariates_start-lockdown-1
	num_sero       = length(seroprev_mean)
	num_covariates = size(covariates, 2)
    cases_start    = num_obs-num_case_obs
	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
    num_time_steps = predict ? num_total_days : num_obs

    ############# 2.) time varying reproduction number
	R0        ~ truncated(Normal(3., 1.), 1., 5.)
	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
	σ_rt      ~ truncated(Normal(0.05, .03), 0, .3)
	latent_Rt ~ RandomWalk(num_Rt_steps, σ_rt, invlink(R1))
	α         ~ filldist(Exponential(0.2), num_covariates)

    Rt                                = TV(undef, num_time_steps)
    Rt[1:lockdown]                   .= R0
    Rt[lockdown+1:covariates_start-1] = link.(latent_Rt)
    Rt[covariates_start:num_obs]      = link.(
        latent_Rt[end] .+ covariates * α
    )
    if predict #if prediction
		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
	end

    ############ 3.) infection dynamics
	## 3.1) initialize serial interval
	num_si   = 15
	si       = serial_intervals
	# μ_si   ~ truncated(Normal(5.06, 0.33), 0, Inf)
	# σ_si   = 2.11
	# si     = pdf.( Ref(GammaMeanStd(μ_si, σ_si)), 1:num_si )
	# si    /= sum(si)

	## 3.2) seed epidemic
	τ ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
	T = typeof(τ)
    y ~ truncated(Exponential(τ),T(0),T(1000))

	newly_infected         = TV(undef, num_time_steps)
	cumulative_infected    = TV(undef, num_time_steps)
	newly_infected[1]      = y
	cumulative_infected[1] = zero(V)
	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
	end

    ## 3.3) run disease dynamics
	for t = (num_impute + 1):num_time_steps
		# Update cumulatively infected
		cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		# Adjusts for portion of population that is susceptible
		St = max(population - cumulative_infected[t], 0) / population
		# effective number of infectious individuals
		effectively_infectious = sum(newly_infected[τ] * si[t - τ] for τ = (t - 1):-1:max(t-num_si,1))
		# number of new infections (unobserved)
		newly_infected[t] = St * Rt[t] * effectively_infectious
	end

	########### 4.) derive observables
	## 4.1) initialize delay distributions
	num_i2c = 40
	num_i2h = 40
	num_i2d = 60

	μ_i2c ~ truncated(Normal(4.5, 1.), 3, 7)
    μ_i2h ~ truncated(Normal(12., 3.), 9, 16)
    μ_i2d ~ truncated(Normal(20., 4.), 10, 30)

	ϕ_i2c = 5.41
	ϕ_i2h = 5.41
	ϕ_i2d = 14.26

	i2c = pdf.( Ref(NegativeBinomial2(μ_i2c, ϕ_i2c)), 1:num_i2c )
	i2h = pdf.( Ref(NegativeBinomial2(μ_i2h, ϕ_i2h)), 1:num_i2h )
	i2d = pdf.( Ref(NegativeBinomial2(μ_i2d, ϕ_i2d)), 1:num_i2d )

	i2c /= sum(i2c)
	i2h /= sum(i2h)
	i2d /= sum(i2d)

	### day of week effect
	weekdayeffect_simplex ~ Dirichlet([7,7,7,7,7,3,1])
	num_days = 7
	weekdayeffect = TV(undef, num_days)
	for i in 1:num_days
		weekdayeffect[i] = weekdayeffect_simplex[i] * num_days
	end

	### holiday effect
	holidayeffect ~ Beta(1,1)

    ### 4.2) derive observables from infections
	ifr       ~ truncated(Normal(6/1000, 3/1000), 1/1000, 10/1000)
	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
    iar       ~ Beta(1,10)

    expected_daily_cases     = TV(undef, num_time_steps)
    expected_daily_hospit    = TV(undef, num_time_steps)
    expected_daily_deaths    = TV(undef, num_time_steps)
	expected_seropositive    = TV(undef, num_time_steps)

	expected_daily_cases[1]  = 1e-15 * newly_infected[1]
    expected_daily_hospit[1] = 1e-15 * newly_infected[1]
    expected_daily_deaths[1] = 1e-15 * newly_infected[1]
	for t = 2:num_time_steps
		periodiceffect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
		expected_daily_cases[t]  = iar * sum(newly_infected[τ] * i2c[t - τ] for τ = (t - 1):-1:max(t-num_i2c,1)) * periodiceffect
        expected_daily_hospit[t] = ihr * sum(newly_infected[τ] * i2h[t - τ] for τ = (t - 1):-1:max(t-num_i2h,1))
        expected_daily_deaths[t] = ifr * sum(newly_infected[τ] * i2d[t - τ] for τ = (t - 1):-1:max(t-num_i2d,1))
	end

	serodelay = 16 # 14d self deferral + 16d seroconversion (see Erikstrup et al. https://doi.org/10.1093/cid/ciaa849)
	expected_seropositive[1:serodelay]    .= zero(V)
	expected_seropositive[serodelay+1:end] = cumulative_infected[1:end-serodelay]

    ########### 4.) compare model to observations

    ## 4.3) observe deaths
	ϕ_d   ~ truncated(Normal(20, 10), 0, Inf)
    ts_d  = epidemic_start:num_obs
	μs_d  = expected_daily_deaths[ts_d]
	!all( 0 .< μs_d .< population ) && (@error "expected_daily_deaths"; Turing.@addlogprob! -Inf; return)
	deaths[ts_d] ~ arraydist(NegativeBinomial2.(μs_d, ϕ_d))

    ## 4.3) seroprevalence study
    select_seroprev = 2:4
	seroprev_σ ~ arraydist( InverseGamma2.(seroprev_std[ select_seroprev ], Ref(.1)) )
    seroprev_μ = expected_seropositive[ seroprev_idx[ select_seroprev ] ]
	seroprev_mean[select_seroprev] ~ MvNormal(seroprev_μ, seroprev_σ) # select only early observations

    return (
		expected_daily_cases  = expected_daily_cases,
        expected_daily_deaths = expected_daily_deaths,
        expected_daily_hospit = expected_daily_hospit,
        Rt = Rt,
		expected_seropositive = expected_seropositive,
		iar = iar,
    )
end


module MultDep

using Hecke
import Base.*

function multiplicative_group_mod_units_fac_elem(A::Array{nf_elem, 1})
  k = parent(A[1])
  @assert all(i->parent(i) === k, A)
  cp = coprime_base(A)
  M = sparse_matrix(FlintZZ)
  for a = A
    T = Tuple{Int, fmpz}[]
    for i = 1:length(cp)
      p = cp[i]
      v = valuation(a, p)
      if v != 0
        push!(T, (i, fmpz(v)))
      end
#      isone(I) && break
    end
    push!(M, sparse_row(FlintZZ, T))
  end
  h, t = Hecke.hnf_kannan_bachem(M, Val{true}, truncate = true)
  return h, t, cp
end

function *(O1::NfAbsOrd, O2::NfAbsOrd)
  k = nf(O1)
  @assert k === nf(O2)
  b1 = basis(O1, k)
  n = length(b1)
  b2 = basis(O2, k)
  p = [x*y for (x,y) in Base.Iterators.ProductIterator((b1, b2))]
  d = reduce(lcm, [denominator(x) for x = p])
  M = zero_matrix(FlintZZ, n*n, n)
  z = fmpz()
  for i = 1:n*n
    a = p[i]*d
    Hecke.elem_to_mat_row!(M, i, z, a)
  end
  h = hnf(M)
  b = nf_elem[]
  for i=1:n
    push!(b, Hecke.elem_from_mat_row(k, h, i, d))
  end
  return Hecke.Order(k, b)
end

mutable struct GeIdeal
  a::NfAbsOrdIdl
  function GeIdeal(a::NfAbsOrdIdl)
    o =order(a)
    if o.ismaximal == 1
      return new(a)
    end
    #keep track of known maximality and use this to speed things up
    o = ring_of_multipliers(a)
    if o === order(a)
      return new(a)
    else
      return new(extend(a, o))
    end
  end
  function GeIdeal(a::nf_elem)
    o = Hecke.any_order(parent(a))
    d = denominator(a, o)
    return new(o(a*d)*o), new(o(d)*o)
  end
end

import Hecke.gcd, Hecke.isone, Hecke.*, Hecke.gcd_into!, Hecke.copy, Hecke.divexact,
       Hecke.isunit, Hecke.coprime_base, Hecke.valuation

function make_compatible!(a::GeIdeal, b::GeIdeal)
  o1 = order(a.a)
  o2 = order(b.a)
  if o1 === o2
    return
  end
  o = o1*o2
  a.a = extend(a.a, o)
  b.a = extend(b.a, o)
end

#isunit, divexact, gcd, isone, *, gcd_into!, copy
isone(a::GeIdeal) = isone(a.a)
isunit(a::GeIdeal) = isone(a)
copy(a::GeIdeal) = GeIdeal(a.a)

function gcd(a::GeIdeal, b::GeIdeal)
  make_compatible!(a, b)
  return GeIdeal(a.a + b.a)
end

function gcd_into!(a::GeIdeal, b::GeIdeal, c::GeIdeal)
  make_compatible!(b, c)
  a.a = b.a + c.a
  return a
end

function *(a::GeIdeal, b::GeIdeal)
  make_compatible!(a, b)
  return GeIdeal(a.a * b.a)
end

function divexact(a::GeIdeal, b::GeIdeal)
  make_compatible!(a, b)
  return GeIdeal(divexact(a.a, b.a))
end

function coprime_base(A::Array{nf_elem, 1})
  c = Array{GeIdeal, 1}()
  for a = A
    n,d = GeIdeal(a)
    isone(n) || push!(c, n)
    isone(d) || push!(c, d)
  end
  return coprime_base(c)
end

function valuation(a::nf_elem, p::GeIdeal)
  return valuation(a, p.a)
end

end


"""
Hartree term: for a decaying potential V the energy would be

1/2 ∫ρ(x)ρ(y)V(x-y) dxdy

with the integral on x in the unit cell and of y in the whole space.
For the Coulomb potential with periodic boundary conditions, this is rather

1/2 ∫ρ(x)ρ(y) G(x-y) dx dy

where G is the Green's function of the periodic Laplacian with zero
mean (-Δ G = sum_{R} 4π δ_R, integral of G zero on a unit cell).
"""
struct Hartree
    scaling_factor::Real  # to scale by an arbitrary factor (useful for exploration)
end
Hartree(; scaling_factor=1) = Hartree(scaling_factor)
(hartree::Hartree)(basis) = TermHartree(basis, hartree.scaling_factor)

struct TermHartree <: Term
    basis::PlaneWaveBasis
    scaling_factor::Real  # scaling factor, absorbed into poisson_green_coeffs
    # Fourier coefficients of the Green's function of the periodic Poisson equation
    poisson_green_coeffs
end
function TermHartree(basis::PlaneWaveBasis{T}, scaling_factor) where T
    # Solving the Poisson equation ΔV = -4π ρ in Fourier space
    # is multiplying elementwise by 4π / |G|^2.
    poisson_green_coeffs = 4T(π) ./ [sum(abs2, G)
                                     for G in G_vectors_cart(basis)]

    # Zero the DC component (i.e. assume a compensating charge background)
    poisson_green_coeffs[1] = 0
    TermHartree(basis, scaling_factor, scaling_factor .* poisson_green_coeffs)
end

@timing "ene_ops: hartree" function ene_ops(term::TermHartree, ψ, occ; ρ, kwargs...)
    basis = term.basis
    T = eltype(basis)
    pot_fourier = term.poisson_green_coeffs .* ρ.fourier
    potential = real(G_to_r(basis, pot_fourier))  # TODO optimize this
    E = real(dot(pot_fourier, ρ.fourier) / 2)

    ops = [RealSpaceMultiplication(basis, kpoint, potential) for kpoint in basis.kpoints]
    (E=E, ops=ops)
end

function compute_kernel(term::TermHartree; kwargs...)
    @assert term.basis.model.spin_polarization in (:none, :spinless, :collinear)
    vc_G = term.poisson_green_coeffs
    # Note that `real` here: if omitted, will result in high-frequency noise of even FFT grids
    K = real(G_to_r_matrix(term.basis) * Diagonal(vec(vc_G)) * r_to_G_matrix(term.basis))

    # Hartree kernel is independent of spin, so to apply it to (ρtot, ρspin)^T
    # and obtain the same contribution to Vα and Vβ the operator has the block structure
    #     ( K 0 )
    #     ( K 0 )
    n_spin = term.basis.model.n_spin_components
    n_spin == 1 ? K : [K 0I; K 0I]
end

function apply_kernel(term::TermHartree, dρ::RealFourierArray, dρspin=nothing; kwargs...)
    @assert term.basis.model.spin_polarization in (:none, :spinless, :collinear)
    kernel = from_fourier(dρ.basis, term.poisson_green_coeffs .* dρ.fourier)
    n_spin = term.basis.model.n_spin_components
    n_spin == 1 ? (kernel, ) : (kernel, kernel)  # Hartree term does not care about spin
end


"""
Module: SymbolicDiff (Symbolic Operation for Arithmetic)
"""

"""
seval(f, (dvar1, dver2), env, cache)
Return the second derivative of expr f with respect to dvar1 and dver2
"""

function seval(f, dvar::Tuple{Symbol,Symbol})
    seval(f, dvar, globalenv, SymbolicCache())
end

function seval(f, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv)
    seval(f, dvar, env, SymbolicCache())
end

function seval(f, dvar::Tuple{Symbol,Symbol}, cache::SymbolicCache)
    seval(f, dvar, globalenv, cache)
end

###

function seval(f, dvar::Tuple{SymbolicVariable{Tv},SymbolicVariable{Tv}}) where Tv
    seval(f, (dvar[1].var, dvar[2].var))
end

function seval(f, dvar::Tuple{SymbolicVariable{Tv},SymbolicVariable{Tv}}, env::SymbolicEnv) where Tv
    seval(f, (dvar[1].var, dvar[2].var), env)
end

function seval(f, dvar::Tuple{SymbolicVariable{Tv},SymbolicVariable{Tv}}, cache::SymbolicCache) where Tv
    seval(f, (dvar[1].var, dvar[2].var), cache)
end

function seval(f, dvar::Tuple{SymbolicVariable{Tv},SymbolicVariable{Tv}}, env::SymbolicEnv, cache::SymbolicCache) where Tv
        seval(f, (dvar[1].var, dvar[2].var), env, cache)
end

###

function seval(f::SymbolicValue{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    0
end

function seval(f::SymbolicVariable{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    0
end

function seval(f::AbstractNumberSymbolic{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    (dvar[1] in f.params) || (dvar[2] in f.params) || return Tv(0)
    get(cache, (f,dvar)) do
        retval = _eval(Val(f.op), f, dvar, env, cache)
        cache[(f,dvar)] = retval
    end
end

"""
_eval(::Val{xx}, dvar, f, env, cache)

Dispached function to evaluate the second derivative of f
"""

function _eval(::Val{:+}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    args = [seval(x, dvar, env, cache) for x = f.args]
    +(args...)
end

function _eval(::Val{:-}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    args = [seval(x, dvar, env, cache) for x = f.args]
    -(args...)
end

function _eval(::Val{:*}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    args = [seval(x, env, cache) for x = f.args]
    dargs_a = [seval(x, dvar[1], env, cache) for x = f.args]
    dargs_b = [seval(x, dvar[2], env, cache) for x = f.args]
    dargs_ab = [seval(x, dvar, env, cache) for x = f.args]

    ret = dargs_ab[1]
    s = args[1]
    s_a = dargs_a[1]
    s_b = dargs_b[1]
    for i = 2:length(args)
        ret *= args[i]
        ret += s * dargs_ab[i]
        ret += s_a * dargs_b[i]
        ret += s_b * dargs_a[i]
        (i == length(args)) && break
        s_a *= args[i]
        s_b *= args[i]
        s_a += s * dargs_a[i]
        s_b += s * dargs_b[i]
        s *= args[i]
    end
    ret
end

function _eval(::Val{:/}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x,y = [seval(x, env, cache) for x = f.args]
    dx_a,dy_a = [seval(x, dvar[1], env, cache) for x = f.args]
    dx_b,dy_b = [seval(x, dvar[2], env, cache) for x = f.args]
    dx_ab,dy_ab = [seval(x, dvar, env, cache) for x = f.args]
    ((dx_ab * y - dx_a * dy_b  - dx_b * dy_a - x * dy_ab) * y + 2 * x * dy_a * dy_b) / y^3
end

function _eval(::Val{:^}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x,y = [seval(x, env, cache) for x = f.args]
    dx_a,dy_a = [seval(x, dvar[1], env, cache) for x = f.args]
    dx_b,dy_b = [seval(x, dvar[2], env, cache) for x = f.args]
    dx_ab,dy_ab = [seval(x, dvar, env, cache) for x = f.args]
    f = x^y
    f_a = f * (dy_a * x * log(x) + y * dx_a) / x
    f_b = f * (dy_b * x * log(x) + y * dx_b) / x
    f_ab = (f_b * (dy_a * x * log(x) + y * dx_a) - f_a * dx_b + f * (dy_ab * x * log(x) + dy_a * dx_b * (1 + log(x)) + dx_a * dy_b + y * dx_ab)) / x
    f_ab
end

function _eval(::Val{:exp}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x, = [seval(x, env, cache) for x = f.args]
    dx_a, = [seval(x, dvar[1], env, cache) for x = f.args]
    dx_b, = [seval(x, dvar[2], env, cache) for x = f.args]
    dx_ab, = [seval(x, dvar, env, cache) for x = f.args]
    exp(x) * (dx_b * dx_a + dx_ab)
end

function _eval(::Val{:log}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x, = [seval(x, env, cache) for x = f.args]
    dx_a, = [seval(x, dvar[1], env, cache) for x = f.args]
    dx_b, = [seval(x, dvar[2], env, cache) for x = f.args]
    dx_ab, = [seval(x, dvar, env, cache) for x = f.args]
    (dx_ab * x - dx_a * dx_b) / x^2
end

function _eval(::Val{:sqrt}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x, = [seval(x, env, cache) for x = f.args]
    dx_a, = [seval(x, dvar[1], env, cache) for x = f.args]
    dx_b, = [seval(x, dvar[2], env, cache) for x = f.args]
    dx_ab, = [seval(x, dvar, env, cache) for x = f.args]
    sqrt(x) * (dx_ab * 2 - dx_a * dx_b) / (4*x)
end

function _eval(::Val{:sum}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    dx_ab, = [seval(x, dvar, env, cache) for x = f.args]
    sum(dx_ab)
end

function _eval(::Val{:dot}, f::SymbolicExpression{Tv}, dvar::Tuple{Symbol,Symbol}, env::SymbolicEnv, cache::SymbolicCache)::Tv where Tv
    x,y = [seval(x, env, cache) for x = f.args]
    dx_a,dy_a = [seval(x, dvar[1], env, cache) for x = f.args]
    dx_b,dy_b = [seval(x, dvar[2], env, cache) for x = f.args]
    dx_ab,dy_ab = [seval(x, dvar, env, cache) for x = f.args]
    dot(x,dy_ab) + dot(dx_b,dy_a) + dot(dx_a,dy_b) + dot(dx_ab,y)
end


using QuDynamics.QuTiP
begin

# Defining the system - Hamiltonian, initial state, time steps
# Hamiltonian of the system
H = sigmax

# Initial state of the system
init_state = statevec(1, FiniteBasis(2))
init_state_dm = init_state*init_state'

# Time steps
tlist = 0.:0.1:2*pi

# Collapse operators
c_ops = [lowerop(2)]

# Operators whose expectation values are to computed
# in the evolved states
e_ops = [sqrt(2)*lowerop(2) + sqrt(3)*raiseop(2), raiseop(2), lowerop(2)]

# Test set involving Schrodinger like equations
se = QuSchrodingerEq(sigmax)

# Evaluating states using the general iterator version
evolved_states_se = Array(QuBase.AbstractQuArray, length(tlist)-1)
i = 1
for (t, psi) in QuStateEvolution(sigmax, init_state, tlist, QuExpmV())
    evolved_states_se[i] = psi
    i = i + 1
end

ind = rand(1:(length(tlist)-1))

# Evaluating using sesolve by passing Hamiltonian and initial state vector
schro_direct = sesolve(sigmax, init_state, tlist)

# Evaluating using sesolve by passing Schrodinger type equation
schro_eq = sesolve(se, init_state, tlist)

# Tests comparing the states from all the above constructs
@assert schro_direct[2][ind] == evolved_states_se[ind]
@assert schro_eq[2][ind] == evolved_states_se[ind]
@assert schro_direct[2][ind] == schro_eq[2][ind]

# Evaluating the expectation values
schro_direct_expect = sesolve(sigmax, init_state, tlist, e_ops)[2]
schro_eq_expect = sesolve(se, init_state, tlist, e_ops)[2]

ind_1 = rand(1:(length(e_ops)))
ind_2 = rand(1:length(tlist)-1)

@assert schro_direct_expect[ind_2, ind_1].im == schro_eq_expect[ind_2, ind_1].im

# Test set involving LiouvillevonNeumann like equations
lvne = QuLiouvillevonNeumannEq(sigmax)

# Evaluating states using the general iterator version
evolved_states_lvne = Array(QuBase.AbstractQuArray, length(tlist)-1)
i = 1
for (t, psi) in QuStateEvolution(sigmax, init_state_dm, tlist, QuExpmV())
    evolved_states_lvne[i] = psi
    i = i + 1
end

ind=rand(1:(length(tlist)-1))

# Evaluating using sesolve by passing Hamiltonian and initial density matrix
lvne_direct = sesolve(sigmax, init_state_dm, tlist)

# Evaluating using sesolve by passing LiouvillevonNeumann type equation
lvne_eq = sesolve(lvne, init_state_dm, tlist)

# Tests comparing the states from all the above constructs
@assert lvne_direct[2][ind] == evolved_states_lvne[ind]
@assert lvne_eq[2][ind] == evolved_states_lvne[ind]
@assert lvne_direct[2][ind] == lvne_eq[2][ind]

# Evaluating the expectation values
expectation_lvne = QuDynamics.QuTiP.eops_helper(e_ops, evolved_states_lvne)

# Evaluating using sesolve
lvne_direct_expect = sesolve(sigmax, init_state_dm, tlist, e_ops)[2]
lvne_eq_expect = sesolve(lvne, init_state_dm, tlist, e_ops)[2]

ind_1 = rand(1:(length(e_ops)))
ind_2 = rand(1:length(tlist)-1)

@assert lvne_direct_expect[ind_2, ind_1].im == expectation_lvne[ind_2, ind_1].im
@assert lvne_eq_expect[ind_2, ind_1].im  == expectation_lvne[ind_2, ind_1].im
@assert lvne_direct_expect[ind_2, ind_1].im == lvne_eq_expect[ind_2, ind_1].im

# Test set involving Lindblad Master equations
lme = QuLindbladMasterEq(sigmax, c_ops)

# Evaluating states using the general iterator version
evolved_states_lme = Array(QuBase.AbstractQuArray, length(tlist)-1)
i = 1
for (t, psi) in QuStateEvolution(lme, init_state_dm, tlist, QuExpmV())
    evolved_states_lme[i] = psi
    i = i + 1
end

ind=rand(1:(length(tlist)-1))

# Evaluating using mesolve by passing Hamiltonian and initial density matrix
lme_direct = mesolve(sigmax, init_state_dm, tlist, c_ops)

# Evaluating using mesolve by passing Lindblad Master equation
lme_eq = mesolve(lme, init_state_dm, tlist)

# Tests comparing the states from all the above constructs
@assert lme_direct[2][ind] == evolved_states_lme[ind]
@assert lme_eq[2][ind] == evolved_states_lme[ind]
@assert lme_direct[2][ind] == lme_eq[2][ind]

# Evaluating the expectation values
expectation_lme = QuDynamics.QuTiP.eops_helper(e_ops, evolved_states_lme)

# Evaluating using mesolve
lme_direct_expect = mesolve(sigmax, init_state_dm, tlist, c_ops, e_ops)[2]
lme_eq_expect = mesolve(lme, init_state_dm, tlist, e_ops)[2]

ind_1 = rand(1:(length(e_ops)))
ind_2 = rand(1:length(tlist)-1)

@assert lme_direct_expect[ind_2, ind_1].im == expectation_lme[ind_2, ind_1].im
@assert lme_eq_expect[ind_2, ind_1].im == expectation_lme[ind_2, ind_1].im
@assert lme_direct_expect[ind_2, ind_1].im == lme_eq_expect[ind_2, ind_1].im

# Testing the expectation using various mcsolve constructs
mcs = mcsolve(H, init_state, tlist, c_ops)
mcs_lme = mcsolve(lme, init_state, tlist)
mcs_lme_expect = mcsolve(lme, init_state_dm, tlist, e_ops)
mcs_expect = mcsolve(H, init_state, tlist, c_ops,  e_ops)

@test_approx_eq_eps mcs[2][ind_2] mcs_lme[2][ind_2] 1e-1
@test_approx_eq_eps trace(mcs[2][ind_2]*e_ops[ind_1]) mcs_expect[2][ind_2, ind_1] 1e-1
@test_approx_eq_eps trace(mcs[2][ind_2]*e_ops[ind_1]) mcs_lme_expect[2][ind_2, ind_1] 1e-1
@test_approx_eq_eps mcs_expect[2][ind_2, ind_1] mcs_lme_expect[2][ind_2, ind_1] 1e-1
end