text:
using Test
using Distributions
using StatsFuns

@testset "loglik list (individual)" begin

    d = Distributions.Gamma(1.0, 2.0)

    μμ = rand(d, 5)
    σσ = rand(d, 5)

    for (μ, σ) in zip(μμ, σσ)
        l = Distributions.LogNormal(μ, σ)
        y = rand(l, 25)

        # LogNormal
        r = LRMoE.LogNormalExpert(μ, σ)

        # Y = hcat(y, y, y, y)
        # model = [LRMoE.LogNormalExpert(μ, σ) LRMoE.LogNormalExpert(0.5*μ, 0.6*σ) LRMoE.LogNormalExpert(1.5*μ, 2.0*σ)]

        Y = hcat(fill(0, length(y)), y, y, fill(Inf, length(y)), fill(0, length(y)), 0.80.*y, 1.25.*y, fill(Inf, length(y)))
        model = [LRMoE.LogNormalExpert(μ, σ) LRMoE.LogNormalExpert(0.5*μ, 0.6*σ) LRMoE.LogNormalExpert(1.5*μ, 2.0*σ);
                LRMoE.ZILogNormalExpert(0.4, μ, σ) LRMoE.LogNormalExpert(0.5*μ, 0.6*σ) LRMoE.ZILogNormalExpert(0.80, 1.5*μ, 2.0*σ)]

        # Test size
        @test size(expert_ll_pos_list(Y, model))[1] == 2
        @test size(expert_ll_pos_list(Y, model)[1])[1] == 25
        @test size(expert_ll_pos_list(Y, model)[1])[2] == 3

        @test size(expert_tn_pos_list(Y, model))[1] == 2
        @test size(expert_tn_pos_list(Y, model)[1])[1] == 25
        @test size(expert_tn_pos_list(Y, model)[1])[2] == 3

        @test size(expert_tn_bar_pos_list(Y, model))[1] == 2
        @test size(expert_tn_bar_pos_list(Y, model)[1])[1] == 25
        @test size(expert_tn_bar_pos_list(Y, model)[1])[2] == 3

        @test size(expert_ll_list(Y, model))[1] == 2
        @test size(expert_ll_list(Y, model)[1])[1] == 25
        @test size(expert_ll_list(Y, model)[1])[2] == 3

        @test size(expert_tn_list(Y, model))[1] == 2
        @test size(expert_tn_list(Y, model)[1])[1] == 25
        @test size(expert_tn_list(Y, model)[1])[2] == 3

        @test size(expert_tn_bar_list(Y, model))[1] == 2
        @test size(expert_tn_bar_list(Y, model)[1])[1] == 25
        @test size(expert_tn_bar_list(Y, model)[1])[2] == 3

        # 1*1 model
        y = rand(l, 1)
        Y = [0 y y Inf]
        model = reshape([LRMoE.LogNormalExpert(μ, σ)], (1, 1))
        @test expert_ll_pos_list(Y, model)[1] ≈ [Distributions.logpdf.(l, Y[2])]
        @test expert_tn_pos_list(Y, model)[1] ≈ [0.0]
        @test expert_tn_bar_pos_list(Y, model)[1] ≈ [-Inf]
        @test expert_ll_list(Y, model)[1] ≈ [Distributions.logpdf.(l, Y[2])]
        @test expert_tn_list(Y, model)[1] ≈ [0.0]
        @test expert_tn_bar_list(Y, model)[1] ≈ [-Inf]

        y = rand(l, 1)
        Y = [0 0.80*y 1.25*y Inf]
        model = reshape([LRMoE.LogNormalExpert(μ, σ)], (1, 1))
        @test expert_ll_pos_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[3]) + log1mexp(Distributions.logcdf.(l, Y[2]) - Distributions.logcdf.(l, Y[3]))]
        @test expert_tn_pos_list(Y, model)[1] ≈ [0.0]
        @test expert_tn_bar_pos_list(Y, model)[1] ≈ [-Inf]
        @test expert_ll_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[3]) + log1mexp(Distributions.logcdf.(l, Y[2]) - Distributions.logcdf.(l, Y[3]))]
        @test expert_tn_list(Y, model)[1] ≈ [0.0]
        @test expert_tn_bar_list(Y, model)[1] ≈ [-Inf]

        y = rand(l, 1)
        Y = [0.50*y 0.80*y 1.25*y 2.00*y]
        model = reshape([LRMoE.LogNormalExpert(μ, σ)], (1, 1))
        @test expert_ll_pos_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[3]) + log1mexp.(Distributions.logcdf.(l, Y[2]) - Distributions.logcdf.(l, Y[3]))]
        @test expert_tn_pos_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))]
        @test expert_tn_bar_pos_list(Y, model)[1] ≈ log1mexp.([Distributions.logcdf.(l, Y[4]) + log1mexp(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))])
        @test expert_ll_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[3]) + log1mexp.(Distributions.logcdf.(l, Y[2]) - Distributions.logcdf.(l, Y[3]))]
        @test expert_tn_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))]
        @test expert_tn_bar_list(Y, model)[1] ≈ log1mexp.([Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))])

        # 2*1 model
        p = rand(Distributions.Uniform(0.0, 1.0), 1)[1]
        y = rand(l, 2) .+ 0.5
        Y = [0.0 0.80*y[1] 1.25*y[1] Inf 0.0 0.80*y[2] 1.25*y[2] Inf]
        model = reshape([LRMoE.LogNormalExpert(μ, σ);
                 LRMoE.ZILogNormalExpert(p, μ, σ)], (2, 1))
        @test expert_ll_pos_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[3]) + log1mexp.(Distributions.logcdf.(l, Y[2]) - Distributions.logcdf.(l, Y[3]))]
        @test expert_ll_pos_list(Y, model)[2] ≈ [Distributions.logcdf.(l, Y[7]) + log1mexp.(Distributions.logcdf.(l, Y[6]) - Distributions.logcdf.(l, Y[7]))]

        @test expert_tn_pos_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))]
        @test expert_tn_pos_list(Y, model)[2] ≈ [Distributions.logcdf.(l, Y[8]) + log1mexp.(Distributions.logcdf.(l, Y[5]) - Distributions.logcdf.(l, Y[8]))]

        @test expert_tn_bar_pos_list(Y, model)[1] ≈ log1mexp.([Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))])
        @test expert_tn_bar_pos_list(Y, model)[2] ≈ log1mexp.([Distributions.logcdf.(l, Y[8]) + log1mexp.(Distributions.logcdf.(l, Y[5]) - Distributions.logcdf.(l, Y[8]))])

        @test expert_ll_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[3]) + log1mexp.(Distributions.logcdf.(l, Y[2]) - Distributions.logcdf.(l, Y[3]))]
        @test expert_ll_list(Y, model)[2] ≈ [log(1-p) .+ Distributions.logcdf.(l, Y[7]) + log1mexp.(Distributions.logcdf.(l, Y[6]) - Distributions.logcdf.(l, Y[7]))]

        @test expert_tn_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))]
        @test expert_tn_list(Y, model)[2] ≈ [Distributions.logcdf.(l, Y[8]) + log1mexp.(Distributions.logcdf.(l, Y[5]) - Distributions.logcdf.(l, Y[8]))]

        @test expert_tn_bar_list(Y, model)[1] ≈ log1mexp.([Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))])
        @test expert_tn_bar_list(Y, model)[2] ≈ log1mexp.([Distributions.logcdf.(l, Y[8]) + log1mexp.(Distributions.logcdf.(l, Y[5]) - Distributions.logcdf.(l, Y[8]))])

        p = rand(Distributions.Uniform(0.0, 1.0), 1)[1]
        y = rand(l, 2) .+ 0.5
        Y = [0.0 0.80*y[1] 1.25*y[1] Inf 0.50*y[2] 0.80*y[2] 1.25*y[2] 2.50*y[2]]
        model = reshape([LRMoE.LogNormalExpert(μ, σ);
                LRMoE.ZILogNormalExpert(p, μ, σ)], (2, 1))
        @test expert_ll_pos_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[3]) + log1mexp.(Distributions.logcdf.(l, Y[2]) - Distributions.logcdf.(l, Y[3]))]
        @test expert_ll_pos_list(Y, model)[2] ≈ [Distributions.logcdf.(l, Y[7]) + log1mexp.(Distributions.logcdf.(l, Y[6]) - Distributions.logcdf.(l, Y[7]))]

        @test expert_tn_pos_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))]
        @test expert_tn_pos_list(Y, model)[2] ≈ [Distributions.logcdf.(l, Y[8]) + log1mexp.(Distributions.logcdf.(l, Y[5]) - Distributions.logcdf.(l, Y[8]))]

        @test expert_tn_bar_pos_list(Y, model)[1] ≈ log1mexp.([Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))])
        @test expert_tn_bar_pos_list(Y, model)[2] ≈ log1mexp.([Distributions.logcdf.(l, Y[8]) + log1mexp.(Distributions.logcdf.(l, Y[5]) - Distributions.logcdf.(l, Y[8]))])

        @test expert_ll_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[3]) + log1mexp.(Distributions.logcdf.(l, Y[2]) - Distributions.logcdf.(l, Y[3]))]
        @test expert_ll_list(Y, model)[2] ≈ [log(1-p) .+ Distributions.logcdf.(l, Y[7]) + log1mexp.(Distributions.logcdf.(l, Y[6]) - Distributions.logcdf.(l, Y[7]))]

        @test expert_tn_list(Y, model)[1] ≈ [Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))]
        @test expert_tn_list(Y, model)[2] ≈ [log(1-p) .+ Distributions.logcdf.(l, Y[8]) + log1mexp.(Distributions.logcdf.(l, Y[5]) - Distributions.logcdf.(l, Y[8]))]

        @test expert_tn_bar_list(Y, model)[1] ≈ log1mexp.([Distributions.logcdf.(l, Y[4]) + log1mexp.(Distributions.logcdf.(l, Y[1]) - Distributions.logcdf.(l, Y[4]))])
        @test expert_tn_bar_list(Y, model)[2] ≈ [log.(p .+ (1-p) .* exp.(log1mexp.(Distributions.logcdf.(l, Y[8]) + log1mexp.(Distributions.logcdf.(l, Y[5]) - Distributions.logcdf.(l, Y[8])))))]

        # 2*3 model
        p1 = rand(Distributions.Uniform(0.0, 1.0), 1)[1]
        p2 = rand(Distributions.Uniform(0.0, 1.0), 1)[1]
        y = rand(l, 100, 2) .+ 0.5
        Y = [fill(0.0, length(y[:,1])) 0.80.*y[:,1] 1.25.*y[:,1] fill(Inf, length(y[:,1])) 0.50.*y[:,2] 0.80.*y[:,2] 1.25.*y[:,2] 2.50.*y[:,2]]
        model = [LRMoE.LogNormalExpert(μ, σ) LRMoE.ZILogNormalExpert(p1, 1.2*μ, 0.8*σ);
                LRMoE.ZILogNormalExpert(p2, 2*μ, 0.5*σ) LRMoE.LogNormalExpert(0.9*μ, 0.75*σ)]
        l1 = Distributions.LogNormal(μ, σ)
        l2 = Distributions.LogNormal(1.2*μ, 0.8*σ)
        l3 = Distributions.LogNormal(2*μ, 0.5*σ)
        l4 = Distributions.LogNormal(0.9*μ, 0.75*σ)

        @test expert_ll_pos_list(Y, model)[1] ≈ hcat(Distributions.logcdf.(l1, Y[:,3]) + log1mexp.(Distributions.logcdf.(l1, Y[:,2]) - Distributions.logcdf.(l1, Y[:,3])),
                                                    Distributions.logcdf.(l2, Y[:,3]) + log1mexp.(Distributions.logcdf.(l2, Y[:,2]) - Distributions.logcdf.(l2, Y[:,3])))
        @test expert_ll_pos_list(Y, model)[2] ≈ hcat(Distributions.logcdf.(l3, Y[:,7]) + log1mexp.(Distributions.logcdf.(l3, Y[:,6]) - Distributions.logcdf.(l3, Y[:,7])),
                                                    Distributions.logcdf.(l4, Y[:,7]) + log1mexp.(Distributions.logcdf.(l4, Y[:,6]) - Distributions.logcdf.(l4, Y[:,7])))

        @test expert_tn_pos_list(Y, model)[1] ≈ hcat(fill(0.0, length(Y[:,2])),
                                                    Distributions.logcdf.(l2, Y[:,4]) + log1mexp.(Distributions.logcdf.(l2, Y[:,1]) - Distributions.logcdf.(l2, Y[:,4])))
        @test expert_tn_pos_list(Y, model)[2] ≈ hcat(Distributions.logcdf.(l3, Y[:,8]) + log1mexp.(Distributions.logcdf.(l3, Y[:,5]) - Distributions.logcdf.(l3, Y[:,8])),
                                                    Distributions.logcdf.(l4, Y[:,8]) + log1mexp.(Distributions.logcdf.(l4, Y[:,5]) - Distributions.logcdf.(l4, Y[:,8])))

        # @test expert_tn_bar_pos_list(Y, model)[1] ≈ hcat(fill(-Inf, length(Y[:,2])),
        #                                             log1mexp.(Distributions.logcdf.(l2, Y[:,4]) + log1mexp.(Distributions.logcdf.(l2, Y[:,1]) - Distributions.logcdf.(l2, Y[:,4]))))
        # @test expert_tn_bar_pos_list(Y, model)[2] ≈ hcat(log1mexp.(Distributions.logcdf.(l3, Y[:,8]) + log1mexp.(Distributions.logcdf.(l3, Y[:,5]) - Distributions.logcdf.(l3, Y[:,8]))),
        #                                             log1mexp.(Distributions.logcdf.(l4, Y[:,8]) + log1mexp.(Distributions.logcdf.(l4, Y[:,5]) - Distributions.logcdf.(l4, Y[:,8]))))

        # @test expert_ll_list(Y, model)[1] ≈ hcat(Distributions.logcdf.(l1, Y[:,3]) + log1mexp.(Distributions.logcdf.(l1, Y[:,2]) - Distributions.logcdf.(l1, Y[:,3])),
        #                                         log(1-p1) .+ Distributions.logcdf.(l2, Y[:,3]) + log1mexp.(Distributions.logcdf.(l2, Y[:,2]) - Distributions.logcdf.(l2, Y[:,3])))
        # @test expert_ll_list(Y, model)[2] ≈ hcat(log(1-p2) .+ Distributions.logcdf.(l3, Y[:,7]) + log1mexp.(Distributions.logcdf.(l3, Y[:,6]) - Distributions.logcdf.(l3, Y[:,7])),
        #                                             Distributions.logcdf.(l4, Y[:,7]) + log1mexp.(Distributions.logcdf.(l4, Y[:,6]) - Distributions.logcdf.(l4, Y[:,7])))

        @test expert_tn_list(Y, model)[1] ≈ hcat(fill(0.0, length(Y[:,2])),
                                                fill(0.0, length(Y[:,2])))

        # @test expert_tn_list(Y, model)[2] ≈ hcat(log(1-p2) .+ Distributions.logcdf.(l3, Y[:,8]) .+ log1mexp.(Distributions.logcdf.(l3, Y[:,5]) - Distributions.logcdf.(l3, Y[:,8])),
        #                                         Distributions.logcdf.(l4, Y[:,8]) + log1mexp.(Distributions.logcdf.(l4, Y[:,5]) - Distributions.logcdf.(l4, Y[:,8])))

        # @test isapprox(expert_tn_list(Y, model)[2], hcat(log(1-p2) .+ Distributions.logcdf.(l3, Y[:,8]) .+ log1mexp.(Distributions.logcdf.(l3, Y[:,5]) - Distributions.logcdf.(l3, Y[:,8])),
        #                                         Distributions.logcdf.(l4, Y[:,8]) + log1mexp.(Distributions.logcdf.(l4, Y[:,5]) - Distributions.logcdf.(l4, Y[:,8]))), atol = 1e-03)

        @test expert_tn_bar_list(Y, model)[1] ≈ hcat(fill(-Inf, length(Y[:,2])),
                                                    fill(-Inf, length(Y[:,2])))

        # @test expert_tn_bar_list(Y, model)[2] ≈ hcat(log.(p2.+ (1-p2) .* (1 .- exp.(Distributions.logcdf.(l3, Y[:,8]) + log1mexp.(Distributions.logcdf.(l3, Y[:,5]) - Distributions.logcdf.(l3, Y[:,8]))))),
        #                                             log1mexp.(Distributions.logcdf.(l4, Y[:,8]) + log1mexp.(Distributions.logcdf.(l4, Y[:,5]) - Distributions.logcdf.(l4, Y[:,8]))))

        # @test isapprox(expert_tn_bar_list(Y, model)[2], hcat(log.(p2.+ (1-p2) .* (1 .- exp.(Distributions.logcdf.(l3, Y[:,8]) + log1mexp.(Distributions.logcdf.(l3, Y[:,5]) - Distributions.logcdf.(l3, Y[:,8]))))),
        #                                             log1mexp.(Distributions.logcdf.(l4, Y[:,8]) + log1mexp.(Distributions.logcdf.(l4, Y[:,5]) - Distributions.logcdf.(l4, Y[:,8])))), atol = 1e-03)

    end

