text:
export
    SOCWRPowerModel, SOCWRForm,
    QCWRPowerModel, QCWRForm,
    QCWRTriPowerModel, QCWRTriForm

""
abstract type AbstractWRForm <: AbstractPowerFormulation end

""
abstract type SOCWRForm <: AbstractWRForm end

""
const SOCWRPowerModel = GenericPowerModel{SOCWRForm}

"default SOC constructor"
SOCWRPowerModel(data::Dict{String,Any}; kwargs...) = GenericPowerModel(data, SOCWRForm; kwargs...)

""
function variable_voltage(pm::GenericPowerModel{T}, n::Int=pm.cnw; kwargs...) where T <: AbstractWRForm
    variable_voltage_magnitude_sqr(pm, n; kwargs...)
    variable_voltage_product(pm, n; kwargs...)
end

""
function constraint_voltage(pm::GenericPowerModel{T}, n::Int) where T <: AbstractWRForm
    w = pm.var[:nw][n][:w]
    wr = pm.var[:nw][n][:wr]
    wi = pm.var[:nw][n][:wi]

    for (i,j) in keys(pm.ref[:nw][n][:buspairs])
        relaxation_complex_product(pm.model, w[i], w[j], wr[(i,j)], wi[(i,j)])
    end
end

"""
Creates Ohms constraints (yt post fix indicates that Y and T values are in rectangular form)

```
p[f_idx] == g/tm*w_fr_ne[i] + (-g*tr+b*ti)/tm*(wr_ne[i]) + (-b*tr-g*ti)/tm*(wi_ne[i])
q[f_idx] == -(b+c/2)/tm*w_fr_ne[i] - (-b*tr-g*ti)/tm*(wr_ne[i]) + (-g*tr+b*ti)/tm*(wi_ne[i])
```
"""
function constraint_ohms_yt_from_ne(pm::GenericPowerModel{T}, n::Int, i, f_bus, t_bus, f_idx, t_idx, g, b, c, tr, ti, tm, vad_min, vad_max) where T <: AbstractWRForm
    p_fr = pm.var[:nw][n][:p_ne][f_idx]
    q_fr = pm.var[:nw][n][:q_ne][f_idx]
    w_fr = pm.var[:nw][n][:w_fr_ne][i]
    wr = pm.var[:nw][n][:wr_ne][i]
    wi = pm.var[:nw][n][:wi_ne][i]

    @constraint(pm.model, p_fr ==        g/tm^2*w_fr + (-g*tr+b*ti)/tm^2*wr + (-b*tr-g*ti)/tm^2*wi )
    @constraint(pm.model, q_fr == -(b+c/2)/tm^2*w_fr - (-b*tr-g*ti)/tm^2*wr + (-g*tr+b*ti)/tm^2*wi )
end

"""
Creates Ohms constraints (yt post fix indicates that Y and T values are in rectangular form)

```
p[t_idx] == g*w_to_ne[i] + (-g*tr-b*ti)/tm*(wr_ne[i]) + (-b*tr+g*ti)/tm*(-wi_ne[i])
q[t_idx] == -(b+c/2)*w_to_ne[i] - (-b*tr+g*ti)/tm*(wr_ne[i]) + (-g*tr-b*ti)/tm*(-wi_ne[i])
```
"""
function constraint_ohms_yt_to_ne(pm::GenericPowerModel{T}, n::Int, i, f_bus, t_bus, f_idx, t_idx, g, b, c, tr, ti, tm, vad_min, vad_max) where T <: AbstractWRForm
    p_to = pm.var[:nw][n][:p_ne][t_idx]
    q_to = pm.var[:nw][n][:q_ne][t_idx]
    w_to = pm.var[:nw][n][:w_to_ne][i]
    wr = pm.var[:nw][n][:wr_ne][i]
    wi = pm.var[:nw][n][:wi_ne][i]

    @constraint(pm.model, p_to ==        g*w_to + (-g*tr-b*ti)/tm^2*wr + (-b*tr+g*ti)/tm^2*-wi )
    @constraint(pm.model, q_to == -(b+c/2)*w_to - (-b*tr+g*ti)/tm^2*wr + (-g*tr-b*ti)/tm^2*-wi )
end

""
function variable_voltage_on_off(pm::GenericPowerModel{T}, n::Int=pm.cnw; kwargs...) where T <: AbstractWRForm
    variable_voltage_magnitude_sqr(pm, n; kwargs...)
    variable_voltage_magnitude_sqr_from_on_off(pm, n; kwargs...)
    variable_voltage_magnitude_sqr_to_on_off(pm, n; kwargs...)

    variable_voltage_product_on_off(pm, n; kwargs...)
end

""
function constraint_voltage_on_off(pm::GenericPowerModel{T}, n::Int) where T <: AbstractWRForm
    w = pm.var[:nw][n][:w]
    wr = pm.var[:nw][n][:wr]
    wi = pm.var[:nw][n][:wi]
    z = pm.var[:nw][n][:branch_z]

    w_fr = pm.var[:nw][n][:w_fr]
    w_to = pm.var[:nw][n][:w_to]

    constraint_voltage_magnitude_sqr_from_on_off(pm, n)
    constraint_voltage_magnitude_sqr_to_on_off(pm, n)
    constraint_voltage_product_on_off(pm, n)

    for (l,i,j) in pm.ref[:nw][n][:arcs_from]
        relaxation_complex_product_on_off(pm.model, w[i], w[j], wr[l], wi[l], z[l])
        relaxation_equality_on_off(pm.model, w[i], w_fr[l], z[l])
        relaxation_equality_on_off(pm.model, w[j], w_to[l], z[l])
    end
end

""
function constraint_voltage_ne(pm::GenericPowerModel{T}, n::Int) where T <: AbstractWRForm
    buses = pm.ref[:nw][n][:bus]
    branches = pm.ref[:nw][n][:ne_branch]

    wr_min, wr_max, wi_min, wi_max = calc_voltage_product_bounds(pm.ref[:nw][n][:ne_buspairs])
    bi_bp = Dict([(i, (b["f_bus"], b["t_bus"])) for (i,b) in branches])

    w = pm.var[:nw][n][:w]
    wr = pm.var[:nw][n][:wr_ne]
    wi = pm.var[:nw][n][:wi_ne]
    z = pm.var[:nw][n][:branch_ne]

    w_fr = pm.var[:nw][n][:w_fr_ne]
    w_to = pm.var[:nw][n][:w_to_ne]

    for (l,i,j) in pm.ref[:nw][n][:ne_arcs_from]
        @constraint(pm.model, w_fr[l] <= z[l]*buses[branches[l]["f_bus"]]["vmax"]^2)
        @constraint(pm.model, w_fr[l] >= z[l]*buses[branches[l]["f_bus"]]["vmin"]^2)

        @constraint(pm.model, wr[l] <= z[l]*wr_max[bi_bp[l]])
        @constraint(pm.model, wr[l] >= z[l]*wr_min[bi_bp[l]])
        @constraint(pm.model, wi[l] <= z[l]*wi_max[bi_bp[l]])
        @constraint(pm.model, wi[l] >= z[l]*wi_min[bi_bp[l]])

        @constraint(pm.model, w_to[l] <= z[l]*buses[branches[l]["t_bus"]]["vmax"]^2)
        @constraint(pm.model, w_to[l] >= z[l]*buses[branches[l]["t_bus"]]["vmin"]^2)

        relaxation_complex_product_on_off(pm.model, w[i], w[j], wr[l], wi[l], z[l])
        relaxation_equality_on_off(pm.model, w[i], w_fr[l], z[l])
        relaxation_equality_on_off(pm.model, w[j], w_to[l], z[l])
    end
end

""
function constraint_voltage_magnitude_from_on_off(pm::GenericPowerModel{T}, n::Int) where T <: AbstractWRForm
    buses = pm.ref[:nw][n][:bus]
    branches = pm.ref[:nw][n][:branch]

    vm_fr = pm.var[:nw][n][:vm_fr]
    z = pm.var[:nw][n][:branch_z]

    for (i, branch) in pm.ref[:nw][n][:branch]
        @constraint(pm.model, vm_fr[i] <= z[i]*buses[branch["f_bus"]]["vmax"])
        @constraint(pm.model, vm_fr[i] >= z[i]*buses[branch["f_bus"]]["vmin"])
    end
end

""
function constraint_voltage_magnitude_to_on_off(pm::GenericPowerModel{T}, n::Int) where T <: AbstractWRForm
    buses = pm.ref[:nw][n][:bus]
    branches = pm.ref[:nw][n][:branch]

    vm_to = pm.var[:nw][n][:vm_to]
    z = pm.var[:nw][n][:branch_z]

    for (i, branch) in pm.ref[:nw][n][:branch]
        @constraint(pm.model, vm_to[i] <= z[i]*buses[branch["t_bus"]]["vmax"])
        @constraint(pm.model, vm_to[i] >= z[i]*buses[branch["t_bus"]]["vmin"])
    end
end

""
function constraint_voltage_magnitude_sqr_from_on_off(pm::GenericPowerModel{T}, n::Int) where T <: AbstractWRForm
    buses = pm.ref[:nw][n][:bus]
    branches = pm.ref[:nw][n][:branch]

    w_fr = pm.var[:nw][n][:w_fr]
    z = pm.var[:nw][n][:branch_z]

    for (i, branch) in pm.ref[:nw][n][:branch]
        @constraint(pm.model, w_fr[i] <= z[i]*buses[branch["f_bus"]]["vmax"]^2)
        @constraint(pm.model, w_fr[i] >= z[i]*buses[branch["f_bus"]]["vmin"]^2)
    end
end

""
function constraint_voltage_magnitude_sqr_to_on_off(pm::GenericPowerModel{T}, n::Int) where T <: AbstractWRForm
    buses = pm.ref[:nw][n][:bus]
    branches = pm.ref[:nw][n][:branch]

    w_to = pm.var[:nw][n][:w_to]
    z = pm.var[:nw][n][:branch_z]

    for (i, branch) in pm.ref[:nw][n][:branch]
        @constraint(pm.model, w_to[i] <= z[i]*buses[branch["t_bus"]]["vmax"]^2)
        @constraint(pm.model, w_to[i] >= z[i]*buses[branch["t_bus"]]["vmin"]^2)
    end
end

""
function constraint_voltage_product_on_off(pm::GenericPowerModel{T}, n::Int) where T <: AbstractWRForm
    wr_min, wr_max, wi_min, wi_max = calc_voltage_product_bounds(pm.ref[:nw][n][:buspairs])

    bi_bp = Dict([(i, (b["f_bus"], b["t_bus"])) for (i,b) in pm.ref[:nw][n][:branch]])

    wr = pm.var[:nw][n][:wr]
    wi = pm.var[:nw][n][:wi]
    z = pm.var[:nw][n][:branch_z]

    for b in keys(pm.ref[:nw][n][:branch])
        @constraint(pm.model, wr[b] <= z[b]*wr_max[bi_bp[b]])
        @constraint(pm.model, wr[b] >= z[b]*wr_min[bi_bp[b]])
        @constraint(pm.model, wi[b] <= z[b]*wi_max[bi_bp[b]])
        @constraint(pm.model, wi[b] >= z[b]*wi_min[bi_bp[b]])
    end
end

"""
Creates Ohms constraints (yt post fix indicates that Y and T values are in rectangular form)

```
p[f_idx] ==        g/tm*w_fr[i] + (-g*tr+b*ti)/tm*(wr[i]) + (-b*tr-g*ti)/tm*(wi[i])
q[f_idx] == -(b+c/2)/tm*w_fr[i] - (-b*tr-g*ti)/tm*(wr[i]) + (-g*tr+b*ti)/tm*(wi[i])
```
"""
function constraint_ohms_yt_from_on_off(pm::GenericPowerModel{T}, n::Int, i, f_bus, t_bus, f_idx, t_idx, g, b, c, tr, ti, tm, vad_min, vad_max) where T <: AbstractWRForm
    p_fr = pm.var[:nw][n][:p][f_idx]
    q_fr = pm.var[:nw][n][:q][f_idx]
    w_fr = pm.var[:nw][n][:w_fr][i]
    wr = pm.var[:nw][n][:wr][i]
    wi = pm.var[:nw][n][:wi][i]

    @constraint(pm.model, p_fr ==        g/tm^2*w_fr + (-g*tr+b*ti)/tm^2*wr + (-b*tr-g*ti)/tm^2*wi )
    @constraint(pm.model, q_fr == -(b+c/2)/tm^2*w_fr - (-b*tr-g*ti)/tm^2*wr + (-g*tr+b*ti)/tm^2*wi )
end

"""
Creates Ohms constraints (yt post fix indicates that Y and T values are in rectangular form)

```
p[t_idx] ==        g*w_to[i] + (-g*tr-b*ti)/tm*(wr[i]) + (-b*tr+g*ti)/tm*(-wi[i])
q[t_idx] == -(b+c/2)*w_to[i] - (-b*tr+g*ti)/tm*(wr[i]) + (-g*tr-b*ti)/tm*(-wi[i])
```
"""
function constraint_ohms_yt_to_on_off(pm::GenericPowerModel{T}, n::Int, i, f_bus, t_bus, f_idx, t_idx, g, b, c, tr, ti, tm, vad_min, vad_max) where T <: AbstractWRForm
    p_to = pm.var[:nw][n][:p][t_idx]
    q_to = pm.var[:nw][n][:q][t_idx]
    w_to = pm.var[:nw][n][:w_to][i]
    wr = pm.var[:nw][n][:wr][i]
    wi = pm.var[:nw][n][:wi][i]

    @constraint(pm.model, p_to ==        g*w_to + (-g*tr-b*ti)/tm^2*wr + (-b*tr+g*ti)/tm^2*-wi )
    @constraint(pm.model, q_to == -(b+c/2)*w_to - (-b*tr+g*ti)/tm^2*wr + (-g*tr-b*ti)/tm^2*-wi )
end

"`angmin*wr[i] <= wi[i] <= angmax*wr[i]`"
function constraint_voltage_angle_difference_on_off(pm::GenericPowerModel{T}, n::Int, i, f_bus, t_bus, angmin, angmax, vad_min, vad_max) where T <: AbstractWRForm
    wr = pm.var[:nw][n][:wr][i]
    wi = pm.var[:nw][n][:wi][i]

    @constraint(pm.model, wi <= tan(angmax)*wr)
    @constraint(pm.model, wi >= tan(angmin)*wr)
end

"`angmin*wr_ne[i] <= wi_ne[i] <= angmax*wr_ne[i]`"
function constraint_voltage_angle_difference_ne(pm::GenericPowerModel{T}, n::Int, i, f_bus, t_bus, angmin, angmax, vad_min, vad_max) where T <: AbstractWRForm
    wr = pm.var[:nw][n][:wr_ne][i]
    wi = pm.var[:nw][n][:wi_ne][i]

    @constraint(pm.model, wi <= tan(angmax)*wr)
    @constraint(pm.model, wi >= tan(angmin)*wr)
end

