text:
"""
COCOS Structure

`cocos::Int`           = COCOS ID number\\
`exp_Bp::Int`          = 0 or 1, depending if psi is already divided by 2pi or not, respectively\\
`sigma_Bp::Int`        = +1 or -1, depending if psi is increasing or decreasing with Ip and B0 positive\\
`sigma_RpZ::Int`       = +1 or -1, depending if (R,phi,Z) is right-handed or (R,Z,phi), respectively\\
`sigma_rhotp::Int`     = +1 or -1, depending if (rho, theta, phi) is right-handed or (rho,phi,theta), repectively\\
`sign_q_pos::Int`      = +1 or -1, depending if q is positive or negative with Ip and B0 positive\\
`sign_pprime_pos::Int` = +1 or -1, depending if dp/dpsi is positive or negative with Ip and B0 positive
"""
struct COCOS
    cocos::Int           # COCOS ID number
    exp_Bp::Int          # 0 or 1, depending if psi is already divided by 2pi or not, respectively
    sigma_Bp::Int        # +1 or -1, depending if psi is increasing or decreasing with Ip and B0 positive
    sigma_RpZ::Int       # +1 or -1, depending if (R,phi,Z) is right-handed or (R,Z,phi), respectively
    sigma_rhotp::Int     # +1 or -1, depending if (rho, theta, phi) is right-handed or (rho,phi,theta), repectively
    sign_q_pos::Int      # +1 or -1, depending if q is positive or negative with Ip and B0 positive
    sign_pprime_pos::Int # +1 or -1, depending if dp/dpsi is positive or negative with Ip and B0 positive
end

cocos(CC::COCOS) = CC

function Base.show(io::IO, CC::COCOS)
    println(io, "COCOS = $(CC.cocos)")
    println(io, " e_Bp  = $(CC.exp_Bp)")
    println(io, " σ_Bp  = $(CC.sigma_Bp)")

    rpz = Dict(1=>"(R,Φ,Z)", -1=>"(R,Z,Φ)")
    rpz_dir = Dict(1=>"CCW", -1=>"CW")
    rhotp = Dict(1=>"(ρ,θ,Φ)", -1=>"(ρ,Φ,θ)")
    rhotp_dir = Dict(1=>"CW", -1=>"CCW")

    println(io, " σ_RΦZ = $(rpz[CC.sigma_RpZ]): $(CC.sigma_RpZ)")
    println(io, " σ_ρθΦ = $(rhotp[CC.sigma_rhotp]): $(CC.sigma_rhotp)")
    println(io, " Φ from top: $(rpz_dir[CC.sigma_RpZ])")
    println(io, " θ from front: $(rhotp_dir[CC.sigma_RpZ*CC.sigma_rhotp])")

    inc = Dict(1=>"Increasing", -1=>"Decreasing")
    println(io, " ψ_ref: $(inc[CC.sigma_Bp]) assuming +Ip, +B0")
    println(io, " sign(q) = $(CC.sign_q_pos) assuming +Ip, +B0")
    print(io, " sign(p') = $(CC.sign_pprime_pos) assuming +Ip, +B0")
end

Base.broadcastable(CC::COCOS) = (CC,)

"""
    cylindrical_cocos(c::COCOS, r, phi, z) -> NTuple{3}

Returns the right-handed cylindrical coordinate according to the provided COCOS.

`cocos.sigma_RpZ = +1 -> (r, phi, z)`\\
`cocos.sigma_RpZ = -1 -> (r, z, phi)`

"""
function cylindrical_cocos(c::COCOS, r, phi, z)
    if c.sigma_RpZ == 1
        return (r, phi, z)
    else
        return (r, z, phi)
    end
end

"""
    cylindrical_cocos_indices(c::COCOS) -> NTuple{3}

Returns the indices of the r, phi, and z coordinates, respectively, relative to the provided COCOS.

`cocos.sigma_RpZ = +1 -> (1,2,3)`\\
`cocos.sigma_RpZ = -1 -> (1,3,2)`

"""
function cylindrical_cocos_indices(c::COCOS)
    if c.sigma_RpZ == 1
        return (1,2,3)
    else
        return (1,3,2)
    end
end

"""
    poloidal_cocos(c::COCOS, rho, theta, phi) -> NTuple{3}

Returns the right-handed poloidal coordinate according to the provided COCOS.

cocos.sigma_rhotp = +1 -> (rho, theta, phi)\\
cocos.sigma_rhotp = -1 -> (rho, phi, theta)
"""
function poloidal_cocos(c::COCOS, rho, theta, phi)
    if c.sigma_rhotp == 1
        return (rho, theta, phi)
    else
        return (rho, phi, theta)
    end
    return _poloidal_cocos(Val(c.sigma_rhotp), rho, theta, phi)
end

"""
    poloidal_cocos_indices(c::COCOS) -> NTuple{3}

Returns the indices of the rho, theta, and phi coordinates, respectively, relative to the provided COCOS.

cocos.sigma_rhotp = +1 -> (1,2,3)\\
cocos.sigma_rhotp = -1 -> (1,3,2)
"""
function poloidal_cocos_indices(c::COCOS, rho, theta, phi)
    if c.sigma_rhotp == 1
        return (1,2,3)
    else
        return (1,3,2)
    end
end

"""
    cocos(cocos_ID) -> COCOS

Returns `COCOS` structure given the `cocos_ID` number
"""
function cocos(cocos_in)

    exp_Bp = cocos_in >= 11 ? 1 : 0

    if cocos_in ∈ (1,11)
        # ITER, Boozer are COCOS=11
        return COCOS(cocos_in,exp_Bp,1,1,1,1,-1)
    elseif cocos_in ∈ (2,12)
        # CHEASE, ONETWO, Hinton-Hazeltine, LION is COCOS=2
        return COCOS(cocos_in,exp_Bp,1,-1,1,1,-1)
    elseif cocos_in ∈ (3,13)
        # Freidberg, CAXE, KINX, EFIT are COCOS=3
        # EU-ITM up to end of 2011 is COCOS=13
        return COCOS(cocos_in,exp_Bp,-1,1,-1,-1,1)
    elseif cocos_in ∈ (4,14)
        return COCOS(cocos_in,exp_Bp,-1,-1,-1,-1,1)
    elseif cocos_in ∈ (5,15)
        return COCOS(cocos_in,exp_Bp,1,1,-1,-1,-1)
    elseif cocos_in ∈ (6,16)
        return COCOS(cocos_in,exp_Bp,1,-1,-1,-1,-1)
    elseif cocos_in ∈ (7,17)
        return COCOS(cocos_in,exp_Bp,-1,1,1,1,1)
    elseif cocos_in ∈ (8,18)
        return COCOS(cocos_in,exp_Bp,-1,-1,1,1,1)
    else
        throw(ArgumentError("COCOS = $cocos_in does not exist"))
    end
end

"""
    check_cocos(sigma_B0, sigma_Ip, sigma_F, sigma_pprime, sigma_q, sigma_dpsi, cc::COCOS; verbose = false) -> Bool

Returns True if equilibrium quantities are consistant with given COCOS

Arguments:\\
`sigma_B0` - Toroidal magnetic field sign\\
`sigma_Ip` - Plasma current sign\\
`sigma_F` - Poloidal current sign\\
`sigma_pprime` - Pressure gradient sign\\
`sigma_dpsi` - dpsi/dr sign\\
`cc::Union{Int,COCOS}` - COCOS structure or ID
"""
function check_cocos(sigma_B0, sigma_Ip, sigma_F, sigma_pprime,
                     sigma_q, sigma_dpsi,
                     cc::Union{Int,COCOS}; verbose=false)

    cc = cocos(cc)

    valid = true
    if sigma_q*cc.sigma_rhotp*sigma_Ip*sigma_B0 < 0
        verbose && @warn "sigma_q($sigma_q) ≠ sigma_rhotp($(cc.sigma_rhotp))*sigma_Ip($sigma_Ip)*sigma_B0($sigma_B0)"
        valid = false
    end

    if sigma_F*sigma_B0 < 0
        verbose && @warn "Signs of F and B0 are not consistant"
        valid = false
    end

    if sigma_dpsi*cc.sigma_Bp*sigma_Ip < 0
        if sigma_dpsi > 0
            verbose && @warn "psi should be decreasing with sign(Ip) = $(sigma_Ip) for COCOS = $(cc.cocos)"
        else
            verbose && @warn "psi should be increasing with sign(Ip) = $(sigma_Ip) for COCOS = $(cc.cocos)"
        end
        valid = false
    elseif sigma_pprime*sigma_Ip*cc.sigma_Bp > 0
        verbose && @warn "sign(pprime) should be $(-sigma_Ip*cc.sigma_Bp)"
        valid = false
    end

    return valid
end

"""
    check_cocos(B0, Ip, F::AbstractVector, pprime::AbstractVector, q::AbstractVector, psi::AbstracVector, cc::COCOS; verbose = false) -> Bool

Returns True if equilibrium quantities are consistant with given COCOS

Arguments:\\
`B0` - Toroidal magnetic field\\
`Ip` - Plasma current\\
`F::AbstractVector` - Poloidal current as a function of `psi`\\
`pprime::AbstracVector` - Pressure gradient w.r.t. `psi` as a function of `psi`\\
`psi::AbstractVector` - Poloidal flux\\
`cc::Union{Int,COCOS}` - COCOS structure or ID
"""
function check_cocos(B0, Ip, F::AbstractVector, pprime::AbstractVector,
                     q::AbstractVector, psi::AbstractVector,
                     cc::Union{Int,COCOS}; kwargs...)

    sigma_B0 = sign(B0)
    sigma_Ip = sign(Ip)
    sigma_F = sign(sum(F))
    sigma_pprime = sign(sum(pprime))
    sigma_q = sign(q[end])
    sigma_dpsi = sign(psi[end] - psi[1])

    return check_cocos(sigma_B0, sigma_Ip, sigma_F, sigma_pprime, sigma_q, sigma_dpsi, cc; kwargs...)

end

"""
    transform_cocos(cc_in::Union{Int,COCOS}, cc_out::Union{Int,COCOS};
                    sigma_Ip = nothing, sigma_B0=nothing,
                    ld = (1,1), lB = (1,1), exp_mu0 = (0,0)) -> Dict

Returns a dictionary of the multiplicative factors to transform COCOS from `cc_in` to `cc_out`

Arguments:\\
`sigma_Ip::Union{NTuple{2,Int},Nothing}` - A tuple of the (Input, Output) current sign or nothing\\
`sigma_B0::Union{NTuple{2,Int},Nothing}` - A tuple of the (Input, Output) toroidal field sign or nothing\\
`ld::NTuple{2,<:Real}` - A tuple of the (Input, Output) length scale factor. Default = (1,1)\\
`lB::NTuple{2,<:Real}` - A tuple of the (Input, Output) magnetic field scale factor. Default = (1,1)\\
`exp_mu0::NTuple{2,<:Real}` - A tuple of the (Input, Output) mu0 exponent (0, 1). Default = (0,0)
"""
function transform_cocos(cc_in::COCOS, cc_out::COCOS;
                         sigma_Ip::Union{NTuple{2,Int},Nothing} = nothing,
                         sigma_B0::Union{NTuple{2,Int},Nothing} = nothing,
                         ld::NTuple{2,<:Real} = (1,1),
                         lB::NTuple{2,<:Real} = (1,1),
                         exp_mu0::NTuple{2,<:Real} = (0,0))

    ld_eff = ld[2]/ld[1]
    lB_eff = lB[2]/lB[1]
    exp_mu0_eff = exp_mu0[2] - exp_mu0[1]

    sigma_RpZ_eff = cc_in.sigma_RpZ * cc_out.sigma_RpZ

    if sigma_Ip === nothing
        sigma_Ip_eff = cc_in.sigma_RpZ * cc_out.sigma_RpZ
    else
        sigma_Ip_eff = sigma_Ip[1]*sigma_Ip[2]
    end

    if sigma_B0 === nothing
        sigma_B0_eff = cc_in.sigma_RpZ * cc_out.sigma_RpZ
    else
        sigma_B0_eff = sigma_B0[1]*sigma_B0[2]
    end

    sigma_Bp_eff = cc_in.sigma_Bp * cc_out.sigma_Bp
    exp_Bp_eff = cc_out.exp_Bp - cc_in.exp_Bp
    sigma_rhotp_eff = cc_in.sigma_rhotp * cc_out.sigma_rhotp

    mu0 = 4*pi*1e-7

    transforms = Dict()
    transforms["R"]        = ld_eff
    transforms["Z"]        = ld_eff
    transforms["P"]        = (lB_eff^2)/(mu0^exp_mu0_eff)
    transforms["PSI"]      = lB_eff * ld_eff^2 * sigma_Ip_eff * sigma_Bp_eff * ((2pi)^exp_Bp_eff) * ld_eff^2 * lB_eff
    transforms["ψ"]        = transforms["PSI"]
    transforms["TOR"]      = lB_eff * ld_eff^2 * sigma_B0_eff
    transforms["Φ"]        = transforms["TOR"]
    transforms["PPRIME"]   = (lB_eff/((ld_eff^2)*(mu0^exp_mu0_eff))) * sigma_Ip_eff * sigma_Bp_eff / ((2pi)^exp_Bp_eff)
    transforms["F_FPRIME"] = lB_eff * sigma_Ip_eff * sigma_Bp_eff / ((2pi)^exp_Bp_eff)
    transforms["B"]        = lB_eff * sigma_B0_eff
    transforms["F"]        = sigma_B0_eff * ld_eff * lB_eff
    transforms["I"]        = sigma_Ip_eff * ld_eff * lB_eff / (mu0^exp_mu0_eff)
    transforms["J"]        = sigma_Ip_eff * lB_eff/((mu0^exp_mu0_eff)*ld_eff)
    transforms["Q"]        = sigma_Ip_eff * sigma_B0_eff * sigma_rhotp_eff

    return transforms
end

function transform_cocos(cc_in::Union{Int,COCOS}, cc_out::Union{Int,COCOS}; kwargs...)
    transform_cocos(cocos(cc_in), cocos(cc_out); kwargs...)
end

