text:
using Revise
using NNFEM
using PyPlot

using Revise
using ADCME
using PyCall
using Test
using MAT
using NNFEM
using PyPlot
using JLD2
using LinearAlgebra
using Distributions, Random
using ForwardDiff
using DelimitedFiles

np = pyimport("numpy")
"""
Force load on one edge of the plate,
The edge is [0,L], which is discretized to ne elements, with porder-th order polynomials
type: a string of load type, which can be "constant", or "Gaussian"
args: an array,
      for Constant load, it has p1 and p2, in tangential and normal direction
      for "Gaussian" load, it has p, x0, and σ, the force is in the normal direction

"""
function ComputeLoad(L, ne, porder, ngp, type,  args)
    @assert(ngp <= 4)
    dx = L/ne

    xx = np.linspace(0, L, ne*porder + 1)

    Ft, Fn = zeros(Float64, ne*porder + 1), zeros(Float64, ne*porder + 1)

    elem_coords = zeros(Float64, porder + 1, 2)

    # construct pressure load function
    if type == "Constant"
        pt, pn = args
        f_px, f_py = x->pt, x->pn
    elseif type == "Gaussian"
        pn, x0, σ = args
        f_px, f_py = x->0, x-> pn * 1.0/(sqrt(2*pi*σ^2)) * exp.(-0.5*(x .- x0).^2/σ^2)
    else
        error("Force load type is not recogonized ", type)
    end

    # assemble the force

    @show ne
    for e = 1:ne
        if porder == 1
            loc_id = [e, e+1]
        elseif porder == 2
            loc_id = [2*(e-1)+1, 2*(e-1)+3, 2*(e-1)+2]
        else
            error("porder error porder ==", porder)
        end

        loc_xx = xx[loc_id]         #x-coordinates of element nodes
        elem_coords[:,1] .= loc_xx

        #return list
        weights, hs = get1DElemShapeData( elem_coords, ngp)

        for igp = 1:ngp
            gp_xx = loc_xx' * hs[igp]   #x-coordinates of element Gaussian points
            Ft[loc_id] += f_px(gp_xx) * hs[igp] * weights[igp]
            Fn[loc_id] += f_py(gp_xx) * hs[igp] * weights[igp]

        end
    end

    return Ft, Fn
end

