text:
using BinaryBuilder

src_version = v"3.7.2"  # also change in raw script string below

# Collection of sources required to build GEOS
sources = [
    "http://download.osgeo.org/geos/geos-$src_version.tar.bz2" =>
    "2166e65be6d612317115bfec07827c11b403c3f303e0a7420a2106bc999d7707",
]

# Bash recipe for building across all platforms
script = raw"""
# use GCC ar instead of the LLVM one on OSX,
# see https://github.com/JuliaPackaging/BinaryBuilder.jl/issues/388
if [[ ${target} == *darwin* ]]; then
    export AR=/opt/${target}/bin/${target}-ar
fi

cd $WORKSPACE/srcdir
cd geos-3.7.2/
./configure --prefix=$prefix --host=$target
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 = supported_platforms()
platforms = expand_gcc_versions(platforms)

# The products that we will ensure are always built
products(prefix) = [
    LibraryProduct(prefix, "libgeos_c", :libgeos),
    LibraryProduct(prefix, "libgeos", :libgeos_cpp)
]

# Dependencies that must be installed before this package can be built
dependencies = []

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


text:
"""
    Atmosphere

Functions concerning atmospheric conditions: air properites, velocities, etc.

# properties

Calculate values for air density, dynamic viscosity, and speed of sound.

## Methods
- atmospherefit
- temp_presdrela
- temp_presnasa
- density
- viscosity
- speedofsound

# wind models

"""
module Atmosphere

include("properties.jl")

end #module Atmosphere

ext: jl


text:
"""
    SurfaceFluxes

  Surface flux functions, e.g., for buoyancy flux,
  friction velocity, and exchange coefficients.

## Sub-modules
  - module Byun1990
  - module Nishizawa2018

## Interface
  - [`compute_buoyancy_flux`](@ref) computes the buoyancy flux
  - In addition, each sub-module has the following functions:
    - [`compute_MO_len`](@ref) computes the Monin-Obukhov length
    - [`compute_friction_velocity`](@ref) computes the friction velocity
    - [`compute_exchange_coefficients`](@ref) computes the exchange coefficients

## References

  - Ref. Byun1990:
    - Byun, Daewon W. "On the analytical solutions of flux-profile
      relationships for the atmospheric surface layer." Journal of
      Applied Meteorology 29.7 (1990): 652-657.
      https://doi.org/10.1175/1520-0450(1990)029<0652:OTASOF>2.0.CO;2

  - Ref. Wyngaard1975:
    - Wyngaard, John C. "Modeling the planetary boundary layer-Extension
      to the stable case." Boundary-Layer Meteorology 9.4 (1975): 441-460.
      https://link.springer.com/article/10.1007/BF00223393

  - Ref. Nishizawa2018:
    - Nishizawa, S., and Y. Kitamura. "A Surface Flux Scheme Based on the
      Monin-Obukhov Similarity for Finite Volume Models." Journal of
      Advances in Modeling Earth Systems 10.12 (2018): 3159-3175.
      https://doi.org/10.1029/2018MS001534

  - Ref. Businger1971:
    - Businger, Joost A., et al. "Flux-profile relationships in the
      atmospheric surface layer." Journal of the atmospheric Sciences
      28.2 (1971): 181-189.

"""
module SurfaceFluxes

using RootSolvers
using ..MoistThermodynamics
using ..PlanetParameters

# export compute_buoyancy_flux

"""
    compute_buoyancy_flux(shf, lhf, T_b, qt_b, ql_b, qi_b, alpha0_0)

Computes buoyancy flux given sensible heat flux `shf`,
latent heat flux `lhf`, surface boundary temperature `T_b`,
total specific humidity `qt_b`, liquid specific humidity `ql_b`,
ice specific humidity `qi_b` and specific `alpha0_0`.
"""
function compute_buoyancy_flux(shf, lhf, T_b, qt_b, ql_b, qi_b, alpha0_0)
    cp_ = cp_m(PhasePartition(qt_b, ql_b, qi_b))
    lv = latent_heat_vapor(T_b)
    temp1 = (molmass_ratio - 1)
    temp2 = (shf + temp1 * cp_ * T_b * lhf / lv)
    return (grav * alpha0_0 / cp_ / T_b * temp2)
end

module Byun1990

using RootSolvers
using ...MoistThermodynamics
using ...PlanetParameters

""" Computes ψ_m for stable case. See Eq. 12 Ref. Byun1990 """
ψ_m_stable(ζ, ζ_0, β_m) = -β_m * (ζ - ζ_0)

""" Computes ψ_h for stable case. See Eq. 13 Ref. Byun1990 """
ψ_h_stable(ζ, ζ_0, β_h) = -β_h * (ζ - ζ_0)

""" Computes ψ_m for unstable case. See Eq. 14 Ref. Byun1990 """
function ψ_m_unstable(ζ, ζ_0, γ_m)
    x(ζ) = sqrt(sqrt(1 - γ_m * ζ))
    temp(ζ, ζ_0) = log1p((ζ - ζ_0) / (1 + ζ_0))  # log((1 + ζ)/(1 + ζ_0))
    ψ_m = (
        2 * temp(x(ζ), x(ζ_0)) + temp(x(ζ)^2, x(ζ_0)^2) - 2 * atan(x(ζ)) +
        2 * atan(x(ζ_0))
    )
    return ψ_m
end

""" Computes ψ_h for unstable case. See Eq. 15 Ref. Byun1990 """
function ψ_h_unstable(ζ, ζ_0, γ_h)
    y(ζ) = sqrt(1 - γ_h * ζ)
    ψ_h = 2 * log1p((y(ζ) - y(ζ_0)) / (1 + y(ζ_0))) # log((1 + y(ζ))/(1 + y(ζ_0)))
    return ψ_h
end

"""
    compute_MO_len(u, flux)

Computes the Monin-Obukhov length (Eq. 3 Ref. Byun1990)
"""
compute_MO_len(u, flux) = -u^3 / (flux * k_Karman)

"""
    compute_friction_velocity(u_ave, flux, z_0, z_1, β_m, γ_m, tol_abs, iter_max)

Computes roots of friction velocity equation (Eq. 10 in Ref. Byun1990)

`u_ave = u_* ( ln(z/z_0) - ψ_m(z/L, z_0/L) ) /κ        Eq. 10 in Ref. Byun1990`
"""
function compute_friction_velocity(
    u_ave,
    flux,
    z_0,
    z_1,
    β_m,
    γ_m,
    tol_abs,
    iter_max,
)

    ustar_0 = u_ave * k_Karman / log(z_1 / z_0)
    ustar = ustar_0
    let u_ave = u_ave, flux = flux, z_0 = z_0, z_1 = z_1, β_m = β_m, γ_m = γ_m

        # use neutral condition as first guess
        stable = z_1 / compute_MO_len(ustar_0, flux) >= 0
        function compute_ψ_m(u)
            L_MO = compute_MO_len(u, flux)
            ζ = z_1 / L_MO
            ζ_0 = z_0 / L_MO
            return stable ? ψ_m_stable(ζ, ζ_0, β_m) : ψ_m_unstable(ζ, ζ_0, γ_m)
        end
        function compute_u_ave_over_ustar(u)
            return (log(z_1 / z_0) - compute_ψ_m(u)) / k_Karman # Eq. 10 in Ref. Byun1990
        end
        compute_ustar(u) = u_ave / compute_u_ave_over_ustar(u)

        if (abs(flux) > 0)
            ustar_1 = compute_ustar(ustar_0)
            sol = RootSolvers.find_zero(
                u -> u_ave - u * compute_u_ave_over_ustar(u),
                ustar_0,
                ustar_1,
                SecantMethod(),
                CompactSolution(),
                tol_abs,
                iter_max,
            )
            ustar = sol.root
        end

    end

    return ustar
end

"""
    compute_exchange_coefficients(Ri, z_b, z_0, γ_m, γ_h, β_m, β_h, Pr_0)

Computes exchange transfer coefficients:
 - C_D  momentum exchange coefficient      (Eq. 36)
 - C_H  thermodynamic exchange coefficient (Eq. 37)
 - L_mo Monin-Obukhov length               (re-arranged Eq. 3)
"""
function compute_exchange_coefficients(Ri, z_b, z_0, γ_m, γ_h, β_m, β_h, Pr_0)
    logz = log(z_b / z_0)
    zfactor = z_b / (z_b - z_0) * logz
    s_b = Ri / Pr_0
    if Ri > 0
        temp = ((1 - 2 * β_h * Ri) - sqrt(1 + 4 * (β_h - β_m) * s_b))
        ζ = zfactor / (2 * β_h * (β_m * Ri - 1)) * temp                    # Eq. 19 in Ref. Byun1990
        L_mo = z_b / ζ                                             # LHS of Eq. 3 in Byun1990
        ζ_0 = z_0 / L_mo
        ψ_m = ψ_m_stable(ζ, ζ_0, β_m)
        ψ_h = ψ_h_stable(ζ, ζ_0, β_h)
    else
        Q_b = 1 / 9 * (1 / (γ_m^2) + 3 * γ_h / γ_m * s_b^2)           # Eq. 31 in Ref. Byun1990
        P_b = 1 / 54 * (-2 / (γ_m^3) + 9 / γ_m * (-γ_h / γ_m + 3) * s_b^2) # Eq. 32 in Ref. Byun1990
        crit = Q_b^3 - P_b^2
        if crit < 0
            T_b = cbrt(sqrt(-crit) + abs(P_b))                     # Eq. 34 in Ref. Byun1990
            ζ = zfactor * (1 / (3 * γ_m) - (T_b + Q_b / T_b))              # Eq. 29 in Ref. Byun1990
        else
            θ_b = acos(P_b / sqrt(Q_b^3))                            # Eq. 33 in Ref. Byun1990
            ζ = zfactor * (-2 * sqrt(Q_b) * cos(θ_b / 3) + 1 / (3 * γ_m))  # Eq. 28 in Ref. Byun1990
        end
        L_mo = z_b / ζ                                             # LHS of Eq. 3 in Byun1990
        ζ_0 = z_0 / L_mo
        ψ_m = ψ_m_unstable(ζ, ζ_0, γ_m)
        ψ_h = ψ_h_unstable(ζ, ζ_0, γ_h)
    end
    cu = k_Karman / (logz - ψ_m)                  # Eq. 10 in Ref. Byun1990, solved for u^*
    cth = k_Karman / (logz - ψ_h) / Pr_0            # Eq. 11 in Ref. Byun1990, solved for h^*
    C_D = cu^2                                                 # Eq. 36 in Byun1990
    C_H = cu * cth                                               # Eq. 37 in Byun1990
    return C_D, C_H, L_mo
end

end # Byun1990 module

module Nishizawa2018
using RootSolvers
using ...MoistThermodynamics
using ...PlanetParameters

""" Computes R_z0 expression, defined after Eq. 15 Ref. Nishizawa2018 """
compute_R_z0(z_0, Δz) = 1 - z_0 / Δz

""" Computes f_m in Eq. A7 Ref. Nishizawa2018 """
compute_f_m(ζ) = sqrt(sqrt(1 - 15 * ζ))

""" Computes f_h in Eq. A8 Ref. Nishizawa2018 """
compute_f_h(ζ) = sqrt(1 - 9 * ζ)

""" Computes ψ_m in Eq. A3 Ref. Nishizawa2018 """
function compute_ψ_m(ζ, L, a)
    f_m = compute_f_m(ζ)
    return L >= 0 ? -a * ζ :
           log((1 + f_m)^2 * (1 + f_m^2) / 8) - 2 * atan(f_m) + π / 2
end

""" Computes ψ_h in Eq. A4 Ref. Nishizawa2018 """
compute_ψ_h(ζ, L, a, Pr) =
    L >= 0 ? -a * ζ / Pr : 2 * log((1 + compute_f_h(ζ)) / 2)

""" Computes Ψ_m in Eq. A5 Ref. Nishizawa2018 """
function compute_Ψ_m(ζ, L, a, tol)
    if ζ < tol
        return ζ >= 0 ? -a * ζ / 2 : -15 * ζ / 8 # Computes Ψ_m in Eq. A13 Ref. Nishizawa2018
    else
        f_m = compute_f_m(ζ)
        # Note that "1-f^3" in Ref. Nishizawa2018 is a typo, it is supposed to be "1-f_m^3".
        # This was confirmed by communication with the author.
        return L >= 0 ? -a * ζ / 2 :
               log((1 + f_m)^2 * (1 + f_m^2) / 8) - 2 * atan(f_m) + π / 2 - 1 +
               (1 - f_m^3) / (12 * ζ)
    end
end

""" Computes Ψ_h in Eq. A6 Ref. Nishizawa2018 """
function compute_Ψ_h(ζ, L, a, Pr, tol)
    if ζ < tol
        return ζ >= 0 ? -a * ζ / (2 * Pr) : -9 * ζ / 4 # Computes Ψ_h in Eq. A14 Ref. Nishizawa2018
    else
        f_h = compute_f_h(ζ)
        return L >= 0 ? -a * ζ / (2 * Pr) :
               2 * log((1 + f_h) / 2) + 2 * (1 - f_h) / (9 * ζ)
    end
end

"""
    compute_MO_len(u, θ, flux)

Computes Monin-Obukhov length. Eq. 3 Ref. Nishizawa2018
"""
compute_MO_len(u, θ, flux) = -u^3 * θ / (k_Karman * grav * flux)

"""
    compute_friction_velocity(u_ave, θ, flux, Δz, z_0, a, Ψ_m_tol, tol_abs, iter_max)

Computes friction velocity, in Eq. 12 in
Ref. Nishizawa2018, by solving the
non-linear equation:

`u_ave = ustar/κ * ( ln(Δz/z_0) - Ψ_m(Δz/L) + z_0/Δz * Ψ_m(z_0/L) + R_z0 [ψ_m(z_0/L) - 1] )`

