text:
using SyntheticDatasets
using DataFrames
using Test

@testset "SkLearn Generators" begin
    samples = 20000
    features = 20

    data = SyntheticDatasets.make_blobs(centers = [-1 1;-0.5 0.75],
                                        cluster_std = 0.225,
                                        n_samples = 20000,
                                        center_box = (-1.5, 1.5))

    @test size(data)[1] == samples
    @test size(data)[2] == 3

    samples = 20000
    data = SyntheticDatasets.make_moons(n_samples = 20000)

    @test size(data)[1] == samples
    @test size(data)[2] == 3

    data = SyntheticDatasets.make_s_curve(  n_samples = samples,
                                            noise = 2.2,
                                            random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == 4

    data = SyntheticDatasets.make_circles(n_samples = samples)

    @test size(data)[1] == samples
    @test size(data)[2] == 3

    data = SyntheticDatasets.make_regression(   n_samples = samples,
                                                n_features = features,
                                                noise = 2.2,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 1

    data = SyntheticDatasets.make_regression(   n_samples = samples,
                                                n_features = features,
                                                n_targets = 2,
                                                noise = 2.2,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 2

    data = SyntheticDatasets.make_regression(   n_samples = samples,
                                                n_features = features,
                                                n_targets = 3,
                                                noise = 2.2,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 3

    data = SyntheticDatasets.make_classification(   n_samples = samples,
                                                    n_features = features,
                                                    n_classes = 1)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 1

    data = SyntheticDatasets.make_friedman1(n_samples = samples,
                                            n_features = features)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 1

    data = SyntheticDatasets.make_friedman2(n_samples = samples)

    @test size(data)[1] == samples
    @test size(data)[2] == 5

    data = SyntheticDatasets.make_friedman3(n_samples = samples)

    @test size(data)[1] == samples
    @test size(data)[2] == 5

    data = SyntheticDatasets.make_low_rank_matrix(  n_samples = samples,
                                                    n_features = features,
                                                    effective_rank = 10,
                                                    tail_strength = 0.5,
                                                    random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features

    data = SyntheticDatasets.make_swiss_roll(   n_samples = samples,
                                                noise = 2.2,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == 4

    data = SyntheticDatasets.make_hastie_10_2(  n_samples = samples,
                                                random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == 11

    data = SyntheticDatasets.make_gaussian_quantiles(   n_samples = samples,
                                                        n_features = features,
                                                        random_state = 5)

    @test size(data)[1] == samples
    @test size(data)[2] == features + 1

    @test_throws DimensionMismatch SyntheticDatasets.make_gaussian_quantiles(   n_samples = 300,
                                                                                n_features = 4,
                                                                                random_state = 5,
                                                                                mean = [1, 2, 3])
end

@testset "Matlab Generators" begin
    samples = 20000

    data = SyntheticDatasets.make_twospirals(n_samples = samples,
                                                 noise = 2.2)
    @test size(data)[1] == samples

    data = SyntheticDatasets.make_halfkernel(n_samples = samples)
    @test size(data)[1] == samples

    data = SyntheticDatasets.make_outlier(n_samples = samples)
    @test size(data)[1] == samples
    @test size(data)[2] == 3
end


text:
using Random

#####################
# General logic
#####################

"""
Variables are represented as 32-bit unsigned integers
"""
const Var = UInt32 # variable ids

"""
Literals are represented as 32-bit signed integers.
Positive literals are positive integers identical to their variable. Negative literals are their negations. Integer 0 should not be used to represent literals.
"""
const Lit = Int32 # variable with a positive or negative sign

"Convert a variable to the corresponding positive literal"
@inline var2lit(v::Var)::Lit = convert(Lit,v)

"Convert a literal its variable, removing the sign of the literal"
@inline lit2var(l::Lit)::Var = convert(Var,abs(l))

var(v::Int) = convert(Var, v)
lit(v::Int) = convert(Lit, v)

#####################
# General circuits
#####################

"Root of the circuit node hierarchy"
abstract type ΔNode <: DagNode end

"Any circuit represented as a bottom-up linear order of nodes"
const Δ = AbstractVector{<:ΔNode}

"A circuit node that has an origin of type O"
abstract type DecoratorΔNode{O<:ΔNode} <: ΔNode end

"Any circuit that has an origin represented as a bottom-up linear order of nodes"
const DecoratorΔ{O} = AbstractVector{<:DecoratorΔNode{O}}

#####################
# General traits
#####################

"""
A trait hierarchy denoting types of nodes
`GateType` defines an orthogonal type hierarchy of node types, not circuit types, so we can dispatch on node type regardless of circuit type.
See @ref{https://docs.julialang.org/en/v1/manual/methods/#Trait-based-dispatch-1}
"""
abstract type GateType end

abstract type LeafGate <: GateType end
abstract type InnerGate <: GateType end

"A trait denoting literal leaf nodes of any type"
struct LiteralGate <: LeafGate end

"A trait denoting constant leaf nodes of any type"
struct ConstantGate <: LeafGate end

"A trait denoting conjuction nodes of any type"
struct ⋀Gate <: InnerGate end

"A trait denoting disjunction nodes of any type"
struct ⋁Gate <: InnerGate end

@inline GateType(instance::ΔNode) = GateType(typeof(instance))

# map gate type traits to graph node traits
import ..Utils.NodeType # make available for extension
@inline NodeType(::Type{N}) where {N<:ΔNode} = NodeType(GateType(N))
@inline NodeType(::LeafGate) = Leaf()
@inline NodeType(::InnerGate) = Inner()

#####################
# node methods
#####################

# following methods should be defined for all types of circuits

"Get the node type at the root of the corresponding hierarchy."
@inline @Base.pure node_type(c::Δ) = node_type(eltype(c))
@inline @Base.pure node_type(n::ΔNode) = node_type(typeof(n))

import ..Utils.children # make available for extension

"Get the logical literal in a given literal leaf node"
@inline literal(n::ΔNode)::Lit = literal(GateType(n), n)
@inline literal(::LiteralGate, n::ΔNode)::Lit =
    error("Each `LiteralGate` should implement a `literal` method. It is missing from $(typeof(n)).")

"Get the logical constant in a given constant leaf node"
@inline constant(n::ΔNode)::Bool = literal(GateType(n), n)
@inline constant(::ConstantGate, n::ΔNode)::Bool =
    error("Each `ConstantGate` should implement a `constant` method.  It is missing from $(typeof(n)).")

# next bunch of methods are derived from literal, constant, children, and the traits

"Get the logical variable in a given literal leaf node"
@inline variable(n::ΔNode)::Var = variable(GateType(n), n)
@inline variable(::LiteralGate, n::ΔNode)::Var = lit2var(literal(n))

"Get the sign of the literal leaf node"
@inline positive(n::ΔNode)::Bool = positive(GateType(n), n)
@inline positive(::LiteralGate, n::ΔNode)::Bool = literal(n) >= 0
@inline negative(n::ΔNode)::Bool = !positive(n)

"Is the circuit syntactically equal to true?"
@inline is_true(n::ΔNode)::Bool = is_true(GateType(n), n)
@inline is_true(::GateType, n::ΔNode)::Bool = false
@inline is_true(::ConstantGate, n::ΔNode)::Bool = (constant(n) == true)

"Is the circuit syntactically equal to false?"
@inline is_false(n::ΔNode)::Bool = is_false(GateType(n), n)
@inline is_false(::GateType, n::ΔNode)::Bool = false
@inline is_false(::ConstantGate, n::ΔNode)::Bool = (constant(n) == false)

"Construct a true node in the hierarchy of node n"
true_like(n) = conjoin_like(n)

"Construct a false node in the hierarchy of node n"
false_like(n) = disjoin_like(n)

"Get the origin of the given decorator circuit node"
@inline (origin(n::DecoratorΔNode{O})::O) where {O<:ΔNode} =
    n.origin
"Get the origin of the given decorator circuit"
@inline origin(circuit::DecoratorΔ) =
    lower_element_type(map(n -> n.origin, circuit))

"Get the first origin of the given decorator circuit node of the given type"
@inline (origin(n::DecoratorΔNode{<:O}, ::Type{O})::O) where {O<:ΔNode} =
    origin(n)
@inline (origin(n::DecoratorΔNode, ::Type{T})::T) where {T<:ΔNode} =
    origin(origin(n),T)
"Get the first origin of the given decorator circuit of the given node type"
@inline origin(circuit::DecoratorΔ, ::Type{T}) where T =
    lower_element_type(map(n -> origin(n,T), circuit)) # TODO can we do this with eltype of the circuit type?

"Get the origin of the origin of the given decorator circuit node"
@inline (grand_origin(n::DecoratorΔNode{<:DecoratorΔNode{O}})::O) where {O} =
    n.origin.origin
"Get the origin of the origin the given decorator circuit"
@inline grand_origin(circuit::DecoratorΔ) =
    origin(origin(circuit))

"Conjoin nodes in the same way as the example"
@inline function conjoin_like(example::ΔNode, arguments::ΔNode...)
    conjoin_like(example, collect(arguments))
end

"Disjoin nodes in the same way as the example"
@inline function disjoin_like(example::ΔNode, arguments::ΔNode...)
    disjoin_like(example, collect(arguments))
end

ext: jl


text:
module VarReplay

using Turing: CACHERESET, CACHEIDCS, CACHERANGES
using Turing: Sampler, realpart, dualpart, vectorize, reconstruct, reconstruct!, SimplexDistribution
using Distributions

import Base: string, isequal, ==, hash, getindex, setindex!, push!, show, isempty
import Turing: link!, invlink!, link, invlink

export VarName, VarInfo, uid, sym, getlogp, set_retained_vns_del_by_spl!, resetlogp!, is_flagged, unset_flag!, setgid!, copybyindex,
       setorder!, updategid!, acclogp!, istrans, link!, invlink!, setlogp!, getranges, getrange, getvns, cleandual!, getval
export string, isequal, ==, hash, getindex, setindex!, push!, show, isempty

###########
# VarName #
###########
struct VarName
  csym      ::    Symbol        # symbol generated in compilation time
  sym       ::    Symbol        # variable symbol
  indexing  ::    String        # indexing
  counter   ::    Int           # counter of same {csym, uid}
end

# NOTE: VarName should only be constructed by VarInfo internally due to the nature of the counter field.

uid(vn::VarName) = (vn.csym, vn.sym, vn.indexing, vn.counter)
Base.hash(vn::VarName) = hash(uid(vn))

isequal(x::VarName, y::VarName) = hash(uid(x)) == hash(uid(y))
==(x::VarName, y::VarName)      = isequal(x, y)

Base.string(vn::VarName) = "{$(vn.csym),$(vn.sym)$(vn.indexing)}:$(vn.counter)"
Base.string(vns::Vector{VarName}) = replace(string(map(vn -> string(vn), vns)), "String" => "")

sym(vn::VarName) = Symbol("$(vn.sym)$(vn.indexing)")  # simplified symbol

cuid(vn::VarName) = (vn.csym, vn.sym, vn.indexing)    # the uid which is only available at compile time

copybyindex(vn::VarName, indexing::String) = VarName(vn.csym, vn.sym, indexing, vn.counter)

###########
# VarInfo #
###########

mutable struct VarInfo
  idcs        ::    Dict{VarName,Int}
  vns         ::    Vector{VarName}
  ranges      ::    Vector{UnitRange{Int}}
  vals        ::    Vector{Real}
  rvs         ::    Dict{Union{VarName,Vector{VarName}},Any}
  dists       ::    Vector{Distributions.Distribution}
  gids        ::    Vector{Int}
  logp        ::    Real
  pred        ::    Dict{Symbol,Any}
  num_produce ::    Int           # num of produce calls from trace, each produce corresponds to an observe.
  orders      ::    Vector{Int}   # observe statements number associated with random variables
  flags       ::    Dict{String,Vector{Bool}}

  VarInfo() = begin
    vals  = Vector{Real}()
    rvs   = Dict{Union{VarName,Vector{VarName}},Any}()
    logp  = zero(Real)
    pred  = Dict{Symbol,Any}()
    flags = Dict{String,Vector{Bool}}()
    flags["del"] = Vector{Bool}()
    flags["trans"] = Vector{Bool}()

    new(
      Dict{VarName, Int}(),
      Vector{VarName}(),
      Vector{UnitRange{Int}}(),
      vals,
      rvs,
      Vector{Distributions.Distribution}(),
      Vector{Int}(),
      logp,
      pred,
      0,
      Vector{Int}(),
      flags
    )
  end
end

const VarView = Union{Int,UnitRange,Vector{Int},Vector{UnitRange}}

getidx(vi::VarInfo, vn::VarName) = vi.idcs[vn]

getrange(vi::VarInfo, vn::VarName) = vi.ranges[getidx(vi, vn)]
getranges(vi::VarInfo, vns::Vector{VarName}) = union(map(vn -> getrange(vi, vn), vns)...)

getval(vi::VarInfo, vn::VarName)       = view(vi.vals, getrange(vi, vn))
setval!(vi::VarInfo, val, vn::VarName) = vi.vals[getrange(vi, vn)] = val

getval(vi::VarInfo, vns::Vector{VarName}) = view(vi.vals, getranges(vi, vns))

getval(vi::VarInfo, vview::VarView)                      = view(vi.vals, vview)
setval!(vi::VarInfo, val::Any, vview::VarView)           = vi.vals[vview] = val
setval!(vi::VarInfo, val::Any, vview::Vector{UnitRange}) = length(vview) > 0 ? (vi.vals[[i for arr in vview for i in arr]] = val) : nothing

getall(vi::VarInfo)            = vi.vals
setall!(vi::VarInfo, val::Any) = vi.vals = val

getsym(vi::VarInfo, vn::VarName) = vi.vns[getidx(vi, vn)].sym

getdist(vi::VarInfo, vn::VarName) = vi.dists[getidx(vi, vn)]

getgid(vi::VarInfo, vn::VarName) = vi.gids[getidx(vi, vn)]

setgid!(vi::VarInfo, gid::Int, vn::VarName) = vi.gids[getidx(vi, vn)] = gid

istrans(vi::VarInfo, vn::VarName) = is_flagged(vi, vn, "trans")
settrans!(vi::VarInfo, trans::Bool, vn::VarName) = trans ? set_flag!(vi, vn, "trans") : unset_flag!(vi, vn, "trans")

getlogp(vi::VarInfo) = vi.logp
setlogp!(vi::VarInfo, logp::Real) = vi.logp = logp
acclogp!(vi::VarInfo, logp::Any) = vi.logp += logp
resetlogp!(vi::VarInfo) = setlogp!(vi, zero(Real))

isempty(vi::VarInfo) = isempty(vi.idcs)

# X -> R for all variables associated with given sampler
link!(vi::VarInfo, spl::Sampler) = begin
  vns = getvns(vi, spl)
  if ~istrans(vi, vns[1])
    for vn in vns
      dist = getdist(vi, vn)
      setval!(vi, vectorize(dist, link(dist, reconstruct(dist, getval(vi, vn)))), vn)
      settrans!(vi, true, vn)
    end
  else
    @warn("[Turing] attempt to link a linked vi")
  end
end

# R -> X for all variables associated with given sampler
invlink!(vi::VarInfo, spl::Sampler) = begin
  vns = getvns(vi, spl)
  if istrans(vi, vns[1])
    for vn in vns
      dist = getdist(vi, vn)
      setval!(vi, vectorize(dist, invlink(dist, reconstruct(dist, getval(vi, vn)))), vn)
      settrans!(vi, false, vn)
    end
  else
    @warn("[Turing] attempt to invlink an invlinked vi")
  end
end

function cleandual!(vi::VarInfo)
  for i = 1:length(vi.vals)
    vi.vals[i] = realpart(vi.vals[i])
  end
  vi.logp = realpart(getlogp(vi))
end

vns(vi::VarInfo) = Set(keys(vi.idcs))            # get all vns
syms(vi::VarInfo) = map(vn -> vn.sym, vns(vi))  # get all symbols

# The default getindex & setindex!() for get & set values
# NOTE: vi[vn] will always transform the variable to its original space and Julia type
Base.getindex(vi::VarInfo, vn::VarName) = begin
  @assert haskey(vi, vn) "[Turing] attempted to replay unexisting variables in VarInfo"
  dist = getdist(vi, vn)
  # if isa(dist, SimplexDistribution) || isa(dist, MvNormal) # Reduce memory allocation for distributions with simplex constraints
  #   if vn in keys(vi.rvs)
  #     r = vi.rvs[vn]
  #     reconstruct!(r, dist, getval(vi, vn))
  #   else
  #     r = reconstruct(dist, getval(vi, vn))
  #     r_real = similar(r, Real)
  #     r_real[:] = r
  #     vi.rvs[vn] = r_real
  #     r = r_real
  #   end
  #   istrans(vi, vn) ?
  #     invlink!(r, dist, r) :
  #     r
  # else
    istrans(vi, vn) ?
      invlink(dist, reconstruct(dist, getval(vi, vn))) :
      reconstruct(dist, getval(vi, vn))
  # end
end

Base.setindex!(vi::VarInfo, val::Any, vn::VarName) = setval!(vi, val, vn)

Base.getindex(vi::VarInfo, vns::Vector{VarName}) = begin
  @assert haskey(vi, vns[1]) "[Turing] attempted to replay unexisting variables in VarInfo"
  dist = getdist(vi, vns[1])
  # if isa(dist, SimplexDistribution) # Reduce memory allocation for distributions with simplex constraints
  #   if vns in keys(vi.rvs)
  #     r = vi.rvs[vns]
  #     reconstruct!(r, dist, getval(vi, vn))
  #   else
  #     r = reconstruct(dist, getval(vi, vns), length(vns))
  #     r_real = similar(r, Real)
  #     r_real[:] = r
  #     vi.rvs[vns] = r_real
  #     r = r_real
  #   end
  #   istrans(vi, vns[1]) ?
  #     invlink!(r, dist, r) :
  #     r
  # else
    istrans(vi, vns[1]) ?
      invlink(dist, reconstruct(dist, getval(vi, vns), length(vns))) :
      reconstruct(dist, getval(vi, vns), length(vns))
  # end
end

# NOTE: vi[vview] will just return what insdie vi (no transformations applied)
Base.getindex(vi::VarInfo, vview::VarView)            = getval(vi, vview)
Base.setindex!(vi::VarInfo, val::Any, vview::VarView) = setval!(vi, val, vview)

Base.getindex(vi::VarInfo, spl::Sampler)            = getval(vi, getranges(vi, spl))
Base.setindex!(vi::VarInfo, val::Any, spl::Sampler) = setval!(vi, val, getranges(vi, spl))

Base.getindex(vi::VarInfo, spl::Nothing)            = getall(vi)
Base.setindex!(vi::VarInfo, val::Any, spl::Nothing) = setall!(vi, val)

Base.keys(vi::VarInfo) = keys(vi.idcs)

Base.haskey(vi::VarInfo, vn::VarName) = haskey(vi.idcs, vn)

Base.show(io::IO, vi::VarInfo) = begin
  vi_str = """
  /=======================================================================
  | VarInfo
  |-----------------------------------------------------------------------
  | Varnames  :   $(string(vi.vns))
  | Range     :   $(vi.ranges)
  | Vals      :   $(vi.vals)
  | RVs       :   $(vi.rvs)
  | GIDs      :   $(vi.gids)
  | Orders    :   $(vi.orders)
  | Logp      :   $(vi.logp)
  | #produce  :   $(vi.num_produce)
  | flags     :   $(vi.flags)
  \\=======================================================================
  """
  print(io, vi_str)
end

# Add a new entry to VarInfo
push!(vi::VarInfo, vn::VarName, r::Any, dist::Distributions.Distribution, gid::Int) = begin

  @assert ~(vn in vns(vi)) "[push!] attempt to add an exisitng variable $(sym(vn)) ($(vn)) to VarInfo (keys=$(keys(vi))) with dist=$dist, gid=$gid"

  val = vectorize(dist, r)

  vi.idcs[vn] = length(vi.idcs) + 1
  push!(vi.vns, vn)
  l = length(vi.vals); n = length(val)
  push!(vi.ranges, l+1:l+n)
  append!(vi.vals, val)
  push!(vi.dists, dist)
  push!(vi.gids, gid)
  push!(vi.orders, vi.num_produce)
  push!(vi.flags["del"], false)
  push!(vi.flags["trans"], false)

  vi
end

setorder!(vi::VarInfo, vn::VarName, index::Int) = begin
  if vi.orders[vi.idcs[vn]] != index
    vi.orders[vi.idcs[vn]] = index
  end
  vi
end

# This method is use to generate a new VarName with the right count
VarName(vi::VarInfo, csym::Symbol, sym::Symbol, indexing::String) = begin
  # TODO: update this method when implementing the sanity check
  VarName(csym, sym, indexing, 1)
end
VarName(vi::VarInfo, syms::Vector{Symbol}, indexing::String) = begin
  # TODO: update this method when implementing the sanity check
    VarName(syms[1], syms[2], indexing, 1)
end

#################################
# Utility functions for VarInfo #
#################################

# expand!(vi::VarInfo) = begin
#   push!(vi.vals, realpart(vi.vals[end])); vi.vals[end], vi.vals[end-1] = vi.vals[end-1], vi.vals[end]
#   push!(vi.trans, deepcopy(vi.trans[end]))
#   push!(vi.logp, zero(Real))
# end
#
# shrink!(vi::VarInfo) = begin
#   pop!(vi.vals)
#   pop!(vi.trans)
#   pop!(vi.logp)
# end
#
# last!(vi::VarInfo) = begin
#   vi.vals = vi.vals[end:end]
#   vi.trans = vi.trans[end:end]
#   vi.logp = vi.logp[end:end]
# end

# Get all indices of variables belonging to gid or 0
getidcs(vi::VarInfo) = getidcs(vi, nothing)
getidcs(vi::VarInfo, spl::Nothing) = filter(i -> vi.gids[i] == 0 || vi.gids[i] == 0, 1:length(vi.gids))
getidcs(vi::VarInfo, spl::Sampler) = begin
  # NOTE: 0b00 is the sanity flag for
  #         |\____ getidcs   (mask = 0b10)
  #         \_____ getranges (mask = 0b01)
  if ~haskey(spl.info, :cache_updated) spl.info[:cache_updated] = CACHERESET end
  if haskey(spl.info, :idcs) && (spl.info[:cache_updated] & CACHEIDCS) > 0
    spl.info[:idcs]
  else
    spl.info[:cache_updated] = spl.info[:cache_updated] | CACHEIDCS
    spl.info[:idcs] = filter(i ->
      (vi.gids[i] == spl.alg.gid || vi.gids[i] == 0) && (isempty(spl.alg.space) || is_inside(vi.vns[i], spl.alg.space)),
      1:length(vi.gids)
    )
  end
end

is_inside(vn::VarName, space::Set)::Bool = begin
  if vn.sym in space
    true
  else
    exprs = filter(el -> isa(el, Expr), space)
    strs = Set((replace(string(ex), r"\(|\)" => "") for ex in exprs))
    vn_str = string(vn.sym) * vn.indexing
    valid = filter(str -> occursin(str, vn_str), strs)
    length(valid) > 0
  end
end

# Get all values of variables belonging to gid or 0
getvals(vi::VarInfo) = getvals(vi, nothing)
getvals(vi::VarInfo, spl::Union{Nothing, Sampler}) = view(vi.vals, getidcs(vi, spl))

# Get all vns of variables belonging to gid or 0
getvns(vi::VarInfo) = getvns(vi, nothing)
getvns(vi::VarInfo, spl::Union{Nothing, Sampler}) = view(vi.vns, getidcs(vi, spl))

# Get all vns of variables belonging to gid or 0
getranges(vi::VarInfo, spl::Sampler) = begin
  if ~haskey(spl.info, :cache_updated) spl.info[:cache_updated] = CACHERESET end
  if haskey(spl.info, :ranges) && (spl.info[:cache_updated] & CACHERANGES) > 0
    spl.info[:ranges]
  else
    spl.info[:cache_updated] = spl.info[:cache_updated] | CACHERANGES
    spl.info[:ranges] = union(map(i -> vi.ranges[i], getidcs(vi, spl))...)
  end
end

# NOTE: this function below is not used anywhere but test files.
#       we can safely remove it if we want.
getretain(vi::VarInfo, spl::Union{Nothing, Sampler}) = begin
  gidcs = getidcs(vi, spl)
  if vi.num_produce == 0 # called at begening of CSMC sweep for non reference particles
    UnitRange[map(i -> vi.ranges[gidcs[i]], length(gidcs):-1:1)...]
  else
    retained = [idx for idx in 1:length(vi.orders) if idx in gidcs && vi.orders[idx] > vi.num_produce]
    UnitRange[map(i -> vi.ranges[i], retained)...]
  end
end

#######################################
# Rand & replaying method for VarInfo #
#######################################

is_flagged(vi::VarInfo, vn::VarName, flag::String) = vi.flags[flag][getidx(vi, vn)]
set_flag!(vi::VarInfo, vn::VarName, flag::String) = vi.flags[flag][getidx(vi, vn)] = true
unset_flag!(vi::VarInfo, vn::VarName, flag::String) = vi.flags[flag][getidx(vi, vn)] = false

set_retained_vns_del_by_spl!(vi::VarInfo, spl::Sampler) = begin
  gidcs = getidcs(vi, spl)
  if vi.num_produce == 0
    for i = length(gidcs):-1:1
      vi.flags["del"][gidcs[i]] = true
    end
  else
    retained = [idx for idx in 1:length(vi.orders) if idx in gidcs && vi.orders[idx] > vi.num_produce]
    for i = retained
      vi.flags["del"][i] = true
    end
  end
end

updategid!(vi::VarInfo, vn::VarName, spl::Sampler) = begin
  if ~isempty(spl.alg.space) && getgid(vi, vn) == 0 && getsym(vi, vn) in spl.alg.space
    setgid!(vi, spl.alg.gid, vn)
  end
end

end


text:
using DDR2import
using Test

@testset "DDR2import.jl" begin
    tests = ["Implemented\\SO6", "Implemented\\T5", "Implemented\\Exp2",
    "Implemented\\Allftplus", "Implemented\\Ase", "Implemented\\Are",
    "Implemented\\Gar", "Implemented\\Gsl", "Implemented\\Frp",
    "Implemented\\Sls", "Implemented\\Sid", "Implemented\\Star",
    "Implemented\\Spc", "Implemented\\Mot", "Implemented\\Narp",
    "Implemented\\Nnpt", "Implemented\\Cost", "Implemented\\Crco",
    "Implemented\\Ur", "Implemented\\Arp", "Implemented\\Ntfv",
    "util", "Implemented\\Country", "Implemented\\Operator",
    "Implemented\\Routes", "Implemented\\Runway", "Implemented\\For"]
    # tests = ["Implemented\\For"]

    for t in tests
        include("$(t).jl")
    end
end


text:
package ln

struct Triangle
    V1::Vector
    V2::Vector
    V3::Vector
    Box::Box
end

type Triangle struct {
	V1, V2, V3 Vector
	Box        Box
}

func NewTriangle(v1, v2, v3 Vector) *Triangle {
	t := Triangle{}
	t.V1 = v1
	t.V2 = v2
	t.V3 = v3
	t.UpdateBoundingBox()
	return &t
}

func (t *Triangle) UpdateBoundingBox() {
	min := t.V1.Min(t.V2).Min(t.V3)
	max := t.V1.Max(t.V2).Max(t.V3)
	t.Box = Box{min, max}
}

func (t *Triangle) Compile() {
}

func (t *Triangle) BoundingBox() Box {
	return t.Box
}

func (t *Triangle) Contains(v Vector, f float64) bool {
	return false
}

func (t *Triangle) Intersect(r Ray) Hit {
	e1x := t.V2.X - t.V1.X
	e1y := t.V2.Y - t.V1.Y
	e1z := t.V2.Z - t.V1.Z
	e2x := t.V3.X - t.V1.X
	e2y := t.V3.Y - t.V1.Y
	e2z := t.V3.Z - t.V1.Z
	px := r.Direction.Y*e2z - r.Direction.Z*e2y
	py := r.Direction.Z*e2x - r.Direction.X*e2z
	pz := r.Direction.X*e2y - r.Direction.Y*e2x
	det := e1x*px + e1y*py + e1z*pz
	if det > -EPS && det < EPS {
		return NoHit
	}
	inv := 1 / det
	tx := r.Origin.X - t.V1.X
	ty := r.Origin.Y - t.V1.Y
	tz := r.Origin.Z - t.V1.Z
	u := (tx*px + ty*py + tz*pz) * inv
	if u < 0 || u > 1 {
		return NoHit
	}
	qx := ty*e1z - tz*e1y
	qy := tz*e1x - tx*e1z
	qz := tx*e1y - ty*e1x
	v := (r.Direction.X*qx + r.Direction.Y*qy + r.Direction.Z*qz) * inv
	if v < 0 || u+v > 1 {
		return NoHit
	}
	d := (e2x*qx + e2y*qy + e2z*qz) * inv
	if d < EPS {
		return NoHit
	}
	return Hit{t, d}
}

func (t *Triangle) Paths() Paths {
	return Paths{
		{t.V1, t.V2},
		{t.V2, t.V3},
		{t.V3, t.V1},
	}
}


text:
using Rubin
using Test

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


text:
import ExaPF: PowerSystem
import ExaPF.PowerSystem: ParsePSSE, ParseMAT
using DelimitedFiles

const PS = PowerSystem

"""
    parse_file(datafile::String)

Parse MATPOWER or PSSE instances using ExaPF's parsers.
Return full dataset as `Dict{String, Array{Float64, 2}`,
with entries

- "bus": specifications for all buses in the network
- "branch": specifications for all branches in the network
- "gen": specifications for all generators in the network
- "costs": costs coefficients.
- "baseMVA": baseMVA of the network

"""
function parse_file(datafile)
    if endswith(datafile, ".raw")
        data_raw = ParsePSSE.parse_raw(datafile)
        data = ParsePSSE.raw_to_exapf(data_raw)
    elseif endswith(datafile, ".m")
        data_mat = ParseMAT.parse_mat(datafile)
        data = ParseMAT.mat_to_exapf(data_mat)
    else
        error("Unsupported format in file $(datafile): supported extensions are " *
              "Matpower (.m) or PSSE (.raw)")
    end
    return data
end

struct Bus
    bus_i::Int
    bustype::Int
    Pd::Float64
    Qd::Float64
    Gs::Float64
    Bs::Float64
    area::Int
    Vm::Float64
    Va::Float64
    baseKV::Float64
    zone::Int
    Vmax::Float64
    Vmin::Float64
end

struct Line
    from::Int
    to::Int
    r::Float64
    x::Float64
    b::Float64
    rateA::Float64
    rateB::Float64
    rateC::Float64
    ratio::Float64 #TAP
    angle::Float64 #SHIFT
    status::Int
    angmin::Float64
    angmax::Float64
end
Line() = Line(0,0,0.,0.,0.,0.,0.,0.,0.,0.,0,0.,0.)

mutable struct Gener
    # .gen fields
    bus::Int
    Pg::Float64
    Qg::Float64
    Qmax::Float64
    Qmin::Float64
    Vg::Float64
    mBase::Float64
    status::Int
    Pmax::Float64
    Pmin::Float64
    Pc1::Float64
    Pc2::Float64
    Qc1min::Float64
    Qc1max::Float64
    Qc2min::Float64
    Qc2max::Float64
    ramp_agc::Float64
    scen_agc::Float64 #ramping factor under contingency
    # .gencost fields
    gentype::Int
    startup::Float64
    shutdown::Float64
    n::Int
    coeff::Array
    alpha::Float64
end

# Update
struct OPFData
    buses::Array{Bus}
    lines::Array{Line}
    generators::Array{Gener}
    bus_ref::Int
    baseMVA::Float64
    Ybus::SparseMatrixCSC{Complex{Float64},Int} # bus-admittance matrix
    BusIdx::Dict{Int,Int}    #map from bus ID to bus index
    BusGenerators::Dict{Int, Array{Int}}     #list of generators for each bus (Array of Array)
    Pd::Array                #2d array of active power demands over a time horizon
    Qd::Array                #2d array of reactive power demands
end

"""
    RawData

Specifies the ACOPF instance data.

- `baseMVA`: imported with ExaPF parser
- `bus_arr`: imported with ExaPF parser
- `branch_arr`: imported with ExaPF parser
- `gen_arr`: imported with ExaPF parser
- `costgen_arr`: imported with ExaPF parser
- `pd_arr`: read from `.Pd` file
- `qd_arr`: read from `.Qd` file
- `ctgs_arr`: read from `.Ctgs` file
"""
mutable struct RawData
    baseMVA::Float64
    bus_arr::Array{Float64, 2}
    branch_arr::Array{Float64, 2}
    gen_arr::Array{Float64, 2}
    costgen_arr::Array{Float64, 2}
    pd_arr::Array{Float64, 2}
    qd_arr::Array{Float64, 2}
    ctgs_arr
end

##
## Consider hard-coding a multiperiod load profile:
## multiplier_hourly = [1.0, 1.08199882, 1.11509159, 1.13441867, 1.14755413, 1.15744821, 1.15517453, 1.15771819, 1.15604058, 1.15952428, 1.18587814, 1.22488686, 1.24967452, 1.22919352, 1.16712273, 1.08527516, 0.98558458, 0.89784627, 0.88171420, 0.86553416, 0.8489768, 0.84243794, 0.84521173, 0.90866251]
## multiplier_minute = [1.0, 0.99961029, 0.99930348, 0.99907022, 0.99890119, 0.99878704, 0.99871845, 0.99868608, 0.99868059, 0.99869265, 0.99871293, 0.99873216, 0.99874886, 0.99877571, 0.99882691, 0.99891664, 0.9990591, 0.99926107, 0.99949027, 0.99970169, 0.99985026, 0.99989093, 0.99978076, 0.99951609, 0.99912816, 0.99864944, 0.99811239, 0.99754946, 0.99698767, 0.99644131, 0.99592287, 0.99544481, 0.99501961, 0.99465395, 0.99431417, 0.9939494, 0.99350871, 0.99294118, 0.99219837, 0.99130363, 0.99036111, 0.98947927, 0.98876657, 0.98833145, 0.98823171, 0.98837066, 0.98862223, 0.98886034, 0.98895895, 0.98880304, 0.9883816, 0.9877406, 0.98692657, 0.98598606, 0.98496561, 0.98391176, 0.98287107, 0.98189006, 0.98101529, 0.9802933]
## Can cyclically repeat the multiplier_hourly profile
##

function RawData(case_name, scen_file::String="")
    data = parse_file(case_name)
    bus_arr = data["bus"]
    branch_arr = data["branch"]
    gen_arr = data["gen"]
    costgen_arr = data["cost"]
    baseMVA = data["baseMVA"][1]

    # Sanity checks
    @assert size(gen_arr, 1) == size(costgen_arr, 1)

    pd_arr = Array{Float64, 2}(undef, 0, 0)
    qd_arr = Array{Float64, 2}(undef, 0, 0)
    ctgs_arr = Array{Int64, 2}(undef, 0, 0)

    if isfile(scen_file * ".Pd")
        pd_arr = readdlm(scen_file * ".Pd")
    end
    if isfile(scen_file * ".Qd")
        qd_arr = readdlm(scen_file * ".Qd")
    end
    if isfile(scen_file * ".Ctgs")
        ctgs_arr = readdlm(scen_file * ".Ctgs", Int)
    end
    return RawData(baseMVA, bus_arr, branch_arr, gen_arr, costgen_arr, pd_arr, qd_arr, ctgs_arr)
end

# UTILS
ctgs_loaddata(raw::RawData, n) = raw.ctgs_arr[1:n]
function check_loads(loads, loads_ref; rtol=1e-2)
    return maximum((loads .- loads_ref) ./ max.(1.0, loads_ref)) < rtol
end

#
"""
    opf_loaddata(raw::RawData;
                 time_horizon_start::Int=1,
                 time_horizon_end::Int=0,
                 load_scale::Float64=1.0,
                 ramp_scale::Float64=0.0,
                 corr_scale::Float64=0.1,
                 lineOff=Line())

Loads the multi-period ACOPF instance data from `raw`
with the time horizon defined to be
[`time_horizon_start`,  `time_horizon_end`].
Note that `time_horizon_end = 0` indicates as many
as possible (the number of columns in `raw.pd_arr`).

All loads in all time periods will be multiplied by `load_scale`.
The `ramp_scale` is the factor which multiplies ``p_{g}^{max}``
to get generator ramping ``r_g``.
The `corr_scale` is the factor which multiplies ``r_g``
to get generator ramping for corrective control.
These are set in `ModelInfo`.  See [Model parameters](@ref).

`lineOff` is a transmission line that can be deleted to
represent a contingency.
"""
function opf_loaddata(raw::RawData;
                      time_horizon_start::Int=1,
                      time_horizon_end::Int=0,
                      load_scale::Float64=1.0,
                      ramp_scale::Float64=0.0,
                      corr_scale::Float64=0.1,
                      lineOff=Line())
    #
    # load buses
    #
    ncols_bus = size(raw.bus_arr, 2)
    bus_arr = raw.bus_arr
    num_buses = size(bus_arr, 1)
    buses = Array{Bus}(undef, num_buses)
    bus_ref = -1
    for i in 1:num_buses
        @assert bus_arr[i,1] > 0  #don't support nonpositive bus ids
        bus_arr[i,9] *= pi / 180.0 # ANIRUDH: Bus is an immutable struct. Modify bus_arr itself
        buses[i] = Bus(bus_arr[i, 1:ncols_bus]...)
        # buses[i].Va *= pi/180 # ANIRUDH: See previous comment
        if buses[i].bustype == PS.REF_BUS_TYPE
            if bus_ref > 0
                error("More than one reference bus present in the data")
            end
            bus_ref = i
        end
    end

    #
    # load branches/lines
    #
    branch_arr = raw.branch_arr
    num_lines = size(branch_arr, 1)
    lines_on = findall((branch_arr[:,11] .> 0) .& ((branch_arr[:,1].!=lineOff.from) .| (branch_arr[:,2].!=lineOff.to)) )
    num_on   = length(lines_on)

    if lineOff.from> 0 && lineOff.to > 0
        # println("opf_loaddata: was asked to remove line from,to=", lineOff.from, ",", lineOff.to)
    end
    if length(findall(branch_arr[:,11].==0))>0
        # println("opf_loaddata: ", num_lines-length(findall(branch_arr[:,11].>0)), " lines are off and will be discarded (out of ", num_lines, ")")
    end

    lines = Array{Line}(undef, num_on)
    ncols_lines = size(branch_arr, 2)

    lit = 0
    has_voltage_angle_bounds = false
    for i in lines_on
        @assert branch_arr[i,11] == 1  #should be on since we discarded all other
        lit += 1
        lines[lit] = Line(branch_arr[i, 1:ncols_lines]...)
        if lines[lit].angmin > -360 || lines[lit].angmax < 360
            has_voltage_angle_bounds = true
        end

    end
    @assert lit == num_on
    (has_voltage_angle_bounds) && @warn("Bounds of voltage angles are still to be implemented.", maxlog=1)

    #
    # load generators
    #
    gen_arr = raw.gen_arr
    costgen_arr = raw.costgen_arr
    num_gens = size(gen_arr, 1)
    ncols_gens = size(gen_arr, 2)

    gens_on = findall(gen_arr[:, 8].!=0); num_on = length(gens_on)
    num_off = num_gens-num_on
    if num_off > 0
        @warn("loaddata: $(num_off) generators are off and will be discarded (out of $(num_gens))", maxlog=1)
    end

    baseMVA = raw.baseMVA
    generators = Array{Gener}(undef, num_on)
    R = 0.04 # Droop regulation
    for (i, git) in enumerate(gens_on)

        generators[i] = Gener(0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,Array{Int}(undef, 0),0)

        generators[i].bus      = gen_arr[git,1]
        generators[i].Pg       = gen_arr[git,2] / baseMVA
        generators[i].Qg       = gen_arr[git,3] / baseMVA
        generators[i].Qmax     = gen_arr[git,4] / baseMVA
        generators[i].Qmin     = gen_arr[git,5] / baseMVA
        generators[i].Vg       = gen_arr[git,6]
        generators[i].mBase    = gen_arr[git,7]
        generators[i].status   = gen_arr[git,8]
        @assert generators[i].status==1
        generators[i].Pmax     = gen_arr[git,9]  / baseMVA
        generators[i].Pmin     = gen_arr[git,10] / baseMVA
        generators[i].Pc1      = gen_arr[git,11]
        generators[i].Pc2      = gen_arr[git,12]
        generators[i].Qc1min   = gen_arr[git,13]
        generators[i].Qc1max   = gen_arr[git,14]
        generators[i].Qc2min   = gen_arr[git,15]
        generators[i].Qc2max   = gen_arr[git,16]
        generators[i].ramp_agc = gen_arr[git,9] * ramp_scale  / baseMVA
        generators[i].scen_agc = gen_arr[git,9] * ramp_scale * corr_scale / baseMVA
        generators[i].gentype  = costgen_arr[git,1]
        generators[i].startup  = costgen_arr[git,2]
        generators[i].shutdown = costgen_arr[git,3]
        generators[i].n        = costgen_arr[git,4]
        generators[i].alpha    = -((1/R)*generators[i].Pmax)
        if generators[i].gentype == 1
            generators[i].coeff = costgen_arr[git,5:end]
            error("Piecewise linear costs remains to be implemented.")
        else
            if generators[i].gentype == 2
                generators[i].coeff = costgen_arr[git,5:end]
                #println(generators[i].coeff, " ", length(generators[i].coeff), " ", generators[i].coeff[2])
            else
                error("Invalid generator cost model in the data.")
            end
        end
    end

    # build a dictionary between buses ids and their indexes
    busIdx = PS.get_bus_id_to_indexes(bus_arr)
    # Remove deactivated branches before building admittance matrix
    topology = PS.makeYbus(bus_arr, branch_arr[lines_on, :], baseMVA, busIdx)
    Ybus = topology.ybus
    # generators at each bus
    gen_active = gen_arr[gens_on, :]
    BusGeners = PS.get_bus_generators(bus_arr, gen_active, busIdx)

    # demands for multiperiod OPF
    Pd = raw.pd_arr
    Qd = raw.qd_arr

    pd_scen = Pd[:, 1]
    qd_scen = Qd[:, 1]
    pd_ref = raw.bus_arr[:, 3]
    qd_ref = raw.bus_arr[:, 4]
    if !check_loads(pd_scen, pd_ref) || !check_loads(qd_scen, qd_ref)
        @warn("Large discrepancy observed between scenarios and MATPOWER's data", maxlog=1)
    end

    if time_horizon_end > 0
        Pd = Pd[:,time_horizon_start:time_horizon_end] .* load_scale
        Qd = Qd[:,time_horizon_start:time_horizon_end] .* load_scale
    end

    return OPFData(buses, lines, generators, bus_ref, baseMVA, Ybus, busIdx, BusGeners, Pd, Qd)
end

function computeAdmitances(lines, buses, baseMVA; lossless::Bool=false, fixedtaps::Bool=false, zeroshunts::Bool=false)
    nlines = length(lines)
    YffR=Array{Float64}(undef, nlines)
    YffI=Array{Float64}(undef, nlines)
    YttR=Array{Float64}(undef, nlines)
    YttI=Array{Float64}(undef, nlines)
    YftR=Array{Float64}(undef, nlines)
    YftI=Array{Float64}(undef, nlines)
    YtfR=Array{Float64}(undef, nlines)
    YtfI=Array{Float64}(undef, nlines)

    for i in 1:nlines
        @assert lines[i].status == 1
        Ys = 1/((lossless ? 0.0 : lines[i].r) + lines[i].x*im)
        #assign nonzero tap ratio
        tap = (lines[i].ratio==0) ? 1.0 : lines[i].ratio
        fixedtaps && (tap = 1.0)

        #add phase shifters
        if (!lossless)
            tap *= exp(lines[i].angle * pi/180 * im)
        end

        Ytt = Ys + lines[i].b/2*im
        Yff = Ytt / (tap*conj(tap))
        Yft = -Ys / conj(tap)
        Ytf = -Ys / tap

        #split into real and imag parts
        YffR[i] = real(Yff); YffI[i] = imag(Yff)
        YttR[i] = real(Ytt); YttI[i] = imag(Ytt)
        YtfR[i] = real(Ytf); YtfI[i] = imag(Ytf)
        YftR[i] = real(Yft); YftI[i] = imag(Yft)

        # if lossless
        #   if !iszero(lines[i].r)
        #     println("warning: lossless assumption changes r from ", lines[i].r, " to 0 for line ", lines[i].from, " -> ", lines[i].to)
        #   end
        #   if !iszero(lines[i].angle)
        #     println("warning: lossless assumption changes angle from ", lines[i].angle, " to 0 for line ", lines[i].from, " -> ", lines[i].to)
        #   end
        # end
    end

    nbuses = length(buses)
    YshR = Array{Float64}(undef, nbuses)
    YshI = Array{Float64}(undef, nbuses)
    for i in 1:nbuses
        YshR[i] = (lossless ? 0.0 : (buses[i].Gs / baseMVA))
        YshI[i] = buses[i].Bs / baseMVA
        zeroshunts && (YshI[i] = 0)
        # if lossless && !iszero(buses[i].Gs)
        #   println("warning: lossless assumption changes Gshunt from ", buses[i].Gs, " to 0 for bus ", i)
        # end
    end

    @assert 0==length(findall(isnan.(YffR)))+length(findall(isinf.(YffR)))
    @assert 0==length(findall(isnan.(YffI)))+length(findall(isinf.(YffI)))
    @assert 0==length(findall(isnan.(YttR)))+length(findall(isinf.(YttR)))
    @assert 0==length(findall(isnan.(YttI)))+length(findall(isinf.(YttI)))
    @assert 0==length(findall(isnan.(YftR)))+length(findall(isinf.(YftR)))
    @assert 0==length(findall(isnan.(YftI)))+length(findall(isinf.(YftI)))
    @assert 0==length(findall(isnan.(YtfR)))+length(findall(isinf.(YtfR)))
    @assert 0==length(findall(isnan.(YtfI)))+length(findall(isinf.(YtfI)))
    @assert 0==length(findall(isnan.(YshR)))+length(findall(isinf.(YshR)))
    @assert 0==length(findall(isnan.(YshI)))+length(findall(isinf.(YshI)))
    if lossless
        @assert 0==length(findall(!iszero, YffR))
        @assert 0==length(findall(!iszero, YttR))
        @assert 0==length(findall(!iszero, YftR))
        @assert 0==length(findall(!iszero, YtfR))
        @assert 0==length(findall(!iszero, YshR))
    end

    return YffR, YffI, YttR, YttI, YftR, YftI, YtfR, YtfI, YshR, YshI
end


text:
using CompScienceMeshes
using BEAST

fn = joinpath(dirname(@__FILE__),"assets","torus.msh")
m = CompScienceMeshes.read_gmsh_mesh(fn)

X = raviartthomas(m)
Y = buffachristiansen(m)

N = BEAST.NCross()
K = Maxwell3D.doublelayer(gamma=0.0)

verts = skeleton(m,0)
edges = skeleton(m,1)
faces = skeleton(m,2)
Λ = connectivity(verts, edges)
Σ = connectivity(faces, edges)

M = assemble(K+0.5N,Y,X)

using LinearAlgebra
s = svdvals(M)

using Plots
plot(log10.(s))

x = nullspace(M, 1.1*s[end])

norm(x)
norm(M*x)
norm(Σ'*Λ)


text:
@testset "MArray" begin
    @testset "Inner Constructors" begin
        @test MArray{(1,),Int,1,1}((1,)).data === (1,)
        @test MArray{(1,),Float64,1,1}((1,)).data === (1.0,)
        @test MArray{(2,2),Float64,2,4}((1, 1.0, 1, 1)).data === (1.0, 1.0, 1.0, 1.0)
        @test isa(MArray{(1,),Int,1,1}(), MArray{(1,),Int,1,1})
        @test isa(MArray{(1,),Int,1}(), MArray{(1,),Int,1,1})
        @test isa(MArray{(1,),Int}(), MArray{(1,),Int,1,1})

        # Bad input
        @test_throws Exception MArray{(2,),Int,1,2}((1,))
        @test_throws Exception MArray{(1,),Int,1,1}(())

        # Bad parameters
        @test_throws Exception MArray{(1,),Int,1,2}((1,))
        @test_throws Exception MArray{(1,),Int,2,1}((1,))
        @test_throws Exception MArray{(1,),1,1,1}((1,))
        @test_throws Exception MArray{(2,),Int,1,1}((1,))
    end

    @testset "Outer constructors and macro" begin
        @test MArray{(1,),Int,1}((1,)).data === (1,)
        @test MArray{(1,),Int}((1,)).data === (1,)
        @test MArray{(1,)}((1,)).data === (1,)

        @test MArray{(2,2),Int,2}((1,2,3,4)).data === (1,2,3,4)
        @test MArray{(2,2),Int}((1,2,3,4)).data === (1,2,3,4)
        @test MArray{(2,2)}((1,2,3,4)).data === (1,2,3,4)

        @test ((@MArray [1])::MArray{(1,)}).data === (1,)
        @test ((@MArray [1,2])::MArray{(2,)}).data === (1,2)
        @test ((@MArray Float64[1,2,3])::MArray{(3,)}).data === (1.0, 2.0, 3.0)
        @test ((@MArray [1 2])::MArray{(1,2)}).data === (1, 2)
        @test ((@MArray Float64[1 2])::MArray{(1,2)}).data === (1.0, 2.0)
        @test ((@MArray [1 ; 2])::MArray{(2,1)}).data === (1, 2)
        @test ((@MArray Float64[1 ; 2])::MArray{(2,1)}).data === (1.0, 2.0)
        @test ((@MArray [1 2 ; 3 4])::MArray{(2,2)}).data === (1, 3, 2, 4)
        @test ((@MArray Float64[1 2 ; 3 4])::MArray{(2,2)}).data === (1.0, 3.0, 2.0, 4.0)

        @test ((@MArray [i*j*k for i = 1:2, j = 2:3, k = 3:4])::MArray{(2,2,2)}).data === (6, 12, 9, 18, 8, 16, 12, 24)
        @test ((@MArray Float64[i*j*k for i = 1:2, j = 2:3, k =3:4])::MArray{(2,2,2)}).data === (6.0, 12.0, 9.0, 18.0, 8.0, 16.0, 12.0, 24.0)

        @test (ex = macroexpand(:(@MArray [1 2; 3])); isa(ex, Expr) && ex.head == :error)

        @test ((@MArray zeros(2,2,1))::MArray{(2,2,1), Float64}).data === (0.0, 0.0, 0.0, 0.0)
        @test ((@MArray ones(2,2,1))::MArray{(2,2,1), Float64}).data === (1.0, 1.0, 1.0, 1.0)
        @test ((@MArray eye(2))::MArray{(2,2), Float64}).data === (1.0, 0.0, 0.0, 1.0)
        @test ((@MArray eye(2,2))::MArray{(2,2), Float64}).data === (1.0, 0.0, 0.0, 1.0)
        @test isa(@MArray(rand(2,2,1)), MArray{(2,2,1), Float64})
        @test isa(@MArray(randn(2,2,1)), MArray{(2,2,1), Float64})

        @test ((@MArray zeros(Float32, 2, 2, 1))::MArray{(2,2,1),Float32}).data === (0.0f0, 0.0f0, 0.0f0, 0.0f0)
        @test ((@MArray ones(Float32, 2, 2, 1))::MArray{(2,2,1),Float32}).data === (1.0f0, 1.0f0, 1.0f0, 1.0f0)
        @test ((@MArray eye(Float32, 2))::MArray{(2,2), Float32}).data === (1.0f0, 0.0f0, 0.0f0, 1.0f0)
        @test ((@MArray eye(Float32, 2, 2))::MArray{(2,2), Float32}).data === (1.0f0, 0.0f0, 0.0f0, 1.0f0)
        @test isa(@MArray(rand(Float32, 2, 2, 1)), MArray{(2,2,1), Float32})
        @test isa(@MArray(randn(Float32, 2, 2, 1)), MArray{(2,2,1), Float32})
    end

    @testset "Methods" begin
        m = @MArray [11 13; 12 14]

        @test isimmutable(m) == false

        @test m[1] === 11
        @test m[2] === 12
        @test m[3] === 13
        @test m[4] === 14

        @test Tuple(m) === (11, 12, 13, 14)

        @test size(m) === (2, 2)
        @test size(typeof(m)) === (2, 2)
        @test size(MArray{(2,2),Int,2}) === (2, 2)
        @test size(MArray{(2,2),Int}) === (2, 2)
        @test size(MArray{(2,2)}) === (2, 2)

        @test size(m, 1) === 2
        @test size(m, 2) === 2
        @test size(typeof(m), 1) === 2
        @test size(typeof(m), 2) === 2

        @test length(m) === 4
    end

    @testset "setindex!" begin
        v = @MArray [1,2,3]
        v[1] = 11
        v[2] = 12
        v[3] = 13
        @test v.data === (11, 12, 13)

        m = @MArray [0 0; 0 0]
        m[1] = 11
        m[2] = 12
        m[3] = 13
        m[4] = 14
        @test m.data === (11, 12, 13, 14)
    end
end


text:
export Ball

struct Ball{K} <: AbstractSimpleShape{K}
    c::SFloat{K}  # center of ball
    r::Float  # radius
    Ball{K}(c,r) where {K} = new(c,r)  # suppress default outer constructor
end

Ball(c::AbsVecReal, r::Real) = Ball{length(c)}(c,r)

Base.:(==)(s1::Ball, s2::Ball) = s1.c==s2.c && s1.r==s2.r
Base.isapprox(s1::Ball, s2::Ball) = s1.c≈s2.c && s1.r≈s2.r
Base.hash(s::Ball, h::UInt) = hash(s.c, hash(s.r, hash(:Ball, h)))

level(x::SReal{K}, s::Ball{K}) where {K} = norm(x-s.c) / s.r - 1.0
center(s::Ball) = s.c
bounds(s::Ball) = (s.c .- s.r, s.c .+ s.r)

max_stars_repo_name: wsshin/LevelSetShapes

max_issues_repo_name: wsshin/LevelSetShapes

max_forks_repo_name: wsshin/LevelSetShapes


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.7.1
#   kernelspec:
#     display_name: Julia 1.2.0
#     language: julia
#     name: julia-1.2
# ---

# # Simulation of COVID-19/pandemic influenza outbreaks in school
#
# ## Preparation

include("SchoolOutbreak.jl")

# +
using Distributions, Plots, InvertedIndices, LinearAlgebra, StatsBase
gr(fontfamily="Helvetica",foreground_color_legend = nothing,background_color_legend = nothing, titlefontsize=9,grid=false, tick_direction=:out)

classes=[2,4,2,1]
sizes=[40,20,20,40]

nclassmat=repeat(classes,1,6)
nstudentsmat=repeat(sizes,1,6).*nclassmat
struct ImproperUniform<: Distribution{Univariate,Discrete} end
Distributions.cdf(d::ImproperUniform, x::Integer)=1.0
Distributions.pdf(d::ImproperUniform, x::Integer)=1.0

rcomdist=Uniform(0,365)
hhfoi=-Inf
Gammameansd(mean,sd)=Gamma(mean^2/sd^2,sd^2/mean)
#gtimedist=Gammameansd(3,1)
gtimedist=Gammameansd(1.7,1) # H3N2 serial interval from Vink et al 2014 (https://doi.org/10.1093/aje/kwu209)
N0=[30,30,30,30,30] # normalizing constant to avoid correlation between β and γ
C0=[0,6*3,3,3,3]

sus_covlabels=[:Grade, :Vaccine_curr, :NPI_mask,:NPI_wash]
inf_covlabels=[:Grade, :Vaccine_curr, :NPI_mask,:NPI_wash,:duringbreak]

strata=vcat.(1:4,sortslices.(repeat.((x->reduce(hcat,vec(collect.(x)))).(Iterators.product.(Ref(1:6),range.(1,classes,step=1)).|>collect),1,sizes),dims=2))
strata=vcat.(strata,repeat.(Ref((1:2)'),1,size.(strata,2).÷2))
strata=reduce(hcat,strata)
parameters=SchoolOutbreak.makeparameters(ones(5),ones(5)/2,ones(5)/2,0.01,
    ones(size(sus_covlabels)),ones(size(inf_covlabels)),gtimedist,20,rcomdist,strata,nclassmat,nstudentsmat,N0,C0,Inf)
@time students=SchoolOutbreak.studentsfromdata(parameters,false,fill(typemax(Int),size(strata,2)),eachcol(strata),
    Ref(zeros(length(sus_covlabels))),Ref(zeros(length(inf_covlabels))),nclassmat,nstudentsmat,-Inf,1);
@time simstudents=SchoolOutbreak.Students.([filter(x->x.stratum[1]==schoolid,students) for schoolid in 1:4],Ref(parameters));
# -

using Mamba
chain1 = read("Mastumoto_chain1.jls", ModelChains)

# preparation for simulation
function parametersfromchain(chain, iter)
    (β=SchoolOutbreak.posterior(chain,:βs,iter).samples[:,[1;1:4]]|>vec,
    γ=SchoolOutbreak.posterior(chain,:γs,iter).samples[:,[1;1:3;3]]|>vec,
    δ=SchoolOutbreak.posterior(chain,:δs,iter).samples[:,[1;1:3;3]]|>vec,
    rcom=SchoolOutbreak.posterior(chain,:rcom,iter).samples|>first|>fill,
    suscoef=SchoolOutbreak.posterior(chain,:suscoef,iter).samples|>vec,
    infcoef=SchoolOutbreak.posterior(chain,:infcoef,iter).samples|>vec)
end
nsample=500
@time updateparameters=parametersfromchain.(Ref(chain1),(1:nsample).*(length(chain1.range)÷(2*nsample)));
(x->x.rcom.=0).(updateparameters)
(x->x.members[1].isinfected=true).(simstudents)
(x->x.members[1].onset=1).(simstudents)

# ## Parameter settings

# +
# Adjust R0
function R0set!(parameters,R0,nclasses=2,classsize=30)
    parameters.β.*=R0/sum((parameters.β./(classsize/30).^parameters.γ./(nclasses./[3,3,3,3,3]).^parameters.δ).*[0,5*nclasses*classsize,(nclasses-1)*classsize,classsize/2,classsize/2])

end
R0set!.(updateparameters,2);

function R0(students)
    sum((x->x.pairwiseβ[]/exp(x.pairwiselogN[]*x.γ[]+x.pairwiselogM[]*x.δ[])).(students.pairs),dims=2)
end

# +
###### infection hazard
using RCall
@rimport distr

## COVID-19
incperiod=distr.Lnorm(1.434065,0.6612) # He et al.
infprof=distr.Norm(0.53,sqrt(7)) # Aschroft et al.
gtime=incperiod+infprof

## Pandemic influenza
#incperiod=distr.Lnorm(0.34,0.42) # Lessler et al. 2009 (https://dx.doi.org/10.1016%2FS1473-3099(09)70069-6)
#gtime=((mean,sd)->distr.Gammad(mean^2/sd^2, sd^2/mean))(1.7,1)
#

gtimevec=distr.p(gtime)((0:20))|>RCall.unsafe_vec|>diff
gtimevec=gtimevec./sum(gtimevec)

# set hazard function
symratio=(0:0.25:1)'
hazards_symsc=gtimevec.*(1 .- symratio.*(distr.p(incperiod)((1:20).-1)|>RCall.unsafe_vec))
surv=(0:0.1:0.5)'
hazards_randsc=[hazard.*((1 .-surv).^(0:19)) for hazard in eachcol(hazards_symsc)]

# enbed generation time distribution in parameters
function setgtime!(parameters,newpdf)
    SchoolOutbreak.updateArray!(parameters.pdfgtime,[0;newpdf;0])
    SchoolOutbreak.updateArray!(parameters.cdfgtime,[0;0;cumsum(newpdf)])
end
setgtime!(parameters,hazards_randsc[3][:,2]); # 50% symptomatic + 10% testing rate
# -

# ## Outbreak simulation

incidence(outbreak)=length(outbreak)
incidence(outbreak::AbstractArray{<:AbstractVector})=incidence.(outbreak)
incidence(outbreak::AbstractArray{<:AbstractMatrix})=sum.(outbreak)
function getgradeclass(student) return((student.stratum[2],student.stratum[3])) end
function reorderedclass(class,gradeorder,classorder) # returns reordered class id
    grnum=findfirst(x->x==class[1],gradeorder)
    clnum=findfirst(x->x==class[2],classorder[class[1]])
    (grnum,clnum)
end
function updateorders!(gradeorder,classorder,newstudents)
    gcs=getgradeclass.(newstudents)
    push!(gradeorder,filter(x->!(x in gradeorder),unique(getindex.(gcs,1)))...) # push nonregistered grades
    for grade in 1:6
        push!(classorder[grade],unique(getindex.(filter(x->(x[1]==grade)&&!(x[2] in classorder[grade]),gcs),2))...)
    end
end
function incidencebyclass_reorder!(members,idsvector,ngrades,nclasses, closuretime=nothing)
    gradeorder=Int[]
    classorder=[Int[] for x in 1:6]
    reorderedincidence=[((members, ids)->
        begin
            casesbyclass=zeros(Int,ngrades,nclasses)
            updateorders!(gradeorder,classorder, members[ids])
            reorderedgcs=reorderedclass.(getgradeclass.(members[ids]),Ref(gradeorder),Ref(classorder))
            for gc in reorderedgcs
                casesbyclass[CartesianIndex(gc)].+=1
            end
            casesbyclass
        end)(members,ids)
    for ids in idsvector]
    if isnothing(closuretime) return(reorderedincidence) end # return reordered incidence only if closuretime is not passed
    reorderedclosuretime=[zeros(Int,ngrades,nclasses) for t in 1:length(idsvector)]
    for (g,closure,corder) in zip(1:6,collect(eachrow(closuretime)),classorder)
        if g in gradeorder
            setindex!.(reorderedclosuretime, (getindex.(closure[corder],t) for t in 1:length(idsvector)), findfirst(x->x==g,gradeorder), Ref(corder))
        end
    end
    (incidence=reorderedincidence, closuretime=reorderedclosuretime)
end
function readoutbreakbyclass_reorder(outbreak, students::Vector{<:SchoolOutbreak.Students}, ngrades,nclasses, closuretime=nothing)
    reordered=incidencebyclass_reorder!.(getfield.(students,:members),outbreak,ngrades,nclasses, closuretime)
    if isnothing(closuretime) return(reordered) end
    (incidence=getindex.(reordered,:incidence), closuretime=getindex.(reordered,:closuretime))
end
function outbreakscenarios!(simstudents,updateparameters,R0,gtime)
    R0set!.(updateparameters,Ref(R0))
    setgtime!(simstudents[1].parameters,gtime)
    outbreaks=SchoolOutbreak.simulateoutbreaks!.(Ref(simstudents), updateparameters, Ref(1:360))
    R0set!.(updateparameters,Ref(1 ./R0)) # to fix the breaking change when R0 is a vector
    readoutbreakbyclass_reorder.(outbreaks,Ref(simstudents),6,Ref([2,4,2,1]))
end

gtimes=[hazards_symsc[:,1],hazards_symsc[:,3], hazards_randsc[3][:,2]]|>permutedims
@time scenarios=outbreakscenarios!.(Ref(simstudents),Ref(updateparameters),[1.8,1.2,0.8],gtimes);
gtimes_dayoff=[hazards_symsc[:,1].*repeat(0:1,size(hazards_symsc,1))[1:size(hazards_symsc,1)],hazards_symsc[:,3].*repeat(0:1,size(hazards_symsc,1))[1:size(hazards_symsc,1)],hazards_symsc[:,3].*repeat([1,1,0],size(hazards_symsc,1))[1:size(hazards_symsc,1)]]|>permutedims
@time scenarios_dayoff=outbreakscenarios!.(Ref(simstudents),Ref(updateparameters),[1.8,1.2,0.8],gtimes_dayoff);
classisolation=[[fill(0.5,3);1;1], [fill(0.5,3);1.2;1.2], [fill(0.1,3);1;1], [fill(0.1,3);1.4;1.4]]|>permutedims
@time scenarios_classisolation=outbreakscenarios!.(Ref(simstudents),Ref(updateparameters),[1.8,1.2,0.8].*classisolation,Ref(gtimes[1]));

# ## Simulation results

Base.:+(x::AbstractRange,y::Number)=x.+y
Base.:-(x::AbstractRange,y::Number)=x.-y
outbreakmaps=[(x->reduce(hcat,vec.(x'))).(mean(scenario))./[40,20,20,40] for scenario in scenarios]
heatmaps=Plots.heatmap.(reduce(vcat,vec(outbreakmaps)),clim=((0,0.2/40)),color=cgrad([:white,:red,:purple,:indigo,:black]),legend=false)
yticks!.(heatmaps,repeat(Ref((1:6) .-0.5).*[2,4,2,1].+0.5,9),Ref(string.(1:6).*""),tickfontsize=6)
Plots.plot!.(heatmaps,repeat(Ref(1:5).*[2,4,2,1].+0.5,9),linetype=:hline,color=:black,linealpha=0.2)
Plots.plot(heatmaps...,layout=(9,4),size=(800,1200),left_margin=2mm,bottom_margin=-2mm, xlabel="time",ylabel="classes/grades",guidefontsize=6,tickfontsize=6)

outbreakmaps=[(x->reduce(hcat,vec.(x'))).(mean(scenario))./[40,20,20,40] for scenario in scenarios_dayoff]
heatmaps=Plots.heatmap.(reduce(vcat,vec(outbreakmaps)),clim=((0,0.2/40)),color=cgrad([:white,:red,:purple,:indigo,:black]),legend=false)
yticks!.(heatmaps,repeat(Ref((1:6) .-0.5).*[2,4,2,1].+0.5,9),Ref(string.(1:6).*""),tickfontsize=6)
Plots.plot!.(heatmaps,repeat(Ref(1:5).*[2,4,2,1].+0.5,9),linetype=:hline,color=:black,linealpha=0.2)
Plots.plot(heatmaps...,layout=(9,4),size=(800,1200),left_margin=2mm,bottom_margin=-2mm, xlabel="time",ylabel="classes/grades",guidefontsize=6,tickfontsize=6)

outbreakmaps=[(x->reduce(hcat,vec.(x'))).(mean(scenario))./[40,20,20,40] for scenario in scenarios_classisolation]
heatmaps=Plots.heatmap.(reduce(vcat,vec(outbreakmaps)),clim=((0,0.2/40)),color=cgrad([:white,:red,:purple,:indigo,:black]),legend=false)
yticks!.(heatmaps,repeat(Ref((1:6) .-0.5).*[2,4,2,1].+0.5,12),Ref(string.(1:6).*""),tickfontsize=6)
Plots.plot!.(heatmaps,repeat(Ref(1:5).*[2,4,2,1].+0.5,12),linetype=:hline,color=:black,linealpha=0.2)
Plots.plot(heatmaps...,layout=(12,4),size=(800,1600),left_margin=2mm,bottom_margin=-2mm, xlabel="time",ylabel="classes/grades",guidefontsize=6,tickfontsize=6)

# ## Responsive management strategy simulation

Base.:^(x,y::AbstractDict)=Dict(keys(y).=>x.^values(y))
detectbinom(d::AbstractDict,prob)=Dict(keys(d).=>rand.(Binomial.(values(d),prob)))
struct Strategy{V,F,D}
    detectdist::V
    detectprob::F
    closurelength::Integer
    closuretime::D
end
Broadcast.broadcastable(m::Strategy) = Ref(m)
function getclosuretime(strategy::Strategy,student::SchoolOutbreak.Student)
    strategy.closuretime[student.stratum[2:3]...]
end
function canceltransmission!(student,frompairs,strategy,currenttime)
    infectedpairs=filter(x->x.isinfected,frompairs)
    if length(infectedpairs)==0 return(nothing) end
    infectiondates=getfield.(getfield.(infectedpairs,:from),:onset).+getfield.(infectedpairs,:serialinterval)

    fromid=(infectiondates.==currenttime)
    if getclosuretime(strategy,student)[currenttime]!=0 || all(getindex.(getclosuretime.(strategy,getfield.(infectedpairs[fromid],:from)),currenttime).!=0)
        student.onset=minimum([filter(x->x>currenttime,infectiondates);typemax(Int)])
        student.isinfected=length(filter(x->x>currenttime,infectiondates))!=0 #sum(fromid)!=length(infectiondates)
    end
end
function (strategy::Strategy)(students,currenttime)
    infectorsid=findall(x->x.onset==currenttime,students.members)
    if length(infectorsid)==0 return(nothing) end # abort if no infected person
    infectors=students.members[infectorsid]
    frompairs=eachcol(@view students.pairs[:,infectorsid])

    canceltransmission!.(infectors, frompairs, strategy, currenttime) # cut transmissions during closure

    detectedsdict=detectbinom(countmap(getfield.(infectors,:stratum)),strategy.detectprob)
    for classdetected in detectedsdict
        if last(classdetected)!=0 # if at least one case was detected
            closurefrom=currenttime+minimum(rand(strategy.detectdist,last(classdetected)))
            closuretimes=strategy.closuretime[first(classdetected)[2:3]...]
            if closurefrom>length(closuretimes)||closuretimes[closurefrom]!=0 continue end # if already closed, no change
            closuretimes[closurefrom+1:min(closurefrom+strategy.closurelength,length(closuretimes))].=reverse(1:min(strategy.closurelength,length(closuretimes)-closurefrom))
            if closurefrom+strategy.closurelength<length(closuretimes) && closuretimes[closurefrom+strategy.closurelength+1]!=0
                @views closuretimes[closurefrom+1:min(closurefrom+closuretimes[closurefrom+1],length(closuretimes))].=0
            end
        end
    end
    (detectedsdict, countmap(getfield.(infectors,:stratum)),strategy.detectprob)
end
function outbreakscenarios!(simstudents,updateparameters,R0,gtime,strategy)
    R0set!.(updateparameters,Ref(R0))
    setgtime!(simstudents[1].parameters,gtime)
    outbreaks=SchoolOutbreak.simulateoutbreaks!.(Ref(simstudents), updateparameters, Ref(1:360), statuscheck! = strategy)
    R0set!.(updateparameters,Ref(1 ./R0)) # to fix the breaking change when R0 is a vector
    incidences=(x->getindex.(x,:incidence)).(outbreaks)
    closuretimes=(x->getindex.(x,:closuretime)).(outbreaks)
    reordered=readoutbreakbyclass_reorder.(incidences,Ref(simstudents),6,Ref([2,4,2,1]),closuretimes)
    (incidence=getindex.(reordered,:incidence), closuretime=getindex.(reordered,:closuretime))
end

# +
gtimes=[hazards_symsc[:,1]]
ipdist=LogNormal(1.434065,0.6612) #COVID-19
#ipdist=LogNormal(0.34,0.42) #Flu

function detectdistribution(ipdist, testingrate, symptomaticratio, trange=0:19)
    probmass=1 .-(1 .-symptomaticratio*cdf.(ipdist,trange)).*(1-testingrate).^(trange)
    probmass=[0;probmass]|>diff
    (dist=Categorical(normalize(probmass,1)), prob=sum(probmass))
end
detect=(x->[first.(x),last.(x)])(detectdistribution.(ipdist, [0,0.1], 0.5)) # replace 0.5 with 0.25 for lower symptomatic ratio
strategies=[Strategy.(detect[1],detect[2], 14, Ref([zeros(Int,360) for i in 1:6, j in 1:classsize])) for classsize in (2,4,2,1)] # replace 14 with 7 for flu
@time scenarios_strategies_notest=outbreakscenarios!.(Ref(simstudents),Ref(updateparameters),[1.8,1.2,0.8].*[1 classisolation],gtimes,Ref(first.(strategies)));
@time scenarios_strategies_10test=outbreakscenarios!.(Ref(simstudents),Ref(updateparameters),[1.8,1.2,0.8].*[1 classisolation],gtimes,Ref(last.(strategies)));
# -

using JLD2
#@save "closurestrategy_covid.jld2" scenarios_strategies_notest scenarios_strategies_10test

# ## Simulation results
# ### No regular testing

outbreakmaps=[(x->reduce(hcat,vec.(x'))).(mean(scenario))./[40,20,20,40] for scenario in first.(scenarios_strategies_notest)]
heatmaps=Plots.heatmap.(reduce(vcat,vec(outbreakmaps)),clim=((0,0.2/40)),color=cgrad([:white,:red,:purple,:indigo,:black]),legend=false)
yticks!.(heatmaps,repeat(Ref((1:6) .-0.5).*[2,4,2,1].+0.5,15),Ref(string.(1:6).*""),tickfontsize=6)
Plots.plot!.(heatmaps,repeat(Ref(1:5).*[2,4,2,1].+0.5,15),linetype=:hline,color=:black,linealpha=0.2)
Plots.plot(heatmaps...,layout=(15,4),size=(800,2000),left_margin=2mm,bottom_margin=-2mm, xlabel="time",ylabel="classes/grades",guidefontsize=6,tickfontsize=6)

closuremaps=[(x->reduce(hcat,vec.(x'))).(mean(scenario)) for scenario in last.(scenarios_strategies_notest)]
heatmaps=Plots.heatmap.(reduce(vcat,vec(closuremaps)),clim=((0,0.2)),color=cgrad([:white,:skyblue,:blue,:indigo,:black]),legend=false)
yticks!.(heatmaps,repeat(Ref((1:6) .-0.5).*[2,4,2,1].+0.5,15),Ref(string.(1:6).*""),tickfontsize=6)
Plots.plot!.(heatmaps,repeat(Ref(1:5).*[2,4,2,1].+0.5,15),linetype=:hline,color=:black,linealpha=0.2)
Plots.plot(heatmaps...,layout=(15,4),size=(800,2000),left_margin=2mm,bottom_margin=-2mm, xlabel="time",ylabel="classes/grades",guidefontsize=6,tickfontsize=6)

# ### With 10% regular testing

outbreakmaps=[(x->reduce(hcat,vec.(x'))).(mean(scenario))./[40,20,20,40] for scenario in first.(scenarios_strategies_10test)]
heatmaps=Plots.heatmap.(reduce(vcat,vec(outbreakmaps)),clim=((0,0.2/40)),color=cgrad([:white,:red,:purple,:indigo,:black]),legend=false)
yticks!.(heatmaps,repeat(Ref((1:6) .-0.5).*[2,4,2,1].+0.5,15),Ref(string.(1:6).*""),tickfontsize=6)
Plots.plot!.(heatmaps,repeat(Ref(1:5).*[2,4,2,1].+0.5,15),linetype=:hline,color=:black,linealpha=0.2)
Plots.plot(heatmaps...,layout=(15,4),size=(800,2000),left_margin=2mm,bottom_margin=-2mm, xlabel="time",ylabel="classes/grades",guidefontsize=6,tickfontsize=6)

closuremaps=[(x->reduce(hcat,vec.(x'))).(mean(scenario)) for scenario in last.(scenarios_strategies_10test)]
heatmaps=Plots.heatmap.(reduce(vcat,vec(closuremaps)),clim=((0,0.2)),color=cgrad([:white,:skyblue,:blue,:indigo,:black]),legend=false)
yticks!.(heatmaps,repeat(Ref((1:6) .-0.5).*[2,4,2,1].+0.5,15),Ref(string.(1:6).*""),tickfontsize=6)
Plots.plot!.(heatmaps,repeat(Ref(1:5).*[2,4,2,1].+0.5,15),linetype=:hline,color=:black,linealpha=0.2)
Plots.plot(heatmaps...,layout=(15,4),size=(800,2000),left_margin=2mm,bottom_margin=-2mm, xlabel="time",ylabel="classes/grades",guidefontsize=6,tickfontsize=6)


text:
module ElementaryChemistry

include("GassLaws.jl")
include("conversions/TempratureConversion.jl")
include("conversions/PressureConversion.jl")

export charles_law, boyles_law_solve_pressure, boyles_law_solve_volume, gay_lussacs_law, avogadros_law, absolute_temperature,
grahams_law_velocity_ratio
end


text:
export GAEvo, evaluate, populate

"""

GAEvo implements a classic genetic algorithm. To do this, a new subtype of
AbstractEvolution is created, GAEvo, and the following functions are
defined for this type:

evaluate
populate

the function generation can also be defined. Without specification, the default
generation function (which is void) is applied.
"""
mutable struct GAEvo{T} <: AbstractEvolution
    config::NamedTuple
    logger::CambrianLogger
    population::Array{T}
    fitness::Function
    gen::Int
end

function GAEvo{T}(cfg::NamedTuple, fitness::Function;
                  logfile=string("logs/", cfg.id, ".csv")) where T
    logger = CambrianLogger(logfile)
    population = initialize(T, cfg)
    GAEvo(cfg, logger, population, fitness, 0)
end

evaluate(e::GAEvo) = fitness_evaluate(e, e.fitness)

"""
populate(e::GAEvo)
"""
function ga_populate(e::AbstractEvolution)
    new_pop = Array{Individual}(undef, 0)
    if e.config.n_elite > 0
        sort!(e.population)
        append!(new_pop,
                e.population[(length(e.population)-e.config.n_elite+1):end])
    end
    for i in (e.config.n_elite+1):e.config.n_population
        p1 = selection(e.population)
        child = deepcopy(p1)
        if e.config.p_crossover > 0 && rand() < e.config.p_crossover
            parents = vcat(p1, [selection(e.population) for i in 2:e.config.n_parents])
            child = crossover(parents...)
        end
        if e.config.p_mutation > 0 && rand() < e.config.p_mutation
            child = mutate(child)
        end
        push!(new_pop, child)
    end
    e.population = new_pop
end

function populate(e::GAEvo)
    ga_populate(e)
end


text:
#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

function componentquery!(state::SmartsParserState)
    """ Start <- Component ('.' Component)*
    """
    if read(state) == '\0'
        return # Empty query
    end
    component!(state)
    while !state.done && read(state) == '.'
        forward!(state)
        state.root = state.node + 1
        component!(state)
    end
end

function component!(state::SmartsParserState)
    """ Component <- '(' Fragment ')' / Fragment
    """
    if read(state) == '('
        # Connectivity restriction
        forward!(state)
        push!(state.connectivity, [state.root])
        fragment!(state)
        c2 = read(state)
        @assert c2 == ')' "unexpected token: $(c2) at $(state.pos)"
        forward!(state)
    else
        fragment!(state)
    end
end

function fragment!(state::SmartsParserState)
    """ Fragment <- Group
    """
    if read(state) == '\0'
        return # Empty molecule
    end
    group!(state, nothing)
    # Validity check
    if state.node + 1 == state.root
        throw(ErrorException(
            "unexpected token: $(read(state)) at $(state.pos)"))
    elseif length(state.ringlabel) > 0
        label = collect(keys(state.ringlabel))[1]
        throw(ErrorException("unclosed ring: $(label)"))
    end
end

function group!(state::SmartsParserState, bond)
    """ Group <- Atom ((Bond? Group) / Chain)* Chain
    """
    a = atom!(state)
    if isempty(a)
        # ex. CC((CC)C)C
        throw(ErrorException(
            "unexpected token: branch starts with '(' at $(state.pos)"))
    end
    state.node += 1
    push!(state.nodeattrs, popfirst!(a))
    if bond !== nothing
        # Connect branch
        push!(state.edges, (state.branch, state.node))
        push!(state.edgeattrs, bond)
    end
    center = state.node
    for h in a
        state.node += 1
        push!(state.nodeattrs, h)
        push!(state.edges, (center, state.node))
        push!(state.edgeattrs, SmilesBond())
    end
    state.branch = center
    while true
        if read(state) == '('
            forward!(state)
            buf = state.branch
            b = something(bond!(state), defaultbond(state))
            group!(state, b)
            state.branch = buf
            c = read(state)
            @assert c == ')' "unexpected token: $(c) at $(state.pos)"
            forward!(state)

            # ex. CC(C)(C) should be CC(C)C but acceptable
            raw"""
            if state.done || read(state) in ")."
                throw(ErrorException(
                    "unexpected token: branch ends with ')' at $(state.pos)"))
            end
            """
        else
            ccen = chain!(state)
            state.branch = ccen
            if state.done || read(state) in ")."
                break
            end
        end
    end
end

function chain!(state::SmartsParserState)
    """ Chain <- (Bond? (Atom / RingLabel))+
    """
    u = state.branch
    while !state.done
        # Bond?
        if read(state) == '.'
            if !state.allow_disconnected
                throw(ErrorException(
                    "unexpected token: disconnected at $(state.pos)"))
            elseif lookahead(state, 1) != '('
                # Disconnected
                forward!(state)
                b = :disconn
            else
                b = bond!(state)
            end
        else
            b = bond!(state)
        end

        # RingLabel
        if isdigit(read(state)) || read(state) == '%'
            if read(state) == '%'
                forward!(state)
                start = state.pos
                while isdigit(lookahead(state, 1))
                    forward!(state)
                end
                num = parse(Int, SubString(state.input, start, state.pos))
                forward!(state)
            else
                num = parse(Int, read(state))
                forward!(state)
            end
            if num in keys(state.ringlabel)
                (v, rb) = state.ringlabel[num]
                b = something(b, rb, defaultbond(state))
                delete!(state.ringlabel, num) # Ring label is reusable
                push!(state.edges, (u, v))
                push!(state.edgeattrs, b)
            else
                state.ringlabel[num] = (u, b)
            end
            continue
        end

        # Atom
        a = atom!(state)
        if isempty(a)
            c = read(state)
            @assert c in "().\0" "unexpected token: $(c) at $(state.pos)"
            if b == :disconn
                throw(ErrorException(
                    "unexpected token: $(read(state)) at $(state.pos)"))
            end
            break
        else
            state.node += 1
            push!(state.nodeattrs, popfirst!(a))
        end
        if b == :disconn
            if state isa SmartsParser
                for conn in state.connectivity
                    if conn[1] == state.root
                        push!(conn, state.node)
                    end
                end
            end
        else
            b = something(b, defaultbond(state))
            push!(state.edges, (u, state.node))
            push!(state.edgeattrs, b)
        end
        center = state.node
        for h in a
            state.node += 1
            push!(state.edges, (center, state.node))
            push!(state.edgeattrs, SmilesBond())
            push!(state.nodeattrs, h)
        end
        u = center
    end
    return u
end


text:
"""
    AbstractSphere{𝔽} <: AbstractEmbeddedManifold{𝔽,DefaultIsometricEmbeddingType}

An abstract type to represent a unit sphere that is represented isometrically in the embedding.
"""
abstract type AbstractSphere{𝔽} <: AbstractEmbeddedManifold{𝔽,DefaultIsometricEmbeddingType} end

@doc raw"""
    Sphere{n,𝔽} <: AbstractSphere{𝔽}

The (unit) sphere manifold $𝕊^{n}$ is the set of all unit norm vectors in $𝔽^{n+1}$.
The sphere is represented in the embedding, i.e.

````math
𝕊^{n} := \bigl\{ p \in 𝔽^{n+1}\ \big|\ \lVert p \rVert = 1 \bigr\}
````

where $𝔽\in\{ℝ,ℂ,ℍ\}$. Note that compared to the [`ArraySphere`](@ref), here the
argument `n` of the manifold is the dimension of the manifold, i.e. $𝕊^{n} ⊂ 𝔽^{n+1}$, $n\in ℕ$.

The tangent space at point $p$ is given by

````math
T_p𝕊^{n} := \bigl\{ X ∈ 𝔽^{n+1}\ |\ \Re(⟨p,X⟩) = 0 \bigr \},
````

where $𝔽\in\{ℝ,ℂ,ℍ\}$ and $⟨\cdot,\cdot⟩$ denotes the inner product in the
embedding $𝔽^{n+1}$.

For $𝔽=ℂ$, the manifold is the complex sphere, written $ℂ𝕊^n$, embedded in $ℂ^{n+1}$.
$ℂ𝕊^n$ is the complexification of the real sphere $𝕊^{2n+1}$.
Likewise, the quaternionic sphere $ℍ𝕊^n$ is the quaternionification of the real sphere
$𝕊^{4n+3}$.
Consequently, $ℂ𝕊^0$ is equivalent to $𝕊^1$ and [`Circle`](@ref), while $ℂ𝕊^1$ and $ℍ𝕊^0$
are equivalent to $𝕊^3$, though with different default representations.

This manifold is modeled as a special case of the more general case, i.e. as an embedded
manifold to the [`Euclidean`](@ref), and several functions like the [`inner`](@ref inner(::Euclidean, ::Any...)) product
and the [`zero_tangent_vector`](@ref zero_tangent_vector(::Euclidean, ::Any...)) are inherited from the embedding.

# Constructor

    Sphere(n[, field=ℝ])

Generate the (real-valued) sphere $𝕊^{n} ⊂ ℝ^{n+1}$, where `field` can also be used to
generate the complex- and quaternionic-valued sphere.
"""
struct Sphere{N,𝔽} <: AbstractSphere{𝔽} end
Sphere(n::Int, field::AbstractNumbers = ℝ) = Sphere{n,field}()

@doc raw"""
    ArraySphere{T<:Tuple,𝔽} <: AbstractSphere{𝔽}

The (unit) sphere manifold $𝕊^{n₁,n₂,...,nᵢ}$ is the set of all unit (Frobenius) norm elements of
$𝔽^{n₁,n₂,...,nᵢ}$, where $𝔽\in\{ℝ,ℂ,ℍ\}. The generalized sphere is
represented in the embedding, and supports arbitrary sized arrays or in other words arbitrary
tensors of unit norm. The set formally reads

````math
𝕊^{n_1, n_2, …, n_i} := \bigl\{ p \in 𝔽^{n_1, n_2, …, n_i}\ \big|\ \lVert p \rVert = 1 \bigr\}
````

where $𝔽\in\{ℝ,ℂ,ℍ\}$. Setting $i=1$ and $𝔽=ℝ$  this  simplifies to unit vectors in $ℝ^n$, see
[`Sphere`](@ref) for this special case. Note that compared to this classical case,
the argument for the generalized case here is given by the dimension of the embedding.
This means that `Sphere(2)` and `ArraySphere(3)` are the same manifold.

The tangent space at point $p$ is given by

````math
T_p 𝕊^{n_1, n_2, …, n_i} := \bigl\{ X ∈ 𝔽^{n_1, n_2, …, n_i}\ |\ \Re(⟨p,X⟩) = 0 \bigr \},
````

where $𝔽\in\{ℝ,ℂ,ℍ\}$ and $⟨\cdot,\cdot⟩$ denotes the (Frobenius) inner product in the
embedding $𝔽^{n_1, n_2, …, n_i}$.

This manifold is modeled as an embedded manifold to the [`Euclidean`](@ref), i.e.
several functions like the [`inner`](@ref inner(::Euclidean, ::Any...)) product and the
[`zero_tangent_vector`](@ref zero_tangent_vector(::Euclidean, ::Any...)) are inherited from the embedding.

# Constructor

    ArraySphere(n₁,n₂,...,nᵢ; field=ℝ)

Generate sphere in $𝔽^{n_1, n_2, …, n_i}$, where $𝔽$ defaults to the real-valued case $ℝ$.
"""
struct ArraySphere{N,𝔽} <: AbstractSphere{𝔽} where {N<:Tuple} end
function ArraySphere(n::Vararg{Int,I}; field::AbstractNumbers = ℝ) where {I}
    return ArraySphere{Tuple{n...},field}()
end

"""
    check_manifold_point(M::AbstractSphere, p; kwargs...)

Check whether `p` is a valid point on the [`AbstractSphere`](@ref) `M`, i.e. is a point in
the embedding of unit length.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_manifold_point(M::AbstractSphere, p; kwargs...)
    mpv = invoke(
        check_manifold_point,
        Tuple{(typeof(get_embedding(M))),typeof(p)},
        get_embedding(M),
        p;
        kwargs...,
    )
    mpv === nothing || return mpv
    if !isapprox(norm(p), 1.0; kwargs...)
        return DomainError(
            norm(p),
            "The point $(p) does not lie on the $(M) since its norm is not 1.",
        )
    end
    return nothing
end

"""
    check_tangent_vector(M::AbstractSphere, p, X; check_base_point = true, kwargs... )

Check whether `X` is a tangent vector to `p` on the [`AbstractSphere`](@ref) `M`, i.e.
after [`check_manifold_point`](@ref)`(M,p)`, `X` has to be of same dimension as `p`
and orthogonal to `p`.
The optional parameter `check_base_point` indicates, whether to call
[`check_manifold_point`](@ref)  for `p` or not.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_tangent_vector(M::AbstractSphere, p, X; check_base_point = true, kwargs...)
    if check_base_point
        mpe = check_manifold_point(M, p; kwargs...)
        mpe === nothing || return mpe
    end
    mpv = invoke(
        check_tangent_vector,
        Tuple{typeof(get_embedding(M)),typeof(p),typeof(X)},
        get_embedding(M),
        p,
        X;
        check_base_point = false, # already checked above
        kwargs...,
    )
    mpv === nothing || return mpv
    if !isapprox(abs(real(dot(p, X))), 0.0; kwargs...)
        return DomainError(
            abs(dot(p, X)),
            "The vector $(X) is not a tangent vector to $(p) on $(M), since it is not orthogonal in the embedding.",
        )
    end
    return nothing
end

function decorated_manifold(M::AbstractSphere{𝔽}) where {𝔽}
    return Euclidean(representation_size(M)...; field = 𝔽)
end
get_embedding(M::AbstractSphere{𝔽}) where {𝔽} = decorated_manifold(M)

@doc raw"""
    distance(M::AbstractSphere, p, q)

Compute the geodesic distance betweeen `p` and `q` on the [`AbstractSphere`](@ref) `M`.
The formula is given by the (shorter) great arc length on the (or a) great circle
both `p` and `q` lie on.

````math
d_{𝕊}(p,q) = \arccos(\Re(⟨p,q⟩)).
````
"""
distance(::AbstractSphere, p, q) = acos(clamp(real(dot(p, q)), -1, 1))

@doc raw"""
    exp(M::AbstractSphere, p, X)

Compute the exponential map from `p` in the tangent direction `X` on the [`AbstractSphere`](@ref)
`M` by following the great arc eminating from `p` in direction `X`.

````math
\exp_p X = \cos(\lVert X \rVert_p)p + \sin(\lVert X \rVert_p)\frac{X}{\lVert X \rVert_p},
````
where $\lVert X \rVert_p$ is the [`norm`](@ref norm(::AbstractSphere,p,X)) on the
tangent space at `p` of the [`AbstractSphere`](@ref) `M`.
"""
exp(::AbstractSphere, ::Any...)

function exp!(M::AbstractSphere, q, p, X)
    θ = norm(M, p, X)
    q .= cos(θ) .* p .+ usinc(θ) .* X
    return q
end

flat!(M::AbstractSphere, ξ::CoTFVector, p, X::TFVector) = copyto!(ξ, X)

function get_basis(M::Sphere{n,ℝ}, p, B::DiagonalizingOrthonormalBasis{ℝ}) where {n}
    A = zeros(n + 1, n + 1)
    A[1, :] = transpose(p)
    A[2, :] = transpose(B.frame_direction)
    V = nullspace(A)
    κ = ones(n)
    if !iszero(B.frame_direction)
        # if we have a nonzero direction for the geodesic, add it and it gets curvature zero from the tensor
        V = cat(B.frame_direction / norm(M, p, B.frame_direction), V; dims = 2)
        κ[1] = 0 # no curvature along the geodesic direction, if x!=y
    end
    Ξ = [V[:, i] for i in 1:n]
    return CachedBasis(B, κ, Ξ)
end

@doc raw"""
    get_coordinates(M::AbstractSphere{ℝ}, p, X, B::DefaultOrthonormalBasis)

Represent the tangent vector `X` at point `p` from the [`AbstractSphere`](@ref) `M` in
an orthonormal basis by rotating the hyperplane containing `X` to a hyperplane whose
normal is the $x$-axis.

Given $q = p λ + x$, where $λ = \operatorname{sgn}(⟨x, p⟩)$, and $⟨⋅, ⋅⟩_{\mathrm{F}}$
denotes the Frobenius inner product, the formula for $Y$ is
````math
\begin{pmatrix}0 \\ Y\end{pmatrix} = X - q\frac{2 ⟨q, X⟩_{\mathrm{F}}}{⟨q, q⟩_{\mathrm{F}}}.
````
"""
get_coordinates(::AbstractSphere{ℝ}, p, X, ::DefaultOrthonormalBasis)

function get_coordinates!(M::AbstractSphere{ℝ}, Y, p, X, ::DefaultOrthonormalBasis{ℝ})
    n = manifold_dimension(M)
    p1 = p[1]
    cosθ = abs(p1)
    λ = nzsign(p1, cosθ)
    pend, Xend = view(p, 2:(n + 1)), view(X, 2:(n + 1))
    factor = λ * X[1] / (1 + cosθ)
    Y .= Xend .- pend .* factor
    return Y
end

@doc raw"""
    get_vector(M::AbstractSphere{ℝ}, p, X, B::DefaultOrthonormalBasis)

Convert a one-dimensional vector of coefficients `X` in the basis `B` of the tangent space
at `p` on the [`AbstractSphere`](@ref) `M` to a tangent vector `Y` at `p` by rotating the
hyperplane containing `X`, whose normal is the $x$-axis, to the hyperplane whose normal is
`p`.

Given $q = p λ + x$, where $λ = \operatorname{sgn}(⟨x, p⟩)$, and $⟨⋅, ⋅⟩_{\mathrm{F}}$
denotes the Frobenius inner product, the formula for $Y$ is
````math
Y = X - q\frac{2 \left\langle q, \begin{pmatrix}0 \\ X\end{pmatrix}\right\rangle_{\mathrm{F}}}{⟨q, q⟩_{\mathrm{F}}}.
````
"""
get_vector(::AbstractSphere{ℝ}, p, X, ::DefaultOrthonormalBasis)

function get_vector!(M::AbstractSphere{ℝ}, Y, p, X, ::DefaultOrthonormalBasis{ℝ})
    n = manifold_dimension(M)
    p1 = p[1]
    cosθ = abs(p1)
    λ = nzsign(p1, cosθ)
    pend = view(p, 2:(n + 1))
    pX = dot(pend, X)
    factor = pX / (1 + cosθ)
    Y[1] = -λ * pX
    Y[2:(n + 1)] .= X .- pend .* factor
    return Y
end

@doc raw"""
    injectivity_radius(M::AbstractSphere[, p])

Return the injectivity radius for the [`AbstractSphere`](@ref) `M`, which is globally $π$.

    injectivity_radius(M::Sphere, x, ::ProjectionRetraction)

Return the injectivity radius for the [`ProjectionRetraction`](@ref) on the
[`AbstractSphere`](@ref), which is globally $\frac{π}{2}$.
"""
injectivity_radius(::AbstractSphere) = π
injectivity_radius(::AbstractSphere, ::ExponentialRetraction) = π
injectivity_radius(::AbstractSphere, ::ProjectionRetraction) = π / 2
injectivity_radius(::AbstractSphere, ::Any) = π
injectivity_radius(::AbstractSphere, ::Any, ::ExponentialRetraction) = π
injectivity_radius(::AbstractSphere, ::Any, ::ProjectionRetraction) = π / 2
eval(
    quote
        @invoke_maker 1 Manifold injectivity_radius(
            M::AbstractSphere,
            rm::AbstractRetractionMethod,
        )
    end,
)

@doc raw"""
    inverse_retract(M::AbstractSphere, p, q, ::ProjectionInverseRetraction)

Compute the inverse of the projection based retraction on the [`AbstractSphere`](@ref) `M`,
i.e. rearranging $p+X = q\lVert p+X\rVert_2$ yields
since $\Re(⟨p,X⟩) = 0$ and when $d_{𝕊^2}(p,q) ≤ \frac{π}{2}$ that

````math
\operatorname{retr}_p^{-1}(q) = \frac{q}{\Re(⟨p, q⟩)} - p.
````
"""
inverse_retract(::AbstractSphere, ::Any, ::Any, ::ProjectionInverseRetraction)

function inverse_retract!(::AbstractSphere, X, p, q, ::ProjectionInverseRetraction)
    return (X .= q ./ real(dot(p, q)) .- p)
end

@doc raw"""
    log(M::AbstractSphere, p, q)

Compute the logarithmic map on the [`AbstractSphere`](@ref) `M`, i.e. the tangent vector,
whose geodesic starting from `p` reaches `q` after time 1.
The formula reads for $x ≠ -y$

````math
\log_p q = d_{𝕊}(p,q) \frac{q-\Re(⟨p,q⟩) p}{\lVert q-\Re(⟨p,q⟩) p \rVert_2},
````

and a deterministic choice from the set of tangent vectors is returned if $x=-y$, i.e. for
opposite points.
"""
log(::AbstractSphere, ::Any...)

function log!(M::AbstractSphere, X, p, q)
    cosθ = clamp(real(dot(p, q)), -1, 1)
    if cosθ ≈ -1 # appr. opposing points, return deterministic choice from set-valued log
        fill!(X, 0)
        if p[1] ≈ 1
            X[2] = 1
        else
            X[1] = 1
        end
        copyto!(X, X .- real(dot(p, X)) .* p)
        X .*= π / norm(X)
    else
        θ = acos(cosθ)
        X .= (q .- cosθ .* p) ./ usinc(θ)
    end
    return project!(M, X, p, X)
end

@doc raw"""
    manifold_dimension(M::AbstractSphere)

Return the dimension of the [`AbstractSphere`](@ref) `M`, respectively i.e. the
dimension of the embedding -1.
"""
manifold_dimension(M::AbstractSphere) = manifold_dimension(get_embedding(M)) - 1

"""
    mean(
        S::AbstractSphere,
        x::AbstractVector,
        [w::AbstractWeights,]
        method = GeodesicInterpolationWithinRadius(π/2);
        kwargs...,
    )

Compute the Riemannian [`mean`](@ref mean(M::Manifold, args...)) of `x` using
[`GeodesicInterpolationWithinRadius`](@ref).
"""
mean(::AbstractSphere, ::Any...)

function Statistics.mean!(
    S::AbstractSphere,
    p,
    x::AbstractVector,
    w::AbstractVector;
    kwargs...,
)
    return mean!(S, p, x, w, GeodesicInterpolationWithinRadius(π / 2); kwargs...)
end

function mid_point!(S::Sphere, q, p1, p2)
    q .= p1 .+ p2
    project!(S, q, q)
    return q
end

"""
    normal_tvector_distribution(S::Sphere{n,ℝ}, p, σ)

Generate a distribution in the tangent space at `p` by generating a
normal distribution in ambient space with standard deviation `σ`
projected to the tangent space at `p`.
"""
function normal_tvector_distribution(S::Sphere{n,ℝ}, p, σ) where {n}
    d = Distributions.MvNormal(zero(p), σ)
    return ProjectedFVectorDistribution(TangentBundleFibers(S), p, d, project!, p)
end

@doc raw"""
    project(M::AbstractSphere, p)

Project the point `p` from the embedding onto the [`Sphere`](@ref) `M`.

````math
\operatorname{proj}(p) = \frac{p}{\lVert p \rVert},
````
where $\lVert\cdot\rVert$ denotes the usual 2-norm for vectors if $m=1$ and the Frobenius
norm for the case $m>1$.
"""
project(::AbstractSphere, ::Any)

project!(::AbstractSphere, q, p) = (q .= p ./ norm(p))

@doc raw"""
    project(M::AbstractSphere, p, X)

Project the point `X` onto the tangent space at `p` on the [`Sphere`](@ref) `M`.

````math
\operatorname{proj}_{p}(X) = X - \Re(⟨p, X⟩)p
````
"""
project(::AbstractSphere, ::Any, ::Any)

project!(S::AbstractSphere, Y, p, X) = (Y .= X .- real(dot(p, X)) .* p)

@doc raw"""
    representation_size(M::AbstractSphere)

Return the size points on the [`AbstractSphere`](@ref) `M` are represented as, i.e., the
representation size of the embedding.
"""
@generated representation_size(M::ArraySphere{N}) where {N} = size_to_tuple(N)
@generated representation_size(M::Sphere{N}) where {N} = (N + 1,)

@doc raw"""
    retract(M::AbstractSphere, p, X, ::ProjectionRetraction)

Compute the retraction that is based on projection, i.e.

````math
\operatorname{retr}_p(X) = \frac{p+X}{\lVert p+X \rVert_2}
````
"""
retract(::AbstractSphere, ::Any, ::Any, ::ProjectionRetraction)

function retract!(M::AbstractSphere, q, p, X, ::ProjectionRetraction)
    q .= p .+ X
    return project!(M, q, q)
end

Base.show(io::IO, ::Sphere{n,𝔽}) where {n,𝔽} = print(io, "Sphere($(n), $(𝔽))")
function Base.show(io::IO, ::ArraySphere{N,𝔽}) where {N,𝔽}
    return print(io, "ArraySphere($(join(N.parameters, ", ")); field = $(𝔽))")
end

"""
    uniform_distribution(M::Sphere{n,ℝ}, p) where {n}

Uniform distribution on given [`Sphere`](@ref) `M`. Generated points will be of
similar type as `p`.
"""
function uniform_distribution(M::Sphere{n,ℝ}, p) where {n}
    d = Distributions.MvNormal(zero(p), 1.0)
    return ProjectedPointDistribution(M, d, project!, p)
end

@doc raw"""
    vector_transport_to(M::AbstractSphere, p, X, q, ::ParallelTransport)

Compute the parallel transport on the [`Sphere`](@ref) of the tangent vector `X` at `p`
to `q`, provided, the [`geodesic`](@ref) between `p` and `q` is unique. The formula reads

````math
P_{p←q}(X) = X - \frac{\Re(⟨\log_p q,X⟩_p)}{d^2_𝕊(p,q)}
\bigl(\log_p q + \log_q p \bigr).
````
"""
vector_transport_to(::AbstractSphere, ::Any, ::Any, ::Any, ::Any, ::ParallelTransport)

function vector_transport_to!(::AbstractSphere, Y, p, X, q, ::ParallelTransport)
    m = p .+ q
    mnorm2 = real(dot(m, m))
    factor = 2 * real(dot(X, q)) / mnorm2
    Y .= X .- m .* factor
    return Y
end