############################################################################################
"""
$(TYPEDEF)

Default arguments for PMCMC constructor.

# Fields
$(TYPEDFIELDS)
"""
struct ParticleMetropolis{P<:ParticleFilter,M<:MCMC} <: PMCMCKernel
    "Particle Filter kernel to estimate latent trajectory."
    pf::P
    "MCMC kernel to sample continuous model parameter."
    mcmc::M
    function ParticleMetropolis(pf::P, mcmc::M) where {P<:ParticleFilter,M<:MCMC}
        @argcheck isa(pf.tune.referencing, Marginal) "Cannot use Conditional Filter in Particle Metropolis setting, use Marginal referencing"
        #!NOTE: Only create warning, but let user decide if they want proceed.
        #@argcheck isa(mcmc.kernel, Metropolis) "Only gradient free mcmc kernels allowed as tuning is based on log-target approximation from particle filter instead of objective(val). For gradient based mcmc kernels, use ParticleGibbs instead.")
        if isa(mcmc.tune.stepsize.adaption, BaytesCore.UpdateTrue)
            println("If no objective(val) specified, it is recommended to use a fixed stepsize for the mcmc kernel, as otherwise the kernel will be tuned based on acceptance rate of mcmc kernel.")
        end
        return new{P,M}(pf, mcmc)
    end
end
function ParticleMetropolis(
    filter::F, mcmc::M; kwargs...
) where {F<:ParticleFilterConstructor,M<:MCMCConstructor}
    return PMCMC(ParticleMetropolis, filter, mcmc; kwargs...)
end

############################################################################################
"""
$(SIGNATURES)
Propose new parameter with pmcmc sampler. If update=true, objective function will be updated with input model and data.

# Examples
```julia
```

"""
function propose!(
    _rng::Random.AbstractRNG,
    pmcmc::ParticleMetropolis,
    model::ModelWrapper,
    data::D,
    temperature::F = model.info.reconstruct.default.output(1.0),
    update::U=BaytesCore.UpdateTrue(),
) where {D,F<:AbstractFloat,U<:BaytesCore.UpdateBool}
    ## Compute initial logposterior and save initial model value
    ℓpostₜ =
        pmcmc.pf.particles.ℓobjective.cumulative +
        ModelWrappers.log_prior_with_transform(model, pmcmc.mcmc.tune.tagged)
    val = deepcopy(model.val)
    ## Update Objective with new model parameter from other MCMC samplers and/or new/latent data
    objective = Objective(model, data, pmcmc.mcmc.tune.tagged, temperature)
    ## Update Kernel with current objective/configs
    update!(pmcmc.mcmc.kernel, objective, update)
    ## Make MCMC Proposal step
    resultᵖ, divergent, accept, sampler_statistic = propagate(
        _rng, pmcmc.mcmc.kernel, pmcmc.mcmc.tune, objective
    )
    ## Make PF proposal step with proposed parameter
    if !divergent
        ModelWrappers.unflatten_constrain!(
            objective.model, pmcmc.mcmc.tune.tagged, resultᵖ.θᵤ
        )
    end
    _, pf_diagnostics = propose!(_rng, pmcmc.pf, objective.model, objective.data, temperature, update)
    ## Compute acceptance rate
    ℓpostₜᵖ =
        pmcmc.pf.particles.ℓobjective.cumulative +
        ModelWrappers.log_prior_with_transform(objective.model, pmcmc.mcmc.tune.tagged)
    acceptᵖ = BaytesCore.AcceptStatistic(_rng, model.info.reconstruct.default.output(ℓpostₜᵖ - ℓpostₜ))
    if acceptᵖ.accepted
        ## If accepted, update MCMC kernel
        pmcmc.mcmc.kernel.result = resultᵖ
    else
        ## Else set model back to initial value
        model.val = val
    end
    #!NOTE: Would need to acceptᵖ here, but makes sampler very noisy. Better to print to either use fixed stepsize or assign objective(val).
    update!(pmcmc.mcmc.tune, pmcmc.mcmc.kernel.result, accept.rate)
    mcmc_diagnostics = MCMCDiagnostics(
        BaytesCore.BaseDiagnostics(
            pmcmc.mcmc.kernel.result.ℓθᵤ,
            objective.temperature,
            predict(_rng, objective),
            pmcmc.mcmc.tune.iter.current
        ),
        sampler_statistic,
        divergent,
        acceptᵖ,
        generate(_rng, objective, pmcmc.mcmc.tune.generated),
    )
    ## Assign base diagnostics - ℓobjective and predictions are taken from particle filter
    diagnostics = BaytesCore.BaseDiagnostics(
        pf_diagnostics.base.ℓobjective, pf_diagnostics.base.temperature, pf_diagnostics.base.prediction,
        mcmc_diagnostics.base.iter
    )
    ## Return pmcmc output
    return model.val, PMCMCDiagnostics(diagnostics, pf_diagnostics, mcmc_diagnostics)
end

############################################################################################
#export
export ParticleMetropolis, propose!


# v0.12 deprecations
@deprecate Dropout(p, dims) Dropout(p; dims=dims)
@deprecate InstanceNorm(λ, β, γ, μ, σ², ϵ, momentum, active=nothing) InstanceNorm(λ, β, γ, μ, σ², ϵ, momentum, true, true, active, length(β))
@deprecate BatchNorm(λ, β, γ, μ, σ², ϵ, momentum, active=nothing) BatchNorm(λ, β, γ, μ, σ², ϵ, momentum, true, true, active, length(β))
@deprecate GroupNorm(G, λ, β, γ, μ, σ², ϵ, momentum, active=nothing) GroupNorm(G, λ, β, γ, μ, σ², ϵ, momentum, true, true, active, length(β))
@deprecate outdims(f, inputsize) outputsize(f, inputsize)
@deprecate Conv(; weight,  bias, activation=identity, kws...) Conv(weight, bias, activation; kws...)
@deprecate ConvTranspose(; weight, bias, activation=identity, kws...) ConvTranspose(weight, bias, activation; kws...)
@deprecate DepthwiseConv(; weight, bias, activation=identity, kws...) DepthwiseConv(weight, bias, activation; kws...)

function Base.getproperty(a::Dense, s::Symbol)
  if s === :W
    Base.depwarn("field name dense.W is deprecated in favour of dense.weight", :Dense)
    return getfield(a, :weight)
  elseif s === :b
    Base.depwarn("field name dense.b is deprecated in favour of dense.bias", :Dense)
    return getfield(a, :bias)
  end
  return getfield(a, s)
end

function ones(dims...)
  Base.depwarn("Flux.ones(size...) is deprecated, please use Flux.ones32(size...) or Base.ones(Float32, size...)", :ones)
  Base.ones(Float32, dims...)
end
ones(T::Type, dims...) = Base.ones(T, dims...)

function zeros(dims...)
  Base.depwarn("Flux.zeros(size...) is deprecated, please use Flux.zeros32(size...) or Base.zeros(Float32, size...)", :zeros)
  Base.zeros(Float32, dims...)
end
zeros(T::Type, dims...) = Base.zeros(T, dims...)

ones32(::Type, dims...) = throw(ArgumentError("Flux.ones32 is always Float32, use Base.ones to specify the element type"))
zeros32(::Type, dims...) = throw(ArgumentError("Flux.zeros32 is always Float32, use Base.zeros to specify the element type"))


using Base.Test
using MeshCat
using GeometryTypes
using CoordinateTransformations
using Colors
using MeshIO, FileIO
using Polyhedra
using CDDLib

@testset "MeshCat" begin
    include("server_client.jl")
    include("visualizer.jl")
    include("notebook.jl")

    @testset "deprecations" begin
        # Deprecated in v0.0.2
        @testset "cylinder" begin
            l = 1.0
            r = 2.0
            c = HyperCylinder(l, r)
            @test c.origin == Point(0., 0, 0)
            @test c.extremity == Point(0., 0, l)
            @test radius(c) == r
        end
    end
end


"""
$(SIGNATURES)

Helper function for `convert_inter_html` that processes an extracted block given a latex context
`lxc` and returns the processed html that needs to be plugged in the final html.
"""
function convert_block(β::AbstractBlock, lxcontext::LxContext)::AS
    # case for special characters / html entities
    β isa HTML_SPCH     && return ifelse(isempty(β.r), β.ss, β.r)
    # Return relevant interpolated string based on case
    βn = β.name
    βn ∈  MD_HEADER        && return convert_header(β)
    βn == :CODE_INLINE     && return html_code_inline(content(β) |> htmlesc)
    βn == :CODE_BLOCK_LANG && return convert_code_block(β.ss)
    βn == :CODE_BLOCK_IND  && return convert_indented_code_block(β.ss)
    βn == :CODE_BLOCK      && return html_code(strip(content(β)), "{{fill lang}}")
    βn == :ESCAPE          && return chop(β.ss, head=3, tail=3)
    βn == :FOOTNOTE_REF    && return convert_footnote_ref(β)
    βn == :FOOTNOTE_DEF    && return convert_footnote_def(β, lxcontext)
    βn == :LINK_DEF        && return ""

    # Math block --> needs to call further processing to resolve possible latex
    βn ∈ MATH_BLOCKS_NAMES && return convert_math_block(β, lxcontext.lxdefs)

    # Div block --> need to process the block as a sub-element
    if βn == :DIV
        ct, _ = convert_md(content(β) * EOS, lxcontext.lxdefs;
                           isrecursive=true, has_mddefs=false)
        name = chop(otok(β).ss, head=2, tail=0)
        return html_div(name, ct)
    end

    # default case, ignore block (should not happen)
    return ""
end
convert_block(β::LxCom, λ::LxContext) = resolve_lxcom(β, λ.lxdefs)

"""
MATH_BLOCKS_PARENS

Dictionary to keep track of how math blocks are fenced in standard LaTeX and how these fences need
to be adapted for compatibility with KaTeX. Each tuple contains the number of characters to chop
off the front and the back of the maths expression (fences) as well as the KaTeX-compatible
replacement.
For instance, `\$ ... \$` will become `\\( ... \\)` chopping off 1 character at the front and the
back (`\$` sign).
"""
const MATH_BLOCKS_PARENS = Dict{Symbol, Tuple{Int,Int,String,String}}(
    :MATH_A     => ( 1,  1, "\\(", "\\)"),
    :MATH_B     => ( 2,  2, "\\[", "\\]"),
    :MATH_C     => ( 2,  2, "\\[", "\\]"),
    :MATH_ALIGN => (13, 11, "\\[\\begin{aligned}", "\\end{aligned}\\]"),
    :MATH_EQA   => (16, 14, "\\[\\begin{array}{c}", "\\end{array}\\]"),
    :MATH_I     => ( 4,  4, "", "")
    )

"""
$(SIGNATURES)

Helper function for the math block case of `convert_block` taking the inside of a math block,
resolving any latex command in it and returning the correct syntax that KaTeX can render.
"""
function convert_math_block(β::OCBlock, lxdefs::Vector{LxDef})::String
    # try to find the block out of `MATH_BLOCKS_PARENS`, if not found, error
    pm = get(MATH_BLOCKS_PARENS, β.name) do
        throw(MathBlockError("Unrecognised math block name."))
    end

    # convert the inside, decorate with KaTex and return, also act if
    # if the math block is a "display" one (with a number)
    inner = chop(β.ss, head=pm[1], tail=pm[2])
    htmls = IOBuffer()
    if β.name ∉ [:MATH_A, :MATH_I]
        # NOTE: in the future if allow equation tags, then will need an `if`
        # here and only increment if there's no tag. For now just use numbers.

        # increment equation counter
        PAGE_EQREFS[PAGE_EQREFS_COUNTER] += 1

        # check if there's a label, if there is, add that to the dictionary
        matched = match(r"\\label{(.*?)}", inner)

        if !isnothing(matched)
            name   = refstring(matched.captures[1])
            write(htmls, "")
            inner  = replace(inner, r"\\label{.*?}" => "")
            # store the label name and associated number
            PAGE_EQREFS[name] = PAGE_EQREFS[PAGE_EQREFS_COUNTER]
        end
    end
    # assemble the katex decorators, the resolved content etc
    write(htmls, pm[3], convert_md_math(inner * EOS, lxdefs, from(β)), pm[4])
    return String(take!(htmls))
