text:
function poisbin_fft!(S::AbstractArray{T,1}, y::AbstractArray{Complex{T},1},
                      z::AbstractArray{Complex{T},1},
                      p::AbstractArray{T,1}) where T <: Real
  n = length(p)
  omega = 2pi/(n+1)
  @inbounds for k in 0:n
    for m in 1:n
      y[k+1] *= p[m] * exp(im * omega * k) + (1-p[m])
    end
  end
  _dft!(z,y)
  S .= real.(z) / (n+1)
end

function poisbin_fft(p::AbstractArray{T,1}) where T <: Real
  n = length(p)
  S = Vector{T}(undef,n+1)
  y = ones(Complex{T}, n+1)
  z = Vector{Complex{T}}(undef,n+1)

  poisbin_fft!(S, y, z, p)
  return S
end

#implementation from Distributions.jl to compare
function poisbin_fft_cf!(S::AbstractArray{T,1}, y::AbstractArray{Complex{T},1},
                         z::AbstractArray{Complex{T},1},
                         p::AbstractArray{T,1}) where T<: Real
  n = length(p)
  y[1] = one(Complex{T}) / (n+1)
  omega = 2 * one(T) / (n+1)
  kmax = ceil(Int, n/2)
  @inbounds for k in 1:kmax
    logz = zero(T)
    argz = zero(T)
    for j in 1:n
      zjl = 1 - p[j] + p[j] * cospi(omega*k) + im * p[j] * sinpi(omega*k)
      logz += log(abs(zjl))
      argz += atan(imag(zjl), real(zjl))
    end
    dl = exp(logz)
    y[k+1] = dl * cos(argz) / (n+1) + im * dl * sin(argz) / (n+1)
    if n + 1 - k > k
      y[n + 1 - k + 1] = conj(y[k+1])
    end
  end
  _dft!(z, y)
  S .= real.(z)
end

function poisbin_fft_cf(p::AbstractArray{T,1}) where T <: Real
  n = length(p)
  S = Vector{T}(undef,n+1)
  y = ones(Complex{T}, n+1)
  z = Vector{Complex{T}}(undef,n+1)

  poisbin_fft_cf!(S, y, z, p)
  return S
end

function poisbin_sum_taub!(S::AbstractArray{T}, p::AbstractArray{T}) where T <: Real
  n = length(p)
  fill!(S,zero(T))
  S[1] = 1-p[1]
  S[2] = p[1]
  @inbounds for col in 2:n
    for r in 1:col
        row = col - r + 1
        S[row+1] = (1-p[col])*S[row+1] + p[col] * S[row]
    end
    S[1] *= 1-p[col]
  end
end

function poisbin_sum_taub(p::AbstractArray{T,1}) where T <: Real
  n = length(p)
  S = Vector{T}(undef,n+1)
  poisbin_sum_taub!(S,p)
  return S
end

function poisbin_sum_taub_log!(S::AbstractArray{T}, p::AbstractArray{T}) where T <: Real
  n = length(p)
  fill!(S,-Inf)
  S[1] = log(1-p[1])
  S[2] = log(p[1])
  @inbounds for col in 2:n
    for r in 1:col
        row = col - r + 1
        Sr = S[row] + log(p[col])
        Sr1 = S[row+1] + log(1-p[col])
        if (Sr1 > Sr) && (Sr1 > zero(T))
          S[row+1] = Sr1 + log1p(exp(Sr - Sr1))
        elseif (Sr >= Sr1) && (Sr > zero(T))
          S[row+1] = Sr + log1p(exp(Sr1-Sr))
        else
          S[row+1] = log(exp(Sr1) + exp(Sr))
        end
    end
    S[1] += log(1-p[col])
  end
end

function poisbin_sum_taub_log(p::AbstractArray{T,1}) where T <: Real
  n = length(p)
  S = Vector{T}(undef,n+1)
  poisbin_sum_taub_log!(S,p)
  return S
end

function poisbin_sum_taub_dervs_2(p::AbstractVector{T}) where T <: Real
  n = length(p)
  S = Vector{T}(undef,n+1)
  H = Array{T,3}(undef,n,n,n+1)
  poisbin_sum_taub!(S, p)
  poisbin_sum_taub_dervs_2!(H, p)
  return S, H
end

function poisbin_sum_taub_dervs_2!(H::AbstractArray{T,3}, p::AbstractVector{T}) where T <: Real
  n = length(p)
  pj=zero(T)
  @inbounds for j in 1:n
    pj = p[j]
    p[j] = zero(T)
    for k in j:n
      pk = p[k]
      p[k] = zero(T)
      @views poisbin_sum_taub!(H[j,k,:], p)
      H[k,j,:] .= H[j,k,:]
      p[k] = pk
    end
    p[j] = pj
  end
end

function poisbin_dc_fft!(S::AbstractArray{T,1}, tempS::AbstractArray{T,2},
                          p::AbstractArray{T,1}, si::AbstractArray{T,1},
                          group_sizes::AbstractArray{D,1},
                          group_start_idx::AbstractArray{D,1}) where {T <: Real, D <: Integer}
  n = length(p)
  M = size(tempS)[2]

  #convolve initial subsets
  @inbounds for g in 1:M
    @views poisbin_sum_taub!(tempS[1:(group_sizes[g]+1),g],
                             p[group_start_idx[g]:(group_start_idx[g]+group_sizes[g]-1)])
    group_sizes[g] += 1
  end

  while M > 1
    next_avail_col = 1
    @inbounds for g in 1:2:M
      m = group_sizes[g] + group_sizes[g+1] - 1
      si .= zero(T)
      @views _filt!(S[1:m], tempS[1:group_sizes[g],g], tempS[1:m,g+1], si[1:(group_sizes[g]-1)])
      @views copyto!(tempS[1:m,next_avail_col], S[1:m])
      group_sizes[next_avail_col] = m
      next_avail_col += 1
    end
    M = div(M,2)
  end
end

function poisbin_dc_fft(p::AbstractArray{T,1}, k::D=2) where {T <: Real, D <: Integer}
  n = length(p)
  k = min(floor(D, log2(n)), k)
  M = 2^k
  L = n/M
  r = rem(n,M) / M

  group_sizes = [fill(floor(D, L), D(M*(1-r))); fill(ceil(D, L), D(M*r))]
  group_start_idx = cumsum(group_sizes) .- (group_sizes .- 1)

  S = Vector{T}(undef,n+1)
  tempS = zeros(T, n+1,M)
  si = zeros(T, n)

  poisbin_dc_fft!(S, tempS, p, si, group_sizes, group_start_idx)
  return S
end

function poisbin_chen!(S::AbstractArray{T,1}, P::AbstractArray{T,1},
                       p::AbstractArray{T,1}) where T <: Real
  n = length(p)
  S[1] = one(T)
  @inbounds for i in 1:n
    S[1] *= (1 - p[i])
    for j in 1:n
      P[j] += (p[i] / (1 - p[i])) ^ j
    end
  end
  @inbounds for i in 2:(n+1)
    k = i-1
    for j in 1:k
      S[i] += (-one(T))^(j-1) * S[k-j+1] * P[j]
    end
    S[i] /= k
  end
end

function poisbin_chen(p::AbstractArray{T,1}) where T <: Real
  n = length(p)
  S = zeros(T, n+1)
  P = zeros(T, n)
  poisbin_chen!(S,P,p)
  return S
end


text:
export optimize_measurement

"""
[`LocalSignaling`](@ref) scenario:

    optimize_measurement(
        scenario :: LocalSignaling,
        game :: BellGame,
        ρ_states :: Vector{<:AbstractMatrix}
    )

Finds the measurement that optimizes the score of the [`BellGame`](@ref) against
the set of quantum states `ρ_states`.
The optimization is performed with the following semi-definite program:

```math
\\begin{aligned}
&\\max_{\\{\\Pi_y\\}} \\sum_{x,y} G_{x,y} \\text{Tr}[\\Pi_y \\rho_x] \\\\
&s.t. \\quad \\sum_y \\Pi_y = \\mathbb{I}, \\quad \\Pi_y \\geq 0
\\end{aligned}
```
"""
function optimize_measurement(
    scenario::LocalSignaling,
    game::BellGame,
    ρ_states::Vector{<:AbstractMatrix};
    atol=1e-7::Float64
) :: Dict
    if scenario.X != length(ρ_states)
        throw(DomainError(scenario, "expected length of `ρ_states` is $(scenario.X)), but got $(length(ρ_states)) instead"))
    elseif size(ρ_states[1]) != (scenario.d,scenario.d)
        throw(DomainError(ρ_states, "dimension of `ρ_states` is not $(scenario.d)"))
    end

    states = map(ρ -> ρ isa State ? ρ : State(ρ, atol=atol), ρ_states)

    norm_game_vector = convert(Vector{Int64}, game)
    norm_bound = norm_game_vector[end]
    norm_game = reshape(norm_game_vector[1:(end-1)], (scenario.Y-1, scenario.X))

    # add povm variables and constraints
    Π_vars = map(i -> HermitianSemidefinite(scenario.d), 1:scenario.Y)
    constraints = (sum(map(Π_y -> real(Π_y), Π_vars)) == Matrix{Float64}(I, scenario.d, scenario.d))
    constraints += (sum(map(Π_y -> imag(Π_y), Π_vars)) == zeros(Float64, scenario.d, scenario.d))

    # sum up the state contibutions for each row
    H_y = map(row_id -> sum(norm_game[row_id,:] .* states), 1:scenario.Y-1)

    # add the objective
    objective = maximize(real(tr(sum(Π_vars[1:end-1] .* H_y))), constraints)

    # optimize model
    solve!(objective, SCS.Optimizer(verbose=0))

    # parse/return results
    score = objective.optval
    violation = score - norm_bound
    Π_opt = _opt_vars_to_povm(map(Π_y -> Π_y.value, Π_vars))

    Dict(
        "violation" => violation,
        "povm" => Π_opt,
        "game" => game,
        "scenario" => scenario,
        "states" => states
    )
end

"""
[`BipartiteNonSignaling`](@ref) scenario:

    optimize_measurement(
        game :: BellGame,
        scenario :: BipartiteNonSignaling,
        ρ_AB :: AbstractMatrix;
        A_POVMs :: Vector{<:AbstractVector{<:AbstractMatrix}},
    ) :: Dict

Find Bob's measurement which optimizes a `BellGame`'s score for the shared quantum
state `ρ_AB` and POVM measurement applied by Alice.
The following semi-definite program optimizes the Bob's POVM:

```math
\\begin{aligned}
    &\\max_{\\{\\Pi_b^y\\}} \\sum_{a,b,x,y} G_{a,b,x,y}\\text{Tr}[(\\Pi_a^x \\otimes \\Pi_b^y)\\rho_{AB}] \\\\
    &s.t. \\quad \\sum_b \\Pi_b^y = \\mathbb{I},\\quad \\Pi_b^y \\geq 0 \\quad \\forall\\; y
\\end{aligned}
```

    optimize_measurement(
        game :: BellGame,
        scenario :: BipartiteNonSignaling,
        ρ_AB :: AbstractMatrix;
        B_POVMs :: Vector{<:AbstractVector{<:AbstractMatrix}},
    ) :: Dict

Find Alice's measurement which optimizes a `BellGame`'s score for the shared quantum
state `ρ_{AB}` and POVM measurement applied by Bob.
The following semi-definite program optimizes the Alice's POVM:

```math
\\begin{aligned}
&\\max_{\\{\\Pi_a^x\\}} \\sum_{a,b,x,y} G_{a,b,x,y}\\text{Tr}[(\\Pi_a^x \\otimes \\Pi_b^y)\\rho_{AB}] \\\\
&s.t. \\quad \\sum_a \\Pi_a^x = \\mathbb{I},\\quad \\Pi_a^x \\geq 0 \\quad \\forall \\;x
\\end{aligned}
```
"""
function optimize_measurement(
    scenario::BipartiteNonSignaling,
    game::BellGame,
    ρ_AB :: AbstractMatrix;
    A_POVMs=Vector{POVM}(undef,0) :: Vector{<:AbstractVector{<:AbstractMatrix}},
    B_POVMs=Vector{POVM}(undef,0) :: Vector{<:AbstractVector{<:AbstractMatrix}},
    atol=1e-7::Float64
) :: Dict
    ρ = ρ_AB isa State ? ρ_AB : State(ρ_AB, atol=atol)

    if length(A_POVMs) > 0
        return _optimize_measurement_B(scenario, game, ρ, A_POVMs, atol=atol)
    elseif length(B_POVMs) > 0
        return _optimize_measurement_A(scenario, game, ρ, B_POVMs, atol=atol)
    else
        throw(DomainError((A_POVMs,B_POVMs), "either `A_POVMs` or `B_POVMs` must be specified."))
    end
end

