text:
#=

This file contains functions that help convert a parameters YAML file into a vSmartMOM_Parameters
object. This process includes converting data types, specifying default values, validating
the input, and producing the output object.

=#

"Check that a field exists in yaml file"
function check_yaml_field(dict::Dict{Any, Any}, full_keys::Array{String}, curr_keys::Array{String}, final_type::Type, valid_options::Array{String})

    # Should have at least one key
    @assert length(curr_keys) >= 1

    # Only one key (that should exist)
    if length(curr_keys) == 1
        @assert curr_keys[1] in keys(dict) "Missing key in parameters yaml: $(join(full_keys, '/'))"
        @assert dict[curr_keys[1]] isa final_type "Improper type for $(join(full_keys, '/')); must be a $(final_type)"
        if length(valid_options) > 0
            @assert dict[curr_keys[1]] in valid_options "Field $(join(full_keys, '/')) must be one of $(valid_options)"
        end
        return true

    # Recursive case, check sub-dictionary
    else
        @assert curr_keys[1] in keys(dict) "Missing key in parameters yaml: $(join(full_keys, '/'))"
        return check_yaml_field(dict[curr_keys[1]], full_keys, curr_keys[2:end], final_type, valid_options)
    end
end

"Given a list of aerosol parameter dictionaries, validate all of them"
function validate_aerosols(aerosols::Union{Array{Dict{Any, Any}}, Vector{Any}})

    @assert !isempty(aerosols) "Aerosols list shouldn't be empty if scattering block is included"

    fields = [(["τ_ref"], Real),
              (["μ"], Real),
              (["σ"], Real),
              (["nᵣ"], Real),
              (["nᵢ"], Real),
              (["p₀"], Real),
              (["p₀"], Real)]

    for aerosol in aerosols
        for i in 1:length(fields)
            key, elty = fields[i][1:2]
            @assert check_yaml_field(aerosol, key, key, elty, Array{String}([]))
        end
    end
end

"Convert the input dictionary of aerosols into a list of RT_aerosols"
function aerosol_params_to_obj(aerosols::Union{Array{Dict{Any, Any}}, Vector{Any}}, FT)

    rt_aerosol_obj_list = RT_Aerosol[]

    for aerosol in aerosols

        size_distribution = LogNormal(FT(aerosol["μ"]), FT(aerosol["σ"]))

        new_aerosol_obj = Aerosol(size_distribution,
                                  FT(aerosol["nᵣ"]),
                                  FT(aerosol["nᵢ"]))

        new_rt_aerosol_obj = RT_Aerosol(new_aerosol_obj, FT(aerosol["τ_ref"]), FT(aerosol["p₀"]), FT(aerosol["σp"]))

        push!(rt_aerosol_obj_list, new_rt_aerosol_obj)
    end

    return rt_aerosol_obj_list
end

"Check that the vmr's in the atmospheric profile match the molecules in the parameters"
function validate_vmrs(molecules::Array, vmr::Dict)
    for molec in unique(vcat(molecules...))
        @assert molec in keys(vmr) "$(molec) listed as molecule in parameters yaml, but no vmr given in atmospheric profile"
        @assert vmr[molec] isa Real || vmr[molec] isa Vector "The vmr for $(molec) in the atmospheric profile must either be a real-valued number, or an array of nodal points from surface to 0hPa (TOA)"
    end
end

"Given a parameter dictionary from a YAML file, validate the dictionary"
function validate_yaml_parameters(params)

    # All the fields that *should* exist and their types
    fields = [

              # radiative_transfer group
              (["radiative_transfer", "spec_bands"], Array{String}),
              (["radiative_transfer", "surface"], Array{String}),
              (["radiative_transfer", "quadrature_type"], String),
              (["radiative_transfer", "polarization_type"], String),
              (["radiative_transfer", "max_m"], Integer),
              (["radiative_transfer", "Δ_angle"], Real),
              (["radiative_transfer", "l_trunc"], Integer),
              (["radiative_transfer", "depol"], Real),
              (["radiative_transfer", "float_type"], String),
              (["radiative_transfer", "architecture"], String),

              # geometry group
              (["geometry", "sza"], Real),
              (["geometry", "vza"], Array{<:Real}),
              (["geometry", "vaz"], Array{<:Real}),
              (["geometry", "obs_alt"], Real),

              # atmospheric_profile group
              (["atmospheric_profile", "T"], Array{<:Real}),
              (["atmospheric_profile", "p"], Array{<:Real}),
              (["atmospheric_profile", "profile_reduction"], Union{Integer, Nothing}),

              # absorption group
              (["absorption", "molecules"], Array),
              (["absorption", "vmr"], Dict),
              (["absorption", "broadening"], String, ["Voigt()", "Doppler()", "Lorentz()"]),
              (["absorption", "CEF"], String, ["HumlicekWeidemann32SDErrorFunction()"]),
              (["absorption", "wing_cutoff"], Integer),

              # scattering group
              (["scattering", "aerosols"], Array),
              (["scattering", "r_max"], Real),
              (["scattering", "nquad_radius"], Real),
              (["scattering", "λ_ref"], Real),
              (["scattering", "decomp_type"], String, ["NAI2()", "PCW()"]),
              ]

    # Check that every field exists and is correctly typed
    for i in 1:length(fields)
        key, elty = fields[i][1:2]
        valid_options = length(fields[i]) == 3 ? fields[i][3] : Array{String}([])

        # Check if the section exists in the YAML file and is an optional section
        if ("absorption" in keys(params) && key[1] == "absorption") ||
           ("scattering" in keys(params) && key[1] == "scattering") ||
           (key[1] != "absorption" && key[1] != "scattering")
            @assert check_yaml_field(params, key, key, elty, valid_options)
        end
    end

    # Check the aerosols separately if specified
    if "scattering" in keys(params)
        validate_aerosols(params["scattering"]["aerosols"])
    end
end

"Given a path to a YAML parameters file, load all the parameters into a vSmartMOM_Parameters struct"
function parameters_from_yaml(file_path)

    # Load the YAML file
    params_dict = YAML.load_file(file_path)

    # Validate to make sure it has all necessary fields
    validate_yaml_parameters(params_dict)

    # #########################################################
    # Evaluate type-specifying fields to get their proper types
    # #########################################################

    # radiative_transfer group
    spec_bands = convert.(Array, map(x -> collect(eval(Meta.parse(x))), params_dict["radiative_transfer"]["spec_bands"]))
    BRDF_per_band = map(x -> eval(Meta.parse(x)), params_dict["radiative_transfer"]["surface"])
    quadrature_type = eval(Meta.parse(params_dict["radiative_transfer"]["quadrature_type"]))
    polarization_type = eval(Meta.parse(params_dict["radiative_transfer"]["polarization_type"]))
    FT = eval(Meta.parse(params_dict["radiative_transfer"]["float_type"]))
    architecture = eval(Meta.parse(params_dict["radiative_transfer"]["architecture"]))

    # atmospheric_profile group
    T = convert.(FT, params_dict["atmospheric_profile"]["T"]) # Level
    p = convert.(FT, params_dict["atmospheric_profile"]["p"]) # Boundaries
    q = "q" in keys(params_dict["atmospheric_profile"]) ? # Specific humidity, if it's specified.
            convert.(Float64, params_dict["atmospheric_profile"]["q"]) : zeros(length(T)) # Otherwise 0

    # absorption group
    if "absorption" in keys(params_dict)
        molecules = Array(params_dict["absorption"]["molecules"])
        vmr = convert(Dict{String, Union{Real, Vector}}, params_dict["absorption"]["vmr"])
        validate_vmrs(params_dict["absorption"]["molecules"], vmr)
        broadening_function = eval(Meta.parse(params_dict["absorption"]["broadening"]))
        CEF = eval(Meta.parse(params_dict["absorption"]["CEF"]))
        wing_cutoff = params_dict["absorption"]["wing_cutoff"]

        absorption_params = AbsorptionParameters(molecules, vmr, broadening_function,
                                                 CEF, wing_cutoff)
    else
        absorption_params = nothing
    end

    # scattering group
    if "scattering" in keys(params_dict)

        # Get aerosols from types
        aerosols = aerosol_params_to_obj(params_dict["scattering"]["aerosols"], FT)
        r_max = FT(params_dict["scattering"]["r_max"])
        nquad_radius = params_dict["scattering"]["nquad_radius"]
        λ_ref = FT(params_dict["scattering"]["λ_ref"])
        decomp_type = eval(Meta.parse(params_dict["scattering"]["decomp_type"]))

        scattering_params = ScatteringParameters(aerosols, r_max, nquad_radius,
                                                 λ_ref, decomp_type)
    else
        scattering_params = nothing
    end

    return vSmartMOM_Parameters(
                                # radiative_transfer group
                                spec_bands,
                                BRDF_per_band,
                                quadrature_type,
                                polarization_type,
                                params_dict["radiative_transfer"]["max_m"],
                                FT(params_dict["radiative_transfer"]["Δ_angle"]),
                                params_dict["radiative_transfer"]["l_trunc"],
                                FT(params_dict["radiative_transfer"]["depol"]),
                                FT,
                                architecture,

                                # geometry group
                                FT(params_dict["geometry"]["sza"]),
                                convert.(FT, params_dict["geometry"]["vza"]),
                                convert.(FT, params_dict["geometry"]["vaz"]),
                                FT(params_dict["geometry"]["obs_alt"]),

                                # atmospheric_profile group
                                T, p, q,
                                params_dict["atmospheric_profile"]["profile_reduction"],

                                # absorption group
                                absorption_params,

                                # scattering group
                                scattering_params
                                );
end


text:
module _Simple

include("binarysplitting.jl")

const c3_over_24 = Int64(640320)^3 ÷ 24
const digits_per_term = log10(c3_over_24 / (6 * 2 * 6))
const guard_terms = 5

# Compute int(pi * 10^digits)
# This is done using Chudnovsky's series with binary splitting
pi(digits::T) where {T<:Integer} = begin
    # how many terms to compute
    n = ceil(T, digits / digits_per_term) + guard_terms

    # calculate P(0, N), Q(0, N) and T(0, N)
    p, q, t = binarysplitting(zero(n), n)

    one_squared = big(10)^(2*digits)
    sqrtC = isqrt(10005 * one_squared)

    (q * 426880 * sqrtC) ÷ t
end

#
# Source: http://www.craig-wood.com/nick/articles/pi-chudnovsky/
#
# Computes the terms for binary splitting the Chudnovsky infinite series
#
# a(a) = +/- (13591409 + 545140134*a)
# p(a) = (6*a-5)*(2*a-1)*(6*a-1)
# b(a) = 1
# q(a) = a*a*a*C3_OVER_24
#
compute(a, b) = compute(convert(BigInt, a), convert(BigInt, b))
compute(a::BigInt, b::BigInt) = begin
    if a == 0
        p = q = big(1)
    else
        a2 = a * 2
        a6 = a2 * 3
        p = (a6 - 5) * (a2 - 1) * (a6 - 1)
        q = a^3 * c3_over_24
    end

    # t(a) = p(a) * a(a)
    t = p * (13591409 + 545140134 * a)
    isodd(a) && (t = -t)

    p, q, t
end

combine(am, mb) = begin
    p₁, q₁, t₁ = am
    p₂, q₂, t₂ = mb

    p = p₁ * p₂
    q = q₁ * q₂
    t = q₂ * t₁ + p₁ * t₂

    p, q, t
end

end

ext: jl


text:
@testset "Laguerre integration" begin
    # forward transform
    # this is just Lnh for n=2.
    # need to get a2 = 1 and other coeffs 0.
    f(x) = (x^2 - 4 * x + 2) * exp(-x / 2) / 2
    N = 3
    an = laguerre_transform(f, N)
    anexact = [0; 0; 1; 0]
    @test all(isapprox.(an, anexact, atol = 1e-8))

    # now inverse transform
    x = 0:0.1:10 |> collect
    fv = inverse_laguerre_transform(an, x)
    @test norm(fv - f.(x)) < 1e-10

    # test integration using Laguerre-Gauss-Radau quadrature.
    f1(x) = sin(x)^2 * exp(-x)  # -> 2/5
    N = 30
    @test isapprox(LGRquad(f1, N), 2 / 5, atol = 1e-3, rtol = 1e-4)
end


text:
"""
    GlmResp

The response vector and various derived vectors in a generalized linear model.
"""
struct GlmResp{V<:FPVector,D<:UnivariateDistribution,L<:Link} <: ModResp
    "`y`: response vector"
    y::V
    d::D
    "`devresid`: the squared deviance residuals"
    devresid::V
    "`eta`: the linear predictor"
    eta::V
    "`mu`: mean response"
    mu::V
    "`offset:` offset added to `Xβ` to form `eta`.  Can be of length 0"
    offset::V
    "`wts:` prior case weights.  Can be of length 0."
    wts::V
    "`wrkwt`: working case weights for the Iteratively Reweighted Least Squares (IRLS) algorithm"
    wrkwt::V
    "`wrkresid`: working residuals for IRLS"
    wrkresid::V
end

