# ============================================================
# PROBABILISTIC PREDICTIONS

# -----------------------------------------------------
# LogLoss

struct LogLoss{R <: Real} <: Measure
    tol::R
end
LogLoss(;eps=eps(), tol=eps) = LogLoss(tol)

metadata_measure(LogLoss;
                 instances                = ["log_loss", "cross_entropy"],
                 target_scitype           = Vec{<:Finite},
                 prediction_type          = :probabilistic,
                 orientation              = :loss,
                 reports_each_observation = true,
                 is_feature_dependent     = false,
                 supports_weights         = false,
                 distribution_type        = UnivariateFinite)

const CrossEntropy = LogLoss
@create_aliases LogLoss

@create_docs(LogLoss,
body=
"""
Since the score is undefined in the case that the true observation is
predicted to occur with probability zero, probablities are clipped
between `tol` and `1-tol`, where `tol` is a constructor key-word
argument.

If `sᵢ` is the predicted probability for the true class `yᵢ` then
the score for that example is given by

    -log(clamp(sᵢ, tol), 1 - tol)

A score is reported for every observation.
""",
scitype=DOC_FINITE)

# for single observation:
_cross_entropy(d::UnivariateFinite{S,V,R,P}, y, tol) where {S,V,R,P} =
    -log(clamp(pdf(d, y), P(tol), P(1) - P(tol)))

# multiple observations:
function (c::LogLoss)(ŷ::Vec{<:UnivariateFinite},
                           y::Vec)
    check_dimensions(ŷ, y)
    check_pools(ŷ, y)
    return broadcast(_cross_entropy, ŷ, y, c.tol)
end
# performant in case of UnivariateFiniteArray:
function (c::LogLoss)(ŷ::UnivariateFiniteVector{S,V,R,P},
                           y::Vec) where {S,V,R,P}
    check_dimensions(ŷ, y)
    check_pools(ŷ, y)
    return -log.(clamp.(broadcast(pdf, ŷ, y), P(c.tol), P(1) - P(c.tol)))
end

# -----------------------------------------------------
# BrierScore

struct BrierScore <: Measure end

metadata_measure(BrierScore;
                 human_name = "Brier score (a.k.a. quadratic score)",
                 instances                = ["brier_score",],
                 target_scitype           = Vec{<:Finite},
                 prediction_type          = :probabilistic,
                 orientation              = :score,
                 reports_each_observation = true,
                 is_feature_dependent     = false,
                 supports_weights         = true,
                 distribution_type        = UnivariateFinite)

@create_aliases BrierScore

@create_docs(BrierScore,
body=
"""
If `p(y)` is the predicted probability for a *single*
observation `y`, and `C` all possible classes, then the corresponding
Brier score for that observation is given by

``2p(y) - \\left(\\sum_{η ∈ C} p(η)^2\\right) - 1``

*Warning.* `BrierScore()` is a "score" in the sense that bigger is
better (with `0` optimal, and all other values negative). In Brier's
original 1950 paper, and many other places, it has the opposite sign,
despite the name. Moreover, the present implementation does not treat
the binary case as special, so that the score may differ, in that
case, by a factor of two from usage elsewhere.
""",
scitype=DOC_FINITE)

# calling on single observations (no checks):
function _brier_score(d::UnivariateFinite{S,V,R,P}, y) where {S,V,R,P}
    levels = classes(d)
    pvec = broadcast(pdf, d, levels)
    offset = P(1) + sum(pvec.^2)
    return P(2) * pdf(d, y) - offset
end

# calling on multiple observations:
function (::BrierScore)(ŷ::Vec{<:UnivariateFinite},
                        y::Vec,
                        w::Union{Nothing,Vec{<:Real}}=nothing)
    check_dimensions(ŷ, y)
    w === nothing || check_dimensions(w, y)

    check_pools(ŷ, y)
    unweighted = broadcast(_brier_score, ŷ, y)

    if w === nothing
        return unweighted
    end
    return w.*unweighted
end

# Performant version in case of UnivariateFiniteArray:
function (::BrierScore)(
    ŷ::UnivariateFiniteVector{S,V,R,P},
    y::Vec,
    w::Union{Nothing,Vec{<:Real}}=nothing) where {S,V,R,P<:Real}

    check_dimensions(ŷ, y)
    w === nothing || check_dimensions(w, y)

    isempty(y) && return P(0)

    check_pools(ŷ, y)

    probs = pdf(ŷ, classes(first(ŷ)))
    offset = P(1) .+ vec(sum(probs.^2, dims=2))

    unweighted = P(2) .* broadcast(pdf, ŷ, y) .- offset

    if w === nothing
        return unweighted
    end
    return w.*unweighted
end

# -----------------------------------------------------
# BrierLoss

struct BrierLoss <: Measure end

metadata_measure(BrierLoss;
                 human_name = "Brier loss (a.k.a. quadratic loss)",
                 instances                = ["brier_loss",],
                 target_scitype           = Vec{<:Finite},
                 prediction_type          = :probabilistic,
                 orientation              = :score,
                 reports_each_observation = true,
                 is_feature_dependent     = false,
                 supports_weights         = true,
                 distribution_type        = UnivariateFinite)

@create_aliases BrierLoss

@create_docs(BrierLoss,
body=
"""
If `p(y)` is the predicted probability for a *single*
observation `y`, and `C` all possible classes, then the corresponding
Brier score for that observation is given by

``\\left(\\sum_{η ∈ C} p(η)^2\\right) - 2p(y) + 1``

*Warning.* In Brier's original 1950 paper, what is implemented here is
called a "loss". It is, however, a "score" in the contemporary use of
that term: smaller is better (with `0` optimal, and all other values
positive).  Note also the present implementation does not treat the
binary case as special, so that the loss may differ, in that case, by
a factor of two from usage elsewhere.
""",
scitype=DOC_FINITE)

# calling on single observations (no checks):
function _brier_loss(d::UnivariateFinite{S,V,R,P}, y) where {S,V,R,P}
    levels = classes(d)
    pvec = broadcast(pdf, d, levels)
    offset = P(1) + sum(pvec.^2)
    return P(2) * pdf(d, y) - offset
end

(m::BrierLoss)(ŷ::Vec{<:UnivariateFinite},
               y::Vec,
               w::Union{Nothing,Vec{<:Real}}=nothing) =
                   - brier_score(ŷ, y, w)

const INVARIANT_LABEL =
    "This metric is invariant to class reordering."
const VARIANT_LABEL =
    "This metric is *not* invariant to class re-ordering"

# =============================================================
# DETERMINISTIC FINITE PREDICTIONS

# ---------------------------------------------------
# misclassification rate

struct MisclassificationRate <: Measure end

metadata_measure(MisclassificationRate;
                 instances  = ["misclassification_rate", "mcr"],
                 target_scitype           = Vec{<:Finite},
                 prediction_type          = :deterministic,
                 orientation              = :loss,
                 reports_each_observation = false,
                 is_feature_dependent     = false,
                 supports_weights         = true)

const MCR = MisclassificationRate
@create_aliases MCR

@create_docs(MisclassificationRate,
body=
"""
A confusion matrix can also be passed as argument.
$INVARIANT_LABEL
""",
scitype=DOC_FINITE)

# calling behaviour:
(::MCR)(ŷ::Vec{<:CategoricalValue},
        y::Vec{<:CategoricalValue}) = mean(y .!= ŷ)
(::MCR)(ŷ::Vec{<:CategoricalValue},
        y::Vec{<:CategoricalValue},
        w::Vec{<:Real}) = sum((y .!= ŷ) .* w) / length(y)
(::MCR)(cm::ConfusionMatrixObject) = 1.0 - sum(diag(cm.mat)) / sum(cm.mat)

# -------------------------------------------------------------
# accuracy

struct Accuracy <: Measure end

metadata_measure(Accuracy;
                 instances = ["accuracy",],
                 target_scitype           = Vec{<:Finite},
                 prediction_type          = :deterministic,
                 orientation              = :score,
                 reports_each_observation = false,
                 is_feature_dependent     = false,
                 supports_weights         = true),

@create_aliases Accuracy

@create_docs(Accuracy,
body=
"""
Accuracy is proportion of correct predictions `ŷ[i]` that match the
ground truth `y[i]` observations. $INVARIANT_LABEL
""",
scitype=DOC_FINITE)

# calling behaviour:
(::Accuracy)(args...) = 1.0 - misclassification_rate(args...)
(::Accuracy)(m::ConfusionMatrixObject) = sum(diag(m.mat)) / sum(m.mat)

# -----------------------------------------------------------
# balanced accuracy

struct BalancedAccuracy <: Measure end

metadata_measure(BalancedAccuracy;
                 instances = ["balanced_accuracy", "bacc", "bac"],
                 target_scitype           = Vec{<:Finite},
                 prediction_type          = :deterministic,
                 orientation              = :score,
                 reports_each_observation = false,
                 is_feature_dependent     = false,
                 supports_weights         = true)

const BACC = BalancedAccuracy
@create_aliases BACC

@create_docs(BalancedAccuracy,
body=
"""
Balanced accuracy compensates standard [`Accuracy`](@ref) for class imbalance.
See [https://en.wikipedia.org/wiki/Precision_and_recall#Imbalanced_data](https://en.wikipedia.org/wiki/Precision_and_recall#Imbalanced_data).
$INVARIANT_LABEL
""",
scitype=DOC_FINITE)

# calling behavior:
function (::BACC)(ŷ::Vec{<:CategoricalValue},
                  y::Vec{<:CategoricalValue})
    class_count = Dist.countmap(y)
    ŵ = 1.0 ./ [class_count[yi] for yi in y]
    return sum( (ŷ .== y) .* ŵ ) / sum(ŵ)
end

function (::BACC)(ŷ::Vec{<:CategoricalValue},
                  y::Vec{<:CategoricalValue},
                  w::Vec{<:Real})
    levels_ = levels(y)
    ŵ = similar(w)
    @inbounds for i in eachindex(w)
        ŵ[i] = w[i] / sum(w .* (y .== y[i]))
    end
    return sum( (ŷ .== y) .* ŵ ) / sum(ŵ)
end

# ==================================================================
## DETERMINISTIC BINARY PREDICTIONS - ORDER-INDEPENDENT

# ------------------------------------------------------------------
# Matthew's correlation

struct MatthewsCorrelation <: Measure end

metadata_measure(MatthewsCorrelation;
                 instances = ["matthews_correlation", "mcc"],
                 target_scitype           = Vec{<:Finite{2}},
                 prediction_type          = :deterministic,
                 orientation              = :score,
                 reports_each_observation = false,
                 is_feature_dependent     = false,
                 supports_weights         = false)
const MCC = MatthewsCorrelation
@create_aliases MCC

@create_docs(MatthewsCorrelation,
body=
"""
[https://en.wikipedia.org/wiki/Matthews_correlation_coefficient](https://en.wikipedia.org/wiki/Matthews_correlation_coefficient)
$INVARIANT_LABEL
""",
scitype=DOC_FINITE_BINARY)

# calling behaviour:
function (::MCC)(cm::ConfusionMatrixObject{C}) where C
    # http://rk.kvl.dk/introduction/index.html
    # NOTE: this is O(C^3), there may be a clever way to
    # speed this up though in general this is only used for low  C
    num = 0
    @inbounds for k in 1:C, l in 1:C, m in 1:C
        num += cm[k,k] * cm[l,m] - cm[k,l] * cm[m,k]
    end
    den1 = 0
    den2 = 0
    @inbounds for k in 1:C
        a = sum(cm[k, :])
        b = sum(cm[setdiff(1:C, k), :])
        den1 += a * b
        a = sum(cm[:, k])
        b = sum(cm[:, setdiff(1:C, k)])
        den2 += a * b
    end
    mcc = num / sqrt(den1 * den2)

    isnan(mcc) && return 0
    return mcc
end

(m::MCC)(ŷ::Vec{<:CategoricalValue},
         y::Vec{<:CategoricalValue}) =
             _confmat(ŷ, y, warn=false) |> m

# ---------------------------------------------------------
# AreaUnderCurve

#. Implementation drawn from
# https://www.ibm.com/developerworks/community/blogs/jfp/entry/Fast_Computation_of_AUC_ROC_score?lang=en
# but this link is now broken. Author contacted here:
# https://www.kaggle.com/c/microsoft-malware-prediction/discussion/76013.

struct AreaUnderCurve <: Measure end

metadata_measure(AreaUnderCurve;
                 human_name = "area under the ROC",
                 instances = ["area_under_curve", "auc"],
                 target_scitype           = Vec{<:Finite{2}},
                 prediction_type          = :probabilistic,
                 orientation              = :score,
                 reports_each_observation = false,
                 is_feature_dependent     = false,
                 supports_weights         = false)

const AUC = AreaUnderCurve
@create_aliases AreaUnderCurve

@create_docs(AreaUnderCurve,
body=
"""
Returns the area under the ROC ([receiver operator
characteristic](https://en.wikipedia.org/wiki/Receiver_operating_characteristic))
$INVARIANT_LABEL
""",
scitpye = DOC_FINITE_BINARY)

