text:
# ## Basis function definitions

mutable struct PGBF
    expn::Float64
    x::Float64
    y::Float64
    z::Float64
    I::Int64
    J::Int64
    K::Int64
    norm::Float64
end

function pgbf(expn,x=0,y=0,z=0,I=0,J=0,K=0,norm=1)
    p = PGBF(expn,x,y,z,I,J,K,norm)
    normalize!(p)
    return p
end

function amplitude(bf::PGBF,x,y,z)
    dx,dy,dz = x-bf.x,y-bf.y,z-bf.z
    r2 = dist2(dx,dy,dz)
    return bf.norm*(dx^bf.I)*(dy^bf.J)*(dz^bf.K)*exp(-bf.expn*r2)
end

function normalize!(pbf::PGBF)
    pbf.norm /= sqrt(overlap(pbf,pbf))
end

mutable struct CGBF
    x::Float64
    y::Float64
    z::Float64
    I::Int64
    J::Int64
    K::Int64
    norm::Float64
    pgbfs::Array{PGBF,1}
    coefs::Array{Float64,1}
end

cgbf(x=0,y=0,z=0,I=0,J=0,K=0) = CGBF(x,y,z,I,J,K,1.0,PGBF[],Float64[])

function amplitude(bf::CGBF,x,y,z)
    s = 0
    for (c,pbf) in primitives(bf)
        s += c*amplitude(pbf,x,y,z)
    end
    return bf.norm*s
end

function normalize!(bf::CGBF)
    bf.norm /= sqrt(overlap(bf,bf))
end

primitives(a::CGBF) = zip(a.coefs,a.pgbfs)

function push!(cbf::CGBF,expn,coef)
    Base.push!(cbf.pgbfs,pgbf(expn,cbf.x,cbf.y,cbf.z,cbf.I,cbf.J,cbf.K))
    Base.push!(cbf.coefs,coef)
    normalize!(cbf)
end

function contract(f,a::CGBF,b::CGBF)
    s = 0
    for (ca,abf) in primitives(a)
        for (cb,bbf) in primitives(b)
            s += ca*cb*f(abf,bbf)
        end
    end
    return a.norm*b.norm*s
end

function contract(f,a::CGBF,b::CGBF,c::CGBF,d::CGBF)
    s = 0
    for (ca,abf) in primitives(a)
        for (cb,bbf) in primitives(b)
            for (cc,cbf) in primitives(c)
                for (cd,dbf) in primitives(d)
                    s += ca*cb*cc*cd*f(abf,bbf,cbf,dbf)
                end
            end
        end
    end
    return a.norm*b.norm*c.norm*d.norm*s
end

function test_pgbf()
    @testset "test_pgbf" begin
    s = pgbf(1.0)
    px = pgbf(1.0,0,0,0,1,0,0)
    @test amplitude(s,0,0,0) ≈ 0.71270547
    @test amplitude(px,0,0,0) ≈ 0
    end # testset
end
function test_cgbf()
    @testset "test_cgbf" begin
    c = cgbf(0.0,0.0,0.0)
    push!(c,1,1)
    @test amplitude(c,0,0,0) ≈ 0.71270547
    c2 = cgbf(0,0,0)
    push!(c2,1,0.2)
    push!(c2,0.5,0.2)
    @test overlap(c2,c2) ≈ 1
    end # testset
end


text:
module errors

	function _next_perm(v::Int64)
	    v = unsigned(v);
	    t = (v | unsigned(v - 1)) + 1;
	    t = unsigned(t);
	    w = t | ((div((t & -t), (v & -v)) >> 1) - 1);
	    return convert(Int64, w)
	end

	function _all_bitstrings_of_weight(n, w)
		K = binomial(n, w);
		perm = 2^w-1;
		perms = [];
		for k=1:K
		    perms = [perms; perm];
		    perm = _next_perm(perm);
		end

		E = zeros(n, K);
		for k=1:K
		    E[:, k] = digits(perms[k], 2, n)
		end
		return convert(Array{Bool, 2}, E)
	end

	function _all_bitstrings_up_to_weight(n, w)
		w_0 = convert(Array{Bool, 2}, zeros(n, 1))
		w_1 = convert(Array{Bool, 2}, eye(n))
		# All errors up to weight w_x
		if w > 1
			# All weight 1 errors
			out = [w_0 w_1]
			# All other weights
			for j=2:w
				out = [out _all_bitstrings_of_weight(n, j)]
			end
		elseif w == 1
			out = [w_0 w_1]
		elseif w == 0
			out = w_0
		end
		return out
	end

	# Generate all errors up to and including weights w_x and w_z
	function _errors(n, w_x, w_z)
		# All X errors up to weight w_x
		out_x = _all_bitstrings_up_to_weight(n, w_x)'

		# All Z errors up to weight w_z
		out_z = _all_bitstrings_up_to_weight(n, w_z)'

		# Take the outer product of the errors
		n_x = size(out_z)[1];
		n_z = size(out_x)[1];
		rep_x = repeat(out_x, inner=[n_x; 1]);
		rep_z = repeat(out_z, outer=[n_z; 1]);
		return (rep_x, rep_z)
	end
	export _errors, _all_bitstrings_up_to_weight, _all_bitstrings_of_weight

end

ext: jl


text:
using JSServe, FileIO, Colors, Observables
using JSServe.DOM
using JSServe: App
JSServe.browser_display()
dir = mktempdir()
for i in 1:5
    save(joinpath(dir, "img$i.jpg"), rand(RGB{Colors.N0f8}, 200, 200))
end

app = App() do
    files = JSServe.Asset.(joinpath.(dir, filter(x-> endswith(x, ".jpg"), readdir(dir))))
    slider = JSServe.Slider(1:length(files))
    image_obs = Observable(first(files))
    img = DOM.img(src=image_obs, style="height: 200px;")
    on(slider) do idx
        image_obs[] = files[idx]
    end
    return DOM.div(slider, img)
end

display(app)


text:
# Include this file, that is,
# julia> include("Setup.jl")
# in order to use the package without installing the package...
function addpath!(this)
  new_path = joinpath(@__DIR__, this)
  if !any(LOAD_PATH .== new_path)
    push!(LOAD_PATH, new_path)
  end
end

addpath!("src")
nothing


text:
using LowRankModels

m,n,k = 10,20,3
Y = rand(3,20)
A = rand(10,20)

ry = Regularizer[fixed_latent_features(Y[:,i]) for i=1:n]
glrm = GLRM(A, QuadLoss(), SimplexConstraint(), ry, k+1)
X, Yp, ch = fit!(glrm)

@assert(Yp[1:k,end] == Y)

m,n,k = 10,20,3
Y = rand(3,20)
A = rand(10,20)

ry = Regularizer[fixed_last_latent_features(Y[:,i]) for i=1:n]
glrm = GLRM(A, QuadLoss(), SimplexConstraint(), ry, k+1)
X, Yp, ch = fit!(glrm)

@assert(Yp[2:end,:] == Y)


text:
# Copyright 2019 Tobias Frilling
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
    plot_matrix(m; <keyword arguments>)

Visualize the given U-, P- or U*-matrix.

# Arguments
- `projection`: A projection of data points to ESOM neuron to show on the map.
- `classes`: The class of each data point. Used to color the dots for the projection.
- `normalize`: Whether to normalize the given matrix. Default: `true`.
- `colorStyle`: Either `:umatrix` or `:pmatrix`. Default: `:umatrix`.
"""
@userplot Plot_Matrix

@recipe function f(p::Plot_Matrix;
                   settings = defaultSettings,
                   projection = Dict(),
                   key = 1:length(projection),
                   classes = Dict(),
                   normalize = true,
                   colorStyle = :umatrix)

    m = p.args[1]
    if normalize
        (min_m, max_m) = quantile(m[:], [.01, .99])
        m = (m .- min_m) ./ (max_m - min_m)
    end
    (rows, columns) = size(m)

    grid := false
    framestyle := :none
    yflip --> true
    colorbar --> true
    aspect_ratio --> 1

    @series begin
        seriestype := :contourf
        (min_m, max_m) = quantile(m[:], [.01, .99])
        levels --> round(Int, max_m / max(min_m, .05))
        seriescolor --> if colorStyle == :pmatrix colormap_pmatrix
                        else colormap_umatrix end
        linewidth --> 0.1
        if settings.toroid
            [m m; m m]
        else
            m
        end
    end

    if !isempty(projection)
        seriestype := :scatter
        widen := false
        legend := false
        colorbar_entry := false
        markersize --> 2.5
        markerstrokewidth --> 0.4

        if isempty(classes)
            classes[1] = keys(projection) |> collect
        end
        colors = distinguishable_colors(length(classes), RGB(1,0,1))

        for (class, indices) in classes
            @series begin
                markercolor := colors[class]
                v = [projection[key[index]] for index in indices]
                if settings.toroid
                    v = vcat(v, map(i -> i + CartesianIndex(rows, 0), v))
                    v = vcat(v, map(i -> i + CartesianIndex(0, columns), v))
                end
                getindex.(v,2), getindex.(v,1)
            end
        end
    end

end

colormap_umatrix = begin
    local c = ["#3C6DF0", "#3C6DF0", "#3C6DF0", "#006602", "#006A02", "#006D01",
               "#007101", "#007501", "#007901", "#007C00", "#008000", "#068103",
               "#118408", "#0B8305", "#17860A", "#1D870D", "#228810", "#288A12",
               "#2E8B15", "#348D18", "#398E1A", "#3F8F1D", "#45911F", "#4A9222",
               "#509325", "#569527", "#5C962A", "#61982C", "#67992F", "#6D9A32",
               "#729C34", "#789D37", "#7E9F39", "#84A03C", "#89A13F", "#8FA341",
               "#95A444", "#9AA547", "#A0A749", "#A6A84C", "#ACAA4E", "#B1AB51",
               "#B7AC54", "#BDAE56", "#C3AF59", "#C8B15B", "#CEB25E", "#CBAF5C",
               "#C8AC59", "#C5A957", "#C3A654", "#C0A352", "#BDA050", "#BA9D4D",
               "#B7994B", "#B49648", "#B29346", "#AF9044", "#AC8D41", "#A98A3F",
               "#A6873C", "#A3843A", "#A08138", "#9E7E35", "#9B7B33", "#987830",
               "#95752E", "#92722B", "#8F6E29", "#8C6B27", "#8A6824", "#876522",
               "#84621F", "#815F1D", "#7E5C1B", "#7B5918", "#795616", "#765313",
               "#714E0F", "#6C480B", "#674307", "#6F4D15", "#785822", "#806230",
               "#896D3E", "#91774C", "#998159", "#A28C67", "#AA9675", "#B3A183",
               "#BBAB90", "#C3B59E", "#CCC0AC", "#D4CABA", "#DDD5C7", "#E5DFD5",
               "#E7E1D8", "#E9E4DB", "#EBE6DE", "#ECE8E1", "#EEEAE4", "#F0EDE7",
               "#F2EFEA", "#F4F1ED", "#F6F4F0", "#F8F6F3", "#F9F8F6", "#FBFAF9",
               "#FDFDFC", "#FFFFFF", "#FFFFFF", "#FEFEFE", "#FEFEFE", "#FEFEFE",
               "#FDFDFD", "#FDFDFD", "#FDFDFD", "#FCFCFC", "#FCFCFC", "#FCFCFC",
               "#FBFBFB", "#FBFBFB", "#FBFBFB", "#FAFAFA", "#FAFAFA", "#FAFAFA",
               "#F9F9F9", "#F9F9F9", "#FFFFFF", "#FFFFFF"]
    (s -> parse(RGB, s)).(c) |> ColorGradient
end

colormap_pmatrix = begin
    local c = ["#FFFFFF", "#FFFFF7", "#FFFFEF", "#FFFFE7", "#FFFFDF", "#FFFFD7",
               "#FFFFCF", "#FFFFC7", "#FFFFBF", "#FFFFB7", "#FFFFAF", "#FFFFA7",
               "#FFFF9F", "#FFFF97", "#FFFF8F", "#FFFF87", "#FFFF80", "#FFFF78",
               "#FFFF70", "#FFFF68", "#FFFF60", "#FFFF58", "#FFFF50", "#FFFF48",
               "#FFFF40", "#FFFF38", "#FFFF30", "#FFFF28", "#FFFF20", "#FFFF18",
               "#FFFF10", "#FFFF08", "#FFFF00", "#FFFA00", "#FFF400", "#FFEF00",
               "#FFEA00", "#FFE400", "#FFDF00", "#FFDA00", "#FFD400", "#FFCF00",
               "#FFCA00", "#FFC500", "#FFBF00", "#FFBA00", "#FFB500", "#FFAF00",
               "#FFAA00", "#FFA500", "#FF9F00", "#FF9A00", "#FF9500", "#FF8F00",
               "#FF8A00", "#FF8500", "#FF8000", "#FF7A00", "#FF7500", "#FF7000",
               "#FF6A00", "#FF6500", "#FF6000", "#FF5A00", "#FF5500", "#FF5000",
               "#FF4A00", "#FF4500", "#FF4000", "#FF3A00", "#FF3500", "#FF3000",
               "#FF2B00", "#FF2500", "#FF2000", "#FF1B00", "#FF1500", "#FF1000",
               "#FF0B00", "#FF0500", "#FF0000", "#FA0000", "#F40000", "#EF0000",
               "#EA0000", "#E40000", "#DF0000", "#DA0000", "#D40000", "#CF0000",
               "#CA0000", "#C50000", "#BF0000", "#BA0000", "#B50000", "#AF0000",
               "#AA0000", "#A50000", "#9F0000", "#9A0000", "#950000", "#8F0000",
               "#8A0000", "#850000", "#800000", "#7A0000", "#750000", "#700000",
               "#6A0000", "#650000", "#600000", "#5A0000", "#550000", "#500000",
               "#4A0000", "#450000", "#400000", "#3A0000", "#350000", "#300000",
               "#2B0000", "#250000", "#200000", "#1B0000", "#150000", "#100000",
               "#0B0000", "#050000"]
    (s -> parse(RGB, s)).(c) |> ColorGradient
end


text:
function add_compat_section!(project::Dict)
    if !haskey(project, "compat")
        project["compat"] = Dict{Any, Any}()
    end
    return project
end

function get_random_string()
    return string(utc_to_string(now_localzone()), "-", rand(UInt32))::String
end

function generate_pr_title_parenthetical(keep_or_drop::Symbol,
                                         parenthetical_in_pr_title::Bool)::String
    if parenthetical_in_pr_title
        if keep_or_drop == :keep
            return " (keep existing compat)"
        elseif keep_or_drop == :drop
            return " (drop existing compat)"
        else
            return ""
        end
    else
        return ""
    end
end


text:
function make_model(grid, ν; timestepper = :RungeKutta3, advection = WENO5())

    v_bcs, w_bcs = no_slip_bc(grid)

    return IncompressibleModel(timestepper = timestepper,
                                 advection =  advection,
                                      grid = grid,
                                  buoyancy = BuoyancyTracer(),
                                   tracers = :b,
                       boundary_conditions = (v=v_bcs, w=w_bcs),
                                   closure = IsotropicDiffusivity(ν=ν, κ=0))
end


text:
# Some help functions to generate interesting test cases
circle_segment(N) = sin.(acos.(range(-1, stop=1, length=N)))
linear_trend(N) = (0:N-1)

""" brute_force_multi(g, M, t=1:length(g); tol=1e-3)
    Finds the solution to the cardinality constrained problem for
    m=1:M segments using brute force with possible breakpoints on grid t.
    `t` has to be supplied if `g` is not `AbstractArray`