end

@testset "loglik list (aggredated) " begin
    d = Distributions.Gamma(1.0, 2.0)

    μμ = rand(d, 5)
    σσ = rand(d, 5)

    for (μ, σ) in zip(μμ, σσ)
        l = Distributions.LogNormal(μ, σ)
        y = rand(l, 25)

        # LogNormal
        r = LRMoE.LogNormalExpert(μ, σ)

        X = rand(Uniform(-1, 1), 25, 10)
        α = rand(Uniform(-1, 1), 3, 10)
        α[3,:] .= 0.0

        Y = hcat(fill(0, length(y)), y, y, fill(Inf, length(y)), fill(0, length(y)), 0.80.*y, 1.25.*y, fill(Inf, length(y)))
        model = [LRMoE.LogNormalExpert(μ, σ) LRMoE.LogNormalExpert(0.5*μ, 0.6*σ) LRMoE.LogNormalExpert(1.5*μ, 2.0*σ);
                LRMoE.ZILogNormalExpert(0.4, μ, σ) LRMoE.LogNormalExpert(0.5*μ, 0.6*σ) LRMoE.ZILogNormalExpert(0.80, 1.5*μ, 2.0*σ)]

        # Test size
        @test size(LogitGating(α, X))[1] == 25
        @test size(LogitGating(α, X))[2] == 3

        gate = LogitGating(α, X)
        ll_ls = loglik_np(Y, gate, model)
        # Test number
        @test isa(ll_ls.ll, Number)
        @test exp.(ll_ls.gate_expert_tn) + exp.(ll_ls.gate_expert_tn_bar) ≈ fill(1.0, length(ll_ls.gate_expert_tn))

    end
end

max_stars_repo_name: sparktseung/JuliaLRMoE

max_issues_repo_name: sparktseung/JuliaLRMoE

max_forks_repo_name: sparktseung/JuliaLRMoE


text:
module NeuroBootcamp

using PyCall
const animation = PyNULL()
function __init__()
    # silence the old "QApplication: invalid style override passed, ignoring
    #  it." warning...
    Base.eval(Main, :(ENV["QT_STYLE_OVERRIDE"] = ""))
    copy!(animation, pyimport("matplotlib.animation"))
end

using PyPlot

include("./LifConfig.jl")
using .LifConfig

include("./Lif.jl")
using .Lif

include("./Networks.jl")
using .Networks

include("./Histogram/Histogram.jl")

include("./Spk/Spk.jl")
using .Spk

include("./SpkPlot/SpkPlot.jl")
using .SpkPlot

import Base.run

export build_demo, build_network, SquareWave, SineWave, WhiteNoise, run_sim
export plot_raster, plot_raster!, plot_cycle_mean, plot_cycle_mean!, plot_xcorr,
    psth, get_xcorr, f1xfm, cycle_mean

abstract type Stimulus end
# ============================================================================ #
mutable struct ChunkArray
    nel::Int
    k::Int
    last::Float64
    d::Vector{Float64}
    # ------------------------------------------------------------------------ #
    function ChunkArray(n::Int)
        self = new()
        self.nel = n
        self.k = 1
        self.d = fill(NaN, n)
        self.last = NaN
        return self
    end
    # ------------------------------------------------------------------------ #
end
# ---------------------------------------------------------------------------- #
function add!(x::ChunkArray, v::Float64)
    x.last = v;
    if x.k < x.nel
        x.d[x.k] = v
        x.k += 1
    else
        x.d[1:end-1] = x.d[2:end]
        x.d[end] = v
    end
end
# ============================================================================ #
mutable struct LiveDemo
    net::LIFNetwork
    speed::Int
    fig::PyPlot.Figure
    ax::PyCall.PyObject
end
"""
    LiveDemo(net::Network, stim::Stimulus, speed::Int=5)
"""
function LiveDemo(net::LIFNetwork, speed::Int=5)
    h = default_figure()
    demo = LiveDemo(net, speed, h, h.axes[1])
    close(h)
    return demo
end
# ---------------------------------------------------------------------------- #
function reset!(demo::LiveDemo)
    # if we always close the figure we avoid the "not responding to keypresses"
    # issue *AND* we get focus back on the figure when run() exits...
    close(demo.fig)
    demo.fig = default_figure()
    demo.ax = demo.fig.axes[1]
    Networks.reset!(demo.net)
end
# ============================================================================ #
"""
    build_network([
            (1=>3, .8),
            (2=>3, .6)
        ], xi=0.0)
"""
function build_network(inp::Vector{Tuple{Pair{T,T}, F}}, xi::Real=0.0) where {T<:Integer, F<:Real}
    ncell = 0
    for x in inp
        ncell = max(ncell, maximum(x[1]))
    end

    net = LIFNetwork(ncell, xi)

    for x in net
        Lif.set!(x, :threshold, 1.0)
        Lif.set!(x, :rm, 1.0)
        Lif.set!(x, :tau, 4.5)
        Lif.set!(x, :vspike, 4.0)
    end

    connect!(net)
    for x in inp
        connect!(net, Tuple(x[1]), (Synapses.Static, x[2], 2.0, 2.0))
    end

    return net
end
# ============================================================================ #
"""
    build_demo([
            (1=>3, .8),
            (2=>3, .6)
        ], xi=0.0)
"""
function build_demo(inp::Vector{Tuple{Pair{T,T}, F}}, xi::Real=0.0) where {T<:Integer, F<:Real}
    return LiveDemo(build_network(inp, xi), 5)
end
"""
    build_demo([1=>3, 2=>3], xi=0.0)
"""
function build_demo(pairs::Vector{Pair{T,T}}, xi::Real=0.0) where {T<:Integer}
    return build_demo([(x, 1.8) for x in pairs], xi)
end
"""
    build_demo(2, xi=0.0)
"""
function build_demo(ncell::Integer, xi::Real=0.0)
    return build_demo([(x=>ncell, 1.8) for x in 1:(ncell-1)], xi)
end
# ============================================================================ #
mutable struct SquareWave <: Stimulus
    amp::Vector{Float64}
    on::Float64
    off::Float64
    tlast::Float64
    state::Bool
end
"""
    SquareWave(demo::LiveDemo, amp::Float64=1.0, on::Float64=20.0,
        off::Float64=30.0)
"""
function SquareWave(demo::LiveDemo, amp::Float64=1.0, on::Float64=20.0,
    off::Float64=30.0)

    return SquareWave([amp; zeros(Float64, length(demo.net)-1)], 20.0, 20.0,
        0.0, false)
end
function reset!(sq::SquareWave)
    sq.state = false
    sq.tlast = 0.0
end
function getstim(sq::SquareWave, id::Integer, t::Time)
    tnext = sq.tlast + (sq.state ? sq.on : sq.off)
    if t >= tnext
        sq.tlast = t
        sq.state = !sq.state
    end
    if sq.state && (0 < id <= length(sq.amp))
        return sq.amp[id]
    else
        return 0.0
    end
end
# ============================================================================ #
mutable struct SineWave <: Stimulus
    amp::Vector{Float64}
    freq::Vector{Float64}
end
"""
    SineWave(demo::LiveDemo, amp::Float64=1.0, freq::Float64=25.0)
"""
function SineWave(demo::LiveDemo, amp::Float64=1.0, freq::Float64=25.0)
    # frequency needs to be converted to cycles-per-ms for simulations
    return SineWave([amp; zeros(Float64, length(demo.net)-1)],
        fill(freq/1000.0, length(demo.net)))
end
reset!(sw::SineWave) = nothing
function getstim(sw::SineWave, id::Integer, t::Time)
    if 0 < id <= length(sw.amp)
        return sw.amp[id] * ((0.5 * sin(2.0*pi*sw.freq[id]*t)) + 0.5)
    end
    return 0.0
end
# ============================================================================ #
mutable struct WhiteNoise <: Stimulus
    last_frame::Float64
    last_value::Float64
    ifi::Float64
    mu::Float64
    sigma::Float64
    # history::Vector{Float64}
end
"""
    WhiteNoise(mu::Real=0.7, sigma::Real=0.3, ifi::Real=4.0)
"""
function WhiteNoise(mu::Real=0.7, sigma::Real=0.3, ifi::Real=4.0)
    return WhiteNoise(-Inf, sigma * randn() + mu, ifi, mu, sigma)#, Float64[])
end
function reset!(wn::WhiteNoise)
    wn.last_frame = -Inf
end
function getstim(wn::WhiteNoise, id::Integer, t::Time)
    if id == 1
        if t >= wn.last_frame + wn.ifi
            wn.last_frame = t
            wn.last_value = wn.sigma * randn() + wn.mu
            # push!(wn.history, wn.last_value)
        end
        return wn.last_value
    else
        return 0.0
    end
end
# ============================================================================ #
mutable struct Keyboard <: Stimulus
    state::Vector{Bool}
    amp::Vector{Float64}
    inc::Float64
end
# ---------------------------------------------------------------------------- #
Keyboard(demo::LiveDemo) = Keyboard(length(demo.net))
Keyboard(n::Integer) = Keyboard(falses(n), ones(Float64, n), 0.1)
Keyboard(a::AbstractArray{T,1}) where {T<:Real} = Keyboard(falses(a), a, 0.1)
# ============================================================================ #
function keypressed(s::Keyboard, key::String)
    if key in ["1", "2", "3"]
        id = parse(Int8, key)
        if 0 < id <= length(s.state)
            s.state[id] = !s.state[id]
        end
    elseif key == "down"
        s.amp .-= s.inc
    elseif key == "up"
        s.amp .+= s.inc
    end
end
# ============================================================================ #
function reset!(s::Keyboard)
    for k in eachindex(s.state)
        s.state[k] = false
    end
end
# ============================================================================ #
function getstim(s::Keyboard, id::Integer, t::Time)
    if 0 < id <= length(s.state)
        return s.state[id] ? s.amp[id] : 0.0
    else
        return 0.0
    end
end
# ============================================================================ #
keypressed(s::Stimulus, key::String) = nothing
# ============================================================================ #
"""
    run(demo::LiveDemo, duration::Real=+Inf)
    run(demo::LiveDemo, stim::Stimulus, duration::Real=+Inf)
"""
function run(demo::LiveDemo, duration::Real=+Inf)
    return run(demo, Keyboard(length(demo.net)), duration)
