"""
Creates Ohms constraints for shiftable PSTs / OLTCs

```
p[f_idx] == g*vm_tap[f_idx]^2 + (-g)*(vm_tap[f_idx]*vm_tap[t_idx]*cos((va[f_bus] - va_shift[f_idx]) - (va[t_bus] - va_shift[t_idx]))) + (-b)*(vm_tap[f_idx]*vm_tap[t_idx]*sin((va[f_bus] - va_shift[f_idx]) - (va[t_bus] - va_shift[t_idx]))))
q[f_idx] == -(b+c/2)*vm_tap[f_idx]^2 - (-b)*(vm_tap[f_idx]*vm_tap[t_idx]*cos((va[f_bus] - va_shift[f_idx]) - (va[t_bus] - va_shift[t_idx]))) + (-g)*(vm_tap[f_idx]*vm_tap[t_idx]*sin((va[f_bus] - va_shift[f_idx]) - (va[t_bus] - va_shift[t_idx]))))
vm_tap[f_idx] * tap_min <= vm[f_bus] <= vm_tap[f_idx] * tap_max
```
"""
function constraint_variable_transformer_y_from(pm::GenericPowerModel{T}, n::Int, cnd::Int, f_bus, t_bus, f_idx, t_idx, g, b, c, g_shunt, tap_min, tap_max)  where T <: PowerModels.AbstractACPForm
    p_fr = PowerModels.var(pm, n, cnd, :p, f_idx)
    q_fr = PowerModels.var(pm, n, cnd, :q, f_idx)
    vm_fr = PowerModels.var(pm, n, cnd, :vm, f_bus)
    vm_to = PowerModels.var(pm, n, cnd, :vm, t_bus)
    va_fr = PowerModels.var(pm, n, cnd, :va, f_bus)
    va_to = PowerModels.var(pm, n, cnd, :va, t_bus)
    va_shift_fr = PowerModels.var(pm, n, cnd, :va_shift, f_idx)
    va_shift_to = PowerModels.var(pm, n, cnd, :va_shift, t_idx)
    vm_tap_fr = PowerModels.var(pm, n, cnd, :vm_tap, f_idx)
    vm_tap_to = PowerModels.var(pm, n, cnd, :vm_tap, t_idx)

    @NLconstraint(pm.model, p_fr == (g + g_shunt / 2)*vm_tap_fr^2 + (-g)*(vm_tap_fr*vm_tap_to*cos((va_fr - va_shift_fr) - (va_to - va_shift_to))) + (-b)*(vm_tap_fr*vm_tap_to*sin((va_fr - va_shift_fr) - (va_to - va_shift_to))))
    @NLconstraint(pm.model, q_fr == -(b+c/2)*vm_tap_fr^2 - (-b)*(vm_tap_fr*vm_tap_to*cos((va_fr - va_shift_fr) - (va_to - va_shift_to))) + (-g)*(vm_tap_fr*vm_tap_to*sin((va_fr - va_shift_fr) - (va_to - va_shift_to))))
    @constraint(pm.model, vm_tap_fr * tap_min <= vm_fr)
    @constraint(pm.model, vm_fr <= vm_tap_fr * tap_max)
end
"""
Creates Ohms constraints for shiftable PSTs / OLTCs

```
p[t_idx] == g*vm_tap[t_idx]^2 + (-g)*(vm_tap[t_idx]*vm_tap[f_idx]*cos((va[t_bus] - va_shift[t_idx]) - (va[f_bus] - va_shift[f_idx]))) + (-b)*(vm_tap[t_idx]*vm_tap[f_idx]*sin((va[t_bus] - va_shift[t_idx]) - (va[f_bus] - va_shift[f_idx]))))
q[t_idx] == -(b+c/2)*vm_tap[t_idx]^2 - (-b)*(vm_tap[t_idx]*vm_tap[f_idx]*cos((va[t_bus] - va_shift[t_idx]) - (va[f_bus] - va_shift[f_idx]))) + (-g)*(vm_tap[t_idx]*vm_tap[f_idx]*sin((va[t_bus] - va_shift[t_idx]) - (va[f_bus] - va_shift[f_idx]))))
vm_tap[t_idx] * tap_min <= vm[t_bus] <= vm_tap[t_idx] * tap_max
```
"""
function constraint_variable_transformer_y_to(pm::GenericPowerModel{T}, n::Int, cnd::Int, f_bus, t_bus, f_idx, t_idx, g, b, c, g_shunt, tap_min, tap_max) where T <: PowerModels.AbstractACPForm
    p_to = PowerModels.var(pm, n, cnd :p, t_idx)
    q_to = PowerModels.var(pm, n, cnd, :q, t_idx)
    vm_fr = PowerModels.var(pm, n, cnd, :vm, f_bus)
    vm_to = PowerModels.var(pm, n, cnd, :vm, t_bus)
    va_fr = PowerModels.var(pm, n, cnd, :va, f_bus)
    va_to = PowerModels.var(pm, n, cnd, :va, t_bus)
    va_shift_fr = PowerModels.var(pm, n, cnd, :va_shift, f_idx)
    va_shift_to = PowerModels.var(pm, n, cnd, :va_shift, t_idx)
    vm_tap_fr = PowerModels.var(pm, n, cnd, :vm_tap, f_idx)
    vm_tap_to = PowerModels.var(pm, n, cnd, :vm_tap, t_idx)

    @NLconstraint(pm.model, p_to == (g + g_shunt / 2)*vm_tap_to^2 + (-g)*(vm_tap_to*vm_tap_fr*cos((va_to - va_shift_to) - (va_fr - va_shift_fr))) + (-b)*(vm_tap_to*vm_tap_fr*sin((va_to - va_shift_to) - (va_fr - va_shift_fr))))
    @NLconstraint(pm.model, q_to == -(b+c/2)*vm_tap_to^2 - (-b)*(vm_tap_to*vm_tap_fr*cos((va_to - va_shift_to) - (va_fr - va_shift_fr))) + (-g)*(vm_tap_to*vm_tap_fr*sin((va_to - va_shift_to) - (va_fr - va_shift_fr))))
    @constraint(pm.model, vm_tap_to * tap_min <= vm_to)
    @constraint(pm.model, vm_to <= vm_tap_to * tap_max)
end
"""
Links voltage magnitudes of not tappable transformers with node voltage magnitudes

```
vm_tap[f_idx] * tap == vm[f_bus]
vm_tap[t_idx] * tap == vm[t_bus]
```
"""
function constraint_link_voltage_magnitudes(pm::GenericPowerModel{T}, n::Int, cnd::Int, f_bus, t_bus, f_idx, t_idx, tap_fr, tap_to) where T <: PowerModels.AbstractACPForm
    vm_fr = PowerModels.var(pm, n, cnd, :vm, f_bus)
    vm_to = PowerModels.var(pm, n, cnd, :vm, t_bus)
    vm_tap_fr = PowerModels.var(pm, n, cnd, :vm_tap, f_idx)
    vm_tap_to = PowerModels.var(pm, n, cnd, :vm_tap, t_idx)
    @constraint(pm.model, vm_tap_fr * tap_fr == vm_fr)
    @constraint(pm.model, vm_tap_to * tap_to == vm_to)
