text:
# Executable

using CSV, DataFrames, Plots
using AutomatonEdification.Clusters, AutomatonEdification.CentroidLinearClassifierImpl, AutomatonEdification.LinearClassifiers

SYN_CENTER_RADIUS = 20
SYN_CENTER_PERTURB_RADIUS = 10.0
SYN_VECTOR_NDIMS = 2
NUM_POSITIVE_SYN_POINTS = 1000
NUM_NEGATIVE_SYN_POINTS = 1000

function write_clusters_to_csv(ndims::Integer,
                               num_positive_points::Integer,
                               num_negative_points::Integer)
    # Generate two clusters of random data points, where each coordinate is an
    # integer in the range [-SYN_CENTER_RADIUS, SYN_CENTER_RADIUS].
    positive_center = rand(-SYN_CENTER_RADIUS:SYN_CENTER_RADIUS, ndims)
    negative_center = rand(-SYN_CENTER_RADIUS:SYN_CENTER_RADIUS, ndims)

    # Generate clusters of points uniformly distributed around each center.
    positive_cluster =
        Clusters.create_uniform_cluster_nsphere(positive_center,
                                               SYN_CENTER_PERTURB_RADIUS,
                                               num_positive_points)
    negative_cluster =
        Clusters.create_uniform_cluster_nsphere(negative_center,
                                               SYN_CENTER_PERTURB_RADIUS,
                                               num_negative_points)

    CSV.write("positive_data.csv", DataFrame(positive_cluster, :auto),
              writeheader=false)
    CSV.write("negative_data.csv", DataFrame(negative_cluster, :auto),
              writeheader=false)

    return positive_center, negative_center
end

function read_clusters_from_csv()
    positive_data =
        Matrix(CSV.read("positive_data.csv", header=false, DataFrame))
    negative_data =
        Matrix(CSV.read("negative_data.csv", header=false, DataFrame))

    return positive_data, negative_data
end

positive_center, negative_center =
    write_clusters_to_csv(SYN_VECTOR_NDIMS,
                          NUM_POSITIVE_SYN_POINTS, NUM_NEGATIVE_SYN_POINTS)
positive_data, negative_data = read_clusters_from_csv()

positive_centroid = Clusters.compute_centroid(positive_data)
negative_centroid = Clusters.compute_centroid(negative_data)

# Train the model.
model = CentroidLinearClassifierImpl.train_model(positive_data, negative_data)

println("Classifier trained!")
println("`positive_center`, `negative_center` are the original unperturbed centers.")
println("`positive_centroid`, `negative_centroid` are the respective perturbed centroids.")
println("`positive_data`, `negative_data` contain the full datasets.")
println("`model` has a function `classify` that can be applied to a vector to classify it.")
println("The vector must be of dimension: ", SYN_VECTOR_NDIMS)
println()
println("positive_centroid: ", positive_centroid)
println("negative_centroid: ", negative_centroid)

# Please don't try to plot things with weird dimensionality.
if (SYN_VECTOR_NDIMS == 2 || SYN_VECTOR_NDIMS == 3)
    # arr[:,k] selects everything along the kth column, i.e. a vector containing the
    # kth coordinate of all elements.
    plt = plot([positive_data[:,n] for n in 1:size(positive_data, 2)]...,
            label = "positive_data",
            seriestype = :scatter,
            showaxis = :show,
            aspect_ratio = :equal)
    plot!(plt, [negative_data[:,n] for n in 1:size(negative_data, 2)]...,
            label = "negative_data",
            seriestype = :scatter,
            showaxis = :show,
            aspect_ratio = :equal)

    # 3D plotting is kind of weird. Only try this for 2D.
    if (SYN_VECTOR_NDIMS == 2)
        boundary = SYN_CENTER_RADIUS + SYN_CENTER_PERTURB_RADIUS
        xns = [-boundary:1:boundary for n in 1:SYN_VECTOR_NDIMS]
        contour!(xns...,
                 (xns...) -> LinearClassifiers.classify([xns...], model),
                 levels=[0])
    end
end

display(plt)

max_stars_repo_name: brianxie/automaton-edification

max_issues_repo_name: brianxie/automaton-edification

max_forks_repo_name: brianxie/automaton-edification


text:
mutable struct SparseSymmetricBlockMatrix{T<:Number}
    blocks::Dict{Any,SparseSymmetricMatrix{T}}
end

SparseSymmetricBlockMatrix(T::Type) = SparseSymmetricBlockMatrix(Dict{Any,SparseSymmetricMatrix{T}}())

SparseSymmetricBlockMatrix() = SparseSymmetricBlockMatrix(Float64)

blocks(m::SparseSymmetricBlockMatrix) = m.blocks

function setindex!{T<:Number}(bm::SparseSymmetricBlockMatrix{T}, v::T, bi, i, j)
    if haskey(blocks(bm), bi)
        blocks(bm)[bi][i, j] = v
    else
        m = SparseSymmetricMatrix(T)
        m[i, j] = v
        blocks(bm)[bi] = m
    end
end

function setindex!{T<:Number}(bm::SparseSymmetricBlockMatrix{T}, m::SparseSymmetricMatrix{T}, bi)
    blocks(bm)[bi] = m
end

function getindex{T<:Number}(bm::SparseSymmetricBlockMatrix{T}, bi, i, j)
    get(blocks(bm), bi, SparseSymmetricMatrix(T))[i, j]
end

getindex{T}(bm::SparseSymmetricBlockMatrix{T}, bi) = get(blocks(bm), bi, SparseSymmetricMatrix(T))

start(m::SparseSymmetricBlockMatrix) = start(blocks(m))
next(m::SparseSymmetricBlockMatrix, state) = next(blocks(m), state)
done(m::SparseSymmetricBlockMatrix, state) = done(blocks(m), state)

copy(m::SparseSymmetricBlockMatrix) = SparseSymmetricBlockMatrix(copy(blocks(m)))

function *{T}(x::T, m::SparseSymmetricBlockMatrix{T})
    p = SparseSymmetricBlockMatrix(T)
    for k in keys(blocks(m))
        p[k] = x * m[k]
    end
    p
end

function show{T<:Number}(io::IO, m::SparseSymmetricBlockMatrix{T})
    println(io, typeof(m))
    for (blockindex, matrix) in blocks(m)
        for ((i, j), v) in matrix.entries
            println(io, "$blockindex\t$i\t$j\t$v")
        end
    end
end

==(A::SparseSymmetricBlockMatrix, B::SparseSymmetricBlockMatrix) = blocks(A) == blocks(B)

function isapprox(A::SparseSymmetricBlockMatrix, B::SparseSymmetricBlockMatrix)
    for (bi, bm) in A
        if !isapprox(bm, B[bi])
            return false
        end
    end
    for (bi, bm) in B
        if !isapprox(bm, A[bi])
            return false
        end
    end
    true
end

ext: jl


text:
module MPRebinSpectra

using DataFrames, DataFramesMeta, Random, Distributions, StatsPlots, ColorSchemes

include("Line.jl")
export  get_line_params,
        get_line_point,
        get_integral_linear,
        plot_lines

include("Volume.jl")
export  get_area,
        get_row_volume,
        get_total_volume,
        get_segment_volume,
        get_volume_matrix

include("Rebin2D.jl")
export  rebin,
        rebin2D,
        normalize2D!,
        get_cdf,
        get_cdf!

end # module MPRebinSpectra


text:
# Sampler and sampler state types

abstract type AbstractSampler end
abstract type MetropolisSampler <: AbstractSampler end
abstract type NonMetropolisSampler <: AbstractSampler end

abstract type ZerothOrderMetropolisSampler <: MetropolisSampler end
abstract type FirstOrderMetropolisSampler <: MetropolisSampler end
abstract type SecondOrderMetropolisSampler <: MetropolisSampler end

abstract type ZerothOrderNonMetropolisSampler <: NonMetropolisSampler end
abstract type FirstOrderNonMetropolisSampler <: NonMetropolisSampler end
abstract type SecondOrderNonMetropolisSampler <: NonMetropolisSampler end

abstract type AbstractSamplerState end
abstract type MetropolisSamplerState <: AbstractSamplerState end
abstract type NonMetropolisSamplerState <: AbstractSamplerState end

abstract type ZerothOrderMetropolisSamplerState <: MetropolisSamplerState end
abstract type FirstOrderMetropolisSamplerState <: MetropolisSamplerState end
abstract type SecondOrderMetropolisSamplerState <: MetropolisSamplerState end

abstract type ZerothOrderNonMetropolisSamplerState <: NonMetropolisSamplerState end
abstract type FirstOrderNonMetropolisSamplerState <: NonMetropolisSamplerState end
abstract type SecondOrderNonMetropolisSamplerState <: NonMetropolisSamplerState end

abstract type AbstractRecycler end

max_stars_repo_name: gideonsimpson/JuBasicMD

max_issues_repo_name: gideonsimpson/JuBasicMD

max_forks_repo_name: gideonsimpson/JuBasicMD


text:
function score(str::AbstractString)
    rank = Dict('a'=>1, 'e'=>1, 'i'=>1, 'o'=>1, 'u'=>1, 'l'=>1,
                'n'=>1, 'r'=>1, 's'=>1, 't'=>1, 'd'=>2, 'g'=>2,
                'b'=>3, 'c'=>3, 'm'=>3, 'p'=>3, 'f'=>4, 'h'=>4,
                'v'=>4, 'w'=>4, 'y'=>4, 'k'=>5, 'j'=>8, 'x'=>8,
                'q'=>10, 'z'=>10)
    length(str) == 0 && return 0
    mapreduce(x->get(rank, x, 0), +, lowercase(str))
end


text:
abstract type MosiModel{T <: MosiVector} end

name(::T) where T <: MosiModel = string(T)
vectype(::Type{<:MosiModel{T}}) where T = T
vectype(::MosiModel{T}) where T = T
is_3d(model::Union{MosiModel, Type{<:MosiModel}}) = vectype(model) ≡ Vector3
natoms(model::MosiModel) = model.N
constraints(::MosiModel, rs) = Float64[]
constraint_gradients(::MosiModel{T}, rs) where T = Vector{T}[]
constraint_gradients(::MosiModel{T}, rs, i) where T = T[]

potential_energy_function(::MosiModel, rs) = error("Unimplemented")
force_function(::MosiModel, rs) = error("Unimplemented")
force_function(::MosiModel, rs, i) = error("Unimplemented")
potential_energy_gradients(model::MosiModel, rs) = -force_function(model, rs)
potential_energy_gradients(model::MosiModel, rs, i) = -force_function(model, rs, i)

pbc_box(model::MosiModel) = model.box
has_pbc(model::MosiModel{T}) where T = pbc_box(model) ≢ zero(T)
distance_function(model_or_system::Union{MosiModel, MosiSystem}) = if has_pbc(model_or_system)
    pbc_distance(pbc_box(model_or_system))
else
    default_distance
end

struct UnknownModel <: MosiModel{MosiVector} end
name(::UnknownModel) = "UnknownModel"


text:
@inline function ExaTronKernels.dnrm2(n::Int,x,incx::Int, tx)

    @synchronize
    v = 0.0
    for i in 1:n
        @inbounds v += x[i]*x[i]
    end

    @synchronize
    v = sqrt(v)
    @synchronize

    return v
end


text:
__precompile__()

module mbContinuation
export PC, tangent, incrementalLoading

using mbPred, mbNewton