end
# ---------------------------------------------------------------------------- #
function run(demo::LiveDemo, stimgen::Stimulus, duration::Real=+Inf)

    TMAX = 50.0 # max time of x axis in MS

    dt = 0.05 # timestep in MS

    npt = round(Int64, TMAX / dt)

    ncell = length(demo.net)

    if isinf(duration)
        ncall = round(Int64, 200.0 / (dt * demo.speed))
        repeat = true
    else
        ncall = round(Int64, duration / (dt * demo.speed))
        repeat = false
    end

    TNOW = 0.0

    time = ChunkArray(npt)

    data = ChunkArray[ChunkArray(npt) for k in 1:ncell]
    stim = ChunkArray[ChunkArray(npt) for k in 1:ncell]

    reset!(demo)
    reset!(stimgen)

    sca(demo.ax)
    xlim(0, TMAX)
    ylim(-2.0, 4.0)

    local col::Vector{String}
    if ncell == 2
        col = ["blue", "red"]
    elseif ncell == 3
        col = ["blue", "cyan", "red"]
    elseif ncell == 4
        col = ["blue", "cyan", "lightblue", "red"]
    else
        error("Not enough colors!")
    end

    data_line = Vector{PyCall.PyObject}(undef, ncell)
    stim_line = Vector{PyCall.PyObject}(undef, ncell)
    for k = 1:ncell
        data_line[k] = plot([], [], color=col[k], linewidth=3)[1]
        stim_line[k] = plot([], [], color=col[k], linewidth=3)[1]
    end

    ts = [Vector{Float64}() for x in 1:ncell]

    fstim(id::Integer, t::Time) = getstim(stimgen, id, t)

    # ------------------------------------------------------------------------ #
    function run_demo(k)
        for kt = 1:demo.speed
            add!(time, TNOW)
            vm, spk = Networks.update!(demo.net, fstim, dt, TNOW)
            for kc = 1:ncell
                offset = .2 * (kc - 1)
                add!(data[kc], vm[kc] - offset)
                add!(stim[kc], fstim(kc, TNOW) - (1.8 - offset))
                if spk[kc]
                    push!(ts[kc], TNOW * 0.001)
                end
            end
            TNOW += dt
        end

        for kc = 1:ncell
            data_line[kc].set_data(time.d, data[kc].d)
            stim_line[kc].set_data(time.d, stim[kc].d)
        end

        if TNOW > TMAX
            demo.ax.set_xlim(time.d[1], time.last)
        end

        demo.fig.canvas.draw()
    end
    # ------------------------------------------------------------------------ #

    anim = animation.FuncAnimation(demo.fig, run_demo, frames=ncall,
        interval=10, repeat=repeat, blit=false)

    # ------------------------------------------------------------------------ #
    function key_press(event)
        if event.key == " "
            anim.event_source.stop()
        elseif event.key == "enter"
            anim.event_source.start()
        else
            keypressed(stimgen, event.key)
        end
    end
    # ------------------------------------------------------------------------ #

    if repeat
        plt.connect("key_press_event", key_press)
    end

    return ts
end
# ============================================================================ #

include("./RunSim.jl")

end # end module

ext: jl


text:
using QuantEcon
using Optim

# utility and marginal utility functions
u(x) = log(x)
du(x) = 1 / x

"""
Income fluctuation problem

##### Fields

- `r::Real` : Strictly positive interest rate
- `R::Real` : The interest rate plus 1 (strictly greater than 1)
- `bet::Real` : Discount rate in (0, 1)
- `b::Real` :  The borrowing constraint
- `Pi::Matrix{T} where T<:Real` : Transition matrix for `z`
- `z_vals::Vector{T} where T<:Real` : Levels of productivity
- `asset_grid::AbstractArray` : Grid of asset values

"""
struct ConsumerProblem{T <: Real}
    r::T
    R::T
    bet::T
    b::T
    Pi::Matrix{T}
    z_vals::Vector{T}
    asset_grid::AbstractArray
end

function ConsumerProblem(;r=0.01,
                         bet=0.96,
                         Pi=[0.6 0.4; 0.05 0.95],
                         z_vals=[0.5, 1.0],
                         b=0.0,
                         grid_max=16,
                         grid_size=50)
    R = 1 + r
    asset_grid = linspace(-b, grid_max, grid_size)

    ConsumerProblem(r, R, bet, b, Pi, z_vals, asset_grid)
end

"""
Apply the Bellman operator for a given model and initial value.

##### Arguments

- `cp::ConsumerProblem` : Instance of `ConsumerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output
- `;ret_policy::Bool(false)`: Toggles return of value or policy functions

##### Returns

None, `out` is updated in place. If `ret_policy == true` out is filled with the
policy function, otherwise the value function is stored in `out`.

"""
function bellman_operator!(cp::ConsumerProblem,
                          V::Matrix,
                          out::Matrix;
                          ret_policy::Bool=false)

    # simplify names, set up arrays
    R, Pi, bet, b = cp.R, cp.Pi, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    z_idx = 1:length(z_vals)

    # value function when the shock index is z_i
    vf = interp(asset_grid, V)

    opt_lb = 1e-8

    # solve for RHS of Bellman equation
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)

            function obj(c)
                EV = dot(vf.(R*a+z-c, z_idx), Pi[i_z, :]) # compute expectation
                return -u(c)  - bet * EV
            end
            res = optimize(obj, opt_lb, R.*a.+z.+b)
            c_star = Optim.minimizer(res)

            if ret_policy
                out[i_a, i_z] = c_star
            else
               out[i_a, i_z] = - Optim.minimum(res)
            end

        end
    end
    out
end

bellman_operator(cp::ConsumerProblem, V::Matrix; ret_policy=false) =
    bellman_operator!(cp, V, similar(V); ret_policy=ret_policy)

"""
Extract the greedy policy (policy function) of the model.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
get_greedy!(cp::ConsumerProblem, V::Matrix, out::Matrix) =
    update_bellman!(cp, V, out, ret_policy=true)

get_greedy(cp::ConsumerProblem, V::Matrix) =
    update_bellman(cp, V, ret_policy=true)

"""
The approximate Coleman operator.

Iteration with this operator corresponds to policy function
iteration. Computes and returns the updated consumption policy
c.  The array c is replaced with a function cf that implements
univariate linear interpolation over the asset grid for each
possible value of z.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `c::Matrix`: Current guess for the policy function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
function coleman_operator!(cp::ConsumerProblem, c::Matrix, out::Matrix)
    # simplify names, set up arrays
    R, Pi, bet, b = cp.R, cp.Pi, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    z_idx = 1:length(z_vals)
    gam = R * bet

    # policy function when the shock index is z_i
    cf = interp(asset_grid, c)

    # compute lower_bound for optimization
    opt_lb = 1e-8

    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)
            function h(t)
                cps = cf.(R*a+z-t, z_idx) # c' for each z'
                expectation = dot(du.(cps), Pi[i_z, :])
                return abs(du(t) - max(gam * expectation, du(R*a+z+b)))
            end
            opt_ub = R*a + z + b  # addresses issue #8 on github
            res = optimize(h, min(opt_lb, opt_ub - 1e-2), opt_ub,
                           method=Optim.Brent())
            out[i_a, i_z] = Optim.minimizer(res)
        end
    end
    out
end

"""
Apply the Coleman operator for a given model and initial value

See the specific methods of the mutating version of this function for more
details on arguments
"""
coleman_operator(cp::ConsumerProblem, c::Matrix) =
    coleman_operator!(cp, c, similar(c))

function initialize(cp::ConsumerProblem)
    # simplify names, set up arrays
    R, bet, b = cp.R, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    shape = length(asset_grid), length(z_vals)
    V, c = Array{Float64}(shape...), Array{Float64}(shape...)

    # Populate V and c
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)
            c_max = R*a + z + b
            c[i_a, i_z] = c_max
            V[i_a, i_z] = u(c_max) ./ (1 - bet)
        end
    end

    return V, c
end


text:
# Jl Plots

# - Plots
x = ["Nerds","Hackers","Scientists"]
y = [0.4,0.35,0.25]
Plots.pie(x, y, title="The Julia Community",l = 0.5)

# Pie - makie

using CairoMakie

f, ax, plt = CairoMakie.pie([π/2, 2π/3, π/4],
                normalize=false,
                offset = π/2,
                color = [:orange, :purple, :green],
                axis = (autolimitaspect = 1,),
                )

max_stars_repo_name: Ellon-M/visualizations

max_issues_repo_name: Ellon-M/visualizations

max_forks_repo_name: Ellon-M/visualizations


text:
# NestedNameSpecifier
function getPrefix(x::NestedNameSpecifier)
    @check_ptrs x
    return NestedNameSpecifier(clang_NestedNameSpecifier_getPrefix(x))
end

function containsErrors(x::NestedNameSpecifier)
    @check_ptrs x
    return clang_NestedNameSpecifier_containsErrors(x)
end

function dump(x::NestedNameSpecifier)
    @check_ptrs x
    return clang_NestedNameSpecifier_dump(x)
end


text:
type GaussQuadratureRule
    weights::Vector{Float64}
    points::Vector{Vector{Float64}}
end

weights(qr::GaussQuadratureRule) = qr.weights
points(qr::GaussQuadratureRule) = qr.points

function integrate(qr::GaussQuadratureRule, f::Function)
    I = 0.0
    for (w, x) in zip(weights(qr), points(qr))
        I += w * f(x)
    end
    return I
end

# Maybe we could stage the coming three function and create ndim loops.
function make_quadrule(ref::ReferenceLine, order::Int)
    p, w = gausslegendre(order)
    weights = Array(Float64, order)
    points = Array(Vector{Float64}, order)
    count = 1
    for i = 1:order
        points[count] = [p[i]]
        weights[count] = w[i]
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

function make_quadrule(ref::ReferenceSquare, order::Int)
    p, w = gausslegendre(order)
    weights = Array(Float64, order^2)
    points = Array(Vector{Float64}, order^2)
    count = 1
    for i = 1:order, j = 1:order
        points[count] = [p[i], p[j]]
        weights[count] = w[i] * w[j]
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

function make_quadrule(ref::ReferenceCube, order::Int)
    p, w = gausslegendre(order)
    weights = Array(Float64, order^3)
    points = Array(Vector{Float64}, order^3)
    count = 1
    for i = 1:order, j = 1:order, k = 1:order
        points[count] = [p[i], p[j], p[k]]
        weights[count] = w[i] * w[j] * w[k]
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

function make_quadrule(ref::ReferenceTriangle, order::Int)
    p1, w1 = gaussjacobi(order, 0.0, 0.0)
    p2, w2 = gaussjacobi(order, 1.0, 0.0)

    affine(a, b) = [(1+a) * (1-b) / 4, (1+b) / 2]

    SCALE = 1/8
    weights = Array(Float64, order^2)
    points = Array(Vector{Float64}, order^2)
    count = 1
    for i = 1:order, j = 1:order
        weights[count] = w1[i] * w2[j] * SCALE
        points[count] = affine(p1[i], p2[j])
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

function make_quadrule(ref::ReferenceTetrahedron, order::Int)
    p1, w1 = gaussjacobi(order, 0.0, 0.0)
    p2, w2 = gaussjacobi(order, 1.0, 0.0)
    p3, w3 = gaussjacobi(order, 2.0, 0.0)
    error("Bugged!")
    function affine(a, b, c)
       a = (a+1) / 2
       b = (b+1) / 2
       c = (c+1) / 2
       [a*b*c, a*b*(1-c), a*(1-b)]
     end

    SCALE = 1/48
    weights = Array(Float64, order^3)
    points = Array(Vector{Float64}, order^3)
    count = 1
    for i = 1:order, j = 1:order, k = 1:order
        weights[count] = w1[i] * w2[j] * w2[k] * SCALE
        points[count] = affine(p1[i], p2[j], p3[k])
        count += 1
    end
    GaussQuadratureRule(weights, points)
end

max_stars_repo_name: JuliaTagBot/JFET

max_issues_repo_name: JuliaTagBot/JFET

max_forks_repo_name: JuliaTagBot/JFET


text:
function myrouwenhorst(N::Integer,ρ::Real,σ::Real)

    """
    Consider the follwoing process
    z_{t+1} = ρz_t  + ϵ_{t+1}

    Arguments:
    N::Int64 :grid size
    ρ::Float64 :persistence of AR(1) process
    σ::Float64 :standard diviation of random components

    Return:
    z::array{float64,1} :discretized grid
    Π::array{float64,2} :transition matrix

    """

    #####################################################
    # step 2. create state grid
    #####################################################
    σ_z = σ/sqrt(1.0-ρ^2)
    ψ = sqrt(N-1)* σ_z # end points
    # create shock grid
    zgrid = collect(linspace(-ψ,ψ,N))

    #####################################################
    # step 3. compute trandition matrix
    #####################################################

    p = (1.0+ρ)/2.0
    Π   =  [p 1-p;1-p p]
    @inbounds for i in 3:N
        zero_v = zeros(i-1,1)
        zero_v_long = zeros(1,i)
        Π = p.*[Π zero_v; zero_v_long] +
             (1-p).*[zero_v Π; zero_v_long] +
             (1-p).*[zero_v_long; Π zero_v] +
             p.*[zero_v_long; zero_v Π]
    ######################################################
    # step 4. devide matrix by 2 except the top and bottom
    ######################################################
        @views Π[2:end-1,:] ./=  2.0
    end

    return zgrid, Π
end

max_stars_repo_name: tokupooh/Quant-Macro

max_issues_repo_name: tokupooh/Quant-Macro

max_forks_repo_name: tokupooh/Quant-Macro


text:
nonlinear_solve!(u, u_prev, δu, model::Model{DIM,2,D,V,E,P,S}, restart_flag) where {DIM,D,V,E,P,S<:AbstractLinearSolver} = @error "Please choose a non linear solver"

