# # Classification of penguin species
#
# ## Packages

using CairoMakie
using CalibrationErrors
using DataFrames
using Distances
using Distributions
using PalmerPenguins
using Query

using Random

using CairoMakie.AbstractPlotting.ColorSchemes: Dark2_8

# ## Data
#
# In this example we study the calibration of different models that classify three penguin
# species based on measurements of their bill and flipper lengths.
#
# We use the [Palmer penguins dataset](https://allisonhorst.github.io/palmerpenguins/) to
# to train and validate the models.

penguins = dropmissing(DataFrame(PalmerPenguins.load()))

f = Figure()
ax = Axis(f[1, 1]; xlabel="bill length [mm]", ylabel="flipper length [mm]")
for (i, (key, df)) in enumerate(pairs(groupby(penguins, :species)))
    scatter!(
        df.bill_length_mm, df.flipper_length_mm; color=(Dark2_8[i], 0.8), label=key.species
    )
end
Legend(f[1, 2], ax, "species")
#!jl save("./figures/penguins.svg", f);

#!jl # ![](./figures/penguins.svg)

# We split the data randomly into a training and validation dataset. The training dataset
# contains around 60% of the samples.

Random.seed!(1234)
idxs = shuffle(axes(penguins, 1))
k = floor(Int, 0.6 * length(idxs))
train_idxs = @view idxs[1:k]
val_idxs = @view idxs[(k + 1):end]

train_penguins = penguins[train_idxs, :]
val_penguins = penguins[val_idxs, :];

#-

f = Figure()
ax = Axis(f[1, 1]; xlabel="bill length [mm]", ylabel="flipper length [mm]")
for (i, df) in enumerate((train_penguins, val_penguins))
    for (j, (key, subdf)) in enumerate(pairs(groupby(df, :species)))
        scatter!(
            subdf.bill_length_mm,
            subdf.flipper_length_mm;
            color=(Dark2_8[j], 0.8),
            marker=i == 1 ? :circle : :diamond,
        )
    end
end
group_marker = [
    MarkerElement(; marker=m, color=:black, strokecolor=:transparent, markersize=20) for
    m in (:circle, :diamond)
]
group_color = [PolyElement(; color=Dark2_8[i], strokecolor=:transparent) for i in 1:3]
Legend(
    f[1, 2],
    [group_marker, group_color],
    [["training", "validation"], string.(levels(penguins.species))],
    ["dataset", "species"],
)
#!jl save("./figures/penguins_datasets.svg", f);

#!jl # ![](./figures/penguins_datasets.svg)

# ## Fitting normal distributions
#
# For each species, we fit independent normal distributions to the observations of the bill
# and flipper length in the training data, using maximum likelihood estimation.

penguins_fit = @from i in train_penguins begin
    @group i by i.species into g
    @select {
        species = key(g),
        proportion = length(g) / nrow(train_penguins),
        bill = fit(Normal, g.bill_length_mm),
        flipper = fit(Normal, g.flipper_length_mm),
    }
    @collect DataFrame
end

# We plot the estimated normal distributions.

function xrange(dists, alpha=0.0001)
    xmin = minimum(Base.Fix2(quantile, alpha), dists)
    xmax = maximum(Base.Fix2(quantile, 1 - alpha), dists)
    return range(xmin, xmax; length=1_000)
end

function plot_normal_fit(dists, species, xlabel)
    f = Figure()
    Axis(f[1, 1]; xlabel=xlabel, ylabel="density")
    xs = xrange(dists)
    plots = map(enumerate(dists)) do (i, dist)
        ys = pdf.(dist, xs)
        l = lines!(xs, ys; color=Dark2_8[i])
        b = band!(xs, 0, ys; color=(Dark2_8[i], 0.2))
        return [l, b]
    end
    Legend(f[1, 2], plots, species, "species")
    return f
end

plot_normal_fit(penguins_fit.bill, penguins_fit.species, "bill length [mm]")
#!jl save("./figures/normal_fit_bill.svg", current_figure());

#!jl # ![](./figures/normal_fit_bill.svg)

plot_normal_fit(penguins_fit.flipper, penguins_fit.species, "flipper length [mm]")
#!jl save("./figures/normal_fit_flipper.svg", current_figure());

#!jl # ![](./figures/normal_fit_flipper.svg)

# ## Naive Bayes classifier
#
# Let us assume that the bill and flipper length are conditionally independent given the
# penguin species. Then Bayes' theorem implies that
# ```math
# \begin{aligned}
# \mathbb{P}(\mathrm{species} \,|\, \mathrm{bill}, \mathrm{flipper})
# &= \frac{\mathbb{P}(\mathrm{species}) \mathbb{P}(\mathrm{bill}, \mathrm{flipper} \,|\, \mathrm{species})}{\mathbb{P}(\mathrm{bill}, \mathrm{flipper})} \\
# &= \frac{\mathbb{P}(\mathrm{species}) \mathbb{P}(\mathrm{bill} \,|\, \mathrm{species}) \mathbb{P}(\mathrm{flipper} \,|\, \mathrm{species})}{\mathbb{P}(\mathrm{bill}, \mathrm{flipper})}.
# \end{aligned}
# ```
# This predictive model is known as
# [naive Bayes classifier](https://en.wikipedia.org/wiki/Naive_Bayes_classifier).
#
# In the section above, we estimated $\mathbb{P}(\mathrm{species})$,
# $\mathbb{P}(\mathrm{bill} \,|\, \mathrm{species})$, and
# $\mathbb{P}(\mathrm{flipper} \,|\, \mathrm{species})$ for each penguin species from
# the training data. For the conditional distributions we used a Gaussian approximation.

function predict_naive_bayes_classifier(fit, data)
    ## Compute unnormalized probabilities
    z =
        log.(permutedims(fit.proportion)) .+
        logpdf.(permutedims(fit.bill), data.bill_length_mm) .+
        logpdf.(permutedims(fit.flipper), data.flipper_length_mm)

    ## Normalize probabilities
    u = maximum(z; dims=2)
    z .= exp.(z .- u)
    sum!(u, z)
    z ./= u

    return DataFrame(z, fit.species)
end

train_predict = predict_naive_bayes_classifier(penguins_fit, train_penguins)
val_predict = predict_naive_bayes_classifier(penguins_fit, val_penguins);

# ## Evaluation
#
# We evaluate the probabilistic predictions of the naive Bayes classifier that we just
# trained. It is easier to work with a numerical encoding of the true penguin species and a
# corresponding vector of predictions.

train_species = convert(Vector{Int}, indexin(train_penguins.species, names(train_predict)))
train_probs = collect(Vector{Float64}, eachrow(train_predict))

val_species = convert(Vector{Int}, indexin(val_penguins.species, names(val_predict)))
val_probs = collect(Vector{Float64}, eachrow(val_predict));

# ### Log-likelihood
#
# We compute the average log-likelihood of the training and validation data. It is
# equivalent to the negative cross-entropy.

function mean_loglikelihood(species, probs)
    return mean(log(p[s]) for (s, p) in zip(species, probs))
end

mean_loglikelihood(train_species, train_probs)

#-

mean_loglikelihood(val_species, val_probs)

# ### Brier score
#
# The average log-likelihood is also equivalent to the
# [logarithmic score](https://sites.stat.washington.edu/raftery/Research/PDF/Gneiting2007jasa.pdf).
# The Brier score is another strictly proper scoring rule that can be used for evaluating
# probabilistic predictions.

function brier_score(species, probs)
    return mean(
        sum(abs2(pi - (i == s)) for (i, pi) in enumerate(p)) for
        (s, p) in zip(species, probs)
    )
end

brier_score(train_species, train_probs)

#-

brier_score(val_species, val_probs)

# ### Expected calibration error
#
# As all proper scoring rules, the logarithmic and the Brier score can be [decomposed in
# three terms that quantify the sharpness and calibration of the predictive model and the
# irreducible uncertainty of the targets that is inherent to the prediction
# problem](https://doi.org/10.1002/qj.456). The calibration term in this decomposition is
# the expected calibration error (ECE)
# ```math
# \mathbb{E} d\big(P_X, \mathrm{law}(Y \,|\, P_X)\big)
# ```
# with respect to the score divergence $d$.
#
# Scoring rules, however, include also the sharpness and the uncertainty term. Thus models
# can trade off calibration for sharpness and therefore scoring rules are not suitable for
# specifically evaluating calibration of predictive models.
#
# The score divergence to the logarithmic and the Brier score are the Kullback-Leibler (KL)
# divergence
# ```math
# d\big(P_X, \mathrm{law}(Y \,|\, P_X)\big) = \sum_{y} \mathbb{P}(Y = y \,|\, P_X)
# \log\big(\mathbb{P}(Y = y \,|\, P_X) / P_X(\{y\})\big)
# ```
# and the squared Euclidean distance
# ```math
# d\big(P_X, \mathrm{law}(Y \,|\, P_X)\big) = \sum_{y} \big(P_X - \mathrm{law}(Y \,|\, P_X)\big)^2(\{y\}),
# ```
# respectively. The KL divergence is defined only if $\mathrm{law}(Y \,|\, P_X)$ is
# absolutely continuous with respect to $P_X$, i.e., if $P_X(\{y\}) = 0$ implies
# $\mathbb{P}(Y = y \,|\, P_X) = 0$.

# We estimate the ECE by binning the probability simplex of predictions $P_X$ and computing
# the weighted average of the distances between the mean prediction and the distribution of
# targets in each bin.
#
# One approach is to use bins of uniform size.

ece = ECE(UniformBinning(10), (μ, y) -> kl_divergence(y, μ))
ece(train_probs, train_species)

#-

ece(val_probs, val_species)

# For the squared Euclidean distance we obtain:

ece = ECE(UniformBinning(10), SqEuclidean())
ece(train_probs, train_species)

#-

ece(val_probs, val_species)

# Alternatively, one can use a data-dependent binning scheme that tries to split the
# predictions in a way that minimizes the variance in each bin.
#
# With the KL divergence we get:

ece = ECE(MedianVarianceBinning(5), (μ, y) -> kl_divergence(y, μ))
ece(train_probs, train_species)

#-

ece(val_probs, val_species)

# For the squared Euclidean distance we obtain:

ece = ECE(MedianVarianceBinning(5), SqEuclidean())
ece(train_probs, train_species)

#-

ece(val_probs, val_species)

# We see that the estimates (of the same theoretical quantity!) are highly dependent on the
# chosen binning scheme.

# ### Kernel calibration error
#
# As an alternative to the ECE, we estimate the kernel calibration error (KCE). We keep it
# simple here, and use the tensor product kernel
# ```math
# k\big((\mu, y), (\mu', y')\big) = \delta_{y,y'} \exp{\bigg(-\frac{{\|\mu - \mu'\|}_2^2}{2\nu^2} \bigg)}
# ```
# with length scale $\nu > 0$ for predictions $\mu,\mu'$ and corresponding targets $y, y'$.
# For simplicity, we estimate length scale $\nu$ with the median heuristic.

distances = pairwise(SqEuclidean(), train_probs)
λ = sqrt(median(distances[i] for i in CartesianIndices(distances) if i[1] < i[2]))
kernel = (GaussianKernel() ∘ ScaleTransform(inv(λ))) ⊗ WhiteKernel();

# We obtain the following biased estimates of the squared KCE (SKCE):

skce = BiasedSKCE(kernel)
skce(train_probs, train_species)

#-

skce(val_probs, val_species)

# Similar to the biased estimates of the ECE, the biased estimates of the SKCE are always
# non-negative. The unbiased estimates can be negative as well, in particular if the model
# is (close to being) calibrated:

skce = UnbiasedSKCE(kernel)
skce(train_probs, train_species)

#-

skce(val_probs, val_species)

# When the datasets are large, the quadratic sample complexity of the standard biased and
# unbiased estimators of the SKCE can become prohibitive. In these cases, one can resort to
# an estimator that averages estimates of non-overlapping blocks of samples. This estimator
# allows to trade off computational cost for increased variance.
#
# Here we consider the extreme case of blocks with two samples, which yields an estimator
# with linear sample complexity:

skce = BlockUnbiasedSKCE(kernel, 2)
skce(train_probs, train_species)

#-

skce(val_probs, val_species)


include("Lanczos_Ritz_Utilities.jl")
include("../../src/iterative_lanczos.jl")

### ============================================================================
### Lanczos vector project onto the extreme ritz value's eigenvalues
###