where `L` is a non-linear function of `ustar` (see `compute_MO_len`).
"""
function compute_friction_velocity(
    u_ave,
    θ,
    flux,
    Δz,
    z_0,
    a,
    Ψ_m_tol,
    tol_abs,
    iter_max,
)
    ustar_0 = u_ave * k_Karman / log(Δz / z_0)
    ustar = ustar_0
    let u_ave = u_ave,
        θ = θ,
        flux = flux,
        Δz = Δz,
        z_0 = z_0,
        a = a,
        Ψ_m_tol = Ψ_m_tol,
        tol_abs = tol_abs,
        iter_max = iter_max
        # Note the lowercase psi (ψ) and uppercase psi (Ψ):
        Ψ_m_closure(ζ, L) = compute_Ψ_m(ζ, L, a, Ψ_m_tol)
        ψ_m_closure(ζ, L) = compute_ψ_m(ζ, L, a)
        function compute_u_ave_over_ustar(u)
            L = compute_MO_len(u, θ, flux)
            R_z0 = compute_R_z0(z_0, Δz)
            temp1 = log(Δz / z_0)
            temp2 = -Ψ_m_closure(Δz / L, L)
            temp3 = z_0 / Δz * Ψ_m_closure(z_0 / L, L)
            temp4 = R_z0 * (ψ_m_closure(z_0 / L, L) - 1)
            return (temp1 + temp2 + temp3 + temp4) / k_Karman
        end
        compute_ustar(u) = u_ave / compute_u_ave_over_ustar(u)
        ustar_1 = compute_ustar(ustar_0)
        sol = RootSolvers.find_zero(
            u -> u_ave - u * compute_u_ave_over_ustar(u),
            ustar_0,
            ustar_1,
            SecantMethod(),
            CompactSolution(),
            tol_abs,
            iter_max,
        )
        ustar = sol.root
    end
    return ustar
end

"""
    compute_exchange_coefficients(z, F_m, F_h, a, u_star, θ, flux, Pr)

Computes exchange transfer coefficients:

  - K_D  momentum exchange coefficient
  - K_H  thermodynamic exchange coefficient
  - L_mo Monin-Obukhov length
"""
function compute_exchange_coefficients(z, F_m, F_h, a, u_star, θ, flux, Pr)

    L_mo = compute_MO_len(u_star, θ, flux)
    ψ_m = compute_ψ_m(z / L_mo, L_mo, a)
    ψ_h = compute_ψ_h(z / L_mo, L_mo, a, Pr)

    K_m = -F_m * k_Karman * z / (u_star * ψ_m) # Eq. 19 in Ref. Nishizawa2018
    K_h = -F_h * k_Karman * z / (Pr * θ * ψ_h) # Eq. 20 in Ref. Nishizawa2018

    return K_m, K_h, L_mo
end

end # Nishizawa2018 module

end # SurfaceFluxes module

max_stars_repo_name: trontrytel/CLIMA

max_issues_repo_name: trontrytel/CLIMA

max_forks_repo_name: trontrytel/CLIMA


text:
gc_ctr(lat::Array{Float64,1}, lon::Array{Float64,1}) =  (atan.(tan.(deg2rad.(lat)).*0.9933056), deg2rad.(lon))
gc_unwrap!(t::Array{Float64,1}) = (t[t .< 0] .+= (2.0*Float64(π)); return t)

@doc """
    G = gcdist(src, rec)

  Compute great circle distance, azimuth, and backazimuth from single source `s`
with coordinates `[s_lat, s_lon]` to receivers `r` with coordinates `[r_lat r_lon].`

  For a single source, pass `src` as a Float64 vector of the form `[s_lat, s_lon]`;
gcdist will return an Array{Float64,2} of the form

    [Δ₁   θ₁   β₁
     Δ₂   θ₂   β₂
     ⋮    ⋮    ⋮
     Δn   θn   βn]

for receivers `1:n`.

  For multiple sources, pass `src` as an Array{Float64,2} with each row
containing one (lat, lon) pair. This returns a three-dimensional matrix where
each two-dimensional slice takes the form

      [Δᵢ₁   θᵢ₁   βᵢ₁
       ⋮     ⋮     ⋮
       Δᵢn   θᵢn   βᵢn]

for source `i` at receivers `1:n`.
""" gcdist
function gcdist(src::Array{Float64,1}, rec::Array{Float64,2})
  N = size(rec, 1)
  lat_src = repeat([src[1]], N)
  lon_src = repeat([src[2]], N)
  lat_rec = rec[:,1]
  lon_rec = rec[:,2]

  ϕ1, λ1 = gc_ctr(lat_src, lon_src)
  ϕ2, λ2 = gc_ctr(lat_rec, lon_rec)
  Δϕ = ϕ2 - ϕ1
  Δλ = λ2 - λ1

  a = sin.(Δϕ/2.0) .* sin.(Δϕ/2.0) + cos.(ϕ1) .* cos.(ϕ2) .* sin.(Δλ/2.0) .* sin.(Δλ/2.0)
  Δ = 2.0 .* atan.(sqrt.(a), sqrt.(1.0 .- a))
  A = atan.(sin.(Δλ).*cos.(ϕ2), cos.(ϕ1).*sin.(ϕ2) - sin.(ϕ1).*cos.(ϕ2).*cos.(Δλ))
  B = atan.(-1.0.*sin.(Δλ).*cos.(ϕ1), cos.(ϕ2).*sin.(ϕ1) - sin.(ϕ2).*cos.(ϕ1).*cos.(Δλ))

  # convert to degrees
  return rad2deg.(hcat(Δ, gc_unwrap!(A), gc_unwrap!(B)))
end
gcdist(lat0::Float64, lon0::Float64, rec::Array{Float64,2}) = gcdist([lat0, lon0], rec)
gcdist(lat0::Float64, lon0::Float64, lat1::Float64, lon1::Float64) = gcdist([lat0, lon0], [lat1 lon1])

function gcdist(src::Array{Float64,2}, rec::Array{Float64,2})
  N_src = size(src,1)
  N_rec = size(rec,1)
  G = Array{Float64,3}(undef, N_rec, 3, N_src)
  for i = 1:N_src
    G[:,:,i] = gcdist(src[i,:], rec)
  end
  return G
end
gcdist(src::Array{Float64,2}, rec::Array{Float64,1}) = gcdist(src, [rec[1] rec[2]])


text:
include("battlesnake.jl")
include("algorithms.jl")
include("minimax.jl")
include("maxn.jl")
include("agent.jl")
using Joseki, JSON, HTTP

### Create some endpoints

function index(req::HTTP.Request)
    response_dict = Dict("name" => "solide julia", "head" => "shades", "tail" => "bolt", "apiversion" => "1")
    simple_json_responder(req, response_dict)
end

function start_req(req::HTTP.Request)
    req.response.status = 200
    return req.response
end

function end_req(req::HTTP.Request)
    req.response.status = 200
    return req.response
end

function move_req(req::HTTP.Request)
    j = try
        body_as_dict(req)
    catch err
        return error_responder(req, "I was expecting a json request body!")
    end
    # get gamemode later when we have implemented maxN
    simulation_time = j["game"]["timeout"]
    board = j["board"]
    height = board["height"]
    width = board["width"]

    food_json = board["food"]
    food_pos = Position[]
    for f in food_json
        push!(food_pos, Position(f["x"] + 1, f["y"] + 1))
    end
    food = Food(food_pos)

    you = j["you"]
    snek_body = Position[]
    for part in you["body"]
        pos = Position(part["x"] + 1, part["y"] + 1)
        # if !(pos in snek_body)
        push!(snek_body, pos)
        # end
    end
    snek = Snake(snek_body, false, you["health"])

    snakes = Snake[]
    push!(snakes, snek)
    json_snakes = board["snakes"]
    for sn in json_snakes
        if sn["id"] == you["id"]
            continue
        end
        body = Position[]
        for part in sn["body"]
            pos = Position(part["x"] + 1, part["y"] + 1)
            #if !(pos in body)
            push!(body, pos)
            #end
        end
        snake = Snake(deepcopy(body), false, sn["health"])
        push!(snakes, snake)
    end

    board = Boardstate(snakes, food, height, width)
    action = move(board, 1, 0.6 * simulation_time)

    response_dict = Dict("move" => dir2str[action])
    simple_json_responder(req, response_dict)
end

dir2str = Dict(
    up => "up",
    down => "down",
    right => "right",
    left => "left"
)

### Create and run the server

# Make a router and add routes for our endpoints.
endpoints = [
    (index, "GET", "/"),
    (index, "GET", "/index"),
    (start_req, "POST", "/start"),
    (move_req, "POST", "/move"),
    (end_req, "POST", "/end"),

]
r = Joseki.router(endpoints)

# Fire up the server
HTTP.serve(r, "0.0.0.0", 8001, verbose=false)

max_stars_repo_name: rhotertj/solide-schlange

max_issues_repo_name: rhotertj/solide-schlange

max_forks_repo_name: rhotertj/solide-schlange


text:
using DataFrames

include("abstractfeatures.jl")

"""
    SpeciesFeatureDescriptor(feature_name, encode_f, decode_f, categorical, contextual, length, encodable_elements)

Construct a feature object that encodes features associated with individual atoms that depend upon their local environment in some way (if your feature is defined only by elemental identity, you should use ElementFeatureDescriptor!)

## Arguments
- `name::String`: the name of the feature
- `encode_f::Function`: a function that takes in <:AbstractAtoms and returns encoded values of this feature for the atoms in that structure
- `decode_f::Function`: inverse function to `encode_f`, takes in encoded feature and returns value (for categorical) or range of values (for continuous-valued) of the feature
- `categorical::Bool`: flag for whether the feature is categorical or continuous-valued
- `length::Int`: length of encoded vector
- `encodable_elements::Vector{String}`: list of elements (by symbol) that can be encoded by this feature
"""
struct SpeciesFeatureDescriptor <: AbstractAtomFeatureDescriptor
    name::String
    length::Integer
    encoder_decoder::AbstractCodec
    categorical::Bool
    encodable_elements::Vector{String}
end

# pretty printing, short version
Base.show(io::IO, af::SpeciesFeatureDescriptor) = print(io, "AtomFeature $(af.name)")

# pretty printing, long version
function Base.show(io::IO, ::MIME"text/plain", af::SpeciesFeatureDescriptor)
    st = "AtomFeature $(af.name):\n   categorical: $(af.categorical)\n   encoded length: $(af.length)"
    print(io, st)
end

# TODO: add way to get range/list of possible values for feature...
encodable_elements(f::SpeciesFeatureDescriptor) = f.encodable_elements

function (f::SpeciesFeatureDescriptor)(a::AbstractAtoms)
    @assert all([el in f.encodable_elements for el in a.elements]) "Feature $(f.name) cannot encode some element(s) in this structure!"
    f.encode_f(a)
end

# TODO: some Weave stuff needed here?


text:
@testset "mock keyword" begin
    f(; x=1) = x

    function f_param()
        x = 2
        return @mock f(; x=x)
    end

    @test f_param() == 2

    function f_kw()
        x = 3
        return @mock f(x=x)
    end

    @test f_kw() == 3

    @static if VERSION >= v"1.5"
        function f_req_param()
            x = 4
            return @mock f(; x)  # Testing keyword parameter that is just a `Symbol`
        end

        @test f_req_param() == 4
    end
end


text:
export contour

"""
    contour(f; options...)

Compute the contour 𝐶(𝑓) of the amoeba ``\\mathcal{A}(f)``.

## Example
```julia
@polyvar x y

contour(x^2+y^2+1)

# custom domain
contour(x^2+y^2+1, domain=(-5, 5, -5, 5))
```

## Options
* `domain`: A tuple in the form `(xmin, xmax, ymin, ymax)` which defines a section Ω
for which the contour 𝐶(𝑓) is computed.
* `membership_options`: The options for the membership test
* `res=(600,600)`: The resolution with which starting points are sampled
* `samples_off_axis=2*MP.maxdegree(p)^2`: The number of sample points per off-axis.
"""
function contour(f::MP.AbstractPolynomial;
    domain=amoeba_carcase_domain_heuristic(f, factor=2.0),
    res::NTuple{2, Int}=(600, 600),
    samples_off_axis::Int=2*MP.maxdegree(f),
    membership_options=MembershipTestOptions(ntries=10, maxiters=30, tol=1e-8))

    grid = Grid2D(domain..., res...)

    _contour2D(ContourFiber2D(f), grid, samples_off_axis, membership_options)
end

# only internally used
"""
    grid_contour(F::ContourFiber2D, grid::Grid2D; kwargs...)

