text:
# This file includes interpolant for interplation of sampled inputs.

"""
    $TYPEDEF

Constructs a linnear interpolant that interpolates between the poinsts `(tinit, coefinit)` and `(tfinal, coeffinal)`.

# Fields

    $TYPEDFIELDS
"""
mutable struct Interpolant{TMB, INB, ITP}
    "Buffer to record time"
    timebuf::TMB
    "Buffer to record data"
    databuf::INB
    "Internal interpolating function"
    itp::ITP
    function Interpolant(nt::Int, nd::Int)
        timebuf = Buffer(nt)
        databuf = nd == 1 ? Buffer(nt) : Buffer(nd, nt)
        itp = [interpolation(zeros(1), zeros(1)) for i in 1 : nd]
        new{typeof(timebuf), typeof(databuf), typeof(itp)}(timebuf, databuf, itp)
    end
end

show(io::IO, interpolant::Interpolant) = print(io, "Interpolant(timebuf:$(interpolant.timebuf), ",
    "databuf:$(interpolant.databuf), itp:$(interpolant.itp))")

# Callling interpolant.
# Syntax for u(t).
(interp::Interpolant)(t) = ndims(buf) == 1 ? interp.itp(t) : map(f -> f(t), interp.itp)

# Syntax for u[idx](t)
getindex(interp::Interpolant, idx::Vararg{Int, N}) where N = interp.itp[idx...]

# Update of interpolant. That is, reinterpolation.
"""
    $SIGNATURES

Updates `interpolant` using the data in `timebuf` and `databuf` of `interpolant`.
"""
update!(interpolant::Interpolant{T1, <:AbstractVector, T2}) where {T1,T2} = interpolant.itp[1] = _update!(interpolant)
update!(interpolant::Interpolant{T1, <:AbstractMatrix, T2}) where {T1,T2} = interpolant.itp = _update!(interpolant)
_update!(interpolant) = interpolation(content(interpolant.timebuf, flip=true), content(interpolant.databuf, flip=true))

interpolation(tt, uu::AbstractMatrix) = map(row -> interpolation(tt, row), eachrow(uu))
interpolation(tt, uu::AbstractVector) = CubicSplineInterpolation(getranges(tt, uu)...; extrapolation_bc=Line())

function getranges(tt, uu)
    length(tt) < 2 && return (range(tt[1], length=2, step=eps()), range(uu[1], length=2, step=eps()))
    return range(tt[1], tt[end], length=length(tt)), uu
end


text:
using CPLEX, Base.Test

@testset "SOS" begin
    @testset "SOS 01" begin
        m = CPLEX.CplexMathProgModel(CPX_PARAM_SCRIND=0)
        CPLEX.loadproblem!(m,
            [1 2 3 0 0 0 0 0 -1 0; 0 0 0 5 4 7 2 1 0 -1], # constraint matrix
            [0, 0, 0, 0, 0, 0, 0, 0, -Inf, -Inf], # variable lb
            [2, 2, 2, 2, 2, 2, 2, 2, Inf, Inf], # variable ub
            vcat(zeros(8), ones(2)), # objective vector
            [0, 0], # constraint lb
            [0, 0], # constraint ub
            :Max)
        @test CPLEX.get_prob_type(m.inner) == :LP

        CPLEX.setvartype!(m, vcat(fill(:Bin, 8), fill(:Cont, 2)))
        @test CPLEX.get_prob_type(m.inner) == :MILP

        CPLEX.addsos1!(m, [1,2,3], [1.0,2.0,3.0])
        CPLEX.addsos2!(m, [4,5,6,7,8], [5.0, 4.0, 7.0, 2.0, 1.0])
        @test CPLEX.get_prob_type(m.inner) == :MILP

        CPLEX.optimize!(m)

        @test CPLEX.getobjval(m) ≈ 15.0
        sol = CPLEX.getsolution(m)
        @test sol[9] ≈  3.0
        @test sol[10] ≈ 12.0

        CPLEX.setvartype!(m, fill(:Cont, 10))
        @test CPLEX.get_prob_type(m.inner) == :MILP

        CPLEX.optimize!(m)

        @test CPLEX.getobjval(m) ≈ 30.0
        sol = CPLEX.getsolution(m)
        @test sol[9] ≈ 6.0
        @test sol[10] ≈ 24.0
    end

    @testset "SOS 02" begin
        m = CPLEX.CplexMathProgModel(CPX_PARAM_SCRIND=0)
        CPLEX.loadproblem!(m, Matrix{Float64}(0,3), [-Inf, -Inf, -Inf], [1,1,2], [2,1,1], Float64[], Float64[], :Max)
        @test CPLEX.get_prob_type(m.inner) == :LP

        CPLEX.addsos1!(m, [1,2], [1.0,2.0])
        CPLEX.addsos1!(m, [1,3], [1.0,2.0])
        @test CPLEX.get_prob_type(m.inner) == :MILP

        CPLEX.optimize!(m)
        @test CPLEX.getobjval(m) ≈ 3.0
    end
end

ext: jl


text:
module CartesianSphericalHarmonics

# Write your package code here.
using LinearAlgebra
using Reexport

@reexport using MultivariatePolynomials
@reexport using TypedPolynomials

export YLMNorm, Schmidt, Laplace, Nonorm, Full
export rlm, ylm
#normalization factor
abstract type YLMNorm{T} end

struct Schmidt{T} <: YLMNorm{T}; end
struct Laplace{T} <: YLMNorm{T}; end
struct Nonorm{T} <: YLMNorm{T}; end
struct Full{T} <: YLMNorm{T}; end

function ylmcoeff(N::Schmidt{T},l::Integer, m::Integer) where T

	k = one(T)
	for i in (l-m+1):(l+m)
	  k *= i
	end
	fac = sqrt(1/k)
	if m!=0
		fac*=sqrt(2*one(T))
	end
	return fac
end

function ylmcoeff(N::Laplace{T}, l::Integer, m::Integer) where T

  k = one(T)
  for i in (l-m+1):(l+m)
    k *= i
  end

  return sqrt((2*l+1) / (4*T(pi)*k))
end

function ylmcoeff(N::Full{T}, l::Integer, m::Integer) where T
	fac = m==0 ? one(T) : sqrt(2*one(T))
	return ylmcoeff(Laplace{T}(),l,m)*sqrt(4T(pi))*fac
end

function ylmcoeff(N::Nonorm{T}, l::Integer, m::Integer) where T
 return one(T)
end

function cossinpoly(m::Integer, x::Variable, y::Variable)

  sum = zero(x*y)
  for j in 0:div(m,2)
	  k = 2j
    sum += ((-1)^j)*binomial(m, k)*y^Int(k)*x^Int(m-k)
  end
  return sum
end

function sinsinpoly(m::Integer, x::Variable, y::Variable)

  sum = zero(x*y)
  for j in 0:div((m-1),2)
	  k = 2j+1
    sum += ((-1)^j)*binomial(m, k)*y^Int(k)*x^Int(m-k)
  end
  return sum
end

#
function legendre_assoc(::Type{T},z,l,m) where T
	p = (z^2 - 1)^l

	for i = 1:l+abs(m)
	c = i <= l ? 1/(2one(T)*i) : one(T)
	p = c*differentiate(p, z)
	end

	p *= (-1)^m
	return p
end
# #legendre polynomial:
# P(l::BigInt,x) = differentiate((x^2 - 1)^l,x,l)
#
# #associated legendre polynomial (without sin(θ)^m)
# P(l::BigInt,m::BigInt,x) = 1//(big(2)^l*factorial(l))*differentiate(P(l,x),x,m)

"""
    ylm(l::Integer, m::Integer, x::Variable, y::Variable, z::Variable)
*Description:*  Calculation of the spherical harmonic for a given order (l,m) in Cartesian coordinates\\

*Input:*  `l`       - Degree of the spherical harmonic\\
          `m`       - Order of the spherical harmonic\\
          `x, y, z` - Cartesian coordinates\\

*Output:*  Spherical harmonic polynomial
"""
function ylm(l::Integer, m::Integer, x::Variable, y::Variable, z::Variable;
			 norm::YLMNorm{T}=Nonorm{Rational{Int}}(),real::Bool=true) where T

	if abs(m) > l
		throw(DomainError(m,"-l <= m <= l expected, but m = $m and l = $l."))
	end

	p = legendre_assoc(T,z,l,m)
	# p = (z^2 - 1)^l

	# for i = 1:l+abs(m)
	#   c = i <= l ? 1/(2one(T)*i) : one(T)
	#   p = c*differentiate(p, z)
	# end

	# p *= (-1)^m

	if real
		if m > 0
			return ylmcoeff(norm, l, m)*cossinpoly(m,x,y)*p
		elseif m < 0
			return ylmcoeff(norm, l, abs(m))*sinsinpoly(abs(m),x,y)*p
		else
			return ylmcoeff(norm, l, 0)*p
		end
	else
		return ylmcoeff(norm, l, m)*(cossinpoly(m,x,y)+im*sinsinpoly(m,x,y))*p
	end
end

# multiplying r^l*ylm(x,y,z)
function rlm(l::Integer, m::Integer, x::Variable, y::Variable, z::Variable;
			   norm::YLMNorm{T}=Nonorm{Rational{Int}}(),real::Bool=true) where T
	p = ylm(l,m,x,y,z;norm=norm,real=real)
	tout = []

	for t in terms(p)
		deg = degree(monomial(t)) # degree of monomial
		degR = l-deg # degree of r: l-deg
		push!(tout,(x^2+y^2+z^2)^div(degR,2)*t) # r² replaced by x²+y²+z²
	end

	return polynomial(tout)
end

end


text:
if Sys.islinux()
    py"""
    import tensorflow as tf
    import socket
    if socket.gethostname() != "Dolores":
        libFwiOp = tf.load_op_library('../Ops/FWI/build/libFwiOp.so')
    else:
        libFwiOp = tf.load_op_library('../Ops/FWI/build_dolores/libFwiOp.so')
    @tf.custom_gradient
    def fwi_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname):
        misfit = libFwiOp.fwi_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        def grad(dy):
            return libFwiOp.fwi_op_grad(dy, tf.constant(1.0,dtype=tf.float64),λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        return misfit, grad
    def fwi_obs_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname):
        misfit = libFwiOp.fwi_obs_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        return misfit
    """
    elseif Sys.isapple()
    py"""
    import tensorflow as tf
    libFwiOp = tf.load_op_library('../Ops/FWI/build/libFwiOp.so')
    @tf.custom_gradient
    def fwi_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname):
        misfit = libFwiOp.fwi_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        def grad(dy):
            return libFwiOp.fwi_op_grad(dy, tf.constant(1.0,dtype=tf.float64),λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        return misfit, grad
    def fwi_obs_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname):
        misfit = libFwiOp.fwi_obs_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        return misfit
    """
    elseif Sys.iswindows()
    py"""
    import tensorflow as tf
    libFwiOp = tf.load_op_library('../Ops/FWI/build/libFwiOp.so')
    @tf.custom_gradient
    def fwi_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname):
        misfit = libFwiOp.fwi_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        def grad(dy):
            return libFwiOp.fwi_op_grad(dy, tf.constant(1.0,dtype=tf.float64),λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        return misfit, grad
    def fwi_obs_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname):
        misfit = libFwiOp.fwi_obs_op(λ,μ,ρ,stf,gpu_id,shot_ids,para_fname)
        return misfit
    """
    end

    fwi_op = py"fwi_op"
    fwi_obs_op = py"fwi_obs_op"


text:
##### Beginning of file

import RepoSnapshots
import Test # stdlib

include("test-snapshot-utils.jl")
include("test-utils.jl")

include("set-up-test-modules.jl")
include("test-package-directory.jl")
include("test-version.jl")

##### End of file


text:
# using Pkg
# Pkg.add("Plots")
# Pkg.add("PyPlot")
import Plots as plt
plt.pyplot()

function F(v::AbstractArray)::AbstractArray
    f = 0.01
    return f * [
    v[2]
    -v[1]
  ]
end

function interpolate(v::AbstractArray, x_D::AbstractArray, x_A::AbstractArray, dx::Float64, I::Int)::AbstractArray
    v_inter = fill([0.0;0.0], (I, I))
    for i_D in 1:I, j_D in 1:I
        x = x_D[i_D, j_D]
        (i0, j0) = floor.(Integer, x / dx)
        (i0, j0) = (i0 + 1, j0 + 1)  # since index not zero based
        (i1, j1) = (i0 + 1, j0)
        (i2, j2) = (i0, j0 + 1)
        if x[2] - j0 * dx > 0.5  # interpolation from upper triangle
            (i0, j0) = (i0 + 1, j0 + 1)
        end
        # open boundaries: shift back to the nearest at boundary
        di = min(i0, i1, i2)
        if di < 1
            i0 += -di + 1
            i1 += -di + 1
            i2 += -di + 1
        end
        di = max(i0 - I, i1 - I, i2 - I)
        if di > 0
            i0 += -di
            i1 += -di
            i2 += -di
        end
        dj = min(j0, j1, j2)
        if dj < 1
            j0 += -dj + 1
            j1 += -dj + 1
            j2 += -dj + 1
        end
        dj = max(j0 - I, j1 - I, j2 - I)
        if dj > 0
            j0 += -dj
            j1 += -dj
            j2 += -dj
        end
        t =  [x_A[i1,j1] - x_A[i0,j0] x_A[i2,j2] - x_A[i0,j0]] \ (x - x_A[i0,j0])
        t0 = 1.0 - t[1] - t[2]
        v_inter[i_D,j_D] = t0 * v[i0,j0] + t[1] * v[i1,j1] + t[2] * v[i2,j2]
    end
    return v_inter
end

function plotField(x_A, v::AbstractArray)
    (x1, x2) = extractCoords(x_A)
    (v1, v2) = extractCoords(v)
    plot = plt.quiver(x1, x2, quiver=(v1, v2))
    display(plot)
end

# extracts field as x and y coordinate vectors
# field: number[2,][n,n]
function extractCoords(field::AbstractArray)
    field = reshape(deepcopy(field), (1, :))
    field = [field[1] field[2:end]...]
    x1 = vec(field[1,begin:end])
    x2 = vec(field[2,begin:end])
    return (x1, x2)
end

function plotParticles(x_P::AbstractArray, v_P::AbstractArray, dt::Float64, I::Int)
    x_P = x_P + dt * v_P
    (x1, x2) = extractCoords(x_P)
    center = floor(Integer, (I/2)^2) # only plot the center particle
    display(plt.scatter!(x1[center:center], x2[center:center], xlims=(0, 6), ylims=(0, 6), label=""))
end

