# Defines BLIS Objects and related methods.
#

export BliObj, BliObjBase

" BliObjBase resembles obj_t of BLIS' object-based interface. "
mutable struct BliObjBase
    # Basic fields.
    root::Ptr{BliObjBase}
    off::NTuple{2, BliDim}
    dim::NTuple{2, BliDim}
    diag_off::BliDoff
    info1::BliObjBits
    info2::BliObjBits
    enum_size::BliSiz
    buffer::Ptr{Nothing}
    rs::BliInc
    cs::BliInc
    is::BliInc

    # Bufferless scalar.
    scalar::BliAtomic

    # Pack-related fields.
    m_padded::BliDim
    n_padded::BliDim
    ps::BliInc
    pd::BliInc
    m_panel::BliDim
    n_panel::BliDim

    # Empty inner-constructor only,
    #  resembling a plain C struct.
    BliObjBase() = begin
        obj = new(Ptr{Nothing}(),
                  (0, 0),
                  (0, 0),
                  0, 0, 0, 0,
                  Ptr{Nothing}(), 0, 0, 0,
                  0,
                  0, 0, 0, 0,
                  0, 0)
    end
end

"""
Julia object BliObj packs obj_t with a reference to buffer array.
This is to preserve the base array from being recycled by GC module.
"""
struct BliObj
    obj::BliObjBase
    # Referee of obj.buffer, which itself is also a reference.
    # Purpose is to keep data alive.
    data::RefValue
end

" Object from strided matrix (view). "
BliObj(A::StridedMatrix{T}) where {T<:BliCompatibleType} = begin
    # Query data format.
    mA, nA = size(A)
    rsA, csA = strides(A)
    dtA = ctype_to_bli_num[T]

    # Initialize object with reference.
    objA = BliObj(BliObjBase(), Ref(A))
    # Pass address of A to object initializer.
    bli_obj_create_with_attached_buffer!(dtA, mA, nA,
                                         unsafe_convert(Ptr{Nothing}, A), rsA, csA,
                                         objA.obj)

    objA
end

" Object from strided vector. "
BliObj(x::StridedVector{T}) where {T<:BliCompatibleType} = begin
    mx, = size(x)
    incx, = strides(x)
    dt = ctype_to_bli_num[T]

    objx = BliObj(BliObjBase(), Ref(x))
    # Creation of vector buffer is roughly the same as matrix.
    bli_obj_create_with_attached_buffer!(dt, mx, 1,
                                         unsafe_convert(Ptr{Nothing}, x), incx, mx*incx,
                                         objx.obj)

    objx
end

" Object from scalar number. "
BliObj(γ::T) where {T<:BliCompatibleType} = begin
    dtγ = ctype_to_bli_num[T]

    objγ = BliObj(BliObjBase(), Ref(nothing))
    # bli_obj_create_1x1_with_attached_buffer!(dtγ,
    #                                          unsafe_convert(Ptr{Nothing}, objγ.data),
    #                                          objγ.obj)
    # Instead of something like above, use internal scalar storage.
    bli_obj_scalar_init_detached!(dtγ, objγ.obj)
    bli_setsc!(real(γ), imag(γ), objγ.obj)

    objγ
end

#----------------------------------------------------------
# Object-creation methods are exposed for the in-principle
#  consistency with BLIS C API and for advanced usages.
# Usually object constructors should be enough for Julia.
#
# NOTE: bli_obj_create, bli_obj_free, bli_obj_alloc_buffer,
#       bli_obj_create_1x1 and bli_obj_create_conf_to.
#       are not interfaced as Julia is a GC language.
#

bli_obj_create_without_buffer!(dt::BliNum,
                               m::BliDim,
                               n::BliDim,
                               obj::BliObjBase) = ccall(dlsym(libblis, :bli_obj_create_without_buffer),
                                                        Nothing,
                                                        (BliNum,
                                                         BliDim,
                                                         BliDim,
                                                         Ptr{BliObjBase}),
                                                        dt, m, n, Ref(obj))

bli_obj_attach_buffer!(p::Ptr{Nothing},
                       rs::BliInc,
                       cs::BliInc,
                       is::BliInc,
                       obj::BliObjBase) = ccall(dlsym(libblis, :bli_obj_attach_buffer),
                                                Nothing,
                                                (Ptr{Nothing},
                                                 BliInc,
                                                 BliInc,
                                                 BliInc,
                                                 Ptr{BliObjBase}),
                                                p, rs, cs, is, Ref(obj))

bli_obj_create_with_attached_buffer!(dt::BliNum,
                                     m::BliDim,
                                     n::BliDim,
                                     p::Ptr{Nothing},
                                     rs::BliInc,
                                     cs::BliInc,
                                     obj::BliObjBase) = ccall(dlsym(libblis, :bli_obj_create_with_attached_buffer),
                                                              Nothing,
                                                              (BliNum,
                                                               BliDim,
                                                               BliDim,
                                                               Ptr{Nothing},
                                                               BliInc,
                                                               BliInc,
                                                               Ptr{BliObjBase}),
                                                              dt, m, n, p, rs, cs, Ref(obj))

bli_obj_create_1x1_with_attached_buffer!(dt::BliNum,
                                         p::Ptr{Nothing},
                                         obj::BliObjBase) = ccall(dlsym(libblis, :bli_obj_create_1x1_with_attached_buffer),
                                                                  Nothing,
                                                                  (BliNum,
                                                                   Ptr{Nothing},
                                                                   Ptr{BliObjBase}),
                                                                  dt, p, Ref(obj))

bli_obj_scalar_init_detached!(dt::BliNum,
                              obj::BliObjBase) = ccall(dlsym(libblis, :bli_obj_scalar_init_detached),
                                                       Nothing,
                                                       (BliNum,
                                                        Ptr{BliObjBase}),
                                                       dt, Ref(obj))

bli_setsc!(zeta_r::Real,
           zeta_i::Real,
           obj::BliObjBase) = ccall(dlsym(libblis, :bli_setsc),
                                    Nothing,
                                    (Float64,
                                     Float64,
                                     Ptr{BliObjBase}),
                                    Float64(zeta_r), Float64(zeta_i), Ref(obj))

#----------------------------------------------------------------
# These methods replicates BLIS' C marcos for setting obj_t.info.
# They are exposed in a minimum fashion.

bli_obj_set_info_bits!(dest::BliObjBits,
                       src::BliObjBits,
                       obj::BliObjBase) = begin
    obj.info1 = UInt32( ( obj.info1 & ~dest ) | src )

    nothing
end

bli_obj_set_conjtrans!(t::BliTrans,
                       obj::BliObjBase) = bli_obj_set_info_bits!(BLIS_CONJTRANS_BITS,
                                                                 BliObjBits(t.enum),
                                                                 obj)

bli_obj_set_onlytrans!(t::BliTrans,
                       obj::BliObjBase) = bli_obj_set_info_bits!(BLIS_TRANS_BIT,
                                                                 BliObjBits(t.enum),
                                                                 obj)

bli_obj_set_conj!(c::BliConj,
                  obj::BliObjBase) = bli_obj_set_info_bits!(BLIS_CONJ_BIT,
                                                            BliObjBits(c.enum),
                                                            obj)

bli_obj_set_uplo!(ul::BliUpLo,
                  obj::BliObjBase) = bli_obj_set_info_bits!(BLIS_UPLO_BITS,
                                                            BliObjBits(ul.enum),
                                                            obj)

bli_obj_set_diag!(d::BliDiag,
                  obj::BliObjBase) = bli_obj_set_info_bits!(BLIS_UNIT_DIAG_BIT,
                                                            BliObjBits(d.enum),
                                                            obj)

bli_obj_set_invert_diag!(id::BliInvDiag,
                         obj::BliObjBase) = bli_obj_set_info_bits!(BLIS_INVERT_DIAG_BIT,
                                                                   BliObjBits(id.enum),
                                                                   obj)

bli_obj_set_struc!(st::BliStruc,
                   obj::BliObjBase) = bli_obj_set_info_bits!(BLIS_STRUC_BITS,
                                                             BliObjBits(st.enum),
                                                             obj)


module AxPlot

using PyPlot
using PyCall
using LinearAlgebra
using Flux: Tracker
using Requires

include("main.jl")
include("utils.jl")

function __init__()
    @require InferGMM="e4654f38-ff99-11e8-3938-77f3ff3d0ead" include("gmm.jl")
end

end


using LinearAlgebra, Test, WaveFD

@testset "compressedio, 3D Array, T=$(T), C=$(C), style=$(style)" for (T,C) in ((Float32, UInt32), (Float64, UInt32), (Float32, Float32), (Float64, Float64), (Float64, Float32)), style in (:file, :buffer)
    nz = 512
    ny = 256
    nx = 128

    bz = 32
    by = 32
    bx = 32

    x = -1 .+ 2 .* rand(T,nz,ny,nx)
    y = -1 .+ 2 .* rand(T,nz,ny,nx)

    C,P = WaveFD.comptype(C, T)
    comp = WaveFD.Compressor(T, P, C, (nz,ny,nx), (bz,by,bx), 1e-2, 2, false)
    open(comp)
    io = style == :file ? open("test.bin","w") : IOBuffer()
    WaveFD.compressedwrite(io, comp, 1, x)
    WaveFD.compressedwrite(io, comp, 2, y)
    style == :file && close(io)

    xx = similar(x)
    yy = similar(y)

    io = style == :file ? open("test.bin") : io
    WaveFD.compressedread!(io, comp, 1, xx)
    WaveFD.compressedread!(io, comp, 2, yy)
    close(io)

    snr_x = 10*log10(norm(x)^2 / norm(x .- xx)^2)
    snr_y = 10*log10(norm(y)^2 / norm(y .- yy)^2)
    err_x = norm(x .- xx) / length(x)
    err_y = norm(y .- yy) / length(y)
    if C == UInt32
        @test snr_x > 40
        @test snr_y > 40
    else
        @test err_x < 1e-6
        @test err_y < 1e-6
    end
    close(comp)
end

@testset "compressedio, 3D SubArray, T=$(T), C=$(C), style=$(style)" for (T,C) in ((Float32, UInt32), (Float64, UInt32), (Float32, Float32), (Float64, Float64), (Float64, Float32)), style in (:file, :buffer)
    nz = 512
    ny = 256
    nx = 128

    bz = 32
    by = 32
    bx = 32

    x = -1 .+ 2 .* rand(T,nz,ny,nx)
    y = -1 .+ 2 .* rand(T,nz,ny,nx)

    rng = (10:nz-10,10:ny-10,10:nx-10)

    x_sub = @view x[rng...]
    y_sub = @view y[rng...]

    nz, ny, nx = size(x_sub)

    C,P = WaveFD.comptype(C, T)
    comp = WaveFD.Compressor(T, P, C, (nz,ny,nx), (bz,by,bx), 1e-2, 2, true)
    open(comp)

    io = style == :file ? open("test.bin", "w") : IOBuffer()
    WaveFD.compressedwrite(io, comp, 1, x, rng)
    WaveFD.compressedwrite(io, comp, 2, y, rng)
    style == :file && close(io)

    xx = copy(x)
    yy = copy(y)

    xx_sub = @view xx[rng...]
    yy_sub = @view yy[rng...]

    xx_sub .= 0.0
    yy_sub .= 0.0

    io = style == :file ? open("test.bin") : io
    WaveFD.compressedread!(io, comp, 1, xx, rng)
    WaveFD.compressedread!(io, comp, 2, yy, rng)
    close(io)

    snr_x = 10*log10(norm(x)^2 / norm(x .- xx)^2)
    snr_y = 10*log10(norm(y)^2 / norm(y .- yy)^2)
    err_x = norm(x .- xx) / length(x)
    err_y = norm(y .- yy) / length(y)
    if C == UInt32
        @test snr_x > 40
        @test snr_y > 40
    else
        @test err_x < 1e-6
        @test err_y < 1e-6
    end
    close(comp)
end

@testset "compressedio, 2D Array, T=$(T), C=$(C), style=$(style)" for (T,C) in ((Float32, UInt32), (Float64, UInt32), (Float32, Float32), (Float64, Float64), (Float64, Float32)), style in (:file, :buffer)
    nz = 512
    nx = 128

    bz = 32
    bx = 32

    x = -1 .+ 2 .* rand(T,nz,nx)
    y = -1 .+ 2 .* rand(T,nz,nx)

    C,P = WaveFD.comptype(C, T)
    comp = WaveFD.Compressor(T, P, C, (nz,nx), (bz,bx), 1e-2, 2, false)
    open(comp)

    io = style == :file ? open("test.bin", "w") : IOBuffer()
    WaveFD.compressedwrite(io, comp, 1, x)
    WaveFD.compressedwrite(io, comp, 2, y)
    style == :file && close(io)

    xx = similar(x)
    yy = similar(y)

    io = style == :file ? open("test.bin") : io
    WaveFD.compressedread!(io, comp, 1, xx)
    WaveFD.compressedread!(io, comp, 2, yy)
    close(io)

    snr_x = 10*log10(norm(x)^2 / norm(x .- xx)^2)
    snr_y = 10*log10(norm(y)^2 / norm(y .- yy)^2)
    err_x = norm(x .- xx) / length(x)
    err_y = norm(y .- yy) / length(y)
    if C == UInt32
        @test snr_x > 40
        @test snr_y > 40
    else
        @test err_x < 1e-6
        @test err_y < 1e-6
    end
    close(comp)
end