function tangent(A)
	local t = [A;ones(1,size(A,2))] \ [zeros(size(A,1)); 1.0]
	return flipsign(t / norm(t), det([A;t']))
end

function PC(H, J, V₀, pred; init=Void, callback=Void, predC=predEps(1e-10), h₀=1.0, κ₀=.5, δ₀=1.0, α₀=pi/36, dir=true)
	local V,h; V,h = deepcopy(V₀),h₀
	V = mbNewton.newton(H, J, V, predC)

	init≠Void && init()
	callback≠Void && callback(V)

	local T′ = Void
	while pred(V)
		# P-step
		local T = tangent(J(V))
		local W = V + (dir?h:-h) * T

		# step-length adaption
		local κ,δ,α

		T′==Void && (T′=T)
		α = acos(clamp(dot(T′, T), -1.0, 1.0))
		T′=T

		α>.95pi && break #TODO fix... bifurcation point detection...

		local X = Any[]
		W = mbNewton.newton(H, J, W, predCount(2); callback=(W,H,J)->push!(X,W))
		δ = norm(X[1]-X[2])
		κ = norm(X[3]-X[2]) / δ

		local f = clamp(max(sqrt(κ/κ₀), sqrt(δ/δ₀), α/α₀), .5, 2.0)
		h /= f

		# C-step
		f≥2.0 && continue
		V = mbNewton.newton(H, J, W, predC)

		callback≠Void && callback(V)
	end

	return V
end

#NOTE untested!
function incrementalLoading(H, J, V₀, R; init=Void, callback=Void, predC=predEps(1e-10))
	local V = deepcopy(V₀)

	init≠Void && init()
	callback≠Void && callback(V)

	for ℵ in R
		V = mbNewton.newton(H, J, V, predC)
		callback≠Void && callback(V)
	end

	return V
end

end #module


text:
# -*- coding: utf-8 -*-
# # Decaying Surface QG turbulence
#
# A simulation of decaying surface quasi-geostrophic turbulence.
# We reproduce here the initial value problem for an elliptical
# vortex as done by Held et al. 1995, _J. Fluid Mech_.
#
# An example of decaying barotropic quasi-geostrophic turbulence over topography.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

using Pkg
pkg"add GeophysicalFlows, Plots, Printf, Random, Statistics"

# ## Let's begin
#
# Let's load `GeophysicalFlows.jl` and some other needed packages.

using GeophysicalFlows, Plots, Printf, Random

using Statistics: mean
using Random: seed!

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)

# ## Numerical parameters and time-stepping parameters

      n = 256                       # 2D resolution = n²
stepper = "FilteredETDRK4"          # timestepper
     dt = 0.03                      # timestep
     tf = 60                        # length of time for simulation
 nsteps = Int(tf / dt)              # total number of time-steps
 nsubs  = round(Int, nsteps/100)    # number of time-steps for intermediate logging/plotting (nsteps must be multiple of nsubs)

# ## Physical parameters

 L = 2π        # domain size
 ν = 1e-19     # hyper-viscosity coefficient
nν = 4         # hyper-viscosity order

# ## Problem setup
# We initialize a `Problem` by providing a set of keyword arguments. In this
# example numerical instability due to accumulation of buoyancy variance at high
# wavenumbers is taken care with the `FilteredTimestepper` we picked.
prob = SurfaceQG.Problem(dev; nx=n, Lx=L, dt=dt,
    stepper=stepper, ν=ν, nν=nν)

# Let's define some shortcuts.
sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
x, y = grid.x, grid.y

# ## Setting initial conditions
#
# We initialize the buoyancy equation with an elliptical vortex.
X, Y = gridpoints(grid)
b₀ = @. exp(-(X^2 + 4*Y^2))

SurfaceQG.set_b!(prob, b₀)

# Let's plot the initial condition. Note that when plotting, we decorate the variable to be
# plotted with `Array()` to make sure it is brought back on the CPU when `vars` live on the GPU.
heatmap(x, y, Array(vars.b'),
     aspectratio = 1,
               c = :deep,
            clim = (0, 1),
           xlims = (-grid.Lx/2, grid.Lx/2),
           ylims = (-grid.Ly/2, grid.Ly/2),
          xticks = -3:3,
          yticks = -3:3,
          xlabel = "x",
          ylabel = "y",
           title = "buoyancy bₛ",
      framestyle = :box)

# ## Diagnostics

# Create Diagnostics; `buoyancy_variance`, `kinetic_energy` and `buoyancy_dissipation`
# functions were imported at the top.
B  = Diagnostic(SurfaceQG.buoyancy_variance, prob; nsteps=nsteps)
KE = Diagnostic(SurfaceQG.kinetic_energy, prob; nsteps=nsteps)
Dᵇ = Diagnostic(SurfaceQG.buoyancy_dissipation, prob; nsteps=nsteps)
diags = [B, KE, Dᵇ] # A list of Diagnostics types passed to `stepforward!`. Diagnostics are updated every timestep.

# ## Output

# We choose folder for outputing `.jld2` files and snapshots (`.png` files).
# Define base filename so saved data can be distinguished from other runs
base_filename = string("SurfaceQG_decaying_n_", n)
# We choose folder for outputing `.jld2` files and snapshots (`.png` files).
datapath = "./"
plotpath = "./"

dataname = joinpath(datapath, base_filename)
plotname = joinpath(plotpath, base_filename)

# Do some basic file management,
if !isdir(plotpath); mkdir(plotpath); end
if !isdir(datapath); mkdir(datapath); end

# and then create Output.
get_sol(prob) = sol # extracts the Fourier-transformed solution
get_u(prob) = irfft(im * grid.l .* sqrt.(grid.invKrsq) .* sol, grid.nx)
out = Output(prob, dataname, (:sol, get_sol), (:u, get_u))

# ## Visualizing the simulation

# We define a function that plots the buoyancy field and the time evolution of kinetic energy
# and buoyancy variance.

function plot_output(prob)
  b = prob.vars.b

  pb = heatmap(x, y, Array(b'),
       aspectratio = 1,
                 c = :deep,
              clim = (0, 1),
             xlims = (-grid.Lx/2, grid.Lx/2),
             ylims = (-grid.Ly/2, grid.Ly/2),
            xticks = -3:3,
            yticks = -3:3,
            xlabel = "x",
            ylabel = "y",
             title = "buoyancy bₛ",
        framestyle = :box)

  pKE = plot(1,
             label = "kinetic energy ∫½(uₛ²+vₛ²)dxdy/L²",
         linewidth = 2,
            legend = :bottomright,
             alpha = 0.7,
             xlims = (0, tf),
             ylims = (0, 1e-2),
            xlabel = "t")

  pb² = plot(1,
             label = "buoyancy variance ∫bₛ²dxdy/L²",
         linecolor = :red,
            legend = :bottomright,
         linewidth = 2,
             alpha = 0.7,
             xlims = (0, tf),
             ylims = (0, 2e-2),
            xlabel = "t")

  layout = @layout [a{0.5w} Plots.grid(2, 1)]
  p = plot(pb, pKE, pb², layout=layout, size = (900, 500))

  return p
end
nothing # hide

# ## Time-stepping the `Problem` forward and create animation by updating the plot.

startwalltime = time()

p = plot_output(prob)

anim = @animate for j = 0:round(Int, nsteps/nsubs)
  if j % (500 / nsubs) == 0
    cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy])

    log1 = @sprintf("step: %04d, t: %.1f, cfl: %.3f, walltime: %.2f min",
          clock.step, clock.t, cfl, (time()-startwalltime)/60)

    log2 = @sprintf("buoyancy variance: %.2e, buoyancy variance dissipation: %.2e",
              B.data[B.i], Dᵇ.data[Dᵇ.i])

    println(log1)

    println(log2)
  end

  p[1][1][:z] = Array(vars.b)
  p[1][:title] = "buoyancy, t=" * @sprintf("%.2f", clock.t)
  push!(p[2][1], KE.t[KE.i], KE.data[KE.i])
  push!(p[3][1], B.t[B.i], B.data[B.i])

  stepforward!(prob, diags, nsubs)
  SurfaceQG.updatevars!(prob)
end

mp4(anim, "sqg_ellipticalvortex.mp4", fps=14)

# Let's see how all flow fields look like at the end of the simulation.

pu = heatmap(x, y, Array(vars.u'),
     aspectratio = 1,
               c = :balance,
            clim = (-maximum(abs.(vars.u)), maximum(abs.(vars.u))),
           xlims = (-L/2, L/2),
           ylims = (-L/2, L/2),
          xticks = -3:3,
          yticks = -3:3,
          xlabel = "x",
          ylabel = "y",
           title = "uₛ(x, y, t=" * @sprintf("%.2f", clock.t) * ")",
      framestyle = :box)

pv = heatmap(x, y, Array(vars.v'),
     aspectratio = 1,
               c = :balance,
            clim = (-maximum(abs.(vars.v)), maximum(abs.(vars.v))),
           xlims = (-L/2, L/2),
           ylims = (-L/2, L/2),
          xticks = -3:3,
          yticks = -3:3,
          xlabel = "x",
          ylabel = "y",
           title = "vₛ(x, y, t=" * @sprintf("%.2f", clock.t) * ")",
      framestyle = :box)

pb = heatmap(x, y, Array(vars.b'),
     aspectratio = 1,
               c = :deep,
            clim = (0, 1),
           xlims = (-L/2, L/2),
           ylims = (-L/2, L/2),
          xticks = -3:3,
          yticks = -3:3,
          xlabel = "x",
          ylabel = "y",
           title = "bₛ(x, y, t=" * @sprintf("%.2f", clock.t) * ")",
      framestyle = :box)

layout = @layout [a{0.5h}; b{0.5w} c{0.5w}]

plot_final = plot(pb, pu, pv, layout=layout, size = (800, 800))


text:
using ProtoBuf: readproto, PipeBuffer
using CodecZlib: ZlibDecompressorStream
using Dates: unix2datetime, DateTime

"""
    readpbf(filename)

`readpbf` has only one argument `filename`, taking a string of the pbf-file path and name.
It returns an object containing the OSM data.
"""
function readpbf(filename::String)::Map
    osmdata = Map()
    blobheader = OSMPBF.BlobHeader()
    blob = OSMPBF.Blob()
    open(filename, "r") do f
        readnext!(f, blobheader, blob)
        @assert blobheader._type == "OSMHeader"
        processheader!(osmdata, readblock!(blob, OSMPBF.HeaderBlock()))
        while !eof(f)
            readnext!(f, blobheader, blob)
            @assert blobheader._type == "OSMData"
            processblock!(osmdata, readblock!(blob, OSMPBF.PrimitiveBlock()))
        end
    end
    return osmdata
end

function readnext!(f, blobheader::OSMPBF.BlobHeader, blob::OSMPBF.Blob)
    n = ntoh(read(f, UInt32))
    readproto(PipeBuffer(read(f, n)), blobheader)
    readproto(PipeBuffer(read(f, blobheader.datasize)), blob)
end

function readblock!(blob::OSMPBF.Blob, block:: Union{OSMPBF.HeaderBlock,OSMPBF.PrimitiveBlock})
    @assert xor(isempty(blob.raw), isempty(blob.zlib_data))
    if !isempty(blob.raw)
        readproto(PipeBuffer(blob.raw), block)
    elseif !isempty(blob.zlib_data)
        readproto(
            ZlibDecompressorStream(IOBuffer(blob.zlib_data)),
            block
        )
    else
        DomainError("Unsupported blob data format")
    end
end

function processheader!(osmdata::Map, header::OSMPBF.HeaderBlock)
    if hasproperty(header, :bbox)
        osmdata.meta[:bbox] = BBox(
            round(1e-9 * header.bbox.bottom, digits=7),
            round(1e-9 * header.bbox.left, digits=7),
            round(1e-9 * header.bbox.top, digits=7),
            round(1e-9 * header.bbox.right, digits=7)
        )
    end
    if hasproperty(header, :osmosis_replication_timestamp)
        osmdata.meta[:writenat] = unix2datetime(header.osmosis_replication_timestamp)
    end
    if hasproperty(header, :osmosis_replication_sequence_number)
        osmdata.meta[:sequencenumber] = header.osmosis_replication_sequence_number
    end
    if hasproperty(header, :osmosis_replication_base_url)
        osmdata.meta[:baseurl] = header.osmosis_replication_base_url
    end
    if hasproperty(header, :writingprogram)
        osmdata.meta[:writingprogram] = header.writingprogram
    end
end

function processblock!(osmdata::Map, primblock::OSMPBF.PrimitiveBlock)
    lookuptable =  Base.transcode.(String, primblock.stringtable.s)
    latlonparameter = Dict(
        :lat_offset =>  primblock.lat_offset,
        :lon_offset =>  primblock.lon_offset,
        :granularity => primblock.granularity
    )
    for primgrp in primblock.primitivegroup
        # Possible extension: callback functions for the selecton of specific elements (e.g. for routing).
        merge!(osmdata.nodes, extractnodes(primgrp, lookuptable))
        if hasproperty(primgrp, :dense)
            merge!(osmdata.nodes,  extractdensenodes(primgrp, lookuptable, latlonparameter))
        end
        merge!(osmdata.ways, extractways(primgrp, lookuptable))
        merge!(osmdata.relations, extractrelations(primgrp, lookuptable))
    end
end

function extractnodes(primgrp::OSMPBF.PrimitiveGroup, lookuptable::Vector{String})::Dict{Int64,Node}
    nodes = Dict{Int64,Node}()
    for n in primgrp.nodes
        @assert length(n.keys) == length(n.vals)
        if length(n.keys) > 0
            tags = Dict{Symbol,String}()
            for (k, v) in zip(n.keys, n.vals)
                tags[Symbol(lookuptable[k + 1])] = lookuptable[v + 1]
            end
            nodes[n.id] = Node(LatLon(n.lat, n.lon), tags)
        else
            nodes[n.id] = Node(LatLon(n.lat, n.lon), nothing)
        end
    end
    return nodes
end

function extractdensenodes(primgrp::OSMPBF.PrimitiveGroup, lookuptable::Vector{String}, latlonparameter::Dict)::Dict{Int64,Node}
    ids = cumsum(primgrp.dense.id)
    lats = round.(1e-9 * (latlonparameter[:lat_offset] .+ latlonparameter[:granularity] .* cumsum(primgrp.dense.lat)), digits=7)
    lons = round.(1e-9 * (latlonparameter[:lon_offset] .+ latlonparameter[:granularity] .* cumsum(primgrp.dense.lon)), digits=7)
    @assert length(ids) == length(lats) == length(lons)
    # extract tags
    @assert primgrp.dense.keys_vals[end] == 0
    # decode tags i: node id index, kv: key-value index, k: key index, v: value index
    i = 1; kv = 1
    tags = Dict{Int64,Dict{Symbol,String}}()
    while kv <= length(primgrp.dense.keys_vals)
        k = primgrp.dense.keys_vals[kv]
        if k == 0
            # move to next node
            i += 1; kv += 1
        else
            # continue with current note
            @assert kv < length(primgrp.dense.keys_vals)
            v = primgrp.dense.keys_vals[kv + 1]
            id = ids[i]
            if !haskey(tags, id)
                tags[id] =  Dict{Symbol,String}()
            end
            tags[id][Symbol(lookuptable[k + 1])] = lookuptable[v + 1]
            kv += 2
        end
    end
    # assemble Node objects
    nodes = Dict{Int64,Node}()
    for (id, lat, lon) in zip(ids, lats, lons)
        nodes[id] = Node(LatLon(lat, lon), get(tags, id, nothing))
    end
    return nodes
end

function extractways(primgrp::OSMPBF.PrimitiveGroup, lookuptable::Vector{String})::Dict{Int64,Way}
    ways = Dict{Int64,Way}()
    for w in primgrp.ways
        if length(w.keys) > 0
            tags = Dict{Symbol,String}()
            for (k, v) in zip(w.keys, w.vals)
                tags[Symbol(lookuptable[k + 1])] = lookuptable[v + 1]
            end
            ways[w.id] = Way(cumsum(w.refs), tags)
        else
            ways[w.id] = Way(cumsum(w.refs), nothing)
        end
    end
    return ways
end

function extractrelations(primgrp::OSMPBF.PrimitiveGroup, lookuptable::Vector{String})::Dict{Int64,Relation}
    relations = Dict{Int64,Relation}()
    for r in primgrp.relations
        if length(r.keys) > 0
            tags = Dict{Symbol,String}()
            for (k, v) in zip(r.keys, r.vals)
                tags[Symbol(lookuptable[k + 1])] = lookuptable[v + 1]
            end
            relations[r.id] = Relation(cumsum(r.memids), membertype.(r.types), lookuptable[r.roles_sid .+ 1], tags)
        else
            relations[r.id] = Relation(cumsum(r.memids), membertype.(r.types), lookuptable[r.roles_sid .+ 1], nothing)
        end
    end
    return relations
end

function membertype(i)
    if i == 0; :node elseif i == 1; :way else; :relation end
end


text:
import ECOS

include("LPBenchmark.jl")

function run_ecos(fname::String, fsol::String="")

    ecos = ECOS.Optimizer()
    cache = MOIU.UniversalFallback(MOIU.Model{Float64}())
    cached = MOIU.CachingOptimizer(cache, ecos)
    bridged = MOIB.full_bridge_optimizer(cached, Float64)

    # Read problem
    load_problem!(bridged, fname)

    # Set some parameters
    MOI.set(bridged, MOI.Silent(), false)

    # Solve the problem
    # Note: this will (i) convert the data to ECOS' format (ii) solve
    # ECOS' internal solve time (excluding conversion) is displayed in log
    run(bridged)

    # Display timing info
    return nothing
end

if abspath(PROGRAM_FILE) == @__FILE__
    run_ecos(ARGS[1])
end

max_stars_repo_name: mtanneau/LPBenchmarks

max_issues_repo_name: mtanneau/LPBenchmarks

max_forks_repo_name: mtanneau/LPBenchmarks


text:
using LaborAPI
using Test

@testset "LaborAPI.jl" begin
    # Write your tests here.
end


text:
module Spectra

export power_spectrum_1d, power_spectrum_2d, power_spectrum_3d

using FFTW
using MPI

using ..ConfigTypes
using ..Mesh.Grids
using ..MPIStateArrays

include("power_spectrum_les.jl")
include("power_spectrum_gcm.jl")

end


text:
"""
    _orientation(x::Tuple{<:Number, <:Number}, y::Tuple{<:Number, <:Number}, z::Tuple{<:Number, <:Number})

Utility function for `convex_hull`. Given an ordered triplet, this function returns if three
points are collinear, oriented in clock-wise or anticlock-wise direction.
"""
function _orientation(x::Tuple{<:Number, <:Number}, y::Tuple{<:Number, <:Number}, z::Tuple{<:Number, <:Number})

    a = ((y[2] - x[2]) * (z[1] - y[1]) - (y[1] - x[1]) * (z[2] - y[2]))

    if isapprox(a, 0, atol=1e-6)
        return 0 # collinear
    elseif a > 0
        return 1 # clock-wise
    else
        return 2 # counterclock-wise
    end

end

"""
    _lt_filter(a::Tuple{<:Number, <:Number}, b::Tuple{<:Number, <:Number})

Utility function for sorting step in `convex_hull`. Given two points, `a` and `b`, this function
returns true if `a` has larger polar angle (counterclock-wise direction) than `b` w.r.t. first point `chull_p1`.
"""
function _lt_filter(a::Tuple{<:Number, <:Number}, b::Tuple{<:Number, <:Number})
    o = QCO._orientation(chull_p1, a, b)

    if o == 0
        if LA.norm(collect(chull_p1 .- b))^2 >= LA.norm(collect(chull_p1 .- a))^2
            return true
        else
            return false
        end
    elseif o == 2
        return true
    else
        return false
    end

end

"""
    convex_hull(points::Vector{T}) where T<:Tuple{Number, Number}

Graham's scan algorithm to compute the convex hull of a finite set of `n` points in a plane
with time complexity `O(n*log(n))`. Given a vector of tuples of co-ordinates, this function returns a
vector of tuples of co-ordinates which form the convex hull of the given set of points.

Sources: https://doi.org/10.1016/0020-0190(72)90045-2
         https://en.wikipedia.org/wiki/Graham_scan
"""
function convex_hull(points::Vector{T}) where T<:Tuple{Number, Number}

    num_points = size(points)[1]

    if num_points == 0
        Memento.error(_LOGGER, "Atleast one point is necessary for evaluating the convex hull")
    elseif num_points <= 2
        return points
    end

    min_y = points[1][2]
    min = 1

    # Bottom-most or else the left-most point when tie
    for i in 2:num_points
        y = points[i][2]
        if (y < min_y) || (y == min_y && points[i][1] < points[min][1])
            min_y = y
            min = i
        end
    end

    # Placing the bottom/left-most point at first position
    points[1], points[min] = points[min], points[1]

    global chull_p1 = points[1]

    points = Base.sort(points, lt = QCO._lt_filter)

    # If two or more points are collinear with chull_p1, remove all except the farthest from chull_p1
    arr_size = 2
    for i in 2:num_points
        while (i < num_points) && (QCO._orientation(chull_p1, points[i], points[i+1]) == 0)
            i += 1
        end
        points[arr_size] = points[i]
        arr_size += 1
    end

    if arr_size <= 3
        return []
    end

    chull = []
    append!(chull, [points[1]])
    append!(chull, [points[2]])
    append!(chull, [points[3]])

    for i in 4:(arr_size-1)
        while (size(chull)[1] > 1) && (QCO._orientation(chull[end - 1], chull[end], points[i]) != 2)
            pop!(chull)
        end
        append!(chull, [points[i]])
    end

    return chull
end


text:
export plottrial, trialaxes

iv2t_s(iv) = (minimum(iv)/(1s), maximum(iv)/(1s))

"""
    plottrial(dfr::DataFrameRate, td::TrialData; axsniff, axtiming)
    plottrial(signal::AbstractVector, td::TrialData; fs, axsniff, axtiming)

Plot a trial in the specified axes. See [`trialaxes`](@ref) to create them.
"""
function plottrial(signal::AbstractVector, td::TrialData; fs, axsniff, axtiming, showlegend=true)
    axsniff.cla()
    axtiming.cla()
    sniff = signal[td.trialrange]
    t = (0:length(sniff)-1)/fs
    axsniff.plot(t * (1Hz), sniff, color="green")
    axsniff.set_xlabel("Time (s)")
    axsniff.set_ylabel("Chest expansion (V)")
    xl = axsniff.get_xlim()
    axtiming.plot(iv2t_s(td.soundinterval), (1, 1), color="blue")
    tcue = convert(Float64, td.cuetime / (1s))
    axtiming.plot(tcue, 1.5, "kv")
    axtiming.plot(iv2t_s(td.odorinterval), (2, 2), color="cyan")
    axtiming.text(iv2t_s(td.odorinterval)[end], 2, string(td.odordirection) * (is_lick(td) ? "+" : "-"))
    tlick = td.lickonsets / (1s)
    if !isempty(tlick)
        axtiming.plot([tlick'; tlick'], [fill(2.75, length(tlick))'; fill(3.25, length(tlick))'], color="red")
    end
    axtiming.set_xlim(xl)
    axtiming.set_ylim((0.5, 3.75))
    axtiming.set_xticks([])
    axtiming.set_yticks([])
    axsniff.set_frame_on(false)
    axtiming.set_frame_on(false)
    axtiming.invert_yaxis()
    if showlegend
        axtiming.legend(("sound", "cue", "odor", "licks"))
    end
    return nothing
end

plottrial(dfr::DataFrameRate, td::TrialData; kwargs...) = plottrial(dfr.sniff, td; fs=dfr.fs, kwargs...)

"""
    axsniff, axtiming = plottrial(dfr::DataFrameRate, td::TrialData, fig::PyPlot.Figure)

Plot a trial in the given `fig`ure.
"""
function plottrial(dfr::DataFrameRate, td::TrialData, fig::PyPlot.Figure)
    axsniff, axtiming = trialaxes(fig)
    plottrial(dfr, td; axsniff, axtiming)
    return axsniff, axtiming
end

"""
    axsniff, axtiming = trialaxes(fig)

Create axes for plotting trials.
"""
function trialaxes(fig::PyPlot.Figure)
    gs = PyPlot.plt.GridSpec(2, 1, fig, height_ratios=[1, 3])
    return fig.add_subplot(gs[2]), fig.add_subplot(gs[1])
end


text:
using Images
using ImageSegmentation
using ImageFiltering
using CSV
using DataFrames, DataFramesMeta
using Dates
using Statistics, StatsBase, Random
using Plots, StatsPlots, Plots.PlotMeasures

function to_array(echodf)
    echo_arr = unstack(echodf, :Interval, :Layer_depth_min, :Sv_mean)
    echo_arr = Matrix(echo_arr[:, 2:end])
    echo_arr[ismissing.(echo_arr) .| (echo_arr .< -90)] .= -90
    return disallowmissing(echo_arr)
end

function find_blobs(echo, threshold)
    bw = echo .> threshold
    markers = label_components(bw)
    w = watershed(echo, markers)
    return labels_map(w)
end

function segment_stats(label, labelmap, echo)
    indices = findall(x -> x == label, labelmap)
    bottom, top = extrema(i.I[2] for i in indices)
    left, right = extrema(i.I[1] for i in indices)
    height = top - bottom + 1
    @assert height > 0 println("bottom: $bottom, top: $top")
    width = right - left + 1
    values = [echo[i] for i in indices]
    npixels = length(values)
    μ = mean(values)
    σ = std(values)
    return (bottom=bottom, top=top, left=left, right=right,
        height=height, width=width, npixels=npixels,
        mean=μ, std=σ, total=μ*npixels)
end

function process_day(df)
    echo_arr = to_array(df)
    datetimes = sort(unique(df.datetime))
    depths = sort(unique(df.Layer_depth_min))
    markers = find_blobs(echo_arr, -60)
    schools = DataFrame(segment_stats(i, markers, exp10.(echo_arr/10))
        for i in unique(markers))
    schools[!, :school_id] .= unique(markers)
    schools = @transform(schools,
        # school_id .= unique(markers),
        :start_time = datetimes[:left],
        # top and bottom are reversed to go from image to depth coordinates
        :bottom = depths[:top],
        :top = depths[:bottom])
    schools = @transform(schools,
        :middle = (:bottom .+ :top) ./ 2,
        :date = fill(first(df.date), nrow(schools)))

    schools = @subset(schools, 2 .< :npixels .< 1000, :width .<= 12)

    # now calculate center of mass, ignoring regions classified as schools
    sv = exp10.(echo_arr/10)
    ii = in(schools.school_id).(markers)
    sv[ii] .= 0
    cm = vec(sum(sv .* depths', dims=2) ./ sum(sv, dims=2))
    layertop = depths[vec(mapslices(ping -> findfirst(cumsum(ping) .> 1e-6), sv, dims=2))]
    layerbase = reverse(depths)[vec(mapslices(ping -> findfirst(cumsum(reverse(ping)) .> 1e-6), sv, dims=2))]
    center_of_mass = DataFrame(date=first(df.date), datetime=datetimes,
        cm=cm, layertop=layertop, layerbase=layerbase)
    return schools, center_of_mass
end

function get_random_color(seed)
    Random.seed!(seed)
    rand(RGB{N0f8})
end

#######################################################################################
# This section illustrates the school-detection procedure on a three-day subset of data.
#######################################################################################
# Reading in a month of acoustic data and formatting it
echo = CSV.read(joinpath(@__DIR__, "../data/monthly/deimos-2019-11.csv"),
    DataFrame, normalizenames=true)

echo = @chain echo begin
    @transform(:datetime = Date.(string.(:Date_M), dateformat"yyyymmdd") .+ :Time_M .- Hour(8))
    @transform(:hour = hour.(:datetime),
              :date = Date.(:datetime))
end

# Selecting three days with some schools
echo_sub = @subset(echo,
    :datetime .> DateTime(2019, 11, 5, 0),
    :datetime .< DateTime(2019, 11, 8, 0),
    :Layer_depth_min .<300)
# convert the data frame to a 2d array
echo_arr = to_array(echo_sub)
datetimes = sort(unique(echo_sub.datetime))
depths = sort(unique(echo_sub.Layer_depth_min))

# locate bright regions on the echogram and label them
markers = find_blobs(echo_arr, -60)
heatmap(map(i->get_random_color(i), markers'))

# calculate statistics for each individual region, storing them in a DataFrame
schools = DataFrame(segment_stats(i, markers, exp10.(echo_arr/10))
  for i in unique(markers))
schools.school_id = unique(markers)

schools = @transform(schools,
  :start_time = datetimes[:left],
  # top and bottom are reversed to go from image to depth coordinates
  :bottom = depths[:top],
  :top = depths[:bottom])
schools = @transform(schools,
  :middle = (:bottom .+ :top) ./ 2,
  :date = fill(first(echo_sub.date), nrow(schools)))

# Drop schools that are too small or too big
schools = @subset(schools, 2 .< :npixels .< 1000, :width .< 12)

# Blank out schools from the echogram
sv = exp10.(echo_arr/10)
ii = in(schools.school_id).(markers)
sv[ii] .= 0
# Echogram with the schools blanked out
heatmap(10log10.(sv'), yflip=true)

# Calculate time series of several layer metrics, ignoring regions classified as schools
cm = vec(sum(sv .* depths', dims=2) ./ sum(sv, dims=2))
center_of_mass = DataFrame(date=first(echo_sub.date), datetime=datetimes, cm=cm)
layertop = depths[vec(mapslices(ping -> findfirst(cumsum(ping) .> 1e-6), sv, dims=2))]
layerbase = reverse(depths)[vec(mapslices(ping -> findfirst(cumsum(reverse(ping)) .> 1e-6), sv, dims=2))]

heatmap(1:432, depths, 10log10.(sv)', yflip=true)
plot!(cm, color=:white, label = "")
plot!(layertop, color=:green, label="")
plot!(layerbase, color=:blue, label="")

## Process all DEIMOS data
schools_list = []
cm_list = []
files = readdir(joinpath(@__DIR__, "../data/monthly"), join=true)

for f in files
    println("Reading $(basename(f))...")
    echo = CSV.read(f, DataFrame, normalizenames=true)
    echo = @chain echo begin
        @subset(15 .<= :Layer_depth_min .<= 500)
        @transform(:datetime = Date.(string.(:Date_M), dateformat"yyyymmdd") .+ :Time_M .- Hour(8))
        @transform(:hour = hour.(:datetime),
                  :date = Date.(:datetime))
    end
    println("Detecting schools...")
    schools = []
    center_of_mass = []
    for echo_day in groupby(echo, :date)
        schools_day, center_of_mass_day = process_day(echo_day)
        push!(schools, schools_day)
        push!(center_of_mass, center_of_mass_day)
    end
    schools = vcat(schools...)
    schools[!, :school_id] .= string.(schools.date) .* "-S" .* string.(schools.school_id)
    center_of_mass = vcat(center_of_mass...)

    push!(schools_list, schools)
    push!(cm_list, center_of_mass)
end
schools = vcat(schools_list...)
center_of_mass = @chain vcat(cm_list...) begin
    @transform(:hour = hour.(:datetime))
end

histogram(log10.(schools.total))
histogram(schools.bottom, bins=200)
histogram(@subset(schools, :width .<= 2).bottom, bins=100)
schools = @subset(schools, :bottom .< 200)

schools_summary = @chain schools begin
    groupby(:date)
    @combine(:nschools = length(:school_id),
        :sa = sum(:total),
        :mean_depth = mean(:middle),
        :max_depth = maximum(:bottom))
end

alldates = DataFrame(date=first(schools.date):Day(1):last(schools.date))
schools_summary = leftjoin(alldates, schools_summary, on=:date)
schools_summary = @chain schools_summary begin
    @transform(:nschools = replace(:nschools, missing => 0),
               :sa = replace(:sa, missing => 0))
    @transform(:Sa = 10log10.(:sa))
    @orderby(:date)
end

@df schools_summary plot(:date, :sa)
@df schools_summary plot(:date, :nschools)
@df schools_summary plot(:date, [:mean_depth, :max_depth])
@df schools_summary scatter(:nschools, :Sa)

@df center_of_mass plot(:datetime, :cm)

center_of_mass = @chain center_of_mass begin
    @transform(:cm_smooth = mapwindow(mean, :cm, 6*3+1),
               :layertop_smooth = mapwindow(mean, :layertop, 6*3+1),
               :layerbase_smooth = mapwindow(mean, :layerbase, 6*3+1))
end

@df center_of_mass[10_000:11_000,:] plot(:datetime,
    [:cm, :cm_smooth, :layertop, :layertop_smooth, :layerbase, :layerbase_smooth])

CSV.write(joinpath(@__DIR__, "..\\data\\schools.csv"), schools)
CSV.write(joinpath(@__DIR__, "..\\data\\schools_summary.csv"), schools_summary)
CSV.write(joinpath(@__DIR__, "..\\data\\center_of_mass.csv"), center_of_mass)

schools = CSV.read(joinpath(@__DIR__, "..\\data\\schools.csv"), DataFrame)
schools_summary = CSV.read(joinpath(@__DIR__, "..\\data\\schools_summary.csv"), DataFrame)
center_of_mass = CSV.read(joinpath(@__DIR__, "..\\data\\center_of_mass.csv"), DataFrame)

## Calculating migration statistics

noon(hour) = hour in [12, 13, 14]
midnight(hour) = hour in [0, 1, 2]

function hourlabel(hour)
    if noon(hour)
        return "noon"
    elseif midnight(hour)
        return "midnight"
    else
        return "other"
    end
end

migration = @chain center_of_mass begin
    @transform(:hourlabel = hourlabel.(:hour))
    @subset(:hourlabel .!= "other")
    groupby([:date, :hourlabel])
    @combine(:cm=median(:cm_smooth),
        :layertop=median(:layertop_smooth))
end

migration = innerjoin(unstack(migration, :date, :hourlabel, :cm, renamecols=x->Symbol(:cm_, x)),
     unstack(migration, :date, :hourlabel, :layertop, renamecols=x->Symbol(:layertop_, x)),
     on=:date)
migration = @chain migration begin
    @transform(:Δcm = :cm_noon .- :cm_midnight)
    @transform(:Δlayertop = :layertop_noon .- :layertop_midnight)
    leftjoin(schools_summary, on=:date)
    @transform(:nschools = replace(:nschools, missing => 0),
               :sa = replace(:sa, missing => 0),
               :Sa = replace(:Sa, -Inf => missing))
end

@df migration plot(:date, [:Δcm, :Δlayertop], labels=["ΔCM" "ΔLayertop"])

@df @subset(center_of_mass, Date(2019, 4, 1) .< :date .< Date(2019, 5, 31)) plot(:datetime, :cm_smooth)

plot(plot(migration.date, migration.layertop_noon, ylabel="Migration"),
    plot(migration.date, migration.sa, ylabel="schools"), layout=(2,1))

@df migration scatter(:nschools, [:cm_noon, :layertop_noon], labels=["CM" "Layer top"])
@df migration scatter(:Sa, [:cm_noon, :layertop_noon], labels=["CM" "Layer top"])

@df migration scatter(:nschools, [:Δcm, :Δlayertop], labels=["CM" "Layer top"])
@df migration scatter(:Sa, [:Δcm, :Δlayertop], labels=["CM" "Layer top"])

## Regression modeling
using GLM

# CMigration, center of mass
m1 = lm(@formula(Δcm ~ nschools), migration)
r2(m1)
m2 = lm(@formula(Δcm ~ Sa), migration)
r2(m2)
m3 = lm(@formula(Δcm ~ Sa + nschools), migration)
r2(m3)

# Migration, layer top
m1 = lm(@formula(Δlayertop ~ nschools), migration)
r2(m1)
m2 = lm(@formula(Δlayertop ~ Sa), migration)
r2(m2)
m3 = lm(@formula(Δlayertop ~ Sa + nschools), migration)
r2(m3)

p1 = @df migration scatter(:Sa, :Δlayertop, legend=false, markerstrokewidth=0, color=:grey,
    xlabel="School backscatter (dB re m² minute)", ylabel="Migration magnitude (m)",
    title="A", dpi=600)
plot!(p1, Sa -> coef(m2)[1] + coef(m2)[2] * Sa, -58, -12, color=:black)

# Noon CM depth
m1 = lm(@formula(cm_noon ~ nschools), migration)
r2(m1)
m2 = lm(@formula(cm_noon ~ Sa), migration)
r2(m2)
m3 = lm(@formula(cm_noon ~ Sa + nschools), migration)
r2(m3)

# Noon layer top depth
m1 = lm(@formula(layertop_noon ~ nschools), migration)
r2(m1)
m2 = lm(@formula(layertop_noon ~ Sa), migration)
r2(m2)
m3 = lm(@formula(layertop_noon ~ Sa + nschools), migration)
r2(m3)

p2 = @df migration scatter(:Sa, :layertop_noon, legend=false, yflip=true,
    markerstrokewidth=0, color=:grey, title="B",
    xlabel="School backscatter (dB re m² minute)", ylabel="Depth at noon (m)", dpi=600)
plot!(p2, Sa -> coef(m2)[1] + coef(m2)[2] * Sa, -58, -12, color=:black)

p = plot(p1, p2, size=(1000, 400), dpi=600, title_loc=:left, margin=5mm)
savefig(p, joinpath(@__DIR__, "../graphics/schools_vs_layers.png"))

## Check against oceanography too

oceanography = CSV.read(joinpath(@__DIR__, "../data/oceanography.csv"), DataFrame)
oceanography = @chain oceanography begin
    @subset(.! ismissing.(:date))
    groupby(:date)
    @combine(:temperature = mean(skipmissing(:temperature)),
             :upwelling = mean(skipmissing(:upwelling)),
             :sealevel = mean(skipmissing(:sealevel_lowpass)))
    @transform(upwelling = replace(:upwelling, NaN => missing))
    rightjoin(migration, on=:date)
endlibrary(lubridate)
library(viridis)
library(scales)

@df oceanography scatter(:temperature, :layertop_noon, legend=false,
    xlabel="SST (°C)", ylabel="Noon depth (m)", dpi=300)
savefig(joinpath(@__DIR__, "../graphics/temperature_vs_layertop.png"))
@df oceanography scatter(:sealevel, :layertop_noon, legend=false,
    xlabel="Sea level (m)", ylabel="Noon depth (m)", dpi=300)
savefig(joinpath(@__DIR__, "../graphics/sealevel_vs_layertop.png"))
@df oceanography scatter(:upwelling, :layertop_noon, legend=false,
    xlabel="Upwelling (kg m⁻¹ s⁻¹)", ylabel="Noon depth (m)", dpi=300)
savefig(joinpath(@__DIR__, "../graphics/upwelling_vs_layertop.png"))

mtemp = lm(@formula(layertop_noon ~ temperature), oceanography)
r2(mtemp)
msealevel = lm(@formula(layertop_noon ~ sealevel), oceanography)
r2(msealevel)
mupwelling = lm(@formula(layertop_noon ~ upwelling), dropmissing(oceanography))
r2(mupwelling)

max_stars_repo_name: ElOceanografo/PelagicFear

max_issues_repo_name: ElOceanografo/PelagicFear

max_forks_repo_name: ElOceanografo/PelagicFear


text:
using AdvancedHMC, Distributions, ForwardDiff

# Choose parameter dimensionality and initial parameter value
D = 10; initial_θ = rand(D)

# Define the target distribution
ℓπ(θ) = logpdf(MvNormal(zeros(D), ones(D)), θ)

# Set the number of samples to draw and warmup iterations
n_samples, n_adapts = 2_000, 1_000

# Define a Hamiltonian system
metric = DiagEuclideanMetric(D)
hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)

# Define a leapfrog solver, with initial step size chosen heuristically
initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
integrator = Leapfrog(initial_ϵ)

# Define an HMC sampler, with the following components
#   - multinomial sampling scheme,
#   - generalised No-U-Turn criteria, and
#   - windowed adaption for step-size and diagonal mass matrix
proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator)
adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))

# Run the sampler to draw samples from the specified Gaussian, where
#   - `samples` will store the samples
#   - `stats` will store diagnostic statistics for each sample
samples, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; progress=true)


text:
## return the right Cholesky factor from the LDLt decomposition of A⁻¹
function Tinvt{T}(p::Pedigree{T})
    n = length(p.sire)
    cpi = Array(Int32,n + 1)            # column pointers
    cpi[1] = one(Int32)
    rvi = sizehint(Int32[],n + n>>1)    # row values of T⁻¹
    for i in 1:n
        r1 = min(p.sire[i],p.dam[i])
        r2 = max(p.sire[i],p.dam[i])
        if r2 == 0
            cpi[i+1] = cpi[i]
        elseif r1 == 0
            cpi[i+1] = cpi[i] + one(Int32)
            push!(rvi,r2)
        else
            cpi[i+1] = cpi[i] + convert(Int32,2)
            push!(rvi,r1)
            push!(rvi,r2)
        end
    end
    Triangular(SparseMatrixCSC(n,n,cpi,rvi,fill(-0.5,(length(rvi,)))),:L,true)
end

## return the right Cholesky factor, Lt, from the LLt factorization of A
function Ltrans{T}(p::Pedigree{T})
    n = length(p.sire)
    tri = Tinvt(p).data
    cpi = tri.colptr
    rvi = tri.rowval
    cpl = sizehint(Int32[1],n+1)             # column pointers for L'
    rvl = sizehint(Int32[], 8n)              # row values for L'
    nvl = sizehint(Float64[], 8n)            # nonzero values in L'
    anc = sizehint(IntSet([]),n)             # set of ancestors
    for j in 1:n
        empty!(anc)
        l2parents = 0.0   # accumulator for sum of squares of L cols of parents
        for k in cpi[j]:(cpi[j+1]-1)    # each parent in the pedigree
            for i in cpl[rvi[k]]:(cpl[rvi[k]+1]-1)
                push!(anc,rvl[i])       # add the parents ancestors to the set
                l2parents += abs2(nvl[i])
            end
        end
        push!(anc,j)                    # the animal counts as one of its ancestors
        push!(cpl,cpl[end] + length(anc))
        rv = collect(anc)
        append!(rvl,rv)
        lrv = length(rv)
        cc = zeros(eltype(tri),(lrv,))
        cc[lrv] = sqrt(1.0 - 0.25*l2parents)
        for k in cpi[j]:(cpi[j+1]-1)    # each parent in the pedigree
            for i in cpl[rvi[k]]:(cpl[rvi[k]+1]-1)  # add 0.5*each parent's column
                cc[searchsortedfirst(rv,rvl[i])] += 0.5*nvl[i]
            end
        end
        append!(nvl,cc)
    end
    SparseMatrixCSC(n,n,cpl,rvl,nvl)
end

function inbreeding(p::Pedigree)
    Lt = Ltrans(p)
    inb = zeros(Lt.n)
    nvl = Lt.nzval
    cpt = Lt.colptr
    for i in 1:Lt.n
        for k in cpt[i]:(cpt[i+1]-1)
            inb[i] += abs2(nvl[k])
        end
        inb[i] -= one(eltype(inb))
    end
    inb,Lt
end


text:
import Base.copy

export
    CostFunction,
    QuadraticQuatCost,
    RBCost,
    QuatLQRCost,
    SatDiffCost,
    ErrorQuadratic

#*********************************#
#       COST FUNCTION CLASS       #
#*********************************#

abstract type CostFunction end

#######################################################
#              COST FUNCTION INTERFACE                #
#######################################################

"""
$(TYPEDEF)
Cost function of the form
    1/2xₙᵀ Qf xₙ + qfᵀxₙ +  ∫ ( 1/2xᵀQx + 1/2uᵀRu + xᵀHu + qᵀx + rᵀu ) dt from 0 to tf
R must be positive definite, Q and Qf must be positive semidefinite

Constructor use any of the following constructors:
```julia
QuadraticCost(Q, R, H, q, r, c)
QuadraticCost(Q, R; H, q, r, c)
QuadraticCost(Q, q, c)
```
Any optional or omitted values will be set to zero(s).
"""
mutable struct QuadraticCost{TQ,TR,TH,Tq,Tr,T} <: CostFunction
    Q::TQ                 # Quadratic stage cost for states (n,n)
    R::TR                 # Quadratic stage cost for controls (m,m)
    H::TH                 # Quadratic Cross-coupling for state and controls (n,m)
    q::Tq                 # Linear term on states (n,)
    r::Tr                 # Linear term on controls (m,)
    c::T                                 # constant term
    function QuadraticCost(Q::TQ, R::TR, H::TH,
            q::Tq, r::Tr, c::T; checks=true) where {TQ,TR,TH,Tq,Tr,T}
        @assert size(Q,1) == length(q)
        @assert size(R,1) == length(r)
        @assert size(H) == (length(r), length(q))
        if checks
            if !isposdef(Array(R))
                @warn "R is not positive definite"
            end
            if !ispossemidef(Array(Q))
                err = ArgumentError("Q must be positive semi-definite")
                throw(err)
            end
        end
        new{TQ,TR,TH,Tq,Tr,T}(Q,R,H,q,r,c)
    end
end

state_dim(cost::QuadraticCost) = length(cost.q)
control_dim(cost::QuadraticCost) = length(cost.r)

# Constructors
function QuadraticCost(Q,R; H=similar(Q,size(R,1), size(Q,1)), q=zeros(size(Q,1)),
        r=zeros(size(R,1)), c=0.0, checks=true)
    QuadraticCost(Q,R,H,q,r,c, checks=checks)
end

function QuadraticCost(Q::Union{Diagonal{T,SVector{N,T}}, SMatrix{N,N,T}},
            R::Union{Diagonal{T,SVector{M,T}}, SMatrix{M,M,T}};
            H=(@SMatrix zeros(T,M,N)),
            q=(@SVector zeros(T,N)),
            r=(@SVector zeros(M)),
            c=0.0, checks=true) where {T,N,M}
    QuadraticCost(Q,R,H,q,r,c, checks=checks)
end

function QuadraticCost(Q,q,c; checks=true)
    QuadraticCost(Q,zeros(0,0),zeros(0,size(Q,1)),q,zeros(0),c,checks=checks)
end

"""
$(SIGNATURES)
Cost function of the form
``(x-x_f)^T Q (x_x_f) + u^T R u``
R must be positive definite, Q must be positive semidefinite
"""
function LQRCost(Q::AbstractArray, R::AbstractArray,
        xf::AbstractVector, uf=(@SVector zeros(size(R,1))); checks=true)
    H = @SMatrix zeros(size(R,1),size(Q,1))
    q = -Q*xf
    r = -R*uf
    c = 0.5*xf'*Q*xf + 0.5*uf'R*uf
    return QuadraticCost(Q, R, H, q, r, c, checks=checks)
end

"""
$(SIGNATURES)
Cost function of the form
``(x-x_f)^T Q (x-x_f)``
Q must be positive semidefinite
"""
function LQRCostTerminal(Qf::AbstractArray,xf::AbstractVector)
    qf = -Qf*xf
    cf = 0.5*xf'*Qf*xf
    return QuadraticCost(Qf,zeros(0,0),zeros(0,size(Qf,1)),qf,zeros(0),cf)
end

# Cost function methods
function stage_cost(cost::QuadraticCost, x::AbstractVector, u::AbstractVector)
    0.5*x'cost.Q*x + 0.5*u'*cost.R*u + cost.q'x + cost.r'u + cost.c #+ u'*cost.H*x
end

function stage_cost(cost::QuadraticCost, xN::AbstractVector{T}) where T
    0.5*xN'cost.Q*xN + cost.q'*xN + cost.c
end

function gradient!(E::AbstractExpansion, cost::QuadraticCost, x, u)
    E.x .= cost.Q*x + cost.q + cost.H'*u
    E.u .= cost.R*u + cost.r + cost.H*x
    return nothing
    mul!(E.x, cost.Q, x)
    E.x .+= cost.q
    mul!(E.x, cost.H', u, 1.0, 1.0)

    mul!(E.u, cost.R, u)
    E.u .+= cost.r
    mul!(E.u, cost.H, x, 1.0, 1.0)
    return nothing
end

function hessian!(E::AbstractExpansion, cost::QuadraticCost, x, u)
    E.xx .= cost.Q
    E.uu .= cost.R
    E.ux .= cost.H
    return nothing
end

# Additional Methods
function Base.show(io::IO, cost::QuadraticCost)
    print(io, "QuadraticCost{...}")
end

import Base: +
"Add two Quadratic costs of the same size"
function +(c1::QuadraticCost, c2::QuadraticCost)
    @assert state_dim(c1) == state_dim(c2)
    @assert control_dim(c1) == control_dim(c2)
    QuadraticCost(c1.Q + c2.Q, c1.R + c2.R, c1.H + c2.H,
                  c1.q + c2.q, c1.r + c2.r, c1.c + c2.c)
end

function copy(cost::QuadraticCost)
    return QuadraticCost(copy(cost.Q), copy(cost.R), copy(cost.H), copy(cost.q), copy(cost.r), copy(cost.c))
end

############################################################################################
#                                 QUADRATIC QUATERNION COST FUNCTION
############################################################################################

struct QuadraticQuatCost{T,N,M,N4} <: CostFunction
    Q::Diagonal{T,SVector{N,T}}
    R::Diagonal{T,SVector{M,T}}
    q::SVector{N,T}
    r::SVector{M,T}
    c::T
    w::T
    q_ref::SVector{4,T}
    q_ind::SVector{4,Int}
    Iq::SMatrix{N,4,T,N4}
    function QuadraticQuatCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}},
            q::SVector{N,T}, r::SVector{M,T}, c::T, w::T,
            q_ref::SVector{4,T}, q_ind::SVector{4,Int}) where {T,N,M}
        Iq = @MMatrix zeros(N,4)
        for i = 1:4
            Iq[q_ind[i],i] = 1
        end
        Iq = SMatrix{N,4}(Iq)
        return new{T,N,M,N*4}(Q, R, q, r, c, w, q_ref, q_ind, Iq)
    end
end

state_dim(::QuadraticQuatCost{T,N,M}) where {T,N,M} = N
control_dim(::QuadraticQuatCost{T,N,M}) where {T,N,M} = M

function QuadraticQuatCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}};
        q=(@SVector zeros(N)), r=(@SVector zeros(M)), c=zero(T), w=one(T),
        q_ref=(@SVector [1.0,0,0,0]), q_ind=(@SVector [4,5,6,7])) where {T,N,M}
    QuadraticQuatCost(Q, R, q, r, c, q_ref, q_ind)
end

function stage_cost(cost::QuadraticQuatCost, x::SVector, u::SVector)
    stage_cost(cost, x) + 0.5*u'cost.R*u + cost.r'u
end

function stage_cost(cost::QuadraticQuatCost, x::SVector)
    J = 0.5*x'cost.Q*x + cost.q'x + cost.c
    q = x[cost.q_ind]
    dq = cost.q_ref'q
    J += cost.w*min(1+dq, 1-dq)
end

function gradient(cost::QuadraticQuatCost{T,N,M}, x::SVector, u::SVector) where {T,N,M}
    Qx = cost.Q*x + cost.q
    q = x[cost.q_ind]
    dq = cost.q_ref'q
    if dq < 0
        Qx += cost.w*cost.Iq*cost.q_ref
    else
        Qx -= cost.w*cost.Iq*cost.q_ref
    end
    Qu = cost.R*u + cost.r
    return Qx, Qu
end

function hessian(cost::QuadraticQuatCost, x::SVector{N}, u::SVector{M}) where {N,M}
    Qxx = cost.Q
    Quu = cost.R
    Qux = @SMatrix zeros(M,N)
    return Qxx, Quu, Qux
end

function QuatLQRCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}}, xf,
        uf=(@SVector zeros(M)); w=one(T), quat_ind=(@SVector [4,5,6,7])) where {T,N,M}
    r = -R*uf
    q = -Q*xf
    c = 0.5*xf'Q*xf + 0.5*uf'R*uf
    q_ref = xf[quat_ind]
    return QuadraticQuatCost(Q, R, q, r, c, w, q_ref, quat_ind)
end

function change_dimension(cost::QuadraticQuatCost, n, m)
    n0,m0 = state_dim(cost), control_dim(cost)
    Q_diag = diag(cost.Q)
    R_diag = diag(cost.R)
    q = cost.q
    r = cost.r
    if n0 != n
        dn = n - n0  # assumes n > n0
        pad = @SVector zeros(dn)
        Q_diag = [Q_diag; pad]
        q = [q; pad]
    end
    if m0 != m
        dm = m - m0  # assumes m > m0
        pad = @SVector zeros(dm)
        R_diag = [R_diag; pad]
        r = [r; pad]
    end
    QuadraticQuatCost(Diagonal(Q_diag), Diagonal(R_diag), q, r, cost.c, cost.w,
        cost.q_ref, cost.q_ind)
end

function (+)(cost1::QuadraticQuatCost, cost2::QuadraticCost)
    @assert state_dim(cost1) == state_dim(cost2)
    @assert control_dim(cost1) == control_dim(cost2)
    @assert norm(cost2.H) ≈ 0
    QuadraticQuatCost(cost1.Q + cost2.Q, cost1.R + cost2.R,
        cost1.q + cost2.q, cost1.r + cost2.r, cost1.c + cost2.c,
        cost1.w, cost1.q_ref, cost1.q_ind)
end

(+)(cost1::QuadraticCost, cost2::QuadraticQuatCost) = cost2 + cost1

struct ErrorQuadratic{Rot,N,M} <: CostFunction
    model::RigidBody{Rot}
    Q::Diagonal{Float64,SVector{12,Float64}}
    R::Diagonal{Float64,SVector{M,Float64}}
    r::SVector{M,Float64}
    c::Float64
    x_ref::SVector{N,Float64}
    q_ind::SVector{4,Int}
end

state_dim(::ErrorQuadratic{Rot,N,M}) where {Rot,N,M} = N
control_dim(::ErrorQuadratic{Rot,N,M}) where {Rot,N,M} = M

function ErrorQuadratic(model::RigidBody{Rot}, Q::Diagonal{T,<:SVector{12}},
        R::Diagonal{T,<:SVector{M}},
        x_ref::SVector{N}, u_ref=(@SVector zeros(T,M)); r=(@SVector zeros(T,M)), c=zero(T),
        q_ind=(@SVector [4,5,6,7])) where {T,N,M,Rot}
    r += -R*u_ref
    c += 0.5*u_ref'R*u_ref
    return ErrorQuadratic{Rot,N,M}(model, Q, R, r, c, x_ref, q_ind)
end

function stage_cost(cost::ErrorQuadratic, x::SVector)
    dx = state_diff(cost.model, x, cost.x_ref)
    return 0.5*dx'cost.Q*dx + cost.c
end

function stage_cost(cost::ErrorQuadratic, x::SVector, u::SVector)
    stage_cost(cost, x) + 0.5*u'cost.R*u + cost.r'u
end

function cost_expansion(cost::ErrorQuadratic{Rot}, model::AbstractModel,
        z::KnotPoint{T,N,M}, G) where {T,N,M,Rot<:UnitQuaternion}
    x,u = state(z), control(z)
    model = cost.model
    Q = cost.Q
    q = orientation(model, x)
    q_ref = orientation(model, cost.x_ref)
    dq = SVector(q_ref\q)
    err = state_diff(model, x, cost.x_ref)
    dx = @SVector [err[1],  err[2],  err[3],
                    dq[1],   dq[2],   dq[3],   dq[4],
                   err[7],  err[8],  err[9],
                   err[10], err[11], err[12]]
    G = state_diff_jacobian(model, dx) # n × dn

    # Gradient
    dmap = inverse_map_jacobian(model, dx) # dn × n
    Qx = G'dmap'Q*err
    Qu = cost.R*u

    # Hessian
    ∇jac = inverse_map_∇jacobian(model, dx, Q*err)
    Qxx = G'dmap'Q*dmap*G + G'∇jac*G + ∇²differential(model, x, dmap'Q*err)
    Quu = cost.R
    Qux = @SMatrix zeros(M,N-1)
    return Qxx, Quu, Qux, Qx, Qu
