# Classification Metrics and Plots
# This files contains various metrics used in Classifiaction problems

# One Hot encode
function onehot_encode(y, labels)
   onehot_arr = zeros(length(labels), length(y))
   for i in 1:length(y)
       onehot_arr[Int(y[i]) + 1, i] = 1
   end
   onehot_arr
end

"""
    bin_to_cat(y_pred, y_true)

Function to convert binary type of data to categorical with two categories. Return `y_pred` and `y_true` of shape `(2, length(y_pred))` as tuple. Utility function to support performance metrics like `Precision`, `Recall` etc, where the function first need to be converted to categorical form before applying metric.
"""
function bin_to_cat(y_pred, y_true)
    @assert length(y_pred) == length(y_true)
    y_pred_ = zeros(2, length(y_pred))
    y_pred_[2, :] = y_pred
    y_pred_[1, :] = 1 .- y_pred
    return (y_pred_, onehot_encode(y_true, 0:1))
end

# Onecold
function onecold(y)
    onecold_arr = zeros(size(y,2))
    argmax_vec = argmax(y, dims = 1)
    for i in 1:size(y,2)
        onecold_arr[i] = argmax_vec[i].I[1]
    end
    return onecold_arr
end

"""
    confusion_matrix(y_pred, y_true)

Function to create a confusion_matrix for classification problems based on provided `y_pred` and `y_true`. Expects `y_true`, to be onehot_enocded already.
"""
function confusion_matrix(y_pred, y_true)
    @assert size(y_pred) == size(y_true)
    label_count = size(y_true, 1)
    ŷ = onehot_encode(onecold(y_pred), 1:label_count)
    return ŷ * transpose(y_true)
end

"""
    TFPN(y_pred, y_true)

Returns `Confusion Matrix` and `True Positive`, `True Negative`, `False Positive` and `False Negative` for each class based on `y_pred` and `y_true`. Expects `y_true`, to be onehot_enocded already.
"""
function TFPN(y_pred, y_true)
    @assert size(y_pred) == size(y_true)
    label_count = size(y_true, 1)
    TP = zeros(label_count) ; TN = zeros(label_count)
    FP = zeros(label_count) ; FN = zeros(label_count)
    ConfusionMatrix = confusion_matrix(y_pred, y_true)
    for c in 1 : label_count
        TP[c] = ConfusionMatrix[c,c]
        FP[c] = sum(ConfusionMatrix[:,c]) - TP[c]
        FN[c] = sum(ConfusionMatrix[c,:]) - TP[c]
        TN[c] = sum(ConfusionMatrix) - TP[c] - FP[c] - FN[c]
    end
    return ConfusionMatrix,TP, TN, FP, FN
end

"""
    binary_accuracy(y_pred, y_true; threshold=0.5)

Calculates Averaged Binary Accuracy based on `y_pred` and `y_true`. Argument `threshold` is used to specify the minimum predicted probability `y_pred` required to be labelled as `1`. Default value set as `0.5`.
"""
function binary_accuracy(y_pred, y_true; threshold=0.5)
    @assert size(y_pred) == size(y_true)
    return sum((y_pred .>= threshold) .== y_true) / size(y_true, 1)
end

"""
    categorical_accuracy(y_pred, y_true)

Calculates Averaged Categorical Accuracy based on `y_pred` and `y_true`.
"""
function categorical_accuracy(y_pred, y_true)
    @assert size(y_pred) == size(y_true)
    return sum(onecold(y_pred) .== onecold(y_true)) / size(y_true, 2)
end

"""
    sparse_categorical(y_pred, y_true)

Calculated Sparse Categorical Accuracy based on `y_pred` and `y_true`. It evaluates the maximal true value is equal to the index of the maximal predicted value. Here, `y_true` is expected to provide only an integer (start from `0` index) as label for each data element (ie. not one hot encoded).
"""
function sparse_categorical(y_pred, y_true)
    @assert size(y_pred, 2) == length(y_true)
    return sum(onecold(y_pred) .== (y_true .+ 1) ) / size(y_true, 1)
end

"""
    top_k_categorical(y_pred, y_true; k=3)

Evaluates if the index of true value is equal to any of the indices of top k predicted values. Default value of `k` set to `3`.
"""
function top_k_categorical(y_pred, y_true; k=3)
    @assert size(y_pred) == size(y_true)
    count = 0
    sparse_y = onecold(y_true)
    for i in 1:size(y_true, 2)
        top_k = partialsortperm(y_pred[:,i], 1:k, lt= >)
        for j in 1:k
            if top_k[j] == sparse_y[i]
                count+=1
                break
            end
        end
    end
    return count / length(sparse_y)
end

"""
    top_k_sparse_categorical(y_pred, y_true; k=3)

Evaluates if the true value is equal to any of the indices of top k predicted values. Default value of `k` set to `3`. Similar to `sparse_categorical`, expects the `y_true` to provide only an integer (start from `0` index) as label for each data element (ie. not one hot encoded).
"""
function top_k_sparse_categorical(y_pred, y_true; k=3)
    count = 0
    for i in 1:length(y_true)
        top_k = partialsortperm(y_pred[:,i], 1:k, lt= >)
        for j in 1:k
            if top_k[j] == y_true[i] + 1
                count+=1
                break
            end
        end
    end
    return count / length(y_true)
end

# Calculate Instances for each class in y
function calc_instances(y)
   instances = zeros(size(y, 1))
   sparse_y = onecold(y)
   for i in 1:length(instances)
       instances[i] = sum(sparse_y .== i)
   end
   return instances
end

"""
    precision(y_pred, y_true; avg_type="macro", sample_weights=nothing)

Computes the precision of the predictions with respect to the labels.

# Arguments
 - `y_pred`: predicted values.
 - `y_true`: ground truth values on the basis of which predicted values are to be assessed.
 - `avg_type="macro"`: Type of average to be used while calculating precision of multiclass models. Can take values as `macro`, `micro` and `weighted`. Default set to `macro`.
 - `sample_weights`: Class weights to be provided when `avg_type` is set to `weighted`. Useful in case of imbalanced classes.

"""
function precision(y_pred, y_true; avg_type="macro", sample_weights=nothing)
    _,TP, TN, FP, FN = TFPN(y_pred, y_true)
    # Macro-averaged Precision
    if avg_type == "macro"
        return mean(TP ./ (TP .+ FP .+ eps(eltype(TP))))
    # Micro-averaged Precision
    elseif avg_type == "micro"
        return mean(TP) / (mean(TP) + mean(FP))
    # Weighted-Averaged Precision
    elseif avg_type == "weighted"
        weights = []
        if sample_weights != nothing
            weights = sample_weights
        else
            weights = calc_instances(y_true) / size(y_true, 2)
        end
        return mean((TP ./ (TP .+ FP .+ eps(eltype(TP)))) .* weights)
    end
end

"""
    recall(y_pred, y_true; avg_type="macro", sample_weights=nothing)

Computes the recall of the predictions with respect to the labels.

# Arguments
 - `y_pred`: predicted values.
 - `y_true`: ground truth values on the basis of which predicted values are to be assessed.
 - `avg_type="macro"`: Type of average to be used while calculating precision of multiclass models. Can take values as `macro`, `micro` and `weighted`. Default set to `macro`.
 - `sample_weights`: Class weights to be provided when `avg_type` is set to `weighted`. Useful in case of imbalanced classes.

Aliases: `sensitivity` and `detection_rate`
"""
function recall(y_pred, y_true; avg_type="macro", sample_weights=nothing)
    _,TP, TN, FP, FN = TFPN(y_pred, y_true)
    # Macro-averaged Precision
    if avg_type == "macro"
        return mean(TP ./ (TP .+ FN .+ eps(eltype(TP))))
    # Micro-averaged Precision
    elseif avg_type == "micro"
        return mean(TP) / (mean(TP) + mean(FN))
    # Weighted-Averaged Precision
    else
        weights = []
        if sample_weights != nothing
            weights = sample_weights
        else
            weights = calc_instances(y_true) / size(y_true, 2)
        end
        return mean((TP ./ (TP .+ FN .+ eps(eltype(TP)))) .* weights)
    end
end
const sensitivity = recall
const detection_rate = recall

"""
    f_beta_score(y_pred, y_true; β=1, avg_type="macro", sample_weights=nothing)

Compute fbeta score. The F_beta score is the weighted harmonic mean of precision and recall, reaching its optimal value at 1 and its worst value at 0.

# Arguments
 - `y_pred`: predicted values.
 - `y_true`: ground truth values on the basis of which predicted values are to be assessed.
 - `β=1`: the weight of precision in the combined score. If `β<1`, more weight given to `precision`, while `β>1` favors recall.
 - `avg_type="macro"`: Type of average to be used while calculating precision of multiclass models. Can take values as `macro`, `micro` and `weighted`. Default set to `macro`.
 - `sample_weights`: Class weights to be provided when `avg_type` is set to `weighted`. Useful in case of imbalanced classes.

"""
function f_beta_score(y_pred, y_true; β=1, avg_type="macro", sample_weights=nothing)
    recall_ = recall(y_pred, y_true, avg_type=avg_type, sample_weights=sample_weights)
    precision_ = precision(y_pred, y_true, avg_type=avg_type, sample_weights=sample_weights)
    return (1 + β^2) * precision_ * recall_ / (precision_ + (β^2) * recall_ + eps(eltype(y_pred)))
end

"""
    specificity(y_pred, y_true; avg_type="macro", sample_weights=nothing)

Computes the specificity of the predictions with respect to the labels.

# Arguments
 - `y_pred`: predicted values.
 - `y_true`: ground truth values on the basis of which predicted values are to be assessed.
 - `avg_type="macro"`: Type of average to be used while calculating precision of multiclass models. Can take values as `macro`, `micro` and `weighted`. Default set to `macro`.
 - `sample_weights`: Class weights to be provided when `avg_type` is set to `weighted`. Useful in case of imbalanced classes.
"""
function specificity(y_pred, y_true; avg_type="macro", sample_weights=nothing)
    _, TP, TN, FP, FN = TFPN(y_pred, y_true)
    # Macro-averaged Precision
    if avg_type == "macro"
        return mean(TN ./ (TN .+ FP .+ eps(eltype(TP))))
    # Micro-averaged Precision
    elseif avg_type == "micro"
        return mean(TN) / (mean(TN) + mean(FP))
    # Weighted-Averaged Precision
    else
        weights = []
        if sample_weights != nothing
            weights = sample_weights
        else
            weights = calc_instances(y_true) / size(y_true, 2)
        end
        return mean((TN ./ (TN .+ FP .+ eps(eltype(TP)))) .* weights)
    end
end

"""
    false_alarm_rate(y_pred, y_true; avg_type="macro", sample_weights=nothing)

Computes the false_alarm_raye of the predictions with respect to the labels as `1 - specificity(y_pred, y_true, avg_type, sample_weights)`

# Arguments
 - `y_pred`: predicted values.
 - `y_true`: ground truth values on the basis of which predicted values are to be assessed.
 - `avg_type="macro"`: Type of average to be used while calculating precision of multiclass models. Can take values as `macro`, `micro` and `weighted`. Default set to `macro`.
 - `sample_weights`: Class weights to be provided when `avg_type` is set to `weighted`. Useful in case of imbalanced classes.

See also: [`specificity`](@ref)
"""
function false_alarm_rate(y_pred, y_true; avg_type="macro", sample_weights=nothing)
    return 1 - specificity(y_pred, y_true, avg_type, sample_weights)
end

"""
    cohen_kappa(y_pred, y_true)

Measures the agreement between two raters (predicted and ground truth, here) who each classify N items into C mutually exclusive categories, using the observed data to calculate the probabilities of each observer randomly seeing each category.
If the raters are in complete agreement then κ = 1. If there is no agreement among the raters other than what would be expected by chance,
κ = 0.

Ref: [Cohen's Kappa](https://en.wikipedia.org/wiki/Cohen%27s_kappa)
"""
function cohen_kappa(y_pred, y_true)
    _, tp, tn, fp, fn = TFPN(y_pred, y_true)
    mrg_a = ((tp .+ fn) .* (tp .+ fp)) ./ (tp .+ fn .+ fp .+ tn)
    mrg_b = ((fp .+ tn) .* (fn .+ tn)) ./ (tp .+ fn .+ fp .+ tn)
    expec_agree = (mrg_a .+ mrg_b) ./ (tp .+ fn .+ fp .+ tn)
    obs_agree = (tp .+ tn) ./ (tp .+ fn .+ fp .+ tn)
    cohens_kappa = mean((obs_agree .- expec_agree) ./ (1 .- expec_agree))
    return cohens_kappa
end

"""
    statsfromTFPN(TP, TN, FP, FN)

Computes statistics in case of binary classification or one-vs-all statsitics in case of multiclass classification.

# Arguments:
 - `TP`: true positive values
 - `TN`: true negative values
 - `FP`: false positive values
 - `FN`: false negative values

Return the result stats as a dictionary.
"""
function statsfromTFPN(TP, TN, FP, FN)
    Confusion_Matrix = reshape([TP, FP, FN, TN], 2, 2)
    Precision = TP / (TP + FP + eps(eltype(TP)))
    Recall = TP / (TP + FN + eps(eltype(TP)))
    Specificity = TN / (TN + FP + eps(eltype(TP)))
    F1_score = 2 * (Precision * Recall) / (Precision + Recall + eps(eltype(TP)))
    Accuracy = (TP + TN) / (TP + TN + FP + FN)
    False_alarm_rate = 1 - Specificity
    return Dict(:Confusion_Matrix => Confusion_Matrix,
                :Precision => Precision, :Recall => Recall,
                :Specificity => Specificity, :F1_score => F1_score,
                :Accuracy => Accuracy)
