text:
import Observables
import AbstractPlotting

export Frame
export update

"""
    Represents a frame.

fields: circle, s2tos3map, s2tos2map, color, s3rotation, config, segments and observable.
"""
mutable struct Frame <: Sprite
    circle::S¹
    s2tos3map::Any
    s2tos2map::Any
    color::Any
    s3rotation::S³
    config::Biquaternion
    segments::Int
    observable::Tuple{Observables.Observable{Array{Float64,2}},
                      Observables.Observable{Array{Float64,2}},
                      Observables.Observable{Array{Float64,2}}}
end

"""
    Frame(scene, circle, s2tos3map, s2tos2map, color,
          [s3rotation, [config, [segments, [transparency]]]])

Construct a Frame with the given `scene`, `circle` in the fiber space, function `f` that
takes the base space to itself, S³ rotation `s3rotation`, configuration `config`, the number
of `segments`, `color` and `transparency`.
"""
function Frame(scene::AbstractPlotting.Scene,
               circle::S¹,
               s2tos3map,
               s2tos2map,
               color::Any;
               s3rotation::S³ = Quaternion(1, 0, 0, 0),
               config::Biquaternion = Biquaternion(ℝ³(0, 0, 0)),
               segments::Int = 36,
               transparency::Bool = false)
    frame = constructframe(circle,
                           s2tos3map,
                           s2tos2map,
                           s3rotation = s3rotation,
                           config = config,
                           segments = segments)
    observable = buildsurface(scene, frame, color, transparency = transparency)
    Frame(circle, s2tos3map, s2tos2map, color, s3rotation, config, segments, observable)
end

"""
    update(frame, circle)

Update a Frame by changing its observable with the given `frame` and `circle` in the fiber
space.
"""
function update(frame::Frame, circle::S¹)
    frame.circle = circle
    value = constructframe(frame.circle,
                           frame.s2tos3map,
                           frame.s2tos2map,
                           s3rotation = frame.s3rotation,
                           config = frame.config,
                           segments = frame.segments)
    updatesurface(value, frame.observable)
end

"""
    update(frame, s2tos3map, s2tos2map)

Update a Frame by changing its observable with the given `frame` and map `s2tos3map` from
the base space into the total space, f: S² → S³, and also the map `s2tos2map` from the base
space into itself, f: S² → S².
"""
function update(frame::Frame, s2tos3map::Any, s2tos2map::Any)
    frame.s2tos3map = s2tos3map
    frame.s2tos2map = s2tos2map
    value = constructframe(frame.circle,
                           frame.s2tos3map,
                           frame.s2tos2map,
                           s3rotation = frame.s3rotation,
                           config = frame.config,
                           segments = frame.segments)
    updatesurface(value, frame.observable)
end

"""
    update(frame, s3rotation)

Update a Frame by changing its observable with the given `frame` and S³ rotation
`s3rotation`.
"""
function update(frame::Frame, s3rotation::S³)
    frame.s3rotation = s3rotation
    value = constructframe(frame.circle,
                           frame.s2tos3map,
                           frame.s2tos2map,
                           s3rotation = frame.s3rotation,
                           config = frame.config,
                           segments = frame.segments)
    updatesurface(value, frame.observable)
end

"""
    update(frame, config)

Update a Frame by changing its observable with the given `frame` and configuration `config`.
"""
function update(frame::Frame, config::Biquaternion)
    frame.config = config
    value = constructframe(frame.circle,
                           frame.s2tos3map,
                           frame.s2tos2map,
                           s3rotation = frame.s3rotation,
                           config = frame.config,
                           segments = frame.segments)
    updatesurface(value, frame.observable)
end


text:
module TestComponents
using Dash
export html_div, html_a, dcc_input, dcc_graph, html_h1, html_img

function html_div(; kwargs...)
    available_props = Symbol[:children, :id, :n_clicks, :n_clicks_timestamp, :key, :role, :accessKey, :className, :contentEditable, :contextMenu, :dir, :draggable, :hidden, :lang, :spellCheck, :style, :tabIndex, :title, :loading_state]
    wild_props = Symbol[Symbol("data-"), Symbol("aria-")]
    return Component("html_div", "Div", "dash_html_components", available_props, wild_props; kwargs...)
end

html_div(children::Any; kwargs...) = html_div(;kwargs..., children = children)
html_div(children_maker::Function; kwargs...) = html_div(children_maker(); kwargs...)

function html_h1(; kwargs...)
    available_props = Symbol[:children, :id, :n_clicks, :n_clicks_timestamp, :key, :role, :accessKey, :className, :contentEditable, :contextMenu, :dir, :draggable, :hidden, :lang, :spellCheck, :style, :tabIndex, :title, :loading_state]
    wild_props = Symbol[Symbol("data-"), Symbol("aria-")]
    return Component("html_h1", "H1", "dash_html_components", available_props, wild_props; kwargs...)
end

html_h1(children::Any; kwargs...) = html_h1(;kwargs..., children = children)
html_h1(children_maker::Function; kwargs...) = html_h1(children_maker(); kwargs...)

function html_img(; kwargs...)
    available_props = Symbol[:children, :id, :n_clicks, :n_clicks_timestamp, :key, :role, :alt, :crossOrigin, :height, :sizes, :src, :srcSet, :useMap, :width, :accessKey, :className, :contentEditable, :contextMenu, :dir, :draggable, :hidden, :lang, :spellCheck, :style, :tabIndex, :title, :loading_state]
    wild_props = Symbol[Symbol("data-"), Symbol("aria-")]
    return Component("html_img", "Img", "dash_html_components", available_props, wild_props; kwargs...)
end

html_img(children::Any; kwargs...) = html_img(;kwargs..., children = children)
html_img(children_maker::Function; kwargs...) = html_img(children_maker(); kwargs...)

function html_a(; kwargs...)
    available_props = Symbol[:children, :id, :n_clicks, :n_clicks_timestamp, :key, :role, :download, :href, :hrefLang, :media, :rel, :shape, :target, :accessKey, :className, :contentEditable, :contextMenu, :dir, :draggable, :hidden, :lang, :spellCheck, :style, :tabIndex, :title, :loading_state]
    wild_props = Symbol[Symbol("data-"), Symbol("aria-")]
    return Component("html_a", "A", "dash_html_components", available_props, wild_props; kwargs...)
end

html_a(children::Any; kwargs...) = html_a(;kwargs..., children = children)
html_a(children_maker::Function; kwargs...) = html_a(children_maker(); kwargs...)

function dcc_input(; kwargs...)
    available_props = Symbol[:id, :value, :style, :className, :debounce, :type, :autoComplete, :autoFocus, :disabled, :inputMode, :list, :max, :maxLength, :min, :minLength, :multiple, :name, :pattern, :placeholder, :readOnly, :required, :selectionDirection, :selectionEnd, :selectionStart, :size, :spellCheck, :step, :n_submit, :n_submit_timestamp, :n_blur, :n_blur_timestamp, :loading_state, :persistence, :persisted_props, :persistence_type]
    wild_props = Symbol[]
    return Component("dcc_input", "Input", "dash_core_components", available_props, wild_props; kwargs...)
end

function dcc_graph(; kwargs...)
    available_props = Symbol[:id, :responsive, :clickData, :clickAnnotationData, :hoverData, :clear_on_unhover, :selectedData, :relayoutData, :extendData, :restyleData, :figure, :style, :className, :animate, :animation_options, :config, :loading_state]
    wild_props = Symbol[]
    return Component("dcc_graph", "Graph", "dash_core_components", available_props, wild_props; kwargs...)
end
end

ext: jl


text:
#--------------------------
# Total Raditiave Forcing.
#--------------------------

@defcomp rf_total begin

    α             = Parameter()             # Scaling factor to capture uncertainty in total aerosol radiative forcing.
    rf_O₃         = Parameter(index=[time]) # Radiative forcing from tropospheric ozone (Wm⁻²).
    rf_CO₂        = Parameter(index=[time]) # Radiative forcing from carbon dioxide (Wm⁻²).
    rf_CH₄        = Parameter(index=[time]) # Radiative forcing from methane (Wm⁻²).
    rf_CH₄_H₂O    = Parameter(index=[time]) # Radiative forcing from stratospheric water vapor due to methane oxidiation (Wm⁻²).
    rf_aerosol    = Parameter(index=[time]) # Radiative forcing from aerosols (Wm⁻²).
    rf_exogenous  = Parameter(index=[time]) # Radiative forcing from other (exogenous) sources (Wm⁻²).

    total_forcing = Variable(index=[time])  # Total radiative forcing from all sources (Wm⁻²).

    function run_timestep(p, v, d, t)

        if is_first(t)
            # Set initial radiative forcing to 0.
            v.total_forcing[t] = 0.0
        else
            # Calculate total radiative forcing (with aerosol contribution scaled by α).
            v.total_forcing[t] = p.rf_CO₂[t] + p.rf_CH₄[t] + p.rf_O₃[t] + p.rf_CH₄_H₂O[t] + (p.α * p.rf_aerosol[t]) + p.rf_exogenous[t]
        end
    end
end


text:
todisplay = true

using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using CAGD
using SparseArrays
if todisplay  using ViewerGL; GL = ViewerGL  end

include("./updatedLarFunctions.jl");

include("./cellular_complex.jl");
include("./CSG_example.jl");
include("./cycles_boundaries.jl");
include("./planar_algebra_2D.jl");
include("./solid_algebra_3D.jl");
include("./space_arrangement_2d.jl");
include("./space_arrangement.jl");
include("./subdivision_of_2_cells.jl");


text:
using DifferentialEquations, RecursiveArrayTools
using NLopt
Xiang2015Bounds = Tuple{Float64, Float64}[(9, 11), (20, 30), (2, 3)]

g1 = @ode_def_nohes LorenzExample begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ=>10.0 ρ=>28.0 β=>2.6666

r0 = [0.1;0.0;0.0]
tspan = (0.0,4.0)
prob = ODEProblem(g1,r0,tspan)
tspan2 = (0.0,3.0)
prob_short = ODEProblem(g1,r0,tspan2)

dt = 0.001
tf = 4.0
tinterval = 0:dt:tf
t  = collect(tinterval)

h = 0.01
M = 300
tstart = 0.0
tstop = tstart + M * h
tinterval_short = 0:h:tstop
t_short = collect(tinterval_short)

# Generate Data
data_short = convert(Array,solve(prob_short,Euler(),tstops=t_short))
data = convert(Array,solve(prob,Euler(),tstops=t))

# Use BlackBoxOptim
obj_short = build_loss_objective(prob_short,Euler(),L2Loss(t_short,data_short),tstops=t_short,dense=false)
res1 = bboptimize(obj_short;SearchRange = Xiang2015Bounds, MaxSteps = 11e3)

# Use NLopt
opt = Opt(:GN_ORIG_DIRECT_L, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:GN_CRS2_LM, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:GN_ISRES, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LN_BOBYQA, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LN_NELDERMEAD, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LD_SLSQP, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

#####################
# Fails to converge
#####################

opt = Opt(:GN_ESCH, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LN_COBYLA, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LN_NEWUOA_BOUND, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LN_PRAXIS, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LN_SBPLX, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LD_MMA, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LD_LBFGS, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LD_TNEWTON_PRECOND_RESTART, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LD_VAR2, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

########

#### Now let's solve the longer version

obj = build_loss_objective(prob,Euler(),L2Loss(t,data),tstops=t,dense=false)
# res1 = bboptimize(obj;SearchRange = Xiang2015Bounds, MaxSteps = 8e3)

opt = Opt(:GN_ORIG_DIRECT_L, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:GN_CRS2_LM, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:GN_ISRES, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LN_BOBYQA, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LN_NELDERMEAD, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])

opt = Opt(:LD_SLSQP, 3)
lower_bounds!(opt,[9.0,20.0,2.0])
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,[9.0,20.0,2.0])


text:
# Example for basic 2D modeling:
# The receiver positions and the source wavelets are the same for each of the four experiments.
# Author: MAthias Loubout, mlouboutin3@gatech.edu
# Date: August 2021
#

using JUDI4Cloud

creds = "/home/mloubout/research/azure/clusterless_creds.json"
init_culsterless(2; credentials=creds, vm_size="Standard_E2s_v3", pool_name="PoolTest", verbose=1)
# Set up model structure
n = (120, 100)   # (x,y,z) or (x,z)
d = (10., 10.)
o = (0., 0.)

# Velocity [km/s]
v = ones(Float32,n) .+ 0.5f0
v0 = ones(Float32,n) .+ 0.5f0
v[:,Int(round(end/2)):end] .= 3.5f0
rho = (v0 .+ .5f0) ./ 2

# Slowness squared [s^2/km^2]
m = (1f0 ./ v).^2
m0 = (1f0 ./ v0).^2
dm = vec(m - m0)

