"""
    FiniteTempBasisSet

Type for holding IR bases and sparse-sampling objects.

An object of this type holds IR bases for fermions and bosons
and associated sparse-sampling objects.

# Fields
- basis_f::FiniteTempBasis: Fermion basis
- basis_b::FiniteTempBasis: Boson basis
- beta::Float64: Inverse temperature
- wmax::Float64: Cut-off frequency
- tau::Vector{Float64}: Sampling points in the imaginary-time domain
- wn_f::Vector{Int}: Sampling fermionic frequencies
- wn_b::Vector{Int}: Sampling bosonic frequencies
- smpl_tau_f::TauSampling: Sparse sampling for tau & fermion
- smpl_tau_b::TauSampling: Sparse sampling for tau & boson
- smpl_wn_f::MatsubaraSampling: Sparse sampling for Matsubara frequency & fermion
- smpl_wn_b::MatsubaraSampling: Sparse sampling for Matsubara frequency & boson
- sve_result::Tuple{PiecewiseLegendrePoly,Vector{Float64},PiecewiseLegendrePoly}: Results of SVE
"""
struct FiniteTempBasisSet{
    TSf<:TauSampling,MSf<:MatsubaraSampling,TSb<:TauSampling,MSb<:MatsubaraSampling
}
    basis_f::_DEFAULT_FINITE_TEMP_BASIS
    basis_b::_DEFAULT_FINITE_TEMP_BASIS
    smpl_tau_f::TSf # TauSampling{Float64,_DEFAULT_FINITE_TEMP_BASIS,Float64,Float64}
    smpl_tau_b::TSb # TauSampling{Float64,_DEFAULT_FINITE_TEMP_BASIS,Float64,Float64}
    smpl_wn_f::MSf # MatsubaraSampling{Int,_DEFAULT_FINITE_TEMP_BASIS,ComplexF64,Float64}
    smpl_wn_b::MSb # MatsubaraSampling{Int,_DEFAULT_FINITE_TEMP_BASIS,ComplexF64,Float64}
end

"""
    FiniteTempBasisSet(β, wmax, ε; sve_result=compute_sve(LogisticKernel(β * wmax); ε))

Create basis sets for fermion and boson and
associated sampling objects.
Fermion and bosonic bases are constructed by SVE of the logistic kernel.
"""
function FiniteTempBasisSet(
    β::AbstractFloat, wmax::AbstractFloat, ε;
    sve_result=compute_sve(LogisticKernel(β * wmax); ε),
)
    # Create bases using the given sve results
    basis_f = FiniteTempBasis(fermion, β, wmax, ε; sve_result)
    basis_b = FiniteTempBasis(boson, β, wmax, ε; sve_result)

    return FiniteTempBasisSet(
        basis_f, basis_b,
        TauSampling(basis_f), TauSampling(basis_b),
        MatsubaraSampling(basis_f), MatsubaraSampling(basis_b),
    )
end

getbeta(bset::FiniteTempBasisSet) = getbeta(bset.basis_f)
getwmax(bset::FiniteTempBasisSet) = getwmax(bset.basis_f)

function Base.getproperty(bset::FiniteTempBasisSet, d::Symbol)
    if d === :tau
        # return getfield(bset, :smpl_tau_f).sampling_points
        return bset.smpl_tau_f.sampling_points
    elseif d === :wn_f
        # return getfield(bset, :smpl_wn_f).sampling_points
        return bset.smpl_wn_f.sampling_points
    elseif d === :wn_b
        # return getfield(bset, :smpl_wn_b).sampling_points
        return bset.smpl_wn_b.sampling_points
    elseif d === :sve_result
        # return getfield(bset, :basis_f).sve_result
        return bset.basis_f.sve_result
    else
        return getfield(bset, d)
    end
end

function Base.propertynames(::FiniteTempBasisSet, private::Bool=false)
    return (:tau, :wn_f, :wn_b, :sve_result, propertynames(FiniteTempBasisSet, private)...)
end

function Base.show(io::IO, b::FiniteTempBasisSet)
    return print(io, "FiniteTempBasisSet: beta=$(getbeta(b)), wmax=$(getwmax(b))")
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM, Test

# Test that if 3d continuum problem have some elements we don't know how to
# deal with, raise error with clear message

elements = [Element(Seg3, (1, 2, 3))]
problem = Problem(Elasticity, "test seg3", 3)
add_elements!(problem, elements)
@test_throws ErrorException assemble!(problem, 0.0)