# core algorithm:
function _auc(::Type{P}, ŷ::Vec{<:UnivariateFinite},
              y::Vec) where P<:Real # type of probabilities
    lab_pos = classes(first(ŷ))[2] # 'positive' label
    scores  = pdf.(ŷ, lab_pos)     # associated scores
    y_sort  = y[sortperm(scores)]  # sort by scores
    n       = length(y)
    n_neg   = 0  # to keep of the number of negative preds
    auc     = P(0)
    @inbounds for i in 1:n
        # y[i] == lab_p --> it's a positive label in the ground truth
        # in that case increase the auc by the cumulative sum
        # otherwise increase the number of negatives by 1
        δ_auc, δ_neg = ifelse(y_sort[i] == lab_pos, (n_neg, 0), (0, 1))
        auc   += δ_auc
        n_neg += δ_neg
    end
    n_pos = n - n_neg
    return auc / (n_neg * n_pos)
end

# calling behaviour:
(::AUC)(ŷ::Vec{<:UnivariateFinite}, y::Vec) = _auc(Float64, ŷ, y)

# performant version for UnivariateFiniteVector:
(::AUC)(ŷ::Vec{<:UnivariateFinite{S,V,R,P}}, y::Vec) where {S,V,R,P} =
    _auc(P, ŷ, y)

# ==========================================================================
# DETERMINISTIC BINARY PREDICTIONS - ORDER DEPENDENT

const CM2 = ConfusionMatrixObject{2}

# --------------------------------------------------------------------------
# FScore

struct FScore{T<:Real} <: Measure
    β::T
    rev::Union{Nothing,Bool}
end

FScore(; β=1.0, rev=nothing) = FScore(β, rev)

metadata_measure(FScore;
                 human_name = "F-Score",
                 instances = ["f1score",],
                 target_scitype           = Vec{<:Finite{2}},
                 prediction_type          = :deterministic,
                 orientation              = :score,
                 reports_each_observation = false,
                 is_feature_dependent     = false,
                 supports_weights         = false)

@create_aliases FScore

@create_docs(FScore,
body=
"""
This is the one-parameter generalization, ``F_β``, of the F-measure or
balanced F-score.

[https://en.wikipedia.org/wiki/F1_score](https://en.wikipedia.org/wiki/F1_score)

Constructor signature: `FScore(; β=1.0, rev=true)`.

By default, the second element of `levels(y)` is designated as
`true`. To reverse roles, specify `rev=true`.
""",
scitype=DOC_ORDERED_FACTOR_BINARY,
footer="Constructor signature: `FScore(β=1.0, rev=false)`. ")

# calling on conf matrix:
function (score::FScore)(m::CM2)
    β = score.β
    β2   = β^2
    prec = precision(m)
    rec  = recall(m)
    return (1 + β2) * (prec * rec) / (β2 * prec + rec)
end

# calling on vectors:
(m::FScore)(ŷ, y) = _confmat(ŷ, y; rev=m.rev) |> m

# -------------------------------------------------------------------------
# TruePositive and its cousins - struct and metadata declerations

const TRUE_POSITIVE_AND_COUSINS =
    (:TruePositive, :TrueNegative, :FalsePositive, :FalseNegative,
     :TruePositiveRate, :TrueNegativeRate, :FalsePositiveRate,
     :FalseNegativeRate, :FalseDiscoveryRate, :Precision,
     :NegativePredictiveValue)

for M in TRUE_POSITIVE_AND_COUSINS
    ex = quote
        struct $M <: Measure rev::Union{Nothing,Bool} end
        $M(; rev=nothing) = $M(rev)
    end
    eval(ex)
end

metadata_measure.((FalsePositive, FalseNegative);
    target_scitype           = Vec{<:Finite{2}},
    prediction_type          = :deterministic,
    orientation              = :loss,
    reports_each_observation = false,
    aggregation              = Sum(),
    is_feature_dependent     = false,
    supports_weights         = false)

metadata_measure.((FalsePositiveRate, FalseNegativeRate, FalseDiscoveryRate);
    target_scitype           = Vec{<:Finite{2}},
    prediction_type          = :deterministic,
    orientation              = :loss,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false)

metadata_measure.((TruePositive, TrueNegative);
    target_scitype           = Vec{<:Finite{2}},
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    aggregation              = Sum(),
    is_feature_dependent     = false,
    supports_weights         = false)

metadata_measure.((TruePositiveRate, TrueNegativeRate, Precision,
                   NegativePredictiveValue);
    target_scitype           = Vec{<:Finite{2}},
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false)

# adjustments:
instances(::Type{<:TruePositive}) = ["true_positive", "truepositive"]
human_name(::Type{<:TruePositive})  = "number of true positives"

instances(::Type{<:TrueNegative}) = ["true_negative", "truenegative"]
human_name(::Type{<:TrueNegative}) = "number of true negatives"

instances(::Type{<:FalsePositive}) = ["false_positive", "falsepositive"]
human_name(::Type{<:FalsePositive}) = "number of false positives"

instances(::Type{<:FalseNegative}) = ["false_negative", "falsenegative"]
human_name(::Type{<:FalseNegative}) = "number of false negatives"

instances(::Type{<:TruePositiveRate}) =
    ["true_positive_rate", "truepositive_rate",
     "tpr", "sensitivity", "recall", "hit_rate"]
human_name(::Type{<:TruePositiveRate}) =
    "true positive rate (a.k.a recall)"

instances(::Type{<:TrueNegativeRate}) =
    ["true_negative_rate", "truenegative_rate", "tnr",
     "specificity", "selectivity"]

instances(::Type{<:FalsePositiveRate}) =
    ["false_positive_rate", "falsepositive_rate",
     "fpr", "fallout"]
                               "."
instances(::Type{<:FalseNegativeRate}) =
    ["false_negative_rate", "falsenegative_rate", "fnr", "miss_rate"]
                               "."
instances(::Type{<:FalseDiscoveryRate}) =
    ["false_discovery_rate", "falsediscovery_rate", "fdr"]

instances(::Type{<:NegativePredictiveValue}) =
    ["negative_predictive_value", "negativepredictive_value", "npv"]

instances(::Type{<:Precision}) =
    ["positive_predictive_value", "ppv", "positivepredictive_value", "precision"]
human_name(::Type{<:Precision}) =
    "precision (a.k.a. positive predictive value)"

# ---------------------------------------------------------------------
# TruePositive and its cousins - doc-string building and alias creation

for M in TRUE_POSITIVE_AND_COUSINS
    eval(quote
         $M == Precision || @create_aliases $M # precision handled separately

         @create_docs($M,
         body=
         """
         Assigns `false` to first element of `levels(y)`. To reverse roles,
         use `$(name($M))(rev=true)`.
         """,
         scitype=DOC_ORDERED_FACTOR_BINARY)
         end)
end

# type aliases:
const TNR = TrueNegativeRate
const Specificity = TrueNegativeRate
const TPR = TruePositiveRate
const Recall = TPR
const FPR = FalsePositiveRate
const FNR = FalseNegativeRate
const FDR = FalseDiscoveryRate
const NPV = NegativePredictiveValue
const PPV = Precision

# special case of precision; cannot generate alias's automatically due
# to conflict with Base.precision:
const positive_predictive_value = Precision()
const ppv = Precision()
const positivepredictive_value = Precision()

# ----------------------------------------------------------------------
# TruePositive and its cousins - helper functions for confusion matrices

_tp(m::CM2) = m[2,2]
_tn(m::CM2) = m[1,1]
_fp(m::CM2) = m[2,1]
_fn(m::CM2) = m[1,2]

_tpr(m::CM2) = _tp(m) / (_tp(m) + _fn(m))
_tnr(m::CM2) = _tn(m) / (_tn(m) + _fp(m))
_fpr(m::CM2) = 1 - _tnr(m)
_fnr(m::CM2) = 1 - _tpr(m)

_fdr(m::CM2) = _fp(m) / (_tp(m) + _fp(m))
_npv(m::CM2) = _tn(m) / (_tn(m) + _fn(m))

# ----------------------------------------------------------------------
# TruePositive and its cousins - calling behaviour

# NOTE: here we assume the CM was constructed a priori with the
# proper ordering so the field `rev` in the measure is ignored

# on confusion matrices:
(::TruePositive)(m::CM2)  = _tp(m)
(::TrueNegative)(m::CM2)  = _tn(m)
(::FalsePositive)(m::CM2) = _fp(m)
(::FalseNegative)(m::CM2) = _fn(m)
(::TPR)(m::CM2) = _tpr(m)
(::TNR)(m::CM2) = _tnr(m)
(::FPR)(m::CM2) = _fpr(m)
(::FNR)(m::CM2) = _fnr(m)
(::FDR)(m::CM2) = _fdr(m)
(::NPV)(m::CM2) = _npv(m)
(::Precision)(m::CM2) = 1.0 - _fdr(m)

# on vectors (ŷ, y):
for M_ex in TRUE_POSITIVE_AND_COUSINS
    local M = eval(M_ex)
    (m::M)(ŷ, y) = _confmat(ŷ, y; rev=m.rev) |> m
end

# special `precision` case (conflict with Base.precision):
Base.precision(m::CM2) = m |> Precision()
Base.precision(ŷ, y)   = _confmat(ŷ, y) |> Precision()

# =================================================================
#MULTICLASS AND ORDER INDEPENDENT

const CM = ConfusionMatrixObject{N} where N

abstract type MulticlassAvg end
struct MacroAvg <: MulticlassAvg end
struct MicroAvg <: MulticlassAvg end
struct NoAvg <: MulticlassAvg end

const macro_avg = MacroAvg()
const micro_avg = MicroAvg()
const no_avg    = NoAvg()

const DS_AVG_RET = "Options for `average` are: `no_avg`, `macro_avg` "*
    "(default) and `micro_avg`. Options for `return_type`, "*
    "applying in the `no_avg` case, are: `LittleDict` (default) or "*
    "`Vector`. "

const DS_RET = "Options for `return_type` are: "*
    "`LittleDict`(default) or "*
    "`Vector`. "

const CLASS_W = "An optional `AbstractDict`, denoted `class_w` above, "*
    "keyed on `levels(y)`, specifies class weights. It applies if "*
    "`average=macro_avg` or `average=no_avg`."

"""
    MulticlassFScore(; β=1.0, average=macro_avg, return_type=LittleDict)

One-parameter generalization, ``F_β``, of the F-measure or balanced F-score for
multiclass observations.

    MulticlassFScore()(ŷ, y)
    MulticlassFScore()(ŷ, y, class_w)

Evaluate the default score on multiclass observations, `ŷ`, given
ground truth values, `y`. $DS_AVG_RET $CLASS_W

For more information, run `info(MulticlassFScore)`.

"""
struct MulticlassFScore{T<:Real,
                        M<:MulticlassAvg,
                        U<:Union{Vector, LittleDict}} <:Measure
    β::T
    average::M
    return_type::Type{U}
end

MulticlassFScore(; β=1.0, average=macro_avg, return_type=LittleDict) =
    MulticlassFScore(β, average, return_type)

metadata_measure(MulticlassFScore;
                 instances = ["macro_f1score", "micro_f1score",
                              "multiclass_f1score"],
                 target_scitype           = Vec{<:Finite{N}} where N,
                 prediction_type          = :deterministic,
                 orientation              = :score,
                 reports_each_observation = false,
                 is_feature_dependent     = false,
                 supports_weights         = false,
                 supports_class_weights   = true)

MLJModelInterface.docstring(::Type{<:MulticlassFScore}) =
    "Multiclass F_β score; aliases: " *
    "`macro_f1score=MulticlassFScore()`, "*
    "`multiclass_f1score=MulticlassFScore()` " *
    "`micro_f1score=MulticlassFScore(average=micro_avg)`."

const micro_f1score      = MulticlassFScore(average=micro_avg)
const macro_f1score      = MulticlassFScore(average=macro_avg)
const multiclass_f1score = MulticlassFScore(average=macro_avg)

for M in (:MulticlassTruePositive, :MulticlassTrueNegative,
          :MulticlassFalsePositive, :MulticlassFalseNegative)
    ex = quote
        struct $M{U<:Union{Vector, LittleDict}} <: Measure
            return_type::Type{U}
        end
#        $M(return_type::Type{U}) where {U} = $M(return_type)
        $M(; return_type=LittleDict) = $M(return_type)
    end
    eval(ex)
end

const _mtp_vec = MulticlassTruePositive(return_type=Vector)
const _mfn_vec = MulticlassFalseNegative(return_type=Vector)
const _mfp_vec = MulticlassFalsePositive(return_type=Vector)
const _mtn_vec = MulticlassTrueNegative(return_type=Vector)

for M in (:MulticlassTruePositiveRate, :MulticlassTrueNegativeRate,
          :MulticlassFalsePositiveRate, :MulticlassFalseNegativeRate,
          :MulticlassFalseDiscoveryRate, :MulticlassPrecision,
          :MulticlassNegativePredictiveValue)
    ex = quote
        struct $M{T<:MulticlassAvg, U<:Union{Vector, LittleDict}} <: Measure
            average::T
            return_type::Type{U}
        end
        # @ablaom says next line looks redundant:
        $M(average::T, return_type::Type{U}) where {T, U} =
            $M(average, return_type)
        $M(; average=macro_avg,
           return_type=LittleDict) where T<:MulticlassAvg where
           U<:Union{Vector, LittleDict} = $M(average, return_type)
    end
    eval(ex)
end

