text:
using Dates
using Logging
using JSON

push!(LOAD_PATH, abspath(@__DIR__))
using BenchmarkHelper

log_level =  Logging.Debug # Logging.Info | Logging.Warn | Logging.Error
global_logger(SimpleLogger(stdout, log_level))

# setup global variables
EVENT_FILE = get(ENV, "GITHUB_EVENT_PATH", "")
if isempty(EVENT_FILE)
    @error "No github event file specified."
    exit(0)
end

EVENT_DATA = JSON.Parser.parsefile(EVENT_FILE)
@debug EVENT_DATA

# Extract event information

if EVENT_DATA["action"] != "created" || !haskey(EVENT_DATA, "comment")
    @debug "Not an event for comment creation"
    exit(0)
end

user = EVENT_DATA["comment"]["user"]["login"]
should_run, tags, branches = BenchmarkHelper.benchmarks_info(EVENT_DATA)
job_id =  Dates.format(Dates.now(), "YmmddHHMMSS")

if !should_run
    exit(0)
end

if (branches != nothing && !isempty(branches))
    @info "benchmark target branches:", branches
    body = [
        "Hi @$user, I just got a benchmark command from you,"
        "I will run a benchmark(JOB_ID=$(job_id)) on branches $branches as soon as possible."
    ]
    BenchmarkHelper.reply_comment(EVENT_DATA, join(body, "\n"))
else
    @info "No benchmarking command found in the comment."
    body = [
        "Hi @$user, I just got a benchmark command from you,"
        "But I can't determine on which branches to run,"
        "Please specify the branches in your command."
    ]
    BenchmarkHelper.reply_comment(EVENT_DATA, join(body, "\n"))
    exit(0)
end

# Run benchmarks on current event
BenchmarkHelper.set_benchmark_info("", job_id)
BenchmarkHelper.run_benchmarks(tags, branches)

# Comment the benchmark results
BenchmarkHelper.reply_comment(EVENT_DATA, BenchmarkHelper.get_results(job_id, user))


text:
# AUTO GENERATED FILE - DO NOT EDIT

export 'feffery'_antdmentions

"""
    'feffery'_antdmentions(;kwargs...)

An AntdMentions component.

Keyword arguments:
- `id` (String; optional)
- `autoSize` (optional): . autoSize has the following type: Bool | lists containing elements 'minRows', 'maxRows'.
Those elements have the following types:
  - `minRows` (Real; optional)
  - `maxRows` (Real; optional)
- `className` (String; optional)
- `defaultValue` (String; optional)
- `disabled` (Bool; optional)
- `loading_state` (optional): . loading_state has the following type: lists containing elements 'is_loading', 'prop_name', 'component_name'.
Those elements have the following types:
  - `is_loading` (Bool; optional): Determines if the component is loading or not
  - `prop_name` (String; optional): Holds which property is loading
  - `component_name` (String; optional): Holds the name of the component that is loading
- `options` (optional): . options has the following type: Array of lists containing elements 'label', 'value'.
Those elements have the following types:
  - `label` (String; optional)
  - `value` (String; optional)s
- `placement` (a value equal to: 'top', 'bottom'; optional)
- `prefix` (String; optional)
- `selectedOptions` (Array of Strings; optional)
- `style` (Dict; optional)
- `value` (String; optional)
"""
function 'feffery'_antdmentions(; kwargs...)
        available_props = Symbol[:id, :autoSize, :className, :defaultValue, :disabled, :loading_state, :options, :placement, :prefix, :selectedOptions, :style, :value]
        wild_props = Symbol[]
        return Component("'feffery'_antdmentions", "AntdMentions", "feffery_antd_components", available_props, wild_props; kwargs...)
end

ext: jl

max_stars_repo_name: ReilYoung/feffery-antd-components

max_issues_repo_name: ReilYoung/feffery-antd-components

max_forks_repo_name: ReilYoung/feffery-antd-components


text:
#=make_scenarios.jl
inputs:
number of timesteps in simulation
value options
probabilities for each value
interval length
TF should we pick a subselection of the scenarios
how many scenarios if so

outputs:
vdr[t,o]
p[o]
this is to make the performance of DR in different
time periods independent.
=#

using StatsBase
function make_scenarios(n_timesteps,v_og,p_og,int_length; randsel = true, nrand = 200)

    Tp = convert(Int64,floor(n_timesteps/int_length))
    n_omega = length(p_og)
    n_new_scenarios = convert(Int64,n_omega^Tp)

    if randsel & (nrand < n_new_scenarios)
        s = sample(1:n_new_scenarios,nrand,replace=false)
    else
        s = 1:n_new_scenarios
    end

    #iterate through rows
    prob_array2 = Array{Float64}(length(s),Tp)
    for i in 1:length(s) #row
        for j in 1:Tp #col
            ind = mod(fld(s[i]-1,n_omega^(Tp-j)),n_omega)+1
            prob_array2[i,j] = p_og[ind]
        end
    end

    # can use prod() to get row-wise products.
    ptemp = prod(prob_array2,2)
    p = ptemp/sum(ptemp)
    # p = round.(p, 10)

    vdr = fill(0.0,(n_timesteps,length(s)))
    for j in 1:length(s) #cols
        for i in 1:Tp #row chunk
            # identify which rows belong to this time period
            if i == Tp
                last_period = n_timesteps
            else
                last_period = int_length*i #t_firsts[i+1]-1
            end
            rows = (int_length*(i-1) + 1):last_period

            ind = mod(fld(s[j]-1,n_omega^(Tp-i)),n_omega)+1
            println(ind)
            vdr[rows,j] = v_og[ind]
        end
    end

    return(vdr,p)
end

max_stars_repo_name: pajalevi/Stoch-UC-model-for-DR

max_issues_repo_name: pajalevi/Stoch-UC-model-for-DR

max_forks_repo_name: pajalevi/Stoch-UC-model-for-DR


text:
"""
    NStream{FT,I} <: AbstractOpticalPropsArry{FT,I}

Holds extinction optical depth `tau`, `ssa`, and phase function moments `p` with leading dimension nmom.

# Fields
$(DocStringExtensions.FIELDS)
"""
struct NStream{FT,I} <: AbstractOpticalPropsArry{FT,I}
  base
  tau#::Array{FT,3}
  ssa#::Array{FT,3}
  p#::Array{FT,4}
end

"""
    delta_scale!(...)

    class(NStream), intent(inout) :: this
    real(FT), dimension(:,:,:), optional,
                                 intent(in   ) :: for_
    character(128)                             :: err_message
"""
delta_scale!(this::NStream, for_) = "delta_scale_nstr: Not yet implemented"

"""

  class(NStream), intent(in) :: this
  character(len=128)                       :: err_message

  integer :: varSizes(3)
"""
function validate!(this::NStream{FT}) where FT

  #
  # Array allocation status, sizing
  #
  varSizes =   [size(this.tau, 1), size(this.tau, 2), size(this.tau, 3)]
  if      !all([size(this.ssa, 1), size(this.ssa, 2), size(this.ssa, 3)] == varSizes) ||
          !all([size(this.p,   2), size(this.p,   3), size(this.p,   4)] == varSizes)
    error("validate: arrays not sized consistently")
  end
  #
  # Valid values
  #
  any_vals_less_than(this.tau,  FT(0)) && error("validate: tau values out of range")
  any_vals_outside(this.ssa,  FT(0), FT(1)) && error("validate: ssa values out of range")
  any_vals_outside(this.p[1,:,:,:], FT(-1), FT(1)) && error("validate: p(1,:,:,:)  = g values out of range")
end

# ------------------------------------------------------------------------------------------
#
#  Routines for array classes: subsetting of optical properties arrays along x (col) direction
#
# Allocate class, then arrays; copy. Could probably be more efficient if
#   classes used pointers internally.
#
# This set takes start position and number as scalars
#
# ------------------------------------------------------------------------------------------

"""
  class(OneScalar), intent(inout) :: full
  integer,                      intent(in   ) :: start, n
  class(AbstractOpticalPropsArry), intent(inout) :: subset
  character(128)                              :: err_message

  integer :: ncol, nlay, ngpt, nmom
"""
function subset_range!(full::OneScalar{FT}, start, n, subset) where FT
  ncol = get_ncol(full)
  nlay = get_nlay(full)
  ngpt = get_ngpt(full)
  if start < 1 || start + n-1 > get_ncol(full)
     error("optical_props%subset: Asking for columns outside range")
  end

  init!(subset, full)
  # Seems like the deallocation statements should be needed under Fortran 2003
  #   but Intel compiler doesn't run without them

  if subset isa OneScalar
      alloc!(subset, n, nlay)
  elseif subset isa TwoStream
      alloc!(subset, n, nlay)
      subset.ssa[1:n,:,:] .= FT(0)
      subset.g[1:n,:,:] .= FT(0)
  elseif subset isa NStream
      if allocated(subset.p)
        nmom = get_nmom(subset)
      else
        nmom = 1
      end
      alloc!(subset, nmom, n, nlay)
      subset.ssa[1:n,:,:] .= FT(0)
      subset.p[:,1:n,:,:] .= FT(0)
  else
    error("Uncaught case in subset_range!(full::OneScalar{FT}")
  end
  extract_subset!(ncol, nlay, ngpt, full.tau, start, start+n-1, subset.tau)

end
# ------------------------------------------------------------------------------------------

