text:
module Inference

using ..Core
using ..Utilities
using DynamicPPL: Metadata, _tail, VarInfo, TypedVarInfo,
    islinked, invlink!, getlogp, tonamedtuple, VarName, getsym, vectorize,
    settrans!, _getvns, getdist, CACHERESET,
    Model, Sampler, SampleFromPrior, SampleFromUniform,
    Selector, DefaultContext, PriorContext,
    LikelihoodContext, MiniBatchContext, set_flag!, unset_flag!, NamedDist, NoDist,
    getspace, inspace
using Distributions, Libtask, Bijectors
using DistributionsAD: VectorOfMultivariate
using LinearAlgebra
using ..Turing: PROGRESS, Turing
using StatsFuns: logsumexp
using Random: AbstractRNG
using DynamicPPL
using AbstractMCMC: AbstractModel, AbstractSampler
using DocStringExtensions: TYPEDEF, TYPEDFIELDS

import AbstractMCMC
import AdvancedHMC; const AHMC = AdvancedHMC
import AdvancedMH; const AMH = AdvancedMH
import AdvancedPS
import BangBang
import ..Core: getchunksize, getADbackend
import DynamicPPL: get_matching_type,
    VarName, _getranges, _getindex, getval, _getvns
import EllipticalSliceSampling
import Random
import MCMCChains
import StatsBase: predict

export  InferenceAlgorithm,
        Hamiltonian,
        GibbsComponent,
        StaticHamiltonian,
        AdaptiveHamiltonian,
        SampleFromUniform,
        SampleFromPrior,
        MH,
        ESS,
        Emcee,
        Gibbs,      # classic sampling
        GibbsConditional,
        HMC,
        SGLD,
        SGHMC,
        HMCDA,
        NUTS,       # Hamiltonian-like sampling
        DynamicNUTS,
        IS,
        SMC,
        CSMC,
        PG,
        Prior,
        assume,
        dot_assume,
        observe,
        dot_observe,
        resume,
        predict,
        isgibbscomponent

#######################
# Sampler abstraction #
#######################
abstract type AbstractAdapter end
abstract type InferenceAlgorithm end
abstract type ParticleInference <: InferenceAlgorithm end
abstract type Hamiltonian{AD} <: InferenceAlgorithm end
abstract type StaticHamiltonian{AD} <: Hamiltonian{AD} end
abstract type AdaptiveHamiltonian{AD} <: Hamiltonian{AD} end

getchunksize(::Type{<:Hamiltonian{AD}}) where AD = getchunksize(AD)
getADbackend(::Hamiltonian{AD}) where AD = AD()

# Algorithm for sampling from the prior
struct Prior <: InferenceAlgorithm end

"""
    mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)

Decide if a proposal ``x'`` with log probability ``\\log p(x') = logp_proposal`` and
log proposal ratio ``\\log k(x', x) - \\log k(x, x') = log_proposal_ratio`` in a
Metropolis-Hastings algorithm with Markov kernel ``k(x_t, x_{t+1})`` and current state
``x`` with log probability ``\\log p(x) = logp_current`` is accepted by evaluating the
Metropolis-Hastings acceptance criterion
```math
\\log U \\leq \\log p(x') - \\log p(x) + \\log k(x', x) - \\log k(x, x')
```
for a uniform random number ``U \\in [0, 1)``.
"""
function mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)
    # replacing log(rand()) with -randexp() yields test errors
    return log(rand()) + logp_current ≤ logp_proposal + log_proposal_ratio
end

######################
# Default Transition #
######################

struct Transition{T, F<:AbstractFloat}
    θ  :: T
    lp :: F
end

function Transition(vi::AbstractVarInfo, nt::NamedTuple=NamedTuple())
    theta = merge(tonamedtuple(vi), nt)
    lp = getlogp(vi)
    return Transition{typeof(theta), typeof(lp)}(theta, lp)
end

metadata(t::Transition) = (lp = t.lp,)

DynamicPPL.getlogp(t::Transition) = t.lp

# Metadata of VarInfo object
metadata(vi::AbstractVarInfo) = (lp = getlogp(vi),)