Instead of returning a list of contour points it fills a bitmap with the nearest contour
points.
"""
function grid_contour(p::MP.AbstractPolynomial; res=600, kwargs...)
    grid_contour(p, Grid2D(p, res); kwargs...)
end

function grid_contour(p::MP.AbstractPolynomial, grid::Grid2D;
    every_main_axis::Int=10, samples_off_axis::Int=8, options=MembershipTestOptions(ntries=3, maxiters=40, tol=1e-8))

    fiber = ContourFiber2D(p)

    _contour2D_grid(fiber, grid, every_main_axis, samples_off_axis, options)
end


text:
@testset "Sanity Row/Array" begin
    @test distance_row([ ], [ ]) == 0

    @testset "Sanity Row/Array/Left" begin
        @test distance_row([ 1 ], [ ])     == 1
        @test distance_row([ 1 2 ], [ ])   == 2
        @test distance_row([ 1 2 3 ], [ ]) == 3
    end

    @testset "Sanity Row/Array/Right" begin
        @test distance_row([ ], [ 1 ])     == 1
        @test distance_row([ ], [ 1 2 ])   == 2
        @test distance_row([ ], [ 1 2 3 ]) == 3
    end

    @testset "Sanity Row/Array/Both" begin
        @test distance_row([ 1 ], [ 1 ])         == 0
        @test distance_row([ 1 ], [ 2 ])         == 1
        @test distance_row([ 1 2 ], [ 2 2 ])     == 1
        @test distance_row([ 1 2 3 ], [ 2 2 3 ]) == 1
        @test distance_row([ 1 2 3 ], [ 1 2 2 ]) == 1
    end
end

max_stars_repo_name: jlegare/LevenshteinToolkit

max_issues_repo_name: jlegare/LevenshteinToolkit

max_forks_repo_name: jlegare/LevenshteinToolkit


text:
using JOLI

export helmholtz_system, ComputationalGrid, odn_to_grid

mutable struct ComputationalGrid{I<:Integer,F<:AbstractFloat}
    phys_to_comp_grid::joAbstractOperator
    comp_to_phys_grid::joAbstractOperator

    comp_o::AbstractArray{F,1}
    comp_d::AbstractArray{F,1}
    comp_n::AbstractArray{I,1}
    comp_n_nopml::AbstractArray{I,1}
end

function odn_to_grid(comp_grid::ComputationalGrid{I,F}) where {F<:AbstractFloat,I<:Integer}
    return odn_to_grid(comp_grid.comp_o,comp_grid.comp_d,comp_grid.comp_n)
end

function helmholtz_system(v::AbstractArray{F,1},model::Model{I,F},freq::Union{F,Complex{F}},opts::PDEopts{I,F}) where {I<:Integer,F<:AbstractFloat}

    ndims = (length(opts.comp_n)==2 || opts.comp_n[3]==1) ? 2 : 3
    lsopts = deepcopy(opts.lsopts)
    if ndims==2
        opts.comp_n = opts.comp_n[1:2]
        opts.comp_d = opts.comp_d[1:2]
        opts.comp_o = opts.comp_o[1:2]
    end
    v = reshape(v,tuple(opts.comp_n...))
    dt = opts.comp_d
    if model.unit == "m/s"
        vmax = maximum(vec(v))
    elseif model.unit == "s2/m2"
        vmax = maximum(vec(v).^(-1/2))
    elseif model.unit == "s2/km2"
        vmax = maximum(1e3*(vec(v).^(-1/2)))
    else
        throw(ArgumentError())
    end
    npml = convert(Array{I,1},ceil.(vmax./(real(freq)*dt)))
    npml = repmat(npml',2,1)
    npml = min.(npml,opts.npml)
    # npml is a 2x3 matrix corresponding to
    # in 2D
    # [ #pml pts lower z  | # pml pts lower x  | 0 ]
    # [ #pml pts upper z  | # pml pts upper x  | 0 ]

    # in 3D
    # [ #pml pts lower x  | # pml pts lower y  | # pml pts lower z ]
    # [ #pml pts upper x  | # pml pts upper y  | # pml pts upper z ]

    nt_nopml = opts.comp_n
    nt_pml = nt_nopml + sum(npml,1)'
    nt_pml = vec(nt_pml)
    ot_pml = opts.comp_o - vec(npml[1,:]) .* vec(dt)
    (Pext,Ppad) = get_pad_ext_ops(nt_nopml,npml,ndims)
    phys_to_comp = Pext
    if opts.cut_pml
        comp_to_phys = Ppad'
    else
        comp_to_phys = Pext'
    end
    v_pml = reshape(Pext*vec(v),tuple(nt_pml...))
    comp_grid = ComputationalGrid{I,F}(phys_to_comp,comp_to_phys,vec(ot_pml),vec(dt),vec(nt_pml),vec(nt_nopml))
    N_system = prod(nt_pml)

    # Set up system matrix
    if ndims==2
        if opts.pde_scheme==helm2d_chen9p
            (H,dH,ddH) = helm2d_chen2013(nt_pml,dt,npml,freq,v_pml,model.f0,model.unit)
        elseif opts.pde_scheme==helm2d_std7
            (H,dH,ddH) = helm2d_std7(nt_pml,dt,npml,freq,v_pml,model.f0,model.unit)
        end
        if lsopts.solver==:lufact
            opH = joInvertibleMatrix(H)
        else
            error("Unimplemented method $(opts.lsopts.solver) for 2D")
        end
        P = :identity
    elseif ndims==3
        if opts.pde_scheme==helm3d_operto27
            (wn,dwn,ddwn) = param_to_wavenum(v_pml,freq,model.unit)
            if !opts.implicit_matrix
                H = helm3d_operto_matrix(wn,dt,nt_pml,freq,npml)
            else
                Hmvp = (x;forw=true)->helm3d_operto_mvp_mt(wn,dt,nt_pml,freq,npml,reshape(x,nt_pml...),forw_mode=forw)
                H = joLinearFunctionFwdCT(N_system,N_system,x->Hmvp(x,forw=true),x->Hmvp(x,forw=false),Complex{F},Complex{F})
                dHmvp = (x;forw=true)->helm3d_operto_mvp(dwn,dt,nt_pml,freq,npml,reshape(x,nt_pml...),forw_mode=forw,deriv_mode=true)
                dH = joLinearFunctionFwdCT(N_system,N_system,x->dHmvp(x,forw=true),x->dHmvp(x,forw=false),Complex{F},Complex{F})
                ddHmvp = (x;forw=true)->helm3d_operto_mvp(ddwn,dt,nt_pml,freq,npml,reshape(x,nt_pml...),forw_mode=forw,deriv_mode=true)
                ddH = joLinearFunctionFwdCT(N_system,N_system,x->ddHmvp(x,forw=true),x->ddHmvp(x,forw=false),Complex{F},Complex{F})
            end

        elseif opts.pde_scheme==helm3d_std9

        end
        if lsopts.precond==:mlgmres
            P = MLGMRES(H,vec(v),comp_grid,model,freq,opts)
        elseif lsopts.precond==:vgmres
            P = VGMRES(H,vec(v),comp_grid,model,freq,opts)
        else
            P = :identity
        end
        lsopts.precond = P
        if lsopts.solver==:lufact
            opH = joInvertibleMatrix(H)
        else
            if opts.implicit_matrix
                  opH = joLinearFunctionCT(N_system,N_system,
                                         x->Hmvp(x,forw=true),
                                         x->Hmvp(x,forw=false),
                                         x->linearsolve(H,x,[],lsopts,forw_mode=true),
                                         x->linearsolve(H,x,[],lsopts,forw_mode=false),
                                         Complex{F},Complex{F})
            else
                opH = H
            end
        end
    end

    T = u-> joLinearFunctionFwdCT( N_system,N_system,
                                  dm->dH*(dm.*u),
                                  z->real(conj(u).*(dH'*z)),F,Complex{F},fMVok=true)
    DTadj = (u,dm,du)->joLinearFunctionFwdT(prod(nt_pml),prod(nt_pml),
                                            z->real(conj(u).*(dm.*(ddH'*z)) + conj(du).*(dH'*z)),
                                            @joNF,F,Complex{F},fMVok=true)

    return (opH,comp_grid,T,DTadj,P)
end

function get_pad_ext_ops(nt_nopml,npml,ndims)
    if ndims == 2
        Pext = joKron(joExtend(nt_nopml[2],:border,pad_lower=npml[1,2],pad_upper=npml[2,2]),
                      joExtend(nt_nopml[1],:border,pad_lower=npml[1,1],pad_upper=npml[2,1]))
        Ppad = joKron(joExtend(nt_nopml[2],:zeros,pad_lower=npml[1,2],pad_upper=npml[2,2]),
                      joExtend(nt_nopml[1],:zeros,pad_lower=npml[1,1],pad_upper=npml[2,1]))
    else

        Pext = joKron(joExtend(nt_nopml[3],:border,pad_lower=npml[1,3],pad_upper=npml[2,3]),
                      joExtend(nt_nopml[2],:border,pad_lower=npml[1,2],pad_upper=npml[2,2]),
                      joExtend(nt_nopml[1],:border,pad_lower=npml[1,1],pad_upper=npml[2,1]))
        Ppad = joKron(joExtend(nt_nopml[3],:zeros,pad_lower=npml[1,3],pad_upper=npml[2,3],DDT=Complex{Float64}),
                      joExtend(nt_nopml[2],:zeros,pad_lower=npml[1,2],pad_upper=npml[2,2],DDT=Complex{Float64}),
                      joExtend(nt_nopml[1],:zeros,pad_lower=npml[1,1],pad_upper=npml[2,1],DDT=Complex{Float64}))
    end
    return (Pext,Ppad)
end


text:
"""
    ThreadedScans.partitioned_hillis_steele!(op, xs::AbstractVector) -> xs

Compute inclusive scan.  Intermediate reductions for `ntasks` chunks are
computed in parallel and then merged using prefix sum algorithm by Hillis and
Steele.  It is marginally better than other methods especially when the number
of worker threads is large.