@testset "compressedio 2D SubArray, T=$(T), C=$(C), style=$(style)" for (T,C) in ((Float32, UInt32), (Float64, UInt32), (Float32, Float32), (Float64, Float64), (Float64, Float32)), style in (:file, :buffer)
    nz = 512
    nx = 128

    bz = 32
    bx = 32

    x = -1 .+ 2 .* rand(T,nz,nx)
    y = -1 .+ 2 .* rand(T,nz,nx)

    rng = (10:nz-10,10:nx-10)

    x_sub = @view x[rng...]
    y_sub = @view y[rng...]

    nz, nx = size(x_sub)

    C,P = WaveFD.comptype(C, T)
    comp = WaveFD.Compressor(T, P, C, (nz,nx), (bz,bx), 1e-2, 2, true)
    open(comp)

    io = style == :file ? open("test.bin", "w") : IOBuffer()
    WaveFD.compressedwrite(io, comp, 1, x, rng)
    WaveFD.compressedwrite(io, comp, 2, y, rng)
    style == :file && close(io)

    xx = copy(x)
    yy = copy(y)

    xx_sub = @view xx[rng...]
    yy_sub = @view yy[rng...]

    xx_sub[:] .= 0.0
    yy_sub[:] .= 0.0

    io = style == :file ? open("test.bin") : io
    WaveFD.compressedread!(io, comp, 1, xx, rng)
    WaveFD.compressedread!(io, comp, 2, yy, rng)
    close(io)

    snr_x = 10*log10(norm(x)^2 / norm(x .- xx)^2)
    snr_y = 10*log10(norm(y)^2 / norm(y .- yy)^2)
    err_x = norm(x .- xx) / length(x)
    err_y = norm(y .- yy) / length(y)
    if C == UInt32
        @test snr_x > 40
        @test snr_y > 40
    else
        @test err_x < 1e-6
        @test err_y < 1e-6
    end
    close(comp)
end

@testset "compressedio, exception" for file in ("test.bin",)
    io = open(file,"w")
    @test_throws ErrorException WaveFD.compressedwrite_exception(io,0,1000)
    rm(file)
end

nothing


# Constraints

"""
    supports_constraint(model::ModelLike, ::Type{F}, ::Type{S})::Bool where {F<:AbstractFunction,S<:AbstractSet}

Return a `Bool` indicating whether `model` supports `F`-in-`S` constraints, that is,
`copy_to(model, src)` does not return `CopyUnsupportedConstraint` when `src` contains `F`-in-`S` constraints.
If `F`-in-`S` constraints are only not supported in specific circumstances, e.g. `F`-in-`S` constraints cannot be combined with another type of constraint, it should still return `true`.
"""
supports_constraint(model::ModelLike, ::Type{<:AbstractFunction}, ::Type{<:AbstractSet}) = false

"""
    struct UnsupportedConstraint{F<:AbstractFunction, S<:AbstractSet} <: UnsupportedError
        message::String # Human-friendly explanation why the attribute cannot be set
    end

An error indicating that constraints of type `F`-in-`S` are not supported by
the model, i.e. that [`supports_constraint`](@ref) returns `false`.
"""
struct UnsupportedConstraint{F<:AbstractFunction, S<:AbstractSet} <: UnsupportedError
    message::String # Human-friendly explanation why the attribute cannot be set
end
UnsupportedConstraint{F, S}() where {F, S} = UnsupportedConstraint{F, S}("")

element_name(::UnsupportedConstraint{F, S}) where {F, S} = "`$F`-in-`$S` constraint"

"""
    struct AddConstraintNotAllowed{F<:AbstractFunction, S<:AbstractSet} <: NotAllowedError
        message::String # Human-friendly explanation why the attribute cannot be set
    end

An error indicating that constraints of type `F`-in-`S` are supported (see
[`supports_constraint`](@ref)) but cannot be added.
"""
struct AddConstraintNotAllowed{F<:AbstractFunction, S<:AbstractSet} <: NotAllowedError
    message::String # Human-friendly explanation why the attribute cannot be set
end
AddConstraintNotAllowed{F, S}() where {F, S} = AddConstraintNotAllowed{F, S}("")

operation_name(::AddConstraintNotAllowed{F, S}) where {F, S} = "Adding `$F`-in-`$S` constraints"

"""
    struct ScalarFunctionConstantNotZero{T, F, S} <: Exception
        constant::T
    end

An error indicating that the constant part of the function in the constraint
`F`-in-`S` is nonzero.
"""
struct ScalarFunctionConstantNotZero{T, F, S} <: Exception
    constant::T
end

function Base.showerror(io::IO,
                    err::ScalarFunctionConstantNotZero{T, F, S}) where {T, F, S}
    print(io, "In `$F`-in-`$S` constraint: Constant $(err.constant) of the ",
          "function is not zero. The function constant should be moved to the ",
          "set. You can use `MOI.Utilities.add_scalar_constraint` which does ",
          "this automatically.")
end

"""
    add_constraint(model::ModelLike, func::F, set::S)::ConstraintIndex{F,S} where {F,S}

Add the constraint ``f(x) \\in \\mathcal{S}`` where ``f`` is defined by `func`, and ``\\mathcal{S}`` is defined by `set`.

    add_constraint(model::ModelLike, v::VariableIndex, set::S)::ConstraintIndex{SingleVariable,S} where {S}
    add_constraint(model::ModelLike, vec::Vector{VariableIndex}, set::S)::ConstraintIndex{VectorOfVariables,S} where {S}

Add the constraint ``v \\in \\mathcal{S}`` where ``v`` is the variable (or vector of variables) referenced by `v` and ``\\mathcal{S}`` is defined by `set`.

* An [`UnsupportedConstraint`](@ref) error is thrown if `model` does not support
  `F`-in-`S` constraints,
* a [`AddConstraintNotAllowed`](@ref) error is thrown if it supports `F`-in-`S`
  constraints but it cannot add the constraint(s) in its current state and
* a [`ScalarFunctionConstantNotZero`](@ref) error may be thrown if
  `func` is an `AbstractScalarFunction` with nonzero constant and `set`
  is [`EqualTo`](@ref), [`GreaterThan`](@ref), [`LessThan`](@ref) or
  [`Interval`](@ref).
* a [`LowerBoundAlreadySet`](@ref) error is thrown if `F` is a
  [`SingleVariable`](@ref) and a constraint was already added to this variable
  that sets a lower bound.
* a [`UpperBoundAlreadySet`](@ref) error is thrown if `F` is a
  [`SingleVariable`](@ref) and a constraint was already added to this variable
  that sets an upper bound.
"""
function add_constraint(model::ModelLike, func::AbstractFunction,
                        set::AbstractSet)
    throw_add_constraint_error_fallback(model, func, set)
end

# throw_add_constraint_error_fallback checks whether func and set are both
# scalar or both vector. If it is the case, it calls
# `correct_throw_add_constraint_error_fallback`
function throw_add_constraint_error_fallback(model::ModelLike,
                                             func::AbstractScalarFunction,
                                             set::AbstractScalarSet;
                                             kwargs...)
    correct_throw_add_constraint_error_fallback(model, func, set; kwargs...)
end
function throw_add_constraint_error_fallback(model::ModelLike,
                                             func::AbstractVectorFunction,
                                             set::AbstractVectorSet;
                                             kwargs...)
    correct_throw_add_constraint_error_fallback(model, func, set; kwargs...)
end
function throw_add_constraint_error_fallback(model::ModelLike,
                                             func::AbstractScalarFunction,
                                             set::AbstractVectorSet;
                                             kwargs...)
    error("Cannot add a constraint of the form `ScalarFunction`-in-`VectorSet`")
end
function throw_add_constraint_error_fallback(model::ModelLike,
                                             func::AbstractVectorFunction,
                                             set::AbstractScalarSet;
                                             kwargs...)
    error("Cannot add a constraint of the form `VectorFunction`-in-`ScalarSet`")
end

# func and set are both scalar or both vector
function correct_throw_add_constraint_error_fallback(model::ModelLike,
                                                     func::AbstractFunction,
                                                     set::AbstractSet;
                                                     error_if_supported=AddConstraintNotAllowed{typeof(func), typeof(set)}())
    if supports_constraint(model, typeof(func), typeof(set))
        throw(error_if_supported)
    else
        throw(UnsupportedConstraint{typeof(func), typeof(set)}())
    end
end

# convenient shorthands TODO: document
add_constraint(model::ModelLike, v::VariableIndex, set::AbstractScalarSet) = add_constraint(model, SingleVariable(v), set)
add_constraint(model::ModelLike, v::Vector{VariableIndex}, set::AbstractVectorSet) = add_constraint(model, VectorOfVariables(v), set)

"""
    add_constraints(model::ModelLike, funcs::Vector{F}, sets::Vector{S})::Vector{ConstraintIndex{F,S}} where {F,S}

Add the set of constraints specified by each function-set pair in `funcs` and `sets`. `F` and `S` should be concrete types.
This call is equivalent to `add_constraint.(model, funcs, sets)` but may be more efficient.
"""
function add_constraints end

# default fallback
add_constraints(model::ModelLike, funcs, sets) = add_constraint.(model, funcs, sets)

"""
    LowerBoundAlreadySet{S1, S2}

Error thrown when setting a `SingleVariable`-in-`S2` when a
`SingleVariable`-in-`S1` has already been added and the sets `S1`, `S2` both
set a lower bound, i.e. they are [`EqualTo`](@ref), [`GreaterThan`](@ref),
[`Interval`](@ref), [`Semicontinuous`](@ref) or [`Semiinteger`](@ref).
"""
struct LowerBoundAlreadySet{S1, S2}
    vi::VariableIndex
end

function Base.showerror(io::IO, err::LowerBoundAlreadySet{S1, S2}) where {S1, S2}
    print(io, typeof(err), ": Cannot add `SingleVariable`-in`", S2,
          "` constraint for variable ", err.vi, " as a `SingleVariable`-in`",
          S1, "` constraint was already set for this variable and both",
          " constraints set a lower bound.")
end

"""
    UpperBoundAlreadySet{S1, S2}

Error thrown when setting a `SingleVariable`-in-`S2` when a
`SingleVariable`-in-`S1` has already been added and the sets `S1`, `S2` both
set an upper bound, i.e. they are [`EqualTo`](@ref), [`LessThan`](@ref),
[`Interval`](@ref), [`Semicontinuous`](@ref) or [`Semiinteger`](@ref).
"""
struct UpperBoundAlreadySet{S1, S2}
    vi::VariableIndex
end

function Base.showerror(io::IO, err::UpperBoundAlreadySet{S1, S2}) where {S1, S2}
    print(io, typeof(err), ": Cannot add `SingleVariable`-in`", S2,
          "` constraint for variable ", err.vi, " as a `SingleVariable`-in`",
          S1, "` constraint was already set for this variable and both",
          " constraints set an upper bound.")
end

"""
## Transform Constraint Set

    transform(model::ModelLike, c::ConstraintIndex{F,S1}, newset::S2)::ConstraintIndex{F,S2}

Replace the set in constraint `c` with `newset`. The constraint index `c`
will no longer be valid, and the function returns a new constraint index with
the correct type.

Solvers may only support a subset of constraint transforms that they perform
efficiently (for example, changing from a `LessThan` to `GreaterThan` set). In
addition, set modification (where `S1 = S2`) should be performed via the
`modify` function.

Typically, the user should delete the constraint and add a new one.

### Examples

If `c` is a `ConstraintIndex{ScalarAffineFunction{Float64},LessThan{Float64}}`,

```julia
c2 = transform(model, c, GreaterThan(0.0))
transform(model, c, LessThan(0.0)) # errors
```
"""
function transform end

# default fallback
function transform(model::ModelLike, c::ConstraintIndex, newset)
    f = get(model, ConstraintFunction(), c)
    delete(model, c)
    add_constraint(model, f, newset)
end


using HCIToolbox
using StableRNGs
using Statistics
using Test

rng = StableRNG(8799)

@testset "HCIToolbox.jl" begin
    include("morphology.jl")
    include("masking.jl")
    include("angles.jl")
    include("inject.jl")
    include("scaling.jl")
    include("geometry.jl")
    include("profiles.jl")
end


"""
    struct WeightingPotential{T, N, S, AT} <: AbstractArray{T, N}

Weighting potential for a given [`Contact`](@ref) which is a unitless potential.

## Parametric types
* `T`: Element type of `data`.
* `N`: Dimension of the `grid` and `data` array.
* `S`: Coordinate system (`Cartesian` or `Cylindrical`).
* `AT`: Axes type.

## Fields
* `data::Array{T, N}`: Array containing the values of the weighting potential at the discrete points of the `grid`.
* `grid::Grid{T, N, S, AT}`: [`Grid`](@ref) defining the discrete points for which the weighting potential is determined.
"""
struct WeightingPotential{T, N, S <: AbstractCoordinateSystem, AT} <: AbstractArray{T, N}
    data::Array{T, N}
    grid::Grid{T, N, S, AT}
end

@inline size(wpot::WeightingPotential{T, N, S}) where {T, N, S} = size(wpot.data)
@inline length(wpot::WeightingPotential{T, N, S}) where {T, N, S} = length(wpot.data)
@inline getindex(wpot::WeightingPotential{T, N, S}, I::Vararg{Int, N}) where {T, N, S} = getindex(wpot.data, I...)
@inline getindex(wpot::WeightingPotential{T, N, S}, i::Int) where {T, N, S} = getindex(wpot.data, i)
@inline getindex(wpot::WeightingPotential{T, N, S}, s::Symbol) where {T, N, S} = getindex(wpot.grid, s)

function WeightingPotential(nt::NamedTuple)
    grid = Grid(nt.grid)
    T = typeof(ustrip(nt.values[1]))
    S = get_coordinate_system(grid)
    N = get_number_of_dimensions(grid)
    WeightingPotential{T, N, S, typeof(grid.axes)}( nt.values, grid)