(function main()
    L =  10.0 # 100.0
    dx = 1.0
    dt = 0.1
    T = 600.0
    I = ceil(Integer, L / dx)
    v = fill([10.0; 0.0], (I, I))
    x_A = [dx * [i - 1; j - 1] for i = 1:I, j = 1:I]

    x_D = deepcopy(x_A)
    x_P = deepcopy(x_A)
    lastPlot = 0.0
    for t = 0.0:dt:T
        for n = 1:5
            x_D = x_A - dt * interpolate(v, x_D, x_A, dx, I)
        end
        v_D = interpolate(v, x_D, x_A, dx, I)
        v = v_D + dt * F.(v_D)
        if t - lastPlot > 5.0
            # plotField(x_A, v)
            v_P = interpolate(v, x_P, x_A, dx, I)
            plotParticles(x_P, v_P, dt, I)
            lastPlot = t
            sleep(1) # TODO
        end
    end
end)()

max_stars_repo_name: applied-math-coding/article-snippets

max_issues_repo_name: applied-math-coding/article-snippets

max_forks_repo_name: applied-math-coding/article-snippets


text:
data = parse.(Int, split(read(joinpath(@__DIR__, "./6.txt"), String), ","))

data = [3,4,3,1,2]

function part1(fish)
  for day in 1:80
    new_fish = fill(8,length(filter(x->x==0,fish)))
    fish = map(x->x==0 ? 6 : x - 1, fish)
    append!(fish,new_fish)
  end

  length(fish)
end

function inc_day(current_counts,days)
  next_counts = Dict{Int,Int}()
  if days > 0
    for i in [1,2,3,4,5,6,7,8,0]
      if i == 0 && haskey(current_counts,i)
        next_counts[8] = current_counts[i]
        if haskey(next_counts,6)
          next_counts[6] += current_counts[i]
        else
          next_counts[6] = current_counts[i]
        end
      elseif i > 0 && haskey(current_counts,i)
        next_counts[i - 1] = current_counts[i]
      end
    end
    return inc_day(next_counts, days - 1)
  else
    return current_counts
  end
end

function part2(fish,days)
  timer_counts = Dict{Int,Int}()
  for n in 1:5
    timer_counts[n] = length(filter(x->x==n,fish))
  end
  return sum(values(inc_day(timer_counts,days)))

end

part1(data)

part2(data,256)


text:
using FIRITools
using Plots

const alt = FIRITools.ALTITUDE

# Compare mean and median profiles
day_mean = firi(chi=(0, 95))
day_med = FIRITools.quantile(0.5, chi=(0, 95))
day_30 = FIRITools.quantile(0.3, chi=(0, 95))
day_70 = FIRITools.quantile(0.7, chi=(0, 95))

plot(day_mean, alt/1000, label="mean", xscale=:log10, linewidth=2, ylims=(55, 110),
     legend=:topleft, title="Day", ylabel="altitude (km)", xlabel="Ne (m⁻³)")
plot!(day_med, alt/1000, label="median", xscale=:log10, linewidth=2)
plot!(day_30, alt/1000, label="30%", xscale=:log10)
plot!(day_70, alt/1000, label="70%", xscale=:log10)
savefig("day_mean.png")

night_mean = firi(chi=(100, 180))
night_med = FIRITools.quantile(0.5, chi=(100, 180))
night_30 = FIRITools.quantile(0.3, chi=(100, 180))
night_70 = FIRITools.quantile(0.7, chi=(100, 180))

plot(night_mean, alt/1000, label="mean", xscale=:log10, linewidth=2, ylims=(55, 110),
     legend=:topleft, title="Night", ylabel="altitude (km)", xlabel="Ne (m⁻³)")
plot!(night_med, alt/1000, label="median", xscale=:log10, linewidth=2)
plot!(night_30, alt/1000, label="30%", xscale=:log10)
plot!(night_70, alt/1000, label="70%", xscale=:log10)
savefig("night_mean.png")

# Compare mean and median around interpolation points
day_mean = firi(chi=(0, 95), lat=(30, 45))
day_med = FIRITools.quantile(0.5, chi=(0, 95), lat=(30, 45))
day_30 = FIRITools.quantile(0.3, chi=(0, 95), lat=(30, 45))
day_70 = FIRITools.quantile(0.7, chi=(0, 95), lat=(30, 45))

plot(day_mean, alt/1000, label="mean", xscale=:log10, linewidth=2, ylims=(55, 110),
     legend=:topleft, title="Day, limited latitude", ylabel="altitude (km)", xlabel="Ne (m⁻³)")
plot!(day_med, alt/1000, label="median", xscale=:log10, linewidth=2)
plot!(day_30, alt/1000, label="30%", xscale=:log10)
plot!(day_70, alt/1000, label="70%", xscale=:log10)
savefig("day_mean_latzoom.png")

day_mean = firi(chi=(30, 45))
day_med = FIRITools.quantile(0.5, chi=(30, 45))
day_30 = FIRITools.quantile(0.3, chi=(30, 45))
day_70 = FIRITools.quantile(0.7, chi=(30, 45))

plot(day_mean, alt/1000, label="mean", xscale=:log10, linewidth=2, ylims=(55, 110),
     legend=:topleft, title="Day, limited SZA", ylabel="altitude (km)", xlabel="Ne (m⁻³)")
plot!(day_med, alt/1000, label="median", xscale=:log10, linewidth=2)
plot!(day_30, alt/1000, label="30%", xscale=:log10)
plot!(day_70, alt/1000, label="70%", xscale=:log10)
savefig("day_mean_chizoom.png")

# Check to see if interpolation is acting as expected

real_lats = unique(FIRITools.HEADER[!,"Lat, deg"])
lats = 0:5:90
interp_lat = setdiff(lats, real_lats)
N = length(lats)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10, legend=:topleft,
         ylims=(55, 110), legendtitle="Latitude")
for la in real_lats
    prof = firi(45, la)
    plot!(p, prof, alt/1000, label=la, color=cmap[findfirst(isequal(la), lats)])
end
for la in interp_lat
    prof = firi(45, la)
    plot!(p, prof, alt/1000, label=nothing, color=cmap[findfirst(isequal(la), lats)],
          linestyle=:dash)
end
display(p)
savefig("lat.png")

real_sza = unique(FIRITools.HEADER[!,"Chi, deg"])
chis = 0:5:130
interp_sza = setdiff(chis, real_sza)
# interp_sza = [105, 110, 115, 120]
N = length(chis)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10,
         ylims=(55, 110), legend=:topleft, legendtitle="SZA")
for sza in real_sza
    prof = firi(sza, 45)
    plot!(p, prof, alt/1000, label=sza, color=cmap[findfirst(isequal(sza), chis)])
end

for sza in interp_sza
    prof = firi(sza, 45)
    plot!(p, prof, alt/1000, label=nothing, color=cmap[findfirst(isequal(sza), chis)],
          linestyle=:dash)
end
display(p)
savefig(p, "sza.png")

real_sza = [75, 80, 85, 90, 95, 100]
chis = 75:100
interp_sza = setdiff(chis, real_sza)
N = length(chis)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10,
         ylims=(55, 110), legend=:topleft, legendtitle="SZA")
for sza in real_sza
    prof = firi(sza, 45)
    plot!(p, prof, alt/1000, label=sza, color=cmap[findfirst(isequal(sza), chis)])
end

for sza in interp_sza
    prof = firi(sza, 45)
    plot!(p, prof, alt/1000, label=nothing, color=cmap[findfirst(isequal(sza), chis)],
          linestyle=:dash)
end
display(p)
savefig("sza_zoom.png")

###
# With extrapolation to ground

real_lats = unique(FIRITools.HEADER[!,"Lat, deg"])
lats = 0:5:90
interp_lat = setdiff(lats, real_lats)
N = length(lats)
cmap = palette(:rainbow, N, rev=true)
newalt = 0:110

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10, legend=:topleft,
         ylims=(0, 110), yticks=0:20:110, legendtitle="Latitude", xticks=exp10.([0, 3, 6, 9, 12]))
for la in real_lats
    prof = FIRITools.extrapolate(firi(45, la), newalt*1e3)
    plot!(p, prof, newalt, label=la, color=cmap[findfirst(isequal(la), lats)])
end
for la in interp_lat
    prof = FIRITools.extrapolate(firi(45, la), newalt*1e3)
    plot!(p, prof, newalt, label=nothing, color=cmap[findfirst(isequal(la), lats)],
          linestyle=:dash)
end
display(p)
savefig("lat_extrap.png")

real_sza = unique(FIRITools.HEADER[!,"Chi, deg"])
chis = 0:5:130
interp_sza = setdiff(chis, real_sza)
# interp_sza = [105, 110, 115, 120]
N = length(chis)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10,
         ylims=(0, 110), yticks=0:20:110, legend=:topleft, legendtitle="SZA", xticks=exp10.([0, 3, 6, 9, 12]))
for sza in real_sza
    prof = FIRITools.extrapolate(firi(sza, 45), newalt*1e3)
    plot!(p, prof, newalt, label=sza, color=cmap[findfirst(isequal(sza), chis)])
end

for sza in interp_sza
    prof = FIRITools.extrapolate(firi(sza, 45), newalt*1e3)
    plot!(p, prof, newalt, label=nothing, color=cmap[findfirst(isequal(sza), chis)],
          linestyle=:dash)
end
display(p)
savefig(p, "sza_extrap.png")

# Using median
real_sza = unique(FIRITools.HEADER[!,"Chi, deg"])
chis = 0:5:130
interp_sza = setdiff(chis, real_sza)
# interp_sza = [105, 110, 115, 120]
N = length(chis)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10,
         ylims=(0, 110), yticks=0:20:110, legend=:topleft, legendtitle="SZA", xticks=exp10.([0, 3, 6, 9, 12]))
for sza in real_sza
    prof = FIRITools.extrapolate(FIRITools.quantile(sza, 45, 0.5), newalt*1e3)
    plot!(p, prof, newalt, label=sza, color=cmap[findfirst(isequal(sza), chis)])
end

for sza in interp_sza
    prof = FIRITools.extrapolate(FIRITools.quantile(sza, 45, 0.5), newalt*1e3)
    plot!(p, prof, newalt, label=nothing, color=cmap[findfirst(isequal(sza), chis)],
          linestyle=:dash)
end
display(p)
savefig(p, "sza_median_extrap.png")


