### Thermal Generation Formulations

abstract type AbstractThermalFormulation <: AbstractDeviceFormulation end

abstract type AbstractThermalDispatchForm <: AbstractThermalFormulation end

struct ThermalUnitCommitment <: AbstractThermalFormulation end

struct ThermalDispatch <: AbstractThermalDispatchForm end

struct ThermalRampLimited <: AbstractThermalDispatchForm end

struct ThermalDispatchNoMin <: AbstractThermalDispatchForm end

# Variables for Thermal Generation

"""
This function add the variables for power generation output to the model
"""
function activepower_variables(ps_m::CanonicalModel,
                               devices::Array{T,1},
                               time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen}

    add_variable(ps_m, devices, time_range, :Pth, false, :var_active)

    return

end

"""
This function add the variables for power generation output to the model
"""
function reactivepower_variables(ps_m::CanonicalModel,
                                 devices::Array{T,1},
                                 time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen}

    add_variable(ps_m, devices, time_range, :Qth, false, :var_reactive)

    return

end

"""
This function add the variables for power generation commitment to the model
"""
function commitment_variables(ps_m::CanonicalModel,
                              devices::Array{T,1},
                              time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen}

    add_variable(ps_m, devices, time_range, :on_th, true)
    add_variable(ps_m, devices, time_range, :start_th, true)
    add_variable(ps_m, devices, time_range, :stop_th, true)

    return

end

"""
This function adds the active power limits of generators when there are no CommitmentVariables
"""
function activepower_constraints(ps_m::CanonicalModel,
                                 devices::Array{T,1},
                                 device_formulation::Type{D},
                                 system_formulation::Type{S},
                                 time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                      D <: AbstractThermalDispatchForm,
                                                                      S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, g.tech.activepowerlimits) for g in devices]

    device_range(ps_m, range_data, time_range, :thermal_active_range, :Pth)

    return

end

"""
This function adds the active power limits of generators when there are CommitmentVariables
"""
function activepower_constraints(ps_m::CanonicalModel,
                                 devices::Array{T,1},
                                 device_formulation::Type{D},
                                 system_formulation::Type{S},
                                 time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                      D <: AbstractThermalFormulation,
                                                                      S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, g.tech.activepowerlimits) for g in devices]

    device_semicontinuousrange(ps_m, range_data, time_range, :thermal_active_range, :Pth, :on_th)

    return

end

"""
This function adds the reactive  power limits of generators when there are CommitmentVariables
"""
function reactivepower_constraints(ps_m::CanonicalModel,
                                   devices::Array{T,1},
                                   device_formulation::Type{D},
                                   system_formulation::Type{S},
                                   time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                        D <: AbstractThermalDispatchForm,
                                                                        S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, g.tech.reactivepowerlimits) for g in devices]

    device_range(ps_m, range_data , time_range, :thermal_reactive_range, :Qth)

    return

end

"""
This function adds the reactive power limits of generators when there CommitmentVariables
"""
function reactivepower_constraints(ps_m::CanonicalModel,
                                   devices::Array{T,1},
                                   device_formulation::Type{D},
                                   system_formulation::Type{S},
                                   time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                        D <: AbstractThermalFormulation,
                                                                        S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, g.tech.reactivepowerlimits) for g in devices]

    device_semicontinuousrange(ps_m, range_data , time_range, :thermal_reactive_range, :Qth, :on_th)

    return

end

"""
This function adds the active power limits of generators when there are no CommitmentVariables
"""
function activepower_constraints(ps_m::CanonicalModel,
                                 devices::Array{T,1},
                                 device_formulation::Type{PSI.ThermalDispatchNoMin},
                                 system_formulation::Type{S},
                                 time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                      S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, (min = 0.0, max=g.tech.activepowerlimits.max)) for g in devices]

    device_range(ps_m, range_data, time_range, :thermal_active_range, :Pth)

    return

end

"""
This function adds the reactive  power limits of generators when there are CommitmentVariables
"""
function reactivepower_constraints(ps_m::CanonicalModel,
                                   devices::Array{T,1},
                                   device_formulation::Type{PSI.ThermalDispatchNoMin},
                                   system_formulation::Type{S},
                                   time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                        S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, (min = 0.0, max=g.tech.reactivepowerlimits.max)) for g in devices]

    device_range(ps_m, range_data , time_range, :thermal_reactive_range, :Qth)

    return

end

### Constraints for Thermal Generation without commitment variables ####

"""
This function adds the Commitment Status constraint when there are CommitmentVariables
"""
function commitment_constraints(ps_m::CanonicalModel,
                                devices::Array{T,1},
                                device_formulation::Type{D},
                                system_formulation::Type{S},
                                time_range::UnitRange{Int64},
                                parameters::Bool) where {T <: PSY.ThermalGen,
                                                                     D <: AbstractThermalFormulation,
                                                                     S <: PM.AbstractPowerFormulation}

    if !(:thermal_status in keys(ps_m.initial_conditions))
        @info("Initial status conditions not provided. This can lead to unwanted results")
        status_init(ps_m, devices, parameters)
    end

    device_commitment(ps_m,
                      ps_m.initial_conditions[:thermal_status],
                      time_range, :commitment_th,
                      (:start_th, :stop_th, :on_th))

    return

end

# ramping constraints

function _get_data_for_rocc(devices::Array{T,1}) where {T <: PSY.ThermalGen}

    set_name = Vector{String}(undef, length(devices))
    ramp_params = Vector{UpDown}(undef, length(devices))
    minmax_params = Vector{MinMax}(undef, length(devices))

    idx = eachindex(devices)
    i, state = iterate(idx)
    for g in devices
        if !isnothing(g.tech.ramplimits)
            set_name[i] = g.name
            ramp_params[i] = g.tech.ramplimits
            minmax_params[i] = g.tech.activepowerlimits
            y = iterate(idx, state)
            y === nothing && (i += 1; break)
            i, state = y
        end
    end

    deleteat!(set_name, i:last(idx))
    deleteat!(ramp_params, i:last(idx))
    deleteat!(minmax_params, i:last(idx))

    return set_name, ramp_params, minmax_params

end

"""
This function adds the ramping limits of generators when there are CommitmentVariables
"""
function ramp_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalFormulation,
                                                               S <: PM.AbstractPowerFormulation}

    data = _get_data_for_rocc(devices)

    if !isempty(data[2])
        if !(:thermal_output in keys(ps_m.initial_conditions))
            @info("Initial conditions for rate of change not provided. This can lead to unwanted results")
            output_init(ps_m, devices, parameters)
        end

        @assert length(data[2]) == length(ps_m.initial_conditions[:thermal_output])
        # Here goes the reactive power ramp limits
        device_mixedinteger_rateofchange(ps_m,
                                        data,
                                        ps_m.initial_conditions[:thermal_output],
                                        time_range,
                                        :ramp_thermal, (:Pth, :start_th, :stop_th))
    else
        @warn "Data doesn't contain generators with ramp limits, consider adjusting your formulation"
    end

    return

end

function ramp_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalDispatchForm,
                                                               S <: PM.AbstractPowerFormulation}

    data = _get_data_for_rocc(devices)

    if !isempty(data[2])
        if !(:thermal_output in keys(ps_m.initial_conditions))
            @info("Initial conditions for rate of change not provided. This can lead to unwanted results")
            output_init(ps_m, devices, parameters)
        end
        # Here goes the reactive power ramp limits
        device_linear_rateofchange(ps_m,
                                   (data[1], data[2]),
                                   ps_m.initial_conditions[:thermal_output], time_range,
                                   :ramp_thermal, :Pth)
    else
        @warn "Data doesn't contain generators with ramp limits, consider adjusting your formulation"
    end

    return

end

"""
This function adds the ramping limits of generators when there are CommitmentVariables
"""
function ramp_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalFormulation,
                                                               S <: PM.AbstractActivePowerFormulation}

    data = _get_data_for_rocc(devices)

    if !isempty(data[2])
        if !(:thermal_output in keys(ps_m.initial_conditions))
            @info("Initial conditions for rate of change not provided. This can lead to unwanted results")
            output_init(ps_m, devices, parameters)
        end

        @assert length(data[2]) == length(ps_m.initial_conditions[:thermal_output])

        device_mixedinteger_rateofchange(ps_m,
                                        data,
                                        ps_m.initial_conditions[:thermal_output],
                                        time_range,
                                        :ramp_thermal, (:Pth, :start_th, :stop_th))

    else
        @warn "Data doesn't contain generators with ramp limits, consider adjusting your formulation"
    end

    return

end

function ramp_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalDispatchForm,
                                                               S <: PM.AbstractActivePowerFormulation}

    data = _get_data_for_rocc(devices)

    if !isempty(data[2])
        if !(:thermal_output in keys(ps_m.initial_conditions))
            @info("Initial conditions for rate of change not provided. This can lead to unwanted results")
            output_init(ps_m, devices, parameters)
        end
        device_linear_rateofchange(ps_m,
                                    (data[1], data[2]),
                                    ps_m.initial_conditions[:thermal_output], time_range,
                                    :ramp_thermal, :Pth)
    else
        @warn "Data doesn't contain generators with ramp limits, consider adjusting your formulation"
    end

    return

end

# time constraints

function _get_data_for_tdc(devices::Array{T,1}) where {T <: PSY.ThermalGen}

    set_name = Vector{String}(undef, length(devices))
    time_params = Vector{UpDown}(undef, length(devices))

    idx = eachindex(devices)
    i, state = iterate(idx)
    for g in devices
        if !isnothing(g.tech.timelimits)
            set_name[i] = g.name
            time_params[i] = g.tech.timelimits
            y = iterate(idx, state)
            y === nothing && (i += 1; break)
            i, state = y
        end
    end

    deleteat!(set_name, i:last(idx))
    deleteat!(time_params, i:last(idx))

    return set_name, time_params
end

function time_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool = false) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalFormulation,
                                                               S <: PM.AbstractPowerFormulation}

    duration_data = _get_data_for_tdc(devices)

    if !isempty(duration_data)
        if !(:thermal_duration_on in keys(ps_m.initial_conditions))
            @info("Initial conditions for time up/down not provided. This can lead to unwanted results")
            duration_init(ps_m, devices, parameters)
        end
        @warn("Currently Min-up/down time limits only work if time resoultion they are defined at matches the simulation's time resoultion")
        device_duration_retrospective(ps_m,
                                      duration_data,
                                      ps_m.initial_conditions[:thermal_duration_on],
                                      ps_m.initial_conditions[:thermal_duration_off],
                                      time_range,
                                      :duration, (:on_th, :start_th, :stop_th))
    else
        @warn "Data doesn't contain generators with time-up/down limits, consider adjusting your formulation"
    end

    return

end

# thermal generation cost