"""
    identify_cocos(sigma_B0, sigma_Ip, sigma_q, sigma_dpsi, clockwise_phi) -> List of possible COCOS IDs

Utility function to identify COCOS coordinate system. If multiple COCOS are possible, then all are returned.

Arguments:\\
`sigma_B0` - (+1,-1) sign of the toroidal magnetic field\\
`sigma_Ip` - (+1,-1) sign of toroidal plasma current\\
`sigma_q`  - (+1,-1) sign of the safety factor (q) within the plasma\\
`sigma_dpsi` -  +1 if psi is increasing, -1 if psi is decreasing\\
`clockwise_phi::Bool` - (optional) [True, False] if phi angle is defined clockwise or not. This is required to identify odd Vs even COCOS. Note that this cannot be determined from the output of a code.
"""
function identify_cocos(sigma_B0, sigma_Ip, sigma_q, sigma_dpsi,
                        clockwise_phi::Union{Bool,Nothing} = nothing)

    if clockwise_phi === nothing
        sigma_rpz = clockwise_phi
    elseif clockwise_phi
        sigma_rpz = -1
    else
        sigma_rpz = +1
    end

    # return both even and odd COCOS if clockwise_phi is not provided
    if sigma_rpz === nothing
        tmp  = identify_cocos(sigma_B0, sigma_Ip, sigma_q, sigma_dpsi, true)
        tmp2  = identify_cocos(sigma_B0, sigma_Ip, sigma_q, sigma_dpsi, false)
        return (tmp..., tmp2...)
    end

    sigma_Bp = sigma_dpsi / sigma_Ip
    sigma_rhotp = sigma_q / (sigma_Ip * sigma_B0)

    sigma2cocos = Dict(
        (+1, +1, +1) => 1,  # +Bp, +rpz, +rtp
        (+1, -1, +1) => 2,  # +Bp, -rpz, +rtp
        (-1, +1, -1) => 3,  # -Bp, +rpz, -rtp
        (-1, -1, -1) => 4,  # -Bp, -rpz, -rtp
        (+1, +1, -1) => 5,  # +Bp, +rpz, -rtp
        (+1, -1, -1) => 6,  # +Bp, -rpz, -rtp
        (-1, +1, +1) => 7,  # -Bp, +rpz, +rtp
        (-1, -1, +1) => 8)  # -Bp, -rpz, +rtp

    return (sigma2cocos[(sigma_Bp, sigma_rpz, sigma_rhotp)], sigma2cocos[(sigma_Bp, sigma_rpz, sigma_rhotp)] + 10)
end

"""
    identify_cocos(B0, Ip, q, psi, clockwise_phi, a) -> List of possible COCOS IDs

Utility function to identify COCOS coordinate system. If multiple COCOS are possible, then all are returned.

Arguments:\\
`B0` - toroidal magnetic field (with sign)\\
`Ip` - plasma current (with sign)\\
`q`  -  safety factor profile (with sign) as function of psi\\
`psi::AbstractVector` -  Vector of poloidal fluxs from the magnetic axis to the plasma boundary\\
`clockwise_phi::Bool` - (optional) [True, False] if phi angle is defined clockwise or not. This is required to identify odd Vs even COCOS. Note that this cannot be determined from the output of a code.\\
`a::AbstractVector` - (optional) flux surfaces minor radius as function of psi. This is required to identify 2*pi term in psi definition
"""
function identify_cocos(B0, Ip, q::AbstractVector, psi::AbstractVector,
                        clockwise_phi::Union{Bool,Nothing} = nothing,
                        a::Union{AbstractVector,Nothing} = nothing)

    sigma_B0 = sign(B0)
    sigma_Ip = sign(Ip)
    sigma_q = sign(q[1])
    sigma_dpsi = sign(psi[2]-psi[1])

    ccs = identify_cocos(sigma_B0, sigma_Ip, sigma_q, sigma_dpsi, clockwise_phi)

    # identify 2*pi term in psi definition based on q estimate
    if a !== nothing
        index = argmin(abs.(q))
        if index == 1
            index += 1
        end

        q_estimate = abs.((pi * B0 * (a .- a[1]).^2) / (psi .- psi[1]))

        if abs(q_estimate[index] - q[index]) < abs(q_estimate[index] / (2 * pi) - q[index])
            eBp = 1
        else
            eBp = 0
        end

        return filter(x -> x > 10*eBp, ccs)
    else
        # return COCOS<10 as well as COCOS>10 if a is not provided
        return ccs
    end
end

# ----- GEQDSK Interface -----

"""
    check_cocos(g::GEQDSKFIle, cc::Union{Int,COCOS}) -> Bool

Checks if `GEQDSKFile` is consistant with the given cocos, `cc`
"""
function check_cocos(g::GEQDSKFile, cc::Union{Int,COCOS}; kwargs...)
    return check_cocos(g.bcentr, g.current, g.fpol, g.pprime, g.qpsi, g.psi, cocos(cc); kwargs...)
end

"""
    identify_cocos(g::GEQDSKFIle; clockwise_phi=nothing) -> List of possible COCOS IDs

Identifies possible `GEQDSKFile` COCOS. A unique identification requires setting the `clockwise_phi` keyword.
"""
function identify_cocos(g::GEQDSKFile; clockwise_phi = nothing)
    return filter(x -> x < 10, identify_cocos(g.bcentr, g.current, g.qpsi, g.psi, clockwise_phi, nothing))
end

"""
    cocos(g::GEQDSKFile; kwargs...) -> COCOS

Identifies and returns `GEQDSKFile` `COCOS`. A unique identification requires setting the `clockwise_phi` keyword.
"""
function cocos(g::GEQDSKFile; kwargs...)
    cc = identify_cocos(g; kwargs...)
    if length(cc) > 1
        @error "Unable to determine unique COCOS. Try providing clockwise_phi::Bool keyword. Possibilities are $cc"
    end
    return cocos(cc[1])
end

"""
    transform_cocos(g::GEQDSKFile, cc_in::Union{Int,COCOS}, cc_out::Union{Int,COCOS}; kwargs...) -> GEQDSKFile

Transforms the given `GEQDSKFile` with `COCOS=cc_in`, and returns a `GEQDSKFile` with `COCOS=cc_out`.
"""
function transform_cocos(g::GEQDSKFile, cc_in::Union{Int,COCOS}, cc_out::Union{Int,COCOS}; kwargs...)
    T = transform_cocos(cc_in, cc_out; kwargs...)

    g_new = GEQDSKFile(g.file*" w/ cocos = $(cocos(cc_out).cocos)", g.nw, g.nh,
                       g.r*T["R"], g.z*T["Z"], g.rdim*T["R"], g.zdim*T["Z"],
                       g.rleft*T["R"], g.zmid*T["Z"], g.nbbbs, g.rbbbs*T["R"], g.zbbbs*T["Z"],
                       g.limitr, g.rlim*T["R"], g.zlim*T["Z"], g.rcentr*T["R"], g.bcentr*T["B"],
                       g.rmaxis*T["R"], g.zmaxis*T["Z"], g.simag*T["PSI"], g.sibry*T["PSI"],
                       g.psi*T["PSI"], g.current*T["I"], g.fpol*T["F"],
                       g.pres*T["P"], g.ffprim*T["F_FPRIME"], g.pprime*T["PPRIME"],
                       g.qpsi*T["Q"], g.psirz*T["PSI"])

    return g_new
end


text:
# test the add_measure_variable extensions
@testset "Expansion Extensions" begin
    # setup the model testing
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1], supports = [0, 1])
    @infinite_parameter(m, pars[1:2] in [0, 1], supports = [0, 1])
    @variable(m, x, Infinite(par))
    @variable(m, y, Infinite(par, pars))
    @variable(m, x0, Point(x, 0))
    tm = optimizer_model(m)
    @variable(tm, a)
    @variable(tm, b)
    @variable(tm, c)
    @variable(tm, d)
    data = transcription_data(tm)
    data.infvar_mappings[x] = [a, b, c]
    data.infvar_supports[x] = [(0.,), (0.5,), (1.,)]
    data.infvar_lookup[x] = Dict{Vector{Float64}, Int}([0] => 1, [0.5] => 2, [1] => 3)
    data.infvar_mappings[y] = [a, b, c, d]
    data.infvar_supports[y] = [(0., [0., 0.]), (0., [1., 1.]), (1., [0., 0.]), (1., [1., 1.])]
    data.infvar_lookup[y] = Dict{Vector{Float64}, Int}([0, 0, 0] => 1, [0, 1, 1] => 2,
                                                       [1, 0, 0] => 3, [1, 1, 1] => 4)
    data.finvar_mappings[x0] = a
    key = Val(:TransData)
    IOTO.set_parameter_supports(tm, m)
    # test add_measure_variable for point variables
    @testset "add_measure_variable (Point)" begin
        info = JuMP.VariableInfo(true, 0., true, 0., true, 0., true, 0., true, true)
        var = PointVariable(info, x, [1.])
        vref = GeneralVariableRef(m, -1, PointVariableIndex)
        @test InfiniteOpt.add_measure_variable(tm, var, key) == vref
        @test transcription_variable(vref) == c
    end
    # test add_measure_variable for semi_infinite variables
    @testset "add_measure_variable (SemiInfinite)" begin
        var = SemiInfiniteVariable(y, Dict{Int, Float64}(1 => 0), [2, 3], [2])
        vref = GeneralVariableRef(m, -1, SemiInfiniteVariableIndex)
        @test InfiniteOpt.add_measure_variable(tm, var, key) == vref
        @test data.semi_infinite_vars == [var]
        @test a in transcription_variable(vref)
        @test b in transcription_variable(vref)
        @test sort!(supports(vref)) == [([0., 0.], ), ([1., 1.], )]
    end
    # test delete_semi_infinite_variable extension
    @testset "delete_semi_infinite_variable" begin
        vref = SemiInfiniteVariableRef(m, SemiInfiniteVariableIndex(-1))
        @test InfiniteOpt.delete_semi_infinite_variable(tm, vref, key) isa Nothing
    end
end

ext: jl


text:
using KernelAbstractions: @index, @kernel, NoneEvent

using Oceananigans.TimeSteppers:  store_field_tendencies!

using Oceananigans: prognostic_fields
using Oceananigans.Grids: AbstractGrid
using Oceananigans.Architectures: device_event

using Oceananigans.Utils: launch!

import Oceananigans.TimeSteppers: store_tendencies!

""" Store source terms for `η`. """
@kernel function _store_free_surface_tendency!(Gη⁻, grid, Gη⁰)
    i, j = @index(Global, NTuple)
    @inbounds Gη⁻[i, j, 1] = Gη⁰[i, j, 1]
end

store_free_surface_tendency!(free_surface, model, barrier) = NoneEvent()

function store_free_surface_tendency!(::ExplicitFreeSurface, model, barrier)

    event = launch!(model.architecture, model.grid, :xy,
                    _store_free_surface_tendency!,
                    model.timestepper.G⁻.η,
                    model.grid,
                    model.timestepper.Gⁿ.η,
                    dependencies = barrier)

    return event
end

""" Store previous source terms before updating them. """
function store_tendencies!(model::HydrostaticFreeSurfaceModel)

    barrier = device_event(model.architecture)

    prognostic_field_names = keys(prognostic_fields(model))
    three_dimensional_prognostic_field_names = filter(name -> name != :η, prognostic_field_names)

    events = []

    for field_name in three_dimensional_prognostic_field_names

        field_event = launch!(model.architecture, model.grid, :xyz,
                              store_field_tendencies!,
                              model.timestepper.G⁻[field_name],
                              model.grid,
                              model.timestepper.Gⁿ[field_name],
                              dependencies = barrier)

        push!(events, field_event)
    end

    η_event = store_free_surface_tendency!(model.free_surface, model, barrier)
    push!(events, η_event)

    wait(device(model.architecture), MultiEvent(Tuple(events)))

    return nothing
end


text:
using OrderedCollections
using YaoLocations: plain, AbstractLocations

function block(f, it)
    ret = Expr(:block)
    for xs in it
        push!(ret.args, f(xs...))
    end
    return ret
end

function flatten_block(ex)
    ex isa Expr || return ex
    ex.head === :block && return _flatten_block(ex)
    return Expr(ex.head, map(flatten_block, ex.args)...)
end

function _flatten_block(ex)
    ex isa Expr || return ex
    ex.head === :block || return ex
    ret = Expr(:block)
    for stmt in ex.args
        if stmt isa Expr && stmt.head === :block
            append!(ret.args, flatten_block(stmt).args)
        else
            push!(ret.args, flatten_block(stmt))
        end
    end
    return ret
end

function eval_interp(m::Module, ex)
    ex isa Expr || return ex
    ex.head === :$ && return Base.eval(m, ex.args[1])
    return Expr(ex.head, map(x->eval_interp(m, x), ex.args)...)
end

struct Hoistable
    origin::OrderedDict{Any, Any}
    hoisted::OrderedDict{Any, Any}
end

function Hoistable(X::Matrix)
    origin = OrderedDict()
    for i in axes(X, 1), j in axes(X, 2)
        origin[(i, j)] = X[i, j]
    end
    return Hoistable(origin, OrderedDict())
end

function hoist!(h::Hoistable, key, others...)
    key = isempty(others) ? key : (key, others...)
    haskey(h.hoisted, key) && return h.hoisted[key]
    @gensym hoisted
    h.hoisted[key] = hoisted
    return hoisted
end

function Base.getindex(h::Hoistable, key, others...)
    key = isempty(others) ? key : (key, others...)
    haskey(h.hoisted, key) && return h.hoisted[key]
    return h.origin[key]
end

function hoisted_expr(h::Hoistable)
    ret = Expr(:block)
    for (key, sym) in h.hoisted
        push!(ret.args, :($sym = $(h.origin[key])))
    end
    return ret
end

struct VarDefs
    mangled::OrderedDict{Symbol, Symbol}
    symbols::OrderedDict{Symbol, Any}
end

VarDefs() = VarDefs(OrderedDict{Symbol, Symbol}(), OrderedDict{Symbol, Any}())

function def!(defs::VarDefs, sym, expr; force=false)
    mangled = getfield(defs, :mangled)
    symbols = getfield(defs, :symbols)
    force || haskey(symbols, sym) && return defs
    symbols[sym] = expr
    mangled[sym] = gensym(sym)
    return defs