function GlmResp(y::V, d::D, l::L, η::V, μ::V, off::V, wts::V) where {V<:FPVector, D, L}
    n  = length(y)
    nη = length(η)
    nμ = length(μ)
    lw = length(wts)
    lo = length(off)

    # Check y values
    checky(y, d)

    # Lengths of y, η, and η all need to be n
    if !(nη == nμ == n)
        throw(DimensionMismatch("lengths of η, μ, and y ($nη, $nμ, $n) are not equal"))
    end

    # Lengths of wts and off can be either n or 0
    if lw != 0 && lw != n
        throw(DimensionMismatch("wts must have length $n or length 0 but was $lw"))
    end
    if lo != 0 && lo != n
        throw(DimensionMismatch("offset must have length $n or length 0 but was $lo"))
    end

    return GlmResp{V,D,L}(y, d, similar(y), η, μ, off, wts, similar(y), similar(y))
end

function GlmResp(y::V, d::D, l::L, off::V, wts::V) where {V<:FPVector,D,L}
    η   = similar(y)
    μ   = similar(y)
    r   = GlmResp(y, d, l, η, μ, off, wts)
    initialeta!(r.eta, d, l, y, wts, off)
    updateμ!(r, r.eta)
    return r
end

deviance(r::GlmResp) = sum(r.devresid)

"""
    cancancel(r::GlmResp{V,D,L})

Returns `true` if dμ/dη for link `L` is the variance function for distribution `D`

When `L` is the canonical link for `D` the derivative of the inverse link is a multiple
of the variance function for `D`.  If they are the same a numerator and denominator term in
the expression for the working weights will cancel.
"""
cancancel(::GlmResp) = false
cancancel(::GlmResp{V,D,LogitLink}) where {V,D<:Union{Bernoulli,Binomial}} = true
cancancel(::GlmResp{V,D,NegativeBinomialLink}) where {V,D<:NegativeBinomial} = true
cancancel(::GlmResp{V,D,IdentityLink}) where {V,D<:Normal} = true
cancancel(::GlmResp{V,D,LogLink}) where {V,D<:Poisson} = true

"""
    updateμ!{T<:FPVector}(r::GlmResp{T}, linPr::T)

Update the mean, working weights and working residuals, in `r` given a value of
the linear predictor, `linPr`.
"""
function updateμ! end

function updateμ!(r::GlmResp{T}, linPr::T) where T<:FPVector
    isempty(r.offset) ? copyto!(r.eta, linPr) : broadcast!(+, r.eta, linPr, r.offset)
    updateμ!(r)
    if !isempty(r.wts)
        map!(*, r.devresid, r.devresid, r.wts)
        map!(*, r.wrkwt, r.wrkwt, r.wts)
    end
    r
end

function updateμ!(r::GlmResp{V,D,L}) where {V<:FPVector,D,L}
    y, η, μ, wrkres, wrkwt, dres = r.y, r.eta, r.mu, r.wrkresid, r.wrkwt, r.devresid

    @inbounds for i in eachindex(y, η, μ, wrkres, wrkwt, dres)
        μi, dμdη = inverselink(L(), η[i])
        μ[i] = μi
        yi = y[i]
        wrkres[i] = (yi - μi) / dμdη
        wrkwt[i] = cancancel(r) ? dμdη : abs2(dμdη) / glmvar(r.d, μi)
        dres[i] = devresid(r.d, yi, μi)
    end
end

function _weights_residuals(yᵢ, ηᵢ, μᵢ, omμᵢ, dμdηᵢ, l::LogitLink)
    # LogitLink is the canonical link function for Binomial so only wrkresᵢ can
    # possibly fail when dμdη==0 in which case it evaluates to ±1.
    if iszero(dμdηᵢ)
        wrkresᵢ = ifelse(yᵢ == 1, one(μᵢ), -one(μᵢ))
    else
        wrkresᵢ = ifelse(yᵢ == 1, omμᵢ, yᵢ - μᵢ) / dμdηᵢ
    end
    wrkwtᵢ = μᵢ*omμᵢ

    return wrkresᵢ, wrkwtᵢ
end

function _weights_residuals(yᵢ, ηᵢ, μᵢ, omμᵢ, dμdηᵢ, l::ProbitLink)
    # Since μomμ will underflow before dμdη for Probit, we can just check the
    # former to decide when to evaluate with the tail approximation.
    μomμᵢ = μᵢ*omμᵢ
    if iszero(μomμᵢ)
        wrkresᵢ = 1/abs(ηᵢ)
        wrkwtᵢ  = dμdηᵢ
    else
        wrkresᵢ = ifelse(yᵢ == 1, omμᵢ, yᵢ - μᵢ) / dμdηᵢ
        wrkwtᵢ  = abs2(dμdηᵢ)/μomμᵢ
    end

    return wrkresᵢ, wrkwtᵢ
end

function _weights_residuals(yᵢ, ηᵢ, μᵢ, omμᵢ, dμdηᵢ, l::CloglogLink)
    if yᵢ == 1
        wrkresᵢ = exp(-ηᵢ)
    else
        emη = exp(-ηᵢ)
        if iszero(emη)
            # Diverges to -∞
            wrkresᵢ = -typeof(wrkresᵢ)(Inf)
        elseif isinf(emη)
            # converges to -1
            wrkresᵢ = -one(emη)
        else
            wrkresᵢ = (yᵢ - μᵢ)/omμᵢ*emη
        end
    end

    wrkwtᵢ  = exp(2*ηᵢ)/expm1(exp(ηᵢ))
    # We know that both limits are zero so we'll convert NaNs
    wrkwtᵢ = ifelse(isnan(wrkwtᵢ), zero(wrkwtᵢ), wrkwtᵢ)

    return wrkresᵢ, wrkwtᵢ
end

# Fallback for remaining link functions
function _weights_residuals(yᵢ, ηᵢ, μᵢ, omμᵢ, dμdηᵢ, l::Link01)
    wrkresᵢ = ifelse(yᵢ == 1, omμᵢ, yᵢ - μᵢ)/dμdηᵢ
    wrkwtᵢ  = abs2(dμdηᵢ)/(μᵢ*omμᵢ)

    return wrkresᵢ, wrkwtᵢ
end

function updateμ!(r::GlmResp{V,D,L}) where {V<:FPVector,D<:Union{Bernoulli,Binomial},L<:Link01}
    y, η, μ, wrkres, wrkwt, dres = r.y, r.eta, r.mu, r.wrkresid, r.wrkwt, r.devresid

    @inbounds for i in eachindex(y, η, μ, wrkres, wrkwt, dres)
        yᵢ, ηᵢ = y[i], η[i]
        μᵢ, omμᵢ, dμdηᵢ = inverselink(L(), ηᵢ)
        μ[i] = μᵢ
        # For large values of ηᵢ the quantities dμdη and μomμ will underflow.
        # The ratios defining (yᵢ - μᵢ)/dμdη and dμdη^2/μomμ have fairly stable
        # tail behavior so we can switch algorithm to avoid 0/0. The behavior
        # is specific to the link function so _weights_residuals dispatches to
        # robust versions for LogitLink and ProbitLink
        wrkres[i], wrkwt[i] = _weights_residuals(yᵢ, ηᵢ, μᵢ, omμᵢ, dμdηᵢ, L())
        dres[i] = devresid(r.d, yᵢ, μᵢ)
    end
end

function updateμ!(r::GlmResp{V,D,L}) where {V<:FPVector,D<:NegativeBinomial,L<:NegativeBinomialLink}
    y, η, μ, wrkres, wrkwt, dres = r.y, r.eta, r.mu, r.wrkresid, r.wrkwt, r.devresid

    @inbounds for i in eachindex(y, η, μ, wrkres, wrkwt, dres)
        θ = r.d.r # the shape parameter of the negative binomial distribution
        μi, dμdη, μomμ = inverselink(L(θ), η[i])
        μ[i] = μi
        yi = y[i]
        wrkres[i] = (yi - μi) / dμdη
        wrkwt[i] = dμdη
        dres[i] = devresid(r.d, yi, μi)
    end
end

"""
    wrkresp(r::GlmResp)

The working response, `r.eta + r.wrkresid - r.offset`.
"""
wrkresp(r::GlmResp) = wrkresp!(similar(r.eta), r)

"""
    wrkresp!{T<:FPVector}(v::T, r::GlmResp{T})

Overwrite `v` with the working response of `r`
"""
function wrkresp!(v::T, r::GlmResp{T}) where T<:FPVector
    broadcast!(+, v, r.eta, r.wrkresid)
    isempty(r.offset) ? v : broadcast!(-, v, v, r.offset)
end

abstract type AbstractGLM <: LinPredModel end

mutable struct GeneralizedLinearModel{G<:GlmResp,L<:LinPred} <: AbstractGLM
    rr::G
    pp::L
    fit::Bool
end

function coeftable(mm::AbstractGLM; level::Real=0.95)
    cc = coef(mm)
    se = stderror(mm)
    zz = cc ./ se
    p = 2 * ccdf.(Ref(Normal()), abs.(zz))
    ci = se*quantile(Normal(), (1-level)/2)
    levstr = isinteger(level*100) ? string(Integer(level*100)) : string(level*100)
    CoefTable(hcat(cc,se,zz,p,cc+ci,cc-ci),
              ["Estimate","Std. Error","z value","Pr(>|z|)","Lower $levstr%","Upper $levstr%"],
              ["x$i" for i = 1:size(mm.pp.X, 2)], 4)
end

function confint(obj::AbstractGLM; level::Real=0.95)
    hcat(coef(obj),coef(obj)) + stderror(obj)*quantile(Normal(),(1. -level)/2.)*[1. -1.]
end

deviance(m::AbstractGLM) = deviance(m.rr)

function loglikelihood(m::AbstractGLM)
    r   = m.rr
    wts = r.wts
    y   = r.y
    mu  = r.mu
    d   = r.d
    ll  = zero(eltype(mu))
    if length(wts) == length(y)
        ϕ = deviance(m)/sum(wts)
        @inbounds for i in eachindex(y, mu, wts)
            ll += loglik_obs(d, y[i], mu[i], wts[i], ϕ)
        end
    else
        ϕ = deviance(m)/length(y)
        @inbounds for i in eachindex(y, mu)
            ll += loglik_obs(d, y[i], mu[i], 1, ϕ)
        end
    end
    ll
end

dof(x::GeneralizedLinearModel) = dispersion_parameter(x.rr.d) ? length(coef(x)) + 1 : length(coef(x))

function _fit!(m::AbstractGLM, verbose::Bool, maxiter::Integer, minstepfac::Real,
               atol::Real, rtol::Real, start)

    # Return early if model has the fit flag set
    m.fit && return m

    # Check arguments
    maxiter >= 1       || throw(ArgumentError("maxiter must be positive"))
    0 < minstepfac < 1 || throw(ArgumentError("minstepfac must be in (0, 1)"))

    # Extract fields and set convergence flag
    cvg, p, r = false, m.pp, m.rr
    lp = r.mu

    # Initialize β, μ, and compute deviance
    if start == nothing || isempty(start)
        # Compute beta update based on default response value
        # if no starting values have been passed
        delbeta!(p, wrkresp(r), r.wrkwt)
        linpred!(lp, p)
        updateμ!(r, lp)
        installbeta!(p)
    else
        # otherwise copy starting values for β
        copy!(p.beta0, start)
        fill!(p.delbeta, 0)
        linpred!(lp, p, 0)
        updateμ!(r, lp)
    end
    devold = deviance(m)

    for i = 1:maxiter
        f = 1.0 # line search factor
        local dev

        # Compute the change to β, update μ and compute deviance
        try
            delbeta!(p, r.wrkresid, r.wrkwt)
            linpred!(lp, p)
            updateμ!(r, lp)
            dev = deviance(m)
        catch e
            isa(e, DomainError) ? (dev = Inf) : rethrow(e)
        end

        # Line search
        ## If the deviance isn't declining then half the step size
        ## The rtol*dev term is to avoid failure when deviance
        ## is unchanged except for rouding errors.
        while dev > devold + rtol*dev
            f /= 2
            f > minstepfac || error("step-halving failed at beta0 = $(p.beta0)")
            try
                updateμ!(r, linpred(p, f))
                dev = deviance(m)
            catch e
                isa(e, DomainError) ? (dev = Inf) : rethrow(e)
            end
        end
        installbeta!(p, f)

        # Test for convergence
        verbose && println("Iteration: $i, deviance: $dev, diff.dev.:$(devold - dev)")
        if devold - dev < max(rtol*devold, atol)
            cvg = true
            break
        end
        @assert isfinite(dev)
        devold = dev
    end
    cvg || throw(ConvergenceException(maxiter))
    m.fit = true
    m
end