function nonlinear_solve!(u::Vector, u_prev::Vector, δu::Vector, model::Model{DIM,2,Nothing,TV,TE,TP,TS}, restart_flag) where {DIM,TV,TE,TP<:ViscousDruckerPrager,TS<:AbstractNonLinearSolver}

    # Unpack some model fields
    grid, dh, dbc, mp, states, res, K, clock, solver = model.grid, model.dofhandler, model.dirichlet_bc, model.material_properties, model.material_state, model.RHS, model.K, model.clock, model.solver

    println("Δt in nlsolve : ", clock.Δt)

    # number of base functions per element
    nbasefuncs = getnbasefunctions.(model.cellvalues_tuple)

    # respective fields global dofs
    dofs_field2 = get_field_dofs(dh.field_names[2],model)
    dofs_disp = get_field_dofs(:u,model)
    #dofs_disp = [i for i in eachindex(u) if i ∉ dofs_D] # first is faster

    # flags for restarting nonlinear iterations
    divergence_flag = false # flag when diverging residual norm
    max_iter_flag = false # flag when reaching max number of NL iterations

    # some initializations
    divergence_count = 0 # used to trigger reset of non linear iterations if residual norm diverges
    norm_prev = 0.0

    # apply boundary dirichlet boundary conditions to u :
    update!(dbc, clock.current_time) # evaluates the Dirichlet-bcs at time t
    apply!(u, dbc)  # set the prescribed dbcs values in the solution vector

    newton_itr = -1 # initialize non linear iteration count
    while true; newton_itr += 1
        @timeit "Newton iter" begin
            #@timeit "assemble"
            if newton_itr == 0
                tt = @elapsed doassemble_elast!(res, K, model, nbasefuncs, u, u_prev)
            else
                tt = @elapsed doassemble_AD!(res, K, model, nbasefuncs, u, u_prev)
                K_mat = Array(K)
                println("AD Jacobian L2-norm conditioning = ", cond(K_mat)) # SLOOOW
                println("AD Jacobian L2-norm u-u block conditioning = ", cond(K_mat[dofs_disp,dofs_disp]))
                println("AD Jacobian L2-norm P-P block conditioning = ", cond(K_mat[dofs_field2,dofs_field2]))
            end
            println("assemble time : ", tt)
            # compute residual norm
            norm_res = norm(res[JuAFEM.free_dofs(dbc)])

            # print current Newton iteration
            print("\nIteration: $newton_itr \tresidual: $(@sprintf("%.8f", norm_res))\n\n")

            #### EXIT CHECKS ####
            if (norm_res < solver.atol)
                break
            elseif (solver.max_iter_atol > 0) & (newton_itr == solver.max_iter_number) & (norm_res < solver.max_iter_atol)
                (norm_res >= solver.atol) && @warn("last newton iteration, accepted residual is ",norm_res)
                break # if
            end

            # if too many successive divergence, reset the nonlinear iterations from last time iteration solution u_converged using a lower timestep
            if newton_itr == 0
                norm_prev = norm_res
            else
                if norm_res >= norm_prev
                    divergence_count += 1
                else
                    divergence_count = 0
                end
                norm_prev = norm_res
            end

            divergence_flag = (divergence_count == solver.max_div_iter)
            max_iter_flag = (newton_itr == solver.max_iter_number)

            if divergence_flag | max_iter_flag
                divergence_flag && (@info "$(solver.max_div_iter) successive non linear iterations were diverging. Restart with decreased timestep")
                max_iter_flag   && (@info "reached maximum number of non linear iterations : $(solver.max_iter_number). Restart with decreased timestep")
                restart_flag = true
                return restart_flag
            end

            @timeit "apply_dbc" apply_zero!(K, res, dbc)
            @timeit "linear_solve" δu .= solver.linear_solver(K,-res,model)

            # P = Diagonal(K)
            # invP = sparse(inv(Array(K)))
            # println("Preconditioned AD Jacobian L2-norm conditioning = ", cond(Array(invP*K)))
            # @timeit "linear_solve" δu .= solver.linear_solver(invP*K,Array(-invP*res),model)

            # solution correction
            α = 1.0 # correction factor
            if newton_itr == 0
                u .+= δu
            else
                u .+= α.*δu
                (α != 1) && println("α = $alpha")
            end
        end
        ###### TEST
        # vtk_grid("TEST_u-logD_iter$(newton_itr)", model.dofhandler) do vtkfile
        #     u2 = copy(u)
        #     u2[dofs_field2] .= exp.(u2[dofs_field2]) # coming back to damage from log(damage)
        #     vtk_point_data(vtkfile, model.dofhandler, u2)
        # end
        ######
    end
end

function nonlinear_solve!(u::Vector, u_prev::Vector, δu::Vector, model::Model{DIM,2,TD,TV,TE,TP,TS}, restart_flag) where {DIM,TD,TV,TE,TP,TS<:AbstractNonLinearSolver}

    # Unpack some model fields
    grid, dh, dbc, mp, states, res, K, clock, solver = model.grid, model.dofhandler, model.dirichlet_bc, model.material_properties, model.material_state, model.RHS, model.K, model.clock, model.solver

    if clock.iter in (2,3)
        # println("ue cell 1 début nonlinearsolve : ")
        # display(u[1:22])
    end
    # number of base functions per element
    nbasefuncs = getnbasefunctions.(model.cellvalues_tuple)

    # respective fields global dofs
    dofs_field2 = get_field_dofs(dh.field_names[2],model)
    dofs_disp = get_field_dofs(:u,model)
    #dofs_disp = [i for i in eachindex(u) if i ∉ dofs_D] # first is faster

    # flags for restarting nonlinear iterations
    divergence_flag = false # flag when diverging residual norm
    max_iter_flag = false # flag when reaching max number of NL iterations

    # some initializations
    divergence_count = 0 # used to trigger reset of non linear iterations if residual norm diverges
    norm_prev = 0.0

    # apply boundary dirichlet boundary conditions to u :
    update!(dbc, clock.current_time) # evaluates the Dirichlet-bcs at time t
    apply!(u, dbc)  # set the prescribed dbcs values in the solution vector

    if clock.iter in (2,3)
        # println("ue cell 1 nonlinearsolve apres apply dbc on u : ")
        # display(u[1:22])
    end

    newton_itr = -1 # initialize non linear iteration count
    while true; newton_itr += 1
        @timeit "Newton iter" begin

            # print current Newton iteration
            print("\nNEWTON ITERATION: $newton_itr\n\n")

            #@timeit "assemble"
            if (newton_itr == 0) | (restart_flag == false)
                tt = @elapsed doassemble_elast!(res, K, model, nbasefuncs, u, u_prev)
            elseif restart_flag == :activate_damage
                #tt = @elapsed doassemble_elast!(res, K, model, nbasefuncs, u, u_prev) #TODO put back
                tt = @elapsed doassemble_AD!(res, K, model, nbasefuncs, u, u_prev)
                #K_mat = Array(K)
                #println("AD Jacobian L2-norm conditioning = ", cond(K_mat)) # SLOOOW
                #println("AD Jacobian L2-norm u-u block conditioning = ", cond(K_mat[dofs_disp,dofs_disp]))
                #println("AD Jacobian L2-norm D-D block conditioning = ", cond(K_mat[dofs_field2,dofs_field2]))
            end
            println("assemble time : ", tt)

            #### Max D TEST :
            # if model.material_properties[1].damage isa Damage
            #     logD = u[dofs_field2]
            #     logD_prev = u_prev[dofs_field2]
            #     max_δD = maximum(exp.(logD).-exp.(logD_prev))
            #     @assert all(logD.-logD_prev .>= 0)
            #     print("\t max ΔD = ", max_δD, "\n")
            # end
            ####

            ### SOLVE ###

            if (newton_itr == 0) | (restart_flag == false)
                println("apply_zero and solve")
                @timeit "apply_dbc" apply_zero!(K, res, dbc)
                @timeit "linear_solve" δu .= solver.linear_solver(K,-res,model)

                ##### TEST #####
                println("δD elast extrema (should be zero) = ", extrema(exp.(δu[dofs_field2])))
                ################
            elseif restart_flag == :activate_damage

                ### test Δt wrt to previous stress solution
                # u_exp = copy(u)
                # u_exp[dofs_field2] .= exp.(u_exp[dofs_field2])
                # Δt_max_damage = get_damage_constrained_Δt(model,u_exp,0.3)
                # println("Δt max damage with elastic stress = ",Δt_max_damage)

                # ################################################
                # ### Compute Jacobian using finite difference ###
                # ################################################
                # f!(res,u) = doassemble_res!(res, model, nbasefuncs, u, u_prev)
                # #f_trailing_args!(res,u,model,nbasefuncs,u_prev) = doassemble_res!(res, model, nbasefuncs, u, u_prev)
                # f(u) = (res = similar(u) ; doassemble_res!(res, model, nbasefuncs, u, u_prev) ; return res)
                # print(" hand made Jac time : ")
                # @time J = compute_jacobian(f!, u ; eps = 1e-6)
                # # dofs_tuple = (dofs_disp, dofs_D),
                # println("hand made Jacobian L2-norm conditioning = ", cond(Array(J)))
                #
                # #uu_jac =
                # #DD_jac =
                # # print(" FiniteDiff Jac time : ")
                # # @time begin
                # #     # @eval using FiniteDifferences #, FiniteDiff, SparsityDetection, SparseDiffTools
                # #     # input = rand(length(u))
                # #     # output = similar(input)
                # #     # sparsity_pattern = jacobian_sparsity(f!,output,input)
                # #     # sparsejac = Float64.(sparse(sparsity_pattern))
                # #     # colors = matrix_colors(sparsejac)
                # #     # J2 = FiniteDiff.JacobianCache(u,colorvec=colors,sparsity=sparsejac)
                # #     # J2 = FiniteDiff.finite_difference_jacobian(f,u)
                # #     # fdm = FiniteDifferences.central_fdm(2,1)
                # #     # J2 = FiniteDifferences.jacobian(fdm, f, u)[1]
                # #
                # #     # println("FiniteDiff Jacobian L2-norm conditioning = ", cond(J2))
                # #     # println("max difference 2 jacs = ", maximum(abs.(J.-J2)))
                # #     # J = sparse(J2)
                # # end
                #
                #
                # ### Apply zeros accordingly to dbcs ###
                # @timeit "apply_dbc" apply_zero!(J, res, dbc) # TODO : maybe after preconditioning
                #
                # println("presence of nans : ", any(isnan.(J)))
                # println("maximum((J'.-J)./J) : ", maximum([J[i,j].-J[j,i]./J[i,j] for i in length(res), j in length(res) if !isnan(J[i,j])]))
                # println("positive definiteness of sym(J) : ", isposdef(Symmetric(J)))
                # #@timeit "linear_solve" δu .= Symmetric(K) \ -res
                #
                # ## Preconditioning :
                # # using Preconditioners
                # # P = DiagonalPreconditioner(J) # similar to the following
                # P = Diagonal(J)
                # invP = inv(P)
                # println("Preconditioned Jacobian L2-norm conditioning = ", cond(Array(invP*J)))
                #@timeit "linear_solve" δu .= solver.linear_solver(invP*J,Array(-invP*res),model)

                # ################################################
                # ################################################

                @timeit "apply_dbc" apply_zero!(K, res, dbc)
                @timeit "linear_solve" δu .= solver.linear_solver(K,-res,model)
                # P = Diagonal(K)
                # invP = sparse(inv(Array(P)))
                #@timeit "linear_solve" δu .= solver.linear_solver(invP*K,Array(-invP*res),model)
                #println("Preconditioned AD Jacobian L2-norm conditioning = ", cond(Array(invP*K)))

                ##### TEST #####
                println("δD damaged extrema damaged iter = ", extrema(exp.(δu[dofs_field2])))
                # TODO try to set negative δD to zero, remove it after test
                # for i in dofs_D
                #     if δu[i] < 0
                #         δu[i] = 0.0
                #     end
                # end
                # println("δD damaged extrema corrected = ", extrema(δu[dofs_D]))
                ################
            end
            # @timeit "linear_solve" begin
            #     ## TODO better
            #     Pl = Preconditioners.AMGPreconditioner(Symmetric(K))
            #     IterativeSolvers.cg!(δu, Symmetric(K), res, Pl = Pl, solver.linear_solver.kwargs...)
            # end

            # displacement and damage corrections
            if (newton_itr == 0) | (restart_flag == false)
                u[dofs_disp] .+= δu[dofs_disp]
                #u .+= δu
            elseif restart_flag == :activate_damage
                D = exp.(u[dofs_field2])
                δD = exp.(δu[dofs_field2])

                max_D_increment = 0.02
                α = minimum(min.(Ref(1.0),max_D_increment./δD, 0.99 .- D)) # ensure that D increment is low enough
                u[dofs_disp] .+= α.*δu[dofs_disp]
                u[dofs_field2] .= log.(D .+ α.*δD) # damage increment is incremented at converged iteration only
                println("α = ", α)
                println("min/max δD newton_iter_$(newton_itr) : ", extrema(δD))
                println("min/max effective α*δD newton_iter_$(newton_itr) : ", extrema(α.*δD))
            end
            # Make sure that the elastic solve didn't affect Damage
            # if newton_itr == 0
            #     dofs_D = get_field_dofs(:D,model)
            #     u[dofs_D] .= u_prev[dofs_D]
            # end

            # compute residual norm
            free_dofs = JuAFEM.free_dofs(dbc)
            free_disp_dofs = free_dofs[in.(free_dofs,Ref(dofs_disp))]
            if (newton_itr == 0) | (restart_flag == false) # norm res on displacement
                norm_res = norm(res[free_disp_dofs])
            else # norm res on displacement and damage
                norm_res = norm(res[free_dofs])
                norm_res_disp = norm(res[free_disp_dofs])
                norm_res_damage = norm(res[dofs_field2])
            end

            print("\ttotal residual: $(@sprintf("%.8f", norm_res))\n")
            if !(newton_itr == 0) & !(restart_flag == false)
                print("\tdisplacement residual: $(@sprintf("%.8f", norm_res_disp))\n")
                print("\tdamage residual: $(@sprintf("%.8f", norm_res_damage))\n\n")
            end
            #### EXIT CHECKS ####

            if (norm_res < solver.atol)
                if restart_flag == :activate_damage
                    restart_flag = false

                    # increment damage
                    # D = exp.(u[dofs_field2])
                    # δD = exp.(δu[dofs_field2])
                    # u[dofs_field2] .= log.(min.(D .+ δD, 0.99))

                    println("convergence with damage growth, restart_flag set to false, should allow time stepping")
                    break
                elseif restart_flag == false
                    #restart_flag = :activate_damage #good line # TODO put it back
                    # uncomment previous line to enable damage growth
                    println("convergence without damage growth, restart_flag set to :activate_damage, should restart time iteration")
                    break
                end
            elseif (solver.max_iter_atol > 0) & (newton_itr == solver.max_iter_number) & (norm_res < solver.max_iter_atol)
                if restart_flag == :activate_damage
                    (norm_res >= solver.atol) && @warn("last newton iteration, accepted residual is ",norm_res)
                    break
                end
            end

            # if too many successive divergence, reset the nonlinear iterations from last time iteration solution u_converged using a lower timestep
            if newton_itr == 0
                norm_prev = norm_res
            else
                if norm_res >= norm_prev
                    divergence_count += 1
                else
                    divergence_count = 0
                end
                norm_prev = norm_res
            end

            divergence_flag = (divergence_count == solver.max_div_iter)
            max_iter_flag = (newton_itr == solver.max_iter_number)

            if divergence_flag | max_iter_flag
                divergence_flag && (@info "$(solver.max_div_iter) successive non linear iterations were diverging. Restart with decreased timestep")
                max_iter_flag   && (@info "reached maximum number of non linear iterations : $(solver.max_iter_number). Restart with decreased timestep")
                restart_flag = true
                return restart_flag
            end

        end
        ###### TEST
        # vtk_grid("TEST/TEST_u-logD_iter$(newton_itr)_timeit$(clock.iter)", model.dofhandler) do vtkfile
        #     #u2 = copy(u)
        #     #u2[dofs_field2] .= exp.(u2[dofs_field2]) # coming back to damage from log(damage)
        #     vtk_point_data(vtkfile, model.dofhandler, u)
        # end
        # println("vtk test saved at ")
        ######
    end
    return restart_flag