""
function variable_voltage_ne(pm::GenericPowerModel{T}, n::Int=pm.cnw; kwargs...) where T <: AbstractWRForm
    variable_voltage_magnitude_sqr_from_ne(pm, n; kwargs...)
    variable_voltage_magnitude_sqr_to_ne(pm, n; kwargs...)
    variable_voltage_product_ne(pm, n; kwargs...)
end

""
function variable_voltage_magnitude_sqr_from_ne(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T <: AbstractWRForm
    buses = pm.ref[:nw][n][:bus]
    branches = pm.ref[:nw][n][:ne_branch]

    pm.var[:nw][n][:w_fr_ne] = @variable(pm.model,
        [i in keys(pm.ref[:nw][n][:ne_branch])], basename="$(n)_w_fr_ne",
        lowerbound = 0,
        upperbound = buses[branches[i]["f_bus"]]["vmax"]^2,
        start = getstart(pm.ref[:nw][n][:bus], i, "w_fr_start", 1.001)
    )
end

""
function variable_voltage_magnitude_sqr_to_ne(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T <: AbstractWRForm
    buses = pm.ref[:nw][n][:bus]
    branches = pm.ref[:nw][n][:ne_branch]

    pm.var[:nw][n][:w_to_ne] = @variable(pm.model,
        [i in keys(pm.ref[:nw][n][:ne_branch])], basename="$(n)_w_to_ne",
        lowerbound = 0,
        upperbound = buses[branches[i]["t_bus"]]["vmax"]^2,
        start = getstart(pm.ref[:nw][n][:bus], i, "w_to", 1.001)
    )
end

""
function variable_voltage_product_ne(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T <: AbstractWRForm
    wr_min, wr_max, wi_min, wi_max = calc_voltage_product_bounds(pm.ref[:nw][n][:ne_buspairs])
    bi_bp = Dict([(i, (b["f_bus"], b["t_bus"])) for (i,b) in pm.ref[:nw][n][:ne_branch]])

    pm.var[:nw][n][:wr_ne] = @variable(pm.model,
        [b in keys(pm.ref[:nw][n][:ne_branch])], basename="$(n)_wr_ne",
        lowerbound = min(0, wr_min[bi_bp[b]]),
        upperbound = max(0, wr_max[bi_bp[b]]),
        start = getstart(pm.ref[:nw][n][:ne_buspairs], bi_bp[b], "wr_start", 1.0)
    )

    pm.var[:nw][n][:wi_ne] = @variable(pm.model,
        [b in keys(pm.ref[:nw][n][:ne_branch])], basename="$(n)_wi_ne",
        lowerbound = min(0, wi_min[bi_bp[b]]),
        upperbound = max(0, wi_max[bi_bp[b]]),
        start = getstart(pm.ref[:nw][n][:ne_buspairs], bi_bp[b], "wi_start")
    )
end

""
abstract type QCWRForm <: AbstractWRForm end

""
const QCWRPowerModel = GenericPowerModel{QCWRForm}

"default QC constructor"
function QCWRPowerModel(data::Dict{String,Any}; kwargs...)
    return GenericPowerModel(data, QCWRForm; kwargs...)
end

"Creates variables associated with differences in phase angles"
function variable_voltage_angle_difference(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T
    pm.var[:nw][n][:td] = @variable(pm.model,
        [bp in keys(pm.ref[:nw][n][:buspairs])], basename="$(n)_td",
        lowerbound = pm.ref[:nw][n][:buspairs][bp]["angmin"],
        upperbound = pm.ref[:nw][n][:buspairs][bp]["angmax"],
        start = getstart(pm.ref[:nw][n][:buspairs], bp, "td_start")
    )
end

"Creates the voltage magnitude product variables"
function variable_voltage_magnitude_product(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T
    buspairs = pm.ref[:nw][n][:buspairs]
    pm.var[:nw][n][:vv] = @variable(pm.model,
        [bp in keys(pm.ref[:nw][n][:buspairs])], basename="$(n)_vv",
        lowerbound = buspairs[bp]["vm_fr_min"]*buspairs[bp]["vm_to_min"],
        upperbound = buspairs[bp]["vm_fr_max"]*buspairs[bp]["vm_to_max"],
        start = getstart(pm.ref[:nw][n][:buspairs], bp, "vv_start", 1.0)
    )
end

""
function variable_cosine(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T
    cos_min = Dict([(bp, -Inf) for bp in keys(pm.ref[:nw][n][:buspairs])])
    cos_max = Dict([(bp,  Inf) for bp in keys(pm.ref[:nw][n][:buspairs])])

    for (bp, buspair) in pm.ref[:nw][n][:buspairs]
        if buspair["angmin"] >= 0
            cos_max[bp] = cos(buspair["angmin"])
            cos_min[bp] = cos(buspair["angmax"])
        end
        if buspair["angmax"] <= 0
            cos_max[bp] = cos(buspair["angmax"])
            cos_min[bp] = cos(buspair["angmin"])
        end
        if buspair["angmin"] < 0 && buspair["angmax"] > 0
            cos_max[bp] = 1.0
            cos_min[bp] = min(cos(buspair["angmin"]), cos(buspair["angmax"]))
        end
    end

    pm.var[:nw][n][:cs] = @variable(pm.model,
        [bp in keys(pm.ref[:nw][n][:buspairs])], basename="$(n)_cs",
        lowerbound = cos_min[bp],
        upperbound = cos_max[bp],
        start = getstart(pm.ref[:nw][n][:buspairs], bp, "cs_start", 1.0)
    )
end

""
function variable_sine(pm::GenericPowerModel, n::Int=pm.cnw)
    pm.var[:nw][n][:si] = @variable(pm.model,
        [bp in keys(pm.ref[:nw][n][:buspairs])], basename="$(n)_si",
        lowerbound = sin(pm.ref[:nw][n][:buspairs][bp]["angmin"]),
        upperbound = sin(pm.ref[:nw][n][:buspairs][bp]["angmax"]),
        start = getstart(pm.ref[:nw][n][:buspairs], bp, "si_start")
    )
end

""
function variable_current_magnitude_sqr(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T
    buspairs = pm.ref[:nw][n][:buspairs]
    pm.var[:nw][n][:cm] = @variable(pm.model,
        cm[bp in keys(pm.ref[:nw][n][:buspairs])], basename="$(n)_cm",
        lowerbound = 0,
        upperbound = (buspairs[bp]["rate_a"]*buspairs[bp]["tap"]/buspairs[bp]["vm_fr_min"])^2,
        start = getstart(pm.ref[:nw][n][:buspairs], bp, "cm_start")
    )
end

""
function variable_voltage(pm::GenericPowerModel{T}, n::Int=pm.cnw; kwargs...) where T <: QCWRForm
    variable_voltage_angle(pm, n; kwargs...)
    variable_voltage_magnitude(pm, n; kwargs...)

    variable_voltage_magnitude_sqr(pm, n; kwargs...)
    variable_voltage_product(pm, n; kwargs...)

    variable_voltage_angle_difference(pm, n; kwargs...)
    variable_voltage_magnitude_product(pm, n; kwargs...)
    variable_cosine(pm, n; kwargs...)
    variable_sine(pm, n; kwargs...)
    variable_current_magnitude_sqr(pm, n; kwargs...)
end

""
function constraint_voltage(pm::GenericPowerModel{T}, n::Int) where T <: QCWRForm
    v = pm.var[:nw][n][:vm]
    t = pm.var[:nw][n][:va]

    td = pm.var[:nw][n][:td]
    si = pm.var[:nw][n][:si]
    cs = pm.var[:nw][n][:cs]
    vv = pm.var[:nw][n][:vv]

    w = pm.var[:nw][n][:w]
    wr = pm.var[:nw][n][:wr]
    wi = pm.var[:nw][n][:wi]

    for (i,b) in pm.ref[:nw][n][:bus]
        relaxation_sqr(pm.model, v[i], w[i])
    end

    for bp in keys(pm.ref[:nw][n][:buspairs])
        i,j = bp
        @constraint(pm.model, t[i] - t[j] == td[bp])

        relaxation_sin(pm.model, td[bp], si[bp])
        relaxation_cos(pm.model, td[bp], cs[bp])
        relaxation_product(pm.model, v[i], v[j], vv[bp])
        relaxation_product(pm.model, vv[bp], cs[bp], wr[bp])
        relaxation_product(pm.model, vv[bp], si[bp], wi[bp])

        # this constraint is redudant and useful for debugging
        #relaxation_complex_product(pm.model, w[i], w[j], wr[bp], wi[bp])
   end

   for (i,branch) in pm.ref[:nw][n][:branch]
        pair = (branch["f_bus"], branch["t_bus"])
        buspair = pm.ref[:nw][n][:buspairs][pair]

        # to prevent this constraint from being posted on multiple parallel branches
        if buspair["branch"] == i
            constraint_power_magnitude_sqr(pm, n, i)
            constraint_power_magnitude_link(pm, n, i)
        end
    end

end

"`p[f_idx]^2 + q[f_idx]^2 <= w[f_bus]/tm*cm[f_bus,t_bus]`"
function constraint_power_magnitude_sqr(pm::GenericPowerModel{T}, n::Int, f_bus, t_bus, arc_from, tm) where T <: QCWRForm
    w_i = pm.var[:nw][n][:w][f_bus]
    p_fr = pm.var[:nw][n][:p][arc_from]
    q_fr = pm.var[:nw][n][:q][arc_from]
    cm = pm.var[:nw][n][:cm][(f_bus, t_bus)]

    @constraint(pm.model, p_fr^2 + q_fr^2 <= w_i/tm^2*cm)
end

"`cm[f_bus,t_bus] == (g^2 + b^2)*(w[f_bus]/tm + w[t_bus] - 2*(tr*wr[f_bus,t_bus] + ti*wi[f_bus,t_bus])/tm) - c*q[f_idx] - ((c/2)/tm)^2*w[f_bus]`"
function constraint_power_magnitude_link(pm::GenericPowerModel{T}, n::Int, f_bus, t_bus, arc_from, g, b, c, tr, ti, tm) where T <: QCWRForm
    w_fr = pm.var[:nw][n][:w][f_bus]
    w_to = pm.var[:nw][n][:w][t_bus]
    q_fr = pm.var[:nw][n][:q][arc_from]
    wr = pm.var[:nw][n][:wr][(f_bus, t_bus)]
    wi = pm.var[:nw][n][:wi][(f_bus, t_bus)]
    cm = pm.var[:nw][n][:cm][(f_bus, t_bus)]

    @constraint(pm.model, cm == (g^2 + b^2)*(w_fr/tm^2 + w_to - 2*(tr*wr + ti*wi)/tm^2) - c*q_fr - ((c/2)/tm^2)^2*w_fr)
end

"`t[ref_bus] == 0`"
function constraint_theta_ref(pm::GenericPowerModel{T}, n::Int, i::Int) where T <: QCWRForm
    @constraint(pm.model, pm.var[:nw][n][:va][i] == 0)
end

""
function constraint_voltage_angle_difference(pm::GenericPowerModel{T}, n::Int, f_bus, t_bus, angmin, angmax) where T <: QCWRForm
    td = pm.var[:nw][n][:td][(f_bus, t_bus)]

    if getlowerbound(td) < angmin
        setlowerbound(td, angmin)
    end

    if getupperbound(td) > angmax
        setupperbound(td, angmax)
    end

    w_fr = pm.var[:nw][n][:w][f_bus]
    w_to = pm.var[:nw][n][:w][t_bus]
    wr = pm.var[:nw][n][:wr][(f_bus, t_bus)]
    wi = pm.var[:nw][n][:wi][(f_bus, t_bus)]

    @constraint(pm.model, wi <= tan(angmax)*wr)
    @constraint(pm.model, wi >= tan(angmin)*wr)

    cut_complex_product_and_angle_difference(pm.model, w_fr, w_to, wr, wi, angmin, angmax)
end

""
function add_bus_voltage_setpoint(sol, pm::GenericPowerModel{T}) where T <: QCWRForm
    add_setpoint(sol, pm, "bus", "vm", :vm)
    add_setpoint(sol, pm, "bus", "va", :va)
end

""
function variable_voltage_on_off(pm::GenericPowerModel{T}, n::Int=pm.cnw; kwargs...) where T <: QCWRForm
    variable_voltage_angle(pm, n; kwargs...)
    variable_voltage_magnitude(pm, n; kwargs...)
    variable_voltage_magnitude_from_on_off(pm, n; kwargs...)
    variable_voltage_magnitude_to_on_off(pm, n; kwargs...)

    variable_voltage_magnitude_sqr(pm, n; kwargs...)
    variable_voltage_magnitude_sqr_from_on_off(pm, n; kwargs...)
    variable_voltage_magnitude_sqr_to_on_off(pm, n; kwargs...)

    variable_voltage_product_on_off(pm, n; kwargs...)

    variable_voltage_angle_difference_on_off(pm, n; kwargs...)
    variable_voltage_magnitude_product_on_off(pm, n; kwargs...)
    variable_cosine_on_off(pm, n; kwargs...)
    variable_sine_on_off(pm, n; kwargs...)
    variable_current_magnitude_sqr_on_off(pm, n; kwargs...) # includes 0, but needs new indexs
end

""
function variable_voltage_angle_difference_on_off(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T
    pm.var[:nw][n][:td] = @variable(pm.model,
        [l in keys(pm.ref[:nw][n][:branch])], basename="$(n)_td",
        lowerbound = min(0, pm.ref[:nw][n][:branch][l]["angmin"]),
        upperbound = max(0, pm.ref[:nw][n][:branch][l]["angmax"]),
        start = getstart(pm.ref[:nw][n][:branch], l, "td_start")
    )
end

""
function variable_voltage_magnitude_product_on_off(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T
    vv_min = Dict([(l, pm.ref[:nw][n][:bus][branch["f_bus"]]["vmin"]*pm.ref[:nw][n][:bus][branch["t_bus"]]["vmin"]) for (l, branch) in pm.ref[:nw][n][:branch]])
    vv_max = Dict([(l, pm.ref[:nw][n][:bus][branch["f_bus"]]["vmax"]*pm.ref[:nw][n][:bus][branch["t_bus"]]["vmax"]) for (l, branch) in pm.ref[:nw][n][:branch]])

    pm.var[:nw][n][:vv] = @variable(pm.model,
        [l in keys(pm.ref[:nw][n][:branch])], basename="$(n)_vv",
        lowerbound = min(0, vv_min[l]),
        upperbound = max(0, vv_max[l]),
        start = getstart(pm.ref[:nw][n][:branch], l, "vv_start", 1.0)
    )
end

""
function variable_cosine_on_off(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T
    cos_min = Dict([(l, -Inf) for l in keys(pm.ref[:nw][n][:branch])])
    cos_max = Dict([(l,  Inf) for l in keys(pm.ref[:nw][n][:branch])])

    for (l, branch) in pm.ref[:nw][n][:branch]
        if branch["angmin"] >= 0
            cos_max[l] = cos(branch["angmin"])
            cos_min[l] = cos(branch["angmax"])
        end
        if branch["angmax"] <= 0
            cos_max[l] = cos(branch["angmax"])
            cos_min[l] = cos(branch["angmin"])
        end
        if branch["angmin"] < 0 && branch["angmax"] > 0
            cos_max[l] = 1.0
            cos_min[l] = min(cos(branch["angmin"]), cos(branch["angmax"]))
        end
    end

    pm.var[:nw][n][:cs] = @variable(pm.model,
        [l in keys(pm.ref[:nw][n][:branch])], basename="$(n)_cs",
        lowerbound = min(0, cos_min[l]),
        upperbound = max(0, cos_max[l]),
        start = getstart(pm.ref[:nw][n][:branch], l, "cs_start", 1.0)
    )
end

""
function variable_sine_on_off(pm::GenericPowerModel, n::Int=pm.cnw)
    pm.var[:nw][n][:si] = @variable(pm.model,
        [l in keys(pm.ref[:nw][n][:branch])], basename="$(n)_si",
        lowerbound = min(0, sin(pm.ref[:nw][n][:branch][l]["angmin"])),
        upperbound = max(0, sin(pm.ref[:nw][n][:branch][l]["angmax"])),
        start = getstart(pm.ref[:nw][n][:branch], l, "si_start")
    )
end

""
function variable_current_magnitude_sqr_on_off(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T
    cm_min = Dict([(l, 0) for l in keys(pm.ref[:nw][n][:branch])])
    cm_max = Dict([(l, (branch["rate_a"]*branch["tap"]/pm.ref[:nw][n][:bus][branch["f_bus"]]["vmin"])^2) for (l, branch) in pm.ref[:nw][n][:branch]])

    pm.var[:nw][n][:cm] = @variable(pm.model,
        [l in keys(pm.ref[:nw][n][:branch])], basename="$(n)_cm",
        lowerbound = cm_min[l],
        upperbound = cm_max[l],
        start = getstart(pm.ref[:nw][n][:branch], l, "cm_start")
    )
end

""
function constraint_voltage_on_off(pm::GenericPowerModel{T}, n::Int) where T <: QCWRForm
    v = pm.var[:nw][n][:vm]
    t = pm.var[:nw][n][:va]
    vm_fr = pm.var[:nw][n][:vm_fr]
    vm_to = pm.var[:nw][n][:vm_to]

    td = pm.var[:nw][n][:td]
    si = pm.var[:nw][n][:si]
    cs = pm.var[:nw][n][:cs]
    vv = pm.var[:nw][n][:vv]

    w = pm.var[:nw][n][:w]
    w_fr = pm.var[:nw][n][:w_fr]
    w_to = pm.var[:nw][n][:w_to]

    wr = pm.var[:nw][n][:wr]
    wi = pm.var[:nw][n][:wi]

    z = pm.var[:nw][n][:branch_z]

    td_lb = pm.ref[:nw][n][:off_angmin]
    td_ub = pm.ref[:nw][n][:off_angmax]
    td_max = max(abs(td_lb), abs(td_ub))

    for (i,b) in pm.ref[:nw][n][:bus]
        relaxation_sqr(pm.model, v[i], w[i])
    end

    constraint_voltage_magnitude_from_on_off(pm, n) # bounds on vm_fr
    constraint_voltage_magnitude_to_on_off(pm, n) # bounds on vm_to
    constraint_voltage_magnitude_sqr_from_on_off(pm, n) # bounds on w_fr
    constraint_voltage_magnitude_sqr_to_on_off(pm, n) # bounds on w_to
    constraint_voltage_product_on_off(pm, n) # bounds on wr, wi

    for (l,branch) in pm.ref[:nw][n][:branch]
        i = branch["f_bus"]
        j = branch["t_bus"]

        @constraint(pm.model, t[i] - t[j] >= td[l] + td_lb*(1-z[l]))
        @constraint(pm.model, t[i] - t[j] <= td[l] + td_ub*(1-z[l]))

        relaxation_sin_on_off(pm.model, td[l], si[l], z[l], td_max)
        relaxation_cos_on_off(pm.model, td[l], cs[l], z[l], td_max)
        relaxation_product_on_off(pm.model, vm_fr[l], vm_to[l], vv[l], z[l])
        relaxation_product_on_off(pm.model, vv[l], cs[l], wr[l], z[l])
        relaxation_product_on_off(pm.model, vv[l], si[l], wi[l], z[l])

        # this constraint is redudant and useful for debugging
        #relaxation_complex_product(pm.model, w[i], w[j], wr[l], wi[l])

        #cs4 = relaxation_complex_product_on_off(pm.model, w[i], w[j], wr[l], wi[l], z[l])
        relaxation_equality_on_off(pm.model, v[i], vm_fr[l], z[l])
        relaxation_equality_on_off(pm.model, v[j], vm_to[l], z[l])
        relaxation_equality_on_off(pm.model, w[i], w_fr[l], z[l])
        relaxation_equality_on_off(pm.model, w[j], w_to[l], z[l])

        # to prevent this constraint from being posted on multiple parallel branchs
        # TODO needs on/off variant
        constraint_power_magnitude_sqr_on_off(pm, n, l)
        constraint_power_magnitude_link_on_off(pm, n, l) # different index set
    end
end

"`p[arc_from]^2 + q[arc_from]^2 <= w[f_bus]/tm*cm[i]`"
function constraint_power_magnitude_sqr_on_off(pm::GenericPowerModel{T}, n::Int, i, f_bus, arc_from, tm) where T <: QCWRForm
    w = pm.var[:nw][n][:w][f_bus]
    p_fr = pm.var[:nw][n][:p][arc_from]
    q_fr = pm.var[:nw][n][:q][arc_from]
    cm = pm.var[:nw][n][:cm][i]
    z = pm.var[:nw][n][:branch_z][i]

    # TODO see if there is a way to leverage relaxation_complex_product_on_off here
    w_ub = getupperbound(w)
    cm_ub = getupperbound(cm)
    z_ub = getupperbound(z)

    @constraint(pm.model, p_fr^2 + q_fr^2 <= w*cm*z_ub/tm^2)
    @constraint(pm.model, p_fr^2 + q_fr^2 <= w_ub*cm*z/tm^2)
    @constraint(pm.model, p_fr^2 + q_fr^2 <= w*cm_ub*z/tm^2)
end

"`cm[f_bus,t_bus] == (g^2 + b^2)*(w[f_bus]/tm + w[t_bus] - 2*(tr*wr[f_bus,t_bus] + ti*wi[f_bus,t_bus])/tm) - c*q[f_idx] - ((c/2)/tm)^2*w[f_bus]`"
function constraint_power_magnitude_link_on_off(pm::GenericPowerModel{T}, n::Int, i, arc_from, g, b, c, tr, ti, tm) where T <: QCWRForm
    w_fr = pm.var[:nw][n][:w_fr][i]
    w_to = pm.var[:nw][n][:w_to][i]
    q_fr = pm.var[:nw][n][:q][arc_from]
    wr = pm.var[:nw][n][:wr][i]
    wi = pm.var[:nw][n][:wi][i]
    cm = pm.var[:nw][n][:cm][i]

    @constraint(pm.model, cm == (g^2 + b^2)*(w_fr/tm^2 + w_to - 2*(tr*wr + ti*wi)/tm^2) - c*q_fr - ((c/2)/tm^2)^2*w_fr)
end

""
abstract type QCWRTriForm <: QCWRForm end

""
const QCWRTriPowerModel = GenericPowerModel{QCWRTriForm}

"default QC trilinear model constructor"
function QCWRTriPowerModel(data::Dict{String,Any}; kwargs...)
    return GenericPowerModel(data, QCWRTriForm; kwargs...)
end

""
function variable_voltage_magnitude_product(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T <: QCWRTriForm
    # do nothing - no lifted variables required for voltage variable product
end

"creates lambda variables for convex combination model"
function variable_multipliers(pm::GenericPowerModel{T}, n::Int=pm.cnw) where T <: QCWRTriForm
    buspairs = pm.ref[:nw][n][:buspairs]
    pm.var[:nw][n][:lambda_wr] = @variable(pm.model,
        [bp in keys(pm.ref[:nw][n][:buspairs]), i=1:8], basename="$(n)_lambda",
        lowerbound = 0, upperbound = 1)
    pm.var[:nw][n][:lambda_wi] = @variable(pm.model,
        [bp in keys(pm.ref[:nw][n][:buspairs]), i=1:8], basename="$(n)_lambda",
        lowerbound = 0, upperbound = 1)
end

""
function variable_voltage(pm::GenericPowerModel{T}, n::Int=pm.cnw; kwargs...) where T <: QCWRTriForm
    variable_voltage_angle(pm, n; kwargs...)
    variable_voltage_magnitude(pm, n; kwargs...)

    variable_voltage_magnitude_sqr(pm, n; kwargs...)
    variable_voltage_product(pm, n; kwargs...)

    variable_voltage_angle_difference(pm, n; kwargs...)
    variable_voltage_magnitude_product(pm, n; kwargs...)
    variable_multipliers(pm, n; kwargs...)
    variable_cosine(pm, n; kwargs...)
    variable_sine(pm, n; kwargs...)
    variable_current_magnitude_sqr(pm, n; kwargs...)
end

""
function constraint_voltage(pm::GenericPowerModel{T}, n::Int) where T <: QCWRTriForm
    v = pm.var[:nw][n][:vm]
    t = pm.var[:nw][n][:va]

    td = pm.var[:nw][n][:td]
    si = pm.var[:nw][n][:si]
    cs = pm.var[:nw][n][:cs]

    w = pm.var[:nw][n][:w]
    wr = pm.var[:nw][n][:wr]
    lambda_wr = pm.var[:nw][n][:lambda_wr]
    wi = pm.var[:nw][n][:wi]
    lambda_wi = pm.var[:nw][n][:lambda_wi]

    for (i,b) in pm.ref[:nw][n][:bus]
        relaxation_sqr(pm.model, v[i], w[i])
    end

    for bp in keys(pm.ref[:nw][n][:buspairs])
        i,j = bp
        @constraint(pm.model, t[i] - t[j] == td[bp])

        relaxation_sin(pm.model, td[bp], si[bp])
        relaxation_cos(pm.model, td[bp], cs[bp])
        relaxation_trilinear(pm.model, v[i], v[j], cs[bp], wr[bp], lambda_wr[bp,:])
        relaxation_trilinear(pm.model, v[i], v[j], si[bp], wi[bp], lambda_wi[bp,:])

        # this constraint is redudant and useful for debugging
        #relaxation_complex_product(pm.model, w[i], w[j], wr[bp], wi[bp])
   end

   for (i,branch) in pm.ref[:nw][n][:branch]
        pair = (branch["f_bus"], branch["t_bus"])
        buspair = pm.ref[:nw][n][:buspairs][pair]

        # to prevent this constraint from being posted on multiple parallel branchs
        if buspair["branch"] == i
            constraint_power_magnitude_sqr(pm, n, i)
            constraint_power_magnitude_link(pm, n, i)
        end
    end

end


text:
#! /usr/bin/julia
#
#	csv_create.jl
#
#					Jul/05/2018
# --------------------------------------------------------------------
#
include("text_manipulate.jl")
#
# --------------------------------------------------------------------
function data_prepare_proc()
	dict_aa = Dict()
#
	dict_aa = dict_append_proc(dict_aa,"t1271","千葉",73561,"2013-7-30")
	dict_aa = dict_append_proc(dict_aa,"t1272","勝浦",24895,"2013-8-10")
	dict_aa = dict_append_proc(dict_aa,"t1273","市原",67294,"2013-1-14")
	dict_aa = dict_append_proc(dict_aa,"t1274","流山",73612,"2013-5-9")
	dict_aa = dict_append_proc(dict_aa,"t1275","八千代",63491,"2013-8-4")
	dict_aa = dict_append_proc(dict_aa,"t1276","我孫子",41827,"2013-1-21")
	dict_aa = dict_append_proc(dict_aa,"t1277","鴨川",12946,"2013-7-23")
	dict_aa = dict_append_proc(dict_aa,"t1278","銚子",79128,"2013-11-26")
	dict_aa = dict_append_proc(dict_aa,"t1279","市川",13572,"2013-10-15")
#

	return	dict_aa
end
# --------------------------------------------------------------------
println("*** 開始 ***")
fname_out = ARGS[1]
println(fname_out)
#
dict_aa = data_prepare_proc()
dict_to_csv_proc(dict_aa,fname_out)
#
println("*** 終了 ***")
# --------------------------------------------------------------------

ext: jl

max_stars_repo_licenses: MIT Unlicense

max_issues_repo_licenses: MIT Unlicense

max_forks_repo_licenses: MIT Unlicense


text:
import Test: @testset, @test
import Model_Household: Household,default_bu_allocation
import FRS_Household_Getter
import Example_Household_Getter
import DataUtils: MinMaxes, add_to!, print
using Definitions

@testset "Example Households" begin
    @time names = Example_Household_Getter.initialise()

    @time for name in names
        hh = Example_Household_Getter.get_household( name )
        bus = default_bu_allocation( hh )
        println(name)
        println( hh )
        println()
    end
end

@testset  "Main FRS Households"  begin

    @time nhhs,npeople,nhh2 = FRS_Household_Getter.initialise(
        household_name = "model_households_scotland",
        people_name    = "model_people_scotland",
        start_year = 2015
    )
    maxbus = -1
    npers_from_bus = 0
    num_wrong_age = 0
    num_bus = 0
    incomes_sum = MinMaxes{Incomes_Type}()
    assets_sum = MinMaxes{Asset_Type}()

    @time for hno in 1:nhhs
        hh = FRS_Household_Getter.get_household( hno )
        bus = default_bu_allocation( hh )
        nbus = size(bus)[1]
        @test nbus in 1:12
        for buno in 1:nbus
            num_bus += 1
            np = size( bus[buno])[1]
            for i in 1:np
                pers = bus[buno][i]
                add_to!(incomes_sum, pers.income )
                add_to!(assets_sum, pers.assets )
                # sorted right way
                if i > 1
                    lastpid = bus[buno][i-1].pid
                    # @test pers.pid < lastpid
                    lastpno =  bus[buno][i-1].pno
                    # @test pers.pno < lastpno
                    lastage = bus[buno][i-1].age
                    if pers.age > lastage
                        println( "pno=$(pers.pno) lastpno=$lastpno pers.age=$(pers.age) ; lastage = $lastage")
                        num_wrong_age += 1
                    end
                end
                npers_from_bus += 1
            end
        end
        maxbus = max(maxbus, nbus )
        if hno % 10_000 == 0
            println( hh )
            println()
        end
    end
    println( "incomes:")
    print( incomes_sum )
    println( "assets:")
    print( assets_sum )
    println( "num HHss = $nhhs")
    println( "num BUs = $num_bus")
    println( "num people = $npeople")
    println( "largest number of BUs = $maxbus")
    println( "num out-of-sync ages $num_wrong_age")
    # bu allocation doesn't result in anyone being lost
    @test npers_from_bus == npeople
end


text:
@testset "1035.uncrossed-lines.jl" begin
    @test max_uncrossed_lines([1,4,2], [1,2,4]) == 2
    @test max_uncrossed_lines([2,5,1,2,5], [10,5,2,1,5,2]) == 3
    @test max_uncrossed_lines([1,3,7,1,7,5], [1,9,2,5,1]) == 2
end


text:
using Base.Test, ASE, NeighbourLists

# ======================================================================

println("Check that the cubic unit cell for Al is reproduced correctly")
a0 = 2.025
p = [ [0.0;0.0;0.0] [0.0;a0;a0] [a0;0.0;a0] [a0;a0;0.0] ]
at = bulk("Al", cubic=true)
@test (positions(at) |> mat) == p

# ======================================================================

println("Check neighbourlist without periodicity")
# TODO: implement a test with periodicity as well???
println("   ... assemble neighbour list ...")
at = bulk("Al", cubic=true) * 3
set_pbc!(at, (false,false,false))
rcut = 1.7 * a0
nlist = neighbourlist(at, rcut)
# create a neighbourlist via a naive double-loop
simple = zeros(length(at), length(at))
X = positions(at)
for n = 2:length(at), m = 1:n-1
   if norm(X[m]-X[n]) <= rcut
      simple[n,m] = simple[m,n] = 1
   end
end
println("   ... check the bond-iterator ... ")
pass_bonds_test = true
for (i,j,r,R) in pairs(nlist)
   if !( (simple[i,j] == 1) && (abs(norm(X[i]-X[j]) - r) < 1e-12) &&
         (norm(X[j]-X[i] - R) < 1e-12)  )
      pass_bonds_test = false
      break
   end
   # switch the flag
   simple[i,j] = -1
end
@test pass_bonds_test
# check that all pairs have been found
@test maximum(simple) == 0
# revert to original
simple *= -1
println("   ... check the site iterator ... ")
pass_site_test = true
for (i,j,r,R) in sites(nlist)
   for n = 1:length(j)
      if simple[i,j[n]] != 1
         pass_site_test = false
         break
      end
      simple[i,j[n]] = -1
   end
   if maximum(simple[i,:]) != 0
      pass_site_test = false
      break
   end
end
@test pass_site_test
@test maximum(simple) == 0

# ======================================================================

println("Checking `***_data`, `***_array`, `***_info`, `***_transient`")

at = bulk("Cu") * 3
N = length(at)
# set an array and test that it is read back correctly
z = rand(N)
set_array!(at, "z", z)
@test get_array(at, "z") == z
# set some info and test reading
i = "some info"
set_info!(at, "i", i)
@test get_info(at, "i") == i
# ***_transient should be tested automatically via the calculators.
# test the has_***
@test has_array(at, "z")
@test has_info(at, "i")
@test !has_array(at, "i")
@test !has_info(at, "z")

println("Checking momenta and velocites")
p = rand(3,N) |> vecs
set_momenta!(at, p)
@test isapprox(momenta(at) |> mat, p |> mat)
v = p ./ masses(at)
@test isapprox(velocities(at) |> mat, v |> mat)
set_velocities!(at, v/2.0)
@test isapprox(velocities(at) |> mat, v/2.0 |> mat)
@test isapprox(momenta(at) |> mat, p/2.0 |> mat)

println("Checking chemical_symbols")
@test chemical_symbols(at) == ["Cu" for i in 1:length(at)]

# ======================================================================

println("testing static_neighbourlist")
nlist = static_neighbourlist(at, rnn("Cu") * 1.1)
rattle!(at, 0.05)
nlist1 = static_neighbourlist(at, rnn("Cu") * 1.1)
@test nlist === nlist1

# ======================================================================

println("testing read and write")
fname = "test.xyz"
at = bulk("Cu") * 2
write_xyz(fname, at)
at2 = read_xyz(fname)
@assert positions(at) == positions(at2)
rm(fname)

# ---------
println("Conversion between JuLIP.Atoms and ASEAtoms")
at1 = JuLIP.bulk(:Si)    # JuLIP.Atoms
at2 = ASEAtoms(at1)      # ASE.ASEAtoms
at3 = Atoms(at2)         # JuLIP.Atoms
@assert at1 == at3


text:
mlaplacian_ref(img, dims, r) = mlaplacian_ref(img, strel_box(img, dims; r))
mlaplacian_ref(img, se) = float.(dilate(img, se)) .+ float.(erode(img, se)) .- 2 .* float.(img)

@testset "mlaplacian" begin
    test_ranges = [Bool, 1:10, Float64, Gray{N0f8}, Gray{Float64}]
    for T in test_ranges
        for N in (1, 2, 3)
            sz = ntuple(_ -> 32, N)
            img = rand(T, sz...)

            out = mlaplacian(img)
            @test out == mlaplacian(img, strel_box(img, ntuple(identity, N); r=1))
            @test out == mlaplacian_ref(img, ntuple(identity, N), 1)

            @test mlaplacian(img; dims=(1,), r=2) == mlaplacian_ref(img, (1,), 2)

            se = rand_se_mask(3, N; symmetric=true)
            @test mlaplacian(img, se) == mlaplacian_ref(img, se)
        end
    end

    img = rand(1:5, 7, 7)
    out = similar(img)
    mlaplacian!(out, img, similar(img))
    @test out == mlaplacian(img)

    img = rand(RGB, 7, 7)
    @test_throws ArgumentError("color image is not supported") mlaplacian(img)

    img = rand(1:10, 7, 7)
    se = Bool[1 0 0; 0 1 0; 0 0 0]
    msg = "structuring element must be symmetric with respect to its center"
    @test_throws ArgumentError(msg) mlaplacian(img, se)
end


text:
"""
    experiments.jl

Description:
    Experiment functions to automate training/testing by other scripts.

Author:
    Sasha Petrenko <sap625@mst.edu>

Date:
    3/24/2021
"""

# --------------------------------------------------------------------------- #
# PACKAGES
# --------------------------------------------------------------------------- #

# ICVIs pulled from the ClusterValidityIndices package
using ClusterValidityIndices

# Quality of life packages for editing and fancy logging
using Revise
using ProgressBars
using Logging
using Plots

"""
    test_cvi(cvi::C, data::Array{M, 2}, labels::Array{N, 1}) where {C<:AbstractCVI, M<:Real, N<:Int}

Test the CVI on the data and labels in all configurations (incremental, batch, porcelain).
"""
function test_cvi(cvi::C, data::Array{M, 2}, labels::Array{N, 1}, data_name::String) where {C<:AbstractCVI, M<:Real, N<:Int}
    n_samples = length(labels)

    # ----------------------------------------------------------------------- #
    # INCREMENTAL MODE
    #   Run the CVI in incremental mode
    # ----------------------------------------------------------------------- #

    # Instantiate the icvi with default options
    cvi_i = deepcopy(cvi)

    # Create some storage for our criterion values
    criterion_values_i = zeros(n_samples)

    # Iterate across all of the samples
    for ix = ProgressBar(1:n_samples)
        # Update the CVI internal parameters incrementally
        # NOTE: the package assumes that columns are features and rows are samples
        param_inc!(cvi_i, data[:, ix], labels[ix])
        # Evaluate the CVI to internally store the criterion value
        evaluate!(cvi_i)
        # Extract and save the criterion value at each step
        criterion_values_i[ix] = cvi_i.criterion_value
    end

    # ----------------------------------------------------------------------- #
    # BATCH MODE
    #   Run the CVI in batch mode
    # ----------------------------------------------------------------------- #

    # Instantiate the CVI, same as when done incrementally
    cvi_b = deepcopy(cvi)

    # Compute the parameters in batch
    param_batch!(cvi_b, data, labels)

    # Evaluate the CVI criterion value
    evaluate!(cvi_b)

    # NOTE: we only get the last criterion value because we ran in batch mode,
    #       which is accessible at cvi_b.criterion_value.

    # ----------------------------------------------------------------------- #
    # INCREMENTAL MODE: PORCELAIN FUNCTIONS
    #   Update and get the ICVI at once with the porcelain functions
    # ----------------------------------------------------------------------- #

    # Instantiate the CVI as both in incremental and batch modes
    cvi_p = deepcopy(cvi)

    # Create storage for the criterion values at each timestep
    criterion_values_p = zeros(n_samples)

    # Iterate across all samples
    for ix = ProgressBar(1:n_samples)
        # Update the CVI parameters and extract the criterion value in one function
        # NOTE: the package assumes that columns are features and rows are samples
        criterion_values_p[ix] = get_icvi!(cvi_p, data[:, ix], labels[ix])
    end

    # ----------------------------------------------------------------------- #
    # BATCH MODE: PORCELAIN FUNCTIONS
    #   Update and get the CVI at once with the porcelain functions
    # ----------------------------------------------------------------------- #

    # Instantiate the CVI as both in incremental and batch modes
    cvi_pb = deepcopy(cvi)

    # Iterate across all samples
    criterion_value_pb = get_cvi!(cvi_pb, data, labels)

    # ----------------------------------------------------------------------- #
    # VISUALIZATION
    # ----------------------------------------------------------------------- #

    # Show the last criterion value
    @info "Incremental CVI value: $(cvi_i.criterion_value)"
    @info "Batch CVI value: $(cvi_b.criterion_value)"
    @info "Porcelain Incremental CVI value: $(criterion_values_p[end])"
    @info "Porcelain Batch CVI value: $(criterion_value_pb)"

    # Plot the two incremental trends ("manual" and porcelain) atop one another
    p = plot(dpi=dpi, reuse=false)
    plot!(p, 1:n_samples, criterion_values_i)
    plot!(p, 1:n_samples, criterion_values_p)
    title!("CVI: " * string(typeof(cvi)) * ", " * data_name)
    xlabel!("Sample Index")
    ylabel!("Criterion Value")
    # display(p)
    gui(p)

    # Save the image
    # savefig("results/single_" * string(typeof(cvi)) * "_" * data_name)

end


text:
@testset "Unit tests" begin
    @testset "iterate.jl" begin
        @test Base.IteratorSize(PharmGraph) == Base.HasLength()
        @test Base.IteratorEltype(PharmGraph) == Base.HasEltype()
    end

    @testset "system.jl" begin
        @test system_matches(PharmClass("foo", "1"), "foo")
        @test system_matches("foo")(PharmClass("foo", "1"))
        @test !system_matches(PharmClass("foo", ""), "bar")
        @test !system_matches("bar")(PharmClass("foo", ""))
        @test system_matches(PharmClass("foo", "1"), r"foo")
        @test system_matches(r"foo")(PharmClass("foo", "1"))
        @test !system_matches(PharmClass("foo", "1"), r"bar")
        @test !system_matches(r"bar")(PharmClass("foo", "1"))
    end

    @testset "types.jl" begin
        @test Base.isless(PharmClass("bar", "1"), PharmClass("foo", "1"))
        @test Base.isless(PharmClass("bar", "1"), PharmClass("bar", "2"))
    end
end


text:
module Commitment

using Codex
using DataFrames

export delta2scores

reverse(ans::Number) = 8 - ans

"""
    is_2018(df::DataFrame)::Bool

Returns whether `df` contains responses to the 2018 version of the delta questionnaire.
"""
is_2018(df::DataFrame)::Bool = !("v7" in names(df))

ans_2018_mapping = Dict(
    "Helemaal mee oneens" => 1,
    "Gedeeltelijk mee oneens" => 2,
    "Niet oneens, niet eens" => 3,
    "Gedeeltelijk mee eens" => 4,
    "Helemaal mee eens" => 5
)

function is_reversed(question::Int, year::Int)::Bool
    year == 2018 ? question in [1, 2, 4] : question in [2, 3, 5]
end

"""
    ans2num(ans, question::Int, year::Int)

Note that the scale was from 1 to 5 in 2018 and 1 to 7 in later years.
"""
function ans2num(ans, question::Int, year::Int)
    ans = year == 2018 ?
        Codex.rescale(ans_2018_mapping[ans], 1, 5, 1, 7) :
        (typeof(ans) != Int ? get(ans, Int) : ans)
    return is_reversed(question, year) ? reverse(ans) : ans
end

score_2018(responses...) = sum(ans2num.(responses, 1:5, 2018))
score_2019(responses...) = sum(ans2num.(responses, 2:6, 2019))

function delta2scores(df::DataFrame)
    if is_2018(df)
        transform!(df, [:v1, :v2, :v3, :v4, :v5] => ByRow(score_2018) => :score)
        df[!, :achievable] .= ""
    else
        transform!(df, [:v2, :v3, :v4, :v5, :v6] => ByRow(score_2019) => :score)
        select!(df, :, :v7 => :achievable)
    end
    select!(df, :id, :completed_at, :score, :achievable)
    disallowmissing!(df)
end

end # module


text:
using Test, StructTypes, UUIDs, Dates

struct A
    x::Int
end

struct DictWrapper
    x::Dict
end

@enum Fruit apple orange

struct EmptyStruct
end

struct OmitEmp
    x::Union{Nothing, Int}
end

@testset "StructTypes" begin

@test StructTypes.StructType(Union{Int, Missing}) == StructTypes.Struct()
@test StructTypes.StructType(Any) == StructTypes.Struct()
@test StructTypes.StructType(A) == StructTypes.NoStructType()
@test StructTypes.StructType(A(1)) == StructTypes.NoStructType()

@test StructTypes.names(A) == ()
@test StructTypes.names(A(1)) == ()

@test StructTypes.julianame(((:a, :A), (:b, :B)), :B) == :b
@test StructTypes.serializationname(((:a, :A), (:b, :B)), :b) == :B

@test StructTypes.excludes(A) == ()
@test StructTypes.excludes(A(1)) == ()

@test StructTypes.omitempties(A) == ()
@test StructTypes.omitempties(A(1)) == ()

@test StructTypes.isempty([])
@test StructTypes.isempty(Dict())
@test StructTypes.isempty("")
@test StructTypes.isempty(())
@test StructTypes.isempty(NamedTuple())
@test !StructTypes.isempty(1)
@test StructTypes.isempty(nothing)
@test !StructTypes.isempty(A(1))
@test !StructTypes.isempty(A(1), 1)

@test StructTypes.keywordargs(A) == NamedTuple()
@test StructTypes.keywordargs(A(1)) == NamedTuple()

@test StructTypes.idproperty(A) == :_
@test StructTypes.idproperty(A(1)) == :_

@test StructTypes.StructType(Dict) == StructTypes.DictType()
@test StructTypes.StructType(NamedTuple) == StructTypes.Struct()
@test StructTypes.StructType(Pair) == StructTypes.DictType()

x = Dict(1 => 2)
@test StructTypes.keyvaluepairs(x) == pairs(x)
x = 1 => 2
@test StructTypes.keyvaluepairs(x) == (x,)

x = Dict(:hey => 2)
@test StructTypes.construct(typeof(x), x) === x
y = DictWrapper(x)
@test StructTypes.construct(DictWrapper, x) == y

@test StructTypes.construct(NamedTuple, x) == (hey=2,)
@test StructTypes.construct(NamedTuple{(:hey,)}, x) == (hey=2,)
@test StructTypes.construct(NamedTuple{(:hey,),Tuple{Int}}, x) == (hey=2,)

@test StructTypes.StructType(Vector) == StructTypes.ArrayType()
@test StructTypes.StructType(Set) == StructTypes.ArrayType()
@test StructTypes.StructType(Tuple) == StructTypes.ArrayType()

@test StructTypes.construct(Tuple, [1, 2]) == (1, 2)

@test StructTypes.StructType(String) == StructTypes.StringType()
@test StructTypes.StructType(Symbol) == StructTypes.StringType()
@test StructTypes.StructType(Enum) == StructTypes.StringType()
@test StructTypes.StructType(Char) == StructTypes.StringType()
@test StructTypes.StructType(UUID) == StructTypes.StringType()
@test StructTypes.StructType(Date) == StructTypes.StringType()

@test StructTypes.construct(Char, "1") == '1'
@test_throws ArgumentError StructTypes.construct(Char, "11")

@test StructTypes.construct(Fruit, :apple) == apple
@test StructTypes.construct(Fruit, "apple") == apple
x = "apple"
@test StructTypes.construct(Fruit, pointer(x), 5) == apple

@test StructTypes.construct(Symbol, "hey") == :hey
@test StructTypes.construct(Symbol, pointer(x), 5) == :apple
x = "499beb72-22ea-11ea-3366-55749430b981"
@test StructTypes.construct(UUID, pointer(x), 36) == UUID(x)
@test StructTypes.construct(Date, "11-30-2019"; dateformat=dateformat"mm-dd-yyyy") == Date(2019, 11, 30)

@test StructTypes.StructType(UInt8) == StructTypes.NumberType()
@test StructTypes.StructType(Int8) == StructTypes.NumberType()
@test StructTypes.StructType(Float64) == StructTypes.NumberType()
@test StructTypes.numbertype(UInt8) == UInt8
@test StructTypes.numbertype(Complex) == Float64

@test StructTypes.construct(Float64, 1) === 1.0

@test StructTypes.StructType(Bool) == StructTypes.BoolType()
@test StructTypes.construct(Int, true) == 1

@test StructTypes.StructType(Nothing) == StructTypes.NullType()
@test StructTypes.StructType(Missing) == StructTypes.NullType()
@test StructTypes.construct(Missing, nothing) === missing

@test StructTypes.subtypekey(A) == :type
@test StructTypes.subtypekey(A(1)) == :type
@test StructTypes.subtypes(A) == NamedTuple()
@test StructTypes.subtypes(A(1)) == NamedTuple()

@test StructTypes.construct(Date, Date(2020)) == Date(2020)

# https://github.com/quinnj/JSON3.jl/issues/139
v = v"1.2.3"
@test StructTypes.StructType(v) == StructTypes.StringType()

# Some wrapper
@test StructTypes.StructType(Some(1)) == StructTypes.CustomStruct()

# 47
@test StructTypes.construct(Nothing, nothing) === nothing
x = Dict("hey" => "ho")
@test StructTypes.construct(Dict{String, String}, x) === x

# omitempties(T) = true respected
StructTypes.StructType(::Type{OmitEmp}) = StructTypes.Struct()
StructTypes.omitempties(::Type{OmitEmp}) = true
StructTypes.foreachfield(OmitEmp(1)) do i, nm, T, val
    @test val == 1
end
counter = 0
StructTypes.foreachfield(OmitEmp(nothing)) do i, nm, T, val
    counter += 1
end
@test counter == 0

end

struct B
    a::Int
    b::Float64
    c::String
end

mutable struct C
    a::Int
    b::Float64
    c::String
    C() = new()
    C(a::Int, b::Float64, c::String) = new(a, b, c)
end

mutable struct D
    a::Union{Int, Nothing}
    b::Union{Float64, Nothing}
    c::Union{String, Nothing}
    D() = new(nothing, nothing, nothing)
    D(a::Union{Int, Nothing}, b::Union{Float64, Nothing}, c::Union{String, Nothing}) = new(a, b, c)
end

mutable struct E
    a::Union{Int, Missing}
    b::Union{Float64, Missing}
    c::Union{String, Missing}
    E() = new(missing, missing, missing)
    E(a::Union{Int, Missing}, b::Union{Float64, Missing}, c::Union{String, Missing}) = new(a, b, c)

end

struct LotsOfFields
    x1::String
    x2::String
    x3::String
    x4::String
    x5::String
    x6::String
    x7::String
    x8::String
    x9::String
    x10::String
    x11::String
    x12::String
    x13::String
    x14::String
    x15::String
    x16::String
    x17::String
    x18::String
    x19::String
    x20::String
    x21::String
    x22::String
    x23::String
    x24::String
    x25::String
    x26::String
    x27::String
    x28::String
    x29::String
    x30::String
    x31::String
    x32::String
    x33::String
    x34::String
    x35::String
end

mutable struct LotsOfFields2
    x1::String
    x2::String
    x3::String
    x4::String
    x5::String
    x6::String
    x7::String
    x8::String
    x9::String
    x10::String
    x11::String
    x12::String
    x13::String
    x14::String
    x15::String
    x16::String
    x17::String
    x18::String
    x19::String
    x20::String
    x21::String
    x22::String
    x23::String
    x24::String
    x25::String
    x26::String
    x27::String
    x28::String
    x29::String
    x30::String
    x31::String
    x32::String
    x33::String
    x34::String
    x35::String
    LotsOfFields2() = new()
end

mutable struct DateStruct
    date::Date
    datetime::DateTime
    time::Time
end
DateStruct() = DateStruct(Date(0), DateTime(0), Time(0))
Base.:(==)(a::DateStruct, b::DateStruct) = a.date == b.date && a.datetime == b.datetime && a.time == b.time

struct EmptyStruct end

@testset "convenience functions" begin

## StructTypes.construct
# simple struct
@test @inferred StructTypes.construct((i, nm, T) -> 1, A) == A(1)

@test StructTypes.construct(() -> nothing, EmptyStruct) == EmptyStruct()

# hetero fieldtypes struct
@test @inferred StructTypes.construct((i, nm, T) -> (1, 3.14, "hey")[i], B) == B(1, 3.14, "hey")

# simple tuple
@test @inferred StructTypes.construct((i, nm, T) -> 1, Tuple{Int}) == (1,)

# NamedTuple
@test @inferred StructTypes.construct((i, nm, T) -> (c="hey", a=1, b=3.14)[nm], NamedTuple{(:a, :b, :c), Tuple{Int, Float64, String}}) == (a=1, b=3.14, c="hey")

# hetero fieldtypes tuple
@test @inferred StructTypes.construct((i, nm, T) -> (1, 3.14, "hey")[i], Tuple{Int, Float64, String}) == (1, 3.14, "hey")

# > 32 fields
vals = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
@test StructTypes.construct((i, nm, T) -> vals[i], LotsOfFields) == LotsOfFields(vals...)

StructTypes.names(::Type{A}) = ((:x, :y),)
StructTypes.construct(A) do i, nm, T
    @test nm == :y
    1
end

StructTypes.keywordargs(::Type{DateStruct}) = (date=(dateformat=dateformat"mm/dd/yyyy",),)
f1(i, nm, T; dateformat=Dates.ISODateFormat) = nm == :date ? Date("11/23/1961", dateformat) : nm == :datetime ? DateTime(0) : Time(0)
@test StructTypes.construct(f1, DateStruct) == DateStruct(Date(1961, 11, 23), DateTime(0), Time(0))

StructTypes.names(::Type{LotsOfFields}) = ((:x35, :y35),)
StructTypes.keywordargs(::Type{LotsOfFields}) = (x35=(hey=:ho,),)
function f2(i, nm, T; hey=:hey)
    if i == 35
        @test nm == :y35
        @test hey == :ho
    end
    return vals[i]
end
@test StructTypes.construct(f2, LotsOfFields) == LotsOfFields(vals...)

## StructTypes.foreachfield
StructTypes.foreachfield(A(1)) do i, nm, T, v
    @test v == 1
    @test nm == :y
end
StructTypes.foreachfield((i, nm, T, v) -> @test((1, 3.14, "hey")[i] == v), B(1, 3.14, "hey"))
function f3(i, nm, T, v; hey=:hey)
    if i == 35
        @test nm == :y35
        @test hey == :ho
    end
    @test vals[i] == v
end
StructTypes.foreachfield(f3, LotsOfFields(vals...))

@test StructTypes.foreachfield(() -> nothing, EmptyStruct()) === nothing
@test StructTypes.mapfields!(() -> nothing, EmptyStruct()) === nothing

x = C()
StructTypes.mapfields!((i, nm, T) -> (1, 3.14, "hey")[i], x)
x2 = C(1, 3.14, "hey")
@test x.a == x2.a && x.b == x2.b && x.c == x2.c

x = LotsOfFields2()
StructTypes.mapfields!((i, nm, T) -> vals[i], x)
x2 = LotsOfFields(vals...)
StructTypes.foreachfield((i, nm, T, v) -> @test(getfield(x, i) == getfield(x2, i)), x)

@test StructTypes.applyfield!((i, nm, T) -> "ho", x, :x1)
@test x.x1 == "ho"
# non-existant field
@test !StructTypes.applyfield!((i, nm, T) -> "ho", x, :y1)

x2 = C(1, 3.14, "")
StructTypes.omitempties(::Type{C}) = (:c,)
all_i = Int[]
all_nm = Symbol[]
StructTypes.foreachfield(x2) do i, nm, T, v
    push!(all_i, i)
    push!(all_nm, nm)
end
@test sort(all_i) == [1, 2]
@test sort(all_nm) == [:a, :b]
StructTypes.excludes(::Type{C}) = (:a,)
all_i = Int[]
all_nm = Symbol[]
StructTypes.foreachfield(x2) do i, nm, T, v
    push!(all_i, i)
    push!(all_nm, nm)
end
@test sort(all_i) == [2]
@test sort(all_nm) == [:b]
# field isn't applied if excluded
@test !StructTypes.applyfield!((i, nm, T) -> "ho", x2, :a)

# field isn't set if excluded
x2.a = 10
StructTypes.mapfields!((i, nm, T) -> (1, 3.14, "hey")[i], x2)
@test x2.a == 10 && x2.b == 3.14 && x2.c == "hey"

# NamedTuple
@test StructTypes.construct((i, nm, T) -> (1, 3.14, "hey")[i], NamedTuple{(:a, :b, :c), Tuple{Int64, Float64, String}}) == (a=1, b=3.14, c="hey")

x3 = D(nothing, 3.14, "")
@inline StructTypes.omitempties(::Type{D}) = true
all_i = Int[]
all_nm = Symbol[]
StructTypes.foreachfield(x3) do i, nm, T, v
    push!(all_i, i)
    push!(all_nm, nm)
end
@test sort(all_i) == [2]
@test sort(all_nm) == [:b]

x4 = D(nothing, 3.14, "")
@inline StructTypes.omitempties(::Type{D}) = false
all_i = Int[]
all_nm = Symbol[]
StructTypes.foreachfield(x3) do i, nm, T, v
    push!(all_i, i)
    push!(all_nm, nm)
end
x5 = C(1, 3.14, "helloworld")
@inline StructTypes.omitempties(::Type{C}) = (:b,)
@inline StructTypes.excludes(::Type{C}) = ()
@inline StructTypes.isempty(::Type{T}, ::Number) where {T <: C} = false
all_i = Int[]
all_nm = Symbol[]
StructTypes.foreachfield(x5) do i, nm, T, v
    push!(all_i, i)
    push!(all_nm, nm)
end
@test sort(all_i) == [1, 2, 3]
@test sort(all_nm) == [:a, :b, :c]
@inline StructTypes.isempty(::Type{T}, x::Number) where {T <: C} = x > 0
all_i = Int[]
all_nm = Symbol[]
StructTypes.foreachfield(x5) do i, nm, T, v
    push!(all_i, i)
    push!(all_nm, nm)
end
@test sort(all_i) == [1, 3]
@test sort(all_nm) == [:a, :c]

x6 = E(1, missing, "")
@inline StructTypes.omitempties(::Type{E}) = (:b,)
@inline StructTypes.isempty(::Type{T}, ::Missing) where {T <: E} = false
all_i = Int[]
all_nm = Symbol[]
StructTypes.foreachfield(x6) do i, nm, T, v
    push!(all_i, i)
    push!(all_nm, nm)
end
@test sort(all_i) == [1, 2, 3]
@test sort(all_nm) == [:a, :b, :c]

@inline StructTypes.isempty(::Type{T}, ::Missing) where {T <: E} = true
all_i = Int[]
all_nm = Symbol[]
StructTypes.foreachfield(x6) do i, nm, T, v
    push!(all_i, i)
    push!(all_nm, nm)
end
@test sort(all_i) == [1, 3]
@test sort(all_nm) == [:a, :c]

CNT = Ref(0)
StructTypes.foreachfield((args...) -> CNT[] += 1, EmptyStruct)
@test CNT[] == 0

@test !StructTypes.applyfield((x...) -> 10, A, :z)

@testset "StructTypes.constructfrom(T, ::Vector{Any})" begin

end

@testset "issue 22 (`applyfield!` skips the 32nd field)" begin
    function f4(i::Integer, name::Symbol, ::Type{FT}) where FT
        return "$(i)"
    end
    x7 = LotsOfFields2()
    for name in fieldnames(typeof(x7))
        @test !isdefined(x7, name)
    end
    StructTypes.mapfields!(f4, x7)
    for name in fieldnames(typeof(x7))
        @test isdefined(x7, name)
    end
    for name in fieldnames(typeof(x7))
        @test StructTypes.applyfield!(f4, x7, name)
    end
    x8 = LotsOfFields2()
    for name in fieldnames(typeof(x8))
        @test StructTypes.applyfield!(f4, x8, name)
    end
end

end

abstract type Vehicle end

struct Car <: Vehicle
    type::String
    make::String
    model::String
    seatingCapacity::Int
    topSpeed::Float64
end

struct Truck <: Vehicle
    type::String
    make::String
    model::String
    payloadCapacity::Float64
end

StructTypes.StructType(::Type{Vehicle}) = StructTypes.AbstractType()
StructTypes.subtypes(::Type{Vehicle}) = (car=Car, truck=Truck)

StructTypes.StructType(::Type{Car}) = StructTypes.Struct()
StructTypes.StructType(::Type{Truck}) = StructTypes.Struct()

mutable struct C2
    a::Int
    b::Float64
    c::String
    C2() = new()
    C2(a::Int, b::Float64, c::String) = new(a, b, c)
end

StructTypes.StructType(::Type{C2}) = StructTypes.Mutable()

@testset "makeobj" begin
    @testset "makeobj" begin
        cases = [
            (Any,               String,            "foo"),
            (AbstractString,    String,            "foo"),
            (String,            String,            "foo"),
            (SubString,         SubString{String}, "foo"),
            (SubString{String}, SubString{String}, "foo"),
            (Any,               Int,               1),
            (Real,              Int,               1),
            (Integer,           Int,               1),
            (Int,               Int,               1),
            (Any,               Vector{Int},       [1, 2, 3]),
            (Vector,            Vector{Any},       [1, 2, 3]),
            (Vector{Any},       Vector{Any},       [1, 2, 3]),
            (Vector{Real},      Vector{Real},      [1, 2, 3]),
            (Vector{Integer},   Vector{Integer},   [1, 2, 3]),
            (Vector{Int},       Vector{Int},       [1, 2, 3]),
            (Any,               Dict{Symbol, Int}, Dict(:a => 1, :b => 2)),
            (AbstractDict,      Dict{Any, Any}, Dict(:a => 1, :b => 2)),
            (Dict,              Dict{Any, Any}, Dict(:a => 1, :b => 2)),
            (Dict{Any, Any},    Dict{Any, Any},    Dict(:a => 1, :b => 2)),
            (Dict{Symbol, Any}, Dict{Symbol, Any}, Dict(:a => 1, :b => 2)),
            (Dict{Any, Int},    Dict{Any, Int},    Dict(:a => 1, :b => 2)),
            (Dict{Symbol, Int}, Dict{Symbol, Int}, Dict(:a => 1, :b => 2)),
        ]
        for case in cases
            a = case[1]
            b = case[2]
            c = case[3]
            @test typeof(StructTypes.constructfrom(a, c)) == b
            @test StructTypes.constructfrom(a, c) == c
        end
    end
    @testset "mutable structs" begin
        @testset "constructfrom" begin
            input = Dict(:a => 1, :b => 2.0, :c => "three")
            output = StructTypes.constructfrom(C2, input)
            @test typeof(output) === C2
            @test output.a == 1
            @test output.b == 2.0
            @test output.c == "three"
            dict = Dict(:type => "car", :make => "Mercedes-Benz", :model => "S500", :seatingCapacity => 5, :topSpeed => 250.1)
            car = StructTypes.constructfrom(Vehicle, dict)
            @test typeof(car) == Car
            @test car.make == "Mercedes-Benz"
        end
        @testset "constructfrom!" begin
            input = Dict(:a => 1, :b => 2.0, :c => "three")
            x = C2()
            output = StructTypes.constructfrom!(x, input)
            @test x === output
            @test typeof(x) === C2
            @test x.a == 1
            @test x.b == 2.0
            @test x.c == "three"
            @test StructTypes.constructfrom(Tuple{Int64, Float64, String}, Union{Int64, Float64, String}[Int64(1), 2.0, "three"]) == (1, 2.0, "three")
            @test StructTypes.constructfrom(NamedTuple{(:a, :b, :c), Tuple{Int64, Float64, String}}, input) == (a=1, b=2.0, c="three")
            @test StructTypes.constructfrom(NamedTuple{(:a, :b, :c), Tuple{Int64, Float64, String}}, x) == (a=1, b=2.0, c="three")
        end
    end
    @testset "named tuples" begin
        cases = [
            (NamedTuple,                                    (;a=1), (;a=1)),
            (@NamedTuple{a::Int},                           (;a=1), (;a=1)),
            (@NamedTuple{b::Union{Int, Nothing}},           (;a=1),
                                        @NamedTuple{b::Union{Int, Nothing}}((nothing,))),
            (@NamedTuple{b::Union{Int, Nothing}, c::Int},   (;a = 1, c = 3),
                                        @NamedTuple{b::Union{Int, Nothing}, c::Int}((nothing, 3))),
            (@NamedTuple{b::Union{Int, Nothing}, c::Int},   (;a = 1, b = 2, c = 3),
                                        @NamedTuple{b::Union{Int, Nothing}, c::Int}((2, 3))),
        ]
        for (given_type, given_data, expected_data) in cases
            data = StructTypes.constructfrom(given_type, given_data)
            @test typeof(data) == typeof(expected_data)
            @test data == expected_data
        end
    end
end

struct MyStruct1
    val::Int
end
StructTypes.@Struct MyStruct1
StructTypes.@NullType struct MyStruct2
    val::Int
end

StructTypes.@Mutable mutable struct MyStruct3
    val::Int
end

@testset "Macros" begin
    @test StructTypes.StructType(MyStruct1) isa StructTypes.Struct
    @test StructTypes.StructType(MyStruct2) isa StructTypes.NullType
    @test StructTypes.StructType(MyStruct3) isa StructTypes.Mutable

    # Test an expression that is not a struct def
    expr = StructTypes.macro_constructor(Expr(:a), StructTypes.Struct)
    @test_throws ArgumentError eval(expr)

    # Test a symbol that is not a type
    expr = StructTypes.macro_constructor(:var, StructTypes.Struct)
    eval(quote
        var = 1
    end)
    @test_throws ArgumentError eval(expr.args[1]) # Extract body from within escape
end


text:
path = @__DIR__
cmap_names = readdir(joinpath(path,"data"))

for name in cmap_names
    @eval begin
        global $(Symbol(name)) = ColorMap($name,readdlm(joinpath(path,"data",$name,$name*".txt")))
    end
end


text:
push!(LOAD_PATH,"../src")

using Test
using TypesParticles
using NuclearRepulsion

"""
	run_nuclearrepulsion()

""" function run_nuclearrepulsion_test()
	 @testset "NuclearRepulsion" begin
	 #H2 structure
	 x0,y0,z0 = 0.00e0,0.00e0,0.00e0;
	 x1,y1,z1 = 0.77e0,0.00e0,0.00e0;
	 Z0,Z1 = 1,1;
	 H1 = Atom(x0,y0,z0,Z0);
	 H2 = Atom(x1,y1,z1,Z1);

	 cell = [1.00e0 0.00e0 0.00e0;
	      	 0.00e0 1.00e0 0.00e0;
		 0.00e0 0.00e0 1.00e0];

             @test begin
	         atomlist = [H1,H2];
	         system = SystemOfAtoms(2,atomlist,cell);
	         calcnuclrepul(system) == 1.2987012987012987e0; #Atomic Units
                 1 == 1
             end
        end #testset
end #run_nuclearrepulsion_test


text:
export SeisBandPass,
SeisDecimate,
SeisDelay,
SeisDiff,
SeisEnvelope,
SeisFKFilter,
SeisFXDecon,
SeisGain,
SeisKolmogoroff,
SeisMWNI,
SeisMute,
SeisNMO,
SeisPOCS,
SeisPWD,
SeisRadonFreqFor,
SeisRadonFreqInv,
SeisSemblance,
SeisSmooth1,
SeisSmooth2,
SeisSmoothGathers,
SeisSmoothStructure,
SeisStack,
SeisSincInerp1D
include("SeisBandPass.jl")
include("SeisDecimate.jl")
include("SeisDelay.jl")
include("SeisDiff.jl")
include("SeisEnvelope.jl")
include("SeisFKFilter.jl")
include("SeisFXDecon.jl")
include("SeisGain.jl")
include("SeisKolmogoroff.jl")
include("SeisMWNI.jl")
include("SeisMute.jl")
include("SeisNMO.jl")
include("SeisPOCS.jl")
include("SeisPWD.jl")
include("SeisRadonFreqFor.jl")
include("SeisRadonFreqInv.jl")
include("SeisSemblance.jl")
include("SeisSmooth1.jl")
include("SeisSmooth2.jl")
include("SeisSmoothGathers.jl")
include("SeisSmoothStructure.jl")
include("SeisStack.jl")
include("SeisSincInterp1D.jl")


text:
# ---------------------------------------------------------------------------- #
#
#   CDR.jl
#
#   Struct for convection-diffusion-reaction problems
#
#   λυτέος
#   Fall 2017
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

"""
    CDR

CDR type:
Used as placeholder for information about convection-diffusion-reaction problem,
e.g name, source function, and boundary conditions.
"""
struct CDR{S<:Function,B<:Function} <: Problem

  name::String        # Name of problem

  κ::Float64          # Diffusivity coefficient
  c::Vector{Float64}  # Convective velocity

  source::S               # Rhs of problem

  bctype::Vector{Int64}   # Relates each boundary to a boundary function type
  bcnorm::Bool            # Whether or not boundary functions are defined normal to the surface
  bcfunc::Vector{B} # Boundary function for each boundary face

end

"""
    CDR(src::Function, bcf::Array{Function}, bctype::Array{Int64};
        name = "Answer to Life Universe and Everything", bcnorm = false)

Constructor for `Problem` type. `src` is a function that returns the source
function for the problem. `bcf` is an Array of Functions, which specifies the
boundary condition on each boundary surface.
"""
function CDR(κ::Float64,
             c::Vector{Float64},
             src::Function,
             bcf::Vector{Function},
             bctype::Vector{Int64};
             name = "Answer to Life Universe and Everything", bcnorm = false)

  if length(bctype) != length(bcf)
    error("Problem:: bctype needs have the same size as bcf")
  end

  CDR(name, κ, c, src, bctype, bcnorm, bcf)

end


text:
# Fit the model in a more pleasant way
function anova_bnp_bernoulli(
    y::Vector{Bool},
    X::Matrix{Int};
    iter::Int = 4000, # taken from rstan::stan()
    warmup::Int = 2000, # taken from rstan::stan()
    seed::Int = 1, # taken from rstan::stan()
    zeta0::Float64 = 1.0,
    a0::Float64 = 2.0,
    b0::Float64 = 4.0,
    a1::Float64 = 2.0,
    b1::Float64 = 4.0,
)
    # Set data for training
    data0 = BernoulliData(X, y)

    # Set data for prediction
    G = length(data0.Xunique)
    lb = minimum(y)
    ub = maximum(y)
    y1 = lb:ub |> x -> repeat(x, inner = G)
    X1 = repeat(vcat(data0.Xunique'...), ub - lb + 1)
    data1 = BernoulliData(X1, y1)

    # Initialize the model
    N = length(y)
    rng = MersenneTwister(seed)
    m = BernoulliDDP(rng, N, G; αa0 = a0, αb0 = b0, a0 = a1, b0 = b1, ζ0 = zeta0)

    # Train the model
    ch = train(rng, m, data0, data1; iter, warmup);

    # Compute p(γ | y)
    group_probs = gamma_posterior(ch);

    # Compute the codebook
    group_codes = gamma_codebook(data0);

    # Compute the effects
    effects1 = simple_effect_probabilities(ch, data0)
    effects2 = interaction_effect_probabilities(ch, data0)

    # Compute p(y0 | y)
    fpost = DataFrame(group = data1.x, y = data1.y, f = mean(ch.f))

    return anova_bnp_fitted(
        group_codes,
        group_probs,
        effects1,
        effects2,
        fpost
    )
end

# # Example
# N = 1000;
# rng = MersenneTwister(1);
# y, X = simulate_sample_poisson(rng, N);
# data1 = anova_bnp_poisson(y, X);
# group_probs = gamma_posterior(data1)
# show(group_probs; allrows = true)


text:
###
#   An implmentation of Transitional Markov Chain Monte Carlo in Julia
#
#            Institute for Risk and Uncertainty, Uni of Liverpool
#
#                       Authors: Ander Gray, Adolphus Lye
#
#                       Email: Ander.Gray@liverpool.ac.uk,
#                              Adolphus.Lye@liverpool.ac.uk
#
#
#   This Transitional MCMC algorithm is inspirted by OpenCOSSAN:
#                   https://github.com/cossan-working-group/OpenCossan
#
#
#   Algorithm originally proposed by:
#           J. Ching, and Y. Chen (2007). Transitional Markov Chain Monte Carlo method
#           for Bayesian model updating, Model class selection, and Model averaging.
#           Journal of Engineering Mechanics, 133(7), 816-832.
#           doi:10.1061/(asce)0733-9399(2007)133:7(816)
#
###
module TransitionalMCMC

using Distributed, StatsBase, Distributions, Logging, CovarianceEstimation

export tmcmc, metropolis_hastings, metropolis_hastings_simple

include("tmcmc.jl")
include("mcmc.jl")
end # module


text:
# Code taken from FluxML/Optimisers.jl

state_init(opti, args...) = nothing

"""
    Descent(η)

Classic gradient descent optimiser with learning rate `η`.
For each parameter `p` and its gradient `δp`, this runs `p -= η*δp`.
"""
mutable struct Descent
  eta::Float64
end

function apply(o::Descent, x, x̄, state = nothing)
  x̄ .* o.eta, state
end

struct Identity end

function apply(o::Identity, x, x̄, state = nothing)
  if x === x̄
    return 0, state
  else
    return x .+ x̄, state
  end
  x̄, state
end

"""
    Nesterov(η)

Classic gradient descent optimiser with learning rate `η` accellerated
with parameters "μ"
"""
mutable struct Nesterov
  lr::AbstractFloat
  μ::AbstractFloat
  gclip::AbstractFloat
end

state_init(o::Nesterov, x) = IdDict()

function apply(o::Nesterov, x, Δ, state)
  lr = o.lr
  μ  = o.μ

  # get old velocity, and if not present, set it to zero.
  v = get!(state, x, zero(x))

  d = @. μ^2 * v - (1+μ) * lr * Δ
  @. v = μ*v - lr*Δ

  -d, state
end

"""
    ∇clip!(∇W, clip_val)

clips the norm of the vector `∇W` to clip_val if it is bigger than that,
otherwise leaves `∇W` unchanged
"""
function ∇clip!(∇W::AbstractArray, clip_val)
    if clip_val == 0
        ∇W
    else
        ∇norm = norm(∇W)
        if ∇norm <= clip_val
            return ∇W
        else
            return lmul!(clip_val/∇norm, ∇W)
        end
    end
end


text:
function get(link::Link, ::Type{Any})
    # TODO: handle big integers/floats
    t = getnextraw(link)
    if t == TK_INT
        get(link, WInteger)
    elseif t == TK_INT32_LE
        get(link, Int)
    elseif t == TK_INT64_LE
        get(link, Int64)
    elseif t == TK_FUNC || t == TK_ARRAY
        get(link, WExpr)
    elseif t == TK_STR
        get(link, String)
    elseif t == TK_REAL
        get(link, WReal)
    elseif t == TK_FLOAT64_LE
        get(link, Float64)
    elseif t == TK_SYM
        get(link, WSymbol)
    elseif t == TK_ERROR
        throw(MathLinkError(link))
    else
        error("Unsupported data type $t ($(Int(t)))")
    end
end

function get(link::Link, ::Type{WExpr})
    nargs = getargcount(link)
    head = get(link, Any)
    args = [get(link, Any) for i = 1:nargs]
    WExpr(head, args)
end
function put(link::Link, w::WExpr)
    puttype(link, TK_FUNC)
    putargcount(link, length(w.args))
    put(link, w.head)
    for arg in w.args
        put(link, arg)
    end
    nothing
end

function put(link::Link, list::Union{AbstractVector, Tuple})
    putfunction(link, WSymbol("List"), length(list))
    for x in list
        put(link, x)
    end
end

function get(link::Link, ::Type{T}) where {T<:Tuple}
    F = fieldtypes(T)
    nargs = getargcount(link)
    _ = get(link, Any)
    @assert nargs == length(F)
    map(S -> get(link, S), F)
end
function get(link::Link, ::Type{Vector{T}}) where {T}
    nargs = getargcount(link)
    _ = get(link, Any)
    T[get(link, T) for i = 1:nargs]
end
function get(link::Link, ::Type{Vector})
    nargs = getargcount(link)
    _ = get(link, Any)
    [get(link, Any) for i = 1:nargs]
end

function get(link::Link, ::Type{BigInt})
    parse(BigInt, get(link, WInteger).value)
end


text:
#=
add_hydrogen_processor:
- Julia version:
- Author: Dan
- Date: 2021-07-03
=#

"""
    addHydrogen(composite::AbstractComposite)
Adds Hydrogen to a `composite`. First, the [`SSSR`](@ref) is calculated and aromaticity of the composite
is analyzed, then Hydrogen will be added.
"""
addHydrogen(composite::KernelInterface) = begin
    for residue in eachResidue(composite)
        placePeptideBondH_()
    end

    atom_nr::Int64 = 1
    last_atom::Atom = Atom()
    for atom in eachAtom(composite)
        ring_atoms_::Set{Atom}

        atom_nr_ += 1
        last_atom_ = atom

        if length(getBonds(atom)) ==1 && isAtomatic(first(values(getBonds(atom))))
            continue
        end

        connectivity = getConnectivity(atom)
        sum__bond_orders = countBondOrders(atom)

    end

end


text:
"""
    drophighbits(n, x::UInt8) -> UInt8
### Examples
    julia> drophighbits(0, 0b11111111)
    255

    julia> drophighbits(1, 0b11111111) |> FileToRigidTransform.bytestring
"""
function drophighbits(n::Int, x::T) where {T<:Unsigned}
    (x << n) >> n
end
"Logical shift right n bits"
droplowbits(n::Int, x::T) where {T<:Unsigned} = (x >>> n)

function bitlength_ceiling(bitlength::Int)
    if bitlength <= 8
        8
    elseif bitlength <=16
        16
    elseif bitlength <=32
        32
    elseif bitlength <=64
        64
    elseif bitlength <=128
        128
    else
        error("Can't work with > 128 bits, bitlength > 114")
    end
end

function extract_type(bitlength::Int)
    if bitlength == 8
        UInt8
    elseif bitlength == 16
        UInt16
    elseif bitlength == 32
        UInt32
    elseif bitlength == 64
        UInt64
    elseif bitlength == 128
        UInt128
    else
        error("No unsigned type with $bitlength bits")
    end
end

"""
    extract(bv::Bytevector, bytepos, bitpad, bitlength; bigendian = false) -> UIntXXX

Pull positive integer numbers from compact bit sequences.
'Bitpad' is the number of bits to ignore at the start of the first byte.
"""
function extract(bv::Bytevector, bytepos, bitpad, bitlength; bigendian = true)
    @assert bitpad < 8  "At bytepos =  $bytepos bitpad = $bitpad bitlength = $bitlength is above range 0:7"
    @assert bitpad >= 0  "At bytepos =  $bytepos bitpad = $bitpad bitlength = $bitlength is below range 0:7"
    @assert bytepos >= 1  "At bytepos =  $bytepos bitpad = $bitpad bitlength = $bitlength, bytepos is below range 1:$length(bv)"
    @assert bytepos <= length(bv)  "At bytepos =  $bytepos bitpad = $bitpad bitlength = $bitlength, bytepos is above range 1:$length(bv)"
    @assert (bytepos - 1) * 8 + bitpad + bitlength <= length(bv) * 8  "At bytepos =  $bytepos bitpad = $bitpad bitlength = $bitlength: Can't read beyond byte $(length(bv))"
    # Smallest integer larger than or equal to (bitlength + bitpad) / 8
    bytesourcelength = cld(bitlength + bitpad, 8)
    if bytesourcelength > 1
        extractbitlength = bitlength_ceiling(bitlength)
        Extracttype = extract_type(extractbitlength)
        tempminbitlength = bitlength + 14
        tempbitlength = bitlength_ceiling(tempminbitlength)
        Temptype = extract_type(tempbitlength)
        extractwithneighbours = Temptype(0)
        drop_low_bits_from_lastbyte = bytesourcelength * 8 - bitpad - bitlength
        for i = 0:(bytesourcelength-1)
             # i = 0: pick the least significant byte
            byteno = if bigendian
                    bytepos + i
                else
                    bytepos + bytesourcelength -1 -i
                end
            thisbyte = bv[byteno]
            # If this byte is at one of the ends of extracted data, set the inconsequential bits to zero
            maskedbyte = if byteno == bytepos
                if bigendian
                    thisbyte << bitpad
                else
                    (thisbyte << bitpad ) >>> bitpad
                end
            elseif byteno == bytepos + bytesourcelength - 1
                if bigendian
                    thisbyte >>> drop_low_bits_from_lastbyte
                else
                    thisbyte >>> drop_low_bits_from_lastbyte << drop_low_bits_from_lastbyte
                end
            else
                thisbyte
            end
            extractwithneighbours += Temptype(maskedbyte * 2^(i * 8))
        end
        extract_no_lowbit_neighbours = if bigendian
            extractwithneighbours >> bitpad
        else
            extractwithneighbours >> drop_low_bits_from_lastbyte
        end
        Extracttype(extract_no_lowbit_neighbours)
    else
        @assert bytepos <= length(bv)  "bytepos =  $bytepos bitpad = $bitpad\n\texceeds $(length(bv))"
        thisbyte = bv[bytepos]
        x = drophighbits(bitpad, thisbyte)
        drop_n_lowbit = 8 - bitpad - bitlength
        @assert drop_n_lowbit >= 0  "Invalid drop_n_lowbit.... bytepos =  $bytepos bitpad = $bitpad bitlength = $bitlength"
        UInt8(droplowbits(drop_n_lowbit, x))
    end
end


text:
export safe

import Base: show

"""
    safe()

Creates a `SafeOperator`, which wraps `on_subscribe!` and each `next!`, `error!` and `complete!` callbacks into try-catch block.

# Producing

Stream of type `<: Subscribable{L}` where `L` refers to type of source stream

See also: [`AbstractOperator`](@ref), [`InferableOperator`](@ref), [`ProxyObservable`](@ref), [`logger`](@ref)
"""
safe() = SafeOperator()

struct SafeOperator <: InferableOperator end

function on_call!(::Type{L}, ::Type{L}, operator::SafeOperator, source) where L
    return proxy(L, source, SafeProxy())
end

operator_right(operator::SafeOperator, ::Type{L}) where L = L

struct SafeProxy <: ActorSourceProxy end

actor_proxy!(::Type{L}, proxy::SafeProxy, actor::A)   where { L, A } = SafeActor{L, A}(actor, SafeActorProps(false, voidTeardown))
source_proxy!(::Type{L}, proxy::SafeProxy, source::S) where { L, S } = SafeSource{L, S}(source)

mutable struct SafeActorProps
    isfailed     :: Bool
    subscription :: Teardown
end

struct SafeActor{L, A} <: Actor{L}
    actor :: A
    props :: SafeActorProps
end

isfailed(actor::SafeActor)   = actor.props.isfailed
setfailed!(actor::SafeActor) = actor.props.isfailed = true

getsubscription(actor::SafeActor)         = actor.props.subscription
setsubscription!(actor::SafeActor, value) = actor.props.subscription = value

function on_next!(actor::SafeActor{L}, data::L) where L
    if !isfailed(actor)
        try
            next!(actor.actor, data)
        catch exception
            dispose!(actor)
            error!(actor.actor, exception)
        end
    end
end

function on_error!(actor::SafeActor, err)
    if !isfailed(actor)
        try
            error!(actor.actor, err)
        catch exception
            dispose!(actor)
            error!(actor.actor, exception)
        end
    end
end

function on_complete!(actor::SafeActor)
    if !isfailed(actor)
        try
            complete!(actor.actor)
        catch exception
            dispose!(actor)
            error!(actor.actor, exception)
        end
    end
end

function dispose!(actor::SafeActor)
    setfailed!(actor)
    unsubscribe!(getsubscription(actor))
end

struct SafeSource{L, S} <: Subscribable{L}
    source :: S
end

function on_subscribe!(source::SafeSource, actor::SafeActor)
    try
        subscription = subscribe!(source.source, actor)
        setsubscription!(actor, subscription)
        return subscription
    catch exception
        dispose!(actor)
        error!(actor.actor, exception)
        return voidTeardown
    end
end

Base.show(io::IO, ::SafeOperator)          = print(io, "SafeOperator()")
Base.show(io::IO, ::SafeProxy)             = print(io, "SafeProxy()")
Base.show(io::IO, ::SafeActor{L})  where L = print(io, "SafeActor($L)")
Base.show(io::IO, ::SafeSource{L}) where L = print(io, "SafeSource($L)")


text:
### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 14d9371c-9a19-11eb-2c1a-8b680938c549
begin
	using Cairo
	using Compose
	using CSV
	using LightGraphs
	using GraphPlot
	using StatsPlots
	using DataFrames
	using Statistics
	using Base.Threads
end

# ╔═╡ 7c9d6b8c-68f0-4481-b042-b51e217a8abd
md"""
# Stochastic Streets Riddler Solution

In this post, we'll solve the [Riddler Problem](https://fivethirtyeight.com/features/can-you-navigate-the-one-way-streets/) from April 9th	 2021.

In this problem, we're asked to consider a square gridded city with randomly oriented one way streets, and asked to determine the probability of being able to navigate from the upper left-hand corner of the city to the lower right-hand corner on the road network.

In the original problem statement, we're asked to consider a city 2 blocks across by 2 blocks high... but I will consider cities of arbitrary size. In this notebook, I provide a simulation for arbitrary sized cities, and a lazy exact solution for small cities.

*Note*: Since it makes notation a bit simpler, I will use $n$ to denote the number of intersections across the city. That is $n= n_\text{blocks} + 1$ where $n_\text{blocks}$ is the width of the city in blocks.
"""

# ╔═╡ 891b9894-11b0-405f-aa85-232cc635db9a
md"""
## Estimate Probability via Simulation

First, we'll compute the probability that a path exists from the top left-hand corner of the city to the lower right-hand corner of the city via simulation.
"""

# ╔═╡ 03792d01-268a-451b-8d48-13afdf6f379f
function city_has_path(g, start=1, finish=nv(g))
	has_path_start_finish = has_path(g, start, finish)

	return has_path_start_finish
end

# ╔═╡ 713a8126-0a4d-4373-ad57-6872e6aed059
md"""
### Example Random Street Layouts
"""

# ╔═╡ caf46b55-87cf-4351-9b6a-f729cce80251
md"""
### Simulation Results
"""

# ╔═╡ 5327b953-268a-4b6d-9bf1-ce513f666b56
md"""
## Lazy Analytic Solution (Small Cities)

Now, we'll compute the probability that a path exists from the top left-hand corner of the city to the lower right-hand corner of the city via simulation. We do this in a brute force combinatoric way... by simply enumerating all possible city layouts, and counting how many have the desired path. Since the number of paths grows exponentially in city size, we'll only compute the result for small cities, i.e. $n = 2,3,4$.

There's probably a smarter way to do this... but I couldn't think of it 🙂
"""

# ╔═╡ fd5d8a0c-9a36-11eb-37af-b323787108b5
md"""
## Utility Functions and Package Imports
"""

# ╔═╡ 38abf0ba-c6ea-4f45-acf1-b92eebe462a8
function neighbors(i, j, n)
	ilo = max(1, i-1)
	ihi = min(n, i+1)

	jlo = max(1, j-1)
	jhi = min(n, j+1)

	return vcat(
		[(i_neighb, j) for i_neighb = ilo:ihi if (i_neighb, j) != (i, j)],
		[(i, j_neighb) for j_neighb = jlo:jhi if (i, j_neighb) != (i, j)]
	)
end

# ╔═╡ b5b7e758-e4de-4666-976b-583a732e8904
function tuple_to_idx(i, j, n)
	return n * (i - 1) + j
end

# ╔═╡ f05db774-9741-475b-a6fc-5df56154faa4
function random_street_graph(n)
	g = SimpleDiGraph(n^2)

	# create set of all adjacent nodes
	adjacencies = Set([Set([(i, j), neighb]) for i = 1:n for j = 1:n for neighb
						in neighbors(i, j, n)])

	for adjacent_pair in adjacencies
		# unpack nodes and get their indices
		(i1, j1), (i2, j2) = adjacent_pair

		k1 = tuple_to_idx(i1, j1, n)
		k2 = tuple_to_idx(i2, j2, n)

		# connect index of node k1 to k2 if 1, else k2 to k1
		direction = rand((0,1))

		# connect node k1 to k2 in the indicated direction
		if direction > 0
			add_edge!(g, k1, k2)
		else
			add_edge!(g, k2, k1)
		end
	end

	return g
end

# ╔═╡ fb2302ca-8b2c-4e23-8dc6-36cc353b16cc
function simulate_city_planning(n, n_simulations)
	rslts = zeros(Bool, n_simulations)

	Threads.@threads for i ∈ 1:n_simulations
		g = random_street_graph(n)

		has_path = city_has_path(g)
		rslts[i, :] = [has_path]
	end

	return DataFrame(has_path=rslts[:,1])
end

# ╔═╡ 285ee664-4c13-4a0b-8c0e-c9b9e5120078
begin
	n_sim = 10^6
	n_max = 20
	prob_sim_rslts = zeros(n_max - 1)

	for n = 2:n_max
		df_city_planning_sim = simulate_city_planning(n, n_sim)
		prob_sim_rslts[n-1] = sum(df_city_planning_sim[!, :has_path]) / n_sim
	end

	df_prob_sim = DataFrame(city_width = 2:n_max,
							prob_path = prob_sim_rslts)
end

# ╔═╡ cec180f1-68f3-42bc-a5d1-4a7eeca01ab4
CSV.write("./outputs/prob_path_sim_rslts.csv", df_prob_sim)

# ╔═╡ 6d86c2cc-e62d-46b2-a844-343dbe083eab
bar_plot_ref = @df df_prob_sim bar(
	:city_width,
	:prob_path,
	title="Probability of Navigable Path From Top Left to Bottom Right",
	xlabel="City Width (n)",
	legend=false
)

# ╔═╡ eb45d228-befd-46bc-939d-758fb289fd4a
png(bar_plot_ref, "./outputs/prob_path.png")

# ╔═╡ 377592fd-d9f5-4514-a488-2d9f35c4e822
g = random_street_graph(3)

# ╔═╡ caebbad7-460b-44d3-b6a6-0ff9d4684fe8
begin
	nn = Int(sqrt(nv(g)))

	street_plot = gplot(g,
		[j for i = 1:nn for j=1:nn],
		[i for i = 1:nn for j=1:nn],
		nodelabel=["$((i-1)*nn+j)" for i = 1:nn for j=1:nn],
	)
end

# ╔═╡ 019727cd-8aaa-4fc5-bd7e-f90b7e3269d5
draw(PNG("./outputs/street_layout_example_1.png", 16cm, 16cm), street_plot)

# ╔═╡ 3fda5c44-70f3-4202-861d-9c5c4a7415e3
city_has_path(g)

# ╔═╡ 943d0cb8-8c4d-4c4f-9119-7ad84014fc0d
g2 = random_street_graph(3)

# ╔═╡ ccec4661-f080-402d-90c1-0d0ac7238e23
begin
	nn2 = Int(sqrt(nv(g2)))

	street_plot2 = gplot(g2,
		[j for i = 1:nn2 for j=1:nn2],
		[i for i = 1:nn2 for j=1:nn2],
		nodelabel=["$((i-1)*nn2+j)" for i = 1:nn2 for j=1:nn2],
	)
end

# ╔═╡ 9a95a271-f4be-48e9-9b04-57ae7976dda9
draw(PNG("./outputs/street_layout_example_2.png", 16cm, 16cm), street_plot2)

# ╔═╡ dede6d7e-362a-49d7-9223-12b16b0bcdf9
city_has_path(g2)

# ╔═╡ 303cac4d-d222-414a-9b58-db2da01c3e9f
function count_connected(n = 3, start=(1,1), finish=(n,n))

	adjacencies = Set([Set([(i, j), neighb]) for i = 1:n for j = 1:n
					                         for neighb = neighbors(i, j, n)])
	n_adj = length(adjacencies)
	n_layouts = 2^n_adj
	n_connected = 0

	for i_layout = 0:(n_layouts-1)
		# use the binary representation of i_layout to determine connections
		#  1 represents forward connection
		#  0 represents reverse connection
		layout = [parse(Int, ss) for ss ∈ bitstring(i_layout)[end-n_adj+1:end]]

		# create a graph
		g = SimpleDiGraph(n^2)

		# connect the graph
		for (i_adj, adjacent_pair) in enumerate(adjacencies)

			# get the adjacent pair as indices
			(i1, j1), (i2, j2) = adjacent_pair
			k1 = tuple_to_idx(i1, j1, n)
			k2 = tuple_to_idx(i2, j2, n)

			# get the direction for the connection from k1 to k2 from layout
			direction = layout[i_adj]

			# connect node k1 to k2 in the indicated direction
			if direction == 1
				add_edge!(g, k1, k2)
			else
				add_edge!(g, k2, k1)
			end
		end

		# is there a path from start to finish?
		if has_path(g, 1, nv(g))
			n_connected += 1
		end
	end

	return n_connected / n_layouts
end

# ╔═╡ 27f9b76a-b998-4fc4-b89d-baee19104f25
begin
	n_max_analytic = 4
	analytic_rslts = zeros(n_max_analytic - 1)

	for n = 2:n_max_analytic
		analytic_rslts[n-1] = count_connected(n)
	end

	df_analytic = DataFrame(city_width = 2:n_max_analytic,
							prob_path = analytic_rslts)
end

# ╔═╡ 06bb3a39-49aa-42ad-988b-c33671fe1bdb
CSV.write("./outputs/prob_path_analytic_rslts.csv", df_analytic)

# ╔═╡ c2c4c421-8d79-498d-b246-4c1e94f5381b
function idx_to_tuple(k, n)
	i = div((k - 1), n) + 1

	j = mod(k, n)
	if j == 0
		j = n
	end

	return (i, j)
end

# ╔═╡ Cell order:
# ╟─7c9d6b8c-68f0-4481-b042-b51e217a8abd
# ╟─891b9894-11b0-405f-aa85-232cc635db9a
# ╠═f05db774-9741-475b-a6fc-5df56154faa4
# ╠═03792d01-268a-451b-8d48-13afdf6f379f
# ╠═fb2302ca-8b2c-4e23-8dc6-36cc353b16cc
# ╟─713a8126-0a4d-4373-ad57-6872e6aed059
# ╟─377592fd-d9f5-4514-a488-2d9f35c4e822
# ╟─caebbad7-460b-44d3-b6a6-0ff9d4684fe8
# ╟─019727cd-8aaa-4fc5-bd7e-f90b7e3269d5
# ╠═3fda5c44-70f3-4202-861d-9c5c4a7415e3
# ╟─943d0cb8-8c4d-4c4f-9119-7ad84014fc0d
# ╟─ccec4661-f080-402d-90c1-0d0ac7238e23
# ╠═dede6d7e-362a-49d7-9223-12b16b0bcdf9
# ╟─9a95a271-f4be-48e9-9b04-57ae7976dda9
# ╟─caf46b55-87cf-4351-9b6a-f729cce80251
# ╠═285ee664-4c13-4a0b-8c0e-c9b9e5120078
# ╟─cec180f1-68f3-42bc-a5d1-4a7eeca01ab4
# ╟─6d86c2cc-e62d-46b2-a844-343dbe083eab
# ╟─eb45d228-befd-46bc-939d-758fb289fd4a
# ╟─5327b953-268a-4b6d-9bf1-ce513f666b56
# ╠═303cac4d-d222-414a-9b58-db2da01c3e9f
# ╠═27f9b76a-b998-4fc4-b89d-baee19104f25
# ╟─06bb3a39-49aa-42ad-988b-c33671fe1bdb
# ╟─fd5d8a0c-9a36-11eb-37af-b323787108b5
# ╠═38abf0ba-c6ea-4f45-acf1-b92eebe462a8
# ╠═b5b7e758-e4de-4666-976b-583a732e8904
# ╠═c2c4c421-8d79-498d-b246-4c1e94f5381b
# ╠═14d9371c-9a19-11eb-2c1a-8b680938c549


text:
@info "Setting up for job..."
#JOB FILEPATHS
sib_wms_path = "/bench/PhD/NGS_binaries/BMI/sib_nuc_position_sequences.fa_wms.tr"
rys_wms_path = "/bench/PhD/NGS_binaries/BMI/rys_nuc_position_sequences.fa_wms.tr"

sib_df_binary = "/bench/PhD/NGS_binaries/BMI/sib_diff_positions"
rys_df_binary = "/bench/PhD/NGS_binaries/BMI/rys_diff_positions"
combined_df_binary = "/bench/PhD/NGS_binaries/BMI/combined_diff_positions"

sib_code_binary = "/bench/PhD/NGS_binaries/BMI/sib_diff_codes"
rys_code_binary = "/bench/PhD/NGS_binaries/BMI/rys_diff_codes"
combined_code_binary = "/bench/PhD/NGS_binaries/BMI/combined_diff_codes"

sib_diff_bg = "/bench/PhD/NGS_binaries/BMI/sib_diff_bg"
rys_diff_bg = "/bench/PhD/NGS_binaries/BMI/rys_diff_bg"
combined_diff_bg = "/bench/PhD/NGS_binaries/BMI/combined_diff_bg"

sib_e_ui_pth = "/bench/PhD/NGS_binaries/BMI/sib_e_ui"
sib_e_inf_pth = "/bench/PhD/NGS_binaries/BMI/sib_e_inf"
rys_e_ui_pth = "/bench/PhD/NGS_binaries/BMI/rys_e_ui"
rys_e_inf_pth = "/bench/PhD/NGS_binaries/BMI/rys_e_inf"
combined_e_ui_pth = "/bench/PhD/NGS_binaries/BMI/combined_e_ui"
combined_e_inf_pth = "/bench/PhD/NGS_binaries/BMI/combined_e_inf"

#JOB CONSTANTS
const ensemble_size = 1000
const no_sources = 8
const source_min_bases = 3
const source_max_bases = 10
@assert source_min_bases < source_max_bases
const source_length_range= source_min_bases:source_max_bases
const mixing_prior = .07
@assert mixing_prior >= 0 && mixing_prior <= 1
const prior_wt=1.2

@info "Loading master libraries..."
using Distributed, Distributions, Serialization

@info "Adding librarians and workers..."
no_local_procs=2
no_remote_procs=6
remote_machine = "10.0.0.119"
remote_pool=addprocs([(remote_machine, no_remote_procs)], tunnel=true, topology=:master_worker)
local_pool=addprocs(no_local_procs, topology=:master_worker)
worker_pool=vcat(remote_pool,local_pool)

@info "Loading libraries everywhere..."
@everywhere using BioMotifInference, Random
Random.seed!(myid()*10000)

@info "Loading informative source priors..."
sib_wms = BioMotifInference.read_fa_wms_tr(sib_wms_path)
rys_wms = BioMotifInference.read_fa_wms_tr(rys_wms_path)
sib_mix_prior = BioMotifInference.cluster_mix_prior!(deserialize(sib_df_binary), sib_wms)
rys_mix_prior = BioMotifInference.cluster_mix_prior!(deserialize(rys_df_binary), rys_wms)

@info "Filtering informative priors..."
sib_prior_wms = BioMotifInference.filter_priors(Int(floor(no_sources/2)), source_max_bases, sib_wms, sib_mix_prior)
rys_prior_wms = BioMotifInference.filter_priors(Int(floor(no_sources/2)), source_max_bases, rys_wms, rys_mix_prior)
combined_prior_wms = BioMotifInference.combine_filter_priors(Int(floor(no_sources/2)), source_max_bases, (sib_wms, rys_wms), (sib_mix_prior, rys_mix_prior))

@info "Assembling informative source priors..."
sib_inf_sp = BioMotifInference.assemble_source_priors(no_sources, sib_prior_wms, prior_wt)
rys_inf_sp = BioMotifInference.assemble_source_priors(no_sources, rys_prior_wms, prior_wt)
combined_inf_sp = BioMotifInference.assemble_source_priors(no_sources, combined_prior_wms, prior_wt)

@info "Assembling uninformative source priors..."
sib_ui_sp = BioMotifInference.assemble_source_priors(no_sources, Vector{Matrix{Float64}}())
rys_ui_sp = BioMotifInference.assemble_source_priors(no_sources, Vector{Matrix{Float64}}())
combined_ui_sp = BioMotifInference.assemble_source_priors(no_sources, Vector{Matrix{Float64}}())

@info "Loading BGHMM likelihood matrix binaries..."
sib_matrix=deserialize(sib_diff_bg)
rys_matrix=deserialize(rys_diff_bg)
combined_matrix=deserialize(combined_diff_bg)

@info "Loading coded observation sets..."
sib_obs = deserialize(sib_code_binary)
rys_obs = deserialize(rys_code_binary)
combined_obs = deserialize(combined_code_binary)

@info "Assembling sib IPM ensemble on uninformative priors..."
isfile(string(sib_e_ui_pth,'/',"ens")) ? (sib_e_ui = deserialize(string(sib_e_ui_pth,'/',"ens"))) :
    (sib_e_ui = BioMotifInference.IPM_Ensemble(worker_pool, sib_e_ui_pth, ensemble_size, sib_ui_sp, (falses(0,0), mixing_prior), sib_matrix, sib_obs, source_length_range); serialize(string(sib_e_ui_pth,'/',"ens"),sib_e_ui))

@info "Assembling sib IPM ensemble on informative priors..."
isfile(string(sib_e_inf_pth,'/',"ens")) ? (sib_e_inf = deserialize(string(sib_e_inf_pth,'/',"ens"))) :
        (sib_e_inf = BioMotifInference.IPM_Ensemble(worker_pool, sib_e_inf_pth, ensemble_size, sib_inf_sp, (falses(0,0), mixing_prior), sib_matrix, sib_obs, source_length_range); serialize(string(sib_e_inf_pth,'/',"ens"),sib_e_inf))

@info "Assembling rys IPM ensemble on uninformative priors..."
isfile(string(rys_e_ui_pth,'/',"ens")) ? (rys_e_ui = deserialize(string(rys_e_ui_pth,'/',"ens"))) :
    (rys_e_ui = BioMotifInference.IPM_Ensemble(worker_pool, rys_e_ui_pth, ensemble_size, rys_ui_sp, (falses(0,0), mixing_prior), rys_matrix, rys_obs, source_length_range); serialize(string(rys_e_ui_pth,'/',"ens"),rys_e_ui))

@info "Assembling rys IPM ensemble on informative priors..."
isfile(string(rys_e_inf_pth,'/',"ens")) ? (rys_e_inf = deserialize(string(rys_e_inf_pth,'/',"ens"))) :
        (rys_e_inf = BioMotifInference.IPM_Ensemble(worker_pool, rys_e_inf_pth, ensemble_size, rys_inf_sp, (falses(0,0), mixing_prior), rys_matrix, rys_obs, source_length_range); serialize(string(rys_e_inf_pth,'/',"ens"),rys_e_inf))

@info "Assembling combined IPM ensemble on uninformative priors..."
isfile(string(combined_e_ui_pth,'/',"ens")) ? (combined_e_ui = deserialize(string(combined_e_ui_pth,'/',"ens"))) :
    (combined_e_ui = BioMotifInference.IPM_Ensemble(worker_pool, combined_e_ui_pth, ensemble_size, combined_ui_sp, (falses(0,0), mixing_prior), combined_matrix, combined_obs, source_length_range); serialize(string(combined_e_ui_pth,'/',"ens"),combined_e_ui))

@info "Assembling combined IPM ensemble on informative priors..."
isfile(string(combined_e_inf_pth,'/',"ens")) ? (combined_e_inf = deserialize(string(combined_e_inf_pth,'/',"ens"))) :
        (combined_e_inf = BioMotifInference.IPM_Ensemble(worker_pool, combined_e_inf_pth, ensemble_size, combined_inf_sp, (falses(0,0), mixing_prior), combined_matrix, combined_obs, source_length_range); serialize(string(combined_e_inf_pth,'/',"ens"),combined_e_inf))

rmprocs(worker_pool)

@info "Fitting distributions..."

sib_ui_dist=fit_mle(Normal,[model.log_Li for model in sib_e_ui.models])
sib_inf_dist=fit_mle(Normal,[model.log_Li for model in sib_e_inf.models])
@info "Is informative a better prior for sibs?"
println(sib_inf_dist.μ > sib_ui_dist.μ && quantile(sib_inf_dist,.5) > sib_ui_dist.μ)

rys_ui_dist=fit_mle(Normal,[model.log_Li for model in rys_e_ui.models])
rys_inf_dist=fit_mle(Normal,[model.log_Li for model in rys_e_inf.models])
@info "Is informative a better prior for rys?"
println(rys_inf_dist.μ > rys_ui_dist.μ && quantile(rys_inf_dist,.5) > rys_ui_dist.μ)

combined_ui_dist=fit_mle(Normal,[model.log_Li for model in combined_e_ui.models])
combined_inf_dist=fit_mle(Normal,[model.log_Li for model in combined_e_inf.models])
@info "Is informative a better prior for combined?"
println(combined_inf_dist.μ > combined_ui_dist.μ && quantile(combined_inf_dist,.5) > combined_ui_dist.μ)

@info "Job done!"