function cost_function(ps_m::CanonicalModel,
                       devices::Array{T,1},
                       device_formulation::Type{D},
                       system_formulation::Type{S}) where {T <: PSY.ThermalGen,
                                                           D <: AbstractThermalDispatchForm,
                                                           S <: PM.AbstractPowerFormulation}

    add_to_cost(ps_m, devices, :Pth, :variablecost)

    return

end

function cost_function(ps_m::CanonicalModel,
                       devices::Array{T,1},
                       device_formulation::Type{D},
                       system_formulation::Type{S}) where {T <: PSY.ThermalGen,
                                                           D <: AbstractThermalFormulation,
                                                           S <: PM.AbstractPowerFormulation}

    #Variable Cost component
    add_to_cost(ps_m, devices, :Pth, :variablecost)

    #Commitment Cost Components
    add_to_cost(ps_m, devices, :start_th, :startupcost)
    add_to_cost(ps_m, devices, :stop_th, :shutdncost)
    add_to_cost(ps_m, devices, :on_th, :fixedcost)

    return

end


"""
Representation of a trajectory, ie a Hamiltonian with a discrete integrator that
also checks for divergence.
"""
struct TrajectoryNUTS{Tf,S}
    # "Hamiltonian."
    # H::TH
    "Log density of z (negative log energy) at initial point."
    π₀::Tf
    "Stepsize for leapfrog."
    ϵ::Tf
    "Smallest decrease allowed in the log density."
    min_Δ::Tf
    "Turn statistic configuration."
    turn_statistic_configuration::S
end

function move(tree::Tree{D,T,L}, H::Hamiltonian, trajectory::TrajectoryNUTS, z::PhasePoint{D,T,L}, fwd::Bool) where {D,T,L}
    @unpack ϵ = trajectory
    leapfrog(tree, H, z, fwd ? ϵ : -ϵ)
end

###
### proposals
###

"""
$(SIGNATURES)
Random boolean which is `true` with the given probability `exp(logprob)`, which can be `≥ 1`
in which case no random value is drawn.
"""
function rand_bool_logprob(rng::AbstractRNG, logprob)
    logprob ≥ 0 || (randexp(rng, Float64) > -logprob)
end

function calculate_logprob2(::TrajectoryNUTS, is_doubling, ω₁, ω₂, ω)
    biased_progressive_logprob2(is_doubling, ω₁, ω₂, ω)
end

function combine_proposals(rng, tree::Tree, ::TrajectoryNUTS, z₁, z₂, logprob2::Real, is_forward::Bool)
    z, flag = rand_bool_logprob(rng, logprob2) ? (z₂, z₁.flag) : (z₁, z₂.flag)
    # is_doubling || free_z!(tree, flag)
    free_z!(tree, flag)
    z
end
# function combine_proposals(
    # rng, tree::Tree, ::TrajectoryNUTS, z₁, z₂, logprob2::Real, is_forward, guard₁, guard₂
# )
    # z, flag, guard = rand_bool_logprob(rng, logprob2) ? (z₂, z₁.flag, guard₁) : (z₁, z₂.flag, guard₂)
    # guard || free_z!(tree, flag)
    # z
# end

###
### statistics for visited nodes
###

struct AcceptanceStatistic{T}
    """
    Logarithm of the sum of metropolis acceptances probabilities over the whole trajectory
    (including invalid parts).
    """
    log_sum_α::T
    "Total number of leapfrog steps."
    steps::Int32
end

function combine_acceptance_statistics(A::AcceptanceStatistic, B::AcceptanceStatistic)
    AcceptanceStatistic(logaddexp(A.log_sum_α, B.log_sum_α), A.steps + B.steps)
end

"""
$(SIGNATURES)
Acceptance statistic for a leaf. The initial leaf is considered not to be visited.
"""
function leaf_acceptance_statistic(Δ, is_initial)
    is_initial ? AcceptanceStatistic(oftype(Δ, -Inf), zero(Int32)) : AcceptanceStatistic(Δ < zero(Δ) ? Δ : zero(Δ), one(Int32))
end

"""
$(SIGNATURES)
Return the acceptance rate (a `Real` betwen `0` and `1`).
"""
acceptance_rate(A::AcceptanceStatistic) = (a = exp(A.log_sum_α) / A.steps; a < one(a) ? a : one(a))

combine_visited_statistics(::TrajectoryNUTS, v, w) = combine_acceptance_statistics(v, w)

###
### turn analysis
###

"Statistics for the identification of turning points. See Betancourt (2017, appendix)."
struct GeneralizedTurnStatistic{D,T,L}
    p♯₋::FlaggedVector{D,T,L}
    p♯₊::FlaggedVector{D,T,L}
    ρ::FlaggedVector{D,T,L}
end

"""
A dummy turn statistic to use for convenient type stability; returns null ptrs that should hopefully segfault if accessed
(we'd only access them accidentally in case of a bug!)
"""
function dummy_turn_statistic(::Tree{D,T,L}) where {D,T,L}
    dummy_flag = zero(UInt32)
    dummy_ptr = reinterpret(Ptr{T}, zero(UInt))
    GeneralizedTurnStatistic(
        FlaggedVector{D,T,L}(dummy_ptr, dummy_flag),
        FlaggedVector{D,T,L}(dummy_ptr, dummy_flag),
        FlaggedVector{D,T,L}(dummy_ptr, dummy_flag)
    )
end

function leaf_turn_statistic(tree::Tree{D,T,L}, ::Val{:generalized}, H, z::PhasePoint{D,T,L}) where {D,T,L}
    p♯ = calculate_p♯(tree, H, z)
    GeneralizedTurnStatistic(p♯, p♯, FlaggedVector{D,T,L}(z.p, 0x00000000))
end

function combine_turn_statistics(
    tree::Tree, ::TrajectoryNUTS,
    x::GeneralizedTurnStatistic{D,T,L},
    y::GeneralizedTurnStatistic{D,T,L}
) where {D,T,L}
    ρₓ = x.ρ
    ρʸ = y.ρ
    # @show x.p♯₊.flag, y.p♯₋.flag
    if ρₓ.flag == 0x00000000
        # We are in depth 1; position allocated as part of a PhasePoint, not as part of a turn statistic
        # Therefore we cannot free either of them, and must allocate a new vector to store the results in.
        ρ = undefined_Σρ(tree)
        # Additionally, x.p♯₊ == x.p♯₋ and y.p♯₊ == y.p♯₋, therefore we cannot free them.
    else # we are at a depth of 2 or greater.
        # We can therefore reuse a ρ; arbitrarily, we reuse xρ
        ρ = ρₓ
        # and free yρ as well as the two discarded p♯
        free_Σρ!(tree, ρʸ.flag)
        free_ρ♯!(tree, x.p♯₊.flag)
        free_ρ♯!(tree, y.p♯₋.flag)
    end
    @avx for l in eachindex(ρ)
        ρ[l] = ρₓ[l] + ρʸ[l]
    end
    # x.p♯₊.flag == x.p♯₋.flag || free_ρ♯!(tree, x.p♯₊.flag)
    # y.p♯₊.flag == y.p♯₋.flag || free_ρ♯!(tree, y.p♯₋.flag)
    GeneralizedTurnStatistic(x.p♯₋, y.p♯₊, ρ)
end

function is_turning(::TrajectoryNUTS, τ::GeneralizedTurnStatistic{D,T,L}) where {D,T,L}
    # Uses the generalized NUTS criterion from Betancourt (2017).
    @unpack p♯₋, p♯₊, ρ = τ
    # @argcheck p♯₋ ≢ p♯₊ "internal error: is_turning called on a leaf"
    d♯₋ = zero(T)
    d♯₊ = zero(T)
    @avx for d ∈ eachindex(ρ)
        ρᵈ = ρ[d]
        d♯₋ += ρᵈ * p♯₋[d]
        d♯₊ += ρᵈ * p♯₊[d]
    end
    # Current version always calculates both dot products; alternative:
    # dot(p♯₋, ρ) < 0 || dot(p♯₊, ρ) < 0
    # Calculating both in one loop is faster
    # ( 2D fma + 3D loads vs 2D fma + 4D loads; at most 2/4 ops/cycle
    #   can be loads, so the loads cannot keep up with the 2/4 fma/cycle
    #   and are the bottleneck ),
    # but it disallows conditional checking to skip second dot product.
    #
    # Which is faster on average depends on probability of turning
    # probability of turning will be fairly low for most models.
    (d♯₋ < zero(T)) | (d♯₊ < zero(T))
end

###
### leafs
###

function leaf(tree::Tree, H::Hamiltonian, trajectory::TrajectoryNUTS, z, is_initial)
    @unpack π₀, min_Δ, turn_statistic_configuration = trajectory
    # @show logdensity(H, z), π₀
    Δ = is_initial ? zero(π₀) : logdensity(H, z) - π₀
    isdiv = Δ < min_Δ
    v = leaf_acceptance_statistic(Δ, is_initial)
    if isdiv
        # @show H
        # @show z
        # @show typeof(z)
        (z, Δ, dummy_turn_statistic(tree)), v, true
    else
        τ = leaf_turn_statistic(tree, turn_statistic_configuration, H, z)
        (z, Δ, τ), v, false
    end
end

####
#### NUTS interface
####

"""
$(TYPEDEF)
Implementation of the “No-U-Turn Sampler” of Hoffman and Gelman (2014), including subsequent
developments, as described in Betancourt (2017).
# Fields
$(FIELDS)
"""
struct NUTS{S}
    "Maximum tree depth."
    max_depth::Int
    "Threshold for negative energy relative to starting point that indicates divergence."
    min_Δ::Float64
    """
    Turn statistic configuration. Currently only `Val(:generalized)` (the default) is
    supported.
    """
    turn_statistic_configuration::S
    function NUTS(; max_depth = DEFAULT_MAX_TREE_DEPTH, min_Δ = -1000.0,
                  turn_statistic_configuration::S = Val{:generalized}()) where {S}
        # @argcheck 0 < max_depth ≤ MAX_DIRECTIONS_DEPTH
        # @argcheck min_Δ < 0
        new{S}(Int(max_depth), Float64(min_Δ), turn_statistic_configuration)
    end
end

"""
$(TYPEDEF)
Diagnostic information for a single tree built with the No-U-turn sampler.
# Fields
Accessing fields directly is part of the API.
$(FIELDS)
"""
struct TreeStatisticsNUTS # now 32 bytes; fits in ymm register
    "Log density (negative energy)."
    π::Float64
    "Acceptance rate statistic."
    acceptance_rate::Float64
    "Reason for termination. See [`InvalidTree`](@ref) and [`REACHED_MAX_DEPTH`](@ref)."
    termination::InvalidTree
    "Depth of the tree."
    depth::Int32
    "Number of leapfrog steps evaluated."
    steps::Int32
    # "Directions for tree doubling (useful for debugging)."
    # directions::Directions
end

