text:
"""
```
TwoAssetHANK{T} <: AbstractCTModel{T}
```

The `TwoAssetHANK` type defines the structure of the two-asset HANK model
originally written in MATLAB by Ben Moll and Se Hyoun Ahn.

### Fields

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

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

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

#### Inputs to Measurement and Equilibrium Condition Equations

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

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

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

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

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

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

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

* `pseudo_observables::OrderedDict{Symbol,Int}`: Maps each pseudo-observable to
  a row in the model's pseudo-measurement equation matrices.

#### Model Specifications and Settings

* `spec::String`: The model specification identifier, \"one_asset_hank\", cached
  here for filepath computation.

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

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

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

#### Other Fields

* `rng::MersenneTwister`: Random number generator. Can be is seeded to ensure
  reproducibility in algorithms that involve randomness (such as
  Metropolis-Hastings).

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

* `observable_mappings::OrderedDict{Symbol,Observable}`: A dictionary that
  stores data sources, series mnemonics, and transformations to/from model units.
  DSGE.jl will fetch data from the Federal Reserve Bank of
  St. Louis's FRED database; all other data must be downloaded by the
  user. See `load_data` and `Observable` for further details.

* `pseudo_observable_mappings::OrderedDict{Symbol,PseudoObservable}`: A
  dictionary that stores names and transformations to/from model units. See
  `PseudoObservable` for further details.
"""
mutable struct TwoAssetHANK{T} <: AbstractCTModel{T}
    parameters::ParameterVector{T}                          # vector of all time-invariant model parameters
    steady_state::ParameterVector{T}                        # model steady-state values
    keys::OrderedDict{Symbol,Int}                           # human-readable names for all the model
                                                            # parameters and steady-states

    endogenous_states::OrderedDict{Symbol,Vector{Int}}      # these fields used to create matrices in the
    exogenous_shocks::OrderedDict{Symbol,Int}               # measurement & equilibrium condition equations
    expected_shocks::OrderedDict{Symbol,Vector{Int}}        #
    equilibrium_conditions::OrderedDict{Symbol,Vector{Int}} #
    endogenous_states_augmented::OrderedDict{Symbol,Int}    #
    observables::OrderedDict{Symbol,Int}                    #
    pseudo_observables::OrderedDict{Symbol,Int}             #

    spec::String                                            # Model specification number (eg "m990")
    subspec::String                                         # Model subspecification (eg "ss0")
    settings::Dict{Symbol,Setting}                          # Settings/flags for computation
    test_settings::Dict{Symbol,Setting}                     # Settings/flags for testing mode
    rng::MersenneTwister                                    # Random number generator
    testing::Bool                                           # Whether we are in testing mode or not

    observable_mappings::OrderedDict{Symbol, Observable}
    pseudo_observable_mappings::OrderedDict{Symbol, PseudoObservable}
end

description(m::TwoAssetHANK) = "Julia implementation of the two-asset HANK model
originally written in MATLAB by Ben Moll and Se Hyoun Ahn."

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

Arguments:
`m:: TwoAssetHANK`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and equilibrium conditions.
"""
function init_model_indices!(m::TwoAssetHANK)

    n_v = get_setting(m, :n_v)
    n_g = get_setting(m, :n_g)
    n_p = get_setting(m, :n_p)
    n_Z = get_setting(m, :n_Z)

    # Endogenous states
    endogenous_states = collect([:value_function, :distribution, :K, :r_b, :Y, :C, :Z])

    # Exogenous shocks
    exogenous_shocks = collect([:mp_shock])

    # Expectations shocks
    expected_shocks = collect([:E_V, :E_π])

    # Equilibrium conditions
    equilibrium_conditions = collect([:value_function, :distribution, :K, :r_b, :Y, :C, :Z])

    # Additional states added after solving model
    # Lagged states and observables measurement error
    endogenous_states_augmented = []

    # Observables
    observables = keys(m.observable_mappings)

    # Pseudo-observables
    pseudo_observables = keys(m.pseudo_observable_mappings)

    # Assign dimensions
    m.endogenous_states[:value_function]  = collect(1:n_v)
    m.endogenous_states[:distribution]    = collect(n_v + 1:n_v + n_g)
    m.endogenous_states[:K]               = [n_v + n_g + 1]
    m.endogenous_states[:r_b]             = [n_v + n_g + 2]

    m.endogenous_states[:Y]               = [n_v + n_g + 3]
    m.endogenous_states[:C]               = [n_v + n_g + 4]

    m.endogenous_states[:Z]               = [n_v + n_g + n_p + 1]

    m.expected_shocks[:E_V]               = m.endogenous_states[:value_function]
    #m.expected_shocks[:E_π]               = m.endogenous_states[:inflation]

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

function TwoAssetHANK(subspec::String="ss0";
                      custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                      testing = false)

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

    # Initialize empty model
    m = TwoAssetHANK{Float64}(
            # Model parameters and steady state values # before 2nd arg was Vector{Float64}()
            Vector{AbstractParameter{Float64}}(), Vector{Float64}(),
            OrderedDict{Symbol,Int}(),

            # Model indices
            OrderedDict{Symbol,Vector{Int}}(), OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,Vector{Int}}(), OrderedDict{Symbol,Vector{Int}}(),
            OrderedDict{Symbol,Int}(),         OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,Int}(),

            spec,
            subspec,
            settings,
            test_settings,
            rng,
            testing,
            OrderedDict{Symbol,Observable}(),
            OrderedDict{Symbol,PseudoObservable}())

    # Initialize parameters
    init_parameters!(m)

    # Set settings
    model_settings!(m)
    default_test_settings!(m)
    adjust_parameters!(m)

    init_model_indices!(m)

    for custom_setting in values(custom_settings)
        m <= custom_setting
    end

    # Test to see if necessary values are defined, throw error otherwise
#=    try
        @assert typeof(get_setting(m, :reduce_state_vars)) == Bool
        @assert typeof(get_setting(m, :reduce_v)) == Bool
    catch
        error("Need to specify in settings whether to perform distribution
               and/or value function reduction")
    end

    if get_setting(m, :reduce_state_vars) || get_setting(m, :reduce_v)
        try
            @assert get_setting(m, :n_jump_vars) > 0
            @assert get_setting(m, :n_state_vars) > 0
            @assert get_setting(m, :n_state_vars_unreduce) >= 0
        catch
            error("Need to enter values in settings for n_jump_vars, n_state_vars,
                   and/or n_state_vars_unreduce in Settings. These indicate the
                   number of jump and aggregate state variables.")
        end
    end

    if get_setting(m, :reduce_state_vars)
        try
            @assert get_setting(m, :krylov_dim) > 0
        catch
            error("Need to specify in settings desired dimension of
                   Krylov subspace, krylov_dim.")
        end
    end

    if get_setting(m, :reduce_v)
        try
            @assert get_setting(m, :n_prior) > 0 && typeof(get_setting(m, :n_prior)) == Int
            @assert get_setting(m, :n_post) > 0 && typeof(get_setting(m, :n_post)) == Int
            @assert in(:knots_dict, keys(m.settings))
            @assert in(:spline_grid, keys(m.settings))
        catch
            error("Need to specify in settings n_prior, n_post, knots_dict
                   and/or spline_grid for value function reduction.
                   See reduction.jl for their uses.")
        end
    end
=#
    # Compute steady state
#    steadystate!(m)
    return m
end

"""
```
init_parameters!(m::TwoAssetHANK)
```