function StatsBase.fit!(m::AbstractGLM;
                        verbose::Bool=false,
                        maxiter::Integer=30,
                        minstepfac::Real=0.001,
                        atol::Real=1e-6,
                        rtol::Real=1e-6,
                        start=nothing,
                        kwargs...)
    if haskey(kwargs, :maxIter)
        Base.depwarn("'maxIter' argument is deprecated, use 'maxiter' instead", :fit!)
        maxiter = kwargs[:maxIter]
    end
    if haskey(kwargs, :minStepFac)
        Base.depwarn("'minStepFac' argument is deprecated, use 'minstepfac' instead", :fit!)
        minstepfac = kwargs[:minStepFac]
    end
    if haskey(kwargs, :convTol)
        Base.depwarn("'convTol' argument is deprecated, use `atol` and `rtol` instead", :fit!)
        rtol = kwargs[:convTol]
    end
    if !issubset(keys(kwargs), (:maxIter, :minStepFac, :convTol))
        throw(ArgumentError("unsupported keyword argument"))
    end
    if haskey(kwargs, :tol)
        Base.depwarn("`tol` argument is deprecated, use `atol` and `rtol` instead", :fit!)
        rtol = kwargs[:tol]
    end

    _fit!(m, verbose, maxiter, minstepfac, atol, rtol, start)
end

function StatsBase.fit!(m::AbstractGLM,
                        y;
                        wts=nothing,
                        offset=nothing,
                        dofit::Bool=true,
                        verbose::Bool=false,
                        maxiter::Integer=30,
                        minstepfac::Real=0.001,
                        atol::Real=1e-6,
                        rtol::Real=1e-6,
                        start=nothing,
                        kwargs...)
    if haskey(kwargs, :maxIter)
        Base.depwarn("'maxIter' argument is deprecated, use 'maxiter' instead", :fit!)
        maxiter = kwargs[:maxIter]
    end
    if haskey(kwargs, :minStepFac)
        Base.depwarn("'minStepFac' argument is deprecated, use 'minstepfac' instead", :fit!)
        minstepfac = kwargs[:minStepFac]
    end
    if haskey(kwargs, :convTol)
        Base.depwarn("'convTol' argument is deprecated, use `atol` and `rtol` instead", :fit!)
        rtol = kwargs[:convTol]
    end
    if !issubset(keys(kwargs), (:maxIter, :minStepFac, :convTol))
        throw(ArgumentError("unsupported keyword argument"))
    end
    if haskey(kwargs, :tol)
        Base.depwarn("`tol` argument is deprecated, use `atol` and `rtol` instead", :fit!)
        rtol = kwargs[:tol]
    end

    r = m.rr
    V = typeof(r.y)
    r.y = copy!(r.y, y)
    isa(wts, Nothing) || copy!(r.wts, wts)
    isa(offset, Nothing) || copy!(r.offset, offset)
    initialeta!(r.eta, r.d, r.l, r.y, r.wts, r.offset)
    updateμ!(r, r.eta)
    fill!(m.pp.beta0, 0)
    m.fit = false
    if dofit
        _fit!(m, verbose, maxiter, minstepfac, atol, rtol, start)
    else
        m
    end
end

"""
    fit(GeneralizedLinearModel, X, y, d, [l = canonicallink(d)]; <keyword arguments>)

Fit a generalized linear model to data. `X` and `y` can either be a matrix and a
vector, respectively, or a formula and a data frame. `d` must be a
`UnivariateDistribution`, and `l` must be a [`Link`](@ref), if supplied.

# Keyword Arguments
- `dofit::Bool=true`: Determines whether model will be fit
- `wts::Vector=similar(y,0)`: prior case weights. Can be length 0.
- `offset::Vector=similar(y,0)`: offset added to `Xβ` to form `eta`.  Can be of
length 0
- `verbose::Bool=false`: Display convergence information for each iteration
- `maxiter::Integer=30`: Maximum number of iterations allowed to achieve convergence
- `atol::Real=1e-6`: Convergence is achieved when the relative change in
deviance is less than `max(rtol*dev, atol)`.
- `rtol::Real=1e-6`: Convergence is achieved when the relative change in
deviance is less than `max(rtol*dev, atol)`.
- `minstepfac::Real=0.001`: Minimum line step fraction. Must be between 0 and 1.
- `start::AbstractVector=nothing`: Starting values for beta. Should have the
same length as the number of columns in the model matrix.
"""
function fit(::Type{M},
    X::Union{Matrix{T},SparseMatrixCSC{T}},
    y::V,
    d::UnivariateDistribution,
    l::Link = canonicallink(d);
    dofit::Bool = true,
    wts::V      = similar(y, 0),
    offset::V   = similar(y, 0),
    fitargs...) where {M<:AbstractGLM,T<:FP,V<:FPVector}

    # Check that X and y have the same number of observations
    if size(X, 1) != size(y, 1)
        throw(DimensionMismatch("number of rows in X and y must match"))
    end

    rr = GlmResp(y, d, l, offset, wts)
    res = M(rr, cholpred(X), false)
    return dofit ? fit!(res; fitargs...) : res
end

fit(::Type{M},
    X::Union{Matrix,SparseMatrixCSC},
    y::AbstractVector,
    d::UnivariateDistribution,
    l::Link=canonicallink(d); kwargs...) where {M<:AbstractGLM} =
        fit(M, float(X), float(y), d, l; kwargs...)

"""
    glm(F, D, args...; kwargs...)

Fit a generalized linear model to data. Alias for `fit(GeneralizedLinearModel, ...)`.
See [`fit`](@ref) for documentation.
"""
glm(F, D, args...; kwargs...) = fit(GeneralizedLinearModel, F, D, args...; kwargs...)

GLM.Link(mm::AbstractGLM) = mm.l
GLM.Link(r::GlmResp{T,D,L}) where {T,D,L} = L()
GLM.Link(r::GlmResp{T,D,L}) where {T,D<:NegativeBinomial,L<:NegativeBinomialLink} = L(r.d.r)
GLM.Link(m::GeneralizedLinearModel) = Link(m.rr)

Distributions.Distribution(r::GlmResp{T,D,L}) where {T,D,L} = D
Distributions.Distribution(m::GeneralizedLinearModel) = Distribution(m.rr)

"""
    dispersion(m::AbstractGLM, sqr::Bool=false)

Return the estimated dispersion (or scale) parameter for a model's distribution,
generally written σ for linear models and ϕ for generalized linear models.
It is, by definition, equal to 1 for the Bernoulli, Binomial, and Poisson families.

If `sqr` is `true`, the squared dispersion parameter is returned.
"""
function dispersion(m::AbstractGLM, sqr::Bool=false)
    r = m.rr
    if dispersion_parameter(r.d)
        wrkwt, wrkresid = r.wrkwt, r.wrkresid
        s = sum(i -> wrkwt[i] * abs2(wrkresid[i]), eachindex(wrkwt, wrkresid)) / dof_residual(m)
        sqr ? s : sqrt(s)
    else
        one(eltype(r.mu))
    end
end

"""
    predict(mm::AbstractGLM, newX::AbstractMatrix; offset::FPVector=Vector{eltype(newX)}(0))

Form the predicted response of model `mm` from covariate values `newX` and, optionally,
an offset.
"""
function predict(mm::AbstractGLM, newX::AbstractMatrix;
                 offset::FPVector=eltype(newX)[])
    eta = newX * coef(mm)
    if !isempty(mm.rr.offset)
        length(offset) == size(newX, 1) ||
            throw(ArgumentError("fit with offset, so `offset` kw arg must be an offset of length `size(newX, 1)`"))
        broadcast!(+, eta, eta, offset)
    else
        length(offset) > 0 && throw(ArgumentError("fit without offset, so value of `offset` kw arg does not make sense"))
    end
    mu = [linkinv(Link(mm), x) for x in eta]
end

# A helper function to choose default values for eta
function initialeta!(eta::AbstractVector,
                    dist::UnivariateDistribution,
                    link::Link,
                    y::AbstractVector,
                    wts::AbstractVector,
                    off::AbstractVector)

    n  = length(y)
    lw = length(wts)
    lo = length(off)

    if lw == n
        @inbounds @simd for i = eachindex(y, eta, wts)
            μ      = mustart(dist, y[i], wts[i])
            eta[i] = linkfun(link, μ)
        end
    elseif lw == 0
        @inbounds @simd for i = eachindex(y, eta)
            μ      = mustart(dist, y[i], 1)
            eta[i] = linkfun(link, μ)
        end
    else
        throw(ArgumentError("length of wts must be either $n or 0 but was $lw"))
    end

    if lo == n
        @inbounds @simd for i = eachindex(eta, off)
            eta[i] -= off[i]
        end
    elseif lo != 0
        throw(ArgumentError("length of off must be either $n or 0 but was $lo"))
    end

    return eta
end

# Helper function to check that the values of y are in the allowed domain
function checky(y, d::Distribution)
    if any(x -> !insupport(d, x), y)
        throw(ArgumentError("y must be in the support of D"))
    end
    return nothing
end
function checky(y, d::Binomial)
    for yy in y
        0 ≤ yy ≤ 1 || throw(ArgumentError("$yy in y is not in [0,1]"))
    end
    return nothing
end


text:
# initialise evotree
function init_evotree_gpu(params::EvoTypes{T,U,S},
    X::AbstractMatrix, Y::AbstractVector) where {T,U,S}

    K = 1
    levels = nothing
    X = convert(Matrix{T}, X)

    if typeof(params.loss) == Logistic
        Y = CuArray(T.(Y))
        μ = [logit(mean(Y))]
    elseif typeof(params.loss) == Poisson
        Y = CuArray(T.(Y))
        μ = fill(log(mean(Y)), 1)
    # elseif typeof(params.loss) == Softmax
    #     if typeof(Y) <: AbstractCategoricalVector
    #         levels = CategoricalArray(CategoricalArrays.levels(Y))
    #         K = length(levels)
    #         μ = zeros(T, K)
    #         Y = MLJModelInterface.int.(Y)
    #     else
    #         levels = CategoricalArray(sort(unique(Y)))
    #         K = length(levels)
    #         μ = zeros(T, K)
    #         Y = UInt32.(Y)
    #     end
    elseif typeof(params.loss) == Gaussian
        K = 2
        Y = CuArray(T.(Y))
        μ = [mean(Y), log(std(Y))]
    else
        Y = CuArray(T.(Y))
        μ = [mean(Y)]
    end

    # initialize preds
    X_size = size(X)
    pred = CUDA.zeros(T, K, X_size[1])
    pred .= CuArray(μ)

    bias = TreeGPU(CuArray(μ))
    evotree = GBTreeGPU([bias], params, Metric(), K, levels)

    # initialize gradients and weights
    δ𝑤 = CUDA.ones(T, 2 * K + 1, X_size[1])

    # binarize data into quantiles
    edges = get_edges(X, params.nbins)
    X_bin = CuArray(binarize(X, edges))

    𝑖_ = UInt32.(collect(1:X_size[1]))
    𝑗_ = UInt32.(collect(1:X_size[2]))
    𝑗 = zeros(eltype(𝑗_), ceil(Int, params.colsample * X_size[2]))

    # initializde histograms
    nodes = [TrainNodeGPU(X_size[2], params.nbins, K, T) for n in 1:2^params.max_depth - 1]
    nodes[1].𝑖 = CUDA.zeros(eltype(𝑖_), ceil(Int, params.rowsample * X_size[1]))
    out = CUDA.zeros(UInt32, length(nodes[1].𝑖))
    left = CUDA.zeros(UInt32, length(nodes[1].𝑖))
    right = CUDA.zeros(UInt32, length(nodes[1].𝑖))

    # store cache
    cache = (params = deepcopy(params),
        X = CuArray(X), X_bin = X_bin, Y = Y, K = K,
        nodes = nodes,
        pred = pred,
        𝑖_ = 𝑖_, 𝑗_ = 𝑗_, 𝑗 = 𝑗, 𝑖 = Array(nodes[1].𝑖),
        out = out, left = left, right = right,
        δ𝑤 = δ𝑤,
        edges = edges)

    cache.params.nrounds = 0

    return evotree, cache
end

function grow_evotree!(evotree::GBTreeGPU{T}, cache) where {T}

    # initialize from cache
    params = evotree.params
    X_size = size(cache.X_bin)
    δnrounds = params.nrounds - cache.params.nrounds

    # loop over nrounds
    for i in 1:δnrounds
        # select random rows and cols
        sample!(params.rng, cache.𝑖_, cache.𝑖, replace=false, ordered=true)
        sample!(params.rng, cache.𝑗_, cache.𝑗, replace=false, ordered=true)
        cache.nodes[1].𝑖 .= CuArray(cache.𝑖)

        # build a new tree
        update_grads_gpu!(params.loss, cache.δ𝑤, cache.pred, cache.Y)
        # # assign a root and grow tree
        tree = TreeGPU(params.max_depth, evotree.K, zero(T))
        grow_tree_gpu!(tree, cache.nodes, params, cache.δ𝑤, cache.edges, CuVector(cache.𝑗), cache.out, cache.left, cache.right, cache.X_bin, cache.K);
        push!(evotree.trees, tree)
        # update predctions
        predict!(params.loss, cache.pred, tree, cache.X, cache.K)
    end # end of nrounds
    cache.params.nrounds = params.nrounds
    # return model, cache
    return evotree
end