"""
function brute_force_multi(g, M, t=1:length(g); tol=1e-3)

    if length(g) <= M
        @warn "Specified M too large. Setting M=length(g)-1."
        M = length(g)-1
    end

    ℓ, V_N = EllZeroTrendFiltering.get_transition_costs(g, t, true, tol=tol)

    I_vec = Vector{Vector{Int64}}(undef, M)
    Y_vec = Vector{Vector{Float64}}(undef, M)
    f_vec = Vector{Float64}(undef, M)

    for m=1:M
        (I_bf, Y_bf, f_bf) = brute_force_search(ℓ, V_N, m)
        I_vec[m] = t[I_bf]
        Y_vec[m] = Y_bf
        f_vec[m] = f_bf
    end

    return I_vec, Y_vec, f_vec
end


text:
using JuMP, Clp

rpomdp = Baby3RPOMDP()
dim = n_observations(rpomdp)
uncset = observation(rpomdp,true,true)
cdfl = cumsum(pdf.(uncset.lower))
cdfu = cumsum(pdf.(uncset.upper))

coeff = [5,2,3]
m = Model(solver = ClpSolver())
@variable(m, 0 <= p[i = 1:dim] <= 1)
@objective(m, Min, dot(coeff,p))
@constraint(m, sum(p) == 1)
@constraint(m, cdfl[1] <= p[1] <= cdfu[1])
for i = 2:dim
    pl = cdfl[i] - cdfu[i-1]
    pu = cdfu[i] - cdfl[i-1]
    @constraint(m, pl <= p[i] <= pu)
end
JuMP.solve(m)
argmin = getvalue(p)
obj = getobjectivevalue(m)

rpomdp = Baby3RPOMDP()
dist_t = transition(rpomdp,true,true)
dist = observation(rpomdp,true,true)
dim = n_observations(rpomdp)
cdfl = cumsum(pdf.(uncset.lower))
cdfu = cumsum(pdf.(uncset.upper))
pt = [0.5,0.6]
α = [1,2]
m = Model(solver = ClpSolver())
@variable(m, 0 <= p[i = 1:dim] <= 1)
@objective(m, Min, dot(α, p .* pt))
@constraint(m, cdfl[zind] <= p <= cdfu[zind])
JuMP.solve(m)
argmin = getvalue(p)
obj = getobjectivevalue(m)

dist = observation(prob,a,sp)
dim = n_observations(rpomdp)
function robustdist(uncset::PBox, dim::Int, α::Vector{Float64})
    cdfl = cumsum(pdf.(uncset.lower))
    cdfu = cumsum(pdf.(uncset.upper))
    m = Model(solver = ClpSolver())
    @variable(m, 0 <= p[i = 1:dim] <= 1)
    @objective(m, Min, dot(α, p .* pt))
    @constraint(m, sum(p) == 1)
    @constraint(m, cdfl[1] <= p[1] <= cdfu[1])
    for i = 2:dim
        pl = cdfl[i] - cdfu[i-1]
        pu = cdfu[i] - cdfl[i-1]
        @constraint(m, pl <= p[i] <= pu)
    end
    JuMP.solve(m)
    argmin = getvalue(p)
    obj = getobjectivevalue(m)
end


text:
# 1
wavelet_matrices(ws, J, L; ξ0=maximum(ws) / 2, c=1) =
    [wavelet_matrix(ws, j, θ; ξ0, L, c) for j in 0:J - 1, θ in range(0., length=L, step=2π / L)]

# 2
e_θ(l, L) = (cos(π / 2 + 2π * l / L), sin(π / 2 + 2π * l / L))
τ_θ_j(l, j_, L) = round.(Int, 2^j_ .* e_θ(l, L))

# ! k > 1 doesn't work if not in Γ_H

function default_Γ_H(J, L, K)
    Γ_H = NamedTuple{
        (:j, :l, :θ, :k, :j_, :l_, :θ_, :k_, :τ_),
        Tuple{Int64, Int64, Float64, Int64, Int64,
        Int64, Float64, Int64, Tuple{Int64, Int64}}
    }[]
    p = Iterators.product(
        0:L - 1,
        0:L - 1,
        0:J - 1,
        0:J - 1,
        [false, true],
        0:K,
        0:K
    )

    for (l, l_, j, j_, isshift, k, k_) in p
        l, l_, j, j_, isshift, k, k_
        if l_ > l
            continue
        end
        if (j > j_) || (j < j_ - 2)
            continue
        end
        if k > min(k_, 1)
            continue
        end

        if isshift
            τ_ = τ_θ_j(l, j_, L)
        else
            τ_ = (0, 0)
        end

        θ = 2π * l / L
        θ_ = 2π * l_ / L
        dl = abs(l - l_)

        if j == j_
            if k == 0 && k_ > 1
                continue
            end
            if k == 1
                if k_ != 1
                    continue
                elseif dl > 2
                    continue
                end
            end
        else
            if k == 0 && k_ == 4
                continue
            end
            if k == 1
                if k_ != 2^(j_ - j)
                    continue
                elseif dl > 2
                    continue
                end
            end
        end

        push!(Γ_H, (;j, l, θ, k, j_, l_, θ_, k_, τ_))
    end
    Γ_H
end

# 3
"Fourier transform of δ(x - x0)"
# Δ(w, x0) = exp(-2π * im * dot(x0, w))
"Fourier transform of point process"
# M(Ws, μ::AbstractMatrix) = [sum(Δ(w, x) for x in eachcol(μ)) for w in Ws]

function Δ(x0, ws)
    a = x0[1] .* ws
    b = x0[2] .* ws
    -2pi * (a' .+ b) .|> cis
end

M(x, ws) = sum(Δ(xi, ws) for xi in eachcol(x))

# main
struct WaveletParams{Tw}
    s::Float64
    N::Int
    J::Int
    L::Int
    K::Int
    σ::Float64

    ws::Tw
    Ws::Matrix{Vector{Float64}}
    FN::Function

    Ψs::Matrix{Matrix{Float64}}

    Γ_H::Vector{NamedTuple{
        (:j, :l, :θ, :k, :j_, :l_, :θ_, :k_, :τ_),
        Tuple{Int64, Int64, Float64, Int64, Int64,
        Int64, Float64, Int64, Tuple{Int64, Int64}}}
        }
    jlk::Array{Tuple{Int64, Int64, Int64}, 3}
end

function WaveletParams(s, N, J, L, K, σ, Γ_H = default_Γ_H(J, L, K))
    dx = 2s / N
    xs = range(-s; length=N, step=dx)
    ws = fftfreq(N, 1 / dx)
    Ws = [[w1, w2] for w1 in ws, w2 in ws]

    g = MvNormal([0, 0], σ^2 * I)
    gaus = [pdf(g, [x, y]) for x in xs, y in xs]
    Gaus = gaus |> ifftshift |> fft |> real

    # FN(μ) = M(Ws, μ)
    FN(μ) = M(μ, ws)

    Ψs = [Ψ .* Gaus for Ψ in wavelet_matrices(ws, J, L)]

    jlk = collect(Iterators.product(1:J, 1:L, 1:K + 1))

    WaveletParams(s, N, J, L, K, σ, ws, Ws, FN, Ψs, Γ_H, jlk)
end

phase_harmonics(z::Complex, k::Int) = abs(z) * exp(im * angle(z) * k)

function wavelet_phase_harmonics(μ, wp, vs)
    M = wp.FN(μ)
    # cs = [ifft(M .* Ψ) for Ψ in wp.Ψs]
    cs = map(Ψ -> ifft(M .* Ψ), wp.Ψs)
    m = ifft(M) .- vs[1]
    # csk = [phase_harmonics.(cs[j, l], k - 1) .- vs[2][j, l, k] for (j, l, k) in wp.jlk]
    csk = map(wp.jlk) do jlk
        (j, l, k) = jlk
        phase_harmonics.(cs[j, l], k - 1) .- vs[2][j, l, k]
    end
    m, csk
end

function v_λ_k_all(μ, wp::WaveletParams)
    m, csk = wavelet_phase_harmonics(μ, wp, (0., zeros(wp.J, wp.L, wp.K + 1)))
    mean(m), mean.(csk)
end

function K_μ(csk; j, l, k, j_, l_, k_, τ_)
    A = csk[j + 1, l + 1, k + 1]
    B = circshift(csk[j_ + 1, l_ + 1, k_ + 1], τ_)
    mean(@. A * conj(B))
end

function K_μ_0(m, csk)
    map(c -> mean(m .* conj.(c)), csk)
end

K_μ(csk, p) = K_μ(csk; p.j, p.l, p.k, p.j_, p.l_, p.k_, p.τ_)

function K_all(μ, wp::WaveletParams, vs)
    m, csk = wavelet_phase_harmonics(μ, wp, vs)
    K_main = map(p -> K_μ(csk, p), wp.Γ_H)
    K_0 = K_μ_0(m, csk)[:]
    [K_main; K_0]
end

function lossW(x; wp, vs, K0, weights, index)
    k = (K_all(x, wp, vs) .* weights)[index]
    k0 = (K0 .* weights)[index]
    sqL2dist(k, k0)
end
lossW(x, p) = lossW(x; p...)

function lossW_params(x, wp, weights=ones(length(wp.Γ_H) + length(wp.jlk)), index=:)
    vs = v_λ_k_all(x, wp)
    K0 = K_all(x, wp, vs)
    (; wp, vs, K0, weights, index)
end


text:
function read_struct_tag!(io::IO, v::Int64=0)
  s = read(io, UInt8)
  s == 0x53 || (close(io); error("damaged or scrambled SUDS file; can't continue."))
  m = read(io, UInt8)
  SB.sid = read(io, Int16)
  SB.nbs = read(io, Int32)
  SB.nbx = read(io, Int32)

  (v > 0) && println("suds_structtag: machine code = ", Char(m),
                      ", ID = ",          SB.sid,
                      ", struct size = ", SB.nbs, " B",
                      ", data size = ",   SB.nbx, " B",
                      )

  return nothing
end

function staident!(io::IO, v::Int64=0)
  read!(io, SB.hdr)
  fill!(SB.id, 0x00)

  # Set ID string
  j = 1
  for i = 1:2
    if SB.hdr[i] != 0x00
      SB.id[j] = SB.hdr[i]
      j += 1
    end
  end
  SB.id[j] = 0x2e
  j += 1
  for i = 5:9
    if SB.hdr[i] != 0x00
      SB.id[j] = SB.hdr[i]
      j += 1
    end
  end
  SB.id[j] = 0x2e
  SB.id[j+1] = 0x2e
  SB.id[j+2] = SB.hdr[10]
  j += 2
  setfield!(SB, :id_str, unsafe_string(pointer(getfield(SB, :id)), j))

  # Print new ID to STDOUT if v > 1
  (v > 1) && println( "id  = ",         SB.id_str,
                      ", inst code = ", Int16(SB.hdr[12]) << 8 | Int16(SB.hdr[11]))
  return nothing
end

function read_chansetentry!(S::GphysData, io::IO, v::Int64)
  SB.C.inst    = read(io, Int32)
  SB.C.stream  = read(io, Int16)
  SB.C.chno    = read(io, Int16)
  staident!(io, 0)
  if v > 2
    println("id = ",          SB.id_str,
            ", inst code = ", Int16(SB.hdr[12]) << 8 | Int16(SB.hdr[11]),
            ", inst = ",      SB.C.inst,
            ", stream = ",    SB.C.stream,
            ", chno = ",      SB.C.chno)
  end
  return nothing
end

# #  STAT_IDENT:  Station identification
# read_1!(S::GphysData, io::IO, v::Int64, full::Bool) = staident!(io, v)
#
# #  STRUCTTAG:  Structure to identify structures when archived together
# read_2!(S::GphysData, io::IO, v::Int64, full::Bool) = read_struct_tag!(io)

# STATIONCOMP:  Generic station component information
function read_5!(S::GphysData, io::IO, v::Int64, full::Bool)
  staident!(io, 0)
  SB.S05.az             = read(io, Int16)    # azimuth, N°E
  SB.S05.inc            = read(io, Int16)    # incidence, from vertical
  SB.S05.lat            = read(io, Float64)  # latitude, N = +
  SB.S05.lon            = read(io, Float64)  # longitude, E = +
  SB.S05.ele            = read(io, Float32)  # elevation, meters
  read!(io, SB.S05.codes)
  read!(io, SB.S05.gain)
  read!(io, SB.S05.a2d)
  SB.t_i32              = read(io, Int32)    # date/time values became effective
  read!(io, SB.S05.t_corr)

  # Post-read processing
  SB.data_type  = SB.S05.codes[9]
  sensor_type   = SB.S05.codes[8]

  # Does this station exist in S?
  i = findid(SB.id_str, S)
  if i == 0
    loc = GeoLoc( "",
                  SB.S05.lat,
                  SB.S05.lon,
                  Float64(SB.S05.ele),
                  0.0,
                  Float64(SB.S05.az),
                  Float64(SB.S05.inc)
                  )
    misc = if full
            Dict{String, Any}(
              "ic"             => Int16(SB.hdr[12]) << 8 | Int16(SB.hdr[11]),
              "enclosure"      => Char(SB.S05.codes[1]),
              "annotation"     => Char(SB.S05.codes[2]),
              "recorder"       => Char(SB.S05.codes[3]),
              "rockclass"      => Char(SB.S05.codes[4]),
              "rocktype"       => Int16(SB.S05.codes[6]) << 8 | Int16(SB.S05.codes[5]),
              "sitecondition"  => Char(SB.S05.codes[7]),
              "sensor_type"    => Char(sensor_type),
              "data_type"      => Char(SB.data_type),
              "data_units"     => Char(SB.S05.codes[10]),
              "polarity"       => Char(SB.S05.codes[11]),
              "st_status"      => Char(SB.S05.codes[12]),
              "max_gain"       => SB.S05.gain[1],
              "clip_value"     => SB.S05.gain[2],
              "con_mvolts"     => SB.S05.gain[3],
              "a2d_ch"         => SB.S05.a2d[1],
              "a2d_gain"       => SB.S05.a2d[2],
              "dt_eff"         => SB.t_i32,
              "clock_corr"     => SB.S05.t_corr[1],
              "sta_delay"      => SB.S05.t_corr[2]
              )
          else
            Dict{String, Any}(
              "data_type"      => Char(SB.data_type),
              "sensor_type"    => Char(sensor_type),
              )
          end

    push!(S, SeisChannel( id = SB.id_str,
                          loc = loc,
                          misc = misc,
                          gain = Float64(SB.S05.gain[1]),
                          units = sensor_types[sensor_type]
                          )
         )
    i = S.n
  end
  return i
end

# MUXDATA:  Header for (possibly) multiplexed data
function read_6!(S::GphysData, io::IO, v::Int64, full::Bool)
  SB.T.net      = read(io, UInt16)
  skip(io, 2)
  SB.t_f64      = read(io, Float64)
  SB.t_i16      = read(io, Int16)
  SB.T.n_ch     = read(io, Int16)
  SB.fs         = read(io, Float32)
  SB.data_type  = read(io, UInt8)
  skip(io, 3)
  SB.T.ns       = read(io, Int32)
  SB.nz         = read(io, Int32)

  if v > 1
    println("net = ",         Char(SB.T.net & 0x00ff), Char(SB.T.net >> 8),
            ", btime = ",     u2d(SB.t_f64),
            ", loctime = ",   SB.t_i16,
            ", n = ",         SB.T.n_ch,
            ", fs = ",        SB.fs,
            ", data_type = ", Char(SB.data_type),
            ", n sweeps = ",  SB.T.ns,
            ", nx = ",        SB.nz
            )
  end

  # return (net, data_type, btime, loctime, numchans, fs, data_type, nx)
  return nothing
end

#=  DESCRIPTRACE:  Descriptive information about a seismic trace.
                   Normally followed by waveform data =#
function read_7!(S::GphysData, io::IO, v::Int64, full::Bool)
  staident!(io, 0)
  SB.t_f64       = read(io, Float64)
  SB.t_i16       = read(io, Int16)
  SB.data_type   = read(io, UInt8)
  SB.T.desc      = read(io, UInt8)
  skip(io, 4)
  SB.nz          = read(io, Int32)
  SB.fs          = read(io, Float32)
  skip(io, 16)
  SB.t_f64      += read(io, Float64)
  SB.rc          = read(io, Float32)

  if v > 2
    tl = SB.t_i16
    println("id = ",          SB.id_str,
            ", inst code = ", Int16(SB.hdr[12]) << 8 | Int16(SB.hdr[11]),
            ", begin = ",     u2d(SB.t_f64 + tl*60),
            " (GMT ", tl < 0 ? "-" : "+", div(tl, 60), ")",
            ", data_type = ", Char(SB.data_type),
            ", nx = ",        SB.nx,
            ", fs = ",        SB.fs + SB.rc)
  end

  # return (id, data_type, nx, Float64(fs+rc), btime+tc)
  return nothing
end

# FEATURE:  Observed phase arrival time, amplitude, and period
function read_10!(S::GphysData, io::IO, v::Int64, full::Bool)
  staident!(io, 0)
  SB.P.pc           = read(io, Int16)
  SB.P.onset        = read(io, UInt8)
  SB.P.fm           = read(io, UInt8)
  SB.P.snr          = read(io, Int16)
  skip(io, 4)
  SB.P.gr           = read(io, Int16)
  SB.t_f64          = read(io, Float64)
  SB.P.amp          = read(io, Float32)
  SB.t_f32          = read(io, Float32)
  SB.t_i32          = read(io, Int32)
  skip(io, 4)

  phase = pick_types[SB.P.pc]
  (v > 1) && println(
            "id  = ",           SB.id_str,
            ", inst code = ",   Int16(SB.hdr[12]) << 8 | Int16(SB.hdr[11]),
            ", phase: ",        phase, Char(SB.P.onset), Char(SB.P.fm),
            ", t = ",           u2d(SB.t_f64),
            ", A = ",           SB.P.amp,
            ", τ = ",           SB.t_f32,
            ", SNR = ",         SB.P.snr,
            ", gain range = ",  SB.P.gr)

  if isa(S, EventTraceData)
    i = findid(SB.id_str, S)
    if i == 0
      push!(S, SeisChannel(id = SB.id_str))
      i = S.n
    end
    S.pha[i][phase] = SeisPha(amp = Float64(SB.P.amp),
                              tt = SB.t_f64,
                              pol = Char(SB.P.fm),
                              )
  end

  return nothing
end

#  ORIGIN: Information about a specific solution for a given event
function read_14!(S::GphysData, io::IO, v::Int64, full::Bool)
  SB.H.evno         = read(io, Int32)
  SB.H.auth         = read(io, Int16)
  read!(io, SB.H.chars)
  SB.H.reg          = read(io, Int32)
  SB.H.ot           = read(io, Float64)
  SB.H.lat          = read(io, Float64)
  SB.H.lon          = read(io, Float64)
  read!(io, SB.H.floats)
  read!(io, SB.H.model)
  SB.H.gap          = read(io, Int16)
  SB.H.d_min        = read(io, Float32)
  read!(io, SB.H.shorts)
  read!(io, SB.H.mag)
  SB.t_i32          = read(io, Int32)
  if v > 1
    println("evno = ",      SB.H.evno,
            ", auth = ",    SB.H.auth,
            ", codes = ",   map(Char, SB.H.chars),
            ", reg = ",     SB.H.reg,
            ", ot = ",      u2d(SB.H.ot),
            ", origin = ",  SB.H.lat, "N, ", SB.H.lon, "E, z ", SB.H.floats[1], " km",
            ", δx = ",      SB.H.floats[2], " km",
            ", δz = ",      SB.H.floats[3], " km",
            ", rms = ",     SB.H.floats[4],
            ", model = ",   String(copy(SB.H.model)),
            ", Δ = ",       SB.H.gap, "∘",
            ", d_min = ",   SB.H.d_min, " km",
            ", shorts = ",  SB.H.shorts,
            ", mag = ",     SB.H.mag,
            ", t_eff = ",   u2d(Float64(SB.t_i32))
            )
  end
  return nothing
end

# COMMENT:  Comment tag to be followed by the bytes of comment
function read_20!(S::GphysData, io::IO, v::Int64, full::Bool)
  read!(io, SB.comm_i)
  L = read(io, Int16)
  skip(io, 2)
  if v > 1
    println("structure ref. ID = ",   SB.comm_i[1],
            ", item in struct ID = ", SB.comm_i[2],
            ", L = ", L)
  end
  SB.comm_s = String(copy(read(io, L)))
  for i in 1:S.n
    if haskey(S.misc[i], "comment")
      append!(S.misc[i]["comment"], SB.comm_s)
    else
      S.misc[i]["comment"] = [SB.comm_s]
    end
  end
  (v > 2) && (printstyled("comment: \n", color=:green); println(SB.comm_s); printstyled("--\n", color=:green))
  return nothing
end

# TRIGGERS:  Earthquake detector trigger statistics
function read_25!(S::GphysData, io::IO, v::Int64, full::Bool)
  staident!(io, 0)
  shorts    = read!(io, zeros(Int16, 6))
  trig_time = read(io, Float64)
  if v > 1
    println("id = ",          SB.id_str,
            ", inst code = ", Int16(SB.hdr[12]) << 8 | Int16(SB.hdr[11]),
            ", shorts = ", shorts,
            ", btime = ", u2d(trig_time))
  end
  return nothing
end

# TRIGSETTING:  Settings for earthquake trigger system
function read_26!(S::GphysData, io::IO, v::Int64, full::Bool)
  net       = read!(io, zeros(UInt8, 4))
  btime     = read(io, Float64)
  shorts    = read!(io, zeros(Int16, 6))
  t_sweep   = read(io, Float32)
  t_aper    = read(io, Float32)
  alg       = read(io, UInt8)
  skip(io, 3)
  if v > 1
    println("net = ", String(copy(net)),
            ", btime = ", u2d(btime),
            ", shorts = ", shorts,
            ", t_sweep = ", t_sweep,
            ", t_aperture = ", t_aper,
            ", alg = ", Char(alg))
  end

  return nothing
end

# EVENTSETTING:  Settings for earthquake trigger system
function read_27!(S::GphysData, io::IO, v::Int64, full::Bool)
  net       = read!(io, zeros(UInt8, 4))
  btime     = read(io, Float64)
  shorts    = read!(io, zeros(Int16, 4))
  dur_min   = read(io, Float32)
  dur_max   = read(io, Float32)
  alg       = read(io, UInt8)
  skip(io, 3)
  if v > 1
    println("net = ", String(copy(net)),
            ", btime = ", u2d(btime),
            ", shorts = ", shorts,
            ", dur_min = ", dur_min,
            ", dur_max = ", dur_max,
            ", alg = ", Char(alg))
  end

  return nothing
end

# DETECTOR
function read_28!(S::GphysData, io::IO, v::Int64, full::Bool)
  alg           = read(io, UInt8)
  event_type    = read(io, UInt8)
  net_node_id   = read!(io, Array{UInt8,1}(undef, 10))
  version       = read(io, Float32)
  event_num     = read(io, Int32)
  skip(io, 4)

  if v > 1
    println("algorithm = ", Char(alg),
            ", event_type = ", Char(event_type),
            ", net_node_id = ", String(net_node_id),
            ", version = ", version,
            ", event_num = ", event_num
            )

  end
  return nothing
end

# ATODINFO
function read_29!(S::GphysData, io::IO, v::Int64, full::Bool)
  io_addr       = read(io, Int16)
  dev_id        = read(io, Int16)
  dev_flags     = read(io, UInt16)
  ext_bufs      = read(io, Int16)
  ext_mux       = read(io, Int16)
  timing_src    = read(io, UInt8)
  trig_src      = read(io, UInt8)

  if v > 1
    println("io_addr = ", io_addr,
            ", dev_id = ", dev_id,
            ", dev_flags = ", bitstring(dev_flags),
            ", ext_bufs = ", ext_bufs,
            ", ext_mux = ", ext_mux,
            ", timing_src = ", timing_src,
            ", trig_src = ", trig_src
            )

  end
  return nothing
end

# TIMECORRECTION:  Time correction information
function read_30!(S::GphysData, io::IO, v::Int64, full::Bool)
  staident!(io, 0)
  SB.tc          = read(io, Float64)
  SB.rc          = read(io, Float32)
  SB.sync_code   = Char(read(io, UInt8))
  skip(io,1)
  SB.t_i32       = read(io, Int32)
  skip(io, 2)

  if v > 1
    println("tc = ",          SB.tc,
            ", rc = ",        SB.rc,
            ", sync_code = ", SB.sync_code,
            ", t_eff = ",     SB.t_i32)
  end

  # Check: is the second part of ID "IRIG" (i.e., a local time clock)?
  SB.irig = false
  id_b = SB.id
  i1 = findfirst(SB.id.==0x2e)+1
  i2 = findnext(SB.id.==0x2e, i1+1)-1
  if SB.id[i1:i2] == [0x49, 0x52, 0x49, 0x47] # "IRIG"
    SB.irig = true
  end

  # return (round(Int64, tc*1.0e6), rc, t_eff*1000000, SB.id_str, netflg)
  return nothing
end

#= INSTRUMENT: Instrument hardware settings, mainly PADS related
                    added by R. Banfill, Jan 1991 =#
function read_31!(S::GphysData, io::IO, v::Int64, full::Bool)
  staident!(io, 0)
  serial    = read(io, Int16)
  ncmp      = read(io, Int16)
  chno      = read(io, Int16)
  sens_type = read(io, UInt8)
  data_type = read(io, UInt8)
  void_samp = read(io, Int32)
  floats    = read!(io, zeros(Float32,10))
  t_eff     = read(io, Int32)
  pre_evt   = read(io, Float32)
  trig_num  = read(io, Int16)
  study     = read(io, 6)
  sn_serial = read(io, Int16)
  if v > 1
    println("id = ",            SB.id_str,
            ", inst code = ",   Int16(SB.hdr[12]) << 8 | Int16(SB.hdr[11]),
            ", serial = ",      serial,
            ", ncmp = ",        ncmp,
            ", chno = ",        chno,
            ", sens_type = ",   Char(sens_type),
            ", data_type = ",   Char(data_type),
            ", void_samp = ",   void_samp,
            ", floats = ",      floats,
            ", t_eff = ",       u2d(t_eff),
            ", pre_evt = ",     pre_evt,
            ", trig_num = ",    trig_num,
            ", study = ",       String(study),
            ", sn_serial = ",   sn_serial
            )
  end
  return nothing
end

# CHANSET:  Associate stations and components with sets...???
function read_32!(S::GphysData, io::IO, v::Int64, full::Bool)
  SB.C.typ        = read(io, Int16)
  SB.C.n          = read(io, Int16)
  read!(io, SB.C.sta)
  skip(io, 1)
  SB.C.tu         = read(io, Int32)
  SB.C.td         = read(io, Int32)
  if v > 1
    println("type = ",    SB.C.typ,
            ", n = ",     SB.C.n,
            ", sta = ",   String(copy(SB.C.sta)),
            ", t_up = ",  u2d(SB.C.tu),
            ", t_dn = ",  u2d(SB.C.td)
            )
    (v > 2) && printstyled("channel entries:\n", color=:green)
  end

  for i = 1:SB.C.n
    read_chansetentry!(S, io, v)
  end

  return nothing
end

# # CHANSETENTRY
# read_33!(S::GphysData, io::IO, v::Int64, full::Bool) = read_chansetentry!(S, io, v)
#
function suds_support()
  println("\nCurrent support for SUDS structures\n")
  printstyled("CODE  STRUCTURE       \n", color=:green, bold=true)
  printstyled("====  =========       \n", color=:green, bold=true)
  for i in 1:33
    str = lpad(i, 4) * "  " * suds_codes[i] * "\n"
    if i in unsupported
      printstyled(str, color=:red)
    elseif i in disp_only
      printstyled(str, color=11)
    else
      printstyled(str, color=:green, bold=true)
    end
  end
  printstyled("\n(supported = ")
  printstyled("GREEN", color=:green, bold=true)
  printstyled(", logging only = ")
  printstyled("YELLOW", color=11)
  printstyled(", unsupported = ")
  printstyled("RED", color=:red)
  printstyled(")\n")
  return nothing
end


text:
#!/usr/bin/env julia
using Plots, Random, Printf, Plots.Measures

function initCondGen() #get random initial conditions for mass/radius, position, and velocity
    function getMass(nBodies) #generate random masses that better reflect actual stellar populations
        mList=zeros(nBodies)
        N=(0.5^(-1.3)-150^(-1.3))/1.3 #crude approximation of IMF integral assuming alpha = 2.3, stellar mass range of 0.5:150 solar masses
        rescale=1e6
        max=floor(Int,N*rescale)
        for i=1:nBodies
            intTarget=rand(0:max,1)[1]/rescale
            m=(0.5^(-1.3)-intTarget*1.3)^(-1/1.3) #just algebra from above
            mList[i]=round(m,digits=2)
        end
        return mList
    end
    m=rand(1:1500,3)./10 #3 random masses between 0.1 and 150 solar masses, uniform distribution
    #m=getMass(3) #get mass from IMF -- this way is kind of boring...so not using it, but left here in case I change my mind?
    rad=m.^0.8 #3 radii based on masses in solar units
    m=m.*2e30 #convert to SI kg
    rad=rad.*7e8 #convert to SI m
    pos1=rand(-10:10,3) #random initial coordinates x & y for first body, AU
    function genPos2(pos1)
        accept2=false
        while accept2==false
            pos2=rand(-10:10,3) #random initial coordinates for second body, AU
            dist21=sqrt((pos1[1]-pos2[1])^2+(pos1[2]-pos2[2])^2+(pos1[3]-pos2[3])^2)
            if (dist21*1.5e11)>(rad[1]+rad[2]) #they aren't touching
                accept2=true
                return pos2
            end
        end
    end
    pos2=genPos2(pos1)
    function genPos3(pos1,pos2)
        accept3=false
        while accept3==false
            pos3=rand(-10:10,3) #random initial coordinates for third body, AU
            dist31=sqrt((pos1[1]-pos3[1])^2+(pos1[2]-pos3[2])^2+(pos1[3]-pos3[3])^2)
            dist32=sqrt((pos2[1]-pos3[1])^2+(pos2[2]-pos3[2])^2+(pos2[3]-pos3[3])^2)
            if (dist31*1.5e11)>(rad[1]+rad[3]) && (dist32*1.5e11)>(rad[2]+rad[3]) #3rd isn't touching either
                accept3=true
                return pos3
            end
        end
    end
    pos3=genPos3(pos1,pos2)
    pos=[pos1[1],pos1[2],pos1[3],pos2[1],pos2[2],pos2[3],pos3[1],pos3[2],pos3[3]].*1.5e11 #convert accepted positions to SI, m
    v=rand(-7e3:7e3,9) #random xyz velocities with mag between -10 & 10 km/s, totally arbitrary...
    #r=[x1,y1,x2,y2,x3,y3,v1x,v1y,v2x,v2y,v3x,v3y]
    r=[pos[1],pos[2],pos[3],pos[4],pos[5],pos[6],pos[7],pos[8],pos[9],v[1],v[2],v[3],v[4],v[5],v[6],v[7],v[8],v[9]]
    open("initCond.txt","w") do f #save initial conditions to file in folder where script is run
        write(f,"m1=$(@sprintf("%.1f",(m[1]/2e30))) m2=$(@sprintf("%.1f",(m[2]/2e30))) m3=$(@sprintf("%.1f",(m[3]/2e30))) (solar masses)\nv1x=$(v[1]/1e3) v1y=$(v[2]/1e3) v1z=$(v[3]/1e3) v2x=$(v[4]/1e3) v2y=$(v[5]/1e3) v2z=$(v[6]/1e3) v3x=$(v[7]/1e3) v3y=$(v[8]/1e3) v3z=$(v[9]/1e3) (km/s)\nx1=$(pos1[1]) y1=$(pos1[2]) z1=$(pos1[3]) x2=$(pos2[1]) y2=$(pos2[2]) z2=$(pos2[3]) x3=$(pos3[1]) y3=$(pos3[2]) z3=$(pos3[3]) (AU from center)")
    end
    return r, rad, m
end

function dR(r,m) #function we will use RK4 on to approximate solution
    G=6.67408313131313e-11# Nm^2/kg^2
    M1,M2,M3=m[1],m[2],m[3] #kg
    x1,x2,x3=r[1],r[4],r[7] #m
    y1,y2,y3=r[2],r[5],r[8] #m
    z1,z2,z3=r[3],r[6],r[9]
    c1,c2,c3=G*M1,G*M2,G*M3 #Nm^2/kg
    r1_2=sqrt((x1-x2)^2+(y1-y2)^2+(z1-z2)^2) #distance from 1->2, m
    r1_3=sqrt((x1-x3)^2+(y1-y3)^2+(z1-z3)^2) #distance from 1->3, m
    r2_3=sqrt((x2-x3)^2+(y2-y3)^2+(z2-z3)^2) #distance from 2->3, m

    v1X,v2X,v3X=r[10],r[13],r[16] #these are our change in position after dt (dr/dt*dt=dr)
    v1Y,v2Y,v3Y=r[11],r[14],r[17] #m after * dt
    v1Z,v2Z,v3Z=r[12],r[15],r[18]

    #get change in velocity from accelerations (d^2r/dt^2*dt=dv/dt*dt=dv)
    dx1=-(c2*(x1-x2)/(r1_2^3))-(c3*(x1-x3)/(r1_3^3)) #d^2x/dt^2 for 1, m/s after * dt
    dx2=-(c1*(x2-x1)/(r1_2^3))-(c3*(x2-x3)/(r2_3^3)) #d^2x/dt^2 for 2, m/s
    dx3=-(c1*(x3-x1)/(r1_3^3))-(c2*(x3-x2)/(r2_3^3)) #d^2x/dt^2 for 3, m/s
    dy1=-(c2*(y1-y2)/(r1_2^3))-(c3*(y1-y3)/(r1_3^3)) #d^2y/dt^2 for 1, m/s
    dy2=-(c1*(y2-y1)/(r1_2^3))-(c3*(y2-y3)/(r2_3^3)) #d^2y/dt^2 for 2, m/s
    dy3=-(c1*(y3-y1)/(r1_3^3))-(c2*(y3-y2)/(r2_3^3)) #d^2y/dt^2 for 3, m/s
    dz1=-(c2*(z1-z2)/(r1_2^3))-(c3*(z1-z3)/(r1_3^3)) #d^2y/dt^2 for 1, m/s
    dz2=-(c1*(z2-z1)/(r1_2^3))-(c3*(z2-z3)/(r2_3^3)) #d^2y/dt^2 for 2, m/s
    dz3=-(c1*(z3-z1)/(r1_3^3))-(c2*(z3-z2)/(r2_3^3)) #d^2y/dt^2 for 3, m/s

    return [v1X,v1Y,v1Z,v2X,v2Y,v2Z,v3X,v3Y,v3Z,dx1,dy1,dz1,dx2,dy2,dz2,dx3,dy3,dz3]
end

function gen3Body(stopCond=[10,100],numSteps=10000) #default stop conditions of 10 yrs and 100 AU sep
    tStop=stopCond[1]*365*24*3600 #convert to SI s
    sepStop=stopCond[2]*1.5e11 #convert to SI m
    stop=false
    currentT=0
    t=range(0,stop=tStop,length=(numSteps+1)) #+1 because I don't want 0 to count
    stepSize=tStop/numSteps
    x1=zeros(length(t))
    y1=zeros(length(t))
    z1=zeros(length(t))
    x2=zeros(length(t))
    y2=zeros(length(t))
    z2=zeros(length(t))
    x3=zeros(length(t))
    y3=zeros(length(t))
    z3=zeros(length(t))
    r,rad,m=initCondGen()
    min12=rad[1]+rad[2]
    min13=rad[1]+rad[3]
    min23=rad[2]+rad[3]
    i=1
    stopT=maximum(t)
    collisionBool=false
    #implement RK4 to model solutions to differential equations
    while stop==false
        if currentT==stopT || currentT>stopT #in case of rounding error or something
            stop=true
        elseif i>numSteps+1 #inf loop failsafe
            stop=true
            println("error: shouldn't have gotten here")
        else
            x1[i]=r[1]
            y1[i]=r[2]
            z1[i]=r[3]
            x2[i]=r[4]
            y2[i]=r[5]
            z2[i]=r[6]
            x3[i]=r[7]
            y3[i]=r[8]
            z3[i]=r[9]

            k1=stepSize*dR(r,m)
            k2=stepSize*dR(r.+0.5.*k1,m)
            k3=stepSize*dR(r.+0.5.*k2,m)
            k4=stepSize*dR(r.+k3,m)
            r+=(k1.+2.0*k2.+2.0.*k3.+k4)./6

            #check separation after each dt step
            sep12=sqrt((x1[i]-x2[i])^2+(y1[i]-y2[i])^2+(z1[i]-z2[i])^2)
            sep13=sqrt((x1[i]-x3[i])^2+(y1[i]-y3[i])^2+(z1[i]-z3[i])^2)
            sep23=sqrt((x3[i]-x2[i])^2+(y3[i]-y2[i])^2+(z3[i]-z2[i])^2)

            if sep12<min12 || sep13<min13 || sep23<min23 || sep12>sepStop || sep13>sepStop || sep23>sepStop
                if sep12<min12 || sep13<min13 || sep23<min23
                    collisionBool=true
                else
                    collisionBool=false
                end
                stop=true #stop if collision happens or body is ejected
                t=range(0,stop=currentT,length=i) #t should match pos vectors
                x1=x1[1:i] #don't want trailing zeros
                y1=y1[1:i]
                z1=z1[1:i]
                x2=x2[1:i]
                y2=y2[1:i]
                z2=z2[1:i]
                x3=x3[1:i]
                y3=y3[1:i]
                z3=z3[1:i]
            end
            i+=1
            currentT+=stepSize #next step
        end
    end
    return [x1,y1,z1,x2,y2,z2,x3,y3,z3], t, m, rad, collisionBool
end

function getInteresting3Body(minTime=0) #in years, defaults to 0
    #sometimes random conditions result in a really short animation where things
    #just crash into each other/fly away, so this function throws away those
    yearSec=365*24*3600
    interesting=false
    i=1
    while interesting==false
        plotData,t,m,rad,collisionBool=gen3Body([60,150],600000)
        if (maximum(t)/yearSec)>minTime #only return if simulation runs for longer than minTime
            println(maximum(t)/yearSec) #tell me how many years we are simulating
            open("cron_log.txt","a") do f #for cron logging, a flag = append
                write(f,"$(maximum(t)/yearSec)\n")
            end
            return plotData,t,m,rad,collisionBool
            interesting=true
        elseif i>5000 #computationally expensive so don't want to go forever
            interesting=true #render it anyways I guess because sometimes it's fun?
            println("did not find interesting solution in number of tries allotted, running anyways")
            println(maximum(t)/yearSec) #how many years simulation runs for
            open("cron_log.txt","a") do f #for cron logging
                write(f,"$(maximum(t)/yearSec)\n")
            end
            return plotData,t,m,rad,collisionBool
        end
        i+=1
    end
end

function getLims(pos,padding) #determines plot limits at each frame, padding in units of pos
    x=[pos[1],pos[4],pos[7]]
    xMin=minimum(x)
    xMax=maximum(x)
    dx=xMax-xMin
    y=[pos[2],pos[5],pos[8]]
    yMin=minimum(y)
    yMax=maximum(y)
    dy=yMax-yMin
    z=[pos[3],pos[6],pos[9]]
    zMin=minimum(z)
    zMax=maximum(z)
    dz=zMax-zMin
    dList=[dx,dy,dz]
    if maximum(dList)==dx
        #use x for square
        xlims=[xMin-padding,xMax+padding]
        ylims=[yMin-padding,yMin+dx+padding]
        zlims=[zMin-padding,zMin+dx+padding]
    elseif maximum(dList)==dy
        #use y for square
        xlims=[xMin-padding,xMin+dy+padding]
        ylims=[yMin-padding,yMax+padding]
        zlims=[zMin-padding,zMin+dy+padding]
    else
        #use z for cube
        xlims=[xMin-padding,xMin+dz+padding]
        ylims=[yMin-padding,yMin+dz+padding]
        zlims=[zMin-padding,zMax+padding]
    end
    return xlims,ylims,zlims
end

function getColors(m,c) #places colors of objects according to mass/size
    #c=[:biggest,:medium,:smallest] (order of input colors)
    maxM=maximum(m)
    minM=minimum(m)
    colors=[:blue,:blue,:blue] #testing
    if m[1]==maxM
        colors[1]=c[1]
        if m[2]==minM
            colors[2]=c[3]
            colors[3]=c[2]
        else
            colors[3]=c[3]
            colors[2]=c[2]
        end
    elseif m[2]==maxM
        colors[2]=c[1]
        if m[1]==minM
            colors[1]=c[3]
            colors[3]=c[2]
        else
            colors[3]=c[3]
            colors[1]=c[2]
        end
    else
        colors[3]=c[1]
        if m[1]==minM
            colors[1]=c[3]
            colors[2]=c[2]
        else
            colors[2]=c[3]
            colors[1]=c[2]
        end
    end
    return colors
end

function makeCircleVals(r,center=[0,0])
    xOffset=center[1]
    yOffset=center[2]
    xVals=[r*cos(i)+xOffset for i=0:(pi/64):(2*pi)]
    yVals=[r*sin(i)+yOffset for i=0:(pi/64):(2*pi)]
    return xVals,yVals
end

plotData,t,m,rad,collisionBool=getInteresting3Body(15)
c=[:DodgerBlue,:Gold,:Tomato] #most massive to least massive, also roughly corresponds to temp
colors=getColors(m,c)
#adding fake stars
numStars=2500
starsX=zeros(numStars)
starsY=zeros(numStars)
starsZ=zeros(numStars)
for i=1:numStars
    num=rand(-200:200,3) #box size is 70 AU but we need some extra padding for movement
    starsX[i]=num[1]
    starsY[i]=num[2]
    starsZ[i]=num[3]
end
stars = [starsX,starsY,starsZ]
global frameNum=1
stop=length(t)
if collisionBool==true
    stop=length(t)-600
end

function makePanel(i1,i2,i,rad,plotData,stars,x,y,lims)
    p=plot(plotData[i1][1:33:i]./1.5e11,plotData[i2][1:33:i]./1.5e11,label="",linewidth=2,linecolor=colors[1],linealpha=max.((1:33:i) .+ 10000 .- i,2500)/10000) #plot orbits up to i
    p=plot!(plotData[i1+3][1:33:i]./1.5e11,plotData[i2+3][1:33:i]./1.5e11,label="",linewidth=2,linecolor=colors[2],linealpha=max.((1:33:i) .+ 10000 .- i,2500)/10000) #linealpha argument causes lines to decay
    p=plot!(plotData[i1+6][1:33:i]./1.5e11,plotData[i2+6][1:33:i]./1.5e11,label="",linewidth=2,linecolor=colors[3],linealpha=max.((1:33:i) .+ 10000 .- i,2500)/10000) #example: alpha=max.((1:i) .+ 100 .- i,0) causes only last 100 to be visible
    p=scatter!(stars[i1],stars[i2],markercolor=:white,markersize=:1,label="") #fake background stars
    star1=makeCircleVals(rad[1],[plotData[i1][i],plotData[i2][i]]) #generate circles with appropriate sizes for each star
    star2=makeCircleVals(rad[2],[plotData[i1+3][i],plotData[i2+3][i]]) #at current positions
    star3=makeCircleVals(rad[3],[plotData[i1+6][i],plotData[i2+6][i]])
    p=plot!(star1[1]./1.5e11,star1[2]./1.5e11,label="$(@sprintf("%.1f", m[1]./2e30))",color=colors[1],fill=true) #plot star circles with labels
    p=plot!(star2[1]./1.5e11,star2[2]./1.5e11,label="$(@sprintf("%.1f", m[2]./2e30))",color=colors[2],fill=true)
    p=plot!(star3[1]./1.5e11,star3[2]./1.5e11,label="$(@sprintf("%.1f", m[3]./2e30))",color=colors[3],fill=true)
    p=plot!(background_color=:black,background_color_legend=:transparent,foreground_color_legend=:transparent,
        background_color_outside=:white,aspect_ratio=:equal,legendtitlefontcolor=:white,legendfontfamily="Courier") #formatting for plot frame
    p=plot!(xlabel="$x: AU",ylabel="$y: AU",title="$x$y plane",
        legend=:best,xaxis=("$x: AU",(lims[i1][1],lims[i1][2]),font(9,"Courier")),yaxis=("$y: AU",(lims[i2][1],lims[i2][2]),font(9,"Courier")),
        grid=false,titlefont=font(14,"Courier"),size=(720,721),legendfontsize=8,legendtitle="Mass (in solar masses)",legendtitlefontsize=8,legendtitlefont="Courier") #add in axes/title/legend with formatting
    return p
end

for i=1:333:stop #this makes animation scale ~1 sec/year with other conditions
    GR.inline("png") #added to eneable cron/jobber compatibility, also this makes frames generate WAY faster? Prior to adding this when run from cron/jobber frames would stop generating at 408 for some reason.
    gr(legendfontcolor = plot_color(:white)) #legendfontcolor=:white plot arg broken right now (at least in this backend)
    print("$(@sprintf("%.2f",i/length(t)*100)) % complete\r") #output percent tracker
    pos=[plotData[1][i],plotData[2][i],plotData[3][i],plotData[4][i],plotData[5][i],plotData[6][i],plotData[7][i],plotData[8][i],plotData[9][i]] #current pos
    lims=getLims(pos./1.5e11,15) #convert to AU, 10 AU padding
    pXY=makePanel(1,2,i,rad,plotData,stars,"x","y",lims)
    pXZ=makePanel(1,3,i,rad,plotData,stars,"x","z",lims)
    pYZ=makePanel(2,3,i,rad,plotData,stars,"y","z",lims)
    title = plot(title = "Random Three-Body Problem: t = $(@sprintf("%0.2f",t[i]/365/24/3600)) years after start", grid = false, showaxis = false,titlefontcolor=:black,titlefont=(font(18,"Courier")),
    background_color=:white,foreground_color=:white,background_color_outside=:white,ticks=false,background_color_legend=:white,foreground_color_legend=:white,background_color_subplot=:white,background_color_inside=:white,foreground_color_subplot=:white)
    P = plot(title,pXY,pXZ,pYZ,layout=@layout([A{0.05h}; [B C D]]),size=(720*3,720),left_margin=5mm,bottom_margin=5mm)
    png(P,@sprintf("tmpPlots/frame_%06d.png",frameNum))
    global frameNum+=1
    closeall() #close plots
end
if collisionBool==true #this condition makes 2 seconds of slo-mo right before the collision
    println("making collision cam")
    for i=1:10:600
        GR.inline("png") #added to eneable cron/jobber compatibility, also this makes frames generate WAY faster? Prior to adding this when run from cron/jobber frames would stop generating at 408 for some reason.
        gr(legendfontcolor = plot_color(:white)) #legendfontcolor=:white plot arg broken right now (at least in this backend)
        print("$(@sprintf("%.2f",i/600*100)) % complete\r") #output percent tracker
        I = length(plotData[1])-(600-i)
        pos=[plotData[1][end-(600-i)],plotData[2][end-(600-i)],plotData[3][end-(600-i)],plotData[4][end-(600-i)],plotData[5][end-(600-i)],plotData[6][end-(600-i)],plotData[7][end-(600-i)],plotData[8][end-(600-i)],plotData[9][end-(600-i)]] #current pos
        lims=getLims(pos./1.5e11,15) #convert to AU, 15 AU padding
        pXY=makePanel(1,2,I,rad,plotData,stars,"x","y",lims)
        pXZ=makePanel(1,3,I,rad,plotData,stars,"x","z",lims)
        pYZ=makePanel(2,3,I,rad,plotData,stars,"y","z",lims)#add in axes/title/legend with formatting
        title = plot(title = "Random Three-Body Problem: t = $(@sprintf("%0.2f",t[i]/365/24/3600)) years after start", grid = false, showaxis = false,titlefontcolor=:black,titlefont=(font(18,"Courier")),
        background_color=:white,foreground_color=:white,background_color_outside=:white,ticks=false,background_color_legend=:white,foreground_color_legend=:white,background_color_subplot=:white,background_color_inside=:white,foreground_color_subplot=:white)
        P = plot(title,pXY,pXZ,pYZ,layout=@layout([A{0.05h}; [B C D]]),size=(720*3,720),left_margin=5mm,bottom_margin=5mm)
        P=annotate!((lims[1][1]+(lims[1][2]-lims[1][1])/40,lims[2][2]-(lims[2][2]-lims[2][1])/40,Plots.text("COLLISION CAM (slo-mo x 33)",10,"Courier",:orange,:left,:bold)))
        png(P,@sprintf("tmpPlots/frame_%06d.png",frameNum))
        global frameNum+=1
        closeall() #close plots
    end
end

run(`ffmpeg -y -framerate 30 -i "tmpPlots/frame_%06d.png" -c:v libx264 -preset slow -coder 1 -movflags +faststart -g 15 -crf 18 -pix_fmt yuv420p -profile:v high -y -bf 2 "/home/kirk/Documents/3Body/3Body3Frame_fps30.mp4"`)


text:
struct UnionFind
    parent::Vector{Int} # Parent index of each index, when -1, it is root
    size::Vector{Int} # group size of each index
    UnionFind(N) = new(fill(-1, N), fill(1, N))
end

function root!(uf::UnionFind, x::Int)
    if uf.parent[x] == -1
        return x
    else
        # Path compuression
        # Set the parent of the values traced from x as the root of x.
        return uf.parent[x] = root!(uf, uf.parent[x])
    end
end

function unite!(uf::UnionFind, x::Int, y::Int)
    x_root = root!(uf, x)
    y_root = root!(uf, y)
    x_root == y_root && return false  # these has already same root
    if uf.size[x_root] < uf.size[y_root] # union by size
        uf.parent[x_root] = y_root
        uf.size[y_root] += uf.size[x_root]
        uf.size[x_root] = uf.size[y_root]
    else
        uf.parent[y_root] = x_root
        uf.size[x_root] += uf.size[y_root]
        uf.size[y_root] = uf.size[x_root]
    end
    return true
end

function members(uf::UnionFind, x::Int)
    x_root = root!(uf, x)
    return [x for x in 1:length(uf.parent) if root!(uf, x) == x_root]
end

function all_group_members(uf::UnionFind)
    group_members = Dict{Int, Vector{Int}}()
    for x in 1:length(uf.parent)
        x_root = root!(uf, x)
        haskey(group_members, x_root) || (group_members[x_root] = Int[])
        push!(group_members[x_root], x)
    end
    return group_members
end

roots(uf::UnionFind)=[x for x in 1:length(uf.parent) if uf.parent[x] == -1]
issame(uf::UnionFind, x::Int, y::Int)=root!(uf, x) == root!(uf, y)
size(uf::UnionFind, x::Int)=uf.size[x]
group_size(uf::UnionFind)=length(roots(uf))


text:
using Documenter, HigherPrecision

makedocs(
    format = :html,
    sitename = "HigherPrecision.jl",
    #modules = [HigherPrecision],
    pages = [
        "Introduction" => "index.md",
        "Compute mode" => "computemode.md",
        "Reference" => "reference.md"
        ]
)

deploydocs(
    repo   = "github.com/saschatimme/HigherPrecision.jl.git",
    target = "build",
    julia = "0.6",
    osname = "linux",
    deps   = nothing,
    make   = nothing
)


text:
@defcomp co2cycle begin
    e_globalCO2emissions=Parameter(index=[time],unit="Mtonne/year")
    e0_globalCO2emissions=Parameter(unit="Mtonne/year", default=38191.0315797948)
    c_CO2concentration=Variable(index=[time],unit="ppbv")
    pic_preindustconcCO2=Parameter(unit="ppbv", default=278000.)
    exc_excessconcCO2=Variable(unit="ppbv")
    c0_CO2concbaseyr=Parameter(unit="ppbv", default=395000.)
    re_remainCO2=Variable(index=[time],unit="Mtonne")
    re_remainCO2base=Variable(unit="Mtonne")
    renoccf_remainCO2wocc=Variable(index=[time],unit="Mtonne")
    air_CO2fractioninatm=Parameter(unit="%", default=62.00)
    stay_fractionCO2emissionsinatm=Parameter(default=0.3)
    tea_CO2emissionstoatm=Variable(index=[time],unit="Mtonne/year")
    teay_CO2emissionstoatm=Variable(index=[time],unit="Mtonne/t")
    ccf_CO2feedback=Parameter(unit="%/degreeC", default=9.66666666666667)
    ccfmax_maxCO2feedback=Parameter(unit="%", default=53.3333333333333)
    cea_cumCO2emissionsatm=Variable(index=[time],unit="Mtonne")
    ce_0_basecumCO2emissions=Parameter(unit="Mtonne", default=2050000.)
    y_year=Parameter(index=[time],unit="year")
    y_year_0=Parameter(unit="year")
    res_CO2atmlifetime=Parameter(unit="year", default=73.3333333333333)
    den_CO2density=Parameter(unit="Mtonne/ppbv", default=7.8)
    rt_g0_baseglobaltemp=Parameter(unit="degreeC", default=0.735309967925382)
    rt_g_globaltemperature=Parameter(index=[time],unit="degreeC")

    function run_timestep(p, v, d, t)

        if is_first(t)
            #CO2 emissions gain calculated based on PAGE 2009
            gain=p.ccf_CO2feedback*p.rt_g0_baseglobaltemp
            #eq.6 from Hope (2006) - emissions to atmosphere depend on the sum of natural and anthropogenic emissions
            tea0=p.e0_globalCO2emissions*p.air_CO2fractioninatm/100
            v.tea_CO2emissionstoatm[t]=(p.e_globalCO2emissions[t])*p.air_CO2fractioninatm/100
            v.teay_CO2emissionstoatm[t]=(v.tea_CO2emissionstoatm[t]+tea0)*(p.y_year[t]-p.y_year_0)/2
            #adapted from eq.1 in Hope(2006) - calculate excess concentration in base year
            v.exc_excessconcCO2=p.c0_CO2concbaseyr-p.pic_preindustconcCO2
            #Eq. 2 from Hope (2006) - base-year remaining emissions
            v.re_remainCO2base=v.exc_excessconcCO2*p.den_CO2density
            #PAGE 2009 initial remaining emissions without CO2 feedback
            renoccf0_remainCO2wocc=v.re_remainCO2base/(1+gain/100)
            #eq. 8 from Hope (2006) - baseline cumulative emissions to atmosphere
            ceabase=p.ce_0_basecumCO2emissions*p.air_CO2fractioninatm/100
            #eq.9 from Hope(2006) - cumulative emissions in atmosphere
            v.cea_cumCO2emissionsatm[t]=ceabase+v.teay_CO2emissionstoatm[t]
            #eq.11 from Hope (2006) - anthropogenic remaining emissions
            v.renoccf_remainCO2wocc[t]=p.stay_fractionCO2emissionsinatm*ceabase*
                (1-exp(-(p.y_year[t]-p.y_year_0)/
                p.res_CO2atmlifetime))+renoccf0_remainCO2wocc*
                exp(-(p.y_year[t]-p.y_year_0)/p.res_CO2atmlifetime)+
                v.teay_CO2emissionstoatm[t]*exp(-(p.y_year[t]-p.y_year_0)/
                (2*p.res_CO2atmlifetime))
            #Hope 2009 - remaining emissions with CO2 feedback
            v.re_remainCO2[t]=v.renoccf_remainCO2wocc[t]*(1+gain/100)
        else
            #CO2 emissions gain calculated based on PAGE 2009
            gain=min(p.ccf_CO2feedback*p.rt_g_globaltemperature[t-1],p.ccfmax_maxCO2feedback)
            #eq.6 from Hope (2006) - emissions to atmosphere depend on the sum of natural and anthropogenic emissions
            v.tea_CO2emissionstoatm[t]=(p.e_globalCO2emissions[t])*p.air_CO2fractioninatm/100
            #eq.7 from Hope (2006) - total emissions over time period
            v.teay_CO2emissionstoatm[t]=(v.tea_CO2emissionstoatm[t]+v.tea_CO2emissionstoatm[t-1])*
                (p.y_year[t]-p.y_year[t-1])/2
            #eq.9 from Hope(2006) - cumulative emissions in atmosphere
            v.cea_cumCO2emissionsatm[t]=v.cea_cumCO2emissionsatm[t-1]+v.teay_CO2emissionstoatm[t]
            #eq.11 from Hope (2006) - anthropogenic remaining emissions
            v.renoccf_remainCO2wocc[t]=p.stay_fractionCO2emissionsinatm*v.cea_cumCO2emissionsatm[t-1]*
                (1-exp(-(p.y_year[t]-p.y_year[t-1])/
                p.res_CO2atmlifetime))+v.renoccf_remainCO2wocc[t-1]*
                exp(-(p.y_year[t]-p.y_year[t-1])/p.res_CO2atmlifetime)+
                v.teay_CO2emissionstoatm[t]*exp(-(p.y_year[t]-p.y_year[t-1])/
                (2*p.res_CO2atmlifetime))
            #Hope 2009 - remaining emissions with CO2 feedback
            v.re_remainCO2[t]=v.renoccf_remainCO2wocc[t]*(1+gain/100)
        end
        #eq.11 from Hope(2006) - CO2 concentration
        v.c_CO2concentration[t]=p.pic_preindustconcCO2+v.exc_excessconcCO2 * v.re_remainCO2[t]/v.re_remainCO2base
    end
end


text:
include("../src/TeukEvolution.jl")

import .TeukEvolution as TE

@time TE.launch(ARGS[1])


text:
using CapacityExpansion
using Clp
exact_scenarios=Dict{String,OptResult}()
for (state, years) in [["GER_1", [2016]],["CA_1", [2016]]]
       # laod data
       ts_input_data = load_timeseries_data_provided(state; T=24, years=years) #CEP
       cep_data = load_cep_data_provided(state)
       ## CLUSTERING ##
       ts_clust_data = run_clust(ts_input_data;method="hierarchical",representation="centroid",n_init=1,n_clust=3)
       ts_full_data = run_clust(ts_input_data;method="hierarchical",representation="centroid",n_init=1,n_clust=30)
       ## OPTIMIZATION ##
       optimizer=Clp.Optimizer
       optimizer_config=Dict{Symbol,Any}(:LogLevel => 0)
       exact_scenarios["$state-$years-co2"] = run_opt(ts_clust_data.clust_data,cep_data,optimizer;descriptor="co2",limit_emission=Dict{String,Number}("CO2/electricity"=>1000),optimizer_config=optimizer_config)
       exact_scenarios["$state-$years-slack"] = run_opt(ts_clust_data.clust_data,cep_data,optimizer;descriptor="slack",lost_load_cost=Dict{String,Number}("electricity"=>1e6), lost_emission_cost=Dict{String,Number}("CO2"=>700),optimizer_config=optimizer_config)
       exact_scenarios["$state-$years-ex"] = run_opt(ts_clust_data.clust_data,cep_data,optimizer;descriptor="ex",infrastructure=Dict{String,Array}("existing"=>["all"]),optimizer_config=optimizer_config)
       exact_scenarios["$state-$years-simple"] = run_opt(ts_clust_data.clust_data,cep_data,optimizer;descriptor="simple storage",storage_type="simple",conversion=true,optimizer_config=optimizer_config)
       exact_scenarios["$state-$years-seasonal"] = run_opt(ts_clust_data.clust_data,cep_data,optimizer;descriptor="seasonal storage",storage_type="seasonal",conversion=true,optimizer_config=optimizer_config)
       design_result=run_opt(ts_clust_data.clust_data,cep_data,optimizer;descriptor="des&op",optimizer_config=optimizer_config)
       exact_scenarios["$state-$years-des&op"] = run_opt(ts_full_data.clust_data,cep_data,design_result.config,get_cep_design_variables(design_result),optimizer;lost_load_cost=Dict{String,Number}("electricity"=>1e6), lost_emission_cost=Dict{String,Number}("CO2"=>700))
end
for (state, years) in [["GER_18", [2016]],["CA_14", [2016]]]
       # laod data
       ts_input_data = load_timeseries_data_provided(state; T=24, years=years) #CEP
       cep_data = load_cep_data_provided(state)
       ## CLUSTERING ##
       ts_clust_data = run_clust(ts_input_data;method="hierarchical",representation="centroid",n_init=1,n_clust=3)
       ## OPTIMIZATION ##
       optimizer=Clp.Optimizer
       optimizer_config=Dict{Symbol,Any}(:LogLevel => 0)
       exact_scenarios["$state-$years-trans"]=run_opt(ts_clust_data.clust_data,cep_data,optimizer;descriptor="trans",transmission=true,optimizer_config=optimizer_config)
end

@save normpath(joinpath(dirname(@__FILE__),"cep_exact_data.jld2")) exact_scenarios


text:
module BitTwiddlingConvenienceFunctions

using Static

intlog2(N::I) where {I <: Integer} = (8sizeof(I) - one(I) - leading_zeros(N)) % I
intlog2(::Type{T}) where {T} = intlog2(sizeof(T))
nextpow2(W::T) where {T<:Base.BitInteger} = (one(T) << (T(8sizeof(W)) - leading_zeros((W - one(T)))))
prevpow2(W::T) where {T<:Base.BitInteger} = (one(T) << (((T(8sizeof(W))) - one(T)) - leading_zeros(W)))
@generated nextpow2(::StaticInt{N}) where {N} = Expr(:call, Expr(:curly, :StaticInt, nextpow2(N)))
@generated prevpow2(::StaticInt{N}) where {N} = Expr(:call, Expr(:curly, :StaticInt, prevpow2(N)))
@generated intlog2(::StaticInt{N}) where {N} = Expr(:call, Expr(:curly, :StaticInt, intlog2(N)))

end


text:
"""Functions for computing radio models."""

abstract type AbstractRadioModel{T<:Real} <: AbstractForegroundModel end

"""
    Radio_Sehgal2009{T}(model parameters...)

