text:
"""
    PayloadSerde

Interface to implement for creating a new compression scheme for the payload.
"""
abstract type PayloadSerde end

"""
    serialize_payload(payload_serde::PayloadSerde, weights::Vector{Float32})

Serialize `weights` by applying the transformation defined by the
`payload_serde`.

!!! info

    This method is an interface to implement when creating a new `PayloadSerde`.
"""
function serialize_payload(::PayloadSerde, ::Vector{Float32})::Vector{UInt8}
    # Nothing - this is an interface to implement...
end

"""
    deserialize_payload(payload_serde::PayloadSerde, data::Vector{UInt8}, from::String)

Deserialize `data` by applying the inverse transformation of the serialization
process. `from` is the URL from which the data are coming, which may be an empty
string if unused by the `payload_serde`.

!!! info

    This method is an interface to implement when creating a new `PayloadSerde`.
"""
function deserialize_payload(::PayloadSerde, ::Vector{UInt8}, ::String)::Vector{Float32}
    # Nothing - this is an interface to implement...
end


text:
module Matching
using  ..Helper
export matcher, compare_trees

include("Environment.jl")
include("Function.jl")
include("Comparison.jl")

import .Function: matcher

end

ext: jl


text:
using ToyFHE
using OffsetArrays
using Test

params = BFVParams(
    7, # plaintext modulus
    ; eval_mult_count = 3
)

plain = zero(plaintext_space(params))
plain[0] = 2

kp1 = keygen(params)
ek = keygen(EvalMultKey, kp1.priv)

c1 = encrypt(kp1, plain)
b1 = invariant_noise_budget(kp1.priv, c1)
c1squared = c1*c1
b2 = invariant_noise_budget(kp1.priv, c1squared)
@test b2 < b1
cswitch1 = keyswitch(ek, c1squared)
bswitch1 = invariant_noise_budget(kp1.priv, cswitch1)
cswitchmul = cswitch1*c1
bswitchmul = invariant_noise_budget(kp1.priv, cswitchmul)
@test bswitchmul < bswitch1 < b1
cswitch2 = keyswitch(ek, cswitchmul)
bswitch2 = invariant_noise_budget(kp1.priv, cswitch2)
cswitchmul2 = cswitch2*c1
bswitchmul2 = invariant_noise_budget(kp1.priv, cswitchmul2)
@test bswitchmul2 < bswitch2 < bswitch1

# This is mostly a test that our heuristics are doing a decent job of chosing
# parameters.
@test 1 < bswitchmul2 < 10


text:
#################
# GET QUEUE URL #
#################

function get_scatter_queue(job_id::JobId=get_job_id())
    return sqs_get_queue_with_retries(
        get_aws_config(),
        string("banyan_", job_id, "_scatter.fifo"),
    )
end

function get_gather_queue(job_id::JobId=get_job_id())
    return sqs_get_queue_with_retries(
        get_aws_config(),
        string("banyan_", job_id, "_gather.fifo"),
    )
end

function get_execution_queue(job_id::JobId=get_job_id())
    return sqs_get_queue_with_retries(
        get_aws_config(),
        string("banyan_", job_id, "_execution.fifo"),
    )
end

###################
# RECEIVE MESSAGE #
###################

function get_next_message(queue, p=nothing; delete = true)
    m = sqs_receive_message(queue)
    while (isnothing(m))
        m = sqs_receive_message(queue)
        # @debug "Waiting for message from SQS"
        if !isnothing(p)
            next!(p)
        end
    end
    if delete
        sqs_delete_message(queue, m)
    end
    return m[:message]
end

function receive_next_message(queue_name, p=nothing)
    content = get_next_message(queue_name, p)
    if startswith(content, "JOB_READY")
        response = Dict{String,Any}(
            "kind" => "JOB_READY"
        )
    elseif startswith(content, "EVALUATION_END")
        # @debug "Received evaluation end"
        response = Dict{String,Any}(
            "kind" => "EVALUATION_END",
            "end" => endswith(content, "MESSAGE_END")
        )
        # Print out logs that were outputed by job on cluster. Will be empty if
        # `print_logs=false` for the job. Remove "EVALUATION_END" at start and
        #  chop off "MESSAGE_END" at the end
        if !isnothing(p) && !p.done
            finish!(p)
        end
        tail = endswith(content, "MESSAGE_END") ? 11 : 0
        println(chop(content, head=14, tail=tail))
        response
    elseif startswith(content, "JOB_FAILURE")
        if !isnothing(p) && !p.done
            finish!(p, spinner='✗')
        end
        # @debug "Job failed"
        # Print job logs. Will be empty if `print_logs=false` for the job. Remove
        # "JOB_FAILURE" and "JOB_END" from the message content. Note that logs
        # are streamed in multiple parts, due to SQS message limits.
        tail = endswith(content, "MESSAGE_END") ? 11 : 0
        println(chop(content, head=11, tail=tail))
        # Destroy job when last part of log is received.
        if endswith(content, "MESSAGE_END")
            # We have to destroy the job here because we could be receiving
            # this message as a result of the executor actually crashing. So a
            # new job entirely will have to be launched. Fortunately, the
            # provisioned nodes should stick around for a bit so it should
            # only be a couple of minutes before the job is back up and
            # running.
            destroy_job(failed=true) # This will reset the `current_job_id` and delete from `jobs`
            error("Job failed; see preceding output")
        end
        Dict{String,Any}("kind" => "JOB_FAILURE")
    else
        # @debug "Received scatter or gather request"
        JSON.parse(content)
    end
end

# Used by Banyan/src/pfs.jl, intended to be called from the executor
function receive_from_client(value_id)
    # Send scatter message to client
    send_message(
        get_gather_queue(),
        JSON.json(Dict("kind" => "SCATTER_REQUEST", "value_id" => value_id))
    )
    # Receive response from client
    m = JSON.parse(get_next_message(get_scatter_queue()))
    v = from_jl_value_contents(m["contents"])
    v
end

################
# SEND MESSAGE #
################

function send_message(queue_name, message)
    sqs_send_message(
        queue_name,
        message,
        (:MessageGroupId, "1"),
        (:MessageDeduplicationId, generate_message_id()),
    )
end

function send_to_client(value_id, value)
    send_message(
        get_gather_queue(),
        JSON.json(
            Dict(
                "kind" => "GATHER",
                "value_id" => value_id,
                "contents" => to_jl_value_contents(value)
            )
        )
    )
end

###########################
# GET MESSAGES FROM QUEUE #
###########################

##########################
# SEND MESSAGES TO QUEUE #
##########################

max_stars_repo_name: banyan-team/banyan-julia

max_issues_repo_name: banyan-team/banyan-julia

max_forks_repo_name: banyan-team/banyan-julia


text:
# This test follows the tutorial at:
# https://dslpitt.org/genie/wiki/SMILE_Tutorial_1:_Creating_a_Bayesian_Network

using Smile

# create a net
net = Network()

# add a success node
success = add_node(net, DSL_CPT, "Success")
somenames = IdArray()
add(somenames, "Success")
add(somenames, "Failure")
set_number_of_outcomes(definition(get_node(net, success)), somenames)

# add a forecast node
forecast = add_node(net, DSL_CPT, "Forecast")
flush(somenames)
add(somenames, "Good")
add(somenames, "Moderate")
add(somenames, "Poor")
set_number_of_outcomes(definition(get_node(net, forecast)), somenames)

# add an arc from "Success" to "Forecast" to represent conditional independence
add_arc(net, success, forecast)

# add the distributions
theprobs = DoubleArray()
set_size(theprobs, 2)
set_at(theprobs, 0, 0.2)
set_at(theprobs, 1, 0.8)
set_definition(definition(get_node(net, success)), theprobs)

thecoordinates = SysCoordinates(definition(get_node(net, forecast)))
set_unchecked_value(thecoordinates, 0.4); next(thecoordinates)
set_unchecked_value(thecoordinates, 0.4); next(thecoordinates)
set_unchecked_value(thecoordinates, 0.2); next(thecoordinates)
set_unchecked_value(thecoordinates, 0.1); next(thecoordinates)
set_unchecked_value(thecoordinates, 0.3); next(thecoordinates)
set_unchecked_value(thecoordinates, 0.6);

# write to file
filename = tempname() * ".xdsl"
write_file(net, filename)
@test isfile(filename)
rm(filename)

filename = tempname() * ".dsl"
write_file(net, filename)
@test isfile(filename)
rm(filename)


text:
macro forwardrule(x)
    :(sqrt(eps(eltype($x))) * max(one(eltype($x)), abs($x)))
end

macro centralrule(x)
    :(cbrt(eps(eltype($x))) * max(one(eltype($x)), abs($x)))
end

