text:
# download the NMRData repository at https://github.com/AI4DBiological-Systems/NMRData

import NMRDataSetup

using FFTW
import PyPlot
import BSON

PyPlot.close("all")
fig_num = 1

PyPlot.matplotlib["rcParams"][:update](["font.size" => 22, "font.family" => "serif"])

### user inputs.

solvent_ppm_guess = 4.7
solvent_window_ppm = 0.1

# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/NRC"
# project_name = "NRC-glucose-2018"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/NRC/misc/glucose/Sep-25-2018"

# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/NRC"
# project_name = "NRC-dmem-2012"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/NRC/misc/dmem_medium/Oct-22-2012"

# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/BMRB-600-100mM"
# project_name = "D-(+)-Glucose"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/BMRB/glucose-600-bmse000855_1"

save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/NRC"
project_name = "NRC-4_amino_acid-Jan2022-1"
experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/NRC/NRC_4_amino_acid_mixture_Jan_2022/1"

# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "bmse000297_ethanol"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/misc/bmse000297_ethanol"

# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/BMRB-500-100mM"
# project_name = "D-(+)-Glucose"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/BMRB/similar_settings/BMRB-500-100mM/D-(+)-Glucose"
#
# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "bmse000038_glutamine"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/misc/bmse000038_glutamine"
#
# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "gissmo_900_leucine"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/misc/gissmo_900_leucine"
#
# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "gissmo_900_leucine_1"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/misc/gissmo_900_leucine/1"
#
# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "bmse000867_serine"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/misc/bmse000867_serine"
#
# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/BMRB-500-100mM"
# project_name = "L-Isoleucine"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/BMRB/similar_settings/BMRB-500-100mM/L-Isoleucine"
#
# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "bmse000860_valine"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/misc/bmse000860_valine"
#
# save_dir = "/home/roy/MEGAsync/outputs/NMR/experiments/misc"
# project_name = "bmse000915_methionine"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/experiments_1D1H/misc/bmse000915_methionine"

# we're on tyrosine.

### end inputs.

## load.
isdir(save_dir) || mkdir(save_dir); # make save folder if it doesn't exist.

s_t, S, hz2ppmfunc, ppm2hzfunc, ν_0ppm, fs, SW, α_0ppm, β_0ppm, λ_0ppm, Ω_0ppm,
    results_0ppm, dic, α_solvent, β_solvent, λ_solvent, Ω_solvent,
    results_solvent = NMRDataSetup.loadspectrum(experiment_full_path;
    solvent_ppm = solvent_ppm_guess,
    solvent_window_ppm = solvent_window_ppm)

## store.
save_folder_path = joinpath(save_dir, project_name)
isdir(save_folder_path) || mkpath(save_folder_path); # make save folder if it doesn't exist.

save_path = joinpath(save_folder_path, "experiment.bson")
BSON.bson(save_path,
s_t = s_t,
fs = fs,
SW = SW,
ν_0ppm = ν_0ppm,
α_0ppm = α_0ppm,
β_0ppm = β_0ppm,
λ_0ppm = λ_0ppm,
dic = dic)

## visualize.
hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs
ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

offset_Hz = ν_0ppm - (ppm2hzfunc(0.3)-ppm2hzfunc(0.0))

N = length(s_t)
DFT_s = fft(s_t)
U_DFT, U_y, U_inds = NMRDataSetup.getwraparoundDFTfreqs(N, fs, offset_Hz)

# shift the DFT so that around 0 ppm (around `offset_Hz` Hz`) appears first in the frequency positions.
S_U = DFT_s[U_inds]
P_y = hz2ppmfunc.(U_y)

q = uu->NMRDataSetup.evalcomplexLorentzian(uu, α_0ppm, β_0ppm, λ_0ppm, 2*π*ν_0ppm)*fs
q_U = q.(U_y)

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P_y, real.(S_U), label = "data")
PyPlot.plot(P_y, real.(q_U), "--", label = "estimated 0ppm resonance component")
PyPlot.gca().invert_xaxis()

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("real part of data spectrum (shifted Discrete Fourier Transform)")

# It is difficult to fit resonance components perfectly in the frequency domain.
inds = findall(xx->(-0.1<xx<0.1), P_y)
P_y_near0 = P_y[inds]
S_U_near0 = S_U[inds]
q_U_near0 = q_U[inds]

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P_y_near0, real.(S_U_near0), label = "data: interpolated")
PyPlot.plot(P_y_near0, real.(q_U_near0), "--", label = "estimated 0ppm resonance component")
PyPlot.plot(P_y_near0, real.(S_U_near0), "x", label = "data: DFT points")

PyPlot.gca().invert_xaxis()

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("real part of data spectrum (shifted Discrete Fourier Transform)")

###normalized.

# # by model: normalize by the 0 ppm resonance component's estimate. sensitive to frequency mis-alignment.
# # Assumes the component has an α of 9, which corresponds to the 9 protons for DSS' 0ppm component.
# c = NMRDataSetup.evalcomplexLorentzian(ν_0ppm, 9.0, β_0ppm, λ_0ppm, 2*π*ν_0ppm)
# Z = q(ν_0ppm)/c

# by data.
# find closest U_y to ν_0ppm, then normalized according to S_U instead of q(blah freq)
val, ind = findmin( abs.(U_y .- ν_0ppm) )
Z = abs(S_U[ind])

y = S_U ./ Z

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P_y, abs.(y))

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("abs")
PyPlot.title("magnitude of normalized data spectrum s.t. 0ppm peak has magnitude 1")


text:
__precompile__()
module Utils
# with β version, please import SeisDownload.jl from the src directory as follows

include("downloadfunc.jl")
using .DownloadFunc

using SeisIO, Printf, Dates, JLD2, FileIO, Distributed

export get_starttimelist, get_timestamplist, get_stationlist, testdownload, convert_tmpfile, printparams, initlogo

"""

get_starttimelist(st::DateTime, et::DateTime, unittime::Float64)
calculate start time list for parallel downloading

    st: start time
    et: end time
    unittime: unit time in Second

    this function returns
    stlist: list of start time

    e.g.
    st = DateTime(2019,1,1,0,0,0)
    et = DateTime(2019,1,1,12,0,0)
    unittime = 3600

    stlist = get_starttimelist(st, et, unittime)

"""
function get_starttimelist(st::DateTime, et::DateTime, unittime::Real)

    reftime = st
    stlist = []

    while reftime < et
        push!(stlist, string(reftime))
        reftime += Dates.Second(float(unittime))
    end

    return stlist
end

"""
    get_timestamplist(stlist::Array{DateTime, 1})

    returns list of timestamp: Format = "Year_Julianday_Starttime".

"""
function get_timestamplist(stlist::Array{Any,1})

    timestamplist = []
    for stid = 1:length(stlist)
        yj = parse(Int64,stlist[stid][1:4])
        dj = md2j(yj, parse(Int64,stlist[stid][6:7]), parse(Int64,stlist[stid][9:10]))
        groupname    = string(yj)*"."*string(dj)*"."*stlist[stid][11:19] #Year_Julianday_Starttime
        push!(timestamplist, groupname)
    end

    return timestamplist

end

"""
    get_stationlist(network::Array{String, 1}, station::Array{String, 1}, location::Array{String, 1}, channel::Array{String, 1})

    returns list of request strings:

"""
function get_stationlist(network::Array{String, 1}, station::Array{String, 1}, location::Array{String, 1}, channel::Array{String, 1})

    stationlist = []
    for networkid = 1:length(network)
        for stationid = 1:length(station)
            for locationid = 1:length(location)
                for channelid = 1:length(channel)
                    requeststr = @sprintf("%s.%s.%s.%s", network[networkid], station[stationid], location[locationid], channel[channelid])
                    push!(stationlist, requeststr)
                end
            end
        end
    end

    return stationlist

end

"""
    testdownload(InputDict::Dict{String,Any} numofitr::Int64)

    print stats of download and return max_num_of_processes_per_parallelcycle

# Output
 -`max_num_of_processes_per_parallelcycle`: maximum number of processes for one request

"""
function testdownload(InputDict::Dict{String,Any}, numofitr::Int64)

    DownloadType    = InputDict["DownloadType"]

    trial_id          = 1
    InputDict_test = deepcopy(InputDict)
    test_suceededflag = false
    println("-------TEST DOWNLOAD START-----------")

    if DownloadType == "Noise" || DownloadType == "noise"

        while !test_suceededflag && trial_id < length(InputDict["starttimelist"])
            # select test request
            for j = 1:length(InputDict["stationinfo"]["stationlist"])
                InputDict_test["stationinfo"]["stationlist"] = [InputDict["stationinfo"]["stationlist"][j]]

                global t1 = @elapsed global dlerror = seisdownload_NOISE(trial_id, InputDict_test, testdownload=true) #[s]

                dl = [dlerror[i] for i in 1:length(dlerror)]
                if issubset(0, dl)
                    test_suceededflag = true
                    break;
                end
            end
            trial_id += 1
        end

    elseif  DownloadType == "Earthquake" || DownloadType == "earthquake"

        while !test_suceededflag && trial_id < length(InputDict["starttimelist"])
            # select test request
            for j = 1:length(InputDict["stationinfo"]["stationlist"])
                InputDict_test["stationinfo"]["stationlist"] = InputDict["stationinfo"]["stationlist"][j]

                global t1 = @elapsed global dlerror = seisdownload_EARTHQUAKE(trial_id, InputDict_test) #[s]

                dl = [dlerror[i] for i in 1:length(dlerror)]
                if issubset(0, dl)
                    test_suceededflag = true
                    break;
                end
            end
            trial_id += 1
        end
    end

    if !test_suceededflag
        error("All requests you submitted with input dictionary was failed. Please check the station availability in your request.")
    end

    estimated_downloadtime = now() + Second(round(1.0 * t1 * length(InputDict["DLtimestamplist"]) * numofitr / nprocs()))

    #println(mem_per_requestid)
    #println(max_num_of_processes_per_parallelcycle)
    println("-------DOWNLOAD STATS SUMMARY--------")

    println(@sprintf("Number of processes is %d.", nprocs()))

    # evaluate total download size by searching tmp directory
    tmppath = InputDict["tmppath"]
    s = read(`du -s -k $tmppath`, String)
    hdduse = parse(Int, split(s)[1])

    totaldownloadsize = hdduse * numofitr * length(InputDict["stationinfo"]["stationlist"])
    if totaldownloadsize < 1024 * 1024 # less than 1 GB
        totaldownloadsize = totaldownloadsize / 1024 #[MB]
        sizeunit = "MB"
    else
        totaldownloadsize = totaldownloadsize / 1024 / 1024
        sizeunit = "GB"
    end

    println(@sprintf("Total download size will be %4.2f [%s].", 0.8 * totaldownloadsize, sizeunit)) #0.8: considering compression efficiency
    #NOTE: due to the first compiling time, it is hard to estimate accurate computational time.
    #println(@sprintf("Download will finish at %s.", round(estimated_downloadtime, Dates.Second(1))))
    #println("*We have a time lag with downloading time above, like in 10 minutes or so.*")
    println("*This estimation also changes if some download requests fail and are skipped.*")
    println("-------START DOWNLOADING-------------")

    return nothing

end

"""
convert_tmpfile(InputDict::Dict)

convert temporal file in "./seisdownload_tmp" to prescribed format.
It has salvage mode, which allows to compile the temporal files in the case of failing during the download.
"""
function convert_tmpfile(InputDict::Dict; salvage::Bool=false)

    println("-------START CONVERTING-------------")
    paths = ls(InputDict["tmppath"])
    fopath = InputDict["fopath"]
    fmt = InputDict["outputformat"]

    #save HEADER information
    if fmt == "JLD2"
        file = jldopen(fopath, "w")
        file["info/DLtimestamplist"] = InputDict["DLtimestamplist"];
        file["info/starttime"]       = InputDict["starttime"]
        file["info/endtime"]         = InputDict["endtime"]
        file["info/DL_time_unit"]    = InputDict["DL_time_unit"]

    end

    stationlist     = []
    DLtimestamplist = []
    varnamelist     = []

    @simd for path in paths
        #println(path)
        try
            S = rseis(path)[1]
            #println(S)

            for ii = 1:S.n #loop at each seis channel

                # make station list
                staid = S[ii].id
                if isempty(filter(x -> x==staid, stationlist))
                    push!(stationlist, staid)
                end

                # save data (checking whether it's already in the jld2 because it causes an error)
                #parse info
                s_str = string(u2d(S[ii].t[1,2]*1e-6))

                # if time id is not matched with DLtimestamplist, this download is discarded for consistency even if it has some data.
                # here we also indirectly check if download_margin is correctly manipulated.
                if !isempty(filter(x -> x==s_str[1:19], InputDict["starttimelist"]))

                    # select output format
                    if fmt == "JLD2"
                        yj = parse(Int64, s_str[1:4])
                        mj = parse(Int64, s_str[6:7])
                        dj = parse(Int64, s_str[9:10])
                        tj = string(s_str)[11:19]

                        djm2j = md2j(yj, mj, dj)
                        groupname = string(yj)*"."*string(djm2j)*"."*tj #Year_Julianday_Starttime
                        varname = joinpath(groupname, staid)

                        if isempty(filter(x -> x==varname, varnamelist))
                            push!(varnamelist, varname)
                            file[varname] = S[ii]
                        end
                        # @info "save data $varname"

                    elseif fmt == "ASDF"
                        write_hdf5(fopath, S[ii])

                    else
                        error("output format in $fmt is not implemented yet.")
                    end

                end
            end

            if !InputDict["Istmpfilepreserved"]
    			rm(path)
    		end

        catch y
            #println(y)
        end
    end

    if fmt == "JLD2"
        # save final station list
        file["info/stationlist"] = stationlist
        JLD2.close(file)
    end

    return nothing
