text:
using ModelObjectsLH, ModelParams, Test

mdl = ModelParams;

function find_test()
    @testset "Find" begin
        m = init_test_model()
        @test check_own_fixed_params(m, get_pvector(m))
        @test check_own_calibrated_params(m, get_pvector(m))

        m2 = init_test_model();
        @test ModelParams.params_equal(m, m2)
    end
end

function change_values_test()
    @testset "Change values by hand" begin
        m = init_test_model();
        yNew = m.o2.y .+ 0.5;
        ModelParams.change_value!(m, :o2, :y, yNew);
        @test ModelParams.get_value(m, :o2, :y) .≈ yNew
        @test m.o2.y .≈ yNew

        pvec = get_pvector(m.o2);
        @test pvec isa ParamVector;
        p = pvec[1];
        @test p isa Param;
        if default_value(p) isa Real
            lb = -10.0;
            ub = 20.0;
        else
            sz = size(default_value(p));
            lb = fill(-10.0, sz);
            ub = fill(20.0, sz);
        end
        set_bounds!(m.o2, p.name; lb, ub);
        @test p.lb == lb;
        @test p.ub == ub;

    end
end

function compare_params_test()
    @testset "Compare params" begin
        m1 = init_test_model();
        m2 = init_test_model();
        pMiss1, pMiss2, pDiff = compare_params(m1, m2);
        @test isempty(pMiss1);
        @test isempty(pMiss2);
        @test isempty(pDiff);

        o2 = find_only_object(m1, :o2);
        allowedV = [get_object_id(o2)];
        @test check_params_match(m1, m2, allowedV);

        mdl.change_value!(get_pvector(o2), :a, 3.3);
        pMiss1, pMiss2, pDiff = compare_params(m1, m2);
        @test isempty(pMiss1);
        @test isempty(pMiss2);
        @test pDiff[get_object_id(o2)][:a] == [:value]
        @test check_params_match(m1, m2, allowedV);

        o3 = find_only_object(m1, :obj3);
        mdl.change_value!(get_pvector(o3), :x, 0.888);
        @test !check_params_match(m1, m2, allowedV);
    end
end

# function dict_test()
#     @testset "Dict/Vector" begin
#         m = init_test_model();
#         isCalibrated = true;

#         d1 = make_dict(m.o1.pvec, isCalibrated);
#         d2 = make_dict(m.o2.pvec, isCalibrated);

#         # Sync values for one object
#         ModelParams.sync_own_values!(m.o2);
#         @test ModelParams.check_calibrated_params(m.o2, m.o2.pvec);
#         @test ModelParams.check_fixed_params(m.o2, m.o2.pvec);

#         # The same step-by-step. Only needed for testing
#         vv = ModelParams.make_guess(m.o1.pvec, isCalibrated);
#         # update ++++++
#         # d11, nUsed1 = vector_to_dict(m.o1.pvec, vv, isCalibrated);
#         # @test d11[:x] ≈ d1[:x]
#         # @test d11[:y] ≈ d1[:y]

#         ModelParams.sync_own_from_vector!(m, vv);
#         # @test nUsed11 == nUsed1

#         # copy into param vector; then sync with model object
#         ModelParams.set_values_from_dict!(m.o1.pvec, d11);
#         ModelParams.set_own_values_from_pvec!(m.o1, isCalibrated);
#         @test m.o1.x ≈ d1[:x]
#         @test m.o1.y ≈ d1[:y]

#         # Sync calibrated model values with param vector
#         ModelParams.set_own_values_from_pvec!(m.o1, isCalibrated);
#         # also sync the non-calibrated default values
#         ModelParams.set_own_default_values!(m.o1, false);
#         @test ModelParams.check_calibrated_params(m.o1, m.o1.pvec);
#         @test ModelParams.check_fixed_params(m.o1, m.o1.pvec);
#     end
# end

function param_tables_test()
    @testset "Parameter table" begin
        m = init_test_model();
        d = param_tables(m, true);
        @test isa(d, Dict{ObjectId, ModelParams.ParamTable})

        pvecV = collect_pvectors(m);
        objIdV = [get_object_id(pvec)  for (_, pvec) ∈ pvecV];
        descrV = ["Description j"  for j = 1 : length(pvecV)];
        lineV = latex_param_table(m, true, objIdV, descrV);
        @test isa(lineV, Vector{String})
        @test startswith(lineV[1], "\\multicolumn")
	end
end

# function sync_test()
#     @testset "Sync" begin
#         m = init_test_model();

# 	end
# end

@testset "Model Objects" begin
    find_test();
    compare_params_test();
    change_values_test();
    param_tables_test();
    # sync_test();
end

# -----------------


text:
##### Show tendencies

export show_tendencies

using PrettyTables

first_param(t::TendencyDef{TT, PV}) where {TT, PV} = PV

function format_tend(tend_type, include_params)
    t = "$(nameof(typeof(tend_type)))"
    return include_params ? t * "{$(first_param(tend_type))}" : t
end

format_tends(tend_types, include_params) =
    "(" * join(format_tend.(tend_types, include_params), ", ") * ")"

"""
    show_tendencies(bl; include_params = false)

Show a table of the tendencies for each
prognostic variable for the given balance law.

Using `include_params = true` will include the type
parameters for each of the tendency fluxes
and sources.

Requires definitions for
 - [`prognostic_vars`](@ref)
 - [`eq_tends`](@ref)
for the balance law.
"""
function show_tendencies(bl; include_params = false)
    ip = include_params
    prog_vars = prognostic_vars(bl)
    if !isempty(prog_vars)
        header = [
            "Equation" "Flux{FirstOrder}" "Flux{SecondOrder}" "Source"
            "(Y_i)" "(F_1)" "(F_2)" "(S)"
        ]
        eqs = collect(string.(typeof.(prog_vars)))
        fmt_tends(tt) = map(prog_vars) do pv
            format_tends(eq_tends(pv, bl, tt), ip)
        end |> collect
        F1 = fmt_tends(Flux{FirstOrder}())
        F2 = fmt_tends(Flux{SecondOrder}())
        S = fmt_tends(Source())
        data = hcat(eqs, F1, F2, S)
        @warn "This table is temporarily incomplete"
        println("\nPDE: ∂_t Y_i + (∇•F_1(Y))_i + (∇•F_2(Y,G)))_i = (S(Y,G))_i")
        pretty_table(
            data,
            header,
            header_crayon = crayon"yellow bold",
            subheader_crayon = crayon"green bold",
            crop = :none,
        )
        println("")
    else
        msg = "Defining `prognostic_vars` and\n"
        msg *= "`eq_tends` for $(nameof(typeof(bl))) will\n"
        msg *= "enable printing a table of tendencies."
        @info msg
    end
    return nothing
end

ext: jl

max_forks_repo_name: jm-c/CLIMA


text:
###
# Distribution only based on its laplace transform function `f` and exponential tilting `c²`.
# The sampling is from Ridout, M. S. (2009). Generating random numbers from a distribution specified by its laplace transform. Statistics and Computing, 19(4):439.
# And the inverse laplace operation is done via Grassmann, W. K. (Ed.). (2000). Computational Probability. International Series in #Operations Research & Management Science. doi:10.1007/978-1-4757-4828-4
struct LaplaceTransformDistribution{T,TAlg} <:
       Distributions.ContinuousUnivariateDistribution
    f::Function # Laplace transform of the pdf
    c²::T # Exponential tilting parameter
    alg::TAlg # Algorithm to compute the inverse Laplace transform
    function LaplaceTransformDistribution{T,TAlg}(
        f::Function, c²::T, alg::TAlg
    ) where {T<:Real,TAlg}
        _check_f(f) || error("The function passed is not valid") # Do a series of checks on f
        c² >= 0 || error("c² has to a be non-negative real")
        return new{T,TAlg}(f, c², alg)
    end
end

function LaplaceTransformDistribution(
    f::Function, c²::T=0.0, alg::TAlg=BromwichInverseLaplace()
) where {T,TAlg}
    return LaplaceTransformDistribution{T,TAlg}(f, c², alg)
end

function _check_f(f)
    return true # TODO Add tests for complete monotonicity / PDR
end
_gradf(d::LaplaceTransformDistribution, x::Real) = first(ForwardDiff.gradient(d.f, [x]))
function _gradlogf(d::LaplaceTransformDistribution, x::Real)
    return first(ForwardDiff.gradient(log ∘ d.f, [x]))
end
function _hessianlogf(d::LaplaceTransformDistribution, x::Real)
    return first(ForwardDiff.hessian(log ∘ d.f, [x]))
end

Distributions.pdf(dist::LaplaceTransformDistribution, x::Real) = apply_f(dist, x)
function Distributions.mean(dist::LaplaceTransformDistribution)
    return _gradf(dist, dist.c²) / dist.f(dist.c²)
end
function Distributions.var(dist::LaplaceTransformDistribution)
    return _hessianlogf(dist, dist.c²) / dist.f(dist.c²) - mean(dist)^2
end

function Random.rand(dist::LaplaceTransformDistribution)
    return first(rand(dist, 1))
end

# Sampling from Ridout (09)
function Random.rand(dist::LaplaceTransformDistribution, n::Int)
    nTries = 100
    i = 0
    while i < nTries
        try
            if i > 10
                @warn "This is taking more than 10 tries! Trie $i/$nTries!"
            end
            return laptrans(dist; n=n)
        catch e
            if e isa AssertionError
                @warn "$(e.msg), Trying again"
                i += 1
                continue
            else
                rethrow(e)
            end
        end
    end
    @error "Sampler failed to converge"
end

struct BromwichInverseLaplace{T}
    A::Int
    l::Int
    m::Int
    n::Int
    expA2l::T
    ijπl::Vector{Complex{T}}
    expijπl::Vector{Complex{T}}
    coefs::Vector{T}
    altern::Vector{Int}
    b::Vector{T}
    s::Vector{T}
    function BromwichInverseLaplace(A::Int=19, l::Int=1, m::Int=11, n::Int=28)
        @assert A > 0 "A should be a positive integer"
        @assert l > 0 "l should be a positive integer"
        @assert m > 0 "m should be a positive integer"
        @assert n > 0 "n should be a positive integer"
        T = typeof(exp(A / (2 * l)))
        ijπl = 1im .* (1:l) .* π / l
        return new{T}(
            A,
            l,
            m,
            n,
            exp(A / (2 * l)),
            ijπl,
            exp.(ijπl),
            binomial.(m, 0:m) ./ (2^m),
            (-1) .^ (0:(n + m)),
            zeros(n + m + 1),
            zeros(n + m + 1),
        )
    end
end

@inline function apply_f(dist::LaplaceTransformDistribution, t)
    if t == 0
        return zero(t)
    else
        invlaplace(dist.f, t, dist.alg)
    end
end

@inline function apply_F(dist::LaplaceTransformDistribution, t)
    if t == 0
        return zero(t)
    else
        invlaplace(s -> dist.f(s) / s, t, dist.alg)
    end
end

function invlaplace(f::Function, t::Real, A::Int=19, l::Int=1, m::Int=11, n::Int=28)
    # n is the burnin, m is the number of components, A and l aare discretization parameters
    alg = BromwichInverseLaplace(A, l, m, n)
    return invlaplace(f, t, alg)
end

function invlaplace(f::Function, t::Real, alg::BromwichInverseLaplace)
    scaled_t = 1 / (2 * alg.l * t)
    alg.b[1] =
        f(alg.A * scaled_t) +
        2 * sum(real.(f.(alg.A * scaled_t .+ alg.ijπl ./ t) .* alg.expijπl))
    @inbounds for k in 1:(alg.n + alg.m)
        alg.b[k + 1] =
            2 * sum(
                real.(
                    f.(alg.A * scaled_t .+ alg.ijπl ./ t .+ 1im * π * k / t) .* alg.expijπl
                ),
            )
    end
    acoeff = alg.expA2l * scaled_t
    cumsum!(alg.s, acoeff .* alg.b .* alg.altern)
    return dot(view(alg.s, (alg.n + 1):length(alg.s)), alg.coefs)
end

# Laplace transform implemented from "Computational Probability Grassmann 2000"
function laptrans(
    dist::LaplaceTransformDistribution;
    n::Int=10,
    jmax::Int=500,
    kmax::Int=1000,
    b::Real=2.0,
    τ::Real=1e-7,
)
    # Step 1
    u = sort!(rand(n))
    # Step 2
    xmax = 0.1
    j = 0
    while apply_F(dist, xmax) < u[end] && j < jmax
        xmax = b * xmax
        j += 1
    end
    @assert j != jmax "Xmax search failed to converge"
    x = zeros(n)
    # Step 3
    for i in 1:n
        x_L = i == 1 ? 0.0 : x[i - 1]
        x_U = xmax
        k = 0
        t = x_L
        while abs(apply_F(dist, t) - u[i]) > τ && k < kmax
            k += 1
            t = t - (apply_F(dist, t) - u[i]) / apply_f(dist, t)
            if t < x_L || t > x_U
                t = 0.5 * (x_L + x_U)
            end
            if apply_F(dist, t) <= u[i]
                x_L = t
            else
                x_U = t
            end
        end
        @assert k != kmax "Modified Newton-Raphson procedure failed to converge for i=$i, x[$i]=$t"
        x[i] = t
    end
    # Step 4
    return shuffle!(x)
end


text:
# Return the index of the maximum entry of a given vector transformed by f.
Base.argmax(f::Any, v::AbsVec) = argmax(f, v, 1:length(v))

# Among specified indices, return the index of the maximum entry of a given vector
# transformed by f.
function Base.argmax(f::Any, v::AbsVec, indv::AbsVecInteger)
    ind = 0  # return 0 if indv is empty
    val = -Inf
    for n = indv
        if f(v[n]) ≥ val  # ≥ rather than > to prevent returning 0 when v has -Inf
            val = f(v[n])
            ind = n
        end
    end

    return ind
end

# Return the index of the minimum entry of a given vector transformed by f
Base.argmin(f::Any, v::AbsVec) = argmin(f, v, 1:length(v))

# Among specified indices, return the index of the minimum entry of a given vector
# transformed by f.
function Base.argmin(f::Any, v::AbsVec, indv::AbsVecInteger)
    ind = 0  # return 0 if v is empty
    val = Inf
    for n = indv
        if f(v[n]) ≤ val  # ≤ rather than < to prevent returning 0 when v has -Inf
            val = f(v[n])
            ind = n
        end
    end

    return ind
end


text:
using Base.Test

const MODE_LOG_FILE = tempname()
println("Path to MODE_LOG_FILE: $MODE_LOG_FILE")