"""
  subset_range!(...)

  class(TwoStream), intent(inout) :: full
  integer,                      intent(in   ) :: start, n
  class(AbstractOpticalPropsArry), intent(inout) :: subset
  character(128)                              :: err_message

  integer :: ncol, nlay, ngpt, nmom
"""
function subset_range!(full::TwoStream{FT}, start, n, subset) where FT

  ncol = get_ncol(full)
  nlay = get_nlay(full)
  ngpt = get_ngpt(full)
  if start < 1 || start + n-1 > get_ncol(full)
     error("optical_props%subset: Asking for columns outside range")
  end

  init!(subset, full)

  if subset isa OneScalar # TODO: check logic
    alloc!(subset, n, nlay)
    extract_subset!(ncol, nlay, ngpt, full.tau, full.ssa, start, start+n-1, subset.tau)
  elseif subset isa TwoStream
    alloc!(subset, n, nlay)
    extract_subset!(ncol, nlay, ngpt, full.tau, start, start+n-1, subset.tau)
    extract_subset!(ncol, nlay, ngpt, full.ssa, start, start+n-1, subset.ssa)
    extract_subset!(ncol, nlay, ngpt, full.g  , start, start+n-1, subset.g  )
  elseif subset isa NStream
    if allocated(subset.p)
      nmom = get_nmom(subset)
    else
      nmom = 1
    end
    alloc!(subset, nmom, n, nlay)
    extract_subset!(ncol, nlay, ngpt, full.tau, start, start+n-1, subset.tau)
    extract_subset!(ncol, nlay, ngpt, full.ssa, start, start+n-1, subset.ssa)
    subset.p[1,1:n,:,:] .= full.g[start:start+n-1,:,:]
    subset.p[2:end,:, :,:] .= FT(0) # TODO Verify this line
  else
    error("Uncaught case in subset_range!(full::TwoStream{FT}")
  end
end

"""
    subset_range!(full::NStream, start, n, subset)

    class(NStream), intent(inout) :: full
    integer,                      intent(in   ) :: start, n
    class(AbstractOpticalPropsArry), intent(inout) :: subset
    character(128)                              :: err_message

    integer :: ncol, nlay, ngpt, nmom
"""
function subset_range!(full::NStream, start, n, subset)

  ncol = get_ncol(full)
  nlay = get_nlay(full)
  ngpt = get_ngpt(full)
  if(start < 1 || start + n-1 > get_ncol(full))
     error("optical_props%subset: Asking for columns outside range")
  end

  init!(subset, full)

  if subset isa OneScalar # TODO: check logic
    alloc!(subset, n, nlay)
    extract_subset!(ncol, nlay, ngpt, full.tau, full.ssa, start, start+n-1, subset.tau)
  elseif subset isa TwoStream
    alloc!(subset, n, nlay)
    extract_subset!(ncol, nlay, ngpt, full.tau, start, start+n-1, subset.tau)
    extract_subset!(ncol, nlay, ngpt, full.ssa, start, start+n-1, subset.ssa)
    subset.g[1:n,:,:] .= full.p[1,start:start+n-1,:,:]
  elseif subset isa NStream
    alloc!(subset, nmom, n, nlay)
    extract_subset!(      ncol, nlay, ngpt, full.tau, start, start+n-1, subset.tau)
    extract_subset!(      ncol, nlay, ngpt, full.ssa, start, start+n-1, subset.ssa)
    extract_subset!(nmom, ncol, nlay, ngpt, full.p  , start, start+n-1, subset.p  )
  else
    error("Uncaught case in subset_range!(full::NStream, start, n, subset)")
  end
end

"""
    class(NStream), intent(in   ) :: this
    integer                                     :: get_nmom
"""
get_nmom(this::NStream) = size(this.p, 1)

"""
    combine_and_reorder_nstr!(...)

Combine absoprtion and Rayleigh optical depths for total tau, ssa, p
using Rayleigh scattering phase function

integer, intent(in) :: ncol, nlay, ngpt, nmom
real(FT), dimension(ngpt,nlay,ncol), intent(in ) :: tau_abs, tau_rayleigh
real(FT), dimension(ncol,nlay,ngpt), intent(inout) :: tau, ssa
real(FT), dimension(ncol,nlay,ngpt,nmom),
                                     intent(inout) :: p
# -----------------------
integer :: icol, ilay, igpt, imom
real(FT) :: t
# -----------------------
"""
function combine_and_reorder_nstr!(ncol, nlay, ngpt, nmom, tau_abs, tau_rayleigh, tau, ssa, p)
  FT = eltype(tau_abs)

  for icol in 1:ncol
    for ilay in 1:nlay
      for igpt in 1:ngpt
        t = tau_abs[igpt,ilay,icol] + tau_rayleigh[igpt,ilay,icol]
        tau[icol,ilay,igpt] = t
        if (t > FT(2) * realmin(FT))
          ssa[icol,ilay,igpt] = tau_rayleigh[igpt,ilay,icol] / t
        else
          ssa[icol,ilay,igpt] = FT(0)
        end
        for imom = 1:nmom
          p[imom,icol,ilay,igpt] = FT(0)
        end
        if (nmom >= 2)
          p[2,icol,ilay,igpt] = FT(0.1)
        end
      end
    end
  end
end


text:
"""
    supports_in_set(supports::Union{Number, Vector{<:Number}},
                    set::AbstractInfiniteSet)::Bool

Used to check if `supports` are in the domain of `set`. Returns `true` if
`supports` are in domain of `set` and returns `false` otherwise.
This is primarily an internal method for performing checks but can be extended
for user-defined set types. Extending this is optional, but recommended where
possible. Note by fallback, this returns `true` for unrecognized set types such
that an error won't be thrown.
"""
function supports_in_set end

# IntervalSet
function supports_in_set(supports::Union{Number, Vector{<:Number}},
                         set::IntervalSet)::Bool
    min_support = minimum(supports)
    max_support = maximum(supports)
    if min_support < set.lower_bound || max_support > set.upper_bound
        return false
    end
    return true
end

# UnivariateDistribution set
function supports_in_set(supports::Union{Number, Vector{<:Number}},
                         set::DistributionSet{<:Distributions.UnivariateDistribution})::Bool
    min_support = minimum(supports)
    max_support = maximum(supports)
    check1 = min_support < minimum(set.distribution)
    check2 = max_support > maximum(set.distribution)
    if check1 || check2
        return false
    end
    return true
end

# Fallback
function supports_in_set(supports::Union{Number, Vector{<:Number}},
                         set::AbstractInfiniteSet)::Bool
    return true
end

"""
    JuMP.has_lower_bound(set::AbstractInfiniteSet)::Bool

Return `Bool` indicating if `set` has a lower bound that can be determined. This
should be extended for user-defined infinite sets. It defaults to `false`
for unrecognized set types.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP)
julia> set = InfiniteSet(0, 1);

julia> has_lower_bound(set)
true
```
"""
function JuMP.has_lower_bound(set::AbstractInfiniteSet)::Bool # fallback
    return false
end

# IntervalSet
JuMP.has_lower_bound(set::IntervalSet)::Bool = true

# DistributionSet (Univariate)
JuMP.has_lower_bound(set::DistributionSet{<:Distributions.UnivariateDistribution})::Bool = true

"""
    JuMP.lower_bound(set::AbstractInfiniteSet)::Number

Return the lower bound of `set` if one exists. This should be extended for
user-defined infinite sets if appropriate. Errors if `JuMP.has_lower_bound`
returns `false`. Extensions are enabled by `JuMP.has_lower_bound(set)` and
`JuMP.lower_bound(set)`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP)
julia> set = InfiniteSet(0, 1);

julia> lower_bound(set)
0.0
```
"""
function JuMP.lower_bound(set::AbstractInfiniteSet) # fallback
    type = typeof(set)
    error("`JuMP.lower_bound` not defined for infinite set of type $type.")
end

# IntervalSet
JuMP.lower_bound(set::IntervalSet)::Number = set.lower_bound

# DistributionSet (Univariate)
function JuMP.lower_bound(set::DistributionSet{<:Distributions.UnivariateDistribution})::Number
    return minimum(set.distribution)
end

"""
    JuMP.set_lower_bound(set::AbstractInfiniteSet,
                         lower::Number)::AbstractInfiniteSet

Set and return the lower bound of `set` if such an aoperation makes sense. Errors
if the type of `set` does not support this operation or has not been extended.
User-defined set types should extend this if appropriate.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP)
julia> set = InfiniteSet(0, 1);

julia> set_lower_bound(set, 0.5)
[0.5, 1]
```
"""
function JuMP.set_lower_bound(set::AbstractInfiniteSet, lower::Number) # fallback
    type = typeof(set)
    error("`JuMP.set_lower_bound` not defined for infinite set of type $type.")
end

# IntervalSet
function JuMP.set_lower_bound(set::IntervalSet, lower::Number)::IntervalSet
    return IntervalSet(lower, set.upper_bound)
end

# DistributionSet
function JuMP.set_lower_bound(set::DistributionSet, lower::Number)
    error("Cannot set the lower bound of a distribution, try using " *
          "`Distributions.Truncated` instead.")
end

"""
    JuMP.has_upper_bound(set::AbstractInfiniteSet)::Bool

Return `Bool` indicating if `set` has a upper bound that can be determined. This
should be extended for user-defined infinite sets. It defaults to `false`
for unrecognized set types.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP)
julia> set = InfiniteSet(0, 1);

julia> has_upper_bound(set)
true
```
"""
function JuMP.has_upper_bound(set::AbstractInfiniteSet)::Bool # fallback
    return false
end

# IntervalSet
JuMP.has_upper_bound(set::IntervalSet)::Bool = true

# DistributionSet (Univariate)
JuMP.has_upper_bound(set::DistributionSet{<:Distributions.UnivariateDistribution})::Bool = true