end

"""
    classwise_stats(y_pred, y_true)

Computes statistics for each of the class for multiclass classification based on provided `y_pred` and `y_true`.

Return the result stats as a dictionary.
"""
function classwise_stats(y_pred, y_true)
    _, TP, TN, FP, FN = TFPN(y_pred, y_true)
    ClasswiseStats = Dict()
    for i in 1:size(y_true, 1)
        ClasswiseStats[i] = StatsfromTFPN(TP[i], TN[i], FP[i], FN[i])
    end
    return ClasswiseStats
end

"""
    global_stats(y_pred, y_true; avg_type="macro")

Computes the overall statistics based on provided `y_pred` and `y_true`. `avg_type` allows to specify the type of average to be used while evaluating the stats. Currently, it can take values as "macro" or "micro".

Return the result stats as a dictionary.
"""
function global_stats(y_pred, y_true; avg_type="macro")
    confusion_matrix_, TP, TN, FP, FN = TFPN(y_pred, y_true)
    if avg_type == "macro"
        precision = mean(TP ./ (TP .+ FP .+ eps(eltype(TP))))
        recall = mean(TP ./ (TP .+ FN .+ eps(eltype(TP))))
        f1_score = 2 * precision * recall / (precision + recall + eps(eltype(TP)))
        specificity = mean(TN ./ (TN .+ FP .+ eps(eltype(TP))))
        accuracy = categorical_accuracy(y_pred, y_true)
        false_alarm_rate = 1 - specificity
        return Dict("Confusion_Matrix" => confusion_matrix_,
                "Precision" => precision, "Recall" => recall,
                "Specificity" => specificity, "F1_score" => f1_score,
                "Accuracy" => accuracy, "False_alarm_rate" => false_alarm_rate)
    elseif avg_type == "micro"
        precision = mean(TP) / (mean(TP) + mean(FP))
        recall = mean(TP) / (mean(TP) + mean(FN))
        f1_score = 2 * precision * recall / (precision + recall + eps(eltype(TP)))
        specificity = mean(TN)/ (mean(TN) + mean(FP))
        accuracy = categorical_accuracy(y_pred, y_true)
        false_alarm_rate = 1 - specificity
        return Dict("Confusion_Matrix" => confusion_matrix_,
                "Precision" => precision, "Recall" => recall,
                "Specificity" => specificity, "F1_score" => f1_score,
                "Accuracy" => accuracy, "False_alarm_rate" => false_alarm_rate)
    end
    # TODO: add weighted stats option as above functions
end

# TODO
# Concordance and Discordance

# Receiver Operating Characterstic (ROC) Curve

# AUC


"""
    Bounds.Ellipsoid([T=Float64], N)
    Bounds.Ellipsoid(center::AbstractVector, A::AbstractMatrix)

An `N`-dimensional ellipsoid defined by

```math
(x - center)^T A (x - center) = 1
```

where `size(center) == (N,)` and `size(A) == (N,N)`.

This implementation follows the algorithm presented in Mukherjee et al. (2006).[^2]

[^2]: Pia Mukherjee, et al., 2006, ApJ 638 L51 ["A Nested Sampling Algorithm for Cosmological Model Selection"](https://iopscience.iop.org/article/10.1086/501068)
"""
mutable struct Ellipsoid{T} <: AbstractBoundingSpace{T}
    center::Vector{T}
    A::Matrix{T}
    axes::Matrix{T}
    axlens::Vector{T}
    volume::T
end

function Ellipsoid(center::AbstractVector, A::AbstractMatrix)
    axes, axlens = decompose(A)
    Ellipsoid(center, A, axes, axlens, _volume(A))
end
Ellipsoid(ndim::Integer) = Ellipsoid(Float64, ndim)
Ellipsoid(T::Type, ndim::Integer) = Ellipsoid(zeros(T, ndim), diagm(0 => ones(T, ndim)))
Ellipsoid{T}(center::AbstractVector, A::AbstractMatrix) where {T} = Ellipsoid(T.(center), T.(A))

Base.broadcastable(e::Ellipsoid) = (e,)

Base.ndims(ell::Ellipsoid) = length(ell.center)

# Returns the volume of an ellipsoid given its axes matrix
_volume(A::AbstractMatrix{T}) where {T} = T(volume_prefactor(size(A, 1))) / sqrt(det(A))
volume(ell::Ellipsoid) = ell.volume

# Returns the principal axes
axes(ell::Ellipsoid) = ell.axes

decompose(A::AbstractMatrix) = decompose(Symmetric(A))  # ensure that eigen() always returns real values

function decompose(A::Symmetric)
    E = eigen(A)
    axlens = @. 1 / sqrt(E.values)
    axes = E.vectors * Diagonal(axlens)
    return axes, axlens
end

# axes and axlens
decompose(ell::Ellipsoid) = ell.axes, ell.axlens

# Scale to new volume
function scale!(ell::Ellipsoid, factor)
    # linear factor
    f = factor^(1 / ndims(ell))
    ell.A ./= f^2
    ell.axes .*= f
    ell.axlens .*= f
    ell.volume *= factor
    return ell
end

function endpoints(ell::Ellipsoid)
    axes, axlens = decompose(ell)
        # find major axis
    major_axis = axes[:, argmax(axlens)]
    return ell.center .- major_axis, ell.center .+ major_axis
end

function Base.in(x::AbstractVector, ell::Ellipsoid)
    d = x .- ell.center
    return dot(d, ell.A * d) ≤ 1.0
end

randoffset(rng::AbstractRNG, ell::Ellipsoid{T}) where {T} = axes(ell) * randball(rng, T, ndims(ell))
Base.rand(rng::AbstractRNG, ell::Ellipsoid) = ell.center .+ randoffset(rng, ell)

fit(E::Type{<:Ellipsoid}, x::AbstractMatrix{S}; pointvol = 0) where {S} = fit(E{float(S)}, x; pointvol = pointvol)

