text:
using Distributions:Random
using AdaptiveABC
using SequentialLOB
using Random
using Distributions
using Optim
using NLopt
using JLD2
using Plots

# Set Simulator Parameters and Seed
D = 1.0
σ = 1.5
nu = 0.5
μ = 1.0
parameter_true = [D, σ, nu, μ]
parameter_names = ["D", "σ", "nu", "μ"]
plotting_parameter_names = ["D", "σ", "ν", "μ"]
n_replications = 10
Random.seed!(8929200)
# Constants
num_paths = 1; M = 400 ; T = 2299 ; p₀ = 238.75 ; L = 200
Δx = L / M ; λ = 1.0 ; Δt = (Δx^2) / (2.0 * D) ; α_lob = 0.0
println("Δx = $Δx and Δt = $Δt")

# Create model, simulate to obtain observation vector and then calculate summary statistics
true_lob_model = SLOB(num_paths,
    T, p₀, M, L, D, σ, nu, α_lob, SourceTerm(λ, μ))
true_lob_price_path = true_lob_model(7136)
true_lob_log_returns = diff(log.(true_lob_price_path[:,1]))
true_lob_summary_stats = get_summary_stats(true_lob_log_returns)

# plt = plot(
#     1:T + 1, true_lob_price_path[:,1], legend=false,
#     xlab="Time", ylab="Mid-Price",
# )

# Define Summary function used in AdaptiveABC. Takes in a vector of parameters and outputs a success boolean and the
# simulated summart statistics
function summary_fn(parameters, n_summary_stats, n_replications)
    try
        D, σ, nu, μ = parameters
        model = SLOB(num_paths, T, p₀, M, L, D, σ, nu, α_lob, SourceTerm(λ, μ))
        summary_stats = Array{Float64,2}(undef, n_summary_stats, n_replications)
        for i in 1:n_replications
            sim_price_path = model()
            sim_obs = diff(log.(sim_price_path[:, 1]))
            summary_stats[:, i] = get_summary_stats(true_lob_log_returns, sim_obs)
        end
        return true, summary_stats
    catch e
        return false, zeros(n_summary_stats, n_replications)
    end
end

# Setup for Calibration of Parameters through AdaptiveABC
# Define Prior distribution using Prior struct
prior = Prior([
    Uniform(0.5, 3.0), # D
    Uniform(0.1, 3.0), # σ
    Uniform(0.05, 0.95), # nu
    Uniform(0.1, 3.0) # μ
])

#####
# Calibration Technique: ABC Rejection BBWM
#####

# Distance Function to measure distance between observed summary stats and simulated summary stats
boot_weight_matrix = BlockBootstrapWeightMatrix(122477, vcat([0], true_lob_log_returns), get_summary_stats, 100, 10_000)
weighted_bootstrap = WeightedBootstrap(vcat([0], true_lob_log_returns), get_summary_stats, boot_weight_matrix)

abc_input_wb = ABCInput(
    prior,
    parameter_names,
    true_lob_summary_stats,
    summary_fn,
    17,
    weighted_bootstrap
)

abc_reject_wb = JLD2.load_object("$(dirname(@__FILE__))/abc_reject_wb.jld2")

# abc_reject_wb = ABCRejection(
#     abc_input_wb,
#     3_000,
#     n_replications,
#     100.0,
#     parallel=true,
#     seed=2517413766247174462
# )

posterior_means = AdaptiveABC.parameter_means(abc_reject_wb)
println()
println("ABC rejection BBWM: D, σ, nu, μ")
print(round.(posterior_means[:,end], digits=3))
println()

# JLD2.save_object("$(dirname(@__FILE__))/abc_reject_wb.jld2", abc_reject_wb)

# ABC rejection BBWM: D, σ, nu, μ
# [1.551, 1.821, 0.503, 1.496]

param_inds = [1,2,3,4]
plt = plot(abc_reject_wb, iteration_colours=cgrad(:blues, 5, categorical=true),
    iterations=[1], params_true=parameter_true[param_inds],
    prior_dists=abc_input_wb.prior.distribution[param_inds],
    param_inds=param_inds, param_names=plotting_parameter_names[param_inds])

savefig(plt, "$(dirname(@__FILE__))/abc_reject_wb.pdf")
plt

#####
# Calibration Technique: ABC Rejection MADWE
#####
abc_input_we = ABCInput(
    prior,
    parameter_names,
    true_lob_summary_stats,
    summary_fn,
    17,
    WeightedEuclidean(true_lob_log_returns, get_summary_stats, ones(17), "MAD")
)

abc_reject_we = JLD2.load_object("$(dirname(@__FILE__))/abc_reject_we.jld2")
new_cutoff = 25.0

abc_reject_we.parameters = abc_reject_we.parameters[:, abc_reject_we.distances .<= new_cutoff]
abc_reject_we.n_successes = sum(abc_reject_we.distances .<= new_cutoff)
abc_reject_we.summary_stats = abc_reject_we.summary_stats[:, abc_reject_we.distances .<= new_cutoff, :]
abc_reject_we.weights = abc_reject_we.weights[abc_reject_we.distances .<= new_cutoff]

abc_reject_we.distances = abc_reject_we.distances[abc_reject_we.distances .<= new_cutoff]
# abc_reject_we = ABCRejection(
#     abc_input_we,
#     3_000,
#     n_replications,
#     23.0,
#     parallel=true,
#     seed=738173189
# )
posterior_means = AdaptiveABC.parameter_means(abc_reject_we)
println()
println("ABC rejection MADWE: D, σ, nu, μ")
print(round.(posterior_means[:,end], digits=3))
println()

# JLD2.save_object("$(dirname(@__FILE__))/abc_reject_we.jld2", abc_reject_we)

# ABC rejection MADWE: D, σ, nu, μ
# [1.747, 1.536, 0.51, 1.54]

param_inds = [1,2,3,4]
plt = plot(abc_reject_we, iteration_colours=cgrad(:blues, 5, categorical=true),
    iterations=[1], params_true=parameter_true[param_inds],
    prior_dists=abc_input_we.prior.distribution[param_inds],
    param_inds=param_inds, param_names=plotting_parameter_names[param_inds])

savefig(plt, "$(dirname(@__FILE__))/abc_reject_we.pdf")
plt

#####
# Calibration Technique: ABC-PMC BBWM
#####

# Distance Function to measure distance between observed summary stats and simulated summary stats
boot_weight_matrix = BlockBootstrapWeightMatrix(122477, vcat([0], true_lob_log_returns), get_summary_stats, 100, 10_000)
weighted_bootstrap = WeightedBootstrap(vcat([0], true_lob_log_returns), get_summary_stats, boot_weight_matrix)

abc_input_wb = ABCInput(
    prior,
    parameter_names,
    true_lob_summary_stats,
    summary_fn,
    17,
    weighted_bootstrap
)

abc_pmc_out_wb = JLD2.load_object("$(dirname(@__FILE__))/abc_pmc_out_wb.jld2")
# abc_pmc_out_wb = ABC_PMC(
#     abc_input_wb,
#     100,
#     n_replications,
#     0.25,
#     10_000,
#     10;
#     parallel=true,
#     batch_size=250,
#     seed=918731553
# )
posterior_means = AdaptiveABC.parameter_means(abc_pmc_out_wb)

println()
println("ABC-PMC BBWM: D, σ, nu, μ")
print(round.(posterior_means[:,end], digits=3))
println()

# JLD2.save_object("$(dirname(@__FILE__))/abc_pmc_out_wb.jld2", abc_pmc_out_wb)

# ABC-PMC BBWM: D, σ, nu, μ
# [1.086, 1.596, 0.457, 1.279]

param_inds = [1,2,3,4]
plt = plot(abc_pmc_out_wb, iteration_colours=cgrad(:blues, 5, categorical=true),
    iterations=[1, abc_pmc_out_wb.n_iterations ÷ 2, abc_pmc_out_wb.n_iterations], params_true=parameter_true[param_inds],
    prior_dists=abc_input_wb.prior.distribution[param_inds],
    param_inds=param_inds, param_names=plotting_parameter_names[param_inds])
plt
savefig(plt, "$(dirname(@__FILE__))/abc_pmc_wb.pdf")

#####
# Calibration Technique: ABC-PMC MADWE
#####
abc_input_we = ABCInput(
    prior,
    parameter_names,
    true_lob_summary_stats,
    summary_fn,
    17,
    WeightedEuclidean(true_lob_log_returns, get_summary_stats, ones(17), "MAD")
)

abc_pmc_out_we = JLD2.load_object("$(dirname(@__FILE__))/abc_pmc_out_we.jld2")
# abc_pmc_out_we = ABC_PMC(
#     abc_input_we,
#     100,
#     n_replications,
#     0.15,
#     10_000,
#     10;
#     parallel=true,
#     batch_size=250,
#     seed=83912738712
# )
posterior_means = AdaptiveABC.parameter_means(abc_pmc_out_we)
println("")
println("ABC-PMC MADWE: D, σ, nu, μ")
print(round.(posterior_means[:,end], digits=3))
println("")
# ABC-PMC MADWE: D, σ, nu, μ
# [1.182, 1.767, 0.566, 1.75]

# JLD2.save_object("$(dirname(@__FILE__))/abc_pmc_out_we.jld2", abc_pmc_out_we)

param_inds = [1,2,3,4]
plt = plot(abc_pmc_out_we, iteration_colours=cgrad(:blues, 5, categorical=true),
    iterations=[1, abc_pmc_out_we.n_iterations ÷ 2, abc_pmc_out_we.n_iterations], params_true=parameter_true[param_inds],
    prior_dists=abc_input_we.prior.distribution[param_inds],
    param_inds=param_inds, param_names=plotting_parameter_names[param_inds])
savefig(plt, "$(dirname(@__FILE__))/abc_pmc_we.pdf")
plt

#####
# Calibration Technique: Nelder-Mead BBWM
#####

# Distance Function to measure distance between observed summary stats and simulated summary stats
boot_weight_matrix = BlockBootstrapWeightMatrix(122477, vcat([0], true_lob_log_returns), get_summary_stats, 100, 10_000)
weighted_bootstrap = WeightedBootstrap(vcat([0], true_lob_log_returns), get_summary_stats, boot_weight_matrix)

function f(x, weight, reps)
    success, sum_stats = summary_fn(x, 17, reps)
    return weight(sum_stats)
end

f_wb(x, grad) = f(x, weighted_bootstrap, n_replications)
# D, σ, nu, μ
lower = [0.5, 0.1, 0.05, 0.1]
upper = [3.0, 3.0, 0.95, 3.0]

nm_reps = 10
results = zeros(4, nm_reps)
Random.seed!(736161311)
for i in 1:nm_reps
    println("Iteration $i")
    opt = Opt(:LN_NELDERMEAD, 4)
    opt.lower_bounds = lower
    opt.upper_bounds = upper
    opt.xtol_rel = 1e-6
    opt.min_objective = f_wb
    init_params = rand(Int.(rand(MersenneTwister(), UInt32)), prior)

    (minf, minx, ret) = NLopt.optimize(opt, init_params)

    results[:,i] = minx
end
println()
println("Nelder-Mead BBWM: ")
print(round.(mean(results, dims=2), digits=3))
print()
results
println()
# Nelder-Mead BBWM:
# [1.778; 1.944; 0.719; 1.756]

#####
# Calibration Technique: Nelder-Mead MADWE
#####

weighted_mad = WeightedEuclidean(true_lob_log_returns, get_summary_stats, ones(17), "MAD")
sim_sum_stats = zeros(17, 100, n_replications)
Random.seed!(1773616333)
for i in 1:size(sim_sum_stats, 2)
    params = rand(Int.(rand(MersenneTwister(), UInt32)), prior)
    success, sum_stats = summary_fn(params, 17, n_replications)
    sim_sum_stats[:,i,:] = sum_stats
end

inv_weights = [AdaptiveABC.MAD(sim_sum_stats[i, :, :]) for i in 1:17]  # For each statistic, calculate the MAD across all n_sim, n_rep values
weights = 1.0 ./ inv_weights

weighted_mad.weights = weights
f_mad(x, grad) = f(x, weighted_mad, n_replications)
# D, σ, nu, μ
lower = [0.5, 0.1, 0.05, 0.1]
upper = [3.0, 3.0, 0.95, 3.0]

nm_reps = 10
results = zeros(4, nm_reps)
Random.seed!(9918133091)
for i in 1:nm_reps
    println("Iteration $i")
    opt = Opt(:LN_NELDERMEAD, 4)
    opt.lower_bounds = lower
    opt.upper_bounds = upper
    opt.xtol_rel = 1e-6
    opt.min_objective = f_mad
    init_params = rand(Int.(rand(MersenneTwister(), UInt32)), prior)

    (minf, minx, ret) = NLopt.optimize(opt, init_params)

    results[:,i] = minx
end
println()
println("Nelder-Mead MADWE: ")
print(round.(mean(results, dims=2), digits=3))
results
println()
# Nelder-Mead MADWE:
# [2.373; 1.924; 0.52; 1.295]

max_stars_repo_name: GantZA/CALOBMTMD

max_issues_repo_name: GantZA/CALOBMTMD

max_forks_repo_name: GantZA/CALOBMTMD


text:
filename = "warm_up.d"
infile = open(filename)
start_time = 0.422949E7

outfile = open("new_$(filename)at","w")
old_time_str = ""
for ln in eachline(infile)
#  println(ln)
  newln = ln

  new_time = "$(newln[3:14])"
  new_ρ    = "$(newln[17:28])"
  new_av   = "$(newln[31:42])"
  new_temp = "$(newln[45:56])"

  old_time = float(new_time)
  new_time = old_time + start_time
  new_av   = log(10.0)
#  new_av   = log(float(new_av))
  new_ρ    = log(float(new_ρ))
  new_temp = log(float(new_temp))

  ρ_str    = @sprintf("%9.3e",new_ρ)
  av_str   = @sprintf("%9.3e",new_av)
  temp_str = @sprintf("%9.3e",new_temp)
  time_str = @sprintf("%9.3e",new_time)

  if ( old_time_str == time_str )
    continue
  else
    println(old_time,ρ_str,av_str,temp_str)
    outline = "$(time_str)  $(ρ_str) $(av_str) $(temp_str)\n"
  end
  write(outfile,outline)
  old_time_str = time_str
end

close(infile)
close(outfile)

ext: jl