"""
    JuMP.upper_bound(set::AbstractInfiniteSet)::Number

Return the upper bound of `set` if one exists. This should be extended for
user-defined infinite sets if appropriate. Errors if `JuMP.has_upper_bound`
returns `false`. Extensions are enabled by `JuMP.has_upper_bound(set)` and
`JuMP.upper_bound(set)`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP)
julia> set = InfiniteSet(0, 1);

julia> upper_bound(set)
1.0
```
"""
function JuMP.upper_bound(set::AbstractInfiniteSet) # fallback
    type = typeof(set)
    error("`JuMP.upper_bound` not defined for infinite set of type $type.")
end

# IntervalSet
JuMP.upper_bound(set::IntervalSet)::Number = set.upper_bound

# DistributionSet (Univariate)
function JuMP.upper_bound(set::DistributionSet{<:Distributions.UnivariateDistribution})::Number
    return maximum(set.distribution)
end

"""
    JuMP.set_upper_bound(set::AbstractInfiniteSet,
                         upper::Number)::AbstractInfiniteSet

Set and return the upper bound of `set` if such an aoperation makes sense. Errors
if the type of `set` does not support this operation or has not been extended.
User-defined set types should extend this if appropriate.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP)
julia> set = InfiniteSet(0, 1);

julia> set_upper_bound(set, 0.5)
[0, 0.5]
```
"""
function JuMP.set_upper_bound(set::AbstractInfiniteSet, upper::Number) # fallback
    type = typeof(set)
    error("`JuMP.set_lower_bound` not defined for infinite set of type $type.")
end

# IntervalSet
function JuMP.set_upper_bound(set::IntervalSet, upper::Number)
    return IntervalSet(set.lower_bound, upper)
end

# DistributionSet
function JuMP.set_upper_bound(set::DistributionSet, lower::Number)
    error("Cannot set the upper bound of a distribution, try using " *
          "`Distributions.Truncated` instead.")
end

"""
    generate_support_values(set::AbstractInfiniteSet; [num_supports::Int = 10,
                            sig_fig::Int = 5])

Generate `num_supports` support values with `sig_figs` significant digits in
accordance with `set` and return them. `IntervalSet`s generate supports
uniformly and `DistributionSet`s generate them randomly accordingly to the
underlyign distribution. Extensions that employ user-defined infinite set types
should extend according to the new set type. Errors if the `set` is a type that
has not been explicitly extended. This is intended as an internal method to be
used by [`generate_and_add_supports!`](@ref) and [`build_parameter`](@ref).
"""
function generate_support_values(set::AbstractInfiniteSet; num_supports::Int = 0,
                                 sig_fig::Int = 1)
    type = typeof(set)
    error("Unable to generate support values for unrecognized infinite set " *
          "type $type")
end

# IntervalSet
function generate_support_values(set::IntervalSet; num_supports::Int = 10,
                         sig_fig::Int = 5)::Vector
    lb = set.lower_bound
    ub = set.upper_bound
    new_supports = round.(collect(range(lb, stop = ub, length = num_supports)),
                          sigdigits = sig_fig)
    return new_supports
end

# DistributionSet
function generate_support_values(set::DistributionSet; num_supports::Int = 10,
                         sig_fig::Int = 5)::Array
    new_supports = round.(Distributions.rand(set.distribution, num_supports),
                          sigdigits = sig_fig)
    return new_supports
end


text:
configs = FunbArray([
    Funb( sin, [(-1,1)],[Float32, Float64], [10] );
    Funb( atan, [(-1,1), (-1,1)],[Float32, Float64],[10, 10] );
    Funb( *, [(-1, 1), (-1, 1), (-1, 1)], [Float32, Float64], [(10,10), (10,10)] );
    ])

benchmark!(configs)

# Test IO
AcuteBenchmark.save("test.jld2", configs)
configs_loaded = AcuteBenchmark.load("test.jld2")

bar(configs)

bar(configs, uniqueType=true)

bar(configs => configs, uniqueType=true)

bar(configs => configs, uniqueType=true, dimAnnotation=false, uniqueDim=true)

configs2 = Funb( sin, [(-1,1)],[Float32, Float64], [10 30 50 100 200 500] );

benchmark!(configs2)

dimplot(configs2)


text:
export DualOptimizer, dual_optimizer

dual_optimizer(optimizer_constructor) = () -> DualOptimizer(MOI.instantiate(optimizer_constructor))

struct DualOptimizer{T, OT <: MOI.ModelLike} <: MOI.AbstractOptimizer
    dual_problem::DualProblem{T, OT}

    function DualOptimizer{T, OT}(dual_problem::DualProblem{T, OT}) where {T, OT <: MOI.ModelLike}
        return new{T, OT}(dual_problem)
    end
end

"""
    DualOptimizer(dual_optimizer::OT) where {OT <: MOI.ModelLike}

The DualOptimizer finds the solution for a problem by solving its dual representation.
It builds the dual model internally and solve it using the `dual_optimizer` as solver.
Primal results are obtained by querying dual results of the internal problem solved
by `dual_optimizer`. Analogously, dual results are obtained by querying primal results
of the internal problem.

The user can define the model providing the `DualOptimizer` and the solver of its choice.

Example:

```julia
julia> using Dualization, JuMP, GLPK

julia> model = Model(dual_optimizer(GLPK.Optimizer))
A JuMP Model
Feasibility problem with:
Variables: 0
Model mode: AUTOMATIC
CachingOptimizer state: EMPTY_OPTIMIZER
Solver name: Dual model with GLPK attached
```
"""
function DualOptimizer(dual_optimizer::OT) where {OT <: MOI.ModelLike}
    return DualOptimizer{Float64}(dual_optimizer)
end

function DualOptimizer{T}(dual_optimizer::OT) where {T, OT <: MOI.ModelLike}
    dual_problem = DualProblem{T}(MOIB.full_bridge_optimizer(MOIU.CachingOptimizer(MOIU.UniversalFallback(DualizableModel{T}()), dual_optimizer), T))
    # discover the type of MOIU.CachingOptimizer(DualizableModel{T}(), dual_optimizer)
    OptimizerType = typeof(dual_problem.dual_model)
    return DualOptimizer{T, OptimizerType}(dual_problem)
end

function DualOptimizer()
    return error("DualOptimizer must have a solver attached")
end

function MOI.supports(::DualOptimizer,
                      ::MOI.ObjectiveSense)
    return true
end
function MOI.supports(optimizer::DualOptimizer{T},
                      ::MOI.ObjectiveFunction{F}) where {T, F}
    # If the objective function is `MOI.SingleVariable` or `MOI.ScalarAffineFunction`,
    # a `MOI.ScalarAffineFunction` is set as objective function for the dual problem.
    # If it is `MOI.ScalarQuadraticFunction` , a `MOI.ScalarQuadraticFunction` is set as objective function for the dual problem.
    G = F <: MOI.ScalarQuadraticFunction ? MOI.ScalarQuadraticFunction{T} : MOI.ScalarAffineFunction{T}
    return supported_obj(F) && MOI.supports(optimizer.dual_problem.dual_model, MOI.ObjectiveFunction{G}())
end

function MOI.supports_constraint(
    optimizer::DualOptimizer{T},
    F::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}},
    S::Type{<:MOI.AbstractScalarSet}) where T
    D = _dual_set_type(S)
    if D === nothing
        return false
    end
    if D <: MOI.AbstractVectorSet # The dual of `EqualTo` is `Reals`
        return MOI.supports_add_constrained_variables(optimizer.dual_problem.dual_model, D)
    else
        return MOI.supports_add_constrained_variable(optimizer.dual_problem.dual_model, D)
    end
end

function MOI.supports_constraint(
    optimizer::DualOptimizer{T},
    F::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}},
    S::Type{<:MOI.AbstractVectorSet}) where T
    D = _dual_set_type(S)
    if D === nothing
        return false
    end
    return MOI.supports_add_constrained_variables(optimizer.dual_problem.dual_model, D)
end

# TODO add this when constrained variables are implemented
#function MOI.supports_add_constrained_variables(
#    optimizer::DualOptimizer{T}, S::Type{MOI.Reals}) where T
#    return MOI.supports_constraint(optimizer.dual_problem.dual_model,
#                                   MOI.ScalarAffineFunction{T},
#                                   MOI.EqualTo{T}) # If it was `MOI.Zeros`, we would not need this method as special case of the one below
#end
#function MOI.supports_add_constrained_variables(
#    optimizer::DualOptimizer{T}, S::Type{<:MOI.AbstractVectorSet}) where T
#    D = _dual_set_type(S)
#    if D === nothing
#        return false
#    end
#    return MOI.supports_constraint(optimizer.dual_problem.dual_model,
#                                   MOI.VectorAffineFunction{T}, D)
#end

function MOI.copy_to(dest::DualOptimizer, src::MOI.ModelLike; kwargs...)
    # Dualize the original problem
    dualize(src, dest.dual_problem)

    # Identity IndexMap
    idx_map = MOIU.IndexMap()

    for vi in MOI.get(src, MOI.ListOfVariableIndices())
        setindex!(idx_map, vi, vi)
    end

    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        for con in MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
            setindex!(idx_map, con, con)
        end
    end
    return idx_map
end

function MOI.optimize!(optimizer::DualOptimizer)
    return MOI.optimize!(optimizer.dual_problem.dual_model)
end

function MOI.is_empty(optimizer::DualOptimizer)
    return (MOI.is_empty(optimizer.dual_problem.dual_model)) && is_empty(optimizer.dual_problem.primal_dual_map)
end

function MOI.empty!(optimizer::DualOptimizer)
    MOI.empty!(optimizer.dual_problem.dual_model)
    empty!(optimizer.dual_problem.primal_dual_map)
    return
end

# MOI.get auxiliary functions
function get_ci_dual_problem(optimizer::DualOptimizer, vi::VI)
    return optimizer.dual_problem.primal_dual_map.primal_var_dual_con[vi]