metadata_measure.((MulticlassFalsePositive, MulticlassFalseNegative);
    target_scitype           = Vec{<:Finite{N}} where N,
    prediction_type          = :deterministic,
    orientation              = :loss,
    reports_each_observation = false,
    aggregation               = Sum(),
    is_feature_dependent     = false,
    supports_weights         = false,
    supports_class_weights   = false)

metadata_measure.((MulticlassFalsePositiveRate, MulticlassFalseNegativeRate,
                   MulticlassFalseDiscoveryRate);
    target_scitype           = Vec{<:Finite{N}} where N,
    prediction_type          = :deterministic,
    orientation              = :loss,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false,
    supports_class_weights   = true)

metadata_measure.((MulticlassTruePositive, MulticlassTrueNegative);
    target_scitype           = Vec{<:Finite{N}} where N,
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    aggregation              = Sum(),
    is_feature_dependent     = false,
    supports_weights         = false,
    supports_class_weights   = false)

metadata_measure.((MulticlassTrueNegativeRate, MulticlassNegativePredictiveValue);
    target_scitype           = Vec{<:Finite{N}} where N,
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false,
    supports_class_weights   = true)

metadata_measure.((MulticlassTruePositiveRate, MulticlassPrecision);
    target_scitype           = Vec{<:Finite{N}} where N,
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false,
    supports_class_weights   = true)

MMI.docstring(::Type{<:MulticlassTruePositive})  =
    "Number of true positives; " *
    "aliases: `multiclass_true_positive`, `multiclass_truepositive`."
instances(::Type{<:MulticlassTruePositive})  =
    ["multiclass_true_positive", "multiclass_truepositive"]
MMI.docstring(::Type{<:MulticlassTrueNegative})  =
    "Number of true negatives; " *
    "aliases: `multiclass_true_negative`, `multiclass_truenegative`."
instances(::Type{<:MulticlassTrueNegative})  =
    ["multiclass_true_negative", "multiclass_truenegative"]
MMI.docstring(::Type{<:MulticlassFalsePositive}) =
    "Number of false positives; " *
    "aliases: `multiclass_false_positive`, `multiclass_falsepositive`."
instances(::Type{<:MulticlassFalsePositive}) =
    ["multiclass_false_positive", "multiclass_falsepositive"]
MMI.docstring(::Type{<:MulticlassFalseNegative}) =
    "Number of false negatives; " *
    "aliases: `multiclass_false_negative`, `multiclass_falsenegative`."
instances(::Type{<:MulticlassFalseNegative}) =
    ["multiclass_false_negative", "multiclass_falsenegative"]

MMI.docstring(::Type{<:MulticlassTruePositiveRate}) =
    "multiclass true positive rate; aliases: " *
    "`multiclass_true_positive_rate`, `multiclass_tpr`, " *
    "`multiclass_sensitivity`, `multiclass_recall`, " *
    "`multiclass_hit_rate`, `multiclass_truepositive_rate`, "
instances(::Type{<:MulticlassTruePositiveRate}) =
    ["multiclass_true_positive_rate", "multiclass_tpr",
    "multiclass_sensitivity", "multiclass_recall",
    "multiclass_hit_rate", "multiclass_truepositive_rate"]
MMI.docstring(::Type{<:MulticlassTrueNegativeRate}) =
    "multiclass true negative rate; aliases: " *
    "`multiclass_true_negative_rate`, `multiclass_tnr` " *
    " `multiclass_specificity`, `multiclass_selectivity`, " *
    "`multiclass_truenegative_rate`."
instances(::Type{<:MulticlassTrueNegativeRate}) =
    ["multiclass_true_negative_rate", "multiclass_tnr",
    "multiclass_specificity", "multiclass_selectivity",
    "multiclass_truenegative_rate"]
MMI.docstring(::Type{<:MulticlassFalsePositiveRate}) =
                       "multiclass false positive rate; aliases: " *
                       "`multiclass_false_positive_rate`, `multiclass_fpr` " *
                       "`multiclass_fallout`, `multiclass_falsepositive_rate`."
instances(::Type{<:MulticlassFalsePositiveRate}) =
    ["multiclass_false_positive_rate", "multiclass_fpr",
     "multiclass_fallout", "multiclass_falsepositive_rate"]
MMI.docstring(::Type{<:MulticlassFalseNegativeRate}) =
    "multiclass false negative rate; aliases: " *
    "`multiclass_false_negative_rate`, `multiclass_fnr`, " *
    "`multiclass_miss_rate`, `multiclass_falsenegative_rate`."
instances(::Type{<:MulticlassFalseNegativeRate}) =
    ["multiclass_false_negative_rate", "multiclass_fnr",
    "multiclass_miss_rate", "multiclass_falsenegative_rate"]
MMI.docstring(::Type{<:MulticlassFalseDiscoveryRate}) =
    "multiclass false discovery rate; "*
    "aliases: `multiclass_false_discovery_rate`, " *
    "`multiclass_falsediscovery_rate`, `multiclass_fdr`."
instances(::Type{<:MulticlassFalseDiscoveryRate}) =
    ["multiclass_falsediscovery_rate", "multiclass_fdr"]
MMI.docstring(::Type{<:MulticlassNegativePredictiveValue}) =
    "multiclass negative predictive value; aliases: " *
    "`multiclass_negative_predictive_value`, " *
    "`multiclass_negativepredictive_value`, `multiclass_npv`."
instances(::Type{<:MulticlassNegativePredictiveValue}) =
    ["multiclass_negative_predictive_value",
    "multiclass_negativepredictive_value", "multiclass_npv"]
MMI.docstring(::Type{<:MulticlassPrecision}) =
  "multiclass positive predictive value (aka precision);"*
  " aliases: `multiclass_positive_predictive_value`, `multiclass_ppv`, " *
  "`multiclass_positivepredictive_value`, " *
  "`multiclass_recall`."
instances(::Type{<:MulticlassPrecision}) =
    ["multiclass_positive_predictive_value", "multiclass_ppv",
     "multiclass_positivepredictive_value", "multiclass_recall"]

const W_KEY_MISMATCH = "Encountered target with levels different from the " *
                       "keys of user-specified dictionary of class weights."
const W_PROMOTE_WARN = "Using macro averaging instead of micro averaging, as "*
    "class weights specified. "

# ----------------------------------------------------
# MulticlassTruePositive

"""
    MulticlassTruePositive(; return_type=LittleDict)

$(docstring(MulticlassTruePositive()))

    MulticlassTruePositive()(ŷ, y)

Number of true positives for multiclass observations `ŷ` and ground
truth `y`, using default return type. $DS_RET

For more information, run `info(MulticlassTruePositive)`.

"""
function MulticlassTruePositive end
const multiclass_true_positive  = MulticlassTruePositive()
const multiclass_truepositive   = MulticlassTruePositive()
const mtp = MulticlassTruePositive()

# ----------------------------------------------------
# MulticlassTrueNegative

"""
    MulticlassTrueNegative(; return_type=LittleDict)

$(docstring(MulticlassTrueNegative()))

    MulticlassTrueNegative()(ŷ, y)

Number of true negatives for multiclass observations `ŷ` and ground truth
`y`, using default return type. $DS_RET

For more information, run `info(MulticlassTrueNegative)`.

"""
function MulticlassTrueNegative end
const multiclass_true_negative  = MulticlassTrueNegative()
const multiclass_truenegative   = MulticlassTrueNegative()
const mtn = MulticlassTrueNegative()

# ----------------------------------------------------
# MulticlassFalsePositive

"""
    MulticlassFalsePositive(; return_type=LittleDict)

$(docstring(MulticlassFalsePositive()))

    MulticlassFalsePositive()(ŷ, y)

Number of false positives for multiclass observations `ŷ` and ground
truth `y`, using default return type. $DS_RET

For more information, run `info(MulticlassFalsePositive)`.

"""
function MulticlassPositive end
const multiclass_false_positive = MulticlassFalsePositive()
const multiclass_falsepositive  = MulticlassFalsePositive()
const mfp = MulticlassFalsePositive()

# ----------------------------------------------------
# MulticlassFalseNegative

"""
    MulticlassFalseNegative(; return_type=LittleDict)

$(docstring(MulticlassFalseNegative()))

    MulticlassFalseNegative()(ŷ, y)

Number of false negatives for multiclass observations `ŷ` and ground
truth `y`, using default return type. $DS_RET

For more information, run `info(MulticlassFalseNegative)`.

"""
function MulticlassNegative end
const multiclass_false_negative = MulticlassFalseNegative()
const multiclass_falsenegative  = MulticlassFalseNegative()
const mfn = MulticlassFalseNegative()

# ----------------------------------------------------
# MulticlassTruePositiveRate

"""
    MulticlassTruePositiveRate(; average=macro_avg, return_type=LittleDict)

$(docstring(MulticlassTruePositiveRate()))

    MulticlassTruePositiveRate(ŷ, y)
    MulticlassTruePositiveRate(ŷ, y, class_w)

True positive rate (a.k.a. sensitivity, recall, hit rate) for
multiclass observations `ŷ` and ground truth `y`, using default
averaging and return type. $DS_AVG_RET $CLASS_W

For more information, run `info(MulticlassTruePositiveRate)`.

"""
function MulticlassTruePositiveRate end
const multiclass_true_positive_rate = MulticlassTruePositiveRate()
const multiclass_truepositive_rate  = MulticlassTruePositiveRate()
const multiclass_tpr                = MulticlassTruePositiveRate()
const multiclass_sensitivity        = MulticlassTruePositiveRate()
const multiclass_hit_rate           = MulticlassTruePositiveRate()
const MTPR                          = MulticlassTruePositiveRate
const multiclass_recall             = MulticlassTruePositiveRate()
const MulticlassRecall              = MulticlassTruePositiveRate

# ----------------------------------------------------
# MulticlassTrueNegativeRate

"""
    MulticlassTrueNegativeRate(; average=macro_avg, return_type=LittleDict)

$(docstring(MulticlassTrueNegativeRate()))

    MulticlassTrueNegativeRate()(ŷ, y)
    MulticlassTrueNegativeRate()(ŷ, y, class_w)

True negative rate for multiclass observations `ŷ` and ground truth
`y`, using default averaging and return type. $DS_AVG_RET $CLASS_W

For more information, run `info(MulticlassTrueNegativeRate)`.

"""
function MulticlassTrueNegativeRate end
const multiclass_true_negative_rate = MulticlassTrueNegativeRate()
const multiclass_truenegative_rate  = MulticlassTrueNegativeRate()
const multiclass_tnr                = MulticlassTrueNegativeRate()
const multiclass_specificity        = MulticlassTrueNegativeRate()
const multiclass_selectivity        = MulticlassTrueNegativeRate()
const MulticlassSpecificity         = MulticlassTrueNegativeRate
const MTNR                          = MulticlassTrueNegativeRate

# ----------------------------------------------------
# MulticlassFalsePositiveRate

"""
    MulticlassFalsePositiveRate(; average=macro_avg, return_type=LittleDict)

$(docstring(MulticlassFalsePositiveRate()))

    MulticlassFalsePositiveRate()(ŷ, y)
    MulticlassFalsePositiveRate()(ŷ, y, class_w)

False positive rate for multiclass observations `ŷ` and ground truth
`y`, using default averaging and return type.  $DS_AVG_RET $CLASS_W

For more information, run `info(MulticlassFalsePositiveRate)`.

"""
function MulticlassFalsePositiveRate end
const multiclass_false_positive_rate = MulticlassFalsePositiveRate()
const multiclass_falsepositive_rate  = MulticlassFalsePositiveRate()
const multiclass_fpr                 = MulticlassFalsePositiveRate()
const MFPR                           = MulticlassFalsePositiveRate
const multiclass_fallout             = MFPR()

# ----------------------------------------------------
# MulticlassFalseNegativeRate

"""
    MulticlassFalseNegativeRate(; average=macro_avg, return_type=LittleDict)

$(docstring(MulticlassFalseNegativeRate()))

    MulticlassFalseNegativeRate()(ŷ, y)
    MulticlassFalseNegativeRate()(ŷ, y, class_w)

False negative rate for multiclass observations `ŷ` and ground truth
`y`, using default averaging and return type.  $DS_AVG_RET $CLASS_W

For more information, run `info(MulticlassFalseNegativeRate)`.

"""
function MulticlassFalseNegativeRate end
const multiclass_false_negative_rate = MulticlassFalseNegativeRate()
const multiclass_falsenegative_rate  = MulticlassFalseNegativeRate()
const multiclass_fnr                 = MulticlassFalseNegativeRate()
const MFNR                           = MulticlassFalseNegativeRate
const multiclass_miss_rate           = MFNR()

# ----------------------------------------------------
# MulticlassFalseDiscoveryRate

"""
    MulticlassFalseDiscoveryRate(; average=macro_avg, return_type=LittleDict)

$(docstring(MulticlassFalseDiscoveryRate()))

    MulticlassFalseDiscoveryRate()(ŷ, y)
    MulticlassFalseDiscoveryRate()(ŷ, y, class_w)

False discovery rate for multiclass observations `ŷ` and ground truth
`y`, using default averaging and return type.  $DS_AVG_RET $CLASS_W

For more information, run `info(MulticlassFalseDiscoveryRate)`.

"""
function MulticlassFalseDiscoveryRate end
const multiclass_false_discovery_rate = MulticlassFalseDiscoveryRate()
const multiclass_falsediscovery_rate  = MulticlassFalseDiscoveryRate()
const multiclass_fdr                  = MulticlassFalseDiscoveryRate()
const MFDR                            = MulticlassFalseDiscoveryRate