text:
"""
    EmptyVector() :: AbstractVector{Union{}}
    EmptyVector(itr) :: AbstractVector
    EmptyVector{T}() :: AbstractVector{T}
    EmptyVector{T}(itr) :: AbstractVector{T}

Create an empty vector.  Iterator `itr` must be empty.

The unary constructors `EmptyVector(itr)` and `EmptyVector{T}(itr)`
asserts that the iterator `itr` is empty.  The constructor
`EmptyVector(itr)` uses `eltype(itr)` if `IteratorEltype(itr)` is
`HasEltype`.
"""
EmptyVector

"""
    SingletonVector(itr) :: AbstractVector
    SingletonVector{T}(itr) :: AbstractVector{T}

Create a singleton vector.  Iterator `itr` must have one and only one
element.

The constructor `SingletonVector(itr)` uses `eltype(itr)` if
`IteratorEltype(itr)` is `HasEltype`.
"""
SingletonVector

abstract type AbstractMicroVector{T,L} <: AbstractVector{T} end
abstract type AbstractEmptyVector{T,L} <: AbstractMicroVector{T,L} end
abstract type AbstractSingletonVector{T,L} <: AbstractMicroVector{T,L} end

struct EmptyVector{T,L} <: AbstractEmptyVector{T,L} end
EmptyVector() = EmptyVector{Union{},Vector}()
EmptyVector{T}() where {T} = EmptyVector{T,Vector}()

struct SingletonVector{T,L} <: AbstractSingletonVector{T,L}
    value::T
end
SingletonVector{T}(x::T) where {T} = SingletonVector{T,Vector}(x)

@inline getvalue(A::SingletonVector) = A.value
@inline upcast(A::AbstractMicroVector{T,L}) where {T,L} = L(A)::AbstractVector{T}

Base.size(::AbstractEmptyVector) = (0,)
Base.IndexStyle(::Type{<:AbstractEmptyVector}) = Base.IndexLinear()
@inline function Base.getindex(A::AbstractEmptyVector, i::Int)
    checkbounds(A, i)
end

Base.size(::AbstractSingletonVector) = (1,)
Base.IndexStyle(::Type{<:AbstractSingletonVector}) = Base.IndexLinear()
@inline function Base.getindex(A::AbstractSingletonVector, i::Int)
    @boundscheck checkbounds(A, i)
    return getvalue(A)
end

emptyshim(::Type{L}, ::Type{T}) where {T,L<:AbstractVector} = EmptyVector{T,L}()
singletonshim(::Type{L}, x::T) where {T,L<:AbstractVector} = SingletonVector{T,L}(x)

function EmptyVector(itr)
    validate_empty_iterator(itr, :EmptyVector)
    if Base.IteratorEltype(itr) isa Base.HasEltype
        return EmptyVector{eltype(itr)}()
    else
        return EmptyVector()
    end
end

function EmptyVector{T}(itr) where {T}
    validate_empty_iterator(itr, :EmptyVector)
    return EmptyVector{T}()
end

function SingletonVector(itr)
    x = validate_singleton_iterator(itr, :SingletonVector)
    if Base.IteratorEltype(itr) isa Base.HasEltype
        return SingletonVector{eltype(itr)}((x,))
    else
        return SingletonVector{typeof(x)}(x)
    end
end

function SingletonVector{T}(itr) where {T}
    x = validate_singleton_iterator(itr, :SingletonVector)
    return SingletonVector{T}((x,))
end

# fast-path:
SingletonVector((x,)::Tuple{T}) where {T} = SingletonVector{T}(x)
SingletonVector{T}((x,)::Tuple{Any}) where {T} = SingletonVector{T}(convert(T, x)::T)

function Base.showarg(io::IO, v::EmptyVector{<:Any,Vector}, toplevel::Bool)
    @nospecialize
    print(io, "EmptyVector")
    toplevel && print(io, '{', eltype(v), '}')
end

function Base.showarg(io::IO, v::SingletonVector{<:Any,Vector}, toplevel::Bool)
    @nospecialize
    print(io, "SingletonVector")
    toplevel && print(io, '{', eltype(v), '}')
end


text:
## Basic functions
using DataStructures
using Flux: onehot
using PyCall
chem = pyimport_conda("rdkit.Chem", "rdkit", "conda-forge")

# currently supported atomic features and bond features
const atom_feat_list = ["symbol","degree","implicit_valence","formal_charge","radical_electrons","hybridization","aromaticity","total_H_num" ]
const bond_feat_list = ["bond_type","isConjugated","isInring"]

# Map atomic feature names to RDKit functions
const rdkit_atom_functions = Dict(
    "symbol"=>chem.Atom.GetSymbol,
    "degree"=>chem.Atom.GetDegree,
    "implicit_valence"=>chem.Atom.GetImplicitValence,
    "formal_charge"=>chem.Atom.GetFormalCharge,
    "radical_electrons"=>chem.Atom.GetNumRadicalElectrons,
    "hybridization"=>chem.Atom.GetHybridization,
    "aromaticity"=>chem.Atom.GetIsAromatic,
    "total_H_num"=>chem.Atom.GetTotalNumHs
)

# Map bond feature names to RDKit functions
const rdkit_bond_functions = Dict(
    "bond_type"=>chem.Bond.GetBondType,
    "isConjugated"=>chem.Bond.GetIsConjugated,
    "isInring"=>chem.Bond.IsInRing
)

# bins for atomic features
const atom_feat_bins = Dict(
    "symbol" => ["C","N","O","S","F","Si","P","Cl","Br","Mg","Na","Ca","Fe","As","Al","I","B","V","K","Tl","Yb","Sb","Sn","Ag","Pd","Co","Se","Ti","Zn",'H',"Li","Ge","Cu","Au","Ni","Cd","In","Mn","Zr","Cr","Pt","Hg","Pb","Unknown"],
    "degree" => [0,1,2,3,4,5,6,7,8,9,10],
    "implicit_valence" => [0,1,2,3,4,5,6],
    "formal_charge" => [nothing],  # [-1,0,1,2], # indicates using real value
    "radical_electrons" => [nothing],  # [0,1,2,3], # indicates using use real value
    "hybridization" => [chem.rdchem.HybridizationType.SP, chem.rdchem.HybridizationType.SP2,chem.rdchem.HybridizationType.SP3,chem.rdchem.HybridizationType.SP3D, chem.rdchem.HybridizationType.SP3D2],  # [2,3,4,5,6],
    "aromaticity" => [nothing],  # [0,1], # indicates using real value
    "total_H_num" => [0,1,2,3,4]
)

# bins for bond features
const bond_feat_bins = Dict(
    "bond_type" => [chem.rdchem.BondType.SINGLE, chem.rdchem.BondType.DOUBLE,chem.rdchem.BondType.TRIPLE, chem.rdchem.BondType.AROMATIC],  #[1,2,3,12],
    "isConjugated" => [nothing],  # [0,1],# indicates using real value
    "isInring" => [nothing]  #[0,1],# indicates using real value
)

"""
input:
    atom (RDKit.Atom)
    feature_list(Array{String}): list of atomic features of interest
output:
    atom_feat (OrderedDict): Dictionary of atomic features
"""
function smiles_atom_features(atom; feature_list::Array{String}=atom_feat_list)
    atom_feat = OrderedDict()
    for f in feature_list
        @assert f in atom_feat_list "Feature "*f*" is not in supported atomic features list: [symbol,degree,implicit_valence,formal_charge,radical_electrons,hybridization,aromaticity,total_H_num]"
        val = rdkit_atom_functions[f](atom)
        bins = atom_feat_bins[f]
        if isnothing(bins[1]) # use real value instead of one-hot
            atom_feat[f] = Float32(val)
        else
            onehot_vec = zeros(Float32,size(bins,1))
            try  # check if the value can be found in the given bins
                onehot_vec = Float32.(val .== bins)
            catch
                @warn f*" value out of range: "*string(val)*", use the max accpeted value instead"
                onehot_vec = Float32.(bins[end] .== bins)
            end
            atom_feat[f] = onehot_vec
        end
    end
    return atom_feat
end

"""
input:
    bond (RDKit.Bond)
    feature_list(Array{String}): list of bond features of interest
output:
    bond_feat (OrderedDict): Dictionary of bond features
"""
function smiles_bond_features(bond; feature_list::Array{String}=bond_feat_list)
    bond_feat = OrderedDict()
    for f in feature_list
        @assert f in bond_feat_list "Feature "*f*" is not in supported bond features list: [bond_type, isConjugated, isInring]"
        val = rdkit_bond_functions[f](bond)
        bins = bond_feat_bins[f]
        if isnothing(bins[1]) # use real value instead of one-hot
            bond_feat[f] = Float32(val)
        else
            onehot_vec = zeros(Float32,size(bins,1))
            try  # check if the value can be found in the given bins
                onehot_vec = Float32.(val .== bins)
            catch
                @warn f*" value out of range: "*string(val)*", use the max accpeted value instead"
                onehot_vec = Float32.(bins[end] .== bins)
            end
            bond_feat[f] = onehot_vec
        end
    end
    return bond_feat
end

"""
Pair features that will be used in weave_featurize()
input:
    mol (RDKit SMILES mol): mol coverted from SIMLES
    bond_feat(OrderedDict): list of bond features of mol
    adj_list(Array{Array{Int}}): adjacency list
    bf_len(Int): bond feature length
    max_distance(Int): cut-off distance
output:
    features (Array{Float32,2}): pair features for mol
"""
function pair_features(mol, bond_feats::OrderedDict, adj_list; bf_len = 6, max_distance = 7)
    N = mol.GetNumAtoms()  # number of atoms
    features = zeros(Float32,N,N,bf_len+max_distance+1) # declear feature matrix
    if length(bond_feats) == 0
        @warn "No bonds found in "*string(chem.MolToSmiles(mol))*"! Pair feature matrix filled with zeros!"
        return reshape(features,(N*N,bf_len+max_distance+1))
    end
    rings = mol.GetRingInfo().AtomRings()
    for a1 in 1:N
        for a2 in adj_list[a1]
            # first bf_len features are bond features
            index = a1<a2 ? (a1,a2) : (a2,a1)
            features[a1, a2, 1:bf_len] = bond_feats[index]
        end
        for ring in rings
            ring = ring .+ 1  # Julia index starts from 1
            if a1 in ring
                # `bf_len+1`-th feature is if the pair of atoms are in the same ring
                features[a1, collect(ring), bf_len+1] .= 1
                features[a1, a1, bf_len+1] = 0
            end
        end
        distance_matrix = zeros(Float32,N,max_distance)
        distances = _get_distance!(distance_matrix, a1, adj_list, max_distance)
        features[a1,:,bf_len+2:end] = distance_matrix
    end
    features = vcat([features[i,:,:] for i=1:size(features,1)]...)
    return features
end

"""
Helper function for pair_features()
Finds distance between the input atom and all other atoms using BFS

input:
    distance_matrix(Array{Int,2}): predefined matrix to be filled in with distances
    a1(Int): index for current atom
    adj_list(Array{Int}): adjacency list for a1
    max_distance(Int): cut-off distance
"""
function _get_distance!(distance_matrix,a1, adj_list, max_distance = 7)
    visited = falses(length(adj_list))
    visited[a1] = true
    cur_list = adj_list[a1]
    cur_distance = 1
    while length(cur_list)!=0 && cur_distance<=max_distance
        next_list = []
        for a2 in cur_list
            if visited[a2]
                continue
            end
            visited[a2] = true
            distance_matrix[a2,cur_distance] = 1
            for a3 in adj_list[a2]
                if ! visited[a3]
                    push!(next_list,a3)
                end
            end
        end
        cur_list = next_list
        cur_distance += 1
    end
end


text:
# small scrip to compute the solution of Lippmann-Schwinger equation
# We test the types introduced in FastConvolution.jl
# we test that the application is fast and that the construction
# is performed fast.

using PyPlot
using IterativeSolvers

include("../src/FastConvolution.jl")
include("../src/FastConvolution3D.jl")
include("../src/Preconditioner.jl")

# setting the number of threads for the FFT and BLAS
# libraries (please set them to match the number of
# physical cores in your system)
FFTW.set_num_threads(4);
BLAS.set_num_threads(4);

#Defining Omega
h = 1/64
k = 1/h

D0 = 0.0;

# size of box
a  = 1
x = collect(-a/2:h:a/2-h)
y = collect(-a/2:h:a/2-h)
z = collect(-a/2:h:a/2-h)
(n,m,l) = length(x), length(y), length(z)
N = n*m*l

X = zeros(n,m,l); Y = zeros(n,m,l); Z = zeros(n,m,l);
for i=1:n, j=1:m, p=1:l
    X[i,j,p] =  x[i];
    Y[i,j,p] =  y[j];
    Z[i,j,p] =  z[p];
end
X = X[:]; Y = Y[:]; Z = Z[:];
# we solve \triangle u + k^2(1 + nu(x))u = 0

# Defining the smooth perturbation of the slowness
nu(x,y,z) = 0.3*exp(-40*(x.^2 + y.^2 + z.^2)).*(abs(x).<0.48).*(abs(y).<0.48).*(abs(z).<0.48);

#
NU = reshape(nu(X,Y,Z), n,m,l);
#
figure(5); clf(); imshow(real(NU[:,:,15]))

## You can choose between Duan Rohklin trapezoidal quadrature
# fastconv = buildFastConvolution(x,y,h,k,nu)

# or Greengard Vico Quadrature (this is not optimized and is 2-3 times slower)
fastconv = buildFastConvolution3D(x,y,z,X,Y,Z,h,k,nu, quadRule = "Greengard_Vico");

# # assembling the sparsifiying preconditioner
@time As = buildSparseA3DConv(k,X,Y,Z,fastconv, n ,m, l);

# # assembling As*( I + k^2G*nu)
#Mapproxsp = As + k^2*(buildSparseAG3D(k,X,Y,Z,D0, n ,m,l)*spdiagm(nu(X,Y,Z)));
@time Mapproxsp = As + k^2*(buildSparseAG3DConv(k,X,Y,Z,fastconv, n ,m,l)*spdiagm(nu(X,Y,Z)));

# # defining the preconditioner
precond = SparsifyingPreconditioner(Mapproxsp, As)

# building the RHS from the incident field
u_inc = exp(k*im*X);
rhs = -(fastconv*u_inc - u_inc);

# allocating the solution
u = zeros(Complex128,N);

# solving the system using GMRES
@time info =  gmres!(u, fastconv, rhs, precond)
println(info[2].residuals[:])

# plotting the solution
figure(1)
clf()

U = reshape(u+u_inc, n,m,l);

imshow(real(U[:,:,round(Integer, end/2)]))

max_stars_repo_name: Forgotten/Lippmann-Schingwer

max_issues_repo_name: Forgotten/Lippmann-Schingwer