end

"""
$(SIGNATURES)

Helper function for the case of a header block (H1, ..., H6).
"""
function convert_header(β::OCBlock)::String
    hk       = lowercase(string(β.name))
    title, _ = convert_md(content(β) * EOS; isrecursive=true, has_mddefs=false)
    # check if the header has appeared before
    rstitle  = refstring(title)
    level    = parse(Int, hk[2])
    occur    = (hv[3] for hv ∈ values(PAGE_HEADERS) if hv[2] == rstitle)
    occur    = isempty(occur) ? 0 : maximum(occur)
    rstitle  = ifelse(occur==0, rstitle, "$(rstitle)_$(occur+1)")
    # save in list of headers
    PAGE_HEADERS[length(PAGE_HEADERS)+1] = (title, rstitle, occur+1, level)
    # return the title
    return html_hk(hk, html_ahref_key(rstitle, title); id=rstitle)
end

"""
$(SIGNATURES)

Helper function for the code block case of `convert_block`.
"""
function convert_code_block(ss::SubString)::String
    fencer = ifelse(startswith(ss, "`````"), "`````", "```")
    reg    = Regex("$fencer([a-z-]*)(\\:[a-zA-Z\\\\\\/-_\\.]+)?\\s*\\n?((?:.|\\n)*)$fencer")
    m      = match(reg, ss)
    lang  = m.captures[1]
    rpath = m.captures[2]
    code  = m.captures[3]

    if isnothing(rpath)
        return html_code(code, lang)
    end
    if lang!="julia"
        @warn "Eval of non-julia code blocks is not yet supported."
        return html_code(code, lang)
    end
    # path currently has an indicative `:` we don't care about
    rpath = rpath[2:end]

    # Here we have a julia code block that was provided with a script path
    # It will consequently be
    #   1. written to script file unless it's already there
    #   2. evaled (unless a file was there and output file is present), redirect out
    #   3. inserted after scrapping out lines (see resolve_lx_input)
    path = resolve_assets_rpath(rpath; canonical=true)

    endswith(path, ".jl") || (path *= ".jl")

    out_path, fname = splitdir(path)
    out_path = mkpath(joinpath(out_path, "output"))
    out_name = splitext(fname)[1] * ".out"
    out_path = joinpath(out_path, out_name)

    # > 1.b check whether the file already exists and if so compare content
    do_eval = !isfile(path) || read(path, String) != code || !isfile(out_path)

    if do_eval
        write(path, "# This file was generated by JuDoc, do not modify it. # hide\n", code)
        # step 2: execute the code while redirecting the output to file (note that
        # everything is ran in one big sequence of scripts so in particular if there
        # are 3 code blocks on the same page then the second and third one can use
        # whatever is defined or loaded in the first (it also has access to scope of
        # other pages but it's really not recommended to exploit that)
        open(out_path, "w") do outf
            redirect_stdout(outf)  do
                try
                    Main.include(path)
                catch e
                    print("There was an error running the code: $(e.error).")
                end
            end
        end
    end

    # step 3, insertion of code stripping of "hide" lines.
    return resolve_lx_input_hlcode(rpath, "julia")
end

"""
$(SIGNATURES)

Helper function for the indented code block case of `convert_block`.
"""
function convert_indented_code_block(ss::SubString)::String
    # 1. decrease indentation of all lines (either frontal \n\t or \n⎵⎵⎵⎵)
    code = replace(ss, r"\n(?:\t| {4})" => "\n")
    # 2. return; lang is a LOCAL_PAGE_VARS that is julia by default and can be set
    sc = strip(code)
    isempty(sc) && return ""
    return html_code(sc, "{{fill lang}}")
end

"""
$(SIGNATURES)

Helper function to convert a `[^1]` into a html sup object with appropriate ref and backref.
"""
function convert_footnote_ref(β::Token)::String
    # β.ss is [^id]; extract id
    id = string(match(r"\[\^(.*?)\]", β.ss).captures[1])
    # add it to the list of refs unless it's been seen before
    pos = 0
    for (i, pri) in enumerate(PAGE_FNREFS)
        if pri == id
            pos = i
            break
        end
    end
    if pos == 0
        push!(PAGE_FNREFS, id)
        pos = length(PAGE_FNREFS)
    end
    return html_sup("fnref:$id", html_ahref(url_curpage() * "#fndef:$id", "[$pos]"; class="fnref"))
end

"""
$(SIGNATURES)

Helper function to convert a `[^1]: ...` into a html table for the def.
"""
function convert_footnote_def(β::OCBlock, lxcontext::LxContext)::String
    # otok(β) is [^id]:
    id = match(r"\[\^(.*?)\]:", otok(β).ss).captures[1]
    pos = 0
    for (i, pri) in enumerate(PAGE_FNREFS)
        if pri == id
            pos = i
            break
        end
    end
    if pos == 0
        # this was never referenced before, so probably best not to show it
        return ""
    end
    # need to process the content which could contain stuff
    ct, _ = convert_md(content(β) * EOS, lxcontext.lxdefs;
                       isrecursive=true, has_mddefs=false)
    """

            $(html_ahref(url_curpage() * "#fnref:$id", "[$pos]"))
            $(ct)

    """
end


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

using JuliaFEM
using JuliaFEM.Preprocess
using JuliaFEM.Postprocess
using JuliaFEM.Abaqus
using JuliaFEM.Testing

@testset "parse abaqus inp file to AbaqusModel" begin
    fn = Pkg.dir("JuliaFEM") * "/test/testdata/cube_tet4.inp"
    model = abaqus_read_model(fn)

    @test length(model.properties) == 1
    section = first(model.properties)
    @test section.element_set == :CUBE
    @test section.material_name == :MAT

    @test haskey(model.materials, :MAT)
    material = model.materials[:MAT]
    @test isapprox(first(material.properties).E, 208.0e3)

    @test length(model.steps) == 1
    step = first(model.steps)
    @test length(step.boundary_conditions) == 2

    bc = step.boundary_conditions[1]
    @test bc.data[1] == [:SYM12, 3]
    @test bc.data[2] == [:SYM23, 1]
    @test bc.data[3] == [:SYM13, 2]

    load = step.boundary_conditions[2]
    @test load.data[1] == [:LOAD, :P, 1.00000]
end

#=
@testset "given abaqus model solve field" begin
    fn = Pkg.dir("JuliaFEM") * "/test/testdata/cube_tet4.inp"
    model = abaqus_read_model(fn)
    problems = model()
    body = first(problems)
    info(body("displacement", 0.0))
    result = XDMF()
    xdmf_new_result!(result, body, 0.0)
    xdmf_save_field!(result, body, 0.0, "displacement"; field_type="Vector")
    xdmf_save!(result, "/tmp/cube_tet4.xmf")
end
=#


# MIT License
# Copyright (c) 2017: Xavier Gandibleux, Anthony Przybylski, Gauthier Soleilhac, and contributors.
using Printf
varname_generic(m::Model, col::Integer) = "VAR$(col)"

function varname_given(m::Model, col::Integer)
    # TODO: deal with non-ascii characters?
    name = getname(m, col)
    for (pat, sub) in [("[", "_"), ("]", ""), (",", "_")]
        name = replace(name, pat => sub)
    end
    name
end

function writeMOP(m, fname::AbstractString, genericnames=false)
    varname = genericnames ? varname_generic : varname_given

    f = open(fname, "w")

    write(f,"NAME   vOptModel\n")

    numRows = length(m.linconstr)

    md = getvOptData(m)
    objlincoef = [prepAffObjective(m, f) for f in md.objs]

    write(f,"OBJSENSE\n")
    if md.objSenses[1] == :Min
        write(f," MIN\n")
    else
        write(f," MAX\n")
    end

    for i = 2:length(md.objs)
        if md.objSenses[i] != md.objSenses[1]
            println("Warning, MOP does not support different optimisation senses. Flipping objective coefficients.")
            objlincoef[i] = -objlincoef[i]
        end
    end

    # Objective and constraint names
    write(f,"ROWS\n")

    for i = 1:length(md.objs)
        @printf(f," N  OBJ%d\n", i)
    end

    hasrange = false
    for c in 1:numRows
        rowsense = JuMP.sense(m.linconstr[c])
        if rowsense == :(<=)
            senseChar = 'L'
        elseif rowsense == :(==)
            senseChar = 'E'
        elseif rowsense == :(>=)
            senseChar = 'G'
        else
            hasrange = true
            senseChar = 'E'
        end
        @printf(f," %c  CON%d\n",senseChar,c)
    end

    A = JuMP.prepConstrMatrix(m)
    colptr = A.colptr
    rowval = A.rowval
    nzval = A.nzval

    # Output each column
    inintegergroup = false
    write(f,"COLUMNS\n")
    for col in 1:m.numCols
        t = m.colCat[col]
        (t == :SemiCont || t == :SemiInt) && error("The MPS file writer does not currently support semicontinuous or semi-integer variables")
        if (t == :Bin || t == :Int) && !inintegergroup
            @printf(f,"    MARKER    'MARKER'                 'INTORG'\n")
            inintegergroup = true
        elseif (t == :Cont || t == :Fixed) && inintegergroup
            @printf(f,"    MARKER    'MARKER'                 'INTEND'\n")
            inintegergroup = false
        end
        for ind in colptr[col]:(colptr[col+1]-1)
            @printf(f,"    %s  CON%d  ",varname(m,col),rowval[ind])
            (Base.Grisu).print_shortest(f,nzval[ind])
            println(f)
        end
        for obj = 1:length(md.objs)
            @printf(f,"    %s  OBJ%d  ",varname(m,col), obj)
            (Base.Grisu).print_shortest(f,objlincoef[obj][col])
            println(f)
        end
    end
    if inintegergroup
        @printf(f,"    MARKER    'MARKER'                 'INTEND'\n")
    end

    # RHSs
    write(f,"RHS\n")
    for c in 1:numRows
        rowsense = JuMP.sense(m.linconstr[c])
        if rowsense != :range
            @printf(f,"    RHS    CON%d    ",c)
            (Base.Grisu).print_shortest(f,JuMP.rhs(m.linconstr[c]))
        else
            @printf(f,"    RHS    CON%d    ",c)
            (Base.Grisu).print_shortest(f,m.linconstr[c].lb)
        end
        println(f)
    end

    # RANGES
    if hasrange
        write(f,"RANGES\n")
        for c in 1:numRows
            rowsense = JuMP.sense(m.linconstr[c])
            if rowsense == :range
                @printf(f,"    rhs    CON%d    ",c)
                (Base.Grisu).print_shortest(f,m.linconstr[c].ub-m.linconstr[c].lb)
                println(f)
            end
        end
    end

    # BOUNDS
    write(f,"BOUNDS\n")
    for col in 1:m.numCols
        if m.colLower[col] == 0
            if m.colUpper[col] != Inf
                # Default lower 0, and an upper
                @printf(f,"  UP BOUND %s ", varname(m,col))
                (Base.Grisu).print_shortest(f, m.colUpper[col])
                println(f)
            else
                # Default bounds. Explicitly state for solvers like Gurobi. See issue #792
                @printf(f,"  PL BOUND %s", varname(m,col))
                println(f)
            end
        elseif m.colLower[col] == -Inf && m.colUpper[col] == +Inf
            # Free
            @printf(f, "  FR BOUND %s\n", varname(m,col))
        elseif m.colLower[col] != -Inf && m.colUpper[col] == +Inf
            # No upper, but a lower
            @printf(f, "  PL BOUND %s\n  LO BOUND %s ",varname(m,col),varname(m,col))
            (Base.Grisu).print_shortest(f,m.colLower[col])
            println(f)
        elseif m.colLower[col] == -Inf && m.colUpper[col] != +Inf
            # No lower, but a upper
            @printf(f,"  MI BOUND %s\n  UP BOUND %s ",varname(m,col),varname(m,col))
            (Base.Grisu).print_shortest(f,m.colUpper[col])
            println(f)
        else
            # Lower and upper
            @printf(f, "  LO BOUND %s ",varname(m,col))
            (Base.Grisu).print_shortest(f,m.colLower[col])
            println(f)
            @printf(f, "  UP BOUND %s ",varname(m,col))
            (Base.Grisu).print_shortest(f,m.colUpper[col])
            println(f)
        end
    end

    write(f,"ENDATA\n")
    close(f)
    nothing