# Keyword arguments
* `ntasks = Threads.nthreads()`: number of tasks
"""
function ThreadedScans.partitioned_hillis_steele!(
    op,
    xs;
    ntasks::Integer = Threads.nthreads(),
    spin::Union{Nothing,Integer} = nothing,
)
    if length(xs) <= 2
        if length(xs) == 2
            xs[begin+1] = op(xs[begin], xs[begin+1])
        end
        return xs
    end
    if ntasks < 2
        scan!(op, xs)
        return xs
    end
    nchunks = min(length(xs), ntasks + 1)
    ntasks = nchunks - 1
    @assert ntasks > 1

    npromises = ntasks + foldl_hillis_steele!((_, x) -> x.stop, nothing, ntasks)::Int
    promises = OneWayPromise{eltype(xs)}[OneWayPromise{eltype(xs)}() for _ in 1:npromises]

    chunks = partitionto(xs, ntasks + 1)
    @assert length(chunks) == ntasks + 1
    onerror() = foreach(close, promises)
    spawn_foreach(1:ntasks; onerror = onerror) do i
        _partitioned_hillis_steele!(op, chunks, i, ntasks, promises, spin)
    end

    return xs
end

function foldl_hillis_steele!(rf, acc, ntasks::Integer)
    stop = 0
    nright = 1   # number of non-take!-ing tasks
    while nright < ntasks
        offset = stop
        stop = stop + ntasks - nright
        # `rf` manipulates `promises[offset+1:stop]`:
        acc = rf(acc, (pow2 = nright, offset = offset, stop = stop))
        nright *= 2
    end
    return acc
end

function _partitioned_hillis_steele!(op, chunks, i, ntasks, promises, spin)
    if i == 1
        accr = scan!(op, chunks[i])
    else
        accr = foldl(op, chunks[i])
    end
    put!(promises[i], accr)

    # Hillis-Steele prefix scan for the "left lane":
    il = if i == 1
        ntasks  # == length(chunks) - 1
    else
        i - 1
    end

    accl = foldl_hillis_steele!(take!(promises[il]), ntasks) do accl, x
        if il <= ntasks - x.pow2
            put!(promises[ntasks+x.offset+il], accl)
        end
        j = il - x.pow2
        if j > 0
            accl = op(accl, take!(promises[ntasks+x.offset+j], spin))
        end
        accl
    end

    scan!(op, accl, chunks[il+1])
    return
end


text:
using Pkg, Random
Pkg.activate(".")

module EnzymeMFABC
include("../example.jl")
end

using FileIO, JLD2
using Statistics

Random.seed!(333)
MFBSL_output = EnzymeMFABC.do_mfbsl.([8_000, 16_000, 32_000, 64_000];
    N₀ = 2000,
    δ = 1e8,
    N_batch = 1,
    GD_batch = 100,
)

c = map(MFBSL_output) do (S, μ)
	sum(S.c_lo .+ S.Σc_hi)
end
v = map(MFBSL_output) do (S, μ)
    Ḡ = mean(EnzymeMFABC.G, S)
	w_mf = S.L_lo .+ (S.ΣL_mf./S.μ)
	sum((w_mf .* (EnzymeMFABC.G.(S.θ) .- Ḡ)).^2) / sum(w_mf)^2
end

S, μ, μ_t, ν_t = MFBSL_output[end]

interesting = findall(p->(!iszero(p.ΣL_mf2)), S)
w_mf = S.L_lo .+ (S.ΣL_mf ./ S.μ)
w_mf_sparse = hcat(interesting, w_mf[interesting])

save("MFBSL_output.jld2", Dict(
    "c" => c,
    "v" => v,
    "w_mf" => w_mf_sparse,
    "μ" => μ,
    "μ_t" => μ_t[100:100:end,:],
    "ν_t" => ν_t[100:100:end,:],
))

max_stars_repo_name: tpprescott/mf-lf

max_issues_repo_name: tpprescott/mf-lf

max_forks_repo_name: tpprescott/mf-lf


text:
using ScatterNNlib
using StaticArrays: @MMatrix, @MArray
using Zygote
using CUDA
using Test

cuda_tests = [
    "cuda/cuarray",
    "cuda/grad",
]

tests = [
    "scatter/densearray",
    "scatter/staticarray",
    "grad",
    "gather",
]

if CUDA.functional()
    append!(tests, cuda_tests)
else
    @warn "CUDA unavailable, not testing CUDA support"
end

@testset "ScatterNNlib.jl" begin
    for t in tests
        include("$(t).jl")
    end
end


text:
"compute the kl divergence with mask where p is already the log(p)"
logkldivergence(q::Abstract3DTensor{T},
                logp::Abstract3DTensor{T},
                mask) where T =
                    sum(reshape(sum(sum(q .* (log.(q .+ eps(q[1])) .- logp); dims=1) .* mask; dims=2), :) ./ reshape(sum(mask; dims=2), :))

function logkldivergence(q, logp, mask)
    kld = (q .* (log.(q .+ eps(q[1])) .- logp)) #handle gpu broadcast error
    sum(kld .* mask) / sum(mask)
end

"compute the cross entropy with mask where p is already the log(p)"
logcrossentropy(q::Abstract3DTensor{T},
                logp::Abstract3DTensor{T},
                mask) where T =
                    -sum(reshape(sum(sum(q .* logp; dims=1) .* mask; dims=2), :) ./ reshape(sum(mask; dims=2), :))

function logcrossentropy(q, logp, mask)
    ce = q .* logp
    -sum(ce .* mask) / sum(mask)
end

"compute the kl divergence where p is already the log(p)"
logkldivergence(q::Abstract3DTensor{T},
                logp::Abstract3DTensor{T}) where T =
                    sum(reshape(sum(sum(q .* (log.(q .+ eps(q[1])) .- logp); dims=1); dims=2), :) ./ size(q, 2))

function logkldivergence(q, logp)
    kld = (q .* (log.(q .+ eps(q[1])) .- logp)) #handle gpu broadcast error
    sum(kld) / size(kld, 2)
end

"compute the cross entropy where p is already the log(p)"
logcrossentropy(q::Abstract3DTensor{T},
                logp::Abstract3DTensor{T}) where T =
                    -sum(reshape(sum(sum(q .* logp; dims=1); dims=2), :) ./ size(q, 2))

function logcrossentropy(q, logp)
    ce = q .* logp
    -sum(ce) / size(ce, 2)
end


text:
const IDLIST = (
# SAFE ARGUMENT CHECKS
    ("fast_check1", false),
    ("fast_check2", false),
    ("fast_check3", false),

# FOR-LOOP MACROS
    # x = op(a, x)
    ("scalarr1_for", true),
    ("scalarr1_for_inc1", true),
    # x = op(a, y)
    ("scalarr1_foroop", true),
    ("scalarr1_foroop_inceq", true),
    ("scalarr1_foroop_inc1", true),
    ("scalarr1_foroop_inc1ieq", true),
    # x = op(x, y)
    ("arr2xy_for", true),
    ("arr2xy_for_inceq", true),
    ("arr2xy_for_inc1", true),
    ("arr2xy_for_inc1ieq", true),
    # x = op(y, z)
    ("arr2yz_foroop", true),
    ("arr2yz_foroop_inceq", true),
    ("arr2yz_foroop_inc1", true),
    ("arr2yz_foroop_inc1ieq", true),
    # x = x + a*y
    ("addarrscal_for", false),
    ("addarrscal_for_inceq", false),
    ("addarrscal_for_inc1", false),
    ("addarrscal_for_inc1ieq", false),
    # x = y + a*z
    ("addarrscal_foroop", false),
    ("addarrscal_foroop_inceq", false),
    ("addarrscal_foroop_inc1", false),
    ("addarrscal_foroop_inc1ieq", false),
    # x = y
    ("copy_foroop", false),
    ("copy_foroop_inceq", false),
    ("copy_foroop_inc1", false),
    ("copy_foroop_inc1ieq", false),
    # x = a
    ("fill_foroop", false),
    ("fill_foroop_inc1", false),

# C LIBRARY MACROS
    # x = y
    ("c_memcpy", false),
    # x = a
    ("c_memset", false),

# BLAS MACROS
    # x = a*x
    ("blas_scale", false),
    # x = y
    ("blas_copy", false),
    # x = x + a*y
    ("blas_axpy", false),

# SET MACROS
    ("set1_inc1", false),
    ("set2_inceq", false),
    ("set2_inc1", false),
    ("set2_inc1ieq", false),
    ("set3_inceq", false),
    ("set3_inc1", false),
    ("set3_inc1ieq", false),
)

# name convention, binary operations:
# op(a, x) = scalarr1, op(x, a) = arr1scal
# op(x, y) = arr2xy, op(y, x) = arr2yx
# foroop for same variable on both sides of =

# "compile" the source srcfile to destfile
function makeFAO()
    destfile = "FastArrayOps.jl"   # warning: gets overwritten
    srcfile = "FastArrayOps_src.jl"
    dir = "templates/"
    opid = "\$op"
    preid = "##!"
    postid = "\n"
    fext = ".jl"
    oplist = ("*","/","+","-")

    file = open(destfile, "w")
    f = readall(srcfile)

    for (id, opstr) in IDLIST
        idfs = readall(string(dir, id, fext))
        idfs = string("# src: ", string(id, fext), "\n", idfs)
        if opstr
            for op in oplist
                idfsop = replace(idfs, opid, op)
                idstr = string(preid, id, "#", op, postid)
                if searchindex(f, idstr) != 0
                    f = replace(f, idstr, idfsop)
                end
            end
        else
            f = replace(f, string(preid, id, postid), idfs)
        end
    end

    idmiss = searchindex(f, preid)
    idmiss != 0 && error("missing ids found : ", f[idmiss:idmiss+30])

    write(file, f)
    close(file)
end

makeFAO()


text:
@testset "Spatial Coordinates Tests" begin
    @testset "Indexing and Basics" begin
        subsel = C[Coord(Lat(0..8))]
        @test hasdim(subsel, Coord)
        @test size(subsel, Coord) < size(C, Coord)
        @test all(0 ≤ c[2] ≤ 8 for c in gnv(dims(subsel, Coord)))

        @test length(collect(spatialidxs(C))) == size(C, Coord)

        Atran = transform_to_coord(A)
        @test hasdim(Atran, Coord)
        @test lons == sort!(unique!([c[1] for c in gnv(dims(Atran, Coord))]))

        # subsel2 = C[Coord(Lat(8..0))] # DimensionalData.jl is not allowing this to work
        # @test subsel2 == subsel
    end

    @testset "Zonal mean" begin
        zC = zonalmean(C)
        @test hasdim(zC, Lat)
        @test gnv(dims(zC, Lat)) == lats
        @test maximum(zC) < 1 # because we aggregate over cosine Lon, the result is small
        # TODO: test with weights
    end

    @testset "Hemispheric stuff" begin
        # TODO: Would be useful to test this hemispheric stuff with coordinates
        # that include 0 latitude and that do not cover symmetric latitdues
        # even though it is super uncommon
        nhi, shi = ClimateBase.hemisphere_indices(gnv(dims(C, Coord)))
        @test length(nhi) == length(shi)
        @test sort!(vcat(nhi, shi)) == 1:size(C, Coord)

        Cnh, Csh = hemispheric_functions(C) # symmetric, non-zero latitudes
        @test hasdim(Cnh, Coord) && hasdim(Csh, Coord)
        @test size(Cnh, Coord) == size(Csh, Coord)

        # hemispheric_functions makes latitudes the same by convention:
        @test gnv(dims(Cnh, Coord)) == gnv(dims(Csh, Coord))
        @test minimum(uniquelats(Cnh)[2]) ≥ 0
        @test minimum(uniquelats(Csh)[2]) ≥ 0

        Cnh, Csh = hemispheric_means(C)
        @test !hasdim(Cnh, Coord) && !hasdim(Csh, Coord)
        @test hasdim(Cnh, Time)
        # TODO: test actual numeric values
    end

end


text:
using OptimalTransport

using Distances
using ForwardDiff
using LogExpFunctions
using PythonOT: PythonOT

using LinearAlgebra
using Random
using Test

const POT = PythonOT

Random.seed!(100)

@testset "sinkhorn_stabilized.jl" begin
    # size of source and target
    M = 250
    N = 200

    # create two random histograms
    μ = normalize!(rand(M), 1)
    ν = normalize!(rand(N), 1)

    # create random cost matrix
    C = pairwise(SqEuclidean(), rand(1, M), rand(1, N); dims=2)

    # regularization parameter
    ε = 0.01

    @testset "example" begin
        # compute optimal transport plan and optimal transport cost
        γ = sinkhorn(μ, ν, C, ε, SinkhornStabilized(); maxiter=5_000, rtol=1e-9)
        c = sinkhorn2(μ, ν, C, ε, SinkhornStabilized(); maxiter=5_000, rtol=1e-9)

        # check that plan and cost are consistent
        @test c ≈ dot(γ, C)

        # compare with POT
        γ_pot = POT.sinkhorn(
            μ, ν, C, ε; method="sinkhorn_stabilized", numItermax=5_000, stopThr=1e-9
        )
        c_pot = POT.sinkhorn2(
            μ, ν, C, ε; method="sinkhorn_stabilized", numItermax=5_000, stopThr=1e-9
        )[1]
        @test γ_pot ≈ γ rtol = 1e-6
        @test c_pot ≈ c rtol = 1e-7

        # compute optimal transport cost with regularization term
        c_w_regularization = sinkhorn2(
            μ, ν, C, ε, SinkhornStabilized(); maxiter=5_000, regularization=true
        )
        @test c_w_regularization ≈ c + ε * sum(x -> iszero(x) ? x : x * log(x), γ)

        # ensure that provided plan is used and correct
        c2 = sinkhorn2(similar(μ), similar(ν), C, rand(), SinkhornStabilized(); plan=γ)
        @test c2 ≈ c
        c2_w_regularization = sinkhorn2(
            similar(μ), similar(ν), C, ε, SinkhornStabilized(); plan=γ, regularization=true
        )
        @test c2_w_regularization ≈ c_w_regularization

        # batches of histograms
        d = 10
        for (size2_μ, size2_ν) in
            (((), (d,)), ((1,), (d,)), ((d,), ()), ((d,), (1,)), ((d,), (d,)))
            # generate batches of histograms
            μ_batch = repeat(μ, 1, size2_μ...)
            ν_batch = repeat(ν, 1, size2_ν...)

            # compute optimal transport plan and check that it is consistent with the
            # plan for individual histograms
            γ_all = sinkhorn(
                μ_batch, ν_batch, C, ε, SinkhornStabilized(); maxiter=5_000, rtol=1e-9
            )
            @test size(γ_all) == (M, N, d)
            @test all(view(γ_all, :, :, i) ≈ γ for i in axes(γ_all, 3))

            # compute optimal transport cost and check that it is consistent with the
            # cost for individual histograms
            c_all = sinkhorn2(
                μ_batch, ν_batch, C, ε, SinkhornStabilized(); maxiter=5_000, rtol=1e-9
            )
            @test size(c_all) == (d,)
            @test all(x ≈ c for x in c_all)
        end
    end

    @testset "infinite absorption tolerance" begin
        # compute optimal transport plan with infinite absorption tolerance
        alg = SinkhornStabilized(; absorb_tol=Inf)
        γ = sinkhorn(μ, ν, C, ε, alg; maxiter=100)
        c = sinkhorn(μ, ν, C, ε, alg; maxiter=100)

        # compare with regular Sinkhorn algorithm
        γ_sinkhorn = sinkhorn(μ, ν, C, ε; maxiter=100)
        c_sinkhorn = sinkhorn(μ, ν, C, ε; maxiter=100)
        @test γ_sinkhorn ≈ γ
        @test c_sinkhorn ≈ c
    end

    @testset "AD" begin
        # compute gradients with respect to source and target marginals separately and
        # together
        ForwardDiff.gradient(zeros(N)) do xs
            return sinkhorn2(μ, softmax(xs), C, ε, SinkhornStabilized())
        end
        ForwardDiff.gradient(zeros(M)) do xs
            return sinkhorn2(softmax(xs), ν, C, ε, SinkhornStabilized())
        end
        ForwardDiff.gradient(zeros(M + N)) do xs
            return sinkhorn2(
                softmax(xs[1:M]), softmax(xs[(M + 1):end]), C, ε, SinkhornStabilized()
            )
        end
    end

    @testset "deprecations" begin
        γ = sinkhorn(μ, ν, C, ε, SinkhornStabilized(; absorb_tol=10))
        @test (@test_deprecated sinkhorn_stabilized(μ, ν, C, ε; absorb_tol=10)) == γ

        γ = sinkhorn(μ, ν, C, ε, SinkhornStabilized(); atol=1e-6)
        @test (@test_deprecated sinkhorn_stabilized(μ, ν, C, ε; tol=1e-6)) == γ

        γ = sinkhorn(μ, ν, C, ε, SinkhornStabilized())
        u, v = @test_deprecated sinkhorn_stabilized(μ, ν, C, ε; return_duals=true)
        @test exp.(-(C .- u .- v') / ε) ≈ γ
    end
end


text:
using DiffEqCallbacks
using Test

# write your own tests here
@time begin
@time @testset "AutoAbstol" begin include("autoabstol_tests.jl") end
@time @testset "TerminateSteadyState tests" begin include("terminatesteadystate_test.jl") end
@time @testset "StepsizeLimiter tests" begin include("stepsizelimiter_tests.jl") end
@time @testset "Function Calling tests" begin include("funccall_tests.jl") end
@time @testset "Saving tests" begin include("saving_tests.jl") end
@time @testset "PresetTime tests" begin include("preset_time.jl") end
@time @testset "Iterative tests" begin include("iterative_tests.jl") end
@time @testset "Periodic tests" begin include("periodic_tests.jl") end
@time @testset "Manifold tests" begin include("manifold_tests.jl") end
@time @testset "Domain tests" begin include("domain_tests.jl") end
end


text:
using Plots

# Model
include_model("biped")

model = free_time_model(model)

# Visualize
# - Pkg.add any external deps from visualize.jl
include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)

function get_q⁺(x)
	view(x, model.nq .+ (1:model.nq))
end

function ellipse_traj(x_start, x_goal, z, T)
	dist = x_goal - x_start
	a = 0.5 * dist
	b = z

	z̄ = 0.0

	x = range(x_start, stop = x_goal, length = T)

	z = sqrt.(max.(0.0, (b^2) * (1.0 .- ((x .- (x_start + a)).^2.0) / (a^2.0))))

	return x, z
end

# Horizon
T = 31
Tm = 31 #convert(Int, floor(0.5 * T))

# Time step
tf = 0.5 * 1.25
h = tf / (T - 1)

# Configurations
# 1: x pos
# 2: z pos
# 3: torso angle (rel. to downward vertical)
# 4: thigh 1 angle (rel. to downward vertical)
# 5: calf 1 (rel. to downward vertical)
# 6: thigh 2 (rel. to downward vertical)
# 7: calf 2 (rel. to downward vertical)

perm = @SMatrix [1.0 0.0 0.0 0.0 0.0 0.0 0.0;
				 0.0 1.0 0.0 0.0 0.0 0.0 0.0;
				 0.0 0.0 1.0 0.0 0.0 0.0 0.0;
			     0.0 0.0 0.0 0.0 0.0 1.0 0.0;
			     0.0 0.0 0.0 0.0 0.0 0.0 1.0;
			     0.0 0.0 0.0 1.0 0.0 0.0 0.0;
			     0.0 0.0 0.0 0.0 1.0 0.0 0.0]

function initial_configuration(model, θ_torso, θ_thigh_1, θ_leg_1, θ_thigh_2)
    q1 = zeros(model.nq)
    q1[3] = θ_torso
    q1[4] = θ_thigh_1
    q1[5] = θ_leg_1
    z1 = model.l_thigh1 * cos(q1[4]) + model.l_calf1 * cos(q1[5])
    q1[6] = θ_thigh_2
    q1[7] = -1.0 * acos((z1 - model.l_thigh2 * cos(q1[6])) / model.l_calf2)
	# q1[6] = θ_thigh_1
	# q1[7] = θ_leg_1
	q1[2] = z1

    # p1 = kinematics_2(model, q1, body = :calf_1, mode = :ee)
    # p2 = kinematics_2(model, q1, body = :calf_2, mode = :ee)
    # @show stride = abs(p1[1] - p2[1])
	#
    # q1[1] = -0.75#-1.0 * p1[1]
	#
    # qM = copy(q1)
    # qM[4] = q1[6]
    # qM[5] = q1[7]
    # qM[6] = q1[4]
    # qM[7] = q1[5]
    # qM[1] = abs(p2[1])
	#
    # pM_1 = kinematics_2(model, qM, body = :calf_1, mode = :ee)
    # pM_2 = kinematics_2(model, qM, body = :calf_2, mode = :ee)
	#
    # qT = copy(q1)
    # qT[1] = 0.75#q1[1] + 1.0#2 * stride
	#
    # pT_1 = kinematics_2(model, qT, body = :calf_1, mode = :ee)
    # pT_2 = kinematics_2(model, qT, body = :calf_2, mode = :ee)

    return q1
end

# q1 = initial_configuration(model, 0.0, pi / 75.0, 0.0, -pi / 100.0)
q1 = initial_configuration(model, 0.0, pi / 25.0, 0.0, -pi / 100.0)
pf1 = kinematics_2(model, q1, body = :calf_1, mode = :ee)
pf2 = kinematics_2(model, q1, body = :calf_2, mode = :ee)

strd = (pf1 - pf2)[1]

qT = Array(perm * copy(q1))
qT[1] += strd
q_ref = linear_interpolation(q1, qT, T+1)
x_ref = configuration_to_state(q_ref)
visualize!(vis, model, q_ref, Δt = h)

# pt1 = kinematics_3(model, q1, body = :foot_1, mode = :toe)
# ph1 = kinematics_3(model, q1, body = :foot_1, mode = :heel)
#
# pt2 = kinematics_3(model, q1, body = :foot_2, mode = :toe)
# ph2 = kinematics_3(model, q1, body = :foot_2, mode = :heel)

T_fix = 0
zh = 0.025
xf2_el, zf2_el = ellipse_traj(pf2[1], pf2[1] + strd, zh, Tm - T_fix)
xf1_el, zf1_el = ellipse_traj(pf1[1], pf1[1] + strd, zh, Tm - T_fix)

zf2 = [[zf2_el[1] for t = 1:T_fix]..., zf2_el...]#, [zf2_el[end] for t = 1:Tm-1]...]
xf2 = [[xf2_el[1] for t = 1:T_fix]..., xf2_el...]#, [xf2_el[end] for t = 1:Tm-1]...]
zf1 = [[zf1_el[1] for t = 1:Tm]...]#-1 + T_fix]...]#, zf1_el...]
xf1 = [[xf1_el[1] for t = 1:Tm]...]#-1 + T_fix]..., xf1_el...]

p1_ref = [[xf1[t]; zf1[t]] for t = 1:T]
p2_ref = [[xf2[t]; zf2[t]] for t = 1:T]

plot(hcat(p1_ref...)', legend = :topleft)
plot!(hcat(p2_ref...)')

# Control
# u = (τ1..7, λ1..4, β1..8, ψ1..4, η1...8, s1)
# τ1: torso angle
# τ2: thigh 1 angle
# τ3: calf 1
# τ4: thigh 2
# τ5: calf 2
# τ6: foot 1
# τ7: foot 2

# ul <= u <= uu
# u1 = initial_torque(model, q1, h)[model.idx_u] # gravity compensation for current q
_uu = Inf * ones(model.m)
_uu[model.idx_u] .= Inf
_uu[end] = 2.0 * h
_ul = zeros(model.m)
_ul[model.idx_u] .= -Inf
_ul[end] = 0.75 * h
ul, uu = control_bounds(model, T, _ul, _uu)

qL = [-Inf; -Inf; q1[3] - pi / 50.0; q1[4:end] .- pi / 6.0; -Inf; -Inf; q1[3] - pi / 50.0; q1[4:end] .- pi / 6.0]
qU = [Inf; q1[2] + 0.001; 0.0; q1[4:end] .+ pi / 6.0; Inf; Inf; q1[3:end] .+ pi / 6.0]

xl, xu = state_bounds(model, T,
    qL, qU,
    x1 = [Inf * ones(model.nq); q1],
	xT = [Inf * ones(model.nq); qT[1]; Inf * ones(model.nq - 1)])

# Objective
include_objective(["velocity", "nonlinear_stage", "control_velocity"])

x0 = configuration_to_state(q_ref)

# penalty on slack variable
obj_penalty = PenaltyObjective(1.0e4, model.m-1)

# quadratic tracking objective
# Σ (x - xref)' Q (x - x_ref) + (u - u_ref)' R (u - u_ref)
q_penalty = 0.0 * ones(model.nq)
# q_penalty[1] = 1.0
# q_penalty[1] = 1.0
# q_penalty[2] = 1.0
# q_penalty[3] = 10.0
# q_penalty[8] = 10.0
# q_penalty[9] = 10.0
x_penalty = [q_penalty; q_penalty]
obj_control = quadratic_time_tracking_objective(
    [Diagonal(x_penalty) for t = 1:T],
    [Diagonal([1.0e-1 * ones(model.nu)..., 1.0e-3 * ones(model.nc)..., 1.0e-2 * ones(model.nb)..., 1.0e-8 * ones(model.m - model.nu - model.nc - model.nb - 1)..., 0.0]) for t = 1:T-1],
    [x_ref[end] for t = 1:T],
    [[zeros(model.nu); zeros(model.m - model.nu)] for t = 1:T-1],
	1.0)

obj_ctrl_vel = control_velocity_objective(Diagonal([1.0e-1 * ones(model.nu); 1.0e-2 * ones(model.nc + model.nb); zeros(model.m - model.nu - model.nc - model.nb)]))

# quadratic velocity penalty
q_v = 1.0e-3 * ones(model.nq)
# q_v[3] = 100.0
# q_v[3:7] .= 1.0e-3
# q_v[8:9] .= 100.0
obj_velocity = velocity_objective(
    [Diagonal(q_v) for t = 1:T-1],
    model.nq,
    h = h)

function l_foot_height(x, u, t)
	J = 0.0
	q1 = view(x, 1:7)
	q2 = view(x, 7 .+ (1:7))
	#
	if true#t >= Tm || t <= T_fix + 1
		pq1 = kinematics_2(model, q1, body = :calf_1, mode = :ee)
		pq2 = kinematics_2(model, q2, body = :calf_1, mode = :ee)
		v = (pq2 - pq1) ./ h
		p_avg = 0.5 * (pq1[1] + pq2[1])
		p_max = max(pq1[1], pq2[1])
		J += 1000.0 * sum((p1_ref[t] - kinematics_2(model, q1, body = :calf_1, mode = :ee)).^2.0)
		J += 1000.0 * sum((p1_ref[t] - kinematics_2(model, q2, body = :calf_1, mode = :ee)).^2.0)
		# J += 1000.0 * v' * v
		# J += 1000.0 * (x_ref[end][1] - q1[1])^2.0
	end

	if true#t <= Tm + T_fix || t == T
		pq1 = kinematics_2(model, q1, body = :calf_2, mode = :ee)
		pq2 = kinematics_2(model, q2, body = :calf_2, mode = :ee)
		v = (pq2 - pq1) ./ h
		p_avg = 0.5 * (pq1[1] + pq2[1])
		p_max = max(pq1[1], pq2[1])

		J += 1000.0 * sum((p2_ref[t] - kinematics_2(model, q1, body = :calf_2, mode = :ee)).^2.0)
		J += 1000.0 * sum((p2_ref[t] - kinematics_2(model, q2, body = :calf_2, mode = :ee)).^2.0)
		# J += 1000.0 * v' * v
		# J += 1000.0 * (p_avg - q1[1])^2.0
	end
	# #
	# # if t < T
	# # 	J += 1.0e-3 * ((q2[1] - q1[1]) / max(1.0e-6, u[model.m]) - 1.0)^2.0
	# # end
	pq1 = kinematics_2(model, q1, body = :calf_1, mode = :ee)
	pq2 = kinematics_2(model, q1, body = :calf_2, mode = :ee)
	p_avg = 0.5 * (pq1[1] + pq2[1])
	J += 1000.0 * (p_avg - q1[1])^2.0

	return J
end

l_foot_height(x) = l_foot_height(x, nothing, T)

obj_foot_height = nonlinear_stage_objective(l_foot_height, l_foot_height)

obj = MultiObjective([obj_penalty,
                      # obj_control,
                      # obj_velocity,
					  obj_ctrl_vel,
					  obj_foot_height])
# Constraints
include_constraints(["contact", "loop", "free_time", "stage"])

function pinned1!(c, x, u, t)
    q1 = view(x, 1:7)
	# q2 = view(x, 9 .+ (1:9))
    c[1:2] = p1_ref[t] - kinematics_2(model, q1, body = :calf_1, mode = :ee)
	# c[3:4] = p2_ref[t] - kinematics_3(model, q1, body = :foot_2, mode = :com)
	nothing
end

function pinned2!(c, x, u, t)
    q1 = view(x, 1:7)
    # c[1:2] = p1_ref[t] - kinematics_3(model, q, body = :foot_1, mode = :com)
	c[1:2] = p2_ref[t] - kinematics_2(model, q1, body = :calf_2, mode = :ee)
	nothing
end

function pinned3!(c, x, u, t)
	if t == T
	    q1 = view(x, 7 .+ (1:7))
	    # c[1:2] = p1_ref[t] - kinematics_3(model, q, body = :foot_1, mode = :com)
		c[1:2] = p2_ref[t] - kinematics_2(model, q1, body = :calf_2, mode = :ee)
	else
		c[1:2] .= 0.0
	end
	nothing
end

p2_ref[T]
# cc = zeros(4)
# xx = zeros(model.n)
# uuu = zeros(model.m)
# pinned!(cc, xx, uuu, 1)
# p1_ref[1]
n_stage = 2
t_idx1 = vcat([t for t = 1:Tm]...)#+T_fix]..., T)
t_idx2 = vcat([1:T_fix])#, [t for t = Tm:T]...)
t_idx3 = vcat([t for t = 1:Tm]...)#, [t for t = Tm:T]...)

con_pinned1 = stage_constraints(pinned1!, n_stage, (1:0), t_idx1)
# con_pinned2 = stage_constraints(pinned2!, n_stage, (1:0), t_idx2)
con_pinned3 = stage_constraints(pinned3!, n_stage, (1:0), t_idx3)

# constraints!(zeros(con_pinned.n), zeros(prob.num_var), con_pinned, model, prob.prob.idx, h, T)
con_loop = loop_constraints(model, collect([(2:7)...,(9:14)...]), 1, T, perm = Array(cat(perm, perm, dims = (1,2))))
con_contact = contact_constraints(model, T)
con_free_time = free_time_constraints(T)
con = multiple_constraints([con_contact,
	con_free_time,
	con_loop,
	con_pinned1])#,
	# con_pinned3])#,
	# con_pinned2])
	# con_pinned3])#,
	# con_pinned2])

con_pinned3.t_idx
# con = multiple_constraints([con_free_time, con_contact, con_pinned1])

# Problem
prob = trajectory_optimization_problem(model,
               obj,
               T,
               xl = xl,
               xu = xu,
               ul = ul,
               uu = uu,
               con = con)

# trajectory initialization
u0 = [[1.0e-2 * randn(model.nu); 0.01 * randn(model.m - model.nu - 1); h] for t = 1:T-1] # random controls

# Pack trajectories into vector
z0 = pack(x0, u0, prob)
# Solve
# NOTE: run multiple times to get good trajectory
# include_snopt()
@time z̄, info = solve(prob, copy(z0),
    nlp = :ipopt,
	max_iter = 2000,
    tol = 1.0e-3, c_tol = 1.0e-3, mapl = 5,
    time_limit = 60 * 3)

# @time z̄, info = solve(prob, copy(z̄ .+ 0.01 * randn(prob.num_var)),
#     nlp = :SNOPT7,
# 	max_iter = 1000,
#     tol = 1.0e-3, c_tol = 1.0e-3, mapl = 5,
#     time_limit = 60 * 3)
# @load joinpath(@__DIR__, "walker_gait_3.jld2") z̄ x̄ ū q̄ τ̄ λ̄ b̄ h̄

@show check_slack(z̄, prob)
x̄, ū = unpack(z̄, prob)
_tf, _t, h̄ = get_time(ū)
q = state_to_configuration(x̄)
u = [u[model.idx_u] for u in ū]
γ = [u[model.idx_λ] for u in ū]
b = [u[model.idx_b] for u in ū]
ψ = [u[model.idx_ψ] for u in ū]
η = [u[model.idx_η] for u in ū]
h̄ = mean(h̄)

t = 1
# ψ[t] .* friction_cone(model, γ[t], b[t])
# model.idx_ψ
# γ[t]
# b[t]
# ψ[t]
# friction_cone(model, γ[t], b[t])

norm(hcat([γ[t] .* ϕ_func(model, q[t+2]) for t = 1:T-1]...),Inf)
norm(hcat([b[t] .* η[t] for t = 1:T-1]...), Inf)
function friction_cone(model::Biped, λ, b)
	# @show model.μ
	return @SVector [model.μ * λ[1] - sum(b[1:2]),
					 model.μ * λ[2] - sum(b[3:4])]
end
norm(hcat([ψ[t] .* friction_cone(model, γ[t], b[t]) for t = 1:T-1]...), Inf)

@save joinpath(@__DIR__, "biped_gait_alt.jld2") z̄ x̄ ū h̄ q u γ b
# @load joinpath(@__DIR__, "walker_gait.jld2") z̄ x̄ ū q̄ τ̄ λ̄ b̄ h̄

maximum([norm(fd(model, x̄[t+1], x̄[t], ū[t], zeros(model.d), h̄, t)) for t = 1:T-1])
(q[end][1] - q[1][1]) / _tf
vis = Visualizer()
render(vis)
visualize!(vis, model,
	[[x̄[1][1:model.nq] for i = 1:10]...,
	 state_to_configuration(x̄)...,
	 [x̄[end][model.nq .+ (1:model.nq)] for i = 1:10]...], Δt = h̄[1])

plot(hcat(u...)', linetype = :steppost)
plot(hcat([γt[1:2] for γt in γ]...)', linetype = :steppost)

plot!(hcat([bt[1:2] for bt in b]...)', linetype = :steppost)
plot!(hcat([bt[3:4] for bt in b]...)', linetype = :steppost)

_pf1 = [kinematics_2(model, qt, body = :calf_1, mode = :ee) for qt in q]
_pf2 = [kinematics_2(model, qt, body = :calf_2, mode = :ee) for qt in q]

# using Plots

plot(hcat(p1_ref...)', title = "foot 1",
 	legend = :topleft, label = ["x" "z"], color = :black, width = 2.0)
plot!(hcat(_pf1...)', legend = :topleft, color = :red, width = 1.0)

plot(hcat(p2_ref...)', title = "foot 2",
	legend = :topleft, label = ["x" "z"], color = :black, width = 2.0)
plot!(hcat(_pf2...)', title = "foot 2", legend = :bottomright,
	color = :red, width = 1.0)

# # # # @save joinpath(pwd(), "examples/trajectories/walker_steps.jld2") z̄
# # # plot(hcat(ū...)[1:model.nu, :]', linetype = :steppost)
function get_q_viz(q̄; N = 4)
	q_viz = [q̄...]
	shift_vec = zeros(model.nq)
	shift_vec[1] = q̄[end][1]
	for i = 1:N
		# println(shift)
		# shift_vec[1] = strd
		#
		q_update = [q + shift_vec for q in q̄[2:end]]
		push!(q_viz, q_update...)
		shift_vec[1] = q_update[end][1]
	end

	return q_viz
end

q_viz = get_q_viz(q)
render(vis)
visualize!(vis, model,
	q_viz,
	Δt = h̄[1])

# function traj_concat(q̄, ū; N = 3)
# 	u_viz = [ū...]
# 	q_viz = [q̄...]
#
# 	shift_vec = zeros(model.nq)
# 	shift_vec[1] = q̄[end][1]
#
# 	for i = 1:N
# 		push!(u_viz, ū...)
# 		# println(shift)
# 		# shift_vec[1] = strd
# 		#
# 		q_update = [q + shift_vec for q in q̄[3:end]]
# 		push!(q_viz, q_update...)
# 		shift_vec[1] = q_update[end][1]
# 	end
#
# 	return q_viz, u_viz
# end
#
# q_cat, u_cat = traj_concat(q̄, ū, N = 25)
# plot(hcat(q_cat...)', label = "")
# plot(hcat(u_cat...)[1:model.nu, :]', label = "")
# x_cat = configuration_to_state(q_cat)
# u_cat
#
# qp = ones(model.nq)
# qp[1] = 1.0
# qp[2] = 1.0
# Q = [Diagonal([qp; qp]) for t = 1:length(x_cat)]
# R = [Diagonal(ones(model.m)) for t = 1:length(u_cat)]
#
# K, P = tvlqr(model, x_cat, u_cat, nothing, Q, R)
#
# plot(hcat([vec(k) for k in K]...)', label = "")
#
# K_walker = K[1:T-1]
# x̄_walker = x̄[1:T]
# ū_walker = ū[1:T-1]
# h̄_walker = h̄[1:T-1]
#
# @save joinpath(@__DIR__, "walker_lqr.jld2") K_walker x̄_walker ū_walker h̄_walker
# # (q̄[end][1] - q̄[1][1]) / tf_
# #
# # h̄[1]
# #
# model_sim = Walker{Discrete, FixedTime}(n, m, d,
# 			  g_world, μ_world,
# 			  l_torso, d_torso, m_torso, J_torso,
# 			  l_thigh, d_thigh, m_thigh, J_thigh,
# 			  l_calf, d_calf, m_calf, J_calf,
# 			  l_foot, d_foot, m_foot, J_foot,
# 			  l_thigh, d_thigh, m_thigh, J_thigh,
# 			  l_calf, d_calf, m_calf, J_calf,
# 			  l_foot, d_foot, m_foot, J_foot,
# 			  qL, qU,
# 			  uL, uU,
# 			  nq,
# 			  nu,
# 			  nc,
# 			  nf,
# 			  nb,
# 			  ns,
# 			  idx_u,
# 			  idx_λ,
# 			  idx_b,
# 			  idx_ψ,
# 			  idx_η,
# 			  idx_s,
# 			  joint_friction)
#
# x_sim = [x̄[1]]
# q_sim = [x̄[1][1:model.nq], x̄[2][model.nq .+ (1:model.nq)]]
# include(joinpath(pwd(), "src/contact_simulator/simulator.jl"))
#
# for t = 1:T-1
# 	_x = step_contact(model_sim, x_sim[end], ū[t][1:model.nu], zeros(model.d), h̄,
# 	        tol_c = 1.0e-5, tol_opt = 1.0e-5, tol_s = 1.0e-4, nlp = :ipopt)
# 	push!(x_sim, _x)
# 	push!(q_sim, x_sim[end][model.nq .+ (1:model.nq)])
# end
#
# plot(hcat(q...)[:, 1:length(q_sim)]', color = :black, width = 2.0)
# plot!(hcat(q_sim...)', color = :red, width = 1.0)
# vis = Visualizer()
# render(vis)
# # open(vis)
# visualize!(vis, model, q_sim, Δt = h̄[1])

max_issues_repo_name: thowell/DirectMotionPlanning


text:
#include("../functions_decompositions/decomposeX.jl")

function greedy_lowrank_EigenAlign(A,B,c1,c2,c3,iters)
  Uk,Vk,Wk,W1,W2 = decomposeX_balance_allfactors(A,B,iters,c1,c2,c3)
  U1,V1 = split_balanced_decomposition(Uk,Wk,Vk);

  k = size(U1,2)
  n = size(U1,1)

  ma = zeros(Int64,n)
  mb = zeros(Int64,n)

  ma_ids = collect(1:n)
  mb_ids = collect(1:n)

  U1t = U1';
  V1t = V1';
  lastid = n;

  for i = 1:n
    curmax = 0
    curma = 0
    curmb = 0
    for mi = ma_ids
      u = U1t[:,mi]
      for mj = mb_ids

        v = V1t[:,mj]
        m = dot(u,v)
        if m > curmax
          curmax = m
          curma = mi
          curmb = mj
        end
      end
    end

    if curma == 0
      lastid = i-1
      break;
    end
    @show curma
    @show curmb

    ma[i] = curma
    mb[i] = curmb
    ma_ids = setdiff(ma_ids,curma)
    mb_ids = setdiff(mb_ids,curmb)
  end
  ma = ma[1:lastid]
  mb = mb[1:lastid]

  return ma,mb
end

function greedy_lowrank(U1,V1)

  k = size(U1,2)
  n = size(U1,1)

  ma = zeros(Int64,n)
  mb = zeros(Int64,n)

  ma_ids = collect(1:n)
  mb_ids = collect(1:n)

  U1t = U1';
  V1t = V1';
  lastid = n;

  for i = 1:n
    curmax = 0
    curma = 0
    curmb = 0
    for mi = ma_ids
      u = U1t[:,mi]
      for mj = mb_ids

        v = V1t[:,mj]
        m = dot(u,v)
        if m > curmax
          curmax = m
          curma = mi
          curmb = mj
        end
      end
    end

    if curma == 0
      lastid = i-1
      break;
    end
    @show curma
    @show curmb
    ma[i] = curma
    mb[i] = curmb
    ma_ids = setdiff(ma_ids,curma)
    mb_ids = setdiff(mb_ids,curmb)
  end

  ma = ma[1:lastid]
  mb = mb[1:lastid]

  return ma,mb
end

max_stars_repo_name: chuckcol/LambdaTAME

max_issues_repo_name: chuckcol/LambdaTAME

max_forks_repo_name: chuckcol/LambdaTAME


text:
dir = []

function next(side, kol)
    return HorizonSide((Int(side) + kol) % 4)
end

function set(side1, side2)
    while (!isborder(robot, side1) || !isborder(robot, side2))
        while !isborder(robot, side1)
            move!(robot, side1)
            push!(dir, side1)
        end
        while !isborder(robot, side2)
            move!(robot, side2)
            push!(dir, side2)
        end
    end
end

function goAround(side)
    num = 0
    kol = 0
    while (isborder(robot, side) && !isborder(robot, next(side, 1)))
        move!(robot, next(side, 1))
        push!(dir, next(side, 1))
        kol += 1
    end
    if !isborder(robot, side)
        move!(robot, side)
        push!(dir, side)
        num += 1
        while isborder(robot, next(side, 3))
            move!(robot, side)
            push!(dir, side)
            num += 1
        end
    end
    for i = 1:kol
        move!(robot, next(side, 3))
        push!(dir, next(side, 3))
    end
    return num
end

function drawLine(side)
    num = 0
    while true
        if !isborder(robot, side)
            putmarker!(robot)
            move!(robot, side)
            push!(dir, side)
            num += 1
        else
            putmarker!(robot)
            num += goAround(side)
        end
        if (num == 11 || (isborder(robot, side) && isborder(robot, next(side, 1))))
            break
        end
    end
    putmarker!(robot)
end

function f()
    set(Sud, West)

    side = Ost
    while true
        drawLine(side)
        if isborder(robot, Nord)
            break
        end
        move!(robot, Nord)
        push!(dir, Nord)
        side = next(side, 2)
    end
    while length(dir) > 0
        move!(robot, next(pop!(dir), 2))
    end
end

f()

max_stars_repo_name: TheXlebushek/mirea-progs

max_issues_repo_name: TheXlebushek/mirea-progs

max_forks_repo_name: TheXlebushek/mirea-progs


text:
const CENTI_POISE_TO_PASCAL_SECOND = 1000.0
const MOL_TO_KMOL = 1000.0
const KELVIN_TO_RANKINE = 5/9
const PASCAL_TO_PSI = 1.450377377302092e-04

"""
    single_phase_label(mixture, cond)