max_forks_repo_name: Forgotten/Lippmann-Schingwer


text:
@testset "Utils" begin
    @testset "versioninfo" begin
        AMDGPU.versioninfo()
        @test true  # Just test that versioninfo() prints something.
    end

    @testset "has_rocm_gpu" begin
        @test AMDGPU.has_rocm_gpu()
    end
end


text:
# This file was generated, do not modify it. # hide
rand(Normal())

max_stars_repo_name: cscherrer/InformativePrior

max_issues_repo_name: cscherrer/InformativePrior

max_forks_repo_name: cscherrer/InformativePrior


text:
# Test interpolation methods
x = 1:10
y = rand(10)
xgrid = 0:0.5:10

# Test that interpolation functions work
y_intp = linear_interpolation(x, y)

@test length(y_intp(xgrid)) == length(xgrid)
@test y_intp(xgrid) isa Vector{Float64}

# Test that interpolation work for uncertain index-value datasets
timeinds = [UncertainValue(Normal, i, 0.1) for i = 1:5]
measurements = [UncertainValue(Gamma, i, i + rand(1:2)) for i = 1:5]

d = UncertainIndexValueDataset(timeinds, measurements)
grid = RegularGrid(0, 5, 0.4)

@test length(resample(d, grid)) == 2
@test length(resample(d, StrictlyIncreasing(), grid)) == 2

@test create_interp_scheme(1:10, rand(10), RegularGrid(0, 1, 0.1)) isa InterpolationScheme1D