# """
# Helper function for optimizing Bob's POVM
# """
function _optimize_measurement_B(
    scenario::BipartiteNonSignaling,
    game::BellGame,
    ρ_AB :: AbstractMatrix,
    A_POVMs :: Vector{<:AbstractVector{<:AbstractMatrix}};
    atol=1e-7::Float64
) :: Dict
    if !(scenario.X == length(A_POVMs))
        throw( DomainError(
            (scenario.X, " != length(A_POVMs)"),
            "A distinct POVM is not specified for input `X` of `BipartiteNonSignaling` scenario"
        ))
    elseif !(all(i -> scenario.A == length(A_POVMs[i]), 1:scenario.X ))
        throw( DomainError(
            A_POVMs, "Each POVM must have $(scenario.A) outputs."
        ))
    end

    Π_Ax = map(Π -> Π isa POVM ? Π : POVM(Π, atol=atol), A_POVMs)

    ρ_dim = size(ρ_AB,1)
    Π_A_dim = size(A_POVMs[1][1],1)

    # The dimension of Bob's outputs
    Π_B_dim = (ρ_dim % Π_A_dim == 0) ? Int64(ρ_dim / Π_A_dim) : throw(DomainError(
        ρ_dim, "`size(ρ_Ab)` must equal `size(Π_A)*size(ΠB)`."
    ))

    # Bob's POVM variables to optimize
    B_POVMs = map(y -> map(b -> HermitianSemidefinite(Π_B_dim), 1:scenario.B), 1:scenario.Y)

    # Objective function
    problem = maximize(real(
        sum(x -> sum(y -> sum(a -> sum(b ->
            game[(a-1)*scenario.B+b,(x-1)*scenario.Y+y]*tr(kron(Π_Ax[x][a],B_POVMs[y][b])*ρ_AB.M),
        1:scenario.B), 1:scenario.A), 1:scenario.Y), 1:scenario.X)
    ))

    # Completeness constraints for Bob's POVM
    for y in 1:scenario.Y
        problem.constraints += sum(real.(B_POVMs[y])) == Matrix{Float64}(I, Π_B_dim, Π_B_dim)
        problem.constraints += sum(imag.(B_POVMs[y])) == zeros(Float64, Π_B_dim, Π_B_dim)
    end

    # optimize model
    solve!(problem, SCS.Optimizer(verbose=0))

    # parse/return results
    score = problem.optval
    violation = score - game.β

    Π_B_opt = map(y -> _opt_vars_to_povm(map(Π_b -> Π_b.value, B_POVMs[y])), 1:scenario.Y)

    Dict(
        "violation" => violation,
        "score" => score,
        "game" => game,
        "scenario" => scenario,
        "state" => ρ_AB,
        "A_POVMs" => Π_Ax,
        "B_POVMs" => Π_B_opt
    )
end

# """
# Helper function for optimizing Alice's POVM
# """
function _optimize_measurement_A(
    scenario::BipartiteNonSignaling,
    game::BellGame,
    ρ_AB :: AbstractMatrix,
    B_POVMs :: Vector{<:AbstractVector{<:AbstractMatrix}};
    atol=1e-7::Float64
) :: Dict
    if !(scenario.Y == length(B_POVMs))
        throw( DomainError(
            (scenario.Y, " != length(B_POVMs)"),
            "A distinct POVM is not specified for input `Y` of `BipartiteNonSignaling` scenario"
        ))
    elseif !(all(i -> scenario.B == length(B_POVMs[i]), 1:scenario.Y ))
        throw( DomainError(
            B_POVMs, "Each POVM must have $(scenario.B) outputs."
        ))
    end

    Π_By = map(Π -> Π isa POVM ? Π : POVM(Π, atol=atol), B_POVMs)

    ρ_dim = size(ρ_AB,1)
    Π_B_dim = size(B_POVMs[1][1],1)

    # dimension of Alice's POVM
    Π_A_dim = (ρ_dim % Π_B_dim == 0) ? Int64(ρ_dim / Π_B_dim) : throw(DomainError(
        ρ_dim, "`size(ρ_Ab)` must equal `size(Π_A)*size(ΠB)`."
    ))

    # Alice's POVM variables to optimize
    A_POVMs = map(x -> map(a -> HermitianSemidefinite(Π_A_dim), 1:scenario.A), 1:scenario.X)

    # objective function
    problem = maximize(real(
        sum(x -> sum(y -> sum(a -> sum(b ->
            game[(a-1)*scenario.B+b,(x-1)*scenario.Y+y]*tr(kron(A_POVMs[x][a],Π_By[y][b])*ρ_AB.M),
        1:scenario.B), 1:scenario.A), 1:scenario.Y), 1:scenario.X)
    ))

    # adding completeness constraints
    for x in 1:scenario.X
        problem.constraints += sum(real.(A_POVMs[x])) == Matrix{Float64}(I, Π_A_dim, Π_A_dim)
        problem.constraints += sum(imag.(A_POVMs[x])) == zeros(Float64, Π_A_dim, Π_A_dim)
    end

    # optimize model
    solve!(problem, SCS.Optimizer(verbose=0))

    # parse/return results
    score = problem.optval
    violation = score - game.β

    Π_A_opt = map( x -> _opt_vars_to_povm(map(Π_a -> Π_a.value, A_POVMs[x])), 1:scenario.X)

    Dict(
        "violation" => violation,
        "score" => score,
        "game" => game,
        "scenario" => scenario,
        "state" => ρ_AB,
        "A_POVMs" => Π_A_opt,
        "B_POVMs" => Π_By
    )
end

# """
# Helper function for converting POVM optimization variables to a POVM
# """
function _opt_vars_to_povm(povm::Vector{<:AbstractMatrix}; atol=1) :: POVM
    Π = is_povm(povm,atol=atol) ? POVM(povm,atol=atol) : throw(DomainError(povm, "not a povm"))

    Π
end

ext: jl


text:
using Models
using Test

@testset "Models.jl" begin
    # Write your tests here.
end


text:
# Simple Wright-Fisher model for neutral continuous variable evolution
# Example runs with parallel map over trais.
# Example run:  julia -L NeutralEvolution.jl nrun.jl examples/sn_example1
# Example run:  julia -L NeutralEvolution.jl nrun.jl examples/sn_example1 <seed>
# Example run:  julia -p 4 -L NeutralEvolution.jl nrun.jl examples/sn_example1
#    where seed is a random number seed.
# When run with a specific seed, the output should agree running the full simulation
#    for neutral multiplicative error with a single attribute (with the same seed).
export simple_neutral_type, simple_neutral_init, cummulative_neutral_type, cummulative_neutral_init,
      print_simple_neutral_params, simple_neutral_simulation, ces, accumulate_results, writeheader,
      writerows, writeheader_populations, writerows_populations, coef_var, inverse_fitness

try
  using Dates
  using Printf
catch
end

mutable struct simple_neutral_type
  simtype::Int64
  N::Int64   # popsize
  mutstddev::Float64
  ngens::Int64
  initial_value::Float64
  num_trials::Int64
  record_interval::Int64
  use_population::Bool     # must be true if N>1
  save_populations::Bool    # Save the final population for each trial as a column of the CSV file
  log_error::Bool           # use log normally distributed error to work in log space
  wright_fisher_copy::Bool  # use Wright-Fisher copy to enable population level drift (default true)
  conformist_probability::Float64   # Probability of copying the mean population value rather than Wright-Fisher copy
  neutral::Bool             # if true, no selection
  fit_slope::Float64        # Slope of the inverse fitness function.  0.0 means neutral
  fit_funct::Function       # Fitness function
  average_attr_mean::Float64
  average_attr_median::Float64
  average_attr_coef_var::Float64
  attr_mean_history::Vector{Float64}
  attr_median_history::Vector{Float64}
  attr_coef_var_history::Vector{Float64}
  saved_population::Vector{Float64}
end

function inverse_fitness( a::Float64, initial_value::Float64, fit_slope::Float64=1.0 )
  return 1.0/(abs(fit_slope*a-initial_value)+1.0)
end

function simple_neutral_init( simtype::Int64,N::Int64, mutstddev::Float64, ngens::Int64, initial_value::Float64, num_trials::Int64,
    record_interval::Int64, use_population::Bool=true, save_population::Bool=false,
    log_error::Bool=false, wright_fisher_copy::Bool=true, conformist_probability::Float64=0.0, neutral::Bool=true,
    fit_slope::Float64=1.0, fit_funct::Function=inverse_fitness)
  num_results = Int(ceil(ngens/record_interval))
  saved_pop_empty = try  Array{Float64,1}(undef,0) catch; Array{Float64,1}(0) end
  sn = simple_neutral_type( simtype, N, mutstddev, ngens, initial_value, num_trials, record_interval, use_population, save_population,
    log_error, wright_fisher_copy, conformist_probability, neutral, fit_slope, inverse_fitness,
    0.0, 0.0, 0.0,
    fill(0.0,num_results),
    fill(0.0,num_results),
    fill(0.0,num_results),
    saved_pop_empty)
  #println("log_error: ",sn.log_error)
  sn
end

function print_simple_neutral_params( sn::simple_neutral_type )
  global use_population
  println("N: ",sn.N,"  mutstddev: ",sn.mutstddev," ngens: ",sn.ngens,"  initial_value: ",sn.initial_value,"  num_trials: ",sn.num_trials,
        " record_interval: ",sn.record_interval )
  if isdefined(:use_population)
    print("  use_population: ",sn.use_population)
  end
  if isdefined(:save_populations)
    print("  save_populations: ",sn.save_populations)
  end
  if isdefined(:conformist_probability)
    println("  conformist_probability: ",sn.conformist_probability)
  end
  if isdefined(:fit_slope)
    println("  fit_slope: ",sn.fit_slope)
  end
  println()
  #println("log_error: ",sn.log_error,"  wright_fisher_copy: ",sn.wright_fisher_copy)
end

mutable struct cummulative_neutral_type
  simtype::Int64
  N::Int64  # popsize
  mutstddev::Float64
  ngens::Int64
  initial_value::Float64
  num_trials::Int64
  record_interval::Int64
  use_population::Bool     # must be true if N>1
  save_populations::Bool    # Save the final population for each trial as a column of the CSV file
  log_error::Bool
  wright_fisher_copy::Bool
  conformist_probability::Float64
  neutral::Bool
  fit_slope::Float64
  count_trials::Int64
  gens_recorded::Vector{Int64}
  attr_mean_sum::Vector{Float64}
  attr_mean_sum_sq::Vector{Float64}
  attr_median_sum::Vector{Float64}
  attr_median_sum_sq::Vector{Float64}
  attr_coef_var_sum::Vector{Float64}
  attr_coef_var_sum_sq::Vector{Float64}
  saved_populations::Array{Float64,2}
end

function cummulative_neutral_init( sn::simple_neutral_type )
  num_results = Int(ceil(sn.ngens/sn.record_interval))
  saved_pops_empty = try  Array{Float64,2}(undef,0,0) catch; Array{Float64,2}(0,0) end
  csn = cummulative_neutral_type( sn.simtype, sn.N, sn.mutstddev, sn.ngens, sn.initial_value, sn.num_trials, sn.record_interval,
      sn.use_population, sn.save_populations, sn.log_error, sn.wright_fisher_copy, sn.conformist_probability, sn.neutral, sn.fit_slope,
      0, fill(0,num_results), fill(0.0,num_results), fill(0.0,num_results), fill(0.0,num_results),
      fill(0.0,num_results), fill(0.0,num_results), fill(0.0,num_results), saved_pops_empty)
  if sn.save_populations
    csn.saved_populations = fill(0.0, (sn.N, sn.num_trials) )
  end
  csn
end

function print_cummulative_neutral( csn::cummulative_neutral_type )
  println("N: ",csn.N,"  mutstddev: ",csn.mutstddev," ngens: ",csn.ngens," record_interval: ",csn.record_interval,
      "  count_trials: ",csn.count_trials )
  println("mean sum   : ",csn.attr_mean_sum)
  println("mean sum sq: ",csn.attr_mean_sum_sq)
  println("cvar sum   : ",csn.attr_coef_var_sum)
  println("cvar sum sq: ",csn.attr_coef_var_sum_sq)
end

function simple_neutral_simulation( sn::simple_neutral_type )
  int_burn_in = 0
  cumm_attr_mean = 0.0
  cumm_attr_coef_var = 0.0
  sum_gens = 0
  record_index = 1
  pop = fill( sn.initial_value, sn.N )
  for g = 1:(int_burn_in + sn.ngens )
    i = 1
    #println("before copy g: ",g,"  pop: ",pop)
    new_pop = mutate_pop( sn, pop )
    #println("after mutate g: ",g,"  new_pop: ",new_pop)
    if sn.neutral || sn.fit_slope == 0.0   # neutral, no selection
      if sn.wright_fisher_copy && (sn.neutral || sn.fit_slope == 0.0)
        #pop = [ new_pop[ sn.N>1 ? rand(1:sn.N) : 1 ] for j = 1:sn.N ]
        if sn.N > 1
          r = rand(1:sn.N,sn.N)
        else
          r = fill(1,sn.N)
        end
        #println("r: ",r)
        pop = new_pop[r]
      else
        #pop = deepcopy(new_pop)   # not sure if deepcopy is necessary
        pop = new_pop
      end
    elseif !sn.neutral && sn.fit_slope > 0.0
      pop = propsel_funct( new_pop, sn.fit_funct, sn.initial_value, sn.fit_slope )
    end
    #println("after WF g: ",g,"  pop: ",pop)
    if sn.conformist_probability > 0.0  # unclear if conformist copy should happen after wf copy or propsel
      conformist_copy!( pop, sn )
    end
    #println("after CC g: ",g,"  pop: ",pop)
    cumm_attr_mean += mean(pop)
    cumm_attr_coef_var += coef_var(pop)
    sum_gens += 1
    if (g-1) % sn.record_interval == 0
      sn.attr_mean_history[record_index] = mean(pop)
      sn.attr_median_history[record_index] = median(pop)
      sn.attr_coef_var_history[record_index] = coef_var(pop)
      #println("g: ",g,"  attr_mean: ",sn.attr_mean_history[record_index],"  attr_coef_var: ",sn.attr_coef_var_history[record_index])
      record_index += 1
    end
    if g == (int_burn_in + sn.ngens)
      #println(" pop: ",pop)
      sn.saved_population = pop
    end
    #pop = deepcopy(new_pop)
  end # for g
  @assert sn.ngens == sum_gens
  cumm_attr_mean /= sn.ngens
  cumm_attr_coef_var /= sn.ngens
  #println("cumm_attr_mean: ",cumm_attr_mean)
  #println("cumm_attr_coef_var: ",cumm_attr_coef_var)
  sn.average_attr_mean = cumm_attr_mean
  sn.average_attr_coef_var = cumm_attr_coef_var
  return sn
end

function wright_fisher_copy_function( pop::Vector{Float64}, sn::simple_neutral_type )
  if sn.N > 1
    r = rand(1:sn.N,sn.N)
  else
    r = fill(1,sn.N)
  end
  pop[r]
#=
  if sn.conformist_probability > 0.0
    pop_mean = mean(pop)
    for i = 1:sn.N
      if rand() < sn.conformist_probability
        new_pop[i] = pop_mean
      end
    end
  end
  return new_pop