using DataFrames
using SpecialFunctions
using LinearAlgebra
using Struve

version = v"0.4.0"

# constants
const OPEN = 1
const CLOSE = 0
const HALF = 0.5
const HEAD = 0
const LAST = 1

const GMM = 1.4  # specific head ratio
const PR = 0.72  # Prandtl number

const type_keywords = Dict("SPLIT"=>:split,
    "TONEHOLE"=>:split,
    "|"=>:split,
    "VALVE_OUT"=>:branch,
    "BRANCH"=>:branch,
    "<"=>:branch,
    "MERGE"=>:merge,
    "VALVE_IN"=>:merge,
    ">"=>:merge,
    "OPEN_END"=>:open,
    "CLOSED_END"=>:close,
    "INSERT"=>:insert,
    "@"=>:insert
    )

const group_keywords = Dict("MAIN"=>:main,
    "["=>:main,
    "END_MAIN"=>:endmain,
    "]"=>:endmain,
    "GROUP"=>:group,
    "{"=>:group,
    "END_GROUP"=>:endgroup,
    "}"=>:endgroup
    )

type_keys = keys(type_keywords)
group_keys = keys(group_keywords)

"""
Men struct which contains basic unit for expressing mensur structure.
"""
mutable struct Men
    # data to store
    df::Float64 # diameter at input
    db::Float64 # diameter at output
    r::Float64  # length of mensur cell
    cm::AbstractString # comment
    gr::AbstractString # group name

    # calculated values
    zi::ComplexF64 # impedance at input
    zo::ComplexF64 # impedance at output
    vi::Array{ComplexF64,1} # (p,U) vector at input
    vo::Array{ComplexF64,1} # (p,U) vector at output
    tm::Array{ComplexF64,2} # transmission matrix

    # child info
    cname::String
    ctype::String
    cratio::Float64

    # list reference
    prev::Union{Men,Nothing}
    next::Union{Men,Nothing}
    child::Union{Men,Nothing}
    parent::Union{Men,Nothing}

    # incomplete constructor
    Men() = new()
end

function isbranch(t::String)
    return(type_keywords[t] == :branch)
end

function issplit(t::String)
    return(type_keywords[t] == :split)
end

function isaddon(t::String)
    return(type_keywords[t] == :addon)
end

function ismerge(t::String)
    return(type_keywords[t] == :merge)
end

"""
create new mensur
"""
function men_create(df::Number=0.0, db::Number=0.0, r::Number=0.0, cm::AbstractString="", gr::AbstractString="")
    men = Men()

    men.df = df
    men.db = db
    men.r = r
    men.cm = cm
    men.gr = gr

    men.zi = complex(0.0)
    men.zo = complex(0.0)
    men.vi=zeros(ComplexF64,2)
    men.vo=zeros(ComplexF64,2)
    men.tm=zeros(ComplexF64,2,2)

    men.cname = ""
    men.ctype = ""
    men.cratio = 0.0

    men.prev = nothing
    men.next = nothing
    men.parent = nothing
    men.child = nothing

    return(men)
end

function men_append!(prev::Men, curr::Men)
    prev.next = curr
    curr.prev = prev
    return(curr)
end

"""insert ins after cur"""
function men_insert!(cur::Men, ins::Men)
    nxt = cur.next
    inse = men_end(ins)

    cur.next = ins
    ins.prev = cur

    nxt.prev = inse
    inse.next = nxt
end

function men_setchildinfo!(men::Men, t::String, n::String,r::Float64)
    men.ctype=t
    men.cname=n
    men.cratio=r
end

function men_top(curr::Men)
    c = curr
    if c != nothing
        while c.prev != nothing
            c = c.prev
        end
    end
    return(c)
end

function men_end(curr::Men)
    c = curr
    if c != nothing
        while c.next != nothing
            c = c.next
        end
    end
    return(c)
end

function eat_comment(s)
    m = match(r"#", s)
    if m != nothing
        ss = s[1:m.offset-1]
    else
        ss = s
    end
    return(ss)
end

function getfbr(cur::Men)
    return(cur.df, cur.db,cur.r,cur.cm)
end

