text:
if basename(pwd()) == "aoc"
    cd("2017/20")
end
include("../../helpers/inherit.jl")

using StaticArrays: MVector
using StatsBase: countmap

struct Particle
    id::Int
    position::MVector{3}
    velocity::MVector{3}
    acceleration::MVector{3}
end
@accessors Particle

Particle(id::Integer, xs::Vararg{Vector}) = Particle(id, MVector{3}.(xs)...)
function Particle(id::Integer, str::AbstractString)
    Particle(id, (parse.(Int, m) for m in eachmatch(r"(-?\d+),(-?\d+),(-?\d+)", str))...)
end

function advance!(particle::Particle)
    particle.velocity .+= particle.acceleration
    particle.position .+= particle.velocity
    particle
end

manhattan(xs) = sum(abs.(xs))

function sorter(left::Particle, right::Particle)
    lₐ = manhattan(left.acceleration)
    rₐ = manhattan(right.acceleration)
    if lₐ == rₐ
        lᵥ = manhattan(left.velocity)
        rᵥ = manhattan(right.velocity)
        if lᵥ == rᵥ
            manhattan(left.position) <= manhattan(right.position)
        else
            lᵥ < rᵥ
        end
    else
        lₐ < rₐ
    end
end

function loadparticles(filename::AbstractString)
    [Particle(id - 1, line) for (id, line) in enumerate(eachline(filename))]
end

part1(filename::AbstractString) = part1(loadparticles(filename))
part1(particles) = id(first(sort(particles, lt = sorter)))

@assert part1("example.txt") == 0

part2(filename::AbstractString) = part2(loadparticles(filename))

function part2(particles, limit = 25)
    len = length(particles)
    lastchange = 0
    while lastchange <= limit
        foreach(advance!, particles)
        sort!(particles, by = position)
        collisions = sort!(collect(keys(filter!(kv -> kv[2] > 1, countmap(position.(particles))))))
        if isempty(collisions)
            lastchange += 1
        else
            filter!(particle -> !insorted(position(particle), collisions), particles)
            diff = len - length(particles)
            len -= diff
            println(diff, " particle", isone(diff) ? "" : "s", " destroyed. ", len, " particle", isone(len) ? "" : "s", " remain", isone(len) ? "s" : "", ".")
            lastchange = 0
        end
    end
end

max_stars_repo_name: stellartux/advent-of-code

max_issues_repo_name: stellartux/advent-of-code

max_forks_repo_name: stellartux/advent-of-code


text:
function getVAFresult(simulation::Simulation, rng::AbstractRNG=Random.GLOBAL_RNG; read_depth=100.0,
    detectionlimit=5/read_depth, cellularity=1.0)

    trueVAF = getallelefreq(simulation)
    sampledVAF = sampledallelefreq(trueVAF, rng, read_depth=read_depth,
        detectionlimit=detectionlimit, cellularity=cellularity)
    freq, freqp = subclonefreq(simulation.output)

    return VAFResult(
        read_depth,
        cellularity,
        detectionlimit,
        simulation.input,
        trueVAF,
        sampledVAF,
        freq,
        freqp
    )
end

function getVAFresult(multisim::MultiSimulation, rng::AbstractRNG=Random.GLOBAL_RNG; read_depth=100.0,
    detectionlimit=5/read_depth, cellularity=1.0)

    trueVAFs = Array{Float64, 1}[]
    sampledVAFs = Array{Float64, 1}[]
    freqs = Array{Float64, 1}[]
    freqps = Array{Float64, 1}[]

    for moduletracker in multisim
        trueVAF = getallelefreq(moduletracker)
        sampledVAF = sampledallelefreq(trueVAF, rng, read_depth=read_depth,
            detectionlimit=detectionlimit, cellularity=cellularity)
        freq, freqp = subclonefreq(moduletracker)
        push!(trueVAFs, trueVAF)
        push!(sampledVAFs, sampledVAF)
        push!(freqs, freq)
        push!(freqps, freqp)
    end

    return VAFResultMulti(
        read_depth,
        cellularity,
        detectionlimit,
        simulation.input,
        trueVAFs,
        sampledVAFs,
        freqs,
        freqps
    )
end

function getallelefreq(simulation::Simulation)
    return getallelefreq(simulation.output, simulation.input.ploidy)
end

function getallelefreq(moduletracker::ModuleTracker, ploidy)
    mutations = cellsconvert(moduletracker.cells).mutations
    return getallelefreq(mutations, moduletracker.Nvec[end], ploidy)
end

function getallelefreq(mutations, N, ploidy=2)
    allelefreq = counts(mutations,minimum(mutations):maximum(mutations))
    # idx = f .> 0.01 #should this be f/(2*cellnum) .> 0.01, i.e. only include freq > 1% ??
    allelefreq = map(Float64, allelefreq)
    allelefreq ./= (ploidy * N) #correct for ploidy
end

function cellsconvert(cells)
    #convert from array of cell types to one array with mutations and one array with cell fitness

    clonetype = zeros(Int64,length(cells))
    mutations = Int64[]
    # sizehint!(mutations, length(cells) * 10) #helps with performance to provide vector size

    for i in 1:length(cells)
        append!(mutations,cells[i].mutations)
        clonetype[i] = cells[i].clonetype
    end

    return (;mutations, clonetype)
end

"""
    sampledhist(trueVAF::Array{Float64,1}, cellnum::Int64; <keyword arguments>)

Create synthetic experimental data by smapling from the true VAF distribution
according to experimental constraints.

### Keyword arguments
- `detectionlimit = 5/read_depth`: minimum VAF which can be detected
- `read_depth = 100`: expected number of reads for each allele
- `cellularity = 1.0`:

"""
function sampledallelefreq(trueVAF::Array{Float64, 1}, rng::AbstractRNG;
    read_depth=100.0, detectionlimit=5/read_depth, cellularity=1.0)

    VAF = trueVAF * cellularity
    filter!(x -> x > detectionlimit, VAF)
    # Pois(np) ≈ B(n,p) when n >> p, thus we can sample from Poisson(read_depth)
    # rather than Binomial(cellnum, read_depth/cellnum), length(VAF)) as D << N^2
    depth = rand(rng, Poisson(read_depth), length(VAF))
    sampalleles = map((n, p) -> rand(rng, Binomial(n, p)), depth, VAF)
    sampledVAF = sampalleles ./ depth
    return sampledVAF
end

"""
    gethist(VAF::Array{Float64,1}; fmin = 0.0, fmax = 1.0, fstep = 0.001)

Fit VAF data into bins `fmin`:`fstep`:`fmax` and return DataFrame with columns
:VAF, :freq [= ``m(f)``], :cumfreq [= ``M(f)``].

Note ``M(f)`` is the cummulative number of mutations with frequency ``f``, i.e. the
number of mutations with frequency in ``(f, 1)``, while ``m(f)`` is the number of mutations with
frequency in ``(f - fstep, f)``.

"""
function gethist(VAF::Array{Float64,1}; fmin = 0.0, fmax = 1, fstep = 0.001)
    x = fmin:fstep:fmax
    y = fit(Histogram, VAF, x, closed=:right).weights
    dfhist = DataFrame(VAF = x[2:end], freq = y)
    dfhist = addcumfreq!(dfhist,:freq)
    return dfhist
end

function addcumfreq!(df,colname)
    df[!, Symbol(:cum,colname)] = reverse(cumsum(reverse(df[!, colname])))
    return df
end

function subclonefreq(moduletracker)
    #get proportion of cells in each subclone
    clonesize = getclonesize(moduletracker)
    clonefreqp = clonesize[2:end]/sum(clonesize)
    clonefreq = copy(clonefreqp)
    if length(clonefreqp) > 1
        clonefreq, subclonalmutations =
            calculateclonefreq!(clonefreq, subclonalmutations, moduletracker.subclones)
    end
    return clonefreq, clonefreqp
end

ext: jl


text:
# Factor Graph OS type utilities
#  IIF methods should direclty detect extended types from user import
# of convert in their namespace

"""
    savefgjld(fgl::FactorGraph; file::AbstractString="tempfg.jld")

Save mostly complete Factor Graph type by converting complicated FunctionNodeData
types to 'Packed' types using user supplied converters. Ground truth can also be
saved and recovered by the associated loadjld(file="tempfg.jld") method.
"""
function savejld(fgl::FactorGraph;
      file::AbstractString="tempfg.jld",
      groundtruth=nothing)
  fgs = encodefg(fgl)
  if groundtruth == nothing
    @save file fgs
  else
    @save file fgs groundtruth
  end
  return file
end

"""
    loadjld(file="tempfg.jld")

Opposite of savejld(fg, gt=gt, file="tempfg.jl") to load data from file. This function
uses the unpacking converters for converting all PackedInferenceType to FunctorInferenceType.
"""
function loadjld(;file::AbstractString="tempfg.jld")
  dd = JLD.load(file)
  # fgs = jldopen(file,"r") do file
  #   read(file, "fgs")
  # end
  fgs, gt = nothing, nothing
  if haskey(dd, "fgs")
    fgs = dd["fgs"]
  else
    error("No factor graph (fgs) data found in this file, only found $(keys(dd))")
  end
  if haskey(dd, "groundtruth")
    gt = dd["groundtruth"]
    println("Also found ground truth data")
  end
  fgd = decodefg(fgs)
  return fgd, gt
end

function ls(fgl::FactorGraph, lbl::Symbol; api::DataLayerAPI=dlapi, ring::Int=1)
  # TODO ring functionality must still be implemented
  lsa = Symbol[]
  # v = nothing
  if haskey(fgl.IDs, lbl)
    id = fgl.IDs[lbl]
  else
    return lsa
  end
  # this is unnecessary
  v = getVert(fgl,id, api=api)
  for outn in api.outneighbors(fgl, v)
    # if outn.attributes["ready"] = 1 && outn.attributes["backendset"]=1
      push!(lsa, Symbol(outn.label))
    # end
  end
  return lsa
end
ls(fgl::FactorGraph, lbl::T) where {T <: AbstractString} = ls(fgl, Symbol(lbl))

function ls(fgl::FactorGraph)
  k = collect(keys(fgl.IDs))
  l = Int[]
  x = Int[]
  for kk in k
    kstr = string(kk)
    val = parse(Int,kstr[2:end]) # kk
    if kstr[1] == 'l'
      push!(l,val)
    elseif kstr[1] == 'x'
      push!(x,val)
    end
  end
  l = sort(l)
  x = sort(x)
  ll = Array{Symbol,1}(length(l))
  xx = Array{Symbol,1}(length(x))
  for i in 1:length(l)
    ll[i] = Symbol(string("l",l[i]))
  end
  for i in 1:length(x)
    xx[i] = Symbol(string("x",x[i]))
  end
  return xx,ll
end

function lsf(fgl::FactorGraph, lbl::Symbol; api::DataLayerAPI=dlapi)
  lsa = Symbol[]
  # v = Union{}
  if haskey(fgl.fIDs, lbl)
    id = fgl.fIDs[lbl]
  else
    return lsa
  end
  v = getVert(fgl, id, api=api) # fgl.g.vertices[id] #fgl.f[id]
  for outn in api.outneighbors(fgl, v) # out_neighbors(v, fgl.g)
    push!(lsa, Symbol(outn.label))
  end
  return lsa
end
lsf{T <: AbstractString}(fgl::FactorGraph, lbl::T) = lsf(fgl,Symbol(lbl))

function lsf{T <: FunctorInferenceType}(fgl::FactorGraph,
      mt::Type{T};
      api::DataLayerAPI=dlapi  )
  #
  syms = Symbol[]
  for (fsym,fid) in fgl.fIDs
    if typeof(getfnctype(fgl, fid, api=api))==T
      push!(syms, fsym)
    end
  end
  return syms
end

function ls2(fgl::FactorGraph, vsym::Symbol)
  xxf = ls(fgl, vsym)
  xlxl = Symbol[]
  for xf in xxf
    xx = lsf(fgl,xf)
    xlxl = union(xlxl, xx)
  end
  xlxl = setdiff(xlxl, [vsym])
  return xlxl
end

hasOrphans(fg) = sum(length.(ls.(fg, [ls(fg)[1];ls(fg)[2]])) .== 0) > 0

"""
    landmarks(fgl::FactorGraph, vsym::Symbol)

Return Vector{Symbol} of landmarks attached to vertex vsym in fgl.
"""
function landmarks(fgl::FactorGraph, vsym::Symbol)
  fsyms = ls(fgl, vsym)
  lms = Symbol[]
  for fs in fsyms
    for varv = lsf(fgl, fs)
      if string(varv)[1] == 'l'
        push!(lms, varv)
      end
    end
  end
  lms
end

function evalLikelihood(fg::FactorGraph, sym::Symbol, point::Vector{Float64})
  p = getVertKDE(fg, sym)
  Ndim(p) == length(point) ? nothing : error("point (dim=$(length(point))) must have same dimension as belief (dim=$(Ndim(p)))")
  evaluateDualTree(p, reshape(point,:,1))[1]
end