function fit(E::Type{<:Ellipsoid{R}}, x::AbstractMatrix{S}; pointvol = 0) where {R,S}
    T = float(promote_type(R, S))
    x = T.(x)
    ndim, npoints = size(x)

    # single element is an n-sphere with pointvol volume
    if npoints == 1
        pointvol > 0 || error("Cannot compute bounding ellipsoid with one point without a valid pointvol (got $pointvol)")
        d = log(pointvol) - log(volume_prefactor(ndim))
        r = exp(d / ndim)
        A = diagm(0 => fill(1 / r^2, ndim))
        return Ellipsoid(vec(x), A)
    end
    # get estimators
    center, cov = mean_and_cov(x, 2)
    delta = x .- center
    # Covariance is smaller than r^2 by a factor of 1/(n+2)
    cov .*= ndim + 2
    # Ensure cov is nonsingular
    targetprod = (npoints * pointvol / volume_prefactor(ndim))^2
    make_eigvals_positive!(cov, targetprod)

    # get transformation matrix. Note: use pinv to avoid error when cov is all zeros
    A = pinv(cov)

    # calculate expansion factor necessary to bound each points
    f = diag(delta' * (A * delta))
    fmax = maximum(f)

    # try to avoid round-off errors s.t. furthest point obeys
    # x^T A x < 1 - √eps
    flex = 1 - sqrt(eps(T))
    if fmax > flex
        A .*= flex / fmax
    end

    ell = E(vec(center), A)

    if pointvol > 0
        minvol = npoints * pointvol
        vol = volume(ell)
        vol < minvol && scale!(ell, minvol / vol)
    end

    return ell
end

# ---------------------------------------------
# Helper functions

"""
    volume_prefactor(::Integer)

Volume constant for an n-dimensional sphere:

for n even:      (2pi)^(n    /2) / (2 * 4 * ... * n)
for n odd :  2 * (2pi)^((n-1)/2) / (1 * 3 * ... * n)
"""
function volume_prefactor(n::Integer)
    f, range = iseven(n) ? (1.0, 2:2:n) : (2.0, 3:2:n)
    for i in range
        f *= 2π / i
    end
    return f
end

# sample N samples from unit D-dimensional ball
randball(T::Type, D::Integer, N::Integer) = randball(GLOBAL_RNG, T, D, N)
function randball(rng::AbstractRNG, T::Type, D::Integer, N::Integer)
    z = randn(rng, T, D, N)
    z .*= rand(rng, T, 1, N).^(1 ./ D) ./ sqrt.(sum(p->p^2, z, dims = 1))
    return z
end

# sample from unit D-dimensional ball
randball(T::Type, D::Integer) = randball(GLOBAL_RNG, T, D)
function randball(rng::AbstractRNG, T::Type, D::Integer)
    z = randn(rng, T, D)
    z .*= rand(rng)^(1 / D) / sqrt(sum(p->p^2, z))
    return z
end

function make_eigvals_positive!(cov::AbstractMatrix, targetprod)
    E = eigen(cov)
    mask = E.values .< 1e-10
    if any(mask)
        nzprod = prod(E.values[.!mask])
        nzeros = count(mask)
        E.values[mask] .= (targetprod / nzprod)^(1 / nzeros)
        cov .= E.vectors * Diagonal(E.values) / E.vectors
    end
    return cov
end

make_eigvals_positive(cov::AbstractMatrix, targetprod) = make_eigvals_positive!(copy(cov), targetprod)


@syms 𝑥
@polyvar 𝑦

function poly(eq, v::Pair)
    α, β = value(first(v)), last(v)
    terms = collect_powers(eq, α)

    if !all(isone, denominator.(collect(keys(terms))))
        error("fractional power not supported")
    end

    sum([val*β^numerator(k) for (k,val) in terms]; init=0)
end

function poly(eq, β)
    xs = get_variables(eq)

    if length(xs) > 1
        error("more than one implicit variable.")
    elseif length(xs) == 0
        return nothing
    else
        return poly(eq, xs[1] => β)
    end
end

poly(eq) = poly(eq, 𝑦)

function sym(p::AbstractPolynomialLike, v::Pair)
    β, α = first(v), value(last(v))
    sum([α^maxdegree(t,β)*c for (t,c) in zip(terms(p), coefficients(p))]; init=0)
end

sym(p::AbstractPolynomialLike, α) = sym(p, var(p) => α)
sym(p::AbstractPolynomialLike) = sym(p, var(p) => 𝑥)

###############################################################################

# pox (power-of-x) is a symbolic function to keep track of the powers of x
# pox(k,n) means k*x^n
@syms pox(k, n)

is_pox(x) = istree(x) && operation(x)==pox
is_not_pox(x) = !is_pox(x)

get_coef(p) = is_pox(p) ? arguments(p)[1] : p
get_power(p) = is_pox(p) ? arguments(p)[2] : 0

replace_x(eq, x) = substitute(eq, Dict(x => pox(1,1)))

iscomplex(x) = x isa Complex

count_rule1 = @rule ^(pox(~k, ~n1), ~n2) => isequal(~k,1) ? pox(1, ~n1 * ~n2) : pox(^(~k,~n2), ~n1 * ~n2)
count_rule2 = @rule pox(~k1, ~n1) * pox(~k2, ~n2) => pox(~k1 * ~k2, ~n1 + ~n2)
count_rule3 = @acrule pox(~k, ~n) * ~u::is_not_pox => pox(~k * ~u, ~n)

"""
    collect_powers separates the powers of x in eq (a polynomial) and returns
    a dictionary of power => term
"""
function collect_powers(eq, x)
    eq = expand(expand_derivatives(eq))
    eq = replace_x(eq, x)
    #eq = Prewalk(PassThrough(count_rule1))(eq)
    eq = Fixpoint(Prewalk(PassThrough(Chain([count_rule1, count_rule2, count_rule3]))))(eq)

    if !istree(eq)
        return Dict{Any, Any}(0 => eq)
    elseif is_pox(eq)
        return Dict{Any, Any}(get_power(eq) => get_coef(eq))
    else
        eqs = Dict{Any, Any}()
        for term in arguments(eq)
            n = get_power(term)
            if haskey(eqs, n)
                eqs[n] = eqs[n] + get_coef(term)
            else
                eqs[n] = get_coef(term)
            end
        end

        return eqs
    end
end


"""
    Xtal._is_linearly_independent(vecs::AbstractMatrix{<:Number}) -> Bool
    Xtal._is_linearly_independent(vecs::AbstractVector{<:Number}...) -> Bool

Determines whether a set of vectors is linearly independent.

This function always returns `false` if the first dimension of the matrix is less than the second.
"""
function is_linearly_independent(vecs::AbstractMatrix{<:Number})
    size(vecs)[1] < size(vecs)[2] || return false
    return LinearAlgebra.rank(vecs) == minimum(size(vecs))
end

# Same thing but for sets of vectors
function is_linearly_independent(vecs::AbstractVector{<:Number}...)
    return is_linearly_independent(hcat(vecs...))
end

#=
"""
    _to_vectors(a::AbstractArray{T,N}) where {T,N} -> Vector{...Vector{T}}

Converts an array to a nest of vectors within vectors.
"""
function _to_vectors(a::AbstractArray{T,N}) where {T,N}

end
=#


#=

   The Euro coin problem.

   From Think Bayes, page 33ff
   """
   A statistical statement appeared in "The Guardian" on Friday January 4, 2002:
      When spun on edge 250 times, a Belgian one-euro coin
      came up heads 140 times and tails 110. 'It looks very
      suspicious to me,' said Barry Blight, a statistics lecturer
      at the London School of Economics. 'If the coin were
      unbiased, the chance of getting a result as extreme as
      that would be less than 7%.'

   But do these data give evidence that the coin is biased rather than fair?
   """

   Continues on page 41:
   """
   Exercise 4.1. Suppose that instead of observing coin tosses directly, you measure
   the outcome using an instrument that is not always correct. Specifically, suppose
   there is a probability y that an actual heads is reported as tails, or actual tails re-
   ported as heads.

   Write a class that estimates the bias of a coin given a series of outcomes and the
   value of y .
   How does the spread of the posterior distribution depend on y ?
   """

   mean(prob): 0.5213151098904925
   Distributions of variable probLt0_5 (num:0)
   0.00000 =>    7020  (0.702000)
   1.00000 =>    2980  (0.298000)

   Distributions of variable error (num:0)
   0.00000 =>    7051  (0.705100)
   1.00000 =>    2949  (0.294900)

   Cf euro_coin_problem.jl
   Cf ~/blog/euro_coin_unreliable_measurements.blog
      ~/webppl/euro_coin_unreliable_measurements.wppl
=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function euro_coin_problem_unreliable_measurements()
    n = 250

    # Probability of throwing head
    prob ~ Beta(2,2)

    # We measure incorrect with probability 0.2
    error ~ flip(0.2)

    throwCoin ~ filldist(flip(prob),n)

    coin = tzeros(n)
    for i in 1:n
        coin[i] ~ Dirac(error == true ? 1-throwCoin[i] : throwCoin[i])
    end

    sum250 ~ Dirac(sum([coin[i] == true ? 1 : 0 for i in 1:n]))
    true ~ Dirac(sum250 == 140)
    probLt0_5 ~ Dirac(prob < 0.5)

end

model = euro_coin_problem_unreliable_measurements()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

println("mean(prob): ", mean(chns[:prob]))
show_var_dist_pct(chns, :probLt0_5)
show_var_dist_pct(chns, :error)


# # Approximate Inference on Sparse GPs using VI

# Loading the necessary packages and setting seed.

using AbstractGPs, Plots, Random
Random.seed!(1234);

# Loading toy regression
# [dataset](https://github.com/GPflow/docs/blob/master/doc/source/notebooks/basics/data/regression_1D.csv)
# taken from GPFlow examples.

x = [0.8658165855998895, 0.6661700880180962, 0.8049218148148531, 0.7714303440386239,
    0.14790478354654835, 0.8666105548197428, 0.007044577166530286, 0.026331737288148638,
    0.17188596617099916, 0.8897812990554013, 0.24323574561119998, 0.028590102134105955];
y = [1.5255314337144372, 3.6434202968230003, 3.010885733911661, 3.774442382979625,
    3.3687639483798324, 1.5506452040608503, 3.790447985799683, 3.8689707574953,
    3.4933565751758713, 1.4284538820635841, 3.8715350915692364, 3.7045949061144983];
scatter(x, y, xlabel="x", ylabel="y")

# Split the observations into train and test set.

(x_train, y_train) = (x[begin:8], y[begin:8]);
(x_test, y_test) = (x[9:end], y[9:end]);

# Instantiate the kernel.

k = Matern52Kernel()

# Instantiate a Gaussian Process with the given kernel `k`.

f = GP(k);

# Instantiate a `FiniteGP`, a finite dimentional projection at the inputs of the dataset
# observed under Gaussian Noise with $\sigma = 0.001$ .

fx = f(x_train, 0.001);

# Calculating the exact posterior over `f` given `y`. The GP's kernel currently has some
# arbitrary fixed parameters.

p_fx = posterior(fx, y_train)

logpdf(p_fx(x_test), y_test)

# ## Variational Inference

# Sanity check for the Evidence Lower BOund (ELBO) implemented according to
# M. K. Titsias's _Variational learning of inducing variables in sparse Gaussian processes_

elbo(fx, y_train, f(rand(7)))

# We will be using `Optim.jl` package's `LBFGS` algorithm to maximize the given ELBO.

using Optim

# Create a helper function for optimization. It takes in the parameters (both variational
# and model parameters) and returns the negative of the ELBO.

function optim_function(params; x=x_train, y=y_train)
    kernel = ScaledKernel(
        transform(
            Matern52Kernel(),
            ScaleTransform(exp.(params[1]))
            ),
        exp.(params[2])
    )
    f = GP(kernel)
    fx = f(x, 0.1)
    return -elbo(fx, y, f(params[3:end]))
end

# Initialize the parameters (Varitational and Model parameters)
x0 = rand(7)

# Sanity check for the helper function. We intend to minimize the result of this function.

-optim_function(x0)

# Optimize using `Optim.jl` package's `LBFGS` algorithm.

opt = optimize(optim_function, x0, LBFGS())

# Optimal parameters:

opt.minimizer

# ELBO with optimal parameters. We see that there is significant improvement when compared
# to the initial parameters.

-optim_function(opt.minimizer; x=x_test, y=y_test)

# Visualize the posterior.

opt_kernel = ScaledKernel(
    transform(
        Matern52Kernel(),
        ScaleTransform(exp.(opt.minimizer[1]))
        ),
    exp.(opt.minimizer[2])
)
opt_f = GP(opt_kernel)
opt_fx = opt_f(x_train, 0.1)
ap = approx_posterior(VFE(), opt_fx, y_train, opt_f(opt.minimizer[3:end]));

# Average log-marginal-probability of data with posterior kernel parameter samples sampled
# using ESS. We can observe that there is significant improvement over exact posterior with
# default kernel parameters.

logpdf(ap(x_test), y_test)

plt = plot(ap, 0:0.001:1, label="Approx Posterior")
plot!(plt, p_fx, 0:0.001:1, label="Exact Posterior")
scatter!(
    plt,
    opt.minimizer[3:end],
    mean(rand(ap(opt.minimizer[3:end], 0.1), 100), dims=2),
    label="Pseudo-points"
)
scatter!(plt, x_train, y_train, label="Train Data")
scatter!(plt, x_test, y_test, label="Test Data")
plt


"""
	transferoperator_exact(t::AbstractTriangulation)

Compute the transfer operator from a triangulation.
"""
function transferoperator_exact(t::AbstractTriangulation)
    n_simplices = size(t.simplex_inds, 1)
    dim = size(t.points, 2)

    #=
    Mismatch for the markovity of the Markov matrix is at most δ if ϵ tolerance
	as below
    =#
    δ::Float64 = 1/10^5
    ϵ::Float64 = δ/n_simplices

    #=
    Tolerance for similary of convex expansion coefficients of simplex vertices in
    simplexintersection function.
    =#
    convex_params_tol::Float64 = 1/10^12

    TO = zeros(Float64, n_simplices, n_simplices) # intersecting volumes

    for i in 1:n_simplices
        imvol = t.volumes_im[i]
        for j in 1:n_simplices
            vol = t.volumes[j]
            if vol * imvol > 0 && (vol/imvol) > ϵ
                # Intersecting volume between these two simplices
                TO[i, j] = simplexintersection(
                                transpose(t.points[t.simplex_inds[j, :], :]),
                                transpose(t.impoints[t.simplex_inds[i, :], :])
                            ) / imvol
            end
        end
    end
    return ExactSimplexTransferOperator(TO)
end

"""
    transferoperator_exact_p(t::AbstractTriangulation)

Estimate the transfer operator from a triangulation of a state space.
Will run in parallel if `nprocs() > 1`.
"""
function transferoperator_exact_p(t::AbstractTriangulation)

    n_simplices = size(t.simplex_inds, 1)
    dim = size(t.points, 2)

    #=
    Mismatch for the markovity of the Markov matrix is at most δ if ϵ
    tolerance as below.
    =#
    δ::Float64 = 1/10^5
    ϵ::Float64 = δ/n_simplices

    #=
    Tolerance for similary of convex expansion coefficients of simplex vertices
    in simplexintersection function.
    =#
    convex_params_tol::Float64 = 1/10^12

    TO = SharedArray{Float64}(n_simplices, n_simplices)

    @sync @distributed for i in 1:n_simplices
        imvol = t.volumes_im[i]
        for j in 1:n_simplices
            vol = t.volumes[j]
            if vol * imvol > 0 && (vol/imvol) > ϵ
                # Intersecting volume between these two simplices
                TO[i, j] = simplexintersection(
                                transpose(t.points[t.simplex_inds[j, :], :]),
                                transpose(t.impoints[t.simplex_inds[i, :], :])
                            ) / imvol
            end
        end
    end

    return ExactSimplexTransferOperator(Array(TO))
end

transferoperator_triang_exact = transferoperator_exact
export transferoperator_triang_exact


#=
   Logistic regresiion
   Inspired by WebPPL's examples/logisticRegression.wppl

   Summary Statistics (for PG(15))
      parameters       mean       std   naive_se      mcse         ess      rhat   ess_per_sec
          Symbol    Float64   Float64    Float64   Float64     Float64   Float64       Float64

               m     1.0844    0.5870     0.0059    0.0077   6361.4279    1.0009      153.4390
               b     0.1315    0.9292     0.0093    0.0111   6431.4331    0.9999      155.1275
           sigma     0.8779    0.8513     0.0085    0.0107   6452.7250    0.9999      155.6411
            y[1]   -10.7253    6.0611     0.0606    0.0801   6408.9867    1.0008      154.5861
            y[2]    -5.3599    3.2295     0.0323    0.0439   6465.7579    1.0007      155.9555
            y[3]     2.5165    1.7809     0.0178    0.0200   6562.8112    1.0010      158.2964
            y[4]     6.6892    3.7534     0.0375    0.0454   6516.1303    1.0008      157.1705
            y[5]    11.0158    5.9711     0.0597    0.0746   6263.6413    1.0010      151.0804
  labels_post[1]     0.0163    0.1266     0.0013    0.0013   8339.5233    0.9999      201.1511
  labels_post[2]     0.0552    0.2284     0.0023    0.0026   8486.3592    1.0001      204.6928
  labels_post[3]     0.8459    0.3611     0.0036    0.0041   8013.9274    1.0001      193.2977
  labels_post[4]     0.9707    0.1687     0.0017    0.0018   8543.2209    1.0002      206.0643
  labels_post[5]     0.9873    0.1120     0.0011    0.0012   9349.4013    0.9999      225.5096

  The labels_post are the posterior predictive and are fairly close to labels values: [0, 0, 1, 1, 1]

=#

using Turing, StatsPlots, DataFrames

include("jl_utils.jl")

@model function logistic_regression(x, labels)
    m ~ Normal(0, 1)
    b ~ Normal(0, 1)
    sigma ~ Gamma(1, 1)

    n = length(x)
    y = tzeros(n)
    sigmoid = tzeros(n)
    for i in 1:n
        y[i] ~ Normal(m * x[i] + b, sigma)
        sigmoid[i] = 1 / (1 + exp(-1 * y[i]))
        labels[i] ~ Bernoulli(sigmoid[i])
    end

    # Post predictive of labels
    labels_post = tzeros(n)
    for i in 1:n
        labels_post[i] ~ Bernoulli(sigmoid[i])
    end
end

xs = [-10, -5, 2, 6, 10]
labels = [0, 0, 1, 1, 1]
model = logistic_regression(xs,labels)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)


module _K2En
export K2En
using Statistics: std
    """
        K2, Ci = K2En(Sig)

    Returns the Kolmogorov entropy estimates `K2` and the correlation
    integrals `Ci` for `m` = [1,2] estimated from the data sequence `Sig`
    using the default parameters: embedding dimension = 2, time delay = 1,
    r = 0.2*SD(`Sig`), logarithm = natural

        K2, Ci = K2En(Sig::AbstractArray{T,1} where T<:Real; m::Int=2, tau::Int=1, r::Real=0.2*std(Sig,corrected=false), Logx::Real=exp(1))

    Returns the Kolmogorov entropy estimates `K2` for dimensions = [1,...,`m`]
    estimated from the data sequence `Sig` using the 'keyword' arguments:

    # Arguments:
    `m`     - Embedding Dimension, a positive integer\n
    `tau`   - Time Delay, a positive integer\n
    `r`     - Radius, a positive scalar \n
    `Logx`  - Logarithm base, a positive scalar\n

    # See also `DistEn`, `XK2En`, `MSEn`

    # References:
        [1] Peter Grassberger and Itamar Procaccia,
            "Estimation of the Kolmogorov entropy from a chaotic signal."
            Physical review A 28.4 (1983): 2591.

        [2] Lin Gao, Jue Wang  and Longwei Chen
            "Event-related desynchronization and synchronization
            quantification in motor-related EEG by Kolmogorov entropy"
            J Neural Eng. 2013 Jun;10(3):03602

    """
    function K2En(Sig::AbstractArray{T,1} where T<:Real; m::Int=2, tau::Int=1, r::Real=0.2*std(Sig,corrected=false), Logx::Real=exp(1))

    N = size(Sig)[1]
    (N>10) ? nothing : error("Sig:   must be a numeric vector")
    (m > 0) ? nothing : error("m:     must be an integer > 0")
    (tau > 0) ? nothing : error("tau:   must be an integer > 0")
    (r>0) ? nothing :  error("r:     must be a positive scalar > 0")
    (Logx>0) ? nothing :  error("Logx:     must be a positive number > 0")

    m = m+1
    Zm = zeros(N,m)
    Ci = zeros(m)
    for n = 1:m
        N2 = N-(n-1)*tau
        Zm[1:N2,n] = Sig[(n-1)*tau + 1:N]
        Norm = Inf*ones(N2-1,N2-1)
        for k = 1:N2-1
            Temp = repeat(transpose(Zm[k,1:n]),N2-k,1) - Zm[k+1:N2,1:n]
            Norm[k,k:N2-1] = sqrt.(sum(Temp.^2, dims=2))
        end
        Ci[n] = 2*sum(Norm .< r)/(N2*(N2-1))
    end

    K2 = log.(Logx, Ci[1:m-1]./Ci[2:m])/tau
    K2[isinf.(K2)] .= NaN

    return K2, Ci
    end