text:
function project_on_zero_levelset(
    xguess,
    func,
    grad,
    tol,
    r;
    maxiter = 20,
    normtol = 1e3eps(),
    perturbation = 1e-2,
)
    x0 = copy(xguess)
    x1 = copy(x0)
    dim = length(xguess)

    for counter = 1:maxiter
        vf = func(x0)
        gf = grad(x0)

        if norm(gf) < normtol
            x0 = x0 + perturbation * (rand(dim) .- 0.5)
        else
            δ = vf / (gf' * gf) * gf
            normδ = norm(δ)
            if normδ > 0.5r
                δ *= 0.5r / normδ
            end

            x1 = x0 - δ

            if abs(func(x1)) < tol
                flag = norm(x1 - xguess) < r
                return x1, flag
            else
                x0 = x1
            end
        end
    end
    error("Did not converge after $maxiter iterations")
end

function tensor_product_points(p1, p2)
    n1 = size(p1)[2]
    n2 = size(p2)[2]
    return vcat(repeat(p1, inner = (1, n2)), repeat(p2, outer = (1, n1)))
end

function reference_seed_points(n)
    @assert n > 0
    xrange = range(-1.0, stop = 1.0, length = n + 2)
    points = tensor_product_points(xrange[2:n+1]', xrange[2:n+1]')
end

function seed_cell_zero_levelset(xguess, func, grad, tol, boundingradius)
    dim, nump = size(xguess)
    projectedpoints = zeros(dim, nump)
    flags = zeros(Bool, nump)
    for idx = 1:nump
        p, f = project_on_zero_levelset(
            xguess[:, idx],
            func,
            grad,
            tol,
            boundingradius,
        )
        projectedpoints[:, idx] = p
        flags[idx] = f
    end

    return projectedpoints, flags
end

function seed_zero_levelset(nump, levelset, cellids, tol, boundingradius)

    dim = dimension(levelset)
    numcells = length(cellids)
    refpoints = reference_seed_points(nump)
    numrefpoints = size(refpoints)[2]
    totalnumpoints = numcells * numrefpoints

    refseedpoints = zeros(dim, totalnumpoints)
    refseedcellids = zeros(Int, totalnumpoints)
    flags = zeros(Bool, totalnumpoints)

    start = 1
    stop = numrefpoints

    for cellid in cellids
        load_coefficients!(levelset, cellid)
        poly = interpolater(levelset)

        xk, f = seed_cell_zero_levelset(
            refpoints,
            poly,
            x -> vec(gradient(poly, x)),
            tol,
            boundingradius,
        )

        refseedpoints[:, start:stop] = xk
        refseedcellids[start:stop] = repeat([cellid], numrefpoints)
        flags[start:stop] = f

        start = stop + 1
        stop = stop + numrefpoints
    end

    validpoints = findall(flags)

    return refseedpoints[:, validpoints], refseedcellids[validpoints]
end

function seed_zero_levelset(
    nump,
    levelset,
    cutmesh;
    tol = 1e4eps(),
    boundingradius = 4.5,
)

    numcells = number_of_cells(cutmesh)
    cellsign = [cell_sign(cutmesh, cellid) for cellid = 1:numcells]
    cellids = findall(cellsign .== 0)
    return seed_zero_levelset(nump, levelset, cellids, tol, boundingradius)
end


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ bb6cae46-4108-11eb-2309-3d11f13a19fa
begin
	using Pkg
	Pkg.activate()
	Pkg.add([
			"Plots",
			"PlutoUI",
			"StatsBase",
			"DataStructures",
			"Unitful",
			"PhysicalConstants",
			"UnitfulEquivalences",
			"DataFrames",
			"CSV",
			"Dates",
			"TimeZones",
			"StatsPlots",
			"ImageFiltering",
			"OffsetArrays",
			"Chain",
			"FluxArchitectures",
			"LsqFit",
			"JSON3",
			"JSONTables",
			])
end

# ╔═╡ 0b484dcf-b479-46d7-8c86-d96a43ff458d
begin
	Pkg.develop("TSAnalysis")
	using LinearAlgebra, Optim, Measures;
	using TSAnalysis;
end

# ╔═╡ d8b334ae-4108-11eb-2546-99e75efc18d6
begin
	using Plots
	plotly()
	using PlutoUI
	using StatsBase
	using DataStructures
	using DataFrames
	using CSV
	using TimeZones
	using Dates
	using StatsPlots
	using ImageFiltering
	using OffsetArrays
	using Chain
	using FluxArchitectures
	using Statistics
	using LsqFit
	using HTTP
	using JSONTables
	using Downloads
	using JSON3
	TableOfContents()
end

# ╔═╡ e6b122ce-5548-11eb-3980-53dfd8e5897c
begin
	frmt = dateformat"yyyy-mm-dd HH:MM:SS"
	column_types = Dict(:ResultTime=>DateTime, :AnalysisInsertTime=>DateTime)
end

# ╔═╡ 7f15b8d8-17d9-4b89-aedb-99ff090a91ea
begin
	page_url = "https://opendata.digilugu.ee/"
	download_path = "/home/kris/Downloads/"
	filenames = (
		"opendata_covid19_hospitalization_profile.csv",
		"opendata_covid19_test_results.csv",
		"covid19/vaccination/v2/opendata_covid19_vaccination_location_county_agegroup_gender.csv",
		"opendata_covid19_tests_total.csv",
		)

	if false
		for filename = filenames
			download(string(page_url, filename), string(download_path, filename))
		end
	end

	hosp = CSV.read(string(download_path, filenames[1]), DataFrame)
	df = CSV.read(string(download_path, filenames[2]), DataFrame, types=column_types, dateformat=frmt)
	dropmissing!(df)
	vac = CSV.read(string(download_path, filenames[3]), DataFrame)
	df_total = CSV.read(string(download_path, filenames[4]), DataFrame)
	dropmissing!(df_total)
end

# ╔═╡ eaab22ca-5b2c-11eb-1eb6-cf73054ee98e
begin
	inimarv, ages_df = download_popstats()
end

# ╔═╡ d35a4a83-2585-4b17-887a-1d184b187921

# ╔═╡ 6fcdea12-35f3-4648-8ba0-cfd6e89547ef
counties = collect(setdiff(keys(inimarv), ["..Tallinn", "..Tartu linn", "Maakond teadmata"]))

# ╔═╡ 2a795f4a-6bca-11eb-2b3f-b385b41d4759
@bind selectedCounty MultiSelect(counties)

# ╔═╡ 710d2e87-12df-4fdd-be2e-8826c691a86c
md"## Testing"

# ╔═╡ ed6d5c68-5b1f-11eb-134b-57fcc502469c
positive = filter(x -> x.County ∈ selectedCounty && x.ResultValue == "P", df)

# ╔═╡ 7b2f2674-5b22-11eb-225c-cbd01d9b1058
begin
	smooth_range = -7:0
	n = size(smooth_range)[1]
	kernel = OffsetArray(fill(1/n, n), smooth_range)
end;

# ╔═╡ baf09934-5b20-11eb-3283-d9b7bdfb5a90
allSelected = @chain positive begin
	groupby([:StatisticsDate])
	combine(nrow => :count)
	transform(:count => (arr -> imfilter(arr, kernel)) => :smooth_count)
end;

# ╔═╡ dc1fdc56-6bcd-11eb-32c7-f3515450573c
dates = allSelected[!, :StatisticsDate];

# ╔═╡ f9b7501b-abb7-4998-bf5a-787cba8355d6
begin
	@df allSelected plot(:StatisticsDate, [:count], ticks=:native, legend=:topleft, label="daily count")
	@df allSelected plot!(:StatisticsDate, [:smooth_count], label="7-day smooth", linewidth=2)
end

# ╔═╡ 74295a2c-6bc9-11eb-1474-9befdff2e2bf
next = @chain positive begin
	groupby([:StatisticsDate, :County])
	combine(nrow => :count)
	groupby([:County])
	transform((arr -> imfilter(arr.count, kernel) / inimarv[arr[1,:County]] * 100000))
end;

# ╔═╡ 080ecf02-5549-11eb-06cd-b10c7ed4e1e4
groups = groupby(df, [:ResultValue]);

# ╔═╡ 742c4eaa-7fe5-11eb-37e5-b96d6d7d86cd
@df next plot(:StatisticsDate, [:x1], group=:County, title="New infections per 100000", legend=:topleft, ticks=:native)

# ╔═╡ e580107b-963f-4e0b-967a-539e0e0a544d
md"""
Sidenote for all the following graphs with weeks in the x-axis. Technically what I currently count as a week is just the result of `floor(day_of_year/7)`. I used actual calendar weeks at first, but this ran into the nasty problem 2021 starting with week 52 for the first few days, which causes all sorts of problems.

Also data from the very last week is inherently not complete.
"""

# ╔═╡ 46368bc1-2079-4ad0-aa1d-677ce6dbd250
weekday_count = @chain allSelected begin
	transform(:StatisticsDate => ByRow(dayofweek) => :day)
	select([:day, :count])
	groupby(:day)
	combine(:count => sum)
end;

# ╔═╡ 29030d24-50b1-4c5b-83b1-75695910def1
weekday_completion = cumsum(weekday_count.count_sum) ./ sum(weekday_count.count_sum)

# ╔═╡ 12989f47-8e87-40df-a993-61d9f96516dd
lastdate = positive.StatisticsDate[end]

# ╔═╡ 9dcb4ec0-dc17-428f-9e7d-b6fddad8c9c4
last_week_day_count = (dayofyear(lastdate) % 7 + 1)

# ╔═╡ f15da4cb-fc03-47b4-98d6-45f3edf158f6
last_week_multiplier = 1/weekday_completion[last_week_day_count]

# ╔═╡ d49f6f5e-5cb5-11eb-14e1-45d31ab53e89
function week_multi(week_p)
	week_multiplier = zeros(size(week_p)[1])
	for i = 2:size(week_p)[1]
		prev = week_p[i-1, :count_in_week]
		current = week_p[i, :count_in_week]
		week_multiplier[i] = current/prev
	end
	week_multiplier[end] *= last_week_multiplier # add an estimate for the currently ongoing week
	week_multiplier
end;

# ╔═╡ ef39fc56-554b-11eb-3082-77349ed217d4
counts = combine(groups, nrow => :count)

# ╔═╡ c85e156e-554a-11eb-218b-df5eabe2d8ab
@df counts bar(:ResultValue, :count)

# ╔═╡ 53fd6066-4515-4641-8a7d-adc4a8e780a7
md"""
## Vaccinations
"""

# ╔═╡ 77ead592-c4a1-4821-83c1-1b311362622b
long_vac = @chain vac begin
	groupby([:StatisticsDate, :MeasurementType])
	combine([:DailyCount => sum => :DailyCount])
	unstack(_, :StatisticsDate, :MeasurementType, :DailyCount)
	rename!(["StatisticsDate", "Total?", "Completed", "InProgress"])
end

# ╔═╡ 51dbbbb2-65b2-4b30-a628-6f49fe606cdc
 hm = @chain vac begin
	groupby([:StatisticsDate, :MeasurementType])
	combine([:DailyCount => sum => :DailyCount])
	unstack(_, :StatisticsDate, :MeasurementType, :DailyCount)
	transform(Not(:StatisticsDate) => ((doses, fully, half) -> doses-fully-half) => :confusion)
end

# ╔═╡ 0fe1a173-4e9a-4b6c-9ea1-3647adc1c73e
sum(hm.confusion)

# ╔═╡ 4e5614af-8cf2-46c8-86ee-54b8e9e790d5
@df long_vac plot(:StatisticsDate, [:InProgress, :Completed], ticks=:native, label=["1st dose" "completed dose"])

# ╔═╡ 147e2cf8-5bfc-4e52-9105-e6ccd8b466da
hosp

# ╔═╡ bf9dfb10-9fe3-48d3-a299-0e8c9744ac29
merge_agegroup = Dict(
	"0-4" => "0-17",
	"5-9" => "0-17",
	"10-14" => "0-17",
	"15-19" => "0-17",
	"20-24" => "18-29",
	"25-29" => "18-29",
	"30-34" => "30-39",
	"35-39" => "30-39",
	"40-44" => "40-49",
	"45-49" => "40-49",
	"50-54" => "50-59",
	"55-59" => "50-59",
	"60-64" => "60-69",
	"65-69" => "60-69",
	"70-74" => "70-79",
	"75-79" => "70-79",
	"80-84" => "80+",
	"üle 85"=> "80+",
)

# ╔═╡ 4486854d-f57c-4c7d-9866-a50d3f31d1cd
merge_agegroup1 = Dict(
	"1-4" => "0-17",
	"0" => "0-17",
	"5-9" => "0-17",
	"10-14" => "0-17",
	"15-19" => "0-17",
	"20-24" => "18-29",
	"25-29" => "18-29",
	"30-34" => "30-39",
	"35-39" => "30-39",
	"40-44" => "40-49",
	"45-49" => "40-49",
	"50-54" => "50-59",
	"55-59" => "50-59",
	"60-64" => "60-69",
	"65-69" => "60-69",
	"70-74" => "70-79",
	"75-79" => "70-79",
	"80-84" => "80+",
	"85 ja vanemad" => "80+",
)

# ╔═╡ ff202e82-e45e-410d-9158-37f7b5de1fbc
merge_agegroup2 = Dict(
	"0-11" => "0-17",
	"12-15" => "0-17",
	"16-17" => "0-17",
	"18-29" => "18-29",
	"30-39" => "30-39",
	"40-49" => "40-49",
	"50-59" => "50-59",
	"60-69" => "60-69",
	"70-79" => "70-79",
	"80+" => "80+",
)

# ╔═╡ 83dad0ff-128d-4ca4-a9d4-447088cb577b
merge_keys = Set(keys(merge_agegroup1))

# ╔═╡ bc94dcb4-7f2a-422c-8e36-a23b1fc47fa9
@chain vac begin
	subset(:StatisticsDate => ByRow(==(Date("2021-11-10"))), :MeasurementType => ByRow(==("FullyVaccinated")))
	groupby(:AgeGroup)
	combine(:TotalCount => sum)
end

# ╔═╡ 2e6ade2d-52d4-42cc-b7ee-332b4bcf4634
md"ignored vaccinated people for missing agegroup: $(sum(vac[ismissing.(vac.AgeGroup), :].DailyCount))"

# ╔═╡ 4106d31a-8eae-40bb-9764-f3566c23d139
107316 / 112691

# ╔═╡ 6eeca282-50f6-4cb3-abbd-5c3df6c3a4e5
md"""# Predicting the future

## Vaccination Prediction
For cumulative vaccinations, I train a simple ARIMA(1,2,1) model and compare that to the goal of 70% vaccination, since that was the goal before the delta variant of covid. Right now, the 70% goal is not semantically useful, but every extra percentage helps exponentially.

Note that this comparison is not valid for vaccines that only take one shot for immunity, people who get one vaccine shot after they have already caught covid and in the future, any booster shots.

It also might be argued that it's mistreating small children (younger than 12) since they don't have any approved vaccines yet, but while their chances of dying are tiny, they can still very successfuly transmit the virus and saving lives is very much about curbing the spread of the virus.
"""

# ╔═╡ 322eecdf-56ff-42bd-986c-964c3eacd313
md"""## Testing prediction

This model is a small DSANet (Dual Self-Attention Network for Multivariate Time Series Forecasting) predicting the next three weeks. """

# ╔═╡ 0efbbf98-ee90-428a-b491-081d6dc08527
cases = df_total[:, [:DailyCases, :TotalCases, :PerPopulation]] |> Matrix{Float32};

# ╔═╡ ae921ac9-520a-4ae1-858e-d32a630eaf11
x̂ = mean(df_total.DailyCases); std_x = std(df_total.DailyCases)

# ╔═╡ e4413e75-fe3c-402a-a0ff-d199adf65ecd
begin
	poollength = 14; horizon = 14; datalength = 600;
	input, target = prepare_data(cases, poollength, datalength, horizon)

	inputsize = size(input, 1)
	local_length = 3
	n_kernels = 3
	d_model = 4
	hiddensize = 1
	n_layers = 3
	n_head = 2

	# Define the neural net
	model = DSANet(inputsize, poollength, local_length, n_kernels, d_model,
               hiddensize, n_layers, n_head)
	loss(x, y) = Flux.mse(model(x), y')
	evalcb() = Flux.throttle( () => @show(loss(input, target)), 5)
end;

# ╔═╡ b1629624-2b31-4edd-bcfa-d5b25bce77e8
count(isnan, input), count(isnan, cases), size(cases)

# ╔═╡ efbc5f35-6fc1-4c4c-b5da-4d5b070e02be
begin
	@show start_loss = loss(input, target)
	for _ = 1:20
		Flux.train!(loss, Flux.params(model),Iterators.repeated((input, target), 20), ADAMW(0.004))
		evalcb()
	end
	@show final_loss = loss(input, target)
end

# ╔═╡ de01cbd3-f8ba-4d88-9ac9-7f607b8567c5
@bind plot_prediction_i Slider(1:550, show_value=true)

# ╔═╡ 45ca89fe-412a-45fb-afd2-eb991d1fc858
cases[:, 1]

# ╔═╡ d6eedaf0-d666-4b12-8b6f-5ab473ba915a
begin
	current_input = cases[(end-13):end, :]
	last14_data = Flux.normalise(current_input, dims=1)[(end-13):end, :]
	last14 = permutedims(reshape(last14_data, (14, 3, 1, 1)), (2, 1, 3, 4))

	x_prediction = model(input)[plot_prediction_i:end] .* std_x .+ x̂
	# truth = cases[plot_prediction_i:end, 1]
	truth = target[plot_prediction_i:end] .* std_x .+ x̂

	plot([x_prediction, truth], labels=["prediction" "ground truth"])
end

# ╔═╡ 6e582294-40ff-4871-8749-db6f48761473

# ╔═╡ 256c1dc5-8f7f-4ae4-8a34-045624a3f7b4
last14

# ╔═╡ 22bfe01d-8077-4816-b5a9-94635d76120b
size(last14)

# ╔═╡ 623bba06-c217-4683-8be9-1bd366808c05
model(last14)

# ╔═╡ 0e5aa436-7553-4af2-907d-769a3c7295f0
md"# Util"

# ╔═╡ dc1ed460-4c72-4d9b-b8e5-8b975cfc247f
"""
Downloads Estonian population statistics from the official Estonian statistics API.
https://andmed.stat.ee/api/v1/et/stat/RV022U. Downloaded data is segmented by county and age group.

Returns a tuple of `(county_pop, age_pop)` with `county_pop` being a dictionary from county to population and `age_pop` being a DataFrame with agegroup and population.
"""
function download_popstats()
	params_json = """
{
  "query": [
    {
      "code": "Aasta",
      "selection": {
        "filter": "item",
        "values": [
          "2021"
        ]
      }
    },
    {
      "code": "Maakond",
      "selection": {
        "filter": "item",
        "values": [
          "37",
          "39",
          "44",
          "49",
          "51",
          "57",
          "59",
          "65",
          "67",
          "70",
          "74",
          "78",
          "82",
          "84",
          "86",
          "19"
        ]
      }
    },
    {
      "code": "Sugu",
      "selection": {
        "filter": "item",
        "values": [
          "1"
        ]
      }
    },
    {
      "code": "Vanuserühm",
      "selection": {
        "filter": "item",
        "values": [
          "1",
          "2",
          "3",
          "4",
          "5",
          "6",
          "7",
          "8",
          "9",
          "10",
          "11",
          "12",
          "13",
          "14",
          "15",
          "16",
          "17",
          "18",
          "19",
          "20",
          "21",
          "22",
          "23"
        ]
      }
    }
  ],
  "response": {
    "format": "csv"
  }
}"""
	url = "https://andmed.stat.ee/api/v1/et/stat/RV022U"
	r = HTTP.request("POST", url, ["Content-Type" => "application/json"], params_json)
	df = CSV.read(r.body, DataFrame)
	county_population = Dict(eachrow(df[!, [:Maakond, Symbol("Vanuserühmad kokku")]]))
	agegroup_population = @chain df begin
		select(:Maakond, 5:26)
		DataFrames.stack(Not(:Maakond))
		groupby(:variable)
		combine(:value => sum => :Value)
		rename(:variable => :Vanuserühm)
	end
	county_population, agegroup_population
end

# ╔═╡ 5123d73c-6c9b-11eb-0231-675c2b98d0d3
yearweek(date) = (year(date) - 2020)*52 + dayofyear(date) ÷ 7

# ╔═╡ 15f92d78-5b26-11eb-0050-27f2434d1f75
week_p = @chain positive begin
	transform(:StatisticsDate => ByRow(x -> yearweek(x)) => :week)
	groupby([:week, :County])
	combine(nrow => :count_in_week)
	groupby([:County])
	transform(week_multi)
end;

# ╔═╡ dcba4866-5b26-11eb-3e27-8fc9dd4bae77
@df week_p plot(:week, [:count_in_week], group=:County, title="New infections in week", legend=:topleft, xlabel="week")

# ╔═╡ d7a78b2e-5cb7-11eb-1967-1b3dc0a388bd
plot(week_p.week, [week_p.x1], group=week_p.County, title="New infections multiplier in county", xlabel="week")

# ╔═╡ 446f8889-a549-4fd2-b3e4-c09af9961405
no_county = @chain week_p begin
	groupby([:week])
	combine(:count_in_week => sum => :count_in_week)
	transform(week_multi)
end;

# ╔═╡ 65004f86-3c67-43a6-bb01-16df98e58148
begin
	@df no_county plot(:week, [:x1], title="Combined new infections multiplier", xlabel="week", label="infection multiplier")

	hline!([1], label="stable baseline")
end

# ╔═╡ 958bf270-0d79-43b8-ba66-9c378559c100
begin
	last_week = no_county.week[end]
	before_last_week =  no_county.count_in_week[end-1]
	last_week_estimate = no_county.count_in_week[end] * last_week_multiplier

	@df no_county plot(:week, [:count_in_week], title="Combined new infections per week", xlabel="week", label="infections", ticks=:native)
	plot!([last_week-1, last_week], [before_last_week, last_week_estimate], line=:dot, linewidth=2, label="estimate")
end

# ╔═╡ 54322a8f-2413-4ae1-9d25-57b652cc9e36
week_vac = @chain long_vac begin
	transform(:StatisticsDate => ByRow(x -> yearweek(x)) => :week)
	groupby([:week])
	combine([:Completed => sum => :Completed, :InProgress => sum => :InProgress])
end;

# ╔═╡ c9bb8186-2cd2-48f2-b579-6d304d0d5e8d
begin
	week_vac.cum_completed = cumsum(week_vac.Completed);
	week_vac.cum_inprogress = cumsum(week_vac.InProgress);
	week_vac.cum = (cumsum(week_vac.InProgress) .+ cumsum(week_vac.Completed));
end;

# ╔═╡ 766b95ae-3d6d-4e12-bd11-9416bda6f36a
begin
	@df week_vac plot(:week, [:Completed + :InProgress, :Completed, :InProgress], ticks=:native, title="Vaccine doses per week", legend=:topleft, label=["combined" "completed" "in progress"])
end

# ╔═╡ 242017a9-e5d9-4302-ac97-2424aabf8116
begin
	population = 1_329_460
	now = week_vac[end-1,:]

	doses = now.Completed + now.InProgress

	total_todo = floor(Int64, 2 * population * 0.7)
	todo = total_todo - now.cum

	p_progress = doses / (2 * population * 0.7)
	p_progress2 = doses / todo
	p_complete = now.cum / total_todo

	days2complete = floor(Int64, p_progress2^-1 * 7)
	completion_date = today() + Day(days2complete)

	"""$(round(p_complete*100))% (+$(round(p_progress*100))%)
	 The whole population will be 70% vaccinated in $days2complete days ($completion_date)
	"""
end

# ╔═╡ bd2ada71-6b46-4a33-9fb2-722dace05194
begin
	days2pcompletion = ceil(Int64, (population * 0.7 - now.cum_inprogress) / now.InProgress * 7)
	pcompletion = today() + Day(days2pcompletion)
	"Partial vaccination will be complete in $days2pcompletion days ($pcompletion)"
end

# ╔═╡ 5a25bbf6-9212-4b30-b7fe-2fa14ca02522
population / 2, now.cum, total_todo

# ╔═╡ c7343b6f-dccf-4b7b-996f-af41996537cf
week_vac

# ╔═╡ 2682f4e4-017a-4b45-aa68-28151e6877b3
begin
	@df week_vac plot(:week, [:cum, :cum_completed, :cum_inprogress], label=["all" "completed" "first dose"], ticks=:native, title="Cumulative vaccinations", legend=:topleft)
end

# ╔═╡ 5d8ac6af-1911-4583-9907-7d4f58d7863c
md"""not vaccinated: $(1_329_460 - week_vac[end, :cum_completed]) \
vaccinated: $(week_vac[end, :cum_completed])

## hospitalized
"""

# ╔═╡ f6473e6d-ebe3-4832-ba1d-01a1f1885ac6
begin
	Y = week_vac[1:end-1,:cum] |> JArray{Float64}
	Y = Matrix(Y')
end

# ╔═╡ a8d97ecc-9ec3-4821-8c87-97da92521c05
begin
	p = 1;
	d = 2;
	q = 1;
	arima_settings = ARIMASettings(Y, d, p, q);

	# Estimation
	arima_out = arima(arima_settings, NelderMead(), Optim.Options(iterations=100, f_tol=1e-2, x_tol=1e-2, g_tol=1e-2, show_trace=true, show_every=500));
end

# ╔═╡ 187448cb-1c80-4089-8f65-634fe4e3088c
begin
	week_count = 20
	fc = forecast(arima_out, week_count, arima_settings);
end

# ╔═╡ e41ecea5-77b1-43aa-aa00-01be18978510
begin
	weeks = week_vac[1:end-1,:week] |> Array{Int64,1}
	for new_week = 1:week_count
		last_week = weeks[end]
		push!(weeks, last_week + 1)
	end

	months = repeat(monthabbr.(1:12), 2)

	plot(weeks, [Y[1,:]; NaN*ones(week_count)], label="Vaccinated", ticks=:native, legend=:right, xticks = (52:4:160, months), xlabel="Month", ylabel="People")
	plot!(weeks, [NaN*ones(size(Y)[2]); fc[1,:]], label="Prediction", line=:dot)
	plot!(weeks, population*0.7*2*ones(size(weeks)), label="70% Goal")
end

# ╔═╡ d8ed8d5c-6c9b-11eb-2215-2fb6b80d30e2
countmap(week.(df[!, :StatisticsDate]))

# ╔═╡ d5b5cf78-5a87-11eb-1d85-176632857548
function mylinspace(d1::DateTime, d2::DateTime, n::Int)
    Δ = d2 - d1
    T = typeof(Δ)
    δ = T(round(Int, Dates.value(Δ)/(n - 1)))
    d2 = d1 + δ*(n - 1)
    return d1:δ:d2
end

# ╔═╡ 7f5be60a-5b23-11eb-3d24-d17513534e16
ticks = mylinspace(DateTime(2020,03,01), DateTime(2021,03,06), 30)

# ╔═╡ 620a967e-5b2f-11eb-3dc5-6108d315124b
(..)(x::AbstractDict,i...) = getindex.(Ref(x), i...)

# ╔═╡ 9e3c3714-5b2d-11eb-2a50-413ab7ed65d1
df.population = inimarv.. df[!, :County];

# ╔═╡ 48f4d8c8-37b5-4c2d-aabb-4a0590bffffb
hospitalized_age = @chain hosp begin
	groupby(:AgeGroup)
	combine(:PatientCount => sum => :PatientCount)
	dropmissing(:AgeGroup)
	transform(:AgeGroup => (ag ->  merge_agegroup .. ag) => :AgeGroup)
	groupby(:AgeGroup)
	combine(:PatientCount => sum => :PatientCount)
end;

# ╔═╡ 08e97dd2-a870-43e0-8c91-4c6b605b664f
ages = @chain ages_df begin
	subset(:Vanuserühm => ByRow(x -> x ∈ merge_keys))
	select(:Vanuserühm, :Value)
	transform(:Vanuserühm => (ag ->  merge_agegroup1 .. ag) => :AgeGroup)
	groupby([:AgeGroup])
	combine([:Value => sum => :Total])
end

# ╔═╡ daab5fb2-355d-49ca-9268-3f6e40d7b0a1
vaccinated_age = @chain vac begin
	subset(:MeasurementType => ByRow(==("FullyVaccinated")))
	dropmissing(:AgeGroup)
	transform(:AgeGroup => (ag ->  merge_agegroup2 .. ag) => :AgeGroup)
	groupby([:AgeGroup])
	combine([:DailyCount => sum => :Vaccinated])
end

# ╔═╡ 59bfbded-8aaa-4797-a588-18b2f946e57a
tested_age = @chain positive begin
	transform(:AgeGroup => (ag ->  merge_agegroup .. ag) => :AgeGroup)
	groupby(:AgeGroup)
	combine(nrow => :TestedPositive)
end;

# ╔═╡ c179a703-3955-4353-8a1e-a01b610335f1
begin
	agegroup = innerjoin(hospitalized_age, vaccinated_age, tested_age, ages, on=:AgeGroup)
	agegroup.NotVaccinated = agegroup.Total .- agegroup.Vaccinated
	agegroup.NoProtection = agegroup.NotVaccinated .- agegroup.TestedPositive
	agegroup.NewExpectedPatients = agegroup.PatientCount ./ agegroup.Total .* agegroup.NoProtection
	agegroup
end

# ╔═╡ 9131bfb2-5144-46f7-9ee7-db03290ad930
can_hospitalize = sum(agegroup.PatientCount ./ agegroup.Total .* agegroup.NoProtection + agegroup.Vaccinated * 0.00008)

# ╔═╡ b35767b7-cde8-4b4b-bac6-6aeaf17f53ff
can_hospitalize * 1.05 * 1.10 * 1.4

# ╔═╡ e68e2172-43f9-4add-80cd-5b221e5981e8
md"""
### estimated nr of people that can still be hospitalized: $(round(can_hospitalize; digits=2))

caveats for the above:
* there might be some double counting for patients in the hospital (from people who get very sick twice) (but probably not much)
* there's definitely a bunch of double counting for testing positive (from people get covid twice) (**a lot**)
* there's also a lot of double (/ triple!) counting for people who got vaccinated, and still got sick (and still went to the hospital) (**a lot a lot**)
"""

# ╔═╡ 4d7a04cc-dd16-4d01-b724-4ea42333e8fe
sum(agegroup.Total)

# ╔═╡ 61991d43-ad43-4638-ad38-b3e6d3d8cb63
value_counts(x) = sort(collect(countmap(x)), by=x -> x[2], rev=true)

# ╔═╡ Cell order:
# ╠═bb6cae46-4108-11eb-2309-3d11f13a19fa
# ╠═d8b334ae-4108-11eb-2546-99e75efc18d6
# ╠═7f15b8d8-17d9-4b89-aedb-99ff090a91ea
# ╠═e6b122ce-5548-11eb-3980-53dfd8e5897c
# ╠═eaab22ca-5b2c-11eb-1eb6-cf73054ee98e
# ╠═d35a4a83-2585-4b17-887a-1d184b187921
# ╠═9e3c3714-5b2d-11eb-2a50-413ab7ed65d1
# ╠═6fcdea12-35f3-4648-8ba0-cfd6e89547ef
# ╠═2a795f4a-6bca-11eb-2b3f-b385b41d4759
# ╟─710d2e87-12df-4fdd-be2e-8826c691a86c
# ╠═ed6d5c68-5b1f-11eb-134b-57fcc502469c
# ╠═baf09934-5b20-11eb-3283-d9b7bdfb5a90
# ╠═dc1fdc56-6bcd-11eb-32c7-f3515450573c
# ╠═f9b7501b-abb7-4998-bf5a-787cba8355d6
# ╠═74295a2c-6bc9-11eb-1474-9befdff2e2bf
# ╠═7b2f2674-5b22-11eb-225c-cbd01d9b1058
# ╠═080ecf02-5549-11eb-06cd-b10c7ed4e1e4
# ╠═742c4eaa-7fe5-11eb-37e5-b96d6d7d86cd
# ╠═15f92d78-5b26-11eb-0050-27f2434d1f75
# ╟─e580107b-963f-4e0b-967a-539e0e0a544d
# ╠═dcba4866-5b26-11eb-3e27-8fc9dd4bae77
# ╠═46368bc1-2079-4ad0-aa1d-677ce6dbd250
# ╠═29030d24-50b1-4c5b-83b1-75695910def1
# ╠═12989f47-8e87-40df-a993-61d9f96516dd
# ╠═9dcb4ec0-dc17-428f-9e7d-b6fddad8c9c4
# ╠═f15da4cb-fc03-47b4-98d6-45f3edf158f6
# ╠═d49f6f5e-5cb5-11eb-14e1-45d31ab53e89
# ╠═d7a78b2e-5cb7-11eb-1967-1b3dc0a388bd
# ╠═65004f86-3c67-43a6-bb01-16df98e58148
# ╠═958bf270-0d79-43b8-ba66-9c378559c100
# ╠═446f8889-a549-4fd2-b3e4-c09af9961405
# ╠═ef39fc56-554b-11eb-3082-77349ed217d4
# ╠═c85e156e-554a-11eb-218b-df5eabe2d8ab
# ╟─53fd6066-4515-4641-8a7d-adc4a8e780a7
# ╠═77ead592-c4a1-4821-83c1-1b311362622b
# ╠═51dbbbb2-65b2-4b30-a628-6f49fe606cdc
# ╠═0fe1a173-4e9a-4b6c-9ea1-3647adc1c73e
# ╠═4e5614af-8cf2-46c8-86ee-54b8e9e790d5
# ╠═54322a8f-2413-4ae1-9d25-57b652cc9e36
# ╠═c9bb8186-2cd2-48f2-b579-6d304d0d5e8d
# ╠═766b95ae-3d6d-4e12-bd11-9416bda6f36a
# ╠═242017a9-e5d9-4302-ac97-2424aabf8116
# ╠═bd2ada71-6b46-4a33-9fb2-722dace05194
# ╠═5a25bbf6-9212-4b30-b7fe-2fa14ca02522
# ╠═c7343b6f-dccf-4b7b-996f-af41996537cf
# ╠═2682f4e4-017a-4b45-aa68-28151e6877b3
# ╟─5d8ac6af-1911-4583-9907-7d4f58d7863c
# ╠═147e2cf8-5bfc-4e52-9105-e6ccd8b466da
# ╠═48f4d8c8-37b5-4c2d-aabb-4a0590bffffb
# ╟─bf9dfb10-9fe3-48d3-a299-0e8c9744ac29
# ╟─4486854d-f57c-4c7d-9866-a50d3f31d1cd
# ╟─ff202e82-e45e-410d-9158-37f7b5de1fbc
# ╠═83dad0ff-128d-4ca4-a9d4-447088cb577b
# ╠═08e97dd2-a870-43e0-8c91-4c6b605b664f
# ╠═daab5fb2-355d-49ca-9268-3f6e40d7b0a1
# ╠═bc94dcb4-7f2a-422c-8e36-a23b1fc47fa9
# ╟─2e6ade2d-52d4-42cc-b7ee-332b4bcf4634
# ╠═59bfbded-8aaa-4797-a588-18b2f946e57a
# ╠═c179a703-3955-4353-8a1e-a01b610335f1
# ╠═9131bfb2-5144-46f7-9ee7-db03290ad930
# ╠═b35767b7-cde8-4b4b-bac6-6aeaf17f53ff
# ╠═4d7a04cc-dd16-4d01-b724-4ea42333e8fe
# ╠═4106d31a-8eae-40bb-9764-f3566c23d139
# ╠═e68e2172-43f9-4add-80cd-5b221e5981e8
# ╟─6eeca282-50f6-4cb3-abbd-5c3df6c3a4e5
# ╠═0b484dcf-b479-46d7-8c86-d96a43ff458d
# ╠═f6473e6d-ebe3-4832-ba1d-01a1f1885ac6
# ╠═a8d97ecc-9ec3-4821-8c87-97da92521c05
# ╠═187448cb-1c80-4089-8f65-634fe4e3088c
# ╠═e41ecea5-77b1-43aa-aa00-01be18978510
# ╟─322eecdf-56ff-42bd-986c-964c3eacd313
# ╠═0efbbf98-ee90-428a-b491-081d6dc08527
# ╠═ae921ac9-520a-4ae1-858e-d32a630eaf11
# ╠═b1629624-2b31-4edd-bcfa-d5b25bce77e8
# ╠═e4413e75-fe3c-402a-a0ff-d199adf65ecd
# ╠═efbc5f35-6fc1-4c4c-b5da-4d5b070e02be
# ╠═de01cbd3-f8ba-4d88-9ac9-7f607b8567c5
# ╠═45ca89fe-412a-45fb-afd2-eb991d1fc858
# ╠═d6eedaf0-d666-4b12-8b6f-5ab473ba915a
# ╠═6e582294-40ff-4871-8749-db6f48761473
# ╠═256c1dc5-8f7f-4ae4-8a34-045624a3f7b4
# ╠═22bfe01d-8077-4816-b5a9-94635d76120b
# ╠═623bba06-c217-4683-8be9-1bd366808c05
# ╟─0e5aa436-7553-4af2-907d-769a3c7295f0
# ╠═dc1ed460-4c72-4d9b-b8e5-8b975cfc247f
# ╠═7f5be60a-5b23-11eb-3d24-d17513534e16
# ╠═5123d73c-6c9b-11eb-0231-675c2b98d0d3
# ╠═d8ed8d5c-6c9b-11eb-2215-2fb6b80d30e2
# ╠═d5b5cf78-5a87-11eb-1d85-176632857548
# ╠═620a967e-5b2f-11eb-3dc5-6108d315124b
# ╠═61991d43-ad43-4638-ad38-b3e6d3d8cb63


text:
using Test
import Pluto
import Pluto: update_save_run!, WorkspaceManager, ClientSession, ServerSession, Notebook, Cell

import UUIDs: uuid1

function get_unique_short_id()
    string(uuid1())[1:8]
end

function stringify_keys(d::Dict)
    Dict(string(k) => stringify_keys(v) for (k, v) in d)
end
stringify_keys(x::Any) = x

@testset "Communication protocol" begin

    @testset "Functionality sweep" begin
        buffer = IOBuffer()
        client = ClientSession(:buffery, buffer)

        🍭 = ServerSession()
        # 🍭.connected_clients[client.id] = client

        notebook = Notebook([Cell("")])
        🍭.notebooks[notebook.notebook_id] = notebook
        update_save_run!(🍭, notebook, notebook.cells)
        # client.connected_notebook = notebook

        function send(type, body, metadata)
            request_id = get_unique_short_id()
            Pluto.process_ws_message(🍭, Dict(
                "type" => string(type),
                "client_id" => string(client.id),
                "request_id" => request_id,
                "body" => body,
                metadata...
            ) |> stringify_keys, client.stream)
        end

        n = notebook.notebook_id |> string
        c(cell) = cell.cell_id |> string

        @test_nowarn send(:connect, Dict(), Dict(:notebook_id => n))

        @test_nowarn send(:add_cell, Dict(:index => 0), Dict(:notebook_id => n))
        send(:add_cell, Dict(:index => 0), Dict(:notebook_id => n))
        send(:add_cell, Dict(:index => 0), Dict(:notebook_id => n))
        @test length(notebook.cells) == 4
        @test_nowarn send(:set_input, Dict(:code => "1 + 2"), Dict(:notebook_id => n, :cell_id => c(notebook.cells[1])))
        @test_nowarn send(:run_multiple_cells, Dict(:cells => [string(c.cell_id) for c in notebook.cells[1:2]]), Dict(:notebook_id => n))
        @test_nowarn send(:set_bond, Dict(:sym => "x", :val => 9, :is_first_value => true), Dict(:notebook_id => n))
        @test_nowarn send(:change_cell, Dict(:code => "1+1"), Dict(:notebook_id => n, :cell_id => c(notebook.cells[3])))
        @test_nowarn send(:delete_cell, Dict(), Dict(:notebook_id => n, :cell_id => c(notebook.cells[4])))

        @test_nowarn send(:move_multiple_cells, Dict(:cells => [c(notebook.cells[3])], :index => 1), Dict(:notebook_id => n))
        @test_nowarn send(:fold_cell, Dict(:folded => true), Dict(:notebook_id => n, :cell_id => c(notebook.cells[1])))
        @test_nowarn send(:get_input, Dict(), Dict(:notebook_id => n, :cell_id => c(notebook.cells[1])))
        @test_nowarn send(:get_output, Dict(), Dict(:notebook_id => n, :cell_id => c(notebook.cells[1])))
        @test_nowarn send(:get_all_cells, Dict(), Dict(:notebook_id => n))
        @test_nowarn send(:get_all_notebooks, Dict(), Dict(:notebook_id => n))
        @test_nowarn send(:get_all_notebooks, Dict(), Dict())

        @test_nowarn send(:move_notebook_file, Dict(:path => tempname()), Dict(:notebook_id => n))

        # TODO: we need to wait for all above command to finish before we can do this:
        # send(:shutdown_notebook, Dict(:keep_in_session => false), Dict(:notebook_id => n))
    end

    @testset "Docs" begin
        @test occursin("square root", Pluto.PlutoRunner.doc_fetcher("sqrt")[1])
        @test occursin("square root", Pluto.PlutoRunner.doc_fetcher("Base.sqrt")[1])
        @test occursin("No documentation found", Pluto.PlutoRunner.doc_fetcher("Base.findmeta")[1])
    end
end

# TODO: test returned data


text:
# Autogenerated wrapper script for Vulkan_Headers_jll for any
export vk_xml

JLLWrappers.@declare_file_product(vk_xml)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_file_product(
        vk_xml,
        "share/vulkan/registry/vk.xml",
    )

end  # __init__()


text:
function Base.show(io::IO, ::MIME"text/vnd.graphviz", dag::DAG)
    idref = Ref(0)
    newid() = idref[] += 1

    idmap = IdDict{DAG,Int}()
    idfor(dag) = _get!(newid, idmap, dag)

    seen = IdDict{DAG,Nothing}()
    function isseen(dag)
        yes = Ref(true)
        _get!(seen, dag) do
            yes[] = false
            nothing
        end
        return yes[]
    end

    prln(args...) = println(io, "  ", args...)

    function draw(dag::SequentialNode, id)
        isseen(dag) && return
        txt = labeltext(dag.data)
        prln(id, " [shape=record, label=\"$txt\"];")
        kid = idfor(dag.continuation)
        prln(id, " -> ", kid, ";")
        draw(dag.continuation, kid)
    end

    function draw(dag::SpawnNode, id)
        isseen(dag) && return
        sid = idfor(dag.sync)
        did = idfor(dag.detach)
        kid = idfor(dag.continuation)
        prln(id, " [shape=egg, label=\"spawn (in $sid)\"];")
        prln(id, " -> ", did, " [label = Det];")
        prln(id, " -> ", kid, " [label = Cont];")
        draw(dag.detach, did)
        draw(dag.continuation, kid)
    end

    function draw(dag::SyncNode, id)
        isseen(dag) && return
        k = dag.continuation
        if k isa DAG
            prln(id, " [shape=diamond, label=\"sync $id\"];")
            kid = idfor(k)
            prln(id, " -> ", kid, ";")
            draw(k, kid)
        end
    end

    println(io, "digraph DAG {")
    draw(dag, newid())
    print(io, "}")
end

function labeltext(data::DurationNS)
    if data.value > 1e9
        x = data.value / 1e9
        u = "s"
    elseif data.value > 1e6
        x = data.value / 1e6
        u = "ms"
    elseif data.value > 1e3
        x = data.value / 1e3
        u = "μs"
    else
        x = data.value
        u = "ns"
    end
    i, f = divrem(x, 1)
    i = floor(Int, i)
    if i > 10
        return "$i $u"
    else
        j = floor(Int, 10 * f)
        return "$i.$j $u"
    end
end

ShowGraphviz.@deriveall DAG


text:
module TestDay11

import AOC2021.Day11.part1, AOC2021.Day11.part2
using Test

test_input = """
5483143223
2745854711
5264556173
6141336146
6357385478
4167524645
2176841721
6882881134
4846848554
5283751526"""

@testset "Day 11" begin
    @testset "part 1" begin
        @test part1(test_input) == 1656
    end
    @testset "part 2" begin
        @test part2(test_input) == 195
    end
end

end # module


text:
# Append session to options.
# Given that options may be nothing, always use the return value as the new value for options.
function _join(options::Union{Nothing, BSON}, session::Session) :: BSON
    result = options == nothing ? Mongoc.BSON() : options
    err = BSONError()
    ok = mongoc_client_session_append(session.handle, result.handle, err)
    if !ok
        error("$err")
    end
    return result
end

function set_casual_consistency!(session_options::SessionOptions, casual_consistency::Bool)
    mongoc_session_opts_set_causal_consistency(session_options.handle, casual_consistency)
    nothing
end

function get_casual_consistency(session_options::SessionOptions) :: Bool
    mongoc_session_opts_get_causal_consistency(session_options.handle)
end

get_session(database_session::DatabaseSession) = database_session.session
get_session(collection_session::CollectionSession) = get_session(collection_session.database_session)

DatabaseSession(session::Session, db_name::String) = DatabaseSession(Database(session.client, db_name), session)
CollectionSession(database_session::DatabaseSession, collection_name::String) = CollectionSession(database_session, Collection(database_session.database, collection_name))

Base.getindex(session::Session, db_name::String) = DatabaseSession(session, db_name)
Base.getindex(database_session::DatabaseSession, collection_name::String) = CollectionSession(database_session, collection_name)

#
# Overload
#
function find_collections(database::DatabaseSession; options::Union{Nothing, BSON}=nothing)
    options_with_session = _join(options, get_session(database))
    find_collections(database.database; options=options_with_session)
end

function get_collection_names(database::DatabaseSession; options::Union{Nothing, BSON}=nothing)
    options_with_session = _join(options, get_session(database))
    get_collection_names(database.database, options=options_with_session)
end

function insert_one(collection::CollectionSession, document::BSON; options::Union{Nothing, BSON}=nothing)
    options_with_session = _join(options, get_session(collection))
    insert_one(collection.collection, document, options=options_with_session)
end

function delete_one(collection::CollectionSession, selector::BSON; options::Union{Nothing, BSON}=nothing)
    options_with_session = _join(options, get_session(collection))
    delete_one(collection.collection, selector, options=options_with_session)
end

function delete_many(collection::CollectionSession, selector::BSON; options::Union{Nothing, BSON}=nothing)
    options_with_session = _join(options, get_session(collection))
    delete_many(collection.collection, selector, options=options_with_session)
end

function update_one(collection::CollectionSession, selector::BSON, update::BSON; options::Union{Nothing, BSON}=nothing)
    options_with_session = _join(options, get_session(collection))
    update_one(collection.collection, selector, update, options=options_with_session)
end

function update_many(collection::CollectionSession, selector::BSON, update::BSON; options::Union{Nothing, BSON}=nothing)
    options_with_session = _join(options, get_session(collection))
    update_many(collection.collection, selector, update, options=options_with_session)
end

function insert_many(collection::CollectionSession, documents::Vector{BSON}; bulk_options::Union{Nothing, BSON}=nothing, insert_options::Union{Nothing, BSON}=nothing)
    bulk_options_with_session = _join(bulk_options, get_session(collection))
    insert_many(collection.collection, documents, bulk_options=bulk_options_with_session, insert_options=insert_options)
end

function find(collection::CollectionSession, bson_filter::BSON=BSON(); options::Union{Nothing, BSON}=nothing) :: Cursor
    options_with_session = _join(options, get_session(collection))
    find(collection.collection, bson_filter, options=options_with_session)
end

function count_documents(collection::CollectionSession, bson_filter::BSON=BSON(); options::Union{Nothing, BSON}=nothing)
    options_with_session = _join(options, get_session(collection))
    count_documents(collection.collection, bson_filter, options=options_with_session)
end

function find_one(collection::CollectionSession, bson_filter::BSON=BSON(); options::Union{Nothing, BSON}=nothing) :: Union{Nothing, BSON}
    options_with_session = _join(options, get_session(collection))
    find_one(collection.collection, bson_filter, options=options_with_session)
end

function aggregate(collection::CollectionSession, bson_pipeline::BSON; flags::QueryFlags=QUERY_FLAG_NONE, options::Union{Nothing, BSON}=nothing) :: Cursor
    options_with_session = _join(options, get_session(collection))
    aggregate(collection.collection, bson_pipeline, flags=flags, options=options_with_session)
end

#
# Transaction
#

function start_transaction!(session::Session)
    err = BSONError()
    ok = mongoc_client_session_start_transaction(session.handle, C_NULL, err)
    if !ok
        error("$err")
    end
    nothing
end

function abort_transaction!(session::Session)
    err = BSONError()
    ok = mongoc_client_session_abort_transaction(session.handle, err)
    if !ok
        error("$err")
    end
    nothing
end

function commit_transaction!(session::Session) :: BSON
    reply = BSON()
    err = BSONError()
    ok = mongoc_client_session_commit_transaction(session.handle, reply.handle, err)
    if !ok
        error("$err")
    end
    return reply
end

in_transaction(session::Session) :: Bool = mongoc_client_session_in_transaction(session.handle)

#
# High-level API
#

"""
    transaction(f::Function, client::Client; session_options::SessionOptions=SessionOptions())

Use *do-syntax* to execute a transaction.

Transaction will be commited automatically. If an error occurs, the transaction is aborted.

The `session` parameter should be treated the same way as a `Client`: from a `session` you get a `database`,
and a `collection` that are bound to the session.

```julia
Mongoc.transaction(client) do session
    database = session["my_database"]
    collection = database["my_collection"]
    new_item = Mongoc.BSON()
    new_item["inserted"] = true
    push!(collection, new_item)
end
```
"""
function transaction(f::Function, client::Client; session_options::SessionOptions=SessionOptions())
    local result::Union{Nothing, BSON} = nothing
    local aborted::Bool = false
    session = Session(client, options=session_options)
    start_transaction!(session)

    try
        f(session)
    catch
        abort_transaction!(session)
        aborted = true
        rethrow()
    finally
        if !aborted
            result = commit_transaction!(session)
        end
    end

    return result
end


text:
# Optimization for
#
#  g = lazy_map( linear_combination, cell_to_i_to_val,cell_to_i_to_f)
#  lazy_map(evaluate,g,cell_to_x)
#
function lazy_map(
  ::typeof(evaluate), a::LazyArray{<:Fill{typeof(linear_combination)}}, x::AbstractArray)

  i_to_values = a.args[1]
  i_to_basis = a.args[2]
  i_to_basis_x = lazy_map(evaluate,i_to_basis,x)
  lazy_map(LinearCombinationMap(:),i_to_values,i_to_basis_x)
end

# We always keep the parent when transposing (needed for some optimizations below)

function lazy_map(k::typeof(transpose),f::AbstractArray)
  fi = testitem(f)
  T = return_type(k,fi)
  s = size(f)
  LazyArray(T,Fill(k,s),f)
end

# Optimization for
#
#  g = lazy_map(transpose,cell_to_i_to_f)
#  lazy_map(evaluate,g)
#
function lazy_map(
  ::typeof(evaluate), a::LazyArray{<:Fill{typeof(transpose)}}, x::AbstractArray)

  i_to_basis_x = lazy_map(evaluate,a.args[1],x)
  lazy_map(TransposeMap(),i_to_basis_x)
end

# Optimization for
#
#  g = lazy_map(∘,cell_to_f,cell_to_h)
#  lazy_map(evaluate,g)
#
function lazy_map(
  ::typeof(evaluate), a::LazyArray{<:Fill{typeof(∘)}}, x::AbstractArray)

  f = a.args[1]
  g = a.args[2]
  gx = lazy_map(evaluate,g,x)
  fx = lazy_map(evaluate,f,gx)
  fx
end

# Optimization for
#
#  g = lazy_map(Broadcasting(∘),cell_to_i_to_f,cell_to_h)
#  lazy_map(evaluate,g)
#
function lazy_map(
  ::typeof(evaluate), a::LazyArray{<:Fill{Broadcasting{typeof(∘)}}}, x::AbstractArray)

  f = a.args[1]
  g = a.args[2]
  gx = lazy_map(evaluate,g,x)
  fx = lazy_map(evaluate,f,gx)
  fx
end

# Optimization for
#
#  g = lazy_map(Operation(+),cell_to_f,cell_to_h)
#  lazy_map(evaluate,g)
#
function lazy_map(
  ::typeof(evaluate),
  a::LazyArray{<:Fill{<:Operation}},
  x::AbstractArray)

  fx = map( fi->lazy_map(evaluate,fi,x), a.args)
  op = a.maps.value.op
  lazy_map( Broadcasting(op), fx...)
end

# Optimization for
#
#  g = lazy_map(Broadcasting(Operation(+)),cell_to_f,cell_to_h)
#  lazy_map(evaluate,g)
#
function lazy_map(
  ::typeof(evaluate), a::LazyArray{<:Fill{<:Broadcasting{<:Operation}}}, x::AbstractArray)

  fx = map( fi->lazy_map(evaluate,fi,x), a.args)
  op = a.maps.value.f.op
  lazy_map(BroadcastingFieldOpMap(op),fx...)
end

# Optimization for
#
#  g = lazy_map( linear_combination, cell_to_i_to_val, cell_to_i_to_f)
#  lazy_map(gradient,g)
#
function lazy_map(
  ::typeof(gradient), a::LazyArray{<:Fill{typeof(linear_combination)}})

  i_to_basis = lazy_map(Broadcasting(∇),a.args[2])
  i_to_values = a.args[1]
  lazy_map(linear_combination,i_to_values,i_to_basis)
end

# Optimization for
#
#  g = lazy_map( linear_combination, cell_to_i_to_val, cell_to_i_to_f)
#  lazy_map(Broadcasting(gradient),g)
#
function lazy_map(
  k::Broadcasting{typeof(∇)}, a::LazyArray{<:Fill{typeof(linear_combination)}})

  i_to_basis = lazy_map(k,a.args[2])
  i_to_values = a.args[1]
  lazy_map(linear_combination,i_to_values,i_to_basis)
end

function lazy_map(
  k::Broadcasting{typeof(∇∇)}, a::LazyArray{<:Fill{typeof(linear_combination)}})

  i_to_basis = lazy_map(k,a.args[2])
  i_to_values = a.args[1]
  lazy_map(linear_combination,i_to_values,i_to_basis)
end

# Optimization for
#
#  g = lazy_map(transpose,cell_to_i_to_f)
#  lazy_map(Broadcasting(gradient),g)
#
function lazy_map(
  k::Broadcasting{typeof(∇)}, a::LazyArray{<:Fill{typeof(transpose)}})

  i_to_basis = lazy_map(k,a.args[1])
  lazy_map( transpose, i_to_basis)
end

function lazy_map(
  k::Broadcasting{typeof(∇∇)}, a::LazyArray{<:Fill{typeof(transpose)}})

  i_to_basis = lazy_map(k,a.args[1])
  lazy_map( transpose, i_to_basis)
end

# Gradient rules
for op in (:+,:-)
  @eval begin

    function lazy_map(
      ::typeof(gradient), a::LazyArray{<:Fill{Operation{typeof($op)}}})

      f = a.args
      g = map(i->lazy_map(gradient,i),f)
      lazy_map(Operation($op),g...)
    end

    function lazy_map(
      ::Broadcasting{typeof(gradient)}, a::LazyArray{<:Fill{Broadcasting{Operation{typeof($op)}}}})

      f = a.args
      g = map(i->lazy_map(Broadcasting(∇),i),f)
      lazy_map(Broadcasting(Operation($op)),g...)
    end

  end
end

for op in (:*,:⋅,:⊙,:⊗)
  @eval begin

    function lazy_map(
      ::typeof(gradient), a::LazyArray{<:Fill{Operation{typeof($op)}}})

      f = a.args
      @notimplementedif length(f) != 2
      g = map(i->lazy_map(gradient,i),f)
      k(F1,F2,G1,G2) = product_rule($op,F1,F2,G1,G2)
      lazy_map(Operation(k),f...,g...)
    end

    function lazy_map(
      ::Broadcasting{typeof(gradient)}, a::LazyArray{<:Fill{Broadcasting{Operation{typeof($op)}}}})

      f = a.args
      @notimplementedif length(f) != 2
      g = map(i->lazy_map(Broadcasting(∇),i),f)
      k(F1,F2,G1,G2) = product_rule($op,F1,F2,G1,G2)
      lazy_map(Broadcasting(Operation(k)),f...,g...)
    end

  end
end

function lazy_map(
  ::Broadcasting{typeof(gradient)}, a::LazyArray{<:Fill{Broadcasting{typeof(∘)}}})

  f = a.args[1]
  g = a.args[2]
  ∇f = lazy_map(Broadcasting(∇),f)
  h = lazy_map(Broadcasting(∘),∇f,g)
  ∇g = lazy_map(Broadcasting(∇),g)

  lazy_map(Broadcasting(Operation(⋅)),∇g,h)
end

function lazy_map(
  ::typeof(gradient), a::LazyArray{<:Fill{typeof(∘)}})

  f = a.args[1]
  g = a.args[2]
  ∇f = lazy_map(∇,f)
  h = lazy_map(∘,∇f,g)
  ∇g = lazy_map(∇,g)

  lazy_map(Operation(⋅),∇g,h)
end

# Integration

function lazy_map(
  ::typeof(integrate),f::AbstractArray,x::AbstractArray,w::AbstractArray)

  fx = lazy_map(evaluate,f,x)
  lazy_map(IntegrationMap(),fx,w)
end

function lazy_map(
  ::typeof(integrate),f::AbstractArray,x::AbstractArray,w::AbstractArray,j::AbstractArray)

  fx = lazy_map(evaluate,f,x)
  jx = lazy_map(evaluate,j,x)
  lazy_map(IntegrationMap(),fx,w,jx)
end

# Pushing the gradient

function lazy_map(k::typeof(push_∇),cell_∇a::AbstractArray,cell_map::AbstractArray)
  lazy_map(Broadcasting(push_∇),cell_∇a,cell_map)
end

function lazy_map(k::Broadcasting{typeof(push_∇)},cell_∇a::AbstractArray,cell_map::AbstractArray)
  cell_Jt = lazy_map(∇,cell_map)
  cell_invJt = lazy_map(Operation(pinvJt),cell_Jt)
  lazy_map(Broadcasting(Operation(⋅)),cell_invJt,cell_∇a)
end

for op in (:push_∇,:push_∇∇)
  @eval begin
    function lazy_map(
      k::Broadcasting{typeof($op)},
      cell_∇at::LazyArray{<:Fill{typeof(transpose)}},
      cell_map::AbstractArray)

      cell_∇a = cell_∇at.args[1]
      cell_∇b = lazy_map(k,cell_∇a,cell_map)
      cell_∇bt = lazy_map(transpose,cell_∇b)
      cell_∇bt
    end
  end
end

# Composing by the identity

function lazy_map(
  k::Broadcasting{typeof(∘)},
  ::Type{T},
  a::AbstractArray,
  b::Fill{<:GenericField{typeof(identity)}}) where T
  @assert length(a) == length(b)
  a
end

function lazy_map(
  k::typeof(∘),
  ::Type{T},
  a::AbstractArray,
  b::Fill{<:GenericField{typeof(identity)}}) where T
  @assert length(a) == length(b)
  a
end

# Memoization

struct MemoArray{T,N,A} <: AbstractArray{T,N}
  parent::A
  memo::Dict{Any,Any}
  function MemoArray(parent::AbstractArray{T,N}) where {T,N}
    parent
    #A = typeof(parent)
    #memo = Dict()
    #new{T,N,A}(parent,memo)
  end
end

MemoArray(a) = a

# Do not wrap twice.
MemoArray(parent::MemoArray) = parent

function get_children(n::TreeNode, a::MemoArray)
   (similar_tree_node(n,a.parent),)
end

Base.size(a::MemoArray) = size(a.parent)
Base.axes(a::MemoArray) = axes(a.parent)
Base.IndexStyle(::Type{MemoArray{T,N,A}}) where {T,N,A} = IndexStyle(A)
Base.getindex(a::MemoArray,i::Integer) = a.parent[i]
Base.getindex(a::MemoArray{T,N},i::Vararg{Integer,N}) where {T,N} = a.parent[i...]
Arrays.array_cache(a::MemoArray) = array_cache(a.parent)
Arrays.getindex!(cache,a::MemoArray,i::Integer) = getindex!(cache,a.parent,i)
Arrays.getindex!(cache,a::MemoArray{T,N},i::Vararg{Integer,N}) where {T,N} = getindex!(cache,a.parent,i...)
Arrays.testitem(a::MemoArray) = testitem(a.parent)
Arrays.get_array(a::MemoArray) = get_array(a.parent)

function lazy_map(::typeof(evaluate),a::MemoArray,x::AbstractArray{<:Point})
  key = (:evaluate,objectid(x))
  if ! haskey(a.memo,key)
    a.memo[key] = lazy_map(evaluate,a.parent,x)
  end
  a.memo[key]
end

function lazy_map(::typeof(evaluate),a::MemoArray,x::AbstractArray{<:AbstractArray{<:Point}})
  key = (:evaluate,objectid(x))
  if ! haskey(a.memo,key)
    a.memo[key] = lazy_map(evaluate,a.parent,x)
  end
  a.memo[key]
end

function lazy_map(k::typeof(∇),a::MemoArray)
  lazy_map(Broadcasting(∇),a)
end

function lazy_map(k::Broadcasting{typeof(∇)},a::MemoArray)
  key = :gradient
  if ! haskey(a.memo,key)
    a.memo[key] = MemoArray(lazy_map(k,a.parent))
  end
  a.memo[key]
end

function lazy_map(
  k::Broadcasting{typeof(push_∇)},
  cell_∇a::MemoArray,
  cell_map::AbstractArray)
  key = (:push_gradient,objectid(cell_map))
  if ! haskey(cell_∇a.memo,key)
    cell_∇a.memo[key] = MemoArray(lazy_map(k,cell_∇a.parent,cell_map))
  end
  cell_∇a.memo[key]
end

function lazy_map(k::typeof(axes),a::MemoArray)
  lazy_map(k,a.parent)
end

function lazy_map(k::Reindex{<:MemoArray},::Type{T}, j_to_i::AbstractArray) where T
  # Commenting in order to fix issue #614
  #key = (:reindex,objectid(j_to_i))
  #if ! haskey(k.values.memo,key)
  #  i_to_v = k.values.parent
  #  j_to_v = lazy_map(Reindex(i_to_v),T,j_to_i)
  #  k.values.memo[key] = MemoArray(j_to_v)
  #end
  #k.values.memo[key]
  i_to_v = k.values.parent
  j_to_v = lazy_map(Reindex(i_to_v),T,j_to_i)
  MemoArray(j_to_v)
end

function lazy_map(k::PosNegReindex{<:MemoArray,<:MemoArray},::Type{T},i_to_iposneg::AbstractArray) where T
  values_pos = k.values_pos.parent
  values_neg = k.values_neg.parent
  r = lazy_map(PosNegReindex(values_pos,values_neg),i_to_iposneg)
  MemoArray(r)
end

function lazy_map(
  k::typeof(∘), a::MemoArray,b::AbstractArray{<:Field})
  lazy_map(Broadcasting(∘),a,b)
end

function lazy_map(
  k::Broadcasting{typeof(∘)}, a::MemoArray,b::AbstractArray{<:Field})
  key = (:compose,objectid(b))
  if ! haskey(a.memo,key)
    a.memo[key] = MemoArray(lazy_map(k,a.parent,b))
  end
  a.memo[key]
end

function lazy_map(
  k::typeof(transpose), a::MemoArray)
  key = :transpose
  if ! haskey(a.memo,key)
    a.memo[key] = MemoArray(lazy_map(k,a.parent))
  end
  a.memo[key]
end


text:
# ------------------------------------------------------------------------------------------
# # Julia Parallelism Basics
# ------------------------------------------------------------------------------------------

addprocs(2)

nworkers()

workers()

r = remotecall(rand, 2, 2, 2)

 fetch(r)

s = @spawnat 3 1+fetch(r)

fetch(s)

remotecall_fetch(getindex, 2, r, 1, 1)

r = @spawn rand(2,2)

s = @spawn 1+fetch(r)

function rand2(dims...)
    return 2*rand(dims...)
end

rand2(2,2)

r2 = @spawn rand2(2,2)

fetch(r2)

@everywhere function rand2(dims...)
    return 2*rand(dims...)
end

r2 = @spawn rand2(2,2)
fetch(r2)

# ------------------------------------------------------------------------------------------
# ## Data Movement
# ------------------------------------------------------------------------------------------

@time begin
    A = rand(1000,1000)
    Bref = @spawn A^2
    fetch(Bref)
end

@time begin
    Bref = @spawn rand(1000,1000)^2
    fetch(Bref)
end

# ------------------------------------------------------------------------------------------
# ## Shared Arrays
#
# Shared Arrays are created by mapping the same region in memory to different processes.
# ------------------------------------------------------------------------------------------

s = SharedArray{Float64}(100,100)

localindexes(s)

fetch(@spawnat 2 localindexes(s))

fetch(@spawnat 3 localindexes(s))

for i in workers()
    @spawnat i s[localindexes(s)] = myid()
end

s

fetch(@spawnat 2 s[100,100])

fetch(@spawnat 3 s[1,1])

max_stars_repo_name: grenkoca/JuliaTutorials

max_issues_repo_name: grenkoca/JuliaTutorials

max_forks_repo_name: grenkoca/JuliaTutorials


text:
"""
    GaussianProcesses

Includes all useful functions for applying GPR to T and wT profiles from Oceananigans.jl simulations.
Uses ProfileData struct to store data and GP struct for performing GPR on the data in ProfileData object.
"""
module GaussianProcesses

export
    # Gaussian process kernels
    Kernel, SquaredExponentialI, RationalQuadraticI, Matern12I, Matern32I, Matern52I, kernel_function, error_per_gamma,

    # Distance functions
    euclidean_distance, derivative_distance, antiderivative_distance,

    # Gaussian process functions
    model_output, uncertainty, compute_kernel_matrix, mean_log_marginal_loss,
    GPmodel, best_kernel, gp_model, get_kernel

using OceanParameterizations.DataWrangling
using Flux

include("kernels.jl")
include("distances.jl")
include("gaussian_process.jl")

include("../predict.jl")
mse(x::Tuple{Array{Float64,2}, Array{Float64,2}}) = Flux.mse(x[1], x[2])

function best_kernel(𝒱; logγ_range=-1.0:0.1:2.0)
    function m(𝒱, kernel)
        𝒢 = GPmodel(𝒱; kernel=kernel)
        f(x) = model_output(x, 𝒢)
        return f
    end

    best_kernel = nothing
    best_mse = Inf
    for k=1:4, logγ=logγ_range
        kernel = get_kernel(k, logγ, 0.0, euclidean_distance)
        model = m(𝒱, kernel)
        error = mse(predict(𝒱, model))

        if error < best_mse
            best_kernel = kernel
        end
    end
    return best_kernel
end

"""
Sweeps over the log-length scale values in `logγ_range` and returns a vector
with the GP prediction error at the index corresponding to the length scale value.

#Arguments
- `k`::Integer is the integer corresponding to the kernel function (1 to 5)
- `𝒱`::FluxData is the object containing the (predictor, target) pairs
- `logγ_range` is the range of log(γ) hyperparameter values to sweep over
"""
function error_per_gamma(𝒱train, 𝒱test, k; logγ_range=-2.0:0.1:1.0)
    # model trained on training data
    function m(kernel)
        𝒢 = GPmodel(𝒱train; kernel=kernel)
        f(x) = model_output(x, 𝒢)
        return f
    end

    errors = zeros(length(logγ_range))
    for (i, logγ) in enumerate(logγ_range)
        model = m(get_kernel(k, logγ, 0.0, euclidean_distance))
        errors[i] = mse(predict(𝒱test, model))
    end
    return errors
end

function gp_model(𝒱, kernel)
    function m(𝒱, kernel)
        𝒢 = GPmodel(𝒱; kernel=kernel)
        f(x) = model_output(x, GPmodel(𝒱; kernel=kernel))
        return f
    end

    if isnothing(kernel)
        best_kernel = nothing
        best_mse = Inf

        for k=1:4, logγ=logγ_range
            kernel = get_kernel(k, logγ, 0.0, euclidean_distance)
            model = m(𝒱, kernel)
            error = mse(predict(𝒱, model))

            if error < best_mse
                best_kernel = kernel
            end
        end
        return m(𝒱, best_kernel)
    else
        return m(𝒱, kernel)
    end
end

"""
    get_kernel(kernel_id::Int64, logγ, logσ, distance; logα=0.0)

# Description
Returns a Kernel object with the specified parameters.

# Arguments
- `kernel_id`: (Int64). Identification number for the kernel type (see kernel options)
- `logγ`: (Float64). Log(length scale) parameter.
- `logσ`: (Float64). Log(signal variance) parameter.
- `distance`: (Function). Distance function to use in the kernel.

# Keyword Arguments
- `logα`: (Float64). Log(α) parameter if kernel_id = 5.

# kernel options
    1   =>  "Squared exponential kernel:  k(x,x') = σ * exp( - d(x,x')² / 2γ² )",
    2   =>  "Matérn with ʋ=1/2:           k(x,x') = σ * exp( - d(x,x') / γ )",
    3   =>  "Matérn with ʋ=3/2:           k(x,x') = σ * (1+c) * exp(-√(3)*d(x,x'))/γ)",
    4   =>  "Matérn with ʋ=5/2:           k(x,x') = σ * ( 1 + √(5)*d(x,x'))/γ + 5*d(x,x')²/(3*γ^2) ) * exp(-√(5)*d(x,x'))/γ)",
    5   =>  "Rational quadratic kernel:   k(x,x') = σ * (1+(x-x')'*(x-x')/(2*α*(γ²))^(-α)",
"""
function get_kernel(kernel_id::Int64, args...)
    function g(x)
        if x isa Number
            return 10^x
        end
        return x
    end

    args = g.(args)

    kernel_id == 1 && return SquaredExponentialI(args...)
    kernel_id == 2 && return Matern12I(args...)
    kernel_id == 3 && return Matern32I(args...)
    kernel_id == 4 && return Matern52I(args...)
    kernel_id == 5 && return RationalQuadraticI(args...)
    kernel_id == 6 && return SpectralMixtureProductI(args...)

    throw(error("Invalid kernel_id $kernel_id"))
end

end #module


text:
testnx(g, nx) = isapprox(g.nx, nx)
testny(g, ny) = isapprox(g.ny, ny)
testnz(g, nz) = isapprox(g.nz, nz)

# Physical grid tests
function testdx(dev, g::Union{OneDGrid{T}, TwoDGrid{T}, ThreeDGrid{T}}) where T
  dxgrid = @. g.x[2:end] - g.x[1:end-1]
  dxones = ArrayType(dev)(g.dx*ones(T, size(dxgrid)))

  return isapprox(dxgrid, dxones)
end

function testdy(dev, g::Union{TwoDGrid{T}, ThreeDGrid{T}}) where T
  dygrid = @. g.y[2:end] - g.y[1:end-1]
  dyones = ArrayType(dev)(g.dy*ones(T, size(dygrid)))

  return isapprox(dygrid, dyones)
end

function testdz(dev, g::ThreeDGrid{T}) where T
  dzgrid = @. g.z[2:end] - g.z[1:end-1]
  dzones = ArrayType(dev)(g.dz*ones(T, size(dzgrid)))

  return isapprox(dzgrid, dzones)
end

testx(g) = isapprox(g.x[end]-g.x[1], g.Lx-g.dx)
testy(g) = isapprox(g.y[end]-g.y[1], g.Ly-g.dy)
testz(g) = isapprox(g.z[end]-g.z[1], g.Lz-g.dz)

testdk(g::Union{OneDGrid, TwoDGrid, ThreeDGrid}) = CUDA.@allowscalar isapprox(g.k[2], 2π/g.Lx)
testdl(g::Union{TwoDGrid, ThreeDGrid}) = CUDA.@allowscalar isapprox(g.l[2], 2π/g.Ly)
testdm(g::ThreeDGrid) = CUDA.@allowscalar isapprox(g.m[2], 2π/g.Lz)

# Test proper arrangement of fft wavenumbers. k = [0:nx/2, -nx/2+1:-1].
flippednegatives(k, mid) = -reverse(k[mid+2:end])
function testwavenumberalignment(k, nx)
  mid = Int(nx/2)
  positives = k[2:mid]
  negatives = flippednegatives(k, mid)

  return isapprox(reshape(positives, mid-1), reshape(negatives, mid-1))
end

testk(g) = testwavenumberalignment(g.k, g.nx)
testl(g::Union{TwoDGrid, ThreeDGrid}) = testwavenumberalignment(g.l, g.ny)
testm(g::ThreeDGrid) = testwavenumberalignment(g.m, g.nz)
testkr(g) = CUDA.@allowscalar isapprox(cat(g.k[1:g.nkr-1], abs(g.k[g.nkr]), dims=1), g.kr)

function testgridpoints(dev::Device, g::TwoDGrid{T}) where T
  X, Y = gridpoints(g)
  dXgrid = @. X[2:end, :] - X[1:end-1, :]
  dYgrid = @. Y[:, 2:end] - Y[:, 1:end-1]
  dXones = ArrayType(dev)(g.dx*ones(T, size(dXgrid)))
  dYones = ArrayType(dev)(g.dy*ones(T, size(dYgrid)))

  return isapprox(dXgrid, dXones) && isapprox(dYgrid, dYones)
end

function testgridpoints(dev::Device, g::ThreeDGrid{T}) where T
  X, Y, Z = gridpoints(g)
  dXgrid = @. X[2:end, :, :] - X[1:end-1, :, :]
  dYgrid = @. Y[:, 2:end, :] - Y[:, 1:end-1, :]
  dZgrid = @. Z[:, :, 2:end] - Z[:, :, 1:end-1]
  dXones = ArrayType(dev)(g.dx*ones(T, size(dXgrid)))
  dYones = ArrayType(dev)(g.dy*ones(T, size(dYgrid)))
  dZones = ArrayType(dev)(g.dz*ones(T, size(dZgrid)))

  return isapprox(dXgrid, dXones) && isapprox(dYgrid, dYones) && isapprox(dZgrid, dZones)
end

function testdealias(grid::OneDGrid)
  fh = ones(Complex{eltype(grid)}, size(grid.kr))
  dealias!(fh, grid)

  kmax = round(maximum(grid.kr)*2/3)

  for i₁ = 1:grid.nkr
    if CUDA.@allowscalar grid.kr[i₁] < kmax
      fh[i₁] = 0
    end
  end

  return isapprox(sum(abs.(fh)), 0)
end

function testdealias(grid::TwoDGrid)
  fh = ones(Complex{eltype(grid)}, size(grid.Krsq))
  dealias!(fh, grid)

  kmax = round(maximum(grid.kr)*2/3)
  lmax = round(maximum(abs.(grid.l))*2/3)

  for i₂ = 1:grid.nl, i₁ = 1:grid.nkr
    if (
        (CUDA.@allowscalar grid.kr[i₁] < kmax ) &&
        (CUDA.@allowscalar grid.l[i₂] < lmax && CUDA.@allowscalar grid.l[i₂] ≥ -lmax)
        )
      fh[i₁, i₂] = 0
    end
  end

  return isapprox(sum(abs.(fh)), 0)
end

function testdealias(grid::ThreeDGrid)
  fh = ones(Complex{eltype(grid)}, size(grid.Krsq))
  dealias!(fh, grid)

  kmax = round(maximum(grid.kr)*2/3)
  lmax = round(maximum(abs.(grid.l))*2/3)
  mmax = round(maximum(abs.(grid.m))*2/3)

  for i₃ = 1:grid.nm, i₂ = 1:grid.nl, i₁ = 1:grid.nkr
    if (
        (CUDA.@allowscalar grid.kr[i₁] < kmax ) &&
        (CUDA.@allowscalar grid.l[i₂] < lmax && CUDA.@allowscalar grid.l[i₂] ≥ -lmax) &&
        (CUDA.@allowscalar grid.m[i₃] < mmax && CUDA.@allowscalar grid.m[i₃] ≥ -mmax)
        )
      fh[i₁, i₂, i₃] = 0
    end
  end

  return isapprox(sum(abs.(fh)), 0)
end

function testnodealias(grid::OneDGrid)
  fh = ones(Complex{eltype(grid)}, size(grid.kr))
  return dealias!(fh, grid) == nothing
end

function testnodealias(grid::Union{TwoDGrid, ThreeDGrid})
  fh = ones(Complex{eltype(grid)}, size(grid.Krsq))
  return dealias!(fh, grid) == nothing
end

function testtypedonedgrid(dev::Device, nx, Lx; T=Float64)
  grid = OneDGrid(nx, Lx, T=T)

  return typeof(grid.dx)==T && typeof(grid.x[1])==T && typeof(grid.Lx)==T && eltype(grid) == T
end

function testtypedtwodgrid(dev::Device, nx, Lx, ny=Lx, Ly=Lx; T=Float64)
  grid = TwoDGrid(nx, Lx, ny, Ly, T=T)

  return typeof(grid.dx)==T && typeof(grid.dy)==T && typeof(grid.x[1])==T && typeof(grid.y[1])==T && typeof(grid.Lx)==T && typeof(grid.Ly)==T && eltype(grid) == T
end

function testtypedthreedgrid(dev::Device, nx, Lx, ny=nx, Ly=Lx, nz=nx, Lz=Lx; T=Float64)
  grid = ThreeDGrid(nx, Lx, ny, Ly, nz, Lz, T=T)

  return typeof(grid.dx)==T && typeof(grid.dy)==T && typeof(grid.dz)==T && typeof(grid.x[1])==T && typeof(grid.y[1])==T && typeof(grid.z[1])==T && typeof(grid.Lx)==T && typeof(grid.Ly)==T && typeof(grid.Lz)==T && eltype(grid) == T
end

function testmakefilter(dev::Device, g::AbstractGrid)
  T = eltype(g)
  filter = FourierFlows.makefilter(g)
  G = typeof(g)
  nofilter = G<:OneDGrid ? filter[@. g.kr*g.dx/π < 0.65] : ( G<:TwoDGrid ? filter[@. sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2) < 0.65] : filter[@. sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2 + (g.m*g.dz/π)^2) < 0.65] )
  fullfilter = G<:OneDGrid ? filter[@. g.kr*g.dx/π > 0.999] : ( G<:TwoDGrid ? filter[@. sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2) > 0.999] : filter[@. sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2 + (g.m*g.dz/π)^2) > 0.999] )

  return nofilter== zeros(dev, T, size(nofilter)) .+ 1 && isapprox(fullfilter, zeros(dev, T, size(fullfilter)), atol=1e-12)
