text:
"""
Various utilities for manipulating spectra

Author: Eric Ford
Created: August 2020
Contact: https://github.com/eford/
"""

import NaNMath

"""
    `apply_doppler_boost!(spectrum, doppler_factor)`
    `apply_doppler_boost!(spectra, df)`

Apply Doppler boost to spectra's λ's and update its metadata[:doppler_factor], so it will know how to undo the transform.
# Arguments:
- `spectrum::AbstractSpectra`: spectrum to be boosted
- `doppler_factor`: boost factor (1 = noop)

or:
- `spectra`: spectra to be boosted
- `df`: DataFrame provides `:drift` and `:ssb_rv` (in m/s) for calculating the Doppler boost for each spectrum

Returns spectrum/spectra with λ boosted
TODO: Improve documentation formatting.  This can serve as a template.
"""
function apply_doppler_boost! end

function apply_doppler_boost!(spectra::AS,doppler_factor::Real) where {AS<:AbstractSpectra}
    if doppler_factor == one(doppler_factor) return spectra end
    #println("# t= ",time, " doppler_factor= ",doppler_factor)
    spectra.λ .*= doppler_factor
    if hasproperty(spectra.metadata,:doppler_factor)
        spectra.metadata[:doppler_factor] /= doppler_factor
    else
        spectra.metadata[:doppler_factor] = 1/doppler_factor
    end
    return spectra
end

global have_issued_ssb_warning = false
global have_issued_drift_warning = false
global have_issued_diffext_warning = false
function apply_doppler_boost!(spectra::AS, dict::AbstractDict ) where { AS<:AbstractSpectra }
    global have_issued_ssb_warning, have_issued_drift_warning, have_issued_diffext_warning, have_issued_multiple_ssb_warning
    local doppler_factor = one(eltype(spectra.λ))
    #= NEID drift model is now included in wavelengths provided
    if !haskey(dict,:drift_rv) && !haskey(dict,:drift_z ) && !have_issued_drift_warning
        @info "apply_doppler_boost! didn't find :drift_rv or :drivft_z to apply."
        have_issued_drift_warning = true
    end
    if  haskey(dict,:drift_z)
        doppler_factor *= calc_doppler_factor(z=dict[:drift_z])
    elseif  haskey(dict,:drift_rv)
        doppler_factor *= calc_doppler_factor(dict[:drift_rv])
    end
    =#
    if !haskey(dict,:ssb_rv) && !haskey(dict,:ssbz) && !haskey(dict,:ssb_rv_kmps) && !have_issued_ssb_warning
         @info "apply_doppler_boost! didn't find :ssbz, :ssb_rv or :ssb_rv_kmps to apply."
         have_issued_ssb_warning = true
    end
    if  haskey(dict,:ssbz)
        doppler_factor   *= calc_doppler_factor(z=dict[:ssbz])
        if  haskey(dict,:ssb_rv) || haskey(dict,:ssb_rv_kmps)
            @info "apply_doppler_boost! found ssbz and either ssb_rv or ssb_rv_kmps. Defaulting to only use ssbz."
            have_issued_multiple_ssb_warning = true
        end
    elseif  haskey(dict,:ssb_rv)
        doppler_factor   *= calc_doppler_factor(dict[:ssb_rv])
        if  haskey(dict,:ssb_rv_kmps)
            @info "apply_doppler_boost! found ssb_rv (in m/s) and ssb_rv_kmps (in km/s). Defaulting to only use ssb_rv."
            have_issued_multiple_ssb_warning = true
        end
    elseif  haskey(dict,:ssb_rv_kmps)
        doppler_factor   *= calc_doppler_factor(dict[:ssb_rv_kmps]*1000.0)
    end
    #=
    # Now plan to apply correction at end, rather than to wavelength
    if !haskey(dict,:diff_ext_rv ) && !have_issued_diffext_warning
        @info "apply_doppler_boost! didn't find :diff_ext_rv to apply."
        have_issued_diffext_warning = true
    end
    if  haskey(dict,:diff_ext_rv)  doppler_factor   *= calc_doppler_factor.(dict[:diff_ext_rv])  end
    =#
    if  haskey(dict,:binaryRV)  doppler_factor   *= calc_doppler_factor.(dict[:binaryRV])  end
    apply_doppler_boost!(spectra,doppler_factor)
    return spectra
end

function apply_doppler_boost!(spectra::AbstractArray{AS}, df::DataFrame ) where { AS<:AbstractSpectra }
    @assert size(spectra,1) == size(df,1)
    global have_issued_ssb_warning, have_issued_drift_warning, have_issued_diffext_warning
    local doppler_factor = ones(size(spectra))
    #if !hasproperty(df,:drift) @info "apply_doppler_boost! didn't find :drift to apply."   end
    #if  hasproperty(df,:drift)        doppler_factor .*= calc_doppler_factor.(df[!,:drift])          end
    #= NEID drift model is now included in wavelengths provided
    if !haskey(dict,:drift_rv) && !haskey(dict,:drift_z ) && !have_issued_drift_warning
        @info "apply_doppler_boost! didn't find :drift_rv or :drivft_z to apply."
        have_issued_drift_warning = true
    end
    if  haskey(dict,:drift_z)
        doppler_factor *= calc_doppler_factor(z=dict[:drift_z])
    elseif  haskey(dict,:drift_rv)
        doppler_factor *= calc_doppler_factor(dict[:drift_rv])
    end
    =#
    if !hasproperty(df,:ssb_rv) && !hasproperty(df,:ssbz)  && !have_issued_ssb_warning
        @info "apply_doppler_boost! didn't find :ssb_rv or :ssbz to apply."
        have_issued_ssb_warning = true
    end
    if  hasproperty(df,:ssbz)
        doppler_factor   .*= calc_doppler_factor.(z=df[!,:ssbz])
    elseif  hasproperty(df,:ssb_rv)
        doppler_factor   .*= calc_doppler_factor.(df[!,:ssb_rv])
    end
    #= Now plan to apply correction at end, rather than to wavelength
    if !hasproperty(df,:drift) && !have_issued_diffext_warning
        @info "apply_doppler_boost! didn't find :diff_ext_rv to apply."
        have_issued_diffext_warning = true
    end
    if  hasproperty(df,:diff_ext_rv)  doppler_factor   .*= calc_doppler_factor.(df[!,:diff_ext_rv])  end
    =#
    map(x->apply_doppler_boost!(x[1],x[2]), zip(spectra,doppler_factor) );
end

"""
`calc_snr(flux, var)`
`calc_snr(spectrum, pixels, order_idx)`

Calculate total SNR in (region of) spectra.
"""
function calc_snr end

function calc_snr(flux::AbstractArray{T1},var::AbstractArray{T2}) where {T1<:Real, T2<:Real}
    @assert size(flux) == size(var)
    sqrt(NaNMath.sum( max.(0.0,flux.^2 ./ var)))
end

function calc_snr(flux::Real,var::Real)
    flux / sqrt(max(0.0,var))
end

function calc_snr(spectrum::ST, pixels::AR, order::Integer) where { ST<:AbstractSpectra2D, AR<:AbstractRange{Int64}, AA1<:AbstractArray{Int64,1} } #, AAR<:AbstractArray{AR,1} }
    flux = view(spectrum.flux,pixels,order)
    var = view(spectrum.var,pixels,order)
    #=
    #if all(isnan.(flux) .|| isnan.(var) )
    if all( isnan.( flux.^2 ./ var ) )
        println("# all NaNs in order = ", order, " pixels = ", pixels)
        return 0
    end
    =#
    calc_snr( view(spectrum.flux,pixels,order), view(spectrum.var,pixels,order) )
end

function calc_snr(chunk::AC) where { AC<:AbstractChunkOfSpectrum }
    calc_snr( chunk.flux, chunk.var )
end

""" Normalize spectrum, multiplying fluxes by scale_fac. """
function normalize_spectrum!(spectrum::ST, scale_fac::Real) where { ST<:AbstractSpectra }
    @assert 0 < scale_fac < Inf
    @assert !isnan(scale_fac^2)
    spectrum.flux .*= scale_fac
    spectrum.var .*= scale_fac^2
    return spectrum
end

""" Normalize each spectrum based on sum of fluxes in chunk_timeseries region of each spectrum. """
function normalize_spectra!(chunk_timeseries::ACLT, spectra::AS) where { ACLT<:AbstractChunkListTimeseries, ST<:AbstractSpectra, AS<:AbstractArray{ST} }
    @assert length(chunk_timeseries) == length(spectra)
    for t in 1:length(chunk_timeseries)
        scale_fac = calc_normalization(chunk_timeseries.chunk_list[t])
        # println("# t= ",t, " scale_fac= ", scale_fac)
        normalize_spectrum!(spectra[t], scale_fac)
    end
    return chunk_timeseries
end

""" Return instrument associated with spectrum """
function get_inst(spectrum::AS) where { AS<:AbstractSpectra2D }
    return spectrum.inst
end

""" Return instrument associated with first spectrum in array """
function get_inst(spectra::AAS) where { AS<:AbstractSpectra2D, AAS<:AbstractVector{AS} }
    @assert length(spectra)>=1
    get_inst(first(spectra))
end

""" Return the largest minimum wavelength and smallest maximum wavelength of a spectrum.
"""
function get_λ_range(data::CLT) where { CLT<:AbstractSpectra }
   λmin = minimum(data.λ)
   λmax = maximum(data.λ)
   return (min = λmin, max = λmax)
end

""" Return the largest minimum wavelength and smallest maximum wavelength across an array of spectra.
Calls get_λ_range(AbstractSpectra2D) that should be specialized for each instrument. """
function get_λ_range(data::ACLT) where { CLT<:AbstractSpectra, ACLT<:AbstractArray{CLT} }
   λminmax = get_λ_range.(data)
   λmin = maximum(map(p->p[1],λminmax))
   λmax = minimum(map(p->p[2],λminmax))
   return (min = λmin, max = λmax)
end