end
"""
Aggregated unit representation in the nodes
```
sum(p[a] for a in bus_arcs) + sum(p_dc[a_dc] for a_dc in bus_arcs_dc) == pnode - gs*v^2
sum(q[a] for a in bus_arcs) + sum(q_dc[a_dc] for a_dc in bus_arcs_dc) == qnode + bs*v^2
```
"""
function constraint_kcl_shunt_aggregated(pm::GenericPowerModel{T}, n::Int, cnd::Int, i, bus_arcs, bus_arcs_dc, bus_gs, bus_bs) where T <: PowerModels.AbstractACPForm
    vm = PowerModels.var(pm, n, cnd, :vm, i)
    p = PowerModels.var(pm, n, cnd, :p)
    q = PowerModels.var(pm, n, cnd, :q)
    pnode = PowerModels.var(pm, n, cnd, :pnode, i)
    qnode = PowerModels.var(pm, n, cnd, :qnode, i)
    p_dc = PowerModels.var(pm, n, cnd, :p_dc)
    q_dc = PowerModels.var(pm, n, cnd, :q_dc)

    @constraint(pm.model, sum(p[a] for a in bus_arcs) + sum(p_dc[a_dc] for a_dc in bus_arcs_dc) == pnode - sum(gs for gs in values(bus_gs))*vm^2)
    @constraint(pm.model, sum(q[a] for a in bus_arcs) + sum(q_dc[a_dc] for a_dc in bus_arcs_dc) == qnode + sum(bs for bs in values(bus_bs))*vm^2)
end
function constraint_redispatch_active_power_gen(pm::GenericPowerModel{T}, n::Int, cnd::Int, i::Int, pref::AbstractFloat) where T <: PowerModels.AbstractACPForm
    pg_delta = PowerModels.var(pm, n, cnd, :pg_delta, i)
    pg = PowerModels.var(pm, n, cnd, :pg, i)

    if pm.setting["relax_absolute_value"] == false
        @NLconstraint(pm.model, pg_delta == abs(pg - pref))
    else
        @constraint(pm.model, pg_delta >= pg - pref)
        @constraint(pm.model, pg_delta >= -(pg - pref))
    end
end
function constraint_redispatch_reactive_power_gen(pm::GenericPowerModel{T}, n::Int,  cnd::Int, i::Int, qref::AbstractFloat) where T <: PowerModels.AbstractACPForm
    qg_delta = PowerModels.var(pm, n, cnd, :qg_delta, i)
    qg = PowerModels.var(pm, n, cnd, :qg, i)

    if pm.setting["relax_absolute_value"] == false
        @NLconstraint(pm.model, qg_delta == abs(qg - qref))
    else
        @constraint(pm.model, qg_delta >= qg - qref)
        @constraint(pm.model, qg_delta >= -(qg - qref))
    end
end
function constraint_second_stage_redispatch_active_power_gen(pm::GenericPowerModel{T}, n::Int, cnd::Int, i::Int, first_stage_network_id) where T <: PowerModels.AbstractACPForm
    pg_delta = PowerModels.var(pm, n, cnd, :pg_delta, i)
    pg = PowerModels.var(pm, n, cnd, :pg, i)
    pg_first_stage = PowerModels.var(pm, first_stage_network_id, cnd, :pg, i)

    if pm.setting["relax_absolute_value"] == false
        @NLconstraint(pm.model, pg_delta == abs(pg - pg_first_stage))
    else
        @constraint(pm.model, pg_delta >= pg - pg_first_stage)
        @constraint(pm.model, pg_delta >= -(pg - pg_first_stage))
    end
end
function constraint_second_stage_redispatch_reactive_power_gen(pm::GenericPowerModel{T}, n::Int, cnd::Int, i::Int, first_stage_network_id) where T <: PowerModels.AbstractACPForm
    qg_delta = PowerModels.var(pm, n, cnd, :qg_delta, i)
    qg = PowerModels.var(pm, n, cnd, :qg, i)
    qg_first_stage = PowerModels.var(pm, first_stage_network_id, cnd, :qg, i)

    if pm.setting["relax_absolute_value"] == false
        @NLconstraint(pm.model, qg_delta == abs(qg - qg_first_stage))
    else
        @constraint(pm.model, qg_delta >= qg - qg_first_stage)
        @constraint(pm.model, qg_delta >= -(qg - qg_first_stage))
    end
end
function constraint_redispatch_active_power_load(pm::GenericPowerModel{T}, n::Int, cnd::Int, i::Int, pref::AbstractFloat) where T <: PowerModels.AbstractACPForm
    pl_delta = PowerModels.var(pm, n, cnd, :pl_delta, i)
    pl = PowerModels.var(pm, n, cnd, :pl, i)

    if pm.setting["relax_absolute_value"] == false
        @NLconstraint(pm.model, pl_delta == abs(pl - pref))
    else
        @constraint(pm.model, pl_delta >= pl - pref)
        @constraint(pm.model, pl_delta >= -(pl - pref))
    end
end
function constraint_redispatch_reactive_power_load(pm::GenericPowerModel{T}, n::Int, cnd::Int, i::Int, qref::AbstractFloat) where T <: PowerModels.AbstractACPForm
    ql_delta = PowerModels.var(pm, n, cnd, :ql_delta, i)
    ql = PowerModels.var(pm, n, cnd, :ql, i)

    if pm.setting["relax_absolute_value"] == false
        @NLconstraint(pm.model, ql_delta == abs(ql - qref))
    else
        @constraint(pm.model, ql_delta >= ql - qref)
        @constraint(pm.model, ql_delta >= -(ql - qref))
    end
end
function constraint_second_stage_redispatch_active_power_load(pm::GenericPowerModel{T}, n::Int, cnd::Int, i::Int, first_stage_network_id)  where T <: PowerModels.AbstractACPForm
    pl_delta = PowerModels.var(pm, n, cnd, :pl_delta, i)
    pl = PowerModels.var(pm, n, cnd, :pl, i)
    pl_first_stage = PowerModels.var(pm, first_stage_network_id, cnd, :pl, i)

    if pm.setting["relax_absolute_value"] == false
        @NLconstraint(pm.model, pl_delta == abs(pl - pl_first_stage))
    else
        @constraint(pm.model, pl_delta >= pl - pl_first_stage)
        @constraint(pm.model, pl_delta >= -(pl - pl_first_stage))
    end
end
function constraint_second_stage_redispatch_reactive_power_load(pm::GenericPowerModel{T}, n::Int, cnd::Int, i::Int, first_stage_network_id)  where T <: PowerModels.AbstractACPForm
    ql_delta = PowerModels.var(pm, n, cnd, :ql_delta, i)
    ql = PowerModels.var(pm, n, cnd, :ql, i)
    ql_first_stage = PowerModels.var(pm, first_stage_network_id, cnd, :ql, i)

    if pm.setting["relax_absolute_value"] == false
        @NLconstraint(pm.model, ql_delta == abs(ql - ql_first_stage))
    else
        @constraint(pm.model, ql_delta >= ql - ql_first_stage)
        @constraint(pm.model, ql_delta >= -(ql - ql_first_stage))
    end