function BoundaryCondition(tid, nx, ny, porder=2, Lx = 1.0, Ly = 0.5; force_scale=5.0)
    nnodes, neles = (nx*porder + 1)*(ny*porder + 1), nx*ny

    x = np.linspace(0.0, Lx, nx*porder + 1)
    y = np.linspace(0.0, Ly, ny*porder + 1)

    X, Y = np.meshgrid(x, y)
    nodes = zeros(nnodes,2)
    nodes[:,1], nodes[:,2] = X'[:], Y'[:]
    ndofs = 2

    EBC, g = zeros(Int64, nnodes, ndofs), zeros(nnodes, ndofs)
    FBC, fext = zeros(Int64, nnodes, ndofs), zeros(nnodes, ndofs)
    gt = nothing

    #todo only change 203

    # setting EBC:
    # tid = 1XX   |==|-->
    # tid = 2XX
    if div(tid,100)==1
        EBC[collect(1:nx*porder+1), :] .= -1 # fix bottom
        FBC[collect((nx*porder+1)*ny*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -2 # force on the top
    elseif div(tid,100)==2
        EBC[collect(1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -1 # fix left
        FBC[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -2 # force on the right
    elseif div(tid,100)==3
        EBC[collect(1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -1 # fix left
        FBC[collect(2:nx*porder+1), :] .= -2 # force on the bottom
    end

    #
    if tid==100
        F1, F2 = ComputeLoad(Lx, nx, porder, ngp, "Constant",  [0, P1])
        fext[collect((nx*porder+1)*ny*porder+1:(nx*porder+1)*(ny*porder+1)), 1] .= F1
        fext[collect((nx*porder+1)*ny*porder+1:(nx*porder+1)*(ny*porder+1)), 2] .= F2

    elseif tid == 200
        F1, F2 = ComputeLoad(Ly, ny, porder, ngp, "Constant",  [P1, 0])
        fext[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 1] .= F1
        fext[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 2] .= F2

    elseif tid == 300

        F1, F2 = ComputeLoad(Lx, nx, porder, ngp, "Gaussian",  [Lx*P2, Lx*5.0/6.0, 0.2*Lx])
        fext[collect(2:nx*porder+1), 1] .= F1[2:end]
        fext[collect(2:nx*porder+1), 2] .= F2[2:end]
    else
        error("tid = $tid is not understood")
    end

    dof_to_active = findall(FBC[:].==-2)
    ft = t->fext[:][dof_to_active]*sin(π*t/(T))

    npoints = (nx+1)*(ny+1)
    node_to_point = zeros(Int64, nnodes) .- 1
    pointid = 1
    for j = 1:2:(ny*porder + 1)
        for i = 1:2:(nx*porder + 1)
            node_to_point[(j-1)*(nx*porder + 1) + i] = pointid
            pointid += 1
        end
    end

    return nodes, EBC, g, gt, FBC, fext, ft, npoints, node_to_point
end

force_scale = 2.0 #50
tid =200
time = 0.03
fiber_size = 1
porder = 2
T = 0.2
NT = 200

P1 = 44800 /force_scale #gcm/ms^2 compress/pull
P2 =  4480 /force_scale #gcm/ms^2 bend
P3 = 16800 /force_scale

@show P1, P2, force_scale
ngp = 3

testtype = "PlaneStress"
prop = Dict("name"=> testtype, "rho"=> 0.800, "E"=>2.0e6, "nu"=>0.2)

nxc, nyc = 10,5
nx, ny =  nxc*fiber_size, nyc*fiber_size
Lx, Ly = 0.1, 0.05 #m
nodes, EBC, g, gt, FBC, fext, ft, npoints, node_to_point = BoundaryCondition(tid, nx, ny, porder, Lx, Ly; force_scale=force_scale)

function Edge_func(x, y, t, tid)
    zero = similar(x) * 0.0
    if (tid == 100)
        f1 = @. zero
        f2 = @. P1*sin(π*t/(T)) + zero
    end
    if (tid == 200)
        f1 = @. P1*sin(π*t/(T)) + zero
        f2 = @. zero
    end
    if (tid == 300)
        pn, x0, σ = [Lx*P2, Lx*5.0/6.0, 0.2*Lx]
        f1 = @. zero
        f2 = @. pn * 1.0/(sqrt(2*pi*σ^2)) * exp.(-0.5*(x .- x0).^2/σ^2)*sin(π*t/(T))
    end

    return [f1 f2]
end

if tid == 100
    Edge_Traction_Data = zeros(Int64, nx, 3)
    for i = 1:nx
        Edge_Traction_Data[i, :] .= i+(ny-1)*nx, 3, 100
    end
elseif tid == 200
    Edge_Traction_Data = zeros(Int64, ny, 3)
    for i = 1:ny
        Edge_Traction_Data[i, :] .= nx+(i-1)*nx, 2, 200
    end
elseif tid == 300
    Edge_Traction_Data = zeros(Int64, nx, 3)
    for i = 1:nx
        Edge_Traction_Data[i, :] .= i, 1, 300
    end
end

elements = []
for j = 1:ny
    for i = 1:nx
        n = (nx*porder+1)*(j-1)*porder + (i-1)porder+1
        #element (i,j)
        if porder == 1
            #   4 ---- 3
            #
            #   1 ---- 2

            elnodes = [n, n + 1, n + 1 + (nx + 1), n + (nx + 1)]
        elseif porder == 2
            #   4 --7-- 3
            #   8   9   6
            #   1 --5-- 2
            elnodes = [n, n + 2, n + 2 + 2*(2*nx+1),  n + 2*(2*nx+1), n+1, n + 2 + (2*nx+1), n + 1 + 2*(2*nx+1), n + (2*nx+1), n+1+(2*nx+1)]
        else
            error("polynomial order error, porder= ", porder)
        end

        coords = nodes[elnodes,:]
        push!(elements,FiniteStrainContinuum(coords,elnodes, prop, 3))
    end
end

ndofs = 2
domain = Domain(nodes, elements, ndofs, EBC, g, FBC, fext, Edge_Traction_Data)
setGeometryPoints!(domain, npoints, node_to_point)

state = zeros(domain.neqs)
∂u = zeros(domain.neqs)
globdat = GlobalData(state,zeros(domain.neqs), zeros(domain.neqs),∂u, domain.neqs, gt, ft, nothing, Edge_func)

assembleMassMatrix!(globdat, domain)
updateStates!(domain, globdat)

fold = zeros(Float64, size(domain.fext))
ID = domain.ID
nodal_force = ft(time) # user defined time-dependent boundary
dof_id = 0
#update fext for active nodes (length of neqs)
for idof = 1:domain.ndims
    for inode = 1:domain.nnodes
        if (domain.FBC[inode, idof] == -2)
            global dof_id
            dof_id += 1
            fold[ID[inode, idof]] = nodal_force[dof_id]
        end
    end
end

fnew = getEdgeForce(domain, globdat, time)

@show norm(fold), norm(fnew)
@show norm(fold - fnew)


text:
function ObservationDataSet(arg0::RinexHeader, arg1::SatelliteSystem, arg2::jint, arg3::AbsoluteDate, arg4::jdouble, arg5::List)
    return ObservationDataSet((RinexHeader, SatelliteSystem, jint, AbsoluteDate, jdouble, List), arg0, arg1, arg2, arg3, arg4, arg5)
end

function get_date(obj::ObservationDataSet)
    return jcall(obj, "getDate", AbsoluteDate, ())
end

function get_header(obj::ObservationDataSet)
    return jcall(obj, "getHeader", RinexHeader, ())
end

function get_observation_data(obj::ObservationDataSet)
    return jcall(obj, "getObservationData", List, ())
end

function get_prn_number(obj::ObservationDataSet)
    return jcall(obj, "getPrnNumber", jint, ())
end

function get_rcvr_clk_offset(obj::ObservationDataSet)
    return jcall(obj, "getRcvrClkOffset", jdouble, ())
end

function get_satellite_system(obj::ObservationDataSet)
    return jcall(obj, "getSatelliteSystem", SatelliteSystem, ())
end

ext: jl


text:
mutable struct DriverBehavior
    desired_speed        :: Float64
    desired_acceleraion  :: Float64
    desired_deceleration :: Float64
    curve_speed          :: Float64
    observe_speed_limits :: Float64
    distance_keeping     :: Float64
    lane_keeping         :: Float64
    speed_keeping        :: Float64
    lane_change_dynamic  :: Float64
    urge_to_overtake     :: Float64
    keep_right_rule      :: Float64
    respond_to_tailgate  :: Float64
    foresight_distance   :: Float64
    steering_distance    :: Float64
    use_of_indicator     :: Float64
    obey_traffic_signs   :: Bool
    obey_traffic_lights  :: Bool
end

const HASTY_DRIVER       = DriverBehavior(0.85,0.80,0.80,0.80,0.35,0.35,0.7,0.7,0.75,0.85,0.5,0.5,0.50,0.7,0.30,true,true)
const BRISK_DRIVER       = DriverBehavior(0.70,0.65,0.65,0.60,0.45,0.40,0.5,0.5,0.60,0.60,0.5,0.5,0.55,0.5,0.65,true,true)
const DEFAULT_DRIVER     = DriverBehavior(0.50,0.50,0.50,0.50,0.50,0.50,0.5,0.5,0.50,0.50,0.5,0.5,0.50,0.5,0.50,true,true)
const COMFORTABLE_DRIVER = DriverBehavior(0.45,0.40,0.40,0.40,0.60,0.55,0.5,0.5,0.45,0.40,0.5,0.5,0.60,0.4,0.75,true,true)
const INSECURE_DRIVER    = DriverBehavior(0.40,0.35,0.65,0.35,0.90,0.65,0.3,0.3,0.35,0.20,0.5,0.5,0.30,0.3,0.20,true,true)

function Base.copy(d::DriverBehavior)
    DriverBehavior(
        d.desired_speed,
        d.desired_acceleraion,
        d.desired_deceleration,
        d.curve_speed,
        d.observe_speed_limits,
        d.distance_keeping,
        d.lane_keeping,
        d.speed_keeping,
        d.lane_change_dynamic,
        d.urge_to_overtake,
        d.keep_right_rule,
        d.respond_to_tailgate,
        d.foresight_distance,
        d. steering_distance,
        d. use_of_indicator,
        d.obey_traffic_signs,
        d.obey_traffic_lights
    )
end

max_stars_repo_name: sisl/VirtualTestDrive

max_issues_repo_name: sisl/VirtualTestDrive

max_forks_repo_name: sisl/VirtualTestDrive


text:
import ChainRulesCore: frule, rrule
using LinearAlgebra

function compute_gn_jac(u_n::AbstractVector{T}) where T
    n = length(u_n) - 1
    g_n = zero(u_n)
    # using `similar` in case u_n is e.g., `CuArray`
    dgdu = fill!(similar(u_n, n + 1, n + 1), zero(T))

    ## First: calculate the a_n terms (eqn. 10)
    a_n = zero(g_n)
    dadu = copy(dgdu)
    # constant term
    a_n[begin] = one(T)
    for i in 1:n
        bcoeff = one(T) # binomial coefficient
        for j in 0:i
            r = (-1)^j * bcoeff
            a_n[begin + j] -= r * u_n[begin + i]
            dadu[begin + i, begin + j] -= r
            bcoeff *= (i - j) / (j + 1)
        end
    end

    ## Second: calculate the g_n terms (eqn. 15)
    for j in n:-1:2
        if j ≥ n - 1
            g_n[begin + j] = a_n[begin + j] / (j + 2)
            for i in 1:n
                dgdu[begin + i, begin + j] = dadu[begin + i, begin + j] / (j + 2)
            end
        else
            g_n[begin + j] = a_n[begin + j] / (j + 2) + g_n[begin + j + 2]
            for i in 1:n
                dgdu[begin + i, begin + j] = dadu[begin + i, begin + j] / (j + 2) + dgdu[begin + i, begin + j + 2]
            end
        end
    end
    if n ≥ 3
        g_n[begin + 1] = a_n[begin + 1] + 3 * g_n[begin + 3]
        for i in 1:n
            dgdu[begin + i, begin + 1] = dadu[begin + i, begin + 1] + 3 * dgdu[begin + i, begin + 3]
        end
    elseif n ≥ 1
        g_n[begin + 1] = a_n[begin + 1]
        for i in 1:n
            dgdu[begin + i, begin + 1] = dadu[begin + i, begin + 1]
        end
    end
    if n ≥ 2
        g_n[begin] = a_n[begin] + 2 * g_n[begin + 2]
        for i in 1:n
            dgdu[begin + i, begin] = dadu[begin + i, begin] + 2 * dgdu[begin + i, begin + 2]
        end
    else
        g_n[begin] = a_n[begin]
        for i in 1:n
            dgdu[begin + i, begin] = dadu[begin + i, begin]
        end
    end

    return g_n, dgdu
end

function compute_uniform_grad(b::T, r; r2, b2, sqarea, fourbrinv) where T
    if b ≤ 1 - r
        flux = π * (1 - r2)
        kap0 = convert(T, π)
        kck = zero(T)
        kite_area2 = zero(T)
        ∇s0 = SA[zero(T), -2 * π * r]
    else
        kite_area2 = sqrt(sqarea)
        r2m1 = (r - 1) * (r + 1)
        kap0 = atan(kite_area2, r2m1 + b2)
        Πmkap1 = atan(kite_area2, r2m1 - b2)
        kck = kite_area2 * fourbrinv
        flux = Πmkap1 - r2 * kap0 + 0.5 * kite_area2
        ∇s0 = SA[kite_area2 / b, -2 * r * kap0]
    end

    return (flux, kap0, kite_area2, kck), ∇s0
end

function compute_linear_grad(b::T, r; k2, kc, kc2, onembmr2, onembpr2, onembmr2inv, sqonembmr2, r2=r^2, b2=b^2, br=b * r, fourbr=4 * br, sqbr=sqrt(br)) where T
    if iszero(b) # case 10
        Λ1 = -2 * π * sqrt(1 - r2)^3
        Eofk = 0.5 * π
        Em1mKdm = 0.25 * π

        ∇s1 = SA[zero(T), r * Λ1]
    elseif b == r
        if r == 0.5 # case 6
            Λ1 = π - 4 / 3
            Eofk = one(T)
            Em1mKdm = one(T)
            ∇s1 = SA[2 / 3, -2]
        elseif r < 0.5 # case 5
            m = 4 * r2
            Eofk = cel(m, one(T), one(T), 1 - m)
            Em1mKdm = cel(m, one(T), one(T), zero(T))
            Λ1 = π + 2 / 3 * ((2 * m - 3) * Eofk - m * Em1mKdm)
                # + (b - r) * 4 * r * (Eofk - 2 * Em1mKdm)
            ∇s1 = SA[-4 / 3 * r * (Eofk - 2 * Em1mKdm), -4 * r * Eofk]
        else # case 7
            m = 4 * r2
            minv = inv(m)
            Eofk = cel(minv, one(T), one(T), 1 - minv)
            Em1mKdm = cel(minv, one(T), one(T), zero(T))
            Λ1 = π + ((2 * m - 3) * Em1mKdm - m * Eofk) / (3 * r)
            #  -(b - r) * 2 * (2 * Eofk - Em1mKdm)
            ∇s1 = SA[2 / 3 * (2 * Eofk - Em1mKdm), -2 * Em1mKdm]
        end
    else
        if b + r > 1 # case 2, case 8
            Πofk, Eofk, Em1mKdm = cel(k2, kc, (b - r)^2 * kc2, zero(T), one(T), one(T), 3 * kc2 * (b - r) * (b + r), kc2, zero(T))
            Λ1 = onembmr2 * (Πofk + (-3 + 6 * r2 + 2 * br) * Em1mKdm - fourbr * Eofk) / (3 * sqbr)
            ∇s1 = SA[2 * r * onembmr2 * (-Em1mKdm + 2 * Eofk) / (sqbr * 3), -2 * r * onembmr2 * Em1mKdm / sqbr]
        elseif b + r < 1 # case 3 case 9
            bmrdbpr = (b - r) / (b + r)
            μ = 3 * bmrdbpr * onembmr2inv
            p = bmrdbpr^2 * onembpr2 * onembmr2inv
            Πofk, Eofk, Em1mKdm = cel(inv(k2), kc, p, 1 + μ, one(T), one(T), p + μ, kc2, zero(T))
            Λ1 = 2 * sqonembmr2 * (onembpr2 * Πofk - (4 - 7 * r2 - b2) * Eofk) / 3
            ∇s1 = SA[-4 / 3 * r * sqonembmr2 * (Eofk - 2 * Em1mKdm), -4 * r * sqonembmr2 * Eofk]
        else # case 4
            Λ1 = 2 * acos(1 - 2 * r) - 2 * π * (r > 0.5) - (4 / 3 * (3 + 2 * r - 8 * r2) + 8 * (r + b - 1) * r) * sqrt(r * (1 - r))
            Eofk = one(T)
            Em1mKdm = one(T)
            g = -8 * r * sqrt(r * (1 - r))
            ∇s1 = SA[-g / 3, g]
        end
    end

    flux = ((1 - T(r > b)) * 2π - Λ1) / 3

    return (flux, Eofk, Em1mKdm), ∇s1
end

function compute_quadratic_grad(b, r; s0, r2, b2, kap0, kite_area2, k2, ∇s0)
    r2pb2 = r2 + b2
    η2 = r2 * (r2pb2 + b2)
    if k2 > 1
        four_pi_eta = 2 * π * (η2 - 1)
        ∇η = 4 * π * SA[2 * b * r2, 2 * r * r2pb2]
    else
        Πmkap1 = atan(kite_area2, (r - 1) * (r + 1) - b2)
        four_pi_eta = 2 * (-Πmkap1 + η2 * kap0 - 0.25 * kite_area2 * (1 + 5 * r2 + b2))
        ∇η = SA[2 / b * (4 * b2 * r2 * kap0 - (1 + r2pb2) * kite_area2),
                8 * r * (r2pb2 * kap0 - kite_area2)]
    end
    s2 = 2 * s0 + four_pi_eta
    ∇s2 = 2 * ∇s0 + ∇η
    return  s2, ∇s2
end

function compute_grad(ld::PolynomialLimbDark, b::S, r) where S
    T = float(S)
    bcut = 1e-3
    n = ld.n_max
    dfdg = zero(ld.g_n)

    ## check for trivial cases
    if b ≥ 1 + r || iszero(r)
        # completely unobscured
        return one(T), dfdg, zero(T), zero(T)
    elseif r ≥ 1 + b
        # completely obscured
        return zero(T), dfdg, zero(T), zero(T)
    end

    r2 = r^2
    b2 = b^2

    if iszero(b)
        onemr2 = 1 - r2
        sqrt1mr2 = sqrt(onemr2)

        # annular ellipse
        facd = -2 * r
        if ld.n_max == 0
            flux = ld.g_n[begin] * onemr2
            dfdr = ld.g_n[begin] * facd
        else
            flux = ld.g_n[begin] * onemr2 + 2 / 3 * ld.g_n[begin + 1] * sqrt1mr2^3
            dfdr = ld.g_n[begin] * facd + ld.g_n[begin + 1] * facd * sqrt1mr2
        end
        fac = 2 * r2 * onemr2
        for i in 2:ld.n_max
            flux -= ld.g_n[begin + i] * fac
            dfdr += ld.g_n[begin + i] * facd * (2 * onemr2 - i * r2)
            dfdg[begin + i] -= fac
            fac *= sqrt1mr2
            facd *= sqrt1mr2
        end
        dfdr *= π * ld.norm
        dfdg[1] = (onemr2 - flux) * π * ld.norm
        dfdg[2] = 2 / 3 * (sqrt1mr2^3 - flux) * π * ld.norm
        return flux * π * ld.norm, dfdg * ld.norm, zero(T), dfdr
    end

    # take a moment to calculate values used repeatedly in the
    # numerical routines ahead
    onembmr2 = (r - b + 1) * (1 - r + b)
    onembmr2inv = inv(onembmr2)
    sqonembmr2 = sqrt(onembmr2)
    br = b * r
    fourbr = 4 * br
    fourbrinv = inv(fourbr)
    sqbr = sqrt(br)
    sqbrinv = inv(sqbr)
    onembpr2 = (1 - r - b) * (1 + b + r)
    sqarea = sqarea_triangle(one(T), r, b)
    k2 = max(zero(T), onembmr2 * fourbrinv)
    k = sqrt(k2)
    omemr2 = (1 - r) * (1 + r)
    onemr2mb2 = omemr2 - b2
    onemr2pb2 = omemr2 + b2
    if k2 ≥ 1
        if k2 > 2
            kc2 = 1 - inv(k2)
        else
            kc2 = onembpr2 * onembmr2inv
        end
    else
        if k2 > 0.5
            kc2 = (r - 1 + b) * (b + r + 1) * fourbrinv
        else
            kc2 = 1 - k2
        end
    end
    kc = sqrt(abs(kc2))

    ## Compute uniform term

    (s0, kap0, kite_area2, kck), ∇s0 = compute_uniform_grad(b, r; b2, r2, sqarea, fourbrinv)

    flux = ld.g_n[begin] * s0
    ∇flux = ld.g_n[begin] * ∇s0
    dfdg[begin] = s0

    if ld.n_max == 0
        dfdg[begin] -= flux * ld.norm * π
        dfdb, dfdr = ∇flux * ld.norm
        return flux * ld.norm, dfdg * ld.norm, dfdb, dfdr
    end

    ## compute linear term
    (s1, Eofk, Em1mKdm), ∇s1 = compute_linear_grad(b, r; k2, kc, kc2, r2, b2, br, fourbr, sqbr, onembmr2, onembpr2, onembmr2inv, sqonembmr2)

    flux += ld.g_n[begin + 1] * s1
    ∇flux += ld.g_n[begin + 1] * ∇s1
    dfdg[begin + 1] = s1

    if ld.n_max == 1
        dfdg[begin] -= flux * ld.norm * π
        dfdg[begin + 1] -= flux * ld.norm * π * 2 / 3
        dfdb, dfdr = ∇flux * ld.norm
        return flux * ld.norm, dfdg * ld.norm, dfdb, dfdr
    end

    ## calculate quadratic term
    s2, ∇s2 = compute_quadratic_grad(b, r; s0, r2, b2, kap0, kite_area2, k2, ∇s0)
    flux += ld.g_n[begin + 2] * s2
    ∇flux = ∇flux + ld.g_n[begin + 2] * ∇s2
    dfdg[begin + 2] = s2
    if ld.n_max == 2
        dfdg[begin] -= flux * ld.norm * π
        dfdg[begin + 1] -= flux * ld.norm * π * 2 / 3
        dfdb, dfdr = ∇flux * ld.norm
        return flux * ld.norm, dfdg * ld.norm, dfdb, dfdr
    end

    ## higher order (N > 3) terms require solving Mn and Nn integrals
    if k2 < 0.5 && ld.n_max > 3
        downwardM!(ld.Mn, ld.Mn_coeff; n_max=ld.n_max, sqonembmr2, sqbr, onemr2mb2, k, k2, sqarea, kap0, Eofk, Em1mKdm, kite_area2)
        if b < bcut
            Nn_lower!(ld.Nn, ld.Nn_coeff; Mn=ld.Mn, n_max=ld.n_max, onembpr2, sqonembmr2, kap0, kc, k2, sqbr, k, Eofk, Em1mKdm)
        end
    else
        upwardM!(ld.Mn; sqbr, n_max=ld.n_max, sqonembmr2, onemr2mb2, sqarea, k2, kap0, Eofk, Em1mKdm, kite_area2)
        if b < bcut
            Nn_raise!(ld.Nn, ld.Mn; n_max=ld.n_max, onembpr2, kap0, kc, k2, sqbr, k, Eofk, Em1mKdm)
        end
    end

    # compute remaining terms
    for n in 3:ld.n_max
        pofgn_M = 2 * r2 * ld.Mn[begin + n] - n / (n + 2) *
                       (onemr2mb2 * ld.Mn[begin + n] + sqarea * ld.Mn[begin + n - 2])
        flux -= ld.g_n[begin + n] * pofgn_M
        dfdg[begin + n] = -pofgn_M
        dpdr_M = 2 * r * ((n + 2) * ld.Mn[begin + n] - n * ld.Mn[begin + n - 2])

        if b < bcut
            dpdb_M = n * (ld.Mn[begin + n - 2] * (2 * r^3 + b^3 - b - 3 * r2 * b) + b * ld.Mn[begin + n] - 4 * r^3 * ld.Nn[begin + n - 2])
        else
            dpdb_M = n / b * ((ld.Mn[begin + n] - ld.Mn[begin + n - 2]) * (r2 + b2) + (b2 - r2)^2 * ld.Mn[begin + n - 2])
        end
        ∇flux -= ld.g_n[begin + n] * SA[dpdb_M, dpdr_M]

    end
    dfdg[begin] -= flux * ld.norm * π
    dfdg[begin + 1] -= flux * ld.norm * π * 2 / 3
    dfdb, dfdr = ∇flux * ld.norm
    return flux * ld.norm, dfdg * ld.norm, dfdb, dfdr
end

function frule((_, Δld, Δb, Δr), ::typeof(compute), ld::LD, b, r) where {LD <: PolynomialLimbDark}
    f, dfdg, dfdb, dfdr = compute_grad(ld, b, r)
    ∂g_n = dot(dfdg, Δld.g_n)
    return f, ∂g_n + dfdb * Δb + dfdr * Δr
end

function rrule(::typeof(compute), ld::LD, b, r) where {LD <: PolynomialLimbDark}
    f, dfdg, dfdb, dfdr = compute_grad(ld, b, r)
    function compute_pullback(Δf)
        ∂ld = Tangent{LD}(g_n=dfdg * Δf)
        ∂b = dfdb * Δf
        ∂r = dfdr * Δf
        return NoTangent(), ∂ld, ∂b, ∂r
    end
    return f, compute_pullback
end

function frule((_, Δu_n), ::Type{<:PolynomialLimbDark}, u::AbstractVector{S}; maxiter=100) where S
    T = float(S)
    # add constant u_0 term
    n_max = length(u)
    u_n = pushfirst!(copy(u), -one(S))
    # get Green's basis coefficients
    g_n, ∇g_n = compute_gn_jac(u_n)
    # compute series expansion for M_n and N_n
    Mn_coeff = compute_Mn_coeffs(T, n_max; maxiter=maxiter)
    Nn_coeff = compute_Nn_coeffs(T, n_max; maxiter=maxiter)

    # calculate flux normalization factor, which only depends on first two terms
    if length(g_n) > 1
        norm = inv(π * (g_n[1] + 2 / 3 * g_n[2]))
    else
        norm = inv(π * g_n[1])
    end

    # pre-allocate temp arrays
    Mn = similar(g_n)
    Nn = similar(g_n)

    Ω = PolynomialLimbDark(n_max, u_n, g_n, Mn_coeff, Nn_coeff, norm, Mn, Nn)
    ∂g_n = @views ∇g_n[begin + 1:end, :]' * Δu_n
    ∂Ω = Tangent{typeof(Ω)}(g_n=∂g_n)
    return Ω, ∂Ω
end

function rrule(::Type{<:PolynomialLimbDark}, u::AbstractVector{S}; maxiter=100) where S
    T = float(S)
    # add constant u_0 term
    n_max = length(u)
    u_n = pushfirst!(copy(u), -one(S))
    # get Green's basis coefficients
    g_n, ∇g_n = compute_gn_jac(u_n)
    # compute series expansion for M_n and N_n
    Mn_coeff = compute_Mn_coeffs(T, n_max; maxiter=maxiter)
    Nn_coeff = compute_Nn_coeffs(T, n_max; maxiter=maxiter)

    # calculate flux normalization factor, which only depends on first two terms
    if length(g_n) > 1
        norm = inv(π * (g_n[1] + 2 / 3 * g_n[2]))
    else
        norm = inv(π * g_n[1])
    end

    # pre-allocate temp arrays
    Mn = similar(g_n)
    Nn = similar(g_n)

    Ω = PolynomialLimbDark(n_max, u_n, g_n, Mn_coeff, Nn_coeff, norm, Mn, Nn)
    ∂g_n = @view ∇g_n[begin + 1:end, :]

    proj = ProjectTo(u)

    function PolynomialLimbDark_pullback(Δld)
        ∂u = ∂g_n * Δld.g_n
        return NoTangent(), proj(∂u)
    end
    return Ω, PolynomialLimbDark_pullback
end


text:
@testset "Function tests" begin
    w = MOI.VariableIndex(0)
    x = MOI.VariableIndex(1)
    y = MOI.VariableIndex(2)
    z = MOI.VariableIndex(3)
    @testset "operate vcat" begin
        v = MOI.VectorOfVariables([y, w])
        f = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2, 4],
                                                           [x, z]), 5)
        g = MOI.VectorAffineFunction(MOI.VectorAffineTerm.([3, 1],
                                                           MOI.ScalarAffineTerm.([5, 2],
                                                                                [y, x])),
                                                           [3, 1, 4])
        wf = MOI.SingleVariable(w)
        xf = MOI.SingleVariable(x)
        @test MOIU.promote_operation(vcat, Int, typeof(wf), typeof(f),
                                     typeof(v), Int, typeof(g), typeof(xf),
                                     Int) == MOI.VectorAffineFunction{Int}
        F = MOIU.operate(vcat, Int, wf, f, v, 3, g, xf, -4)
        @test F.terms == MOI.VectorAffineTerm.([1, 2, 2, 3, 4, 8, 6, 9],
                                               MOI.ScalarAffineTerm.([1, 2, 4, 1, 1, 5, 2, 1],
                                                                     [w, x, z, y, w, y, x, x]))
        @test F.constants == [0, 5, 0, 0, 3, 3, 1, 4, 0, -4]
    end
    @testset "MultirowChange construction" begin
        chg1 = MOI.MultirowChange(w, [(Int32(2), 2.0), (Int32(1), 3.0)])
        chg2 = MOI.MultirowChange(w, [(Int64(2), 2.0), (Int64(1), 3.0)])
        @test chg1.variable == chg2.variable
        @test chg1.new_coefficients == chg2.new_coefficients
    end
    @testset "VectorAffineTerm/VectorQuadraticTerm construction" begin
        scalaraffine = MOI.ScalarAffineTerm(2.0, z)
        @test MOI.VectorAffineTerm(Int32(3), scalaraffine) === MOI.VectorAffineTerm(Int64(3), scalaraffine)
        @test MOI.VectorAffineTerm{Float64}(Int32(3), scalaraffine) === MOI.VectorAffineTerm(Int64(3), scalaraffine)
        scalarquad = MOI.ScalarQuadraticTerm(2.0, y, z)
        @test MOI.VectorQuadraticTerm(Int32(3), scalarquad) === MOI.VectorQuadraticTerm(Int64(3), scalarquad)
        @test MOI.VectorQuadraticTerm{Float64}(Int32(3), scalarquad) === MOI.VectorQuadraticTerm(Int64(3), scalarquad)
    end
    @testset "evalvariables" begin
        # We do tests twice to make sure the function is not modified
        vals = Dict(w=>0, x=>3, y=>1, z=>5)
        fsv = MOI.SingleVariable(z)
        @test MOI.output_dimension(fsv) == 1
        @test MOIU.evalvariables(vi -> vals[vi], fsv) ≈ 5
        @test MOIU.evalvariables(vi -> vals[vi], fsv) ≈ 5
        fvv = MOI.VectorOfVariables([x, z, y])
        @test MOI.output_dimension(fvv) == 3
        @test MOIU.evalvariables(vi -> vals[vi], fvv) ≈ [3, 5, 1]
        @test MOIU.evalvariables(vi -> vals[vi], fvv) ≈ [3, 5, 1]
        fsa = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x), MOI.ScalarAffineTerm(3.0, z), MOI.ScalarAffineTerm(2.0, y)], 2.0)
        @test MOI.output_dimension(fsa) == 1
        @test MOIU.evalvariables(vi -> vals[vi], fsa) ≈ 22
        @test MOIU.evalvariables(vi -> vals[vi], fsa) ≈ 22
        fva = MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 1, 2], MOI.ScalarAffineTerm.([1.0, 3.0, 2.0], [x, z, y])), [-3.0, 2.0])
        @test MOI.output_dimension(fva) == 2
        @test MOIU.evalvariables(vi -> vals[vi], fva) ≈ [12, 7]
        @test MOIU.evalvariables(vi -> vals[vi], fva) ≈ [12, 7]
        fsq = MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm.(1.0, [x, y]),
                                          MOI.ScalarQuadraticTerm.(1.0, [x, w, w], [z, z, y]), -3.0)
        @test MOI.output_dimension(fsq) == 1
        @test MOIU.evalvariables(vi -> vals[vi], fsq) ≈ 16
        @test MOIU.evalvariables(vi -> vals[vi], fsq) ≈ 16
        fvq = MOI.VectorQuadraticFunction(MOI.VectorAffineTerm.([2, 1], MOI.ScalarAffineTerm.(1.0, [x, y])),
                                          MOI.VectorQuadraticTerm.([1, 2, 2], MOI.ScalarQuadraticTerm.(1.0, [x, w, w], [z, z, y])), [-3.0, -2.0])
        @test MOI.output_dimension(fvq) == 2
        @test MOIU.evalvariables(vi -> vals[vi], fvq) ≈ [13, 1]
        @test MOIU.evalvariables(vi -> vals[vi], fvq) ≈ [13, 1]
    end
    @testset "mapvariables" begin
        fsq = MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm.(1.0, [x, y]),
                                          MOI.ScalarQuadraticTerm.(1.0, [x, w, w], [z, z, y]), -3.0)
        gsq = MOIU.mapvariables(Dict(x => y, y => z, w => w, z => x), fsq)
        sats = MOI.ScalarAffineTerm.(1.0, [y, z])
        sqts = MOI.ScalarQuadraticTerm.(1.0, [y, w, w], [x, x, z])
        @test gsq.affine_terms == sats
        @test gsq.quadratic_terms == sqts
        @test gsq.constant == -3.
        fvq = MOI.VectorQuadraticFunction(MOI.VectorAffineTerm.([2, 1], MOI.ScalarAffineTerm.(1.0, [x, y])),
                                          MOI.VectorQuadraticTerm.([1, 2, 2], MOI.ScalarQuadraticTerm.(1.0, [x, w, w], [z, z, y])), [-3.0, -2.0])
        gvq = MOIU.mapvariables(Dict(x => y, y => z, w => w, z => x), fvq)
        @test gvq.affine_terms == MOI.VectorAffineTerm.([2, 1], sats)
        @test gvq.quadratic_terms == MOI.VectorQuadraticTerm.([1, 2, 2], sqts)
        @test MOIU.constant(gvq) == [-3., -2.]
    end
    @testset "Conversion VectorOfVariables -> VectorAffineFunction" begin
        f = MOI.VectorAffineFunction{Int}(MOI.VectorOfVariables([z, x, y]))
        @test f isa MOI.VectorAffineFunction{Int}
        @test f.terms == MOI.VectorAffineTerm.(1:3, MOI.ScalarAffineTerm.(ones(Int, 3), [z, x, y]))
        @test all(iszero.(MOIU.constant(f)))
        f = MOI.VectorAffineFunction{Float64}(MOI.VectorOfVariables([x, w]))
        @test f isa MOI.VectorAffineFunction{Float64}
        @test f.terms == MOI.VectorAffineTerm.(1:2, MOI.ScalarAffineTerm.(1.0, [x, w]))
        @test all(iszero.(MOIU.constant(f)))
    end
    @testset "Iteration and indexing on VectorOfVariables" begin
        f = MOI.VectorOfVariables([z, w, x, y])
        it = MOIU.eachscalar(f)
        @test length(it) == 4
        @test eltype(it) == MOI.SingleVariable
        @test collect(it) == [MOI.SingleVariable(z), MOI.SingleVariable(w), MOI.SingleVariable(x), MOI.SingleVariable(y)]
        @test it[2] == MOI.SingleVariable(w)
        @test it[end] == MOI.SingleVariable(y)
    end
    @testset "Indexing on VectorAffineFunction" begin
        f = MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 1, 3, 2, 2, 1, 3, 1, 2],
                                                           MOI.ScalarAffineTerm.([1, 7, 2, 9, 3, 1, 6, 4, 1],
                                                                                [x, y, z, z, y, z, x, x, y])),
                                     [2, 7, 5])
        it = MOIU.eachscalar(f)
        @test length(it) == 3
        @test eltype(it) == MOI.ScalarAffineFunction{Int}
        g = it[2]
        @test g isa MOI.ScalarAffineFunction
        @test g.terms == MOI.ScalarAffineTerm.([1, 9, 3, 1], [x, z, y, y])
        @test g.constant == 7
        g = it[1]
        @test g isa MOI.ScalarAffineFunction
        @test g.terms == MOI.ScalarAffineTerm.([7, 1, 4], [y, z, x])
        @test g.constant == 2
        g = it[end]
        @test g isa MOI.ScalarAffineFunction
        @test g.terms == MOI.ScalarAffineTerm.([2, 6], [z, x])
        @test g.constant == 5
        h = it[[3, 1]]
        @test h isa MOI.VectorAffineFunction
        @test h.terms == MOI.VectorAffineTerm.([1, 1, 2, 2, 2], MOI.ScalarAffineTerm.([2, 6, 7, 1, 4], [z, x, y, z, x]))
        @test MOIU.constant(h) == [5, 2]
        F = MOIU.operate(vcat, Int, it[[1, 2]], it[3])
        @test F isa MOI.VectorAffineFunction{Int}
        @test F.terms == MOI.VectorAffineTerm.([1, 1, 1, 2, 2, 2, 2, 3, 3], MOI.ScalarAffineTerm.([7, 1, 4, 1, 9, 3, 1, 2, 6], [y, z, x, x, z, y, y, z, x]))
        @test MOIU.constant(F) == MOIU.constant(f)
    end
    @testset "Indexing on VectorQuadraticFunction" begin
        f = MOI.VectorQuadraticFunction(MOI.VectorAffineTerm.([2, 1, 3, 2, 2],
                                                              MOI.ScalarAffineTerm.([1, 7, 2, 9, 3],
                                                                                [x, y, z, z, y])),
                                        MOI.VectorQuadraticTerm.([2, 3, 1, 2],
                                                                 MOI.ScalarQuadraticTerm.([1, 6, 4, 3],
                                                                                [z, x, x, y],
                                                                                [y, z, z, y])),
                                        [2, 7, 5])
        it = MOIU.eachscalar(f)
        @test length(it) == 3
        @test eltype(it) == MOI.ScalarQuadraticFunction{Int}
        g = it[2]
        @test g isa MOI.ScalarQuadraticFunction
        @test g.affine_terms == MOI.ScalarAffineTerm.([1, 9, 3], [x, z, y])
        @test g.quadratic_terms == MOI.ScalarQuadraticTerm.([1, 3], [z, y], [y, y])
        @test g.constant == 7
        g = it[end]
        @test g isa MOI.ScalarQuadraticFunction
        @test g.affine_terms == MOI.ScalarAffineTerm.([2], [z])
        @test g.quadratic_terms == MOI.ScalarQuadraticTerm.([6], [x], [z])
        @test g.constant == 5
    end
    @testset "Scalar" begin
        @testset "Affine" begin
            @testset "promote_operation" begin
                @test MOIU.promote_operation(+, Float64, MOI.SingleVariable,
                                             MOI.SingleVariable) == MOI.ScalarAffineFunction{Float64}
                @test MOIU.promote_operation(+, Float64,
                                             MOI.ScalarAffineFunction{Float64},
                                             Float64) == MOI.ScalarAffineFunction{Float64}
                @test MOIU.promote_operation(+, Int,
                                             MOI.ScalarAffineFunction{Int},
                                             MOI.ScalarAffineFunction{Int}) == MOI.ScalarAffineFunction{Int}
            end
            @testset "Comparison tolerance" begin
                @test MOIU.operate(+, Float64, MOI.SingleVariable(x),
                                   MOI.SingleVariable(z)) + 1.0 ≈
                      MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1, 1e-7, 1], [x, y, z]), 1.0) atol=1e-6
                f1 = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x), MOI.ScalarAffineTerm(1e-7, y)], 1.0)
                f2 = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 1.0)
                @test f1 ≈ f2 atol=1e-6
                fdiff = f1 - f2
                MOIU.canonicalize!(fdiff)
                @test !MOIU.isapprox_zero(fdiff, 1e-8)
                @test MOIU.isapprox_zero(fdiff, 1e-6)
            end
            @testset "canonical" begin
                f = MOIU.canonical(MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2, 1, 3, -2, -3], [y, x, z, x, z]), 5))
                @test MOI.output_dimension(f) == 1
                @test f.terms == MOI.ScalarAffineTerm.([-1, 2], [x, y])
                @test f.constant == 5
                f = MOIU.canonical(MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1, 3, 1, 2, -3, 2],
                                                                                [w, y, w, x,  x, z]), 2) +
                                   MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-1, -2, -2, 3, 2],
                                                                                [ y,  z,  w, x, y]), 3))
                @test f.terms == MOI.ScalarAffineTerm.([2, 4], [x, y])
                @test f.constant == 5
            end
            f = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2, 4], [x, y]),
                                         5)
            @testset "convert" begin
                @test_throws InexactError convert(MOI.SingleVariable, f)
                quad_f = MOI.ScalarQuadraticFunction(f.terms,
                                                     MOI.ScalarQuadraticTerm{Int}[],
                                                     f.constant)
                @test convert(MOI.ScalarQuadraticFunction{Int}, f) ≈ quad_f
                g = convert(MOI.ScalarAffineFunction{Float64}, MOI.SingleVariable(x))
                @test convert(MOI.SingleVariable, g) == MOI.SingleVariable(x)
            end
            @testset "operate" begin
                f = MOIU.canonical(MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1, 3, 1, 2, -3, 2],
                                                                                [w, y, w, x,  x, z]), 2) +
                                   MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-1, -2, -2, 3, 2],
                                                                                [ y,  z,  w, x, y]), 3))
                @test f ≈ MOI.SingleVariable(x) + MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1, 4],
                                                                                [x, y]), 5)
                @test f ≈ MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1, 4], [x, y]), 5) + MOI.SingleVariable(x)
                @test f ≈ MOI.SingleVariable(x) - MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-1, -4], [x, y]), -5)
                @test f ≈ MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([3, 4], [x, y]), 5) - MOI.SingleVariable(x)
            end
            @testset "modification" begin
                f = MOIU.modifyfunction(f, MOI.ScalarConstantChange(6))
                @test f.constant == 6
                g = deepcopy(f)
                @test g ≈ f
                f = MOIU.modifyfunction(f, MOI.ScalarCoefficientChange(y, 3))
                @test !(g ≈ f)
                @test g.terms == MOI.ScalarAffineTerm.([2, 4], [x, y])
                @test f.terms == MOI.ScalarAffineTerm.([2, 3], [x, y])
                f = MOIU.modifyfunction(f, MOI.ScalarCoefficientChange(x, 0))
                @test f.terms == MOI.ScalarAffineTerm.([3], [y])
            end
        end
        @testset "Quadratic" begin
            @testset "promote_operation" begin
                @test MOIU.promote_operation(+, Int,
                                             MOI.ScalarQuadraticFunction{Int},
                                             MOI.ScalarQuadraticFunction{Int}) == MOI.ScalarQuadraticFunction{Int}
                @test MOIU.promote_operation(+, Int,
                                             MOI.ScalarQuadraticFunction{Int},
                                             MOI.ScalarAffineFunction{Int}) == MOI.ScalarQuadraticFunction{Int}
                @test MOIU.promote_operation(+, Int,
                                             MOI.ScalarAffineFunction{Int},
                                             MOI.ScalarQuadraticFunction{Int}) == MOI.ScalarQuadraticFunction{Int}
                @test MOIU.promote_operation(*, Int,
                                             MOI.SingleVariable,
                                             MOI.SingleVariable) == MOI.ScalarQuadraticFunction{Int}
                @test MOIU.promote_operation(*, Float64,
                                             MOI.SingleVariable,
                                             MOI.ScalarAffineFunction{Float64}) == MOI.ScalarQuadraticFunction{Float64}
                @test MOIU.promote_operation(*, Int,
                                             MOI.ScalarAffineFunction{Int},
                                             MOI.SingleVariable) == MOI.ScalarQuadraticFunction{Int}
                @test MOIU.promote_operation(*, Float64,
                                             MOI.ScalarAffineFunction{Float64},
                                             MOI.ScalarAffineFunction{Float64}) == MOI.ScalarQuadraticFunction{Float64}
                @test MOIU.promote_operation(/, Float64,
                                             MOI.ScalarQuadraticFunction{Float64},
                                             Float64) == MOI.ScalarQuadraticFunction{Float64}
            end
            fx = MOI.SingleVariable(x)
            fy = MOI.SingleVariable(y)
            f = 7 + 3fx + 1fx * fx + 2fy * fy + 3fx * fy
            MOIU.canonicalize!(f)
            @test MOI.output_dimension(f) == 1
            @testset "isapprox_zero" begin
                @test !MOIU.isapprox_zero(f, 1e-8)
                # Test isapprox_zero with zero terms
                @test MOIU.isapprox_zero(0 * f, 1e-8)
                g = 1.0fx * fy - (1 + 1e-6) * fy * fx
                MOIU.canonicalize!(g)
                @test MOIU.isapprox_zero(g, 1e-5)
                @test !MOIU.isapprox_zero(g, 1e-7)
            end
            @testset "convert" begin
                @test_throws InexactError convert(MOI.SingleVariable, f)
                @test_throws InexactError convert(MOI.ScalarAffineFunction{Int},
                                                  f)
                g = convert(MOI.ScalarQuadraticFunction{Float64}, fx)
                @test convert(MOI.SingleVariable, g) == fx
            end
            @testset "operate" begin
                @test f ≈ 7 + (fx + 2fy) * (1fx + fy) + 3fx
                @test f ≈ 7 + MOIU.operate(*, Int, fx, fx) + 3fx * (fy + 1) + 2fy * fy
                @test f ≈ (fx + 2) * (fx + 1) + (fy + 1) * (2fy + 3fx) + (5 - 3fx - 2fy)
                @test f ≈ begin
                    MOI.ScalarQuadraticFunction([MOI.ScalarAffineTerm(3, x)],
                                                MOI.ScalarQuadraticTerm.([1], [x], [x]), 4) +
                    MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm{Int}[],
                                            MOI.ScalarQuadraticTerm.([2, 3], [y, x], [y, y]), 3)
                end
                @test f ≈ begin
                    MOI.ScalarQuadraticFunction([MOI.ScalarAffineTerm(3, x)],
                                                MOI.ScalarQuadraticTerm.([1], [x], [x]), 10) -
                    MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm{Int}[],
                                                MOI.ScalarQuadraticTerm.([-2, -3], [y, x], [y, y]), 3)
                end
                @test f ≈ begin
                    MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(3, x)], 5) +
                    MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm{Int}[],
                                                MOI.ScalarQuadraticTerm.([1, 2, 3], [x, y, x], [x, y, y]), 2)
                end
                @test f ≈ begin
                    MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(3, x)], 5) -
                    MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm{Int}[],
                                                MOI.ScalarQuadraticTerm.([-1, -2, -3], [x, y, x], [x, y, y]), -2)
                end
                @test f ≈ begin
                    MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm{Int}[],
                                                MOI.ScalarQuadraticTerm.([1, 2, 3], [x, y, x], [x, y, y]), 2) +
                    MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(3, x)], 5)
                end
                @test f ≈ begin
                    MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm{Int}[],
                                                MOI.ScalarQuadraticTerm.([1, 2, 3], [x, y, x], [x, y, y]), 12) -
                    MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(-3, x)], 5)
                end
                @test f ≈ begin
                    MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm.([2], [x]),
                                                MOI.ScalarQuadraticTerm.([1, 2, 3], [x, y, x], [x, y, y]), 7) +
                    MOI.SingleVariable(x)
                end
                @test f ≈ MOI.ScalarQuadraticFunction(MOI.ScalarAffineTerm.([3], [x]),
                                                      MOI.ScalarQuadraticTerm.([1, 2, 3], [x, y, x], [x, y, y]), 10) - 3
            end
            @testset "modification" begin
                f = MOIU.modifyfunction(f, MOI.ScalarConstantChange(9))
                @test f.constant == 9
                f = MOIU.modifyfunction(f, MOI.ScalarCoefficientChange(y, 0))
                @test f.affine_terms == MOI.ScalarAffineTerm.([3], [x])
                g = deepcopy(f)
                @test f ≈ g
                f = MOIU.modifyfunction(f, MOI.ScalarCoefficientChange(y, 2))
                @test !(f ≈ g)
                @test g.affine_terms == MOI.ScalarAffineTerm.([3], [x])
                @test f.affine_terms == MOI.ScalarAffineTerm.([3, 2], [x, y])
            end
        end
    end
    @testset "Vector" begin
        @testset "Affine" begin
            f = MOIU.canonical(MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 1, 2,  1,  1,  2, 2,  2, 2, 1, 1,  2, 1,  2],
                                                                              MOI.ScalarAffineTerm.([3, 2, 3, -3, -1, -2, 3, -2, 1, 3, 5, -2, 0, -1],
                                                                                [x, x, z,  y,  y,  x, y,  z, x, y, y,  x, x,  z])), [5, 7]))
            @test MOI.output_dimension(f) == 2
            @test f.terms == MOI.VectorAffineTerm.([1, 1, 2], MOI.ScalarAffineTerm.([2, 4, 3], [x, y, y]))
            @test MOIU.constant(f) == [5, 7]
            f = MOIU.modifyfunction(f, MOI.VectorConstantChange([6, 8]))
            @test MOIU.constant(f) == [6, 8]
            g = deepcopy(f)
            @test f ≈ g
            f = MOIU.modifyfunction(f, MOI.MultirowChange(y, [(2, 9)]))
            @test !(f ≈ g)
            @test f.terms == MOI.VectorAffineTerm.([1, 1, 2], MOI.ScalarAffineTerm.([2, 4, 9], [x, y, y]))
            @test g.terms == MOI.VectorAffineTerm.([1, 1, 2], MOI.ScalarAffineTerm.([2, 4, 3], [x, y, y]))
            f = MOIU.modifyfunction(f, MOI.MultirowChange(y, [(1, 0)]))
            @test f.terms == MOI.VectorAffineTerm.([1, 2], MOI.ScalarAffineTerm.([2, 9], [x, y]))
        end
        @testset "Quadratic" begin
            f = MOI.VectorQuadraticFunction(MOI.VectorAffineTerm.([1, 2, 2], MOI.ScalarAffineTerm.([3, 1, 2], [x, x, y])), MOI.VectorQuadraticTerm.([1, 1, 2], MOI.ScalarQuadraticTerm.([1, 2, 3], [x, y, x], [x, y, y])), [7, 3, 4])
            @test MOI.output_dimension(f) == 3
            f = MOIU.modifyfunction(f, MOI.VectorConstantChange([10, 11, 12]))
            @test MOIU.constant(f) == [10, 11, 12]
            f = MOIU.modifyfunction(f, MOI.MultirowChange(y, [(2, 0), (1, 1)]))
            @test f.affine_terms == MOI.VectorAffineTerm.([1, 2, 1], MOI.ScalarAffineTerm.([3, 1, 1], [x, x, y]))
            g = deepcopy(f)
            f = MOIU.modifyfunction(f, MOI.MultirowChange(x, [(1, 0), (3, 4)]))
            @test f.affine_terms == MOI.VectorAffineTerm.([2, 1, 3], MOI.ScalarAffineTerm.([1, 1, 4], [x, y, x]))
            @test g.affine_terms == MOI.VectorAffineTerm.([1, 2, 1], MOI.ScalarAffineTerm.([3, 1, 1], [x, x, y]))
        end
    end
    @testset "Conversion to canonical form" begin
        function isapprox_ordered(f1::T, f2::T) where {T <: Union{MOI.ScalarAffineFunction, MOI.VectorAffineFunction}}
            ((MOI.term_indices.(f1.terms) == MOI.term_indices.(f2.terms)) &&
             (MOI._constant(f1) ≈ MOI._constant(f2)) &&
             (MOI.coefficient.(f1.terms) ≈ MOI.coefficient.(f2.terms)))
        end
        function test_canonicalization(f::T, expected::T) where {T <: Union{MOI.ScalarAffineFunction, MOI.VectorAffineFunction}}
            @test MOIU.iscanonical(expected)
            g = @inferred(MOIU.canonical(f))
            @test isapprox_ordered(g, expected)
            @test MOIU.iscanonical(g)
            @test MOIU.iscanonical(expected)
            @test isapprox_ordered(MOIU.canonical(g), g)
            @test MOIU.canonical(g) !== g
            @test @allocated(MOIU.canonicalize!(f)) == 0
        end
        @testset "ScalarAffine" begin
            @test MOIU.iscanonical(MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex.(Int[])), 1.0))
            @test !MOIU.iscanonical(MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([0.0], MOI.VariableIndex.([1])), 1.0))
            @test !MOIU.iscanonical(MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 2.0], MOI.VariableIndex.([2, 1])), 2.0))
            @test !MOIU.iscanonical(MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 0.0], MOI.VariableIndex.([1, 2])), 2.0))

            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex[]), 1.5),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex[]), 1.5),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([0.0], [MOI.VariableIndex(1)]), -2.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex[]), -2.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([0.0], [MOI.VariableIndex(2)]), -2.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex[]), -2.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([3.0], [MOI.VariableIndex(2)]),  1.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([3.0], [MOI.VariableIndex(2)]),  1.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([3.0, 4.0], MOI.VariableIndex.([2, 1])), 0.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([4.0, 3.0], MOI.VariableIndex.([1, 2])), 0.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 0.1], MOI.VariableIndex.([1, 1])), 5.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.1], MOI.VariableIndex.([1])), 5.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 0.5, 2.0], MOI.VariableIndex.([1, 3, 1])), 5.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([3.0, 0.5], MOI.VariableIndex.([1, 3])), 5.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 0.5, 2.0], MOI.VariableIndex.([1, 3, 1])), 5.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([3.0, 0.5], MOI.VariableIndex.([1, 3])), 5.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 2.0, 3.0, 4.0, 5.0], MOI.VariableIndex.([1, 3, 1, 3, 2])), 5.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([4.0, 5.0, 6.0], MOI.VariableIndex.([1, 2, 3])), 5.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, -1.0], MOI.VariableIndex.([7, 7])), 5.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex.(Int[])), 5.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, -1.0, 0.5], MOI.VariableIndex.([7, 7, 2])), 5.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([0.5], MOI.VariableIndex.([2])), 5.0),
                )
            test_canonicalization(
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, -0.5, 0.5, -0.5], MOI.VariableIndex.([7, 7, 2, 7])), 5.0),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([0.5], MOI.VariableIndex.([2])), 5.0),
                )
        end
        @testset "VectorAffine" begin
            @test MOIU.iscanonical(MOI.VectorAffineFunction(MOI.VectorAffineTerm.(Int[], MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex.(Int[]))), Float64[]))
            @test !MOIU.iscanonical(MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1], MOI.ScalarAffineTerm.([0.0], MOI.VariableIndex.([1]))), [1.0]))
            @test !MOIU.iscanonical(MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1, 2], MOI.ScalarAffineTerm.([0.0, 1.0], MOI.VariableIndex.([1, 2]))), [1.0]))
            @test !MOIU.iscanonical(MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 1], MOI.ScalarAffineTerm.([1.0, 1.0], MOI.VariableIndex.([1, 2]))), [1.0]))
            @test !MOIU.iscanonical(MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1, 1], MOI.ScalarAffineTerm.([1.0, -1.0], MOI.VariableIndex.([1, 1]))), [1.0]))
            @test !MOIU.iscanonical(MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1, 1], MOI.ScalarAffineTerm.([1.0, 1.0], MOI.VariableIndex.([1, 1]))), [1.0]))

            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.(Int[], MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex.(Int[]))), Float64[]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.(Int[], MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex.(Int[]))), Float64[])
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([3], MOI.ScalarAffineTerm.([0.0], MOI.VariableIndex.([5]))), [1.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.(Int[], MOI.ScalarAffineTerm.(Float64[], MOI.VariableIndex.(Int[]))), [1.0])
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([3, 2], MOI.ScalarAffineTerm.([1.0, 2.0], MOI.VariableIndex.([5, 6]))), [1.0, 2.0, 3.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 3], MOI.ScalarAffineTerm.([2.0, 1.0], MOI.VariableIndex.([6, 5]))), [1.0, 2.0, 3.0])
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([3, 3], MOI.ScalarAffineTerm.([1.0, 2.0], MOI.VariableIndex.([6, 5]))), [1.0, 2.0, 3.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([3, 3], MOI.ScalarAffineTerm.([2.0, 1.0], MOI.VariableIndex.([5, 6]))), [1.0, 2.0, 3.0])
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([3, 3], MOI.ScalarAffineTerm.([1.0, 2.0], MOI.VariableIndex.([1, 1]))), [1.0, 2.0, 3.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([3], MOI.ScalarAffineTerm.([3.0], MOI.VariableIndex.([1]))), [1.0, 2.0, 3.0])
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1, 1, 1], MOI.ScalarAffineTerm.([1.0, 2.0, 3.0], MOI.VariableIndex.([1, 2, 1]))), [4.0, 5.0, 6.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1, 1], MOI.ScalarAffineTerm.([4.0, 2.0], MOI.VariableIndex.([1, 2]))), [4.0, 5.0, 6.0]),
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1, 1, 1], MOI.ScalarAffineTerm.([1.0, 2.0, 3.0], MOI.VariableIndex.([2, 1, 1]))), [4.0, 5.0, 6.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([1, 1], MOI.ScalarAffineTerm.([5.0, 1.0], MOI.VariableIndex.([1, 2]))), [4.0, 5.0, 6.0]),
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 3, 3], MOI.ScalarAffineTerm.([1.0, 2.0, 3.0], MOI.VariableIndex.([1, 1, 1]))), [4.0, 5.0, 6.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 3], MOI.ScalarAffineTerm.([1.0, 5.0], MOI.VariableIndex.([1, 1]))), [4.0, 5.0, 6.0]),
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 3, 3], MOI.ScalarAffineTerm.([1.0, -3.0, 3.0], MOI.VariableIndex.([1, 1, 1]))), [4.0, 5.0, 6.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2], MOI.ScalarAffineTerm.([1.0], MOI.VariableIndex.([1]))), [4.0, 5.0, 6.0]),
                )
            test_canonicalization(
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2, 3, 3, 3], MOI.ScalarAffineTerm.([1.0, 3.0, -1.0, -2.0], MOI.VariableIndex.([1, 1, 1, 1]))), [4.0, 5.0, 6.0]),
                MOI.VectorAffineFunction(MOI.VectorAffineTerm.([2], MOI.ScalarAffineTerm.([1.0], MOI.VariableIndex.([1]))), [4.0, 5.0, 6.0]),
                )
        end
    end