# Setup info and model structure
nsrc = 2	# number of sources
model = Model(n, d, o, m)
model0 = Model(n, d, o, m0)

# Set up receiver geometry
nxrec = 120
xrec = range(50f0, stop=1150f0, length=nxrec)
yrec = 0f0
zrec = range(50f0, stop=50f0, length=nxrec)

# receiver sampling and recording time
timeR = 1000f0   # receiver recording time [ms]
dtR = 2f0    # receiver sampling interval [ms]

# Set up receiver structure
recGeometry = Geometry(xrec, yrec, zrec; dt=dtR, t=timeR, nsrc=nsrc)

# Set up source geometry (cell array with source locations for each shot)
xsrc = convertToCell(range(400f0, stop=800f0, length=nsrc))
ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
zsrc = convertToCell(range(200f0, stop=200f0, length=nsrc))

# source sampling and number of time steps
timeS = 1000f0  # ms
dtS = 2f0   # ms

# Set up source structure
srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtS, t=timeS)

# setup wavelet
f0 = 0.01f0     # kHz
wavelet = ricker_wavelet(timeS, dtS, f0)
q = judiVector(srcGeometry, wavelet)

# Set up info structure for linear operators
ntComp = get_computational_nt(srcGeometry, recGeometry, model)
info = Info(prod(n), nsrc, ntComp)

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

# Write shots as segy files to disk
opt = JUDI.Options(optimal_checkpointing=false, isic=false, subsampling_factor=2, dt_comp=1.0)

# Setup operators
Pr = judiProjection(info, recGeometry)
F = judiModeling(info, model; options=opt)
F0 = judiModeling(info, model0; options=opt)
Ps = judiProjection(info, srcGeometry)
J = judiJacobian(Pr*F0*adjoint(Ps), q)

# Nonlinear modeling
dobs = Pr*F*adjoint(Ps)*q
# # Adjoint
qad = Ps*adjoint(F)*adjoint(Pr)*dobs

# Linearized modeling
dD = J*dm
# Adjoint jacobian
rtm = adjoint(J)*dD

# # evaluate FWI objective function
f, g = fwi_objective(model0, q, dobs; options=opt)

# evaluate LSRTM objective function
fj, gj = lsrtm_objective(model0, q, dD, dm; options=opt)
fjn, gjn = lsrtm_objective(model0, q, dobs, dm; nlind=true, options=opt)

# By extension, lsrtm_objective is the same as fwi_objecive when `dm` is zero
# And with computing of the residual. Small noise can be seen in the difference
# due to floating point roundoff errors with openMP, but running with
# OMP_NUM_THREAS=1 (no parllelism) produces the exact (difference == 0) same result
# gjn2 == g
fjn2, gjn2 = lsrtm_objective(model0, q, dobs, 0f0.*dm; nlind=true, options=opt)


text:
using Flux
import Zygote
using CUDA
using TensorOperations

include("model_utils.jl")

# https://github.com/FluxML/Flux.jl/issues/160
function weight_params(m::Chain, ps=Flux.Params())
    map((l)->weight_params(l, ps), m.layers)
    ps
end
weight_params(m::Dense, ps=Flux.Params()) = push!(ps, m.W)
weight_params(m::Conv, ps=Flux.Params()) = push!(ps, m.weight)
weight_params(m::ConvTranspose, ps=Flux.Params()) = push!(ps, m.weight)
weight_params(m, ps=Flux.Params()) = ps

function test_weight_params()
    # get model
    weight_params(model)
end

struct Equivariant{T}
    w1::AbstractArray{T}
    w2::AbstractArray{T}
    w3::AbstractArray{T}
    w4::AbstractArray{T}
    w5::AbstractArray{T}
    # FIXME how Flux decides ch is not a trainable parameter?
    ch::Pair{<:Integer,<:Integer}
end
Flux.@functor Equivariant

function weight_params(m::Equivariant, ps=Flux.Params())
    push!(ps, m.w1)
    push!(ps, m.w2)
    push!(ps, m.w3)
    push!(ps, m.w4)
    push!(ps, m.w5)
end

mynfan(dims...) = prod(dims[1:end-2]) .* (dims[end-1], dims[end])
my_glorot_uniform(dims...) = (rand(Float32, dims...) .- 0.5f0) .* sqrt(24.0f0 / sum(mynfan(dims...)))

function Equivariant(ch::Pair{<:Integer,<:Integer})
    # FIXME NOW init function
    return Equivariant(
        my_glorot_uniform(ch[1],ch[2]),
        my_glorot_uniform(ch[1],ch[2]),
        my_glorot_uniform(ch[1],ch[2]),
        my_glorot_uniform(ch[1],ch[2]),
        Float32.(zeros(ch[2])),
        ch)
end

"""
Y = w1 X + w2 11ᵀ X + w3 X 11ᵀ + w4 11ᵀ X 11ᵀ + w5 11ᵀ
"""
function eqfn(X::AbstractArray,
              w1::AbstractArray, w2::AbstractArray, w3::AbstractArray,
              w4::AbstractArray, w5::AbstractArray)
    d = size(X,1)
    if typeof(X) <: CuArray
        # FIXME performance hell!!!
        one = CUDA.ones(d,d)
    else
        # support CPU as well
        # DEPRECATED convert to Float32, when moving this to GPU
        one = ones(Float32, d, d)
    end
    # FIXME CUTENSOR_STATUS_ARCH_MISMATCH error: cutensor only supports RTX20
    # series (with compute capability 7.0+) see:
    # https://developer.nvidia.com/cuda-gpus
    @tensor X1[a,b,ch2,batch] := X[a,b,ch1,batch] * w1[ch1,ch2]

    # TODO how to do max pooling?
    @tensor X2[a,c,ch2,batch] := one[a,b] * X[b,c,ch1,batch] * w2[ch1,ch2]
    @tensor X3[a,c,ch2,batch] := X[a,b,ch1,batch] * one[b,c] * w3[ch1,ch2]

    @tensor X4[a,d,ch2,batch] := one[a,b] * X[b,c,ch1,batch] * one[c,d] * w4[ch1,ch2]
    @tensor X5[a,b,ch2] := one[a,b] * w5[ch2]
    # FIXME broadcasting X5
    Y = X1 + (X2 ./ d) + (X3 ./ d) + (X4 ./ (d * d)) .+ X5
end

# from https://github.com/mcabbott/TensorGrad.jl
Zygote.@adjoint function eqfn(X::AbstractArray,
                              w1::AbstractArray, w2::AbstractArray, w3::AbstractArray,
                              w4::AbstractArray, w5::AbstractArray)
    d = size(X,1)
    if typeof(X) <: CuArray
        # FIXME performance hell!!!
        one = CUDA.ones(d,d)
    else
        one = ones(Float32, d, d)
    end
    eqfn(X, w1, w2, w3, w4, w5), function (ΔY)
        # ΔY is FillArray.Fill, and this is not handled in @tensor. Convert it
        # to normal array here. FIXME performance
        # ΔY = Array(ΔY)
        #
        # FIXME However, this will change CuArray to regular Array. Removing
        # this would work for CuArray, but not on CPU. I probably won't
        # calculate gradient on CPU anyway.
        @tensor ΔX1[a,b,ch1,batch] := ΔY[a,b,ch2,batch] * w1[ch1,ch2]

        @tensor ΔX2[a,c,ch1,batch] := one[a,b] * ΔY[b,c,ch2,batch] * w2[ch1,ch2]
        @tensor ΔX3[a,c,ch1,batch] := ΔY[a,b,ch2,batch] * one[b,c] * w3[ch1,ch2]
        @tensor ΔX4[a,d,ch1,batch] := one[a,b] * ΔY[b,c,ch2,batch] * one[c,d] * w4[ch1,ch2]

        # DEBUG new formula for ΔX[2:4]
        # @tensor ΔX2[a,b,ch1,batch] := ΔY[a,b,ch2,batch] * w[ch1,ch2]
        # @tensor ΔX3[a,b,ch1,batch] := ΔY[a,b,ch2,batch] * w[ch1,ch2]
        # @tensor ΔX4[a,b,ch1,batch] := ΔY[a,b,ch2,batch] * γ[ch1,ch2]
        # ΔX5 is 0

        # DEBUG even newer NOT WORKING
        # delta = w1 .+ w2 .+ w3 .+ w4 .+ reshape(w5, 1,:)
        # @tensor ΔX[a,b,ch1,batch] := ΔY[a,b,ch2,batch] * delta[ch1,ch2]

        @tensor Δw1[ch1,ch2] := X[a,b,ch1,batch] * ΔY[a,b,ch2,batch]
        @tensor Δw2[ch1,ch2] := one[a,b] * X[b,c,ch1,batch] * ΔY[a,c,ch2,batch]
        @tensor Δw3[ch1,ch2] := X[a,b,ch1,batch] * one[b,c] * ΔY[a,c,ch2,batch]
        @tensor Δw4[ch1,ch2] := one[a,b] * X[b,c,ch1,batch] * one[c,d] * ΔY[a,d,ch2,batch]
        # FIXME should I normalize Δb?
        @tensor Δw5[ch2,batch] := one[a,b] * ΔY[a,b,ch2,batch]
        # FIXME can I just do a mean here?
        Δw5 = dropdims(mean(Δw5, dims=2), dims=2)

        return (
            ΔX1
            + ΔX2 ./ d + ΔX3 ./ d + ΔX4 ./ (d*d),
            # + ΔX2 + ΔX3 + ΔX4,
            # ΔX,
            Δw1,
            Δw2 ./ d,
            Δw3 ./ d,
            Δw4 ./ (d*d),
            Δw5)
    end
end

function (l::Equivariant)(X)
    # X = X[:,:,:,:]
    eqfn(X, l.w1, l.w2, l.w3, l.w4, l.w5)
end

struct Reshape
    target
end

function (l::Reshape)(x)
    reshape(x, l.target...)
end

function fc_model(; d, ch=1, z=1024, reg=false, nlayer=3)
    nlayer >= 2 || error("nlayer must be >= 2")

    layers = []

    push!(layers, Dense(d * d * ch, z, relu))
    if reg push!(layers, Dropout(0.5)) end
    for i in 1:(nlayer-2)
        push!(layers, Dense(z, z, relu))
        if reg push!(layers, Dropout(0.5)) end
    end
    push!(layers, Dense(z, d * d))

    Chain(Reshape([d*d*ch, :]),
          layers...,
          Reshape([d,d,:]))
end

struct DimAdd end
function (l::DimAdd)(x)
    reshape(x, size(x)[1:end-1]...,1,size(x)[end])
end
struct DimDrop end
function (l::DimDrop)(x)
    # reshape(x, size(x)[1:end-2]...,size(x)[end])
    dropdims(x, dims=3)
end

function eq_model(; z=300, reg=false, nlayer=3, ch=1)
    nlayer >= 2 || error("nlayer must be >= 2")

    layers = []

    # The ch is either 1 or 2.
    # - ch==1: the input is a matrix, COR, COV, maxCOV, etc
    # - ch==2: the first channel is the COR matrix, and the second channel is the variance in the diagonal.
    push!(layers, Equivariant(ch=>z))
    push!(layers, LeakyReLU())
    if reg push!(layers, Dropout(0.5)) end
    for i in 1:(nlayer-2)
        push!(layers, Equivariant(z=>z))
        push!(layers, LeakyReLU())
        # FIXME whether this is the correct drop
        if reg push!(layers, Dropout(0.5)) end
    end
    # FIXME since all previous layer has LeakyReLU activation, the output is
    # largely positive. This layer does not have ReLU. But is this one alone
    # sufficient for producing many negative values? Should I use multiple
    # layers?
    push!(layers, Equivariant(z=>1))

    Chain(
        # CAUTION I'm removing the DimAdd function, and expect the input to have d*d*ch*batch
#         DimAdd(),
        layers...,
        # IMPORTANT drop the second-to-last dim 1
        DimDrop())
end

# function cnn_model(ch=1)
#     Chain(
# #         DimAdd(),
#           Conv((3,3), ch=>32, relu, pad=(1,1)),
#           BatchNorm(32),
#           # FIXME channel size
#           # FIXME number of layers
#           # FIXME normalization
#           Conv((3,3), 32=>32, relu, pad=(1,1)),
#           BatchNorm(32),
#           Conv((3,3), 32=>32, relu, pad=(1,1)),
#           BatchNorm(32),
#           Conv((3,3), 32=>32, relu, pad=(1,1)),
#           BatchNorm(32),
#           Conv((3,3), 32=>32, relu, pad=(1,1)),
#           BatchNorm(32),
#           # CAUTION no relu in the last layer
#           Conv((3,3), 32=>1, pad=(1,1)),
#           DimDrop())
# end