end

function nonlinear_solve!(u::Vector, uprev::Vector, δu::Vector, model::Model{DIM,1,TD,TV,TE,TP,TS}, restart_flag) where {DIM,TD,TV,TE,TP,TS<:AbstractNonLinearSolver}

    # Unpack some model fields
    grid, dh, dbc, mp, states, K, res, clock, solver = model.grid, model.dofhandler, model.dirichlet_bc, model.material_properties, model.material_state, model.K, model.RHS, model.clock, model.solver

    println("Δt in nlsolve : ", clock.Δt)

    # number of base functions per element
    nbasefuncs = getnbasefunctions(model.cellvalues_tuple[1])

    # flags for restarting nonlinear iterations
    divergence_flag = false # flag when diverging residual norm
    max_iter_flag = false # flag when reaching max number of NL iterations

    # some initializations
    divergence_count = 0 # used to trigger reset of non linear iterations if residual norm diverges
    norm_prev = 0.0

    # apply boundary dirichlet boundary conditions to u :
    update!(dbc, clock.current_time) # evaluates the D-bndc at time t
    apply!(u, dbc)  # set the prescribed values in the solution vector

    newton_itr = -1 # initialize non linear iteration count
    while true; newton_itr += 1
        @timeit "Newton iter" begin
            #@timeit "assemble"
            tt = @elapsed doassemble!(model::Model{DIM,1,TD,TV,TE,TP}, nbasefuncs, u ; noplast = (newton_itr == 0))
            println("assemble time : ", tt)
            # compute residual norm
            norm_res = norm(res[JuAFEM.free_dofs(dbc)])

            # print current Newton iteration :
            if newton_itr > 0
                print("\nIteration: $newton_itr \t residual: $(@sprintf("%.8f", norm_res))\n")
            else
                print("First elastic iteration: $newton_itr \t residual: $(@sprintf("%.8f", norm_res))\n")
            end

            #### Max D TEST :
            if TD <: Damage
                max_ΔD = 0
                for cell_states in states
                    cell_temp_D = mean([state.temp_D for state in cell_states])
                    cell_D = mean([state.D for state in cell_states])
                    max_ΔD = max(max_ΔD,cell_temp_D - cell_D)
                end
                print("\t max ΔD = ", max_ΔD, "\n")
            end
            ####

            #### EXIT CHECKS ####
            if (norm_res < solver.atol)
                break
            elseif (solver.max_iter_atol > 0) & (newton_itr == solver.max_iter_number) & (norm_res < solver.max_iter_atol)
                (norm_res >= solver.atol) && @warn("last newton iteration, accepted residual is ",norm_res)
                break # if
            end

            # if too many successive divergence, reset the nonlinear iterations from last time iteration solution u_converged using a lower timestep
            if newton_itr == 0
                norm_prev = norm_res
            else
                if norm_res >= norm_prev
                    divergence_count += 1
                else
                    divergence_count = 0
                end
                norm_prev = norm_res
            end

            divergence_flag = (divergence_count == solver.max_div_iter)
            max_iter_flag = (newton_itr == solver.max_iter_number)

            if divergence_flag | max_iter_flag
                divergence_flag && (@info "$(solver.max_div_iter) successive non linear iterations were diverging. Restart with decreased timestep")
                max_iter_flag   && (@info "reached maximum number of non linear iterations : $(solver.max_iter_number). Restart with decreased timestep")
                restart_flag = true
                return restart_flag
            end

            ### Linear Solve for δu ###
            @timeit "apply_dbc" apply_zero!(K, res, dbc)
            #@timeit "linear_solve" δu .= Symmetric(K) \ -res
            @timeit "linear_solve" δu .= solver.linear_solver(K,-res,model)
            # @timeit "linear_solve" begin
            #     ## TODO better
            #     Pl = Preconditioners.AMGPreconditioner(Symmetric(K))
            #     IterativeSolvers.cg!(δu, Symmetric(K), res, Pl = Pl, solver.linear_solver.kwargs...)
            # end

            # displacement correction

            α = 1.0 # correction factor
            if newton_itr == 0
                u .+= δu
            else
                u .+= α.*δu
                (α != 1) && println("α = $alpha")
            end

        end
    end
    return restart_flag
end

# function output_values(model::Model{dim,1,D,V,E,P}) where {dim,D,V,E,P<:Plasticity}
#     grid, dh, dbc, mp, states, K, res, solver = model.grid, model.dofhandler, model.dirichlet_bc, model.material_properties, model.material_state, model.K, model.RHS, model.solver
#
#     τ_values = zeros(getncells(grid))
#     F_values = zeros(getncells(grid))
#     p_values = zeros(getncells(grid))
#     σxx_values = zeros(getncells(grid))
#     σyy_values = zeros(getncells(grid))
#     σzz_values = zeros(getncells(grid))
#     ϵxx_values = zeros(getncells(grid))
#     ϵyy_values = zeros(getncells(grid))
#     ϵzz_values = zeros(getncells(grid))
#     ϵkk_values = zeros(getncells(grid))
#     γ_values = zeros(getncells(grid))
#     norm_ep_values = zeros(getncells(grid))
#     I2D = SymmetricTensor{2,3}(δ) # second order identity tensor
#     for (el, cell_states) in enumerate(states)
#         for state in cell_states
#             #el == 1 && show(state.σ)
#             p = -vol(state.temp_σ)[1]
#             s = state.temp_σ + p*I2D
#             p_values[el] += p
#             ϕ = mp[el].plasticity.ϕ
#             C = mp[el].plasticity.C
#             τ_yield = sind(ϕ)*p + cosd(ϕ)*C
#             F_values[el] += get_τ(s,mp[el].plasticity) - τ_yield
#             τ_values[el] += get_τ(s,mp[el].plasticity)
#             norm_ep_values[el] += norm(state.temp_ϵᵖ)
#             σxx_values[el] += state.temp_σ[1,1]
#             σyy_values[el] += state.temp_σ[2,2]
#             σzz_values[el] += state.temp_σ[3,3]
#             ϵxx_values[el] += state.temp_ϵ[1,1]
#             ϵyy_values[el] += state.temp_ϵ[2,2]
#             ϵzz_values[el] += state.temp_ϵ[3,3]
#             ϵkk_values[el] += tr(state.temp_ϵ)
#             dev_ϵ_el = dev(state.temp_ϵ)
#             γ_values[el] += sqrt(2*dev_ϵ_el ⊡ dev_ϵ_el)
#         end
#         p_values[el] /= length(cell_states)
#         τ_values[el] /= length(cell_states)
#         norm_ep_values[el] /= length(cell_states)
#         F_values[el] /= length(cell_states)
#         σxx_values[el] /= length(cell_states)
#         σyy_values[el] /= length(cell_states)
#         σzz_values[el] /= length(cell_states)
#         ϵxx_values[el] /= length(cell_states)
#         ϵyy_values[el] /= length(cell_states)
#         ϵzz_values[el] /= length(cell_states)
#         ϵkk_values[el] /= length(cell_states)
#         γ_values[el] /= length(cell_states)
#     end
#     return (p = p_values, τ = τ_values, norm_ep = norm_ep_values, σxx = σxx_values, σyy = σyy_values, σzz = σzz_values, ϵxx = ϵxx_values, ϵyy = ϵyy_values, ϵzz = ϵzz_values, ϵkk = ϵkk_values, γ = γ_values)
# end
#
# function output(vals, u, dh, dbc, clock)
#     vtk_grid("Rheology_$(clock.iter)", dh) do vtkfile
#          vtk_point_data(vtkfile, dh, u) # displacement field
#          vtk_point_data(vtkfile, dbc) # dirichlet boundary conditions
#          vtk_cell_data(vtkfile, vals.p, "pressure")
#          vtk_cell_data(vtkfile, vals.τ, "tau")
#          vtk_cell_data(vtkfile, vals.norm_ep, "norm_ep")
#          vtk_cell_data(vtkfile, vals.σxx, "sigma xx")
#          vtk_cell_data(vtkfile, vals.σyy, "sigma yy")
#          vtk_cell_data(vtkfile, vals.σzz, "sigma zz")
#          vtk_cell_data(vtkfile, vals.ϵxx, "epsilon xx")
#          vtk_cell_data(vtkfile, vals.ϵyy, "epsilon yy")
#          vtk_cell_data(vtkfile, vals.ϵzz, "epsilon zz")
#          vtk_cell_data(vtkfile, vals.ϵkk, "volumetric strain")
#          vtk_cell_data(vtkfile, vals.γ, "dev strain invariant")
#          println("converged newton_itr saved")
#     end
# end


text:
dark_palette = [
    colorant"#FE4365", # red
    colorant"#eca25c", # orange
    colorant"#3f9778", # green
    colorant"#005D7F" # blue
]

dark_bg = colorant"#191919"

_themes[:custom] = PlotTheme(
    legend=false,
    bg=dark_bg,
    bginside=dark_bg,
    bglegend=dark_bg,
    fg=colorant"#ADB2B7",
    fgtext=colorant"#FFFFFF",
    fgguide=colorant"#FFFFFF",
    fglegend=colorant"#FFFFFF",
    legendfontcolor=colorant"#FFFFFF",
    legendtitlefontcolor=colorant"#FFFFFF",
    titlefontcolor=colorant"#FFFFFF",
    palette=expand_palette(dark_bg, dark_palette; lchoices=[57], cchoices=[100]),
    colorgradient=:fire
)


text:
using CryptoGroups: Group, PGroup, Hash

Base.write(io::IO, tree::Tree) = write(io, encode(tree))
Base.write(f::AbstractString, tree::Tree) = write(f, encode(tree))

Base.string(tree::Tree) = bytes2hex(encode(tree))

function unmarshal_full_public_key(g::Group, tree::Tree)

    G = typeof(g)

    g′, y = convert(Tuple{G, G}, tree)

    @assert g′ == g

    return y
end

marshal_full_public_key(g::G, y::G) where G <: Group = Tree((g, y))

function unmarshal_publickey(tree::Tree)

    g = unmarshal(BigInt, tree.x[1])
    c1, c2 = convert(Tuple{BigInt, BigInt}, tree.x[2])

    @assert value(g) == c1

    G = typeof(g)
    y = convert(G, c2)

    return y, g
end

function marshal_publickey(g::G, y::G) where G <: PGroup

    group_spec = marshal(g)

    p = modulus(g)

    ### It may be actaully be forcing to use the same bytelength as prime order
    gleaf = Leaf(g, L = bitlength(p))
    yleaf = Leaf(y, L = bitlength(p))

    public_key = Tree((gleaf, yleaf))
    tree = Tree((group_spec, public_key))

    return tree
end

function marshal_privatekey(g::PGroup, s::BigInt) #where G <: PGroup
    group_spec = marshal(g)

    q = order(g)

    @assert s < q "Secret key must be with in the order of the group"

    sleaf = Leaf(s, bytelength(q))

    tree = Tree((group_spec, sleaf))

    return tree
end

function unmarshal_privatekey(tree::Tree)

    g = unmarshal(BigInt, tree.x[1])
    s = convert(BigInt, tree.x[2])

    return (s, g) ### The group can often be omited when not needed.
end

function map_hash_name(x::AbstractString)
    if x == "SHA-256"
        return "sha256"
    elseif x == "SHA-384"
        return "sha384"
    elseif x == "SHA-512"
        return "sha512"
    else
        error("No corepsonding mapping for $x implemented")
    end
end

function map_hash_name_back(x::AbstractString)
    if x == "sha256"
        return "SHA-256"
    elseif x == "sha384"
        return "SHA-384"
    elseif x == "sha512"
        return "SHA-512"
    else
        error("No corepsonding mapping for $x implemented")
    end
end

map_hash_name_back(x::Hash) = map_hash_name_back(x.spec)