Initializes the model's parameters, as well as empty values for the steady-state
parameters (in preparation for `steadystate!(m)` being called to initialize
those).
"""
function init_parameters!(m::TwoAssetHANK)
    # Preferences
    m <= parameter(:coefrra, 1.0, fixed=true,
                   description="Relative risk aversion",
                   tex_label="coefrra")

    m <= parameter(:rrho, 0.014526858925819, fixed=true,
                   description="Discount factor (quarterly)",
                   tex_label="rrho")

    m <= parameter(:ddeath, 1.0 / (4.0 * 45.0), fixed=true,
                   description="Death rate (quarterly)",
                   tex_label="ddeath")

    m <= parameter(:ggamma, 1.0, fixed=true,
                   description="Risk aversion",
                   tex_label="ggama")

    # Technology
    m <= parameter(:aalpha, 0.4, fixed=true,
                   description="Capital share",
                   tex_label="aalpha")

    m <= parameter(:nnu_aggZ, 0.05, fixed=true,
                   description="nnu_aggZ",
                   tex_label="nnu_aggZ")

    m <= parameter(:ssigma_aggZ, 0.007 / (1.0 - m[:aalpha]), fixed=true,
                   description="ssigma_aggZ",
                   tex_label="ssigma_aggZ")

    m <= parameter(:ddelta, 0.075 / 4.0 , fixed=true,
                   description="Depreciation rate (quarterly)",
                   tex_label="ddelta")

    # Portfolio adjustment cost function
    m <= parameter(:chi0, 0.008198815318426, fixed=true,
                   description="Portfolio adjustment cost function",
                   tex_label="chi0")

    m <= parameter(:chi1, 0.072974260992256, fixed=true,
                   description="Portfolio adjustment cost function",
                   tex_label="chi1")

    m <= parameter(:chi2, 1.3583017016896, fixed=true,
                   description="Portfolio adjustment cost function",
                   tex_label="chi2")

    m <= parameter(:a_lb, 0.02, fixed=true,
                   description="Portfolio adjustment cost function",
                   tex_label="a_lb")

    # Tax code
    m <= parameter(:trans, 0.1, fixed=true,
                   description="Transfer to households",
                   tex_label="trans")

    m <= parameter(:tau_I, 0.30, fixed=true,
                   description="Tax rate on wage income",
                   tex_label="tau_I")

    # Deposit share
    m <= parameter(:xxi, 0.0, fixed=true,
                   description="Deposit share",
                   tex_label="xxi")

    # Perfect annuity markets
    m <= parameter(:pam, 1.0, fixed=true,
                   description="Perfect annuity markets",
                   tex_label="pam")

    # Liquid return
    m <= parameter(:r_b_SS, 0.005, fixed=true,
                   description="Steady-state savings rate, liquid assets",
                   tex_label="r_b_SS")
    m <= parameter(:borrwedge_SS, 0.020137162949122, fixed=true,
                   description="Steady-state borrowing wedge, liquid assets",
                   tex_label="borrwedge_SS")
    m <= parameter(:r_b_borr_SS, m[:r_b_SS] + m[:borrwedge_SS], fixed=true,
                   description="Steady-state borrowing rate, liquid assets",
                   tex_label="r_b_borr_SS")
    m <= parameter(:pphi, 2.0, fixed=true,
                   description="Supply elasticity (only relevant when r_b_phi = 1)",
                   tex_label="pphi")

    # Stochastic Properties of aggregate TFP shock
    m <= parameter(:nnu_aggZ, 0.25, fixed=true,
                   description="Persistence parameter in Ornstein-Uhlenbeck process",
                   tex_label="nnu_aggZ")
    m <= parameter(:sigma_aggZ, 0.007 / (1 - m[:aalpha]), fixed=true,
                   description="Volatility parameter in Ornstein-Uhlenbeck process",
                   tex_label="sigma_aggZ")

    # Household effects of TFP shock
    m <= parameter(:kappa, 10.0, fixed=true,
                   description="Household effects of TFP shock",
                   tex_label="kappa")

    # Steady State Parameters
    m <= SteadyStateParameter(:KL_SS, NaN,
                              description = "")
    m <= SteadyStateParameter(:r_a_SS, NaN,
                              description = "")
    m <= SteadyStateParameter(:w_SS, NaN,
                              description = "")
    m <= SteadyStateParameter(:K_SS, NaN,
                              description = "")
    m <= SteadyStateParameterGrid(:u_SS, Vector{Float64}(),
                                  description = "")
    m <= SteadyStateParameterGrid(:c_SS, Vector{Float64}(),
                                  description = "")
    m <= SteadyStateParameterGrid(:d_SS, Vector{Float64}(),
                                  description = "")
    m <= SteadyStateParameterGrid(:V_SS, Vector{Float64}(),
                                  description = "")
    m <= SteadyStateParameterGrid(:g_SS, Vector{Float64}(),
                                  description = "")
    m <= SteadyStateParameterGrid(:dab, Vector{Float64}(),
                                  description = "")
    m <= SteadyStateParameterGrid(:dab_g, Vector{Float64}(),
                                  description = "")
m <= SteadyStateParameter(:C_SS, NaN,
                          description = "")
m <= SteadyStateParameter(:C_Var_SS, NaN,
                          description = "")
m <= SteadyStateParameter(:I_SS, NaN,
                          description = "")
m <= SteadyStateParameter(:B_SS, NaN,
                          description = "")
m <= SteadyStateParameter(:Y_SS, NaN,
                          description = "")
m <= SteadyStateParameter(:n_SS, NaN,
                          description = "")
m <= SteadyStateParameterGrid(:earn_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameter(:earn_Var_SS, NaN,
                          description = "")
m <= SteadyStateParameterGrid(:IcF_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameterGrid(:IcB_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameterGrid(:Ic0_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameterGrid(:IcFB_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameterGrid(:IcBF_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameterGrid(:IcBB_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameterGrid(:Ic00_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameterGrid(:IcBB_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameterGrid(:WHTM_indicator, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameter(:WHTM_SS, NaN,
                          description = "")
m <= SteadyStateParameter(:C_WHTM_SS, NaN,
                          description = "")
m <= SteadyStateParameterGrid(:PHTM_indicator, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameter(:PHTM_SS, NaN,
                          description = "")
m <= SteadyStateParameter(:C_PHTM_SS, NaN,
                          description = "")
m <= SteadyStateParameterGrid(:NHTM_indicator, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameter(:NHTM_SS, NaN,
                          description = "")
m <= SteadyStateParameter(:C_NHTM_SS, NaN,
                          description = "")
m <= SteadyStateParameterGrid(:gg_SS, Vector{Float64}(),
                              description = "")
m <= SteadyStateParameter(:illiquid_wedge, NaN,
                          description = "")
m <= SteadyStateParameterGrid(:vars_SS, Vector{Float64}(),
                              description = "")
end

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

Calculates the model's steady-state values. `steadystate!(m)` must be called whenever
the parameters of `m` are updated.
"""
function steadystate!(m::TwoAssetHANK)
    # Read in parameters
    aalpha = m[:aalpha].value
    ddelta = m[:ddelta].value
    ddeath = m[:ddeath].value
    rrho   = m[:rrho].value
    chi0   = m[:chi0].value
    chi1   = m[:chi1].value
    chi2   = m[:chi2].value
    a_lb   = m[:a_lb].value
    pam    = m[:pam].value
    xxi    = m[:xxi].value
    ggamma = m[:ggamma].value
    tau_I  = m[:tau_I].value
    trans  = m[:trans].value

    # Set liquid rates
    r_b_SS       = m[:r_b_SS].value
    r_b_borr_SS  = m[:r_b_borr_SS].value
    borrwedge_SS = m[:borrwedge_SS].value

    lambda   = get_setting(m, :lambda)
    K_liquid = get_setting(m, :K_liquid)

    aggregate_variables        = get_setting(m, :aggregate_variables)
    distributional_variables   = get_setting(m, :distributional_variables)
    distributional_variables_1 = get_setting(m, :distributional_variables_1)

    # Read in approximation parameters
    maxit_HJB  = get_setting(m, :maxit_HJB)
    crit_HJB   = get_setting(m, :crit_HJB)
    Delta      = get_setting(m, :Delta)

    maxit_HIS  = get_setting(m, :maxit_HIS)
    crit_HIS   = get_setting(m, :crit_HIS)
    start_HIS  = get_setting(m, :start_HIS)

    maxit_KFE  = get_setting(m, :maxit_KFE)
    crit_KFE   = get_setting(m, :crit_KFE)
    Delta_KFE  = get_setting(m, :Delta_KFE)

    maxit_KL   = get_setting(m, :maxit_KL)
    crit_KL    = get_setting(m, :crit_KL)
    relax_KL   = get_setting(m, :relax_KL)

    # Read in grids
    I       = get_setting(m, :I)
    J       = get_setting(m, :J)
    a_g     = get_setting(m, :a_g)
    b_g     = get_setting(m, :b_g)
    I_g     = get_setting(m, :I_g)
    J_g     = get_setting(m, :J_g)
    N       = get_setting(m, :N)
    a       = get_setting(m, :a)
    b       = get_setting(m, :b)
    y       = get_setting(m, :y)
    y_dist  = get_setting(m, :y_dist)
    y_mean  = get_setting(m, :y_mean)
    KL      = get_setting(m, :KL_0)

    interp_decision = get_setting(m, :interp_decision)

    # Set liquid rates
    r_b      = r_b_SS
    r_b_borr = r_b_borr_SS

    # Compute prices associated with initial guess of KL
    w	= (1 - aalpha) * (KL ^ aalpha)
    r_a	= aalpha * (KL ^ (aalpha - 1)) - ddelta

    a_grid, a_g_grid, b_grid, b_g_grid, y_grid, y_g_grid, r_b_grid, r_b_g_grid, daf_grid, daf_g_grid, dab_grid, dab_g_grid, dab_tilde_grid, dab_g_tilde_grid, dab_g_tilde_mat, dab_g_tilde, dbf_grid, dbf_g_grid, dbb_grid, dbb_g_grid = set_grids(a, b, a_g, b_g, vec(y), r_b, r_b_borr)
    r_b_vec   = r_b .* (b   .>= 0) + r_b_borr .* (b   .< 0)

    # Construct problem functions
    util, deposit, cost = construct_problem_functions(ggamma, chi0, chi1, chi2, a_lb)

    # Initial consumption and value function
    c_0 = (1-xxi) * w * y_grid .+ (r_a + ddeath*pam) .* a_grid +
        (r_b_borr + ddeath * pam) .* b_grid .+ trans - tau_I * w * y_grid
    V_0	= 1/(rrho + ddeath) * util.(c_0)

    # Initial distribution
    gg0 = zeros(I_g, J_g, N)
    gg0[b .== 0, 1, :] = vec(y_dist)
    gg0 = gg0 ./ sum(gg0)
    gg0 = gg0 ./ dab_g_tilde_grid 	# ensure integration to 1
    gg0 = vec(gg0) #, I_g*J_g*N, 1)
    gg = gg0

    aau = Array{Float64}(undef, 0, 0)
    bbu = Array{Float64}(undef, 0, 0)
    ccu = Array{Float64}(undef, 0, 0)

    Vn1 = Array{Float64}(undef, I, J, N)
    Vn  = Array{Float64}(undef, I, J, N)

    gg_tilde = Array{Float64}(undef, 0, 0)

    g = Array{Float64}(undef, 0, 0)
    c = Array{Float64}(undef, I, J, N)
    s = Array{Float64}(undef, I, J, N)
    u = Array{Float64}(undef, I, J, N)
    d = Array{Float64}(undef, I, J, N)
    c_g = Array{Float64}(undef, 0, 0)
    s_g = Array{Float64}(undef, 0, 0)
    d_g = Array{Float64}(undef, 0, 0)

    IcF = BitArray{3}(undef, I,J,N)
    IcB = BitArray{3}(undef, I,J,N)
    Ic0 = BitArray{3}(undef, I,J,N)
    IcFB = BitArray{3}(undef, I,J,N)
    IcBF = BitArray{3}(undef, I,J,N)
    IcBB = BitArray{3}(undef, I,J,N)
    Ic00 = BitArray{3}(undef, I,J,N)

    K_supply  = 0.
    L_supply  = 0.
    KL_supply = 0.
    Vamin     = 0.
    Vbmin     = 1e-8

    #----------------------------------------------------------------
    # Iterate on KL to find steady state
    #----------------------------------------------------------------
    for ii = 1 : maxit_KL
	    # Derive aggregates, given KL
	    w   = (1 - aalpha) * (KL ^ aalpha)
 	    r_a = aalpha * (KL ^ (aalpha - 1)) - ddelta

	    # Store current value function
	    Vn = V_0

	    #----------------------------------------------------------------
	    # Solve HJB
	    #----------------------------------------------------------------
        @time for nn = 1 : maxit_HJB
            perm_c =  ddeath * pam # NOTE: perm?
            c0_c   = ((1-xxi) - tau_I) * w

            for i=1:I, j=1:J, n=1:N
                c0 = c0_c * real(y[n]) + b[i] * (r_b_vec[i] + perm_c) + trans

                # ---- Liquid Assets, Forward + Backward Difference ----
                VaF = (j==J) ? 0.0 : max((Vn[i,j+1,n] - Vn[i,j,n]) / (a[j+1]-a[j]), Vamin)
                VaB = (j==1) ? 0.0 : max((Vn[i,j,n] - Vn[i,j-1,n]) / (a[j]-a[j-1]), Vamin)

                # ---- Illiquid Assets, Forward + Backward Difference ----
                VbF = (i==I) ? 0.0 : max((Vn[i+1,j,n] - Vn[i,j,n]) / (b[i+1]-b[i]), Vbmin)
                VbB = (i==1) ? 0.0 : max((Vn[i,j,n] - Vn[i-1,j,n]) / (b[i]-b[i-1]), Vbmin)

                # Decisions conditional on a particular direction of derivative
                cF = (i==I) ? 0.0 : VbF ^ (-1 / ggamma)
                cB = (i==1) ? c0  : VbB ^ (-1 / ggamma)

                sF = (i==I) ? 0.0 : c0 - cF
                sB = (i==1) ? 0.0 : c0 - cB

                if (i==1 && j > 1) VbB = util(cB) end

                #----------------------------------------------------------------
                # Consumption  & Savings Decision
                #----------------------------------------------------------------
                # Which direction to use
                Hc0 = util(c0)
                HcF = (i==I) ? -1e12 : util(cF) + VbF * sF
                HcB = util(cB) + VbB * sB

                validF = (sF > 0)
                validB = (sB < 0)

                IcF[i,j,n] = validF * max(!validB, (HcF >= HcB)) * (HcF >= Hc0)
                IcB[i,j,n] = validB * max(!validF, (HcB >= HcF)) * (HcB >= Hc0)
                Ic0[i,j,n] = 1 - IcF[i,j,n] - IcB[i,j,n]

                c[i,j,n] = IcF[i,j,n] * cF + IcB[i,j,n] * cB + Ic0[i,j,n] * c0
                s[i,j,n] = IcF[i,j,n] * sF + IcB[i,j,n] * sB

                #----------------------------------------------------------------
                # Deposit Decision
                #----------------------------------------------------------------
                dFB = (i == 1 || j == J) ? 0.0 : deposit(VaF, VbB, max(a[j], a_lb))
                dBF = (i == I || j == 1) ? 0.0 : deposit(VaB, VbF, max(a[j], a_lb))
                dBB = (j == 1)           ? 0.0 : deposit(VaB, VbB, max(a[j], a_lb))

                HdFB = (i == 1 || j == J) ? -1e12 : VaF * dFB - VbB * (dFB + cost(dFB, a[j]))
                validFB = (dFB > 0) * (HdFB > 0)

                HdBB = (j == 1) ? -1.0e12 : VaB * dBB - VbB * (dBB + cost(dBB, a[j]))
                validBB = (dBB > -cost(dBB, a[j])) .* (dBB <= 0) * (HdBB > 0)

                HdBF = (i == I || j == 1) ? -1e12 : VaB * dBF - VbF * (dBF + cost(dBF, a[j]))
                validBF = (dBF <= -cost(dBF, a[j])) * (HdBF > 0)

                IcFB[i,j,n] = validFB * max(!validBF, HdFB >= HdBF) * max(!validBB, HdFB >= HdBB)
                IcBF[i,j,n] = max(!validFB, HdBF >= HdFB) * validBF * max(!validBB, HdBF >= HdBB)
                IcBB[i,j,n] = max(!validFB, HdBB >= HdFB) * max(!validBF, HdBB >= HdBF) * validBB
                Ic00[i,j,n] = !validFB * !validBF * !validBB

                d[i,j,n] = IcFB[i,j,n] * dFB + IcBF[i,j,n] * dBF + dBB * IcBB[i,j,n]
            end
            u = util.(c)

            # Interpolate
            d_g = reshape(interp_decision * vec(d), I_g, J_g, N)
            s_g = reshape(interp_decision * vec(s), I_g, J_g, N)
            c_g = reshape(interp_decision * vec(c), I_g, J_g, N)

            aa, bb, aau, bbu = transition(ddeath, pam, xxi, w, chi0, chi1, chi2, a_lb, r_a,
                                          vec(y), d, d_g, s, s_g, a, a_g, b, b_g)
            A = aa + bb

            #------------------------------------------------------------
            # Update value function
            #------------------------------------------------------------
            Vn1 = Array{Float64}(undef, I,J,N)
            for kk = 1:N
                Ak          = A[1+(kk-1)*(I*J):kk*(I*J), 1+(kk-1)*(I*J):kk*(I*J)]
                Bk          = (1 + Delta*(rrho + ddeath) -
                               Delta*lambda[kk,kk]) * my_speye(I*J) - Delta*Ak
                uk_stacked  = reshape(u[:,:,kk], I * J, 1)
                Vk_stacked  = reshape(Vn[:,:,kk], I * J, 1)
                indx_k      = 1:N .!= kk
                Vkp_stacked = sum(broadcast(*,  lambda[kk, indx_k]',
                                          reshape(Vn[:,:,indx_k],I*J,N-1)), dims=2)
                qk          = Delta * uk_stacked + Vk_stacked + Delta * Vkp_stacked
                Vn1k_stacked = Bk \ qk
                Vn1[:,:,kk]  = reshape(Vn1k_stacked, I, J, 1)
            end

            # Howard improvement step
            if nn >= start_HIS
                for jj = 1:maxit_HIS
                    Vn2 = Array{Float64}(undef, I, J, N)
                    for kk = 1:N
                        uk_stacked = reshape(u[:,:,kk],I*J,1)
                        Vk_stacked = reshape(Vn1[:,:,kk],I*J,1)
                        Ak = A[1+(kk-1)*(I*J):kk*(I*J), 1+(kk-1)*(I*J):kk*(I*J)]
                        Bk = (1 + Delta*(rrho + ddeath) -
                              Delta*lambda[kk,kk]) * my_speye(I*J) - Delta*Ak
                        indx_k         = 1:N .!= kk
                        Vkp_stacked = sum(broadcast(*,  lambda[kk, indx_k]',
                                          reshape(Vn1[:,:,indx_k],I*J,N-1)), dims=2)
                        qk             = Delta*uk_stacked + Vk_stacked + Delta*Vkp_stacked
                        Vn2k_stacked = Bk \ qk
                        Vn2[:,:,kk] = reshape(Vn2k_stacked,I,J,1)
                    end
                    VHIS_delta = Vn2 - Vn1
                    Vn1  = Vn2
                    dist = maximum(abs.(VHIS_delta))

                    if dist < crit_HIS
                        break
                    end
                end
            end

            # Check for convergence
            V_Delta = Vn1 - Vn
            Vn      = Vn1

            dist = maximum(abs.(V_Delta))
            @show dist

            if dist < crit_HJB
                 println("Value Function Converged, Iteration = ", nn)
                break
            end
        end

        # Store value function
        V_0 = Vn

        #----------------------------------------------------------------
        # Compute stationary distribution and update aggregate KL
        #----------------------------------------------------------------

        # Find new stationary distribution associated with decision rules
        A   = aau + bbu
        λ0  = lambda - diagm(0 => diag(lambda))  # transition matrix with diagonal killed
        λ0p = λ0'

        gg_tilde = dab_g_tilde_mat * gg
        gg1      = Array{Float64,2}(undef, I_g * J_g, N)
        g        = Array{Float64,3}(undef, I_g, J_g, N)

        K_supply  = 0.
        L_supply  = 0.
        KL_supply = 0.

        # Iterate
        for nn = 1:maxit_KFE

            gg_tilde = dab_g_tilde_mat * gg
            gg1      = Array{Float64,2}(undef, I_g * J_g, N)

            for kk = 1:N
                Ak = A[1+(kk-1)*(I_g*J_g):kk*(I_g*J_g),
                       1+(kk-1)*(I_g*J_g):kk*(I_g*J_g)]

                death_inflow = zeros(I_g, J_g)
                death_inflow[b .== 0, 1] .= sum(gg_tilde[1+(kk-1)*(I_g*J_g):kk*(I_g*J_g)])

                death_inflow = reshape(death_inflow, I_g*J_g, 1)

                gk_sum = sum(repeat(λ0p[kk,:]', I_g*J_g,1) .* reshape(gg_tilde,I_g*J_g,N), dims=2)
                gg1[:,kk] = (my_speye(I_g*J_g) - Delta_KFE * Ak' - Delta_KFE *
                             (lambda[kk,kk] - ddeath) *
                             my_speye(I_g*J_g))\(gg_tilde[1+(kk-1)*(I_g*J_g):kk*(I_g*J_g)] +
                                                 Delta_KFE*gk_sum + Delta_KFE*ddeath*death_inflow)
            end

            gg1     = reshape(gg1, I_g*J_g*N,1)
            gg1     = gg1 ./ sum(gg1)
            gg1     = dab_g_tilde_mat \ gg1

            dist = maximum(abs.(gg1-gg))

            if mod(nn,100) == 0
                println("Iteration ", nn, ", distance is ", dist)
            end
            if dist < crit_KFE
                gg = gg1
                g  = reshape(gg, I_g, J_g, N)

                println("Distribution Converged, Iteration = ", nn)
                break
            end
            gg = gg1
            g  = reshape(gg, I_g, J_g, N)
        end

        #------------------------------------------------------------
        # Update guess of KL
        #------------------------------------------------------------
        # Capital supply
        if K_liquid == 1
            K_supply = sum(g .* (a_g_grid + b_g_grid) .* dab_tilde_grid)
        else
            K_supply = sum(g .* a_g_grid .* dab_g_tilde_grid)
        end

        # Labor supply
        L_supply = sum(y_g_grid .* g .* dab_g_tilde_grid)

        # Capital-Labor ratio
        KL_supply = K_supply / L_supply

        # Check for convergence and update
        gap = KL - KL_supply
        println("The current gap is ", gap)

        if abs(gap) > crit_KL
            KL = relax_KL * KL_supply + (1 - relax_KL) * KL
        else
            println("I have found the steady state, Iteration = ", ii)
            break
        end

    end

    #compute_savings()
    ccu                  = kron(lambda, my_speye(I_g * J_g))
    A                    = aau + bbu + ccu
    dab_small            = reshape(dab_g_tilde_grid[:,:,1], I_g*J_g, 1)
    loc                  = findall(!iszero, b .== 0)
    dab_small            = dab_small ./ dab_small[loc] * ddeath
    dab_small[loc]      .= 0.0
    death_process        = -ddeath * my_speye(I_g*J_g)
    death_process[loc,:] = vec(dab_small)
    death_process        = kron(my_speye(N), death_process)

    g_new = A' * gg_tilde
    g_new = dab_g_tilde_mat \ g_new
    g_new = g_new + death_process * gg
    g_new = reshape(g_new ,I_g, J_g, N)
    g_new_a = sum(sum(g_new, dims=1), dims=3)
    save_a = dot(g_new_a, a_g)
    #g_new_b = sum(sum(g_new, dims=2), dims=3)
    #save_b = dot(vec(g_new_b), b_g)
    #compute_savings()

    # Rename variables in steady state
    m[:KL_SS]           = real(KL)
    m[:r_a_SS]          = aalpha * (m[:KL_SS].value ^ (aalpha - 1)) - ddelta
    m[:w_SS]            = (1 - aalpha) * (m[:KL_SS].value ^ aalpha)
    m[:K_SS]            = real(m[:KL_SS].value * L_supply)
    m[:u_SS]            = u
    m[:c_SS]            = c
    m[:d_SS]            = d
    m[:V_SS]            = Vn
    m[:g_SS]            = g
    m[:dab]             = dab_tilde_grid
    m[:dab_g]           = dab_g_tilde_grid
    m[:C_SS]            = sum(g .* c_g .* dab_g_tilde_grid)
    m[:C_Var_SS]        = sum(g .* log.(c_g).^2 .* dab_g_tilde_grid) -
                                 sum(g .* log.(c_g) .* dab_g_tilde_grid)^2
    m[:I_SS]            = save_a
    m[:B_SS]            = sum(g .* b_g_grid .* dab_g_tilde_grid)
    m[:Y_SS]            = real((K_supply ^ aalpha) * (L_supply ^ (1 - aalpha)))
    m[:n_SS]            = real(L_supply)
    m[:earn_SS]         = log.((1-tau_I) * w * y_g_grid + b_g_grid .*
                               (r_b_g_grid .+ ddeath*pam) .+ trans .+ a_g_grid .*
                               (r_a .+ ddeath*pam))
    m[:earn_Var_SS]     = sum(g .* m[:earn_SS].value .^ 2 .* dab_g_tilde_grid) -
                                sum(g .* m[:earn_SS].value .* dab_g_tilde_grid)^2
    m <= Setting(:IcF_SS,  IcF)
    m <= Setting(:IcB_SS,  IcB)
    m <= Setting(:Ic0_SS,  Ic0)
    m <= Setting(:IcFB_SS, IcFB)
    m <= Setting(:IcBF_SS, IcBF)
    m <= Setting(:IcBB_SS, IcBB)
    m <= Setting(:Ic00_SS, Ic00)

    a_g_0pos = get_setting(m, :a_g_0pos)
    b_g_0pos = get_setting(m, :b_g_0pos)

    ###
    # Consumption by hand-to-mouth status
    ###
    WHTM_indicator      = zeros(I_g, J_g, N)
    WHTM_indicator[b_g_0pos:b_g_0pos+1, a_g_0pos+2:end,:] .= 1.0

    m[:WHTM_indicator] = WHTM_indicator
    m[:WHTM_SS]        = sum(g[:] .* WHTM_indicator[:] .* dab_g_tilde_grid[:])
    m[:C_WHTM_SS]      = sum(WHTM_indicator[:] .* c_g[:] .* g[:] .* dab_g_tilde_grid[:])

    PHTM_indicator     = zeros(I_g, J_g, N)
    PHTM_indicator[b_g_0pos:b_g_0pos+1,a_g_0pos:a_g_0pos+2:end,:] .= 1.0
    m[:PHTM_indicator] = PHTM_indicator
    m[:PHTM_SS]        = sum(g[:] .* PHTM_indicator[:] .* dab_g_tilde_grid[:])
    m[:C_PHTM_SS]      = sum(c_g[:] .* g[:] .* PHTM_indicator[:] .* dab_g_tilde_grid[:])

    NHTM_indicator      = zeros(I_g,J_g,N)
    NHTM_indicator[(b_g_0pos+3):end,2:end,:] .= 1.0

    m[:NHTM_indicator] = NHTM_indicator
    m[:NHTM_SS]        = sum(g[:] .* NHTM_indicator[:] .* dab_g_tilde_grid[:])
    m[:C_NHTM_SS]      = sum(c_g[:] .* g[:] .* NHTM_indicator[:] .* dab_g_tilde_grid[:]) /
        m[:NHTM_SS].value
    m[:r_a_grid]       = repeat([r_a], I_g, J_g, N)

    n_v = get_setting(m, :n_v)
    n_g = get_setting(m, :n_g)
    n_p = get_setting(m, :n_p)
    n_Z = get_setting(m, :n_Z)

    # Collect variables into vector
    vars_SS                  = zeros(n_v + n_g + n_p + n_Z, 1)
    vars_SS[1:n_v]           = reshape(m[:V_SS].value, I*J*N, 1)

    gg_SS                    = reshape(m[:g_SS].value, I_g*J_g*N, 1)
    m[:gg_SS]                = gg_SS

    vars_SS[n_v+1:n_v+n_g,1] = gg_SS[1:I_g*J_g*N-1]
    vars_SS[n_v+n_g+1,1]     = m[:K_SS].value
    vars_SS[n_v+n_g+2,1]     = m[:r_b_SS].value

    if aggregate_variables == 1
        vars_SS[n_v+n_g+3,1]         = m[:Y_SS].value
        vars_SS[n_v+n_g+4,1]         = m[:C_SS].value
    elseif distributional_variables == 1
        vars_SS[n_v+n_g+3,1]        = m[:C_Var_SS].value
        vars_SS[n_v+n_g+4,1]        = m[:earn_Var_SS].value
    elseif distributional_variables_1 == 1
        vars_SS[n_v+n_g+3,1]        = m[:C_WHTM_SS].value
        vars_SS[n_v+n_g+4,1]        = m[:C_PHTM_SS]
        #vars_SS[n_v+n_g+3,1]        = C_NHTM_SS
    end
    vars_SS[n_v+n_g+n_p+1,1] = 0.0

    # Compute illiquid wedge
    m[:illiquid_wedge] = m[:r_a_SS].value - m[:r_b_SS].value

    # Save SS variables
    m[:vars_SS] = vars_SS

    return m
end

"""
```
function adjust_parameters!(m::TwoAssetHANK)
```
Function adjusts model parameters depending on what settings are provided.
"""
function adjust_parameters!(m::TwoAssetHANK)
    if get_setting(m, :para_new) == 0
        m[:rrho]   = 0.014526858925819 # discount factor (quarterly)
        m[:ddeath] = 1. / (4. * 45)    # death rate (quarterly)
        # Technology
        m[:aalpha] = 0.4     # capital share
        m[:ddelta] = 0.02    # depreciation rate (quarterly)

        # Portfolio adjustment cost function
        m[:chi0] = 0.070046
        m[:chi1] = 0.535791
        m[:chi2] = 0.971050
        m[:a_lb] = 0.2407

        # Tax code
        m[:trans] = 0.430147477922432 # transfer to households
        m[:tau_I] = 0.25              # tax rate on wage income

        # Liquid return
        m[:borrwedge_SS] = 0.019875
        m[:r_b_borr_SS]  = m[:r_b_SS] + m[:borrwedge_SS]

        # Stochastic Properties of aggregate TFP shock
        if !get_setting(m, :permanent)
            m[:nnu_aggZ]   = 0.05  # persistence parameter in Ornstein-Uhlenbeck process
            m[:sigma_aggZ] = 0.007 # volatility parameter in Ornstein-Uhlenbeck process
        else
            m[:nnu_aggZ]   = 2.15 # persistence parameter in Ornstein-Uhlenbeck process
            m[:sigma_aggZ] = 0.04 # volatility parameter in Ornstein-Uhlenbeck process
        end
    else

        m[:rrho]   = 0.014495841360662
        m[:ddelta] = 0.075 / 4
        m[:a_lb]   = 0.02
        m[:trans] = 0.1
        m[:tau_I] = 0.3
        m[:xxi]   = 0.0

        # Stochastic properties of aggregate TFP shock
        if !get_setting(m, :permanent)
            m[:nnu_aggZ]    = 0.25
            m[:ssigma_aggZ] = 0.007 / (1.0 - m[:aalpha])
        else
            m[:nnu_aggZ]    = 1.265
            m[:ssigma_aggZ] = 0.0282
        end
    end

    # Household effects of TFP shock
    if get_setting(m, :asymm)
        if get_setting(m, :y_size) == 2
            m[:kappa] = 10.0
        elseif get_setting(m, :y_size) == 33
            m[:kappa] = 2.0
        end
    else
        m[:kappa] = 0.0
    end

    # Necessary settings for permanent shock
    if get_setting(m, :permanent)
        m <= Setting(:reduceDist_hor, 400)
        m <= Setting(:r_b_fix, true)
    end
end

function model_settings!(m::TwoAssetHANK)
    default_settings!(m)

    # Productivity shock properties
    m <= Setting(:permanent, false,
                 "Permanent or transitory shock. For permanent, use reduceDist_hor=400, r_b_fix=1.")
    m <= Setting(:asymm,     false,
                 "Does effect of shock vary according to position in the wealth distribution?")

    # Solution approach
    m <= Setting(:reduceDistribution, true, "Reduction of g using observability matrix")
    m <= Setting(:reduceV,            true, "Reduction of v using splines")
    m <= Setting(:reduceDist_hor,     100,
                 "Maximal horizon for observability matrix")
    m <= Setting(:krylov_dim, 100, "Maximal horizon for observability matrix") # same as above

    # Liquid asset suply
    m <= Setting(:r_b_fix,  true, "Adjustment procedure 1: r_b is fixed")
    m <= Setting(:r_b_phi,  false,
                 "Adjustment procedure 2: r_b is supplied elastically, with elasticity phi")
    m <= Setting(:B_fix,    false, "The total amount of the liquid asset is fixed at B")
    m <= Setting(:K_liquid, false, "K_liquid")

    # Model size
    m <= Setting(:y_size, 30, "Model size: number of income states. Supported sizes: 2, 30, 33")

    # Distributional variables
    m <= Setting(:aggregate_variables,        true, "Aggregate variables")
    m <= Setting(:distributional_variables,   false, "Distributional_variables")
    m <= Setting(:distributional_variables_1, false, "Distributional_variables_1")

    # Parameterization
    m <= Setting(:para_new, 1, "Parameterization")

    # Grids
    m <= Setting(:agrid_new, 0, "agrid_new = 0")
    m <= Setting(:bgrid_new, 0, "bgrid_new = 0")
    m <= Setting(:ygrid_new, 0, "ygrid_new = 0")

    # Create income grid
    lambda, y, y_mean, y_dist = create_y_grid(get_setting(m, :y_size),
                                              get_setting(m, :ygrid_new))
    m <= Setting(:lambda, lambda, "Transition probabilities")
    m <= Setting(:y, y, "Grid of income shocks")
    m <= Setting(:y_mean, y_mean, "Mean grid of income shocks")
    m <= Setting(:y_dist, y_dist, "Stationary distribution of income shocks")
    m <= Setting(:N, length(y), "N")

    # Create illiquid asset grid
    m <= Setting(:J,    40, "Number of illiquid grid points (a)")
    m <= Setting(:J_g,  get_setting(m, :agrid_new) == 0 ? 45  : 40, "")
    m <= Setting(:amax, get_setting(m, :agrid_new) == 0 ? 200 : 2000,
                 "Maximum illiquid asset grid value")
    m <= Setting(:amin, 0, "Minimum illiquid asset grid value")
    m <= Setting(:a_,   0, "Lower bound on illiquid assets")

    a, a_g, a_g_0pos = create_a_grid(get_setting(m, :agrid_new), get_setting(m, :J),
                                     get_setting(m, :J_g), get_setting(m, :amin),
                                     get_setting(m, :amax))
    m <= Setting(:a,        a, "a grid")
    m <= Setting(:a_g,      a_g, "a_g grid")
    m <= Setting(:a_g_0pos, a_g_0pos[1], "Points in illiquid asset grid equal to 0")

    # Create liquid asset grid
    m <= Setting(:I,   50, "Number of liquid grid points")
    m <= Setting(:I_g, 50, "I_g")

    I_neg, bmin, bmax, b, b_g, b_g_0pos = create_b_grid(get_setting(m, :agrid_new),
                                                        get_setting(m, :I), get_setting(m, :I_g))
    m <= Setting(:I_neg,    I_neg, "Number of liquid grid points on negative part of real line")
    m <= Setting(:I_pos,    get_setting(m, :I) - I_neg,
                 "Number of liquid grid points on positive part of real line")
    m <= Setting(:bmin,     bmin, "Lower bound on liquid assets")
    m <= Setting(:bmax,     bmax, "Upper bound on liquid assets")
    m <= Setting(:b,        b, "b")
    m <= Setting(:b_g,      b_g, "b_g")
    m <= Setting(:b_g_0pos, b_g_0pos[1], "Points in liquid asset grid equal to 0")

    m <= Setting(:interp_decision, kron(my_speye(get_setting(m, :N)), interp(b_g, a_g, b, a)),
                                        "interp_decision")

    m <= Setting(:n_v,      get_setting(m, :I) * get_setting(m, :J) * get_setting(m, :N),
                 "Number of jump variables (value function entries)")
    m <= Setting(:n_g,      get_setting(m, :I_g) * get_setting(m, :J_g) * get_setting(m, :N) - 1,
                 "Number of endogenous state variables (distribution)")
    m <= Setting(:n_p,     (get_setting(m, :aggregate_variables) == 1) ||
                           (get_setting(m, :distributional_variables) == 1) ? 4 : 2, "n_p")
    m <= Setting(:n_Z,      1, "Number of aggregate shocks")
    m <= Setting(:nVars,    get_setting(m, :n_v) + get_setting(m, :n_g) + get_setting(m, :n_p) +
                            get_setting(m, :n_Z), "nVars")
    m <= Setting(:nEErrors, get_setting(m, :n_v), "nEErrors")
    m <= Setting(:n_v_full, get_setting(m, :n_v), "n_v_Full")
    m <= Setting(:n_g_full, get_setting(m, :n_g), "n_g_Full")

    # Our syntax
    m <= Setting(:n_jump_vars,  get_setting(m, :I) * get_setting(m, :J) * get_setting(m, :N),
                 "Number of jump variables (value function entries)")
    m <= Setting(:n_state_vars, get_setting(m, :I_g)*get_setting(m, :J_g)*get_setting(m, :N)-1,
                 "Number of endogenous state variables (distribution)")
    m <= Setting(:n_state_vars_unreduce, 1, "Number of aggregate shocks")

    # Steady state approximation
    m <= Setting(:maxit_HJB, 100,  "Max number of iterations for HJB")
    m <= Setting(:crit_HJB,  1e-6, "Tolerance for HJB error")
    m <= Setting(:Delta,     1e6,  "Multiplier for implicit upwind scheme of HJB")

    m <= Setting(:maxit_HIS, 10,   "Max allowable number of Howard improvement steps")
    m <= Setting(:start_HIS, 2,    "When in HJB loop should Howard improvement steps begin?")
    m <= Setting(:crit_HIS,  1e-5, "Critical value HIS")

    m <= Setting(:maxit_KFE, 1000, "Max allowable number of KFE iterations")
    m <= Setting(:crit_KFE,  1e-7, "Tolerance for KFE error")
    m <= Setting(:Delta_KFE, 1000, "Step size KFE")

    m <= Setting(:maxit_KL,  1000, "Max number of iterations over capital-labor ratio")
    m <= Setting(:crit_KL,   1e-5, "Tolerance for KL error")
    m <= Setting(:relax_KL,  0.1,  "Updating speed (must lie within [0,1])")

    # Approximation Parameters
    m <= Setting(:KL_0, 43.8800, "Initial guess of capital to labor ratio")

    # TODO: Inserted from old code; not certain necessary
    #=

    m <= Setting(:n_static_relations, 5,
                 "Number of static relations: bond-market clearing, labor market clearing, consumption, output, total assets")
    m <= Setting(:n_vars, get_setting(m, :n_jump_vars) + get_setting(m, :n_state_vars)
                 + get_setting(m, :n_static_relations),
                 "Number of variables, total")

    # Reduction
    m <= Setting(:n_knots, 12, "Number of knot points")
    m <= Setting(:c_power, 1, "Amount of bending of knot point locations to make them nonuniform")
    m <= Setting(:n_post, length(get_setting(m, :zz)[1,:]),
                 "Number of dimensions that need to be approximated by spline basis")
    m <= Setting(:n_prior, 1,
                 "Number of dimensions approximated by spline basis that
                  were not used to compute the basis matrix")

    knots = collect(range(get_setting(m, :amin), stop=get_setting(m, :amax),
                          length=get_setting(m, :n_knots)-1))
    knots = (get_setting(m, :amax) - get_setting(m, :amin)) /
        (2^get_setting(m, :c_power)-1) * ((knots .- get_setting(m, :amin))
             / (get_setting(m, :amax) - get_setting(m, :amin)) .+ 1) .^
             get_setting(m, :c_power) .+ get_setting(m, :amin) .- (get_setting(m, :amax) -
             get_setting(m, :amin)) / (2 ^ get_setting(m, :c_power) - 1)
    m <= Setting(:knots_dict, Dict(1 => knots),
                 "Location of knot points for each dimension for value function reduction")
    m <= Setting(:krylov_dim, 20, "Krylov reduction dimension")
    m <= Setting(:reduce_state_vars, true, "Reduce state variables or not") #turned off reduction
    m <= Setting(:reduce_v, true, "Reduce value function or not") #turned off reduction
    m <= Setting(:spline_grid, get_setting(m, :a), "Grid of knot points for spline basis")
=#
    # Sampling method
    m <= Setting(:sampling_method, :SMC, "Set sampling method to SMC")

    # SMC Settings
    m <= Setting(:n_particles, 3000)
    m <= Setting(:n_Φ, 200)
    m <= Setting(:λ, 2.0)
    m <= Setting(:n_smc_blocks, 1)
    m <= Setting(:use_parallel_workers, true)
    m <= Setting(:step_size_smc, 0.5)
    m <= Setting(:n_mh_steps_smc, 1)
    m <= Setting(:resampler_smc, :polyalgo)
    m <= Setting(:target_accept, 0.25)
    m <= Setting(:mixture_proportion, .9)
    m <= Setting(:tempering_target, 0.95)
    m <= Setting(:resampling_threshold, .5)
    m <= Setting(:use_fixed_schedule, true)
    m <= Setting(:smc_iteration, 0)

    # Track lag
    m <= Setting(:track_lag, false, "Add first lag when constructing measurement equation")

    # Forecast
    m <= Setting(:use_population_forecast, true,
                 "Whether to use population forecasts as data")
    m <= Setting(:forecast_zlb_value, 0.13,
        "Value of the zero lower bound in forecast periods, if we choose to enforce it")

    # Simulating states
    m <= Setting(:state_simulation_freq, 2,
                 "How many states you want to simulate between states + 1")
end


text:
using Test

using EvenTotalPayment

# https://www.extension.iastate.edu/agdm/wholefarm/html/c5-93.html
@testset "Smoke" begin
    p = even_total_payment(10000.0, 0.07, 20)
    @test round(Int, p.payment) == 944
    @test round(Int, p.principals[10]) == 448
    @test round(Int, p.interests[10]) == 495
    @test round(Int, total_interest(p)) == 8879
    @test round(Int, total_payment(p)) == 18879
end

# 微博高利贷
@testset "Weibo Loan Shark" begin
    # 根据微博消费贷规则，借款8000元按等额本息半年还清，每月需还款1476.78元。(王高飞：“加点利息怎么
    # 了？”)
    rate = estimate_interest_rate(8000.0, 6, 1476.78)
    @test round(Int, rate * 100 * 12) == 36  # 年化百分比利率（Annual percentage rate）高达36%
    # 有效年利率（Effective annual rate）
    ear = round((1.0 + rate) ^ 12 - 1.0; sigdigits=3)
    println("Effective annual rate = $(ear * 100)%")
end

ext: jl

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions


text:
module StringManipulation

import Base: convert, String, @kwdef

################################################################################
#                                  Constants
################################################################################

const _CSI = "\x1b["
# Regex that removes all ANSI escape sequences.
const _REGEX_ANSI = r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])"
# Escape sequence that reset all decorations.
const _RESET_DECORATIONS = _CSI * "0m"

################################################################################
#                                  Structures
################################################################################

# Enumetation to store the state in `Decoration`.
@enum DecorationState begin
    inactive  = 0
    active    = 1
    unchanged = 2
end

export Decoration

"""
    struct Decoration

Structure to hold the current decoration of a string.
"""
@kwdef struct Decoration
    foreground::String         = ""
    background::String         = ""
    bold::DecorationState      = unchanged
    underline::DecorationState = unchanged
    reset::Bool                = false
    reversed::DecorationState  = unchanged
end

const _DEFAULT_DECORATION = Decoration()
const _RESET_DECORATION = Decoration(reset = true)

################################################################################
#                                   Includes
################################################################################

include("./constants.jl")

include("./alignment.jl")
include("./ansi.jl")
include("./crop.jl")
include("./decorations.jl")
include("./highlighting.jl")
include("./search.jl")
include("./state.jl")
include("./split.jl")
include("./view.jl")
include("./width.jl")

# The environment variable `STRING_MANIPULATION_NO_PRECOMPILATION` is used to
# disable the precompilation directives. This option must only be used inside
# Github Actions to improve the coverage results.
if Base.VERSION >= v"1.4.2" && !haskey(ENV, "STRING_MANIPULATION_NO_PRECOMPILATION")
    # This try/catch is necessary in case the precompilation statements do not
    # exists. In this case, StringManipulation.jl will work correctly but
    # without the optimizations.
    try
        include("../precompilation/precompile_StringManipulation.jl")
        include("../precompilation/precompile_StringManipulation_manual.jl")
        _precompile_()
        _precompile_manual_()
    catch
    end
end

end # module


text:
"""
    Deterministic Training Conditional (DTC) covariance strategy.
"""
struct DeterminTrainCondStrat{M<:AbstractMatrix} <: SparseStrategy
    inducing::M
end
SubsetOfRegsStrategy(dtc::DeterminTrainCondStrat) = SubsetOfRegsStrategy(dtc.inducing)
function init_precompute(covstrat::DeterminTrainCondStrat, X, y, k)
    SoR = SubsetOfRegsStrategy(covstrat)
    return init_precompute(SoR, X, y, k)
end

function alloc_cK(covstrat::DeterminTrainCondStrat, nobs)
    SoR = SubsetOfRegsStrategy(covstrat)
    return alloc_cK(SoR, nobs)
end
function update_cK!(cK::SubsetOfRegsPDMat, x::AbstractMatrix, kernel::Kernel,
                    logNoise::Real, data::KernelData, covstrat::DeterminTrainCondStrat)
    SoR = SubsetOfRegsStrategy(covstrat)
    return update_cK!(cK, x, kernel, logNoise, data, SoR)
end
function dmll_kern!(dmll::AbstractVector, gp::GPBase, precomp::SoRPrecompute, covstrat::DeterminTrainCondStrat)
    SoR = SubsetOfRegsStrategy(covstrat)
    return dmll_kern!(dmll, gp, precomp, SoR)
end
function dmll_noise(gp::GPE, precomp::SoRPrecompute, covstrat::DeterminTrainCondStrat)
    SoR = SubsetOfRegsStrategy(covstrat)
    return dmll_noise(gp, precomp, SoR)
end

"""
    Deterministic Training Conditional (DTC) multivariate normal predictions.

    See Quiñonero-Candela and Rasmussen 2005, equations 20b.
        μ_DTC = μ_SoR
        Σ_DTC = Σxx - Qxx + Σ_SoR

    where μ_DTC and Σ_DTC are the predictive mean and covariance
    functions for the Subset of Regressors approximation.
"""
function predictMVN(xpred::AbstractMatrix,
                    xtrain::AbstractMatrix, ytrain::AbstractVector,
                    kernel::Kernel, meanf::Mean,
                    alpha::AbstractVector,
                    covstrat::DeterminTrainCondStrat, Ktrain::AbstractPDMat)
    SoR = SubsetOfRegsStrategy(covstrat)
    μ_SoR, Σ_SoR = predictMVN(xpred, xtrain, ytrain, kernel, meanf, alpha, SoR, Ktrain)
    inducing = covstrat.inducing
    Kuu = Ktrain.Kuu

    sparsedata = SparseKernelData(kernel, inducing, xpred, xpred)
    Qxx = getQaa(Ktrain, kernel, xpred, sparsedata)
    densedata = KernelData(kernel, xpred, xpred)
    Σxx = cov(kernel, xpred, xpred, densedata)

    Σ_DTC = Σxx - Qxx + Σ_SoR
    LinearAlgebra.copytri!(Σ_DTC, 'U')
    return μ_SoR, Σ_DTC
end

function DTC(x::AbstractMatrix, inducing::AbstractMatrix, y::AbstractVector, mean::Mean, kernel::Kernel, logNoise::Real)
    covstrat = DeterminTrainCondStrat(inducing)
    GPE(x, y, mean, kernel, logNoise, covstrat)
end


text:
push!(LOAD_PATH, dirname(@__DIR__))

using Documenter
using StructArrays

makedocs(
    sitename = "StructArrays",
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
    modules = [StructArrays],
    pages = [
            "Overview"=>"index.md",
            "Example usage"=>"examples.md",
            "Some counterintuitive behaviors"=>"counterintuitive.md",
            "Advanced techniques"=>"advanced.md",
            "Index"=>"reference.md",
            ]
)

# Documenter can also automatically deploy documentation to gh-pages.
# See "Hosting Documentation" and deploydocs() in the Documenter manual
# for more information.
deploydocs(
    repo = "github.com/JuliaArrays/StructArrays.jl.git",
    push_preview = true,
)


text:
#############################################################
#
# Users Functions
#
#############################################################

function get_verify_credentials(; options=Dict{AbstractString, AbstractString}())

    r = get_oauth("https://api.twitter.com/1.1/account/verify_credentials.json", options)

    #return to_USERS(r)
    return r.status == 200 ? to_USERS(Requests.json(r)) : r

end

function post_account_settings(; options=Dict{AbstractString, AbstractString}())

    r = post_oauth("https://api.twitter.com/1.1/account/settings.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

function post_update_delivery_device(device::AbstractString; options=Dict{AbstractString, AbstractString}())

    #Add required parameter(s) to options dict
    options["device"] = device

    r = post_oauth("https://api.twitter.com/1.1/account/update_delivery_device.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

function post_update_profile(; options=Dict{AbstractString, AbstractString}())

    r = post_oauth("https://api.twitter.com/1.1/account/update_profile.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

function post_update_profile_background(; options=Dict{AbstractString, AbstractString}())

    r = post_oauth("https://api.twitter.com/1.1/account/update_profile_background_image.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

function post_update_profile_colors(; options=Dict{AbstractString, AbstractString}())

    r = post_oauth("https://api.twitter.com/1.1/account/update_profile_colors.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

#Need to enforce type around image as base64 encoded, not sure how
function post_update_profile_image(image; options=Dict{AbstractString, AbstractString}())

    #Add required parameter(s) to options dict
    options["image"] = image

    r = post_oauth("https://api.twitter.com/1.1/account/update_profile_image.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

function get_blocks_list(; options=Dict{AbstractString, AbstractString}())

    r = get_oauth("https://api.twitter.com/1.1/blocks/list.json", options)

    #return to_USERS(r["users"])
    return r.status == 200 ? to_USERS(Requests.json(r)["users"]) : r

end

function post_blocks_create(; options=Dict{AbstractString, AbstractString}())

    r = post_oauth("https://api.twitter.com/1.1/blocks/create.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

function post_blocks_destroy(; options=Dict{AbstractString, AbstractString}())

    r = post_oauth("https://api.twitter.com/1.1/blocks/destroy.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

function get_users_lookup(; options=Dict{AbstractString, AbstractString}())

    r = get_oauth("https://api.twitter.com/1.1/users/lookup.json", options)

    #return to_USERS(r)
    return r.status == 200 ? to_USERS(Requests.json(r)) : r

end

function get_users_show(; options=Dict{AbstractString, AbstractString}())

    r = get_oauth("https://api.twitter.com/1.1/users/show.json", options)

    #return to_USERS(r)
    return r.status == 200 ? to_USERS(Requests.json(r)) : r

end

function get_users_search(q::AbstractString; options=Dict{AbstractString, AbstractString}())

    #Add required parameter(s) to options dict
    options["q"] = q

    r = get_oauth("https://api.twitter.com/1.1/users/search.json", options)

    #return to_USERS(r)
    return r.status == 200 ? to_USERS(Requests.json(r)) : r

end

function get_users_contributees(; options=Dict{AbstractString, AbstractString}())

    r = get_oauth("https://api.twitter.com/1.1/users/contributees.json", options)

    #return to_USERS(r)
    return r.status == 200 ? to_USERS(Requests.json(r)) : r

end

function get_users_contributors(; options=Dict{AbstractString, AbstractString}())

    r = get_oauth("https://api.twitter.com/1.1/users/contributors.json", options)

    #return to_USERS(r)
    return r.status == 200 ? to_USERS(Requests.json(r)) : r

end

function post_account_remove_profile_banner(; options=Dict{AbstractString, AbstractString}())

    r = post_oauth("https://api.twitter.com/1.1/account/remove_profile_banner.json", options)

    return r.status == 200 ? Requests.json(r) : r

end

#Need to enforce type around image as base64 encoded, not sure how
function post_account_update_profile_banner(banner::AbstractString; options=Dict{AbstractString, AbstractString}())

    #Add required parameter(s) to options dict
    options["banner"] = banner

    r = post_oauth("https://api.twitter.com/1.1/account/update_profile_banner.json", options)

    return r.status == 200 ? Requests.json(r) : r

end


text:
"""
Model

Specifies the unconstrained program for minimizing a smooth objective function.

A model is constructed with a `name` and may be incrementally loaded.
Once a model is loaded, the model becomes final, meaning the objective and gradient
function may no longer be modified. A model creates a directory storing logged
information throughout the model life. The model's directory is found relative
to the user's current working directory, identified by the models name used as
the relative portion of the directory path.
"""
mutable struct Model
    name::String
    objective::Union{Function,Missing}
    gradient::Union{Function,Missing}
    initial_iterate::Union{AbstractArray,Missing}
    formula::Union{String,Missing}

    dimension::Union{Int,Missing}
    directory::String
    final::Bool

    function Model(
        name::String;
        objective = missing,
        gradient = missing,
        initial_iterate = missing,
        formula = missing,
    )
        dimension = isa(initial_iterate, Missing) ? missing : length(initial_iterate)
        directory = mkpath(joinpath(pwd(), name))
        final = !isa(objective, Missing) && !isa(gradient, Missing)

        new(
            name,
            objective,
            gradient,
            initial_iterate,
            formula,
            dimension,
            directory,
            final,
        )
    end
end

"""
    name(m::Model)

Access the name associated with model `m`, as assigned during the construction of `m`.
"""
name(m::Model) = getfield(m, :name)

"""
    objective(m::Model)

The objective function of model `m`, having an unassigned default value of `missing`.

See `objective!` to load a model's objective.
"""
objective(m::Model) = getfield(m, :objective)

"""
    gradient(m::Model)

The gradient function of model `m`, having an unassigned default value of `missing`.

See `gradient!` to load a `m`'s gradient.
"""
gradient(m::Model) = getfield(m, :gradient)

"""
    initial_iterate(m::Model)

The initial iterate of model `m`, having an unassigned default value of `missing`.
Setting the initial iterate of a model assigns the model a dimension.

See `initial_iterate!` to load `m`'s starting location.
"""
initial_iterate(m::Model) = getfield(m, :initial_iterate)

"""
    formula(m::Model)

The formula of model `m`, having an unassigned default value of `missing`.

See `formula!` to set `m`'s formula.
"""
formula(m::Model) = getfield(m, :formula)

"""
    dimension(m::Model)

The dimension of model `m` which is set by specifying an initial iterate.

See: initial_iterate!
"""
dimension(m::Model) = getfield(m, :dimension)

"""
    directory(m::Model)

The directory path of model `m` holding information related to `m`.
"""
directory(m::Model) = getfield(m, :directory)

"""
    final(m::Model)

The final status of model `m`, occuring when the model's gradient
and objective function are no longer missing. A finalized model may no
longer modify it's previously defined gradient or objective.
"""
final(m::Model) = getfield(m, :final)

"""
    objective!(m::Model, f)

Assign model `m` the objective function `f`.

If the model is final, the `objective!` call simply returns without modifying the model.
"""
objective!(m::Model, f) = !final(m) && begin

    setfield!(m, :final, !isa(gradient(m), Missing))

    setfield!(m, :objective, f)
end

"""
    gradient!(m::Model, ∇f!)

Assign model `m` the in-place gradient function `∇f!`, of the form

```julia
∇f!(out, x) = (out .= ∇f(x))
```

where `∇f(x)` is the steepest descent direction at `x` in stored the place
of the inputed buffer `out`.

If the model is final, the `gradient!` call simply returns without modifying the model.
"""
gradient!(m::Model, ∇f!) =
    !final(m) && begin

        setfield!(m, :final, !isa(objective(m), Missing))

        setfield!(m, :gradient, ∇f!)
    end

"""
    initial_iterate!(m::Model, x0)

Assign model `m` an initial iterate.
"""
initial_iterate!(m::Model, x0) = begin

    setfield!(m, :dimension, length(x0))

    setfield!(m, :initial_iterate, x0)
end

"""
    formula!(m::Model, f)

Assign model `m` an escaped ``\\LaTeX`` string represents `m`'s objective function.
The formula is used to create visual plots showing a simulation trace of `m`.

## Example
```julia
julia> dot_formila = "\$ f(x) = x⋅x \$";
julia> formula!(m, dot_formila);
```
Here, the euclidean squared distance formula now represents `m`'s objective formula.
"""
formula!(m::Model, formula) = setfield!(m, :formula, formula)

"""
    obj(m::Model, x)

Evaluate the objective function of model `m` at `x`.
"""
obj(m::Model, x) = begin
    isa(dimension(m), Missing) && return missing
    @lencheck dimension(m) x
    objective(m)(x)
end

"""
    grad!(m::Model, out, x)

Evaluate the in-place gradient function of model `m` at `x`, storing the steepest descent
direction in the place of the passed input buffer `out`.
"""
grad!(m::Model, out, x) = begin
    isa(gradient(m), Missing) && return missing
    @lencheck dimension(m) out x
    return gradient(m)(out, x)
end

"""
    grad(m::Model, x)

Evaluates the gradient function of model `m` at `x`, allocating similar storage
for the output as given by the input `x`.
"""
grad(m::Model, x) = begin
    isa(gradient(m), Missing) && return missing
    @lencheck dimension(m) x
    grad!(m, similar(x), x)
end

"""
    hessAD(m::Model, x)

The dense Hessian matrix of model `m`'s objective function at the point `x`.
The computation uses the `ForwardDiff.jacobian` forward-mode AD routine on the model's
gradient function.
"""
hessAD(m::Model, x) = begin
    isa(gradient(m), Missing) && return missing
    jacobian(gradient(m), similar(x), x)
end

"""
    hess_sample(m::Model, x, dx)

The hessian vector-product of model `m`'s objective function at the point `x`
with the vector `dx.` The `gHS` routine uses a more effiecient scheme to compute
the hessian samples with ForwardDiff's `Dual` numbers.
"""
hess_sample(m::Model, x, dx) = begin
    isa(gradient(m), Missing) && return missing
    hessAD(m, x) * dx
end

Base.getproperty(m::Model, s::Symbol) = @restrict Model

Base.propertynames(m::Model) = ()


text:
@testset "missings" begin
    @testset "any(ismissing, x)" begin
        # any(ismissing, a::AbstractArray)
        @test any(ismissing, Any[missing, 1])
        @test !any(ismissing, [1, 2])
        @test !any(ismissing, repeat([1, 2], outer = [1, 2]))
        @test !any(ismissing, repeat([1, 2], outer = [1, 2, 2]))

        # any(ismissing, da::DataArray)
        @test !any(ismissing, DataArray([1, 2], falses(2)))
        @test !any(ismissing, DataArray(repeat([1, 2], outer = [1, 2]), falses(2, 2)))
        da = DataArray(repeat([1, 2], outer = [1, 2, 2]), falses(2, 2, 2))
        @test !any(ismissing, da)
        da[2] = missing
        @test any(ismissing, da)

        # any(ismissing, pda::PooledDataArray)
        @test !any(ismissing, PooledDataArray([1, 2], falses(2)))
        @test !any(ismissing, PooledDataArray(repeat([1, 2], outer = [1, 2]), falses(2, 2)))
        pda = PooledDataArray(repeat([1, 2], outer = [1, 2, 2]), falses(2, 2, 2))
        @test !any(ismissing, pda)
        pda[2] = missing
        @test any(ismissing, pda)
    end

    @testset "all(ismissing, x)" begin
        # all(ismissing, a::AbstractArray)
        @test all(ismissing, Any[missing, missing])
        @test !all(ismissing, Any[missing, 1])
        @test !all(ismissing, [1, 2])
        @test !all(ismissing, repeat([1, 2], outer = [1, 2]))
        @test !all(ismissing, repeat([1, 2], outer = [1, 2, 2]))

        # all(ismissing, da::DataArray)
        @test !all(ismissing, DataArray([1, 2], falses(2)))
        @test !all(ismissing, DataArray(repeat([1, 2], outer = [1, 2]), falses(2, 2)))
        da = DataArray(repeat([1, 2], outer = [1, 2, 2]), falses(2, 2, 2))
        da[1] = missing
        @test !all(ismissing, da)
        da[:] = missing
        @test all(ismissing, da)

        # all(ismissing, da::PooledDataArray)
        @test !all(ismissing, PooledDataArray([1, 2], falses(2)))
        @test !all(ismissing, PooledDataArray(repeat([1, 2], outer = [1, 2]), falses(2, 2)))
        pda = PooledDataArray(repeat([1, 2], outer = [1, 2, 2]), falses(2, 2, 2))
        pda[1] = missing
        @test !all(ismissing, pda)
        pda[:] = missing
        @test all(ismissing, pda)
    end

    dv = DataArray(collect(1:6), fill(false, 6))
    a = dv[.!ismissing.(dv)]
    @test collect(Missings.fail(dv)) == a
    @test collect(skipmissing(dv)) == a
    @test collect(Missings.replace(dv, 4)) == a

    dv[[1, 2, end]] = missing

    @testset "promotion" for (T1, T2) in ((Int, Float64),
                                          (Dates.Minute, Dates.Second))
        @eval begin
            @test promote_type($T1, Union{$T2, Missing})              == Union{$T2, Missing}
            @test promote_type(Union{$T1, Missing}, $T2)              == Union{$T2, Missing}
            @test promote_type(Union{$T1, Missing}, Union{$T2, Missing}) == Union{$T2, Missing}
        end
    end

    a = dv[.!ismissing.(dv)]
    @test_throws MissingException for v in Missings.fail(dv); end
    @test collect(skipmissing(dv)) == a
    @test collect(Missings.replace(dv, 4)) == [4, 4, a..., 4]

end


text:
module coexist

end # module


text:
#using XfromProjections.curve_utils
#using XfromProjections.snake_forward
#using XfromProjections
using Logging
using IterTools
using LinearAlgebra
include("./utils.jl")
include("../../phantoms/sperm_phantom.jl")
include("../snake_forward.jl")
include("../snake.jl")
include("../curve_utils.jl")

cwd = @__DIR__
savepath = normpath(joinpath(@__DIR__, "result"))
!isdir(savepath) && mkdir(savepath)

#object to store center and radius of circle in complex plane
struct circle
    center::Complex{Float64}
    radius::Float64
end

#return the parametric equation for a circle in the complex plane
function get_parametric_circle_equation(c::circle)
    return t -> c.radius*ℯ^(t*im) + c.center
end

images, tracks = get_sperm_phantom(301,grid_size=0.1)

detmin, detmax = -38.0, 38.0
grid = collect(detmin:0.1:detmax)
bins = collect(detmin:0.125:detmax)
angles = collect(range(0.0,π,length=6))[1:end-1]
iterations = 1000
α = 0.01
β = 0.01
step_size = 0.2
r(s) = 1.0

num_points = 62
t = range(0, 2π, length=num_points)
c = circle((0.0-15.0*im), 15.0)
γ = get_parametric_circle_equation(c)
compl_circ = γ.(t)
outline_templ = cat(real(compl_circ), imag(compl_circ), dims=2)

frames2reconstruct = collect(1:10:300)
reconstructions = zeros(num_points,2,length(frames2reconstruct))
for (iter, frame_nr) in Base.Iterators.reverse(enumerate(frames2reconstruct))
    @info iter frame_nr

    #Get projection
    @info "making forward projection for frame: " frame_nr
    outline, normals = get_outline(tracks[frame_nr], r)
    sinogram = parallel_forward(outline, angles, bins)

    #Add noise
    @info "adding gaussian noise at level 0.01"
    rho = 0.01
    e = randn(size(sinogram));
    e = rho*norm(sinogram)*e/norm(e);
    sinogram = sinogram + e;

    @info "reconstructing"
    reconstruction = vedrana(outline_templ,angles,bins,sinogram, iterations, α, β, step_size)
    reconstructions[:,:,iter] = reconstruction

    @info "plotting"
    heatmap(grid, grid, Gray.(images[:,:,frame_nr]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=true)
    plot!(reconstruction[:,1], reconstruction[:,2], aspect_ratio=:equal, linewidth=2, label=frame_nr)
    cd(savepath)
    savefig(@sprintf "result_original_vedrana_1%d" frame_nr)
end

global ps = AbstractPlot[]
for (iter, frame_nr) in enumerate(frames2reconstruct)
    best_recon = reconstructions[:,:,iter]

    l = @layout [a b c]
    p = heatmap(grid, grid, Gray.(images[:,:,frame_nr]), yflip=false, aspect_ratio=:equal, framestyle=:none, legend=false)
    plot!(best_recon[:,1],best_recon[:,2], aspect_ratio=:equal, linewidth=5)
    push!(ps,p)
    if length(ps) == 3
        plot(ps[1], ps[2], ps[3], layout = l, size=(2000,600), linewidth=5)
        savefig(@sprintf "result_all_original_vedrana_%d_%d" length(angles) frame_nr)
        global ps = AbstractPlot[]
    end
end
cd(cwd)

max_stars_repo_name: QPImaging/xfromProjections

max_issues_repo_name: QPImaging/xfromProjections

max_forks_repo_name: QPImaging/xfromProjections


text:
"""
    evaluate_field_array(f::AbstractArray,x::AbstractArray) -> AbstractArray

Evaluates the fields in the array `f` at all the vector of points in the
array of vector of points `x` and returns the result as a lazy array.

The result is numerically equivalent to

    map(evaluate_field,f,x)
"""
function evaluate_field_array(a::AbstractArray,x::AbstractArray)
  _evaluate_field_array(a,x)
end

function _evaluate_field_array(a::AbstractArray,x::AbstractArray)
  k = Eval()
  apply(k,a,x)
end

struct Eval <: Kernel end

function kernel_cache(k::Eval,a,x)
  field_cache(a,x)
end

function apply_kernel!(cache,k::Eval,a,x)
  evaluate_field!(cache,a,x)
end

function kernel_return_type(k::Eval,a,x)
  field_return_type(a,x)
end

function _evaluate_field_array(a::AbstractArray{<:Field},x::AbstractArray)
  apply(a,x)
end

"""
    evaluate_field_arrays(f::Tuple,x::AbstractArray) -> Tuple

Equivalent to

    tuple((evaluate_field_array(fi,x) for fi in f)...)
"""
function evaluate_field_arrays(f::Tuple,x::AbstractArray)
  _evaluate_field_arrays(x,f...)
end

function _evaluate_field_arrays(x,a,b...)
  ax = evaluate_field_array(a,x)
  bx = evaluate_field_arrays(b,x)
  (ax,bx...)
end

function _evaluate_field_arrays(x,a)
  ax = evaluate_field_array(a,x)
  (ax,)
end

"""
    evaluate(a::AbstractArray{<:Field},x::AbstractArray)

Equivalent to

    evaluate_field_array(a,x)

But only for arrays `a` whose element type inherits from `Field`. If this is not the case,
use `evaluate_field_array(a,x)` instead.
"""
function evaluate(a::AbstractArray{<:Field},x::AbstractArray)
  evaluate_field_array(a,x)
end

# Optimized version for arrays of fields obtained from a kernel
# and other arrays
function evaluate_field_array(
  a::AppliedArray{T,N,F,<:Fill} where {T,N,F},x::AbstractArray)
  kernel_evaluate(a.g.value,x,a.f...)
end

"""
    kernel_evaluate(k,x,f...)

Function to control the evaluation of the field resulting from the operation `apply(k,f...)`, where `k`
is a kernel and `f...` contains several arrays of fields.

By default, this function is implemented as

    g = apply(k,f...)
    evaluate_field_array(g,x)

However, it can be rewritten for specific kernels in order to improve performance and simplify
the underlying operation tree.

"""
function kernel_evaluate(k,x,f...)
  a = apply(k,f...)
  _evaluate_field_array(a,x)
end

"""
    field_array_gradient(a::AbstractArray)

Returns an array containing the gradients of the fields in the array `a`.
Numerically equivalent to

    map(field_gradient,a)
"""
function field_array_gradient(a::AbstractArray)
  _field_array_gradient(a)
end

function _field_array_gradient(a::AbstractArray)
  k = Grad()
  apply(k,a)
end

struct Grad <: Kernel end

@inline apply_kernel!(::Nothing,k::Grad,x) = field_gradient(x)

function field_array_gradient(
  a::AppliedArray{T,N,F,<:Fill} where {T,N,F})
  apply_gradient(a.g.value,a.f...)
end

"""
    apply_gradient(k,f...)

Function to control the generation of the gradient the field resulting from the operation `apply(k,f...)`, where `k`
is a kernel and `f...` contains several arrays of fields.

By default, it returns the array obtained as

    a = apply(k,f...)
    field_array_gradient(a)

However, it can be rewritten for specific kernels in order to improve performance and simplify
the underlying operation tree.

"""
function apply_gradient(k,f...)
  a = apply(k,f...)
  _field_array_gradient(a)
end

"""
    gradient(f::AbstractArray{<:Field})

Equivalent to

    field_array_gradient(f)

but only for arrays whose element type is `<:Field`. Use function `field_array_gradient` otherwise.
"""
function gradient(f::AbstractArray{<:Field})
  field_array_gradient(f)
end

"""
    field_array_gradients(f...)

Equivalent to

    map(field_array_gradient,f)
"""
function field_array_gradients(f...)
  map(field_array_gradient,f)
end

"""
    field_array_cache(a::AbstractArray,x::AbstractArray) -> Tuple

Returns the caches needed to perform the following iteration

    ca, cfi, cx = field_array_cache(a,x)

    for i in length(a)
      fi = getindex!(ca,a,i)
      xi = getindex!(cx,x,i)
      fxi = evaluate!(cfi,fi,xi)
    end
"""
function field_array_cache(a::AbstractArray,x::AbstractArray)
  ca = array_cache(a)
  fi = testitem(a)
  xi = testitem(x)
  cfi = field_cache(fi,xi)
  cx = array_cache(x)
  (ca,cfi,cx)
end

"""
    function test_array_of_fields(
      a::AbstractArray,
      x::AbstractArray,
      v::AbstractArray,
      cmp::Function=(==);
      grad = nothing)

Function to test an array of fields `a`. The array `v` is the expected result when calling
`evaluate_field_array(a,x)`. The entries in the computed array and the expected one are compared
with the `cmp` function. The key-word argument `grad` is optional. If present, it should contain
the expected result of

    ∇a = field_array_gradient(a)
    evaluate_field_array(∇a,x)
"""
function test_array_of_fields(
  a::AbstractArray,
  x::AbstractArray,
  v::AbstractArray,
  cmp::Function=(==);
  grad = nothing)

  ax = evaluate_field_array(a,x)
  test_array(ax,v,cmp)

  ca, cfi, cx = field_array_cache(a,x)

  t = true
  for i in 1:length(a)
    fi = getindex!(ca,a,i)
    xi = getindex!(cx,x,i)
    fxi = evaluate_field!(cfi,fi,xi)
    vi = v[i]
    ti = cmp(fxi,vi)
    t = t && ti
  end
  @test t

  if grad != nothing
    g = field_array_gradient(a)
    test_array_of_fields(g,x,grad,cmp)
  end

end

"""
    apply_to_field_array(k,f::AbstractArray...)
    apply_to_field_array(::Type{T},k,f::AbstractArray...) where T

Returns an array of fields numerically equivalent to

    map( (x...) -> apply_kernel_to_field(k,x...), f )

"""
function apply_to_field_array(
  k,f::AbstractArray...)
  v = Valued(k)
  apply(v,f...)
end

function apply_to_field_array(
  ::Type{T},k,f::AbstractArray...) where T
  v = Valued(k)
  apply(T,v,f...)
end

struct Valued{K} <: Kernel
  k::K
  function Valued(k)
    new{typeof(k)}(k)
  end
end

@inline function apply_kernel!(cache,k::Valued,x...)
  apply_kernel_to_field(k.k,x...)
end

function kernel_evaluate(k::Valued,x,f...)
  fx = evaluate_field_arrays(f,x)
  a = apply(k.k,fx...)
end

# More optimizations

function evaluate_field_array(a::AppendedArray,b::AppendedArray)
  if (length(a.a) == length(b.a)) && (length(a.b) == length(b.b))
    c_a = evaluate_field_array(a.a,b.a)
    c_b = evaluate_field_array(a.b,b.b)
    lazy_append(c_a,c_b)
  else
    _evaluate_field_array(a,b)
  end
end

function evaluate_field_array(a::AppendedArray,b::AbstractArray)
  n = length(a.a)
  _b = lazy_append(lazy_split(b,n)...)
  evaluate_field_array(a,_b)
end

function field_array_gradient(a::AppendedArray)
  c_a = field_array_gradient(a.a)
  c_b = field_array_gradient(a.b)
  lazy_append(c_a,c_b)
end


text:
module TupleStructs
export tostruct

"""tostruct: creates an expression ex that when eval'ed creates type structname
The field types are taken sequentially from t.types.
The field names are generated sequentially as "fi" where i is the index into types.

A classic example from the manual:

```
    eval(tostruct(:Point, Tuple{Float64, Float64})
    p =  Point(1.5, 2.0)
    @assert p.f1 == 1.5
    @assert p.f2 == 2.0
```

Keep in mind that types cannot be recreated in the same scope.
See tests for more examples.
"""
function tostruct(structname::Symbol, t::Type)
   types = t.types
   ex = quote
       type $structname
       end
   end
   for (i,ty) in enumerate(types)
       fieldname = symbol("f$i")
       fielddef = :($fieldname::$ty)
       push!(ex.args[2].args[3].args, fielddef)
   end
   return ex
end

end


text:
include("synthesize.jl")
include("wavgen.jl")

using .Synthesize, .WavGen

## Examples

music1 = Music("C C G G A A G-2 F F E E D D C-2", 100)

music2 = Music("E' E' F' G' G' F' E' D' C' C' D' E' E'-1.5 D'-0.5 D' 0 E' E' F' G' G' F' E' D' C' C' D' E' D'-1.5 C'-0.5 C' 0 D' D' E' C' D' E'-0.5 F'-0.5 E' C' D' E'-0.5 F'-0.5 E' D' C' D' G E'-2 E' F' G' G' F' E' D' C' C' D' E' D'-1.5 C'-0.5 C'", 140)

music3 = Music("E'-0.75 F'-0.25 G' E'-0.5 C'-0.5 F#'-0.5 G'-0.5 E'-0.5 C'-0.5 H-0.5 C'-0.5 C#'-0.5 D'-0.5 G C'-0.75 H-0.25 A-0.5 H-0.25 C'-0.25 D'-0.25 E'-0.25 F'-0.25 G'-0.25 G#'-0.25 A'-0.25 G'-0.25 F'-0.25 E'-0.25 D'-0.25 E'-0.25 F'-0.25 C'-1.5 D'-0.25 C'-0.25 H-0.5", 86)

writeToWav("beethoven.wav", musicFromNotes(music2, organ), 44100)
writeToWav("mozart.wav", musicFromNotes(music3, flute), 44100)

max_stars_repo_name: agostonsipos/synthesizer

max_issues_repo_name: agostonsipos/synthesizer

max_forks_repo_name: agostonsipos/synthesizer


text:
## reductions ##

###### Functors ######

# Note that functors are merely used as internal machinery to enhance code reuse.
# They are not exported.
# When function arguments can be inlined, the use of functors can be removed.

abstract Func{N}

type IdFun <: Func{1} end
type AbsFun <: Func{1} end
type Abs2Fun <: Func{1} end
type ExpFun <: Func{1} end
type LogFun <: Func{1} end

type AddFun <: Func{2} end
type MulFun <: Func{2} end
type AndFun <: Func{2} end
type OrFun <: Func{2} end
type MaxFun <: Func{2} end
type MinFun <: Func{2} end

evaluate(::IdFun, x) = x
evaluate(::AbsFun, x) = abs(x)
evaluate(::Abs2Fun, x) = abs2(x)
evaluate(::ExpFun, x) = exp(x)
evaluate(::LogFun, x) = log(x)
evaluate(f::Callable, x) = f(x)

evaluate(::AddFun, x, y) = x + y
evaluate(::MulFun, x, y) = x * y
evaluate(::AndFun, x, y) = x & y
evaluate(::OrFun, x, y) = x | y
evaluate(::MaxFun, x, y) = scalarmax(x, y)
evaluate(::MinFun, x, y) = scalarmin(x, y)
evaluate(f::Callable, x, y) = f(x, y)

###### Generic (map)reduce functions ######

# r_promote: promote x to the type of reduce(op, [x])
r_promote(op, x) = x
r_promote(::AddFun, x) = x + zero(x)
r_promote(::MulFun, x) = x * one(x)

## foldl && mapfoldl

function mapfoldl_impl(f, op, v0, itr, i)
    if done(itr, i)
        return v0
    else
        (x, i) = next(itr, i)
        v = evaluate(op, v0, evaluate(f, x))
        while !done(itr, i)
            (x, i) = next(itr, i)
            v = evaluate(op, v, evaluate(f, x))
        end
        return v
    end
end

mapfoldl(f, op, v0, itr) = mapfoldl_impl(f, op, v0, itr, start(itr))

function mapfoldl(f, op::Function, v0, itr)
    is(op, +) ? mapfoldl(f, AddFun(), v0, itr) :
    is(op, *) ? mapfoldl(f, MulFun(), v0, itr) :
    is(op, &) ? mapfoldl(f, AndFun(), v0, itr) :
    is(op, |) ? mapfoldl(f, OrFun(), v0, itr) :
    mapfoldl_impl(f, op, v0, itr, start(itr))
end

function mapfoldl(f, op, itr)
    i = start(itr)
    done(itr, i) && error("Argument is empty.")
    (x, i) = next(itr, i)
    v0 = evaluate(f, x)
    mapfoldl_impl(f, op, v0, itr, i)
end

foldl(op, v0, itr) = mapfoldl(IdFun(), op, v0, itr)
foldl(op, itr) = mapfoldl(IdFun(), op, itr)

## foldr & mapfoldr

function mapfoldr_impl(f, op, v0, itr, i::Integer)
    if i == 0
        return v0
    else
        x = itr[i]
        v = evaluate(op, evaluate(f, x), v0)
        while i > 1
            x = itr[i -= 1]
            v = evaluate(op, evaluate(f, x), v)
        end
        return v
    end
end

mapfoldr(f, op, v0, itr) = mapfoldr_impl(f, op, v0, itr, endof(itr))
mapfoldr(f, op, itr) = (i = endof(itr); mapfoldr_impl(f, op, evaluate(f, itr[i]), itr, i-1))

foldr(op, v0, itr) = mapfoldr(IdFun(), op, v0, itr)
foldr(op, itr) = mapfoldr(IdFun(), op, itr)

## reduce & mapreduce

# mapreduce_***_impl require ifirst < ilast
function mapreduce_seq_impl(f, op, A::AbstractArray, ifirst::Int, ilast::Int)
    @inbounds fx1 = evaluate(f, A[ifirst])
    @inbounds fx2 = evaluate(f, A[ifirst+=1])
    @inbounds v = evaluate(op, fx1, fx2)
    while ifirst < ilast
        @inbounds fx = evaluate(f, A[ifirst+=1])
        v = evaluate(op, v, fx)
    end
    return v
end

function mapreduce_pairwise_impl(f, op, A::AbstractArray, ifirst::Int, ilast::Int, blksize::Int)
    if ifirst + blksize > ilast
        return mapreduce_seq_impl(f, op, A, ifirst, ilast)
    else
        imid = (ifirst + ilast) >>> 1
        v1 = mapreduce_seq_impl(f, op, A, ifirst, imid)
        v2 = mapreduce_seq_impl(f, op, A, imid+1, ilast)
        return evaluate(op, v1, v2)
    end
end

mapreduce(f, op, itr) = mapfoldl(f, op, itr)
mapreduce(f, op, v0, itr) = mapfoldl(f, op, v0, itr)
mapreduce_impl(f, op, A::AbstractArray, ifirst::Int, ilast::Int) =
    mapreduce_seq_impl(f, op, A, ifirst, ilast)

# handling empty arrays
mr_empty(f, op, T) = error("Reducing over an empty array is not allowed.")
mr_empty(::IdFun, op::AddFun, T) = r_promote(op, zero(T))
mr_empty(::AbsFun, op::AddFun, T) = r_promote(op, abs(zero(T)))
mr_empty(::Abs2Fun, op::AddFun, T) = r_promote(op, abs2(zero(T)))
mr_empty(::IdFun, op::MulFun, T) = r_promote(op, one(T))
mr_empty(::AbsFun, op::MaxFun, T) = abs(zero(T))
mr_empty(::Abs2Fun, op::MaxFun, T) = abs2(zero(T))
mr_empty(f, op::AndFun, T) = true
mr_empty(f, op::OrFun, T) = false

function _mapreduce{T}(f, op, A::AbstractArray{T})
    n = length(A)
    if n == 0
        return mr_empty(f, op, T)
    elseif n == 1
        return r_promote(op, evaluate(f, A[1]))
    elseif n < 16
        @inbounds fx1 = evaluate(f, A[1])
        @inbounds fx2 = evaluate(f, A[2])
        s = evaluate(op, fx1, fx2)
        i = 2
        while i < n
            @inbounds fx = evaluate(f, A[i+=1])
            s = evaluate(op, s, fx)
        end
        return s
    else
        return mapreduce_impl(f, op, A, 1, n)
    end
end

mapreduce(f, op, A::AbstractArray) = _mapreduce(f, op, A)
mapreduce(f, op, a::Number) = evaluate(f, a)

function mapreduce(f, op::Function, A::AbstractArray)
    is(op, +) ? _mapreduce(f, AddFun(), A) :
    is(op, *) ? _mapreduce(f, MulFun(), A) :
    is(op, &) ? _mapreduce(f, AndFun(), A) :
    is(op, |) ? _mapreduce(f, OrFun(), A) :
    _mapreduce(f, op, A)
end

reduce(op, v0, itr) = mapreduce(IdFun(), op, v0, itr)
reduce(op, itr) = mapreduce(IdFun(), op, itr)
reduce(op, a::Number) = a

###### Specific reduction functions ######

## sum

function mapreduce_seq_impl(f, op::AddFun, a::AbstractArray, ifirst::Int, ilast::Int)
    @inbounds if ifirst + 6 >= ilast  # length(a) < 8
        i = ifirst
        s = evaluate(f, a[i]) + evaluate(f, a[i+1])
        i = i+1
        while i < ilast
            s += evaluate(f, a[i+=1])
        end
        return s

    else # length(a) >= 8, manual unrolling
        s1 = evaluate(f, a[ifirst]) + evaluate(f, a[ifirst + 4])
        s2 = evaluate(f, a[ifirst + 1]) + evaluate(f, a[ifirst + 5])
        s3 = evaluate(f, a[ifirst + 2]) + evaluate(f, a[ifirst + 6])
        s4 = evaluate(f, a[ifirst + 3]) + evaluate(f, a[ifirst + 7])
        i = ifirst + 8
        il = ilast - 3
        while i <= il
            s1 += evaluate(f, a[i])
            s2 += evaluate(f, a[i+1])
            s3 += evaluate(f, a[i+2])
            s4 += evaluate(f, a[i+3])
            i += 4
        end
        while i <= ilast
            s1 += evaluate(f, a[i])
            i += 1
        end
        return s1 + s2 + s3 + s4
    end
end

# Note: sum_seq uses four accumulators, so each accumulator gets at most 256 numbers
sum_pairwise_blocksize(f) = 1024

# This appears to show a benefit from a larger block size
sum_pairwise_blocksize(::Abs2Fun) = 4096

mapreduce_impl(f, op::AddFun, A::AbstractArray, ifirst::Int, ilast::Int) =
    mapreduce_pairwise_impl(f, op, A, ifirst, ilast, sum_pairwise_blocksize(f))

sum(f::Union(Function,Func{1}), a) = mapreduce(f, AddFun(), a)
sum(a) = mapreduce(IdFun(), AddFun(), a)
sum(a::AbstractArray{Bool}) = countnz(a)
sumabs(a) = mapreduce(AbsFun(), AddFun(), a)
sumabs2(a) = mapreduce(Abs2Fun(), AddFun(), a)

# Kahan (compensated) summation: O(1) error growth, at the expense
# of a considerable increase in computational expense.
function sum_kbn{T<:FloatingPoint}(A::AbstractArray{T})
    n = length(A)
    if n == 0
        return sumzero(T)
    end
    c = zero(T)
    s = A[1] + c
    for i in 2:n
        @inbounds Ai = A[i]
        t = s + Ai
        if abs(s) >= abs(Ai)
            c += ((s-t) + Ai)
        else
            c += ((Ai-t) + s)
        end
        s = t
    end
    s + c
end

## prod

prod(f::Union(Function,Func{1}), a) = mapreduce(f, MulFun(), a)
prod(a) = mapreduce(IdFun(), MulFun(), a)

prod(A::AbstractArray{Bool}) =
    error("use all() instead of prod() for boolean arrays")

## maximum & minimum

function mapreduce_impl(f, op::MaxFun, A::AbstractArray, first::Int, last::Int)
    # locate the first non NaN number
    v = evaluate(f, A[first])
    i = first + 1
    while v != v && i <= last
        @inbounds v = evaluate(f, A[i])
        i += 1
    end
    while i <= last
        @inbounds x = evaluate(f, A[i])
        if x > v
            v = x
        end
        i += 1
    end
    v
end

function mapreduce_impl(f, op::MinFun, A::AbstractArray, first::Int, last::Int)
    # locate the first non NaN number
    v = evaluate(f, A[first])
    i = first + 1
    while v != v && i <= last
        @inbounds v = evaluate(f, A[i])
        i += 1
    end
    while i <= last
        @inbounds x = evaluate(f, A[i])
        if x < v
            v = x
        end
        i += 1
    end
    v
end

maximum(f::Union(Function,Func{1}), a) = mapreduce(f, MaxFun(), a)
minimum(f::Union(Function,Func{1}), a) = mapreduce(f, MinFun(), a)

maximum(a) = mapreduce(IdFun(), MaxFun(), a)
minimum(a) = mapreduce(IdFun(), MinFun(), a)

maxabs(a) = mapreduce(AbsFun(), MaxFun(), a)
minabs(a) = mapreduce(AbsFun(), MinFun(), a)

## extrema

extrema(r::Range) = (minimum(r), maximum(r))
extrema(x::Real) = (x, x)

function extrema(itr)
    s = start(itr)
    done(itr, s) && error("argument is empty")
    (v, s) = next(itr, s)
    while v != v && !done(itr, s)
        (x, s) = next(itr, s)
        v = x
    end
    vmin = v
    vmax = v
    while !done(itr, s)
        (x, s) = next(itr, s)
        if x > vmax
            vmax = x
        elseif x < vmin
            vmin = x
        end
    end
    return (vmin, vmax)
end

## all & any

function mapfoldl(f, ::AndFun, itr)
    for x in itr
        if !evaluate(f, x)
            return false
        end
    end
    return true
end

function mapfoldl(f, ::OrFun, itr)
    for x in itr
        if evaluate(f, x)
            return true
        end
    end
    return false
end

function mapreduce_impl(f, op::AndFun, A::AbstractArray, ifirst::Int, ilast::Int)
    while ifirst <= ilast
        @inbounds x = A[ifirst]
        if !evaluate(f, x)
            return false
        end
        ifirst += 1
    end
    return true
end

function mapreduce_impl(f, op::OrFun, A::AbstractArray, ifirst::Int, ilast::Int)
    while ifirst <= ilast
        @inbounds x = A[ifirst]
        if evaluate(f, x)
            return true
        end
        ifirst += 1
    end
    return false
end

all(a) = mapreduce(IdFun(), AndFun(), a)
any(a) = mapreduce(IdFun(), OrFun(), a)

all(pred::Union(Function,Func{1}), a) = mapreduce(pred, AndFun(), a)
any(pred::Union(Function,Func{1}), a) = mapreduce(pred, OrFun(), a)

## in & contains

immutable EqX{T} <: Func{1}
    x::T
end
EqX{T}(x::T) = EqX{T}(x)
evaluate(f::EqX, y) = (y == f.x)

in(x, itr) = any(EqX(x), itr)

const ∈ = in
∉(x, itr)=!∈(x, itr)
∋(itr, x)= ∈(x, itr)
∌(itr, x)=!∋(itr, x)

function contains(itr, x)
    depwarn("contains(collection, item) is deprecated, use in(item, collection) instead", :contains)
    in(x, itr)
end

function contains(eq::Function, itr, x)
    for y in itr
        if eq(y, x)
            return true
        end
    end
    return false
end

## countnz & count

function count(pred::Union(Function,Func{1}), itr)
    n = 0
    for x in itr
        if evaluate(pred, x)
            n += 1
        end
    end
    return n
end

function count(pred::Union(Function,Func{1}), a::AbstractArray)
    n = 0
    for i = 1:length(a)
        @inbounds if evaluate(pred, a[i])
            n += 1
        end
    end
    return n
end

type NotEqZero <: Func{1} end
evaluate(NotEqZero, x) = (x != 0)

countnz(a) = count(NotEqZero(), a)

max_stars_repo_name: jakebolewski/julia

max_issues_repo_name: jakebolewski/julia

max_forks_repo_name: jakebolewski/julia


text:
using JLD, SparseArrays, Test, Random

include("src/DC.jl")
using .DC

@testset "Exactness on trees" begin
    dir = "test/Ising_trees/"
    ls = readdir(dir)
    filenames = replace.(ls[findall(occursin.("model",ls))], "model.jld"=>"")
    @testset "Ising open-chains $i" for i in 1:length(filenames)
        model = load(string(dir,filenames[i],"model.jld"));
        true_moments = load(string(dir,filenames[i],"moments.jld"),"moments");
        N = length(model["h"])
        βvec = [0.1,0.3]
        @testset "β $β" for β in βvec
            psi = DC.IsingFG(N,β*model["h"],β*sparse(model["J"]), :EnergyFun);
            ϕ , res, it, epsv = DC.density_consistency(psi,N,verbose = false, epsconv = 1e-10,ρ = 0.9)
            @test isapprox(true_moments[β][1], ϕ.µt, atol = 1e-7)
        end
    end
end

@testset "Equivalence of Ising Constructors" begin
    dir = "test/random_graphs/"
    ls = readdir(dir)
    filenames = replace.(ls[findall(occursin.("model",ls))], "model.jld"=>"")
    model = load(string(dir,filenames[3],"model.jld"));
    N = length(model["h"])
    β = 0.2
    psiF = DC.IsingFG(N,β*model["h"],β*sparse(model["J"]), :EnergyFun);
    psiPI = DC.IsingFG(N,β*model["h"],β*sparse(model["J"]), :IsingPair);
    ϕF , res, it, epsv = DC.density_consistency(psiF,N,verbose = false, epsconv = 1e-10, ρ = 0.9)
    ϕI , res, it, epsv = DC.density_consistency(psiPI,N,verbose = false, epsconv = 1e-10, ρ = 0.9)
    #@testset "fields $x" for x in fieldnames(typeof(ϕF))
    #    @test isapprox(getfield(ϕF,x),getfield(ϕI,x),atol = 1e-9)
    #end
    @test isapprox(ϕF.µt,ϕI.µt,atol = 1e-7)
end

max_stars_repo_name: giovact/DensityConsistency

max_issues_repo_name: giovact/DensityConsistency

max_forks_repo_name: giovact/DensityConsistency


text:
import CUDA.CUDNN: cudnnActivationForwardAD
using CUDA.CUDNN: cudnnActivationBackward, scalingParameter, handle
using AutoGrad: AutoGrad, @primitive1, value

@primitive1(
    (cudnnActivationForwardAD(x; activationDesc, alpha, xDesc, beta, yDesc, y),
     _dy, _y),
    ((x, y, dy, dx) = (value(x), value(_y), value(_dy), similar(x));
     if alpha[] != 1; y = y ./ alpha[]; end;
     cudnnActivationBackward(handle(), activationDesc, alpha, yDesc, y, yDesc, dy, xDesc, x, beta, xDesc, dx);
     dx))


text:
#=

    Seguir tutorial de instalação do Julia demonstrado no README.md
    Como rodar esse script: $ julia bhaskara.jl

    Objetivo desse script: realizar uma operação básica de bhaskara.
    Referência: https://brasilescola.uol.com.br/matematica/formula-bhaskara.htm

=#

#=
    Declaração de variáveis a, b e c respectivamente
    Como visto, variáveis no Julia não são tipadas como em algumas linguagens,
    ou seja, você pode declará-las e o inicializador delas assumirá o valor
    que você colocar como atribuição.
    Ex.: Em a = 1 (a será um inteiro) / Em a = 'Olá' (a será uma string)
=#
a = 1
b = 0
c = -16

#   Definindo uma variável chamada 'delta' que assumirá o papel de
#   uma função onde algo será executado e o resultado será armazenado
delta = function (a, b, c)
    return b * b - 4 * a * c
end

#   Definindo uma variável chamada 'raiz1' que assumirá o papel de
#   uma função onde algo será executado e o resultado será armazenado
raiz1 = function (a, b, c)
    return (-b + Base.Math.sqrt(delta(a, b, c))) / 2 * a;
end

#   Definindo uma variável chamada 'raiz2' que assumirá o papel de
#   uma função onde algo será executado e o resultado será armazenado
raiz2 = function (a, b, c)
    return (-b - Base.Math.sqrt(delta(a, b, c))) / 2 * a;
end

#   Printando resultado das funções atribuídos a variável respectiva.
#   Comentário para demonstrar o valor que deve ser apresentado.
#   Seria interessante colocar um assert.equals para ficar mais nítida a compreensão.

println(delta(a, b, c)) # Deve imprimir: 64
println(raiz1(a, b, c)) # Deve imprimir: 4
println(raiz2(a, b, c)) # Deve imprimir: -4