"""
$(SIGNATURES)
No-U-turn Hamiltonian Monte Carlo transition, using Hamiltonian `H`, starting at evaluated
log density position `Q`, using stepsize `ϵ`. Parameters of `algorithm` govern the details
of tree construction.
Return two values, the new evaluated log density position, and tree statistics.
"""
function sample_tree(rng, tree::Tree, algorithm::NUTS, H::Hamiltonian, z::PhasePoint, ϵ;
                          p = nothing, directions = rand(rng, Directions))
    if p === nothing
        rand_p!(rng, z.p, H.κ)
    else
        @unpack Q, flag = z
        z = PhasePoint(Q, p, flag)
    end
    @unpack max_depth, min_Δ, turn_statistic_configuration = algorithm
    trajectory = TrajectoryNUTS(logdensity(H, z), ϵ, min_Δ, turn_statistic_configuration)
    ζ, v, termination, depth = sample_trajectory(rng, tree, H, trajectory, z, max_depth, directions)
    tree_statistics = TreeStatisticsNUTS(logdensity(H, ζ), acceptance_rate(v), termination, depth, v.steps)#, directions)
    ζ, tree_statistics
end


# unary # 2^2 - 2 = 2 definitions
@inline fmap(f::F, x::Tuple{X}) where {F,X} = (f(first(x)),)
@inline fmap(f::F, x::NTuple) where {F} = (f(first(x)), fmap(f, Base.tail(x))...)

# binary # 2^3 - 2 = 6 definitions
@inline fmap(f::F, x::Tuple{X}, y::Tuple{Y}) where {F,X,Y} = (f(first(x), first(y)),)
@inline fmap(f::F, x::Tuple{X}, y) where {F,X} = (f(first(x), y),)
@inline fmap(f::F, x, y::Tuple{Y}) where {F,Y} = (f(x, first(y)),)
@inline fmap(f::F, x::Tuple{Vararg{Any,N}}, y::Tuple{Vararg{Any,N}}) where {F,N} = (f(first(x), first(y)), fmap(f, Base.tail(x), Base.tail(y))...)
@inline fmap(f::F, x::Tuple, y) where {F} = (f(first(x), y), fmap(f, Base.tail(x), y)...)
@inline fmap(f::F, x, y::Tuple) where {F} = (f(x, first(y)), fmap(f, x, Base.tail(y))...)

fmap(f::F, x::Tuple{X}, y::Tuple) where {F,X} = throw("Dimension mismatch.")
fmap(f::F, x::Tuple, y::Tuple{Y}) where {F,Y} = throw("Dimension mismatch.")
fmap(f::F, x::Tuple, y::Tuple) where {F} = throw("Dimension mismatch.")

@generated function fmap(f::F, x::Vararg{Any,N}) where {F,N}
    q = Expr(:block, Expr(:meta, :inline))
    t = Expr(:tuple)
    U = 1
    call = Expr(:call, :f)
    syms = Vector{Symbol}(undef, N)
    istup = Vector{Bool}(undef, N)
    for n ∈ 1:N
        syms[n] = xₙ = Symbol(:x_, n)
        push!(q.args, Expr(:(=), xₙ, Expr(:ref, :x, n)))
        istup[n] = ist = (x[n] <: Tuple)
        if ist
            U = length(x[n].parameters)
            push!(call.args, Expr(:ref, xₙ, 1))
        else
            push!(call.args, xₙ)
        end
    end
    push!(t.args, call)
    for u ∈ 2:U
        call = Expr(:call, :f)
        for n ∈ 1:N
            xₙ = syms[n]
            if istup[n]
                push!(call.args, Expr(:ref, xₙ, u))
            else
                push!(call.args, xₙ)
            end
        end
        push!(t.args, call)
    end
    push!(q.args, t); q
end

for op ∈ [:vsub, :vabs, :vfloor, :vceil, :vtrunc, :vround, :vsqrt, :vnot, :vleading_zeros, :vtrailing_zeros, :vsub_fast]
    @eval @inline $op(v1::VecUnroll{N,W,T}) where {N,W,T} = VecUnroll(fmap($op, v1.data))
end
# only for `Float32` and `Float64`
@inline inv_approx(v::VecUnroll{N,W,T}) where {N,W,T<:Union{Float32,Float64}} = VecUnroll(fmap(inv_approx, v.data))

@inline vreinterpret(::Type{T}, v::VecUnroll) where {T<:Number} = VecUnroll(fmap(vreinterpret, T, v.data))
for op ∈ [:vadd,:vsub,:vmul,:vand,:vor,:vxor,:vlt,:vle,:vgt,:vge,:veq,:vne,:vadd_fast,:vsub_fast,:vmul_fast]
    @eval begin
        @inline $op(v1::VecUnroll, v2::VecUnroll) = VecUnroll(fmap($op, v1.data, v2.data))
        @inline $op(v1::VecUnroll{N,W,T,V}, ::StaticInt{M}) where {N,W,T,V,M} = VecUnroll(fmap($op, v1.data, vbroadcast(Val{W}(), T(M))))
        @inline $op(::StaticInt{M}, v1::VecUnroll{N,W,T,V}) where {N,W,T,V,M} = VecUnroll(fmap($op, vbroadcast(Val{W}(), T(M)), v1.data))

        @inline $op(v1::VecUnroll{N,W,T,V}, ::StaticInt{1}) where {N,W,T,V} = VecUnroll(fmap($op, v1.data, One()))
        @inline $op(::StaticInt{1}, v1::VecUnroll{N,W,T,V}) where {N,W,T,V} = VecUnroll(fmap($op, One(), v1.data))

        @inline $op(v1::VecUnroll{N,W,T,V}, ::StaticInt{0}) where {N,W,T,V} = VecUnroll(fmap($op, v1.data, Zero()))
        @inline $op(::StaticInt{0}, v1::VecUnroll{N,W,T,V}) where {N,W,T,V} = VecUnroll(fmap($op, Zero(), v1.data))
    end
end
for op ∈ [:vmax,:vmax_fast,:vmin,:vmin_fast,:vcopysign]
    @eval begin
        @inline $op(v1::VecUnroll, v2::VecUnroll) = VecUnroll(fmap($op, v1.data, v2.data))
    end
end
for op ∈ [:vgt,:vge,:vlt,:vle,:veq,:vne,:vmax,:vmax_fast,:vmin,:vmin_fast]
    @eval begin
        @inline $op(v::VecUnroll{N,W,T,V}, s::Union{NativeTypes,AbstractSIMDVector}) where {N,W,T,V} = VecUnroll(fmap($op, v.data, vbroadcast(Val{W}(), s)))
        @inline $op(s::Union{NativeTypes,AbstractSIMDVector}, v::VecUnroll{N,W,T,V}) where {N,W,T,V} = VecUnroll(fmap($op, vbroadcast(Val{W}(), s), v.data))
    end
end

for op ∈ [:vrem, :vshl, :vashr, :vlshr, :vdiv, :vfdiv, :vrem_fast, :vfdiv_fast]
    @eval begin
        @inline $op(v1::VecUnroll, v2::VecUnroll) = VecUnroll(fmap($op, v1.data, v2.data))
    end
end
for op ∈ [:vrem, :vand, :vor, :vxor, :vshl, :vashr, :vlshr, :vlt, :vle,:vgt,:vge,:veq,:vne]
    @eval begin
        @inline $op(vu::VecUnroll, i::MM) = $op(vu, Vec(i))
        @inline $op(i::MM, vu::VecUnroll) = $op(Vec(i), vu)
    end
end
for op ∈ [:vshl, :vashr, :vlshr]
    @eval @inline $op(m::Mask, vu::VecUnroll) = $op(Vec(m), vu)
end
for op ∈ [:rotate_left,:rotate_right,:funnel_shift_left,:funnel_shift_right]
    @eval begin
        @inline $op(v1::VecUnroll{N,W,T,V}, v2::R) where {N,W,T,V,R<:IntegerTypes} = VecUnroll(fmap($op, v1.data, promote_type(V, R)(v2)))
        @inline $op(v1::R, v2::VecUnroll{N,W,T,V}) where {N,W,T,V,R<:IntegerTypes} = VecUnroll(fmap($op, promote_type(V, R)(v1), v2.data))
        @inline $op(v1::VecUnroll, v2::VecUnroll) = VecUnroll(fmap($op, v1.data, v2.data))
    end
end

for op ∈ [:vfma, :vmuladd, :vfma_fast, :vmuladd_fast, :vfnmadd, :vfmsub, :vfnmsub, :vfnmadd_fast, :vfmsub_fast, :vfnmsub_fast,
          :vfmadd231, :vfnmadd231, :vfmsub231, :vfnmsub231, :ifmahi, :ifmalo]
    @eval begin
        @inline $op(v1::VecUnroll{N,W}, v2::VecUnroll{N,W}, v3::VecUnroll{N,W}) where {N,W} = VecUnroll(fmap($op, v1.data, v2.data, v3.data))
    end
end
@inline vifelse(v1::VecUnroll{N,W,<:Boolean}, v2::T, v3::T) where {N,W,T<:NativeTypes} = VecUnroll(fmap(vifelse, v1.data, Vec{W,T}(v2), Vec{W,T}(v3)))
@inline vifelse(v1::VecUnroll{N,W,<:Boolean}, v2::T, v3::T) where {N,W,T<:Real} = VecUnroll(fmap(vifelse, v1.data, v2, v3))
@inline vifelse(v1::Vec{W,Bool}, v2::VecUnroll{N,W,T}, v3::Union{NativeTypes,AbstractSIMDVector,StaticInt}) where {N,W,T} = VecUnroll(fmap(vifelse, Vec{W,T}(v1), v2.data, Vec{W,T}(v3)))
@inline vifelse(v1::Vec{W,Bool}, v2::Union{NativeTypes,AbstractSIMDVector,StaticInt}, v3::VecUnroll{N,W,T}) where {N,W,T} = VecUnroll(fmap(vifelse, Vec{W,T}(v1), Vec{W,T}(v2), v3.data))
@inline vifelse(v1::VecUnroll{N,W,<:Boolean}, v2::VecUnroll{N,W,T}, v3::Union{NativeTypes,AbstractSIMDVector,StaticInt}) where {N,W,T} = VecUnroll(fmap(vifelse, v1.data, v2.data, Vec{W,T}(v3)))
@inline vifelse(v1::VecUnroll{N,W,<:Boolean}, v2::Union{NativeTypes,AbstractSIMDVector,StaticInt}, v3::VecUnroll{N,W,T}) where {N,W,T} = VecUnroll(fmap(vifelse, v1.data, Vec{W,T}(v2), v3.data))
@inline vifelse(v1::Vec{W,Bool}, v2::VecUnroll{N,W,T}, v3::VecUnroll{N,W,T}) where {N,W,T} = VecUnroll(fmap(vifelse, Vec{W,T}(v1), v2.data, v3.data))
@inline vifelse(v1::VecUnroll{N,W,<:Boolean}, v2::VecUnroll{N,W,T}, v3::VecUnroll{N,W,T}) where {N,W,T} = VecUnroll(fmap(vifelse, v1.data, v2.data, v3.data))
@inline function vifelse(v1::VecUnroll{N,W,<:Boolean}, v2::VecUnroll{N,W}, v3::VecUnroll{N,W}) where {N,W}
    v4, v5 = promote(v2, v3)
    VecUnroll(fmap(vifelse, v1.data, v4.data, v5.data))