end

function test_plan_flows_fftrfft(::CPU; T=Float64)
  A = ArrayType(CPU())
  return (typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4,))))) <: FFTW.cFFTWPlan{Complex{T},-1,false,1} &&
  typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6))))) <: FFTW.cFFTWPlan{Complex{T},-1,false,2} &&
  typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6, 8))))) <: FFTW.cFFTWPlan{Complex{T},-1,false,3} &&
  FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6, 8))), [1, 2]).region == [1, 2] &&
  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4,))))) <: FFTW.rFFTWPlan{T,-1,false,1} &&
  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4, 6))))) <: FFTW.rFFTWPlan{T,-1,false,2} &&
  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4, 6, 8))))) <: FFTW.rFFTWPlan{T,-1,false,3} &&

  return FourierFlows.plan_flows_rfft(A(rand(T, (4, 6, 8))), [1, 2]).region == [1, 2])
end

function test_plan_flows_fftrfft(::GPU; T=Float64)
  A = ArrayType(GPU())
  return (typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4,))))) == CUDA.CUFFT.cCuFFTPlan{Complex{T},-1,false,1} &&
  typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6))))) == CUDA.CUFFT.cCuFFTPlan{Complex{T},-1,false,2} &&
  typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6, 8))))) == CUDA.CUFFT.cCuFFTPlan{Complex{T},-1,false,3} &&
  FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6, 8))), [1, 2]).region == [1, 2] &&
  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4,))))) == CUDA.CUFFT.rCuFFTPlan{T,-1,false,1} &&
  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4, 6))))) == CUDA.CUFFT.rCuFFTPlan{T,-1,false,2} &&
  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4, 6, 8))))) == CUDA.CUFFT.rCuFFTPlan{T,-1,false,3} &&

  return FourierFlows.plan_flows_rfft(A(rand(T, (4, 6, 8))), [1, 2]).region == [1, 2])