end
Base.convert(T::Type{WeightingPotential}, x::NamedTuple) = T(x)

function NamedTuple(wpot::WeightingPotential{T, 3}) where {T}
    return (
        grid = NamedTuple(wpot.grid),
        values = wpot.data,
    )
end
Base.convert(T::Type{NamedTuple}, x::WeightingPotential) = T(x)


function build_robot!(vis::Visualizer, model::WalledCartpole; name::Symbol=:WalledCartpole,
		r=0.05, α=1.0)
	r = convert(Float32, r)
	r_contact = convert(Float32, r * 1.5)

	wall_thickness = 0.15
	body_mat = MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, α))
	contact_mat = MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, α))
	wall_mat = MeshPhongMaterial(color = RGBA(0.7, 0.7, 0.7, 1.0))

	default_background!(vis)

	link = Cylinder(Point3f0(0.0), Point3f0(0.0, 0.0, model.l), r/2)
	setobject!(vis[name][:robot]["link"], link, body_mat)
	vec_base = Vec(2r_contact, 4r, 1.5r)
	vec_shaft = Vec(3model.w + 2r_contact + 2wall_thickness, 0.5r, 0.5r)
	setobject!(vis[name][:robot]["linkbase"], Rect(-0.5 * vec_base, vec_base), body_mat)
	setobject!(vis[name][:robot]["shaft"], Rect(-0.5 * vec_shaft, vec_shaft), body_mat)
	setobject!(vis[name][:robot]["contact"], Sphere(Point3f0(0.0), r_contact), contact_mat)

	setobject!(vis[name][:env]["wall1"]["init"],
		Rect(Vec(0, 0, 0),Vec(wall_thickness, wall_thickness, 0.75)), wall_mat)
	setobject!(vis[name][:env]["wall2"]["init"],
		Rect(Vec(0, 0, 0),Vec(wall_thickness, wall_thickness, 0.75)), wall_mat)
	settransform!(vis[name][:env]["wall1"]["init"], Translation([-wall_thickness - model.w - 1.5 * r; -wall_thickness/2; 3r]))
	settransform!(vis[name][:env]["wall2"]["init"], Translation([model.w + 1.5 * r;                   -wall_thickness/2; 3r]))
	return nothing
end

function set_robot!(vis::Visualizer, model::WalledCartpole, q::AbstractVector;
		name::Symbol=:WalledCartpole)

	pt = [_kinematics(model, q, mode =  :tip)[1], 0.0, _kinematics(model, q, mode =  :tip)[2]]
    pb = [_kinematics(model, q, mode = :base)[1], 0.0, _kinematics(model, q, mode = :base)[2]]

	settransform!(vis[name][:robot]["link"], cable_transform(pt, pb))
	settransform!(vis[name][:robot]["linkbase"], Translation(pb))
	settransform!(vis[name][:robot]["contact"], Translation(pt))

	xw1 = q[3]
	xw2 = q[4]
	settransform!(vis[name][:env]["wall1"], Translation([xw1, 0.0, 0.0]))
	settransform!(vis[name][:env]["wall2"], Translation([xw2, 0.0, 0.0]))

	return nothing
end


include("train.jl")

stack = @stack MetaTac stack_input=true begin
  Conv(1024, relu, stride = 3, padding = 0)
  Dense(512, relu)
end

model = NeuralModel(MetaTac, stack) |> to_gpu |> Async
player = MCTSPlayer(model, power = 250)

mcts_opponents = [MCTSPlayer(power = p) for p in [10, 50, 100, 200, 500]]
self_opponents = [MCTSPlayer(model, power = p, name = "current$p") for p in [10, 50]]

loss = Loss(policy = 0.5, reg = 1e-5)

with_contest( train_self!
            , player
            , loss = loss
            , playings = 50
            , batchsize = 100
            , branching = 0.02
            , iterations = 10
            , epochs = 100
            , augment = true
            , opponents = [mcts_opponents; self_opponents]
            , length = 1000
            , interval = 5
            )


module ReliablePowerSimulations

#################################################################################
# Exports
export ThermalStandardUCOutages
export ThermalDispatchOutages
export HydroDispatchRunOfRiverOutage
export HydroDispatchReservoirStorageOutage
export HydroDispatchPumpedStoragewReservationOutage
export BookKeepingwReservationOutage
export EndOfPeriodEnergyTargetOutage
export RenewableOutageDispatch

#################################################################################
# Imports
using PowerSystems
import InfrastructureSystems
import Dates
import PowerSimulations
import PowerModels
import JuMP
import ParameterJuMP

const PSY = PowerSystems
const IS = InfrastructureSystems
const PM = PowerModels
const PSI = PowerSimulations
const PJ = ParameterJuMP

#################################################################################
# Includes
include("./core/variables.jl")
include("./core/initial_condition_types.jl")
include("./core/update_initial_conditions.jl")

include("device_models/devices/common/constraints_structs.jl")
include("device_models/devices/common/duration_constraint.jl")
include("device_models/devices/common/outage_constraint.jl")

include("./device_models/devices/thermal_generation.jl")
include("./device_models/devices/renewable_generation.jl")
include("./device_models/devices/hydro_generation.jl")
include("./device_models/devices/storage.jl")

include("./device_models/device_constructors/thermalgeneration_constructor.jl")
include("./device_models/device_constructors/hydrogeneration_constructor.jl")
include("./device_models/device_constructors/renewablegeneration_constructor.jl")
include("./device_models/device_constructors/storage_constructor.jl")

end # module


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    struct PartitionedSampling <: AbstractSamplingAlgorithm

*Experimental feature, not part of stable public API.*

A sampling algorithm that partitions parameter space into multiple subspaces and
samples/integrates them independently ([Caldwell et al.](https://arxiv.org/abs/2008.03098)).

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct PartitionedSampling{
    TR<:AbstractDensityTransformTarget,
    S<:AbstractSamplingAlgorithm,
    E<:AbstractSamplingAlgorithm,
    I<:IntegrationAlgorithm,
    P<:SpacePartitioningAlgorithm
} <: AbstractSamplingAlgorithm
    trafo::TR = PriorToUniform()
    npartitions::Integer = 10
    sampler::S = MCMCSampling(strict=false)
    exploration_sampler::E = MCMCSampling(nchains=30, nsteps = 800, strict=false)
    partitioner::P = KDTreePartitioning()
    integrator::I = AHMIntegration()
    nmax_resampling::AbstractFloat = 5
end

export PartitionedSampling

function bat_sample_impl(rng::AbstractRNG, target::PosteriorDensity, algorithm::PartitionedSampling)

    density_notrafo = convert(AbstractDensity, target)
    shaped_density, trafo = bat_transform(algorithm.trafo, density_notrafo)
    density = unshaped(shaped_density)

    @info "Generating Exploration Samples"
    exploration_samples = bat_sample(density, algorithm.exploration_sampler).result

    @info "Constructing Partition Tree"
    partition_tree, cost_values = partition_space(exploration_samples, algorithm.npartitions, algorithm.partitioner)
    # Convert 'partition_tree' structure into a set of truncated posteriors:
    posteriors_array = convert_to_posterior(density, partition_tree, extend_bounds = algorithm.partitioner.extend_bounds)

    @info "Sampling Subspaces"
    iterator_subspaces = [
        [subspace_ind, posteriors_array[subspace_ind], algorithm.sampler] for subspace_ind in Base.OneTo(algorithm.npartitions)]
    samples_subspaces_run = pmap(inp -> sample_subspace(inp...), iterator_subspaces)

    unconv_mask = [samples_subspace.isvalid for samples_subspace in samples_subspaces_run] # returns "false" if subspace was not converged during tuning cycle
    unconv_ind = findall(x->x==false, unconv_mask)
    rep_sspace = !isempty(unconv_ind) # perform resampling if "true"

    if algorithm.nmax_resampling > 0
        samples_subspaces = samples_subspaces_run[unconv_mask]
    else
        samples_subspaces = samples_subspaces_run
    end

    rec_level = 1
    resampled_trees = []

    while rep_sspace && rec_level <= algorithm.nmax_resampling

        @info "Re-Sampling Subspaces: #$rec_level"

        posteriors_array_run = []

        for (ind, rep_ind) in enumerate(unconv_ind)

            exploration_samples_rep = bat_sample(
                samples_subspaces_run[rep_ind].samples,
                OrderedResampling(nsamples=algorithm.exploration_sampler.nsteps)
            ).result

            partition_tree_rep, _ = partition_space(exploration_samples_rep, 2, algorithm.partitioner)
            push!(resampled_trees, partition_tree_rep)
            posteriors_rep_array = convert_to_posterior(posteriors_array[rep_ind], partition_tree_rep, extend_bounds = algorithm.partitioner.extend_bounds)

            append!(posteriors_array_run, posteriors_rep_array)

        end

        iterator_subspaces = [
            [subspace_ind, posteriors_array_run[subspace_ind], algorithm.sampler] for subspace_ind in Base.OneTo(length(posteriors_array_run))]
        samples_subspaces_run = pmap(inp -> sample_subspace(inp...), iterator_subspaces)

        unconv_mask = [samples_subspace.isvalid for samples_subspace in samples_subspaces_run]
        unconv_ind = findall(x->x==false, unconv_mask)
        rep_sspace = !isempty(unconv_ind)

        rec_level += 1
        if rep_sspace && (rec_level > algorithm.nmax_resampling)
            @warn "Convergence is not reached. Try to increase the number of resampling cycles."
            append!(samples_subspaces, samples_subspaces_run)
        else
            append!(samples_subspaces, samples_subspaces_run[unconv_mask])
        end

        posteriors_array = posteriors_array_run

    end

    @info "Integrating Subspaces"
    samples_subspaces = pmap(inp -> integrate_subspace(inp, algorithm.integrator), samples_subspaces)

    @info "Combining Samples"
    samples = deepcopy(samples_subspaces[1].samples)
    info = deepcopy(samples_subspaces[1].info)
    # Save indices from different subspaces:
    info.samples_ind[1] = 1:length(samples)
    for subspace in samples_subspaces[2:end]
        start_ind, stop_ind = length(samples)+1, length(samples)+length(subspace.samples)
        subspace.info.samples_ind[1] = start_ind:stop_ind
        append!(samples, subspace.samples)
        append!(info, subspace.info)
    end

    samples_trafo = varshape(shaped_density).(samples)
    samples_notrafo = inv(trafo).(samples_trafo)

    return (
        result = samples_notrafo, result_trafo = samples_trafo, trafo = trafo,
        info = info,
        exp_samples = exploration_samples, part_tree = partition_tree,
        cost_values = cost_values,
        resampled_trees = resampled_trees
    )
end

function sample_subspace(
    space_id::Integer,
    posterior::PosteriorDensity,
    sampling_algorithm::A,
) where {N<:NamedTuple, A<:AbstractSamplingAlgorithm, I<:IntegrationAlgorithm}

    @info "Sampling subspace #$space_id"
    sampling_wc_start = Dates.Time(Dates.now())
    sampling_cpu_time = CPUTime.@CPUelapsed begin
        sampling_output = bat_sample(posterior, sampling_algorithm)
        samples_subspace = sampling_output.result
        isvalid = sampling_output.isvalid
    end
    sampling_wc_stop = Dates.Time(Dates.now())

    info_subspace = TypedTables.Table(
            sampling_cpu_time = [sampling_cpu_time],
            sampling_wc = [Dates.value(sampling_wc_start):Dates.value(sampling_wc_stop)],
            worker_id = [Distributed.myid()],
            n_threads = [Threads.nthreads()],
            samples_ind = [0:0],
            sum_weights = [sum(samples_subspace.weight)],
        )

    return (samples = samples_subspace, info = info_subspace, isvalid=isvalid)
end

function integrate_subspace(
    sampling_reuslt::N,
    integration_algorithm::I,
) where {N<:NamedTuple, A<:AbstractSamplingAlgorithm, I<:IntegrationAlgorithm}

    samples_subspace = sampling_reuslt.samples

    integration_wc_start = Dates.Time(Dates.now())
    integration_cpu_time = CPUTime.@CPUelapsed begin
        # ToDo: Use samples_subspace_trafo for integration instead of samples_subspace?
        integras_subspace = bat_integrate(samples_subspace, integration_algorithm).result
    end
    integration_wc_stop = Dates.Time(Dates.now())

    samples_subspace_reweighted = DensitySampleVector(
        (
            samples_subspace.v,
            samples_subspace.logd,
            integras_subspace.val .* samples_subspace.weight ./ sum(samples_subspace.weight),
            samples_subspace.info,
            samples_subspace.aux
        )
    )

    info_subspace = TypedTables.Table(
            density_integral = [integras_subspace],
            integration_cpu_time = [integration_cpu_time],
            integration_wc = [Dates.value(integration_wc_start):Dates.value(integration_wc_stop)],
        )

    return (samples = samples_subspace_reweighted, info = TypedTables.Table(info_subspace, sampling_reuslt.info))
end

function convert_to_posterior_resampled(
    posterior::PosteriorDensity,
    partition_tree::SpacePartTree,
    posterior_subspace::PosteriorDensity;
    extend_bounds::Bool=true
)

    # Note: this function is very inefficient
    # To improve it, truncation of truncated densities should be defined, or use space transformaton

    if extend_bounds
        # Exploration samples might not always cover properly tails of the distribution.
        # We will extend boudnaries of the partition tree with original bounds which are:
        lo_bounds = getprior(posterior_subspace).bounds.vol.lo
        hi_bounds = getprior(posterior_subspace).bounds.vol.hi
        extend_tree_bounds!(partition_tree, lo_bounds, hi_bounds)
    end

    #Get flattened rectangular parameter bounds from tree
    subspaces_rect_bounds = get_tree_par_bounds(partition_tree)

    n_params = size(subspaces_rect_bounds[1])[1]

    posterior_array = map(
        x -> begin
            bounds = HyperRectBounds(x[:,1], x[:,2])
            prior_dist = BAT.truncate_density(getprior(posterior), bounds)
            PosteriorDensity(getlikelihood(posterior), prior_dist)
        end,
    subspaces_rect_bounds)

    return posterior_array
end

function convert_to_posterior(posterior::PosteriorDensity, partition_tree::SpacePartTree; extend_bounds::Bool=true)

    if extend_bounds
        # Exploration samples might not always cover properly tails of the distribution.
        # We will extend boudnaries of the partition tree with original bounds which are:
        vol = spatialvolume(var_bounds(posterior))
        lo_bounds = vol.lo
        hi_bounds = vol.hi
        extend_tree_bounds!(partition_tree, lo_bounds, hi_bounds)
    end

    #Get flattened rectangular parameter bounds from tree
    subspaces_rect_bounds = get_tree_par_bounds(partition_tree)

    posterior_array = map(subspaces_rect_bounds) do x
        bounds = StructArray{Interval}((x[:,1], x[:,2]))
        BAT.truncate_density(posterior, bounds)
    end

    return posterior_array
end


function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:PPO},
    ::Val{:Pendulum},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_PPO_Pendulum_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)
    rng = StableRNG(seed)
    inner_env = PendulumEnv(T = Float32, rng = rng)
    action_space = action_space(inner_env)
    low = action_space.low
    high = action_space.high
    ns = length(state(inner_env))

    N_ENV = 8
    UPDATE_FREQ = 2048
    env = MultiThreadEnv([
        PendulumEnv(T = Float32, rng = StableRNG(hash(seed + i))) |>
        ActionTransformedEnv(x -> clamp(x * 2, low, high)) for i in 1:N_ENV
    ])

    init = glorot_uniform(rng)

    agent = Agent(
        policy = PPOPolicy(
            approximator = ActorCritic(
                actor = GaussianNetwork(
                    pre = Chain(
                        Dense(ns, 64, relu; initW = glorot_uniform(rng)),
                        Dense(64, 64, relu; initW = glorot_uniform(rng)),
                    ),
                    μ = Chain(Dense(64, 1, tanh; initW = glorot_uniform(rng)), vec),
                    σ = Chain(Dense(64, 1; initW = glorot_uniform(rng)), vec),
                ),
                critic = Chain(
                    Dense(ns, 64, relu; initW = glorot_uniform(rng)),
                    Dense(64, 64, relu; initW = glorot_uniform(rng)),
                    Dense(64, 1; initW = glorot_uniform(rng)),
                ),
                optimizer = ADAM(3e-4),
            ) |> cpu,
            γ = 0.99f0,
            λ = 0.95f0,
            clip_range = 0.2f0,
            max_grad_norm = 0.5f0,
            n_epochs = 10,
            n_microbatches = 32,
            actor_loss_weight = 1.0f0,
            critic_loss_weight = 0.5f0,
            entropy_loss_weight = 0.00f0,
            dist = Normal,
            rng = rng,
            update_freq = UPDATE_FREQ,
        ),
        trajectory = PPOTrajectory(;
            capacity = UPDATE_FREQ,
            state = Matrix{Float32} => (ns, N_ENV),
            action = Vector{Float32} => (N_ENV,),
            action_log_prob = Vector{Float32} => (N_ENV,),
            reward = Vector{Float32} => (N_ENV,),
            terminal = Vector{Bool} => (N_ENV,),
        ),
    )

    stop_condition = StopAfterStep(50_000)
    total_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    hook = ComposedHook(
        total_reward_per_episode,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info(
                    "training",
                    actor_loss = agent.policy.actor_loss[end, end],
                    critic_loss = agent.policy.critic_loss[end, end],
                    loss = agent.policy.loss[end, end],
                )
                for i in 1:length(env)
                    if is_terminated(env[i])
                        @info "training" reward = total_reward_per_episode.rewards[i][end] log_step_increment =
                            0
                        break
                    end
                end
            end
        end,
    )

    Experiment(agent, env, stop_condition, hook, "# Play Pendulum with PPO")