end

"""
printparams(param::Dict)

print parameters
"""
function printparams(param::Dict)
    printstyled("-----------Input Parameters-----------\n"; color=:cyan, bold=true)
    for key in keys(param)
        if length(string(param["$key"])) > 60
            param_str = string(param["$key"])[1:30]*"..."
        else
            param_str = string(param["$key"])
        end
        println(@sprintf("%-24s = %-10s", key, param_str))
    end
end

"""
initlogo()

print initial logo
"""
function initlogo()

    print("

      _____        _       _____                          _                    _
     / ____|      (_)     |  __ \\                        | |                  | |
    | (___    ___  _  ___ | |  | |  ___ __      __ _ __  | |  ___    __ _   __| |
     \\___ \\  / _ \\| |/ __|| |  | | / _ \\\\ \\ /\\ / /| '_ \\ | | / _ \\  / _` | / _` |
     ____) ||  __/| |\\__ \\| |__| || (_) |\\ V  V / | | | || || (_) || (_| || (_| |
    |_____/  \\___||_||___/|_____/  \\___/  \\_/\\_/  |_| |_||_| \\___/  \\__,_| \\__,_|
                      _         _  _
                     | |       | |(_)           |
    __      __       | | _   _ | | _   __ _     | v1.2 (Last update 08/15/2019)
    \\ \\ /\\ / /   _   | || | | || || | / _` |    | © Kurama Okubo
     \\ V  V /_  | |__| || |_| || || || (_| |    |
      \\_/\\_/(_)  \\____/  \\__,_||_||_| \\__,_|    |

")

    println("Job start running at "*string(now())*"\n")

end

end

ext: jl


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "FLANN"
version = v"1.9.1"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://github.com/mariusmuja/flann/archive/$version.tar.gz", "b23b5f4e71139faa3bcb39e6bbcc76967fbaf308c4ee9d4f5bfbeceaa76cc5d3"),
    DirectorySource("./bundled")
]

# Bash recipe for building across all platforms
script = raw"""

cd $WORKSPACE/srcdir/flann-*

#CMake doesn't work straight from clone, see https://github.com/mariusmuja/flann/issues/369 for source of workaround

atomic_patch -p1 "$WORKSPACE/srcdir/patches/cmake_empty_target.patch"

touch src/cpp/empty.cpp

cmake -DCMAKE_INSTALL_PREFIX=$prefix \
      -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
      -DCMAKE_BUILD_TYPE=Release \
      -DBUILD_EXAMPLES=OFF \
      -DBUILD_TESTS=OFF \
      -DBUILD_DOC=OFF \
      -DBUILD_PYTHON_BINDINGS=OFF \
      -DBUILD_MATLAB_BINDINGS=OFF

make -j${nproc}
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = expand_cxxstring_abis(supported_platforms())

# The products that we will ensure are always built
products = [
    LibraryProduct("libflann_cpp", :libflann_cpp),
    LibraryProduct("libflann", :libflann)
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
    Dependency(PackageSpec(name="CompilerSupportLibraries_jll", uuid="e66e0078-7015-5450-92f7-15fbd957f2ae")),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
function plot_mean(fig, gp, ax, x, y, smoothing,ticks)
    xall = fill(0, length(unique(x)))
    yall = zeros(size(xall))
    xpos = fill(0, length(unique(x)))
    ypos = zeros(length(xpos))
    last = first(x)
    idx = 1
    for (new, b) in zip(x, y)
        idx += new !== last
        last = new
        if  b > 0
            ypos[idx] += b
            xpos[idx] += 1
        end
        yall[idx] += b
        xall[idx] += 1
    end
    yall ./= xall
    ypos ./= [n == 0 ? one(n) : n for n in xpos]
    if smoothing
        yall = accumulate!((o,n) -> o*0.8+n*0.2, yall,yall)
        ypos = accumulate!((o,n) -> o*0.8+n*0.2, ypos,ypos)
    end
    l = lines!(ax,ticks, yall, color = :red)
    l2 = lines!(ax,ticks, ypos, color = :green)
    gp[2,1] = Legend(fig, [l,l2], ["mean", "⊕ mean"], tellwidth = false, tellhight = false, valign = :top , halign = :right)
    nothing
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule GTK4_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("GTK4")
JLLWrappers.@generate_main_file("GTK4", UUID("6ebb71f1-8434-552f-b6b1-dc18babcca63"))
end  # module GTK4_jll


text:
SUITE["intersections"] = BenchmarkGroup(["microbenchmark", "intersections"]) # add tags for eventual filtering

add_benchmark!(SUITE["intersections"], [
  quote
    s1 = Segment((0.0,0.0), (1.0,0.0))
    s2 = Segment((0.5,0.0), (2.0,0.0))
    s1 ∩ s2
  end
])


text:
tests = ["tests"]

for t in tests
    fp = string(t, ".jl")
    println("* running $fp ...")
    include(fp)
end


text:
function readdata(path::AbstractString)
	a = readlines(path)
	if a[1][1] == '>'
		fasta = []
		dna = []
		for i in a
			if i[1] == '>'
				push!(fasta,i)
				push!(dna,"")
			else
				dna[end] *= i
			end
		end
		return fasta,dna
	else
		entity = ""
		for i in a
			entity *= i
		end
		return entity

	end
end

max_stars_repo_name: Ved-Mahajan/JustForFun

max_issues_repo_name: Ved-Mahajan/JustForFun

max_forks_repo_name: Ved-Mahajan/JustForFun


text:
function QUBO_base(n, weight = 1)
    N = n^2
    M = zeros(N, N)
    for k in 1:n, j in ((k-1)*n+1):(k*n), i in ((k-1)*n+1):(k*n)
        M[j,i] = weight
    end
    return M - 2 * weight * I
end


text:
using Snopt
using DelimitedFiles
using PyPlot
import ForwardDiff

# set up boundary constraint wrapper function
function boundary_wrapper(x)
    # include relevant globals
    # global boundary_center
    # global boundary_radius
    global boundary_vertices
    global boundary_normals

    # get number of turbines
    nturbines = Int(length(x)/2)

    # extract x and y locations of turbines from design variables vector
    turbine_x = x[1:nturbines]
    turbine_y = x[nturbines+1:end]

    # get and return boundary distances
    return ff.ray_trace_boundary(boundary_vertices,boundary_normals,turbine_x,turbine_y)
end

# set up spacing constraint wrapper function
function spacing_wrapper(x)
    # include relevant globals
    global rotor_diameter

    # get number of turbines
    nturbines = Int(length(x)/2)

    # extract x and y locations of turbines from design variables vector
    turbine_x = x[1:nturbines]
    turbine_y = x[nturbines+1:end]

    # get and return spacing distances
    return 2.0*rotor_diameter[1] .- ff.turbine_spacing(turbine_x,turbine_y)
end

# set up objective wrapper function
function aep_wrapper(x)
    # include relevant globals
    global turbine_z
    global rotor_diameter
    global hub_height
    global turbine_yaw
    global ct_model
    global generator_efficiency
    global cut_in_speed
    global cut_out_speed
    global rated_speed
    global rated_power
    global windresource
    global power_models
    global model_set
    global rotor_points_y
    global rotor_points_z
    global obj_scale

    # get number of turbines
    nturbines = Int(length(x)/2)

    # extract x and y locations of turbines from design variables vector
    turbine_x = x[1:nturbines]
    turbine_y = x[nturbines+1:end]

    # calculate AEP
    AEP = obj_scale*ff.calculate_aep(turbine_x, turbine_y, turbine_z, rotor_diameter,
                hub_height, turbine_yaw, ct_model, generator_efficiency, cut_in_speed,
                cut_out_speed, rated_speed, rated_power, windresource, power_models, model_set,
                rotor_sample_points_y=rotor_points_y,rotor_sample_points_z=rotor_points_z)

    # return the objective as an array
    return [AEP]
end

# set up optimization problem wrapper function
function wind_farm_opt(x)

    # calculate spacing constraint value and jacobian
    spacing_con = spacing_wrapper(x)
    ds_dx = ForwardDiff.jacobian(spacing_wrapper, x)

    # calculate boundary constraint and jacobian
    boundary_con = boundary_wrapper(x)
    db_dx = ForwardDiff.jacobian(boundary_wrapper, x)

    # combine constaint values and jacobians into overall constaint value and jacobian arrays
    c = [spacing_con; boundary_con]
    dcdx = [ds_dx; db_dx]

    # calculate the objective function and jacobian (negative sign in order to maximize AEP)
    AEP = -aep_wrapper(x)[1]
    dAEP_dx = -ForwardDiff.jacobian(aep_wrapper,x)

    # set fail flag to false
    fail = false

    # return objective, constraint, and jacobian values
    return AEP, c, dAEP_dx, dcdx, fail
end

# import model set with wind farm and related details
include("./model_sets/model_set_6.jl")

# scale objective to be between 0 and 1
obj_scale = 1E-11

# set wind farm boundary parameters
# boundary_center = [0.0,0.0]
# boundary_radius = 300.0
boundary_vertices = ([0 0; 1 0; 1 .75; .75 .75; .75 1; 0 1] .- .5).*500 # Utah-shape boundary
boundary_normals = [0 1.0; -1 0; 0 -1; -1 0; 0 -1; 1 0]

# set globals for use in wrapper functions
global model_set
global rotor_points_y
global rotor_points_z
global turbine_z
global ambient_ti
global rotor_diameter
# global boundary_center
# global boundary_radius
global boundary_vertices
global boundary_normals
global obj_scale

# initialize design variable array
x = [copy(turbine_x);copy(turbine_y)]

# report initial objective value
println("starting objective value: ", aep_wrapper(x)[1])

# add initial turbine location to plot
for i = 1:length(turbine_x)
    plt.gcf().gca().add_artist(plt.Circle((turbine_x[i],turbine_y[i]), rotor_diameter[1]/2.0, fill=false,color="C0"))
end

# set general lower and upper bounds for design variables\
lb = zeros(length(x)) .+ minimum(boundary_vertices) # create a square that completely encloses the wind farm boundary
ub = zeros(length(x)) .+ maximum(boundary_vertices)

# set up options for SNOPT
options = Dict{String, Any}()
options["Derivative option"] = 1
options["Verify level"] = 1
options["Major optimality tolerance"] = 1e-12
options["Major iteration limit"] = 1e8
options["Summary file"] = "snopt_summary.out"
options["Print file"] = "snopt_print.out"

# run and time optimization
t1 = time()
xopt, fopt, info = snopt(wind_farm_opt, x, lb, ub, options)
t2 = time()
clkt = t2-t2

# print optimization results
println("Finished in : ", clkt, " (s)")
println("info: ", info)
println("end objective value: ", aep_wrapper(xopt))

# extract final turbine locations
turbine_x = copy(xopt[1:nturbines])
turbine_y = copy(xopt[nturbines+1:end])

# add final turbine locations to plot
for i = 1:length(turbine_x)
    plt.gcf().gca().add_artist(plt.Circle((turbine_x[i],turbine_y[i]), rotor_diameter[1]/2.0, fill=false,color="C1", linestyle="--"))
end

# add wind farm boundary to plot
plt.gcf().gca().plot([boundary_vertices[:,1];boundary_vertices[1,1]],[boundary_vertices[:,2];boundary_vertices[1,2]], color="C2")

# set up and show plot
axis("square")
xlim(minimum(boundary_vertices) - (maximum(boundary_vertices)-minimum(boundary_vertices))/5, maximum(boundary_vertices) + (maximum(boundary_vertices)-minimum(boundary_vertices))/5)
ylim(minimum(boundary_vertices) - (maximum(boundary_vertices)-minimum(boundary_vertices))/5, maximum(boundary_vertices) + (maximum(boundary_vertices)-minimum(boundary_vertices))/5)
plt.show()


text:
module DiffOpt

using Random
using LinearAlgebra
using SparseArrays
using IterativeSolvers: lsqr

import BlockDiagonals

using MathOptInterface
const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities

using MathOptSetDistances
const MOSD = MathOptSetDistances

using MatrixOptInterface
const MatOI = MatrixOptInterface

include("gen_random_problem.jl")
include("utils.jl")
include("moi_wrapper_methods.jl")
include("MOI_wrapper.jl")

export diff_optimizer, Optimizer, backward, backward_quad, backward_conic
export is_equality  # just for reference sake
export generate_lp, generate_qp

end # module


text:
# Sample a gas fitness DGP and estimate the parameters for different lenghts of
# time series
#using DynNets
using JLD2
using Hwloc
counts = Hwloc.histmap(Hwloc.topology_load())
ncores = counts[:Core]
if nworkers()<ncores
    addprocs(ncores  - nworkers())
end
@everywhere using DynNets
#
N_est = nworkers()
T=200
Nvals = [100  ]#[15,25,35,40,55]
Nmax= maximum(Nvals)
N_N = length(Nvals)
est_times = SharedArray{Float64}(zeros(N_N,N_est))
est_w = SharedArray{Float64}(zeros(N_N,N_est,Nmax))
est_A = SharedArray{Float64}(zeros(N_N,N_est,Nmax))
est_B = SharedArray{Float64}(zeros(N_N,N_est,Nmax))
est_conv_flag = SharedArray{Bool,2}(N_N,N_est)
est_flag = SharedArray{Bool,2}(N_N,N_est)

real_A = SharedArray{Float64}(zeros(N_N,N_est,Nmax))
real_B = SharedArray{Float64}(zeros(N_N,N_est,Nmax))
 @progress "ind" for ind = 1:N_N
    N = Nvals[N_N+1-ind]
    NTV =N
    indTvPars = falses(N);indTvPars[1:NTV] = true

    @sync @parallel for n = 1:N_est
        sim_A = rand((0.02:0.00001:0.8),NTV)
        sim_B = rand((0.4:0.00001:0.999),NTV)
        real_A[ind,n,1:NTV] = sim_A
        real_B[ind,n,1:NTV] = sim_B
        Y_T,fit_T =  DynNets.sampleParDynNetsBin1(T,"Fit-Multi-gas",
                                                    indTvPar = [trues(NTV);falses(N-NTV)],
                                                    α = 5,β=1,
                                                    Agas =sim_A , Bgas = sim_B  )

        degsT = squeeze(sum(Y_T,2),2)

        model_data = DynNets.SdErgm1(degsT,[zeros(i)for i in [N,NTV,NTV]], indTvPars)
        #estimate
        est_times[ind,n] = @elapsed hat_pars,hat_inds,conv_flag = DynNets.estimate(model_data)
        est_w[ind,n,1:N] = hat_pars[1]
        est_A[ind,n,1:NTV] = hat_pars[2]
        est_B[ind,n,1:NTV] = hat_pars[3]
        est_conv_flag[ind,n] = conv_flag
        est_flag[ind,n] = true
        #println([T,n])
    end

    # save
    save_path = "./data/estimatesTest/EstVarN_allTV_T200.jld"
    @save(save_path)
    println([N,sum(est_times[ind,:])])

end

#


text:
@testset "Array constructors" begin
    @testset "undef" begin
    # We can't use test_rrule here (as it's currently implemented) because the elements of
    # the array have arbitrary values. The only thing we can do is ensure that we're getting
    # `ZeroTangent`s back, and that the forwards pass produces the correct thing still.
    # Issue: https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/202
        val, pullback = rrule(Array{Float64}, undef, 5)
        @test size(val) == (5, )
        @test val isa Array{Float64, 1}
        @test pullback(randn(5)) == (NoTangent(), NoTangent(), NoTangent())
    end
    @testset "from existing array" begin
        # fwd
        test_frule(Array, randn(2, 5))
        test_frule(Array, Diagonal(randn(5)))
        # rev
        test_rrule(Array, randn(2, 5))
        test_rrule(Array, Diagonal(randn(5)))
        test_rrule(Matrix, Diagonal(randn(5)))
        test_rrule(Matrix, transpose(randn(4)))
        test_rrule(Array{ComplexF64}, randn(3))
    end
end
@testset "AbstractArray constructors" begin
    # These are what float(x) calls, but it's trivial with floating point numbers:
    test_frule(AbstractArray{Float32}, rand(3); atol=0.01)
    test_frule(AbstractArray{Float32}, Diagonal(rand(4)); atol=0.01)
    # rev
    test_rrule(AbstractArray{Float32}, rand(3); atol=0.01)
    test_rrule(AbstractArray{Float32}, Diagonal(rand(4)); atol=0.01)
    # Check with integers:
    rrule(AbstractArray{Float64}, [1, 2, 3])[2]([1, 10, 100]) == (NoTangent(), [1.0, 10.0, 100.0])
end

@testset "vect" begin
    test_rrule(Base.vect)
    @testset "homogeneous type" begin
        test_rrule(Base.vect, (5.0, ), (4.0, ))
        test_rrule(Base.vect, 5.0, 4.0, 3.0)
        test_rrule(Base.vect, randn(2, 2), randn(3, 3))
    end
    @testset "inhomogeneous type" begin
        # fwd
        test_frule(Base.vect, 5.0, 3f0)
        # rev
        test_rrule(
            Base.vect, 5.0, 3f0;
            atol=1e-6, rtol=1e-6,
        ) # tolerance due to Float32.
        test_rrule(Base.vect, 5.0, randn(3, 3); check_inferred=false)
        test_rrule(Base.vect, (5.0, 4.0), (y=randn(3),); check_inferred=false)
    end
    @testset "_instantiate_zeros" begin
        # This is an internal function also used for `cat` etc.
        @eval using ChainRules: _instantiate_zeros
        # Check these hit the fast path, unrealistic input so that map would fail:
        @test _instantiate_zeros((true, 2 , 3.0), ()) == (1, 2, 3)
        @test _instantiate_zeros((1:2, [3, 4]), ()) == (1:2, 3:4)
    end
end

@testset "copyto!" begin
    test_frule(copyto!, rand(5), rand(5))
    test_frule(copyto!, rand(10), 3, rand(5))
    test_frule(copyto!, rand(10), 2, rand(5), 2)
    test_frule(copyto!, rand(10), 2, rand(5), 2, 4)
end

@testset "reshape" begin
    # Forward
    test_frule(reshape, rand(4, 3), 2, :)
    test_frule(reshape, rand(4, 3), axes(rand(6, 2)))
    @test_skip test_frule(reshape, Diagonal(rand(4)), 2, :) # https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/239

    # Reverse
    test_rrule(reshape, rand(4, 5), (2, 10))
    test_rrule(reshape, rand(4, 5), 2, 10)
    test_rrule(reshape, rand(4, 5), 2, :)
    test_rrule(reshape, rand(4, 5), axes(rand(10, 2)))
    # structured
    test_rrule(reshape, transpose(rand(4)), :)
    test_rrule(reshape, adjoint(rand(ComplexF64, 4)), :)
    @test rrule(reshape, adjoint(rand(ComplexF64, 4)), :)[2](rand(4))[2] isa Adjoint{ComplexF64}
    @test rrule(reshape, Diagonal(rand(4)), (2, :))[2](ones(2,8))[2] isa Diagonal
    @test_skip test_rrule(reshape, Diagonal(rand(4)), 2, :)  # DimensionMismatch("second dimension of A, 22, does not match length of x, 16")
    @test_skip test_rrule(reshape, UpperTriangular(rand(4,4)), (8, 2)) # https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/239
end

@testset "dropdims" begin
    # fwd
    test_frule(dropdims, rand(4, 1); fkwargs=(; dims=2))
    # rev
    test_rrule(dropdims, rand(4, 1); fkwargs=(; dims=2))
    test_rrule(dropdims, transpose(rand(4)); fkwargs=(; dims=1))
    test_rrule(dropdims, adjoint(rand(ComplexF64, 4)); fkwargs=(; dims=1))
    @test rrule(dropdims, adjoint(rand(ComplexF64, 4)); dims=1)[2](rand(4))[2] isa Adjoint{ComplexF64}
end

@testset "permutedims + PermutedDimsArray" begin
    # Forward
    test_frule(permutedims, rand(5))
    test_frule(permutedims, rand(3, 4), (2, 1))
    test_frule(permutedims!, rand(4,3), rand(3, 4), (2, 1))
    test_frule(PermutedDimsArray, rand(3, 4, 5), (3, 1, 2))

    # Reverse
    test_rrule(permutedims, rand(5))
    test_rrule(permutedims, rand(3, 4), (2, 1))
    test_rrule(permutedims, Diagonal(rand(5)), (2, 1))
    # Note BTW that permutedims(Diagonal(rand(5))) does not use the rule at all

    @test invperm((3, 1, 2)) != (3, 1, 2)
    test_rrule(permutedims, rand(3, 4, 5), (3, 1, 2))

    @test_skip test_rrule(PermutedDimsArray, rand(3, 4, 5), (3, 1, 2))  # https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/240
    x = rand(2, 3, 4)
    dy = rand(4, 2, 3)
    @test rrule(permutedims, x, (3, 1, 2))[2](dy)[2] == rrule(PermutedDimsArray, x, (3, 1, 2))[2](dy)[2]
end

@testset "repeat" begin
    # forward
    test_frule(repeat, rand(4), 2)
    test_frule(repeat, rand(2, 3); fkwargs = (inner=(1,2), outer=(1,3)))

    # reverse
    test_rrule(repeat, rand(4, ))
    test_rrule(repeat, rand(4, 5))
    test_rrule(repeat, rand(4, 5); fkwargs = (outer=(1,2),))
    test_rrule(repeat, rand(4, 5); fkwargs = (inner=(1,2), outer=(1,3)))

    test_rrule(repeat, rand(4, ), 2)
    test_rrule(repeat, rand(4, 5), 2)
    test_rrule(repeat, rand(4, 5), 2, 3)
    test_rrule(repeat, rand(1,2,3), 2,3,4; check_inferred=VERSION>v"1.6")
    test_rrule(repeat, rand(0,2,3), 2,0,4; check_inferred=VERSION>v"1.6")
    test_rrule(repeat, rand(1,1,1,1), 2,3,4,5; check_inferred=VERSION>v"1.6")

    # These need Julia 1.6
    test_rrule(repeat, rand(4, 5); fkwargs = (inner=(2,4), outer=(1,1,1,3)))
    test_rrule(repeat, rand(1,2,3), 2,3)
    test_rrule(repeat, rand(1,2,3), 2,3,4,2)
    test_rrule(repeat, fill(1.0), 2)
    test_rrule(repeat, fill(1.0), 2, 3)

    # These fail for other v1.0 related issues (add!!)
    # v"1.0": fill(1.0) + fill(1.0) != fill(2.0)
    # v"1.6: fill(1.0) + fill(1.0) == fill(2.0) # Expected
    test_rrule(repeat, fill(1.0); fkwargs = (inner=2,))
    test_rrule(repeat, fill(1.0); fkwargs = (inner=2, outer=3,))

    @test rrule(repeat, [1,2,3], 4)[2](ones(12))[2] == [4,4,4]
    @test rrule(repeat, [1,2,3], outer=4)[2](ones(12))[2] == [4,4,4]

end

@testset "hcat" begin
    # forward
    test_frule(hcat, randn(3, 2), randn(3))
    test_frule(hcat, randn(), randn(1,3))

    # reverse
    test_rrule(hcat, randn(3, 2), randn(3), randn(3, 3))
    test_rrule(hcat, rand(), rand(1,2), rand(1,2,1))
    test_rrule(hcat, rand(3,1,1,2), rand(3,3,1,2))

    # mix types
    test_rrule(hcat, rand(2, 2), rand(2, 2)')
end

@testset "reduce hcat" begin
    mats = [randn(3, 2), randn(3, 1), randn(3, 3)]
    test_frule(reduce, hcat, mats)
    test_rrule(reduce, hcat, mats)

    vecs = [rand(3) for _ in 1:4]
    test_frule(reduce, hcat, vecs)
    test_rrule(reduce, hcat, vecs)

    mix = AbstractVecOrMat[rand(4,2), rand(4)]  # this is weird, but does hit the fast path
    test_rrule(reduce, hcat, mix)

    adjs = vec([randn(2, 4), randn(1, 4), randn(3, 4)]')  # not a Vector
    # test_rrule(reduce, hcat, adjs ⊢ map(m -> rand(size(m)), adjs))
    dy = 1 ./ reduce(hcat, adjs)
    @test rrule(reduce, hcat, adjs)[2](dy)[3] ≈ rrule(reduce, hcat, collect.(adjs))[2](dy)[3]

    # mix types
    mats = [randn(2, 2), rand(2, 2)']
    test_rrule(reduce, hcat, mats)
end

@testset "vcat" begin

    # forward
    test_frule(vcat, randn(), randn(3), rand())
    test_frule(vcat, randn(3, 1), randn(3))

    # reverse
    test_rrule(vcat, randn(2, 4), randn(1, 4), randn(3, 4))
    test_rrule(vcat, rand(), rand())
    test_rrule(vcat, rand(), rand(3), rand(3,1,1))
    test_rrule(vcat, rand(3,1,2), rand(4,1,2))

    # mix types
    test_rrule(vcat, rand(2, 2), rand(2, 2)')
end

@testset "reduce vcat" begin
    mats = [randn(2, 4), randn(1, 4), randn(3, 4)]
    test_frule(reduce, vcat, mats)
    test_rrule(reduce, vcat, mats)

    vecs = [rand(2), rand(3), rand(4)]
    test_frule(reduce, vcat, vecs)
    test_rrule(reduce, vcat, vecs)

    mix = AbstractVecOrMat[rand(4,1), rand(4)]
    test_rrule(reduce, vcat, mix)
end

@testset "cat" begin
    # forward
    test_frule(cat, rand(2, 4), rand(1, 4); fkwargs=(dims=1,))
    test_frule(cat, rand(), rand(2,3); fkwargs=(dims=(1,2),))

    # reverse
    test_rrule(cat, rand(2, 4), rand(1, 4); fkwargs=(dims=1,))
    test_rrule(cat, rand(2, 4), rand(2); fkwargs=(dims=Val(2),))
    test_rrule(cat, rand(), rand(2, 3); fkwargs=(dims=[1,2],))
    test_rrule(cat, rand(1), rand(3, 2, 1); fkwargs=(dims=(1,2),), check_inferred=false) # infers Tuple{Zero, Vector{Float64}, Any}

    test_rrule(cat, rand(2, 2), rand(2, 2)'; fkwargs=(dims=1,))
end

@testset "hvcat" begin
    # forward
    test_frule(hvcat, 2, rand(6)...)

    # reverse
    test_rrule(hvcat, 2, rand(ComplexF64, 6)...)
    test_rrule(hvcat, (2, 1), rand(), rand(1,1), rand(2,2))
    test_rrule(hvcat, 1, rand(3)' ⊢ rand(1,3), transpose(rand(3)) ⊢ rand(1,3))
    test_rrule(hvcat, 1, rand(0,3), rand(2,3), rand(1,3,1))

    # mix types (adjoint and transpose)
    test_rrule(hvcat, 1, rand(3)', transpose(rand(3)) ⊢ rand(1,3))
end

@testset "reverse" begin
    @testset "Tuple" begin
        test_frule(reverse, Tuple(rand(10)))
        @test_skip test_rrule(reverse, Tuple(rand(10)))  # Ambiguity in isapprox, https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/229
    end
    @testset "Array" begin
        # Forward
        test_frule(reverse, rand(5))
        test_frule(reverse, rand(5), 2, 4)
        test_frule(reverse, rand(5), fkwargs=(dims=1,))
        test_frule(reverse, rand(3,4), fkwargs=(dims=2,))
        test_frule(reverse, rand(3,4))
        test_frule(reverse, rand(3,4,5), fkwargs=(dims=(1,3),))

        test_frule(reverse!, rand(5))
        test_frule(reverse!, rand(5), 2, 4)
        test_frule(reverse!, rand(3,4), fkwargs=(dims=2,))

        # Reverse
        test_rrule(reverse, rand(5))
        test_rrule(reverse, rand(5), 2, 4)
        test_rrule(reverse, rand(5), fkwargs=(dims=1,))

        test_rrule(reverse, rand(3,4), fkwargs=(dims=2,))
        test_rrule(reverse, rand(3,4))
        test_rrule(reverse, rand(3,4,5), fkwargs=(dims=(1,3),))

        # Structured
        y, pb = rrule(reverse, Diagonal([1,2,3]))
        # We only preserve structure in this case if given structured tangent (no ProjectTo)
        @test unthunk(pb(Diagonal([1.1, 2.1, 3.1]))[2]) isa Diagonal
        @test unthunk(pb(rand(3, 3))[2]) isa AbstractArray
    end
end

@testset "circshift" begin
    # Forward
    test_frule(circshift, rand(10), 1)
    test_frule(circshift, rand(10), (1,))
    test_frule(circshift, rand(3,4), (-7,2))

    test_frule(circshift!, rand(10), rand(10), 1)
    test_frule(circshift!, rand(3,4), rand(3,4), (-7,2))

    # Reverse
    test_rrule(circshift, rand(10), 1)
    test_rrule(circshift, rand(10) .+ im, -2)
    test_rrule(circshift, rand(10), (1,))
    test_rrule(circshift, rand(3,4), (-7,2))
end

@testset "fill" begin
    # Forward
    test_frule(fill, 12.3, 4)
    test_frule(fill, 5.0, (6, 7))

    test_frule(fill!, rand(2, 3), rand())

    # Reverse
    test_rrule(fill, 44.4, 4)
    test_rrule(fill, 55 + 0.5im, 5)
    test_rrule(fill, 3.3, (3, 3, 3))
end

@testset "filter" begin
    @testset "Array" begin
        # Random numbers will confuse finite differencing here, as it may perturb across the boundary.
        x5 = [0.0, 1.0, 0.2, 0.9, 0.7]
        x34 = Float64[-113  124   -37   12
                        96  -89   103  119
                        91  -21  -110   10]

        # Forward
        test_frule(filter, >(0.5) ⊢ NoTangent(), x5)
        test_frule(filter, <(0), x34)
        test_frule(filter, >(100), x5)

        # Reverse
        test_rrule(filter, >(0.5) ⊢ NoTangent(), x5)  # Without ⊢, MethodError: zero(::Base.Fix2{typeof(>), Float64}) -- https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/231
        test_rrule(filter, <(0), x34)
        test_rrule(filter, >(100), x5)  # fixed in https://github.com/JuliaDiff/ChainRulesCore.jl/pull/534
        @test unthunk(rrule(filter, >(100), x5)[2](Int[])[3]) == zero(x5)
    end
end

@testset "findmin & findmax" begin
    # Forward
    test_frule(findmin, rand(10))
    test_frule(findmax, rand(10))
    @test @inferred(frule((nothing, rand(3,4)), findmin, rand(3,4))) isa Tuple{Tuple{Float64, CartesianIndex}, Tangent}
    @test @inferred(frule((nothing, rand(3,4)), findmin, rand(3,4), dims=1)) isa Tuple{Tuple{Matrix, Matrix}, Tangent}
    @test_skip test_frule(findmin, rand(3,4)) # error from test_approx(actual::CartesianIndex{2}, expected::CartesianIndex{2}
    @test_skip test_frule(findmin, rand(3,4), output_tangent = (rand(), NoTangent()))
    @test_skip test_frule(findmin, rand(3,4), fkwargs=(dims=1,))
    # These skipped tests might be fixed by https://github.com/JuliaDiff/FiniteDifferences.jl/issues/188

    # Reverse
    test_rrule(findmin, rand(10), output_tangent = (rand(), false))
    test_rrule(findmax, rand(10), output_tangent = (rand(), false))
    test_rrule(findmin, rand(5,3))
    test_rrule(findmax, rand(5,3))
    @test [0 0; 0 5] == @inferred unthunk(rrule(findmax, [1 2; 3 4])[2]((5.0, nothing))[2])
    @test [0 0; 0 5] == @inferred unthunk(rrule(findmax, [1 2; 3 4])[2]((5.0, NoTangent()))[2])

    # Reverse with dims:
    @test [0 0; 5 6] == @inferred unthunk(rrule(findmax, [1 2; 3 4], dims=1)[2](([5 6], nothing))[2])
    @test [5 0; 6 0] == @inferred unthunk(rrule(findmin, [1 2; 3 4], dims=2)[2]((hcat([5,6]), nothing))[2])
    test_rrule(findmin, rand(3,4), fkwargs=(dims=1,), output_tangent = (rand(1,4), NoTangent()))
    test_rrule(findmin, rand(3,4), fkwargs=(dims=2,))

    # Second derivatives
    test_frule(ChainRules._zerolike_writeat, rand(2, 2), 5.0, :, CartesianIndex(2, 2))
    test_rrule(ChainRules._zerolike_writeat, rand(2, 2), 5.0, :, CartesianIndex(2, 2))
    @test_skip test_rrule(ChainRules._zerolike_writeat, rand(2, 2), 5.0, 1, [CartesianIndex(2, 1) CartesianIndex(2, 2)]  ⊢ NoTangent())  # MethodError: no method matching isapprox(::Matrix{Float64}, ::Float64; rtol=1.0e-9, atol=1.0e-9)
    y, bk = rrule(ChainRules._zerolike_writeat, rand(2, 2), 5.0, 1, [CartesianIndex(2, 1) CartesianIndex(2, 2)])
    @test y == [0 0; 5 5]
    @test bk([1 2; 3 4]) == (NoTangent(), NoTangent(), [3 4], NoTangent(), NoTangent())
end

@testset "$imum" for imum in [maximum, minimum]
    # Forward
    test_frule(imum, rand(10))
    test_frule(imum, rand(3,4))
    test_frule(imum, rand(3,4), fkwargs=(dims=1,))
    test_frule(imum, [rand(2) for _ in 1:3])
    test_frule(imum, [rand(2) for _ in 1:3, _ in 1:4]; fkwargs=(dims=1,))

    # Reverse
    test_rrule(imum, rand(10))
    test_rrule(imum, rand(3,4))
    test_rrule(imum, rand(3,4), fkwargs=(dims=1,))
    test_rrule(imum, rand(3,4,5), fkwargs=(dims=(1,3),))

    # Arrays of arrays
    test_rrule(imum, [rand(2) for _ in 1:3]; check_inferred=false)
    test_rrule(imum, [rand(2) for _ in 1:3, _ in 1:4]; fkwargs=(dims=1,), check_inferred=false)

    # Case which attains max twice -- can't use FiniteDifferences for this
    res = imum == maximum ? [0,1,0,0,0,0] : [1,0,0,0,0,0]
    @test res == @inferred unthunk(rrule(imum, [1,2,1,2,1,2])[2](1.0)[2])

    # Structured matrix -- NB the minimum is a structral zero here
    @test unthunk(rrule(imum, Diagonal(rand(3) .+ 1))[2](5.5)[2]) isa Diagonal
    @test unthunk(rrule(imum, UpperTriangular(rand(3,3) .+ 1))[2](5.5)[2]) isa UpperTriangular{Float64}
    @test_skip test_rrule(imum, Diagonal(rand(3) .+ 1)) # MethodError: no method matching zero(::Type{Any}), from fill!(A::SparseArrays.SparseMatrixCSC{Any, Int64}, x::Bool)
end

@testset "extrema" begin
    test_rrule(extrema, rand(10), output_tangent = (rand(), rand()))
    test_rrule(extrema, rand(3,4), fkwargs=(dims=1,), output_tangent = collect(zip(rand(1,4), rand(1,4))))
    # Case where both extrema are the same index, to check accumulation:
    test_rrule(extrema, rand(1), output_tangent = (rand(), rand()))
    test_rrule(extrema, rand(1,1), fkwargs=(dims=2,), output_tangent = hcat((rand(), rand())))
    test_rrule(extrema, rand(3,1), fkwargs=(dims=2,), output_tangent = collect(zip(rand(3,1), rand(3,1))))
    # Double-check the forward pass
    A = randn(3,4,5)
    @test extrema(A, dims=(1,3)) == rrule(extrema, A, dims=(1,3))[1]
    B = hcat(A[:,:,1], A[:,:,1])
    @test extrema(B, dims=2) == rrule(extrema, B, dims=2)[1]
end


text:
using HDF5
using DocOpt
using HPAT

function generate_files(path,N)
    A = rand(Float64,N)
    hdf5_file = path*"1D_large.hdf5"
    csv_file = path*"1D_large.csv"
    if isfile(hdf5_file)
        rm(hdf5_file)
    end
    if isfile(csv_file)
        rm(csv_file)
    end
    h5write(hdf5_file,"/labels", A)
    writecsv(csv_file,A)
end

function generate_files_small(path)
    N=1000
    A = rand(Float64,N)
    hdf5_file = path*"1D_small.hdf5"
    csv_file = path*"1D_small.csv"
    if isfile(hdf5_file)
        rm(hdf5_file)
    end
    if isfile(csv_file)
        rm(csv_file)
    end
    h5write(hdf5_file,"/labels", A)
    writecsv(csv_file,A)
end

function main()
    doc = """generate random 1D array data.

Usage:
  generate_1D_array.jl -h | --help
  generate_1D_array.jl [--instances=<instances>] [--path=<path>]

Options:
  -h --help                  Show this screen.
  --instances=<instances>    Specify number of instances; defaults to 2000000.
  --path=<path>              Specify output path for generated files; defaults to HPAT's default data path.
"""

    arguments = docopt(doc)

    if (arguments["--instances"] != nothing)
        instances = parse(Int, arguments["--instances"])
    else
        instances = 10^6
    end

    if (arguments["--path"] != nothing)
        path = arguments["--path"]
    else
        path = HPAT.getDefaultDataPath()
    end

    generate_files_small(path)
    generate_files(path, instances)
end

main()


text:
import Base: @propagate_inbounds,getindex, setindex!,iterate,size,length,push!,
              collect,view

export BodyList, getrange, numpts

struct BodyList
    list :: Vector{Body}
end

BodyList() = BodyList(Body[])

@propagate_inbounds getindex(A::BodyList, i::Int) = A.list[i]
@propagate_inbounds setindex!(A::BodyList, v::Body, i::Int) = A.list[i] = v
@propagate_inbounds getindex(A::BodyList, I...) = A.list[I...]
@propagate_inbounds setindex!(A::BodyList, v, I...) = A.list[I...] = v

iterate(A::BodyList) = iterate(A.list)
iterate(A::BodyList,I) = iterate(A.list,I)
size(A::BodyList) = size(A.list)
length(A::BodyList) = length(A.list)
numpts(A::BodyList) = mapreduce(b -> length(b.x),+,A)

numpts(A::Body{N}) where {N} = N

push!(bl::BodyList,b::Body) = push!(bl.list,b)

"""
    collect(bl::bodylist) -> Vector{Float64}, Vector{Float64}

Collect the inertial-space coordinates of all of the Lagrange points comprising
the bodies in body list `bl` and return each assembled set of coordinates as a vector.
"""
function collect(bl::BodyList)
    xtmp = Float64[]
    ytmp = Float64[]
    for b in bl
        append!(xtmp,b.x)
        append!(ytmp,b.y)
    end
    return xtmp,ytmp
end

collect(body::Body) = collect(BodyList([body]))

"""
    getrange(bl::BodyList,i::Int) -> Range

Return the range of indices in the global set of Lagrange point data corresponding to body `i` in body list `bl`.
"""
function getrange(bl::BodyList,i::Int)
    i <= length(bl) || error("Unavailable body")
    first = 1
    j = 1
    while j < i
        first += length(bl[j])
        j += 1
    end
    last = first+length(bl[i])-1
    return first:last
end

"""
    view(f::AbstractVector,bl::BodyList,i::Int) -> SubArray

Provide a view of the range of values in vector `f` corresponding to the Lagrange
points of the body with index `i` in body list `bl`.
"""
function Base.view(f::AbstractVector,bl::BodyList,i::Int)
    length(f) == numpts(bl) || error("Inconsistent size of data for viewing")
    return view(f,getrange(bl,i))
end

"""
    sum(f::AbstractVector,bl::BodyList,i::Int) -> Real

Compute a sum of the elements of vector `f` corresponding to body `i` in body
list `bl`.
"""
Base.sum(f::AbstractVector,bl::BodyList,i::Int) = sum(view(f,bl,i))


text:
pkg_dict = Pkg.installed()  # get dictionary of installed package names to version numbers

if !haskey(pkg_dict, "MPI")
  Pkg.clone("MPI")
  Pkg.build("MPI")
end

if !(haskey(pkg_dict, "FactCheck"))
  Pkg.clone("FactCheck")
  start_dir = pwd()
  cd(Pkg.dir("FactCheck"))
  run(`git checkout e3739d5fdf0e54bc1e74957c060c693cd8ce9cd6`)
  cd(start_dir)
  Pkg.build("FactCheck")
end

if !haskey(ENV, "PETSC_DIR")  && !haskey(ENV, "PETSC_ARCH")
 run(`./install_petsc.sh`)
end


text:
using Pkg
Pkg.activate(pwd())

# # Soft Local Scope

for i in 1:2
    t = 1 + i
    @show t
end

t

#+

for j in 1:5
    for i in 1:2
        @show i + j
    end
    i
end

#+

s = 0

for i = 1:10
    t = 1 + i # new local variable t
    s = t # assign a new value to the global variable
end

s

#+

code = """
s = 0
for i = 1:10
    t = 1 + i # new local variable t
    s = t # new local variable s and warning
end
s
""";

include_string(Main, code)

#+

code_local = """
s = 0
for i = 1:10
    t = 1 + i # new local variable t
    local s = t # assigning a new value to the global variable
end
s
"""

code_global = """
s = 0
for i = 1:10
    t = 1 + i # new local variable t
    global s = t # assigning a new value to the global variable
end
s
"""

#+

include_string(Main, code_global)
include_string(Main, code_local)


text:
export rollout

using StatsBase: wsample

function rollout(policy; epsilon = 0.1, kwargs...)
    rewards = Float32[]
    state = Vector{Float32}[]
    actions = Int[]

    env = CartPoleEnv(;kwargs...)

    action_space = get_action_space(env)

    action_spacec = collect(action_space)

    obs = observe(env)

    while true
        push!(state, obs.state)
        if rand() < epsilon
            action = rand(action_space)
        else
            pos = policy(obs.state)
            if any(isnan, pos)
                error("wtf")
            end
            action = wsample(action_spacec, pos)
        end
        push!(actions, action)
        env(action)
        obs = observe(env)
        push!(rewards, obs.reward)

        get_terminal(obs) && break
    end
    state, rewards, actions
end

max_stars_repo_name: xiaodaigh/my-rl-gym

max_issues_repo_name: xiaodaigh/my-rl-gym

max_forks_repo_name: xiaodaigh/my-rl-gym


text:
# using Pkg
# Pkg.develop(path=joinpath(@__DIR__, "../../DojoEnvironments"))
# Pkg.develop(path=joinpath(@__DIR__, "../.."))
# Pkg.activate(joinpath(@__DIR__, ".."))
# Pkg.instantiate()

# ## visualizer
vis = Visualizer()
open(vis)

# ## setup
using Dojo
using IterativeLQR
using LinearAlgebra
using FiniteDiff
using DojoEnvironments
using Plots
using StaticArrays

################################################################################
# Continuation
################################################################################
function reset!(env::Environment; rtol=1e-4, btol=1e-3, undercut=2.0)
    env.opts_step.rtol = rtol
    env.opts_step.btol = btol
    env.opts_step.undercut = 5.0
    env.opts_grad.rtol = rtol
    env.opts_grad.btol = btol
    env.opts_grad.undercut = undercut
    return nothing
end

function continuation_callback!(solver::Solver, env::Environment; ρ=1.5)
    # contact smoothness continuation
    env.opts_step.rtol = max(1e-6, env.opts_step.rtol/ρ)
    env.opts_step.btol = max(1e-4, env.opts_step.btol/ρ)
    env.opts_grad.rtol = max(1e-6, env.opts_grad.rtol/ρ)
    env.opts_grad.btol = max(1e-4, env.opts_grad.btol/ρ)

    # visualize current policy
    ū = solver.problem.actions
    x̄ = IterativeLQR.rollout(model, x1, ū)
    DojoEnvironments.visualize(env, x̄)

    println("r_tol $(scn(env.opts_grad.rtol))  " *
        "κ_tol $(scn(env.opts_grad.btol))")
    return nothing
end

################################################################################
# ## system
################################################################################
gravity = -9.81
timestep = 0.01
friction_coefficient = 0.4
damper = 1.0*1.5
spring = 0.0*5.0
env = DojoEnvironments.get_environment(:quadruped,
    representation=:minimal,
    timestep=timestep,
    contact_body=false,
    gravity=gravity,
    friction_coefficient=friction_coefficient,
    # damper=damper,
    # spring=spring,
    infeasible_control=true,
    vis=vis)

# ## dimensions
n = env.num_states
m = env.num_inputs
nu_infeasible = 6

# ## template
include(joinpath(@__DIR__, "../../DojoEnvironments/src",
    "quadruped/methods/template.jl"))

################################################################################
# ## simulation test
################################################################################
mech = get_mechanism(:quadruped,
    contact_body=false,
    timestep=timestep,
    gravity=gravity,
    friction_coefficient=friction_coefficient,
    damper=damper,
    limits=false,
    spring=spring)

initialize!(mech, :quadruped, body_position=[0,0,0.0])
u_hover = 0.0*[0.02;0;0.6; 0;0;0; zeros(12)]
function ctrl!(m, k; u=u_hover)
    nu = input_dimension(m)
    set_input!(m, SVector{nu}(u))
end

storage = simulate!(mech, 0.27, ctrl!,
    record=true,
    verbose=true,
    opts=SolverOptions(rtol=1e-5, btol=1e-4, undercut=5.0, verbose=false),
    )
Dojo.visualize(mech, storage, vis=env.vis)

# z_rest = get_maximal_state(mech)

################################################################################
# ## reference trajectory
################################################################################
N = 1
initialize!(env.mechanism, :quadruped)
xref = quadruped_trajectory(env.mechanism,
    r=0.00,
    z=0.25;
    Δx=-0.01,
    Δfront=0.05,
    width_scale=0.0,
    height_scale=1.0,
    N=1,
    Ncycles=N)
zref = [minimal_to_maximal(env.mechanism, x) for x in xref]
DojoEnvironments.visualize(env, xref)

# ## horizon
T = length(zref)

################################################################################
# ## ILQR problem
################################################################################
# ## model
dyn = IterativeLQR.Dynamics(
    (y, x, u, w) -> dynamics(y, env, x, u, w),
    (dx, x, u, w) -> dynamics_jacobian_state(dx, env, x, u, w),
    (du, x, u, w) -> dynamics_jacobian_input(du, env, x, u, w),
    n, n, m)

model = [dyn for t = 1:T-1]

# ## rollout
x1 = deepcopy(xref[1])
ū = [u_hover for t = 1:T-1]

x̄ = IterativeLQR.rollout(model, x1, ū)
DojoEnvironments.visualize(env, x̄)

# ## objective
############################################################################
qt = [0.3; 0.05; 0.05;
    5e-2 * ones(3);
    1e-0 * ones(3);
    1e-0 * ones(3);
    fill([2, 1e-0], 12)...]
ots = [(x, u, w) -> transpose(x - xref[t]) * Diagonal(timestep * qt) * (x - xref[t]) +
# transpose(u - u_hover) * Diagonal(timestep * 0.01 * ones(m)) * (u - u_hover) for t = 1:T-1]
    transpose(u) * Diagonal(timestep * 1.0 * ones(m)) * u for t = 1:T-1]
oT = (x, u, w) -> transpose(x - xref[end]) * Diagonal(timestep * qt) * (x - xref[end])

cts = [IterativeLQR.Cost(ot, n, m) for ot in ots]
cT = IterativeLQR.Cost(oT, n, 0)
obj = [cts..., cT]

# ## constraints
############################################################################
ul = -1.0 * 1e-3*ones(nu_infeasible)
uu = +1.0 * 1e-3*ones(nu_infeasible)

function contt(x, u, w)
    [
        1e-1 * (ul - u[1:nu_infeasible]);
        1e-1 * (u[1:nu_infeasible] - uu);
    ]
end

function goal(x, u, w)
    Δ = 1e-0 * (x - xref[end])
    return Δ
end

con_policyt = IterativeLQR.Constraint(contt, n, m, indices_inequality=collect(1:2nu_infeasible))
con_policyT = IterativeLQR.Constraint(goal, n, 0)

cons = [[con_policyt for t = 1:T-1]..., con_policyT]

# ## solver
options = Options(line_search=:armijo,
        max_iterations=50,
        max_dual_updates=12,
        min_step_size=1e-5,
        objective_tolerance=1e-3,
        lagrangian_gradient_tolerance=1e-3,
        constraint_tolerance=1e-4,
        initial_constraint_penalty=1e-1,
        scaling_penalty=10.0,
        max_penalty=1e4,
        verbose=true)

s = IterativeLQR.Solver(model, obj, cons, options=options)

IterativeLQR.initialize_controls!(s, ū)
IterativeLQR.initialize_states!(s, x̄)

# ## solve
local_callback!(solver::IterativeLQR.Solver) = continuation_callback!(solver, env)
reset!(env, rtol=1e-3, btol=1e-3)
@time IterativeLQR.constrained_ilqr_solve!(s, augmented_lagrangian_callback! = local_callback!)

# ## solution
x_sol, u_sol = IterativeLQR.get_trajectory(s)

# ## visualize
x_view = [[x_sol[1] for t = 1:15]..., x_sol..., [x_sol[end] for t = 1:15]...]
DojoEnvironments.visualize(env, x_view)

################################################################################
# TVLQR
################################################################################
N = 50
x_tv = [fill(x_sol[1:end-1], N)...; [x_sol[end]]]
u_tv = [fill(u_sol, N)...;]

env.opts_grad = SolverOptions(rtol=1e-4, btol=1e-3)
K_tv, P_tv = tvlqr(x_tv, u_tv, env;
        q_tracking=[10;10;10;
            1e0 * ones(3);
            1e0 * ones(3);
            1e0 * ones(3);
            fill([1e0, 1e0], 12)...],
        r_tracking=env.mechanism.timestep * 50 * ones(length(u_tv[1])))
env.opts_grad = SolverOptions(rtol=1e-5, btol=1e-5)

# K_tv, P_tv = tvlqr(x_tv, u_tv, env;
#         q_tracking=[0.3; 0.05; 0.05;
#             5e-1 * ones(3);
#             1e-0 * ones(3);
#             1e-0 * ones(3);
#             fill([2, 1e-0], 12)...],
#         r_tracking=env.mechanism.timestep * 1e1 * ones(length(u_tv[1])))

nu = input_dimension(mech)
nx = minimal_dimension(mech)
plot(hcat([reshape(K, nu*nx) for K in K_tv]...)', legend=false)

eigvals(P_tv[1])
cond(P_tv[1])
K_tv[1][18,35]
K_tv[1][18,36]

K_tv[1][17,33]
K_tv[1][17,34]

plot(hcat(x_sol...)')
plot(hcat(u_sol...)')

plot(hcat(x_tv...)')
plot(hcat(u_tv...)')

################################################################################
# Test Policy
################################################################################

# initialize!(mech, :quadruped, body_position=[0,0,0.])
# set_maximal_state!(mech, z_rest)
set_maximal_state!(mech, deepcopy(zref[1]))
function ctrl!(mechanism, k)
    nu = input_dimension(mechanism)
    x = get_minimal_state(mechanism)
    u = u_sol[1]/1.0 + 1.0*K_tv[1] * (x_sol[1] - x)
    u = [zeros(6); u[7:18]]
    set_input!(mechanism, SVector{nu}(u))
end

storage = simulate!(mech, 5.0, ctrl!,
    record=true,
    verbose=true,
    opts=SolverOptions(rtol=1e-5, btol=1e-4, undercut=5.0, verbose=false),
    )
Dojo.visualize(mech, storage, vis=env.vis, build=false)

################################################################################
# CIMPC compat
################################################################################

mutable struct TVLQRPolicy114{T}
    K::Vector{Matrix{T}}
    x::Vector{Vector{T}}
    u::Vector{Vector{T}}
    timestep::T
    H::Int
end

policy = TVLQRPolicy114(K_tv[1:T-1], x_sol[1:T-1], u_sol[1:T-1], timestep, T-1)

function exec_policy(p::TVLQRPolicy114{T}, x::Vector{T}, t::T) where {T}
    timestep = p.timestep
    i = floor(t/timestep) % H + 1
    u = p.u[i] + p.K[i] * (p.x[i] - x)
    return u / timestep # force
end

JLD2.jldsave(joinpath(@__DIR__, "tvlqr_policies", "standing_tvlqr_policy.jld2"),
    policy=policy,
    K=K_tv[1:T-1],
    x=x_sol[1:T-1],
    u=u_sol[1:T-1],
    timestep=timestep,
    H=T-1)

file = JLD2.jldopen(joinpath(@__DIR__, "tvlqr_policies", "standing_tvlqr_policy.jld2"))
policy = file["policy"]
K = file["K"]
x = file["x"]
u = file["u"]
timestep = file["timestep"]
H = file["H"]
JLD2.close(file)

# H = 4
# for t in 0:0.0035:0.3
#     @show floor(t/timestep) % H + 1
# end


text:
function compute()

	N=1000
	for a=1:N
		for b=a+1:N
			c=N-a-b
			if a^2+b^2==c^2
				return a*b*c

			end
		end
	end
end

print(compute())


text:
using DataFrames, CSV, Dates
using JLD

# import data
df_covid_raw = DataFrame(CSV.File("data/covid19-canada.csv"))

# look at BC only
df_covid_BC = select(filter(:prname => ==("British Columbia"),df_covid_raw),
                    :date,:numconf,:numactive, :numdeaths, :numrecover)

# population of BC in 2021 is about 4.9 mil
# assume s0 is 4.9 mil [estimated pop of BC in 2021]
init_pop = 4900000
compute_s0(I,R) = init_pop - (I + R)
compute_R0(D,R) = D+R
# treat deaths as recovered
transform!(df_covid_BC,[:numdeaths,:numrecover] => ByRow(compute_R0) => :numrecover)
transform!(df_covid_BC,[:numactive,:numrecover] => ByRow(compute_s0) => :S)
rename!(df_covid_BC, :numactive => :I)
rename!(df_covid_BC, :numrecover => :R)
rename!(df_covid_BC, :date => :t)
select!(df_covid_BC,:t,:S,:I,:R)

# alternative: March 19, 20201 third wave started

# select the initial date
# choose 2020-03-26; no missing data
filter!(:t => >=(Date(2020,8,10)), df_covid_BC)

# before vaccine is administered: Dec 2020
df_covid_BC_pre = filter(:t => <(Date(2021,1,11)) ,df_covid_BC)

# after vaccine
df_covid_BC_post = filter(:t => >=(Date(2021,1,11)),df_covid_BC)

# save the output
CSV.write("data/covid_19_bc_all.csv",df_covid_BC)
CSV.write("data/covid19_bc_pre_vaccine.csv",df_covid_BC_pre)
CSV.write("data/covid19_bc_post_vaccine.csv",df_covid_BC_post)

# save as JLD
# save("data/covid_bc.jld","pre",df_covid_BC_pre,"post",df_covid_BC_post)

# before vaccine is administered: Dec 2020
df_covid_BC_second_wave = filter(:t => <(Date(2021,1,19)) ,df_covid_BC)

# after vaccine
df_covid_BC_third_wave = filter(:t => >=(Date(2021,3,19)),df_covid_BC)

# save the output
CSV.write("data/covid19_bc_second_wave.csv",df_covid_BC_second_wave)
CSV.write("data/covid19_bc_third_wave.csv",df_covid_BC_third_wave)


text:
using Flux: conv
using DSP: conv

# 4 = length(size(x))
join(x) = cat(x...; dims = 4)

prod0__(x, y) = join([x .* y[:,:,:,i:i] for i in 1:size(y,4)])
prod_0_(x, y) = join([x[:,:,:,i:i].*y for i in 1:size(x,4)])
prod__0(x, y) =
    join([sum([x[:, :, :, i] .* y[:, :, :, i] for i = 1:size(x, 4)])])
prod111(x,y) =
     join([
    x[:, :, :, 2].*y[:, :, :, 3] - x[:, :, :, 3].* y[:, :, :, 2],
    x[:, :, :, 3].* y[:, :, :, 1] - x[:, :, :, 1].* y[:, :, :, 3],
    x[:, :, :, 1].* y[:, :, :, 2] - x[:, :, :, 2].* y[:, :, :, 1],
    ])

function myconv(x,y)
    DSP.conv(x,y)
    # Flux.conv(cat(x,dims=5),cat(y,dims=5);pad=(size(y).-1).÷2)[:,:,:,1,1]
end
conv0__ =
    (x, y) -> join([myconv(x[:, :, :, 1], y) for y in eachslice(y, dims = 4)])
conv_0_ =
    (x, y) -> join([myconv(x, y[:, :, :, 1]) for y in eachslice(x, dims = 4)])
conv111 =
    (x, y) -> join([
        myconv(x[:, :, :, 2], y[:, :, :, 3]) - myconv(x[:, :, :, 3], y[:, :, :, 2]),
        myconv(x[:, :, :, 3], y[:, :, :, 1]) - myconv(x[:, :, :, 1], y[:, :, :, 3]),
        myconv(x[:, :, :, 1], y[:, :, :, 2]) - myconv(x[:, :, :, 2], y[:, :, :, 1]),
    ])
conv__0 =
    (x, y) -> join(sum(myconv.(eachslice(x, dims = 4), eachslice(y, dims = 4))))
conv121 =
    (x, y) -> join([
        myconv(x[:, :, :, 1], -y[:, :, :, 3] / sqrt(3) .- y[:, :, :, 5]) +
        myconv(x[:, :, :, 2], y[:, :, :, 2]) +
        myconv(x[:, :, :, 3], y[:, :, :, 1]),
        myconv(x[:, :, :, 1], y[:, :, :, 2]) +
        myconv(x[:, :, :, 2], 2 / sqrt(3) * y[:, :, :, 3]) +
        myconv(x[:, :, :, 3], y[:, :, :, 4]),
        myconv(x[:, :, :, 1], y[:, :, :, 1]) +
        myconv(x[:, :, :, 2], y[:, :, :, 4]) +
        myconv(x[:, :, :, 3], -y[:, :, :, 3] / sqrt(3) .+ y[:, :, :, 5]),
    ])

function FieldProd(l1, l2, l)
    # l1, l2, l = ranks
    if l1 == 0
        return  prod0__
    elseif l2 == 0
        return  prod_0_
    elseif l1 == l2 && l == 0
        return  prod__0
    elseif l1 == l2 ==l == 1
        return  prod111
    end

    # (x, y) -> yield(f(x, y), x.grid)
    # (x,y)->yield
end

function FieldConv(l1, l2, l)
    if l1 == 0
        f = conv0__
    elseif l2 == 0
        f = conv_0_
    elseif l1 == l2 && l == 0
        f = conv__0
    elseif l1 == l2 == l == 1
        f = conv111
    elseif (l1, l2, l) == (1, 2, 1)
        f = conv121
    end
    f
    # (x, y) -> tfconv(f,x,y)
end

"""
    field_prod(x::AbstractArray, y::AbstractArray; rank = nothing)
"""
function field_prod(x::AbstractArray, y::AbstractArray; rank = nothing)
    lx, ly = field_rank.([x, y])
    FieldProd(lx, ly, rank === nothing ? lx + ly : rank)(x, y)
end

"""
"""
function field_conv(x::AbstractArray, y::AbstractArray; rank = nothing)
    lx, ly = field_rank.([x, y])
    FieldConv(lx, ly, rank === nothing ? lx + ly : rank)(x, y)
end

"""
    function field_norm(field::AbstractArray)
"""
function fieldnorm(x::AbstractArray;squared=false)
# sqrt.(sum([x[:,:,:,i:i] .^ 2 for i in 1:size(x,4)]))
if  size(x,4)==1
    return x
end
r=sum([x[:,:,:,i:i] .^ 2 for i in 1:size(x,4)])
sqrt.(r.+.01).-.1
# sqrt.(sum([x .^ 2 for x in eachslice(field,dims=4)]))
# cat(sqrt.(sum([x .^ 2 for x in eachslice(field,dims=4)])),dims=4)
end

max_stars_repo_name: aced-differentiate/equiv-docs

max_issues_repo_name: aced-differentiate/equiv-docs

max_forks_repo_name: aced-differentiate/equiv-docs


text:
##################################################
# Computation of the NR diffusion coefficients
##################################################

##################################################
# Preliminary useful functions
##################################################

"""
    logCoulomb(iBath)

Coulomb logarithm of the family indexed by iBath in the source file.
Defined as the logarithm of the ratio of the mass of the SMBH over the mass of an element of the family iBath.

# Remark
Using the convention where log is the Neperian logarithm (in base e).

"""
function logCoulomb(iBath::Int64)
    m = tabmBath[iBath]
    Q = mBH/m
    return log(Q)
end

"""
    prefac_J(E)

Numerical prefactor J(E) used to convert the PDF f(E,R) to the number density N(E,R), where R = j^2.
Depends on the binding energy E, which is the opposite of the gravitational energy.

Explicit formula is N(E,R) = J(E)*f(E,R).
"""
function prefac_J(E::Float64)
    return (2.0)^(1/2)*pi^3*(G*mBH)^3*E^(-5/2)
end

"""
    E_bind(a)

Binding energy of the system, which is the opposite of the gravitational energy.
Defined as `E(a) = G*mBH / (2a)`, where `a` is the semi-major axis of the orbit, G is the gravitational constant and mBH the mass of the SMBH.
"""
function E_bind(a::Float64)
    return (G*mBH)/(2.0*a)
end

"""
    j_to_R(j)

Converts reduced angular momentum j = J/Jc(a) into R = j^2, where J is the angular momentum and Jc(a) the angular momentum of a circular orbit.
"""
function j_to_R(j::Float64)
    return j^2
end

"""
    DF_E(E,iBath)

R-averaged distribution function in binding enery E, defined in E space, for the family indexed by iBath in the source file.
R is the squared reduced angular momentum j^2.
"""
function DF_E(E::Float64,iBath::Int64)
    E0    = E_bind(taba0Bath[iBath])
    NBar  = tabNbarBath[iBath]
    gamma = tabgammaBath[iBath]
    return NBar/prefac_J(E) * (G*mBH)/(2.0*E^2) * (E0/E)^(2.0-gamma)
end

"""
    gamma_prefactor(iBath)

Gamma prefactor which appears in Fi integrals, for the family indexed by iBath in the source file.
Used in the computation of the non-resonant diffusion coefficients.
"""
function gamma_prefactor(iBath::Int64)
    m = tabmBath[iBath]
    return 4.0*pi*(G*m)^2 * logCoulomb(iBath)
end

"""
    xpm(R)

xplus and xminus quantities which appear in the bounds of the integrals involved in the computation of the non-resonant diffusion coefficients.

# Remark
x = E(r)/gravitational potential(r)
"""
function xpm(R::Float64)
    return (1.0+sqrt(1.0-R))/2.0,(1.0-sqrt(1.0-R))/2.0
end

"""
    inverse_reduced_E_bind(Ep,E)

Reduced binding energy used in the computation of integrals Fi.
Defined as Ep over E.
"""
function inverse_reduced_E_bind(Ep::Float64,E::Float64)
    return Ep/E
end

##################################################
# Bounds in A_k, B_k and C_k (xp>1/s) integrals
##################################################

"""
    _beta(R,s)

Bound in Ck integral when x_plus > 1/s, where s is the reduced binding energy and R is the squared reduced angular momentum.
"""
function _beta(R::Float64,s::Float64)
    xp,xm = xpm(R)
    res = 1.0/(xm*s)
    res = sqrt(res)
    return acosh(res)
end

"""
    _eta(R)

Bound in Ak integral when xp <= 1/s, where s is the reduced binding energy and R is the squared reduced angular momentum.
"""
function _eta(R::Float64)
    xp,xm = xpm(R)
    res = 1.0/(2.0*xm)
    res = sqrt(res)
    return acosh(res)
end

"""
    _delta(R)

Bound in Bk integral when xp <= 1/s, where s is the reduced binding energy and R is the squared reduced angular momentum.
"""
function _delta(R::Float64)
    xp,xm = xpm(R)
    res = 1.0/(2.0*xp)
    res = sqrt(res)
    return acos(res)
end

##################################################
# Computation of the NR diffusion coefficients.
##################################################

"""
    _fB(theta,s,R)

Computes all seven integrands within Bk integrals for 1<=k<=7.
Returns a 7-uple, ordered from k=1 to k=7.
See `Conrad et Kulsrud (1978)`.

# Remarks
- All the integrands are computed at the same time.
- Avoids multiple evalution of cos(theta), xplus and xminus.
- Useful when xplus <= 1/s.

# Arguments
- `theta::Float64`: Integration variable.
- `s    ::Float64`: Reduced binding energy.
- `R    ::Float64`: Squared reduced angular momentum j^2.
"""
function _fB(theta::Float64,s::Float64,R::Float64)
    xp,xm = xpm(R)
    ym    = xm/xp
    csTh  = cos(theta)

    fB1 = csTh^3/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(1)/sqrt(1.0-xp*csTh^2)^(1)
    fB2 = csTh^5/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(1)/sqrt(1.0-xp*csTh^2)^(3)
    fB3 = csTh^7/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(1)/sqrt(1.0-xp*csTh^2)^(1)
    fB4 = csTh/sqrt(csTh^2-ym)   * sqrt(1.0-s*xp*csTh^2)^(3)/sqrt(1.0-xp*csTh^2)^(1)
    fB5 = csTh^3/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(3)/sqrt(1.0-xp*csTh^2)^(3)
    fB6 = csTh^5/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(3)/sqrt(1.0-xp*csTh^2)^(5)
    fB7 = csTh^7/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(3)/sqrt(1.0-xp*csTh^2)^(3)
    return fB1,fB2,fB3,fB4,fB5,fB6,fB7
end

"""
    _fAC(theta,s,R)

Computes all 14 integrands within Ak and Ck integrals for 1<=k<=7.
Returns a 7-uple, ordered from k=1 to k=7.
See `Conrad et Kulsrud (1978)`.

# Remarks
- All the integrands are computed at the same time.
- Avoids multiple evalution of cos(theta), xplus and xminus.
- Useful when xplus <= 1/s for Ak.
- Useful when xplus > 1/s for Ck.

# Arguments
- `theta::Float64`: Integration variable.
- `s    ::Float64`: Reduced binding energy.
- `R    ::Float64`: Squared reduced angular momentum j^2.
"""
function _fAC(theta::Float64,s::Float64,R::Float64)
    xp,xm = xpm(R)
    ym    = xp/xm
    cshTh = cosh(theta)

    fAC1 = cshTh^3/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(1)/sqrt(1.0-xm*cshTh^2)^(1)
    fAC2 = cshTh^5/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(1)/sqrt(1.0-xm*cshTh^2)^(3)
    fAC3 = cshTh^7/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(1)/sqrt(1.0-xm*cshTh^2)^(1)
    fAC4 = cshTh  /sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(3)/sqrt(1.0-xm*cshTh^2)^(1)
    fAC5 = cshTh^3/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(3)/sqrt(1.0-xm*cshTh^2)^(3)
    fAC6 = cshTh^5/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(3)/sqrt(1.0-xm*cshTh^2)^(5)
    fAC7 = cshTh^7/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(3)/sqrt(1.0-xm*cshTh^2)^(3)
    return fAC1,fAC2,fAC3,fAC4,fAC5,fAC6,fAC7
end

"""
    _C_NR(s,R,nbC=500)

Computes all seven Ck integrals for 1<=k<=7 used for the computation of the non-resonance diffusion coefficients.
Midpoint sampling rule is used for integration.
Returns a 7-uple, ordered from k=1 to k=7.
See `Conrad et Kulsrud (1978)`.

# Remarks
- All the integrands are computed at the same time.
- Avoids multiple evalution of cos(theta), xplus and xminus, etc.

# Arguments
- `s  ::Float64`: Reduced binding energy.
- `R  ::Float64`: Squared reduced angular momentum j^2.
- `nbC::Float64`: Sampling number of the computed integrals. Set as 500 by default.
"""
function _C_NR(s::Float64,R::Float64,nbC::Int64=500)
    xp,xm = xpm(R)
    C1,C2,C3,C4,C5,C6,C7 = 0.0,0.0,0.0,0.0,0.0,0.0,0.0

    if xp>1/s
        supC = _beta(R,s)
        stepC = supC/nbC
        prefactorC = 4.0*stepC/pi
    else
        supA = _eta(R)
        supB = _delta(R)
        stepA = supA/nbC
        stepB = supB/nbC
        prefactorA = 4.0*stepA/pi
        prefactorB = 4.0*stepB/pi
        A1,A2,A3,A4,A5,A6,A7 = 0.0,0.0,0.0,0.0,0.0,0.0,0.0
        B1,B2,B3,B4,B5,B6,B7 = 0.0,0.0,0.0,0.0,0.0,0.0,0.0
    end
    for jTheta=1:nbC
        if xp>1/s
            thetaC = (jTheta-0.5)*stepC
            fAC1,fAC2,fAC3,fAC4,fAC5,fAC6,fAC7 = _fAC(thetaC,s,R)
            C1 += fAC1
            C2 += fAC2
            C3 += fAC3
            C4 += fAC4
            C5 += fAC5
            C6 += fAC6
            C7 += fAC7
        else
            thetaA = (jTheta-0.5)*stepA
            thetaB = (jTheta-0.5)*stepB
            fAC1,fAC2,fAC3,fAC4,fAC5,fAC6,fAC7 = _fAC(thetaA,s,R)
            fB1,fB2,fB3,fB4,fB5,fB6,fB7 = _fB(thetaB,s,R)
            A1 += fAC1
            A2 += fAC2
            A3 += fAC3
            A4 += fAC4
            A5 += fAC5
            A6 += fAC6
            A7 += fAC7
            B1 += fB1
            B2 += fB2
            B3 += fB3
            B4 += fB4
            B5 += fB5
            B6 += fB6
            B7 += fB7
        end
    end
    if xp>1/s
        C1 *= prefactorC*xm
        C2 *= prefactorC*xm^2
        C3 *= prefactorC*xm^3
        C4 *= prefactorC
        C5 *= prefactorC*xm
        C6 *= prefactorC*xm^2
        C7 *= prefactorC*xm^3
    else
        A1 *= prefactorA*xm
        A2 *= prefactorA*xm^2
        A3 *= prefactorA*xm^3
        A4 *= prefactorA
        A5 *= prefactorA*xm
        A6 *= prefactorA*xm^2
        A7 *= prefactorA*xm^3
        B1 *= prefactorB*xp
        B2 *= prefactorB*xp^2
        B3 *= prefactorB*xp^3
        B4 *= prefactorB
        B5 *= prefactorB*xp
        B6 *= prefactorB*xp^2
        B7 *= prefactorB*xp^3
        C1 = A1 + B1
        C2 = A2 + B2
        C3 = A3 + B3
        C4 = A4 + B4
        C5 = A5 + B5
        C6 = A6 + B6
        C7 = A7 + B7
    end
    return C1,C2,C3,C4,C5,C6,C7
end

"""
    _F_NR_0(E,iBath)

Analytic evaluation of the F0 integral, at binding energy E and for the family indexed by iBath in the source file.
See `Conrad et Kulsrud (1978)`.
"""
function _F_NR_0(E::Float64,iBath::Int64)
    a0 = taba0Bath[iBath]
    E0 = E_bind(a0)
    NBar  = tabNbarBath[iBath]
    gamma = tabgammaBath[iBath]
    gamma_pref = gamma_prefactor(iBath)
    return 4.0*pi*gamma_pref* a0 * 1.0/prefac_J(E)*NBar/(gamma-0.5) * (E0/E)^(3.0-gamma)
end

"""
    _F_NR(E,R,iBath,nbC=500)

Computes all seven Fk integrals for 1<=k<=7 used for the computation of the non-resonance diffusion coefficients.
Midpoint sampling rule is used for integration.
Returns a 7-uple, ordered from k=1 to k=7.
See `Conrad et Kulsrud (1978)`.

# Remarks
- All the integrands are computed at the same time.
- Avoids multiple evalution of cos(theta), xplus and xminus, etc.

# Arguments
- `E    ::Float64`: Binding energy.
- `R    ::Float64`: Squared reduced angular momentum j^2.
- `iBath::Float64`: Index in the source file of the considered family.
- `nbC  ::Float64`: Sampling number of the computed integrals. Set as 500 by default.
"""
function _F_NR(E::Float64,R::Float64,iBath::Int64,nbF::Int64=500)
    xp,xm = xpm(R)
    inf = E
    sup = E/xm
    step = E/nbF * (1.0/xm - 1.0)
    F1,F2,F3,F4,F5,F6,F7 = 0.0,0.0,0.0,0.0,0.0,0.0,0.0
    for kEnergy=1:nbF
        Ep = inf+(kEnergy-0.5)*step
        s  = inverse_reduced_E_bind(Ep,E)
        C1,C2,C3,C4,C5,C6,C7 = _C_NR(s,R)
        DF = DF_E(Ep,iBath)
        F1 += DF*C1
        F2 += DF*C2
        F3 += DF*C3
        F4 += DF*C4
        F5 += DF*C5
        F6 += DF*C6
        F7 += DF*C7
    end
    prefactor = 4.0*pi*gamma_prefactor(iBath)*step
    F1 *= prefactor
    F2 *= prefactor
    F3 *= prefactor
    F4 *= prefactor
    F5 *= prefactor
    F6 *= prefactor
    F7 *= prefactor
    return F1,F2,F3,F4,F5,F6,F7
end

##################################################
# Diffusion coefficients
##################################################

"""
    DNR_EJ(E,j,iBath,m_test=0.0)

Non-resonant diffusion coefficients in (E,R) coordinates and Bar-Or&Alexander (2016) convention.
Returns the 5-uple `(DE, DEE/E, DJ/Jc, DJJ/Jc^2, DEJ/Jc)` for the family indexed by iBath.

# Arguments
- `E     ::Float64`: Binding energy.
- `j     ::Float64`: Reduced angular momentum `j=J/Jc`.
- `iBath ::Float64`: Index in the source file of the considered family.
- `m_test::Float64`: Mass of a test particle.
"""
function DNR_EJ(E::Float64,j::Float64,iBath::Int64,m_test::Float64=0.0)
    R = j^2
    F1,F2,F3,F4,F5,F6,F7 = _F_NR(E,R,iBath)
    F0 = _F_NR_0(E,iBath)
    m = tabmBath[iBath]
    D_E = -F0 + m_test/m * F1
    D_EE_over_E = 4/3*(F0 + F4)
    D_J_over_Jc = 1.0/(j*E) * ( (5.0 - 3.0*j^2)/(12.0)*F0 - j^2*(m_test + m)/(2.0*m)*F2 + F3 - 1/3*F7 )
    D_JJ_over_JcSQ = 1.0/E* ( (5.0 - 3.0*j^2)/6.0*F0 + 1/2*j^2*F6 - 1/2*j^2*F2 + 2.0*F3 - 2/3*F7 )
    D_EJ_over_Jc = -2/3*j*(F0 + F5)
    return D_E,D_EE_over_E,D_J_over_Jc,D_JJ_over_JcSQ,D_EJ_over_Jc
end

"""
    DNR_J_red(a,j)

Total non-resonant reduced diffusion coefficient DJ/Jc in (E,R) coordinates and Bar-Or&Alexander (2016) convention.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_J_red(a::Float64,j::Float64)
    res = 0.0
    E = E_bind(a)
    for iBath=1:nbBath
        res += DNR_EJ(E,j,iBath)[3]
    end
    return res
end

"""
    DNR_JJ_red(a,j)

Total non-resonant reduced diffusion coefficient DJJ/Jc^2 in (E,R) coordinates and Bar-Or&Alexander (2016) convention.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_JJ_red(a::Float64,j::Float64)
    res = 0.0
    E = E_bind(a)
    for iBath=1:nbBath
        res += DNR_EJ(E,j,iBath)[4]
    end
    return res
end

"""
    DNR_red(a,j,iBath)

Non-resonant reduced diffusion coefficients in (a,j) coordinates and Bar-Or&Alexander (2016) convention.
Returns the 5-uple `(Da, Daa, Dj, Djj, Daj)` for the family indexed by iBath.

# Arguments
- `a     ::Float64`: Semi-major axis.
- `j     ::Float64`: Reduced angular momentum `j=J/Jc`.
- `iBath ::Float64`: Index in the source file of the considered family.
"""
function DNR_red(a::Float64,j::Float64,iBath::Int64)
    E = E_bind(a)
    cst = 2.0/(G*mBH)
    D_E,D_EE_over_E,D_J_over_Jc,D_JJ_over_JcSQ,D_EJ_over_Jc = DNR_EJ(E,j,iBath)
    D_a = cst*a^2*(-D_E + D_EE_over_E)
    D_aa = cst*a^3*D_EE_over_E
    D_j = D_J_over_Jc + j/(2.0*E)*D_E + 1.0/(2.0*E)*D_EJ_over_Jc - j/(8.0*E)*D_EE_over_E
    D_jj = D_JJ_over_JcSQ + j/E*D_EJ_over_Jc + j^2/(4.0*E)*D_EE_over_E
    D_aj = -cst*a^2*(D_EJ_over_Jc + j/2*D_EE_over_E)
    return D_a,D_aa,D_j,D_jj,D_aj
end

"""
    DNR_aj(a,j)

Total non-resonant reduced diffusion coefficients in (a,j) coordinates and Bar-Or&Alexander (2016) convention.
Returns the 5-uple `(Da, Daa, Dj, Djj, Daj)`.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_aj(a::Float64,j::Float64) #Wrapped function to give total D_NR_jj
    D_a,D_aa,D_j,D_jj,D_aj = 0.0,0.0,0.0,0.0,0.0
    for iBath=1:nbBath
        Di_a,Di_aa,Di_j,Di_jj,Di_aj = DNR_red(a,j,iBath)
        D_a  += Di_a
        D_aa += Di_aa
        D_j  += Di_j
        D_jj += Di_jj
        D_aj += Di_aj
    end
    return D_a,D_aa,D_j,D_jj,D_aj
end

"""
    DNR_jj(a,j)

Total non-resonant reduced diffusion coefficient Djj in (a,j) coordinates and Bar-Or&Alexander (2016) convention.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_jj(a::Float64,j::Float64)
    return DNR_aj(a,j)[4]
end

"""
    DNR_j(a,j)

Total non-resonant reduced diffusion coefficient Dj in (a,j) coordinates and Bar-Or&Alexander (2016) convention.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_j(a::Float64,j::Float64)
    return DNR_aj(a,j)[3]
end


text:
##############################################################################
#
# In this test problem we consider a quadratic penalty method.
# This example features an algorithm with the 3 steps:
# the penalization - the unconstrained min - the 1d min
#
# Note that there is no optimization of the evaluations here.
# The penalization gives an approximation of the gradients, multipliers...
#
# Note the use of a structure for the algorithmic parameters which is
# forwarded to all the 3 steps. If a parameter is not mentioned, then the default
# entry in the algorithm will be taken.
#
#############################################################################

#include("uncons.jl")

##############################################################################
#
# Quadratic penalty algorithm
# fill_in! used instead of update! (works but usually more costly in evaluations)
# subproblems are solved via Newton method
#
#############################################################################
function penalty(stp::NLPStopping; rho0 = 1.0, rho_min = 1e-10, rho_update = 0.5, prms = nothing)

  #algorithm's parameters
  rho = rho0

  #First call to the stopping
  #Becareful here, some entries must be filled in first.
  fill_in!(stp, stp.current_state.x)
  OK = start!(stp)

  #prepare the subproblem stopping:
  sub_nlp_at_x = NLPAtX(stp.current_state.x)
  sub_pb = ADNLPModel(
    x ->
      obj(stp.pb, x) +
      1 / rho * norm(max.(cons(stp.pb, x) - stp.pb.meta.ucon, 0.0))^2 +
      1 / rho * norm(max.(-cons(stp.pb, x) + stp.pb.meta.lcon, 0.0))^2,
    x0,
  )
  sub_stp =
    NLPStopping(sub_pb, sub_nlp_at_x, main_stp = stp, optimality_check = unconstrained_check)

  #main loop
  while !OK

    #solve the subproblem
    reinit!(sub_stp)
    sub_stp.meta.atol = min(rho, sub_stp.meta.atol)
    global_newton(sub_stp, prms)

    #Update all the entries of the State
    fill_in!(stp, sub_stp.current_state.x)

    #Either stop! is true OR the penalty parameter is too small
    if rho < rho_min
      stp.meta.suboptimal = true
    end
    OK = stop!(stp)

    @show stp.meta.nb_of_stop, OK, rho

    #update the penalty parameter if necessary
    if !OK
      rho = rho * rho_update
      sub_stp.pb = ADNLPModel(
        x ->
          obj(stp.pb, x) +
          1 / rho * norm(max.(cons(stp.pb, x) - stp.pb.meta.ucon, 0.0))^2 +
          1 / rho * norm(max.(-cons(stp.pb, x) + stp.pb.meta.lcon, 0.0))^2,
        x0,
      )
    end
  end

  return stp
end

##############################################################################
#
# Quadratic penalty algorithm: buffer function
#
#############################################################################
function penalty(stp::NLPStopping, prms)

  #extract required values in the prms file
  r0 = :rho0 ∈ fieldnames(typeof(prms)) ? prms.rho0 : 1.0
  rm = :rho_min ∈ fieldnames(typeof(prms)) ? prms.rho_min : 1e-10
  ru = :rho_update ∈ fieldnames(typeof(prms)) ? prms.rho_update : 0.5

  return penalty(stp, rho0 = r0, rho_min = rm, ru = 0.5, prms = prms)
end

##############################################################################
#
# Algorithmic parameters structure
#
#############################################################################

mutable struct Param

  #parameters for the penalty
  rho0::Float64 #initial value of the penalty parameter
  rho_min::Float64 #smallest possible parameter
  rho_update::Float64 #update of the penalty parameter

  #parameters of the unconstrained minimization
  armijo_prm::Float64 #Armijo parameter
  wolfe_prm::Float64 #Wolfe parameter
  onedsolve::Function #1D solver
  ls_func::Function

  #parameters of the 1d minimization
  back_update::Float64 #backtracking update

  function Param(;
    rho0::Float64 = 1.0,
    rho_min::Float64 = sqrt(eps(Float64)),
    rho_update::Float64 = 0.5,
    armijo_prm::Float64 = 0.01,
    wolfe_prm::Float64 = 0.99,
    onedsolve::Function = backtracking_ls,
    ls_func::Function = (x, y) -> armijo(x, y, τ₀ = armijo_prm),
    back_update::Float64 = 0.5,
  )
    return new(rho0, rho_min, rho_update, armijo_prm, wolfe_prm, onedsolve, ls_func, back_update)
  end
end


text:
export cholinv, cholsqrt

using LinearAlgebra
using PositiveFactorizations

import LinearAlgebra: BlasInt

cholinv(x)           = inv(fastcholesky(x))
cholinv(x::Diagonal) = Diagonal(inv.(diag(x)))
cholinv(x::Real)     = inv(x)

function cholinv(x::AbstractMatrix{T}) where { T <: LinearAlgebra.BlasFloat }
    y = fastcholesky(x)
    LinearAlgebra.inv!(y)
    return y.factors
end

cholsqrt(x)           = Matrix(fastcholesky(x).L)
cholsqrt(x::Diagonal) = Diagonal(sqrt.(diag(x)))
cholsqrt(x::Real)     = sqrt(x)

chollogdet(x)           = logdet(fastcholesky(x))
chollogdet(x::Diagonal) = logdet(x)
chollogdet(x::Real)     = logdet(x)

function cholinv_logdet(x)
    # calculate cholesky decomposition
    y = fastcholesky(x)

    # return inverse and log-determinant
    return inv(y), logdet(y)
end
function cholinv_logdet(x::AbstractMatrix{T}) where { T <: LinearAlgebra.BlasFloat }
    # calculate cholesky decomposition
    y = fastcholesky(x)

    # calculate logdeterminant of cholesky decomposition
    ly = logdet(y)

    # calculate inplace inverse of A and store in y.factors
    LinearAlgebra.inv!(y)

    # return inverse and log-determinant
    return y.factors, ly
end
cholinv_logdet(x::Diagonal) = Diagonal(inv.(diag(x))), mapreduce(z -> log(z), +, diag(x))
cholinv_logdet(x::Real)     = inv(x), log(abs(x))

function fastcholesky(mat::AbstractMatrix)
    A = copy(mat)
    C = fastcholesky!(A)
    if !isposdef(C)
        return cholesky(PositiveFactorizations.Positive, Hermitian(mat))
    else
        return C
    end
end

function fastcholesky!(A::AbstractMatrix)
    n = LinearAlgebra.checksquare(A)
    @inbounds for col=1:n
        @inbounds @simd for idx in 1:col-1
            A[col, col] -= A[col, idx]^2;
        end
        if A[col,col] <= 0
            return Cholesky(A, 'L', convert(BlasInt, -1))
        end
        A[col, col] = sqrt(A[col, col])

        @inbounds for row in col+1: n
            @inbounds @simd for idx in 1:col-1
                A[row, col] -= A[row, idx]*A[col, idx]
            end
            A[row, col] /= A[col, col]
        end
    end
    return Cholesky(A, 'L', convert(BlasInt, 0))
end

function fastcholesky!(A::AbstractMatrix{T}) where { T <: LinearAlgebra.BlasFloat }
    # blocked version (https://arxiv.org/pdf/1812.02056.pdf)

    # step size
    s = 250

    n = LinearAlgebra.checksquare(A)
    z = 1
    @inbounds for c in 1:n

        if c == z + s
            BLAS.gemm!('N', 'T', -one(T), view(A, c:n, z:c-1), view(A, c:n, z:c-1), one(T), view(A, c:n, c:n)) # replace with syrk once julia bug has been fixed
            z = c
        end

        @inbounds for k in z:c-1
            A[c,c] -= A[c,k]^2
        end
        if A[c,c] <= 0
            return Cholesky(A, 'L', convert(BlasInt, -1))
        end
        A[c,c] = sqrt(A[c,c])
        @inbounds for i in c+1:n
            @inbounds for k in z:c-1
                A[i,c] -= A[i,k]*A[c,k]
            end
            A[i,c] /= A[c,c]
        end
    end

    return Cholesky(A, 'L', convert(BlasInt, 0))
end


text:
using Test
include("module-3-exercise-2.jl")

@testset "$(get_module_3_excersise_2_author_name()): problem 2 tests" begin
    result = problem_2()
    @test get_module_3_excersise_2_author_name() != "YOUR NAME HERE"
    @test isapprox(2, result[1]; atol = 0.05)
    @test isapprox(1, result[2]; atol = 0.05)
    @test isapprox(3, result[1] + result[2]; atol = 0.05)
    @test isapprox(5, 4*result[1] - 3*result[2]; atol = 0.05)
end


text:
using KernelCompiler
using Test

# Can't run inside a testset :/
parent(x) = child(x)
child(x) = x+1

thunk1 = KernelCompiler.jit(parent, Tuple{Int})
@test thunk1 == KernelCompiler.jit(parent, Tuple{Int})

@test thunk1(2) == 3
child(x) = x+2

thunk2 = KernelCompiler.jit(parent, Tuple{Int})
@test thunk1 != thunk2
@test thunk2(2) == 4

parent() = child()
child() = 3

thunk1 = KernelCompiler.jit(parent, Tuple{})
@test thunk1 == KernelCompiler.jit(parent, Tuple{})

@test thunk1() == 3
child() = 4

thunk2 = KernelCompiler.jit(parent, Tuple{})
@test thunk1 != thunk2

@test thunk2() == 4


text:
module TestReplaceComp

using Test
using Mimi
import Mimi:
    compdefs, compname, compdef, components, comp_id, external_param_conns, external_params

@defcomp X begin
    x = Parameter(index = [time])
    y = Variable(index = [time])
    function run_timestep(p, v, d, t)
        v.y[t] = 1
    end
end

@defcomp X_repl begin
    x = Parameter(index = [time])
    y = Variable(index = [time])
    function run_timestep(p, v, d, t)
        v.y[t] = 2
    end
end

@defcomp bad1 begin
    x = Parameter()                 # parameter has same name but different dimensions
    y = Variable(index = [time])
end

@defcomp bad2 begin
    x = Parameter(index = [time])
    z = Variable(index = [time])    # different variable name
end

@defcomp bad3 begin
    z = Parameter()                 # different parameter name
    y = Variable(index = [time])
end

@defcomp bad4 begin
    x = Parameter{Symbol}(index = [time])   # different datatype
    y = Variable()                          # different variable dimensions
end

# 1. Test scenario where the replacement works

m = Model()
set_dimension!(m, :time, 2000:2005)
add_comp!(m, X)                         # Original component X
set_param!(m, :X, :x, zeros(6))
@test_throws ErrorException replace_comp!(m, X_repl, :X) # test that the old function name now errors
replace!(m, :X => X_repl)               # Replace X with X_repl
run(m)
@test length(components(m)) == 1        # Only one component exists in the model
@test m[:X, :y] == 2 * ones(6)          # Successfully ran the run_timestep function from X_repl

# 2. Test bad internal incoming parameter

m = Model()
set_dimension!(m, :time, 2000:2005)
add_comp!(m, X, :first)                                     # Add two components
add_comp!(m, X, :second)
connect_param!(m, :second => :x, :first => :y)              # Make an internal connection with a parameter with a time dimension
@test_throws ErrorException replace!(m, :second => bad1)    # Cannot make reconnections because :x in bad1 has different dimensions
@test_throws ErrorException replace_comp!(m, bad1.comp_id, :second, reconnect = false) # Test that the old function name now errors
replace!(m, :second => bad1, reconnect = false)             # Can replace without reconnecting
second = compdef(m, :second)
@test second.comp_id.comp_name == :bad1                     # Successfully replaced

# 3. Test bad internal outgoing variable

m = Model()
set_dimension!(m, :time, 2000:2005)
add_comp!(m, X, :first)                                  # Add two components
add_comp!(m, X, :second)
connect_param!(m, :second => :x, :first => :y)           # Make an internal connection from a variable with a time dimension
@test_throws ErrorException replace!(m, :first => bad2)  # Cannot make reconnections because bad2 does not have a variable :y
replace!(m, :first => bad2, reconnect = false)           # Can replace without reconnecting
first = compdef(m, :first)
@test first.comp_id.comp_name == :bad2                   # Successfully replaced

# 4. Test bad external parameter name

m = Model()
set_dimension!(m, :time, 2000:2005)
add_comp!(m, X)
set_param!(m, :X, :x, zeros(6))                     # Set external parameter for :x

# Replaces with bad3, but warns that there is no parameter by the same name :x
@test_logs(
    # (:warn, r".*parameter x no longer exists in component.*"),
    replace!(m, :X => bad3)
)

@test compname(compdef(m, :X)) == :bad3            # The replacement was still successful
@test length(external_param_conns(m)) == 0         # The external parameter connection was removed
@test length(external_params(m)) == 1              # The external parameter still exists

# 5. Test bad external parameter dimensions

m = Model()
set_dimension!(m, :time, 2000:2005)
add_comp!(m, X)
set_param!(m, :X, :x, zeros(6))                         # Set external parameter for :x
@test_throws ErrorException replace!(m, :X => bad1)     # Cannot reconnect external parameter, :x in bad1 has different dimensions

# 6. Test bad external parameter datatype

m = Model()
set_dimension!(m, :time, 2000:2005)
add_comp!(m, X)
set_param!(m, :X, :x, zeros(6))                         # Set external parameter for :x
@test_throws ErrorException replace!(m, :X => bad4)  # Cannot reconnect external parameter, :x in bad4 has different datatype

# 7. Test component name that doesn't exist

m = Model()
set_dimension!(m, :time, 2000:2005)
add_comp!(m, X)
@test_throws ErrorException replace!(m, :Z => X_repl)    # Component Z does not exist in the model, cannot be replaced

# 8. Test original position placement functionality

m = Model()
set_dimension!(m, :time, 2000:2005)
add_comp!(m, X, :c1)
add_comp!(m, X, :c2)
add_comp!(m, X, :c3)

replace!(m, :c3 => X_repl)        # test replacing the last component
@test compdef(m, :c3).comp_id == X_repl.comp_id

replace!(m, :c2 => X_repl)        # test replacing not the last one
@test compdef(m, :c2).comp_id == X_repl.comp_id

# 9. Test that any previous set parameters are preserved, even in the presence of default values

@defcomp A begin
    p1 = Parameter(default=1)
end

@defcomp B begin
    p1 = Parameter(default=2)
end

m = Model()
set_dimension!(m, :time, 10)
add_comp!(m, A)
set_param!(m, :A, :p1, 3)
replace!(m, :A => B)
run(m)
@test m[:A, :p1] == 3

end # module


text:
@with_kw mutable struct HalfspaceParameters{T <: AbstractFloat} <: AbstractPolicyParameters
    α::T = 0.
    a::Vector{T}
end

struct Halfspace <: AbstractProx
    params::HalfspaceParameters{Float64}
    norm::Float64

    function (::Type{Halfspace})(; kwargs...)
        params = HalfspaceParameters(; kwargs...)
        @unpack a = params
        return new(params,a⋅a)
    end
end

function prox!(halfspace::Halfspace,step::Real,xprev::AbstractVector,gcurr::AbstractVector,xcurr::AbstractVector)
    @unpack α,a = halfspace.params
    x = xprev - step*gcurr
    diff = a⋅x - α
    xcurr .= x - (max(diff,zero(eltype(a)))/halfspace.norm)*a
end


text:
@testset "Non-commutative GramMatrix" begin
    @ncpolyvar x y
    P = GramMatrix([2 3 4;
                       3 4 5;
                       4 5 6], [x*y, x^2, y^2])
    @test P.Q.Q == [4, 3, 5, 2, 4, 6]
    P = GramMatrix((i,j) -> i + j, [x*y, x^2, y^2])
    @test P.Q.Q == [4, 3, 5, 2, 4, 6]
    p = polynomial(P)
    @test p.a == [4, 3, 5, 4, 3, 2, 6, 5, 4]
    @test p.x == monovec([x^4, x^3*y, x^2*y^2, x*y^3, x*y*x^2, x*y*x*y, y^4, y^2*x^2, y^2*x*y])
    @inferred GramMatrix(Matrix{Float64}(0, 0), typeof(x)[]) == 0
    @test GramMatrix(Matrix{Float64}(0, 0), typeof(x)[]) isa AbstractPolynomialLike{Float64}
    @test GramMatrix(Matrix{Float64}(0, 0), typeof(x)[]) == 0
    P = GramMatrix((i,j) -> ((i,j) == (1,1) ? 2 : 0), [x*y, x^2, y^2])
    Q = GramMatrix([0 1; 1 0], [x^2, y^2])
    @test P != Q
end