end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""


@doc raw"""
    LogEuclideanMetric <: RiemannianMetric

The LogEuclidean Metric consists of the Euclidean metric applied to all elements after mapping them
into the Lie Algebra, i.e. performing a matrix logarithm beforehand.
"""
struct LogEuclideanMetric <: RiemannianMetric end

@doc raw"""
    distance(M::MetricManifold{SymmetricPositiveDefinite{N},LogEuclideanMetric}, p, q)

Compute the distance on the [`SymmetricPositiveDefinite`](@ref) manifold between
`p` and `q` as a [`MetricManifold`](@ref) with [`LogEuclideanMetric`](@ref).
The formula reads

```math
    d_{\mathcal P(n)}(p,q) = \lVert \operatorname{Log} p - \operatorname{Log} q \rVert_{\mathrm{F}}
```

where $\operatorname{Log}$ denotes the matrix logarithm and
$\lVert\cdot\rVert_{\mathrm{F}}$ denotes the matrix Frobenius norm.
"""
function distance(
    M::MetricManifold{ℝ,SymmetricPositiveDefinite{N},LogEuclideanMetric},
    p,
    q,
) where {N}
    return norm(log(Symmetric(p)) - log(Symmetric(q)))
end


using  ForwardDiff,GenericLinearAlgebra,LinearAlgebra

mutable struct svdmat
    U::Any
    S::Any
    V::Any
end
export svd_lowrank

function svd_lowrank(A::Matrix, l::Int, q::Int)
    #= note:: The implementation is based on the Algorithm 5.1 from
    Halko et al, 2009.=#
    Q = get_approximate_basis(A, l, q)
    Q = conj(Q')
    B = Q * A
    U, S, V = svd(B)
    V = conj(V')
    U = Q * U
    return svdmat(U, S, V)
end

function get_approximate_basis(A::Matrix, l::Int, q::Int)
    #= note:: The implementation is based on the Algorithm 4.4 from
    Halko et al, 2009. =#
    m, n = size(A)
    omega = randn(n, l)
    Y = A * omega
    Q, R = qr(Y)
    Y_H = conj(Y')

    for i = 1:q
        Q, R = qr(Y_H * Q)
        Q, R = qr(Y * Q)
    end
    return Q
end


"""
```julia
using DifferentialEquations, Plots
function lorenz!(du,u,p,t)
 du[1] = 10.0*(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz!,u0,tspan)
sol = solve(prob)
plot(sol)
```
"""
module DifferentialEquations

  using Reexport

  @reexport using DiffEqBase
  @reexport using DiffEqNoiseProcess
  @reexport using RecursiveArrayTools

  @reexport using SteadyStateDiffEq
  @reexport using StochasticDiffEq
  @reexport using OrdinaryDiffEq
  @reexport using BoundaryValueDiffEq
  using Sundials
  @reexport using DelayDiffEq

  @reexport using DiffEqCallbacks
  @reexport using DiffEqJump

  @reexport using DiffEqFinancial
  @reexport using MultiScaleArrays
  @reexport using DiffEqPhysics

  @reexport using DimensionalPlotRecipes
  @reexport using ParameterizedFunctions

  using LinearAlgebra

  import DiffEqBase: solve

  include("default_solve.jl")
  include("default_arg_parsing.jl")
  include("ode_default_alg.jl")
  include("sde_default_alg.jl")
  include("dae_default_alg.jl")
  include("dde_default_alg.jl")
  include("discrete_default_alg.jl")
  include("rode_default_alg.jl")
  include("steady_state_default_alg.jl")
  include("bvp_default_alg.jl")

end # module


#!/usr/bin/julia

# Date: 21 August 2016
# https://github.com/trizen

# An efficient algorithm for computing factorial of a large number, modulus a larger number.
# Algorithm from: http://stackoverflow.com/questions/9727962/fast-way-to-calculate-n-mod-m-where-m-is-prime

function facmod(n::Int64, mod::Int64)

    f = 1
    for k in n+1:mod-1
       f *= k
       f %= mod
    end

    (-1 * invmod(f, mod) + mod) % mod
end

println(facmod(100_000_000, 100_000_037))


# Jonas August, 2/2022

#=
This problem was asked by Two Sigma.

Using a function rand7() that returns an integer from 1 to 7 (inclusive) with uniform probability, implement a function rand5() that returns an integer from 1 to 5 (inclusive).
=#

rand7() = rand(1:7)

function rand5()
    x = 6
    while x > 5  # Keep sampling uniformly on 1:7 until in range 1:5
        x = rand7()
    end
    x
end

sample = [rand5() for _ in 1:50000]

" Calculate 5 bin histogram. "
function hist5(sample::Vector{Int})
    h = zeros(Int, 5)
    for x in sample
        h[x] += 1
    end
    h
end

# Test uniformity of probability: All bins should have similar counts.
println(hist5(sample))


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 9e38d21a-f785-11ea-1ce0-5d313a4e9a7e
using Pkg, DrWatson

# ╔═╡ 9e393624-f785-11ea-077b-cb4f64059bb3
begin
	@quickactivate "StatisticalRethinkingTuring"
	using StatisticalRethinking
end

# ╔═╡ 4e27a698-f785-11ea-1ded-a548d1953a14
md"## Fig3.1t.jl"

# ╔═╡ 9e3b92a2-f785-11ea-090f-a78cbded0a8d
begin
	N = 10000
	p_grid = range(0, stop=1, length=N)
	prior = ones(length(p_grid))
	likelihood = pdf.(Binomial.(9, p_grid), 6)
	posterior = likelihood .* prior
	posterior = posterior / sum(posterior)
	samples = sample(p_grid, Weights(posterior), length(p_grid));
end;

# ╔═╡ 9e43adc0-f785-11ea-2b6a-f11004009094
begin
	p1 = scatter(samples, ylim=(0, 1), xlab="Sample number",
	  ylab="Proportion water(p)", leg=false)
	p2 = density(samples, xlim=(0.0, 1.0), ylim=(0.0, 3.0),
	  xlab="Proportion water (p)",
	  ylab="Density", leg=false)
end;

# ╔═╡ d13a98ea-f785-11ea-0cb5-e51b20ddcdac
plot(p1, p2, layout=(1,2))

# ╔═╡ 9e520366-f785-11ea-1d50-8f7ab0c29c9c
md"## End of Fig3.1t.jl"

# ╔═╡ Cell order:
# ╟─4e27a698-f785-11ea-1ded-a548d1953a14
# ╠═9e38d21a-f785-11ea-1ce0-5d313a4e9a7e
# ╠═9e393624-f785-11ea-077b-cb4f64059bb3
# ╠═9e3b92a2-f785-11ea-090f-a78cbded0a8d
# ╠═9e43adc0-f785-11ea-2b6a-f11004009094
# ╠═d13a98ea-f785-11ea-0cb5-e51b20ddcdac
# ╟─9e520366-f785-11ea-1d50-8f7ab0c29c9c


export eigen2x2
function eigen2x2(A)
	Tp = eltype(A)
	a = A[1, 1]; b = A[1, 2]; c = A[2, 1]; d = A[2, 2];
	T = tr(A); D = det(A);

	L₁ = T / 2 + sqrt((T^2 / 4) - D)
	L₂ = T / 2 - sqrt((T^2 / 4) - D)
	values = [L₁, L₂]

	global vectors = []

	if !iszero(c)
		push!(vectors, [L₁ - d, c])
		push!(vectors, [L₂ - d, c])
	elseif !iszero(b)
		push!(vectors, [b, L₁ - a])
		push!(vectors, [b, L₂ - a])
	elseif iszero(b) && iszero(c)
		push!(vectors, [one(Tp), zero(Tp)])
		push!(vectors, [zero(Tp), one(Tp)])
	end

	return values, vectors
end


using LinearAlgebra

function polar_force(force,angle)
    return force.*[sin(deg2rad(angle)), cos(deg2rad(angle))]
end

function moments(forces, locations)
    sum = zeros(3)
    for force_loc in zip(forces, locations)
        force, location = force_loc
        force = vcat(force,[0])
        location = vcat(location,[0])
        res = cross(force, location)
        print(res)
        sum+=res
    end
    sum
end

#forces = [polar_force(718.4, 180 - 30), polar_force(879.54, 45), polar_force(100, -90)]
forces = [[0, -2000], [0, -1200], [0, 15600], [0, -12400]]
location = [[0, 0], [6, 0], [10, 0], [12, 0]]
print(moments(forces, location))


# find expression of projections on cones and their derivatives here:
#   https://stanford.edu/~boyd/papers/pdf/cone_prog_refine.pdf

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Zeros) where {T}

projection of vector `v` on zero cone i.e. K = {0}
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Zeros) where {T}
    return FillArrays.Zeros{T}(size(v))
end

"""
    projection_on_set(::AbstractDistance, ::MOI.Reals, v::Array{T}) where {T}

projection of vector `v` on real cone i.e. K = R
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Reals) where {T}
    return v
end

function projection_on_set(::DefaultDistance, v::T, set::MOI.EqualTo) where {T}
    return zero(T) + set.value
end

function projection_on_set(::DefaultDistance, v::T, set::MOI.LessThan) where {T}
    return min(v, MOI.constant(set))
end

function projection_on_set(::DefaultDistance, v::T, set::MOI.GreaterThan) where {T}
    return max(v, MOI.constant(set))
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonnegatives) where {T}

projection of vector `v` on Nonnegative cone i.e. K = R^n+
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonnegatives) where {T}
    return max.(v, zero(T))
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonpositives) where {T}

projection of vector `v` on Nonpositive cone i.e. K = R^n-
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonpositives) where {T}
    return min.(v, zero(T))
end

"""
    projection_on_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, ::MOI.SecondOrderCone) where {T}

projection of vector `v` on second order cone i.e. K = {(t, x) ∈ R+ × Rn |  ||x|| ≤ t }
"""
function projection_on_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, ::MOI.SecondOrderCone) where {p, T}
    t = v[1]
    x = v[2:length(v)]
    norm_x = LinearAlgebra.norm(x, p)
    if norm_x <= t
        return copy(v)
    elseif norm_x <= -t
        return zeros(T, size(v))
    end
    result = zeros(T, size(v))
    result[1] = one(T)
    result[2:length(v)] = x / norm_x
    result *= (norm_x + t) / 2
    return result
end

function projection_on_set(::DefaultDistance, v::AbstractVector{T}, cone::MOI.SecondOrderCone) where {T}
    return projection_on_set(NormedEpigraphDistance{2}(), v, cone)
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.PositiveSemidefiniteConeTriangle) where {T}

projection of vector `v` on positive semidefinite cone i.e. K = S^n⨥
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.PositiveSemidefiniteConeTriangle) where {T}
    dim = isqrt(2*length(v))
    X = unvec_symm(v, dim)
    λ, U = LinearAlgebra.eigen(X)
    D = LinearAlgebra.Diagonal(max.(λ, 0))
    return vec_symm(U * D * U')
end

"""
    unvec_symm(x, dim)

Returns a dim-by-dim symmetric matrix corresponding to `x`.

`x` is a vector of length dim*(dim + 1)/2, corresponding to a symmetric matrix
```
X = [ X11  X12 ... X1k
      X21  X22 ... X2k
      ...
      Xk1  Xk2 ... Xkk ],
```
where `vec(X) = (X11, X12, X22, X13, X23, X33, ..., Xkk)`.

### Note on inner products

Note that the scalar product for the symmetric matrix in its vectorized form is
the sum of the pairwise product of the diagonal entries plus twice the sum of
the pairwise product of the upper diagonal entries; see [p. 634, 1].
Therefore, this transformation breaks inner products:
```
dot(unvec_symm(x, dim), unvec_symm(y, dim)) != dot(x, y).
```

### References

[1] Boyd, S. and Vandenberghe, L.. *Convex optimization*. Cambridge university press, 2004.
"""
function unvec_symm(x, dim=isqrt(2length(x)))
    X = zeros(eltype(x), dim, dim)
    idx = 1
    for i in 1:dim
        for j in 1:i
            X[j,i] = X[i,j] = x[idx]
            idx += 1
        end
    end
    return X
end

"""
    vec_symm(X)