""" Extract the metadata from a time series of spectra and return it as an array. """
function make_vec_metadata_from_spectral_timeseries(spec_arr::AA) where { AS<:AbstractSpectra, AA<:AbstractArray{AS,1} }
    map(s->s.metadata,spec_arr)
end

""" Return DataFrame with information about which pixels and wavelengths to use from each order"""
function get_order_info(all_spectra::AAS; orders_to_use::Union{UnitRange,AbstractVector{T1} } = min_order(get_inst(all_spectra)):max_order(get_inst(all_spectra))) where { AS<:AbstractSpectra2D, AAS<:AbstractVector{AS}, T1<:Integer }
  @assert length(orders_to_use) >= 1

  inst = get_inst(all_spectra)
  num_orders = length(orders_to_use)
  #pixel_ranges = Array{UnitRange,1}(undef,num_orders)
  order_info = Any[]
  for order in orders_to_use #min_order(inst):max_order(inst)
    this_order_info = Dict{Symbol,Any}(:order=>order)
    pixel_range = get_pixel_range(inst,order)
    min_pixels_in_order = 60
    #=
    pixel_range = min_col_default(inst,order):max_col_default(inst,order)
    if typeof(inst) <: AnyEXPRES &&
        ( EXPRES.max_col_excalibur(inst, order) - EXPRES.min_col_excalibur(inst, order) >= min_pixels_in_order )

        minc = max(first(pixel_range), EXPRES.min_col_excalibur(inst,order), EXPRES.min_col_nonnan(inst,order))
        maxc = min(last(pixel_range), EXPRES.max_col_excalibur(inst,order), EXPRES.max_col_nonnan(inst,order))
        pixel_range = minc:maxc
      end
    =#
    if length(pixel_range) > min_pixels_in_order
        this_order_info[:pixel_range] = pixel_range
        (this_order_info[:λ_min], this_order_info[:λ_max], this_order_info[:mean_Δv]) = get_shared_wavelength_range_for_order(all_spectra,order, pixels_to_use=pixel_range )
        push!(order_info,this_order_info)
    end
  end
  return DataFrame(order_info)
end

function calc_mean_Δv(λ::AbstractVector{T1}, λ_min::Real, λ_max::Real ) where { T1<:Real }
    idx_min = searchsortedfirst(λ,λ_min)
    idx_max = searchsortedlast(λ,λ_max)
    @assert 1 <= idx_min <= length(λ)
    @assert 1 <= idx_max <= length(λ)
    Δv = log(λ[idx_max]/λ[idx_min])/(idx_max-idx_min) * RvSpectMLBase.speed_of_light_mps
end

function get_shared_wavelength_range_for_order(spectra::AbstractVector{AS}, order::Integer; pixels_to_use::AR = min_col_default(get_inst(spectra),order):max_col_default(get_inst(spectra),order),
             boost_factor::AA1 = ones(length(spectra)), verbose::Bool=false ) where {
             AS<:AbstractSpectra2D, AR<:AbstractRange{Int64}, T1<:Real, AA1<:AbstractVector{T1} }
        num_obs = length(spectra)
        @assert 1 <= num_obs # <= max_num_spectra
        if verbose   println("# order = ", order, " pixels_to_use = ", pixels_to_use)   end
        λ_min = maximum(minimum(spectra[t].λ[pixels_to_use,order])/boost_factor[t] for t in 1:num_obs)
        λ_max = minimum(maximum(spectra[t].λ[pixels_to_use,order])/boost_factor[t] for t in 1:num_obs)
        mean_Δv = sum(map(t->calc_mean_Δv(view(spectra[t].λ, pixels_to_use, order), λ_min, λ_max),1:num_obs)) / num_obs
        return (λ_min=λ_min, λ_max=λ_max, mean_Δv=mean_Δv)
        #=
        @assert spacing == :Log || spacing == :Linear
        if spacing == :Log
            @warn "There's some issues with end points exceeding the bounds.  Round off error?  May cause bounds errors."
        end
        =#
end

function discard_large_metadata(data::Union{ACLT,AS, AAS}) where { ACLT<:AbstractChunkListTimeseries, AS<:AbstractSpectra, AAS<:AbstractArray{AS} }
    if typeof(data.inst) <: AnyTheoreticalInstrument
        # Nothing to do
    else
        inst_module = get_inst_module(data.inst)
        inst_module.discard_large_metadata(data)
    end

    #=
    if typeof(data.inst) <: AnyEXPRES
        inst_module.discard_blaze(data)
        inst_module.discard_continuum(data)
        inst_module.discard_tellurics(data)
        inst_module.discard_pixel_mask(data)
        inst_module.discard_excalibur_mask(data)
    end
    if typeof(data.inst) <: AnyNEID
        #discard_blaze(data)
        #discard_continuum(data)
        #discard_tellurics(data)
        #discard_pixel_mask(data)
        #discard_excalibur_mask(data)
    end
    =#
end

function discard_blaze(metadata::Dict{Symbol,Any} )
    delete!(metadata,:blaze)
end

function discard_blaze(data::AST) where { AST<:AbstractSpectra }
   discard_blaze(data.metadata)
end

function discard_blaze(data::ACLT) where { CLT<:AbstractSpectra, ACLT<:AbstractArray{CLT} }
   map(spectra->discard_blaze(spectra.metadata),data)
end

function discard_blaze(data::CLT) where { CLT<:AbstractChunkListTimeseries }
   map(discard_blaze,data.metadata)
end

function discard_continuum(metadata::Dict{Symbol,Any} )
    delete!(metadata,:continuum)
end

function discard_continuum(data::AST) where { AST<:AbstractSpectra }
   discard_continuum(data.metadata)
end

function discard_continuum(data::ACLT) where { CLT<:AbstractSpectra, ACLT<:AbstractArray{CLT} }
   map(spectra->discard_continuum(spectra.metadata),data)
end

function discard_continuum(data::CLT) where { CLT<:AbstractChunkListTimeseries }
   map(discard_continuum,data.metadata)
end

function discard_tellurics(metadata::Dict{Symbol,Any} )
    delete!(metadata,:tellurics)
end

function discard_tellurics(data::AST) where { AST<:AbstractSpectra }
   discard_tellurics(data.metadata)
end

function discard_tellurics(data::ACLT) where { CLT<:AbstractSpectra, ACLT<:AbstractArray{CLT} }
   map(spectra->discard_tellurics(spectra.metadata),data)
end

function discard_tellurics(data::CLT) where { CLT<:AbstractChunkListTimeseries }
   map(discard_tellurics,data.metadata)
end

function discard_pixel_mask(metadata::Dict{Symbol,Any} )
    delete!(metadata,:pixel_mask)
end

function discard_pixel_mask(data::AST) where { AST<:AbstractSpectra }
   discard_pixel_mask(data.metadata)
end

function discard_pixel_mask(data::ACLT) where { CLT<:AbstractSpectra, ACLT<:AbstractArray{CLT} }
   map(spectra->discard_pixel_mask(spectra.metadata),data)
end

function discard_pixel_mask(data::CLT) where { CLT<:AbstractChunkListTimeseries }
   map(discard_pixel_mask,data.metadata)
end

function discard_excalibur_mask(metadata::Dict{Symbol,Any} )
    delete!(metadata,:excalibur_mask)
end

function discard_excalibur_mask(data::AST) where { AST<:AbstractSpectra }
   discard_excalibur_mask(data.metadata)
end

function discard_excalibur_mask(data::ACLT) where { CLT<:AbstractSpectra, ACLT<:AbstractArray{CLT} }
   map(spectra->discard_excalibur_mask(spectra.metadata),data)
end

function discard_excalibur_mask(data::CLT) where { CLT<:AbstractChunkListTimeseries }
   map(discard_excalibur_mask,data.metadata)
end

max_stars_repo_name: RvSpectML/RvSpectMLTemplate

max_issues_repo_name: RvSpectML/RvSpectMLTemplate

max_forks_repo_name: RvSpectML/RvSpectMLTemplate


text:
#= addition.jl
Philip Barrett, pobarrett@gmail.com
30may2016, Chicago

Tests the addition functions for the polygon class
=#

## Binary set sum
Z = [ 1.0 -1; 1 1; -1 1; -1 -1 ]
a = Polygon( pts=Z )
h = Polygon( pts=[ 0.0 1 ; 1 0 ; -1 0 ; 0 -1 ] )
NN = 40
exactdirs = vcat([ [ cos(2*(i-1)*pi/NN) sin(2*(i-1)*pi/NN) ] for i in 1:NN ]...)
    # The union of the face directions.  Will give an exact result.
exact = setSum( a, h, exactdirs )
exactinner = setSum( a, h, exactdirs, false )

@test deeDoop(exact.pts) ≈ [ -2 1; -2.0 -1; -1 -2; 1 -2; 2.0 -1; 2 1; 1 2; -1 2 ]
@test deeDoop(exact.pts) ≈ [ exactinner.pts[ end, :]; exactinner.pts[ 1:(end-1), :] ]
    # Arises because we have a vertical left edge.  Annoying.

## Array set sum
exactarray = setSum( [a, h], exactdirs )
@test exact.pts == exactarray.pts

ext: jl


text:
using BioXMLSciMLTest, SBMLBioModelsRepository

# sbml_test_suite()
fns = get_sbml_suite_fns()
# @show fns
@test !isempty(fns)

files_to_sciml(fns; pmap=true)
df = results_df()


text:
const CONE_ORDER = Dict(
    MOI.Zeros=>1,
    MOI.Nonnegatives=>2,
    MOI.SecondOrderCone=>3,
    MOI.PositiveSemidefiniteConeTriangle=>4,
    MOI.ExponentialCone=>5,
    MOI.DualExponentialCone=>6,
    MOI.PowerCone=>7
)

const SUPPORTED_SOLVERS = Dict(
    "ECOS" => ECOS.Optimizer,
    "Hypatia" => Hypatia.Optimizer,
    "SCS" => SCS.Optimizer
)