end

function cost_expansion(cost::ErrorQuadratic, model::AbstractModel,
        z::KnotPoint{T,N,M}, G) where {T,N,M}
    x,u = state(z), control(z)
    model = cost.model
    q = orientation(model, x)
    q_ref = orientation(model, cost.x_ref)
    err = state_diff(model, x, cost.x_ref)
    dx = err
    G = state_diff_jacobian(model, dx) # n × n

    # Gradient
    dmap = inverse_map_jacobian(model, dx) # n × n
    Qx = G'dmap'cost.Q*err
    Qu = cost.R*u + cost.r

    # Hessian
    Qxx = G'dmap'cost.Q*dmap*G
    Quu = cost.R
    Qux = @SMatrix zeros(M,N)
    return Qxx, Quu, Qux, Qx, Qu
end

function change_dimension(cost::ErrorQuadratic, n, m)
    n0,m0 = state_dim(cost), control_dim(cost)
    Q_diag = diag(cost.Q)
    R_diag = diag(cost.R)
    r = cost.r
    if n0 != n
        dn = n - n0  # assumes n > n0
        pad = @SVector zeros(dn) # assume the new states don't have quaternions
        Q_diag = [Q_diag; pad]
    end
    if m0 != m
        dm = m - m0  # assumes m > m0
        pad = @SVector zeros(dm)
        R_diag = [R_diag; pad]
        r = [r; pad]
    end
    ErrorQuadratic(cost.model, Diagonal(Q_diag), Diagonal(R_diag), r, cost.c,
        cost.x_ref, cost.q_ind)
