text:
"""
```
get_forecast_input_file(m, input_type)
```

Compute the appropriate forecast input filenames for model `m` and
forecast input type `input_type`.

The default input files for each `input_type` can be overriden by adding entries
to the `Dict{Symbol, String}` returned from
`forecast_input_file_overrides(m)`. For example:

```
overrides = forecast_input_file_overrides(m)
overrides[:mode] = \"path/to/input/file.h5\"
```
"""
function get_forecast_input_file(m, input_type;
                                 filestring_addl::Vector{String} = Vector{String}(undef, 0))
    overrides = forecast_input_file_overrides(m)
    if haskey(overrides, input_type)
        override_file = overrides[input_type]
        if ispath(override_file)
            return override_file
        else
            error("Input file $override_file does not exist")
        end
    elseif input_type == :subset
        # If input_type == :subset, also check for existence of overrides[:full]
        return get_forecast_input_file(m, :full)
    end

    if input_type == :mode || input_type == :mode_draw_shocks
        return rawpath(m,"estimate","paramsmode.h5", filestring_addl)
    elseif input_type == :mean
        return workpath(m,"estimate","paramsmean.h5", filestring_addl)
    elseif input_type == :init || input_type == :init_draw_shocks
        return ""
    elseif input_type == :prior
        return ""
    elseif input_type in [:full, :subset]
        if get_setting(m, :sampling_method) == :MH
            return rawpath(m, "estimate", "mhsave.h5", filestring_addl)
        elseif get_setting(m, :sampling_method) == :SMC
            return rawpath(m, "estimate", "smcsave.h5", filestring_addl)
        else
            throw("Invalid sampling method specification. Change in setting :sampling_method")
        end
    else
        throw(ArgumentError("Invalid input_type: $(input_type)"))
    end
end

"""
```
get_forecast_filename(m, input_type, cond_type, output_var;
    pathfcn = rawpath, forecast_string = "", fileformat = :jld2)

get_forecast_filename(directory, filestring_base, input_type, cond_type,
    output_var; forecast_string = "", fileformat = :jld2)
```

### Notes

- If `input_type == :subset`, then the `forecast_string` is also appended to the
filenames. If in this case `forecast_string` is empty, `get_forecast_filename` throws
an error.

- In the second method, `directory` should be a string of the form
  `\"\$saveroot/m990/ss2/forecast/raw/\"`. (Note that a `pathfcn` is therefore
  not required.) `filestring_base` should be equivalent to the result of
  `filestring_base(m)`.
"""
function get_forecast_filename(m::AbstractDSGEModel, input_type::Symbol,
                               cond_type::Symbol, output_var::Symbol;
                               pathfcn::Function = rawpath,
                               forecast_string::String = "",
                               fileformat::Symbol = :jld2)

    # First, we need to make sure we get all of the settings that have been printed to this filestring
    directory = pathfcn(m, "forecast")
    base      = filestring_base(m)

    # Now we can find the filestring we are looking for
    get_forecast_filename(directory, base, input_type, cond_type, output_var;
                          forecast_string = forecast_string, fileformat = fileformat)
end

function get_forecast_filename(directory::String, filestring_base::Vector{String},
                               input_type::Symbol, cond_type::Symbol, output_var::Symbol;
                               forecast_string::String = "", fileformat = :jld2)

    # Base file name
    file_name = String("$output_var.$fileformat")

    # Gather all of the file strings into an array
    filestring_addl = get_forecast_filestring_addl(input_type, cond_type,
                                                   forecast_string = forecast_string)

    return savepath(directory, file_name, filestring_base, filestring_addl)
end

function get_forecast_filestring_addl(input_type, cond_type; forecast_string::String = "")

    filestring_addl = Vector{String}()
    push!(filestring_addl, String("para=" * String(input_type)))
    push!(filestring_addl, String("cond=" * abbrev_symbol(cond_type)))
    if isempty(forecast_string)
        if input_type == :subset
            error("Must supply a nonempty forecast_string if input_type = subset")
        end
    else
        push!(filestring_addl, String("fcid=" * forecast_string))
    end

    filestring_addl
end

"""
```
get_forecast_output_files(m, input_type, cond_type, output_vars;
    forecast_string = "", fileformat = :jld2)

get_forecast_output_files(directory, filestring_base, input_type, cond_type,
    output_vars; forecast_string = "", fileformat = :jld2)
```

Compute the appropriate output filenames for model `m`, forecast input
type `input_type`, and conditional type `cond_type`, for each output variable in
`output_vars`. Returns a dictionary of file names with one entry for each output_var.

### Arguments

- See `forecast_one` for descriptions of other non-keyword arguments.

### Keyword Arguments

- `forecast_string::String`: subset identifier for when `input_type = :subset`
- `fileformat::Symbol`: file extension, without a period. Defaults to
  `:jld2`, though `:h5` is another common option.

### Notes

See `get_forecast_filename` for more information.
"""
function get_forecast_output_files(m::AbstractDSGEModel, input_type::Symbol,
                                   cond_type::Symbol, output_vars::Vector{Symbol};
                                   forecast_string::String = "",
                                   fileformat::Symbol = :jld2)

    directory = rawpath(m, "forecast")
    base      = filestring_base(m)

    get_forecast_output_files(directory, base, input_type, cond_type, output_vars,
                              forecast_string = forecast_string,
                              fileformat = fileformat)
end

function get_forecast_output_files(directory::String, filestring_base::Vector{String},
                                   input_type::Symbol, cond_type::Symbol, output_vars::Vector{Symbol};
                                   forecast_string::String = "", fileformat::Symbol = :jld2)

    output_files    = Dict{Symbol, String}()
    for var in remove_meansbands_only_output_vars(output_vars)
        output_files[var] = get_forecast_filename(directory, filestring_base,
                                                  input_type, cond_type, var,
                                                  forecast_string = forecast_string,
                                                  fileformat = fileformat)
    end

    return output_files
end

"""
```
write_forecast_outputs(m, output_vars, forecast_output_files, forecast_output;
    df = DataFrame(), block_number = Nullable{Int64}(), block_inds = 1:0,
    verbose = :low)
```

Writes the elements of `forecast_output` indexed by `output_vars` to file, given
`forecast_output_files`, which maps `output_vars` to file names.

If `output_vars` contains `:histobs`, data must be passed in as `df`.
"""
function write_forecast_outputs(m::AbstractDSGEModel, input_type::Symbol,
                                output_vars::Vector{Symbol},
                                forecast_output_files::Dict{Symbol, String},
                                forecast_output::Dict{Symbol, Array{Float64}};
                                df::DataFrame = DataFrame(),
                                block_number::Nullable{Int64} = Nullable{Int64}(),
                                block_inds::AbstractRange{Int64} = 1:0,
                                subset_inds::AbstractRange{Int64} = 1:0,
                                verbose::Symbol = :low)

    for var in output_vars
        prod = get_product(var)
        if prod in [:histut, :hist4q, :forecastut, :bddforecastut, :forecast4q, :bddforecast4q]
            # These are computed and saved in means and bands, not
            # during the forecast itself
            continue
        end
        filepath = forecast_output_files[var]
        # Write forecast metadata to a jld2 and the raw forecast output
        # to an h5. The data in the HDF5 will be transferred to the jld2
        # and the h5 file will be deleted when combine_raw_forecast_output_and_metadata
        # is executed.
        if isnull(block_number) || get(block_number) == 1
            JLD2.jldopen(filepath, "w") do file
                write_forecast_metadata(m, file, var)
            end
            h5open(replace(filepath, "jld2" => "h5"), "w") do file
                if var == :histobs
                    # :histobs just refers to data, so we only write one draw
                    # (as all draws would be the same)
                    @assert !isempty(df) "df cannot be empty if trying to write :histobs"
                    df1 = df[date_mainsample_start(m) .<= df[:date] .<= date_mainsample_end(m), :]
                    data = df_to_matrix(m, df1)

                    # Must call missing2nan since you cannot write Missing types to HDF5 files
                    write(file, "arr", missing2nan(data))
                else
                    # Otherwise, pre-allocate HDF5 dataset which will contain
                    # all draws
                    if !isnull(block_number) && get(block_number) == 1
                        # Determine forecast output size
                        dims  = get_forecast_output_dims(m, input_type, var; subset_inds = subset_inds)
                        block_size = forecast_block_size(m)
                        chunk_dims = collect(dims)
                        chunk_dims[1] = block_size

                        # Initialize dataset
                        #pfile = file #.plain
                        dset = HDF5.d_create(file, "arr", datatype(Float64), dataspace(dims...), "chunk", chunk_dims)
                    end
                end
            end
        end

        if var != :histobs
            h5open(replace(filepath, "jld2" => "h5"), "r+") do file
                if isnull(block_number)
                    write(file, "arr", Array{Float64}(forecast_output[var]))
                else
                    write_forecast_block(file, Array{Float64}(forecast_output[var]), block_inds)
                end
            end
        end
        println(verbose, :high, " * Wrote $(basename(filepath))")
    end
end

"""
```
write_forecast_metadata(m::AbstractDSGEModel, file::JLDFile, var::Symbol)
```

Write metadata about the saved forecast output `var` to `filepath`.

Specifically, we save dictionaries mapping dates, as well as state, observable,
pseudo-observable, and shock names, to their respective indices in the saved
forecast output array. The saved dictionaries include:

- `date_indices::Dict{Date, Int}`: saved for all forecast outputs except IRFs
- `state_names::Dict{Symbol, Int}`: saved for `var in [:histstates, :forecaststates, :shockdecstates]`
- `observable_names::Dict{Symbol, Int}`: saved for `var in [:forecastobs, :shockdecobs]`
- `observable_revtransforms::Dict{Symbol, Symbol}`: saved identifiers for reverse transforms used for observables
- `pseudoobservable_names::Dict{Symbol, Int}`: saved for `var in [:histpseudo, :forecastpseudo, :shockdecpseudo]`
- `pseudoobservable_revtransforms::Dict{Symbol, Symbol}`: saved identifiers for reverse transforms used for pseudoobservables
- `shock_names::Dict{Symbol, Int}`: saved for `var in [:histshocks, :forecastshocks, :shockdecstates, :shockdecobs, :shockdecpseudo]`

Note that we don't save dates or transformations for impulse response
functions.
"""
function write_forecast_metadata(m::AbstractDSGEModel, file::JLD2.JLDFile, var::Symbol)

    prod  = get_product(var)
    class = get_class(var)

    # Write date range
    if prod != :irf
        dates = if prod == :hist
            quarter_range(date_mainsample_start(m), date_mainsample_end(m))
        elseif prod in [:forecast, :bddforecast]
            quarter_range(date_forecast_start(m), date_forecast_end(m))
        elseif prod in [:shockdec, :dettrend, :trend]
            quarter_range(date_shockdec_start(m), date_shockdec_end(m))
        elseif prod == :decomp
            quarter_range(date_mainsample_start(m), date_forecast_end(m))
        end

        date_indices = Dict(d::Date => i::Int for (i, d) in enumerate(dates))
        write(file, "date_indices", date_indices)
    end

    # Write state names
    if class == :states
        state_indices = merge(m.endogenous_states, m.endogenous_states_augmented)
        @assert length(state_indices) == n_states_augmented(m) # assert no duplicate keys
        write(file, "state_indices", state_indices)
        rev_transforms = Dict{Symbol,Symbol}(x => Symbol("identity") for x in keys(state_indices))
        write(file, "state_revtransforms", rev_transforms)
    end

    # Write observable names and transforms
    if class == :obs
        write(file, "observable_indices", m.observables)
        rev_transforms = if prod != :irf
            Dict{Symbol,Symbol}(x => Symbol(m.observable_mappings[x].rev_transform) for x in keys(m.observables))
        else
            Dict{Symbol,Symbol}(x => Symbol("identity") for x in keys(m.observables))
        end
        write(file, "observable_revtransforms", rev_transforms)
    end

    # Write pseudo-observable names and transforms
    if class == :pseudo
        write(file, "pseudoobservable_indices", m.pseudo_observables)
        rev_transforms = if prod != :irf
            Dict{Symbol,Symbol}(x => Symbol(m.pseudo_observable_mappings[x].rev_transform)
                                for x in keys(m.pseudo_observables))
        else
            Dict{Symbol,Symbol}(x => Symbol("identity") for x in keys(m.pseudo_observables))
        end
        write(file, "pseudoobservable_revtransforms", rev_transforms)
    end

    # Write shock names and transforms
    if class in [:shocks, :stdshocks] || prod in [:shockdec, :irf, :decompshockdec]
        write(file, "shock_indices", m.exogenous_shocks)
        if class in [:shocks, :stdshocks]
            rev_transforms = Dict{Symbol,Symbol}(x => Symbol("identity") for x in keys(m.exogenous_shocks))
            write(file, "shock_revtransforms", rev_transforms)
        end
    end
end

"""
```
write_forecast_block(file::JLDFile, arr::Array, block_number::Int,
    block_inds::AbstractRange{Int64})
```

Writes `arr` to the subarray of `file` indicated by `block_inds`.
"""
function write_forecast_block(file, arr::Array,
                              block_inds::AbstractRange{Int64})
    #pfile = file #.plain
    dataset = HDF5.d_open(file, "arr")
    dims = size(dataset)
    ndims = length(dims)
    dataset[block_inds, fill(Colon(), ndims-1)...] = arr #pfile was dataset before
    set_dims!(dataset, dims)