end


text:
rabit_finalize() = ccall((:RabitFinalize, _jl_libxgboost), Nothing, ())
function rabit_init()
    argv = Ref{Ptr{Nothing}}()
    ccall((:RabitInit, _jl_libxgboost), Nothing, (Cint, Ref{Ptr{Nothing}}), 0, argv)
end

function rabit_is_distributed()
    ret = ccall((:RabitIsDistributed, _jl_libxgboost), Cint, ())
    return (ret != Cint(0))
end

rabit_get_rank() = ccall((:RabitGetRank, _jl_libxgboost), Cint, ())

rabit_get_world_size() = ccall((:RabitGetWorldSize, _jl_libxgboost), Cint, ())

rabit_get_version_number() = ccall((:RabitVersionNumber, _jl_libxgboost), Cint, ())


text:
import Dates
import Pkg
import Tar

function get_bool_from_env(name::AbstractString, default_value::Bool)
    value = get(ENV, name, "$(default_value)") |> strip |> lowercase
    result = parse(Bool, value)::Bool
    return result
end

const is_buildkite         = get_bool_from_env("BUILDKITE",                  false)
const always_save_rr_trace = get_bool_from_env("JULIA_ALWAYS_SAVE_RR_TRACE", false)

function get_from_env(name::AbstractString)
    if is_buildkite
        value = ENV[name]
    else
        value = get(ENV, name, "")
    end
    result = convert(String, strip(value))::String
    return result