end

macro def(defs, ex::Expr)
    ex.head === :(=) || error("expect an assignment expr")
    quote
        $def!($defs, $(QuoteNode(ex.args[1])), $(Expr(:quote, ex.args[2])))
        $(ex.args[1]) = $defs.$(ex.args[1])
    end |> esc
end

macro defnew(defs, ex::Expr)
    ex.head === :(=) || error("expect an assignment expr")
    quote
        $def!($defs, $(QuoteNode(ex.args[1])), $(Expr(:quote, ex.args[2])); force=true)
        $(ex.args[1]) = $defs.$(ex.args[1])
    end |> esc
end

function var_exprs(defs::VarDefs)
    mangled = getfield(defs, :mangled)
    symbols = getfield(defs, :symbols)
    isempty(symbols) && return
    ret = Expr(:block)
    for (sym, expr) in symbols
        push!(ret.args, Expr(:(=), mangled[sym], expr))
    end
    return ret
end

function Base.getproperty(defs::VarDefs, name::Symbol)
    mangled = getfield(defs, :mangled)
    haskey(mangled, name) && return mangled[name]
    error("$name is not defined")
end

function Base.show(io::IO, defs::VarDefs)
    mangled = getfield(defs, :mangled)
    symbols = getfield(defs, :symbols)
    println(io, "global variables:")
    for (k, (sym, expr)) in enumerate(symbols)
        print(io, "  ", mangled[sym], ": ", sym, " = ", expr)
        if k != length(symbols)
            println(io)
        end
    end
end

@inline function literal_pow(::Val{:im}, p::Int)
    m = mod(p, 4)
    return if m == 0
        Complex{Int}(1, 0)
    elseif m == 1
        Complex{Int}(0, 1)
    elseif m == 2
        Complex{Int}(-1, 0)
    else
        Complex{Int}(0, -1)
    end
end
# fast pow of ``exp(im * pi / d)^x``
@generated function literal_pow(::Val{:exp}, ::Val{d}, x::Int) where d
    root = Expr(:if)
    curr = root
    for k in 0:(2d-2)
        push!(curr.args, :(m == $k))
        push!(curr.args, exp(im * π * k / d))
        next = Expr(:elseif)
        push!(curr.args, next)
        curr = next
    end

    push!(curr.args, exp(im * π * (2d-1) / d))

    return quote
        m = mod(x, 2d)
        $root
    end
end

function sym_mul(a, b)
    @match (a, b) begin
        ( 0, B) => 0
        ( 1, B) => B
        (-1, B::Number) => -B
        (-1, B) => :(-$B)

        (A,  0) => 0
        (A,  1) => A
        (A::Number, -1) => -A
        (A, -1) => :(-$A)

        _ => :($a * $b)
    end
end

function sym_plus(a, b)
    @match (a, b) begin
        (0, 0) => 0
        (0, B) => B
        (A, 0) => A
        _ => :($a+$b)
    end
end

function nexprs(f, n::Int)
    ex = Expr(:block)
    for k in 1:n
        push!(ex.args, f(k))
    end
    return ex
end

function isperm(x::Matrix)
    all(count(x->!isa(x, Number) || !iszero(x), col) == 1 for col in eachcol(x))
end

function isdiag(x::Matrix)
    for i in 1:size(x, 1), j in 1:size(x, 2)
        e = x[i, j]
        if i != j
            e isa Number && iszero(e) && continue
            return false
        end
    end
    return true
end


text:
export JuliaSolver,getJuliaSolver,copySolver,solveLinearSystem,solveLinearSystem!,setupSolver

## This did not work in 0.7. Check this in the future.
# import Base.\
# function \{T1,T2}(A::Base.SparseArrays.UMFPACK.UmfpackLU{T1},R::SparseMatrixCSC{T2})

	# n,nrhs = size(R)
	# X = zeros(promote_type(T1,T2),n,nrhs)
	# for k=1:nrhs
		# X[:,k] = A\full(vec(R[:,k]))
	# end
	# return X
# end

"""
mutable struct jInvLinearSolvers.JuliaSolver<: AbstractSolver

Fields:

	Ainv         - holds factorization (LU or Cholesky)
	sym          - 0=unsymmetric, 1=symm. pos def, 2=general symmetric
	isTransposed - flag whether A comes transposed or not
	doClear      - flag to clear factorization
	facTime      - cumulative time for factorizations
	nSolve       - number of solves
	solveTime    - cumnulative time for solves
	nFac         - number of factorizations performed

Example:

	Ainv = getJuliaSolver()
"""
mutable struct JuliaSolver<: AbstractSolver
	Ainv
	sym::Int # 0 = unsymmetric, 1 = symmetric s.t A = A';
	isTransposed::Int
	doClear::Int
	facTime::Real
	nSolve::Int
	solveTime::Real
	nFac::Int
end

"""
function jInv.LinearSolvers.getJuliaSolver

Constructor for JuliaSolver

Optional Keyword Arguments

	Ainv = []
	sym = 0
	isTransposed = 0
	doClear = 0
"""
function getJuliaSolver(;Ainv = [],sym = 0,isTransposed = 0, doClear = 0)
	return JuliaSolver(Ainv,sym,isTransposed,doClear,0.0,0,0.0,0);
end

solveLinearSystem(A,B,param::JuliaSolver,doTranspose::Int=0) = solveLinearSystem!(A,B,[],param,doTranspose)

function setupSolver(A::SparseMatrixCSC,param::JuliaSolver)
	tt = time_ns();
	if param.sym==1 && isreal(A)
		param.Ainv = cholesky(A)
	elseif param.sym==2 && isreal(A)
		param.Ainv = ldlt(A)
	else
		param.Ainv = lu(A);
	end
	param.facTime+= (tt-time_ns())/1e+9;
	param.nFac+=1
	return param;
end