# ----------------------------------------------------
# MulticlassPrecision

"""
    MulticlassPrecision(; average=macro_avg, return_type=LittleDict)

$(docstring(MulticlassPrecision()))

    MulticlassPrecision()(ŷ, y)
    MulticlassPrecision()(ŷ, y, class_w)

Precision for multiclass observations `ŷ` and ground truth `y`, using
default averaging and return type. $DS_AVG_RET $CLASS_W

For more information, run `info(MulticlassPrecision)`.

"""
function MulticlassPrecision end
const multiclass_precision                 = MulticlassPrecision()
const multiclass_ppv                       = MulticlassPrecision()
const multiclass_positive_predictive_value = MulticlassPrecision()
const multiclass_positivepredictive_value  = MulticlassPrecision()
const MPPV                                 = MulticlassPrecision

# ----------------------------------------------------
# MulticlassNegativePredictiveValue

"""
    MulticlassNegativePredictiveValue(; average=macro_avg, return_type=LittleDict)

$(docstring(MulticlassNegativePredictiveValue()))

    MulticlassNegativePredictiveValue()(ŷ, y)
    MulticlassNegativePredictiveValue()(ŷ, y, class_w)

Negative predictive value for multiclass observations `ŷ` and ground truth
`y`, using default averaging and return type. $DS_AVG_RET $CLASS_W

For more information, run `info(MulticlassNegativePredictiveValue)`.

"""
function MulticlassNegativePredictiveValue end
const multiclass_npv                       = MulticlassNegativePredictiveValue()
const multiclass_negative_predictive_value = MulticlassNegativePredictiveValue()
const multiclass_negativepredictive_value  = MulticlassNegativePredictiveValue()
const MNPV                                 = MulticlassNegativePredictiveValue

# -----------------------------------------------------
## INTERNAL FUCNTIONS ON MULTICLASS CONFUSION MATRIX

_mtp(m::CM, return_type::Type{Vector}) = diag(m.mat)
_mtp(m::CM, return_type::Type{LittleDict}) =
    LittleDict(m.labels, diag(m.mat))

_mfp(m::CM, return_type::Type{Vector}) =
    (col_sum = vec(sum(m.mat, dims=2)); col_sum .-= diag(m.mat))

_mfp(m::CM, return_type::Type{LittleDict}) =
    (col_sum = vec(sum(m.mat, dims=2)); col_sum .-= diag(m.mat);
     LittleDict(m.labels, col_sum))

_mfn(m::CM, return_type::Type{Vector}) =
    (row_sum = vec(collect(transpose(sum(m.mat, dims=1))));
     row_sum .-= diag(m.mat))

_mfn(m::CM, return_type::Type{LittleDict}) =
    (row_sum = vec(collect(transpose(sum(m.mat, dims=1))));
     row_sum .-= diag(m.mat); LittleDict(m.labels, row_sum))

function _mtn(m::CM, return_type::Type{Vector})
    _sum  = sum(m.mat, dims=2)
    _sum .= sum(m.mat) .- (_sum .+= sum(m.mat, dims=1)'.+ diag(m.mat))
    return vec(_sum)
end

function _mtn(m::CM, return_type::Type{LittleDict})
    _sum  = sum(m.mat, dims=2)
    _sum .= sum(m.mat) .- (_sum .+= sum(m.mat, dims=1)'.+ diag(m.mat))
    return LittleDict(m.labels, vec(_sum))
end

@inline function _mean(x::Vec{<:Real})
    for i in eachindex(x)
        @inbounds x[i] = ifelse(isnan(x[i]), zero(eltype(x)), x[i])
    end
    return mean(x)
end

@inline function _class_w(level_m::Vec{<:String},
                          class_w::AbstractDict{<:Any, <:Real})
    class_w_labels = levels(keys(class_w))
    string.(class_w_labels) == level_m || throw(ArgumentError(W_KEY_MISMATCH))
    return [class_w[l] for l in class_w_labels]
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            average::NoAvg, return_type::Type{Vector})
    return vec(a ./ (a + b))
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            average::NoAvg, return_type::Type{LittleDict})
    return LittleDict(m.labels, _mc_helper(m, a, b, average, Vector))
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            average::MacroAvg, return_type)
    return _mean(_mc_helper(m, a, b, no_avg, Vector))
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            average::MicroAvg, return_type)
    a_sum, b_sum = sum(a), sum(b)
    return a_sum / (a_sum + b_sum)
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            class_w::AbstractDict{<:Any, <:Real},
                            average::NoAvg, return_type::Type{Vector})
    level_w = _class_w(m.labels, class_w)
    return _mc_helper(m, a, b, no_avg, return_type) .* level_w
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            class_w::AbstractDict{<:Any, <:Real},
                            average::MacroAvg, return_type::Type{Vector})
    return _mean(_mc_helper(m, a, b, class_w, no_avg, return_type))
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            class_w::AbstractDict{<:Any, <:Real},
                            average::MicroAvg, return_type)
    @warn W_PROMOTE_WARN
    return _mc_helper(m, a, b, class_w, macro_avg, Vector)
end

@inline function _mc_helper_b(m::CM, helper_name,
                              class_w::AbstractDict{<:Any, <:Real},
                              average::NoAvg, return_type::Type{Vector})
    level_w = _class_w(m.labels, class_w)
    return (1 .- helper_name(m, no_avg, return_type)) .* level_w
end

@inline function _mc_helper_b(m::CM, helper_name,
                              class_w::AbstractDict{<:Any, <:Real},
                              average::NoAvg, return_type::Type{LittleDict})
    level_w = _class_w(m.labels, class_w)
    return LittleDict(m.labels, ((1 .- helper_name(m, no_avg, Vector)) .* level_w))
end

@inline function _mc_helper_b(m::CM, helper_name,
                              class_w::AbstractDict{<:Any, <:Real},
                              average::MacroAvg, return_type)
    return _mean(_mc_helper_b(m, helper_name, class_w, no_avg, Vector))
end

@inline function _mc_helper_b(m::CM, helper_name,
                              class_w::AbstractDict{<:Any, <:Real},
                              average::MicroAvg, return_type)
    @warn W_PROMOTE_WARN
    return _mc_helper_b(m, helper_name, class_w, macro_avg, Vector)
end

@inline function _mc_helper_b(m::CM, helper_name, average::NoAvg,
                              return_type::Type{LittleDict})
    return LittleDict(m.labels, 1.0 .- helper_name(m, average, Vector))
end

@inline function _mc_helper_b(m::CM, helper_name, average::NoAvg,
                              return_type::Type{Vector})
    return 1.0 .- helper_name(m, average, Vector)
end

@inline function _mc_helper_b(m::CM, helper_name, average::MacroAvg,
                              return_type)
    return 1.0 .- helper_name(m, average, Vector)
end

@inline function _mc_helper_b(m::CM, helper_name, average::MicroAvg,
                              return_type)
    return 1.0 .- helper_name(m, average, Vector)
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            class_w::AbstractDict{<:Any, <:Real},
                            average::NoAvg, return_type::Type{LittleDict})
    level_w = _class_w(m.labels, class_w)
    return LittleDict(m.labels, _mc_helper(m, a, b, class_w, no_avg, Vector))
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            class_w::AbstractDict{<:Any, <:Real},
                            average::MacroAvg, return_type::Type{U}) where U
    return _mean(_mc_helper(m, a, b, class_w, no_avg, Vector))
end

@inline function _mc_helper(m::CM, a::Vec{<:Real}, b::Vec{<:Real},
                            class_w::AbstractDict{<:Any, <:Real},
                            average::MicroAvg, return_type::Type{U}) where U
    @warn W_PROMOTE_WARN
    return _mc_helper(m, a, b, class_w, macro_avg, return_type)
end

function _mtpr(m::CM, average::A, return_type::Type{U}) where {A, U}
    mtp_val, mfn_val = _mtp_vec(m), _mfn_vec(m)
    return _mc_helper(m, mtp_val, mfn_val, average, return_type)
end

function _mtpr(m::CM, class_w::AbstractDict{<:Any, <:Real}, average::A,
               return_type::Type{U}) where {A, U}
    mtp_val, mfn_val = _mtp_vec(m), _mfn_vec(m)
    return _mc_helper(m, mtp_val, mfn_val, class_w, average, return_type)
end

function _mtnr(m::CM, average::A, return_type::Type{U}) where {A, U}
    mtn_val, mfp_val = _mtn_vec(m), _mfp_vec(m)
    return _mc_helper(m, mtn_val, mfp_val, average, return_type)
end

function _mtnr(m::CM, class_w::AbstractDict{<:Any, <:Real}, average::A,
               return_type::Type{U}) where {A, U}
    mtn_val, mfp_val = _mtn_vec(m), _mfp_vec(m)
    return _mc_helper(m, mtn_val, mfp_val, class_w, average, return_type)
end

_mfpr(m::CM, average::A, return_type::Type{U}) where {A, U} =
    _mc_helper_b(m, _mtnr, average, return_type)

function _mfpr(m::CM, class_w::AbstractDict{<:Any, <:Real}, average::A,
               return_type::Type{U}) where {A, U}
    return _mc_helper_b(m, _mtnr, class_w, average, return_type)
end

_mfnr(m::CM, average::A, return_type::Type{U}) where {A, U} =
    _mc_helper_b(m, _mtpr, average, return_type)

function _mfnr(m::CM, class_w::AbstractDict{<:Any, <:Real}, average::A,
               return_type::Type{U}) where {A, U}
    return _mc_helper_b(m, _mtpr, class_w, average, return_type)
end

function _mfdr(m::CM, average::A, return_type::Type{U}) where {A, U}
    mfp_val, mtp_val = _mfp_vec(m), _mtp_vec(m)
    return _mc_helper(m, mfp_val, mtp_val, average, return_type)
end

function _mfdr(m::CM, class_w::AbstractDict{<:Any, <:Real}, average::A,
               return_type::Type{U}) where {A, U}
    mfp_val, mtp_val = _mfp_vec(m), _mtp_vec(m)
    return _mc_helper(m, mfp_val, mtp_val, class_w, average, return_type)
end

function _mnpv(m::CM, average::A, return_type::Type{U}) where {A, U}
    mtn_val, mfn_val = _mtn_vec(m), _mfn_vec(m)
    return _mc_helper(m, mtn_val, mfn_val, average, return_type)
end

function _mnpv(m::CM, class_w::AbstractDict{<:Any, <:Real}, average::A,
               return_type::Type{U}) where {A, U}
    mtn_val, mfn_val = _mtn_vec(m), _mfn_vec(m)
    return _mc_helper(m, mtn_val, mfn_val, class_w, average, return_type)
end

## CALLABLES ON MULTICLASS CONFUSION MATRIX

(p::MulticlassTruePositive)(m::CM)  = _mtp(m, p.return_type)
(n::MulticlassTrueNegative)(m::CM)  = _mtn(m, n.return_type)
(p::MulticlassFalsePositive)(m::CM) = _mfp(m, p.return_type)
(n::MulticlassFalseNegative)(m::CM) = _mfn(m, n.return_type)

(r::MTPR)(m::CM) = _mtpr(m, r.average, r.return_type)
(r::MTPR)(m::CM, w::AbstractDict{<:Any, <:Real}) =
    _mtpr(m, w, r.average, r.return_type)

(r::MTNR)(m::CM) = _mtnr(m, r.average, r.return_type)
(r::MTNR)(m::CM, w::AbstractDict{<:Any, <:Real}) =
    _mtnr(m, w, r.average, r.return_type)

(r::MFPR)(m::CM) = _mfpr(m, r.average, r.return_type)
(r::MFPR)(m::CM, w::AbstractDict{<:Any, <:Real}) =
    _mfpr(m, w, r.average, r.return_type)

(r::MFNR)(m::CM) = _mfnr(m, r.average, r.return_type)
(r::MFNR)(m::CM, w::AbstractDict{<:Any, <:Real}) =
    _mfnr(m, w, r.average, r.return_type)

(r::MFDR)(m::CM) = _mfdr(m, r.average, r.return_type)
(r::MFDR)(m::CM, w::AbstractDict{<:Any, <:Real}) =
    _mfdr(m, w, r.average, r.return_type)

(v::MNPV)(m::CM) = _mnpv(m, v.average, v.return_type)
(v::MNPV)(m::CM, w::AbstractDict{<:Any, <:Real}) =
    _mnpv(m, w, v.average, v.return_type)

(p::MulticlassPrecision)(m::CM) =
    _mc_helper_b(m, _mfdr, p.average, p.return_type)
function (p::MulticlassPrecision)(m::CM, class_w::AbstractDict{<:Any, <:Real})
    return _mc_helper_b(m, _mfdr, class_w, p.average, p.return_type)
end

@inline function _fs_helper(m::CM, β::Real, rec::Vec{<:Real}, prec::Vec{<:Real},
                    average::NoAvg, return_type::Type{LittleDict})
    β2 = β^2
    return LittleDict(m.labels, (1 + β2) .* (prec .* rec) ./ (β2 .* prec .+ rec))
end