end

function (+)(cost1::ErrorQuadratic, cost2::QuadraticCost)
    @assert control_dim(cost1) == control_dim(cost2)
    @assert norm(cost2.H) ≈ 0
    @assert norm(cost2.q) ≈ 0
    if state_dim(cost2) == 13
        rm_quat = @SVector [1,2,3,4,5,6,8,9,10,11,12,13]
        Q2 = Diagonal(diag(cost2.Q)[rm_quat])
    else
        Q2 = cost2.Q
    end
    ErrorQuadratic(cost1.model, cost1.Q + Q2, cost1.R + cost2.R,
        cost1.r + cost2.r, cost1.c + cost2.c,
        cost1.x_ref, cost1.q_ind)
end

(+)(cost1::QuadraticCost, cost2::ErrorQuadratic) = cost2 + cost1

struct SatDiffCost{Rot} <: CostFunction
    model::RigidBody{Rot}
    Q1::Diagonal{Float64,SVector{3,Float64}}
    Q2::Diagonal{Float64,SVector{3,Float64}}
    R::Diagonal{Float64,SVector{3,Float64}}
    q_ref::SVector{4,Float64}
    ω_ref::SVector{3,Float64}
end

function stage_cost(cost::SatDiffCost{Rot}, x::SVector) where Rot
    ω = @SVector [x[1],x[2],x[3]]
    q = Dynamics.orientation(cost.model, x)
    q0 = Rot(UnitQuaternion(cost.q_ref))

    dω = ω - cost.ω_ref
    dq = q ⊖ q0
    return 0.5*(dq'cost.Q2*dq + dω'cost.Q1*dω)