end

"""
```
combine_raw_forecast_output_and_metadata(m, forecast_output_files; verbose = :low)
```
Writes the raw forecast output data (`arr`) saved in the temporary h5 file to the
jld2 file containing the rest of the forecast metadata. The intermediary h5 step exists because
jld2 does not support chunked memory assignment in the same way that jld and h5 permitted previously.
"""
function combine_raw_forecast_output_and_metadata(m::AbstractDSGEModel,
                                                  forecast_output_files::Dict{Symbol, String};
                                                  verbose::Symbol = :low)
    for jld_file_path in values(forecast_output_files)
        h5_file_path = replace(jld_file_path, "jld2" => "h5")
        arr = h5read(h5_file_path, "arr")
        JLD2.jldopen(jld_file_path, "r+") do file
            write(file, "arr", arr)
        end
        # Remove the h5 containing the raw forecast output when the data has been transferred.
        rm(h5_file_path)
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            jld_file_name = replace(jld_file_path, rawpath(m, "forecast")*"/" => "")
            h5_file_name = replace(h5_file_path, rawpath(m, "forecast")*"/" => "")
            println("Wrote raw forecast output to $jld_file_name and removed $h5_file_name")
        end
    end
end

"""
```
read_forecast_metadata(file::JLDFile)
```

Read metadata from forecast output files. This includes dictionaries mapping
dates, as well as state, observable, pseudo-observable, and shock names, to
their respective indices in the saved forecast output array. Depending on the
`output_var`, the saved dictionaries might include:

- `date_indices::Dict{Date, Int}`: not saved for IRFs
- `state_indices::Dict{Symbol, Int}`
- `observable_indices::Dict{Symbol, Int}`
- `pseudoobservable_indices::Dict{Symbol, Int}`
- `shock_indices::Dict{Symbol, Int}`
- `state_revtransforms::Dict{Symbol, Symbol}`: states are not transformed, so
  all values are `:identity`
- `observable_revtransforms::Dict{Symbol, Symbol}`
- `pseudoobservable_revtransforms::Dict{Symbol, Symbol}`
- `shock_revtransforms::Dict{Symbol, Symbol}`: shocks are not transformed, so
  all values are `:identity`
"""
function read_forecast_metadata(file::JLD2.JLDFile)
    metadata = Dict{Symbol, Any}()
    for field in setdiff(keys(file), "arr")
        metadata[Symbol(field)] = read(file, field)
    end
    return metadata
end

function read_forecast_output(m::AbstractDSGEModel, input_type::Symbol, cond_type::Symbol,
                              output_var::Symbol, var_name::Symbol,
                              shock_name::Nullable{Symbol} = Nullable{Symbol}();
                              forecast_string::String = "")
    # Get filename
    filename = get_meansbands_input_file(m, input_type, cond_type, output_var,
                                         forecast_string = forecast_string)

    # Determine class and product
    product = get_product(output_var)
    class   = get_class(output_var)

    # Get index corresponding to var_name
    class_long = get_class_longname(class)
    indices = FileIO.load(filename, "$(class_long)_indices")
    var_ind = indices[var_name]

    # Read forecast output
    fcast_series = if isnull(shock_name)
        read_forecast_series(filename, product, var_ind)
    else
        # Get indices corresponding to shock_name
        shock_name = get(shock_name)
        shock_indices = FileIO.load(filename, "shock_indices")
        shock_ind = shock_indices[shock_name]

        read_forecast_series(filename, var_ind, shock_ind)
    end

    JLD2.jldopen(filename, "r") do file
        # The `fcast_output` for trends only is of size `ndraws` x `nvars`. We
        # need to use `repeat` below because population adjustments will be
        # different in each period. Now we have something of size `ndraws` x
        # `nvars` x `nperiods`
        if product == :trend
            nperiods = length(read(file, "date_indices"))
            fcast_series = repeat(fcast_series, outer = (1, nperiods))
        end

        # Parse transform
        class_long = get_class_longname(class)
        transforms = read(file, string(class_long) * "_revtransforms")
        transform = parse_transform(transforms[var_name])

        fcast_series, transform
    end
end

"""
```
read_forecast_series(filepath, product, var_ind)
```

Read only the forecast output for a particular variable (e.g. for a particular
observable) and possibly a particular shock. Result should be a matrix of size
`ndraws` x `nperiods`.
"""
function read_forecast_series(filepath::String, product::Symbol, var_ind::Int)

    dataset = FileIO.load(filepath, "arr")
    ndims = length(size(dataset))

    # Trends are ndraws x nvars
    if product == :trend
        whole = FileIO.load(filepath, "arr")
        if ndims == 1 # one draw
            arr = whole[var_ind, Colon()]
            arr = reshape(arr, (1, 1))
        elseif ndims == 2 # many draws
            arr = whole[Colon(), var_ind]
            arr = reshape(arr, (length(arr),1))
        end

    # Other products are ndraws x nvars x nperiods
    elseif product in [:hist, :histut, :hist4q, :forecast, :forecastut, :forecast4q,
                       :bddforecast, :bddforecastut, :bddforecast4q, :dettrend,
                       :decompdata, :decompnews, :decomppara, :decompdettrend, :decomptotal]
        inds_to_read = if ndims == 2 # one draw
            whole = FileIO.load(filepath, "arr")
            arr = whole[var_ind, Colon()]
            arr = reshape(arr, (1, length(arr)))
        elseif ndims == 3 # many draws
            whole = FileIO.load(filepath, "arr")
            arr = whole[Colon(), var_ind, Colon()]
        end
    else
        error("Invalid product: $product for this method")
    end

    return arr
end

function read_forecast_series(filepath::String, var_ind::Int, shock_ind::Int)

    dataset = FileIO.load(filepath, "arr")
    ndims = length(size(dataset))

    if ndims == 3 # one draw
        whole = FileIO.load(filepath, "arr")
        arr = whole[var_ind, :, shock_ind]
        arr = reshape(arr, 1, length(arr))
    elseif ndims == 4 # many draws
        whole = FileIO.load(filepath, "arr")
        arr = whole[:, var_ind, :, shock_ind]
    end

    return arr
end


text:
# HACK: This global variable stores information about the root locus diagrams
# that are open, so that it is possible to interact with them. This was the only
# way I could think to handle such case. This dictionary is indexed by the
# pointer of the figure canvas. In python, they create a class for each open
# figure and the variables are stored in an instantiation of this class. In
# julia, I do not know any good way to do this.
global rlocus_dict = Dict{Any, Any}()

"""
### function plot_rlocus(root_locus, root_locus_k, p, z, mp, mz)

Plot the root locus.

##### Args

* root_locus: Array with the closed-loop poles.
* root_locus_k: Array with the gains used to compute the closed-loop poles.
* p: Poles of the system.
* z: Zeros of the system.
* mp: Poles of the minimum realization of the system.
* mz: Zeros of the minimum realization of the system.

"""

function plot_rlocus(root_locus, root_locus_k, p, z, mp, mz)
    if plot_system == :pyplot
        # Create the figure.
        # ==================
        fig = plt.figure()
        ax  = plt.gca()

        # Plot axes.
        plt.axhline(0.0, color="black", linewidth=alw, zorder = 1)
        plt.axvline(0.0, color="black", linewidth=alw, zorder = 1)

        # Configure title and labels.
        plt.title("Root Locus")
        plt.xlabel("Real Axis")
        plt.ylabel("Imaginary Axis")

        # Add the root locus information to the global dictionary.
        rlocus_dict[fig[:canvas]] =
            (root_locus, root_locus_k)

        fig[:canvas][:mpl_connect]("button_press_event", rlocus_on_mouse_press)

        # Get real and imaginary parts of root locus.
        rl_real = map(x->real(x), root_locus)
        rl_imag = map(y->imag(y), root_locus)

        for i = 1:length(root_locus[1])
            rl_real_i = map(x->x[i], rl_real)
            rl_imag_i = map(x->x[i], rl_imag)

            # Get the next color.
            color = pybuiltin("next")(ax[:_get_lines][:prop_cycler])["color"]

            ax[:plot](rl_real_i, rl_imag_i, linewidth=plw, color=color)
        end

        # Mark poles and zeros.
        ( length(p) > 0 ) &&
        ax[:scatter](map(x->real(x), p), map(y->imag(y), p),
                     s=mksz, marker="x", linewidths=plw,
                     edgecolor="red", color="red")

        ( length(z) > 0 ) &&
        ax[:scatter](map(x->real(x), z), map(y->imag(y), z),
                     s=mksz, marker="o", linewidths=plw,
                     edgecolor="blue", facecolor="none")

        # Asymptotes
        # ==========

        # Compute the number of asymptotes.
        num_asymptotes = length(mp) - length(mz)

        σ_a = (sum(mp) - sum(mz))/num_asymptotes
        θ_a = [ (2*m+1)*pi/num_asymptotes for m in 0:(num_asymptotes-1) ]

        # Get the X and Y axis limits.
        xl = ax[:get_xlim]()
        yl = ax[:get_ylim]()

        # Get the limits in polar coordinates.
        θl_1 = mod(atan2(yl[2], real(xl[2]-σ_a)), 2*pi)
        θl_2 = mod(atan2(yl[2], real(xl[1]-σ_a)), 2*pi)
        θl_3 = mod(atan2(yl[1], real(xl[1]-σ_a)), 2*pi)
        θl_4 = mod(atan2(yl[1], real(xl[2]-σ_a)), 2*pi)

        for i = 1:length(θ_a)
            # Start point.
            x_0 = real(σ_a)
            y_0 = imag(σ_a)

            if (θ_a[i] <= θl_1) || (θ_a[i] >= θl_4)
                x_f = xl[2]
                y_f = y_0 + (x_f - x_0)*tan(θ_a[i])
            elseif (θ_a[i] >= θl_2) && (θ_a[i] <= θl_3)
                x_f = xl[1]
                y_f = y_0 + (x_f - x_0)*tan(θ_a[i])
            elseif (θ_a[i] > θl_1) && (θ_a[i] < θl_2)
                y_f = yl[2]
                x_f = x_0 + (y_f - y_0)/tan(θ_a[i])
            else
                y_f = yl[1]
                x_f = x_0 + (y_f - y_0)/tan(θ_a[i])
            end

            ax[:plot]([x_0;x_f], [y_0;y_f], linestyle="--", color="#BBBBBB")
        end
    end

    nothing
end

"""
### function rlocus_on_mouse_press(event)

Function called if the used click on a root locus plot. This is only available
for PyPlot

##### Args

* event: Click event.

"""

function rlocus_on_mouse_press(event)
    rlocus_data = rlocus_dict[event[:canvas]]

    root_locus   = rlocus_data[1]
    root_locus_k = rlocus_data[2]

    # The click point inside the plot.
    x = event[:xdata]
    y = event[:ydata]

    # If `x`  or `y` is `nothing`, then the user clicked outside the plot area.
    ( (x == nothing) || (y == nothing) ) && return

    # If there is no point ploted in the root locus, then nothing should be
    # done.
    (isempty(root_locus))    && return

    # If the system does not have any closed-loop poles, then nothing should be
    # done.
    (isempty(root_locus[1])) && return

    # Get the closest point in the graphic to the click point.
    ind1  = indmin(map(i->minimum(abs.(i - Complex(x,y))), root_locus))
    ind2  = indmin(abs.(root_locus[ind1] - Complex(x,y)))
    poles = root_locus[ind1]
    gain  = root_locus_k[ind1]

    # Print the information about the poles.
    println("")
    print((@bold)*(@green))
    println("Closed-Loop Poles")
    print(@creset)
    println("=================")
    println("")
    print("Selected gain: ")
    print(@treset)
    println("$gain")
    print(@bold)
    println("Closed-loop poles (selected one is marked in yellow):")
    println(@treset)
    pole_info(poles; highlight_ind=ind2)

    nothing
end

ext: jl


text:
# __BEGIN_LICENSE__
#
# ThreeDeconv.jl
#
# Copyright (c) 2018, Stanford University
#
# All rights reserved.
#
# Redistribution and use in source and binary forms for academic and other
# non-commercial purposes with or without modification, are permitted provided
# that the following conditions are met:
#
# * Redistributions of source code, including modified source code, must retain
#   the above copyright notice, this list of conditions and the following
#   disclaimer.
#
# * Redistributions in binary form or a modified form of the source code must
#   reproduce the above copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other materials provided with
#   the distribution.
#
# * Neither the name of The Leland Stanford Junior University, any of its
#   trademarks, the names of its employees, nor contributors to the source code
#   may be used to endorse or promote products derived from this software without
#   specific prior written permission.
#
# * Where a modified version of the source code is redistributed publicly in
#   source or binary forms, the modified source code must be published in a freely
#   accessible manner, or otherwise redistributed at no charge to anyone
#   requesting a copy of the modified source code, subject to the same terms as
#   this agreement.
#
# THIS SOFTWARE IS PROVIDED BY THE TRUSTEES OF THE LELAND STANFORD JUNIOR
# UNIVERSITY "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE LELAND STANFORD JUNIOR
# UNIVERSITY OR ITS TRUSTEES BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
# OF SUCH DAMAGE.
#
# __END_LICENSE__