function men_bykeyword(cur::Men, wd::Array{String,1})
    w = wd[1]
    if length(wd) > 1
        name = wd[2]
    end
    if length(wd) > 2
        ratio = float(eval(Meta.parse(wd[3])))
    end
    df,db,r,c = getfbr(cur)
    gr = cur.gr

    men = nothing
    if type_keywords[w] == :close
        men = men_create(0.0,0.0,0.0,"",gr )
    elseif type_keywords[w] == :open
        men = men_create(db,0.0,0.0,"",gr )
    else
        if type_keywords[w] == :insert
            ratio = 1.0
        end
        men = men_create(db,db,0.0,"",gr )
        men_setchildinfo!(men,w,name,ratio)
    end
    return(men)
end

function men_setchild!(cur::Men, sid::Men)
    cur.child = sid
    sid.parent = cur
end

function men_resolvechild!(mentable)
    cur = mentable["MAIN"]

    while cur != nothing
        cnm = cur.cname
        typ = cur.ctype
        if cnm != ""
            if type_keywords[typ] == :insert
                sid = mentable[cnm]
                men_insert!(cur,sid)
            else
                if type_keywords[typ] == :merge
                    sid = men_end(mentable[cnm])
                else
                    sid = mentable[cnm]
                end
                men_setchild!(cur,sid)
            end
        end
        cur = cur.next
    end
end

"""Build mensur from text lines"""
function men_build(lns)
    mentable = Dict{String,Men}()
    grouplist = []
    groupnames = []
    cur = nothing

    for ln in lns
        # @show ln
        ln = eat_comment(ln)
        ln = replace(ln,r"[\t\s]+"=>"") # remove all white spaces

        if ln != ""
            if occursin("=",ln)
                # evaluate var definition line
                eval(Meta.parse(ln))
            else
                # normal df,db,r,comment line
                wd = String.(split(ln,","))

                w = wd[1] # df,db,r or group definition
                if w in group_keys
                    # group definitions
                    if group_keywords[w] == :endmain
                        grouplist = []
                        cur = nothing
                    elseif group_keywords[w] == :endgroup
                        pop!(grouplist)
                        if length(grouplist) == 0
                            cur = nothing
                        end
                    elseif group_keywords[w] == :main
                        "MAIN" in groupnames ? throw(AssertionError("Duplicating MAIN.")) : nothing
                        push!(grouplist,"MAIN")
                        push!(groupnames,"MAIN")
                    elseif group_keywords[w] == :group
                        push!(grouplist,wd[2])
                        gname = join(grouplist,":")
                        gname in groupnames ? throw(AssertionError("group name $group_name is not unique.")) : nothing
                        push!(groupnames,gname)
                    end
                elseif w in type_keys
                    # BRANCH, MERGE, SPLIT, etc
                    men = men_bykeyword(cur,wd)
                    cur = men_append!(cur,men)
                elseif length(wd) > 2
                    # normal "df,db,r,cmt" like string
                    df = eval(Meta.parse(wd[1]))
                    db = eval(Meta.parse(wd[2]))
                    r = eval(Meta.parse(wd[3]))
                    cmt = length(wd) > 3 ? wd[4] : ""

                    men = men_create(df*0.001, db*0.001,r*0.001, cmt, groupnames[end])

                    # set mensur group dictionary
                    if cur == nothing
                        cur = men
                        if length(mentable) == 0
                            push!(mentable,"MAIN"=>cur)
                        else
                            push!(mentable,groupnames[end] => cur)
                        end
                    else
                        cur = men_append!(cur,men)
                    end
                else
                    throw(AssertionError("string $ln is invalid."))
                end
            end
        end
    end
    # men_printtable(mentable) # debug
    men_resolvechild!(mentable)
    return(mentable)
end

function readxmen(path)
    lns = readlines(path)

    mentable = []
    if length(lns) > 0
        mentable = men_build(lns)
    end
    return(mentable)
end

function men_print(men::Men)
    gr = men.group
    if gr != "MAIN"
        println("GROUP,",gr)
    end

    while men != nothing
        df,db,r,c = getfbr(men)
        if men.child != nothing
            println(men.ctype,",",men.cname,",",men.cratio)
        else
            println(df*1000.0,",",db*1000.0,",",r*1000.0,",",c)
        end
        men = men.next
    end
    if gr != "MAIN"
        println("END_GROUP")
    end
end

function men_printtable(mentable)
    for (key, men) in pairs(mentable)
        if key == "MAIN"
            println(key)
            men_print(men)
            println("END_MAIN")
        else
            men_print(men)
        end
    end
end