Returns a vectorized representation of a symmetric matrix `X`.
`vec(X) = (X11, X12, X22, X13, X23, X33, ..., Xkk)`

### Note on inner products

Note that the scalar product for the symmetric matrix in its vectorized form is
the sum of the pairwise product of the diagonal entries plus twice the sum of
the pairwise product of the upper diagonal entries; see [p. 634, 1].
Therefore, this transformation breaks inner products:
```
dot(vec_symm(X), vec_symm(Y)) != dot(X, Y).
```

### References

[1] Boyd, S. and Vandenberghe, L.. *Convex optimization*. Cambridge university press, 2004.

"""
function vec_symm(X)
    return X[LinearAlgebra.triu(trues(size(X)))]
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.ExponentialCone) where {T}

projection of vector `v` on closure of the exponential cone
i.e. `cl(Kexp) = {(x,y,z) | y e^(x/y) <= z, y>0 } U {(x,y,z)| x <= 0, y = 0, z >= 0}`.

References:
* [Proximal Algorithms, 6.3.4](https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf)
by Neal Parikh and Stephen Boyd.
* [Projection, presolve in MOSEK: exponential, and power cones]
(https://docs.mosek.com/slides/2018/ismp2018/ismp-friberg.pdf) by Henrik Friberg
* [Projection onto the exponential cone: a univariate root-finding problem]
(https://docs.mosek.com/whitepapers/expcone-proj.pdf) by Henrik Friberg
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::MOI.ExponentialCone; tol=1e-8) where {T}
    _check_dimension(v, s)

    if _in_exp_cone(v; dual=false)
        return SVector{3}(v)
    end
    if _in_exp_cone(-v; dual=true)
        # if in polar cone Ko = -K*
        return zeros(SVector{3,T})
    end
    if v[1] <= 0 && v[2] <= 0
        return @SVector([v[1], 0, max(v[3],0)])
    end

    return _exp_cone_proj_case_4(v; tol=tol)
end

function _in_exp_cone(v::AbstractVector{T}; dual=false, tol=1e-8) where {T}
    if dual
        return (
            (isapprox(v[1], 0, atol=tol) && v[2] >= 0 && v[3] >= 0) ||
            (v[1] < 0 && v[1]*exp(v[2]/v[1]) + ℯ * v[3] >= tol)
        )
    else
        return (
            (v[1] <= 0 && isapprox(v[2], 0, atol=tol) && v[3] >= 0) ||
            (v[2] > 0 && v[2] * exp(v[1] / v[2]) - v[3] <= tol)
        )
    end
end

function _exp_cone_proj_case_4(v::AbstractVector{T}; tol=1e-8) where {T}
    # Try Heuristic solutions [Friberg 2021, Lemma 5.1]
    # vp = proj onto primal cone, vd = proj onto polar cone
    vp = SVector{3,T}(min(v[1], 0), zero(T), max(v[3], 0))
    vd = SVector{3,T}(zero(T), min(v[2], 0), min(v[3], 0))
    if v[2] > 0
        zp = max(v[3], v[2]*exp(v[1]/v[2]))
        if zp - v[3] < norm(vp - v)
            vp = SVector{3,T}(v[1], v[2], zp)
        end
    end
    if v[1] > 0
        zd = min(v[3], -v[1]*exp(v[2]/v[1] - 1))
        if v[3] - zd < norm(vd - v)
            vd = SVector{3,T}(v[1], v[2], zd)
        end
    end

    # Check if heuristics above approximately satisfy the optimality conditions
    opt_norm = norm(vp + vd - v)
    opt_ortho = abs(dot(vp, vd))
    if norm(v - vp) < tol || norm(v - vd) < tol || (opt_norm < tol && opt_ortho < tol)
        return vp
    end

    # Failure of heuristics -> non heuristic solution
    # Ref: https://docs.mosek.com/slides/2018/ismp2018/ismp-friberg.pdf, p47-48
    # Thm: h(x) is smooth, strictly increasing, and changes sign on domain
    r, s, t = v[1], v[2], v[3]
    h(x) = (((x-1)*r + s) * exp(x) - (r - x*s)*exp(-x))/(x^2 - x + 1) - t

    # Note: won't both be Inf by case 3 of projection
    lb = r > 0 ? 1 - s/r : -Inf
    ub = s > 0 ? r/s : Inf

    # Deal with ±Inf bounds
    if isinf(lb)
        lb = min(ub-0.125, -0.125)
        for _ in 1:10
            h(lb) < 0 && break
            ub = lb
            lb *= 2
        end
    end
    if isinf(ub)
        ub = max(lb+0.125, 0.125)
        for _ in 1:10
            h(ub) > 0 && break
            lb = ub
            ub *= 2
        end
    end

    # Check bounds
    if !(h(lb) < 0 && h(ub) > 0)
        error("Failure to find bracketing interval for exp cone projection.")
    end

    x = _bisection(h, lb, ub)
    if x === nothing
        error("Failure in root-finding for exp cone projection with boundaries ($lb, $ub).")
    end

    return ((x - 1) * r + s)/(x^2 - x + 1) * SVector{3,T}(x, 1, exp(x))
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.DualExponentialCone) where {T}

projection of vector `v` on the dual exponential cone
i.e. `Kexp^* = {(u,v,w) | u < 0, -u*exp(v/u) <= ew } U {(u,v,w)| u == 0, v >= 0, w >= 0}`.

References:
* [Proximal Algorithms, 6.3.4](https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf)
by Neal Parikh and Stephen Boyd.
* [Projection, presolve in MOSEK: exponential, and power cones](https://docs.mosek.com/slides/2018/ismp2018/ismp-friberg.pdf)
by Henrik Friberg
"""
function projection_on_set(d::DefaultDistance, v::AbstractVector{T}, ::MOI.DualExponentialCone) where {T}
    p = projection_on_set(d, -v, MOI.ExponentialCone())
    return SVector{3,T}(v[1] + p[1], v[2] + p[2], v[3] + p[3])
end

"""
    projection_on_set(::DefaultDistance, v::AbstractVector{T}, sets::Array{<:MOI.AbstractSet})

Projection onto `sets`, a product of sets
"""
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, sets::Array{<:MOI.AbstractSet}) where {T}
    length(v) == length(sets) || throw(DimensionMismatch("Mismatch between value and set"))
    return reduce(vcat, (projection_on_set(DefaultDistance(), v[i], sets[i]) for i in eachindex(sets)))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Zeros) where {T}

derivative of projection of vector `v` on zero cone i.e. K = {0}^n
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Zeros) where {T}
    return FillArrays.Zeros(length(v), length(v))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Reals) where {T}

derivative of projection of vector `v` on real cone i.e. K = R^n
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Reals) where {T}
    return FillArrays.Eye(length(v))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::T, ::MOI.EqualTo)
"""
function projection_gradient_on_set(::DefaultDistance, ::T, ::MOI.EqualTo) where {T}
    return zero(T)
end

function projection_gradient_on_set(::DefaultDistance, v::T, s::MOI.LessThan) where {T}
    return oneunit(T) * (v <= MOI.constant(s))
end

function projection_gradient_on_set(::DefaultDistance, v::T, s::MOI.GreaterThan) where {T}
    return oneunit(T) * (v >= MOI.constant(s))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonnegatives) where {T}

derivative of projection of vector `v` on Nonnegative cone i.e. K = R^n+
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonnegatives) where {T}
    y = (sign.(v) .+ one(T))/2
    return LinearAlgebra.Diagonal(y)
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonpositives) where {T}

derivative of projection of vector `v` on Nonpositives cone i.e. K = R^n-
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.Nonpositives) where {T}
    y = @. (-sign(v) + one(T))/2
    return LinearAlgebra.Diagonal(y)
end

"""
    projection_gradient_on_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, ::MOI.SecondOrderCone) where {T}

derivative of projection of vector `v` on second order cone i.e. K = {(t, x) ∈ R+ × Rn |  ||x|| ≤ t }
"""
function projection_gradient_on_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, ::MOI.SecondOrderCone) where {p,T}
    n = length(v)
    t = v[1]
    x = v[2:n]
    norm_x = LinearAlgebra.norm(x, p)
    if norm_x <= t
        return Matrix{T}(LinearAlgebra.I,n,n)
    elseif norm_x <= -t
        return zeros(T, n, n)
    end
    result = [
        norm_x     x';
        x          (norm_x + t)*Matrix{T}(LinearAlgebra.I,n-1,n-1) - (t/(norm_x^2))*(x*x')
    ]
    result ./= (2 * norm_x)
    return result
end

function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, cone::MOI.SecondOrderCone) where {T}
    return projection_gradient_on_set(NormedEpigraphDistance{2}(), v, cone)
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, cone::MOI.PositiveSemidefiniteConeTriangle) where {T}

derivative of projection of vector `v` on positive semidefinite cone i.e. K = S^n⨥
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.PositiveSemidefiniteConeTriangle) where {T}
    n = length(v)
    dim = isqrt(2n)
    X = unvec_symm(v, dim)
    λ, U = LinearAlgebra.eigen(X)
    Tp = promote_type(T, Float64)

    # if all the eigenvalues are >= 0
    if all(λi ≥ zero(λi) for λi in λ)
        return Matrix{Tp}(LinearAlgebra.I, n, n)
    end

    # k is the number of negative eigenvalues in X minus ONE
    k = count(λi < 1e-4 for λi in λ)

    y = zeros(Tp, n)
    D = zeros(Tp, n, n)

    for idx in 1:n
        # set eigenvector
        y[idx] = 1

        # defining matrix B
        X̃ = unvec_symm(y, dim)
        B = U' * X̃ * U

        for i in 1:size(B)[1] # do the hadamard product
            for j in 1:size(B)[2]
                if (i <= k && j <= k)
                    @inbounds B[i, j] = 0
                elseif (i > k && j <= k)
                    λpi = max(λ[i], zero(Tp))
                    λmj = -min(λ[j], zero(Tp))
                    @inbounds B[i, j] *= λpi / (λmj + λpi)
                elseif (i <= k && j > k)
                    λmi = -min(λ[i], zero(Tp))
                    λpj = max(λ[j], zero(Tp))
                    @inbounds B[i, j] *= λpj / (λmi + λpj)
                end
            end
        end
        @inbounds D[idx, :] = vec_symm(U * B * U')
        # reset eigenvector
        @inbounds y[idx] = 0
    end
    return D
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.ExponentialCone) where {T}

derivative of projection of vector `v` on closure of the exponential cone,
i.e. `cl(Kexp) = {(x,y,z) | y e^(x/y) <= z, y>0 } U {(x,y,z)| x <= 0, y = 0, z >= 0}`.

References:
* [Solution Refinement at Regular Points of Conic Problems](https://stanford.edu/~boyd/papers/cone_prog_refine.html)
by Enzo Busseti, Walaa M. Moursi, and Stephen Boyd
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, s::MOI.ExponentialCone) where {T}
    _check_dimension(v, s)

    if _in_exp_cone(v; dual=false)
        return SMatrix{3,3,T}(I)
    end
    if _in_exp_cone(-v; dual=true)
        # if in polar cone Ko = -K*
        return zeros(SMatrix{3,3,T})
    end
    if v[1] <= 0 && v[2] <= 0
        return @SMatrix(T[
            1 0 0
            0 0 0
            0 0 (v[3] >= 0)
        ])
    end

    z1, z2, z3 = _exp_cone_proj_case_4(v)
    nu = z3 - v[3]
    rs = z1/z2
    exp_rs = exp(rs)

    mat = inv(@SMatrix([
        1+nu*exp_rs/z2     -nu*exp_rs*rs/z2       0     exp_rs;
        -nu*exp_rs*rs/z2   1+nu*exp_rs*rs^2/z2    0     (1-rs)*exp_rs;
        0                  0                      1     -1
        exp_rs             (1-rs)*exp_rs          -1    0
    ]))
    return SMatrix{3,3}(@view(mat[1:3,1:3]))
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.DualExponentialCone) where {T}

derivative of projection of vector `v` on the dual exponential cone,
i.e. `Kexp^* = {(u,v,w) | u < 0, -u*exp(v/u) <= ew } U {(u,v,w)| u == 0, v >= 0, w >= 0}`.

References:
* [Solution Refinement at Regular Points of Conic Problems]
(https://stanford.edu/~boyd/papers/cone_prog_refine.html)
by Enzo Busseti, Walaa M. Moursi, and Stephen Boyd
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.DualExponentialCone) where {T}
    # from Moreau decomposition: x = P_K(x) + P_-K*(x)
    return I - projection_gradient_on_set(DefaultDistance(), -v, MOI.ExponentialCone())
end

"""
    projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, sets::AbstractVector{<:MOI.AbstractSet})

Derivative of the projection of vector `v` on product of `sets`
projection_gradient_on_set[i,j] = ∂projection_on_set[i] / ∂v[j] where `projection_on_set` denotes projection of `v` on `cone`