function optimize(optimizer::Optimizer, options::DeconvolutionOptions)
    tic = time()

    state = optimizer.initial_state

    metric = optimizer.check_state(state, 0, tic)
    trace = OptimizationTrace{typeof(metric)}()
    push!(trace, metric)

    if options.show_trace
        print_header(optimizer.method)
        show(metric)
    end

    if options.x0 != []
        state.x .= typeof(state.x)(options.x0)
    end

    converged = false
    for iter in 1:options.max_iters

        optimizer.update_state!(state)

        if iter % options.check_every == 0

            metric = optimizer.check_state(state, iter, tic)
            push!(trace, metric)

            if options.show_trace
                show(metric)
            end

            if optimizer.check_convergence(metric)
                @printf "Optimization converged.\n"
                converged = true
                break
            end

            if !isnan(options.time_limit) && metric.elapsed_time > options.time_limit
                @printf "Time out.\n"
                break
            end

        end

    end

    return OptimizationResult(state.x, trace, converged)
end


text:
# old interface, deprecated
using DataFrames

open(dbpath::AbstractString) = DBInterface.connect(DuckDB.DB, dbpath)

connect(db::DB) = DBInterface.connect(db)

disconnect(con::Connection) = DBInterface.close!(con)
close(db::DB) = DBInterface.close!(db)

toDataFrame(res::QueryResult) = res.df
toDataFrame(con::Connection, sql::AbstractString) = toDataFrame(DBInterface.execute(con, sql))

function appendDataFrame(input_df::DataFrame, con::Connection, table::AbstractString, schema::String = "main")
    register_data_frame(con, input_df, "__append_df")
    DBInterface.execute(con, "INSERT INTO \"$schema\".\"$table\" SELECT * FROM __append_df")
    return unregister_data_frame(con, "__append_df")
end

appendDataFrame(input_df::DataFrame, db::DB, table::AbstractString, schema::String = "main") =
    appendDataFrame(input_df, db.main_connection, table, schema)

max_stars_repo_name: lokax/duckdb

max_issues_repo_name: lokax/duckdb

max_forks_repo_name: lokax/duckdb


text:
export Segment,
    Segments,
    intersect!,
    is_lower_end,
    is_upper_end,
    contains,
    find_leftmost,
    find_rightmost,
    contains,
    is_singular,
    min_x,
    min_y,
    max_x,
    max_y,
    trivial_miss

struct Segment{T<:Float64}
    p::Point{T}
    q::Point{T}
    slope::T
    function Segment(p::Point{T}, q::Point{T}) where {T<:Float64}
        if p > q
            p, q = q, p
        end
        slope = (q.x - p.x) / (q.y - p.y)
        return new{T}(p, q, slope)
    end
end

Base.isless(s::Segment, t::Segment) = (s.p.y > t.q.y) | ((s.p.y == t.p.y) && (s.p.x < t.p.x))
Base.:(==)(s::Segment, t::Segment) = (s.p == t.p) && (s.q == t.q)

get_x(segment::Segment, y) = segment.p.x + segment.slope * (y - segment.p.y)
get_y(segment::Segment, x) = segment.p.y + (x - segment.p.x) / segment.slope

min_x(segment::Segment) = min(segment.p.x, segment.q.x)
min_y(segment::Segment) = min(segment.p.y, segment.q.y)
max_x(segment::Segment) = max(segment.p.x, segment.q.x)
max_y(segment::Segment) = max(segment.p.y, segment.q.y)

Segment(px, py, qx, qy) = Segment(Point(px, py), Point(qx, qy))

"""
Checks if the segment s is purely horizontal or vertical.
"""
function is_singular(s::Segment)
    if (s.p.x == s.q.x) || (s.p.y == s.q.y)
        return true
    else
        return false
    end
end

function trivial_miss(s1::Segment, s2::Segment)
    if s1.slope ≈ s2.slope rtol=1e-10
        return true
    elseif max_x(s1) < min_x(s2)
        return true
    elseif max_x(s2) < min_x(s1)
        return true
    elseif max_y(s1) < min_y(s2)
        return true
    elseif max_y(s2) < min_y(s1)
        return true
    end
    return false
end

"""
Checks for the intersection of two segments s1, s2.
"""
function Base.intersect!(
    s1::Segment{T},
    s2::Segment{T},
    A::Matrix{T},
    b::Vector{T},
    tol=1e-9
) where {T<:AbstractFloat}
    if is_singular(s1) || is_singular(s2)
        return false, Point(0.0, 0.0)
    end
    if trivial_miss(s1, s2)
        return false, Point(0.0, 0.0)
    end
    A[1, 1] = s1.q.x - s1.p.x
    A[1, 2] = s2.p.x - s2.q.x
    A[2, 1] = s1.q.y - s1.p.y
    A[2, 2] = s2.p.y - s2.q.y
    b[1] = s2.p.x - s1.p.x
    b[2] = s2.p.y - s1.p.y
    sol = 0.0
    try
        sol = A \ b
    catch
        @warn "Singular matrix. Check for edge cases!"
        return false, Point(0.0, 0.0)
    end
    if (-tol < sol[1] < 1+tol) && (-tol < sol[2] < 1+tol)
        intersection =
            Point(s1.p.x + sol[1] * (s1.q.x - s1.p.x), s1.p.y + sol[1] * (s1.q.y - s1.p.y))
        return true, intersection
    else
        return false, Point(0.0, 0.0)
    end
end

function Base.intersect!(s1::Segment{T}, s2::Segment{T}) where {T<:AbstractFloat}
    A = zeros((2, 2))
    b = zeros(2)
    return intersect!(s1, s2, A, b)
end

is_lower_end(segment::Segment, Point::Point) = (segment.q == Point)
is_upper_end(segment::Segment, Point::Point) = (segment.p == Point)
function Base.contains(segment::Segment, point::Point, tol=1e-9)
    if is_lower_end(segment, point) | is_upper_end(segment, point)
        return false
    end
    y = get_y(segment, point.x)
    if y ≈ point.y
        atol = tol
        return true
    else
        return false
    end
end

function find_leftmost(segment_set, y, tol=1e-9)
    ret = nothing
    xmin = Inf
    for segment in segment_set
        x = get_x(segment, y - tol)
        if x < xmin
            xmin = x
            ret = segment
        end
    end
    return ret, xmin
end

function find_rightmost(segment_set, y, tol=1e-9)
    ret = nothing
    xmax = 0
    for segment in segment_set
        x = get_x(segment, y - tol)
        if x > xmax
            xmax = x
            ret = segment
        end
    end
    return ret, xmax
end


text:
"""
    correction_hull(A::IntervalMatrix{T}, t, p) where {T}

Compute the correction term for the convex hull of a point and its linear map
with an interval matrix in order to contain all trajectories of a linear system.

### Input

- `A` -- interval matrix
- `t` -- non-negative time value
- `p` -- order of the approximation

### Output

An interval matrix representing the correction term.

### Algorithm

See Theorem 3 in [1].

[1] M. Althoff, O. Stursberg, M. Buss. Reachability Analysis of Linear Systems
with Uncertain Parameters and Inputs. CDC 2007.
"""
function correction_hull(A::IntervalMatrix{T}, t, p) where {T}
    F = _exp_remainder(A, t, p)
    _correction_loop!(F, A, A, t, p)
    return F
end

# general matrix type
function correction_hull(A::AbstractMatrix, t, p)
    return correction_hull(IntervalMatrix(A), t, p)
end

"""
    input_correction(A::IntervalMatrix{T}, t, p) where {T}

Compute the *input correction matrix* for discretizing an inhomogeneous affine
dynamical system with an interval matrix and an input domain not containing the
origin.

### Input

- `A` -- interval matrix
- `t` -- non-negative time value
- `p` -- order of the Taylor approximation

### Output

An interval matrix representing the correction matrix.

### Algorithm

See Proposition 3.4 in [1].

[1] M. Althoff. Reachability analysis and its application to the safety
assessment of automonous cars. 2010.
"""
function input_correction(A::IntervalMatrix{T}, t, p) where {T}
    n = checksquare(A)

    E = _exp_remainder(A, t, p; n=n)
    F = E / opnorm(A, Inf)

    id = IntervalMatrix(Interval(one(T)) * I, n)
    _correction_loop!(F, A, id, t, p)

    return F
end

# inputs:
# F:  interval matrix sum (modified)
# A:  interval matrix
# Aⁱ: interval matrix factor in iteration i to start with; examples:
#     if Aⁱ == A, then this function will multiply with Aⁱ
#     if Aⁱ == id, then this function will multiply with Aⁱ⁻¹
# t:  time step
# p:  Taylor-approximation order
# outputs:
# F is modified by adding \\sum_{i=2}^p [...., 0] * t^i * (Aⁱ * A^i)
function _correction_loop!(F, A::IntervalMatrix{T}, Aⁱ, t, p) where {T}
    i! = 1
    tⁱ = t
    @inbounds for i in 2:p
        tⁱ *= t
        left = (one(T) / i^(i/(i-1)) - one(T) / i^(1/(i-1))) * tⁱ
        itv = Interval(left, zero(T))
        Aⁱ *= A
        i! *= i
        F .+= itv * Aⁱ / i!
    end
    return F
end


text:
struct ReferenceBus end

function attach_frequency_reference!(inputs::SimulationInputs, d::PSY.DynamicGenerator)
    inputs.global_vars[:ω_sys_index] = inputs.global_index[PSY.get_name(d)][:ω]
    return
end

function attach_frequency_reference!(inputs::SimulationInputs, d::PSY.DynamicInverter)
    inputs.global_vars[:ω_sys_index] = inputs.global_index[PSY.get_name(d)][:ω_oc]
    return
end

function attach_frequency_reference!(inputs::SimulationInputs, ::PSY.Source)
    inputs.global_vars[:ω_sys_index] = 0
    return
end

function attach_frequency_reference!(input::SimulationInputs, d::PSY.StaticInjection)
    return attach_frequency_reference!(input, PSY.get_dynamic_injector(d))
end

function set_frequency_reference!(inputs::SimulationInputs, sys::PSY.System)
    set_frequency_reference!(ReferenceBus, inputs, sys)
end

function set_frequency_reference!(
    ::Type{ReferenceBus},
    inputs::SimulationInputs,
    sys::PSY.System,
)
    ref_devices = PSY.get_components(
        PSY.StaticInjection,
        sys,
        x ->
            PSY.get_bustype(PSY.get_bus(x)) == PSY.BusTypes.REF &&
                !isa(x, PSY.ElectricLoad),
    )
    if length(ref_devices) > 1
        throw(
            IS.ConflictingInputsError(
                "More than one source or generator in the REF Bus is not supported in ReferenceBus model",
            ),
        )
    elseif length(ref_devices) < 1
        throw(
            IS.ConflictingInputsError(
                "ReferenceBus model requires at least one bus of type BusTypes.REF with a DynamicInjection or Source device connected to it",
            ),
        )
    else
        attach_frequency_reference!(inputs, first(ref_devices))
    end
    return
end


text:
module JLL

using GAP
using GAP_pkg_datastructures_jll

function __init__()
    # ensure GAP finds kernel extensions or other binaries
    sopath = joinpath(GAP_pkg_datastructures_jll.find_artifact_dir(), "lib", "gap")
    @debug "GAP package 'datastructures' sopath = " * sopath
    setproperty!(GAP.Globals.DirectoriesPackageProgramsOverrides, :datastructures, GapObj(sopath))
end

end # module


text:
using Knet, Test, Base.Iterators, Printf, LinearAlgebra, Random,  IterTools, CUDA
Knet.atype() = KnetArray{Float32}
include("/scratch/users/bberkay14/workfolder/language_to_logic/seq2tree/seq2tree/atis/tree.jl")
include("main.jl")
include("/scratch/users/bberkay14/workfolder/language_to_logic/seq2tree/seq2tree/atis/util.jl")
using CUDA
Knet.atype() = KnetArray{Float32}
using JLD2

function convert_to_string(idx_list, form_manager)
    w_list = []
    for i in range(1, stop=length(idx_list))
        push!(w_list, form_manager.get_idx_symbol(Int(idx_list[i])))
    end
    return join(w_list, " ")
end