end

nextline(f) = split(chomp(readline(f)), ' ', keepempty=false)
function parseMOP(fname::AbstractString; solver=JuMP.UnsetSolver())

    m = vModel(solver = solver)

    open(fname) do f
        ln = nextline(f)
        while ln[1] != "NAME"
            ln = nextline(f)
        end

        objSense = :Min
        ln = nextline(f)
        if ln[1] == "OBJSENSE"
            ln = nextline(f)
            if ln[1] == "MAXIMIZE" || ln[1] == "MAX"
                objSense = :Max
            end
            ln = nextline(f)
        end

        if ln[1] == "ROWS"
            ln = nextline(f)
            DicoExpr = Dict{AbstractString,Tuple{JuMP.AffExpr, Symbol}}()
            objOrder = Vector{AbstractString}()
            while ln[1] != "COLUMNS"
                push!(DicoExpr, ln[2] => ( zero(JuMP.AffExpr), Symbol(ln[1]) ) )
                ln[1] == "N" && push!(objOrder, ln[2])
                ln = nextline(f)
            end
        end

        if ln[1] == "COLUMNS"
            ln = nextline(f)
            DicoVar = Dict{AbstractString,JuMP.Variable}()
            isInt = false
            while ln[1] != "RHS"

                if ln[3] == "'INTORG'"
                    isInt = true
                elseif ln[3] == "'INTEND'"
                    isInt = false
                else
                    if !(ln[1] in keys(DicoVar))

                        _cat = isInt ? (:Int) : (:Cont)
                        v = @variable(m, category=_cat, basename=ln[1], lowerbound=0.0)

                        # v = @variable(m, basename=ln[1], lowerbound=0.0)
                        # if isInt
                        #     m.colCat[v.col] = :Int
                        # end

                        push!(DicoVar, ln[1] => v)
                    end

                    expr, expr_type = DicoExpr[ln[2]]
                    var = DicoVar[ln[1]]
                    coeff = parse(Float64, ln[3])

                    push!(expr, coeff, var)

                end
                ln = nextline(f)
            end
        end

        DicoCstr = Dict{AbstractString, Tuple{JuMP.LinearConstraint, Symbol}}()
        DicoObj  = Dict{AbstractString, JuMP.QuadExpr}()
        for (k,v) in DicoExpr
            expr, expr_type = v
            if expr_type == :N
                push!(DicoObj, k => convert(QuadExpr, expr))
            else
                push!(DicoCstr, k => (LinearConstraint(expr, -Inf, Inf), expr_type) )
            end
        end

        md = m.ext[:vOpt]
        for key in objOrder
            obj = DicoObj[key]
            push!(md.objs, obj)
            push!(md.objSenses, objSense)
        end

        if ln[1] == "RHS"
            ln = nextline(f)
            while length(ln) > 1

                constraint, constraint_type = DicoCstr[ln[2]]
                if constraint_type == :E
                    constraint.lb = parse(Float64, ln[3])
                    constraint.ub = parse(Float64, ln[3])
                elseif constraint_type == :G
                    constraint.lb = parse(Float64, ln[3])
                elseif constraint_type == :L
                    constraint.ub = parse(Float64, ln[3])
                end
                ln = nextline(f)
            end
        end

        for (k,v) in DicoCstr
            cstr, cstr_type = v
            JuMP.addconstraint(m, cstr)
        end

        if ln[1] == "RANGES"
            ln = nextline(f)
            while length(ln) > 1

                constraint, constraint_type = DicoCstr[ln[2]]
                r = parse(Float64, ln[3])
                if constraint_type == :G
                    constraint.ub = constraint.lb + abs(r)
                elseif constraint_type == :L
                    constraint.lb = constraint.ub - abs(r)
                else
                    if r >= 0
                        constraint.ub = constraint.lb + r
                    else
                        constraint.lb = constraint.ub + r
                    end
                end
                ln = nextline(f)
            end
        end

        fill!(m.colLower, 0.0)

        if ln[1] == "BOUNDS"
            ln = nextline(f)
            while length(ln) > 1

                boundType, varName = ln[1], ln[3]
                var = DicoVar[varName]
                if boundType == "LO"
                    val = parse(Float64, ln[4])
                    JuMP.setlowerbound(var, val)
                elseif boundType == "UP"
                    val = parse(Float64, ln[4])
                    JuMP.setupperbound(var, val)
                elseif boundType == "FX"
                    val = parse(Float64, ln[4])
                    JuMP.setupperbound(var, val)
                    JuMP.setlowerbound(var, val)
                elseif boundType == "FR"
                    JuMP.setupperbound(var, Inf)
                    JuMP.setlowerbound(var, -Inf)
                elseif boundType == "MI"
                    JuMP.setlowerbound(var, -Inf)
                elseif boundType == "PL"
                    JuMP.setupperbound(var, Inf)
                elseif boundType == "BV"
                    setcategory(var, :Bin)
                elseif boundType == "LI"
                    val = parse(Float64, ln[4])
                    setlowerbound(var, val)
                elseif boundType == "UI"
                    val = parse(Float64, ln[4])
                    setupperbound(var, val)
                elseif boundType == "SC"
                    error("Semi-continuous variables are not supported by vOptGeneric")
                end
                ln = nextline(f)
            end
        end
    end

    return m
end


module TestPhysChem

using SoilTracers.PhysChem
using Base.Test

bindings = [
    :c2k,
    :k2c,
    :diffus_air,
    :diffus_water,
    :diffus_soil_air,
    :diffus_soil_water,
    :diffus_soil,
    :solub_co2,
    :solub_cos,
    :water_density,
    :water_dissoc,
    :hydrolysis_cos]

private_bindings = [
    :diffus_air_stp,
    :diffus_water_stp,
    :soil_shape_params]

@time @testset "PhysChem module bindings" begin
    for binding in bindings
        @test isdefined(binding)
    end

    # private bindings should not be exported
    for private_binding in private_bindings
        @test !isdefined(private_binding)
    end
end

end


export AbstractPreconditioner
export AbstractPreconditionerCache
export AbstractPreconditionerType
export PCinitialize!, PCapply!
export PCside, PCright, PCleft

"""
    AbstractPreconditioner

This is an abstract type representing an abstract preconditioner.

The available concrete implementations are:

- [`Solvent.Identity`](@ref)
- [`Solvent.Jacobi`](@ref)
"""
abstract type AbstractPreconditioner end

abstract type AbstractPreconditionerCache end
abstract type AbstractPreconditionerType end

abstract type PCside end
struct PCright <: PCside end
struct PCleft <: PCside end

"""
    function Preconditioner(
        pc_type::AbstractPreconditionerType,
        linearoperator!,
        Q,
    )

Construct preconditioner for `linearoperator!`.
"""
mutable struct Preconditioner{
    pcType <: AbstractPreconditionerType,
    OP,
    pccType <: AbstractPreconditionerCache,
    stype <: PCside,
} <: AbstractPreconditioner
    pc_type::pcType
    linop!::OP
    cache::pccType
    pc_side::stype

    function Preconditioner(
        pc_type::AbstractPreconditionerType,
        linearoperator!,
        Q,
    )
        pccache = cache(pc_type, linearoperator!, Q)
        OP = typeof(linearoperator!)
        pcType = typeof(pc_type)
        pccType = typeof(pccache)
        stype = typeof(pc_type.pc_side)

        return new{pcType, OP, pccType, stype}(
            pc_type,
            linearoperator!,
            pccache,
            pc_type.pc_side,
        )
    end
end

"""
    PCinitialize!(
        pc::AbstractPreconditioner,
        Q,
        Qrhs,
        args...,
    )

Perform precomputations for `pc`.
"""
function PCinitialize!(
    pc::AbstractPreconditioner,
    Q,
    Qrhs,
    args...,
)
    PCinitialize!(pc.pc_type, pc, Q, Qrhs, args...)
end

"""
    PCapply!(
        pc::AbstractPreconditioner,
        W,
        Q,
        args...,
    )

Apply preconditioner `pc` to `Q` and store the result in `W`.
"""
function PCapply!(
    pc::AbstractPreconditioner,
    W,
    Q,
    args...,
)
    PCapply!(pc.pc_type, pc, W, Q, args...)
end

include("identity.jl")
include("jacobi.jl")


function monotone_forecast_variance(forec::StateSpaceModels.Forecast)
    for i in 2:length(forec.covariance)
        if all(diag(forec.covariance[i]) .< diag(forec.covariance[i - 1]))
            return false
        end
    end
    return true
end

function expected_value_of_scenarios(scenarios::Array{T,3}) where T
    expected_values = Vector{Vector{T}}(undef, size(scenarios, 1))
    for t in 1:length(expected_values)
        expected_values[t] = vec(mean(scenarios[t, :, :]; dims=2))
    end
    return expected_values
end

function test_scenarios_adequacy_with_forecast(
    forec::StateSpaceModels.Forecast, scenarios::Array{T,3}; rtol::T=5e-3, atol::T=NaN
) where T
    # Test expected values
    if isnan(atol)
        for t in 1:length(forec.expected_value)
            @test forec.expected_value[t] ≈ expected_value_of_scenarios(scenarios)[t] rtol =
                rtol
        end
    else
        for t in 1:length(forec.expected_value)
            @test forec.expected_value[t] ≈ expected_value_of_scenarios(scenarios)[t] atol =
                atol
        end
    end
    return nothing
end


#=
Functions related to exchanging ghost data.
=#

"""
    ExchangeGroup

A description of information that needs to be communicated between processes during
a ghost exchange. It includes the ID of the involved processes, the included elements,
and the variable and component indices. The ghost exchange function will use this
to set up the data and path to communicate.
"""
mutable struct ExchangeGroup
    members::Array{Int,1} # ranks of the members of the group

    ghost_elements_only::Bool   # True if the exchange is for ghost element pairs
    exchange_all_data::Bool     # True if all variables/components are exchanged

    elements::Array{Int,2}  # 1xN list of global indices of elements to share. 2xN Pairs if ghost_elements_only.
    variables::Array{Int,1} # indices of variables to share (empty if exchange_all_data=true)
    components::Array{Array{Int,1},1} # components of each variable to share (empty if exchange_all_data=true)
end

"""

"""
function build_exchange_groups(grid; members=[], variables=[], components=[], ghost_elements_only=true, exchange_all_data=true)
    # This list of exchange groups will be returned
    ex_groups = [];

    if ghost_elements_only
        if exchange_all_data
            # If ghost_elements_only and exchange_all_data, this matches the ghost info from the grid
            ngroups = grid.num_neighbor_partitions;
            for gi=1:ngroups
                neighbor = grid.neighboring_partitions[gi];

                push!(ex_groups, ExchangeGroup([neighbor], ghost_elements_only, exchange_all_data, grid.ghost_index[gi], [], []));
            end
        else
            # If exchange_all_data = false, this matches the ghost info from the grid, but specifies some components
            ngroups = grid.num_neighbor_partitions;
            for gi=1:ngroups
                neighbor = grid.neighboring_partitions[gi];

                push!(ex_groups, ExchangeGroup([neighbor], ghost_elements_only, exchange_all_data, grid.ghost_index[gi], variables, components));
            end
        end
    else # all local elements are shared with ranks that have same partition
        if exchange_all_data

        else

        end
    end

    return ex_groups;
end

#=
Exchange ghosts.
This does both send and receive.
Values are sent/received directly from the variable.values arrays

NOTE: This is currently only set up for the case where each process has a distinct partition.
For other strategies, this needs to be modified.
=#
function exchange_ghosts(var, grid, tag=0)
    tag = config.num_procs + Int(round(tag));
    if config.solver_type == FV
        return exchange_ghosts_fv(var, grid, tag);
    else
        printerr("Ghost exchange not ready for "*config.solver_type, fatal=true)
    end