end


function UCST_mix(model::EoSModel,T;v0=nothing)
    if v0 == nothing
        v0 = x0_UCST_mix(model,T)
    end
    f! = (F,x) -> Obj_UCST_mix(model, F, x[2], exp10(x[1]), T)
    r  = Solvers.nlsolve(f!,v0[1:end],LineSearch(Newton()))
    sol = Solvers.x_sol(r)
    z_c = FractionVector(sol[2])
    V_c = exp10(sol[1])
    p_c = pressure(model, V_c, T, z_c)
    return (p_c, V_c, z_c)
end
function x0_UCST_mix(model::EoSModel,T)
    V  = x0_volume_liquid(model,T,[0.5,0.5])
    return [log10(V),0.5]
end

function Obj_UCST_mix(model::EoSModel,F,z,V,T)
    z    = FractionVector(z)
    f(x) = eos(model,V,T,x)
    H(x) = ForwardDiff.hessian(f,x)/8.134/T
    L(x) = det(H(x))
    dL(x) = ForwardDiff.gradient(L,x)
    M(x) = [H(x)[1:end-1,:];transpose(dL(x))]
    F[1] = L(z)
    F[2] = det(M(z))
    return F
end


function CHISQUAREHEUR(trans::SparseMatrixCSC{Float64, Int64},
                       capacity::Array{Int64,1},
                       Q::Array{Int64,2},
                       sig::Float64)
    if CHECKCAPACITY(Q,capacity) == 1
        # Number of transactions
        J = size(trans,1)

        # Number of products
        I = size(trans,2)

        # Chi Square Test
        ## We start the chi-square heuristic by performing chi-square tests
        ## for all unique SKU-combinations. For each undirected SKU-pair we
        ## evaluate a chi-square test and either accept or reject the null
        ## hypothesis of independence. Our goal is a matrix dep with the
        ## “net-benefit” of all unique positive dependent SKU-combinations
        ## and a matrix nor with all "independent contributions" of the unique
        ## SKU-combinations. More details in our article.
        sum_cond_sku = Int.(vec(sum(trans,dims=1)))
        dep, nor = HYOPTHESISCHI(Q::Array{Int64,2},
                                 I::Int64,
                                 J::Int64,
                                 sig::Float64,
                                 sum_cond_sku::Array{Int64,1})

        ## Create the sku-warehouse allocation matrix
        X = Array{Bool,2}(undef,I,size(capacity,1)) .= false

        ## Determine the sum of the coappearances for each SKU on the base of
        ## the matrices dep and nor
        sum_dep = vec(sum(dep,dims = 1))
        sum_nor = vec(sum(nor,dims = 1))

        ## Calculate the weight of each warehouse. It shows us the density of
        ## the independent coappearances in each warehouse if we were to allocate
        ## all SKUs simply according to the highest independent coappearances.
        weight = WHWEIGHT(capacity::Array{Int64,1},sum_nor::Array{Float64,1})

        ## Copy the capacity to have an array that keeps the left-over capacity
        cap_left = copy(capacity)

        ## In the heuristic we compare the split-delivery minimising potential
        ## of SKU allocations to maximise dependent coappearances (SKUs with
        ## positive dependencies are stored in the same warehouse) with the
        ## maximisation of independent coappearances (SKUs with high independent
        ## coappearances are stored in warehouses with high capacities).
        ## Start the allocation:
        ## If all SKUs are allocated once, continue after the while loop, else
        while sum(X) < I
        ## select the SKU i with the highest coappearance not being allocated
        ## to the warehouses yet. In addition, select the first sorted
        ## warehouse k which has unused capacity left.
        ## Then, check for each warehouse with free space whether SKU i has
        ## significant dependencies to other already allocated SKUs.
        ## If dependencies exist check whether they are expected to dominate.
        ## If so, allocate SKU i to to the warehouse with the dependencies
        ## to complement the positive dependent SKUs in the warehouse.
        ## Otherwise allocate SKU i to warehouse k as the independent
        ## coappearances are expected to dominate.
        ## If no dependencies to yet allocated SKUs exist, select the warehouse
        ## with the highest unused capacity. If the dependet coappearances are
        ## expected to to dominate later, we allocate SKU i to the warehouse
        ## with the highest unused capacity. This maximises the likelihood to
        ## allocate all significant SKUs from a SKU-cluster into one warehouse.
        ## Otherwise we allocate it to warehouse k to maximise the independent
        ## coappearances in the allocation.
        i,k  = SELECTIK(sum_dep::Array{Float64,1},
                            sum_nor::Array{Float64,1},
                            weight,
                            cap_left::Array{Int64,1},
                            X::Array{Bool,2},
                            dep::Array{Float64,2})
        ALLOCATEONE!(X::Array{Bool,2},
                        sum_dep::Array{Float64,1},
                        sum_nor::Array{Float64,1},
                        cap_left::Array{Int64,1},
                        i::Int64,
                        k::Int64)

        ## Check for all unallocated SKUs whether they have positive
        ## dependencies to the SKUs in the warehouse k the last SKU was allocated
        ## to. If so, check whether the dependencies are expected to dominate the
        ## independent coapperances. If yes, allocate the corresponding SKUs to
        ## the warehouse k.
        ADDDEPENDENT!(X::Array{Bool,2},
                        cap_left::Array{Int64,1},
                        k::Int64,
                        dep::Array{Float64,2},
                        nor::Array{Float64,2},
                        sum_dep::Array{Float64,1},
                        sum_nor::Array{Float64,1})
        FILLLAST!(X::Array{Bool,2},cap_left::Array{Int64,1})
        end
        if sum(cap_left) > 0
            ## First, remove every so far assigned dependent SKU-pair from the coappearance
            ## matrix dep to prevent the allocation bias described in our article.
            Qs = REMOVEALLOC(X::Array{Bool,2},
                            cap_left::Array{Int64,1},
                            nor::Array{Float64,2},
                            dep::Array{Float64,2})

            ## Afterwards allocate the SKUs with the highest potential allocation value to
            ## each warehouse with leftover storage space until it is full. If no SKU
            ## with coappearances is found and there is still storage space left, terminate the
            ## algorithm as further allocations pose no benefit.
            FILLUP!(X::Array{Bool,2},
                        Qs::Array{Float64,2},
                        cap_left::Array{Int64,1})
        end
    end
    ## return the resulting allocation matrix
    X = convert(Matrix{Int64},X)
    return X::Array{Int64,2}
end


# ---
# title: 827. Making A Large Island
# id: problem827
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Depth-first Search
# link: <https://leetcode.com/problems/making-a-large-island/description/>
# hidden: true
# ---
#
# In a 2D grid of `0`s and `1`s, we change at most one `0` to a `1`.
#
# After, what is the size of the largest island? (An island is a 4-directionally
# connected group of `1`s).
#
# **Example 1:**
#
#
#
#     Input: [[1, 0], [0, 1]]
#     Output: 3
#     Explanation: Change one 0 to 1 and connect two 1s, then we get an island with area = 3.
#
#
# **Example 2:**
#
#
#
#     Input: [[1, 1], [1, 0]]
#     Output: 4
#     Explanation: Change the 0 to 1 and make the island bigger, only one island with area = 4.
#
# **Example 3:**
#
#
#
#     Input: [[1, 1], [1, 1]]
#     Output: 4
#     Explanation: Can't change any 0 to 1, only one island with area = 4.
#
#
#
# Notes:
#
#   * `1 <= grid.length = grid[0].length <= 50`.
#   * `0 <= grid[i][j] <= 1`.
#
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ aeae4550-720d-11eb-3ef5-1f23537ec599
using GMT