function do_generate( l::L2L,  enc_w_list, word_manager, form_manager, opt,  checkpoint)
    # initialize the rnn state to all zeros

    prev_c  = torch.zeros((1, encoder.hidden_size), requires_grad=False)
    prev_h  = torch.zeros((1, encoder.hidden_size), requires_grad=False)
    prev_c  = param(1, enc_max_len; init=zeros)
    prev_h  = param(1, enc_max_len; init=zeros)

    # reversed order
    push!(enc_w_list, get_symbol_idx(word_manager,"<S>"))
    insert!(enc_w_list, 1, get_symbol_idx(word_manager, "<E>"))
    end_w = length(enc_w_list)
    enc_outputs = torch.zeros((1, end_w, encoder.hidden_size), requires_grad=False)
    enc_outputs = param(1, end_w, checkpoint["opt"]["rnn_size"] ; init=zeros)

    for i in range(end_w, 1, step=-1)

        cur_input = Knet.KnetArray(enc_w_list[i])

        prev_c, prev_h = encode(l, cur_input, prev_c, prev_h)
        enc_outputs[:, i, :] = prev_h
    end
    # decode
    queue_decode = []
    #queue_decode.append({"s": (prev_c, prev_h), "parent":0, "child_index":1, "t": Tree()})
    push!(queue_decode, Dict("s"=> (prev_c, prev_h), "parent"=>0, "child_index"=>1, "t"=> Tree()))

    head = 1
    while head <= size(queue_decode,1) and head <=100
        s = queue_decode[head-1]["s"]
        parent_h = s[1]
        t = queue_decode[head-1]["t"]
        if head == 1
            prev_word = KnetArray(get_symbol_idx(form_manager, "<S>"))

        else
            prev_word = KnetArray(get_symbol_idx(form_manager, "("))
        end

        i_child = 1
        while True
            curr_c, curr_h = decode(l, prev_word,  parent_h)
            prediction = attention_decode(l, enc_outputs, curr_h)

            s = (curr_c, curr_h)
            _prev_word = maximum(prediction, dims=2)

            prev_word = _prev_word
            if Int.(prev_word[1]) == get_symbol_idx(form_manager, "<E>") || t.num_children >= checkpoint["opt"]["dec_seq_length"]
                break
            elseif Int64.(prev_word[1]) == get_symbol_idx(form_manager, "<N>")
                #print("we predicted N");exit()
                push!(queue_decode, Dict("s"=> (deepcopy(s[1]), deepcopy(s[2])), "parent"=> head, "child_index"=>i_child, "t"=> Tree()))
                #deepcopy(s[1])
                #push!(queue_decode, Dict("s" => (s[0].clone(), s[1].clone()), "parent": head, "child_index":i_child, "t": Tree()))
                add_child(t, Int.(prev_word[1]))
            else
                add_child(t, Int.(prev_word[1]))
            end
            i_child = i_child + 1
        end
        head = head + 1
    # refine the root tree (TODO, what is this doing?)
    end
    for i in range(size(queue_decode,1), stop=1, step=-1)
        cur = queue_decode[i]
        queue_decode[cur["parent"]]["t"].children[cur["child_index"]] = cur["t"]
    end
    return to_list(queue_decode[1]["t"], form_manager)
end

s = ArgParseSettings()
@add_arg_table m begin
    "--gpuid"
        arg_type = Int
        help = "which gpu to use. -1 = use CPU"
        default = 0
    "--temperature"
        help = "temperature of sampling"
        arg_type = Int
        default = 0
    "--sample"
        help = "0 to use max at each timestep (-beam_size=1), 1 to sample at each timestep, 2 to beam search"
        arg_type = Int
        default = 0
    "--beam_size"
        help = "beam size"
        arg_type = Int
        default = 20
    "--display"
        help = "whether display on console"
        arg_type = Int
        default = 1
    "--data_dir"
        help = "data path"
        arg_type = String
        default = "../data/"
    "-input"
        help = "input data filename"
        arg_type = String
        default = "test.t7"
    "--output"
        help =  "input data filename"
        arg_type = String
        default = "output/seq2seq_output.txt"
    "--model"
        help = "model checkpoint to use for sampling"
        arg_type = String
        default = "checkpoint_dir/model_seq2seq"
    "-seed"
        help = "torch manual random number generator seed"
        arg_type = Int
        default = 123
end

args = parse_args(s)

# load the model checkpoint
checkpoint = Knet.load(args["model"])
l_model = checkpoint["lang2logic_model"]

# initialize the vocabulary manager to display text
#managers = pkl.load( open("{}/map.pkl".format(args.data_dir), "rb" ) )
@load  (string(checkpoint["data_dir"]) * "/map.jld2") managers
word_manager, form_manager = managers
# load data
#data = pkl.load(open("{}/test.pkl".format(args.data_dir), "rb"))
@load  (string(checkpoint["data_dir"]) * "/test.jld2") managers
reference_list = []
candidate_list = []
#with open(args.output, "w") as output:
open(args["output"], "w") do output
    # TODO change when running full -- this is to just reproduce the error
    #for i in range(30,50):
    #for i in range(278,280):
    for i in range(1, stop=length(data))
        #print("example {}\n".format(i))
        x = data[i]
        reference = x[2]
        candidate = do_generate(l_model, x[1], word_manager, form_manager, args, checkpoint)
        candidate = [Int(c) for c in candidate]

        num_left_paren = sum(1 for c in candidate if form_manager.idx2symbol[Int(c)] == "(")

        num_right_paren = sum(1 for c in candidate if form_manager.idx2symbol[Int(c)]== ")")
        diff = num_left_paren - num_right_paren
        #print(diff)
        if diff > 0
            for i in range(1, stop= diff)
                push!(candidate, form_manager.symbol2idx[")"])
            end
        elseif diff < 0
            candidate = candidate[1:diff]
        end
        ref_str = convert_to_string(reference, form_manager)
        cand_str = convert_to_string(candidate, form_manager)

        push!(reference_list, reference)
        push!(candidate_list, candidate)
        # print to console
        if args["display"] > 0
            print("results: ")
            print(ref_str)
            print(cand_str)
            print(' ')
        #output.write("{}\n".format(ref_str))
        write(output, string(ref_str) *"\n")
        #output.write("{}\n".format(cand_str))
        write(output, string(cand_str) *"\n")
    end
    val_acc = compute_tree_accuracy(candidate_list, reference_list, form_manager)
    print("ACCURACY = " * string(val_acc) * "\n")
    #output.write("ACCURACY = {}\n".format(val_acc))
    write(output,"ACCURACY = " * string(val_acc) *"\n")

end


text:
function factorize(n::Integer)
    fs = [1, n]
    for i in 2:isqrt(n)
        d, r = divrem(n, i)
        if iszero(r)
            push!(fs, d, i)
        end
    end
    unique!(sort!(fs))
end

function part1(i = 33100000 ÷ 10)
    Iterators.peel(Iterators.dropwhile(n -> sum(factorize(n)) < i, Iterators.countfrom(1)))[1]
end

function part2(i = 33100000 ÷ 11)
    Iterators.peel(Iterators.dropwhile(n -> sum(
            filter(factorize(n)) do x
                50x >= n
            end
        ) < i, Iterators.countfrom(1)))[1]
end

max_stars_repo_name: stellartux/advent-of-code

max_issues_repo_name: stellartux/advent-of-code

max_forks_repo_name: stellartux/advent-of-code


text:
module BA_4B

using NBodyIPs.FastPolys
using StaticArrays

import NBodyIPs.BAInvariants: invariants, invariants_d, invariants_ed, tdegrees, rtdegrees

tdegrees(::Val{4}) =
   (1, 1, 2, 2, 3, 3), (0, 2, 2, 3, 3, 3, 3, 3, 3, 4, 4, 6)

# this is a terrible hack that mostly likely will not work as we want.
# see Slack discussion for more details.
rtdegrees(::Val{4}) =
   ( (1,0), (0,1), (2,0), (1,1), (0, 3), (3, 1) ),
   ( (0,0), (1,1), (0,2), (3,0), (2,1), (1,2), (2,1), (1,2), (0,3), (0,4), (2, 2), (4,2) )

const G_BA_4B = [
[ 1, 2, 3, 4, 5, 6 ]
,[ 2, 3, 1, 6, 4, 5 ]
,[ 3, 1, 2, 5, 6, 4 ]
,]
simplex_permutations(x::SVector{6}) = [x[G_BA_4B[i]] for i=1:3]
# Primary invariants for BA_4B
 # : definitions at the beginning of the file
const P1_1 = (1,3,2,)

const P2_1 = (4,5,6,)

const P3_1 = (1,3,2,)

const P4_1 = (1,3,2,)
const P4_2 = (4,5,6,)

const P5_1 = (4,5,6,)

const P6_1 = (1,3,2,4,6,5)
const P6_2 = (1,3,2,5,4,6)

# Irreducible secondaries for group BA_4B
 # : definitions at the beginning of the file
const IS1_1 = (1,3,2,)
const IS1_2 = (5,6,4,)

const IS2_1 = (4,5,6,)

const IS3_1 = (1,3,2,)
const IS3_2 = (3,2,1,)

const IS4_1 = (1,1,2,)
const IS4_2 = (2,3,3,)
const IS4_3 = (4,5,6,)

const IS5_1 = (4,5,6,)
const IS5_2 = (1,3,2,)

const IS6_1 = (1,3,2,)
const IS6_2 = (5,6,4,)

const IS7_1 = (1,3,2,)
const IS7_2 = (4,5,4,)
const IS7_3 = (5,6,6,)

const IS8_1 = (4,5,6,)
const IS8_2 = (6,4,5,)

# Primary invariants for BA_4B
 # : definitions of the types at the beginning of the file
const P1 = Val((P1_1,))
const P2 = Val((P2_1,))
const P3 = Val((P3_1,))
const P4 = Val((P4_1,P4_2,))
const P5 = Val((P5_1,))
const P6 = Val((P6_1,P6_2))
# Irreducible secondaries for group BA_4B
 # : definitions of the types at the beginning of the file
const IS1 = Val((IS1_1,IS1_2,))
const IS2 = Val((IS2_1,))
const IS3 = Val((IS3_1,IS3_2,))
const IS4 = Val((IS4_1,IS4_2,IS4_3,))
const IS5 = Val((IS5_1,IS5_2,))
const IS6 = Val((IS6_1,IS6_2,))
const IS7 = Val((IS7_1,IS7_2,IS7_3,))
const IS8 = Val((IS8_1,IS8_2,))

function invariants(x1::SVector{6, T}) where {T}
   x2 = x1.*x1
   x3 = x2.*x1
   #------------------------------------------------
   # Primaries
   #------------------------------------------------

# Primary invariants for BA_4B
 # : what goes in the function for the evaluation
P1 = fpoly((x1,) , BA_4B.P1)
P2 = fpoly((x1,) , BA_4B.P2)
P3 = fpoly((x2,) , BA_4B.P3)
P4 = fpoly((x1,x1,) , BA_4B.P4)
P5 = fpoly((x3,) , BA_4B.P5)
P6 = fpoly((x2,x1,) , BA_4B.P6)

   #------------------------------------------------
   # Irreducible secondaries
   #------------------------------------------------

# Irreducible secondaries for group BA_4B
 # : what goes in the function for the evaluation
IS1 = fpoly((x1,x1,) , BA_4B.IS1)
IS2 = fpoly((x2,) , BA_4B.IS2)
IS3 = fpoly((x2,x1,) , BA_4B.IS3)
IS4 = fpoly((x1,x1,x1,) , BA_4B.IS4)
IS5 = fpoly((x2,x1,) , BA_4B.IS5)
IS6 = fpoly((x2,x1,) , BA_4B.IS6)
IS7 = fpoly((x1,x1,x1,) , BA_4B.IS7)
IS8 = fpoly((x2,x1,) , BA_4B.IS8)

   #------------------------------------------------
   # All secondaries
   #------------------------------------------------

SEC1  = 1
SEC2  = IS1
SEC3  = IS2
SEC4  = IS3
SEC5  = IS4
SEC6  = IS5
SEC7  = IS6
SEC8  = IS7
SEC9  = IS8
SEC10  = IS1*IS2
SEC11  = IS2^2
SEC12  = IS4^2

return (@SVector [P1,P2,P3,P4,P5,P6,]), (@SVector [SEC1,SEC2,SEC3,SEC4,SEC5,SEC6,SEC7,SEC8,SEC9,SEC10,SEC11,SEC12,])
 end

function invariants_d(x1::SVector{6, T}) where {T}
   x2 = x1.*x1
   x3 = x2.*x1

   dx1 = @SVector ones(6)
   dx2 = 2 * x1
   dx3 = 3 * x2
   #------------------------------------------------
   # Primaries
   #------------------------------------------------

# Primary invariants for BA_4B
 # : what goes in the function for the derivatives
dP1 = fpoly_d((x1,),(dx1,) , BA_4B.P1)
dP2 = fpoly_d((x1,),(dx1,) , BA_4B.P2)
dP3 = fpoly_d((x2,),(dx2,) , BA_4B.P3)
dP4 = fpoly_d((x1,x1,),(dx1,dx1,) , BA_4B.P4)
dP5 = fpoly_d((x3,),(dx3,) , BA_4B.P5)
dP6 = fpoly_d((x2,x1,),(dx2,dx1,) , BA_4B.P6)

   #------------------------------------------------
   # Irreducible secondaries
   #------------------------------------------------

# Irreducible secondaries for group BA_4B
 # : what goes in the function for the evaluation
IS1 = fpoly((x1,x1,) , BA_4B.IS1)
IS2 = fpoly((x2,) , BA_4B.IS2)
IS3 = fpoly((x2,x1,) , BA_4B.IS3)
IS4 = fpoly((x1,x1,x1,) , BA_4B.IS4)
IS5 = fpoly((x2,x1,) , BA_4B.IS5)
IS6 = fpoly((x2,x1,) , BA_4B.IS6)
IS7 = fpoly((x1,x1,x1,) , BA_4B.IS7)
IS8 = fpoly((x2,x1,) , BA_4B.IS8)