end

@inline veq(v::VecUnroll{N,W,T}, x::AbstractIrrational) where {N,W,T} = v == vbroadcast(Val{W}(), T(x))
@inline veq(x::AbstractIrrational, v::VecUnroll{N,W,T}) where {N,W,T} = vbroadcast(Val{W}(), T(x)) == v

@inline vunsafe_trunc(::Type{T}, v::VecUnroll) where {T<:Real} = VecUnroll(fmap(vunsafe_trunc, T, v.data))
@inline vrem(v::VecUnroll, ::Type{T}) where {T<:Real} = VecUnroll(fmap(vrem, v.data, T))
@inline vrem(v::VecUnroll{N,W}, ::Type{VecUnroll{N,W,T,V}}) where {N,W,T,V} = VecUnroll(fmap(vrem, v.data, V))

@inline (::Type{VecUnroll{N,W,T,V}})(vu::VecUnroll{N,W,T,V}) where {N,W,T,V<:AbstractSIMDVector{W,T}} = vu
@inline (::Type{VecUnroll{N,W,T,VT}})(vu::VecUnroll{N,W,S,VS})  where {N,W,T,VT<:AbstractSIMDVector{W,T},S,VS<:AbstractSIMDVector{W,S}} = VecUnroll(fmap(convert, Vec{W,T}, vu.data))

function collapse_expr(N, op)
    N += 1
    t = Expr(:tuple); s = Vector{Symbol}(undef, N)
    for n ∈ 1:N
        s_n = s[n] = Symbol(:v_, n)
        push!(t.args, s_n)
    end
    q = quote
        $(Expr(:meta,:inline))
        $t = data(vu)
    end
    while N > 1
        for n ∈ 1:N >>> 1
            push!(q.args, Expr(:(=), s[n], Expr(:call, op, s[n], s[n + (N >>> 1)])))
        end
        isodd(N) && push!(q.args, Expr(:(=), s[1], Expr(:call, op, s[1], s[N])))
        N >>>= 1
    end
    q
end
@generated collapse_add(vu::VecUnroll{N}) where {N} = collapse_expr(N, :vadd)
@generated collapse_mul(vu::VecUnroll{N}) where {N} = collapse_expr(N, :vmul)
@generated collapse_max(vu::VecUnroll{N}) where {N} = collapse_expr(N, :max)
@generated collapse_min(vu::VecUnroll{N}) where {N} = collapse_expr(N, :min)
@generated collapse_and(vu::VecUnroll{N}) where {N} = collapse_expr(N, :&)
@generated collapse_or(vu::VecUnroll{N}) where {N} = collapse_expr(N, :|)
@inline vsum(vu::VecUnroll) = vsum(collapse_add(vu))
@inline vsum(s, vu::VecUnroll) = vsum(s, collapse_add(vu))
@inline vprod(vu::VecUnroll) = vprod(collapse_mul(vu))
@inline vprod(s, vu::VecUnroll) = vprod(s, collapse_mul(vu))
@inline vmaximum(vu::VecUnroll) = vmaximum(collapse_max(vu))
@inline vminimum(vu::VecUnroll) = vminimum(collapse_min(vu))
@inline vall(vu::VecUnroll) = vall(collapse_and(vu))
@inline vany(vu::VecUnroll) = vany(collapse_or(vu))


#!/usr/bin/env julia

#
#  This is the main file for data assimilation runs.
#
#
# 1. loads up a configuration file,
# 2. obtains data from a WRF model,
# 3. construct covariate vectors
# 3. reads in observations and metadata for a list of stations,
# 4. runs the moisture model and the assimilation mechanism.
#
#

using Calendar
import Calendar.CalendarTime

using Storage
import Storage.setup_tag, Storage.spush, Storage.next_frame, Storage.flush_frame

using Stations
import Stations.Station, Stations.Observation, Stations.load_station_info,
       Stations.load_station_data, Stations.build_observation_data, Stations.register_to_grid, Stations.nearest_grid_point, Stations.obs_value, Stations.obs_station_id

using Kriging
import Kriging.trend_surface_model_kriging

using NetCDF
using WRF

using FM
import FM.FMModel, FM.advance_model, FM.kalman_update

function create_fm_variables(file_name)

    # first extract the dimensions
    nc = NetCDF.open(file_name)
    we_dim = nc.dim["west_east"]
    sn_dim = nc.dim["south_north"]
    t_dim = nc.dim["Time"]
    NetCDF.close(nc)

    var_names = ASCIIString[ "FM1", "FM10", "FM100" ]
    long_names = ["1hr fuel moisture", "10hr fuel moisture", "100hr fuel moisture"]
    for i in 1:3
        if !haskey(nc.vars, var_names[i])
            atts = {"coordinates"=>"XLONG XLAT","units"=>"kg/kg","stagger"=>"","MemoryOrder"=>"XY ","description"=>long_names[i]}
            NetCDF.nccreate(file_name, var_names[i], atts, we_dim, sn_dim, t_dim)
            println("INFO: created variable $(var_names[i])")
        else
            println("INFO: variable $(var_names[i]) already existed")
        end
    end
end

function main(args)

    # the arguments passed to the julia program do not include program name
    if length(args) != 1
        println("Usage: julia run_data_assimilation.jl cfg_file")
        exit(1)
    end

    t1 = Calendar.now()
    println("INFO: run_data_assimilation.jl started on $t1")

    ### Read configuration file and setup the system
    println("INFO: reading in config file $(args[1])")
    cfg = evalfile(args[1])

    # create the output directory if it does not exist
    println("INFO: output directory is ", cfg["output_dir"])
    !ispath(cfg["output_dir"]) && mkdir(cfg["output_dir"])

    # configure Storage mechanism
    Storage.sopen(cfg["output_dir"], "moisture_model_v2_diagnostics.txt", "frame")

    # setup Storage & output policies for interesting quantities
    setup_tag("mt", true, true)

    setup_tag("fm1_model_state_assim", false, false)
    setup_tag("fm1_model_var", false, false)
    setup_tag("fm10_model_state", false, false)
    setup_tag("fm10_model_state_assim", false, false)
    setup_tag("fm10_model_var", false, false)
    setup_tag("fm100_model_state_assim", false, false)
    setup_tag("fm100_model_var", false, false)

    setup_tag("fm10_model_na_state", false, false)
    setup_tag("fm10_model_deltas", false, false)

    setup_tag("kriging_beta", true, true)
    setup_tag("kriging_xtx_cond", false, true)
    setup_tag("kriging_field", false, false)
    setup_tag("kriging_variance", false, false)
    setup_tag("kriging_sigma2_eta", true, true)
    setup_tag("kriging_iters", false, true)
    setup_tag("kriging_subzero_s2_estimates", false, true)
    setup_tag("model_raws_mae", false, true)
    setup_tag("model_raws_mae_assim", false, true)
    setup_tag("model_na_raws_mae", false, true)

    # co-located model/model_na/kriging field/observation
    setup_tag("kriging_obs", false, false)
    setup_tag("kriging_obs_station_ids", false, false)
    setup_tag("kriging_obs_ngp", false, false)
    setup_tag("kriging_errors", true, true)

    setup_tag("kalman_gain_fm10", false, false)

    ### Load WRF output data
    t1 = Calendar.now()
    println("INFO: configuration complete, loading WRF data.")

    # read in data from the WRF output file pointed to by cfg
    w = WRF.load_wrf_data(cfg["wrf_output"], ["HGT"])

    # the terrain height need not be stored for all time points
    WRF.slice_field(w, "HGT")

    # extract WRF fields
    lat, lon = WRF.lat(w), WRF.lon(w)
    dsize = size(lat)
    wtm = WRF.times(w)
    println("INFO: WRF grid size is $(dsize[1]) x $(dsize[2]) and found $(length(wtm)) timepoints.")

    # retrieve equilibria and rain (these are already precomputed)
    Ed, Ew = WRF.field(w, "Ed"), WRF.field(w, "Ew")
    rain = WRF.field(w, "RAIN")
    hgt = WRF.field(w, "HGT")
    T = WRF.interpolated_field(w, "T2")
    P = WRF.interpolated_field(w, "PSFC")

    t2 = Calendar.now()
    println("INFO: WRF output loaded, sliced and diced [$(t2-t1)].")

    ### Load observation data from stations
    io = open(join([cfg["station_info_dir"], cfg["station_info"]], "/"), "r")
    station_ids = filter(x -> x[1] != '#', map(x -> strip(x), readlines(io)))
    close(io)

    # load each station from its info and observation files
    stations = Station[]
    for sid in station_ids
        s = load_station_info(join([cfg["station_info_dir"], string(sid, ".info")], "/"))
        load_station_data(s, join([cfg["station_data_dir"], string(sid, ".obs")], "/"))
	register_to_grid(s, lat, lon)