end

function stage_cost(cost::SatDiffCost, x::SVector, u::SVector)
    J = stage_cost(cost, x) + 0.5*u'cost.R*u
end

function cost_expansion(cost::SatDiffCost{Rot}, model::AbstractModel,
        z::KnotPoint{T,N,M}, G) where {T,N,M,Rot<:UnitQuaternion}
    x,u = state(z), control(z)
    Q = cost.Q2  # cost for quaternion
    ω = @SVector [x[1],x[2],x[3]]
    q = @SVector [x[1],x[2],x[3],x[4]]
    q = Dynamics.orientation(cost.model, x)
    q0 = Rot(UnitQuaternion(cost.q_ref))

    dω = ω - cost.ω_ref
    dq = q0\q
    err = q ⊖ q0
    G = ∇differential(dq) #Lmult(dq)*Vmat()'

    # Gradient
    Qω = cost.Q1*ω
    dmap = inverse_map_jacobian(dq) #jacobian(CayleyMap,dq)
    Qq = G'dmap'Q*err
    Qx = [Qω; Qq]
    Qu = cost.R*u

    # Hessian
    Qωω = cost.Q1
    ∇jac = inverse_map_∇jacobian(dq, Q*err)
    Qqq = G'dmap'Q*dmap*G + G'∇jac*G
    Qxx = @SMatrix [
        Qωω[1,1] 0 0 0 0 0;
        0 Qωω[2,2] 0 0 0 0;
        0 0 Qωω[3,3] 0 0 0;
        0 0 0 Qqq[1,1] Qqq[1,2] Qqq[1,3];
        0 0 0 Qqq[2,1] Qqq[2,2] Qqq[2,3];
        0 0 0 Qqq[3,1] Qqq[3,2] Qqq[3,3];
    ]
    Quu = cost.R
    Qux = @SMatrix zeros(M,N-1)
    return Qxx, Quu, Qux, Qx, Qu
end

