text:
export @ordered_piecewise

function generate_ordered_piecewise(bps, funcs)
    ndiv = length(bps)
    str = "x -> begin\n"
    str *= "if x < $(bps[1])\n"
    str*= "return (" * string(funcs[1]) * ")(x)\n"
    for i in 2:ndiv
        str*= "elseif x < $(bps[i])\n"
        str*= "return (" * string(funcs[i]) * ")(x)\n"
    end
    str*= "else\n"
    str*= "return (" * string(funcs[end]) * ")(x)\n"
    #str*="return"
    str*= "end\nend"
    return str |> Meta.parse
end

function collect_conditions(ex::Expr; var = :x, check = false)
    block = ex.args
    nstatements = length(block) ÷ 2
    conditions = []
    functions = Expr[]
    for (i, line) in enumerate(block)
        if !(line isa LineNumberNode)
            cond, func = split_condition(line)
            check && check_condition(var, cond)
            push!(conditions, cond)
            push!(functions, func)
        end
    end
    return conditions, functions
end

function check_condition(var, condition)
    @assert any(condition.args .|> string .== string(var)) "Conditions must be defined in terms of $(string(var))"
end

function split_condition(ex::Expr)
    @assert ex.args[1] == :(=>) "Piecewise function definition must be performed with :(=>) operator"
    condition = ex.args[2]
    func = ex.args[3]
    return condition, func
end

numbers(condition) = filter(x -> x isa Number, condition.args)

function parse_block(block::Expr)
    breakpoints, functions = collect_conditions(block)
    @assert breakpoints[end] == :_ "Otherwise condition must be denoted by _ (underscore) character"
    @views breakpoints = [b for b in breakpoints[1:end-1]]

    return breakpoints, functions
end

macro ordered_piecewise(block::Expr)
    breakpoints, functions = parse_block(block)
    expr = generate_ordered_piecewise(breakpoints, functions)
    N = length(breakpoints)
    expr_strings = functions .|> prettify_expr
    functions = eval.(functions)
    quote
        OrderedPiecewiseFunction{$N}($(esc(functions)), $(esc(breakpoints)), $(esc(expr)), $(esc(expr_strings)))
    end
end

macro ordered_piecewise(functions::Expr, breakpoints::Expr)
    fs, bps = functions.args, breakpoints.args
    bps = [b for b in bps]
    expr = generate_ordered_piecewise(bps, fs)
    N = length(bps)
    expr_strings = fs .|> prettify_expr

    functions = eval.(fs)

    quote
        OrderedPiecewiseFunction{$N}(($(esc(functions))), $(esc(bps)), $(esc(expr)), $(esc(expr_strings)))
    end
end

macro piecewise_polynomial(block::Expr)
    breakpoints, polynomials = parse_block(block)
    N = length(breakpoints)
    polys = [StaticPolynomial(p.args[3]) for p in polynomials]

    quote
        PiecewisePolynomial{$N}($(esc(polys)), $(esc(breakpoints)))
    end
end

macro piecewise_polynomial(polynomials::Expr, breakpoints::Expr)
    polys, bps = polynomials.args, breakpoints.args
    bps = [b for b in bps]
    N = length(bps)
    polys = [StaticPolynomial(p.args[3]) for p in polys]
    quote
        PiecewisePolynomial{$N}($(esc(polys)), $(esc(bps)))
    end
end

function prettify_expr(expr::Expr)
    args = expr.args[2].args
    cleaned = []
    for a in args
        if !(a isa LineNumberNode)
            push!(cleaned, a)
        end
    end
    #@show isa.(args, LineNumberNode)
    expr.args[2].args = cleaned
    str = split(string(expr), '\n')[1:end-1] .|> strip |> foldr'(*)# |> remove_spaces
    str = replace(str, "->begin" => " -> ")
    str = replace(str, " ^ " => "^")
    return str
end


text:
module Optimization

using MicroLogging
using KrylovMethods

using DeformableRegistration
using DeformableRegistration.Types

# include("../experimental/aspin.jl")

export checkStoppingCriteria, ArmijoLineSearch, optimizeGaussNewton, optimizeGaussNewtonAugmentedLagrangian

# function augmentedLagrangian(x, λ, μ, c::Function)
#     m = length(x)
#     F = - λ'*c(x)[1] + μ/2 * c(x)[1]'*c(x)[1]
#     dF = - c(x)[2]*λ   + μ * c(x)[2]'*c(x)[1]
#     d2F = spzeros(m,m)#pdiagm(-λ, 0)  +  μ * c(x)[2]'*c(x)[2] + μ * spdiagm(c(x)[1],0)
#     return [F, dF, d2F]
# end

function augmentedLagrangian(x, λ, μ, c::Function)
    m = length(x)
    cx, ∇c = c(x)
    F = - λ' * cx + μ/2 * cx' * cx
    ∇F = - ∇c * λ   + μ * ∇c' * cx
    ∇2F(y) = μ * ∇c' * ∇c * y
    return [F, ∇F, ∇2F]
end