function cnn_model(ch=1, midch=32, kernel=(3,3), pad=(1,1))
    Chain(
          Conv(kernel, ch=>midch, relu, pad=pad),
          BatchNorm(midch),
          Conv(kernel, midch=>midch, relu, pad=pad),
          BatchNorm(midch),
          Conv(kernel, midch=>midch, relu, pad=pad),
          BatchNorm(midch),
          Conv(kernel, midch=>midch, relu, pad=pad),
          BatchNorm(midch),
          Conv(kernel, midch=>midch, relu, pad=pad),
          BatchNorm(midch),
          # CAUTION no relu in the last layer
          Conv(kernel, midch=>1, pad=pad),
          DimDrop())
end

function test()
    size(cnn_model()(randn(32,32,100)))
    size(cnn_model()(randn(16,16,100)))
    size(cnn_model()(randn(8,8,100)))
    size(cnn_model()(randn(9,9,100)))
end

fc_model_fn(d) = fc_model(d=d, z=1024, nlayer=3)
fc_dropout_model_fn(d) = fc_model(d=d, z=1024, reg=true, nlayer=3)

deep_fc_model_fn(d) = fc_model(d=d, z=1024, nlayer=6)
deep_fc_dropout_model_fn(d) = fc_model(d=d, z=1024, reg=true, nlayer=6)

eq_model_fn() = eq_model(z=300, reg=false, nlayer=3)
eq_dropout_model_fn() = eq_model(z=300, reg=true, nlayer=3)

# TODO wide model
deep_eq_model_fn() = eq_model(z=300, reg=false, nlayer=6)
deep_eq_dropout_model_fn() = eq_model(z=300, reg=true, nlayer=6)

# CAUTION ch=2!!!
eq2_model_fn() = eq_model(z=300, reg=false, nlayer=6, ch=2)

max_stars_repo_name: lihebi/DAG-EQ

max_issues_repo_name: lihebi/DAG-EQ

max_forks_repo_name: lihebi/DAG-EQ


text:
# -*- encoding: utf-8 -*-
#
# The MIT License (MIT)
#
# Copyright © 2021 Matteo Foglieni and Riccardo Gervasoni
#

@testset "test_Hit" begin
     sphere = Sphere()
     ray1 = Ray(Point(0, 0, 2), -VEC_Z)

     intersection1 = ray_intersection(sphere, ray1)
     @test typeof(intersection1) == HitRecord
     @test HitRecord(
          Point(0.0, 0.0, 1.0),
          Normal(0.0, 0.0, 1.0),
          Vec2d(0.0, 0.0),
          1.0,
          ray1
          ) ≈ intersection1

     ray2 = Ray( Point(3, 0, 0), -VEC_X )
     intersection2 = ray_intersection(sphere, ray2)
     @test typeof(intersection2) == HitRecord
     @test HitRecord(
          Point(1.0, 0.0, 0.0),
          Normal(1.0, 0.0, 0.0),
          Vec2d(0.0, 0.5),
          2.0,
          ray2
          ) ≈ intersection2

     @test typeof( ray_intersection(sphere, Ray(Point(0, 10, 2), -VEC_Z ) ) ) == Nothing
end

@testset "test_InnerHit" begin
     sphere = Sphere()

     ray = Ray(Point(0, 0, 0), VEC_X)
     intersection = ray_intersection(sphere, ray)

     @test typeof(intersection) == HitRecord
     @test HitRecord(
          Point(1.0, 0.0, 0.0),
          Normal(-1.0, 0.0, 0.0),
          Vec2d(0.0, 0.5),
          1.0,
          ray
          ) ≈ intersection
end

@testset "test_Transformation" begin
     sphere = Sphere(translation(Vec(10.0, 0.0, 0.0)))

     ray1 = Ray(Point(10, 0, 2), -VEC_Z)
     intersection1 = ray_intersection(sphere, ray1)
     @test typeof(intersection1) == HitRecord
     @test HitRecord(
          Point(10.0, 0.0, 1.0),
          Normal(0.0, 0.0, 1.0),
          Vec2d(0.0, 0.0),
          1.0,
          ray1
          ) ≈ intersection1

     ray2 = Ray(Point(13, 0, 0), -VEC_X)
     intersection2 = ray_intersection(sphere, ray2)
     @test typeof(intersection2) == HitRecord
     @test HitRecord(
          Point(11.0, 0.0, 0.0),
          Normal(1.0, 0.0, 0.0),
          Vec2d(0.0, 0.5),
          2.0,
          ray2
          ) ≈ intersection2

     # Check if the sphere failed to move by trying to hit the untransformed shape
     @test typeof( ray_intersection(sphere, Ray( Point(0, 0, 2), -VEC_Z ) ) ) == Nothing

     # Check if the *inverse* transformation was wrongly applied
     @test typeof( ray_intersection(sphere, Ray( Point(-10, 0, 0), -VEC_Z ) ) ) == Nothing
end

@testset "test_Normals" begin
     sphere = Sphere(scaling(Vec(2.0, 1.0, 1.0)))
     ray = Ray(Point(1.0, 1.0, 0.0), Vec(-1.0, -1.0, 0.0))
     intersection = ray_intersection(sphere, ray)

     @test intersection.normal ≈ Normal(1.0, 4.0, 0.0)
end

@testset "test_Normal_direction" begin
     # Scaling a sphere by -1 keeps the sphere the same but reverses its
     # reference frame
     sphere = Sphere(scaling(Vec(-1.0, -1.0, -1.0)))

     ray = Ray(Point(0.0, 2.0, 0.0), -VEC_Y)
     intersection = ray_intersection(sphere, ray)

     @test intersection.normal ≈ Normal(0.0, 1.0, 0.0)
end

@testset "test_UV_Coordinates" begin
     sphere = Sphere()

     # The first four rays hit the unit sphere at the
     # points P1, P2, P3, and P4.
     #
     #                    ^ y
     #                    | P2
     #              , - ~ * ~ - ,
     #          , '       |       ' ,
     #        ,           |           ,
     #       ,            |            ,
     #      ,             |             , P1
     # -----*-------------+-------------*---------> x
     #   P3 ,             |             ,
     #       ,            |            ,
     #        ,           |           ,
     #          ,         |        , '
     #            ' - , _ * _ ,  '
     #                    | P4
     #
     # P5 and P6 are aligned along the x axis and are displaced
     # along z (ray5 in the positive direction, ray6 in the negative
     # direction).

     ray1 = Ray(Point(2.0, 0.0, 0.0), -VEC_X)
     @test ray_intersection(sphere, ray1).surface_point ≈ Vec2d(0.0, 0.5)

     ray2 = Ray(Point(0.0, 2.0, 0.0), -VEC_Y)
     @test ray_intersection(sphere, ray2).surface_point ≈ Vec2d(0.25, 0.5)

     ray3 = Ray(Point(-2.0, 0.0, 0.0), VEC_X)
     @test ray_intersection(sphere, ray3).surface_point ≈ Vec2d(0.5, 0.5)

     ray4 = Ray(Point(0.0, -2.0, 0.0), VEC_Y)
     @test ray_intersection(sphere, ray4).surface_point ≈ Vec2d(0.75, 0.5)

     ray5 = Ray(Point(2.0, 0.0, 0.5), -VEC_X)
     @test ray_intersection(sphere, ray5).surface_point ≈ Vec2d(0.0, 1/3)

     ray6 = Ray(Point(2.0, 0.0, -0.5), -VEC_X)
     @test ray_intersection(sphere, ray6).surface_point ≈ Vec2d(0.0, 2/3)
end


text:
#=
This file is auto-generated. Do not edit.
=#

#! format: off

"""
    mutable struct FiveMassShaft <: Shaft
        H::Float64
        H_hp::Float64
        H_ip::Float64
        H_lp::Float64
        H_ex::Float64
        D::Float64
        D_hp::Float64
        D_ip::Float64
        D_lp::Float64
        D_ex::Float64
        D_12::Float64
        D_23::Float64
        D_34::Float64
        D_45::Float64
        K_hp::Float64
        K_ip::Float64
        K_lp::Float64
        K_ex::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
        internal::InfrastructureSystemsInternal
    end

Parameters of 5 mass-spring shaft model.
 It contains a High-Pressure (HP) steam turbine, Intermediate-Pressure (IP)
 steam turbine, Low-Pressure (LP) steam turbine, the Rotor and an Exciter (EX) mover.

# Arguments
- `H::Float64`: Rotor inertia constant in MWs/MVA, validation range: `(0, nothing)`
- `H_hp::Float64`: High pressure turbine inertia constant in MWs/MVA, validation range: `(0, nothing)`
- `H_ip::Float64`: Intermediate pressure turbine inertia constant in MWs/MVA, validation range: `(0, nothing)`
- `H_lp::Float64`: Low pressure turbine inertia constant in MWs/MVA, validation range: `(0, nothing)`
- `H_ex::Float64`:  Exciter inertia constant in MWs/MVA, validation range: `(0, nothing)`
- `D::Float64`: Rotor natural damping in pu, validation range: `(0, nothing)`
- `D_hp::Float64`: High pressure turbine natural damping in pu, validation range: `(0, nothing)`
- `D_ip::Float64`: Intermediate pressure turbine natural damping in pu, validation range: `(0, nothing)`
- `D_lp::Float64`: Low pressure turbine natural damping in pu, validation range: `(0, nothing)`
- `D_ex::Float64`: Exciter natural damping in pu, validation range: `(0, nothing)`
- `D_12::Float64`: High-Intermediate pressure turbine damping, validation range: `(0, nothing)`
- `D_23::Float64`: Intermediate-Low pressure turbine damping, validation range: `(0, nothing)`
- `D_34::Float64`: Low pressure turbine-Rotor damping, validation range: `(0, nothing)`
- `D_45::Float64`: Rotor-Exciter damping, validation range: `(0, nothing)`
- `K_hp::Float64`: High pressure turbine angle coefficient, validation range: `(0, nothing)`
- `K_ip::Float64`: Intermediate pressure turbine angle coefficient, validation range: `(0, nothing)`
- `K_lp::Float64`: Low pressure turbine angle coefficient, validation range: `(0, nothing)`
- `K_ex::Float64`: Exciter angle coefficient, validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: The states are:
	δ: rotor angle,
	ω: rotor speed,
	δ_hp: rotor angle of high pressure turbine,
	ω_hp: rotor speed of high pressure turbine,
	δ_ip: rotor angle of intermediate pressure turbine,
	ω_ip: rotor speed of intermediate pressure turbine,
	δ_lp: rotor angle of low pressure turbine,
	ω_lp: rotor speed of low pressure turbine,
	δ_ex: rotor angle of exciter,
	ω_lp: rotor speed of exciter
- `n_states::Int`: FiveMassShaft has 10 states
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct FiveMassShaft <: Shaft
    "Rotor inertia constant in MWs/MVA"
    H::Float64
    "High pressure turbine inertia constant in MWs/MVA"
    H_hp::Float64
    "Intermediate pressure turbine inertia constant in MWs/MVA"
    H_ip::Float64
    "Low pressure turbine inertia constant in MWs/MVA"
    H_lp::Float64
    " Exciter inertia constant in MWs/MVA"
    H_ex::Float64
    "Rotor natural damping in pu"
    D::Float64
    "High pressure turbine natural damping in pu"
    D_hp::Float64
    "Intermediate pressure turbine natural damping in pu"
    D_ip::Float64
    "Low pressure turbine natural damping in pu"
    D_lp::Float64
    "Exciter natural damping in pu"
    D_ex::Float64
    "High-Intermediate pressure turbine damping"
    D_12::Float64
    "Intermediate-Low pressure turbine damping"
    D_23::Float64
    "Low pressure turbine-Rotor damping"
    D_34::Float64
    "Rotor-Exciter damping"
    D_45::Float64
    "High pressure turbine angle coefficient"
    K_hp::Float64
    "Intermediate pressure turbine angle coefficient"
    K_ip::Float64
    "Low pressure turbine angle coefficient"
    K_lp::Float64
    "Exciter angle coefficient"
    K_ex::Float64
    ext::Dict{String, Any}
    "The states are:
	δ: rotor angle,
	ω: rotor speed,
	δ_hp: rotor angle of high pressure turbine,
	ω_hp: rotor speed of high pressure turbine,
	δ_ip: rotor angle of intermediate pressure turbine,
	ω_ip: rotor speed of intermediate pressure turbine,
	δ_lp: rotor angle of low pressure turbine,
	ω_lp: rotor speed of low pressure turbine,
	δ_ex: rotor angle of exciter,
	ω_lp: rotor speed of exciter"
    states::Vector{Symbol}
    "FiveMassShaft has 10 states"
    n_states::Int
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function FiveMassShaft(H, H_hp, H_ip, H_lp, H_ex, D, D_hp, D_ip, D_lp, D_ex, D_12, D_23, D_34, D_45, K_hp, K_ip, K_lp, K_ex, ext=Dict{String, Any}(), )
    FiveMassShaft(H, H_hp, H_ip, H_lp, H_ex, D, D_hp, D_ip, D_lp, D_ex, D_12, D_23, D_34, D_45, K_hp, K_ip, K_lp, K_ex, ext, [:δ, :ω, :δ_hp, :ω_hp, :δ_ip, :ω_ip, :δ_lp, :ω_lp, :δ_ex, :ω_ex], 10, InfrastructureSystemsInternal(), )
end

function FiveMassShaft(; H, H_hp, H_ip, H_lp, H_ex, D, D_hp, D_ip, D_lp, D_ex, D_12, D_23, D_34, D_45, K_hp, K_ip, K_lp, K_ex, ext=Dict{String, Any}(), states=[:δ, :ω, :δ_hp, :ω_hp, :δ_ip, :ω_ip, :δ_lp, :ω_lp, :δ_ex, :ω_ex], n_states=10, internal=InfrastructureSystemsInternal(), )
    FiveMassShaft(H, H_hp, H_ip, H_lp, H_ex, D, D_hp, D_ip, D_lp, D_ex, D_12, D_23, D_34, D_45, K_hp, K_ip, K_lp, K_ex, ext, states, n_states, internal, )
end

# Constructor for demo purposes; non-functional.
function FiveMassShaft(::Nothing)
    FiveMassShaft(;
        H=0,
        H_hp=0,
        H_ip=0,
        H_lp=0,
        H_ex=0,
        D=0,
        D_hp=0,
        D_ip=0,
        D_lp=0,
        D_ex=0,
        D_12=0,
        D_23=0,
        D_34=0,
        D_45=0,
        K_hp=0,
        K_ip=0,
        K_lp=0,
        K_ex=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`FiveMassShaft`](@ref) `H`."""