# ╔═╡ 0745ac40-720d-11eb-1acd-2d4920b5c475
md"# Contourf examples

GMT does not actually have a *contourf* module like Matlab for example, but we can obtain the same result using *grdview*, *grdcontour* and *pscontour*. However, to make things the Julia wrapper wrapped up a module called *contourf* that makes it really easy to use. To show how it works let's start by creating an example grid and a CPT.
"

# ╔═╡ b2cdb8a0-720d-11eb-200b-cd2c056c8375
G = GMT.peaks();

# ╔═╡ d747f28e-720d-11eb-0e68-a7242445823b
C = makecpt(T=(-7,9,2));

# ╔═╡ df241fc0-720d-11eb-13d2-c311aa8f5518
md"Now if we pass those two to the *contourf* module we get an annotated plot where the annotations come from the color CPT.
"

# ╔═╡ 0faaea20-720e-11eb-0ff8-1105205be52c
contourf(G, C, show=true)

# ╔═╡ 29a06c20-720e-11eb-312a-ff244a4d7a71
md"If we want to just draw some contours and not annotate them, we pass an array with the contours to be drawn.
"

# ╔═╡ 5ca0c200-720e-11eb-3a76-5db53780ef22
contourf(G, C, contour=[-2, 0, 2, 5], show=true)

# ╔═╡ 6a4f90c0-720e-11eb-29c7-1529f5116d73
md"### What if one has an *x,y,z* file instead of a grid?
That is also simple, let's simulate it with synthetic data.
"

# ╔═╡ a9d59eae-720e-11eb-050a-d1737498f9a0
d = [0 2 5; 1 4 5; 2 0.5 5; 3 3 9; 4 4.5 5; 4.2 1.2 5; 6 3 1; 8 1 5; 9 4.5 5];

# ╔═╡ d12ce0e0-720e-11eb-05c4-3b65ee0d2da2
contourf(d, limits=(-0.5,9.5,0,5), pen=0.25, labels=(line=(:min,:max),), show=1)

# ╔═╡ b22ad6fe-720f-11eb-2de4-37064019d3f6
md"In the above since we did not specify a CPT the program picked the GMT's default one. But if we want use another one it's only a question of creating and passed it in.
"

# ╔═╡ f583bede-720f-11eb-0ff0-617f73c13447
cpt = makecpt(range=(0,10,1), cmap=:batlow);

# ╔═╡ feda3000-720f-11eb-2b2f-d546d65f4abe
contourf(d, contours=cpt, limits=(-0.5,9.5,0,5), pen=0.25, labels=(line=(:min,:max),), show=true)

# ╔═╡ Cell order:
# ╟─0745ac40-720d-11eb-1acd-2d4920b5c475
# ╠═aeae4550-720d-11eb-3ef5-1f23537ec599
# ╠═b2cdb8a0-720d-11eb-200b-cd2c056c8375
# ╠═d747f28e-720d-11eb-0e68-a7242445823b
# ╟─df241fc0-720d-11eb-13d2-c311aa8f5518
# ╠═0faaea20-720e-11eb-0ff8-1105205be52c
# ╠═29a06c20-720e-11eb-312a-ff244a4d7a71
# ╠═5ca0c200-720e-11eb-3a76-5db53780ef22
# ╟─6a4f90c0-720e-11eb-29c7-1529f5116d73
# ╠═a9d59eae-720e-11eb-050a-d1737498f9a0
# ╠═d12ce0e0-720e-11eb-05c4-3b65ee0d2da2
# ╟─b22ad6fe-720f-11eb-2de4-37064019d3f6
# ╠═f583bede-720f-11eb-0ff0-617f73c13447
# ╠═feda3000-720f-11eb-2b2f-d546d65f4abe


# ---
# title: Renaming and transforming variables
# cover: assets/new_columns_on_the_fly.png
# description: Transforming data before generating the plot.
# author: "[Pietro Vertechi](https://github.com/piever)"
# ---

using AlgebraOfGraphics, CairoMakie
set_aog_theme!() #src

df = (x=rand(100), y=rand(100), z=rand(100), c=rand(["a", "b"], 100))
layers = linear() + mapping(color=:z)

plt = data(df) * layers * mapping(:x => (x -> x^2) => L"x^2", :y => L"This variable is called $y$")
fg = draw(plt)

# Use a `Tuple` to pass combine several columns into a unique operation.

plt = data(df) * layers * mapping(:x, (:x, :y, :z) => (+) => L"the new variable $x + y + z$", layout=:c)
fg = draw(plt)

# save cover image #src
mkpath("assets") #src
save("assets/new_columns_on_the_fly.png", fg) #src


using DoubleExponentialFormulas
using LinearAlgebra: norm
using Test

quades32 = QuadES(Float32)
quades64 = QuadES(Float64)
quadesBF = QuadES(BigFloat)

# Test Base.show works without error
let
    io = IOBuffer()
    @test (Base.show(io, MIME("text/plain"), quades32); true)
    @test (Base.show(io, MIME("text/plain"), quades64); true)
    @test (Base.show(io, MIME("text/plain"), quadesBF); true)
end

# Test [0, ∞] interval (with Exp-Sinh quadrature)
let
    f(x::AbstractFloat) = exp(-x)
    expect = 1

    I, E = quades32(f)
    @test I isa Float32
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I)))*norm(I)

    I, E = quades64(f)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I)))*norm(I)

    rtol = 1e-30
    I, E = quadesBF(f, rtol=rtol)
    @test I isa BigFloat
    @test isapprox(I, expect, rtol=10rtol)
    @test E ≤ rtol*norm(I)
end

# Test [0, ∞] interval (with Exp-Sinh quadrature)
let
    f(x::AbstractFloat) = -exp(-x)*log(x)
    expect = BigFloat("5.77215664901532860606512090082e-1")  # γ constant

    I, E = quades32(f)
    @test I isa Float32
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I)))*norm(I)

    I, E = quades64(f)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I)))*norm(I)

    rtol = 1e-30
    I, E = quadesBF(f, rtol=rtol)
    @test I isa BigFloat
    @test isapprox(I, expect, rtol=10rtol)
    @test E ≤ rtol*norm(I)
end

# Test [0, ∞] interval (with Exp-Sinh quadrature)
let
    f(x::AbstractFloat) = 2/(1 + x^2)
    expect = π

    I, E = quades32(f)
    @test I isa Float32
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I)))*norm(I)

    I, E = quades64(f)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I)))*norm(I)

    rtol = 1e-30
    I, E = quadesBF(f, rtol=rtol)
    @test I isa BigFloat
    @test isapprox(I, expect, rtol=10rtol)
    @test E ≤ rtol*norm(I)