function cost_expansion(cost::SatDiffCost{Rot}, model::AbstractModel,
        z::KnotPoint{T,N,M}, G) where {T,N,M,Rot}
    x,u = state(z), control(z)
    Q = cost.Q2  # cost for quaternion
    ω = @SVector [x[1],x[2],x[3]]
    q = @SVector [x[1],x[2],x[3],x[4]]
    q = Dynamics.orientation(cost.model, x)
    q0 = Rot(UnitQuaternion(cost.q_ref))

    dω = ω - cost.ω_ref
    dq = q0\q
    err = q ⊖ q0
    G = ∇differential(dq) #Lmult(dq)*Vmat()'

    # Gradient
    Qω = cost.Q1*ω
    Qq = G'Q*err
    Qx = [Qω; Qq]
    Qu = cost.R*u

    # Hessian
    Qωω = cost.Q1
    ∇jac = ∇²differential(dq, Q*err)
    Qqq = G'Q*G + ∇jac
    Qxx = @SMatrix [
        Qωω[1,1] 0 0 0 0 0;
        0 Qωω[2,2] 0 0 0 0;
        0 0 Qωω[3,3] 0 0 0;
        0 0 0 Qqq[1,1] Qqq[1,2] Qqq[1,3];
        0 0 0 Qqq[2,1] Qqq[2,2] Qqq[2,3];
        0 0 0 Qqq[3,1] Qqq[3,2] Qqq[3,3];
    ]
    Quu = cost.R
    Qux = @SMatrix zeros(M,N)
    return Qxx, Quu, Qux, Qx, Qu
end

#
# """
# $(TYPEDEF)
# Cost function of the form
#     ℓf(xₙ) + ∫ ℓ(x,u) dt from 0 to tf
# """
# struct GenericCost <: CostFunction
#     ℓ::Function             # Stage cost
#     ℓf::Function            # Terminal cost
#     expansion::Function     # 2nd order Taylor Series Expansion of the form,  Q,R,H,q,r = expansion(x,u)
#     n::Int                  #                                                     Qf,qf = expansion(xN)
#     m::Int
# end
#
# """
# $(SIGNATURES)
# Create a Generic Cost, specifying the gradient and hessian of the cost function analytically
#
# # Arguments
# * hess: multiple-dispatch function of the form,
#     Q,R,H = hess(x,u) with sizes (n,n), (m,m), (m,n)
#     Qf = hess(xN) with size (n,n)
# * grad: multiple-dispatch function of the form,
#     q,r = grad(x,u) with sizes (n,), (m,)
#     qf = grad(x,u) with size (n,)
#
# """
# function GenericCost(ℓ::Function, ℓf::Function, grad::Function, hess::Function, n::Int, m::Int)
#     @warn "Use GenericCost with caution. It is untested and not likely to work"
#     function expansion(x::Vector{T},u::Vector{T}) where T
#         Q,R,H = hess(x,u)
#         q,r = grad(x,u)
#         return Q,R,H,q,r
#     end
#     expansion(xN) = hess(xN), grad(xN)
#     GenericCost(ℓ,ℓf, expansion, n,m)
# end
#
# """ $(TYPEDEF)
# This is an experimental cost function type that allows a cost function to be evaluated
#     on only a portion of the state and control. Right now, the implementation assumes
#     there is no coupling between the state and control.
#
# It should be noted that for `QuadraticCost`s it is likely more efficient to simply make a new
#     `QuadraticCost` that has zeros in the right places, and then add the cost functions together.
#
# # Constructor
# ```julia
# IndexedCost(cost, ix::UnitRange, iu::UnitRange)
# ```
# """
# struct IndexedCost{iX,iU,C} <: CostFunction
#     cost::C
# end
#
# function IndexedCost(cost::C, ix::UnitRange, iu::UnitRange) where C<:CostFunction
#     if C <: QuadraticCost
#         if norm(cost.H) != 0
#             throw(ErrorException("IndexedCost of functions with x-u coupling not implemented"))
#         end
#     else
#         @warn "IndexedCost will only work for costs without x-u coupling (Qux = 0)"
#     end
#     IndexedCost{ix,iu,C}(cost)
# end
#
# @generated function stage_cost(costfun::IndexedCost{iX,iU}, x::SVector{N}, u::SVector{M}) where {iX,iU,N,M}
#     ix = SVector{length(iX)}(iX)
#     iu = SVector{length(iU)}(iU)
#     quote
#         x0 = x[$ix]
#         u0 = u[$iu]
#         stage_cost(costfun.cost, x0, u0)
#     end
# end
#
# @generated function stage_cost(costfun::IndexedCost{iX,iU}, x::SVector{N}) where {iX,iU,N}
#     ix = SVector{length(iX)}(iX)
#     quote
#         x0 = x[$ix]
#         stage_cost(costfun.cost, x0)
#     end
# end
#
# @generated function gradient(costfun::IndexedCost{iX,iU}, x::SVector{N}, u::SVector{M}) where {iX,iU,N,M}
#     l1x = iX[1] - 1
#     l2x = N-iX[end]
#     l1u = iU[1] - 1
#     l2u = M-iU[end]
#     quote
#         x = x[$iX]
#         u = u[$iU]
#         Qx, Qu = gradient(costfun.cost, x, u)
#         Qx = [@SVector zeros($l1x); Qx; @SVector zeros($l2x)]
#         Qu = [@SVector zeros($l1u); Qu; @SVector zeros($l2u)]
#         return Qx, Qu
#     end
# end
#
#
# @generated function hessian(costfun::IndexedCost{iX,iU}, x::SVector{N}, u::SVector{M}) where {iX,iU,N,M}
#     l1x = iX[1] - 1
#     l2x = N-iX[end]
#     l1u = iU[1] - 1
#     l2u = M-iU[end]
#     quote
#         x = x[$iX]
#         u = u[$iU]
#         Qxx, Quu, Qux  = hessian(costfun.cost, x, u)
#         Qxx1 = Diagonal(@SVector zeros($l1x))
#         Qxx2 = Diagonal(@SVector zeros($l2x))
#         Quu1 = Diagonal(@SVector zeros($l1u))
#         Quu2 = Diagonal(@SVector zeros($l2u))
#
#         Qxx = blockdiag(Qxx1, Qxx, Qxx2)
#         Quu = blockdiag(Quu1, Quu, Quu2)
#         Qux = @SMatrix zeros(M,N)
#         Qxx, Quu, Qux
#     end
# end
#
function SparseArrays.blockdiag(Qs::Vararg{<:Diagonal})
    Diagonal(vcat(diag.(Qs)...))
end

function SparseArrays.blockdiag(Qs::Vararg{<:AbstractMatrix})
    # WARNING: this is slow and is only included as a fallback
    cat(Qs...,dims=(1,2))
end
#
# function change_dimension(cost::CostFunction,n,m)
#     n0,m0 = state_dim(cost), control_dim(cost)
#     ix = 1:n0
#     iu = 1:m0
#     IndexedCost(cost, ix, iu)
# end
#
function change_dimension(cost::QuadraticCost, n, m)
    n0,m0 = state_dim(cost), control_dim(cost)
    @assert n >= n0
    @assert m >= m0

    ix = 1:n0
    iu = 1:m0

    Q_ = Diagonal(@SVector zeros(n-n0))
    R_ = Diagonal(@SVector zeros(m-m0))
    H1 = @SMatrix zeros(m0, n-n0)
    H2 = @SMatrix zeros(m-m0, n)
    q_ = @SVector zeros(n-n0)
    r_ = @SVector zeros(m-m0)
    c = cost.c

    # Insert old values
    Q = blockdiag(cost.Q, Q_)
    R = blockdiag(cost.R, R_)
    H = [cost.H H1]
    H = [H; H2]
    q = [cost.q; q_]
    r = [cost.r; r_]
    QuadraticCost(Q,R,H,q,r,c,checks=false)
end


text:
InitialValues.@def push!! [x]

# # Custom definition for `monoid!!(_, Init(monoid!!))`
#
# Using the default approach `InitialValues.@def_monoid monoid!!` was
# problematic because defining `monoid!!(::CustomType, x)` introduced
# method ambiguities.  It could be solved semi-automatically by
# `@disambiguate append!! CustomType` but it is very ugly and very
# problematic for the extensibility of the method.
#
# Furthermore, the definition
#
#     append!!(::InitAppend!!, src) = src
#
# that is generated by `@def_monoid` is not appropriate if `src`
# should be consumed by the time `append!!` is called (e.g.,
# https://discourse.julialang.org/t/38845/3).  This is mostly the case
# since the returned value of `monoid!!(init, src)` would be used as
# the first argument (i.e., it would be mutated) in the "next
# iteration" in reduction code idioms.  That is to say, `monoid!!`
# should not assume the ownership of the second argument.
#
# The method ambiguity problem is solved by introducing a "dispatch
# pipeline"; i.e., a `CustomType` implementer can overload
# `__monoid!!__` instead of `monoid!!`.
#
# The memory ownership problem is solved by defining InitialValues.jl
# interface manually.  For example, `append!!(::InitAppend!!, src)`
# now calls `copymutable(src)`.
#
# ## Alternative approach
#
# For `append!!`, an alternative approach considered is to create a
# custom `EmptyCollection` type that behaves like "an" identity
# element of `append!!`.  This may be possible by allowing
# `InitialValues.Init(::typeof(monoid!!))` overload.  This approach
# has some nice properties:
#
# * No need to define `append!!(::Any, ::InitType)` as it would work
#   via the default implementation of `append!!` based on iterator.
#   (The flip-side is that a bug could be hidden if the init object is
#   fed to other functions like `collect`. But this is not a serious
#   problem.)
#
# * By using the trait mechanism and defining `NoBang.append`,
#   `append!!` would behave as expected without any "surface"
#   dispatches.
#
# A *major downside* of this approach is that the fold code does not
# have a canonical way to check if `monoid!!` is ever called.  If
# overloading `InitialValues.Init` were supported, the fold
# implementers cannot simply call `acc isa InitialValue` to check if
# `acc = monoid!!(acc, x)` was called at least once.

const InitAppend!! = InitialValues.GenericInitialValue{typeof(append!!)}
append!!(::InitAppend!!, src) = copyappendable(src)
append!!(dest, ::InitAppend!!) = dest
append!!(dest::InitAppend!!, ::InitAppend!!) = dest
InitialValues.hasinitialvalue(::Type{typeof(append!!)}) = true

copyappendable(src) = Base.copymutable(src)

const InitMergeWith!!{F} = InitialValues.GenericInitialValue{MergeWith!!{F}}
(f::MergeWith!!{F})(dest::InitMergeWith!!{F}, src) where {F} = copymergeable(src)
(f::MergeWith!!{F})(dest, ::InitMergeWith!!{F}) where {F} = dest
(f::MergeWith!!{F})(dest::InitMergeWith!!{F}, src::InitMergeWith!!{F}) where {F} = dest
InitialValues.hasinitialvalue(::Type{MergeWith!!{F}}) where {F} = true

const InitMerge!! = InitialValues.GenericInitialValue{typeof(merge!!)}
merge!!(dest::InitMerge!!, src) = copymergeable(src)
merge!!(dest, ::InitMerge!!) = dest
merge!!(dest::InitMerge!!, src::InitMerge!!) = dest
merge!!(::Base.Callable, dest::InitMerge!!) = dest  # disambiguation
InitialValues.hasinitialvalue(::Type{typeof(merge!!)}) = true

copymergeable(src) = Dict(src)
copymergeable(src::NamedTuple) = src

const InitUnion!! = InitialValues.GenericInitialValue{typeof(union!!)}
union!!(dest::InitUnion!!, src) = copyunionable(src)
union!!(dest, ::InitUnion!!) = dest
union!!(dest::InitUnion!!, src::InitUnion!!) = dest
InitialValues.hasinitialvalue(::Type{typeof(union!!)}) = true

copyunionable(src::AbstractVector) = Base.copymutable(src)
copyunionable(src) = Set(src)


text:
function variable_converter_filter_voltage(pm::_PM.AbstractBFModel; kwargs...)
    variable_converter_filter_voltage_magnitude_sqr(pm; kwargs...)
    variable_conv_transformer_current_sqr(pm; kwargs...)
end

function variable_converter_internal_voltage(pm::_PM.AbstractBFModel; kwargs...)
    variable_converter_internal_voltage_magnitude_sqr(pm; kwargs...)
    variable_conv_reactor_current_sqr(pm; kwargs...)
end
"""
Converter transformer constraints
```
p_tf_fr + ptf_to ==  rtf*itf
q_tf_fr + qtf_to ==  xtf*itf
p_tf_fr^2 + qtf_fr^2 <= w/tm^2 * itf
wf == w/tm^2 -2*(rtf*ptf_fr + xtf*qtf_fr) + (rtf^2 + xtf^2)*itf
```
"""
function constraint_conv_transformer(pm::_PM.AbstractBFQPModel, n::Int,  i::Int, rtf, xtf, acbus, tm, transformer)
    w = _PM.var(pm, n,  :w, acbus)
    itf = _PM.var(pm, n,  :itf_sq, i)
    wf = _PM.var(pm, n,  :wf_ac, i)

    ptf_fr = _PM.var(pm, n,  :pconv_tf_fr, i)
    qtf_fr = _PM.var(pm, n,  :qconv_tf_fr, i)
    ptf_to = _PM.var(pm, n,  :pconv_tf_to, i)
    qtf_to = _PM.var(pm, n,  :qconv_tf_to, i)

    if transformer
        JuMP.@constraint(pm.model,   ptf_fr + ptf_to ==  rtf*itf)
        JuMP.@constraint(pm.model,   qtf_fr + qtf_to ==  xtf*itf)
        JuMP.@constraint(pm.model,   ptf_fr^2 + qtf_fr^2 <= w/tm^2 * itf)
        JuMP.@constraint(pm.model,   wf == w/tm^2 -2*(rtf*ptf_fr + xtf*qtf_fr) + (rtf^2 + xtf^2)*itf)
    else
        JuMP.@constraint(pm.model, ptf_fr + ptf_to == 0)
        JuMP.@constraint(pm.model, qtf_fr + qtf_to == 0)
        JuMP.@constraint(pm.model, wf == w )
    end
end

function constraint_conv_transformer(pm::_PM.AbstractBFConicModel, n::Int,  i::Int, rtf, xtf, acbus, tm, transformer)
    w = _PM.var(pm, n,  :w, acbus)
    itf = _PM.var(pm, n,  :itf_sq, i)
    wf = _PM.var(pm, n,  :wf_ac, i)

    ptf_fr = _PM.var(pm, n,  :pconv_tf_fr, i)
    qtf_fr = _PM.var(pm, n,  :qconv_tf_fr, i)
    ptf_to = _PM.var(pm, n,  :pconv_tf_to, i)
    qtf_to = _PM.var(pm, n,  :qconv_tf_to, i)

    if transformer
        JuMP.@constraint(pm.model,   ptf_fr + ptf_to ==  rtf*itf)
        JuMP.@constraint(pm.model,   qtf_fr + qtf_to ==  xtf*itf)
        JuMP.@constraint(pm.model,   [w/(sqrt(2)*tm), itf/(sqrt(2)*tm), ptf_fr, qtf_fr] in JuMP.RotatedSecondOrderCone())
        JuMP.@constraint(pm.model,   wf == w/tm^2 -2*(rtf*ptf_fr + xtf*qtf_fr) + (rtf^2 + xtf^2)*itf)
    else
        JuMP.@constraint(pm.model, ptf_fr + ptf_to == 0)
        JuMP.@constraint(pm.model, qtf_fr + qtf_to == 0)
        JuMP.@constraint(pm.model, wf == w/tm^2 )
    end