end


"""
Bridges
`GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.CLOSED_COUNTED_VALUES, T}`
to `GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}`.
"""
struct GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T} <: MOIBC.AbstractBridge
    cons_domain::Vector{MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.Membership}}
    con_gcv::MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}}
end

function MOIBC.bridge_constraint(
    ::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}},
    model,
    f::MOI.VectorOfVariables,
    s::CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.CLOSED_COUNTED_VALUES, T},
) where {T}
    return MOIBC.bridge_constraint(
        GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
        model,
        MOI.VectorAffineFunction{T}(f),
        s,
    )
end

function MOIBC.bridge_constraint(
    ::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}},
    model,
    f::MOI.VectorAffineFunction{T},
    s::CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.CLOSED_COUNTED_VALUES, T},
) where {T}
    f_scalars = MOIU.scalarize(f)
    f_sought = f_scalars[(s.dimension + s.n_values + 1):(s.dimension + 2 * s.n_values)]

    cons_domain = MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.Membership}[
        MOI.add_constraint(
            model,
            MOIU.vectorize(
                [
                    f_scalars[i],
                    f_sought...
                ]
            ),
            CP.Membership(s.n_values),
        )
        for i in 1:s.dimension
    ]

    con_gcv = MOI.add_constraint(
        model,
        f,
        CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}(s.dimension, s.n_values),
    )

    return GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge(cons_domain, con_gcv)
end

function MOI.supports_constraint(
    ::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}},
    ::Union{Type{MOI.VectorOfVariables}, Type{MOI.VectorAffineFunction{T}}},
    ::Type{CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.CLOSED_COUNTED_VALUES, T}},
) where {T}
    return true
end

function MOIB.added_constrained_variable_types(::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}}) where {T}
    return Tuple{Type}[]
end

function MOIB.added_constraint_types(::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}}) where {T}
    return [
        (MOI.VectorAffineFunction{T}, CP.Membership),
        (MOI.VectorAffineFunction{T}, CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}),
    ]
end

function MOI.get(
    b::GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.VectorAffineFunction{T}, CP.Membership,
    },
) where {T}
    return length(b.cons_domain)
end

function MOI.get(
    ::GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.VectorAffineFunction{T}, CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T},
    },
) where {T}
    return 1
end

function MOI.get(
    b::GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{T}, CP.Membership,
    },
) where {T}
    return copy(b.cons_domain)
end

function MOI.get(
    b::GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{T}, CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T},
    },
) where {T}
    return [b.con_gcv]
end


#
# This file is part of module
#   ModiaMedia (ModiaMedia/src/ModiaMedia.jl)
#
using LinearAlgebra

const excludeEnthalpyOfFormation = true                         # If true, enthalpy of formation Hf is not included in specific enthalpy hc
const referenceEnthalpy          = ReferenceEnthalpy_UserDefined   # Choice of reference enthalpy
const h_offset                   = 0.0                          # User defined offset for reference enthalpy, if ReferenceEnthalpy = ReferenceEnthalpy_UserDefined

"""
    medium = SimpleIdealGasMedium(; mediumName     = nothing,
                                    reference_p    = 101325,
                                    reference_T    = 298.15,
                                    p_default      = 101325,
                                    T_default      = 293.15,
                                    fluidConstants = nothing,
                                    data           = nothing)

Generate a `SimpleIdealGasMedium <: PureSubstance` medium object.
"""

# what is the deal with the baseProperties symbol?
struct IdealMoistAirMedium <: CondensingGases
    infos::FluidInfos
    fluidConstants::SVector{2,IdealGasFluidConstants}
    fluidLimits::FluidLimits
    data::SVector{2,SingleGasNasaData}

    function IdealMoistAirMedium(; mediumName="MoistAir",
                                    substanceNames=["water", "air"],
                                    reference_p=101325,
                                    reference_T=298.15,
                                    reference_X=[0.01, 0.99],
                                    p_default=101325,
                                    T_default=293.15,
                                    X_default=[0.01, 0.99],
                                    fluidConstants=[mediumDict["H2O"].fluidConstants[1],
                                                    mediumDict["Air"].fluidConstants[1]],
                                    data=[mediumDict["H2O"].data,
                                          mediumDict["Air"].data])

        infos = FluidInfos(mediumName           = mediumName,
                           substanceNames       = substanceNames,
                           extraPropertiesNames = fill("",0),
                           ThermoStates         = IndependentVariables_pTX,
                           baseProperties       = :BaseProperties_SingleGasNasa,
                           singleState          = false,
                           reducedX             = true,
                           fixedX               = false,
                           reference_p          = reference_p,
                           reference_T          = reference_T,
                           reference_X          = reference_X,
                           p_default            = p_default,
                           T_default            = T_default,
                           X_default            = X_default,
                           h_default            = h_pTX(SVector{2,SingleGasNasaData}(data),p_default,T_default,X_default))

        # these limits should be fixed after the specific enthalpy can be calculated below 273.15K; the high end also has problems as well.
        fluidLimits = FluidLimits(TMIN = max(mediumDict["H2O"].fluidLimits.TMIN, mediumDict["Air"].fluidLimits.TMIN),
                                  TMAX = min(mediumDict["H2O"].fluidLimits.TMAX, mediumDict["Air"].fluidLimits.TMAX),
                                  HMIN = h_pTX(SVector{2,SingleGasNasaData}(data),p_default,280.0,X_default),
                                  HMAX = h_pTX(SVector{2,SingleGasNasaData}(data),p_default,600.0,X_default))

        new(infos, SVector{2,IdealGasFluidConstants}(fluidConstants), fluidLimits, SVector{2,SingleGasNasaData}(data))
    end
end

function saturationPressure(TSat)
    # this is only valid to 273.15 K (need to add the ice saturation pressure and spline both together)
    TCritical=647.096
    pCritical=22.064e6
    r1=(1 - TSat/TCritical)

    a = [-7.85951783,1.84408259,-11.7866497,22.6807411,-15.9618719,1.80122502]
    n = [1.0,1.5,3.0,3.5,4.0,7.5]
    pSat = exp(((a[1]*r1^n[1] + a[2]*r1^n[2] + a[3]*r1^n[3] + a[4]*r1^n[4]
     + a[5]*r1^n[5] + a[6]*r1^n[6])*TCritical)/TSat)*pCritical

     return pSat
end

# enthalpy of water only for liquid condition; need to splice with the enthalpy of ice to get number valid over whole range
enthalpyOfWater(T) = 4200.0

function h_pTX(data::SVector{2,SingleGasNasaData},p,T,X)
    # steam.MM/dryAir.MM
    k_mair = data[1].MM/data[2].MM

    p_steam_sat = saturationPressure(T)
    X_sat = min(p_steam_sat*k_mair/max(100*eps(), p-p_steam_sat)*(1-X[1]), 1.0)
    X_liquid = max(X[1]-X_sat, 0.0);
    X_steam = X[1]-X_liquid;
    X_air = 1-X[1];

    h = dot([h_T(data[1], T, excludeEnthalpyOfFormation, referenceEnthalpy, (46479.819 + 2501014.5));
             h_T(data[2], T, excludeEnthalpyOfFormation, referenceEnthalpy, 25104.684)],
            [X_steam; X_air])
        + enthalpyOfWater(T)*X_liquid

    return h
end

function u_pTX(data::SVector{2,SingleGasNasaData},p,T,X)
    # steam.MM/dryAir.MM
    k_mair = data[1].MM/data[2].MM

    p_steam_sat = saturationPressure(T)
    X_sat = min(p_steam_sat*k_mair/max(100*eps(), p-p_steam_sat)*(1-X[1]), 1.0)
    X_liquid = max(X[1]-X_sat, 0.0);
    X_steam = X[1]-X_liquid;
    X_air = 1-X[1];
    R_gas = data[2].R*X_air/(1-X_liquid) + data[1].R*X_steam/(1-X_liquid);

    u = h_pTX(data,p,T,X) - R_gas*T
    return u
end

p_dTX(data::SVector{2,SingleGasNasaData},d,T,X) = d*dot([data[1].R; data[2].R], X)*T

setState_pTX(m::IdealMoistAirMedium,p,T,X) = ThermodynamicState_pTX(p,T,X)
setState_dTX(m::IdealMoistAirMedium,d,T,X) = ThermodynamicState_pTX(p_dTX(m.data,d,T,X),T,X)

pressure(               m::IdealMoistAirMedium, state::ThermodynamicState_pTX)::Float64 = state.p
temperature(            m::IdealMoistAirMedium, state::ThermodynamicState_pTX)::Float64 = state.T
gasConstant(            m::IdealMoistAirMedium, state::ThermodynamicState_pTX)::Float64 = m.data[2].R*(1-state.X[1]) + m.data[1].R*state.X[1]
density(                m::IdealMoistAirMedium, state::ThermodynamicState_pTX)::Float64 = state.p/(gasConstant(m,state)*state.T)
specificEnthalpy(       m::IdealMoistAirMedium, state::ThermodynamicState_pTX)::Float64 = h_pTX(m.data,state.p,state.T,state.X)
specificInternalEnergy( m::IdealMoistAirMedium, state::ThermodynamicState_pTX)::Float64 = u_pTX(m.data,state.p,state.T,state.X)


#####
##### `AbstractDifferential`
#####