end

function get_ci_dual_problem(optimizer::DualOptimizer, ci::CI)
    return optimizer.dual_problem.primal_dual_map.primal_con_dual_con[ci]
end

function get_primal_ci_constant(optimizer::DualOptimizer, ci::CI)
    return first(get_primal_ci_constants(optimizer, ci))
end

function get_primal_ci_constants(optimizer::DualOptimizer, ci::CI)
    return optimizer.dual_problem.primal_dual_map.primal_con_constants[ci]
end

function get_vi_dual_problem(optimizer::DualOptimizer, ci::CI)
    return first(get_vis_dual_problem(optimizer, ci))
end

function get_vis_dual_problem(optimizer::DualOptimizer, ci::CI)
    return optimizer.dual_problem.primal_dual_map.primal_con_dual_var[ci]
end

function MOI.get(optimizer::DualOptimizer, ::MOI.SolverName)
    return "Dual model with "*MOI.get(optimizer.dual_problem.dual_model, MOI.SolverName()) * " attached"
end

function MOI.get(optimizer::DualOptimizer, ::MOI.VariablePrimal, vi::VI)
    return -MOI.get(optimizer.dual_problem.dual_model,
                    MOI.ConstraintDual(), get_ci_dual_problem(optimizer, vi))
end

function MOI.get(optimizer::DualOptimizer, ::MOI.ConstraintDual,
                 ci::CI{F,S}) where {F <: MOI.AbstractScalarFunction, S <: MOI.AbstractScalarSet}
    return MOI.get(optimizer.dual_problem.dual_model,
                   MOI.VariablePrimal(), get_vi_dual_problem(optimizer, ci))
end

function MOI.get(optimizer::DualOptimizer, ::MOI.ConstraintDual,
                 ci::CI{F,S}) where {F <: MOI.AbstractVectorFunction, S <: MOI.AbstractVectorSet}
    return MOI.get.(optimizer.dual_problem.dual_model,
                    MOI.VariablePrimal(), get_vis_dual_problem(optimizer, ci))
end

function MOI.get(optimizer::DualOptimizer, ::MOI.ConstraintPrimal,
                 ci::CI{F,S}) where {F <: MOI.AbstractScalarFunction, S <: MOI.AbstractScalarSet}
    primal_ci_constant = get_primal_ci_constant(optimizer, ci)
    # If it has no key than there is no dual constraint
    if !haskey(optimizer.dual_problem.primal_dual_map.primal_con_dual_con, ci)
        return -primal_ci_constant
    end
    ci_dual_problem = get_ci_dual_problem(optimizer, ci)
    return MOI.get(optimizer.dual_problem.dual_model, MOI.ConstraintDual(), ci_dual_problem) - primal_ci_constant
end

function MOI.get(optimizer::DualOptimizer{T}, ::MOI.ConstraintPrimal,
                 ci::CI{F,S}) where {T, F <: MOI.AbstractVectorFunction, S <: MOI.AbstractVectorSet}
    # If it has no key than there is no dual constraint
    if !haskey(optimizer.dual_problem.primal_dual_map.primal_con_dual_con, ci)
        # The number of dual variable associated with the primal constraint is the ci dimension
        ci_dimension = length(get_vis_dual_problem(optimizer, ci))
        return zeros(T, ci_dimension)
    end
    ci_dual_problem = get_ci_dual_problem(optimizer, ci)
    return MOI.get(optimizer.dual_problem.dual_model, MOI.ConstraintDual(), ci_dual_problem)
end

function MOI.get(optimizer::DualOptimizer, ::MOI.TerminationStatus)
    return dual_status(MOI.get(optimizer.dual_problem.dual_model, MOI.TerminationStatus()))
end

function dual_status(term::MOI.TerminationStatusCode)
    if term == MOI.INFEASIBLE
        return MOI.DUAL_INFEASIBLE
    elseif term == MOI.DUAL_INFEASIBLE
        return MOI.INFEASIBLE
    elseif term == MOI.ALMOST_INFEASIBLE
        return MOI.ALMOST_DUAL_INFEASIBLE
    elseif term == MOI.ALMOST_DUAL_INFEASIBLE
        return MOI.ALMOST_INFEASIBLE
    end
    return term
end

function MOI.get(optimizer::DualOptimizer, ::MOI.ObjectiveValue)
    return MOI.get(optimizer.dual_problem.dual_model, MOI.DualObjectiveValue())
end

function MOI.get(optimizer::DualOptimizer, ::MOI.DualObjectiveValue)
    return MOI.get(optimizer.dual_problem.dual_model, MOI.ObjectiveValue())
end

function MOI.get(optimizer::DualOptimizer, ::MOI.PrimalStatus)
    return MOI.get(optimizer.dual_problem.dual_model, MOI.DualStatus())
end

function MOI.get(optimizer::DualOptimizer, ::MOI.DualStatus)
    return MOI.get(optimizer.dual_problem.dual_model, MOI.PrimalStatus())
end

function MOI.set(optimizer::DualOptimizer, attr::MOI.AbstractOptimizerAttribute, value)
    return MOI.set(optimizer.dual_problem.dual_model, attr, value)
end
function MOI.get(optimizer::DualOptimizer, attr::Union{MOI.AbstractModelAttribute, MOI.AbstractOptimizerAttribute})
    return MOI.get(optimizer.dual_problem.dual_model, attr)
end


text:
using JLD
using SharedArrays
using Distributed
include("../examples/solenoid.jl")
function obj_fun(u)
	return u[1]^2 + u[2]^2
end
function obj_fun_erg_avg(s)
	nSteps = 10000
	nRunup = 500
	u = rand(3)
	J = 0.
	for i = 1:nRunup
		u = next(u,s)
	end
	for i = 1:nSteps
		J += obj_fun(u)/nSteps
		u = next(u,s)
	end
	return J
end
function get_Javg_vs_s(ind)
	s = [1.0, 4.0, 0]
	n_pts = 20
	n_rep = 16000
	s_ind = LinRange(1.,2.0,n_pts)
	J = zeros(n_pts)
	J_proc = SharedArray{Float64}(n_rep)
	for i = 1:n_pts
		@show s_ind[i]
		s[ind] = s_ind[i]
		#s[3] = s_ind[i]
		J_proc .= 0.
		t = @distributed for n=1:n_rep
			J_proc[n] = obj_fun_erg_avg(s)/n_rep
		end
		wait(t)
		J[i] = sum(J_proc)
	end
	save("../data/obj_erg_avg/solenoid/r2_s$(ind)_sens.jld",
		 "s$ind", s_ind,
		"J", J)
end

max_stars_repo_name: nishaChandramoorthy/linearResponse

max_issues_repo_name: nishaChandramoorthy/linearResponse

max_forks_repo_name: nishaChandramoorthy/linearResponse


text:
import Statistics

line = read("input.txt", String)
numbers = parse.(Int, split(line, ','))

function sumto(n)
    return (n * (n + 1)) ÷ 2
end

# argmin x ∑i=1-N (∑j=1-|s_i−x| j) = 1/2n ∑i=1-N -sign(x_i-x) - mean(x)
position = convert(Int, floor(Statistics.mean(numbers)))
println(sum(@. sumto(abs(numbers - position))))


text:
module TracerAdvDiff

export
  ConstantDiffProblem,
  set_c!,
  updatevars!

using
  FourierFlows,
  FFTW

using LinearAlgebra: mul!, ldiv!

abstract type AbstractTracerParams <: AbstractParams end
abstract type AbstractConstDiffParams <: AbstractParams end
abstract type AbstractSteadyFlowParams <: AbstractParams end

# --
# Problems
# --

"""
    Problem(; parameters...)

Construct a constant diffusivity problem with steady or time-varying flow.
"""
noflow(args...) = 0.0 # used as defaults for u, v functions in Problem()

Problem(; kwargs...) = ConstDiffProblem(; kwargs...) # only problem defined for now