end

# Test [0, ∞] interval (with Exp-Sinh quadrature)
# Dirichlet integral
let
    f(x::AbstractFloat) = 2*sin(x)/x
    expect = π

    I, E = quades32(f)
    @test I isa Float32
    @test_broken I ≈ expect
    @test_skip E ≤ sqrt(eps(typeof(I)))*norm(I)

    I, E = quades64(f)
    @test I isa Float64
    @test_broken I ≈ expect
    @test_skip E ≤ sqrt(eps(typeof(I)))*norm(I)

    # FIXME: Something wrong. Infinite loop?
    # rtol = 1e-17
    # I, E = quadesBF(f, rtol=rtol)
    # @test I isa BigFloat
    # @test_broken isapprox(I, expect, rtol=10rtol)
    # @test_skip E ≤ rtol*norm(I)
end

# Test `f(xk)*wk = 0*Inf` case
let
    f(x::AbstractFloat) = x*exp(-x^2)
    expect = 1/BigFloat(2)

    I, E = quades32(f)
    @test I isa Float32
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I)))*norm(I)

    I, E = quades64(f)
    @test I isa Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(typeof(I)))*norm(I)

    rtol = 1e-30
    I, E = quadesBF(f, rtol=rtol)
    @test I isa BigFloat
    @test isapprox(I, expect, rtol=10rtol)
    @test E ≤ rtol*norm(I)
end

# Test non-scalar output with QuadES
let
    f(x::AbstractFloat) = [exp(-x), 2exp(-x)]
    expect = [1, 2]

    I, E = quades32(f)
    @test eltype(I) == Float32
    @test I ≈ expect
    @test E ≤ sqrt(eps(eltype(I)))*norm(I)

    I, E = quades64(f)
    @test eltype(I) == Float64
    @test I ≈ expect
    @test E ≤ sqrt(eps(eltype(I)))*norm(I)

    rtol = 1e-30
    I, E = quadesBF(f, rtol=rtol)
    @test eltype(I) == BigFloat
    @test isapprox(I, expect, rtol=10rtol)
    @test E ≤ rtol*norm(I)
end


export ChimeraLattice

struct ChimeraLattice <: AbstractLattice
    Nx::Int
    Ny::Int
end

vertices(lt::ChimeraLattice) = 1:8*lt.Nx*lt.Ny
bravais_size(cl::ChimeraLattice) = (cl.Nx, cl.Ny)
function bravais_size(cl::ChimeraLattice, i::Int)
    if i==1
        cl.Nx
    elseif i==2
        cl.Ny
    else
        DimensionMismatch("expect dimension 1 or 2, got $i")
    end
end
Base.size(lt::ChimeraLattice) = (8, lt.Nx, lt.Ny)

gap_x(lt::ChimeraLattice) = 0.5
gap_y(lt::ChimeraLattice) = 0.5
unit(lt::ChimeraLattice) = 1/(max(lt.Nx*(2+gap_x(lt)), lt.Ny*(4+gap_y(lt))))

# for beautiful printing.
function Base.getindex(lt::ChimeraLattice, k::Int, i::Int, j::Int)
    ki = (k-1) ÷ 4
    kj = mod1(k, 4)-1
    x = ((i-1)*(gap_x(lt) + 2)+ki+0.5)*unit(lt)
    y = ((j-1)*(gap_y(lt) + 4)+kj+0.5)*unit(lt)
    return (x+0.01*kj, y+0.02*ki)
end

function Base.getindex(lt::ChimeraLattice, i::Int)
    CI = CartesianIndices(size(lt))
    lt[CI[i].I...]
end

function bonds(lt::ChimeraLattice)
    bonds = Tuple{Int,Int}[]
    red_bond!(lt, 1, bonds)
    redblack_bond!(lt, 1, bonds)
    for i=2:lt.Nx
        # BLACK
        black_bond!(lt, i-1, bonds)
        # Contract with RED
        red_bond!(lt, i, bonds)
        redblack_bond!(lt, i, bonds)
    end
    return bonds
end

function red_bond!(lt, i::Int, bonds)
    CI = LinearIndices(lt |> size)
    for j=1:lt.Ny-1
        for k = 1:4
            push!(bonds, (CI[k, i, j], CI[k, i, j+1]))
        end
    end
    return bonds
end

function redblack_bond!(lt, i::Int, bonds)
    CI = LinearIndices(lt |> size)
    for j=1:lt.Ny
        for k2=5:8
            for k1=1:4
                push!(bonds, (CI[k2, i, j], CI[k1, i, j]))
            end
        end
    end
    return bonds
end

function black_bond!(lt, i::Int, bonds)
    CI = LinearIndices(lt |> size)
    for j=1:lt.Ny
        for k = 5:8
            push!(bonds, (CI[k,i,j], CI[k,i+1,j]))
        end
    end
    return bonds
end


using BenchmarkTools
using FinancialMonteCarlo, CUDA
CUDA.allowscalar(false)
S0 = 100.0f0;
K = 100.0f0;
r = 0.02f0;
T = 1.0f0;
d = 0.01f0;
D = 90.0f0;

Nsim = 10000;
Nstep = 30;
sigma = 0.2f0;
mc = MonteCarloConfiguration(Nsim, Nstep);
mc_2 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.CudaMode());
toll = 1e-3;

rfCurve = ZeroRate(r);

FwdData = Forward(T)
EUData = EuropeanOption(T, K)
AMData = AmericanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model = BlackScholesProcess(sigma, Underlying(S0, d));

@show "CUDA_2 fwd"
@btime FwdPrice = pricer(Model, rfCurve, mc_2, FwdData);
@show "CUDA_2 eu"
@btime EuPrice = pricer(Model, rfCurve, mc_2, EUData);


using Plots

s_date = 0
e_date = 31

function kwt_alfa(degalfa_p)
    eff = 0.15
    area = 30

    teta_p = deg2rad(170)
    alfa_p = deg2rad(degalfa_p)

    lat = deg2rad(59.3)

    delta_s(n) = deg2rad(-23.44) * cos((2pi / 365) * n)

    omega(t) = deg2rad((15 * t) - 180)

    alfa_s(n, t) = asin(sin(lat) * sin(delta_s(n)) + cos(lat) * cos(delta_s(n)) * cos(omega(t)))

    teta_s(n, t) = pi + (omega(t) / abs(omega(t))) * acos((sin(lat) * sin(alfa_s(n, t)) - sin(delta_s(n))) / (cos(lat) * cos(alfa_s(n, t))))

    I_0 = 1360
    I(n, t) = 1.1 * I_0 * 0.7^((1 / (abs(sin(alfa_s(n, t))) + sin(alfa_s(n, t))) / 2)^0.678)  # underlig abs medelvärde för att få 0 om negativt värde

    I_p(n, t) = I(n, t) * (cos(teta_p - teta_s(n, t)) * cos(alfa_p - alfa_s(n, t)) + sin(alfa_s(n, t)) * sin(alfa_p) * (1 - cos(teta_p - teta_s(n, t))))

    weather_loss = 1 - (1821 / (365 * 12))

    P(x) = (*(((abs(I_p(floor(x / 24), x)) + I_p(floor(x / 24), x)) / 2), weather_loss, eff, area))

    step = 0.1

    x = 24*s_date:step:24*(e_date)+24
    x = filter(x -> x % 12 != 0, x)

    wh = 0

    for i in x
        wh += step * P(i)
    end

    return wh / 1000