function joint_mensur(men::Men)
    if isbranch(men.ctype)
        e = men_end(men.child)
        if e.parent != nothing && ismerge(e.parent.ctype)
            return(e.parent)
        else
            return(nothing)
        end
    else
        return(nothing)
    end
end

"""
create calculation parameters

Example:
    initcalcparam(minfreq,maxfreq,stepfreq,temperature,radiation)
    initcalcparam(0.0,2000.0,2.5,"PIPE/BAFFLE/NONE")
"""
function initcalcparam(;minfreq=0.0, maxfreq=2000.0, stepfreq=2.5, temperature=24.0, radiation="PIPE")
    c0 = 331.45 * sqrt(temperature / 273.16 + 1)
    rho = 1.2929 * (273.16 / (273.16 + temperature))
    rhoc0 = rho * c0
    mu = (18.2 + 0.0456*(temperature - 25)) * 1.0e-6  # viscosity constant. Linear approximation from Scientific Dictionary.
    nu = mu/rho  # dynamic viscous constant.
    dmp = (1+(GMM-1)/sqrt(PR)) # wall dumping factor

    Dict(
        :minfreq=>minfreq,
        :maxfreq=>maxfreq,
        :stepfreq=>stepfreq,
        :temperature=>temperature,
        :radiation=>radiation,
        :c0=>c0,
        :rho =>rho,
        :rhoc0 => rhoc0,
        :mu => mu,
        :nu => nu,
        :dmp => dmp
    )
end

function radimp(wf::Float64, dia::Float64; c0::Float64, rhoc0::Float64, radiation::String,kwd...)
    if wf > 0
        if dia == 0.0
            zr = complex(Inf) # closed end
        else
            if radiation == "NONE"
                zr = complex(0)  # simple open end impedance
            else
                s = dia^2*pi/4.0
                k = wf/c0
                x = k*dia

                re = rhoc0/s*(1 - besselj1(x)/x*2)  # 1st order bessel function.
                img = rhoc0/s*Struve.H(1,x)/x*2  # 1st order struve function.

                if radiation == "BAFFLE"
                    zr = complex(re,img)
                elseif radiation == "PIPE"
                    # real is about 0.5 times and imaginary is 0.7 times when without frange.
                    zr = complex(0.5*re,0.7*img)
                end
            end
        end
        return(zr)
    end
end

function calc_transmission(wf::Float64,men::Men; c0::Float64,rhoc0::Float64,nu::Float64,dmp::Float64,kwd...)
    df = men.df
    db = men.db
    r = men.r
    d = (df + db)*0.5
    aa = dmp * sqrt(2*wf*nu)/c0/d  # wall dumping factor
    k = sqrt((wf/c0)*(wf/c0 - 2*complex(-1,1)*aa))  # complex wave number including wall dumping
    x = k * r
    cc = cos(x)
    ss = sin(x)

    tm = zeros(ComplexF64,2,2)

    if df != db
        # taper
        r1 = df*0.5
        r2 = db*0.5
        dr = r2-r1

        tm[1,1] = (r2*x*cc - dr*ss)/(r1*x)
        tm[1,2] = im*rhoc0*ss/(pi*r1*r2)
        tm[2,1] = -im*pi*(dr*dr*x*cc - (dr*dr + x*x*r1*r2)*ss)/(x*x*rhoc0)
        tm[2,2] = (r1*x*cc + dr*ss)/(r2*x)
    else
        # straight
        s1 = pi/4*df*df
        tm[1,1] = tm[2,2] = cc
        tm[1,2] = im*rhoc0*ss/s1
        tm[2,1] = im*s1*ss/rhoc0
    end

    return(tm)
end

function zo2zi(tm::Array{ComplexF64,2},zo::ComplexF64)
    if !isinf(zo)
        zi = (tm[1,1]*zo + tm[1,2])/(tm[2,1]*zo + tm[2,2])
    else
        if tm[2,1] != 0
            zi = tm[1,1]/tm[2,1]
        else
            zi = complex(Inf)
        end
    end
    return(zi)
end

function transmission_matrix(men1::Men, men2::Union{Men,Nothing})
    if men2 == nothing
        men = men_end(men1)
        men = men.prev
    else
        men = men2
    end

    m = Matrix{ComplexF64}(I,2,2) # eye
    while men != nothing && men != men1
        m *= men.tm
        men = men.prev
    end
    m *= men1.tm

    return(m)
end