Define CIB model parameters. Defaults are from Viero et al. 2013.

```@example
model = CIBModel{Float32}(a_0=0.4)
```
"""
Base.@kwdef struct Radio_Sehgal2009{T<:Real} <: AbstractRadioModel{T}
    nside::Int64    = 4096
    min_redshift::T = 0.0
    max_redshift::T = 5.0
    min_mass::T     = 1e13
    box_size::T     = 40000

    # these coefficients are shared for type I and II
    a_0::T   = -1.0
    a_1_dist::Distributions.Uniform{T} = Distributions.Uniform(T(-0.12), T(0.07))
    a_2_dist::Distributions.Uniform{T} = Distributions.Uniform(T(-0.34), T(0.99))
    a_3_dist::Distributions.Uniform{T} = Distributions.Uniform(T(-0.75), T(-0.23))

    I_R_int::T = 10^(-2.6)
    I_γ::T = 6.0
    I_N_0::T = 1.0
    I_M_0::T = 4e13
    I_α::T = 0.1
    I_L_b::T = 10^(24.0)
    I_m::T = -1.55
    I_n::T = 0.0
    I_δ::T = 3.0
    I_z_p::T = 0.8

    II_R_int::T = 10^(-2.8)
    II_γ::T = 8.0
    II_N_0::T = 0.015
    II_M_0::T = 3e15
    II_α::T = 0.1
    II_L_b::T = 10^(27.5)
    II_m::T = -1.6
    II_n::T = -0.65
    II_z_p::T = 1.3
    II_σ_l::T = 0.4
    II_σ_h::T = 0.73

    I_L_min::T = 2.5e23
    II_L_min::T = 4e23

    # physical constants
    phys_h::T     = 6.62606957e-27      # erg.s
    phys_c::T     = 3e+10               # cm/s
    phys_k::T     = 1.3806488e-16       # erg/K
    phys_Msun::T  = 2e33                # g
    phys_Mpc::T   = 3.086e24            # cm
end

pdf_norm1(x, μ, σ) = exp( -(x-μ)^2 / (2 * σ^2) );

function FR_I_redshift_evolution(z, model)
    if z < model.I_z_p
        return (1+z)^model.I_δ
    else
        return (1+model.I_z_p)^model.I_δ
    end
end

function FR_II_redshift_evolution(z, model)
    norm = pdf_norm1(0, model.II_z_p, model.II_σ_l)
    if z < model.II_z_p
        return pdf_norm1(z, model.II_z_p, model.II_σ_l) / norm
    else
        return pdf_norm1(z, model.II_z_p, model.II_σ_h) / norm
    end
end

"""
Populate halos with radio sources according to the HOD in Sehgal et al. 2009.