Li's method for single-phase labeling of a mixture. Estimate of pure vapor/liquid.

Returns a vapor fraction that is either 1.0 (=pure vapor) or 0.0 (=pure liquid).
"""
function single_phase_label(mixture, cond)
    # Li's method for phase labeling.
    z = cond.z
    T_c = 0.0
    V_c = 0.0
    @inbounds for (i, m) in enumerate(mixture.properties)
        V = m.V_c*z[i]
        T_c += m.T_c*V
        V_c += V
    end
    return Float64(cond.T < T_c)
end

"""
    lbc_viscosity(eos, p, T, ph; <keyword arguments>)

Compute the viscosity of a mixture using the [Lohrenz-Bray-Clark](https://doi.org/10.2118/915-PA) correlation.
"""
function lbc_viscosity(eos, p, temperature, ph::FlashedPhase{T}; coeff = (0.1023, 0.023364, 0.058533, -0.040758, 0.0093324), shift = -1e-4) where T
    z = ph.mole_fractions
    properties = eos.mixture.properties
    mw_mix, P_pc, T_pc, V_pc = pseudo_critical_properties(properties, z)
    mu_atm = atmospheric_mu_estimate(properties, z, temperature)
    e_mix = mixture_viscosity_parameter(mw_mix, T_pc, P_pc)
    # From Jossi et al
    # coeffs = [0.1023, 0.023364, 0.058533, -0.040758, 0.0093724, 1e-4]
    # From LBC paper
    # coeffs = [0.1023, 0.023364, 0.058533, -0.040758, 0.0093324, -1e-4]
    # Compute reduced density
    V = molar_volume(eos, p, temperature, ph)
    rho_r = V_pc/V
    corr = zero(T)
    for (i, c) in enumerate(coeff)
        corr += c*rho_r^(i-1)
    end
    mu = mu_atm + (corr^4 + shift)/e_mix
    return mu::T
end

function atmospheric_mu_estimate(props, z::V, temperature) where V<:AbstractVector{T} where T
    a = zero(T)
    b = zero(T)
    @inbounds for (prop, zi) in zip(props, z)
        mw = prop.mw
        T_c = prop.T_c
        p_c = prop.p_c
        # Add contributions to atmospheric mu
        T_r = temperature/T_c
        e_i = mixture_viscosity_parameter(mw, T_c, p_c)
        if T_r > 1.5
            mu_i = 17.78e-5*(4.58*T_r - 1.67)^0.625
        else
            mu_i = 34e-5*T_r^(0.94)
        end
        tmp = sqrt(mw)*zi
        a += tmp*mu_i/e_i
        b += tmp
    end
    return a/b
end

@inline function mixture_viscosity_parameter(molar_mass, T_c, p_c)
    A = 5.4402*(T_c*KELVIN_TO_RANKINE)^(1/6)
    B = sqrt(MOL_TO_KMOL*molar_mass)*(p_c*PASCAL_TO_PSI)^(2/3)
    return CENTI_POISE_TO_PASCAL_SECOND*A/B
end

function pseudo_critical_properties(props, z::V) where V<:AbstractVector{T} where T
    P_pc = zero(T)
    T_pc = zero(T)
    Vc = zero(T)
    mwc = zero(T)
    @inbounds for (prop, zi) in zip(props, z)
        # Accumulate weighted properties
        mwc += zi*prop.mw
        P_pc += zi*prop.p_c
        T_pc += zi*prop.T_c
        Vc += zi*prop.V_c
    end
    return (mwc, P_pc, T_pc, Vc)
end

@inline function molar_volume(R, p, T, Z)
    return R*T*Z/p
end


text:
using Pkg

ssh = true
if ssh
    git = "git@github.com:SebastianAment/"
else
    git = "https://github.com/SebastianAment/"
end

add(x) = Pkg.add(Pkg.PackageSpec(url = git * x * ".git"))

add("LazyInverse.jl")
add("LinearAlgebraExtensions.jl")


text:
#!/usr/bin/env julia
include("post_process_defs.jl")

function isDir(path)
  ftype = strip(split(read(`file $path`, String), ":")[2])
  return ftype == "directory"
end

function lsRec(path)
  return split(read(`find $path`, String))
end

#= Test
lsRec("results/sim3")
=#

function dirContains(word, dir)::Bool
  allFiles = lsRec(dir)
  for f in allFiles
    if occursin(word, f)
      return true
    end
  end
  return false
end

### MAIN ###
RESULTS_DIR = length(ARGS) >= 1 ? ARGS[1] : "results/sim3/"
MCMC_OUTPUT = filter(d -> occursin(".jld2", d), lsRec(RESULTS_DIR))
# TODO: make retrieve4
# Then make postprocess4 with this extra line
#MCMC_OUTPUT = filter(d -> occursin("N_factor10000", d), MCMC_OUTPUT) # Run this additional filter
successes = []
failures = []

@time for mcmc in MCMC_OUTPUT
  println("Processing: $mcmc")
  try
    @time post_process(mcmc)
    append!(successes, [mcmc])
  catch
    append!(failures, [mcmc])
  end
end

println("Successes")
[println(s) for s in successes];
println()

println("Failures")
[println(f) for f in failures];
println()


text:
# ---
# title: 210. Course Schedule II
# id: problem210
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Depth-first Search, Breadth-first Search, Graph, Topological Sort
# link: <https://leetcode.com/problems/course-schedule-ii/description/>
# hidden: true
# ---
#
# There are a total of `n` courses you have to take labelled from `0` to `n -
# 1`.
#
# Some courses may have `prerequisites`, for example, if `prerequisites[i] =
# [ai, bi]` this means you must take the course `bi` before the course `ai`.
#
# Given the total number of courses `numCourses` and a list of the
# `prerequisite` pairs, return the ordering of courses you should take to finish
# all courses.
#
# If there are many valid answers, return **any** of them. If it is impossible
# to finish all courses, return **an empty array**.
#
#
#
# **Example 1:**
#
#
#
#     Input: numCourses = 2, prerequisites = [[1,0]]
#     Output: [0,1]
#     Explanation: There are a total of 2 courses to take. To take course 1 you should have finished course 0. So the correct course order is [0,1].
#
#
# **Example 2:**
#
#
#
#     Input: numCourses = 4, prerequisites = [[1,0],[2,0],[3,1],[3,2]]
#     Output: [0,2,1,3]
#     Explanation: There are a total of 4 courses to take. To take course 3 you should have finished both courses 1 and 2. Both courses 1 and 2 should be taken after you finished course 0.
#     So one correct course order is [0,1,2,3]. Another correct ordering is [0,2,1,3].
#
#
# **Example 3:**
#
#
#
#     Input: numCourses = 1, prerequisites = []
#     Output: [0]
#
#
#
#
# **Constraints:**
#
#   * `1 <= numCourses <= 2000`
#   * `0 <= prerequisites.length <= numCourses * (numCourses - 1)`
#   * `prerequisites[i].length == 2`
#   * `0 <= ai, bi < numCourses`
#   * `ai != bi`
#   * All the pairs `[ai, bi]` are **distinct**.
#
#
## @lc code=start
using LeetCode

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


text:
import Pkg;
Pkg.instantiate();

using PackageCompiler

create_app(joinpath(@__DIR__, ".."), "build"; executables = ["monkey" => "julia_main"])


text:
"""
    OILMM(fs, U, S, D)

An Orthogonal Instantaneous Linear Mixing Model (OILMM) -- a distribution over vector-
valued functions. Let `p` be the number of observed outputs, and `m` the number of latent
processes, then

# Arguments:
- fs: a length-`m` vector of Gaussian process objects from Stheno.jl.
- U: a `p x m` matrix with mutually orthonormal columns.
- S: an `m x m` `Diagonal` matrix. Same dim. as `fs`. Positive entries.
- D: an `m x m` `Diagonal` matrix, variance of noise on each latent process. Same size as
    `S`. Positive entries.

We recommend constructing `U` and `S` from the `svd` of some other matrix. e.g.
```julia
H = randn(p, m)
U, s, _ = svd(H)
S = Diagonal(H)
```
"""
struct OILMM{
    Tfs<:AbstractVector{<:AbstractGP},
    TU<:AbstractMatrix{<:Real},
    TS<:Diagonal{<:Real},
    TD<:Diagonal{<:Real},
} <: AbstractGP
    fs::Tfs
    U::TU
    S::TS
    D::TD
end

noise_var(Σ::Diagonal{<:Real, <:Fill}) = FillArrays.getindex_value(Σ.diag)

function unpack(fx::FiniteGP{<:OILMM, <:MOInput, <:Diagonal{<:Real, <:Fill}})
    fs = fx.f.fs
    S = fx.f.S
    U = fx.f.U
    D = fx.f.D
    σ² = noise_var(fx.Σy)
    x = fx.x.x

    # Check that the number of outputs requested agrees with the model.
    fx.x.out_dim == size(U, 1) || throw(error("out dim of x != out dim of f."))
    return fs, S, U, D, σ², x
end

reshape_y(y::AbstractVector{<:Real}, N::Int) = ColVecs(reshape(y, N, :)')

# Note that `cholesky` exploits the diagonal structure of `S`.
function project(
    S::Diagonal{T},
    U::AbstractMatrix{T},
    Y::ColVecs{T},
    σ²::T,
    D::Diagonal{T},
) where {T<:Real}

    # Compute the projection of the data.
    Yproj = cholesky(S).U \ U' * Y.X

    # Compute the projected noise, which is a matrix of size `size(Yproj)`.
    ΣT = repeat(diag(σ² * inv(S) + D), 1, size(Yproj, 2))

    return Yproj, ΣT
end

# Compute the regularisation term in the log marginal likelihood. See e.g. appendix A.4.
function regulariser(
    S::Diagonal{T},
    U::AbstractMatrix{T},
    σ²::T,
    Y::ColVecs{T},
) where {T<:Real}
    n = length(Y)
    p, m = size(U)
    return -(n * (logdet(cholesky(S)) + (p - m) * log(2π * σ²)) +
        sum(abs2, (I - U * U') * Y.X) / σ²) / 2
end

"""
    rand_latent(rng::AbstractRNG, fx::FiniteGP{<:OILMM})

Sample from the latent (noiseless) process.

See also `rand`.
"""
function rand_latent(rng::AbstractRNG, fx::FiniteGP{<:OILMM})
    fs, S, U, D, σ², x = unpack(fx)

    # Generate from the latent processes.
    X = hcat(map((f, d) -> rand(rng, f(x, d)), fs, D.diag)...)

    # Transform latents into observed space.
    return vec(U * cholesky(S).U * X')
end

"""
    rand(rng::AbstractRNG, fx::FiniteGP{<:OILMM})

Sample from the OILMM, including the observation noise.
Follows generative structure of model 2 from [1].
Follows the AbstractGPs.jl API.

See also `rand_latent`.

[1] - Bruinsma et al 2020.
"""
function AbstractGPs.rand(rng::AbstractRNG, fx::FiniteGP{<:OILMM})

    # Sample from the latent process.
    F = rand_latent(rng, fx)

    # Generate iid noise and add to each output.
    return F .+ sqrt(noise_var(fx.Σy)) .* randn(rng, size(F))
end

"""
    denoised_marginals(fx::FiniteGP{<:OILMM})

Returns the marginal distribution over the OILMM without the IID noise components.

See also `marginals`.
"""
function denoised_marginals(fx::FiniteGP{<:OILMM})
    fs, S, U, D, σ², x = unpack(fx)

    # Compute the marginals over the independent latents.
    fs_marginals = reduce(hcat, map(f -> marginals(f(x)), fs))
    M_latent = mean.(fs_marginals)'
    V_latent = var.(fs_marginals)'

    # Compute the latent -> observed transform.
    H = U * cholesky(S).U

    # Compute the means.
    M = H * M_latent

    # Compute the variances.
    V = abs2.(H) * V_latent

    # Package everything into independent Normal distributions.
    return AbstractGPs.Normal.(vec(M'), sqrt.(vec(V')))
end

# See AbstractGPs.jl API docs.
function AbstractGPs.mean_and_var(fx::FiniteGP{<:OILMM})
    fs, S, U, D, σ², x = unpack(fx)

    # Compute the marginals over the independent latents.
    fs_marginals = hcat(map(f -> marginals(f(x)), fs)...)
    M_latent = mean.(fs_marginals)'
    V_latent = var.(fs_marginals)'

    # Compute the latent -> observed transform.
    H = U * cholesky(S).U

    # Compute the means.
    M = H * M_latent

    # Compute the variances.
    V = abs2.(H) * (V_latent .+ D.diag) .+ σ²

    # Package everything into independent Normal distributions.
    return vec(M'), vec(V')
end

AbstractGPs.mean(fx::FiniteGP{<:OILMM}) = mean_and_var(fx)[1]

AbstractGPs.var(fx::FiniteGP{<:OILMM}) = mean_and_var(fx)[2]

# See AbstractGPs.jl API docs.
function AbstractGPs.logpdf(fx::FiniteGP{<:OILMM}, y::AbstractVector{<:Real})
    fs, S, U, D, σ², x = unpack(fx)

    # Projection step.
    Y = reshape_y(y, length(x))
    Yproj, ΣT = project(S, U, Y, σ², D)

    # Latent process log marginal likelihood calculation.
    y_rows = collect(eachrow(Yproj))
    ΣT_rows = collect(eachrow(ΣT))
    lmls_latents = map((f, s, y) -> logpdf(f(x, collect(s)), collect(y)), fs, ΣT_rows, y_rows)

    return regulariser(S, U, σ², Y) + sum(lmls_latents)
end

# See AbstractGPs.jl API docs.
function AbstractGPs.posterior(fx::FiniteGP{<:OILMM}, y::AbstractVector{<:Real})
    fs, S, U, D, σ², x = unpack(fx)

    # Projection step.
    Y = reshape_y(y, length(x))
    Yproj, ΣT = project(S, U, Y, σ², D)

    # Condition each latent process on the projected observations.
    y_rows = collect(eachrow(Yproj))
    ΣT_rows = collect(eachrow(ΣT))
    fs_posterior = map((f, s, y) -> posterior(f(x, collect(s)), collect(y)), fs, ΣT_rows, y_rows)
    return OILMM(fs_posterior, U, S, D)
end


text:
"""
    Fermi.CoupledCluster.RCCSD

Fermi struct that holds information about RCCSD wavefunctions

# Fields

    CorrelationEnergy   CCSD correlation energy
    T1                  T1 amplitudes
    T2                  T2 amplitudes

_struct tree:_

**RCCSD** <: AbstractCCWavefunction <: AbstractCorrelatedWavefunction <: AbstractWavefunction
"""
struct RCCSD{T} <: AbstractCCWavefunction
    GuessEnergy::T
    CorrelationEnergy::T
    T1::_T where _T <: Fermi.AbstractTensor
    T2::_T where _T <: Fermi.AbstractTensor
end

# Most general function. It defines the precision and call a precision-specific function
"""
    Fermi.CoupledCluster.RCCSD()

Compute a RCCSD wave function using Fermi.CurrentOptions data.
"""
function RCCSD()
    prec = select_precision(Fermi.CurrentOptions["precision"])
    RCCSD{prec}()
end

function RCCSD(refwfn::RHF)
    prec = select_precision(Fermi.CurrentOptions["precision"])
    alg = select_algorithm(Fermi.CurrentOptions["cc_alg"])
    dummy = RCCSD{prec}(0.0,0.0,Fermi.MemTensor(zeros(prec,0,0)),Fermi.MemTensor(zeros(prec,0,0,0,0)))
    RCCSD{prec}(refwfn, dummy, alg)
end

function RCCSD(guess::RCCSD{T}) where T <: AbstractFloat
    prec = eltype(guess.T2.data)
    RCCSD{prec}(guess)
end

function RCCSD{T}() where T <: AbstractFloat
    dummy = RCCSD{T}(0.0,0.0,Fermi.MemTensor(zeros(T,0,0)),Fermi.MemTensor(zeros(T,0,0,0,0)))
    RCCSD{T}(dummy)
end

"""
    Fermi.CoupledCluster.RCCSD{T}()

Compute a RCCSD wave function for a given precision T (Float64 or Float32)
"""
function RCCSD{T}(guess::RCCSD{Tb}) where { T <: AbstractFloat,
                                           Tb <: AbstractFloat }
    alg = select_algorithm(Fermi.CurrentOptions["cc_alg"])
    RCCSD{T}(guess,alg)
end

# implementation specific functions (look at one of these for guidance implementing a new method!)
include("CTF.jl")
include("DF-CTF.jl")

# main kernel (ccsd iteration logic)
include("kernel.jl")


text:
type XPTFile
    fname::String
    io::IO
    info::XPTInfo

    function XPTFile(fname, io, info)
        xpt = new(fname, io, info)
        finalizer(xpt, closexpt)
        xpt
    end
end
function show(io::IO, xpt::XPTFile)
    println(io, "File name: ", xpt.fname)
    show(io, xpt.info)
end

function XPTFile(fname::String)
    io = open(fname)
    info = XPTInfo(io)
    XPTFile(fname, io, info)
end

function XPTFile(io::IO)
    info = XPTInfo(io)
    XPTFile(io.name, io, info)
end

closexpt(xpt::XPTFile) = close(xpt.io)

function readdf(xpt::XPTFile)
    seek(xpt.io, ROW_LENGTH * 3)
    dsinfo = DSInfo(xpt.io)
    readrecords(xpt.io, dsinfo.namestrs)
end

function readrecords(io::IO, namestrs::Array{NameStr, 1})
    obs_header=readbytes(io, ROW_LENGTH)
    @assert obs_header[1:48] == OBS_HEADER_RECORD

    reclen = mapreduce(ns-> ns.nlng, +, namestrs)

    vararray = {}
    for ns in namestrs
        push!(vararray, ns.ntype == 1? DataArray(Float64[]) : ASCIIString[])
    end

    pdv= Uint8[]
    nextline = Uint8[]

    done=false

    while !done
        if isempty(nextline)
            if eof(io)
                done = true
            else
                nextline = readbytes(io, ROW_LENGTH)
                if nextline[1:48] == MEMBER_HEADER_RECORD
                    skip(io, -ROW_LENGTH)
                    done=true
                end
            end
        else
            append!(pdv, splice!(nextline, 1:min(reclen-size(pdv, 1), size(nextline, 1))))

            assert(size(pdv, 1) <= reclen)

            if size(pdv, 1) == reclen
                for j in 1:size(namestrs, 1)
                    i1 = namestrs[j].npos + 1
                    i2 = namestrs[j].nlng + i1 - 1
                    if namestrs[j].ntype == 1
                        num=zeros(Uint8, 8)
                        num[1:namestrs[j].nlng] = pdv[i1:i2]
                        if ENDIAN_BOM == 0x04030201
                            reverse!(num)
                        end
                        push!(vararray[j], IBM2Float64(num))
                    else
                        push!(vararray[j], ascii(pdv[i1:i2]))
                    end
                end
                pdv = Uint8[]
            end
        end
    end

    names = map(ns -> symbol(strip(lowercase(ns.nname))), namestrs)

    DataFrame(vararray, names)
end


text:
""" Category of (possibly infinite) sets and functions.

This module defines generic types for the category of sets ([`SetOb`](@ref),
[`SetFunction`](@ref)), as well as a few basic concrete types, such as a wrapper
type to view Julia types as sets ([`TypeSet`](@ref)). Extensive support for
finite sets is provided by another module, [`FinSets`](@ref).
"""
module Sets
export SetOb, TypeSet, PredicatedSet, SetFunction, ConstantFunction

using AutoHashEquals
using FunctionWrappers: FunctionWrapper

using ...GAT, ..FreeDiagrams, ..Limits
using ...Theories: Category
import ...Theories: dom, codom, id, compose, ⋅, ∘
import ..Limits: limit, universal

# Data types
############

""" Abstract type for object in the category **Set**.

The type parameter `T` is the element type of the set.

Note: This type is more abstract than the built-in Julia types `AbstractSet` and
`Set`, which are intended for data structures for finite sets. Those are
encompassed by the subtype [`FinSet`](@ref).
"""
abstract type SetOb{T} end

Base.eltype(::Type{<:SetOb{T}}) where T = T

""" A Julia data type regarded as a set.
"""
struct TypeSet{T} <: SetOb{T} end

TypeSet(T::Type) = TypeSet{T}()

Base.show(io::IO, ::TypeSet{T}) where T = print(io, "TypeSet($T)")

""" Abstract type for morphism in the category **Set**.

Every instance of `SetFunction{<:SetOb{T},<:SetOb{T′}}` is callable with
elements of type `T`, returning an element of type `T′`.

Note: This type would be better called simply `Function` but that name is
already taken by the base Julia type.
"""
abstract type SetFunction{Dom <: SetOb, Codom <: SetOb} end

SetFunction(f::Function, args...) = SetFunctionCallable(f, args...)
SetFunction(::typeof(identity), args...) = SetFunctionIdentity(args...)

show_type(io::IO, ::Type{<:SetFunction}) = print(io, "SetFunction")

""" Function in **Set** defined by a callable Julia object.
"""
@auto_hash_equals struct SetFunctionCallable{
    T,T′,Dom<:SetOb{T},Codom<:SetOb{T′}} <: SetFunction{Dom,Codom}
  # Field `func` is usually a `Function` but can be any Julia callable.
  func::FunctionWrapper{T′,Tuple{T}}
  dom::Dom
  codom::Codom

  function SetFunctionCallable(f, dom::Dom, codom::Codom) where
      {T,T′,Dom<:SetOb{T},Codom<:SetOb{T′}}
    new{T,T′,Dom,Codom}(FunctionWrapper{T′,Tuple{T}}(f), dom, codom)
  end
end

function (f::SetFunctionCallable{T,T′})(x::T)::T′ where {T,T′}
  f.func(x)
end

function Base.show(io::IO, f::F) where F <: SetFunctionCallable
  func = f.func.obj[] # Deference FunctionWrapper
  show_type(io, F)
  print(io, "($(nameof(func)), $(f.dom), $(f.codom))")
end

""" Identity function in **Set**.
"""
@auto_hash_equals struct SetFunctionIdentity{Dom} <: SetFunction{Dom,Dom}
  dom::Dom
end

function SetFunctionIdentity(dom::SetOb, codom::SetOb)
  dom == codom || error("Domain mismatch in identity function: $dom != $codom")
  SetFunctionIdentity(dom)
end

codom(f::SetFunctionIdentity) = f.dom

(f::SetFunctionIdentity)(x) = x

function Base.show(io::IO, f::F) where F <: SetFunctionIdentity
  show_type(io, F)
  print(io, "(identity, $(f.dom))")
end

""" Function in **Set** taking a constant value.
"""
@auto_hash_equals struct ConstantFunction{T,Value<:T,Dom,Codom<:SetOb{T}} <:
    SetFunction{Dom,Codom}
  value::Value
  dom::Dom
  codom::Codom
end

ConstantFunction(value::T, dom::SetOb) where T =
  ConstantFunction(value, dom, TypeSet{T}())

(f::ConstantFunction)(x) = f.value

# Predicated sets
#################

""" Set defined by a predicate (boolean-valued function) on a Julia data type.
"""
struct PredicatedSet{T} <: SetOb{T}
  predicate::FunctionWrapper{Bool,Tuple{T}}

  PredicatedSet{T}(f) where T = new{T}(FunctionWrapper{Bool,Tuple{T}}(f))
end

PredicatedSet(T::Type, f) = PredicatedSet{T}(f)

function (s::PredicatedSet{T})(x::T)::Bool where {T}
  s.predicate(x)
end

function Base.show(io::IO, s::PredicatedSet{T}) where T
  func = s.predicate.obj[] # Deference FunctionWrapper
  print(io, "PredicatedSet($T, $(nameof(func)))")
end

const PredicatedFunction{T,T′} =
  SetFunctionCallable{T,T′,PredicatedSet{T},PredicatedSet{T′}}

function (f::PredicatedFunction{T,T′})(x::T)::T′ where {T,T′}
  dom(f)(x) || error("Predicate not satisfied on input: $x")
  y = f.func(x)
  codom(f)(y) || error("Predicate not satisfied on output: $y")
  y
end

# Category of sets
##################

""" Category of sets and functions.
"""
@instance Category{SetOb, SetFunction} begin
  dom(f::SetFunction) = f.dom
  codom(f::SetFunction) = f.codom

  id(A::SetOb) = SetFunction(identity, A, A)

  function compose(f::SetFunction, g::SetFunction)
    codom(f) == dom(g) ||
      error("Domain mismatch in composition: $(codom(f)) != $(dom(g))")
    compose_id(f, g)
  end
end

@inline compose_id(f::SetFunction, g::SetFunction) = do_compose(f, g)
@inline compose_id(f::SetFunction, ::SetFunctionIdentity) = f
@inline compose_id(::SetFunctionIdentity, g::SetFunction) = g
@inline compose_id(f::SetFunctionIdentity, ::SetFunctionIdentity) = f

do_compose(f::SetFunction, g::SetFunction) = SetFunction(g∘f, dom(f), codom(g))
do_compose(f::SetFunction, c::ConstantFunction) =
  ConstantFunction(c.value, dom(f), codom(c))
do_compose(c::ConstantFunction, f::SetFunction) =
  ConstantFunction(f(c.value), dom(c), codom(f))
do_compose(c::ConstantFunction, d::ConstantFunction) =
  ConstantFunction(d.value, dom(c), codom(d))

@cartesian_monoidal_instance SetOb SetFunction

# Limits
########

limit(Xs::EmptyDiagram{<:TypeSet}) =
  Limit(Xs, SMultispan{0}(TypeSet(Nothing)))

universal(lim::Terminal{TypeSet{Nothing}}, span::SMultispan{0,<:SetOb}) =
  ConstantFunction(nothing, apex(span), ob(lim))

function limit(Xs::ObjectPair{<:TypeSet})
  X1, X2 = Xs
  X = TypeSet(Tuple{eltype(X1),eltype(X2)})
  π1, π2 = SetFunction(first, X, X1), SetFunction(last, X, X2)
  Limit(Xs, Span(π1, π2))
end

function universal(lim::BinaryProduct{<:TypeSet}, span::Span{<:SetOb})
  f, g = span
  SetFunction(x -> (f(x),g(x)), apex(span), ob(lim))
end

function limit(Xs::DiscreteDiagram{<:TypeSet})
  X = TypeSet(Tuple{map(eltype, Xs)...})
  πs = [ SetFunction(x -> getindex(x, i), X, Xi) for (i, Xi) in enumerate(Xs) ]
  Limit(Xs, Multispan(X, πs))
end

function universal(lim::Product{<:TypeSet}, span::Multispan{<:SetOb})
  @assert length(cone(lim)) == length(span)
  fs = Tuple(legs(span))
  SetFunction(x -> map(f -> f(x), fs), apex(span), ob(lim))
end

function limit(cospan::Multicospan{<:TypeSet})
  eltype(apex(cospan)) == Nothing ? product(feet(cospan)) :
    error("Pullbacks of TypeSets that are not products are not supported")
end

end


text:
module SpawnerRecruitModels

export
    AbstractSpawnerRecruitModel,
    BevertonHoltSpawnerRecruitModel,
    RickerSpawnerRecruitModel,
    LudwigWaltersSpawnerRecruitModel,
    CushingSpawnerRecruitModel,
    DerisoSchnuteSpawnerRecruitModel,
    ShepherdSpawnerRecruitModel,
    GammaSpawnerRecruitModel,
    recruit,
    max_recruits,
    max_spawnrecruits

abstract type AbstractSpawnerRecruitModel end

"""
    BevertonHoltSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel

Parameters:

- α: density-independent productivity parameter
- β: density-dependence parameter

Follows equation 3.6 (p. 88) of Quinn and Deriso. Asymptotic recruitment with
increasing spawner biomass.
"""
struct BevertonHoltSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T

    function BevertonHoltSpawnerRecruitModel(α::T, β::T) where T
        new{T}(α, β)
    end
end

function recruit(spawners, srmod::BevertonHoltSpawnerRecruitModel)
    srmod.α * spawners / (1 + srmod.β * spawners)
end

function max_recruits(srmod::BevertonHoltSpawnerRecruitModel)
    srmod.α / srmod.β
end

function max_spawnrecruits(srmod::BevertonHoltSpawnerRecruitModel)
    (Inf, max_recruits(srmod))
end

"""
    RickerSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- β: density-dependence parameter

Parameterized as in equation 3.8 of Quinn and Deriso. Dome-shaped recruitment
curve.
"""
struct RickerSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels
    α::T
    β::T

    function RickerSpawnerRecruitModel(α::T, β::T) where T
        new{T}(α, β)
    end
end

function recruit(spawners, srmod::RickerSpawnerRecruitModel)
    srmod.α * spawners * exp(-srmod.β * spawners)
end

function max_recruits(srmod::RickerSpawnerRecruitModel)
    srmod.α / (srmod.β * ℯ)
end

function max_spawnrecruits(srmod::RickerSpawnerRecruitModel)
    (1 / srmod.β, max_recruits(srmod))
end

"""
    LudwigWaltersSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- β: density-dependence parameter
- γ: density-dependence power parameter

Generalization of Ricker spawner recruit model. Parameterized as in equation 3.9
of Quinn and Deriso. Dome-shaped recruitment curve.
"""
struct LudwigWaltersSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T
    γ::T

    function LudwigWaltersSpawnerRecruitModel(α::T, β::T, γ::T) where T
        new{T}(α, β, γ)
    end
end

function recruit(spawners, srmod::LudwigWaltersSpawnerRecruitModel)
    srmod.α * spawners * exp(-srmod.β * spawners ^ srmod.γ)
end

function max_recruits(srmod::LudwigWaltersSpawnerRecruitModel)
    srmod.α * (srmod.β * srmod.γ) ^ -(1 / srmod.γ)
end

function max_spawnrecruits(srmod::LudwigWaltersSpawnerRecruitModel)
    (exp(-1 / srmod.γ), max_recruits(srmod))
end

"""
    CushingSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- γ: index of density dependence

Generalization of Ricker spawner recruit model. Parameterized as in equation 3.12
of Quinn and Deriso. Dome-shaped recruitment curve.
"""
struct CushingSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    γ::T

    function CushingSpawnerRecruitModel(α::T, γ::T) where T
        new{T}(α, γ)
    end
end

function recruit(spawners, srmod::CushingSpawnerRecruitModel)
    srmod.α * spawners ^ srmod.γ
end

max_recruits(srmod::CushingSpawnerRecruitModel) = Inf

max_spawnrecruits(srmod::CushingSpawnerRecruitModel) = (Inf, Inf)

"""
    DerisoSchnuteSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- β: optimality parameter
- γ: recruitment limitation or skewness parameter

Generalization of Beverton-Holt and Ricker spawner recruit models. Parameterized
as in equation 3.20 of Quinn and Deriso. May be asymptotic or dome-shaped
recruitment curve, depending on the value of γ.
"""
struct DerisoSchnuteSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T
    γ::T

    function DerisoSchnuteSpawnerRecruitModel(α::T, β::T, γ::T) where T
        new{T}(α, β, γ)
    end
end

function recruit(spawners, srmod::DerisoSchnuteSpawnerRecruitModel)
    srmod.α * spawners * (1 - srmod.β * srmod.γ * spawners) ^ (1 / srmod.γ)
end

function max_recruits(srmod::DerisoSchnuteSpawnerRecruitModel)
    (srmod.α / srmod.β) * (1 + srmod.γ) ^ -((1 + srmod.γ) / srmod.γ)
end

function max_spawnrecruits(srmod::DerisoSchnuteSpawnerRecruitModels)
    (1 / (srmod.β * (1 + srmod.γ)), max_recruits(srmod))
end

"""
    ShepherdSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- γ: index of density dependence

Generalization of Beverton-Holt, Ricker, and Cushing spawner recruit models.
Parameterized as in equation 3.21 of Quinn and Deriso. Unlimited recruitment for
γ < 1 as in the Cushing model, asymptotic recruitment (Beverton-Holt)
recruitment at γ = 1, and dome-shaped recruitment at γ > 1.
"""
struct ShepherdSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T
    γ::T

    function ShepherdSpawnerRecruitModel(α::T, β::T, γ::T) where T
        new{T}(α, β, γ)
    end
end

function recruit(spawners, srmod::ShepherdSpawnerRecruitModel)
    srmod.α * spawners / (1 + β * spawners ^ srmod.γ)
end

function max_recruits(srmod::ShepherdSpawnerRecruitModel)
    if srmod.γ > 1
        γm1 = srmod.γ - 1
        max_rec = srmod.α / (srmod.β * srmod.γ) * (srmod.β * γm1) ^ (γm1 / srmod.γ)
    elseif srmod.γ == 1
        max_rec = max_recruits(BevertonHoltSpawnerRecruitModel(srmod.α, srmod.β))
    else
        max_rec = Inf
    end
    max_rec
end

function max_spawnrecruits(srmod::ShepherdSpawnerRecruitModel)
    if srmod.γ > 1
        spawnmax = (srmod.β * (srmod.γ - 1)) ^ (-1 / srmod.γ)
    elseif srmod.γ ≤ 1
        spawnmax = Inf
    end
    (spawnmax, max_recruits(srmod))
end

"""
    GammaSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel

Parameters:

- α: productivity parameter
- β:
- γ:

Takes the form of an unnomalized gamma function. Contains the Ricker (γ = 1) and
Cushing (β = 0) as special cases. Approaches Beverton-Holt as γ → 0
"""
struct GammaSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T
    γ::T

    function GammaSpawnerRecruitModel(α::T, β::T, γ::T) where T
        γ > 0 || DomainError("γ must be positive for a valid spawner-recruit model")
        new{T}(α, β, γ)
    end
end

function recruit(spawners, srmod::GammaSpawnerRecruitModel)
    srmod.α * spawners ^ srmod.γ * exp(-srmod.β * spawners)
end

function max_recruits(srmod::GammaSpawnerRecruitModel)
    srmod.α * (srmod.γ / srmod.β) ^ srmod.γ * exp(-srmod.γ)
end

function max_spawnrecruits(srmod::GammaSpawnerRecruitModel)
    (srmod.γ / srmod.β, max_recruits(srmod))
end

end


text:
"""
    Filter([f])

Uses a function `f` to identify values, rows, columns or slices of data that should be
removed during an `apply` call. The default function `f` will removing `missing`s, or any
rows, columns or slices containing `missing`s.
"""
struct Filter{F<:Function}
    func::F
end

Filter() = Filter(_keep)

_keep(x) = !ismissing(x)
_keep(x::Union{Tuple, AbstractArray, IterTools.PropertyValues}) = all(_keep, x)

apply!(data::Vector, f::Filter) = Base.filter!(f.func, data)
function apply!(data::Vector{<:NamedTuple}, f::Filter; dims=:rows)
    d = dim(data, dims)
    d == 1 || throw(ArgumentError("Rowtables only support in-place filtering rowwise."))
    return Base.filter!(r -> f.func(propertyvalues(r)), data)
end

apply(data::Vector, f::Filter) = Base.filter(f.func, data)
function apply(data::Vector{<:NamedTuple}, f::Filter; dims=:rows)
    d = dim(data, dims)
    if d == 1
        return Base.filter(r -> f.func(propertyvalues(r)), data)
    else
        return materializer(data)(apply(Tables.columns(data), f; dims=dims))
    end
end

function apply(data::AbstractArray{Union{T, Missing}}, f::Filter; dims) where T
    d = dim(data, dims)
    mask = map(f.func, eachslice(data; dims=d))
    # use selectdim to reduce along dimension d using our mask, but call copy
    # because we don't want to return a view and collect may change the return type
    return copy(selectdim(data, d, mask))
end

apply(data::AbstractArray, f::Filter; kwargs...) = disallowmissing(data)

function apply(table, f::Filter; dims)
    istable(table) || throw(MethodError(apply, (table, f)))

    d = dim(table, dims)
    filtered = if d == 1
        Iterators.filter(Tables.rows(table)) do r
            f.func(propertyvalues(r))
        end
    else
        cols = Tables.columns(table)

        cnames = Iterators.filter(propertynames(cols)) do cname
            f.func(getproperty(cols, cname))
        end

        TableOperations.select(table, cnames...)
    end

    return materializer(table)(filtered)
end


text:
import DiffEqBase
import PyCall

PyCall.PyObject(::typeof(DiffEqBase.solve)) =
    PyCall.pyfunctionret(DiffEqBase.solve,Any,Vararg{PyCall.PyAny})

function DiffEqBase.numargs(f::PyCall.PyObject)
    inspect = PyCall.pyimport("inspect")
    PyCall.hasproperty(f,:py_func) ? _f = f.py_func : _f = f
    nargs = length(first(inspect.getfullargspec(_f)))
    if inspect.ismethod(_f)
        # `f` is a bound method (i.e., `self.f`) but `getfullargspec`
        # includes `self` in the `args` list.  So, we are subtracting
        # 1 manually here:
        return nargs - 1
    else
        return nargs
    end
end

max_stars_repo_name: jlapeyre/diffeqpy

max_issues_repo_name: jlapeyre/diffeqpy

max_forks_repo_name: jlapeyre/diffeqpy