# Evaluate the likelihood of an Array{2} of points on the marginal belief of some variable
# note the dimensions must match
function evalLikelihood(fg::FactorGraph, sym::Symbol, points::Array{Float64,2})
  p = getVertKDE(fg, sym)
  Ndim(p) == size(points,1) ? nothing : error("points (dim=$(size(points,1))) must have same dimension as belief (dim=$(Ndim(p)))")
  evaluateDualTree(p, (points))
end

#


text:
using MOCNeutronTransport
@testset "Quadrilateral_2D" begin
    for F in [Float32, Float64]
        @testset "Constructors" begin
            p₁ = Point_2D(F, 0)
            p₂ = Point_2D(F, 1)
            p₃ = Point_2D(F, 1, 1)
            p₄ = Point_2D(F, 0, 1)
            quad = Quadrilateral_2D(SVector(p₁, p₂, p₃, p₄))
            @test quad.points == SVector(p₁, p₂, p₃, p₄)

            # single constructor
            quad = Quadrilateral_2D(p₁, p₂, p₃, p₄)
            @test quad.points == SVector(p₁, p₂, p₃, p₄)
        end

        @testset "Methods" begin
            p₁ = Point_2D(F, 0)
            p₂ = Point_2D(F, 1)
            p₃ = Point_2D(F, 1, 1)
            p₄ = Point_2D(F, 0, 1)
            quad = Quadrilateral_2D(SVector(p₁, p₂, p₃, p₄))

            # interpolation
            @test quad(0, 0) ≈ p₁
            @test quad(1, 0) ≈ p₂
            @test quad(1, 1) ≈ p₃
            @test quad(0, 1) ≈ p₄
            @test quad(1//2, 1//2) ≈ Point_2D(F, 1//2, 1//2)

            # area
            a = area(quad)
            @test typeof(a) == F
            @test a == F(1)

            # in
            p = Point_2D(F, 1//2, 1//10)
            @test p ∈  quad
            p = Point_2D(F, 1//2, -1//10)
            @test p ∉ quad

            # 3 intersections
            l = LineSegment_2D(p₃, p₁)
            ipoints, points = intersect(l, quad)
            @test ipoints == 3
            @test points[1] ≈ p₁
            @test points[2] ≈ p₃
            @test points[3] ≈ p₃

            # 2 intersections
            l = LineSegment_2D(Point_2D(F, 0, 1//2), Point_2D(F, 1, 1//2))
            ipoints, points = intersect(l, quad)
            @test ipoints == 2
            @test points[1] ≈ Point_2D(F, 1, 1//2)
            @test points[2] ≈ Point_2D(F, 0, 1//2)

            # 0 intersections
            l = LineSegment_2D(Point_2D(F, -1, -1), Point_2D(F, 2, -1))
            ipoints, points = intersect(l, quad)
            @test ipoints == 0
        end
    end
end

max_stars_repo_name: KyleVaughn/MOCNeutronTransport

max_issues_repo_name: KyleVaughn/MOCNeutronTransport

max_forks_repo_name: KyleVaughn/MOCNeutronTransport


text:
# # Amplifier circuit
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/OpAmp.ipynb)
#
#md # !!! note "Overview"
#md #     System type: linear continuous system with piecewise-linear input\
#md #     State dimension: 2\
#md #     Application domain: Electrical Engineering
#
# ## Operational amplifiers
#
# Operational amplifiers ("op-amp") are electronic devices to amplify an input
# voltage signal. They are widely used to construct filters that remove a
# desired range of frequencies from the input signal.

# The schematic diagram of an op-amp is shown below. It has two terminals:
# - input side (left)
# - output terminal (right)

# ![](https://github.com/JuliaReach/ReachabilityAnalysis.jl/blob/master/examples/OpAmp/opamp.png?raw=true)

# The output voltage is ``e_o = K(e_B - e_A)``, where ``K`` is the voltage gain
# of the op-amp. ``K`` is usually vey large (the order of ``10^5 V / V``).

# In the *ideal op-amp* scenario, the following assumptions hold:
# - The input terminals of the op-amp draw negligible current.
# - The voltage difference at the input terminals ``e_B - e_A`` is zero.
# - The gain ``K`` is infinite.

# ## Inverting amplifier

# Consider the following circuit, which consists of an ideal op-amp in the so-called inverting
# configuration.

# ![](https://raw.githubusercontent.com/JuliaReach/ReachabilityAnalysis.jl/master/examples/OpAmp/opamp_c.png)

# It is not hard to show [^KLU15] that the output satisfies the following ODE:

# ```math
#    \dfrac{de_o(t)}{dt} = -\dfrac{1}{R_2 C}e_o(t) - \dfrac{1}{R_1C}e_{in}(t).
# ```
# The circuit exhibits a transient behavior because of the capacitor connected in parallel
# to the op-amp. In the stationary limit ``t \to \infty``, ``\dfrac{de_o(t)}{dt} = 0`` and
# ``e_o = - \frac{R_2}{R_1} e_{in}``. The gain is thus ``- \frac{R_2}{R_1}`` and the minus sign
# is the reason that this circuit is called inverting amplifier.

# The output voltage ``e_o(t)`` depends on the input voltage ``e_{in}(t)``.
# In order to highlight some typical assumptions on the inputs, in the rest
# of this notebook we model the input voltage in different ways:

#  1. **Non-deterministic input** with bounded range. This is the most general case,
#     in which ``e_{in}(t) \in E_{in}``, with ``E_{in} \subseteq \mathbb{R}`` is an interval
#     for the set of admissible inputs.
#     This approach guarantees that the flowpipe covers all possible trajectories for any
#     bounded, measurable, function ``e_{in}(t)`` satisfying the constraint ``e_{in}(t) \in E_{in}`` for all
#     ``t`` of the given time span.
#
#  2. **Constant input,** a special case such that
#     ``e_{in}(t) \in E_{in}`` and ``e_{in}(t)`` is *constant* along trajectories.
#
#  3. **Time-varying input** i.e. such that it is given by a (possibly time-varying)
#     predefined function ``e_{in}(t)`` that is given by the solution of some ODE.
#     In particular, we consider linear and exponential input functions.
#
# We will consider one by one the different approaches described above in the simple
# inverting amplifier circuit.

# ## Non-deterministic input

# It is convenient to define constants ``α`` and ``β`` such that
#
# ``\dfrac{de_o(t)}{dt} = α e_o(t) + β e_{in}(t)``.
#
# In this section we assume that ``e_{in}(t)`` can be any (measurable) function
# whose range is contained in the set ``E_{in}``.

using ReachabilityAnalysis, Plots, LaTeXStrings #!jl

function opamp_nondet(; X0 = Interval(0.0, 0.0),
                        R₁ = 2., R₂ = 6., C = 1.e-3,
                        Ein = Interval(1.9, 2.1))

    α = hcat(-1/(R₂ * C))
    β = hcat(-1/(R₁ * C))

    ## continuous evolution
    invariant = Universe(1) # no state constraint
    sys = @system(eout' = α*eout + β*ein, ein ∈ Ein, eout ∈ invariant)

    ## initial-value problem
    return @ivp(sys, x(0) ∈ X0)
end

#-
sol_nondet = solve(opamp_nondet(), T=0.1, alg=INT(δ=1e-4));

#md # !!! tip "Performance tip"
#md #     We have chosen the algorithm `INT` as it is the best choice in terms of efficiency for one-dimensional ODEs.

fig = plot()
plot!(fig, sol_nondet, vars=(0, 1), xlab=L"t", ylab=L"e_{out}", title="Solution for non-deterministic input", lw=0.2)

# The solution `sol_nondet` assumed that the initial state is at the origin.
# Now we solve for three different initial conditions of increasing width.

fig = plot(xlab=L"t", ylab=L"e_{out}(t)", title="Solution for non-deterministic input")

Δt = 0 .. 0.04

sol = solve(opamp_nondet(X0=Interval(-1.0, 1.0)), T=0.1, alg=INT(δ=1e-4));
plot!(fig, sol(Δt), vars=(0, 1), lab="X0 = -1 .. 1", lw=0.2)

sol = solve(opamp_nondet(X0=Interval(-0.5, 0.5)), T=0.1, alg=INT(δ=1e-4));
plot!(fig, sol(Δt), vars=(0, 1), lab="X0 = -0.5 .. 0.5", lw=0.2)

sol = solve(opamp_nondet(), T=0.1, alg=INT(δ=1e-4));
plot!(fig, sol(Δt), vars=(0, 1), lab="X0 = 0", lw=0.2)

fig

#md # !!! tip "Visualization tip"
#md #     The order in which we have defined the plots is such that the flowpipe
#md #     with largest width comes first, so it doesn't overlap with those
#md #     which are strictly included in it. This way, the final plot more clearly
#md #     reveals the different flowpipes by their color.

# ## Constant input

# We model the input voltage ``e_{in}(t)`` with a family of ODEs of the form
#
# ``\dfrac{d e_{in}(t)}{dt} = \gamma e_{in}(t) + \delta``.
#
# There are two distinct behaviors depending on the parameters ``\gamma`` and ``\delta``:

#  - ``\gamma = \delta = 0``: ``e_{in}(t) = e_{in, 0}`` i.e. the input is *constant* over the given domain.
#  - ``\gamma = 0``: linear with rate ``\delta``, it is given by ``e_{in}(t) = \delta t + e_{in, 0}``.
#  - ``\gamma \neq 0``: exponential with rate ``\gamma``, it is given by ``e_{in}(t) = (e_{in, 0} + \delta / \gamma)e^{\gamma t} - \delta / \gamma``.

# Moreover, it is natural to assume that the input saturates at a given threshold value.
# The following function, `opamp_with_saturation`, models the system as a
# hybrid automaton with two modes: in the first mode, the input signal ``e_{in}(t)``
# varies according to the specified dynamics, and in the second mode, the input
# signal saturates and ceases to increase (or decrease, depding on the signs chosen).

# In this section we consider the constant input scenario,
# - ``e_{in}(0) ∈ E_{in, 0}``, ``e_{in}(t)`` constant.
# - ``\delta = 0``
# - ``\gamma = 0``

# In the following section we consider the other cases.

using Symbolics

function opamp_with_saturation(; X0 = BallInf(zeros(2), 0.0),
                                 R₁ = 2., R₂ = 6., C = 1.e-3,
                                 γ = 100., δ = 100., Es = 2.)
    var = @variables eₒ ein
    α = hcat(-1/(R₂ * C))
    β = hcat(-1/(R₁ * C))

    ## transition graph
    automaton = LightAutomaton(2)
    add_transition!(automaton, 1, 2, 1)

    ## modes
    A = [α β; 0. γ]
    b = [0.; δ]
    invariant = HalfSpace(ein <= Es, var)
    mode1 = @system(x' = Ax + b, x ∈ invariant)
    A = [α β; 0. 0.];
    mode2 = @system(x' = Ax, x ∈ Universe(2))
    modes = [mode1, mode2]

    ## transition mode1 -> mode2 (saturation)
    guard = Hyperplane(ein == Es, var)
    t1 = @map(x -> x, dim:2, x ∈ guard)
    resetmaps = [t1]

    H = HybridSystem(automaton, modes, resetmaps, [AutonomousSwitching()])

    ## initial condition is at the origin in mode 1
    initial_condition = [(1, X0)]
    return IVP(H, initial_condition)
end

# Observe that in the model defined above, the first coordinate of the state-space
# corresponds to the output ``e_o(t)``, while the second coordinate, ``e_{in}(t)``,
# corresponds to the input signal ``e_{in}(t)``.

# We specify that the initial condition is such that ``e_o(0)`` is at the origin,
# and ``e_{in}(0) ∈ 1.4 .. 1.6``.

X0 = Hyperrectangle([0., 1.5], [0.0, 0.1]);

# The value of the saturation constant ``Es`` is chosen to be sufficiently high
# such that the system doesn't transition to mode 2.

prob_const = opamp_with_saturation(X0=X0, γ = 0., δ = 0., Es = 3.)
sol_const = solve(prob_const, T=0.1, alg=BOX(δ=1e-4));

#-

plot(sol_const, vars=(0, 2), xlab=L"t", ylab=L"e_{in}(t)", title="Constant input signal", lw=0.2)
plot!(x->x, x->1.5, 0., 0.1, line=2, color="red", linestyle=:dash, legend=nothing)

#-

plot(sol_const, vars=(0, 1), xlab=L"t", ylab=L"e_{out}(t)", title="Solution for constant input", lw=0.2)

# ## Time-varying input

# For the time-varying with a prescribed function case, we consider two scenarios:

# - Linearly increasing input voltage: ``\delta \neq 0``, ``\gamma = 0``.
# - Exponentially increasing input voltage (up to saturation at `Es = 1.0`): ``\delta \neq 0``, ``\gamma < 0``.

# For both cases we consider that set of initial states is the origin:

X0 = Singleton(zeros(2));

#-

# We solve both instances by choosing som values of ``δ`` and ``γ``:

## linearly increasing input signal
prob_lin = opamp_with_saturation(X0=X0, γ=0., δ=100., Es=1.);
sol_lin = solve(prob_lin, T=0.1, alg=BOX(δ=1e-4));

## exponentially increasing input signal
prob_exp = opamp_with_saturation(X0=X0, γ=-100., δ=100., Es=3.);
sol_exp = solve(prob_exp, T=0.1, alg=BOX(δ=1e-4));

#-
plot(sol_lin, vars=(0, 2), xlab=L"t", ylab=L"e_{in}(t)", title="Linear input w/saturation", lw=0.2)
#-
plot(sol_exp, vars=(0, 2), xlab=L"t", ylab=L"e_{in}(t)", title="Exponential input", lw=0.2)

#-
plot(sol_lin, vars=(0, 1), xlab=L"t", ylab=L"e_{out}(t)", title="Solution for linear input", lw=0.2)
#-
plot(sol_exp, vars=(0, 1), xlab=L"t", ylab=L"e_{out}(t)", title="Solution for exponential input", lw=0.2)

# ## References

# [^KLU15]: Kluever, C. A. (2015). Dynamic systems: modeling, simulation, and control. John Wiley & Sons.


text:
DiffEqBase.@def iipnlcachefields begin
  nlcache = alg.nlsolve.cache
  @unpack κ,tol,max_iter,min_iter,new_W = nlcache
  z = similar(u,axes(u))
  dz = similar(u,axes(u)); tmp = similar(u,axes(u)); b = similar(u,axes(u))
  fsalfirst = zero(rate_prototype)
  k = zero(rate_prototype)
  uToltype = real(uBottomEltypeNoUnits)
  ηold = one(uToltype)

  nf = f isa SplitFunction && alg isa SplitAlgorithms ? f.f1 : f
  islin = (f isa ODEFunction && islinear(f.f)) || (f isa SplitFunction && islinear(f.f1.f))
  # check if `nf` is linear
  if islin && alg.nlsolve isa NLNewton
    # get the operator
    J = nf.f
    W = WOperator(f.mass_matrix, dt, J)
    du1 = rate_prototype
    uf = nothing
    jac_config = nothing
    linsolve = alg.linsolve(Val{:init},nf,u)
    z₊ = z
  elseif alg.nlsolve isa NLNewton
    if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
      W = WOperator(f, dt)
      J = nothing # is J = W.J better?
    else
      J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
      W = similar(J)
    end
    du1 = zero(rate_prototype)
    # if the algorithm specializes on split problems the use `nf`
    uf = DiffEqDiffTools.UJacobianWrapper(nf,t,p)
    jac_config = build_jac_config(alg,nf,uf,du1,uprev,u,tmp,dz)
    linsolve = alg.linsolve(Val{:init},uf,u)
    z₊ = z
  elseif typeof(alg.nlsolve) <: NLFunctional
    J = nothing
    W = nothing
    du1 = rate_prototype
    uf = nothing
    jac_config = nothing
    linsolve = nothing
    z₊ = similar(z)
  end

  if κ != nothing
    κ = uToltype(nlcache.κ)
  else
    κ = uToltype(1//100)
  end
  if tol != nothing
    tol = uToltype(nlcache.tol)
  else
    tol = uToltype(min(0.03,first(reltol)^(0.5)))
  end
  _nlsolve = alg.nlsolve
end
DiffEqBase.@def oopnlcachefields begin
  nlcache = alg.nlsolve.cache
  @unpack κ,tol,max_iter,min_iter,new_W = nlcache
  z = uprev
  nf = f isa SplitFunction && alg isa SplitAlgorithms ? f.f1 : f
  if alg.nlsolve isa NLNewton
    # only use `nf` if the algorithm specializes on split eqs
    uf = DiffEqDiffTools.UDerivativeWrapper(nf,t,p)
  else
    uf = nothing
  end
  islin = (f isa ODEFunction && islinear(f.f)) || (f isa SplitFunction && islinear(f.f1.f))
  if (islin || DiffEqBase.has_jac(f)) && typeof(alg.nlsolve) <: NLNewton
    # get the operator
    J = islin ? nf.f : f.jac(uprev, p, t)
    if !isa(J, DiffEqBase.AbstractDiffEqLinearOperator)
      J = DiffEqArrayOperator(J)
    end
    W = WOperator(f.mass_matrix, dt, J)
  else
    W = typeof(u) <: Number ? u : Matrix{uEltypeNoUnits}(undef, 0, 0) # uEltype?
  end
  uToltype = real(uBottomEltypeNoUnits)
  ηold = one(uToltype)

  if κ != nothing
    κ = uToltype(nlcache.κ)
  else
    κ = uToltype(1//100)
  end
  if tol != nothing
    tol = uToltype(nlcache.tol)
  else
    tol = uToltype(min(0.03,first(reltol)^(0.5)))
  end
  z₊,dz,tmp,b,k = z,z,z,z,rate_prototype
  _nlsolve = oop_nlsolver(alg.nlsolve)
end

mutable struct ImplicitEulerCache{uType,rateType,uNoUnitsType,J,W,UF,JC,F,N} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  uprev2::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  z::uType
  dz::uType
  b::uType
  tmp::uType
  atmp::uNoUnitsType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolve::N
end

u_cache(c::ImplicitEulerCache)    = (c.uprev2,c.z,c.dz)
du_cache(c::ImplicitEulerCache)   = (c.k,c.fsalfirst)

function alg_cache(alg::ImplicitEuler,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  atmp = similar(u,uEltypeNoUnits,axes(u))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1,1,ηold,z₊,dz,tmp,b,k))
  ImplicitEulerCache(u,uprev,uprev2,du1,fsalfirst,k,z,dz,b,tmp,atmp,J,W,uf,jac_config,linsolve,nlsolve)
end

mutable struct ImplicitEulerConstantCache{F,N} <: OrdinaryDiffEqConstantCache
  uf::F
  nlsolve::N
end

function alg_cache(alg::ImplicitEuler,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1,1,ηold,z₊,dz,tmp,b,k))
  ImplicitEulerConstantCache(uf,nlsolve)
end

mutable struct ImplicitMidpointConstantCache{F,N} <: OrdinaryDiffEqConstantCache
  uf::F
  nlsolve::N
end

function alg_cache(alg::ImplicitMidpoint,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1//2,1//2,ηold,z₊,dz,tmp,b,k))
  ImplicitMidpointConstantCache(uf,nlsolve)
end

mutable struct ImplicitMidpointCache{uType,rateType,J,W,UF,JC,F,N} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  z::uType
  dz::uType
  b::uType
  tmp::uType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolve::N
end

u_cache(c::ImplicitMidpointCache)    = (c.z,c.dz)
du_cache(c::ImplicitMidpointCache)   = (c.k,c.fsalfirst)

function alg_cache(alg::ImplicitMidpoint,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1//2,1//2,ηold,z₊,dz,tmp,b,k))
  ImplicitMidpointCache(u,uprev,du1,fsalfirst,k,z,dz,b,tmp,J,W,uf,jac_config,linsolve,nlsolve)
end

mutable struct TrapezoidConstantCache{F,uType,tType,N} <: OrdinaryDiffEqConstantCache
  uf::F
  uprev3::uType
  tprev2::tType
  nlsolve::N
end

function alg_cache(alg::Trapezoid,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,
                   uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  uprev3 = u
  tprev2 = t
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1//2,1,ηold,z₊,dz,tmp,b,k))
  TrapezoidConstantCache(uf,uprev3,tprev2,nlsolve)
end

mutable struct TrapezoidCache{uType,rateType,uNoUnitsType,J,W,UF,JC,tType,F,N} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  uprev2::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  z::uType
  dz::uType
  b::uType
  tmp::uType
  atmp::uNoUnitsType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  uprev3::uType
  tprev2::tType
  nlsolve::N
end

u_cache(c::TrapezoidCache)    = (c.uprev2,c.z,c.dz)
du_cache(c::TrapezoidCache)   = (c.k,c.fsalfirst)

function alg_cache(alg::Trapezoid,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  atmp = similar(u,uEltypeNoUnits,axes(u))
  uprev3 = similar(u)
  tprev2 = t
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1//2,1,ηold,z₊,dz,tmp,b,k))
  TrapezoidCache(u,uprev,uprev2,du1,fsalfirst,k,z,dz,b,tmp,atmp,J,W,uf,jac_config,linsolve,uprev3,tprev2,nlsolve)
end

mutable struct TRBDF2ConstantCache{F,Tab,N} <: OrdinaryDiffEqConstantCache
  uf::F
  nlsolve::N
  tab::Tab
end

function alg_cache(alg::TRBDF2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,
                   uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  tab = TRBDF2Tableau(uToltype,real(tTypeNoUnits))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,tab.d,tab.γ,ηold,z₊,dz,tmp,b,k))
  TRBDF2ConstantCache(uf,nlsolve,tab)
end

mutable struct TRBDF2Cache{uType,rateType,uNoUnitsType,J,W,UF,JC,Tab,F,N} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  zprev::uType
  zᵧ::uType
  z::uType
  dz::uType
  b::uType
  tmp::uType
  atmp::uNoUnitsType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolve::N
  tab::Tab
end

u_cache(c::TRBDF2Cache)    = (c.zprev,c.zᵧ,c.z,c.dz)
du_cache(c::TRBDF2Cache)   = (c.k,c.fsalfirst)

function alg_cache(alg::TRBDF2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  atmp = similar(u,uEltypeNoUnits,axes(u)); zprev = similar(u,axes(u));
  zᵧ = similar(u,axes(u))
  tab = TRBDF2Tableau(uToltype,real(tTypeNoUnits))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,tab.d,tab.γ,ηold,z₊,dz,tmp,b,k))
  TRBDF2Cache(u,uprev,du1,fsalfirst,k,zprev,zᵧ,z,dz,b,tmp,atmp,J,
              W,uf,jac_config,linsolve,nlsolve,tab)
end

mutable struct SDIRK2ConstantCache{F,N} <: OrdinaryDiffEqConstantCache
  uf::F
  nlsolve::N
end

function alg_cache(alg::SDIRK2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,
                   uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1,1,ηold,z₊,dz,tmp,b,k))

  SDIRK2ConstantCache(uf,nlsolve)
end

mutable struct SDIRK2Cache{uType,rateType,uNoUnitsType,J,W,UF,JC,F,N} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  z₁::uType
  z₂::uType
  dz::uType
  b::uType
  tmp::uType
  atmp::uNoUnitsType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolve::N
end

u_cache(c::SDIRK2Cache)    = (c.z₁,c.z₂,c.dz)
du_cache(c::SDIRK2Cache)   = (c.k,c.fsalfirst)

function alg_cache(alg::SDIRK2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  z₁ = similar(u,axes(u)); z₂ = z;
  atmp = similar(u,uEltypeNoUnits,axes(u))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1,1,ηold,z₊,dz,tmp,b,k))

  SDIRK2Cache(u,uprev,du1,fsalfirst,k,z₁,z₂,dz,b,tmp,atmp,J,
              W,uf,jac_config,linsolve,nlsolve)
end

mutable struct SSPSDIRK2ConstantCache{F,N} <: OrdinaryDiffEqConstantCache
  uf::F
  nlsolve::N
end

function alg_cache(alg::SSPSDIRK2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,
                   uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1//4,1//1,ηold,z₊,dz,tmp,b,k))

  SSPSDIRK2ConstantCache(uf,nlsolve)
end

mutable struct SSPSDIRK2Cache{uType,rateType,J,W,UF,JC,F,N} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  z₁::uType
  z₂::uType
  dz::uType
  b::uType
  tmp::uType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolve::N
end

u_cache(c::SSPSDIRK2Cache)    = (c.z₁,c.z₂,c.dz)
du_cache(c::SSPSDIRK2Cache)   = (c.k,c.fsalfirst)

function alg_cache(alg::SSPSDIRK2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  z₁ = similar(u,axes(u)); z₂ = z;
  atmp = similar(u,uEltypeNoUnits,axes(u))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,1//4,1//1,ηold,z₊,dz,tmp,b,k))

  SSPSDIRK2Cache(u,uprev,du1,fsalfirst,k,z₁,z₂,dz,b,tmp,J,
                W,uf,jac_config,linsolve,nlsolve)
end

mutable struct Kvaerno3ConstantCache{UF,Tab,N} <: OrdinaryDiffEqConstantCache
  uf::UF
  nlsolve::N
  tab::Tab
end

function alg_cache(alg::Kvaerno3,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,
                   uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  tab = Kvaerno3Tableau(uToltype,real(tTypeNoUnits))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,tab.γ,2tab.γ,ηold,z₊,dz,tmp,b,k))

  Kvaerno3ConstantCache(uf,nlsolve,tab)
end

mutable struct Kvaerno3Cache{uType,rateType,uNoUnitsType,J,W,UF,JC,Tab,F,N} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  z₁::uType
  z₂::uType
  z₃::uType
  z₄::uType
  dz::uType
  b::uType
  tmp::uType
  atmp::uNoUnitsType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolve::N
  tab::Tab
end

u_cache(c::Kvaerno3Cache)    = (c.z₁,c.z₂,c.z₃,c.z₄,c.dz)
du_cache(c::Kvaerno3Cache)   = (c.k,c.fsalfirst)

function alg_cache(alg::Kvaerno3,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  z₁ = similar(u,axes(u)); z₂ = similar(u,axes(u));
  z₃ = similar(u,axes(u)); z₄ = z;
  atmp = similar(u,uEltypeNoUnits,axes(u))
  tab = Kvaerno3Tableau(real(uBottomEltypeNoUnits),real(tTypeNoUnits))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,tab.γ,2tab.γ,ηold,z₊,dz,tmp,b,k))

  Kvaerno3Cache(u,uprev,du1,fsalfirst,k,z₁,z₂,z₃,z₄,dz,b,tmp,atmp,J,
                W,uf,jac_config,linsolve,nlsolve,tab)
end

mutable struct Cash4ConstantCache{F,N,Tab} <: OrdinaryDiffEqConstantCache
  uf::F
  nlsolve::N
  tab::Tab
end

function alg_cache(alg::Cash4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,
                   uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  tab = Cash4Tableau(uToltype,real(tTypeNoUnits))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,tab.γ,tab.γ,ηold,z₊,dz,tmp,b,k))

  Cash4ConstantCache(uf,nlsolve,tab)
end

mutable struct Cash4Cache{uType,rateType,uNoUnitsType,J,W,UF,JC,N,Tab,F} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  z₁::uType
  z₂::uType
  z₃::uType
  z₄::uType
  z₅::uType
  dz::uType
  b::uType
  tmp::uType
  atmp::uNoUnitsType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolve::N
  tab::Tab
end

u_cache(c::Cash4Cache)    = (c.z₁,c.z₂,c.z₃,c.z₄,c.z₅,c.dz)
du_cache(c::Cash4Cache)   = (c.k,c.fsalfirst)

function alg_cache(alg::Cash4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  z₁ = similar(u,axes(u)); z₂ = similar(u,axes(u));
  z₃ = similar(u,axes(u)); z₄ = similar(u,axes(u))
  z₅ = z
  atmp = similar(u,uEltypeNoUnits,axes(u))
  tab = Cash4Tableau(uToltype,real(tTypeNoUnits))
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,tab.γ,tab.γ,ηold,z₊,dz,tmp,b,k))

  Cash4Cache(u,uprev,du1,fsalfirst,k,z₁,z₂,z₃,z₄,z₅,dz,b,tmp,atmp,J,
              W,uf,jac_config,linsolve,nlsolve,tab)
end

mutable struct Hairer4ConstantCache{F,N,Tab} <: OrdinaryDiffEqConstantCache
  uf::F
  nlsolve::N
  tab::Tab
end

function alg_cache(alg::Union{Hairer4,Hairer42},u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,
                   uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  @oopnlcachefields
  if typeof(alg) <: Hairer4
    tab = Hairer4Tableau(uToltype,real(tTypeNoUnits))
  else
    tab = Hairer42Tableau(uToltype,real(tTypeNoUnits))
  end
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,tab.γ,tab.γ,ηold,z₊,dz,tmp,b,k))

  Hairer4ConstantCache(uf,nlsolve,tab)
end

mutable struct Hairer4Cache{uType,rateType,uNoUnitsType,J,W,UF,JC,Tab,F,N} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  du1::rateType
  fsalfirst::rateType
  k::rateType
  z₁::uType
  z₂::uType
  z₃::uType
  z₄::uType
  z₅::uType
  dz::uType
  b::uType
  tmp::uType
  atmp::uNoUnitsType
  J::J
  W::W
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolve::N
  tab::Tab
end

u_cache(c::Hairer4Cache)    = (c.z₁,c.z₂,c.z₃,c.z₄,c.z₅,c.dz)
du_cache(c::Hairer4Cache)   = (c.k,c.fsalfirst)

function alg_cache(alg::Union{Hairer4,Hairer42},u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
                   tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  @iipnlcachefields
  z₁ = similar(u,axes(u)); z₂ = similar(u,axes(u));
  z₃ = similar(u,axes(u)); z₄ = similar(u,axes(u))
  z₅ = z
  dz = similar(u,axes(u))
  atmp = similar(u,uEltypeNoUnits,axes(u))

  if typeof(alg) <: Hairer4
    tab = Hairer4Tableau(uToltype,real(tTypeNoUnits))
  else # Hairer42
    tab = Hairer42Tableau(uToltype,real(tTypeNoUnits))
  end
  nlsolve = typeof(_nlsolve)(NLSolverCache(κ,tol,min_iter,max_iter,10000,new_W,z,W,tab.γ,tab.γ,ηold,z₊,dz,tmp,b,k))

  Hairer4Cache(u,uprev,du1,fsalfirst,k,z₁,z₂,z₃,z₄,z₅,dz,b,tmp,atmp,J,
               W,uf,jac_config,linsolve,nlsolve,tab)
end


text:
"""
Set of Julia functions to assist with interfacing between the Julia MPC and the
Python neural network model (mats).
"""

using PyCall
push!(PyVector(pyimport("sys")."path"), "../scripts")
python_utils = pyimport("python_utils")

struct Node
    x               ::Vector
    y               ::Vector
    θ               ::Vector
    v               ::Vector
    first_timestep  ::Int64
    last_timestep   ::Int64
    type            ::String
    id              ::Int64
end

struct Scene{T}
    # information about agents in scene
    robot        ::Node
    nodes        ::Vector{Node}

    # possibly useful metadata
    dt           ::T
    timesteps    ::Int64
    node_ids     ::Vector{String}
    x_offset     ::Float64
    y_offset     ::Float64
end

mutable struct PredictionSettings
    mats
    hyperparams
    env
    num_modes
end

mutable struct Obstacle
    ps        ::Vector
    active    ::Bool
end

function Scene(env, scene_num::Int64)
    # information about agents in scene
    robot = nothing
    non_robot_nodes = Vector{Node}()
    x_offset = env.scenes[scene_num].x_min
    y_offset = env.scenes[scene_num].y_min
    for node in env.scenes[scene_num].nodes
        if node.is_robot
            x = node.data.data[:, 1] .+ x_offset
            y = node.data.data[:, 2] .+ y_offset
            θ = node.data.data[:, 9]
            v = node.data.data[:, 11]
            first_timestep = node.first_timestep
            last_timestep = node.last_timestep
            type = node.type.name
            id = node.id
            robot = Node(x, y, θ, v, first_timestep, last_timestep, type, id)
        else
            x = node.data.data[:, 1] .+ x_offset
            y = node.data.data[:, 2] .+ y_offset
            θ = []
            v = []
            first_timestep = node.first_timestep
            last_timestep = node.last_timestep
            type = node.type.name
            id = node.id
            push!(non_robot_nodes, Node(x, y, θ, v, first_timestep, last_timestep, type, id))
        end
    end

    # possibly useful metadata
    dt = env.scenes[scene_num].dt
    timesteps = env.scenes[scene_num].timesteps
    node_ids = [string(node.id) for node in non_robot_nodes]

    Scene(robot, non_robot_nodes, dt, timesteps, node_ids, x_offset, y_offset)
end

function load_model(model_path, env; ts=100)
    mats, hyperparams =
        pycall(python_utils.load_model, NTuple{2, PyObject}, model_path, env, ts);

    mats.set_environment(env)
    mats.set_annealing_params()
    python_utils.calculate_scene_graphs(env, hyperparams)
    mats, hyperparams
end

function predicted_dynamics(pred_settings, scene_num, timestep)
    prediction_info, dynamics_dict = python_utils.predict(pred_settings.mats,
                                                          pred_settings.hyperparams,
                                                          pred_settings.env.scenes[scene_num],
                                                          timestep-1,
                                                          pred_settings.num_modes)

    pred_horizon = pred_settings.hyperparams.get("prediction_horizon")
    num_modes = pred_settings.num_modes

    Aps = [[dynamics_dict["A"][ts, mode, :, :] for ts in 1:pred_horizon] for mode in 1:num_modes]
    Bps = [[dynamics_dict["B"][ts, mode, :, :] for ts in 1:pred_horizon] for mode in 1:num_modes]
    gps = [[dynamics_dict["affine_terms"][ts, mode, :, :] for ts in 1:pred_horizon] for mode in 1:num_modes]

    state_dim = 4 #TODO: state dimension is hardcoded and assumed uniform
    q0 = zeros((length(prediction_info) + 1) * state_dim) # the plus one accounts for robot state
    ordered_node_ids = Array{String, 1}(undef, length(prediction_info))
    for (k, v) in prediction_info
        node_idx = v["node_idx"]
        state_vector_idx = node_idx * state_dim + 1
        q0[state_vector_idx:state_vector_idx+state_dim-1] = v["current_state"]
        ordered_node_ids[v["node_idx"]] = split(k, "/")[2]
    end

    #TODO: hard coded indices
    q_robot = [env.scenes[scene_num].robot.data.data[timestep, 1], # x
               env.scenes[scene_num].robot.data.data[timestep, 2], # y
               env.scenes[scene_num].robot.data.data[timestep, 9], # heading
               env.scenes[scene_num].robot.data.data[timestep, 11]] # v
    q0[1:4] = q_robot

    Aps, Bps, gps, q0, ordered_node_ids
end

#TODO: hard coded indices
#TODO: check direction of accel vector so there's nothing weird happening with accel norm
function get_recorded_robot_controls(pred_settings, scene_num, timestep)
    ω = pred_settings.env.scenes[scene_num].robot.data.data[timestep:end, 10]
    a = pred_settings.env.scenes[scene_num].robot.data.data[timestep:end, 12]
    [ω a]
end

#TODO: hardcoded indices
function predict_future_states(q0, us, As, Bs, gs, mode)
    pred_horizon = pred_settings.hyperparams.get("prediction_horizon")
    state_dim = size(As[1][1])[1]
    q_pred = zeros(state_dim, pred_horizon+1)
    q_pred[:, 1] = q0

    for k in 1:pred_horizon
        q_pred[:, k+1] =  gs[mode][k] + As[mode][k] * q_pred[:, k] + Bs[mode][k] * us[k, :]
    end

    q_pred
end

function init_node_obstacles!(scene, vals)
    for node_id in scene.node_ids
        vals.p_obs[node_id] = Obstacle([[[0., 0.] for i = 1:vals.obs_horizon]
            for j = 1:vals.n_modes], 0)
    end
end

#TODO: hardcoded indices
function update_obstacles_from_predictions!(q_pred, node_ids, vals, scene; iter=nothing)
    #####
    obstacle_heading = Dict()
    for (obs_id, obs) in vals.p_obs
        obstacle_heading[obs_id] = []
    end
    #####

    for mode = 1:vals.n_modes
        # node_idx = 2 # start at 2 to avoid robot state
        for (node_id, obs) in vals.p_obs
            if node_id in node_ids
                node_idx = findall(node_ids .== node_id)[1] + 1 #TODO: fix this convoluted mess (+1 from julia vs. python indexing)
                if length(findall(node_ids .== node_id)) > 1
                    println("Check update_obstacles_from_predictions!().")
                end
                row_idx = (node_idx - 1) * 4 + 1
                vals.p_obs[node_id].ps[mode] = [[q_pred[mode][row_idx, k] + scene.x_offset, q_pred[mode][row_idx+1, k] + scene.y_offset] for k = 1:vals.obs_horizon]
                vals.p_obs[node_id].active = 1

                #####
                push!(obstacle_heading[node_id], q_pred[mode][row_idx+2, :])
                #####

                node_idx += 1
            else
                vals.p_obs[node_id].active = 0
            end
        end
    end

    ####
    filename = ""
    if iter == nothing
        filename = "data/obstacle_heading_iteration_1"
    else
        filename = "data/obstacle_heading_iteration_"*string(iter)
    end
    np.save(filename, obstacle_heading)
    ####

end

#TODO: hardcoded indices
function extract_positions_from_state_pred!(q_pred, node_ids, vals)
    num_nodes = length(node_ids)
    horizon = size(q_pred)[2]

    for node_idx = 1:num_nodes
        row_idx = (node_idx - 1) * 4 + 1
        ps_node = []
        for k = 1:horizon
            push!(ps_node, [q_pred[row_idx, k], q_pred[row_idx+1, k]])
        end
        vals.p_obs[node_id] = ps_node
    end
end

"""
Sets positions for static obstacles. The variable p_obs is a list of the
static obstacle position vectors for each "mode"
(i.e., [[x_1, y_1], [x_2, y_2], ... , [x_n_modes, y_n_modes]]).
"""
function add_static_obstacle!(obstacles, vals)
    @assert length(obstacles) <= vals.n_modes
    @assert length(vals.p_obs) < vals.n_obs

    idx = length(vals.p_obs) + 1
    n_obs_modes = length(obstacles)
    vals.p_obs[idx] = [[obstacles[i] for j = 1:vals.obs_horizon] for i = 1:n_obs_modes]
    if n_obs_modes < vals.n_modes
        append!(vals.p_obs[idx], [[[0, 0] for i = 1:vals.obs_horizon]
            for j = 1:(vals.n_modes - n_obs_modes)])
    end
end

max_stars_repo_name: StanfordASL/MATS

max_issues_repo_name: StanfordASL/MATS

max_forks_repo_name: StanfordASL/MATS


text:
function quickgroupby{T, S<:Number}(by::AbstractVector{T}, val::AbstractVector{S})
  j = length(by)
  if j == 1
    return by
  elseif j == 2
    return by
  elseif j  <= 0
    warn("what?")
  end

  pivot = by[1]

  lowergroupi = 1
  i = 2

  while true
    while i <= j
      byi = by[i]
      if byi < pivot
        i += 1
      elseif byi > pivot
        break
      else
         lowergroupi += 1
         by[lowergroupi], by[i] = by[i], by[lowergroupi]
         val[lowergroupi], val[i] = val[i], val[lowergroupi]
         i += 1
      end
    end

    while i <= j
      byj = by[j]
      if pivot < byj
        j -= 1
      elseif pivot > byj
        break
      else
        lowergroupi += 1
        by[lowergroupi], by[j], by[i] = by[j], by[i], by[lowergroupi-1]
        val[lowergroupi], val[j], val[i] = val[j], val[i], val[lowergroupi-1]
        j -= 1
        i += 1
      end
    end

    if i > j
      # print(by[i])
      # print(by[j])
      break
    end
    by[i], by[j] = by[j], by[i]
    val[i], val[j] = val[j], val[i]
    i += 1
    j -= 1
  end

  # println(by)
  # println(lowergroupi+1)
  # println(i)
  # println(j)

  if lowergroupi == length(by)
    return by
  elseif i > length(by) || ((lowergroupi + 1) > (i-1))
    lby = quickgroupby(by[(lowergroupi+1):end],val[(lowergroupi+1):end])
    return vcat(by[1:lowergroupi], lby)
  else
    lby = quickgroupby(by[(lowergroupi+1):(i-1)],val[(lowergroupi+1):(i-1)])
    rby = quickgroupby(by[i:end],val[i:end])
    return vcat(by[1:lowergroupi], lby, rby)
  end
end
#
# x = rand(1:3, 10)
# by = copy(x)
# val = copy(x)
# quickgroupby(by,val)
# const N = 10_000_000
# const K = 100
# addprocs()
# id6 = rand(1:(N/K), N)
# sid6 = SharedArray(id6)
# v1 = similar(id6)
# @time quickgroupby(id6, v1)
#
# using FastGroupBy
# @time sumby(id6,v1)


text:
module RocketCombineLatestObservableTest

using Test
using Rocket

include("../test_helpers.jl")

@testset "CombineLatestObservable" begin

    println("Testing: combineLatest")

    @testset begin
        @test_throws ErrorException combineLatest()
    end

    run_testset([
        (
            source = combineLatest(of(1), from(1:5)),
            values = @ts([ (1, 1), (1, 2), (1, 3), (1, 4), (1, 5), c ]),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(of(1) |> async(0), from(1:5)),
            values = @ts([ (1, 5) ] ~ c),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(of(1) |> async(0), from(1:5) |> async(0)),
            values = @ts([ (1, 1) ] ~ [ (1, 2) ] ~ [ (1, 3) ] ~ [ (1, 4) ] ~ [ (1, 5) ] ~ c),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(from(1:5), of(2.0)),
            values = @ts([ (5, 2), c ]),
            source_type = Tuple{Int, Float64}
        ),
        (
            source = combineLatest(from(1:5) |> async(0), of(2.0)),
            values = @ts([ (1, 2.0) ] ~ [ (2, 2.0) ] ~ [ (3, 2.0) ] ~ [ (4, 2.0) ] ~ [ (5, 2.0) ] ~ c),
            source_type = Tuple{Int, Float64}
        ),
        (
            source = combineLatest(completed(Int), of(1)),
            values = @ts(c),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(completed(Int), completed(Int)),
            values = @ts(c),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(faulted(Float64, "err"), completed(Int)),
            values = @ts(e("err")),
            source_type = Tuple{Float64, Int}
        ),
        (
            source = combineLatest(completed(Int), faulted(Float64, "err")),
            values = @ts(c),
            source_type = Tuple{Int, Float64}
        ),
        (
            source = combineLatest(faulted(Float64, "err1"), faulted(Float64, "err2")),
            values = @ts(e("err1")),
            source_type = Tuple{Float64, Float64}
        ),
        (
            source = combineLatest(of(1), from(1:5), strategy = PushNew()),
            values = @ts([ (1, 1), (1, 2), (1, 3), (1, 4), (1, 5), c ]),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(of(1) |> async(0), from(1:5), strategy = PushNew()),
            values = @ts([ (1, 5) ] ~ c),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(from(1:5), of(2.0), strategy = PushNew()),
            values = @ts([ (5, 2), c ]),
            source_type = Tuple{Int, Float64}
        ),
        (
            source = combineLatest(from(1:5) |> async(0), of(2.0), strategy = PushNew()),
            values = @ts([ (1, 2.0) ] ~ [ (2, 2.0) ] ~ [ (3, 2.0) ] ~ [ (4, 2.0) ] ~ [ (5, 2.0) ] ~ c),
            source_type = Tuple{Int, Float64}
        ),
        (
            source = combineLatest(completed(Int), of(1), strategy = PushNew()),
            values = @ts(c),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(completed(Int), completed(Int), strategy = PushNew()),
            values = @ts(c),
            source_type = Tuple{Int, Int}
        ),
        (
            source = combineLatest(faulted(Float64, "err"), completed(Int), strategy = PushNew()),
            values = @ts(e("err")),
            source_type = Tuple{Float64, Int}
        ),
        (
            source = combineLatest(completed(Int), faulted(Float64, "err"), strategy = PushNew()),
            values = @ts(c),
            source_type = Tuple{Int, Float64}
        ),
        (
            source = combineLatest(faulted(Float64, "err1"), faulted(Float64, "err2"), strategy = PushNew()),
            values = @ts(e("err1")),
            source_type = Tuple{Float64, Float64}
        ),
    ])

    somenumbers    = [ [ 1 for i in 1:50 ]..., [ 1.0 for i in 1:50 ]... ]
    expectedoutput = tuple(somenumbers...)

    combined100 = combineLatest(map(n -> of(n), somenumbers)...)

    run_testset([
        (
            source = combined100,
            values = @ts([ tuple([ [ 1 for i in 1:50 ]..., [ 1.0 for i in 1:50 ]... ]...), c ]),
            source_type = Tuple{ map(n -> typeof(n), somenumbers)... }
        )
    ])

    @testset begin
        s1 = Subject(Int)
        s2 = Subject(Float64)
        s3 = Subject(String)
        s4 = of(5)
        s5 = from(1:10)

        latest = combineLatest(s1, s2, s3, s4, s5)
        values = Vector{Any}()
        subscription = subscribe!(latest, lambda(
            on_next     = (d) -> push!(values, d),
            on_error    = (e) -> push!(values, e),
            on_complete = ()  -> push!(values, "completed")
        ))

        @test values == [ ]

        next!(s1, 1)

        @test values == [ ]

        next!(s2, 2.0)

        @test values == [ ]

        next!(s3, "Hello")

        @test values == [ (1, 2.0, "Hello", 5, 10) ]

        next!(s1, 2)

        @test values == [ (1, 2.0, "Hello", 5, 10), (2, 2.0, "Hello", 5, 10) ]

        next!(s3, "Hello, world!")

        @test values == [ (1, 2.0, "Hello", 5, 10), (2, 2.0, "Hello", 5, 10), (2, 2.0, "Hello, world!", 5, 10)  ]

        complete!(s1);

        next!(s2, 3.0)

        @test values == [ (1, 2.0, "Hello", 5, 10), (2, 2.0, "Hello", 5, 10), (2, 2.0, "Hello, world!", 5, 10), (2, 3.0, "Hello, world!", 5, 10)  ]

        complete!(s2)

        @test values == [ (1, 2.0, "Hello", 5, 10), (2, 2.0, "Hello", 5, 10), (2, 2.0, "Hello, world!", 5, 10), (2, 3.0, "Hello, world!", 5, 10)  ]

        unsubscribe!(subscription)
    end

    @testset begin
        s1 = Subject(Int)
        s2 = Subject(Float64)
        s3 = Subject(String)
        s4 = of(5)
        s5 = from(1:10)

        latest = combineLatest(s1, s2, s3, s4, s5, strategy = PushNew())
        values = Vector{Any}()
        subscription = subscribe!(latest, lambda(
            on_next     = (d) -> push!(values, d),
            on_error    = (e) -> push!(values, e),
            on_complete = ()  -> push!(values, "completed")
        ))

        @test values == [ ]

        next!(s1, 1)

        @test values == [ ]

        next!(s2, 2.0)

        @test values == [ ]

        next!(s3, "Hello")

        @test values == [ (1, 2.0, "Hello", 5, 10) ]

        next!(s1, 2)

        @test values == [ (1, 2.0, "Hello", 5, 10) ]

        next!(s3, "Hello, world!")

        @test values == [ (1, 2.0, "Hello", 5, 10)  ]

        complete!(s1);

        next!(s2, 3.0)

        @test values == [ (1, 2.0, "Hello", 5, 10), (2, 3.0, "Hello, world!", 5, 10)  ]

        complete!(s2)

        @test values == [ (1, 2.0, "Hello", 5, 10), (2, 3.0, "Hello, world!", 5, 10)  ]

        next!(s3, "Something")

        @test values == [ (1, 2.0, "Hello", 5, 10), (2, 3.0, "Hello, world!", 5, 10), (2, 3.0, "Something", 5, 10)  ]

        complete!(s3)

        @test values == [ (1, 2.0, "Hello", 5, 10), (2, 3.0, "Hello, world!", 5, 10), (2, 3.0, "Something", 5, 10), "completed"  ]

        unsubscribe!(subscription)
    end

    @testset begin
        s1 = Subject(Int)
        s2 = Subject(Float64)
        s3 = Subject(String)

        latest = combineLatest(s1, s2, s3, strategy = PushNewBut{1}())
        values = Vector{Any}()
        subscription = subscribe!(latest, lambda(
            on_next     = (d) -> push!(values, d),
            on_error    = (e) -> push!(values, e),
            on_complete = ()  -> push!(values, "completed")
        ))

        @test values == [ ]

        next!(s1, 1)

        @test values == [ ]

        next!(s2, 2.0)

        @test values == [ ]

        next!(s3, "Hello")

        @test values == [ (1, 2.0, "Hello") ]

        next!(s2, 3.0)

        @test values == [ (1, 2.0, "Hello") ]

        next!(s3, "Hello, world!")

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!") ]

        next!(s1, 2)

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!") ]

        next!(s1, 3)

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!") ]

        next!(s1, 4)

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!") ]

        next!(s2, 4.0)

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!") ]

        next!(s3, "")

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!"), (4, 4.0, "") ]

        complete!(s1);

        next!(s2, 5.0)

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!"), (4, 4.0, "") ]

        next!(s3, "foo")

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!"), (4, 4.0, ""), (4, 5.0, "foo") ]

        complete!(s2)

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!"), (4, 4.0, ""), (4, 5.0, "foo") ]

        next!(s3, "bar")

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!"), (4, 4.0, ""), (4, 5.0, "foo"), (4, 5.0, "bar") ]

        complete!(s3)

        @test values == [ (1, 2.0, "Hello"), (1, 3.0, "Hello, world!"), (4, 4.0, ""), (4, 5.0, "foo"), (4, 5.0, "bar"), "completed" ]

        unsubscribe!(subscription)
    end

    @testset begin
        s1 = Subject(Int)
        s2 = Subject(Float64)
        s3 = Subject(String)

        latest = combineLatest(s1, s2, s3, strategy = PushStrategy(BitArray([ false, true, false ])))
        values = Vector{Any}()
        subscription = subscribe!(latest, lambda(
            on_next     = (d) -> push!(values, d),
            on_error    = (e) -> push!(values, e),
            on_complete = ()  -> push!(values, "completed")
        ))

        @test values == [ ]

        next!(s1, 1)

        @test values == [ ]

        next!(s2, 2.0)

        @test values == [ ]

        next!(s3, "Hello")

        @test values == [ (1, 2.0, "Hello") ]

        next!(s2, 3.0)

        @test values == [ (1, 2.0, "Hello") ]

        next!(s3, "Hello, world!")

        @test values == [ (1, 2.0, "Hello") ]

        next!(s1, 2)

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!") ]

        next!(s1, 3)

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!") ]

        next!(s1, 4)

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!") ]

        next!(s2, 4.0)

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!") ]

        next!(s3, "")

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!"), (4, 4.0, "") ]

        complete!(s1);

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!"), (4, 4.0, "") ]

        next!(s2, 5.0)

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!"), (4, 4.0, "") ]

        next!(s3, "1")

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!"), (4, 4.0, ""), (4, 5.0, "1") ]

        complete!(s3)

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!"), (4, 4.0, ""), (4, 5.0, "1") ]

        next!(s2, 6.0)

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!"), (4, 4.0, ""), (4, 5.0, "1"), (4, 6.0, "1") ]

        complete!(s2)

        @test values == [ (1, 2.0, "Hello"), (2, 3.0, "Hello, world!"), (4, 4.0, ""), (4, 5.0, "1"), (4, 6.0, "1"), "completed" ]

        unsubscribe!(subscription)
    end

end

end


text:
let
	@test BayesNets.ndgrid([2,1]) == [2,1]
	@test elementwise_isapprox(BayesNets.ndgrid([1.0, 2.0, 3.0]), [1.0, 2.0, 3.0])
	@test BayesNets.ndgrid([1,2], [3,4]) == (
											[1 1
											 2 2],
											[3 4
											 3 4])

	# TODO: make this test more rigorous
	@test isa(BayesNets.ndgrid([1], [2], [3]), Tuple{Array{Int64,3},Array{Int64,3},Array{Int64,3}})

	# TODO: make this test more rigorous
	@test BayesNets.ndgrid_fill!(Array{Int}(undef, 2), [1,2], 1, 2) == [1,2]

	@test BayesNets.meshgrid([1,2]) == (
										[1 2
										 1 2],
										[1 1
										 2 2])
	@test BayesNets.meshgrid([1,2], [3,4]) == (
												[1 2
												 1 2],
												[3 3
												 4 4])

	# TODO: make this test more rigorous
	@test isa(BayesNets.meshgrid([1], [2], [3]), Tuple{Array{Int64,3},Array{Int64,3},Array{Int64,3}})
end


text:
function  mock_solver_unit_tests()

    mock_tests()

end

function mock_tests()

    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    rvrp_sol = RVRP.solve(data, RVRP.MockSolver())
    @test typeof(rvrp_sol) == RVRP.RvrpSolution
    @test rvrp_sol.cost == 644.0

    for act_idx in 1:length(rvrp_sol.routes[4].sequence)
        @test rvrp_sol.routes[4].sequence[act_idx].id == string("act_", act_idx)
    end

    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    data.vehicle_sets[1].nb_of_vehicles_range = RVRP.FlexibleRange(soft_range = RVRP.Range(0, 3))
    rvrp_sol = RVRP.solve(data, RVRP.MockSolver())
    @test rvrp_sol.unassigned_request_ids == ["req_4", "req_5", "req_6", "req_7", "req_8", "req_9", "req_10", "req_11", "req_12", "req_13", "req_14", "req_15"]

    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    data.vehicle_sets = RVRP.HomogeneousVehicleSet[]
    rvrp_sol = RVRP.solve(data, RVRP.MockSolver())
    @test rvrp_sol.routes == RVRP.Route[]
    @test rvrp_sol.unassigned_request_ids == ["req_1", "req_2", "req_3", "req_4", "req_5", "req_6", "req_7", "req_8", "req_9", "req_10", "req_11", "req_12", "req_13", "req_14", "req_15"]

    # Empty instance
    data = RVRP.RvrpInstance()
    rvrp_sol = RVRP.solve(data, RVRP.MockSolver())
    @test rvrp_sol.cost == 0.0
    @test rvrp_sol.routes == RVRP.Route[]
    @test rvrp_sol.unassigned_request_ids == String[]

end


text:
using Gadfly

scatter = readdlm("scatter.dat")
scatter = [scatter; readdlm("scatter1.dat")]

w = 0.02

p = plot(
    layer(
        x = scatter[:,1].*(w/π),
        y = scatter[:,2],
        Geom.line,
        Geom.point,
        color = ["Transmission Probability"]
    ),
    layer(
        x = scatter[:,1].*(w/π),
        y = scatter[:,3],
        Geom.line,
        Geom.point,
        color = ["Reflection Probability"]
    ),
    Guide.xlabel("Multiple of Resonant Wave Number, k(w/π)"),
    Guide.ylabel("Probability"),
    Guide.title("Transmission probability vs. momentum")
)

draw(SVG("scatter.svg", 20cm, 12cm), p)


text:
using Pkg
Pkg.activate(".")

using ArgParse
using Profile
using dlgo
using dlgo.Agent
using DelimitedFiles
using DataStructures

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--board_size", "-b"
            help = "The game will have dimensions BOARD_SIZExBOARD_SIZE"
            arg_type = Int
            default = 9
        "--num_games", "-n"
            help = "Number of games to play"
            arg_type = Int
            default = 1
        "--komi", "-k"
            help = "White player handicap for scoring"
            arg_type = Float64
            default = 6.5
        "--verbose", "-v"
            help = "print results of game to command line"
            arg_type = Bool
            default = true
        "--profile", "-p"
            help = "profile shows time spent per call"
            arg_type = Bool
            default = false
        "--num_rounds", "-m"
            help = "number of mcts rollouts per decision"
            arg_type = Int
            default = 10
        "--print_every", "-e"
            help = "only if verbose, print every?"
            arg_type = Int
            default = 1
    end
    return parse_args(s, as_symbols=true)
end

num_rounds = 100
komi = 6.5
board_size = 9

function main(board_size::Int, num_games::Int, komi::Float64, num_rounds::Int, verbose::Bool, print_every::Int)
    verbose && println("Playing ", num_games, " game(s) with board size ", board_size, "x", board_size)

    tree = MCTSTree()
    agent = MCTSAgent(num_rounds, komi)

    for it in 1:num_games
        root = MCTSNode(new_game(board_size), tree)

        nmoves = 0
        node = root
        while !is_over(node.state)
            if is_leaf(node)
                expand_leaf!(agent, node, tree)
            end
            node = select_move(agent, node)
            nmoves += 1
        end
        result = compute_game_result(node.state, komi)
        verbose && println("Finished game ", it, " in " , nmoves, " moves with result ", result)
        (it % print_every == 0) && verbose && print_board(node.state.board)

        # traverse save the tree and save data
        update_tree!(tree, root)

        # generate and save tree output data
        featdata = []
        outputdata = []
        to_visit = Queue{MCTSNode}()  # bfs traversal to update storage
        enqueue!(to_visit, root)
        while !isempty(to_visit)
            vnode = dequeue!(to_visit)
            if !is_leaf(vnode)
                feats = features(vnode.state)
                push!(featdata, feats)
                probs = search_probabilities(vnode, 250, .03)
                N = num_visits(vnode)
                Nb = num_wins(vnode)[black]
                Nw = num_wins(vnode)[white]
                wb = N / Nb
                ww = N / Nw
                push!(outputdata, [vec(probs); [N, Nb, Nw, wb, ww]])
                for child in vnode.children
                    enqueue!(to_visit, child)
                end
            end
        end
        featmat = vcat([f' for f in featdata]...)
        outputmat = vcat([o' for o in outputdata]...)
        open("test/testdata/features.csv", "w") do io
            writedlm(io, featmat, ',')
        end
        open("test/testdata/target.csv", "w") do io
            writedlm(io, outputmat, ',')
        end
    end

    # flush data
end

args = parse_commandline()
@time begin
    if args[:profile]
        @profile main(args[:board_size], args[:num_games], args[:komi], args[:num_rounds], args[:verbose], args[:print_every])
        Profile.print(format=:flat)
    else
        main(args[:board_size], args[:num_games], args[:komi], args[:num_rounds], args[:verbose], args[:print_every])
    end
end


text:
using ChangePrecision
using Test, LinearAlgebra, Statistics

≡(x::T,y::T) where {T<:Number} = isequal(x, y)
≡(x::T,y::T) where {T} = x == y
≡(x, y) = false

@testset "basic tests" begin
    for T in (Float16, Float32, BigFloat)
        @test T(1)/T(3) ≡ @changeprecision(T, 1/3)
        @test T(1) ≡ @changeprecision(T, float(1))
        @test parse(T, "1.234") ≡ @changeprecision(T, 1.234)
        if T != BigFloat
            @test @changeprecision(T, rand()) isa T
            @test @changeprecision(T, rand(Float64)) isa Float64
        end
        @test @changeprecision(T, ones(2,3)) ≡ ones(T,2,3)
        @test @changeprecision(T, ones()) ≡ ones(T)
        @test @changeprecision(T, ones(Float64, 2,3)) ≡ ones(Float64,2,3)
        @test @changeprecision(T, sin(1)) ≡ sin(T(1))
        @test @changeprecision(T, sqrt(2)) ≡ sqrt(T(2))
        @test @changeprecision(T, Inf) ≡ T(Inf)
        @test @changeprecision(T, NaN) ≡ T(NaN)
    end
end

@testset "rationals" begin
    @test @changeprecision(Float32, sqrt(2//3)) === sqrt(Float32(2//3))
    @test @changeprecision(Float32, inv(2//3)) === 3//2
    @test @changeprecision(Float32, (2//3) / 2) === 1//3
    @test @changeprecision(Float32, 2 / (4//3)) === 3//2
    @test @changeprecision(Float32, (2//3) / (4//3)) === 1//2
    @test @changeprecision(Float32, (2//3) \ (4//3)) === 2//1
end

@testset "complex" begin
    @test @changeprecision(Float32, abs(1+1im)) === sqrt(Float32(2))
    @test @changeprecision(Float32, angle(1+0im)) === Float32(0)
end

const euler = VERSION < v"0.7.0-DEV.1592" ? e : ℯ # changed in JuliaLang/julia#23427

@testset "irrational" begin
    @test @changeprecision(Float32, sqrt(pi)) === sqrt(Float32(pi))
    @test @changeprecision(Float32, pi/3) === Float32(pi)/3
    @test @changeprecision(Float32, 2pi) === Float32(pi)*2
    @test @changeprecision(Float32, pi+2) === Float32(pi)+2
    @test @changeprecision(Float32, -pi) === -Float32(pi)
    @test @changeprecision(Float32, 2Float64(pi)) === Float64(pi)*2
    @test @changeprecision(Float32, 2*2*2*2*2*Float32(pi)) === @changeprecision(Float32,Float32(pi)*2*2*2*2*2)=== 32*Float32(pi)
    @test @changeprecision(Float32, pi^2) === Float32(pi)^2
    @test @changeprecision(Float32, pi^2) === Float32(pi)^2
    @test @changeprecision(Float32, euler^2) === @changeprecision(Float32, euler^(3-1)) === exp(Float32(2))
end

@testset "powers" begin
    @test @changeprecision(Float32, 3^2) === @changeprecision(Float32, 3^(3-1)) === 9
    @test @changeprecision(Float32, 3^(1//2)) === sqrt(Float32(3))
    @test @changeprecision(Float32, 3^(1+2im)) === Float32(3)^(1+2im)
    @test @changeprecision(Float32, 0.5^2) === 0.25f0
end

@testset "arrays" begin
    @test @changeprecision(Float32, mean([1,3])) === @changeprecision(Float32, median([1,3])) === 2.0f0
    @test @changeprecision(Float32, norm([1,3],Inf)) === 3.0f0
    @test @changeprecision(Float32, varm([1,3],2)) === @changeprecision(Float32, varm([1,3],2.0)) === 2.0f0
    @test @changeprecision(Float32, varm([1,3], Float64(2))) === 2.0
    @test @changeprecision(Float32, [2,4]/2) ≡ @changeprecision(Float32, 2\[2,4]) ≡ Float32[1,2]
    @test @changeprecision(Float32, mean([1//1,3//1])) === 2//1
end

@testset "matrices" begin
    @test @changeprecision(Float32, inv(I+zeros(Float64,2,2))) ≡ I+zeros(Float64,2,2)
    @test @changeprecision(Float32, (I+zeros(Int,2,2)) \ (1:2)) ≡ Float32[1,2]
    @test @changeprecision(Float32, [1,2]' / (I+zeros(Int,2,2))) ≡ Float32[1,2]'
    @test @changeprecision(Float32, eigvals(I+zeros(Int,2,2))) ≡ @changeprecision(Float32, eigvals(I+zeros(Rational{Int},2,2))) ≡ eigvals(I+zeros(Float32,2,2))
    @test @changeprecision(Float32, norm(I+zeros(Int,2,2))) ≡ @changeprecision(Float32, norm(I+zeros(Rational{Int},2,2))) ≡ norm(I+zeros(Float32,2,2))
end

@testset "Dot macro (#9)" begin
    @test @changeprecision(Float32, @. 1 + [2.,3]) ≡Float32[3,4]
end

module Foo
using ChangePrecision
@changeprecision Float32 include("foo.jl")
end
import .Foo
@testset "include" begin
    @test Foo.foo(1) === 1.0f0 / 3
    @test Foo.foo(1.0) === 1.0 / 3
end

@testset "bigfloat" begin
    @changeprecision BigFloat foobar(x) = pi * x * 0.1
    for p in (50, 100, 1000)
        setprecision(p) do
            @test foobar(3) ≡ big(pi) * 3 * parse(BigFloat, "0.1")
        end
    end
end


text:
# f: linear frequency with standard unit (THz)
# c: speed of light, 2.99792458e-4 (m * THz)

export transmit, transmit!, reflect, reflect!, propagate, propagate!, ϵr2nκ!

transmit(n::Real, κ::Real; forward::Bool=true) = forward ? transmit(1.00027, 0.0, n, κ) : transmit(n, κ, 1.00027, 0.0)

function transmit(ni::Real, κi::Real, nj::Real, κj::Real)
    nipnj = ni + nj
    κipκj = κi + κj

    m = 2.0 * apy2(ni, κi) / apy2(nipnj, κipκj)
    θ = atan(κi * nipnj - ni * κipκj, ni * nipnj + κi * κipκj)
    return m, θ
end

reflect(n::Real, κ::Real; forward::Bool=true) = forward ? reflect(1.00027, 0.0, n, κ) : reflect(n, κ, 1.00027, 0.0)

function reflect(ni::Real, κi::Real, nj::Real, κj::Real)
    nipnj = ni + nj
    κipκj = κi + κj
    nimnj = ni - nj
    κimκj = κi - κj

    x = nimnj * nipnj + κimκj * κipκj
    y = nipnj * κimκj - nimnj * κipκj

    m = flipsign(apy2(nimnj, κimκj) / apy2(nipnj, κipκj), x)
    iszero(y) && return m, 0.0
    return m, atan(abs(y), abs(x))
end

propagate(d::Real, f::Real; forward::Bool=true) = forward ? (1.0, 1.00027 * d * f * 𝚷 / 2.99792458e-4) : (1.0, -1.00027 * d * f * 𝚷 / 2.99792458e-4)

function propagate(n::Real, κ::Real, f::Real, d::Real)
    dk = d * f * 𝚷 / 2.99792458e-4
    return exp(-κ * dk), n * dk
end

for fname in (:transmit, :reflect)
    fname! = Symbol(fname, '!')
    @eval begin
        function $(fname!)(rs::VecIO, θs::VecIO, n::VecI, κ::VecI; forward::Bool=true)
            @inbounds for i in eachindex(rs)
                rs[i], θs[i] = $(fname)(n[i], κ[i]; forward)
            end
        end
        function $(fname!)(rs::VecIO, θs::VecIO, n1::VecI, κ1::VecI, n2::VecI, κ2::VecI)
            @inbounds for i in eachindex(rs)
                rs[i], θs[i] = $(fname)(n1[i], κ1[i], n2[i], κ2[i])
            end
        end
    end
end

function propagate!(rs::VecIO, θs::VecIO, fs::VecI, d::Real; forward::Bool=true)
    ϕ = forward ? 1.00027 * d * 𝚷 : -1.00027 * d * 𝚷
    @simd for i in eachindex(θs)
        @inbounds rs[i], θs[i] = 1.0, ϕ * fs[i] / 2.99792458e-4
    end
end

function propagate!(rs::VecIO, θs::VecIO, n::VecI, κ::VecI, fs::VecI, d::Real)
    @inbounds for i in eachindex(rs)
        tmpdk = d * fs[i] * 𝚷 / 2.99792458e-4
        rs[i] = exp(-κ[i] * tmpdk)
        θs[i] = n[i] * tmpdk
    end
end

function ϵr2nκ!(ϵ1::VecIO, ϵ2::VecIO)
    @inbounds for i in eachindex(ϵ1)
        e1 = ϵ1[i]
        e2 = ϵ2[i]
        # modulus
        em = apy2(e1, e2)
        ϵ1[i] = sqrt(0.5 * (em + e1))
        ϵ2[i] = sqrt(0.5 * (em - e1))
    end
end


text:
module _DistEn2D
export DistEn2D
using StatsBase: fit, Histogram, skewness
    """
        Dist2D = DistEn2D(Mat)

    Returns the bidimensional distribution entropy estimate (`Dist2D`)
    estimated for the data matrix (`Mat`) using the default parameters:
    time delay = 1, histogram binning method = "sturges", logarithm = natural,
    template matrix size = [floor(H/10) floor(W/10)], (where H and W represent
    the height (rows) and width (columns) of the data matrix `Mat`) \n
    ** The minimum number of rows and columns of Mat must be > 10.**

        Dist2D = DistEn2D(Mat::AbstractArray{T,2} where T<:Real; m::Union{Int,Tuple{Int,Int}}=floor.(Int, size(Mat)./10), tau::Int=1,
                            Bins::Union{Int,String}="Sturges", Logx::Real=2, Norm::Int=2, Lock::Bool=true)

    Returns the bidimensional distribution entropy (`Dist2D`) estimate for
    the data matrix (`Mat`) using the specified 'keyword' arguments:

    # Arguments:
    `m`     - Template submatrix dimensions, an integer scalar (i.e. the same
              height and width) or a two-element tuple of integers
              [height, width] with a minimum value > 1.
              [default: [floor(H/10) floor(W/10)]] \n
    `tau`   - Time Delay, a positive integer       [default: 1]  \n
    `Bins`  - Histogram bin selection method for distance distribution,
              an integer > 1 indicating the number of bins, or one of the
              following strings {`"sturges", "sqrt", "rice", "doanes"``}
              [default: 'sturges']  \n
    `Logx`  - Logarithm base, a positive scalar        [default: natural]\n
              ** enter 0 for natural logarithm.**\n
    `Norm`  - Normalisation of `Dist2D` value, one of the following integers:
              [0]  no normalisation.
              [1]  normalises values of data matrix (`Mat`) to range [0 1].
              [2]  normalises values of data matrix (`Mat`) to range [0 1],
                   and normalises the distribution entropy value (`Dist2D`)
                   w.r.t the number of histogram bins.  [default]
              [3]  normalises the distribution entropy value
                   w.r.t the number of histogram bins, without normalising
                   data matrix values. \n
    `Lock`  - By default, DistEn2D only permits matrices with a maximum
              size of 128 x 128 to prevent memory errors when storing data on RAM.
              e.g. For Mat = [200 x 200], m = 3, and tau = 1, DistEn2D
              creates a vector of 753049836 elements. To enable matrices
              greater than [128 x 128] elements, set `Lock` to false.
              [default: 'true']
              `WARNING: unlocking the permitted matrix size may cause your Julia
               IDE to crash.`

    # See also `DistEn`, `XDistEn`, `SampEn2D`, `FuzzEn2D`, `MSEn`

    # References:
        [1] Hamed Azami, Javier Escudero and Anne Humeau-Heurtier,
            "Bidimensional distribution entropy to analyze the irregularity
            of small-sized textures."
            IEEE Signal Processing Letters
            24.9 (2017): 1338-1342.

    """
    function DistEn2D(Mat::AbstractArray{T,2} where T<:Real;
        m::Union{Int,Tuple{Int,Int}}=floor.(Int, size(Mat)./10), tau::Int=1,
        Bins::Union{Int,String}="Sturges", Logx::Real=2, Norm::Int=2, Lock::Bool=true)

    Logx == 0  ? Logx = exp(1) : nothing

    NL, NW = size(Mat)
    ((NL > 128 || NW > 128) && Lock) ?
        error("To prevent memory errors, matrix width & length must have <= 128 elements.
            To estimate DistEn2D  for the current matrix ($NL,$NW) change Lock to 'false'.
            Caution: unlocking the safe matrix size may cause the Julia IDE to crash.") :
            nothing

    length(m)==1 ?  (mL = m; mW = m) : (mL = m[1];  mW = m[2])
    (NL > 10 && NW > 10) ? nothing :
    error("Number of rows and columns in Mat must be > 10")
    (minimum(m) > 1) ? nothing : error("m:   must be an integer > 1")
    (tau >0) ? nothing : error("tau:   must be an integer > 0")
    (Logx>0) ? nothing : error("Logx:  must be a positive number > 0")
    if typeof(Bins)<:Int
        (Bins>1) ? nothing :
        error("Bins:  must be an integer > 1 (or name of binning method)")
    elseif typeof(Bins)<:String
        (lowercase(Bins) in ["sturges","sqrt","rice","doanes"]) ? nothing :
        error("Bins:    must be one of the following strings
                'sturges', 'sqrt', 'rice', 'doanes' (or an integer >1)")
    end
    (Norm in [0, 1, 2, 3]) ? nothing :  error("Norm:  must be an integer in the range [0 3]")

    Norm in [1, 2] ? Mat = (Mat.-minimum(Mat))./maximum(Mat.-minimum(Mat)) : nothing
    NL = NL - (mL-1)*tau
    NW = NW - (mW-1)*tau
    X = zeros(mL,mW,NL*NW)
    p = 0
    for k = 1:NL
        for n = 1:NW
            p += 1
            X[:,:,p] = Mat[k:tau:(mL-1)*tau+k,n:tau:(mW-1)*tau+n]
        end
    end

    p = size(X,3)
    p != NL*NW ? @warn("Potential error with submatrix division.") : nothing
    Ny = Int(p*(p-1)/2)
    Ny > 300000000 ? @warn("Number of pairwise distance calculations is $Ny") : nothing

    Y = zeros(Ny)
    for k = 1:p-1
        Ix = Int.([(k-1)*(p - k/2)+1, k*(p-((k+1)/2))])
        Y[Ix[1]:Ix[2]] = maximum(abs.(X[:,:,k+1:end] .- X[:,:,k]),dims=(1,2))
    end

    if eltype(Bins)<:Char
        if lowercase(Bins) == "sturges"
            Bx = ceil(log2(Ny) + 1)
        elseif lowercase(Bins) == "rice"
            Bx = ceil(2*(Ny^(1/3)))
        elseif lowercase(Bins) == "sqrt"
            Bx = ceil(sqrt(Ny))
        elseif lowercase(Bins) == "doanes"
            sigma = sqrt(6*(Ny-2)/((Ny+1)*(Ny+3)))
            Bx = ceil(1+log2(Ny)+log2(1+abs(skewness(convert(Array{Float64,1},Y))/sigma)))
        else
            error("Please enter a valid binning method")
        end
    else
        Bx = Bins
    end

    By = collect(range(minimum(Y),maximum(Y),length=Int(Bx+1)))
    By[end] += 1; By[1] -= 1

    Ppi = fit(Histogram, Y[:], By).weights/Ny
    if round(sum(Ppi),digits=6) != 1
        @warn("Potential error estimating probabilities (p = $(sum(Ppi)))")
        Ppi = Ppi[Ppi.!=0]
    elseif any(Ppi.==0)
        print("Note: $(sum(Ppi.==0))/$(length(Ppi)) bins were empty \n")
        Ppi = Ppi[Ppi.!=0]
    end
    Dist2D = -sum(Ppi.*log.(Logx, Ppi))
    Norm >= 2  ? Dist2D /= (log(Logx, Bx)) : nothing

    return Dist2D
    end

end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""

#= Y = []
for k = 1:p-1
    append!(Y,maximum(abs.(X[:,:,k+1:end] .- X[:,:,k]),dims=(1,2)))
end =#


text:
import ..update!, ..select_action, ..reset!

using ..Buffers:Experience

using StaticArrays
using Random
using LinearAlgebra:Diagonal, UniformScaling

using Ramnet.Models.AltDiscriminators:AdaptiveFunctionalDiscriminator
using Ramnet.Optimizers
import Ramnet.Loss

mutable struct AdaptiveContinousPolicyLoss <: Loss.AbstractLoss
    start_cov::Float64
    end_cov::Float64
    cov_decay::Int
    update_count::Int
end

function AdaptiveContinousPolicyLoss(start_cov::Float64, end_cov::Float64, cov_decay::Int)
    AdaptiveContinousPolicyLoss(start_cov, end_cov, cov_decay, 0)
end

struct AdaptiveContinousPolicy{OS,AS,T <: Real,O <: StaticArray{Tuple{OS},T,1},A <: StaticArray{Tuple{AS},T,1},E <: AbstractEncoder{T}}
    opt::AdaptiveOptimizer{AS,AdaptiveContinousPolicyLoss}
    f::AdaptiveFunctionalDiscriminator{AS,T,E}
    rng::MersenneTwister
end

function AdaptiveContinousPolicy(::Type{O}, ::Type{A}, n::Int, encoder::E; λ::Float64=1.0, start_cov::Float64, end_cov::Float64, cov_decay::Int, μ::Float64=0.1, η::Float64=0.1, epochs::Int=1, partitioner::Symbol=:uniform_random, seed::Union{Nothing,Int}=nothing) where {AS,OS,T <: Real,O <: StaticArray{Tuple{OS},T,1},A <: StaticArray{Tuple{AS},T,1},E <: AbstractEncoder{T}}
    AdaptiveContinousPolicy{OS,AS,T,O,A,E}(
        AdaptiveOptimizer{AS}(AdaptiveContinousPolicyLoss(start_cov, end_cov, cov_decay); λ, meta_rate=μ, initial_rate=η),
        AdaptiveFunctionalDiscriminator{AS}(OS, n, λ, encoder, partitioner; seed),
        isnothing(seed) ? MersenneTwister() : MersenneTwister(seed)
    )
end

function cov(l::AdaptiveContinousPolicyLoss)
    α = l.update_count / l.cov_decay
    UniformScaling((1 - α) * l.start_cov + α * l.end_cov)
end

function Loss.grad(l::AdaptiveContinousPolicyLoss, action, pred_action)
    inv(cov(l)) * (pred_action - action)
end

function Loss.hessian(l::AdaptiveContinousPolicyLoss, y_true, y_pred)
    inv(cov(l))
end

# TODO: Validate action (must be ether -1 or +1)
function update!(p::AdaptiveContinousPolicy{OS,AS,T,O,A,E}, observation::O, action::A, G) where {OS,AS,T,O,A,E}
    gradient, trace_update = Optimizers.grad(p.opt, p.f, observation, action)
    rate = SizedVector{AS}(learning_rate(p.opt))

    # @show rate
    # @show gradient
    # @show typeof(gradient)
    # @show rate .* G * gradient
    # @show rate .* trace_update
    # @show typeof(gradient)
    train!(p.f, observation, G * rate .* gradient, rate .* SizedVector{AS}(trace_update))
    p.opt.loss.update_count = p.opt.loss.update_count == p.opt.loss.cov_decay ? p.opt.loss.cov_decay : p.opt.loss.update_count + 1

    return nothing
end

# function update!(p::AdaptiveContinousPolicy{OS,AS,T,O,A,E}, experiences::AbstractVector{Experience{O,A}}, critic) where {OS,AS,T,O,A,E}
#     gradients = Matrix{Float64}(undef, AS, length(experiences))
#     observations = Matrix{T}(undef, OS, length(experiences))

#     for i in eachindex(experiences)

#         observation = experiences[i].observation
#         action = experiences[i].action
#         value = experiences[i].value
#         advantage = value - critic(observation)
#         gradients[:, i] = Optimizers.grad(p.opt, p, observation, action, advantage)
#         observations[:, i] = observation
#     end

#     train!(p.f, observations, SizedMatrix{AS,length(experiences)}(learning_rate(p.opt) * gradients))
# p.update_count = p.update_count == p.cov_decay ? p.cov_decay : p.update_count + 1

#     nothing
# end

function select_action(p::AdaptiveContinousPolicy{OS,AS,T,O,A,E}, observation::O) where {OS,AS,T,O,A,E}
return SizedVector{AS,T}(cov(p.opt.loss) * randn(p.rng, AS) + predict(p.f, observation)[1])
end

function reset!(p::AdaptiveContinousPolicy)
    Ramnet.reset!(p.f)

    return nothing
end


text:
@testset "FeedbackTrees" begin
    @testset "constructor" begin
        # 1. can construct a FeedbackTree like a normal chain
        c = FeedbackTree(Dense(10, 5), relu)
        @test length(c.layers) == 2
        # 2. can construct a FeedbackTree with forkers and mergers
        c = FeedbackTree(Merger("f1", identity, +), Dense(10, 10, relu), Splitter("f1"), Dense(10, 1))
        @test length(c.layers) == 4
    end # @testset "constructor"

    @testset "apply" begin
        h = Dict("f1" => rand(10))
        x = randn(10)
        # apply empty tree
        c = FeedbackTree()
        @test c(h, x) == (h, x)
        # apply tree without feedback
        c = FeedbackTree(x -> (2.0 .* x), x -> relu.(x), x -> x'x)
        state, y = c(Dict(), x)
        @test y ≈ relu.(2.0 .* x)'relu.(2.0 .* x)
        @test state isa Dict{String, Any}
        # apply tree with feedback
        l1 = Dense(10, 10, relu)
        l2 = Dense(10, 1)
        c = FeedbackTree(Merger("f1", identity, +), l1, Splitter("f1"), l2)
        state, y = c(h, x)
        @test y ≈ l2(h["f1"])
        @test haskey(state, "f1")
        @test state["f1"] ≈ l1(h["f1"] + x)
    end # @testset "apply"

    @testset "params" begin
        l1 = Dense(10, 10)
        l2 = Dense(7, 10)
        c = FeedbackTree(l1, l2)
        @test hcat(params(c)...) == hcat(params(l1)..., params(l2)...)
    end # @testset "params"

    @testset "recur" begin
        h = Dict("f1" => randn(10))
        x = randn(10)
        l1 = Dense(10, 10, relu)
        l2 = Dense(10, 1)
        c = FeedbackTree(Merger("f1", identity, +), l1, Splitter("f1"), l2)
        c = Flux.Recur(c, h)
        y = c(x)
        @test y ≈ l2(h["f1"])
        @test haskey(c.state, "f1")
        @test c.state["f1"] ≈ l1(h["f1"] + x)
    end # @testset "recur"
end # @testset "FeedbackTrees"


text:
# All values will be English names so that the `yaml` files will be human-readable.
abstract type VsAbstractID{T} <: SimpleTypeWrapper{T} end
macro _id_type_wrapper(name, i18n_key)
    i18n_key = string(i18n_key)
    convert_functions = esc(quote
        $name(id::String) = $name(Symbol(id))
        Base.convert(::Type{$name}, id::String) = $name(id)
    end)
    i18n_function = esc(:(get_i18n_namespace(::Type{$name}) = $i18n_key))
    quote
        @type_wrapper $name Symbol nothing VsAbstractID
        $convert_functions
        $i18n_function
    end
end
@_id_type_wrapper PokemonID pokemon
@_id_type_wrapper ItemID item
@_id_type_wrapper MoveID move
@_id_type_wrapper AbilityID ability

SimpleI18n.i18n(id::VsAbstractID; language = nothing) = get_i18n(
    get_i18n_namespace(typeof(id)), id,
    lang = language, default = id
)

@enum Gender begin
    GENDER_NULL
    GENDER_MALE
    GENDER_FEMALE
end
VsRecorderBase.enum_prefix(::Type{Gender}) = "GENDER_"

@enum PokemonType begin
    TYPE_BIRD = 0
    TYPE_NORMAL
    TYPE_FIGHTING
    TYPE_FLYING
    TYPE_POISON
    TYPE_GROUND
    TYPE_ROCK
    TYPE_BUG
    TYPE_GHOST
    TYPE_STEEL
    TYPE_FIRE
    TYPE_WATER
    TYPE_GRASS
    TYPE_ELECTRIC
    TYPE_PSYCHIC
    TYPE_ICE
    TYPE_DRAGON
    TYPE_DARK
    TYPE_FAIRY
end
VsRecorderBase.enum_prefix(::Type{PokemonType}) = "TYPE_"

const STATS_FIELDS = (:hp, :attack, :defense, :special_attack, :special_defense, :speed)
const DEX_STATS_FIELDS = (:hp, :atk, :def, :spa, :spd, :spe)

@enum Nature begin
    NATURE_HARDY = 0
    NATURE_LONELY
    NATURE_ADAMANT
    NATURE_NAUGHTY
    NATURE_BRAVE
    NATURE_BOLD
    NATURE_DOCILE
    NATURE_IMPISH
    NATURE_LAX
    NATURE_RELAXED
    NATURE_MODEST
    NATURE_MILD
    NATURE_BASHFUL
    NATURE_RASH
    NATURE_QUIET
    NATURE_CALM
    NATURE_GENTLE
    NATURE_CAREFUL
    NATURE_QUIRKY
    NATURE_SASSY
    NATURE_TIMID
    NATURE_HASTY
    NATURE_JOLLY
    NATURE_NAIVE
    NATURE_SERIOUS
end
VsRecorderBase.enum_prefix(::Type{Nature}) = "NATURE_"
Nature(s::AbstractString) = enum_from_string(s, Nature)
SimpleI18n.i18n(id::Nature; language = nothing) = get_i18n("nature", enum_to_string(id), lang = language)

const EnumIDTypes = Union{PokemonType, Nature}

const NATURE_NEUTURAL_DEFAULT = NATURE_HARDY
function get_nature_effect(nature::Nature)
    up, down = divrem(Int64(nature), 5)
    up ≡ down && return (nothing, nothing)
    (DEX_STATS_FIELDS[up + 1], DEX_STATS_FIELDS[down + 1])
end
get_nature_effect(::Missing) = (nothing, nothing)

function get_nature(up, down)
    a = findfirst(==(up), DEX_STATS_FIELDS)
    b = findfirst(==(down), DEX_STATS_FIELDS)
    a ≡ b && return NATURE_NEUTURAL_DEFAULT
    Nature((a - 1) * 5 + (b - 1))
end


text:
# Memory structures
abstract type Memory end

#=
    Persistent represents persistent memory, such as global variables, local
    variables and constants.
=#
struct Persistent <: Memory
    type1::Vector{Union{Int64, Nothing}}
    type2::Vector{Union{Float64, Nothing}}
    type3::Vector{Union{String, Nothing}}
end

#=
    Temporary represents all temporary variables.
=#
struct Temporary <: Memory
    type1::Vector{Union{Int64, Nothing}}
    type2::Vector{Union{Float64, Nothing}}
    type3::Vector{Union{Bool, Nothing}}
    type4::Vector{Union{UInt16, Nothing}}
end

#=
    GlobalMem represents the global memory of the virtual machine, consisting of
    global variables and all the constants.
=#
struct GlobalMem
    variables::Persistent
    constants::Persistent
end

#=
    MemoryObj represents the local memory that will go into the MemoryStack,
    which includes local persistent variables and temporary variables.
=#
struct MemoryObj
    persistent::Persistent
    temporary::Temporary
end

#=
    getMemory(mem, type):
        Given a memory instance, it will return the specific memory vector to
        modify depending on the type received.
=#
function getMemory(mem::Memory, type::Char)
    type === '1' && return mem.type1
    type === '2' && return mem.type2
    type === '3' && return mem.type3
    mem.type4
end

#=
    fetch(mem, address, type):
        Given a memory instance, an address and a type, it will attempt to fetch
        the value inside a memory array and return it.
=#
function fetch(mem::Memory, address::UInt16, type::Char)
    memory = getMemory(mem, type)
    length(memory) === false || length(memory) < address && return false
    memory[address]
end

#=
    store(mem, value, address, type):
        Given a memory instance, a value, an address and a type, it will push or
        modify the value inside a memory array that is represented by the address.
=#
function store(mem::Memory, value::Any, address::UInt16, type::Char)
    memory = getMemory(mem, type)
    if length(memory) >= address
        memory[address] = value
        return
    end
    if length(memory) + 1 == address
        push!(memory, value)
        return
    end
    while length(memory) + 1 < address
        push!(memory, nothing)
    end
    push!(memory, value)
end

# Static values
ranges = [1000:2999,
          3000:4999,
          5000:5999,
          6000:7999,
          8000:9999,
          10000:10999,
          11000:15999,
          16000:20999,
          21000:25999,
          26000:27999,
          28000:29999,
          30000:30999,
          31000:31999]

operators = ["++", "--", "^", "*", "/", "//", "+", "-", "<", "<=", ">", ">=", "==", "¬=", "and", "or", "<-", "<+>"]
stats = ["Size", "Mean", "Median", "Mode", "Variance", "Std", "Range", "Sum", "Min", "Max"]
xplots = ["Histogram", "Violin", "Box"]
xyplots = ["Bar", "Scatter"]

max_stars_repo_name: CharlesGovea/Alice

max_issues_repo_name: CharlesGovea/Alice

max_forks_repo_name: CharlesGovea/Alice


text:
"""
    rstar(rng=Random.GLOBAL_RNG, classifier, chains::Chains; kwargs...)

Compute the ``R^*`` convergence diagnostic of the MCMC `chains` with the `classifier`.

The keyword arguments supported here are the same as those in `rstar` for arrays of samples
and chain indices.

# Examples

```jldoctest rstar; setup = :(using Random; Random.seed!(100))
julia> using MLJBase, MLJXGBoostInterface, Statistics

julia> chains = Chains(fill(4.0, 100, 2, 3));
```

One can compute the distribution of the ``R^*`` statistic with the probabilistic classifier.

```jldoctest rstar
julia> distribution = rstar(XGBoostClassifier(), chains);

julia> isapprox(mean(distribution), 1; atol=0.1)
true
```

For deterministic classifiers, a single ``R^*`` statistic is returned.

```jldoctest rstar
julia> @pipeline XGBoostClassifier name = XGBoostDeterministic operation = predict_mode;

julia> value = rstar(XGBoostDeterministic(), chains);

julia> isapprox(value, 1; atol=0.2)
true
```
"""
function MCMCDiagnosticTools.rstar(
    classif::MLJModelInterface.Supervised, chn::Chains; kwargs...
)
    return MCMCDiagnosticTools.rstar(Random.GLOBAL_RNG, classif, chn; kwargs...)
end

function MCMCDiagnosticTools.rstar(
    rng::Random.AbstractRNG, classif::MLJModelInterface.Supervised, chn::Chains; kwargs...
)
    nchains = size(chn, 3)
    nchains <= 1 && throw(DimensionMismatch())

    # collect data
    x = Array(chn)
    y = repeat(chains(chn); inner = size(chn,1))

    return MCMCDiagnosticTools.rstar(rng, classif, x, y; kwargs...)
end


text:
mutable struct KernelRegression{K, I, O, C}
    kernel::K
    input::I
    output::O
    cs::Union{Nothing, O}
    cost::C
end

"""
    KernelRegression(X::AbstractArray, Y::AbstractArray;
        kernel=gaussiankernel)

Return a `KernelRegression` object, to perform Kernel Ridge regression.
`X` and `Y` are training input and output. `Y` can be either a vector or a matrix.
`kernel` is required to have a method `kernel(u, v)`, where `u` and `v` are matrices
(datapoints are columns).
See [`additivegaussiankernel`](@ref) for an example. `cost` is used to regularize.

Use `fit!(kr::KernelRegression)` to fit the `KernelRegression` object, and
`predict(kr::KernelRegression, input)` to get the model prediction on a dataset.
If no `input` is given, it will default to the training data.
"""
function KernelRegression(X::AbstractArray, Y::AbstractArray;
    kernel=gaussiankernel, cost=0)

    return KernelRegression(kernel, X, Y, nothing, cost)
end

function fit!(kr::KernelRegression)

    input, output, kernel, cost = kr.input, kr.output, kr.kernel, kr.cost
    m = kernel(input', input')
    m += cost * size(input, 1) * I

    kr.cs = m \ output
    return kr
end

function predict(kr::KernelRegression, input=kr.input)
    return kr.kernel(input', kr.input') * kr.cs
end