# TODO: specify for each solver
SOLVER_OPTION_SET = [:max_iters, :eps, :warm_start, :scale, :verbose]

function solve_and_diff(
    A::AbstractMatrix{S}, b::AbstractVector{S}, c::Vector{S}, cone_prod::Vector{T}; kwargs...
) where {T <: MOI.AbstractVectorSet, S <: AbstractFloat}
    kwargs = Dict(kwargs)
    m,n = size(A)

    # Arguments check
    (length(b) == m && length(c) == n) || throw(DimensionMismatch("Mismatch between A, b, c"))
    all([typeof(c) <: SUPPORTED_INPUT_SETS for c in cone_prod]) || throw(ArgumentError("Unsupported cones"))
    all([MOI.dimension(c) > 0 for c in cone_prod]) || throw(ArgumentError("Cones must have nonzero dimension"))
    sum([MOI.dimension(c) for c in cone_prod]) == m || throw(DimensionMismatch("Mismatch between dimension of c and cone"))
    if haskey(kwargs, :solver) && kwargs[:solver] in keys(SUPPORTED_SOLVERS)
        optimizer = SUPPORTED_SOLVERS[kwargs[:solver]]
        # TODO: check cone_prod for each solver
    elseif haskey(kwargs, :solver)
        throw(ArgumentError("Invalid solver"))
    else
        optimizer = SCS.Optimizer
    end

    warm_start = haskey(kwargs, :warm_start) ? kwargs[:warm_start] : nothing
    if haskey(kwargs, :method) && kwars[:method] == "dense"
        use_lsqr = false
    else
        use_lsqr = true
    end

    options = Dict()
    for option in SOLVER_OPTION_SET
        if haskey(kwargs, option)
            options[option] = kwargs[option]
        end
    end
    options[:verbose] = haskey(options, :verbose) ? options[:verbose] : false

    return _solve_and_diff(A, b, c, cone_prod, optimizer, use_lsqr; options...)
end