function flowargs(u)
  umethods = methods(u)
  length(umethods.ms) == 1 ? umethods.ms[1].nargs-1 : error("Either define steadyflow or use a flow function
    with only one argument.")
end

function ConstDiffProblem(;
    nx = 128,
    Lx = 2π,
    ny = nx,
    Ly = Lx,
  grid = TwoDGrid(nx, Lx, ny, Ly),
   kap = 0.1,
   eta = kap,
     u = noflow,
     v = noflow,
    dt = 0.01,
  stepper = "RK4",
  steadyflow = nothing
  )

  if steadyflow==nothing # deduce whether flow is time-dependent
    if (typeof(u) <: Function && flowargs(u) > 2) || (typeof(v) <: Function && flowargs(v) > 2)
      steadyflow = false
    else
      steadyflow = true
    end
  end

  if steadyflow; pr = TracerAdvDiff.ConstDiffSteadyFlowParams(eta, kap, u, v, grid)
  else;          pr = TracerAdvDiff.ConstDiffParams(eta, kap, u, v)
  end

  vs = TracerAdvDiff.Vars(grid)
  eq = TracerAdvDiff.Equation(pr, grid)
  ts = FourierFlows.autoconstructtimestepper(stepper, dt, eq.LC, grid)

  FourierFlows.Problem(grid, vs, pr, eq, ts)
end

# --
# Params
# --

"""
    ConstDiffParams(eta, kap, kaph, nkaph, u, v)
    ConstDiffParams(eta, kap, u, v)

Returns the params for constant diffusivity problem with time-varying flow.
"""
struct ConstDiffParams{T} <: AbstractConstDiffParams
  eta::T                 # Constant isotropic horizontal diffusivity
  kap::T                 # Constant isotropic vertical diffusivity
  kaph::T                # Constant isotropic hyperdiffusivity
  nkaph::Int             # Constant isotropic hyperdiffusivity order
  u::Function            # Advecting x-velocity
  v::Function            # Advecting y-velocity
end
ConstDiffParams(eta, kap, u, v) = ConstDiffParams(eta, kap, 0eta, 0, u, v)

"""
    ConstDiffSteadyFlowParams(eta, kap, kaph, nkaph, u, v, g)
    ConstDiffSteadyFlowParams(eta, kap, u, v, g)

Returns the params for constant diffusivity problem with time-steady flow.
"""
struct ConstDiffSteadyFlowParams{T} <: AbstractSteadyFlowParams
  eta::T                 # Constant horizontal diffusivity
  kap::T                 # Constant vertical diffusivity
  kaph::T                # Constant isotropic hyperdiffusivity
  nkaph::Int             # Constant isotropic hyperdiffusivity order
  u::Array{T,2}          # Advecting x-velocity
  v::Array{T,2}          # Advecting y-velocity
end

function ConstDiffSteadyFlowParams(eta, kap, kaph, nkaph, u, v, g)
  if typeof(u) <: Function; ugrid = u.(g.X, g.Y)
  else;                     ugrid = u
  end
  if typeof(v) <: Function; vgrid = v.(g.X, g.Y)
  else;                     vgrid = v
  end
  ConstDiffSteadyFlowParams(eta, kap, kaph, nkaph, ugrid, vgrid)
end

ConstDiffSteadyFlowParams(eta, kap, u, v, g) = ConstDiffSteadyFlowParams(eta, kap, 0eta, 0, u, v, g)

# --
# Equations
# --

"""
    Equation(p, g)

Returns the equation for constant diffusivity problem with params p and grid g.
"""
function Equation(p::ConstDiffParams, g)
  LC = zero(g.Kr)
  @. LC = -p.eta*g.kr^2 - p.kap*g.l^2 - p.kaph*g.KKrsq^p.nkaph
  FourierFlows.Equation{typeof(LC[1, 1]),2}(LC, calcN!)
end

function Equation(p::ConstDiffSteadyFlowParams, g)
  LC = zero(g.Kr)
  @. LC = -p.eta*g.kr^2 - p.kap*g.l^2 - p.kaph*g.KKrsq^p.nkaph
  FourierFlows.Equation{typeof(LC[1, 1]),2}(LC, calcN_steadyflow!)
end

# --
# Vars
# --

# Construct Vars types
 physicalvars = [:c, :cx, :cy]
transformvars = [:ch, :cxh, :cyh]

eval(FourierFlows.structvarsexpr(:Vars, physicalvars, transformvars))

"""
    Vars(g)

Returns the vars for constant diffusivity problem on grid g.
"""
function Vars(g)
  @createarrays typeof(g.Lx) (g.nx, g.ny) c cx cy
  @createarrays Complex{typeof(g.Lx)} (g.nkr, g.nl) ch cxh cyh
  Vars(c, cx, cy, ch, cxh, cyh)
end

# --
# Solvers
# --

"""
    calcN!(N, sol, t, s, v, p, g)

Calculate the advective terms for a tracer equation with constant diffusivity and time-varying flow.
"""
function calcN!(N, sol, t, s, v, p::AbstractConstDiffParams, g)
  @. v.cxh = im*g.kr*sol
  @. v.cyh = im*g.l*sol

  ldiv!(v.cx, g.rfftplan, v.cxh) # destroys v.cxh when using fftw
  ldiv!(v.cy, g.rfftplan, v.cyh) # destroys v.cyh when using fftw

  @. v.cx = -p.u(g.X, g.Y, s.t)*v.cx - p.v(g.X, g.Y, s.t)*v.cy # copies over v.cx so v.cx = N in physical space
  mul!(N, g.rfftplan, v.cx)
  nothing
end

"""
    calcN_steadyflow!(N, sol, t, s, v, p, g)

Calculate the advective terms for a tracer equation with constant diffusivity and time-constant flow.
"""
function calcN_steadyflow!(N, sol, t, s, v, p::AbstractSteadyFlowParams, g)
  @. v.cxh = im*g.kr*sol
  @. v.cyh = im*g.l*sol

  ldiv!(v.cx, g.rfftplan, v.cxh) # destroys v.cxh when using fftw
  ldiv!(v.cy, g.rfftplan, v.cyh) # destroys v.cyh when using fftw

  @. v.cx = -p.u*v.cx - p.v*v.cy # copies over v.cx so v.cx = N in physical space
  mul!(N, g.rfftplan, v.cx)
  nothing
end

# --
# Helper functions
# --

"""
    updatevars!(v, s, g)

Update the vars in v on the grid g with the solution in s.sol.
"""
function updatevars!(s, v, g)
  v.ch .= s.sol
  ch1 = deepcopy(v.ch)
  ldiv!(v.c, g.rfftplan, ch1)
  nothing
end

updatevars!(prob) = updatevars!(prob.state, prob.vars, prob.grid)

"""
    set_c!(s, v, g, c)
    set_c!(s, v, g, c::Function)
    set_c!(prob, c)

Set the solution s.sol as the transform of c and update variables v
on the grid g.
"""
function set_c!(s, v, g, c)
  mul!(s.sol, g.rfftplan, c)
  updatevars!(s, v, g)
  nothing
end

function set_c!(s, v, g, c::Function)
  cgrid = c.(g.X, g.Y)
  mul!(s.sol, g.rfftplan, cgrid)
  updatevars!(s, v, g)
  nothing
end

set_c!(prob, c) = set_c!(prob.state, prob.vars, prob.grid, c)

end # module


text:
# Local Tangent Space Alignment (LTSA)
# ---------------------------
# Principal Manifolds and Nonlinear Dimension Reduction via Local Tangent Space Alignment,
# Zhang, Zhenyue; Hongyuan Zha (2004),  SIAM Journal on Scientific Computing 26 (1): 313–338.
# doi:10.1137/s1064827502419154.

#### LTSA type
struct LTSA{T <: Real} <: AbstractDimensionalityReduction
    k::Int
    λ::AbstractVector{T}
    proj::Projection{T}

    LTSA{T}(k::Int, λ::AbstractVector{T}, proj::Projection{T}) where T = new(k, λ, proj)
end

## properties
outdim(R::LTSA) = size(R.proj, 1)
eigvals(R::LTSA) = R.λ
neighbors(R::LTSA) = R.k

## show
summary(io::IO, R::LTSA) = print(io, "LTSA(outdim = $(outdim(R)), neighbors = $(neighbors(R)))")

## interface functions
function fit(::Type{LTSA}, X::AbstractMatrix{T}; maxoutdim::Int=2, k::Int=12) where {T<:Real}
    n = size(X, 2)

    # Construct NN graph
    D, E = find_nn(X, k)
    S = ones(k)./sqrt(k)
    B = spzeros(T, n,n)
    for i=1:n
        II = @view E[:,i]

        # re-center points in neighborhood
        μ = mean(X[:, II], dims=2)
        δ_x = X[:, II] .- μ

        # Compute orthogonal basis H of θ'
        θ_t = svd(δ_x).V[:,1:maxoutdim]

        # Construct alignment matrix
        G = hcat(S, θ_t)
        B[II, II] .+= diagm(0 => fill(one(T), k)) .- G*transpose(G)
    end

    # Align global coordinates
    λ, V = decompose(B, maxoutdim)
    return LTSA{T}(k, λ, transpose(V))
end

transform(R::LTSA) = R.proj


text:
"""
    constructtestfields_2layer(gr)

Constructs flow fields for a 2-layer problem with parameters such that
    q1 = Δψ1 + 25*(ψ2-ψ1),
    q2 = Δψ2 + 25/4*(ψ1-ψ2).
"""
function constructtestfields_2layer(gr)
  x, y = gridpoints(gr)
  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers

  # a set of streafunctions ψ1 and ψ2, ...
  ψ1 = @. 1e-3 * ( 1/4*cos(2k0*x)*cos(5l0*y) + 1/3*cos(3k0*x)*cos(3l0*y) )
  ψ2 = @. 1e-3 * (     cos(3k0*x)*cos(4l0*y) + 1/2*cos(4k0*x)*cos(2l0*y) )

  # ... their corresponding PVs q1, q2,
  q1 = @. 1e-3 * ( 1/6 *(   75cos(4k0*x)*cos(2l0*y) +   2cos(3k0*x)*( -43cos(3l0*y) + 75cos(4l0*y) ) - 81cos(2k0*x)*cos(5l0*y) ) )
  q2 = @. 1e-3 * ( 1/48*( -630cos(4k0*x)*cos(2l0*y) + 100cos(3k0*x)*(    cos(3l0*y) - 15cos(4l0*y) ) + 75cos(2k0*x)*cos(5l0*y) ) )

  # ... and various derived fields, e.g., ∂ψ1/∂x,
  ψ1x = @. 1e-3 * (     -k0/2*sin(2k0*x)*cos(5l0*y) -       k0*sin(3k0*x)*cos(3l0*y) )
  ψ2x = @. 1e-3 * (      -3k0*sin(3k0*x)*cos(4l0*y) -      2k0*sin(4k0*x)*cos(2l0*y) )
  Δψ2 = @. 1e-3 * ( -25*k0*l0*cos(3k0*x)*cos(4l0*y) - 10*k0*l0*cos(4k0*x)*cos(2l0*y) )

  q1x = @. 1e-3 * ( 1/6 *( -4k0*75sin(4k0*x)*cos(2l0*y) - 3k0*2sin(3k0*x)*( -43cos(3l0*y) + 75cos(4l0*y) ) + 2k0*81sin(2k0*x)*cos(5l0*y) ) )
  q2x = @. 1e-3 * ( 1/48*( 4k0*630sin(4k0*x)*cos(2l0*y) - 3k0*100sin(3k0*x)*(    cos(3l0*y) - 15cos(4l0*y) ) - 2k0*75sin(2k0*x)*cos(5l0*y) ) )

  Δq1 = @. 1e-3 * (k0*l0)*( 1/6 *( -20* 75cos(4k0*x)*cos(2l0*y) +   2cos(3k0*x)*( +18*43cos(3l0*y) - 25*75cos(4l0*y) ) +29*81cos(2k0*x)*cos(5l0*y) ) )
  Δq2 = @. 1e-3 * (k0*l0)*( 1/48*( +20*630cos(4k0*x)*cos(2l0*y) + 100cos(3k0*x)*(    -18cos(3l0*y) + 25*15cos(4l0*y) ) -29*75cos(2k0*x)*cos(5l0*y) ) )

  return ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2
end

"""
    constructtestfields_3layer(gr)

Constructs flow fields for a 3-layer problem with parameters such that
    q1 = Δψ1 + 20ψ2 - 20ψ1,
    q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3,
    q2 = Δψ2 + 12ψ2 - 12ψ3.
"""
function constructtestfields_3layer(gr)
  x, y = gridpoints(gr)
  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers

  # a set of streafunctions ψ1, ψ2, ψ3, ...
  ψ1 = @. 1e-3 * ( 1/4*cos(2k0*x)*cos(5l0*y) + 1/3*cos(3k0*x)*cos(3l0*y) )
  ψ2 = @. 1e-3 * (     cos(3k0*x)*cos(4l0*y) + 1/2*cos(4k0*x)*cos(2l0*y) )
  ψ3 = @. 1e-3 * (     cos(1k0*x)*cos(3l0*y) + 1/2*cos(2k0*x)*cos(2l0*y) )

  Δψ1 = @. -1e-3 * ( 1/4*((2k0)^2+(5l0)^2)*cos(2k0*x)*cos(5l0*y) + 1/3*((3k0)^2+(3l0)^2)*cos(3k0*x)*cos(3l0*y) )
  Δψ2 = @. -1e-3 * (     ((3k0)^2+(4l0)^2)*cos(3k0*x)*cos(4l0*y) + 1/2*((4k0)^2+(2l0)^2)*cos(4k0*x)*cos(2l0*y) )
  Δψ3 = @. -1e-3 * (     ((1k0)^2+(3l0)^2)*cos(1k0*x)*cos(3l0*y) + 1/2*((2k0)^2+(2l0)^2)*cos(2k0*x)*cos(2l0*y) )

  # ... their corresponding PVs q1, q2, q3,
  q1 = @. Δψ1 + 20ψ2 - 20ψ1
  q2 = @. Δψ2 + 20ψ1 - 44ψ2 + 24ψ3
  q3 = @. Δψ3        + 12ψ2 - 12ψ3

  return ψ1, ψ2, ψ3, q1, q2, q3
end

"""
    test_pvtofromstreamfunction_2layer()

Tests the pvfromstreamfunction function that gives qh from ψh. To do so, it
constructs a 2-layer problem with parameters such that
    q1 = Δψ1 + 25*(ψ2-ψ1), q2 = Δψ2 + 25/4*(ψ1-ψ2).
Then given ψ1 and ψ2 it tests if pvfromstreamfunction gives the expected
q1 and q2. Similarly, that streamfunctionfrompv gives ψ1 and ψ2 from q1 and q2.
"""
function test_pvtofromstreamfunction_2layer()
  n, L = 128, 2π
  gr = TwoDGrid(n, L)

  nlayers = 2      # these choice of parameters give the
  f0, g = 1, 1     # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultilayerQG.Problem(nlayers=nlayers, nx=n, Lx=L, f0=f0, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  vs.psih[:, :, 1] .= rfft(ψ1)
  vs.psih[:, :, 2] .= rfft(ψ2)

  MultilayerQG.pvfromstreamfunction!(vs.qh, vs.psih, pr.S, gr)
  MultilayerQG.invtransform!(vs.q, vs.qh, pr)

  vs.qh[:, :, 1] .= rfft(q1)
  vs.qh[:, :, 2] .= rfft(q2)

  MultilayerQG.streamfunctionfrompv!(vs.psih, vs.qh, pr.invS, gr)
  MultilayerQG.invtransform!(vs.psi, vs.psih, pr)

  isapprox(q1, vs.q[:, :, 1], rtol=rtol_multilayerqg) && isapprox(q2, vs.q[:, :, 2], rtol=rtol_multilayerqg) && isapprox(ψ1, vs.psi[:, :, 1], rtol=rtol_multilayerqg) && isapprox(ψ2, vs.psi[:, :, 2], rtol=rtol_multilayerqg)
end

"""
    test_pvtofromstreamfunction_3layer()

Tests the pvfromstreamfunction function that gives qh from ψh. To do so, it
constructs a 3-layer problem with parameters such that
  q1 = Δψ1 + 20ψ2 - 20ψ1, q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3, q3 = Δψ3 + 12ψ2 - 12ψ3.
Then given ψ1, ψ2, and ψ2 it tests if pvfromstreamfunction gives the expected
q1, q2, and q3. Similarly, that streamfunctionfrompv gives ψ1, ψ2, and ψ2 from
q1, q2, and q3.
"""
function test_pvtofromstreamfunction_3layer()
  n, L = 128, 2π
  gr = TwoDGrid(n, L)

  nlayers = 3      # these choice of parameters give the
  f0, g = 1, 1     # desired PV-streamfunction relations
  H = [0.25, 0.25, 0.5]  # q1 = Δψ1 + 20ψ2 - 20ψ1,
  ρ = [4.0, 5.0, 6.0]    # q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3,
                         # q3 = Δψ3        + 12ψ2 - 12ψ3.

  prob = MultilayerQG.Problem(nlayers=nlayers, nx=n, Lx=L, f0=f0, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, ψ3, q1, q2, q3 = constructtestfields_3layer(gr)

  vs.psih[:, :, 1] .= rfft(ψ1)
  vs.psih[:, :, 2] .= rfft(ψ2)
  vs.psih[:, :, 3] .= rfft(ψ3)

  MultilayerQG.pvfromstreamfunction!(vs.qh, vs.psih, pr.S, gr)
  MultilayerQG.invtransform!(vs.q, vs.qh, pr)

  vs.qh[:, :, 1] .= rfft(q1)
  vs.qh[:, :, 2] .= rfft(q2)
  vs.qh[:, :, 3] .= rfft(q3)

  MultilayerQG.streamfunctionfrompv!(vs.psih, vs.qh, pr.invS, gr)
  MultilayerQG.invtransform!(vs.psi, vs.psih, pr)

  isapprox(q1, vs.q[:, :, 1], rtol=rtol_multilayerqg) && isapprox(q2, vs.q[:, :, 2], rtol=rtol_multilayerqg) && isapprox(q3, vs.q[:, :, 3], rtol=rtol_multilayerqg) && isapprox(ψ1, vs.psi[:, :, 1], rtol=rtol_multilayerqg) && isapprox(ψ2, vs.psi[:, :, 2], rtol=rtol_multilayerqg) && isapprox(ψ3, vs.psi[:, :, 3], rtol=rtol_multilayerqg)
end

"""
    test_mqg_nonlinearadvection(dt, stepper; kwargs...)

Tests the advection term by timestepping a test problem with timestep dt and
timestepper identified by the string stepper. The test 2-layer problem is
derived by picking a solution q1f, q2f (with streamfunctions ψ1f, ψ2f) for which
the advection terms J(ψn, qn) are non-zero. Next, a forcing Ff is derived such
that a solution to the problem forced by this Ff is then qf.
(This solution may not be realized, at least at long times, if it is unstable.)
"""
function test_mqg_nonlinearadvection(dt, stepper; n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  tf = 0.5
  nt = round(Int, tf/dt)

  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f0, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  β = 0.35

  U1, U2 = 0.1, 0.05
  u1 = @. 0.5sech(gr.y/(Ly/15))^2
  u2 = @. 0.02cos(3l0*gr.y)
  uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) ))
  uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) ))

  U = zeros(ny, nlayers)
  U[:, 1] = u1 .+ U1
  U[:, 2] = u2 .+ U2

  μ, ν, nν = 0.1, 0.05, 1

  η0, σx, σy = 1.0, Lx/25, Ly/20
  η = @. η0*exp( -(x+Lx/8)^2/(2σx^2) -(y-Ly/8)^2/(2σy^2) )
  ηx = @. -(x+Lx/8)/(σx^2) * η

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  Ff1 = FourierFlows.jacobian(ψ1, q1, gr)     + (β .- uyy1 .-   25*(U2.+u2.-U1.-u1) ).*ψ1x + (U1.+u1).*q1x - ν*Δq1
  Ff2 = FourierFlows.jacobian(ψ2, q2 + η, gr) + (β .- uyy2 .- 25/4*(U1.+u1.-U2.-u2) ).*ψ2x + (U2.+u2).*(q2x + ηx) + μ*Δψ2 - ν*Δq2

  Ff = zeros(gr.nx, gr.ny, nlayers)
  Ff[:, :, 1] .= Ff1
  Ff[:, :, 2] .= Ff2

  Ffh = zeros(Complex{Float64}, gr.nkr, gr.nl, nlayers)
  Ffh[:, :, 1] .= rfft(Ff1)
  Ffh[:, :, 2] .= rfft(Ff2)

  function calcFq!(Fqh, sol, t, cl, v, p, g)
    Fqh .= Ffh
    nothing
  end

  prob = MultilayerQG.Problem(nlayers=nlayers, nx=nx, ny=ny, Lx=Lx, Ly=Ly, f0=f0,
          g=g, H=H, ρ=ρ, U=U, eta=η, β=β, μ=μ, ν=ν, nν=nν, calcFq=calcFq!)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  qf = zeros(gr.nx, gr.ny, nlayers)
  qf[:, :, 1] .= q1
  qf[:, :, 2] .= q2

  ψf = zeros(gr.nx, gr.ny, nlayers)
  ψf[:, :, 1] .= ψ1
  ψf[:, :, 2] .= ψ2

  MultilayerQG.set_q!(prob, qf)

  stepforward!(prob, round(Int, nt))
  MultilayerQG.updatevars!(prob)

  isapprox(vs.q, qf, rtol=rtol_multilayerqg) && isapprox(vs.psi, ψf, rtol=rtol_multilayerqg)