The optional rng parameter provides an array of random number generators, one
for each thread.
"""
function hod_sehgal(
        halo_mass, redshift, model::Radio_Sehgal2009{T}) where T

    N_halos = size(halo_mass,1)
    nsources_I = Array{Int32}(undef, N_halos)
    nsources_II = Array{Int32}(undef, N_halos)

    # compute poisson mean, then draw it
    Threads.@threads for i = 1:N_halos
        I_HON = model.I_N_0 * (halo_mass[i] / model.I_M_0)^model.I_α
        I_HON *= FR_I_redshift_evolution(redshift[i], model)
        nsources_I[i] = rand(Distributions.Poisson(Float64.(I_HON)))
        II_HON = model.II_N_0 * (halo_mass[i] / model.II_M_0)^model.II_α
        II_HON *= FR_II_redshift_evolution(redshift[i], model)
        nsources_II[i] = rand(Distributions.Poisson(Float64.(II_HON)))
    end

    return nsources_I, nsources_II
end

function sehgal_LFn0(rand_unit::T, m, Lb, Lmin) where T <: Real
    crossover = log( Lb/Lmin )
    x = rand_unit * (-(1/m)+log(Lb/Lmin))
    if x < crossover
        return exp(x) * Lmin
    else
        return Lb * (1+m * x + m * log(Lmin/Lb))^(1/m)
    end
end

function sehgal_LF(rand_unit::T, m, n, Lb, Lmin) where T
    crossover = (T(1) - Lb^(-n) * Lmin^n)/n
    x = (rand_unit * (m-Lb^(-n) * Lmin^n * m-n)/(m * n))
    if x < crossover
        return (Lmin^n + Lb^n * n * x)^(1/n)
    else
        return Lb * (1 + (m * (-1 + (Lmin/Lb)^n + n * x))/n)^(1/m)
    end
end

"""
Fills the result array with draws from the luminosity function.
"""
function sehgal_LF!(result::Array{T,1}, m, n, Lb::T, Lmin,
    rand_array::Array{T,1}) where T

    if n ≈ 0.0
        Threads.@threads for i = 1:size(result,1)
            result[i] = sehgal_LFn0(
                rand_array[i],
                Float64(m), Float64(Lb), Float64(Lmin))
        end
    else
        Threads.@threads for i = 1:size(result,1)
            result[i] = sehgal_LF(
                rand_array[i],
                Float64(m), Float64(n), Float64(Lb), Float64(Lmin))
        end
    end
end

B(cosθ, β) = ( (1-β*cosθ)^(-2) + (1+β*cosθ)^(-2) ) / 2

function get_core_lobe_lum(L_beam::T, ν_Hz, R_int, γ, a_0, a_1, a_2, a_3, cosθ) where T
    β = sqrt(1-γ^(-2))
    R_obs = R_int * B(cosθ, β)
    L_int = L_beam * (1 + R_int) / (1 + R_obs)
    L_l_int = L_int / (1 + R_int)
    L_c_beam = R_obs * L_l_int

    # now put in the frequency dependence. a_0 = 0.0
    lν = log10(ν_Hz / T(1e9) ) # normed to 151 MHz
    f_core =  T(10) ^ ( a_1 * lν + a_2 * lν^2 + a_3 * lν^3 ) # a_0 +

    lν_norm = T(log10(151e6 / 1e9 )) # normed to 151 MHz
    norm_core =  10 ^ ( a_1 * lν_norm + a_2 * lν_norm^2 + a_3 * lν_norm^3 )

    f_lobe = (ν_Hz / T(151e6) ).^T(-0.8)
    return T(L_c_beam * f_core * 0.3f0 / norm_core), T(L_l_int * f_lobe)
end

function draw_coeff!(a_coeff, model::Radio_Sehgal2009)
    a1_task = Threads.@spawn rand!(model.a_1_dist, @view a_coeff[1,:])
    a2_task = Threads.@spawn rand!(model.a_2_dist, @view a_coeff[2,:])
    a3_task = Threads.@spawn rand!(model.a_3_dist, @view a_coeff[3,:])

    wait(a1_task)
    wait(a2_task)
    wait(a3_task)
end

"""
Produce a source catalog from a model and halo catalog.
"""
function generate_sources(
        # model parameters
        model::AbstractRadioModel{T}, cosmo::Cosmology.FlatLCDM{T},
        # halo arrays
        halo_pos_inp::AbstractArray{TH,2}, halo_mass_inp::AbstractArray{TH,1};
        verbose=true) where {T, TH}

    # make sure halo inputs are the model type
    halo_pos = convert(Array{T,2}, halo_pos_inp)
    halo_mass = convert(Array{T,1}, halo_mass_inp)

    res = Resolution(model.nside)

    verbose && println("Culling halos below mass $(model.min_mass).")
    mass_cut = halo_mass .> model.min_mass
    halo_pos = halo_pos[:, mass_cut]
    halo_mass = halo_mass[mass_cut]
    N_halos = size(halo_mass, 1)

    verbose && println("Allocating for $(N_halos) halos.")
    dist, redshift, hp_ind = get_basic_halo_properties(halo_pos, model, cosmo, res)
    θ, ϕ = get_angles(halo_pos)

    verbose && println("Populating HOD.")
    nsources_I, nsources_II = hod_sehgal(
        halo_mass, redshift, model)

     # total number of sources
    total_n_I, total_n_II = sum(nsources_I), sum(nsources_II)

    # the last element of cumsat_I is the total number of type I sources
    verbose && println("Drawing spectral coefficients.")
    a_coeff_I = Array{T, 2}(undef, 3, total_n_I)
    a_coeff_II = Array{T, 2}(undef, 3, total_n_II)
    draw_coeff!(a_coeff_I, model)
    draw_coeff!(a_coeff_II, model)

    verbose && println("Drawing from luminosity function.")
    # generate random numbers for LF
    rand_buffer_I = Array{T, 1}(undef, total_n_I)
    rand_buffer_II = Array{T, 1}(undef, total_n_II)
    threaded_rand!(rand_buffer_I)
    threaded_rand!(rand_buffer_II)

    # draw luminosities for both populations
    L_I_151 = Array{T, 1}(undef, total_n_I)
    L_II_151 = Array{T, 1}(undef, total_n_II)
    sehgal_LF!(L_I_151, model.I_m, model.I_n, model.I_L_b,
        model.I_L_min, rand_buffer_I)
    sehgal_LF!(L_II_151, model.II_m, model.II_n, model.II_L_b,
        model.II_L_min, rand_buffer_II)

    verbose && println("Drawing for impact parameter.")
    # reuse rand buffers for the impact parameters
    threaded_rand!(rand_buffer_I)
    threaded_rand!(rand_buffer_II)
    cosθ_I = rand_buffer_I
    cosθ_II = rand_buffer_II

    return (
        redshift=redshift, halo_mass=halo_mass, halo_pos=halo_pos,
        halo_hp_ind=hp_ind, θ=θ, ϕ=ϕ,
        dist=dist, nsources_I=nsources_I, nsources_II=nsources_II,
        a_coeff_I=a_coeff_I, a_coeff_II=a_coeff_II,
        L_I_151=L_I_151, L_II_151=L_II_151,
        cosθ_I=cosθ_I, cosθ_II=cosθ_II,
        total_n_I=total_n_I, total_n_II=total_n_II
    )
end

function paint!(result_map::HealpixMap{T_map,RingOrder},
                nu_obs::T, model::AbstractRadioModel, sources;
                return_fluxes=false) where {T_map, T}

    pixel_array = result_map.pixels
    fill!(pixel_array, zero(T))  # prepare the frequency map

    flux_to_Jy = ustrip(u"Jy", 1u"W/Hz*Mpc^-2")

    source_offset_I = generate_subhalo_offsets(sources.nsources_I)
    source_offset_II = generate_subhalo_offsets(sources.nsources_II)
    N_halo = size(sources.halo_mass, 1)
    total_n_sat_I = size(sources.L_I_151, 1)
    total_n_sat_II = size(sources.L_II_151, 1)

    flux_I = Array{T, 1}(undef, total_n_sat_I)
    flux_II = Array{T, 1}(undef, total_n_sat_II)
    redshift_I = Array{T, 1}(undef, total_n_sat_I)
    redshift_II = Array{T, 1}(undef, total_n_sat_II)
    θ_I = Array{T, 1}(undef, total_n_sat_I)
    θ_II = Array{T, 1}(undef, total_n_sat_II)
    ϕ_I = Array{T, 1}(undef, total_n_sat_I)
    ϕ_II = Array{T, 1}(undef, total_n_sat_II)

    Threads.@threads for i_halo in 1:N_halo
        nu = (1 + sources.redshift[i_halo]) * nu_obs
        hp_ind = sources.halo_hp_ind[i_halo]

        # Generate FR I fluxes
        for source_index in 1:sources.nsources_I[i_halo]
            # index of satellite in satellite arrays
            i_sat = source_offset_I[i_halo] + source_index
            L_c, L_l = get_core_lobe_lum(
                sources.L_I_151[i_sat], nu, model.I_R_int, model.I_γ,
                model.a_0,
                sources.a_coeff_I[1,i_sat], sources.a_coeff_I[2,i_sat],
                sources.a_coeff_I[3,i_sat], sources.cosθ_I[i_sat])

            flux_I[i_sat] = l2f(L_c+L_l, sources.dist[i_halo],
                sources.redshift[i_halo]) * flux_to_Jy
            redshift_I[i_sat] = sources.redshift[i_halo]
            θ_I[i_sat] = sources.θ[i_halo]
            ϕ_I[i_sat] = sources.ϕ[i_halo]
            pixel_array[hp_ind] += flux_I[i_sat]
        end

        # Generate FR II fluxes
        for source_index in 1:sources.nsources_II[i_halo]
            # index of satellite in satellite arrays
            i_sat = source_offset_II[i_halo] + source_index
            L_c, L_l = get_core_lobe_lum(
                sources.L_II_151[i_sat], nu, model.II_R_int, model.II_γ,
                model.a_0,
                sources.a_coeff_II[1,i_sat], sources.a_coeff_II[2,i_sat],
                sources.a_coeff_II[3,i_sat], sources.cosθ_II[i_sat])

            flux_II[i_sat] = l2f(L_c+L_l, sources.dist[i_halo],
                sources.redshift[i_halo]) * flux_to_Jy
            redshift_II[i_sat] = sources.redshift[i_halo]
            θ_II[i_sat] = sources.θ[i_halo]
            ϕ_II[i_sat] = sources.ϕ[i_halo]
            pixel_array[hp_ind] += flux_II[i_sat]
        end
    end

    # divide by healpix pixel size
    per_pixel_steradian = 1 / nside2pixarea(result_map.resolution.nside)
    pixel_array .*= per_pixel_steradian

    # maps are in Jansky per steradian, fluxes are in Jansky
    if return_fluxes
        return flux_I, redshift_I, θ_I, ϕ_I, flux_II, redshift_II, θ_II, ϕ_II
    end

end

max_stars_repo_name: xzackli/XGPaint

max_issues_repo_name: xzackli/XGPaint

max_forks_repo_name: xzackli/XGPaint


text:
"""
    Reprint(fn) -- apply the lambda function when printed