function child_impedance(wf::Float64,men::Men; params...)
    if issplit(men.ctype)
        # split (tonehole) type.
        # @show men.child
        impedance!(wf,men.child; params...)  # recursive call for input impedance
        if men.cratio == 0
            men.zo = men.next.zi
        else
            z1 = men.child.zi / men.cratio  # adjust blending ratio
            z2 = men.next.zi
            if z1 == complex(0.0) && z2 == complex(0.0)
                z = complex(0.0)
            else
                z = z1*z2/(z1+z2)
            end
            men.zo = z
        end
    elseif isbranch(men.ctype) && men.cratio > 0
        # multiple tube connection
        impedance!(wf, men.child;params...)
        m = transmission_matrix(men.child, nothing)
        jnt = joint_mensur(men)
        n = transmission_matrix(men.next, jnt)

        # section area adjustment
        if men.cratio == 1
            m[1,2] = complex(Inf)
        else
            m[1,2] /= (1 - men.cratio)
        end
        m[2,1] *= (1 - men.cratio)

        if men.cratio == 0
            n[1,2] = complex(Inf)
        else
            n[1,2] /= men.cratio
        end
        n[2,1] *= men.cratio

        z2 = jnt.next.zi
        dv = (m[2,2]*n[1,2] + m[1,2]*n[2,2] + (
            (m[1,2] + n[1,2])*(m[2,1] + n[2,1]) - (m[1,1] - n[1,1])*(m[2,2] - n[2,2]))*z2)
        if dv != complex(0.0)
            z = (m[1,2]*n[1,2] + (m[1,2]*n[1,1] + m[1,1]*n[1,2])*z2)/dv
        else
            z = 0
        end
        men.zo = z
    elseif isaddon(men.ctype) && men.cratio > 0
        # this routine will not called until 'ADDON(LOOP)' type of connection is implemented.
        impedance!(wf, men.child,params)
        m = transmission_matrix(men.child, nothing)
        z1 = m[1,2]/(m[1,2]*m[2,1]-(1-m[1,1])*(1-m[2,2]))
        z2 = men.next.zi
        if men.cratio == 0
            men.zo = men.next.zi
        elseif men.cratio == 1
            men.zo = z1
        else
            z1 /= men.cratio
            z2 /= (1 - men.cratio)
            if z1 == complex(0.0) && z2 == complex(0.0)
                z = complex(0.0)
            else
                z = z1*z2/(z1+z2)
            end
            men.zo = z
        end
    end
end

"""recursively called function to calculate impedance at each mensur node"""
function calc_impedance!(wf::Float64,men::Men; params...)
    if men.child != nothing
        child_impedance(wf, men; params...)
    elseif men.next != nothing
        men.zo = men.next.zi
    end

    if men.r > 0.0
        men.tm = calc_transmission(wf,men;params...)
        men.zi = zo2zi(men.tm, men.zo)
    else
        # length 0
        men.zi = men.zo
    end
end

"""calculate impedance at given frequency """
function impedance!(wf::Float64, men::Men; params...)
    if wf > 0
        cur = men_end(men)
        # end impedance
        cur.zo = radimp(wf,cur.df; params...)

        while cur != men
            calc_impedance!(wf,cur; params...)
            cur = cur.prev
        end
        calc_impedance!(wf,men; params...)
        return(men.zi)
    else
        return(complex(0.0))
    end
end

function input_impedance(mentable;params...)
    minfreq = params[:minfreq]
    maxfreq = params[:maxfreq]
    stepfreq = params[:stepfreq]

    imped = DataFrame()

    frq = range(minfreq, stop=maxfreq, step=stepfreq)
    imped[:frq] = frq
    imped[:imp] = zeros(ComplexF64,length(frq))

    men = mentable["MAIN"]

    ss = men.df^2 * pi/4

    for i in 1:length(frq)
        wf = frq[i]*2*pi
        imped[:imp][i] = ss * impedance!(wf,men;params...)
    end

    return(imped)
end

"Calculate input impedance for given xmensur file."
function calcimp(fpath::String; params...)
    mentable = readxmen(fpath)
    prm = initcalcparam(;params...)
    imped = input_impedance(mentable;prm...)
    imped[:mag] = 20log10.(abs.(imped[:imp]))
    return(imped)
end

"Reuse read data"
function calcimp(mtable::Dict{String,Men};params...)
    prm = initcalcparam(;params...)
    imped = input_impedance(mtable;prm...)
    imped[:mag] = 20log10.(abs.(imped[:imp]))
    return(imped)