function PerformAndProject(n=64; offset=20, exact=false)
    A = Diagonal(LinRange(-1, 1, n).^3)
    il = IterativeLanczos(A, ones(n))
    il.reorthogonalize = exact
    for _ in 1: offset - 1
        il()
    end
    ProjOnRitz = Vector()
    for II in offset:n - 1
        q = il()
        T = il|>GetTMatrix
        Q = il|>GetQMatrix

        s = eigvecs(T)[:, end:-1: (end - 2)]
        push!(ProjOnRitz, q'*Q*s)
    end

return vcat(ProjOnRitz...) end

# ----- Floats --------
ProjOnRitz = PerformAndProject()
ProjOnRitz = ProjOnRitz .|> abs
Iters = 20:63
fig1 = plot(
    Iters,
    ProjOnRitz[:, 1],
    yaxis=:log,
    legend=:bottomleft, size=(750, 500),
    label="\$q^T_kQ_ks_1\$",
    dpi=300, leftmargin=5*Plots.mm
)
plot!(fig1, Iters, ProjOnRitz[:, 2], linestyle=:dashdot, label="\$q^T_kQ_ks_2\$")
plot!(fig1, Iters, ProjOnRitz[:, 3], linestyle=:dash, label="\$q^T_kQ_ks_3\$")
title!("Lanczos vec proj onto Ritz Vectors (Floats)")
yaxis!("\$|q^T_kQS|\$")
xaxis!("iterations")
fig1|>display
savefig(fig1, "$(@__DIR__)/plots/lanczos_proj_on_ritz_float.png")

# ------ Exact --------
ProjOnRitz = PerformAndProject(exact=true)
ProjOnRitz = ProjOnRitz .|> abs
Iters = 20:63
fig1 = plot(
    Iters,
    ProjOnRitz[:, 1],
    yaxis=:log,
    legend=:bottomleft, size=(750, 500),
    label="\$q^T_kQ_ks_1\$",
    dpi=300, leftmargin=5*Plots.mm
)

plot!(fig1, Iters, ProjOnRitz[:, 2], linestyle=:dashdot, label="\$q^T_kQ_ks_2\$")
plot!(fig1, Iters, ProjOnRitz[:, 3], linestyle=:dash, label="\$q^T_kQ_ks_3\$")
title!("Lanczos vec proj on ritz vectors (largest 3) (Exact)")
yaxis!("\$|q^T_kQS|\$")
xaxis!("iterations")
fig1|>display
savefig(fig1, "$(@__DIR__)/plots/lanczos_proj_on_ritz_exact.png")


module TensorCore

using LinearAlgebra

export ⊙, hadamard, hadamard!
export ⊗, tensor, tensor!
export ⊡, boxdot, boxdot!

"""
    hadamard(a, b)
    a ⊙ b

For arrays `a` and `b`, perform elementwise multiplication.
`a` and `b` must have identical `axes`.

`⊙` can be passed as an operator to higher-order functions.

# Examples
```jldoctest; setup=:(using TensorCore)
julia> a = [2, 3]; b = [5, 7];

julia> a ⊙ b
2-element Array{$Int,1}:
 10
 21

julia> a ⊙ [5]
ERROR: DimensionMismatch("Axes of `A` and `B` must match, got (Base.OneTo(2),) and (Base.OneTo(1),)")
[...]
```

See also `hadamard!(y, a, b)`.
"""
function hadamard(A::AbstractArray, B::AbstractArray)
    @noinline throw_dmm(axA, axB) = throw(DimensionMismatch("Axes of `A` and `B` must match, got $axA and $axB"))

    axA, axB = axes(A), axes(B)
    axA == axB || throw_dmm(axA, axB)
    return map(*, A, B)
end
const ⊙ = hadamard

"""
    hadamard!(dest, A, B)

Similar to `hadamard(A, B)` (which can also be written `A ⊙ B`), but stores its results in
the pre-allocated array `dest`.
"""
function hadamard!(dest::AbstractArray, A::AbstractArray, B::AbstractArray)
    @noinline function throw_dmm(axA, axB, axdest)
        throw(DimensionMismatch("`axes(dest) = $axdest` must be equal to `axes(A) = $axA` and `axes(B) = $axB`"))
    end

    axA, axB, axdest = axes(A), axes(B), axes(dest)
    ((axdest == axA) & (axdest == axB)) || throw_dmm(axA, axB, axdest)
    @simd for I in eachindex(dest, A, B)
        @inbounds dest[I] = A[I] * B[I]
    end
    return dest
end

"""
    tensor(A, B)
    A ⊗ B

Compute the tensor product of `A` and `B`.
If `C = A ⊗ B`, then `C[i1, ..., im, j1, ..., jn] = A[i1, ... im] * B[j1, ..., jn]`.

For vectors `v` and `w`, the Kronecker product is related to the tensor product by
`kron(v,w) == vec(w ⊗ v)` or `w ⊗ v == reshape(kron(v,w), (length(w), length(v)))`.

# Examples
```jldoctest; setup=:(using TensorCore)
julia> a = [2, 3]; b = [5, 7, 11];

julia> a ⊗ b
2×3 Array{$Int,2}:
 10  14  22
 15  21  33
```
See also `tensor!(Y,A,B)`.
"""
tensor(A::AbstractArray, B::AbstractArray) = [a*b for a in A, b in B]
const ⊗ = tensor

const CovectorLike{T} = Union{Adjoint{T,<:AbstractVector},Transpose{T,<:AbstractVector}}
function tensor(u::AbstractArray, v::CovectorLike)
    # If `v` is thought of as a covector, you might want this to be two-dimensional,
    # but thought of as a matrix it should be three-dimensional.
    # The safest is to avoid supporting it at all. See discussion in #35150.
    error("`tensor` is not defined for co-vectors, perhaps you meant `*`?")
end
function tensor(u::CovectorLike, v::AbstractArray)
    error("`tensor` is not defined for co-vectors, perhaps you meant `*`?")
end
function tensor(u::CovectorLike, v::CovectorLike)
    error("`tensor` is not defined for co-vectors, perhaps you meant `*`?")
end

"""
    tensor!(dest, A, B)

Similar to `tensor(A, B)` (which can also be written `A ⊗ B`), but stores its results in
the pre-allocated array `dest`.
"""
function tensor!(dest::AbstractArray, A::AbstractArray, B::AbstractArray)
    @noinline function throw_dmm(axA, axB, axdest)
        throw(DimensionMismatch("`axes(dest) = $axdest` must concatenate `axes(A) = $axA` and `axes(B) = $axB`"))
    end

    axA, axB, axdest = axes(A), axes(B), axes(dest)
    axes(dest) == (axA..., axB...) || throw_dmm(axA, axB, axdest)
    if IndexStyle(dest) === IndexCartesian()
        for IB in CartesianIndices(axB)
            @inbounds b = B[IB]
            @simd for IA in CartesianIndices(axA)
                @inbounds dest[IA,IB] = A[IA]*b
            end
        end
    else
        i = firstindex(dest)
        @inbounds for b in B
            @simd for a in A
                dest[i] = a*b
                i += 1
            end
        end
    end
    return dest
end

export boxdot, ⊡, boxdot!

"""
    boxdot(A,B) = A ⊡ B    # \\boxdot

Generalised matrix multiplication: Contracts the last dimension of `A` with
the first dimension of `B`, for any `ndims(A)` & `ndims(B)`.
If both are vectors, then it returns a scalar `== sum(A .* B)`.

# Examples
```jldoctest; setup=:(using TensorCore)
julia> A = rand(3,4,5); B = rand(5,6,7);

julia> size(A ⊡ B)
(3, 4, 6, 7)

julia> typeof(rand(5) ⊡ rand(5))
Float64

julia> try B ⊡ A catch err println(err) end
DimensionMismatch("neighbouring axes of `A` and `B` must match, got Base.OneTo(7) and Base.OneTo(3)")
```
This is the same behaviour as Mathematica's function `Dot[A, B]`.
It is not identicaly to Python's `numpy.dot(A, B)`, which contracts with the second-last
dimension of `B` instead of the first, but both keep all the other dimensions.
Unlike Julia's `LinearAlgebra.dot`, it does not conjugate `A`, so these two agree only
for real-valued vectors.

When interacting with `Adjoint` vectors, this always obeys `(x ⊡ y)' == y' ⊡ x'`,
and hence may sometimes return another `Adjoint` vector. (And similarly for `Transpose`.)

```jldoctest; setup=:(using TensorCore)
julia> M = rand(5,5); v = rand(5);

julia> typeof(v ⊡ M')
Array{Float64,1}

julia> typeof(M ⊡ v')  # adjoint of the previous line
Adjoint{Float64,Array{Float64,1}}

julia> typeof(v' ⊡ M')  # same as *, and equal to adjoint(M ⊡ v)
Adjoint{Float64,Array{Float64,1}}

julia> typeof(v' ⊡ v)
Float64
```
See also `boxdot!(Y,A,B)`, which is to `⊡` as `mul!` is to `*`.
"""
function boxdot(A::AbstractArray, B::AbstractArray)
    Amat = _squash_left(A)
    Bmat = _squash_right(B)

    axA, axB = axes(Amat,2), axes(Bmat,1)
    axA == axB || _throw_dmm(axA, axB)

    return _boxdot_reshape(Amat * Bmat, A, B)
end

const ⊡ = boxdot

@noinline _throw_dmm(axA, axB) = throw(DimensionMismatch("neighbouring axes of `A` and `B` must match, got $axA and $axB"))

_squash_left(A::AbstractArray) = reshape(A, :,size(A,ndims(A)))
_squash_left(A::AbstractMatrix) = A

_squash_right(B::AbstractArray) = reshape(B, size(B,1),:)
_squash_right(B::AbstractVecOrMat) = B

function _boxdot_reshape(AB::AbstractArray, A::AbstractArray{T,N}, B::AbstractArray{S,M}) where {T,N,S,M}
    ax = ntuple(i -> i<N ? axes(A, i) : axes(B, i-N+2), Val(N+M-2))
    reshape(AB, ax) # some cases don't come here, so this doesn't really support OffsetArrays
end

# These can skip final reshape:
_boxdot_reshape(AB::AbstractVecOrMat, A::AbstractMatrix, B::AbstractVecOrMat) = AB

# These produce scalar output:
function boxdot(A::AbstractVector, B::AbstractVector)
    axA, axB = axes(A,1), axes(B,1)
    axA == axB || _throw_dmm(axA, axB)
    if eltype(A) <: Number
        return transpose(A)*B
    else
        return sum(a*b for (a,b) in zip(A,B))
    end
end

# Multiplication by a scalar:
boxdot(A::AbstractArray, b::Number) = A*b
boxdot(a::Number, B::AbstractArray) = a*B
boxdot(a::Number, b::Number) = a*b

using LinearAlgebra: AdjointAbsVec, TransposeAbsVec, AdjOrTransAbsVec

# Adjont and Transpose, vectors or almost (returning a scalar)
boxdot(A::AdjointAbsVec, B::AbstractVector) = A * B
boxdot(A::TransposeAbsVec, B::AbstractVector) = A * B

boxdot(A::AbstractVector, B::AdjointAbsVec) = A ⊡ vec(B)
boxdot(A::AbstractVector, B::TransposeAbsVec) = A ⊡ vec(B)

boxdot(A::AdjointAbsVec, B::AdjointAbsVec) = adjoint(adjoint(B) ⊡ adjoint(A))
boxdot(A::AdjointAbsVec, B::TransposeAbsVec) = vec(A) ⊡ vec(B)
boxdot(A::TransposeAbsVec, B::AdjointAbsVec) = vec(A) ⊡ vec(B)
boxdot(A::TransposeAbsVec, B::TransposeAbsVec) = transpose(transpose(B) ⊡ transpose(A))

# ... with a matrix (returning another such)
boxdot(A::AdjointAbsVec, B::AbstractMatrix) = A * B
boxdot(A::TransposeAbsVec, B::AbstractMatrix) = A * B

boxdot(A::AbstractMatrix, B::AdjointAbsVec) = (B' ⊡ A')'
boxdot(A::AbstractMatrix, B::TransposeAbsVec) = transpose(transpose(B) ⊡ transpose(A))

# ... and with higher-dim (returning a plain array)
boxdot(A::AdjointAbsVec, B::AbstractArray) = vec(A) ⊡ B
boxdot(A::TransposeAbsVec, B::AbstractArray) = vec(A) ⊡ B

boxdot(A::AbstractArray, B::AdjointAbsVec) = A ⊡ vec(B)
boxdot(A::AbstractArray, B::TransposeAbsVec) = A ⊡ vec(B)

"""
    boxdot!(Y, A, B, α=1, β=0)

In-place version of `boxdot`, i.e. `Y .= (A ⊡ B) .* β .+ Y .* α`.
Like 5-argument `mul!`, the use of `α, β` here requires Julia 1.3 or later.
"""
function boxdot! end

if VERSION < v"1.3" # Then 5-arg mul! isn't defined

    function boxdot!(Y::AbstractArray, A::AbstractArray, B::AbstractArray)
        szY = prod(size(A)[1:end-1]), prod(size(B)[2:end])
        mul!(reshape(Y, szY), _squash_left(A), _squash_right(B))
        Y
    end

    boxdot!(Y::AbstractArray, A::AbstractArray, B::AdjOrTransAbsVec) = boxdot!(Y, A, vec(B))

else

    function boxdot!(Y::AbstractArray, A::AbstractArray, B::AbstractArray, α::Number=true, β::Number=false)
        szY = prod(size(A)[1:end-1]), prod(size(B)[2:end])
        mul!(reshape(Y, szY), _squash_left(A), _squash_right(B), α, β)
        Y
    end

    # For boxdot!, only where mul! behaves differently:
    boxdot!(Y::AbstractArray, A::AbstractArray, B::AdjOrTransAbsVec,
        α::Number=true, β::Number=false) = boxdot!(Y, A, vec(B))

end

"""
    TensorCore._adjoint(A)

This extends `adjoint` to understand higher-dimensional arrays, always reversing the
order of dimensions. On Julia 1.5 and later, the symbol `'` can be overloaded locally
as `var"'"`, as shown below.

Then `(x ⊡ y)' == y' ⊡ x'` holds for `x` and `y` arrays of any dimension.

# Examples
```jldoctest; setup=:(using TensorCore)
julia> T3 = rand(3,4,5); v = rand(5);

julia> size(T3 ⊡ v')
(3, 4)

julia> let var"'" = TensorCore._adjoint
         v ⊡ T3' ≈ (T3 ⊡ v')'
       end
true
```
"""
_adjoint(x) = adjoint(x)
_adjoint(x::AbstractVecOrMat) = adjoint(x)
_adjoint(x::AbstractArray{T,N}) where {T<:Number,N} = conj(PermutedDimsArray(x, ntuple(i -> N-i+1, N)))
_adjoint(x::AbstractArray{T,N}) where {T,N} = adjoint.(PermutedDimsArray(x, ntuple(i -> N-i+1, N)))

_transpose(x) = transpose(x)
_transpose(x::AbstractVecOrMat) = transpose(x)
_transpose(x::AbstractArray{T,N}) where {T<:Number,N} = PermutedDimsArray(x, ntuple(i -> N-i+1, N))
_transpose(x::AbstractArray{T,N}) where {T,N} = transpose.(PermutedDimsArray(x, ntuple(i -> N-i+1, N)))

end


using SymPy: @vars, Sym, solve
using SymPy: simplify
using Zygote

dot(x,y) = sum(x .* y)

function generate_vector(v::String, n::Int; start::Int = 1, kwargs::String = "")
    data::Vector{Sym} = []
    for i ∈ start:n
        seqdigit = [c for c in string(i)]
        v_s = "$v" * join(seqdigit)
        eval(Meta.parse("@vars $v_s $kwargs"))
        push!(data, eval(Meta.parse("$(v_s)")))
    end
    return data
end

D = 2 # Dimension of data
N = 3 # Num of data

function generate_matrix(
    v::String,
    row::Int,
    col::Int;
    rstart::Int = 1,
    cstart::Int = 1,
    kwargs::String = "",
)::Matrix{Sym}
    X_ = Vector{Sym}[]
    for n = rstart:row
        xn = "$v$n"
        eval(Meta.parse("xn=generate_vector(\"x$n\", $col, start = $cstart, kwargs=\"$kwargs\")"))
        push!(X_, eval(Meta.parse("xn")))
    end
    Xᵀ = Matrix{Sym}(hcat(X_...))
    X = transpose(Xᵀ)
    X
end

X = generate_matrix("x", N, D, cstart = 0, kwargs = "real=true")
y = generate_vector("y", N, kwargs = "real=true")
w = generate_vector("w", D, start = 0, kwargs = "real=true")

# Numerical
# X = rand(N, D + 1)
# y = rand(N)

X[:, 1] .= 1.0 # set xn0 = 1 for n in 1:N

ŷ(X) = X * w
E(X) = 1 / 2 * (dot(y, y) - 2 * dot(y, ŷ(X)) + dot(ŷ(X), ŷ(X)))

@info "calc-gradient"

∇E = gradient(Params([w])) do
    E(X)
end

@info "solve"
solution = solve(∇E[w], w)
solution = [solution[k] for k in w]
@info "solve-with-norma-equation"
# solve ŵ with normal equation
theoretical = inv(Matrix{eltype(X)}(X' * X)) * Matrix{eltype(X)}(X') * y

@info "check isequal"
if eltype(theoretical) == Sym
    @assert all(solution .- theoretical .|> simplify .== 0)
else
    solution=Array{Float64}(solution)
    @assert all(isapprox.(solution ,theoretical, atol=1e-6))
end
@info "done"


export MinimalDominatingSet

struct MinimalDominatingSet end

"""
    dominating_set(g, MinimalDominatingSet(); seed=-1)

Find a set of vertices that consitute a dominating set (all vertices in `g` are either adjacent to a vertex
in the set or is a vertex in the set) and it is not possible to delete a vertex from the set
without sacrificing the dominating property.

### Implementation Notes
Initially, every vertex is in the dominating set.
In some random order, we check if the removal of a vertex from the set will destroy the
dominating property. If no, the vertex is removed from the dominating set.

### Performance
Runtime: ``\\mathcal{O}(|V|+|E|)``
Memory: ``\\mathcal{O}(|V|)``

### Optional Arguments
- If `seed >= 0`, a random generator is seeded with this value.
"""
function dominating_set(
    g::AbstractGraph{T},
    alg::MinimalDominatingSet;
    seed::Int=-1
    ) where T <: Integer

    nvg = nv(g)
    in_dom_set = trues(nvg)
    length_ds = Int(nvg)
    dom_degree = degree(g)
    @inbounds @simd for v in vertices(g)
        dom_degree[v] -= (has_edge(g, v, v) ? 1 : 0)
    end

    for v in randperm(getRNG(seed), nvg)
    	(dom_degree[v] == 0) && continue #It is not adjacent to any dominating vertex
    	#Check if any vertex is depending on v to be dominated
        dependent = findfirst(u -> !in_dom_set[u] && dom_degree[u] <= 1, neighbors(g, v))

        (dependent != nothing) && continue
        in_dom_set[v] = false
        length_ds -= 1
        dom_degree[neighbors(g, v)] .-= 1
    end

    return LightGraphs.findall!(in_dom_set, Vector{T}(undef, length_ds))
end


# [test/rkfixed_tests.jl]
using Test

@testset "RK Fixed Tests" begin

    using JLD2
    import ODEHybrid.rkfixed, ODEHybrid.ODE_SET

    @testset "Example from Code Comments" begin
        a = [0   0   0  0;
            0.5 0   0  0;
            0   0.5 0  0;
            0   0   1  0];
        b = [1 2 2 1]/6;
        c = [0 0.5 0.5 1];

        ts = [0.0 10.0]
        x0 = [1.0; 0.0]
        dt = 0.1

        # mat"""
        # [$t_mat, $x_mat] = rkfixed(@(t,x) [-x(2); x(1)], $ts, [1; 0], $dt, $a, $b, $c);
        # % rkfixed(@(t,x) [-x(2); x(1)], $ts, $x0, $dt, $a, $b, $c)
        # """

        @load "solutions/rkfixed_example_from_code.jld2" t_mat x_mat;

        t_jul, x_jul = rkfixed( (t, x) -> [-x[2]; x[1]], ts, x0, dt, a, b, c)

        @test isapprox(x_mat, x_jul, atol = 1e-12)
        @test isapprox(t_mat, t_jul, atol = 1e-12)
    end;

    @testset "Default tableau values" begin
        ts = [0.0 10.0]
        x0 = [1.0; 0.0]
        dt = 0.1

        # mat"""
        # [$t_mat, $x_mat] = rkfixed(@(t,x) [-x(2); x(1)], $ts, $x0, $dt);
        # """

        @load "solutions/rkfixed_default_tab.jld2" t_mat x_mat

        t_jul, x_jul = rkfixed( (t, x) -> [-x[2]; x[1]], ts, x0, dt)

        @test isapprox(x_mat, x_jul, atol = 1e-12)
        @test isapprox(t_mat, t_jul, atol = 1e-12)
    end;

    @testset "Function with time" begin
        ts = [-5.0 5.0]
        x0 = [5.0; 2.5; 0.0; -2.5; 5.0]
        dt = 0.25

        # mat"""
        # [$t_mat, $x_mat] = rkfixed(@(t,x) [-t * x(2); x(3); -x(4); x(5); -x(1)], $ts, $x0, $dt);
        # """

        @load "solutions/rkfixed_func_with_time.jld2" t_mat x_mat

        t_jul, x_jul = rkfixed( (t, x) -> [-t * x[2]; x[3]; -x[4]; x[5]; -x[1]], ts, x0, dt)

        @test isapprox(x_mat, x_jul, atol = 1e-12)
        @test isapprox(t_mat, t_jul, atol = 1e-12)
    end;

    @testset "Negative time step" begin
        ts = [12.0 -5.0]
        x0 = [1.0; -1.0; 1.0; -1.0]
        dt = -0.25

        # mat"""
        #     [$t_mat, $x_mat] = rkfixed(@(t,x) [-t * x(2); t * x(3); -t * x(4); t * x(1)], $ts, $x0, $dt);
        # """

        @load "solutions/rkfixed_neg_time_step.jld2" t_mat x_mat

        t_jul, x_jul = rkfixed( (t, x) -> [-t * x[2]; t * x[3]; -t * x[4]; t * x[1]], ts, x0, dt)

        @test isapprox(x_mat, x_jul, atol = 1e-12)
        @test isapprox(t_mat, t_jul, atol = 1e-12)
    end;

    # Plot currently supressed
    @testset "ODE_SET with output function" begin
        ts = [-10.0 0]
        x0 = [1.0; -1.0]
        dt = 0.25

        # mat"""
        #     ode = @(t, x) [-x(2); x(1)];
        #     % options = odeset('OutputFcn', @odeplot);
        #     options = odeset();
        #     options.MaxStep = $dt;
        #     [$t_mat, $x_mat] = rkfixed(ode, $ts, $x0, options)
        # """

        @load "solutions/rkfixed_odeset_with_output.jld2" t_mat x_mat

        options = ODE_SET()
        # options.OutputFcn = odeplot
        options.MaxStep = dt

        t_jul, x_jul = rkfixed( (t, x) -> [-x[2]; x[1]], ts, x0, options)

        @test isapprox(x_mat, x_jul, atol = 1e-12)
        @test isapprox(t_mat, t_jul, atol = 1e-12)
    end;

end;


using Makie
using AbstractPlotting: textslider
using Distributions: Normal, pdf

function single()
    s_μ, obs_μ = textslider(range(-1., stop=1., length=101), "μ", start=0)
    σ = 1
    #=
    pos = lift(obs_μ) do μ
            d = Normal(μ, σ)
            [Point2f0(x, pdf(d, x)) for x in -5.:0.02:5.]
        end
    =#

    pos = lift(μ ->begin
                        d = Normal(μ, σ)
                        [Point2f0(x, pdf(d, x)) for x in -5.:0.02:5.]
                   end,
                obs_μ
           )

    p = lines(pos)
    RecordEvents(
        hbox(
            p,
            vbox(s_μ),
            parent=Scene()
        ),
        "output"
    )
end

function main()
    s_μ, obs_μ = textslider(range(-1., stop=1., length=101), "μ", start=0)
    s_σ, obs_σ = textslider(range(0.1, stop=5.1, length=101), "σ", start=1)

    #=
    params = lift(tuple, obs_μ, obs_σ)
    pos = lift(params) do (μ, σ)
        d = Normal(μ, σ)
        map(LinRange(-5, 5, 100)) do x
            Point2f0(x, pdf(d, x))
        end
    end
    =#

    #=
    pos = lift((μ, σ) -> begin
                            d = Normal(μ, σ)
                            [Point2f0(x, pdf(d, x))
                                for x in -5.:0.02:5.]
                         end,
               obs_μ, obs_σ)
    =#

    pos = lift(lift(tuple, obs_μ, obs_σ)) do (μ, σ)
            d = Normal(μ, σ)
            [Point2f0(x, pdf(d, x)) for x in -5.:0.02:5.]
        end
    p = lines(pos)

    RecordEvents(
        hbox(
            p,
            vbox(s_μ, s_σ),
            parent=Scene()
        ),
        "output"
    )

end

main()


# An implementation of CGS for the solution of the square linear system Ax = b.
#
# This method is described in
#
# Y. Saad, Iterative methods for sparse linear systems.
# PWS Publishing Company, Boston, USA, 1996.
#
# P. Sonneveld, CGS, A Fast Lanczos-Type Solver for Nonsymmetric Linear systems.
# SIAM Journal on Scientific and Statistical Computing, 10(1), pp. 36--52, 1989.
#
# Alexis Montoison, <alexis.montoison@polymtl.ca>
# Montreal, October 2018.

export cgs

"""Solve the consistent linear system Ax = b using conjugate gradient squared algorithm.

From "Iterative Methods for Sparse Linear Systems (Y. Saad)" :

«The method is based on a polynomial variant of the conjugate gradients algorithm.
Although related to the so-called bi-conjugate gradients (BCG) algorithm,
it does not involve adjoint matrix-vector multiplications, and the expected convergence
rate is about twice that of the BCG algorithm.

The Conjugate Gradient Squared algorithm works quite well in many cases.
However, one difficulty is that, since the polynomials are squared, rounding errors
tend to be more damaging than in the standard BCG algorithm. In particular, very
high variations of the residual vectors often cause the residual norms computed
to become inaccurate.

TFQMR and BICGSTAB were developed to remedy this difficulty.»

This implementation allows a right preconditioner M.
"""
function cgs(A :: AbstractLinearOperator, b :: AbstractVector{T};
             M :: AbstractLinearOperator=opEye(size(A,1)),
             atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6,
             itmax :: Int=0, verbose :: Bool=false) where {T <: Number}

  m, n = size(A)
  m == n || error("System must be square")
  length(b) == m || error("Inconsistent problem size")
  verbose && @printf("CGS: system of size %d\n", n)

  # Initial solution x₀ and residual r₀.
  x = zeros(T, n) # x₀
  r = copy(b)     # r₀
  # Compute ρ₀ = < r₀,r₀ > and residual norm ‖r₀‖₂.
  ρ = @kdot(n, r, r)
  rNorm = sqrt(ρ)
  rNorm == 0 && return x, SimpleStats(true, false, [rNorm], [], "x = 0 is a zero-residual solution")

  iter = 0
  itmax == 0 && (itmax = 2*n)

  rNorms = [rNorm;]
  ε = atol + rtol * rNorm
  verbose && @printf("%5d  %7.1e\n", iter, rNorm)

  # Set up workspace.
  u = copy(r)  # u₀
  p = copy(r)  # p₀
  q = zeros(n) # q₋₁

  # Stopping criterion.
  solved = rNorm ≤ ε
  tired = iter ≥ itmax
  status = "unknown"

  while !(solved || tired)

    y = M * p                 # yₘ = M⁻¹pₘ
    v = A * y                 # vₘ = Ayₘ
    σ = @kdot(n, v, b)        # σₘ = < AM⁻¹pₘ,r₀ >
    α = ρ / σ                 # αₘ = ρₘ / σₘ
    @. q = u - α * v          # qₘ = uₘ - αₘ * AM⁻¹pₘ
    @kaxpy!(n, 1.0, q, u)     # uₘ₊½ = uₘ + qₘ
    z = M * u                 # zₘ = M⁻¹uₘ₊½
    @kaxpy!(n, α, z, x)       # xₘ₊₁ = xₘ + αₘ * M⁻¹(uₘ + qₘ)
    w = A * z                 # wₘ = AM⁻¹(uₘ + qₘ)
    @kaxpy!(n, -α, w, r)      # rₘ₊₁ = rₘ - αₘ * AM⁻¹(uₘ + qₘ)
    ρ_next = @kdot(n, r, b)   # ρₘ₊₁ = < rₘ₊₁,r₀ >
    β = ρ_next / ρ            # βₘ = ρₘ₊₁ / ρₘ
    @. u = r + β * q          # uₘ₊₁ = rₘ₊₁ + βₘ * qₘ
    @kaxpby!(n, 1.0, q, β, p) # pₘ₊₁ = uₘ₊₁ + βₘ * (qₘ + βₘ * pₘ)
    @kaxpby!(n, 1.0, u, β, p)

    # Update ρ.
    ρ = ρ_next # ρₘ ← ρₘ₊₁

    # Update iteration index.
    iter = iter + 1

    # Compute residual norm ‖rₘ‖₂.
    rNorm = @knrm2(n, r)
    push!(rNorms, rNorm)

    # Update stopping criterion.
    solved = rNorm ≤ ε
    tired = iter ≥ itmax
    verbose && @printf("%5d  %7.1e\n", iter, rNorm)
  end
  verbose && @printf("\n")

  status = tired ? "maximum number of iterations exceeded" : "solution good enough given atol and rtol"
  stats = SimpleStats(solved, false, rNorms, T[], status)
  return (x, stats)
end


# Bi-objective linear problem (bilp)
#
# Example 6.3 (from Steuer, 1985), page 154 of
# Multicriteria Optimization (2nd edt), M. Ehrgott, Springer 2005.

# ---- Packages to use
using vOptGeneric, JuMP, GLPK

# ---- setting the model + values
bilp = vModel( GLPK.Optimizer )

@variable( bilp, x1 >= 0 )
@variable( bilp, x2 >= 0 )

@addobjective( bilp, Min, 3*x1 + x2 )
@addobjective( bilp, Min, -x1 - 2*x2 )

@constraint( bilp, cst1, x2 <= 3 )
@constraint( bilp, cst2, 3*x1 - x2 <= 6 )

# ---- Invoking the solver (lexicographic method)
vSolve( bilp, method=:lex )

# ---- Querying the results
Y_N = getY_N( bilp )

# ---- Displaying the results for lex(1,2) and lex(2,1)
for s = 1:2
    print("X = [")
    print(round(getvOptData(bilp).X_E[s][1],digits=5))
    print(" ")
    print(round(getvOptData(bilp).X_E[s][2],digits=5))
    print("] | Z = [")
    print(round(Y_N[s][1], digits=5))
    print(" ")
    print(round(Y_N[s][2], digits=5))
    println("]")
end


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

"""
Return the Lar.Struct of octahedron
"""
function octahedron(l)
    a = l * (sqrt(3) / 2)
    h = sqrt(a^2 - ((l / 2)^2))
    #h=0.70710678118654
    V = [0 l 0 l l/2; 0 0 l l l/2; 0 0 0 0 h]
    VV = [[1], [2], [3], [4], [5]]
    EV = [[1, 2], [2, 3], [3, 4], [4, 1], [1, 5], [2, 5], [3, 5], [4, 5]]
    FV = [[1, 2, 5], [1, 3, 5], [3, 4, 5], [2, 4, 5]]
    CV = [[1, 2, 3, 4, 5]]
    tetrahedron = V, EV, FV, CV
    return Lar.Struct([tetrahedron, Lar.t(0, l, 0), Lar.r(π, 0, 0), tetrahedron])
end
"""
Return the Lar.Struct of an irregular octahedron whit h as height
"""
function octahedron(l, h)
    V = [0 l 0 l l / 2; 0 0 l l l / 2; 0 0 0 0 h]
    VV = [[1], [2], [3], [4], [5]]
    EV = [[1, 2], [3, 4], [4, 5], [1, 5], [1, 5], [2, 5], [3, 5], [4, 5]]
    FV = [[1, 2, 5], [1, 3, 5], [3, 4, 5], [2, 4, 5]]
    CV = [[1, 2, 3, 4, 5]]
    tetrahedron = V, EV, FV, CV
    return Lar.Struct([tetrahedron, Lar.t(0, l, 0), Lar.r(π, 0, 0), tetrahedron])
end
"""
Return the Lar.Struct of a rectangular base octahedron whit h as height
"""
function rectangularbase_octahedron(ab, da, h)
    V = [0 ab 0 ab ab / 2; 0 0 da da da / 2; 0 0 0 0 h]
    V2 = [0 ab 0 ab ab / 2; 0 0 da da da / 2; 0 0 0 0 -h]
    VV = [[1], [2], [3], [4], [5]]
    EV = [[1, 2], [3, 4], [4, 5], [1, 5], [1, 5], [2, 5], [3, 5], [4, 5]]
    FV = [[1, 2, 5], [1, 3, 5], [3, 4, 5], [2, 4, 5]]
    CV = [[1, 2, 3, 4, 5]]
    tetrahedron = V, EV, FV, CV
    tetrahedron2 = V2, EV, FV, CV
    return Lar.Struct([tetrahedron, tetrahedron2])
end


using LinearAlgebra: I, norm, qr
using Parameters: @with_kw

function _orthogonal_complement_matrix( Y, p = Inf )
    Q, _ = qr(Y)
    Z = Q[:, size(Y,2) + 1 : end]
    if size(Z,2) > 0
        Z ./= norm.( eachcol(Z), p )'
    end
    return Z
end

# stateful (Y,Z are stored so that I have access to them later)
@with_kw mutable struct AffinelyIndependentPointFilter{F <: AbstractFloat, VF <: AbstractVector{F}, SV <: AbstractVector{VF}, I<:Real}
	x_0 :: VF

	seeds :: SV = Vector{VF}

	shifted_seeds :: SV = [ s .- x_0 for s in seeds]

	candidate_indices :: Vector{Int} = collect(eachindex(seeds))

	n :: Int = length(x_0)

	Y :: Matrix{F} = Matrix{eltype(x_0)}(undef, n, 0)
	Z :: Matrix{F} = Matrix{eltype(x_0)}(I(n))

	p :: I = Inf 	# which vector norm to use

	pivot_val :: F = 1e-3

	return_indices :: Bool = false

	@assert n > 0 "`x_0` must not be empty and `n` must be positive."
	@assert n <= length(x_0) "Maximum number must be lower than `length(x_0)`."
end

function reset!( filter :: AffinelyIndependentPointFilter{F,VF,SF} ) where{F,VF,SF}
	filter.candidate_indices = collect(eachindex(filter.seeds))
	filter.Y = Matrix{F}(undef, filter.n, 0)
	filter.Z = Matrix{F}(I(filter.n))
	return nothing
end

function Base.iterate( filter :: AffinelyIndependentPointFilter{F,VF,SV} ) :: Union{Nothing, Tuple{<:Union{Int,VF}, Int}} where{F,VF,SV}
	isempty(filter.shifted_seeds) && return nothing
	_, i = findmax( norm.(filter.shifted_seeds, filter.p) )

	if length(filter.candidate_indices) != length(filter.seeds)
		reset!(filter)
	end

	filter.Y = hcat(filter.Y, filter.shifted_seeds[i])
	filter.Z = _orthogonal_complement_matrix( filter.Y, filter.p )

	setdiff!(filter.candidate_indices, i)

	if filter.return_indices
		return i, 1
	else
		return filter.seeds[i], 1
	end
end

function Base.iterate( filter :: AffinelyIndependentPointFilter{F,VF,SV}, num_found :: Int ) :: Union{Nothing,Tuple{<:Union{VF,Int}, Int}} where{F,VF,SV}
	num_found == filter.n && return nothing			# found enough points already
	isempty(filter.candidate_indices) && return nothing	# no more points to search

	Z = filter.Z

	best_val = -F(Inf)
	best_index = -1

	for i ∈ filter.candidate_indices
		x = filter.shifted_seeds[i]
		proj_x_Z = norm( Z*(Z'*x), filter.p )

		if proj_x_Z > best_val
			best_val = proj_x_Z
			best_index = i
			# NOTE we could check `proj_x_Z > filter.pivot_val` and break early here instead of finding the maximizer
		end
	end

	if best_val > filter.pivot_val
		i = best_index
		filter.Y = hcat(filter.Y, filter.shifted_seeds[i])
		filter.Z = _orthogonal_complement_matrix( filter.Y, filter.p )
		setdiff!(filter.candidate_indices, i)

		if filter.return_indices
			return i, num_found + 1
		else
			return filter.seeds[i], num_found + 1
		end
	end

	@debug "No point was sufficiently linearly independent."
	return nothing
end

Base.IteratorSize( ::AffinelyIndependentPointFilter ) = Base.SizeUnknown()
Base.IteratorEltype( :: AffinelyIndependentPointFilter ) = Base.HasEltype()
function Base.eltype( filter :: AffinelyIndependentPointFilter{F,VF,SV} ) where {F,VF,SV}
	if filter.return_indices
		return Int
	else
		return VF
	end
end


using LightGraphs
using LinearAlgebra: UpperTriangular

"""
UniDirectionalBeta
"""
function beta_graph(beta::Vector; seed::Integer=-1)
    n = length(beta)
    coefM = UpperTriangular(beta' .+ beta)
    g = SimpleGraph(n)
    for i in 1:n
        for j in i:n
            expcoef_ij = exp(coefM[i,j])
            if rand() < (expcoef_ij / (1 + expcoef_ij))
                add_edge!(g, i, j)
            end
        end
    end
    return g
end

"""
BiDirectionalBeta
"""
function beta_graph(alpha::Vector, beta::Vector; seed::Integer=-1)
    n = length(beta)
    @assert length(alpha) == n "alpha and beta must be the same length"
    coefM = alpha' .+ beta
    g = SimpleDiGraph(n)
    for i in 1:n
        for j in 1:n
            expcoef_ij = exp(coefM[i,j])
            if rand() < (expcoef_ij / (1 + expcoef_ij))
                add_edge!(g, i, j)
            end
        end
    end
    return g
end

"""
UniDirectionalSparseBeta
"""
function sparse_beta_graph(beta::Vector, mu::Real; seed::Integer=-1)
    @assert all(beta .>= 0) "beta must be non-negative"
    @assert 0 in beta "beta must contain at least one zero"
    n = length(beta)
    coefM = UpperTriangular(beta' .+ beta)
    g = SimpleGraph(n)
    for i in 1:n
        for j in i:n
            expcoef_ij = exp(mu + coefM[i,j])
            if rand() < (expcoef_ij / (1 + expcoef_ij))
                add_edge!(g, i, j)
            end
        end
    end
    return g
end


using FractionalCalculus
using Test

@testset "FractionalCalculus" begin
    include("./Derivative.jl")
    include("./Integral.jl")
end


######### CmdStan batch program example  ###########

using CmdStan, StatsPlots

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  eightschools ="
  data {
    int<lower=0> J; // number of schools
    real y[J]; // estimated treatment effects
    real<lower=0> sigma[J]; // s.e. of effect estimates
  }
  parameters {
    real mu;
    real<lower=0> tau;
    real eta[J];
  }
  transformed parameters {
    real theta[J];
    for (j in 1:J)
      theta[j] <- mu + tau * eta[j];
  }
  model {
    eta ~ normal(0, 1);
    y ~ normal(theta, sigma);
  }
  "

  schools8data = Dict("J" => 8,
      "y" => [28,  8, -3,  7, -1,  1, 18, 12],
      "sigma" => [15, 10, 16, 11,  9, 11, 10, 18],
      "tau" => 25
    )

  global stanmodel, rc, chns, cnames
  stanmodel = Stanmodel(name="schools8", model=eightschools,
    output_format=:mcmcchains);
  rc, chn, cnames = stan(stanmodel, schools8data, ProjDir, CmdStanDir=CMDSTAN_HOME)

  if rc == 0

    chns = set_section(chn, Dict(
      :parameters => ["mu", "tau"],
      :thetas => ["theta.$i" for i in 1:8],
      :etas => ["eta.$i" for i in 1:8],
      :internals => ["lp__", "accept_stat__", "stepsize__", "treedepth__", "n_leapfrog__",
        "divergent__", "energy__"]
      )
    )

    if isdefined(Main, :StatsPlots)
      p1 = plot(chns)
      savefig(p1, "traceplot.pdf")
      #p2 = plot(chns, [:thetas])
      #savefig(p2, "thetas.pdf")
    end

    show(chns)
    println("\n")
    summarize(chns, sections=[:thetas])

  end
end # cd


using LowRankModels
using FactCheck

TOL = 1e-3

facts("Classification Losses") do

  context("logistic loss") do
    l = LogisticLoss()
    l1 = 1.31326168
    l0 = 0.3132616875
    # 1 is true
    # -1 and 0 are both false
    # anything else is an error
    @fact evaluate(l, 1, true) --> roughly(l0, TOL)
    @fact evaluate(l, 1, false) --> roughly(l1, TOL)
    @fact evaluate(l, -1, true) --> roughly(l1, TOL)
    @fact evaluate(l, -1, false) --> roughly(l0, TOL)
    @fact evaluate(l, 1, 1) --> roughly(l0, TOL)
    @fact evaluate(l, 1, -1) --> roughly(l1, TOL)
    @fact evaluate(l, 1, 0) --> roughly(l1, TOL)
    @fact evaluate(l, -1, 1) --> roughly(l1, TOL)
    @fact evaluate(l, -1, -1) --> roughly(l0, TOL)
    @fact evaluate(l, -1, 0) --> roughly(l0, TOL)

    @fact evaluate(3*l, 1, false) --> roughly(3*l1, TOL)
  end

  context("hinge loss") do
    l = HingeLoss()
    # 1 is true
    # -1 and 0 are both false
    # anything else is an error
    @fact evaluate(l, 1, true) --> roughly(0, TOL)
    @fact evaluate(l, 1, false) --> roughly(2, TOL)
    @fact evaluate(l, -1, true) --> roughly(2, TOL)
    @fact evaluate(l, -1, false) --> roughly(0, TOL)
    @fact evaluate(l, 1, 1) --> roughly(0, TOL)
    @fact evaluate(l, 1, -1) --> roughly(2, TOL)
    @fact evaluate(l, 1, 0) --> roughly(2, TOL)
    @fact evaluate(l, -1, 1) --> roughly(2, TOL)
    @fact evaluate(l, -1, -1) --> roughly(0, TOL)
    @fact evaluate(l, -1, 0) --> roughly(0, TOL)

    @fact evaluate(3*l, 1, false) --> roughly(3*2, TOL)

    @fact grad(l, -1, true) --> roughly(-1, TOL)
    @fact grad(l, 2, true) --> roughly(0, TOL)
    @fact grad(l, -2, false) --> roughly(0, TOL)
    @fact grad(l, 2, false) --> roughly(1, TOL)
  end

end


using MieScatter
using Test

@testset "Cross sections" begin

    # These tests are made compared to the FORTRAN77 Mie code
    # written by Karri Muinonen to the precision of the output
    # of that code.

    # Test computation of scattering, extinction and absorbsion
    # cross sections.

    # Size parameter 2, refractive index 1.33
    S, Qsca, Qext, Qback = compute_mie(2.0, 1.33, 181)
    Qabs = Qext-Qsca
    @test isapprox(Qsca, 0.71295 ; atol = 6)
    @test isapprox(Qext, 0.71295 ; atol = 6)
    @test isapprox(Qabs, 0.0 ; atol = 10)

    # Size parameter 2, refractive index 1.6
    S, Qsca, Qext, Qback = compute_mie(2.0, 1.6, 181)
    Qabs = Qext-Qsca
    @test isapprox(Qsca, 2.4834 ; atol = 5)
    @test isapprox(Qext, 2.4834 ; atol = 5)
    @test isapprox(Qabs, 0.0 ; atol = 10)

    # Size parameter 20, refractive index 1.33
    S, Qsca, Qext, Qback = compute_mie(20.0, 1.33, 181)
    Qabs = Qext-Qsca
    @test isapprox(Qsca, 0.21401E+01 ; atol = 6)
    @test isapprox(Qext, 0.21401E+01 ; atol = 6)
    @test isapprox(Qabs, 0.66613E-14 ; atol = 6)

    # Size parameter 20, refractive index 1.6
    S, Qsca, Qext, Qback = compute_mie(20.0, 1.6, 181)
    Qabs = Qext-Qsca
    @test isapprox(Qsca, 0.26115E+01 ; atol = 6)
    @test isapprox(Qext, 0.26115E+01 ; atol = 6)
    @test isapprox(Qabs, 0.39968E-14 ; atol = 6)

end # cross section testset

@testset "Return types and NaN" begin
    S, Qsca, Qext, Qback = compute_mie(20.0, 1.6, [pi/2])
    @test typeof(S) == Array{Float64,2}
    @test size(S) == (1,4)
    @test typeof(Qsca) == Float64
    @test typeof(Qext) == Float64
    @test typeof(Qback) == Float64
    @test isnan(Qext)
    @test isnan(Qback)
    @test !isnan(Qsca)
end


using Plots
gr()

using LaTeXStrings

using Unitful
using Unitful: N, mm, inch, m, cm, MPa
Unitful.preferunits(u"mm, s"...)

using DataFrames

using LinearAlgebra

θ = LinRange(0, 2π, 100)
y = sin.(θ)

p1 = plot(θ, y,
	title = L"sin(\theta)",
	xlabel = L"\theta",
	ylabel = L"Amplitude",
	xlim = [0, 2π],
	ylim = [-1, 1],
	legend = false
);

display(p1)

p2 = scatter(θ, y,
	alpha = 0.5,
	color = :jet1,
	marker_z = y,
	marker = :hexagon,
	markersize = 5,
	legend = false
);
display(p2)

F = [1000, 3200, -4000, 4300, -2750, 6630]N
r = [300mm, 400mm, 10inch, 15cm, 0.188m, 27cm] .|> mm

df = DataFrame(r=r, F=F)


"""
    sum_of_sines(domain, c, ω, ϕ)

Create initial condition function `u₀` and an antiderivative `U₀`.
The function is a sum of sines of amplitudes `c`, frequencies `ω` and phase-shifts `ϕ`.
"""
function sum_of_sines(domain, c, ω, ϕ)
    φ(ω, ϕ, x) = sin(ω * x - ϕ)
    Φ(ω, ϕ, x) = ω ≈ 0 ? -sin(ϕ) * x : -cos(ω * x - ϕ) / ω
    u₀(x) = sum(c .* φ.(ω, ϕ, x))
    U₀(x) = sum(c .* Φ.(ω, ϕ, x))
    u₀, U₀
end


using Statistics
include("dsp.jl")

function sdnn(nn_interval::Vector{<:Real})

    sdnn = std(nn_interval)
    return sdnn
end

function sdann_asdnn(nn_interval::Vector{<:Real}, index_array::Vector{<:Integer},
    duration::Integer, min_number_of_interval::Integer = 3)

    if length(nn_interval) < min_number_of_interval
        return nothing, nothing
    end

    segment_index = get_segment_index(index_array, duration)

    average_array::Vector{Float64} = []
    std_array::Vector{Float64} = []

    start_idx = 1
    for end_idx in segment_index
        landmark = @view nn_interval[start_idx:end_idx]

        if length(landmark) >= min_number_of_interval
            push!(average_array, mean(landmark))
            push!(std_array, std(landmark))
        end
        start_idx = end_idx
    end

    if length(average_array) >= min_number_of_interval

        sdann = std(average_array)
        asdnn = mean(std_array)
        return sdann, asdnn
    end

    return nothing, nothing
end


########################################################################
# Atoms of the Positive Simplex
########################################################################

"""
    PosSimplexAtom(n, z)

Atom in the positive simplex of dimension n.
"""
struct PosSimplexAtom <: AbstractAtom
    n::Int64
    z::Vector{Float64}
    function PosSimplexAtom(z)
        n = length(z)
        return new(n, z)
    end
end

"""
Multiply a Positive simplex atom by a linear map.
"""
Base.:(*)(M::AbstractLinearOp, a::PosSimplexAtom) = M*a.z
Base.vec(a::PosSimplexAtom) = a.z
Base.length(a::PosSimplexAtom) = a.n

########################################################################
# Atomic set for the positive simplex
########################################################################

"""
    PosSimplex(c, n, maxrank=n)

Atomic set defined by positive simplex in `n` variables.
The atomic set takes an optional parameters:

`maxrank` is the maximum dimension of the face exposed by a vector.
"""
struct PosSimplex <: AbstractAtomicSet
    c::Vector{Float64}
    n::Int64
    maxrank::Int64
    function PosSimplex(c, maxrank)
        n = length(c)
        minimum(c) > 0 || throw(DomainError(c,"c must be in non-negative orthant"))
        n ≥ maxrank ≥ 1 || throw(DomainError(maxrank,"maxrank must be ≥ 1"))
        return new(c, n, maxrank)
    end
end

function PosSimplex(c::Vector{Float64}; maxrank = length(findall(!iszero,c)) )
    index_c = findall(x -> x != 0, c)
    maxrank <= length(index_c) || throw(DomainError(maxrank,"maxrank should be less than sparstiy level of c"))
    return PosSimplex(c, maxrank)
end

PosSimplex(n::Int64; maxrank=n) = PosSimplex(ones(n), maxrank)

"""
    gauge(A::PosSimplex, x::Vector)

Gives the sum `dot(c, x)`.
"""
function gauge(A::PosSimplex, x::Vector)
    for i in eachindex(x)
        if (A.c[i] == 0 && x[i] != 0) || x[i] < 0
            return Inf
        end
    end
    return dot(A.c,x)
end
"""
    support(A::PosSimplex, z::Vector)

Gives the inf-norm of `c ⊙ z`.
"""
function support(A::PosSimplex, z::Vector)
    sup = 0.
    for i in eachindex(z)
        if A.c[i] != 0
            sup = max(sup, z[i]/A.c[i])
        end
    end

    return sup
end
"""
    expose(A::PosSimplex, z::Vector)

A non-overwriting version of [`expose!`](@ref).
"""
expose(A::PosSimplex, z; kwargs...) = expose!(A::PosSimplex, copy(z); kwargs...)

"""
    expose!(A::PosSimplex, z::Vector; tol=1e-12)

Obtain an atom in the face exposed by the vector `z`.
The vector `z` is overwritten. If `norm(z,Inf)<tol`, then
`z` is returned untouched.
"""

function expose!(A::PosSimplex, z::Vector; tol=1e-1)
    # z = Float64.(z)
    zsup = support(A, z)
    if maximum(z) < -1e-12
        return PosSimplexAtom(zero(z))
    end

    index_c = findall(x -> x != 0, A.c)
    index_c_comp = findall(x -> x == 0, A.c)

    if zsup < 1e-12
        nnz = 1
    else
        nnz = 0
    end

    # for i in eachindex(index_c_comp)
    #     z[index_c_comp[i]] = 0
    # end
    for i in eachindex(index_c)
        val = (A.c[index_c[i]]).^(-1) * z[index_c[i]]
        if abs(zsup - val) ≤ tol*zsup
            z[index_c[i]] = 1/A.c[index_c[i]]
            nnz += 1
        else
            z[index_c[i]] = zero(eltype(z))
        end
    end
    return PosSimplexAtom(z ./= nnz)
end

Base.length(A::PosSimplex) = A.n
rank(A::PosSimplex) = A.maxrank
atom_name(A::PosSimplex) = "Positive Simplex"
atom_description(A::PosSimplex) = "{ x ∈ ℝⁿ | ⟨c, x⟩ ≤ 1, x ≥ 0}"
atom_parameters(A::PosSimplex) = "c = cost vector, n = $(length(A)); maxrank = $(A.maxrank)"

########################################################################
# Face of the positive simplex.
########################################################################

struct PosSimplexFace{LM <: LinearMap} <: AbstractFace
    n::Int64
    k::Int64
    S::LM
end

PosSimplexFace(S::LinearMap) = PosSimplexFace(size(S)...,S)

"""
    face(A, z)

Return a face of the atomic set `A` exposed by the vector `z`. The dimension of
the exposed face is limited by `k=maxrank(A)`. If the dimension is being limited
by `k`, then the routine returns a set of rank `k` atoms that define the subset
of the exposed face.
"""
function face(A::PosSimplex, z::Vector; rTol=1e-1)

    t = support(A, z)
    a = expose(A, z, tol=rTol*t)
    idx = findall(!iszero, vec(a))

    k = length(idx)
    val = [(A.c[i]).^(-1) for i in idx]
    n = length(A)
    if support(A,z) < 1e-12
        S = sparse(idx, collect(1:k), val, n, k + 1)
    else
        S = sparse(idx, collect(1:k), val, n, k)
    end
    S = LinearMap(S)
    return PosSimplexFace(S)
end

"""
Given a face and a set of weights, reveal the corresponding
point on the face.
"""
Base.:(*)(F::PosSimplexFace, c::Vector) = F.S*c
Base.:(*)(M::AbstractLinearOp, F::PosSimplexFace) = M*F.S
Base.:(*)(λ::Real, F::PosSimplexFace) = λ*F.S
Base.length(F::PosSimplexFace) = F.n
rank(F::PosSimplexFace) = F.k
vec(F::PosSimplexFace) = F.S
face_name(F::PosSimplexFace) = "Face of positive simplex"
face_parameters(F::PosSimplexFace) = "rank = $(rank(F)); n = $(length(F))"


using PolynomialBasis
using Test
# using Revise
using CutCellDG
include("../useful_routines.jl")

polyorder = 1
numqp = 2

basis = LagrangeTensorProductBasis(2,1)
dim,numpts = size(interpolation_points(basis))
mesh = CutCellDG.CGMesh([0.0, 0.0], [3.0, 1.0], [3, 1], numpts)

x0 = [1.5,0.0]
normal = [1.,0.]
distancefunc(x) = plane_distance_function(x,normal,x0)
levelset = CutCellDG.LevelSet(distancefunc,mesh,basis)
CutCellDG.load_coefficients!(levelset,1)
@test allapprox(levelset.interpolater.coeffs,[-1.5,-1.5,-0.5,-0.5]')
@test levelset([0.,0.]) ≈ -1.0

CutCellDG.update_coefficients!(levelset,2,[-0.6,-0.6,0.4,0.4])
@test allapprox(CutCellDG.coefficients(levelset,2),[-0.6,-0.6,0.4,0.4])
@test allapprox(CutCellDG.coefficients(levelset,1),[-1.5,-1.5,-0.6,-0.6])
@test allapprox(CutCellDG.coefficients(levelset,3),[0.4,0.4,1.5,1.5])

basis = HermiteTensorProductBasis(2)
numqp = required_quadrature_order(3)
quad = tensor_product_quadrature(2,numqp)
dim,numpts = size(interpolation_points(basis))
mesh = CutCellDG.CGMesh([0.,0.],[1.,1.],[2,2],numpts)
x0 = [0.75,0.0]
normal = [1.0,0.0]
distancefunc(x) = plane_distance_function(x,normal,x0)

levelset = CutCellDG.LevelSet(distancefunc,mesh,basis)

coeffs1 = CutCellDG.coefficients(levelset,1)
testcoeffs = [-0.75  -0.75  -0.25  -0.25
               0.0    0.0    0.0    0.0
               0.25   0.25   0.25   0.25
               0.0    0.0    0.0    0.0]
testcoeffs = vec(testcoeffs)
@test allapprox(coeffs1,testcoeffs,1e3eps())

x0 = [1.0,0.5]
normal = [1.0,1.0]/sqrt(2)

function testcubic(v)
    x,y = v
    return 8x^3*y^3 + 4x^2*y^3 + 17*x*y + 11
end

levelset = CutCellDG.LevelSet(testcubic,mesh,basis)
testcellid = 4
CutCellDG.load_coefficients!(levelset,testcellid)
cellmap = CutCellDG.cell_map(mesh,testcellid)
xrange = range(-1,stop=1,length=5)
testp = ImplicitDomainQuadrature.tensor_product_points(xrange',xrange')
spatialtestp = cellmap(testp)

vals = vec(mapslices(levelset,testp,dims=1))
testvals = vec(mapslices(testcubic,spatialtestp,dims=1))
@test allapprox(vals,testvals,1e4eps())


@testset "no linesearch" begin

function f_nolin!(fvec, x)
    fvec[1] = (x[1]+3)*(x[2]^3-7)+18
    fvec[2] = sin(x[2]*exp(x[1])-1)
end

r = nlsolve(f_nolin!, [ -0.5; 1.4], autodiff = :forward, method = :newton, linesearch = LineSearches.Static())
@test converged(r)

end


"""
    Basis

Abstract supertype of various basis sets that vector spherical harmonics may be decomposed in.
"""
abstract type Basis end
@doc raw"""
    SphericalCovariant <: Basis

The spherical covariant basis ``\chi_\mu`` for ``\mu\in\{-1,0,1\}``
"""
struct SphericalCovariant <: Basis end
@doc raw"""
    Polar <: Basis

The spherical polar basis ``\hat{r}``, ``\hat{\theta}``, ``\hat{\phi}``.
"""
struct Polar <: Basis end
@doc raw"""
    HelicityCovariant  <: Basis

The helicity basis ``\mathbf{e}_\mu`` for ``\mu\in\{-1,0,1\}``
"""
struct HelicityCovariant <: Basis end
@doc raw"""
    Cartesian  <: Basis

The Cartesian basis ``\hat{x}``, ``\hat{y}``, ``\hat{z}``
"""
struct Cartesian <: Basis end

Base.broadcastable(B::Basis) = Ref(B)

_basisinds(::Union{HelicityCovariant, SphericalCovariant}) = -1:1
_basisinds(::Union{Cartesian, Polar}) = 1:3

#= Define matrices that convert between components C1 and C2 in bases B1 and B2 as
    M * C1 = C2

The corresponding functions are named B1_B2_conversion.
These matrices are conjugates of those that convert between the bases, so conj.(M) * B1 = B2
=#
const HelicitySphericalPolarConversionMatrix = SMatrix{3,3}([
                                        0       1   0
                                        1/√2    0   -1/√2
                                        -im/√2   0    -im/√2
                                        ])

helicity_polar_conversion(θ, ϕ) = HelicitySphericalPolarConversionMatrix
polar_helicity_conversion(θ, ϕ) = helicity_polar_conversion(θ, ϕ)'

function helicity_spherical_conversion(θ, ϕ)
    cisϕ = cis(ϕ)
    sinθ, cosθ = sincos(θ)
    invsqrt2 = 1/√2
    sin²θby2 = (1 - cosθ)/2
    cos²θby2 = (1 + cosθ)/2
    SMatrix{3,3}((
        cos²θby2 * cisϕ,
        -sinθ * invsqrt2,
        sin²θby2 * conj(cisϕ),

        sinθ * invsqrt2 * cisϕ,
        cosθ,
        -sinθ * invsqrt2 * conj(cisϕ),

        sin²θby2 * cisϕ,
        sinθ * invsqrt2,
        cos²θby2 * conj(cisϕ),
        ))
end
spherical_helicity_conversion(θ, ϕ) = helicity_spherical_conversion(θ, ϕ)'

function helicity_cartesian_conversion(θ, ϕ)
    sinθ, cosθ = sincos(θ)
    sinϕ, cosϕ = sincos(ϕ)
    invsqrt2 = 1/√2

    SMatrix{3,3}((
        (cosθ * cosϕ + im * sinϕ) * invsqrt2,
        (cosθ * sinϕ - im * cosϕ) * invsqrt2,
        -sinθ * invsqrt2,

        sinθ * cosϕ,
        sinθ * sinϕ,
        cosθ,

        -(cosθ * cosϕ - im * sinϕ) * invsqrt2,
        -(cosθ * sinϕ + im * cosϕ) * invsqrt2,
        sinθ * invsqrt2,
        ))
end
cartesian_helicity_conversion(θ, ϕ) = helicity_cartesian_conversion(θ, ϕ)'

function spherical_polar_conversion(θ, ϕ)
    invsqrt2 = 1/√2
    invsqrt2cisϕ = invsqrt2 * cis(ϕ)
    sinθ, cosθ = sincos(θ)
    normsinθcisϕ = sinθ * invsqrt2cisϕ
    normcosθcisϕ = cosθ * invsqrt2cisϕ
    SMatrix{3,3}((
        conj(normsinθcisϕ),
        conj(normcosθcisϕ),
        -im * conj(invsqrt2cisϕ),

        cosθ,
        -sinθ,
        0,

        -normsinθcisϕ,
        -normcosθcisϕ,
        -im * invsqrt2cisϕ,
        ))
end
polar_spherical_conversion(θ, ϕ) = spherical_polar_conversion(θ, ϕ)'

const SphericalCartesianConversionMatrix = SMatrix{3,3}((
        1/√2,
        -im/√2,
        0,

        0,
        0,
        1,

        -1/√2,
        -im/√2,
        0,
        ))

spherical_cartesian_conversion(θ, ϕ) = SphericalCartesianConversionMatrix
cartesian_spherical_conversion(θ, ϕ) = spherical_cartesian_conversion(θ, ϕ)'

function cartesian_polar_conversion(θ, ϕ)
    sinθ, cosθ = sincos(θ)
    sinϕ, cosϕ = sincos(ϕ)
    SMatrix{3,3}((
        sinθ*cosϕ,
        cosθ*cosϕ,
        -sinϕ,

        sinθ*sinϕ,
        cosθ*sinϕ,
        cosϕ,

        cosθ,
        -sinθ,
        0,
        ))
end
polar_cartesian_conversion(θ, ϕ) = cartesian_polar_conversion(θ, ϕ)'

"""
    basisconversionmatrix(B1, B2, θ, ϕ)

Return the matrix that converts the components of a vector from the basis `B1` to the basis `B2`
at the point `(θ,ϕ)`.

!!! note
    For the complex bases `SphericalCovariant` and `HelicityCovariant`, the matrix transforms the
    contravariant components between bases.

# Examples
```jldoctest
julia> v = [1,0,0] # x
3-element Vector{Int64}:
 1
 0
 0

julia> θ, ϕ = pi/2, 0
(1.5707963267948966, 0)

julia> M = VectorSphericalHarmonics.basisconversionmatrix(Cartesian(), Polar(), θ, ϕ)
3×3 StaticArrays.SMatrix{3, 3, Float64, 9} with indices SOneTo(3)×SOneTo(3):
  1.0          0.0   6.12323e-17
  6.12323e-17  0.0  -1.0
 -0.0          1.0   0.0

julia> M * v # polar coordinates along r, θ, ϕ
3-element StaticArrays.SVector{3, Float64} with indices SOneTo(3):
 1.0
 6.123233995736766e-17
 0.0
```
"""
basisconversionmatrix(::T, ::T, θ, ϕ) where {T<:Basis} = I
basisconversionmatrix(::HelicityCovariant, ::SphericalCovariant, θ, ϕ) = helicity_spherical_conversion(θ, ϕ)
basisconversionmatrix(::SphericalCovariant, ::HelicityCovariant, θ, ϕ) = spherical_helicity_conversion(θ, ϕ)
basisconversionmatrix(::HelicityCovariant, ::Polar, θ, ϕ) = helicity_polar_conversion(θ, ϕ)
basisconversionmatrix(::Polar, ::HelicityCovariant, θ, ϕ) = polar_helicity_conversion(θ, ϕ)
basisconversionmatrix(::HelicityCovariant, ::Cartesian, θ, ϕ) = helicity_cartesian_conversion(θ, ϕ)
basisconversionmatrix(::Cartesian, ::HelicityCovariant, θ, ϕ) = cartesian_helicity_conversion(θ, ϕ)

basisconversionmatrix(::SphericalCovariant, ::Polar, θ, ϕ) = spherical_polar_conversion(θ, ϕ)
basisconversionmatrix(::Polar, ::SphericalCovariant, θ, ϕ) = polar_spherical_conversion(θ, ϕ)
basisconversionmatrix(::SphericalCovariant, ::Cartesian, θ, ϕ) = spherical_cartesian_conversion(θ, ϕ)
basisconversionmatrix(::Cartesian, ::SphericalCovariant, θ, ϕ) = cartesian_spherical_conversion(θ, ϕ)

basisconversionmatrix(::Cartesian, ::Polar, θ, ϕ) = cartesian_polar_conversion(θ, ϕ)
basisconversionmatrix(::Polar, ::Cartesian, θ, ϕ) = polar_cartesian_conversion(θ, ϕ)


# Numerically stable expit
@inline function expit(t::T) where T
    if t >= 0.0
        return one(T) / (one(T) + exp(-t))
    else
        e_t = exp(t)
        return e_t / (one(T) + e_t)
    end
end

# A numerically robust function to evaluate -log(1 + exp(-t))
# See: http://fa.bianp.net/blog/2019/evaluate_logistic/
@inline function log1pexp(t::T)::T where T
    if t < -33.3
        return t
    elseif t <= -18.0
        return t - exp(t)
    elseif t <= 37.0
        return -log1p(exp(-t))
    else
        return -exp(-t)
    end
end

"Compute Fenchel transform of f(x) = log(1 + exp(-x))."
function logloss(λ::T) where T <: Real
    # TODO: numerically robust version
    if λ == 0.0 || λ == -1.0
        return 0.0
    elseif -1.0 < λ < 0.0
        return (1.0 + λ) * log(1.0 + λ) - λ * log(-λ)
    else
        return Inf
    end
end

abstract type AbstractScaler end

"Scale dataset before optimization."
scale!(::AbstractScaler, X::AbstractArray{T, 2}) where T = nothing

"""
Scale each feature in dataset by standard deviation.

```math
μ_j = mean(x_j)
σ_j = std(x_j)
x_j^+ =  (x_j - μ_j) / σ_j

```

"""
struct NormalScaler <: AbstractScaler end

function scale!(::NormalScaler, X::Array{T, 2}) where T
    n, d = size(X)
    μ = mean(X, dims=1)
    σ = std(X, dims=1)

    @inbounds for i in 1:d
        X[:, i] .= (X[:, i] .- μ[i]) ./ σ[i]
    end
end

function scale!(::NormalScaler, X::AbstractSparseMatrix{T, Int}) where T
    n, d = size(X)
    σ = std(X, dims=1)
    @inbounds for i in 1:d
        X[:, i] ./= σ[i]
    end
end

"""
Format labels `y = [y_1, ..., y_n]` to ensure that `y_i ∈ { -1, 1}`
for all index i.

# Examples

```julia
    y = [-2.0, 1.0, -1.0, 2.0]
    format_label!(y)

```
"""
function format_label!(y::AbstractVector)
    set_reference = [-1.0, 1.0]
    elts = unique(y)
    if union(elts, set_reference) != set_reference
        intersets = intersect(elts, set_reference)
        replacements = setdiff(set_reference, intersets)
        # For all elements not in reference set
        count = 1
        for elt in setdiff(elts, set_reference)
            y[y .== elt] .= replacements[count]
            count += 1
        end
    end
    return nothing
end

# Return upper index (ordered by rows) of a pxp triangular matrix
function triul(p)
    index = Int[]
    for i in 1:p
        for j in i:p
            push!(index, i + (j - 1) * p)
        end
    end
    return index
end


@taylorize function lotka_volterra!(du, u, p, t)
    u1u2 = u[1] * u[2]
    du[1] = 3.0 * (u[1] - u1u2)
    du[2] = u1u2 - u[2]
    return du
end

function lotka_volterra_hybrid(; nsplit=1,
                                 ε = 0.008,
                                 ε_ext=1e-4, # threshold for the outer approximation
                                 n_int=50)   # number of directions for the inner approximation

    # generate external / internal polytopic approximations of the guard
    B = Ball2([1.0, 1.0], 0.15) # "exact"
    B_ext = overapproximate(B, ε_ext) # outer approximation
    B_int = underapproximate(B, PolarDirections(n_int)) # inner approximation
    B_int = tohrep(convert(VPolygon, B_int)) # cast to Hrep
    B_intᶜ = complement(B_int)

    # define modes
    aut = LightAutomaton(3)
    outside = @system(x' = lotka_volterra!(x), dim: 2, x ∈ B_intᶜ)
    inside = @system(x' = lotka_volterra!(x), dim: 2, x ∈ B_ext)
    outside_unconstrained = @system(x' = lotka_volterra!(x), dim: 2, x ∈ Universe(2))

    # define the transition graph
    add_transition!(aut, 1, 2, 1)
    add_transition!(aut, 2, 3, 2)
    T_out_in = @map(x -> x, dim:2, x ∈ B_ext)
    T_in_out = @map(x -> x, dim:2, x ∈ B_intᶜ)

    # initial-value problem
    H = HybridSystem(automaton=aut, modes=[outside, inside, outside_unconstrained],
                                           resetmaps=[T_out_in, T_in_out])

    # initial states with splitting
    X0 = Hyperrectangle(low=[1.3-ε, 1.], high=[1.3+ε, 1.])
    X0s = split(X0, [nsplit, 1])
    X0st = [(X0s_i, 1) for X0s_i in X0s]
    return InitialValueProblem(H, X0st)
end

#=
# TEST RUN:

prob = lotka_volterra_hybrid(nsplit=10, ε_ext=1e-4, n_int = 30, ε = 0.008);
@time sol = solve(prob,
                  tspan=(0.0, 3.64),
                  alg=TMJets(abstol=1e-12, orderT=7, orderQ=1, adaptive=true, disjointness=RA.ZonotopeEnclosure()),
                  max_jumps=2,
                  intersect_source_invariant=false,
                  intersection_method=RA.BoxIntersection(),
                  clustering_method=RA.BoxClustering(),
                  disjointness_method=RA.BoxEnclosure());
solz = overapproximate(sol, Zonotope);

B = Ball2([1.0, 1.0], 0.15) # "exact"
B_ext = overapproximate(B, 1e-6) # outer approximation
plot(solz, vars=(1, 2))
plot!(B_ext, ratio=1)
=#


using DelimitedFiles

# This is a function to simulate a "chaos game"
function chaos_game(n::Int, shape_points)

    # Initializing the output array and the initial point
    output_points = zeros(n,2)
    point = [rand(), rand()]

    for i = 1:n
        output_points[i,:] .= point
        point = 0.5*(rand(shape_points) .+ point)
    end

    return output_points

end

# This will generate a Sierpinski triangle with a chaos game of n points for an
# initial triangle with three points on the vertices of an equilateral triangle:
#     A = (0.0, 0.0)
#     B = (0.5, sqrt(0.75))
#     C = (1.0, 0.0)
# It will output the file sierpinski.dat, which can be plotted after
shape_points = [[0.0, 0.0],
                [0.5, sqrt(0.75)],
                [1.0, 0.0]]
output_points = chaos_game(10000, shape_points)
writedlm("sierpinski.dat", output_points)


# Endogenous Grid Points with IID Income
# Greg Kaplan 2017
# Translated by Tom Sweeney Dec 2020

using Random, Interpolations, NLsolve, Plots
include("discrete_normal.jl")
include("lininterp1.jl")

# PARAMETERS

## preferences
risk_aver = 2
beta = 0.95

## returns
r = 0.03
R = 1+r

## income risk: discretized N(mu,sigma^2)
mu_y = 1
sd_y = 0.2
ny = 5

## asset grids
na = 50
amax = 50
borrow_lim = 0
agrid_par = 0.5 # 1 for linear, 0 for L-shaped

## computation
max_iter = 1000
tol_iter = 1.0e-6
Nsim = 50000
Tsim = 500

## mpc options
mpcamount1 = 1.0e-10 # approximate thoeretical mpc
mpcamount2 = 0.10 # one percent of average income: approx $500

# OPTIONS
Display = 1
DoSimulate = 1
MakePlots = 1
ComputeMPC = 1

# DRAW RANDOM NUMBERS
Random.seed!(2020)
yrand = rand(Nsim,Tsim)

# SET UP GRIDS

## assets
agrid = range(0,1,length=na)
agrid = agrid.^(1 ./ agrid_par)
agrid = borrow_lim .+ (amax.-borrow_lim).*agrid

## income: disretize normal distribution
width = nlsolve(x -> discrete_normal(ny,mu_y,sd_y,x...)[1],[2.0]).zero
temp, ygrid, ydist = discrete_normal(ny,mu_y,sd_y,width...)
ycumdist = cumsum(ydist)

# UTILITY FUNCTION

u(c) = (c.^(1-risk_aver)-1)./(1-risk_aver)
u1(c) = c.^(-risk_aver)
u1inv(u) = u.^(-1 ./risk_aver)

# INITIALIZE CONSUMPTION FUNCTION

conguess = zeros(na,ny)
for iy = 1:ny
    conguess[:,iy] = r.*agrid.+ygrid[iy]
end

# ITERATE ON EULER EQUATION WITH ENDOGENOUS GRID POINTS

con = copy(conguess)

iter = 0
cdiff = 1000

while iter <= max_iter && cdiff>tol_iter
    iter = iter + 1
    global sav = zeros(na,ny)

    conlast = copy(con)

    emuc = u1(conlast)*ydist
    muc1 = beta.*R.*emuc
    con1 = u1inv(muc1)

    ## loop over income
    ass1 = zeros(na,ny)
    for iy = 1:ny

        ass1[:,iy] = (con1 .+ agrid .-ygrid[iy])./R

        ## loop over current period ssets
        for ia  = 1:na
            if agrid[ia]<ass1[1,iy] # borrowing constraint binds
                sav[ia,iy] = borrow_lim
            else # borrowing constraint does not bind;
                sav[ia,iy] = lininterp1(ass1[:,iy],agrid,agrid[ia])
            end
        end
        con[:,iy] = R.*agrid .+ ygrid[iy] - sav[:,iy]
    end

    cdiff = maximum(abs.(con - conlast))
    if Display>=1
        println("Iteration no. " * string(iter), " max con fn diff is " * string(cdiff))
    end
end

# SIMULATE
if DoSimulate==1

    yindsim = zeros(Int,Nsim,Tsim)
    asim = zeros(Nsim,Tsim)

    ## create interpolating function
    savinterp = Array{Any}(undef,ny)
    for iy = 1:ny
        savinterp[iy] = interpolate((agrid,), sav[:,iy], Gridded(Linear()))
    end

    ## loop over time periods
    for it = 1:Tsim
        if Display>=1 && mod(it,100)==0
            println("Simulating, time period " * string(it))
        end

        ## income realization: note we vectorize simulations at once because
        ## of matlab, in other languages we would loop over individuals
        yindsim[yrand[:,it].<=ycumdist[1],it] .= 1
        for iy = 2:ny
            yindsim[(yrand[:,it].> ycumdist[iy-1]) .& (yrand[:,it].<=ycumdist[iy]),it] .= iy
        end

        ## asset choice
        if it<Tsim
            for iy = 1:ny
                asim[yindsim[:,it].==iy,it+1] = savinterp[iy](asim[yindsim[:,it].==iy,it])
            end
        end
    end

    ## assign actual income values
    ysim = ygrid[yindsim]

end

# MAKE PLOTS
if MakePlots==1

    ## consumption policy function
    p1 = plot(agrid, [con[:,1] con[:,ny]], xlims=(0,amax), title="Consumption", color=[:blue :red], label=["Lowest income state" "Highest income state"])
    display(p1)

    ## savings policy function
    p2 = plot(agrid, [sav[:,1].-agrid[:,1] sav[:,ny].-agrid[:,1]], xlims=(0,amax), title="Savings", color=[:blue :red], legend=false)
    plot!(agrid, zeros(na,1), color=:black, lw=0.5)
    display(p2)

    ## nice zoom
    xlimits = (0,1)
    xlimind = trues(na)
    if minimum(agrid) < xlimits[1]
        xlimind = xlimind .& (agrid.>=maximum(agrid[agrid<xlimits[1]]))
    elseif minimum(agrid) > xlimits[2]
        xlimind .= 0
    end
    if maximum(agrid) > xlimits[2]
        xlimind = xlimind .& (agrid.<=minimum(agrid[agrid.>xlimits[2]]))
    elseif maximum(agrid) < xlimits[1]
        xlimind .= 0
    end

    ## consumption policy function: zoomed in
    p3 = plot(agrid[xlimind], [con[xlimind,1] con[xlimind,ny]], xlims=xlimits, title="Consumption: Zoomed", marker=:circle, color=[:blue :red], linewidth=2, legend=false)
    display(p3)

    ## savings policy function: zoomed in
    p4 = plot(agrid[xlimind], [sav[xlimind,1].-agrid[xlimind] sav[xlimind,ny].-agrid[xlimind]], xlims=xlimits, title="Savings: Zoomed (a'-a)", marker=:circle, color=[:blue :red], linewidth=2, legend=false)
    plot!(agrid, zeros(na,1), color=:black, lw=0.5)
    display(p4)

    ## income distribution
    p5 = histogram(ysim[:,Tsim], bins=[2*ygrid[1]-ygrid[2];ygrid].+(ygrid[2]-ygrid[1])/2, title="Income distribution", color=RGB(0,0.5,0.5), linecolor=:blue, legend=false)
    display(p5)

    ## asset distribution
    p6 = histogram(asim[:,Tsim], nbins=100, title="Asset distribution", color=RGB(.7,.7,.7), linecolor=:black, legend=false)
    display(p6)

    ## convergence check
    p7 = plot(1:Tsim, mean(asim,dims=1)', title="Mean Asset Convergence", xlabel="Time Period", color=:black, lw=1.5, legend=false)
    display(p7)

    ## asset distribution statistics
    aysim = asim[:,Tsim]./mean(ysim[:,Tsim])
    println("Mean assets (relative to mean income): " * string(mean(aysim)))
    println("Fraction borrowing constrained: " * string(sum(aysim.==borrow_lim)/Nsim * 100) * '%')
    println("10th Percentile: " * string(quantile(aysim,.1)))
    println("50th Percentile: " * string(quantile(aysim,.5)))
    println("90th Percentile: " * string(quantile(aysim,.9)))
    println("99th Percentile: " * string(quantile(aysim,.99)))
end

# COMPUTE MPCs
if ComputeMPC==1

    ## theoretical mpc lower bound
    mpclim = R*((beta*R)^(-1/risk_aver))-1

    coninterp = Array{Any}(undef,ny)
    mpc1 = zeros(na,ny)
    mpc2 = zeros(na,ny)

    for iy = 1:ny
        ## create interpolating function
        coninterp[iy] = extrapolate(interpolate((agrid,), con[:,iy], Gridded(Linear())), Line())

        mpc1[:,iy] = ( coninterp[iy](agrid.+mpcamount1) - con[:,iy] ) ./ mpcamount1
        mpc2[:,iy] = ( coninterp[iy](agrid.+mpcamount2) - con[:,iy] ) ./ mpcamount2

    end

    ## mpc functions
    p8 = plot([agrid agrid],[mpc1[:,1] mpc2[:,1]], xlims=(0,10), linestyle=[:solid :dash], color=:blue, label=["Lowest income state: amount 1" "Lowest income state: amount 2"])
    plot!([agrid agrid],[mpc1[:,ny] mpc2[:,ny]], linestyle=[:solid :dash], color=:red, label=["Highest income state: amount 1" "Highest income state: amount 2"])
    plot!(agrid,mpclim.*ones(size(agrid)),color=:black,linestyle=:dot,linewidth=2,label="Theoretical MPC limit = " * string(mpclim))
    display(p8)

    ## mpc distribution
    mpc1sim = zeros(Nsim,1)
    mpc2sim = zeros(Nsim,1)
    for iy = 1:ny
        mpc1sim[yindsim[:,Tsim].==iy] = ( coninterp[iy](asim[yindsim[:,Tsim].==iy,Tsim].+mpcamount1) - coninterp[iy](asim[yindsim[:,Tsim].==iy,Tsim]) ) ./ mpcamount1
        mpc2sim[yindsim[:,Tsim].==iy] = ( coninterp[iy](asim[yindsim[:,Tsim].==iy,Tsim].+mpcamount2) - coninterp[iy](asim[yindsim[:,Tsim].==iy,Tsim]) ) ./ mpcamount2
    end

    p9 = histogram(mpc1sim,bins=0:0.02:1.5, title="MPC distribution",color=RGB(.7,.7,.7),legend=false)
    display(p9)

    # mpc distribution statistics
    println("Mean MPC amount 1: " * string(mean(mpc1sim)))
    println("Mean MPC amount 2: " * string(mean(mpc2sim)))
end


# ---
# title: 793. Preimage Size of Factorial Zeroes Function
# id: problem793
# author: Indigo
# date: 2021-01-29
# difficulty: Hard
# categories: Binary Search
# link: <https://leetcode.com/problems/preimage-size-of-factorial-zeroes-function/description/>
# hidden: true
# ---
#
# Let `f(x)` be the number of zeroes at the end of `x!`. (Recall that `x! = 1 *
# 2 * 3 * ... * x`, and by convention, `0! = 1`.)
#
# For example, `f(3) = 0` because 3! = 6 has no zeroes at the end, while `f(11)
# = 2` because 11! = 39916800 has 2 zeroes at the end. Given `K`, find how many
# non-negative integers `x` have the property that `f(x) = K`.
#
#
#
#     **Example 1:**
#     Input: K = 0
#     Output: 5
#     Explanation: 0!, 1!, 2!, 3!, and 4! end with K = 0 zeroes.
#
#     **Example 2:**
#     Input: K = 5
#     Output: 0
#     Explanation: There is no x such that x! ends in K = 5 zeroes.
#
#
# **Note:**
#
#   * `K` will be an integer in the range `[0, 10^9]`.
#
#
## @lc code=start
using LeetCode

function preimage_size_fzf(K::Int)
    zeta(x) = x > 0 ? x ÷ 5 + zeta(x ÷ 5) : 0
    lo, hi = K, 5 * K + 1
    while lo < hi
        mid = (lo + hi) ÷ 2
        z_mid = zeta(mid)
        if z_mid == K
            return 5
        elseif z_mid < K
            lo = mid + 1
        else
            hi = mid
        end
    end
    return 0
end
## @lc code=end


using AbstractLogic
using Test
using Suppressor

@test LogicalCombo() |> nfeasible == 0

logicset = @suppress logicalparse("a,b,c,d,e in 1:2")
@test logicset[1,1] == 1
@test logicset[4^3,3] == 2
@test logicset[:,:a] == logicset[:,1]

logicset = @suppress logicalparse("a.1,a.2,a.3,b.1,b.2,b.3 in 1:3 || {{j}}.1 != {{j}}.2, {{j}}.3 &&& {{j}}.2 != {{j}}.3")
@test nfeasible(logicset) == 36
@test @suppress nfeasible(logicalparse("a,b,c in 1:3 || {{i}}!={{>i}}")) == nfeasible(logicalparse("a,b,c in unique"))

logicset = @suppress logicalparse("a,b,c in 1:3 || {{i}}!={{!i}}")
@test AbstractLogic.operatoreval("a^=b",logicset) |> nfeasible == 0
@test AbstractLogic.operatoreval("a|=b",logicset) |> nfeasible == 0
@test AbstractLogic.operatoreval("a!|=b",logicset) |> nfeasible == 6
@test AbstractLogic.operatoreval("a<=b",logicset) |> nfeasible == 3
@test AbstractLogic.operatoreval("a<<b",logicset) |> nfeasible == 1
@test AbstractLogic.operatoreval("a>=b",logicset) |> nfeasible == 3
@test AbstractLogic.operatoreval("a>>b",logicset) |> nfeasible == 1
@test AbstractLogic.operatoreval("a % 3 ==0",logicset) |> nfeasible == 2

@test @suppress logicalparse("{{<3|>3}}=1", logicset) |> nfeasible  == 0

@test @suppress search("{{i+4}}=1", logicset)[1] === missing

@test @suppress logicalparse("a,b ∈ 0:1; a==a === b==b; a == b !=> b != a") |> nfeasible == 4

x = @suppress logicalparse("a,b in 1:5; true &&& true ||| false ==> true <== false ||| true ^^^ false")
@test nfeasible(x) == 25

x = @suppress logicalparse("a,b in 1:5; true &&&& true |||| false ===> true <=== false |||| true ^^^^ false")
@test nfeasible(x) == 25

logicset = @suppress logicalparse("a, b in 1; a != b")
@test @suppress( search("=1", logicset)) === missing
@test  @suppress(checkfeasible("a=1", logicalparse("a, b in 1; a != b")) === missing)

@test AbstractLogic.definelogicalset("a ∈ 1") |> nfeasible == 1
@test AbstractLogic.definelogicalset("a,b,c ∈ 1,2,3 || unique") |> nfeasible == 6

@test expand(LogicalCombo(), [:a => 1:2, :b => 1:3]) |> nfeasible == 6

logicset = @suppress logicalparse("a, b, c in 1:3; a > b")
@test @suppress !dependenton("a ⊥ b", logicset)
@test @suppress dependenton("a !⊥ b", logicset)

@test @suppress !dependenton("a ⊥ b", logicset)
@test @suppress dependenton("a !⊥ b", logicset)
@test @suppress dependenton("a ⊥ c", logicset)

@test @suppress !setcompare("a ⊂ b", logicset)
@test @suppress setcompare("a !⊂ b", logicset)
@test @suppress setcompare("c ⊃ b", logicset)
@test @suppress setcompare("b !⊃ c", logicset)
@test @suppress setcompare("a,b ⋂ b,c", logicset)
@test @suppress !setcompare("a !⋂ c", logicset)

@test logicset |> showfeasible == logicset[:,:,true]

@test @suppress(logicalparse("a,b,c,d,e in 1:2; {{i}}!={{i+1}}") |> nfeasible) == 2
@test @suppress(logicalparse("a,b,c,d,e in 1:2; {{i}}!={{i-1}}") |> nfeasible) == 2
@test @suppress(logicalparse("a,b in 1:2; {{i}}!={{>i}}") |> nfeasible) == 2

@test @capture_out(help()) |> length > 500
@test @capture_out(help("")) |> length > 500
@test @capture_out(help("in")) |> length > 250
@test @capture_out(help("back")) |> length > 50
@test @capture_out(help("!=")) |> length > 50
@test @capture_out(help("!==")) |> length > 430
@test @capture_out(help("!===")) |> length > 440
@test @capture_out(help("{{i}}")) |> length > 370

@test @suppress(checkfeasible("{{i}}=1", logicalparse("a,b in 1:2")))[1] == .25
@test @capture_out(logicalparse("a,b in 1:2; error"))[(end-16):end] == "not interpreted!\n"

x = @suppress logicalparse("a,b in 1:2; a==b")
@test @suppress(AbstractLogic.back(x) |> nfeasible) == 4

@test @suppress nfeasible(logicalparse(["a,b in 1:2","a==b"])) == nfeasible(x)

@test @suppress(logicalparse("a, b ∈ 0:1; a | b ||| a & b")) |> nfeasible == 3


function rgammanm(z, n, m=0)
    if n==m
        return one(z)
    elseif n < m
        return one(z) / rgammanm(z, m, n)
    end

    rg = one(z)
    for i = (n-1):-1:m
        rg = rg * (z + i)
    end

    return rg
end

function rgammaxy(x, y, n)
    rg = x/y * gamma(x)/gamma(y)

    for i = 1:(n-1)
        rg = rg * (x+i)/(y+i)
    end
    return rg
end

function rgammaxnym(x, n, y, m)

    if x == y
        return rgammanm(x, n, m)
    elseif n == m
        return rgammaxy(x, y, n)
    elseif n < m
        return one(x)/rgammaxnym(y, m, x, n)
    end

    rg = x/y * gamma(x)/gamma(y)

    for i = (m-1):-1:1
        rg = rg * (x+i)/(y+i)
    end

    for i = (n-1):-1:m
        rg = rg * (x+i)
    end

    return rg
end


using Distributions, Plots; pyplot()

include("../data/mvParams.jl")
biNorm = MvNormal(meanVect,covMat)

N = 10^3
points = rand(MvNormal(meanVect,covMat),N)

support = 15:0.5:40
z = [ pdf(biNorm,[x,y]) for y in support, x in support ]

p1 = scatter(points[1,:], points[2,:], ms=0.5, c=:black, legend=:none)
p1 = contour!(support, support, z,
		levels=[0.001, 0.005, 0.02], c=[:blue, :red, :green],
		xlims=(15,40), ylims=(15,40), ratio=:equal, legend=:none,
		xlabel="x", ylabel="y")
p2 = surface(support, support, z, lw=0.1, c=cgrad([:blue, :red]),
		 legend=:none, xlabel="x", ylabel="y",camera=(-35,20))

plot(p1, p2, size=(800, 400))


### 5 point calculation

function SparseArray(t::Type,f::Function,indicies::Tuple{Int,Int})
    n,m = indicies

    V = Array{t}(0)
    I = Array{Int}(0)
    J = Array{Int}(0)

    for i in 1:n
        for j in 1:m
            val = f(i,j)
            if val!=0
                push!(I,i)
                push!(J,j)
                push!(V,val)
            end
        end
    end

    return sparse(I,J,V)
end

function d1(i,j,y)
    x(p) = y(i+p) - y(i)

    j==i-2 ? -((x(-1)*x(1)*x(2))/(x(-2)*(x(-2)-x(-1))*(x(-2)-x(1))*(x(-2)-x(2)))) :
        j==i-1 ? (x(-2)*x(1)*x(2))/((x(-2)-x(-1))*x(-1)*(x(-1)-x(1))*(x(-1)-x(2))) :
        j==i ? -(1/x(-1))-1/x(1)-1/x(2)-1/x(-2) :
        j==i+1 ? (x(-2)*x(-1)*x(2))/((x(-2)-x(1))*x(1)*(x(1)-x(-1))*(x(1)-x(2))) :
        j==i+2 ? (x(-2)*x(-1)*x(1))/((x(-2)-x(2))*x(2)*(x(2)-x(-1))*(x(2)-x(1))) :
        0
end

function d2(i,j,y)
    x(p) = y(i+p) - y(i)

    j==i-2 ? (2*(x(1)*x(2)+x(-1)*(x(1)+x(2))))/(x(-2)*(x(-2)-x(-1))*(x(-2)-x(1))*(x(-2)-x(2))) :
        j==i-1 ? -((2*(x(1)*x(2)+x(-2)*(x(1)+x(2))))/((x(-2)-x(-1))*x(-1)*(x(-1)-x(1))*(x(-1)-x(2)))) :
        j==i ? (2*(x(1)*x(2)+x(-1)*(x(1)+x(2))+x(-2)*(x(-1)+x(1)+x(2))))/(x(-2)*x(-1)*x(1)*x(2)) :
        j==i+1 ? -((2*(x(-1)*x(2)+x(-2)*(x(-1)+x(2))))/((x(-2)-x(1))*x(1)*(x(1)-x(-1))*(x(1)-x(2)))) :
        j==i+2 ? -((2*(x(-1)*x(1)+x(-2)*(x(-1)+x(1))))/((x(-2)-x(2))*x(2)*(x(2)-x(-1))*(x(2)-x(1)))) :
        0
end

### There is a bug if I do put d1,d2 inside a function

function Hamiltonian(α,II,G,p0,L,nn)

    sigmax = [0 1;1 0]
    sigmay = [0 -1im;1im 0]
    sigmaz = [1 0;0 -1]

    Ix,Iy = II
    Gx,Gy = G
    p0x,p0y = p0

    ### One could make theese nicer
    xx(i) = i<= nn//2 ? -(L+1)^((nn/2-i+1)/(nn/2+1)) + 1 : (L+1)^((i-nn/2)/(nn/2+1)) - 1
    x(i) = i<= nn//2 ? -(L+1)^((nn/2-i+1)/(nn/2+1)) + 1 - xx(nn/2)/2 : (L+1)^((i-nn/2)/(nn/2+1)) - 1 - xx(nn/2+1)/2

    ### Laplacian operator

    D2 = SparseArray(Float64,(i,j)->d2(i,j,x),(nn,nn))
    Lap = Gx*kron(D2,speye(nn)) + Gy*kron(speye(nn),D2)

    ### Potential

    V = Float64[]
    for i in 1:nn
        for j in 1:nn
            x_ = x(i)
            y_ = x(j)
            r = sqrt(x_^2+y_^2)
            push!(V,-α/r)
        end
    end
    V = spdiagm(V)

    ### Weyl coupling

    D1 = SparseArray(Float64,(i,j)->d1(i,j,x),(nn,nn))

    Hm = kron(-1/2*Lap + V,eye(2)) + Ix*kron(1im*kron(D1,speye(nn)) + p0x*speye(nn^2),sigmax) + Iy*kron(1im*kron(speye(nn),D1) + p0y*speye(nn^2),sigmay)
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ b7433dc1-846b-4173-ad51-f7598cd263a1
using FFTW, Plots, SampledSignals

# ╔═╡ 37f59728-aeae-4b15-ba8c-a90368dfa8f0
begin
    fs = 48000
	freqs = Array(LinRange(0, 2*pi*200, fs))
end

# ╔═╡ cfc3b26d-d2da-4128-9ef2-a9494c4c25e5
function plotfreqs(tone, title)
	freqs = fft(tone) |> fftshift
	domain = fftfreq(length(tone), fs) |> fftshift
	plot(domain, 10*log10.(abs.(freqs)/fs), title=title, xlimit=(0, 2000), ylimit=(-50, 0))
end

# ╔═╡ b70c60fd-b29a-46f3-8daf-e6d7277156ca
function plotfreqs!(tone, title)
	freqs = fft(tone) |> fftshift
	domain = fftfreq(length(tone), fs) |> fftshift
	plot!(domain, 10*log10.(abs.(freqs)/fs), title=title, xlimit=(0, 2000), ylimit=(-50, 0))
end

# ╔═╡ baafd0a2-b449-11eb-0654-172c64b81fd6
md"# The Harmonic Series
Periodic sounds have a base frequency that is the reciprocal of the period,
and also contain components (overtones) of integer multiples of the base frequency.
The amplitudes of the overtones determine the timbre of the sound.
"

# ╔═╡ 7a4f0f06-6b37-4a3c-81e4-8004462da33a
function note(t)
	t = t/pi%2-1
	t = 3*sqrt(3)/2*t*(t*t-1)
end

# ╔═╡ ef073029-a63e-4625-aecf-0cde96af5686
tone = SampleBuf(note.(freqs), fs)

# ╔═╡ 9997ac3d-3cae-4bb5-8290-6ca53f431533
plotfreqs(tone, "Frequencies of a single 200 Hz note")

# ╔═╡ dcd90ce9-dfc0-42ed-990a-55710c32b34e
md"## Intervals
If we play two unrelated notes at the same time we get two harmonic series"

# ╔═╡ 7f363c48-a194-4494-bb59-d5e00db70bd2
begin
   ratio=1+rand()
   plotfreqs(note.(freqs), "First note")
   plotfreqs!(note.(ratio*freqs), "Second note")
end

# ╔═╡ 9e45b6c3-a674-4b1f-b93e-11b4a99a22da
dissonant = SampleBuf((note.(freqs) + note.(ratio*freqs))/2, fs)

# ╔═╡ 07312f0d-59d9-45ee-b60a-c6ecfc2bc0df
# The frequencies of the combined note is the sum of each notes frequencies
plotfreqs(dissonant, "Two notes frequencies")

# ╔═╡ bfd2d6b6-2863-41ee-9b6a-c99ed19f8d7d
md"""## Octave
If we use frequencies in a $1:2$ ratio the harmonic series line up.
Every frequency is an overtone of the lower note.
This single harmonic series has more pronounced even overtones.
Because there is a single harmonic series,
these two notes combined are equivalent to a single note
with a different timbre.
With real world sounds different timbres, the sound source locations,
and innacuracies of analog tuning make this equivalence inexact.
"""

# ╔═╡ 3813605c-e5c8-4bd3-9813-f553d22d47b2
begin
   plotfreqs(note.(freqs), "")
   plotfreqs!(note.(2*freqs), "Octave")
end

# ╔═╡ 6a831a92-a222-48e8-abfb-72bc9f3e0fb9
octave = SampleBuf(note.(freqs)/2 + note.(2*freqs)/2, fs)

# ╔═╡ 0b9c89bf-291e-4b93-8632-f2d5117d25a4
md"If we take a tone with only odd harmonics we can construct an octave with no harmonics in common but that still sounds like a single note because of the shared harmonic series. The notes have to be scaled so that the curve of the harmonics line up to sound cohesive."

# ╔═╡ a4882ae0-32ed-4db4-ba0c-165081128f2d
function odd(t)
	t = t/pi%2-1
	t = 4*t*(abs(t)-1)
end

# ╔═╡ e3faa985-c01a-471c-ae10-0aaaee8a6b43
odd_octave = SampleBuf(odd.(freqs) + odd.(2*freqs)/8 + odd.(4*freqs)/64, fs)

# ╔═╡ 6049ed44-b19f-41a9-b017-860e80d62017
begin
    plotfreqs(odd.(freqs), "")
    plotfreqs!(odd.(2*freqs)/8, "")
	plotfreqs!(odd.(4*freqs)/64, "Odd harmonics combined")
end

# ╔═╡ ccb62252-fcf1-4fc8-b486-1df69ea0c1e3
md"Any $1:n$ ratio gives a similar lining up like the octave.
Unison $1:1$ being a special case that gives the exact original sound but louder.
However, since for natural timbres the higher harmonics sound less loud,
higher intervals harmonics blend less and stand out more from their root note.
If the higher note is made softer they start to blend again.
For example a $1:3$ ratio, or tritive sounds like this."

# ╔═╡ 14af9192-4b8a-49c3-9775-b4dfef343780
tritive = SampleBuf(note.(freqs)/2 + note.(3*freqs)/2, fs)

# ╔═╡ 02a98eb5-4021-4ea0-a0f8-b43c811d87b2
blend_tritive = SampleBuf(note.(freqs)/2 + note.(3*freqs)/8, fs)

# ╔═╡ 31b44424-b3f9-4ea0-8d99-28738f1fc857
md"## Fifth
Though only integer ratios will give harmonics that line up exactly
any rational ratio will have overlap.
A $2:3$ ratio, the simplest non-integer interval, gives a fifth.
The fifth is an octave below a tritave
and so the fifth and the tritave have a similar quality.
Going up a fifth $(3/2)$ is the inverse of going down a fifth $(2/3)$
a characteristic shared by all intervals."

# ╔═╡ 82582e9b-fcc0-49ae-9f3b-558e70d8f0b8
fifth = SampleBuf(note.(freqs)/2 + note.(3/2*freqs)/2, fs)

# ╔═╡ db7c4951-972b-43db-a647-6e61c8ebdf03
begin
   plotfreqs(note.(freqs), "")
   plotfreqs!(note.(3/2*freqs), "Fifth")
end

# ╔═╡ 53773643-1b58-4eee-ac11-bb30b7535fec
md"Every second harmonic of the higher note lines up
with every third harmonic of the lower note.
Unlike playing notes in unison or octaves,
the two notes have distinct pitches,
but unlike a random interval it's consonant.
"

# ╔═╡ 151b3663-524e-465f-9426-3f85689e9e27
md"## The Fundamental
Because of how the overtones line up,
every overtone is part of the harmonic series
of a note that isn't being played.
for the fifth this is one octave below the lower note.
This corresponds to the period of the combined sounds,
which repeats whenever the period of both notes line up.

For an interval $n:m$ that is rational and in reduced terms,
the fundamental is as $1$ in the ratio: $1:n:m$.
So the fundamental is at the inteval 1/n from the first note
and 1/m from the second.
The ratio $1:n$ is a measure of consonance and is equivalent to the porportion of harmonics in the second note that are also in the first.
Alternately the ratio $1:m$ is is another measure and is equivalent to the porportion
of harmonics of the first note present in the second.
Normally the lowest note of a collection
is chosen as the root, the reference point for calculating intervals,
but for chord inversions using a different note is often more convenient.
This measure extends naturally to collections of more than 2 notes."

# ╔═╡ 24e5caba-ce1c-4436-8ca2-647b5438658a
begin
    plotfreqs(note.(freqs), "")
    plotfreqs!(note.(3/2*freqs), "")
    plotfreqs!(note.(freqs/2), "Fifth with fundamental")
end

# ╔═╡ 8adb4c0a-a29e-4cee-bb49-ee285f90ff70
md"""## The Common Overtone
Like the fundamental of a collection of notes
is the highest note that has all of the notes as harmonics,
the dual to the fundamental is the common overtone:
The lowest note that is a harmonic of all of the notes.
The harmonic series of this note contains the frequencies
common to all the harmonic series of the other notes.

For an interval $n:m$ that is rational and in reduced terms,
the ratio with the dual fundamental is $n:m:nm$.
This is the least common multiple of $n,m,$ dual to the fundamental
which is the greatest common denominator.
The ratios $n:nm = 1:m$, and $m:nm = 1:n$ so the
common overtone gives the same measures of consonance as the
fundamental but compared with the opposite notes.
For collections of more than two notes this difference leads to
collections with simple ratios to the fundamental being different
from collections with simple ratios to their common overtone."""

# ╔═╡ 8378910d-3a2c-441e-a3b3-716d819eced9
begin
    plotfreqs(note.(freqs), "")
    plotfreqs!(note.(3/2*freqs), "")
    plotfreqs!(note.(freqs*3), "Fifth with common overtone")
end

# ╔═╡ 732ab799-970e-4ae7-a81d-cda4c79873f4
md" # Fourth
The fourth $3:4$ is an octave equivalent to the fifth.
Taking a fifth and raising the lower note by an octave
(or equivalently lowering the higher note by an octave)
gives a fourth."

# ╔═╡ df7e7e47-fe34-4250-bb0b-670fee0b56de
fourth = SampleBuf(note.(freqs)/2 + note.(freqs*4/3)/2, fs)

# ╔═╡ 518f1a8b-977b-4008-b1b7-667f0475b42c
begin
    plotfreqs(note.(freqs), "")
    plotfreqs!(note.(4/3*freqs), "")
    plotfreqs!(note.(freqs/3), "Fourth with fundamental")
end

# ╔═╡ 849e0964-2c80-4fa6-9320-5fa0c595d79a
begin
    plotfreqs(note.(freqs), "")
    plotfreqs!(note.(4/3*freqs), "")
    plotfreqs!(note.(freqs*4), "Fourth with common overtone")
end

# ╔═╡ 06b01905-7249-4492-a5b2-b9b43f1908c1
md"Other harmonic intervals work by similar rules.
Each has an integer ratio, where the interval going up is the dual of the interval going down. When treating octaves as equivalent intervals can be normalized by shifting them by octaves into the range $[1, 2]$.
Each interval $m/n$ in $[1,2]$ has an inverse also in $[1,2]$: $2(n/m)$.
An interval and it's 'octave inverse' combine to give an octave
and they will have similar qualities."

# ╔═╡ Cell order:
# ╠═b7433dc1-846b-4173-ad51-f7598cd263a1
# ╠═37f59728-aeae-4b15-ba8c-a90368dfa8f0
# ╠═cfc3b26d-d2da-4128-9ef2-a9494c4c25e5
# ╠═b70c60fd-b29a-46f3-8daf-e6d7277156ca
# ╟─baafd0a2-b449-11eb-0654-172c64b81fd6
# ╠═7a4f0f06-6b37-4a3c-81e4-8004462da33a
# ╠═ef073029-a63e-4625-aecf-0cde96af5686
# ╟─9997ac3d-3cae-4bb5-8290-6ca53f431533
# ╟─dcd90ce9-dfc0-42ed-990a-55710c32b34e
# ╟─7f363c48-a194-4494-bb59-d5e00db70bd2
# ╠═9e45b6c3-a674-4b1f-b93e-11b4a99a22da
# ╟─07312f0d-59d9-45ee-b60a-c6ecfc2bc0df
# ╟─bfd2d6b6-2863-41ee-9b6a-c99ed19f8d7d
# ╟─3813605c-e5c8-4bd3-9813-f553d22d47b2
# ╠═6a831a92-a222-48e8-abfb-72bc9f3e0fb9
# ╟─0b9c89bf-291e-4b93-8632-f2d5117d25a4
# ╠═a4882ae0-32ed-4db4-ba0c-165081128f2d
# ╠═e3faa985-c01a-471c-ae10-0aaaee8a6b43
# ╠═6049ed44-b19f-41a9-b017-860e80d62017
# ╟─ccb62252-fcf1-4fc8-b486-1df69ea0c1e3
# ╠═14af9192-4b8a-49c3-9775-b4dfef343780
# ╠═02a98eb5-4021-4ea0-a0f8-b43c811d87b2
# ╟─31b44424-b3f9-4ea0-8d99-28738f1fc857
# ╠═82582e9b-fcc0-49ae-9f3b-558e70d8f0b8
# ╟─db7c4951-972b-43db-a647-6e61c8ebdf03
# ╟─53773643-1b58-4eee-ac11-bb30b7535fec
# ╟─151b3663-524e-465f-9426-3f85689e9e27
# ╟─24e5caba-ce1c-4436-8ca2-647b5438658a
# ╟─8adb4c0a-a29e-4cee-bb49-ee285f90ff70
# ╟─8378910d-3a2c-441e-a3b3-716d819eced9
# ╟─732ab799-970e-4ae7-a81d-cda4c79873f4
# ╠═df7e7e47-fe34-4250-bb0b-670fee0b56de
# ╟─518f1a8b-977b-4008-b1b7-667f0475b42c
# ╠═849e0964-2c80-4fa6-9320-5fa0c595d79a
# ╟─06b01905-7249-4492-a5b2-b9b43f1908c1


include("header.jl")

#=
# Example calculation in a pipe system
In this notebook we will see an example of how to solve a problem in a pipe
system, using the 1-d steady-flow energy equation.

In our calculations, we will use the units that we discussed in notebook 1.0
in order to take advantage of the useful tools that come with them.
=#

# ### Set up the module
using MAE103

#=
## Example: Find the flow rate
In this example (Example 8.10 in the textbook), we are to determine the change of
flow rate that occurs after we modify a design for a fume hood. The fume hood
involves a fan that pumps air and noxious gas out of an enclosed region,
through a duct, and out to the atmosphere where it can safely mix with the
ambient air. The existing design involves a fan with a short duct of diameter
8 inches, and the losses correpond to a loss coefficient of $K_L = 5$. This
leads to a flow rate of $\dot{Q}_0 = 9$ ft$^3$/s, which is safely within
the regulated range of 6 to 12 ft$^3$/s.

The new design involves a long 100 ft pipe of galvanized iron (leading to
major loss from viscosity) and a new total loss coefficient of $K_L = 10$.
We wish to see if the new design's volume flow rate falls within the
regulated range. We will assume that the fan's head does not change in
the redesign.
=#
#-
#=
First, let's put in the parameters. We treat the gas as air, since any noxious
gas species are assumed to be at low concentrations.

Galvanized iron has a typical roughness of 0.0005 ft.
=#
Q0 = VolumeFlowRate(9u"ft^3/s")
KL0 = 5
D = Diameter(8u"inch")
μ = Viscosity(Air)
ρ = Density(Air)
ϵ = Height(0.0005u"ft")
g = Gravity()
#=
Now let's calculate the flow velocity in the original design
=#
A = Area(π/4*D^2)
V0 = Velocity(Q0/A)
value(V0,u"ft/s")

#=
And now, we can calculate the fan head $h_p$, since this is equal
to the kinetic energy of the exiting flow plus the head loss. (Pressures
are equal and ambient in the fume hood and at the exit, and elevation
changes are negligible.

$$h_p = \dfrac{V^2}{2g} + K_L\dfrac{V^2}{2g}$$
=#
hp = Head((KL0+1)*V0^2/(2*g))
value(hp,u"ft")

#=
Now we will analyze the modified system. We add 100 ft of duct after the
fan and increase the loss coefficient to 10:
=#
L = Length(100u"ft")
KL = 10.0

#=
The energy equation now becomes

$$h_p = \dfrac{V^2}{2g} + K_L\dfrac{V^2}{2g} + \dfrac{fL}{D}\dfrac{V^2}{2g}$$

The unknown in this equation is $V$, but $f$ (the friction factor) is also
unknown. Let's solve for $V$, pretending that we know $f$:

$$V = \left(\dfrac{2gh_p}{1 + K_L + fL/D}\right)^{1/2}$$

To get $f$, we need Reynolds number and roughness coefficient.
And to get Reynolds number, we need $V$. So, we clearly must iterate.

Our approach will be
0. Guess a value for $f$, based on $\epsilon/D$ and a Reynolds number close to $\infty$.
1. Calculate $V$ from the energy equation with the current guess of $f$.
2. Calculate $Re_D = \rho V D/\mu$ from $V$.
3. Calculate a new value of $f$. Check if it is equal to the guess used in 1. If
yes, we stop. If not, then return to step 1 with this $f$ as our new guess.

Let's prepare ourselves for this iteration by defining some equations. The
roughness coefficient:
=#
eD = RoughnessCoefficient(ϵ/D)

# Here is $V$ as a function of $f$ (and the other quantities, which don't change).
Vfromf(f) = Velocity(sqrt(2*g*hp/(1+KL+f*L/D)))

# Here is $Re_D$ as a function of $V$ (and other quantities that don't change)
ReD(V) = ReynoldsNumber(ρ*V*D/μ)

#=
To do the iteration, we will define a simple function that takes in
an initial guess for $f$, iterates using the procedure above, and then returns
the correct value.

In the function, we create a quantity `f_old`, which simply hold our old guess
for $f$, and `f` will hold our next guess. We will compare these in each iteration,
checking whether they are nearlyequal. We'll say they're equal if they are closer
than $10^{-8}$. To ensure that the test fails at first, we initialize `f_old` to
infinity. This will force it to proceed into the iteration loop:
=#
function iterate_f(f0::FrictionFactor)
  f = f0
  f_old = Inf
  while abs(f - f_old) > 1e-8
      f_old = f
      f = FrictionFactor(ReD(Vfromf(f_old)),eD)
  end
  return f
end

#=
Now, to get $f$ from $Re_D$ and $\epsilon/D$, we use the function
`FrictionFactor(Re,eD)`, which solves the Colebrook equation for $f$.
Here, we use it to find our initial guess for $f$, treating $Re_D$ as very large.
=#
f0 = FrictionFactor(ReynoldsNumber(1e10),eD)

f = iterate_f(f0)

# We converged on the final $f$ value! The actual velocity is thus
V = Vfromf(f)
value(Vfromf(f),u"ft/s")

# And the final Reynolds number is
ReD(V)

# And finally, the volume flow rate after making the change to the setup is
Q = VolumeFlowRate(V*A)
value(Q,u"ft^3/s")

#=
So this shows that the flow rate is too small to meet the specifications.
We would need to decrease the length, $L$, or reduce the losses in $K_L$
to make it meet the specifications.
=#