text:
const CS_MODE_SINGLE = 0x01
const CS_MODE_DUAL = 0x02
const CS_MODE_QUAD = 0x04
const CS_MODE_OCT = 0x08
const CS_MASKED_MODE = Float32(0x00000000)
const CS_MODE_RESERVED = 0x10
const CS_MODE_STREAMING = 0x20
const CS_MODE_HW_SUPER_RES = 0x40
const CS_MODE_POWER_ON = 0x80
const CS_MODE_PRETRIG_MULREC = 0x0200
const CS_MODE_REFERENCE_CLK = 0x0400
const CS_MODE_CS3200_CLK_INVERT = 0x0800
const TIMESTAMP_GCLK = 0x00
const TIMESTAMP_FREERUN = 0x10
const TIMESTAMP_DEFAULT = TIMESTAMP_GCLK | TIMESTAMP_FREERUN
const TIMESTAMP_MCLK = 0x01
const TIMESTAMP_SEG_RESET = 0x00
const TxMODE_DATA_ANALOGONLY = 0x00
const TxMODE_DEFAULT = TxMODE_DATA_ANALOGONLY
const TxMODE_SLAVE = 0x80000000
const TxMODE_RAWDATA = 0x40000000
const TxMODE_DATA_FLOAT = 0x01
const TxMODE_TIMESTAMP = 0x02
const TxMODE_DATA_16 = 0x04
const TxMODE_DATA_ONLYDIGITAL = 0x08
const TxMODE_DATA_32 = 0x10
const TxMODE_DATA_FFT = 0x30
const TxMODE_DATA_INTERLEAVED = 0x40
const TxMODE_SEGMENT_TAIL = 0x80
const TxFORMAT_OCTAL_INTERLEAVED = 0x12345678
const TxFORMAT_QUAD_INTERLEAVED = 0x11223344
const TxFORMAT_DUAL_INTERLEAVED = 0x11221122
const TxFORMAT_COBRA_DUAL_INTERLEAVED = 0x11112222
const TxFORMAT_NON_INTERLEAVED = 0x11111111
const TxFORMAT_SINGLE = 0x11111111
const TxFORMAT_STACKED = Float32(0x07ffffff)
const CS_CHAN_1 = 1
const CS_CHAN_2 = 2
const CS_FILTER_OFF = 0
const CS_FILTER_ON = 1
const CS_COUPLING_DC = 0x01
const CS_COUPLING_AC = 0x02
const CS_COUPLING_MASK = 0x03
const CS_DIFFERENTIAL_INPUT = 0x04
const CS_DIRECT_ADC_INPUT = 0x08
const CS_REAL_IMP_1M_OHM = 1000000
const CS_REAL_IMP_50_OHM = 50
const CS_GAIN_100_V = 100000
const CS_GAIN_40_V = 40000
const CS_GAIN_20_V = 20000
const CS_GAIN_10_V = 10000
const CS_GAIN_8_V = 8000
const CS_GAIN_5_V = 5000
const CS_GAIN_4_V = 4000
const CS_GAIN_2_V = 2000
const CS_GAIN_1_V = 1000
const CS_GAIN_800_MV = 800
const CS_GAIN_500_MV = 500
const CS_GAIN_400_MV = 400
const CS_GAIN_200_MV = 200
const CS_GAIN_100_MV = 100
const CS_GAIN_50_MV = 50
const CS_GAIN_CMOS = 2500
const CS_GAIN_TTL = 1500
const CS_GAIN_PECL = 3000
const CS_GAIN_LVDS = 0
const CS_GAIN_ECL = -2000
const CS_TRIG_COND_NEG_SLOPE = 0
const CS_TRIG_COND_POS_SLOPE = 1
const CS_TRIG_COND_PULSE_WIDTH = 2
const CS_MAX_TRIG_COND = 3
const CS_RELATION_OR = 0
const CS_RELATION_AND = 1
const CS_EVENT1 = 2
const CS_EVENT2 = 3
const CS_TRIG_ENGINES_DIS = 0
const CS_TRIG_ENGINES_EN = 1
const CS_TRIG_SOURCE_DISABLE = 0
const CS_TRIG_SOURCE_CHAN_1 = 1
const CS_TRIG_SOURCE_CHAN_2 = 2
const CS_TRIG_SOURCE_EXT = -1
const CS_BOARD_INFO = 1
const CS_SYSTEM = 3
const CS_CHANNEL = 4
const CS_TRIGGER = 5
const CS_ACQUISITION = 6
const CS_PARAMS = 7
const CS_FIR_CONFIG = 8
const CS_EXTENDED_BOARD_OPTIONS = 9
const CS_TIMESTAMP_TICKFREQUENCY = 10
const CS_CHANNEL_ARRAY = 11
const CS_TRIGGER_ARRAY = 12
const CS_SMT_ENVELOPE_CONFIG = 13
const CS_SMT_HISTOGRAM_CONFIG = 14
const CS_FFT_CONFIG = 15
const CS_FFTWINDOW_CONFIG = 16
const CS_MULREC_AVG_COUNT = 17
const CS_TRIG_OUT_CFG = 18
const CS_GET_SEGMENT_COUNT = 19
const CS_CURRENT_CONFIGURATION = 1
const CS_ACQUISITION_CONFIGURATION = 2
const CS_ACQUIRED_CONFIGURATION = 3
const ACTION_COMMIT = 1
const ACTION_START = 2
const ACTION_FORCE = 3
const ACTION_ABORT = 4
const ACTION_CALIB = 5
const ACTION_RESET = 6
const ACTION_COMMIT_COERCE = 7
const ACTION_TIMESTAMP_RESET = 8
const ACQ_STATUS_READY = 0
const ACQ_STATUS_WAIT_TRIGGER = 1
const ACQ_STATUS_TRIGGERED = 2
const ACQ_STATUS_BUSY_TX = 3
const ACQ_STATUS_BUSY_CALIB = 4
const CSDEMO8_BOARDTYPE = 0x0801
const CSDEMO12_BOARDTYPE = 0x0802
const CSDEMO14_BOARDTYPE = 0x0803
const CSDEMO16_BOARDTYPE = 0x0804
const CSDEMO32_BOARDTYPE = 0x0805
const CSDEMO_BT_MASK = 0x0800
const CSDEMO_BT_FIRST_BOARD = CSDEMO8_BOARDTYPE
const CSDEMO_BT_LAST_BOARD = CSDEMO32_BOARDTYPE
const CS8500_BOARDTYPE = 0x01
const CS82G_BOARDTYPE = 0x02
const CS12100_BOARDTYPE = 0x03
const CS1250_BOARDTYPE = 0x04
const CS1220_BOARDTYPE = 0x05
const CS1210_BOARDTYPE = 0x06
const CS14100_BOARDTYPE = 0x07
const CS1450_BOARDTYPE = 0x08
const CS1610_BOARDTYPE = 0x09
const CS1602_BOARDTYPE = 0x0a
const CS3200_BOARDTYPE = 0x0b
const CS85G_BOARDTYPE = 0x0c
const CS14200_BOARDTYPE = 0x0d
const CS14105_BOARDTYPE = 0x0e
const CS12400_BOARDTYPE = 15
const CS14200v2_BOARDTYPE = 16
const SNA142_BOARDTYPE = 17
const CS_COMBINE_FIRST_BOARD = CS14200_BOARDTYPE
const CS_COMBINE_LAST_BOARD = SNA142_BOARDTYPE
const CS16x1_BOARDTYPE = 0x20
const CS16x2_BOARDTYPE = 0x21
const CS14x1_BOARDTYPE = 0x22
const CS14x2_BOARDTYPE = 0x23
const CS12x1_BOARDTYPE = 0x24
const CS12x2_BOARDTYPE = 0x25
const CS16xyy_BT_FIRST_BOARD = CS16x1_BOARDTYPE
const CS16xyy_LAST_BOARD = CS12x2_BOARDTYPE
const CS16xyy_BASEBOARD = CS16x1_BOARDTYPE
const CS8220_BOARDTYPE = 0x40
const CS8221_BOARDTYPE = 0x41
const CS8222_BOARDTYPE = 0x42
const CS8223_BOARDTYPE = 0x43
const CS8224_BOARDTYPE = 0x44
const CS8225_BOARDTYPE = 0x45
const CS8226_BOARDTYPE = 0x46
const CS8227_BOARDTYPE = 0x47
const CS8228_BOARDTYPE = 0x48
const CS8229_BOARDTYPE = 0x49
const CS822x_BOARDTYPE_MASK = 0x40
const CS8240_BOARDTYPE = 0x50
const CS8241_BOARDTYPE = 0x51
const CS8242_BOARDTYPE = 0x52
const CS8243_BOARDTYPE = 0x53
const CS8244_BOARDTYPE = 0x54
const CS8245_BOARDTYPE = 0x55
const CS8246_BOARDTYPE = 0x56
const CS8247_BOARDTYPE = 0x57
const CS8248_BOARDTYPE = 0x58
const CS8249_BOARDTYPE = 0x59
const CS824x_BOARDTYPE_MASK = 0x50
const CS8280_BOARDTYPE = 0x60
const CS8281_BOARDTYPE = 0x61
const CS8282_BOARDTYPE = 0x62
const CS8283_BOARDTYPE = 0x63
const CS8284_BOARDTYPE = 0x64
const CS8285_BOARDTYPE = 0x65
const CS8286_BOARDTYPE = 0x66
const CS8287_BOARDTYPE = 0x67
const CS8288_BOARDTYPE = 0x68
const CS8289_BOARDTYPE = 0x69
const CS828x_BOARDTYPE_MASK = 0x60
const CS8320_BOARDTYPE = 0x80
const CS8321_BOARDTYPE = 0x81
const CS8322_BOARDTYPE = 0x82
const CS8323_BOARDTYPE = 0x83
const CS8324_BOARDTYPE = 0x84
const CS8325_BOARDTYPE = 0x85
const CS8326_BOARDTYPE = 0x86
const CS8327_BOARDTYPE = 0x87
const CS8328_BOARDTYPE = 0x88
const CS8329_BOARDTYPE = 0x89
const CS832x_BOARDTYPE_MASK = 0x80
const CS8340_BOARDTYPE = 0x90
const CS8341_BOARDTYPE = 0x91
const CS8342_BOARDTYPE = 0x92
const CS8343_BOARDTYPE = 0x93
const CS8344_BOARDTYPE = 0x94
const CS8345_BOARDTYPE = 0x95
const CS8346_BOARDTYPE = 0x96
const CS8347_BOARDTYPE = 0x97
const CS8348_BOARDTYPE = 0x98
const CS8349_BOARDTYPE = 0x99
const CS834x_BOARDTYPE_MASK = 0x90
const CS8380_BOARDTYPE = 0xa0
const CS8381_BOARDTYPE = 0xa1
const CS8382_BOARDTYPE = 0xa2
const CS8383_BOARDTYPE = 0xa3
const CS8384_BOARDTYPE = 0xa4
const CS8385_BOARDTYPE = 0xa5
const CS8386_BOARDTYPE = 0xa6
const CS8387_BOARDTYPE = 0xa7
const CS8388_BOARDTYPE = 0xa8
const CS8389_BOARDTYPE = 0xa9
const CS838x_BOARDTYPE_MASK = 0xa0
const CS8210_BOARDTYPE = 0xb0
const CS8211_BOARDTYPE = 0xb1
const CS8212_BOARDTYPE = 0xb2
const CS8213_BOARDTYPE = 0xb3
const CS8214_BOARDTYPE = 0xb4
const CS8215_BOARDTYPE = 0xb5
const CS8216_BOARDTYPE = 0xb6
const CS8217_BOARDTYPE = 0xb7
const CS8218_BOARDTYPE = 0xb8
const CS8219_BOARDTYPE = 0xb9
const CS821x_BOARDTYPE_MASK = 0xb0
const CS8310_BOARDTYPE = 0xc0
const CS8311_BOARDTYPE = 0xc1
const CS8312_BOARDTYPE = 0xc2
const CS8313_BOARDTYPE = 0xc3
const CS8314_BOARDTYPE = 0xc4
const CS8315_BOARDTYPE = 0xc5
const CS8316_BOARDTYPE = 0xc6
const CS8317_BOARDTYPE = 0xc7
const CS8318_BOARDTYPE = 0xc8
const CS8319_BOARDTYPE = 0xc9
const CS831x_BOARDTYPE_MASK = 0xc0
const CS8410_BOARDTYPE = 0xd0
const CS8412_BOARDTYPE = 0xd2
const CS8420_BOARDTYPE = 0xd4
const CS8422_BOARDTYPE = 0xd6
const CS8440_BOARDTYPE = 0xd8
const CS8442_BOARDTYPE = 0xda
const CS8480_BOARDTYPE = 0xdc
const CS8482_BOARDTYPE = 0xde
const CS84xx_BOARDTYPE_MASK = 0xd0
const CS8_BT_MASK = 0xd0
const CS8_BT_FIRST_BOARD = CS8220_BOARDTYPE
const CS8_BT_LAST_BOARD = CS8482_BOARDTYPE
const CS8xxx_BASEBOARD = CS8_BT_FIRST_BOARD
const CS_ZAP_FIRST_BOARD = CS8410_BOARDTYPE
const CS_ZAP_LAST_BOARD = CS8482_BOARDTYPE
const CS_ZAP_BASEBOARD = CS_ZAP_FIRST_BOARD
const CSxyG8_BOARDTYPE_MASK = 0x0100
const CS22G8_BOARDTYPE = 0x0100
const CS21G8_BOARDTYPE = 0x0101
const CS11G8_BOARDTYPE = 0x0102
const LAB11G_BOARDTYPE = 0x0103
const CSxyG8_FIRST_BOARD = CS22G8_BOARDTYPE
const CSxyG8_LAST_BOARD = LAB11G_BOARDTYPE
const CS_BASE8_BOARDTYPE = 0x0400
const CS_COBRAX_BOARDTYPE_MASK = 0x0480
const CSX11G8_BOARDTYPE = 0x0481
const CSX21G8_BOARDTYPE = 0x0482
const CSX22G8_BOARDTYPE = 0x0483
const CSX13G8_BOARDTYPE = 0x0484
const CSX23G8_BOARDTYPE = 0x0485
const CSX14G8_BOARDTYPE = 0x0486
const CSX_NTS_BOARDTYPE = 0x0487
const CS12500_BOARDTYPE = 0x0500
const CS121G_BOARDTYPE = 0x0501
const CS122G_BOARDTYPE = 0x0502
const CS14250_BOARDTYPE = 0x0510
const CS14500_BOARDTYPE = 0x0511
const CS141G_BOARDTYPE = 0x0512
const CSJD_FIRST_BOARD = CS12500_BOARDTYPE
const CSJD_LAST_BOARD = CS141G_BOARDTYPE
const CSUSB_BOARDTYPE = 0x1001
const CSUSB_BT_FIRST_BOARD = CSUSB_BOARDTYPE
const CSUSB_BT_LAST_BOARD = CSUSB_BOARDTYPE
const PLX_BASED_BOARDTYPE = 0x4000
const CPCI_BOARDTYPE = 0x8000
const CS1610C_BOARDTYPE = CPCI_BOARDTYPE | CS1610_BOARDTYPE
const CS14100C_BOARDTYPE = CPCI_BOARDTYPE | CS14100_BOARDTYPE
const CS82GC_BOARDTYPE = CPCI_BOARDTYPE | CS82G_BOARDTYPE
const CS85GC_BOARDTYPE = CPCI_BOARDTYPE | CS85G_BOARDTYPE
const CS3200C_BOARDTYPE = CPCI_BOARDTYPE | CS3200_BOARDTYPE
const CS_TIMEOUT_DISABLE = -1
const ACQ_EVENT_TRIGGERED = 0
const ACQ_EVENT_END_BUSY = 1
const ACQ_EVENT_END_TXFER = 2
const NUMBER_ACQ_EVENTS = 3
const CAPS_SAMPLE_RATES = 0x00010000
const CAPS_INPUT_RANGES = 0x00020000
const CAPS_IMPEDANCES = 0x00030000
const CAPS_COUPLINGS = 0x00040000
const CAPS_ACQ_MODES = 0x00050000
const CAPS_TERMINATIONS = 0x00060000
const CAPS_FLEXIBLE_TRIGGER = 0x00070000
const CAPS_BOARD_TRIGGER_ENGINES = 0x00080000
const CAPS_TRIGGER_SOURCES = 0x00090000
const CAPS_FILTERS = 0x000a0000
const CAPS_MAX_SEGMENT_PADDING = 0x000b0000
const CAPS_DC_OFFSET_ADJUST = 0x000c0000
const CAPS_CLK_IN = 0x000d0000
const CAPS_TRIG_ENGINES_PER_CHAN = 0x00200000
const CAPS_MULREC = 0x00400000
const CAPS_TRIGGER_RES = 0x00410000
const CAPS_MIN_EXT_RATE = 0x00420000
const CAPS_SKIP_COUNT = 0x00430000
const CAPS_MAX_EXT_RATE = 0x00440000
const CAPS_TRANSFER_EX = 0x00450000
const CS_NO_FILTER = Float32(0x0fffffff)
const CS_NO_EXTTRIG = Float32(0x0fffffff)

# Skipping MacroDefinition: CS_FAILED ( x ) ( ( x ) < 0 ? TRUE : FALSE )
# Skipping MacroDefinition: CS_SUCCEEDED ( x ) ( ( x ) < 0 ? FALSE : TRUE )

const CS_USING_DEFAULT_ACQ_DATA = 2
const CS_USING_DEFAULT_CHANNEL_DATA = 4
const CS_USING_DEFAULT_TRIGGER_DATA = 8
const CS_USING_DEFAULT_APP_DATA = 16
const CS_USING_DEFAULTS = 32
const ACQUISITION_DATA = 1
const CHANNEL_DATA = 2
const TRIGGER_DATA = 3
const APPLICATION_DATA = 4
const DEFAULT_TRANSFER_START_POSITION = 0
const DEFAULT_TRANSFER_LENGTH = 4096
const DEFAULT_TRANSFER_SEGMENT_START = 1
const DEFAULT_TRANSFER_SEGMENT_COUNT = 1
const DEFAULT_PAGE_SIZE = 32768

const uInt8 = Cuchar
const uInt16 = UInt16
const uInt32 = Culong
const int8 = UInt8
const int16 = Int16
const int32 = Clong
const bool = UInt8
const int64 = Clonglong
const real32 = Cfloat
const real64 = Cdouble
const cschar = UInt8
const CSHANDLE = uInt32
const PCSHANDLE = Ptr{CSHANDLE}
const RMHANDLE = uInt32
const PRMHANDLE = Ptr{RMHANDLE}
const DRVHANDLE = uInt32
const PDRVHANDLE = Ptr{DRVHANDLE}

# begin enum _CSFILETYPES
const _CSFILETYPES = Cint
const cftNone = 0 |> Int32
const cftBinaryData = 1 |> Int32
const cftSigFile = 2 |> Int32
const cftAscii = 3 |> Int32
# end enum _CSFILETYPES

const CsFileType = Cvoid

# begin enum _CSFILEMODES
const _CSFILEMODES = Cint
const cfmRead = 1 |> Int32
const cfmWrite = 2 |> Int32
const cfmReadWrite = 3 |> Int32
# end enum _CSFILEMODES

const CsFileMode = Cvoid

# begin enum _CSFILEATTACH
const _CSFILEATTACH = Cint
const fmNew = 0 |> Int32
# end enum _CSFILEATTACH

const CsFileAttach = Cvoid

mutable struct CSBOARDINFO
    u32Size::uInt32
    u32BoardIndex::uInt32
    u32BoardType::uInt32
    strSerialNumber::NTuple{32, cschar}
    u32BaseBoardVersion::uInt32
    u32BaseBoardFirmwareVersion::uInt32
    u32AddonBoardVersion::uInt32
    u32AddonBoardFirmwareVersion::uInt32
    u32AddonFwOptions::uInt32
    u32BaseBoardFwOptions::uInt32
    u32AddonHwOptions::uInt32
    u32BaseBoardHwOptions::uInt32