function _solve_and_diff(A, b, c, cone_prod, optimizer, use_lsqr; options...)
    m,n = size(A)
    typeof(A) <: SparseMatrixCSC && dropzeros!(A)
    x_star, y_star, s_star = solve_opt_problem(A, b, c, cone_prod, optimizer; options...)

    Q = spzeros(m+n+1,m+n+1)
    Q[1:n,n+1:n+m]      .= A'
    Q[n+1:n+m,1:n]      .= -A
    Q[n+1:n+m,m+n+1]    .= b
    Q[m+n+1,n+1:n+m]    .= -b
    Q[1:n,m+n+1]        .= c
    Q[m+n+1,1:n]        .= -c

    # precompute M
    u, v, w = (x_star, y_star - s_star, 1.0)
    M = (Q - I) * dpi_z(u, v, w, cone_prod) + I

    function pullback(dx)
        # Assume that dy = ds = 0
        dz = [
            dx;
            zeros(length(y_star));
            -x_star' * dx
        ]

        g = use_lsqr ? lsqr(-M', dz, atol=1e-10, btol=1e-10) : -M' \ dz

        # TODO: make more efficient by only pulling needed cols/rows
        dQ = g * pi_z(u, v, w, cone_prod)'
        dA = dQ[1:n,n+1:n+m]'  - dQ[n+1:n+m,1:n]    # dQ12' - dQ21
        db = dQ[n+1:n+m,m+n+1] - dQ[m+n+1,n+1:n+m]  # dQ23 - dQ32'
        dc = dQ[1:n,m+n+1]     - dQ[m+n+1,1:n]      # dQ13 - dQ31'
        return dA, db, dc
    end

    function pushforward(dA, db, dc)
        # TODO: Could we compute a matrix vector product instead of full matrix?
        DπKdual_v = d_project_onto_cone(v, [MOI.dual_set(c) for c in cone_prod])

        dQ = spzeros(n+m+1, n+m+1)
        dQ[1:n,n+1:n+m]      .= dA'
        dQ[n+1:n+m,1:n]      .= -dA
        dQ[n+1:n+m,m+n+1]    .= db
        dQ[m+n+1,n+1:n+m]    .= -db #dont use tranpose bc of broadcast
        dQ[1:n,m+n+1]        .= dc
        dQ[m+n+1,1:n]        .= -dc

        g = dQ*pi_z(u, v, w, cone_prod)

        dz = use_lsqr ? lsqr(-M, g, atol=1e-10, btol=1e-10) : -M \ g
        @views du, dv, dw = dz[1:n], dz[n+1:n+m], dz[n+m+1]
        dx = du .- dw*x_star
        dy = DπKdual_v*dv - dw*y_star
        ds = DπKdual_v*dv - dv - dw*s_star
        return dx, dy, ds
    end

    return x_star, y_star, s_star, pushforward, pullback
end

function reorder_opt_problem_scs(A, b, cones)
    # sort the cones so that they're in the correct order
    sort_ixs = sortperm(cones, by=cone->CONE_ORDER[typeof(cone)])
    sorted_cones = cones[sort_ixs]

    f = 0  # the number of zero cones
    l = 0  # the number of nonneg cones
    q = Int[]  # the array of SOC sizes
    s = Int[]  # the array of SDC sizes
    ep = 0  # the number of exponential cones
    ed = 0  # the number of dual exponential cones
    p = Float64[]  # the array of power cone parameters

    # create the parameters for SCS
    for cone in sorted_cones
        if typeof(cone) <: MOI.Zeros
            f += MOI.dimension(cone)
        elseif typeof(cone) <: MOI.Nonnegatives
            l += MOI.dimension(cone)
        elseif typeof(cone) <: MOI.SecondOrderCone
            push!(q, MOI.dimension(cone))
        elseif typeof(cone) <: MOI.PositiveSemidefiniteConeTriangle
            push!(s, cone.side_dimension)
        elseif typeof(cone) <: MOI.ExponentialCone
            ep += 1
        elseif typeof(cone) <: MOI.DualExponentialCone
            ep += 1
        elseif typeof(cone) <: MOI.PowerCone
            push!(p, MOI.dimension(a))
        elseif typeof(cone) <: MOI.DualPowerCone
            push!(p, -MOI.dimension(a))
        end
    end
    cone_dict = Dict(:f=>f, :l=>l, :q=>q, :s=>s, :ep=>ep, :ed=>ed, :p=>p)

    original_ixs = []
    start_ix = 1
    for cone in cones
        end_ix = start_ix + MOI.dimension(cone) - 1
        push!(original_ixs, start_ix:end_ix)
        start_ix = end_ix + 1
    end

    index_map = [original_ixs[i] for i in sort_ixs]
    index_map = collect(Iterators.flatten(index_map))
    reverse_map = sortperm(index_map)

    return A[index_map, :], b[index_map], cone_dict, reverse_map
end

function solve_opt_problem(A, b, c, cone_prod, optimizer_factory; options...)
    # TODO: potentially play with the scale depending on the cones?
    #       e.g. see https://github.com/jump-dev/Convex.jl/issues/104
    #       Essentially want to have the primal and dual res down at the same rate
    #           primal slower -> inc. scale
    #           dual slower -> dec. scale
    m,n = size(A)

    if optimizer_factory == SCS.Optimizer
        A, b = rewrite_sdps_to_row_major(A, b, cone_prod)
        A_reordered, b_reordered, cone_dict, reverse_map = reorder_opt_problem_scs(A, b, cone_prod)
        f, l, q, s = cone_dict[:f], cone_dict[:l], cone_dict[:q], cone_dict[:s]
        ep, ed, p = cone_dict[:ep], cone_dict[:ed], cone_dict[:p]

        # TODO: clean up argument passing
        # TODO: for some reason this fails when b is a sparse vector -- see SCS.jl
        scs_sol = SCS_solve(SCS.IndirectSolver, m, n,
                            A_reordered, b_reordered, c,
                            f, l, q, s, ep, ed, p;
                            options...)
        # TODO: additional solution checks
        #/* SCS returns one of the following integers: */
        #define SCS_INFEASIBLE_INACCURATE (-7)
        #define SCS_UNBOUNDED_INACCURATE (-6)
        #define SCS_SIGINT (-5)
        #define SCS_FAILED (-4)
        #define SCS_INDETERMINATE (-3)
        #define SCS_INFEASIBLE (-2) /* primal infeasible, dual unbounded   */
        #define SCS_UNBOUNDED (-1)  /* primal unbounded, dual infeasible   */
        #define SCS_UNFINISHED (0)  /* never returned, used as placeholder */
        #define SCS_SOLVED (1)
        #define SCS_SOLVED_INACCURATE (2)
        if !(scs_sol.ret_val in [1,2])
            error("Model not solved correctly.
                   SCS termination status: $(scs_sol.ret_val)\n
                   Info: $(scs_sol.info)")
        elseif scs_sol.ret_val == 2
            @warn "Maximum iterations reached. Set max_iters to increase accuracy"
        end

        return scs_sol.x, scs_sol.y[reverse_map], scs_sol.s[reverse_map]
    end

    #TODO: interface directly with other optimizer solvers

    # If another optimizer is used, construct the problem with JuMP
    # NOTE: this has a high computational overhead
    model = Model()
    set_optimizer(model, optimizer_factory)
    # TODO: allow for attributes in general optimizer, e.g.
    # optimizer_with_attributes(SCS.Optimizer, "eps" => 1e-6, etc...)
    @variable(model, x[1:n])
    @variable(model, s[1:m])
    @objective(model, Min, c'*x)
    con = @constraint(model, A*x + s .== b)
    curr = 1
    for cone in cone_prod
        @constraint(model, s[curr:curr+MOI.dimension(cone)-1] in cone)
        curr += MOI.dimension(cone)
    end
    optimize!(model)
    if !(termination_status(model) in [MOI.OPTIMAL, MOI.ALMOST_OPTIMAL])
        error("Model not solved correctly.
               Termination status: $(termination_status(model))")
    elseif termination_status(model) == MOI.ALMOST_OPTIMAL
        @warn "Maximum iterations reached. Set max_iters to increase accuracy"
    end

    # Ensure dual has the correct sign
    y = isapprox(A'* dual.(con) + c, zeros(length(c)), atol=1e-5) ? dual.(con) : -dual.(con)
    return value.(x), y, value.(s)
end

function pi_z(u, v, w, cone_prod; dual=false)
    cone = dual ? [MOI.dual_set(c) for c in cone_prod] : cone_prod
    return vcat(
        u,
        project_onto_cone(v, cone),
        max(w, 0.0)
    )
end

function dpi_z(u, v, w, cone_prod; dual=false)
    cone = dual ? [MOI.dual_set(c) for c in cone_prod] : cone_prod
    return blockdiag(
        sparse(Matrix{Float64}(I, length(u), length(u))),
        d_project_onto_cone(v, cone),
        sparse(Matrix{Float64}(I, 1,1))
    )
end


text: ConstructionBase.setproperties(x::ComponentVector, patch::NamedTuple) = ComponentVector(x; patch...)


text:
"""
	iterate_until_collisionless(f)

Sample `f` until the result is collisionless.
"""
function iterate_until_collisionless(f)

    """samples f until the result is collisionless

	as it takes a function, example usage :

		random_mode_occupation_collisionless(n::Int,m::Int) = iterate_until_collisionless(() -> random_mode_occupation(n,m))
	"""

    while true

        state = f()
        if is_collisionless(state)
            return state
        end

    end

end

"""
	fill_arrangement(occupation_vector)
	fill_arrangement(r::ModeOccupation)
	fill_arrangement(input::Input)

Convert a mode occupation list to a mode assignement.
"""
function fill_arrangement(occupation_vector)

    """from a mode occupation list to a mode assignment list (Tichy : s to d(s))

    For instance if vect{s} = (2,0,1) then vect{d} = (1,1,3) """

    arrangement = zeros(eltype(occupation_vector), sum(occupation_vector))

    position = 1
    for i in 1:length(occupation_vector)
        for j in 1:occupation_vector[i]
            arrangement[position] = i
            position += 1
        end
    end

    arrangement

end

fill_arrangement(r::ModeOccupation) = fill_arrangement(r.state)

fill_arrangement(inp::Input) = fill_arrangement(inp.r)

"""
	random_occupancy(n::Int, m::Int)

Return a vector of size `m`	with `n` randomly placed ones.
"""
function random_occupancy(n::Int, m::Int)

	""" returns a vector of size m with n randomly placed ones """
	if n > m
		throw(ArgumentError("Implemented at most one photon per mode"))
	else
		occupancy_vector = shuffle(append!(ones(n), zeros(m-n)))
		occupancy_vector = Int.(occupancy_vector)
	end

	return occupancy_vector
end

"""
	random_mode_occupation(n::Int, m::Int)

Create a [`ModeOccupation`](@ref) from a mode occupation list of `n` ramdomly placed ones
among `m` sites.
"""
random_mode_occupation(n::Int, m::Int) = ModeOccupation(random_occupancy(n,m))

"""
	random_mode_occupation_collisionless(n::Int, m::Int)

Create a [`ModeOccupation`](@ref) from a random mode occupation that is likely collisionless.
"""
function random_mode_occupation_collisionless(n::Int, m::Int)

	n<=m ? iterate_until_collisionless(() -> random_mode_occupation(n,m)) : error("n>m cannot make for collisionless mode occupations")

end

"""
	at_most_one_photon_per_bin(occupancy_vector::Vector{Int})
	check_at_most_one_particle_per_mode(occ)

Check wether `occupancy_vector` contains more than one photon per site.
"""
function at_most_one_photon_per_bin(occupancy_vector::Vector{Int})
	all(in([0,1]).(occupancy_vector))
end

check_at_most_one_particle_per_mode(occ) = at_most_one_photon_per_bin(occ) ? nothing : error("more than one input per mode")

"""
	occupancy_vector_to_partition(occupancy_vector)
	occupancy_vector_to_mode_occupancy(occupancy_vector)

Return a partition of occupied modes from an `occupancy_vector`.
"""
function occupancy_vector_to_partition(occupancy_vector)

	partition = []

	check_at_most_one_particle_per_mode(occupancy_vector)
	# a partition cannot be defined by taking twice the same mode, it would lead to
	# hidden errors in the results (double counting)

	for (mode, selected) in enumerate(occupancy_vector)
		if selected == 1
			append!(partition, mode)
		end
	end
	partition

end

"""
	occupancy_vector_to_mode_occupancy(occupancy_vector)

Return a partition of occupied modes from an `occupancy_vector`.
"""
function occupancy_vector_to_mode_occupancy(occupancy_vector)
	occupancy_vector_to_partition(occupancy_vector)
end

"""
	scattering_matrix(U::Matrix, input_state::Vector{Int}, output_state::Vector{Int})
	scattering_matrix(U::Interferometer, r::ModeOccupation, s::ModeOccupation)
	scattering_matrix(U::Interferometer, i::Input, o::FockDetection)

Return the submatrix of `U` whose rows and columns are respectively defined by
`input_state` and `output_state`.

!!! note "Reference"
	[http://arxiv.org/abs/quant-ph/0406127v1](http://arxiv.org/abs/quant-ph/0406127v1)
"""
function scattering_matrix(U::Matrix, input_state::Vector{Int}, output_state::Vector{Int})

    """
    U = interferometer matrix, size m*m
    input_state = input occupation number vector (s[i] is the number of photons in the ith input)
    output_state = output occupation number vector (r[i] is the number of photons in the ith output)
    n photons

    follows http://arxiv.org/abs/quant-ph/0406127v1
    """

    m = size(U,1)
    n = sum(input_state)

    if length(input_state) != m || length(output_state) != m
        throw(DimensionMismatch())
    elseif sum(input_state) != sum(output_state)
        error("photon number not the same at input and output")
    end

    index_input = fill_arrangement(input_state)
    index_output = fill_arrangement(output_state)

	try
		U[index_input,index_output]
	catch err
		if isa(MethodError, err)
			copy(U)[index_input,index_output]
		else
			println(err)
		end
	end

end

scattering_matrix(interf::Interferometer, r::ModeOccupation, s::ModeOccupation) = scattering_matrix(interf.U,r.state,s.state)
scattering_matrix(interf::Interferometer, i::Input, o::FockDetection) = scattering_matrix(interf.U,i.r,o.s)

function vector_factorial(occupancy_vector)

    """returns the function mu at the denominator of event probabilities"""

    prod(factorial.(occupancy_vector))

end

vector_factorial(r::ModeOccupation) = vector_factorial(r.state)
vector_factorial(i::Input) = vector_factorial(i.r)
vector_factorial(o::FockDetection) = vector_factorial(o.s)

"""
	bosonic_amplitude(U, input_state, output_state, permanent=ryser)
	process_amplitude(U, input_state, output_state, permanent=ryser)

Compute the probability amplitude to go from `input_state` to `output_state`
through the interferomter `U` in the [`Bosonic`](@ref) case.
"""
function bosonic_amplitude(U, input_state, output_state, permanent = ryser)

    """event amplitude"""

    permanent(scattering_matrix(U, input_state, output_state))/sqrt(vector_factorial(input_state) * vector_factorial(output_state))
end

function process_amplitude(U, input_state, output_state, permanent = ryser)

	bosonic_amplitude(U, input_state, output_state, permanent)
end

"""
	bosonic_probability(U, input_state, output_state)
	process_probability(U, input_state, output_state)

Compute the probability to go from `input_state` to `output_state`
through the interferometer `U` in the [`Bosonic`](@ref) case.
"""
function bosonic_probability(U, input_state, output_state)

	"""bosonic process_probability"""

    abs(process_amplitude(U, input_state, output_state))^2

end

function process_probability(U, input_state, output_state)
	#@warn "obsolete function, use bosonic_probability or probability"
	bosonic_probability(U, input_state, output_state)
end

"""
	distinguishable_probability(U, input_state, output_state, permanent=ryser)
	process_probability_distinguishable(U, input_state, output_state, permanent=ryser)

Compute the probability to go from `input_state` to `output_state` through
the interferomter `U` in the [`Distinguishable`](@ref) case.
"""
function distinguishable_probability(U, input_state, output_state, permanent = ryser)

	"""distinguishable (or classical) process_probability"""

	permanent(abs.(scattering_matrix(U, input_state, output_state)).^2)/sqrt(vector_factorial(input_state) * vector_factorial(output_state))

end

function process_probability_distinguishable(U, input_state, output_state, permanent = ryser)

	#@warn "obsolete function, use distinguishable_probability or probability"

	distinguishable_probability(U, input_state, output_state, permanent)
end

### need to implement partial distinguishability processs probabilities ###

"""
	process_probability_partial(U, S, input_state, output_state)
	process_probability_partial(interf::Interferometer, input_state::Input{TIn} where {TIn<:PartDist},output_state::FockDetection)

Compute the probability to go from `input_state` to `output_state` through the
interferometer `U` in the [`PartDist`](@ref) case where partial distinguishable is described
by the [`GramMatrix`](@ref) `S`.

!!! note "Reference"
    [https://arxiv.org/abs/1410.7687](https://arxiv.org/abs/1410.7687)
"""
function process_probability_partial(U, S, input_state,output_state)

    """computes the partially distinguishable process probability according to Tichy's tensor permanent https://arxiv.org/abs/1410.7687"""

    n = size(S,1)

    @argcheck sum(input_state) == sum(output_state) "particles not conserved"

    @argcheck sum(input_state) == n "S matrix doesnt have the same number of photons as the input"

	M = scattering_matrix(U, input_state, output_state)
	W = Array{eltype(U)}(undef, (n,n,n))

	for ss in 1:n
	    for rr in 1:n
	        for j in 1:n
	            W[ss,rr,j] = M[ss, j] * conj(M[rr, j]) * S[rr, ss]
	        end
	    end
	end

    1/(vector_factorial(input_state) * vector_factorial(output_state)) * ryser_tensor(W)

end

process_probability_partial(interf::Interferometer, input_state::Input{TIn} where {TIn<:PartDist},output_state::FockDetection) = process_probability_partial(interf.U, input_state.G.S, input_state.r.state,output_state.s.state)

"""
	compute_probability(ev::Event{TIn, TOut}) where {TIn<:InputType, TOut<:FockDetection}

Given an [`Event`](@ref), gives the probability to get the outcome `TOut` when `TIn`
passes though the interferometer `ev.interferometer`.
"""
function compute_probability!(ev::Event{TIn,TOut}) where {TIn<:InputType, TOut<:FockDetection}

	check_probability_empty(ev)

	ev.proba_params.precision = eps()
	ev.proba_params.failure_probability = 0

	if TIn == Bosonic

		ev.proba_params.probability = bosonic_probability(ev.interferometer.U, ev.input_state.r.state, ev.output_measurement.s.state)

	elseif TIn == Distinguishable

		ev.proba_params.probability = distinguishable_probability(ev.interferometer.U, ev.input_state.r.state, ev.output_measurement.s.state)

	else
		ev.proba_params.probability = process_probability_partial(ev.interferometer, ev.input_state, ev.output_measurement)
	end

	ev.proba_params.probability = clean_proba(ev.proba_params.probability)

end

"""
	output_mode_occupation(n::Int, m::Int)

Return all possible configurations of `n` photons among `m` modes.
"""
function output_mode_occupation(number_photons, number_modes)

	nlist = collect(1:number_modes)

	for i = 1:number_photons-1
		sublist = []
		for j = 1:length(nlist)
			for k = 1:number_modes
				push!(sublist, vcat(nlist[j], k))
			end
		end
		nlist = sublist
	end

	nlist

end

"""
	check_suppression_law(event)

Check if the event is suppressed according to the [rule](https://arxiv.org/pdf/1002.5038.pdf).
"""
function check_suppression_law(event)

	if mod(sum(event), length(event)) != 0
		return true
	else
		return false
	end

end

### this is an old function that needs to be cleaned

# function partition_probability_distribution_distinguishable(part, U)
# 	"""generates a vector giving the probability to have k photons in
# 	the partition part for the interferometer U by the random walk method
# 	discussed in a 22/02/21 email
# 	part is written like (1,2,4) if it is the output modes 1, 2 and 4"""
#
# 	function proba_photon_in_partition(i, part, U)
# 		"""returns the probability that the photon i ends up in the set of outputs part, for distinguishable particles only"""
#
# 		sum([abs(U[i, j])^2 for j in part]) # note the inversion line column
# 	end
#
#
# 	function walk(probability_vector, walk_number, part, U)
#
# 		new_probability_vector = similar(probability_vector)
# 		n = size(U)[1]
# 		proba_this_walk = proba_photon_in_partition(walk_number, part, U)
#
# 		for i in 1 : n+1
# 			new_probability_vector[i] = (1-proba_this_walk) * probability_vector[i] + proba_this_walk * probability_vector[i != 1 ? i-1 : n+1]
# 		end
#
# 		new_probability_vector
# 	end
#
# 	n = size(U)[1]
# 	probability_vector = zeros(n+1)
# 	probability_vector[1] = 1
#
# 	for walk_number in 1 : n
# 		probability_vector = walk(probability_vector, walk_number, part, U)
# 	end
#
# 	if !(isapprox(sum(probability_vector), 1., rtol = 1e-8))
# 		println("WARNING : probabilites do not sum to one ($(sum(probability_vector)))")
# 	end
#
# 	probability_vector
# end

function is_collisionless(r)
    all(r .<= 1)
end

is_collisionless(r::ModeOccupation) = is_collisionless(r.state)
is_collisionless(i::Input) = is_collisionless(i.r.state)

first_modes_array(n::Int,m::Int) = first_modes(n,m).state


text:
using Pumas, LinearAlgebra, Plots, Query

iv = @model begin
  @param   begin
    tvcl ∈ RealDomain(lower=0)
    tvv ∈ RealDomain(lower=0)
    Ω ∈ PDiagDomain(2)
    σ_prop ∈ RealDomain(lower=0)
  end

  @random begin
    η ~ MvNormal(Ω)
  end

  @pre begin
    CL = tvcl * (Wt/70)^0.75 * exp(η[1])
    V  = tvv * (Wt/70) * exp(η[2])
  end

  @covariates Wt

  @dynamics ImmediateAbsorptionModel

  @derived begin
      cp := @. (Central / V)
      dv ~ @. Normal(cp, sqrt(cp^2*σ_prop))
    end
end

ev = DosageRegimen(2000, time=0)
choose_covariates() = (wt = rand(55:80))
pop =  Population(map(i -> Subject(id=i, evs=ev, cvs = (Wt = choose_covariates(),)),1:24))
param = (
  tvcl = 11.5,
  tvv  = 50,
  Ω = Diagonal([0.04,0.04]),
  σ_prop = 0.001
  )

sd_obstimes = [0, 0.25, 0.5, 0.75, 1, 2, 4, 8,
                    12, 16, 20, 21, 24]
obs = simobs(iv, pop, param, obstimes=sd_obstimes)
iv_sim = DataFrame(obs)|>
          @filter(!(_.time == 0 && _.evid == 0)) |>
          @mutate(cmt = 1) |>
          @mutate(route = "iv") |>
          DataFrame

timeu = u"hr"
concu = u"mg/L"
amtu  = u"mg"

iv_nca = read_nca(iv_sim, id=:id, time=:time, conc=:dv, amt=:amt,
    route=:route,timeu=timeu, concu=concu, amtu=amtu,llq=0.4concu)

NCA.auc(iv_nca)
NCA.auc_extrap_percent(iv_nca)
report = NCAReport(iv_nca)

iv_sim_7percBLQ = iv_sim |>
                  @filter()
sd_obstimes_7percBLQ = [0, 0.25, 0.5, 0.75, 1, 2, 4, 8, 12, 16, 20]
sd_obstimes_18percBLQ = [0, 0.25, 0.5, 0.75, 1, 2, 4, 8, 12, 16, 20, 21, 24]

plot(obs,  yscale = :log10)


text:
x = rand(1:100, 5)
println("Sum of ", x, " = ", sum(x))

A = rand(-100:100, 3, 5)

sum(A, dims=1)

sum(A, dims=2)

sum(A, dims=1)[:]

x = rand(1:100, 5)
display(x)
prod(x)                         # Product of all elements

display(A)
display(maximum(A, dims=1))     # Largest element in each column
display(minimum(A, dims=1))     # Smallest element in each column

x = 1:5
display(cumsum(x))                   # Cumulative sum, that is, entry n is the sum of x_1,...,x_n
display(cumprod(x))                  # Cumulative product

A = reshape(1:6, 2, 3)

cumsum(A, dims=1)                    # Cumulative sum along dimension 1 (that is, column-wise)

cumprod(A, dims=2)                   # Cumulative product along dimension 2

function taylor_cos(x,n)
    factors = [ -x^2 / ((2k-1) * 2k) for k = 1:n ]
    y = 1 + sum(cumprod(factors))
end

println(taylor_cos(10, 50)) # Taylor approximation
println(cos(10))            # true value

xx = -10:0.01:10
yy = cos.(xx)
yy_taylor = [ taylor_cos(x,n) for x in xx, n in 0:5 ]

using PyPlot
plot(xx, yy, linewidth=2, color="k")
plot(xx, yy_taylor)
grid(true)
axis([-10,10,-1.2,1.2]);

x = 1:2:1000                 # Odd numbers
println(503 in x)            # True, 503 is in the list
println(1000 ∈ x)            # False, 1000 is not in the list

println(all(x .< 500))
println(all(x .> 0))

println(any(x .== 503))
println(any(x .== 1000))

idx = findfirst(x .== 503)    # Index in x with the value 503

ind = findall(@. x % 97 == 0)
println("Indices in x with numbers that are multiples of 97: ", ind)

count(@. x % 97 == 0)


text:
using Test
using TwoBit
using BioSequences
using BioGenerics
using FormatSpecimens
using Documenter

@testset "TwoBit" begin
    buffer = IOBuffer()
    writer = TwoBit.Writer(buffer, ["chr1", "chr2"])
    chr1 = dna"ACGTNN"
    chr2 = dna"N" ^ 100 * dna"ACGT" ^ 100 * dna"N" ^ 100
    write(writer, TwoBit.Record("chr1", chr1))
    write(writer, TwoBit.Record("chr2", chr2))
    seekstart(buffer)
    reader = TwoBit.Reader(buffer)
    @test length(reader) == 2
    @test TwoBit.seqnames(reader) == ["chr1", "chr2"]
    @test TwoBit.sequence(reader["chr1"]) == chr1
    @test TwoBit.sequence(reader["chr2"]) == chr2
    @test TwoBit.sequence(reader["chr1"]) == chr1
    @test_throws KeyError reader["chr10"]
    @test_throws BoundsError reader[3]

    function check_2bit_parse(filename)
        stream = open(TwoBit.Reader, filename)
        @test eltype(stream) === TwoBit.Record
        # read from a stream
        for record in stream
            @test hassequence(record) == TwoBit.hassequence(record) == true
            @test TwoBit.sequence(record) == TwoBit.sequence(BioSequences.LongDNA{4}, record)
        end
        close(stream)

        # round trip
        buffer = IOBuffer()
        reader = open(TwoBit.Reader, filename)
        writer = TwoBit.Writer(buffer, TwoBit.seqnames(reader))
        expected_entries = TwoBit.Record[]
        for (name, record) in zip(TwoBit.seqnames(reader), reader)
            write(writer, TwoBit.Record(name, TwoBit.sequence(record), TwoBit.maskedblocks(record)))
            push!(expected_entries, record)
        end

        read_entries = TwoBit.Record[]
        seekstart(buffer)
        for record in TwoBit.Reader(buffer)
            push!(read_entries, record)
        end

        return expected_entries == read_entries
    end

    path = path_of_format("2bit")
    goodfiles = list_valid_specimens("2bit")

    for specimen in goodfiles
        filepath = joinpath(path, filename(specimen))
        @test check_2bit_parse(filepath)
    end

    # Include doctests.
    DocMeta.setdocmeta!(TwoBit, :DocTestSetup, :(using TwoBit); recursive=true)
    doctest(TwoBit; manual = false)
end


text:
using Roots

"""
	EulerExact(ρ1, u1, p1, ρ4, u4, p4, tEnd, n)

Classical Gas Exact Riemann Solver for shock tube problems
This programs is based on the code of
P. Wesseling. PRINCIPLES OF COMPUTATIONAL FLUID DYNAMICS Springer-Verlag, Berlin etc., 2001.
ISBN 3-540-67853-0
NOTE:
A Cavitation Check is incorporated in the code. It further prevents plotting for
 possible but physically unlikely case of expansion shocks.
# Arguments
- `ρ1::Float64`: left density at t=0.
- `u1::Float64`: left velocity at t=0.
- `p1::Float64`: left pressure at t=0.
- `ρ4::Float64`: right density at t=0.
- `u4::Float64`: right velocity at t=0.
- `p4::Float64`: right pressure at t=0.
- `tEnd::Float64`: final solution time.
- `n::Integer64`: the gas Degree of freedom.
Coded by Manuel Diaz, IAM, NTU 03/09/2011.
Migrated by Hongyang Zhou from MATLAB to Julia, 11/05/2019
"""
function EulerExact(ρ1, u1, p1, ρ4, u4, p4, tEnd, n)

   # Gamma values
   γ = (n+2)/n; α = (γ+1)/(γ-1)

   # Assumed structure of exact solution
   #
   #    \         /      |con |       |s|
   #     \   f   /       |tact|       |h|
   # left \  a  /  state |disc| state |o| right
   # state \ n /    2    |cont|   3   |c| state
   #   1    \ /          |tinu|       |k|   4
   #         |           |ity |       | |

   PRL      = p4/p1
   cright   = √(γ*p4/ρ4)
   cleft    = √(γ*p1/ρ1)
   CRL      = cright/cleft
   MachLeft = (u1-u4)/cleft

   # Basic shock tube relation equation (10.51)
   f(P) = (1 + MachLeft*(γ-1)/2 - (γ-1)*CRL*(P-1)/√(2*γ*(γ-1+(γ+1)*P)))^(2*γ/
      (γ-1))/P - PRL

   # solve for P = p34 = p3/p4
   p34 = find_zero(f, 3)

   p3 = p34*p4
   ρ3 = ρ4*(1+α*p34)/(α+p34)
   ρ2 = ρ1*(p34*p4/p1)^(1/γ)
   u2 = u1 -u4 + (2/(γ-1))*cleft*(1-(p34*p4/p1)^((γ-1)/(2*γ)))
   c2 = √(γ*p3/ρ2)
   spos = 0.5 + tEnd*cright*√((γ-1)/(2*γ)+(γ+1)/(2*γ)*p34) + tEnd*u4

   x0 = 0.5
   conpos=x0 + u2*tEnd + tEnd*u4	   # Position of contact discontinuity
   pos1 = x0 + (u1 - cleft)*tEnd	   # Start of expansion fan
   pos2 = x0 + (u2 + u4 - c2)*tEnd	# End of expansion fan

   # Plot structures
   x = 0:0.002:1
   xsize = size(x)
   p = zeros(xsize)
   ux= zeros(xsize)
   ρ = zeros(xsize)
   Mach = zeros(xsize)
   cexact = zeros(xsize)

   for i = 1:length(x)
      if x[i] ≤ pos1
         ρ[i], p[i], ux[i] = ρ1, p1, u1
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      elseif x[i] ≤ pos2
         p[i] = p1*(1+(pos1-x[i])/(cleft*α*tEnd))^(2*γ/(γ-1))
         ρ[i] = ρ1*(1+(pos1-x[i])/(cleft*α*tEnd))^(2/(γ-1))
         ux[i] = u1 + (2/(γ+1))*(x[i]-pos1)/tEnd
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      elseif x[i] ≤ conpos
         ρ[i], p[i], ux[i] = ρ2, p3, u2+u4
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      elseif x[i] ≤ spos
         ρ[i], p[i], ux[i] = ρ3, p3, u2+u4
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      else
         ρ[i], p[i], ux[i] = ρ4, p4, u4
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      end
   end
   entro = @. log(p/ρ^γ) # entropy
   e = @. p/((γ-1)*ρ)	 # internal energy
   t = @. 2/n*e          # temperature

   return x, ρ, ux, p, e, t, Mach, entro
end

# Vector x version
function EulerExact(ρ1, u1, p1, ρ4, u4, p4, tEnd, n, x)

   # Gamma values
   γ=(n+2)/n; α=(γ+1)/(γ-1)

   # Assumed structure of exact solution
   #
   #    \         /      |con |       |s|
   #     \   f   /       |tact|       |h|
   # left \  a  /  state |disc| state |o| right
   # state \ n /    2    |cont|   3   |c| state
   #   1    \ /          |tinu|       |k|   4
   #         |           |ity |       | |

   PRL      = p4/p1
   cright   = √(γ*p4/ρ4)
   cleft    = √(γ*p1/ρ1)
   CRL      = cright/cleft
   MachLeft = (u1-u4)/cleft

   # Basic shock tube relation equation (10.51)
   f(P) = (1 + MachLeft*(γ-1)/2 - (γ-1)*CRL*(P-1)/√(2*γ*(γ-1+(γ+1)*P)))^(2*γ/
      (γ-1))/P - PRL

   # solve for P = p34 = p3/p4
   p34 = find_zero(f, 3)

   p3 = p34*p4
   ρ3 = ρ4*(1+α*p34)/(α+p34)
   ρ2 = ρ1*(p34*p4/p1)^(1/γ)
   u2 = u1 -u4 + (2/(γ-1))*cleft*(1-(p34*p4/p1)^((γ-1)/(2*γ)))
   c2 = √(γ*p3/ρ2)
   spos = 0.5 + tEnd*cright*√((γ-1)/(2*γ)+(γ+1)/(2*γ)*p34) + tEnd*u4

   x0 = 0.5
   conpos=x0 + u2*tEnd + tEnd*u4	   # Position of contact discontinuity
   pos1 = x0 + (u1 - cleft)*tEnd	   # Start of expansion fan
   pos2 = x0 + (u2 + u4 - c2)*tEnd	# End of expansion fan

   # Plot structures
   xsize = size(x)
   p = zeros(xsize)
   ux= zeros(xsize)
   ρ = zeros(xsize)
   Mach = zeros(xsize)
   cexact = zeros(xsize)

   for i = 1:length(x)
      if x[i] ≤ pos1
         ρ[i], p[i], ux[i] = ρ1, p1, u1
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      elseif x[i] ≤ pos2
         p[i] = p1*(1+(pos1-x[i])/(cleft*α*tEnd))^(2*γ/(γ-1))
         ρ[i] = ρ1*(1+(pos1-x[i])/(cleft*α*tEnd))^(2/(γ-1))
         ux[i] = u1 + (2/(γ+1))*(x[i]-pos1)/tEnd
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      elseif x[i] ≤ conpos
         ρ[i], p[i], ux[i] = ρ2, p3, u2+u4
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      elseif x[i] ≤ spos
         ρ[i], p[i], ux[i] = ρ3, p3, u2+u4
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      else
         ρ[i], p[i], ux[i] = ρ4, p4, u4
         cexact[i] = √(γ*p[i]/ρ[i])
         Mach[i] = ux[i]/cexact[i]
      end
   end
   entro = @. log(p/ρ^γ) # entropy
   e = @. p/((γ-1)*ρ)	 # internal energy
   t = @. 2/n*e          # temperature

   return x, ρ, ux, p, e, t, Mach, entro
end


text:
@reexport using Distributions
import Distributions: mean, var, logpdf, pdf, cdf, maximum, minimum, insupport, quantile
export ContinuousUnivariateCopula, marginal_pdf_constants, mvsk_to_absm, conditional_pdf_constants, crossterm_res, gvc_vec_continuous, update_res!
using LinearAlgebra: BlasReal, copytri!

struct ContinuousUnivariateCopula{
    DistT <: ContinuousUnivariateDistribution,
    T     <: Real
    } <: ContinuousUnivariateDistribution
    d  :: DistT
    μ :: T
    σ2 :: T
    c0 :: T
    c1 :: T
    c2 :: T
    c  :: T # normalizing constant
end

"""
    ContinuousUnivariateCopula(d, c0, c1, c2)
The distribution with density `c * P(x = x) * (c0 + c1 * x + c2 * x^2)`, where `f`
is the density of the base distribution `d` and `c` is the normalizing constant.
"""
function ContinuousUnivariateCopula(
    d  :: DistT,
    c0 :: T,
    c1 :: T,
    c2 :: T) where {DistT <: ContinuousUnivariateDistribution, T <: Real}
    μ = mean(d)
    σ2 = var(d)
    c  = inv(c0 + c1 * μ + c2 * (σ2 + abs2(μ)))
    Tc = typeof(c)
    ContinuousUnivariateCopula(d, Tc(μ), Tc(σ2), Tc(c0), Tc(c1), Tc(c2), c)
end

"""
    pdf_constants(Γ::Matrix{<:Real}, res::Vector{<:Real}, i::Int64, dist::ContinuousUnivariateDistribution)
This function will fill out the appropriate constants, c0, c1, c2 for each conditional distribution to form the `ContinuousUnivariateCopula` structure.
"""
function pdf_constants(Γ::Matrix{T}, res::Vector{T}, i::Int64, dist::ContinuousUnivariateDistribution) where T <: Real
    μ = Distributions.mean(dist)
    σ2 = Distributions.var(dist)
    c_0 = μ^2 * inv(σ2)
    c__0 = μ * inv(sqrt(σ2)) * sum(crossterm_res(res, i, Γ))
    c_1 = -2μ * inv(σ2)
    c__1 = inv(sqrt(σ2)) * sum(crossterm_res(res, i, Γ))
    c_2 = inv(σ2)
    storage = zeros(i-1)
    # first multiply Γ[1:i-1, 1:i-1] *  res[1:i-1]
    mul!(storage, Γ[1:i-1, 1:i-1], res[1:i-1])
    # then multiply to get transpose(res[1:i-1]) *  Γ[1:i-1, 1:i-1] *  res[1:i-1]
    c0 = 1 + 0.5 * dot(storage, res[1:i-1]) +  0.5 * tr(Γ[i+1:end, i+1:end]) + 0.5 * Γ[i, i] * c_0  - c__0
    c1 = 0.5 * Γ[i, i] * c_1  + c__1
    c2 = 0.5 * Γ[i, i] * c_2
    ContinuousUnivariateCopula(dist, c0, c1, c2)
end

"""
    pdf_constants(Γ::Matrix{<:Real}, dist::DiscreteUnivariateDistribution)
This function will fill out the appropriate constants, c0, c1, c2 for the univariate marginal distribution to form the `DiscreteUnivariateCopula` structure.
"""
function pdf_constants(Γ::Matrix{T}, i::Int64, dist::ContinuousUnivariateDistribution) where T <: Real
    μ = Distributions.mean(dist)
    σ2 = Distributions.var(dist)
    c_0 = μ^2 * inv(σ2)
    c_1 = -2μ * inv(σ2)
    c_2 = inv(σ2)
    c0 = 1  +  0.5 * (tr(Γ) - Γ[i, i]) + 0.5 * Γ[i, i] * c_0
    c1 = 0.5 * Γ[i, i] * c_1
    c2 = 0.5 * Γ[i, i] * c_2
    ContinuousUnivariateCopula(dist, c0, c1, c2)
end

"""
    pdf_constants(γ::T, dist::ContinuousUnivariateDistribution)
This function will fill out the appropriate constants, c0, c1, c2 for the univariate marginal distribution to form the `ContinuousUnivariateCopula` structure.
When the number of observations in the cluster is 1
"""
function pdf_constants(γ::T, dist::ContinuousUnivariateDistribution) where T <: Real
    μ = Distributions.mean(dist)
    σ2 = Distributions.var(dist)
    c_0 = μ^2 * inv(σ2)
    c_1 = -2μ * inv(σ2)
    c_2 = inv(σ2)
    c0 = 1  + 0.5 * γ * c_0
    c1 = 0.5 * γ * c_1
    c2 = 0.5 * γ * c_2
    ContinuousUnivariateCopula(dist, c0, c1, c2)
end

"""
    crossterm_res(res::Vector{<:Real}, i::Int64, Γ::Matrix{<:Real}; all = false)
This function will compute the crossterm involving the residual values for constants c0 and c1 in the conditional densities.
Default all = false (conditional density):sum_{j = 1}^{i-1} γ_ij * r_j, but if all = true then will output (marginal density): sum_{j = 1}^{i} γ_ij * r_j
"""
function crossterm_res(res::Vector{T}, i::Integer, Γ::Matrix{T}) where {T<: BlasReal}
    results = zeros(i-1)
    if i == 1
        return results
    else
        for j in 1:i - 1
            results[j] = res[j] * Γ[i, j]
        end
    end
    results
 end

"""
    update_res!(Y::Real, res:Real, gc_obs::Union{ContinuousUnivariateCopula{<:ContinuousUnivariateDistribution, <:Real}, DiscreteUnivariateCopula{<:DiscreteUnivariateDistribution, <:Real}})
This function will update the residual value, given the base distributions mean and variance. This step is necessary when constructing the constants, c0, c1, c2, in the conditional density.
"""
 function update_res!(
    Y::T,
    res::T,
    gc_obs::Union{ContinuousUnivariateCopula{D, T}, DiscreteUnivariateCopula{D, T}}) where {T <: BlasReal, D}
    res = (Y - gc_obs.μ) * inv(sqrt(gc_obs.σ2))
 end

minimum(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}) = minimum(d.d)
maximum(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}) = maximum(d.d)
insupport(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}) = insupport(d.d)

"""
    mvsk_to_absm(μ, σ², sk, kt)
Convert mean `μ`, variance `σ²`, skewness `sk`, and kurtosis `kt` to first four
moments about zero. See formula at <https://en.wikipedia.org/wiki/Central_moment#Relation_to_moments_about_the_origin>.
"""
function mvsk_to_absm(μ, σ², sk, kt)
    σ = sqrt(σ²)
    m1 = μ # E[X]
    m2 = σ² + abs2(μ) # Var(X) + E[X]^2 = E[X^2]
    m3 = sk * σ^3 + 3μ * m2 - 2μ^3 # E[X^3] ; where sk = m3 - 3μ(m2) + 2μ^3
    m4 = kt * abs2(σ²) + 4μ * m3 - 6 * abs2(μ) * m2 + 3μ^4 # E[X^4]; where kt = m4 - 4μ (m3) + 6μ^2 (m2) -3μ^4
    m1, m2, m3, m4
end

"""
    mean(d::ContinuousUnivariateCopula)
Theoretical mean under the copula model.
"""
function mean(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula})
    μ, σ², sk, kt = Distributions.mean(d.d), Distributions.var(d.d), Distributions.skewness(d.d), Distributions.kurtosis(d.d, false) # proper kurtosis (un-corrected) when false
    m1, m2, m3, _ = mvsk_to_absm(μ, σ², sk, kt)
    d.c * (d.c0 * m1 + d.c1 * m2 + d.c2 * m3)
end

"""
    var(d::ContinuousUnivariateCopula)
Theoretical variance under the copula model.
"""
function var(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula})
    μ, σ², sk, kt = mean(d.d), var(d.d), skewness(d.d), kurtosis(d.d, false)
    _, m2, m3, m4 = mvsk_to_absm(μ, σ², sk, kt)
    d.c * (d.c0 * m2 + d.c1 * m3 + d.c2 * m4) - abs2(mean(d)) # E[Y_k^2] - E[Y_k]^2