end

"""
Converter reactor constraints
```
p_pr_fr + ppr_to == rc*ipr
q_pr_fr + qpr_to == xc*ipr
p_pr_fr^2 + qpr_fr^2 <= wf * ipr
wc == wf -2*(rc*ppr_fr + xc*qpr_fr) + (rc^2 + xc^2)*ipr
```
"""
function constraint_conv_reactor(pm::_PM.AbstractBFQPModel, n::Int,  i::Int, rc, xc, reactor)
    pconv_ac = _PM.var(pm, n,  :pconv_ac, i)
    qconv_ac = _PM.var(pm, n,  :qconv_ac, i)
    ppr_to = - pconv_ac
    qpr_to = - qconv_ac
    ppr_fr = _PM.var(pm, n,  :pconv_pr_fr, i)
    qpr_fr = _PM.var(pm, n,  :qconv_pr_fr, i)

    wf = _PM.var(pm, n,  :wf_ac, i)
    ipr = _PM.var(pm, n,  :irc_sq, i)
    wc = _PM.var(pm, n,  :wc_ac, i)

    if reactor
        JuMP.@constraint(pm.model, ppr_fr + ppr_to == rc*ipr)
        JuMP.@constraint(pm.model, qpr_fr + qpr_to == xc*ipr)
        JuMP.@constraint(pm.model, ppr_fr^2 + qpr_fr^2 <= wf * ipr)
        JuMP.@constraint(pm.model, wc == wf -2*(rc*ppr_fr + xc*qpr_fr) + (rc^2 + xc^2)*ipr)

    else
        JuMP.@constraint(pm.model, ppr_fr + ppr_to == 0)
        JuMP.@constraint(pm.model, qpr_fr + qpr_to == 0)
        JuMP.@constraint(pm.model, wc == wf)
    end
end

function constraint_conv_reactor(pm::_PM.AbstractBFConicModel, n::Int,  i::Int, rc, xc, reactor)
    pconv_ac = _PM.var(pm, n,  :pconv_ac, i)
    qconv_ac = _PM.var(pm, n,  :qconv_ac, i)
    ppr_to = - pconv_ac
    qpr_to = - qconv_ac
    ppr_fr = _PM.var(pm, n,  :pconv_pr_fr, i)
    qpr_fr = _PM.var(pm, n,  :qconv_pr_fr, i)

    wf = _PM.var(pm, n,  :wf_ac, i)
    ipr = _PM.var(pm, n,  :irc_sq, i)
    wc = _PM.var(pm, n,  :wc_ac, i)

    if reactor
        JuMP.@constraint(pm.model, ppr_fr + ppr_to == rc*ipr)
        JuMP.@constraint(pm.model, qpr_fr + qpr_to == xc*ipr)
        JuMP.@constraint(pm.model, [wf/sqrt(2), ipr/sqrt(2), ppr_fr, qpr_fr] in JuMP.RotatedSecondOrderCone())
        JuMP.@constraint(pm.model, wc == wf -2*(rc*ppr_fr + xc*qpr_fr) + (rc^2 + xc^2)*ipr)

    else
        JuMP.@constraint(pm.model, ppr_fr + ppr_to == 0)
        JuMP.@constraint(pm.model, qpr_fr + qpr_to == 0)
        JuMP.@constraint(pm.model, wc == wf)
    end
end
"""
Links converter power & current
```
pconv_ac[i]^2 + pconv_dc[i]^2 <= wc[i] * iconv_ac_sq[i]
pconv_ac[i]^2 + pconv_dc[i]^2 <= (Umax)^2 * (iconv_ac[i])^2
```
"""
function constraint_converter_current(pm::_PM.AbstractBFQPModel, n::Int,  i::Int, Umax, Imax)
    wc = _PM.var(pm, n,  :wc_ac, i)
    pconv_ac = _PM.var(pm, n,  :pconv_ac, i)
    qconv_ac = _PM.var(pm, n,  :qconv_ac, i)
    iconv = _PM.var(pm, n,  :iconv_ac, i)
    iconv_sq = _PM.var(pm, n,  :iconv_ac_sq, i)

    JuMP.@constraint(pm.model, pconv_ac^2 + qconv_ac^2 <=  wc * iconv_sq)
    JuMP.@constraint(pm.model, pconv_ac^2 + qconv_ac^2 <= (Umax)^2 * iconv^2)
    JuMP.@constraint(pm.model, iconv^2 <= iconv_sq)
    JuMP.@constraint(pm.model, iconv_sq <= iconv*Imax)
end

function constraint_converter_current(pm::_PM.AbstractBFConicModel, n::Int,  i::Int, Umax, Imax)
    wc = _PM.var(pm, n,  :wc_ac, i)
    pconv_ac = _PM.var(pm, n,  :pconv_ac, i)
    qconv_ac = _PM.var(pm, n,  :qconv_ac, i)
    iconv = _PM.var(pm, n,  :iconv_ac, i)
    iconv_sq = _PM.var(pm, n,  :iconv_ac_sq, i)

    JuMP.@constraint(pm.model, [wc/sqrt(2), iconv_sq/sqrt(2), pconv_ac, qconv_ac] in JuMP.RotatedSecondOrderCone())
    JuMP.@constraint(pm.model, [Umax * iconv/sqrt(2), Umax * iconv/sqrt(2), pconv_ac, qconv_ac] in JuMP.RotatedSecondOrderCone())
    JuMP.@constraint(pm.model, [iconv_sq, iconv/sqrt(2), iconv/sqrt(2)] in JuMP.RotatedSecondOrderCone())
    JuMP.@constraint(pm.model, iconv_sq <= iconv*Imax)
end

############### TNEP Constraints #######################################
function variable_converter_filter_voltage_ne(pm::_PM.AbstractBFModel; kwargs...)
    variable_converter_filter_voltage_magnitude_sqr_ne(pm; kwargs...)
    variable_conv_transformer_current_sqr_ne(pm; kwargs...)
end

function variable_converter_internal_voltage_ne(pm::_PM.AbstractBFModel; kwargs...)
    variable_converter_internal_voltage_magnitude_sqr_ne(pm; kwargs...)
    variable_conv_reactor_current_sqr_ne(pm; kwargs...)
end

"""
Converter transformer constraints

```
p_tf_fr + ptf_to ==  rtf*itf
q_tf_fr + qtf_to ==  xtf*itf
p_tf_fr^2 + qtf_fr^2 <= w/tm^2 * itf
wf == w/tm^2 -2*(rtf*ptf_fr + xtf*qtf_fr) + (rtf^2 + xtf^2)*itf
```
"""
function constraint_conv_transformer_ne(pm::_PM.AbstractBFQPModel, n::Int, i::Int, rtf, xtf, acbus, tm, transformer)
    w = _PM.var(pm, n, :w, acbus)
    itf = _PM.var(pm, n, :itf_sq_ne, i)
    wf = _PM.var(pm, n, :wf_ac_ne, i)
    w_du = _PM.var(pm, n, :w_du, i)
    ptf_fr = _PM.var(pm, n, :pconv_tf_fr_ne, i)
    qtf_fr = _PM.var(pm, n, :qconv_tf_fr_ne, i)
    ptf_to = _PM.var(pm, n, :pconv_tf_to_ne, i)
    qtf_to = _PM.var(pm, n, :qconv_tf_to_ne, i)
    z = _PM.var(pm, n, :conv_ne)[i]

    if transformer
        JuMP.@constraint(pm.model,   ptf_fr + ptf_to ==  rtf*itf)
        JuMP.@constraint(pm.model,   qtf_fr + qtf_to ==  xtf*itf)
        JuMP.@constraint(pm.model,   ptf_fr^2 + qtf_fr^2 <= w_du/tm^2 * itf)
        JuMP.@constraint(pm.model,   wf  == w_du/tm^2 -2*(rtf*ptf_fr + xtf*qtf_fr) + (rtf^2 + xtf^2)*itf)
        _IM.relaxation_equality_on_off(pm.model, w, w_du, z)
        JuMP.@constraint(pm.model, w_du >= z*JuMP.lower_bound(w))
        JuMP.@constraint(pm.model, w_du <= z*JuMP.upper_bound(w))
    else
        JuMP.@constraint(pm.model, ptf_fr + ptf_to == 0)
        JuMP.@constraint(pm.model, qtf_fr + qtf_to == 0)
        JuMP.@constraint(pm.model, wf == w_du)
        _IM.relaxation_equality_on_off(pm.model, w, w_du, z)
        JuMP.@constraint(pm.model, w_du >= z*JuMP.lower_bound(w))
        JuMP.@constraint(pm.model, w_du <= z*JuMP.upper_bound(w))
    end
end

function constraint_conv_transformer_ne(pm::_PM.AbstractBFConicModel, n::Int, i::Int, rtf, xtf, acbus, tm, transformer)
    w = _PM.var(pm, n, :w, acbus)
    itf = _PM.var(pm, n, :itf_sq_ne, i)
    wf = _PM.var(pm, n, :wf_ac_ne, i)
    w_du = _PM.var(pm, n, :w_du, i)
    ptf_fr = _PM.var(pm, n, :pconv_tf_fr_ne, i)
    qtf_fr = _PM.var(pm, n, :qconv_tf_fr_ne, i)
    ptf_to = _PM.var(pm, n, :pconv_tf_to_ne, i)
    qtf_to = _PM.var(pm, n, :qconv_tf_to_ne, i)
    z = _PM.var(pm, n, :conv_ne)[i]

    if transformer
        JuMP.@constraint(pm.model,   ptf_fr + ptf_to ==  rtf*itf)
        JuMP.@constraint(pm.model,   qtf_fr + qtf_to ==  xtf*itf)
        JuMP.@constraint(pm.model,   [w_du/(sqrt(2)*tm), itf/(sqrt(2)*tm), ptf_fr, qtf_fr] in JuMP.RotatedSecondOrderCone())
        JuMP.@constraint(pm.model,   wf  == w_du/tm^2 -2*(rtf*ptf_fr + xtf*qtf_fr) + (rtf^2 + xtf^2)*itf)
        _IM.relaxation_equality_on_off(pm.model, w, w_du, z)
        JuMP.@constraint(pm.model, w_du >= z*JuMP.lower_bound(w))
        JuMP.@constraint(pm.model, w_du <= z*JuMP.upper_bound(w))
    else
        JuMP.@constraint(pm.model, ptf_fr + ptf_to == 0)
        JuMP.@constraint(pm.model, qtf_fr + qtf_to == 0)
        JuMP.@constraint(pm.model, wf == w_du )
        _IM.relaxation_equality_on_off(pm.model, w, w_du, z)
        JuMP.@constraint(pm.model, w_du >= z*JuMP.lower_bound(w))
        JuMP.@constraint(pm.model, w_du <= z*JuMP.upper_bound(w))
    end
end

"""
Converter reactor constraints

```
p_pr_fr + ppr_to == rc*ipr
q_pr_fr + qpr_to == xc*ipr
p_pr_fr^2 + qpr_fr^2 <= wf * ipr
wc == wf -2*(rc*ppr_fr + xc*qpr_fr) + (rc^2 + xc^2)*ipr
```
"""
function constraint_conv_reactor_ne(pm::_PM.AbstractBFQPModel, n::Int, i::Int, rc, xc, reactor)
    pconv_ac = _PM.var(pm, n, :pconv_ac_ne, i)
    qconv_ac = _PM.var(pm, n, :qconv_ac_ne, i)
    ppr_to = - pconv_ac
    qpr_to = - qconv_ac
    ppr_fr = _PM.var(pm, n, :pconv_pr_fr_ne, i)
    qpr_fr = _PM.var(pm, n, :qconv_pr_fr_ne, i)

    wf = _PM.var(pm, n, :wf_ac_ne, i)
    ipr = _PM.var(pm, n, :iconv_ac_sq_ne, i)
    wc = _PM.var(pm, n, :wc_ac_ne, i)

    if reactor
        JuMP.@constraint(pm.model, ppr_fr + ppr_to == rc*ipr)
        JuMP.@constraint(pm.model, qpr_fr + qpr_to == xc*ipr)
        JuMP.@constraint(pm.model, ppr_fr^2 + qpr_fr^2 <= wf * ipr)
        JuMP.@constraint(pm.model, wc  == wf  -2*(rc*ppr_fr + xc*qpr_fr) + (rc^2 + xc^2)*ipr)
    else
        JuMP.@constraint(pm.model, ppr_fr + ppr_to == 0)
        JuMP.@constraint(pm.model, qpr_fr + qpr_to == 0)
        JuMP.@constraint(pm.model, wc  == wf )
    end
end

function constraint_conv_reactor_ne(pm::_PM.AbstractBFConicModel, n::Int, i::Int, rc, xc, reactor)
    pconv_ac = _PM.var(pm, n, :pconv_ac_ne, i)
    qconv_ac = _PM.var(pm, n, :qconv_ac_ne, i)
    ppr_to = - pconv_ac
    qpr_to = - qconv_ac
    ppr_fr = _PM.var(pm, n, :pconv_pr_fr_ne, i)
    qpr_fr = _PM.var(pm, n, :qconv_pr_fr_ne, i)

    wf = _PM.var(pm, n, :wf_ac_ne, i)
    ipr = _PM.var(pm, n, :irc_sq_ne, i)
    wc = _PM.var(pm, n, :wc_ac_ne, i)

    if reactor
        JuMP.@constraint(pm.model, ppr_fr + ppr_to == rc*ipr)
        JuMP.@constraint(pm.model, qpr_fr + qpr_to == xc*ipr)
        JuMP.@constraint(pm.model, [wf/sqrt(2), ipr/sqrt(2), ppr_fr, qpr_fr] in JuMP.RotatedSecondOrderCone())
        JuMP.@constraint(pm.model, wc  == wf  -2*(rc*ppr_fr + xc*qpr_fr) + (rc^2 + xc^2)*ipr)
    else
        JuMP.@constraint(pm.model, ppr_fr + ppr_to == 0)
        JuMP.@constraint(pm.model, qpr_fr + qpr_to == 0)
        JuMP.@constraint(pm.model, wc  == wf )
    end
end