"""
mutable struct Reprint
    content::Function
end

Base.print(io::IO, r::Reprint) = r.content(io)

"""
    Render(data) - printed object shows its text/html
"""
struct Render{T}
    content::T
end

Base.print(io::IO, r::Render) =
    show(io, MIME"text/html"(), r.content)

"""
    Bypass(data) - printed object passes though EscapeProxy unescaped
"""

mutable struct Bypass{T}
    content::T
end

Base.print(io::IO, x::Bypass) = print(io, x.content)

"""
    EscapeProxy(io) - wrap an `io` to perform HTML escaping

This is a transparent proxy that performs HTML escaping so that objects
that are printed are properly converted into valid HTML values. As a
special case, objects wrapped with `Bypass` are not escaped, and
bypass the proxy.

# Examples
```julia-repl
julia> ep = EscapeProxy(stdout);
julia> print(ep, "A&B")
A&B
julia> print(ep, Bypass("<tag/>"))
<tag/>
```
"""
struct EscapeProxy{T<:IO} <: IO
    io::T
end

EscapeProxy(io::EscapeProxy) = io

Base.print(ep::EscapeProxy, h::Reprint) = h.content(ep)
Base.print(ep::EscapeProxy, w::Render) =
    show(ep.io, MIME"text/html"(), w.content)
Base.print(ep::EscapeProxy, x::Bypass) = print(ep.io, x.content)

function Base.write(ep::EscapeProxy, octet::UInt8)
    if octet == Int('&')
        write(ep.io, "&")
    elseif octet == Int('<')
        write(ep.io, "<")
    elseif octet == Int('"')
        write(ep.io, """)
    elseif octet == Int('\'')
        write(ep.io, "'")
    else
        write(ep.io, octet)
    end