end

function my_exit(process::Base.Process)
    wait(process)

    @info(
        "",
        process.exitcode,
        process.termsignal,
    )

    # Pass the exit code back up
    if process.termsignal != 0
        ccall(:raise, Cvoid, (Cint,), process.termsignal)

        # If for some reason the signal did not cause an exit, we'll exit manually.
        # We need to make sure that we exit with a non-zero exit code.
        if process.exitcode != 0
            exit(process.exitcode)
        else
            exit(1)
        end
    end
    exit(process.exitcode)
end

if Base.VERSION < v"1.6"
    throw(ErrorException("The `$(basename(@__FILE__))` script requires Julia 1.6 or greater"))
end

if length(ARGS) < 1
    throw(ErrorException("Usage: julia $(basename(@__FILE__)) [command...]"))
end

@info "We will run the command under rr"

const build_number                      = get_from_env("BUILDKITE_BUILD_NUMBER")
const job_name                          = get_from_env("BUILDKITE_STEP_KEY")
const commit_full                       = get_from_env("BUILDKITE_COMMIT")
const commit_short                      = first(commit_full, 10)
const JULIA_TEST_RR_TIMEOUT             = get(ENV,  "JULIA_TEST_RR_TIMEOUT", "120")
const timeout_minutes                   = parse(Int, JULIA_TEST_RR_TIMEOUT)
const JULIA_TEST_NUM_CORES              = get(ENV,  "JULIA_TEST_NUM_CORES", "8")
const julia_test_num_cores_int          = parse(Int, JULIA_TEST_NUM_CORES)
const num_cores = min(
    8,
    Sys.CPU_THREADS,
    julia_test_num_cores_int + 1,
)