=#
end

function conformist_copy!( pop::Vector{Float64}, sn::simple_neutral_type )
  for i = 1:sn.N
    if rand() < sn.conformist_probability
      pop[i] = mean(pop)
    end
  end
end

function mutate_pop( sn::simple_neutral_type, pop::Vector{Float64} )
  i = 1
  for p in pop
    rn = randn()
    #println("rn: ",rn)
    mult = 1.0 + sn.mutstddev*rn
    while mult <= 1.e-6
      rn = randn()
      #println("rn: ",rn)
      mult = 1.0 + sn.mutstddev*rn
    end
    #println("mult: ",mult)
    if sn.log_error
      pop[i] = p+log(mult)
    else
      pop[i] = p*mult
      @assert pop[i] > 0.0
    end
    i+= 1
  end
  return pop
end

# Single individual copy error simulation according to Eerkins and Lipo
# Note that s is the size (not the log of size)
function ces( sn::simple_neutral_type )
  cumm_attr_mean = 0.0
  cumm_attr_median = 0.0
  sum_gens = 0
  record_index = 1
  s = sn.initial_value
  for g = 1:sn.ngens
    rn = randn()
    #println("rn: ",rn)
    mult = 1.0+sn.mutstddev*rn
    #println("mult: ",mult)
    while mult < 1.e-6
      rn = randn()
      #println("rn: ",rn)
      mult = 1.0+sn.mutstddev*rn
    end
    if sn.log_error
      s += log(mult)
    else
      s *= mult
      @assert s > 0.0
    end
    #println("s: ",s)
    sum_gens += 1
    cumm_attr_mean += s
    cumm_attr_median += s
    if (g-1) % record_interval == 0
      sn.attr_mean_history[record_index] = s
      sn.attr_median_history[record_index] = s
      sn.attr_coef_var_history[record_index] = 0.0
      record_index += 1
    end
  end # for g
  @assert sn.ngens == sum_gens
  cumm_attr_mean /= sn.ngens
  sn.average_attr_mean = cumm_attr_mean
  cumm_attr_median /= sn.ngens
  sn.average_attr_median = cumm_attr_median
  sn.average_attr_coef_var = 0.0
  return sn
end

# single individual copy error simulation according to Hamilton and Buchanan
# S is log of size
function ces_log( sn::simple_neutral_type )
  cumm_attr_mean = 0.0
  cumm_attr_median = 0.0
  sum_gens = 0
  record_index = 1
  S = sn.initial_value
  for g = 1:sn.ngens
    S += log(1.0+sn.mutstddev*randn())
    sum_gens += 1
    #cumm_attr_mean += S
    cumm_attr_mean += exp(S)
    cumm_attr_median += exp(S)
    if (g-1) % record_interval == 0
      #sn.attr_mean_history[record_index] = S
      sn.attr_mean_history[record_index] = exp(S)
      sn.attr_median_history[record_index] = exp(S)
      sn.attr_coef_var_history[record_index] = 0.0
      record_index += 1
    end
  end # for g
  @assert sn.ngens == sum_gens
  cumm_attr_mean /= sn.ngens
  sn.average_attr_mean = cumm_attr_mean
  cumm_attr_median /= sn.ngens
  sn.average_attr_median = cumm_attr_median
  sn.average_attr_coef_var = 0.0
  return sn
end

function accumulate_results( sn::simple_neutral_type, csn::cummulative_neutral_type )
  num_results = Int(ceil(sn.ngens/sn.record_interval))
  csn.N = sn.N
  csn.wright_fisher_copy=sn.wright_fisher_copy
  csn.conformist_probability=sn.conformist_probability
  record_index = 1
  for g = 1:sn.ngens
    if (g-1) % sn.record_interval == 0
      csn.gens_recorded[record_index] = g
      record_index += 1
    end
  end
  for i = 1:num_results
    csn.attr_mean_sum[i] += sn.attr_mean_history[i]
    csn.attr_mean_sum_sq[i] += sn.attr_mean_history[i]^2
    csn.attr_median_sum[i] += sn.attr_median_history[i]
    csn.attr_median_sum_sq[i] += sn.attr_median_history[i]^2
    csn.attr_coef_var_sum[i] += sn.attr_coef_var_history[i]
    csn.attr_coef_var_sum_sq[i] += sn.attr_coef_var_history[i]^2
  end
  csn.count_trials += 1
  if csn.save_populations
    csn.saved_populations[:,csn.count_trials] = sn.saved_population
  end
end

function writeheader( stream::IO, csn::cummulative_neutral_type )
  param_strings = [
    "# $(string(Dates.today()))",
    "# $((csn.simtype==3) ? "neutral continuous var model simtype==3" : "invalid simtype")",
    #"# N=$(csn.N)",
    "# mutation stddev=$(csn.mutstddev)",
    "# ngens=$(csn.ngens)",
    "# num_trials=$(csn.num_trials)",
    "# initial_value=$(csn.initial_value)",
    "# record_interval=$(csn.record_interval)",
    "# use_population=$(csn.use_population)",
    "# save_populations=$(csn.save_populations)",
    "# log_error=$(csn.log_error)",
    "# fit_slope=$(csn.fit_slope)",
    "# neutral=$(csn.neutral)"
    #"# wright_fisher_copy=$(csn.wright_fisher_copy)"
    #"# conformist_probability=$(csn.conformist_probability)"
    ]
  head_strings = [
    "N",
    "wf_copy",
    "conf_prob",
    "Gen",
    "average mean",
    "stddev of mean",
    "average median",
    "stddev of median",
    "average coef var",
    "stddev of coef var"
    ]
  write(stream, join(param_strings, "\n"), "\n")
  write(stream, join(head_strings, ","), "\n")
end

function writeheader_populations( stream::IO, csn::cummulative_neutral_type )
  param_strings = [
    "# $(string(Dates.today()))",
    "# $((csn.simtype==3) ? "neutral continuous var model simtype==3" : "invalid simtype")",
    "# N=$(csn.N)",
    "# mutation stddev=$(csn.mutstddev)",
    "# ngens=$(csn.ngens)",
    "# num_trials=$(csn.num_trials)",
    "# initial_value=$(csn.initial_value)",
    "# record_interval=$(csn.record_interval)",
    "# use_population=$(csn.use_population)",
    "# save_populations=$(csn.save_populations)",
    "# log_error=$(csn.log_error)",
    #"# wright_fisher_copy=$(csn.wright_fisher_copy)",
    "# conformist_probability=$(csn.conformist_probability)"
    ]
  head_strings = vcat(["wf_copy"],[ "P$i" for i = 1:csn.num_trials ])
  write(stream, join(param_strings, "\n"), "\n")
  write(stream, join(head_strings, ","), "\n")
end

function writerows_populations( stream::IO, csn::cummulative_neutral_type )
  for j = 1:csn.N
    @printf(stream,"%6s,",csn.wright_fisher_copy ? "true" : "false")
    for i = 1:(csn.num_trials-1)
      @printf(stream,"%.4f,",csn.saved_populations[j,i])
    end
    @printf(stream,"%.4f\n",csn.saved_populations[j,csn.num_trials])
  end
end

function writerows( stream::IO, csn::cummulative_neutral_type )
  num_results = Int(ceil(csn.ngens/csn.record_interval))
  means = csn.attr_mean_sum/csn.count_trials
  mean_squared_list = map(x->x^2, csn.attr_mean_sum)
  mean_sq_list = csn.attr_mean_sum .* csn.attr_mean_sum
  @assert reduce(&, mean_squared_list .== mean_sq_list )
  mean_stdvs = sqrt.(csn.attr_mean_sum_sq/(csn.count_trials-1) - mean_squared_list/csn.count_trials/(csn.count_trials-1))
  medians = csn.attr_median_sum/csn.count_trials
  median_squared_list = map(x->x^2, csn.attr_median_sum)
  median_sq_list = csn.attr_median_sum .* csn.attr_median_sum
  @assert reduce(&, median_squared_list .== median_sq_list )
  median_stdvs = sqrt.(csn.attr_median_sum_sq/(csn.count_trials-1) - median_squared_list/csn.count_trials/(csn.count_trials-1))
  coef_vars = csn.attr_coef_var_sum/csn.count_trials
  if csn.N > 1
    coef_var_squared_list = map(x->x^2, csn.attr_coef_var_sum)
    #println("coef_var_squared_list: ",coef_var_squared_list)
    coef_var_sq_list = csn.attr_coef_var_sum .* csn.attr_coef_var_sum
    #println("coef_var_squ_list: ",coef_var_sq_list)
    @assert reduce(&, coef_var_squared_list .== coef_var_sq_list )
    cv_stdevs = sqrt.(csn.attr_coef_var_sum_sq/(csn.count_trials-1) - coef_var_squared_list/csn.count_trials/(csn.count_trials-1))
  else
    cv_stdevs = fill(0.0,num_results)
  end
  for i = 1:num_results
    values = Number[
      csn.N,
      csn.wright_fisher_copy,
      csn.conformist_probability,
      csn.gens_recorded[i],
      means[i],
      mean_stdvs[i],
      medians[i],
      median_stdvs[i],
      coef_vars[i],
      cv_stdevs[i]
      ]
    write(stream, join(values,","), "\n")
  end
end

function coef_var( lst )
  return std(lst)/mean(lst)
end


text:
module AED

using Distributed, SharedArrays, DataFrames, CSV
using Statistics, StatsBase, Distributions
using LinearAlgebra, Optim, Roots, ForwardDiff

using SplineHazard
import SplineHazard.Spline.Dierckx.Spline1D

abstract type AbstractData end
abstract type AbstractParam end
abstract type AbstractFixedParam end
abstract type AbstractFit end
abstract type AbstractPrior end

struct Model{T <: AbstractParam, S <: AbstractFixedParam}
    param::T
    fparam::S
    generate_data::Function
end

include("SA.jl")
include("Dirac.jl")
include("MCMCUtil.jl")

import .SA: grid_search
import .MCMCUtil: MH_single, AM_RR_single, updateMean, updateVar

include("models/BiPHSpline.jl")
include("models/WeibullPH.jl")
include("models/DoubleWeibull.jl")
include("models/BiPHWeibull.jl")

include("forward_sim.jl")
include("objective.jl")

include("gsd_trial.jl")
include("aed_trial.jl")

include("sim.jl")

"""
A Julia package for Bayesian adaptive enrichment designs with subgroup selection
"""
AED

end


text:
using SpecialFunctions, Distributions

a,b = 0.2, 0.7
x = 0.75

betaAB1 = beta(a,b)
betaAB2 = (gamma(a)gamma(b))/gamma(a+b)
betaAB3 = (factorial(a-1)factorial(b-1))/factorial(a+b-1)
betaPDFAB1 = pdf(Beta(a,b),x)
betaPDFAB2 = (1/beta(a,b))*x^(a-1) * (1-x)^(b-1)

println("beta($a,$b)    = $betaAB1,\t$betaAB2,\t$betaAB3 ")
println("betaPDF($a,$b) at $x = $betaPDFAB1,\t$betaPDFAB2")

max_stars_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_issues_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_forks_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook


text:
# This file inludes fractal integration tools.

export integrate

"""
    $SIGNATURES

Calculates numerical integration of `pts` using the fractal integration method `method` that may be

* [`Interp1D`](@ref)
* [`HInterp1D`](@ref)
* [`Interp2D`](@ref)
* [`HInterp2D`](@ref)

"""
integrate(pts::AbstractVector{<:AbstractVector{<:Real}}, method::AbstractInterp) =
    integrate(map(pnt -> Point(pnt...), pts), method)

function integrate(pts::PointVector, method::AbstractInterp)
    tess = tessellate(pts, method)
    transforms = gettransforms(pts, method)
    val = evaluate(transforms, method, pts)
end

# Interp1D
function evaluate(transforms::AbstractVector{<:Transformation}, method::Interp1D, pts::PointVector{2})
    x  = getindex.(pts, 1)
    (a11, a21, _, a22), (b1, b2) = extract(transforms)

    # J   = abs.(a11)  # Taking abs is not true, |⋅| denotes determinant, not abs
    J   = a11
    K1  = (x[end]^2 - x[1]^2) / 2
    K2  = x[end] - x[1]

    num = sum(J .* (a21 * K1 + b2 * K2))
    denum = 1 - sum(J .* a22)
    num / denum
end

# HInterp1D
function evaluate(transforms::AbstractVector{<:Transformation}, method::HInterp1D, pts::PointVector{3})
    x = getindex.(pts, 1)
    (a11, a21, a31, _, a22, a32, _, a23, a33), (b1, b2, b3) = extract(transforms)

    # J   = abs.(a11) # Taking abs is not true, |⋅| denotes determinant, not abs
    J   = abs.(a11)
    K1  = (x[end]^2 - x[1]^2) / 2
    K2  = (x[end] - x[1])

    W11 = sum(a22 .* J);    W12 = sum(a23 .* J);    Λ1  = sum((a21 * K1 + b2 * K2) .* J);
    W21 = sum(a32 .* J);    W22 = sum(a33 .* J);    Λ2  = sum((a31 * K1 + b3 * K2) .* J);

    A = [1 - W11    -W12;
         -W21       1 - W22]
    b = [Λ1, Λ2]
    I = A \ b
    I[1]    # Hidden integration value I[2] is not returned.
end

# Interp2D
function evaluate(transforms::AbstractVector{<:Transformation}, method::Interp2D, pts::PointVector{3})
    (x1, y1), (x2, y2), (x3, y3) = getboundary(pts, method)

    k11 = x2 - x1;      k12 = x3 - x1;      l1 = x1
    k21 = y2 - y1;      k22 = y3 - y1;      l2 = y1

    # JT = abs(k11 * k22 - k21 * k12) # Taking abs is not true, |⋅| denotes determinant, not abs
    JT = k11 * k22 - k21 * k12
    Δ1 = k11 + k12 + 3 * l1
    Δ2 = k21 + k22 + 3 * l2

    (a11, a21, a31, a12, a22, a32, _, _, a33), (b1, b2, b3) = extract(transforms)
    # JL = abs.(a11 .* a22 - a21 .* a12) # Taking abs is not true, |⋅| denotes determinant, not abs
    JL = a11 .* a22 - a21 .* a12

    num = JT / 6 * sum(JL .* (a31 * Δ1 + a32 * Δ2 + 3 * b3))
    denum = 1 - sum(JL .* a33)
    num / denum