end

function Base.unsafe_write(ep::EscapeProxy, input::Ptr{UInt8}, nbytes::UInt)
    written = 0
    last = cursor = input
    final = input + nbytes
    while cursor < final
        ch = unsafe_load(cursor)
        if ch == Int('&')
            written += unsafe_write(ep.io, last, cursor - last)
            written += unsafe_write(ep.io, pointer("&"), 5)
            cursor += 1
            last = cursor
            continue
        end
        if ch == Int('<')
            written += unsafe_write(ep.io, last, cursor - last)
            written += unsafe_write(ep.io, pointer("<"), 4)
            cursor += 1
            last = cursor
            continue
        end
        if ch == Int('\'')
            written += unsafe_write(ep.io, last, cursor - last)
            written += unsafe_write(ep.io, pointer("'"), 6)
            cursor += 1
            last = cursor
            continue
        end
        if ch == Int('"')
            written += unsafe_write(ep.io, last, cursor - last)
            written += unsafe_write(ep.io, pointer("""), 6)
            cursor += 1
            last = cursor
            continue
        end
        cursor += 1
    end
    if last < final
        written += unsafe_write(ep.io, last, final - last)
    end
    return written
end


text:
## General purpose functions

#User is expected to pass functions
#x_k(p0, alpha, mhatk)  : provides a solution with data mhatk
#c_ik(x)
#These are often collected into arrays of functions
# xs			: [x_k for k = 1:K ]
# c_k  			: [c_ik for i = 1:d ]
# cs			: (d, K) Matrix with elements c_ik

#General purpose functions
###
function shrink(phat_k, p0, alpha, Nhat_k)
	if Nhat_k <= 0
		return p0
	end
	@. alpha * p0 / (Nhat_k + alpha) + Nhat_k * phat_k / (Nhat_k + alpha)
end

#need to be careful because some problem smay have zero data
function get_GM_anchor(mhats)
	Nhats = vec(sum(mhats, dims=1))
	non_zero_indx = Nhats .> 0
	vec(mean(mhats[:, non_zero_indx] ./ Nhats[non_zero_indx]', dims=2))
end

function sim_path(p_k, N::Int)
	@assert N > 0 "Something weird"
    mhats_k = zeros(length(p_k))
    for j in rand(Categorical(p_k), N)
        mhats_k[j] += 1
    end
    mhats_k
end

#ps is a matrix of pk
function sim_path(ps, Nhats)
    K = size(ps, 2)
    out = zeros(size(ps))
    for k = 1:K
        out[:, k] = sim_path(ps[:, k], Nhats[k])
    end
    out
end

#empirical bayes moment-matched estimates.
#returns p0, alpha0
#VG this method does not seem to consistenty yield alpha in [0, 1]
function eb_mm_estimates(mhats)
	Nhats = sum(mhats, dims=1)
	p0 = mean(mhats ./ Nhats, dims=2)
	K = size(mhats, 2)

	C0 = sum(mhats.^2)/K
	sq_norm_p0 = norm(p0)^2
	Nbar = mean(Nhats)
	var_n = mean( @.(Nhats^2 - Nhats) )

	alpha0 = C0 - sq_norm_p0 * var_n - Nbar
	@assert alpha0 > 0 "Moment matching fails"
	alpha0 /= (1 - sq_norm_p0) / var_n
	alpha0 = 1 / alpha0 - 1

	return p0, alpha0
end

#solves an approximation to the Apriori MSE equation for alpha
#for large K this is like AlphaOR for MSE
#for now, super lazy and just search a grid.
#returns AlphaEstimate, minimizingAlphaIndex
function mse_estimates(mhats, supps, p0, alpha_grid)
	K = size(mhats, 2)
	Nhats = sum(mhats, dims=1)
	phats = mhats ./ Nhats
	mu0s = vec(p0' * supps)

	#summary statistics
	muhats = Vector{Float64}(undef, K)
	sigmas = Vector{Float64}(undef, K)
	mse0 = Vector{Float64}(undef, K)
	for k = 1:K
		muhats[k] = dot(phats[:, k], supps[:, k])
		sigmas[k] = dot(phats[:, k], supps[:, k].^2) - muhats[k]^2
		mse0[k] = dot(phats[:, k], @.((supps[:, k] - mu0s[k])^2) )
	end

	#worker function
	function mse(a)
		out = 0.
		for k = 1:K
			out += (Nhats[k] - a^2) / (Nhats[k] + a)^2 * sigmas[k]
			out += ( a / (a + Nhats[k]) )^2 * mse0[k]
		end
		out/K
	end
	out = map(mse, alpha_grid)
	alpha_grid[argmin(out)], argmin(out)
end

function z_k(x_k, c_k, p0, alpha, mhat_k, ps_k, lam_k, lamavg)
	x = x_k(p0, alpha, mhat_k)
	cs = map(c_ik -> c_ik(x), c_k)
	lam_k/lamavg * dot(ps_k, cs)
end

#actually returns zLOO_k * N * lamavg
function zLOO_k_unsc(x_k, c_k, p0, alpha, mhat_k)
	out = 0.
	mhatloo = mhat_k[:] #copy data
	#only compute for terms where mhat_k[i] > 0
	#compute a base solution to resuse as necessary
	x_base = x_k(p0, alpha, mhat_k)
	for i = 1:length(mhat_k)
		#correct previous toggle
		if i > 1 && mhat_k[i - 1] > 0
			mhatloo[i - 1] += 1
		end

		if mhat_k[i] > 0
			mhatloo[i] -= 1
		    x = x_k(p0, alpha, mhatloo)
		else
			x = x_base
		end
		out += mhat_k[i] * c_k[i](x)
	end
	out
end

#average true performance
function zbar(xs, cs, p0, alpha, mhats, ps, lams)
	lamavg = mean(lams)
	K = size(cs, 2)
	out = 0.
	#VG maybe change this to a more numerically stable way to do avg?
	for k = 1:K
		out += z_k(xs[k], cs[:, k], p0, alpha, mhats[:, k], ps[:, k], lams[k], lamavg)
	end
	out/K
end

#full-info for scaling/comparison
#uses the fact that data scale doesn't matter.
zstar(xs, cs, ps, lams) = zbar(xs, cs, ps[:, 1], 0., ps, ps, lams)

function zLOObar_unsc(xs, cs, p0, alpha, mhats)
	K = size(cs, 2)
	out = 0.
	for k = 1:K
		out += zLOO_k_unsc(xs[k], cs[:, k], p0, alpha, mhats[:, k])
	end
	out /K
end

#returns alphaOR, minimizingAlphaIndex, curveInAlpha
function oracle_alpha(xs, cs, mhats, ps, lams, p0, alpha_grid)
	alphaOR = 0.
	jstar = -1
	best_val = Inf
	out = zeros(length(alpha_grid))
	for (j, alpha) in enumerate(alpha_grid)
		out[j] = zbar(xs, cs, p0, alpha, mhats, ps, lams)
		if out[j] < best_val
			jstar = j
			alphaOR = alpha
			best_val = out[j]
		end
	end
	return alphaOR, jstar, out
end

#returns alphaLOO, jstar, and cuveInAlpha_unsc
function loo_alpha(xs, cs, mhats, p0, alpha_grid)
	alphaLOO = 0.
	jstar = -1
	best_val = Inf
	out = zeros(length(alpha_grid))
	for (j, alpha) in enumerate(alpha_grid)
		out[j] = zLOObar_unsc(xs, cs, p0, alpha, mhats)
		if out[j] < best_val
			jstar = j
			alphaLOO = alpha
			best_val = out[j]
		end
	end
	return alphaLOO, jstar, out
end

###
#### Deprecated
#Convenience functions to generate a sequence of newsvendor problems
#  with common support, potentially different service levels
#  assumes supp is ordered vector of Reals
# returns cs, xs
# function genNewsvendors(supp, ss, K)
# 	@assert issorted(supp) "Supp must be an ordered vector"

# 	#Generic computation of the sth quantile
# 	function x_k(p0, alpha, mhat_k, s)
# 	    const Nhat_k = sum(mhat_k)
# 	    palpha = JS.shrink(mhat_k./Nhat_k, p0, alpha, Nhat_k)
# 	    indx = quantile(Categorical(palpha), s)
# 	    supp[indx]
# 	end

# 	function c_ik(i, x, s)
# 	    supp[i] > x ? s/(1-s) * (supp[i] - x) : (x - supp[i])
# 	end

# 	xs  = [(p0, alpha, mhat_k)-> x_k(p0, alpha, mhat_k, ss[k]) for k = 1:K]
# 	cs = [x->c_ik(i, x, ss[k]) for i = 1:length(supp), k = 1:K]
# 	cs, xs
# end

###
# Ideally should have a broadcast implementation that combines
# previous two...
function genNewsvendorsDiffSupp(supps, s, K)
	#Generic computation of the sth quantile
	function x_k(p0, alpha, mhat_k, k, s)
	    Nhat_k = sum(mhat_k)
	    palpha = JS.shrink(mhat_k./Nhat_k, p0, alpha, Nhat_k)
	    if !isapprox(sum(palpha), 1) || sum(mhat_k .< 0) > 0
	    	println("Nhat_k:\t", Nhat_k)
	    	println("alpha:\t", alpha)
	    	println("mhat_k:\n", mhat_k)
	    	println("phat_k:\n", palpha)
	    end

	    indx = quantile(Categorical(palpha), s)
	    supps[indx, k]
	end

	function c_ik(i, k, x, s)
	    supps[i, k] > x ? s/(1-s) * (supps[i, k] - x) : (x - supps[i, k])
	end

	xs  = [(p0, alpha, mhat_k)-> x_k(p0, alpha, mhat_k, k, s) for k = 1:K]
	cs = [x->c_ik(i, k, x, s) for i = 1:size(supps, 1), k = 1:K]
	cs, xs
end

#### Helper function for binning data
#dat_vec is vector of actual realizations
#computes a grid of size d
#returns
	#d vector of lefthand binpoints
	#d vector of counts
	#raw histogram object for safety
#assumes NA's already filtered from calcualtion
function bin_data(dat_vec, d; TOL = .001)
	#Explicitly make bins to prevent histogram from rounding
	l = minimum(dat_vec)
	u = maximum(dat_vec)
	bin_width = (u - l)/d
    bin_edges = collect(range(l, stop=u, length=d + 1))
    bin_edges[end] += TOL * bin_width  #inflate the last one a little...
    bin_edges[1] -= TOL * bin_width
    dat_hist = fit(Histogram, dat_vec, bin_edges, closed =:left)
    bin_edges[1:d], dat_hist.weights, dat_hist
end

max_stars_repo_name: CausalML/DataPooling

max_issues_repo_name: CausalML/DataPooling

max_forks_repo_name: CausalML/DataPooling


text:
"""
	msfm()

	Calculate the shortest distance from a list of
	points to all other pixels in an image, using the
	Multistencil Fast Marching Method (MSFM).

	This method gives more accurate
	distances by using second order derivatives and cross neighbours.

	T=msfm2d(F, SourcePoints, UseSecond, UseCross)

	# inputs,
	  F: The speed image. The speed function must always be larger
				than zero (min value 1e-8), otherwise some regions will
				never be reached because the time will go to infinity.
	  SourcePoints : A list of starting points [2 x N] (distance zero)
	  UseSecond : Boolean Set to true if not only first but also second
	               order derivatives are used (default)
	  UseCross : Boolean Set to true if also cross neighbours
	               are used (default)
	# outputs
	  T : Image with distance from SourcePoints to all pixels
"""
function msfm(
    speedimage::AbstractArray{T,2},
    SourcePoints::AbstractArray{T},
    usesecond::Bool=true,
    usecross::Bool=true,
    Ed=false::Bool,
) where {T}
    distanceimage = zeros(T, size(speedimage))

    Frozen = falses(size(speedimage)) # Pixels which are processed and have a final distance are frozen

    msfm!(distanceimage, Frozen, speedimage, SourcePoints, usesecond, usecross, Ed)
    return distanceimage
end

function msfm!(
    distanceimage,
    Frozen,
    speedimage::AbstractArray{T,2},
    SourcePoints::AbstractArray{T},
    usesecond::Bool=true,
    usecross::Bool=true,
    Ed=false::Bool,
) where {T}
    fill!(distanceimage, T(-Inf))
    # Augmented Fast Marching (For skeletonize)

    # Euclidian distance image
    if Ed
        Y = zeros(T, size(speedimage))
    end

    # Store neighbours of the (segmented) region
    neg_free = length(speedimage)
    neg_pos = 0
    neg_list_1 = zeros(T, neg_free)
    neg_list_x = fill(1, neg_free)
    neg_list_y = fill(1, neg_free)
    if Ed
        neg_list_4 = zeros(T, neg_free)
    end

    # (There are 3 pixel classes:
    #   - frozen (processed)
    #   - narrow band (boundary) (in list to check for the next pixel with smallest distance)
    #   - far (not yet used)

    # Neighbours
    ne = [
        -1 0
        1 0
        0 -1
        0 1
    ]
    Tt = T(0)
    Ty = T(0)

    # set all starting points to distance zero and frozen
    for z in 1:size(SourcePoints, 2)
        # starting point
        x = round(Int, SourcePoints[1, z])
        y = round(Int, SourcePoints[2, z])
        # Set starting point to frozen and distance to zero
        Frozen[x, y] = true
        distanceimage[x, y] = 0
    end

    # Add all neighbours of the starting points to narrow list
    @inbounds for z in 1:size(SourcePoints, 2)
        # starting point
        x = round(Int, SourcePoints[1, z])
        y = round(Int, SourcePoints[2, z])

        for k in 1:4
            # Location of neighbour
            i = x + ne[k, 1]
            j = y + ne[k, 2]
            # Check if current neighbour is not yet frozen and inside the
            # picture
            if (
                (i > 0) &&
                (j > 0) &&
                (i <= size(speedimage, 1)) &&
                (j <= size(speedimage, 2)) &&
                (~Frozen[i, j])
            )
                Tt = 1 / max(speedimage[i, j], eps(T))
                Ty = T(1)
                # Update distance in neighbour list or add to neighbour list
                if distanceimage[i, j] > 0
                    if neg_list_1[distanceimage[i, j]] > Tt
                        neg_list_1[distanceimage[i, j]] = Tt
                    end
                    if Ed
                        neg_list_4[distanceimage[i, j]] = min(
                            Ty, neg_list_4[distanceimage[i, j]]
                        )
                    end
                else
                    neg_pos = neg_pos + 1
                    neg_list_1[neg_pos] = Tt
                    neg_list_x[neg_pos] = i
                    neg_list_y[neg_pos] = j
                    if Ed
                        neg_list_4[neg_pos] = Ty
                    end
                    distanceimage[i, j] = neg_pos
                end
            end
        end
    end

    # Reuse variables:
    Tpatch = zeros(T, 5, 5)
    Order = zeros(Int8, 1, 4)
    Tm = zeros(T, 1, 4)
    Tm2 = zeros(T, 1, 4)
    Coeff = zeros(T, 3)
    TT = zeros(T, 2)
    TT2 = zeros(T, 2)

    # Loop through all pixels of the image
    for itt in 1:length(speedimage)
        # Get the pixel from narrow list (boundary list) with smallest
        # distance value and set it to current pixel location

        if neg_pos == 0
            break
        end
        (t, index) = findmin(view(neg_list_1, 1:neg_pos))
        x = neg_list_x[index]
        y = neg_list_y[index]
        Frozen[x, y] = true
        distanceimage[x, y] = neg_list_1[index]

        if Ed
            Y[x, y] = neg_list_4[index]
        end

        # Remove min value by replacing it with the last value in the array
        if index < neg_pos
            neg_list_1[index] = neg_list_1[neg_pos]
            neg_list_x[index] = neg_list_x[neg_pos]
            neg_list_y[index] = neg_list_y[neg_pos]
            if Ed
                neg_list_4[index] = neg_list_4[neg_pos]
            end
            x2 = neg_list_x[index]
            y2 = neg_list_y[index]
            distanceimage[x2, y2] = index
        end
        neg_pos = neg_pos - 1
        # Loop through all 4 neighbours of current pixel
        for k in 1:4
            # Location of neighbour
            i = x + ne[k, 1]
            j = y + ne[k, 2]

            # Check if current neighbour is not yet frozen and inside the
            # picture
            if (i > 0) &&
               (j > 0) &&
               (i <= size(speedimage, 1)) &&
               (j <= size(speedimage, 2)) &&
               (~Frozen[i, j])
                Tt = calculatedistance!(
                    distanceimage,
                    Tpatch,
                    Order,
                    Coeff,
                    Tm,
                    Tm2,
                    TT,
                    TT2,
                    speedimage[i, j],
                    size(speedimage),
                    i,
                    j,
                    usesecond,
                    usecross,
                    Frozen,
                )
                if Ed
                    Ty = calculatedistance!(
                        Y,
                        Tpatch,
                        Order,
                        Coeff,
                        Tm,
                        Tm2,
                        TT,
                        TT2,
                        speedimage[i, j],
                        size(speedimage),
                        i,
                        j,
                        usesecond,
                        usecross,
                        Frozen,
                    )
                end

                # Update distance in neighbour list or add to neighbour list
                if distanceimage[i, j] > 0
                    neg_list_1[round(Int, distanceimage[i, j])] = min(
                        Tt, neg_list_1[round(Int, distanceimage[i, j])]
                    )
                    if Ed
                        neg_list_4[round(Int, distanceimage[i, j])] = min(
                            Ty, neg_list_4[round(Int, distanceimage[i, j])]
                        )
                    end
                else
                    neg_pos = neg_pos + 1
                    neg_list_1[neg_pos] = Tt
                    neg_list_x[neg_pos] = i
                    neg_list_y[neg_pos] = j
                    if Ed
                        neg_list_4[neg_pos] = Ty
                    end
                    distanceimage[i, j] = neg_pos
                end
            end
        end
    end
    return distanceimage
end # End function msfm

function calculatedistance(
    TI::AbstractArray{T}, Fij, sizeF, i, j, usesecond, usecross, Frozen
) where {T}
    Tpatch = zeros(T, 5, 5)
    Order = zeros(Int8, 1, 4)
    Tm = zeros(T, 1, 4)
    Tm2 = zeros(T, 1, 4)
    Coeff = zeros(T, 3)
    TT = zeros(T, 2)
    TT2 = zeros(T, 2)
    return calculatedistance!(
        TI,
        Tpatch,
        Order,
        Coeff,
        Tm,
        Tm2,
        TT,
        TT2,
        Fij,
        sizeF,
        i,
        j,
        usesecond,
        usecross,
        Frozen,
    )
end

function calculatedistance!(
    TI::AbstractArray{T},
    Tpatch,
    Order,
    Coeff,
    Tm,
    Tm2,
    TT,
    TT2,
    Fij,
    sizeF,
    i,
    j,
    usesecond,
    usecross,
    Frozen,
) where {T}
    fill!(Tpatch, T(Inf))
    @inbounds for nx in -2:2
        for ny in -2:2
            i_n = i + nx
            j_n = j + ny
            if i_n > 0 && j_n > 0 && i_n <= sizeF[1] && j_n <= sizeF[2] && Frozen[i_n, j_n]
                Tpatch[nx + 3, ny + 3] = TI[i_n, j_n]
            end
        end
    end

    # The values in order is 0 if no neighbours in that direction
    # 1 if 1e order derivatives is used and 2 if second order
    # derivatives are used
    fill!(Order, 0)

    # Make 1e order derivatives in x and y direction
    fill!(Tm, 0)
    Tm[1] = min(Tpatch[2, 3], Tpatch[4, 3])
    if isfinite(Tm[1])
        Order[1] = 1
    end

    Tm[2] = min(Tpatch[3, 2], Tpatch[3, 4])
    if isfinite(Tm[2])
        Order[2] = 1
    end

    # Make 1e order derivatives in cross directions
    if usecross
        Tm[3] = min(Tpatch[2, 2], Tpatch[4, 4])
        if isfinite(Tm[3])
            Order[3] = 1
        end
        Tm[4] = min(Tpatch[2, 4], Tpatch[4, 2])
        if isfinite(Tm[4])
            Order[4] = 1
        end
    end

    # Make 2e order derivatives
    if usesecond
        fill!(Tm2, 0)
        # pixels with a pixeldistance 2 from the center must be
        # lower in value otherwise use other side or first order
        ch1 = (Tpatch[1, 3] < Tpatch[2, 3]) && isfinite(Tpatch[2, 3])
        ch2 = (Tpatch[5, 3] < Tpatch[4, 3]) && isfinite(Tpatch[4, 3])

        if ch1 && ch2
            Tm2[1] = min(
                (4 * Tpatch[2, 3] - Tpatch[1, 3]) / 3, (4 * Tpatch[4, 3] - Tpatch[5, 3]) / 3
            )
            Order[1] = 2
        elseif ch1
            Tm2[1] = (4 * Tpatch[2, 3] - Tpatch[1, 3]) / 3
            Order[1] = 2
        elseif ch2
            Tm2[1] = (4 * Tpatch[4, 3] - Tpatch[5, 3]) / 3
            Order[1] = 2
        end

        ch1 = (Tpatch[3, 1] < Tpatch[3, 2]) && isfinite(Tpatch[3, 2])
        ch2 = (Tpatch[3, 5] < Tpatch[3, 4]) && isfinite(Tpatch[3, 4])

        if ch1 && ch2
            Tm2[2] = min(
                (4 * Tpatch[3, 2] - Tpatch[3, 1]) / 3, (4 * Tpatch[3, 4] - Tpatch[3, 5]) / 3
            )
            Order[2] = 2
        elseif ch1
            Tm2[2] = (4 * Tpatch[3, 2] - Tpatch[3, 1]) / 3
            Order[2] = 2
        elseif ch2
            Tm2[2] = (4 * Tpatch[3, 4] - Tpatch[3, 5]) / 3
            Order[2] = 2
        end

        if usecross
            ch1 = (Tpatch[1, 1] < Tpatch[2, 2]) && isfinite(Tpatch[2, 2])
            ch2 = (Tpatch[5, 5] < Tpatch[4, 4]) && isfinite(Tpatch[4, 4])
            if ch1 && ch2
                Tm2[3] = min(
                    (4 * Tpatch[2, 2] - Tpatch[1, 1]) / 3,
                    (4 * Tpatch[4, 4] - Tpatch[5, 5]) / 3,
                )
                Order[3] = 2
            elseif ch1
                Tm2[3] = (4 * Tpatch[2, 2] - Tpatch[1, 1]) / 3
                Order[3] = 2
            elseif ch2
                Tm2[3] = (4 * Tpatch[4, 4] - Tpatch[5, 5]) / 3
                Order[3] = 2
            end

            ch1 = (Tpatch[1, 5] < Tpatch[2, 4]) && isfinite(Tpatch[2, 4])
            ch2 = (Tpatch[5, 1] < Tpatch[4, 2]) && isfinite(Tpatch[4, 2])
            if ch1 && ch2
                Tm2[4] = min(
                    (4 * Tpatch[2, 4] - Tpatch[1, 5]) / 3,
                    (4 * Tpatch[4, 2] - Tpatch[5, 1]) / 3,
                )
                Order[4] = 2
            elseif ch1
                Tm2[4] = (4 * Tpatch[2, 4] - Tpatch[1, 5]) / 3
                Order[4] = 2
            elseif ch2
                Tm2[4] = (4 * Tpatch[4, 2] - Tpatch[5, 1]) / 3
                Order[4] = 2
            end
        end
    else
        fill!(Tm2, 0)
    end

    # Calculate the distance using x and y direction
    Coeff[1] = 0
    Coeff[2] = 0
    Coeff[3] = -1 / (max(Fij^2, eps(T)))
    for t in 1:2
        if Order[t] == 1
            Coeff[1] += 1
            Coeff[2] -= 2 * Tm[t]
            Coeff[3] += Tm[t]^2
        elseif Order[t] == 2
            Coeff[1] += 2.25
            Coeff[2] -= 2.25 * 2 * Tm2[t]
            Coeff[3] += 2.25 * Tm2[t]^2
        end
    end

    roots!(TT, Coeff)
    # Calculate the distance using the cross directions
    if usecross
        # Coeff[1] += 0
        # Coeff[2] += 0
        Coeff[3] += -1 / (max(Fij^2, eps(T)))
        for t in 3:4
            if Order[t] == 1
                Coeff[1] += 0.5
                Coeff[2] -= Tm[t]
                Coeff[3] += 0.5 * Tm[t]^2
            elseif Order[t] == 2
                Coeff[1] += 0.5 * 2.25
                Coeff[2] -= 2.25 * Tm2[t]
                Coeff[3] += 0.5 * 2.25 * Tm2[t]^2
            end
        end
        roots!(TT2, Coeff)
        # Select minimum distance value of both stensils
        if ~isempty(TT2)
            TT = min(maximum(TT), maximum(TT2))
        end
    else
        TT = maximum(TT)
    end

    # Upwind condition check, current distance must be larger
    # then direct neighbours used in solution
    # DirectNeigbInSol=Tm(isfinite(Tm));
    # if(nnz(DirectNeigbInSol>=Tt)>0) # Will this ever happen?
    #    Tt=min(DirectNeigbInSol)+(1/(max(Fij,eps)));
    # end
    return minimum(TT)
end

function roots!(z, Coeff::AbstractArray{T}) where {T}
    a = Coeff[1]
    b = Coeff[2]
    c = Coeff[3]
    d = max(b^2 - 4.0 * a * c, zero(T))
    if a != 0
        z[1] = (-b - sqrt(d)) / 2a
        z[2] = (-b + sqrt(d)) / 2a
    else
        z[1] = 2c / (-b - sqrt(d))
        z[2] = 2c / (-b + sqrt(d))
    end
end

function roots(Coeff::AbstractArray{T}) where {T}
    z = zeros(T, 2, 1)
    roots!(z, Coeff)
    return z
end