# grow a single tree - grow through all depth
function grow_tree_gpu!(
    tree::TreeGPU{T},
    nodes,
    params::EvoTypes{T,U,S},
    δ𝑤::AbstractMatrix{T},
    edges,
    𝑗, out, left, right,
    X_bin::AbstractMatrix, K) where {T,U,S}

    n_next = [1]
    n_current = copy(n_next)
    depth = 1

    # reset nodes
    @threads for n in eachindex(nodes)
        nodes[n].h .= 0
        nodes[n].∑ .= 0
        nodes[n].gain = -Inf
        fill!(nodes[n].gains, -Inf)
    end

    # initialize summary stats
    nodes[1].∑ .= vec(sum(δ𝑤[:, nodes[1].𝑖], dims=2))
    nodes[1].gain = get_gain(params.loss, Array(nodes[1].∑), params.λ, K) # should use a GPU version?

    # grow while there are remaining active nodes - TO DO histogram substraction hits issue on GPU
    while length(n_current) > 0 && depth <= params.max_depth
        offset = 0 # identifies breakpoint for each node set within a depth
        if depth < params.max_depth
            for n_id ∈ 1:length(n_current)
                n = n_current[n_id]
                if n_id % 2 == 0
                    if n % 2 == 0
                        nodes[n].h .= nodes[n >> 1].h .- nodes[n + 1].h
                        CUDA.synchronize()
                    else
                        nodes[n].h .= nodes[n >> 1].h .- nodes[n - 1].h
                        CUDA.synchronize()
                    end
                else
                    update_hist_gpu!(params.loss, nodes[n].h, δ𝑤, X_bin, nodes[n].𝑖, 𝑗, K)
                end
            end
        end

        # grow while there are remaining active nodes
        for n ∈ sort(n_current)
            if depth == params.max_depth || nodes[n].∑[end] <= params.min_weight
                pred_leaf_gpu!(params.loss, tree.pred, n, Array(nodes[n].∑), params)
            else
                update_gains_gpu!(params.loss, nodes[n], 𝑗, params, K)
                best = findmax(nodes[n].gains)
                if best[2][1] != params.nbins && best[1] > nodes[n].gain + params.γ
                    tree.gain[n] = best[1]
                    tree.cond_bin[n] = best[2][1]
                    tree.feat[n] = best[2][2]
                    tree.cond_float[n] = edges[tree.feat[n]][tree.cond_bin[n]]
                end
                # println("node: ", n, " | best: ", best, " | nodes[n].gain: ", nodes[n].gain)
                tree.split[n] = tree.cond_bin[n] != 0
                if !tree.split[n]
                    pred_leaf_gpu!(params.loss, tree.pred, n, Array(nodes[n].∑), params)
                    popfirst!(n_next)
                else
                    _left, _right = split_set_threads_gpu!(out, left, right, nodes[n].𝑖, X_bin, tree.feat[n], tree.cond_bin[n], offset)
                    nodes[n << 1].𝑖, nodes[n << 1 + 1].𝑖 = _left, _right
                    offset += length(nodes[n].𝑖)
                    # println("length(_left): ", length(_left), " | length(_right): ", length(_right))
                    # println("best: ", best)
                    update_childs_∑_gpu!(params.loss, nodes, n, best[2][1], best[2][2])
                    nodes[n << 1].gain = get_gain(params.loss, Array(nodes[n << 1].∑), params.λ, K)
                    nodes[n << 1 + 1].gain = get_gain(params.loss, Array(nodes[n << 1 + 1].∑), params.λ, K)

                    if length(_right) >= length(_left)
                        push!(n_next, n << 1)
                        push!(n_next, n << 1 + 1)
                    else
                        push!(n_next, n << 1 + 1)
                        push!(n_next, n << 1)
                    end
                    popfirst!(n_next)
                end
            end
        end
        n_current = copy(n_next)
        depth += 1
    end # end of loop over active ids for a given depth

    return nothing
end