Find expression of projections on cones and their derivatives here: https://stanford.edu/~boyd/papers/pdf/cone_prog_refine.pdf
"""
function projection_gradient_on_set(::DefaultDistance, v::AbstractVector{T}, sets::AbstractVector{<:MOI.AbstractSet}) where {T}
    length(v) == length(sets) || throw(DimensionMismatch("Mismatch between value and set"))
    return BlockDiagonal([projection_gradient_on_set(DefaultDistance(), v[i], sets[i]) for i in eachindex(sets)])
end

"""
    projection_on_set(::DefaultDistance, V::AbstractVector{T}, s::NormBallNuclear{T}) where {T}

projection of matrix `V` onto nuclear norm ball
"""
function projection_on_set(d::DefaultDistance, V::AbstractMatrix{T}, s::NormNuclearBall{T}) where {T}
    U, sing_val, Vt = LinearAlgebra.svd(V)
    if (sum(sing_val) <= s.radius)
        return V
    end
    sing_val_proj = projection_on_set(d, sing_val, ProbabilitySimplex(length(sing_val), s.radius))
    return U * Diagonal(sing_val_proj) * Vt'
end

# initial implementation in FrankWolfe.jl
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::ProbabilitySimplex{T}) where {T}
    _check_dimension(v, s)
    # TODO: allocating a ton, should implement the recent non-sorting alg
    n = length(v)
    if sum(v) ≈ s.radius && all(>=(0), v)
        return v
    end
    rev = v .- maximum(v)
    u = sort(rev, rev=true)
    cssv = cumsum(u)
    rho = sum(eachindex(u)) do idx
        u[idx] * idx > (cssv[idx] - s.radius)
    end - 1
    theta = (cssv[rho+1] - s.radius) / (rho + 1)
    w = clamp.(rev .- theta, 0.0, Inf)
    return w
end

function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::StandardSimplex{T}) where {T}
    _check_dimension(v, s)
    n = length(v)
    if sum(v) ≤ s.radius && all(>=(0), v)
        return v
    end
    x = copy(v)
    sum_pos = zero(T)
    for idx in eachindex(x)
        if x[idx] < 0
            x[idx] = 0
        else
            sum_pos += x[idx]
        end
    end
    # at least one positive element
    if sum_pos > 0
        @. x = x / sum_pos * s.radius
    end
    return x
end

function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::NormInfinityBall{T}) where {T}
    if norm(v, Inf) <= s.radius
        return v
    end
    return clamp.(v, -s.radius, s.radius)
end

function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::NormTwoBall{T}) where {T}
    nv = norm(v)
    if nv <= s.radius
        return v
    end
    return v .* s.radius ./ nv
end

# inspired by https://github.com/MPF-Optimization-Laboratory/ProjSplx.jl
function projection_on_set(::DefaultDistance, v::AbstractVector{T}, s::NormOneBall{T}) where {T}
    n = length(v)
    if norm(v, 1) ≤ τ
        return v
    end
    u = abs.(v)
    # simplex projection
    bget = false
    s_indices = sortperm(u, rev=true)
    tsum = zero(τ)

    @inbounds for i in 1:n-1
        tsum += u[s_indices[i]]
        tmax = (tsum - τ) / i
        if tmax ≥ u[s_indices[i+1]]
            bget = true
            break
        end
    end
    if !bget
        tmax = (tsum + u[s_indices[n]] - τ) / n
    end

    @inbounds for i in 1:n
        u[i] = max(u[i] - tmax, 0)
        u[i] *= sign(v[i])
    end
    return u
end


"""
    src

A general module name. In this demo, it contains:
2. DiscreteMarkovChain, a type to define a discrete Markov chain
3. stationary_distribution, a function to solve stationary distribution of a `DiscreteMarkovChain`
4. approx_ar1, approximates an AR(1) process with a discrete Markov Chain
5. golden_section, Golden Section search
6. golden_section_improved, improved 3-point Golden Section search
7. DiscreteMarkovChain, discrete state Markov Chain
"""
module src
    import Distributions
    # ----------------------
    export approx_ar1
    export DiscreteMarkovChain, stationary_distribution
    export DiscreteTimeAR1
    export golden_section, golden_section_improved

# ------------------ TYPE: DISCRETE MARKOV CHAIN
"""
    DiscreteMarkovChain( P::Matrix{T} where T <: Real ; w::Vector = Array(1:size(T,1)) )

immutable type for discrete markov chains,
including:
1. K::Int, the number of states
2. w::Vector, state values, can be anything
2. P::Matrix{T} where T <: Real, transition matrix, square matrix, P[i,j], prob from state i to j
"""
mutable struct DiscreteMarkovChain <: Any
    K::Int  # number of states
    w::Vector # state values, can be anything
    P::Matrix{T} where T <: Real  # transition matrix
    function DiscreteMarkovChain( P::Matrix{T} where T <: Real ; w::Vector = Array(1:size(P,1)) )
        (size(P)[1] == size(P)[2])  ?  nothing  :  throw(ErrorException("non-square transition matrix received"))
        all(isapprox.( sum(P,dims = 2), 1.0))  ?  nothing  : throw(ErrorException("row-sum == 1 transition matrix received!"))
        # adjust P, making row-sum EXACT 1
        local newP = copy(P)
        for x in 1:size(P,1)
            newP[x,:] ./= sum(newP[x,:])
        end
        return new( size(P,1), w, newP )
    end
end
# -----
"""
    stationary_distribution( mc::DiscreteMarkovChain )

get the stationary distribution of a given discrete Markov Chain.
returns a `Vector{Float64}`, elements are in the order of `mc.w`.
"""
function stationary_distribution( mc::DiscreteMarkovChain )
    local tmpP = copy(mc.P)
    while true
        (sum(abs.(tmpP .- tmpP^2)) < 1E-8) ? break : nothing
        tmpP ^= 2
    end # while
    return tmpP[1,:]::Vector{Float64}
end # stationary_distribution

abstract type AbstractARProcess <: Any end
# ------------------------------- AR(1) process
"""
    DiscreteTimeAR1( ρ::Real, σ::Real ; Z̄::Real = 0.0, Z0::Union{Real,Distributions.UnivariateDistribution} = 0.0 )

An AR(1) process defined as: `` Z_{t+1} = (1-\\rho)\\cdot \\bar{Z} + \\rho Z_t + \\epsilon_t ``.
where `ρ` is the auto-reg coefficient, `Z̄` is the mean value, `σ` is the std of the error term ``\\epsilon_t``
which follows i.i.d. `N(0,σ)`, and `Z0` is the initial value/distribution of this process.
It can be a constant, or an instance of one of the subtypes of `Distributions.UnivariateDistribution`, e.g. `Normal()`.
"""
struct DiscreteTimeAR1 <: AbstractARProcess
    ρ::Real
    Z̄::Real
    Z0::Union{Real,Distributions.UnivariateDistribution}
    σ::Real
    # ---------------
    function DiscreteTimeAR1( ρ::Real, σ::Real ; Z̄::Real = 0.0, Z0::Union{Real,Distributions.UnivariateDistribution} = 0.0 )
        @assert(isfinite(ρ), "rho must be finite"); @assert(isfinite(Z̄), "the mean value must be finite")
        @assert(0<σ<Inf, "sigma must be non-negative and finite")
        isa(Z0, Real) ?  @assert(isfinite(Z0), "Z0 must be finite") : nothing
        # -----
        return new(ρ,Z̄,Z0,σ)
    end # AR1()
end # AR1

# ------------------ APPROXIMATE AR(1) WITH MARKOV CHAIN
"""
    approx_ar1(ρ::Real, σ::Real ; m::Int = 9, λ::Int = 3, Z̄::Real = 0.0 )

approximates a zero-mean AR(1) process with a discrete finite-state Markov Chain
whose states are equally divided.
The zero-mean AR(1) process is defined as:
`` Z_{t+1} = (1-\\rho)\\cdot \\bar{Z} + \\rho Z_t + \\epsilon_t ``,
where `ρ` is the auto-regression coefficient, `σ` is the standard error of ``\\epsilon_t``,
`m` is the number of the new Markov chain's states, `Z̄` is the mean of `Zt`.
`λ>0` is the absolute maximum deviation of the Markov chain's states,
it is defined as `max(z) = λσ`, where `z` is the state value of the new Markov Chain.

returns an instance of `DiscreteMarkovChain`.

This algorithm is from `(Tauchen, 1986)` and developed by `(Heer & Maussner, 2015)`.
In most cases, `m=9` and `λ=3` are enough to approximate an AR(1).
"""
function approx_ar1(ρ::Real, σ::Real ; m::Int = 9, λ::Real = 3, Z̄::Real = 0.0 )
    # assertions
    abs(ρ) < 1  ?  nothing : @warn("the AR(1) process is not stationary becuase abs(ρ) >=1")
    @assert(0 < σ < Inf, "the standard error must be greater than 0 and finite")
    @assert(1 < m, "there must be at least two states to make a Markov chain")
    @assert(0 < λ < Inf, "the absolute deviation must be greater than 0 and finite")
    @assert(isfinite(Z̄), "the mean value of Zt must be finite")
    # 1. states
    local w::Vector = zeros(m)
    w[1] = -λ * σ / √(1-ρ^2)  # minimum state value (lower bound, finite)
    local steplen::Real = -2 * w[1] / (m-1)  # step length between two states
    for x in 2:m; w[x] = w[x-1] + steplen; end;  # fill
    # 2. transition matrix
    local P::Matrix = zeros(m,m)
    local tmpNorm = Distributions.Normal() # std norm distrib
    for x in 1:m
        P[x,1] = Distributions.cdf(tmpNorm, (w[1]-ρ*w[x])/σ + 0.5 * steplen / σ )
        if m>2
            for y in 2:(m-1)
                P[x,y] = Distributions.cdf(tmpNorm, (w[y]-ρ*w[x])/σ + 0.5 * steplen / σ ) -
                    Distributions.cdf(tmpNorm, (w[y]-ρ*w[x])/σ - 0.5 * steplen / σ )
            end # for
        end # if
        P[x,m] = 1 - sum(P[x,1:(m-1)])
    end # x
    return DiscreteMarkovChain(P,w = w .+ Z̄)::DiscreteMarkovChain # add mean value to state
end # approx_ar1
# -------------------
approx_ar1(ar1::DiscreteTimeAR1 ; m::Int = 9, λ::Real = 3) = begin
    return approx_ar1(ar1.ρ, ar1.σ, m = m, λ = λ, Z̄ = ar1.Z̄)::DiscreteMarkovChain
end # approx_ar1

# ----------------------- Linear interpolation & expolation
"""
    linear_function_interpolation( x0::Real, xsamples::Vector{T} where T <: Real, fsamples::Vector{T} where T <: Real )

linear interpolation on one-parameter function `f(x)` which is defined by a vector of sample `x` points
and a vector of sample `f(x)` corresponding to sample `x`;
`xsamples` should be increasingly sorted.
returns the function value evaluated at the given point `x0`.
if `x0<findmin(x)` or `x0>findmax(x)`, this function uses the linear expolation
of minimum/maximum sample `x` and the second minimum/maximum sample `x`.
"""
function linear_function_interpolation( x0::Real, xsamples::Vector{T} where T <: Real, fsamples::Vector{T} where T <: Real )
    issorted(xsamples) ? nothing : throw(ErrorException("xsamples should be pre-sorted"))
    local N = length(fsamples); local xmax = xsamples[end]; local xmin = xsamples[1]
    local x0_location = (x0 - xmin) / (xmax - xmin) * (N - 1) + 1  # locating x0
    local nearest_xsample_loc = floor(Int,x0_location)  # the nearest sample x point (on the left side) to interpolate x0
    local dist_from_nearest_xsample = x0_location - nearest_xsample_loc # the fractional distance between x0 and the nearest x point
    # NOTE: relationship: x[near] < x0 < x[near+1]

    # case: if touching the upper bound
    if (x0 > xmax) | (nearest_xsample_loc == N)  # to prevent BoundsError at index N
        res = fsamples[end]
        return res::Float64
    # case: if touching the lower bound
    elseif (x0 <= xmin) | (nearest_xsample_loc == 1)
        res = fsamples[1] - (1 - x0_location) * ( fsamples[2] - fsamples[1] )
        return res::Float64
    # case: normal case
    else
        res = (1 - dist_from_nearest_xsample) * fsamples[nearest_xsample_loc] +
            dist_from_nearest_xsample * fsamples[nearest_xsample_loc + 1]
        return res::Float64
    end # if

end # linear_function_interpolation

# ------------------------------ conventional Golden Section for 1-parameter function
"""
    golden_section(f::Function, lb::Real, ub::Real ; atol::Real = 1E-8, maxiter::Int = 5000)

conventional Golden Section search for MINIMIZATION problem.
`f` must only receive one `Real` type location parameter;
`lb` is the finite lower bound of searching; `ub` is the finite upper bound of searching;
`atol` is the absolute error to converge,
and `maxiter` sets the maximum loops to prevent dead loop.
returns a `Tuple` (in order) consisting of optimal x, optimal `f` value, a flag indicating if converged,
and in which round the algorithm ends.

testing code:
```
g(x::Real) = (x-1)^2
BenchmarkTools.@benchmark tmp = src.golden_section(g,-1.0,3.0,atol=1E-12)
```