"""
The subtypes of `AbstractDifferential` define a custom \"algebra\" for chain
rule evaluation that attempts to factor various features like complex derivative
support, broadcast fusion, zero-elision, etc. into nicely separated parts.

All subtypes of `AbstractDifferential` implement the following operations:

`add(a, b)`: linearly combine differential `a` and differential `b`

`mul(a, b)`: multiply the differential `a` by the differential `b`

`Base.conj(x)`: complex conjugate of the differential `x`

`extern(x)`: convert `x` into an appropriate non-`AbstractDifferential` type for
use outside of `ChainContext`.

Valid arguments to these operations are `T` where `T<:AbstractDifferential`, or
where `T` has proper `+` and `*` implementations.

Additionally, all subtypes of `AbstractDifferential` support `Base.iterate` and
`Base.Broadcast.broadcastable(x)`.
"""
abstract type AbstractDifferential end

"""
    extern(x)

Return `x` converted to an appropriate non-`AbstractDifferential` type, for use
with external packages that might not handle `AbstractDifferential` types.

Note that this function may return an alias (not necessarily a copy) to data
wrapped by `x`, such that mutating `extern(x)` might mutate `x` itself.
"""
@inline extern(x) = x

@inline Base.conj(x::AbstractDifferential) = x

#=
This `AbstractDifferential` algebra has a monad-y "fallthrough" implementation;
each step handles an element of the algebra before dispatching to the next step.
This way, we don't need to implement promotion/conversion rules between subtypes
of `AbstractDifferential` to resolve potential ambiguities.
=#

const PRECEDENCE_LIST = [:wirtinger, :casted, :zero, :dne, :one, :thunk, :fallback]

global defs = Expr(:block)

let previous_add_name = :add, previous_mul_name = :mul
    for name in PRECEDENCE_LIST
        next_add_name = Symbol(string(:add_, name))
        next_mul_name = Symbol(string(:mul_, name))
        push!(defs.args, quote
            @inline $(previous_add_name)(a, b) = $(next_add_name)(a, b)
            @inline $(previous_mul_name)(a, b) = $(next_mul_name)(a, b)
        end)
        previous_add_name = next_add_name
        previous_mul_name = next_mul_name
    end
end

eval(defs)

@inline add_fallback(a, b) = a + b

@inline mul_fallback(a, b) = a * b

@inline add(x) = x

@inline mul(x) = x

#####
##### `Wirtinger`
#####

"""
    Wirtinger(primal::Union{Number,AbstractDifferential},
              conjugate::Union{Number,AbstractDifferential})

Returns a `Wirtinger` instance representing the complex differential:

```
df = ∂f/∂z * dz + ∂f/∂z̄ * dz̄
```

where `primal` corresponds to `∂f/∂z * dz` and `conjugate` corresponds to `∂f/∂z̄ * dz̄`.

The two fields of the returned instance can be accessed generically via the
[`wirtinger_primal`](@ref) and [`wirtinger_conjugate`](@ref) methods.
"""
struct Wirtinger{P,C} <: AbstractDifferential
    primal::P
    conjugate::C
    function Wirtinger(primal::Union{Number,AbstractDifferential},
                       conjugate::Union{Number,AbstractDifferential})
        return new{typeof(primal),typeof(conjugate)}(primal, conjugate)
    end
    function Wirtinger(primal, conjugate)
        error("`Wirtinger` only supports elements of type <: Union{Number,AbstractDifferential} for now")
    end
end

wirtinger_primal(x::Wirtinger) = x.primal
wirtinger_primal(x) = x

wirtinger_conjugate(x::Wirtinger) = x.conjugate
wirtinger_conjugate(::Any) = Zero()

extern(x::Wirtinger) = error("`Wirtinger` cannot be converted into an external type.")

Base.Broadcast.broadcastable(w::Wirtinger) = Wirtinger(broadcastable(w.primal),
                                                       broadcastable(w.conjugate))

Base.iterate(x::Wirtinger) = (x, nothing)
Base.iterate(::Wirtinger, ::Any) = nothing

Base.conj(x::Wirtinger) = error("`conj(::Wirtinger)` not yet defined")

function add_wirtinger(a::Wirtinger, b::Wirtinger)
    return Wirtinger(add(a.primal, b.primal), add(a.conjugate, b.conjugate))
end

add_wirtinger(a::Wirtinger, b) = add(a, Wirtinger(b, Zero()))
add_wirtinger(a, b::Wirtinger) = add(Wirtinger(a, Zero()), b)

function mul_wirtinger(a::Wirtinger, b::Wirtinger)
    error("""
          cannot multiply two Wirtinger objects; this error likely means a
          `WirtingerRule` was inappropriately defined somewhere. Multiplication
          of two Wirtinger objects is not defined because chain rule application
          often expands into a non-commutative operation in the Wirtinger
          calculus. To put it another way: simply given two Wirtinger objects
          and no other information, we can't know "locally" which components to
          conjugate in order to implement the chain rule. We could pick a
          convention; for example, we could define `a::Wirtinger * b::Wirtinger`
          such that we assume the chain rule application is of the form `f_a ∘ f_b`
          instead of `f_b ∘ f_a`. However, picking such a convention is likely to
          lead to silently incorrect derivatives due to commutativity assumptions
          in downstream generic code that deals with the reals. Thus, ChainRules
          makes this operation an error instead.
          """)
end

mul_wirtinger(a::Wirtinger, b) = Wirtinger(mul(a.primal, b), mul(a.conjugate, b))
mul_wirtinger(a, b::Wirtinger) = Wirtinger(mul(a, b.primal), mul(a, b.conjugate))

#####
##### `Casted`
#####

"""
TODO
"""
struct Casted{V} <: AbstractDifferential
    value::V
end

cast(x) = Casted(x)
cast(f, args...) = Casted(broadcasted(f, args...))

extern(x::Casted) = materialize(broadcasted(extern, x.value))

Base.Broadcast.broadcastable(x::Casted) = x.value

Base.iterate(x::Casted) = iterate(x.value)
Base.iterate(x::Casted, state) = iterate(x.value, state)

Base.conj(x::Casted) = cast(conj, x.value)

add_casted(a::Casted, b::Casted) = Casted(broadcasted(add, a.value, b.value))
add_casted(a::Casted, b) = Casted(broadcasted(add, a.value, b))
add_casted(a, b::Casted) = Casted(broadcasted(add, a, b.value))

mul_casted(a::Casted, b::Casted) = Casted(broadcasted(mul, a.value, b.value))
mul_casted(a::Casted, b) = Casted(broadcasted(mul, a.value, b))
mul_casted(a, b::Casted) = Casted(broadcasted(mul, a, b.value))

#####
##### `Zero`
#####

"""
TODO
"""
struct Zero <: AbstractDifferential end

extern(x::Zero) = false

Base.Broadcast.broadcastable(::Zero) = Ref(Zero())

Base.iterate(x::Zero) = (x, nothing)
Base.iterate(::Zero, ::Any) = nothing

add_zero(::Zero, ::Zero) = Zero()
add_zero(::Zero, b) = b
add_zero(a, ::Zero) = a

mul_zero(::Zero, ::Zero) = Zero()
mul_zero(::Zero, ::Any) = Zero()
mul_zero(::Any, ::Zero) = Zero()

#####
##### `DNE`
#####

"""
TODO
"""
struct DNE <: AbstractDifferential end

extern(x::DNE) = error("`DNE` cannot be converted into an external type.")

Base.Broadcast.broadcastable(::DNE) = Ref(DNE())

Base.iterate(x::DNE) = (x, nothing)
Base.iterate(::DNE, ::Any) = nothing

add_dne(::DNE, ::DNE) = DNE()
add_dne(::DNE, b) = b
add_dne(a, ::DNE) = a

mul_dne(::DNE, ::DNE) = DNE()
mul_dne(::DNE, ::Any) = DNE()
mul_dne(::Any, ::DNE) = DNE()

#####
##### `One`
#####

"""
TODO
"""
struct One <: AbstractDifferential end

extern(x::One) = true

Base.Broadcast.broadcastable(::One) = Ref(One())

Base.iterate(x::One) = (x, nothing)
Base.iterate(::One, ::Any) = nothing

add_one(a::One, b::One) = add(extern(a), extern(b))
add_one(a::One, b) = add(extern(a), b)
add_one(a, b::One) = add(a, extern(b))

mul_one(::One, ::One) = One()
mul_one(::One, b) = b
mul_one(a, ::One) = a

#####
##### `Thunk`
#####

"""
TODO
"""
struct Thunk{F} <: AbstractDifferential
    f::F
end

macro thunk(body)
    return :(Thunk(() -> $(esc(body))))
end

@inline extern(x::Thunk{F}) where {F} = x.f()

Base.Broadcast.broadcastable(x::Thunk) = broadcastable(extern(x))

@inline function Base.iterate(x::Thunk)
    externed = extern(x)
    element, state = iterate(externed)
    return element, (externed, state)
end

@inline function Base.iterate(::Thunk, (externed, state))
    element, new_state = iterate(externed, state)
    return element, (externed, new_state)
end

Base.conj(x::Thunk) = @thunk(conj(extern(x)))

add_thunk(a::Thunk, b::Thunk) = add(extern(a), extern(b))
add_thunk(a::Thunk, b) = add(extern(a), b)
add_thunk(a, b::Thunk) = add(a, extern(b))

mul_thunk(a::Thunk, b::Thunk) = mul(extern(a), extern(b))
mul_thunk(a::Thunk, b) = mul(extern(a), b)
mul_thunk(a, b::Thunk) = mul(a, extern(b))

#####
##### misc.
#####

"""
    Wirtinger(primal::Real, conjugate::Real)

Return `add(primal, conjugate)`.

Actually implementing the Wirtinger calculus generally requires that the
summed terms of the Wirtinger differential (`∂f/∂z * dz` and `∂f/∂z̄ * dz̄`) be
stored individually. However, if both of these terms are real-valued, then
downstream Wirtinger propagation mechanisms resolve to the same mechanisms as
real-valued calculus, so that the terms' sum can be eagerly computed and
propagated without requiring a special `Wirtinger` representation

This method primarily exists as an optimization.
"""
function Wirtinger(primal::Union{Real,DNE,Zero,One},
                   conjugate::Union{Real,DNE,Zero,One})
    return add(primal, conjugate)
end


using Molly
using DelimitedFiles

data_dir = normpath(dirname(pathof(Molly)), "..", "data")
ff_dir = joinpath(data_dir, "force_fields")
openmm_dir = joinpath(data_dir, "openmm_6mrr")

ff = OpenMMForceField(joinpath.(ff_dir, ["ff99SBildn.xml", "tip3p_standard.xml", "his.xml"])...)
velocities = SVector{3}.(eachrow(readdlm(joinpath(openmm_dir, "velocities_300K.txt"))))u"nm * ps^-1"
s = System(joinpath(data_dir, "6mrr_equil.pdb"), ff; velocities=velocities)
simulator = VelocityVerlet(dt=0.0005u"ps", coupling=AndersenThermostat(300.0u"K", 1.0u"ps"))
n_steps = 500

parallel = false
simulate!(s, simulator, 5; parallel=parallel)
@time simulate!(s, simulator, n_steps; parallel=parallel)

parallel = true
simulate!(s, simulator, 5; parallel=parallel)
@time simulate!(s, simulator, n_steps; parallel=parallel)


julia> # Load DataFrame package

julia> require("DataFrames")

julia> using DataFrames

julia>

julia> # Load a very simple lm() function

julia> require(Pkg.dir("DataFrames", "demo", "lm.jl"))

julia>

julia> # Load a CSV file into a DataFrame

julia> df = read_table(Pkg.dir("DataFrames", "demo", "toy_example.csv"))
6x3 DataFrame:
          A       B   C
[1,]    2.5   "One" 3.0
[2,]    3.6   "One" 5.0
[3,]    3.5   "Two" 3.0
[4,]    4.5   "Two" 5.0
[5,]    4.5 "Three" 3.0
[6,]    5.5 "Three" 5.0

julia>

julia> # Run a simple linear model that predicts A using B and C

julia> lm_fit = lm(:(A ~ B + C), df)
ArgumentError("float64(String): invalid number format")
 in float64 at string.jl:1209
 in float at string.jl:1222
 in model_matrix at /Users/johnmyleswhite/.julia/DataFrames/src/formula.jl:141
 in lm at /Users/johnmyleswhite/.julia/DataFrames/demo/lm.jl:66

julia>

julia> # Print out a summary of the results

julia> print(lm_fit)
lm_fit not defined

julia>

julia> #

julia> # Behind The Scenes

julia> #

julia>

julia> # Generate a Formula object

julia> f = Formula(:(A ~ B + C))
Formula([:A],[:(+(B,C))])

julia>

julia> # Generate a ModelFrame object

julia> mf = model_frame(f, df)
ModelFrame(6x3 DataFrame:
          A       B   C
[1,]    2.5   "One" 3.0
[2,]    3.6   "One" 5.0
[3,]    3.5   "Two" 3.0
[4,]    4.5   "Two" 5.0
[5,]    4.5 "Three" 3.0
[6,]    5.5 "Three" 5.0
,[1],Formula([:A],[:(+(B,C))]))

julia>

julia> # Generate a ModelMatrix object

julia> mm = model_matrix(mf)
ArgumentError("float64(String): invalid number format")
 in float64 at string.jl:1209
 in float at string.jl:1222
 in model_matrix at /Users/johnmyleswhite/.julia/DataFrames/src/formula.jl:141

julia>


export freqCont, freqContExtendedSources, freqContExtendedSourcesSS, freqCont;
using JLD
using Multigrid.ParallelJuliaSolver
using Statistics