end

function maxer()
    values = Float64[]
    for i in 0:1:90
        push!(values, kwt_alfa(i))
    end
    return findmax(values)[2] - 1
end

maxvalues = Int[]
# january
push!(maxvalues, maxer())

#february
s_date = 32
e_date = 59
push!(maxvalues, maxer())

#march
s_date = 60
e_date = 90
push!(maxvalues, maxer())

#april
s_date = 91
e_date = 120
push!(maxvalues, maxer())

#may
s_date = 121
e_date = 151
push!(maxvalues, maxer())

#june
s_date = 152
e_date = 181
push!(maxvalues, maxer())

#july
s_date = 182
e_date = 212
push!(maxvalues, maxer())

#august
s_date = 213
e_date = 243
push!(maxvalues, maxer())

#september
s_date = 244
e_date = 273
push!(maxvalues, maxer())

#october
s_date = 274
e_date = 304
push!(maxvalues, maxer())

#november
s_date = 305
e_date = 334
push!(maxvalues, maxer())

#december
s_date = 335
e_date = 365
push!(maxvalues, maxer())

plot(1:12, maxvalues, label = "", xticks = 1:12)

xlabel!("Optimal elevationsvinkel[grader]")
ylabel!("Månad")

savefig("blub.svg")


using Test
using ConvDiffMIPDECO
using jInv.Mesh
using LinearAlgebra

domain = [0 1. 0 1 0 1]
n  = [7 8 9]
M2D  = getRegularMesh(domain[1:4],n[1:2])

idg,idgn,idint = getGhostIndices(n[1:2])

xn = getNodalGrid(getPaddedMesh(M2D))

@test norm(xn[idg[1],1].- (-M2D.h[1]/2))<1e-15
@test norm(xn[idg[2],2].- (-M2D.h[2]/2))<1e-15
@test norm(xn[idg[3],1].- (1+M2D.h[1]/2))<1e-15
@test norm(xn[idg[4],2].- (1+M2D.h[2]/2))<1e-15
@test norm(xn[idgn[1],1].- (M2D.h[1]/2))<1e-15
@test norm(xn[idgn[2],2].- (M2D.h[2]/2))<1e-15
@test norm(xn[idgn[3],1].- (1-M2D.h[1]/2))<1e-15
@test norm(xn[idgn[4],2].- (1-M2D.h[2]/2))<1e-15
@test length(idgn[1])==n[2]
@test length(idgn[2])==n[1]
@test length(idgn[3])==n[2]
@test length(idgn[4])==n[1]
@test length(idg[1])==n[2]
@test length(idg[2])==n[1]
@test length(idg[3])==n[2]
@test length(idg[4])==n[1]
@test length(idint)==prod(n[1:2])

M3D  = getRegularMesh(domain,n)

idg,idgn,idint = getGhostIndices(n)

xn = getNodalGrid(getPaddedMesh(M3D))

@test norm(xn[idg[1],1].- (-M3D.h[1]/2))  < 1e-15
@test norm(xn[idg[2],2].- (-M3D.h[2]/2))  < 1e-15
@test norm(xn[idg[3],3].- (-M3D.h[3]/2))  < 1e-15
@test norm(xn[idg[4],1].- (1+M3D.h[1]/2)) < 1e-15
@test norm(xn[idg[5],2].- (1+M3D.h[2]/2)) < 1e-15
@test norm(xn[idg[6],3].- (1+M3D.h[3]/2)) < 1e-15
@test norm(xn[idgn[1],1].- (M3D.h[1]/2))  < 1e-15
@test norm(xn[idgn[2],2].- (M3D.h[2]/2))  < 1e-15
@test norm(xn[idgn[3],3].- (M3D.h[3]/2))  < 1e-15
@test norm(xn[idgn[4],1].- (1-M3D.h[1]/2))  < 1e-15
@test norm(xn[idgn[5],2].- (1-M3D.h[2]/2))  < 1e-15
@test norm(xn[idgn[6],3].- (1-M3D.h[3]/2))  < 1e-15

@test length(idgn[1])==prod(n[2:3])
@test length(idgn[2])==prod(n[[1,3]])
@test length(idgn[3])==prod(n[1:2])
@test length(idgn[4])==prod(n[2:3])
@test length(idgn[5])==prod(n[[1,3]])
@test length(idgn[6])==prod(n[1:2])

@test length(idg[1])==prod(n[2:3])
@test length(idg[2])==prod(n[[1,3]])
@test length(idg[3])==prod(n[1:2])
@test length(idg[4])==prod(n[2:3])
@test length(idg[5])==prod(n[[1,3]])
@test length(idg[6])==prod(n[1:2])

@test length(idint)==prod(n)


for N in [Float32, Float64, Rational{Int}]
    # Example 6 from [1].
    #
    # [1] Tommaso Dreossi, Thao Dang, and Carla Piazza. *Reachability computation for polynomial dynamical systems.*
    #      Formal Methods in System Design 50.1 (2017): 1-38.
    D = N[-1 0 0; -1 -1 0; 0 0 -1]
    c = N[-0.80, -0.95, 0.0, 0.85, 1.0, 0.0]
    P = HParallelotope(D, c)

    # test getter functions
    @test directions(P) == D
    @test offset(P) == c
    @test dim(P) == 3

    # check constructor's size check
    @test_throws AssertionError HParallelotope(D, vcat(c, c))

    # test computation of the list of constraints
    @test constraints_list(P) == [HalfSpace(D[1, :], c[1]),
                                  HalfSpace(D[2, :], c[2]),
                                  HalfSpace(D[3, :], c[3]),
                                  HalfSpace(-D[1, :], c[4]),
                                  HalfSpace(-D[2, :], c[5]),
                                  HalfSpace(-D[3, :], c[6])]

    P = HParallelotope(N[1 0; 0 1], N[1, 1, 1, 1])

    # test center
    @test center(P) == N[0, 0]

    # test vertices functions
    @test base_vertex(P) == N[-1, -1]
    @test extremal_vertices(P) == [N[1, -1], N[-1, 1]]

    # test generators getters
    @test genmat(P) == N[1 0; 0 1]

    # random parallelotope
    rand(HParallelotope)
end

for N in [Float32, Float64]
    # conversion from zonotope
    Z = Zonotope(N[0, 0], N[1 0; 0 1])
    @test convert(HParallelotope, Z) ==
        HParallelotope(N[0 -1; 1 0], N[1, 1, 1, 1])
    Z = Zonotope(N[0, 0], N[1 0 1; 0 1 1])
    @test_throws AssertionError convert(HParallelotope, Z)
end

for N in [Float64]
    P = HParallelotope(N[1 0; 0 1], N[1, 1, 1, 1])
    @test collect(generators(P)) == [N[1, 0], N[0, 1]]
end