# Irreducible secondaries for group BA_4B
 # : what goes in the function for the derivatives
dIS1 = fpoly_d((x1,x1,),(dx1,dx1,) , BA_4B.IS1)
dIS2 = fpoly_d((x2,),(dx2,) , BA_4B.IS2)
dIS3 = fpoly_d((x2,x1,),(dx2,dx1,) , BA_4B.IS3)
dIS4 = fpoly_d((x1,x1,x1,),(dx1,dx1,dx1,) , BA_4B.IS4)
dIS5 = fpoly_d((x2,x1,),(dx2,dx1,) , BA_4B.IS5)
dIS6 = fpoly_d((x2,x1,),(dx2,dx1,) , BA_4B.IS6)
dIS7 = fpoly_d((x1,x1,x1,),(dx1,dx1,dx1,) , BA_4B.IS7)
dIS8 = fpoly_d((x2,x1,),(dx2,dx1,) , BA_4B.IS8)

   #------------------------------------------------
   # All secondaries
   #------------------------------------------------

dSEC1   = @SVector zeros(6)
dSEC2  = dIS1
dSEC3  = dIS2
dSEC4  = dIS3
dSEC5  = dIS4
dSEC6  = dIS5
dSEC7  = dIS6
dSEC8  = dIS7
dSEC9  = dIS8
dSEC10  =  + dIS1*IS2 + dIS2*IS1
dSEC11  =  + dIS2*2IS2
dSEC12  =  + dIS4*2IS4

return (dP1,dP2,dP3,dP4,dP5,dP6,), (dSEC1,dSEC2,dSEC3,dSEC4,dSEC5,dSEC6,dSEC7,dSEC8,dSEC9,dSEC10,dSEC11,dSEC12,)
 end

function invariants_ed(x1::SVector{6, T}) where {T}
   x2 = x1.*x1
   x3 = x2.*x1

   dx1 = @SVector ones(6)
   dx2 = 2 * x1
   dx3 = 3 * x2
   #------------------------------------------------
   # Primaries
   #------------------------------------------------

# Primary invariants for BA_4B
 # : what goes in the function for the evaluation and derivatives
P1, dP1 = fpoly_ed((x1,),(dx1,) , BA_4B.P1)
P2, dP2 = fpoly_ed((x1,),(dx1,) , BA_4B.P2)
P3, dP3 = fpoly_ed((x2,),(dx2,) , BA_4B.P3)
P4, dP4 = fpoly_ed((x1,x1,),(dx1,dx1,) , BA_4B.P4)
P5, dP5 = fpoly_ed((x3,),(dx3,) , BA_4B.P5)
P6, dP6 = fpoly_ed((x2,x1,),(dx2,dx1,) , BA_4B.P6)

   #------------------------------------------------
   # Irreducible secondaries
   #------------------------------------------------

# Irreducible secondaries for group BA_4B
 # : what goes in the function for the evaluation and derivatives
IS1, dIS1 = fpoly_ed((x1,x1,),(dx1,dx1,) , BA_4B.IS1)
IS2, dIS2 = fpoly_ed((x2,),(dx2,) , BA_4B.IS2)
IS3, dIS3 = fpoly_ed((x2,x1,),(dx2,dx1,) , BA_4B.IS3)
IS4, dIS4 = fpoly_ed((x1,x1,x1,),(dx1,dx1,dx1,) , BA_4B.IS4)
IS5, dIS5 = fpoly_ed((x2,x1,),(dx2,dx1,) , BA_4B.IS5)
IS6, dIS6 = fpoly_ed((x2,x1,),(dx2,dx1,) , BA_4B.IS6)
IS7, dIS7 = fpoly_ed((x1,x1,x1,),(dx1,dx1,dx1,) , BA_4B.IS7)
IS8, dIS8 = fpoly_ed((x2,x1,),(dx2,dx1,) , BA_4B.IS8)

   #------------------------------------------------
   # All secondaries
   #------------------------------------------------

SEC1  = 1
SEC2  = IS1
SEC3  = IS2
SEC4  = IS3
SEC5  = IS4
SEC6  = IS5
SEC7  = IS6
SEC8  = IS7
SEC9  = IS8
SEC10  = IS1*IS2
SEC11  = IS2^2
SEC12  = IS4^2

dSEC1   = @SVector zeros(6)
dSEC2  = dIS1
dSEC3  = dIS2
dSEC4  = dIS3
dSEC5  = dIS4
dSEC6  = dIS5
dSEC7  = dIS6
dSEC8  = dIS7
dSEC9  = dIS8
dSEC10  =  + dIS1*IS2 + dIS2*IS1
dSEC11  =  + dIS2*2IS2
dSEC12  =  + dIS4*2IS4

return (@SVector [P1,P2,P3,P4,P5,P6,]),
       (@SVector [SEC1,SEC2,SEC3,SEC4,SEC5,SEC6,SEC7,SEC8,SEC9,SEC10,SEC11,SEC12,]),
       (@SVector [dP1,dP2,dP3,dP4,dP5,dP6,]),
       (@SVector [dSEC1,dSEC2,dSEC3,dSEC4,dSEC5,dSEC6,dSEC7,dSEC8,dSEC9,dSEC10,dSEC11,dSEC12,])
end

end

max_stars_repo_name: cortner/ManyBodyIPs

max_issues_repo_name: cortner/ManyBodyIPs

max_forks_repo_name: cortner/ManyBodyIPs


text:
HJI_cache_filename = joinpath(@__DIR__(), "BicycleCAvoid.jld2")
if !isfile(HJI_cache_filename)
    download("https://github.com/schmrlng/Pigeon.jl/releases/download/v0.2.0/BicycleCAvoid.jld2", HJI_cache_filename)
end

max_stars_repo_name: Michelle-NYX/CollisionAvoidance

max_issues_repo_name: Michelle-NYX/CollisionAvoidance

max_forks_repo_name: Michelle-NYX/CollisionAvoidance


text:
@testset "base interface" begin

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

@testset "method caching" begin

import InteractiveUtils: _dump_function

# #17057 fallout
@eval @noinline post17057_child(i) = sink(i)
@eval function post17057_parent(arr::Ptr{Int64})
    i = post17057_child(0)
    unsafe_store!(arr, i, i)
end

# bug: default module activation segfaulted on NULL child function if cached=false
params = Base.CodegenParams(cached=false)
_dump_function(post17057_parent, Tuple{Ptr{Int64}},
               #=native=#false, #=wrapper=#false, #=strip=#false,
               #=dump_module=#true, #=syntax=#:att, #=optimize=#false,
               params)

end

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

end


text:
module PanelGeometry

## Package imports
#==========================================================================================#

import Base: +, -, zero

using StaticArrays
using LinearAlgebra
using Rotations
using CoordinateTransformations
using SplitApplyCombine

import ..MathTools: affine_2D, rotation, inverse_rotation, structtolist, sine_spacing, cosine_spacing, partition
import ..AeroMDAO: collocation_point

## Panel setup
#==========================================================================================#

abstract type AbstractPanel end

include("2d_panels.jl")
include("3d_panels.jl")

# Methods on panels
p1(p :: AbstractPanel) = p.p1
p2(p :: AbstractPanel) = p.p2
p3(p :: AbstractPanel) = p.p3
p4(p :: AbstractPanel) = p.p4

transform_panel(panel_1 :: AbstractPanel, panel_2 :: AbstractPanel) = transform_panel(panel_1, collocation_point(panel_2))

distance(panel_1 :: AbstractPanel, panel_2 :: AbstractPanel) = norm(collocation_point(panel_2) - collocation_point(panel_1))

# split_panels(panels :: Vector{<: Panel}) = collect.(span(panel -> panel_location(panel) == "upper", panels))

end


text:
@testset "Utilities" begin

    empty = RegexDict{Number}()
    @test 3.2 == get(empty, "foo", 3.2)
    @test 3.2 == get(empty, :foo, 3.2)
    @test_throws BoundsError empty["foo"]
    @test_throws BoundsError empty[:foo]

    dict = RegexDict((r"^Hello", 52), (r"bar", 32), ("foo", 45), ("^baz", 66))

    @test 52 == get(dict, "Hello", -1)
    @test 52 == get(dict, "Hello World", -1)
    @test -1 == get(dict, "not Hello", -1)
    @test 52 == dict["Hello"]
    @test 52 == dict["Hello World"]

    @test 32 == get(dict, "bar", -1)
    @test 32 == get(dict, "qux bar", -1)
    @test 32 == get(dict, "bar qux", -1)
    @test 32 == dict["bar"]
    @test 32 == dict["qux bar"]
    @test 32 == dict["bar qux"]

    @test 45 == get(dict, "foo", -1)
    @test 45 == get(dict, "foo qux", -1)
    @test 45 == get(dict, "qux foo", -1)

    @test 66 == get(dict, "baz", -1)
    @test 66 == get(dict, "baz World", -1)
    @test -1 == get(dict, "not baz", -1)

    @test_throws BoundsError dict["asdf"]

end


text:
# ## model
include("trajopt_model_wall_4feet.jl")
include(joinpath(@__DIR__, "..", "..", "..", "src/dynamics/centroidal_quadruped_wall/visuals.jl"))
@load joinpath(@__DIR__, "wall_stand_FL_4.jld2") qm um γm bm ψm ηm μm hm

vis = Visualizer()
render(vis)

# ## horizon
T = 21
h = 0.05

# ## centroidal_quadruped
s = get_simulation("centroidal_quadruped_wall", "flat_3D_lc", "flat")
model = s.model
env = s.env
nx = 2 * model.nq
nc = model.nc
nu = model.nu + nc + 4 * nc + nc + 4 * nc + 1
nθ = 93

# ## model
d1 = DTO.Dynamics((y, x, u, w) -> centroidal_quadruped_dynt(
	model, env, [h], y, x, u, w), nx + nθ, nx, nu)

dt = DTO.Dynamics((y, x, u, w) -> centroidal_quadruped_dynt(
	model, env, [h], y, x, u, w), nx + nθ , nx + nθ , nu)

dyn = [d1, [dt for t = 2:T-1]...]

# ## initial conditions
body_height = 0.3
foot_x = 0.17
foot_y = 0.17

function nominal_configuration(model::CentroidalQuadrupedWall)
    qm[end]
end

function final_configuration(model::CentroidalQuadrupedWall)
    x_shift = 0.0
    [
        qm[end][1:9];
        0.25;-foot_y; body_height;
       -foot_x + x_shift; foot_y; 0.0;
       -foot_x + x_shift;-foot_y; 0.0;
    ]
end

q1 = nominal_configuration(model)
qM = nominal_configuration(model)
qT = final_configuration(model)
q_ref = nominal_configuration(model)

visualize!(vis, model, [qT], Δt=h)

x1 = [q1; q1]
xM = [qM; qM]
xT = [qT; qT]
x_ref = [q_ref; q_ref]

q1[10]
rad = qT[10] - q1[10]
θ_rad = range(1.0 * π, stop = 0.5 * π, length=T)
foot_arc_x = rad * cos.(θ_rad) .+ q1[10] .+ rad
foot_arc_z = rad * sin.(θ_rad) .* body_height ./ rad# .+ q1[9]

qT[10]
plot(foot_arc_x, foot_arc_z, aspect_ratio=:equal)

# ## objective
obj = DTO.Cost{Float64}[]
for t = 1:T
    if t == 1
        push!(obj, DTO.Cost((x, u, w) -> begin
            J = 0.0
            J += 0.5 * transpose(x[1:6] - x_ref[1:6]) * Diagonal(10000.0 * ones(6)) * (x[1:6] - x_ref[1:6])

            J += 0.5 * 10000.0 * (x[10] - foot_arc_x[t])^2
            J += 0.5 * 10000.0 * (x[11] + foot_y)^2
            J += 0.5 * 10000.0 * (x[12] - foot_arc_z[t])^2

            J += 0.5 * transpose(u) * Diagonal([1.0 * ones(model.nu); zeros(nu - model.nu)]) * u
            J += 10000.0 * u[end] # slack
            J += 0.5 * transpose(u[model.nu + 8 .+ (1:40)]) * Diagonal(1.0e-3 * ones(40)) * u[model.nu + 8 .+ (1:40)]

            return J
        end,
        nx, nu))
    elseif t == T
        push!(obj, DTO.Cost((x, u, w) -> begin
            J = 0.0
            J += 0.5 * transpose(x[1:6] - xT[1:6]) * Diagonal(10000.0 * ones(6)) * (x[1:6] - xT[1:6])

            J += 0.5 * 10000.0 * (x[10] - foot_arc_x[t])^2
            J += 0.5 * 10000.0 * (x[11] + foot_y)^2
            J += 0.5 * 10000.0 * (x[12] - foot_arc_z[t])^2

            J -= 100.0 * x[36 + model.nu + 5]
            J -= 100.0 * x[36 + model.nu + 6]
            return J
        end, nx + nθ, 0))
    else
        push!(obj, DTO.Cost((x, u, w) -> begin
            J = 0.0

            u_prev = x[nx .+ (1:93)]
            w = (u - u_prev) ./ h
            J += 0.5 * 1.0 * dot(w[1:end-1], w[1:end-1])

            J += 0.5 * transpose(x[1:6] - x_ref[1:6]) * Diagonal(10000.0 * ones(6)) * (x[1:6] - x_ref[1:6])

            J += 0.5 * 10000.0 * (x[10] - foot_arc_x[t])^2
            J += 0.5 * 10000.0 * (x[11] + foot_y)^2
            J += 0.5 * 10000.0 * (x[12] - foot_arc_z[t])^2

            J += 0.5 * transpose(u) * Diagonal([1.0 * ones(model.nu); zeros(nu - model.nu)]) * u
            J += 10000.0 * u[end] # slack
            J += 0.5 * transpose(u[model.nu + 8 .+ (1:40)]) * Diagonal(1.0e-3 * ones(40)) * u[model.nu + 8 .+ (1:40)]

            return J
        end, nx + nθ, nu))
    end