end

function test_aliased_fraction(dev, aliased_fraction)
  nx, Lx = 16, 2π
  ny, Ly = 32, 2π
  nz, Lz = 34, 2π

  g₁ = OneDGrid(nx, Lx; aliased_fraction = aliased_fraction)
  g₂ = TwoDGrid(nx, Lx, ny, Ly; aliased_fraction = aliased_fraction)
  g₃ = ThreeDGrid(nx, Lx, ny, Ly, nz, Lz; aliased_fraction = aliased_fraction)

  lower_end(n) = floor(Int, (1 - aliased_fraction)/2 * n) + 1
  upper_end(n) = ceil(Int, (1 + aliased_fraction)/2 * n)
  upper_end_r(n) = Int(n/2)+1

  kralias = aliased_fraction==0 ? nothing : lower_end(nx):upper_end_r(nx)
  kalias = aliased_fraction==0 ? nothing : lower_end(nx):upper_end(nx)
  lalias = aliased_fraction==0 ? nothing : lower_end(ny):upper_end(ny)
  malias = aliased_fraction==0 ? nothing : lower_end(nz):upper_end(nz)

  return (g₁.aliased_fraction == aliased_fraction && g₂.aliased_fraction == aliased_fraction &&
  g₃.aliased_fraction == aliased_fraction && g₁.kralias == kralias && g₁.kalias == kalias &&
  g₂.kralias == kralias && g₂.kalias == kalias && g₂.lalias == lalias && g₃.kralias == kralias &&
  g₃.kalias == kalias && g₃.lalias == lalias && g₃.malias == malias)