end

"""
    test_mqg_linearadvection(dt, stepper; kwargs...)

Tests the advection term of the linearized equations by timestepping a test
problem with timestep dt and timestepper identified by the string stepper.
The test 2-layer problem is derived by picking a solution q1f, q2f (with
streamfunctions ψ1f, ψ2f) for which the advection terms J(ψn, qn) are non-zero.
Next, a forcing Ff is derived such that a solution to the problem forced by this
Ff is then qf. (This solution may not be realized, at least at long times, if it
is unstable.)
"""
function test_mqg_linearadvection(dt, stepper; n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  tf = 0.5
  nt = round(Int, tf/dt)

  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f0, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  β = 0.35

  U1, U2 = 0.1, 0.05
  u1 = @. 0.5sech(gr.y/(Ly/15))^2
  u2 = @. 0.02cos(3l0*gr.y)
  uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) ))
  uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) ))

  U = zeros(ny, nlayers)
  U[:, 1] = u1 .+ U1
  U[:, 2] = u2 .+ U2

  μ, ν, nν = 0.1, 0.05, 1

  η0, σx, σy = 1.0, Lx/25, Ly/20
  η = @. η0*exp( -(x+Lx/8)^2/(2σx^2) -(y-Ly/8)^2/(2σy^2) )
  ηx = @. -(x+Lx/8)/(σx^2) * η

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  Ff1 = (β .- uyy1 .-   25*(U2.+u2.-U1.-u1) ).*ψ1x + (U1.+u1).*q1x - ν*Δq1
  Ff2 = FourierFlows.jacobian(ψ2, η, gr) + (β .- uyy2 .- 25/4*(U1.+u1.-U2.-u2) ).*ψ2x + (U2.+u2).*(q2x + ηx) + μ*Δψ2 - ν*Δq2

  Ff = zeros(gr.nx, gr.ny, nlayers)
  Ff[:, :, 1] .= Ff1
  Ff[:, :, 2] .= Ff2

  Ffh = zeros(Complex{Float64}, gr.nkr, gr.nl, nlayers)
  Ffh[:, :, 1] .= rfft(Ff1)
  Ffh[:, :, 2] .= rfft(Ff2)

  function calcFq!(Fqh, sol, t, cl, v, p, g)
    Fqh .= Ffh
    nothing
  end

  prob = MultilayerQG.Problem(nlayers=nlayers, nx=nx, ny=ny, Lx=Lx, Ly=Ly, f0=f0,
          g=g, H=H, ρ=ρ, U=U, eta=η, β=β, μ=μ, ν=ν, nν=nν, calcFq=calcFq!, linear=true)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  qf = zeros(gr.nx, gr.ny, nlayers)
  qf[:, :, 1] .= q1
  qf[:, :, 2] .= q2

  ψf = zeros(gr.nx, gr.ny, nlayers)
  ψf[:, :, 1] .= ψ1
  ψf[:, :, 2] .= ψ2

  MultilayerQG.set_q!(prob, qf)

  stepforward!(prob, round(Int, nt))
  MultilayerQG.updatevars!(prob)

  isapprox(vs.q, qf, rtol=rtol_multilayerqg) && isapprox(vs.psi, ψf, rtol=rtol_multilayerqg)