end

"""
    logpdf(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}, x::Real)
Theoretical log pdf under the copula model.
"""
function logpdf(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}, x::Real)
    log(d.c * (d.c0 + d.c1 * x + d.c2 * abs2(x))) + logpdf(d.d, x)
end

"""
    pdf(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}, x::Real)
Theoretical pdf under the copula model.
"""
function pdf(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}, x::Real)
    d.c * pdf(d.d, x) * (d.c0 + d.c1 * x + d.c2 * abs2(x))
end

"""
    cdf(d::Union{ContinuousUnivariateCopula{Normal{<:Real}}, x::Real)
Theoretical cdf for the Normal base distribution derived under the copula model.
"""
function cdf(d::ContinuousUnivariateCopula{Normal{T},T}, x::Real) where T <: Real
    μ, σ    = d.d.μ, d.d.σ
    z       = (x - μ) / σ
    result  = (d.c0 + d.c1 * μ + d.c2 * abs2(μ)) * cdf(Normal(), z)
    result += (d.c1 + 2d.c2 * μ) * σ / sqrt(2π) *
        ifelse(z < 0, - ccdf(Chi(2), -z), cdf(Chi(2), z) - 1)
    result += d.c2 * abs2(σ) / 2 * ifelse(z < 0, ccdf(Chi(3), -z), cdf(Chi(3), z) + 1)
    result *= d.c