"""
Links converter power & current

```
pconv_ac[i]^2 + pconv_dc[i]^2 <= wc[i] * iconv_ac_sq[i]
pconv_ac[i]^2 + pconv_dc[i]^2 <= (Umax)^2 * (iconv_ac[i])^2
```
"""
function constraint_converter_current_ne(pm::_PM.AbstractBFQPModel, n::Int, i::Int, Umax, Imax)
    wc = _PM.var(pm, n, :wc_ac_ne, i)
    pconv_ac = _PM.var(pm, n, :pconv_ac_ne, i)
    qconv_ac = _PM.var(pm, n, :qconv_ac_ne, i)
    iconv = _PM.var(pm, n, :iconv_ac_ne, i)
    iconv_sq = _PM.var(pm, n, :iconv_ac_sq_ne, i)

    JuMP.@constraint(pm.model,      pconv_ac^2 + qconv_ac^2 <=  wc * iconv_sq)
    JuMP.@constraint(pm.model, pconv_ac^2 + qconv_ac^2 <= (Umax)^2 * iconv^2)
    JuMP.@constraint(pm.model, iconv^2 <= iconv_sq)

    JuMP.@constraint(pm.model, iconv_sq <= iconv*Imax)
end
function constraint_converter_current_ne(pm::_PM.AbstractBFConicModel, n::Int, i::Int, Umax, Imax)
    wc = _PM.var(pm, n, :wc_ac_ne, i)
    pconv_ac = _PM.var(pm, n, :pconv_ac_ne, i)
    qconv_ac = _PM.var(pm, n, :qconv_ac_ne, i)
    iconv = _PM.var(pm, n, :iconv_ac_ne, i)
    iconv_sq = _PM.var(pm, n, :iconv_ac_sq_ne, i)
    irc_sq = _PM.var(pm, n, :irc_sq_ne, i)

    JuMP.@constraint(pm.model, [wc/sqrt(2), iconv_sq/sqrt(2), pconv_ac, qconv_ac] in JuMP.RotatedSecondOrderCone())
    JuMP.@constraint(pm.model, [Umax * iconv/sqrt(2), Umax * iconv/sqrt(2), pconv_ac, qconv_ac] in JuMP.RotatedSecondOrderCone())

    JuMP.@constraint(pm.model, iconv_sq <= iconv*Imax)
    JuMP.@constraint(pm.model, iconv_sq == irc_sq)
end


text:
# Chapter 17. Multiple Dispatch
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#chap17

using Printf

# Methods
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#_methods
function printtime(mytime::MyTime)
    @printf("%02d:%02d:%02d\n", mytime.hour, mytime.minute, mytime.second)
end

# Exercise 17-1

function timetosecond(time::MyTime)
    seconds = time.hour*3600 + time.minute*60 + time.second
end

function secondtotime(totalsecond::Int64)
    (minute, second) = divrem(totalsecond, 60)
    (hour, minute) = divrem(minute, 60)
    return MyTime(hour, minute, second)
end

# Constructors
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#constructor
function MyTime(mytime::MyTime)
    MyTime(mytime.hour, mytime.minute, mytime.second)
end

function Base.show(io::IO, time::MyTime)
    @printf(io, "%02d:%02d:%02d", time.hour, time.minute, time.second)
end

import Base.+
function +(t1::MyTime, t2::MyTime)
    seconds = timetosecond(t1) + timetosecond(t2)
    secondtotime(seconds)
end


text:
#############################################################################
#Copyright (c) 2017 Charles Le Losq
#
#The MIT License (MIT)
#
#Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the #Software without restriction, including without limitation the rights to use, copy, #modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, #and to permit persons to whom the Software is furnished to do so, subject to the #following conditions:
#
#The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, #INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
#############################################################################

"""
	mass()
Return a DataFrame containing the mass of the elements from the IUPAS Periodic Table of the Elements, version 2007.

Call the elements by symbols, for instance db[:Si] to get the mass of silicium.

Forcused at the present mopment on the elements I use everyday... Any contribution is welcome!

"""

function mass()

	df = DataFrame()

	# column 1
	df[:H] = 1.00794
	df[:Li] = 6.941
	df[:Na] = 22.98976928
	df[:K] = 39.0983
	df[:Rb] = 85.4678
	df[:Cs] = 132.9054519
	df[:Fr] = 223

	# column 2
	df[:Be] = 9.012182
	df[:Mg] = 24.3050
	df[:Ca] = 40.078
	df[:Sr] = 87.62
	df[:Ba] = 137.327
	df[:Ra] = 226

	# column 4
	df[:Ti] = 47.867

	# column 6
	df[:Cr] = 51.9961

	#column 7
	df[:Mn] = 54.938045

	# column 8
	df[:Fe] = 55.845

	# column 10
	df[:Ni] = 58.6934

	# column 13
	df[:Al] = 26.9815386

	# column 14
	df[:Si] = 28.0855

	# column 16
	df[:O] = 15.9994

	return df

end


text:
module Day11

using Advent: save_input, read_charmatrix

const directions = map(CartesianIndex, ((0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (-1, -1), (1, -1), (-1, 1)))

function count_occupied(A, I, limit)
    R = CartesianIndices(A)
    Ifirst, Ilast = first(R), last(R)
    I1 = oneunit(I)
    count(directions) do dir
        @inbounds for scale in Base.OneTo(limit)
            coord = muladd(dir, scale, I)
            checkbounds(Bool, A, coord) || return false
            A[coord] == '#' && return true
            A[coord] == 'L' && return false
        end
        false
    end
end

# the maximum number of neighbors a passenger will tolerate within distance
function stabilize(curr, tolerance, distance=maximum(size(curr)))
    prev = nothing
    while prev != curr
        prev = deepcopy(curr)
        for I in CartesianIndices(curr)
            prev[I] == 'L' && count_occupied(prev, I, distance) == 0         && (curr[I] = '#')
            prev[I] == '#' && count_occupied(prev, I, distance) >= tolerance && (curr[I] = 'L')
        end
    end
    count(==('#'), curr)
end

solve_a(filename) = stabilize(read_charmatrix(filename), 4, 1)
solve_b(filename) = stabilize(read_charmatrix(filename), 5)

solve(filename = save_input(2020, 11)) = solve_a(filename), solve_b(filename)
solve_big() = solve("day11_big.in")

end

max_stars_repo_name: AnAverageHuman/competitive

max_issues_repo_name: AnAverageHuman/competitive

max_forks_repo_name: AnAverageHuman/competitive


text:
module ServiceTest

using Kroki.Service:
  DEFAULT_ENDPOINT,
  DockerComposeExecutionError,
  ENDPOINT,
  EXECUTE_DOCKER_COMPOSE,
  executeDockerCompose,
  setEndpoint!,
  start!,
  status,
  stop!,
  update!
using SimpleMock
using Test: @testset, @test, @test_logs, @test_skip

# Helper function to temporarily replace `EXECUTE_DOCKER_COMPOSE` with a
# `Mock`. Used to gain control over `docker-compose` behavior for local service
# instance management tests
function mockExecuteDockerCompose(f::Function, mock::Mock)
  EXECUTE_DOCKER_COMPOSE[] = mock
  f(mock)
  EXECUTE_DOCKER_COMPOSE[] = executeDockerCompose
end

@testset "Service" begin
  @testset "`ENDPOINT`" begin
    # Keep track of the original endpoint, so it can later be restored after
    # all tests have finished running. This should be robust with respect to
    # errors in the inner `TestSet`s
    original_endpoint = ENDPOINT[]

    @testset "`setEndpoint!`" begin
      @testset "with `endpoint` adjusts `ENDPOINT` to `endpoint`" begin
        expected_endpoint = "http://specific.endpoint.jl"

        setEndpoint!(expected_endpoint)

        @test ENDPOINT[] === expected_endpoint
      end

      @testset "without `endpoint` adjusts `ENDPOINT` to" begin
        @testset "`KROKI_ENDPOINT` environment variable if set" begin
          expected_endpoint = "http://from.environment.jl"
          withenv("KROKI_ENDPOINT" => expected_endpoint) do
            setEndpoint!()
            @test ENDPOINT[] === expected_endpoint
          end
        end

        @testset "`DEFAULT_ENDPOINT` otherwise" begin
          # Force the `KROKI_ENDPOINT` environment variable to be unset, as it
          # may be set outside the scope of the tests which would cause this
          # test to fail
          withenv("KROKI_ENDPOINT" => nothing) do
            setEndpoint!()
            @test ENDPOINT[] === DEFAULT_ENDPOINT
          end
        end
      end

      @testset "logs `ENDPOINT` updates only if changed" begin
        expected_endpoint = "http://logged.endpoint.jl"

        @test_logs (:info, "Setting Kroki service endpoint to $(expected_endpoint).") setEndpoint!(
          expected_endpoint,
        )
        @test_logs setEndpoint!(expected_endpoint)
      end
    end

    # Restore the original endpoint
    ENDPOINT[] = original_endpoint
  end

  @testset "local instance management" begin
    @static if Sys.which("docker-compose") !== nothing
      @testset "wraps `docker-compose` with local service definitions" begin
        status_report = executeDockerCompose("ps")

        # Verify known pieces of `docker-compose ps` output are returned
        @test occursin("Name", status_report)
        @test occursin("Command", status_report)
        @test occursin("State", status_report)
        @test occursin("Ports", status_report)
      end
    else
      @test_skip "Tests requiring `docker-compose` are skipped"
    end

    @testset "`executeDockerCompose` throws descriptive errors indicating" begin
      @testset "dependencies are missing" begin
        try
          # Breaking the path should ensure neither Docker, nor Docker Compose
          # can be found
          withenv(() -> executeDockerCompose("ps"), "PATH" => nothing)
        catch exception
          @test exception isa ErrorException

          error_message = sprint(showerror, exception)
          @test startswith(error_message, "Missing dependencies!")
          @test occursin("Docker Compose", error_message)
        end
      end

      # These tests require `docker-compose` to be available, as they
      # specifically test errors in the interface from Kroki.jl to it.
      #
      # This construct ensures the tests are always run if possible and clearly
      # marked broken on unsupported platforms
      @static if Sys.which("docker-compose") !== nothing
        @testset "to file an issue on indications of errors in Kroki.jl" begin
          try
            executeDockerCompose(["--non-existent", "ps"])
          catch exception
            @test exception isa DockerComposeExecutionError

            error_message = sprint(showerror, exception)
            @test occursin("docker-compose", error_message)
            @test occursin("file an issue", error_message)
            @test occursin("The reported error was", error_message)
          end
        end
      else
        @test_skip "Tests requiring `docker-compose` are skipped"
      end
    end

    @testset "`start!` starts Kroki service components" begin
      # The `start!` function adjusts the `ENDPOINT` so the started Kroki
      # service is immediately ready for use. This needs to be restored
      original_endpoint = ENDPOINT[]

      @testset "updating the `ENDPOINT` by default" begin
        mockExecuteDockerCompose(Mock()) do _executeDockerCompose
          # Set the `ENDPOINT` to a known value specific to the test, so it can
          # be detected that `start!` modifies it's value
          unstarted_instance_endpoint = "http://unstarted.instance.jl"
          setEndpoint!(unstarted_instance_endpoint)
          @test ENDPOINT[] === unstarted_instance_endpoint

          # The following explicitly uses `match_mode=:any` to prevent having
          # to specify log messages caused by changes to `ENDPOINT`
          returned = @test_logs (:info, "Starting Kroki service components.") match_mode =
            :any start!()

          # Ensure nothing gets returned from a call to `start!` instead of
          # `Process`es from the `docker-compose` execution
          @test returned === nothing
          @test called_with(_executeDockerCompose, ["up", "--detach"])

          @test ENDPOINT[] == "http://localhost:8000"
        end
      end

      @testset "while not updating `ENDPOINT` if asked" begin
        mockExecuteDockerCompose(Mock()) do _executeDockerCompose
          # Set the `ENDPOINT` to a known value specific to the test, so it can
          # be detected that `start!` modifies it's value
          non_modified_endpoint = "http://non.modified.endpoint.jl"
          setEndpoint!(non_modified_endpoint)
          @test ENDPOINT[] === non_modified_endpoint

          # Instruct `start!` to _not_ update the `ENDPOINT`
          start!(false)

          @test ENDPOINT[] === non_modified_endpoint
        end
      end

      ENDPOINT[] = original_endpoint
    end

    @testset "`status` reports individual Kroki service component status" begin
      status_mock = Mock((cmd::Vector{String}) -> (any(cmd .== "status=running") ? """
                                                                                core
                                                                                blockdiag
                                                                                mermaid
                                                                                """ : "bpmn"))

      mockExecuteDockerCompose(status_mock) do _executeDockerCompose
        service_statuses = status()

        @test service_statuses.core
        @test service_statuses.blockdiag
        @test !service_statuses.bpmn
        @test service_statuses.mermaid

        @test called_with(
          _executeDockerCompose,
          ["ps", "--filter", "status=running", "--services"],
        )
        @test called_with(
          _executeDockerCompose,
          ["ps", "--filter", "status=stopped", "--services"],
        )
      end
    end

    @testset "`stop!` stops any running Kroki service components" begin
      # The `stop!` function adjusts the `ENDPOINT` to ensure it doesn't point
      # to a stopped local Kroki service
      original_endpoint = ENDPOINT[]

      mockExecuteDockerCompose(Mock()) do _executeDockerCompose
        # Set the `ENDPOINT` to a known value specific to the test, so it can
        # be detected that `stop!` modifies it's value
        local_instance_endpoint = "http://local.instance.jl"
        setEndpoint!(local_instance_endpoint)
        @test ENDPOINT[] === local_instance_endpoint

        # The following explicitly uses `match_mode=:any` to prevent having to
        # specify log messages caused by changes to `ENDPOINT`
        returned = @test_logs (:info, "Stopping Kroki service components.") match_mode =
          :any stop!()

        # Ensure nothing gets returned from a call to `stop!` instead of
        # `Process`es from the `docker-compose` execution
        @test returned === nothing
        @test called_with(_executeDockerCompose, "stop")
        @test called_with(_executeDockerCompose, ["rm", "--force"])

        # The `ENDPOINT` should have been updated to the 'default' (which can
        # be either the `KROKI_ENDPOINT` environment variable of the
        # `DEFAULT_ENDPOINT`
        @test ENDPOINT[] !== local_instance_endpoint
      end

      @testset "optionally without cleaning up containers" begin
        mockExecuteDockerCompose(Mock()) do _executeDockerCompose
          # Ensure nothing gets returned from a call to `stop!` instead of
          # `Process`es from the `docker-compose` execution
          @test stop!(false) === nothing
          @test called_once_with(_executeDockerCompose, "stop")
        end
      end

      ENDPOINT[] = original_endpoint
    end

    @testset "`update!` pulls Kroki service component Docker images" begin
      mockExecuteDockerCompose(Mock()) do _executeDockerCompose
        # Ensure nothing gets returned from a call to `update!` instead of
        # `Process`es from the `docker-compose` execution
        @test update!() === nothing
        @test called_with(_executeDockerCompose, ["pull", "--quiet"])
      end
    end
  end
end

end