#        println("STATION: $(s.id), $(s.loc), ngp is $(s.ngp) with lat $(lat[s.ngp[1], s.ngp[2]]) and lon $(lon[s.ngp[1], s.ngp[2]]).")
        push!(stations, s)
    end

    # build the observation data from stations
    obs_fm10 = build_observation_data(stations, "FM")
    obs_times = keys(obs_fm10)

    t3 = Calendar.now()
    println("INFO: Station data loaded and preprocessed [$(t3 - t2)].")

    ### Initialize model

    # number of simulated fuel components
    Nf = 3

    # construct initial conditions (FIXME: can we do better here?)
    E = squeeze(0.5 * (Ed[:,:,2] + Ew[:,:,2]), 3)

    # set up parameters
    Q = diagm(cfg["Q"])
    P0 = diagm(cfg["P0"])
    mV = zeros(Float64, dsize)
    pred = zeros(Float64, dsize)
    mresV = zeros(Float64, dsize)
    mid = zeros(Int32, dsize)
    Kg = zeros(Float64, (dsize[1], dsize[2], 9))
    K = zeros(Float64, dsize)
    V = zeros(Float64, dsize)

    # prepare static & time-varying covariates
    cov_ids = cfg["covariates"]
    st_covar_map = [:lon => lon,
                    :lat => lat,
                    :elevation => hgt,
                    :constant => ones(Float64, dsize) ]
    dyn_covar_map = [:temperature => T, :pressure => P, :rain => rain]
    Xd3 = length(cov_ids) + 1
    X = zeros(Float64, (dsize[1], dsize[2], Xd3))
    Xr = zeros(Float64, (dsize[1], dsize[2], Xd3))
    for i in 2:Xd3
        cov_id = cov_ids[i-1]
        if haskey(st_covar_map, cov_id)
            println("INFO: processing static covariate $cov_id.")
            v = st_covar_map[cov_id]
            Xr[:,:,i] = v
        elseif haskey(dyn_covar_map, cov_id)
            println("INFO: found dynamic covariate $(cov_id).")
        else
            error("ERROR: unknown covariate $(cov_id) encountered, fatal.")
        end
    end
    println("INFO: there are $Xd3 covariates (including model state).")

    t1 = Calendar.now()
    println("INFO: starting simulation at $t1 ...")
    dt = (wtm[2] - wtm[1]).millis / 1000
    assim_time_win = cfg["assimilation_time_window"]
    println("INFO: time step from WRF is $dt s, assimilation time window is $assim_time_win s.")

    # construct model grid from fuel parameters
    Tk = [ 1.0, 10.0, 100.0 ]
    models = Array(FMModel, dsize)
    models_na = Array(FMModel, dsize)
    for i in 1:dsize[1]
    	for j in 1:dsize[2]
            geo_loc = (lat[i,j], lon[i,j])
	    models[i,j] = FMModel(geo_loc, Nf, E[i,j], P0, Tk)
	    models_na[i,j] = FMModel(geo_loc, Nf, E[i,j], P0, Tk)
	end
    end

    # create the new variables in the wrfout file for storing fuel moisture
    println("INFO: creating FM variables in wrfout file")
    create_fm_variables(cfg["wrf_output"])
    nc = NetCDF.open(cfg["wrf_output"], NetCDF.NC_WRITE)

    ###  Run the model and data assimilation
    for t in 2:length(wtm)
    	mt = wtm[t]
        spush("mt", mt)

        # run the model update (in parallel if possible)
	for i in 1:dsize[1]
	    for j in 1:dsize[2]
	        advance_model(models[i,j], Ed[i, j, t-1], Ew[i, j, t-1], rain[i, j, t-1], dt, Q)
		advance_model(models_na[i,j], Ed[i, j, t-1], Ew[i, j, t-1], rain[i, j, t-1], dt, Q)
	    end
	end

        # store the model state in an array (and store in output frame)
        fm10_model_state = [ models[i,j].m_ext[2] for i=1:dsize[1], j=1:dsize[2] ]
        fm10_model_na_state = [ models_na[i,j].m_ext[2] for i=1:dsize[1], j=1:dsize[2] ]
        fm10_model_var = [ models[i,j].P[2,2] for i=1:dsize[1], j=1:dsize[2] ]

        spush("fm10_model_state", fm10_model_state)
        spush("fm10_model_na_state", fm10_model_na_state)
        spush("fm10_model_var", fm10_model_var)

        # if observation data for this timepoint is available
        obs_i = Observation[]
        tm_valid_now = filter(x -> abs((mt - x).millis) / 1000.0 <= assim_time_win/2, obs_times)

        # gather all observations
        for tvn in tm_valid_now append!(obs_i, obs_fm10[tvn]) end

        # exclude zero observations - must be sensor failure
        obs_i = filter(x -> obs_value(x) > 0, obs_i)

        # if there are no valid observations, continue with next time step, else run kriging
        if length(obs_i) > 0

            # set the current fm10 model state as the covariate
            X[:,:,1] = fm10_model_state

#            fm10_norm = sum(fm10_model_state.^2)^0.5
            println("INFO: assimilating $(length(obs_i)) observations.")

            # loop over additional covariates
            for i in 2:Xd3
                cov_id = cov_ids[i-1]
                if haskey(st_covar_map, cov_id)
                    # just copy and rescale corresponding static covariate
                    X[:,:,i] = Xr[:,:,i]
                elseif haskey(dyn_covar_map, cov_id)
                    # retrieve the field pointed to by the dynamic covariate id
                    F = dyn_covar_map[cov_id]
                    X[:,:,i] = squeeze(F[:,:,t], 3)
                else
                    error("FATAL: found unknown covariate.")
                end
            end

            # store diagnostic information
            ngp_list = map(x -> nearest_grid_point(x), obs_i)
            stat_ids = map(x -> obs_station_id(x), obs_i)
            m_at_obs = Float64[X[i, j, 1] for (i,j) in  ngp_list]
            m_na_at_obs = Float64[models_na[i,j].m_ext[2] for (i,j) in ngp_list]
            raws = Float64[obs_value(o) for o in obs_i]

            spush("model_raws_mae", mean(abs(m_at_obs - raws)))
            spush("model_na_raws_mae", mean(abs(m_na_at_obs - raws)))

            spush("kriging_obs", raws)
            spush("kriging_obs_station_ids", stat_ids)
            spush("kriging_obs_ngp", ngp_list)

            # compute the kriging estimates and fill in pre-allocated arrays
            trend_surface_model_kriging(obs_i, X, K, V)

            # push diagnostic outputs
            spush("kriging_field", K)
            spush("kriging_variance", V)

            # execute the Kalman update at each grid point
            Kp = zeros(1)
            Vp = zeros(1,1)
            Kg = zeros(Float64, dsize)
            fuel_types = [2]
            for i in 1:dsize[1]
                for j in 1:dsize[2]
                    Kp[1] = K[i,j]
                    Vp[1,1] = V[i,j]
                    Kg[i,j] = kalman_update(models[i,j], Kp, Vp, fuel_types)[1,1]
                end
            end

            # push the fm10 model state after the assimilation
            fm10_model_state = [ models[i,j].m_ext[2] for i=1:dsize[1], j=1:dsize[2] ]
            spush("fm10_model_state_assim", fm10_model_state)

            # retrieve adjustments to time constants and to equilibria
            fm10_adj = zeros(Float64, (6, dsize[1], dsize[2]))
            for i in 1:dsize[1]
                for j in 1:dsize[2]
                    fm10_adj[:,i,j] = models[i,j].m_ext[Nf+1:2*Nf+3]
                end
            end
            fm10_adj_max = [ max(abs(fm10_adj[i,:,:])) for i in 1:6 ]
            spush("fm10_model_deltas", fm10_adj_max)

            # gather model values at ngp points after assimilation
            m_at_obs = Float64[fm10_model_state[i, j] for (i,j) in  ngp_list]
            spush("model_raws_mae_assim", mean(abs(m_at_obs - raws)))

            spush("kalman_gain_fm10", Kg)

            # move to the next storage frame
            next_frame()

        end # if there is anything to assimilate

        # store current assimilated state in the wrfout file
        fm_stor = zeros(Float64, (dsize[1], dsize[2], 1))
        fm_stor[:,:,1] = [ models[i,j].m_ext[1] for i=1:dsize[1], j=1:dsize[2] ]
        NetCDF.putvar(nc, "FM1", [1, 1, t], fm_stor)
        fm_stor[:,:,1] = [ models[i,j].m_ext[2] for i=1:dsize[1], j=1:dsize[2] ]
        NetCDF.putvar(nc, "FM10", [1, 1, t], fm_stor)
        fm_stor[:,:,1] = [ models[i,j].m_ext[3] for i=1:dsize[1], j=1:dsize[2] ]
        NetCDF.putvar(nc, "FM100", [1, 1, t], fm_stor)

    end # for each time point

    # Close down the storage system
    Storage.sclose()

    NetCDF.close(nc)

    t2 = Calendar.now()
    println("INFO: simulation completed at $t2 after $(t2-t1).")

end

main(ARGS)


export MPIElasticPropagatorParams, MPIElasticSource, MPIElasticReceiver, MPIElasticPropagator, compute_PML_Params!,
        MPIElasticPropagatorSolver, MPISimulatedObservation!, MPIElasticReceiver, extract_local_patch
@with_kw mutable struct MPIElasticPropagatorParams
    # number of grids along x,y axis and time steps
    NX::Int64 = 0
    NY::Int64 = 0
    M::Int64 = 0
    N::Int64 = 0
    n::Int64 = 300
    NSTEP::Int64 = 2000 * 2

    # size of grid cell and time step
    DELTAX::Float64 = 10.
    DELTAY::Float64 = 10.
    DELTAT::Float64 = 2.e-3 / 2

    # PML boundary conditon
    f0::Float64 = 5. # source frequency
    vp_ref::Float64 = 2000.
    USE_PML_XMIN::Bool = true
    USE_PML_XMAX::Bool = true
    USE_PML_YMIN::Bool = true
    USE_PML_YMAX::Bool = true
    NPOINTS_PML::Int64 = 12
    NPOWER::Float64 = 2.
    K_MAX_PML::Float64 = 1.
    ALPHA_MAX_PML::Float64 = 2. * π * ( f0/2.)
    Rcoef::Float64 = 0.001 # Relative reflection coefficient

    # Auxilliary Data
    ax::Union{Missing, Array{Float64, 2}} = missing
    bx::Union{Missing, Array{Float64, 2}} = missing
    kx::Union{Missing, Array{Float64, 2}} = missing
    ay::Union{Missing, Array{Float64, 2}} = missing
    by::Union{Missing, Array{Float64, 2}} = missing
    ky::Union{Missing, Array{Float64, 2}} = missing

    # display params
    IT_DISPLAY::Int64 = 0

    # mpi
    II::Int64 = -1
    JJ::Int64 = -1

    # indices
    i_1j::Union{Missing, Array{Int64, 1}} = missing
    i_2j::Union{Missing, Array{Int64, 1}} = missing
    i1j::Union{Missing, Array{Int64, 1}} = missing
    i2j::Union{Missing, Array{Int64, 1}} = missing
    ij_1::Union{Missing, Array{Int64, 1}} = missing
    ij_2::Union{Missing, Array{Int64, 1}} = missing
    ij1::Union{Missing, Array{Int64, 1}} = missing
    ij2::Union{Missing, Array{Int64, 1}} = missing
    ij::Union{Missing, Array{Int64, 1}} = missing
    i1j1::Union{Missing, Array{Int64, 1}} = missing

end

"""
    MPIElasticSource

The input is the global index.

`srci`, `srcj`: nsrc
`srcv`: nt × nsrc, each row corresponds to a source time function
"""
mutable struct MPIElasticSource
    srci::Union{Missing, PyObject}
    srcj::Union{Missing, PyObject}
    srctype::Union{Missing,PyObject}
    srcv::Union{Missing, PyObject}
    function MPIElasticSource(param::MPIElasticPropagatorParams, srci::Array{Int64,1},
         srcj::Array{Int64,1}, srctype::Array{Int64, 1}, srcv::Union{Array{Float64,2}, PyObject})
        II, JJ = param.II, param.JJ
        nsrc = length(srci)
        srcv = constant(srcv)
        srcv = srcv'
        local_srci = Int64[]
        local_srcj = Int64[]
        local_srcv = PyObject[]
        local_srctype = Int64[]
        for i = 1:nsrc
            if ((II - 1) * param.n + 1 <= srci[i] <= II * param.n) &&
                    ((JJ - 1) * param.n + 1 <= srcj[i] <= JJ * param.n)
                    push!(local_srci, srci[i] - (II - 1) * param.n + 2)
                    push!(local_srcj, srcj[i] - (JJ - 1) * param.n + 2)
                    push!(local_srcv, srcv[i])
                    push!(local_srctype, srctype[i])
            end
        end
        local_srcv = hcat(local_srcv...)
        if length(local_srci)>0
            local_srci, local_srcj, local_srctype, local_srcv = convert_to_tensor([local_srci, local_srcj, local_srctype, local_srcv], [Int64, Int64, Int64, Float64])
            new(local_srci, local_srcj, local_srctype, local_srcv)
        else
            new(missing, missing, missing, missing)
        end
    end