@inline function _fs_helper(m::CM, β::Real, rec::Vec{<:Real}, prec::Vec{<:Real},
                    average::NoAvg, return_type::Type{Vector})
    β2 = β^2
    return (1 + β2) .* (prec .* rec) ./ (β2 .* prec .+ rec)
end

@inline function _fs_helper(m::CM, β::Real, rec::Vec{<:Real}, prec::Vec{<:Real},
                    average::MacroAvg, return_type::Type{U}) where U
    return _mean(_fs_helper(m, β, rec, prec, no_avg, Vector))
end

@inline function _fs_helper(m::CM, β::Real, rec::Real, prec::Real,
                            average::MicroAvg, return_type::Type{U}) where U
    β2 = β^2
    return (1 + β2) * (prec * rec) / (β2 * prec + rec)
end

function (f::MulticlassFScore)(m::CM)
    if f.average == micro_avg
        rec = MulticlassRecall(; average=f.average, return_type=Vector)(m)
        f.β == 1.0 && return rec
        return _fs_helper(m, f.β, rec, rec, f.average, f.return_type)
    end
    rec = MulticlassRecall(; average=no_avg, return_type=Vector)(m)
    prec = MulticlassPrecision(; average=no_avg, return_type=Vector)(m)
    return _fs_helper(m, f.β, rec, prec, f.average, f.return_type)
end

@inline function _fs_helper(m::CM, w::AbstractDict{<:Any, <:Real}, β::Real,
                    average::NoAvg, return_type::Type{LittleDict})
    level_w = _class_w(m.labels, w)
    return LittleDict(m.labels,
                      MulticlassFScore(β=β,
                                       average=no_avg,
                                       return_type=Vector)(m) .* level_w)
end

@inline function _fs_helper(m::CM, w::AbstractDict{<:Any, <:Real}, β::Real,
                    average::NoAvg, return_type::Type{Vector})
    level_w = _class_w(m.labels, w)
    return MulticlassFScore(β=β,
                            average=no_avg,
                            return_type=Vector)(m) .* level_w
end

@inline function _fs_helper(m::CM, w::AbstractDict{<:Any, <:Real}, β::Real,
                            average::MacroAvg, return_type::Type{U}) where U
    return _mean(_fs_helper(m, w, β, no_avg, Vector))
end

@inline function _fs_helper(m::CM, w::AbstractDict{<:Any, <:Real}, β::Real,
                            average::MicroAvg, return_type::Type{U}) where U
    @warn W_PROMOTE_WARN
    return _fs_helper(m, w, β, macro_avg, return_type)
end

function (f::MulticlassFScore)(m::CM, class_w::AbstractDict{<:Any, <:Real})
    return _fs_helper(m, class_w, f.β, f.average, f.return_type)
end

## Callables on vectors

for M in (MulticlassTruePositive, MulticlassTrueNegative,
          MulticlassFalsePositive, MulticlassFalseNegative)
    (m::M)(ŷ, y) = m(_confmat(ŷ, y, warn=false))
end

for M in (MTPR, MTNR, MFPR, MFNR, MFDR, MulticlassPrecision, MNPV,
          MulticlassFScore)
    @eval (m::$M)(ŷ, y) = m(_confmat(ŷ, y, warn=false))
    @eval (m::$M)(ŷ, y, class_w::AbstractDict{<:Any, <:Real}) =
                          m(_confmat(ŷ, y, warn=false), class_w)
end

## ROC COMPUTATION

"""
    _idx_unique_sorted(v)

Internal function to return the index of unique elements in `v` under the
assumption that the vector `v` is sorted in decreasing order.
"""
function _idx_unique_sorted(v::Vec{<:Real})
    n    = length(v)
    idx  = ones(Int, n)
    p, h = 1, 1
    cur  = v[1]
    @inbounds while h < n
        h     += 1                  # head position
        cand   = v[h]               # candidate value
        cand   < cur || continue    # is it new? otherwise skip
        p     += 1                  # if new store it
        idx[p] = h
        cur    = cand               # and update the last seen value
    end
    p < n && deleteat!(idx, p+1:n)
    return idx
end

"""
    fprs, tprs, ts = roc_curve(ŷ, y) = roc(ŷ, y)

Return the ROC curve for a two-class probabilistic prediction `ŷ` given the
ground  truth `y`. The true positive rates, false positive rates over a range
of thresholds `ts` are returned. Note that if there are `k` unique scores,
there are correspondingly  `k` thresholds and `k+1` "bins" over which the FPR
and TPR are constant:

* `[0.0 - thresh[1]]`
* `[thresh[1] - thresh[2]]`
* ...
* `[thresh[k] - 1]`

consequently, `tprs` and `fprs` are of length `k+1` if `ts` is of length `k`.

To draw the curve using your favorite plotting backend, do `plot(fprs, tprs)`.
"""
function roc_curve(ŷ::Vec{<:UnivariateFinite},
                   y::Vec{<:CategoricalValue})

    n       = length(y)
    lab_pos = levels(y)[2]
    scores  = pdf.(ŷ, lab_pos)
    ranking = sortperm(scores, rev=true)

    scores_sort = scores[ranking]
    y_sort_bin  = (y[ranking] .== lab_pos)

    idx_unique = _idx_unique_sorted(scores_sort)
    thresholds = scores_sort[idx_unique]

    # detailed computations with example:
    # y = [  1   0   0   1   0   0   1]
    # s = [0.5 0.5 0.2 0.2 0.1 0.1 0.1] thresh are 0.5 0.2, 0.1 // idx [1, 3, 5]
    # ŷ = [  0   0   0   0   0   0   0] (0.5 - 1.0] # no pos pred
    # ŷ = [  1   1   0   0   0   0   0] (0.2 - 0.5] # 2 pos pred
    # ŷ = [  1   1   1   1   0   0   0] (0.1 - 0.2] # 4 pos pred
    # ŷ = [  1   1   1   1   1   1   1] [0.0 - 0.1] # all pos pre

    idx_unique_2 = idx_unique[2:end]   # [3, 5]
    n_ŷ_pos      = idx_unique_2 .- 1   # [2, 4] implicit [0, 2, 4, 7]

    cs   = cumsum(y_sort_bin)          # [1, 1, 1, 2, 2, 2, 3]
    n_tp = cs[n_ŷ_pos]                 # [1, 2] implicit [0, 1, 2, 3]
    n_fp = n_ŷ_pos .- n_tp             # [1, 2] implicit [0, 1, 2, 4]

    # add end points
    P = sum(y_sort_bin) # total number of true positives
    N = n - P           # total number of true negatives

    n_tp = [0, n_tp..., P] # [0, 1, 2, 3]
    n_fp = [0, n_fp..., N] # [0, 1, 2, 4]

    tprs = n_tp ./ P  # [0/3, 1/3, 2/3, 1]
    fprs = n_fp ./ N  # [0/4, 1/4, 2/4, 1]

    return fprs, tprs, thresholds
end

const roc = roc_curve


# This file is a part of EmpiricalDistributions.jl, licensed under the MIT License (MIT).

_bin_left(edge::AbstractVector{<:Real}) = view(edge, firstindex(edge):lastindex(edge)-1)
_bin_right(edge::AbstractVector{<:Real}) = view(edge, firstindex(edge)+1:lastindex(edge))
_bin_centers(edge::AbstractVector{<:Real}) = (_bin_left(edge) .+ _bin_right(edge)) ./ 2

_bin_widths(edge::AbstractVector{<:Real}) = _bin_right(edge) .- _bin_left(edge)

# Ranges need special treatment, since subtraction of integer ranges results in "step cannot be zero" error
function _bin_widths(edge::AbstractRange{<:Real})
    ref = step(edge)
    T = typeof(ref)
    stp = zero(T)
    n = length(eachindex(edge)) - 1
    StepRangeLen(ref, stp, n)
end

function _ratio(a::Real, b::Real)
    R = float(promote_type(typeof(a), typeof(b)))
    a == b ?  one(R) : convert(R, a / b)
end

function _linear_interpol(x_lo::Real, x_hi::Real, y_lo::Real, y_hi::Real, x::Real)
    T = promote_type(typeof(y_lo), typeof(y_hi), typeof(x))
    w_hi = T(_ratio(x - x_lo, x_hi - x_lo))
    w_lo = one(w_hi) - w_hi
    y = y_lo * w_lo + y_hi * w_hi
    @assert y_lo <= y <= y_hi
    y
end

function _linear_interpol(X::AbstractVector{<:Real}, Y::AbstractVector{<:Real}, x::Real)
    i_lo, i_hi = _find_idxs_lohi(X, x)
    x_lo, x_hi = X[i_lo], X[i_hi]
    y_lo, y_hi = Y[i_lo], Y[i_hi]
    _linear_interpol(x_lo, x_hi, y_lo, y_hi, x)
end

function _find_bin(edge::AbstractVector{<:Real}, closed_left::Bool, x::Real)
    if closed_left == false
        searchsortedfirst(edge, x) - 1
    else
        searchsortedlast(edge, x)
    end
end

function _find_bin(edges::NTuple{N,AbstractVector{<:Real}}, closed_left::Bool, x::NTuple{N,Real}) where N
    map(_find_bin, edges, map(_ -> closed_left, x), x)
end

function _find_idxs_lohi(X::AbstractVector{<:Real}, x::Real)
    x_min, x_max = first(X), last(X)
    x_min <= x <= x_max || throw(DomainError(x, "require value between x_min and x_max"))
    idxs = searchsorted(X, x)
    i1, i2 = first(idxs), last(idxs)
    i1, i2 = first(idxs), last(idxs)
    from, to = min(i1, i2), max(i1, i2)
    @assert checkbounds(Bool, X, from:to)
    (from, to)
end

function _mean(h::StatsBase.Histogram{<:Real, N}) where {N}
    @assert h.isdensity # Implementation requires normalized histogram

    Y = h.weights
    X = _bin_centers.(h.edges)
    W = _bin_widths.(h.edges)

    mean_est = mapreduce(
        idx -> map(getindex, X, idx.I) .* Y[idx] .* prod(map(getindex, W, idx.I)),
        (a, b) -> a .+ b,
        CartesianIndices(Y)
    ) ./ sum(idx -> Y[idx] .* prod(map(getindex, W, idx.I)), CartesianIndices(Y))
end

_findmaxidx_tuple_or_int(A::AbstractVector{<:Real}) = findmax(A)[2]
_findmaxidx_tuple_or_int(A::AbstractArray{<:Real}) = findmax(A)[2].I

function _mode(h::StatsBase.Histogram)
    @assert h.isdensity # Implementation requires normalized histogram

    maxidx = _findmaxidx_tuple_or_int(h.weights)
    mode_corner1 = map(getindex, h.edges, maxidx)
    mode_corner2 = map(getindex, h.edges, maxidx .+ 1)
    (mode_corner1 .+ mode_corner2) ./ 2
end

function _var(h::StatsBase.Histogram{<:Real,N}, mean_est::NTuple{N,Real}) where {N}
    @assert h.isdensity # Implementation requires normalized histogram

    Y = h.weights
    X = _bin_centers.(h.edges)
    W = _bin_widths.(h.edges)

    mean_est = mapreduce(
        idx -> (map(getindex, X, idx.I) .- mean_est).^2 .* Y[idx] .* prod(map(getindex, W, idx.I)),
        (a, b) -> a .+ b,
        CartesianIndices(Y)
    ) ./ sum(idx -> Y[idx] .* prod(map(getindex, W, idx.I)), CartesianIndices(Y))
end

function _cov(h::StatsBase.Histogram{<:Real,N}, mean_est::NTuple{N,Real}) where {N}
    @assert h.isdensity # Implementation requires normalized histogram

    Y = h.weights
    X = _bin_centers.(h.edges)
    W = _bin_widths.(h.edges)

    dim_idxs = 1:N

    [
        mapreduce(
            idx -> (X[i][idx.I[i]] - mean_est[i]) * (X[j][idx.I[j]] - mean_est[j]) .* Y[idx] .* prod(map(getindex, W, idx.I)),
            (a, b) -> a .+ b,
            CartesianIndices(Y)
        ) ./ sum(idx -> Y[idx] .* prod(map(getindex, W, idx.I)), CartesianIndices(Y))
        for i in dim_idxs, j in dim_idxs
    ]
end

_rand_uniform(rng::AbstractRNG, T::Type{<:Real}, lo::Real, hi::Real) = fma(hi - lo, rand(rng, T), lo)

# For tests of adapt()
struct _AdaptToArray end
Adapt.adapt_storage(::Val{_AdaptToArray}, xs::AbstractArray) = convert(Array, xs)


module MultiFieldFESpacesWithLinearConstraintsTests

using Gridap.Algebra
using Gridap.Arrays
using Gridap.Fields
using Gridap.Geometry
using Gridap.FESpaces
using Gridap.TensorValues
using Gridap.MultiField
using Gridap.CellData
using Test

domain = (0,1,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)

labels = get_face_labeling(model)
add_tag_from_tags!(labels,"dirichlet_1",[1,2,5])
add_tag_from_tags!(labels,"neumann_1",[6,7,8])
add_tag_from_tags!(labels,"dirichlet_2",[1,3,7])
add_tag_from_tags!(labels,"neumann_2",[5,6,8])