get_H(value::FiveMassShaft) = value.H
"""Get [`FiveMassShaft`](@ref) `H_hp`."""
get_H_hp(value::FiveMassShaft) = value.H_hp
"""Get [`FiveMassShaft`](@ref) `H_ip`."""
get_H_ip(value::FiveMassShaft) = value.H_ip
"""Get [`FiveMassShaft`](@ref) `H_lp`."""
get_H_lp(value::FiveMassShaft) = value.H_lp
"""Get [`FiveMassShaft`](@ref) `H_ex`."""
get_H_ex(value::FiveMassShaft) = value.H_ex
"""Get [`FiveMassShaft`](@ref) `D`."""
get_D(value::FiveMassShaft) = value.D
"""Get [`FiveMassShaft`](@ref) `D_hp`."""
get_D_hp(value::FiveMassShaft) = value.D_hp
"""Get [`FiveMassShaft`](@ref) `D_ip`."""
get_D_ip(value::FiveMassShaft) = value.D_ip
"""Get [`FiveMassShaft`](@ref) `D_lp`."""
get_D_lp(value::FiveMassShaft) = value.D_lp
"""Get [`FiveMassShaft`](@ref) `D_ex`."""
get_D_ex(value::FiveMassShaft) = value.D_ex
"""Get [`FiveMassShaft`](@ref) `D_12`."""
get_D_12(value::FiveMassShaft) = value.D_12
"""Get [`FiveMassShaft`](@ref) `D_23`."""
get_D_23(value::FiveMassShaft) = value.D_23
"""Get [`FiveMassShaft`](@ref) `D_34`."""
get_D_34(value::FiveMassShaft) = value.D_34
"""Get [`FiveMassShaft`](@ref) `D_45`."""
get_D_45(value::FiveMassShaft) = value.D_45
"""Get [`FiveMassShaft`](@ref) `K_hp`."""
get_K_hp(value::FiveMassShaft) = value.K_hp
"""Get [`FiveMassShaft`](@ref) `K_ip`."""
get_K_ip(value::FiveMassShaft) = value.K_ip
"""Get [`FiveMassShaft`](@ref) `K_lp`."""
get_K_lp(value::FiveMassShaft) = value.K_lp
"""Get [`FiveMassShaft`](@ref) `K_ex`."""
get_K_ex(value::FiveMassShaft) = value.K_ex
"""Get [`FiveMassShaft`](@ref) `ext`."""
get_ext(value::FiveMassShaft) = value.ext
"""Get [`FiveMassShaft`](@ref) `states`."""
get_states(value::FiveMassShaft) = value.states
"""Get [`FiveMassShaft`](@ref) `n_states`."""
get_n_states(value::FiveMassShaft) = value.n_states
"""Get [`FiveMassShaft`](@ref) `internal`."""
get_internal(value::FiveMassShaft) = value.internal

"""Set [`FiveMassShaft`](@ref) `H`."""
set_H!(value::FiveMassShaft, val) = value.H = val
"""Set [`FiveMassShaft`](@ref) `H_hp`."""
set_H_hp!(value::FiveMassShaft, val) = value.H_hp = val
"""Set [`FiveMassShaft`](@ref) `H_ip`."""
set_H_ip!(value::FiveMassShaft, val) = value.H_ip = val
"""Set [`FiveMassShaft`](@ref) `H_lp`."""
set_H_lp!(value::FiveMassShaft, val) = value.H_lp = val
"""Set [`FiveMassShaft`](@ref) `H_ex`."""
set_H_ex!(value::FiveMassShaft, val) = value.H_ex = val
"""Set [`FiveMassShaft`](@ref) `D`."""
set_D!(value::FiveMassShaft, val) = value.D = val
"""Set [`FiveMassShaft`](@ref) `D_hp`."""
set_D_hp!(value::FiveMassShaft, val) = value.D_hp = val
"""Set [`FiveMassShaft`](@ref) `D_ip`."""
set_D_ip!(value::FiveMassShaft, val) = value.D_ip = val
"""Set [`FiveMassShaft`](@ref) `D_lp`."""
set_D_lp!(value::FiveMassShaft, val) = value.D_lp = val
"""Set [`FiveMassShaft`](@ref) `D_ex`."""
set_D_ex!(value::FiveMassShaft, val) = value.D_ex = val
"""Set [`FiveMassShaft`](@ref) `D_12`."""
set_D_12!(value::FiveMassShaft, val) = value.D_12 = val
"""Set [`FiveMassShaft`](@ref) `D_23`."""
set_D_23!(value::FiveMassShaft, val) = value.D_23 = val
"""Set [`FiveMassShaft`](@ref) `D_34`."""
set_D_34!(value::FiveMassShaft, val) = value.D_34 = val
"""Set [`FiveMassShaft`](@ref) `D_45`."""
set_D_45!(value::FiveMassShaft, val) = value.D_45 = val
"""Set [`FiveMassShaft`](@ref) `K_hp`."""
set_K_hp!(value::FiveMassShaft, val) = value.K_hp = val
"""Set [`FiveMassShaft`](@ref) `K_ip`."""
set_K_ip!(value::FiveMassShaft, val) = value.K_ip = val
"""Set [`FiveMassShaft`](@ref) `K_lp`."""
set_K_lp!(value::FiveMassShaft, val) = value.K_lp = val
"""Set [`FiveMassShaft`](@ref) `K_ex`."""
set_K_ex!(value::FiveMassShaft, val) = value.K_ex = val
"""Set [`FiveMassShaft`](@ref) `ext`."""
set_ext!(value::FiveMassShaft, val) = value.ext = val


text:
module IOStrings

using Base: @propagate_inbounds

"""
`IOString` is an abstract string that accumulates bytes from an IO stream.

`IOString(::IO)` creates a new `IOString` connected to an `IO` stream.

`pump(::IOString)` reads bytes from the `IO` stream into the string.

`incomplete(::IOString)` is true until the `IO` stream reaches `eof()`.

`pump(::Function, ::IOString)` runs a `Function` until it returns or throws
a recoverable error. When a recoverable error is thrown, more data is read
from the `IO` stream and the `Function` is restarted (the function is assumed
to be idempotent or free of side effects).

`recoverable(e) -> Bool` defines which errors are recoverable.

It the following example, `parse_item` is called repeatedly to parse items
from the string `s` at index `i`. When the parser gets to the end of the string
and throws an error, the `IOString` is pumped to retrieve more data and the
parsing continues.
```
IOString.recoverable(::MyUnexpectedEndOfInput) = true
s = IOString(socket)
i = 1
while !eof(socket)
    i, v = pump(()->MyParser.parse_item(s, i), s)
    println(v)
end
```
"""
struct IOString{T <: IO} <: AbstractString
    io::T
    buf::IOBuffer
end

struct IncompleteError <: Exception
end

const ASCII_ETB = 0x17

function IOString(io::T) where T <: IO
    ios = IOString{T}(io, IOBuffer())
    Base.ensureroom(ios.buf, 1)
    ios.buf.data[1] = ASCII_ETB
    @assert incomplete(ios)
    return ios
end

Base.convert(::Type{Base.String}, s::IOString) =
    unsafe_string(pointer(s), ncodeunits(s))

Base.String(s::IOString) = convert(Base.String, s)

Base.IteratorSize(::Type{IOString}) = Base.SizeUnknown()
Base.ncodeunits(s::IOString) = s.buf.size
Base.codeunit(s::IOString) = UInt8
@propagate_inbounds Base.codeunit(s::IOString, i::Integer) = s.buf.data[i]

Base.checkbounds(s::IOString, I::Union{Integer,AbstractArray}) =
    checkbounds(Bool, s, I) ? nothing :
              incomplete(s) ? throw(IncompleteError()) :
                              throw(BoundsError(s, I))

Base.pointer(s::IOString) = pointer(s.buf.data)
Base.pointer(s::IOString, i) = pointer(s.buf.data, i)

incomplete(s::IOString) = codeunit(s, ncodeunits(s) + 1) == ASCII_ETB

recoverable(e) = false
recoverable(::IncompleteError) = true

@inline function pump(f::Function, s::IOString)
    while true
        try
            return f()
        catch e
            if recoverable(e)
                pump(s)
            else
                rethrow(e)
            end
        end
    end
end

function pump(s::IOString)
    if eof(s.io)
        Base.ensureroom(s.buf, 1)
        s.buf.data[s.buf.size + 1] = 0x00
        @assert !incomplete(s)
    else
        write(s.buf, readavailable(s.io))
        Base.ensureroom(s.buf, 1)
        s.buf.data[s.buf.size + 1] = ASCII_ETB
        @assert incomplete(s)
    end
end

end # module IOString


text:
import Base: .+, .-, .*, ./, .^, max, min, .==, .>, .>=, .<, .<=, +, -

# Broadcasting binary functions (uses the same list as same size arrays)
cuda12 = cuda11

# Broadcast max/min haven't been defined:
max(a::Array,b::Array)=broadcast(max,a,b)
min(a::Array,b::Array)=broadcast(min,a,b)

function vbroadcast_shape(x,y)
    nz = max(ndims(x),ndims(y))
    dz = ones(Int,nz)
    xdims = ydims = xsame = ysame = xlast = ylast = 0; zlen = 1
    for i=1:nz
        if size(x,i) > 1
            xdims += 1; xlast = i
            dz[i] = size(x,i)
        end
        if size(y,i) > 1
            ydims += 1; ylast = i
            if dz[i] == 1
                dz[i] = size(y,i)
            else
                dz[i] == size(y,i) || throw(DimensionMismatch("arrays could not be broadcast to a common size"))
            end
        end
        xsame += (dz[i] == size(x,i))
        ysame += (dz[i] == size(y,i))
        zlen *= dz[i]
    end
    xsame == nz || xdims <= 1 || error("Only vector broadcasting supported")
    ysame == nz || ydims <= 1 || error("Only vector broadcasting supported")
    if xdims == 0
        sx = zlen; nx = 1
    elseif xdims == 1
        sx = prod(dz[1:xlast-1]); nx=dz[xlast]
    elseif xsame == nz
        sx = 1; nx=zlen
    else
        error("Broadcasting error")
    end
    if ydims == 0
        sy = zlen; ny = 1
    elseif ydims == 1
        sy = prod(dz[1:ylast-1]); ny=dz[ylast]
    elseif ysame == nz
        sy = 1; ny=zlen
    else
        error("Broadcasting error")
    end
    return (tuple(dz...), sx, nx, sy, ny)
end