macro hessianrule(x)
    :(eps(eltype($x))^(1//4) * max(one(eltype($x)), abs($x)))
end

macro complexrule(x)
    :(eps($x))
end


text:
using Random

function mpi_init()
    !MPI.Initialized() && MPI.Init()
    comm = MPI.COMM_WORLD
    rank = MPI.Comm_rank(comm)
    size = MPI.Comm_size(comm)
    return comm, rank, size
end

"""
    get_indexes(manager, action, size, unsolved=[true, false], original=[true, false], merged=[false])

Retrieve the indexes used to process either instances or decompositions.

# Examples (results depends on the database)
```julia-repl
julia> manager = ExperimentManager()
julia> get_indexes(manager, :solve, 4)
4-element Array{Array{Int64,1},1}:
 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
 [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]
 [23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33]
 [34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44]

julia> get_indexes(manager, :solve, 4, [true])
4-element Array{Array{Int64,1},1}:
 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
 [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]
 [23, 24, 25, 26, 27, 28, 29, 30, 31, 32]
 [33, 34, 35, 36, 37, 38, 39, 40, 41, 42]

julia> get_indexes(manager, :gen, 4)
4-element Array{Array{Int64,1},1}:
 [1, 2, 3, 4, 5, 6]
 [7, 8, 9, 10, 11]
 [12, 13, 14, 15, 16]
 [17, 18, 19, 20, 21]
```
"""
function get_indexes(manager, action, size, unsolved=[true, false], original=[true, false], merged=[true, false], combined=[true, false], select=nothing)
    indexes = []
    if action in [:solve, :merge, :combine]
        decompositions_index = db.getdecompositions_index(manager.decompositions;
                                                          unsolved=unsolved,
                                                          original=original,
                                                          merge=merged,
                                                          combined=combined,
                                                          select=select)
        indexes = split_index(length(decompositions_index), size)
        indexes = [decompositions_index[i[begin]:i[end]] for i in indexes]
    elseif action in [:gen, :genoriginal]
        paths_matpower = db.getmatpowerpath_all(manager.instances)
        indexes = split_index(length(paths_matpower), size)
        indexes = [[i[begin]:i[end]...] for i in indexes]
    end
    return indexes
end

function balance_decomposition(decompositions, chunk_size)
    decompositions = vcat([sort(decomposition[n:n + chunk_size], by=d->d["features"]["chordal_graph"]["nv"]) for n in 1:chunk_size:length(decompositions) - chunk_size]...)
    return decompositions
end

function process_decomposition(manager, decompositions, action, rank=nothing;
                               mpi_infos=nothing,
                               merge_type=:molzahn,
                               merge_treshold_type=:nv,
                               merge_treshold=0.1,
                               combine_type=:intersect,
                               neighbors_g=db.getgraphs_dict())
    for (i, decomposition) in enumerate(decompositions)
        instance_name = decomposition["_id"]["instance_name"]
        oid = decomposition["_id"]["oid"]
        if action == :merge
            @info "[$rank] Merging $(instance_name): $(oid), $(decomposition["nb_added_edges"])"
            genmerge(manager, instance_name, oid;
                     merge_type=merge_type,
                     merge_treshold_type=merge_treshold_type,
                     merge_treshold=merge_treshold)
            @info "[$rank] $i / $(length(decompositions)) $(instance_name): $(oid), $(decomposition["nb_added_edges"]) merged"
        elseif action == :solve
            @info "[$rank] Solving $(instance_name): $(oid), $(decomposition["nb_added_edges"])"
            solve_decomposition(manager, instance_name, oid, true, mpi_infos)
            @info "[$rank] $i / $(length(decompositions)) $(instance_name): $(oid), $(decomposition["nb_added_edges"]) solved"
        elseif action == :combine
            @info "[$rank] Combining $(instance_name): $(oid), $(decomposition["nb_added_edges"])"
            results = gencombine(manager, instance_name, oid; combine_type=combine_type, neighbors_g=neighbors_g)
            if isnothing(results)
                @info "[$rank] $i / $(length(decompositions)) $(instance_name): $(oid), $(decomposition["nb_added_edges"]) could not be combined"
            else
                @info "[$rank] $i / $(length(decompositions)) $(instance_name): $(oid), $(decomposition["nb_added_edges"]) combined"
            end
        end
    end
end

function process_instances(manager, paths_matpower, action, rank=nothing;
                           configfile=nothing,
                           kwargs...)
    for (i, path) in enumerate(paths_matpower)
        @info "[$rank] $i / $(length(paths_matpower)) Generating $(split(basename(path), ".")[1])"
        if action == :genoriginal
            genoriginal(manager, path; kwargs...)
        elseif action == :gen
            # Here we could save some operations by parsing the JSON file only once.
            # Maybe I'll change it later but for the moment I find it clearer to do it this way
            # and I don't think it add too much computation time (I didn't test so it's only guess until I test)
            gendecomposition_json(manager, path, configfile; kwargs...)
        end
        @info "[$rank] $i / $(length(paths_matpower)) $(basename(path)) generated"
    end
end

function mpirun(action;
                port="mongodb://localhost:27017",
                db_name="choleskyexp",
                configfile::String="./config.json",
                mpi_infos::AbstractDict=Dict(),
                unsolved=[true],
                original=[false],
                merged=[true, false],
                combined=[true, false],
                finalize=true,
                merge_treshold_type=:nv,
                merge_treshold=0.1,
                merge_type::Symbol=:molzahn,
                combine_type::Symbol=:intersect,
                seed=1234,
                kwargs...)
    Random.seed!(seed)
    comm, rank, size = mpi_init()
    manager = ExperimentManager(port=port, db=db_name)
    neighbors_g = action == :combine ? db.getgraphs_dict(manager.decompositions) : nothing
    if rank == 0
        @info "Size: $size"
        select = action == :combine ? (x -> x["features"]["solver"]["cholesky_diff"] < 0) : nothing
        indexes = get_indexes(manager, action, size, unsolved, original, merged, combined, select)
        for (i, index) in enumerate(indexes)
            @info "[$rank] Rank $(i - 1): $(index[begin]) -> $(index[end]) | $(length(index))"
            MPI.Isend(index, i - 1, 0, comm)
        end
        if length(indexes) < size
            for i in length(indexes) + 1:size
                MPI.Isend(Vector{Int}(undef,0), i - 1, 0, comm)
            end
        end
    end
    MPI.Barrier(comm)
    @info "[$rank] Retrieving the number of instances/decompositions to process"
    status = MPI.Probe(0, 0, comm)
    count = MPI.Get_count(status, Int)
    if count == 0
        @info "[$rank] Exciting because no instances/decompositions to process"
        MPI.Finalize()
        return
    end
    indexes = Array{Int}(undef, count)
    MPI.Irecv!(indexes, 0, 0, comm)
    @info "[$rank] Processing"
    if action in [:merge, :solve, :combine]
        decompositions = collect(manager.decompositions)[indexes]
        sort!(decompositions, by=d -> d["features"]["chordal_graph"]["nv"])
        process_decomposition(manager, decompositions, action, rank;
                              merge_type=merge_type,
                              merge_treshold_type=merge_treshold_type,
                              merge_treshold=merge_treshold,
                              combine_type=combine_type,
                              neighbors_g=neighbors_g,
                              mpi_infos=mpi_infos)
    elseif action in [:gen, :genoriginal]
        paths_matpower = db.getmatpowerpath_all(manager.instances)[indexes]
        process_instances(manager, paths_matpower, action, rank;
                          configfile,
                          kwargs...)
    end
    @info "[$rank] Process done."
    finalize && MPI.Finalize()
end

max_stars_repo_name: charlyalizadeh/OPFdecgen

max_issues_repo_name: charlyalizadeh/OPFdecgen

max_forks_repo_name: charlyalizadeh/OPFdecgen


text:
module Nanomsg

const libnanomsg = "libnanomsg"

# domains
const AF_SP     = Cint(1)
const AF_SP_RAW = Cint(2)

# protocols
const PAIR = Cint(1 << 4 + 0)
const PUB  = Cint(2 << 4 + 0)
const SUB  = Cint(2 << 4 + 1)
const REQ  = Cint(3 << 4 + 0)
const REP  = Cint(3 << 4 + 1)

const SUB_SUBSCRIBE   = Cint(1)
const SUB_UNSUBSCRIBE = Cint(2)

# send/recv flags
const DONTWAIT = Cint(1)

macro genfunc(ex)
    @assert ex.head == :(::)
    call = ex.args[1]
    @assert call.head == :call
    name = call.args[1]
    rettype = ex.args[2]
    argtypes = Expr(:tuple, [arg.args[2] for arg in call.args[2:end]]...)
    argnames = [arg.args[1] for arg in call.args[2:end]]
    esc(quote
        function $(name)($(argnames...))
            return ccall(($(Expr(:quote, Symbol("nn_", name))), libnanomsg), $(rettype), $(argtypes), $(argnames...))
        end
    end)
end

# Generate C binding functions.
@genfunc socket(domain::Cint, protocol::Cint)::Cint
@genfunc close(s::Cint)::Cint
@genfunc setsockopt(s::Cint, level::Cint, option::Cint, optval::Ptr{Void}, optvallen::Csize_t)::Cint
@genfunc getsockopt(s::Cint, level::Cint, option::Cint, optval::Ptr{Void}, optvallen::Ptr{Csize_t})::Cint
@genfunc bind(s::Cint, addr::Ptr{UInt8})::Cint
@genfunc connect(s::Cint, addr::Ptr{UInt8})::Cint
@genfunc shutdown(s::Cint, how::Cint)::Cint
@genfunc send(s::Cint, buf::Ptr{Void}, len::Csize_t, flags::Cint)::Cint
@genfunc recv(s::Cint, buf::Ptr{Void}, len::Csize_t, flags::Cint)::Cint
@genfunc sendmsg(s::Cint, msghdr::Ptr{Void}, flags::Cint)::Cint
@genfunc recvmsg(s::Cint, msghdr::Ptr{Void}, flags::Cint)::Cint

@genfunc allocmsg(size::Csize_t, type_::Cint)::Ptr{Void}
@genfunc reallocmsg(msg::Ptr{Void}, size::Csize_t)::Ptr{Void}
@genfunc freemsg(msg::Ptr{Void})::Cint

# @genfunc poll(fds::Ptr{Void}, nfds::Cint, timeout::Cint)::Cint

@genfunc errno()::Cint
@genfunc strerror(errnum::Cint)::Ptr{UInt8}

@genfunc get_statistic(s::Cint, stat::Cint)::UInt64

@genfunc device(s1::Cint, s2::Cint)::Cint

@genfunc term()::Void

setsockopt(s, level, option, optval) = setsockopt(s, level, option, pointer(optval), sizeof(optval))
send(s, data, flags=0) = send(s, pointer(data), sizeof(data), flags)
recv(s, data, flags=0) = recv(s, pointer(data), sizeof(data), flags)

end # module


text:
using AccuracyAtTop: find_threshold, find_kth, find_quantile

s = collect(1:10)
y = s .>= 6

@testset "auxiliary threshold functions" begin
    @testset "maximum" begin
        @test find_threshold(All, findmax, y, s) == (10, 10)
        @test find_threshold(Neg, findmax, y, s) == (5, 5)
        @test find_threshold(Pos, findmax, y, s) == (10, 10)

        @test find_threshold(Maximum(All), y, s) == (10, 10)
        @test find_threshold(Maximum(Neg), y, s) == (5, 5)
        @test find_threshold(Maximum(Pos), y, s) == (10, 10)
    end

    @testset "minimum" begin
        @test find_threshold(All, findmin, y, s) == (1, 1)
        @test find_threshold(Neg, findmin, y, s) == (1, 1)
        @test find_threshold(Pos, findmin, y, s) == (6, 6)

        @test find_threshold(Minimum(All), y, s) == (1, 1)
        @test find_threshold(Minimum(Neg), y, s) == (1, 1)
        @test find_threshold(Minimum(Pos), y, s) == (6, 6)
    end

    @testset "$(k) all samples" for k in 1:10
        @test find_kth(s, k, false) == (k, k)
        @test find_kth(s, k, true) == (11 - k, 11 - k)

        @test find_threshold(All, find_kth, y, s, k, false) == (k, k)
        @test find_threshold(All, find_kth, y, s, k, true) == (11 - k, 11 - k)

        @test find_threshold(Kth(k, All; rev = false), y, s) == (k, k)
        @test find_threshold(Kth(k, All; rev = true), y, s) == (11 - k, 11 - k)
    end

    @testset "$(k) pos/neg samples" for k in 1:5
        @test find_threshold(Neg, find_kth, y, s, k, false) == (k, k)
        @test find_threshold(Neg, find_kth, y, s, k, true) == (6 - k, 6 - k)
        @test find_threshold(Kth(k, Neg; rev = false), y, s) == (k, k)
        @test find_threshold(Kth(k, Neg; rev = true), y, s) == (6 - k, 6 - k)

        @test find_threshold(Pos, find_kth, y, s, k, false) == (5 + k, 5 + k)
        @test find_threshold(Pos, find_kth, y, s, k, true) == (11 - k, 11 - k)
        @test find_threshold(Kth(k, Pos; rev = false), y, s) == (5 + k, 5 + k)
        @test find_threshold(Kth(k, Pos; rev = true), y, s) == (11 - k, 11 - k)
    end

    @testset "$(τ)-quantile all samples" for (k, τ) in zip(vcat(1, 1:10), 0:0.1:1)
        @test find_quantile(s, τ, false) == (k, k)
        @test find_threshold(All, find_quantile, y, s, τ, false) == (k, k)
        @test find_threshold(Quantile(τ, All; rev = false), y, s) == (k, k)
    end

    @testset "$(τ)-quantile all samples reversed" for (k, τ) in zip(vcat(10:-1:1, 1), 0:0.1:1)
        @test find_quantile(s, τ, true) == (k, k)
        @test find_threshold(All, find_quantile, y, s, τ, true) == (k, k)
        @test find_threshold(Quantile(τ, All; rev = true), y, s) == (k, k)
    end

    @testset "$(τ)-quantile pos/neg samples" for (k, τ) in zip(vcat(1, 1:5), 0:0.2:1)
        @test find_threshold(Neg, find_quantile, y, s, τ, false) == (k, k)
        @test find_threshold(Quantile(τ, Neg; rev = false), y, s) == (k, k)

        @test find_threshold(Pos, find_quantile, y, s, τ, false) == (k + 5, k + 5)
        @test find_threshold(Quantile(τ, Pos; rev = false), y, s) == (k + 5, k + 5)
    end

    @testset "$(τ)-quantile pos/neg samples reversed" for (k, τ) in zip(vcat(5:-1:1, 1), 0:0.2:1)
        @test find_threshold(Neg, find_quantile, y, s, τ, true) == (k, k)
        @test find_threshold(Quantile(τ, Neg; rev = true), y, s) == (k, k)

        @test find_threshold(Pos, find_quantile, y, s, τ, true) == (k + 5, k + 5)
        @test find_threshold(Quantile(τ, Pos; rev = true), y, s) == (k + 5, k + 5)
    end
end

s = rand(4, 100)
y = rand(Bool, 4, 100)

function test_indices(tp, y, s)
    t, ind = find_threshold(tp, y, s)

    @testset "$tp for $i-th row" for i in 1:size(s, 1)
        @test t[i] == s[i, ind[i]]
    end
end

@testset "thresholds rrule" begin
    @testset "$tp" for tp in Maximum.([All, Neg, Pos])
        test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
        test_indices(tp, y, s)
    end

    @testset "$tp" for tp in Minimum.([All, Neg, Pos])
        test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
        test_indices(tp, y, s)
    end

    @testset "$K" for K in [1, 2, 5, 10, 15]
        @testset "$rev" for rev in [true, false]
            @testset "$tp" for tp in Kth.(K, [All, Neg, Pos]; rev)
                test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
                test_indices(tp, y, s)
            end
        end
    end

    @testset "$τ" for τ in [0.05, 0.1, 0.15, 0.2, 0.5]
        @testset "$rev" for rev in [true, false]
            @testset "$tp" for tp in Quantile.(τ, [All, Neg, Pos]; rev)
                test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
                test_indices(tp, y, s)
            end
        end
    end

    @testset "$τ" for τ in [0.05, 0.1, 0.15, 0.2, 0.5]
        sampler() = τ
        @testset "$rev" for rev in [true, false]
            @testset "$tp" for tp in SampledQuantile.(sampler, [All, Neg, Pos]; rev)
                test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
                test_indices(tp, y, s)
            end
        end
    end
end


text:
@testset "Multivectors" begin
    @test is_homogeneous(mv_1)
    @test grades(mv_1) == [1]
    @test grades(mv_2) == [1,2,3]
    @test all(nonzero_blades(mv_1) .== Any[1v1, 1v2])
    @test all(nonzero_blades(mv_2) .== Any[1v1, 1v12, 1v123])
    @test mv_1[2] == 1
    @test mv_1[3] == 1
    @test mv_2[CartesianIndex(1,2)] == 1
    @test mv_2[CartesianIndex(1,2,3)] == 1
    @test sum(nonzero_blades(mv_1)) == sum(blades(mv_1)) == mv_1
    @test sum(nonzero_blades(mv_2)) == sum(blades(mv_2)) == mv_2

    @test 1v2 + 1v12 == Multivector([0, 0, 1, 0, 1, 0, 0, 0])
    @test 1v1 + mv_1 == Multivector([0, 2, 1, 0, 0, 0, 0, 0])
    @test v1 + v2 == mv_1
end


text:
using Symbolics
using Test

using Random

@variables x y z

@test_throws ArgumentError semipolynomial_form(x,[x],0)

d, r = semipolynomial_form(x, [x], 1)

@test d == Dict(x=>1)
@test r == 0

d, r = semipolynomial_form(x + sin(x) + 1 + y, [x], 1)

@test d == Dict(x=>1)
@test iszero(r - sin(x) - 1 - y)

d, r = semipolynomial_form(x^2+1+y, [x], 1)

@test isempty(d)
@test iszero(r - (x^2+1+y))

d, r = semipolynomial_form((x+2)^12, [x], 1)

@test d == Dict(x => 24576)
@test iszero(r + 24576x - (x+2)^12)

@syms a b c

const components = [2, a, b, c, x, y, z, (1+x), (1+y)^2, z*y, z*x]

function verify(t, d, wrt, nl)
    try
        iszero(t - (isempty(d) ? nl : sum(k*v for (k, v) in d) + nl))
    catch err
        println("""Error verifying semi-pf result for $t
                 wrt = $wrt
                 d = $d
                 nl = $nl""")
        rethrow(err)
    end
end

seed = 0

function trial()
    global seed += 1
    Random.seed!(666+seed)
    n = rand(2:length(components)-1)
    l, r = rand(components, n), vcat(rand(components, n), [1, 0, 1//2, 1.5])

    t = *(map(1:rand(1:3)) do _
        pow = rand([1,1,1,1,1,1,1,1,2,3])
        nterms = rand(2:5)
        sum(rand(l, nterms) .* rand(r, nterms)) ^ pow
    end...)

    @show t

    for _ = 1:4
        wrt = unique(rand([a,b,c,x,y,z], rand(1:6)))
        for deg=Any[1,2,3,4,Inf]
            if deg == 1
                A, c = semilinear_form([t], wrt)
                res = iszero(A*wrt + c - [t])
                if !res
                    println("Semi-linear form is wrong: [$t]  w.r.t   $wrt ")
                    @show A c
                end
            elseif deg == 2
                A,B,v2, c = semiquadratic_form([t], wrt)
                res = iszero(A * wrt + B * v2 + c - [t])
                if !res
                    println("Semi-quadratic form is wrong: $t  w.r.t   $wrt")
                    @show A B v2 c
                end
            else
                if isfinite(deg)
                    d, nl = semipolynomial_form(t, wrt, deg)
                    @test all(x->Symbolics.pdegree(x) <= deg, keys(d))
                    for x in wrt
                        d2, enl = semipolynomial_form(expand(nl), wrt, Inf)
                        elim = all(x->Symbolics.pdegree(x)>deg, keys(d2))
                        if !elim
                            println("Imperfect elimination:")
                            @show t wrt deg nl expand(nl)
                        end
                        @test elim
                    end
                else
                    d, nl = polynomial_coeffs(t, wrt)
                end

                res = verify(t, d, wrt, nl)
                if !res
                    println("""Semi-poly form is wrong: $t  w.r.t   $wrt  deg=$deg
                            Result: $d + $nl""")
                end
            end

            @test res
        end
    end
end

for i=1:20
    @testset "fuzz semi-polynomial-form ($i/20)" begin
        trial()
    end
end


text:
# Autogenerated wrapper script for marble_jll for x86_64-w64-mingw32-cxx03
export astro, marble_qt, marblewidget

using Qt_jll
using Libglvnd_jll
using OpenSSL_jll
JLLWrappers.@generate_wrapper_header("marble")
JLLWrappers.@declare_library_product(astro, "libastro.dll")
JLLWrappers.@declare_executable_product(marble_qt)
JLLWrappers.@declare_library_product(marblewidget, "libmarblewidget-qt5.dll")
function __init__()
    JLLWrappers.@generate_init_header(Qt_jll, Libglvnd_jll, OpenSSL_jll)
    JLLWrappers.@init_library_product(
        astro,
        "bin\\libastro.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_executable_product(
        marble_qt,
        "bin\\marble-qt.exe",
    )

    JLLWrappers.@init_library_product(
        marblewidget,
        "bin\\libmarblewidget-qt5.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
### Struct variable
struct PowerFlowDC
    main::Array{Float64,2}
    flow::Array{Float64,2}
    generation::Array{Float64,2}
end

### DC power flow
@inbounds function rundcpf(system, num, settings, info)
    printstyled("Algorithm: DC power flow\n"; bold = true)

    ########## Pre-processing ##########
    busi, type, Pload, Gshunt, Tini, Pgeni, genOn, resistance, reactance,
    charging, transTap, transShift, branchOn = read_dcsystem(system)

    main = fill(0.0, num.Nbus, 6)
    flow = fill(0.0, num.Nbranch, 5)
    gene = fill(0.0, num.Ngen, 2)
    Va, Pinj, Pbus, Pij, Pji, Pgen = write_dcsystem(system, num, main, flow, gene)

  algtime = @elapsed begin
    ########## Convert in integers ##########
    geni = convert(Array{Int64,1}, system.generator[:, 1])
    fromi = convert(Array{Int64,1}, system.branch[:, 2])
    toi = convert(Array{Int64,1}, system.branch[:, 3])

    ########## Numbering and slack bus ##########
    numbering = false
    bus = collect(1:num.Nbus)
    slack = 0
    for i = 1:num.Nbus
        if bus[i] != busi[i]
            numbering = true
        end
        if type[i] == 3
            slack = bus[i]
        end
    end
    if numbering
        println("The new bus numbering is running.")
    end
    if slack == 0
        slack = 1
        println("The slack bus is not found. Slack bus is the first bus.")
    end

    gen_bus = renumber(geni, busi, bus, numbering)
    from = renumber(fromi, busi, bus, numbering)
    to = renumber(toi, busi, bus, numbering)

    ########## Generation ##########
    for (k, i) in enumerate(gen_bus)
        if genOn[k] == 1
            Pbus[i] += Pgeni[k] / system.basePower
            Pgen[k] = Pgeni[k]
        end
    end

    ########## Ybus matrix ##########
    Ybus, admitance, Pshift = ybusdc(system, num, bus, from, to, branchOn, transTap, reactance, transShift)

    ######### Remove row/column of the slack bus, but keep column rank ##########
    rowSlack = Ybus[:, slack]
    for i in Ybus.colptr[slack]:(Ybus.colptr[slack + 1] - 1)
        row = Ybus.rowval[i]
        Ybus[row, slack] = 0.0; Ybus[slack, row] = 0.0
    end
    Ybus[slack, slack] = 1.0

    ########## Solve the system ##########
    b = Pbus - Pshift - (Pload + Gshunt) ./ system.basePower
    Ti = ls(Ybus, b, settings.solve)

    Ti[slack] = 0.0
    Ti = (pi / 180) * Tini[slack] .+ Ti

    ########## Post-processing ##########
    for i = 1:num.Nbranch
        if branchOn[i] == 1
            Pij[i] = admitance[i] * (Ti[from[i]] - Ti[to[i]] - (pi / 180) * transShift[i]) * system.basePower
            Pji[i] = -Pij[i]
        end
    end

    for i = 1:num.Nbus
        Va[i] = (180 / pi) * Ti[i]
        Pbus[i] = Pbus[i] * system.basePower
        Pinj[i] = Pbus[i] - Pload[i]
        if i == slack
            I = 0.0
            for j in Ybus.colptr[i]:(Ybus.colptr[i + 1] - 1)
                row = Ybus.rowval[j]
                I += rowSlack[row] * Ti[row]
            end
            Pinj[slack] = (I + Pshift[slack]) * system.basePower + Gshunt[slack]
            Pbus[slack] = Pinj[slack] + Pload[slack]
        end
    end

    flag = true
    tempslack = 0
    for i = 1:num.Ngen
        if genOn[i] == 1
            if gen_bus[i] == slack && flag == false
                Pgen[tempslack] -= Pgen[i]
            end
            if gen_bus[i] == slack && flag == true
                Pgen[i] = Pbus[slack]
                tempslack = i
                flag = false
            end
        end
    end
 end # algtime

    ########## Results ##########
    results = PowerFlowDC(main, flow, gene)
    header, group = results_flowdc(system, num, settings, results, slack, algtime)
    if !isempty(settings.save)
        savedata(results, system; info = info, group = group, header = header, path = settings.save)
    end

    return results
end

### Read data
function read_dcsystem(system)
    busi = @view(system.bus[:, 1])
    type = @view(system.bus[:, 2])
    Pload = @view(system.bus[:, 3])
    Gshunt = @view(system.bus[:, 5])
    Tini = @view(system.bus[:, 9])

    Pgeni = @view(system.generator[:, 2])
    genOn = @view(system.generator[:, 8])

    resistance = @view(system.branch[:, 4])
    reactance = @view(system.branch[:, 5])
    charging = @view(system.branch[:, 6])
    transTap = @view(system.branch[:, 10])
    transShift = @view(system.branch[:, 11])
    branchOn = @view(system.branch[:, 12])

    return busi, type, Pload, Gshunt, Tini, Pgeni, genOn,
    resistance, reactance, charging, transTap, transShift, branchOn
end

### Write data
function write_dcsystem(system, num, main, flow, gene)
    for i = 1:num.Nbus
        main[i, 1] = system.bus[i, 1]
        main[i, 5] = system.bus[i, 3]
        main[i, 6] = system.bus[i, 5]
    end
    Va = @view(main[:, 2])
    Pinj = @view(main[:, 3])
    Pbus = @view(main[:, 4])

    for i = 1:num.Nbranch
        flow[i, 1] = system.branch[i, 1]
        flow[i, 2] = system.branch[i, 2]
        flow[i, 3] = system.branch[i, 3]
    end
    Pij = @view(flow[:, 4])
    Pji = @view(flow[:, 5])

    for i = 1:num.Ngen
        gene[i, 1] = system.generator[i, 1]
    end
    Pgen = @view(gene[:,2])

    return Va, Pinj, Pbus, Pij, Pji, Pgen
end

max_stars_repo_name: mcosovic/JuliaGrid

max_issues_repo_name: mcosovic/JuliaGrid

max_forks_repo_name: mcosovic/JuliaGrid


text:
#
# Module to define Jacobian and hessian functions
# all of these should match the implemented conversion
#

#=
# get the partial derivative of the ith input parameter
ith_partial{N}(X::SMatrix{3,N}, i) = SVector(X[1,i], X[2,i], X[3,i])

ith_partial{N}(X::SMatrix{4,N}, i) = SVector(X[1,i], X[2,i], X[3,i], X[4,i])

# reformat to produce the usual 3×3 rotation matrix in this case
ith_partial{N}(X::SMatrix{9,N}, i) = @SMatrix([X[1,i]   X[4,i]   X[7,i];
                                               X[2,i]   X[5,i]   X[8,i];
                                               X[3,i]   X[6,i]   X[9,i]])
=#

#######################################################
# Jacobians for transforming to / from rotation matrices
# (only to rotation matrix is implemented - the other way seems weird)
#######################################################

"""
    jacobian(::Type{output_param}, R::input_param)
Returns the jacobian for transforming from the input rotation parameterization
to the output parameterization, centered at the value of R.

    jacobian(R::rotation_type, X::AbstractVector)
Returns the jacobian for rotating the vector X by R.
"""
function jacobian(::Type{RotMatrix},  q::QuatRotation)
    w = q.q.s
    x = q.q.v1
    y = q.q.v2
    z = q.q.v3

    # let q = s * qhat where qhat is a unit quaternion and  s is a scalar,
    # then R = RotMatrix(q) = RotMatrix(s * qhat) = s * RotMatrix(qhat)

    # get R(q)
    # R = q[:] # FIXME: broken with StaticArrays 0.4.0 due to https://github.com/JuliaArrays/StaticArrays.jl/issues/128
    R = SVector(Tuple(q))

    # solve d(s*R)/dQ (because its easy)
    dsRdQ = @SMatrix [ 2*w   2*x   -2*y   -2*z ;
                       2*z   2*y    2*x    2*w ;
                      -2*y   2*z   -2*w    2*x ;

                      -2*z   2*y    2*x   -2*w ;
                       2*w  -2*x    2*y   -2*z ;
                       2*x   2*w    2*z    2*y ;

                       2*y   2*z    2*w    2*x ;
                      -2*x  -2*w    2*z    2*y ;
                       2*w  -2*x   -2*y    2*z ]

    # get s and dsdQ
    # s = 1
    dsdQ = @SVector [2*w, 2*x, 2*y, 2*z]

    # now R(q) = (s*R) / s
    # so dR/dQ = (s * d(s*R)/dQ - (s*R) * ds/dQ) / s^2
    #          = (d(s*R)/dQ - R*ds/dQ) / s

    # now R(q) = (R(s*q)) / s   for scalar s, because RotMatrix(s * q) = s * RotMatrix(q)
    #
    # so dR/dQ = (dR(s*q)/dQ*s - R(s*q) * ds/dQ) / s^2
    #          = (dR(s*q)/dQ*s - s*R(q) * ds/dQ) / s^2
    #          = (dR(s*q)/dQ   - R(q) * ds/dQ) / s

    jac = dsRdQ - R * transpose(dsdQ)

    # now reformat for output.  TODO: is the the best expression?
    # return Vec{4,Mat{3,3,T}}(ith_partial(jac, 1), ith_partial(jac, 2), ith_partial(jac, 3), ith_partial(jac, 4))

end

# derivatives of R w.r.t a MRP
function jacobian(::Type{RotMatrix},  X::MRP)

    # get the derivatives of the quaternion w.r.t to the mrp
    dQdX = jacobian(QuatRotation,  X)

    # get the derivatives of the rotation matrix w.r.t to the mrp
    dRdQ = jacobian(RotMatrix,  QuatRotation(X))

    # and return
    return dRdQ * dQdX

end

#######################################################
# Jacobians for transforming Quaternion <-> MRP
#
#######################################################

function jacobian(::Type{QuatRotation},  X::MRP)

    # differentiating
    # q = Quaternion((1-alpha2) / (alpha2 + 1), 2*X.x / (alpha2 + 1),   2*X.y  / (alpha2 + 1), 2*X.z / (alpha2 + 1), true)
    # q *= sgn(q.w)

    # derivatives of alpha2
    vspq = SVector(X.x, X.y, X.z)
    alpha2 = X.x * X.x + X.y * X.y + X.z * X.z
    dA2dX = 2 * vspq

    # f = (1-alpha2) / (alpha2 + 1);
    den2 = (alpha2 + 1) * (alpha2 + 1)
    dQ1dX = (-dA2dX * (alpha2 + 1) - dA2dX * (1-alpha2)) / den2

    # do the on diagonal terms
    # f = 2*x / (alpha2 + 1) => g = 2*x, h = alpha2 + 1
    # df / dx = (dg * h - dh * g) / (h^2)
    dQiDi = 2 * ((alpha2 + 1) .- dA2dX .* vspq) / den2

    # do the off entries
    # f = 2x / (alpha2 + 1)
    dQxDi = -2 * vspq[1] * dA2dX / den2
    dQyDi = -2 * vspq[2] * dA2dX / den2
    dQzDi = -2 * vspq[3] * dA2dX / den2

    # assemble it all
    dQdX = @SMatrix [ dQ1dX[1]  dQ1dX[2]  dQ1dX[3] ;
                      dQiDi[1]  dQxDi[2]  dQxDi[3] ;
                      dQyDi[1]  dQiDi[2]  dQyDi[3] ;
                      dQzDi[1]  dQzDi[2]  dQiDi[3] ]

    return dQdX
end

#
# Jacobian converting from a Quaternion to an MRP
#
function jacobian(::Type{MRP}, q::QuatRotation{T}) where T
    w = q.q.s
    x = q.q.v1
    y = q.q.v2
    z = q.q.v3

    den = 1 + w
    scale = 1 / den
    dscaledQw = -(scale * scale)
    dSpqdQw = SVector(x, y, z) * dscaledQw
    J0 = @SMatrix [ dSpqdQw[1]  scale   zero(T) zero(T) ;
                    dSpqdQw[2]  zero(T) scale   zero(T) ;
                    dSpqdQw[3]  zero(T) zero(T) scale ]

    # Need to project out norm component of Quat
    dQ = @SVector [w, x, y, z]
    return J0 - (J0*dQ)*dQ'
end

#######################################################
# Jacobian for rotating a 3 vectors
#
#######################################################

# Note: this is *not* projected into the orthogonal matrix tangent space.
# can do this by projecting each 3×3 matrix (row of 9) by (jacobian[i] - r * jacabian[i]' * r) / 2   (for i = 1:3)
function jacobian(r::RotMatrix{3}, X::AbstractVector)
    @assert length(X) === 3
    T = promote_type(eltype(r), eltype(X))
    Z = zero(T)

    @inbounds return @SMatrix [ X[1] Z    Z     X[2]  Z     Z     X[3]  Z     Z    ;
                                Z    X[1] Z     Z     X[2]  Z     Z     X[3]  Z    ;
                                Z    Z    X[1]  Z     Z     X[2]  Z     Z     X[3] ]
end

@inline function d_cross(u::AbstractVector{T}) where T
    @assert length(u) === 3
    @inbounds return @SMatrix [ zero(T)  -u[3]      u[2]    ;
                                u[3]      zero(T)  -u[1]    ;
                               -u[2]      u[1]      zero(T) ]
end

# TODO: should this be jacobian(:rotate, q,  X)   # or something?
function jacobian(q::QuatRotation, X::AbstractVector)
    w = q.q.s
    x = q.q.v1
    y = q.q.v2
    z = q.q.v3

    @assert length(X) === 3
    T = eltype(q)

    # derivatives ignoring the scaling
    q_im = SVector(2*x, 2*y, 2*z)
    dRdQr  = SVector{3}(2 * w * X + cross(q_im, X))
    dRdQim = -X * q_im' + dot(X, q_im) * one(SMatrix{3,3,T}) + q_im * X' - 2*w * d_cross(X)

    dRdQs = hcat(dRdQr, dRdQim)

    # include normalization (S, s = norm of quaternion)
    dSdQ = SVector(2*w, 2*x, 2*y, 2*z)     # h(x)

    # and finalize with the quotient rule
    Xo = q * X           # N.B. g(x) = s * Xo, with dG/dx = dRdQs
    Xom = Xo * transpose(dSdQ)
    return dRdQs -  Xom
end

function jacobian(spq::MRP, X::AbstractVector)
    dQ = jacobian(QuatRotation, spq)
    q = QuatRotation(spq)
    return jacobian(q, X) * dQ
end

#=
#######################################################
# Hessians for transforming Quaternion <-> SpQuat
#
#######################################################

#
# 2nd derivative of the SpQuat - > Quaternion transformation
#
"""
1) hessian(::Type{output_param}, R::input_param)
Returns the 2nd order partial derivatives for transforming from the input rotation parameterization to the output parameterization, centered at the value of R.
The output is an N vector of DxD matrices, where N and D are the number of parameters in the output and input parameterizations respectively.
2) hessian(R::rotation_type, X::AbstractVector)
Returns the 2nd order partial derivatives for rotating the vector X by R.
The output is an 3 vector of DxD matrices, where D is the number of parameters of the rotation parameterization.
"""
function hessian(::Type{Quaternion},  X::SpQuat)

    # make it match q = Quaternion(X) which puts the return in the domain with q.w >= 0
    q = spquat_to_quat_naive(X)
    s = sgn(q.w)

    # state with the hessian of the first Quaternion term
    # Q[1] = (1-alpha2) / (alpha2 + 1)
    #
    # let A = 1 + alpha2
    # let B = 1 - alpha2
    #
    # dQ1 / dx = (dB/dx * A - B * dA/dx) / (A^2)
    #
    # let C = (dB/dx * A - B * dA/dx)
    # so  dQ1/dx = C / (A^2)
    #
    # then
    # ddQ1/dxdy = (dC/dy*A^2 - C * 2*A * dA/dy) / (A^4)
    #           = (dC/dy*A   - 2 * C * dA/dy) / (A^3)
    #
    # with
    # dC/dy = dB/dxdy * A + dB/dx*dA/dy - dB/dy*dA/dx - B * dA/dxdy

    # calculate C
    vspq = SVector(X)
    mag = sum(X.x .* X.x + X.y .* X.y + X.z .* X.z)
    A = 1 + mag
    B = 1 - mag
    dAdX = 2 * vspq
    dBdX = -2 * vspq
    C = dBdX * A - B * dAdX

    # dC/dxx terms (middle two terms cancel)
    T = eltype(A)
    dCdxx = one(Mat{3,3,T}) * (-2 * (A + B))

    # dC/dxy terms (dB/dxdy and dA/dxdy) are zero
    # put it all together
    A3 = A*A*A
    d2Q1 = s * (dCdxx * A - 2 * C *  dAdX') / A3

    #
    # now "pure" 2nd derivatives
    #
    # Qx = 2*x / (alpha2 + 1) = 2*x/A
    # dQxdx = (2 * A - 2*x*dA/dx) / (A^2)
    #       = C / A^2
    # with C = 2 * A - 2*x*dA/dX
    #
    # dCdx =  2*dA/dX - (2*dA/dX + 2*x*2) = 2*(dA/dX - dAdX - 2*x)
    #      = -4x
    #
    # d2Qxdxx = (dC/dX*A*A - C*2*A*dA/dX) / (A^4)
    #         = (dC/dX*A - 2*C*dA/dX) / (A^3)
    # d2Qxdxx = (-4*A*vspq - 2*(2*A - 2*vspq.*dAdX) .* dAdX) / A3
    #         =  (-4*A*vspq - 4*dAdX .* (A - vspq.*dAdX)) / A3
    d2Qdxx    = s * -4 * (A*vspq + dAdX .* (A - vspq.*dAdX)) / A3

    #
    # now the mixed 2nd derivatives
    #
    # start at
    # Qx    = 2*x / (alpha2 + 1) = 2*x/A
    # dQxdy = (-2*x*dAdy) / (A^2)
    # d2Qxdxy = (-2*dA/dy*A^2 + 2*x*dA/dy*2*A*dA/dx) / (A^4)
    #         = (-2*dA/dy*A + 2*x*dA/dy*2*dA/dx) / (A^3)
    #         = 2*dA/dy*(2*x*dA/dx - A) / (A^3)
    d2Qxdxy    = s * 2 * dAdX * (2*vspq.*dAdX - A)' / A3

    #
    # now the other 2nd derivatives
    #

    # Qx = 2*x / (alpha2 + 1) = 2*x/A
    # dQxdy =  (-2*x*dAdXy) / (A^2)
    # dQxdyz = -(-2*x*dAdXy*2*A*dAdXz) / (A^4)
    #        =  (4*x*dAdXy.*dAdXz) / (A3)
    d2Qxdyz = s * (4 * vspq[1] * dAdX[2] * dAdX[3]) / A3
    d2Qydxz = s * (4 * vspq[2] * dAdX[1] * dAdX[3]) / A3
    d2Qzdxy = s * (4 * vspq[3] * dAdX[1] * dAdX[2]) / A3

    # Qx = 2*x / (alpha2 + 1) = 2*x/A
    # dQxdy =   (-2*x*dAdXy) / (A^2)
    # let C =   -2 * x * dAdXy
    # dCdy  =   -4 * x

    # dQxdyy =  (dCdy * A^2 + 2*x*dAdXy*2*A*dAdXy) / (A^4)
    #        =  (dCdy * A   + 2*x*dAdXy*2*dAdXy) / (A^3)
    #        =  (-4*x * A   + 2*x*dAdXy*2*dAdXy) / (A^3)
    #        =  4*x*(dAdXy^2 - A) / (A^3)
    d2Qxdyy   = s * 4 * vspq * (dAdX.*dAdX - A)' / A3

    #
    # And form it all
    #
    hess = SVector{4, Mat{3,3,T}}(

      Tuple(d2Q1),

      #d2Qx
      Tuple(@fsa([d2Qdxx[1]       d2Qxdxy[2, 1]   d2Qxdxy[3, 1];
                  d2Qxdxy[2, 1]   d2Qxdyy[1,2]    d2Qxdyz;
                  d2Qxdxy[3, 1]   d2Qxdyz         d2Qxdyy[1, 3]])),

      #d2Qy
      Tuple(@fsa([d2Qxdyy[2, 1]   d2Qxdxy[1, 2]   d2Qydxz;
                  d2Qxdxy[1, 2]   d2Qdxx[2]       d2Qxdxy[3, 2];
                  d2Qydxz         d2Qxdxy[3, 2]   d2Qxdyy[2, 3]])),
      #d2Qz
      Tuple(@fsa([d2Qxdyy[3, 1]   d2Qzdxy         d2Qxdxy[1, 3];
                  d2Qzdxy         d2Qxdyy[3, 2]   d2Qxdxy[2, 3];
                  d2Qxdxy[1, 3]   d2Qxdxy[2, 3]   d2Qdxx[3]]))

           )
end

#
# 2nd derivative of the Quaternion -> SpQuat transformation
#
function hessian(::Type{SpQuat},  X::Quaternion)

    # always apply to the quaternion in the domain with the real part >= 0
    s = sgn(X.s)
    Xim = SVector(s * X.v1, s * X.v2, s * X.v3)

    # A = (1 - qs) / (1 + qs)
    # dAdQ1 = 2 / (qs + 1)^2
    # d2AdQ1 = 4 / (qs +1 )^3
    d = (1 + s * X.s); d2 = d*d; d3 = d2*d
    dAdQs = -2 / d2
    d2AdQs = 4 / d3

    # N.B. only partials of the real part are non-zero (i.e. the output is sparse)
    # partials w.r.t the imaginary parts of the quaternion are all the same
    dSdQxQs = dAdQs / 2

    # now do d/dQ1dQ1 terms
    d2Sx_dQsQs = Xim/2  * d2AdQs

    #
    # And build it
    #
    T = typeof(dSdQxQs)
    z = zero(T)

    hess = SVector{3, Mat{4,4,T}}(

        # d2Sx
        Tuple(@fsa([d2Sx_dQsQs[1]  dSdQxQs  z        z;
                    dSdQxQs        z        z        z;
                    z              z        z        z;
                    z              z        z        z])),

        # d2Sy
        Tuple(@fsa([d2Sx_dQsQs[2]  z        dSdQxQs  z;
                    z              z        z        z;
                    dSdQxQs        z        z        z;
                    z              z        z        z])),

        # d2Sz
        Tuple(@fsa([d2Sx_dQsQs[3]  z        z        dSdQxQs;
                    z              z        z        z;
                    z              z        z        z;
                    dSdQxQs        z        z        z]))

        )

end

#######################################################
# Jacobian for rotating one Quaternion by another
#
#######################################################

#
# get the jacobian of quaternion multiplication w.r.t. the right side quaternion
# each column is the ith partial
#
# TODO: should this be jacobian{T}(:*, const_q::Quaternion{T},  variable_q::Quaternion{T})   # or something?
function jacobian(const_q::Quaternion,  variable_q::Quaternion, right_variable::Type{Val{true}}=Val{true})

    @fsa([const_q.w   -const_q.x     -const_q.y    -const_q.z;
          const_q.x   const_q.w      -const_q.z     const_q.y;
          const_q.y   const_q.z      const_q.w     -const_q.x;
          const_q.z  -const_q.y      const_q.x     const_q.w])

end

function jacobian(variable_q::Quaternion, const_q::Quaternion, right_variable::Type{Val{false}})

    @fsa([ const_q.w   -const_q.x     -const_q.y    -const_q.z;
           const_q.x   const_q.w       const_q.z    -const_q.y;
           const_q.y  -const_q.z      const_q.w      const_q.x;
           const_q.z   const_q.y     -const_q.x     const_q.w])

end

=#

#=

#######################################################
# Hessian for rotating a 3 vector
#
#######################################################

function hessian(q::Quaternion, X::AbstractVector)

    s = norm(q)
    T = typeof(s)
    z = zero(T)

    # f(x)
    Xo = rotate(q, X)

    #
    # first derivative without removing the scaling (g(x) = s * f(x))
    #
    q_im = SVector(2*q.x, 2*q.y, 2*q.z)
    dRdQr  = 2 * q.w * X + cross(q_im, X)
    dRdQim = -X * q_im' + dot(X, q_im) * one(Mat{3,3,T}) + q_im * X' - 2* q.w * d_cross(X)

    #
    # second derivative ignoring the scaling
    #
    X2 = 2*X; Xm2 = -X2
    d2dri = d_cross(Xm2)

    # hessian for the X coord
    d2R1s = @fsa([X2[1]        d2dri[1,1]   d2dri[1,2]   d2dri[1,3];
                  d2dri[1,1]   X2[1]        X2[2]        X2[3];
                  d2dri[1,2]   X2[2]        Xm2[1]       z;
                  d2dri[1,3]   X2[3]        z            Xm2[1]      ])

    # hessian for the Y coord
    d2R2s = @fsa([X2[2]        d2dri[2,1]   d2dri[2,2]   d2dri[2,3];
                  d2dri[2,1]   Xm2[2]       X2[1]        z;
                  d2dri[2,2]   X2[1]        X2[2]        X2[3];
                  d2dri[2,3]   z            X2[3]        Xm2[2]      ])

    # hessian for the Z coord
    d2R3s = @fsa([X2[3]        d2dri[3,1]   d2dri[3,2]   d2dri[3,3];
                  d2dri[3,1]   Xm2[3]       z            X2[1];
                  d2dri[3,2]   z            Xm2[3]       X2[2];
                  d2dri[3,3]   X2[1]        X2[2]        X2[3]      ])

    #
    # Now the scaling part, s(x) = sqrt(x' * x)
    #
    dSdQ = SVector(q) / s
    dSdQt = dSdQ'

    s2 = s*s
    d2SdQ = (one(Mat{4,4,T}) * s -  SVector(q) * dSdQ') / (s2)

    # and combine them
    dSdQ_2 = 2 * dSdQ * dSdQt
    gd1 = SVector(dRdQr[1], dRdQim[1, 1], dRdQim[1, 2], dRdQim[1, 3])
    gd2 = SVector(dRdQr[2], dRdQim[2, 1], dRdQim[2, 2], dRdQim[2, 3])
    gd3 = SVector(dRdQr[3], dRdQim[3, 1], dRdQim[3, 2], dRdQim[3, 3])

    hess = SVector{3, Mat{4,4,T}}(

        # d2X[1]
        Tuple((d2R1s - Xo[1]*d2SdQ)/s + (Xo[1] * dSdQ_2 - gd1 * dSdQt - dSdQ * gd1')/s2),

        # d2X[2]
        Tuple((d2R2s - Xo[2]*d2SdQ)/s + (Xo[2] * dSdQ_2 - gd2 * dSdQt - dSdQ * gd2')/s2),

        # d2X[3]
        Tuple((d2R3s - Xo[3]*d2SdQ)/s + (Xo[3] * dSdQ_2 - gd3 * dSdQt - dSdQ * gd3')/s2)

        )
end

function hessian(spq::SpQuat, X::AbstractVector)

    # converting to a Quaternion
    j1 = jacobian(Quaternion, spq)
    j1t = j1'
    h1 = hessian(Quaternion, spq)

    # and rotating
    q = Quaternion(spq)
    j2 = jacobian(q, X)
    h2 = hessian(q, X)

    # build them
    hess = SVector{3, Mat{3, 3, eltype(j2)}}(

        # d2X[1]
        Tuple((j2[1,1] * h1[1] + j2[1,2] * h1[2] + j2[1,3] * h1[3] + j2[1,4] * h1[4]) + j1t * h2[1] * j1),

        # d2X[2]
        Tuple((j2[2,1] * h1[1] + j2[2,2] * h1[2] + j2[2,3] * h1[3] + j2[2,4] * h1[4]) + j1t * h2[2] * j1),

        # d2X[3]
        Tuple((j2[3,1] * h1[1] + j2[3,2] * h1[2] + j2[3,3] * h1[3] + j2[3,4] * h1[4]) + j1t * h2[3] * j1)

        )

end
=#


text:
function cluster(period_labels::AbstractVector,
                 load::AbstractVector{Float64},
                 wind::AbstractVector{Float64},
                 pv::AbstractVector{Float64},
                 n_clusters::Int)

    n_labels = length(period_labels)
    n_load   = length(load)
    n_wind   = length(wind)
    n_pv     = length(pv)

    (n_labels == n_load == n_wind == n_pv) ||
        error("Time series inputs must have matching length")

    freqs = StatsBase.countmap(period_labels)
    periods = collect(keys(freqs))
    counts = values(freqs)
    period_ntimesteps = first(counts)
    all(x -> x == period_ntimesteps, counts) ||
        error("Each period must contain the same number of timesteps")

    n_periods = length(periods)
    period_key = Dict(zip(keys(freqs), 1:n_periods))

    # Create feature vectors for each period
    features = Matrix{Float64}(undef, 3 * period_ntimesteps, n_periods)
    idxs = 1:period_ntimesteps
    for p in keys(freqs)
        p_idx = period_key[p]
        include_timesteps = findall(isequal(p), period_labels)
        features[idxs, p_idx] = load[include_timesteps]
        features[period_ntimesteps .+ idxs, p_idx] = wind[include_timesteps]
        features[2*period_ntimesteps .+ idxs, p_idx] = pv[include_timesteps]
    end

    # Normalize feature vectors
    features .-= Statistics.mean(features, dims=2)
    stds = Statistics.std(features, dims=2)
    for i in 1:size(features, 1)
        stds[i] > 0 && (features[i, :] ./= stds[i])
    end

    # Create period distance matrix
    feature_dists = [LinearAlgebra.norm(ci .- cj, 2)
                     for ci in eachcol(features), cj in eachcol(features)]

    result = Clustering.kmedoids(feature_dists, n_clusters)
    medoid_periods = getindex.(Ref(periods), result.medoids)
    cluster_nperiods = Dict(zip(medoid_periods, result.counts))
    period_clusters = Dict(zip(periods, result.assignments))

    return cluster_nperiods, period_clusters

end

function find_representative_days(simulation_dir::String,
                                  base_dir::String,
                                  n_clusters::Int64)

    representative_days_file = joinpath(base_dir, "representative_days_data_$(n_clusters).csv")

    # Use the same representative days if file already exists, otherwise create new clusters
    if isfile(representative_days_file)
        representative_days_data = read_data(representative_days_file)
        clusters = Dict(representative_days_data[i, "Date"] => representative_days_data[i, "Weight"] for i in 1:DataFrames.nrow(representative_days_data))
    else
        net_load_data = read_data(joinpath(simulation_dir, "timeseries_data_files", "Net Load Data", "load_n_vg_data.csv"))
        labels = [Dates.Date(net_load_data[i, :Year], net_load_data[i, :Month], net_load_data[i, :Day]) for i in 1:DataFrames.nrow(net_load_data)]
        load = vec(sum(Matrix(net_load_data[:, r"load"]), dims=2))
        wind = zeros(length(load))
        pv = zeros(length(load))

        for i in names(net_load_data)
            if occursin("WIND", i) || occursin("WT", i)
                wind += net_load_data[:, i]
            elseif occursin("PV", i) || occursin("PVe", i)
                pv += net_load_data[:, i]
            end
        end
        clusters, allperiods = cluster(labels, load, wind, pv, n_clusters)
        representative_days_data = DataFrames.DataFrame(Date = Dates.Date[], Weight = Int64[])
        for i in keys(clusters)
            push!(representative_days_data, (i, clusters[i]))
        end
        CSV.write(representative_days_file, representative_days_data)
    end
    return clusters
end


text:
"""
	grdtrend(cmd0::String="", arg1=nothing, arg2=nothing; kwargs...)

reads a 2-D grid file and fits a low-order polynomial trend to these data by
[optionally weighted] least-squares.

Full option list at [`grdtrend`]($(GMTdoc)grdtrend.html)

Parameters
----------

- **N** | **model** :: [Type => Str | Number]

    Sets the number of model parameters to fit.
    ($(GMTdoc)grdtrend.html#n)
- **D** | **diff** :: [Type => Str | []]

    Compute the difference (input data - trend). Optionaly provide a file name to save result on disk.
    ($(GMTdoc)grdtrend.html#d)
- **T** | **trend** :: [Type => Str | []]

    Compute the trend surface. Optionaly provide a file name to save result on disk.
    ($(GMTdoc)grdtrend.html#t)
- **W** | **weights** :: [Type => Str]

    If weight.nc exists, it will be read and used to solve a weighted least-squares problem.
    ($(GMTdoc)grdtrend.html#w)
- $(GMT.opt_R)
- $(GMT.opt_V)
"""
function grdtrend(cmd0::String="", arg1=nothing, arg2=nothing; kwargs...)

	length(kwargs) == 0 && return monolitic("grdtrend", cmd0, arg1, arg2)

	d = KW(kwargs)

	cmd, = parse_R("", d)
	cmd = parse_V_params(cmd, d)
	cmd = parse_these_opts(cmd, d, [[:D :diff], [:T :trend]])
	opt_N = add_opt("", "N", d, [:N :model], (n="", n_model="", robust="_+r"), false, true)
	if (opt_N == "")  error("The 'model' parameter is mandatory")  end
	cmd *= opt_N

	cmd, got_fname, arg1 = find_data(d, cmd0, cmd, arg1)
	if (isa(arg1, Array{<:Number}))  arg1 = mat2grid(arg1)  end

	if ((val = find_in_dict(d, [:W :weights])[1]) !== nothing)
		if (isa(val, String) || isa(val, Symbol))
			cmd *= " -W" * arg2str(val)
		else
			if (isa(val, Tuple) && length(val) == 2 && (isa(val, GMTgrid) || isa(val, Array{GMT.GMTgrid,1})))
				val = val[1];	cmd *= "+s"
			end
			cmd, N_used = put_in_slot(cmd, val, 'W', [arg1, arg2])
			(N_used == 1) ? arg1 = val : arg2 = val
		end
	end

	if (occursin("-D", cmd) && occursin("-T", cmd))
		@warn("Usage error, both difference and trend were required. Ignoring the trend request.")
	elseif (!occursin("-D", cmd) && !occursin("-T", cmd))
		cmd *= " -T" 			# No -T -or -D provided so default to -T
	end

	return common_grd(d, "grdtrend " * cmd, arg1, arg2)
end

# ---------------------------------------------------------------------------------------------------
grdtrend(arg1, arg2=nothing, cmd0::String=""; kw...) = grdtrend(cmd0, arg1, arg2; kw...)


text:
#
# Denoise an SPD Example with Cyclic Proximal Point applied to the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data as for the corresponding CP algorithm
#
using Manopt
using Images, CSV, DataFrames, LinearAlgebra, JLD2

#
# Settings
ExportResult = true
ExportOrig = true
ExportResultVideo = false
ExportTable = true
resultsFolder = "src/examples/Total_Variation/S2_TV/"
experimentName = "WhirlCPPA"
if !isdir(resultsFolder)
    mkdir(resultsFolder)
end
#
# Manifold & Data
f = artificialS2WhirlImage(64)
pixelM = Sphere(2);

if ExportOrig
      renderAsymptote(resultsFolder*experimentName*"-orig.asy", asyExportS2Data; data=f) # (5)
end
#
# Parameters
α = 1.5
maxIterations = 4000
#
# Build Problem for L2-TV
M = Power(pixelM,size(f))
d = length(size(f))
iRep = [Integer.(ones(d))...,d]
fidelity(x) = 1/2*distance(M,x,f)^2
Λ(x) = forwardLogs(M,x) # on T_xN
prior(x) = norm(norm.(Ref(pixelM),getValue(repeat(x,iRep...)), getValue(Λ(x)) ), 1)
#
# Setup and Optimize
cost(x) = fidelity(x) + α*prior(x)
proximalMaps = [(λ,x) -> proxDistance(M,λ,f,x,2), (λ,x) -> proxTV(M,α*λ,x,1)]
x0 = f
@time y, yRec = cyclicProximalPoint(M,cost,proximalMaps,x0;
    debug = [:Iteration," | ", DebugProximalParameter()," | ", :Change, " | ", :Cost, "\n",100,:Stop],
    record = [:Iteration, :Iterate, :Cost],
    stoppingCriterion = stopAfterIteration(maxIterations),
    λ = i -> π/(2*i)
)
#
# Results
if ExportResult
  renderAsymptote(resultsFolder*experimentName*"-result-$(maxIterations)-α$(replace(string(α), "." => "-")).asy", asyExportS2Data; data=y, render=4) #(6)
end
if ExportTable
    A = cat( [ y[1] for y in yRec], [y[3] for y in yRec]; dims=2 )
    CSV.write(string(resultsFolder*experimentName*"-Result.csv"),
        DataFrame(A), writeheader=false);
    save(resultsFolder*experimentName*"-CostValue.jld2",
      Dict("compareCostFunctionValue" => last(yRec)[3])
    )
end


text:
#=
composir_iterator:
- Julia version:
- Author: Dan
- Date: 2021-06-15
=#

export
    CompositeIterator, iterate_over_subtype, SystemIterator, ChainIterator, ResidueIterator,
    AtomIterator

abstract type CompositeIterator end
using Base: SizeUnknown
import Base.collect

function iterate_over_subtype(target_type::T) where T <: AbstractComposite
    end

mutable struct SystemIterator <:CompositeIterator
    node::AbstractComposite
end

mutable struct ChainIterator <: CompositeIterator
    node::AbstractComposite
end

mutable struct ResidueIterator <: CompositeIterator
    node::AbstractComposite
end

mutable struct AtomIterator <: CompositeIterator
    node::AbstractComposite
end

const iterator_to_composite = Dict{Type{T} where T<: CompositeIterator,Type{S} where S <: AbstractComposite}(
        SystemIterator => System,
        ChainIterator => Chain,
        ResidueIterator => Residue,
        AtomIterator => Atom
    )

const composite_to_iterator = Dict{Type{T} where T<: AbstractComposite,Type{S} where S <: CompositeIterator}(
        System => SystemIterator,
        Chain => ChainIterator,
        Residue => ResidueIterator,
        Atom => AtomIterator
    )

Base.length(::T) where T <: CompositeIterator = SizeUnknown()
Base.eltype(::T) where T <: CompositeIterator = iterator_to_composite[T]

#usage: for item in ChainIterator(someNode) ... end
"""
    Base.iterate(s :: T) where T <: CompositeIterator
Iterator that returns all Nodes of the respective Type in the entire tree.\\
Preferred functions with enhanced speed: [`collectAtoms`](@ref),[`collectResidues`](@ref),[`collectChains`](@ref),
[`collectSystems`](@ref)
Usage:\\
(`node` is of Type `AbstractComposite` or one of the `KERNEL` types.)
```
for system in SystemIterator(node)  #substiute `SystemIterator` for `ChainIterator`,`ResidueIterator`,`AtomIterator`
    #...
end
```
"""
Base.iterate(s :: T) where T <: CompositeIterator = begin

    while s.node.parent_ !== nothing
        s.node = s.node.parent_
    end
    stack = T[s]
    travel = T[]

    cur_type = typeof(s.node)
    skip_types = nothing

    if T == SystemIterator
        if cur_type != Composite
            return (s.node,(T[],s))
        else
            return (s.node.first_child_,(T[],T(s.node.first_child_)))
        end
    elseif T == ChainIterator
        skip_types = SA[System]
        stop_types = SA[Residue,Atom]
    elseif T == ResidueIterator
        skip_types = SA[System, Chain]
        stop_types = SA[Atom]
    elseif T == AtomIterator
        skip_types = SA[System,Chain, Residue]
        stop_types = SA[]
    end

    while !isempty(stack)
        cur = pop!(stack)
        #see if first is here
        if cur.node.first_child_ !== nothing
            if !(typeof(cur.node.first_child_) in skip_types)
                if typeof(cur.node.first_child_) in stop_types
                    nothing
                else
                    push!(travel, T(cur.node.first_child_))
                end
            else
                push!(stack, T(cur.node.first_child_))
            end
        end
        #see if nex tis here
        if cur.node.next_ !== nothing
            push!(stack, T(cur.node.next_))
        end
    end
    isempty(travel) && return nothing
    cp = last(travel)
    pop!(travel)
    return (cp.node,(travel,cp))
end

function Base.iterate(_ :: T, state::Tuple{Vector{T},Union{Nothing,T}}) where T <: CompositeIterator
    stack = state[1]

    if state[2] !== nothing
        cp = state[2]
        if cp.node.first_child_ !== nothing &&
            typeof(cp.node.first_child_) == iterator_to_composite[T]
            push!(stack, T(cp.node.first_child_))
        end
        if cp.node.next_ != nothing
            push!(stack,T(cp.node.next_))
        end

    end
    stack = state[1]
    isempty(stack) && return nothing
    cp = last(stack)
    if cp.node.first_child_ !== nothing &&
            typeof(cp.node.first_child_) == iterator_to_composite[T]
        stack[end] = T(cp.node.first_child_)
        return (cp.node,(stack,nothing))
    end
    if cp.node.next_ !== nothing
        stack[end] = T(cp.node.next_)
        return (cp.node, (stack,nothing))
    else # backtracks
        while !isempty(stack)
            cp = last(stack)
            if cp.node.next_ !== nothing
                stack[end] = cp
                return (cp.node, (stack,nothing))
            end
            pop!(stack)
            return (cp.node,(stack,nothing))
        end
        return nothing
    end

end

Base.collect(iter::T) where T<:CompositeIterator = begin
    recursive_collect(iter.node, iterator_to_composite[T])
end

"""
    collectAtoms(node::AbstractComposite)
Collects all the `Atom`s of the subtree rooted in `node`.
"""
collectAtoms(node::AbstractComposite) = begin
    recursive_collect(node,Atom)
end

"""
    collectResidues(node::AbstractComposite)
Collects all the `Residue`s of the subtree rooted in `node`.
"""
collectResidues(node::AbstractComposite) = begin
    recursive_collect(node,Residue)
end

"""
    collectChains(node::AbstractComposite)
Collects all the `Chain`s of the subtree rooted in `node`.
"""
collectChains(node::AbstractComposite) = begin
    recursive_collect(node,Chain)
end

"""
    collectSystems(node::AbstractComposite)
Collects all the `System`s of the subtree rooted in `node`.
"""
collectSystems(node::AbstractComposite) = begin
    recursive_collect(node,System)
end


text:
export to_positive_definite, to_positive_definite_softplus, to_positive_definite_square,
    upper_triangular, to_positive_definite_softplus_inv

# Copy-Paste some util from Stheno.jl
diag_At_A(A::AbstractVecOrMat) = vec(sum(abs2.(A); dims=1))
diag_Xt_invA_X(A::Cholesky, X::AbstractVecOrMat) = diag_At_A(A.U' \ X)

"""
    to_positive_definite(U::UpperTriangular)

Exponentiate the diagonal of `U` so that `U'U` is positive definite.
"""
function to_positive_definite(U::UpperTriangular)
    new_data = copy(U.data)
    new_data[diagind(new_data)] .= exp.(diag(new_data)) .+ 1e-12
    return UpperTriangular(new_data)
end

# Custom adjoint necessary to handle mutation.
ZygoteRules.@adjoint function to_positive_definite(U::UpperTriangular)
    out = to_positive_definite(U)
    return out, function(Δ::AbstractMatrix)
        out_back = copy(Δ)
        out_back[diagind(out_back)] .= diag(Δ) .* diag(out)
        return (out_back,)
    end
end

"""
    to_positive_definite_softplus(U::UpperTriangular)

Exponentiate the diagonal of `U` so that `U'U` is positive definite.
"""
function to_positive_definite_softplus(U::UpperTriangular)
    new_data = copy(U.data)
    new_data[diagind(new_data)] .= softplus.(diag(new_data)) .+ 1e-12
    return UpperTriangular(new_data)
end

ZygoteRules.@adjoint function to_positive_definite_softplus(U::UpperTriangular)
    out = to_positive_definite_softplus(U)
    return out, function(Δ::AbstractMatrix)
        out_back = copy(Δ)
        out_back[diagind(out_back)] .= diag(Δ) .* logistic.(diag(U))
        return (out_back,)
    end
end

"""
    to_positive_definite_softplus_inv(U::UpperTriangular)

Inverse of to_positive_definite_softplus.
"""
function to_positive_definite_softplus_inv(U::UpperTriangular)
    new_data = copy(U.data)
    new_data[diagind(new_data)] .= invsoftplus.(diag(new_data) .- 1e-12)
    return UpperTriangular(new_data)
end

function frule(::typeof(to_positive_definite_softplus_inv), U::UpperTriangular)
    V = to_positive_definite_softplus_inv(U)
    function to_positive_definite_softplus_inv_pushforward(Δself, ΔU::AbstractMatrix)
        return to_positive_definite_softplus_inv_pushforward(Δself, UpperTriangular(ΔU))
    end
    function to_positive_definite_softplus_inv_pushforward(Δself, ΔU::UpperTriangular)
        new_data = copy(ΔU.data)
        new_data[diagind(new_data)] .= diag(new_data) .* exp.(diag(U) .- diag(V))
        return UpperTriangular(new_data)
    end
    return V, to_positive_definite_softplus_inv_pushforward
end

"""
    to_positive_definite_square(U::UpperTriangular)

Square the diagonal of `U` so that `U'U` is positive definite.
"""
function to_positive_definite_square(U::UpperTriangular)
    new_data = copy(U.data)
    new_data[diagind(new_data)] .= abs2.(diag(new_data)) .+ 1e-12
    return UpperTriangular(new_data)
end

ZygoteRules.@adjoint function to_positive_definite_square(U::UpperTriangular)
    out = to_positive_definite_square(U)
    return out, function(Δ::AbstractMatrix)
        out_back = copy(Δ)
        out_back[diagind(out_back)] .= 2 .* diag(Δ) .* diag(U)
        return (out_back,)
    end
end

"""
    upper_triangular(A::Matrix)

Exactly the same as the regular UpperTriangular function, except the adjoint returns a
dense matrix, rather than another UpperTriangular matrix. This is helpful when working with
(at least) the AMSGrad implementation in Flux.
"""
upper_triangular(A::Matrix) = UpperTriangular(A)

ZygoteRules.@adjoint function upper_triangular(A::Matrix)
    return upper_triangular(A), Δ->(collect(UpperTriangular(Δ)),)
end

#
# frule for cholesky
#

chol(A::Symmetric) = cholesky(A).U

function frule(::typeof(chol), A::Symmetric{T, Matrix{T}} where {T<:Real})
    U = chol(A)
    function chol_pushforward(Δself, ΔA::Symmetric)
        T = U' \ (U' \ ΔA')'
        T = UpperTriangular(T)
        T[diagind(T)] ./= 2
        return T * U
    end
    return U, chol_pushforward
end

#
# rrule for cholesky inv
#

ZygoteRules.@adjoint function inv(C::Cholesky{<:Real})
    Cinv = inv(C)
    return Cinv, function(Δ)
        return ((factors=-(C.U' \ (Δ + Δ')) * Cinv,),)
    end
end

# #
# # Helper functionality to make reverse-mode compute the correct thing...
# #

# chol_inv(S::Symmetric{<:Real}) = inv(cholesky(S))

# # Specialised rrule for natural_to_standard. This does things in a slightly strange way.
# ZygoteRules.@adjoint function chol_inv(S::Symmetric{<:Real})
#     Sinv = inv(cholesky(S))
#     return Sinv, Δ::AbstractMatrix{<:Real} -> (-Sinv'Δ * Sinv',)
# end

#
# cholesky specifically for Symmetric matrices.
#

# Implementation due to Seeger, Matthias, et al. "Auto-differentiating linear algebra."
ZygoteRules.@adjoint function cholesky(Σ::Symmetric{T, <:StridedMatrix{T}} where {T<:Real})
  C = cholesky(Σ)
  return C, function(Δ::NamedTuple)
    U, Ū = C.U, Δ.factors
    ΔΣ = Ū * U'
    ΔΣ = LinearAlgebra.copytri!(ΔΣ, 'U')
    ΔΣ = ldiv!(U, ΔΣ)
    BLAS.trsm!('R', 'U', 'T', 'N', one(eltype(Σ)), U.data, ΔΣ)
    ΔΣ ./= 2
    return (Symmetric(ΔΣ),)
  end
end


text:
"""Evaluates a set of SCM configurations for a single parameter vector."""

using ArgParse
using CalibrateEDMF
using CalibrateEDMF.Pipeline
src_dir = dirname(pathof(CalibrateEDMF))
using CalibrateEDMF.HelperFuncs
using JLD2

# Read iteration number of ensemble to be recovered
s = ArgParseSettings()
@add_arg_table s begin
    "--version"
    help = "Calibration process number"
    arg_type = Int
    "--job_dir"
    help = "Job output directory"
    arg_type = String
    default = "output"
    "--mode"
    help = "Forward model evaluation mode: `train` or `validation`"
    arg_type = String
    default = "train"
end
parsed_args = parse_args(ARGS, s)
outdir_path = parsed_args["job_dir"]
include(joinpath(outdir_path, "config.jl"))

config = get_config()
version = parsed_args["version"]
mode = parsed_args["mode"]
versioned_model_eval(version, outdir_path, mode, config)