ENV["JULIA_RRCAPTURE_NUM_CORES"] = "$(num_cores)"

@info(
    "",
    build_number,
    job_name,
    commit_full,
    commit_short,
    timeout_minutes,
    num_cores,
)

const dumps_dir       = joinpath(pwd(), "dumps")
const temp_parent_dir = joinpath(pwd(), "temp_for_rr")

mkpath(dumps_dir)
mkpath(temp_parent_dir)

proc = nothing

mktempdir(temp_parent_dir) do dir
    Pkg.activate(dir)
    Pkg.add("rr_jll")
    Pkg.add("Zstd_jll")

    rr_jll = Base.require(Base.PkgId(Base.UUID((0xe86bdf43_55f7_5ea2_9fd0_e7daa2c0f2b4)), "rr_jll"))
    zstd_jll = Base.require(Base.PkgId(Base.UUID((0x3161d3a3_bdf6_5164_811a_617609db77b4)), "Zstd_jll"))
    rr(func) = Base.invokelatest(rr_jll.rr, func; adjust_LIBPATH=false)
    rr() do rr_path
        capture_script_path = joinpath(dir, "capture_output.sh")
        loader = Sys.WORD_SIZE == 64 ? "/lib64/ld-linux-x86-64.so.2" : "/lib/ld-linux.so.2"
        open(capture_script_path, "w") do io
            write(io, """
            #!/bin/bash

            $(rr_path) record --nested=detach "\$@" > >(tee -a $(dir)/stdout.log) 2> >(tee -a $(dir)/stderr.log >&2)
            """)
        end
        chmod(capture_script_path, 0o755)

        new_env = copy(ENV)
        new_env["_RR_TRACE_DIR"] = joinpath(dir, "rr_traces")
        new_env["RR_LOG"]="all:debug"
        new_env["RR_LOG_BUFFER"]="100000"
        new_env["JULIA_RR"] = capture_script_path
        t_start = time()
        global proc = run(setenv(`$(rr_path) record --num-cores=$(num_cores) $ARGS`, new_env), (stdin, stdout, stderr); wait=false)

        # Start asynchronous timer that will kill `rr`
        @async begin
            sleep(timeout_minutes * 60)

            # If we've exceeded the timeout and `rr` is still running, kill it.
            if isopen(proc)
                println(stderr, "\n\nProcess timed out (with a timeout of $(timeout_minutes) minutes). Signalling `rr` for force-cleanup!")
                kill(proc, Base.SIGTERM)

                # Give `rr` a chance to cleanup and upload.
                # Note: this time period includes the time to upload the `rr` trace files
                # as Buildkite artifacts, so make sure it is long enough to allow the
                # uploads to finish.
                cleanup_minutes = 30
                sleep(cleanup_minutes * 60)

                if isopen(proc)
                    println(stderr, "\n\n`rr` failed to cleanup and upload within $(cleanup_minutes) minutes, killing and exiting immediately!")
                    kill(proc, Base.SIGKILL)
                    exit(1)
                end
            end
        end

        # Wait for `rr` to finish, either through naturally finishing its run, or `SIGTERM`.
        wait(proc)
        process_failed = !success(proc)

        if process_failed || always_save_rr_trace || is_buildkite
            println(stderr, "`rr` returned $(proc.exitcode), packing and uploading traces...")

            if !isdir(joinpath(dir, "rr_traces"))
                println(stderr, "No `rr_traces` directory!  Did `rr` itself fail?")
                exit(1)
            end

            # Clean up non-traces
            rm(joinpath(dir, "rr_traces", "latest-trace"))
            rm(joinpath(dir, "rr_traces", "cpu_lock"))

            # Create a directory for the pack files to go
            pack_dir = joinpath(dir, "pack")
            mkdir(pack_dir)

            # Pack all traces
            trace_dirs = [joinpath(dir, "rr_traces", f) for f in readdir(joinpath(dir, "rr_traces"))]
            filter!(isdir, trace_dirs)
            run(ignorestatus(`$(rr_path) pack --pack-dir=$pack_dir $(trace_dirs)`))

            # Tar it up
            mkpath(dumps_dir)
            date_str = Dates.format(Dates.now(), Dates.dateformat"yyyy_mm_dd_HH_MM_SS")
            dst_file_name = string(
                "rr",
                "--build_$(build_number)",
                "--$(job_name)",
                "--commit_$(commit_short)",
                "--$(date_str)",
                ".tar.zst",
            )
            dst_full_path = joinpath(dumps_dir, dst_file_name)
            zstd_jll.zstdmt() do zstdp
                tarproc = open(`$(zstdp) -o $(dst_full_path)`, "w")
                Tar.create(dir, tarproc)
                close(tarproc.in)
            end

            @info "The `rr` trace file has been saved to: $(dst_full_path)"
            if is_buildkite
                @info "Since this is a Buildkite run, we will upload the `rr` trace file."
                cd(dumps_dir) do
                    run(`buildkite-agent artifact upload $(dst_file_name)`)
                end
            end
        end

    end