function cuda12def(f, j=f, o...)
    J=Symbol(j)
    for S in (32,64)
        T = Symbol("Float$S")
        F11 = "$(f)_$(S)_11"
        F12 = "$(f)_$(S)_12"
        @eval begin
            function $J(x::KnetArray{$T},y::KnetArray{$T})
                if size(x)==size(y)
                    z = similar(x)
                    ccall(($F11,$libknet8),Void,(Cint,$Ptr{$T},Ptr{$T},Ptr{$T}),length(z),x,y,z)
                    return z
                else
                    (dz,sx,nx,sy,ny) = vbroadcast_shape(x,y)
                    z = similar(x,dz)
                    ccall(($F12,$libknet8),Void,(Cint,$Ptr{$T},Cint,Cint,Ptr{$T},Cint,Cint,Ptr{$T}),length(z),x,sx,nx,y,sy,ny,z)
                    return z
                end
            end
        end
    end
end

#if isdefined(:libknet8)
    for f in cuda12
        isa(f,Tuple) || (f=(f,))
        cuda12def(f...)
    end
#end


text:
module TestMatrixComps
using CustomTest
using Base.Test
using ControlSystems

A = [-0.21 0.2; 0.2 -0.21]
B = 0.01*eye(2)
C = eye(2)
D = 0
sys = ss(A,B,C,D)
sysr, G = balreal(sys)

@test_approx_eq gram(sysr, :c) G
@test_approx_eq gram(sysr, :o) G
@test_approx_eq sort(pole(sysr)) sort(pole(sys))

sysb,T = ControlSystems.balance_statespace(sys)
Ab,Bb,Cb,T = ControlSystems.balance_statespace(A,B,C)

@test_approx_eq sysb.A Ab

@test_approx_eq ControlSystems.balance_transform(A,B,C) ControlSystems.balance_transform(sys)

W = eye(2)
@test_approx_eq covar(sys, W) [0.002560975609756 0.002439024390244; 0.002439024390244 0.002560975609756]
D2 = eye(2)
@test_approx_eq covar(ss(A,B,C,D2, 1), W) [1.000110108378310 -0.000010098377310; -0.000010098377310 1.000110108378310]
# Direct term means infinite covariance
@test_approx_eq covar(ss(A,B,C,D2), W) [Inf Inf; Inf Inf]

# No noise on second output should give finite variance
@test_approx_eq covar(ss(A,B,C,[1 0; 0 0]), W) [Inf Inf; Inf 0.002560975609756]

# Unstable system has inf covar
@test covar(ss(eye(2),B,C,0), W) == [Inf Inf; Inf Inf]

# Discrete system can have direct term
@test_approx_eq covar(ss(A,B,C,D2,0.1),W) [1.00011010837831 -1.0098377309782909e-5; -1.0098377309782909e-5 1.00011010837831]

end


text:
using ThreeBodyDecay

const mΛb = 5.61960
const mπ = 0.14
# intermediate resonances
const mΣb = 5.81065; const ΓΣb = 0.005
const mΣb_x = 5.83032; const ΓΣb_x = 0.0094

const mΛb2S = 6.3; # just a peak of the plot

tbs = ThreeBodySystem(mπ,mΛb,mπ; m0=mΛb2S,
    two_jps = ([0, 1, 0, 1], ['-','+','-','+']))  # 1/2+ 0- 0- 1/2+
dpp = randomPoint(tbs)

# lineshape
dc_Σb1 = DecayChainLS(1, σ->BW(σ,mΣb,ΓΣb); tbs=tbs, two_s = 1, parity='+')
dc_Σb3 = DecayChainLS(3, σ->BW(σ,mΣb,ΓΣb); tbs=tbs, two_s = 1, parity='+')

dc_Σb_x1 = DecayChainLS(1, σ->BW(σ,mΣb_x,ΓΣb_x); tbs=tbs, two_s = 3, parity='+')
dc_Σb_x3 = DecayChainLS(3, σ->BW(σ,mΣb_x,ΓΣb_x); tbs=tbs, two_s = 3, parity='+')

dc_f0 = DecayChainLS(2, σ->1.0; tbs=tbs, two_s = 0, parity='+')

full_a(σs,two_λs) = sum(amplitude(σs,two_λs, ch) for ch in [dc_Σb1, dc_Σb_x1, dc_Σb3, dc_Σb_x3])
full_a(dpp) = full_a(dpp.σs, dpp.two_λs)
#
total_I(σs) = sum(abs2(full_a(σs,two_λs)) for two_λs in Iterators.product(
    -tbs.two_js[1]:2:tbs.two_js[1],
    -tbs.two_js[2]:2:tbs.two_js[2],
    -tbs.two_js[3]:2:tbs.two_js[3],
    -tbs.two_js[4]:2:tbs.two_js[4]))

σ3v,σ1v = flatDalitzPlotSample31(tbs)
σ2v = [gσ2(σ3,σ1,tbs.msq) for (σ3,σ1) in zip(σ3v,σ1v)]

cala = [full_a(DalitzPlotPoint31(σ3,σ1,tbs)) for (σ3,σ1) in zip(σ3v,σ1v)]
cali = abs2.(cala)

using Plots
histogram2d(σ3v, σ2v, bins=50)
histogram2d(σ3v, σ2v, weights=cali, bins=50)

let
    σ3v = range(tbs.mthsq[3],tbs.sthsq[3], length=150)
    σ2v = range(tbs.mthsq[2],tbs.sthsq[2], length=150)
    cali = [let dpp = DalitzPlotPoint23(σ2,σ3,tbs)
        Kibble(dpp.σs, tbs.msq) < 0 ? total_I(dpp.σs) : NaN
        end for (σ3,σ2) in Iterators.product(σ3v,σ2v)]
    heatmap(σ2v,σ3v,cali)
end

max_stars_repo_name: mmikhasenko/ThreeBodyDecay

max_issues_repo_name: mmikhasenko/ThreeBodyDecay

max_forks_repo_name: mmikhasenko/ThreeBodyDecay


text:
"""
    initial_population(individues::Integer, lower_limits::AbstractArray, upper_limits::AbstractArray)
Return a `m*n` matrix which `m` is the number of `individues` and `n` the number of parameters. Each individue (row) has `n` parameters
with upper and lower limits given by `lower_limits, upper_limits` arrays.

# Examples
```julia-repl
julia> initial_population(4,[-1,-5],[1,5])
4×3 Array{Float64,2}:
 -0.997462  -3.461379
  0.056423   4.564237
  0.167943  -4.467215
  0.695743   0.561379
```
"""
function initial_population(individues::Integer, lower_limits::AbstractArray, upper_limits::AbstractArray)

    #Error checks
    if length(lower_limits) != length(upper_limits)
        throw(ErrorException("The limits must have the same length: lower and upper limits of length $(length(lower_limits)) and $(length(upper_limits))."))
    end

    return row_mul( (upper_limits .- lower_limits), rand(individues,length(lower_limits)) ) .+ lower_limits'

end


text:
#==============================================================================
This script reconstructs the binary matrix and frequency vector from experimental
measurement data stored in MAT file realdata.mat.

In addition, the dependence between the frequency vector and the reconstruction
error is examined with synthetic data which is generated with the given binary
matrix, different frequency vectors, and additive Gaussian noise.

Parallelization can be enabled, e.g., by starting julia with "julia -p 16".
Additional packages may be needed to install for the visualization.

This is the set of experiments presented (with nT=200 and higher resolution) in
Section 4.4 of the article:

L. Mustonen, X. Gao, A. Santana, R. Mitchell, Y. Vigfusson, and L. Ruthotto,
A Bayesian framework for molecular strain identification from mixed diagnostic samples
https://arxiv.org/abs/1803.02916
==============================================================================#

@everywhere using StrainRecon
using PyPlot
using jInv.Utils
using jInv.Mesh
using jInvVis
using MAT
srand(2) # for reproducibility

# load data from MAT file
data = matread("realdata.mat")
M_true = data["M_true"]
w_true = vec(data["w_true"])
d_meas = vec(data["d_meas"])

m = length(d_meas) # number of measurements
n = 3 # number of strains in the reconstruction
R = 10 # variance of the measurement noise is assumed to be 0.01
param = getStrainReconParam(d_meas,n,R=R)
nT = 200 # number of trials in BCD

# compute sample standard deviation of the error
errstd = sqrt(var(M_true*w_true-d_meas))
println("Sample standard deviation of the measurement error is $errstd")

# compute MAP estimates
MrecoBCD,wrecoBCD = solveMAP(param,maxTry=nT)
MrecoMIQP,wrecoMIQP = solveConvexMAP(param)

# compute posterior means and variances
Mcm,wcm,dummy,Mvar,wvar = wrapperIntegratePosterior(param)

# show results
imcolor = "gray"
imlimits = [.5, m+.5, n+.5, .5]