end

"""
    cdf(d::Union{ContinuousUnivariateCopula{Gamma{<:Real}}, x::Real)
Theoretical cdf for the Gamma base distribution derived under the copula model.
"""
function cdf(d::ContinuousUnivariateCopula{Gamma{T},T}, x::Real) where T <: Real
    α, θ = params(d.d)
    result  = d.c0 * cdf(d.d, x)
    result += d.c1 * (StatsFuns.gamma(α + 1)/ StatsFuns.gamma(α)) * cdf(Gamma(α + 1, θ), x)
    result += d.c2 * (StatsFuns.gamma(α + 2)/ StatsFuns.gamma(α)) * cdf(Gamma(α + 2, θ), x)
    result *= d.c
end

"""
    cdf(d::Union{ContinuousUnivariateCopula{Exponential{<:Real}}, x::Real)
Theoretical cdf for the Exponential base distribution derived under the copula model.
"""
function cdf(d::ContinuousUnivariateCopula{Exponential{T},T}, x::Real) where T <: Real
    θ = params(d.d)[1]
    normalizing_c1 = θ * (StatsFuns.gamma(2)/ StatsFuns.gamma(1))
    normalizing_c2 = θ^2 * (StatsFuns.gamma(3)/ StatsFuns.gamma(1))
    result  = d.c0 * cdf(d.d, x)
    result += d.c1 * normalizing_c1 * cdf(Gamma(2, θ), x)
    result += d.c2 * normalizing_c2 * cdf(Gamma(3, θ), x)
    result *= d.c