end

function actualnext(men::Men)
    if men.ctype == "BRANCH" && men.cratio > HALF
        m = men.child
    elseif men.next != nothing
        m = men.next
    else
        m = men.parent
    end
    m
end

function slice_mensur!(men::Men, step::Float64)
    if step > 0
        while men != nothing
            if men.r > step
                if men.db != men.df
                    db2 = (men.db - men.df)/men.r * step + men.df
                else
                    db2 = men.db
                end
                r2 = men.r - step
                nw = men_create(db2, men.db, r2, men.cm, men.gr)
                men_insert!(men, nw)
                # change self data
                men.r = step
                men.db = db2
            end
            men = actualnext(men)
        end
    end
end

function calc_pressure(p::Float64,men::Men)
    v = complex([p,0.0]) # initial pressure, volume speed pair
    x = 0.0 # length from top
    prs = DataFrame(x=[],D=[],p=[],U=[]) # x, dia, pressure, volume velosity
    push!(prs,(x,men.df, v[1],v[2]))

    while men != nothing
        # @show men.tm
        if men.child != nothing
            push!(prs, (x,0.0, v[1],v[2]))
        end
        if men.r > 0
            v = inv(men.tm)*v
        end
        x += men.r
        push!(prs,(x,men.df,v[1],v[2]))
        men = actualnext(men)
    end
    return prs
end

function calcprs(mentable, frequency::Float64; slice=1.0, temperature=24.0, pressure = 2.0e-2)
    slice_mensur!(mentable["MAIN"], slice/1000.0) # unit conv here

    # calc impedance first
    prm = initcalcparam(temperature=temperature)
    wf = 2pi*frequency
    men = mentable["MAIN"]
    impedance!(wf,men;prm...)

    prs = calc_pressure(pressure,men) # return (x,p,U) dataframe
    prs[:x] .*= 1000.0 # x unit conv
    prs[:D] .*= 1000.0 # dia unit conv
    return prs
end

"""
    calcprs(filepath, frequency; kwargs...)

    Calculate pressure distribution for given frequency.
    Mensur will be sliced by `slice=` value (default 1mm).
    Pressure at end must be supplied (default = 2.0e-2 Pa = 60dBSPL).

    Returns dataframe of (x, D, p, U),
    where x:length from top, D,p,U are diameter,pressure,volume velosity at x.
"""
function calcprs(fpath::String,frequency::Float64; params...)
    mentable = readxmen(fpath)
    return calcprs(mentable,frequency; params...)
end


# 2D LSRTM/FWI parallelization test with 2 sources and 2 vintages
#
# Ziyi Yin, ziyi.yin@gatech.edu
# August 2021

parsed_args = parse_commandline()

nlayer = parsed_args["nlayer"]
tti = parsed_args["tti"]
fs =  parsed_args["fs"]

### Model
model, model0, dm = setup_model(tti, 4)
q, srcGeometry, recGeometry, info = setup_geom(model; nsrc=4)
q1 = q[[1,4]]
q2 = q[[2,3]]
srcGeometry1 = subsample(srcGeometry,[1,4])
srcGeometry2 = subsample(srcGeometry,[2,3])
recGeometry1 = subsample(recGeometry,[1,4])
recGeometry2 = subsample(recGeometry,[2,3])
info1 = subsample(info,[1,4])
info2 = subsample(info,[2,3])

opt = Options(sum_padding=true, free_surface=fs)
F1 = judiModeling(info1, model, srcGeometry1, recGeometry1; options=opt)
F2 = judiModeling(info2, model, srcGeometry2, recGeometry2; options=opt)

# Observed data
dobs1 = F1*q1
dobs2 = F2*q2