end

function exchange_ghosts_fv(var, grid, tag)
    n_neighbors = grid.num_neighbor_partitions;
    # How big does each buffer need to be?
    buffer_size = zeros(Int, n_neighbors);
    send_arrays = fill(zeros(0), n_neighbors);
    recv_arrays = fill(zeros(0), n_neighbors);
    requests = Vector{MPI.Request}(undef, n_neighbors * 2);
    if typeof(var) <: Array
        # multiple variables being solved for simultaneously
        dofs_per_node = 0;
        for vi=1:length(var)
            dofs_per_node += var[vi].total_components;
        end
    else
        # one variable
        dofs_per_node = var.total_components;
    end
    for ni=1:n_neighbors
        buffer_size[ni] = dofs_per_node * grid.ghost_counts[ni]; # number of values to send to neighbor ni
        send_arrays[ni] = zeros(buffer_size[ni]);
        recv_arrays[ni] = zeros(buffer_size[ni]);

        # Fill the send buffer
        next_ind = 1;
        if typeof(var) <: Array
            for vi=1:length(var)
                for ei=1:grid.ghost_counts[ni]
                    for comp_i=1:var[vi].total_components
                        send_arrays[ni][next_ind] = var[vi].values[comp_i,grid.ghost_index[ni][2,ei]];
                        next_ind += 1;
                    end
                end
            end
        else
            for ei=1:grid.ghost_counts[ni]
                for comp_i=1:var.total_components
                    send_arrays[ni][next_ind] = var.values[comp_i,grid.ghost_index[ni][2,ei]];
                    next_ind += 1;
                end
            end
        end
    end

    # send ghosts
    if config.use_mpi
        for ni=1:grid.num_neighbor_partitions
            requests[ni*2-1] = MPI.Isend(send_arrays[ni], grid.neighboring_partitions[ni], tag, MPI.COMM_WORLD);
        end
    end
    # receive ghosts
    if config.use_mpi
        for ni=1:grid.num_neighbor_partitions
            requests[ni*2] = MPI.Irecv!(recv_arrays[ni], grid.neighboring_partitions[ni], tag, MPI.COMM_WORLD);
        end

        # wait for it
        stats = MPI.Waitall!(requests);
    end

    # Put the received ghost values in place
    for ni=1:n_neighbors
        next_ind = 1;
        if typeof(var) <: Array
            for vi=1:length(var)
                for ei=1:grid.ghost_counts[ni]
                    for comp_i=1:var[vi].total_components
                        var[vi].values[comp_i,grid.ghost_index[ni][1,ei]] = recv_arrays[ni][next_ind];
                        next_ind += 1;
                    end
                end
            end
        else
            for ei=1:grid.ghost_counts[ni]
                for comp_i=1:var.total_components
                    var.values[comp_i,grid.ghost_index[ni][1,ei]] = recv_arrays[ni][next_ind];
                    next_ind += 1;
                end
            end
        end
    end
end


# Define high-level structs and their zero initializations

"""
Struct for modelling parameters specific to each supersource.
Each worker needs a vector of this struct, as it models multiple superspources.
"""
mutable struct P_x_worker_x_pw_x_ss{N}
    iss::Int64
    wavelets::Vector{NamedStack{Vector{Float64}}} # [ .. for it in 1:nt]
    ssprayw::NamedStack{Vector{Vector{Float64}}}
    records::NamedStack{Matrix{Float64}}
    rinterpolatew::NamedStack{Vector{Vector{Float64}}}
    sindices::NamedStack{Vector{T1}} where T1<:CartesianIndices{N} # contains indices for each source, each named field
    rindices::NamedStack{Vector{T2}} where T2<:CartesianIndices{N}# contains indices for each receiver, each named field
    boundary::Vector{Array{Float64,3}}
    snaps::NamedVector{
        Array{Data.Number,N},
        Vector{Array{Data.Number,N}},
        Tuple{OrderedDict{String,Int64}},
    }
    illum::Matrix{Float64}
    grad_mod::NamedStack{Matrix{Float64}} # w.r.t different coeffs
end

function iszero_boundary(pa)
    result = false
    @sync begin
        for (ip, p) in enumerate(procs(pa.p))
            @async remotecall_wait(p) do
                pap = localpart(pa.p)
                for issp = 1:length(pap.ss)
                    pass = pap.ss[issp]
                    for i = 1:length(pass.boundary)
                        result = result | iszero(pass.boundary[i])
                    end
                end
            end
        end
    end
    return result
end
function initialize_boundary!(pa)
    # zero out results stored per worker
    for ipw = 1:pa.c.ic[:npw]
        @sync begin
            for (ip, p) in enumerate(procs(pa.p))
                @async remotecall_wait(p) do
                    pap = localpart(pa.p)[ipw]
                    for issp = 1:length(pap.ss)
                        pass = pap.ss[issp]
                        for i = 1:length(pass.boundary)
                            fill!(pass.boundary[i], 0.0)
                        end
                    end
                end
            end
        end
    end
end

"""
Initialize between each simulation
"""
function initialize!(pa::P_x_worker_x_pw_x_ss)
    fill!.(pa.records, 0.0)
    fill!.(pa.snaps, 0.0)
    fill!(pa.illum, 0.0)
    fill!.(pa.grad_mod, 0.0)
end

"""
Parameters specific to each worker, not necessarily for every supersource.
This struct contains velocity and stress fields simulated by each worker.
Again, note that a single worker can take care of multiple supersources.
* T1==Matrix{Float64} for 2D simulation
"""
mutable struct P_x_worker_x_pw{N,Q<:Data.Array{N}}
    ss::Vector{P_x_worker_x_pw_x_ss{N}} # supersource specific arrays
    # 2D arrays of p, vx, vz, their previous snapshots, and
    # x derivatives of p, vx, vz
    # z derivatives of p, vx, vz
    w1::NamedStack{NamedStack{Q}} # p, vx, vz on GPU or CPU
    wr::NamedStack{Array{Data.Number,N}} # same as above but only for receiver fields when GPU to facilitate faster scalar indexing!?
    w2::NamedStack{NamedStack{Q}} # required for attenuation, where third dimension is nsls (only used for 2D simulation right now)
    memory_pml::NamedStack{Q} # memory variables for CPML
    born_svalue_stack::Array{Float64,N} # used for born modeling # only used for 2-D simulation
end

# P_x_worker=Vector{P_x_worker_x_pw}

function initialize!(pap::Vector{P_x_worker_x_pw{N,B}}) where {N,B}
    reset_w2!(pap)
    for pap_x_pw in pap
        initialize!.(pap_x_pw.ss)
    end

end

# reset wavefields for every worker
function reset_w2!(pap::Vector{P_x_worker_x_pw{N,B}}) where {N,B}
    for pap_x_pw in pap
        fill!(pap_x_pw.born_svalue_stack, 0.0)
        for w in pap_x_pw.w1
            fill!.(w, 0.0)
        end
        for w in pap_x_pw.w2
            fill!.(w, 0.0)
        end
        fill!.(pap_x_pw.memory_pml, 0.0)
    end
end

"""
Modelling parameters common for all supersources
# Keyword Arguments that are modified by the method (some of them are returned as well)

* `gradient::Vector{Float64}=Medium(medium.mgrid)` : gradient medium modified only if `gmodel_flag`
* `TDout::Vector{Records.TD}=[Records.TD_zeros(rfields,tgridmod,ageom[ip]) for ip in 1:length(findn(rflags))]`
* `illum::Array{Float64,2}=zeros(length(medium.mgrid[1]), length(medium.mgrid[2]))` : source energy if `illum_flag`
* `boundary::Array{Array{Float64,4},1}` : stored boundary values for first propagating wavefield
* `snaps::NamedArray{Array{Float64}}` :snapshots saved at `tsnaps`

# Return (in order)

* modelled data for each propagating wavefield as `Vector{TD}`
* stored boundary values of the first propagating wavefield as `Array{Array{Float64,4},1}` (use for backpropagation)
* final conditions of the first propagating wavefield as `Array{Float64,4}` (use for back propagation)
* gradient model as `Medium`
* stored snaps shots at tsnaps as Array{Float64,4}

"""
mutable struct P_common{T,N,Q1<:Data.Array{1},Q2<:Data.Array{N}}
    jobname::Symbol
    attrib_mod::T
    activepw::Vector{Int64}
    exmedium::Medium
    medium::Medium
    ageom::Vector{AGeom}
    srcwav::Vector{SrcWav}
    pml_edges::Vector{Symbol}
    sfields::Vector{Vector{Symbol}}
    sflags::Vector{Int64}
    rfields::Vector{Symbol}
    rflags::Vector{Int64}
    fc::NamedStack{Float64}
    ic::NamedStack{Int64}
    pml::NamedStack{NamedStack{Q1}} # e.g., pml[:x][:a], pml[:z][:b]
    mod::NamedStack{Q2} # e.g., mod[:KI], mod[:K]
    mod3::NamedStack{Array{Float64,3}} # (only used for 2-D attenuation modelling, so fixed)
    #=
    	modKI::Matrix{Float64}
    	modK::Matrix{Float64} # just storing inv(modKI) for speed
    	modrr::Matrix{Float64}
    	modrhovxI::Matrix{Float64}
    	modrhovzI::Matrix{Float64}
    	=#
    δmod::NamedStack{Array{Float64,N}}
    δmodall::Vector{Float64}
    #=
    	δmodKI::Matrix{Float64}
    	δmodrr::Matrix{Float64}
    	δmodrhovxI::Matrix{Float64}
    	δmodrhovzI::Matrix{Float64}
    	δmod::Vector{Float64} # perturbation vector (KI, rhoI)
    	=#
    gradient::Vector{Float64}  # output gradient vector w.r.t (KI, rhoI)
    grad_mod::NamedStack{SharedArrays.SharedArray{Float64,N}}
    #=
    	grad_modKI_stack::SharedArrays.SharedArray{Float64,2} # contains gmodKI
    	grad_modrhovxI_stack::SharedArrays.SharedArray{Float64,2}
    	grad_modrhovzI_stack::SharedArrays.SharedArray{Float64,2}
    	grad_modrr_stack::Array{Float64,2}
    	=#
    illum_flag::Bool
    illum_stack::SharedArrays.SharedArray{Float64,N}
    backprop_flag::Int64
    snaps_field::Symbol
    itsnaps::NamedVector{
        Int64,
        Vector{Int64},
        Tuple{OrderedCollections.OrderedDict{String,Int64}},
    }
    gmodel_flag::Bool
    bindices::NamedStack{Int64}
    #=
    	ibx0::Int64
    	ibz0::Int64
    	ibx1::Int64
    	ibz1::Int64
    	isx0::Int64
    	isz0::Int64
    	=#
    datamat::SharedArrays.SharedArray{Float64,3}
    data::Vector{Records}
    # attenuation related parameters
    verbose::Bool
end

function initialize!(pac::P_common)
    fill!(pac.gradient, 0.0)
    fill!.(pac.grad_mod, 0.0)
    # fill!(pac.grad_mod[:rhovxI],0.0)
    # fill!(pac.grad_mod[:rhovzI],0.0)
    # fill!(pac.grad_mod[:rhoI],0.0)
    fill!(pac.illum_stack, 0.0)
    for dat in pac.data
        fill!(dat, 0.0)
    end
    fill!(pac.datamat, 0.0)
end

"""
A single struct combining parameters for all workers.
"""
mutable struct PFdtd{T,N,B}
    sschunks::Vector{UnitRange{Int64}} # how supersources are distributed among workers
    p::DistributedArrays.DArray{Vector{P_x_worker_x_pw{N,B}},1,Vector{P_x_worker_x_pw{N,B}}} # distributed parameters among workers
    c::P_common{T,N} # common parameters
end


"""
Abstract supertype for organ parameters.
Extend to add additional components to organ parameters.
"""
abstract type AbstractParams end