trian = Triangulation(model)
btrian1 = BoundaryTriangulation(model,"neumann_1")
btrian2 = BoundaryTriangulation(model,"neumann_2")
strian = SkeletonTriangulation(model)

quad = CellQuadrature(trian,2)
bquad1 = CellQuadrature(btrian1,2)
bquad2 = CellQuadrature(btrian2,2)
squad = CellQuadrature(strian,2)

bn1 = get_normal_vector(btrian1)
bn2 = get_normal_vector(btrian2)

V1 = FESpace(
  model=model,valuetype=Float64,reffe=:Lagrangian,order=1,
  conformity=:H1,dirichlet_tags="dirichlet_1")

V2 = FESpace(
  model=model,valuetype=Float64,reffe=:Lagrangian,order=1,
  conformity=:H1,dirichlet_tags="dirichlet_2")

sDOF_to_dof = [1,5,-2]
sDOF_to_dofs = Table([[-1,4],[4,6],[-1,-3]])
sDOF_to_coeffs = Table([[0.5,0.5],[0.5,0.5],[0.5,0.5]])
V1c = FESpaceWithLinearConstraints(sDOF_to_dof, sDOF_to_dofs, sDOF_to_coeffs, V1)

sDOF_to_dof = [4,-2]
sDOF_to_dofs = Table([[2,6],[-1,-3]])
sDOF_to_coeffs = Table([[0.5,0.5],[0.5,0.5]])
V2c = FESpaceWithLinearConstraints(sDOF_to_dof, sDOF_to_dofs, sDOF_to_coeffs, V2)

u1(x) = x[1] + 2*x[2]
u2(x) = 4*x[1] + x[2]
f1(x) = -Δ(u1)(x) + u2(x)
#f2(x) = -Δ(u2)(x)

U1c = TrialFESpace(V1c,u1)
U2c = TrialFESpace(V2c,u2)

U = MultiFieldFESpace([U1c,U2c])
V = MultiFieldFESpace([V1c,V2c])
test_fe_space(U)
test_fe_space(V)
@test has_constraints(U)
@test has_constraints(V)

a(u,v) = ∇(v)⋅∇(u)
b_Γ(v,u,n_Γ) = v*(n_Γ⋅∇(u))

function A(u,v)
  u1,u2 = u
  v1,v2 = v
  a(v1,u1) + v1*u2 + a(v2,u2)
end

function B(v)
  v1,v2 = v
  v1*f1
end

function B1_Γ(v)
  v1,v2 = v
  b_Γ(v1,u1,bn1)
end

function B2_Γ(v)
  v1,v2 = v
  b_Γ(v2,u2,bn2)
end

# This is only to stress skeleton machinery
# since shape functions are continuous
function As(u,v)
  u1,u2 = u
  v1,v2 = v
  jump(v1)*jump(u2) + jump(v2)*jump(u2)
end

function Bs(v)
  v1,v2 = v
  jump(v2)
end

t_Ω = AffineFETerm(A,B,trian,quad)
t1_Γ = FESource(B1_Γ,btrian1,bquad1)
t2_Γ = FESource(B2_Γ,btrian2,bquad2)
t_s = AffineFETerm(As,Bs,strian,squad)

op = AffineFEOperator(U,V,t_Ω,t1_Γ,t2_Γ,t_s)
uh = solve(op)

uh1, uh2 = uh

e1 = u1 - uh1
e2 = u2 - uh2

#using Gridap.Visualization
#v1 = FEFunction(V1,rand(num_free_dofs(V1)))
#v2 = FEFunction(V2,rand(num_free_dofs(V2)))
#v1c = FEFunction(V1c,rand(num_free_dofs(V1c)))
#v2c = FEFunction(V2c,rand(num_free_dofs(V2c)))
#writevtk(trian,"trian",nsubcells=10,cellfields=["v1"=>v1,"v2"=>v2,"v1c"=>v1c,"v2c"=>v2c])

#using Gridap.Visualization
#writevtk(trian,"trian",nsubcells=10,cellfields=["uh1"=>uh1,"uh2"=>uh2,"e1"=>e1,"e2"=>e2])

l2(u) = u*u
e1_l2 = sqrt(sum(integrate(l2(e1),trian,quad)))
e2_l2 = sqrt(sum(integrate(l2(e2),trian,quad)))

tol = 1.0e-9
@test e1_l2 < tol
@test e2_l2 < tol

end # module


function Base.hash(m::AbstractMonomial, u::UInt)
    nnz = count(!iszero, exponents(m))
    if iszero(nnz)
        hash(1, u)
    elseif isone(nnz) && isone(degree(m))
        # We want the has of `m` to match the hash of `variable(m)`
        # so we ignore the exponent one.
        hash(variable(m), u)
    else
        reduce((u, power) -> iszero(power[2]) ? u : hash(power, u), powers(m), init=u)
    end
end

function Base.hash(t::AbstractTerm, u::UInt)
    if iszero(t)
        hash(0, u)
    elseif isone(coefficient(t))
        # We want the has of `t` to match the hash of `monomial(t)`
        # so we ignore the coefficient one.
        hash(monomial(t), u)
    else
        hash(monomial(t), hash(coefficient(t), u))
    end
end

function Base.hash(p::AbstractPolynomial, u::UInt)
    if iszero(p)
        hash(0, u)
    else
        reduce((u, t) -> hash(t, u), terms(p); init=u)
    end
end


module UnitTests

using Test
using PBC

using PBC.Curve: BN, FP, EP, EP2, curve_gen, LIB
using PBC.Config: ORDER, G1, G2, @EP, @EP2
using PBC.Shamir: evalpoly, LagrangeCoeffGenerator, BarycentricWeightGenerator, lagrange_interpolate_c0

@testset "lagrange_interpolate_c0 - trivial" begin
    ids = 1:2
    for c in (BN(1), curve_gen(EP), curve_gen(EP2))
        poly = [c, c] # derive secret coeffs from c
        shares = Dict(x=>evalpoly(poly, x) for x in ids)
        @test lagrange_interpolate_c0(BarycentricWeightGenerator(keys(shares)), shares) == c
        @test lagrange_interpolate_c0(LagrangeCoeffGenerator(keys(shares)), shares) == c
        @test lagrange_interpolate_c0(shares) == c
    end
end

@testset "Shamir: lagrange_interpolate_c0 - complex" begin
    ids = 1:5
    n = length(ids)
    for c in (BN(123), 123 * curve_gen(EP), 123 * curve_gen(EP2))
        poly = [c, 2c, 3c] # derive secret coeffs from c
        t = length(poly)
        shares = Dict(x=>evalpoly(poly, x) for x in ids)
        for i in 1:n-t+1
            r = i:i+t-1
            @test lagrange_interpolate_c0(BarycentricWeightGenerator(keys(shares)), shares) == c
            @test lagrange_interpolate_c0(LagrangeCoeffGenerator(keys(shares)), shares) == c
            @test lagrange_interpolate_c0(shares) == c
        end
    end
end

@testset "Shamir: LagrangeCoeffGenerator" begin
    coeffs = LagrangeCoeffGenerator()
    # generate first coeff
    @test push!(coeffs, 1) == BN(1)

    # generate second coeff
    @test push!(coeffs, 2) == ORDER - 1
    @test push!(coeffs, 1) == BN(2)

    # generate third share
    @test push!(coeffs, 3) == BN(1)
    @test push!(coeffs, 2) == ORDER - 3
    @test push!(coeffs, 1) == BN(3)

    # delete third share
    delete!(coeffs, 3)
    @test length(coeffs) == 2
    @test push!(coeffs, 2) == ORDER - 1
    @test push!(coeffs, 1) == BN(2)
end

@testset "Shamir: BarycentricWeightGenerator" begin
    weights = BarycentricWeightGenerator()
    # generate first coeff
    @test push!(weights, 1) == ORDER - 1

    # generate second coeff
    @test push!(weights, 2) == BN(2)
    @test push!(weights, 1) == ORDER - 1

    # generate third share
    @test push!(weights, 3) == ORDER - 6
    @test push!(weights, 2) == BN(2)
    @test push!(weights, 1) == ORDER - 2

    # delete third share
    delete!(weights, 3)
    @test length(weights) == 2
    @test push!(weights, 2) == BN(2)
    @test push!(weights, 1) == ORDER - 1
end

@testset "PrivateKey" begin
    # generate random with default RNG
    @test rand(PrivateKey).sk < PBC.Config.PRIME

    # from bytes
    sk = PrivateKey(UInt8[1, 2, 3])
    @test sk.sk < PBC.Config.PRIME
    @test sk.sk.dp[1] == 1 << 16 | 2 << 8 | 3
    @test !iszero(sk) && isvalid(sk)

    # serialization
    @test length(Vector{UInt8}(sk)) == PBC.Config.PRIVATE_KEY_SIZE
    @test PrivateKey(Vector{UInt8}(sk)) == sk
end

@testset "PublicKey" begin
    # generate from secret
    sk = PrivateKey(UInt8[1, 2, 3])
    pk = PublicKey(sk)
    @test !isinf(pk) && isvalid(pk)

    # serialization
    @test length(Vector{UInt8}(pk)) == PBC.Config.PUBLIC_KEY_SIZE
    @test PublicKey(Vector{UInt8}(pk)) == pk
end

@testset "Hash" begin
    hash = Hash("foo")
    @test !isinf(hash) && isvalid(hash)
end

@testset "Identity" begin
    @test !signbit(Identity(-1).id)
    @test !iszero(Identity(PublicKey(PrivateKey(UInt8[1, 2, 3]))).id)
end

@testset "PrivateKeyPoly" begin
    poly = PrivateKeyPoly(2)
    @test length(poly) == 3
    # test iterate
    @test length(collect(poly)) == 3
    # test index
    @test isa(poly[1], BN)
end

@testset "PublicKeyPoly" begin
    skpoly = PrivateKeyPoly(2)
    pkpoly = PublicKeyPoly(skpoly)
    @test length(skpoly) == length(pkpoly)
    # test iterate
    @test length(collect(pkpoly)) == 3
    # test index
    @test isa(pkpoly[1], @EP2)
end

@testset "create_share & validate_share" begin
    # create a secret polynomial
    pk1, pk2 = PublicKey(PrivateKey(UInt8[1, 2, 3])), PublicKey(PrivateKey(UInt8[2, 3, 4]))
    skpoly = PrivateKeyPoly(10)
    share1, share2 = PBC.create_share(skpoly, pk1), PBC.create_share(skpoly, pk2)
    @test share1 != share2
    @test isa(share1, PBC.Model.AbstractPrivateKey)

    pkpoly = PublicKeyPoly(skpoly)
    @test PBC.verify_share(pkpoly, pk1, share1)
    @test PBC.verify_share(pkpoly, pk2, share2)
    # TODO: explore if this can be made to work?
    @test_broken PBC.verify_share(pkpoly, pk1+pk2, share1+share2)
end

@testset "batch validate_share" begin
    pk = PublicKey(PrivateKey(UInt8[1, 2, 3]))
    # create a secret polynomial
    skpoly1, skpoly2 = PrivateKeyPoly(10), PrivateKeyPoly(10)
    # create shares for the same id from two different secret polynomials
    share1, share2 = PBC.create_share(skpoly1, pk), PBC.create_share(skpoly2, pk)
    pkpoly1, pkpoly2 = PublicKeyPoly(skpoly1), PublicKeyPoly(skpoly2)
    # make sure that the aggregate secret share validates with the aggregate public key polynomial
    @test PBC.verify_share(pkpoly1+pkpoly2, pk, share1+share2)
end

@testset "share sign & recover" begin
    alice, bob, charlotte  = PublicKey(rand(PrivateKey)), PublicKey(rand(PrivateKey)), PublicKey(rand(PrivateKey))
    # create a secret polynomial of degree 2 => 3 shares
    skpoly = PrivateKeyPoly(2)
    pkpoly = PublicKeyPoly(skpoly)
    # create three private keys from the secret polynomial
    skalice, skbob, skcharlotte = PBC.create_share(skpoly, alice), PBC.create_share(skpoly, bob), PBC.create_share(skpoly, charlotte)
    # the signature shares are collected in a SignatureShares object
    sigshares = SignatureShares()
    # each user signs with each secret key share, thus creates separate signature shares
    sigshares[alice] = PBC.sign(skalice, "foo")
    sigshares[bob] = PBC.sign(skbob, "foo")
    # not enough signatures - the recovered signature is NOT valid
    @test !PBC.verify(Signature(sigshares), PublicKey(pkpoly), "foo")
    # enough signatures - the recovered signature is valid
    sigshares[charlotte] = PBC.sign(skcharlotte, "foo")
    @test PBC.verify(Signature(sigshares), PublicKey(pkpoly), "foo")
end

@testset "sign & verify" begin
    sk = PrivateKey(UInt8[1, 2, 3])
    sig = PBC.sign(sk, "foo")
    @test !isinf(sig) && isvalid(sig)
    @test PBC.verify(sig, PublicKey(sk), "foo")
    @test !PBC.verify(sig, PublicKey(sk), "bar")
end