end

const PCSBOARDINFO = Ptr{CSBOARDINFO}

mutable struct ARRAY_BOARDINFO
    u32BoardCount::uInt32
    aBoardInfo::NTuple{1, CSBOARDINFO}
end

const PARRAY_BOARDINFO = Ptr{ARRAY_BOARDINFO}

mutable struct CSSYSTEMINFO
    u32Size::uInt32
    i64MaxMemory::int64
    u32SampleBits::uInt32
    i32SampleResolution::int32
    u32SampleSize::uInt32
    i32SampleOffset::int32
    u32BoardType::uInt32
    strBoardName::NTuple{32, cschar}
    u32AddonOptions::uInt32
    u32BaseBoardOptions::uInt32
    u32TriggerMachineCount::uInt32
    u32ChannelCount::uInt32
    u32BoardCount::uInt32
    CSSYSTEMINFO() = (new(sizeof(CSSYSTEMINFO),0,0,0,0,0,0,NTuple{32, cschar}(),0,0,0,0,0))
end

const PCSSYSTEMINFO = Ptr{CSSYSTEMINFO}

mutable struct CSACQUISITIONCONFIG
    u32Size::uInt32
    i64SampleRate::int64
    u32ExtClk::uInt32
    u32ExtClkSampleSkip::uInt32
    u32Mode::uInt32
    u32SampleBits::uInt32
    i32SampleRes::int32
    u32SampleSize::uInt32
    u32SegmentCount::uInt32
    i64Depth::int64
    i64SegmentSize::int64
    i64TriggerTimeout::int64
    u32TrigEnginesEn::uInt32
    i64TriggerDelay::int64
    i64TriggerHoldoff::int64
    i32SampleOffset::int32
    u32TimeStampConfig::uInt32
    CSACQUISITIONCONFIG() = (new(sizeof(CSACQUISITIONCONFIG),0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0))
end

const PCSACQUISITIONCONFIG = Ptr{CSACQUISITIONCONFIG}

mutable struct CSCHANNELCONFIG
    u32Size::uInt32
    u32ChannelIndex::uInt32
    u32Term::uInt32
    u32InputRange::uInt32
    u32Impedance::uInt32
    u32Filter::uInt32
    i32DcOffset::int32
    i32Calib::int32
    CSCHANNELCONFIG() = (new(sizeof(CSCHANNELCONFIG),1,0,0,0,0,0,0))
end

const PCSCHANNELCONFIG = Ptr{CSCHANNELCONFIG}

mutable struct ARRAY_CHANNELCONFIG
    u32ChannelCount::uInt32
    aChannel::NTuple{1, CSCHANNELCONFIG}
end

const PARRAY_CHANNELCONFIG = Ptr{ARRAY_CHANNELCONFIG}

mutable struct CSTRIGGERCONFIG
    u32Size::uInt32
    u32TriggerIndex::uInt32
    u32Condition::uInt32
    i32Level::int32
    i32Source::int32
    u32ExtCoupling::uInt32
    u32ExtTriggerRange::uInt32
    u32ExtImpedance::uInt32
    i32Value1::int32
    i32Value2::int32
    u32Filter::uInt32
    u32Relation::uInt32
    CSTRIGGERCONFIG() = (new(sizeof(CSTRIGGERCONFIG),1,0,0,0,0,0,0,0,0,0,0))
end

const PCSTRIGGERCONFIG = Ptr{CSTRIGGERCONFIG}

mutable struct ARRAY_TRIGGERCONFIG
    u32TriggerCount::uInt32
    aTrigger::NTuple{1, CSTRIGGERCONFIG}
end

const PARRAY_TRIGGERCONFIG = Ptr{ARRAY_TRIGGERCONFIG}

mutable struct CSSYSTEMCONFIG
    u32Size::uInt32
    pAcqCfg::PCSACQUISITIONCONFIG
    pAChannels::PARRAY_CHANNELCONFIG
    pATriggers::PARRAY_TRIGGERCONFIG
end

const PCSSYSTEMCONFIG = Ptr{CSSYSTEMCONFIG}

mutable struct CSSAMPLERATETABLE
    i64SampleRate::int64
    strText::NTuple{32, cschar}
end

const PCSSAMPLERATETABLE = Ptr{CSSAMPLERATETABLE}

mutable struct CSRANGETABLE
    u32InputRange::uInt32
    strText::NTuple{32, cschar}
    u32Reserved::uInt32
end

const PCSRANGETABLE = Ptr{Cvoid}

mutable struct CSIMPEDANCETABLE
    u32Imdepdance::uInt32
    strText::NTuple{32, cschar}
    u32Reserved::uInt32
end

const PCSIMPEDANCETABLE = Ptr{Cvoid}

mutable struct CSCOUPLINGTABLE
    u32Coupling::uInt32
    strText::NTuple{32, cschar}
    u32Reserved::uInt32
end

const PCSCOUPLINGTABLE = Ptr{Cvoid}

mutable struct CSFILTERTABLE
    u32LowCutoff::uInt32
    u32HighCutoff::uInt32
end

const PCSFILTERTABLE = Ptr{CSFILTERTABLE}

mutable struct IN_PARAMS_TRANSFERDATA
    u16Channel::uInt16
    u32Mode::uInt32
    u32Segment::uInt32
    i64StartAddress::int64
    i64Length::int64
    pDataBuffer::Ptr{Cvoid}
    hNotifyEvent::Ptr{Cint}
end
IN_PARAMS_TRANSFERDATA() = new(zeros())

const PIN_PARAMS_TRANSFERDATA = Ptr{IN_PARAMS_TRANSFERDATA}

mutable struct OUT_PARAMS_TRANSFERDATA
    i64ActualStart::int64
    i64ActualLength::int64
    i32LowPart::int32
    i32HighPart::int32
end

const POUT_PARAMS_TRANSFERDATA = Ptr{OUT_PARAMS_TRANSFERDATA}

mutable struct IN_PARAMS_TRANSFERDATA_EX
    u16Channel::uInt16
    u32Mode::uInt32
    u32StartSegment::uInt32
    u32SegmentCount::uInt32
    i64StartAddress::int64
    i64Length::int64
    pDataBuffer::Ptr{Cvoid}
    i64BufferLength::int64
    hNotifyEvent::Ptr{Cint}
end

const PIN_PARAMS_TRANSFERDATA_EX = Ptr{Cvoid}

mutable struct OUT_PARAMS_TRANSFERDATA_EX
    u32DataFormat0::uInt32
    u32DataFormat1::uInt32
end

const POUT_PARAMS_TRANSFERDATA_EX = Ptr{OUT_PARAMS_TRANSFERDATA_EX}

mutable struct CALLBACK_DATA
    u32Size::uInt32
    hSystem::CSHANDLE
    u32ChannelIndex::uInt32
    i32Token::int32
end

const PCALLBACK_DATA = Ptr{CALLBACK_DATA}

mutable struct CSTIMESTAMP
    u16Hour::uInt16
    u16Minute::uInt16
    u16Second::uInt16
    u16Point1Second::uInt16
end

const PCSTIMESTAMP = Ptr{CSTIMESTAMP}

mutable struct CSSIGSTRUCT
    u32Size::uInt32
    i64SampleRate::int64
    i64RecordStart::int64
    i64RecordLength::int64
    u32RecordCount::uInt32
    u32SampleBits::uInt32
    u32SampleSize::uInt32
    i32SampleOffset::int32
    i32SampleRes::int32
    u32Channel::uInt32
    u32InputRange::uInt32
    i32DcOffset::int32
    TimeStamp::CSTIMESTAMP
end

const PCSSIGSTRUCT = Ptr{CSSIGSTRUCT}

mutable struct CSDISKFILEHEADER
    cData::NTuple{512, UInt8}
end

const PCSDISKFILEHEADER = Ptr{CSDISKFILEHEADER}

mutable struct CSSEGMENTTAIL
    i64TimeStamp::int64
    i64Reserved0::int64
    i64Reserved1::int64
    i64Reserved2::int64
end

const PCSSEGMENTTAIL = Ptr{CSSEGMENTTAIL}

mutable struct _FILEHEADERSTRUCT
    i64Start::int64
    i64Length::int64
    i64SampleRate::int64
    dTimeStamp::Cdouble
    u32SegmentNumber::uInt32
    u32SampleSize::uInt32
    i32SampleRes::int32
    i32SampleOffset::int32
    u32InputRange::uInt32
    i32DcOffset::int32
    u32SegmentCount::uInt32
    u32AverageCount::uInt32
end

const FileHeaderStruct = Cvoid
const CSAPPLICATIONDATA = Cvoid
const PCSAPPLICATIONDATA = Ptr{Cvoid}


text:
export SecondOrderCone

struct SecondOrderCone{T <: Real} <: ConvexSet
    P::Matrix{T}
    q::Vector{T}
    r::Vector{T}
    s::T
end

function SecondOrder(y)
#
# This function projects the point y = (x,t) onto the second
# order cone {(x,t)|norm(x)<=t}.
#
   n = norm(y[1:end-1])
   if n <= y[end]
      return y
   elseif n <= -y[end]
      return zeros(size(y))
   else
      z = zeros(size(y))
      r = 0.5*(1.0+y[end]/n)
      z[end] = r*n
      z[1:end-1] = r*y[1:end-1]
      return z
   end
end

function AffineProjection(y,A,b,cached=false,C=zeros(size(A)),d=zeros(size(b)))
#
# This function projects the point y into the affine set
# {x|Ax = b} based on the projection formula
# z = y-pinv(A)*(A*y-b). You should call the function via
# z, C, d = AffineProjection(y1,A,b)
# and then call it again with the new argument y2 via
# z = AffineProjection(y2,A,b,true,C,d).
# Passing back the cached values C and d speeds up
# calculation of the projected point z.
#
   if cached
      z = C*y+d
      return z
   else
      pA = pinv(A)
      C = eye(length(y))-pA*A
      d = pA*b
      z = C*y+d
      return (z,C,d)
   end
end

function SecondOrderCone(s::SecondOrderCone{<:Real}, v::Vector{<:Real}, z::Vector{<:Real})
#
# This function projects the point z onto the second-order cone
# constraint {x:||P*x+q||<=r'*x+s}.
#
# Initialize control constants.
#
   rho = 1.0
   rho_inc = 1.05
   rho_max = 10
   eps = 1e-14
   eps_dec = 1.2
   eps_min = 1e-14
   tol = 1e-8
   max_iter = 500
   m = size(s.P,1)
   n = size(s.P,2)
#
# Check whether the point z is already feasible.
#
   u = s.P*z + s.q
   t = s.r'*z + s.s
   if norm(u) <= t
      v = z
   end
#
# Initialize vectors and matrices.
#
   x = [z,u,t]
   y = [z,u,t]
   M = [P -eye(m) zeros(m); r' zeros(m)' -1]
   v = -[s.q,s.s]
   zx = zeros(m+n+1)
#
# Prepare to iterate.
#
   ut = SecondOrder(x[n+1:end])
   x = [x[1:n],ut]
   (y,E,f) = AffineProjection(y,M,v)
   a = sqrt(norm(z-x[1:n])^2+1.0)
   b = sqrt(norm(z-y[1:n])^2+1.0)
   loss = 0.5*(a+b)
   dist = norm(x-y)/sqrt(2.0)
   objective = loss + rho*sqrt(dist^2 + eps)
#
# Enter the alternating projection loop.
#
   for i = 1:max_iter
      println(i-1," ",objective," ",dist," ",rho," ",eps)
#
# Project onto the second-order cone.
#
      c = 0.5/a
      w = 0.5*rho/sqrt(dist^2+eps)
      x[1:n] = (c/(c+w))*z+(w/(c+w))*y[1:n]
      ut = (c/(c+w))*x[n+1:end]+(w/(c+w))*y[n+1:end]
      x[n+1:end] = SecondOrder(ut)
#
# Project onto the affine set.
#
      dist = norm(x-y)/sqrt(2.0)
      d = 0.5/b
      w = 0.5*rho/sqrt(dist^2+eps)
      zx = (d/(d+w))*[z,y[n+1:end]]+(w/(d+w))*x
      y = AffineProjection(zx,M,v,true,E,f)
#
# Update the objective.
#
      a = sqrt(norm(z-x[1:n])^2+1.0)
      b = sqrt(norm(z-y[1:n])^2+1.0)
      loss = 0.5*(a+b)
      dist = norm(x-y)/sqrt(2.0)
      next_objective = loss + rho*sqrt(dist^2 + eps)
#
# Check for convergence.
#
      if abs(next_objective - objective)/(objective + 1.0) < tol && dist < tol
         break
      else
         rho = min(rho_inc*rho,rho_max)
         eps = max(eps/eps_dec,eps_min)
         objective = next_objective
      end
   end
   v = 0.5*(x[1:n]+y[1:n])
end

function project(s::SecondOrderCone{<:Real}, y::Vector{<:Real})
    project!(s, similar(y), y)
end


text:
@with_kw mutable struct DRegularizedData{T <: Real}
    Q::T = 1e10
    Q̃::T = 1e10
    θ::T = -1e10
    σ::T = 1.0
    timestamp::Int = 1
    iterations::Int = 0
    major_iterations::Int = 0
    minor_iterations::Int = 0
    regularizerindex::Int = -1
end

@with_kw mutable struct DRegularizedParameters{T <: Real}
    κ::T = 0.6
    τ::T = 1e-6
    γ::T = 0.9
    σ::T = 1.0
    σ̅::T = 4.0
    σ̲::T = 0.5
    bundle::Int = 1
    log::Bool = true
    autotune::Bool = false
    linearize::Bool = false
end