end

# HInterp2D
function evaluate(transforms::AbstractVector{<:Transformation}, method::HInterp2D, pts::PointVector{4})
    (x1, y1), (x2, y2), (x3, y3) = getboundary(pts, method)
    k11 = x2 - x1;      k12 = x3 - x1;      l1 = x1
    k21 = y2 - y1;      k22 = y3 - y1;      l2 = y1

    # JT = abs(k11 * k22 - k21 * k12) # Taking abs is not true, |⋅| denotes determinant, not abs
    JT = k11 * k22 - k21 * k12
    Δ1 = k11 + k12 + 3 * l1
    Δ2 = k21 + k22 + 3 * l2

    (a11, a21, a31, a41, a12, a22, a32, a42, _, _, a33, a43, _, _, a34, a44), (b1, b2, b3, b4) = extract(transforms)
    # JL = abs.(a11 .* a22 - a21 .* a12) # Taking abs is not true, |⋅| denotes determinant, not abs
    JL = a11 .* a22 - a21 .* a12

    W11 = sum(a33 .* JL);    W12 = sum(a34 .* JL);    Λ1  = JT / 6 * sum((a31 * Δ1 + a32 * Δ2 + 3 * b3) .* JL);
    W21 = sum(a43 .* JL);    W22 = sum(a44 .* JL);    Λ2  = JT / 6 * sum((a41 * Δ1 + a42 * Δ2 + 3 * b4) .* JL);

    A = [1 - W11    -W12;
         -W12       1 - W22]
    b = [Λ1, Λ2]
    I = A \ b
    I[1]    # Hidden integration value I[2] is not returned.
end

function extract(transforms::AbstractVector{<:Transformation})
    As = getfield.(transforms, :A)
    bs = getfield.(transforms, :b)
    A = map(idx -> getindex.(As, idx...), [(i, j) for i in 1 : size(As[1], 1), j in 1 : size(As[1], 2)])
    b = map(idx -> getindex.(bs, idx), 1 : length(bs[1]))
    (A, b)
end


text:
"""
Bridges `CP.Strictly{CP.DoublyLexicographicallyGreaterThan}`
to `CP.Strictly{CP.LexicographicallyGreaterThan}`.
"""
struct DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T} <: MOIBC.AbstractBridge
    con::MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.Strictly{CP.LexicographicallyGreaterThan, T}} # Columns are lexicographically sorted.
    con_transposed::MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.Strictly{CP.LexicographicallyGreaterThan, T}} # Rows are lexicographically sorted.
end

function MOIBC.bridge_constraint(
    ::Type{DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T}},
    model,
    f::MOI.VectorOfVariables,
    s::CP.Strictly{CP.DoublyLexicographicallyGreaterThan, T},
) where {T}
    return MOIBC.bridge_constraint(
        DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T},
        model,
        MOI.VectorAffineFunction{T}(f),
        s,
    )
end

function MOIBC.bridge_constraint(
    ::Type{DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T}},
    model,
    f::MOI.VectorAffineFunction{T},
    s::CP.Strictly{CP.DoublyLexicographicallyGreaterThan, T},
) where {T}
    # Constraint on columns.
    con = MOI.add_constraint(
        model,
        f,
        CP.Strictly{CP.LexicographicallyGreaterThan, T}(CP.LexicographicallyGreaterThan(s.set.row_dim, s.set.column_dim))
    )

    # Constraint on rows.
    f_scalars = MOIU.scalarize(f)
    f_matrix = reshape(f_scalars, s.set.row_dim, s.set.column_dim)
    f_transposed = MOIU.vectorize(vec(f_matrix)) # vec() does the transposition.

    con_transposed = MOI.add_constraint(
        model,
        f_transposed,
        CP.Strictly{CP.LexicographicallyGreaterThan, T}(CP.LexicographicallyGreaterThan(s.set.column_dim, s.set.row_dim))
    )

    return DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge(con, con_transposed)
end

function MOI.supports_constraint(
    ::Type{DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T}},
    ::Union{Type{MOI.VectorOfVariables}, Type{MOI.VectorAffineFunction{T}}},
    ::Type{CP.Strictly{CP.DoublyLexicographicallyGreaterThan, T}},
) where {T}
    return true
end

function MOIB.added_constrained_variable_types(::Type{DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T}}) where {T}
    return Tuple{Type}[]
end

function MOIB.added_constraint_types(::Type{DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T}}) where {T}
    return [
        (MOI.VectorAffineFunction{T}, CP.Strictly{CP.LexicographicallyGreaterThan, T}),
    ]
end

function MOI.get(
    ::DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.VectorAffineFunction{T}, CP.Strictly{CP.LexicographicallyGreaterThan, T},
    },
) where {T}
    return 2
end

function MOI.get(
    b::DoublyStrictlyLexicographicallyLessThan2StrictlyLexicographicallyGreaterThanBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{T}, CP.Strictly{CP.LexicographicallyGreaterThan, T},
    },
) where {T}
    return [b.con, b.con_transposed]
end


text:
let
    using Distributions
    d = Categorical([0.4, 0.5, 0.1])
    c = collect(weighted_iterator(d))
    @test c == [1=>0.4, 2=>0.5, 3=>0.1]
end


text:
using Test, LinearAlgebra, Random
using Gaius
using StructArrays

@testset "ComplexFloat64 Multiplication" begin
    for sz ∈ [10, 50, 100, 200, 400, 1000]
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m)" begin
            C1 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            C2 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            A  = StructArray{ComplexF64}((randn(n, k), randn(n, k)))
            B  = StructArray{ComplexF64}((randn(k, m), randn(k, m)))

            @test blocked_mul!(C1, A, B) ≈ mul!(C2, A, B)
            @test blocked_mul(A, B) ≈ C1
        end
    end
    for sz ∈ [10, 50, 100, 200, 400, 1000]
        n, k, m = shuffle([sz + rand(-5:5), sz + rand(-5:5), 10])
        @testset "($n × $k) × ($k × $m)" begin
            C1 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            C2 = StructArray{ComplexF64}((zeros(n, m), zeros(n, m)))
            A  = StructArray{ComplexF64}((randn(n, k), randn(n, k)))
            B  = StructArray{ComplexF64}((randn(k, m), randn(k, m)))

            @test blocked_mul!(C1, A, B) ≈ mul!(C2, A, B)
            @test blocked_mul(A, B) ≈ C1
        end
    end
end