end

@info "Finished running the command under rr"
my_exit(proc)

max_stars_repo_name: wentasah/julia

max_issues_repo_name: wentasah/julia

max_forks_repo_name: wentasah/julia


text:
using Cliffords, Base.Test
using Compat

import Cliffords: IZ, ZI, XI, IX, YI, IY, II, XX, YY, ZZ, XY, XZ, YX, ZX,
	generators

# Paulis
@test X * Id == X
@test Id * X == X
@test Y * Id == Y
@test Z * Id == Z

@test X * X == Id
@test Y * Y == Id
@test Z * Z == Id

@test X * Y == im*Z
@test Y * Z == im*X
@test Z * X == im*Y
@test Y * X == -im*Z

@test Id < X
@test X < Y
@test Y < Z

@test II < IX
@test IX < XX
@test XI < XX
@test IX < IY
@test XI < YI
@test XX < YY
@test XX < XY
@test YX > XX

# generators
@test generators(X) == [X]
@test generators(Z) == [Z]
@test generators(Y) == [im*X, Z]

@test generators(XX) == [XI, IX]
@test generators(ZZ) == [ZI, IZ]
@test generators(XY) == [im*XI, IX, IZ]
@test generators(YX) == [im*XI, ZI, IX]
@test generators(YY) == [-XI, ZI, IX, IZ]