text:
"""
Executes a transient FSI driver with a constant function in space, linear in time:
  u(x,t) = [1.0, -1.0]^T * t
  v(x,t) = [1.0, -1.0]^T
  p(x,t) = 0.0
"""
function execute(problem::FSIProblem{:analytical};kwargs...)

  # Problem setting
  println("Setting analytical fluid problem parameters")
  # Solid properties
  E_s = _get_kwarg(:E_s,kwargs,1.0)
  ν_s = _get_kwarg(:nu_s,kwargs,0.4)
  ρ_s = _get_kwarg(:rho_s,kwargs,1.0)
  # Fluid properties
  ρ_f = _get_kwarg(:rho_f,kwargs,1.0)
  μ_f = _get_kwarg(:mu_f,kwargs,1.0)
  γ_f = _get_kwarg(:gamma_f,kwargs,1.0)
  # Mesh properties
  n_m = _get_kwarg(:n_m,kwargs,10)
  E_m = _get_kwarg(:E_m,kwargs,1.0)
  ν_m = _get_kwarg(:nu_m,kwargs,-0.1)
  α_m = _get_kwarg(:alpha_m,kwargs,1.0e-5)
  weight_strategy = _get_kwarg(:alpha_m_weight,kwargs,"constant")
  # Time stepping
  t0 = _get_kwarg(:t0,kwargs,0.0)
  tf = _get_kwarg(:tf,kwargs,0.5)
  dt = _get_kwarg(:dt,kwargs,0.1)
  θ  = _get_kwarg(:theta,kwargs,0.5)
  # Post-process
  is_vtk = _get_kwarg(:is_vtk,kwargs,false)

  # Mesh strategy
  strategyName = _get_kwarg(:strategy,kwargs,"laplacian")
  strategy = WeakForms.MeshStrategy{Symbol(strategyName)}()

  # Fluid-Solid coupling type
  couplingName = _get_kwarg(:coupling,kwargs,"strong")
  coupling = WeakForms.Coupling{Symbol(couplingName)}()

  # Define BC functions
  println("Defining Boundary conditions")
  #u(x, t) = 1.0e-2 * VectorValue( x[1]^2*x[2] , -x[1]*x[2]^2) * t
  #v(x, t) = 1.0e-2 * VectorValue( x[1]^2*x[2] , -x[1]*x[2]^2)
  u(x, t) = 1.0e-2 * VectorValue( 1.0 , -1.0) * t
  v(x, t) = 1.0e-2 * VectorValue( 1.0 , -1.0)
  u(t::Real) = x -> u(x,t)
  v(t::Real) = x -> v(x,t)
  ∂tu(t) = x -> VectorValue(ForwardDiff.derivative(t -> get_array(u(x,t)),t))
  ∂tv(t) = x -> VectorValue(ForwardDiff.derivative(t -> get_array(v(x,t)),t))
  ∂tu(x,t) = ∂tu(t)(x)
  ∂tv(x,t) = ∂tv(t)(x)
  T_u = typeof(u)
  T_v = typeof(v)
  @eval ∂t(::$T_u) = $∂tu
  @eval ∂t(::$T_v) = $∂tv
  p(x,t) = 0.0
  p(t::Real) = x -> p(x,t)
  bconds = get_boundary_conditions(problem,strategy,u,v)

  # Define Forcing terms
  I = TensorValue( 1.0, 0.0, 0.0, 1.0 )
  F(t) = x -> ∇(u(t))(x) + I
  J(t) = x -> det(F(t))(x)
  E(t) = x -> 0.5 * ((F(t)(x)')⋅F(t)(x) - I)
  (λ_s, μ_s) = WeakForms.lame_parameters(E_s,ν_s)
  (λ_m, μ_m) = WeakForms.lame_parameters(E_m,ν_m)
  S_SV(t) = x -> 2*μ_s*E(t)(x) + λ_s*tr(E(t)(x))*I
  fv_ST_Ωf(t) = x -> - μ_f*Δ(v(t))(x) + ∇(p(t))(x)
  function fu_closure(t,strategy)
    fu = if typeof(strategy) == WeakForms.MeshStrategy{:laplacian}
      x -> - α_m * Δ(u(t))(x)
    elseif typeof(strategy) == WeakForms.MeshStrategy{:linearElasticity}
      x -> - μ_m * Δ(u(t))(x)
    elseif typeof(strategy) == WeakForms.MeshStrategy{:neoHookean}
      x -> - μ_m * Δ(u(t))(x)
    else
      @notimplemented("The soruce term for $strategy strategy is not implemented")
    end
    return fu
  end
  fu_Ωf(t) = fu_closure(t,strategy)
  fv_Ωf(t) = x -> ρ_f * ∂t(v)(t)(x) - μ_f * Δ(v(t))(x) + ∇(p(t))(x) + ρ_f*( (∇(v(t))(x)')⋅(v(t)(x) - ∂t(u)(t)(x)) )
  fp_Ωf(t) = x -> (∇⋅v(t))(x)
  fu_Ωs(t) = x -> ∂t(u)(t)(x) - v(t)(x)
  fv_Ωs(t) = x -> ρ_s * ∂t(v)(t)(x) #- (∇⋅(F(t)⋅S_SV(t)))(x)  # Divergence of a a doted function not working yet...

  # Discrete model
  println("Defining discrete model")
  domain = (-1,1,-1,1)
  partition = (n_m,n_m)
  model = CartesianDiscreteModel(domain,partition)
  trian = Triangulation(model)
  R = 0.5
  xc = 0.0
  yc = 0.0
  function is_in(coords)
    n = length(coords)
    x = (1/n)*sum(coords)
    d = (x[1]-xc)^2 + (x[2]-yc)^2 - R^2
    d < 1.0e-8
  end
  oldcell_to_coods = get_cell_coordinates(trian)
  oldcell_to_is_in = collect1d(lazy_map(is_in,oldcell_to_coods))
  incell_to_cell = findall(oldcell_to_is_in)
  outcell_to_cell = findall(collect(Bool, .! oldcell_to_is_in))
  model_solid = DiscreteModel(model,incell_to_cell)
  model_fluid = DiscreteModel(model,outcell_to_cell)
  models = Dict(:Ω => model, :Ωf => model_fluid, :Ωs => model_solid)

  # Build fluid-solid interface labelling
  println("Defining Fluid-Solid interface")
  labeling = get_face_labeling(model_fluid)
  new_entity = num_entities(labeling) + 1
  topo = get_grid_topology(model_fluid)
  D = num_cell_dims(model_fluid)
  for d in 0:D-1
    fluid_boundary_mask = collect(Bool,get_isboundary_face(topo,d))
    fluid_outer_boundary_mask = get_face_mask(labeling,"boundary",d)
    fluid_interface_mask = collect(Bool,fluid_boundary_mask .!= fluid_outer_boundary_mask)
    dface_list = findall(fluid_interface_mask)
    for face in dface_list
      labeling.d_to_dface_to_entity[d+1][face] = new_entity
    end
  end
  add_tag!(labeling,"interface",[new_entity])

  # Triangulations
  println("Defining triangulations")
  Tₕ = get_FSI_triangulations(models,coupling)

  # Quadratures
  println("Defining quadratures")
  order = _get_kwarg(:order,kwargs,2)
  dTₕ = get_FSI_measures(Tₕ,order)

  # Test FE Spaces
  println("Defining FE spaces")
  Y_ST, X_ST, Y_FSI, X_FSI = get_FE_spaces(strategy,coupling,models,order,bconds,constraint=:zeromean)

  # Stokes problem for initial solution
  println("Defining Stokes operator")
  op_ST = get_Stokes_operator(X_ST,Y_ST,strategy,dTₕ[:Ωf],μ_f,fv_ST_Ωf(0.0))

  # Setup equation parameters
  mesh_params = Dict{Symbol,Any}(
    :w_strategy=>weight_strategy,
    :α=>α_m,
    :E=>E_m,
    :ν=>ν_m,
  )
  fluid_params = Dict{Symbol,Any}(
    :μ=>μ_f,
    :ρ=>ρ_f,
    :fu=>fu_Ωf,
    :fv=>fv_Ωf,
  )
  solid_params = Dict{Symbol,Any}(
    :ρ=>ρ_s,
    :E=>E_s,
    :ν=>ν_s,
    :fu=>fu_Ωs,
    :fv=>fv_Ωs,
  )
  Γi_params = Dict{Symbol,Any}(
    :μ=>μ_f,
    :γ=>γ_f,
    :dt=>dt
  )
  params = (mesh_params, fluid_params, solid_params, Γi_params)

  # FSI problem
  println("Defining FSI operator")
  op_FSI = get_FSI_operator(X_FSI,Y_FSI,coupling,strategy,Tₕ,dTₕ,params)

  # Setup output files
  folderName = "fsi-results"
  fileName = "fields"
  if !isdir(folderName)
    mkdir(folderName)
  end
  filePath = join([folderName, fileName], "/")

  # Solve Stokes problem
  @timeit "ST problem" begin
    println("Defining Stokes solver")
    xh = solve(op_ST)
    if(is_vtk)
      writePVD(filePath, Tₕ[:Ωf], [(xh, 0.0)])
    end
  end

  # Compute Stokes solution L2-norm
  l2(w) = w⋅w
  eu_ST = u(0.0) - xh[1]
  ev_ST = v(0.0) - xh[2]
  ep_ST = p(0.0) - xh[3]
  eul2_ST = sqrt(∑( ∫(l2(eu_ST))dTₕ[:Ωf] ))
  evl2_ST = sqrt(∑( ∫(l2(ev_ST))dTₕ[:Ωf] ))
  epl2_ST = sqrt(∑( ∫(l2(ep_ST))dTₕ[:Ωf] ))
  println("Stokes L2-norm u: ", eul2_ST)
  println("Stokes L2-norm v: ", evl2_ST)
  println("Stokes L2-norm p: ", epl2_ST)
  @test eul2_ST < 1.0e-10
  @test evl2_ST < 1.0e-10
  @test epl2_ST < 1.0e-10

  # Solve FSI problem
  @timeit "FSI problem" begin
  println("Defining FSI solver")
  xh0 = interpolate_everywhere([u(0.0),v(0.0),p(0.0)],X_FSI(0.0))
  nls = NLSolver(
  show_trace = true,
  method = :newton,
  linesearch = BackTracking(),
  ftol = 1.0e-10,
  iterations = 50
  )
  odes =  ThetaMethod(nls, dt, 0.5)
  solver = TransientFESolver(odes)
  xht = solve(solver, op_FSI, xh0, t0, tf)
  end

  # Compute outputs
  out_params = Dict(
  :u=>u,
  :v=>v,
  :p=>p,
  :filePath=>filePath,
  :is_vtk=>is_vtk
  )
  output = computeOutputs(xht,Tₕ,dTₕ,strategy,out_params)

end

function get_boundary_conditions(problem::FSIProblem{:analytical},strategy::WeakForms.MeshStrategy,u,v)
  boundary_conditions = (
  # Tags
  FSI_Vu_tags = ["boundary"],
  FSI_Vv_tags = ["boundary"],
  ST_Vu_tags = ["boundary","interface"],
  ST_Vv_tags = ["boundary","interface"],
  # Values,
  FSI_Vu_values = [u],
  FSI_Vv_values = [v],
  ST_Vu_values = [u(0.0),u(0.0)],
  ST_Vv_values = [v(0.0),v(0.0)],
  )
end

function get_FE_spaces(problem::FSIProblem{:analytical},strategy::WeakForms.MeshStrategy,model,model_fluid,order,bconds)
  Vu_FSI = TestFESpace(
      model=model,
      valuetype=VectorValue{2,Float64},
      reffe=:Lagrangian,
      order=order,
      conformity =:H1,
      dirichlet_tags=bconds[:FSI_Vu_tags])
  Vv_FSI = TestFESpace(
      model=model,
      valuetype=VectorValue{2,Float64},
      reffe=:Lagrangian,
      order=order,
      conformity =:H1,
      dirichlet_tags=bconds[:FSI_Vv_tags])
  Vu_ST = TestFESpace(
      model=model_fluid,
      #model=model,
      valuetype=VectorValue{2,Float64},
      reffe=:Lagrangian,
      order=order,
      conformity =:H1,
      dirichlet_tags=bconds[:ST_Vu_tags])
  Vv_ST = TestFESpace(
      model=model_fluid,
      #model=model,
      valuetype=VectorValue{2,Float64},
      reffe=:Lagrangian,
      order=order,
      conformity =:H1,
      dirichlet_tags=bconds[:ST_Vv_tags])
  Q = TestFESpace(
      model=model_fluid,
      #model=model,
      valuetype=Float64,
      order=order-1,
      reffe=:Lagrangian,
      constraint=:zeromean,
      conformity=:C0)

  # Trial FE Spaces
  Uu_ST = TrialFESpace(Vu_ST,bconds[:ST_Vu_values])
  Uv_ST = TrialFESpace(Vv_ST,bconds[:ST_Vv_values])
  Uu_FSI = TransientTrialFESpace(Vu_FSI,bconds[:FSI_Vu_values])
  Uv_FSI = TransientTrialFESpace(Vv_FSI,bconds[:FSI_Vv_values])
  P = TrialFESpace(Q)

  # Multifield FE Spaces
  fe_spaces = (
      Y_ST = MultiFieldFESpace([Vu_ST,Vv_ST,Q]),
      X_ST = MultiFieldFESpace([Uu_ST,Uv_ST,P]),
      Y_FSI = MultiFieldFESpace([Vu_FSI,Vv_FSI,Q]),
      X_FSI = TransientMultiFieldFESpace([Uu_FSI,Uv_FSI,P])
  )
end

function computeOutputs(xht,Tₕ,dTₕ,strategy,params)

  # Unpack parameters
  u = params[:u]
  v = params[:v]
  p = params[:p]
  filePath = params[:filePath]
  is_vtk = params[:is_vtk]
  l2(w) = w⋅w

  ## Initialize arrays
  tpl = Real[]
  eupl = Real[]

  ## Loop over steps
  outfiles = paraview_collection(filePath, append=true) do pvd
    for (i,(xh, t)) in enumerate(xht)
      println("STEP: $i, TIME: $t")
      println("============================")

      # Compute errors
      eu = u(t) - xh[1]
      ev = v(t) - xh[2]
      ep = p(t) - xh[3]
      eul2 = sqrt(∑( ∫(l2(eu))dTₕ[:Ω] ))
      evl2 = sqrt(∑( ∫(l2(ev))dTₕ[:Ω] ))
      epl2 = sqrt(∑( ∫(l2(ep))dTₕ[:Ω] ))

      # Write to PVD
      if(is_vtk)
        uh = xh[1]
        vh = xh[2]
        ph = xh[3]
        pvd[t] = createvtk(
        Tₕ[:Ω],
        filePath * "_$t.vtu",
        cellfields = ["uh" => uh, "vh" => vh, "ph" => ph, "euh" => eu]
        )
      end

      # Test checks
      println("FSI L2-norm u: ", eul2)
      println("FSI L2-norm v: ", evl2)
      println("FSI L2-norm p: ", epl2)
      @test eul2 < 1.0e-10
      @test evl2 < 1.0e-10
      @test epl2 < 1.0e-10

      push!(tpl, t)
      push!(eupl, eul2)
    end
  end
  return (tpl, eupl)
end


text:
# This script demonstrates usage of the SMLMData package.
using Revise
using SMLMData
using DataFrames
using CSV

## Create an empty smld structure.
smld_empty = SMLMData.SMLD2D()

## Create an empty smld structure with a specified numberr of localizations.
nlocalizations = 21
smld_init = SMLMData.SMLD2D(nlocalizations)

## Create an smld structure from a properly formatted csv file:
#  The csv file should be organized s.t. each row is a localization, with the
#  data organized as [datasetnum, framenum, x, y, x_se, y_se].
data = DataFrames.DataFrame(CSV.File("C:\\Users\\David\\Documents\\GitHub\\example_data.csv"))
smld = SMLMData.SMLD2D(data)

## Isolate a subregion of `smld`.
smld_sub = SMLMData.isolatesmld(smld, [1, 2, 3])
smld_sub = SMLMData.isolatesmld(smld, smld.connectID.==2)


text:
# SPDX-License-Identifier: X11
# Simple HTTP server for encrypted writeups

import CSV
import HTTP
import JSON
import Nettle
import Sockets

const ROOTDIR = "/srv/http"
const ASSETS = "$ROOTDIR/ctf-writeups-vault"
const PASSWDFILE = "$ASSETS/passwd"
const PASSWD = CSV.File(PASSWDFILE; delim=':', type=String)

const HOST = Sockets.localhost
const PORT = 5334

const HASH_TYPE = "sha256"
const SHA1_FORMAT = r"^[0-9a-f]{40}$"

const BAD_REQUEST = 400
const UNAUTHORISED = 401
const NOT_FOUND = 404
const METHOD_NOT_ALLOWED = 405
const FAIL_RESPONSES = Dict(
  BAD_REQUEST        => """
                        {
                          "msg": "Sorry, I only speak JSON.",
                          "ps": "Do you have the right keys?"
                        }
                        """,
  UNAUTHORISED       => "What's the super secret password?\n",
  NOT_FOUND          => "I couldn't find what you're looking for...\n",
  METHOD_NOT_ALLOWED => "POST requests only, please.\n"
)

function httphandler(http::HTTP.Stream)::Nothing
  @show http.message

  # Allow POST only
  if http.message.method != "POST"
    respond_fail(http, METHOD_NOT_ALLOWED)
    return
  end

  # Parse POSTed data
  j = try
    JSON.parse(String(readavailable(http)))
  catch
    respond_fail(http, BAD_REQUEST)
    return
  end

  # Check that received JSON has required keys
  if !(haskey(j, "file") && haskey(j, "flag") &&
       occursin(SHA1_FORMAT, j["file"])) # Check if filename is valid
    respond_fail(http, BAD_REQUEST)
    return
  # Check if file exists
  elseif !isfile("$ASSETS/" * j["file"])
    respond_fail(http, NOT_FOUND)
    return
  # Validate flag
  elseif !validate(j["file"], j["flag"])
    respond_fail(http, UNAUTHORISED)
    return
  else
    # Send data!
    HTTP.Streams.setstatus(http, 200)
    HTTP.Streams.setheader(http,
      "Content-Type" => "application/octet-stream")

    HTTP.Streams.startwrite(http)
    write(http, read("$ASSETS/" * j["file"]))
    return
  end
end

function validate(filename::AbstractString, flag::AbstractString)::Bool
  for r ∈ PASSWD
    # Iterate until a matching entry is found
    r.filename != filename && continue

    # Compare the salted hash of the flag
    return Nettle.hexdigest(HASH_TYPE, r.salt * flag) == r.hash
  end
end

function respond_fail(http::HTTP.Stream, status::Int)::Nothing
  # Clear stream to clear my error log~
  readavailable(http)

  HTTP.Streams.setstatus(http, status)
  HTTP.Streams.setheader(http,
    "Content-Type" => "text/plain")

  HTTP.Streams.startwrite(http)
  write(http, FAIL_RESPONSES[status])
  return
end

println("Listening...")
HTTP.Servers.listen(httphandler, HOST, PORT)


text:
using Cubature
using Distributions

"""
Demo of robust Bayesian analysis, from Berger 1984
"""
function robustPriorDemo()
    x = 5
    obsVar = 1

    # First let us verify that the Gaussian prior indeed satisfies the prior quartiles
    priorMu = 0
    priorVar = 2.19

    nprior = Normal(priorMu, sqrt(priorVar))
    @assert isapprox(cdf(nprior, -1), .25; atol=.001)
    @assert isapprox(cdf(nprior, 0) - cdf(nprior, -1), .25; atol=.001)
    @assert isapprox(cdf(nprior, 1) - cdf(nprior, 0), .25; atol=.001)
    @assert isapprox(1-cdf(nprior, 1), .25; atol=.001)
    @assert isapprox(median(nprior), 0; atol=.001)

    # Now compute posterior mean using Gaussian prior
    postVar = 1/(1/obsVar + 1/priorVar)
    postMeanNormal = postVar*(priorMu/priorVar + x/obsVar)
    @assert isapprox(postMeanNormal, 3.43; atol=.01)

    # Now let us do the same thing for the Cauchy
    cprior = Cauchy(0, 1)
    @assert isapprox(cdf(cprior, -1), .25; atol=.001)
    @assert isapprox(cdf(cprior, 0) - cdf(cprior, -1), .25; atol=.001)
    @assert isapprox(cdf(cprior, 1) - cdf(cprior, 0), .25; atol=.001)
    @assert isapprox(1-cdf(cprior, 1), .25; atol=.001)
    @assert isapprox(median(cprior), 0; atol=.001)

    # Now let us compute posterior mean using Cauchy

    lik(θ) = pdf(Normal(θ, sqrt(obsVar)), x)
    prior(θ) = pdf(cprior, θ)
    post(θ) = lik(θ).*prior(θ)

    ∞=10

    Z, err = hquadrature(post, -∞, ∞; reltol=1e-4)

    f(θ) = θ.*post(θ)/Z

    postMeanCauchy, err = hquadrature(f, -∞, ∞; reltol=1e-4)
    @assert isapprox(postMeanCauchy, 4.56; atol=.01)

    (postMeanNormal, postMeanCauchy)
end


text:
module GridapGmsh

using Libdl

using Gridap.Arrays
using Gridap.TensorValues
using Gridap.Fields
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.Geometry: max_cells_arround_vertex
using Gridap.Geometry: _fill_cells_around_scratch!
using Gridap.Geometry: _set_intersection!

export GmshDiscreteModel

deps_jl = joinpath(@__DIR__, "..", "deps", "deps.jl")
if !isfile(deps_jl)
  s = """
  Package GridapGmsh not installed properly.
  Run Pkg.build(\"GridapGmsh\"), restart Julia and try again.
  """
  error(s)
end

include(deps_jl)

if GMSH_FOUND
  include(gmsh_jl)
  # Hack taken from MPI.jl
  function __init__()
    @static if Sys.isunix()
      Libdl.dlopen(gmsh.lib, Libdl.RTLD_LAZY | Libdl.RTLD_GLOBAL)
    end
  end
else
    s = """
    Gmsh not found in system paths.
    Install Gmsh or export path to Gmsh and rebuild the project.
    Run Pkg.build(\"GridapGmsh\"), restart Julia and try again.
    """
    @warn s
end

macro check_if_loaded()
  quote
    if ! GMSH_FOUND
      error("GridapGmsh is not loaded or installed properly.")
    end
  end
end

include("GmshDiscreteModels.jl")

end # module


text:
module Model

using ParticleDA

using Random, Distributions, Base.Threads, GaussianRandomFields, HDF5
using ParticleDA.Default_params
using DelimitedFiles
using FieldMetadata

include("llw2d.jl")
using .LLW2d

"""
    ModelParameters()

Parameters for the model. Keyword arguments:

* `nx::Int` : Number of grid points in the x direction
* `ny::Int` : Number of grid points in the y direction
* `x_length::AbstractFloat` : Domain size (m) in the x direction
* `y_length::AbstractFloat` : Domain size (m) in the y direction
* `dx::AbstractFloat` : Distance (m) between grid points in the x direction
* `dy::AbstractFloat` : Distance (m) between grid points in the y direction
* `n_state_var::Int`: Number of variables in the state vector
* `nobs::Int` : Number of observation stations
* `station_filename::String` : Name of input file for station coordinates
* `station_distance_x::Float` : Distance between stations in the x direction [m]
* `station_distance_y::Float` : Distance between stations in the y direction [m]
* `station_boundary_x::Float` : Distance between bottom left edge of box and first station in the x direction [m]
* `station_boundary_y::Float` : Distance between bottom left edge of box and first station in the y direction [m]
* `n_integration_step::Int` : Number of sub-steps to integrate the forward model per time step.
* `time_step::AbstractFloat` : Time step length (s)
* `state_prefix::String` : Prefix of the time slice data groups in output
* `title_da::String` : Suffix of the data assimilated data group in output
* `title_syn::String` : Suffix of the true state data group in output
* `title_grid::String` : Name of the grid data group in output
* `title_stations::String` : Name of the station coordinates data group in output
* `title_params::String` : Name of the parameters data group in output
* `peak_position::Vector{AbstractFloat}` : The [x,y] coordinates (m) of the initial wave peak
* `peak_height::AbstractFloat` : The height (m) of the initial wave peak
* `source_size::AbstractFloat` : Cutoff distance (m) from the peak for the initial wave
* `bathymetry_setup::AbstractFloat` : Bathymetry set-up.
* `lambda::AbstractFloat` : Length scale for Matérn covariance kernel in background noise
* `nu::AbstractFloat` : Smoothess parameter for Matérn covariance kernel in background noise
* `sigma::AbstractFloat` : Marginal standard deviation for Matérn covariance kernel in background noise
* `lambda_initial_state::AbstractFloat` : Length scale for Matérn covariance kernel in initial state of particles
* `nu_initial_state::AbstractFloat` : Smoothess parameter for Matérn covariance kernel in initial state of particles
* `sigma_initial_state::AbstractFloat` : Marginal standard deviation for Matérn covariance kernel in initial state of particles
* `padding::Int` : Min padding for circulant embedding gaussian random field generator
* `primes::Int`: Whether the size of the minimum circulant embedding of the covariance matrix can be written as a product of small primes (2, 3, 5 and 7). Default is `true`.
* `particle_initial_state::String` : Initial state of the particles before noise is added. Possible options are
  * "zero" : initialise height and velocity to 0 everywhere
  * "true" : copy the true initial state
* `absorber_thickness_fraction::Float` : Thickness of absorber for sponge absorbing boundary conditions, fraction of grid size
* `boundary_damping::Float` : damping for boundaries
* `cutoff_depth::Float` : Shallowest water depth
* `obs_noise_std::Float`: Standard deviation of noise added to observations of the true state
* `particle_dump_file::String`: file name for dump of particle state vectors
* `particle_dump_time::Int`: list of (one more more) time steps to dump particle states
"""
Base.@kwdef struct ModelParameters{T<:AbstractFloat}

    nx::Int = 200
    ny::Int = 200
    x_length::T = 400.0e3
    y_length::T = 400.0e3
    dx::T = x_length / (nx - 1)
    dy::T = y_length / (ny - 1)

    n_state_var::Int = 3

    time_step::T = 50.0
    n_integration_step::Int = 50

    station_filename::String = ""
    nobs::Int = 4
    station_distance_x::T = 20.0e3
    station_distance_y::T = 20.0e3
    station_boundary_x::T = 150.0e3
    station_boundary_y::T = 150.0e3

    source_size::T = 3.0e4
    bathymetry_setup::T = 3.0e3
    peak_height = 1.0
    peak_position = [floor(Int, nx / 4) * dx, floor(Int, ny / 4) * dy]

    lambda::Vector{T} = [1.0e4, 1.0e4, 1.0e4]
    nu::Vector{T} = [2.5, 2.5, 2.5]
    sigma::Vector{T} = [1.0, 1.0, 1.0]

    lambda_initial_state::Vector{T} = [1.0e4, 1.0e4, 1.0e4]
    nu_initial_state::Vector{T} = [2.5, 2.5, 2.5]
    sigma_initial_state::Vector{T} = [10.0, 10.0, 10.0]

    padding::Int = 100
    primes::Bool = true

    particle_initial_state::String = "zero"

    absorber_thickness_fraction::T = 0.1
    boundary_damping::T = 0.015
    cutoff_depth::T = 10.0

    state_prefix::String = "data"
    title_avg::String = "avg"
    title_var::String = "var"
    title_syn::String = "syn"
    title_grid::String = "grid"
    title_stations::String = "stations"
    title_params::String = "params"

    obs_noise_std::T = 1.0

    particle_dump_file = "particle_dump.h5"
    particle_dump_time = [-1]
end

function ParticleDA.get_params(T::Type{ModelParameters}, user_input_dict::Dict)

    for key in ("lambda", "nu", "sigma", "lambda_initial_state", "nu_initial_state", "sigma_initial_state")
        if haskey(user_input_dict, key) && !isa(user_input_dict[key], Vector)
            user_input_dict[key] = fill(user_input_dict[key], 3)
        end
    end

    user_input = (; (Symbol(k) => v for (k,v) in user_input_dict)...)
    params = T(;user_input...)

end

function get_obs!(obs::AbstractVector{T},
                  state::AbstractArray{T,3},
                  ist::AbstractVector{Int},
                  jst::AbstractVector{Int},
                  params::ModelParameters) where T

    get_obs!(obs,state,params.nx,ist,jst)

end

# Return observation data at stations from given model state
function get_obs!(obs::AbstractVector{T},
                  state::AbstractArray{T,3},
                  nx::Integer,
                  ist::AbstractVector{Int},
                  jst::AbstractVector{Int}) where T
    @assert length(obs) == length(ist) == length(jst)

    for i in eachindex(obs)
        ii = ist[i]
        jj = jst[i]
        iptr = (jj - 1) * nx + ii
        obs[i] = state[iptr]
    end
end

function tsunami_update!(dx_buffer::AbstractMatrix{T},
                         dy_buffer::AbstractMatrix{T},
                         state::AbstractArray{T,3},
                         model_matrices::LLW2d.Matrices{T},
                         params::ModelParameters) where T

    tsunami_update!(dx_buffer, dy_buffer, state, params.n_integration_step,
                    params.dx, params.dy, params.time_step, model_matrices)

end

# Update tsunami wavefield with LLW2d in-place.
function tsunami_update!(dx_buffer::AbstractMatrix{T},
                         dy_buffer::AbstractMatrix{T},
                         state::AbstractArray{T,3},
                         nt::Int,
                         dx::Real,
                         dy::Real,
                         time_interval::Real,
                         model_matrices::LLW2d.Matrices{T}) where T

    eta_a = @view(state[:, :, 1])
    mm_a  = @view(state[:, :, 2])
    nn_a  = @view(state[:, :, 3])
    eta_f = @view(state[:, :, 1])
    mm_f  = @view(state[:, :, 2])
    nn_f  = @view(state[:, :, 3])

    dt = time_interval / nt

    for it in 1:nt
        # Parts of model vector are aliased to tsunami heiht and velocities
        LLW2d.timestep!(dx_buffer, dy_buffer, eta_f, mm_f, nn_f, eta_a, mm_a, nn_a, model_matrices, dx, dy, dt)
    end

end

struct RandomField{F<:GaussianRandomField,X<:AbstractArray,W<:AbstractArray,Z<:AbstractArray}
    grf::F
    xi::X
    w::W
    z::Z
end

@metadata name ("","","") NTuple{3, String}
@metadata unit ("","","") NTuple{3, String}
@metadata description ("","","") NTuple{3, String}

@name @description @unit struct StateVectors{T<:AbstractArray, S<:AbstractArray}

    particles::T | ("height","vx","vy") | ("Ocean surface height","Ocean surface velocity x-component","Ocean surface velocity y-component") | ("m","m/s","m/s")
    truth::S | ("height","vx","vy") | ("Ocean surface height","Ocean surface velocity x-component","Ocean surface velocity y-component") | ("m","m/s","m/s")

end

struct ObsVectors{T<:AbstractArray,S<:AbstractArray}

    truth::T
    model::S

end

struct StationVectors{T<:AbstractArray}

    ist::T
    jst::T

end

function get_axes(params::ModelParameters)

    return get_axes(params.nx, params.ny, params.dx, params.dy)

end

function get_axes(nx::Int, ny::Int, dx::Real, dy::Real)

    x = range(0, length=nx, step=dx)
    y = range(0, length=ny, step=dy)

    return x,y
end

function init_gaussian_random_field_generator(params::ModelParameters)

    x, y = get_axes(params)
    return init_gaussian_random_field_generator(params.lambda,params.nu, params.sigma, x, y, params.padding, params.primes)

end

# Initialize a gaussian random field generating function using the Matern covariance kernel
# and circulant embedding generation method
# TODO: Could generalise this
function init_gaussian_random_field_generator(lambda::Vector{T},
                                              nu::Vector{T},
                                              sigma::Vector{T},
                                              x::AbstractVector{T},
                                              y::AbstractVector{T},
                                              pad::Int,
                                              primes::Bool) where T

    # Let's limit ourselves to two-dimensional fields
    dim = 2

    function _generate(l, n, s)
        cov = CovarianceFunction(dim, Matern(l, n, σ = s))
        grf = GaussianRandomField(cov, CirculantEmbedding(), x, y, minpadding=pad, primes=primes)
        v = grf.data[1]
        xi = Array{eltype(grf.cov)}(undef, size(v)..., nthreads())
        w = Array{complex(float(eltype(v)))}(undef, size(v)..., nthreads())
        z = Array{eltype(grf.cov)}(undef, length.(grf.pts)..., nthreads())
        RandomField(grf, xi, w, z)
    end

    return [_generate(l, n, s) for (l, n, s) in zip(lambda, nu, sigma)]
end

# Get a random sample from random_field_generator using random number generator rng
function sample_gaussian_random_field!(field::AbstractMatrix{T},
                                       random_field_generator::RandomField,
                                       rng::Random.AbstractRNG) where T

    @. @view(random_field_generator.xi[:,:,threadid()]) = randn((rng,), T)
    sample_gaussian_random_field!(field, random_field_generator, @view(random_field_generator.xi[:,:,threadid()]))

end

# Get a random sample from random_field_generator using random_numbers
function sample_gaussian_random_field!(field::AbstractMatrix{T},
                                       random_field_generator::RandomField,
                                       random_numbers::AbstractArray{T}) where T

    field .= GaussianRandomFields._sample!(@view(random_field_generator.w[:,:,threadid()]),
                                           @view(random_field_generator.z[:,:,threadid()]),
                                           random_field_generator.grf,
                                           random_numbers)

end

# Add a gaussian random field to the height in the state vector of all particles
function add_random_field!(state::AbstractArray{T,4},
                           field_buffer::AbstractArray{T,4},
                           generators::Vector{<:RandomField},
                           rng::Random.AbstractRNG,
                           nvar::Int,
                           nprt::Int) where T

    Threads.@threads for ip in 1:nprt

        for ivar in 1:nvar

            sample_gaussian_random_field!(@view(field_buffer[:, :, 1, threadid()]), generators[ivar], rng)
            @view(state[:, :, ivar, ip]) .+= @view(field_buffer[:, :, 1, threadid()])

        end

    end

end

function add_noise!(vec::AbstractVector{T}, rng::Random.AbstractRNG, params::ModelParameters) where T

    add_noise!(vec, rng, 0.0, params.obs_noise_std)

end

# Add a (mean, std) normal distributed random number to each element of vec
function add_noise!(vec::AbstractVector{T}, rng::Random.AbstractRNG, mean::T, std::T) where T

    d = truncated(Normal(mean, std), 0.0, Inf)
    @. vec += rand((rng,), d)

end

function init_arrays(params::ModelParameters, nprt_per_rank)

    return init_arrays(params.nx, params.ny, params.n_state_var, params.nobs, nprt_per_rank)

end

function init_arrays(nx::Int, ny::Int, n_state_var::Int, nobs::Int, nprt_per_rank::Int)

    # TODO: ideally this will be an argument of the function, to choose a
    # different datatype.
    T = Float64

    state_avg = zeros(T, nx, ny, n_state_var) # average of particle state vectors
    state_var = zeros(T, nx, ny, n_state_var) # variance of particle state vectors

    # Model vector for data assimilation
    #   state[:, :, 1, :]: tsunami height eta(nx,ny)
    #   state[:, :, 2, :]: vertically integrated velocity Mx(nx,ny)
    #   state[:, :, 3, :]: vertically integrated velocity Mx(nx,ny)
    state_particles = zeros(T, nx, ny, n_state_var, nprt_per_rank)
    state_truth = zeros(T, nx, ny, n_state_var) # model vector: true wavefield (observation)
    obs_truth = Vector{T}(undef, nobs)          # observed tsunami height
    obs_model = Matrix{T}(undef, nobs, nprt_per_rank) # forecasted tsunami height

    # station location in digital grids
    ist = Vector{Int}(undef, nobs)
    jst = Vector{Int}(undef, nobs)

    # Buffer array to be used in the tsunami update
    field_buffer = Array{T}(undef, nx, ny, 2, nthreads())

    return StateVectors(state_particles, state_truth), ObsVectors(obs_truth, obs_model), StationVectors(ist, jst), field_buffer
end

function set_initial_state!(states::StateVectors, model_matrices::LLW2d.Matrices{T},
                            field_buffer::AbstractArray{T, 4},
                            rng::Random.AbstractRNG,
                            nprt_per_rank::Int,
                            params::ModelParameters) where T

    # Set true initial state
    LLW2d.initheight!(@view(states.truth[:, :, 1]), model_matrices, params.dx, params.dy,
                      params.source_size, params.peak_height, params.peak_position)

    # Create generator for the initial random field
    x,y = get_axes(params)
    initial_grf = init_gaussian_random_field_generator(params.lambda_initial_state,
                                                       params.nu_initial_state,
                                                       params.sigma_initial_state,
                                                       x,
                                                       y,
                                                       params.padding,
                                                       params.primes)

    # Since states.particles is initially created as `zeros` we don't need to set it to 0 here
    # to get the default behaviour

    if params.particle_initial_state == "true"
        states.particles .= states.truth
    end

    # Add samples of the initial random field to all particles
    add_random_field!(states.particles, field_buffer, initial_grf, rng, params.n_state_var, nprt_per_rank)

end

function set_stations!(stations::StationVectors, params::ModelParameters) where T

    set_stations!(stations.ist,
                  stations.jst,
                  params.station_filename,
                  params.station_distance_x,
                  params.station_distance_y,
                  params.station_boundary_x,
                  params.station_boundary_y,
                  params.dx,
                  params.dy)

end

function set_stations!(ist::AbstractVector, jst::AbstractVector, filename::String, distance_x::T, distance_y::T, boundary_x::T, boundary_y::T, dx::T, dy::T) where T

    if filename != ""
        coords = readdlm(filename, ',', Float64, '\n'; comments=true, comment_char='#')
        ist .= floor.(Int, coords[:,1] / dx)
        jst .= floor.(Int, coords[:,2] / dy)
    else
        LLW2d.set_stations!(ist,jst,distance_x,distance_y,boundary_x,boundary_y,dx,dy)
    end

end

struct ModelData{A,B,C,D,E,F,G,H}
    model_params::A
    states::B
    observations::C
    stations::D
    field_buffer::E
    background_grf::F
    model_matrices::G
    rng::H
end
ParticleDA.get_particles(d::ModelData) = d.states.particles
# TODO: we should probably get rid of `get_truth`: it is only used as return
# value of `particle_filter`, we may just return the whole `model_data`.
ParticleDA.get_truth(d::ModelData) = d.states.truth
ParticleDA.get_stations(d::ModelData) = (nst = d.model_params.nobs,
                                         ist = d.stations.ist,
                                         jst = d.stations.jst)
ParticleDA.get_obs_noise_std(d::ModelData) = d.model_params.obs_noise_std
ParticleDA.get_model_noise_params(d::ModelData) = Matern(d.model_params.lambda[1],
                                                         d.model_params.nu[1],
                                                         σ=d.model_params.sigma[1])

function ParticleDA.set_particles!(d::ModelData, particles::AbstractArray{T}) where T

    d.states.particles .= particles

end
ParticleDA.get_grid_size(d::ModelData) = d.model_params.nx,d.model_params.ny
ParticleDA.get_grid_domain_size(d::ModelData) = d.model_params.x_length,d.model_params.y_length
ParticleDA.get_grid_cell_size(d::ModelData) = d.model_params.dx,d.model_params.dy
ParticleDA.get_n_state_var(d::ModelData) = d.model_params.n_state_var

function init(model_params_dict::Dict, nprt_per_rank::Int, my_rank::Integer, rng::Random.AbstractRNG)

    model_params = ParticleDA.get_params(ModelParameters, get(model_params_dict, "llw2d", Dict()))
    states, observations, stations, field_buffer = init_arrays(model_params, nprt_per_rank)

    background_grf = init_gaussian_random_field_generator(model_params)

    # Set up tsunami model
    model_matrices = LLW2d.setup(model_params.nx,
                                 model_params.ny,
                                 model_params.bathymetry_setup,
                                 model_params.absorber_thickness_fraction,
                                 model_params.boundary_damping,
                                 model_params.cutoff_depth)

    set_stations!(stations, model_params)

    set_initial_state!(states, model_matrices, field_buffer, rng, nprt_per_rank, model_params)

    return ModelData(model_params, states, observations, stations, field_buffer, background_grf, model_matrices, rng)
end

function ParticleDA.update_truth!(d::ModelData, _)
    tsunami_update!(@view(d.field_buffer[:, :, 1, 1]),
                    @view(d.field_buffer[:, :, 2, 1]),
                    d.states.truth, d.model_matrices, d.model_params)

    # Get observation from true synthetic wavefield
    get_obs!(d.observations.truth, d.states.truth, d.stations.ist, d.stations.jst, d.model_params)
    return d.observations.truth
end

function ParticleDA.update_particle_dynamics!(d::ModelData, nprt_per_rank)
    # Update dynamics
    Threads.@threads for ip in 1:nprt_per_rank
        tsunami_update!(@view(d.field_buffer[:, :, 1, threadid()]), @view(d.field_buffer[:, :, 2, threadid()]),
                        @view(d.states.particles[:, :, :, ip]), d.model_matrices, d.model_params)
    end
end

function ParticleDA.update_particle_noise!(d::ModelData, nprt_per_rank)
    # Add process noise
    add_random_field!(d.states.particles,
                      d.field_buffer,
                      d.background_grf,
                      d.rng,
                      d.model_params.n_state_var,
                      nprt_per_rank)
end

function ParticleDA.get_particle_observations!(d::ModelData, nprt_per_rank)
    # get observations
    for ip in 1:nprt_per_rank
        get_obs!(@view(d.observations.model[:,ip]),
                 @view(d.states.particles[:, :, :, ip]),
                 d.stations.ist,
                 d.stations.jst,
                 d.model_params)
    end
    return d.observations.model
end

### Model IO

function write_params(output_filename, params)

    file = h5open(output_filename, "cw")

    if !haskey(file, params.title_params)

        group = create_group(file, params.title_params)

        fields = fieldnames(typeof(params));

        for field in fields

            attributes(group)[string(field)] = getfield(params, field)

        end

    else

        @warn "Write failed, group " * params.title_params * " already exists in " * file.filename * "!"

    end

    close(file)

end

function write_grid(output_filename, params)

    h5open(output_filename, "cw") do file

        if !haskey(file, params.title_grid)

            # Write grid axes
            group = create_group(file, params.title_grid)
            x,y = get_axes(params)
            #TODO: use d_write instead of create_dataset when they fix it in the HDF5 package
            ds_x,dtype_x = create_dataset(group, "x", collect(x))
            ds_y,dtype_x = create_dataset(group, "y", collect(x))
            ds_x[1:params.nx] = collect(x)
            ds_y[1:params.ny] = collect(y)
            attributes(ds_x)["Unit"] = "m"
            attributes(ds_y)["Unit"] = "m"

        else

            @warn "Write failed, group " * params.title_grid * " already exists in " * file.filename * "!"

        end

    end

end

function write_stations(output_filename, ist::AbstractVector, jst::AbstractVector, params::ModelParameters) where T

    h5open(output_filename, "cw") do file

        if !haskey(file, params.title_stations)
            group = create_group(file, params.title_stations)

            for (dataset_name, val) in zip(("x", "y"), (ist .* params.dx, jst .* params.dy))
                ds, dtype = create_dataset(group, dataset_name, val)
                ds[:] = val
                attributes(ds)["Description"] = "Station "*dataset_name*" coordinate"
                attributes(ds)["Unit"] = "m"
            end
        else
            @warn "Write failed, group " * params.title_stations * " already exists in " * file.filename * "!"
        end
    end
end

function write_weights(file::HDF5.File, weights::AbstractVector, unit::String, it::Int, params::ModelParameters)

    group_name = "weights"
    dataset_name = "t" * lpad(string(it),4,'0')

    group, subgroup = ParticleDA.create_or_open_group(file, group_name)

    if !haskey(group, dataset_name)
        #TODO: use d_write instead of create_dataset when they fix it in the HDF5 package
        ds,dtype = create_dataset(group, dataset_name, weights)
        ds[:] = weights
        attributes(ds)["Description"] = "Particle Weights"
        attributes(ds)["Unit"] = unit
        attributes(ds)["Time step"] = it
        attributes(ds)["Time (s)"] = it * params.time_step
    else
        @warn "Write failed, dataset " * group_name * "/" * dataset_name *  " already exists in " * file.filename * "!"
    end

end

function ParticleDA.write_snapshot(output_filename::AbstractString,
                                   d::ModelData,
                                   avg::AbstractArray{T,3},
                                   var::AbstractArray{T,3},
                                   weights::AbstractVector{T},
                                   it::Int) where T

    if it == 0
        # These are written only at the initial state it == 0
        write_grid(output_filename, d.model_params)
        write_params(output_filename, d.model_params)
        write_stations(output_filename, d.stations.ist, d.stations.jst, d.model_params)
    end

    if any(d.model_params.particle_dump_time .== it)
        write_particles(d.model_params.particle_dump_file, d.states, it, d.model_params)
    end

    return ParticleDA.write_snapshot(output_filename, d.states, avg, var, weights, it, d.model_params)
end

function ParticleDA.write_snapshot(output_filename::AbstractString,
                                   states::StateVectors,
                                   avg::AbstractArray{T,3},
                                   var::AbstractArray{T,3},
                                   weights::AbstractVector{T},
                                   it::Int,
                                   params::ModelParameters) where T

    println("Writing output at timestep = ", it)

    h5open(output_filename, "cw") do file

        for (i,(name,desc,unit)) in enumerate(zip(name(states, :truth), description(states, :truth), unit(states, :truth)))

            write_field(file, @view(states.truth[:,:,i]), it, unit, params.title_syn, name, desc, params)
            write_field(file, @view(avg[:,:,i]), it, unit, params.title_avg, name, desc, params)
            write_field(file, @view(var[:,:,i]), it, "("*unit*")^2", params.title_var, name, desc, params)

        end

        write_weights(file, weights, "", it, params)
    end

end

function write_particles(output_filename::AbstractString,
                         states::StateVectors,
                         it::Int,
                         params::ModelParameters) where T

    println("Writing particle states at timestep = ", it)
    nprt = size(states.particles,4)

    h5open(output_filename, "cw") do file

        for iprt = 1:nprt
            group_name = "particle" * string(iprt)

            for (i,(name,desc,unit)) in enumerate(zip(name(states, :particles), description(states, :particles), unit(states, :particles)))

                write_field(file, @view(states.particles[:,:,i,iprt]), it, unit, group_name, name, desc, params)

            end

        end

    end

end

function write_field(file::HDF5.File,
                     field::AbstractMatrix{T},
                     it::Int,
                     unit::String,
                     group::String,
                     dataset::String,
                     description::String,
                     params::ModelParameters) where T

    group_name = params.state_prefix * "_" * group
    subgroup_name = "t" * lpad(string(it),4,'0')
    dataset_name = dataset

    group, subgroup = ParticleDA.create_or_open_group(file, group_name, subgroup_name)

    if !haskey(subgroup, dataset_name)
        #TODO: use d_write instead of create_dataset when they fix it in the HDF5 package
        ds,dtype = create_dataset(subgroup, dataset_name, field)
        ds[:,:] = field
        attributes(ds)["Description"] = description
        attributes(ds)["Unit"] = unit
        attributes(ds)["Time step"] = it
        attributes(ds)["Time (s)"] = it * params.time_step
    else
        @warn "Write failed, dataset " * group_name * "/" * subgroup_name * "/" * dataset_name *  " already exists in " * file.filename * "!"
    end
end

end # module


text:
"""
    ScriptTagProxy(io)

This is a transparent proxy that ensures neither `<!--` nor `</script>`
occur in the output stream.

# Examples
```julia-repl
julia> gp = ScriptTagProxy(stdout);
julia> print(gp, "valid");
valid
julia> print(gp, "</script>")
ERROR: "Content within a script tag must not contain `</script>`"
```
"""
mutable struct ScriptTagProxy{T<:IO} <: IOProxy where {T}
    io::T
    index::Int

    ScriptTagProxy(io::T) where T = new{T}(io::T, 0)
end

"""
    ScriptTag(data)

This object prints `data` unescaped within a `<script>` tag, wrapped in
a `ScriptTagProxy` that guards against invalid script content.
"""
struct ScriptTag
    content
end

Base.print(ep::EscapeProxy, x::ScriptTag) =
    print_script_hook(ScriptTagProxy(ep.io), x.content)

"""
    Script(data)

This object renders Javascript `data` escaped within an attribute.
"""
mutable struct Script
    content::Any
end

Base.print(io::IO, x::Script) =
    print_script_hook(io, x.content)

Base.show(io::IO, ::MIME"text/javascript", sv::Script) =
    print(io, sv)

"""
    js(x)

This method may be implemented to specify a printed representation
suitable for use within a quoted attribute value starting with `on`.
"""
js(x) = Script(x)

"""
    JavaScript(js) - shows `js` as `"text/javascript"`
"""
struct JavaScript
    content
end

Base.show(io::IO, ::MIME"text/javascript", js::JavaScript) =
    print(io, js.content)

"""
    print_script_hook(io, value)

Provides a hook to override `print_script` for custom Javascript
runtimes, such as `Pluto.jl`, to provide their own value marshalling.
"""
print_script_hook(io::IO, value) =
   print_script(io, value)

"""
    print_script(io, value)

Show `value` as `"text/javascript"` to the given `io`, this provides
some baseline functionality for built-in data types.

    - `nothing` becomes `undefined`
    - `missing` becomes `null`
    - `Bool` values are printed as `true` or `false`
    - `AbstractString` and `Symbol` become a double-quoted string
    - `AbstractVector` and `Tuple` become an array
    - `Dict` and `NamedTuple` become a Javascript object, with
       keys converted to string values
    - `AbstractFloat` and `Integer` are printed directly, where
      `NaN` remains `NaN` but `Inf` is printed as `Infinity`

The fallback behavior of `print_script` is to show the object as
`"text/javascript"`. The `Javascript` wrapper will take any string
and let it be printed in this way.
"""
@generated function print_script(io::IO, value)
     # workaround for Julia#18221
     if hasmethod(show, Tuple{IO, MIME{Symbol("text/javascript")}, value})
         return :(show(io, MIME"text/javascript"(), value))
     end
     throw("$(value) is not showable as text/javascript")
end
print_script(io::IO, ::Nothing) =
    print(io, "undefined")
print_script(io::IO, ::Missing) =
    print(io, "null")
print_script(io::IO, value::Union{Bool, Symbol}) =
    print(io, value)

function print_script(io::IO, value::Union{NamedTuple, AbstractDict})
    print(io, '{')
    first = true
    for (k,v) in pairs(value)
        if !first
            print(io, ", ")
        end
        print_script(io, string(k))
        print(io, ": ")
        print_script(io, v)
        first = false
    end
    print(io, '}')
end

function print_script(io::IO, value::Union{Tuple, AbstractVector})
    print(io, '[')
    first = true
    for item in value
        if !first
            print(io, ", ")
        end
        print_script(io, item)
        first = false
    end
    print(io, ']')
end

function print_script(io::IO, value::Union{Integer, AbstractFloat})
    if isfinite(value) || isnan(value)
        print(io, value)
    else
        if value < 0
            print(io, "-")
        end
        print(io, "Infinity")
    end
end

function print_script(io::IO, value::AbstractString)
    final = lastindex(value)
    i = last = 1
    function emit(s::String)
        print(io, SubString(value, last, prevind(value, i)))
        last = nextind(value, i)
        print(io, s)
    end
    print(io, "\"")
    while i <= final
        ch = value[i]
        if ch === '\n'
            emit("\\n")
        elseif ch === '\r'
            emit("\\r")
        elseif ch === '\\'
            emit("\\\\")
        elseif ch === '\"'
            emit("\\\"")
        elseif ch === '\u2028'
            emit("\\u2028")
        elseif ch === '\u2029'
            emit("\\u2029")
        elseif ch === '<' && i+1 <= final
            # escape nested script and comment tags
            nc = value[i+1]
            if nc in ('s', 'S', '!', '/')
                emit("<\\")
            end
        end
        i = nextind(value, i)
    end
    print(io, SubString(value, last, final))
    print(io, "\"")
end

function scan_for_script(index::Int, octet::UInt8)::Int
    if 1 == index
        return octet == Int('!') ? 12 : octet == Int('/') ? 2 : 0
    elseif 2 == index
        return (octet == Int('S') || octet == Int('s')) ? 3 : 0
    elseif 3 == index
        return (octet == Int('C') || octet == Int('c')) ? 4 : 0
    elseif 4 == index
        return (octet == Int('R') || octet == Int('r')) ? 5 : 0
    elseif 5 == index
        return (octet == Int('I') || octet == Int('i')) ? 6 : 0
    elseif 6 == index
        return (octet == Int('P') || octet == Int('p')) ? 7 : 0
    elseif 7 == index
        return (octet == Int('T') || octet == Int('t')) ? 8 : 0
    elseif 8 == index
        if octet == Int('>')
            throw("Content within a script tag must not contain `</script>`")
        end
    elseif 12 == index
        return octet == Int('-') ? 13 : 0
    elseif 13 == index
        if octet == Int('-')
            throw("Content within a script tag must not contain `<!--`")
        end
    else
        @assert false # unreachable?!
    end
    return 0
end

function Base.write(sp::ScriptTagProxy, octet::UInt8)
    if 0 == sp.index
        if octet == Int('<')
            sp.index = 1
        end
        return write(sp.io, octet)
    end
    sp.index = scan_for_script(sp.index, octet)
    return write(sp.io, octet)
end

function Base.unsafe_write(sp::ScriptTagProxy, input::Ptr{UInt8}, nbytes::UInt)
    cursor = input
    index = sp.index
    final = input + nbytes
    while cursor < final
        octet = unsafe_load(cursor)
        if octet == Int('<')
            index = 1
        elseif 0 != index
            index = scan_for_script(index, octet)
        end
        cursor += 1
    end
    sp.index = index
    return unsafe_write(sp.io, input, nbytes)
end


text:
#!/usr/local/bin/julia
# last tested Julia version: 1.6.1
# Hindmarsh-Rose model on a 2D lattice
# FvW 03/2018

using NPZ
using PyCall
using PyPlot
using Statistics
using VideoIO
@pyimport matplotlib.animation as anim

function animate_pyplot(fname, data)
    """
    Animate 3D array as .mp4 using PyPlot, save as `fname`
    array dimensions:
        1: time
        2, 3: space
    """
    nt, nx, ny = size(data)
    vmin = minimum(data)
    vmax = maximum(data)
    # setup animation image
    fig = figure(figsize=(6,6))
    axis("off")
    t = imshow(data[1,:,:], origin="lower", cmap=ColorMap("gray"),
			   vmin=vmin, vmax=vmax)
    tight_layout()
    # frame generator
    println("[+] animate")
    function animate(i)
        (i%100 == 0) && print("    t = ", i, "/", nt, "\r")
        t.set_data(data[i+1,:,:])
    end
    # create animation
    ani = anim.FuncAnimation(fig, animate, frames=nt, interval=10)
    println("\n")
    # save animation
    ani[:save](fname, bitrate=-1,
               extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p"])
    show()
end

function animate_video(fname, data)
    """
    Animate 3D array as .mp4 using VideoIO, save as `fname`
    array dimensions:
        1: time
        2, 3: space
    """
    # BW
    y = UInt8.(round.(255*(data .- minimum(data))/(maximum(data)-minimum(data))))
    # BW inverted
    #y = UInt8.(round.(255 .- 255*(data .- minimum(data))/(maximum(data)-minimum(data))))
    encoder_options = (color_range=2, crf=0, preset="medium")
    framerate=90
    T = size(data,1)
    open_video_out(fname, y[1,end:-1:1,:], framerate=framerate,
                   encoder_options=encoder_options) do writer
        for i in range(2,stop=T,step=1)
            write(writer, y[i,end:-1:1,:])
        end
    end
end

function L(x)
    # Laplace operator
    # periodic boundary conditions
    xU = circshift(x, [-1 0])
    xD = circshift(x, [1 0])
    xL = circshift(x, [0 -1])
    xR = circshift(x, [0 1])
    Lx = xU + xD + xL + xR - 4x
    # non-periodic boundary conditions
    Lx[1,:] .= 0.0
    Lx[end,:] .= 0.0
    Lx[:,1] .= 0.0
    Lx[:,end] .= 0.0
    return Lx
end

function hr2d(N, T, t0, dt, sd, D, a, b, c, d, s, r, x_1, I0, stim, blocks)
	# initialize Hindmarsh-Rose system
	x, y, z = zeros(Float64,N,N), zeros(Float64,N,N), zeros(Float64,N,N)
	dxdt, dydt, dzdt = zeros(Float64,N,N), zeros(Float64,N,N), zeros(Float64,N,N)
	sd_sqrt_dt = sd*sqrt(dt)
	X = zeros(Float64,T,N,N)
	I = zeros(Float64,t0+T,N,N)
	# stimulation protocol
	for st in stim
		t_on, t_off = st[1]
		x0, x1 = st[2]
		y0, y1 = st[3]
		I[t0+t_on:t0+t_off, x0:x1, y0:y1] .= I0
		I[1:t0, x0:x1, y0:y1] .= I0
	end
	# iterate
	for t in range(1, stop=t0+T, step=1)
		(t%100 == 0) && print("    t = ", t, "\r")
		# Hindmarsh-Rose equations
		dx = -a*x.*x.*x + b*x.*x + y - z + I[t,:,:] + D*L(x)
		dy = c .- d*x.*x - y
		dz = r*(s*(x .- x_1) - z)
		# Ito stochastic integration
		x += (dx*dt + sd_sqrt_dt*randn(N,N))
		y += (dy*dt)
		z += (dz*dt)
		# dead block(s):
		for bl in blocks
			x[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
			y[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
			z[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
		end
		(t > t0) && (X[t-t0,:,:] = x)
	end
	println("\n")
	return X
end

function L(x)
	# Laplace operator
    xU = circshift(x, [-1 0])
    xD = circshift(x, [1 0])
    xL = circshift(x, [0 -1])
    xR = circshift(x, [0 1])
    Lx = xU + xD + xL + xR - 4x
	nx = size(x)[1]
	ny = size(x)[2]
    Lx[1,:] .= 0.0
	Lx[nx,:] .= 0.0
	Lx[:,1] .= 0.0
	Lx[:,ny] .= 0.0
    return Lx
end

function main()
    println("[+] Hindmarsh-Rose 2D lattice:")
	N = 64
    T = 15000
	t0 = 2500
	dt = 0.05
	sd = 0.05
	D = 2.0
	a = 1.0
	b = 3.0
	c = 1.0
	d = 5.0
	s = 2.0 # s=1: spiking, s=4: strong accomodation
	r = 0.001
    x_1 = -1.6
	I0 = 3.50 # 1.3
	#I1 = 0.05 # 1.3
    println("[+] Lattice size N: ", N)
	println("[+] Time steps T: ", N)
    println("[+] Warm-up steps t0: ", t0)
	println("[+] Integration time step dt: ", dt)
	println("[+] Noise intensity sd: ", sd)
	println("[+] Diffusion constant D: ", D)
    println("[+] HR parameter a: ", a)
    println("[+] HR parameter b: ", b)
    println("[+] HR parameter c: ", c)
    println("[+] HR parameter d: ", d)
    println("[+] HR parameter s: ", s)
    println("[+] HR parameter r: ", r)
    println("[+] HR parameter x_1: ", x_1)
    println("[+] Stimulation current I0: ", I0)
    #println("[+] Stimulation current I1: ", I1)

	# stim protocol, array of elements [[t0,t1], [x0,x1], [y0,y1]]
	stim = [ [[1,15000], [5,N-5], [5,N-5]] ]
    #stim = []# dead blocks, array of elementy [[x0,x1], [y0,y1]]

	# dead blocks, array of elementy [[x0,x1], [y0,y1]]
	#blocks = [ [[2*n_4,3*n_4], [15,20]], [[2*n_4+10,3*n_4+10], [40,45]] ]
	blocks = []

	# run simulation
	data = hr2d(N, T, t0, dt, sd, D, a, b, c, d, s, r, x_1, I0, stim, blocks)
    println("[+] Data dimensions: ", size(data))

	# plot mean voltage
    m = mean(reshape(data, (T,N*N)), dims=2)
    plot(m, "-k"); show()

	# save data
	I_str = rpad(I0, 4, '0') # stim. current amplitude as 4-char string
	s_str = rpad(s, 4, '0') # noise as 4-char string
	sd_str = rpad(sd, 4, '0') # noisy intensity as 4-char string
    D_str = rpad(D, 4, '0') # diffusion coefficient as 4-char string
	fname1 = string("hr2d_I_", I_str, "_s_", s_str, "_sd_", sd_str,
					"_D_", D_str, ".npy")
    #npzwrite(data_filename, data)
	#println("[+] Data saved as: ", data_filename)

	# video
	fname2 = string("hr2d_I_", I_str, "_s_", s_str, "_sd_", sd_str,
					"_D_", D_str, ".mp4")
	#animate_pyplot(fname2, data) # slow
    animate_video(fname2, data) # fast
	println("[+] Data saved as: ", fname2)
end

main()