"""
Model parameters that vary between organs
"""
@default_kw @flattenable @selectable struct Params{As,Sh,Al,Ma,Tr,Re,Ge,Pr} <: AbstractParams
    # Field               | Default                    | _     | Selectable Types
    assimilation_pars::As | ConstantCAssim()           | _     | Union{Nothing,AbstractAssim}
    scaling_pars::Sh      | nothing                    | _     | Union{Nothing,AbstractScaling}
    allometry_pars::Al    | nothing                    | _     | Union{Nothing,AbstractAllometry}
    maturity_pars::Ma     | nothing                    | _     | Union{Nothing,AbstractMaturity}
    activetrans_pars::Tr  | nothing                    | _     | Union{Nothing,ActiveTranslocation}
    passivetrans_pars::Re | LosslessPassiveTranslocation() | _ | PassiveTranslocation
    germination_pars::Ge  | nothing                    | _     | Union{Nothing,AbstractGermination}
    production_pars::Pr   | nothing                    | _     | Union{Nothing,AbstractProduction}
end

for fn in fieldnames(Params)
    @eval $fn(p::Params) = p.$fn
end

"""
Astract supertype for shared parameters.
Extend to change the components that are shared.
"""
abstract type AbstractSharedParams end

"""
    SharedParams(su_pars, core_pars, resorption_pars, tempcorr_pars, catabolism_pars)

Model parameters shared between organs.

# FieldMetadata macros
- `@default_kw` provides default values and constructors,
- `@selectable` provides the set of types that can be used for the parameter,
  so that they can be selected in a live interface.
"""
@default_kw @selectable struct SharedParams{SU,Co,Fe,Te,Ca} <: AbstractSharedParams
    # Field               | Default                   | Selectable Types
    su_pars::SU           | ParallelComplementarySU() | AbstractSynthesizingUnit
    core_pars::Co         | DEBCore()                 | _
    resorption_pars::Fe   | nothing                   | Union{Nothing,AbstractResorption}
    tempcorr_pars::Te     | nothing                   | Union{Nothing,AbstractTemperatureCorrection}
    catabolism_pars::Ca   | CatabolismCN()            | AbstractCatabolism
end

for fn in fieldnames(SharedParams)
    @eval $fn(p::SharedParams) = p.$fn
end

###########################################################################################
# Variables

"""
Model variables.
Allow storing and accessing variables for use by multiple components.
"""
abstract type AbstractVars end

"""
    PlottableVars()

Plottable model variables. These are vectors witih values for each time-step,
to allow plotting and model introspection.
"""
@udefault_kw @units @plottable struct PlottableVars{F,R,E,T,WP,H,TS} <: AbstractVars
    scaling::F        | [0.0] | _                 | _
    rate::R           | [0.0] | mol*mol^-1*d^-1   | _
    E_ctb::E          | [0.0] | mol/hr            | _
    θE::F             | [0.0] | _                 | _
    temp::T           | [0.0] | K                 | _
    tempcorrection::F | [0.0] | _                 | _
    swp::WP           | [0.0] | kPa               | _
    soilcorrection::F | [0.0] | _                 | _
    height::H         | [0.0] | m                 | _
    tstep::TS         | [1]   | _                 | false
end

"""
    Vars()

Mutable struct to allow storing variables
for use by multiple components.
"""
@udefault_kw @units mutable struct Vars{F,R,E,T,WP,H} <: AbstractVars
    scaling::F        | 0.0   | _
    rate::R           | 0.0   | mol*mol^-1*d^-1
    θE::F             | 0.0   | _
    E_ctb::E          | 0.0   | mol/hr
    temp::T           | 0.0   | K
    tempcorrection::F | 0.0   | _
    swp::WP           | 0.0   | kPa
    soilcorrection::F | 0.0   | _
    height::H         | 0.0   | m
end

tstep(v::PlottableVars) = v.tstep[1]
tstep(v::Vars) = nothing

set_tstep!(v::PlottableVars, val) = v.tstep[1] = val
set_tstep!(v::Vars, val) = nothing

depth(v) = height(v)

build_vars(vars::Vars, tspan) = vars
build_vars(vars::PlottableVars, tspan) = begin
    len = length(tspan)
    len <= length(vars.rate) && return vars

    for fname in fieldnames(typeof(vars))
        varvec = getfield(vars, fname)
        append!(varvec, fill(getfield(vars, fname)[1], len - length(varvec)))
    end
    vars
end

###########################################################################################
# Organs and Organisms

"""
Abstract supertype for organs. Inherit from it if you need to difine
behaviour different to that or [`Organ`](@ref).
"""
abstract type AbstractOrgan{P,S} end

params(o::AbstractOrgan) = o.params
shared(o::AbstractOrgan) = o.shared
vars(o::AbstractOrgan) = o.vars
flux(o::AbstractOrgan) = o.J

#= κ functions allow modular addition of destinations for catabolised flux,
by default maturity and active translocation. If those components are
not used the flux fraction will be allocated to κsoma. =#

κtra(o::AbstractOrgan) = κtra(activetrans_pars(o))
κtra(o::Nothing) = 0.0

κmat(o::AbstractOrgan) = κmat(maturity_pars(o))
κmat(::Nothing) = 0.0

κsoma(o::AbstractOrgan) = oneunit(κtra(o)) - κtra(o) - κmat(o)

# Define `scaling` and `setscaling` etc. methods
for fn in fieldnames(Vars)
    setfn = Symbol.(:set_, fn, :!)
    @eval $fn(o::AbstractOrgan) = $fn(vars(o))
    @eval $setfn(o::AbstractOrgan, x) = $setfn(vars(o), x)
    @eval @inline ($fn)(vars::PlottableVars) = vars.$fn[tstep(vars)]
    @eval @inline ($setfn)(vars::PlottableVars, val) = vars.$fn[tstep(vars)] = val
    @eval @inline ($fn)(vars::Vars) = vars.$fn
    @eval @inline ($setfn)(vars::Vars, val) = vars.$fn = val
end

#= Forward variable and parameter getter/setter methods
so they can be acessesed/set directly from the organ,
without knowning where they are actually stored. =#
tstep(o::AbstractOrgan) = tstep(vars(o))
set_tstep!(o::AbstractOrgan, t) = set_tstep!(vars(o), t)

for fn in fieldnames(DEBCore)
    @eval $fn(p::SharedParams) = $fn(core_pars(p))
    @eval $fn(o::AbstractOrgan{<:Any,<:SharedParams}) = $fn(shared(o))
end

for fn in fieldnames(SharedParams)
    @eval $fn(o::AbstractOrgan{<:Any,<:SharedParams}) = $fn(shared(o))
end

for fn in fieldnames(Params)
    @eval $fn(o::AbstractOrgan{<:Params}) = $fn(params(o))
end

"""
    Organ(params, shared, vars, J)

Basic model components. For a plants, organs might be roots, stem and leaves
"""
struct Organ{P,S,V,F} <: AbstractOrgan{P,S}
    params::P
    shared::S
    vars::V
    J::F
end
"""
    Organ(params, shared, records)

Construct an organ from parameters, shared parameters and
views into records arrays for vaiable and flux matrices.
"""
Organ(params::AbstractParams, shared::AbstractSharedParams, records) = begin
    vars = records.vars
    set_tstep!(vars, 1)
    J = view(records.J, Ti(1))
    Organ(params, shared, vars, J)
end

# Records hold vars and flux, allowing storage of each
# timestep for plotting if PlottableVars are used.
abstract type AbstractRecords end

@plottable struct PlottableRecords{V,F} <: AbstractRecords
    vars::V | true
    J::F    | false
end
PlottableRecords(vars::PlottableVars, J::AbstractArray) =
    PlottableRecords{map(typeof,(vars,J))...}(vars, J)
PlottableRecords(vars::PlottableVars, fluxval::Number, fluxaxes::Tuple, tspan::AbstractRange) = begin
    vars = build_vars(vars, tspan)
    J = build_flux(fluxval, fluxaxes..., tspan)
    PlottableRecords(vars, J)
end

@plottable struct Records{V,F} <: AbstractRecords
    vars::V | false
    J::F    | false
end
Records(vars::Vars, J::AbstractArray) =
    Records{map(typeof,(vars,J))...}(vars, J)
Records(vars::Vars, fluxval::Number, fluxaxes::Tuple) = begin
    J = build_flux(fluxval, fluxaxes...)
    Records{map(typeof, (vars,J))...}(vars, J)
end

build_flux(fluxval, x::Tuple, y::Tuple) = begin
    dims = X(Val(x)), Y(Val(y))
    A = zeros(typeof(fluxval), map(length, dims)...)
    DimensionalArray(A, dims)
end
build_flux(fluxval, x::Tuple, y::Tuple, time::AbstractRange) = begin
    dims = X(Val(x)), Y(Val(y)), Ti(time)
    A = zeros(typeof(fluxval), map(length, dims)...)
    DimensionalArray(A, dims)
end

"""
A an Organism is a model object for modelling the growth
of a single organism.

It can be run as a functor:

```julia
(o::AbstactOrganism)(du, u, p, t) =
...
```

So that it can be passed into DifferentialEquations.jl solvers
if required.

Where `du` is the flux to be updated, `u` is the current state,
`p` are new model paremeters or `nothing`, and `t` is the current
time step. See model.jl for implemntations.

When julia 1.5 is released this will be implemented for any
`AbstactOrganism`, but for now it is only implemented for `Plant`
due to current limitations in Julia.
"""
abstract type AbstractOrganism end

params(o::AbstractOrganism) = o.params
shared(o::AbstractOrganism) = o.shared
records(o::AbstractOrganism) = o.records
organs(o::AbstractOrganism) = o.organs
environment(o::AbstractOrganism) = o.environment
environment_start(o::AbstractOrganism) = o.environment_start
dead(o::AbstractOrganism) = o.dead[]
set_dead!(o::AbstractOrganism, val) = o.dead[] = val

"""
    define_organs(o::AbstractOrganism, t)

Organs are constructed with views of Records and J Arrays at time t
"""
define_organs(o::AbstractOrganism, t) =
    define_organs(params(o), shared(o), records(o), t)
define_organs(params::Tuple, shared, records::Tuple, t) =
    map((p, r) -> Organ(p, shared, r), params, records)

update_organs(organs, t) = update_organs(vars(first(organs)), organs, t)
update_organs(::Vars, organs, t) = organs
update_organs(::PlottableVars, organs, t) = begin
    i = floor(Int, ustrip(t))
    map(organs) do organ
        organ.vars.tstep[1] = i
        J = organ.J
        data = J.data
        @set! data.indices[3] = i
        @set! data.offset1 = (i - 1) * length(data.indices[1]) * length(data.indices[2])
        # Setfield.jl isn't type-stable, DD `rebuild` is better
        @set! organ.J = DimensionalData.rebuild(J, data)
        organ
    end
end

"""
    Plant(params, shared, records, environment, environment_start, dead)
    Plant(states=(:V, :C, :N),
          transformations=(:asi, :gro, :mai, :rej, :res),
          params=(ShootParamsCN(), RootParamsCN()),
          vars=(Vars(), Vars()),
          shared=SharedParams(),
          records=nothing,
          environment=nothing,
          time=0.0hr:1.0hr:8760.0hr,
          environment_start=Ref(1.0hr),
          dead=Ref(false))

Plant model.

`params` are a tuple of Parameters objects, one for each organ.

`shared` is a struct of parameters shared between organs.

`vars`: can be `Vars` or `PlottableVars` or custom struct with additional variables.
  `PlottableVars` will be stored for each timestep for plotting.

`environment` can be `nothing`, or a `MicroclimPoint` or `MicroclimControl`
from Microclimates.jl

`time` determines the timespan over which `PlottableVars` will be constructed.
it isn't used with regular `Vars`.

If 4 state model is used, pass in
```julia
states=(:V, :E, :C, :N),
```

Similarly, if additional components like Maturity or ActiveTranslocation
are used, their state label needs to be passed in:

```julia
transformations=(:asi, :gro, :mai, :mat, :rej, :tra, :res),
```
"""
@flattenable @description mutable struct Plant{P,S,R,O,E,ES,D} <: AbstractOrganism
    params::P             | true  | "Model parameters"
    shared::S             | true  | "Parameters shared between organs"
    records::R            | false | "Plotable variables stored in arrays and sliced for each timestep on demand"
    organs::O             | false | ""
    environment::E        | false | "Environment object, provides environmental variables for each timestep"
    environment_start::ES | false | "Start index of environmental data"
    dead::D               | false | "`Bool` flag: has the plant died"
    Plant(params::P, shared::S, records::R, organs::O, environment::E, environment_start::ES, dead::D) where {P,S,R,O,E,ES,D} = begin
        organs = define_organs(params, shared, records, 0)
        new{P,S,R,typeof(organs),E,ES,D}(params, shared, records, organs, environment, environment_start, dead)
    end