"""
    DRegularized

Functor object for the distributed regularized L-shaped algorithm. Create by supplying `:drd` to the `LShapedSolver` factory function and then pass to a `StochasticPrograms.jl` model, assuming there are available worker cores.

...
# Algorithm parameters
- `κ::Real = 0.6`: Amount of cutting planes, relative to the total number of scenarios, required to generate a new iterate in master procedure.
- `τ::Real = 1e-6`: Relative tolerance for convergence checks.
- `σ::Real = 1.0`: Initial value of regularization parameter. Controls the relative penalty of the deviation from the current major iterate.
- `σ̅::real = 4.0`: Maximum value of the regularization parameter.
- `σ̲::real = 0.5`: Minimum value of the regularization parameter.
- `bundle::Int = 1`: Amount of cutting planes in bundle. A value of 1 corresponds to a multicut algorithm and a value of at least the number of scenarios yields the classical L-shaped algorithm.
- `log::Bool = true`: Specifices if L-shaped procedure should be logged on standard output or not.
- `autotune::Bool = false`: If `true`, heuristic methods are used to set `σ, σ̅` and `σ̲` based on the initial decision.
- `linearize::Bool = false`: If `true`, the quadratic terms in the master problem objective are linearized through a ∞-norm approximation.
...
"""
struct DRegularized{F, T <: Real, A <: AbstractVector, SP <: StochasticProgram, M <: LQSolver, S <: LQSolver} <: AbstractLShapedSolver{F,T,A,M,S}
    stochasticprogram::SP
    solverdata::DRegularizedData{T}

    # Master
    mastersolver::M
    mastervector::A
    c::A
    x::A
    Q_history::A

    # Subproblems
    nscenarios::Int
    subobjectives::Vector{A}
    finished::Vector{Int}

    # Workers
    subworkers::Vector{SubWorker{F,T,A,S}}
    work::Vector{Work}
    decisions::Decisions{A}
    cutqueue::CutQueue{T}
    active_workers::Vector{Future}

    # Trust region
    ξ::A
    Q̃_history::A
    σ_history::A

    # Cuts
    θs::A
    cuts::Vector{SparseHyperPlane{T}}
    θ_history::A

    # Params
    parameters::DRegularizedParameters{T}
    progress::ProgressThresh{T}

    @implement_trait DRegularized RD
    @implement_trait DRegularized Parallel

    function (::Type{DRegularized})(stochasticprogram::StochasticProgram, ξ₀::AbstractVector, mastersolver::MPB.AbstractMathProgSolver, subsolver::SubSolver, F::Bool; kw...)
        if nworkers() == 1
            @warn "There are no worker processes, defaulting to serial version of algorithm"
            return Regularized(stochasticprogram, ξ₀, mastersolver, get_solver(subsolver); kw...)
        end
        first_stage = StochasticPrograms.get_stage_one(stochasticprogram)
        length(ξ₀) != first_stage.numCols && error("Incorrect length of starting guess, has ", length(ξ₀), " should be ", first_stage.numCols)

        T = promote_type(eltype(ξ₀), Float32)
        c_ = convert(AbstractVector{T}, JuMP.prepAffObjective(first_stage))
        c_ *= first_stage.objSense == :Min ? 1 : -1
        ξ₀_ = convert(AbstractVector{T}, copy(ξ₀))
        x₀_ = convert(AbstractVector{T}, copy(ξ₀))
        mastervector = convert(AbstractVector{T}, copy(ξ₀))
        A = typeof(ξ₀_)
        SP = typeof(stochasticprogram)
        msolver = LQSolver(first_stage, mastersolver)
        M = typeof(msolver)
        solver_instance = get_solver(subsolver)
        S = LQSolver{typeof(MPB.LinearQuadraticModel(solver_instance)),typeof(solver_instance)}
        n = StochasticPrograms.nscenarios(stochasticprogram)

        lshaped = new{F,T,A,SP,M,S}(stochasticprogram,
                                    DRegularizedData{T}(),
                                    msolver,
                                    mastervector,
                                    c_,
                                    x₀_,
                                    A(),
                                    n,
                                    Vector{A}(),
                                    Vector{Int}(),
                                    Vector{SubWorker{F,T,A,S}}(undef,nworkers()),
                                    Vector{Work}(undef,nworkers()),
                                    RemoteChannel(() -> DecisionChannel(Dict{Int,A}())),
                                    RemoteChannel(() -> Channel{QCut{T}}(4*nworkers()*n)),
                                    Vector{Future}(undef,nworkers()),
                                    ξ₀_,
                                    A(),
                                    A(),
                                    A(),
                                    Vector{SparseHyperPlane{T}}(),
                                    A(),
                                    DRegularizedParameters{T}(;kw...),
                                    ProgressThresh(1.0, "Distributed RD L-Shaped Gap "))
        # Initialize solver
        init!(lshaped, subsolver)
        return lshaped
    end
end
DRegularized(stochasticprogram::StochasticProgram, mastersolver::MPB.AbstractMathProgSolver, subsolver::MPB.AbstractMathProgSolver, checkfeas::Bool; kw...) = DRegularized(stochasticprogram, rand(decision_length(stochasticprogram)), mastersolver, subsolver, checkfeas; kw...)

function (lshaped::DRegularized)()
    # Reset timer
    lshaped.progress.tfirst = lshaped.progress.tlast = time()
    # Start workers
    init_workers!(lshaped)
    # Start procedure
    while true
        status = iterate!(lshaped)
        if status != :Valid
            close_workers!(lshaped)
            return status
        end
    end
end


text:
type DirectSolver <: Solver
    A
    function DirectSolver(A::AbstractMatrix)
        new(factorize(A))
    end
end

solve!(x, S::DirectSolver) =  A_ldiv_B!(S.A, x)

export DirectSolver, solve!


text:
### A Pluto.jl notebook ###
# v0.12.123123

using Markdown
using InteractiveUtils

# ╔═╡ ec425e6e-6b9c-11eb-0c63-97fade58f6b4
md"Hello!"

# ╔═╡ Cell order:
# ╠═ec425e6e-6b9c-11eb-0c63-97fade58f6b4


text:
# functions to read-in the data