end

mutable struct MPIElasticReceiver
    rcvi::Union{Missing,Array{Int64, 1}}
    rcvj::Union{Missing,Array{Int64, 1}}
    rcvtype::Union{Missing, Array{Int64, 1}}
    rcvv::Union{Missing, PyObject}
    function MPIElasticReceiver(param::MPIElasticPropagatorParams, rcvi::Array{Int64,1},
         rcvj::Array{Int64,1}, rcvtype::Array{Int64, 1})
        II, JJ = param.II, param.JJ
        nrcv = length(rcvi)
        local_rcvi = Int64[]
        local_rcvj = Int64[]
        local_rcvtype = Int64[]
        for i = 1:nrcv
            if ((II - 1) * param.n + 1 <= rcvi[i] <= II * param.n) &&
                    ((JJ - 1) * param.n + 1 <= rcvj[i] <= JJ * param.n)
                    push!(local_rcvi, rcvi[i] - (II - 1) * param.n)
                    push!(local_rcvj, rcvj[i] - (JJ - 1) * param.n)
                    push!(local_rcvtype, rcvtype[i])
            end
        end
        if length(local_rcvi)>0
            new(local_rcvi, local_rcvj, local_rcvtype, missing)
        else
            new(missing, missing, missing, missing)
        end
    end
end

"""
    ElasticPropagator

A structure holding all the information for acoustic simulation
"""
mutable struct MPIElasticPropagator
    param::MPIElasticPropagatorParams
    src::MPIElasticSource
    vx::PyObject
    vy::PyObject
    sigmaxx::PyObject
    sigmayy::PyObject
    sigmaxy::PyObject
    mem::PyObject
end

function get_mpi_id2(a, b, n)
    idx = Int64[]
    for i = 1:length(a)
        for j = 1:length(b)
            push!(idx, (n+4)*(a[i]-1)+b[j])
        end
    end
    idx
end

function extract_local_patch(param::MPIElasticPropagatorParams, v::Union{PyObject, Array{Float64,2}}; deps = missing, tag = 0)
    v = v[(param.II - 1) * param.n + 1:param.II * param.n, (param.JJ - 1) * param.n + 1:param.JJ * param.n]
    v = mpi_halo_exchange2(v, param.M, param.N, deps = deps, tag = tag)
end