end
# function obj1(x, u, w)
# 	J = 0.0
# 	J += 0.5 * transpose(x[1:nx] - x_ref) * Diagonal([1000.0 * ones(6); 100.0 * ones(nx-6)]) * (x[1:nx] - x_ref)
# 	J += 0.5 * transpose(u) * Diagonal([1.0 * ones(model.nu); zeros(nu - model.nu)]) * u
#     J += 1000.0 * u[end] # slack
#     J += 0.5 * transpose(u[model.nu + 8 .+ (1:40)]) * Diagonal(1.0e-3 * ones(40)) * u[model.nu + 8 .+ (1:40)]

# 	return J
# end

# function objt(x, u, w)
# 	J = 0.0

#     u_prev = x[nx .+ (1:93)]
#     w = (u - u_prev) ./ h
#     J += 0.5 * 1.0 * dot(w[1:end-1], w[1:end-1])

# 	J += 0.5 * transpose(x[1:nx] - x_ref) * Diagonal([1000.0 * ones(6); 100.0 * ones(nx-6)]) * (x[1:nx] - x_ref)
# 	J += 0.5 * transpose(u) * Diagonal([1.0 * ones(model.nu); zeros(nu - model.nu)]) * u
#     J += 1000.0 * u[end] # slack
#     J += 0.5 * transpose(u[model.nu + 8 .+ (1:40)]) * Diagonal(1.0e-3 * ones(40)) * u[model.nu + 8 .+ (1:40)]

# 	return J
# end

# function objT(x, u, w)
# 	J = 0.0
# 	J += 0.5 * transpose(x[1:nx] - xT) * Diagonal([1000.0 * ones(6); 1000.0 * ones(nx-6)]) * (x[1:nx] - xT)
#     J -= 100.0 * x[36 + model.nu + 5]
#     return J
# end

# c1 = DTO.Cost(obj1, nx, nu)
# ct = DTO.Cost(objt, nx + nθ , nu)
# cT = DTO.Cost(objT, nx + nθ , 0)
# obj = [c1, [ct for t = 2:T-1]..., cT];

# ## constraints
ql = q1
qu = q1

# initial condition
xl1 = [q1; q1]
xu1 = [q1; q1]
xlt = [-Inf * ones(nx); -Inf * ones(nθ)]
xut = [Inf * ones(nx); Inf * ones(nθ)]

# final condition
xlT = [-Inf * ones(nx); -Inf * ones(nθ)]
xuT = [Inf * ones(nx); Inf * ones(nθ)]

ul = [-Inf * ones(model.nu); zeros(nu - model.nu)]
uu = [Inf * ones(model.nu); Inf * ones(nu - model.nu)]

bnd1 = DTO.Bound(nx, nu, state_lower=xl1, state_upper=xu1, action_lower=ul, action_upper=uu)
bndt = DTO.Bound(nx + nθ , nu, state_lower=xlt, state_upper=xut, action_lower=ul, action_upper=uu)
bndT = DTO.Bound(nx + nθ , 0, state_lower=xlT, state_upper=xuT)
bnds = [bnd1, [bndt for t = 2:T-1]..., bndT];

contact_constraints_equality(model, env, h, rand(nx), rand(nu), zeros(0))
contact_constraints_inequality_1(model, env, h, rand(nx), rand(nu), zeros(0))
contact_constraints_inequality_t(model, env, h, rand(nx + nθ ), rand(nu), zeros(0))
contact_constraints_inequality_T(model, env, h, rand(nx + nθ ), rand(nu), zeros(0))

function constraints_1(x, u, w)
    [
     # equality (32)
     contact_constraints_equality(model, env, h, x, u, w);
     # inequality (56)
     contact_constraints_inequality_1(model, env, h, x, u, w);
     x[6 .+ (1:12)] - q1[6 .+ (1:12)];
     x[18 + 6 .+ (1:12)] - q1[6 .+ (1:12)];
    ]
end

function constraints_t(x, u, w)
    [
     # equality (32)
     contact_constraints_equality(model, env, h, x, u, w);
     # inequality (64)
     contact_constraints_inequality_t(model, env, h, x, u, w);
     x[6 .+ (1:3)] - q1[6 .+ (1:3)];
     x[6 .+ (7:12)] - q1[6 .+ (7:12)];
     x[18 + 6 .+ (1:3)] - q1[6 .+ (1:3)];
     x[18 + 6 .+ (7:12)] - q1[6 .+ (7:12)];
    #  x[18 + 6 .+ (4:12)] - q1[6 .+ (4:12)];
    #  x[18 + 3] - q1[3];
    ]
end

function constraints_T(x, u, w)
    [
     # inequality (16)
     contact_constraints_inequality_T(model, env, h, x, u, w);
     x[6 .+ (1:3)] - qT[6 .+ (1:3)];
     x[6 .+ (7:12)] - qT[6 .+ (7:12)];

     x[18 .+ (1:3)] - qT[1:3];
     x[18 + 6 .+ (1:3)] - qT[6 .+ (1:3)];
     x[18 + 6 .+ (7:12)] - qT[6 .+ (7:12)];
     x[18 + 9 .+ 1] - qT[9 + 1];
     x[18 + 9 .+ 2] - qT[9 + 2];
     x[18 + 9 .+ 3] - qT[9 + 3];
     x[0 .+ (3:6)] - qT[3:6];
     x[18 .+ (3:6)] - qT[3:6];
    ]
end

con1 = DTO.Constraint(constraints_1, nx, nu, indices_inequality=collect(32 .+ (1:56)))
cont = DTO.Constraint(constraints_t, nx + nθ , nu, indices_inequality=collect(32 .+ (1:64)))
conT = DTO.Constraint(constraints_T, nx + nθ , nu, indices_inequality=collect(0 .+ (1:16)))
cons = [con1, [cont for t = 2:T-1]..., conT];

# ## problem
direct_solver = DTO.Solver(dyn, obj, cons, bnds,
    options=DTO.Options(
        tol=1.0e-3,
        constr_viol_tol=1.0e-3,
        max_iter=3000,
        ))

# ## initialize
x_interpolation = [x1, [[x1; zeros(nθ)] for t = 2:T]...]
u_guess = [1.0e-4 * rand(nu) for t = 1:T-1] # may need to run more than once to get good trajectory
DTO.initialize_states!(direct_solver, x_interpolation)
DTO.initialize_controls!(direct_solver, u_guess)

# ## solve
@time DTO.solve!(direct_solver)

# ## solution
x_sol, u_sol = DTO.get_trajectory(direct_solver)
@show x_sol[1]
@show x_sol[T]
maximum([u[nu] for u in u_sol[1:end-1]])

plot([x_sol[t][36 + model.nu + 5] for t = 2:T])
plot([x_sol[t][36 + model.nu + 6] for t = 2:T])

# plot([x[18 + 7] for x in x_sol], [x[18 + 8] for x in x_sol])
x_sol[end][36 + model.nu + 5]
x_sol[end][36 + model.nu + 6]
x_sol[end][36 + model.nu .+ (1:4)]

# ## visualize
vis = Visualizer()
render(vis)
# q_sol = state_to_configuration([x[1:nx] for x in x_sol])
visualize!(vis, model, [x_sol[1][1:18], [x[18 .+ (1:18)] for x in x_sol]...], Δt=h)

N_first = 10
N_last = 20
q_opt = [[x_sol[1][model.nq .+ (1:model.nq)] for t = 1:N_first]..., x_sol[1][1:model.nq], [x[model.nq .+ (1:model.nq)] for x in x_sol]..., [x_sol[end][model.nq .+ (1:model.nq)] for t = 1:N_last]...]
v_opt = [[(x_sol[1][model.nq .+ (1:model.nq)] - x_sol[1][0 .+ (1:model.nq)]) ./ h for t = 1:N_first]..., [(x[model.nq .+ (1:model.nq)] - x[0 .+ (1:model.nq)]) ./ h for x in x_sol]..., [(x_sol[end][model.nq .+ (1:model.nq)] - x_sol[end][0 .+ (1:model.nq)]) ./ h for t = 1:N_last]...]
u_opt = [[u_sol[1][1:model.nu] for t = 1:N_first]..., [u[1:model.nu] for u in u_sol]..., [u_sol[end][1:model.nu] for t = 1:N_last]...]
γ_opt = [[u_sol[1][model.nu .+ (1:8)] for t = 1:N_first]..., [u[model.nu .+ (1:8)] for u in u_sol]..., [u_sol[end][model.nu .+ (1:8)] for t = 1:N_last]...]
b_opt = [[u_sol[1][model.nu + 8 .+ (1:32)] for t = 1:N_first]..., [u[model.nu + 8 .+ (1:32)] for u in u_sol]..., [u_sol[end][model.nu + 8 .+ (1:32)] for t = 1:N_last]...]
ψ_opt = [[u_sol[1][model.nu + 8 + 32 .+ (1:8)] for t = 1:N_first]..., [u[model.nu + 8 + 32 .+ (1:8)] for u in u_sol]..., [u_sol[end][model.nu + 8 + 32 .+ (1:8)] for t = 1:N_last]...]
η_opt = [[u_sol[1][model.nu + 8 + 32 + 8 .+ (1:32)] for t = 1:N_first]..., [u[model.nu + 8 + 32 + 8 .+ (1:32)] for u in u_sol]..., [u_sol[end][model.nu + 8 + 32 + 8 .+ (1:32)] for t = 1:N_last]...]

qm = q_opt
vm = v_opt
um = u_opt
γm = γ_opt
bm = b_opt
ψm = ψ_opt
ηm = η_opt