configure(; modes = ["debug", "info", "warn", "error", "crazy"])
add_truck(Lumberjack.LumberjackTruck(MODE_LOG_FILE, "info"), "infotruck")
add_saw(Lumberjack.Saw(Lumberjack.fn_call_saw, "error"))

# Test limiting trucks and saws to certain log levels.
@noinline caller(mode, msg) = log(mode, msg)
caller("debug", "Message")  # Not logged.
caller("info", "Message")   # Logged.
caller("warn", "Message")   # Logged.
caller("error", "Message")  # Logged with function call.
caller("crazy", "Message")  # Logged with function call.

remove_truck("infotruck")
log_lines = readlines(open(MODE_LOG_FILE, "r"))

fn_regex = "caller@\\Q$(basename(@__FILE__))\\E:\\d+"
@test log_lines[1] == "info: Message\n"
@test log_lines[2] == "warn: Message\n"
@test ismatch(Regex(fn_regex * ".+error: Message"), log_lines[3])
@test ismatch(Regex(fn_regex * ".+crazy: Message"), log_lines[4])

# Clean up.
@test success(`rm $MODE_LOG_FILE`)


text:
using Test
using MatrixProductStates


text:
using Zygote, NNlib, Test, Random, LinearAlgebra, Statistics, FillArrays, FFTW
using Zygote: gradient
using NNlib: conv, ∇conv_data, depthwiseconv
using Base.Broadcast: broadcast_shape

function ngradient(f, xs::AbstractArray...)
  grads = zero.(xs)
  for (x, Δ) in zip(xs, grads), i in 1:length(x)
    δ = sqrt(eps())
    tmp = x[i]
    x[i] = tmp - δ/2
    y1 = f(xs...)
    x[i] = tmp + δ/2
    y2 = f(xs...)
    x[i] = tmp
    Δ[i] = (y2-y1)/δ
  end
  return grads
end

gradcheck(f, xs...) =
  all(isapprox.(ngradient(f, xs...),
                gradient(f, xs...), rtol = 1e-5, atol = 1e-5))

gradtest(f, xs::AbstractArray...) = gradcheck((xs...) -> sum(sin.(f(xs...))), xs...)
gradtest(f, dims...) = gradtest(f, rand.(Float64, dims)...)

Random.seed!(0)