end

"""
    cdf(d::Union{ContinuousUnivariateCopula{Beta{<:Real}}, x::Real)
Theoretical cdf for the Beta base distribution derived under the copula model.
"""
function cdf(d::ContinuousUnivariateCopula{Beta{T},T}, x::Real) where T <: Real
    α, β = params(d.d)
    normalizing_c1 = inv(StatsFuns.gamma(α) * StatsFuns.gamma(α + β + 1)) * (StatsFuns.gamma(α + β) * StatsFuns.gamma(α + 1))
    normalizing_c2 = inv(StatsFuns.gamma(α) * StatsFuns.gamma(α + β + 2)) * (StatsFuns.gamma(α + β) * StatsFuns.gamma(α + 2))
    result  = d.c0 * cdf(d.d, x)
    result += d.c1 * normalizing_c1 * cdf(Beta(α + 1, β), x)
    result += d.c2 * normalizing_c2 * cdf(Beta(α + 2, β), x)
    result *= d.c
end

"""
    quantile(d::Union{ContinuousUnivariateCopula{Normal{<:Real}}, p::Real)
Finds the quantile value for the specified cumulative probability `p`, under the Normal Base distribution of the copula model, using Newtons Method.
"""
# function quantile(d::ContinuousUnivariateCopula{Normal{T}, T}, p::T) where T <: Real
#     Distributions.quantile_newton(d, p, mean(d))
# end