end

"""
    test_mqg_energies(dt, stepper; kwargs...)

Tests the kinetic (KE) and potential (PE) energies function by constructing a
2-layer problem and initializing it with a flow field whose KE and PE are known.
"""
function test_mqg_energies(; dt=0.001, stepper="ForwardEuler", n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers

  nlayers = 2      # these choice of parameters give the
  f0, g = 1, 1     # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultilayerQG.Problem(nlayers=nlayers, nx=nx, ny=ny, Lx=Lx, Ly=Ly, f0=f0, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  qf = zeros(gr.nx, gr.ny, nlayers)
  qf[:, :, 1] .= q1
  qf[:, :, 2] .= q2

  MultilayerQG.set_q!(prob, qf)

  KE, PE = MultilayerQG.energies(prob)

  isapprox(KE[1], 61/640*1e-6, rtol=rtol_multilayerqg) && isapprox(KE[2], 3*1e-6, rtol=rtol_multilayerqg) && isapprox(PE[1], 1025/1152*1e-6, rtol=rtol_multilayerqg) && MultilayerQG.addforcing!(prob.timestepper.RHS₁, sol, cl.t, cl, vs, pr, gr)==nothing
end

"""
    test_mqg_fluxes(dt, stepper; kwargs...)

Tests the lateral and vertical eddy fluxes by constructing a 2-layer problem and
initializing it with a flow field whose fluxes are known.
"""
function test_mqg_fluxes(; dt=0.001, stepper="ForwardEuler", n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 128, 126
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers

  nlayers = 2      # these choice of parameters give the
  f0, g = 1, 1     # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).
  U = zeros(ny, nlayers)
  U[:, 1] = @. sech(gr.y/0.2)^2

  prob = MultilayerQG.Problem(nlayers=nlayers, nx=nx, ny=ny, Lx=Lx, Ly=Ly, f0=f0, g=g, H=H, ρ=ρ, U=U)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1 = @. cos(k0*x)*cos(l0*y)
  ψ2 = @. cos(k0*x+π/10)*cos(l0*y)
  ψ = zeros(gr.nx, gr.ny, nlayers)
  ψ[:, :, 1] .= ψ1
  ψ[:, :, 2] .= ψ2
  MultilayerQG.set_psi!(prob, ψ)
  lateralfluxes, verticalfluxes = MultilayerQG.fluxes(prob)

  isapprox(lateralfluxes[1], 0, atol=1e-12) && isapprox(lateralfluxes[2], 0, atol=1e-12) && isapprox(verticalfluxes[1], -0.04763511558, rtol=1e-6)
end

"""
    test_setqsetpsi(dt, stepper; kwargs...)

Tests the set_q!() and set_psi!() functions that initialize sol with a flow with
given `q` or `psi` respectively.
"""
function test_mqg_setqsetpsi(; dt=0.001, stepper="ForwardEuler", n=64, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 32, 34
  L = 2π
  gr = TwoDGrid(nx, L, ny, L)

  x, y = gridpoints(gr)
  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers

  nlayers = 2      # these choice of parameters give the
  f0, g = 1, 1     # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultilayerQG.Problem(nlayers=nlayers, nx=nx, ny=ny, Lx=L, f0=f0, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  f1 = @. 2cos(k0*x)*cos(l0*y)
  f2 = @.  cos(k0*x+π/10)*cos(2l0*y)
  f = zeros(gr.nx, gr.ny, nlayers)
  f[:, :, 1] .= f1
  f[:, :, 2] .= f2

  ψtest = zeros(size(f))
  MultilayerQG.set_psi!(prob, f)
  @. vs.qh = sol
  MultilayerQG.streamfunctionfrompv!(vs.psih, vs.qh, pr.invS, gr)
  MultilayerQG.invtransform!(ψtest, vs.psih, pr)

  qtest = zeros(size(f))
  MultilayerQG.set_q!(prob, f)
  @. vs.qh = sol
  MultilayerQG.invtransform!(qtest, vs.qh, pr)

  isapprox(ψtest, f, rtol=rtol_multilayerqg) && isapprox(qtest, f, rtol=rtol_multilayerqg)
end

"""
    test_paramsconstructor(; kwargs...)

Tests that `Params` constructor works with both mean flow `U` being a floats
(i.e., constant `U` in each layer) or vectors (i.e., `U(y)` in each layer).
"""
function test_mqg_paramsconstructor(; dt=0.001, stepper="ForwardEuler")
  nx, ny = 32, 34
  L = 2π
  gr = TwoDGrid(nx, L, ny, L)

  nlayers = 2      # these choice of parameters give the
  f0, g = 1, 1     # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  U1, U2 = 0.1, 0.05

  Uvectors = zeros(ny, nlayers)
  Uvectors[:, 1] .= U1
  Uvectors[:, 2] .= U2

  Ufloats = zeros(nlayers)
  Ufloats[1] = U1
  Ufloats[2] = U2

  probUvectors = MultilayerQG.Problem(nlayers=nlayers, nx=nx, ny=ny, Lx=L, f0=f0, g=g, H=H, ρ=ρ, U=Uvectors)
  probUfloats = MultilayerQG.Problem(nlayers=nlayers, nx=nx, ny=ny, Lx=L, f0=f0, g=g, H=H, ρ=ρ, U=Ufloats)

  isapprox(probUfloats.params.U, probUvectors.params.U, rtol=rtol_multilayerqg)
end

function test_mqg_problemtype(T=Float32)
  prob = MultilayerQG.Problem(nlayers=2, T=T)

  (typeof(prob.sol)==Array{Complex{T},3} && typeof(prob.grid.Lx)==T && typeof(prob.grid.x)==Array{T,2} && typeof(prob.vars.u)==Array{T,3})
end


text:
using LineSearches
function bellmanobjective(system::DynamicSystem1D,
                          t,x,a,ω, vt::Function)
    numsamples = length(ω)
    retval = 0.
    @simd for n = 1:numsamples
        retval += system.U(t+1,x,a,ω[n]) + vt(system.f(t+1,x,a,ω[n]))
    end
    return -retval / numsamples # Optim minimizes
end

function bellmanobjective(system::DynamicSystem1D,
                          t,x,a,ω, Iv::AbstractInterpolation)
    vt(x) = Iv[x]
    bellmanobjective(system,t,x,a,ω,vt)
end

function solvebellman!(v, α, system::DynamicSystem1D, xtup, ω)
    K = length(xtup[1])
    T = system.T
    @simd for i = 1:K
        v[i,T+1] = system.Ubar(xtup[1][i])
    end
    ti = T
    t = ti-1
    v[1,ti] = system.Ubar(xtup[1][1]) # Boundary condition
    @simd for i = K:-1:2
        objective(a) = bellmanobjective(system,t,xtup[1][i],a, ω[:,ti], system.Ubar)
        res = optimize(objective, system.amin, system.amax)
        v[i,ti] = -Optim.minimum(res)
        α[i,ti] = Optim.minimizer(res)
    end

    for ti = T-1:-1:1
        Iv = interpolate(xtup, v[:,ti+1], Gridded(Linear()))
        t = ti-1
        v[1,ti] = system.Ubar(xtup[1][1]) # Boundary condition
        @simd for i = K:-1:2
            objective(a) = bellmanobjective(system,t,xtup[1][i],a, ω[:,ti], Iv)
            res = optimize(objective, system.amin, system.amax)
            v[i,ti] = -Optim.minimum(res)
            α[i,ti] = Optim.minimizer(res)
        end
    end
    #α[1,:] = NaN # No policy at x=0
    α[1,:] = system.amax
end

function bellmanobjective(system::DynamicSystem,
                          t,x,a,ω::Tuple, vt::Function)
    numsamples = length(ω[1])
    retval = 0.
    @simd for n = 1:numsamples
        wn = [ωi[n] for ωi in ω]
        retval += system.U(t+1,x,a,wn) + vt(system.f(t+1,x,a,wn))
    end
    return -retval / numsamples
end

function bellmanobjective(system::DynamicSystem,
                          t,x,a,ω::Tuple,Iv::AbstractInterpolation)
    vt(x) = Iv[x...]
    bellmanobjective(system,t,x,a,ω,vt)
end

function solvebellman!(v, α, system::DynamicSystem, xtup, ω::Tuple)
    # TODO: use CartesianIndex stuff to allow any dimension
    #       see http://julialang.org/blog/2016/02/iteration
    @assert system.statedim == 2     # TODO: generalise
    @assert length(system.amin(0,xtup[1])) == 1 # TODO: generalise

    K = [length(xtup[i]) for i = 1:system.statedim]
    T = system.T
    vti = view(v,:,:,T+1)
    for i = 1:K[1]
        @simd for j=1:K[2]
            xij = [xtup[1][i], xtup[2][j]]
            vti[j,i] = system.Ubar(xij)
        end
    end

    ti = T
    vti = view(v,:,:,ti)
    αti = view(α,:,:,ti)
    ωti = tuple([view(ωi,:,ti) for ωi in ω]...)
    t = ti-1
    @show t
    # Boundary condition
    @simd for j = 1:K[2]
        xij = [xtup[1][1], xtup[2][j]] # TODO: CartesianIndex on xtup somehow?
        vti[j,1] = system.Ubar(xij)
    end
    for i = 2:K[1] # TODO: CartesianIndex on vti
        @simd for j = 1:K[2]
            xij = [xtup[1][i], xtup[2][j]] # TODO: CartesianIndex on xtup somehow?
            objective(a) = bellmanobjective(system,t,xij,a, ωti, system.Ubar)
            res = optimize(objective, system.amin(t,xij), system.amax(t,xij))
            #res = optimize(objective, [0.5], LBFGS(linesearch = LineSearches.BackTracking()),
            #               OptimizationOptions(autodiff=true))
            @inbounds vti[j,i] = -Optim.minimum(res)#[1]
            @inbounds αti[j,i] = Optim.minimizer(res)#[1]
        end
    end

    for ti = T-1:-1:1
        # TODO: How to deal with extrapolations?
        Iv = interpolate(xtup, vti', Gridded(Linear()))
        #Iv = extrapolate(Iv,Interpolations.Throw())
        vti = view(v,:,:,ti)
        αti = view(α,:,:,ti)
        ωti = tuple([view(ωi,:,ti) for ωi in ω]...)
        t = ti-1
        @show t
        # Boundary condition
        @simd for j = 1:K[2]
            xij = [xtup[1][1], xtup[2][j]] # TODO: CartesianIndex on xtup somehow?
            vti[j,1] = system.Ubar(xij)
        end
        for i = 2:K[1] # TODO: CartesianIndex on vti
            @simd for j = 1:K[2]
                xij = [xtup[1][i], xtup[2][j]] # TODO: CartesianIndex on xtup somehow?
                objective(a) = bellmanobjective(system,t,xij,a, ωti, Iv)
                res = optimize(objective, system.amin(t,xij), system.amax(t,xij))
                vti[j,i] = -Optim.minimum(res)
                αti[j,i] = Optim.minimizer(res)
                #@show xij
                #@show vti[j,i], αti[j,i]
            end
        end
    end
    @inbounds α[:,1,:] = NaN # No policy at x[1]=0
end


text:
# ---
# title: 1476. Subrectangle Queries
# id: problem1476
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array
# link: <https://leetcode.com/problems/subrectangle-queries/description/>
# hidden: true
# ---
#
# Implement the class `SubrectangleQueries` which receives a `rows x cols`
# rectangle as a matrix of integers in the constructor and supports two methods:
#
# 1.` updateSubrectangle(int row1, int col1, int row2, int col2, int newValue)`
#
#   * Updates all values with `newValue` in the subrectangle whose upper left coordinate is `(row1,col1)` and bottom right coordinate is `(row2,col2)`.
#
# 2.` getValue(int row, int col)`
#
#   * Returns the current value of the coordinate `(row,col)` from the rectangle.
#
#
#
# **Example 1:**
#
#
#
#     **Input**
#     ["SubrectangleQueries","getValue","updateSubrectangle","getValue","getValue","updateSubrectangle","getValue","getValue"]
#     [[[[1,2,1],[4,3,4],[3,2,1],[1,1,1]]],[0,2],[0,0,3,2,5],[0,2],[3,1],[3,0,3,2,10],[3,1],[0,2]]
#     **Output**
#     [null,1,null,5,5,null,10,5]
#     **Explanation**
#     SubrectangleQueries subrectangleQueries = new SubrectangleQueries([[1,2,1],[4,3,4],[3,2,1],[1,1,1]]);
#     // The initial rectangle (4x3) looks like:
#     // 1 2 1
#     // 4 3 4
#     // 3 2 1
#     // 1 1 1
#     subrectangleQueries.getValue(0, 2); // return 1
#     subrectangleQueries.updateSubrectangle(0, 0, 3, 2, 5);
#     // After this update the rectangle looks like:
#     // 5 5 5
#     // 5 5 5
#     // 5 5 5
#     // 5 5 5
#     subrectangleQueries.getValue(0, 2); // return 5
#     subrectangleQueries.getValue(3, 1); // return 5
#     subrectangleQueries.updateSubrectangle(3, 0, 3, 2, 10);
#     // After this update the rectangle looks like:
#     // 5   5   5
#     // 5   5   5
#     // 5   5   5
#     // 10  10  10
#     subrectangleQueries.getValue(3, 1); // return 10
#     subrectangleQueries.getValue(0, 2); // return 5
#
#
# **Example 2:**
#
#
#
#     **Input**
#     ["SubrectangleQueries","getValue","updateSubrectangle","getValue","getValue","updateSubrectangle","getValue"]
#     [[[[1,1,1],[2,2,2],[3,3,3]]],[0,0],[0,0,2,2,100],[0,0],[2,2],[1,1,2,2,20],[2,2]]
#     **Output**
#     [null,1,null,100,100,null,20]
#     **Explanation**
#     SubrectangleQueries subrectangleQueries = new SubrectangleQueries([[1,1,1],[2,2,2],[3,3,3]]);
#     subrectangleQueries.getValue(0, 0); // return 1
#     subrectangleQueries.updateSubrectangle(0, 0, 2, 2, 100);
#     subrectangleQueries.getValue(0, 0); // return 100
#     subrectangleQueries.getValue(2, 2); // return 100
#     subrectangleQueries.updateSubrectangle(1, 1, 2, 2, 20);
#     subrectangleQueries.getValue(2, 2); // return 20
#
#
#
#
# **Constraints:**
#
#   * There will be at most `500` operations considering both methods: `updateSubrectangle` and `getValue`.
#   * `1 <= rows, cols <= 100`
#   * `rows == rectangle.length`
#   * `cols == rectangle[i].length`
#   * `0 <= row1 <= row2 < rows`
#   * `0 <= col1 <= col2 < cols`
#   * `1 <= newValue, rectangle[i][j] <= 10^9`
#   * `0 <= row < rows`
#   * `0 <= col < cols`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
# isdefined(Base, :__precompile__) && __precompile__()

module VariationalInequality

# package code goes here

using JuMP, LinearAlgebra, SparseArrays, Printf
import Ipopt

include("model.jl")
include("algorithms.jl")

export  VIPModel,
        VIPData,
        solveVIP,
        clear_values,
        save_solution,
        gap_function,
        @mapping,
        @innerproduct

end # module