end
Plant(; states=(:V, :C, :N),
        transformations=(:asi, :gro, :mai, :rej, :tra, :res),
        params=(
            Params(assimilation_pars=KooijmanSLAPhotosynthesis()),
            Params(assimilation_pars=NAssim()),
        ),
        vars=(Vars(), Vars()),
        shared=SharedParams(),
        environment=nothing,
        time=0.0hr:1.0hr:8760.0hr,
        records=nothing,
        environment_start=Ref(0.0hr),
        dead=Ref(false)
      ) = begin
    fluxaxes = states, transformations
    fluxval = 1.0mol/hr
    records = build_records(records, vars, fluxval, fluxaxes, time)
    Plant(params, shared, records, nothing, environment, environment_start, dead)
end

build_records(records::AbstractRecords, args...) = records
build_records(records::Nothing, vars::Tuple, fluxval, fluxaxes, tspan) =
    map(vars) do v
        build_records(v, fluxval, fluxaxes, tspan)
    end
build_records(vars::Vars, fluxval::Number, fluxaxes::Tuple, tspan::AbstractRange) =
    Records(vars, fluxval, fluxaxes)
build_records(vars::PlottableVars, fluxval::Number, fluxaxes::Tuple, tspan::AbstractRange) =
    PlottableRecords(vars, fluxval, fluxaxes, tspan)

# Define a SubArray constructor so we can increment the time index
ConstructionBase.constructorof(::Type{A}) where A<:SubArray{T,N,P,I,L} where {T,N,P,I,L} =
    (parent, indices, offset1, stride1) -> begin
        SubArray{eltype(parent),ndims(parent),typeof(parent),typeof(indices),L}(parent, indices, offset1, stride1)
    end


using DataFrames
using Distributions
using Random

function simple_run(individuals_df::DataFrame;
        tracking_prob::Float64=0.0,
        constant_kernel_param::Float64=1.0,
        tracking_delay::Float64=1.0,
        mild_detection_prob=0.0,
        history::Union{Nothing, Vector{Simulation.Event}}=nothing,
        execution_history::Union{Nothing, BitVector}=nothing,
        state_history::Union{Nothing, Vector{Simulation.IndividualState}}=nothing,
        debug_return=false,
        seed=123,
        initial_infections::Integer=100
    )
    rng = MersenneTwister(seed)
    params = Simulation.load_params(
        rng,
        population=individuals_df,

        constant_kernel_param=constant_kernel_param,
        household_kernel_param=1.0,

        backward_tracking_prob=tracking_prob,
        backward_detection_delay=tracking_delay/2,

        forward_tracking_prob=tracking_prob,
        forward_detection_delay=tracking_delay/2,

        testing_time=tracking_delay/2,

        mild_detection_prob = mild_detection_prob
    );
    state = Simulation.SimState(params.progressions |> length, seed=seed)

    sample(rng, 1:length(params.progressions), initial_infections) .|> person_id ->
        push!(state.queue, Simulation.Event(Val(Simulation.OutsideInfectionEvent), 0.0, person_id))

    @time Simulation.simulate!(
        state,
        params,
        history=history,
        execution_history=execution_history,
        state_history=state_history
    )

    return state
end

function simple_run!(state::Simulation.SimState, individuals_df::DataFrame;
        tracking_prob::Float64=0.0,
        constant_kernel_param::Float64=1.0,
        household_kernel_param::Real=0.3,

        tracking_delay::Float64=2.5,
        forward_detection_delay::Union{Nothing,Float64}=nothing,
        backward_detection_delay::Union{Nothing, Float64}=nothing,
        testing_time::Union{Nothing, Float64}=nothing,

        mild_detection_prob=0.0,
        phone_tracking_usage=0.0,
        phone_detection_delay=0.25,
        hospital_kernel_param::Float64=0.0,
        history::Union{Nothing, Vector{Simulation.Event}}=nothing,
        execution_history::Union{Nothing, BitVector}=nothing,
        state_history::Union{Nothing, Vector{Simulation.IndividualState}}=nothing,
        debug_return=false,
        seed=123,
        initial_infections::Integer=100
    )
    Simulation.reset!(state)

    state.rng = MersenneTwister(seed)

    params = Simulation.load_params(
        state.rng,
        population=individuals_df,

        mild_detection_prob = mild_detection_prob,

        constant_kernel_param=constant_kernel_param,
        household_kernel_param=household_kernel_param,
        hospital_kernel_param=hospital_kernel_param,

        backward_tracking_prob=tracking_prob,
        backward_detection_delay= nothing==backward_detection_delay ? tracking_delay/2 : backward_detection_delay,

        forward_tracking_prob=tracking_prob,
        forward_detection_delay= nothing==forward_detection_delay ? tracking_delay/2 : forward_detection_delay,

        testing_time= nothing==testing_time ? tracking_delay/2 : testing_time,
        phone_tracking_usage=phone_tracking_usage,
        phone_detection_delay=phone_detection_delay
    );

    Simulation.initialfeed!(state, initial_infections)

    Simulation.simulate!(
        state,
        params,
        history=history,
        execution_history=execution_history,
        state_history=state_history
    )
    return state
end

function merge_history(history::Vector{Simulation.Event}, execution_history::BitVector, state_history::Vector{Simulation.IndividualState}; pop_last_event=false)
    history = copy(history)
    if pop_last_event
        last_event = pop!(history)
    end
    zip(history, state_history, execution_history)
end

function state2plot(state::Simulation.SimState)
    transmissions = vcat(state.forest.infections...)
    sort!(transmissions, lt=(x,y)->x.time<y.time)
    times = getproperty.(transmissions, :time)
    points = 1:length(transmissions)
    return times, points
end


using Oceananigans: AbstractModel
using Oceananigans.Grids
using Oceananigans.Utils: tupleit

struct KernelComputedField{X, Y, Z, S, A, G, K, C, F, P} <: AbstractField{X, Y, Z, A, G}
                     data :: A
                     grid :: G
                   kernel :: K
      boundary_conditions :: C
    computed_dependencies :: F
               parameters :: P
                   status :: S

    function KernelComputedField{X, Y, Z}(kernel::K, arch, grid;
                                          boundary_conditions = ComputedFieldBoundaryConditions(grid, (X, Y, Z)),
                                          computed_dependencies = (),
                                          parameters::P = nothing,
                                          data = nothing,
                                          recompute_safely = true) where {X, Y, Z, K, P}

        computed_dependencies = tupleit(computed_dependencies)

        if isnothing(data)
            data = new_data(arch, grid, (X, Y, Z))
        end

        # Use FieldStatus if we want to avoid always recomputing
        status = recompute_safely ? nothing : FieldStatus(0.0)

        G = typeof(grid)
        A = typeof(data)
        S = typeof(status)
        F = typeof(computed_dependencies)
        C = typeof(boundary_conditions)

        return new{X, Y, Z, S,
                   A, G, K, C, F, P}(data, grid, kernel, boundary_conditions, computed_dependencies, parameters)
    end
end

"""
    KernelComputedField(X, Y, Z, kernel, model;
                        boundary_conditions = ComputedFieldBoundaryConditions(grid, (X, Y, Z)),
                        computed_dependencies = (),
                        parameters = nothing,
                        data = nothing,
                        recompute_safely = true)

Builds a `KernelComputedField` at `X, Y, Z` computed with `kernel` and `model.architecture` and `model.grid`, with `boundary_conditions`.

`computed_dependencies` are an iterable of `AbstractField`s or other objects on which `compute!` is called prior to launching `kernel`.

`data` is a three-dimensional `OffsetArray` of scratch space where the kernel computation is stored.

If `data=nothing` (the default) then additional memory will be allocated to store the `data` of `KernelComputedField`.

If `isnothing(parameters)`, `kernel` is launched with the function signature

`kernel(data, grid, computed_dependencies...)`

Otherwise, `kernel` is launched with the function signature

`kernel(data, grid, computed_dependencies..., parameters)`

`recompute_safely` (default: `true`) determines whether the `KernelComputedField` is "recomputed" if embedded in the expression
tree of another operation.
    - If `recompute_safely=true`, the `KernelComputedField` is always recomputed.
    - If `recompute_safely=false`, the `KernelComputedField` will not be recomputed if its status is up-to-date.
    - If `data=nothing`, then `recompute_safely` is switched to `false`.

Example
=======

```julia
using KernelAbstractions: @index, @kernel
using Oceananigans.Fields: AveragedField, KernelComputedField, compute!
using Oceananigans.Grids: Center, Face

@inline ψ²(i, j, k, grid, ψ, Ψ) = @inbounds (ψ[i, j, k] - Ψ[i, j, k])^2
@kernel function compute_variance!(var, grid, ϕ, Φ)
    i, j, k = @index(Global, NTuple)

    @inbounds var[i, j, k] = ψ′²(i, j, k, grid, ϕ, Φ)
end

u, v, w = model.velocities

U = AveragedField(u, dims=(1, 2))
V = AveragedField(v, dims=(1, 2))

u′² = KernelComputedField(Face, Center, Center, compute_variance!, model; computed_dependencies=(u, U,))
v′² = KernelComputedField(Center, Face, Center, compute_variance!, model; computed_dependencies=(v, V,))
w′² = KernelComputedField(Center, Center, Face, compute_variance!, model; computed_dependencies=(w, 0,))

compute!(u′²)
compute!(v′²)
compute!(w′²)
```
"""
KernelComputedField(X, Y, Z, kernel, model::AbstractModel; kwargs...) =
    KernelComputedField{X, Y, Z}(kernel, model.architecture, model.grid; kwargs...)

KernelComputedField(X, Y, Z, kernel, arch::AbstractArchitecture, grid::AbstractGrid; kwargs...) =
    KernelComputedField{X, Y, Z}(kernel, arch, grid; kwargs...)

function compute!(kcf::KernelComputedField{X, Y, Z}) where {X, Y, Z}

    for dependency in kcf.computed_dependencies
        compute!(dependency)
    end

    arch = architecture(kcf.data)

    workgroup, worksize = work_layout(kcf.grid,
                                      :xyz,
                                      location=(X, Y, Z),
                                      include_right_boundaries=true)

    compute_kernel! = kcf.kernel(device(arch), workgroup, worksize)

    event = isnothing(kcf.parameters) ?
        compute_kernel!(kcf.data, kcf.grid, kcf.computed_dependencies...) :
        compute_kernel!(kcf.data, kcf.grid, kcf.computed_dependencies..., kcf.parameters)

    wait(device(arch), event)

    fill_halo_regions!(kcf, arch)

    return nothing
end

compute!(field::KernelComputedField{X, Y, Z, <:FieldStatus}, time) where {X, Y, Z} =
    conditional_compute!(field, time)

Adapt.adapt_structure(to, kcf::KernelComputedField) = Adapt.adapt(to, kcf.data)


############################################################
# judiWavefield ##############################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: June 2017

export judiWavefield, judiWavefieldException, judiDFTwavefield, muteWavefield, dump_wavefield

############################################################

mutable struct judiWavefield{vDT<:Number} <: joAbstractLinearOperator{vDT,vDT}
	name::String
	m::Integer
	n::Integer
	info::Info
    dt
	data
end

mutable struct judiWavefieldException <: Exception
	msg :: String
end

############################################################

## outer constructors