mutable struct NormalSwitches{T1} <: AbstractDistributionSwitches{T1}
    pvec :: ParamVector
    # logNormal :: Bool
end

mutable struct Normal{T1} <: AbstractDistributionLH{T1}
    objId :: ObjectId
    switches :: NormalSwitches{T1}
    mean :: T1
    std :: T1
end

## ------------  Generic

Lazy.@forward Normal.switches (
    ModelParams.get_pvector
);

# function exp_draws!(drawV)
#     for (j, x) in enumerate(drawV)
#         drawV[j] = exp(x);
#     end
# end

function draw(u :: Normal{T1}, nDims, rng :: AbstractRNG) where T1
    drawV = u.mean .+ u.std .* randn(rng, T1, nDims...);
    # is_lognormal(u)  &&  exp_draws!(drawV);
    return drawV
end

function draw(u :: Normal{T1}, rng :: AbstractRNG)  where T1
    drawV = u.mean + u.std * rand(rng, T1);
    # is_lognormal(u)  &&  (drawV = exp(drawV));
    return drawV
end

function quantiles(u :: Normal{T1}, pctV)  where T1
    normalDistr = Distributions.Normal(u.mean, u.std);
    qV = Distributions.quantile.(normalDistr, pctV);
    # is_lognormal(u)  &&  exp_draws!(qV);
    return qV
end

# is_lognormal(switches :: NormalSwitches{T1}) where T1 = switches.logNormal;

Base.eltype(::Normal{T1}) where T1 = T1;

isbounded(u :: Normal{T1}) where T1 = false;

Base.show(io :: IO, um :: Normal{T1}) where T1 =
    print(io, "Normal{$T1}($(um.mean), $(um.std))");

# Validate draws
function check_draws(u :: Normal{T1}, drawM) where T1
    isValid = true;
    # if is_lognormal(u)
    #     isValid = isValid  &&  all(x -> x > zero(T1), drawM);
    # end
    return isValid
end

## ------------  Construction

ModelParams.get_pvector(switches :: NormalSwitches{T1}) where T1 = switches.pvec;
# ModelParams.get_pvector(nd :: Normal{T1}) where T1 = nd.switches.pvec;

init_distribution(objId :: ObjectId, switches :: NormalSwitches{T1}) where T1 =
    init_normal(switches);

function init_normal(switches :: NormalSwitches{T1}) where T1
    pvec = get_pvector(switches);
    return Normal{T1}(switches.pvec.objId, switches,
        param_default_value(pvec, :mean),
        param_default_value(pvec, :std))
end

function init_normal_switches(objId :: ObjectId;
    meanVal = 0.0, stdVal = 1.0, calMean = true, calStd = true)
    T1 = typeof(meanVal);
    pMean = init_mean(; meanVal, calMean);
    pStd = init_std(; stdVal, calStd);
    pvec = ParamVector(objId, [pMean, pStd]);
    return NormalSwitches{T1}(pvec)
end

function init_mean(; descr = "Mean", lsym = "μ", meanVal = 0.0,
    meanLb = -10.0, meanUb = 10.0, calMean = true)
    return Param(:mean, descr, lsym, meanVal, meanVal, meanLb, meanUb, calMean);
end

function init_std(; descr = "StdDev", lsym = "σ", stdVal = 1.0,
    stdLb = 0.0, stdUb = 10.0, calStd = true)
    return Param(:std, descr, lsym, stdVal, stdVal, stdLb, stdUb, calStd);
end

make_test_normal_switches() = init_normal_switches(ObjectId(:none);
    meanVal = 1.2, stdVal = 2.3);
make_test_normal() = init_normal(make_test_normal_switches());

# ---------------


__precompile__()

"""

# GeoEfficiency Package
introduce a fast and flexible tool to calculate in batch or individually the `geometrical efficiency`
for a set of common radiation detectors shapes (cylindrical,Bore-hole, Well-type) as seen form
a source. The source can be a point, a disc or even a cylinder.

# Quick Usage
*  geoEff()	: Calculate the geometrical efficiency for one geometrical setup return only the value of the geometrical efficiency.\n
*  calc() 	: Calculate the geometrical efficiency for one geometrical setup and display full information on the console.\n
*  calcN()	: Calculate the geometrical efficiency for geometrical setup(s) and display full information on the console until the user quit.\n
*  batch()	: Calculate the geometrical efficiency using data in the **``$dataDir``** folder in batch mode.

**for more information and updates refer to the repository at [``GitHub.com``](https://github.com/DrKrar/GeoEfficiency.jl/)**

"""
module GeoEfficiency

export # Config
    about,

 # Input_Interface
	getDetectors,
	setSrcToPoint,
	typeofSrc,

 # Physics_Model
	Point,
	source,
	Detector,
	CylDetector,
	BoreDetector,
	WellDetector,

 # Calculations
	geoEff,

 # Output_Interface
	calc,
	calcN,
	batch,
	batchInfo

include("Config.jl")
include("Physics_Model.jl")
include("Input_Interface.jl")
include("Calculations.jl")
include("Output_Interface.jl")

about()
printstyled("""
\n\tBatch mode
\t-  read files from directory `$dataDir`
\t-  save results to directory `$resultdir`
\n\tfor more information see `batch`, `prepare_batch`.
\n"""
,  color=:green)

end #module


# https://adrhill.github.io/ExplainableAI.jl/stable/generated/advanced_lrp/#How-it-works-internally
abstract type AbstractLRPRule end

# Generic LRP rule. Since it uses autodiff, it is used as a fallback for layer types without custom implementation.
function lrp!(Rₖ, rule::R, layer::L, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule,L}
    lrp_autodiff!(Rₖ, rule, layer, aₖ, Rₖ₊₁)
    return nothing
end

function lrp_autodiff!(
    Rₖ::T1, rule::R, layer::L, aₖ::T1, Rₖ₊₁::T2
) where {R<:AbstractLRPRule,L,T1,T2}
    layerᵨ = modify_layer(rule, layer)
    c::T1 = only(
        gradient(aₖ) do a
            z::T2 = layerᵨ(a)
            s = Zygote.@ignore Rₖ₊₁ ./ modify_denominator(rule, z)
            z ⋅ s
        end,
    )
    Rₖ .= aₖ .* c
    return nothing
end