@testset "Float64 Multiplication" begin
    for sz ∈ [10, 50, 100, 200, 400, 1000]
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m)" begin
            C1 = zeros(n, m)
            C2 = zeros(n, m)
            A  = randn(n, k)
            B  = randn(k, m)
            At = copy(A')
            Bt = copy(B')

            @test blocked_mul!(C1, A, B) ≈ mul!(C2, A, B)
            @test blocked_mul(A, B) ≈ C1

            fill!(C1, NaN); @test blocked_mul!(C1, At', B)   ≈ C2
            fill!(C1, NaN); @test blocked_mul!(C1, A,   Bt') ≈ C2
            fill!(C1, NaN); @test blocked_mul!(C1, At', Bt') ≈ C2
        end
    end
    for sz ∈ [10, 50, 100, 200, 400, 1000]
        n, k, m = shuffle([sz + rand(-5:5), sz + rand(-5:5), 10])
        @testset "($n × $k) × ($k × $m)" begin
            C1 = zeros(n, m)
            C2 = zeros(n, m)
            A  = randn(n, k)
            B  = randn(k, m)
            At = copy(A')
            Bt = copy(B')

            @test blocked_mul!(C1, A, B) ≈ mul!(C2, A, B)
            @test blocked_mul(A, B) ≈ C1
            fill!(C1, NaN); @test blocked_mul!(C1, At', B)   ≈ C2
            fill!(C1, NaN); @test blocked_mul!(C1, A,   Bt') ≈ C2
            fill!(C1, NaN); @test blocked_mul!(C1, At', Bt') ≈ C2
        end
    end
end

@testset "Float32 Multiplication" begin
    for sz ∈ [10, 50, 100, 200, 400, 1000]
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m)" begin
            C1 = zeros(Float32, n, m)
            C2 = zeros(Float32, n, m)
            A  = randn(Float32, n, k)
            B  = randn(Float32, k, m)
            At = copy(A')
            Bt = copy(B')

            @test blocked_mul!(C1, A, B) ≈ mul!(C2, A, B)
            @test blocked_mul(A, B) ≈ C1
            fill!(C1, NaN32); @test blocked_mul!(C1, At', B)   ≈ C2
            fill!(C1, NaN32); @test blocked_mul!(C1, A,   Bt') ≈ C2
            fill!(C1, NaN32); @test blocked_mul!(C1, At', Bt') ≈ C2
        end
    end
end

@testset "Int64 Multiplication" begin
    for sz ∈ [10, 50, 100, 200, 400, 1000]
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m)" begin
            C1 = zeros(Int, n, m)
            C2 = zeros(Int, n, m)
            A  = rand(-100:100, n, k)
            B  = rand(-100:100, k, m)
            At = copy(A')
            Bt = copy(B')

            @test blocked_mul!(C1, A, B) ≈ mul!(C2, A, B)
            @test blocked_mul(A, B) ≈ C1
            fill!(C1, typemax(Int64)); @test blocked_mul!(C1, At', B)   ≈ C2
            fill!(C1, typemax(Int64)); @test blocked_mul!(C1, A,   Bt') ≈ C2
            fill!(C1, typemax(Int64)); @test blocked_mul!(C1, At', Bt') ≈ C2
        end
    end
end

@testset "Int32 Multiplication" begin
    for sz ∈ [10, 50, 100, 200, 400, 1000]
        n, k, m = (sz .+ rand(-5:5, 3))
        @testset "($n × $k) × ($k × $m)" begin
            C1 = zeros(Int32, n, m)
            C2 = zeros(Int32, n, m)
            A  = rand(Int32.(-100:100), n, k)
            B  = rand(Int32.(-100:100), k, m)
            At = copy(A')
            Bt = copy(B')

            @test blocked_mul!(C1, A, B) ≈ mul!(C2, A, B)
            @test blocked_mul(A, B) ≈ C1
            fill!(C1, typemax(Int32)); @test blocked_mul!(C1, At', B)   ≈ C2
            fill!(C1, typemax(Int32)); @test blocked_mul!(C1, A,   Bt') ≈ C2
            fill!(C1, typemax(Int32)); @test blocked_mul!(C1, At', Bt') ≈ C2
        end
    end
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct PolicySku <: SwaggerModel
    name::Any # spec type: Union{ Nothing, String } # spec name: name
    tier::Any # spec type: Union{ Nothing, String } # spec name: tier

    function PolicySku(;name=nothing, tier=nothing)
        o = new()
        validate_property(PolicySku, Symbol("name"), name)
        setfield!(o, Symbol("name"), name)
        validate_property(PolicySku, Symbol("tier"), tier)
        setfield!(o, Symbol("tier"), tier)
        o
    end
end # type PolicySku

const _property_map_PolicySku = Dict{Symbol,Symbol}(Symbol("name")=>Symbol("name"), Symbol("tier")=>Symbol("tier"))
const _property_types_PolicySku = Dict{Symbol,String}(Symbol("name")=>"String", Symbol("tier")=>"String")
Base.propertynames(::Type{ PolicySku }) = collect(keys(_property_map_PolicySku))
Swagger.property_type(::Type{ PolicySku }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_PolicySku[name]))}
Swagger.field_name(::Type{ PolicySku }, property_name::Symbol) =  _property_map_PolicySku[property_name]

function check_required(o::PolicySku)
    (getproperty(o, Symbol("name")) === nothing) && (return false)
    true
end

function validate_property(::Type{ PolicySku }, name::Symbol, val)
end


text:
"""
    ActiveSet{AT, R, IT}

Represents an active set of extreme vertices collected in a FW algorithm,
along with their coefficients `(λ_i, a_i)`.
`R` is the type of the `λ_i`, `AT` is the type of the atoms `a_i`.
The iterate `x = ∑λ_i a_i` is stored in x with type `IT`.
"""
struct ActiveSet{AT,R,IT} <: AbstractVector{Tuple{R,AT}}
    weights::Vector{R}
    atoms::Vector{AT}
    x::IT
end

ActiveSet{AT,R}() where {AT,R} = ActiveSet{AT,R,Vector{float(eltype(AT))}}([], [])

ActiveSet{AT}() where {AT} = ActiveSet{AT,Float64,Vector{float(eltype(AT))}}()

function ActiveSet(tuple_values::AbstractVector{Tuple{R,AT}}) where {AT,R}
    n = length(tuple_values)
    weights = Vector{R}(undef, n)
    atoms = Vector{AT}(undef, n)
    @inbounds for idx in 1:n
        weights[idx] = tuple_values[idx][1]
        atoms[idx] = tuple_values[idx][2]
    end
    x = float.(similar(atoms[1]))
    as = ActiveSet{AT,R,typeof(x)}(weights, atoms, x)
    compute_active_set_iterate!(as)
    return as
end

function ActiveSet{AT,R}(tuple_values::AbstractVector{<:Tuple{<:Number,<:Any}}) where {AT,R}
    n = length(tuple_values)
    weights = Vector{R}(undef, n)
    atoms = Vector{AT}(undef, n)
    x = float.(similar(tuple_values[1][2]))
    x .= 0
    @inbounds for idx in 1:n
        weights[idx] = tuple_values[idx][1]
        atoms[idx] = tuple_values[idx][2]
        x .+= weights[idx] * atoms[idx]
    end
    return ActiveSet{AT,R,typeof(x)}(weights, atoms, x)
end

Base.getindex(as::ActiveSet, i) = (as.weights[i], as.atoms[i])
Base.size(as::ActiveSet) = size(as.weights)

function Base.push!(as::ActiveSet, (λ, a))
    push!(as.weights, λ)
    push!(as.atoms, a)
    return as
end

# these two functions do not update the active set iterate

function Base.deleteat!(as::ActiveSet, idx)
    deleteat!(as.weights, idx)
    deleteat!(as.atoms, idx)
    return as
end

function Base.setindex!(as::ActiveSet, tup::Tuple, idx)
    as.weights[idx] = tup[1]
    as.atoms[idx] = tup[2]
    return tup
end

function Base.empty!(as::ActiveSet)
    empty!(as.atoms)
    empty!(as.weights)
    as.x .= 0
    return as
end

function Base.isempty(as::ActiveSet)
    return isempty(as.atoms)
end

"""
Copies an active set, the weight and atom vectors and the iterate.
Individual atoms are not copied.
"""
function Base.copy(as::ActiveSet{AT, R, IT}) where {AT, R, IT}
    return ActiveSet{AT, R, IT}(
        copy(as.weights),
        copy(as.atoms),
        copy(as.x),
    )
end

"""
    active_set_update!(active_set::ActiveSet, lambda, atom)

Adds the atom to the active set with weight lambda or adds lambda to existing atom.
"""
function active_set_update!(active_set::ActiveSet, lambda, atom, renorm=true, idx=nothing)
    # rescale active set
    active_set.weights .*= (1 - lambda)
    # add value for new atom
    if idx === nothing
        idx = find_atom(active_set, atom)
    end
    updating = false
    if idx > 0
        @inbounds active_set.weights[idx] = active_set.weights[idx] + lambda
        updating = true
    else
        push!(active_set, (lambda, atom))
    end
    if renorm
        active_set_cleanup!(active_set, update=false)
        active_set_renormalize!(active_set)
    end
    @. active_set.x = active_set.x * (1 - lambda) + lambda * atom
    return active_set
end

function active_set_update_iterate_pairwise!(active_set::ActiveSet, lambda, fw_atom, away_atom)
    @. active_set.x += lambda * fw_atom - lambda * away_atom
    return active_set
end

function active_set_validate(active_set::ActiveSet)
    return sum(active_set.weights) ≈ 1.0
end

function active_set_renormalize!(active_set::ActiveSet)
    renorm = sum(active_set.weights)
    active_set.weights ./= renorm
    return active_set
end

function weight_from_atom(active_set::ActiveSet, atom)
    idx = find_atom(active_set, atom)
    if idx > 0
        return active_set.weights[idx]
    else
        return nothing
    end
end

"""
    get_active_set_iterate(active_set)

Return the current iterate corresponding. Does not recompute it.
"""
function get_active_set_iterate(active_set)
    return active_set.x
end

function compute_active_set_iterate!(active_set)
    active_set.x .= 0
    for (λi, ai) in active_set
        active_set.x .+= λi * ai
    end
    return active_set.x
end

function active_set_cleanup!(active_set; weight_purge_threshold=1e-12, update=true)
    filter!(e -> e[1] > weight_purge_threshold, active_set)
    if update
        compute_active_set_iterate!(active_set)
    end
    return nothing
end

function find_atom(active_set::ActiveSet, atom)
    @inbounds for idx in eachindex(active_set)
        if _unsafe_equal(active_set.atoms[idx], atom)
            return idx
        end
    end
    return -1
end

"""
    active_set_argmin(active_set::ActiveSet, direction)

Computes the linear minimizer in the direction on the active set.
Returns `(λ_i, a_i, i)`
"""
function active_set_argmin(active_set::ActiveSet, direction)
    val = dot(active_set.atoms[1], direction)
    idx = 1
    temp = 0
    for i in 2:length(active_set)
        temp = fast_dot(active_set.atoms[i], direction)
        if temp < val
            val = temp
            idx = i
        end
    end
    # return lambda, vertex, index
    return (active_set[idx]..., idx)
end

"""
    active_set_argminmax(active_set::ActiveSet, direction)

Computes the linear minimizer in the direction on the active set.
Returns `(λ_min, a_min, i_min, val_min, λ_max, a_max, i_max, val_max, val_max-val_min ≥ Φ)`
"""
function active_set_argminmax(active_set::ActiveSet, direction; Φ=0.5)
    val = Inf
    valM = -Inf
    idx = -1
    idxM = -1
    for i in eachindex(active_set)
        temp_val = fast_dot(active_set.atoms[i], direction)
        if temp_val < val
            val = temp_val
            idx = i
        end
        if valM < temp_val
            valM = temp_val
            idxM = i
        end
    end
    return (active_set[idx]..., idx, val, active_set[idxM]..., idxM, valM, valM - val ≥ Φ)
end

"""
    active_set_initialize!(as, v)

Resets the active set structure to a single vertex `v` with unit weight.
"""
function active_set_initialize!(as::ActiveSet{AT,R}, v) where {AT,R}
    empty!(as)
    push!(as, (one(R), v))
    @. as.x = one(R) * v
    return as
end


text:
# Categorical
function randCategorical(p::T...) where T<:Real
    draw = rand(T)
    cp = zero(T)
    i = zero(Int)
    while cp < draw && i < length(p)
        cp += p[i+=1]
    end
    max(i, 1)
end
randCategorical(p::T...) where T<:Integer = randCategorical(Float64.(p)...)

function logpdfCategorical(i::I, p_::Real...) where I<:Integer
    p = homogenize(p_)
    T = eltype(eltype(p))
    ifelse(1 ≤ i ≤ length(p), @inbounds log(p[i]), -T(Inf))
end

function _logpdfCategorical(i::I, p_::Real...) where I<:Integer
    p = homogenize(p_)
    T = eltype(eltype(p))
    ifelse(1 ≤ i ≤ length(p), @inbounds CUDAnative.log(p[i]), -T(Inf))
end

# Poisson
# TODO eliminar dependencia de StatsFuns:
randPoisson(λ::Real) = convert(Int, poisinvcdf(λ, rand()))

function logpdfPoisson(i::Integer, λ::T) where T<:Real
    x = convert(eltype(T), i)
    iszero(λ) && return ifelse(iszero(x), zero(T), -T(Inf))
    x * log(λ) - λ - lgamma(x + one(T))
end

function _logpdfPoisson(i::Integer, λ::T) where T<:Real
    x = convert(eltype(T), i)
    iszero(λ) && return ifelse(iszero(x), zero(T), -T(Inf))
    x * CUDAnative.log(λ) - λ - CUDAnative.lgamma(x + one(T))
end

# Uniform
randUniform(a::T, b::T) where T<:Real = a + (b - a)rand(T)
randUniform(a::T, b::T) where T<:Integer = Float64(a) + Float64(b - a)rand()

function logpdfUniform(x_::Real, a_::Real, b_::Real)
    x, a, b = promote(x_, a_, b_)
    ifelse(a ≤ x ≤ b, -log(b - a), -typeof(x)(Inf))
end

function _logpdfUniform(x_::Real, a_::Real, b_::Real)
    x, a, b = promote(x_, a_, b_)
    ifelse(a ≤ x ≤ b, -CUDAnative.log(b - a), -typeof(x)(Inf))
end

# Normal
randNormal(μ::T, σ::T) where T<:Real = μ + σ * randn(T)
randNormal(μ::T, σ::T) where T<:Integer = Float64(μ) + Float64(σ) * randn()

function logpdfNormal(x_::Real, μ_::Real, σ_::Real)
    x, μ, σ = promote(x_, μ_, σ_)
    T = eltype(x)
    iszero(σ) && return ifelse(x == μ, T(Inf), -T(Inf))
    -(((x - μ) / σ)^2 + T(log2π))/2 - log(σ)
end

function _logpdfNormal(x_::Real, μ_::Real, σ_::Real)
    x, μ, σ = promote(x_, μ_, σ_)
    T = eltype(x)
    iszero(σ) && return ifelse(x == μ, T(Inf), -T(Inf))
    -(((x - μ) / σ)^2 + T(log2π))/2 - CUDAnative.log(σ)
end

# Exponential
randExponential(λ::Real) = λ * randexp(eltype(λ))
randExponential(λ::Integer) = λ * randexp()

function logpdfExponential(x_::Real, λ_::Real)
    x, λ = promote(x_, λ_)
    ifelse(x < zero(x), -eltype(x)(Inf), log(λ) - λ * x)
end

function _logpdfExponential(x_::Real, λ_::Real)
    x, λ = promote(x_, λ_)
    ifelse(x < zero(x), -eltype(x)(Inf), CUDAnative.log(λ) - λ * x)
end

# Gamma
function _MarsagliaTsang2000(α::T) where T<:Real
    d = α - one(T)/3
    c = one(T) / sqrt(9d)
    while true
        x = randn(T)
        v = (one(T) + c*x)^3
        while v < zero(T)
            x = randn(T)
            v = (one(T) + c*x)^3
        end
        u = rand(T)
        # u < one(F) - F(0.0331)x^4 && return d*v
        log(u) < x^2/2 + d*(one(T) - v + log(v)) && return d*v
    end
end

function randGamma(α::T, θ::T) where T<:Real
    if α < one(T) # use the γ(1+α)*U^(1/α) trick from Marsaglia and Tsang (2000)
        x = θ * _MarsagliaTsang2000(α + one(T)) # Gamma(α + 1, θ)
        e = randexp(T)
        return x * exp(-e / α)
    elseif α == one(T)
        return θ * randexp(eltype(T))
    else
        return θ * _MarsagliaTsang2000(α)
    end
end
randGamma(α::T, θ::T) where T<:Integer = randGamma(Float64(α), Float64(θ))

function logpdfGamma(x_::Real, α_::Real, θ_::Real)
    x, α, θ = promote(x_, α_, θ_)
    T = eltype(x)
    -lgamma(α) - α*log(θ) + (α-one(T))log(x) - x/θ
end

function _logpdfGamma(x_::Real, α_::Real, θ_::Real)
    x, α, θ = promote(x_, α_, θ_)
    T = eltype(x)
    -CUDAnative.lgamma(α) - α*CUDAnative.log(θ) + (α-one(T))CUDAnative.log(x) - x/θ
end


text:
# example
# numcars = 5;
# x = collect(range(0,10,length = 100));
# ys = randboxcars(x, numcars)
# boxwidthstd = 0.4 unit of x's length
function randboxcars(x::Array{T,1}, numcars::Int; boxwidthstd = 0.4) where {T} #
	halfwidthstd = boxwidthstd*0.5;
	boxcentr = rand(x, numcars);
	halfboxwidth = halfwidthstd .* (x[end]-x[1]) .* abs.(randn(numcars));
	ys = fill(0, length(x));
	for i = 1:numcars
		lb, centr, rb = get_rect(x,boxcentr,halfboxwidth,i);
		ys[lb .< x .< rb] .= ys[lb .< x .< rb] .+ 1;
	end
	return ys
end
# number of rows of ys has to be the same as length(x).
function randboxcars(x::Array{T,1}, ys::Array{N,2}; boxwidthstd = 0.4) where {T,N}
	halfwidthstd = boxwidthstd*0.5;
	numcars = size(ys, 2);
	boxcentr = rand(x, numcars);
	halfboxwidth = halfwidthstd .* (x[end]-x[1]) .* abs.(randn(numcars));
	ys = fill(0, (length(x), numcars));
	for i = 1:numcars
		lb, centr, rb = get_rect(x,boxcentr,halfboxwidth,i);
		ys[lb .< x .< rb,i] .= 1;
	end
	return ys
end

function get_rect(x,boxcentr,halfboxwidth,i)
	centr = findlast(x .<= boxcentr[i]);
	rb = x[centr] + halfboxwidth[i]; # right bound
	lb = x[centr] - halfboxwidth[i]; # left bound
	return lb, centr, rb
end


text:
module SQLLoader

using Octo.Repo: ExecuteResult

# db_connect
function db_connect(; kwargs...)
end

# db_disconnect
function db_disconnect()
end

# query
function query(sql::String)
end

function query(prepared::String, vals::Vector)
end

# execute
function execute(sql::String)::ExecuteResult
    ExecuteResult()
end

function execute(prepared::String, vals::Vector)::ExecuteResult
    ExecuteResult()
end

function execute(prepared::String, nts::Vector{<:NamedTuple})::ExecuteResult
    ExecuteResult()
end

end # module Octo.Backends.SQLLoader


text:
#Java types
#JDocument = @jimport edu.stanford.nlp.simple.Document
#JSentence = @jimport edu.stanford.nlp.simple.Sentence
#JStanfordCoreNLP = @jimport edu.stanford.nlp.pipeline.StanfordCoreNLP
#JAnnotation = @jimport edu.stanford.nlp.pipeline.Annotation
#JCoreLabel = @jimport edu.stanford.nlp.ling.CoreLabel
#JArrayList = @jimport java.util.ArrayList
#JTree = @jimport edu.stanford.nlp.trees.Tree
JSemanticGraph = @jimport edu.stanford.nlp.semgraph.SemanticGraph #?
#JHasWord = @jimport edu.stanford.nlp.ling.HasWord
#JTaggedWord = @jimport edu.stanford.nlp.ling.TaggedWord
#JMaxentTagger = @jimport edu.stanford.nlp.tagger.maxent.MaxentTagger
#JStringReader = @jimport java.io.StringReader
#JDependencyParser = @jimport edu.stanford.nlp.parser.nndep.DependencyParser
#JDocumentPreprocessor = @jimport edu.stanford.nlp.process.DocumentPreprocessor
#JGrammaticalStructure = @jimport edu.stanford.nlp.trees.GrammaticalStructure

#Java classes
##Inner
#JSentencesAnnotationClass = classforname("edu.stanford.nlp.ling.CoreAnnotations\$SentencesAnnotation")
#JTokensAnnotationClass = classforname("edu.stanford.nlp.ling.CoreAnnotations\$TokensAnnotation")
JTextAnnotationClass = classforname("edu.stanford.nlp.ling.CoreAnnotations\$TextAnnotation")  #?
JLemmaAnnotationClass = classforname("edu.stanford.nlp.ling.CoreAnnotations\$LemmaAnnotation") #?
JPartOfSpeechAnnotationClass = classforname("edu.stanford.nlp.ling.CoreAnnotations\$PartOfSpeechAnnotation") #?
JNamedEntityTagAnnotationClass = classforname("edu.stanford.nlp.ling.CoreAnnotations\$NamedEntityTagAnnotation") #?
#JTreeAnnotationClass = classforname("edu.stanford.nlp.trees.TreeCoreAnnotations\$TreeAnnotation")
JTreeClass = classforname("edu.stanford.nlp.trees.Tree") #?
#JCollapsedCCProcessedDependenciesAnnotationClass = classforname("edu.stanford.nlp.semgraph.SemanticGraphCoreAnnotations\$CollapsedCCProcessedDependenciesAnnotation")
JCorefChainAnnotationClass = classforname("edu.stanford.nlp.dcoref.CorefCoreAnnotations\$CorefChainAnnotation") #?

##Outer
JMaxentTaggerClass = classforname("edu.stanford.nlp.tagger.maxent.MaxentTagger") #?
JGrammaticalStructureClass = classforname("edu.stanford.nlp.trees.GrammaticalStructure") #?
JDependencyParserClass = classforname("edu.stanford.nlp.parser.nndep.DependencyParser") #?


text:
using RootTips
RootTips.main()


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

# when constructing a without generator infos like `SampledDensity(density, samples)`:
struct UnknownSampleGenerator<: AbstractSampleGenerator end
getalgorithm(sg::UnknownSampleGenerator) = nothing

# for samplers without specific infos, e.g. current ImportanceSamplers:
struct GenericSampleGenerator{A <: AbstractSamplingAlgorithm} <: AbstractSampleGenerator
    algorithm::A
end
getalgorithm(sg::GenericSampleGenerator) = sg.algorithm

"""
    struct IIDSampling <: AbstractSamplingAlgorithm

Sample via `Random.rand`. Only supported for posteriors of type
`Distributions.MultivariateDistribution` and `BAT.DistLikeDensity`.

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct IIDSampling <: AbstractSamplingAlgorithm
    nsamples::Int = 10^5
end
export IIDSampling

function bat_sample_impl(rng::AbstractRNG, target::AnyIIDSampleable, algorithm::IIDSampling)
    shaped_density = convert(DistLikeDensity, target)
    density = unshaped(shaped_density)
    shape = varshape(shaped_density)
    n = algorithm.nsamples

    # ToDo: Parallelize, using hierarchical RNG (separate RNG for each sample)
    v = nestedview(rand(rng, sampler(density), n))
    logd = map(logdensityof(density), v)

    weight = fill(_default_int_WT(1), length(eachindex(logd)))
    info = fill(nothing, length(eachindex(logd)))
    aux = fill(nothing, length(eachindex(logd)))

    unshaped_samples = DensitySampleVector((v, logd, weight, info, aux))

    samples = shape.(unshaped_samples)
    (result = samples,)
end

"""
    struct RandResampling <: AbstractSamplingAlgorithm

Resamples from a given set of samples.

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct RandResampling <: AbstractSamplingAlgorithm
    nsamples::Int = 10^5
end
export RandResampling

function bat_sample_impl(rng::AbstractRNG, posterior::DensitySampleVector, algorithm::RandResampling)
    n = algorithm.nsamples
    orig_idxs = eachindex(posterior)
    weights = FrequencyWeights(float(posterior.weight))
    resampled_idxs = sample(orig_idxs, weights, n, replace=true, ordered=false)

    samples = posterior[resampled_idxs]
    samples.weight .= 1

    (result = samples,)
end

"""
    struct OrderedResampling <: AbstractSamplingAlgorithm

Efficiently resamples from a given series of samples, keeping the order of
samples.

Can be used to efficiently convert weighted samples into samples with unity
weights.

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct OrderedResampling <: AbstractSamplingAlgorithm
    nsamples::Int = 10^5
end
export OrderedResampling

function bat_sample_impl(rng::AbstractRNG, samples::DensitySampleVector, algorithm::OrderedResampling)
    @assert axes(samples) == axes(samples.weight)
    W = samples.weight
    idxs = eachindex(samples)

    n = algorithm.nsamples
    resampled_idxs = Vector{Int}()
    sizehint!(resampled_idxs, n)

    p_factor = n / sum(W)

    for i in eachindex(W)
        w_eff_0 = p_factor * W[i]
        w_eff::typeof(w_eff_0) = w_eff_0
        while w_eff > 0
            rand(rng) < w_eff && push!(resampled_idxs, i)
            w_eff = w_eff - 1
        end
    end

    new_samples = samples[resampled_idxs]
    new_samples.weight .= 1

    (result = new_samples,)
end


text:
struct ModelClosure{M,A} <: AbstractModel{A,B}
    model :: M
    argvals :: A
end

function Base.show(io::IO, cm::ModelClosure)
    println(io, "ModelClosure given")
    println(io, "    arguments    ", keys(argvals(cm)))
    println(io, "    observations ", keys(observations(cm)))
    println(io, model(cm))
end

export argvals
argvals(c::ModelClosure) = c.argvals
argvals(c::ModelPosterior) = c.argvals
argvals(m::AbstractModel) = NamedTuple()

export observations
observations(c::ModelClosure) = c.obs

export observed
function observed(cm::ModelClosure{M,A,O}) where {M,A,O}
    keys(schema(Obs))
end

model(c::ModelClosure) = c.model

ModelClosure(m::AbstractModel) = ModelClosure(m,NamedTuple(), NamedTuple())
model(::Type{<:ModelPosterior{M,A,O}}) where {M,A,O} = type2model(Model{M,A,O})

ModelPosterior(m::Model) = ModelPosterior(m,NamedTuple(), NamedTuple())

(m::AbstractModel)(nt::NamedTuple) = ModelClosure(m)(nt)

(cm::ModelClosure)(nt::NamedTuple) = ModelClosure(cm.model, merge(cm.argvals, nt), cm.obs)

(m::AbstractModel)(;argvals...)= m((;argvals...))

(m::AbstractModel)(args...) = m(NamedTuple{Tuple(m.args)}(args...))
# (m::Model)(args...) = m(NamedTuple{Tuple(m.args)}(args))
(m::Model)(args...) = m(argstype(m)(args))

import Base

Base.:|(m::AbstractModel, nt::NamedTuple) = ModelClosure(m) | nt

Base.:|(cm::ModelClosure, nt::NamedTuple) = ModelClosure(cm.model, cm.argvals, merge(cm.obs, nt))


text:
# functions for passing out info from simulations, similar to the info return from openai gym
# maintained by @zsunberg

"""
    a, ai = action_info(policy, x)

Return a tuple containing the action determined by policy 'p' at state or belief 'x' and information (usually a `NamedTuple`, `Dict` or `nothing`) from the calculation of that action.

By default, returns `nothing` as info.
"""
function action_info(p::Policy, x)
    return action(p, x), nothing
end

"""
    policy, si = solve_info(solver, problem)

Return a tuple containing the policy determined by a solver and information (usually a `NamedTuple`, `Dict` or `nothing`) from the calculation of that policy.

By default, returns `nothing` as info.
"""
function solve_info(s::Solver, problem::Union{POMDP,MDP})
    return solve(s, problem), nothing
end

"""
    bp, i = update_info(updater, b, a, o)

Return a tuple containing the new belief and information (usually a `NamedTuple`, `Dict` or `nothing`) from the belief update.

By default, returns `nothing` as info.
"""
function update_info(up::Updater, b, a, o)
    return update(up, b, a, o), nothing
end

# once POMDPs v0.8 is released, this should be a jldoctest
"""
    add_infonode(ddn::DDNStructure)

Create a new DDNStructure object with a new node labeled :info for returning miscellaneous informationabout a simulation step.

Typically, the object in info is associative (i.e. a `Dict` or `NamedTuple`) with keys corresponding to different pieces of information.

# Example (using POMDPs v0.8)

```julia
using POMDPs, POMDPModelTools, POMDPPolicies, POMDPSimulators, Random

struct MyMDP <: MDP{Int, Int} end

# add the info node to the DDN
POMDPs.DDNStructure(::Type{MyMDP}) = mdp_ddn() |> add_infonode

# the dynamics involve two random numbers - here we record the values for each in info
function POMDPs.gen(m::MyMDP, s, a, rng)
    r1 = rand(rng)
    r2 = randn(rng)
    return (sp=s+a+r1+r2, r=s^2, info=(r1=r1, r2=r2))
end

m = MyMDP()
@show nodenames(DDNStructure(m))
p = FunctionPolicy(s->1)
for (s,info) in stepthrough(m, p, 1, "s,info", max_steps=5, rng=MersenneTwister(2))
    @show s
    @show info
end
```
"""
function add_infonode(ddn) # for DDNStructure, but it is not declared in v0.7.3, so there is not annotation
    add_node(ddn, :info, ConstantDDNNode(nothing), nodenames(ddn))
end

function add_infonode(ddn::POMDPs.DDNStructureV7{nodenames}) where nodenames
    return POMDPs.DDNStructureV7{(nodenames..., :info)}()
end

###############################################################
# Note all generate functions will be deprecated in POMDPs v0.8
###############################################################

if DDNStructure(MDP) isa POMDPs.DDNStructureV7
    """
    Return a tuple containing the next state and reward and information (usually a `NamedTuple`, `Dict` or `nothing`) from that step.

    By default, returns `nothing` as info.
    """
    function generate_sri(p::MDP, s, a, rng::AbstractRNG)
        return generate_sr(p, s, a, rng)..., nothing
    end

    """
    Return a tuple containing the next state, observation, and reward and information (usually a `NamedTuple`, `Dict` or `nothing`) from that step.

    By default, returns `nothing` as info.
    """
    function generate_sori(p::POMDP, s, a, rng::AbstractRNG)
        return generate_sor(p, s, a, rng)..., nothing
    end

    POMDPs.gen(::DDNOut{(:sp,:o,:r,:i)}, m, s, a, rng) = generate_sori(m, s, a, rng)
    POMDPs.gen(::DDNOut{(:sp,:o,:r,:info)}, m, s, a, rng) = generate_sori(m, s, a, rng)
    POMDPs.gen(::DDNOut{(:sp,:r,:i)}, m, s, a, rng) = generate_sri(m, s, a, rng)
    POMDPs.gen(::DDNOut{(:sp,:r,:info)}, m, s, a, rng) = generate_sri(m, s, a, rng)
else
    @deprecate generate_sri(args...) gen(DDNOut(:sp,:r,:info), args...)
    @deprecate generate_sori(args...) gen(DDNOut(:sp,:o,:r,:info), args...)
end


text:
export
    LinearRegressionModel,
    fit!,
    predict

mutable struct LinearRegressionModel
    df::DataFrame
    label::Symbol
    features::Vector{Symbol}
    argv::Vector{Float64}
end

function LinearRegressionModel(df::DataFrame, label::Symbol, features::Vector{Symbol})
    return LinearRegressionModel(df, label, features, rand(length(features)+1))
end

function LinearRegressionModel(df::DataFrame, label::Symbol, features::Symbol)
    return LinearRegressionModel(df, label, [features], rand(2))
end

function g(model::LinearRegressionModel, row_n::Int64)
    xs = collect(model.df[row_n, model.features])
    pushfirst!(xs, 1)
    y = sum(model.argv .* xs)

    return y
end

function loss(model::LinearRegressionModel)
    n = nrow(model.df)
    l = 0
    for i in 1:n
        l += (g(model, i) - model.df[i, model.label])^2
    end

    return l/n
end

function fit!(model::LinearRegressionModel; lr=1e-4, atol::Float64=1e-6, show=false)
    while (l = loss(model)) > atol
        show && println("Loss: $l")
        dl_da = 0
        dl_db = zeros(length(model.features))
        for i in 1:nrow(model.df)
            # intersection
            dl_da += g(model, i) - model.df[i, model.label]

            # sloaps
            for (j, f) in enumerate(model.features)
                dl_db[j] += (g(model, i) - model.df[i, model.label]) * model.df[i, f]
            end
        end

        # intersection
        model.argv[1] -= lr * 2 * dl_da
        # slopes
        for j in 2:length(model.argv)
            model.argv[j] -= lr * 2 * dl_db[j-1]
        end
    end
end

function predict(model::LinearRegressionModel, xs::Vector{<:Real})
    pushfirst!(xs, 1)
    y = sum(model.argv .* xs)

    return y
end

function predict(model::LinearRegressionModel, x::Real)
    return predict(model, [x])
end


text:
#
# Functions to extract and evaluate variables from a simulation solution
#

function get_variable(sim, sol, var_name::String, inputs)
    input_parameter_order = isnothing(inputs) ? nothing : collect(keys(inputs))
    p = isnothing(inputs) ? nothing : collect(values(inputs))

    # Generate the function using PyBaMM
    pybamm = pyimport("pybamm")
    var_str = pybamm.get_julia_function(
        sim.built_model.variables[var_name],
        funcname="var_func",
        input_parameter_order=input_parameter_order
    )
    var_func! = runtime_eval(Meta.parse(var_str))

    # Evaluate and fill in the vector
    # 0D variables only for now
    var = Array{Float64}(undef, length(sol.t))
    out = [0.0]
    for i in 1:length(sol.t)
        # Updating 'out' in-place
        var_func!(out, sol.u[i], p, sol.t[i])
        var[i] = out[1]
    end
    return var
end

get_variable(sim, sol, var_name::String) = get_variable(sim, sol, var_name, nothing)

function get_l2loss_function(sim, var_name, inputs, data)
    input_parameter_order = isnothing(inputs) ? nothing : collect(keys(inputs))
    p = isnothing(inputs) ? nothing : collect(values(inputs))

    # Generate the function using PyBaMM
    pybamm = pyimport("pybamm")
    var_str = pybamm.get_julia_function(
        sim.built_model.variables[var_name],
        funcname="var_func",
        input_parameter_order=input_parameter_order
    )
    var_func! = runtime_eval(Meta.parse(var_str))

    # Evaluate L2 loss
    out = [0.0]

    function loss(sol)
        p = sol.prob.p
        sumsq = 0.0
        for i in 1:length(sol.t)
            # Updating 'out' in-place
            var_func!(out, sol.u[i], p, sol.t[i])
            sumsq += (data[i] - out[1])^2
            # @show i,data[i], out[1]
        end
        sumsq
    end
end


text:
struct ID
    id::BigInt
end

Base.:(==)(a::ID,b::ID) = a.id==b.id
Base.hash(a::ID, h::UInt) = Base.hash(a.id,hash(:ID, h))

a,b,c = ID(1), ID(2), ID(3)

set = Set(ID[a,b,c])

@show a in set

max_stars_repo_name: akels/DemocracyCredit

max_issues_repo_name: akels/DemocracyCredit

max_forks_repo_name: akels/DemocracyCredit


text:
module SequentialWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const AbstractCovarianceMatrixProvider = @jimport org.orekit.estimation.sequential.AbstractCovarianceMatrixProvider
const AbstractIntegratedPropagator = @jimport org.orekit.propagation.integration.AbstractIntegratedPropagator
const Class = @jimport java.lang.Class
const ConstantProcessNoise = @jimport org.orekit.estimation.sequential.ConstantProcessNoise
const CovarianceMatrixProvider = @jimport org.orekit.estimation.sequential.CovarianceMatrixProvider
const DSSTKalmanModel = @jimport org.orekit.estimation.sequential.DSSTKalmanModel
const DSSTPropagator = @jimport org.orekit.propagation.semianalytical.dsst.DSSTPropagator
const EstimatedMeasurement = @jimport org.orekit.estimation.measurements.EstimatedMeasurement
const IntegratedPropagatorBuilder = @jimport org.orekit.propagation.conversion.IntegratedPropagatorBuilder
const Iterable = @jimport java.lang.Iterable
const JString = @jimport java.lang.String
const KalmanEstimation = @jimport org.orekit.estimation.sequential.KalmanEstimation
const KalmanEstimator = @jimport org.orekit.estimation.sequential.KalmanEstimator
const KalmanEstimatorBuilder = @jimport org.orekit.estimation.sequential.KalmanEstimatorBuilder
const KalmanModel = @jimport org.orekit.estimation.sequential.KalmanModel
const KalmanODModel = @jimport org.orekit.estimation.sequential.KalmanODModel
const KalmanObserver = @jimport org.orekit.estimation.sequential.KalmanObserver
const LOFType = @jimport org.orekit.frames.LOFType
const List = @jimport java.util.List
const MatrixDecomposer = @jimport org.hipparchus.linear.MatrixDecomposer
const Measurement = @jimport org.hipparchus.filtering.kalman.Measurement
const MeasurementDecorator = @jimport org.orekit.estimation.sequential.MeasurementDecorator
const NonLinearEvolution = @jimport org.hipparchus.filtering.kalman.extended.NonLinearEvolution
const NonLinearProcess = @jimport org.hipparchus.filtering.kalman.extended.NonLinearProcess
const NumericalPropagator = @jimport org.orekit.propagation.numerical.NumericalPropagator
const Object = @jimport java.lang.Object
const ObservedMeasurement = @jimport org.orekit.estimation.measurements.ObservedMeasurement
const ParameterDriversList = @jimport org.orekit.utils.ParameterDriversList
const PositionAngle = @jimport org.orekit.orbits.PositionAngle
const ProcessEstimate = @jimport org.hipparchus.filtering.kalman.ProcessEstimate
const PropagationType = @jimport org.orekit.propagation.PropagationType
const RealMatrix = @jimport org.hipparchus.linear.RealMatrix
const RealVector = @jimport org.hipparchus.linear.RealVector
const SpacecraftState = @jimport org.orekit.propagation.SpacecraftState
const UnivariateFunction = @jimport org.hipparchus.analysis.UnivariateFunction
const UnivariateProcessNoise = @jimport org.orekit.estimation.sequential.UnivariateProcessNoise

export AbstractCovarianceMatrixProvider
export ConstantProcessNoise
export CovarianceMatrixProvider
export DSSTKalmanModel
export KalmanEstimation
export KalmanEstimator
export KalmanEstimatorBuilder
export KalmanModel
export KalmanODModel
export KalmanObserver
export MeasurementDecorator
export UnivariateProcessNoise
export add_propagation_configuration
export build
export decomposer
export equals
export estimated_measurements_parameters
export estimation_step
export evaluation_performed
export finalize_estimation
export get_class
export get_corrected_measurement
export get_corrected_spacecraft_states
export get_covariance
export get_current_date
export get_current_measurement_number
export get_estimate
export get_estimated_measurements_parameters
export get_estimated_orbital_parameters
export get_estimated_propagation_parameters
export get_estimated_propagators
export get_evolution
export get_initial_covariance_matrix
export get_innovation
export get_lof_cartesian_orbital_parameters_evolution
export get_lof_type
export get_observed_measurement
export get_orbital_parameters_drivers
export get_physical_estimated_covariance_matrix
export get_physical_estimated_state
export get_physical_innovation_covariance_matrix
export get_physical_kalman_gain
export get_physical_measurement_jacobian
export get_physical_state_transition_matrix
export get_position_angle
export get_predicted_measurement
export get_predicted_spacecraft_states
export get_process_noise_matrix
export get_propagation_parameters_drivers
export get_propagation_parameters_evolution
export get_time
export get_value
export hash_code
export notify
export notify_all
export process_measurements
export set_observer
export to_string
export wait

include("abstract_covariance_matrix_provider.jl")
include("constant_process_noise.jl")
include("covariance_matrix_provider.jl")
include("dsst_kalman_model.jl")
include("kalman_estimation.jl")
include("kalman_estimator.jl")
include("kalman_estimator_builder.jl")
include("kalman_model.jl")
include("kalman_observer.jl")
include("kalman_od_model.jl")
include("measurement_decorator.jl")
include("univariate_process_noise.jl")

end


text:
;; ---------------------------------------------------------------------
;;	text_manipulate.jl
;;
;;					Mar/02/2011
;;
;; --------------------------------------------------------------------
(require 'tables)
;; --------------------------------------------------------------------
(defun dict_display_proc_single (key unit_bb)
(format t (concat key "\t"
	(table-ref unit_bb "name") "\t"
	(table-ref unit_bb "population") "\t"
	(table-ref unit_bb "date_mod")))
)

;; --------------------------------------------------------------------
(define (text_read_line_proc out_str delimit dict_aa)
	(string-match delimit out_str)
	(setq id (substring out_str 0 (match-start 0) ))
	(setq str_bb (substring out_str (+ (match-start 0) 1)))
	(string-match delimit str_bb)
	(setq name (substring str_bb 0 (match-start 0)))
	(setq str_cc (substring str_bb (+ (match-start 0) 1)))
	(string-match delimit str_cc)
	(setq population (substring str_cc 0 (match-start 0)))
;	(setq date_mod (substring str_cc (+ (match-start 0) 1)))
	(setq str_dd (substring str_cc (+ (match-start 0) 1)))
	(string-match "\n" str_dd)
	(setq date_mod (substring str_dd 0 (match-start 0)))
	(dict_append_proc dict_aa id name population date_mod)
)

;; --------------------------------------------------------------------
(define (text_read_proc port dict_aa)
	(let ((out_str (read-line port)))
		(if out_str (progn
			(setq delimit "\t")
;			(format standard-output "%1$s" out_str)
			(text_read_line_proc out_str delimit dict_aa)
			(text_read_proc port dict_aa)
			)
		)
	)

)

;; --------------------------------------------------------------------
(defun dict_append_proc (dict_aa id_in name population date_mod)
	(setq unit_pp (make-table string-hash string=))
	(table-set unit_pp "name" name)
	(table-set unit_pp "population" population)
	(table-set unit_pp "date_mod" date_mod)
	(table-set dict_aa id_in unit_pp)
;
	dict_aa)

;; --------------------------------------------------------------------
(defun dict_write_proc_single (fp-file key unit_bb delimit)
(write fp-file (concat key delimit
	(table-ref unit_bb "name") delimit
	(table-ref unit_bb "population") delimit
	(table-ref unit_bb "date_mod") "\n"))
)

;; --------------------------------------------------------------------
(defun text_write_proc (file_out dict_aa)
(setq delimit "\t")
(let ((fp-file (open-file file_out 'write)))
	(table-walk (lambda (key value)
		 (dict_write_proc_single fp-file key value delimit)) dict_aa)
	)
)

;; --------------------------------------------------------------------
(define (dict_update_proc doc_aa id_in population_in)
(format t "%1$s" (tablep doc_aa))
;
(let ((uu (table-ref doc_aa id_in))
	(today (current-time-string nil "%Y-%m-%d")))
	(progn
	(print "*** dict_update_proc *** pppp ***")
	(format t (table-ref uu "population"))
	(table-set uu "population" population_in)
	(table-set uu "date_mod" today)
	(print "*** dict_update_proc *** end ***\n")
	)
	)
)

;; --------------------------------------------------------------------
(define (csv_read_proc port dict_aa)
	(let ((out_str (read-line port)))
		(if out_str (progn
			(setq delimit ",")
			(text_read_line_proc out_str delimit dict_aa)
			(csv_read_proc port dict_aa)
			)
		)
	)

)

;; --------------------------------------------------------------------
(defun csv_write_proc (file_out dict_aa)
(setq delimit ",")
(let ((fp-file (open-file file_out 'write)))
	(table-walk (lambda (key value)
		 (dict_write_proc_single fp-file key value delimit)) dict_aa)
	)
)

;; --------------------------------------------------------------------

max_stars_repo_licenses: MIT Unlicense

max_issues_repo_licenses: MIT Unlicense

max_forks_repo_licenses: MIT Unlicense


text:
module MaxHelpingHandRespawningJellyfish

using ..Ahorn, Maple

@mapdef Entity "MaxHelpingHand/RespawningJellyfish" RespawningJellyfish(x::Integer, y::Integer, bubble::Bool=false, tutorial::Bool=false, respawnTime::Number=2.0, spriteDirectory::String="objects/MaxHelpingHand/glider")

const placements = Ahorn.PlacementDict(
    "Respawning Jellyfish (max480's Helping Hand)" => Ahorn.EntityPlacement(
        RespawningJellyfish
    ),
    "Respawning Jellyfish (Floating) (max480's Helping Hand)" => Ahorn.EntityPlacement(
        RespawningJellyfish,
        "point",
        Dict{String, Any}(
            "bubble" => true
        )
    )
)

function Ahorn.selection(entity::RespawningJellyfish)
    x, y = Ahorn.position(entity)
    sprite = get(entity, "spriteDirectory", "objects/MaxHelpingHand/glider") * "/idle0"

    return Ahorn.getSpriteRectangle(sprite, x, y - 4)
end

function Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::RespawningJellyfish, room::Maple.Room)
    sprite = get(entity, "spriteDirectory", "objects/MaxHelpingHand/glider") * "/idle0"
    Ahorn.drawSprite(ctx, sprite, 0, -4)

    if get(entity, "bubble", false)
        curve = Ahorn.SimpleCurve((-7, -1), (7, -1), (0, -6))
        Ahorn.drawSimpleCurve(ctx, curve, (1.0, 1.0, 1.0, 1.0), thickness=1)
    end
end

end

max_stars_repo_name: Kitty-Cats/MaxHelpingHand

max_issues_repo_name: Kitty-Cats/MaxHelpingHand

max_forks_repo_name: Kitty-Cats/MaxHelpingHand


text:
# Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors    #src
# This Source Code Form is subject to the terms of the Mozilla Public License   #src
# v.2.0. If a copy of the MPL was not distributed with this file, You can       #src
# obtain one at https://mozilla.org/MPL/2.0/.                                   #src

# # The factory schedule example

# This is a Julia translation of part 5 from "Introduction to to Linear
# Programming with Python"
# available at https://github.com/benalexkeen/Introduction-to-linear-programming
#
# For 2 factories (A, B), minimize the cost of production over the course of 12
# months while meeting monthly demand. Factory B has a planned outage during
# month 5.
#
# It was originally contributed by @Crghilardi.

using JuMP
import GLPK
import Test

function example_factory_schedule()
   ## Sets in the problem:
   months, factories = 1:12, [:A, :B]
   ## This function takes a matrix and converts it to a JuMP container so we can
   ## refer to elements such as `d_max_cap[1, :A]`.
   containerize(A::Matrix) = Containers.DenseAxisArray(A, months, factories)
   ## Maximum production capacity in (month, factory) [units/month]:
   d_max_cap = containerize([
         100000	50000;
         110000	55000;
         120000	60000;
         145000	100000;
         160000	0;
         140000	70000;
         155000	60000;
         200000	100000;
         210000	100000;
         197000	100000;
         80000	120000;
         150000	150000;
   ])
   ## Minimum production capacity in (month, factory) [units/month]:
   d_min_cap = containerize([
         20000	20000;
         20000	20000;
         20000	20000;
         20000	20000;
         20000	0;
         20000	20000;
         20000	20000;
         20000	20000;
         20000	20000;
         20000	20000;
         20000	20000;
         20000	20000;
   ])
   ## Variable cost of production in (month, factory) [$/unit]:
   d_var_cost = containerize([
         10	5;
         11	4;
         12	3;
         9	5;
         8	0;
         8	6;
         5	4;
         7	6;
         9	8;
         10	11;
         8	10;
         8	12
   ])
   ## Fixed cost of production in (month, factory) # [$/month]:
   d_fixed_cost = containerize([
         500	600;
         500	600;
         500	600;
         500	600;
         500	0;
         500	600;
         500	600;
         500	600;
         500	600;
         500	600;
         500	600;
         500	600
   ])
   ## Demand in each month [units/month]:
   d_demand = [
      120_000,
      100_000,
      130_000,
      130_000,
      140_000,
      130_000,
      150_000,
      170_000,
      200_000,
      190_000,
      140_000,
      100_000,
   ]
   ## The model!
   model = Model(GLPK.Optimizer)
   ## Decision variables
   @variables(model, begin
      status[m in months, f in factories], Bin
      production[m in months, f in factories], Int
   end)
   ## The production cannot be less than minimum capacity.
   @constraint(
      model,
      [m in months, f in factories],
      production[m, f] >= d_min_cap[m, f] * status[m, f],
   )
   ## The production cannot be more tha maximum capacity.
   @constraint(
      model,
      [m in months, f in factories],
      production[m, f] <= d_max_cap[m, f] * status[m, f],
   )
   ## The production must equal demand in a given month.
   @constraint(model, [m in months], sum(production[m, :]) == d_demand[m])
   ## Factory B is shut down during month 5, so production and status are both
   ## zero.
   fix(status[5, :B], 0.0)
   fix(production[5, :B], 0.0)
   ## The objective is to minimize the cost of production across all time
   ##periods.
   @objective(
      model,
      Min,
      sum(
         d_fixed_cost[m, f] * status[m, f] + d_var_cost[m, f] * production[m, f]
         for m in months, f in factories
      )
   )
   ## Optimize the problem
   optimize!(model)
   ## Check the solution!
   Test.@testset "Check the solution against known optimal" begin
      Test.@test termination_status(model) == MOI.OPTIMAL
      Test.@test objective_value(model) == 12_906_400.0
      Test.@test value.(production)[1, :A] == 70_000
      Test.@test value.(status)[1, :A] == 1
      Test.@test value.(status)[5, :B] == 0
      Test.@test value.(production)[5, :B] == 0
   end
   println("The production schedule is:")
   println(value.(production))
   return
end

example_factory_schedule()

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
mutable struct CustomSet{T} <: AbstractSet{T}
    dict::Dict{T,Nothing}
    CustomSet(a) = new{eltype(a)}(Dict(k => nothing for k in a))
    CustomSet(d::Dict) = new{keytype(d)}(d)
end

Base.length(s::CustomSet) = length(s.dict)

Base.iterate(s::CustomSet) = iterate(keys(s.dict))
Base.iterate(s::CustomSet,el) = iterate(keys(s.dict),el)

Base.in(el,s::CustomSet) = haskey(s.dict,el)

Base.keys(s::CustomSet) = keys(s.dict)

Base.union(l::CustomSet, r::CustomSet) = CustomSet(merge(l.dict,r.dict))
Base.union!(l::CustomSet, r::CustomSet) = begin
    merge!(l.dict,r.dict)
    l
end
Base.intersect(l::CustomSet, r::CustomSet) = CustomSet(intersect(keys(l),keys(r)))
Base.intersect!(l::CustomSet, r::CustomSet) = begin
    l.dict = Dict(k => nothing for k in intersect(keys(l),keys(r)))
    l
end

disjoint(l, r) = l ∩ r |> isempty

Base.push!(s::CustomSet, k) = begin
    s.dict[k] = nothing
    s
end

Base.setdiff(l::CustomSet, r::CustomSet) = CustomSet(setdiff(keys(l),keys(r)))
Base.setdiff!(l::CustomSet, r::CustomSet) = begin
    l.dict = Dict(k => nothing for k in setdiff(keys(l),keys(r)))
    l
end

Base.copy(s::CustomSet) = CustomSet(copy(s.dict))

complement = setdiff
complement! = setdiff!


text:
using Cosmology
using Test, Unitful, UnitfulAstro, QuadGK

# values from http://icosmos.co.uk/

dist_rtol = 1e-6
age_rtol = 2e-4
# Integrating a unitful function would require UnitfulIntegration.jl.  Without using it, we
# strip the units away from the integrand function
integrand(c, z) = 4pi*ustrip(comoving_volume_element(c, z))

@testset "FlatLCDM" begin
    c = cosmology(h=0.7, OmegaM=0.3, OmegaR=0)
    @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1651.9145u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,1,2,rtol=dist_rtol) ≈ 625.3444u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,pi,rtol=dist_rtol) ≈ angular_diameter_dist(c,0,pi,rtol=dist_rtol) rtol = dist_rtol
    @test comoving_radial_dist(c,1,rtol=dist_rtol) ≈ 3303.829u"Mpc" rtol = dist_rtol
    @test comoving_volume(c,1,rtol=dist_rtol) ≈ 151.0571u"Gpc^3" rtol = dist_rtol
    @test quadgk(z -> integrand(c, z), 0, 2.5)[1] ≈ ustrip(comoving_volume(c, 2.5))
    @test luminosity_dist(c,1,rtol=dist_rtol) ≈ 6607.6579u"Mpc" rtol = dist_rtol
    @test distmod(c,1,rtol=dist_rtol) ≈ 44.1002 rtol = dist_rtol
    @test age(c,0,rtol=age_rtol) ≈ 13.4694u"Gyr" rtol = age_rtol
    @test age(c,1,rtol=age_rtol) ≈ 5.7527u"Gyr" rtol = age_rtol
    @test lookback_time(c,1,rtol=age_rtol) ≈ (13.4694-5.7527)u"Gyr" rtol = age_rtol
    @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
end

@testset "OpenLCDM" begin
    c = cosmology(h=0.7, OmegaK=0.1, OmegaM=0.3, OmegaR=0)
    @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1619.9588u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,1,2,rtol=dist_rtol) ≈ 598.9118u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,pi,rtol=dist_rtol) ≈ angular_diameter_dist(c,0,pi,rtol=dist_rtol) rtol = dist_rtol
    @test comoving_radial_dist(c,1,rtol=dist_rtol) ≈ 3209.784u"Mpc" rtol = dist_rtol
    @test comoving_volume(c,1,rtol=dist_rtol) ≈ 140.0856u"Gpc^3" rtol = dist_rtol
    @test quadgk(z -> integrand(c, z), 0, 2.5)[1] ≈ ustrip(comoving_volume(c, 2.5))
    @test luminosity_dist(c,1,rtol=dist_rtol) ≈ 6479.8352u"Mpc" rtol = dist_rtol
    @test distmod(c,1,rtol=dist_rtol) ≈ 44.0578 rtol = dist_rtol
    @test age(c,0,rtol=age_rtol) ≈ 13.064u"Gyr" rtol = age_rtol
    @test age(c,1,rtol=age_rtol) ≈ 5.5466u"Gyr" rtol = age_rtol
    @test lookback_time(c,1,rtol=age_rtol) ≈ (13.064-5.5466)u"Gyr" rtol = age_rtol
    @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
end

@testset "ClosedLCDM" begin
    c = cosmology(h=0.7, OmegaK=-0.1, OmegaM=0.3, OmegaR=0)
    @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1686.5272u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,1,2,rtol=dist_rtol) ≈ 655.6019u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,pi,rtol=dist_rtol) ≈ angular_diameter_dist(c,0,pi,rtol=dist_rtol) rtol = dist_rtol
    @test comoving_radial_dist(c,1,rtol=dist_rtol) ≈ 3408.937u"Mpc" rtol = dist_rtol
    @test comoving_volume(c,1,rtol=dist_rtol) ≈ 163.8479u"Gpc^3" rtol = dist_rtol
    @test quadgk(z -> integrand(c, z), 0, 2.5)[1] ≈ ustrip(comoving_volume(c, 2.5))
    @test luminosity_dist(c,1,rtol=dist_rtol) ≈ 6746.1088u"Mpc" rtol = dist_rtol
    @test distmod(c,1,rtol=dist_rtol) ≈ 44.1453 rtol = dist_rtol
    @test age(c,0,rtol=age_rtol) ≈ 13.925u"Gyr" rtol = age_rtol
    @test age(c,1,rtol=age_rtol) ≈ 5.9868u"Gyr" rtol = age_rtol
    @test lookback_time(c,1,rtol=age_rtol) ≈ (13.925-5.9868)u"Gyr" rtol = age_rtol
    @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
end

@testset "FlatWCDM" begin
    c = cosmology(h=0.7, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
    @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1612.0585u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,1,2,rtol=dist_rtol) ≈ 607.6802u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,pi,rtol=dist_rtol) ≈ angular_diameter_dist(c,0,pi,rtol=dist_rtol) rtol = dist_rtol
    @test comoving_radial_dist(c,1,rtol=dist_rtol) ≈ 3224.1169u"Mpc" rtol = dist_rtol
    @test comoving_volume(c,1,rtol=dist_rtol) ≈ 140.3851u"Gpc^3" rtol = dist_rtol
    @test quadgk(z -> integrand(c, z), 0, 2.5)[1] ≈ ustrip(comoving_volume(c, 2.5))
    @test luminosity_dist(c,1,rtol=dist_rtol) ≈ 6448.2338u"Mpc" rtol = dist_rtol
    @test distmod(c,1,rtol=dist_rtol) ≈ 44.0472 rtol = dist_rtol
    @test age(c,0,rtol=age_rtol) ≈ 13.1915u"Gyr" rtol = age_rtol
    @test age(c,1,rtol=age_rtol) ≈ 5.6464u"Gyr" rtol = age_rtol
    @test lookback_time(c,1,rtol=age_rtol) ≈ (13.1915-5.6464)u"Gyr" rtol = age_rtol
    @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
end

@testset "OpenWCDM" begin
    c = cosmology(h=0.7, OmegaK=0.1, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
    @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1588.0181u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,1,2,rtol=dist_rtol) ≈ 585.4929u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,pi,rtol=dist_rtol) ≈ angular_diameter_dist(c,0,pi,rtol=dist_rtol) rtol = dist_rtol
    @test comoving_radial_dist(c,rtol=dist_rtol,1) ≈ 3147.6227u"Mpc" rtol = dist_rtol
    @test comoving_volume(c,1,rtol=dist_rtol) ≈ 132.0466u"Gpc^3" rtol = dist_rtol
    @test quadgk(z -> integrand(c, z), 0, 2.5)[1] ≈ ustrip(comoving_volume(c, 2.5))
    @test luminosity_dist(c,1,rtol=dist_rtol) ≈ 6352.0723u"Mpc" rtol = dist_rtol
    @test distmod(c,1,rtol=dist_rtol) ≈ 44.0146 rtol = dist_rtol
    @test age(c,0,rtol=age_rtol) ≈ 12.8488u"Gyr" rtol = age_rtol
    @test age(c,1,rtol=age_rtol) ≈ 5.4659u"Gyr" rtol = age_rtol
    @test lookback_time(c,1,rtol=age_rtol) ≈ (12.8488-5.4659)u"Gyr" rtol = age_rtol
    @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
end

@testset "ClosedWCDM" begin
    c = cosmology(h=0.7, OmegaK=-0.1, OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
    @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1637.5993u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,1,2,rtol=dist_rtol) ≈ 632.5829u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,pi,rtol=dist_rtol) ≈ angular_diameter_dist(c,0,pi,rtol=dist_rtol) rtol = dist_rtol
    @test comoving_radial_dist(c,1,rtol=dist_rtol) ≈ 3307.9932u"Mpc" rtol = dist_rtol
    @test comoving_volume(c,1,rtol=dist_rtol) ≈ 149.8301u"Gpc^3" rtol = dist_rtol
    @test quadgk(z -> integrand(c, z), 0, 2.5)[1] ≈ ustrip(comoving_volume(c, 2.5))
    @test luminosity_dist(c,1,rtol=dist_rtol) ≈ 6550.3973u"Mpc" rtol = dist_rtol
    @test distmod(c,1,rtol=dist_rtol) ≈ 44.0813 rtol = dist_rtol
    @test age(c,0,rtol=age_rtol) ≈ 13.5702u"Gyr" rtol = age_rtol
    @test age(c,1,rtol=age_rtol) ≈ 5.8482u"Gyr" rtol = age_rtol
    @test lookback_time(c,1,rtol=age_rtol) ≈ (13.5702-5.8482)u"Gyr" rtol = age_rtol
    @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
end

@testset "Non-Float64" begin
    # Test that FlatLCDM works with non-Float64 (BigFloat in this example)
    c = cosmology(h=0.7, OmegaM=big(0.3), OmegaR=0)
    @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1651.9145u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,1,2,rtol=dist_rtol) ≈ 625.3444u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,pi,rtol=dist_rtol) ≈ angular_diameter_dist(c,0,pi,rtol=dist_rtol) rtol = dist_rtol
    @test comoving_volume_element(c, big(1.41)) ≈ 3.4030879e10u"Mpc^3" rtol = dist_rtol
    # Test that FlatWCDM works with non-Float64 (BigFloat in this example)
    c = cosmology(h=big(0.7), OmegaM=0.3, OmegaR=0, w0=-0.9, wa=0.1)
    @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1612.0585u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,1,2,rtol=dist_rtol) ≈ 607.6802u"Mpc" rtol = dist_rtol
    @test angular_diameter_dist(c,pi,rtol=dist_rtol) ≈ angular_diameter_dist(c,0,pi,rtol=dist_rtol) rtol = dist_rtol
    @test comoving_volume_element(c, big(1.41)) ≈ 3.1378625e10u"Mpc^3" rtol = dist_rtol
end

@testset "Unit conversion" begin
    c = cosmology(h=0.9, OmegaM=0.5, OmegaR=0)
    for u in (u"m", u"pc", u"ly")
        @test unit(luminosity_dist(u, c, 1)) == u
        @test unit(angular_diameter_dist(u, c, 2)) == u
    end
    for u in (u"s", u"yr")
        @test unit(age(u, c, 3)) == u
        @test unit(lookback_time(u, c, 4)) == u
    end
end

@testset "Utilities" begin
    c = cosmology(h = 0.7)
    @test hubble_time(c, 0) ≈ Cosmology.hubble_time0(c)
    @test hubble_dist(c, 0) ≈ Cosmology.hubble_dist0(c)
    @test H(c, 0) ≈ 70u"km/s/Mpc"
end