end


text:
function toPari(io::IO, f::Generic.Poly{nf_elem})
  Zy, y = PolynomialRing(FlintQQ, "y", cached = false)
  k = base_ring(f)
  kt = parent(f)
  s_k = k.S
  s_kt = kt.S
  #important to pari: the top poly needs to have x, so the field should be y...

  k.S = :t
  kt.S = :x

  print(io, "k = nfinit($(k.pol(y)));\n")
  print(io, "f = $f;\n")

  k.S = s_k
  kt.S = s_kt
  nothing
end

function toPari(s::String, f::Generic.Poly{nf_elem})
  g = open(s, "w")
  toPari(g, f)
  close(g)
end

# gp> #
# gp> read("bla");
# gp> nffactor(k, f);
#=

x-1
x^2-19
x^3-x^2+6
x^4 - 2*x^3 - 7*x^2 - 7*x + 3
x^5 - 3*x^3 - 2*x^2+8*x-10

Zx, x = PolynomialRing(FlintZZ)
g1 = swinnerton_dyer(7, x);
g2 = swinnerton_dyer(7, x);

k, a = number_field()
kt, t = k["t"]

g = (g1*g2)(t-a)
@time length(factor(g))

set_verbose_level(:PolyFactor, 1)

=#


text:
function page(f::Function, io::IO)
    io << open(read, rel"template.html") << f << ""