function solveLinearSystem!(A::SparseMatrixCSC,B,X,param::JuliaSolver,doTranspose=0)
	if issparse(B)
		#println("");
		#@warn("jInv: Julia solvers do not support sparse RHSs for now. Check in the future");
		if length(size(B))==1
			B = Vector(B);
		else
			B = Matrix(B);
		end
	end
	if param.doClear == 1
		clear!(param)
	end
	if param.sym==0
		if doTranspose==1 && param.isTransposed==0
			clear!(param);
			A = sparse(A');
			param.isTransposed = 1;
		end
		if doTranspose==0 && param.isTransposed==1
			clear!(param);
		end
	end
	if param.Ainv == []
		param = setupSolver(A,param);
	end

	tt = time_ns()
	U = param.Ainv\B;
	param.solveTime+=(tt-time_ns())/1e+9;
	param.nSolve+=1

	return U, param
end # function solveLinearSystem

function clear!(param::JuliaSolver)
	param.Ainv = [];
	param.isTransposed = 0;
	param.doClear = 0;
end

function copySolver(Ainv::JuliaSolver)
	return getJuliaSolver(sym = Ainv.sym,doClear = Ainv.doClear);
end


text:
# include(pwd()*"/src/SFFM.jl")
using LinearAlgebra, Plots, JLD2, SFFM
#
# ## define the model(s)
include(pwd()*"/examples/meNumerics/discontinuitiesModelDef.jl")

## simulate Ψ paths for comparison
innerNSim = 10^4
NSim = 10^2 * innerNSim

# # we can do this in parallell
using SharedArrays, Distributed
nprocs() < 5 && addprocs(5 - nprocs())
@everywhere include(pwd()*"/src/SFFM.jl")
@everywhere include(pwd()*"/examples/meNumerics/discontinuitiesModelDef.jl")
simsOuter_Psi = SharedArray(zeros(NSim, 5))
simsOuter_1 = SharedArray(zeros(NSim, 5))

t = 10.2

@time @sync @distributed for n = 1:(NSim÷innerNSim)
    IC = (φ = 1 .* ones(Int, innerNSim), X = zeros(innerNSim), Y = zeros(innerNSim))
    simsInner_Psi = SFFM.SimSFFM(
        model,
        SFFM.FirstExitY(0, Inf),
        IC,
    )
    simsOuter_Psi[1+(n-1)*innerNSim:n*innerNSim, :] =
        [simsInner_Psi.t simsInner_Psi.φ simsInner_Psi.X simsInner_Psi.Y simsInner_Psi.n]
    simsInner_1 = SFFM.SimSFFM(
        model,
        SFFM.FixedTime(t),
        IC,
    )
    simsOuter_1[1+(n-1)*innerNSim:n*innerNSim, :] =
        [simsInner_1.t simsInner_1.φ simsInner_1.X simsInner_1.Y simsInner_1.n]
end

sims_Psi = (
    t = simsOuter_Psi[:, 1],
    φ = simsOuter_Psi[:, 2],
    X = simsOuter_Psi[:, 3],
    Y = simsOuter_Psi[:, 4],
    n = simsOuter_Psi[:, 5],
)

sims_1 = (
    t = simsOuter_1[:, 1],
    φ = simsOuter_1[:, 2],
    X = simsOuter_1[:, 3],
    Y = simsOuter_1[:, 4],
    n = simsOuter_1[:, 5],
)

@save pwd()*"/examples/meNumerics/discontinuitiesModelSims_t_"*string(t)*".jld2" sims_Psi sims_1 model
interrupt()

max_stars_repo_name: angus-lewis/SFFM

max_issues_repo_name: angus-lewis/SFFM

max_forks_repo_name: angus-lewis/SFFM


text:
using Optics
# using RefractiveIndex
using Base.Test

# write your own tests here
# materials = Material[]
push!(materials, getdata("main", "au", "babar")) # tabulated nk
push!(materials, getdata("glass", "ami", "amtir-2")) # formula 1 / tabulated k
push!(materials, getdata("glass", "schott", "k7")) # formula 2 / tabulated k
push!(materials, getdata("glass", "hoya", "bac4")) # formula 3 / tabulated k
push!(materials, getdata("main", "Ag3AsS3", "Hulme-e.yml")) # formula 4
push!(materials, getdata("glass", "misc", "soda-lime/Rubin-bronze.yml")) # formula 5
push!(materials, getdata("main", "Ar", "Bideau-Mehu.yml")) # formula 6
push!(materials, getdata("main", "Si", "Edwards.yml")) # formula 7
push!(materials, getdata("main", "AgBr", "Schroter.yml")) # formula 8
push!(materials, getdata("organic", "CH4N2O - urea", "Rosker-e.yml")) # formula 9

# write your own tests here
@testset "Reflection Test" begin
    @test round(snellius(1.0, 1.5, deg2rad(60.0)), 5) == 0.61548
    @test reflection(materials[1], 4.6, deg2rad(60.0)) .* 1e5 .|> round == (99748, 98997)
    Φ = phaseangle(materials[1], 4.6, deg2rad(60.0))
    @test round.(rad2deg.(Φ), 3) == (-178.184, 7.261)
end

@testset "Transmittance Test" begin
    transmittance(n\underbar, d)
end


text:
export scan_shots

function scan_shots!(s::IO, mem_chunk::Int, mem_trace::Int,
                    keys::Array{String,1}, file::String, scan::Array{BlockScan,1}, fl_eof::Bool)

    # Load chunk into memory
    chunk_start = position(s)
    buf = IOBuffer(read(s, mem_chunk))
    eof(s) ? (fl_eof=true) : nothing
    buf_size = position(seekend(buf)); seekstart(buf)
    ntraces = Int(floor(buf_size/mem_trace))
    headers = Array{BinaryTraceHeader,1}(undef, ntraces)

    # Get headers from chunk
    for i in 1:ntraces
        headers[i] = read_traceheader(buf, keys, SegyIO.th_b2s)
        skip(buf, mem_trace-240)
    end

    # Get all requested header vectors
    vals = Dict{String, Array{Int32,1}}()
    for k in keys
        tmp = Int32.([getfield((headers[i]), Symbol(k)) for i in 1:ntraces])
        vals["$k"] = tmp
    end # k

    # Deliminate to find shots
    sx = vals["SourceX"]
    sy = vals["SourceY"]
    #combo = [[view(sx,i) view(sy,i)] for i in 1:ntraces]
    combo = [[sx[i] sy[i]] for i in 1:ntraces]
    part = delim_vector(combo,1)
    fl_eof ? push!(part, length(combo) + 1) : push!(part, ntraces + 1)

    # Summarise each shot
    for shot in 1:length(part)-1
        start_trace = part[shot]
        end_trace = part[shot + 1]-1
        start_byte = mem_trace*(start_trace-1) + chunk_start
        end_byte = mem_trace*(end_trace) + chunk_start
        summary = Dict{String, Array{Int32,1}}()

        for k in keys
            tmp = vals[k][start_trace:end_trace]
            summary["$k"] = [minimum(tmp); maximum(tmp)]
        end

        push!(scan, BlockScan(file, start_byte, end_byte, summary))
    end

    seek(s, scan[end].endbyte)
    close(buf)
    nothing
end


text:
module MPBUtils
# ---------------------------------------------------------------------------------------- #

using Crystalline
using Crystalline: AbstractFourierLattice,
                   matching_littlegroups,
                   DEFAULT_ATOL,
                   TEST_αβγ
using StaticArrays
using Statistics: quantile # for `filling2isoval`
using DocStringExtensions
using Requires
using DelimitedFiles

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

export prepare_mpbcalc,
       prepare_mpbcalc!,
       mpb_calcname,
       write_lgs_to_mpb!,
       lattice_from_mpbparams,
       kvecs_from_mpbparams,
       filling2isoval,
       isoval2filling

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

include("utils.jl")
include("filling2isoval.jl")
include("export2mpb.jl")
include("read_utils.jl")
export read_symdata, extract_multiplicities, extract_individual_multiplicities,
       collect_separable, pick_lgirreps

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

function __init__()
    # plotting utitilities when PyPlot is loaded
    @require PyPlot="d330b81b-6aea-500a-939a-2ce795aea3ee" begin
        include("compat/pyplot.jl")
        export plot_lattice_from_mpbparams
    end
end

# ---------------------------------------------------------------------------------------- #
end


text:
module MathepiaModels
using DifferentialEquations
using Plots

include("Models/odecompartmentsmodel.jl")

export SIRbasic
export SEIRbasic
export SEIARbasic
export SISbasic

end


text:
using Base.Test
rbc = Model(:[c+(k-(1-δ)*k[-1])-Y
			  ϕh*(1-h)^-σh/c^-σc-W
		      k-β*Expect(λ[+1]*R[+1]*k)/λ
],:[    k=(3,12,6)
        η = (1,0.9,0.01,2)
],:[    c = (0,10,.7),
        h = (0,1,0.35),
        k = (0,100,9.5)
],:[    λ = c^-σc
		Y = η*k[-1]^α*h^(1-α)
		W = (1-α)*k[-1]^α*η*h^-α
		R = (1-δ+α*η*(k[-1]/h)^(α-1))
],:[    α = 0.33
		δ = 0.025
		σc = 2.5
		ϕh = 1.66
		σh = 2
		β = 0.99])
solve(rbc)

@test abs(rbc.G(rbc[:k,0]-rbc[:k,-1],[9.575 1]))[1]<1e-3


text:
# caching of neighbor_index calls with site integer and direction integer
type Cached <: Indexing
    indexing::Indexing
    cache::Matrix
end
function Cached(kern::Indexing, lattice::AbstractLattice)
    const I = typeof(kern).parameters[1]
    result = Cached(kern, zeros(I, (length(lattice), length(kern))))
    for site in 1:length(kern), d in 1:length(lattice)
        result.cache[d, site] = neighbor_index(kern, site, lattice, d)
    end
    result
end
function Cached(kern::Cartesian, lattice::AbstractLattice)
    const I = typeof(kern).parameters[1]
    result = Cached(kern, zeros(I, (length(lattice), length(kern))))
    for site in 1:length(kern), d in 1:length(lattice)
        const coords = gridcoords(kern, site)
        if all(coords .> 1) && all(coords .< size(kern))
            result.cache[d, site] = neighbor_index(kern, site, lattice, d)
        end
    end
    result
end

function cached_cartesian(dims::GridCoords, lattice::AbstractLattice)
    if length(dims) != size(lattice.celerities, 1)
        error("Grid and lattice dimensionality do not match")
    end
    Cached(Cartesian(dims), lattice)
end
function cached_periodic(dims::GridCoords, lattice::AbstractLattice)
    if length(dims) != size(lattice.celerities, 1)
        error("Grid and lattice dimensionality do not match")
    end
    Cached(Periodic(dims), lattice)
end

index(kern::Cached, indices::GridCoords) = index(kern.indexing, indices)
neighbor_index(kern::Cached, site::Integer, ::AbstractLattice, d::Integer) = kern.cache[d, site]
function neighbor_index(kern::Cached, site::GridCoords, ::AbstractLattice, d::Integer)
    kern.cache[d, index(kern.indexing, site)]
end

length(kernel::Cached) = length(kernel.indexing)
size(kernel::Cached) = size(kernel.indexing)
gridcoords(kernel::Cached, index::Integer) = gridcoords(kernel.indexing, index)
function index{I <: Integer}(kernel::Cached, indices::Vararg{I})
    index(kernel.indexing, indices)
end


text:
### samplingSolver
include("revampedLinkedListFloatStorage.jl")
export llsInit
export llsAdd
export llsPurge

include("sqLinOpWrapper.jl")
export SqLinOp

include("samplingSolver.jl")
export samplingSDDSolver
export samplingLapSolver
export samplingParams

### KMP
include("KMPSolver.jl")
export KMPSDDSolver
export KMPLapSolver
export KMPParams

### hybridSolver
include("hybridSolver.jl")
export hybridSDDSolver
export hybridLapSolver
export hybridParams

### useful misc
include("condNumber.jl")
export condNumber

include("powerIteration.jl")
export powerIteration

"""A list containing SDD linear system solvers. They take in a SDD matrix plus tol, maxits and maxtime parameters."""
SDDSolvers = [augTreeSolver, KMPSDDSolver, hybridSDDSolver, samplingSDDSolver, AMGSolver]

"""A list containing Laplacian linear system solvers. They take in an adjacency matrix plus tol, maxits and maxtime parameters."""
LapSolvers = [augTreeLapSolver, KMPLapSolver, hybridLapSolver, samplingLapSolver, AMGLapSolver]


text:
function _gml_read_one_graph(gs, dir)
    nodes = [x[:id] for x in gs[:node]]
    if dir
        g = DiGraph(length(nodes))
    else
        g = Graph(length(nodes))
    end
    mapping = Dict{Int,Int}()
    for (i,n) in enumerate(nodes)
        mapping[n] = i
    end
    sds = [(Int(x[:source]), Int(x[:target])) for x in gs[:edge]]
    for (s,d) in (sds)
        add_edge!(g, mapping[s], mapping[d])
    end
    return g
end
function loadgml(io::IO, gname::AbstractString)
    p = GML.parse_dict(readall(io))
    for gs in p[:graph]
        dir = Bool(get(gs, :directed, 0))
        graphname = get(gs, :label, dir ? "digraph" : "graph")

        (gname == graphname) && return _gml_read_one_graph(gs, dir)
    end
    error("Graph $gname not found")
end

function loadgml_mult(io::IO)
    p = GML.parse_dict(readall(io))
    graphs = Dict{AbstractString, SimpleGraph}()
    for gs in p[:graph]
        dir = Bool(get(gs, :directed, 0))
        graphname = get(gs, :label, dir ? "digraph" : "graph")
        graphs[graphname] = _gml_read_one_graph(gs, dir)
    end
    return graphs
end

"""
savegml(f, g, gname = "graph")

Writes a graph `g` with name `gname`
to a file `f` in the
[GML](https://en.wikipedia.org/wiki/Graph_Modelling_Language) format.
"""
function savegml(io::IO, g::SimpleGraph, gname::AbstractString = "")
    println(io, "graph")
    println(io, "[")
    length(gname) > 0 && println(io, "label \"$gname\"")
    is_directed(g) && println(io, "directed 1")
    for i=1:nv(g)
        println(io,"\tnode")
        println(io,"\t[")
        println(io,"\t\tid $i")
        println(io,"\t]")
    end
    for (s,t) in edges(g)
        println(io,"\tedge")
        println(io,"\t[")
        println(io,"\t\tsource $s")
        println(io,"\t\ttarget $t")
        println(io,"\t]")
    end
    println(io, "]")
    return 1
end

"""Writes a dictionary of (name=>graph) to a file `fn`

Returns number of graphs written.
"""
function savegml_mult(io::IO, graphs::Dict)
    ng = 0
    for (gname, g) in graphs
        ng += savegml(io, g, gname)
    end
    return ng
end

filemap[:gml] = (loadgml, loadgml_mult, savegml, savegml_mult)


text:
function handle_macro(@nospecialize(x), state) end
function handle_macro(x::EXPR, state)
    typof(x) !== MacroCall && return
    if typof(x[1]) === MacroName
        state(x[1])
        if _points_to_Base_macro(x[1], :deprecate, state) && length(x) == 3
            if bindingof(x[2]) !== nothing
                return
            elseif CSTParser.is_func_call(x[2])
                # add deprecated method
                # add deprecated function binding and args in new scope
                mark_binding!(x[2], x)
                mark_sig_args!(x[2])
                s0 = state.scope # store previous scope
                state.scope = Scope(s0, x, Dict(), nothing, false)
                setscope!(x, state.scope) # tag new scope to generating expression
                state(x[2])
                state(x[3])
                state.scope = s0
            elseif isidentifier(x[2])
                mark_binding!(x[2], x)
            end
        elseif _points_to_Base_macro(x[1], :enum, state)
            for i = 2:length(x)
                if !(typof(x[i]) === PUNCTUATION)
                    if bindingof(x[i]) !== nothing
                        break
                    end
                    if i == 3 && typof(x[i]) === CSTParser.Begin
                        for j in 1:length(x[i][2])
                            mark_binding!(x[i][2][j], x)
                        end
                        break
                    end
                    mark_binding!(x[i], x)
                end
            end
        elseif _points_to_Base_macro(x[1], :goto, state)
            if length(x) == 2 && typof(x[2]) === CSTParser.IDENTIFIER
                setref!(x[2], Binding(noname, nothing, nothing, EXPR[], nothing, nothing))
            end
        elseif _points_to_Base_macro(x[1], :label, state)
            if length(x) == 2 && typof(x[2]) === CSTParser.IDENTIFIER
                mark_binding!(x[2])
            end
        elseif length(x[1]) == 2 && isidentifier(x[1][2]) && valof(x[1][2]) == "nospecialize"
            for i = 2:length(x)
                if !(typof(x[i]) === PUNCTUATION)
                    if bindingof(x[i]) !== nothing
                        break
                    end
                    mark_binding!(x[i], x)
                end
            end
        elseif _points_to_arbitrary_macro(x[1], :Turing, :model, state) && length(x) == 2 &&
            _binary_assert(x[2], CSTParser.Tokens.EQ) &&
            _expr_assert(x[2][3], CSTParser.Begin, 3) && typof(x[2][3][2]) === CSTParser.Block
            for i = 1:length(x[2][3][2])
                ex = x[2][3][2][i]
                if typof(ex) == CSTParser.BinaryOpCall && kindof(ex[2]) === CSTParser.Tokens.APPROX
                    mark_binding!(ex)
                end
            end
        elseif _points_to_arbitrary_macro(x[1], :JuMP, :variable, state)
            if length(x) < 3
                return
            elseif length(x) >= 5 && typof(x[2]) === PUNCTUATION
                _mark_JuMP_binding(x[5])
            else
                _mark_JuMP_binding(x[3])
            end
        elseif (_points_to_arbitrary_macro(x[1], :JuMP, :expression, state) ||
            _points_to_arbitrary_macro(x[1], :JuMP, :NLexpression, state) ||
            _points_to_arbitrary_macro(x[1], :JuMP, :constraint, state) || _points_to_arbitrary_macro(x[1], :JuMP, :NLconstraint, state)) && length(x) > 1
            if typof(x[2]) === PUNCTUATION
                if length(x) == 8
                    _mark_JuMP_binding(x[5])
                end
            else
                if length(x) == 4
                    _mark_JuMP_binding(x[3])
                end
            end
        end
    elseif typof(x[1]) == CSTParser.GlobalRefDoc && length(x) == 3 && CSTParser.isliteral(x[2]) && kindof(x[2]) === CSTParser.Tokens.TRIPLE_STRING && isidentifier(x[3])
        mark_binding!(x[3])
        setref!(x[3], bindingof(x[3]))
    end
end

function _rem_ref(x::EXPR)
    if typof(x) === CSTParser.Ref && length(x) > 0
        return x[1]
    end
    return x
end

function _mark_JuMP_binding(arg)
    if CSTParser.isidentifier(arg) || typof(arg) === CSTParser.Ref
        mark_binding!(_rem_ref(arg))
    elseif _binary_assert(arg, CSTParser.Tokens.EQEQ) || _binary_assert(arg, CSTParser.Tokens.LESS_EQ)  || _binary_assert(arg, CSTParser.Tokens.GREATER_EQ)
        if CSTParser.isidentifier(arg[1]) || typof(arg[1]) === CSTParser.Ref
            mark_binding!(_rem_ref(arg[1]))
        else
            mark_binding!(_rem_ref(arg[3]))
        end
    elseif typof(arg) === CSTParser.Comparison && length(arg) == 5
        mark_binding!(_rem_ref(arg[3]))
    end
end

function _points_to_Base_macro(x::EXPR, name, state)
    length(x) == 2 && isidentifier(x[2]) && Symbol(valof(x[2])) == name && refof(x[2]) == getsymbolserver(state.server)[:Base][Symbol("@", name)]
end

function _points_to_arbitrary_macro(x::EXPR, module_name, name, state)
    length(x) == 2 && isidentifier(x[2]) && valof(x[2]) == name && haskey(getsymbolserver(state.server), Symbol(module_name)) && haskey(getsymbolserver(state.server)[Symbol(module_name)], Symbol("@", name)) && refof(x[2]) == getsymbolserver(state.server)[Symbol(module_name)][Symbol("@", name)]
end


text:
"""
$(README)

# Exports
$(EXPORTS)
"""
module SpectralDistances

using LinearAlgebra, SparseArrays, Statistics, Printf
using DSP,
    ControlSystems,
    Distances,
    DocStringExtensions,
    DoubleFloats,
    FillArrays,
    Hungarian,
    Lazy,
    LoopVectorization,
    Optim,
    OrdinaryDiffEq,
    PolynomialRoots,
    ProgressMeter,
    QuadGK,
    RecipesBase,
    Roots,
    StaticArrays,
    StatsBase,
    ThreadTools,
    TotalLeastSquares,
    ZygoteRules,
    UnbalancedOptimalTransport,
    UnPack

using SlidingDistancesBase
import SlidingDistancesBase: floattype, lastlength, distance_profile, distance_profile!

import Base.@kwdef
using Base.Threads: nthreads, @threads, threadid

export ls,
    plr,
    polar,
    toreim,
    reflect,
    hungariansort,
    coefficients,
    s1,
    v1,
    v1!,
    n1,
    roots,
    pole,
    twoD,
    threeD,
    precompute,
    assignmentplot,
    roots2poly,
    evalfr,
    Log,
    Identity,
    residues,
    residueweight,
    simplex_residueweight,
    unitweight,
    normalize_energy,
    spectralenergy,
    normalization_factor,
    polyconv,
    domain,
    fitmodel,
    move_real_poles,
    checkroots,
    change_precision,
    embedding,
    distmat,
    distmat_euclidean,
    distmat_euclidean!,
    symmetrize!

export evaluate,
distance_profile,
AbstractDistance,
AbstractRationalDistance,
AbstractSignalDistance,
SymmetricDistance,
CoefficientDistance,
ModelDistance,
EuclideanRootDistance,
HungarianRootDistance,
OptimalTransportRootDistance,
KernelWassersteinRootDistance,
DiscretizedRationalDistance,
WelchOptimalTransportDistance,
WelchLPDistance,
ConvOptimalTransportDistance,
SlidingConvOptimalTransportDistance,
RationalOptimalTransportDistance,
OptimalTransportHistogramDistance,
DiscreteGridTransportDistance,
RationalCramerDistance,
BuresDistance,
EnergyDistance,
CompositeDistance,
discrete_grid_transportplan,
discrete_grid_transportcost

export sinkhorn,
sinkhorn_log,
sinkhorn_log!,
IPOT,
sinkhorn_diff,
sinkhorn_unbalanced,
transport_plan,
sinkhorn_convolutional,
SCWorkspace,
normalize_spectrogram,
normalize_spectrogram!,
mask_filter,
mask_filter!

using UnbalancedOptimalTransport: KL, TV, Balanced, RG
export KL, TV, Balanced, RG

export FitMethod,
AbstractModel,
AR,
ARMA,
PLR,
LS,
TLS,
IRLS,
examplemodels

export TimeEvolution,
Continuous,
Discrete,
SortAssignement,
HungarianAssignement,
DiscreteRoots,
ContinuousRoots,
domain_transform

export interpolator, barycenter, ISA, barycentric_coordinates, barycentric_weighting, barycenter_convolutional, BCWorkspace, BCCWorkspace

export TimeDistance, TimeVaryingAR, TimeWindow

include("eigenvalue_manipulations.jl")
include("ltimodels.jl")
include("losses.jl")
include("interpolations.jl")
include("utils.jl")
include("sinkhorn.jl")
include("barycenter.jl")
include("time.jl")
include("plotting.jl")
include("adjoints.jl")

import Requires
function __init__()
    Requires.@require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
        @eval using BackwardsLinalg
        @inline nograd(x) = Zygote.dropgrad(x)
    end

    Requires.@require GLPK = "60bf3e95-4087-53dc-ae20-288a0d20c6a6" begin
        Requires.@require JuMP = "4076af6c-e467-56ae-b986-b466b2749572" begin
            export ot_jump
            include("jump.jl")
        end
        Requires.@require Convex = "f65535da-76fb-5f13-bab9-19810c17039a" begin
            export ot_convex
            include("convex.jl")
        end
    end

    Requires.@require SCS = "c946c3f1-0d1f-5ce8-9dea-7daa1f7e2d13" begin
        Requires.@require Convex = "f65535da-76fb-5f13-bab9-19810c17039a" begin
            export complete_distmat
            include("complete_distmat.jl")
        end
    end

    Requires.@require Clustering = "aaaa29a8-35af-508c-8bc3-b662a17a0fe5" begin
        export kbarycenters
        include("kbarycenters.jl")
    end

    Requires.@require DynamicAxisWarping = "aaaaaaaa-4a10-5553-b683-e707b00e83ce" begin
        include("dtw.jl")
    end
end
end # module


text:
# For reading information & metadata about the file

function _get_fc(file::String) # convenience function for `get_duration` and `get_start_time`
    v = open(file)
    return v.format_context, v
end
get_duration(fc::AVFormatContextPtr) = fc.duration / AV_TIME_BASE

"""
    get_duration(file::String) -> Float64

Return the duration of the video `file` in seconds (float).
"""
function get_duration(file::String)
    fc, v = _get_fc(file)
    duration = get_duration(fc)
    close(v)
    return duration
end

get_start_time(fc::AVFormatContextPtr) = Dates.unix2datetime(fc.start_time_realtime)

"""
    get_start_time(file::String) -> DateTime

Return the starting date & time of the video `file`. Note that if the starting date & time are missing, this function will return the Unix epoch (00:00 1st January 1970).
"""
function get_start_time(file::String)
    fc, v = _get_fc(file)
    starttime = get_start_time(fc)
    close(v)
    return starttime
end

get_time_duration(fc::AVFormatContextPtr) = (get_start_time(fc), get_duration(fc))

"""
    get_time_duration(file::String) -> (DateTime, Microsecond)

Return the starting date & time as well as the duration of the video `file`. Note that if the starting date & time are missing, this function will return the Unix epoch (00:00 1st January 1970).
"""
function get_time_duration(file::String)
    fc, v = _get_fc(file)
    starttime, duration = get_time_duration(fc)
    close(v)
    return starttime, duration
end

"""
    get_number_frames(file [, streamno])

Query the the container `file` for the number of frames in video stream
`streamno` if applicable, instead returning `nothing` if the container does not
report the number of frames. Will not decode the video to count the number of
frames in a video.
"""
function get_number_frames(file::AbstractString, streamno::Integer = 0)
    streamno >= 0 || throw(ArgumentError("streamno must be non-negative"))
    frame_strs = FFMPEG.exe(`-v error -select_streams v:$(streamno) -show_entries
                            stream=nb_frames -of
                            default=nokey=1:noprint_wrappers=1 $file`,
                           command = FFMPEG.ffprobe, collect = true)
    frame_str = frame_strs[1]
    if occursin("No such file or directory", frame_str)
        error("Could not find file $file")
    elseif occursin("N/A", frame_str)
        return nothing
    end
    return parse(Int, frame_str)
end


text:
# MIT License

# Copyright (c) Microsoft Corporation.

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE

using Documenter
using OpticSim

makedocs(
    sitename = "OpticSim.jl",
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
        assets = [asset("assets/logo.svg", class = :ico, islocal = true)],
    ),
    modules = [OpticSim],
    pages = [
        "Home" => "index.md",
        "Examples" => "examples.md",
        # "Glasses" => cat_pages,
        "Geometry" => [
            "Primitives" => "primitives.md",
            "CSG" => "csg.md"
        ],
        "Optical" => [
            "Systems" => "systems.md",
            "Emitters" => "emitters.md",
            "Interfaces" => "interfaces.md",
            "Lenses" => "lenses.md"
        ],
        "Visualization" => "vis.md",
        "Glass Functions" => "glasscat.md",
        "Optimization" => "optimization.md",
        "Reference" => "ref.md",
        "Roadmap" => "roadmap.md"
    ],
    expandfirst = ["glasscat.md", "systems.md", "vis.md"])

deploydocs(
    repo = "github.com/microsoft/OpticSim.jl.git",
    devbranch = "main",
)

# function children(m::Module)
#     ns = names(m, imported = false, all = true)
#     ms = []

#     for n in ns
#         try
#             x = Core.eval(m, n)
#             if x isa Module
#                 if(x != OpticSim.GlassCat)
#                 println("x $x")
#                 push!(ms, x)
#             end
#         end
#         catch
#         end
#     end
#     return ms
# end

#WARNING: I think this code creates a type for each glass name, which overwrites the definition in the src files where each glass name corresponds to an integer indexing into a glass table. Obviously nothing works after this. Needs major surgery to generate glass documentation at the same time as the documentation for everything else.

# # write a code file for the catalog with docstrings
# io = open(joinpath(@__DIR__, "../src/AGFGlassCatDocs.jl"), "w")
# catalogs = children(OpticSim.GlassCat)
# cat_pages = []
# for catname in catalogs
#     println(catalogs)
#     eval_string = ["module $(nameof(catname))"]
#     escape_catalog_name = replace(string(nameof(catname)), "_" => "\\_")
#     push!(cat_pages, escape_catalog_name => "$(nameof(catname)).md")
#     iomd = open(joinpath(@__DIR__, "src/$(nameof(catname)).md"), "w")
#     write(iomd, "# $escape_catalog_name\n\n")
#     write(iomd, "```@raw html\n")
#     write(iomd, "<style>article p {display:flex;justify-content:space-between;}</style>\n")
#     write(iomd, "```\n")
#     write(iomd, "```@docs\n")
#     catalog_module = eval(catname)
#     glass_names = names(catalog_module, all = true, imported = false)
#     for glass_name in glass_names
#         glass_name_str = string(glass_name)
#         if !occursin("#", glass_name_str) && glass_name_str != "eval" && glass_name_str != "include" && glass_name != nameof(catname) && !in(glass_name,(:MODEL, :MIL, :AGF, :OTHER, :AIR))
#             glass = Core.eval(catalog_module, glass_name)

#             temp = typeof(glass)
# println("type of glass $temp")

#             let io = IOBuffer()
#                 OpticSim.GlassCat.docstring(io, glass)
#                 infostr = String(take!(io))
#                 push!(eval_string, "\"\"\"\n$infostr\n\"\"\"")
#             end
#             push!(eval_string, "function $glass_name_str()\nend")
#             write(iomd, "$catname.$glass_name_str\n")
#         end
#     end
#     write(iomd, "```\n")
#     push!(eval_string, "end") # module
#     eval_string = join(eval_string, "\n")
#     write(io, eval_string * "\n")
#     close(iomd)
# end
# close(io)

# # include source with docstrings
# OpticSim.GlassCat.include(joinpath(@__DIR__, "../src/AGFGlassCatDocs.jl"))

# clean up
# rm(joinpath(@__DIR__, "../src/AGFGlassCatDocs.jl"))
# for p in cat_pages
#     rm(joinpath(@__DIR__, "src/" * p[2]))
# end


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:5] <= 1)

                     @variable(m, -6.488824413359986 <= q <= 6.729439090743391)

                     add_NL_constraint(m, :(1/(1 + exp(-(0.6442211513786371 + 0.6060607309092267*1/(1 + exp(-(-0.25461520796754256 + 0.8980721543089638*$(x[1]) + 0.9381266822727681*$(x[2]) + -0.8446276014249228*$(x[3]) + -0.43741343386290765*$(x[4]) + 0.17400967344335916*$(x[5])))) + -0.2689278778257371*1/(1 + exp(-(-0.32405265497274005 + 0.7689511041538082*$(x[1]) + 0.7170113596676764*$(x[2]) + -0.7954301950243226*$(x[3]) + 0.3706953387574874*$(x[4]) + -0.13209825476725356*$(x[5]))))))) + 1/(1 + exp(-(-0.8503398678327141 + -0.43763464280202813*1/(1 + exp(-(-0.25461520796754256 + 0.8980721543089638*$(x[1]) + 0.9381266822727681*$(x[2]) + -0.8446276014249228*$(x[3]) + -0.43741343386290765*$(x[4]) + 0.17400967344335916*$(x[5])))) + -0.8707322759459402*1/(1 + exp(-(-0.32405265497274005 + 0.7689511041538082*$(x[1]) + 0.7170113596676764*$(x[2]) + -0.7954301950243226*$(x[3]) + 0.3706953387574874*$(x[4]) + -0.13209825476725356*$(x[5]))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
#### Specialized matrix types ####

## Hermitian tridiagonal matrices
immutable SymTridiagonal{T} <: AbstractMatrix{T}
    dv::Vector{T}                        # diagonal
    ev::Vector{T}                        # subdiagonal
    function SymTridiagonal(dv::Vector{T}, ev::Vector{T})
        length(ev)==length(dv)-1 || throw(DimensionMismatch(""))
        new(dv,ev)
    end
end

SymTridiagonal{T}(dv::Vector{T}, ev::Vector{T}) = SymTridiagonal{T}(dv, ev)

function SymTridiagonal{Td,Te}(dv::Vector{Td}, ev::Vector{Te})
    T = promote_type(Td,Te)
    SymTridiagonal(convert(Vector{T}, dv), convert(Vector{T}, ev))
end

SymTridiagonal(A::AbstractMatrix) = diag(A,1)==diag(A,-1)?SymTridiagonal(diag(A), diag(A,1)):throw(DimensionMismatch("matrix is not symmetric; cannot convert to SymTridiagonal"))
full{T}(M::SymTridiagonal{T}) = convert(Matrix{T}, M)
convert{T}(::Type{SymTridiagonal{T}}, S::SymTridiagonal) = SymTridiagonal(convert(Vector{T}, S.dv), convert(Vector{T}, S.ev))
convert{T}(::Type{AbstractMatrix{T}}, S::SymTridiagonal) = SymTridiagonal(convert(Vector{T}, S.dv), convert(Vector{T}, S.ev))
convert{T}(::Type{Matrix{T}}, M::SymTridiagonal{T})=diagm(M.dv)+diagm(M.ev,-1)+conj(diagm(M.ev,1))

size(m::SymTridiagonal) = (length(m.dv), length(m.dv))
size(m::SymTridiagonal, d::Integer) = d<1 ? error("dimension out of range") : (d<=2 ? length(m.dv) : 1)

#Elementary operations
for func in (:copy, :round, :iround, :conj)
    @eval begin
        ($func)(M::SymTridiagonal) = SymTridiagonal(($func)(M.dv), ($func)(M.ev))
    end
end

transpose(M::SymTridiagonal) = M #Identity operation
ctranspose(M::SymTridiagonal) = conj(M)

function diag{T}(M::SymTridiagonal{T}, n::Integer=0)
    absn = abs(n)
    absn==0 ? M.dv : absn==1 ? M.ev : absn<size(M,1) ? zeros(T,size(M,1)-absn) : throw(BoundsError())
end

+(A::SymTridiagonal, B::SymTridiagonal) = SymTridiagonal(A.dv+B.dv, A.ev+B.ev)
-(A::SymTridiagonal, B::SymTridiagonal) = SymTridiagonal(A.dv-B.dv, A.ev-B.ev)
*(A::SymTridiagonal, B::Number) = SymTridiagonal(A.dv*B, A.ev*B)
*(B::Number, A::SymTridiagonal) = A*B
/(A::SymTridiagonal, B::Number) = SymTridiagonal(A.dv/B, A.ev/B)
==(A::SymTridiagonal, B::SymTridiagonal) = (A.dv==B.dv) && (A.ev==B.ev)

## Solver
function \{T<:BlasFloat}(M::SymTridiagonal{T}, rhs::StridedVecOrMat{T})
    if stride(rhs, 1) == 1
        return LAPACK.gtsv!(copy(M.ev), copy(M.dv), copy(M.ev), copy(rhs))
    end
    solve(Tridiagonal(M), rhs)  # use the Julia "fallback"
end

#Wrap LAPACK DSTE{GR,BZ} to compute eigenvalues
eigfact!{T<:BlasFloat}(m::SymTridiagonal{T}) = Eigen(LAPACK.stegr!('V', m.dv, m.ev)...)
eigfact!{T<:BlasFloat}(m::SymTridiagonal{T}, irange::UnitRange) = Eigen(LAPACK.stegr!('V', 'I', m.dv, m.ev, 0.0, 0.0, irange.start, irange.stop)...)
eigfact!{T<:BlasFloat}(m::SymTridiagonal{T}, vl::Real, vu::Real) = Eigen(LAPACK.stegr!('V', m.dv, m.ev, convert(T, vl), convert(T, vu), 0, 0)...)
eigvals!{T<:BlasFloat}(m::SymTridiagonal{T}) = LAPACK.stegr!('N', 'A', m.dv, m.ev, 0.0, 0.0, 0, 0)[1]
eigvals!{T<:BlasFloat}(m::SymTridiagonal{T}, irange::UnitRange) = LAPACK.stegr!('N', 'I', m.dv, m.ev, 0.0, 0.0, irange.start, irange.stop)[1]
eigvals!{T<:BlasFloat}(m::SymTridiagonal{T}, vl::Real, vu::Real) = LAPACK.stegr!('N', 'V', m.dv, m.ev, vl, vu, 0, 0)[1]

#Computes largest and smallest eigenvalue
eigmax(m::SymTridiagonal) = eigvals(m, size(m, 1), size(m, 1))[1]
eigmin(m::SymTridiagonal) = eigvals(m, 1, 1)[1]

#Compute selected eigenvectors only corresponding to particular eigenvalues
eigvecs(m::SymTridiagonal) = eigfact(m)[:vectors]
eigvecs{T<:BlasFloat,Eigenvalue<:Real}(m::SymTridiagonal{T}, eigvals::Vector{Eigenvalue}) = LAPACK.stein!(m.dv, m.ev, eigvals)

###################
# Generic methods #
###################

#Needed for inv_usmani()
type ZeroOffsetVector
    data::Vector
end
getindex (a::ZeroOffsetVector, i) = a.data[i+1]
setindex!(a::ZeroOffsetVector, x, i) = a.data[i+1]=x

#Implements the inverse using the recurrence relation between principal minors
# a, b, c are assumed to be the subdiagonal, diagonal, and superdiagonal of
# a tridiagonal matrix.
#Ref:
#    R. Usmani, "Inversion of a tridiagonal Jacobi matrix",
#    Linear Algebra and its Applications 212-213 (1994), pp.413-414
#    doi:10.1016/0024-3795(94)90414-6
function inv_usmani{T}(a::Vector{T}, b::Vector{T}, c::Vector{T})
    n = length(b)
    θ = ZeroOffsetVector(zeros(T, n+1)) #principal minors of A
    θ[0] = 1
    n>=1 && (θ[1] = b[1])
    for i=2:n
        θ[i] = b[i]*θ[i-1]-a[i-1]*c[i-1]*θ[i-2]
    end
    φ = zeros(T, n+1)
    φ[n+1] = 1
    n>=1 && (φ[n] = b[n])
    for i=n-1:-1:1
        φ[i] = b[i]*φ[i+1]-a[i]*c[i]*φ[i+2]
    end
    α = Array(T, n, n)
    for i=1:n, j=1:n
        sign = (i+j)%2==0 ? (+) : (-)
        if i<j
            α[i,j]=(sign)(prod(c[i:j-1]))*θ[i-1]*φ[j+1]/θ[n]
        elseif i==j
            α[i,i]=                       θ[i-1]*φ[i+1]/θ[n]
        else #i>j
            α[i,j]=(sign)(prod(a[j:i-1]))*θ[j-1]*φ[i+1]/θ[n]
        end
    end
    α
end

#Implements the determinant using principal minors
#Inputs and reference are as above for inv_usmani()
function det_usmani{T}(a::Vector{T}, b::Vector{T}, c::Vector{T})
    n = length(b)
    θa = one(T)
    n==0 && return θa
    θb = b[1]
    for i=2:n
        θb, θa = b[i]*θb-a[i-1]*c[i-1]*θa, θb
    end
    return θb
end

inv(A::SymTridiagonal) = inv_usmani(A.ev, A.dv, A.ev)
det(A::SymTridiagonal) = det_usmani(A.ev, A.dv, A.ev)

function getindex{T}(A::SymTridiagonal{T}, i::Integer, j::Integer)
    (1<=i<=size(A,2) && 1<=j<=size(A,2)) || throw(BoundsError())
    i==j ? A.dv[i] : i==j+1 ? A.ev[j] : i+1==j ? A.ev[i] : zero(T)
end

## Tridiagonal matrices ##
immutable Tridiagonal{T} <: AbstractMatrix{T}
    dl::Vector{T}    # sub-diagonal
    d::Vector{T}     # diagonal
    du::Vector{T}    # sup-diagonal
    du2::Vector{T}   # supsup-diagonal for pivoting
end
function Tridiagonal{T}(dl::Vector{T}, d::Vector{T}, du::Vector{T})
    n = length(d)
    if (length(dl) != n-1 || length(du) != n-1)
        error(string("Cannot make Tridiagonal from incompatible lengths of subdiagonal, diagonal and superdiagonal: (", length(dl), ", ", length(d), ", ", length(du),")"))
    end
    Tridiagonal(dl, d, du, zeros(T,n-2))
end
function Tridiagonal{Tl, Td, Tu}(dl::Vector{Tl}, d::Vector{Td}, du::Vector{Tu})
    Tridiagonal(map(v->convert(Vector{promote_type(Tl,Td,Tu)}, v), (dl, d, du))...)
end

size(M::Tridiagonal) = (length(M.d), length(M.d))
size(M::Tridiagonal, d::Integer) = d<1 ? error("dimension out of range") : (d<=2 ? length(M.d) : 1)

full{T}(M::Tridiagonal{T}) = convert(Matrix{T}, M)
function convert{T}(::Type{Matrix{T}}, M::Tridiagonal{T})
    A = zeros(T, size(M))
    for i = 1:length(M.d)
        A[i,i] = M.d[i]
    end
    for i = 1:length(M.d)-1
        A[i+1,i] = M.dl[i]
        A[i,i+1] = M.du[i]
    end
    A
end
function similar(M::Tridiagonal, T, dims::Dims)
    if length(dims) != 2 || dims[1] != dims[2]
        throw(DimensionMismatch("Tridiagonal matrices must be square"))
    end
    Tridiagonal{T}(similar(M.dl), similar(M.d), similar(M.du), similar(M.du2))
end

# Operations on Tridiagonal matrices
copy!(dest::Tridiagonal, src::Tridiagonal) = Tridiagonal(copy!(dest.dl, src.dl), copy!(dest.d, src.d), copy!(dest.du, src.du), copy!(dest.du2, src.du2))

#Elementary operations
for func in (:copy, :round, :iround, :conj)
    @eval begin
        ($func)(M::Tridiagonal) = Tridiagonal(map(($func), (M.dl, M.d, M.du, M.du2))...)
    end
end

transpose(M::Tridiagonal) = Tridiagonal(M.du, M.d, M.dl)
ctranspose(M::Tridiagonal) = conj(transpose(M))

diag{T}(M::Tridiagonal{T}, n::Integer=0) = n==0 ? M.d : n==-1 ? M.dl : n==1 ? M.du : abs(n)<size(M,1) ? zeros(T,size(M,1)-abs(n)) : throw(BoundsError())
function getindex{T}(A::Tridiagonal{T}, i::Integer, j::Integer)
    (1<=i<=size(A,2) && 1<=j<=size(A,2)) || throw(BoundsError())
    i==j ? A.d[i] : i==j+1 ? A.dl[j] : i+1==j ? A.du[i] : zero(T)
end

###################
# Generic methods #
###################

+(A::Tridiagonal, B::Tridiagonal) = Tridiagonal(A.dl+B.dl, A.d+B.d, A.du+B.du)
-(A::Tridiagonal, B::Tridiagonal) = Tridiagonal(A.dl-B.dl, A.d-B.d, A.du+B.du)
*(A::Tridiagonal, B::Number) = Tridiagonal(A.dl*B, A.d*B, A.du*B)
*(B::Number, A::Tridiagonal) = A*B
/(A::Tridiagonal, B::Number) = Tridiagonal(A.dl/B, A.d/B, A.du/B)

==(A::Tridiagonal, B::Tridiagonal) = (A.dl==B.dl) && (A.d==B.d) && (A.du==B.du)
==(A::Tridiagonal, B::SymTridiagonal) = (A.dl==A.du==B.ev) && (A.d==B.dv)
==(A::SymTridiagonal, B::Tridiagonal) = (B.dl==B.du==A.ev) && (B.d==A.dv)

inv(A::Tridiagonal) = inv_usmani(A.dl, A.d, A.du)
det(A::Tridiagonal) = det_usmani(A.dl, A.d, A.du)

# Elementary operations that mix Tridiagonal and SymTridiagonal matrices
convert(::Type{Tridiagonal}, A::SymTridiagonal) = Tridiagonal(A.ev, A.dv, A.ev)
+(A::Tridiagonal, B::SymTridiagonal) = Tridiagonal(A.dl+B.ev, A.d+B.dv, A.du+B.ev)
+(A::SymTridiagonal, B::Tridiagonal) = Tridiagonal(A.ev+B.dl, A.dv+B.d, A.ev+B.du)
-(A::Tridiagonal, B::SymTridiagonal) = Tridiagonal(A.dl-B.ev, A.d-B.dv, A.du-B.ev)
-(A::SymTridiagonal, B::Tridiagonal) = Tridiagonal(A.ev-B.dl, A.dv-B.d, A.ev-B.du)

convert{T}(::Type{Tridiagonal{T}},M::Tridiagonal) = Tridiagonal(convert(Vector{T}, M.dl), convert(Vector{T}, M.d), convert(Vector{T}, M.du), convert(Vector{T}, M.du2))
convert{T}(::Type{AbstractMatrix{T}},M::Tridiagonal) = convert(Tridiagonal{T}, M)
convert{T}(::Type{Tridiagonal{T}}, M::SymTridiagonal{T}) = Tridiagonal(M)
convert{T}(::Type{SymTridiagonal{T}}, M::Tridiagonal) = M.dl==M.du ? (SymTridiagonal(M.dl, M.d)) :
    error("Tridiagonal is not symmetric, cannot convert to SymTridiagonal")
convert{T}(::Type{SymTridiagonal{T}},M::SymTridiagonal) = SymTridiagonal(convert(Vector{T}, M.dv), convert(Vector{T}, M.ev))

function A_mul_B!(C::AbstractVecOrMat, A::Tridiagonal, B::AbstractVecOrMat)
    size(C,1) == size(B,1) == (nA = size(A,1)) || throw(DimensionMismatch(""))
    size(C,2) == (nB = size(B,2)) || throw(DimensionMismatch(""))
    l = A.dl
    d = A.d
    u = A.du
    @inbounds begin
        for j = 1:nB
            C[1,j] = d[1]*B[1,j] + u[1]*B[2,j]
            for i = 2:nA-1
                C[i,j] = l[i-1]*B[i-1,j] + d[i]*B[i,j] + u[i]*B[i+1,j]
            end
            C[nA,j] = l[nA-1]*B[nA-1,j] + d[nA]*B[nA,j]
        end
    end
    C
end
*(A::Tridiagonal, B::AbstractVecOrMat) = A_mul_B!(similar(B), A, B)

A_ldiv_B!(A::Tridiagonal,B::AbstractVecOrMat) = A_ldiv_B!(lufact!(A), B)
At_ldiv_B!(A::Tridiagonal,B::AbstractVecOrMat) = At_ldiv_B!(lufact!(A), B)
Ac_ldiv_B!(A::Tridiagonal,B::AbstractVecOrMat) = Ac_ldiv_B!(lufact!(A), B)

max_stars_repo_name: porterjamesj/julia

max_issues_repo_name: porterjamesj/julia

max_forks_repo_name: porterjamesj/julia


text:
const s_star = 0.339
const sp = 0.612
const τ1 = 10000
const vi = PauliVec[3][2] ⊗ PauliVec[3][2] ⊗ PauliVec[3][2]

"""
    build_s(stage)

Build the annealing parameter `s` based on `stage`. The stage can be 1, 2 or 3, as defined in HOQST paper.
"""
function build_s(stage)
    if stage == 1
        spfun = (s) -> s <= 0.5 ? 1.0 : 2 * (sp - 1) * s + 2 - sp
        sfun = (s) -> s <= 0.5 ? 2 * (s_star - 1) * s + 1 : s_star
    elseif stage == 2
        spfun = (s) -> sp
        sfun = (s) -> s_star
    elseif stage == 3
        spfun = (s) -> s <= 0.5 ? 2 * (1 - sp) * s + sp : 1.0
        sfun = (s) -> s <= 0.5 ? s_star : 2 * (1 - s_star) * s + 2 * s_star - 1
    else
        throw(ArgumentError("Stage $stage does not exist."))
    end
    spfun, sfun
end

"""
    build_ds(stage)

Build the derivative of the annealing parameter `s` based on `stage`.
"""
function build_ds(stage)
    if stage == 1
        dspfun = (s) -> s <= 0.5 ? 0.0 : 2 * (sp - 1)
        dsfun = (s) -> s <= 0.5 ? 2 * (s_star - 1) : 0.0
    elseif stage == 2
        dspfun = (s) -> 0.0
        dsfun = (s) -> 0.0
    elseif stage == 3
        dspfun = (s) -> s <= 0.5 ? 2 * (1 - sp) : 0.0
        dsfun = (s) -> s <= 0.5 ? 0.0 : 2 * (1 - s_star)
    else
        throw(ArgumentError("Stage $stage does not exists."))
    end
    dspfun, dsfun
end

function build_Hlist(hp; polaron=false)
    polaron == true ?
    [-σi ⊗ σx ⊗ σi - σi ⊗ σi ⊗ σx, build_Hising(hp)] :
    [-σx ⊗ σi ⊗ σi, -σi ⊗ σx ⊗ σi - σi ⊗ σi ⊗ σx, build_Hising(hp)]
end

function build_Hising(hp)
    jp = -1.8
    js = -2.5
    hl = local_field_term([-hp, -jp, 0], [1, 2, 3], 3)
    h2 = two_local_term([jp, js], [[1, 2], [2, 3]], 3)
    hp = hl + h2
end

"""
    build_H(hp, stage; polaron=false)

Build the `stage` stage Hamiltonian given `hp`. `polaron` specifies whether the Hamiltonian is in polaron frame or not.
"""
function build_H(hp, stage; polaron=false)
    if stage == 1
        spfun, sfun = build_s(stage)
        Ap = (s) -> A(spfun(s))
        As = (s) -> A(sfun(s))
        Bs = (s) -> B(sfun(s))
        Hlist = build_Hlist(hp, polaron=polaron)
        polaron == false ? DenseHamiltonian([Ap, As, Bs], Hlist, unit=:ħ) : DenseHamiltonian([As, Bs], Hlist, unit=:ħ)
    elseif stage == 2
        hlist = build_Hlist(hp, polaron=polaron)
        hmat = polaron == false ? A(sp) * hlist[1] + A(s_star) * hlist[2] + B(s_star) * hlist[3] : A(s_star) * hlist[1] + B(s_star) * hlist[2]
        function build_user_eigen(u_cache)
            w, v = eigen(Hermitian(hmat))
            EIGS = function (H, t, lvl)
                w[1:lvl], v[:, 1:lvl]
            end
        end
        DenseHamiltonian([(s) -> 1.0], [hmat], unit=:ħ, EIGS=build_user_eigen)
    elseif stage == 3
        spfun, sfun = build_s(stage)
        Ap(s) = A(spfun(s))
        As(s) = A(sfun(s))
        Bs(s) = B(sfun(s))
        Hlist = build_Hlist(hp, polaron=polaron)
        polaron == false ? DenseHamiltonian([Ap, As, Bs], Hlist, unit=:ħ) : DenseHamiltonian([As, Bs], Hlist, unit=:ħ)
    else
        error("Stage $stage not defined.")
    end
end

"""
    build_dH(hp, stage; polaron=false)

Build the derivative of the `stage` stage Hamiltonian as a function.
"""
function build_dH(hp, stage; polaron=false)
    if stage in [1,2,3]
        spfun, sfun = build_s(stage)
        dspfun, dsfun = build_ds(stage)
        dAp(s) = dA(spfun(s)) * dspfun(s)
        dAs(s) = dA(sfun(s)) * dsfun(s)
        dBs(s) = dB(sfun(s)) * dsfun(s)
        Hlist = build_Hlist(hp, polaron=polaron)
        polaron == false ? (s) -> dAp(s) * Hlist[1] + dAs(s) * Hlist[2] + dBs(s) * Hlist[3] : (s) -> dAs(s) * Hlist[1] + dBs(s) * Hlist[2]
    else
        error("Stage $stage not defined.")
    end
end

"""
    init_lvl(hp)

Given `hp` value, check the initial energy level of the all one state.
"""
function init_lvl(hp)
    w, v = eigen(Hermitian(build_Hising(hp)))
    p = abs2.(v' * vi)
    ind = findfirst((x) -> x ≈ 1.0, p)
end

"""
    build_adiabatic_frame_u0(hp)

Build the adiabatic frame initial state given `hp` value.
"""
function build_adiabatic_frame_u0(hp)
    ilvl = init_lvl(hp)
    u0 = zeros(ComplexF64, 4)
    u0[ilvl] = 1.0
    u0
end

function hp_to_gap(hp)
    2 * B(s_star) * hp
end

function gap_to_hp(g)
    g / 2 / B(s_star)
end


text:
## Show method

## WeaveTpl
function Base.show(io::IO, ::MIME"text/html", x::ImageFile)
    write(io, gif_to_data(x.f, x.caption))
end

function Base.show(io::IO, ::MIME"text/latex", x::ImageFile)
    fname = x.f
    if occursin(r"gif$", fname)
        println(io, "XXX can not include `.gif` file here")
    else
        print(io, """
\\begin{figure}
\\caption{$(x.caption)}
\\includegraphics[width=0.6\\textwidth]{$(x.f)}
\\end{figure}
""")
    end
end

# Show SymPy
## Type piracy
function Base.show(io::IO, ::MIME"text/html", x::T) where {T <: SymPy.SymbolicObject}
    write(io, "
")
    show(io, "text/latex", x)
    write(io, "
")
end

function Base.show(io::IO, ::MIME"text/html", x::Array{T}) where {T <: SymPy.SymbolicObject}
    write(io, "
")
    show(io, "text/latex", x)
    write(io, "
")
end


text:
using DataStructures
using Unitful
using MonteCarloMeasurements
using Distributions

t = range(0.0, stop=10.0, length=100)
uniform1(xmin,xmax) = MonteCarloMeasurements.Particles(      100,Distributions.Uniform(xmin,xmax))
uniform2(xmin,xmax) = MonteCarloMeasurements.StaticParticles(100,Distributions.Uniform(xmin,xmax))
particles1 = uniform1(-0.4, 0.4)
particles2 = uniform2(-0.4, 0.4)
result = OrderedDict{String,Any}()

# ∓ are 100 StaticParticles uniform distribution

result["time"] = t*u"s"
result["phi1"] = [sin(t[i]) + particles1*t[i]/10.0 for i in eachindex(t)]*u"rad"
result["phi2"] = [sin(t[i]) + particles2*t[i]/10.0 for i in eachindex(t)]*u"rad"
result["phi3"] = [sin(t[i]) ∓ 0.4*t[i]/10.0        for i in eachindex(t)]*u"rad"

plot(result, ["phi1", "phi2", "phi3"], figure=1,
     heading="Sine(time) with MonteCarloParticles/StaticParticles (plot area)")

plot(result, ["phi1", "phi2", "phi3"], figure=2,
     heading="Sine(time) with MonteCarloParticles/StaticParticles (plot all runs)")

println("\n... result info:")
printResultInfo(result)


text:
using Alpaca
using RDatasets, Test, Distributions
using Random
using StableRNGs
using CategoricalArrays

rng = StableRNG(1234)

df = dataset("datasets", "iris")
df.binary = vec(Float64.(rand(rng,0:1,size(df,1),1)))
df[!,:SpeciesDummy] = categorical(df[!,:Species])
idx = rand(rng,1:3,size(df,1),1)
a = ["A","B","C"]
df.Random = vec([a[i] for i in idx])
df.RandomCategorical = categorical(df[!,:Random])

# proper interface

# try it natively in R
# @rput df
# R"mod <- feglm(binary ~ SepalWidth | SpeciesDummy , df, family = binomial())"
# R"summary(mod, \"hessian\")"
# binomial/logit
result = Alpaca.feglm(df, Alpaca.@formula(binary ~ SepalWidth),
    Binomial(),
    fe = :SpeciesDummy
    )
@test StatsBase.coef(result) ≈ [-0.684548] atol = 1e-4
# testing display of RegressionResult
@show result
# with clustering
# R"mod <- feglm(binary ~ SepalWidth | SpeciesDummy | SpeciesDummy, df, family = binomial())"
# R"summary(mod, \"cluster\", cluster = ~ SpeciesDummy)"
result = Alpaca.feglm(df, Alpaca.@formula(binary ~ SepalWidth),
    Binomial(),
    fe = :SpeciesDummy, vcov = :(cluster(SpeciesDummy))
    )
@test StatsBase.stderror(result) ≈ [0.191123] atol = 1e-4
# two-way clustering
# R"mod <- feglm(binary ~ SepalWidth | SpeciesDummy | SpeciesDummy + RandomCategorical, df, family = binomial())"
# R"summary(mod, \"cluster\", cluster = ~ SpeciesDummy + RandomCategorical)"
result = Alpaca.feglm(df, Alpaca.@formula(binary ~ SepalWidth),
    Binomial(),
    fe = :SpeciesDummy, vcov = :(cluster(SpeciesDummy + RandomCategorical))
    )
@test StatsBase.stderror(result) ≈ [0.259286] atol = 1e-4
# poisson
# R"mod <- feglm(binary ~ SepalWidth | SpeciesDummy, df, family = poisson())"
# R"summary(mod, \"hessian\")"
result = Alpaca.feglm(df, Alpaca.@formula(binary ~ SepalWidth),
    Poisson(),
    fe = :SpeciesDummy
    )
@test StatsBase.coef(result) ≈ [-0.328779] atol = 1e-4
# two fe's
# R"mod <- feglm(binary ~ SepalWidth | SpeciesDummy + RandomCategorical, df, family = poisson())"
# R"summary(mod, \"hessian\")"
result = Alpaca.feglm(df, Alpaca.@formula(binary ~ SepalWidth),
    Poisson(),
    fe =:(SpeciesDummy + RandomCategorical)
    )
@test StatsBase.coef(result) ≈ [-0.319244] atol = 1e-4

# string interface (not essential)
# rr1 = Alpaca.feglm(df, "binary ~ SepalWidth | SpeciesDummy", "Binomial")
# @test StatsBase.coef(rr1) ≈ [-0.221486] atol = 1e-4
# rr1 = Alpaca.feglm(df, "binary ~ SepalWidth | SpeciesDummy", "Binomial", vcov = :(cluster(SpeciesDummy)))
# @test StatsBase.stderror(rr1) ≈ [0.317681] atol = 1e-4


text:
using PyCall: pyimport_conda, pycall
using Conda

function installpypackage()
	try
		pyimport_conda("sklearn", "scikit-learn")
	catch
		try
			Conda.add("scikit-learn")
			pyimport_conda("sklearn", "scikit-learn")
			pyimport_conda("sklearn.decomposition", "scikit-learn")
		catch
			println("scikit-learn failed to install")
		end
	end
end

installpypackage()


text:
using LAP_julia, Images, TestImages, BenchmarkTools, ComputationalResources, AbstractTrees, PyPlot, LinearAlgebra, TimerOutputs
using ImageDistances: mean

img, imgw, flow = gen_init();
1+1


text:
struct NormalData
    X::Matrix{Int}
    x::Vector{Int}
    y::Vector{Float64}
    Xunique::Vector{Vector{Int}}
    function NormalData(X::Matrix{Int}, y::Vector{Float64})
        x = denserank([X[i, :] for i in 1:size(X, 1)])
        Xunique = sort(unique([X[i, :] for i in 1:size(X, 1)]))
        new(X, x, y, Xunique)
    end
end

NormalData(x::Vector{Int}, y::Vector{Float64}) = NormalData(x[:, :], y)

struct NormalDDP <: AbstractDPM
    parent::DPM
    v0::Float64 # TODO: Replace theta hyperparameters by thetaprior::NormalInverseGamma{Float64}
    r0::Float64
    u0::Float64
    s0::Float64
    v1::Vector{Vector{Float64}}
    r1::Vector{Vector{Float64}}
    u1::Vector{Vector{Float64}}
    s1::Vector{Vector{Float64}}
    gammaprior::Womack
    gamma::Vector{Bool}
    G::Int
    function NormalDDP(
        rng::AbstractRNG,
        N::Int,
        G::Int;
        K0::Int = 5,
        a0::Float64 = 2.0, # TODO: Replace a0, b0 by alphaprior = Gamma(2.0, 0.25) (I must update this in DPMNeal3.jl first)
        b0::Float64 = 4.0,
        v0::Float64 = 2.0, # TODO: Replace (v0, r0, u0, s0) by NormalInverseGamma(u0, 1 / r0, v0, s0)
        r0::Float64 = 1.0,
        u0::Float64 = 0.0,
        s0::Float64 = 1.0,
        ζ0::Float64 = 1.0,
    )
        parent = DPM(rng, N; K0, a0, b0)
        v1 = [v0 * ones(G)]
        r1 = [r0 * ones(G)]
        u1 = [u0 * ones(G)]
        s1 = [s0 * ones(G)]
        gammaprior = Womack(G - 1, ζ0)
        gamma = ones(Bool, G)
        new(parent, v0, r0, u0, s0, v1, r1, u1, s1, gammaprior, gamma, G)
    end
end

function parent_dpm(m::NormalDDP)
    m.parent
end

function add_cluster!(m::NormalDDP)
    @extract m : G v0 r0 u0 s0 v1 r1 u1 s1
    push!(v1, v0 * ones(G))
    push!(r1, r0 * ones(G))
    push!(u1, u0 * ones(G))
    push!(s1, s0 * ones(G))
end

function update_suffstats!(m::NormalDDP, data)
    @extract m : v0 r0 u0 s0 v1 r1 u1 s1 gamma
    @extract data : y x
    d = cluster_labels(m)
    while length(v1) < cluster_capacity(m)
        add_cluster!(m)
    end
    for k in active_clusters(m)
        v1[k] .= v0
        r1[k] .= r0
        u1[k] .= u0
        s1[k] .= s0
    end
    for i = 1:length(y)
        zi = gamma[x[i]] ? x[i] : 1
        di = d[i]
        v1[di][zi] += 1
        rm = r1[di][zi] += 1
        um = u1[di][zi] = ((rm - 1) * u1[di][zi] + y[i]) / rm
        s1[di][zi] += (rm / (rm - 1)) * (y[i] - um)^2
    end
end

function update_suffstats!(m::NormalDDP, data, i::Int, k1::Int, k2::Int)
    @extract data : y x
    @extract m : v1 r1 u1 s1 gamma
    while length(v1) < cluster_capacity(m)
        add_cluster!(m)
    end
    zi = iszero(gamma[x[i]]) ? 1 : x[i]

    # Modify cluster/group di/k2
    v1[k2][zi] += 1
    rm = r1[k2][zi] += 1
    um = u1[k2][zi] = ((rm - 1) * u1[k2][zi] + y[i]) / rm
    s1[k2][zi] += (rm / (rm - 1)) * (y[i] - um)^2

    # Modify cluster/group di/k1
    rm = r1[k1][zi]
    um = u1[k1][zi]
    s1[k1][zi] -= (rm / (rm - 1)) * (y[i] - um)^2
    u1[k1][zi]  = (rm * u1[k1][zi] - y[i]) / (rm - 1)
    v1[k1][zi] -= 1
    r1[k1][zi] -= 1
end

function update_gamma!(rng::AbstractRNG, m::NormalDDP, data)
    @extract m : gammaprior gamma
    A = active_clusters(m)

    # Resample gamma[g], given the other gamma's
    for g = 2:length(gamma)
        # log-odds (numerator)
        gamma[g] = 1
        update_suffstats!(m, data)
        log_num = logpdf(gammaprior, gamma[2:end])
        for k ∈ A, j ∈ (1, g)
            log_num += logmglik(m, j, k)
        end

        # log-odds (denominator)
        gamma[g] = 0
        update_suffstats!(m, data)
        log_den = logpdf(gammaprior, gamma[2:end])
        for k ∈ A, j ∈ (1)
            log_den += logmglik(m, j, k)
        end

        # log-odds and new gamma[g]
        log_odds = log_num - log_den
        gamma[g] = rand(rng) <= 1 / (1 + exp(-log_odds))
    end
end

function update_hyperpars!(rng::AbstractRNG, m::NormalDDP, data)
    update_gamma!(rng, m, data)
end

function logpredlik(m::NormalDDP, data, i::Int, k::Int)
    @extract m : v1 r1 u1 s1 gamma
    @extract data : y x
    d = cluster_labels(m)
    yi = y[i]
    di = d[i]
    zi = iszero(gamma[x[i]]) ? 1 : x[i]

    if di == k
        v̄1 = v1[k][zi]
        r̄1 = r1[k][zi]
        ū1 = u1[k][zi]
        s̄1 = s1[k][zi]
        v̄0 = v̄1 - 1
        r̄0 = r̄1 - 1
        ū0 = (r̄1 * ū1 - yi) / r̄0
        s̄0 = s̄1 - (r̄1 / r̄0) * (yi - ū1)^2
    else
        v̄0 = v1[k][zi]
        r̄0 = r1[k][zi]
        ū0 = u1[k][zi]
        s̄0 = s1[k][zi]
        v̄1 = v̄0 + 1
        r̄1 = r̄0 + 1
        ū1 = (r̄0 * ū0 + yi) / r̄1
        s̄1 = s̄0 + (r̄1 / r̄0) * (yi - ū1)^2
    end

    if iszero(r̄0)
        return - 0.5v̄1 * log(s̄1) + loggamma(v̄1 / 2) - 0.5 * log(π)
    else
        return (
            0.5v̄0 * log(s̄0) -
            0.5v̄1 * log(s̄1) +
            loggamma(v̄1 / 2) -
            loggamma(v̄0 / 2) +
            0.5 * log(r̄0 / r̄1) -
            0.5 * log(π)
        )
    end
end

function logpredlik(m::NormalDDP, train, predict, i::Int, k::Int)
    @extract m : v1 r1 u1 s1 gamma
    @extract predict : y x
    yi = y[i]
    zi = iszero(gamma[x[i]]) ? 1 : x[i]
    v̄0 = v1[k][zi]
    r̄0 = r1[k][zi]
    ū0 = u1[k][zi]
    s̄0 = s1[k][zi]
    v̄1 = v̄0 + 1
    r̄1 = r̄0 + 1
    ū1 = (r̄0 * ū0 + yi) / r̄1
    s̄1 = s̄0 + (r̄1 / r̄0) * (yi - ū1)^2

    if iszero(r̄0)
        return - 0.5v̄1 * log(s̄1) + loggamma(v̄1 / 2) - 0.5 * log(π)
    else
        return (
            0.5v̄0 * log(s̄0) -
            0.5v̄1 * log(s̄1) +
            loggamma(v̄1 / 2) -
            loggamma(v̄0 / 2) +
            0.5 * log(r̄0 / r̄1) -
            0.5 * log(π)
        )
    end
end

function logmglik(m::NormalDDP, j::Int, k::Int)
    @extract m : v0 v1 r0 r1 s0 s1
    return(
        0.5v0 * log(s0) -
        0.5v1[k][j] * log(s1[k][j]) +
        loggamma(v1[k][j] / 2) -
        loggamma(v0 / 2) +
        0.5 * log(r0 / r1[k][j]) -
        0.5 * log(π) * (r1[k][j] - r0)
    )
end


text:
using Distributed, StatsBase, ClusterManagers

addprocs(SlurmManager(parse(Int64, ARGS[1])))
@everywhere begin

    using TransitionalMCMC, Distributions

    # Prior Bounds
    lb, ub  = -5, 5

    # Prior Density and sampler
    logprior(x) = logpdf(Uniform(lb, ub), x[1]) + logpdf(Uniform(lb, ub), x[2])
    priorRnd(Nsamples) = rand(Uniform(lb, ub), Nsamples, 2)

    # Log Likelihood
    function logLik(x)
        sleep(1)
        return -1 * ((x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2)
    end

end

Nsamples = 2000

@time samps, acc = tmcmc(logLik, logprior, priorRnd, Nsamples, 5, 2)

rmprocs(workers())


text:
mutable struct ALRSVI <: GradDescent.Optimizer
    opt_type::String
    t::Int64
    g::AbstractArray
    h::Float64
    τ::Float64
    ρ::Float64
end

""" Construct Adaptive Learning Rate for Stochastic Variational Inference"""
function ALRSVI(;τ::Int=100)
    ALRSVI("Adaptive Learning Rate for Stochastic Variational Inference",0,[0.0],0.0,τ,-1.0)
end

GradDescent.params(opt::ALRSVI) = "ρ=$(opt.ρ)"

function init!(inference::Inference{T},model::SVGP,τ::Int=10) where T
    for n_s in 1:τ
        model.inference.MBIndices = StatsBase.sample(1:model.inference.nSamples,inference.nSamplesUsed,replace=false)
        computeMatrices!(model)
        local_updates!(model)
        natural_gradient!(model)
        if n_s == 1
            for (i,opt) in enumerate(inference.optimizer_η₁)
                opt.g = inference.∇η₁[i]./τ
                opt.h = dot(inference.∇η₁[i],inference.∇η₁[i])/τ
            end
            for (i,opt) in enumerate(inference.optimizer_η₂)
                opt.g = Array(inference.∇η₂[i])./τ
                opt.h = dot(inference.∇η₂[i],inference.∇η₂[i])/τ
            end
        else
            for (i,opt) in enumerate(inference.optimizer_η₁)
                opt.g .+= inference.∇η₁[i]./τ
                opt.h += dot(inference.∇η₁[i],inference.∇η₁[i])/τ
            end
            for (i,opt) in enumerate(inference.optimizer_η₂)
                opt.g .+= Array(inference.∇η₂[i])./τ
                opt.h += dot(inference.∇η₂[i],inference.∇η₂[i])/τ
            end
        end
    end
    for (i,opt) in enumerate(inference.optimizer_η₁)
        opt.τ = τ
        opt.ρ = 1.0#dot(opt.g,opt.g)/opt.h
    end
    for (i,opt) in enumerate(inference.optimizer_η₂)
        opt.τ = τ
        opt.ρ = 1.0#dot(opt.g,opt.g)/opt.h
    end
end

function GradDescent.update(opt::ALRSVI, g_t::AbstractArray{T,N}) where {T<:Real,N}
    # update timestep
    if opt.ρ < 0
        @error "Optimizer has not been initialized externally, it needs a special initialization"
    end
    opt.t += 1
    opt.g .= (1.0-1.0/opt.τ).*opt.g + 1.0/opt.τ .* g_t
    opt.h = (1.0-1.0/opt.τ).*opt.h + 1.0/opt.τ .* dot(g_t,g_t)
    opt.ρ = dot(opt.g,opt.g)/opt.h
    opt.τ = opt.τ*(1-opt.ρ)+1.0
    return opt.ρ * g_t
end


text:
mutable struct Cube
    a::Vector{Int}
    b::Vector{Int}
end

function day22(n=0)
    lines = filter(!isempty, strip.(readlines("AoCdata/AoC_2021_day22.txt")))
    oldlist = Cube[]
    for (j, line) in enumerate(lines)
        n != 0 && j > n && break
        arr = split(line, r"[\s\,][xyz]\=|\.\.|\,")
        newc = Cube(
            [parse.(Int, s) for s in arr[2:2:6]], [parse.(Int, s) + 1 for s in arr[3:2:7]]
        )
        newlist = Cube[]
        for oldc in oldlist
            if any(oldc.b .<= newc.a .|| newc.b .<= oldc.a) # no overlap
                push!(newlist, oldc)
            else
                for i in 1:3
                    if oldc.a[i] < newc.a[i] && newc.b[i] < oldc.b[i]
                        c1, c2 = deepcopy(oldc), deepcopy(oldc)
                        c1.b[i], c2.a[i] = newc.a[i], newc.b[i]
                        push!(newlist, c1, c2)
                        oldc.a[i], oldc.b[i] = newc.a[i], newc.b[i]
                    elseif oldc.a[i] < newc.a[i] <= oldc.b[i]
                        c1 = deepcopy(oldc)
                        c1.b[i] = newc.a[i]
                        push!(newlist, c1)
                        oldc.a[i] = newc.a[i]
                    elseif oldc.a[i] <= newc.b[i] < oldc.b[i]
                        c1 = deepcopy(oldc)
                        c1.a[i] = newc.b[i]
                        push!(newlist, c1)
                        oldc.b[i] = newc.b[i]
                    end
                end
            end
        end
        oldlist = newlist
        first(arr) == "on" && push!(oldlist, newc)
    end
    return sum(prod(abs.(cub.a .- cub.b)) for cub in oldlist)
end

println("Part 1: ", day22(20))
println("Part 2: ", day22())


text:
using Test, Distributions
@testset "Normal" begin
    @test isa(convert(Normal{Float64}, Float16(0), Float16(1)),
              Normal{Float64})
    @test Inf === logpdf(Normal(0, 0), 0)
    @test -Inf === logpdf(Normal(), Inf)
    @test iszero(logcdf(Normal(0, 0), 0))
    @test iszero(logcdf(Normal(), Inf))
    @test -Inf === logccdf(Normal(0, 0), 0)
    @test iszero(logccdf(Normal(eps(), 0), 0))
    @test -Inf === quantile(Normal(), 0)
    @test iszero(quantile(Normal(), 0.5))
    @test Inf === quantile(Normal(), 1)
    @test -Inf === quantile(Normal(0, 0), 0)
    @test iszero(quantile(Normal(0, 0), 0.75))
    @test Inf === quantile(Normal(0, 0), 1)
    @test -Inf === quantile(Normal(0.25, 0), 0)
    @test 0.25 == quantile(Normal(0.25, 0), 0.95)
    @test Inf === quantile(Normal(0.25, 0), 1)
    @test Inf === cquantile(Normal(), 0)
    @test iszero(cquantile(Normal(), 0.5))
    @test -Inf === cquantile(Normal(), 1)
    @test Inf === cquantile(Normal(0, 0), 0)
    @test iszero(cquantile(Normal(0, 0), 0.75))
    @test -Inf === cquantile(Normal(0, 0), 1)
    @test Inf === cquantile(Normal(0.25, 0), 0)
    @test 0.25 == cquantile(Normal(0.25, 0), 0.95)
    @test -Inf === cquantile(Normal(0.25, 0), 1)
    @test -Inf === invlogcdf(Normal(), -Inf)
    @test isnan(invlogcdf(Normal(), NaN))
    @test Inf === invlogccdf(Normal(), -Inf)
    @test isnan(invlogccdf(Normal(), NaN))
end


text:
export softmax, softmax!, ∇softmax, ∇softmax!,
       logsoftmax, logsoftmax!, ∇logsoftmax, ∇logsoftmax!

"""
    softmax(xs) = exp.(xs) ./ sum(exp.(xs))

[Softmax](https://en.wikipedia.org/wiki/Softmax_function) takes
log-probabilities (any real vector) and returns a probability distribution that
sums to 1.

If given a matrix it will treat it as a batch of vectors, with each column
independent.

    julia> softmax([1,2,3.])
    3-element Array{Float64,1}:
      0.0900306
      0.244728
      0.665241
"""
softmax(xs) = softmax!(similar(xs), xs)

function softmax!(out::AbstractVecOrMat{T}, xs::AbstractVecOrMat{T}) where {T}
    @inbounds for j = 1:size(xs, 2)
        # First, store column-wise maximum in the last element of `out`
        out[end, j] = xs[end, j]
        @inbounds for i = 1:(size(xs, 1) - 1)
            out[end, j] = max(out[end, j], xs[i, j])
        end

        # Subtract the column-wise maximums to normalize, take exp()
        # out .= exp(xs .- out[end, :])
        @inbounds for i = 1:size(out, 1)
            out[i, j] = exp(xs[i, j] - out[end, j])
        end

        # Normalize by sum of the entire thing
        # out ./= sum(out, 1)
        s = T(0)
        @inbounds for i = 1:size(out, 1)
            s += out[i, j]
        end
        @inbounds for i = 1:size(out, 1)
            out[i, j] /= s
        end
    end
    return out
end

function ∇softmax!(out::AbstractVecOrMat, Δ::AbstractVecOrMat, xs::AbstractVecOrMat)
    sf = softmax(xs)
    out .= sf .* (Δ .- sum(Δ .*sf, dims = 1))
end

∇softmax(Δ, xs) = ∇softmax!(similar(Δ), Δ, xs)
∇softmax!(Δ, xs) = ∇softmax!(Δ, Δ, xs)

"""
    logsoftmax(xs) = log.(exp.(xs) ./ sum(exp.(xs)))

`logsoftmax(xs)` computes the log of `softmax(xs)`, but in a more numerically stable
way than directly taking the log of the the softmax function, which is commonly used in
computing cross entropy loss.
"""
logsoftmax(xs) = logsoftmax!(similar(xs), xs)
function logsoftmax!(out::AbstractVecOrMat, xs::AbstractVecOrMat)
    for j = 1:size(xs, 2)
        @inbounds begin
            xi_max = xs[1, j]
            for i = 1:size(out, 1)
                xi_max = max(xi_max, xs[i, j])
            end
            s = zero(eltype(out))
            for i = 1:size(out, 1)
                s += exp(xs[i, j] - xi_max)
            end
            for i = 1:size(out, 1)
                out[i, j] = xs[i, j] - log(s) - xi_max
            end
        end
    end
    return out
end
∇logsoftmax(Δ, xs) = ∇softmax(Δ ./ max.(eps(eltype(xs)),softmax(xs)), xs)
∇logsoftmax!(Δ, xs) = ∇softmax!(Δ, Δ, xs)