#########################################
# Default definitions for the interface #
#########################################

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:InferenceAlgorithm},
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, sampler, N;
                                       chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::Prior,
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, SampleFromPrior(), N;
                                       chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, parallel, N, n_chains;
                               kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), parallel, N, n_chains;
                               kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:InferenceAlgorithm},
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.mcmcsample(rng, model, sampler, parallel, N, n_chains;
                                   chain_type=chain_type, progress=progress, kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::Prior,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.sample(rng, model, SampleFromPrior(), parallel, N, n_chains;
                               chain_type=chain_type, progress=progress, kwargs...)
end

##########################
# Chain making utilities #
##########################

"""
    getparams(t)

Return a named tuple of parameters.
"""
getparams(t) = t.θ
getparams(t::VarInfo) = tonamedtuple(TypedVarInfo(t))

function _params_to_array(ts::Vector)
    names = Vector{Symbol}()
    # Extract the parameter names and values from each transition.
    dicts = map(ts) do t
        nms, vs = flatten_namedtuple(getparams(t))
        for nm in nms
            if !(nm in names)
                push!(names, nm)
            end
        end
        # Convert the names and values to a single dictionary.
        return Dict(nms[j] => vs[j] for j in 1:length(vs))
    end
    # names = collect(names_set)
    vals = [get(dicts[i], key, missing) for i in eachindex(dicts),
        (j, key) in enumerate(names)]

    return names, vals
end

function flatten_namedtuple(nt::NamedTuple)
    names_vals = mapreduce(vcat, keys(nt)) do k
        v = nt[k]
        if length(v) == 1
            return [(Symbol(k), v)]
        else
            return mapreduce(vcat, zip(v[1], v[2])) do (vnval, vn)
                return collect(FlattenIterator(vn, vnval))
            end
        end
    end
    return [vn[1] for vn in names_vals], [vn[2] for vn in names_vals]
end

function get_transition_extras(ts::AbstractVector{<:VarInfo})
    valmat = reshape([getlogp(t) for t in ts], :, 1)
    return [:lp], valmat
end

function get_transition_extras(ts::AbstractVector)
    # Extract all metadata.
    extra_data = map(metadata, ts)
    return names_values(extra_data)
end

function names_values(extra_data::AbstractVector{<:NamedTuple{names}}) where names
    values = [getfield(data, name) for data in extra_data, name in names]
    return collect(names), values
end

function names_values(extra_data::AbstractVector{<:NamedTuple})
    # Obtain all parameter names.
    names_set = Set(Symbol[])
    for data in extra_data
        for name in names(data)
            push!(extra_names_set, name)
        end
    end
    extra_names = collect(extra_names_set)

    # Extract all values as matrix.
    values = [
        hasfield(data, name) ? missing : getfield(data, name)
        for data in extra_data, name in extra_names
    ]

    return extra_names, values
end

getlogevidence(transitions, sampler, state) = missing

# Default MCMCChains.Chains constructor.
# This is type piracy (at least for SampleFromPrior).
function AbstractMCMC.bundle_samples(
    ts::Vector,
    model::AbstractModel,
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
    state,
    chain_type::Type{MCMCChains.Chains};
    save_state = false,
    kwargs...
)
    # Convert transitions to array format.
    # Also retrieve the variable names.
    nms, vals = _params_to_array(ts)

    # Get the values of the extra parameters in each transition.
    extra_params, extra_values = get_transition_extras(ts)

    # Extract names & construct param array.
    nms = [nms; extra_params]
    parray = hcat(vals, extra_values)

    # Get the average or final log evidence, if it exists.
    le = getlogevidence(ts, spl, state)

    # Set up the info tuple.
    if save_state
        info = (model = model, sampler = spl, samplerstate = state)
    else
        info = NamedTuple()
    end

    # Conretize the array before giving it to MCMCChains.
    parray = MCMCChains.concretize(parray)

    # Chain construction.
    return MCMCChains.Chains(
        parray,
        nms,
        (internals = extra_params,);
        evidence=le,
        info=info,
    ) |> sort
end

# This is type piracy (for SampleFromPrior).
function AbstractMCMC.bundle_samples(
    ts::Vector,
    model::AbstractModel,
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
    state,
    chain_type::Type{Vector{NamedTuple}};
    kwargs...
)
    return map(ts) do t
        params = map(first, getparams(t))
        return merge(params, metadata(t))
    end
end

function save(c::MCMCChains.Chains, spl::Sampler, model, vi, samples)
    nt = NamedTuple{(:sampler, :model, :vi, :samples)}((spl, model, deepcopy(vi), samples))
    return setinfo(c, merge(nt, c.info))
end

function resume(chain::MCMCChains.Chains, args...; kwargs...)
    return resume(Random.GLOBAL_RNG, chain, args...; kwargs...)
end

function resume(rng::Random.AbstractRNG, chain::MCMCChains.Chains, args...;
                progress=PROGRESS[], kwargs...)
    isempty(chain.info) && error("[Turing] cannot resume from a chain without state info")

    # Sample a new chain.
    return AbstractMCMC.mcmcsample(
        rng,
        chain.info[:model],
        chain.info[:sampler],
        args...;
        resume_from = chain,
        chain_type = MCMCChains.Chains,
        progress = progress,
        kwargs...
    )
end

DynamicPPL.loadstate(chain::MCMCChains.Chains) = chain.info[:samplerstate]

#######################################
# Concrete algorithm implementations. #
#######################################

include("ess.jl")
include("hmc.jl")
include("mh.jl")
include("is.jl")
include("AdvancedSMC.jl")
include("gibbs.jl")
include("gibbs_conditional.jl")
include("../contrib/inference/sghmc.jl")
include("emcee.jl")

################
# Typing tools #
################

for alg in (:SMC, :PG, :MH, :IS, :ESS, :Gibbs, :Emcee)
    @eval DynamicPPL.getspace(::$alg{space}) where {space} = space
end
for alg in (:HMC, :HMCDA, :NUTS, :SGLD, :SGHMC)
    @eval DynamicPPL.getspace(::$alg{<:Any, space}) where {space} = space
end

function get_matching_type(
    spl::Sampler{<:Union{PG, SMC}},
    vi,
    ::Type{TV},
) where {T, N, TV <: Array{T, N}}
    return TArray{T, N}
end

##############
# Utilities  #
##############

DynamicPPL.getspace(spl::Sampler) = getspace(spl.alg)
DynamicPPL.inspace(vn::VarName, spl::Sampler) = inspace(vn, getspace(spl.alg))

"""

    predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)

Execute `model` conditioned on each sample in `chain`, and return the resulting `Chains`.

If `include_all` is `false`, the returned `Chains` will contain only those variables
sampled/not present in `chain`.

# Details
Internally calls `Turing.Inference.transitions_from_chain` to obtained the samples
and then converts these into a `Chains` object using `AbstractMCMC.bundle_samples`.

# Example
```jldoctest
julia> using Turing; Turing.turnprogress(false);
[ Info: [Turing]: progress logging is disabled globally

julia> @model function linear_reg(x, y, σ = 0.1)
           β ~ Normal(0, 1)

           for i ∈ eachindex(y)
               y[i] ~ Normal(β * x[i], σ)
           end
       end;

julia> σ = 0.1; f(x) = 2 * x + 0.1 * randn();

julia> Δ = 0.1; xs_train = 0:Δ:10; ys_train = f.(xs_train);

julia> xs_test = [10 + Δ, 10 + 2 * Δ]; ys_test = f.(xs_test);

julia> m_train = linear_reg(xs_train, ys_train, σ);

julia> chain_lin_reg = sample(m_train, NUTS(100, 0.65), 200);
┌ Info: Found initial step size
└   ϵ = 0.003125

julia> m_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)), σ);

julia> predictions = predict(m_test, chain_lin_reg)
Object of type Chains, with data of type 100×2×1 Array{Float64,3}

Iterations        = 1:100
Thinning interval = 1
Chains            = 1
Samples per chain = 100
parameters        = y[1], y[2]

2-element Array{ChainDataFrame,1}

Summary Statistics
  parameters     mean     std  naive_se     mcse       ess   r_hat
  ──────────  ───────  ──────  ────────  ───────  ────────  ──────
        y[1]  20.1974  0.1007    0.0101  missing  101.0711  0.9922
        y[2]  20.3867  0.1062    0.0106  missing  101.4889  0.9903

Quantiles
  parameters     2.5%    25.0%    50.0%    75.0%    97.5%
  ──────────  ───────  ───────  ───────  ───────  ───────
        y[1]  20.0342  20.1188  20.2135  20.2588  20.4188
        y[2]  20.1870  20.3178  20.3839  20.4466  20.5895

julia> ys_pred = vec(mean(Array(group(predictions, :y)); dims = 1));

julia> sum(abs2, ys_test - ys_pred) ≤ 0.1
true
```
"""
function predict(model::Model, chain::MCMCChains.Chains; kwargs...)
    return predict(Random.GLOBAL_RNG, model, chain; kwargs...)
end
function predict(rng::AbstractRNG, model::Model, chain::MCMCChains.Chains; include_all = false)
    spl = DynamicPPL.SampleFromPrior()

    # Sample transitions using `spl` conditioned on values in `chain`
    transitions = [
        transitions_from_chain(rng, model, chain[:, :, chn_idx]; sampler = spl)
        for chn_idx = 1:size(chain, 3)
    ]

    # Let the Turing internals handle everything else for you
    chain_result = reduce(
        MCMCChains.chainscat, [
            AbstractMCMC.bundle_samples(
                transitions[chn_idx],
                model,
                spl,
                nothing,
                MCMCChains.Chains
            ) for chn_idx = 1:size(chain, 3)
        ]
    )

    parameter_names = if include_all
        names(chain_result, :parameters)
    else
        filter(k -> ∉(k, names(chain, :parameters)), names(chain_result, :parameters))
    end

    return chain_result[parameter_names]
end

"""

    transitions_from_chain(
        [rng::AbstractRNG,]
        model::Model,
        chain::MCMCChains.Chains;
        sampler = DynamicPPL.SampleFromPrior()
    )

Execute `model` conditioned on each sample in `chain`, and return resulting transitions.

The returned transitions are represented in a `Vector{<:Turing.Inference.Transition}`.

# Details

In a bit more detail, the process is as follows:
1. For every `sample` in `chain`
   1. For every `variable` in `sample`
      1. Set `variable` in `model` to its value in `sample`
   2. Execute `model` with variables fixed as above, sampling variables NOT present
      in `chain` using `SampleFromPrior`
   3. Return sampled variables and log-joint

# Example
```julia-repl
julia> using Turing

julia> @model function demo()
           m ~ Normal(0, 1)
           x ~ Normal(m, 1)
       end;

julia> m = demo();

julia> chain = Chains(randn(2, 1, 1), ["m"]); # 2 samples of `m`

julia> transitions = Turing.Inference.transitions_from_chain(m, chain);

julia> [Turing.Inference.getlogp(t) for t in transitions] # extract the logjoints
2-element Array{Float64,1}:
 -3.6294991938628374
 -2.5697948166987845

julia> [first(t.θ.x) for t in transitions] # extract samples for `x`
2-element Array{Array{Float64,1},1}:
 [-2.0844148956440796]
 [-1.704630494695469]
```
"""
function transitions_from_chain(
    model::Turing.Model,
    chain::MCMCChains.Chains;
    kwargs...
)
    return transitions_from_chain(Random.GLOBAL_RNG, model, chain; kwargs...)
end
function transitions_from_chain(
    rng::AbstractRNG,
    model::Turing.Model,
    chain::MCMCChains.Chains;
    sampler = DynamicPPL.SampleFromPrior()
)
    vi = Turing.VarInfo(model)

    transitions = map(1:length(chain)) do i
        c = chain[i]
        md = vi.metadata
        for v in keys(md)
            for vn in md[v].vns
                vn_sym = Symbol(vn)

                # Cannot use `vn_sym` to index in the chain
                # so we have to extract the corresponding "linear"
                # indices and use those.
                # `ks` is empty if `vn_sym` not in `c`.
                ks = MCMCChains.namesingroup(c, vn_sym)

                if !isempty(ks)
                    # 1st dimension is of size 1 since `c`
                    # only contains a single sample, and the
                    # last dimension is of size 1 since
                    # we're assuming we're working with a single chain.
                    val = copy(vec(c[ks].value))
                    DynamicPPL.setval!(vi, val, vn)
                    DynamicPPL.settrans!(vi, false, vn)
                else
                    DynamicPPL.set_flag!(vi, vn, "del")
                end
            end
        end
        # Execute `model` on the parameters set in `vi` and sample those with `"del"` flag using `sampler`
        model(rng, vi, sampler)

        # Convert `VarInfo` into `NamedTuple` and save
        theta = DynamicPPL.tonamedtuple(vi)
        lp = Turing.getlogp(vi)
        Transition(theta, lp)
    end

    return transitions
end

end # module


text:
# Before running this, please make sure to activate and instantiate the environment
# corresponding to [this `Project.toml`](https://raw.githubusercontent.com/alan-turing-institute/MLJTutorials/master/Project.toml) and [this `Manifest.toml`](https://raw.githubusercontent.com/alan-turing-institute/MLJTutorials/master/Manifest.toml)
# so that you get an environment which matches the one used to generate the tutorials:
#
# ```julia
# cd("MLJTutorials") # cd to folder with the *.toml
# using Pkg; Pkg.activate("."); Pkg.instantiate()
# ```

# ## Getting started
using MLJ, RDatasets, PrettyPrinting

import Distributions
const D = Distributions

@load LinearRegressor pkg=MLJLinearModels
@load RidgeRegressor pkg=MLJLinearModels
@load LassoRegressor pkg=MLJLinearModels

hitters = dataset("ISLR", "Hitters")
@show size(hitters)
names(hitters) |> pprint

# The target is `Salary`
y, X = unpack(hitters, ==(:Salary), col->true);

# It has missing values which we will just ignore:
no_miss = .!ismissing.(y)
y = collect(skipmissing(y))
X = X[no_miss, :]
train, test = partition(eachindex(y), 0.5, shuffle=true, rng=424);

using PyPlot

figure(figsize=(8,6))
plot(y, ls="none", marker="o")

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Index", fontsize=14), ylabel("Salary", fontsize=14)

# \figalt{Salary}{ISL-lab-6-g1.svg}
# That looks quite skewed, let's have a look at a histogram:
figure(figsize=(8,6))
hist(y, bins=50, density=true)

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Salary", fontsize=14); ylabel("Density", fontsize=14)

edfit = D.fit_mle(D.Exponential, y)
xx = range(minimum(y), 2500, length=100)
yy = pdf.(edfit, xx)
plot(xx, yy, lw=3, label="Exponential distribution fit")

legend(fontsize=12)

# \figalt{Distribution of salary}{ISL-lab-6-g2.svg}## ### Data preparation## Most features are currently encoded as integers but we will consider them as continuous
Xc = coerce(X, autotype(X, rules=(:discrete_to_continuous,)))
scitype(Xc)

# There're a few features that are categorical which we'll one-hot-encode.
# ## Ridge pipeline# ### Baseline## Let's first fit a simple pipeline with a standardizer, a one-hot-encoder and a basic linear regression:
@pipeline RegPipe(std = Standardizer(),
                  hot = OneHotEncoder(),
                  reg = LinearRegressor())

model = RegPipe()
pipe  = machine(model, Xc, y)
fit!(pipe, rows=train)
ŷ = predict(pipe, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# Let's get a feel for how we're doing
figure(figsize=(8,6))

res = ŷ .- y[test]
stem(res)

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Index", fontsize=14); ylabel("Residual (ŷ - y)", fontsize=14)

ylim([-1300, 1000])

# \figalt{Residuals}{ISL-lab-6-g3.svg}
figure(figsize=(8,6))
hist(res, bins=30, density=true, color="green")

xx = range(-1100, 1100, length=100)
ndfit = D.fit_mle(D.Normal, res)
lfit  = D.fit_mle(D.Laplace, res)

plot(xx, pdf.(ndfit, xx), lw=3, color="orange", label="Normal fit")
plot(xx, pdf.(lfit, xx), lw=3, color="magenta", label="Laplace fit")

legend(fontsize=12)

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Residual (ŷ - y)", fontsize=14); ylabel("Density", fontsize=14)
xlim([-1100, 1100])

# \figalt{Distribution of residuals}{ISL-lab-6-g4.svg}
# ### Basic Ridge## Let's now swap the linear regressor for a Ridge one without specifying the penalty (`1` by default):
pipe.model.reg = RidgeRegressor()
fit!(pipe, rows=train)
ŷ = predict(pipe, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# Ok that's a bit better but surely we can do better with an appropriate selection of the hyperparameter.
# ### Cross validating
# What penalty should you use? Let's do a simple CV to try  to find out:
r  = range(model, :(reg.lambda), lower=1e-2, upper=100_000, scale=:log10)
tm = TunedModel(model=model, ranges=r, tuning=Grid(resolution=50),
                resampling=CV(nfolds=3, rng=4141), measure=rms)
mtm = machine(tm, Xc, y)
fit!(mtm, rows=train)

best_mdl = fitted_params(mtm).best_model
round(best_mdl.reg.lambda, sigdigits=4)

# right, and  with that we get:
ŷ = predict(mtm, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# Let's see:
figure(figsize=(8,6))

res = ŷ .- y[test]
stem(res)

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Index", fontsize=14); ylabel("Residual (ŷ - y)", fontsize=14)

ylim([-1300, 1000])

# \figalt{Ridge residuals}{ISL-lab-6-g5.svg}## You can compare that with the residuals obtained earlier.
# ## Lasso pipeline## Let's do the same as above but using a Lasso model and adjusting the range a bit:
mtm.model.model.reg = LassoRegressor()
mtm.model.range = range(model, :(reg.lambda), lower=500, upper=100_000, scale=:log10)
fit!(mtm, rows=train)

best_mdl = fitted_params(mtm).best_model
round(best_mdl.reg.lambda, sigdigits=4)

# Ok and let's see how that does:
ŷ = predict(mtm, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# Pretty good! and the parameters are reasonably sparse as expected:
coefs, intercept = fitted_params(mtm.fitresult.fitresult.machine)
round.(coefs, sigdigits=2)

# with around 50% sparsity:
sum(coefs .≈ 0) / length(coefs)

# Let's visualise this:
figure(figsize=(8,6))
stem(coefs)

# name of the features including one-hot-encoded ones
all_names = [:AtBat, :Hits, :HmRun, :Runs, :RBI, :Walks, :Years,
             :CAtBat, :CHits, :CHmRun, :CRuns, :CRBI, :CWalks,
             :League__A, :League__N, :Div_E, :Div_W,
             :PutOuts, :Assists, :Errors, :NewLeague_A, :NewLeague_N]

idxshow = collect(1:length(coefs))[abs.(coefs) .> 10]
xticks(idxshow .- 1, all_names[idxshow], rotation=45, fontsize=12)
yticks(fontsize=12)
ylabel("Amplitude", fontsize=14)

# \figalt{Lasso coefficients}{ISL-lab-6-g6.svg}
# ## Elastic net pipeline
@load ElasticNetRegressor pkg=MLJLinearModels

mtm.model.model.reg = ElasticNetRegressor()
mtm.model.range = [range(model, :(reg.lambda), lower=0.1, upper=100, scale=:log10),
                    range(model, :(reg.gamma),  lower=500, upper=10_000, scale=:log10)]
mtm.model.tuning = Grid(resolution=10)
fit!(mtm, rows=train)

best_mdl = fitted_params(mtm).best_model
@show round(best_mdl.reg.lambda, sigdigits=4)
@show round(best_mdl.reg.gamma, sigdigits=4)

# And it's not too bad in terms of accuracy either
ŷ = predict(mtm, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# But the simple ridge regression seems to work best here.

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl

ext: jl

max_stars_repo_name: ven-k/MLJTutorials

max_issues_repo_name: ven-k/MLJTutorials

max_forks_repo_name: ven-k/MLJTutorials


text:
module C

    const depfile = joinpath(@__DIR__, "..", "deps", "deps.jl")
    if isfile(depfile)
        include(depfile)
    else
        error("libspatialindex not properly installed. Please run Pkg.build(\"LibSpatialIndex\")")
    end

    include("include/sidx_config.jl")
    include("include/sidx_api.jl")

end


text:
# Matern 1/2 isotropic covariance Function

"""
    Mat12Iso <: MaternISO

Isotropic Matern 1/2 kernel (covariance)
```math
k(x,x') = σ^2 \\exp(-|x-y|/ℓ)
```
with length scale ``ℓ`` and signal standard deviation ``σ``.
"""
mutable struct Mat12Iso <: MaternIso
    "Length scale"
    ℓ::Float64
    "Signal variance"
    σ2::Float64
    "Priors for kernel parameters"
    priors::Array

    """
        Mat12Iso(ll::Float64, lσ::Float64)

    Create `Mat12Iso` with length scale `exp(ll)` and signal standard deviation `exp(σ)`.
    """
    Mat12Iso(ll::Float64, lσ::Float64) = new(exp(ll), exp(2 * lσ), [])
end

function set_params!(mat::Mat12Iso, hyp::VecF64)
    length(hyp) == 2 || throw(ArgumentError("Matern 1/2 covariance function only has two parameters"))
    mat.ℓ, mat.σ2 = exp(hyp[1]), exp(2 * hyp[2])
end

get_params(mat::Mat12Iso) = Float64[log(mat.ℓ), log(mat.σ2) / 2]
get_param_names(mat::Mat12Iso) = [:ll, :lσ]
num_params(mat::Mat12Iso) = 2

Statistics.cov(mat::Mat12Iso, r::Float64) = mat.σ2 * exp(-r / mat.ℓ)

@inline dk_dll(mat::Mat12Iso, r::Float64) = r / mat.ℓ * cov(mat, r)


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

struct VirtualMachineScaleSetVMsApi <: SwaggerApi
    client::Swagger.Client
end

"""

Run command on a virtual machine in a VM scale set.
Param: resourceGroupName::String (required)
Param: vmScaleSetName::String (required)
Param: instanceId::String (required)
Param: parameters::RunCommandInput (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: RunCommandResult
"""
function _swaggerinternal_virtualMachineScaleSetVMsRunCommand(_api::VirtualMachineScaleSetVMsApi, resourceGroupName::String, vmScaleSetName::String, instanceId::String, parameters, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "POST", RunCommandResult, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/virtualMachineScaleSets/{vmScaleSetName}/virtualmachines/{instanceId}/runCommand", ["azure_auth"], parameters)
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "vmScaleSetName", vmScaleSetName)  # type String
    Swagger.set_param(_ctx.path, "instanceId", instanceId)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    return _ctx
end

function virtualMachineScaleSetVMsRunCommand(_api::VirtualMachineScaleSetVMsApi, resourceGroupName::String, vmScaleSetName::String, instanceId::String, parameters, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = _swaggerinternal_virtualMachineScaleSetVMsRunCommand(_api, resourceGroupName, vmScaleSetName, instanceId, parameters, api_version, subscriptionId; _mediaType=_mediaType)
    Swagger.exec(_ctx)
end

function virtualMachineScaleSetVMsRunCommand(_api::VirtualMachineScaleSetVMsApi, response_stream::Channel, resourceGroupName::String, vmScaleSetName::String, instanceId::String, parameters, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = _swaggerinternal_virtualMachineScaleSetVMsRunCommand(_api, resourceGroupName, vmScaleSetName, instanceId, parameters, api_version, subscriptionId; _mediaType=_mediaType)
    Swagger.exec(_ctx, response_stream)
end

export virtualMachineScaleSetVMsRunCommand


text:
#=
ROS scavenging (detoxification) system
=#
using Parameters
include("common.jl")
include("concentrations.jl")

# SOD(superoxide dismutase) parameters
@with_kw struct SODParams
    K1 = 1200.0  # 2nd order rate constant of SOD (1/mMms)
    K3 = 24.0  # 2nd order rate constant of SOD (1/mMms)
    K5 = 2.4E-4  # 1st order rate constant of SOD (1/ms)
    KI_H2O2 = 0.5  # Inhibition constant of H2O2 (mM)
    ET = 0.0003  # SOD concentration (mM)
    VMAX = ET * K5
end

# Rate of superoxide dismutase (mM/ms),
function vsod(sox, h2o2, K1, K3, K5, KI_H2O2, VMAX)
	f_h2o2 = K3 * _mm_reci(KI_H2O2, h2o2)
	f_sox  = K1 * sox
	denom = K5 * (2 * K1 + f_h2o2) + f_h2o2 * f_sox
	vSOD = 2 * VMAX * f_sox * (K1 + f_h2o2) / denom
end

function vsod(sox, h2o2, pSOD::SODParams)
	@unpack K1, K3, K5, KI_H2O2, VMAX = pSOD
	vsod(sox, h2o2, K1, K3, K5, KI_H2O2, VMAX)
end

# GPX (Glutathione peroxidase) parameters
@with_kw struct GPXParams
    𝚽1 = 5E-5  # Rate constant of GPX
    𝚽2 = 27.0  # Rate constant of GPX
    ET = 9.77E-5  # GPX concentration (mM)
end

@with_kw struct TPXParams
	𝚽1 = 3.83  # Rate constant of TPX
    𝚽2 = 1.85  # Rate constant of TPX
    ET = 3E-3  # TPX concentration (mM)
end

#=
Reaction rate of glutathione peroxidase (GPX) and thioredoxin peroxidase (TPX)
H2O2 + 2GSH -> 2H2O + GSSH
H2O2 + TRXSH2 -> 2H2O + TRXSS
=#
vgpx(h2o2, gsh, ET, 𝚽1, 𝚽2) = ET * h2o2 * gsh / (𝚽1 * gsh + 𝚽2 * h2o2)
vgpx(h2o2, gsh, pGPX::GPXParams) = vgpx(h2o2, gsh, pGPX.ET, pGPX.𝚽1, pGPX.𝚽2)
vtpx(h2o2, trxsh2, ET, 𝚽1, 𝚽2) = vgpx(h2o2, trxsh2, ET, 𝚽1, 𝚽2)
vtpx(h2o2, trxsh2, pTPX::TPXParams) = vtpx(h2o2, trxsh2, pTPX.ET, pTPX.𝚽1, pTPX.𝚽2)

# GR (glutathion reductace) parameters
@with_kw struct GRParams
    K1 = 2.5E-3  # Catalytic constant of GR
    ET = 2.3E-3  # Mitochondiral concentration of GR (mM)
    VMAX = K1 * ET
    KM_NADPH = 6.54E-2  # Michaelis constant for NADPH of GR (mM)
    KM_GSSG = 20.6E-3  # Michaelis constant for oxidized GSH of GR (mM)
    GSH_T = 1.0  # Glutathione pool (mM)
end

# Rate of glutathione reductase (GR) (mM/ms)
# GSSG + NADPH + H+ -> 2GSH + NADP+
vgr(gssg, nadph, VMAX, KM_GSSG, KM_NADPH) = VMAX * _mm(nadph, KM_NADPH) * _mm(gssg, KM_GSSG)
vgr(gsh, nadph, VMAX, KM_GSSG, KM_NADPH, GSH_T) = vgr(0.5 * (GSH_T - gsh), nadph, VMAX, KM_GSSG, KM_NADPH)
function vgr(gsh, nadph, pGR::GRParams)
	@unpack VMAX, KM_GSSG, KM_NADPH, GSH_T = pGR
	vGR = vgr(gsh, nadph, VMAX, KM_GSSG, KM_NADPH, GSH_T)
end

# Thioredoxin reductase (TR) parameters
@with_kw struct TRParams
    K1 = 22.75E-3  # Catalytic constant of GR
    ET = 3.5E-4  # Mitochondiral concentration of GR (mM)
    VMAX = K1 * ET
    KM_NADPH = 0.065  # Michaelis constant for NADPH of GR (mM)
    KM_TRXSS = 0.035  # Michaelis constant for oxidized thioredoxin of TR (mM)
    TRX_T = 0.025  # Thioredoxin pool(mM)
end

# Thioredoxin reductase (TR) (mM/ms)
# TRXSS + NADPH -> TRXSH2 + NADP
vtr(trxsh2, nadph, VMAX, KM_TRXSS, KM_NADPH, TRX_T) = vgr(TRX_T - trxsh2, nadph, VMAX, KM_TRXSS, KM_NADPH)
function vtr(trxsh2, nadph, pTR::TRParams)
	@unpack VMAX, KM_TRXSS, KM_NADPH, TRX_T = pTR
	vTR = vtr(trxsh2, nadph, VMAX, KM_TRXSS, KM_NADPH, TRX_T)
end

#=
Rate of glutathione reductase (GR) (mM/ms)
GSSH + NADPH -> 2GSH + NADP
as well as for thioredoxin reductase (TR) (mM/ms)
TRXSS + NADPH -> TRXSH2 + NADP
=#

# Rate of H2O2 diffusion
vdiff_h2o2(h2o2_m, h2o2_i, C_DIFF) = _vdiff(h2o2_m, h2o2_i, C_DIFF)

# Catalase (CAT)
@with_kw struct CATParams
    K1 = 17.0  # Rate constant of CAT (1/mMms)
    ET = 1E-6  # Total pool of CAT (mM)
    VMAX = ET * K1
    FR = 0.05  # H2O2 inhibition factor of CAT
end

# Rate of intracellular catalase (CAT), scalar version
v_cat(h2o2_i, VMAX, FR) = 2 * VMAX * h2o2_i * exp(-FR * h2o2_i)
v_cat(h2o2_i, pCAT::CATParams) = v_cat(h2o2_i, pCAT.VMAX, pCAT.FR)

# IMAC (Inner mitochondrial anion channel) from Cortassa et al. (2004)
@with_kw struct IMACParams
	A = 1E-3  # Basal IMAC conductance
    B = 1E4  # Activation factor by cytoplasmic superoxide
    KCC_SOX = 0.01  # Michaelis constant for cytoplasmic superoxide of IMAC
    GL = 3.5E-8  # Leak conductance of IMAC
    G_MAX = 3.9085E-6  # Maximal conductance of IMAC
    κ = 0.07  # Steepness factor
    DPSI_OFFSET = 4.0  # Potential at half saturation
    J = 0.1  # Fraction of IMAC conductance
end

#=
Rate of mitochondrial inner membrane anion channel (IMAC) (mM/ms),
as well as superoxide diffusion from mitochondira to cytosol (VtrROS)
=#
# Conductance of IMAC (scalar version)
function _g_imac(sox_i, dpsi, A, B, KCC_SOX, GL, G_MAX, κ, DPSI_OFFSET)
	basal = A + B * _mm(sox_i, KCC_SOX)
	voltFactor = GL + G_MAX / (1 + exp(κ * (DPSI_OFFSET + dpsi)))
	gIMAC = voltFactor * basal
end

# scalar version of IMAC and superoxide pass-through (vTrROS)
function vimac_vtrros(dpsi, sox_m, sox_i, A, B, KCC_SOX, GL, G_MAX, κ, DPSI_OFFSET, J)
    gimac = _g_imac(sox_i, dpsi, A, B, KCC_SOX, GL, G_MAX, κ, DPSI_OFFSET)
	# Workaround negative ROS problem
	sox_i = max(1e-120, sox_i)
	sox_m = max(1e-120, sox_m)

    vTrROS = J * gimac * (dpsi + _esox(sox_i, sox_m))
    vIMAC = gimac * dpsi
    return vIMAC, vTrROS
end

function vimac_vtrros(dpsi, sox_m, sox_i, pIMAC::IMACParams)
	@unpack A, B, KCC_SOX, GL, G_MAX, κ, DPSI_OFFSET, J = pIMAC
	vimac_vtrros(dpsi, sox_m, sox_i, A, B, KCC_SOX, GL, G_MAX, κ, DPSI_OFFSET, J)
end

# nadph_m-producing isocitrate dehydrogenase (IDH2) parameters
@with_kw struct IDH2Params
    KM_H = 0.5  # Dissociation constant for H+ of IDH2
    KM_ISOC = 0.045  # Dissociation constant for isocitrate of IDH2
    KM_NADP = 0.046  # Michealis constant for NADP of IDH2
    KI_NADP = 2E-6  # Inhibition constant for NADP of IDH2
    KM_NADPH = 1.2E-2  # Michealis constant for nadph_m of IDH2
    KM_AKG = 0.08  # Michealis constant for AKG of IDH2
    VF = 8.72E-2  # Max forward rate of IDH2
    VB = 5.45E-3  # Max backward rate of IDH2
end

# Rate of isocitrate dehydrogenase II (producing nadph_m). Scalar parameters.
function vidh2(nadph_m, nadp_m, isoc, akg, h_m, KM_H, KM_ISOC, KM_NADP, KI_NADP, KM_AKG, KM_NADPH, VF, VB)
	f_h = _mm(KM_H, h_m)
    f_isoc = isoc / KM_ISOC
    f_nadp = nadp_m / (KM_NADP * _mm(nadp_m, KI_NADP))
    f_akg = akg / KM_AKG
    f_nadph = nadph_m / KM_NADPH
	denom = (1 + f_isoc + f_akg) * (1 + f_nadp + f_nadph)
	num = (VF * f_isoc * f_nadp - VB * f_akg * f_nadph)
	vIDH2 = num * f_h / denom
end

function vidh2(nadph_m, nadp_m, isoc, akg, h_m, pIDH2::IDH2Params)
	@unpack KM_H, KM_ISOC, KM_NADP, KI_NADP, KM_AKG, KM_NADPH, VF, VB = pIDH2
	vidh2(nadph_m, nadp_m, isoc, akg, h_m, KM_H, KM_ISOC, KM_NADP, KI_NADP, KM_AKG, KM_NADPH, VF, VB)
end

# transhydrogenase (THD) parameters
@with_kw struct THDParams
    KM_NADH = 0.01  # Michealis constant for NADH of THD
    KM_NAD = 0.125  # Michealis constant for NAD of THD
    KM_NADP = 0.02  # Michealis constant for NADP of THD
    KM_NADPH = 0.02  # Michealis constant for nadph_m of THD
    ET = 1E-3  # THD concnetration (mM)
    KF = 1.17474  # Max forward rate of THD
    VF = ET * KF
    KB = 17.2756  # Max backward rate of THD
    VB = ET * KB
    D = 0.5  # Voltage assymetry factor of THD
    X = 0.1  # Voltage dependence factor of THD
end

# Rate of transhydrogenase (producing nadph_m). Scalar version
function vthd(nadh, nad, nadph_m, nadp_m, ΔμH, X, D, KM_NAD, KM_NADH, KM_NADP, KM_NADPH, VF, VB)
	# Corrected an error from both the paper and code from Gauthier et al.
	vNAD = _ra(X * (1.0 - D) * ΔμH)
	fNAD = nad / KM_NAD
	fvNAD = fNAD * vNAD
	fNADH = nadh / KM_NADH
	fNADP = nadp_m / KM_NADP
	vNADP = _ra(X * D * ΔμH)
	fvNADP = vNADP * fNADP
	fNADPH = nadph_m / KM_NADPH
	denom = one(fNAD) + fNAD + fNADH + fNADP + fNADPH + (fNADPH + fvNAD) * (fNADPH + fvNADP)
	num = VF * fNADH * fvNADP - VB * fvNAD * fNADPH
	vTHD = num / denom
end

function vthd(nadh, nad, nadph_m, nadp_m, ΔμH, pTHD::THDParams)
	@unpack X, D, KM_NAD, KM_NADH, KM_NADP, KM_NADPH, VF, VB = pTHD
	vthd(nadh, nad, nadph_m, nadp_m, ΔμH, X, D, KM_NAD, KM_NADH, KM_NADP, KM_NADPH, VF, VB)
end

function ros_system(nadh, nadph_m, sox_m, sox_i, h2o2_m, h2o2_i, gsh_m, gsh_i,
	trxsh2_m, trxsh2_i, dpsi, ΔμH, isoc, akg, h_m, vROS, conc::Concentrations,
	pSODi::SODParams, pSODm::SODParams, pGPXm::GPXParams, pGPXi::GPXParams, pTPXm::TPXParams,
	pTPXi::TPXParams, pGRm::GRParams, pGRi::GRParams, pTRm::TRParams, pTRi::TRParams, pCAT::CATParams,
	pIMAC::IMACParams, pIDH2::IDH2Params, pTHD::THDParams, C_DIFF_H2O2, V_MITO_V_MYO)
	@unpack NAD_T, NADP_M_T, nadph_i = conc
	nad = NAD_T - nadh
	nadp_m = NADP_M_T - nadph_m
	vIDH2 = vidh2(nadph_m, nadp_m, isoc, akg, h_m, pIDH2)
	vTHD = vthd(nadh, nad, nadph_m, nadp_m, ΔμH, pTHD)
	vIMAC, vTrROS = vimac_vtrros(dpsi, sox_m, sox_i, pIMAC)
	vCAT = v_cat(h2o2_i, pCAT)
	vDiffH2O2 = vdiff_h2o2(h2o2_m, h2o2_i, C_DIFF_H2O2)
	vTR_i = vtr(trxsh2_i, nadph_i, pTRi)
	vTR_m = vtr(trxsh2_m, nadph_m, pTRm)
	vGR_i = vgr(gsh_i, nadph_i, pGRi)
	vGR_m = vgr(gsh_m, nadph_m, pGRm)
	vTPX_i = vtpx(h2o2_i, trxsh2_i, pTPXi)
	vTPX_m = vtpx(h2o2_m, trxsh2_m, pTPXm)
	vGPX_i = vgpx(h2o2_i, gsh_i, pGPXi)
	vGPX_m = vgpx(h2o2_m, gsh_m, pGPXm)
	vSOD_i = vsod(sox_i, h2o2_i, pSODi)
	vSOD_m = vsod(sox_m, h2o2_m, pSODm)

	d_nadph_m = vIDH2 + vTHD - vGR_m - vTR_m
	d_sox_m = vROS - vTrROS - vSOD_m
	d_sox_i = V_MITO_V_MYO * vTrROS - vSOD_i
	d_h2o2_m = 0.5 * vSOD_m - vDiffH2O2 - vTPX_m - vGPX_m
	d_h2o2_i = 0.5 * vSOD_i + V_MITO_V_MYO * vDiffH2O2 - vTPX_i - vGPX_i - vCAT
	d_gsh_m  = vGR_m - vGPX_m
	d_gsh_i  = vGR_i - vGPX_i
	d_trxsh2_m = vTR_m - vTPX_m
	d_trxsh2_i = vTR_i - vTPX_i
	return ( d_nadph_m, d_sox_m, d_sox_i, d_h2o2_m, d_h2o2_i, d_gsh_m, d_gsh_i,
		d_trxsh2_m, d_trxsh2_i, vIDH2, vTHD, vIMAC)
end


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Julia 1.4.1
#     language: julia
#     name: julia-1.4
# ---

# # ヒストグラム平坦化

# # 準備

# ## Load Julia Packages

using Images
using TestImages
using Plots

# ## read testimage

c = testimage("c")

# ## Visualize gray scale image histogram

data=reshape(255 .* channelview(c),prod(size(c)))
histogram(data, bins=0:255,normalize=:pdf)

# ## Visualize cumulative distribution function (so called CDF)

# +
i2c = Dict{Int,Int}()
for i in 0:255
    i2c[i] = 0
end

for i in Int.(255 .* channelview(c))
    i2c[i] += 1
end
# -

arr = [i2c[i] for i in 0:255]
cdf = [sum(arr[begin:i]) for i in 1:length(arr)]
plot(cdf./length(c), label="CDF")

# # Histogram Equalization

# - $q$ を画像の輝度の分布を表す確率密度関数とする. ヒストグラムを平坦化することは分布が（なるべく)一様分布に近似されれば良い.つまり累積分布関数が $y=ax$ のように原点と積分範囲の右端(画像であれば255=256-1)で $y=1$ となる一次関数になっていれば良い.
# - 今述べたことを定式化する. $X$ を画像の分布を表す確率変数だとする. $Y$ をヒストグラム平坦化により得られら分布とし滑らかに $Y=f(X)$ と変化したとする. さらに変換が可逆だと仮定する. $p$ を $Y$ の確率密度函数とすれば, 次のことを要請する.
#
# $$
# \frac{y}{255} = \int_{0}^{y} p(\tilde{y})d\tilde{y}
# $$

# 一般に $Y=f(X)$ と変換された確率変数の確率密度関数 $p$ は
#
# $$
# p(y) = \int \delta(y-f(x))q(x)dx
# $$
#
# である. この一般論を上式の右辺に適用すれば $p(y) = \frac{q(x)}{f'(x)}\ \mathrm{where}\ y=f(x)$ が従う.
#
# $$
# \begin{aligned}
# \frac{1}{255} &= \frac{d}{dy}\left(\frac{y}{255}\right) \\
#               &= \frac{d}{dy}\int_{0}^{y} p(\tilde{y})d\tilde{y} \\
#               &= p(y) \\
#               &= \frac{q(x)}{f'(x)}\ \mathrm{where}\ y=f(x) \\
# \end{aligned}
# $$
#
# よって
#
# $$
# f'(x) = 255q(x)
# $$
#
# 要するに欲しい変換は累積分布函数に255(確率変数 $X$ がとる最大値) を乗じたものである.
#
# $$
# f(x) = 255 \int_{0}^x q(\tilde{x}) d\tilde{x}
# $$
#
# これを画像処理に適用することを考えると積分は離散和になる. $q(x)$ は輝度が$x\in [0, 255]$である確率になる. 平たくいうと $x$ での（正規化された）ヒストグラムの値である．これを踏まえて下記のように実装をしてみよう.

n_pixels = length(c)
function hist_eq(p)
    s = 0
    for i in 0:p
        s+=i2c[i]
    end
    floor(255*s/n_pixels)
end

new_c=hist_eq.(255 .* channelview(c))

colorview(Gray, new_c./255)

histogram(new_c[:], normalize=:pdf, bins=0:255)

# +
new_i2c = Dict{Int,Int}()
for i in 0:255
    new_i2c[i] = 0
end

for i in new_c
    new_i2c[i] += 1
end
# -

# # 結果の確認

arr = [new_i2c[i] for i in 0:255]
cdf = [sum(arr[begin:i]) for i in 1:length(arr)]
plot(cdf./length(c))
# 直線になっている様子がわかるはず


text:
# here we define and document the prototypes that are implemented

"""
`AbstractAtoms`

the abstract supertype for storing atomistic
configurations. A basic implementation might simply store a list of positions.
"""
abstract AbstractAtoms

"""
Abstract supertype of neighbourlists.

The standard contructor should be of the form
```
   NeighbourList(at::AbstractAtoms, rcut)
```
where `rcut` should be either a scalar or a vector of cut-offs.
"""
abstract AbstractNeighbourList

abstract AbstractConstraint

"""
`Dofs`: dof vector
"""
typealias Dofs Vector{Float64}

"""
`AbstractCalculator`: the abstract supertype of calculators. These
store model information, and are linked to the implementation of energy,
forces, and so forth.
"""
abstract AbstractCalculator

"""
`Preconditioner`: abstract base type for preconditioners

Preconditioners need to implement the following three functions:
* `update!(P, at::AbstractAtoms)`
* `A_mul_B!(out::Dof, P, x::Dof)`
* `A_ldiv_B!(out::Dof, P, f::Dof)`
"""
abstract Preconditioner <: AbstractMatrix{Float64}

"""
`macro protofun(fsig::Expr)`

generates "function prototypes" which throw an error message
with a little extra information in case a certain function has not been
implemented. This is in some way duplicating default Julia behaviour, but it
has the advantage that we can write documention for non-existing functions and
it also emphasizes that this specific function is part of the abstract atoms
interface.

## Usage

```julia
"Returns number of atoms"
@protofun length(::AbstractAtoms)
```

This will create a function `length(at::AbstractAtoms)`; if somebody
calls `length(at)` on a concrete sub-type of `AbstractAtoms` for which
`length` hasn't been implemented, it will throw an error.
"""
macro protofun(fsig::Expr)
    @assert fsig.head == :call
    fname = fsig.args[1]
    argnames = Any[]
    for idx in 2:length(fsig.args)
        arg = fsig.args[idx]
        if isa(arg, Expr) && arg.head == :kw
            arg = arg.args[1]
        end
        if isa(arg, Symbol)
            push!(argnames, arg)
        elseif isa(arg, Expr) && arg.head == :(::)
            if length(arg.args) != 2
                @gensym s
                insert!(arg.args, 1, s)
            end
            push!(argnames, arg.args[1])
        end
    end
    body = quote
        error(string("JuLIP: ", $fname,
                     ($([:(typeof($(esc(arg)))) for arg in argnames]...),),
                     " ) has no implementation."))
    end
    return Expr(:function, esc(fsig), body)
end

# function defined primarily on AbstractAtoms
export AbstractAtoms,
      positions, get_positions, set_positions!, unsafe_positions,
      get_cell, set_cell!, is_cubic, pbc, get_pbc, set_pbc!,
      set_data!, get_data,
      set_calculator!, calculator, get_calculator!,
      set_constraint!, constraint, get_constraint,
      neighbourlist, cutoff

# length is used for several things
import Base: length, A_ldiv_B!, A_mul_B!, cell

export AbstractCalculator,
      energy, potential_energy, forces, grad

export AbstractNeighbourList,
      sites, bonds

export AbstractConstraint, NullConstraint, dofs

export Preconditioner, preconditioner

# TODO: probably rename Preconditioner to AbstractPreconditioner and
#       AMGPrecon to Preconditioner
# TODO: iterators for angles, dihedrals

"Return number of atoms"
@protofun length(::AbstractAtoms)

"Return copy of positions of all atoms as a `3 x N` array."
@protofun positions(::AbstractAtoms)

"alias for `positions`"
get_positions = positions

"return a reference to positions"
@protofun unsafe_positions(::AbstractAtoms)

"Set positions of all atoms as a `3 x N` array."
@protofun set_positions!(::AbstractAtoms, ::JVecs)

set_positions!(at::AbstractAtoms, p::Matrix) = set_positions!(at, vecs(p))

"get computational cell"
@protofun cell(::AbstractAtoms)

"alias for `cell`"
get_cell = cell

"set computational cell"
@protofun set_cell!(::AbstractAtoms, ::Matrix)

"set periodic boundary conditions"
@protofun set_pbc!(::AbstractAtoms, ::NTuple{3,Bool})

"get array (or tuple) determining which directions are periodic"
@protofun pbc(::AbstractAtoms)

"alias for `pbc`"
get_pbc = pbc

"determines whether a cubic cell is used (i.e. cell is a diagonal matrix)"
is_cubic(a::AbstractAtoms) = isdiag(cell(a))

"""
`set_data!(at, name, value)`:
associate some data with `at`; to be stored in a Dict within `at`;
if `name` is of type `Symbol` or `String` then can also use
 `setindex!`
"""
@protofun set_data!(a::AbstractAtoms, name::Any, value::Any)
Base.setindex!(at::AbstractAtoms, value,
               name::Union{Symbol, AbstractString}) = set_data!(at, name, value)

"""
`get_data(a, name)`:
obtain some data stored with `set_data!`,
if `name` is of type `Symbol` or `String` then can also use
 `getindex`
"""
@protofun get_data(a::AbstractAtoms, name::Any)
Base.getindex(at::AbstractAtoms,
               name::Union{Symbol, AbstractString}) = get_data(at, name)

"delete an atom"
@protofun deleteat!(a::AbstractAtoms, n::Integer)

"alias for `deleteat!`"
delete_atom! = deleteat!

"return an attached calculator"
@protofun calculator(at::AbstractAtoms)

"attach a calculator"
@protofun set_calculator!(at::AbstractAtoms, calc::AbstractCalculator)

"return attached constraint"
@protofun constraint(at::AbstractAtoms)

"attach a constraint"
@protofun set_constraint!(at::AbstractAtoms, cons::AbstractConstraint)

"""
`neighbourlist(a::AbstractAtoms, rcut)`

construct a suitable neighbourlist for this atoms object. `rcut` should
be allowed to be either a scalar or a vector.
"""
@protofun neighbourlist(a::AbstractAtoms, rcut::Float64)

#######################################################################
#  NEIGHBOURLIST
#######################################################################

"""
`sites(::AbstractNeighbourList)` or `sites(::AbstractAtoms, rcut)`

Returns an iterator over atomic sites.
```julia
for (idx, neigs, r, R, S) in sites(nlist)
    # do something at this site
end
```
Here `idx` is the current center-atom, `neigs` a collection
indexing the neighbouring atoms, `r` a vector if distances and
`R` a vector of distance vectors. `S` stores information
about which copy of the cell the neighbours belong to

A quicker way, if `nlist` won't be reused is
```julia
for n, ... in sites(at, rcut)
```

It should be assumed that `neigs, r, R, S` are views and therefore
should not be modified!
"""
@protofun sites(::AbstractNeighbourList)

sites(at::AbstractAtoms, rcut::Float64) = sites(neighbourlist(at, rcut))

"""
`bonds` : iterator over (pair-) bonds

E.g., the a pair potential can be implemented as follows:
```julia
ϕ(r) = r^(-12) - 2.0 * r^(-6)
dϕ(r) = -12 * (r^(-13) - r^(-7))
function lj(at::AbstractAtoms)
   E = 0.0; dE = zeros(3, length(at)) |> vecs
   for (i, j, r, R, S) in bonds(at, 4.1)
      E += ϕ(r)
      dE[j] += (dϕ(r)/r) * R
      dE[i] -= (dϕ(r)/r) * R
   end
   return E, dE
end
```
"""
@protofun bonds(::AbstractNeighbourList)

bonds(at::AbstractAtoms, cutoff::Float64) = bonds(neighbourlist(at, cutoff))

#######################################################################
#     CALCULATOR
#######################################################################

type  NullCalculator <: AbstractCalculator end

"Returns the cut-off radius of the potential."
@protofun cutoff(::AbstractCalculator)

"""
`energy`: can be called in various ways
*  `energy(calc, at)`: base definition
* `energy(at) = energy(calculator(at), at)`: if a calculator is attached to `at`
* `energy(calc, at, const, dof) = energy(calc, dof2at!(at,const,dof))`
* `energy(calc, at, dof) = energy(calc, at, constraint(dof), dof)`
* `energy(at, dof) = energy`

Return the total potential energy of a configuration of atoms `a`, using the calculator
`c`.
"""
@protofun energy(c::AbstractCalculator, a::AbstractAtoms)
energy(at::AbstractAtoms) = energy(calculator(at), at)

"`potential_energy` : alias for `energy`"
potential_energy = energy

# """
# energy difference between two configurations; default is to just compute the
# two energies, but this allows implementation of numerically robust
# differences,w hich can be important for very large problems.
# """
# energy_difference(c::AbstractCalculator, a::AbstractAtoms, aref::AbstractAtoms) =
#    energy(c, a) - energy(c, aref)

"""
Returns the negative gradient of the total energy in the format.
"""
@protofun forces(c::AbstractCalculator, a::AbstractAtoms)
forces(at::AbstractAtoms) = forces(calculator(at), at)

"`grad(c,a) = - forces(c, a)`"
grad(c::AbstractCalculator, a::AbstractAtoms) = scale!(forces(c, a), -1.0)
grad(at::AbstractAtoms) = grad(calculator(at), at)

#######################################################################
#  Constraints
#######################################################################

type NullConstraint <: AbstractConstraint end

"""
* `dofs(cons::AbstractConstraint, vecs::JVecs)`

Take a direction in position space (e.g. a collection of forces)
and project it to a dof-vector
"""
# function dofs end
@protofun dofs(at::AbstractAtoms, cons::AbstractConstraint, v_or_p)
dofs(at::AbstractAtoms, cons::AbstractConstraint) = dofs(at, cons, positions(at))
dofs(at::AbstractAtoms) = dofs(at, constraint(at))

"""
`vecs(cons::AbstractConstraint, vecs::JVecs)`:

Take a dof-vector and reconstruct a direction in position space
"""
@protofun vecs(cons::AbstractConstraint, at::AbstractAtoms, dofs::Dofs)

@protofun positions(cons::AbstractConstraint, at::AbstractAtoms, dofs::Dofs)

set_positions!(cons::AbstractConstraint, at::AbstractAtoms, dofs::Dofs) =
      set_positions!(at, positions(cons, at, dofs))

"""
`project(cons::AbstractConstraint, at::AbstractAtoms)`

take an atomistic configuration (collection of positions) and project
onto the manifold defined by the constraint.
"""
@protofun project!(at::AbstractAtoms, cons::AbstractConstraint)

# a bunch of short-cuts
energy(at::AbstractAtoms, x::Dofs) = energy(set_positions!(at, x))
forces(at::AbstractAtoms, x::Dofs) = dofs(at, constraint(at), forces(set_positions!(at, x)))
grad(at::AbstractAtoms, x::Dofs) = dofs(at, constraint(at), grad(set_positions!(at, x)))
set_positions!(at::AbstractAtoms, dofs::Dofs) = set_positions!(constraint(at), at, dofs)
positions(at::AbstractAtoms, dofs::Dofs) = positions(constraint(at), at, dofs)

#######################################################################
#                     PRECONDITIONER
#######################################################################

"""
`update!(precond::Preconditioner, at::AbstractAtoms)`

Update the preconditioner with the new geometry information.
"""
@protofun update!(precond::Preconditioner, at::AbstractAtoms)
update!(precond::Preconditioner, at::AbstractAtoms, x::Dofs) =
            update!(precond, set_positions!(at, x))
# TODO: this is a bit of a problem with the nice abstract framework we
#       are constructing here; it can easily happen now that we
#       update positions multiple times

"""
Identity preconditioner, i.e., no preconditioner.
"""
type Identity <: Preconditioner
end

A_ldiv_B!(out::Dofs, P::Identity, x::Dofs) = copy!(out, x)
A_mul_B!(out::Dofs, P::Identity, f::Dofs) = copy!(out, f)
update!(P::Identity, at::AbstractAtoms) = P

"construct a preconditioner suitable for this atoms object"
preconditioner(at::AbstractAtoms) = preconditioner(at, calculator(at), constraint(at))
preconditioner(at::AbstractAtoms, calc::AbstractCalculator, con::AbstractConstraint) =
         Identity()


text:
using StochasticDiffEq, DiffEqDevTools, Test
using DiffEqProblemLibrary.SDEProblemLibrary: importsdeproblems; importsdeproblems()
import DiffEqProblemLibrary.SDEProblemLibrary: prob_sde_linear, prob_sde_2Dlinear
prob = prob_sde_linear

Random.seed!(100)
sol1 = solve(prob,SRI(),dt=1//2^(4))
Random.seed!(100)
sol2 = solve(prob,SRIW1(),dt=1//2^(4))

@test sol1[end] ≈ sol2[end]

prob = prob_sde_2Dlinear

Random.seed!(100)
sol1 = solve(prob,SRI(),dt=1//2^(4))
Random.seed!(100)
sol2 = solve(prob,SRIW1(),dt=1//2^(4))

@test sol1[end] ≈ sol2[end]

prob = prob_sde_linear

Random.seed!(100)
integrator1 = init(prob,SRI(),dt=1//2^(4))
step!(integrator1); step!(integrator1)

Random.seed!(100)
integrator2 = init(prob,SRIW1(),dt=1//2^(4))
step!(integrator2); step!(integrator2)

@test integrator1.EEst ≈ integrator2.EEst

prob = prob_sde_2Dlinear

Random.seed!(100)
integrator1 = init(prob,SRI(),dt=1//2^(4))
step!(integrator1); step!(integrator1)

Random.seed!(100)
integrator2 = init(prob,SRIW1(),dt=1//2^(4))
step!(integrator2); step!(integrator2)

@test integrator1.EEst ≈ integrator2.EEst


text:
using GenericInstruments
GI = GenericInstruments
####################################
# Instantiate obj
resmgr   = GI.viOpenDefaultRM()
psu1 = GI.PSU.INSTR(:KeysightE3645A, "GPIB0::5::INSTR")#single ch psu
psu2 = GI.PSU.INSTR(:AgilentE3646A, "GPIB0::6::INSTR")#dual ch psu
####################################
# USAGE
####################################
# Power supply
GI.PSU.connect!(resmgr,psu1)
GI.PSU.connect!(resmgr,psu2)
# IDN
GI.PSU.get_idn(psu1)
GI.PSU.get_idn(psu2)
# Reset
GI.PSU.reset_instr(psu1)
GI.PSU.reset_instr(psu2)
# Set range
GI.PSU.set_range(psu1,ch="1",vrang="low")
GI.PSU.set_range(psu2,ch="1",vrang="high")
# Set voltage
GI.PSU.set_volt(psu1,ch="1",volt=5)
GI.PSU.set_volt(psu2,ch="1",volt=3)
# Set current compliance
GI.PSU.set_compl(psu1,ch="1",crtlim=0.6)
GI.PSU.set_compl(psu2,ch="1",crtlim=1.1)
# Output on/off
GI.PSU.set_outp(psu1,ch="1",st="on")
GI.PSU.set_outp(psu2,ch="1",st="on")
# Disconnect everything
GI.PSU.disconnect!(psu1)
GI.PSU.disconnect!(psu2)


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

using Markdown
using InteractiveUtils

# ╔═╡ 794a5356-6e05-44b4-ad9d-375a1fa3613b
begin
	using PlutoUI
	using DataFrames
	# using Statistics
	using CategoricalArrays
	using Pipe

	using Test
end

# ╔═╡ b3dcf3f4-9d68-11eb-01a6-317fe0e0cd7e
md"""
## Introduction DataFrames / 7 - working with CategoricalArrays.jl

ref. [Introduction to DataFrames, Bogumił Kamiński](https://github.com/bkamins/Julia-DataFrames-Tutorial/blob/master/07_factors.ipynb)

$(html"
© Pascal, April 2021
")
"""

# ╔═╡ d6993a5b-8158-4ce3-accb-8e8c514229b3
PlutoUI.TableOfContents(indent=true, depth=4, aside=true)

# ╔═╡ de97ded0-3952-4188-9a43-a8f383e1ae65
md"""
### Constructor
"""

# ╔═╡ 42476ba8-478d-49e2-94b5-80df7877a20f
x = categorical(["A", "B", "B", "C"])               ## unordered

# ╔═╡ bfffc9aa-9d20-4011-accf-ae379281209c
y = categorical(["A", "B", "B", "C"], ordered=true) ## ordered, by default order is sorting order

# ╔═╡ f61250b5-f3d9-496b-85be-c107685c9ba3
z = categorical(["A","B","B","C", missing])         ## unordered with missings

# ╔═╡ 7694c9d7-655a-4335-8478-8f7095c1805d
c = cut(1:10, 5) ## ordered, into equal counts, possible to rename labels and give custom breaks

# ╔═╡ 31d9f385-5926-4f9a-b085-c16a766063bb
@pipe DataFrame(x=cut(randn(100000), 10)) |>
      groupby(_, :x) |>
      combine(_, nrow)            ## just to make sure cut works right

# ╔═╡ 7109068b-5913-4622-8a74-6c54a9f05d4e
v = categorical([1, 2, 2, 3, 3])  ## contains integers not strings

# ╔═╡ d3ebccf9-b96d-4902-8fcd-bdf74e9e3169
Vector{Union{String, Missing}}(z) ## sometimes you need to convert back to a standard vector

# ╔═╡ ca9927e5-3ff2-4c73-8206-4d7495668312
md"""
### Managing levels
"""

# ╔═╡ 5ace73ce-ca39-4565-90d2-ce6637408f8c
ary = [x, y, z, c, v]

# ╔═╡ 46fe7e92-91d2-44b0-978e-25b2304e01c5
isordered.(ary)                  ## check if categorical array is orderd

# ╔═╡ dcff20c4-e82f-4f9f-be0d-3b5d37bf766a
ordered!(x, true), isordered(x)  ## make x ordered

# ╔═╡ 29d813df-bac4-46f3-b859-78dda9819b6f
ordered!(x, false), isordered(x) ## and unordered again

# ╔═╡ 45d1eb30-d9c6-45ac-b45d-72074c3ae9b7
levels.(ary)                     ## list levels

# ╔═╡ 725a4fdc-c46d-4e4f-894c-36616cbdbaed
unique.(ary)                     ## missing will be included

# ╔═╡ 751906fe-503f-435a-a73e-e4755c22d20c
y[1] < y[2]                      ## can compare as y is ordered

# ╔═╡ 15a37486-be1d-4098-98a8-10482190b81f
@test_throws ArgumentError v[1] < v[2]
## not comparable, v is unordered although it contains integers

# ╔═╡ 7c71af4c-23fd-4a49-b4fe-8415e83dfb7c
y[2] < "A"      ## comparison against type underlying categorical value is possible

# ╔═╡ db7f8d20-7c20-4f5f-8cb3-496a9a9b24b7
@test_throws KeyError y[2] < "Z"
## but it is treated as a level, and thus only valid levels are allowed

# ╔═╡ 775d8043-aac0-410d-98f4-df1fbad111b6
levels!(y, ["C", "B", "A"]) ## you can reorder levels, mostly useful for ordered CategoricalArrays

# ╔═╡ db1c777b-8b9c-4b50-b0d6-1a9093d48c35
y[1] < y[2]    ## observe that the order is changed

# ╔═╡ b3bb8f9e-3b4f-4b9b-a5b2-8f09e7b766e9
y[1] < "B"     ## level ordering is respected also when comparing with an underlying type

# ╔═╡ 6d83f001-47f2-4744-9e5d-8c0750f77cf7
@test_throws ArgumentError levels!(z, ["A", "B"])
## you have to specify all levels that are present

# ╔═╡ c6ad1156-7d24-4de4-ad56-f208180fedb8
levels!(z, ["A", "B"], allowmissing=true)
## unless the underlying array allows for missings and force removal of levels

# ╔═╡ 491c112f-e266-43f8-ad6e-b54af3ff940a
z[1] = "B"; z   ## now z has only "B" entries

# ╔═╡ d451ce95-8272-42da-95de-33ee2d7a7b14
levels(z)       ## but it remembers the levels it had (the reason is mostly performance)

# ╔═╡ e2313fb2-2781-40ea-a1e8-67153dfa710e
begin
	droplevels!(z)  ## this way we can clean it up
	levels(z)
end

# ╔═╡ a4615e9b-d3d9-443a-9146-e411f9bbd749
md"""
### Data manipulation

"""

# ╔═╡ 9065a655-af23-4f3b-88cc-b21f245035a8
x, levels(x)

# ╔═╡ 7a809f76-36fb-4867-a87f-c237251fd6dd
# begin
# 	x[2] = "0"
# 	x, levels(x) ## new level added at the end (works only for unordered)
# end

# ╔═╡ bc223b35-929a-4912-a07b-f5cc9aa6a84f
v, levels(v)

# ╔═╡ 35d2408b-5252-4042-9f0f-edcc9d21b8a6
@test_throws MethodError v[1] + v[2]
## even though the underlying data is Int, we cannot operate on it

# ╔═╡ e88e3cb2-d48f-4543-a575-ef7c3b575102
Vector{Int}(v)   ## you have either to retrieve the data by conversion (may be expensive)

# ╔═╡ b1f576e9-2145-4467-9d39-3e626db25ed4
get(v[1]) + get(v[2])            ## or get a single value

# ╔═╡ 80327e1c-faf3-44ce-aa61-e91796558892
get.(v)                          ## this will work for arrays without missings

# ╔═╡ c0986ddb-feae-4ffc-a054-81f09cb7c6d9
@test_throws MethodError get.(z) ## but will fail on missing values

# ╔═╡ 50a8f8f7-5857-4aa0-9dc6-09378dea2360
passmissing(get).(z)             ## you have to wrap it in passmissing

# ╔═╡ 59edffc4-8b66-4962-896f-3e9749281a87
Vector{Union{String, Missing}}(z) ## or do the conversion

# ╔═╡ bac73d64-1974-4717-a218-b6c650b8e3f9
recode([1, 2, 3, 4, 5, missing], 1=>10)
## recode some values in an array; has also in place recode! equivalent

# ╔═╡ 542d11f5-3941-4de4-b96a-06142e84499b
recode([1, 2, 3, 4, 5, missing], "a", 1=>10, 2=>20)
## here we provided a default value for not mapped recordings

# ╔═╡ f61e7337-db34-4c22-9df9-1b6a1753a261
recode([1, 2, 3, 4, 5, missing], 1=>10, missing=>"missing")
## to recode Missing you have to do it explicitly

# ╔═╡ e08b1c3f-6fc9-468e-9e5a-b4d4012c811c
begin
	t = categorical([1:5; missing])
	t, levels(t)
end

# ╔═╡ dbfb0c05-a389-4bd2-9149-5d6724fc885e
begin
	recode!(t, [1,3]=>2)
	t, levels(t) # note that the levels are dropped after recode
end

# ╔═╡ 646169b5-df7e-4dd8-8c95-3476f4726058
begin
	t₁ = categorical([1, 2, 3], ordered=true)
	levels(recode(t₁, 2=>0, 1=>-1))
	## and if you introduce a new levels they are added at the end in the order of appearance
end

# ╔═╡ 577b2c8c-f059-4e9a-8911-8bab2d266319
begin
	t₂ = categorical([1, 2, 3, 4, 5], ordered=true) # when using default it becomes the last level
	levels(recode(t₂, 300, [1, 2]=>100, 3=>200))
end

# ╔═╡ a103be27-45af-4f86-b25b-143ee811c34d
md"""
### Comparisons
"""

# ╔═╡ 7716455a-6bc5-4fa5-8275-03420ffa697f
begin
	x₀ = categorical([1, 2, 3])
	xs = [x₀, categorical(x₀), categorical(x₀, ordered=true), categorical(x₀, ordered=true)]
	levels!(xs[2], [3, 2, 1])
	levels!(xs[4], [2, 3, 1])
	[a == b for a in xs, b in xs]  ## all are equal - comparison only by contents
end

# ╔═╡ e8a72937-65ac-486f-91f7-dd710b19e1d8
begin
	signature(x::CategoricalArray) = (x, levels(x), isordered(x))
	## this is actually the full 	signature of CategoricalArray

	## all are different, notice that x[1] and x[2] are unordered but have a different
	## order of levels
	[signature(a) == signature(b) for a in xs, b in xs]
end

# ╔═╡ 5455528c-b692-4a44-8004-4992b61163db
@test_throws ArgumentError x₀[1] < x₀[2]
## you cannot compare elements of unordered CategoricalArray

# ╔═╡ 53b2f0f6-70d4-4a26-9f66-9a4d69773b4a
isless(t[1], t[2])  ## but you can do it for an ordered one (use `isless`)

# ╔═╡ 486de780-108e-4291-b251-8869d676ba31
isless(x[1], x[2])  ## isless works within the same CategoricalArray even if it is not ordered

# ╔═╡ d95df360-7718-4e64-994f-eb67271d3867
begin
	y₀ = deepcopy(x₀)  ## but not across categorical arrays
	@test_throws ArgumentError isless(x₀[1], y₀[2])
end

# ╔═╡ 1a593150-62aa-4150-a892-ec7e780fba64
isless(get(x₀[1]), get(y₀[2]))
## you can use get to make a comparison of the contents of CategoricalArray

# ╔═╡ fb305059-cc46-4233-ad75-003a2bb3c649
x[1] == y[2]  ## equality tests works OK across CategoricalArrays

# ╔═╡ 64039ae5-c281-4bd4-9b69-92bf9238b77d
md"""
### Categorical columns in a DataFrame"
"""

# ╔═╡ a3f12968-59ea-4c03-b3dd-4dd22efc61fe
df = DataFrame(x = 1:3, y = 'a':'c', z = ["a","b","c"])

# ╔═╡ ba18cae0-5e7d-4634-915e-9dbc803e9f95
md"""
Convert all string columns to categorical in-place
"""

# ╔═╡ db301499-c92a-475b-88d9-539a5ae8c7b3
transform!(df, names(df, String) => categorical, renamecols=false)

# ╔═╡ 194a7be1-5eac-4473-a35f-59a8dc9088df
describe(df)

# ╔═╡ c54cdbf1-7787-4f0f-b9ff-d4389f4a7c5e
html"""

"""

# ╔═╡ Cell order:
# ╟─b3dcf3f4-9d68-11eb-01a6-317fe0e0cd7e
# ╠═794a5356-6e05-44b4-ad9d-375a1fa3613b
# ╟─d6993a5b-8158-4ce3-accb-8e8c514229b3
# ╟─de97ded0-3952-4188-9a43-a8f383e1ae65
# ╠═42476ba8-478d-49e2-94b5-80df7877a20f
# ╠═bfffc9aa-9d20-4011-accf-ae379281209c
# ╠═f61250b5-f3d9-496b-85be-c107685c9ba3
# ╠═7694c9d7-655a-4335-8478-8f7095c1805d
# ╠═31d9f385-5926-4f9a-b085-c16a766063bb
# ╠═7109068b-5913-4622-8a74-6c54a9f05d4e
# ╠═d3ebccf9-b96d-4902-8fcd-bdf74e9e3169
# ╟─ca9927e5-3ff2-4c73-8206-4d7495668312
# ╠═5ace73ce-ca39-4565-90d2-ce6637408f8c
# ╠═46fe7e92-91d2-44b0-978e-25b2304e01c5
# ╠═dcff20c4-e82f-4f9f-be0d-3b5d37bf766a
# ╠═29d813df-bac4-46f3-b859-78dda9819b6f
# ╠═45d1eb30-d9c6-45ac-b45d-72074c3ae9b7
# ╠═725a4fdc-c46d-4e4f-894c-36616cbdbaed
# ╠═751906fe-503f-435a-a73e-e4755c22d20c
# ╠═15a37486-be1d-4098-98a8-10482190b81f
# ╠═7c71af4c-23fd-4a49-b4fe-8415e83dfb7c
# ╠═db7f8d20-7c20-4f5f-8cb3-496a9a9b24b7
# ╠═775d8043-aac0-410d-98f4-df1fbad111b6
# ╠═db1c777b-8b9c-4b50-b0d6-1a9093d48c35
# ╠═b3bb8f9e-3b4f-4b9b-a5b2-8f09e7b766e9
# ╠═6d83f001-47f2-4744-9e5d-8c0750f77cf7
# ╠═c6ad1156-7d24-4de4-ad56-f208180fedb8
# ╠═491c112f-e266-43f8-ad6e-b54af3ff940a
# ╠═d451ce95-8272-42da-95de-33ee2d7a7b14
# ╠═e2313fb2-2781-40ea-a1e8-67153dfa710e
# ╟─a4615e9b-d3d9-443a-9146-e411f9bbd749
# ╠═9065a655-af23-4f3b-88cc-b21f245035a8
# ╠═7a809f76-36fb-4867-a87f-c237251fd6dd
# ╠═bc223b35-929a-4912-a07b-f5cc9aa6a84f
# ╠═35d2408b-5252-4042-9f0f-edcc9d21b8a6
# ╠═e88e3cb2-d48f-4543-a575-ef7c3b575102
# ╠═b1f576e9-2145-4467-9d39-3e626db25ed4
# ╠═80327e1c-faf3-44ce-aa61-e91796558892
# ╠═c0986ddb-feae-4ffc-a054-81f09cb7c6d9
# ╠═50a8f8f7-5857-4aa0-9dc6-09378dea2360
# ╠═59edffc4-8b66-4962-896f-3e9749281a87
# ╠═bac73d64-1974-4717-a218-b6c650b8e3f9
# ╠═542d11f5-3941-4de4-b96a-06142e84499b
# ╠═f61e7337-db34-4c22-9df9-1b6a1753a261
# ╠═e08b1c3f-6fc9-468e-9e5a-b4d4012c811c
# ╠═dbfb0c05-a389-4bd2-9149-5d6724fc885e
# ╠═646169b5-df7e-4dd8-8c95-3476f4726058
# ╠═577b2c8c-f059-4e9a-8911-8bab2d266319
# ╟─a103be27-45af-4f86-b25b-143ee811c34d
# ╠═7716455a-6bc5-4fa5-8275-03420ffa697f
# ╠═e8a72937-65ac-486f-91f7-dd710b19e1d8
# ╠═5455528c-b692-4a44-8004-4992b61163db
# ╠═53b2f0f6-70d4-4a26-9f66-9a4d69773b4a
# ╠═486de780-108e-4291-b251-8869d676ba31
# ╠═d95df360-7718-4e64-994f-eb67271d3867
# ╠═1a593150-62aa-4150-a892-ec7e780fba64
# ╠═fb305059-cc46-4233-ad75-003a2bb3c649
# ╟─64039ae5-c281-4bd4-9b69-92bf9238b77d
# ╠═a3f12968-59ea-4c03-b3dd-4dd22efc61fe
# ╟─ba18cae0-5e7d-4634-915e-9dbc803e9f95
# ╠═db301499-c92a-475b-88d9-539a5ae8c7b3
# ╠═194a7be1-5eac-4473-a35f-59a8dc9088df
# ╟─c54cdbf1-7787-4f0f-b9ff-d4389f4a7c5e

max_stars_repo_name: pascal-p/julia-notebooks

max_issues_repo_name: pascal-p/julia-notebooks

max_forks_repo_name: pascal-p/julia-notebooks


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

using Markdown
using InteractiveUtils

# ╔═╡ 7fbfbd62-5e18-11ec-36b7-9f00610a9206
begin
  using Ripserer
  using Plots
end

# ╔═╡ cffa6493-e4bd-4769-a0ce-8137fee1fbfe
function noisy_circle(n; r1=1, r2=1, noise=0.1)
    points = NTuple{2,Float64}[]
    for _ in 1:n
        θ = 2π * rand()
        point = (
            r1 * sin(θ) + noise * rand() - noise / 2,
            r2 * cos(θ) + noise * rand() - noise / 2,
        )
        push!(points, point)
    end
    return points
end

# ╔═╡ 153f893a-6c95-4794-b1f6-6014f64f86ea
anim = @animate for _ in 1:200
    points = noisy_circle(100; noise=0)
    result = ripserer(points)

    plt_pts = scatter(
        points;
        legend=false,
        aspect_ratio=1,
        xlim=(-2.2, 2.2),
        ylim=(-2.2, 2.2),
        title="Data",
    )
    plt_diag = plot(result; infinity=3)

    plot(plt_pts, plt_diag; size=(800, 400))
end

# ╔═╡ dcc743e1-5ef7-4df5-a1a6-0ec1bc17a95e

# ╔═╡ Cell order:
# ╠═7fbfbd62-5e18-11ec-36b7-9f00610a9206
# ╠═cffa6493-e4bd-4769-a0ce-8137fee1fbfe
# ╠═153f893a-6c95-4794-b1f6-6014f64f86ea
# ╠═dcc743e1-5ef7-4df5-a1a6-0ec1bc17a95e

max_stars_repo_name: BioMedAI/Pluto-Note

max_issues_repo_name: BioMedAI/Pluto-Note

max_forks_repo_name: BioMedAI/Pluto-Note


text:
"""
    simpsons_rule(x::Vector, y::Vector)

Numerical integration via Simpson's rule.

Lengths of vectors must be equal to each other, and be odd (to have an even number of intervals).
Values in `x` represent abcissa values, while `y` are the integrand counterparts.
"""
function simpsons_rule(x::Vector, y::Vector)
    n = length(y)-1
    n % 2 == 0 || error("`y` length (number of intervals) must be odd")
    length(x)-1 == n || error("`x` and `y` length must be equal")
    h = (x[end]-x[1])/n
    @inbounds @views s = sum(y[1:2:n] .+ 4y[2:2:n] .+ y[3:2:n+1])
    return h/3 * s
end

"""
    trapezoidal_rule(x::Vector, y::Vector)

Numerical integration via the Trapezoidal rule.

Lengths of vectors must be equal to each other.
Values in `x` represent equally-spaced abcissa values, while `y` are the integrand counterparts.
"""
function trapezoidal_rule(x::Vector, y::Vector)
    return (x[2] - x[1]) * ( sum(y) - (y[1] + y[end])/2 )
end

"""
    trapezoidal(fun::Function, n, flim...)

Applies the `trapezoidal_rule` using `n` points over the frequency intervals specified by `flim...`.
"""
function trapezoidal(fun::Function, n, flim...)
    ii = 0.0
    for i in 2:length(flim)
        xi = collect(range(flim[i-1], stop=flim[i], length=n))
        yi = fun.(xi)
        ii += trapezoidal_rule(xi, yi)
    end
    return ii
end

"""
    gauss_interval(integrand::Function, n, fmin, fmax)

Computes Gauss-Legendre integration using `n` nodes and weights over the intervals `[fmin,fmax]`.
"""
function gauss_interval(integrand::Function, n, fmin, fmax)
    xi, wi = gausslegendre(n)
    ifi = @. integrand( (fmax-fmin)/2 * xi + (fmin+fmax)/2 )
    return (fmax-fmin)/2 * dot( wi, ifi )
end

"""
    gauss_intervals(fun::Function, n, flim...)

Computes Gauss-Legendre integration using `n` nodes and weights of the intervals specified in `flim...`.
"""
function gauss_intervals(fun::Function, n, flim...)
    xi, wi = gausslegendre(n)
    ii = 0.0
    for i in 2:length(flim)
        ii += (flim[i]-flim[i-1])/2 * dot( wi, fun.( (flim[i]-flim[i-1])/2 * xi .+ (flim[i]+flim[i-1])/2 ) )
    end
    return ii
end


text:
"""
    hmcr(Mvir; <keyword arguments>)

Compute the expected halo concentration from halo mass using the relation of
Dutton & Maccio 2014 (equations 7, 10-13).

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
"""
function hmcr(Mvir; mdef = default_mdef, cosmo = default_cosmo, z = 0.0)
    mdef == lowercase(strip(mdef))
    if mdef == "200c"
        a = 0.520 + (0.905 - 0.520) * exp(-0.617 * z ^ 1.21)
        b = -0.101 + 0.026 * z
    elseif mdef == "vir"
        a = 0.537 + (1.025 - 0.537) * exp(-0.718 * z ^ 1.08)
        b = -0.097 + 0.024 * z
    else
        throw("only implemented for 200c and vir halo mass definitions")
    end
    # convert Mvir to h-scaled units
    return @. exp10(a + b * (log10(Mvir * cosmo.h) - 12.0))
end

"""
    hmcr_prior(Mvir, cvir; <keyword arguments>)

Compute the probability of drawing a halo with virial parameters (Mvir, cvir),
using the halo mass-concentration relation of Dutton & Maccio 2014.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `sigma_logc::Real = 0.16`: scatter in log_concentration at fixed Mvir (in dex)
"""
function hmcr_prior(Mvir, cvir;
                    mdef = default_mdef, cosmo = default_cosmo, z = 0.0,
                    sigma_logc = 0.16)
    logc_expected = log10(hmcr(Mvir; mdef = mdef, cosmo = cosmo, z = z))
    logc = log10(cvir)
    return log_gauss(logc, logc_expected, sigma_logc)
end

function hmcr_prior(halo::HaloModel;
                    mdef = default_mdef, cosmo = default_cosmo, z = 0.0,
                    sigma_logc = 0.16)
    Mvir = virial_mass(halo; mdef = mdef, z = z, cosmo = cosmo)
    logc_expected = log10(hmcr(Mvir; mdef = mdef, cosmo = cosmo, z =z))
    logc_expected = log10(cvir_expected)
    logc = log10(concentration(halo; mdef = mdef, z = z, cosmo = cosmo))
    return log_gauss(logc, logc_expected, sigma_logc)
end

"""
    shmr(Mvir; <keyword arguments>)

Compute the expected stellar mass from halo mass using the relation of
Rodriguez-Puebla et al. 2017.

See equations 25-33 for the functional form and 49 - 55 for the parameters.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
"""
function shmr(Mvir; z = 0.0, cosmo = default_cosmo, mdef = default_mdef)

    if mdef != "vir"
        @warn("converting $mdef to vir halo mass definition", maxlog=1)
        @warn("assuming an NFW profile with standard concentration", maxlog=1)
        cvir = hmcr(Mvir; z = z, cosmo = cosmo, mdef = mdef)
        halo = NFW_from_virial(Mvir, cvir; z = z, cosmo = cosmo, mdef = mdef)
        Mvir = virial_mass(halo; z = z, cosmo = cosmo, mdef = "vir")
    end

    P(x, y, z) = y * z - x * z / (1.0 + z)
    Q(z) = exp(-4.0 / (1.0 + z)^2)

    ϵ0  =  -1.758
    ϵ1  =  +0.110
    ϵ2  =  -0.061
    ϵ3  =  -0.023
    M00 = +11.548
    M01 =  -1.297
    M02 =  -0.026
    α0  =  +1.975
    α1  =  +0.714
    α2  =  +0.042
    δ0  =  +3.390
    δ1  =  -0.472
    δ2  =  -0.931
    γ0  =  +0.498
    γ1  =  -0.157

    logϵ = ϵ0 + P(ϵ1, ϵ2, z) * Q(z) + P(ϵ3, 0.0, z)
    logM0 = M00 + P(M01, M02, z) * Q(z)
    α = α0 + P(α1, α2, z) * Q(z)
    δ = δ0 + P(δ1, δ2, z) * Q(z)
    γ = γ0 + P(γ1, 0.0, z) * Q(z)

    x = log10(Mvir) - logM0
    g(x) = δ * log10(1.0 + exp(x))^γ / (1.0 + exp(10.0^-x)) - log10(10^(-α * x) + 1.0)
    logMstar = logϵ + logM0 + g(x) - g(0)

    return exp10(logMstar)
end

function shmr(halo::HaloModel; mdef = default_mdef, z = 0.0, cosmo = default_cosmo)
    Mvir = virial_mass(halo; mdef = "vir", z = z, cosmo = cosmo)
    return shmr(Mvir; z = z, cosmo = cosmo, mdef = "vir")
end

"""
    shmr_prior(Mvir, Mstar; <keyword arguments>)

Compute the probability of drawing a halo mass and stellar mass pair (Mvir, Mstar).

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `sigma_logMstar::Real = 0.15` scatter in log_Mstar (in dex)
"""
function shmr_prior(Mvir, Mstar;
                    mdef = default_mdef, cosmo = default_cosmo, z = 0.0,
                    sigma_logMstar = 0.15)
    logMstar_expected = log10(shmr(Mvir; mdef = mdef, z = z, cosmo = cosmo))
    logMstar = log10(Mstar)
    return log_gauss(logMstar, logMstar_expected, sigma_logMstar)
end

function shmr_prior(halo::HaloModel, Mstar;
                    mdef = default_mdef, cosmo = default_cosmo, z = 0.0,
                    sigma_logMstar = 0.15)
    logMstar_expected = log10(shmr(halo; z = z, cosmo = cosmo))
    logMstar = log10(Mstar)
    return log_gauss(logMstar, logMstar_expected, sigma_logMstar)
end

"""
    peak_height(Mvir; <keyword arguments>)

Computes the dimensionless peak height,

```math
\\nu = \\delta_\\mathrm{crit}(z) / \\sigma(M, z)
```

using the approximation of Dutton & Maccio 2014 for a Planck 2013 cosmology.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
"""
function peak_height(Mvir; cosmo = default_cosmo, mdef="200c", z = 0.0)
    if mdef != "200c"
        throw("only implemented for 200c definition")
    end
    m = @. log10(M / 1e12 * cosmo.h)
    ν0 = @. exp10(-0.11 + 0.146 * m + 0.0138 * m^2 + 0.00123 * m^3)
    return @. ν0 * (0.033 + 0.79 * (1.0 + z) + 0.176 * exp(-1.356 * z))
end

"""
    alpha_peak(Mvir; <keyword arguments>)

Compute the Einasto shape parameter from peak height using the relation of Gao+2008.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
"""
function alpha_peak(Mvir; z = 0.0, cosmo = default_cosmo, mdef = "200c")
    if mdef != "200c"
        throw("only implemented for 200c definition")
    end
    ν = peak_height(Mvir, z; cosmo = cosmo, mdef = mdef)
    return @. 0.0095 * ν^2 + 0.155
end

"""
    alpha_peak_prior(Mvir, alpha; <keyword arguments>)

Compute the probability of drawing a halo with the given virial mass and Einasto shape
parameter using the alpha-peak height relation of Gao et al. 2008.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `sigma0::Real = 0.16`: scatter (in dex) of Mvir at redshift z = 0
- `sigmaz::Real = 0.03`: slope with redshift in scatter (in dex) of Mvir
"""
function alpha_peak_prior(Mvir, alpha;
                          mdef = default_mdef, cosmo = default_cosmo, z = 0.0,
                          sigma0 = 0.16, sigmaz = 0.03)
    sigma_logalpha = sigma0 + sigmaz * z
    logalpha_expected = log10(alpha_peak(Mvir; z = z, cosmo = cosmo, mdef = mdef))
    logalpha = log10(alpha)
    return log_gauss(logalpha, logalpha_expected, sigma_logalpha)
end

"""
    abg_from_logshm(logshm)

αβγ scalings with log10(Mstar / Mhalo) from DiCintio+2014a.
"""
function abg_from_logshm(logshm)
    x = logshm
    if x < -4.1 || x > -1.3
        return 1, 3, 1
    end
    α = 2.94 - log10(exp10(-1.08 * (x + 2.33)) + exp10(2.29 * (x + 2.33)))
    β = 4.23 + 1.34 * x + 0.26 * x^2
    γ = -0.06 + log10(exp10(-0.68 * (x + 2.56)) + exp10(x + 2.56))
    return α, β, γ
end


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
@inline calcM(x) = 2^x

struct MyType{N,M} end

function MyType{N}() where N
    @assert N ≥ 0
    M = calcM(N)
    MyType{N,M}()
end

@code_warntype MyType{3}() # not inferred of course

# %%
MyType{3}()

# %%
function MyType{N}() where N
    #T = typeof(N)
    M = 2^N
    MyType{N, M}()
end

@show MyType{3}()

@code_warntype MyType{3}()

# %%
@code_typed MyType{3}()

# %%
using PyCall
sys = pyimport("sys")
sys.stdout = stdout
sys.stderr = stderr

# %%
for i in 1:10
    print(i, "\n"); flush(stdout)
    sleep(0.5)
end

# %%

max_stars_repo_path: 0013/MyType{M, N}.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_path: 0013/MyType{M, N}.jl

max_issues_repo_name: genkuroki/public

max_forks_repo_path: 0013/MyType{M, N}.jl

max_forks_repo_name: genkuroki/public


text:
const ConcreteScheduler = Union{WorkStealingScheduler,PrioritizedScheduler}
const ConcreteThunk = Thunk{ConcreteScheduler}

function create_workers(threadids::AbstractVector{<:Integer}, schedulerref::Ref)
    return map(threadids) do threadid
        t = @task try
            eventloop(schedulerref[])
        catch err
            @error(
                "Unexpected error from `eventloop($(schedulerref[]))`",
                exception = (err, catch_backtrace())
            )
            close(schedulerref[])
            rethrow()
        end
        @check 1 <= threadid <= Threads.nthreads()
        t.sticky = true
        ccall(:jl_set_task_tid, Cvoid, (Any, Cint), t, threadid - 1)
        return Worker(t)
    end
end

function Schedulers.open(@nospecialize(f), scheduler::AbstractScheduler)
    function main()
        try
            return f(scheduler)
        finally
            @record(:main_done)
            # TODO: wait until scheduler is empty? (i.e., the last worker that
            # tried to sleep close the scheduler; it still doesn't wait for the
            # "I/O")
            close(scheduler)
        end
    end
    @record(:open_register_main, scheduler)
    task = generic_task(main, scheduler)
    pushat!(scheduler, 1, taskof(task))
    run(scheduler)
    @check istaskdone(task)
    fetch(task)
end

schedulertypeof(::typeof(Schedulers.workstealing)) = WorkStealingScheduler
schedulertypeof(::typeof(Schedulers.prioritized)) = PrioritizedScheduler

(ctx::SchedulerContextFunction)(;
    workers::Union{Integer,AbstractVector{<:Integer}} = Threads.nthreads(),
    options...,
) = schedulertypeof(ctx)(normalize_threadids(workers); options...)

(ctx::SchedulerContextFunction)(@nospecialize(f); options...) =
    Schedulers.open(f, ctx(; options...))

function start(scheduler::AbstractScheduler)
    unfair = nothing
    for (threadid, worker) in zip(scheduler.threadids, scheduler.workers)
        task = taskof(worker)
        if threadid == Threads.threadid()
            unfair = (; threadid = threadid, task = task)
        else
            Base.schedule(task)
        end
    end
    if unfair !== nothing
        @check unfair.threadid == Threads.threadid()
        yield(unfair.task)
    end
    return scheduler
end

function Base.run(scheduler::AbstractScheduler)
    @record(:run_begin, scheduler)
    start(scheduler)
    waitall(Iterators.map(taskof, scheduler.workers);)
    @record(:run_end, scheduler)
    return
end

function normalize_threadids(workers::Integer)
    @argcheck 1 <= workers <= Threads.nthreads()
    return collect(1:Int(workers))
end

function normalize_threadids(workers::AbstractVector{<:Integer})
    @argcheck !isempty(workers)
    threadids = Int[]
    invalids = nothing
    for (i, n) in pairs(workers)
        if 1 <= n <= Threads.nthreads()
            push!(threadids, n)
        else
            invalids = @something(invalids, Pair{Int,Int}[])
            push!(invalids, i => n)
        end
    end
    if invalids === nothing
        sort!(threadids)
        dups = duplicated_items_in_sorted_iterable(threadids)
        if dups === nothing
            return threadids
        end
        error(str, "Duplicated worker IDs: ", join(dups, ", "))
    end
    str = sprint(invalids) do io, invalids
        print(io, "Invalid worker IDs (must be `0 < _ ≤ nthreads = ")
        print(io, Threads.nthreads())
        print(io, "`): ")
        join(io, ("$n ($i-th element)" for (i, n) in invalids), ", ")
    end
    error(str)
end

struct RescheduleRequest end
# struct WakeupRequest end

struct YieldRequest end
struct WaitRequest end
struct ReturnRequest end
struct ExceptionRequest end

function eventloop(scheduler::AbstractScheduler)
    @record(:eventloop_begin, scheduler)
    threadid = Threads.threadid()
    worker = scheduler.workers[scheduler.workerindices[Threads.threadid()]]
    @check current_task() === taskof(worker)
    @atomic worker.state = WorkerStates.ACTIVE
    unfairspins = max(0, scheduler.unfairspins)
    yieldspins = unfairspins + max(0, scheduler.yieldspins)
    nspins = 0
    while true  # task execution loop
        @record(:try_get_next_task, scheduler, nspins)
        task::Union{Task,Nothing} = get_next_task!(scheduler, nspins > 0)
        if task === nothing
            @record(:no_next_task, scheduler)
        else
            @record(:got_next_task, scheduler, task = Historic.taskid(task))
            @DBG schedulerof(task) === scheduler
        end

        if task === nothing
            if !isopen(scheduler)
                @atomic worker.state = WorkerStates.DONE
                break
            end
            nspins += 1
            if nspins < unfairspins
                pause()
                continue
            elseif nspins < yieldspins
                yield()
                continue
            else
                @DBG @check worker.state == WorkerStates.ACTIVE
                task = trysleep!(scheduler, worker)
                if task === nothing
                    @DBG @check threadid == Threads.threadid()
                    continue
                end
            end
        end
        nspins = 0

        # Try to lock the task
        _, ok = @atomicreplace(
            :acquire_release,
            :monotonic,
            thunkof(task).worker,
            nothing => worker
        )
        if !ok
            push!(scheduler.externalqueue, task)
            wakeall!(scheduler)
            continue
        end

        while true  # request-reply loop
            @DBG @check threadid == Threads.threadid()
            @record(:schedule, scheduler, task = Historic.taskid(task))
            unset_threadid(task)
            @DBG @check Threads.threadid(task) == 0
            request = yieldto(task, RescheduleRequest())
            @atomic :release thunkof(task).worker = nothing
            @record(:requested, scheduler, task = Historic.taskid(task), request)
            @DBG @check threadid == Threads.threadid()
            if request isa YieldRequest
                next = get_next_task!(scheduler, false)
                if next isa Task
                    push!(scheduler, task)
                    task = next
                else
                    # Nothing in the queue. Re-schedule the `task` that just yielded.
                end
            elseif request isa WaitRequest
                break
            elseif request isa Union{ReturnRequest,ExceptionRequest}
                if request isa ExceptionRequest
                    # Let the native scheduler complete the `task`.
                    # See also `Base.wait(task::GenericTask)`.
                    yield(task)
                end
                notifywaiters!(scheduler, task)
                break
            else
                msg = "thread=$threadid: Unexpected request from `yieldto`"
                @error msg request
                close(scheduler)
                errorwith(; scheduler, worker, request) do io, (; request)
                    print(io, msg)
                    println(io, ":")
                    show(io, MIME"text/plain"(), request)
                end
            end
        end  # request-reply loop
    end  # task execution loop
    @record(:eventloop_end, scheduler)
end

function (thunk::Thunk)()
    try
        copy_stack = current_task().sticky
        current_task().sticky = true # always sticky wrt the native scheduler
        @record(:thunk_begin)
        if copy_stack
            msg = "Stack-copying not supported"
            @error(msg)
            error(msg)
        end
        thunk.state = TaskStates.STARTED
        y = thunk.f()
        @record(:thunk_returning)
        thunk.result = y
        thunk.state = TaskStates.DONE
        yieldto(taskof(thunk.worker), ReturnRequest())
    catch err
        # @error(:thunk_error, exception = (err, catch_backtrace()))
        @record(:thunk_error, exception = (err, catch_backtrace()))
        thunk.result = err
        thunk.state = TaskStates.ERROR
        yieldto(taskof(thunk.worker), ExceptionRequest())
        rethrow()
    end
end

function Schedulers.spawn(@nospecialize(f); options...)
    scheduler = Schedulers.current_scheduler()
    task = generic_task(f, scheduler; options...)
    @record(:spawn, task = taskid(task))
    schedule!(scheduler, task)
    return task
end

function Schedulers.spawn(@nospecialize(f), scheduler::AbstractScheduler; options...)
    task = generic_task(f, scheduler; options...)
    @record(:spawn_sch, task = taskid(task), scheduler)
    schedule!(scheduler, task, scheduler !== Schedulers.current_scheduler())
    return task
end

function Schedulers.schedule(task::GenericTask)
    scheduler = thunkof(task).scheduler
    schedule!(scheduler, task, scheduler !== Schedulers.current_scheduler())
    return
end

function schedule!(scheduler, task::Union{Task,GenericTask}, external::AnyBool = FALSE)
    task = taskof(task)
    isopen(scheduler) || errorwith(; scheduler) do io, data
        print(io, "scheduler ", typeof(data.scheduler), " is closed")
    end
    if boolean(external)
        push!(scheduler.externalqueue, task)
    else
        @DBG @check schedulerof(task) === scheduler
        push!(scheduler, task)
    end
    wakeall!(scheduler)
    return
end

function reschedule(task::Union{Task,GenericTask}, current_scheduler)
    scheduler = schedulerof(task)
    schedule!(scheduler, taskof(task), scheduler !== current_scheduler)
end

function trysleep!(scheduler::AbstractScheduler, worker::Worker)
    @atomic worker.state = WorkerStates.PREPARE_WAITING
    task::Union{Task,Nothing} = get_next_task!(scheduler, true)
    task isa Task && @goto done

    state = @atomic :monotonic worker.state
    state == WorkerStates.PREPARE_WAITING || @goto done

    @record(:worker_wait_cas, scheduler)
    @yield_unsafe begin
        state, ok = @atomicreplace(
            worker.state,
            WorkerStates.PREPARE_WAITING => WorkerStates.WAITING
        )
        ok || @goto done

        @record(:worker_wait_begin, scheduler, {yield = false})
    end  # @yield_unsafe
    Base.wait()
    @record(:worker_wait_end, scheduler)

    @label done
    @atomic worker.state = WorkerStates.ACTIVE
    return task
end

function wakeall!(scheduler::AbstractScheduler)
    for worker in scheduler.workers
        state = @atomic worker.state
        @record(:wakeall_check, worker_state = state, task = taskid(worker))
        while state == WorkerStates.PREPARE_WAITING
            state, ok = @atomicreplace(
                worker.state,
                WorkerStates.PREPARE_WAITING => WorkerStates.NOTIFYING
            )
            if ok
                state = WorkerStates.NOTIFYING
                @record(:wakeall_set_notifying, task = taskid(worker))
                break
            end
        end
        while state == WorkerStates.WAITING
            state, ok =
                @atomicreplace(worker.state, WorkerStates.WAITING => WorkerStates.NOTIFYING)
            if ok
                Base.schedule(worker.task)
                @record(:wakeall_scheduled, task = taskid(worker))
                state = WorkerStates.NOTIFYING
                break
            end
        end
    end
end

Schedulers.yield() = yieldto(scheduler_task(), current_task())::RescheduleRequest

function Schedulers.wait()
    @yield_unsafe begin
        @record(:wait_begin, waiter_state = thunkof(current_task()).state, {yield = false})
    end
    yieldto(scheduler_task(), WaitRequest())::RescheduleRequest
    @record(:wait_end, waiter_state = thunkof(current_task()).state)
    @DBG @check Threads.threadid() == Threads.threadid(scheduler_task())
    return
end

function maybe_current_task()
    task::Task = current_task()
    isdefined(task, :code) || return nothing  # root task does not set it
    task.code isa ConcreteThunk && return GenericTask(task)
    return nothing
end

function Schedulers.current_task()
    @return_if_something maybe_current_task()
    errorwith(; task = current_task()) do io, (; task)
        print(
            io,
            "Schedulers.current_task() is called from a task not scheduled with \
            Schedulers.spawn: ",
            task,
        )
    end
end

function Schedulers.current_scheduler()
    task = @something(maybe_current_task(), return nothing)
    return thunkof(task).scheduler
end

function trywait(task::GenericTask)
    thunk::Thunk = thunkof(task)
    head = @atomic thunk.waiter
    head isa Closed && return thunk.state

    # Register the `current_task()` as a waiter of `task`
    waiter = Waiter()
    @DBG @check waiter.task === current_task()
    @DBG @check waiter.state == WaiterStates.INITIAL
    while true
        waiter.next = head
        (head, ok) = @atomicreplace(thunk.waiter, head => waiter)
        head isa Closed && return thunk.state
        ok && break
    end
    @record(:wait_register, node = objid(waiter))

    @record(:wait_cas)
    # If the CAS is successful, there should be no yield point until
    # `Schedulers.wait()` yields.
    @yield_unsafe begin
        (state, ok) =
            @atomicreplace(waiter.state, WaiterStates.INITIAL => WaiterStates.WAITING)
    end  # @yield_unsafe
    if ok
        Schedulers.wait()
        @check waiter.state == WaiterStates.NOTIFYING
    else
        @record(:wait_dontwait)
        @check state == WaiterStates.DONTWAIT
    end
    return thunk.state
end

function Base.wait(task::GenericTask)
    @record(:wait_task_begin, task = taskid(task))
    state = trywait(task)
    @record(:wait_task_end, task = taskid(task), task_state = thunkof(task).state, state)
    if state == TaskStates.DONE
        @record(:wait_done, task = taskid(task))
        return
    elseif state == TaskStates.ERROR
        @record(:wait_error, task = taskid(task))
        # throw(thunk.result)
        Base.wait(task.task)  # use `TaskFailedException` for better stack trace

        # Unreachable:
        threadid = Threads.threadid()
        waiter = current_task()
        msg = "thread=$threadid: Unreachable: `wait` did not throw"
        @error msg waiter task
        errorwith(; threadid, waiter, task) do io, (; waiter, task)
            print(io, msg)
            print(io, "\n  waiter = ", waiter)
            print(io, "\n  task = ", task)
        end
    else
        error("unreachable: invalid state returned from `trywait`: ", state)
    end
end

function notifywaiters!(scheduler::AbstractScheduler, task::Task)
    task_state = thunkof(task).state
    @record(:notify_begin, task = taskid(task), task_state)
    if !isfinished(task_state)
        errorwith(; scheduler, task, task_state) do io, (; task, task_state)
            print(io, "FATAL: `notifywaiters!` called on an unfinished task")
            print(io, "\n  task = ", task)
            print(io, "\n  stateof(task) = ", task_state)
        end
    end
    thunk = thunkof(task)
    waiter::Union{Waiter,Nothing} = @atomicswap thunk.waiter = Closed()
    while true
        waiter === nothing && break
        state = @atomic :monotonic waiter.state
        @record(
            :notify_task,
            task = taskid(task),
            waiter = taskid(waiter),
            node = objid(waiter)
        )
        while state == WaiterStates.INITIAL
            (state, ok) =
                @atomicreplace(waiter.state, WaiterStates.INITIAL => WaiterStates.DONTWAIT)
            if ok
                state = WaiterStates.DONTWAIT
                @record(:notify_dontwait, task = taskid(task), waiter = taskid(waiter))
                break
            end
        end
        while state == WaiterStates.WAITING
            (state, ok) =
                @atomicreplace(waiter.state, WaiterStates.WAITING => WaiterStates.NOTIFYING)
            if ok
                @record(:notify_waking, task = taskid(task), waiter = taskid(waiter))
                # TODO: wake the task with `WakeupRequest()`
                reschedule(taskof(waiter), scheduler)
                break
            end
        end
        waiter = waiter.next
    end
    @record(:notify_end, task = taskid(task))
end

function Base.fetch(task::GenericTask)
    Base.wait(task)
    return thunkof(task).result
end