function ro_prefix(protinfo::AbstractDict; auxsid="default")

    version = protinfo["version"]
    sid = protinfo["sid"]

    s_H = protinfo["rohash"]
    s_PRG = protinfo["prg"]
    s_Gq = protinfo["pgroup"]

    nr = parse(Int32, protinfo["statdist"])
    nv = parse(Int32, protinfo["vbitlenro"])
    ne = parse(Int32, protinfo["ebitlenro"])

    data = (version, sid * "." * auxsid, nr, nv, ne, s_PRG, s_Gq, s_H)

    tree = Tree(data)
    binary = encode(tree)

    rohash = Hash(map_hash_name(protinfo["rohash"]))

    ρ = rohash(binary) ### Which hash function shall be used here?

    return ρ
end


text:
using Distributed

"""
Intended to be used with macros, `macroArguments` extracts both positional
and keyword arguments. The return value is a tuple containing 3 slots:
  * a function or `nothing` if no function argument provided
  * an array of positional arguments
  * a dictionary of keyword arguments. The keys of the dictionary are `Symbol`s
"""
function macroArguments(arguments)::Tuple{Union{Expr,Nothing},Vector{Any},Dict{Symbol,Any}}
  functional = nothing
  positional = []
  keywords = Dict{}()
  for i in eachindex(arguments)
    argument = arguments[i]
    if i == 1 && argument isa Expr && argument.head === :(->)
      functional = argument
    elseif argument isa Expr && argument.head === :(kw)
      kw, arg = argument.args
      keywords[kw] = arg
    else
      push!(positional, argument)
    end
  end
  (functional, positional, keywords)
end

macro anywhere(expr)
  :(@async fetch(@spawnat :any $(expr)) )
end


text:
using Pipe

const data = parse.(Int, readlines())

# Part One
@pipe data |> diff |> count(>(0), _) |> println

# Part Two
@pipe data[begin+3:end] - data[begin:end-3] |> count(>(0), _) |> println

max_stars_repo_name: pLOPeGG/JuliaOfCode

max_issues_repo_name: pLOPeGG/JuliaOfCode

max_forks_repo_name: pLOPeGG/JuliaOfCode


text:
using Clines, SimpleDrawing, Plots

function originals()
    X = Line(-1, 7)

    C1 = Circle(1, 1, 1)
    C2 = Circle(2.5, 1, 0.5)
    C3 = Circle(4, 1, 1)

    pts = [1, 2 + im, 3 + im]
    SS = Cline(pts...)

    return [X, C1, C2, C3, SS]
end

function fourway(F::LFT)
    x = originals()
    y = F.(x)
    newdraw()
    draw.(y[1:4])
    try
        draw(y[5], linestyle = :dash)
    catch
        draw(y[5])
    end
    finish()
end

function fourway(z::Number = 2 + 4im)
    F = LFT(0, 1, 1, -z)
    fourway(F)
end


text:
function reservevariables(m::JuMP.AbstractModel, devices::Array{NamedTuple{(:device, :formulation), Tuple{R,DataType}}}, time_periods::Int64) where {R <: PowerSystems.PowerSystemDevice}

    on_set = [d.device.name for d in devices]

    t = 1:time_periods

    p_rsv = @variable(m, p_rsv[on_set,t] >= 0)

    return p_rsv

end

# headroom constraints
function make_pmax_rsv_constraint(m::JuMP.AbstractModel,t::Int64, device::G, formulation::Type{D}) where {G<:PowerSystems.ThermalGen, D <: AbstractThermalDispatchForm}
    return @constraint(m, m[:p_th][device.name,t] + m[:p_rsv][device.name,t]  <= device.tech.activepowerlimits.max)
end

function make_pmax_rsv_constraint(m::JuMP.AbstractModel,t::Int64, device::G, formulation::Type{D}) where {G<:PowerSystems.ThermalGen, D <: AbstractThermalCommitmentForm}
    return @constraint(m, m[:p_th][device.name,t] + m[:p_rsv][device.name,t] <= device.tech.activepowerlimits.max * m[:on_th][device.name,t])
end

function make_pmax_rsv_constraint(m::JuMP.AbstractModel,t::Int64, device::G, formulation::Type{D}) where {G<:PowerSystems.RenewableGen, D <: AbstractRenewableDispatchForm}
    return @constraint(m, m[:p_re][device.name,t] + m[:p_rsv][device.name,t] <= device.tech.installedcapacity * values(device.scalingfactor)[t])
end

function make_pmax_rsv_constraint(m::JuMP.AbstractModel,t::Int64, device::G, formulation::Type{D}) where {G<:PowerSystems.InterruptibleLoad, D <: AbstractControllableLoadForm}
    return @constraint(m, m[:p_cl][device.name,t] + m[:p_rsv][device.name,t] <= device.maxactivepower * values(device.scalingfactor)[t])
end

# ramp constraints
function make_pramp_rsv_constraint(m::JuMP.AbstractModel,t::Int64, device::G, formulation::Type{D}, timeframe) where {G<:PowerSystems.ThermalGen, D <: AbstractThermalFormulation}
    rmax = device.tech.ramplimits != nothing  ? device.tech.ramplimits.up : device.tech.activepowerlimits.max
    return @constraint(m, m[:p_rsv][device.name,t] <= rmax/60 * timeframe)
end

function make_pramp_rsv_constraint(m::JuMP.AbstractModel,t::Int64, device::G, formulation::Type{D}, timeframe) where {G<:PowerSystems.RenewableGen, D <: AbstractRenewableDispatchForm}
    return
end
function make_pramp_rsv_constraint(m::JuMP.AbstractModel,t::Int64, device::G, formulation::Type{D}, timeframe) where {G<:PowerSystems.InterruptibleLoad, D <: AbstractControllableLoadForm}
    #rmax =  device.maxactivepower * values(device.scalingfactor)[t] #nominally setting load ramp limit to full range within 1 min
    #return @constraint(m, m[:p_rsv][device.name,t] <= rmax/60 * timeframe)
    return
end

function reserves(m::JuMP.AbstractModel, devices::Array{NamedTuple{(:device, :formulation), Tuple{R,DataType}}}, service::PowerSystems.StaticReserve, time_periods::Int64) where {R <: PowerSystems.PowerSystemDevice}

    p_rsv = m[:p_rsv]
    time_index = m[:p_rsv].axes[2]
    name_index = m[:p_rsv].axes[1]

    (length(time_index) != time_periods) ? @error("Length of time dimension inconsistent") : true

    pmin_rsv = JuMP.JuMPArray(Array{ConstraintRef}(undef,length(time_index)), time_index) #minimum system reserve provision
    pmax_rsv = JuMP.JuMPArray(Array{ConstraintRef}(undef, length.(JuMP.axes(p_rsv))), name_index, time_index) #maximum generator reserve provision

    for t in time_index
        pmin_rsv[t] = @constraint(m, sum([p_rsv[name,t] for name in name_index]) >= service.requirement)

        for (ix, name) in enumerate(name_index)
            if name == devices[ix].device.name
                pmax_rsv[name,t] = make_pmax_rsv_constraint(m, t, devices[ix].device,devices[ix].formulation)
            else
                @error "Gen name in Array and variable do not match"
            end
        end

    end

    rmp_devices = [d for d in devices if d.formulation<:PowerSimulations.ThermalDispatch]
    rmp_name_index = [d.device.name for d in rmp_devices]

    pramp_rsv = JuMP.JuMPArray(Array{ConstraintRef}(undef, (length(rmp_name_index),length(time_index))), rmp_name_index, time_index) #maximum generator reserve provision

    for t in time_index
        # TODO: check the units of ramplimits
        for (ix, name) in enumerate(rmp_name_index)
            if name == rmp_devices[ix].device.name
                pramp_rsv[name,t] = make_pramp_rsv_constraint(m, t, rmp_devices[ix].device, rmp_devices[ix].formulation, service.timeframe)
            else
                @error "Gen name in Array and variable do not match"
            end
        end

    end

    JuMP.register_object(m, :RsvProvisionMin, pmin_rsv)
    JuMP.register_object(m, :RsvProvisionMax, pmax_rsv)
    JuMP.register_object(m, :RsvProvisionRamp, pramp_rsv)

    return m

end


text:
struct InitialInvestments{R,G}
    options::Matrix{Int} # Existing (buildable) options (r x g)
    builds::Matrix{Int}  # Existing (dispatchable) units

    function InitialInvestments{}(options::Matrix, builds::Matrix)
        R, G = size(options)
        @assert (R, G) == size(builds)
        new{R,G}(options, builds)
    end

end

struct InitialConditions{R,G1,G2,G3}
    thermal::InitialInvestments{R,G1}
    variable::InitialInvestments{R,G2}
    storage::InitialInvestments{R,G3}
end

mutable struct ResourceInvestments{R,G}

    # Unit state machine:
    # --> vesting --> building --> retired
    #            \      ^     \       ^
    #             \     |      \      |
    #              > holding    > dispatching

    # Parameters

    # Capital costs are considered sunk and so modeled as one-time expenses,
    # although in terms of cash flow they may actually be amortized over time
    optioncost::Matrix{Float64}     # ($/unit, r x g)
    buildcost::Matrix{Float64}
    retirementcost::Matrix{Float64}

    optionleadtime::Matrix{Int}     # investment periods, r x g
    buildleadtime::Matrix{Int}

    newoptionslimit::Matrix{Int} # investment periods, r x g
    newbuildslimit::Matrix{Int}
    newforcedretirements::Matrix{Int}

    # Variables

    newoptions::Matrix{VariableRef} # New options purchased (r x g)
    newholding::Matrix{VariableRef} # Hold a newly-vested option
    newbuilds_vesting::Matrix{VariableRef}  # Excercise a newly-vested option
    newbuilds_holding::Matrix{VariableRef} # Start construction on held project
    newdispatching::Matrix{VariableRef} # Start dispatching newly-constructed units
    newretirements_dispatching::Matrix{VariableRef}  # Retire a dispatchable unit
    newretirements_building::Matrix{VariableRef}  # Retire a newly-constructed unit

    # Expressions

    optionsvested::Matrix{ExpressionRef} # Newly-vested options (r x g)
    buildsstarted::Matrix{ExpressionRef}  # Newly-started construction
    buildsfinished::Matrix{ExpressionRef}  # Newly-completed construction
    newretirements::Matrix{ExpressionRef}  # New retirements

    vesting::Matrix{ExpressionRef} # Options vesting (r x g)
    holding::Matrix{ExpressionRef} # Units buildable (option vested but not used)
    building::Matrix{ExpressionRef} # Units under construction (option used)
    dispatching::Matrix{ExpressionRef} # Units dispatchable (finished construction)
    retired::Matrix{ExpressionRef} # Units retired

    totalforcedretirements::Matrix{ExpressionRef} # Minimum cumulative units retired

    ucap::Matrix{ExpressionRef} # Unforced dispatchable capacity (MW, r x g)

    investmentcosts::Matrix{ExpressionRef} # $, r x g

    # Constraints

    minnewoptions::Matrix{GreaterThanConstraintRef} # r x g
    maxnewoptions::Matrix{LessThanConstraintRef}

    minnewholding::Matrix{GreaterThanConstraintRef}

    minnewbuilds_vesting::Matrix{GreaterThanConstraintRef}
    minnewbuilds_holding::Matrix{GreaterThanConstraintRef}
    maxnewbuilds::Matrix{LessThanConstraintRef}

    minnewdispatching::Matrix{GreaterThanConstraintRef}

    minnewretirements_dispatching::Matrix{GreaterThanConstraintRef}
    minnewretirements_building::Matrix{GreaterThanConstraintRef}

    optionsvested_balance::Matrix{EqualToConstraintRef}
    buildsfinished_balance::Matrix{EqualToConstraintRef}

    minvesting::Matrix{GreaterThanConstraintRef}
    minholding::Matrix{GreaterThanConstraintRef}
    minbuilding::Matrix{GreaterThanConstraintRef}
    mindispatching::Matrix{GreaterThanConstraintRef}
    minretired::Matrix{GreaterThanConstraintRef}

    mintotalretirements::Matrix{GreaterThanConstraintRef}

    function ResourceInvestments{}(
        optioncost::Matrix{Float64}, buildcost::Matrix{Float64},
        retirementcost::Matrix{Float64},
        optionleadtime::Matrix{Int}, buildleadtime::Matrix{Int},
        maxnewoptions::Matrix{Int}, maxnewbuilds::Matrix{Int},
        forcedretirements::Matrix{Int})

        R, G = size(maxnewoptions)
        @assert size(maxnewbuilds) == (R,G)

        @assert size(optioncost) == (R,G)
        @assert size(buildcost) == (R,G)
        @assert size(retirementcost) == (R,G)

        @assert size(optionleadtime) == (R,G)
        @assert size(buildleadtime) == (R,G)

        @assert size(maxnewoptions) == (R,G)

        new{R,G}(optioncost, buildcost, retirementcost,
                 optionleadtime, buildleadtime,
                 maxnewoptions, maxnewbuilds, forcedretirements)

    end

end

function ResourceInvestments(
    tech::AbstractTechnology{G},
    regionlookup::Dict{String,Int},
    resourcetypefolder::String
) where G

    R = length(regionlookup)
    techdata = DataFrame!(CSV.File(joinpath(resourcetypefolder, "parameters.csv"),
                                   types=scenarios_resource_param_types))
    genlookup = Dict(zip(tech.name, 1:G))

    optioncost = zeros(Float64, R, G)
    buildcost = zeros(Float64, R, G)
    retirementcost = zeros(Float64, R, G)

    optionleadtime = zeros(Int, R, G)
    buildleadtime = zeros(Int, R, G)

    newoptionslimit = zeros(Int, R, G)
    newbuildslimit = zeros(Int, R, G)
    newforcedretirements = zeros(Int, R, G)

    for row in eachrow(techdata)

        region_idx = regionlookup[row.region]
        gen_idx = genlookup[row.class]

        optioncost[region_idx, gen_idx] = row.optioncost
        buildcost[region_idx, gen_idx] = row.buildcost
        retirementcost[region_idx, gen_idx] = row.retirementcost

        optionleadtime[region_idx, gen_idx] = row.optionleadtime
        buildleadtime[region_idx, gen_idx] = row.buildleadtime

        newoptionslimit[region_idx, gen_idx] = row.newoptionslimit
        newbuildslimit[region_idx, gen_idx] = row.newbuildslimit
        newforcedretirements[region_idx, gen_idx] = row.newforcedretirements

    end

    return ResourceInvestments(
        optioncost, buildcost, retirementcost,
        optionleadtime, buildleadtime,
        newoptionslimit, newbuildslimit, newforcedretirements)