@testset "simple aggregation" begin
    sk1, sk2, sk3 = PrivateKey(UInt8[1, 2, 3]), PrivateKey(UInt8[2, 3, 4]), PrivateKey(UInt8[3, 4, 5])
    pk1, pk2, pk3 = PublicKey(sk1), PublicKey(sk2), PublicKey(sk3)
    sig1, sig2, sig3 = PBC.sign(sk1, "foo"), PBC.sign(sk2, "foo"), PBC.sign(sk3, "foo")
    @test PBC.verify(sig1 + sig2, pk1 + pk2, "foo")
    @test !PBC.verify(sig2, pk1 + pk2, "foo")
    @test pk1 - pk1 + pk2 == pk1 + pk2 - pk1
    @test PBC.verify(sig2, pk1 - pk1 + pk2, "foo")
end

@testset "advanced aggregation" begin
    sk1, sk2, sk3 = PrivateKey(UInt8[1, 2, 3]), PrivateKey(UInt8[2, 3, 4]), PrivateKey(UInt8[3, 4, 5])
    pk1, pk2, pk3 = PublicKey(sk1), PublicKey(sk2), PublicKey(sk3)
    msg1, msg2, msg3 = "foo", "bar", "baz"
    sig1, sig2, sig3 = PBC.sign(sk1, msg1), PBC.sign(sk2, msg2), PBC.sign(sk3, msg3)
    @test PBC.verify(sig1 + sig2, (pk1 => Hash(msg1), pk2 => Hash(msg2)))
    @test !PBC.verify(sig1 + sig2, (pk1 => Hash(msg1),))
end

@testset "PrivateKey aggregation" begin
    sk1, sk2 = PrivateKey(UInt8[1, 2, 3]), PrivateKey(UInt8[2, 3, 4])
    pk1, pk2 = PublicKey(sk1), PublicKey(sk2)
    @test PBC.sign(sk1 + sk2, "foo") == PBC.sign(sk1, "foo") + PBC.sign(sk2, "foo")
    @test iszero((sk1 - sk1).sk)
end

@testset "invalid aggregation" begin
    sk = PrivateKey(UInt8[1, 2, 3])
    @test_throws ErrorException PBC.verify(PBC.sign(sk, "foo"), [])
end

@testset "Shamir: evalpoly" begin
    for c in (BN(123), curve_gen(EP), curve_gen(EP2))
        poly = [c, 2c, 3c]
        @test evalpoly(poly, 0) == c
        @test evalpoly(poly, 1) == c + 2c + 3c
        @test evalpoly(poly, 2) == c + 2c * 2 + 3c * 2^2
    end
end

end


using GLVisualize.ComposeBackend, Gadfly
using Colors, GLVisualize

if !isdefined(:runtests)
    window = glscreen()
    composebackend = ComposeBackend.GLVisualizeBackend(window)
end

const static_example = true

using Distributions
x1 = rand(40)
y1 = 4.*x1 .+ 2 .+randn(40)
x2 = rand(40)
y2 = -6.*x2 .+ 3 .+ randn(40)
x  = [x1;x2]
y  = [y1;y2]
col = [fill("Slope 4",40); fill("Slope -6",40)]
p = plot(x=x,y=y,colour=col, Geom.point, Geom.smooth(method=:lm))

draw(composebackend, p)

if !isdefined(:runtests)
	renderloop(window)
end


# ---
# title: 437. Path Sum III
# id: problem437
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Tree
# link: <https://leetcode.com/problems/path-sum-iii/description/>
# hidden: true
# ---
#
# You are given a binary tree in which each node contains an integer value.
#
# Find the number of paths that sum to a given value.
#
# The path does not need to start or end at the root or a leaf, but it must go
# downwards (traveling only from parent nodes to child nodes).
#
# The tree has no more than 1,000 nodes and the values are in the range
# -1,000,000 to 1,000,000.
#
# **Example:**
#
#
#
#     root = [10,5,-3,3,2,null,11,3,-2,null,1], sum = 8
#
#           10
#          /  \
#         **5**   **-3**
#        **/** **\**    **\**
#       **3**   **2**   **11**
#      / \   **\**
#     3  -2   **1**
#
#     Return 3. The paths that sum to 8 are:
#
#     1.  5 -> 3
#     2.  5 -> 2 -> 1
#     3. -3 -> 11
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


@info "using LTVModels"
using LTVModels
@info "using Test, LinearAlgebra, Statistics, Random"
using Test, LinearAlgebra, Statistics, Random
@info "using Plots, Optim"
using Plots, Optim