@test gradient(//, 2, 3) === (1//3, -2//9)

@test gradtest((x, W, b) -> σ.(W*x .+ b), 5, (2,5), 2)
@test gradtest((x, W, b) -> σ.(W*x .+ b), (5,3), (2,5), 2)
@test gradtest((x, W, b) -> logσ.(W*x .+ b), 5, (2,5), 2)
@test gradtest((x, W, b) -> logσ.(W*x .+ b), (5,3), (2,5), 2)

@test gradtest((w, x) -> w'*x, randn(10, 2), randn(10))
@test gradtest((w, x) -> transpose(w)*x, randn(5,5), randn(5,5))

@test gradtest(x -> sum(x, dims = (2, 3)), (3,4,5))
@test gradtest(x -> sum(abs2, x), randn(4, 3, 2))
@test gradtest(x -> sum(abs2, x; dims=1), randn(4, 3, 2))
@test gradtest(x -> prod(x, dims = (2, 3)), (3,4,5))
@test gradtest(x -> prod(x), (3,4,5))

@test gradtest(x -> softmax(x).*(1:3), 3)
@test gradtest(x -> softmax(x).*(1:3), (3,5))
@test gradtest(x -> logsoftmax(x).*(1:3), 3)
@test gradtest(x -> logsoftmax(x).*(1:3), (3,5))

@test gradtest(x -> x', rand(5))

@test gradtest(det, (4, 4))
@test gradtest(logdet, map(x -> x*x', (rand(4, 4),))[1])
@test gradtest(x -> logabsdet(x)[1], (4, 4))

@test gradtest(x -> view(x,:,2,:), (3,4,5))
@test gradtest(x -> view(x,1:2,3:4), (3,4))

@testset "conv" begin
  for spatial_rank in (1, 2, 3)
    x = rand(repeat([10], spatial_rank)..., 3, 2)
    w = rand(repeat([3], spatial_rank)..., 3, 3)
    cdims = DenseConvDims(x, w)
    @test gradtest((x, w) -> conv(x, w, cdims), x, w)
    y = conv(x, w, cdims)
    @test gradtest((y, w) -> ∇conv_data(y, w, cdims), y, w)
    dcdims = DepthwiseConvDims(x, w)
    @test gradtest((x, w) -> depthwiseconv(x, w, dcdims), x, w)
  end
end

@testset "pooling" begin
  for spatial_rank in (1, 2)
    x = rand(repeat([10], spatial_rank)..., 3, 2)
    pdims = PoolDims(x, 2)
    @test gradtest(x -> maxpool(x, pdims), x)
    @test gradtest(x -> meanpool(x, pdims), x)
  end
end

@test gradtest(x -> permutedims(x), rand(2))
@test gradtest(x -> permutedims(x), rand(2,3))
@test gradtest(x -> permutedims(x, [3,1,2]), rand(4,5,6))
@test gradtest(x -> PermutedDimsArray(x, (3,1,2)), rand(4,5,6))
let
  y, back = Zygote.pullback(permutedims, randn(3))
  @test first(back(randn(1, 3))) isa Vector
end

@test gradtest(x -> repeat(x; inner=2), rand(5))
@test gradtest(x -> repeat(x; inner=2, outer=3), rand(5))
@test gradtest(x -> repeat(x; inner=(2,2,1), outer=(1,1,3)), rand(5,4,3))

@test gradtest(tr, rand(4, 4))

@testset "fill" begin
  rng, N, M, P = MersenneTwister(123456), 11, 6, 5
  @test gradtest(x->fill(first(x), N), randn(rng, 1))
  @test gradtest(x->fill(first(x), N, M), randn(rng, 1))
  @test gradtest(x->fill(first(x), N, M, P), randn(rng, 1))
end

@testset "dot" begin
  rng = MersenneTwister(123456)
  @test gradtest((x, y)->dot(x[1], y[1]), [randn(rng)], [randn(rng)])
  @test gradtest(dot, randn(rng, 10), randn(rng, 10))
  @test gradtest(dot, randn(rng, 10, 3), randn(rng, 10, 3))
end

@test gradtest(kron, rand(5), rand(3))
@test gradtest(kron, rand(5), rand(3), rand(8))
@test gradtest(kron, rand(5,1), rand(3,1))
@test gradtest(kron, rand(5,1), rand(3,1), rand(8,1))
@test gradtest(kron, rand(5,2), rand(3,2), rand(8,2))

@testset "map" begin
  @test gradtest(xs -> sum(map(x -> x^2, xs)), rand(2,3))
  @test gradtest((xss...) -> sum(map((xs...) -> sqrt(sum(xs.^2)), xss...)), [rand(5) for _ in 1:6]...)
  function foo(y)
    bar = (x) -> x*y
    sum(map(bar, 1:5))
  end
  @test gradtest(foo, 3)
  @test gradient(v -> sum([x for x in v]), [1.1,2.2,3.3]) == ([1, 1, 1],)
end

@testset "mean" begin
  @test gradtest(mean, rand(2, 3))

  @test gradtest(x -> mean(x, dims=1), rand(2, 3))
  @test gradtest(x -> mean(x, dims=2), rand(2, 3))
  @test gradtest(x -> mean(x, dims=3), rand(2, 3, 4))

  @test gradtest(x -> mean(x, dims=[1, 2]), rand(2, 3, 4))
end

@testset "var" begin
  @test gradtest(var, rand(2, 3))
  @test gradtest(x -> var(x, dims=1), rand(2, 3))
  @test gradtest(x -> var(x, dims=2), rand(2, 3))
  @test gradtest(x -> var(x, dims=3), rand(2, 3, 4))
  @test gradtest(x -> var(x, dims=[1, 2]), rand(2, 3, 4))

  @test gradtest(x -> var(x, corrected=false), rand(2, 3))
  @test gradtest(x -> var(x, dims=1, corrected=false), rand(2, 3))
  @test gradtest(x -> var(x, dims=2, corrected=false), rand(2, 3))
  @test gradtest(x -> var(x, dims=3, corrected=false), rand(2, 3, 4))
  @test gradtest(x -> var(x, dims=[1, 2], corrected=false), rand(2, 3, 4))

  @test gradtest(x -> var(x, mean=mean(x)), rand(2, 3))
  @test gradtest(x -> var(x, dims=1, mean=mean(x, dims=1)), rand(2, 3))
  @test gradtest(x -> var(x, dims=2, mean=mean(x, dims=2)), rand(2, 3))
  @test gradtest(x -> var(x, dims=3, mean=mean(x, dims=3)), rand(2, 3, 4))
  @test gradtest(x -> var(x, dims=[1, 2], mean=mean(x, dims=[1, 2])), rand(2, 3, 4))

  @test gradtest(x -> var(x, corrected=false, mean=mean(x)), rand(2, 3))
  @test gradtest(x -> var(x, dims=1, corrected=false, mean=mean(x, dims=1)), rand(2, 3))
  @test gradtest(x -> var(x, dims=2, corrected=false, mean=mean(x, dims=2)), rand(2, 3))
  @test gradtest(x -> var(x, dims=3, corrected=false, mean=mean(x, dims=3)), rand(2, 3, 4))
  @test gradtest(x -> var(x, dims=[1, 2], corrected=false, mean=mean(x, dims=[1, 2])), rand(2, 3, 4))
end

@testset "std" begin
  @test gradtest(std, rand(2, 3))
  @test gradtest(x -> std(x, dims=1), rand(2, 3))
  @test gradtest(x -> std(x, dims=2), rand(2, 3))
  @test gradtest(x -> std(x, dims=3), rand(2, 3, 4))
  @test gradtest(x -> std(x, dims=[1, 2]), rand(2, 3, 4))

  @test gradtest(x -> std(x, corrected=false), rand(2, 3))
  @test gradtest(x -> std(x, dims=1, corrected=false), rand(2, 3))
  @test gradtest(x -> std(x, dims=2, corrected=false), rand(2, 3))
  @test gradtest(x -> std(x, dims=3, corrected=false), rand(2, 3, 4))
  @test gradtest(x -> std(x, dims=[1, 2], corrected=false), rand(2, 3, 4))

  @test gradtest(x -> std(x, mean=mean(x)), rand(2, 3))
  @test gradtest(x -> std(x, dims=1, mean=mean(x, dims=1)), rand(2, 3))
  @test gradtest(x -> std(x, dims=2, mean=mean(x, dims=2)), rand(2, 3))
  @test gradtest(x -> std(x, dims=3, mean=mean(x, dims=3)), rand(2, 3, 4))
  @test gradtest(x -> std(x, dims=[1, 2], mean=mean(x, dims=[1, 2])), rand(2, 3, 4))

  @test gradtest(x -> std(x, corrected=false, mean=mean(x)), rand(2, 3))
  @test gradtest(x -> std(x, dims=1, corrected=false, mean=mean(x, dims=1)), rand(2, 3))
  @test gradtest(x -> std(x, dims=2, corrected=false, mean=mean(x, dims=2)), rand(2, 3))
  @test gradtest(x -> std(x, dims=3, corrected=false, mean=mean(x, dims=3)), rand(2, 3, 4))
  @test gradtest(x -> std(x, dims=[1, 2], corrected=false, mean=mean(x, dims=[1, 2])), rand(2, 3, 4))
end

@testset "maximum" begin
  @test gradtest(maximum, rand(2, 3))

  @test gradtest(x -> maximum(x, dims=1), rand(2, 3))
  @test gradtest(x -> maximum(x, dims=2), rand(2, 3))
  @test gradtest(x -> maximum(x, dims=3), rand(2, 3, 4))

  @test gradtest(x -> maximum(x, dims=[1, 2]), rand(2, 3, 4))

  @test gradient(x -> 1 / maximum(x), [1., 2, 3])[1] == [0, 0, -1/9]
end

@testset "minimum" begin
  @test gradtest(minimum, rand(2, 3))

  @test gradtest(x -> minimum(x, dims=1), rand(2, 3))
  @test gradtest(x -> minimum(x, dims=2), rand(2, 3))
  @test gradtest(x -> minimum(x, dims=3), rand(2, 3, 4))

  @test gradtest(x -> minimum(x, dims=[1, 2]), rand(2, 3, 4))
end

@testset "dropdims" begin
  @test gradtest(x -> dropdims(x, dims = 3), rand(2, 2, 1, 2))
  @test gradtest(x -> dropdims(x, dims = (2, 3)), rand(2, 1, 1, 3))
  @test gradtest(x -> dropdims(x, dims = (1, 2, 3)), rand(1, 1, 1, 3))
end

@testset "$f(::AbstractArray)" for f in (real, conj, imag)
  rng, N = MersenneTwister(123456), 3
  Ts = (Float64, ComplexF64)
  @testset "$f(::Array{$IT})" for IT in Ts
    A = randn(IT, N, N)
    y, back = Zygote.pullback(f, A)
    y2, back2 = Zygote.pullback(x->f.(x), A)
    @test y == y2
    @testset "back(::Array{$BT})" for BT in Ts
      ȳ = randn(BT, N, N)
      @test back(ȳ)[1] == back2(ȳ)[1]
    end
  end
end

@testset "(p)inv" begin
  rng, P, Q = MersenneTwister(123456), 13, 11
  A, B, C = randn(rng, P, Q), randn(rng, P, P), randn(Q, P)
  @test gradtest(pinv, A)
  @test gradtest(inv, B)
  @test gradtest(pinv, C)
end

@testset "multiplication" begin
  @testset "matrix-matrix" begin
    rng, M, P, Q = MersenneTwister(123456), 13, 7, 11
    @test gradtest(*, randn(rng, M, P), randn(rng, P, Q))
    @test gradtest(*, randn(rng, M, P), randn(rng, P))
    @test gradtest(*, randn(rng, M, 1), randn(rng, 1, Q))
    @test gradtest(*, randn(rng, M), randn(rng, 1, Q))
    @test gradtest(*, randn(rng, 10)', randn(rng, 10))

    let
      y, back = Zygote.pullback(*, randn(rng, M, P), randn(rng, P))
      @test last(back(randn(rng, M))) isa Vector
    end
    let
      y, back = Zygote.pullback(*, randn(rng, M), randn(rng, 1, P))
      @test first(back(randn(rng, M, P))) isa Vector
    end
  end
end

@testset "backsolve" begin
  rng, M, P, Q = MersenneTwister(123456), 13, 10, 9
  X, Y, y = randn(rng, P, P), randn(rng, P, Q), randn(rng, P)
  A, B = randn(rng, P, M), randn(P, Q)
  D = collect(Diagonal(randn(rng, P)))
  L = collect(LowerTriangular(randn(rng, P, P)))
  L[diagind(L)] .= 1 .+ 0.01 .* randn(rng, P)
  U = collect(UpperTriangular(randn(rng, P, P)))
  U[diagind(U)] .= 1 .+ 0.01 .* randn(rng, P)

  # \ (Dense square)
  @test gradtest(\, X, Y)
  @test gradtest(\, X, y)

  # \ (Dense rectangular)
  @test gradtest(\, A, Y)
  @test gradtest(\, A, y)
  @test gradtest(\, B, Y)
  @test gradtest(\, B, y)

  # \ (Diagonal)
  @test gradtest(\, D, Y)
  @test gradtest(\, D, y)
  @test gradtest((D, Y)-> Diagonal(D) \ Y, D, Y)
  @test gradtest((D, Y)-> Diagonal(D) \ Y, D, y)

  # \ (LowerTriangular)
  @test gradtest(\, L, Y)
  @test gradtest(\, L, y)
  @test gradtest((L, Y) -> LowerTriangular(L) \ Y, L, Y)
  @test gradtest((L, Y) -> LowerTriangular(L) \ Y, L, y)

  # \ (UpperTriangular)
  @test gradtest(\, U, Y)
  @test gradtest(\, U, y)
  @test gradtest((U, Y) -> UpperTriangular(U) \ Y, U, Y)
  @test gradtest((U, Y) -> UpperTriangular(U) \ Y, U, y)

  # /
  @test gradtest(/, Y', X)
  @test gradtest((y, X)->y' / X, y, X)

  # / (rectangular)
  @test gradtest(/, Y', A')
  @test gradtest((y, A)->y' / A', y, A)
  @test gradtest(/, Y', B')
  @test gradtest((y, A)->y' / A', y, B)

  # / (Diagonal)
  @test gradtest((D, Y) -> Y' / D, D, Y)
  @test gradtest((D, Y) -> Y' / D, D, y)
  @test gradtest((D, Y)-> Y' / Diagonal(D), D, Y)
  @test gradtest((D, Y)-> Y' / Diagonal(D), D, y)

  # / (LowerTriangular)
  @test gradtest((L, Y) -> Y' / L, L, Y)
  @test gradtest((L, Y) -> Y' / L, L, y)
  @test gradtest((L, Y) -> Y' / LowerTriangular(L), L, Y)
  @test gradtest((L, Y) -> Y' / LowerTriangular(L), L, y)

  # / (UpperTriangular)
  @test gradtest((U, Y) -> Y' / U, U, Y)
  @test gradtest((U, Y) -> Y' / U, U, y)
  @test gradtest((U, Y) -> Y' / UpperTriangular(U), U, Y)
  @test gradtest((U, Y) -> Y' / UpperTriangular(U), U, y)

  @testset "Cholesky" begin

    # Check that the forwards pass computes the correct thing.
    @test Zygote.pullback(X->cholesky(X * X' + I) \ Y, X)[1] == cholesky(X * X' + I) \ Y
    @test gradtest(X->cholesky(X * X' + I) \ Y, X)
    @test gradtest(Y->cholesky(X * X' + I) \ Y, Y)
    @test gradtest(X->cholesky(X * X' + I) \ y, X)
    @test gradtest(y->cholesky(X * X' + I) \ y, y)
  end
end

@testset "Symmetric" begin
  @testset "real" begin
    rng, P = MersenneTwister(123456), 7
    A = randn(rng, P, P)
    @testset "uplo=$uplo" for uplo in (:U, :L)
      @test gradtest(x->Symmetric(x, uplo), A)
      y, back = Zygote.pullback(Symmetric, A, uplo)
      @test y isa Symmetric

      @testset "back(::Diagonal)" begin
        D̄ = Diagonal(randn(rng, P))
        @test back(Diagonal(D̄))[1] isa Diagonal
        @test back(Diagonal(D̄))[1] ≈ back(Matrix(D̄))[1]
      end

      @testset "back(::$TTri)" for TTri in (LowerTriangular,UpperTriangular)
        D̄ = TTri(randn(rng, P, P))
        @test back(D̄)[1] isa Matrix
        @test back(D̄)[2] === nothing
        @test back(D̄)[1] ≈ back(Matrix(D̄))[1]
      end
    end
  end

  @testset "complex" begin
    rng, P = MersenneTwister(123456), 7
    Re = randn(rng, P, P)
    Im = randn(rng, P, P)
    A = complex.(Re, Im)

    @testset "gradcheck dense" begin
      for uplo in (:U, :L)
        @test gradcheck(Re,Im) do a, b
          c = Symmetric(complex.(a, b), uplo)
          d = exp.(c)
          sum(real.(d) + imag.(d))
        end
      end
    end

    @testset "uplo=$uplo" for uplo in (:U, :L)
      y, back = Zygote.pullback(Symmetric, A, uplo)
      @test y isa Symmetric

      @testset "back(::Diagonal)" begin
        D̄ = Diagonal(complex.(randn(rng, P), randn(rng, P)))
        @test back(Diagonal(D̄))[1] isa Diagonal
        @test back(Diagonal(D̄))[1] ≈ back(Matrix(D̄))[1]
      end

      @testset "back(::$TTri)" for TTri in (LowerTriangular,UpperTriangular)
        D̄ = TTri(complex.(randn(rng, P, P), randn(rng, P, P)))
        @test back(D̄)[1] isa Matrix
        @test back(D̄)[2] === nothing
        @test back(D̄)[1] ≈ back(Matrix(D̄))[1]
      end
    end
  end
end

@testset "Hermitian" begin
  rng, P = MersenneTwister(123456), 7
  Re = randn(rng, P, P)
  Im = randn(rng, P, P)
  A = complex.(Re, Im)

  @testset "gradcheck dense" begin
    for uplo in (:U, :L)
      @test gradcheck(Re,Im) do a, b
        c = Hermitian(complex.(a, b), uplo)
        d = exp.(c)
        sum(real.(d) + imag.(d))
      end
    end
  end

  @testset "uplo=$uplo" for uplo in (:U, :L)
    y, back = Zygote.pullback(Hermitian, A, uplo)
    _, back_sym = Zygote.pullback(Symmetric, A, uplo)
    @test y isa Hermitian

    @testset "back" begin
      D̄ = randn(rng, P, P)
      @test back(D̄)[1] ≈ back_sym(D̄)[1]
    end

    @testset "back(::Diagonal)" begin
      D̄ = Diagonal(complex.(randn(rng, P), randn(rng, P)))
      @test back(D̄)[1] isa Diagonal
      @test back(D̄)[2] === nothing
      @test back(D̄)[1] ≈ back(Matrix(D̄))[1]
      @test back(real(D̄))[1] ≈ back_sym(real(D̄))[1]
    end

    @testset "back(::$TTri)" for TTri in (LowerTriangular,UpperTriangular)
      D̄ = TTri(complex.(randn(rng, P, P), randn(rng, P, P)))
      @test back(D̄)[1] isa Matrix
      @test back(D̄)[2] === nothing
      @test back(D̄)[1] ≈ back(Matrix(D̄))[1]
      @test back(real(D̄))[1] ≈ back_sym(real(D̄))[1]
    end
  end
end

@testset "diag" begin
  rng, P = MersenneTwister(123456), 10
  A = randn(rng, P, P)
  @test gradtest(diag, A)
end

@testset "Diagonal" begin
  rng, P = MersenneTwister(123456), 10
  d = randn(rng, P)
  @test gradtest(Diagonal, d)
  y, back = Zygote.pullback(Diagonal, d)
  D̄ = randn(rng, P, P)
  @test back(D̄) == back(Diagonal(D̄))
  @test back(D̄) == back((diag=diag(D̄),))
end

@testset "dense + UniformScaling" begin
  rng = MersenneTwister(123456)
  A, λ = randn(rng, 10, 10), randn(rng)
  @test gradtest(A->A + 5I, A)
  @test gradtest(λ->A + λ[1] * I, [λ])
end

@testset "cholesky" begin
  @testset "cholesky - dense" begin
    rng, N = MersenneTwister(123456), 5
    A = randn(rng, N, N)
    @test cholesky(A' * A + I) == first(Zygote.pullback(A->cholesky(A' * A + I), A))
    @test gradtest(A->cholesky(A' * A + I).U, A)
    @test gradtest(A->logdet(cholesky(A' * A + I)), A)
    @test gradtest(B->cholesky(Symmetric(B)).U, A * A' + I)
    @test gradtest(B->logdet(cholesky(Symmetric(B))), A * A' + I)
  end
  @testset "cholesky - scalar" begin
    rng = MersenneTwister(123456)
    y, back = Zygote.pullback(cholesky, 5.0 * ones(1, 1))
    y′, back′ = Zygote.pullback(cholesky, 5.0)
    C̄ = randn(rng, 1, 1)
    @test back′((factors=C̄,))[1] isa Real
    @test back′((factors=C̄,))[1] ≈ back((factors=C̄,))[1][1, 1]
  end
  @testset "cholesky - Diagonal" begin
    rng, N = MersenneTwister(123456), 3
    D = Diagonal(exp.(randn(rng, N)))
    Dmat = Matrix(D)
    y, back = Zygote.pullback(cholesky, Dmat)
    y′, back′ = Zygote.pullback(cholesky, D)
    C̄ = (factors=randn(rng, N, N),)
    @test back′(C̄)[1] isa Diagonal
    @test diag(back′(C̄)[1]) ≈ diag(back(C̄)[1])
  end
end

@testset "lyap" begin
  rng, N = MersenneTwister(6865943), 5
  for i = 1:5
    A = randn(rng, N, N)
    C = randn(rng, N, N)
    @test gradtest(lyap, A, C)
  end
  @test gradcheck(x->lyap(x[1],x[2]),[3.1,4.6])
end

@testset "matrix exponential" begin
  @testset "real dense" begin
    rng, N = MersenneTwister(6865931), 8
    for i = 1:5
      A = randn(rng, N, N)
      @test gradtest(exp, A)
    end
  end

  @testset "complex dense" begin
    rng, N = MersenneTwister(6865931), 8
    for i = 1:5
      A = randn(rng, N, N)
      B = randn(rng, N, N)
      @test gradcheck(A, B) do a,b
        c = complex.(a, b)
        d = exp(c)
        return sum(real.(d) + 2 .* imag.(d))
      end
    end
  end
end

@testset "eigen(::RealHermSymComplexHerm)" begin
  @testset "eigen(::Symmetric{<:Real})" begin
    rng, N = MersenneTwister(123), 7
    A = Symmetric(randn(rng, N, N))
    @test gradtest(collect(A)) do (x)
      d, Q = eigen(Symmetric(x))
      return Q * Diagonal(exp.(d)) * transpose(Q)
    end
    y = Zygote.pullback(eigen, A)[1]
    y2 = eigen(A)
    @test y.values ≈ y2.values
    @test y.vectors ≈ y2.vectors
    @testset "low rank" begin
      U = eigvecs(A)
      A2 = Symmetric(U * Diagonal([randn(rng), zeros(N-1)...]) * U')
      @test_broken gradtest(collect(A2)) do (x)
        d, Q = eigen(Symmetric(x))
        return Q * Diagonal(exp.(d)) * transpose(Q)
      end
    end
  end

  @testset "eigen(::Hermitian{<:Real})" begin
    rng, N = MersenneTwister(456), 7
    A = Hermitian(randn(rng, N, N))
    @test gradtest(collect(A)) do (x)
      d, Q = eigen(Hermitian(x))
      return Q * Diagonal(exp.(d)) * transpose(Q)
    end
    y = Zygote.pullback(eigen, A)[1]
    y2 = eigen(A)
    @test y.values ≈ y2.values
    @test y.vectors ≈ y2.vectors
    @testset "low rank" begin
      U = eigvecs(A)
      A2 = Hermitian(U * Diagonal([randn(rng), zeros(N-1)...]) * U')
      @test_broken gradtest(collect(A2)) do (x)
        d, Q = eigen(Hermitian(x))
        return Q * Diagonal(exp.(d)) * transpose(Q)
      end
    end
  end

  @testset "eigen(::Hermitian{<:Complex})" begin
    rng, N = MersenneTwister(789), 7
    A = Hermitian(randn(rng, ComplexF64, N, N))
    @test gradtest(reim(collect(A))...) do a,b
      d, U = eigen(Hermitian(complex.(a, b)))
      X = U * Diagonal(exp.(d)) * U'
      return real.(X) .+ imag.(X)
    end
    y = Zygote.pullback(eigen, A)[1]
    y2 = eigen(A)
    @test y.values ≈ y2.values
    @test y.vectors ≈ y2.vectors
    @testset "low rank" begin
      U = eigvecs(A)
      A2 = Hermitian(U * Diagonal([randn(rng), zeros(N-1)...]) * U')
      @test_broken gradtest(reim(collect(A2))...) do a,b
        d, U = eigen(Hermitian(complex.(a, b)))
        X = U * Diagonal(exp.(d)) * U'
        return real.(X) .+ imag.(X)
      end
    end
  end
end

@testset "eigvals(::RealHermSymComplexHerm)" begin
  @testset "eigvals(::Symmetric{<:Real})" begin
    rng, N = MersenneTwister(123), 7
    A = Symmetric(randn(rng, N, N))
    @test gradtest(x->eigvals(Symmetric(x)), collect(A))
    @test Zygote.pullback(eigvals, A)[1] ≈ eigvals(A)
  end

  @testset "eigvals(::Hermitian{<:Real})" begin
    rng, N = MersenneTwister(456), 7
    A = Hermitian(randn(rng, N, N))
    @test gradtest(x->eigvals(Hermitian(x)), collect(A))
    @test Zygote.pullback(eigvals, A)[1] ≈ eigvals(A)
  end

  @testset "eigvals(::Hermitian{<:Complex})" begin
    rng, N = MersenneTwister(789), 7
    A, B = randn(rng, N, N), randn(rng, N, N)
    @test gradtest(A, B) do a,b
      c = Hermitian(complex.(a, b))
      return eigvals(c)
    end
    @test Zygote.pullback(eigvals, Hermitian(A))[1] ≈ eigvals(Hermitian(A))
  end
end

using Distances

Zygote.refresh()

@testset "distances" begin
  rng, P, Q, D = MersenneTwister(123456), 10, 9, 8

   # Check sqeuclidean.
  let
    x, y = randn(rng, D), randn(rng, D)
    @test gradtest(x->sqeuclidean(x, y), x)
    @test gradtest(y->sqeuclidean(x, y), y)
  end

   # Check binary colwise.
  let
    X, Y = randn(rng, D, P), randn(rng, D, P)
    @test gradtest(X->colwise(SqEuclidean(), X, Y), X)
    @test gradtest(Y->colwise(SqEuclidean(), X, Y), Y)
  end

   # Check binary pairwise.
  let
    X, Y = randn(rng, D, P), randn(rng, D, Q)
    @test gradtest(X->pairwise(SqEuclidean(), X, Y; dims=2), X)
    @test gradtest(Y->pairwise(SqEuclidean(), X, Y; dims=2), Y)
  end
  let
    Xt, Yt = randn(rng, P, D), randn(rng, Q, D)
    @test gradtest(Xt->pairwise(SqEuclidean(), Xt, Yt; dims=1), Xt)
    @test gradtest(Yt->pairwise(SqEuclidean(), Xt, Yt; dims=1), Yt)
  end

   # Check unary pairwise.
  @test gradtest(X->pairwise(SqEuclidean(), X; dims=2), randn(rng, D, P))
  @test gradtest(Xt->pairwise(SqEuclidean(), Xt; dims=1), randn(rng, P, D))

  @testset "colwise(::Euclidean, X, Y; dims=2)" begin
    rng, D, P = MersenneTwister(123456), 2, 3
    X, Y, D̄ = randn(rng, D, P), randn(rng, D, P), randn(rng, P)
    gradtest((X, Y)->colwise(Euclidean(), X, Y), X, Y)
  end
  @testset "pairwise(::Euclidean, X, Y; dims=2)" begin
    rng, D, P, Q = MersenneTwister(123456), 2, 3, 5
    X, Y, D̄ = randn(rng, D, P), randn(rng, D, Q), randn(rng, P, Q)
    gradtest((X, Y)->pairwise(Euclidean(), X, Y; dims=2), X, Y)
  end
  @testset "pairwise(::Euclidean, X; dims=2)" begin
    rng, D, P = MersenneTwister(123456), 2, 3
    X, D̄ = randn(rng, D, P), randn(rng, P, P)
    gradtest(X->pairwise(Euclidean(), X; dims=2), X)
  end
end

function cat_test(f, A::Union{AbstractVector, AbstractMatrix}...)
  @test gradtest(f, A...)
  Z, back = Zygote.pullback(f, A...)
  Ā = back(randn(size(Z)))
  @test all(map((a, ā)->ā isa typeof(a), A, Ā))
end

@testset "vcat" begin
  # Scalar
  @test gradient((x,y) -> sum(vcat(x,y)), 1,2) == (1,1)
  @test gradient((x,y) -> sum([x;y]), 1,2) == (1,1)

  # Vector-only.
  cat_test(vcat, randn(1))
  cat_test(vcat, randn(10))

  cat_test(vcat, randn(1), randn(1))
  cat_test(vcat, randn(5), randn(1))
  cat_test(vcat, randn(1), randn(6))
  cat_test(vcat, randn(5), randn(6))

  # Matrix-only.
  for c in [1, 2], r1 in [1, 2], r2 in [1, 2]
    cat_test(vcat, randn(r1, c))
    cat_test(vcat, randn(r1, c), randn(r2, c))
  end

  # Matrix-Vector / Vector-Matrix / Vector-Matrix-Vector.
  for r in [1, 2]
    cat_test(vcat, randn(r, 1), randn(r))
    cat_test(vcat, randn(r), randn(r, 1))
    cat_test(vcat, randn(r), randn(r, 1), randn(r))
  end
end

@testset "hcat" begin
  # Scalar
  @test gradient((x,y) -> sum(hcat(x,y)), 1,2) == (1,1)
  @test gradient((x,y) -> sum([x y]), 1,2) == (1,1)
  @test gradient((a,b,c,d) -> sum(sqrt, [a b;c d]), 1,1,1,4) == (0.5, 0.5, 0.5, 0.25)

  # Vector-only.
  for r in [1, 2]
    cat_test(hcat, randn(r))
    cat_test(hcat, randn(r), randn(r))
  end

  # Matrix-only.
  for r in [1, 2], c1 in [1, 2], c2 in [1, 2]
    cat_test(hcat, randn(r, c1), randn(r, c2))
  end

  # Matrix-Vector / Vector-Matrix / Vector-Matrix-Vector.
  for r in [1, 2], c in [1, 2]
    cat_test(hcat, randn(r, c), randn(r))
    cat_test(hcat, randn(r), randn(r, c))
    cat_test(hcat, randn(r), randn(r, c), randn(r))
  end
end

@testset "hvcat" begin
  @test gradient(xs -> hvcat((2,2),xs...)[1,1], [1,2,3,4])[1] == (1,0,0,0)
  @test gradient(xs -> hvcat((2,2),xs...)[2,1], [1,2,3,4])[1] == (0,0,1,0)
  @test gradient(xs -> hvcat((2,2),xs...)[1,2], [1,2,3,4])[1] == (0,1,0,0)
  @test gradient(xs -> hvcat((2,2),xs...)[2,2], [1,2,3,4])[1] == (0,0,0,1)
end

@testset "cat(..., dims = $dim)" for dim in 1:5
  catdim = (x...) -> cat(x..., dims = dim)
  @test gradtest(catdim, rand(5), rand(5))
  @test gradtest(catdim, rand(2,5), rand(2,5), rand(2,5))
  @test gradtest(catdim, rand(2,5,3), rand(2,5,3), rand(2,5,3))
end

@testset "one(s) and zero(s)" begin
  @test Zygote.gradient(x->sum(ones(size(x))), randn(5))[1] isa Nothing
  @test Zygote.gradient(x->sum(one(x)), randn(3, 3))[1] isa Nothing
  @test Zygote.gradient(x->sum(zeros(size(x))), randn(7))[1] isa Nothing
  @test Zygote.gradient(x->sum(zero(x)), randn(3))[1] isa Nothing
end

@testset "fma and muladd" begin
    @test gradcheck(x -> fma(x[1], x[2], x[3]), [2.0, 3.0, 5.0])
    @test gradcheck(x -> muladd(x[1], x[2], x[3]), [2.0, 3.0, 5.0])
end

import StatsFuns

Zygote.refresh()

@testset "xlogx" begin
  @test gradcheck(x->2.5 * StatsFuns.xlogx(x[1]), [1.0])
  @test gradcheck(x->2.5 * StatsFuns.xlogx(x[1]), [2.45])
end

@testset "logistic" begin
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [-5.0])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [-1.0])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [-eps()])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [0.0])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [eps()])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [1.0])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [5.0])
end

@testset "logit" begin
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.1])
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.3])
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.5])
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.7])
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.9])
end

function test_log1pexp(T, xs)
  y = T(4.3)
  for x in xs
    @test gradcheck(x->y * StatsFuns.log1pexp(x[1]), [x])
  end
end

@testset "log1pexp" begin
  @testset "Float64" begin
    @testset "x ∈ (-∞, 18.0)" begin
      test_log1pexp(Float64, [-1000.0, -50.0, -25.0, -10.0, 0.0, 10.0, 18.0 - eps()])
    end
    @testset "x ∈ [18.0, 33.3)" begin
      test_log1pexp(Float64, [18.0, 18.0 + eps(), 33.3 - eps()])
    end
    @testset "x ∈ [33.3, ∞)" begin
      test_log1pexp(Float64, [33.3, 33.3 + eps(), 100.0])
    end
  end
end

@testset "log1psq" begin
  rng = MersenneTwister(123456)
  @testset "Float64" begin
    for x in [-10.0, -5.0, -1.0, -eps(), 0.0, eps(), 1.0, 5.0, 10.0]
      @test gradcheck(x->5.1 * StatsFuns.log1psq(x[1]), [x])
    end
  end
end

@testset "logsumexp" begin
  rng = MersenneTwister(123456)
  @testset "Float64" begin
    @test gradtest(StatsFuns.logsumexp, randn(rng, 1))
    @test gradtest(StatsFuns.logsumexp, randn(rng, 1, 1))
    @test gradtest(StatsFuns.logsumexp, randn(rng, 3))
    @test gradtest(StatsFuns.logsumexp, randn(rng, 3, 4, 5))
  end
end

@testset "* sizing" begin
  @test size(Zygote.gradient((x, y)->sum(x * y), randn(1, 1), randn(1, 10))[1]) == (1, 1)
  @test size(Zygote.gradient((x, y)->sum(x * y), randn(1, 1), randn(1, 10))[2]) == (1, 10)
end

@testset "broadcast" begin
  @test gradient(x -> sum(sin.(x)), Diagonal(randn(3)))[1][2] == 1
end

using Zygote: Buffer

@testset "Buffer" begin
  @test gradient([1, 2, 3]) do x
    b = Buffer(x)
    b[:] = x
    return sum(copy(b))
  end == ([1,1,1],)

  function vstack(xs)
    buf = Buffer(xs, length(xs), 5)
    for i = 1:5
      buf[:, i] = xs
    end
    return copy(buf)
  end

  @test gradient(x -> sum(vstack(x)), [1, 2, 3]) == ([5, 5, 5],)

  buf = Buffer([1, 2, 3])
  buf[1] = 1
  copy(buf)
  @test_throws ErrorException buf[1] = 1
  @test eltype(buf) === Int
  @test length(buf) === 3
  @test ndims(buf) === 1
  @test size(buf) === (3, )
  @test size(buf, 2) === 1
  @test axes(buf) == (1:3, )
  @test axes(buf, 2) == 1:1
  @test eachindex(buf) == 1:3
  @test stride(buf, 2) === 3
  @test strides(buf) === (1, )

  @test gradient([1, 2, 3]) do xs
    b = Zygote.Buffer(xs)
    b .= xs .* 2
    return sum(copy(b))
  end == ([2,2,2],)
end

@testset "FillArrays" begin
  @test gradcheck(x->sum(Fill(x[], (2, 2))), [0.1])
  @test first(Zygote.gradient(sz->sum(Ones(sz)), 6)) === nothing
  @test first(Zygote.gradient(sz->sum(Zeros(sz)), 6)) === nothing
end

@testset "AbstractArray Addition / Subtraction / Negation" begin
  rng, M, N, P = MersenneTwister(123567), 3, 7, 11
  A, B = randn(rng, M, N, P), randn(rng, M, N, P)
  @test gradtest(+, A, B)
  @test gradtest(-, A, B)
  @test gradtest(-, A)
end

@testset "FFTW" begin
  x=[-0.353213 -0.789656 -0.270151; -0.95719 -1.27933 0.223982]
  # gradient of ifft(rfft) must be 1
  @test gradient((x)->real(ifft(fft(x))[1]),x)[1][1] == 1.0+0.0im
  @test gradient((x)->real(fft(ifft(x))[1]),x)[1][1] == 1.0+0.0im

  # check ffts for individual dimensions
  @test gradient((x)->sum(abs.(FFTW.fft(x))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
  @test gradient((x)->abs(sum((FFTW.fft(x)))),x)[1] ≈ gradient((x)->abs(sum(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
  @test gradient((x, dims)->sum(abs.(FFTW.fft(x,dims))),x,(1,2))[1] ≈ gradient((x)->sum(abs.(FFTW.fft(x))),x)[1]
  @test gradient((x)->sum(abs.(FFTW.fft(x,(1,2)))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
  @test gradient((x, dims)->sum(abs.(FFTW.ifft(x,dims))),x,(1,2))[1] ≈ gradient((x)->sum(abs.(FFTW.ifft(x))),x)[1]
  @test gradient((x)->sum(abs.(FFTW.ifft(x,(1,2)))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.ifft(FFTW.ifft(x,1),2))),x)[1]

  @test gradcheck(x->sum(abs.(FFTW.fft(x))), x)
  @test gradcheck(x->sum(abs.(FFTW.ifft(x))), x)
  @test gradcheck(x->sum(abs.(FFTW.fft(x, 1))), x)
  @test gradcheck(x->sum(abs.(FFTW.ifft(x, 1))), x)

end

@testset "FillArrays" begin
  rng, M, N = MersenneTwister(123456), 7, 11
  x, y = randn(rng), randn(rng)
  @test Zygote.gradient(x->sum(Fill(x, N)), x)[1] == N
  @test Zygote.gradient(x->sum(Fill(x, N, 3, 4)), x)[1] == N * 3 * 4
  @test Zygote.gradient((x, y)->sum(Fill(x, N)), x, y) == (N, nothing)

  let
    out, back = Zygote.pullback(sum, Fill(x, N))
    @test back(nothing) isa Nothing
  end

  z = randn(rng, N)
  @test gradtest(x->Fill(first(x), N), [x])
  let
    out, back = Zygote.pullback(x->Fill(x, N), x)
    @test out == Fill(x, N)
    @test first(back(Fill(y, N))) ≈ y * N
  end

  # Test unary broadcasting gradients.
  out, back = Zygote.pullback(x->exp.(x), Fill(x, N))
  @test out isa Fill
  @test out == Fill(exp(x), N)
  @test back(Ones(N))[1] isa Fill
  @test back(Ones(N))[1] == Ones(N) .* exp(x)
  @test back(ones(N))[1] isa Vector
  @test back(ones(N))[1] == ones(N) .* exp(x)
  @test gradtest(x->exp.(Fill(3 * first(x), N)), [x])

  @testset "broadcast + and *" begin
    for sx in [(M, N), (M, 1), (1, N), (1, 1)]
      for sy in [(M, N), (M, 1), (1, N), (1, 1)]
        z = randn(rng, broadcast_shape(sx, sy))

        # Addition
        @test gradtest((x, y)->Fill(first(x), sx...) .+ Fill(first(y), sy...), [x], [y])
        @test gradtest(x->Fill(first(x), sx...) .+ Ones(sy...), [x])
        @test gradtest(x->Fill(first(x), sx...) .+ Zeros(sy...), [x])
        @test gradtest(y->Ones(sx...) .+ Fill(first(y), sy...), [y])
        @test gradtest(y->Zeros(sx...) .+ Fill(first(y), sy...), [y])

        # Multiplication
        @test gradtest((x, y)->Fill(first(x), sx...) .* Fill(first(y), sy...), [x], [y])
        @test gradtest(x->Fill(first(x), sx...) .* Ones(sy...), [x])
        @test gradtest(x->Fill(first(x), sx...) .* Zeros(sy...), [x])
        @test gradtest(y->Ones(sx...) .* Fill(first(y), sy...), [y])
        @test gradtest(y->Zeros(sx...) .* Fill(first(y), sy...), [y])
      end
    end
  end
end

@testset "@nograd" begin
  @test gradient(x -> findfirst(ismissing, x), [1, missing]) == (nothing,)
  @test gradient(x -> findlast(ismissing, x), [1, missing]) == (nothing,)
  @test gradient(x -> findall(ismissing, x)[1], [1, missing]) == (nothing,)
end


text:
using Bijectors, Random
using Test

Random.seed!(123456)

@testset "Interface" begin
    include("interface.jl")
end
include("transform.jl")


text:
function get_data(args)::TunnelData
    case_file = parsed_args["casefile"]
    wall_case_file = parsed_args["wallcasefile"]

    case_data = readdlm(case_file)
    wall_distance = case_data[6]
    wp = []
    for i in 7:2:length(case_data)
        push!(wp, [case_data[i], case_data[i+1]])
        (i+1 == length(case_data)) && (break)
    end
    distances = []
    c_distances = []
    for i in 1:(length(wp)-1)
        push!(distances, dist(wp[i], wp[i+1]))
    end
    c_distances = [distances[1]]
    for i in 2:length(distances)
        last = c_distances[end]
        push!(c_distances, last + distances[i])
    end
    num_turns = length(wp) - 2
    num_segments = length(wp) - 1

    wall_data = readdlm(wall_case_file, ',')
    lw_wp = []
    rw_wp = []
    for i in 1:2:size(wall_data)[1]
        push!(lw_wp, [wall_data[i, 1], wall_data[i+1, 1]])
        push!(rw_wp, [wall_data[i, 2], wall_data[i+1, 2]])
        (i+1 == size(wall_data)[1]) && (break)
    end

    return TunnelData(wp, lw_wp, rw_wp,
        distances, c_distances, num_turns, num_segments, wall_distance)

end


text:
using BitBasis, LegibleLambdas

export @λ, @lambda

addbits!(n::Int) = @λ(register -> addbits!(register, n))
insert_qudits!(loc::Int; nqudits::Int = 1) =
    @λ(register -> insert_qudits!(register, loc; nqudits = nqudits))
YaoAPI.insert_qubits!(loc::Int; nqubits::Int = 1) =
    @λ(register -> insert_qubits!(register, loc; nqubits = nqubits))
YaoAPI.insert_qubits!(reg::AbstractRegister{2}, loc::Int; nqubits::Int = 1) = insert_qudits!(reg, loc; nqudits=nqubits)

nremain(r::AbstractRegister) = nqudits(r) - nactive(r)
nlevel(r::AbstractRegister{D}) where {D} = D
nqubits(r::AbstractRegister{2}) = nqudits(r)

"""
    focus!(locs...) -> f(register) -> register

Lazy version of [`focus!`](@ref), this returns a lambda which requires a register.
"""
focus!(locs::Int...) = focus!(locs)
focus!(locs::NTuple{N,Int}) where {N} = @λ(register -> focus!(register, locs))
focus!(locs::UnitRange) = @λ(register -> focus!(register, locs))

relax!(r::AbstractRegister; to_nactive::Int = nqudits(r)) =
    relax!(r, (); to_nactive = to_nactive)

"""
    relax!(locs::Int...; to_nactive=nqudits(register)) -> f(register) -> register

Lazy version of [`relax!`](@ref), it will be evaluated once you feed a register
to its output lambda.
"""
relax!(locs::Int...; to_nactive::Union{Nothing,Int} = nothing) =
    relax!(locs; to_nactive = to_nactive)

function relax!(locs::NTuple{N,Int}; to_nactive::Union{Nothing,Int} = nothing) where {N}
    lambda = function (r::AbstractRegister)
        if to_nactive === nothing
            return relax!(r, locs; to_nactive = nqudits(r))
        else
            return relax!(r, locs; to_nactive = to_nactive)
        end
    end

    @static if VERSION < v"1.1.0"
        return LegibleLambda("(register->relax!(register, locs...; to_nactive))", lambda)
    else
        return LegibleLambda(:(register -> relax!(register, locs...; to_nactive)), lambda)
    end
end

## Measurement
measure!(postprocess::PostProcess, op, reg::AbstractRegister; kwargs...) =
    measure!(postprocess, op, reg, AllLocs(); kwargs...)
measure!(postprocess::PostProcess, reg::AbstractRegister, locs; kwargs...) =
    measure!(postprocess, ComputationalBasis(), reg, locs; kwargs...)
measure!(postprocess::PostProcess, reg::AbstractRegister; kwargs...) =
    measure!(postprocess, ComputationalBasis(), reg, AllLocs(); kwargs...)
measure!(op, reg::AbstractRegister, args...; kwargs...) =
    measure!(NoPostProcess(), op, reg, args...; kwargs...)
measure!(reg::AbstractRegister, args...; kwargs...) =
    measure!(NoPostProcess(), reg, args...; kwargs...)

measure(op, reg::AbstractRegister; kwargs...) = measure(op, reg, AllLocs(); kwargs...)
measure(reg::AbstractRegister, locs; kwargs...) =
    measure(ComputationalBasis(), reg, locs; kwargs...)
measure(reg::AbstractRegister; kwargs...) =
    measure(ComputationalBasis(), reg, AllLocs(); kwargs...)

# focus! to specify locations, we that we only need to consider full-space measure in the future.
function measure!(
    postprocess::PostProcess,
    op,
    reg::AbstractRegister,
    locs;
    kwargs...,
) where {MODE}
    nbit = nactive(reg)
    focus!(reg, locs)
    res = measure!(postprocess, op, reg, AllLocs(); kwargs...)
    if postprocess isa RemoveMeasured
        relax!(reg; to_nactive = nbit - length(locs))
    else
        relax!(reg, locs; to_nactive = nbit)
    end
    res
end

function measure(op, reg::AbstractRegister, locs; kwargs...) where {MODE}
    nbit = nactive(reg)
    focus!(reg, locs)
    res = measure(op, reg, AllLocs(); kwargs...)
    relax!(reg, locs; to_nactive = nbit)
    res
end

"""
    basis(register) -> UnitRange

Returns an `UnitRange` of the all the bits in the Hilbert space of given register.
"""
BitBasis.basis(r::AbstractRegister) = basis(nqudits(r))

invorder!(r::AbstractRegister) = reorder!(r, Tuple(nactive(r):-1:1))

"""
    partial_tr(locs) -> f(register)

Curried version of `partial_tr(register, locs)`.
"""
partial_tr(locs) = @λ(register -> partial_tr(register, locs))

"""
    select!(b::Integer) -> f(register)

Lazy version of [`select!`](@ref). See also [`select`](@ref).
"""
select!(bits...) = @λ(register -> select!(register, bits...))

ρ(x) = density_matrix(x)


text:
@testset "Test MultiFunction in multiple dimensions" begin
    Nxlist= collect(1:5)
    Blist = [ProHermiteBasis(10); PhyHermiteBasis(10); CstProHermiteBasis(10); CstPhyHermiteBasis(10); CstLinProHermiteBasis(10); CstLinPhyHermiteBasis(10)]
    Nψ = 10
    Ne = 50

    for Nx in Nxlist
        X = randn(Nx, Ne)
        idx = rand(0:10,Nψ, Nx)
        for b in Blist
            for i=1:Nψ
            B = MultiBasis(b, Nx)
            F = MultiFunction(B, idx[i,:])
                for j=1:Ne
                    @test norm(F(X[:,j]) - prod(k-> B.B[idx[i,k]+1](X[k,j]), 1:Nx))<1e-6
                end
            end
        end
    end
end


text:
using Oceananigans
using Oceananigans.Units
using ParameterEstimocean
using ParameterEstimocean: Transformation
using LinearAlgebra
using DataDeps
using Distributions
using CairoMakie
using ElectronDisplay

using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: CATKEVerticalDiffusivity

#####
##### Compile LESbrary
#####

times = [4hours, 24hours, 48hours]
field_names = (:b, :e, :u, :v)
Nensemble = 1000
Δt = 20minutes
closure = CATKEVerticalDiffusivity()

z = [-256.0,
     -238.3,
     -207.2,
     -182.3,
     -162.5,
     -146.5,
     -133.0,
     -121.3,
     -110.5,
     -100.2,
     - 90.0,
     - 80.0,
     - 70.0,
     - 60.0,
     - 50.0,
     - 40.0,
     - 30.0,
     - 20.0,
     - 10.0,
        0.0]

regrid = RectilinearGrid(size=length(z)-1; z, topology=(Flat, Flat, Bounded))

transformation = (b = ZScore(),
                  u = ZScore(),
                  v = ZScore(),
                  e = RescaledZScore(1.0))

cases = ["free_convection",
         "strong_wind_weak_cooling",
         "med_wind_med_cooling",
         "weak_wind_strong_cooling",
         "strong_wind",
         "strong_wind_no_rotation"]

suite = "two_day_suite"
case_path(case) = joinpath("data", suite, case * "_instantaneous_statistics.jld2")

observation_library = Dict()

# Don't optimize u, v for free_convection
observation_library["free_convection"] =
    SyntheticObservations(case_path("free_convection"); transformation, times, regrid,
                          field_names = (:b, :e))

# Don't optimize v for non-rotating cases
observation_library["strong_wind_no_rotation"] =
    SyntheticObservations(case_path("strong_wind_no_rotation"); transformation, times, regrid,
                          field_names = (:b, :e, :u))

# The rest are standard
for case in ["strong_wind", "med_wind_med_cooling", "strong_wind_weak_cooling", "weak_wind_strong_cooling"]
    observation_library[case] = SyntheticObservations(case_path(case); field_names,
                                                      transformation, times, regrid)
end

observations = [observation_library[case] for case in cases]

#####
##### Calibration
#####

bounds_library = Dict()

# Turbulent kinetic energy parameters
bounds_library[:CᵂwΔ]  = ( 2.0,  20.0)
bounds_library[:Cᵂu★]  = ( 2.0,  20.0)
bounds_library[:Cᴰ⁻]   = ( 0.0,   1.0)
bounds_library[:Cᴰʳ]   = (-1.0,  10.0)
bounds_library[:CᴰRiᶜ] = (-4.0,   4.0)
bounds_library[:CᴰRiʷ] = ( 0.0,   2.0)

# Mixing length parameters
#
#   Recall σ = σ⁻ (1 + σʳ * step(x, c, w))
#
bounds_library[:Cᴷu⁻]  = ( 0.0,  10.0)
bounds_library[:Cᴷc⁻]  = ( 0.0,  10.0)
bounds_library[:Cᴷe⁻]  = ( 0.0,  10.0)
bounds_library[:Cᴷuʳ]  = (-1.0,   0.1)
bounds_library[:Cᴷcʳ]  = (-1.0,   0.1)
bounds_library[:Cᴷeʳ]  = (-1.0,   0.1)
bounds_library[:CᴷRiᶜ] = (-4.0,   4.0)
bounds_library[:CᴷRiʷ] = ( 0.0,   2.0)
bounds_library[:Cᵇu]   = ( 0.0,   4.0)
bounds_library[:Cᵇc]   = ( 0.0,   4.0)
bounds_library[:Cᵇe]   = ( 0.0,   4.0)
bounds_library[:Cˢu]   = ( 0.0,   4.0)
bounds_library[:Cˢc]   = ( 0.0,   4.0)
bounds_library[:Cˢe]   = ( 0.0,   4.0)

bounds_library[:Cᴬˢu]  = ( 0.0,   2.0)
bounds_library[:Cᴬˢc]  = ( 0.0,   2.0)
bounds_library[:Cᴬˢe]  = ( 0.0,   2.0)
bounds_library[:Cᴬu]   = ( 0.0,  10.0)
bounds_library[:Cᴬc]   = ( 0.0,  10.0)
bounds_library[:Cᴬe]   = ( 0.0,  10.0)

# Extras
bounds_library[:Cᵇ]    = ( 0.0,   4.0)
bounds_library[:Cˢ]    = ( 0.0,   4.0)
bounds_library[:Cᴰ]    = ( 0.0,   2.0)
bounds_library[:Cᵟu]   = ( 0.0,  10.0)
bounds_library[:Cᵟc]   = ( 0.0,  10.0)
bounds_library[:Cᵟe]   = ( 0.0,  10.0)

prior_library = Dict()

for p in keys(bounds_library)
    prior_library[p] = ScaledLogitNormal(; bounds=bounds_library[p])
end

# No convective adjustment:
constant_Ri_parameters = (:Cᴰ, :CᵂwΔ, :Cᵂu★, :Cᴸᵇ, :Cᴷu⁻, :Cᴷc⁻, :Cᴷe⁻, :Cᵟu, :Cᵟc, :Cᵟe)
variable_Ri_parameters = (:Cᴷuʳ, :Cᴷcʳ, :Cᴷeʳ, :CᴷRiʷ, :CᴷRiᶜ, :Cᴰ, :Cᴸᵇ, :CᵂwΔ, :Cᵂu★)
convective_adjustment_parameters = (:Cᴬc, :Cᴬe) # Cᴬu

# :Cᴸˢ
parameter_names = (:CᵂwΔ,  :Cᵂu★, :Cᴰ⁻, :Cᴰʳ,
                   :Cˢc,   :Cˢu,  :Cˢe,
                   :Cᵇc,   :Cᵇu,  :Cᵇe,
                   :Cᴷc⁻,  :Cᴷu⁻, :Cᴷe⁻,
                   :Cᴷcʳ,  :Cᴷuʳ, :Cᴷeʳ,
                   #:Cᴬc,  :Cᴬu, :Cᴬe, :Cᴬˢc, :Cᴬˢu, :Cᴬˢe,
                   :CᴷRiᶜ, :CᴷRiʷ)

free_parameters = FreeParameters(prior_library, names=parameter_names)

function build_simulation()
    simulation = ensemble_column_model_simulation(observations; closure, Nensemble,
                                                  architecture = CPU(),
                                                  tracers = (:b, :e))

    simulation.Δt = Δt

    Qᵘ = simulation.model.velocities.u.boundary_conditions.top.condition
    Qᵇ = simulation.model.tracers.b.boundary_conditions.top.condition
    N² = simulation.model.tracers.b.boundary_conditions.bottom.condition

    for (case, obs) in enumerate(observations)
        f = obs.metadata.parameters.coriolis_parameter
        view(Qᵘ, :, case) .= obs.metadata.parameters.momentum_flux
        view(Qᵇ, :, case) .= obs.metadata.parameters.buoyancy_flux
        view(N², :, case) .= obs.metadata.parameters.N²_deep
        view(simulation.model.coriolis, :, case) .= Ref(FPlane(f=f))
    end

    return simulation
end

simulation = build_simulation()
calibration = InverseProblem(observations, simulation, free_parameters)
resampler = Resampler(resample_failure_fraction=0.5, acceptable_failure_fraction=1.0)
eki = EnsembleKalmanInversion(calibration; resampler, pseudo_stepping = ConstantConvergence(0.6))

#####
##### Plot utils
#####

Nt = length(times)
observed_data = []

for observation in observations
    time_serieses = observation.field_time_serieses
    names = keys(time_serieses)
    case_data = NamedTuple(n => interior(getproperty(time_serieses, n)[Nt])[1, 1, :] for n in names)
    push!(observed_data, case_data)
end

function get_modeled_case(icase, name, k=1)
    model_time_serieses = calibration.time_series_collector.field_time_serieses
    field = getproperty(model_time_serieses, name)[Nt]
    return view(interior(field), k, icase, :)
end

colorcycle =  [:black, :royalblue1, :darkgreen, :lightsalmon, :seagreen, :magenta2, :red4, :khaki1,   :darkgreen, :bisque4,
               :silver, :lightsalmon, :lightseagreen, :teal, :royalblue1, :darkorchid4]

markercycle = [:rect, :utriangle, :star5, :circle, :cross, :+, :pentagon, :ltriangle, :diamond, :star4]
markercycle = repeat(markercycle, inner=2)
colorcycle = repeat(colorcycle, inner=2)

function make_axes(fig, row=1, label=nothing)
    ax_b = Axis(fig[row, 1], xlabel = "Buoyancy \n[cm s⁻²]", ylabel = "z [m]")
    ax_u = Axis(fig[row, 2], xlabel = "x-velocity \n[cm s⁻¹]")
    ax_v = Axis(fig[row, 3], xlabel = "y-velocity \n[cm s⁻¹]")
    ax_e = Axis(fig[row, 4], xlabel = "Turbulent kinetic energy \n[cm² s⁻²]")
    if !isnothing(label)
        ax_t = Axis(fig[row, 5])
        xlims!(0, 1)
        ylims!(0, 1)
        hidespines!(ax_t)
        hidedecorations!(ax_t)
        text!(ax_t, label, justification=:left, align=(:left, :center), position=(0, 0.5))
    end
    return (ax_b, ax_u, ax_v, ax_e)
end

function plot_fields!(axs, label, color, b, e, u=zeros(size(b)), v=zeros(size(b)); linewidth=2, linestyle=:solid)
    grid = first(values(observation_library)).grid
    z = znodes(Center, grid)
    b, u, v, e = Tuple(Array(f) for f in (b, u, v, e))

    for (q, name) in zip((b, u, v, e), ("b", "u", "v", "e"))
        any(isnan.(q)) && @warn("NaNs found in $label $(name)!")
    end

    ## Note unit conversions below, eg m s⁻¹ -> cm s⁻¹:cyan
    lines!(axs[1], 1e2 * b, z; color, linestyle, label, linewidth)
    lines!(axs[2], 1e2 * u, z; color, linestyle, label, linewidth)
    lines!(axs[3], 1e2 * v, z; color, linestyle, label, linewidth)
    lines!(axs[4], 1e4 * e, z; color, linestyle, label, linewidth)

    return nothing
end

function min_max_parameters(summary)
    names = keys(summary.ensemble_mean)
    Nens = length(summary.parameters)
    parameter_matrix = [summary.parameters[k][name] for name in names, k = 1:Nens]
    θ_min = minimum(parameter_matrix, dims=2)
    θ_max = maximum(parameter_matrix, dims=2)
    return θ_min, θ_max
end

function finitefind(a, val, find)
    b = deepcopy(a)
    b[.!isfinite.(a)] .= val
    return find(b)
end

finitefindmin(a) = finitefind(a, Inf, findmin)
finitefindmax(a) = finitefind(a, -Inf, findmax)

function visualize_parameter_evolution(eki)
    summaries = eki.iteration_summaries
    Niters = length(summaries)
    names = eki.inverse_problem.free_parameters.names
    θ_mean = NamedTuple(name => map(s -> s.ensemble_mean[name], summaries) for name in names)

    k_best(s) = finitefindmin(s.mean_square_errors)[2]
    θ_best = NamedTuple(name => map(s -> s.parameters[k_best(s)][name], summaries) for name in names)

    θ_min_max = [min_max_parameters(s) for s in summaries]
    θ_min = [[θn[1][i] for θn in θ_min_max] for i in 1:length(names)]
    θ_max = [[θn[2][i] for θn in θ_min_max] for i in 1:length(names)]

    θᵢ = NamedTuple(name => first(θ_mean[name]) for name in names)
    Δθ = NamedTuple(name => (θ_mean[name] .- θᵢ[name]) ./ θᵢ[name] for name in names)
    iterations = 0:length(summaries)-1

    fig = Figure(resolution=(1200, 1200))
    ax1 = Axis(fig[1:3, 1], xlabel = "Iteration", ylabel = "Δθ")
    for (i, name) in enumerate(names)
        label = string(name)
        marker = markercycle[i]
        color = colorcycle[i]
        scatterlines!(ax1, iterations, parent(Δθ[name]); marker, color=(color, 0.8), label, linewidth=4)
    end

    fig[1:3, 2] = Legend(fig, ax1)

    Nparts = 3
    Nθpart = floor(Int, length(names) / Nparts)

    for p in 1:Nparts
        axp = Axis(fig[p+3, 1], xlabel = "Iteration", ylabel = "θ")

        if p == Nparts
            np = UnitRange((p-1) * Nθpart + 1, length(names))
        else
            np = UnitRange((p-1) * Nθpart + 1, p * Nθpart)
        end

        partnames = names[np]

        for (n, name) in enumerate(partnames)
            i = (p - 1) * Nθpart + n
            label = string(name)
            marker = markercycle[i]
            color = colorcycle[i]
            scatterlines!(axp, iterations, parent(θ_mean[name]); marker, color=(color, 0.6), label, linewidth=4)
            lines!(axp, iterations, parent(θ_best[name]); color, linewidth=2)
            band!(axp, iterations, parent(θ_min[i]), parent(θ_max[i]), color=(color, 0.3))
        end

        fig[p+3, 2] = Legend(fig, axp)
    end

    display(fig)

    return nothing
end

function plot_latest(eki)
    latest_summary = eki.iteration_summaries[end]
    min_error, k_min = finitefindmin(latest_summary.mean_square_errors)
    max_error, k_max = finitefindmax(latest_summary.mean_square_errors)

    fig = Figure(resolution=(1200, 1200))

    for (c, case) in enumerate(cases)
        label = replace(case, "_" => "\n")
        axs = make_axes(fig, c, label)
        observed = observed_data[c]
        obs = observations[c]

        min_error_data = NamedTuple(n => get_modeled_case(c, n, k_min) for n in keys(obs.field_time_serieses))
        max_error_data = NamedTuple(n => get_modeled_case(c, n, k_max) for n in keys(obs.field_time_serieses))

        plot_fields!(axs, "observed at t = " * prettytime(times[end]), (:gray23, 0.6), observed...; linewidth=4)
        plot_fields!(axs, "min", :navy, min_error_data...)
        plot_fields!(axs, "max", :orangered3, max_error_data...)

        fig[1, 6] = Legend(fig, axs[1])
    end

    display(fig)

    return nothing
end

#####
##### Visualize observations
#####

fig = Figure(resolution=(1200, 1200))
linestyles = [:solid, :dash, :dot, :dashdot, :dashdotdot, :solid]
all_axs = []

for (o, observation) in enumerate(observations)
    axs_label = replace(cases[o], "_" => "\n")
    axs = make_axes(fig, o, axs_label)
    append!(all_axs, axs)
    for (n, t) in enumerate(times)
        linestyle = linestyles[o]
        label = "t = " * prettytime(t)
        names = keys(observation.field_time_serieses)
        data = map(name -> interior(observation.field_time_serieses[name][n])[1, 1, :], names)
        plot_fields!(axs, label, colorcycle[n], data...)
    end
end

display(fig)

#####
##### Calibrate
#####

# Initial state after 0 iterations
plot_latest(eki)
@show eki.iteration_summaries[end]

# Continuously update
for i = 1:20
    @info "Iterating..."
    start_time = time_ns()
    iterate!(eki)
    elapsed = 1e-9 * (time_ns() - start_time)
    @info string("   done. (", prettytime(elapsed), ")")
    @show eki.iteration_summaries[end]
    visualize_parameter_evolution(eki)
    plot_latest(eki)
end

max_stars_repo_name: glwagner/LocalOceanClosureCalibration

max_issues_repo_name: glwagner/LocalOceanClosureCalibration

max_forks_repo_name: glwagner/LocalOceanClosureCalibration


text:
export naiveRounding, objGapRedRounding, massPreservingRounding

"""
naiveRounding(s,thresh=0.5)

simple rounding scheme based on threshold.

"""
naiveRounding(s,thresh=0.5) = (s.>=thresh)

"""
function massPreservingRounding(s)

mass-preserving rounding

"""
function massPreservingRounding(s)
    r   = falses(length(s))
    pos = reverse(sortperm(s))
    for k=1:round(Int64,sum(s))
        r[pos[k]] = 1
    end
    return r
end

"""
function objGapRedRounding(s,pMis,pInv;delta=0.1)

objective gap reduction rounding
"""
function objGapRedRounding(s,pMis,pInv;delta=0.1,out=0)
    r     = zeros(Int64,length(s))
    Jbest = 10000
    sbest = 0
    smin = floor(minimum(s)*10)*.1 + delta
    smax = floor(maximum(s)*10)*.1
    for k=smin:delta:smax
        r = 1.0*naiveRounding(s,k)
        sig,dsig = pInv.modelfun(r)
         Dc,F,dF,d2F,pMis,tMis = computeMisfit(sig,pMis,true)
         R,dR,d2R = computeRegularizer(pInv.regularizer,1.0*r,pInv.mref,pInv.MInv,pInv.alpha)
		 if out>=2
			 println("Obj-gap reduction rounding: cut-off, MisFit, Regularizer, Total $(k), $(F) , $(R) , $(F+R)")
		 end

         if Jbest > F+R
             sbest = k
             Jbest = F+R
         end
     end
     if sbest > 0
         r = naiveRounding(s,sbest)
     end
    return r
end

max_stars_repo_name: JuliaInv/ConvDiffMIPDECO

max_issues_repo_name: JuliaInv/ConvDiffMIPDECO

max_forks_repo_name: JuliaInv/ConvDiffMIPDECO


text:
using BenchmarkTools, MkCell

cellfcc = [-0.5 0 0.5; 0 0.5 0.5; -0.5 0.5 0];
cellcubic = [1.0 0 0; 0 1.0 0; 0 0 1.0];

const suite = BenchmarkGroup()

suite["integration"] = BenchmarkGroup()
suite["integration"]["fcc1"] = @benchmarkable MkCell.cellopt($cellfcc, 1, 2)
suite["integration"]["cubic3"] = @benchmarkable MkCell.cellopt($cellcubic, 3, 2)

# If a cache of tuned parameters already exists, use it, otherwise, tune and cache
# the benchmark parameters. Reusing cached parameters is faster and more reliable
# than re-tuning `suite` every time the file is included.
paramspath = joinpath(dirname(@__FILE__), "params.json")
if isfile(paramspath)
loadparams!(suite, BenchmarkTools.load(paramspath)[1], :evals);
else
tune!(suite)
    BenchmarkTools.save(paramspath, params(suite));
end

# run the suite
println("Warming up")
warmup(suite, verbose=false)
println("Running")
bmrun = run(suite)
println("fcc1 integration:")
display(bmrun["integration"]["fcc1"])
println()
println("cubic3 integration: ")
display(bmrun["integration"]["cubic3"])
println()


text:
using Random
Random.seed!(1)

function montyHall(switchPolicy)
    prize, choice = rand(1:3), rand(1:3)
    if prize == choice
        revealed = rand(setdiff(1:3,choice))
    else
        revealed = rand(setdiff(1:3,[prize,choice]))
    end

    if switchPolicy
        choice = setdiff(1:3,[revealed,choice])[1]
    end
    return choice == prize
end

N = 10^6
println("Success probability with policy I (stay): ",
	sum([montyHall(false) for _ in 1:N])/N)
println("Success probability with policy II (switch): ",
	sum([montyHall(true) for _ in 1:N])/N)

max_stars_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_issues_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_forks_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook


text:
@testset "running_mean" begin
    ndim = 20
    npts = 10000
    x = randn(Float,ndim,npts)
    m = Vector{Float}(undef,ndim)
    VectorDataUtils.running_mean!(m,x)
    @test mean(x,dims=2) ≈ m
    @test mean(x,dims=2) ≈ VectorDataUtils.running_mean(x)
end

@testset "update_running_mean!" begin
    ndim = 20
    npts = 10000
    x = randn(Float,ndim,npts)
    m = zeros(Float,ndim)
    for i ∈ axes(x,2)
        VectorDataUtils.update_running_mean!(m,x[:,i],i)
        # @test mean(x[:,1:i],dims=2) ≈ m
    end
    @test mean(x,dims=2) ≈ m
end

@testset "running_meanvar" begin
    ndim = 20
    npts = 10000
    x = randn(Float,ndim,npts)
    m̂ = Vector{Float}(undef,ndim)
    v̂ = Vector{Float}(undef,ndim)
    VectorDataUtils.running_meanvar!(m̂,v̂,x)
    m,v = VectorDataUtils.running_meanvar(x)
    @test mean(x,dims=2) ≈ m̂
    @test var(x,dims=2) ≈ v̂
    @test mean(x,dims=2) ≈ m
    @test var(x,dims=2) ≈ v
end

@testset "update_running_meanvar!" begin
    ndim = 20
    npts = 10000
    x = randn(Float,ndim,npts)
    m,v = VectorDataUtils.running_meanvar(x[:,1:2])
    # m = mean(x[:,1:2],dims=2)
    # v = var(x[:,1:2],dims=2)
    for i ∈ 3:size(x,2)
        VectorDataUtils.update_running_meanvar!(m,v,x[:,i],i)
        # @test mean(x[:,1:i],dims=2) ≈ m
        # @test var(x[:,1:i],dims=2) ≈ v
    end
    @test mean(x,dims=2) ≈ m
    @test var(x,dims=2) ≈ v
end

@testset "running_meancov" begin
    ndim = 20
    npts = 10000
    x = randn(Float,ndim,npts)
    m̃ = Vector{Float}(undef,ndim)
    C̃ = Matrix{Float}(undef,ndim,ndim)
    VectorDataUtils.running_meancov!(m̃,C̃,x)
    m,C = VectorDataUtils.running_meancov(x)
    m̂ = mean(x,dims=2)
    Ĉ = ((x.-m̂)*(x.-m̂)')/(size(x,2)-1)
    @test m̂ ≈ m̃
    @test Ĉ ≈ C̃
    @test m̂ ≈ m
    @test Ĉ ≈ C
end

@testset "update_running_meancov!" begin
    ndim = 20
    npts = 10000
    x = randn(Float,ndim,npts)
    m,C = VectorDataUtils.running_meancov(x[:,1:2])
    for i ∈ 3:size(x,2)
        VectorDataUtils.update_running_meancov!(m,C,x[:,i],i)
    end
    m̂ = mean(x,dims=2)
    Ĉ = ((x.-m̂)*(x.-m̂)')/(size(x,2)-1)
    @test m̂ ≈ m
    @test Ĉ ≈ C
end


text:
using Reexport

@reexport module SystemModels

    import DynamicalSystems:
        ContinuousDynamicalSystem,
        DiscreteDynamicalSystem,
        AbstractDataset,
        Dataset,
        trajectory,
        columns
    using UncertainData
    import StaticArrays: SVector
    import StatsBase: std
    import SimpleDiffEq
    import SimpleDiffEq: get_dt
    using Distributions
    import Base.rand

    # Abstract model definitions and methods
    include("AbstractSystemModel.jl")
    include("ContinuousSystemModel.jl")
    include("DiscreteSystemModel.jl")

    # Methods to add observational noise after orbits have been obtained
    include("add_noise.jl")

    # Continuous system definitions
    include("continuous/RosslerLorenzUnidir.jl")

    export
    AbstractSystemModel,
    ContinuousSystemModel,
    DiscreteSystemModel,
    rand,
    get_dt,
    get_ui,
    get_nvars,
    add_observational_noise!,
    trajectory,
    RosslerLorenzUnidir

end


text:
####################
#
# Utils
#
####################

# Open a shapefile
function open_shapefile(location::AbstractString)
    shp = open(location) do fd
        read(fd, Shapefile.Handle)
    end
    shp
end

#=
Create a canvas to draw on
The correspondence of Compose.UnitBox fields to the Shapefile.Rect fields are as follows:
x0 := left, y0 := top, width := right-left, height := -(top-bottom)
The UnitBox typically has origin at top left and moves down and right (an image convention).
To change to the convention of having the origin at the bottom left we use a negative height.
=#
function create_canvas(MBR::Shapefile.Rect, convertcoords, img_width)

    # Coordinate conversion
    left, top = convertcoords(MBR.left, MBR.top)
    right, bottom = convertcoords(MBR.right, MBR.bottom)

    # Canvas size
    width, height = right - left, top - bottom
    ratio = width / height
    Compose.set_default_graphic_size(img_width, img_width/ratio)

    # Dimensions of canvas
    # Add 5% padding on border (TODO - make padding a parameter)
    left -= width * 0.05
    width *= 1.1
    top += height * 0.05
    height *= 1.1
    dims = Compose.UnitBox(left, top, width, -height)

    # Return blank canvas
    return context(units = dims)
end

# Calculate a minimum bounding rectangle from a selection of shapes
# Useful for drawing a subset of shapes and not using the "global" MBR
function minBR(shapes::AbstractArray{T, 1}) where {T<:Shapefile.GeoInterface.AbstractGeometry}
    box = shapes[1].MBR
    left, bottom, right, top = box.left, box.bottom, box.right, box.top
    for i = 2:length(shapes)
        (shapes[i].MBR.left < left) && (left = shapes[i].MBR.left)
        (shapes[i].MBR.bottom < bottom) && (bottom = shapes[i].MBR.bottom)
        (shapes[i].MBR.right > right) && (right = shapes[i].MBR.right)
        (shapes[i].MBR.top > top) && (top = shapes[i].MBR.top)
    end
    return Shapefile.Rect(left, bottom, right, top)
end

# Calculate a minimum bounding rectangle from a selection of points
# Point types don't have field MBR so need it's own method
function minBR(shapes::AbstractArray{T, 1}) where {T<:ShpPoint}
    #box = Shapefile.Rect(shapes[1].y, shapes[1].x, shapes[1].y, shapes[1].x)
    left, bottom, right, top = shapes[1].x, shapes[1].y, shapes[1].x, shapes[1].y
    for i = 2:length(shapes)
        (shapes[i].x < left) && (left = shapes[i].x)
        (shapes[i].y < bottom) && (bottom = shapes[i].y)
        (shapes[i].x > right) && (right = shapes[i].x)
        (shapes[i].y > top) && (top = shapes[i].y)
    end
    return Shapefile.Rect(left, bottom, right, top)
end

# Convert shapes from Abstract Type Array{ESRIShape, 1} to concrete type
# No longer needed after changes to Shapefile.jl package
#=
function concrete{T<:Shapefile.GeoInterface.AbstractGeometry}(shapes::AbstractArray{T, 1})
    convert(Array{typeof(shapes[1]), 1}, shapes)
end
=#

# Root and dataset paths
rootpath = dirname(@__FILE__)[1:end-4]
testdatapath = joinpath(rootpath, "test", "testdata")


text:
using Pkg: TOML.parsefile, Types.VersionRange
using DataFrames, CSV

function parse_package(name)
    dir = "$(homedir())/.julia/registries/General/$(name[1:1])/$name"
    version = parsefile("$dir/Versions.toml") |>
        keys |>
        (k -> maximum(VersionNumber, k))
    compat = parsefile("$dir/Compat.toml") |>
        (x -> [ ((name = n,
                  version = VersionRange(v))
                  for (n, v) in v) for (k, v) ∈ x if version ∈ VersionRange(k) ] ) |>
        Iterators.flatten |>
        collect |>
        sort!
    if v"1" ∈ compat[findfirst(x-> isequal("julia", x.name), compat)].version
        package = parsefile("$dir/Package.toml") |>
            (x -> (name = x["name"], uuid = x["uuid"], repo = x["repo"][1:end - 4]))
        output = DataFrame(name = package.name,
                           repository = package.repo,
                           version = version,
                           dependency = getproperty.(compat, :name))
    else
        output = DataFrame([String, String, VersionNumber, String],
                           [:name, :repository, :version, :dependency],
                           0)
    end
    output
end

data = readdir.(string.("$(homedir())/.julia/registries/General/", 'A':'Z')) |>
    Iterators.flatten |>
    collect |>
    (x -> filter!(x -> !any(x ∈ ["julia", ".DS_Store"]) , x)) |>
    (x -> vcat(parse_package.(x)...))
CSV.write("data/data.tsv", data, delim = '\t')

max_stars_repo_name: team-oss/scrape-cran

max_issues_repo_name: team-oss/scrape-cran

max_forks_repo_name: team-oss/scrape-cran


text:
#!/bin/sh
#= La ligne shell suivante est en commentaire multiligne pour julia
exec julia --color=yes --startup-file=no --depwarn=no -- "$0" "$@"
=#

# AUTRE EXEMPLE DE RÉPERTOIRE DE TEST (cf QML.jl)
#   https://github.com/barche/QML.jl/blob/master/test/runtests.jl
#

global const APPDIR = dirname(dirname(realpath(@__FILE__())))
@show APPDIR

# Démarrage des fonctions de chronométrage
include("../src/utils/time.jl")
println("Début des tests en ", ms(), "s")

using Test
using Random
using Printf
using Dates

if !@isdefined(Args)
    include("../src/args.jl")
end

if !@isdefined(APPDIR)
    include("../src/args.jl")
end

include("../src/utils/log.jl")
include("../src/plane.jl")
include("../src/instance.jl")
include("../src/utils/array.jl")
include("../src/solution.jl")

include("../src/solvers/earliest_timing.jl")

include("../src/processing/instance_read_alp.jl")
include("../src/processing/instance_generators.jl")
# include("../src/solution.jl")

@ms using Statistics

@ms using Distributions: NegativeBinomial
@ms using StatsBase: sample, ProbabilityWeights
@ms include("../src/solvers/neighbour_operators.jl")

include("seqata_test_util.jl")

Args.parse_commandline(["test"])

include("../src/usings.jl") # pour les solveurs PL

println("## TEST $(basename(@__FILE__)) ")

@show(ARGS)

#
# Phase 1 : construction de la liste explicite du/des fichiers à tester
# - soit le/les fichier(s) passés en paramètre et aucun autre
# - soit tous les fichiers :
#   - commençant par test-*
#   - sans les fichiers de test de performance *-perf.jl (car lents !)
#   - et sans un liste de fichier explicitement à ignorer
#
global files = Vector{String}()
if !isempty(ARGS)
    # files = copy(ARGS)
    for file in ARGS
        push!(files, basename(file))
    end
else
    # Calcul de la liste des tests par défaut
    # EXEMPLES DE FICHIERS À NE PAS INCLURE E.G. CAR LONGGG.
    # CES FICHIERS DOIVENT ÊTRE TESTÉS EXPLICITEMENT EN LES PASSANT EN PARAMETRE.
    excluded = [
        "test-14b-dmip-cbc.jl", # long (1.5mn)
        "test-14c-dmip-glpk.jl", # trés trés long !!
    ]
    # global old_pwd = pwd() # global sinon warning julia07
    old_pwd = pwd()
    cd(dirname(@__FILE__))
    allfiles = readdir()
    cd(old_pwd)

    for file in allfiles
        print("\nfile=$file ?...")

        # On exclue les fichiers ne commençant pas par test-
        !startswith(file, "test-") && continue

        # On exclue les fichiers explicitement indiqués
        # in(excluded, file) && continue # NE MARCHE PAS !!!!
        (file in excluded) && continue # ok

        # On exclue les fichiers dédié aux performancex (car très lentttts !)
        endswith(file, "-perf.jl") && continue

        print("  => OK")
        # @show file
        push!(files, file)
    end
end

println("Liste des $(length(files)) tests à effectuer :")
for file in files
    println("  $file")
end

#
# Phase 1 : exécution de chaque test dans un contexte indépendant
# - on se déplace éventuellement dans sont sous-répertoire
# - on capture une erreur éventuelle pour ne pas arreter les autres tests
# - on chronomètre chaque test
#
@testset "Tests pour projet SEQATA" begin
    for file in files
        t0 = ms()
        println("\n"*"="^80)
        println("====== Test du fichier $(file)...")
        # global old_pwd = pwd() # global sinon warning julia07
        old_pwd = pwd()
        cd(dirname(@__FILE__))
        try
            @testset "Test $(file)" begin
                # println("include(", abspath(file), ")")
                include(abspath(file))
            end
        catch err
            println()
            # println(err)
            @warn err
            # rethrow(err)
        end
        cd(dirname(old_pwd))
        dt = round(ms()-t0, digits=3)
        println("====== Test du fichier $(file) fait en $(dt)s")
    end
    println("Fin des tests à ", ms(), "s")
end
# La suite ne serait pas exécutée en cas d'erreur
# println("Fin des tests en ", ms(), "s")


text:
using JuMP, Cbc

checkModel = Model(Cbc.Optimizer)

@variable(checkModel, box[1:4], binary = true)
@variable(checkModel, routes[1:3,1:4], binary = true)

cost = [75 150 225 75;
        357 255 153 306;
        84 105 126 105]

@objective(checkModel, Min, 150*sum(box[1:3])+sum(cost.*routes))

@constraint(checkModel, box[4]==1)

for city in 1:4
    @constraint(checkModel, routes[:,city].<=box[city])
end
for region in 1:3
    @constraint(checkModel, sum(routes[region,:])==1)
end

res = JuMP.optimize!(checkModel)

boxes = JuMP.value.(box)
routes = JuMP.value.(routes)

max_stars_repo_name: valenmgama/optimization-tec

max_issues_repo_name: valenmgama/optimization-tec

max_forks_repo_name: valenmgama/optimization-tec


text:
function DSSTAtmosphericDrag(arg0::Atmosphere, arg1::DragSensitive, arg2::jdouble)
    return DSSTAtmosphericDrag((Atmosphere, DragSensitive, jdouble), arg0, arg1, arg2)
end

function DSSTAtmosphericDrag(arg0::Atmosphere, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return DSSTAtmosphericDrag((Atmosphere, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end

function DSSTAtmosphericDrag(arg0::DragForce, arg1::jdouble)
    return DSSTAtmosphericDrag((DragForce, jdouble), arg0, arg1)
end

function get_atmosphere(obj::DSSTAtmosphericDrag)
    return jcall(obj, "getAtmosphere", Atmosphere, ())
end

function get_drag(obj::DSSTAtmosphericDrag)
    return jcall(obj, "getDrag", DragForce, ())
end

function get_events_detectors(obj::DSSTAtmosphericDrag)
    return jcall(obj, "getEventsDetectors", Vector{EventDetector}, ())
end

function get_field_events_detectors(obj::DSSTAtmosphericDrag, arg0::Field)
    return jcall(obj, "getFieldEventsDetectors", Vector{FieldEventDetector}, (Field,), arg0)
end

function get_rbar(obj::DSSTAtmosphericDrag)
    return jcall(obj, "getRbar", jdouble, ())
end

function get_spacecraft(obj::DSSTAtmosphericDrag)
    return jcall(obj, "getSpacecraft", DragSensitive, ())
end


text:
include("Opcodes.jl")
using .OpCodes
using StaticArrays

function run1(data, startwhite = false)
    right =  @SMatrix [0  1;
                     -1  0]
    left = @SMatrix [0 -1;
                      1  0]
    dir =   @SVector [-1, 0]
    pos =   @SVector [0, 0]
    field = Dict{typeof(pos), Int}()
    inp = Channel{Int}(1)
    outp = Channel{Int}(1)
    ic = IntComputer(data, inp, outp)
    @async compute!(ic)
    startwhite && (field[pos] = 1)
    while !ic.finished[]
        put!(inp, get(field, pos, 0))
        field[pos] = take!(outp)
        dir = take!(outp) == 1 ? right * dir : left * dir
        pos = pos + dir
    end
    field
end

# question 1
f = run1(readline("data/Dec11_data.txt"))
length(keys(f))

# question 2 (plus an added arg to the function)
f = run1(readline("data/Dec11_data.txt"), true)

b = fill(NaN, 100, 100)
for (k,v) in f
    b[(k .+ 50)...] = v
end
heatmap(b, ylim = (40,60), xlim = (40, 100), yflip = true)


text:
@testset "SDI Interface" begin
    @test DoubleSDI(PCA(15)) == DoubleSDI(PCA(15), PCA(15))
    @test SliceSDI(PCA(15)) == SliceSDI(PCA(15), PCA(15))
    cube = ones(5, 10, 100, 100)
    angles = zeros(10)
    scales = collect(range(1.2, 1, length=5))
    @test PCA(15)(cube, angles, scales) == SingleSDI(PCA(15))(cube, angles, scales)
end

@testset "Trivial SDI - $(typeof(ALG))" for ALG in (PCA(5), GreeDS(5), Classic())
    cube = ones(5, 10, 100, 100)
    angles = zeros(10)
    scales = collect(range(1.2, 1, length=5))

    reduced= ALG(cube, angles, scales)

    @test size(reduced) == (100, 100)
    @test all(x -> isapprox(x, 0, atol = 1e-9), reduced)
end

@testset "Trivial Double SDI - $(typeof(ALG1)), $(typeof(ALG2))" for ALG1 in (PCA(2), Classic(), GreeDS(2)),
                                                 ALG2 in (PCA(5), GreeDS(5), Classic())
    cube = ones(5, 10, 100, 100)
    angles = zeros(10)
    scales = collect(range(1.2, 1, length=5))

    reduced = DoubleSDI(ALG1, ALG2)(cube, angles, scales)

    @test size(reduced) == (100, 100)
    @test all(x -> isapprox(x, 0, atol = 1e-9), reduced)
end

@testset "Trivial Slice SDI - $(typeof(ALG1)), $(typeof(ALG2))" for ALG1 in (PCA(2), Classic(), GreeDS(2)),
                                                 ALG2 in (PCA(5), GreeDS(5), Classic())
    cube = ones(5, 10, 100, 100)
    angles = zeros(10)
    scales = collect(range(1.2, 1, length=5))

    reduced = SliceSDI(ALG1, ALG2)(cube, angles, scales)

    @test size(reduced) == (100, 100)
    @test all(x -> isapprox(x, 0, atol = 1e-9), reduced)
end


text:
export BlendMode
export MouseWheelScrollEvent
export JoystickMoveEvent
export JoystickConnectEvent
export JoystickButtonEvent
export JoystickAxis
export joystickX
export joystickY
export joystickZ
export joystickR
export joystickU
export joystickV
export joystickPovX
export joystickPovY
export is_connected
export get_button_count
export has_axis
export is_button_pressed
export get_axis_position
export joystick_update
export get_samples
export get_samplecount
export get_points
export get_thickness
export Sound
export Font
export CircleShape
export set_position
export set_radius
export set_fillcolor
export set_outlinecolor
export move
export get_position
export get_radius
export set_origin
export get_origin
export get_fillcolor
export get_outlinecolor
export rotate
export scale
export copy
export set_scale
export get_scale
export set_rotation
export get_rotation
export get_localbounds
export get_globalbounds
export set_outline_thickness
export get_outline_thickness
export get_pointcount
export set_pointcount
export get_texture_rect
export get_texture
export set_texture_rect
export set_texture
export get_point
export Color
export set_point
export ConvexShape
export Drawable
export Image
export save_to_file
export get_size
export set_pixel
export get_pixel
export get_pixels
export flip_vertically
export flip_horizontally
export Line
export draw
export set_points
export set_thickness
export FloatRect
export IntRect
export contains
export intersects
export RectangleShape
export set_size
export RenderStates
export RenderTexture
export set_active
export display
export clear
export set_view
export get_view
export get_default_view
export get_viewport
export set_smooth
export is_smooth
export set_repeated
export is_repeated
export pixel2coords
export coords2pixel
export RenderWindow
export set_framerate_limit
export isopen
export pollevent
export close
export set_vsync_enabled
export capture
export ContextSettings
export WindowStyle
export window_none
export window_resize
export window_defaultstyle
export window_close
export window_fullscreen
export set_title
export waitevent
export set_visible
export set_mousecursor_visible
export set_keyrepeat_enabled
export requestfocus
export hasfocus
export set_icon
export Shader
export destroy
export set_parameter
export shader_is_available
export ShaderFromMemory
export get_color
export Sprite
export set_color
export set_style
export set_charactersize
export set_font
export set_string
export RenderText
export get_style
export get_string
export get_font
export get_charactersize
export TextRegular
export TextBold
export TextUnderlined
export TextItalic
export TextStrikeThrough
export find_character_pos
export Texture
export copy_to_image
export update
export Transform
export PrimitiveType
export VertexArray
export Vertex
export get_vertexcount
export get_vertex
export resize
export append
export set_primitive_type
export get_primitive_type
export getbounds
export VideoMode
export get_desktop_mode
export is_valid
export zoom
export get_center
export reset
export set_viewport
export set_center
export View
export set_listener_upvector
export set_listener_direction
export set_listener_position
export set_global_volume
export get_listener_upvector
export get_listener_direction
export get_listener_position
export get_global_volume
export Music
export set_loop
export play
export pause
export stop
export get_channelcount
export get_samplerate
export set_pitch
export get_pitch
export get_volume
export get_status
export Sound
export set_buffer
export get_loop
export set_volume
export SoundBuffer
export get_duration
export SoundBufferRecorder
export start
export get_sample_rate
export get_buffer
export SoundStatus
export upload
export download
export delete_file
export rename_file
export delete_directory
export create_directory
export parent_directory
export change_directory
export get_directory_listing
export get_working_directory
export keep_alive
export disconnect
export login
export login_anonymous
export connect
export Ftp
export get_message
export isok
export FtpResponse
export get_directory
export FtpDirectoryResponse
export get_name
export get_count
export FtpListingResponse
export FtpTransferMode
export FtpStatus
export HttpRequest
export HttpResponse
export HttpMethod
export Http
export set_host
export send_request
export set_body
export set_http_version
export set_uri
export get_field
export get_body
export get_major_version
export get_minor_version
export set_field
export set_method
export HttpStatus
export IpAddress
export to_string
export Packet
export get_data_size
export read_bool
export read_string
export read_double
export write
export read_int
export SocketStatus
export SOCKET_DONE
export SOCKET_NOT_READY
export SOCKET_PARTIAL
export SOCKET_DISCONNECTED
export SOCKET_ERROR
export accept
export listen
export get_localport
export is_blocking
export set_blocking
export TcpListener
export SocketTCP
export receive
export send
export get_remoteaddress
export get_remoteport
export SocketUDP
export bind
export unbind
export max_datagram_size
export Clock
export get_elapsed_time
export restart
export Time
export as_seconds
export as_microseconds
export as_milliseconds
export seconds
export milliseconds
export microseconds
export Vector2i
export Vector2f
export Vector2u
export to_vec2u
export to_vec2f
export to_vec2i
export distance
export distance_squared
export Vector3f
export Thread
export launch
export wait
export terminate
export Event
export EventType
export get_type
export KeyEvent
export TextEvent
export MouseButtonEvent
export MouseMoveEvent
export MouseWheelEvent
export SizeEvent
export get_key
export get_text
export get_mousebutton
export get_mousemove
export get_mousewheel
export is_key_pressed
export KeyCode
export MouseButton
export get_mousepos
export set_mousepos
export is_mouse_pressed
export make_gif
export screenshot
export Vector2


text:
"""
`Session` has following public fields:

* `path` - path to the session directory
* `identifier` - identifier of session, extracted from directory name
  `ses-<session_id>`. In non-longitudinal study, this will always be "1".
* `files` - a vector of `File`
* `scans_details` - `DataFrame` parsed from `*_scans.tsv`. If the tsv is not exist, this
  will be an empty `DataFrame`

The `Subject` can be initialized by specifying only `path`. Other behavior can also be
tweaked accordingly with the optional parameters:

* `search::Bool` - defaults to `true`. If `true` this will search over `root` to find
  `Subject`, `Session` and `Files`.
* `load_metadata::Bool` - defaults to `true`. If `true`, this expects a JSON-sidecar to
  be present alongside with every BIDS files.
* `require_modality::Bool` - defaults to `true`. If `true` this expects a modality
  presents in every file name, e.g. sub-subtest_ses-1_run-001_**T1w**.nii.gz
* `longitudinal::Bool` - defaults to `true`. If `true` this expects a session directory
  `ses-` exists under subject directory.
* `strict::Bool` - defaults to `true`. If `true`, `BIDSTools` will throw errors on
  invalid BIDS filenames, this can be turned to `false` to not parse those file names
  and display a warning instead.
* `extract_from_full_path::Bool` - defaults to `true`. If `true`, will try to extract
  `sub` and `ses` from full path and append to `entities` if they can't be found in the
  parsed filename.

# Example

```julia-repl
julia> ses = Session("/path/to/bids/root/sub-Subtest/ses-1/")
Session:
    identifier = 1
    total files = 8
```
"""
struct Session
    path::String
    identifier::String
    files::Vector{File}
    scans_detail::DataFrame
end

function Session(
    path::AbstractString;
    search::Bool=true,
    load_metadata::Bool=true,
    require_modality::Bool=true,
    longitudinal::Bool=true,
    strict::Bool=true,
    extract_from_full_path::Bool=true
)
    files = Vector{File}()
    scans_detail = nothing
    identifier = nothing
    # If not longitudinal, the subject dir == ses dir and session id will
    # always be 1 (only consists of 1 session/subject)
    if longitudinal
        sep = Vector{Char}(Base.Filesystem.path_separator)
        @assert startswith(basename(rstrip(path, sep)), "ses-")
        identifier = replace(basename(rstrip(path, sep)), "ses-" => "")
    else
        identifier = "1"
    end
    # Look for _scans.tsv
    for d in readdir(path)
        if isnothing(scans_detail)
            scans_detail = isfile(joinpath(path, d)) && endswith(d, "_scans.tsv") ?
                              CSV.File(joinpath(path, d)) |> DataFrame! : nothing
        end
    end
    if search
        # Search for files
        for d in readdir(path)
            !isdir(joinpath(path, d)) && continue
            for f in readdir(joinpath(path, d))
                # Ignore json-sidecar since it'll go to the metadata
                if isfile(joinpath(path, d, f)) && !endswith(f, ".json")
                    push!(
                        files,
                        File(
                            joinpath(path, d, f),
                            load_metadata=load_metadata,
                            require_modality=require_modality,
                            strict=strict,
                            extract_from_full_path=extract_from_full_path
                        )
                    )
                end
            end
        end
    end
    if isnothing(scans_detail)
        scans_detail = DataFrame()
    end
    Session(path, identifier, files, scans_detail)
end

#-------------------------------------------------------------------------------

"""
    function total_files(session::Session)

Get number of files in a session. The argument `Session` can be changed into `Subject`
or `Layout`.
"""
function total_files(session::Session)
    length(session.files)
end

function Base.show(io::IO, session::Session)
    print(io, """
        Session:
            identifier = $(session.identifier)
            total files = $(total_files(session))""")
end

"""
    function list_scans_detail(session::Session)

Function to pretty print scans detail spreadsheet (_scans.tsv) using PrettyTables
"""
function list_scans_detail(session::Session)
    pretty_table(session.scans_detail, crop= :none)
end

# Docstring of get_files is in File.jl
function get_files(
    session::Session; path::Union{String, Regex, Nothing}=nothing, kws...
)
    result = get_files(session.files; path=path, kws...)
    return result
end


text:
# This file is a part of TextSearch.jl

export compute_bow_corpus, compute_bow, bow, dvec
const BOW = DVEC{UInt64,Int32}

"""
    compute_bow(tokenlist::AbstractVector{S}, bow::BOW=BOW()) where {S<:Symbol}
    compute_bow(tok::Tokenizer, text::Symbol, bow::BOW=BOW())
    compute_bow(tok::Tokenizer, text::AbstractVector, bow::BOW=BOW())

Creates a bag of words from the given text (a string or a list of strings).
If bow is given then updates the bag with the text.
When `config` is given, the text is parsed according to it.
"""
function compute_bow(tokenlist::AbstractVector, bow::BOW=BOW())
    for sym in tokenlist
        bow[sym] = get(bow, sym, 0) + 1
    end

    bow
end

function compute_bow(tok::Tokenizer, text::AbstractString, bow::BOW=BOW())
    compute_bow(tokenize(tok, text), bow)
end

"""
    compute_bow(messages::AbstractVector, bow::BOW=BOW())

Computes a bag of words from messages
"""
function compute_bow(tok::Tokenizer, messages::AbstractVector, bow::BOW=BOW())
    for text in messages
        empty!(tok)
        tokenize(tok, text)
        compute_bow(tok.tokens, bow)
    end

    bow
end

"""
    compute_bow_corpus(tok::Tokenizer, corpus::AbstractVector)

Computes a list of bag of words from a corpus
"""
function compute_bow_corpus(tok::Tokenizer, corpus::AbstractVector, bow::BOW=BOW())
    X = Vector{BOW}(undef, length(corpus))
    for i in eachindex(corpus)
        empty!(tok)
        empty!(bow)
        compute_bow(tok, corpus[i], bow)
        X[i] = copy(bow)
    end

    X
end

"""
    bow(model::Tokenizer, x::AbstractSparseVector)
    bow(model::Tokenizer, x::DVEC{Ti,Tv}) where {Ti<:Integer,Tv<:Number}

Creates a bag of words using the sparse vector `x` and the text model `model`
"""
function bow(m::Tokenizer, x::AbstractSparseVector)
    DVEC{String,eltype{x.nzval}}(decode(m, t) => v for (t, v) in zip(x.nzind, x.nzval))
end

function bow(m::Tokenizer, x::DVEC{Ti,Tv}) where {Ti<:Integer,Tv<:Number}
    DVEC{String,Tv}(decode(m, t) => v for (t, v) in x)
end

"""
    dvec(m::Tokenizer, x::DVEC{Symbol,Tv}, Ti=Int) where {Tv<:Number}

Creates a DVEC sparse vector from a bag of words sparse vector (i.e., with type DVEC{Symbol,Tv}),
using the inverse map `m`
"""
function dvec(m::Tokenizer, x::DVEC{String,Tv}, Ti=Int) where {Tv<:Number}
    DVEC{Ti,Tv}(encode(m, t) => v for (t, v) in x)
end


text:
# SPDX-License-Identifier: X11
# 2020-10-21
# Cream puff (300pt)

N = parse(Int, readline())
A = []

push!(A, 1)
for i ∈ 2:floor(Int, sqrt(N))
  if N % i == 0
    push!(A, i)
    i != N ÷ i && push!(A, N ÷ i)
  end
end
N != 1 && push!(A, N)
sort!(A)

for i ∈ A
  println(i)
end

max_stars_repo_name: jthvai/competitive-programming

max_issues_repo_name: jthvai/competitive-programming

max_forks_repo_name: jthvai/competitive-programming


text:
"""
    ExponentiatedKernel()

Exponentiated kernel.

# Definition

For inputs ``x, x' \\in \\mathbb{R}^d``, the exponentiated kernel is defined as
```math
k(x, x') = \\exp(x^\\top x').
```
"""
struct ExponentiatedKernel <: SimpleKernel end

kappa(::ExponentiatedKernel, xᵀy::Real) = exp(xᵀy)

metric(::ExponentiatedKernel) = DotProduct()

iskroncompatible(::ExponentiatedKernel) = true

Base.show(io::IO, ::ExponentiatedKernel) = print(io, "Exponentiated Kernel")