end

function setup!(
    s::AbstractScenario,
    invtype::Symbol,
    m::Model,
    initconds::InitialInvestments{R,G}
) where {R, G}

    invprob = s.investmentproblem
    Rs = invprob.regionnames
    Gs = getfield(invprob.technologies, invtype).name

    invs = getfield(s.investments, invtype)

    # Variables

    invs.newoptions = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newoptions, "newoptions_$(invtype)_$(s.name)", Rs, Gs)

    invs.newholding = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newholding, "newholding_$(invtype)_$(s.name)", Rs, Gs)

    invs.newbuilds_vesting = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newbuilds_vesting, "newbuilds_vesting_$(invtype)_$(s.name)", Rs, Gs)

    invs.newbuilds_holding = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newbuilds_holding, "newbuilds_holding_$(invtype)_$(s.name)", Rs, Gs)

    invs.newdispatching = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newdispatching, "newdispatching_$(invtype)_$(s.name)", Rs, Gs)

    invs.newretirements_dispatching = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newretirements_dispatching, "newretirements_dispatching_$(invtype)_$(s.name)", Rs, Gs)

    invs.newretirements_building = @variable(m, [1:R, 1:G], Int)
    varnames!(invs.newretirements_building, "newretirements_building_$(invtype)_$(s.name)", Rs, Gs)

    # Expressions

    invs.optionsvested = @expression(m, [r in 1:R, g in 1:G],
        invs.newholding[r,g] + invs.newbuilds_vesting[r,g])

    invs.buildsstarted = @expression(m, [r in 1:R, g in 1:G],
        invs.newbuilds_holding[r,g] + invs.newbuilds_vesting[r,g])

    invs.buildsfinished = @expression(m, [r in 1:R, g in 1:G],
        invs.newretirements_building[r,g] + invs.newdispatching[r,g])

    invs.newretirements = @expression(m, [r in 1:R, g in 1:G],
        invs.newretirements_dispatching[r,g] + invs.newretirements_building[r,g])

    invs.vesting = @expression(m, [r in 1:R, g in 1:G],
        invs.newoptions[r,g] - invs.optionsvested[r,g])

    invs.holding = @expression(m, [r in 1:R, g in 1:G],
        invs.newholding[r,g] - invs.newbuilds_holding[r,g])

    invs.building = @expression(m, [r in 1:R, g in 1:G],
        invs.buildsstarted[r,g] - invs.buildsfinished[r,g])

    invs.dispatching = @expression(m, [r in 1:R, g in 1:G],
        invs.newdispatching[r,g] - invs.newretirements_dispatching[r,g])

    invs.retired = @expression(m, [r in 1:R, g in 1:G],
        invs.newretirements[r,g])

    invs.totalforcedretirements = @expression(m, [r in 1:R, g in 1:G],
        invs.newforcedretirements[r,g])

    if isnothing(s.parent)
        invs.holding .+= initconds.options
        invs.dispatching .+= initconds.builds
    else
        parentinvs = getfield(s.parent.investments, invtype)
        invs.vesting .+= parentinvs.vesting
        invs.holding .+= parentinvs.holding
        invs.building .+= parentinvs.building
        invs.dispatching .+= parentinvs.dispatching
        invs.retired .+= parentinvs.retired
        invs.totalforcedretirements .+= parentinvs.totalforcedretirements
    end

    invs.investmentcosts =
        @expression(m, [r in 1:R, g in 1:G],
                    invs.optioncost[r,g] * invs.newoptions[r,g] +
                    invs.buildcost[r,g] * invs.buildsstarted[r,g] +
                    invs.retirementcost[r,g] * invs.newretirements[r,g])

    invs.minnewoptions = @constraint(m, [r in 1:R, g in 1:G],
        invs.newoptions[r,g] >= 0)

    invs.minnewholding = @constraint(m, [r in 1:R, g in 1:G],
        invs.newholding[r,g] >= 0)

    invs.minnewbuilds_vesting = @constraint(m, [r in 1:R, g in 1:G],
        invs.newbuilds_vesting[r,g] >= 0)

    invs.minnewbuilds_holding = @constraint(m, [r in 1:R, g in 1:G],
        invs.newbuilds_holding[r,g] >= 0)

    invs.minnewdispatching = @constraint(m, [r in 1:R, g in 1:G],
        invs.newdispatching[r,g] >= 0)

    invs.minnewretirements_dispatching = @constraint(m, [r in 1:R, g in 1:G],
        invs.newretirements_dispatching[r,g] >= 0)

    invs.minnewretirements_building = @constraint(m, [r in 1:R, g in 1:G],
        invs.newretirements_building[r,g] >= 0)

    invs.minvesting = @constraint(m, [r in 1:R, g in 1:G],
        invs.vesting[r,g] >= 0)

    invs.minholding = @constraint(m, [r in 1:R, g in 1:G],
        invs.holding[r,g] >= 0)

    invs.minbuilding = @constraint(m, [r in 1:R, g in 1:G],
        invs.building[r,g] >= 0)

    invs.mindispatching = @constraint(m, [r in 1:R, g in 1:G],
        invs.dispatching[r,g] >= 0)

    invs.minretired = @constraint(m, [r in 1:R, g in 1:G],
        invs.retired[r,g] >= 0)

    invs.maxnewoptions =
        @constraint(m, [r in 1:R, g in 1:G],
                    invs.newoptions[r,g] <= invs.newoptionslimit[r,g])

    invs.maxnewbuilds =
        @constraint(m, [r in 1:R, g in 1:G],
                    invs.buildsstarted[r,g] <= invs.newbuildslimit[r,g])

    invs.optionsvested_balance = @constraint(m, [r in 1:R, g in 1:G],
        invs.optionsvested[r,g] ==
        maturing(s, r, g, invtype, :optionleadtime, :newoptions))

    invs.buildsfinished_balance = @constraint(m, [r in 1:R, g in 1:G],
        invs.buildsfinished[r,g] ==
        maturing(s, r, g, invtype, :buildleadtime, :buildsstarted))

    invs.mintotalretirements =
        @constraint(m, [r in 1:R, g in 1:G],
                    invs.retired[r,g] >= invs.totalforcedretirements[r,g])

    return

end

welfare(x::ResourceInvestments) = -sum(x.investmentcosts)

struct Investments{R,G1,G2,G3}
    thermal::ResourceInvestments{R,G1}
    variable::ResourceInvestments{R,G2}
    storage::ResourceInvestments{R,G3}
end

function loadinvestments(
    techs::Technologies, regions::Dict{String,Int}, resourcepath::String)

    thermal =
        ResourceInvestments(techs.thermal, regions,
                            joinpath(resourcepath, "thermal"))
    variable =
        ResourceInvestments(techs.variable, regions,
                            joinpath(resourcepath, "variable"))
    storage =
        ResourceInvestments(techs.storage, regions,
                            joinpath(resourcepath, "storage"))

    return Investments(thermal, variable, storage)

end

welfare(inv::Investments) =
    welfare(inv.thermal) + welfare(inv.variable) + welfare(inv.storage)


text:
@fastmath function middleloop!( ir::Int, rb_tar_idx::Int, rb_src_idx::Int,
                                pssrb::PotentialSimulationSetupRB{T, Cylindrical, 3, Array{T, 3}},
                                update_even_points::Val{even_points},
                                depletion_handling::Val{depletion_handling_enabled},
                                is_weighting_potential::Val{_is_weighting_potential},
                                only2d::Val{only_2d},
                                idx3iseven::Val{idx3_is_even})::Nothing where {T, even_points, depletion_handling_enabled, _is_weighting_potential, only_2d, idx3_is_even}
    @inbounds begin
        inr = ir - 1

        pwwrr               = pssrb.geom_weights[1][1, inr]
        pwwrl               = pssrb.geom_weights[1][2, inr]
        r_inv_pwΔmpr        = pssrb.geom_weights[1][3, inr]
        Δr_ext_inv_r_pwmprr = pssrb.geom_weights[1][4, inr]
        Δr_ext_inv_l_pwmprl = pssrb.geom_weights[1][5, inr]
        Δmpr_squared        = pssrb.geom_weights[1][6, inr]

        line_weights::Array{T, 2} = Array{T, 2}(undef, size(pssrb.potential, 1) - 2, 6)
        # Even though this causes some allocations it is faster than using a predefined array, e.g. stored in pssrb
        # Especially when using multiple threads

        if only_2d
            iφ = 2
            inφ = iφ - 1

            pwwφr        = pssrb.geom_weights[2][1, inφ]
            pwwφl        = pssrb.geom_weights[2][2, inφ]
            pwΔmpφ       = pssrb.geom_weights[2][3, inφ]
            Δφ_ext_inv_r = pssrb.geom_weights[2][4,  iφ]
            Δφ_ext_inv_l = pssrb.geom_weights[2][4, inφ]

            if inr == 1
                pwwφr = T(0.5)
                pwwφl = T(0.5)
                pwΔmpφ = T(2π)
                Δφ_ext_inv_r = inv(pwΔmpφ)
                Δφ_ext_inv_l = Δφ_ext_inv_r
            end
            pwwrr_pwwφr = pwwrr * pwwφr
            pwwrr_pwwφl = pwwrr * pwwφl
            pwwrl_pwwφr = pwwrl * pwwφr
            pwwrl_pwwφl = pwwrl * pwwφl

            Δr_ext_inv_r_pwmprr_pwΔmpφ = Δr_ext_inv_r_pwmprr * pwΔmpφ
            Δr_ext_inv_l_pwmprl_pwΔmpφ = Δr_ext_inv_l_pwmprl * pwΔmpφ
            pwΔmpφ_Δmpr_squared = pwΔmpφ * Δmpr_squared
            r_inv_pwΔmpr_Δφ_ext_inv_r = r_inv_pwΔmpr * Δφ_ext_inv_r
            r_inv_pwΔmpr_Δφ_ext_inv_l = r_inv_pwΔmpr * Δφ_ext_inv_l

            load_weights_for_innerloop!(line_weights, pssrb, iφ, inφ, ir, inr,
                update_even_points, Val(idx3_is_even ? true : false),
                pwwrr, pwwrl, pwwφr, pwwφl,
                pwwrr_pwwφr, pwwrl_pwwφr, pwwrr_pwwφl, pwwrl_pwwφl,
                pwΔmpφ_Δmpr_squared,
                Δr_ext_inv_r_pwmprr_pwΔmpφ, Δr_ext_inv_l_pwmprl_pwΔmpφ,
                r_inv_pwΔmpr_Δφ_ext_inv_r, r_inv_pwΔmpr_Δφ_ext_inv_l,
            )

            innerloop!(line_weights, pssrb, iφ, inφ, ir, inr, rb_tar_idx, rb_src_idx,
                update_even_points, Val(idx3_is_even ? true : false),
                depletion_handling, is_weighting_potential, only2d)
        else
            #=
                Splitting this into two loops over even and uneven iφ
                seems to increase the performance quite a bit
                by hard-coding oscillating type:
                    rφi_is_even_t::Union{Val{true}, Val{false}} = Val(iseven(ir + iφ) ? true : false)
            =#
            for iφ in 2:2:(size(pssrb.potential, 2) - 1)
                inφ = iφ - 1

                pwwφr        = pssrb.geom_weights[2][1, inφ]
                pwwφl        = pssrb.geom_weights[2][2, inφ]
                pwΔmpφ       = pssrb.geom_weights[2][3, inφ]
                Δφ_ext_inv_r = pssrb.geom_weights[2][4,  iφ]
                Δφ_ext_inv_l = pssrb.geom_weights[2][4, inφ]

                if inr == 1
                    pwwφr = T(0.5)
                    pwwφl = T(0.5)
                    pwΔmpφ = T(2π)
                    Δφ_ext_inv_r = inv(pwΔmpφ)
                    Δφ_ext_inv_l = Δφ_ext_inv_r
                end
                pwwrr_pwwφr = pwwrr * pwwφr
                pwwrr_pwwφl = pwwrr * pwwφl
                pwwrl_pwwφr = pwwrl * pwwφr
                pwwrl_pwwφl = pwwrl * pwwφl

                Δr_ext_inv_r_pwmprr_pwΔmpφ = Δr_ext_inv_r_pwmprr * pwΔmpφ
                Δr_ext_inv_l_pwmprl_pwΔmpφ = Δr_ext_inv_l_pwmprl * pwΔmpφ
                pwΔmpφ_Δmpr_squared = pwΔmpφ * Δmpr_squared
                r_inv_pwΔmpr_Δφ_ext_inv_r = r_inv_pwΔmpr * Δφ_ext_inv_r
                r_inv_pwΔmpr_Δφ_ext_inv_l = r_inv_pwΔmpr * Δφ_ext_inv_l

                load_weights_for_innerloop!(line_weights, pssrb, iφ, inφ, ir, inr,
                    update_even_points, Val(idx3_is_even ? true : false),
                    pwwrr, pwwrl, pwwφr, pwwφl,
                    pwwrr_pwwφr, pwwrl_pwwφr, pwwrr_pwwφl, pwwrl_pwwφl,
                    pwΔmpφ_Δmpr_squared,
                    Δr_ext_inv_r_pwmprr_pwΔmpφ, Δr_ext_inv_l_pwmprl_pwΔmpφ,
                    r_inv_pwΔmpr_Δφ_ext_inv_r, r_inv_pwΔmpr_Δφ_ext_inv_l,
                )

                innerloop!(line_weights, pssrb, iφ, inφ, ir, inr, rb_tar_idx, rb_src_idx,
                    update_even_points, Val(idx3_is_even ? true : false),
                    depletion_handling, is_weighting_potential, only2d)
            end
            for iφ in 3:2:(size(pssrb.potential, 2) - 1)
                inφ = iφ - 1

                pwwφr        = pssrb.geom_weights[2][1, inφ]
                pwwφl        = pssrb.geom_weights[2][2, inφ]
                pwΔmpφ       = pssrb.geom_weights[2][3, inφ]
                Δφ_ext_inv_r = pssrb.geom_weights[2][4,  iφ]
                Δφ_ext_inv_l = pssrb.geom_weights[2][4, inφ]

                if inr == 1
                    pwwφr = T(0.5)
                    pwwφl = T(0.5)
                    pwΔmpφ = T(2π)
                    Δφ_ext_inv_r = inv(pwΔmpφ)
                    Δφ_ext_inv_l = Δφ_ext_inv_r
                end
                pwwrr_pwwφr = pwwrr * pwwφr
                pwwrr_pwwφl = pwwrr * pwwφl
                pwwrl_pwwφr = pwwrl * pwwφr
                pwwrl_pwwφl = pwwrl * pwwφl

                Δr_ext_inv_r_pwmprr_pwΔmpφ = Δr_ext_inv_r_pwmprr * pwΔmpφ
                Δr_ext_inv_l_pwmprl_pwΔmpφ = Δr_ext_inv_l_pwmprl * pwΔmpφ
                pwΔmpφ_Δmpr_squared = pwΔmpφ * Δmpr_squared
                r_inv_pwΔmpr_Δφ_ext_inv_r = r_inv_pwΔmpr * Δφ_ext_inv_r
                r_inv_pwΔmpr_Δφ_ext_inv_l = r_inv_pwΔmpr * Δφ_ext_inv_l

                load_weights_for_innerloop!(line_weights, pssrb, iφ, inφ, ir, inr,
                    update_even_points, Val(idx3_is_even ? false : true),
                    pwwrr, pwwrl, pwwφr, pwwφl,
                    pwwrr_pwwφr, pwwrl_pwwφr, pwwrr_pwwφl, pwwrl_pwwφl,
                    pwΔmpφ_Δmpr_squared,
                    Δr_ext_inv_r_pwmprr_pwΔmpφ, Δr_ext_inv_l_pwmprl_pwΔmpφ,
                    r_inv_pwΔmpr_Δφ_ext_inv_r, r_inv_pwΔmpr_Δφ_ext_inv_l,
                )

                innerloop!(line_weights, pssrb, iφ, inφ, ir, inr, rb_tar_idx, rb_src_idx,
                    update_even_points, Val(idx3_is_even ? false : true),
                    depletion_handling, is_weighting_potential, only2d)
            end # φ loop
        end
    end # inbounds