eye(n) = Matrix{Float64}(I,n,n)
@testset "LTVModels" begin
    @info "Testing LTVModels"

    @testset "findpeaks" begin
        @info "Testing findpeaks"

        y = randn(1000)
        peaks = LTVModels.findpeaks(y,doplot=isinteractive(), filterlength=1)
        for p in peaks
            @test y[p] >= y[max(1,p-1)]
            @test y[p] >= y[min(length(y),p+1)]
        end

        y = randn(1000)
        peaks = LTVModels.findpeaks(y,doplot=isinteractive(), filterlength=1, minh=2)
        for p in peaks
            @test y[p] >= 2
            @test y[p] >= 2
        end

        y = randn(1000)
        peaks = LTVModels.findpeaks(y,doplot=isinteractive(), filterlength=1, minw=10)
        dpeaks = diff(peaks)
        @test all(>=(10), dpeaks)

    end

    @testset "KalmanModel" begin
        @info "Testing KalmanModel"

        T = 10000
        A,B,d,n,m,N = LTVModels.testdata(T=T, σ_state_drift=0.001, σ_param_drift=0.001)

        R1          = 0.001*eye(n^2+n*m) # Increase for faster adaptation
        R2          = 10*eye(n)
        P0          = 10000R1
        @time model = KalmanModel(d,R1,R2,P0,extend=true)

        normA  = [norm(A[:,:,t]) for t                = 1:T]
        normB  = [norm(B[:,:,t]) for t                = 1:T]
        errorA = [norm(A[:,:,t]-model.At[:,:,t]) for t = 1:T]
        errorB = [norm(B[:,:,t]-model.Bt[:,:,t]) for t = 1:T]

        @test sum(normA) > 1.8sum(errorA)
        @test sum(normB) > 10sum(errorB)
        @static isinteractive() && plot([normA errorA normB errorB], lab=["normA" "errA" "normB" "errB"], show=false, layout=2, subplot=1, size=(1500,900))#, yscale=:log10)
        @static isinteractive() && plot!(flatten(A), l=(2,:auto), xlabel="Time index", ylabel="Model coefficients", lab="True",subplot=2, c=:red)
        @static isinteractive() && plot!(flatten(model.At), l=(2,:auto), xlabel="Time index", ylabel="Model coefficients", lab="Estimated", subplot=2, c=:blue)
        P = zeros(N,N,T)
        for i=1:T
            P[:,:,i] .= 0.01eye(N)
        end
        prior = KalmanModel(A,B,P) # Use ground truth as prior
        @time model = KalmanModel(model, prior, d,R1,R2,P0,extend=true)
        errorAp = [norm(A[:,:,t]-model.At[:,:,t]) for t = 1:T]
        errorBp = [norm(B[:,:,t]-model.Bt[:,:,t]) for t = 1:T]
        @static isinteractive() && plot!([errorAp errorBp], lab=["errAp" "errBp"], show=true, subplot=1)

        @test all(errorAp .<= errorA) # We expect this since ground truth was used as prior
        @test all(errorBp .<= errorB)

        @testset "No control signal" begin
            @info "Testing No control signal"
            A,B,d,n,m,N = LTVModels.testdata(T=T, σ_state_drift=0.1, σ_param_drift=0.0003, σ_control=0)
            m = 0
            d = iddata(d.x,zeros(0,length(d)), d.x)
            R1          = 0.0003*eye(n^2) # Increase for faster adaptation
            R2          = 1*eye(n)
            P0          = 10000R1
            D = 1
            for D in [1,2]
                model = KalmanModel(d,R1,R2,P0,extend=true, D=D)

                normA  = [norm(A[:,:,t]) for t                = 1:T]
                errorA = [norm(A[:,:,t]-model.At[:,:,t]) for t = 1:T]

                @test sum(normA) > 4sum(errorA)
                @static isinteractive() && plot([normA errorA], lab=["normA" "errA"], show=false, layout=2, subplot=1, size=(1500,900))#, yscale=:log10)
                @static isinteractive() && plot!(flatten(A), l=(2,:auto), xlabel="Time index", ylabel="Model coefficients", lab="True",subplot=2, c=:red)
                @static isinteractive() && plot!(flatten(model.At), l=(2,:auto), xlabel="Time index", ylabel="Model coefficients", lab="Estimated", subplot=2, c=:blue)
            end

        end

        @testset "LTVAutoRegressive" begin
            @info "Testing LTVAutoRegressive"
            ## Generate chirp signal
            T = 2^14
            e = 0.01randn(T)
            fs = 16_000
            t = range(0,step=1/fs, length=T)
            chirp_f1 = LinRange(2000, 2500, T)
            chirp_f2 = 5000 .+ 200sin.(2pi/T .* (1:T))
            y = sin.(2pi .* chirp_f1 .* t )
            y .+= sin.(2pi .* chirp_f2 .* t )
            yn = y + e
            d = iddata(yn)

            n  = 6
            R1 = eye(n) # Increase for faster adaptation
            R2 = [1e5]
            P0 = 1e4R1
            @time model = LTVAutoRegressive(d,R1,R2,P0,extend=true, D=1)
            RS = LTVModels.rootspectrogram(model,fs)
            @test mean(minimum(abs, RS .-  LinRange(2000, 3000, T), dims=2)) < 20 # for some reason, the increase in frequency is double that of the increase in the chirp. This seems correct as a welch spectrogam shows the same.

            @time model = LTVAutoRegressive(d,R1,R2,P0,extend=true, D=2)
            RS = LTVModels.rootspectrogram(model,fs)
            @test mean(minimum(abs, RS .-  LinRange(2000, 3000, T), dims=2)) < 20

        end

    end

    # Tests ========================================================================

    # Iterative solver =============================================================
    @testset "Statespace fit" begin
        @info "Testing Statespace fit"

        # Generate data

        T_       = 400
        x,xm,u,n,m = LTVModels.testdata(T_)

        # model, cost, steps = fit_statespace_gd(xm,u,20, normType = 1, D = 1, lasso = 1e-8, step=5e-3, momentum=0.99, iters=100, reduction=0.1, extend=true);
        # Profile.clear()

        function callback(k)
            model = LTVModels.statevec2model(SimpleLTVModel,k,n,m,true)
            @static isinteractive() && plot(flatten(model.At), l=(2,:auto), xlabel="Time index", ylabel="Model coefficients", show=true)
        end

        d = iddata(x,u,x)
        dn = iddata(xm,u,xm)
        @time model = LTVModels.fit_admm(SimpleLTVModel(dn,extend=true), dn,17,
        iters    = 20000,
        D        = 1,
        zeroinit = true,
        tol      = 1e-5,
        ridge    = 0,
        cb       = callback);
        # using ProfileView
        # ProfileView.view()
        # model, cost, steps = fit_statespace_gd!(model,xm,u,100, normType = 1, D = 1, lasso = 1e-8, step=5e-3, momentum=0.99, iters=1000, reduction=0.1, extend=true);
        y = predict(model,d);
        e = x[:,2:end] - y[:,1:end-1]
        println("RMS error: ",rms(e))

        At,Bt = model.At,model.Bt
        @static isinteractive() && begin
            plot(flatten(At), l=(2,:auto), xlabel="Time index", ylabel="Model coefficients")
            plot!([1,T_÷2-1], [0.95 0.1; 0 0.95][:]'.*ones(2), l=(:dash,:black, 1))
            plot!([T_÷2,T_], [0.5 0.05; 0 0.5][:]'.*ones(2), l=(:dash,:black, 1), grid=false)
            gui()
        end

        # R1          = 0.1*eye(n^2+n*m) # Increase for faster adaptation
        # R2          = 10*eye(n)
        # P0          = 10000R1
        # modelk = KalmanModel(x,u,R1,R2,P0,extend=true)
        # @static isinteractive() && plot!(flatten(modelk.At), l=(2,:auto), lab="Kalman", c=:red)

        # # savetikz("figs/ss.tex", PyPlot.gcf())#, [" axis lines = middle,enlargelimits = true,"])
        #
        # @static isinteractive() && plot(y, lab="Estimated state values", l=(:solid,), xlabel="Time index", ylabel="State value", grid=false, layout=2)
        # @static isinteractive() && plot!(x[2:end,:], lab="True state values", l=(:dash,))
        # @static isinteractive() && plot!(xm[2:end,:], lab="Measured state values", l=(:dot,))
        # # savetikz("figs/states.tex", PyPlot.gcf())#, [" axis lines = middle,enlargelimits = true,"])

        # @static isinteractive() && plot([cost[1:end-1] steps], lab=["Cost" "Stepsize"],  xscale=:log10, yscale=:log10)

        act = activation(model)
        changepoints = [findmax(act)[2]]
        fit_statespace_constrained(dn,changepoints)

        Random.seed!(0)
        T_       = 200
        x,xm,u,n,m = LTVModels.testdata(T_)
        d = iddata(x,u,x)
        dn = iddata(xm,u,xm)

        # TODO: reenable tests below when figured out error with DiffBase
        model2, res = fit_statespace_gd(dn,5000, normType = 1, D = 2, step=0.01, iters=1000, reduction=0.1, extend=true, opt=LBFGS(m=50));
        y2 = predict(model2,d);
        At2,Bt2 = model2.At,model2.Bt
        e2 = x[:,2:end] - y2[:,1:end-1]
        println("RMS error: ",rms(e2))
        @static isinteractive() && plot(flatten(At2), l=(2,:auto), xlabel="Time index", ylabel="Model coefficients")
        @static isinteractive() && plot!([1,T_÷2-1], [0.95 0.1; 0 0.95][:]'.*ones(2), ylims=(-0.1,1), l=(:dash,:black, 1))
        @static isinteractive() && plot!([T_÷2,T_], [0.5 0.05; 0 0.5][:]'.*ones(2), l=(:dash,:black, 1), grid=false)

        model2, res = fit_statespace_gd(dn,1e5, normType = 2, D = 1, step=0.01, momentum=0.99, iters=10000, reduction=0.01, extend=true, lasso=1e-4);
        y2 = predict(model2,d);
        At2,Bt2 = model2.At,model2.Bt
        e2 = x[:,2:end] - y2[:,1:end-1]
        println("RMS error: ",rms(e2))
        @static isinteractive() && plot(flatten(At2), l=(2,:auto), xlabel="Time index", ylabel="Model coefficients")
        @static isinteractive() && plot!([1,T_÷2-1], [0.95 0.1; 0 0.95][:]'.*ones(2), ylims=(-0.1,1), l=(:dash,:black, 1))
        @static isinteractive() && plot!([T_÷2,T_], [0.5 0.05; 0 0.5][:]'.*ones(2), l=(:dash,:black, 1), grid=false)

        @test norm(rms(e2)) < 2
    end

    function test_gmmmodel()
        function toOrthoNormal(Ti)
            local T = deepcopy(Ti)
            U_,S_,V_ = svd(T[1:3,1:3])
            local R = U_*diagm([1,1,sign(det(U_*V_'))])*V_'
            T[1:3,1:3] = R
            return T
        end

        A,B,x,xnew,u,n,m,N = testdata(T=10000, σ_state_drift=0.001, σ_param_drift=0.001)

        model = GMMModel(x,u,xnew, 5,doplot = false, nTries = 5, d1 = 6)

        xnewhat = predict(model,x,u)

        @static isinteractive() && plot(xnew, lab="x", show=true, layout=n)#, yscale=:log10)
        @static isinteractive() && plot!(xnewhat, lab="xhat")

        fx,fu = df(model,x',u')[1:2]
        normA  = [norm(A[:,:,t]) for t = 1:T]
        normB  = [norm(B[:,:,t]) for t = 1:T]
        errorA = [norm(A[:,:,t]-fx[:,:,t]) for t = 1:T]
        errorB = [norm(B[:,:,t]-fu[:,:,t]) for t = 1:T]
        @static isinteractive() && plot([normA errorA normB errorB], lab=["normA" "errA" "normB" "errB"], show=true)#, yscale=:log10)

    end

@testset "bellman" begin
    @info "Testing bellman"
    LTVModels.benchmark_const(100, 2, true) # Dynamic Programming Bellman
    LTVModels.benchmark_ss(100, 2, true)    # Dynamic Programming Bellman
    LTVModels.benchmark_lin(100, 2, true)
end

end


@testset "detect_gradient_orientation" begin
    @info "Test: detect_gradient_orientation"

    @testset "Numerical" begin
        # Points on a compass in the sequence: E, NE, N, NW, W, SW, S, SE
        # expressed in raster coordinates (increasing rows, increasing columns)
        g₁ = [0, -1, -1, -1, 0, 1, 1, 1]
        g₂ = [1, 1, 0, -1, -1, -1, 0, 1]

        # Measuring from East in counter-clockwise manner in degrees.
        orientation_convention = OrientationConvention(in_radians = false,
                                                       is_clockwise = false,
                                                       compass_direction = 'E')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== [0.0,  45.0,  90.0,  135.0,  180.0,  225.0,  270.0,  315.0])

        # Measuring from East in counter-clockwise manner in radians.
        orientation_convention = OrientationConvention(in_radians = true,
                                                       is_clockwise = false,
                                                       compass_direction = 'E')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== deg2rad.([0.0,  45.0,  90.0,  135.0,  180.0,  225.0,  270.0,  315.0]))

        # Measuring from South in counter-clockwise manner in degrees.
        orientation_convention = OrientationConvention(in_radians = false,
                                                       is_clockwise = false,
                                                       compass_direction = 'S')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== [90.0,  135.0,  180.0,  225.0,  270.0,  315.0,  0.0,  45.0])

        # Measuring from South in counter-clockwise manner in radians.
        orientation_convention = OrientationConvention(in_radians = true,
                                                       is_clockwise = false,
                                                       compass_direction = 'S')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== deg2rad.([90.0,  135.0,  180.0,  225.0,  270.0,  315.0,  0.0,  45.0]))

        # Measuring from West in counter-clockwise manner in degrees.
        orientation_convention = OrientationConvention(in_radians = false,
                                                       is_clockwise = false,
                                                       compass_direction = 'W')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== [180.0,  225.0,  270.0,  315.0,  360.0,  45.0,  90.0,  135.0])

        # Measuring from West in counter-clockwise manner in radians.
        orientation_convention = OrientationConvention(in_radians = true,
                                                       is_clockwise = false,
                                                       compass_direction = 'W')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== deg2rad.([180.0,  225.0,  270.0,  315.0,  360.0,  45.0,  90.0,  135.0]))

        # Measuring from N in counter-clockwise manner degrees.
        orientation_convention = OrientationConvention(in_radians = false,
                                                       is_clockwise = false,
                                                       compass_direction = 'N')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== [270.0,  315.0,  0.0,  45.0,  90.0,  135.0,  180.0,  225.0])

        # Measuring from N in counter-clockwise manner in radians.
        orientation_convention = OrientationConvention(in_radians = true,
                                                       is_clockwise = false,
                                                       compass_direction = 'N')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== deg2rad.([270.0,  315.0,  0.0,  45.0,  90.0,  135.0,  180.0,  225.0]))

        # Measuring from East in clockwise manner.
        orientation_convention = OrientationConvention(in_radians = false,
                                                       is_clockwise = true,
                                                       compass_direction = 'E')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== [0.0,  315.0,  270.0,  225.0,  180.0,  135.0,  90.0,  45.0])

        # Measuring from East in clockwise manner.
        orientation_convention = OrientationConvention(in_radians = true,
                                                       is_clockwise = true,
                                                       compass_direction = 'E')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== deg2rad.([0.0,  315.0,  270.0,  225.0,  180.0,  135.0,  90.0,  45.0]))

        # Measuring from South in clockwise manner in degrees.
        orientation_convention = OrientationConvention(in_radians = false,
                                                       is_clockwise = true,
                                                       compass_direction = 'S')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== [270.0,  225.0,  180.0,  135.0,  90.0,  45.0,  0.0,  315.0])

        # Measuring from South in clockwise manner in radians.
        orientation_convention = OrientationConvention(in_radians = true,
                                                       is_clockwise = true,
                                                       compass_direction = 'S')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== deg2rad.([270.0,  225.0,  180.0,  135.0,  90.0,  45.0,  0.0,  315.0]))

        # Measuring from West in clockwise manner in degrees.
        orientation_convention = OrientationConvention(in_radians = false,
                                                       is_clockwise = true,
                                                       compass_direction = 'W')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== [180.0,  135.0,  90.0,  45.0,  0.0,  315.0,  270.0,  225.0])

        # Measuring from West in clockwise manner in radians.
        orientation_convention = OrientationConvention(in_radians = true,
                                                       is_clockwise = true,
                                                       compass_direction = 'W')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== deg2rad.([180.0,  135.0,  90.0,  45.0,  0.0,  315.0,  270.0,  225.0]))

        # Measuring from North in clockwise manner in degrees.
        orientation_convention = OrientationConvention(in_radians = false,
                                                       is_clockwise = true,
                                                       compass_direction = 'N')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== [90.0,  45.0,  0.0,  315.0,  270.0,  225.0,  180.0,  135.0])

        # Measuring from North in clockwise manner in radians.
        orientation_convention = OrientationConvention(in_radians = true,
                                                       is_clockwise = true,
                                                       compass_direction = 'N')
        out = detect_gradient_orientation(g₁, g₂, orientation_convention)
        @test all(out .== deg2rad.([90.0,  45.0,  0.0,  315.0,  270.0,  225.0,  180.0,  135.0]))

        # Verify default setting.
        # Measuring from South in counter-clockwise manner in radians.
        out = detect_gradient_orientation(g₁, g₂)
        @test all(out .== deg2rad.([90.0,  135.0,  180.0,  225.0,  270.0,  315.0,  0.0,  45.0]))
    end

    @testset "Undefined Angles" begin
        # When the gradient magnitude is close to zero the gradient orientation
        # is undefined and we return a "sentinel" value of 360 degrees or
        # 2π.
        out₁ = detect_gradient_orientation([0.0], [0.0], OrientationConvention(in_radians = false))
        out₂ = detect_gradient_orientation([0.0], [0.0], OrientationConvention(in_radians = true))
        @test out₁ == [360]
        @test out₂ == [2π]

        detect_gradient_orientation!(out₂, [0.0], [0.0])
        @test out₂ == [2π]
    end

    @testset "Warning" begin
        convention = OrientationConvention(compass_direction = 'X')
        @test_logs (:warn, "Unrecognised compass_direction... using a default direction of south (S).")  detect_gradient_orientation([0.0], [0.0], convention)
    end

end


"""
    NoncentralT(ν, λ)
"""
struct NoncentralT{T<:Real} <: ContinuousUnivariateDistribution
    ν::T
    λ::T

    function NoncentralT{T}(ν::T, λ::T) where T
        @check_args(NoncentralT, ν > zero(ν))
        @check_args(NoncentralT, isreal(λ))
        new{T}(ν, λ)
    end
end

NoncentralT(ν::T, λ::T) where {T<:Real} = NoncentralT{T}(ν, λ)
NoncentralT(ν::Real, λ::Real) = NoncentralT(promote(ν, λ)...)
NoncentralT(ν::Integer, λ::Integer) = NoncentralT(Float64(ν), Float64(λ))

@distr_support NoncentralT -Inf Inf

### Conversions
convert(::Type{NoncentralT{T}}, ν::S, λ::S) where {T <: Real, S <: Real} = NoncentralT(T(ν), T(λ))
convert(::Type{NoncentralT{T}}, d::NoncentralT{S}) where {T <: Real, S <: Real} = NoncentralT(T(d.ν), T(d.λ))

### Parameters

params(d::NoncentralT) = (d.ν, d.λ)
@inline partype(d::NoncentralT{T}) where {T<:Real} = T

### Statistics

function mean(d::NoncentralT{T}) where T<:Real
    if d.ν > 1
        isinf(d.ν) ? d.λ :
        sqrt(d.ν/2) * d.λ * gamma((d.ν - 1)/2) / gamma(d.ν/2)
    else
        T(NaN)
    end
end

function var(d::NoncentralT{T}) where T<:Real
    d.ν > 2 ? d.ν*(1 + d.λ^2) / (d.ν - 2) - mean(d)^2 : T(NaN)
end

### Evaluation & Sampling

@_delegate_statsfuns NoncentralT ntdist ν λ

## sampling
function rand(rng::AbstractRNG, d::NoncentralT)
    ν = d.ν
    z = randn(rng)
    v = rand(rng, Chisq(ν))
    (z+d.λ)/sqrt(v/ν)
end


@recipe function f(sim::ConvergenceSimulation)
  if ndims(collect(values(sim.errors))[1])>1 #Monte Carlo
    vals = [mean(x,1)' for x in values(sim.errors)]
  else #Deterministic
    vals = [x for x in values(sim.errors)]
  end
  seriestype --> :path
  label  --> reshape([string(key) for key in keys(sim.errors)],1,length(keys(sim.errors)))
  xguide  --> "Convergence Axis"
  yguide  --> "Error"
  xscale --> :log10
  yscale --> :log10
  sim.convergence_axis, vals
end

@recipe function f(shoot::Shootout)
  seriestype --> :bar
  legend := false
  xguide --> "Algorithms"
  yguide --> "Efficiency"
  shoot.names,shoot.effs
end

@recipe function f(wp::WorkPrecision)
  seriestype --> :path
  label -->  wp.name
  lw --> 3
  yguide --> "Time (s)"
  xguide --> "Error"
  xscale --> :log10
  yscale --> :log10
  wp.errors,wp.times
end

@recipe function f(wp_set::WorkPrecisionSet)
  seriestype --> :path
  label -->  reshape(wp_set.names,1,length(wp_set))
  lw --> 3
  yguide --> "Time (s)"
  xguide --> "Error"
  xscale --> :log10
  yscale --> :log10
  errors = Vector{Any}(0)
  times = Vector{Any}(0)
  for i in 1:length(wp_set)
    push!(errors,wp_set[i].errors)
    push!(times,wp_set[i].times)
  end
  errors,times
end

@recipe function f(tab::ODERKTableau;dx=1/100,dy=1/100,xlim=[-6,1],ylim=[-5,5])
  x = xlim[1]:dx:xlim[2]
  y = ylim[1]:dy:ylim[2]
  f = (u,v)-> abs(stability_region(u+v*im,tab))<1
  seriestype --> :contour
  fill --> true
  cbar --> false
  x,y,f
end