"""
optimizeGaussNewtonAugmentedLagrangian(Jfunc, y, yReference, options; constraint)

Augmented lagrangian implementation following Nocelda & Wright, 2006, p. 515 (Framework 17.3)
"""
function optimizeGaussNewtonAugmentedLagrangian(Jfunc::Function,  # objective Function
                             y::Array{Float64,1}, yReference::Array{Float64,1}, options::regOptions;
                             constraint::Function = x -> [0, 0, 0], printFunction=x->[], gradientDescentOnly=false, displayCallback=x->x) #no constraint by default,

    c = constraint
    computeConstraint = ! (c == x -> [0, 0, 0])

    L(x, λ, μ; doDerivative=true, doHessian=true) = (c(y) != [0, 0, 0]) ? Jfunc(x, doDerivative=doDerivative, doHessian=doHessian) +
     augmentedLagrangian(x, λ, μ, c) : Jfunc(x, doDerivative=doDerivative, doHessian=doHessian)

    λ = spzeros(size(y)...)
    μ = 1

    JRef = L(yReference, λ, μ)[1]
    JOld = NaN
    oldNormCy = 99 * norm(c(y)[1], Inf)

    for iter = 1:options.maxIterGaussNewton
     if computeConstraint
         λ = λ - μ * c(y)[1]

         "Enforce reduction of infeasability"
         τReduction = 0.8
         if norm(c(y)[1], Inf) >  τReduction * oldNormCy;
             @debug @sprintf("... norm(c(y)[1], Inf) >  %1.1f * oldNormCy: %3.3e > %3.3e", τReduction, norm(c(y)[1], Inf), oldNormCy)
             @debug "... updating μ = 2 * μ = ", 2*μ
             μ = 2 * μ;
         end
         oldNormCy = norm(c(y)[1], Inf)

     end

     J, dJ, d2J = L(y, λ, μ, doDerivative=true, doHessian=true)

     @debug d2J(-dJ)
     if gradientDescentOnly
         dy = -dJ
         cgIterations = -1
     else
         dy,flag,resvec,cgIterations = KrylovMethods.cg(d2J,-dJ,maxIter=options.maxIterCG, tol=options.tolCG)[1:4] #2/3 of the time of this function is here
     end

        # check descent direction
        if( (dJ'*dy)[1] > 0)
         dy = -dy
         @warn "Changing sign of computed descent direction. This is a suspicious move."
        end

     # armijo line search (LS) method
     stepLength,LSiter,LSfailed = ArmijoLineSearch(x -> L(x, λ, μ), J, dJ, y, dy, tolLS=1e-4)

     if(LSfailed)
         @info "STOPPING after line search failed."
         break
     end

     # output
     # is this a IJulia notebook? -> delete previous iteration's output
     if isdefined(Main, :IJulia) && Main.IJulia.inited
        displayCallback(true)
        @info size(y)
     end
     D, α, S = printFunction(y)
     s = @sprintf("%03d | D %2.2e | S %2.2e | λ!=0: %d | |c(y)|= %2.2e | μ=%2.2e | LSiter: %2d | CGiter: %3d | J/Jref: %2.2e",
         iter, D[1], S[1], sum(λ.!=0), norm(c(y)[1], Inf), μ,  LSiter, cgIterations, J[1]/JRef[1],
         )
     @info s
     @debug "  λ=", λ
     @debug "  c=", c(y)[1]

     # update parameter y
     y = y + stepLength .* dy

     # stopping criteria
     if (iter>1)
         if (checkStoppingCriteria(J, JOld, JRef, dJ, y, stepLength*dy,
                 tolJ = options.stopping["tolJ"],    # tolerance: change of the objective function
                 tolY = options.stopping["tolY"],    # tolerance: change of the variables
                 tolG = options.stopping["tolG"],    # tolerance: change of the gradient
                 tolQ = options.stopping["tolQ"]))   # tolerance: change of quotient)
             break
         end
     end

     # updating JOld
     JOld = J[1]
     @info "\n"
    end

    return y

end

function optimizeGaussNewton(Jfunc::Function,  # objective Function
                             y::Array{Float64,1}, yInitial::Array{Float64,1}, options::regOptions; displayCallback=x->x)

    # Initilization of JRef: reference value of objective Function
    JRef = Jfunc(yInitial)[1]
    JOld = NaN

    for iter = 1:options.maxIterGaussNewton

        # get derivatives of objective function
        J, dJ, d2J = Jfunc(y,doDerivative=true,doHessian=true)

        # solve system of linear equations d2J*dy=-dJ
          # (dy: change of the variables / search direction)
        # number of parameters < 10 => use backslash operator
        # else (large systems) => use conjugate gradient (cg) method

        cgIterations=0

        if(length(y)<10 && !(typeof(d2J) <: Function))
            dy=d2J\-dJ
        else
            dy,flag,resvec,cgIterations = KrylovMethods.cg(d2J,-dJ,maxIter=options.maxIterCG, tol=options.tolCG)[1:4]
        end

	    # check descent direction
	    if( (dJ'*dy)[1] > 0)
            dy = -dy
            @warn "Changing sign of computed descent direction. This is a suspicious move."
	    end

        # armijo line search (LS) method
        stepLength,LSiter,LSfailed = ArmijoLineSearch(Jfunc,J,dJ,y,dy)
        if(LSfailed)
            break
        end

        # is this a IJulia notebook? -> hand over to user provided method to delete previous iteration's output
        if isdefined(Main, :IJulia) && Main.IJulia.inited
           displayCallback(true)
        end

        # output
        if(cgIterations==0)
          s = @sprintf("%3d: J %8.4e     LSiter: %2d    J/Jref: %1.2f \n",iter, J[1], LSiter, J[1]/JRef[1])
          @info s
        else
          s = @sprintf("%3d: J %8.4e     LSiter: %2d     CGiter: %3d     J/Jref: %1.2f \n",iter, J[1], LSiter,cgIterations, J[1]/JRef[1])
          @info s
        end

        # update parameter y
        y = y + stepLength.*dy

        # stopping criteria
        if (iter>1)
            if checkStoppingCriteria(J, JOld, JRef, dJ, y, stepLength*dy,
                    tolJ = options.stopping["tolJ"],    # tolerance: change of the objective function
                    tolY = options.stopping["tolY"],    # tolerance: change of the variables
                    tolG = options.stopping["tolG"],    # tolerance: change of the gradient
                    tolQ = options.stopping["tolQ"])    # tolerance: change of quotient)
                break
            end
        end

        # updating JOld
        JOld = J[1]

    end

    return y

end

function checkArmijoStepValid(Jfunc::Function, Jc::Float64, y::Array{Float64,1},
        stepLength::Float64, dy::Array{Float64,1}, tolLS::Float64,
        dJ::Array{Float64,1})
    @debug "Jfunc(y+stepLength*dy)[1] < (Jc + (stepLength.*tolLS.*(dJ'*dy))[1])" Jfunc(y+stepLength*dy)[1]  (Jc + (stepLength.*tolLS.*(dJ'*dy))[1])
    validStep = Jfunc(y+stepLength*dy)[1] < (Jc + (stepLength.*tolLS.*(dJ'*dy))[1])
    return validStep
end

function ArmijoLineSearch(Jfunc::Function,         # objective function
                          Jc::Float64,             # Jc = Jfunc(y)
                          dJ::Array{Float64,1},    # gradient of the objective funtion
                          y::Array{Float64,1},     # old variables
                          dy::Array{Float64,1};    # change of the variables / search direction
                          tolLS::Float64 = 1e-4)   # tolerance

    stepLength = 1.0; LSiter = 1; LSfailed = false;

    while( true )
        validStep = checkArmijoStepValid(Jfunc, Jc, y, stepLength, dy, tolLS, dJ)
        if validStep
            break
        end
        if(stepLength < 1e-4)
            LSfailed = true
            break
        end

        stepLength = 0.5 * stepLength
        LSiter = LSiter + 1

    end

    if(LSfailed)
        s = @sprintf("   Line search failed after %2d iterations.",LSiter)
        @info s
    end

    return stepLength,LSiter,LSfailed

end

# somehow slower with this signature
# function checkStoppingCriteria(J::Float64,JOld::Float64,JRef::Float64,    # value of the current objective function J(y+dy), the old J(y) and the reference J(yInital)
#                                dJ::Array{Float64,1},             # gradient of the objective funtion
#                                y::Array{Float64,1},              # old variables
#                                dy::Array{Float64,1};             # change of the variables / search direction
#                                tolJ::Float64 = 1e-3,    # tolerance: change of the objective function
#                                tolY::Float64 = 1e-2,    # tolerance: change of the variables
#                                tolG::Float64 = 1e-2,    # tolerance: change of the gradient
#                                tolQ::Float64 = 1e-4     # tolerance: change of quotient
#                                )

function checkStoppingCriteria(J,JOld,JRef,    # value of the current objective function J(y+dy), the old J(y) and the reference J(yInital)
                               dJ,             # gradient of the objective funtion
                               y,              # old variables
                               dy;             # change of the variables / search direction
                               tolJ = 1e-3,    # tolerance: change of the objective function
                               tolY = 1e-2,    # tolerance: change of the variables
                               tolG = 1e-2,    # tolerance: change of the gradient
                               tolQ = 1e-4     # tolerance: change of quotient
                               )

    # Initialize STOP-Array with false
    STOP = Array{Bool}(5); STOP[:] = false

    # Either STOP[1:3] or STOP[4]
    STOP[1] = abs(JOld-J) <= tolJ  * (1+abs(JOld))
    STOP[2] = norm(dy)    <= tolY  * (1+norm(y))
    STOP[3] = norm(dJ)    <= tolG  * (1+abs(JOld))
    STOP[4] = norm(dJ)    <= 1e6 * eps()
    STOP[5] = abs((JOld-J)/(JRef-J)) <= tolQ

    if(all(STOP[1:3]))
        s = @sprintf("STOPPING CRITERIA:\n")
        @info s
        s = @sprintf("   1. |JOld-J| = %3e <= %3e (tolJ * (1 + |JOld|)) \n",abs(JOld-J), tolJ*(1+abs(JOld)) )
        @info s
        s = @sprintf("   2. ||dy||   = %3e <= %3e (tolY * (1 + ||y||))\n",norm(dy),    tolY*(1+norm(y))  )
        @info s
        s = @sprintf("   3. ||dJ||   = %3e <= %3e (tolG * (1 + |JOld|)) \n",norm(dJ),    tolG*(1+abs(JOld)) )
        @info s
    end

    if(STOP[4])
        s = @sprintf("STOPPING CRITERION:\n")
        @info s
        s = @sprintf("   4. ||dJ|| = %3e <= %3e (1e6*eps)\n",norm(dJ), 1e6*eps())
        @info s
    end

    if(STOP[5])
        s = @sprintf("STOPPING CRITERION:\n")
        @info s
        s = @sprintf("   5. |(JOld-J)/(JRef-J)|= %3e <= %3e (tolQ)\n",abs((JOld-J)/(JRef-J)), tolQ)
        @info s
    end

    return all(STOP[1:3]) | STOP[4] | STOP[5]

end

end

ext: jl


text:
target = "cpu"
using CUDA
using CUDA.CUSPARSE
using ExaPF
using LinearAlgebra
using BenchmarkTools
using SparseArrays
using TimerOutputs
import ExaPF: Parse, PowerSystem

# read data
function run_benchmark(datafile)
    to = TimerOutputs.TimerOutput()
    data = Parse.parse_raw(datafile)
    BUS_B, BUS_AREA, BUS_VM, BUS_VA, BUS_NVHI, BUS_NVLO, BUS_EVHI,
    BUS_EVLO, BUS_TYPE = Parse.idx_bus()
    bus = data["BUS"]
    nbus = size(bus, 1)

    # obtain V0 from raw data
    V = Array{Complex{Float64}}(undef, nbus)
    T = Vector
    for i in 1:nbus
        V[i] = bus[i, BUS_VM]*exp(1im * pi/180 * bus[i, BUS_VA])
    end

    # form Y matrix
    Ybus, Yf_br, Yt_br, Yf_tr, Yt_tr = PowerSystem.makeYbus(data);

    Vm = abs.(V)
    Va = angle.(V)
    bus = data["BUS"]
    gen = data["GENERATOR"]
    load = data["LOAD"]
    nbus = size(bus, 1)
    ngen = size(gen, 1)
    nload = size(load, 1)

    ybus_re, ybus_im = ExaPF.Spmat{T}(Ybus)
    SBASE = data["CASE IDENTIFICATION"][1]
    Sbus = PowerSystem.assembleSbus(gen, load, SBASE, nbus)
    pbus = real(Sbus)
    qbus = imag(Sbus)

    ref, pv, pq = PowerSystem.bustypeindex(bus, gen)
    npv = size(pv, 1);
    npq = size(pq, 1);

    F = zeros(Float64, npv + 2*npq)
    println("[CPU] Benchmark function residualFunction_polar!")
    @btime begin
        ExaPF.residualFunction_polar!($F, $Vm, $Va,
            $ybus_re.nzval, $ybus_re.colptr, $ybus_re.rowval,
            $ybus_im.nzval, $ybus_im.colptr, $ybus_im.rowval,
            $pbus, $qbus, $pv, $pq, $nbus)
    end

    J = ExaPF.residualJacobian(V, Ybus, pv, pq)
    F = zeros(Float64, npv + 2*npq)

    # Then, create a JacobianAD object
    coloring = ExaPF.matrix_colors(J)
    jacobianAD = ExaPF.AD.JacobianAD(J, coloring, F, Vm, Va, pv, pq)
    # and compute Jacobian with ForwardDiff
    println("[CPU] Benchmark residualJacobianAD!")
    @btime begin
        ExaPF.AD.residualJacobianAD!(
            $jacobianAD, ExaPF.residualFunction_polar!, $Vm, $Va,
            $ybus_re, $ybus_im, $pbus, $qbus, $pv, $pq, $nbus, $to)
    end

    # Passing to GPU
    target = "cuda"
    gF = F |> CuArray
    gVm = Vm |> CuArray
    gVa = Va |> CuArray
    gpbus = pbus |> CuArray
    gqbus = qbus |> CuArray
    gpv = pv |> CuArray
    gpq = pq |> CuArray
    gybus_re, gybus_im = ExaPF.Spmat{CuVector}(Ybus)
    fill!(gF, 0.0)

    println("[GPU] Benchmark residualFunction_polar!")
    @btime begin
        $gF .= .0
        ExaPF.residualFunction_polar!($gF, $gVm, $gVa,
                $gybus_re.nzval, $gybus_re.colptr, $gybus_re.rowval,
                $gybus_im.nzval, $gybus_im.colptr, $gybus_im.rowval,
                $gpbus, $gqbus, $gpv, $gpq, $nbus)
    end
    gJ = CuSparseMatrixCSR(J)

    # Then, create a JacobianAD object
    ccolors = ExaPF.matrix_colors(J) |> CuArray

    gjacobianAD = ExaPF.AD.JacobianAD(gJ, ccolors, gF, gVm, gVa, gpv, gpq)
    println("[GPU] Benchmark residualJacobianAD!")
    @btime begin
        ExaPF.AD.residualJacobianAD!(
            $gjacobianAD, ExaPF.residualFunction_polar!, $gVm, $gVa,
            $gybus_re, $gybus_im, $gpbus, $gqbus, $gpv, $gpq, $nbus, $to)
    end
end

datafile = joinpath(dirname(@__FILE__), "case14.raw")
run_benchmark(datafile)


text:
# Grid PRS, see grid_prs.jl

"""
    weighted_interaction_graph(
        [rng::Random.AbstractRNG,]
        ising::Ising
    )::SWG.SimpleWeightedGraph

Return a weighted version of `ising.graph` where each edge is attached an independent uniform random variable.

This is a subroutine of [`PRS.generate_sample_grid_prs`](@ref).
"""
function weighted_interaction_graph(
    rng::Random.AbstractRNG,
    ising::Ising
)::SWG.SimpleWeightedGraph
    return uniform_weighted_graph(rng, ising.graph)
end

weighted_interaction_graph(ising::Ising) = weighted_interaction_graph(Random.default_rng(), ising)

@doc raw"""
    initialize_cells(
        ising::Ising{T}
    )::Vector{GraphCellGridPRS{T}} where {T}

Each node of `ising.graph` is considered as a `cell` of type [`PRS.GraphCellGridPRS`](@ref) such that `cell.window` is a [`PRS.GraphNode`](@ref) and `cell.value` initialized to `zero(T)`.

This is a subroutine of [`PRS.generate_sample_grid_prs`](@ref).
"""
function initialize_cells(
    ising::Ising{T}
)::Vector{GraphCellGridPRS{T}} where {T}
    return [GraphCellGridPRS(GraphNode(i), zero(T)) for i in 1:LG.nv(ising.graph)]
end

"""
    generate_sample(
        [rng::Random.AbstractRNG,]
        pp::Ising;
        win::GraphNode
    )

Generate an exact sample from the marginal distribution of `pp` at state indexed by `win.idx`.
"""
function generate_sample(
    rng::Random.AbstractRNG,
    pp::Ising,
    win::GraphNode
)
    return generate_sample(rng, pp, win.idx)
end

# generate_sample(pp::Ising, win::GraphNode) = generate_sample(Random.default_rng(), pp, win)

@doc raw"""
    gibbs_interaction(
        ising::Ising{T},
        xᵢ::GraphCellGridPRS{T},
        xⱼ::GraphCellGridPRS{T}
    )::Real where {T}

Compute the Gibbs interaction of [`PRS.Ising`](@ref)

```math
    \exp(J x_i x_j - |J|)).
```

**Note** the Gibbs interaction is normalized in be in [0, 1] to fit the framework of [MoKr20](@cite) and [FeViYi19](@cite).
"""
function gibbs_interaction(
    ising::Ising{T},
    xᵢ::GraphCellGridPRS{T},
    xⱼ::GraphCellGridPRS{T}
)::Real where {T}
    return exp(ising.J * xᵢ.value * xⱼ.value - abs(ising.J))
end

@doc raw"""
    is_inner_interaction_possible(
        ising::Ising{T},
        xᵢ::GraphCellGridPRS{T},
        xⱼ::GraphCellGridPRS{T}
    ) = true

Assume `xᵢ` and `xⱼ` are neighboring sites in the weighted interaction graph constructed by [`PRS.weighted_interaction_graph`](@ref) from `ising.graph` and already identified in the set of variables to be resampled [`PRS.generate_sample_grid_prs`](@ref).
Given the states of `xᵢ` and `xⱼ`, check whether a new assigment of ``U_{ij}`` (the weight of edge ``\{i,j\}``) can induce the bad event

```math
    U_{ij} > \exp(J x_i x_j - |J|))
```

This is a subroutine of [`PRS.generate_sample_grid_prs`](@ref).
"""
function is_inner_interaction_possible(
    ising::Ising{T},
    xᵢ::GraphCellGridPRS{T},
    xⱼ::GraphCellGridPRS{T}
)::Bool where {T}
    return (sign(ising.J) * xᵢ.value * xⱼ.value) < 0
end

@doc raw"""
    is_outer_interaction_possible(
        ising::Ising{T},
        xᵢ::GraphCellGridPRS{T},
        xⱼ::GraphCellGridPRS{T}
    ) = true

Assume `xᵢ` and `xⱼ` are neighboring sites in the weighted interaction graph constructed by [`PRS.weighted_interaction_graph`](@ref) from `ising.graph`.
Given the state of `xᵢ`, one can always find a new assigment of `xⱼ` and/or ``U_{ij}`` (the weight of edge ``\{i,j\}``) can induce the bad event

```math
    U_{ij} > \exp(J x_i x_j - |J|))
```

This is a subroutine of [`PRS.generate_sample_grid_prs`](@ref).
"""
function is_outer_interaction_possible(
    ising::Ising{T},
    xᵢ::GraphCellGridPRS{T},
    xⱼ::GraphCellGridPRS{T}
)::Bool where {T}
    return true
end


text:
#=
In this example we'll compare the optimization of a relatively simple problem: Logistic Regression.

Consider a binary classification dataset
  $\Omega = \{(x_i, y_i) \subset \mathbb{R}^p\times\{0,1\}:\ i = 1, \dots, n\}$.

The logistic regression problem defines a model
$h_{\beta}(x) = \sigma(\hat{x}^T \beta) = \sigma(\beta_0 + x^T\beta_{1:p})$,
where
$\hat{x} = \begin{bmatrix} 1 \\ x \end{bmatrix}$.
The value of $\beta$ is found by finding the maximum of the log-likelihood function

$$\ell(\beta) = \frac{1}{n} \sum_{i=1}^n y_i \ln \big(h_{\beta}(x_i)\big) + (1 - y_i) \ln\big(1 - h_{\beta}(x_i)\big).$$

We'll input the gradient of this function manually. It is given by

$$\nabla \ell(\beta) = \frac{1}{n} \sum_{i=1}^n \big(y_i - h_{\beta}(x_i)\big) \hat{x}_i = \frac{1}{n} \begin{bmatrix} e^T \\ X^T \end{bmatrix} (y - h_{\beta}(X)), $$

where $e$ is the vector with all components equal to 1.

Download [data.csv](https://github.com/abelsiqueira/julia-tutorial-at-nlesc/raw/main/assets/data.csv) to follow the example.

In Python, we can define the function $-\ell(\beta)$ to minimize, and its derivatives, as follows:

```python
import numpy as np
import pandas as pd
from scipy.optimize import minimize

df = pd.read_csv('data.csv')
X = df.iloc[:,0:-1].values
y = df.iloc[:,-1].values
n, p = X.shape

def sigmoid(t):
    return 1 / (1 + np.exp(-t))

def myfun(beta):
    hbeta = sigmoid(beta[0] + X.dot(beta[1:]))
    return -np.sum(
        y * np.log(hbeta + 1e-8) + (1 - y) * np.log(1 - hbeta + 1e-8)
    ) / n

def myjac(beta):
    hbeta = sigmoid(beta[0] + X.dot(beta[1:]))
    dif = hbeta - y
    return np.concatenate((np.array([np.sum(dif)]), X.T.dot(dif))) / n

beta0 = np.zeros(p + 1)
res = minimize(myfun, beta0, jac=myjac, method='l-bfgs-b', tol=1e-4)
beta = res.x
hbeta = sigmoid(beta[0] + X.dot(beta[1:]))
ypred = np.round(hbeta)
accuracy = sum(y == ypred) / n
print('accuracy = ', accuracy)
```
If you're using Jupyter, I also ask you to run
```
%%timeit
res = minimize(myfun, beta0, jac=myjac, method='l-bfgs-b', tol=1e-4)
```
So that we can compare the time

Here's a similar Julia version
=#
using CSV, DataFrames, Optim, LinearAlgebra

path = joinpath("assets", "data.csv")
df = DataFrame(CSV.File(path))
X = Matrix(df[:,1:end-1])
y = df[:,end]
n, p = size(X)

sigmoid(t) = 1 / (1 + exp(-t))

function myfun(β, X, y)
  @views hβ = sigmoid.(β[1] .+ X * β[2:end])
  out = sum(
    yᵢ * log(ŷᵢ + 1e-8) + (1 - yᵢ) * log(1 - ŷᵢ + 1e-8)
    for (yᵢ, ŷᵢ) in zip(y, hβ)
  )
  return -out / length(y)
end

function myjac(out, β, X, y)
  n = length(y)
  @views δ = (sigmoid.(β[1] .+ X * β[2:end]) - y) / n
  out[1] = sum(δ)
  out[2:end] .= X' * δ
  return out
end

output = optimize(
  β -> myfun(β, X, y),
  (out, β) -> myjac(out, β, X, y),
  zeros(p + 1),
  LBFGS(),
  Optim.Options(
    g_tol = 1e-4,
  ),
)

β = Optim.minimizer(output)
hβ = sigmoid.(β[1] .+ X * β[2:end])
ŷ = round.(hβ)
accuracy = sum(y .== ŷ) / n

#=
The beginning of the codes is similar: read the file, create `X` and `y`.

We then define the functions. Some main differences:
- Instead of writing the objective vectorized, I use a loop over the indices.
- To avoid creating a vector for `β[2:end]`, I use `@views` at the beginning of the line. The macro automatically transforms `β[2:end]` into `view(β, 2:end)`.
- The Julia `myjac` function receives an additional argument for the memory that it should reuse (`gx`) for the output.

Now, the specific part. In Python, we use SciPy's optimize, and in Julia we're using Optim.jl.
The arguments are similar, although Optim's arguments are positional.

We can benchmark Optim's code to compare against SciPy's optimize:
=#
using BenchmarkTools

β₀ = zeros(p + 1)
@benchmark optimize(
  β -> myfun(β, X, y),
  (out, β) -> myjac(out, β, X, y),
  β₀,
  LBFGS(),
  Optim.Options(
    g_tol = 1e-4,
  ),
)
#=
In Julia, we don't have a single optimization library.
In fact, I'm co-creator of the [JuliaSmoothOptimizers](https://juliasmoothoptimizers.github.io) organization, which "competes" against Optim.
We focus on tools for Nonlinear Optimization developers, instead of focusing on users, so our basic interface is a lot more raw.
On the other hand, we claim that our methods are faster, but notice my bias.

NLPModels alone doesn't have a simple user interface, so we create a struct `MyLogisticRegression` to hold the problem and to indicate the objective and gradient.
The struct fields are not obvious and you have to follow the docs of NLPModels to learn how to do it.
But you can check the speed gain.
=#

using JSOSolvers, Logging, NLPModels

struct MyLogisticRegression <: AbstractNLPModel{Float64, Vector{Float64}}
  meta :: NLPModelMeta{Float64, Vector{Float64}}
  counters :: Counters
  X
  y

  function MyLogisticRegression(X, y)
    meta = NLPModelMeta(p+1, x0=zeros(p+1))
    return new(meta, Counters(), X, y)
  end
end

NLPModels.obj(nlp :: MyLogisticRegression, x) = myfun(x, nlp.X, nlp.y)
NLPModels.grad!(nlp :: MyLogisticRegression, x, gx) = myjac(gx, x, nlp.X, nlp.y)

nlp = MyLogisticRegression(X, y)
output = lbfgs(nlp, atol=1e-4, rtol=1e-4)
println(output)

@benchmark with_logger(NullLogger()) do
  lbfgs(nlp, atol=1e-4, rtol=1e-4)
end

# TODO: Check progress of ManualNLPModels (check after 18-Oct)

#=
One noteworthy aspect of this comparison is that Python uses the L-BFGS-B Fortran algorithm, which is a reference implementation of a classic method.
The Julia versions are pure Julia.
=#

max_stars_repo_name: chunggnuhc/julia-tutorial-at-nlesc

max_issues_repo_name: chunggnuhc/julia-tutorial-at-nlesc

max_forks_repo_name: chunggnuhc/julia-tutorial-at-nlesc


text:
# Utility functions

ord = [
	21, 17, 13,
	9, 5, 1,
	22, 18, 14,
	10, 6, 2,
	23, 19, 15,
	11, 7, 3,
	24, 20, 16,
	12, 8, 4]

function estimate_snr(x::AbstractVector; fs=10)
    # Estimate SNR for a 30-102.4s segment
    pow = power(periodogram(x, nfft=1024, window=hanning))
    ps = argmax(pow)
    n = mean(pow[Not(ps, 2ps)]) * 0.73
    w = (pow[ps] * fs / 1024 - n) / n
    return w
end

"""
    moving_stats(x::AbstractVector, L::Int)

Return moving average and variance of x using a window length L
"""
function moving_stats(x::AbstractVector{T}, L::Int) where T
    # Use filters to calculate running mean and var
    avg_filt = ones(T, L) ./ L
    var_filt = ones(T, L) ./ (L - 1) # unbiased

    moving_avg = filt(avg_filt, 1, x)
    moving_var = filt(var_filt, 1, (x .- moving_avg) .^ 2)

    return moving_avg, moving_var
end

"""
    moving_stats(x::AbstractVector, w::AbstractVector{Int})

Return moving average and variance of x using a window lengths w
"""
function moving_stats(x::AbstractVector{T}, w::AbstractVector{Int}) where T
    moving_avg = zeros(T, length(x))
    moving_var = zeros(T, length(x))
    @views for i in eachindex(x)
        win = i < w[i] ? i : w[i]
        moving_avg[i] = mean(x[i-win+1:i])
        moving_var[i] = sum((x[i-win+1:i] .- moving_avg[i-win+1:i]) .^ 2) / (w[i]-1)
    end
    return moving_avg, moving_var
end

"""
    apply2seg(f::Function, x::AbstractMatrix{T} , n::Integer)

Convenience function to apply a function f to segments of length n to the matrix x.
"""
function apply2seg(f::Function, x::AbstractMatrix{T} , n::Integer) where T
    # Assumes a mat comes in and a vec goes out
    # Not much better than a mapreduce but it includes last segment
    ra = 1:n:size(x, 1)-n
    out = Vector{T}(undef, size(x, 1))
    for i in ra
        out[i:i+n-1] = f(@view x[i:i+n-1, :])
    end
    out[ra.stop+n:end] = f(@view x[ra.stop+n:end, :])
    return out
end

"""
    mat_shape(x, n)

Reorders the columns of a matrix to be in the proper sensor order.
"""
function mat_shape(x::AbstractVector, n=3)
    return view(x, mapreduce(i -> ord .+ (24 * i), vcat, 0:n-1))
end

function mat_shape(x::AbstractMatrix, n=3)
    return view(x, :, mapreduce(i -> ord .+ (24 * i), vcat, 0:n-1))
end

"""
	reshape_psm(psm, n_mats)
Takes in a PSM and reshapes it to be in the correct order.

ex. Nx72 -> 9x8xN

ex. Nx24 -> 3x8xN
"""
function reshape_psm(x::AbstractMatrix, n=div(size(x, 2), 24))
	new_shape = reshape(mat_shape(x, n)', 3, 8, n, :)
	return reduce(vcat, eachslice(new_shape, dims=3))
end

"""
    active_sensors(x, thresh)

    Return the indices of sensors whose mean value is greater than thresh.
"""
active_sensors(x::AbstractMatrix, thresh=0.4*2046) = vec(mean(x, dims=1) .> thresh)

"""
    choose_ref(x)

Return the index of the sensor with the greatest power.
"""
function choose_ref(x::AbstractMatrix)
    # Choose
    out = sum(abs2, x, dims=1) |> argmax
    return out[2]
end

extract_ref(x::AbstractMatrix) = x[:, choose_ref(x)]

"""
    polarity_flip(x)

Flips the polarity of sensors based on their PCC with the reference sensor.
"""
polarity_flip(x) = sign.(cor(x, @view x[:, choose_ref(x)]))' .* x

"""
    sfm(x)
Returns the spectral flatness measure of a signal.
"""
function sfm(x)
    s = abs2.(fft(x))
    return geomean(s) / mean(s)
end

"""
    active_sfm(x, n, thresh)

Return the indices of sensors with a mean spectral flatness measure, taken in segments of
length n, above the threshold in dB.
"""
function active_sfm(x::AbstractMatrix, n, thresh=-50)
    fo(sig) = mapreduce(i -> sfm(sig[i:i+n-1]), vcat, 1:n:size(x, 1)-n)
    s = mapreduce(fo, hcat, eachcol(x))
    actives = pow2db.(s) .> thresh
    return vec(mean(actives, dims=1) .> 0.4)
end


text:
using Exfiltrator
using Test

@testset "Exfiltrator.jl" begin
    function f(x)
        @exfiltrate
        @exfiltrate A x
        @exfiltrate B
        @exfiltrate_push! C x
        @exfiltrate_push! D
    end
    f(2)
    f(3)
    @test x == 3
    @test A == 3
    @test B.x == 3
    @test C == [2, 3]
    @test D[1].x == 2
    @test D[2].x == 3
    f(:a)
    @test x == :a
end


text:
using .Gadfly

"""
    shewhart(
        df::AbstractDataFrame, # DataFrame containing columns :Date & :Value (or below)
        nominal::Real, # Nominal value for :Value column
        onesigma::Real; # Nominal 1σ variation in :Value column
        timestamp::Union{Symbol,AbstractString}=:Date, # Alternative to :Date
        value::Union{Symbol,AbstractString}=:Value, # Alternative to :Value
        result::Union{Missing,Function}=missing,
        color::Bool=true # Use only color vs only black and white
    )

Plot a Shewhart-style control chart from the data in `df`.  The `DataFrame` is evaluated using
the `weco(...)` function to compute the WECO rules based on the specified `nominal` (mean value)
and `onesigma` (standard-deviation).

`color` controls whether the data is plotted in colors (green=>meets, red=>fails, red line=>3σ, yellow line=>2σ)
or black-and-white (white fill=>meets, black fill=>fails, long dashed gray line=>3σ, short dash gray line=>2σ)
"""
function shewhart(
    df::AbstractDataFrame, # DataFrame containing columns :Date & :Value (or below)
    nominal::Real, # Nominal value for :Value column
    onesigma::Real; # Nominal 1σ variation in :Value column
    timestamp::Union{Symbol,AbstractString}=:Date, # Alternative to :Date
    value::Union{Symbol,AbstractString}=:Value, # Alternative to :Value
    result::Union{Missing,Function}=missing,
    color::Bool=true, # Use colors vs only black and white
    title::Union{AbstractString,Nothing}=nothing,
)
    df = sort(df, timestamp)
    res = weco(df, nominal, onesigma, timestamp=timestamp, value=value, result=result)
    gt = isnothing(title) ? [] : [ Guide.title(title) ]
    if color
        plot(
            layer(yintercept=[3*onesigma, -3*onesigma].+nominal, Geom.hline(color="red")),
            layer(yintercept=[2*onesigma, -2*onesigma].+nominal, Geom.hline(color="yellow")),
            layer(res, x=timestamp, y=value, color=:Passes, Geom.point),
            Scale.color_discrete_manual("lime", "firebrick", levels = [ true, false ]), gt...
        )
    else
        plot(
            layer(yintercept=[3*onesigma, -3*onesigma].+nominal, Geom.hline(color="gray50", style=:dash)),
            layer(yintercept=[2*onesigma, -2*onesigma].+nominal, Geom.hline(color="gray50", style=:dot)),
            layer(res, x=timestamp, y=value, color=:Passes, Geom.point, Theme(discrete_highlight_color=c->"gray50")),
            Scale.color_discrete_manual("gray80", "black", levels = [ true, false ]), gt...
        )
    end
end


text:
module SwiftObjectStores

using PyCall
using JLD
using FileIO

const service = pyimport("swiftclient.service")

include("service_options.jl")
include("responses.jl")
include("upload.jl")
include("download.jl")
include("management.jl")
include("util.jl")
end # module


text:
meancov(u) = u[1], u[2]
function correct(u::T, v, H; droptol=1e-7) where T
    x, Ppred = meancov(u)
    y, R = meancov(v)
    yres = y - H*x # innovation residual

    S = (H*Ppred*H' + R) # innovation covariance

    # not sure what is the right thing.
    if false # Dense matrix algebra
        K1 = Ppred*H'*inv(Matrix(S)) # Kalman gain
        @show sparsity(K1)
        K = droptol!(sparse(K1), droptol)
        @show sparsity(K)
    elseif true # Sparse cholesky and dense H*Ppred
        K = (cholesky(Hermitian(S))\(H*Ppred))' # Kalman gain # sparse backsolve is missing
        @show sparsity(K')
        droptol!(K', droptol)
        @show sparsity(K')
    else # Krylov method
        #K = Ppred*H'*inv(S)
        Pre = preconditioner(S, 16, 0.02)+0.01I
        K = srsolve(Pre*S, Pre*Ppred*H', tol = 5e-5, droptol=droptol)
    end
    #sparse(SuiteSparse.CHOLMOD.spsolve(0, cholesky(outer(sprand(10,10,0.5))), SuiteSparse.CHOLMOD.Sparse(sprand(10,10,0.1))))
    x = x + K*yres
    P = (I - K*H)*Ppred*(I - K*H)' + K*R*K'
    T((x, P)), yres, S
end


text:
using Catalyst, ModelingToolkit

# naming tests
@parameters k
@variables t, A(t)
rx = Reaction(k, [A], nothing)
function rntest(rn, name)
    @test nameof(rn) == name
    @test isequal(species(rn)[1], ModelingToolkit.unwrap(A))
    @test isequal(parameters(rn)[1], ModelingToolkit.unwrap(k))
    @test reactions(rn)[1] == rx
end

function emptyrntest(rn, name)
    @test nameof(rn) == name
    @test numreactions(rn) == 0
    @test numspecies(rn) == 0
    @test numreactionparams(rn) == 0
end

rn = @reaction_network name begin
    k, A --> 0
end k
rntest(rn, :name)

name = :blah
rn = @reaction_network $name begin
    k, A --> 0
end k
rntest(rn, :blah)

rn = @reaction_network begin
    k, A --> 0
end k
rntest(rn, nameof(rn))

function makern(; name)
    @reaction_network $name begin
        k, A --> 0
    end k
end
@named testnet = makern()
rntest(testnet, :testnet)

rn = @reaction_network name
emptyrntest(rn, :name)

rn = @reaction_network $name
emptyrntest(rn, :blah)

# test variables that appear only in rates and aren't ps
# are categorized as species
rn = @reaction_network begin
    π*k*D*hill(B,k2,B*D*H,n), 3*A  --> 2*C
end k k2 n
@parameters k,k2,n
@variables t,A(t),B(t),C(t),D(t),H(t)
@test issetequal([A,B,C,D,H], species(rn))
@test issetequal([k,k2,n], parameters(rn))


text:
"""BivarateNormals

2-d gaussians and their derivatives.
"""
module BivariateNormals

using StaticArrays
using ..Celeste: Const, @aliasscope, @unroll_loop

# TODO: Many parameter and type names sound specific to "galaxies". However,
# they're actually general to bivariate normals or mixtures thereof.
# Parameters should be renamed to reflect this.

# Hardcoding this might not be needed but it is critical that this is
# compile time constant.
const GAL_SHAPE_IDS_LENGTH = 3

"""
Unpack a rotation-parameterized BVN covariance matrix.

Args:
 - ab: The ratio of the minor to the major axis.
 - angle: Rotation angle (in radians)
 - scale: The major axis.

Returns:
  The 2x2 covariance matrix parameterized by the inputs.
"""
function get_bvn_cov{NumType <: Number}(
    ab::NumType, angle::NumType, scale::NumType)

    if NumType <: AbstractFloat
        @assert 0 < scale
        @assert 0 < ab <= 1.
    end

    cp = cos(angle)
    sp = sin(angle)
    ab_term = (ab ^ 2 - 1)
    scale_squared = scale ^ 2
    off_diag_term = -scale_squared * cp * sp * ab_term
    @SMatrix NumType[scale_squared * (1 + ab_term * (sp ^ 2))  off_diag_term;
                     off_diag_term                             scale_squared * (1 + ab_term * (cp ^ 2))]
end

"""
Pre-allocated memory for quantities related to derivatives of bivariate
normals.
"""
struct BivariateNormalDerivatives{NumType <: Number}

  # Pre-allocated memory for py1, py2, and f when evaluating BVNs
  py1::Array{NumType, 1}
  py2::Array{NumType, 1}
  f_pre::Array{NumType, 1}

  # Derivatives of a bvn with respect to (x, sig).
  bvn_x_d::Array{NumType, 1}
  bvn_sig_d::Array{NumType, 1}
  bvn_xx_h::SizedMatrix{2, 2, NumType, 2}
  bvn_xsig_h::SizedMatrix{2, 3, NumType, 2}
  bvn_sigsig_h::SizedMatrix{3, 3, NumType, 2}

  # intermediate values used in d bvn / d(x, sig)
  dpy1_dsig::Array{NumType, 1}
  dpy2_dsig::Array{NumType, 1}

  # Derivatives of a bvn with respect to (pos, shape)
  bvn_u_d::Array{NumType, 1}
  bvn_uu_h::SizedMatrix{2, 2, NumType, 2}
  bvn_s_d::Array{NumType, 1}
  bvn_ss_h::SizedMatrix{GAL_SHAPE_IDS_LENGTH, GAL_SHAPE_IDS_LENGTH, NumType, 2}
  bvn_us_h::SizedMatrix{2, GAL_SHAPE_IDS_LENGTH, NumType, 2}

  function (::Type{BivariateNormalDerivatives{NumType}}){NumType}()
    py1 = zeros(NumType, 1)
    py2 = zeros(NumType, 1)
    f_pre = zeros(NumType, 1)

    bvn_x_d = zeros(NumType, 2)
    bvn_sig_d = zeros(NumType, 3)
    bvn_xx_h = zeros(NumType, 2, 2)
    bvn_xsig_h = zeros(NumType, 2, 3)
    bvn_sigsig_h = zeros(NumType, 3, 3)

    dpy1_dsig = zeros(NumType, 3)
    dpy2_dsig = zeros(NumType, 3)

    # Derivatives wrt pos.
    bvn_u_d = zeros(NumType, 2)
    bvn_uu_h = zeros(NumType, 2, 2)

    # Shape deriviatives.  Here, s stands for "shape".
    bvn_s_d = zeros(NumType, GAL_SHAPE_IDS_LENGTH)

    # The hessians.
    bvn_ss_h = zeros(NumType, GAL_SHAPE_IDS_LENGTH, GAL_SHAPE_IDS_LENGTH)
    bvn_us_h = zeros(NumType, 2, GAL_SHAPE_IDS_LENGTH)

    new{NumType}(py1, py2, f_pre,
        bvn_x_d, bvn_sig_d, bvn_xx_h, bvn_xsig_h, bvn_sigsig_h,
        dpy1_dsig, dpy2_dsig,
        bvn_u_d, bvn_uu_h, bvn_s_d, bvn_ss_h, bvn_us_h)
  end
end

function zero!(bvn_derivs::BivariateNormalDerivatives{T}) where {T}
    x = zero(T)
    fill!(bvn_derivs.py1, x)
    fill!(bvn_derivs.py2, x)
    fill!(bvn_derivs.f_pre, x)
    fill!(bvn_derivs.bvn_x_d, x)
    fill!(bvn_derivs.bvn_sig_d, x)
    fill!(bvn_derivs.bvn_xx_h, x)
    fill!(bvn_derivs.bvn_xsig_h, x)
    fill!(bvn_derivs.bvn_sigsig_h, x)
    fill!(bvn_derivs.dpy1_dsig, x)
    fill!(bvn_derivs.dpy2_dsig, x)
    fill!(bvn_derivs.bvn_u_d, x)
    fill!(bvn_derivs.bvn_uu_h, x)
    fill!(bvn_derivs.bvn_s_d, x)
    fill!(bvn_derivs.bvn_ss_h, x)
    fill!(bvn_derivs.bvn_us_h, x)
    return bvn_derivs
end

"""
Relevant parameters of a bivariate normal distribution.

Args:
  the_mean: The mean as a 2x1 column vector
  the_cov: The covaraiance as a 2x2 matrix
  weight: A scalar weight

Attributes:
   the_mean: The mean argument
   precision: The inverse of the_cov
   z: The weight times the normalizing constant.
   dsiginv_dsig: The derivative of sigma inverse with respect to sigma.
   major_sd: The standard deviation of the major axis.
"""
struct BvnComponent{NumType <: Number}
    the_mean::SVector{2,NumType}
    precision::SMatrix{2,2,NumType,4}
    z::NumType
    dsiginv_dsig::SMatrix{3,3,NumType,9}
    major_sd::NumType
end

function BvnComponent{T1<:Number,T2<:Number,T3<:Number}(
    the_mean::SVector{2,T1}, the_cov::SMatrix{2,2,T2,4}, weight::T3,
    calculate_siginv_deriv::Bool=true)

    NumType = promote_type(T1,T2,T3)
    c = 1 / (sqrt(det(the_cov)) * 2pi)
    major_sd = sqrt(max( the_cov[1, 1], the_cov[2, 2] ))

    precision = inv(the_cov)

    if calculate_siginv_deriv
        # Derivatives of Sigma^{-1} with respect to sigma.  These are the second
        # derivatives of log|Sigma| with respect to sigma.
        # dsiginv_dsig[a, b] is the derivative of sig^{-1}[a] / d sig[b]

        dsiginv_dsig11 = -precision[1, 1] ^ 2
        dsiginv_dsig12 = -2 * precision[1, 1] * precision[1, 2]
        dsiginv_dsig13 = -precision[1, 2] ^ 2

        dsiginv_dsig21 = -precision[1, 1] * precision[2, 1]
        dsiginv_dsig22 =
          -(precision[1, 1] * precision[2, 2] + precision[1, 2] ^ 2)
        dsiginv_dsig23 = -precision[2, 2] * precision[1, 2]

        dsiginv_dsig31 = -precision[1, 2] ^ 2
        dsiginv_dsig32 = - 2 * precision[2, 2] * precision[2, 1]
        dsiginv_dsig33 = -precision[2, 2] ^ 2

        dsiginv_dsig = @SMatrix NumType[ dsiginv_dsig11 dsiginv_dsig12 dsiginv_dsig13;
                                           dsiginv_dsig21 dsiginv_dsig22 dsiginv_dsig23;
                                           dsiginv_dsig13 dsiginv_dsig32 dsiginv_dsig33 ]

        BvnComponent{NumType}(the_mean, precision, c * weight,
                              dsiginv_dsig, major_sd)
    else
        BvnComponent{NumType}(the_mean, precision, c * weight,
                              zeros(SMatrix{3,3,NumType,9}), major_sd)
    end
end

"""
Return quantities related to the pdf of an offset bivariate normal.

Args:
  - bvn_derivs: BivariateNormalDerivatives to be updated in place
  - bmc: A bivariate normal component
  - x: A 2x1 vector containing a mean offset to be applied to bmc

Returns:
  In bvn_derivs, sets these values in place:
  - py1: The first row of the precision times (x - the_mean)
  - py2: The second row of the precision times (x - the_mean)
  - The density of the bivariate normal times the weight.
"""
function eval_bvn_pdf!{NumType <: Number}(
    bvn_derivs::BivariateNormalDerivatives{NumType},
    bmc::BvnComponent{NumType}, x::SVector{2,Float64})

  bvn_derivs.py1[1] =
    bmc.precision[1,1] * (x[1] - bmc.the_mean[1]) +
    bmc.precision[1,2] * (x[2] - bmc.the_mean[2])
  bvn_derivs.py2[1] =
    bmc.precision[2,1] * (x[1] - bmc.the_mean[1]) +
    bmc.precision[2,2] * (x[2] - bmc.the_mean[2])
  bvn_derivs.f_pre[1] =
    bmc.z * exp(-0.5 * ((x[1] - bmc.the_mean[1]) * bvn_derivs.py1[1] +
                        (x[2] - bmc.the_mean[2]) * bvn_derivs.py2[1]))
end

##################
# Derivatives

"""
Calculate the value, gradient, and hessian of
  -0.5 * x' sigma^-1 x - 0.5 * log|sigma|
with respect to x and sigma.  This assumes that
  bvn_derivs.py1, .py2, and .f_pre have already been populated
  witha call to eval_bvn_pdf!.

Args:
  - elbo_vars: A data structure with pre-allocated intermediate variables.
  - bvn: A bivariate normal component to get derivatives for.
  - calculate_x_hess: Whether to calcualte x Hessian terms.
  - calculate_sigma_hess: Whether to calcualte sigma Hessian terms.
"""
function get_bvn_derivs!{NumType <: Number}(
    bvn_derivs::BivariateNormalDerivatives{NumType},
    bvn::BvnComponent{NumType}, calculate_x_hess::Bool,
    calculate_sigma_hessian::Bool)

  @inbounds begin

    # Gradient with respect to x.
    bvn_x_d = bvn_derivs.bvn_x_d
    bvn_x_d[1] = -bvn_derivs.py1[1]
    bvn_x_d[2] = -bvn_derivs.py2[1]

    if calculate_x_hess
      bvn_xx_h = bvn_derivs.bvn_xx_h

      # Hessian terms involving only x
      bvn_xx_h[1, 1] = -bvn.precision[1, 1]
      bvn_xx_h[2, 2] = -bvn.precision[2, 2]
      bvn_xx_h[1, 2] = bvn_xx_h[2, 1] = -bvn.precision[1 ,2]
    end

    # The first term is the derivative of -0.5 * x' Sigma^{-1} x
    # The second term is the derivative of -0.5 * log|Sigma|
    bvn_sig_d = bvn_derivs.bvn_sig_d
    bvn_sig_d[1] =
      0.5 * bvn_derivs.py1[1] * bvn_derivs.py1[1] - 0.5 * bvn.precision[1, 1]
    bvn_sig_d[2] =
      bvn_derivs.py1[1] * bvn_derivs.py2[1]             - bvn.precision[1, 2]
    bvn_sig_d[3] =
      0.5 * bvn_derivs.py2[1] * bvn_derivs.py2[1] - 0.5 * bvn.precision[2, 2]

    if calculate_sigma_hessian

      # Hessian calculation for terms containing sigma.

      # Derivatives of py1 and py2 with respect to s11, s12, s22 in that order.
      # These are used for the hessian calculations.
      dpy1_dsig = bvn_derivs.dpy1_dsig
      dpy1_dsig[1] = -bvn_derivs.py1[1] * bvn.precision[1,1]
      dpy1_dsig[2] = -bvn_derivs.py2[1] * bvn.precision[1,1] -
                      bvn_derivs.py1[1] * bvn.precision[1,2]
      dpy1_dsig[3] = -bvn_derivs.py2[1] * bvn.precision[1,2]

      dpy2_dsig = bvn_derivs.dpy2_dsig
      dpy2_dsig[1] = -bvn_derivs.py1[1] * bvn.precision[1,2]
      dpy2_dsig[2] = -bvn_derivs.py1[1] * bvn.precision[2,2] -
                      bvn_derivs.py2[1] * bvn.precision[1,2]
      dpy2_dsig[3] = -bvn_derivs.py2[1] * bvn.precision[2,2]

      # Hessian terms involving only sigma
      bvn_sigsig_h = bvn_derivs.bvn_sigsig_h
      for s_ind=1:3
        # Differentiate with respect to s_ind second.
        bvn_sigsig_h[1, s_ind] = #bvn_sigsig_h[s_ind, 1] =
          bvn_derivs.py1[1] * dpy1_dsig[s_ind] - 0.5 * bvn.dsiginv_dsig[1, s_ind]

        # d log|sigma| / dsigma12 is twice lambda12.
        bvn_sigsig_h[2, s_ind] =
          bvn_derivs.py1[1] * dpy2_dsig[s_ind] + bvn_derivs.py2[1] * dpy1_dsig[s_ind] -
          bvn.dsiginv_dsig[2, s_ind]

        bvn_sigsig_h[3, s_ind] = #bvn_sigsig_h[s_ind, 3] =
          bvn_derivs.py2[1] * dpy2_dsig[s_ind] - 0.5 * bvn.dsiginv_dsig[3, s_ind]
      end

      # Hessian terms involving both x and sigma.
      # Note that dpyA / dxB = bvn.precision[A, B]
      bvn_xsig_h = bvn_derivs.bvn_xsig_h
      for x_ind=1:2
        bvn_xsig_h[x_ind, 1] = bvn_derivs.py1[1] * bvn.precision[1, x_ind]
        bvn_xsig_h[x_ind, 2] =
          bvn_derivs.py1[1] * bvn.precision[2, x_ind] +
          bvn_derivs.py2[1] * bvn.precision[1, x_ind]
        bvn_xsig_h[x_ind, 3] = bvn_derivs.py2[1] * bvn.precision[2, x_ind]
      end
    end
  end
end

"""
The derivatives of sigma with respect to the galaxy shape parameters.  In
each case, sigma is [Sigma11, Sigma12, Sigma22], and the galaxy shape
parameters are indexed by GalaxyShapeParams.
 - j: A Jacobian matrix Sigma x GalaxyShapeParams of
      d Sigma / d GalaxyShapeParams
 - t: A Sigma x GalaxyShapeParams x GalaxyShapeParams tensor of second
      derivatives d2 Sigma / d GalaxyShapeParams d GalaxyShapeParams.
"""
struct GalaxySigmaDerivs{NumType <: Number}
    j::SMatrix{3,GAL_SHAPE_IDS_LENGTH,NumType,9}
    t::SArray{Tuple{3,GAL_SHAPE_IDS_LENGTH,GAL_SHAPE_IDS_LENGTH},NumType,3,27}
end

"""
Args:
  - gal_angle: Phi in the notes
  - gal_ab: Rho in the notes
  - gal_scale: Lower case sigma in the notes
  - XiXi: The value of sigma.

Note that nubar is not included.
"""
function GalaxySigmaDerivs{NumType <: Number}(
    gal_angle::NumType, gal_ab::NumType, gal_scale::NumType,
    XiXi::SMatrix{2,2,NumType,4}, nuBar::Float64=1.0, calculate_tensor::Bool=true)

  cos_sin = cos(gal_angle)sin(gal_angle)
  sin_sq  = sin(gal_angle)^2
  cos_sq  = cos(gal_angle)^2

  j = hcat(2 * gal_ab * gal_scale^2 * SVector{3,NumType}(sin_sq, -cos_sin, cos_sq),
           gal_scale^2 * (gal_ab^2 - 1) * SVector{3,NumType}(2cos_sin, sin_sq - cos_sq, -2cos_sin),
           2 * SVector{3,NumType}(XiXi[1], XiXi[2], XiXi[4]) / gal_scale)

  if calculate_tensor
    # Second derivatives.

    t = SArray{Tuple{3,3,3}, NumType, 3, 27}(sin_sq * 2 * gal_scale^2, -cos_sin * 2 * gal_scale^2, cos_sq * 2 * gal_scale^2,
      2cos_sin * 2 * gal_scale^2 * gal_ab, (sin_sq - cos_sq) * 2 * gal_scale^2 * gal_ab, -2cos_sin * 2 * gal_scale^2 * gal_ab,
      2 * j[1, 1]  / gal_scale, 2 * j[2, 1]  / gal_scale, 2 * j[3, 1]  / gal_scale,
      2cos_sin * 2 * gal_scale^2 * gal_ab, (sin_sq - cos_sq) * 2 * gal_scale^2 * gal_ab, -2cos_sin * 2 * gal_scale^2 * gal_ab,
      (cos_sq - sin_sq) * 2 * gal_scale^2 * (gal_ab^2 - 1), 2cos_sin * 2 * gal_scale^2 * (gal_ab^2 - 1), (sin_sq - cos_sq) * 2 * gal_scale^2 * (gal_ab^2 - 1),
      2 * j[1, 2] / gal_scale, 2 * j[2, 2] / gal_scale, 2 * j[3, 2] / gal_scale,
      2 * j[1, 1]  / gal_scale, 2 * j[2, 1]  / gal_scale, 2 * j[3, 1]  / gal_scale,
      2 * j[1, 2] / gal_scale, 2 * j[2, 2] / gal_scale, 2 * j[3, 2] / gal_scale,
      2 * XiXi[1 << (1 - 1)] / gal_scale^2, 2 * XiXi[1 << (2 - 1)] / gal_scale^2, 2 * XiXi[1 << (3 - 1)] / gal_scale^2)

  else
    t = @SArray zeros(NumType, 3, 3, 3)
  end

  GalaxySigmaDerivs(j*nuBar, t*nuBar)
end

GalaxySigmaDerivs{NumType}(::Type{NumType}) = GalaxySigmaDerivs(
                                                     @SMatrix(zeros(NumType,3,GAL_SHAPE_IDS_LENGTH)),
                                                     @SArray( zeros(NumType,3,GAL_SHAPE_IDS_LENGTH,GAL_SHAPE_IDS_LENGTH)))

###################################################
# Transform derivatives into the model parameterization.

"""
Transform the bvn derivatives and hessians from (x) to the
galaxy parameters (pos). Updates bvn_u_d and bvn_uu_h in place.

bvn_x_d and bvn_xx_h should already have been set using get_bvn_derivs!()
"""
function transform_bvn_ux_derivs!{NumType <: Number}(
    bvn_derivs::BivariateNormalDerivatives{NumType},
    wcs_jacobian, calculate_hessian::Bool)

  # Gradient calculations.

  # Note that dxA_duB = -wcs_jacobian[A, B].  (It is minus the jacobian
  # because the object position affects the bvn.the_mean term, which is
  # subtracted from the pixel location as defined in bvn_sf.d.)
  bvn_u_d = bvn_derivs.bvn_u_d
  bvn_x_d = bvn_derivs.bvn_x_d
  bvn_u_d[1] =
    -(bvn_x_d[1] * wcs_jacobian[1, 1] + bvn_x_d[2] * wcs_jacobian[2, 1])
  bvn_u_d[2] =
    -(bvn_x_d[1] * wcs_jacobian[1, 2] + bvn_x_d[2] * wcs_jacobian[2, 2])

  if calculate_hessian
    # Hessian calculations.

    bvn_uu_h = bvn_derivs.bvn_uu_h
    bvn_xx_h = bvn_derivs.bvn_xx_h
    fill!(bvn_uu_h, 0.0)
    # Second derivatives involving only pos.
    # As above, dxA_duB = -wcs_jacobian[A, B] and d2x / du2 = 0.
    # TODO: time consuming **************
    @inbounds for x_id2 in 1:2, x_id1 in 1:2, u_id2 in 1:2
      inner_term = bvn_xx_h[x_id1, x_id2] * wcs_jacobian[x_id2, u_id2]
      @inbounds for u_id1 in 1:u_id2
        bvn_uu_h[u_id1, u_id2] += inner_term * wcs_jacobian[x_id1, u_id1]
      end
    end
    @inbounds bvn_uu_h[2, 1] = bvn_uu_h[1, 2]
  end
end

@generated function fast_fill!{T<:SizedMatrix}(s::T, x)
  quote
    $(Expr(:meta, :inline))
    @unroll_loop for i in 1:$(size(s, 1))
      @unroll_loop for j in 1:$(size(s, 2))
        @inbounds s[i, j] = x
      end
    end
  end
end
fast_fill!{T<:Array}(s::T, x) = fill!(s, x)

# WARNING: HUGE PERFORMANCE HOTSPOT
function transform_bvn_derivs_hessian!{NumType <: Number}(
    bvn_derivs::BivariateNormalDerivatives{NumType},
    sig_sf::GalaxySigmaDerivs{NumType},
    wcs_jacobian)
  @aliasscope begin
    # Hessian calculations.

    bvn_ss_h = bvn_derivs.bvn_ss_h
    bvn_us_h = bvn_derivs.bvn_us_h
    bvn_sig_d = Const(bvn_derivs.bvn_sig_d)
    #wcs_jacobian = Const(wcs_jacobian)
    sig_sf_j = sig_sf.j

    # Manually inlined version of fill!
    fast_fill!(bvn_ss_h, 0.0)
    fast_fill!(bvn_us_h, 0.0)

    # Second derviatives involving only shape parameters.
    # TODO: time consuming **************
    sig_sf_t = sig_sf.t
    @unroll_loop for shape_id2 in 1:GAL_SHAPE_IDS_LENGTH
      @unroll_loop for shape_id1 in 1:shape_id2
        @inbounds @unroll_loop for sig_id1 in 1:3
          bvn_ss_h[shape_id1, shape_id2] += bvn_sig_d[sig_id1] * sig_sf_t[sig_id1, shape_id1, shape_id2]
        end
      end
    end

    bvn_sigsig_h = Const(bvn_derivs.bvn_sigsig_h)
    @unroll_loop for sig_id1 in 1:3
      @unroll_loop for sig_id2 in 1:3
        @inbounds @unroll_loop for shape_id2 in 1:GAL_SHAPE_IDS_LENGTH
          inner_term = bvn_sigsig_h[sig_id1, sig_id2] * sig_sf_j[sig_id2, shape_id2]
          @unroll_loop for shape_id1 in 1:shape_id2
            bvn_ss_h[shape_id1, shape_id2] += inner_term * sig_sf_j[sig_id1, shape_id1]
          end
        end
      end
    end

    @unroll_loop for shape_id2 in 1:GAL_SHAPE_IDS_LENGTH
      @inbounds @unroll_loop for shape_id1 in 1:shape_id2
        bvn_ss_h[shape_id2, shape_id1] = bvn_ss_h[shape_id1, shape_id2]
      end
    end

    # Second derivates involving both a shape term and a pos term.
    # TODO: time consuming **************
    bvn_xsig_h = Const(bvn_derivs.bvn_xsig_h)
    @unroll_loop for shape_id in 1:GAL_SHAPE_IDS_LENGTH
      @unroll_loop for u_id in 1:2
        @unroll_loop for sig_id in 1:3
          @inbounds @unroll_loop for x_id in 1:2
            bvn_us_h[u_id, shape_id] +=
              bvn_xsig_h[x_id, sig_id] * sig_sf_j[sig_id, shape_id] * (-wcs_jacobian[x_id, u_id])
          end
        end
      end
    end
  end
end

"""
Transform all the bvn derivatives from x and sigma to pos and the model
parameters.

These values should already have been set using get_bvn_derivs!():
bvn_x_d
bvn_xx_h
bvn_sig_d
bvn_sigsig_h
bvn_xsig_h
"""
function transform_bvn_derivs!{NumType <: Number}(
    bvn_derivs::BivariateNormalDerivatives{NumType},
    sig_sf::GalaxySigmaDerivs{NumType},
    wcs_jacobian, calculate_hessian::Bool)

  # Transform the pos derivates first.
  # bvn_x_d and bvn_xx_h should already have been set using get_bvn_derivs!()
  transform_bvn_ux_derivs!(bvn_derivs, wcs_jacobian, calculate_hessian)

  # Gradient calculations.

  @aliasscope begin
      # Use the chain rule for the shape derviatives.
      # TODO: time consuming **************
      @aliasscope begin
        bvn_s_d   = bvn_derivs.bvn_s_d
        bvn_sig_d = Const(bvn_derivs.bvn_sig_d)
        sig_sf_j = sig_sf.j

        fast_fill!(bvn_s_d, 0.0)
        @unroll_loop for shape_id in 1:GAL_SHAPE_IDS_LENGTH
          @inbounds @unroll_loop for sig_id in 1:3
            bvn_s_d[shape_id] += bvn_sig_d[sig_id] * sig_sf_j[sig_id, shape_id]
          end
        end
      end

      if calculate_hessian
        transform_bvn_derivs_hessian!(bvn_derivs, sig_sf, wcs_jacobian)
      end
    end
end

end  # module


text:
# Autogenerated wrapper script for Enzyme_jll for x86_64-linux-musl-cxx03-julia_version+1.8.0
export libEnzyme

JLLWrappers.@generate_wrapper_header("Enzyme")
JLLWrappers.@declare_library_product(libEnzyme, "libEnzyme-12.so")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libEnzyme,
        "lib/libEnzyme-12.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using ACTRSimulators, Test, ACTRModels, Random, DataFrames
import ACTRSimulators: start!, press_key!
Random.seed!(8985)
include("task.jl")

scheduler = ACTRScheduler(;model_trace=true, store=true)
task = SimpleTask(;scheduler)
procedural = Procedural()
T = vo_to_chunk() |> typeof
visicon = VisualObject[]
visual_location = VisualLocation(buffer=T[])
visual = Visual(buffer=T[])
motor = Motor()
memory = [Chunk(;animal=:dog), Chunk(;animal=:cat)]
declarative = Declarative(;memory)
actr = ACTR(;scheduler, procedural, visual_location, visual, motor, declarative, visicon)

function can_stop(actr)
    c1(actr) = !actr.visual_location.state.empty
    return (c1,)
end

function stop(actr, task)
    stop!(actr.scheduler)
end

rule1 = Rule(;conditions=can_stop, action=stop, actr, task, name="Stop")
push!(procedural.rules, rule1)
run!(actr, task)
chunk = actr.visual_location.buffer[1]
@test chunk.slots == (color = :black, text = "hello", x = 300.0, y = 300.0)

observed = map(x->x.description, scheduler.complete_events)
expected = [
    "Starting",
    "Present Stimulus",
    "Selected Stop"
]
@test expected == observed


text:
#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

include("../../src/Sparso.jl")
include("../../src/simple-show.jl")
using Sparso

module X
    module Y
        module Z
            module U
                module V
                    module W
                        function f(x)
                            x + 1
                        end
                    end
                end
            end
        end
    end
end

function foo()
    X.Y.Z.U.V.W.f(1)
end

function bar(A, B)
    A * B
end

ast = code_typed(foo, (), optimize = false)
call_args = ast[1].args[3].args[2].args[1].args
println("Call args: ", call_args)
module_name, function_name = Sparso.resolve_call_names(call_args)
println("Module name: ", module_name)
println("Function name: ", function_name)

ast = code_typed(bar, (SparseMatrixCSC, Vector), optimize = false)
call_args = ast[1].args[3].args[2].args[1].args
println("\nCall args: ", call_args)
module_name, function_name = Sparso.resolve_call_names(call_args)
println("Module name: ", module_name)
println("Function name: ", function_name)

max_stars_repo_name: IntelLabs/Sparso

max_issues_repo_name: IntelLabs/Sparso

max_forks_repo_name: IntelLabs/Sparso


text:
module TestProgress
include("preamble.jl")

using Test: collect_test_logs
using Logging: LogLevel

const ProgressLevel = LogLevel(-1)

function test_only_one_progress(f)
    logs, = collect_test_logs(f, min_level = ProgressLevel)
    ps = [get(l.kwargs, :progress, nothing) for l in logs]
    @test any(isa.(ps, Real))
    @test "done" in ps
    @test length(unique([l.id for l in logs])) == 1
end

@testset "foldl" begin
    test_only_one_progress() do
        @test foldl(+, withprogress(1:100; interval=0)) == 5050
    end
end

@testset "sum" begin
    test_only_one_progress() do
        @test sum(withprogress(1:100; interval=0)) == 5050
    end
end

end  # module


text:
# This script can be used to generate reference solutions for new cases
using DelimitedFiles
using Printf
import Juqbox

include("evalGrad.jl")

caseNames =["rabi", "swap02", "flux", "cnot-lab", "cnot2", "cnot3"]
dirName = "reference_solutions"

nCases = length(caseNames)
pass = BitArray(undef, nCases)
pass .= false

q = 0
for case in caseNames
    global q += 1

    println("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
    println("Case: ", case)

     # setup the testcase (assign pcof0, params and wa)
    juliaFile = case * "-setup.jl"
    include(juliaFile);

    # evaluate fidelity & save reference sol
    refFile = dirName * "/" * case * "-ref.jld2"
    evalObjGrad(pcof0, params, wa, refFile, true) # true for saving a reference jld2 file

    println("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")

end

println("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
println("    TEST SUMMARY")

for q in 1:nCases
    println("case=", caseNames[q], " pass=", pass[q])
end


text:
### four vectors
################################################################################

function four_vectors_in_binary_decay(cosθ, ϕ;
        m1sq::Real=error("give mass 1 squared"),
        m2sq::Real=error("give mass 2 squared"),
        m0sq::Real=error("give decay mass squared"))
    E1 = (m0sq+m1sq-m2sq)/(2sqrt(m0sq));
    E2 = (m0sq-m1sq+m2sq)/(2sqrt(m0sq));
    √m0sq < √m1sq + √m2sq && error("!!! -√m0sq + √m1sq + √m2sq = $(-√m0sq + √m1sq + √m2sq)")
    p = sqrt(Kallen(m0sq,m1sq,m2sq))/(2sqrt(m0sq));
    #
    sinθ = sqrt(1-cosθ^2)
    p1 = MVector( p*sinθ*cos(ϕ),  p*sinθ*sin(ϕ),  p*cosθ, E1)
    p2 = MVector(-p*sinθ*cos(ϕ), -p*sinθ*sin(ϕ), -p*cosθ, E2)
    return (p1, p2)
end

function four_vectors_in_binary_decay(p0,cosθ,ϕ;
        m1sq::Real=error("give mass 1 squared"),
        m2sq::Real=error("give mass 2 squared"))
    psq = sum(abs2,p0[1:3])
    m0sq = p0[4]^2-psq
    #
    γ = p0[4]/sqrt(m0sq)
    cosθ0 = psq≈0.0 ? 1.0 : p0[3]/sqrt(psq)
    ϕ0 = atan(p0[2],p0[1])
    #
    p1, p2 = four_vectors_in_binary_decay(cosθ,ϕ; m1sq=m1sq, m2sq=m2sq, m0sq=m0sq)
    boostz!(p1,γ); roty_cos!(p1,cosθ0); rotz!(p1, ϕ0)
    boostz!(p2,γ); roty_cos!(p2,cosθ0); rotz!(p2, ϕ0)
    #
    return (p1, p2)
end

invmasssq(p) = p[4]^2-sum(abs2, p[1:3])

function rotz!(p,θ)
    c, s = cos(θ), sin(θ)
    p[1], p[2] = [c -s; s c]*[p[1], p[2]]
    return
end
function roty!(p,θ)
    c, s = cos(θ), sin(θ)
    p[1], p[3] = [c s; -s c]*[p[1], p[3]]
    return
end
function roty_cos!(p,cosθ)
    c, s = cosθ, sqrt(1-cosθ^2)
    p[1], p[3] = [c s; -s c]*[p[1], p[3]]
    return
end
function roty_cos_inv!(p,cosθ)
    c, s = cosθ, sqrt(1-cosθ^2)
    p[1], p[3] = [c -s; s c]*[p[1], p[3]]
    return
end
function boostz!(p,γ)
    γ, βγ = γ*sign(γ), sqrt(γ^2-1)*sign(γ)
    p[3], p[4] = [γ βγ; βγ γ]*[p[3], p[4]]
    return
end

max_stars_repo_name: mmikhasenko/ThreeBodyDecay

max_issues_repo_name: mmikhasenko/ThreeBodyDecay

max_forks_repo_name: mmikhasenko/ThreeBodyDecay


text:
module MadNLPHSL

import Libdl: dlopen, RTLD_DEEPBIND
import MadNLP: @kwdef, Logger, @debug, @warn, @error,
    AbstractOptions, AbstractLinearSolver, set_options!, SparseMatrixCSC, SubVector, StrideOneVector,
    SymbolicException,FactorizationException,SolveException,InertiaException,
    introduce, factorize!, solve!, improve!, is_inertia, inertia, findIJ, nnz,
    get_tril_to_full, transfer!

include(joinpath("..","deps","deps.jl"))

if @isdefined(libhsl)
    include("mc68.jl")
    include("ma27.jl")
    include("ma57.jl")
    include("ma77.jl")
    include("ma86.jl")
    include("ma97.jl")
    export MadNLPMa27, MadNLPMa57, MadNLPMa77, MadNLPMa86, MadNLPMa97
end

function __init__()
    check_deps()
    try
        @isdefined(libhsl) && dlopen(libhsl,RTLD_DEEPBIND)
    catch e
        println("HSL shared library cannot be loaded")
    end
end

end # module