# For linear layer types such as Dense layers, using autodiff is overkill.
function lrp!(Rₖ, rule::R, layer::Dense, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
    lrp_dense!(Rₖ, rule, layer, aₖ, Rₖ₊₁)
    return nothing
end

function lrp_dense!(Rₖ, rule::R, l, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
    ρW, ρb = modify_params(rule, get_params(l)...)
    ãₖ₊₁ = modify_denominator(rule, ρW * aₖ .+ ρb)
    @tullio Rₖ[j, b] = aₖ[j, b] * ρW[k, j] / ãₖ₊₁[k, b] * Rₖ₊₁[k, b]
    return nothing
end

# Other special cases that are dispatched on layer type:
lrp!(Rₖ, ::AbstractLRPRule, ::DropoutLayer, aₖ, Rₖ₊₁) = (Rₖ .= Rₖ₊₁)
lrp!(Rₖ, ::AbstractLRPRule, ::ReshapingLayer, aₖ, Rₖ₊₁) = (Rₖ .= reshape(Rₖ₊₁, size(aₖ)))

# To implement new rules, we can define two custom functions `modify_params` and `modify_denominator`.
# If this isn't done, the following fallbacks are used by default:
"""
    modify_params(rule, W, b)

Function that modifies weights and biases before applying relevance propagation.
Returns modified weights and biases as a tuple `(ρW, ρb)`.
"""
modify_params(::AbstractLRPRule, W, b) = (W, b) # general fallback

"""
    modify_denominator(rule, d)

Function that modifies zₖ on the forward pass, e.g. for numerical stability.
"""
modify_denominator(::AbstractLRPRule, d) = stabilize_denom(d, 1.0f-9) # general fallback

"""
    modify_layer(rule, layer)

Function that modifies a layer before applying relevance propagation.
Returns a new, modified layer.
"""
modify_layer(::AbstractLRPRule, layer) = layer # skip layers without modify_params
function modify_layer(rule::R, layer::L) where {R<:AbstractLRPRule,L<:Union{Dense,Conv}}
    return set_params(layer, modify_params(rule, get_params(layer)...)...)
end

"""
    ZeroRule()

Constructor for LRP-0 rule. Commonly used on upper layers.
"""
struct ZeroRule <: AbstractLRPRule end

"""
    GammaRule([γ=0.25])

Constructor for LRP-``γ`` rule. Commonly used on lower layers.

Arguments:
- `γ`: Optional multiplier for added positive weights, defaults to 0.25.
"""
struct GammaRule{T} <: AbstractLRPRule
    γ::T
    GammaRule(γ=0.25f0) = new{Float32}(γ)
end
function modify_params(r::GammaRule, W, b)
    T = eltype(W)
    ρW = W + convert(T, r.γ) * relu.(W)
    ρb = b + convert(T, r.γ) * relu.(b)
    return ρW, ρb
end

"""
    EpsilonRule([ϵ=1.0f-6])

Constructor for LRP-``ϵ`` rule. Commonly used on middle layers.

Arguments:
- `ϵ`: Optional stabilization parameter, defaults to `1f-6`.
"""
struct EpsilonRule{T} <: AbstractLRPRule
    ϵ::T
    EpsilonRule(ϵ=1.0f-6) = new{Float32}(ϵ)
end
modify_denominator(r::EpsilonRule, d) = stabilize_denom(d, r.ϵ)

"""
    ZBoxRule(low, high)

Constructor for LRP-``z^{\\mathcal{B}}``-rule.
Commonly used on the first layer for pixel input.

The parameters `low` and `high` should be set to the lower and upper bounds of the input features,
e.g. `0.0` and `1.0` for raw image data.
It is also possible to provide two arrays of that match the input size.
"""
struct ZBoxRule{T} <: AbstractLRPRule
    low::T
    high::T
end

# The ZBoxRule requires its own implementation of relevance propagation.
lrp!(Rₖ, r::ZBoxRule, layer::Dense, aₖ, Rₖ₊₁) = lrp_zbox!(Rₖ, r, layer, aₖ, Rₖ₊₁)
lrp!(Rₖ, r::ZBoxRule, layer::Conv, aₖ, Rₖ₊₁) = lrp_zbox!(Rₖ, r, layer, aₖ, Rₖ₊₁)

_zbox_bound(T, c::Real, in_size) = fill(convert(T, c), in_size)
function _zbox_bound(T, A::AbstractArray, in_size)
    size(A) != in_size && throw(
        ArgumentError(
            "Bounds `low`, `high` of ZBoxRule should either be scalar or match input size.",
        ),
    )
    return convert.(T, A)
end

function lrp_zbox!(Rₖ::T1, r::ZBoxRule, layer::L, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
    T = eltype(aₖ)
    in_size = size(aₖ)
    l = _zbox_bound(T, r.low, in_size)
    h = _zbox_bound(T, r.high, in_size)

    W, b = get_params(layer)
    layer⁺ = set_params(layer, max.(0, W), max.(0, b)) # W⁺, b⁺
    layer⁻ = set_params(layer, min.(0, W), min.(0, b)) # W⁻, b⁻

    c::T1, cₗ::T1, cₕ::T1 = gradient(aₖ, l, h) do a, l, h
        f::T2 = layer(a)
        f⁺::T2 = layer⁺(l)
        f⁻::T2 = layer⁻(h)

        z = f - f⁺ - f⁻
        s = Zygote.@ignore safedivide(Rₖ₊₁, z, 1e-9)
        z ⋅ s
    end
    Rₖ .= aₖ .* c + l .* cₗ + h .* cₕ
    return nothing
end


"Soft Boolean.  Value in [0, 1]"
struct SoftBool{ET <: Real} <: AbstractSoftBool
  logerr::ET
  SoftBool(l::T) where {T} = new{T}(l)
  SoftBool(l::T) where {T <: ForwardDiff.Dual} = new{T}(l) # Resolves type ambiguity (from flux)
end
@invariant 0 <= err(b::SoftBool) <= 1

"Error in [0, 1]"
err(x::SoftBool) = exp(x.logerr)

"Log error in `[-Inf, 0]`"
logerr(x::SoftBool) = x.logerr
Bool(x::SoftBool) = logerr(x) == 0.0
ssofttrue(::Type{T} = Float64) where T = SoftBool(zero(t))
ssoftfalse(::Type{T} = Float64) where T = SoftBool(-inf(T))

# (In)Equalities #

"Kernel return type as function of arguments"
kernelrettype(x::T, y::T) where T = T

softeq(a::Bool, b::Bool) = softeq(float(a), float(b))
softeq(a::Integer, b::Integer) = softeq(float(a), float(b))

"Soft Equality"
function ssofteq(x, y, k = globalkernel())
  r = d(x, y)
  SoftBool(k(r)::typeof(r))
end

"Soft >"
function ssoftgt(x::Real, y::Real, k = globalkernel())
  r = bound_loss(x, y, Inf)
  SoftBool(k(r)::typeof(r))
end

"Soft <"
function ssoftlt(x::Real, y::Real, k = globalkernel())
  r = bound_loss(x, -Inf, y)
  SoftBool(k(r)::typeof(r))
end

# Boolean Operators #
function Base.:&(x::SoftBool, y::SoftBool)
  a = logerr(x)
  b = logerr(y)
  # c = min(a, b)
  c = a + b
  SoftBool(c)
end
Base.:|(x::SoftBool, y::SoftBool) = SoftBool(max(logerr(x), logerr(y)))
# zt: FIXME abstractvector?
Base.all(xs::Vector{<:SoftBool}) = SoftBool(minimum(logerr.(xs)))
Base.all(xs::Vector{<:RandVar}) = RandVar(all, (xs, ))

# Arithmetic #
Base.:*(x::SoftBool{T}, y::T) where T <: Real = SoftBool(x.logerr * y)
Base.:*(x::T, y::SoftBool{T}) where T <: Real = SoftBool(x * y.logerr)

# Show #
Base.show(io::IO, sb::SoftBool) = print(io, "ϵ:$(logerr(sb))")