# taken from https://github.com/JuliaStats/Distributions.jl/blob/ec8f7d82c59348443f799e0fb0394441fab31073/src/quantilealgs.jl#L37
# because we don't have a explicit way to compute mode(d), so newton's not guaranteed to converge
function quantile(d::ContinuousUnivariateCopula{Normal{T}, T}, p::T) where T <: Real
    xs = mean(d)
    tol = 1e-12
    x = xs + (p - cdf(d, xs)) / pdf(d, xs)
    if 0 < p < 1
        x0 = T(xs)
        iter = 0
        while abs(x-x0) > max(abs(x),abs(x0)) * tol && iter ≤ 100 # iterate max 100 times
            x0 = x
            x = x0 + (p - cdf(d, x0)) / pdf(d, x0)
            iter += 1
        end
        return x
    elseif p == 0
        return T(minimum(d))
    elseif p == 1
        return T(maximum(d))
    else
        return T(NaN)
    end
end

"""
    quantile(d::Union{ContinuousUnivariateCopula{Gamma{<:Real}}, p::Real)
Finds the quantile value for the specified cumulative probability `p`, under the Gamma Base distribution of the copula model, using Newtons method.
"""
function quantile(d::ContinuousUnivariateCopula{Gamma{T}, T}, p::T) where T <: Real
    α, θ = params(d.d)
    if α ≥ 1
        Distributions.quantile_newton(d, p, ((d.d.α - 1) * d.d.θ))
    else
        error("Gamma has no mode when shape < 1")
    end
end

"""
    quantile(d::Union{ContinuousUnivariateCopula{Exponential{<:Real}}, p::Real)
Finds the quantile value for the specified cumulative probability `p`, under the Normal Base distribution of the copula model, using Newtons method.
"""
function quantile(d::ContinuousUnivariateCopula{Exponential{T}, T}, p::T) where T <: Real
    Distributions.quantile_newton(d, p, 0.0)
end

"""
    quantile(d::Union{ContinuousUnivariateCopula{Exponential{<:Real}}, p::Real)
Finds the quantile value for the specified cumulative probability `p`, under the Beta Base distribution of the copula model, using Bisection Method.
Since the Beta density has finite extrema, [0, 1], we can use the bisection method to implement the inverse cdf sampling rather than Newtons Method (greater stability).
"""
function quantile(d::ContinuousUnivariateCopula{Beta{T}, T}, p::T) where T <: Real
    min, max = extrema(d.d)
    Distributions.quantile_bisect(d, p, min, max, 1.0e-12)
end


text:
using Aeconomica
using Test
using DataFrames
import Dates

# Fetching without setting a key fails
set_apikey("")
@test_throws ErrorException("You have not set an `apikey`. Please run `set_apikey(\"YOURKEY\")` first.") fetch_series("CPI")

set_apikey(ENV["AECONOMICA_TEST_KEY"])

include("test-fetchseries.jl")
include("test-series_as_at.jl")
include("test-dataset.jl")


text:
@testset "BeadDefinition" begin
    perbead_param_keys= (:k,:kθ)
    global_param_keys= (:d₀,)
    function chain_pe(p,b,n,perbead_params,global_params)
        rp= (p .- b)
        rn= (n .- b)
        dp= sqrt(sum(rp.^2))
        dn= sqrt(sum(rn.^2))
        cosθ= sum(rp .* rn)/dp/dn
        perbead_params.k*1//2*((dn-global_params.d₀)^2) + perbead_params.kθ*(cosθ+1)
    end
    beaddef1= BeadDefinition(chain_pe, perbead_param_keys, global_param_keys, Val(3))
    p= [0.0, 0.0, 0.0]
    b= [0.0, 0.0, 1.0]
    n= [0.0, 0.0, 2.0]
    params= (k=3.5, kθ=4.4)
    globalparams= (d₀=1.0, )
    @test 0.0 == beaddef1.chain_pe(p,b,n,params, globalparams)
    fp, fb, fn = beaddef1.chain_force(p,b,n,params, globalparams)
    @test 0.0 == norm(fp)
    @test 0.0 == norm(fb)
    @test 0.0 == norm(fn)
end

@testset "Chain" begin
    perbead_param_keys= (:k,:kθ)
    global_param_keys= (:d₀,)
    function chain_pe(p,b,n,perbead_params,global_params)
        rp= (p .- b)
        rn= (n .- b)
        dp= sqrt(sum(rp.^2))
        dn= sqrt(sum(rn.^2))
        cosθ= sum(rp .* rn)/dp/dn
        perbead_params.k*1//2*((dn-global_params.d₀)^2) + perbead_params.kθ*(cosθ+1)
    end
    beaddef1= BeadDefinition(chain_pe, perbead_param_keys, global_param_keys, Val(3))
    params= fill((k=3.5, kθ=4.4),10)
    params[10]= (k=0.0, kθ= 0.0)# turn off interation to split ends
    params[1]= (k=3.5, kθ= 0.0)
    pos= zeros(10,3)
    pos[:,3] .= range(0.0, step= 1.0, length=10)
    globalparams= (d₀=1.0, )
    Chain(beaddef1,params, globalparams, pos)
end

@testset "SpecificForce" begin
    #set parameters for the force
    params= (k=10.0,)
    #define potential energy, note x[1] is the 1st bead
    pe(x,time,params)= 1//2 * params.k * sum((x[1]).^2)
    #define the list of beads affected by this force
    #The inner tuple is the chain id and bead id in that chain
    #in this case, The force is external,
    #so only one chain/bead id is listed
    interactions = [((1,1),)
                    ((1,2),)]
    #create the force Val(3) specifies 3d beads
    fdef= SpecificForce(pe, params, interactions, Val(3))
    @test [-10,-10,-10] == fdef.force(([1.0,1.0,1.0],), 0.0, params)[1]

    #now an example of a two bead time dependent interaction
    params2= (k₁=10.0, k₂=2.0)
    #note that pe2 is a regular function so it has local scope
    pe2(x,time,params)= 1//2 * sum((params.k₁, params.k₂, params.k₂*time) .* ((x[1] .- x[2]).^2))
    #define the list of beads affected by this force
    #The inner tuple is the chain id and bead id in that chain.
    #In this case, the force is on two beads,
    #so only two chain/bead ids are listed per row.
    #In this example chain 1 bead 1 is bonded to chain 2 bead 1, and
    #chain 1 bead 2 is bonded to chain 1 bead 1
    interactions2 = [((1,1),(2,1))
                    ((1,2),(1,1))]
    #create the force Val(3) specifies 3d beads
    fdef2= SpecificForce(pe2, params2, interactions2, Val(3))
    #second argument is time
    f1,f2 = fdef2.force(([1.0,1.0,1.0],[0.0,0.0,0.0]), 0.0, params2)
    @test f1 == [-10,-2,0]
    @test f2 == .-f1
    #check force at time 1.0
    f1,f2 = fdef2.force(([1.0,1.0,1.0],[0.0,0.0,0.0]), 1.0, params2)
    @test f1 == [-10,-2,-2]
    @test f2 == .-f1

    #now test that SpecificForce will catch basic pe errors
    pe_bad(x,time,params)= 1//2 * params.k * (x[1][4])^2
    @test_throws BoundsError SpecificForce(pe_bad, params, interactions, Val(3))
end

@testset "ChainSystem" begin
    #putting it all together
    perbead_param_keys= (:k,:kθ)
    global_param_keys= (:d₀,)
    function chain_pe(p,b,n,perbead_params,global_params)
        rp= (p .- b)
        rn= (n .- b)
        dp= sqrt(sum(rp.^2))
        dn= sqrt(sum(rn.^2))
        cosθ= sum(rp .* rn)/dp/dn
        perbead_params.k*1//2*((dn-global_params.d₀)^2) + perbead_params.kθ*(cosθ+1)
    end
    beaddef1= BeadDefinition(chain_pe, perbead_param_keys, global_param_keys, Val(3))
    globalparams= (d₀=1.0, )
    system= ChainSystem(0.0,beaddef1,globalparams)
    #create some chain
    params= fill((k=3.5, kθ=4.4),10)
    params[10]= (k=0.0, kθ= 0.0)# turn off interation to split ends
    params[1]= (k=3.5, kθ= 0.0)
    pos= zeros(10,3)
    pos[:,3] .= range(0.0, step= 1.0, length=10)
    chain1= Chain(beaddef1,params, globalparams, pos)
    system= append(system,chain1)
    #create chain2 translated [1 0 0] from chain1
    chain2= Chain(beaddef1,params, globalparams, pos .+ [1 0 0])
    system= append(system,chain2)
    #specific forces
    params= (k=10.0,)
    pe(x,time,params)= 1//2 * params.k * sum((x[1]).^2)
    interactions = [((1,1),)
                    ((1,2),)]
    fdef= SpecificForce(pe, params, interactions, Val(3))
    system= append(system,fdef)

    params2= (k₁=10.0, k₂=2.0)
    pe2(x,time,params)= 1//2 * sum((params.k₁, params.k₂, params.k₂*time) .* ((x[1] .- x[2]).^2))
    interactions2 = [((1,1),(2,1))
                    ((1,2),(1,1))]
    fdef2= SpecificForce(pe2, params2, interactions2, Val(3))
    system= append(system,fdef2)
end


text:
module MyTestPackage
display(terminal)
# Write your package code here.

end

max_stars_repo_name: blixi/MyTestPackage

max_issues_repo_name: blixi/MyTestPackage

max_forks_repo_name: blixi/MyTestPackage