function readMP(fileAdd::String)
    # read in the file from a specific address
    # input: the address of the Matpower data file
    # output:
    #   busST: the string containing bus information
    #   genST: the string containing generator information
    #   brST: the string containing branch information
    #   cST: the string containing cost information
    #   baseMVA: the base unit

    f = open(fileAdd);
    rawStr = read(f,String);
    close(f);

    # separate the raw strings for bus/gen/branch/cost
    if typeof(match(r"mpc.bus = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr)) != Nothing
        busST = match(r"mpc.bus = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr).captures[1];
    else
        busST = "";
    end
    if typeof(match(r"mpc.gen = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr)) != Nothing
        genST = match(r"mpc.gen = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr).captures[1];
    else
        genST = "";
    end
    if typeof(match(r"mpc.branch = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr)) != Nothing
        brST = match(r"mpc.branch = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr).captures[1];
    else
        brST = "";
    end
    if typeof(match(r"mpc.gencost = \[([0-9\t\ \.\;\r\n\-e]*)\];",rawStr)) != Nothing
        cST = match(r"mpc.gencost = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr).captures[1];
    else
        cST = "";
    end
    if typeof(match(r"mpc.uncertain = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr)) != Nothing
        uST = match(r"mpc.uncertain = \[([0-9\t\ \.\;\n\r\-e]*)\];",rawStr).captures[1];
    else
        uST = ""
    end

    baseMVA = parse(Float64,match(r"mpc.baseMVA = ([0-9\.]+);",rawStr).captures[1]);

    return busST,genST,brST,cST,uST,baseMVA;
end

function parsebusST(busST,baseMVA)
    # parse the busST string collected from readMP
    # input: the strings containing bus information
    # output: List of bus ID, list of generator ID, data variable Vmax/Vmin

    # separate the raw strings line by line
    busSTlbl = collect(eachmatch(r"([0-9\t\ \.\-e]+[\;\r]*\n)",busST));

    # parse each line of the bus data
    Vmax = Dict();
    Vmin = Dict();
    bType = Dict();
    gs = Dict();
    bs = Dict();
    Vmag = Dict();
    Vang = Dict();
    Pd = Dict();
    Qd = Dict();
    IDList = [];
    for bstr in busSTlbl
        bsstr = bstr.captures[1];
        bsstr = strip(strip(bsstr),';');
        bdata = split(bsstr);

        # update the IDList with the ID of the current bus
        cID = parse(Int64,bdata[1]);
        push!(IDList,cID);

        # obtain the type of the bus
        bType[cID] = parse(Int64,bdata[2]);

        # obtain the active/reactive power demand
        Pd[cID] = parse(Float64,bdata[3])/baseMVA;
        Qd[cID] = parse(Float64,bdata[4])/baseMVA;

        # shunt conductance and susceptance
        gs[cID] = parse(Float64,bdata[5])/baseMVA;
        bs[cID] = parse(Float64,bdata[6])/baseMVA;

        # nominal voltage magnitude and angle
        Vmag[cID] = parse(Float64,bdata[8]);
        Vang[cID] = parse(Float64,bdata[9]);

        # obtain the maximum and the minimum of the voltage magnitude
        Vmax[cID] = parse(Float64,bdata[12]);
        Vmin[cID] = parse(Float64,bdata[13]);
    end

    return IDList,Vmax,Vmin,gs,bs,Vmag,Vang,Pd,Qd,bType;
end

function parsegenST(genST,baseMVA)
  # parse the genST string collected from readMP
  # input: the strings containing generator information
  # output: data variable

  # separate the raw strings line by line
    genSTlbl = collect(eachmatch(r"([0-9\t\ \.\-e]+[\;\r]*\n)",genST));

    # parse each line of the generator data
    Pmax = Dict();
    Pmin = Dict();
    Qmax = Dict();
    Qmin = Dict();
    Pg = Dict();
    Qg = Dict();
    RU = Dict();
    RD = Dict();
    L = Dict();
    LR = Dict();
    l = 0;
    genIDList = [];
    for gstr in genSTlbl
        gsstr = gstr.captures[1];
        gsstr = strip(strip(gsstr),';');
        gdata = split(gsstr);
        l += 1;

        cID = parse(Int64,gdata[1]);
        push!(genIDList,l)

        # obtain the maximum/minimum active/reactive power output
        # L is the bus in which the generator is located
        L[l] = cID;
        if !(cID in keys(LR))
            LR[cID] = [];
            push!(LR[cID],l);
        else
            push!(LR[cID],l);
        end
        Pmax[l] = parse(Float64,gdata[9])/baseMVA;
        Pmin[l] = parse(Float64,gdata[10])/baseMVA;
        Qmax[l] = parse(Float64,gdata[4])/baseMVA;
        Qmin[l] = parse(Float64,gdata[5])/baseMVA;

        Pg[l] = parse(Float64,gdata[2])/baseMVA;
        Qg[l] = parse(Float64,gdata[3])/baseMVA;

        RU[l] = parse(Float64,gdata[17]);
        RD[l] = -RU[l];
    end
    return genIDList,L,LR,Pmax,Pmin,Qmax,Qmin,Pg,Qg,RU,RD;
end

function parsebrST(brST,baseMVA)
    # parse the brST string collected from readMP
    # input: the strings containing branch information
    # output: data variable

    # separate the raw strings line by line
    brSTlbl = collect(eachmatch(r"([0-9\t\ \.\-e]+[\;\r]*\n)",brST));

    g = Dict();
    b = Dict();
    bc = Dict();
    angmax = Dict();
    angmin = Dict();
    τ1 = Dict();
    τ2 = Dict();
    σ = Dict();
    m = Dict();
    brList = [];
    brList1 = [];
    rateA = Dict();

    for brstr in brSTlbl
        brsstr = brstr.captures[1];
        brsstr = strip(strip(brsstr),';');
        brdata = split(brsstr);

        fID = parse(Int64,brdata[1]);
        tID = parse(Int64,brdata[2]);
        if !((fID,tID) in brList)
            m[(fID,tID)] = 1;
            m[(tID,fID)] = 1;
            push!(brList,(fID,tID));
            push!(brList,(tID,fID));

            r = parse(Float64,brdata[3]);
            x = parse(Float64,brdata[4]);
            bc1 = parse(Float64,brdata[5]);
            τ = parse(Float64,brdata[9]);
            if τ == 0.0
                τ = 1.0;
            end
            σ1 = parse(Float64,brdata[10]);
            g[(fID,tID,1)] = round(r/(r^2 + x^2),digits=6);
            g[(tID,fID,1)] = round(r/(r^2 + x^2),digits=6);
            b[(fID,tID,1)] = round(-x/(r^2 + x^2),digits=6);
            b[(tID,fID,1)] = round(-x/(r^2 + x^2),digits=6);
            bc[(fID,tID,1)] = bc1;
            bc[(tID,fID,1)] = bc1;
            τ1[(fID,tID,1)] = τ;
            τ1[(tID,fID,1)] = 1.0;
            τ2[(fID,tID,1)] = 1.0;
            τ2[(tID,fID,1)] = τ;
            σ[(fID,tID,1)] = σ1/180*pi;
            σ[(tID,fID,1)] = -σ1/180*pi;

          if parse(Float64,brdata[6]) == 0
              rateA[(fID,tID,1)] = Inf;
              rateA[(tID,fID,1)] = Inf;
          else
              rateA[(fID,tID,1)] = parse(Float64,brdata[6])/baseMVA;
              rateA[(tID,fID,1)] = parse(Float64,brdata[6])/baseMVA;
          end

          angmin[(fID,tID,1)] = parse(Float64,brdata[12]);
          angmax[(fID,tID,1)] = parse(Float64,brdata[13]);
          angmin[(tID,fID,1)] = parse(Float64,brdata[12]);
          angmax[(tID,fID,1)] = parse(Float64,brdata[13]);
          push!(brList1,(fID,tID,1));
          push!(brList1,(tID,fID,1));
        else
            m[(fID,tID)] += 1;
            m[(tID,fID)] += 1;
            r = parse(Float64,brdata[3]);
            x = parse(Float64,brdata[4]);
            bc1 = parse(Float64,brdata[5]);
            τ = parse(Float64,brdata[9]);
            if τ == 0.0
                τ = 1.0;
            end
            σ1 = parse(Float64,brdata[10]);
            g[(fID,tID,m[(fID,tID)])] = round(r/(r^2 + x^2),digits=6);
            g[(tID,fID,m[(tID,fID)])] = round(r/(r^2 + x^2),digits=6);
            b[(fID,tID,m[(fID,tID)])] = round(-x/(r^2 + x^2),digits=6);
            b[(tID,fID,m[(tID,fID)])] = round(-x/(r^2 + x^2),digits=6);
            bc[(fID,tID,m[(fID,tID)])] = bc1;
            bc[(tID,fID,m[(tID,fID)])] = bc1;
            τ1[(fID,tID,m[(fID,tID)])] = τ;
            τ1[(tID,fID,m[(tID,fID)])] = 1.0;
            τ2[(fID,tID,m[(tID,fID)])] = 1.0;
            τ2[(tID,fID,m[(tID,fID)])] = τ;
            σ[(fID,tID,m[(fID,tID)])] = σ1/180*pi;
            σ[(tID,fID,m[(tID,fID)])] = -σ1/180*pi;

            if parse(Float64,brdata[6]) == 0
                rateA[(fID,tID,m[(fID,tID)])] = Inf;
                rateA[(tID,fID,m[(tID,fID)])] = Inf;
            else
                rateA[(fID,tID,m[(fID,tID)])] = parse(Float64,brdata[6])/baseMVA;
                rateA[(tID,fID,m[(tID,fID)])] = parse(Float64,brdata[6])/baseMVA;
            end

            angmin[(fID,tID,m[(fID,tID)])] = parse(Float64,brdata[12]);
            angmax[(fID,tID,m[(fID,tID)])] = parse(Float64,brdata[13]);

            angmin[(tID,fID,m[(tID,fID)])] = parse(Float64,brdata[12]);
            angmax[(tID,fID,m[(tID,fID)])] = parse(Float64,brdata[13]);

            push!(brList1,(fID,tID,m[(fID,tID)]));
            push!(brList1,(tID,fID,m[(tID,fID)]));
        end
    end

    # build the reverse dictionary
    brRev = Dict();
    for k in brList1
        brRev[k] = (k[2],k[1],k[3]);
    end

    return brList1,brRev,g,b,bc,angmax,angmin,rateA,τ1,τ2,σ;
end

function parsecST(cST,genIDList,baseMVA)
    # parse the cST string collected from readMP
    # input: the strings containing cost information
    # output: data variable

    # separate the raw strings line by line
    cSTlbl = collect(eachmatch(r"([0-9\t\ \.\-e]+[\;\r]*\n)",cST));
    cp = Dict();
    cq = Dict();

    if length(genIDList) == length(cSTlbl)
        # only active power
        counter = 0;

        for cstr in cSTlbl
            csstr = cstr.captures[1];
            csstr = strip(strip(csstr),';');
            cdata = split(csstr);

            counter += 1;
            typeC = parse(Int64,cdata[1]);
            startC = parse(Float64,cdata[2]);
            endC = parse(Float64,cdata[3]);
            nC = parse(Float64,cdata[4]);
            paramC = [];
            for i in 1:Int64(nC)
                push!(paramC,parse(Float64,cdata[i+4])*(baseMVA^(Int64(nC) - i)));
            end

            citem = costDataType(typeC,startC,endC,nC,paramC);
            cp[genIDList[counter]] = citem;
        end
    elseif length(genIDList)*2 == length(cSTlbl)
        # both active power and reactive power
        counter = 0;
        for cstr in cSTlbl
            csstr = cstr.captures[1];
            csstr = strip(strip(csstr),';');
            cdata = split(csstr,r"\t");

            counter += 1;
            typeC = parse(Int64,cdata[1]);
            startC = parse(Float64,cdata[2]);
            endC = parse(Float64,cdata[3]);
            nC = parse(Float64,cdata[4]);
            paramC = [];
            for i in 1:Int64(nC)
                push!(paramC,parse(Float64,cdata[i+4])*(baseMVA^(Int64(nC) - i)));
            end

            citem = costDataType(typeC,startC,endC,nC,paramC);
            if counter <= length(genIDList)
                cp[genIDList[counter]] = citem;
            else
                cq[genIDList[counter - length(genIDList)]] = citem;
            end
        end
    else
        error("The number of cost entries does not match the number of generators.")
    end

    return cp,cq;
end

function constFixed(baseMVA,bType,IDList,genIDList,brList,brRev,Vmax,Vmin,L,LR,Pmax,Pmin,Qmax,Qmin,gs,bs,Vmag,Vang,Pd,Qd,Pg,Qg,RU,RD,g,b,bc,angmax,angmin,rateA,τ1,τ2,σ,cp,cq,cz)
    # combine all the data to a struct
    connectPair = [];
    connectDict = Dict();
    branchDict1 = Dict();
    branchDict2 = Dict();
    for i in IDList
        connectDict[i] = [];
        branchDict1[i] = [];
        branchDict2[i] = [];
    end
    for k in brList
        push!(branchDict1[k[1]],k);
        push!(branchDict2[k[2]],k);
        if !((k[1],k[2]) in connectPair)
            push!(connectPair,(k[1],k[2]));
            push!(connectDict[k[1]],k[2]);
        end
    end
    kpDict = Dict();
    for k in brList
        if (k[1],k[2]) in keys(kpDict)
            push!(kpDict[(k[1],k[2])],k);
        else
            kpDict[(k[1],k[2])] = [k];
        end
    end

    # collapse all the lines with high conductance
    clusterList = [];
    for k in brList
        if abs(b[k]) >= 1e4
            inCluster = false;
            for item in clusterList
                if (k[1] in item) && !(k[2] in item)
                    push!(item, k[2]);
                    inCluster = true;
                elseif (k[2] in item) && !(k[1] in item)
                    push!(item, k[1]);
                    inCluster = true;
                elseif (k[2] in item) && (k[1] in item)
                    inCluster = true;
                end
            end
            if !inCluster
                push!(clusterList,[k[1],k[2]]);
            end
        end
    end

    # for each cluster, combine the nodes to the lowest indexed node
    for item in clusterList
        # combine the nodes to the lowest indexed node
        combinedN = minimum(item);
        item1 = copy(item);
        filter!(e -> e != combinedN, item1);
        # remove the node
        filter!(e -> !(e in item1), IDList);
        for i in item1
            # remove the pseudo line
            filter!(e -> !(((e[1] == i)&&(e[2] == combinedN))||((e[1] == combinedN)&&(e[2] == i))), brList);
            for k in keys(brRev)
                if !(k in brList)
                    delete!(brRev,k);
                    delete!(kpDict,(k[1],k[2]));
                    filter!(e -> e!=k, branchDict1[i]);
                    filter!(e -> e!=k, branchDict2[i]);
                    delete!(rateA,k);
                    delete!(g,k);
                    delete!(b,k);
                    delete!(bc,k);
                    delete!(τ1,k);
                    delete!(τ2,k);
                    delete!(σ,k);
                end
            end
            # consolidate the generator information
            if i in keys(LR)
                rmGenList = LR[i];
                if combinedN in keys(LR)
                    append!(LR[combinedN], rmGenList);
                    for gitem in rmGenList
                        L[gitem] = combinedN;
                    end
                else
                    LR[combinedN] = rmGenList;
                    for gitem in rmGenList
                        L[gitem] = combinedN;
                    end
                end
                delete!(LR,i);
            end

            # remove the node information
            Vmax[combinedN] = min(Vmax[combinedN],Vmax[i]);
            Vmin[combinedN] = max(Vmin[i],Vmin[combinedN]);
            gs[combinedN] += gs[i];
            bs[combinedN] += bs[i];
            Pd[combinedN] += Pd[i];
            Qd[combinedN] += Qd[i];
            delete!(Vmax,i);
            delete!(Vmin,i);
            delete!(gs,i);
            delete!(bs,i);
            delete!(Pd,i);
            delete!(Qd,i);
            delete!(Vmag,i);
            delete!(Vang,i);
            delete!(angmax,i);
            delete!(angmin,i);

            # consolidate the arcs
            branchDict1Temp = [];
            for k in branchDict1[i]
                filter!(e -> e != k, brList);
                if (combinedN,k[2]) in keys(kpDict)
                    newbr = (combinedN,k[2],length(kpDict[(combinedN,k[2])]) + 1);
                    push!(brList,newbr);
                    brRev[newbr] = (newbr[2],newbr[1],newbr[3]);
                    kpDict[(newbr[1],newbr[2])] = newbr;
                else
                    newbr = (combinedN,k[2],1);
                    push!(brList,newbr);
                    brRev[newbr] = (newbr[2],newbr[1],newbr[3]);
                    kpDict[(newbr[1],newbr[2])] = newbr;
                end
                filter!(e -> e != k, branchDict2[k[2]]);
                push!(branchDict2[k[2]],newbr);
                rateA[newbr] = rateA[k];
                g[newbr] = g[k];
                b[newbr] = b[k];
                bc[newbr] = bc[k];
                τ1[newbr] = τ1[k];
                τ2[newbr] = τ2[k];
                σ[newbr] = σ[k];
                push!(branchDict1Temp,newbr);
                delete!(brRev,k);
                delete!(kpDict,(k[1],k[2]));
                delete!(rateA,k);
                delete!(g,k);
                delete!(b,k);
                delete!(bc,k);
                delete!(τ1,k);
                delete!(τ2,k);
                delete!(σ,k);
            end
            branchDict1[i] = branchDict1Temp;
            branchDict2Temp = [];
            for k in branchDict2[i]
                filter!(e -> e != k, brList);
                if (k[1],combinedN) in keys(kpDict)
                    newbr = (k[1],combinedN,length(kpDict[(k[1],combinedN)]) + 1);
                    push!(brList,newbr);
                    brRev[newbr] = (newbr[2],newbr[1],newbr[3]);
                    kpDict[(newbr[1],newbr[2])] = newbr;
                else
                    newbr = (k[1],combinedN,1);
                    push!(brList,newbr);
                    brRev[newbr] = (newbr[2],newbr[1],newbr[3]);
                    kpDict[(newbr[1],newbr[2])] = newbr;
                end
                filter!(e -> e != k, branchDict1[k[1]]);
                push!(branchDict1[k[1]],newbr);
                rateA[newbr] = rateA[k];
                g[newbr] = g[k];
                b[newbr] = b[k];
                bc[newbr] = bc[k];
                τ1[newbr] = τ1[k];
                τ2[newbr] = τ2[k];
                σ[newbr] = σ[k];
                push!(branchDict2Temp,newbr);
                delete!(brRev,k);
                delete!(kpDict,(k[1],k[2]));
                delete!(rateA,k);
                delete!(g,k);
                delete!(b,k);
                delete!(bc,k);
                delete!(τ1,k);
                delete!(τ2,k);
                delete!(σ,k);
            end
            branchDict2[i] = branchDict2Temp;
        end
    end

    ii = 0;
    busInd = Dict();
    for i in IDList
        ii += 1;
        busInd[i] = ii;
    end

    # obtain the pairs that are connected
    connectPair = [];
    connectDict = Dict();
    branchDict1 = Dict();
    branchDict2 = Dict();
    for i in IDList
        connectDict[i] = [];
        branchDict1[i] = [];
        branchDict2[i] = [];
    end
    for k in brList
        push!(branchDict1[k[1]],k);
        push!(branchDict2[k[2]],k);
        if !((k[1],k[2]) in connectPair)
            push!(connectPair,(k[1],k[2]));
            push!(connectDict[k[1]],k[2]);
        end
    end
    kpDict = Dict();
    for k in brList
        if (k[1],k[2]) in keys(kpDict)
            push!(kpDict[(k[1],k[2])],k);
        else
            kpDict[(k[1],k[2])] = [k];
        end
    end
    RUadj = Dict();
    RDadj = Dict();
    for iKey in keys(RU)
        RUadj[iKey] = RU[iKey]/baseMVA;
    end
    for iKey in keys(RD)
        RDadj[iKey] = RD[iKey]/baseMVA;
    end
    fData = fixedData(baseMVA,bType,IDList,genIDList,brList,brRev,
                        L,LR,Vmax,Vmin,Pmax,Pmin,Qmax,Qmin,gs,bs,Vmag,Vang,
                        Pd,Qd,Pg,Qg,RUadj,RDadj,g,b,bc,angmax,angmin,rateA,τ1,τ2,σ,cp,cq,cz,
                        busInd,branchDict1,branchDict2,connectPair,connectDict,kpDict);
    return fData
end

function readStatic(fileName::String, cz = 1e7)
    # read in the static network information from a .m file (MatPower)
    busST,genST,brST,cST,uST,baseMVA = readMP(fileName);
    IDList,Vmax,Vmin,gs,bs,Vmag,Vang,Pd,Qd,bType = parsebusST(busST,baseMVA);
    genIDList,Loc,LocRev,Pmax,Pmin,Qmax,Qmin,Pg,Qg,RU,RD = parsegenST(genST,baseMVA);
    brList,brRev,g,b,bc,angmax,angmin,rateA,τ1,τ2,σ = parsebrST(brST,baseMVA);
    cp,cq = parsecST(cST,genIDList,baseMVA);
    fData = constFixed(baseMVA,bType,IDList,genIDList,brList,brRev,Vmax,Vmin,Loc,LocRev,Pmax,Pmin,Qmax,Qmin,
            gs,bs,Vmag,Vang,Pd,Qd,Pg,Qg,RU,RD,g,b,bc,angmax,angmin,rateA,τ1,τ2,σ,cp,cq,cz);

    # specify the line power flow constraint if there is none
    θu = pi/3;
    for k in fData.brList
        if fData.rateA[k] == Inf
            fData.rateA[k] = sqrt(fData.g[k]^2+fData.b[k]^2)*max(fData.Vmax[k[1]],fData.Vmax[k[2]])*sqrt(fData.Vmax[k[1]]^2 + fData.Vmax[k[2]]^2 - 2*fData.Vmax[k[1]]*fData.Vmin[k[2]]*cos(θu));
        end
    end
    return fData;
end

function readDisruption(tfileName,ofileName,fileType,τ = nothing)
    # read in the disruption data: disruption time distribution
    if fileType == "csv"
        # csv file format:
        # First row: time period t
        # Second row: pr{D = t}
        # Third row: Ω
        # Fourth row: pr{B_ω = 0}
        dataT = readdlm(tfileName, ',');
        tDistrn = Dict();
        mt,nt = size(dataT);
        for tInd in 1:nt
            tDistrn[Int64(dataT[1,tInd])] = dataT[2,tInd];
        end

        # detect the dimension of Ω
        dataOme = readdlm(ofileName, ',');
        ωDistrn = Dict();
        ωDict = Dict();
        ωτ = Dict();
        mo,no = size(dataOme);
        for oInd in 1:mo
            scenInd = dataOme[oInd,1];
            ωDistrn[scenInd] = dataOme[oInd,2];
            ωDict[scenInd] = [];
            if τ == nothing
                τr = dataOme[oInd,3];
            else
                τr = τ;
            end
            ωτ[scenInd] = τr;
            for noInd in 4:no
                lineStr = dataOme[oInd,noInd];
                if ',' in lineStr
                    # parse the line
                    startN,endN = match(r"\(([0-9]+),([0-9]+)\)",lineStr).captures;
                    push!(ωDict[scenInd],(parse(Int64,startN),parse(Int64,endN)));
                elseif !(lineStr == "")
                    # parse the node
                    disN = lineStr;
                    push!(ωDict[scenInd],disN);
                end
            end
        end
        if sum(values(ωDistrn)) == 1
            pDistr = probDistrn(tDistrn,ωDistrn,ωDict,ωτ);
            return pDistr;
        else
            ωList = [item for item in keys(ωDistrn)];
            ωDistrn[ωList[1]] = 1 - sum([ωDistrn[item] for item in keys(ωDistrn) if item != ωList[1]]);
            pDistr = probDistrn(tDistrn,ωDistrn,ωDict,ωτ);
            return pDistr;
        end
    else
        println("Currently your file type is not supported");
    end
end

function readDisruption_old(fileName,fileType)
    # read in the disruption data: disruption time distribution
    if fileType == "csv"
        # csv file format:
        # First row: time period t
        # Second row: pr{D = t}
        # Third row: Ω
        # Fourth row: pr{B_ω = 0}
        dataRaw = readdlm(fileName, ',');

        # detect the dimension of T
        dataT = dataRaw[1:2,:];
        emptyInd = findfirst(x -> x=="", dataT[1,:]);
        if emptyInd != nothing
            # if it is not full length
            dataT = dataT[:,1:(emptyInd - 1)];
        end
        tDistrn = Dict();
        mt,nt = size(dataT);
        for tInd in 1:nt
            tDistrn[Int64(dataT[1,tInd])] = dataT[2,tInd];
        end

        # detect the dimension of Ω
        dataOme = dataRaw[3:4,:];
        emptyInd = findfirst(x -> x=="", dataOme[1,:]);
        if emptyInd != nothing
            # if it is not full length
            dataOme = dataOme[:,1:(emptyInd - 1)];
        end
        ωDistrn = Dict();
        mo,no = size(dataOme);
        for oInd in 1:no
            lineStr = dataOme[1,oInd];
            if ',' in lineStr
                # parse the line
                startN,endN = match(r"\(([0-9]+),([0-9]+)\)",lineStr).captures;
                ωDistrn[parse(Int64,startN),parse(Int64,endN)] = dataOme[2,oInd];
            else
                # parse the node
                disN = lineStr;
                ωDistrn[disN] = dataOme[2,oInd];
            end
        end
        if sum(values(ωDistrn)) == 1
            pDistr = probDistrn(tDistrn,ωDistrn);
            return pDistr;
        else
            ωList = [item for item in keys(ωDistrn)];
            ωDistrn[ωList[1]] = 1 - sum([ωDistrn[item] for item in keys(ωDistrn) if item != ωList[1]]);
            pDistr = probDistrn(tDistrn,ωDistrn);
            return pDistr;
        end
    else
        println("Currently your file type is not supported");
    end
end

function readBattery(fileName,fileType,baseMVA = 100)
    # read in the battery information: charging/discharging factor, capacity
    if fileType == "csv"
        # csv file format:
        # First column: node ID
        # Second column: ηd
        # Third column: ηc
        # Fourth column: capacity
        # Fifth column: cost (optional)
        dataRaw = readdlm(fileName, ',');
        mb,nb = size(dataRaw);
        capacity = Dict();
        cost = Dict();
        ηα = Dict();
        ηβ = Dict();
        IDList = [];
        LocDict = Dict();
        bInv = Dict();
        uCap = Dict();

        for i in 1:mb
            ID = Int64(dataRaw[i,1]);
            push!(IDList,ID);
            loc = Int64(dataRaw[i,2]);
            LocDict[ID] = loc;
            capacity[ID] = dataRaw[i,3]/baseMVA;
            cost[ID] = dataRaw[i,4]*baseMVA;
            bInv[ID] = dataRaw[i,5]/baseMVA;
            uCap[ID] = dataRaw[i,6]/baseMVA;
            ηparams = [j for j in dataRaw[i,7:nb] if j != ""];
            ηα[ID] = [];
            ηβ[ID] = [];
            for j in 1:2:length(ηparams)
                push!(ηα[ID],ηparams[j]);
                # ηα must be monotonously decreasing
                push!(ηβ[ID],ηparams[j + 1]);
            end
        end
        bData = batteryData(IDList,LocDict,bInv,ηα,ηβ,capacity,cost,uCap);
    else
        println("Currently your file type is not supported");
    end
    return bData;
end

function readDemand(fileNameP, fileNameQ, fileType)
    # read in the battery information: charging/discharging factor, capacity
    if fileType == "csv"
        # csv file format:
        # First column: node ID
        # Second column: ηd
        # Third column: ηc
        dataRawP = readdlm(fileNameP, ',');
        mp,np = size(dataRawP);
        pd = Dict();
        for i in 1:mp
            pd[Int64(dataRawP[i,1])] = dataRawP[i,2:np];
        end

        dataRawQ = readdlm(fileNameQ, ',');
        mq,nq = size(dataRawQ);
        qd = Dict();
        for i in 1:mq
            qd[Int64(dataRawQ[i,1])] = dataRawQ[i,2:nq];
        end
        T = nq - 1;
        dData = demandData(T,pd,qd);
        return dData;
    else
        println("Currently your file type is not supported");
    end
end

function readInData(i,caseList,T,τ,cz = 1e4,λD = 0)
    fileAdd = "case$(caseList[i])_ieee.m";
    global fData = readStatic(fileAdd,cz);
    disAddt = "testProbReadt_$(caseList[i]).csv";
    disAddo = "testProbReado_$(caseList[i]).csv";
    global pDistr = readDisruption(disAddt,disAddo,"csv",τ);
    pAdd = "testDataP_$(caseList[i]).csv";
    qAdd = "testDataQ_$(caseList[i]).csv";
    global dData = readDemand(pAdd,qAdd,"csv");
    bAdd = "testDataB_$(caseList[i]).csv";
    global bData = readBattery(bAdd,"csv",fData.baseMVA);

    if λD == 0
        global pDistr = modifyT(pDistr,4/T,T);
    else
        global pDistr = modifyT(pDistr,λD,T);
    end
end

function readInData_tau(i,caseList,T,cz = 1e4,λD = 0,suffix = "",τ = nothing)
    fileAdd = "case$(caseList[i])_ieee.m";
    global fData = readStatic(fileAdd,cz);
    disAddt = "testProbReadt_$(caseList[i]).csv";
    disAddo = "testProbReado_$(caseList[i])_tau$(suffix).csv";
    global pDistr = readDisruption(disAddt,disAddo,"csv",τ);
    pAdd = "testDataP_$(caseList[i]).csv";
    qAdd = "testDataQ_$(caseList[i]).csv";
    global dData = readDemand(pAdd,qAdd,"csv");
    bAdd = "testDataB_$(caseList[i]).csv";
    global bData = readBattery(bAdd,"csv",fData.baseMVA);

    if λD == 0
        global pDistr = modifyT(pDistr,4/T,T);
    else
        global pDistr = modifyT(pDistr,λD,T);
    end
end

function readInData_old(T,ωSet,cz = 500,λD = 1/2)
    fileAdd = "case13_ieee_old.m";
    global fData = readStatic(fileAdd,cz);
    disAdd = "testProbRead_96.csv"
    global pDistr = readDisruption(disAdd,"csv");
    pAdd = "testDataP_96_Fixed.csv";
    qAdd = "testDataQ_96_Fixed.csv";
    global dData = readDemand(pAdd,qAdd,"csv");
    bAdd = "testDataB.csv";
    global bData = readBattery(bAdd,"csv",fData.baseMVA);

    for ω in keys(pDistr.ωDistrn)
        if !(ω in ωSet)
            global pDistr = modifyOmega(pDistr,ω);
        end
    end

    if λD == 0
        global pDistr = modifyT(pDistr,4/T,T);
    else
        global pDistr = modifyT(pDistr,λD,T);
    end
end


text:
@testset "model.jl" begin
    @testset "setval! & generated_quantities" begin
        @model function demo1(xs, ::Type{TV} = Vector{Float64}) where {TV}
            m = TV(undef, 2)
            for i in 1:2
                m[i] ~ Normal(0, 1)
            end

            for i in eachindex(xs)
                xs[i] ~ Normal(m[1], 1.)
            end

            return (m, )
        end

        @model function demo2(xs)
            m ~ MvNormal(2, 1.)

            for i in eachindex(xs)
                xs[i] ~ Normal(m[1], 1.)
            end

            return (m, )
        end

        xs = randn(3)
        model1 = demo1(xs)
        model2 = demo2(xs)

        chain1 = sample(model1, MH(), 100)
        chain2 = sample(model2, MH(), 100)

        res11 = generated_quantities(model1, chain1)
        res21 = generated_quantities(model2, chain1)

        res12 = generated_quantities(model1, chain2)
        res22 = generated_quantities(model2, chain2)

        # Check that the two different models produce the same values for
        # the same chains.
        @test all(res11 .== res21)
        @test all(res12 .== res22)
        # Ensure that they're not all the same (some can be, because rejected samples)
        @test any(res12[1:end - 1] .!= res12[2:end])

        test_setval!(model1, chain1)
        test_setval!(model2, chain2)

        # Next level
        @model function demo3(xs, ::Type{TV} = Vector{Float64}) where {TV}
            m = Vector{TV}(undef, 2)
            for i = 1:length(m)
                m[i] ~ MvNormal(2, 1.)
            end

            for i in eachindex(xs)
                xs[i] ~ Normal(m[1][1], 1.)
            end

            return (m, )
        end

        @model function demo4(xs, ::Type{TV} = Vector{Vector{Float64}}) where {TV}
            m = TV(undef, 2)
            for i = 1:length(m)
                m[i] ~ MvNormal(2, 1.)
            end

            for i in eachindex(xs)
                xs[i] ~ Normal(m[1][1], 1.)
            end

            return (m, )
        end

        model3 = demo3(xs)
        model4 = demo4(xs)

        chain3 = sample(model3, MH(), 100)
        chain4 = sample(model4, MH(), 100)

        res33 = generated_quantities(model3, chain3)
        res43 = generated_quantities(model4, chain3)

        res34 = generated_quantities(model3, chain4)
        res44 = generated_quantities(model4, chain4)

        # Check that the two different models produce the same values for
        # the same chains.
        @test all(res33 .== res43)
        @test all(res34 .== res44)
        # Ensure that they're not all the same (some can be, because rejected samples)
        @test any(res34[1:end - 1] .!= res34[2:end])
    end
end


text:
using Oceananigans.Architectures: device
using Oceananigans.Operators: ΔzC, Δzᵃᵃᶠ, Δzᵃᵃᶜ

"""
Compute the vertical integrated volume flux from the bottom to z=0 (i.e. linear free-surface)

    `U^{*} = ∫ [(u^{*})] dz`
    `V^{*} = ∫ [(v^{*})] dz`
"""
### Note - what we really want is RHS = divergence of the vertically integrated volume flux
###        we can optimize this a bit later to do this all in one go to save using intermediate variables.
function compute_vertically_integrated_volume_flux!(free_surface, model)

    event = launch!(model.architecture,
                    model.grid,
                    :xyz,
                    _compute_vertically_integrated_volume_flux!,
                    model.velocities,
                    model.grid,
                    free_surface.barotropic_volume_flux,
                    dependencies=Event(device(model.architecture)))

    return event
end
@kernel function _compute_vertically_integrated_volume_flux!(U, grid, barotropic_volume_flux )
    i, j, k = @index(Global, NTuple)
    # U.w[i, j, 1] = 0 is enforced via halo regions.
    barotropic_volume_flux.u[i, j, 1] = 0.
    barotropic_volume_flux.v[i, j, 1] = 0.
    @unroll for k in 1:grid.Nz
        #### Not sure this will always be Δzᵃᵃᶜ. When we have step bathymetry then it may be
        #### Δzᶠᶜᶜ and Δzᶜᶠᶜ. Not entirely sure the notation works perfectly. For the X direction
        #### volume flux we locate the z direction length at the intersection of a YZ plane that is
        #### on an x face and a XZ plane that passes through a y center. For the Y direction
        #### volume flux we locate the z direction lenght at the intersection of the XZ plane
        #### that is on a y face and the YZ plane that passes through the x center.
        @inbounds barotropic_volume_flux.u[i, j, 1] += U.u[i, j, k]*Δyᶠᶜᵃ(i, j, k, grid)*Δzᵃᵃᶜ(i, j, k, grid)
        @inbounds barotropic_volume_flux.v[i, j, 1] += U.v[i, j, k]*Δxᶜᶠᵃ(i, j, k, grid)*Δzᵃᵃᶜ(i, j, k, grid)
     end
end

"""
Compute volume integrated divergence (need to change name)

 `D = sum.... `
"""
function compute_volume_scaled_divergence!(free_surface, model)
    event = launch!(model.architecture,
                    model.grid,
                    :xy,
                    _compute_volume_scaled_divergence!,
                    model.grid,
                    free_surface.barotropic_volume_flux.u,
                    free_surface.barotropic_volume_flux.v,
                    free_surface.implicit_step_solver.solver.settings.RHS,
                    dependencies=Event(device(model.architecture)))
    return event
end
@kernel function _compute_volume_scaled_divergence!(grid, ut, vt, div)
    # Here we use a integral form that has been multiplied through by volumes to be
    # consistent with the symmetric "A" matrix.
    # The quantities differenced here are transports i.e. normal velocity vectors
    # integrated over an area.
    #
    i, j = @index(Global, NTuple)
    @inbounds div[i, j, 1] = δxᶜᵃᵃ(i, j, 1, grid, ut) + δyᵃᶜᵃ(i, j, 1, grid, vt)
end

function add_previous_free_surface_contribution(free_surface, model, Δt )
   g = model.free_surface.gravitational_acceleration
   event = launch!(model.architecture,
                   model.grid,
                   :xy,
                   _add_previous_free_surface_contribution!,
                   model.grid,
                   free_surface.η.data,
                   free_surface.implicit_step_solver.solver.settings.RHS,
                   g,
                   Δt,
                   dependencies=Event(device(model.architecture)))
    return event
end
@kernel function _add_previous_free_surface_contribution!(grid,η,RHS,g,Δt)
    i, j = @index(Global, NTuple)
    @inbounds RHS[i,j,1] -= Azᶜᶜᵃ(i, j, 1, grid)*η[i,j, 1]/(g*Δt^2)
end

function ∇²_baro_operator( HAx, HAy )
    ## Some of this should probably end up in some operator and grid generic auxilliaries
    Ax_baro = HAx
    Ay_baro = HAy
    @inline Ax_∂xᶠᶜᵃ_baro(i, j, k, grid, c) = Ax_baro[i, j, 1] * ∂xᶠᶜᵃ(i, j, 1, grid, c)
    @inline Ay_∂yᶜᶠᵃ_baro(i, j, k, grid, c) = Ay_baro[i, j, 1] * ∂yᶜᶠᵃ(i, j, 1, grid, c)
    @inline function ∇²_baro(i, j, k, grid, c)
       return  δxᶜᵃᵃ(i, j, 1, grid, Ax_∂xᶠᶜᵃ_baro, c) +
               δyᵃᶜᵃ(i, j, 1, grid, Ay_∂yᶜᶠᵃ_baro, c)
    end
    return  ∇²_baro
end

@kernel function implicit_η!(∇²_baro, Δt, g, grid, f, implicit_η_f)
        i, j = @index(Global, NTuple)
        ### Not sure what to call this function
        ### it is for left hand side operator in
        ### ( ∇ʰ⋅H∇ʰ - 1/gΔt² ) ηⁿ⁺¹ = 1/(gΔt) ∇ʰH U̅ˢᵗᵃʳ - 1/(gΔt²) ηⁿ
        ### written in a discrete finite volume form in which the equation
        ### is arranged to ensure a symmtric form
        ### e.g.
        ###
        ### δⁱÂʷ∂ˣηⁿ⁺¹ + δʲÂˢ∂ʸηⁿ⁺¹ - 1/gΔt² Aᶻηⁿ⁺¹ =
        ###  1/(gΔt)(δⁱÂʷu̅ˢᵗᵃʳ + δʲÂˢv̅ˢᵗᵃʳ) - 1/gΔt² Aᶻηⁿ
        ###
        ### where  ̂ indicates a vertical integral, and
        ###        ̅ indicates a vertical average
        ###
        i, j = @index(Global, NTuple)
        @inbounds implicit_η_f[i, j, 1] =  ∇²_baro(i, j, 1, grid, f) - Azᶜᶜᵃ(i, j, 1, grid)*f[i,j, 1]/(g*Δt^2)
end


text:
# Useful math constants for distributions

import Base.@math_const

@math_const halfπ  1.5707963267948966192 big(π) * 0.5
@math_const twoπ   6.2831853071795864769 big(2.) * π
@math_const sqrt2  1.4142135623730950488 sqrt(big(2.))
@math_const sqrt3  1.7320508075688772935 sqrt(big(3.))
@math_const invπ   0.3183098861837906715 big(1.) / π
@math_const logπ   1.1447298858494001741 log(big(π))
@math_const log2π  1.8378770664093454836 log(big(2.)*π)
@math_const log4π  2.5310242469692907930 log(big(4.)*π)
@math_const sqrt2π 2.5066282746310005024 sqrt(big(2.)*π)
@math_const logtwo 0.6931471805599453094 log(big(2.))
@math_const loghalf -0.6931471805599453094 log(big(0.5))

@math_const sqrthalfπ 1.2533141373155002512 sqrt(big(0.5)*π)


text:
# See https://web.stanford.edu/~lcambier/fast/tuto.php
@testset "Hydro Thermal Scheduling with $solver" for solver in lp_solvers
    num_stages = 5
    numScen = 2
    C = 5
    V = 8
    d = 6
    r = [2, 10]
    x = Matrix{JuMP.Variable}(undef, num_stages, numScen)
    y = Matrix{JuMP.Variable}(undef, num_stages, numScen)
    p = Matrix{JuMP.Variable}(undef, num_stages, numScen)

    models = Vector{JuMP.Model}(undef, num_stages)
	for s in 1:num_stages
        for ξ in 1:(s == 1 ? 1 : numScen)
			if s == 1
				model = StructuredModel(num_scenarios=numScen)
			else
				model = StructuredModel(parent=models[s-1], prob=1/2, same_children_as=(ξ == 1 ? nothing : models[s]), id=ξ, num_scenarios=(s == num_stages ? 0 : numScen))
			end
			x[s, ξ] = @variable(model, lowerbound=0, upperbound=V)
			y[s, ξ] = @variable(model, lowerbound=0)
			p[s, ξ] = @variable(model, lowerbound=0)
			if s > 1
				@constraint(model, x[s, ξ] <= x[s-1, 1] + r[ξ] - y[s, ξ])
			else
                @constraint(model, x[s, ξ] <= Compat.Statistics.mean(r) - y[s, ξ])
			end
			@constraint(model, p[s, ξ] + y[s, ξ] >= d)
			@objective(model, Min, C * p[s, ξ])
			if ξ == 1
				models[s] = model
			end
		end
	end

    forwardcuts = true
    cutmode = StructProg.MultiCutGenerator()
    detectlb = false
    @testset "$(forwardcuts ? "With" : "Without") forward cuts" for forwardcuts in [false, true]
        @testset "With cut mode $cutmode" for cutmode in [StructProg.MultiCutGenerator(), StructProg.AvgCutGenerator()]
            @testset "$(detectlb ? "D" : "Not d")etecting value function lower bound" for detectlb in [false, true]
                sp = SOI.stochasticprogram(models[1], num_stages, solver, AvgCutPruningAlgo(-1), cutmode, detectlb)
                algo = SDDP.Algorithm(K = 16, forwardcuts = forwardcuts, backwardcuts = !forwardcuts)
                sol = SOI.optimize!(sp, algo, SOI.Pereira(2, 0.5) | SOI.IterLimit(10), 0)

                @test SOI.last_result(sol).status == :Optimal
                lb = SOI.last_result(sol).lowerbound
                if forwardcuts
                    @test SOI.niterations(sol) == (detectlb ? 3 : 5)
                    @test lb ≈ (detectlb ? 15 : 18.75)
                else
                    @test SOI.niterations(sol) == (!detectlb && isa(cutmode, StructProg.AvgCutGenerator) ? 5 : 2)
                    if isa(cutmode, StructProg.AvgCutGenerator)
                        if detectlb
                            @test lb ≈ 22.5 # Clp returns an approximate value
                        else
                            @test lb ≈ 23.75
                        end
                    else
                        if detectlb
                            @test lb ≈ 23.75
                        else
                            @test lb ≈ 18.75
                        end
                    end
                end

                StructProg.clear(models[1])
            end
        end
    end
end


text:
function trend_predict(x, m, kernel, transform, range)
    tail = x[(length(x) - m + 1):end]
    in_space = transform(tail)
    slope = kernel(in_space)
    predicted_space = slope * 1:range + space[end]
    inverse(transform(predicted_space))
end

x0 = [1.7, 1.9, 1.8, 1.4, 1.3, 1.4]
trend_trials = [
    [x0, 5, :weibull, :linear, 1, [2.3]],
    [x0, 5, :weibull, :log, 2, [2.4, 2.9]],
    [x0, 5, :exponential, :linear, 3, [2.1, 2.15, 2.18]],
    [x0, 5, :exponential, :log, 2, [2.31, 2.4, 2.47]],
]
for trial in trend_trials
    args = trial[1:5]
    result = trend_predict(args...)
    @test result == trial[6]
end

max_stars_repo_name: adolgert/test-scientific-code

max_issues_repo_name: adolgert/test-scientific-code

max_forks_repo_name: adolgert/test-scientific-code


text:
#
# Parallel thread spliiter
#
splitter(first,n) = first:nthreads():n

"""

```
reduce(output::Number, output_threaded::Vector{<:Number})
```

Functions to reduce the output of common options (vectors of numbers
and vectors of vectors). This function can be overloaded by custom
reduction methods. It always must both receive the `output` variable
as a parameter, and return it at the end.

"""
reduce(output::Number, output_threaded::Vector{<:Number}) = sum(output_threaded)
function reduce(output::AbstractVector, output_threaded::AbstractVector{<:AbstractVector})
  for i in 1:nthreads()
    @. output += output_threaded[i]
  end
  return output
end

#
# Serial version for self-pairwise computations
#
function map_pairwise_serial!(
  f::F, output, box::Box, cl::CellList{N,T};
  show_progress::Bool=false
) where {F,N,T}
  show_progress && (p = Progress(cl.ncwp[1],dt=1))
  for icell in 1:cl.ncwp[1]
    output = inner_loop!(f,box,icell,cl,output)
    show_progress && next!(p)
  end
  return output
end

#
# Parallel version for self-pairwise computations
#
function map_pairwise_parallel!(
  f::F1, output, box::Box, cl::CellList{N,T};
  output_threaded=output_threaded,
  reduce::F2=reduce,
  show_progress::Bool=false
) where {F1,F2,N,T}
  show_progress && (p = Progress(cl.ncwp[1],dt=1))
  @threads for it in 1:nthreads()
    for icell in splitter(it,cl.ncwp[1])
      output_threaded[it] = inner_loop!(f,box,icell,cl,output_threaded[it])
      show_progress && next!(p)
    end
  end
  output = reduce(output,output_threaded)
  return output
end

"""

```
partition!(x::AbstractVector,by)
```

Function that reorders `x` vector by putting in the first positions the
elements with values satisfying `by(el)`. Returns the number of elements
that satisfy the condition.

"""
function partition!(x::AbstractVector,by)
  iswap = 1
  @inbounds for i in eachindex(x)
    if by(x[i])
      if iswap != i
        x[iswap], x[i] = x[i], x[iswap]
      end
      iswap += 1
    end
  end
  return iswap - 1
end

"""

```
project_particles(jc_cartesian,cl,Δc,icell,box)
```

Projects all particles of the cell of cartesian coordinates
`jc_cartesian` into the unnitary vector `Δc` with direction
connecting `icell.center` and the center of cell `jc_cartesian`.
Returns a view of the array of projected particles, with length
equal to the number of particles of cell `jc_cartesian`.

"""
function project_particles(jc_cartesian,cl,Δc,icell,box)
  @unpack nc = box
  @unpack fp = cl
  projected_particles = cl.projected_particles[threadid()]
  jc = cell_linear_index(nc,jc_cartesian)
  pⱼ = fp[jc]
  npcell = cl.npcell[jc]
  j = pⱼ.index
  for jp in 1:npcell
    j_orig = pⱼ.index_original
    xpⱼ = pⱼ.coordinates
    xproj = dot(xpⱼ-icell.center,Δc)
    xreal = pⱼ.real
    projected_particles[jp] = ProjectedParticle(j_orig,xproj,xpⱼ,xreal)
    pⱼ = cl.np[j]
    j = pⱼ.index
  end
  pp = @view(projected_particles[1:npcell])
  return pp
end

function inner_loop!(
  f,box::Box{TriclinicCell},icell,
  cl::CellList{N,T},
  output
) where {N,T}
  @unpack cutoff, cutoff_sq, nc = box
  cell = cl.cwp[icell]

  for neighbour_cell in neighbour_cells_all(box)
    jc_cartesian = cell.cartesian + neighbour_cell
    jc = cell_linear_index(nc,jc_cartesian)

    # Vector connecting cell centers
    if jc == cell.icell
      Δc = SVector{N,T}(ntuple(i->1,N))
    else
      Δc = cell_center(jc_cartesian,box) - cell.center
    end
    Δc = Δc/norm(Δc)
    pp = project_particles(jc_cartesian,cl,Δc,cell,box)

    pᵢ = cl.fp[cell.icell]
    i = pᵢ.index
    while i > 0
      if !pᵢ.real
        pᵢ = cl.np[pᵢ.index]
        i = pᵢ.index
        continue
      end
      xpᵢ = pᵢ.coordinates
      i_orig = pᵢ.index_original

      # Partition pp array according to the current projections. This
      # is faster than it looks, because after the first partitioning, the
      # array will be almost correct, and essentially the algorithm
      # will only run over most elements. Avoiding this partitioning or
      # trying to sort the array before always resulted to be much more
      # expensive.
      xproj = dot(xpᵢ-cell.center,Δc)
      n = partition!(pp, el -> el.xproj - xproj <= cutoff)

      for j in 1:n
        pⱼ = pp[j]
        j_orig = pⱼ.index_original
        if i_orig < j_orig
          xpⱼ = pⱼ.coordinates
          d2 = norm_sqr(xpᵢ - xpⱼ)
          if d2 <= cutoff_sq
            output = f(xpᵢ,xpⱼ,i_orig,j_orig,d2,output)
          end
        end
      end
      pᵢ = cl.np[pᵢ.index]
      i = pᵢ.index
    end
  end

  return output
end

#
# Inner loop for Orthorhombic cells is faster because we can guarantee that
# there are not repeated computations even if running over half of the cells.
#
function inner_loop!(
  f,box::Box{OrthorhombicCell},icell,
  cl::CellList{N,T},
  output
) where {N,T}
  @unpack cutoff_sq = box
  cell = cl.cwp[icell]

  # loop over list of non-repeated particles of cell ic
  pᵢ = cl.fp[cell.icell]
  i = pᵢ.index
  while i > 0
    xpᵢ = pᵢ.coordinates
    pⱼ = cl.np[i]
    j = pⱼ.index # here j starts from i+1
    while j > 0
      xpⱼ = pⱼ.coordinates
      d2 = norm_sqr(xpᵢ - xpⱼ)
      if d2 <= cutoff_sq
        i_orig = pᵢ.index_original
        j_orig = pⱼ.index_original
        output = f(xpᵢ,xpⱼ,i_orig,j_orig,d2,output)
      end
      pⱼ = cl.np[pⱼ.index]
      j = pⱼ.index
    end
    pᵢ = cl.np[pᵢ.index]
    i = pᵢ.index
  end

  for jcell in neighbour_cells(box)
    output = cell_output!(f,box,cell,cl,output,cell.cartesian+jcell)
  end

  return output
end

function cell_output!(
  f,
  box::Box{OrthorhombicCell},
  icell,
  cl::CellList{N,T},
  output,
  jc_cartesian
) where {N,T}
  @unpack cutoff, cutoff_sq = box

  # Vector connecting cell centers
  Δc = cell_center(jc_cartesian,box) - icell.center
  Δc = Δc/norm(Δc)
  pp = project_particles(jc_cartesian,cl,Δc,icell,box)

  # Loop over particles of cell icell
  pᵢ = cl.fp[icell.icell]
  i = pᵢ.index
  while i > 0
    xpᵢ = pᵢ.coordinates
    xproj = dot(xpᵢ-icell.center,Δc)

    # Partition pp array according to the current projections
    n = partition!(pp, el -> el.xproj - xproj <= cutoff)

    # Compute the interactions
    for j in 1:n
      pⱼ = pp[j]
      xpⱼ = pⱼ.coordinates
      d2 = norm_sqr(xpᵢ - xpⱼ)
      if d2 <= cutoff_sq
        i_orig = pᵢ.index_original
        j_orig = pⱼ.index_original
        output = f(xpᵢ,xpⱼ,i_orig,j_orig,d2,output)
      end
    end
    pᵢ = cl.np[i]
    i = pᵢ.index
  end

  return output
end

#
# Serial version for cross-interaction computations
#
function map_pairwise_serial!(
  f::F, output, box::Box,
  cl::CellListPair{N,T};
  show_progress=show_progress
) where {F,N,T}
  show_progress && (p = Progress(length(cl.small),dt=0))
  for i in eachindex(cl.small)
    output = inner_loop!(f,output,i,box,cl)
    show_progress && next!(p)
  end
  return output
end

#
# Parallel version for cross-interaction computations
#
function map_pairwise_parallel!(
  f::F1, output, box::Box,
  cl::CellListPair{N,T};
  output_threaded=output_threaded,
  reduce::F2=reduce,
  show_progress=show_progress
) where {F1,F2,N,T}
  show_progress && (p = Progress(length(cl.small),dt=1))
  @threads for it in 1:nthreads()
    for i in splitter(it,length(cl.small))
      output_threaded[it] = inner_loop!(f,output_threaded[it],i,box,cl)
      show_progress && next!(p)
    end
  end
  output = reduce(output,output_threaded)
  return output
end

#
# Inner loop of cross-interaction computations. If the two sets
# are large, this might(?) be improved by computing two sepparate
# cell lists and using projection and partitioning.
#
function inner_loop!(
  f,output,i,box,
  cl::CellListPair{N,T}
) where {N,T}
  @unpack nc, cutoff_sq = box
  xpᵢ = wrap_to_first(cl.small[i],box)
  ic = particle_cell(xpᵢ,box)
  for neighbour_cell in neighbour_cells_all(box)
    jc = cell_linear_index(nc,neighbour_cell+ic)
    pⱼ = cl.large.fp[jc]
    j = pⱼ.index
    # loop over particles of cell jc
    while j > 0
      xpⱼ = pⱼ.coordinates
      d2 = norm_sqr(xpᵢ - xpⱼ)
      if d2 <= cutoff_sq
        j_orig = pⱼ.index_original
        output = f(xpᵢ,xpⱼ,i,j_orig,d2,output)
      end
      pⱼ = cl.large.np[j]
      j = pⱼ.index
    end
  end
  return output
end