@testset "FWI/LSRTM objective multi-level parallelization test with $(nlayer) layers and tti $(tti) and freesurface $(fs)" begin

    ftol = 1f-5
    dm1 = 2f0*circshift(dm, 10)
    dm2 = 2f0*circshift(dm, 30)

    Jm0, grad = lsrtm_objective([model0, model0], [q1, q2], [dobs1, dobs2], [dm1, dm2]; options=opt, nlind=true)
    Jm01, grad1 = lsrtm_objective(model0, q1, dobs1, dm1; options=opt, nlind=true)
    Jm02, grad2 = lsrtm_objective(model0, q2, dobs2, dm2; options=opt, nlind=true)
    @test isapprox(Jm01+Jm02, Jm0; rtol=ftol)
    @test isapprox(grad1, grad[1]; rtol=ftol)
    @test isapprox(grad2, grad[2]; rtol=ftol)

    _Jm0, _grad = fwi_objective([model0, model0], [q1, q2], [dobs1, dobs2]; options=opt)
    _Jm01, _grad1 = fwi_objective(model0, q1, dobs1; options=opt)
    _Jm02, _grad2 = fwi_objective(model0, q2, dobs2; options=opt)
    @test isapprox(_Jm01+_Jm02, _Jm0; rtol=ftol)
    @test isapprox(_grad1, _grad[1]; rtol=ftol)
    @test isapprox(_grad2, _grad[2]; rtol=ftol)

end


"""
```
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


######################################################################## # Codes for the paper:
# An Anthropocene Map of Genetic Diversity
#
# Andreia Miraldo, Sen Li, Michael K. Borregaard, Alexander Floréz-Rodriguéz, Shyam Gopalakrishnan, Mirneza Risvanovic, Zhiheng Wang, Carsten Rahbek, Katharine A. Marske & David Nogués-Bravo
#
# Science, 2016
# Code in this file by Sen Li and Michael K. Borregaard
#########################################################################

# Load the necessary library and functions
include("Compare_Pairwise_Function.jl")
using DataFrames

"""
A function to calculate the summary statistic for all sites (e.g. grid cell or biome) for one species.

**Parameters**
* 'species':        A string with the name of the species
* 'species_seqs':   A matrix where the rows are aligned genetic sequences, and columns are loci. Basepairs must be coded as 1, 2, 3 or 4, or with a 0 signifying that the locus is absent from the alignment.
* 'sitenames':      A vector of strings with the names of the sites
"""
function calcspecies(species::String, species_seqs::Matrix{Int}, sitenames::Vector{String})
    res = DataFrame(species = String[], cell = String[], seq1 = Int[], seq2 = Int[], length_seq1 = Int[], length_seq2 = Int[], overlap = Float64[], commons = Int[], num_per_bp = Float64[])
    uniq_grid = unique(sitenames, 1)
    for iter_grid = 1:size(uniq_grid,1) # Go through each site in turn
        seqs_grid = species_seqs[findin(sitenames, uniq_grid[iter_grid, :]), :]

        tot_muts = compare_pairwise(seqs_grid)
        lns = size(tot_muts, 1)
        tmp = DataFrame(species = repeat([species], inner= lns), cell = repeat(vec(uniq_grid[iter_grid, :]), inner = lns)) #Expand species and cell names to the length of the resulting DataFrame
        tot_muts = hcat(tmp, tot_muts)

        res = vcat(res, tot_muts)
    end

    for sym in [:seq1, :seq2, :length_seq1, :length_seq2, :overlap, :commons, :num_per_bp]
        res[sym][res[sym] .< 0.] = NA   #Replace empty values with an NA identifier
    end
    res
end

"""
A function to create master data matrices that are used to compute genetic diversity, assess data quality and do sensitivity analyses.