function compute_PML_Params!(param::MPIElasticPropagatorParams)
    NX, NY = param.NX, param.NY
    n = param.n
    r = mpi_rank()
    param.M = div(NX,n)
    param.N = div(NY,n)
    param.II, param.JJ = div(r, param.N)+1, mod(r, param.N)+1
    if mpi_size() != param.M * param.N
        error("MPI Size Mismatch: mpi_size = $(mpi_size()), M = $(param.M), N = $(param.N)")
    end
    @assert NX>0 && NY>0
    @assert mod(NX, n)==0 && mod(NY, n)==0

    # ID
    k1 = 4:n+3
    k2 = 5:n+4
    k_1 = 2:n+1
    k_2 = 1:n
    kk = 3:n+2
    param.i_1j = get_mpi_id2(k_1, kk, n)
    param.i_2j = get_mpi_id2(k_2, kk, n)
    param.i1j = get_mpi_id2(k1, kk, n)
    param.i2j = get_mpi_id2(k2, kk, n)
    param.ij_1 = get_mpi_id2(kk, k_1, n)
    param.ij_2 = get_mpi_id2(kk, k_2, n)
    param.ij1 = get_mpi_id2(kk, k1, n)
    param.ij2 = get_mpi_id2(kk, k2, n)
    param.ij = get_mpi_id2(kk, kk, n)
    param.i1j1 = get_mpi_id2(k1, k1, n)

    ## define profile of absorption in PML region
    NPOINTS_PML = param.NPOINTS_PML
    NPOWER = param.NPOWER
    NX = param.NX
    NY = param.NY
    DELTAX = param.DELTAX
    DELTAY = param.DELTAY
    DELTAT = param.DELTAT
    K_MAX_PML = param.K_MAX_PML
    ALPHA_MAX_PML = param.ALPHA_MAX_PML
    USE_PML_XMIN = param.USE_PML_XMIN
    USE_PML_XMAX = param.USE_PML_XMAX
    USE_PML_YMIN = param.USE_PML_YMIN
    USE_PML_YMAX = param.USE_PML_YMAX
    Rcoef = param.Rcoef

    thickness_PML_x = NPOINTS_PML * DELTAX
    thickness_PML_y = NPOINTS_PML * DELTAY

    # reflection coefficient
    d0_x = - (NPOWER + 1) * param.vp_ref * log(Rcoef) / (2. * thickness_PML_x)
    d0_y = - (NPOWER + 1) * param.vp_ref * log(Rcoef) / (2. * thickness_PML_y)

    xoriginleft = thickness_PML_x
    xoriginright = (NX-0.5)*DELTAX - thickness_PML_x

    a_x = zeros(n)
    a_y = zeros(n)
    a_x_half = zeros(n)
    a_y_half = zeros(n)
    b_x = zeros(n)
    b_y = zeros(n)
    b_x_half = zeros(n)
    b_y_half = zeros(n)
    d_x = zeros(n)
    d_y = zeros(n)
    d_x_half = zeros(n)
    d_y_half = zeros(n)
    K_x = ones(n)
    K_y = ones(n)
    K_x_half = ones(n)
    K_y_half = ones(n)
    alpha_x = zeros(n)
    alpha_y = zeros(n)
    alpha_x_half = zeros(n)
    alpha_y_half = zeros(n)

    for i = 1:n
        xval = DELTAX * (i-1) + (param.II - 1) * param.n * DELTAX
        if (USE_PML_XMIN)

            abscissa_in_PML = xoriginleft - xval
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_x
                d_x[i] = d0_x * abscissa_normalized^NPOWER
                K_x[i] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_x[i] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_x[i] = ALPHA_MAX_PML * abscissa_normalized
            end

            abscissa_in_PML = xoriginleft - (xval + DELTAX/2.)
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_x
                d_x_half[i] = d0_x * abscissa_normalized^NPOWER
                K_x_half[i] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_x_half[i] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_x_half[i] = ALPHA_MAX_PML * abscissa_normalized
            end

        end

        if (USE_PML_XMAX)
            abscissa_in_PML = xval - xoriginright
            if (abscissa_in_PML >= 0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_x
                d_x[i] = d0_x * abscissa_normalized^NPOWER
                K_x[i] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_x[i] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_x[i] = ALPHA_MAX_PML * abscissa_normalized
            end

            abscissa_in_PML = xval + DELTAX/2.0 - xoriginright
            if (abscissa_in_PML >= 0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_x
                d_x_half[i] = d0_x * abscissa_normalized^NPOWER
                K_x_half[i] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_x_half[i] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_x_half[i] = ALPHA_MAX_PML * abscissa_normalized
            end
        end

        if (alpha_x[i] < 0)
            alpha_x[i] = 0
        end
        if (alpha_x_half[i] < 0)
            alpha_x_half[i] = 0
        end

        b_x[i] = exp(- (d_x[i] / K_x[i] + alpha_x[i]) * DELTAT)
        b_x_half[i] = exp(- (d_x_half[i] / K_x_half[i] + alpha_x_half[i]) * DELTAT)

        if (abs(d_x[i]) > 1e-6)
            a_x[i] = d_x[i] * (b_x[i] - 1.) / (K_x[i] * (d_x[i] + K_x[i] * alpha_x[i]))
        end
        if (abs(d_x_half[i]) > 1e-6)
            a_x_half[i] = d_x_half[i] * (b_x_half[i] - 1.) / (K_x_half[i] * (d_x_half[i] + K_x_half[i] * alpha_x_half[i]))
        end
    end

    yoriginbottom = thickness_PML_y
    yorigintop = (NY-0.5)*DELTAY - thickness_PML_y

    for j = 1:n

        yval = DELTAY * (j-1) + (param.JJ - 1) * param.n * DELTAX
        if param.USE_PML_YMIN
            abscissa_in_PML = yoriginbottom - yval
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML_y
                d_y[j] = d0_y * abscissa_normalized^NPOWER
                K_y[j] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_y[j] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_y[j] = ALPHA_MAX_PML * abscissa_normalized
            end
            abscissa_in_PML = yoriginbottom - (yval + DELTAY/2.)
            if abscissa_in_PML >= 0
                abscissa_normalized = abscissa_in_PML / thickness_PML_y
                d_y_half[j] = d0_y * abscissa_normalized^NPOWER
                K_y_half[j] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_y_half[j] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_y_half[j] = ALPHA_MAX_PML * abscissa_normalized
            end
        end

        if param.USE_PML_YMAX
            abscissa_in_PML = yval - yorigintop
            if abscissa_in_PML >= 0
                abscissa_normalized = abscissa_in_PML / thickness_PML_y
                d_y[j] = d0_y * abscissa_normalized^NPOWER
                K_y[j] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_y[j] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_y[j] = ALPHA_MAX_PML * abscissa_normalized
            end

            abscissa_in_PML = yval + DELTAY/2. - yorigintop
            if abscissa_in_PML >= 0
                abscissa_normalized = abscissa_in_PML / thickness_PML_y
                d_y_half[j] = d0_y * abscissa_normalized^NPOWER
                K_y_half[j] = 1. + (K_MAX_PML - 1.) * abscissa_normalized^NPOWER
                # alpha_y_half[j] = ALPHA_MAX_PML * (1. - abscissa_normalized)
                alpha_y_half[j] = ALPHA_MAX_PML * abscissa_normalized
            end
        end

        b_y[j] = exp(- (d_y[j] / K_y[j] + alpha_y[j]) * DELTAT)
        b_y_half[j] = exp(- (d_y_half[j] / K_y_half[j] + alpha_y_half[j]) * DELTAT)
        if abs(d_y[j]) > 1e-6
            a_y[j] = d_y[j] * (b_y[j] - 1.) / (K_y[j] * (d_y[j] + K_y[j] * alpha_y[j]))
        end
        if (abs(d_y_half[j]) > 1e-6)
            a_y_half[j] = d_y_half[j] * (b_y_half[j] - 1.) / (K_y_half[j] * (d_y_half[j] + K_y_half[j] * alpha_y_half[j]))
        end
    end
    # @show norm(d_x-d_y)

    ax = [a_x'; a_x_half']
    bx = [b_x'; b_x_half']
    kx = [K_x'; K_x_half']
    ay = [a_y'; a_y_half']
    by = [b_y'; b_y_half']
    ky = [K_y'; K_y_half']

    param.ax = ax
    param.bx = bx
    param.kx = kx
    param.ay = ay
    param.by = by
    param.ky = ky
    param
end

function MPIElasticPropagatorSolver(param::MPIElasticPropagatorParams, src::MPIElasticSource,
    ρ::Union{PyObject, Array{Float64, 2}}, λ::Union{PyObject, Array{Float64, 2}}, μ::Union{PyObject, Array{Float64, 2}};
    tag_offset::Int64 = 0, dep = nothing)
    ρ = reshape(convert_to_tensor(ρ),(-1,))
    λ = reshape(convert_to_tensor(λ),(-1,))
    μ = reshape(convert_to_tensor(μ),(-1,))
    n = param.n
    @assert size(ρ)==size(λ)==size(μ)==((n+4)^2,)
    @assert param.II>0
    if !ismissing(src.srcv)
        @assert size(src.srcv, 2)==length(src.srci)==length(src.srcj)==length(src.srctype)
        @assert size(src.srcv, 1)==param.NSTEP
    end
    NSTEP = param.NSTEP
    vx_arr = TensorArray(NSTEP+1);
    vx_arr = write(vx_arr, 1, constant(zeros(n^2)))
    vy_arr = TensorArray(NSTEP+1);
    vy_arr = write(vy_arr, 1, constant(zeros(n^2)))
    sigmaxx_arr = TensorArray(NSTEP+1);
    sigmaxx_arr = write(sigmaxx_arr, 1, constant(zeros(n^2)))
    sigmayy_arr = TensorArray(NSTEP+1);
    sigmayy_arr = write(sigmayy_arr, 1, constant(zeros(n^2)))
    sigmaxy_arr = TensorArray(NSTEP+1);
    sigmaxy_arr = write(sigmaxy_arr, 1, constant(zeros(n^2)))
    mem_arr = [TensorArray(NSTEP+1) for i = 1:8]
    for i = 1:8
        mem_arr[i] = write(mem_arr[i], 1, constant(zeros(n^2)))
    end
    ax,bx,kx,ay,by,ky = param.ax,param.bx,param.kx,param.ay,param.by,param.ky
    function condition(i, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arr...)
        i<=NSTEP+1
    end

    if isnothing(dep)
        dep = 0.0
    end

    function body(i, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arr...)
        mem = Array{PyObject}(undef, 8)
        σxx,σyy,σxy,vx,vy =
            read(sigmaxx_arr, i-1), read(sigmayy_arr, i-1), read(sigmaxy_arr, i-1),
            read(vx_arr, i-1), read(vy_arr, i-1)
        for k = 1:8
            m = read(mem_arr[k], i-1)
            mem[k] = _reshape_and_halo_exchange(m, param, 18*i+k-1+tag_offset, k>1 ? mem[k-1][1] + dep : λ[1] + ρ[1] + μ[1] + vy[1] + σxy[1] + dep)
        end

        mpi_idx = 18*i+8+tag_offset

        halo = (op, deps, mpi_idx)->begin
            op = _reshape_and_halo_exchange(op, param, mpi_idx, deps)
            mpi_idx = mpi_idx + 1
            op, mpi_idx
        end
        vx, mpi_idx = halo(vx, mem[8][1], mpi_idx)
        vy, mpi_idx = halo(vy, vx[1], mpi_idx)
        σxx, mpi_idx = halo(σxx, vy[1], mpi_idx)
        σyy, mpi_idx = halo(σyy, σxx[1], mpi_idx)
        σxy, mpi_idx = halo(σxy, σyy[1], mpi_idx)

        srci,srcj,srctype,srcv = ElasticSourceAtTimeT(src, i-1)
        σxx, σyy, σxy,vx, vy, mem, mpi_idx = one_step(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
            ax, bx, ay, by, kx, ky, srci, srcj, srcv, srctype, mpi_idx)

        sigmaxx_arr,sigmayy_arr,sigmaxy_arr,vx_arr,vy_arr =
            write(sigmaxx_arr, i, σxx), write(sigmayy_arr, i, σyy), write(sigmaxy_arr, i, σxy),
                    write(vx_arr, i, vx), write(vy_arr, i, vy)

        mem_arr_ = Array{PyObject}(undef, 8)
        for k = 1:8
            mem_arr_[k] = write(mem_arr[k], i, mem[k])
        end

        if param.IT_DISPLAY>0
            op = tf.cond(tf.equal(tf.math.floormod(i,param.IT_DISPLAY),1), ()->tf.print("time step = ", i-1, " out of $NSTEP"), ()->tf.no_op())
            i = bind(i, op)
        end

        i+1, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arr_...
    end
    i = constant(2, dtype=Int32)
    _, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arrs =
        while_loop(condition, body, [i, vx_arr, vy_arr, sigmaxx_arr, sigmayy_arr, sigmaxy_arr, mem_arr...])

    vx = reshape(stack(vx_arr), (param.NSTEP+1, n, n))
    vy = reshape(stack(vy_arr), (param.NSTEP+1, n, n))
    σxx = reshape(stack(sigmaxx_arr), (param.NSTEP+1, n, n))
    σyy = reshape(stack(sigmayy_arr), (param.NSTEP+1, n, n))
    σxy = reshape(stack(sigmaxy_arr), (param.NSTEP+1, n, n))
    MPIElasticPropagator(
        param, src, vx, vy, σxx, σyy, σxy, mem_arrs
    )
end

function _reshape_and_print(o, m, n, info)
    o = reshape(o, (m, n))
    o = print_tensor(o, info)
    reshape(o, (m*n,))
end

function _reshape_and_halo_exchange(o, param, tag, deps)
    @assert size(o)==(param.n^2,)
    o = reshape(o, (param.n, param.n))
    out = mpi_halo_exchange2(o, param.M, param.N, tag = tag, deps= deps)
    reshape(out, ((param.n+4)^2,))
end

function fw1(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
        ax, bx, ay, by, kx, ky, mpi_idx)
    NX, NY = param.NX, param.NY
    n = param.n
    DELTAX, DELTAY, DELTAT = param.DELTAX, param.DELTAY, param.DELTAT
    i_1j = param.i_1j
    i_2j = param.i_2j
    i1j = param.i1j
    i2j = param.i2j
    ij_1 = param.ij_1
    ij_2 = param.ij_2
    ij1 = param.ij1
    ij2 = param.ij2
    ij = param.ij

    # λ_ = 0.5*(λ[i1j]+λ[ij])
    # μ_ = 0.5*(μ[i1j]+μ[ij])
    λ_ = λ[ij]
    μ_ = μ[ij]
    λμ = λ_ + 2μ_
    value_dvx_dx = (27vx[i1j]-27vx[ij]-vx[i2j]+vx[i_1j])/(24*DELTAX)
    value_dvy_dy = (27*vy[ij]-27*vy[ij_1]-vy[ij1]+vy[ij_2]) / (24*DELTAY)
    # broadcast
    mem[1] =  makevector(mem[1], ij, adbroadcast(bx[2,:], mem[1][ij], 1) + adbroadcast(ax[2,:], value_dvx_dx, 1));
    mem[2] =  makevector(mem[2], ij, adbroadcast(by[1,:], mem[2][ij], 2) + adbroadcast(ay[1,:], value_dvy_dy, 2));
    value_dvx_dx = adbroadcast(value_dvx_dx, kx[2,:], 3) + mem[1][ij]
    value_dvy_dy = adbroadcast(value_dvy_dy, ky[1,:], 4) + mem[2][ij]
    σxx = scatter_add_op(σxx, ij, (λμ * value_dvx_dx + λ_ * value_dvy_dy) * DELTAT )
    σyy = scatter_add_op(σyy, ij, (λμ * value_dvy_dy + λ_ * value_dvx_dx) * DELTAT )

    σxx = _reshape_and_halo_exchange(σxx[ij], param, mpi_idx, σxy[1])
    σyy = _reshape_and_halo_exchange(σyy[ij], param, mpi_idx+1, σxx[1])

    return vx, vy, σxx, σyy,σxy, mem, mpi_idx+2
end

function fw2(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
    ax, bx, ay, by, kx, ky, mpi_idx)
    NX, NY = param.NX, param.NY
    n = param.n
    DELTAX, DELTAY, DELTAT = param.DELTAX, param.DELTAY, param.DELTAT
    i_1j = param.i_1j
    i_2j = param.i_2j
    i1j = param.i1j
    i2j = param.i2j
    ij_1 = param.ij_1
    ij_2 = param.ij_2
    ij1 = param.ij1
    ij2 = param.ij2
    ij = param.ij

    # μ_ = 0.5*(μ[ij]+μ[ij1])
    μ_ = μ[ij]
    value_dvy_dx = (27vy[ij]-27vy[i_1j]-vy[i1j]+vy[i_2j])/(24*DELTAX)
    value_dvx_dy = (27*vx[ij1]-27*vx[ij]-vx[ij2]+vx[ij_1]) / (24*DELTAY)

    # broadcast
    mem[3] = makevector(mem[3], ij,  adbroadcast(bx[1,:], mem[3][ij], 1) + adbroadcast(ax[1, :], value_dvy_dx, 1))
    mem[4] = makevector(mem[4], ij,  adbroadcast(by[2,:], mem[4][ij], 2) + adbroadcast(ay[2, :], value_dvx_dy, 2));

    value_dvy_dx = adbroadcast(value_dvy_dx, kx[1,:], 3) + mem[3][ij]
    value_dvx_dy = adbroadcast(value_dvx_dy, ky[1,:], 4) + mem[4][ij]

    σxy = scatter_add_op(σxy, ij, μ_*(value_dvy_dx + value_dvx_dy) * DELTAT )

    σxy = _reshape_and_halo_exchange(σxy[ij], param, mpi_idx, σyy[1])

    return vx, vy, σxx, σyy, σxy,mem, mpi_idx+1
end

function fw3(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
    ax, bx, ay, by, kx, ky, mpi_idx)
    NX, NY = param.NX, param.NY
    n = param.n
    DELTAX, DELTAY, DELTAT = param.DELTAX, param.DELTAY, param.DELTAT
    i_1j = param.i_1j
    i_2j = param.i_2j
    i1j = param.i1j
    i2j = param.i2j
    ij_1 = param.ij_1
    ij_2 = param.ij_2
    ij1 = param.ij1
    ij2 = param.ij2
    ij = param.ij

    value_dsigmaxx_dx = (27*σxx[ij]-27*σxx[i_1j]-σxx[i1j]+σxx[i_2j]) / (24*DELTAX);
    value_dsigmaxy_dy = (27*σxy[ij]-27*σxy[ij_1]-σxy[ij1]+σxy[ij_2]) / (24*DELTAY);

    mem[5] = makevector(mem[5],ij, adbroadcast(bx[1,:], mem[5][ij], 1) + adbroadcast(ax[1,:], value_dsigmaxx_dx, 1))
    mem[6] = makevector(mem[6],ij, adbroadcast(by[1,:], mem[6][ij], 2) + adbroadcast(ay[1,:], value_dsigmaxy_dy, 2))

    value_dsigmaxx_dx = adbroadcast(value_dsigmaxx_dx, kx[1,:], 3) + mem[5][ij]
    value_dsigmaxy_dy = adbroadcast(value_dsigmaxy_dy, ky[1,:], 4) + mem[6][ij]

    vx = scatter_add_op(vx, ij, (value_dsigmaxx_dx + value_dsigmaxy_dy) * DELTAT / ρ[ij])
    vx = _reshape_and_halo_exchange(vx[ij], param, mpi_idx, σxy[1])

    return vx, vy, σxx, σyy, σxy, mem, mpi_idx+1
end

function fw4(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
    ax, bx, ay, by, kx, ky, mpi_idx)
    NX, NY, n = param.NX, param.NY, param.n
    DELTAX, DELTAY, DELTAT = param.DELTAX, param.DELTAY, param.DELTAT
    i_1j = param.i_1j
    i_2j = param.i_2j
    i1j = param.i1j
    i2j = param.i2j
    ij_1 = param.ij_1
    ij_2 = param.ij_2
    ij1 = param.ij1
    ij2 = param.ij2
    ij = param.ij
    i1j1 = param.i1j1

    # ρ_ = 0.25 * (ρ[ij] + ρ[i1j] + ρ[i1j1] + ρ[ij1]);
    ρ_ = ρ[ij]

    value_dsigmaxy_dx = (27*σxy[i1j]-27*σxy[ij]-σxy[i2j]+σxy[i_1j]) / (24*DELTAX);
    value_dsigmayy_dy = (27*σyy[ij1]-27*σyy[ij]-σyy[ij2]+σyy[ij_1]) / (24*DELTAY);

    mem[7] = makevector(mem[7], ij, adbroadcast(bx[2,:], mem[7][ij], 1) + adbroadcast(ax[2,:] , value_dsigmaxy_dx, 1))
    mem[8] = makevector(mem[8], ij, adbroadcast(by[2,:], mem[8][ij], 2) + adbroadcast(ay[2,:] , value_dsigmayy_dy, 2))

    value_dsigmaxy_dx = adbroadcast(value_dsigmaxy_dx, kx[2,:], 3) + mem[7][ij]
    value_dsigmayy_dy = adbroadcast(value_dsigmayy_dy, ky[2,:], 4) + mem[8][ij]

    vy = scatter_add_op( vy, ij, (value_dsigmaxy_dx + value_dsigmayy_dy) * DELTAT / ρ_)
    vy = _reshape_and_halo_exchange(vy[ij], param, mpi_idx, vx[1])
    return vx, vy, σxx, σyy, σxy,mem, mpi_idx+1
end

function one_step(param::MPIElasticPropagatorParams, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
    ax, bx, ay, by, kx, ky, srci, srcj, srcv, srctype, mpi_idx)
    n = param.n
    @assert size(ρ)==size(λ)==size(μ)==size(vx)==size(vy)==size(σxx)==size(σyy)==size(σxy)==((n+4)^2,)

    add_source = load_op_and_grad("$(@__DIR__)/../deps/CustomOps/build/libADSeismic","add_source", multiple=true)
    vx, vy, σxx, σyy, σxy,mem, mpi_idx = fw1(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
                    ax, bx, ay, by, kx, ky, mpi_idx)
    vx, vy, σxx, σyy, σxy,mem, mpi_idx = fw2(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
                    ax, bx, ay, by, kx, ky, mpi_idx)
    vx, vy, σxx, σyy, σxy,mem, mpi_idx = fw3(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
                    ax, bx, ay, by, kx, ky, mpi_idx)
    vx, vy, σxx, σyy, σxy,mem, mpi_idx = fw4(param, ρ, λ, μ, vx, vy, σxx, σyy, σxy,mem,
                    ax, bx, ay, by, kx, ky, mpi_idx)
    if !ismissing(srci)
        σxx, σyy, σxy,vx, vy  =  add_source(σxx, σyy, σxy,vx, vy, constant(srci),constant(srcj),constant(srctype),constant(n+2),constant(n+2),srcv)
    end
    ij = param.ij
    σxx, σyy, σxy,vx, vy  = σxx[ij], σyy[ij], σxy[ij],vx[ij],vy[ij]
    for i = 1:8
        mem[i] = mem[i][ij]
    end
    return σxx, σyy, σxy,vx, vy, mem, mpi_idx
end

function ElasticSourceAtTimeT(src::MPIElasticSource,i::PyObject)
    if ismissing(src.srci)
        return missing, missing, missing, missing
    else
        return src.srci, src.srcj, src.srctype, src.srcv[i]
    end
end

"""
    SimulatedObservation!(ap::ElasticPropagator, rcv::ElasticReceiver)

Extract and save simulated displacement u into rcv::ElasticReceiver.
"""
function MPISimulatedObservation!(ep::MPIElasticPropagator, rcv::MPIElasticReceiver)
    if ismissing(rcv.rcvi)
        op = group([ep.vx, ep.vy, ep.sigmaxx, ep.sigmayy, ep.sigmaxy])
        rcv.rcvv = op
        return
    end

    N = (ep.param.n)*(ep.param.n)
    get_receive = load_op_and_grad("$(@__DIR__)/../deps/CustomOps/build/libADSeismic", "get_receive")
    u = get_receive(tf.reshape(ep.vx, (ep.param.NSTEP+1, N)),
                    tf.reshape(ep.vy, (ep.param.NSTEP+1, N)),
                    tf.reshape(ep.sigmaxx, (ep.param.NSTEP+1, N)),
                    tf.reshape(ep.sigmayy, (ep.param.NSTEP+1, N)),
                    tf.reshape(ep.sigmaxy, (ep.param.NSTEP+1, N)),
                    convert_to_tensor(rcv.rcvi), convert_to_tensor(rcv.rcvj), convert_to_tensor(rcv.rcvtype),
                        constant(ep.param.n-2), constant(ep.param.n-2))
    rcv.rcvv = u
end


import DiffEqBase
export DiffEqJLSolver, DiffEqJLIMEXSolver

abstract type AbstractDiffEqJLSolver <: AbstractODESolver end

"""
    DiffEqJLSolver(f, RKA, RKB, RKC, Q; dt, t0 = 0)

This is a time stepping object for explicitly time stepping the differential
equation given by the right-hand-side function `f` with the state `Q`, i.e.,

```math
  \\dot{Q} = f(Q, t)
```

via a DifferentialEquations.jl DEAlgorithm, which includes support
for OrdinaryDiffEq.jl, Sundials.jl, and more.
"""
mutable struct DiffEqJLSolver{I} <: AbstractDiffEqJLSolver
    integ::I
    steps::Int

    function DiffEqJLSolver(
        rhs!,
        alg,
        Q,
        args...;
        t0 = 0,
        p = nothing,
        kwargs...,
    )
        prob = DiffEqBase.ODEProblem(
            (du, u, p, t) -> rhs!(du, u, p, t; increment = false),
            Q,
            (float(t0), typemax(typeof(float(t0)))),
            p,
        )
        integ = DiffEqBase.init(
            prob,
            alg,
            args...;
            adaptive = false,
            save_everystep = false,
            save_start = false,
            save_end = false,
            kwargs...,
        )
        new{typeof(integ)}(integ, 0)
    end
end

"""
    DiffEqJLSolver(f, RKA, RKB, RKC, Q; dt, t0 = 0)

This is a time stepping object for explicitly time stepping the differential
equation given by the right-hand-side function `f` with the state `Q`, i.e.,

```math
  \\dot{Q} = f_I(Q, t) + f_E(Q, t)
```

via a DifferentialEquations.jl DEAlgorithm, which includes support
for OrdinaryDiffEq.jl, Sundials.jl, and more.
"""
mutable struct DiffEqJLIMEXSolver{I} <: AbstractDiffEqJLSolver
    integ::I
    steps::Int

    function DiffEqJLIMEXSolver(
        rhs!,
        rhs_implicit!,
        alg,
        Q,
        args...;
        t0 = 0,
        p = nothing,
        kwargs...,
    )
        prob = DiffEqBase.SplitODEProblem(
            (du, u, p, t) -> rhs_implicit!(du, u, p, t; increment = false),
            (du, u, p, t) -> rhs!(du, u, p, t; increment = false),
            Q,
            (float(t0), typemax(typeof(float(t0)))),
            p,
        )
        integ = DiffEqBase.init(
            prob,
            alg,
            args...;
            adaptive = false,
            save_everystep = false,
            save_start = false,
            save_end = false,
            kwargs...,
        )

        new{typeof(integ)}(integ, 0)
    end
end

gettime(solver::AbstractDiffEqJLSolver) = solver.integ.t
getdt(solver::AbstractDiffEqJLSolver) = solver.integ.dt
updatedt!(solver::AbstractDiffEqJLSolver, dt) =
    DiffEqBase.set_proposed_dt!(solver.integ, dt)
updatetime!(solver::AbstractDiffEqJLSolver, t) =
    DiffEqBase.set_t!(solver.integ, t)
isadjustable(solver::AbstractDiffEqJLSolver) = true # Is this isadaptive? Or something different?

"""
    ODESolvers.general_dostep!(Q, solver::AbstractODESolver, p,
                               timeend::Real, adjustfinalstep::Bool)

Use the solver to step `Q` forward in time from the current time, to the time
`timeend`. If `adjustfinalstep == true` then `dt` is adjusted so that the step
does not take the solution beyond the `timeend`.
"""
function general_dostep!(
    Q,
    solver::AbstractDiffEqJLSolver,
    p,
    timeend::Real;
    adjustfinalstep::Bool,
)
    integ = solver.integ

    if first(integ.opts.tstops) !== timeend
        DiffEqBase.add_tstop!(integ, timeend)
    end
    dostep!(Q, solver, p, time)
    solver.integ.t
end

function dostep!(
    Q,
    solver::AbstractDiffEqJLSolver,
    p,
    time,
    slow_δ = nothing,
    slow_rv_dQ = nothing,
    in_slow_scaling = nothing,
)

    integ = solver.integ
    integ.p = p # Can this change?

    rv_Q = realview(Q)
    if integ.u != Q
        integ.u .= Q
        DiffEqBase.u_modified!(integ, true)
        # Will time always be correct?
    end

    DiffEqBase.step!(integ)
    rv_Q .= solver.integ.u
end

function DiffEqJLConstructor(alg)
    constructor =
        (F, Q; dt = 0, t0 = 0) -> DiffEqJLSolver(F, alg, Q; t0 = t0, dt = dt)
    return constructor
end