end

function load_weights_for_innerloop!(line_weights, pssrb::PotentialSimulationSetupRB{T, Cylindrical, 3, Array{T, 3}},
                    iφ, inφ, ir, inr,
                    update_even_points, rφi_is_even_t::Val{rφi_is_even},
                    pwwrr, pwwrl, pwwφr, pwwφl,
                    pwwrr_pwwφr, pwwrl_pwwφr, pwwrr_pwwφl, pwwrl_pwwφl,
                    pwΔmpφ_Δmpr_squared,
                    Δr_ext_inv_r_pwmprr_pwΔmpφ, Δr_ext_inv_l_pwmprl_pwΔmpφ,
                    r_inv_pwΔmpr_Δφ_ext_inv_r, r_inv_pwΔmpr_Δφ_ext_inv_l,
                    ) where {T, rφi_is_even}
    # Prepare weightos for most inner loop of the SOR
    @fastmath @inbounds @simd ivdep for iz in 2:(size(pssrb.potential, 1) - 1)
    # @fastmath @inbounds for iz in 2:(size(pssrb.potential, 1) - 1)
        inz::Int = nidx(iz, update_even_points, rφi_is_even_t)

        pwwzr::T        = pssrb.geom_weights[3][1, inz]
        pwwzl::T        = pssrb.geom_weights[3][2, inz]
        pwΔmpz::T       = pssrb.geom_weights[3][3, inz]
        Δz_ext_inv_r::T = pssrb.geom_weights[3][4, inz + 1]
        Δz_ext_inv_l::T = pssrb.geom_weights[3][4, inz]

        ϵ_rrr::T = pssrb.ϵ_r[  ir,  iφ, inz + 1]
        ϵ_rlr::T = pssrb.ϵ_r[  ir, inφ, inz + 1]
        ϵ_rrl::T = pssrb.ϵ_r[  ir,  iφ, inz ]
        ϵ_rll::T = pssrb.ϵ_r[  ir, inφ, inz ]
        ϵ_lrr::T = pssrb.ϵ_r[ inr,  iφ, inz + 1]
        ϵ_llr::T = pssrb.ϵ_r[ inr, inφ, inz + 1]
        ϵ_lrl::T = pssrb.ϵ_r[ inr,  iφ, inz ]
        ϵ_lll::T = pssrb.ϵ_r[ inr, inφ, inz ]

        pwwφr_pwwzr::T = pwwφr * pwwzr
        pwwφl_pwwzr::T = pwwφl * pwwzr
        pwwφr_pwwzl::T = pwwφr * pwwzl
        pwwφl_pwwzl::T = pwwφl * pwwzl
        pwwrl_pwwzr::T = pwwrl * pwwzr
        pwwrr_pwwzr::T = pwwrr * pwwzr
        pwwrl_pwwzl::T = pwwrl * pwwzl
        pwwrr_pwwzl::T = pwwrr * pwwzl

        # right weight in r: wrr
        wrr::T = ϵ_rrr * pwwφr_pwwzr
        wrr    = muladd(ϵ_rlr, pwwφl_pwwzr, wrr)
        wrr    = muladd(ϵ_rrl, pwwφr_pwwzl, wrr)
        wrr    = muladd(ϵ_rll, pwwφl_pwwzl, wrr)
        # left weight in r: wrr
        wrl::T = ϵ_lrr * pwwφr_pwwzr
        wrl    = muladd(ϵ_llr, pwwφl_pwwzr, wrl)
        wrl    = muladd(ϵ_lrl, pwwφr_pwwzl, wrl)
        wrl    = muladd(ϵ_lll, pwwφl_pwwzl, wrl)
        # right weight in φ: wφr
        wφr::T = ϵ_lrr * pwwrl_pwwzr
        wφr    = muladd(ϵ_rrr, pwwrr_pwwzr, wφr)
        wφr    = muladd(ϵ_lrl, pwwrl_pwwzl, wφr)
        wφr    = muladd(ϵ_rrl, pwwrr_pwwzl, wφr)
        # left weight in φ: wφl
        wφl::T = ϵ_llr * pwwrl_pwwzr
        wφl    = muladd(ϵ_rlr, pwwrr_pwwzr, wφl)
        wφl    = muladd(ϵ_lll, pwwrl_pwwzl, wφl)
        wφl    = muladd(ϵ_rll, pwwrr_pwwzl, wφl)
        # right weight in z: wzr
        wzr::T = ϵ_rrr * pwwrr_pwwφr
        wzr    = muladd(ϵ_rlr, pwwrr_pwwφl, wzr)
        wzr    = muladd(ϵ_lrr, pwwrl_pwwφr, wzr)
        wzr    = muladd(ϵ_llr, pwwrl_pwwφl, wzr)
        # left weight in z: wzr
        wzl::T = ϵ_rrl * pwwrr_pwwφr
        wzl    = muladd(ϵ_rll, pwwrr_pwwφl, wzl)
        wzl    = muladd(ϵ_lrl, pwwrl_pwwφr, wzl)
        wzl    = muladd(ϵ_lll, pwwrl_pwwφl, wzl)

        wrr *= Δr_ext_inv_r_pwmprr_pwΔmpφ * pwΔmpz
        wrl *= Δr_ext_inv_l_pwmprl_pwΔmpφ * pwΔmpz
        wφr *= r_inv_pwΔmpr_Δφ_ext_inv_r * pwΔmpz
        wφl *= r_inv_pwΔmpr_Δφ_ext_inv_l * pwΔmpz
        wzr *= Δz_ext_inv_r * pwΔmpφ_Δmpr_squared
        wzl *= Δz_ext_inv_l * pwΔmpφ_Δmpr_squared

        # pssrb.weights_line[iz-1, 1, Base.Threads.threadid()] = wrr
        # pssrb.weights_line[iz-1, 2, Base.Threads.threadid()] = wrl
        # pssrb.weights_line[iz-1, 3, Base.Threads.threadid()] = wφr
        # pssrb.weights_line[iz-1, 4, Base.Threads.threadid()] = wφl
        # pssrb.weights_line[iz-1, 5, Base.Threads.threadid()] = wzr
        # pssrb.weights_line[iz-1, 6, Base.Threads.threadid()] = wzl

        line_weights[iz-1, 1] = wrr
        line_weights[iz-1, 2] = wrl
        line_weights[iz-1, 3] = wφr
        line_weights[iz-1, 4] = wφl
        line_weights[iz-1, 5] = wzr
        line_weights[iz-1, 6] = wzl
    end # z loop
    nothing
end

function innerloop!(line_weights, pssrb::PotentialSimulationSetupRB{T, Cylindrical, 3, Array{T, 3}},
                    iφ, inφ, ir, inr, rb_tar_idx, rb_src_idx,
                    update_even_points, rφi_is_even_t::Val{rφi_is_even},
                    depletion_handling::Val{depletion_handling_enabled},
                    is_weighting_potential::Val{_is_weighting_potential},
                    only2d::Val{only_2d}) where {T, depletion_handling_enabled, _is_weighting_potential, only_2d, rφi_is_even}
    # Most inner loop of the SOR
    @fastmath @inbounds @simd ivdep for iz in 2:(size(pssrb.potential, 1) - 1)
    # @fastmath @inbounds for iz in 2:(size(pssrb.potential, 1) - 1)
        izr::Int = get_rbidx_right_neighbour(iz, update_even_points, rφi_is_even_t)

        wrr = line_weights[iz-1, 1] #pssrb.weights_line[iz-1, 1, Base.Threads.threadid()]
        wrl = line_weights[iz-1, 2] #pssrb.weights_line[iz-1, 2, Base.Threads.threadid()]
        wφr = line_weights[iz-1, 3] #pssrb.weights_line[iz-1, 3, Base.Threads.threadid()]
        wφl = line_weights[iz-1, 4] #pssrb.weights_line[iz-1, 4, Base.Threads.threadid()]
        wzr = line_weights[iz-1, 5] #pssrb.weights_line[iz-1, 5, Base.Threads.threadid()]
        wzl = line_weights[iz-1, 6] #pssrb.weights_line[iz-1, 6, Base.Threads.threadid()]

        old_potential::T = pssrb.potential[iz, iφ, ir, rb_tar_idx]

        vrr::T = pssrb.potential[     iz,     iφ, ir + 1, rb_src_idx]
        vrl::T = pssrb.potential[     iz,     iφ,    inr, rb_src_idx]
        vφr::T = only_2d ? old_potential : pssrb.potential[ iz, iφ + 1, ir, rb_src_idx]
        vφl::T = only_2d ? old_potential : pssrb.potential[ iz,    inφ, ir, rb_src_idx]
        vzr::T = pssrb.potential[    izr,     iφ,     ir, rb_src_idx]
        vzl::T = pssrb.potential[izr - 1,     iφ,     ir, rb_src_idx]

        new_potential::T = _is_weighting_potential ? zero(T) : (pssrb.q_eff_imp[iz, iφ, ir, rb_tar_idx] + pssrb.q_eff_fix[iz, iφ, ir, rb_tar_idx])

        new_potential = muladd( wrr, vrr, new_potential)
        new_potential = muladd( wrl, vrl, new_potential)
        new_potential = muladd( wφr, vφr, new_potential)
        new_potential = muladd( wφl, vφl, new_potential)
        new_potential = muladd( wzr, vzr, new_potential)
        new_potential = muladd( wzl, vzl, new_potential)

        new_potential *= pssrb.volume_weights[iz, iφ, ir, rb_tar_idx]

        new_potential -= old_potential
        new_potential = muladd(new_potential, pssrb.sor_const[inr], old_potential)

        if depletion_handling_enabled
            if inr == 1 vrl = vrr end
            vmin::T = min(vrr, vrl, vφr, vφl, vzr, vzl)
            vmax::T = max(vrr, vrl, vφr, vφl, vzr, vzl)
            new_point_type = pssrb.point_types[iz, iφ, ir, rb_tar_idx]
            if new_potential < vmin || new_potential > vmax
                new_potential -= pssrb.q_eff_imp[iz, iφ, ir, rb_tar_idx] * pssrb.volume_weights[iz, iφ, ir, rb_tar_idx] * pssrb.sor_const[inr]
                if (pssrb.point_types[iz, iφ, ir, rb_tar_idx] & undepleted_bit == 0 &&
                    pssrb.point_types[iz, iφ, ir, rb_tar_idx] & pn_junction_bit > 0)
                    # pssrb.point_types[iz, iφ, ir, rb_tar_idx] += undepleted_bit
                    new_point_type += undepleted_bit
                end # mark this point as undepleted
            elseif pssrb.point_types[iz, iφ, ir, rb_tar_idx] & undepleted_bit > 0
                new_point_type -= undepleted_bit
            end
            pssrb.point_types[iz, iφ, ir, rb_tar_idx] = new_point_type
        end

        pssrb.potential[iz, iφ, ir, rb_tar_idx]::T = ifelse(pssrb.point_types[iz, iφ, ir, rb_tar_idx] & update_bit > 0, new_potential, old_potential)
    end # z loop
    nothing
end