XII = kron(X,Id,Id)
IXI = kron(Id,X,Id)
IIX = kron(Id,Id,X)
ZII = kron(Z,Id,Id)
IZI = kron(Id,Z,Id)
IIZ = kron(Id,Id,Z)

@test generators(kron(X,X,X)) == [XII, IXI, IIX]
@test generators(kron(Z,Z,Z)) == [ZII, IZI, IIZ]
@test generators(kron(X,Y,Z)) == [im*XII, IXI, IZI, IIZ]

@test Pauli([1 0; 0 1]) == Id
@test Pauli([0 1; 1 0]) == X
@test Pauli([0 -im; im 0]) == Y
@test Pauli([1 0; 0 -1]) == Z

@test complex(paulieye(2)) == eye(4)

# Cliffords

@test Clifford(1/sqrt(2) * [1 1; 1 -1]) == H
@test Clifford(eye(2)) == RI
@test Clifford(eye(4)) == kron(RI,RI)
@test Clifford([1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0]) == CNOT
@test Clifford(diagm([1,1,1,-1])) == CZ
@test Clifford(expm(-im*pi/4*[1 0; 0 -1])) == S

function eq_upto_phase(A, B)
	idx = findfirst(A)
	rel_phase = (B[idx] == 0) ? 1.0 : (A[idx] / B[idx])
	return rel_phase * A ≈ B