subplot(5,1,1)
imshow(M_true',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("true M")
println("true w:")
println(w_true)
println()

subplot(5,1,2)
imshow(MrecoBCD',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("M from BCD")
println("w from BCD:")
println(wrecoBCD)
println()

subplot(5,1,3)
imshow(MrecoMIQP',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("M from MIQP")
println("w from MIQP:")
println(wrecoMIQP)
println()

subplot(5,1,4)
imshow(Mcm',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("posterior mean of M")
println("posterior mean of w:")
println(wcm)
println()

subplot(5,1,5)
imshow(sqrt.(Mvar)',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("posterior standard deviation of M")
println("posterior standard deviation of w:")
println(sqrt.(wvar))
println()

# keep M fixed, vary w and reconstruct from synthetic, noisy measurements
noiseVar_range = [1e-2, 1e-4]
nT = 20 # decrease the number of trials for efficiency
resolution = 64
Ma = getRegularMesh([1/3 1 0 1/2 ],[resolution resolution])
as = getCellCenteredGrid(Ma)
recoerrs = SharedArray{Float64}(Ma.nc,2)
recoerrs[:] = NaN # white background

for noiseVar_idx = 1:2
    @sync @parallel for k=1:size(as,1)
        at = vec([ vec(as[k,:]); 1-sum(as[k,:])])
        if all(at.>=0) && sum(at)<=1 && all(diff(at).<=0)
            d = M_true * at + sqrt(noiseVar_range[noiseVar_idx])*randn(m)
            param = getStrainReconParam(d,n)
            Mreco,wreco = solveMAP(param,maxTry=nT)
            recoerrs[k,noiseVar_idx] = srDistance(BitArray(M_true),at,BitArray(Mreco),wreco,pn=1.0)
        end
    end
end

figure()
subplot(1,2,1)
viewImage2D(recoerrs[:,1],Ma)
xlabel("w1")
ylabel("w2")
axis("image")
title("reconstruction error, noise $(sqrt(noiseVar_range[1]))")
colorbar()
set_cmap("Spectral")
clim(0,17)

subplot(1,2,2)
viewImage2D(recoerrs[:,2],Ma)
xlabel("w1")
ylabel("w2")
axis("image")
title("reconstruction error, noise $(sqrt(noiseVar_range[2]))")
colorbar()
set_cmap("Spectral")
clim(0,17)


text:
"""
Create linear grayscale gradient image of type `Gray`.
"""
function gradient_image(height, width)
    row = reshape(range(0; stop=1, length=width), 1, width)
    grad = Gray.(vcat(repeat(row, height))) # Linear gradient
    img = srgb2linear.(grad) # For printing, compensate for SRGB colorspace
    return grad, img
end

"""
Test dithering algorithm `alg` on linear grayscale gradient from `gradient_image`
and show stacked plot.
"""
function test_on_gradient(alg::AbstractDither)
    srgb, linear = gradient_image(100, 800)
    dither_srgb = dither(srgb, alg)
    dither_linear = dither(linear, alg)

    return mosaicview([srgb, dither_srgb, linear, dither_linear]; ncol=1)
end


text:
#-----------------------------------------------------------------------# common
abstract type HistogramStat{T} <: OnlineStat{T} end

# Index of `edges` that `y` belongs in, depending on if bins are `left` closed and if
# the bin on the end is `closed` instead of half-open.
function binindex(edges::AbstractVector, y, left::Bool, closed::Bool)
    a, b = extrema(edges)
    y < a && return 0
    y > b && return length(edges)
    closed && y == a && return 1
    closed && y == b && return length(edges) - 1
    if left
        if isa(edges, AbstractRange)
            return floor(Int, (y - a) / step(edges)) + 1
        else
            return searchsortedlast(edges, y)
        end
    else
        if isa(edges, AbstractRange)
            return ceil(Int, (y - a) / step(edges))
        else
            return searchsortedfirst(edges, y) - 1
        end
    end
end

# requires: edges(o), midpoints(o), counts(o)
split_candidates(o::HistogramStat) = midpoints(o)
Statistics.mean(o::HistogramStat) = mean(midpoints(o), fweights(counts(o)))
Statistics.var(o::HistogramStat) = var(midpoints(o), fweights(counts(o)); corrected=true)
Statistics.std(o::HistogramStat) = sqrt(var(o))
Statistics.median(o::HistogramStat) = quantile(o, .5)
function Base.extrema(o::HistogramStat)
    x, y = midpoints(o), counts(o)
    x[findfirst(x -> x > 0, y)], x[findlast(x -> x > 0, y)]
end
function Statistics.quantile(o::HistogramStat, p = [0, .25, .5, .75, 1])
    x, y = midpoints(o), counts(o)
    inds = findall(x -> x != 0, y)
    quantile(x[inds], fweights(y[inds]), p)
end

#-----------------------------------------------------------------------# Hist
"""
    Hist(edges; left = true, closed = true)

Create a histogram with bin partition defined by `edges`.

- If `left`, the bins will be left-closed.
- If `closed`, the bin on the end will be closed.
    - E.g. for a two bin histogram ``[a, b), [b, c)`` vs. ``[a, b), [b, c]``

# Example

    o = fit!(Hist(-5:.1:5), randn(10^6))

    # approximate statistics
    using Statistics

    mean(o)
    var(o)
    std(o)
    quantile(o)
    median(o)
    extrema(o)
"""
struct Hist{T, R} <: HistogramStat{T}
    edges::R
    counts::Vector{Int}
    out::Vector{Int}
    left::Bool
    closed::Bool

    function Hist(edges::R, T::Type = eltype(edges); left::Bool=true, closed::Bool=true) where {R<:AbstractVector}
        new{T,R}(edges, zeros(Int, length(edges) - 1), [0,0], left, closed)
    end
end
nobs(o::Hist) = sum(o.counts) + sum(o.out)
value(o::Hist) = (x=o.edges, y=o.counts)

midpoints(o::Hist) = midpoints(o.edges)
counts(o::Hist) = o.counts
edges(o::Hist) = o.edges

function area(o::Hist)
    c = o.counts
    e = o.edges
    if isa(e, AbstractRange)
        return step(e) * sum(c)
    else
        return sum((e[i+1] - e[i]) * c[i] for i in 1:length(c))
    end
end

function pdf(o::Hist, y)
    i = binindex(o.edges, y, o.left, o.closed)
    if i < 1 || i > length(o.counts)
        return 0.0
    else
        return o.counts[i] / area(o)
    end
end

function _fit!(o::Hist, y)
    i = binindex(o.edges, y, o.left, o.closed)
    if 1 ≤ i < length(o.edges)
        o.counts[i] += 1
    else
        o.out[1 + (i > 0)] += 1
    end
end

function _merge!(o::Hist, o2::Hist)
    if o.edges == o2.edges
        for j in eachindex(o.counts)
            o.counts[j] += o2.counts[j]
        end
    else
        @warn("Histogram edges are not aligned.  Merging is approximate.")
        for (yi, wi) in zip(midpoints(o2.edges), o2.counts)
            for k in 1:wi
                _fit!(o, yi)
            end
        end
    end
    for j in eachindex(o.out)
        o.out[j] += o2.out[j]
    end
end

#-----------------------------------------------------------------------# ExpandingHist
const ExpandableRange = Union{StepRange, StepRangeLen, LinRange}

"""
    ExpandingHist(nbins)

An adaptive histogram where the bin edges keep doubling in size in order to contain every observation.
`nbins` must be an even number.  Bins are left-closed and the rightmost bin is closed, e.g.

- `[a, b)`, `[b, c)`, `[c, d]`

# Example
    o = fit!(ExpandingHist(200), randn(10^6))

    using Plots
    plot(o)

# Details

How `ExpandingHist` works is best understood through example.  Suppose we start with a histogram
of edges/counts as follows:

```
|1|2|5|3|2|
```

- Now we observe a data point that is not contained in the bin edges:

```
|1|2|5|3|2|       *
```

- In order to contain the point, the range of the edges doubles in the direction of the new
  data point and adjacent bins merge their counts:

```
|1|2|5|3|2|       *
 \\ / \\ / \\ /      ↓
  ↓   ↓   ↓       ↓
| 3 | 8 | 2 | 0 | 1 |
```

- Note that multiple iterations of bin-doubling may occur until the new point is contained by the
  bin edges.
"""
mutable struct ExpandingHist{T, R <: StepRangeLen} <: HistogramStat{T}
    edges::R
    counts::Vector{Int}
    n::Int
    function ExpandingHist(init::R, T::Type=Number) where {R <: ExpandableRange}
        new{T, R}(init, zeros(Int, length(init) - 1), 0)
    end
end
function ExpandingHist(b::Int)
    @assert iseven(b)
    ExpandingHist(range(0, stop=0, length = b + 1), Number)
end
value(o::ExpandingHist) = (x=o.edges, y=o.counts)

Base.show(io::IO, o::ExpandingHist) = print(io, "ExpandingHist | n = $(nobs(o)) | nbins = $(length(o.counts))")

midpoints(o::ExpandingHist) = midpoints(o.edges)
counts(o::ExpandingHist) = o.counts
edges(o::ExpandingHist) = o.edges

function _fit!(o::ExpandingHist, y)
    isfinite(y) || error("ExpandingHist requires finite values.  Found $y.")
    o.n += 1

    # init
    @inbounds if nobs(o) == 1
        o.edges = range(y, stop=y, length=length(o.edges))
        o.counts[1] = 1
    elseif step(o.edges) == 0.0
        a = o.edges[1]
        if y == a
            o.counts[1] += 1
        elseif y < a
            o.edges = range(y, stop=a, length=length(o.edges))
            o.counts[end] = o.counts[1]
            o.counts[1] = 1
        else
            o.edges = range(a, stop=y, length=length(o.edges))
            o.counts[end] = 1
        end
    else
        expand!(o, y)
        i = binindex(o.edges, y, true, true)
        o.counts[i] += 1
    end
end

function expand!(o::ExpandingHist, y)
    a, b = extrema(o.edges)
    w = b - a
    w = w == 0.0 ? y - a : w
    halfnbin = round(Int, length(o.counts) / 2)
    @inbounds if y > b  # find K such that y <= a + 2^K * w
        K = ceil(Int, log2((y - a) / w))
        C = 2 ^ K
        o.edges = range(a, stop = a + C*w, length=length(o.edges))
        for _ in 1:min(K, halfnbin)
            for i in eachindex(o.counts)
                o.counts[i] = i ≤ halfnbin ? sum(view(o.counts, (2i-1):(2i))) : 0
            end
        end
    elseif y < a # find K such that y >= b - 2^K * w
        K = ceil(Int, log2((b - y) / w))
        C = 2 ^ K
        o.edges = range(b - C*w, stop = b, length=length(o.edges))
        for _ in 1:min(K, halfnbin)
            n = length(o.counts)
            for i in 0:(length(o.counts) - 1)
                rng = (n-2i-1):(n-2i)
                o.counts[end-i] = i+1 ≤ halfnbin ? sum(view(o.counts, rng)) : 0
            end
        end
    end
end


text:
using QOptCtrl

X = pauli_X
Y = pauli_Y
Z = pauli_Z
H = 1/sqrt(2)*[1 1; 1 -1] # Hadamard Gate

ω = 1

prob = GateTransfer(
        (ω/2)*Z, # drift Hamiltonian
        [X, Y], # control Hamiltonians
        20, # number of steps
        4, # transfer time
        H, # tarfet gate
        rand(20, 2) #initial guess
)

init_ctrl_arr,  final_ctrl_arr = OPTIMIZE(prob)
visualise_pulses(init_ctrl_arr, final_ctrl_arr)


text:
# ---
# title: Template Matching
# cover: assets/template_matching.png
# author: Michael Pusterhofer
# date: 2021-06-22
# ---

# This demo shows how to find objects in an image using template matching.

# The main idea is to check the similarity between a search target(template) and a subsection of the image. The subsection is usually the same size as the template. Therefore by using [`mapwindow`](@ref) we can assign a value for every subsection of the image.

# At first we import the following packages.
using ImageCore: Gray
using ImageMorphology: label_components, component_centroids
using ImageFiltering: mapwindow, Fill, imfilter, KernelFactors
using ImageDistances: sqeuclidean
using ImageContrastAdjustment: adjust_histogram, LinearStretching
using TestImages
using Plots: scatter!, plot

# `ImageCore` enables the generation of images, `ImageFiltering` provides the [`mapwindow`](@ref) function and `ImageFeatures` provides functions to label segments of an image. To calculate the similarity `squeclidean` is used. This defines a distance between two images as:
#
# ```math
# \textrm{distance} = \sum_i^{\textrm{number of pixels}} (\textrm{template}_i - \textrm{subsection}_i)^2
# ```
#
# `ImageContrastAdjustment` provides functions to adjust the histogram, which is useful when an image contains values bigger than 1 or smaller than 0.  The `Testimages` package will provide our image and plot is used to overlay a scatter plot onto the image.

# To start we first load our image.
img = testimage("moonsurface")

# Let's say we want to find the medium sized craters in the image based one of the impact areas. For this we generate a template from a subsection of the image and apply a small gaussian blur using [`imfilter`](@ref). The gaussian blur often helps when the search targets are not exactly the same.
template = img[12:22,20:30]
template = imfilter(template,KernelFactors.gaussian((0.5,0.5)))

# Now that we have an image and a template, the next step is to define how we measure the similarity between a section of the image and the template. This can be done in multiple way, but a sum of square distances should work quite well. The `ImageDistance` package provides an already optimized version called `sqeuclidean`, which can be used to define a function for [`mapwindow`](@ref).  Let's call it `SDIFF`.

function SDIFF(template)
  (subsection)->sqeuclidean(subsection, template)
end

# To actually generate our similarity map we use [`mapwindow`](@ref) in the following way.

res = mapwindow(SDIFF(template), img, size(template), border=Fill(1)) .|> Gray
rescaled_map = adjust_histogram(res, LinearStretching())

# If the subsection is located at the border of the image the image has to be extended and in our case we will fill all values outside the image with 1. One thing to keep in mind is that because all of the square differences will be summed up per subsection the resulting sum can be bigger than 1. This will be a problem if we just convert it to an image to check the values. To rescale the values to be between 0 and 1 we can use `adjust_histogram`.

# To find the best locations we have to look for small values on the similarity map. This can be done by comparing if the pixel is below a certain value. Let's chose a value of `0.05`.

threshold = rescaled_map .< 0.05
Gray.(threshold)

# Now we see small blobs at the locations which match our template and we can label the connected regions by `label_components`.  This will enumerate are connected regions and `component_centroids` can be used to get the centroid of each region.  `component_centroids` also return the centroid for the background region, which is at the first position and we will omit it.

centroids = component_centroids(label_components(threshold))[2:end]

# To check if it worked correctly we can overlay the centroids with the original image using the `Plots` package. As the images are stored using the first index for rows we have to reverse the order of the coordinates to match the order of the plotting library.
plot(Gray.(img), size=(512,512))
scatter!(reverse.(centroids), label="centroids", ms=10, alpha=0.5, c=:red, msw=3)

using Plots: savefig #src
savefig("assets/template_matching.png") #src


text:
function resize(v,n)
   vNew = zeros(n);
    if n > size(v,1)
        for i = 1:size(v,1)
            vNew[i] = v[i];
        end
    elseif n < size(v,1)
        for i = 1:n
            vNew[i] = v[i];
        end
    elseif n == 0
        vNew = [];
    else
        vNew = v;
    end
    return vNew;
end


text:
#######################################################################
# Macros for printing information depending on verbosity
#
# ASSUMPTION: the package defines variable `VERBOSE`
#######################################################################

const STATUS_BEGIN = "..."
const STATUS_END   = " completed"

# If verbose, prints information about the start and completion of
# of `action` named `name`
macro status(name, action)
    quote
        @statusb $(esc(name))
        tmp = $(esc(action))
        @statuse $(esc(name))
        tmp
    end
end

# if verbose, prints information `name`
macro status(name)
    :(if VERBOSE ; @info($(esc(name))) end)
end

# If verbose, prints information about the start of `name`
macro statusb(name)
    :(if VERBOSE ; @info($(esc(name)) * $STATUS_BEGIN) end)
end

# If verbose, prints information about the completion of `name`
macro statuse(name)
    :(if VERBOSE ; @info($(esc(name)) * $STATUS_END) end)
end

# If debug mode, performs action
macro debugonly(action)
    :(if DEBUG ; $(esc(action)) end)
end

max_stars_repo_name: prl-julia/julia-sub

max_issues_repo_name: prl-julia/julia-sub

max_forks_repo_name: prl-julia/julia-sub


text:
module WellProductivity

using RecipesBase
using Dates
using DataFrames
using Colors
using Plots
using Statistics
using Distributions
using Interpolations
using Polynomials

export Reg, DeclinationForecast, Declination, OilInflow, OilOutflow, OilOutflowSen, GasInflow, GasOutflow, EspPerformance, EspTornado

include("OilRateRecipe.jl")
include("OilCumRecipe.jl")
include("GasRateRecipe.jl")
include("GasCumRecipe.jl")
include("BswRecipe.jl")
include("Reg.jl")
include("Types.jl")
include("Declination.jl")
include("OilInflow.jl")
include("IPRRecipe.jl")
include("OilOutflow.jl")
include("OilOutflowSen.jl")
include("GasInflow.jl")
include("VFPRecipe.jl")
include("GasOutflow.jl")
include("NodalRecipe.jl")
include("EspPerformance.jl")
include("EspTornado.jl")
include("EspRecipes.jl")

end # module


text:
"""
# Usage

    solve(FDDE::FDDEProblem, h, DelayPECE())

Using the delayed predictor-corrector method to solve the delayed fractional differential equation problem.

### References

```tex
@article{Wang2013ANM,
  title={A Numerical Method for Delayed Fractional-Order Differential Equations},
  author={Zhen Wang},
  journal={J. Appl. Math.},
  year={2013},
  volume={2013},
  pages={256071:1-256071:7}
}
```
"""
struct DelayPECE <: FractionalDiffEqAlgorithm end

function solve(FDDE::FDDEProblem, h, ::DelayPECE)
    @unpack f, ϕ, α, τ, tspan = FDDE
    T = tspan
    t = collect(0:h:T)
    maxn = size(t, 1)
    yp = collect(0:h:T+h)
    y = copy(t)
    y[1] = ϕ(0)

    @fastmath @inbounds @simd for n in 1:maxn-1
        yp[n+1] = 0
        @fastmath @inbounds @simd for j = 1:n
            yp[n+1] = yp[n+1]+b(j-1, n-1, α, h)*f(t[j], y[j], v(ϕ, j, τ, h, y, yp))
        end
        yp[n+1] = yp[n+1]/gamma(α)+ϕ(0)

        y[n+1] = 0

        @fastmath @inbounds @simd for j=1:n
            y[n+1] = y[n+1]+a(j-1, n-1, α, h)*f(t[j], y[j], v(ϕ, j, τ, h, y, yp))
        end

        y[n+1] = y[n+1]/gamma(α)+h^α*f(t[n+1], yp[n+1], v(ϕ, n+1, τ, h, y, yp))/gamma(α+2)+ϕ(0)
    end

    delayedterm = copy(t)
    @fastmath @inbounds @simd for n = 1:maxn-1
        delayedterm[n] = v(ϕ, n, τ, h, y, yp)
    end
    return delayedterm, y
end

function a(j, n, α, h)
    if j == n+1
        result = 1
    elseif j == 0
        result = n^(α+1)-(n-α)*(n+1)^α
    else
        result = (n-j+2)^(α+1) + (n-j)^(α+1) - 2*(n-j+1)^(α+1)
    end
    return result*h^α / (α * (α + 1))
end

function b(j, n, α, h)
    return h^α/α*((n-j+1)^α - (n-j)^α)
end

function v(ϕ, n, τ, h, y, yp)
    if τ >= n*h
        return ϕ((n-1)*h-τ)
    else
        m = floor(Int, τ/h)
        δ = m-τ/h

        if m>1
            return δ*y[n-m+2] + (1-δ)*y[n-m+1]
        elseif m == 1
            return δ*yp[n+1] + (1-δ)*y[n]
        end
    end
end


text:
# abstract model definition
"""
Abstract model.
"""
abstract type Model end

# Model interface: mandatory
"""
    nsites(m::Model)

Number of lattice sites of the given model.
"""
nsites(m::Model) = error("Model $(typeof(m)) doesn't implement `nsites`!")


text:
@testset "c api memory leak test" begin
    if (Sys.islinux())
        function get_memory_usage()
            open("/proc/$(getpid())/statm") do io
                return split(read(io, String))[1]
            end
        end

        file = joinpath(videodir, "annie_oakley.ogg")

        @testset "open file test" begin
            check_size = 10
            usage_vec = Vector{String}(undef, check_size)

            for i in 1:check_size
                f = VideoIO.openvideo(file)
                close(f)
                GC.gc()

                usage_vec[i] = get_memory_usage()
            end

            println(usage_vec)

            @test usage_vec[end-1] == usage_vec[end]
        end

        @testset "open and read file test" begin
            check_size = 10
            usage_vec = Vector{String}(undef, check_size)

            for i in 1:check_size
                f = VideoIO.openvideo(file)
                img = read(f)
                close(f)
                GC.gc()

                usage_vec[i] = get_memory_usage()
            end

            println(usage_vec)

            @test usage_vec[end-1] == usage_vec[end]
        end
    end
end


text:
using KeyedDistributions
using Documenter

DocMeta.setdocmeta!(KeyedDistributions, :DocTestSetup, :(using KeyedDistributions); recursive=true)

makedocs(;
    modules=[KeyedDistributions],
    authors="Invenia Technical Computing Corporation",
    repo="https://github.com/invenia/KeyedDistributions.jl/blob/{commit}{path}#{line}",
    sitename="KeyedDistributions.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://invenia.github.io/KeyedDistributions.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
    checkdocs=:exports,
    strict=true,
)

deploydocs(;
    repo="github.com/invenia/KeyedDistributions.jl",
    devbranch = "main",
    push_preview = true,
)

max_stars_repo_name: invenia/KeyedDistributions

max_issues_repo_name: invenia/KeyedDistributions

max_forks_repo_name: invenia/KeyedDistributions


text:
## Base.peek doesn't support IOBuffer, this fixes that
@inline function get_next_char(p::Ptr{UInt8}, i::Int, len::Int)
    b = unsafe_load(p, i)
    if b < 0x80
        return Char(b)
    end
    c, i = Base.slow_utf8_next(p, b, i, len)
    return c
end
function peek(io::IOBuffer)
    i, len = io.ptr, io.size
    i > len && Base.throw_boundserror(io, i)
    return get_next_char(pointer(io.data), i, len)
end
peek(io) = Base.peek(io)

lastuint8(io::IOBuffer) = io.data[io.ptr-1]

## Similar to readuntil, but stops before adding any char appearing in delim
function readupto(io::IO, delim)
    out = IOBuffer()
    c = '\0'
    while !eof(io)
        c = read(io, Char)
        if c ∈ delim
            break
        end
        write(out, c)
    end
    return String(take!(out)), c
end

function eatws(io::IO)
    eof(io) && return io
    c = peek(io)
    while isspace(c) && !eof(io)
        read(io, Char)
        c = peek(io)
    end
    io
end

# function skiptext(io::IO)
#     c = peek(io)
#     if c == '"'
#         # Skip over text
#         read(io, Char)
#         readuntil(io, '"')
#         eatws(io)
#         c = peek(io)
#     end
#     return c
# end

## Utility to skip the "description" (text) field
function skiptext(io::IO, term)
    c = peek(io)
    while !eof(io)
        c ∈ term && return c
        if c == '"'
            # Skip over text between quotes
            read(io, Char)
            readuntil(io, '"')
            eatws(io)
            return peek(io)
        end
        if isspace(c)
            read(io, Char)
            c = peek(io)
            continue
        end
        readupto(io, (',', '/'))
        eatws(io)
        return eof(io) ? '\0' : peek(io)
    end
    return c
end


text:
abstract Equation{T <: Number}

function do_nothing(x...) end

include("ode.jl")
include("pode.jl")
include("sode.jl")
include("dae.jl")
include("pdae.jl")

# TODO Add functions and vectors for invariants (momentum maps, energy, ...).


text:
################################################################################
#
#  Saturation
#
################################################################################

# TH:
# Let U = <u_1,...,u_n,z> with z a generator for Tor(U)
# For a prime p the group U/U^p is F_p-vector space of dimension
# rank(U) or rank(U) + 1 (depending on the order of z).
# if p divides N(P) - 1 = #F_P*, then F_P*/F_P*^p is a one-dimensional
# F_p-vector space. Thus the canonical map F_p-linear
#               U/U^p -> F_P*/F_P*^p
# can be described by a 1 x (rank(U)) or 1 x (rank(U) + 1) matrix over F_p,
# and can be computed by solving discrete logarithms in F_P
#
function _is_saturated(U::UnitGrpCtx, p::Int, B::Int = 2^30 - 1, proof::Bool = false)
  if proof
    error("Not yet implemented")
  end

  N = 3*unit_rank(order(U))

  @vprint :UnitGroup 1 "Computing $N prime ideals for saturation ...\n"

  primes =  _find_primes_for_saturation(order(U), p, N, B)

  m = _matrix_for_saturation(U, primes[1], p)

  for i in 2:N
    m = vcat(m, _matrix_for_saturation(U, primes[i], p))
  end

  @vprint :UnitGroup 1 "Computing kernel of p-th power map ...\n"
  (K, k) = _right_kernel(m)

  K = transpose(K)
  L = lift(K)
  T = typeof(L[1,1])

  nonzerorows = Array{Int, 1}()

  for j in 1:rows(L)
    if !iszero_row(L, j)
      push!(nonzerorows, j)
    end
  end

  if k == 0
    return (true, zero(nf(order(U))))
  elseif k == 1 && sum(T[ L[nonzerorows[1], i]::T for i in 1:cols(L)-1]) == 0
    # Only one root, which is torsion.
    # We assume that the torsion group is the full torsion group
    return (true, zero(nf(order(U))))
  else
    for j in nonzerorows
      a = U.units[1]^(L[j, 1])
      for i in 2:length(U.units)
        a = a*U.units[i]^L[j, i]
      end

      if gcd(p, U.torsion_units_order) != 1
        a = a*elem_in_nf(U.torsion_units_gen)^L[j, length(U.units) + 1]
      end

      b = evaluate(a)

      @vprint :UnitGroup 1 "Testing/computing root ... \n"

      @vprint :UnitGroup 1 "Bitsize of coefficient $([nbits(elem_in_basis(U.order(b))[i]) for i in 1:degree(U.order)])"

      #has_root, roota = root(b, p)
      has_root, _roota = ispower(U.order(b), p)
      roota = elem_in_nf(_roota)

      if !has_root
        continue
      end

      return (false, roota)
    end
  end

  # try some random linear combination of kernel vectors

  MAX = 10

  for i in 1:MAX

    ra = rand(0:p-1, rows(K))
    v = zero_matrix(base_ring(K), 1, cols(K))
    for j in 1:cols(K)
      for l in 1:rows(K)
        v[1, j] = v[1, j] + ra[l]*K[l,j]
      end
    end

    if v == parent(v)(0)# || sum([v[1, j] for j in 1:rows(K)-1]) == 0
      continue
    end

    v = lift(v)

    a = U.units[1]^(v[1, 1])
    for j in 2:length(U.units)
      a = a*U.units[j]^v[1, j]
    end

    if gcd(p, U.torsion_units_order) != 1
      a = a*elem_in_nf(U.torsion_units_gen)^v[1, length(U.units) + 1]
    end

    b = evaluate(a)

    #has_root, roota = root(b, p)
    has_root, _roota = ispower(U.order(b), p)
    roota = elem_in_nf(_roota)

    if has_root
      return (false, roota)
    end
  end

  return (true, zero(nf(order(U))))
end

# The output will be of type
# elem_type(MatrixSpace(ResidueRing(FlintZZ, p), 1, rank(U) ( + 1))), so
# nmod_mat or fmpz_mod_mat
# THIS FUNCTION IS NOT TYPE STABLE
function _matrix_for_saturation(U::UnitGrpCtx, P::NfOrdIdl, p::Int)
  O = order(U)
  K = nf(O)
  F, mF = ResidueField(O, P)
  mK = extend(mF, K)
  g = _primitive_element(F)

  # We have to add the generator of the torsion group
  if gcd(p, U.torsion_units_order) != 1
    res = zero_matrix(ResidueRing(FlintZZ, p, cached=false), 1, unit_rank(O) + 1)
  else
    res = zero_matrix(ResidueRing(FlintZZ, p, cached=false), 1, unit_rank(O))
  end

  t = K()

  for i in 1:length(U.units)
    u = U.units[i]
    y = one(F)

    # P.gen_two should be P-unformizer
    #println("$(P.gen_one), $b, $(P.gen_two)")

    for b in base(u)
      t = b*anti_uniformizer(P.gen_two)^(valuation(b, P))

      #if mod(denominator(t), minimum(P)) == 0
      #  l = valuation(denominator(t), P)
      #  y = y*(mK(t*elem_in_nf(P.anti_uniformizer)^l)*mF(P.anti_uniformizer)^(-l))^u.fac[b]
      #else
        y = y*mK(t)^u.fac[b]
      #end
    end

    res[1, i] = disc_log(y, g, p)
  end

  if gcd(p, U.torsion_units_order) != 1
    res[1, unit_rank(O) + 1] = disc_log(mF(U.torsion_units_gen), g, p)
  end

  return res
end

# TH:
# This function finds n prime ideals P of O such that p divides N(P) - 1
# Moreover the prime ideals are unramified and min(P) does not divide
# the index of O in the equation order.
#
# The function loops through all prime ideals ordered by the minimum,
# starting at next_prime(st)
function _find_primes_for_saturation(O::NfOrd, p::Int, n::Int,
                                     st::Int = 0)
  res = Array{NfOrdIdl}(n)
  i = 0

  q = st
  while i < n
    q = next_prime(q)

    if mod(index(O), q) == 0 || isramified(O, q)
      continue
    end

    lp = prime_decomposition(O, q)

    j = 1

    while j <= length(lp) && i < n
      Q = lp[j]
      if mod(norm(Q[1]) - 1, p) == 0
        i = i + 1
        res[i] = Q[1]
      end
      j = j + 1
    end
  end

  return res
end

function _refine_with_saturation(c::ClassGrpCtx, u::UnitGrpCtx)
  @vprint :UnitGroup "Enlarging unit group using saturation ... \n"

  b = _validate_class_unit_group(c, u)

  p = 2

  while b > 1
    @vprint :UnitGroup 1 "Saturating at $p ... \n"

    @v_do :UnitGroup 1 pushindent()
    issat, new_unit = _is_saturated(u, p)
    @v_do :UnitGroup 1 popindent()

    while !issat
      #println("I have found a new unit: $new_unit")
      _add_dependent_unit(u, FacElem(new_unit))
      #println("$(u.tentative_regulator)")

      @v_do :UnitGroup 1 pushindent()
      b = _validate_class_unit_group(c, u)
      @v_do :UnitGroup 1 popindent()

      if b == 1
        break
      end

      @v_do :UnitGroup 1 pushindent()
      issat, new_unit = _is_saturated(u, p)
      @v_do :UnitGroup 1 popindent()
    end

    @v_do :UnitGroup 1 pushindent()
    b = _validate_class_unit_group(c, u)
    @v_do :UnitGroup 1 popindent()

    p = next_prime(p)
    if p > b
      break
    end
  end
  return b
end


text:
"""
DifferenceLists, (c) 2018, Bill Burdick (William R. Burdick Jr.)
MIT Licensed (see LICENSE file).
"""

module DifferenceLists

export DL, dl, concat, push, pushfirst, todl, dlconcat

"""
    DL(func)

Given function `func`, construct a difference list.

Difference lists are highly efficient, immutable, concatenate and prepend in constant time, and iterate in time N.

# Examples
```jldoctest
julia> [x for x = dl(1, 2, 3)]
3-element Array{Int64,1}:
 1
 2
 3
```
"""
struct DL
    func
end

"""
    dl()::DL
    dl(items...)::DL

Construct a difference list of `items`.

# Examples
```jldoctest
julia> dl()
dl()

julia> dl(1)
dl(1)

julia> dl(1, 2, 3)
dl(1, 2, 3)

julia> dl(1, dl(2, 3), 4)
dl(1, dl(2, 3), 4)
```
"""
dl() = DL(last -> last)
dl(items...) = todl(items)

"""
    todl(items)

Create a difference list from something you can iterate over

# Examples
```jldoctest
julia> todl([1, 2, 3])
dl(1, 2, 3)
```
"""
todl(items) = DL(last -> nextFor(items, iterate(items), last))
todl(dl::DL) = dl

"""
    nextFor(items, state, last)

Compute the next iteration value for an embedded collection.
"""
nextFor(items, ::Nothing, last) = last
nextFor(items, (item, state), last) = item, (items, state, last)

"""
    push(item, dl::DL)

Push an item onto the end of a difference list.

# Examples
```jldoctest
julia> push(2, push(1, dl(7, 8, 9)))
dl(7, 8, 9, 1, 2)
```
"""
push(dl::DL, items...) = concat(dl, todl(items))

"""
    pushfirst(item, dl::DL)

Push an item onto the front of a difference list.

# Examples
```jldoctest
julia> pushfirst(1, pushfirst(2, dl(7, 8, 9)))
dl(1, 2, 7, 8, 9)
```
"""
pushfirst(dl::DL, items...) = concat(todl(items), dl)

"""
    concat(lists::DL...)::DL

Concatenate difference lists in constant time

See also: [`dl`](@ref)

# Examples
```jldoctest
julia> concat(dl(1, 2), dl(3, 4))
dl(1, 2, 3, 4)

julia> concat(dl(1), dl(2))
dl(1, 2)
```
"""
concat(lists::DL...) = DL(last -> foldr((x, y) -> x.func(y), lists, init=last))
dlconcat(lists...) = DL(last -> foldr((x, y) -> x.func(y), map(todl, lists), init=last))

"""
    (a::DL)(lists::DL...)::DL

A difference list itself can be used as shorthand for concat.

See also: [`dl`](@ref), [`concat`](@ref)

# Examples
```jldoctest
julia> dl(1, 2)(dl(3, 4), dl(5, 6, 7))
dl(1, 2, 3, 4, 5, 6, 7)
```
"""
(a::DL)(lists...) = concat(a, map(todl, lists)...)

# Iteration support
Base.iterate(d::DL) = d.func(nothing)
Base.iterate(::DL, cur::Tuple{Any, Any}) = cur
Base.iterate(::DL, (items, state, last)::Tuple{Any, Any, Any}) = nextFor(items, iterate(items, state), last)
Base.iterate(::DL, ::Nothing) = nothing
Base.IteratorSize(::DL) = Base.SizeUnknown()

# value display support
Base.show(io::IO, dl::DL) = print(io, "dl(", join([sprint(show, x) for x = dl], ", "), ")")

end # module


text:
struct LRTestResult{N}
    nobs::Int
    deviance::NTuple{N, Float64}
    dof::NTuple{N, Int}
    pval::NTuple{N, Float64}
end

_diff(t::NTuple{N}) where {N} = ntuple(i->t[i+1]-t[i], N-1)

"""
    isnested(m1::StatisticalModel, m2::StatisticalModel; atol::Real=0.0)

Indicate whether model `m1` is nested in model `m2`, i.e. whether
`m1` can be obtained by constraining some parameters in `m2`.
Both models must have been fitted on the same data.
"""
function isnested end

"""
    lrtest(mods::StatisticalModel...; atol::Real=0.0)

For each sequential pair of statistical models in `mods...`, perform a likelihood ratio
test to determine if the first one fits significantly better than the next.

A table is returned containing degrees of freedom (DOF),
difference in DOF from the preceding model, deviance, difference in deviance
from the preceding model, and likelihood ratio and p-value for the comparison
between the two models.

Optional keyword argument `atol` controls the numerical tolerance when testing whether
the models are nested.

# Examples

Suppose we want to compare the effects of two or more treatments on some result.
Our null hypothesis is that `Result ~ 1` fits the data as well as
`Result ~ 1 + Treatment`.

```jldoctest
julia> using DataFrames, GLM

julia> dat = DataFrame(Result=[1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 1],
                       Treatment=[1, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2],
                       Other=categorical([1, 1, 2, 1, 2, 1, 3, 1, 1, 2, 2, 1]));

julia> nullmodel = glm(@formula(Result ~ 1), dat, Binomial(), LogitLink());

julia> model = glm(@formula(Result ~ 1 + Treatment), dat, Binomial(), LogitLink());

julia> bigmodel = glm(@formula(Result ~ 1 + Treatment + Other), dat, Binomial(), LogitLink());

julia> lrtest(nullmodel, model, bigmodel)
┌ Warning: Could not check whether models are nested as model type TableRegressionModel does not implement isnested: results may not be meaningful
└ @ StatsModels ~/.julia/dev/StatsModels/src/lrtest.jl:108
Likelihood-ratio test: 3 models fitted on 12 observations
──────────────────────────────────────────────
     DOF  ΔDOF  Deviance  ΔDeviance  p(>Chisq)
──────────────────────────────────────────────
[1]    1         16.3006
[2]    2     1   15.9559    -0.3447     0.5571
[3]    4     2   14.0571    -1.8988     0.3870
──────────────────────────────────────────────

julia> lrtest(bigmodel, model, nullmodel)
┌ Warning: Could not check whether models are nested as model type TableRegressionModel does not implement isnested: results may not be meaningful
└ @ StatsModels ~/.julia/dev/StatsModels/src/lrtest.jl:108
Likelihood-ratio test: 3 models fitted on 12 observations
──────────────────────────────────────────────
     DOF  ΔDOF  Deviance  ΔDeviance  p(>Chisq)
──────────────────────────────────────────────
[1]    4         14.0571
[2]    2    -2   15.9559     1.8988     0.3870
[3]    1    -1   16.3006     0.3447     0.5571
──────────────────────────────────────────────
```
"""
function lrtest(mods::StatisticalModel...; atol::Real=0.0)
    if length(mods) < 2
        throw(ArgumentError("At least two models are needed to perform LR test"))
    end
    T = typeof(mods[1])
    df = dof.(mods)
    forward = df[1] <= df[2]
    if !all(m -> typeof(m) == T, mods)
        throw(ArgumentError("LR test is only valid for models of the same type"))
    end
    if !all(==(nobs(mods[1])), nobs.(mods))
        throw(ArgumentError("LR test is only valid for models fitted on the same data, " *
                            "but number of observations differ"))
    end
    checknested = hasmethod(isnested, Tuple{T, T})
    if forward
        for i in 2:length(mods)
            if df[i-1] >= df[i] ||
                (checknested && !isnested(mods[i-1], mods[i], atol=atol))
                throw(ArgumentError("LR test is only valid for nested models"))
            end
        end
    else
        for i in 2:length(mods)
            if df[i] >= df[i-1] ||
                (checknested && !isnested(mods[i], mods[i-1], atol=atol))
                throw(ArgumentError("LR test is only valid for nested models"))
            end
        end
    end
    if !checknested
        @warn "Could not check whether models are nested as model type " *
            "$(nameof(T)) does not implement isnested: results may not be meaningful"
    end

    dev = deviance.(mods)
    Δdev = _diff(dev)

    Δdf = _diff(df)
    dfr = Int.(dof_residual.(mods))

    if (forward && any(x -> x > 0, Δdev)) || (!forward && any(x -> x < 0, Δdev))
        throw(ArgumentError("Residual deviance must be strictly lower " *
                            "in models with more degrees of freedom"))
    end

    pval = (NaN, chisqccdf.(abs.(Δdf), abs.(Δdev))...)
    return LRTestResult(Int(nobs(mods[1])), dev, df, pval)
end

function Base.show(io::IO, lrr::LRTestResult{N}) where N
    Δdf = _diff(lrr.dof)
    Δdev = _diff(lrr.deviance)

    nc = 6
    nr = N
    outrows = Matrix{String}(undef, nr+1, nc)

    outrows[1, :] = ["", "DOF", "ΔDOF", "Deviance", "ΔDeviance", "p(>Chisq)"]

    outrows[2, :] = ["[1]", @sprintf("%.0d", lrr.dof[1]), " ",
                     @sprintf("%.4f", lrr.deviance[1]), " ", " "]

    for i in 2:nr
        outrows[i+1, :] = ["[$i]", @sprintf("%.0d", lrr.dof[i]),
                           @sprintf("%.0d", Δdf[i-1]),
                           @sprintf("%.4f", lrr.deviance[i]), @sprintf("%.4f", Δdev[i-1]),
                           string(StatsBase.PValue(lrr.pval[i])) ]
    end
    colwidths = length.(outrows)
    max_colwidths = [maximum(view(colwidths, :, i)) for i in 1:nc]
    totwidth = sum(max_colwidths) + 2*5

    println(io, "Likelihood-ratio test: $N models fitted on $(lrr.nobs) observations")
    println(io, '─'^totwidth)

    for r in 1:nr+1
        for c in 1:nc
            cur_cell = outrows[r, c]
            cur_cell_len = length(cur_cell)

            padding = " "^(max_colwidths[c]-cur_cell_len)
            if c > 1
                padding = "  "*padding
            end

            print(io, padding)
            print(io, cur_cell)
        end
        print(io, "\n")
        r == 1 && println(io, '─'^totwidth)
    end
    print(io, '─'^totwidth)
end


text:
function StreamingOemWriter(arg0::Appendable, arg1::TimeScale, arg2::Map)
    return StreamingOemWriter((Appendable, TimeScale, Map), arg0, arg1, arg2)
end

function new_segment(obj::StreamingOemWriter, arg0::Frame, arg1::Map)
    return jcall(obj, "newSegment", StreamingOemWriter_Segment, (Frame, Map), arg0, arg1)
end

function write_header(obj::StreamingOemWriter)
    return jcall(obj, "writeHeader", void, ())
end