μm = model.μ_world
hm = h
timesteps = range(0.0, stop=(h * (length(qm) - 2)), length=(length(qm) - 2))
plot(hcat(qm...)', labels="")
plot(timesteps, hcat(um...)', labels="")
plot(timesteps, hcat(γm...)', labels="")
plot(timesteps, hcat(bm...)', labels="")
plot(timesteps, hcat(ψm...)', labels="")
plot(timesteps, hcat(ηm...)', labels="")

using JLD2

@save joinpath(@__DIR__, "wall_stand_FL_4.jld2") qm um γm bm ψm ηm μm hm
@load joinpath(@__DIR__, "wall_stand_FL_4.jld2") qm um γm bm ψm ηm μm hm

# q_opt = [x_sol[1][1:model.nq], [x[model.nq .+ (1:model.nq)] for x in x_sol]...]
# v_opt = [(x[model.nq .+ (1:model.nq)] - x[0 .+ (1:model.nq)]) ./ h for x in x_sol]
# u_opt = [u[1:model.nu] for u in u_sol]
# λ_opt = [u[model.nu .+ (1:5)] for u in u_sol]
# b_opt = [u[model.nu + 5 .+ (1:20)] for u in u_sol]

# q_opt = [x_sol[1][1:model.nq], [x[model.nq .+ (1:model.nq)] for x in x_sol]...]
# v_opt = [(x[model.nq .+ (1:model.nq)] - x[0 .+ (1:model.nq)]) ./ h for x in x_sol]
# u_opt = [u[1:model.nu] for u in u_sol]
# γ_opt = [u[model.nu .+ (1:5)] for u in u_sol]
# b_opt = [u[model.nu + 5 .+ (1:20)] for u in u_sol]
# ψ_opt = [u[model.nu + 5 + 20 .+ (1:5)] for u in u_sol]
# η_opt = [u[model.nu + 5 + 20 + 5 .+ (1:20)] for u in u_sol]

# qm = q_opt
# vm = v_opt
# um = u_opt
# γm = γ_opt
# bm = b_opt
# ψm = ψ_opt
# ηm = η_opt

# μm = model.μ_world
# hm = h
# timesteps = range(0.0, stop=(h * (length(qm) - 2)), length=(length(qm) - 2))
# plot(hcat(qm...)', labels="")
# plot(timesteps, hcat(um...)', labels="")
# plot(timesteps, hcat(γm...)', labels="")
# plot(timesteps, hcat(bm...)', labels="")
# plot(timesteps, hcat(ψm...)', labels="")
# plot(timesteps, hcat(ηm...)', labels="")

# using JLD2
# @save joinpath(@__DIR__, "wall_stand_FL.jld2") qm um γm bm ψm ηm μm hm
# @load joinpath(@__DIR__, "wall_stand_FL.jld2") qm um γm bm ψm ηm μm hm


text:
using Stheno: AbstractGP, BlockData, blocks
using FiniteDifferences: j′vp

const _rtol = 1e-10
const _atol = 1e-10

_to_psd(A::Matrix{<:Real}) = A * A' + I
_to_psd(a::Vector{<:Real}) = exp.(a) .+ 1
_to_psd(σ::Real) = exp(σ) + 1

Base.length(::Nothing) = 0

function print_adjoints(adjoint_ad, adjoint_fd, rtol, atol)
    @show typeof(adjoint_ad), typeof(adjoint_fd)
    adjoint_ad, adjoint_fd = to_vec(adjoint_ad)[1], to_vec(adjoint_fd)[1]
    println("atol is $atol, rtol is $rtol")
    println("ad, fd, abs, rel")
    abs_err = abs.(adjoint_ad .- adjoint_fd)
    rel_err = abs_err ./ adjoint_ad
    display([adjoint_ad adjoint_fd abs_err rel_err])
    println()
end

# AbstractArrays.
function FiniteDifferences.to_vec(x::ColVecs{<:Real})
    x_vec, back = to_vec(x.X)
    return x_vec, x_vec -> ColVecs(back(x_vec))
end

function FiniteDifferences.to_vec(X::BlockArray)
    X_Array = Array(X)
    x, X_Array_from_vec = to_vec(X_Array)
    function BlockArray_from_vec(x::Vector)
        X_Array = X_Array_from_vec(x)
        return BlockArray(X_Array, axes(X))
    end
    return x, BlockArray_from_vec
end

function FiniteDifferences.to_vec(x::BlockData)
    x_vecs, x_backs = zip(map(to_vec, blocks(x))...)
    sz = cumsum([map(length, x_vecs)...])
    return vcat(x_vecs...), function(v)
        return BlockData([x_backs[n](v[sz[n]-length(blocks(x)[n])+1:sz[n]])
            for n in 1:length(blocks(x))])
    end
end

function FiniteDifferences.to_vec(X::T) where T<:Union{Adjoint,Transpose}
    U = T.name.wrapper
    return vec(Matrix(X)), x_vec->U(permutedims(reshape(x_vec, size(X))))
end

Base.zero(d::Dict) = Dict([(key, zero(val)) for (key, val) in d])
Base.zero(x::Array) = zero.(x)
Base.zero(x::SubArray) = zero.(x)
Base.zero(x::ColVecs) = ColVecs(zero(x.X))

# My version of isapprox
function fd_isapprox(x_ad::Nothing, x_fd, rtol, atol)
    return fd_isapprox(x_fd, zero(x_fd), rtol, atol)
end
function fd_isapprox(x_ad::AbstractArray, x_fd::AbstractArray, rtol, atol)
    return all(fd_isapprox.(x_ad, x_fd, rtol, atol))
end
function fd_isapprox(x_ad::Real, x_fd::Real, rtol, atol)
    return isapprox(x_ad, x_fd; rtol=rtol, atol=atol)
end
function fd_isapprox(x_ad::NamedTuple, x_fd, rtol, atol)
    f = (x_ad, x_fd)->fd_isapprox(x_ad, x_fd, rtol, atol)
    return all([f(getfield(x_ad, key), getfield(x_fd, key)) for key in keys(x_ad)])
end
function fd_isapprox(x_ad::Tuple, x_fd::Tuple, rtol, atol)
    return all(map((x, x′)->fd_isapprox(x, x′, rtol, atol), x_ad, x_fd))
end
function fd_isapprox(x_ad::Dict, x_fd::Dict, rtol, atol)
    return all([fd_isapprox(get(()->nothing, x_ad, key), x_fd[key], rtol, atol) for
        key in keys(x_fd)])
end

function adjoint_test(
    f, ȳ, x...;
    rtol=_rtol,
    atol=_atol,
    fdm=FiniteDifferences.central_fdm(5, 1),
    print_results=false,
)
    # Compute forwards-pass and j′vp.
    y, back = Zygote.pullback(f, x...)
    @timeit to "adj_ad" adj_ad = back(ȳ)
    @timeit to "adj_fd" adj_fd = j′vp(fdm, f, ȳ, x...)

    # Check that forwards-pass agrees with plain forwards-pass.
    @test y ≈ f(x...)

    # Check that ad and fd adjoints (approximately) agree.
    print_results && print_adjoints(adj_ad, adj_fd, rtol, atol)
    @test fd_isapprox(adj_ad, adj_fd, rtol, atol)
end

"""
    abstractgp_interface_tests(
        f::AbstractGP,
        f′::AbstractGP,
        x0::AbstractVector,
        x1::AbstractVector,
        x2::AbstractVector,
        x3::AbstractVector;
        atol=1e-9, rtol=1e-9,
    )

Check that the `AbstractGP` interface is at least implemented for `f` and is
self-consistent. `x0` and `x1` must be valid inputs for `f`. `x2` and `x3` must be a valid
input for `f′`.
"""
function abstractgp_interface_tests(
    f::AbstractGP,
    f′::AbstractGP,
    x0::AbstractVector,
    x1::AbstractVector,
    x2::AbstractVector,
    x3::AbstractVector;
    atol=1e-9, rtol=1e-9,
)
    m = mean(f, x0)
    @test m isa AbstractVector{<:Real}
    @test length(m) == length(x0)

    @assert length(x0) ≠ length(x1)
    @assert length(x0) ≠ length(x2)
    @assert length(x0) == length(x3)

    # Check that binary cov conforms to the API
    K_ff′_x0_x2 = cov(f, f′, x0, x2)
    @test K_ff′_x0_x2 isa AbstractMatrix{<:Real}
    @test size(K_ff′_x0_x2) == (length(x0), length(x2))
    @test K_ff′_x0_x2 ≈ cov(f′, f, x2, x0)'

    # Check that unary cov is consistent with binary cov and conforms to the API
    K_x0 = cov(f, x0)
    @test K_x0 isa AbstractMatrix{<:Real}
    @test size(K_x0) == (length(x0), length(x0))
    @test K_x0 ≈ cov(f, f, x0, x0) atol=atol rtol=rtol
    @test minimum(eigvals(K_x0)) > -atol
    @test K_x0 ≈ K_x0' atol=atol rtol=rtol

    # Check that single-process binary cov is consistent with binary-process binary-cov
    K_x0_x1 = cov(f, x0, x1)
    @test K_x0_x1 isa AbstractMatrix{<:Real}
    @test size(K_x0_x1) == (length(x0), length(x1))
    @test K_x0_x1 ≈ cov(f, f, x0, x1)

    # Check that binary var conforms to the API and is consistent with binary cov
    K_x0_x3_diag = var(f, f′, x0, x3)
    @test K_x0_x3_diag isa AbstractVector{<:Real}
    @test length(K_x0_x3_diag) == length(x0)
    @test K_x0_x3_diag ≈ diag(cov(f, f′, x0, x3)) atol=atol rtol=rtol
    @test K_x0_x3_diag ≈ var(f′, f, x3, x0) atol=atol rtol=rtol

    # Check that unary-binary var is consistent.
    K_x0_x0_diag = var(f, x0, x0)
    @test K_x0_x0_diag isa AbstractVector{<:Real}
    @test length(K_x0_x0_diag) == length(x0)
    @test K_x0_x0_diag ≈ diag(cov(f, x0, x0)) atol=atol rtol=rtol

    # Check that unary var conforms to the API and is consistent with unary cov
    K_x0_diag = var(f, x0)
    @test K_x0_diag isa AbstractVector{<:Real}
    @test length(K_x0_diag) == length(x0)
    @test K_x0_diag ≈ diag(cov(f, x0)) atol=atol rtol=rtol
end


text:
#mpo.jl
"""
    struct MPO{S<:EuclideanSpace} <: AbstractMPS{S}

This type include the Hamiltonian and also other operators which are in the form of MPO.
"""
struct MPO{S<:EuclideanSpace} <: AbstractMPS{S}
    Ws::Vector{<:AbstractTensorMap{S, 2, 2}} # codomain= left*down, domain = up*right
    function MPO{S}(Ws::Vector{<:AbstractTensorMap{S, 2, 2}}) where {S<:EuclideanSpace}
        new(Ws)
    end
end

MPO(Ws) = MPO{spacetype(Ws[1])}(Ws)
"""
    function TransverseFieldIsing(L::Int,J::Real,g::Real)

``H = - J ∑_i σˣ_i σˣ_{i+1}  - g ∑_i σᶻ_i``
This model has a Z2 symmetry which is ``∏_i σᶻ_i``.
"""
function TransverseFieldIsing(L::Int,J::Real,g::Real)
    final_position = 2
    V_virtual = ℤ₂Space(0=>2, 1=>1)
    V_physical = ℤ₂Space(0=>1,1=>1)
    chi_virtual = dim(V_virtual)
    chi_physical = dim(V_physical)
    data = zeros(chi_virtual,chi_physical,chi_physical,chi_virtual)
    Id = [1 0; 0 1]
    sigmax = [0 1; 1 0]
    sigmay = [0 -1im; 1im 0]
    sigmaz = [1 0; 0 -1]
    data[1,:,:,1] = Id
    data[2,:,:,1] = -g*sigmaz
    data[3,:,:,1] = -J*sigmax
    data[2,:,:,3] = sigmax
    data[2,:,:,2] = Id
    W = TensorMap(data, V_virtual ⊗ V_physical, V_physical ⊗ V_virtual)
    return MPO([W for i in 1:L]), final_position
end

"""
    function Heisenberg(N::Int)
"""
function Heisenberg(N::Int)
end


text:
export RationalPoly

# TODO We should take gcd between numerator and denominator
struct RationalPoly{NT <: APL, DT <: APL}
    num::NT
    den::DT
end

# This constructor is called from LinearAlgebra in the method Matrix{T}(s::UniformScaling{Bool}, dims)
RationalPoly{NT,DT}(x::Bool) where {NT,DT} = ifelse(x, one(RationalPoly{NT,DT}), zero(RationalPoly{NT,DT}))

Base.numerator(r::RationalPoly) = r.num
Base.denominator(r::RationalPoly) = r.den

Base.convert(::Type{RationalPoly{NT, DT}}, q::RationalPoly{NT, DT}) where {NT, DT} = q
Base.convert(::Type{RationalPoly{NTout, DTout}}, q::RationalPoly{NTin, DTin}) where {NTout, DTout, NTin, DTin} = convert(NTout, q.num) / convert(DTout, q.den)
#function Base.convert(::Type{RationalPoly{NT, DT}}, p::NT) where {NT, DT}
#    p / one(DT)
#end
function Base.convert(::Type{RationalPoly{NT, DT}}, p::APL) where {NT, DT}
    convert(NT, p) / one(DT)
end
function Base.convert(::Type{RationalPoly{NT, DT}}, α) where {NT, DT}
    convert(NT, α) / one(DT)
    #convert(RationalPoly{NT, DT}, convert(NT, α))
end

Base.inv(r::RationalPoly) = r.den / r.num
Base.inv(p::APL{T}) where T = one(T) / p
Base.:/(r::RationalPoly, p) = r.num / (r.den * p)
Base.:/(r::RationalPoly, s::RationalPoly) = (r.num * s.den) / (s.num * r.den)
function Base.:/(num::NT, den::DT) where {NT <: APL, DT <: APL}
    RationalPoly{NT, DT}(num, den)
end
function Base.:/(num, den::APL)
    constantterm(num, den) / den
end
# Polynomial divided by coefficient is a polynomial not a rational polynomial
# (1/den) * num would not be correct in case of noncommutative coefficients
Base.:/(num::APL, den) = mapcoefficientsnz(α -> α/den, num)

function Base.:+(r::RationalPoly, s::RationalPoly)
    (r.num*s.den + r.den*s.num) / (r.den * s.den)
end
function _plus(r::RationalPoly, p)
    (p*r.den + r.num) / r.den
end
Base.:+(p::APL, r::RationalPoly) = _plus(r, p)
Base.:+(r::RationalPoly, p::APL) = _plus(r, p)
Base.:+(r::RationalPoly, α) = _plus(r, α)
Base.:+(α, r::RationalPoly) = _plus(r, α)
function Base.:-(r::RationalPoly, s::RationalPoly)
    (r.num*s.den - r.den*s.num) / (r.den * s.den)
end
_minus(p, s::RationalPoly) = (p * s.den - s.num) / s.den
_minus(s::RationalPoly, p) = (s.num - p * s.den) / s.den
Base.:-(p::APL, r::RationalPoly) = _minus(p, r)
Base.:-(r::RationalPoly, p::APL) = _minus(r, p)
Base.:-(r::RationalPoly, α) = _minus(r, α)
Base.:-(α, r::RationalPoly) = _minus(α, r)
Base.:-(r::RationalPoly) = (-r.num) / r.den

Base.:*(r::RationalPoly, s::RationalPoly) = (r.num*s.num) / (r.den*s.den)
Base.:*(p::APL, r::RationalPoly)          = (p * r.num) / r.den
Base.:*(r::RationalPoly, p::APL)          = (r.num * p) / r.den
Base.:*(α, r::RationalPoly)               = (α * r.num) / r.den
Base.:*(r::RationalPoly, α)               = (r.num * α) / r.den

Base.zero(r::RationalPoly) = zero(typeof(r))
Base.zero(::Type{RationalPoly{NT, DT}}) where {NT, DT} = zero(NT) / one(DT)
Base.one(r::RationalPoly) = one(typeof(r))
Base.one(::Type{RationalPoly{NT, DT}}) where {NT, DT} = one(NT) / one(DT)


text:
using ArgParse
using Compat
import Compat: UTF8String, ASCIIString, view, readstring, foreach

include("CLI.jl")
include("commands.jl")

function main()
    parsed_args = parse_args(build_arg_table())
    command = parsed_args["%COMMAND%"]
    println(typeof(parsed_args[command]["N"]))

    if command == "help"
        foreach(x -> println(x), help())
    elseif command == "train_model"
        train_model(parsed_args[command]["N"],
        parsed_args[command]["a"],
        parsed_args[command]["hl_size"],
        parsed_args[command]["w"],
        parsed_args[command]["b"],
        parsed_args[command]["C"],
        parsed_args[command]["pos_seq"],
        parsed_args[command]["total_seq"],
        parsed_args[command]["output_file"]
        )
    elseif command == "predict"
        predict(parsed_args[command]["N"],
        parsed_args[command]["a"],
        parsed_args[command]["hl_size"],
        parsed_args[command]["w"],
        parsed_args[command]["b"],
        parsed_args[command]["C"],
        parsed_args[command]["pos_seq"],
        parsed_args[command]["total_seq"],
        parsed_args[command]["predict_seq"],
        parsed_args[command]["output_file"]
        )
    elseif command == "simple_encoder"
        simple_encoder(parsed_args[command]["N"],
            parsed_args[command]["a"],
            parsed_args[command]["output_file"])
    else
        println("Command not recognized")
    end
end

# fire up simulation if run using command line
if !isinteractive()
    main()
end


text:
# # 2D forced-dissipative turbulence budgets
#
#md # This example can viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/literated/twodnavierstokes_stochasticforcing_budgets.ipynb).
#
# A simulation of forced-dissipative two-dimensional turbulence. We solve the
# two-dimensional vorticity equation with stochastic excitation and dissipation in
# the form of linear drag and hyperviscosity. As a demonstration, we compute how
# each of the forcing and dissipation terms contribute to the energy and the
# enstrophy budgets.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add GeophysicalFlows, Random, Printf, Plots"
# ```

# ## Let's begin
# Let's load `GeophysicalFlows.jl` and some other needed packages.
#
using GeophysicalFlows, Random, Printf, Plots
parsevalsum = FourierFlows.parsevalsum

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)
nothing # hide

# ## Numerical, domain, and simulation parameters
#
# First, we pick some numerical and physical parameters for our model.

 n, L  = 256, 2π              # grid resolution and domain length
 ν, nν = 2e-7, 2              # hyperviscosity coefficient and hyperviscosity order
 μ, nμ = 1e-1, 0              # linear drag coefficient
dt, tf = 0.005, 0.2 / μ       # timestep and final time
    nt = round(Int, tf / dt)  # total timesteps
    ns = 4                    # how many intermediate times we want to plot
nothing # hide

# ## Forcing

# We force the vorticity equation with stochastic excitation that is delta-correlated in time
# and while spatially homogeneously and isotropically correlated. The forcing has a spectrum
# with power in a ring in wavenumber space of radius ``k_f`` (`forcing_wavenumber`) and width
# ``δ_f`` (`forcing_bandwidth`), and it injects energy per unit area and per unit time
# equal to ``\varepsilon``. That is, the forcing covariance spectrum is proportional to
# ``\exp{[-(|\bm{k}| - k_f)^2 / (2 δ_f^2)]}``.

forcing_wavenumber = 14.0 * 2π/L  # the forcing wavenumber, `k_f`, for a spectrum that is a ring in wavenumber space
forcing_bandwidth  = 1.5  * 2π/L  # the width of the forcing spectrum, `δ_f`
ε = 0.1                           # energy input rate by the forcing

grid = TwoDGrid(dev, n, L)

K = @. sqrt(grid.Krsq)             # a 2D array with the total wavenumber

forcing_spectrum = @. exp(-(K - forcing_wavenumber)^2 / (2 * forcing_bandwidth^2))
@CUDA.allowscalar forcing_spectrum[grid.Krsq .== 0] .= 0 # ensure forcing has zero domain-average

ε0 = parsevalsum(forcing_spectrum .* grid.invKrsq / 2, grid) / (grid.Lx * grid.Ly)
@. forcing_spectrum *= ε/ε0        # normalize forcing to inject energy at rate ε
nothing # hide

# We reset of the random number generator for reproducibility
if dev==CPU(); Random.seed!(1234); else; CUDA.seed!(1234); end
nothing # hide

# Next we construct function `calcF!` that computes a forcing realization every timestep.
# First we make sure that if `dev=GPU()`, then `CUDA.rand()` function is called for random
# numbers uniformly distributed between 0 and 1.
random_uniform = dev==CPU() ? rand : CUDA.rand

function calcF!(Fh, sol, t, clock, vars, params, grid)
  Fh .= sqrt.(forcing_spectrum) .* exp.(2π * im * random_uniform(eltype(grid), size(sol))) ./ sqrt(clock.dt)

  return nothing
end
nothing # hide

# ## Problem setup
# We initialize a `Problem` by providing a set of keyword arguments. The
# `stepper` keyword defines the time-stepper to be used.
prob = TwoDNavierStokes.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt, stepper="ETDRK4",
                                calcF=calcF!, stochastic=true)
nothing # hide

# Define some shortcuts for convenience.
sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

x, y = grid.x, grid.y
nothing # hide

# First let's see how a forcing realization looks like. Function `calcF!()` computes
# the forcing in Fourier space and saves it into variable `vars.Fh`, so we first need to
# go back to physical space.
#
# Note that when plotting, we decorate the variable to be plotted with `Array()` to make sure
# it is brought back on the CPU when the variable lives on the GPU.
calcF!(vars.Fh, sol, 0.0, clock, vars, params, grid)

heatmap(x, y, Array(irfft(vars.Fh, grid.nx)'),
     aspectratio = 1,
               c = :balance,
            clim = (-200, 200),
           xlims = (-L/2, L/2),
           ylims = (-L/2, L/2),
          xticks = -3:3,
          yticks = -3:3,
          xlabel = "x",
          ylabel = "y",
           title = "a forcing realization",
      framestyle = :box)

# ## Setting initial conditions

# Our initial condition is a fluid at rest.
TwoDNavierStokes.set_ζ!(prob, ArrayType(dev)(zeros(grid.nx, grid.ny)))

# ## Diagnostics

# Create Diagnostics; the diagnostics are aimed to probe the energy and enstrophy budgets.
E  = Diagnostic(TwoDNavierStokes.energy,                               prob, nsteps=nt) # energy
Rᵋ = Diagnostic(TwoDNavierStokes.energy_dissipation_hypoviscosity,     prob, nsteps=nt) # energy dissipation by drag μ
Dᵋ = Diagnostic(TwoDNavierStokes.energy_dissipation_hyperviscosity,    prob, nsteps=nt) # energy dissipation by drag μ
Wᵋ = Diagnostic(TwoDNavierStokes.energy_work,                          prob, nsteps=nt) # energy work input by forcing
Z  = Diagnostic(TwoDNavierStokes.enstrophy,                            prob, nsteps=nt) # enstrophy
Rᶻ = Diagnostic(TwoDNavierStokes.enstrophy_dissipation_hypoviscosity,  prob, nsteps=nt) # enstrophy dissipation by drag μ
Dᶻ = Diagnostic(TwoDNavierStokes.enstrophy_dissipation_hyperviscosity, prob, nsteps=nt) # enstrophy dissipation by drag μ
Wᶻ = Diagnostic(TwoDNavierStokes.enstrophy_work,                       prob, nsteps=nt) # enstrophy work input by forcing
diags = [E, Dᵋ, Wᵋ, Rᵋ, Z, Dᶻ, Wᶻ, Rᶻ] # a list of Diagnostics passed to `stepforward!` will  be updated every timestep.
nothing # hide

# ## Visualizing the simulation

# We define a function that plots the vorticity field and the evolution of
# the diagnostics: energy, enstrophy, and all terms involved in the energy and
# enstrophy budgets. Last, we also check (by plotting) whether the energy and enstrophy
# budgets are accurately computed, e.g., ``\mathrm{d}E/\mathrm{d}t = W^\varepsilon -
# R^\varepsilon - D^\varepsilon``.

function computetendencies_and_makeplot(prob, diags)
  sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  TwoDNavierStokes.updatevars!(prob)

  E, Dᵋ, Wᵋ, Rᵋ, Z, Dᶻ, Wᶻ, Rᶻ = diags

  clocktime = round(μ * clock.t, digits=2)

  dEdt_numerical = (E[2:E.i] - E[1:E.i-1]) / clock.dt # numerical first-order approximation of energy tendency
  dZdt_numerical = (Z[2:Z.i] - Z[1:Z.i-1]) / clock.dt # numerical first-order approximation of enstrophy tendency

  dEdt_computed = Wᵋ[2:E.i] + Dᵋ[1:E.i-1] + Rᵋ[1:E.i-1]
  dZdt_computed = Wᶻ[2:Z.i] + Dᶻ[1:Z.i-1] + Rᶻ[1:Z.i-1]

  residual_E = dEdt_computed - dEdt_numerical
  residual_Z = dZdt_computed - dZdt_numerical

  εᶻ = parsevalsum(forcing_spectrum / 2, grid) / (grid.Lx * grid.Ly)

  pζ = heatmap(x, y, Array(vars.ζ'),
            aspectratio = 1,
            legend = false,
                 c = :viridis,
              clim = (-25, 25),
             xlims = (-L/2, L/2),
             ylims = (-L/2, L/2),
            xticks = -3:3,
            yticks = -3:3,
            xlabel = "μt",
            ylabel = "y",
             title = "∇²ψ(x, y, μt=" * @sprintf("%.2f", μ * clock.t) * ")",
        framestyle = :box)

  pζ = plot(pζ, size = (400, 400))

  t = E.t[2:E.i]

  p1E = plot(μ * t, [Wᵋ[2:E.i] ε.+0*t Dᵋ[1:E.i-1] Rᵋ[1:E.i-1]],
             label = ["energy work, Wᵋ" "ensemble mean energy work, <Wᵋ>" "dissipation, Dᵋ" "drag, Rᵋ = - 2μE"],
         linestyle = [:solid :dash :solid :solid],
         linewidth = 2,
             alpha = 0.8,
            xlabel = "μt",
            ylabel = "energy sources and sinks")

  p2E = plot(μ * t, [dEdt_computed, dEdt_numerical],
           label = ["computed Wᵋ-Dᵋ" "numerical dE/dt"],
       linestyle = [:solid :dashdotdot],
       linewidth = 2,
           alpha = 0.8,
          xlabel = "μt",
          ylabel = "dE/dt")

  p3E = plot(μ * t, residual_E,
           label = "residual dE/dt = computed - numerical",
       linewidth = 2,
           alpha = 0.7,
          xlabel = "μt")

  t = Z.t[2:E.i]

  p1Z = plot(μ * t, [Wᶻ[2:Z.i] εᶻ.+0*t Dᶻ[1:Z.i-1] Rᶻ[1:Z.i-1]],
           label = ["enstrophy work, Wᶻ" "mean enstrophy work, <Wᶻ>" "enstrophy dissipation, Dᶻ" "enstrophy drag, Rᶻ = - 2μZ"],
       linestyle = [:solid :dash :solid :solid],
       linewidth = 2,
           alpha = 0.8,
          xlabel = "μt",
          ylabel = "enstrophy sources and sinks")

  p2Z = plot(μ * t, [dZdt_computed, dZdt_numerical],
         label = ["computed Wᶻ-Dᶻ" "numerical dZ/dt"],
     linestyle = [:solid :dashdotdot],
     linewidth = 2,
         alpha = 0.8,
        xlabel = "μt",
        ylabel = "dZ/dt")

  p3Z = plot(μ * t, residual_Z,
         label = "residual dZ/dt = computed - numerical",
     linewidth = 2,
         alpha = 0.7,
        xlabel = "μt")

  layout = @layout Plots.grid(3, 2)

  pbudgets = plot(p1E, p1Z, p2E, p2Z, p3E, p3Z, layout=layout, size = (900, 1200))

  return pζ, pbudgets
end
nothing # hide

# ## Time-stepping the `Problem` forward

# Finally, we time-step the `Problem` forward in time.

startwalltime = time()
for i = 1:ns
  stepforward!(prob, diags, round(Int, nt/ns))

  TwoDNavierStokes.updatevars!(prob)

  cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy])

  log = @sprintf("step: %04d, t: %.1f, cfl: %.3f, walltime: %.2f min", clock.step, clock.t,
        cfl, (time()-startwalltime)/60)

  println(log)
end

# ## Plot
# And now let's see what we got. First we plot the final snapshot
# of the vorticity field.

pζ, pbudgets = computetendencies_and_makeplot(prob, diags)

pζ

# And finaly the energy and enstrophy budgets.

pbudgets


text:
#=
struct Layer
    solution ::OffsetVector{Dict{Variable, Coefficient}}
end

struct Fiber
    layers ::Vector{Layer}
end
=#