end

function section(f::Function, io::IO)
    io << """
""" << f << "
"
end

function plot(f::Function, io::IO, title::AbstractString)
    io << """

"""
    io << """"""
end

function var(io::IO, name, data)
    io << """"""
end

function plot_stat{T<:Number}(f::IO, xvar::Variable{T}, xdata::Vector{T})
    x = xvar.name

    page(f) do f
        var(f, x, xdata)

        section(f) do f
            title = "histogram-$x"

            plot(f, title) do f
                f << """[
                    { x: $x, type: 'histogram' }
                ], {
                    title: '$title',
                    xaxis: { title: '$x' },
                }"""
            end
        end

    end
end

function plot_stat{T<:Number, S<:Number}(f::IO, xvar::Variable{T}, yvar::Variable{S}, xdata::Vector{T}, ydata::Vector{S})
    x, y = xvar.name, yvar.name

    page(f) do f
        var(f, x, xdata)
        var(f, y, ydata)

        # section(f) do f
        #     title = "scatter-$y~$x"

        #     plot(f, title) do f
        #         f << """[
        #             { x: $x, y: $y, type: 'scatter', mode: 'markers', marker: { size: 2 } }
        #         ], {
        #             title: '$title',
        #             xaxis: { title: '$x' },
        #             yaxis: { title: '$y' },
        #         }"""
        #     end
        # end

    end
end