**Parameters**
* 'foldername' : The name of of a folder containing the data files. There must be files of two types (file extension 'fasta' and file extension 'coords') with  the same filename, e.g. the species names (e.g. folder contents could be 'Bufo_bufo.fasta, Bufo_bufo.coords, Rana_arvalis.fasta, Rana_arvalis.coords', etc.). The .fasta files contain the sequences as an m x n integer matrix, where m is the number of sequences and n is the length of the alignments. The .coords files contain the geographic coordinates of the sequences, as an m x 2 floating point matrix with latitude in the first column and longitude in the second.  TODO: A folder with the anthrome and biome of each sequence
"""
function create_master_matrices(foldername::String, sitesfolder::String)
    species_list = [x[1:(end-6)] for x in filter(st->contains(st, ".fasta"), readdir(foldername))]      #identify unique file names ignoring the extension
    num_files = size(species_list,1)

    #equalarea = latbands = biomes = anthrome = gridcells = DataFrame(species = String[], cell = String[], seq1 = Int[], seq2 = Int[], length_seq1 = Int[], length_seq2 = Int[], overlap = Float64[], commons = Int[], num_per_bp = Float64[])   # Pre-initialize the DataFrame to ensure correct element types
    ret = DataFrame(species = String[], cell = String[], seq1 = Int[], seq2 = Int[], length_seq1 = Int[], length_seq2 = Int[], overlap = Float64[], commons = Int[], num_per_bp = Float64[])   # Pre-initialize the DataFrame to ensure correct element types

    for iter_file = 1:num_files
        # read in the sequence data for one species
        file_name_seq = joinpath(foldername, species_list[iter_file] * ".fasta")
        species_seqs = readdlm(file_name_seq, Int)
        length_seq = size(species_seqs,2)

        # read the coordinates for one species
        file_name_sites = joinpath(foldername, species_list[iter_file] * ".coords")
        species_sites = readdlm(file_name_sites, String)[:,1]
        #if eltype(species_sites) <: AbstractFloat
        #    species_sites = floor(Int, species_sites)
        #end

        ret = vcat(ret, calcspecies(species_list[iter_file], species_seqs, species_sites))
        println(iter_file)

    end
    writetable("pairwise_$(sitesfolder)_$(foldername).csv", ret)
end

function grid_coordinates(foldername = ".", xsize = 1, ysize = xsize, lowerleftoffset = (0,0))
    in("gridcells", readdir(foldername)) && error("There is already a folder named gridcells!")
    mkdir(joinpath(foldername, "gridcells"))
    species_list = [x[1:(end-7)] for x in filter(st->contains(st, ".coords"), readdir(foldername))]
    for species in species_list
        species_coords = readdlm(joinpath(foldername,"$(species).coords"))
        coords_grid = hcat(ceil(Int, (species_coords[:,1] + lowerleftoffset[1]) ./ xsize) .* xsize .- lowerleftoffset[1] .- 0.5xsize, ceil(Int, (species_coords[:,2] + lowerleftoffset[2]) ./ ysize) .* ysize .- lowerleftoffset[2] .- 0.5ysize)
        sitenames = mapslices(x -> "$(x[1])_$(x[2])", coords_grid, 2)
        writedlm(joinpath(foldername, "gridcells", species), sitenames)
    end
end

function latband_coordinates(foldername = ".", bandwidth = 1)
    in("latbands", readdir(foldername)) && error("There is already a folder named latbands!")
    mkdir(joinpath(foldername, "latbands"))
    species_list = [x[1:(end-7)] for x in filter(st->contains(st, ".coords"), readdir(foldername))]
    for species in species_list
        species_coords = readdlm(joinpath(foldername,"$(species).coords"))
        latband = bandwidth * floor(species_coords[:,1]/bandwidth)
        writedlm(joinpath(foldername, "latbands", species), latband)
    end
end

Miraldo et al. 2016/Create_Master_Matrices.jl


@testset "High-level operations on distributed data" begin

    W = addprocs(2)
    @everywhere using GigaSOM

    Random.seed!(1)
    dd = rand(11111, 5)
    buckets = rand([1, 2, 3], 11111)

    di1 = scatter_array(:test1, dd, W[1:1])
    di2 = scatter_array(:test2, dd, W)
    buckets1 = scatter_array(:buckets1, buckets, W[1:1])
    buckets2 = scatter_array(:buckets2, buckets, W)

    dc = gather_array(di1)
    dc[:, 1:2] = asinh.(dc[:, 1:2] ./ 1.23)
    dtransform_asinh(di1, [1, 2], 1.23)
    dtransform_asinh(di2, [1, 2], 1.23)

    @testset "dtransform_asinh" begin
        @test isapprox(gather_array(di1), dc)
        @test isapprox(gather_array(di2), dc)
    end

    unscatter(di1)
    unscatter(di2)
    rmprocs(W)
end


struct JointToLinkPosition{Rn,R3,S,JC} <: TaskMap{Rn,R3,S}
    state_cache::StateCache{S,JC}
    link::RigidBody{S}
end
struct JointToLinkPositionT{Rn,R3,S,JC} <: TaskMapT{Rn,R3,S}
    base_map::JointToLinkPosition{Rn,R3,S,JC}
end
JointToLinkPosition{ℝ{n},R3}(s::StateCache{S,JC}, l::RigidBody{S}) where {n,R3,S,JC} =
    JointToLinkPosition{ℝ{n},R3,S,JC}(s,l)

function task_map_emb(::EmbRep, ::EmbRep, xme, task_map::JointToLinkPosition{ℝ{n},R3}) where n
    state = task_map.state_cache[eltype(xme)]
    set_configuration!(state, xme)
    translation(transform_to_root(state, task_map.link))
end