export FreqContParam, getFreqContParams
"""
	FreqContParam

	parameters for the frequency continuation method

	Input:
		mc    		- current model
		itersNum	- number of GN iters for each frequency continuation step
		originalSources	- original sources before sources extension
		nrcv		- number of receivers
		pInv		- Inverse param
		pMis 		- misfit params (remote)
		windowSize  - How many frequencies to treat at once at the most.
		resultsFilename - a filename for saving the intermediate results according to the GN and continuation (FC) iterations (done in dumpFun())
		dumpFun     - a function for plotting, saving and doing all the things with the intermidiate results.
		Iact		- active set locations
		mback		- active set model values
		Z1			- first part of the extended sources initial guess
		simSrcDim	- dimensions of the simultaneous sources
		alpha1		- Z1 penalty coeficient
		alpha2Orig	- Z2 penalty coeficient
		stepReg		- regularization coeficient on the step size
		mode        - either "1stInit" or anything else.
					  1stInit will use the first group of misfits as an initialization and will not include it together with the next group of misfits.
		startFrom   - a start index for the continuation. Usefull when our run broke in the middle of the iterations.
		endAt		- index of the final frequency for the continuation.Usefull when we want to use only low frequencies.
		cycle       - just for identifying different runs when saving the results.
		method      - "projGN" or "barrierGN"
		FWImethod	- "FWI" or "FWI_ES"
		updateMref	- wheter to update the mref after each GN iteration, relevant for the ES case.

"""
mutable struct FreqContParam
	mc						:: Vector{Float64}
	itersNum				:: Int64
	originalSources			:: SparseMatrixCSC
	nrcv					:: Int64
	pInv					:: InverseParam
	pMis					:: Array{RemoteChannel}
	windowSize				:: Int64
	resultsFilename			:: String
	dumpFun					:: Function
	Iact					:: SparseMatrixCSC
	mback					:: Vector{Float64}
	Z1						:: Union{Int64, Array{ComplexF64, 2}}
	simSrcDim				:: Int64
	alpha1					:: Float64
	alpha2Orig				:: Float64
	stepReg					:: Float64
	mode					:: String
	startFrom				:: Int64
	endAt					:: Int64
	cycle					:: Int64
	method					:: String
	FWImethod				:: String
	updateMref				:: Bool
end

function getFreqContParams(mc,itersNum::Int64,
		originalSources::SparseMatrixCSC,nrcv, pInv::InverseParam,
		pMis::Array{RemoteChannel}, windowSize::Int64, resultsFilename::String,
		dumpFun::Function,Iact,mback;Z1=0,simSrcDim=1,
		alpha1=0.0, alpha2Orig=0.0, stepReg=0.0, mode::String="", startFrom::Int64 = 1,
		endAt=length(pMis),cycle::Int64=0, method::String="projGN",
		FWImethod::String="FWI",updateMref=false)
		return FreqContParam(mc, itersNum, originalSources,
				nrcv, pInv, pMis, windowSize, resultsFilename, dumpFun,
				Iact, mback, Z1, simSrcDim, alpha1, alpha2Orig, stepReg,
				mode, startFrom, endAt, cycle, method, FWImethod, updateMref)
end

function calculateZ2(m,misfitCalc::Function, p::Integer, nsrc::Integer, nfreq::Integer,
	nrec::Integer, nwork::Integer,
	numOfCurrentProblems::Integer, mergedWd::Array, mergedRc::Array, pMis::Array,
	pMisCurrent::Array{MisfitParam}, Z1::Matrix, alpha::Float64)

	rhs = zeros(ComplexF64, (p, nsrc));
	lhs = zeros(ComplexF64, (p,p));

	HinvPsZ1 = computeHinvTRecX(pMis, Z1, m, 0);

	function calculateLhs(meanWd,Z1,HinvP)
		return meanWd .* HinvP' * HinvP;
	end

	function calculateRhs(meanWd,HinvP,mergedRc)
		return meanWd .* HinvP' * mergedRc;
	end

	lhsides = Array{Array{ComplexF64}}(undef,nfreq)
	rhsides = Array{Array{ComplexF64}}(undef,nfreq)

	@sync begin
	 		for k=1:nfreq
		@async begin
			meanWd = mean(mergedWd[k])
			meanWdC = (real(meanWd).^2)
				lhsides[k] = remotecall_fetch(calculateLhs, k % nworkers() + 1,meanWdC,Z1,HinvPsZ1[k]);
				rhsides[k] = remotecall_fetch(calculateRhs, k % nworkers() + 1,meanWdC,HinvPsZ1[k], mergedRc[k]);
	 		end
	 	end
	end

	lhs = sum(lhsides)
	rhs = sum(rhsides)
	lhs += alpha * I;

	return lhs\rhs;
end

function MyPCG(A::Function,b::Array,x::Array,M::Array,numiter)
r = b-A(x);
z = M.*r;
p = copy(z);
norms = [norm(r)];

# use 20 iters if starting residual is high
if norms[1] > 100
	numiter=20;
end
for k=1:numiter
	Ap = A(p);
	alpha = real(dot(z,r)/dot(p,Ap));
	x .+= alpha*p;
	beta = real(dot(z,r));
	r .-= alpha*Ap
	z = M.*r
	nn = norm(r);
	if (nn < 0.1 * norms[1] || nn < 1e-2)
		break
	end
	norms = [norms; nn];
	beta = real(dot(z,r)) / beta;
	p = z + beta*p;
end
return x,norms
end

function misfitCalc(m, Z1,Z2,mergedWd,mergedRc,nfreq,alpha1,alpha2,pMis)
	## HERE WE  NEED TO MAKE SURE THAT Wd is equal in its real and imaginary parts.
	sum = 0.0;
	HinvPs = computeHinvTRecX(pMis, Z1, m, 0);
	for i = 1:nfreq
		res, = SSDFun(HinvPs[i] * Z2,mergedRc[i],mean(mergedWd[i]) .* ones(size(mergedRc[i])));
		sum += res;
	end
	return sum;
end

function objectiveCalc(Z1,Z2,misfit,alpha1,alpha2)
	return misfit + 0.5*alpha1 * norm(Z1)^2 + 0.5*alpha2 * norm(Z2)^2;
end