performance tips:
```
BenchmarkTools.Trial:
  memory estimate:  96 bytes
  allocs estimate:  1
  --------------
  minimum time:     144.933 ns (0.00% GC)
  median time:      150.576 ns (0.00% GC)
  mean time:        181.232 ns (7.65% GC)
  maximum time:     71.892 μs (99.47% GC)
  --------------
  samples:          10000
  evals/sample:     852
```
"""
function golden_section(f::Function, lb::Real, ub::Real ; atol::Real = 1E-8, maxiter::Int = 5000)
    @assert(-Inf < lb < ub < Inf, "lb must be less than ub, both must be finite")
    local goldnum = 0.618033988749895  # the golden number
    local bounds = [lb,ub]  # changable bounds to update
    local left_x = bounds[1] + (1-goldnum) * (bounds[2] - bounds[1])  # left trial
    local right_x = bounds[1] + goldnum * (bounds[2]-bounds[1])  # right trial
    # initial evaluation (NOTE: using type assertion to test if `f` only returns one `Real` value)
    local left_fval::Real = f(left_x)
    local right_fval::Real = f(right_x)
    # test if both trials are defined & finite (because NaN & Inf cannot be operated rationally)
    if isnan(left_fval) | isinf(left_fval) | isnan(left_fval) | isinf(left_fval)
        throw(DomainError("NaN or Inf function value(s) found in the initial evaluation of Golden section search"))
    end # if
    # begin search
    for j in 1:maxiter
        # check convergency
        if abs(bounds[2]-bounds[1]) < atol
            local finalx::Real = (bounds[1] + bounds[2]) / 2
            return ( finalx, f(finalx), true, j )::Tuple
        end # if
        # if not converge, go on to update
        if left_fval > right_fval
            bounds[1] = left_x; bounds[2] = bounds[2]; left_x = right_x
            right_x = bounds[1] + goldnum * (bounds[2] - bounds[1])
            left_fval = f(left_x)
            right_fval = f(right_x)
        else # i.e. if left_fval <= right_fvals
            bounds[1] = bounds[1]; bounds[2] = right_x; right_x = left_x
            left_x = bounds[1] + (1-goldnum) * (bounds[2] - bounds[1])
            left_fval = f(left_x)
            right_fval = f(right_x)
        end # if
    end # for j
    # if loop normally ends, it means that the algorithm did not converge
    local finalx::Real = (left_x+right_x)/2
    return ( finalx, f(finalx), false, maxiter )::Tuple
end # golden_section

# ---------------------------- improved Golden Section search
"""
    golden_section_improved(f::Function, LLB::Float64, LB::Float64, RB::Float64 ;  TOL::Float64=1E-08, ITER=300)

improved Golden Section search, using a third point x_mid between (x_low,x_high) to improve performance.
designed for one-dim function MAXIMIZATION problem

Par:
    1. f [annonymous func]: target to search (a maximization problem), it MUST HAVE ONLY ONE PARAMETER f(x)
    1. LLB [num]: the (initial) very left bound of searching (initial)
    1. LB [num]: the (initial) left bound of searching
    1. RB [num]: the (initial) right bound of searching
    1. TOL [num]: tolerance, 1E-4 or 1E-5 is enough
    1. ITER [int]: maximum iteration times in searching

Ret:
    1. Xmin [num]: the solution

Depend:
    1. func
"""
function golden_section_improved(f::Function, LLB::Real, LB::Real, RB::Real ;  TOL::Float64=1E-08)
    # -------- DATA PROCESS ----------
    # NOTE: in Julia, the golden numebr is an integrated const (before v0.6), but we redefine a golden number for case in v1.0 and later
    GoldenNumber = 1.618033988749895
    r1 = GoldenNumber - 1; r2 = 1 - r1
    x0 = LLB; x3 = RB # initia-lize bounds
    if abs(RB-LB)<=abs(LB-LLB)
        x1 = LB; x2 = LB+r2*(RB-LB)
    else
        x2 = LB; x1 = LB-r2*(LB-LLB)
    end
    # initialization of function value (and turns a maximization to a minimization)
    f1 = - f(x1); f2 = - f(x2)
    # Searching
    while true
        if f2<f1
            x0=x1 # Update the very lower bound
            x1=x2; x2=r1*x1+r2*x3 # new left golden position
            f1=f2; f2= - f(x2)
        else
            x3=x2; x2=x1; x1=r1*x2+r2*x0; f2=f1; f1= - f(x1)
        end # if
        if abs(x3-x0) > TOL * (abs(x1)+abs(x2))
            break
        end # if
    end # while
    # Post-convergence
    Xmin = f1<=f2 ? x1 : x2

    return Xmin::Float64
end

end # src
#


using SMM
using Optim
using HTTP
using GLM
using LinearAlgebra
using Random
using Statistics
using DataFrames
using DataFramesMeta
using CSV

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 1
#:::::::::::::::::::::::::::::::::::::::::::::::::::
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS1-julia-intro/nlsw88.csv"
df = CSV.read(HTTP.get(url).body,DataFrame)
X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
y = df.married.==1

#GMM
function ols_gmm(b, X, y)
    g = y .- X*b
    J = g'*I*g
    return J
end
res =  optimize(b -> ols_gmm(b, X, y), rand(size(X,2)), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000, show_trace=true))
println(res.minimizer)

#closed form solution
res1 = inv(X'*X)*X'*y
println(res1)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 2
#:::::::::::::::::::::::::::::::::::::::::::::::::::
df = dropmissing(df, :occupation)
df[df.occupation.==8 ,:occupation] .= 7
df[df.occupation.==9 ,:occupation] .= 7
df[df.occupation.==10,:occupation] .= 7
df[df.occupation.==11,:occupation] .= 7
df[df.occupation.==12,:occupation] .= 7
df[df.occupation.==13,:occupation] .= 7
X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
y = df.occupation

#MlE
function mlogit(alpha, X, y)

        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

        num = zeros(N,J)
        dem = zeros(N)
        for j=1:J
            num[:,j] = exp.(X*bigAlpha[:,j])
            dem .+= num[:,j]
        end

        P = num./repeat(dem,1,J)

        loglike = -sum( bigY.*log.(P) )

        return loglike
    end

alpha_zero = zeros(6*size(X,2))
alpha_rand = rand(6*size(X,2))
alpha_true = [.1910213,-.0335262,.5963968,.4165052,-.1698368,-.0359784,1.30684,-.430997,.6894727,-.0104578,.5231634,-1.492475,-2.26748,-.0053001,1.391402,-.9849661,-1.398468,-.0142969,-.0176531,-1.495123,.2454891,-.0067267,-.5382892,-3.78975]
alpha_start = alpha_true.*rand(size(alpha_true))
#println(size(alpha_true))
alpha_hat_optim = optimize(a -> mlogit(a, X, y), alpha_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
alpha_hat_mle = alpha_hat_optim.minimizer
println(alpha_hat_mle)

#GMM
function mlogit_gmm(alpha, X, y)
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

    num = zeros(N,J)
    dem = zeros(N)
    for j=1:J
        num[:,j] = exp.(X*bigAlpha[:,j])
        dem .+= num[:,j]
    end

    P = num./repeat(dem,1,J)

    g = bigY .-P
    J = sum(g'*I*g)
    return J
end

alpha_rand = rand(6*size(X,2))
alpha_hat_optim1 = optimize(a -> mlogit(a, X, y), alpha_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
alpha_hat_gmm = alpha_hat_optim1.minimizer
println(alpha_hat_gmm)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 3
#:::::::::::::::::::::::::::::::::::::::::::::::::::
function mlogit_sim(alpha,X,K,J,N)
    bigAlpha = [reshape(alpha,K,J-1) zeros(K)]
    num = zeros(N,J)
    dem = zeros(N)
    for j=1:J
        num[:,j] = exp.(X*bigAlpha[:,j])
        dem .+= num[:,j]
    end

    P = num./repeat(dem,1,J)
    epsilon = rand(N)

    Y = zeros(N)
    for i =1:N
        temp = zeros(J)
        for j = 1:J
            temp[j] = convert(Int,(sum(P[i,j:end])>epsilon[i]))
        end
        Y[i] = sum(temp)
    end
    return Y
end

K = 3
J = 7
N = 10000
X = randn((N,K))
alpha_true = rand((J-1)*K)
println("beta true:",alpha_true)
y_sim = mlogit_sim(alpha_true,X,K,J,N)
alpha_start = alpha_true.*rand(size(alpha_true))
alpha_hat_optim2 = optimize(a -> mlogit(a, X, y_sim), alpha_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
alpha_hat_sim = alpha_hat_optim2.minimizer
println("beta estimated:",alpha_hat_sim)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 4
#:::::::::::::::::::::::::::::::::::::::::::::::::::
MA = SMM.parallelNormal()
dc = SMM.history(MA.chains[1])
dc = dc[dc[:accepted].==true,:]
println(describe(dc))

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 5
#:::::::::::::::::::::::::::::::::::::::::::::::::::
function mlogit_smm(alpha, X, y, D)
    K = size(X,2)
    J = length(unique(y))
    N = length(y)

   # N+1 moments in both model and data
    gmodel = zeros(N+1,D)
    # data moments are just the y vector itself
    # and the variance of the y vector
    gdata  = vcat(y,var(y))

    Random.seed!(1234)

    # simulated model moments

    for d = 1:D
        y_sim = mlogit_sim(alpha,X,K,J,N)
        gmodel[1:end-1,d] = y_sim
        gmodel[  end  ,d] = var(y_sim)
    end
    # criterion function
    err = vec(gdata .- mean(gmodel; dims=2))
    # weighting matrix is the identity matrix
    # minimize weighted difference between data and moments
    J = err'*I*err

    return J
end

alpha_true = [.1910213,-.0335262,.5963968,.4165052,-.1698368,-.0359784,1.30684,-.430997,.6894727,-.0104578,.5231634,-1.492475,-2.26748,-.0053001,1.391402,-.9849661,-1.398468,-.0142969,-.0176531,-1.495123,.2454891,-.0067267,-.5382892,-3.78975]
alpha_start = alpha_true.*rand(size(alpha_true))
D = 2000
#println(size(alpha_true))
alpha_hat_optim = optimize(a -> mlogit_smm(a, X, y,D), alpha_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
alpha_hat_smm = alpha_hat_optim.minimizer
println(alpha_hat_smm)


"""
    hermite(x, p::Integer)

Evaluate the Hermite polynomial of degree `p` at `x` using the three term recursion.
"""
function hermite(x, p::Integer)
    T = typeof( 2x )
    p₀ = one(T)
    p₁ = 2x

    if p <= 0
        return p₀
    elseif p == 1
        return p₁
    end

    for n in 2:p
        p₀, p₁ = p₁, 2x*p₁ - 2*(n-1)*p₀
    end

    p₁
end


# Hello World
# variable substitution with the $ operator
word = "world"
output = "Hello $word !"
println(output)

# tab autocomplete of LateX characters
ℵ = "world"
output = "Hello $ℵ, again!"
println(output)

# this extends to function evaluations as well
α = 2
println("The sine of $α is $(sin(α))")

# sidenote:
# so if printing to a file you can directly direference the evalutioan
# without passing it to an array which can then be printed.

typeof(α)
typeof(sin)
methods(sin)

# built in type promotion
check = int8(10)
typemax(Int8)
typeof(check)
check = check *1232
typeof(check)

# defining numerical functions
# looks and feels very much like real math

f(x) = 2*x + x^2 + 5
f(3.4)

g(x, y) = 34 + 2*x^2 + y^3
g(2, 3)

# passing function results to arrays
arr = Float64[]
for x in 1.3:0.01:3.9 push!(arr, f(x)) end
println(arr)
for x in 1.3:0.01:9.9 push!(arr, f(x)) end
println(arr)

arr2 = [f(i) for i in [1:2:10, 15]]
println(arr2)

using ASCIIPlots
lineplot([1.3:0.01:3.9], [1.3:0.01:3.9])
lineplot([1:2:10], [f(i) for i in [1:2:10]])

# working with imaginary numbers
a = 23
b = (1 + 34im)
c = a+b

typeof(a)
typeof(b)
typeof(c)

# now c can be dereferenced into real and imaginary parts
d = c.re
e = c.im

typeof(d)
typeof(e)

d = c*conj(c)
typeof(d)


"""
```
find_threshold(Entropy(),counts,edges)
```
An algorithm for finding the threshold value for a gray-level histogram using
the entropy of the histogram.

# Output

Returns the point in the `AbstractRange` which corresponds to the threshold bin in the histogram.

# Details

This algorithm uses the entropy of a gray level histogram to produce a threshold
value.

Let ``f_1, f_2, \\ldots, f_I`` be the frequencies in the various bins of the
histogram and ``I`` the number of bins. With ``N = \\sum_{i=1}^{I}f_i``, let
``p_i = \\frac{f_i}{N}`` (``i = 1, \\ldots, I``) denote the probability
distribution of gray levels. From this distribution one derives two additional
distributions. The first defined for discrete values ``1`` to ``s`` and the
other, from ``s+1`` to ``I``. These distributions are

```math
A: \\frac{p_1}{P_s}, \\frac{p_2}{P_s}, \\ldots, \\frac{p_s}{P_s}
\\quad \\text{and} \\quad
B: \\frac{p_{s+1}}{1-P_s}, \\ldots, \\frac{p_n}{1-P_s}
\\quad \\text{where} \\quad
P_s = \\sum_{i=1}^{s}p_i.
```
The entropies associated with each distribution are as follows:

```math
H(A) = \\ln(P_s) + \\frac{H_s}{P_s}
```
```math
H(B) = \\ln(1-P_s) + \\frac{H_n-H_s}{1-P_s}
```
```math
\\quad \\text{where} \\quad
H_s = -\\sum_{i=1}^{s}p_i\\ln{p_i}
\\quad \\text{and} \\quad
H_n = -\\sum_{i=1}^{I}p_i\\ln{p_i}.
```

Combining these two entropy functions we have

```math
\\psi(s) = \\ln(P_s(1-P_s)) + \\frac{H_s}{P_s} + \\frac{H_n-H_s}{1-P_s}.
```
Finding the discrete value ``s`` which maximises the function ``\\psi(s)`` produces
the sought-after threshold value (i.e. the bin which determines the threshold).

See Section 4 of [1] for more details on the derivation of the entropy.

# Options

## Choices for `counts`
You can specify an `AbstractArray` which should be a 1D array of frequencies for
a histogram. You should submit the corresponding `edges` range for the bins of
the histogram. The function will throw an error if it detects that the `edges`
and `counts` have different lengths.

## Choices for `edges`
You can specify an `AbstractRange` which should be the corresponding range for
the bins of the histogram array passed into `counts`.

# Example

```julia

using TestImages, Images

img = testimage("cameraman")
# building a histogram with 256 bins
edges, counts = build_histogram(img, 256)
#=
  The `counts` array stores at index 0 the frequencies that were below the
  first bin edge. Since we are seeking a threshold over the interval
  partitioned by `edges` we need to discard the first bin in `counts`
  so that the dimensions of `edges` and `counts` match.
=#
find_threshold(EntropyThresholding(), counts[1:end], edges)
```

# References
[1] J. N. Kapur, P. K. Sahoo, and A. K. C. Wong, “A new method for gray-level picture thresholding using the entropy of the histogram,” *Computer Vision, Graphics, and Image Processing*, vol. 29, no. 1, p. 140, Jan. 1985.[doi:10.1016/s0734-189x(85)90156-2](https://doi.org/10.1016/s0734-189x%2885%2990156-2)
"""
function find_threshold(algorithm::Entropy, counts::AbstractArray, edges::AbstractRange)
    if length(edges) != length(counts)
        error("the lengths of edges and counts must match")
    end
    nbins = length(counts)
    pdf = counts ./ sum(counts)
    Ψₘₐₓ = 0
    Hₛ = H(pdf)
    Hₙ = -1*Hₛ[nbins]
    Pₛ = cumsum(pdf)
    threshold_bin = 1
    for i in 1:nbins
        Pᵢ = Pₛ[i]
        Hᵢ = -1*Hₛ[i]
        Ψᵢ = log(Pᵢ*(1-Pᵢ)) + (Hᵢ/Pᵢ) + ((Hₙ-Hᵢ)/(1-Pᵢ))
        if Ψᵢ > Ψₘₐₓ
            Ψₘₐₓ = Ψᵢ
            threshold_bin = i
        end
    end
    edges[threshold_bin]
end

function H(pdf::AbstractArray)
    h_dist = zeros(length(pdf))
    for i in 1:length(pdf)
        h_dist[i] = (pdf[i] != 0) ? pdf[i]*log(pdf[i]) : 0.0
    end
    cumsum(h_dist)
end


"""
    local_clustering_coefficient(g, v)
    local_clustering_coefficient(g, vs)

Return the [local clustering coefficient](https://en.wikipedia.org/wiki/Clustering_coefficient)
for node `v` in graph `g`. If a list of vertices `vs` is specified, return a vector
of coefficients for each node in the list.
"""
function local_clustering_coefficient(g::AbstractGraph, v::Integer)
    ntriang, alltriang = local_clustering(g, v)

    return alltriang == 0 ? 0. : ntriang * 1.0 / alltriang
end
local_clustering_coefficient(g::AbstractGraph, vs = vertices(g)) =
    [local_clustering_coefficient(g, v) for v in vs]

@doc_str """
    local_clustering(g, v)
    local_clustering(g, vs)

Return a tuple `(a, b)`, where `a` is the number of triangles in the neighborhood
of `v` and `b` is the maximum number of possible triangles. If a list of vertices
`vs` is specified, return two vectors representing the number of triangles and
the maximum number of possible triangles, respectively, for each node in the list.

This function is related to the local clustering coefficient `r` by ``r=\frac{a}{b}``.
"""
function local_clustering(g::AbstractGraph, v::Integer)
    k = degree(g, v)
    k <= 1 && return (0, 0)
    neighs = neighbors(g, v)
    c = 0
    for i in neighs, j in neighs
        i == j && continue
        if has_edge(g, i, j)
            c += 1
        end
    end
    return is_directed(g) ? (c , k*(k-1)) : (div(c,2) , div(k*(k-1),2))
end
function local_clustering(g::AbstractGraph, vs = vertices(g))
    ntriang = zeros(Int, length(vs))
    nalltriang = zeros(Int, length(vs))
    i = 0
    for (i, v) in enumerate(vs)
        ntriang[i], nalltriang[i] = local_clustering(g, v)
    end
    return ntriang, nalltriang
end

"""
    triangles(g[, v])
    triangles(g, vs)

Return the number of triangles in the neighborhood of node `v` in graph `g`.
If a list of vertices `vs` is specified, return a vector of number of triangles
for each node in the list. If no vertices are specified, return the number
of triangles for each node in the graph.
"""
triangles(g::AbstractGraph, v::Integer) = local_clustering(g, v)[1]
triangles(g::AbstractGraph, vs = vertices(g)) = local_clustering(g, vs)[1]

"""
    global_clustering_coefficient(g)

Return the [global clustering coefficient](https://en.wikipedia.org/wiki/Clustering_coefficient)
of graph `g`.
"""
function global_clustering_coefficient(g::AbstractGraph)
    c = 0
    ntriangles = 0
    for v in vertices(g)
        neighs = neighbors(g, v)
        for i in neighs, j in neighs
            i == j && continue
            if has_edge(g, i, j)
                c += 1
            end
        end
        k = degree(g, v)
        ntriangles += k*(k-1)
    end
    ntriangles == 0 && return 1.
    return c / ntriangles
end


function B(; T=T, λ=λ)

    k_B = 1.380649e-23
    h   = 6.62607015e-34
    c   = 2.99792458e8

    aux1 = (2.0 * h * c^2) / λ^5

    aux2 = (h * c) / (λ * k_B * T)

    aux1 / (exp(aux2) - 1.0)

end

function scaledB(; s=s, T=T, λ=λ)

    s * B(T=T, λ=λ)

end

function fitgalaxy()

    wave = [433,550,790] * 10^-9

    Y = [4.18, 8.60, 20.83]

    # wave = [0.5;4;6] * 1000.0 * 1e-9
    # Y = 1e-8*[B(T=5555.5,λ = 500*1e-9); B(T=5555.5,λ = 4000*1e-9);B(T=5555.5,λ = 6000*1e-9)]

    function loss(T)

        local V = map( w -> B(T=T,λ=w), wave)

        local s = dot(V,Y) / dot(V,V)

        local l = 0.0

        for i in 1:length(Y)
            l += (Y[i] - s*V[i])^2
        end

        return l

    end

    result = optimize(loss, 10, 10_000)

    @show T = result.minimizer

    V = map( w -> B(T=T, λ=w), wave)

    s = dot(V,Y) / dot(V,V)

    λ = collect(LinRange(100, 10_000, 100)) * 1e-9;

    plot(λ, map(λ -> s*B(T=T,λ=λ), λ), "g.")

    plot(wave, Y, "ko")

end

function galaxyposterior()

    Tmin = 1000.0
    Tmax = 10_000.0
    σ(x) = 1.0 / (exp(-x) + 1)

    wave = [433,550,790] * 10^-9

    Y = [4.18, 8.60, 20.83]

    # wave = [0.5;4;6] * 1000.0 * 1e-9
    # Y = 1e-8*[B(T=5555.5,λ = 500*1e-9); B(T=5555.5,λ = 4000*1e-9);B(T=5555.5,λ = 6000*1e-9)]

    function loss(T)

        local V = map( w -> B(T=T,λ=w), wave)

        local s = dot(V,Y) / dot(V,V)

        local l = 0.0

        for i in 1:length(Y)
            l += (Y[i] - s*V[i])^2
        end

        return l

    end

    Trange = 1000:20:8_000

    negloss = map(x -> -loss(x), Trange)

    plot(Trange, negloss, ".")

    Trange[argmax(negloss)], maximum(negloss)

end

# λ = collect(LinRange(10,10_000,1000)) * 1e-9;
# plot(λ, map(λ->B(T=4000.0,λ=λ), λ), "g.")
# plot(λ, map(λ->B(T=5000.0,λ=λ), λ), "b.")


module CoordinateSystem

using LinearAlgebra

export CSys

struct CSys
    O::Vector{Float64}
    x::Vector{Float64}
    y::Vector{Float64}
    z::Vector{Float64}
    T::Matrix{Float64} #transorm_matrix
end

function CSys(o::Vector{Float64},p₁::Vector{Float64},p₂::Vector{Float64})
    v₁=p₁-o
    v₂=p₂-o
    if abs(v₁⋅v₂/norm(v₁)/norm(v₂))==1
        error("Two vectors should not be parallel!")
    end
    x=v₁/norm(v₁)
    z=v₁×v₂
    z=z/norm(z)
    y=z×x
    T=[reshape(x,1,3);
        reshape(y,1,3);
        reshape(z,1,3)]
    CSys(o,x,y,z,T)
end

end


#=  https://www.damtp.cam.ac.uk/research/dd/files/teaching/MB17Gog.pdf#subsubsection.1.3.4  =#
using Plots

δt = 0.001
T = 1
times = 0:δt:T
ntimes = length(times)

N = 10000

S = zeros(size(times))
S[1] = N - 1
I = zeros(size(times))
I[1] = 1
R = zeros(size(times))

# Rate of infection
β = 0.005
# Rate of recovery
ν = 4

# Randomness for σ > 0
σ = 1

for i ∈ 2:ntimes
    δS = - β * I[i-1] * S[i-1] * δt
	δR = ν * I[i-1] * δt
	δI = - δS - δR
    S[i] = max(S[i-1] + δS * exp(σ * randn()), 0)
    I[i] = max(I[i-1] + δI * exp(σ * randn()), 0)
	R[i] = max(R[i-1] + δR * exp(σ * randn()), 0)
	scale = (S[i] + I[i] + R[i])/N
	S[i] /= scale; I[i] /= scale; R[i] /= scale
end

lineplot = plot(times, [S I R], label = ["Susceptible" "Infectious" "Recovered"], color = [:gold :red :blue]);
areaplot_ = areaplot(times, [S I R], label = ["Susceptible" "Infectious" "Recovered"], color = [:gold :red :blue]);

display(lineplot)
display(areaplot_)


# # [Lotka-Volterra Model](@id predation_example)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/predation/lotka-volterra.ipynb)

using AlgebraicPetri

using OrdinaryDiffEq
using Plots

using Catlab
using Catlab.Graphics
using Catlab.WiringDiagrams
using Catlab.CategoricalAlgebra
using Catlab.Programs.RelationalPrograms

display_uwd(ex) = to_graphviz(ex, box_labels=:name, junction_labels=:variable, edge_attrs=Dict(:len=>".75"));

# #### Step 1: Define the building block Petri nets needed to construct the model

birth_petri = Open(PetriNet(1, 1=>(1,1)));
Graph(birth_petri)
#-
predation_petri = Open(PetriNet(2, (1,2)=>(2,2)));
Graph(predation_petri)
#-
death_petri = Open(PetriNet(1, 1=>()));
Graph(death_petri)

# #### Step 2: Generate models using a relational syntax

lotka_volterra = @relation (wolves, rabbits) begin
  birth(rabbits)
  predation(rabbits, wolves)
  death(wolves)
end
display_uwd(lotka_volterra)
#-
lv_dict = Dict(:birth=>birth_petri, :predation=>predation_petri, :death=>death_petri);
lotka_petri = apex(oapply(lotka_volterra, lv_dict))
Graph(lotka_petri)

# Generate appropriate vector fields, define parameters, and visualize solution

u0 = [100, 10];
p = [.3, .015, .7];
prob = ODEProblem(vectorfield(lotka_petri),u0,(0.0,100.0),p);
sol = solve(prob,Tsit5(),abstol=1e-8);
plot(sol)

# #### Step 3: Extend your model to handle more complex phenomena
# such as a small food chain between little fish, big fish, and sharks

dual_lv = @relation (fish, Fish, Shark) begin
  birth(fish)
  predation(fish, Fish)
  death(Fish)
  predation(Fish, Shark)
  death(Shark)
end
display_uwd(dual_lv)
#-
dual_lv_petri = apex(oapply(dual_lv, lv_dict))
Graph(dual_lv_petri)

# Generate a new solver, provide parameters, and analyze results

u0 = [100, 10, 2];
p = [.3, .015, .7, .017, .35];
prob = ODEProblem(vectorfield(dual_lv_petri),u0,(0.0,100.0),p);
sol = solve(prob,Tsit5(),abstol=1e-6);
plot(sol)