function judiWavefield(info,dt::Real,data::Union{Array, PyCall.PyObject, String}; vDT::DataType=Float32)

	# length of vector
	m = info.n * sum(info.nt)
	n = 1
	dataCell = Array{Any}(undef, info.nsrc)
	for j=1:info.nsrc
		dataCell[j] = data
	end
	return judiWavefield{vDT}("judiWavefield",m,n,info,dt,dataCell)
end

function judiWavefield(info,dt::Real,data::Array{Any,1};vDT::DataType=Float32)
	# length of vector
	m = info.n * sum(info.nt)
	n = 1
	return judiWavefield{vDT}("judiWavefield",m,n,info,dt,data)
end

####################################################################
## overloaded Base functions

# conj(jo)
conj(A::judiWavefield{vDT}) where vDT =
	judiWavefield{vDT}("conj("*A.name*")",A.m,A.n,A.info,A.dt,A.data)

# transpose(jo)
transpose(A::judiWavefield{vDT}) where vDT =
	judiWavefield{vDT}(""*A.name*".'",A.n,A.m,A.info,A.dt,A.data)

# ctranspose(jo)
ctranspose(A::judiWavefield{vDT}) where vDT =
	judiWavefield{vDT}(""*A.name*"'",A.n,A.m,A.info,A.dt,A.data)

####################################################################

function vcat(a::judiWavefield{avDT},b::judiWavefield{bvDT}) where {avDT, bvDT}
	m = a.m + b.m
	n = 1
	nsrc = a.info.nsrc + b.info.nsrc
	data = Array{Any}(undef, nsrc)
	nt = Array{Any}(undef, nsrc)
	for j=1:a.info.nsrc
		data[j] = a.data[j]
		nt[j] = a.info.nt[j]
	end
	for j=a.info.nsrc+1:nsrc
		data[j] = b.data[j-a.info.nsrc]
		nt[j] = b.info.nt[j-a.info.nsrc]
	end
	info = Info(a.info.n,nsrc,nt)
	return judiWavefield(info,a.dt,data)
end

# add and subtract, mulitply and divide, norms, dot ...

# DFT operator for wavefields, acts along time dimension
function fft_wavefield(x_in,mode)
	nsrc = x_in.info.nsrc
	if mode==1
		x = judiWavefield(x_in.info,deepcopy(x_in.data); vDT=Complex{Float32})
		for i=1:nsrc
			x.data[i] = convert(Array{Complex{Float32}},x.data[i])
			nx = size(x.data[i],2)
			nz = size(x.data[i],3)
			for j=1:nx
				for k=1:nz
					x.data[i][:,j,k] = fft(x.data[i][:,j,k])
				end
			end
		end
	elseif mode==-1
		x = judiWavefield(x_in.info,deepcopy(x_in.data); vDT=Float32)
		for i=1:nsrc
			nx = size(x.data[i],2)
			nz = size(x.data[i],3)
			for j=1:nx
				for k=1:nz
					x.data[i][:,j,k] = real(ifft(x.data[i][:,j,k]))
				end
			end
			x.data[i] = convert(Array{Float32},real(x.data[i]))
		end
	end
	return x
end

# Sampling mask to extract wavefields from full vector
subsample(u::judiWavefield,srcnum) = judiWavefield(u.info,u.data[srcnum];vDT=eltype(u))

function muteWavefield(u_in::judiWavefield,ts_keep)
	u = deepcopy(u_in)
	for j=1:u.info.nsrc
		idx = ones(size(u.data[j],1));
		idx[ts_keep] = 0
		zero_idx = findall(idx)
		u.data[j][zero_idx,:,:] *= 0.f0
	end
	return u
end

# norm
function norm(a::judiWavefield{avDT}, p::Real=2) where avDT
	np = load_numpy()
    x = 0.f0
    for j=1:a.info.nsrc
        if typeof(a.data[j]) == String
            x += a.dt * sum(np[:abs](np[:load](a.data[j])).^p)
        else
            x += a.dt * sum(np[:abs](a.data[j]["data"]).^p)
        end
    end
    return x^(1.f0/p)
end

# Save wavefield to disk
function dump_wavefield(u::PyObject)
    name = join(["wavefield_", randstring(8), ".dat"])
    u["data"][:dump](name)
    return name
end

function dump_wavefield(u::judiWavefield)
    name = Array{Any}(undef, u.info.nsrc)
    for j=1:u.info.nsrc
        name[j] = join(["wavefield_", randstring(8), ".dat"])
        u.data[j]["data"][:dump](name[j])
    end
    return name
end


export System,ReadInput

function ReadInput(filename::String)
    input = open(filename, "r")
    inputs=JSON.parse(input)
    close(input)
    return inputs
end

struct System
    # User Inputs Simulation Setup
    dt::Float64
    temp::Float64
    steps::Int64
    start_step::Int64
    atoms::Atoms
    thermostat::Array{Thermostat,1}
    integrator::Integrator
    bdinter::Array{BdInter,1}
    field::Field
    velocity::Velocity
    thermologger::Logger
    trjdump::Dump
    restart::Restart
    pbc::BC

    first_step::Array{Bool,1}
    current_step::Array{Int64,1}
    current_temp::Array{Float64,1}
    current_momentum::Array{Float64,1}
    mpi::BalancingMPI

    # Initialize Constructor
    function System(input::Dict,config::Dict,comm::MPI.Comm,root::Int64)

        # local worker ID
        my_rank = MPI.Comm_rank(comm)
        comm_size=MPI.Comm_size(comm)

        # Initialize Simulation Setup
        start_step=1
        dt=0.005
        temp=1.0
        steps=100

        atoms=Atoms()
        pbc=NoPBC()

        bdinter=Array{BdInter,1}(undef,0)
        field=NoField()
        integrator=NoIntegrator()
        thermostat=Array{Thermostat,1}(undef,0)
        velocity=NoVelocity()
        trjdump=NoDump()
        restart=NoRestart()

        first_step=[true]
        current_step=[1]
        current_temp=[0.0]
        current_momentum=[0.0]

        thermologger=NoLogger()

        configuration=Configuration(Dict())
        mpi=BalancingMPI(configuration,comm_size)

        if my_rank==root

            if "dt" in keys(input)
                dt=input["dt"]
                @info "Timestep dt: $(dt)"
            end

            if "temp" in keys(input)
                temp=input["temp"]
                @info "Temperature T: $(temp)"
            end
            if "steps" in keys(input)
                steps=input["steps"]
                @info "Simulation Total Steps: $(steps)"
            end
            configuration=Configuration(config)
            mpi=BalancingMPI(configuration,MPI.Comm_size(comm))

            start_step=configuration.start
            current_step[1]=configuration.start

            pbc=PBC(;xhi=configuration.box[1],yhi=configuration.box[2],zhi=configuration.box[3])
            atoms.rho0[1]=sum(mpi.commsizes)/(pbc.Lx^3)

            types_all=InitTypes(configuration,mpi)

            atoms.num_atomtype[1]=length(unique(types_all))
            atoms.total_N[1]=length(types_all)

            if "velocity verlet" in keys(input)
                integrator=VVIntegrator(dt)
                @info "Velocity Velert Integrator"
            end

            if "gaussian velocity" in keys(input)
                temp=input["gaussian velocity"]["T"]
                seed=input["gaussian velocity"]["Seed"]
                velocity=GaussianVelocity(temp,seed)
                @info "Gaussian Velocity Initialization with Temperature: $(temp) and Seed:  $(seed)"
            elseif "zero velocity" in keys(input)
                velocity=ZeroVelocity()
                @info "Zero Velocity Initialization"
            end
            if "restart" in keys(input)
                if input["restart"]["JSONRestart"]
                    filename=input["restart"]["file"]
                    freq=input["restart"]["freq"]
                    if freq==-1
                        freq=Int(1e10)
                    end
                    restart=JSONRestart(filename,freq)

                end
            end

            masses_all=InitMasses(configuration,mpi)
            coords_all=InitCoords(configuration,mpi)
            forces_all=InitForces(configuration)
            vels_all=InitVels(configuration,mpi,velocity)
            energies_all=InitEnergies(configuration)

            init_atoms!(atoms,
                        types_all,
                        masses_all,
                        coords_all,
                        forces_all,
                        vels_all,
                        energies_all)

            if "harmonic bond" in keys(input)
                bonds_all=InitBonds(configuration,mpi)
                init_bonds!(atoms,bonds_all)
            end

            if "thermo information" in keys(input)
                thermologger=init(NoLogger(),input,start_step)
            end

            if "BerendsenNVT" in keys(input)
                thermostat_=init(BerendsenNVT(),input)
                push!(thermostat,thermostat_)
                @info "BerendsenNVT Thermostat with tau $(thermostat_.τ)"
            elseif "LangevinNVT" in keys(input)
                thermostat_=init(LangevinNVT(),input)
                push!(thermostat,thermostat_)
                @info "LangevinNVT Thermostat with gamma $(thermostat_.gamma)"
            end

            if "harmonic bond" in keys(input)
                bdinter_=init(Harmonic(),input)
                push!(bdinter,bdinter_)
                @info "Employing Harmonic Bond Potential with K= $(input["harmonic bond"]["k"]) and l0= $(input["harmonic bond"]["l0"])"
            end

            if "Canonical field" in keys(input)
                field=init(CanonicalField(),input,configuration)
                @info "Employing Canonical Field with Cell Length $(input["Canonical field"]["Lcell"]), Field Update every $(input["Canonical field"]["update"])"
            end

            if "LAMMPSTrj" in keys(input)
                filename=input["LAMMPSTrj"]["file"]
                freq=input["LAMMPSTrj"]["freq"]
                trjdump=LAMMPSTrj(filename,freq)
            end
        end

        new(# User Inputs Simulation Setup
        dt::Float64,
        temp::Float64 ,
        steps::Int64,
        start_step::Int64,

        atoms::Atoms,
        thermostat::Array{Thermostat,1},
        integrator::Integrator,
        bdinter::Array{BdInter,1},
        field::Field,
        velocity::Velocity,
        thermologger::Logger,
        trjdump::Dump,
        restart::Restart,
        pbc::BC,

        first_step::Array{Bool,1},
        current_step::Array{Int64,1},
        current_temp::Array{Float64,1},
        current_momentum::Array{Float64,1},
        mpi::BalancingMPI
        )
    end
end