function MultOp(HPinv, R, Z2)
	return (HPinv' * R) * Z2
end

function MultOpT(HPinv, R, Z2)
	return HPinv * (R * Z2')
end

function MultAll(m, avgWds, pMis, R, Z1, Z2, alpha, stepReg)
	sum = zeros(ComplexF64, size(R))
	inner = computeHinvTRecX(pMis, R,m, 0);

	function calculateInner(avgWd, inner, Z2)
		return avgWd^2 .* inner * (Z2 * Z2');
	end
	@sync begin
	 		for k=1:length(avgWds)
		@async begin
	 			inner[k] = remotecall_fetch(calculateInner,
	 				k % nworkers() + 1,avgWds[k],inner[k],Z2);
	 		end
	 	end
	end

	outer = computeHinvTRecXarr(pMis, inner,m, 1);
	for i = 1:length(avgWds)
		sum += outer[i]
	end
	eps = 1e-2
	return sum + (alpha./(abs.(Z1) .+ eps * maximum(abs.(Z1)))).*R + stepReg*R;
end

function calculateZ1(m, misfitCalc::Function, nfreq::Integer, mergedWd::Array, mergedRc::Array, pMis::Array, Z1::Matrix,Z2::Matrix, alpha1::Float64,stepReg::Float64)
	rhsElement = computeHinvTRecXarr(pMis, mergedRc, m, 1);
	rhsides = Array{Array{ComplexF64}}(undef,nfreq)

	function calculateRhs(meanWd,Z2,rhsElement)
		return meanWd.*  (rhsElement * Z2');
	end

	@sync begin
	 		for k=1:nfreq
		@async begin
			meanWd = mean(mergedWd[k])
			meanWdC = (real(meanWd).^2)
			rhsides[k] = remotecall_fetch(calculateRhs, k % nworkers() + 1,meanWdC,Z2, rhsElement[k]);
	 		end
	 	end
	end

	rhs = sum(rhsides) + (stepReg) .* Z1

	OP = x-> MultAll(m, mean.(real.(mergedWd)), pMis, x, Z1, Z2, alpha1, stepReg);
	eps = 10.0
	M = (abs.(Z1) .+ eps * maximum(abs.(Z1)));
	Z1, = MyPCG(OP,rhs,Z1,M,5);
	return Z1;
end

function standardGNrun(method, mc, simSrcDim, originalSources, pInv, pMisTemp,
		HIS, resultsFilename, cycle, freqIdx, dumpFun)
pMisTempFetched = map(fetch, pMisTemp);
dobs = map(x -> x.dobs,pMisTempFetched)
wd = map(x-> x.Wd,pMisTempFetched)
nsrc = size(originalSources,2);

if simSrcDim==1
	TEmat = Matrix(1.0I,nsrc,nsrc)
else
	TEmat = rand([-1,1],(nsrc,simSrcDim));
end

if simSrcDim > 1
	reducedSources = originalSources * TEmat;
	reducedDobs = map(x-> x*TEmat, dobs);
	sizeWD = size(reducedDobs[1])
	reducedWd = map(x-> mean(x)/sqrt(simSrcDim) * ones(sizeWD), wd);

	pMisTemp = setSourcesSame(pMisTemp,reducedSources);
	pMisTemp = setDobs(pMisTemp,reducedDobs);
	pMisTemp = setWd(pMisTemp,reducedWd);
end

if resultsFilename == ""
	filename = "";
	hisMatFileName = "";
else
	Temp = splitext(resultsFilename);
	if cycle==0
		filename = string(Temp[1],"_FC",freqIdx,"_GN",Temp[2]);
		hisMatFileName  = string(Temp[1],"_FC",freqIdx);
	else
		filename = string(Temp[1],"_Cyc",cycle,"_FC",freqIdx,"_GN",Temp[2]);
		hisMatFileName  =  string(Temp[1],"_Cyc",cycle,"_FC",freqIdx);
	end
end

# Here we set a dump function for GN for this iteracion of FC
function dumpGN(mc,Dc,iter,pInv,PF)
	dumpFun(mc,Dc,iter,pInv,PF,filename);
end

if method == "projGN"
	mc,Dc,flag,His = projGNCG(mc,pInv,pMisTemp,dumpResults = dumpGN);
elseif method == "barrierGN"
	mc,Dc,flag,His = barrierGNCG(mc,pInv,pMisTemp,rho=1.0,dumpResults = dumpGN);
end

if simSrcDim > 1
	# set full data back
	pMisTemp = setSourcesSame(pMisTemp,originalSources);
	pMisTemp = setDobs(pMisTemp,dobs)
	pMisTemp = setWd(pMisTemp,wd)
end

if hisMatFileName != ""
	file = matopen(string(hisMatFileName,"_HisGN.mat"), "w");
	### PUT THE CONTENT OF "HIS" INTO THE MAT FILE ###
	His.Dc = [];
	write(file,"His",His);
	close(file);
end
His.Dc = []
push!(HIS,His)

clear!(pMisTemp);

return mc,Dc,flag,HIS
end

function ExtendedSourcesGNRun(method, mc, simSrcDim, nrcv, itersNum, originalSources,
		Z1, alpha1, alpha2, stepReg, pInv, pMisTemp, HIS, resultsFilename,
		cycle, freqIdx, dumpFun)
numOfCurrentProblems = length(pMisTemp);
N_nodes = prod(pInv.MInv.n .+ 1);
nsrc = size(originalSources, 2);
p = size(Z1,2);
nwork = nworkers()
nsrc = size(originalSources, 2);
nfreq = length(pMisTemp);
wheres = map(x -> x.where, pMisTemp);
pMisTempFetched = map(fetch, pMisTemp);
mergedDobs = map(x -> x.dobs,pMisTempFetched)
mergedWd = map(x-> x.Wd,pMisTempFetched)

##########################
# Initialize outputs	 #
##########################
Dc = 0;
flag = -1;
HIS = [];
Z2 = 0;
##########################

FafterGN = 0.0;
F_zero_prev = 0.0;
mc_prev = convert(Array{Float16},mc);
for j = 1:itersNum
	println("============================== New ALM Iter ======================================");
	flush(Base.stdout)

	if simSrcDim==1
		TEmat = Matrix(1.0I,nsrc,nsrc)
	else
		TEmat = rand([-1,1],(nsrc,simSrcDim));
	end

	reducedSources = originalSources * TEmat;
	reducedDobs = map(x-> x*TEmat, mergedDobs);
	sizeWD = size(reducedDobs[1])
	reducedWd = map(x-> mean(x)/sqrt(simSrcDim) * ones(sizeWD), mergedWd);

	pMisTemp = setSourcesSame(pMisTemp,reducedSources);
	pMisTemp = setDobs(pMisTemp,reducedDobs);
	pMisTemp = setWd(pMisTemp,reducedWd);

	t1 = time_ns();
	Dc,F_zero, = computeMisfit(mc,pMisTemp,false);
	e1 = time_ns();
	println("Computed Misfit with orig sources : ",F_zero, " [Time: ",(e1 - t1)/1.0e9," sec]");

	mc_prev = convert(Array{Float16},mc);
	pMisHps = pMisTemp

	# Rc = Dc-Dobs, where Dc is the clean (wrt Z1,Z2) simulated data
	mergedRc = Array{Array{ComplexF64}}(undef,nfreq);
	for f = 1:nfreq
		mergedRc[f] = reducedDobs[f] .- fetch(Dc[f]);
	end

	println("Misfit with Zero Z2 (our computation): ", misfitCalc(mc,zeros(ComplexF64, (N_nodes, p)),zeros(ComplexF64, (p, size(TEmat,2))),mergedWd ./ sqrt(simSrcDim),mergedRc,nfreq,alpha1,alpha2, pMisHps))
	mergedRcReduced = mergedRc;

	pMisTempFetched = map(fetch, pMisTemp)
	t1 = time_ns();
	srcNum = simSrcDim == 1 ? nsrc : simSrcDim;
	Z2 = calculateZ2(mc,misfitCalc, p, srcNum, nfreq, nrcv,nwork, numOfCurrentProblems, mergedWd ./ sqrt(simSrcDim), mergedRcReduced, pMisHps, pMisTempFetched, Z1, alpha2);
	e1 = time_ns();

	print("After First Z2 update: ");
	mis = misfitCalc(mc,Z1,Z2,mergedWd ./ sqrt(simSrcDim),mergedRcReduced,nfreq,alpha1,alpha2, pMisHps);
	obj = objectiveCalc(Z1,Z2,mis,alpha1,alpha2);
	initialMis = mis
	initialObj = obj

	println("mis: ",mis,", obj: ",obj,", norm Z2 = ", norm(Z2)^2," norm Z1: ", norm(Z1)^2, ", [Time: ",(e1 - t1)/1.0e9," sec]")

	innerIters = j == 1 ? 5 : 1; # do five in first outer iteration only for better performance
	for iters = 1:innerIters
		###################################################
		#### COMPUTING Z1:
		###################################################
		t1 = time_ns();
		Z1 = calculateZ1(mc,misfitCalc, nfreq, mergedWd ./ sqrt(simSrcDim), mergedRcReduced, pMisHps, Z1, Z2, alpha1, stepReg);
		e1 = time_ns();
		mis = misfitCalc(mc,Z1,Z2,mergedWd ./ sqrt(simSrcDim) ,mergedRcReduced,nfreq,alpha1,alpha2, pMisHps);
		obj = objectiveCalc(Z1,Z2,mis,alpha1,alpha2);
		println("After Z1: mis: ",mis,", obj: ",obj,", norm Z2 = ", norm(Z2)^2," norm Z1: ", norm(Z1)^2, ", [Time: ",(e1 - t1)/1.0e9," sec]")

		###################################################
		#### COMPUTING Z2:
		###################################################
		t1 = time_ns();
		Z2 = calculateZ2(mc,misfitCalc, p, srcNum, nfreq, nrcv,nwork, numOfCurrentProblems, mergedWd ./ sqrt(simSrcDim), mergedRcReduced, pMisHps, pMisTempFetched, Z1, alpha2);
		e1 = time_ns();
		mis = misfitCalc(mc,Z1,Z2,mergedWd ./ sqrt(simSrcDim) ,mergedRcReduced,nfreq,alpha1,alpha2, pMisHps);
		obj = objectiveCalc(Z1,Z2,mis,alpha1,alpha2);
		println("After Z2: mis: ",mis,", obj: ",obj,", norm Z2 = ", norm(Z2)^2," norm Z1: ", norm(Z1)^2, ", [Time: ",(e1 - t1)/1.0e9," sec]")

		# allow extended source to reduce misfit by up to 20%
		if mis < 0.8 * initialMis
			break
		end
	end

	if mis / F_zero > 0.95
			alpha1 = alpha1 / 1.5;
			alpha2 = alpha2 / 1.5;
			println("Ratio mis/F_zero is: ",mis/F_zero,", hence decreasing alphas by 1.5: ",alpha1,",",alpha2);
	end

	if mis / F_zero < 0.3
			alpha1 = alpha1*1.5;
			alpha2 = alpha2*1.5;
			println("Ratio mis/F_zero is: ",mis/F_zero,", hence increasing alpha1 by 1.5: ",alpha1,",",alpha2);
	end

	newSources = originalSources * TEmat + Z1 * Z2;

	pMisTemp = setSourcesSame(pMisTemp,newSources);
	pMisTemp = setDobs(pMisTemp,reducedDobs)
	pMisTemp = setWd(pMisTemp,reducedWd)

	if resultsFilename == ""
		filename = "";
		hisMatFileName = "";
	else
		Temp = splitext(resultsFilename);
		filename = string(Temp[1],"_Cyc",cycle,"_FC",freqIdx,"_",j,"_GN",Temp[2]);
		hisMatFileName  =  string(Temp[1],"_Cyc",cycle,"_FC",freqIdx);
	end

	# Here we set a dump function for GN for this iteracion of FC
	function dumpGN(mc,Dc,iter,pInv,PF)
		dumpFun(mc,Dc,iter,pInv,PF,filename);
	end

	pMisTE = pMisTemp;

	flush(Base.stdout)
	t1 = time_ns();
	if method == "projGN"
		mc,Dc,flag,His = projGNCG(mc,pInv,pMisTE,dumpResults = dumpGN);
	elseif method == "barrierGN"
		mc,Dc,flag,His = barrierGNCG(mc,pInv,pMisTE,rho=1.0,dumpResults = dumpGN);
	end

	e1 = time_ns();
	print("runtime of GN:"); println((e1 - t1)/1.0e9);

	FafterGN = His.F[end];
	println("Computed Misfit with new sources after GN : ",FafterGN);

	if hisMatFileName != ""
		file = matopen(string(hisMatFileName,"_HisGN.mat"), "w");
		### PUT THE CONTENT OF "HIS" INTO THE MAT FILE ###
		His.Dc = [];
		write(file,"His",His);
		close(file);
	end
	His.Dc = []
	push!(HIS,His)
end

pMisTemp = setSourcesSame(pMisTemp,originalSources);
pMisTemp = setDobs(pMisTemp,mergedDobs)
pMisTemp = setWd(pMisTemp,mergedWd)

return mc,Z1,Z2,alpha1,alpha2*simSrcDim,Dc,flag,HIS
end

"""
	function freqCont

	Frequency continuation procedure for running FWI.
	This function runs GaussNewton on misfit functions defined by the frequencies.
"""
function freqCont(freqContParam::FreqContParam)
mc = freqContParam.mc
itersNum = freqContParam.itersNum
originalSources = freqContParam.originalSources
nrcv = freqContParam.nrcv
pInv = freqContParam.pInv
pMis = freqContParam.pMis
windowSize = freqContParam.windowSize
resultsFilename = freqContParam.resultsFilename
dumpFun = freqContParam.dumpFun
Iact = freqContParam.Iact
mback = freqContParam.mback
Z1 = freqContParam.Z1
simSrcDim = freqContParam.simSrcDim
alpha1 = freqContParam.alpha1
alpha2Orig = freqContParam.alpha2Orig
stepReg = freqContParam.stepReg
mode = freqContParam.mode
startFrom = freqContParam.startFrom
endAt = freqContParam.endAt
cycle = freqContParam.cycle
method = freqContParam.method
FWImethod = freqContParam.FWImethod
updateMref = freqContParam.updateMref

Z2 = 0;
Dc = 0;
flag = -1;
HIS = [];
alpha2 = alpha2Orig / simSrcDim;

println("~~~~~~~~~~~~~~~ SEG1  FreqCont: Regs are: ",alpha1,",",alpha2,",",stepReg);
regfun = pInv.regularizer
originalItersNum = itersNum
for freqIdx = startFrom:endAt
	if updateMref
		pInv.mref = copy(mc[:]);
	end
	if mode=="1stInit"
		reqIdx1 = freqIdx;
		if freqIdx > 1
			reqIdx1 = max(2,freqIdx-windowSize+1);
		end
		reqIdx2 = freqIdx;
	else
		reqIdx1 = freqIdx;
		if freqIdx > 1
			reqIdx1 = max(1,freqIdx-windowSize+1);
		end
		reqIdx2 = freqIdx;
	end

	currentProblems = reqIdx1:reqIdx2;
	pMisTemp = pMis[currentProblems];
	println("\n======= New Continuation Stage: selecting continuation batches: ",reqIdx1," to ",reqIdx2,"=======\n");
	if FWImethod == "FWI"
		pInv.maxIter = itersNum;
		mc,Dc,flag,HIS = standardGNrun(method, mc, simSrcDim, originalSources,
		 		pInv, pMisTemp, HIS, resultsFilename, cycle, freqIdx, dumpFun)
	elseif FWImethod == "FWI_ES"
		pInv.maxIter = 1;
		mc,Z1,Z2,alpha1,alpha2,Dc,flag,HIS = ExtendedSourcesGNRun(
				method, mc, simSrcDim, nrcv, itersNum, originalSources,
				Z1,alpha1, alpha2, stepReg, pInv, pMisTemp, HIS,
				resultsFilename, cycle, freqIdx, dumpFun)
	end

end

pInv.regularizer = regfun;
return mc,Z1,Z2,alpha1,alpha2*simSrcDim,Dc,flag,HIS;
end


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

const SMARTS_BOND_SYMBOL = Dict(
    '-' => QueryFormula(:bondorder, 1),
    '=' => QueryFormula(:bondorder, 2),
    '#' => QueryFormula(:bondorder, 3),
    '@' => QueryFormula(:isringbond, true),
    ':' => QueryFormula(:isaromaticbond, true),
    '/' => QueryFormula(:stereo, :up),
    '\\' => QueryFormula(:stereo, :down)
)

defaultbond(state::SmilesParser) = SmilesBond()
defaultbond(state::SmartsParser) = SmartsBond(QueryFormula(:defaultbond, true))

"""
    bondsymbol!(state::SmartsParserState) -> Union{Pair,Nothing}

BondSymbol <- [-=#@:/\\] / '/?' / '\\?'
"""
function bondsymbol!(state::SmartsParserState)
    sym1 = read(state)
    sym2 = lookahead(state, 1)
    if sym1 == '/' && sym2 == '?'
        forward!(state, 2)
        return QueryFormula(:not, QueryFormula(:stereo, :down))
    elseif sym1 == '\\' && sym2 == '?'
        forward!(state, 2)
        return QueryFormula(:not, QueryFormula(:stereo, :up))
    elseif sym1 in keys(SMARTS_BOND_SYMBOL)
        forward!(state)
        return SMARTS_BOND_SYMBOL[sym1]
    end
    # Implicit single bond returns nothing
end

"""
    bond!(state::SmilesParser) -> Union{SmilesBond,Nothing}

Bond <- BondSymbol?
"""
function bond!(state::SmilesParser)
    fml = bondsymbol!(state)
    if fml === nothing
        return
    elseif fml.key == :bondorder
        return SmilesBond(fml.value)
    elseif fml.key == :isaromaticbond
        return SmilesBond(1, true, :unspecified)
    elseif fml.key == :stereo
        return SmilesBond(1, false, fml.value)
    end
    # return nothing
end

"""
    bond!(state::SmartsParser) -> Union{SmartsBond,Nothing}

Bond <- '~' / (BondSymbol / LogicalCond)?
"""
function bond!(state::SmartsParser)
    if read(state) == '~'
        forward!(state)
        return SmartsBond()
    end
    fml = lglowand!(state, bondsymbol!)
    if fml !== nothing
        fml = tidyformula(fml)
        return SmartsBond(fml)
    end
    # return nothing: Invalid bond token or implicit single bond
end


### A Hamiltonian is composed of blocks (kpoints), which have a list of RealFourierOperator
# corresponding to each term
# This is the "high-level" interface, provided for convenience

struct HamiltonianBlock
    basis::PlaneWaveBasis
    kpoint::Kpoint

    # The operators are vectors of RealFourierOperator,
    # not typed because of type invariance issues.
    operators::Vector  # the original list of RealFourierOperator
                       # (as many as there are terms), kept for easier exploration
    optimized_operators::Vector  # the optimized list of RealFourierOperator, to be used for applying
    scratch  # Pre-allocated scratch arrays for fast application
end
function HamiltonianBlock(basis, kpt, operators, scratch)
    HamiltonianBlock(basis, kpt, operators, optimize_operators_(operators), scratch)
end
Base.eltype(block::HamiltonianBlock) = complex(eltype(block.basis))
Base.size(block::HamiltonianBlock) =
    (length(G_vectors(block.kpoint)), length(G_vectors(block.kpoint)))
Base.size(block::HamiltonianBlock, i::Integer) = i < 3 ? size(block)[i] : 1

struct Hamiltonian
    basis::PlaneWaveBasis
    blocks::Vector{HamiltonianBlock}
end

# Loop through bands, IFFT to get ψ in real space, loop through terms, FFT and accumulate into Hψ
# Do this band by band to conserve memory
# As an optimization we special-case nonlocal operators to apply them
# instead on the full block and benefit from BLAS3
@views @timing "Hamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray,
                                                                        H::HamiltonianBlock,
                                                                        ψ::AbstractArray)
    basis = H.basis
    kpt = H.kpoint
    nband = size(ψ, 2)

    # Allocate another scratch array:
    Hψ_fouriers = [similar(Hψ[:, 1]) for tid = 1:Threads.nthreads()]

    # take ψi, IFFT it to ψ_real, apply each term to Hψ_temp and Hψ_real, and add it to Hψi
    Threads.@threads for iband = 1:nband
        tid = Threads.threadid()
        ψ_real = H.scratch.ψ_reals[tid]
        Hψ_real = H.scratch.Hψ_reals[tid]
        Hψ_fourier = Hψ_fouriers[tid]

        Hψ_real .= 0
        Hψ_fourier .= 0
        G_to_r!(ψ_real, basis, kpt, ψ[:, iband])
        for op in H.optimized_operators
            if !(op isa NonlocalOperator)
                apply!((fourier=Hψ_fourier, real=Hψ_real),
                       op,
                       (fourier=ψ[:, iband], real=ψ_real))
            end
        end
        Hψ[:, iband] .= Hψ_fourier
        r_to_G!(Hψ_fourier, basis, kpt, Hψ_real)
        Hψ[:, iband] .+= Hψ_fourier
    end

    # Apply the nonlocal operators
    for op in H.optimized_operators
        if op isa NonlocalOperator
            apply!((fourier=Hψ, real=nothing), op, (fourier=ψ, real=nothing))
        end
    end

    Hψ
end
Base.:*(H::HamiltonianBlock, ψ) = mul!(similar(ψ), H, ψ)

function LinearAlgebra.mul!(Hψ, H::Hamiltonian, ψ)
    for ik = 1:length(H.basis.kpoints)
        mul!(Hψ[ik], H.blocks[ik], ψ[ik])
    end
    Hψ
end
# need `deepcopy` here to copy the elements of the array of arrays ψ (not just pointers)
Base.:*(H::Hamiltonian, ψ) = mul!(deepcopy(ψ), H, ψ)

# Get energies and Hamiltonian
# kwargs is additional info that might be useful for the energy terms to precompute
# (eg the density ρ)
@timing function energy_hamiltonian(basis::PlaneWaveBasis, ψ, occ; kwargs...)
    # it: index into terms, ik: index into kpoints
    ene_ops_arr = [ene_ops(term, ψ, occ; kwargs...) for term in basis.terms]
    energies    = [eh.E for eh in ene_ops_arr]
    operators   = [eh.ops for eh in ene_ops_arr]         # operators[it][ik]
    hks_per_k   = [[blocks[ik] for blocks in operators]  # hks_per_k[ik][it]
                   for ik = 1:length(basis.kpoints)]

    # Preallocated scratch arrays
    T = eltype(basis)
    scratch = (
        ψ_reals=[zeros(complex(T), basis.fft_size...) for tid = 1:Threads.nthreads()],
        Hψ_reals=[zeros(complex(T), basis.fft_size...) for tid = 1:Threads.nthreads()]
    )

    H = Hamiltonian(basis, [HamiltonianBlock(basis, kpt, hks, scratch)
                            for (hks, kpt) in zip(hks_per_k, basis.kpoints)])
    E = Energies(basis.model.term_types, energies)
    (E=E, H=H)
end
function Hamiltonian(basis::PlaneWaveBasis; ψ=nothing, occ=nothing, kwargs...)
    _, H = energy_hamiltonian(basis, ψ, occ; kwargs...)
    H
end

import Base: Matrix, Array
function Matrix(block::HamiltonianBlock)
    sum(Matrix(h) for h in block.optimized_operators)
end
Array(block::HamiltonianBlock) = Matrix(block)

"""
Get the total local potential of the given Hamiltonian, in real space.
"""
function total_local_potential(ham::Hamiltonian)
    @assert ham.basis.model.spin_polarization in (:none, :spinless)
    block = ham.blocks[1]  # all local potentials are the same
    rs = [o for o in block.optimized_operators if o isa RealSpaceMultiplication]
    @assert length(rs) == 1
    rs[1].potential
end


# ******************************************************************************************
# Notices:
#
# Copyright © 2022 United States Government as represented by the Administrator of the
# National Aeronautics and Space Administration.  All Rights Reserved.
#
# Disclaimers
#
# No Warranty: THE SUBJECT SOFTWARE IS PROVIDED "AS IS" WITHOUT ANY WARRANTY OF ANY KIND,
# EITHER EXPRESSED, IMPLIED, OR STATUTORY, INCLUDING, BUT NOT LIMITED TO, ANY WARRANTY
# THAT THE SUBJECT SOFTWARE WILL CONFORM TO SPECIFICATIONS, ANY IMPLIED WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, OR FREEDOM FROM INFRINGEMENT, ANY
# WARRANTY THAT THE SUBJECT SOFTWARE WILL BE ERROR FREE, OR ANY WARRANTY THAT DOCUMENTATION,
# IF PROVIDED, WILL CONFORM TO THE SUBJECT SOFTWARE. THIS AGREEMENT DOES NOT, IN ANY MANNER,
# CONSTITUTE AN ENDORSEMENT BY GOVERNMENT AGENCY OR ANY PRIOR RECIPIENT OF ANY RESULTS,
# RESULTING DESIGNS, HARDWARE, SOFTWARE PRODUCTS OR ANY OTHER APPLICATIONS RESULTING FROM
# USE OF THE SUBJECT SOFTWARE.  FURTHER, GOVERNMENT AGENCY DISCLAIMS ALL WARRANTIES AND
# LIABILITIES REGARDING THIRD-PARTY SOFTWARE, IF PRESENT IN THE ORIGINAL SOFTWARE, AND
# DISTRIBUTES IT "AS IS."
#
# Waiver and Indemnity:  RECIPIENT AGREES TO WAIVE ANY AND ALL CLAIMS AGAINST THE UNITED
# STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS, AS WELL AS ANY PRIOR RECIPIENT.
# IF RECIPIENT'S USE OF THE SUBJECT SOFTWARE RESULTS IN ANY LIABILITIES, DEMANDS, DAMAGES,
# EXPENSES OR LOSSES ARISING FROM SUCH USE, INCLUDING ANY DAMAGES FROM PRODUCTS BASED ON,
# OR RESULTING FROM, RECIPIENT'S USE OF THE SUBJECT SOFTWARE, RECIPIENT SHALL INDEMNIFY AND
# HOLD HARMLESS THE UNITED STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS, AS WELL
# AS ANY PRIOR RECIPIENT, TO THE EXTENT PERMITTED BY LAW.  RECIPIENT'S SOLE REMEDY FOR ANY
# SUCH MATTER SHALL BE THE IMMEDIATE, UNILATERAL TERMINATION OF THIS AGREEMENT.
# ******************************************************************************************

# To be implemented by user.

"""
    GrayBox <: AbstractSimulation

Parent type for user simulation in which environment is exposed.
"""
abstract type GrayBox <: AbstractSimulation end

"""
    reset!(sim::GrayBox)

Reset simulation.
"""
reset!(sim::GrayBox)::Nothing = unimplemented()

"""
    environment(sim::GrayBox)

Return `Environment` object constructed by simulation. `Environment` is an alias of
`Dict{Symbol, Sampleable}`.
"""
environment(sim::GrayBox)::Environment = unimplemented()

"""
    observe(sim::GrayBox)

Return quasi-normalized observation of simulation. Observation should copy, not reference,
simulation state. If simulation is unobservable, leave unimplemented.
"""
observe(sim::GrayBox)::AbstractVector{<:Real} = unimplemented()

"""
    step(sim::GrayBox, x::EnvironmentValue)

Step simulation given an `EnvironmentValue` object. `EnvironmentValue` is an alias of
`Dict{Symbol, Any}`.
"""
step!(sim::GrayBox, x::EnvironmentValue)::Nothing = unimplemented()

"""
    isterminal(sim::GrayBox)

Check whether simulation has finished due to time limit or terminal state, independent of
failure status.
"""
isterminal(sim::GrayBox)::Bool = unimplemented()

"""
   isevent(sim::GrayBox)

Check whether simulation is in event (failure) state. For episodic simulation, check whether
event occurred during episode.
"""
isevent(sim::GrayBox)::Bool = unimplemented()

"""
    distance(sim::GrayBox)

Return custom metric of distance to event. For best results, metric should depend only on
current state. For episodic simulation, return minimum distance to event across entire
episode (miss distance).
"""
distance(sim::GrayBox)::Real = unimplemented()

"""
    flatten(distribution, value)

Flatten environment variable to quasi-normalized array.
"""
flatten(distribution::Any, value::Any)::AbstractVector{<:Real} = unimplemented()
flatten(d::Distribution{Univariate, Continuous}, v::Real) = [(v - mean(d)) / std(d)]
flatten(d::Uniform, v::Real) = [(2 * v - d.a - d.b) / (d.b - d.a)]

"""
    unflatten(distribution, array::AbstractVector{<:Real})

Reconstruct environment variable from quasi-normalized array.
"""
unflatten(distribution::Any, array::AbstractVector{<:Real})::Any = unimplemented()
unflatten(d::Distribution{Univariate, Continuous}, a::AbstractVector{<:Real}) = std(d) * a[] + mean(d)
unflatten(d::Uniform, a::AbstractVector{<:Real}) = (a[] * (d.b - d.a) + d.a + d.b) / 2

"""
    reward(sim::GrayBox, x::EnvironmentValue)

Advanced option: additional reward or heuristic. Returns scalar value or partially-applied
function, allowing efficient calculation of R(s,a,s') = R(s,a)(s'). Should not alter
simulation. See documentation for more details.
"""
reward(sim::GrayBox, x::EnvironmentValue) = 0.0


function get_ylabel(name::String)
    if name == "Phosphorylated_MEKc"
        return "Phosphorylated MEK\n(cytoplasm)"
    elseif name == "Phosphorylated_ERKc"
        return "Phosphorylated ERK\n(cytoplasm)"
    elseif name == "Phosphorylated_RSKw"
        return "Phosphorylated RSK\n(whole cell)"
    elseif name == "Phosphorylated_CREBw"
        return "Phosphorylated CREB\n(whole cell)"
    elseif name == "dusp_mRNA"
        return L"$\it{dusp}$"*" mRNA\nexpression"
    elseif name == "cfos_mRNA"
        return L"$\it{c}$"*"-"*L"$\it{fos}$"*" mRNA\nexpression"
    elseif name == "cFos_Protein"
        return "c-Fos Protein\nexpression"
    elseif name == "Phosphorylated_cFos"
        return "Phosphorylated c-Fos\nProtein expression"
    end
end

function plotFunc_timecourse(Sim::Module, n_file::Vector{Int}, viz_type::String,
                                show_all::Bool, stdev::Bool, simulations_all::Array{Float64, 4})
    if !isdir("./figure/simulation/$viz_type")
        mkdir("./figure/simulation/$viz_type")
    end

    rc("figure", figsize = (20, 8))
    rc("font", family = "Arial")
    rc("mathtext", fontset = "custom")
    rc("mathtext", it = "Arial:italic")
    rc("font", size = 18)
    rc("axes", linewidth = 2)
    rc("xtick.major", width = 2)
    rc("ytick.major", width = 2)
    rc("xtick", direction = "in")
    rc("ytick", direction = "in")
    rc("lines", linewidth = 2.5)
    rc("lines", markersize = 12)

    cmap = ["mediumblue", "red"]
    shape = ["D", "s"]

    subplots_adjust(wspace=0.5, hspace=0.5)

    for (i, name) in enumerate(observables)
        subplot(2, 4, i)
        if show_all
            for j in eachindex(n_file)
                for l in eachindex(Sim.conditions)
                    plot(
                        Sim.t, simulations_all[i, j, :, l]./maximum(simulations_all[i, j, :, :]),
                        color=cmap[l], alpha=0.05
                    )
                end
            end
        end
        if viz_type != "average"
            for l in eachindex(Sim.conditions)
                plot(
                    Sim.t, Sim.simulations[i, :, l]/(maximum(Sim.simulations[i, :, :])),
                    color=cmap[l]
                )
            end
        else
            normalized = Array{Float64, 4}(
                undef, length(observables), length(n_file), length(Sim.t), length(Sim.conditions)
            )
            for j in eachindex(n_file)
                for l in eachindex(Sim.conditions)
                    normalized[i, j, :, l] = simulations_all[i, j, :, l]./maximum(simulations_all[i, j, :, :])
                end
            end
            for l in eachindex(Sim.conditions)
                plot(
                    Sim.t, [mean(filter(!isnan, normalized[i, :, k, l])) for k in eachindex(Sim.t)],
                    color=cmap[l]
                )
            end
            if stdev
                for l in eachindex(Sim.conditions)
                    ymean = [
                        mean(
                            filter(
                                !isnan, normalized[i, :, k, l]
                            )
                        ) for k in eachindex(Sim.t)
                    ]
                    yerr = [
                        std(
                            filter(
                                !isnan, normalized[i, :, k, l]
                            )
                        ) for k in eachindex(Sim.t)
                    ]
                    fill_between(
                        Sim.t, ymean - yerr, ymean + yerr,
                        lw=0, color=cmap[l], alpha=0.1
                    )
                end
            end
        end

        if isassigned(Exp.experiments, i)
            exp_t = Exp.get_timepoint(name)
            if isassigned(Exp.standard_error, i)
                for (l, condition) in enumerate(Sim.conditions)
                    if condition in keys(Exp.experiments[i])
                        exp_data = errorbar(
                            exp_t./60., Exp.experiments[i][condition],
                            yerr=Exp.standard_error[i][condition],
                            ecolor=cmap[l], elinewidth=1, capsize=8,
                            markerfacecolor="None", markeredgecolor=cmap[l],
                            fmt=shape[l], clip_on=false
                        )
                        for capline in exp_data[2]
                            capline.set_clip_on(false)
                        end
                        for barlinecol in exp_data[3]
                            barlinecol.set_clip_on(false)
                        end
                    end
                end
            else
                for (l, condition) in enumerate(Sim.conditions)
                    if condition in keys(Exp.experiments[i])
                        plot(
                            exp_t./60., Exp.experiments[i][condition], shape[l],
                            markerfacecolor="None", markeredgecolor=cmap[l],
                            clip_on=false
                        )
                    end
                end
            end
        end

        xlim(0, 90)
        xticks([0, 30, 60, 90])
        yticks([0, 0.3, 0.6, 0.9, 1.2])
        ylim(0, 1.2)
        xlabel("Time (min)")
        ylabel(get_ylabel(name))
    end

    savefig(
        "./figure/simulation/$viz_type.png", dpi=300, bbox_inches="tight"
    )
end