end
@test eq_upto_phase(complex(RI), eye(2))
@test eq_upto_phase(complex(H), 1/sqrt(2) * [1 1; 1 -1])
@test eq_upto_phase(complex(CNOT), [1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0])

# Cliffords * Paulis
@test H * Id == Id
@test H * X == Z
@test H * Z == X
@test H * Y == -Y

II = kron(Id,Id)

@test kron(RI, H) * II == II
@test kron(RI, H) * IX == IZ
@test kron(RI, H) * IZ == IX
@test kron(RI, H) * IY == -IY

@test CNOT * II == II
@test CNOT * IX == IX
@test CNOT * XI == XX
@test CNOT * YY == -XZ

# Cliffords \ Paulis
@test H \ Id == Id
@test H \ X == Z
@test H \ Z == X
@test H \ Y == -Y

II = kron(Id,Id)

@test kron(RI, H) \ II == II
@test kron(RI, H) \ IX == IZ
@test kron(RI, H) \ IZ == IX
@test kron(RI, H) \ IY == -IY

@test CNOT \ II == II
@test CNOT \ IX == IX
@test CNOT \ XI == XX
@test CNOT \ YY == -XZ

# Cliffords * Cliffords

@test H * S * H == Clifford(@compat(Dict(+X=>+X,+Z=>-Y)),@compat Dict(+X=>+X,+Z=>+Y))
@test H * S * H == Clifford(expm(-im*pi/4*[0 1; 1 0]))
@test RY * RX * RY == SelfInverseClifford(@compat Dict(+X=>+X,+Z=>-Z))
@test kron(RI, H) * CNOT * kron(RI, H) == SelfInverseClifford(@compat Dict(+IZ=>+IZ,+ZI=>+ZI,+XI=>+XZ,+IX=>+ZX))
@test kron(RI, H) * CNOT * kron(RI, H) == CZ
@test S * S == RZ

# 3 CNOTs is a SWAP
CNOT21 = expand(CNOT, [2,1], 2)
@test CNOT * CNOT21 * CNOT == SWAP

# iSWAP
iSWAP = CZ*kron(S,S)*SWAP
@test iSWAP == Clifford([1 0 0 0; 0 0 im 0; 0 im 0 0; 0 0 0 1])

# Cliffords \ Cliffords
@test H \ H == H * H
@test CNOT \ CNOT == CNOT * CNOT
@test RY * RX * RY == RY \ (RX * RY)
@test S \ S == RI
@test S \ RZ == S

# Local Cliffords
@test localclifford(1) == RI

# Weight
for i=1:20
  v = rand(10)
  vX = (v .> .25) .* (v .< .5)
  vY = 2*(v .> .5) .* (v .< .75)
  vZ = 3*(v .> .75)
  r = rand()
  rS = (r > .25) + (r > .5) + (r > .75)
  @test weight(Pauli(vX + vY + vZ,rS)) == sum(v .> .25)
end

#
@test_approx_eq norm(complex(X)-[0 1; 1 0]) 0
@test_approx_eq norm(complex(Y)-[0 -1im;1im 0]) 0
@test_approx_eq norm(complex(Z)-[1 0; 0 -1]) 0


text:
immutable TProblem{T<:Integer,U<:String}
    sd::Array{Array{T,1},1}
    toc::Array{T,2}
    labels::Array{Array{U,1},1}
    tsort::Array{Array{T,2}, 1}
end

function TProblem{T<:Integer,U<:String}(s::Array{T,1},
                                        d::Array{T,1},
                                        toc::Array{T,2},
                                        slab::Array{U,1},
                                        dlab::Array{U,1})
    scnt = length(s)
    dcnt = length(d)
    size(toc) = (scnt,dcnt) || error("Supply, Demand, TOC Size Mismatch")
    length(slab) == scnt || error("Supply Label Size Labels")
    length(dlab) == dcnt || error("Demand Label Size Labels")
    0 <= minimum(s) || error("Negative Supply Value")
    0 <= minimum(d) || error("Negative Demand Value")
    sd = Array{T,1}[]
    push!(sd, s)
    push!(sd, d)
    labels = Array{U,1}[]
    push!(labels, slab)
    push!(labels, dlab)
    tsort = Array{T,2}[]
    push!(tsort, mapslices(sortperm, toc, 2))
    push!(tsort, mapslices(sortperm, toc, 1))
    TProblem(sd, toc, labels, tsort)
end
isbalanced(tp::TProblem) = sum(tp.sd[1]) == sum(tp.sd[2])

type Resource{T<:Integer}
    dim::T
    i::T
    quant::T
    l::T
    m::T
    p::T
    q::T
end
function Resource{T<:Integer}(dim::T, i::T, quant::T)
    zed = zero(T)
    Resource(dim, i, quant, zed, zed, zed, zed)
end

isavailable(r::Resource) = 0 < r.quant
Base.isless(a::Resource, b::Resource) = a.p < b.p || (a.p == b.p && b.q < a.q)

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode