text:
import LowRankModels: parameter_estimate, ObsArray, AbstractGLRM, Loss, QuadReg, GLRM, scale_regularizer!

export GFRM, parameter_estimate, scale_regularizer!

# todo
# * estimate lipshitz constant more reasonably
# * map GFRM to GLRM and back
# * check that PRISMA code calculates the right thing via SDP

type GFRM<:AbstractGLRM
    A                            # The data table
    losses::Array{Loss,1}        # Array of loss functions
    r::ProductRegularizer        # The regularization to be applied to U
    k::Int                       # Estimated rank of solution U (used to improve performance of PRISMA algorithm)
    observed_features::ObsArray  # for each example, an array telling which features were observed
    observed_examples::ObsArray  # for each feature, an array telling in which examples the feature was observed
    U::AbstractArray{Float64,2}  # Representation of data in numerical space. A ≈ U = X'Y
    W::AbstractMatrix            # Representation of data in encoded space, depending on algorithm.
end
                                 # for prisma, W = [? U; U' ?]
                                 # for frank wolfe, W = LowRankOperator(X', Y)

### From GFRMs to GLRMs and back

function GFRM(glrm::GLRM; force=false, use_reg_scale=true, init_prisma=false)
    if !force # error check unless we force a conversion
    	if !isa(glrm.rx, QuadReg)
    		error("I don't know how to convert the regularization on X, $(typeof(glrm.rx)), into a regularizer for a GFRM")
    	end
    	if !(isa(glrm.ry, QuadReg) || all(map(r->isa(r, QuadReg), glrm.ry)))
    		error("I don't know how to convert the regularization on Y, $(glrm.ry), into a regularizer for a GFRM")
    	end
    end

    if use_reg_scale
    	minscale = minimum(map(scale,glrm.ry))
        maxscale = maximum(map(scale,glrm.ry))
        @assert(minscale == maxscale && minscale == scale(glrm.rx),
        	"we only know how to convert glrms with similarly scaled quadratic regularization")
        r = TraceNormReg(scale(glrm.rx)*2)
    else
    	r = TraceNormReg(1)
    end

    # translate the low rank variables X and Y into full rank variables U and W
    if init_prisma
        U = glrm.X'*glrm.Y
        W = [glrm.X glrm.Y]'*[glrm.X glrm.Y]
    else
        m = size(glrm.X,2)
        d = size(glrm.Y,2)
        U = spzeros(0,0)
        W = spzeros(0,0)
    end
    return GFRM(glrm.A, glrm.losses, r, glrm.k,
		    	glrm.observed_features, glrm.observed_examples,
		    	U, W)
end

# if k=0 we use k = numerical_rank(gfrm.U)
function GLRM(gfrm::GFRM, k::Int=0; tol=1e-5)
    # we only know how to convert glrms with similarly scaled quadratic regularization
    if isa(gfrm.r, TraceNormReg)
        r = QuadReg(scale(gfrm.r)/2)
    else
    	error("I don't know how to convert $(typeof(gfrm.r)) into a regularizer for a GLRM")
    end
    u,s,v = svd(gfrm.U)
    if k == 0
    	k = sum(s.>=tol)
    else
    	k = min(k, min(size(gfrm.U)...))
    end
    X = diagm(sqrt(s[1:k]))*u[:,1:k]'
    Y = diagm(sqrt(s[1:k]))*v[:,1:k]'
    return GLRM(gfrm.A, gfrm.losses, r, r, k,
		    	observed_features = gfrm.observed_features,
		    	observed_examples = gfrm.observed_examples,
		    	X = X, Y = Y)
end

parameter_estimate(gfrm::GFRM) = (gfrm.W,)

function scale_regularizer!(gfrm::GFRM, newscale::Number)
    scale!(gfrm.r, newscale)
    return gfrm
end


text:
using CSV, DataFrames, ArgParse

# function parse_commandline()
#   s = ArgParseSettings()
#   s.description = "Merge all atomic csv into one large csv."
#   @add_arg_table s begin
#     "run_dir"
#     help = "Directory which contains each run folder."
#     arg_type = String
#     required = true
#   end
#   return parse_args(s)
# end

function main(args)
    println("--- Merge_csv.jl ---")

    ## Setting working directory
    # input_params = parse_commandline()
    # input_params = Dict("run_dir" => joinpath("Mosek_runs", "pararuns", "07_31-19h20"))
    input_params = Dict("run_dir" => joinpath("07_31-19h20"))
    workdir = input_params["run_dir"]

    info("Working in $workdir")

    run_folders = Set(filter(x->isdir(joinpath(workdir, x)), readdir(workdir)))
    println("Found $(length(run_folders)) folders")

    global_csv = DataFrames.DataFrame()

    n_addedinstances = 1
    for folder in run_folders
        curfolder = joinpath(workdir, folder)
        println("-- working on $curfolder")

        ## Collecting information
        csvfile = joinpath(curfolder, "momentsos_solve.csv")

        out_str = filter(x->ismatch(r"hierarchy.*\.o", x), readdir(curfolder))
        @assert length(out_str) == 1
        out_str = first(out_str)

        err_str = filter(x->ismatch(r"hierarchy.*\.e", x), readdir(curfolder))
        @assert length(err_str) == 1
        err_str = first(err_str)

        instancename = out_str[11:end-2]

        waskilled = readstring(joinpath(curfolder, out_str)) == "Killed\n"
        waserror = readstring(joinpath(curfolder, err_str)) != ""

        # # Has timed out ?
        # hastimedout = true
        # if isfile(csvfile) != 0
        #     hastimedout = false
        # end
        @assert isfile(csvfile)

        if isfile(csvfile)
            ## Reading csv
            cur_csv = CSV.read(csvfile, delim=";")

            cur_csv[:slv_run_status] = "All went well."
            if waskilled
                cur_csv[:slv_run_status] = "Process was killed (by qsub)."
                info("Process was killed (qsub)")
            end
            if waserror
                cur_csv[:slv_run_status] = "Process errored (likely was qdel)."
                info("Process errored (likely was qdel).")
            end

            if eltype(cur_csv[:pb_name]) == Missing
                warn("Empty pb name")
                newcsv = DataFrame()
                for name in names(global_csv)
                    if name == :pb_name
                        newcsv[:pb_name] = instancename
                    else
                        newcsv[name] = cur_csv[name]
                    end
                end
                cur_csv = newcsv
            end

            allowmissing!(cur_csv)

            if size(global_csv) == (0,0)
                global_csv = deepcopy(cur_csv)
                allowmissing!(global_csv)
            elseif names(global_csv) != names(cur_csv)
                # @show names(cur_csv)
                # @show length(names(cur_csv))
                # @show names(global_csv)
                # @show length(names(global_csv))
                warn("Ignoring  $folder :   Names")
            elseif DataFrames.eltypes(global_csv) != DataFrames.eltypes(cur_csv)
                # @show DataFrames.eltypes(global_csv)
                # @show DataFrames.eltypes(cur_csv)
                # @show names cur_csv
                warn("Ignoring  $folder :   Elttypes")
                # append!(global_csv, cur_csv)
            else
                n_addedinstances += 1
                append!(global_csv, cur_csv)
            end
        end

        println("   Size global CSV : $(size(global_csv))")
    end

    csvname = "report_log.csv"
    println("Writing full csv file $(joinpath(workdir, csvname))")

    CSV.write(joinpath(workdir, csvname), global_csv, delim=';')

    nbfine = size(global_csv[global_csv[:slv_run_status].=="All went well.", :], 1)

    info("$nbfine runs were fine.")
    info("$n_addedinstances / $(length(run_folders)) added to final csv.")

    return nothing
end

main(ARGS)

ext: jl


text:
ENV["JULIA_CUDA_MEMORY_POOL"] = "split" # "binned" / "split"

using AlphaZero
using AlphaZero.Log
using Formatting

const REPEAT = 200

#####
##### Profiling utilities
#####

function profile_self_play(network, params, n)
  player = AlphaZero.MctsPlayer(network, params)
  time = @elapsed for i in 1:n
    AlphaZero.play_game(player, player)
  end
  itr = MCTS.inference_time_ratio(player.mcts)
  aed = MCTS.average_exploration_depth(player.mcts)
  return (time / n, itr, aed)
end

function profile_self_play(configs::Vector)
  log  = Logger()
  time = Log.ColType(10, x -> format("{:.1f} min", x * 100 / 60))
  itr  = Log.ColType(5,  x -> format("{}%", round(Int, x * 100)))
  expd = Log.ColType(5,  x -> format("{:.1f}", x))
  tab  = Log.Table([
    ("T100", time, s -> s[1]),
    ("ITR",  itr,  s -> s[2]),
    ("EXPD", expd, s -> s[3])])
  for (title, net, params) in configs
    profile_self_play(net, params, 1) # Compilation
    rep = profile_self_play(net, params, REPEAT)
    Log.table_row(log, tab, rep, [title])
  end
end

#####
##### Main
#####

include("../games.jl")
const GAME = "connect-four"
const SelectedGame = GAME_MODULE[GAME]
using .SelectedGame: Game, Training

function config(nblocks, nfilters, niters, nworkers)
  title = "$nblocks blocks, $nfilters filters, $niters iters, $nworkers workers"
  network = ResNet{Game}(ResNetHP(Training.netparams,
    num_blocks=nblocks, num_filters=nfilters))
  params = MctsParams(Training.params.self_play.mcts,
    num_iters_per_turn=niters)
  return (title, network, params)
end

profile_self_play([
  config( 7,  64, 400, 64),
  config( 7, 128, 400, 64),
  config(10, 128, 400, 64),
  config(10, 128, 320, 64),
])


text:
abstract type Component{T} end
abstract type AbstractBody{T} <: Component{T} end
abstract type AbstractConstraint{T,N} <: Component{T} end

# TODO do id differently?
CURRENTID = -1
getGlobalID() = (global CURRENTID -= 1; return CURRENTID + 1)
resetGlobalID() = (global CURRENTID = -1; return)

Base.show(io::IO, component::Component) = summary(io, component)
function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, constraint::AbstractConstraint{T,N}) where {T,N}
    summary(io, constraint)
    println(io,"")
    println(io, " id:     "*string(constraint.id))
    println(io, " name:   "*string(constraint.name))
    println(io, " active: "*string(constraint.active))
end

Base.eltype(::Type{<:Component{E}}) where {E} = @isdefined(E) ? E : Any
Base.length(::AbstractConstraint{T,N}) where {T,N} = N
getid(component::Component) = component.id
activate!(component::Component) = (component.active = true; return)
deactivate!(component::Component) = (component.active = false; return)

isactive(component::Component) = component.active
isinactive(component::Component) = !isactive(component)

function Base.getindex(dict::UnitDict{Base.OneTo{K},<:Component}, key::String) where K
    for component in dict.values
        component.name == key && return component
    end

    return
end
function Base.getindex(dict::UnitDict{UnitRange{K},<:Component}, key::String) where K
    for component in dict.values
        component.name == key && return component
    end

    return
end


text:
#include("HPDE_options.jl")
##########################################################
"""
$(TYPEDEF)

Data type for PDE equations (i.e., problem)
"""
abstract type AbstractHPDEProblem end

abstract type Dimension end
abstract type _1D <: Dimension end
abstract type _2D <: Dimension end
abstract type _3D <: Dimension end

function checkDim(D::Type{<:Dimension})
    if D === _1D
        return 1
    elseif D === _2D
        return 2
    elseif D === _3D
        return 3
    end
    return 0
end

mutable struct FVMPDEProblem
    """
    Flux function F in x direction

    """
    F::Function

    """
    Flux function G in y direction
    !!! compat
        Will be added in future versions
    """
    G::Function

    """
    Flux function H in z direction
    !!! compat
        Will be added in future versions
    """
    H::Function

    """
    Source function
    !!! compat
        Will be added in future versions
    """
    S::Function

    nvars::Int8

    param::Any

    U0

    U_min
    U_max

    dimension::Type{<:Dimension}

    function FVMPDEProblem(func_F::Function, U0; kwargs...)

        zeroFunc(U) = zeros(size(U))
        F = func_F
        D = _1D
        (G,D) = haskey(kwargs,:func_G) ? (kwargs[:func_G],_2D) : (zeroFunc,D)
        (H,D) = haskey(kwargs,:func_G) && haskey(kwargs,:func_H) ? (kwargs[:func_H],_3D) : (zeroFunc,D)
        S = haskey(kwargs,:func_S) ? kwargs[:func_S] : zeroFunc

        @assert size(F(U0)) == size(U0)
        @assert size(G(U0)) == size(U0)
        @assert size(H(U0)) == size(U0)
        @assert size(S(U0)) == size(U0)

        θ = haskey(kwargs,:parameters) ? kwargs[:parameters] : nothing

        U_min = haskey(kwargs,:ul) ? kwargs[:ul] : -Inf * ones(size(U0))
        U_max = haskey(kwargs,:uh) ? kwargs[:uh] : Inf * ones(size(U0))

        @assert size(U_min) == size(U0)
        @assert size(U_max) == size(U0)

        nvars = size(U0,1)

        #@assert checkDim(D) == size(U0,)

        return new(F, G, H, S, nvars, θ, U0, U_min, U_max, D )
    end

end

    # for 2D it should be Array{Float64,4}
    # for 3D it should be Array{Float64,5}
    # this will be fixed later
F(u) = zeros(size(u))
A = FVMPDEProblem(F,[1])


text:
@testset "Property" begin
    atom = Atom("")
    set_property!(atom, "test", false)
    @test property(atom, "test") == false
    set_property!(atom, "bar", 3.4)
    @test properties_count(atom) == 2
    @test Set(list_properties(atom)) == Set(["bar", "test"])

    residue = Residue("")
    set_property!(residue, "test", false)
    @test property(residue, "test") == false
    @test properties_count(residue) == 1
    @test list_properties(residue) == ["test"]

    frame = Frame()
    set_property!(frame, "test", 1234.567)
    @test property(frame, "test") == 1234.567

    set_property!(frame, "test2", "TESTINGTESTING")
    @test property(frame, "test2") == "TESTINGTESTING"

    set_property!(frame, "test2", [1.0,2.0,3.0])
    @test property(frame, "test2") == [1.0,2.0,3.0]

    lipsum = "Enim possimus non tenetur laborum. Quo ut
    reiciendis in at suscipit in nulla quas. Eos minus animi quis quo enim
    facilis dolores quo.

    Iusto incidunt molestiae distinctio cum officia cupiditate tempore qui. Sunt
    error voluptatibus reprehenderit rerum aut a. Earum illum asperiores est
    accusamus illo sit.

    Magni incidunt enim nam. Nisi adipisci harum corrupti reprehenderit nihil.
    Voluptatem aut non labore et sint nesciunt qui. Perspiciatis atque laborum
    aut."
    set_property!(frame, "very long", lipsum)
    @test property(frame, "very long") == lipsum
    @test properties_count(frame) == 3
    @test Set(list_properties(frame)) == Set(["very long", "test2", "test"])
end


text:
## time-related functions ##

# TODO: check for usleep errors?
@unix_only systemsleep(s::Real) = ccall(:usleep, Int32, (Uint32,), uint32(iround(s*1e6)))
@windows_only systemsleep(s::Real) = (ccall(:Sleep, stdcall, Void, (Uint32,), uint32(iround(s*1e3))); return int32(0))

type TmStruct
    sec::Int32
    min::Int32
    hour::Int32
    mday::Int32
    month::Int32
    year::Int32
    wday::Int32
    yday::Int32
    isdst::Int32
    # on some platforms the struct is 14 words, even though 9 are specified
    _10::Int32
    _11::Int32
    _12::Int32
    _13::Int32
    _14::Int32

    TmStruct(sec, min, hour, mday, month, year, wday, yday, isdst) =
        new(sec, min, hour, mday, month, year, wday, yday, isdst, 0,0,0,0,0)
    TmStruct() = new(0,0,0,0,0,0,0,0,0,0,0,0,0,0)
    function TmStruct(t::Real)
        t = floor(t)
        tm = TmStruct()
        # TODO: add support for UTC via gmtime_r()
        ccall(:localtime_r, Ptr{Void}, (Ptr{Int}, Ptr{Void}), &t, &tm)
        return tm
    end
end

strftime(t) = strftime("%c", t)
strftime(fmt::String, t::Real) = strftime(fmt, TmStruct(t))
function strftime(fmt::String, tm::TmStruct)
    timestr = Array(Uint8, 128)
    n = ccall(:strftime, Int, (Ptr{Uint8}, Int, Ptr{Uint8}, Ptr{Void}),
              timestr, length(timestr), fmt, &tm)
    if n == 0
        return ""
    end
    bytestring(convert(Ptr{Uint8},timestr))
end

strptime(timestr::String) = strptime("%c", timestr)
function strptime(fmt::String, timestr::String)
    tm = TmStruct()
    r = ccall(:strptime, Ptr{Uint8}, (Ptr{Uint8}, Ptr{Uint8}, Ptr{Void}),
              timestr, fmt, &tm)
    # the following would tell mktime() that this is a local time, and that
    # it should try to guess the timezone. not sure if/how this should be
    # exposed in the API.
    # tm.isdst = -1
    if r == C_NULL
        error("invalid arguments")
    end
    tm
end

time(tm::TmStruct) = float64(ccall(:mktime, Int, (Ptr{Void},), &tm))

## process-related functions ##

getpid() = ccall(:jl_getpid, Int32, ())

## network functions ##

function gethostname()
    hn = Array(Uint8, 256)
    @unix_only err=ccall(:gethostname, Int32, (Ptr{Uint8}, Uint), hn, length(hn))
    @windows_only err=ccall(:gethostname, stdcall, Int32, (Ptr{Uint8}, Uint32), hn, length(hn))
    systemerror("gethostname", err != 0)
    bytestring(convert(Ptr{Uint8},hn))
end

## Memory related ##

c_free(p::Ptr) = ccall(:free, Void, (Ptr{Void},), p)
c_malloc(size::Integer) = ccall(:malloc, Ptr{Void}, (Csize_t,), size)
c_realloc(p::Ptr, size::Integer) = ccall(:realloc, Ptr{Void}, (Ptr{Void}, Csize_t), p, size)
c_calloc(num::Integer, size::Integer) = ccall(:calloc, Ptr{Void}, (Csize_t, Csize_t), num, size)

max_stars_repo_name: tobig/julia

max_issues_repo_name: tobig/julia

max_forks_repo_name: tobig/julia


text:
# KineticMonteCarlo
struct KineticMonteCarlo end

"""
    next(alg::KineticMonteCarlo, [rng::AbstractRNG,] rates::AbstractWeights)::Tuple{Float64,Int}

Next stochastic event (`\\Delta t`, index) drawn proportional to probability given in `rates`
"""
function next(alg::KineticMonteCarlo, rng::AbstractRNG, rates::Union{AbstractWeights, AbstractVector})::Tuple{Float64,Int}
    sum_rates = sum(rates)
    if !(sum_rates > 0)
        return Inf, 0
    end
    dtime = next_time(rng, sum_rates)
    event = next_event(rng, rates)
    return dtime, event
end

next(alg::KineticMonteCarlo, rates::AbstractWeights) = next(alg, Random.GLOBAL_RNG, rates)

"""
    next(alg::KineticMonteCarlo, [rng::AbstractRNG,] event_handler::AbstractEventHandlerRate)

Next stochastic event (`\\Delta t`, event type) organized by `event_handler`
fast(to be tested, depends on overhead of EventList) implementation of next_event_rate if defined by EventList object

TOOD: write this as wrapper around above..
"""
function next(alg::KineticMonteCarlo, rng::AbstractRNG, event_handler::AbstractEventHandlerRate)::Tuple{Float64,Int}
    if !(sum(event_handler.list_rate) > 0)
        return Inf, 0
    end
    dtime = next_time(rng, sum(event_handler.list_rate))
    event = next_event(rng, event_handler)
    return dtime, event
end

next(alg::KineticMonteCarlo, event_handler::AbstractEventHandlerRate) = next(alg, Random.GLOBAL_RNG, event_handler)

"""
    next_time([rng::AbstractRNG,] rate::Float64)::Float64

Next stochastic `\\Delta t` for Poisson process with `rate`
"""
function next_time(rng::AbstractRNG, rate::Float64)::Float64
    return randexp(rng) / rate
end

next_time(rate::Float64) = next_time(Random.GLOBAL_RNG, rate)

"""
    next_event(rng::AbstractRNG, cumulated_rates::Vector{T})::Int where {T<:AbstractFloat}

Select a single random index in `1:length(cumulated_rates)` with cumulated probability given in `cumulated_rates`.

# Remarks
Deprecated unless we find a good data structure for (dynamic) cumulated rates
"""
function next_event_cumulative(rng::AbstractRNG, cumulated_rates::Vector{T})::Int where {T <: AbstractFloat}
    theta = rand(rng) * cumulated_rates[end]
    index = MonteCarloX.binary_search(cumulated_rates, theta)
    return index
end

"""
    next_event([rng::AbstractRNG,] rates::AbstractWeights)::Int

Select a single random index in `1:length(rates)` with probability proportional to the entry in `rates`.

# Remarks
This is on average twice as fast as StatsBase.sampling because it can iterate from either beginning or end of rates
"""
function next_event(rng::AbstractRNG, rates::Union{AbstractWeights, AbstractVector})::Int
    sum_rates = sum(rates)
    theta = rand(rng) * sum_rates

    # this is for extra performance in case of large lists (factor of 2)
    if theta < 0.5 * sum_rates
        index = 1
        cumulated_rates = rates[index]
        while cumulated_rates < theta && index < length(rates)
            index += 1
            @inbounds cumulated_rates += rates[index]
        end
        return index

    else
        index = length(rates)
        cumulated_rates_lower = sum_rates - rates[index]
        # cumulated_rates in this case belong to index-1
        while cumulated_rates_lower > theta && index > 1
            index -= 1
            @inbounds cumulated_rates_lower -= rates[index]
        end
        return index

    end
end
# relevent to write this as extra function this for performance issues (compared to keyword argument with preset value on GLOBAL_RNG)
next_event(rates::Union{AbstractWeights, AbstractVector}) = next_event(Random.GLOBAL_RNG, rates)

"""
    next_event([rng::AbstractRNG,] rates::MVector{2,T})::Int where {T <: AbstractFloat}

Select next event for special case of two events only.
"""

function next_event(rng::AbstractRNG, rates::MVector{2,T})::Int where {T <: AbstractFloat}
    if rand(rng) * (rates[1] + rates[2]) < rates[1]
        return 1
    else
        return 2
    end
end
next_event(rates::MVector{2,T}) where T = next_event(Random.GLOBAL_RNG, rates)

"""
    next_event([rng::AbstractRNG,] event_handler::AbstractEventHandlerRate{T})::T where T

Select a single random event with a given probability managed by `event_handler`.

The `event_handler` also manages the case that no valid events are left (e.g.
when all rates are equal to zero). This becomes relevant when using [`advance!`](@ref)
to advance for some time.

See also: [`advance!`](@ref)
"""
function next_event(rng::AbstractRNG, event_handler::AbstractEventHandlerRate{T})::T where T
    ne = length(event_handler)
    if ne > 1
        theta::Float64 = rand(rng) * sum(event_handler.list_rate)
        index_last = last_index(event_handler)
        index_first = first_index(event_handler)

        if theta < 0.5 * sum(event_handler.list_rate)
            index = index_first
            cumulated_rates = event_handler.list_rate[index]
            while cumulated_rates < theta && index < index_last
                index = next_index(event_handler, index)
                @inbounds cumulated_rates += event_handler.list_rate[index]
            end
        else
            index = index_last
            cumulated_rates_lower = sum(event_handler.list_rate) - event_handler.list_rate[index]
            # cumulated_rates in this case belong to id-1
            while cumulated_rates_lower > theta && index > index_first
                index = previous_index(event_handler, index)
                @inbounds cumulated_rates_lower -= event_handler.list_rate[index]
            end
        end
        return index
    elseif ne == 1
        return index = first_index(event_handler)
    else
        return event_handler.noevent
    end
end

next_event(event_handler::AbstractEventHandlerRate) = next_event(Random.GLOBAL_RNG, event_handler)

function next_event(rng::AbstractRNG, event_handler::ListEventRateSimple{T})::T where T
    if length(event_handler) > 0
        index = next_event(rng, event_handler.list_rate)
        return event_handler.list_event[index]
    else
        return event_handler.noevent
    end
end

#Q: need to keep the type stability output?
next_event(event_handler::ListEventRateSimple) = next_event(Random.GLOBAL_RNG, event_handler)

"""
    advance!(alg::KineticMonteCarlo, [rng::AbstractRNG], event_handler::AbstractEventHandlerRate, update!::Function, total_time::T)::T where {T<:Real}

Draw events from `event_handler` and update `event_handler` with `update!` until `total_time` has passed. Return time of last event.
"""
function advance!(alg::KineticMonteCarlo, rng::AbstractRNG, event_handler::AbstractEventHandlerRate, update!::Function, total_time::T)::T where {T <: AbstractFloat}
    time::T = 0
    while time <= total_time
        if length(event_handler) == 0
            println("WARNING: no events left before total_time reached")
            return time
        end
        dt, event = next(alg, rng, event_handler)
        time += dt
        update!(event_handler, event)
    end
    return time
end

function advance!(alg::KineticMonteCarlo, event_handler::AbstractEventHandlerRate, update!::Function, total_time::T)::T where {T <: AbstractFloat}
  advance!(alg, Random.GLOBAL_RNG, event_handler, update!, total_time)
end

"""
    advance!(alg::KineticMonteCarlo, [rng::AbstractRNG], rates::AbstractVector, update!::Function, total_time::T)::T where {T<:Real}

Draw events from `event_handler` and update `event_handler` with `update!` until `total_time` has passed. Return time of last event.
#TODO: Discuss with Martin what a suitable API is here
"""
function advance!(alg::KineticMonteCarlo, rng::AbstractRNG, rates::AbstractVector, update!::Function, total_time::T)::T where {T <: AbstractFloat}
    time::T = 0
    while time <= total_time
        dt, event = next(alg, rng, rates)
        time += dt
        update!(rates, event)
    end
    return time
end


text:
using Test
using Flux
using Flux:params
using GraphTransformer
using GeometricFlux
using LightGraphs

@testset "GraphTransformer.jl" begin
    include("./multi_head_attn.jl")
    include("./graph_transformer_layer.jl")
end

max_stars_repo_name: DwaraknathT/Flux-Graph-Transformer

max_issues_repo_name: DwaraknathT/Flux-Graph-Transformer

max_forks_repo_name: DwaraknathT/Flux-Graph-Transformer


text:
#! Mistake I was making before: Treating each firm as if it was some kind of monopolist of a given size.
#! However, firms have actually equal size!

#! Remove strategic complementarities

"""
    Labor demand by old and sophisticated (OS) firm.
"""
function labor_demand_S(k, par)
    @unpack_Pars_OLG par
    ((α_hat * Y^α_Y * k^δ_hat) / w)^(1 / (1 - α_hat))
end

"""
    Labor demand by old and unsophisticated (OU) firm.
"""
function labor_demand_U(k, par)
    @unpack_Pars_OLG par
    ((α_hat * Y^α_Y * k^δ_hat) / w)^(1 / (1 - α_hat))
end

"""
    Labor demand by young (Y) firm.
"""
function labor_demand_Y(k, par)
    @unpack_Pars_OLG par
    ((α_hat * Y^α_Y * k^δ_hat) / w)^(1 / (1 - α_hat))
end

"""
    How much capital A firms have in place given their production and selling decision.
"""
total_capital(k_b, k_P) = k_b + k_P

"""
    Market clearing condition for traded capital. For now, unsophisticated and young firms buy data from old sophisticated firms. For any unit OU firms ship, only 1-τ units arrive. Needs to be scaled by sector size (γ for A and 1-γ for B).
"""
function market_clearing(k_S, k_U, k_Y, par)
    @unpack_Pars_OLG par
    #! Check whether k_U is positive or negative
    if k_U > 0
        μ * k_Y + μ * (1 - γ) * k_U + (1 - τ) * μ * γ * k_S
    elseif k_U < 0
        μ * k_Y + (1-τ) * μ * (1 - γ) * k_U + (1 - τ) * μ * γ * k_S
    end
end

"""
    Marginal cost of capital production by sophisticated firms.
"""
mc_S(k_I, par) = par.ca_S * par.cb_S * k_I^(par.cb_S - 1)

"""
    Marginal cost of capital production by B firms.
"""
mc_U(k_I, par) = par.ca_U * par.cb_U * k_I^(par.cb_U - 1)

"""
    Yost of capital production by A firms.
"""
c_S(k_I, par) = par.ca_S * k_I^par.cb_S
"""
    Yost of capital production by B firms.
"""
c_U(k_I, par) = par.ca_U * k_I^par.cb_U

"""
    How much capital is produced by A firms if they are indifferent between
    producing one more unit or selling one unit less.

        AND

    The firm does not sell all of its capital! Otherwise producing one more unit has a different payoff (sell and keep the rest instead of keeping it fully).
"""
k_prod_S(par) = (par.p_I / (par.ca_S * par.cb_S))^(1 / (par.cb_S - 1))

"""
    How much capital is produced by B firms if they are indifferent between
    producing one more unit or selling one unit less.

        AND

    The firm does not sell all of its capital! Otherwise producing one more unit has a different payoff (sell and keep the rest instead of keeping it fully).
"""
k_prod_U_buy(par) = ((par.p_I / (1 - par.τ)) / (par.ca_U * par.cb_U))^(1 / (par.cb_U - 1))
k_prod_U_sell(par) = (par.p_I / (par.ca_U * par.cb_U))^(1 / (par.cb_U - 1))

"""
    The price for an intermediate good is downwards sloping in the quantity produced.
"""
function price(Y_i, par::Pars_OLG)
    @unpack_Pars_OLG par
    Y^α_Y / (Y_i)^(1 / σ)
end

# """
#     The price for the B-good. Y_U is the quantity of B-goods.
# """
# function price_U(Y_U, par::Pars_OLG)
#     @unpack_Pars_OLG par
#     (Y / Y_U)^(1 / σ)
# end

"""
    The wage rate that leads to market clearing (l_S + l_U + l_Y = 1). Here, labor demand by each firm needs to be multiplied by the mass of the group of firms.
"""
function wage_rate(k_S, k_U, k_Y; par)
    @unpack_Pars_OLG par
    a = (1 / (1 - α_hat))
    α_hat * Y^α_Y * μ * ((γ * k_S^δ_hat)^a + ((1 - γ) * k_U^δ_hat)^a + (k_Y^δ_hat)^a)^(a^-1)
end

"""
    Profit for both type of firms.
"""
function profit(par)
    @unpack_Pars_OLG par
    k_T_S, k_T_U, k_T_Y = solve_firm_problems(par)

    k_S, l_S = capital_and_labor_S(k_T_S, par)
    k_U, l_U = capital_and_labor_U(k_T_U, par)
    k_Y, l_Y = capital_and_labor_Y(k_T_Y, par)

    k_P_S = k_S - k_T_S
    k_P_U = k_U - k_T_U

    @assert k_T_Y == k_Y

    #@assert k_T_U > 0

    k_P_Y = 0

    profit_S = Y^α_Y * l_S^α_hat * k_S^δ_hat - w * l_S - p_I * k_T_S - c_S(k_P_S, par)
    profit_U = Y^α_Y * l_U^α_hat * k_U^δ_hat - w * l_U - p_I / (1 - τ) * k_T_U - c_U(k_P_U, par)
    profit_Y = Y^α_Y * l_Y^α_hat * k_Y^δ_hat - w * l_Y - p_I / (1 - τ) * k_T_Y

    #@show k_T_S, k_T_U, k_T_Y, k_P_S, k_P_U
    return profit_S, profit_U, profit_Y
end

"""
    Solve the capital problem of each firm
"""
function solve_firm_problems(par)
    k_P_S = k_prod_S(par)
    min_val_S = -k_P_S + 0.00001
    f_S(x) = loss_S(x^2 + min_val_S, par)
    k_T_S = find_zero(f_S, 0.25)^2 + min_val_S

    # This part needs adjusting. I may be that the unsophisticated firm is in different regions!
    k_T_U = solve_U_problem(par)

    f_Y(x) = loss_Y(x^2, par)
    k_T_Y = find_zero(f_Y, 0.25)^2

    return k_T_S, k_T_U, k_T_Y
end

function solve_U_problem(par)
    k_P_U_buy = k_prod_U_buy(par)
    # Here I can improve and avoid setting bounds
    min_val_U_buy = -k_P_U_buy + 0.00001
    f_U_buy(x) = loss_U_buy(x^2 + min_val_U_buy, par)
    k_T_U_buy = find_zero(f_U_buy, 0.25)^2 + min_val_U_buy

    k_P_U_sell = k_prod_U_sell(par)
    min_val_U_sell = -k_P_U_sell + 0.00001
    f_U_sell(x) = loss_U_sell(x^2 + min_val_U_sell, par)
    k_T_U_sell = find_zero(f_U_sell, 0.25)^2 + min_val_U_sell

    #@show k_T_U_buy
    #@show k_T_U_sell

    #! Do some disambiguation
    if k_T_U_buy > 0 # First case: U firms actually buy at the buying price p_I / ( 1 - τ ).
        return k_T_U_buy
    elseif k_T_U_sell < 0 # Second case: U firms actually sell at the selling price p_I.
        return k_T_U_sell
    elseif k_T_U_buy < 0 && k_T_U_sell > 0 # Third case: Unsophisticated firms find it best to self-produce.
        return 0 # Unsophisticated firms do not trade
    end
end

function loss_S(k_T, par)
    @unpack_Pars_OLG par
    k, l = capital_and_labor_S(k_T, par)
    return FOC_K_S(k, l, par)
end

function loss_U_buy(k_T, par)
    @unpack_Pars_OLG par
    k, l = capital_and_labor_U_buy(k_T, par)
    return FOC_K_U_buy(k, l, par)
end

function loss_U_sell(k_T, par)
    @unpack_Pars_OLG par
    k, l = capital_and_labor_U_sell(k_T, par)
    return FOC_K_U_sell(k, l, par)
end

function loss_Y(k_T, par)
    @unpack_Pars_OLG par
    k, l = capital_and_labor_Y(k_T, par)
    return FOC_K_Y(k, l, par)
end

"""
    First order condition for capital employment sophisticated firms.

    Selling one unit less costs p_I.
"""
#! Why not solve analytically
function FOC_K_S(k, l, par)
    @unpack_Pars_OLG par
    p_I - δ_hat * Y^α_Y * l^α_hat * k^(δ_hat - 1)
end

"""
    First order condition for capital employment unsophisticated firms.

    Need to do some disambiguation in the case with three different types of firms:

    1. The unsophisticated firms buys capital from sophisticated firms.
    2. The unsophisticated firm completely self-produces data.
    3. The unsophisticated firm sells capital to young firms.

    Buying one unit more costs p_I / (1-τ).
"""
function FOC_K_U_buy(k, l, par)
    @unpack_Pars_OLG par
    p_I / (1 - τ) - δ_hat * Y^α_Y * l^α_hat * k^(δ_hat - 1)
end

function FOC_K_U_sell(k, l, par)
    @unpack_Pars_OLG par
    p_I - δ_hat * Y^α_Y * l^α_hat * k^(δ_hat - 1)
end

function FOC_K_U_prod(k, l, par)
    @unpack_Pars_OLG par
    ca_U*cb_U*k^(cb_U-1) - δ_hat * Y^α_Y * l^α_hat * k^(δ_hat - 1)
end

"""
    First order condition for capital employment young firms.

    Buying one unit more costs p_I / (1-τ).
"""
function FOC_K_Y(k, l, par)
    @unpack_Pars_OLG par
    p_I / (1 - τ) - δ_hat * Y^α_Y * l^α_hat * k^(δ_hat - 1)
end

"""
    Given how much capital the sophisticated firm buys, how much capital and labor does it employ.
"""
function capital_and_labor_S(k_T, par::Pars_OLG)
    k = total_capital(k_T, k_prod_S(par))
    l = labor_demand_S(k, par)
    return k, l
end

function capital_and_labor_U(k_T, par::Pars_OLG)
    if k_T > 0
        capital_and_labor_U_buy(k_T, par)
    elseif k_T < 0
        capital_and_labor_U_sell(k_T, par)
    elseif k_T == 0
        capital_and_labor_U_prod(par::Pars_OLG)
    end
end

function capital_and_labor_U_buy(k_T, par::Pars_OLG)
    k = total_capital(k_T, k_prod_U_buy(par))
    l = labor_demand_U(k, par)
    return k, l
end

function capital_and_labor_U_sell(k_T, par::Pars_OLG)
    k = total_capital(k_T, k_prod_U_sell(par))
    l = labor_demand_U(k, par)
    return k, l
end

function capital_and_labor_U_prod(par::Pars_OLG)
    f(k) = begin
        l = labor_demand_U(k, par)
        FOC_K_U_sell(k,l,par)
    end

    k = find_zero(f, [0.00001, 20])

    return k, labor_demand_U(k, par)
end

function capital_and_labor_Y(k_T, par::Pars_OLG)
    k = total_capital(k_T, 0)
    l = labor_demand_Y(k, par)
    return k, l
end

# capital_and_labor_U(k_T, par) = total_capital(k_T, k_prod_U(par)), labor_demand_U(k, par)
# capital_and_labor_Y(k_T, par) = total_capital(k_T, 0), labor_demand_Y(k, par)

function find_price_info(par)
    f(p_I) = begin
        @set! par.p_I = p_I^2
        k_T_S, k_T_U, k_T_Y = solve_firm_problems(par)
        return market_clearing(k_T_S, k_T_U, k_T_Y, par) # What one firm sells needs to be equal to what the other firm buys
    end
    find_zero(f, 0.25, tol = 1e-16)^2
end

function solve_agg_Y_iterative_unc(par; tol = 1e-10)
    diff = 10
    Y_guess = par.Y
    w_guess = par.w
    while diff > tol
        @set! par.Y = Y_guess
        #@show par.C
        Y_sol, w_sol = find_agg_Y(par)
        @set! par.w = w_sol
        diff = copy((Y_sol - Y_guess)^2 + (w_sol - w_guess)^2)
        if diff > tol
            Y_guess = copy(Y_sol)
            w_guess = copy(w_sol)
        else
        @set! par.w = w_sol
        @set! par.Y = Y_guess
        @set! par.p_I = find_price_info(par)
        @set! par.k_P_S_ss = k_prod_S(par)

    #@set! par.k_P_U_ss = k_prod_U(par)

        f_S(x) = loss_S(x, par)
        @set! par.k_T_S_ss = find_zero(f_S, [-par.k_P_S_ss + 0.00001, 5])

    #f_U(x) = loss_U(x, par)
    #@set! par.k_T_U_ss = find_zero(f_U, [-par.k_P_U_ss + 0.00001, 5])

        @set! par.k_T_U_ss = solve_U_problem(par)
        k_P_U_ss, l_U_ss = capital_and_labor_U(par.k_T_U_ss, par)
        @set! par.k_P_U_ss = k_P_U_ss

        f_Y(x) = loss_Y(x, par)
        @set! par.k_T_Y_ss = find_zero(f_Y, [0.00001, 5])

        @set! par.k_S_ss = total_capital(par.k_T_S_ss, par.k_P_S_ss)
        @set! par.k_U_ss = total_capital(par.k_T_U_ss, par.k_P_U_ss)
        @set! par.k_Y_ss = total_capital(par.k_T_Y_ss, 0)

        @set! par.l_S_ss = labor_demand_S(par.k_S_ss, par)
        @set! par.l_U_ss = labor_demand_U(par.k_U_ss, par)
        @set! par.l_Y_ss = labor_demand_Y(par.k_Y_ss, par)
        end
    end
    return par
end

firm_prod(l, k; par) = l^par.α_hat * k^par.δ_hat

function find_agg_Y(par)
    @set! par.p_I = find_price_info(par)
    k_T_S, k_T_U, k_T_Y = solve_firm_problems(par)

    k_S, l_S = capital_and_labor_S(k_T_S, par)
    k_U, l_U = capital_and_labor_U(k_T_U, par)
    k_Y, l_Y = capital_and_labor_Y(k_T_Y, par)

    Y_S = firm_prod(l_S, k_S; par)
    Y_U = firm_prod(l_U, k_U; par)
    Y_Y = firm_prod(l_S, k_S; par)

    return agg_prod(Y_S, Y_U, Y_Y, par), wage_rate(k_S, k_U, k_Y; par)
end

function agg_prod(Y_S, Y_U, Y_Y, par)
    @unpack_Pars_OLG par
    a = ((σ - 1) / σ)
    if α_Y > 0 # strategic complementarities
        (μ * Y_Y^a + μ * γ * Y_S^a + μ * (1 - γ) * Y_U^a)^(a^-1)
    elseif α_Y == 0 # no strategic complementarities
        μ * Y_Y^a + μ * γ * Y_S^a + μ * (1 - γ) * Y_U^a
    else
        error("α_Y < 0 was not planned for right now")
    end
end

"""
    Find the mass of entrants
"""
function find_mass(par; tol = 1e-16)

    f(x) = begin
        @set! par.μ = x^2
        par = solve_agg_Y_iterative_unc(par)
        profit_S, profit_U, profit_Y = profit(par)
        exp_profit = profit_Y + par.β * (par.γ * profit_S + (1 - par.γ) * profit_U)
        return (exp_profit - par.f)^2
    end

    find_zero(f, 1.5)^2
end


text:
function _precompile_()
    ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
    precompile(Tuple{typeof(Parsers.parse), Type{Int64}, String})
    precompile(Tuple{typeof(Parsers.parse), Type{Float64}, String})
    precompile(Tuple{typeof(Parsers.parse), Type{Date}, String})
end
_precompile_()


text:
import SeisIO.UW: readuwevt, uwdf, uwdf!, uwpf, uwpf!

printstyled("  UW\n", color=:light_green)
uwf0 = joinpath(path, "SampleFiles/UW/00*W")
uwf1 = joinpath(path, "SampleFiles/UW/99011116541")
uwf2 = joinpath(path, "SampleFiles/UW/94100613522o")
uwf3 = joinpath(path, "SampleFiles/UW/02062915175o")
uwf4 = joinpath(path, "SampleFiles/UW/00012502123W")
uwf5 = joinpath(path, "SampleFiles/UW/02062915205o")

S = verified_read_data("uw", uwf0)
S1 = SeisData()
uwdf!(S1, joinpath(path, "SampleFiles/UW/00012502123W"), false, false, false, 0)
[@test S.src[i] == abspath(uwf0) for i in 1:S.n]
S.src = S1.src
@test S == S1

printstyled("  readuwevt\n", color=:light_green)

# Can we read from pickfile only? datafile only?
printstyled("    read (pickfile, datafile) from pickfile name\n", color=:light_green)
W = readuwevt(uwf1*"o")

printstyled("    read (pickfile, datafile) from datafile name\n", color=:light_green)

redirect_stdout(out) do
  W = readuwevt(uwf1*"W", v=3)
end

for i in ["UW.WWVB..TIM","UW.TCG..TIM","UW.TDH..EHZ","UW.VLM..EHZ"]
  @test !isempty(findall(W.data.id.==i))
  n = findfirst(W.data.id.==i)
  @test ≈(W.data.fs[n], 100.0)
end

# Can we read from filename stub?
printstyled("    read (pickfile, datafile) from filename stub\n", color=:light_green)
W = readuwevt(uwf1)
@test W.hdr.id == W.source.eid == "99011116541"
@test W.hdr.mag.val == 3.0f0
@test occursin("99011116541o", W.hdr.src)
@test W.hdr.ot == DateTime("1999-01-11T16:54:11.96")

S = breaking_seis()
n = S.n
append!(S, convert(SeisData, W.data))
@test S.n == n + W.data.n

δt = μs*(rem(W.hdr.ot.instant.periods.value*1000 - SeisIO.dtconst, 60000000))

i = findfirst(W.data.id.=="UW.TDH..EHZ")
pha = W.data[i].pha["P"].tt
@test ≈(pha + δt, 14.506)

i = findfirst(W.data.id.=="UW.VLM..EHZ")
pha = W.data[i].pha["S"].tt
@test ≈(pha + δt, 24.236)

i = findfirst(W.data.id.=="UW.VFP..EHZ")
@test W.data.misc[i]["dur"] == 116.0

printstyled("    pickfile handling\n", color=:light_green)

# What about when there is no data file?
W = readuwevt(uwf2)
@test W.hdr.mag.val == 0.9f0
@test occursin("94100613522o", W.hdr.src)
@test W.hdr.ot == DateTime("1994-10-06T13:52:39.02")

W = readuwevt(uwf3)
@test W.hdr.id == "041568"

printstyled("    data file with a time correction structure\n", color=:light_green)
redirect_stdout(out) do
  W = readuwevt(uwf4, v=2)
end

printstyled("    pick file with nonnumeric error info\n", color=:light_green)
redirect_stdout(out) do
  W = readuwevt(uwf5, v=2)
end


text:
export Martin

struct Martin <: LakeModel end

@with_kw_noshow mutable struct MartinParameters <: LakeParameters
    @deftype Float64
    nutrients = 0.7 # current nutrient level
    ib = 2e-5 # g⋅m⁻²⋅day⁻¹ immigration rate of bream (9e-3 g⋅m⁻²⋅year⁻¹)
    ip = 2e-5 # g⋅m⁻²⋅day⁻¹ immigration rate of pike (9e-3 g⋅m⁻²⋅year⁻¹)
    r = 7.5e-3 # day⁻¹ maximum growth rate of bream (2.74 year⁻¹)
    H₁ = 0.5 # half saturation constant
    H₂ = 11 # % half saturation constant
    H₃ = 20 # g⋅m⁻² half saturation constant
    H₄ = 15 # g⋅m⁻² half saturation constant
    cb = 7.5e-5 # m⁻²⋅g⁻¹⋅day⁻¹ intraspecific competition constant for bream (0.0274 m⁻²⋅g⁻¹⋅year⁻¹)
    cp = 2.75e-4 # m⁻²⋅g⁻¹⋅day⁻¹ intraspecific competition constant for pike (0.1 m⁻²⋅g⁻¹⋅year⁻¹)
    prmax = 5e-2 # day⁻¹ maximum predation rate of pike (18.25 year⁻¹)
    ce = 0.1 # pike food conversion efficiency to growth
    mp = 2.25e-3 # day⁻¹ mortality rate of pike (0.82 year⁻¹)
    rv = 7e-3 # day⁻¹ vetegation growth rate
    cv = 6e-5 # m² intraspecific competition for vetetation
    mv = 7e-3 # day⁻¹ mortality of vegetation
    pv = 0.0 # day planting rate
    tb = 0.0 # day trawling rate
end

function lake_dynamics!(du, u, p::MartinParameters, t)
    B, P, V = u # g⋅m⁻² bream/pike/vegetation density

    FR = B^2 / (B^2 + p.H₄^2) # fuctional response of pike
    # Wrapped Cauchy distribution, simplifed to a yearly growth cycle between
    # March & July.
    # Expanded: γ*(1-ρ^2)/(1+ρ^2-2ρ*cos(2π*(t-μ)/τ))
    # NOTE: Having some issues with this, it's probably not necessary at all anyway.
    #plant = p.pv > 0.0 ? -p.pv / (p.pv + 19.4872 * cos(2π / 365.0 * (t - 122.0))) : 0.0

    du[1] =
        dB =
            p.ib + p.r * (p.nutrients / (p.nutrients + p.H₁)) * B - p.tb * B - p.cb * B^2 -
            p.prmax * FR * P
    du[2] = dP = p.ip + p.ce * p.prmax * FR * P * (V / (V + p.H₂)) - p.mp * P - p.cp * P^2
    du[3] = dV = p.rv * V + p.pv * V - p.cv * V^2 - p.mv * (V * B^2 / (p.H₃^2 + B^2))
end

"""
    lake_initial_state(nutrients, bream, pike, vegetation, Martin)

Provides initial conditions of the lake with specified values.
"""
function lake_initial_state(
    nutrients::Float64,
    bream::Float64,
    pike::Float64,
    vegetation::Float64,
    ::Type{Martin};
    kwargs...,
)
    ([bream, pike, vegetation], MartinParameters(; nutrients = nutrients, kwargs...))
end

"""
    lake_initial_state(nutrients, <:LakeDefinition, Martin)

A special override that returns interpolated results from precaclulated `Clear` and
`Turbid` bifurcation results. Not as precise as the preset conditions, since the
numbers do not come directly from the analytical solution, but close.
"""
function lake_initial_state(
    nutrients::Float64,
    ::Type{L},
    ::Type{Martin};
    kwargs...,
) where {L<:LakeDefinition}
    initial_state = interpolated_position(nutrients, L)
    (initial_state, MartinParameters(; nutrients = nutrients, kwargs...))
end

LakeParameters(::Type{Martin}, nutr; kwargs...) =
    MartinParameters(; nutrients = nutr, kwargs...)

preset_conditions(::Type{Clear}, ::Type{Martin}) = (0.7, [20.5172, 1.7865, 56.8443])
preset_conditions(::Type{Turbid}, ::Type{Martin}) = (2.5, [83.0128, 0.0414705, 6.40048])
preset_conditions(::Type{X1}, ::Type{Martin}) = (2.2, [36.730, 2.87725, 26.6798]) #deep unstable
preset_conditions(::Type{X2}, ::Type{Martin}) = (1.05, [59.0606, 0.819124, 12.0023]) #unstable
preset_conditions(::Type{X3}, ::Type{Martin}) = (1.05, [64.0559, 0.374008, 10.3631]) #stable

preset_conditions(::Type{S1}, ::Type{Martin}) = (2.0, [79.597, 0.050, 6.928]) #turbid transition, stable
preset_conditions(::Type{S2}, ::Type{Martin}) = (3.5, [85.436, 0.0373, 6.061]) #turbid eutrophic, stable
preset_conditions(::Type{S3}, ::Type{Martin}) = (3.5, [20.0, 3.5, 50.0]) #clear eutrophic, unstable
preset_conditions(::Type{T1}, ::Type{Martin}) = (0.9, [21.647, 2.050, 53.726]) #clear oligotrophic, high pike concentration, stable
preset_conditions(::Type{T2}, ::Type{Martin}) = (2.0, [26.901, 2.809, 41.530]) #clear transition, stable
preset_conditions(::Type{T3}, ::Type{Martin}) = (3.5, [20.0, 3.5, 50.0]) #clear eutrophic unstable

function Base.show(io::IO, ::MIME"text/plain", p::MartinParameters)
    println(io, "Parameters for lake dynamics (Martin) model):")
    println(io, "Nutrient level: $(p.nutrients)")
    println(io, "Immigration rate (g⋅m⁻²⋅day⁻¹) for bream: $(p.ib), pike: $(p.ip)")
    println(
        io,
        "Growth rate (day⁻¹) for bream: $(p.r), Predation rate (day⁻¹) of pike: $(p.prmax)",
    )
    println(
        io,
        "Half sturation constants: H₁ $(p.H₁), H₂ $(p.H₂)%, H₃ $(p.H₃) g⋅m⁻², H₄ $(p.H₄) g⋅m⁻²",
    )
    println(
        io,
        "Intraspecific competition constant (g⋅m⁻²⋅day⁻¹) for bream: $(p.cb), pike: $(p.cp)",
    )
    println(
        io,
        "Pike food conversion efficiency to growth: $(p.ce), Mortalitiy rate (day⁻¹): $(p.mp)",
    )
    println(
        io,
        "Vegetation rates. Growth: $(p.rv) day⁻¹, mortality: $(p.mv) day⁻¹, competition: $(p.cv) m²",
    )
end

function interpolated_position(nutrients::Float64, ::Type{Turbid})
    @assert 1.04 <= nutrients <= 4.0
    [turbid_b(nutrients), turbid_p(nutrients), turbid_v(nutrients)]
end

function interpolated_position(nutrients::Float64, ::Type{Clear})
    @assert 0.0 <= nutrients <= 2.49
    [clear_b(nutrients), clear_p(nutrients), clear_v(nutrients)]
end

# Higher order interpolations were a bit jumpy on the ends. Since we account for this already in the
# grid, there's no real point doing much more than linear interpolation in these data.
const turbid_n = [1.04, 1.06, 1.08, 1.1, 1.12, 1.14, 1.16, 1.18, 1.2, 1.22, 1.24, 1.26, 1.28, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0, 2.1, 2.2, 2.3, 2.4, 2.8, 2.9, 3.0, 3.1, 3.2, 3.3, 3.4, 3.5, 3.6, 3.7, 3.8, 3.9, 4.0]
const turbid_b = LinearInterpolation([62.8433, 64.7811, 65.812, 66.5967154617846, 67.256, 67.8374, 68.3651, 68.8529, 69.3096380295164, 69.7411, 70.1515, 70.5439, 70.9205, 71.2830815128419, 72.927620226924, 74.3575735324184, 75.6261770482514, 76.7652279581881, 77.7965051730699, 78.7361588920137, 79.5967630382622, 80.3884228650353, 81.1194401014282, 81.7967455181591, 82.4261974404837, 84.5557605975222, 85.0087274579361, 85.4355006715746, 85.838305423903, 86.2191189576847, 86.5797047356448, 86.9216409781937, 87.2463446598708, 87.5550918073417, 87.849034761985, 88.1292169339137, 88.3965854697749, 88.6520021756033], turbid_n)
const turbid_p = LinearInterpolation([0.467635, 0.324727, 0.264457, 0.226438572671889, 0.199488, 0.179158, 0.163257, 0.150262, 0.139566501222224, 0.130573, 0.122817, 0.116197, 0.110491, 0.105282876267212, 0.086680414110511, 0.075007876580953, 0.066965656681847, 0.06106181970993, 0.056627123629523, 0.053158110844447, 0.05023878194799, 0.048013362468672, 0.046125185620789, 0.044221413526291, 0.042868037308363, 0.038674849804663, 0.038099615841367, 0.037304090525497, 0.036805912699434, 0.036048380454724, 0.035464861024496, 0.034753892992486, 0.034831406585964, 0.034185163193823, 0.03418612062571, 0.033622695610838, 0.032871258345589, 0.032617613483335], turbid_n)
const turbid_v = LinearInterpolation([10.7297, 10.1524, 9.86357, 9.6515996372934, 9.47861, 9.32976, 9.1976, 9.07784, 8.96776090326038, 8.86554, 8.76989, 8.67984, 8.59468, 8.51383001117352, 8.16073696425393, 7.87085225645183, 7.62612135404186, 7.4157645933494, 7.2325673056776, 7.07136843964007, 6.92831478514543, 6.80044251320062, 6.68542064458636, 6.58138361478245, 6.48681686627753, 6.18128018305646, 6.11902788139653, 6.0612162952174, 6.00738633014564, 5.95714022580627, 5.91013159862657, 5.86605737616413, 5.82465121270593, 5.78567807610134, 5.74892976796485, 5.71422119441633, 5.68138724304635, 5.6502801549636], turbid_n)

const clear_n = [0.0, 0.04, 0.06, 0.08, 0.1, 0.12, 0.14, 0.16, 0.18, 0.2, 0.22, 0.24, 0.26, 0.28, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0, 2.1, 2.2, 2.3, 2.4, 2.42, 2.44, 2.46, 2.48, 2.49]
const clear_b = LinearInterpolation([0.0, 7.19822, 10.2447, 12.7489, 14.3417899783703, 15.1403, 15.6098, 15.9564, 16.2464, 16.5043680973213, 16.7416, 16.9638, 17.1743, 17.3752, 17.5681143534832, 18.4409411744765, 19.2038264444943, 19.8895014439516, 20.5171998487337, 21.0999740586095, 21.6473927297074, 22.1668544296341, 22.6643378784209, 23.1448798085877, 23.6129093438797, 24.0725068928104, 24.5276307687052, 24.9823468293504, 25.4411004976985, 25.9090895977944, 26.3928443866286, 26.9012392076907, 27.4474789812928, 28.0536139494919, 28.7632478716902, 29.6944205304346, 29.9376, 30.2195, 30.5678, 31.0733, 31.6065], clear_n)
const clear_p = LinearInterpolation([0.0111003, 0.0142037, 0.023943, 0.0485642, 0.105562347388399, 0.190455, 0.282928, 0.374336, 0.462298, 0.546230121593651, 0.626104, 0.702078, 0.774371, 0.843219, 0.908850307777514, 1.19586574516796, 1.42898170571722, 1.6226301580547, 1.78649633470377, 1.92730359110776, 2.04986983505402, 2.15774343788958, 2.2535990505578, 2.33949577013192, 2.41704615336661, 2.48753542023639, 2.55200022043938, 2.6112914060315, 2.66611434179448, 2.71706494128768, 2.76465434572004, 2.80932964683044, 2.85149627675782, 2.89154743171501, 2.92990689414219, 2.96718342309789, 2.97461, 2.98208, 2.98966, 2.99752, 3.00196], clear_n)
const clear_v = LinearInterpolation([109.96, 103.287, 92.4176, 82.9578, 77.0475035887902, 74.1649, 72.5014, 71.2896, 70.2871, 69.4037351989893, 68.5992, 67.8522, 67.1506, 66.4863, 65.8539683061925, 63.0572420785788, 60.7016894081954, 58.6565110822716, 56.8443323127489, 55.2131743708244, 53.7257083300374, 52.3539004349192, 51.0759395776379, 49.8743098980489, 48.7344963127633, 47.6440564653745, 46.5919019868462, 45.5676811402653, 44.5611694769215, 43.5615585328629, 42.5564678917703, 41.5303333021048, 40.4613483013786, 39.3146251420862, 38.023089573663, 36.4083114241203, 36.0011, 35.536, 34.9723, 34.1743, 33.3579], clear_n)


text:
module HartreeFock

using LinearAlgebra
using SpecialFunctions
using Printf
using BenchmarkTools

include("constants.jl")
include("Vec3.jl")
include("STOnG.jl")
include("link_coulomb.jl")
include("molecule.jl")
include("gto_integral.jl")
include("hf.jl")

export a0, hartree,
       Vec3, vec3,
       STONG, STO3G,
       Molecule, Atom,
       hartree_fock

end # module


text:
module RootedTrees

using LinearAlgebra

import Base: show, isless, ==, iterate, copy

export rootedtree, RootedTreeIterator

export α, β, γ, σ, order, residual_order_condition, elementary_weight, derivative_weight

export count_trees

"""
    RootedTree

Represents a rooted tree using its level sequence.

Reference:
  Beyer, Terry, and Sandra Mitchell Hedetniemi.
  "Constant time generation of rooted trees."
  SIAM Journal on Computing 9.4 (1980): 706-712.
"""
mutable struct RootedTree{T<:Integer, V<:AbstractVector}
  level_sequence::V
  iscanonical::Bool
end

function RootedTree(level_sequence::AbstractVector, iscanonical=false)
  T = eltype(level_sequence)
  V = typeof(level_sequence)
  RootedTree{T,V}(level_sequence, iscanonical)
end

"""
    rootedtree

Construct a canonical `RootedTree` object from a level sequence.

Reference:
  Beyer, Terry, and Sandra Mitchell Hedetniemi.
  "Constant time generation of rooted trees."
  SIAM Journal on Computing 9.4 (1980): 706-712.
"""
rootedtree(level_sequence::AbstractVector) = canonical_representation(RootedTree(level_sequence))

iscanonical(t::RootedTree) = t.iscanonical
#TODO: Validate rooted tree in constructor?

copy(t::RootedTree) = RootedTree(copy(t.level_sequence), t.iscanonical)

#  #function RootedTree(sequence::Vector{T}, valid::Bool)
#  function RootedTree(sequence::Array{T,1})
#    length(sequence) < 1 && throw(ArgumentError("Rooted trees must have a root, in particular at least one element!"))#
#
#    ## If there is only one element, the sequence must be valid.
#    #if !valid && length(sequence) > 1
#    #  # Test, whether there is exactly one root element at the beginning of sequence, if necessary.
#    #  root = sequence[1]
#    #  for level in sequence[2:end]
#    #    level <= root && throw(ArgumentError("Rooted trees must have exactly one element at root-level at the beginning."))
#    #  end
#    #end
#    # If there is only one element, the sequence must be valid.
#    if length(sequence) > 1
#      # Test, whether there is exactly one root element at the beginning of sequence, if necessary.
#      root = sequence[1]
#      for level in sequence[2:end]
#        level <= root && throw(ArgumentError("Rooted trees must have exactly one element at root-level at the beginning."))
#      end
#    end
#
#    new(sequence)
#  end
#RootedTree{T<:Integer}(sequence::Vector{T}) = RootedTree{T}(sequence, false)

function show(io::IO, t::RootedTree{T}) where {T}
  print(io, "RootedTree{", T, "}: ")
  show(io, t.level_sequence)
end

# comparison

"""
    isless(t1::RootedTree, t2::RootedTree)

Compares two rooted trees using a lexicographical comparison of their level sequences.
"""
function isless(t1::RootedTree, t2::RootedTree)
  isless(t1.level_sequence, t2.level_sequence)
end

function ==(t1::RootedTree, t2::RootedTree)
  t1.level_sequence == t2.level_sequence
end

# generation and canonical representation
"""
    canonical_representation!(t::RootedTree)

Change the representation of the rooted tree `t` to the canonical one, i.e. the
one with lexicographically biggest level sequence.
"""
function canonical_representation!(t::RootedTree)
  subtr = subtrees(t)
  for i in eachindex(subtr)
    canonical_representation!(subtr[i])
  end
  sort!(subtr, rev=true)

  i = 2
  for τ in subtr
    t.level_sequence[i:i+order(τ)-1] = τ.level_sequence[:]
    i += order(τ)
  end
  t.iscanonical = true

  t
end

"""
    canonical_representation(t::RootedTree)

Returns a new tree using the canonical representation of the rooted tree `t`,
i.e. the one with lexicographically biggest level sequence.
"""
function canonical_representation(t::RootedTree)
  canonical_representation!(copy(t))
end

"""
    RootedTreeIterator{T<:Integer}

Iterator over all rooted trees of given `order`. The returned trees are views to
an internal tree modified during the iteration. If the returned trees shall be
stored or modified during the iteration, a `copy` has to be made.
"""
struct RootedTreeIterator{T<:Integer}
  order::T
  t::RootedTree{T,Vector{T}}

  function RootedTreeIterator(order::T) where {T<:Integer}
    new{T}(order, RootedTree(Vector{T}(one(T):order), true))
  end
end

function iterate(iter::RootedTreeIterator{T}) where {T}
  iter.t.level_sequence[:] = one(T):iter.order
  (iter.t, false)
end

function iterate(iter::RootedTreeIterator{T}, state) where {T}
  state && return nothing

  two = iter.t.level_sequence[1] + one(T)
  p = 1
  q = 1
  @inbounds for i in 2:length(iter.t.level_sequence)
    if iter.t.level_sequence[i] > two
      p = i
    end
  end
  p == 1 && return nothing

  level_q = iter.t.level_sequence[p] - one(T)
  @inbounds for i in 1:p
    if iter.t.level_sequence[i] == level_q
      q = i
    end
  end

  @inbounds for i in p:length(iter.t.level_sequence)
    iter.t.level_sequence[i] = iter.t.level_sequence[i - (p-q)]
  end

  (iter.t, false)
end

"""
    count_trees(order)

Counts all rooted trees of `order`.
"""
function count_trees(order)
  order < 1 && throw(ArgumentError("The `order` must be at least one."))

  num = 0
  for _ in RootedTreeIterator(order)
    num += 1
  end
  num
end

# subtrees
struct Subtrees{T<:Integer} <: AbstractVector{RootedTree{T}}
  level_sequence::Vector{T}
  indices::Vector{T}
  iscanonical::Bool

  function Subtrees(t::RootedTree{T}) where {T}
    level_sequence = t.level_sequence
    indices = Vector{T}()

    start = 2
    i = 3
    while i <= length(level_sequence)
      if level_sequence[i] <= level_sequence[start]
        push!(indices, start)
        start = i
      end
      i += 1
    end
    push!(indices, start)

    # in order to get the stopping index for the last subtree
    push!(indices, length(level_sequence)+1)

    new{T}(level_sequence, indices, iscanonical(t))
  end
end

Base.size(s::Subtrees) = (length(s.indices)-1, )
Base.getindex(s::Subtrees, i::Int) = RootedTree(view(s.level_sequence, s.indices[i]:s.indices[i+1]-1), s.iscanonical)

"""
    subtrees(t::RootedTree)

Returns a vector of all subtrees of `t`.
"""
function subtrees(t::RootedTree{T}) where {T}
  subtr = typeof(t)[]

  if length(t.level_sequence) < 2
    return subtr
  end

  start = 2
  i = 3
  while i <= length(t.level_sequence)
    if t.level_sequence[i] <= t.level_sequence[start]
      push!(subtr, RootedTree(t.level_sequence[start:i-1]))
      start = i
    end
    i += 1
  end
  push!(subtr, RootedTree(t.level_sequence[start:end]))
end

# functions on trees

"""
    order(t::RootedTree)

The `order` of a rooted tree, i.e. the length of it's level sequence.
"""
order(t::RootedTree) = length(t.level_sequence)

"""
    σ(t::RootedTree)

The symmetry `σ` of a rooted tree `t`, i.e. the order of the group of automorphisms
on a particular labelling (of the vertices) of `t`.

Reference: Section 301 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function σ(t::RootedTree)
  if order(t) == 1 || order(t) == 2
    return 1
  end

  if !iscanonical(t)
    t = canonical_representation(t)
  end

  subtr = Subtrees(t)
  sym = 1
  num = 1

  @inbounds for i in 2:length(subtr)
    if subtr[i] == subtr[i-1]
      num += 1
    else
      sym *= factorial(num) * σ(subtr[i-1])^num
      num = 1
    end
  end
  sym *= factorial(num) * σ(subtr[end])^num
end

"""
    γ(t::RootedTree)

The density `γ(t)` of a rooted tree, i.e. the product over all vertices of `t`
of the order of the subtree rooted at that vertex.

Reference: Section 301 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function γ(t::RootedTree)
  if order(t) == 1
    return 1
  elseif order(t) == 2
    return 2
  end

  subtr = Subtrees(t)
  den = order(t)
  for τ in subtr
    den *= γ(τ)
  end
  den
end

"""
    α(t::RootedTree)

The number of monotonic labellings of `t` not equivalent under the symmetry group.

Reference: Section 302 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function α(t::RootedTree)
  div(factorial(order(t)), σ(t)*γ(t))
end

"""
    β(t::RootedTree)

The total number of labellings of `t` not equivalent under the symmetry group.

Reference: Section 302 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function β(t::RootedTree)
  div(factorial(order(t)), σ(t))
end

"""
    elementary_weight(t::RootedTree, A::AbstractMatrix, b::AbstractVector, c::AbstractVector)

Compute the elementary weight Φ(`t`) of `t` for the Butcher coefficients
`A, b, c` of a Runge-Kutta method.

Reference: Section 312 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function elementary_weight(t::RootedTree, A::AbstractMatrix, b::AbstractVector, c::AbstractVector)
  T = promote_type(promote_type(eltype(A), eltype(b)), eltype(c))
  elementary_weight(t, Matrix{T}(A), Vector{T}(b), Vector{T}(c))
end

function elementary_weight(t::RootedTree, A::AbstractMatrix{T}, b::AbstractVector{T}, c::AbstractVector{T}) where {T}
  dot(b, derivative_weight(t, A, b, c))
end

"""
    derivative_weight(t::RootedTree, A, b, c)

Compute the derivative weight (ΦᵢD)(`t`) of `t` for the Butcher coefficients
`A, b, c` of a Runge-Kutta method.

Reference: Section 312 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function derivative_weight(t::RootedTree, A, b, c)
  if order(t) == 1
    return zero(c) .+ one(eltype(c))
  end

  subtr = Subtrees(t)
  result = A * derivative_weight(subtr[1], A, b, c)
  for i in 2:length(subtr)
    tmp = A * derivative_weight(subtr[i], A, b, c)
    result = result .* tmp
  end
  return result
end

"""
    residual_order_condition(t::RootedTree, A, b, c)

The residual of the order condition
  `(Φ(t) - 1/γ(t)) / σ(t)`
with elementary weight `Φ(t)`, density `γ(t)`, and symmetry `σ(t)` of the
rooted tree `t` for the Runge-Kutta method with Butcher coefficients
`A, b, c`.

Reference: Section 315 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function residual_order_condition(t::RootedTree, A, b, c)
  ew = elementary_weight(t, A, b, c)
  T = typeof(ew)

  (ew - one(T) / γ(t)) / σ(t)
end

end # module


text:
struct Cell{T}
  cell::T
end

#
# cells
#

celltagprefix() = "c="
celltags(n::Integer) = TagSet(celltagprefix() * string(n))
celltags(n1::Integer, n2::Integer) = TagSet(celltagprefix() * n1 * "|" * n2)

indextagprefix() = "n="
#
# translatecell
#

# TODO: account for shifting by a tuple, for example:
# translatecell(ts"Site,c=1|2", (2, 3)) -> ts"Site,c=3|5"
# TODO: ts"c=10|12" has too many characters
# TODO: ts"c=1|2|3" has too many characters
#

# Determine the cell `n` from the tag `"c=n"`
function getcell(ts::TagSet)
  celltag = tag_starting_with(ts, celltagprefix())
  if isnothing(celltag) #dealing with link legs
    return celltag
  end
  return parse(Int, celltag[(length(celltagprefix()) + 1):end])
end

function getsite(ts::TagSet)
  celltag = tag_starting_with(ts, indextagprefix())
  return parse(Int, celltag[(length(indextagprefix()) + 1):end])
end

function translatecell(ts::TagSet, n::Integer)
  ncell = getcell(ts)
  if isnothing(ncell)
    return ts
  end
  return replacetags(ts, celltags(ncell) => celltags(ncell + n))
end

function translatecell(i::Index, n::Integer)
  ts = tags(i)
  translated_ts = translatecell(ts, n)
  return replacetags(i, ts => translated_ts)
end

function translatecell(is::Union{<:Tuple,<:Vector}, n::Integer)
  return translatecell.(is, n)
end

translatecell(T::ITensor, n::Integer) = ITensors.setinds(T, translatecell(inds(T), n))
translatecell(T::MPO, n::Integer) = translatecell.(T, n)

struct CelledVector{T} <: AbstractVector{T}
  data::Vector{T}
end
ITensors.data(cv::CelledVector) = cv.data

Base.convert(::Type{CelledVector{T}}, v::Vector) where {T} = CelledVector{T}(v)

function CelledVector{T}(::UndefInitializer, n::Integer) where {T}
  return CelledVector(Vector{T}(undef, n))
end

"""
    celllength(cv::CelledVector)

The length of a unit cell of a CelledVector.
"""
celllength(cv::CelledVector) = length(ITensors.data(cv))

# For compatibility with Base
Base.size(cv::CelledVector) = size(ITensors.data(cv))

"""
    cell(cv::CelledVector, n::Integer)

Which unit cell index `n` is in.
"""
function cell(cv::CelledVector, n::Integer)
  _cell = fld1(n, celllength(cv))
  return _cell
end

"""
    cellindex(cv::CelledVector, n::Integer)

Which index in the unit cell index `n` is in.
"""
cellindex(cv::CelledVector, n::Integer) = mod1(n, celllength(cv))

# Get the value at index `n`, where `n` must
# be within the first unit cell
_getindex_cell1(cv::CelledVector, n::Int) = ITensors.data(cv)[n]

# Set the value at index `n`, where `n` must
# be within the first unit cell
_setindex_cell1!(cv::CelledVector, val, n::Int) = (ITensors.data(cv)[n] = val)

# Fallback
translatecell(x, ::Integer) = x

function getindex(cv::CelledVector, n::Int)
  cellₙ = cell(cv, n)
  siteₙ = cellindex(cv, n)
  return translatecell(_getindex_cell1(cv, siteₙ), cellₙ - 1)
end

# Do we need this definition? Maybe uses generic Julia fallback
#getindex(cv::CelledVector, r::AbstractRange) = [cv[n] for n in r]

function Base.firstindex(cv::CelledVector, c::Cell)
  return (c.cell - 1) * celllength(cv) + 1
end

function Base.lastindex(cv::CelledVector, c::Cell)
  return c.cell * celllength(cv)
end

function Base.keys(cv::CelledVector, c::Cell)
  return firstindex(cv, c):lastindex(cv, c)
end

## function Base.eachindex(cv::CelledVector, c::Cell)
##   return firstindex(cv, c):lastindex(cv, c)
## end

getindex(cv::CelledVector, c::Cell) = cv[eachindex(cv, c)]

function setindex!(cv::CelledVector, T, n::Int)
  cellₙ = cell(cv, n)
  siteₙ = cellindex(cv, n)
  _setindex_cell1!(cv, translatecell(T, -(cellₙ - 1)), siteₙ)
  return cv
end

celltags(cell) = TagSet("c=$cell")

#
# TODO: This version accepts a more general translation function between
# unit cells
#

#struct CelledVector{T, F, FINV} <: AbstractVector{T}
#  data::Vector{T}
#  f::F
#end
#
#CelledVector(v::AbstractVector) = CelledVector(v, identity)
#
#cell_length(cv::CelledVector) = length(cv.data)
#
## Determine which cell the index is in
#cell(cv::CelledVector, n::Integer) = fld1(n, cell_length(cv))
#
## Determine the index in the cell where the index sits
#cell_index(cv::CelledVector, n::Integer) = mod1(n, cell_length(cv))
#
## Return the cell and the cell index
#cell_and_cell_index(cv::CelledVector, n::Integer) = fldmod1(n, cell_length(cv))
#
#function getindex(cv::CelledVector, n::Integer)
#  cellₙ, cell_indexₙ = cell_and_cell_index(cv, n)
#  return cv.f(cv.data[cell_indexₙ], cellₙ)
#end
#
#function setindex!(cv::CelledVector, val, n::Integer)
#  cellₙ, cell_indexₙ = cell_and_cell_index(cv, n)
#  # XXX: is this offset generally correct?
#  # It seems to be for "linear" functions.
#  cv.data[cell_indexₙ] = cv.f(val, -cellₙ + 2)
#  return cv
#end


text:
function build_simple_lp!(pre_opt::MOIU.CachingOptimizer)
    optim = MOIB.full_bridge_optimizer(pre_opt, Float64)
    MOI.empty!(optim)
    @test MOI.is_empty(optim)

    # add 10 variables - only diagonal is relevant
    X = MOI.add_variables(optim, 2)

    # add sdp constraints - only ensuring positivenesse of the diagonal
    vov = MOI.VectorOfVariables(X)

    c1 = MOI.add_constraint(optim,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(2.0, X[1]),
            MOI.ScalarAffineTerm(1.0, X[2])
        ], 0.0), MOI.EqualTo(4.0))

    c2 = MOI.add_constraint(optim,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[1]),
            MOI.ScalarAffineTerm(2.0, X[2])
        ], 0.0), MOI.EqualTo(4.0))

    b1 = MOI.add_constraint(optim,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[1])
        ], 0.0), MOI.GreaterThan(0.0))

    b2 = MOI.add_constraint(optim,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[2])
        ], 0.0), MOI.GreaterThan(0.0))

    MOI.set(optim,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], [X[1], X[2]]), 0.0)
        )
    MOI.set(optim, MOI.ObjectiveSense(), MOI.MIN_SENSE)
end

const optimizer_maxiter = MOIU.CachingOptimizer(cache, ProxSDP.Optimizer(max_iter = 1, log_verbose = false))
const optimizer_timelimit = MOIU.CachingOptimizer(cache, ProxSDP.Optimizer(time_limit = 0.0001, log_verbose = false))

@testset "MOI status" begin
    @testset "MOI.OPTIMAL" begin
        build_simple_lp!(optimizer)
        MOI.optimize!(optimizer)
        @test MOI.get(optimizer, MOI.TerminationStatus()) == MOI.OPTIMAL
        MOI.empty!(optimizer)
    end

    @testset "MOI.ITERATION_LIMIT" begin
        build_simple_lp!(optimizer_maxiter)
        MOI.optimize!(optimizer_maxiter)
        @test MOI.get(optimizer_maxiter, MOI.TerminationStatus()) == MOI.ITERATION_LIMIT
        MOI.empty!(optimizer_maxiter)
    end

    # @testset "MOI.TIME_LIMIT" begin
    #     build_simple_lp!(optimizer_timelimit)
    #     MOI.optimize!(optimizer_timelimit)
    #     @test MOI.get(optimizer_timelimit, MOI.TerminationStatus()) == MOI.TIME_LIMIT
    #     MOI.empty!(optimizer_timelimit)
    # end
end


text:
function preprocess(PData::PDataKTR, Hop, g, calls, max_calls)
    ζ, ξ, maxtol, mintol = PData.ζ, PData.ξ, PData.maxtol, PData.mintol
    nshifts = PData.nshifts
    shifts = PData.shifts

    n = length(g)
    gNorm2 = norm(g)
    # precision = max(1e-12, min(0.5, (gNorm2^ζ)))
    # Tolerance used in Assumption 2.6b in the paper ( ξ > 0, 0 < ζ ≤ 1 )
    cgatol = min(maxtol, ξ * gNorm2^(1 + ζ))
    cgatol = max(mintol, cgatol) # add some feasible limit
    cgrtol = min(maxtol, ξ * gNorm2^ζ)
    cgrtol = max(mintol, cgrtol) # add some feasible limit

    nshifts = length(shifts)
    solver = PData.solver
    cg_lanczos!(
        solver,
        Hop,
        -g,
        shifts,
        itmax = min(max_calls - sum(calls), 2 * n),
        atol = cgatol,
        rtol = cgrtol,
        verbose = 0,
        check_curvature = true,
    )

    PData.indmin = 0
    PData.positives .= solver.converged
    for i = 1:nshifts
        PData.xShift[i] .= solver.x[i]
        PData.norm_dirs[i] = norm(solver.x[i])
    end
    PData.shifts .= shifts
    PData.nshifts = nshifts
    PData.OK = sum(solver.converged) != 0 # at least one system was solved

    return PData
end


text:
module NowcastLSTM

include("Functions.jl")

using CSV, DataFrames, Dates, PyCall

for n in [names(CSV); names(DataFrames); names(Dates); names(PyCall)]
        @eval export $n
end

end


text:
@testset "Rates" begin
    @testset "rate types" begin
        rs = Yields.Rate.([0.1, 0.02], Yields.Continuous())
        @test rs[1] == Yields.Rate(0.1, Yields.Continuous())
        @test rs[1] == Yields.Continuous(0.1)
        @test rate(rs[1]) == 0.1
    end

    @testset "constructor" begin
        @test Yields.Continuous(0.05) == Yields.Rate(0.05, Yields.Continuous())
        @test Yields.Periodic(0.02, 2) == Yields.Rate(0.02, Yields.Periodic(2))

        @test Yields.Continuous()(0.05) == Yields.Rate(0.05, Yields.Continuous())
        @test Yields.Periodic(2)(0.02) == Yields.Rate(0.02, Yields.Periodic(2))

        @test Yields.Rate(0.02, 2) == Yields.Rate(0.02, Yields.Periodic(2))
        @test Yields.Rate(0.02, Inf) == Yields.Rate(0.02, Yields.Continuous())

    end

    @testset "rate conversions" begin
        m = Yields.Rate(0.1, Yields.Periodic(2))
        @test convert(Periodic(2),0.1) ≈ m
        @test convert(Periodic(2),m) ≈ m
        @test Periodic(m,2) ≈ m
        @test Periodic(2)(m) ≈ m
        @test convert(Yields.Continuous(), m) ≈ Yields.Rate(0.09758, Yields.Continuous()) atol = 1e-5
        @test Yields.Continuous(m) ≈ Yields.Rate(0.09758, Yields.Continuous()) atol = 1e-5
        @test Yields.Continuous()(m) ≈ Yields.Rate(0.09758, Yields.Continuous()) atol = 1e-5

        c = Yields.Rate(0.09758, Yields.Continuous())
        @test convert(Yields.Continuous(), c) == c
        @test convert(Yields.Continuous(), 0.09758) == c
        @test Yields.Continuous(c) == c
        @test Yields.Continuous()(c) == c
        @test convert(Yields.Periodic(2), c) ≈ Yields.Rate(0.1, Yields.Periodic(2)) atol = 1e-5
        @test Periodic(2)(c) ≈ Yields.Rate(0.1, Yields.Periodic(2)) atol = 1e-5
        @test Periodic(c,2) ≈ Yields.Rate(0.1, Yields.Periodic(2)) atol = 1e-5
        @test convert(Yields.Periodic(2), c) ≈ Yields.Rate(0.1, Yields.Periodic(2)) atol = 1e-5
        @test convert(Yields.Periodic(2), c) ≈ Yields.Rate(0.1, Yields.Periodic(2)) atol = 1e-5
        @test convert(Yields.Periodic(4), m) ≈ Yields.Rate(0.09878030638383972, Yields.Periodic(4)) atol = 1e-5

    end

    @testset "rate equality" begin
        a = Yields.Periodic(0.02, 2)
        a_eq = Yields.Periodic((1+.02/2)^2-1, 1)
        b = Yields.Periodic(0.03, 2)
        c = Yields.Continuous(0.02)

        @test a == a
        @test !(a == a_eq) # not equal due to floating point error
        @test a ≈ a_eq
        @test a != b
        @test ~(a ≈ b)
        @test (a ≈ a)
        @test ~(a ≈ c)

    end

    @testset "discounting and accumulation" for t in [-1.3, 2.46, 6.7]

        unspecified_rate = 0.035
        periodic_rate = Yields.Periodic(0.02, 2)
        continuous_rate = Yields.Continuous(0.03)

        @test discount(unspecified_rate, t) ≈ (1 + 0.035)^(-t)
        @test discount(periodic_rate, t) ≈ (1 + 0.02 / 2)^(-t * 2)
        @test discount(continuous_rate, t) ≈ exp(-0.03 * t)

        @test accumulation(unspecified_rate, t) ≈ (1 + 0.035)^t
        @test accumulation(periodic_rate, t) ≈ (1 + 0.02 / 2)^(t * 2)
        @test accumulation(continuous_rate, t) ≈ exp(0.03 * t)

    end

    @testset "rate over interval" begin

        from = -0.45
        to = 3.4
        rate = 0.15

        @test discount(rate, from, to) ≈ discount(rate, to - from)
        @test accumulation(rate, from, to) ≈ accumulation(rate, to - from)

    end

    @testset "AbstractYield Interface" begin
        c = Continuous(0.03)
        p = Periodic(0.04,2)

        @test zero(c,2) ≈ c
        @test zero(p,2) ≈ p
        @test forward(c,2) ≈ c
        @test forward(p,2) ≈ p

        @test discount(c,2) ≈ exp(-2*0.03)
        @test discount(p,2) ≈ 1 / (1 + .04/2)^(2*2)

        @test discount(c,2) ≈ 1 / accumulation(c,2)
        @test discount(p,2) ≈ 1 / accumulation(p,2)

    end

    @testset "rate algebra" begin

        a = 0.03
        b = 0.02

        @testset "addition" begin
            c(x) = Yields.Continuous(x)
            p(x) = Yields.Periodic(x, 1)

            @test c(a) + b ≈ Yields.Continuous(0.05)
            @test a + c(b) ≈ Yields.Continuous(0.05)

            @test p(a) + b ≈ Yields.Periodic(0.05,1)
            @test a + p(b) ≈ Yields.Periodic(0.05,1)
        end

        @testset "multiplication" begin
            c(x) = Yields.Continuous(x)
            p(x) = Yields.Periodic(x, 1)

            @test c(a) * b ≈ Yields.Continuous(a * b)
            @test a * c(b) ≈ Yields.Continuous(a * b)

            @test p(a) * b ≈ Yields.Periodic(a * b,1)
            @test a * p(b) ≈ Yields.Periodic(a * b,1)
        end

        @testset "division" begin
            c(x) = Yields.Continuous(x)
            p(x) = Yields.Periodic(x, 1)

            @test c(a) / b ≈ Yields.Continuous(a / b)
            @test_throws MethodError a / c(b) ≈ Yields.Continuous(a / b)

            @test p(a) / b ≈ Yields.Periodic(a / b,1)
            @test_throws MethodError a / p(b) ≈ Yields.Periodic(a / b,1)
        end

        @testset "subtraction" begin
            c(x) = Yields.Continuous(x)
            p(x) = Yields.Periodic(x, 1)

            @test c(a) - b ≈ Yields.Continuous(0.01)
            @test a - c(b) ≈ Yields.Continuous(0.01)

            @test p(a) - b ≈ Yields.Periodic(0.01,1)
            @test a - p(b) ≈ Yields.Periodic(0.01,1)
        end

        @testset "Rate and Rate" begin
            r = Yields.Periodic(0.04,2) - Yields.Periodic(0.01,2)
            @test r ≈ Yields.Periodic(0.03,2)
            r = Yields.Periodic(0.04,2) + Yields.Periodic(0.01,2)
            @test r ≈ Yields.Periodic(0.05,2)

            @test Yields.Periodic(0.04,1) > Yields.Periodic(0.03,2)
            @test Yields.Periodic(0.03,1) < Yields.Periodic(0.04,2)
            @test ~(Yields.Periodic(0.04,1) < Yields.Periodic(0.03,2))
            @test ~(Yields.Periodic(0.03,1) > Yields.Periodic(0.04,2))

            @test Yields.Periodic(0.03,1) < Yields.Periodic(0.03,2)
            @test Yields.Periodic(0.03,100) < Yields.Continuous(0.03)
            @test Yields.Periodic(0.03,2) > Yields.Periodic(0.03,1)
            @test Yields.Continuous(0.03) > Yields.Periodic(0.03,100)
        end
    end

end


text:
# The first task is to find the variables to make the mean vector:
find_variables(x) = find_variables!(Symbol[], x) #this is so that we can call the function without the exclamation
function find_variables!(var_names, x::Number) #if the variable name is a number then just return it without doing anything
    return var_names
end
function find_variables!(var_names, x::Symbol) # if the variable name is a symbol then push it to the list of var_names because its a name of a column
    push!(var_names, x)
end

function find_variables!(var_names, x::Expr) # if the variable is a expression object then we have to crawl through each argument of the expression
    # safety checking
    if x.head == :call  # check for + symbol bc we are summing linear combinations within each expression argument
      # pass the remaining expression
      for argument in x.args[2:end] # since the first argument is the :+ call
        find_variables!(var_names, argument) #recursively find the names given in each argument so check if number, if symbol, if expression etc.again agian
    end
end
return var_names
end

function search_variables!(x::Expr, var::Symbol)
    for i in eachindex(x.args)
        if x.args[i] == var # if the argument is one of the variables given then just put it in the right format df[:x1]
            x.args[i] = Meta.parse(string(:input_data_from_user,"[", ":", var, "]"))
        elseif x.args[i] isa Expr # else if the argument is an expression (i.e not a varaible (symbol) or a number) then
            search_variables!(x.args[i], var) #go through this function recursively on each of the arguments of the expression object
        end
    end
    return x
end

function search_variables!(x::Expr, vars...) # this is for when you have more than one variable name found in the string
    for var in vars #goes through each of the variables in the vector vars
        x = search_variables!(x, var) #runs the recursion on each variable in vars
    end
    return x
end

function mean_formula(user_formula_string::String, df::DataFrame)
    global input_data_from_user = df

    users_formula_expression = Meta.parse(user_formula_string)
    if(users_formula_expression isa Expr)
        found_markers = find_variables(users_formula_expression) #store the vector of symbols of the found variables

        dotted_args = map(Base.Broadcast.__dot__, users_formula_expression.args) # adds dots to the arguments in the expression
        dotted_expression = Expr(:., dotted_args[1], Expr(:tuple, dotted_args[2:end]...)) #reformats the exprssion arguments by changing the variable names to tuples of the variable names to keep the dot structure of julia

        julia_interpretable_expression = search_variables!(dotted_expression, found_markers...) #gives me the julia interpretable exprsesion with the dataframe provided

        mean_vector = eval(Meta.parse(string(julia_interpretable_expression))) #evaluates the julia interpretable expression on the dataframe provided
    else
        mean_vector = [users_formula_expression for i in 1:size(df, 1)]
    end
    return mean_vector
end

max_stars_repo_name: sarah-ji/Tutorials

max_issues_repo_name: sarah-ji/Tutorials

max_forks_repo_name: sarah-ji/Tutorials


text:
#add each of the individual melt rate models
#include("./analytic_melt_rate_model.jl")
include("./binfile_melt_rate.jl")
include("./plume_emulator.jl")
include("./pico.jl")
include("./quadratic_melt_rate.jl")
include("./mismip_melt_rate.jl")

##### default temperature and salinity profiles #####
"""
    two_layer_ambient()

Returns a function that, in turn, returns a value v_low below a depth d_low, and a value v_hi above a depth d_hi > d_low, with a linear interpolation between these value.
"""
function two_layer_function(z;v_low,v_hi,d_low,d_hi)
    @assert d_low < d_hi
    if  z< d_low
        return v_low
    elseif z > d_hi
        return v_hi
    else #depth between d_low and d_hi
        return v_low + (v_hi - v_low)/(d_hi - d_low) * (z - d_low)
    end
end

isomip_warm0_salinity(z) = two_layer_function(z, v_low = 34.6, v_hi = 34.0, d_low = -700, d_hi = -300)
isomip_warm0_temp(z) = two_layer_function(z, v_low =1.2, v_hi = -1.0, d_low = -700, d_hi = -300)

struct UniformMeltRate{T <: Real} <: AbstractMeltRate
    m :: T #uniform melt rate applied everywhere
end

UniformMeltRate(; m = 0.0) = UniformMeltRate(m)

"""
    update_melt_rate(melt_rate::UniformMeltRate, fields, grid)

Update the melt rate when for the UniformMeltRate type

"""
function update_melt_rate!(melt_rate::UniformMeltRate, fields, grid)
    @unpack basal_melt = fields.gh
    basal_melt .= melt_rate.m
    return nothing
end


text:
function include_parameters(canonical_model::CanonicalModel,
                            data::Matrix,
                            param_reference::UpdateRef,
                            axs...)

    _add_param_container!(canonical_model, param_reference, axs...)
    param = par(canonical_model, param_reference)

    Cidx = CartesianIndices(length.(axs))

    for idx in Cidx
        param.data[idx] = PJ.add_parameter(canonical_model.JuMPmodel, data[idx])
    end

    return

end

function include_parameters(canonical_model::CanonicalModel,
                            ts_data::Vector{Tuple{String, Int64, Float64, Vector{Float64}}},
                            param_reference::UpdateRef,
                            expression::Symbol,
                            multiplier::Float64 = 1.0)

    time_steps = model_time_steps(canonical_model)
    _add_param_container!(canonical_model, param_reference, (r[1] for r in ts_data), time_steps)
    param = par(canonical_model, param_reference)
    expr = exp(canonical_model, expression)

    for t in time_steps, r in ts_data
        param[r[1], t] = PJ.add_parameter(canonical_model.JuMPmodel, r[4][t]);
        _add_to_expression!(expr, r[2], t, param[r[1], t], r[3] * multiplier)
    end

    return

end

############################ injection expression with parameters ####################################

########################################### Devices ####################################################

function _nodal_expression_param(canonical_model::CanonicalModel,
                                devices::IS.FlattenIteratorWrapper{G},
                                system_formulation::Type{S}) where {G<:PSY.Generator,
                                                                    S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)
    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(devices))
    ts_data_reactive = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(devices))

    for (ix, d) in enumerate(devices)
        bus_number = PSY.get_number(PSY.get_bus(d))
        tech = PSY.get_tech(d)
        name = PSY.get_name(d)
        pf = sin(acos(PSY.get_powerfactor(PSY.get_tech(d))))
        time_series_vector = ones(time_steps[end])
        ts_data_active[ix] = (name, bus_number, PSY.get_rating(tech), time_series_vector)
        ts_data_reactive[ix] = (name, bus_number, PSY.get_rating(tech) * pf, time_series_vector)
    end

    include_parameters(canonical_model,
                    ts_data_active,
                    UpdateRef{G}(Symbol("P_$(G)")),
                    :nodal_balance_active)
    include_parameters(canonical_model,
                    ts_data_reactive,
                    UpdateRef{G}(Symbol("Q_$(G)")),
                    :nodal_balance_reactive)

    return

end

function _nodal_expression_param(canonical_model::CanonicalModel,
                                devices::IS.FlattenIteratorWrapper{G},
                                system_formulation::Type{S}) where {G<:PSY.Generator,
                                                                    S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical_model)
    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(devices))

    for (ix, d) in enumerate(devices)
        bus_number = PSY.get_number(PSY.get_bus(d))
        tech = PSY.get_tech(d)
        name = PSY.get_name(d)
        time_series_vector = ones(time_steps[end])
        ts_data_active[ix] = (name, bus_number, PSY.get_rating(tech), time_series_vector)
    end

    include_parameters(canonical_model,
                    ts_data_active,
                    UpdateRef{G}(Symbol("P_$(G)")),
                    :nodal_balance_active)

    return

end

############################################## Time Series ###################################
function _nodal_expression_param(canonical_model::CanonicalModel,
                                 forecasts::Vector{PSY.Deterministic{G}},
                                 system_formulation::Type{S}) where {G<:PSY.Generator,
                                                                     S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)
    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(forecasts))
    ts_data_reactive = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(forecasts))

    for (ix, f) in enumerate(forecasts)
        device = PSY.get_component(f)
        bus_number = PSY.get_number(PSY.get_bus(device))
        tech = PSY.get_tech(device)
        name = PSY.get_name(device)
        pf = sin(acos(PSY.get_powerfactor(PSY.get_tech(device))))
        time_series_vector = values(PSY.get_timeseries(f))
        ts_data_active[ix] = (name, bus_number, PSY.get_rating(tech), time_series_vector)
        ts_data_reactive[ix] = (name, bus_number, PSY.get_rating(tech) * pf, time_series_vector)
    end

    include_parameters(canonical_model,
                    ts_data_active,
                    UpdateRef{G}(Symbol("P_$(G)")),
                    :nodal_balance_active)
    include_parameters(canonical_model,
                    ts_data_reactive,
                    UpdateRef{G}(Symbol("Q_$(G)")),
                    :nodal_balance_reactive)

    return

end

function _nodal_expression_param(canonical_model::CanonicalModel,
                                forecasts::Vector{PSY.Deterministic{G}},
                                system_formulation::Type{S}) where {G<:PSY.Generator,
                                                                    S<:PM.AbstractActivePowerModel}

    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, length(forecasts))

    for (ix, f) in enumerate(forecasts)
        device = PSY.get_component(f)
        device = PSY.get_component(f)
        bus_number = PSY.get_number(PSY.get_bus(device))
        tech = PSY.get_tech(device)
        name = PSY.get_name(device)
        time_series_vector = values(PSY.get_timeseries(f))
        ts_data_active[ix] = (name, bus_number, PSY.get_rating(tech), time_series_vector)
    end

    include_parameters(canonical_model,
                    ts_data_active,
                    UpdateRef{G}(Symbol("P_$(G)")),
                    :nodal_balance_active)

    return

end

############################ injection expression with fixed values ####################################
########################################### Devices ####################################################
function _nodal_expression_fixed(canonical_model::CanonicalModel,
                                devices::IS.FlattenIteratorWrapper{G},
                                system_formulation::Type{S}) where {G<:PSY.Generator,
                                                                     S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)

    for t in time_steps, d in devices
        bus_number = PSY.get_number(PSY.get_bus(d))
        active_power = PSY.get_rating(PSY.get_tech(d))
        reactive_power = active_power * sin(acos(PSY.get_powerfactor(PSY.get_tech(d))))
        _add_to_expression!(canonical_model.expressions[:nodal_balance_active],
                            bus_number,
                            t,
                            active_power)
        _add_to_expression!(canonical_model.expressions[:nodal_balance_reactive],
                            bus_number,
                            t,
                            reactive_power)
    end

    return

end

function _nodal_expression_fixed(canonical_model::CanonicalModel,
                                    devices::IS.FlattenIteratorWrapper{G},
                                    system_formulation::Type{S}) where {G<:PSY.Generator,
                                                                         S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical_model)

    for t in time_steps, d in devices
        bus_number = PSY.get_number(PSY.get_bus(d))
        active_power = PSY.get_rating(PSY.get_tech(d))
        _add_to_expression!(canonical_model.expressions[:nodal_balance_active],
                            bus_number,
                            t,
                            active_power)
    end

    return

end

############################################## Time Series ###################################
function _nodal_expression_fixed(canonical_model::CanonicalModel,
                                forecasts::Vector{PSY.Deterministic{G}},
                                system_formulation::Type{S}) where {G<:PSY.Generator,
                                                                    S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)

    for f in forecasts
        device = PSY.get_component(f)
        bus_number = PSY.get_number(PSY.get_bus(device))
        active_power = PSY.get_rating(PSY.get_tech(device))
        reactive_power = active_power * sin(acos(PSY.get_powerfactor(PSY.get_tech(device))))
        time_series_vector = values(PSY.get_timeseries(f))
        for t in time_steps
            _add_to_expression!(canonical_model.expressions[:nodal_balance_active],
                                bus_number,
                                t,
                                time_series_vector[t] * active_power)
            _add_to_expression!(canonical_model.expressions[:nodal_balance_reactive],
                                bus_number,
                                t,
                                time_series_vector[t] * reactive_power)
        end
    end

    return

end

function _nodal_expression_fixed(canonical_model::CanonicalModel,
                                forecasts::Vector{PSY.Deterministic{G}},
                                system_formulation::Type{S}) where {G<:PSY.Generator,
                                                                    S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical_model)

    for f in forecasts
        device = PSY.get_component(f)
        bus_number = PSY.get_number(PSY.get_bus(device))
        active_power = PSY.get_rating(PSY.get_tech(device))
        time_series_vector = values(PSY.get_timeseries(f))
        for t in time_steps
            _add_to_expression!(canonical_model.expressions[:nodal_balance_active],
                                bus_number,
                                t,
                                time_series_vector[t] * active_power)
        end
    end

    return

end


text:
# -------- Serializing (Async) Players --------------------------------------- #

function _serialize(p :: Union{IntuitionPlayer, MCTSPlayer})

  # Convert the playing model to cpu
  m = playing_model(p)
  m = isa(m, Async) ? worker_model_to_cpu(m) : to_cpu(m)

  # Temporarily replace the model by DummyModel
  pt = switch_model(p, DummyModel())

  (pt, decompose(m), on_gpu(training_model(p)))

end

function _deserialize(p :: Union{IntuitionPlayer, MCTSPlayer}, dm, gpu)

  # Reconstruct the model and bring it to the GPU if wanted
  m = compose(dm)
  gpu && (m = isa(m, Async) ? worker_model_to_gpu(m) : to_gpu(m))

  # Replace DummyModel by the reconstructed model
  switch_model(p, m)

end

_on_gpu(splayers) = any(p[3] for p in splayers)

# -------- Distributed Calculations ------------------------------------------ #

function ticket_sizes(n, m)
  ns = zeros(Int, m)
  ns .+= floor(Int, n / m)
  ns[1:(n % m)] .+= 1
  ns
end

take_ticket!(tic) = try take!(tic) catch _ 0 end

# TODO: We currently assume that we stay on one machine. Should not be too hard
# to generalize this to work on quite arbitrary cluster-systems.
# For this, we should provide the 'with_workers' function with more information,
# e.g., which workers will land on nodes with how many GPUs.
function with_workers( f :: Function
                     , players
                     , n
                     , args...
                     ; workers = workers()
                     , tickets = length(workers)
                     , kwargs... )

  # Count the number of workers
  m = length(workers)

  # Generate channels to provide worker tickets and receive results
  tic = RemoteChannel(() -> Channel(tickets))
  res = RemoteChannel(() -> Channel(m))

  # Divide the n tasks as evenly as possible and fill the ticket supply
  foreach(s -> put!(tic, s), ticket_sizes(n, tickets))

  # Fill the buffer asynchronously
  task = @async map(_ -> take!(res), 1:tickets)

  # Serialize the players
  splayers = _serialize.(players)

  # Count the GPUs
  gpu = _on_gpu(splayers)
  devices = Knet.cudaGetDeviceCount()
  m > devices && gpu && @info "Multiple workers will share one GPU device" maxlog = 1

  # Start let the workers work on the tickets
  map(1:m) do i

    @spawnat workers[i] begin

      # Make adjustments to the environment, like number of threads or GPU id
      # TODO: This does not seem to work, as Knet.gpu is not known on the process
      # if this call is hidden in set_defaults.
      # set_defaults(i, devices, _on_gpu(splayers))
      if gpu
        Knet.gpu((i-1) % devices)
      end

      # Reconstruct the players
      ps = [_deserialize(p...) for p in splayers]

      # Wait for a ticket. If we get one, carry it out. If there are no tickets
      # left, the ticket channel is closed and we end the loop.
      while (n = take_ticket!(tic)) != 0

        put!(res, f(ps, n, args...; kwargs...))

      end

    end

  end

  # Fetch the data with error handling if something unexpected happens
  data = try

    fetch(task)

  catch exn

    close(res); close(tic)
    interrupt(workers)
    rethrow(exn)

  end

  # Close the channels
  close(res); close(tic)

  # Return the data
  data

end

max_stars_repo_name: roSievers/jtac