module ModelClockSystem

    export ModelAlarm, ModelClock, advanceClock!, setClock!, addAlarm!, clock2str, dt2str, dt2tuple, dropRungAlarm!

    using CFTime, Formatting

    using Dates
    using Dates: CompoundPeriod

    mutable struct ModelAlarm
        name      :: String
        time      :: AbstractCFDateTime
        priority  :: Integer   # There might be multiple alarms ring at the same time. And if order is important, then the higher the number the higher the priority.
        callbacks :: Array{Function, 1}
        recurring :: Union{Nothing, Period, CompoundPeriod, Function}
        done      :: Bool
    end

    function isEarlier(x :: ModelAlarm, y :: ModelAlarm)
        return (x.time < y.time) || ( (x.time == y.time) && (x.priority > y.priority))
    end

    function isEqual(x :: ModelAlarm, y :: ModelAlarm)
        return (x.time == y.time) && (x.priority == y.priority)
    end

    function isLater(x :: ModelAlarm, y :: ModelAlarm)
        return ! ( isEarlier(x, y) ||  isEqual(x, y) )
    end

    mutable struct ModelClock

        name         :: String
        time         :: AbstractCFDateTime
        alarms       :: Array{ModelAlarm, 1}
        alarms_dict  :: Dict
        alarm_ptr    :: Integer

        function ModelClock(
            name :: String,
            time :: AbstractCFDateTime,
        )
            alarms = Array{ModelAlarm}(undef, 0)
            alarms_dict = Dict()
            return new(
                name,
                time + Dates.Second(0),
                alarms,
                alarms_dict,
                0,
            )
        end

    end

    function checkType(time1 :: AbstractCFDateTime, time2 :: AbstractCFDateTime)
        if typeof(time1) != typeof(time2)
            throw(ErrorException("Time type does not match!"))
        end
    end

    function setClock!(
        mc :: ModelClock,
        time :: AbstractCFDateTime,
    )

        checkType(mc.time, time)
        mc.time = time + Second(0)

    end

    function advanceClock!(
        mc :: ModelClock,
        t :: Union{Second, AbstractCFDateTime},
    )
        if typeof(t) <: Second  # inteprete as time interval
            mc.time += t
        elseif typeof(Δt) <: AbstractCFDateTime
            setClock!(mc, t)
        end

        # Only check alarms when there is any
        if length(mc.alarms) > 0
            if mc.alarm_ptr == 0 && mc.time >= mc.alarms[1].time
                mc.alarm_ptr = 1
            end

            while (0 < mc.alarm_ptr <= length(mc.alarms)) && (mc.time >= mc.alarms[mc.alarm_ptr].time)
            #    println("Before ring alarm: mc.alarm_ptr  = $(mc.alarm_ptr)")
                # In case we are at the end of alarm array

                alarm_ptr_tmp = mc.alarm_ptr
                if mc.alarm_ptr < length(mc.alarms)
                    mc.alarm_ptr += 1
                end
                ringAlarm!(mc, mc.alarms[alarm_ptr_tmp])
             #   println("xxx alarm_ptr = $(mc.alarm_ptr)")
                if mc.alarm_ptr == length(mc.alarms)
                    break
                end
            end
        end

    end

    function ringAlarm!(mc :: ModelClock, alarm :: ModelAlarm)
        #println("alarm.time = $(alarm.time). done = $(alarm.done)")
        if ! alarm.done
            println(format("Alarm '{:s}' rings at {:s}", alarm.name, dt2str(alarm.time)))

            for callback in alarm.callbacks
                callback(mc, alarm)
            end

            # IMPORTANT: alarm.done has to be here.
            # If it is after addAlarm!, there is a potential recursive loop because
            # advanceClock! might be called when calling addAlarm! (i.e. ring immediately)
            alarm.done = true

            if alarm.recurring != nothing

                if isa(alarm.recurring, Function)
                    next_time = alarm.recurring(alarm.time)
                else
                    next_time = alarm.time + alarm.recurring
                end
                #println("Current time: $(string(alarm.time))")
                #println("Next alarm: " * string(next_time))
                addAlarm!(
                    mc,
                    alarm.name,
                    next_time,
                    alarm.priority;
                    callback = alarm.callbacks,
                    recurring = alarm.recurring,
                )
            end
        end
    end

    function clock2str(mc :: ModelClock)
        return dt2str(mc.time)
    end

    function dt2str(dt)
        return format("{:04d}-{:02d}-{:02d} {:02d}:{:02d}:{:02d}", dt2tuple(dt)...)
    end

    function dt2tuple(dt)
        return (Dates.year(dt), Dates.month(dt), Dates.day(dt), Dates.hour(dt), Dates.minute(dt), Dates.second(dt))
    end

    function addAlarm!(
        mc   :: ModelClock,
        name :: String,
        time :: AbstractCFDateTime,
        priority :: Integer;
        callback :: Union{Array{Function, 1}, Function, Nothing} = nothing,
        recurring :: Union{ Nothing, Period, CompoundPeriod, Function } = nothing,
    )

        checkType(time, mc.time)

        if callback == nothing
            callbacks = Array{Function, 1}(undef, 0)
        elseif typeof(callback) <: Function
            callbacks = [ callback ]
        else # it is already an array
            callbacks = callback
        end

        alarm = ModelAlarm(
            name,
            time,
            priority,
            callbacks,
            recurring,
            false,
        )

        if mc.alarm_ptr == 0
            mc.alarm_ptr = 1
        end

        if alarm.time < mc.time
            println("alarm.time = ", dt2str(alarm.time))
            throw(ErrorException("Alarm needs to be set in the future."))
        end

        push!(mc.alarms, alarm)
        if ! haskey(mc.alarms_dict, name)
            mc.alarms_dict[name] = []
        end

        if (mc.alarm_ptr > 1) && (isEarlier(alarm, mc.alarms[mc.alarm_ptr]) || isEqual(alarm, mc.alarms[mc.alarm_ptr]))
            mc.alarm_ptr -= 1
        end

        push!(mc.alarms_dict[name], alarm)
        # lt = less than = earlier and higher priority
        sort!(mc.alarms, lt = isEarlier)

        if alarm.time == mc.time
            advanceClock!(mc, Second(0))
        end

    end

    function dropRungAlarm!(
        mc   :: ModelClock,
    )
        while length(mc.alarms) > 0 && mc.alarms[1].done
            alarm = popfirst!(mc.alarms)
            delete!(mc.alarms_dict, alarm.name)
            mc.alarm_ptr -= 1
            if mc.alarm_ptr < 0
                throw(ErrorException("Error happens in alarm counts. Please check."))
            end
        end
    end

    function listAlarms(
        mc :: ModelClock,
    )
        for (i, alarm) in enumerate(mc.alarms)
            println("Alarm[$(i)] = $(alarm.time). [$( (alarm.done) ? "v" : " " )] $( (i==mc.alarm_ptr) ? "*" : "" )")
        end
    end
end


using FFTW;

struct Person
    name::String
    age::Int64
end

rfft(A [, dims])


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Scenario-decomposition evaluation #
# ================================= #
function evaluate_decision(structure::ScenarioDecompositionStructure, decision::AbstractVector)
    return _eval_subproblems(structure, decision)
end

function statistically_evaluate_decision(structure::ScenarioDecompositionStructure, decision::AbstractVector)
    𝔼Q, σ² = _stat_eval_subproblems(structure, decision)
    return 𝔼Q, sqrt(σ²)
end

function _eval_subproblems(structure::ScenarioDecompositionStructure{2,1,Tuple{SP}},
                           decision::AbstractVector) where SP <: ScenarioProblems
    # Update decisions
    map(subprob -> take_decisions!(subprob, all_decision_variables(subprob, 1), decision), subproblems(structure))
    # Cache result
    result = outcome_mean(subproblems(structure), probability.(scenarios(structure)))
    # Revert back to untaken decisions
    map(subprob -> untake_decisions!(subprob, all_decision_variables(subprob, 1)), subproblems(structure))
    # Return evaluation result
    return result
end
function _eval_subproblems(structure::ScenarioDecompositionStructure{2,1,Tuple{SP}},
                           decision::AbstractVector) where SP <: DistributedScenarioProblems
    Qs = Vector{Float64}(undef, nworkers())
    sp = scenarioproblems(structure)
    @sync begin
        for (i,w) in enumerate(workers())
            @async Qs[i] = remotecall_fetch(
                w,
                sp[w-1],
                sp.decisions[w-1],
                decision) do sp, d, x
                    scenarioproblems = fetch(sp)
                    num_scenarios(scenarioproblems) == 0 && return 0.0
                    # Update decisions
                    map(subprob -> take_decisions!(subprob, all_decision_variables(subprob, 1), x), subproblems(scenarioproblems))
                    # Cache result
                    result = outcome_mean(subproblems(scenarioproblems),
                                          probability.(scenarios(scenarioproblems)))
                    # Revert back to untaken decisions
                    map(subprob -> untake_decisions!(subprob, all_decision_variables(subprob, 1)), subproblems(scenarioproblems))
                    # Return evaluation result
                    return result
                end
        end
    end
    return sum(Qs)
end

function _stat_eval_subproblems(structure::ScenarioDecompositionStructure{2,1,Tuple{SP}},
                                decision::AbstractVector) where SP <: ScenarioProblems
    # Update decisions
    map(subprob -> take_decisions!(subprob, all_decision_variables(subprob, 1), x), subproblems(structure))
    # Cache result
    result = welford(subproblems(structure), probability.(scenarios(structure)))
    # Revert back to untaken decisions
    map(subprob -> untake_decisions!(subprob, all_decision_variables(subprob, 1)), subproblems(structure))
    # Return evaluation result
    return result
end
function _stat_eval_subproblems(structure::ScenarioDecompositionStructure{2,1,Tuple{SP}},
                                decision::AbstractVector) where SP <: DistributedScenarioProblems
    partial_welfords = Vector{Tuple{Float64,Float64,Float64,Int}}(undef, nworkers())
    sp = scenarioproblems(structure)
    @sync begin
        for (i,w) in enumerate(workers())
            @async partial_welfords[i] = remotecall_fetch(
                w,
                sp[w-1],
                sp.decisions[w-1],
                x) do sp, d, x
                    scenarioproblems = fetch(sp)
                    num_scenarios(scenarioproblems) == 0 && return 0.0, 0.0, 0.0, 0
                    # Update decisions
                    map(subprob -> take_decisions!(subprob, all_decision_variables(subprob, 1), x), subproblems(scenarioproblems))
                    # Cache result
                    result = welford(subproblems(scenarioproblems), probability.(scenarios(scenarioproblems)))
                    # Revert back to untaken decisions
                    map(subprob -> untake_decisions!(subprob, all_decision_variables(subprob, 1)), subproblems(scenarioproblems))
                    # Return evaluation result
                    return result
                end
        end
    end
    𝔼Q, σ², _ = reduce(aggregate_welford, partial_welfords)
    return 𝔼Q, σ²
end


using RadeonProRender, GeometryBasics, Colors
const RPR = RadeonProRender

function translationmatrix(t::Vec{3,T}) where {T}
    T0, T1 = zero(T), one(T)
    return Mat{4}(T1, T0, T0, T0, T0, T1, T0, T0, T0, T0, T1, T0, t[1], t[2], t[3], T1)
end

context = RPR.Context()
scene = RPR.Scene(context)
matsys = RPR.MaterialSystem(context, 0)

camera = RPR.Camera(context)
lookat!(camera, Vec3f(8, 0, 5), Vec3f(2, 0, 2), Vec3f(0, 0, 1))
RPR.rprCameraSetFocalLength(camera, 45.0)
set!(scene, camera)
set!(context, scene)

env_light = RadeonProRender.EnvironmentLight(context)
image_path = joinpath(@__DIR__, "studio026.exr")
img = RPR.Image(context, image_path)
set!(env_light, img)
setintensityscale!(env_light, 1.1)
push!(scene, env_light)

light = RPR.PointLight(context)
transform!(light, translationmatrix(Vec3f(2, 0, 5)))
f = 20
RPR.setradiantpower!(light, 500 / f, 641 / f, 630 / f)
push!(scene, light)

function add_shape!(scene, context, matsys, mesh; material=RPR.RPR_MATERIAL_NODE_DIFFUSE,
                    color=colorant"green", roughness=0.01)
    rpr_mesh = RadeonProRender.Shape(context, mesh)
    push!(scene, rpr_mesh)
    m = RPR.MaterialNode(matsys, material)
    set!(m, RPR.RPR_MATERIAL_INPUT_COLOR, color)
    set!(m, RPR.RPR_MATERIAL_INPUT_ROUGHNESS, roughness, roughness, roughness, roughness)
    set!(rpr_mesh, m)
    return rpr_mesh, m
end

mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(-10, -10, -1), Vec3f(20, 20, 1));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, -10, 0), Vec3f(0.1, 20, 5));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, -2, 0), Vec3f(5, 0.1, 5));
                       color=colorant"white")
mesh, mat = add_shape!(scene, context, matsys, Rect3f(Vec3f(0, 2, 0), Vec3f(5, 0.1, 5));
                       color=colorant"white")

mesh, mat_sphere = add_shape!(scene, context, matsys, Tesselation(Sphere(Point3f(2, 0, 2), 1.0f0), 100);
                              material=RPR.RPR_MATERIAL_NODE_MICROFACET, roughness=0.2, color=colorant"red")

fb_size = (800, 600)
frame_buffer = RPR.FrameBuffer(context, RGBA, fb_size)
frame_bufferSolved = RPR.FrameBuffer(context, RGBA, fb_size)
set!(context, RPR.RPR_AOV_COLOR, frame_buffer)

set_standard_tonemapping!(context)

begin
    clear!(frame_buffer)
    RPR.rprContextSetParameterByKey1u(context, RPR.RPR_CONTEXT_ITERATIONS, 100)
    RPR.render(context)
    RPR.rprContextResolveFrameBuffer(context, frame_buffer, frame_bufferSolved, false)
    RPR.save(frame_bufferSolved, "test.png")
end