using Polytopes
using LightGraphs
using Test

@testset "Polytopes.jl" begin
    _edges = [[1, 2], [2, 3], [3, 1], [2, 3], [3, 4], [4, 2]]
    _faces = [[1, 2, 3], [4, 5, 6]]
    p = Polytope(_edges, _faces)
    @test rank(p) == 2
    @test nverts(p) == 4
    @test faces(p, 0).graph == SimpleDiGraph(nverts(p))
    @test faces(p, 1) == facets(p) == Polytope([_edges])
    @test faces(p, 2) === p
    @test p.graph == SimpleDiGraph([
        Edge(5, 1),
        Edge(5, 2),
        Edge(6, 2),
        Edge(6, 3),
        Edge(7, 3),
        Edge(7, 1),
        Edge(8, 2),
        Edge(8, 3),
        Edge(9, 3),
        Edge(9, 4),
        Edge(10, 4),
        Edge(10, 2),
        Edge(11, 5),
        Edge(11, 6),
        Edge(11, 7),
        Edge(12, 8),
        Edge(12, 9),
        Edge(12, 10),
    ])
    @test string(p) == "2-polytope with 4 vertices, 6 edges, 2 faces"

    gp = GeometricPolytope(p, map(x -> Vertex(randn(3)), 1:nverts(p)))
    @test embeddim(gp) == 3
    @test gp.polytope == p
    @test string(gp) == "3-dimensional geometric 2-polytope with 4 vertices, 6 edges, 2 faces"
end


# Initial conditions for thermogeostrophic balance for the thermal
# shallow water equations. Adapted for the sphere from the planar
# configuration in section 5.2 of:
#     Eldred, Dubos and Kritsikis, JCP 379, 2018

const H₀ = 5960.0
const U₀ = 20.0
const C₀ = 0.1

function uθ(θϕr)
  θ,ϕ,r = θϕr
  u = U₀*cos(ϕ)
  u
end

# Initial velocity
function u₀(xyz)
  θϕr = xyz2θϕr(xyz)
  u   = uθ(θϕr)
  spherical_to_cartesian_matrix(θϕr)⋅VectorValue(u,0,0)
end

function havg(xyz)
  -1.0*U₀*Ωₑ*xyz[3]*xyz[3]/rₑ/g + H₀
end

# Initial fluid depth
function h₀(xyz)
  θϕr   = xyz2θϕr(xyz)
  θ,ϕ,r = θϕr
  h = havg(xyz)
  Rc = π/8.0
  θc = -π/2.0 # mountain top longitude
  ϕc = +π/6.0 # mountain top latitude
  rc = sqrt((θ - θc)*(θ - θc) + (ϕ - ϕc)*(ϕ - ϕc))
  if rc < Rc
    h = h + 200.0*(1.0 - rc/Rc)
  end
  h
end

function s₀(xyz)
  #h = h₀(xyz)
  h = havg(xyz)
  s = g*(1.0 + C₀*H₀*H₀/h/h)
  s
end

function S₀(xyz)
  h = h₀(xyz)
  s = s₀(xyz)
  S = s*h
  S
end

# Topography
function topography(xyz)
  0.0
end


function fibsq08()
#
#  test polynomial suggested by Goedecker
#
    n = 8;
    p = fib(n)
    p = conv(p,p)
    z = [-.8762862300182460;
         -.6416053894725764-.6063952206057722*im; -.6416053894725764+.6063952206057722*im;
         -.04694032461065018-.9030234661229178*im; -.04694032461065018+.9030234661229178*im;
         .6286732392246423-.7084725692273577*im; .6286732392246423+.7084725692273577*im;
         1.996031179735415];
    z = [z 2*ones(n)]
    p, PolyZeros(z)
end


# ---------------------------------------------------
## CLASSIFICATION METRICS (PROBABILISTIC PREDICTIONS)
# >> CrossEntropy
# >> BriersScore

# -----------------------------------------------------
# cross entropy

struct CrossEntropy{R} <: Measure where R <: AbstractFloat
    eps::R
end
CrossEntropy(;eps=eps()) = CrossEntropy(eps)

metadata_measure(CrossEntropy;
    name                     = "cross_entropy",
    target_scitype           = Vec{<:Finite},
    prediction_type          = :probabilistic,
    orientation              = :loss,
    reports_each_observation = true,
    is_feature_dependent     = false,
    supports_weights         = false,
    docstring                = "Cross entropy loss with probabilities " *
                               "clamped between eps and 1-eps; aliases: " *
                               "`cross_entropy`.",
    distribution_type        = UnivariateFinite)

"""
    cross_entropy

$(docstring(CrossEntropy()))

    CrossEntropy(; eps=eps())
    ce(ŷ, y)

Given an abstract vector of distributions `ŷ` and an abstract vector
of true observations `y`, return the corresponding Cross-Entropy
loss (aka log loss) scores.

Since the score is undefined in the case of the true observation has
predicted probability zero, probablities are clipped between `eps` and
`1-eps` where `eps` can be specified.

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

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

For more information, run `info(cross_entropy)`.
"""
cross_entropy = CrossEntropy()

# for single observation:
_cross_entropy(d, y, eps) = -log(clamp(pdf(d, y), eps, 1 - eps))

function (c::CrossEntropy)(ŷ::Vec{<:UnivariateFinite},
                           y::Vec{<:CategoricalElement})
    check_dimensions(ŷ, y)
    check_pools(ŷ, y)
    return broadcast(_cross_entropy, ŷ, y, c.eps)
end

# -----------------------------------------------------
# brier score

# TODO: support many distributions/samplers D below:

struct BrierScore{D} <: Measure end

# As this measure is parametric, the use of `metadata_measure` is not
# appropriate.

MLJModelInterface.name(::Type{<:BrierScore{D}}) where D =
    "BrierScore{$(string(D))}"
MLJModelInterface.docstring(::Type{<:BrierScore{D}}) where D =
    "Brier proper scoring rule for distributions of type $D; "*
    "aliases: `BrierScore{$D}`"
MLJModelInterface.docstring(::Type{<:BrierScore{<:UnivariateFinite}}) =
    "Brier proper scoring rule for `MultiClass` or `OrderedFactor` data; "*
    "aliases: `BrierScore()`, `BrierScore{UnivariateFinite}`"
MLJModelInterface.docstring(::Type{BrierScore}) =
    "Brier proper scoring rule for various distribution types; " *
    "use `brier_score` for `BrierScore{UnivariateFinite}` "*
"(`Multiclass` or `OrderedFactor` targets)."
MLJModelInterface.target_scitype(::Type{<:BrierScore{D}}) where D =
    AbstractVector{<:Finite}
MLJModelInterface.prediction_type(::Type{<:BrierScore}) = :probabilistic
orientation(::Type{<:BrierScore}) = :score
reports_each_observation(::Type{<:BrierScore}) = true
is_feature_dependent(::Type{<:BrierScore}) = false
MLJModelInterface.supports_weights(::Type{<:BrierScore}) = true
distribution_type(::Type{<:BrierScore{D}}) where D =
    UnivariateFinite

"""
    BrierScore(; distribution=UnivariateFinite)(ŷ, y)

Given an abstract vector of distributions `ŷ` of type `distribution`,
and an abstract vector of true observations `y`, return the
corresponding Brier (aka quadratic) scores.

Currently only `distribution=UnivariateFinite` is supported, which is
applicable to superivised models with `Finite` target scitype. In this
case, 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``

Note that `BrierScore()=BrierScore{UnivariateFinite}` has the alias
`brier_score`.

*Warning.* `BrierScore` defines a true *score* (bigger is better). In
Brier's original 1950 paper, and many other places, it is defined as a
loss, 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.

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

"""
function BrierScore(; distribution=UnivariateFinite)
    distribution == UnivariateFinite ||
        error("Only `UnivariateFinite` Brier scores currently supported. ")
    return BrierScore{distribution}()
end

# For single observations (no checks):

# UnivariateFinite:
function _brier_score(d::UnivariateFinite, y)
    levels = classes(d)
    pvec = broadcast(pdf, d, levels)
    offset = 1 + sum(pvec.^2)
    return 2 * pdf(d, y) - offset
end

# For multiple observations:

# UnivariateFinite:
function (::BrierScore{<:UnivariateFinite})(ŷ::Vec{<:UnivariateFinite},
                                            y::Vec{<:CategoricalElement})
    check_dimensions(ŷ, y)
    check_pools(ŷ, y)
    return broadcast(_brier_score, ŷ, y)
end

function (score::BrierScore{<:UnivariateFinite})(ŷ, y, w::Vec{<:Real})
    check_dimensions(y, w)
    return w .* score(ŷ, y) ./ (sum(w)/length(y))
end

const brier_score = BrierScore()

# ============================================================
## CLASSIFICATION METRICS (DETERMINISTIC PREDICTIONS)

const INVARIANT_LABEL_MULTICLASS = "This metric is invariant to class labelling and can be used for multiclass classification."
const INVARIANT_LABEL_BINARY = "This metric is invariant to class labelling and can be used only for binary classification."
const VARIANT_LABEL_BINARY = "This metric is labelling-dependent and can only be used for binary classification."

# ==============================================================
## MULTICLASS

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

struct MisclassificationRate <: Measure end

metadata_measure(MisclassificationRate;
    name                     = "misclassification_rate",
    target_scitype           = Vec{<:Finite},
    prediction_type          = :deterministic,
    orientation              = :loss,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = true,
    docstring                = "misclassification rate; aliases: " *
                               "`misclassification_rate`, `mcr`.")

"""
    misclassification_rate

$(docstring(MisclassificationRate()))

    misclassification_rate(ŷ, y)
    misclassification_rate(ŷ, y, w)
    misclassification_rate(conf_mat)

Returns the rate of misclassification of the (point) predictions `ŷ`,
given true observations `y`, optionally weighted by the weights
`w`. All three arguments must be abstract vectors of the same length.
A confusion matrix can also be passed as argument.
$INVARIANT_LABEL_MULTICLASS

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

"""
const misclassification_rate = MisclassificationRate()
const mcr = misclassification_rate
const MCR = MisclassificationRate

(::MCR)(ŷ::Vec{<:CategoricalElement},
        y::Vec{<:CategoricalElement}) = mean(y .!= ŷ)

(::MCR)(ŷ::Vec{<:CategoricalElement},
        y::Vec{<:CategoricalElement},
        w::Vec{<:Real}) = sum((y .!= ŷ) .* w) / sum(w)

(::MCR)(cm::ConfusionMatrix) = 1.0 - sum(diag(cm.mat)) / sum(cm.mat)

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

struct Accuracy <: Measure end

metadata_measure(Accuracy;
    name                     = "accuracy",
    target_scitype           = Vec{<:Finite},
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = true,
    docstring                = "Classification accuracy; aliases: `accuracy`.")

"""
    accuracy

$(docstring(Accuracy()))

    accuracy(ŷ, y)
    accuracy(ŷ, y, w)
    accuracy(conf_mat)

Returns the accuracy of the (point) predictions `ŷ`,
given true observations `y`, optionally weighted by the weights
`w`. All three arguments must be abstract vectors of the same length.
$INVARIANT_LABEL_MULTICLASS

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

"""
const accuracy = Accuracy()

(::Accuracy)(args...) = 1.0 - misclassification_rate(args...)
(::Accuracy)(m::ConfusionMatrix) = sum(diag(m.mat)) / sum(m.mat)

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

struct BalancedAccuracy <: Measure end

metadata_measure(BalancedAccuracy;
    name                     = "balanced_accuracy",
    target_scitype           = Vec{<:Finite},
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = true,
    docstring                = "Balanced classification accuracy; aliases: "*
                               "`balanced_accuracy`, `bacc`, `bac`.")

"""
    balanced_accuracy

$(docstring(BalancedAccuracy()))

    balanced_accuracy(ŷ, y [, w])
    balanced_accuracy(conf_mat)

Return the balanced accuracy of the point prediction `ŷ`, given true
observations `y`, optionally weighted by `w`. The balanced accuracy takes
into consideration class imbalance.
All  three arguments must have the same length.
$INVARIANT_LABEL_MULTICLASS

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

"""
const balanced_accuracy = BalancedAccuracy()
const bacc = balanced_accuracy
const bac  = bacc
const BACC = BalancedAccuracy

function (::BACC)(ŷ::Vec{<:CategoricalElement},
                  y::Vec{<:CategoricalElement})
    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{<:CategoricalElement},
                  y::Vec{<:CategoricalElement},
                  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

# ==================================================================
## BINARY AND ORDER-INDEPENDENT

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

struct MatthewsCorrelation <: Measure end

metadata_measure(MatthewsCorrelation;
    name                     = "matthews_correlation",
    target_scitype           = Vec{<:Finite{2}},
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false,
    docstring                = "Matthew's correlation; aliases: " *
                               "`matthews_correlation`, `mcc`")

"""
    matthews_correlation

$(docstring(MatthewsCorrelation()))

    matthews_correlation(ŷ, y)
    matthews_correlation(conf_mat)

Return Matthews' correlation coefficient corresponding to the point
prediction `ŷ`, given true observations `y`.
$INVARIANT_LABEL_MULTICLASS

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

"""
const matthews_correlation = MatthewsCorrelation()
const mcc = matthews_correlation
const MCC = MatthewsCorrelation

function (::MCC)(cm::ConfusionMatrix{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{<:CategoricalElement},
         y::Vec{<:CategoricalElement}) =
             confmat(ŷ, y, warn=false) |> m

# ---------------------------------------------------------
# area under the ROC curve

struct AUC <: Measure end

metadata_measure(AUC;
    name                     = "area_under_curve",
    target_scitype           = Vec{<:Finite},
    prediction_type          = :probabilistic,
    orientation              = :score,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false,
    docstring                = "Area under the ROC curve; "*
       "aliases: `area_under_curve`, `auc`")

"""
    area_under_curve

$(docstring(AUC()))

    area_under_curve(ŷ, y)

Return the area under the receiver operator characteristic (curve),
for probabilistic predictions `ŷ`, given ground truth `y`.
$INVARIANT_LABEL_BINARY

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

"""
const area_under_curve = AUC()
const auc = AUC()

function (::AUC)(ŷ::Vec{<:UnivariateFinite},
                 y::Vec{<:CategoricalElement})
    # implementation drawn from https://www.ibm.com/developerworks/community/blogs/jfp/entry/Fast_Computation_of_AUC_ROC_score?lang=en
    lab_pos = levels(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     = 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

# ==========================================================================
## BINARY AND ORDER DEPENDENT

const CM2 = ConfusionMatrix{2}

# --------------------------------------------------------------------------
# F_β-Score

"""
    FScore{β}(rev=nothing)

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

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

    FScore{β}(ŷ, y)

Evaluate ``F_β`` score on observations ,`ŷ`, given ground truth values, `y`.

By default, the second element of `levels(y)` is designated as
`true`. To reverse roles, use `FScore{β}(rev=true)` instead of
`FScore{β}`.

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

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

FScore{β}(; rev=nothing) where β = FScore{β}(rev)

metadata_measure(FScore;
    target_scitype           = Vec{<:Finite},
    prediction_type          = :deterministic,
    orientation              = :score,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false)

MLJModelInterface.name(::Type{<:FScore{β}}) where β = "FScore{$β}"
MLJModelInterface.name(::Type{FScore})            = "FScore" # for registry
MLJModelInterface.docstring(::Type{<:FScore})       = "F_β score; aliases: " *
                                        "`FScore{β}`, `f1score=FScore{1}`"
const f1score      = FScore{1}()

# -------------------------------------------------------------------------
# truepositive, true_negative, etc

for M in (:TruePositive, :TrueNegative, :FalsePositive, :FalseNegative,
          :TruePositiveRate, :TrueNegativeRate, :FalsePositiveRate,
          :FalseNegativeRate, :FalseDiscoveryRate, :Precision, :NPV)
    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},
    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},
    prediction_type          = :deterministic,
    orientation              = :loss,
    reports_each_observation = false,
    is_feature_dependent     = false,
    supports_weights         = false)

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

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

# adjustments
MMI.name(::Type{<:TruePositive})       = "true_positive"
MMI.docstring(::Type{<:TruePositive})  = "Number of true positives; " *
                                         "aliases: `true_positive`, `truepositive`."
MMI.name(::Type{<:TrueNegative})       = "true_negative"
MMI.docstring(::Type{<:TrueNegative})  = "Number of true negatives; " *
                                         "aliases: `true_negative`, `truenegative`."
MMI.name(::Type{<:FalsePositive})      = "false_positive"
MMI.docstring(::Type{<:FalsePositive}) = "Number of false positives; " *
                                         "aliases: `false_positive`, `falsepositive`."
MMI.name(::Type{<:FalseNegative})      = "false_negative"
MMI.docstring(::Type{<:FalseNegative}) = "Number of false negatives; " *
                                         "aliases: `false_negative`, `falsenegative`."

MMI.name(::Type{<:TruePositiveRate})      = "true_positive_rate"
MMI.docstring(::Type{<:TruePositiveRate}) = "True positive rate; aliases: " *
                               "`true_positive_rate`, `truepositive_rate`, `tpr`, `sensitivity`, " *
                               "`recall`, `hit_rate`."
MMI.name(::Type{<:TrueNegativeRate})      = "true_negative_rate"
MMI.docstring(::Type{<:TrueNegativeRate}) = "true negative rate; aliases: " *
                               "`true_negative_rate`, `truenegative_rate`, `tnr`, `specificity`, " *
                               "`selectivity`."
MMI.name(::Type{<:FalsePositiveRate})      = "false_positive_rate"
MMI.docstring(::Type{<:FalsePositiveRate}) = "false positive rate; aliases: " *
                               "`false_positive_rate`, `falsepositive_rate`, `fpr`, `fallout`."
MMI.name(::Type{<:FalseNegativeRate})      = "false_negative_rate"
MMI.docstring(::Type{<:FalseNegativeRate}) = "false negative rate; aliases: " *
                               "`false_negative_rate`, `falsenegative_rate`, `fnr`, `miss_rate`."
MMI.name(::Type{<:FalseDiscoveryRate})      = "false_discovery_rate"
MMI.docstring(::Type{<:FalseDiscoveryRate}) = "false discovery rate; "*
                               "aliases: `false_discovery_rate`, `falsediscovery_rate`, `fdr`."
MMI.name(::Type{<:NPV})      = "negative_predictive_value"
MMI.docstring(::Type{<:NPV}) = "negative predictive value; aliases: " *
                               "`negative_predictive_value`, `negativepredictive_value`, `npv`."

MMI.name(::Type{<:Precision})         = "positive_predictive_value"
MMI.docstring(::Type{<:Precision})    = "positive predictive value "*
  "(aka precision); aliases: `positive_predictive_value`, `ppv`, `Precision()`, `positivepredictive_value`. "

"""
    true_positive

$(docstring(TruePositive()))

    true_positive(ŷ, y)

Number of true positives for observations `ŷ` and ground truth
`y`. Assigns `false` to first element of `levels(y)`. To reverse roles,
use `TruePositive(rev=true)` instead of `true_positive`.

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

"""
const true_positive = TruePositive()
const tp = TruePositive()
const truepositive  = TruePositive()

"""
    true_negative

$(docstring(TrueNegative()))

    true_negative(ŷ, y)

Number of true negatives for observations `ŷ` and ground truth
`y`. Assigns `false` to first element of `levels(y)`. To reverse roles,
use `TrueNegative(rev=true)` instead of `true_negative`.

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

"""
const true_negative = TrueNegative()
const tn = TrueNegative()
const truenegative  = TrueNegative()

"""
    false_positive

$(docstring(FalsePositive()))

    false_positive(ŷ, y)

Number of false positives for observations `ŷ` and ground truth
`y`. Assigns `false` to first element of `levels(y)`. To reverse roles,
use `FalsePositive(rev=true)` instead of `false_positive`.

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

"""
const false_positive = FalsePositive()
const fp = FalsePositive()
const falsepositive = FalsePositive()

"""
    false_negative

$(docstring(FalseNegative()))

    false_negative(ŷ, y)

Number of false positives for observations `ŷ` and ground truth
`y`. Assigns `false` to first element of `levels(y)`. To reverse roles,
use `FalseNegative(rev=true)` instead of `false_negative`.

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

"""
const false_negative = FalseNegative()
const fn = FalseNegative()
const falsenegative = FalseNegative()

"""
    true_positive_rate

$(docstring(TruePositiveRate()))

    true_positive_rate(ŷ, y)

True positive rate for observations `ŷ` and ground truth `y`. Assigns
`false` to first element of `levels(y)`. To reverse roles, use
`TruePositiveRate(rev=true)` instead of `true_positive_rate`.

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

"""
const true_positive_rate = TruePositiveRate()
const tpr = TruePositiveRate()
const TPR = TruePositiveRate
const truepositive_rate  = TPR()
const recall = TPR()
const Recall = TPR
const sensitivity  = recall
const hit_rate     = recall

"""
    true_negative_rate

$(docstring(TrueNegativeRate()))

    true_negative_rate(ŷ, y)

True negative rate for observations `ŷ` and ground truth `y`. Assigns
`false` to first element of `levels(y)`. To reverse roles, use
`TrueNegativeRate(rev=true)` instead of `true_negative_rate`.

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

"""
const true_negative_rate = TrueNegativeRate()
const tnr = TrueNegativeRate()
const TNR = TrueNegativeRate
const truenegative_rate  = TNR()
const Specificity = TNR
const specificity  = truenegative_rate
const selectivity  = specificity

"""
    false_positive_rate

$(docstring(FalsePositiveRate()))

    false_positive_rate(ŷ, y)

False positive rate for observations `ŷ` and ground truth `y`. Assigns
`false` to first element of `levels(y)`. To reverse roles, use
`FalsePositiveRate(rev=true)` instead of `false_positive_rate`.

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

"""
const false_positive_rate = FalsePositiveRate()
const fpr = FalsePositiveRate()
const FPR = FalsePositiveRate
const falsepositive_rate = FPR()
const fallout      = falsepositive_rate

"""
    false_negative_rate

$(docstring(FalseNegativeRate()))

    false_negative_rate(ŷ, y)

False negative rate for observations `ŷ` and ground truth `y`. Assigns
`false` to first element of `levels(y)`. To reverse roles, use
`FalseNegativeRate(rev=true)` instead of `false_negative_rate`.

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

"""
const false_negative_rate = FalseNegativeRate()
const fnr = FalseNegativeRate()
const FNR = FalseNegativeRate
const falsenegative_rate = FNR()

"""
    false_discovery_rate

$(docstring(FalseDiscoveryRate()))

    false_discovery_rate(ŷ, y)

False discovery rate for observations `ŷ` and ground truth `y`. Assigns
`false` to first element of `levels(y)`. To reverse roles, use
`FalseDiscoveryRate(rev=true)` instead of `false_discovery_rate`.

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

"""
const false_discovery_rate = FalseDiscoveryRate()
const fdr = FalseDiscoveryRate()
const FDR = FalseDiscoveryRate
const falsediscovery_rate = FDR()
const miss_rate    = falsenegative_rate

"""
    negative_predictive_value

$(docstring(NPV()))

    negative_predictive_value(ŷ, y)

Negative predictive value for observations `ŷ` and ground truth
`y`. Assigns `false` to first element of `levels(y)`. To reverse roles,
use `NPV(rev=true)` instead of `negative_predictive_value`.

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

"""
const negative_predictive_value = NPV()
const npv = NPV()
const negativepredictive_value = NPV()

"""
    positive_predictive_value

$(docstring(Precision()))

    positive_predictive_value(ŷ, y)

Positive predictive value for observations `ŷ` and ground truth
`y`. Assigns `false` to first element of `levels(y)`. To reverse roles,
use `Precision(rev=true)` instead of `positive_predictive_value`.

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

"""
const positive_predictive_value = Precision()
const ppv = Precision()
const positivepredictive_value = Precision()
const PPV = Precision

## INTERNAL FUNCTIONS ON CONFUSION MATRIX

_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))

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

(::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)

function (::FScore{β})(m::CM2) where β
    β2   = β^2
    prec = precision(m)
    rec  = recall(m)
    return (1 + β2) * (prec * rec) / (β2 * prec + rec)
end

## Callables on vectors

for M in (TruePositive, TrueNegative, FalsePositive, FalseNegative,
          TPR, TNR, FPR, FNR,
          FDR, Precision, NPV, FScore)
    (m::M)(ŷ, y) = confmat(ŷ, y; rev=m.rev) |> m
end

# specify this as `precision` is in Base and so is ambiguous
Base.precision(m::CM2) = m |> Precision()
Base.precision(ŷ, y)   = confmat(ŷ, y) |> Precision()

## 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

"""
    tprs, fprs, 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{<:CategoricalElement})

    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


function CoordinateDescent3{F<:AbstractFloat}(X::Array{F}, Y::Array{F}, LASSOpath::Array{F}, λ_max::F, w::Array{F}, σ::Array{F})
    ## Number of available lambdas: Granularity of the regulation parameter
    MaxIterλ = length(LASSOpath)

    n, p = size(X)

    ## Preallocating vector β_lasso for speed
    β_lasso = zeros(p, MaxIterλ)

    ## Initializing parameters
    β_tilde = zeros(p,1)

    ## Active set
    A= Int[]

    innerX = X' * X

    ## Calcula os produtos internos das colunas de X com Y
    innerY = sum(X.*repmat(Y,1,p),1)

    bestB = 1
    BICold = 100000000
    teste = 0

    β_ols = zeros(p,1)

    ## First for runs through the LASSO path
    for cont = 1:MaxIterλ
        ## Computing next lambda
        λ = LASSOpath[cont] * λ_max
        lastchange = 0
        it_while = 1

        while (teste != 2)
            ##  Reference active set
            A_anterior = A

            ##  Cycle
            if it_while == 1
                j_range = collect(1:p)'
            else
                #j_range = A_actual
                j_range = collect(1:p)'
            end

            for j = j_range

                aux, λEff, β_ols = CoordinateDescentCore(A, innerX, j, β_tilde, β_ols, n, innerY, λ, σ)

                #updateCandidates{T<:Signed, F<:AbstractFloat}(λEff::F, β_ols::Array{F},
                #                 j::T, A_actual::Array{T}, lastchange::T, β_tilde::Array{F}; w::Array{F} = [])

                A, β_tilde, lastchange, breakThis = updateCandidates(λEff, β_ols, j, A, lastchange, β_tilde, w = w)

                if breakThis == 1
                    break
                end

            end
            ##  Check if the active set has been altered in the last cycle
            teste = convCheck(A_anterior, A)
            it_while = it_while + 1

        end
        teste = 0

        BICnew = n * ( log(var( Y - X*β_tilde )) ) + sum( abs(β_tilde./σ) .> 1e-3 ) * log(n)
        if BICnew < BICold
            BICold = BICnew
            bestB = cont
        end

        β_lasso[:,cont] = β_tilde
    end

    return β_lasso, bestB
end


givensids(P) = [ (idx1, idx2) for idx1 = 1:P-1 for idx2 = (idx1+1):P ]

ngivensangles(n::Integer) = fld(n*(n-1),2)

mat2rotations(args...) = mat2rotations_lowmemory(args...)

function mat2rotations_reference(mtx::AbstractMatrix{T}) where T <: Real
    P = size(mtx, 1)

    ids = givensids(P)
    res = similar(mtx, length(ids))

    for (nr, (idx1, idx2)) in enumerate(ids)
        a = givens(mtx, idx1, idx2, idx1)
        g = a[1]

        res[nr] = atan(g.s, g.c)

        R = Matrix(I,P,P)
        R[g.i1, g.i1] =  g.c
        R[g.i1, g.i2] =  g.s
        R[g.i2, g.i1] = -g.s
        R[g.i2, g.i2] =  g.c

        mtx .= R*mtx
    end
    (res, round.(diag(mtx)))
end

function mat2rotations_lowmemory(mtx::AbstractMatrix{T}) where {T<:Real}
    ids = givensids(size(mtx, 1))
    res = similar(mtx, length(ids))

    mtx = Array(mtx)
    for (nr, (idx1, idx2)) in enumerate(ids)
        a = givens(mtx, idx1, idx2, idx1)
        g = a[1]

        res[nr] = atan(g.s, g.c)

        row1 = mtx[g.i1, :]
        vrw2 = @view mtx[g.i2, :]
        mtx[g.i1, :] =  g.c * row1 + g.s * vrw2
        mtx[g.i2, :] = -g.s * row1 + g.c * vrw2
    end
    (res, round.(diag(mtx)))
end

function rotations2mat(θs::AbstractArray, sig::AbstractArray)
    rotations2mat(θs, sig, round(Integer, (1 + sqrt(1+8*length(θs))) / 2))
end

rotations2mat(θs, sig, P) = rotations2mat_lowmemory(θs, sig, P)

function rotations2mat_reference(θs::AbstractArray{TA}, sig::AbstractArray{TS}, P::Integer) where {TA<:Real,TS<:Number}
    mtx = Matrix{TA}(I,P,P)

    for (nr, (idx1, idx2)) in enumerate(givensids(P))
        c, s = cos(θs[nr]), sin(θs[nr])

        R = Matrix(I,P,P)
        R[idx1, idx1] =  c
        R[idx1, idx2] = -s
        R[idx2, idx1] =  s
        R[idx2, idx2] =  c

        mtx .= mtx*R
    end
    mtx * diagm(0 => sig)
end

function rotations2mat_lowmemory(θs::AbstractArray{TA}, sig::AbstractArray{TS}, P::Integer) where {TA<:Real,TS<:Number}
    mtx = Matrix{TA}(I,P,P)

    for (nr, (idx1, idx2)) in enumerate(givensids(P))
        s, c = sincos(θs[nr])

        col1 = mtx[:, idx1]
        vcl2 = @view mtx[:, idx2]

        mtx[:, idx1] =  c * col1 + s * vcl2
        mtx[:, idx2] = -s * col1 + c * vcl2
    end
    mtx * diagm(0 => sig)
end

graph_one_factorization(P::Integer) = graph_one_factorization(Val(isodd(P)), P)

# When P is odd.
function graph_one_factorization(::Val{true}, P::Integer)
    map(0:P-1) do pivot_idx
        map(1:fld(P,2)) do offset
            ids = [ pivot_idx - offset, pivot_idx + offset ]
            (sort(mod.(ids, P) .+ 1)...,)
        end
    end
end

# When P is even.
function graph_one_factorization(::Val{false}, P::Integer)
    oddgf = graph_one_factorization(P-1)
    [ [ elem..., (idx, P) ] for (idx, elem) in enumerate(oddgf) ]
end

# function mat2rotations_permutated(mtx::AbstractMatrix{T}) where T <: Real
#     P = size(mtx, 1)
#
#     res = similar(mtx, fld(P*(P-1),2))
#     mtx = Array(mtx)
#
#     for (nr, (idx1, idx2)) in enumerate(vcat(graph_one_factorization(P)...))
#         a = givens(mtx, idx1, idx2, idx1)
#         g = a[1]
#
#         res[nr] = atan(g.s, g.c)
#
#         R = Matrix{T}(I,P,P)
#         R[g.i1, g.i1] =  g.c
#         R[g.i1, g.i2] =  g.s
#         R[g.i2, g.i1] = -g.s
#         R[g.i2, g.i2] =  g.c
#         mtx .= R*mtx
#     end
#     (res, round.(diag(mtx)))
# end
#
# function rotations2mat_permutated(θs::AbstractArray{TA}, sig::AbstractArray{TS}, P::Integer) where {TA<:Real,TS<:Number}
#     θrsp = reshape(θs, 2*cld(P,2)-1, fld(P,2))
#     mtx = mapreduce(*, enumerate(graph_one_factorization(P))) do (blk, ids)
#         R = Matrix{TA}(I,P,P)
#         for (nr, (idx1, idx2)) in enumerate(ids)
#             s = sin(θrsp[blk, nr])
#             c = cos(θrsp[blk, nr])
#             R[idx1, idx1] =  c
#             R[idx1, idx2] = -s
#             R[idx2, idx1] =  s
#             R[idx2, idx2] =  c
#         end
#         R
#     end
#     mtx * diagm(0 => sig)
# end


export constructtriad
export constructtorus
export constructsphere
export constructcylinder
export constructbox
export constructwhirl
export constructframe
export getplane

"""
    constructtriad(q [; length])

Construct a triad with the given configuration `q` and `length`.
"""
function constructtriad(q::Biquaternion; length::Float64 = 1.0)
    map(x -> x + gettranslation(q), rotate(ℝ³.([[0; 0; 0], [length; 0; 0],
                                                [0; 0; 0], [0; length; 0],
                                                [0; 0; 0], [0; 0; length]]),
                                           getrotation(q)))
end

"""
    constructtorus(q, r, R)

Construct a torus of revolution with the given configuration `q`, the smaller radius `r`
and the bigger radius `R`.
"""
function constructtorus(q::Biquaternion,
                        r::Real,
                        R::Real;
                        segments::Int = 36)
    array = Array{ℝ³,2}(undef, segments, segments)
    for i in 1:segments
        for j in 1:segments
            ϕ = i * 2pi / (segments - 1)
            θ = j * 2pi / (segments - 1)
            x₁ = (R + r * cos(ϕ)) * cos(θ)
            x₂ = (R + r * cos(ϕ)) * sin(θ)
            x₃ = r * sin(ϕ)
            array[i, j] = ℝ³(x₁, x₂, x₃)
        end
    end
    map(x -> x + gettranslation(q), rotate(array, getrotation(q)))
end

"""
    constructsphere(q, radius)

Construct a sphere with the given configuration `q` and `radius`.
"""
function constructsphere(q::Biquaternion,
                         radius::Real;
                         segments::Int = 36)
    array = Array{ℝ³,2}(undef, segments, segments)
    lspace = collect(range(float(-pi), stop = float(pi), length = segments))
    for i in 1:segments
        for j in 1:segments
            ϕ = lspace[i]
            θ = lspace[j] / 2
            array[i, j] = ℝ³(Cartesian(Geographic(radius, ϕ, θ)))
        end
    end
    map(x -> x + gettranslation(q), rotate(array, getrotation(q)))
end

"""
    constructcylinder(q, height, radius)

Construct a cylinder with the given configuration `q`, `height` and `radius`.
"""
function constructcylinder(q::Biquaternion,
                           height::Float64,
                           radius::Float64;
                           segments::Int = 36)
    array = Array{ℝ³,2}(undef, segments, segments)
    lspace = collect(range(float(-pi), stop = float(pi), length = segments))
    for i in 1:segments
        for j in 1:segments
            α = lspace[j]
            x₁ = radius * cos(α)
            x₂ = radius * sin(α)
            x₃ = (i - 1) * height / (segments - 1)
            array[i, j] = ℝ³(x₁, x₂, x₃)
        end
    end
    map(x -> x + gettranslation(q), rotate(array, getrotation(q)))
end

"""
    constructbox(p, q, s)

Construct a box with the given translation `p`, rotation `q` and scale `s`.
"""
function constructbox(p::ℝ³, q::Quaternion, s::ℝ³)
    x, y, z = vec(s) ./ 2
    points = [ℝ³(-x, -y, -z) ℝ³(x, -y, -z) ℝ³(-x, -y, -z) ℝ³(x, -y, -z);
              ℝ³(-x, -y, +z) ℝ³(x, -y, +z) ℝ³(-x, -y, +z) ℝ³(x, -y, +z);
              ℝ³(-x, +y, +z) ℝ³(x, +y, +z) ℝ³(-x, +y, +z) ℝ³(x, +y, +z);
              ℝ³(-x, +y, -z) ℝ³(x, +y, -z) ℝ³(-x, +y, -z) ℝ³(x, +y, -z)]
    map(x -> x + p, rotate(points, q))
end

"""
    constructwhirl(points,
                   s2tos3map,
                   s2tos2map;
                   [top, [bottom, [s3rotation, [config, [segments [, scale]]]]]])

Construct a whirl with the given `points` in the base space, , a map from the base space to
the total space `s2tos3map`, a conformal map that takes the base space to itself `s2tos2map`
, `top` and `bottom` U(1) actions, S³ rotation `s3rotation`, configuration `config`, the
number of `segments` and `scale`.
"""
function constructwhirl(points::Array{<:S²,1},
                        s2tos3map,
                        s2tos2map;
                        top::S¹ = U1(-pi),
                        bottom::S¹ = U1(pi),
                        s3rotation::S³ = Quaternion(1, 0, 0, 0),
                        config::Biquaternion = Biquaternion(ℝ³(0, 0, 0)),
                        segments::Int = 36,
                        scale::Real = 1.0)
    array = Array{ℝ³,2}(undef, segments, length(points))
    lspace = range(angle(top), stop = angle(bottom), length = segments)
    for (i, α) in enumerate(lspace)
        for (j, p) in enumerate(points)
            array[i, j] = compressedλmap(rotate(S¹action(s2tos3map(s2tos2map(p)), U1(α)),
                                                s3rotation)) * scale
        end
    end
    applyconfig(array, config)
end

"""
    constructframe(circle, s2tos3map, s2tos2map; [s3rotation, [config, [segments]]])

Construct a frame that is like an S³ cross-section under stereographic projection with the
given `circle` that determines the phase angle in the fiber space, a map from the base space
to the total space `s2tos3map`, a conformal map that takes the base space to itself
`s2tos2map`, S³ rotation `s3rotation`, configuration `config` and the number of `segments`.
"""
function constructframe(circle::S¹,
                        s2tos3map,
                        s2tos2map;
                        s3rotation::S³ = Quaternion(1, 0, 0, 0),
                        config::Biquaternion = Biquaternion(ℝ³(0, 0, 0)),
                        segments::Int = 36)
    array = Array{ℝ³,2}(undef, segments, segments)
    lspaceϕ = collect(range(float(-pi), stop = float(pi), length = segments))
    lspaceθ = collect(range(float(pi / 2), stop = float(-pi / 2), length = segments))
    for (i, θ) in enumerate(lspaceθ)
        for (j, ϕ) in enumerate(lspaceϕ)
            p = Geographic(1, ϕ, θ)
            array[i, j] = compressedλmap(rotate(S¹action(s2tos3map(s2tos2map(p)), circle),
                                                s3rotation))
        end
    end
    applyconfig(array, config)
end

getplane(p::Array{Float64}, q::Quaternion, s::Array{Float64}) = begin
    plane = Array{Float64,3}(undef, 2, 2, 3)
    plane[1, 1, :] = vec(rotate(ℝ³([0.0; -1.0; 1.0]), q)) .* s + p
    plane[1, 2, :] = vec(rotate(ℝ³([0.0; 1.0; 1.0]), q)) .* s + p
    plane[2, 1, :] = vec(rotate(ℝ³([0.0; -1.0; -1.0]), q)) .* s + p
    plane[2, 2, :] = vec(rotate(ℝ³([0.0; 1.0; -1.0]), q)) .* s + p
    plane
end


## calculatoin_tools.jl
export pack
export radial_mean, Δ_phase
export moment_proj_normed

"""
    radial_mean(data; maxbin=nothing, bin_step=nothing, pixelsize=nothing)

Calculates the radial mean of a dataset `data`.
Returns a tuple of the radial_mean and the bin_centers.

 ## Arguments:
+ `data`: data to radially average
+ `maxbin`: a maximum bin value
+ `bin_step`: optionally defines the step between the bins

## Examples
```julia-repl
julia> radial_mean([5, 1, 1, 3, 15])
([1.0, 2.0, 10.0], [0.5, 1.5, 2.5])
```
"""
function radial_mean(data; nbins=nothing, bin_step=nothing, offset=CtrFT, scale=nothing)
    if isnothing(scale)
        scale=Tuple(ones(Int,ndims(data)))
    end
    if isnothing(bin_step)
        bin_step=maximum(scale)
    end

    idx = round.(Int, rr(size(data), offset=offset, scale=scale./bin_step)) .+ 1
    if isnothing(nbins)
        nbins = maximum(idx)
    else
        idx[idx.>nbins] .= nbins
    end

    bincenters = bin_step .* collect(0.5:nbins-0.5)

    myevents = zeros(Float64, nbins)
    mysum = zeros(Float64, nbins)
    idxlist = idx[:]
    v = @view mysum[idxlist]
    v .+= data[:]
    v = @view myevents[idxlist]
    v .+= 1
    myevents[myevents .== 0] .= 1
    return mysum ./ myevents, bincenters
end

"""
    Δ_phase(arr, dim)

Calculates the relative phase slope along dimension `dim` of a non-zero array `arr` without wrap-around problems.

Arguments:
+ arr: array of which to evaluate the phase slope
+ dim: dimension along which to evaluate the slope
"""
function Δ_phase(arr, dim)
    no_shift = ((n==dim) ? (1:size(arr,dim)-1) : (:) for n in 1:ndims(arr))
    shift = ((n==dim) ? (2:size(arr,dim)) : (:) for n in 1:ndims(arr))
    return angle.(arr[no_shift...] ./ arr[shift...]) # this is a complex-valued division to obtain the relative phase angle!
end

"""
    moment_proj(data, h=3; pdims=3)

Performs a projection over the `h`-th moment of the data along dimension(s) `pdims`.
See also [`moment_proj_normed`](@ref MicroscopyTools.moment_proj_normed)

## Example
```jldoctest
julia> MicroscopyTools.moment_proj([1 3; 2 4], 2, pdims=(1,))
1×2 Matrix{Float64}:
 0.25  0.25
```
"""
function moment_proj(data, h=3; pdims=3)
    mdata = mean(data, dims=pdims)
    return mean((data .- mdata).^h, dims=pdims)
end

"""
    moment_proj_normed(data, h=3; pdims=3)

Performs a projection over the `h`-th moment of the data along dimension(s) `pdims` normed by the variance.

## Example
```jldoctest
julia> MicroscopyTools.moment_proj_normed([1 3; 2 4], 2, pdims=(1,))
1×2 Matrix{Float64}:
 0.5  0
```
"""
function moment_proj_normed(data, h=3; pdims=3)
    moment_proj(data, h, pdims=pdims) ./ (moment_proj(data,2, pdims=pdims) .^((h-1)/2))
end


using Devectorize
using Gadfly

"""
Linear Regression Cost function with Regularisation

Parameters:
	params      = vector containing rolled up X and Θ
		X       = array of features for each movie
		Θ       = array of weights for each user
	Y           = array of ratings for each movie
	R           = array of the existance of ratings for each movie
	numUsers    = number of users
	numMovies   = number of movies
	numFeatures = number of features
	λ           = regularisation constant

Returns:
	J = linear regression cost of the given weights with the given regularisation constant
"""
function cost(params, Y::Array{Int}, numUsers, numMovies, numFeatures, λ)
	# Unfold params -> X, Θ
	X, Θ = unroll(params, numUsers, numMovies, numFeatures)

	# Calculate cost
	J = 0
 	@simd for i = 1:size(Y)[2]
		@inbounds J += (dot(X[Y[2, i], :][:], Θ[Y[1, i], :][:]) - Y[3, i])^2
	end
	J = 1/2 * J
	# add regularisation
	J += λ/2 * (sum(Θ.^2) + sum(X.^2))

	return J
end

"""
Unroll parameter vector into X and Θ

Parameters:
	params = vector containing rolled up X and Θ
	numUsers    = number of users
	numMovies   = number of movies
	numFeatures = number of features

Returns:
	X = array of features for each movie
	Θ = array of weights for each user
"""
function unroll(params, numUsers, numMovies, numFeatures)
	X = reshape(params[1 : numMovies * numFeatures], numMovies, numFeatures)
	Θ = reshape(params[numMovies * numFeatures + 1 : end], numUsers, numFeatures)
	return X, Θ
end

"""
Roll up X and Θ into single vector

Parameters:
	X = array of features for each movie
	Θ = array of weights for each user

Returns:
	params = vector containing rolled up X and Θ
"""
function rollup(X, Θ)
	return params = vcat(X..., Θ...)
end

function predict(movie, me, others, Y, users)
	meIndices = find(Y[1, :] .== me)
	meMovies = Y[2, meIndices]
	weightSum = 0.0
	summation = 0.0
	otherIndices = Int[]
	tempMeIndices = Int[]
	for other in others
		otherIndices = users[other]
		otherMovies = Y[2, otherIndices]
		jointMovies = intersect(meMovies, otherMovies)
		if length(jointMovies) == 0
			continue
		end
		tempMeIndices = filter(x -> in(Y[2, x], jointMovies), meIndices)
		tempOtherIndices = filter(x -> in(Y[2, x], jointMovies), otherIndices)

		meRatings = Y[3, tempMeIndices]
		meAverage = mean(meRatings)
		otherRating = Y[3, find(Y[2, otherIndices] .== movie)][1]
		otherRatings = Y[3, tempOtherIndices]
		otherAverage = mean(otherRatings)

		temp = sqrt(sum((meRatings - meAverage).^2)) * sqrt(sum((otherRatings - otherAverage).^2))
		if temp == 0
			continue
		end

		@devec weight = sum((meRatings - meAverage).*(otherRatings - otherAverage)) ./ temp

		summation += (otherRating - otherAverage) * weight
		weightSum += weight
	end

	meRatings = Y[3, meIndices]
	meAverage = mean(meRatings)

	return meAverage + summation / weightSum
end

"""
Entry point

Params:
	dataPath = path to() CSVs
"""
function main()
	dataPath = ""
	if ARGS[1] == "small"
		dataPath = "../ml-latest-small/fixed/"
	elseif ARGS[1] == "full"
		dataPath = "../ml-latest/fixed/"
	end

	# Load data from CSV - using wc to speed up loading by preallocating
	# N.B. requires cat and wc to be installed - a shell should not be nessecary though
	#lines = parse(readall(pipeline(`cat $[dataPath]ratings.csv`, `wc -l`)))
	#@time ratings = readdlm(dataPath * "ratings.csv", ',', header=true, dims=(lines, 4))[1]

	#lines = parse(readall(pipeline(`cat $[dataPath]train.csv`, `wc -l`)))
	#@time train = readdlm(dataPath * "train.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]cross.csv`, `wc -l`)))
	@time cross = readdlm(dataPath * "cross.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]test.csv`, `wc -l`)))
	@time test = readdlm(dataPath * "test.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]movies.csv`, `wc -l`)))
	@time movies = readdlm(dataPath * "movies.csv", ',', header=true, dims=(lines, 3))
	moviesHeader = movies[2]
	movies = movies[1]

	lines = parse(readall(pipeline(`cat X-results.csv`, `wc -l`)))
	@time X = readdlm("X-results.csv", ',', header=false, dims=(lines, 18))

	lines = parse(readall(pipeline(`cat Theta-results.csv`, `wc -l`)))
	@time Θ = readdlm("Theta-results.csv", ',', header=false, dims=(lines, 18))

	testUsersData = readcsv("test-users-data.csv")
	users = testUsersData[:, 2]
	usersNumMovies = testUsersData[:, 3]
	usersAvgRatings = testUsersData[:, 4]
	usersRatingsDifference = testUsersData[:, 5]

	#=@printf "users\n"
	usersIndices = Vector{Vector{Int}}()
	f = open("usersIndices.csv")
	i = 0
	@time while (l = readline(f)) != nothing && strip(l) != ""
		i += 1
		if i % 10_000 == 0
			@printf "%i/%i\t%s\n" i 240_000 now()
		end
		l = replace(l, "[", "")
		l = replace(l, "]", "")
		temp = map(parse, split(l, ","))
		if typeof(temp) != Array{Int, 1}
			temp = []
		end
		push!(usersIndices, temp)
	end
	close(f)=#

	#data = [train; cross]
	#data[:, 3] *= 2
	#data = Array{Int}(data[:, 1:3]')
	test[:, 3] *= 2
	test = Array{Int}(test[:, 1:3]')

	#nonEmptyUserIndices = filter(x -> length(usersIndices[x]) > 0, collect(1:size(usersIndices)[1]))
	#realUsers = map(x -> data[1, nonEmptyUserIndices[x][1]], collect(1:size(nonEmptyUserIndices)[1]))
	#reallyNonEmptyUserIndices = usersIndices[nonEmptyUserIndices]

	function userCostModel(user, indices)
		cost = 0
		for i in indices
			cost += (dot(X[test[2, i], :][:], Θ[test[1, i], :][:]) - test[3, i])^2
		end
		return cost / length(indices)
	end

	#=function userCostMemory(user, indices)
		cost = 0
		for i in indices
			movie = test[2, i]
			rating = test[3, i]

			others = find(data[2, :] .== movie)
			others = data[1, others]
			others = intersect(others, realUsers)
			cost += (predict(movie, user, others, data, reallyNonEmptyUserIndices) - rating)^2
		end
		return cost / length(indices)
	end=#

	#memoryCost = zeros(length(realUsers))
	modelCost = zeros(length(users))
	idx = 0
	for user in users
		if idx % 1000 == 0
			@show idx
		end
		idx += 1
		indices = find(test[1, :] .== user)
		try
			modelCost[idx] = userCostModel(user, indices)
			#memoryCost[idx] = userCostMemory(user, indices)
		catch e
			@show e
			continue
		end
	end

	#writecsv("memoryCost.csv", memoryCost)
	writecsv("modelCost.csv", modelCost)
	#writecsv("final-out.csv", hcat(users, usersNumMovies, usersAvgRatings, usersRatingsDifference, memoryCost, modelCost))

	#=@printf "movies\n"
	moviesIndices = Vector{Vector{Int}}()
	f = open("moviesIndices.csv")
	i = 0
	@time while (l = readline(f)) != nothing && strip(l) != ""
		i += 1
		if i % 1_000 == 0
			@printf "%i/%i\t%s\n" i 33_000 now()
		end
		l = replace(l, "[", "")
		l = replace(l, "]", "")
		temp = map(parse, split(l, ","))
		if typeof(temp) != Array{Int, 1}
			temp = []
		end
		push!(moviesIndices, temp)
	end
	close(f)

	@printf "Finished loading, begin processing\n"

	data = [train; cross]
	data[:, 3] *= 2
	data = Array{Int}(data[:, 1:3]')
	test[:, 3] *= 2
	test = Array{Int}(test[:, 1:3]')

	users = union(test[1, :])
	@time usersAvgRatings = pmap(x -> mean(data[2, usersIndices[x]]), users)
	@show size(usersAvgRatings)
	@show size(usersIndices)
	@show size(usersIndices[1])
	@show usersIndices[1]
	@show data[2, usersIndices[1]]
	@show mean(data[2, usersIndices[1]])
	@show map(x -> mean(data[2, usersIndices[x]]), collect(1:5))
	@show size(data[2, usersIndices[1]])
	# Count users, movies and features
	numUsers = length(union(ratings[:,1]))
	numMovies = size(movies)[1]

	numFeatures = 18
	λ = 0.05

	function costTemp(x)
		cost(x, train, numUsers, numMovies, numFeatures, λ)
	end

	#@time moviesIndices = pmap(x -> find(data[2, :] .== x), movies[:, 1])
	@time moviesRatingsCount = pmap(x -> length(moviesIndices[x]), movies[:, 1])
	@time moviesAvgRating = pmap(x -> mean(data[3, moviesIndices[x]]), movies[:, 1])

	users = union(test[1, :])
	@time usersNumMovies = pmap(x -> length(usersIndices[x]), users)
	@time usersAvgRatings = pmap(x -> mean(moviesRatingsCount[data[2, usersIndices[x]]]), users)
	@time usersRatingsDifference = pmap(x -> mean(map(y -> moviesAvgRating[data[2, y]] - data[3, y], usersIndices[x])), users)

	@printf "Done calculating, saving\n"

	writecsv("test-users-data.csv", hcat(users, usersNumMovies, usersAvgRatings, usersRatingsDifference))

	@printf "Done :D\n"=#
end

main()


# @author b33rNc0d3
# @date 2020-06-10
# Using Turing to implement JAGS model from Coursera Bayesian Statistics: Techniques and Models

using Turing
using Distributions
using StatsPlots
using Random

Random.seed!(50)

@model TuringJagsModel(y) = begin
    μ ~ TDist(1.0)
    σ = 1.0
    N = length(y)

    for i = 1 : N
        y[i] ~ Normal(μ, σ)
    end
end;

y = [1.2, 1.4, -0.5, 0.3, 0.9, 2.3, 1.0, 0.1, 1.3, 1.9]

iterations = 1000
ϵ = 0.05
τ = 10

chain = sample( TuringJagsModel(y), HMC(ϵ, τ), iterations)

StatsPlots.plot(chain)

print(chain)

plot(chain)


"""
  τ2λ(τ::Vector{Float64}, λ::Vector{Float64})

Suplement the vector of λ patrameters of Marshall-Olkin copula, given some of those
parameters and a vector fo Kendall's τ correlations. Wroks fo 2,3 variate MO copula
"""
function τ2λ(τ::Vector{Float64}, λ::Vector{Float64})
  if length(τ) == 1
    return [λ[1],λ[2],(λ[1]+λ[2])*τ[1]/(1-τ[1])]
  else
    t = τ./(1 .-τ)
    M = Matrix(1.0I, 3, 3) .-(ones(3,3) .-Matrix(1.0I, 3, 3)) .*t
    fm = hcat([1. 1. 0.; 0. 1. 1.; 1. 0. 1.].*t, .-[1.; 1.; 1.])
    ret = map(x -> (x > 0) ? (x) : (0.000000001), inv(M)*fm*λ)
    [λ[1:3]..., ret..., λ[end]]
  end
end

function moρ2τ(ρ::Float64)
  f(τ) = 1/2 .*sin.(τ[1]*pi/2)+τ[1]/2 - ρ[1]
  fzero(f, -0.999, 0.999)
end

#=
# some Marshall-Olkin copulas helpers

function getmoλ(λ::Vector{Float64}, ind::Vector{Int})
  n = floor(Int, log(2, length(λ)+1))
  s = collect(combinations(1:n))
  λ[[ind == a for a in s]]
end

function setmoλ!(λ::Vector{Float64}, ind::Vector{Int}, a::Float64)
  n = floor(Int, log(2, length(λ)+1))
  s = collect(combinations(1:n))
  λ[[ind == a for a in s]] = a
end
=#


using Plots
#komvos1 komvos2 komvos3
xyz = 0.1*[0 1 -1;
           0 2 -1;
           0 0  0
]; #orismos komvon san xyz
n = size(xyz)[2]; #arithmos komvon

#komvos1_start komvos1_end
numel = [1 2;
         1 3
]; #orismos ton stixion os grammes me toys akraious komvous

#orismos elastikotitas stixion
E = 220e9.*ones(N,1); A = 4e-4.*ones(N,1);

N = size(numel)[1]; #arithmos stoixion

#orismos ton desmeumenon komvon
BC = zeros(3*n,1)
#for 2d uncomment below line
dimconstr = 1:3:3*n; BC[dimconstr] .= 1;
BC[4:end] .= 1; #0==free
fixedofs = findall(x->x==1,BC[:])
freedofs = findall(x->x==0,BC[:])
#boundary conditions
U = zeros(3*n,1); U[4:end].=0;
U[dimconstr].=0; #uncomment for 2d
F = zeros(3*n,1); F[[1 2]]=[-100 100] ;
#orismos olikou mitroou stivarotitas
global K = zeros(3*n,3*n);
#start main loop to calculate stiffness matrix
for i=1:N
    i1 = numel[i,1]
    i2 = numel[i,2]
    #calclulate directions
    x1 = xyz[1,i1]; x2 = xyz[1,i2]; dx = x2-x1;
    y1 = xyz[2,i1]; y2 = xyz[2,i2]; dy = y2-y1;
    z1 = xyz[3,i1]; z2 = xyz[3,i2]; dz = z2-z1;
    L = sqrt(dx^2+dy^2+dz^2);
    nx = dx/L; ny = dy/L; nz = dz/L;
    #direction matrix
    I = [nx^2  nx*ny nx*nz;
         ny*nx ny^2  ny*nz;
         nz*nx nz*ny nz^2];
    k = E[i]*A[i]/L
    Ke = k*[I -I; -I I]
    #println(size(Ke),"\n")
    edofs = [3*i1-2 3*i1-1 3*i1  3*i2-2 3*i2-1 3*i2]
    K[edofs,edofs] = K[edofs,edofs][1,:,1,:] + Ke;
    #println(K,"\n")
end
Kfinal = K[freedofs,freedofs]
rhs = F
rhs = rhs[freedofs]
rhs = rhs - K[freedofs,fixedofs]*U[fixedofs]

Ufree = Kfinal\rhs
println(Ufree)


# Bhattacharyya distances. Much like for KLDivergence we assume the vectors to
# be compared are probability distributions, frequencies or counts rather than
# vectors of samples. Pre-calc accordingly if you have samples.

type BhattacharyyaDist <: SemiMetric end

type HellingerDist <: Metric end

# Bhattacharyya coefficient

function bhattacharyya_coeff{T<:Number}(a::AbstractVector{T}, b::AbstractVector{T})
    n = length(a)
    sqab = zero(T)
    # We must normalize since we cannot assume that the vectors are normalized to probability vectors.
    asum = zero(T)
    bsum = zero(T)

    @simd for i = 1:n
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        sqab += sqrt(ai * bi)
        asum += ai
        bsum += bi
    end

    sqab / sqrt(asum * bsum)
end

bhattacharyya_coeff{T <: Number}(a::T, b::T) = throw("Bhattacharyya coefficient cannot be calculated for scalars")

# Faster pair- and column-wise versions TBD...

# Bhattacharyya distance
evaluate{T<:Number}(dist::BhattacharyyaDist, a::AbstractVector{T}, b::AbstractVector{T}) = -log(bhattacharyya_coeff(a, b))
bhattacharyya(a::AbstractVector, b::AbstractVector) = evaluate(BhattacharyyaDist(), a, b)
evaluate{T <: Number}(dist::BhattacharyyaDist, a::T, b::T) = throw("Bhattacharyya distance cannot be calculated for scalars")
bhattacharyya{T <: Number}(a::T, b::T) = evaluate(BhattacharyyaDist(), a, b)

# Hellinger distance
evaluate{T<:Number}(dist::HellingerDist, a::AbstractVector{T}, b::AbstractVector{T}) = sqrt(1 - bhattacharyya_coeff(a, b))
hellinger(a::AbstractVector, b::AbstractVector) = evaluate(HellingerDist(), a, b)
evaluate{T <: Number}(dist::HellingerDist, a::T, b::T) = throw("Hellinger distance cannot be calculated for scalars")
hellinger{T <: Number}(a::T, b::T) = evaluate(HellingerDist(), a, b)


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ 11c4ccdd-3d86-48ce-b471-3a581902c345
begin
	using Pkg
	Pkg.activate(mktempdir())
	Pkg.Registry.update()
	Pkg.add("Yao")
	Pkg.add("YaoPlots")
end

# ╔═╡ f6fa5d36-1dd4-11eb-1c33-3344c7697e82
using Yao, YaoPlots

# ╔═╡ 3c87d83e-1dca-11eb-2068-6fa8ceb4849a
md"## More Quantum Gates"

# ╔═╡ 5545465e-1dca-11eb-0273-6d32f757b4a2
md"As you might have thought, there do exist gates, other than the previously defined ones."

# ╔═╡ 8def8f9e-1dca-11eb-32cf-fdb2038a8feb
md"### Single qubit gates"

# ╔═╡ 83f65eb6-1dca-11eb-32cb-c577ad1d0b68
md"#### The `` R_𝜑^Z `` gate"

# ╔═╡ 9d8d835c-1dca-11eb-0ab7-01a80ca23e30
md"Passing a qubit through the `` R_𝜑^Z `` is equivalent to multiplying its state vector by `` \begin{bmatrix}1 & 0\\0 & e^{𝜑i}\end{bmatrix} ``.
Remember, `` e^{iθ} = \cos(𝜃) + i\sin(𝜃) ``.

The `` R_𝜑^Z `` gate can be alternatively denoted by, `` \begin{bmatrix}e^{-𝜑i/2} & 0\\0 & e^{𝜑i/2}\end{bmatrix} ``. Its just the original matrix, multiplied by `` e^{-𝜑i/2} ``. We can do this since multiplication by `` e^{-𝜑i/2} `` is not *observable* during measurement as its a complex unit and `` | e^{iθ} | = |\cos(𝜃) + i\sin(𝜃)| = 1 . `` Remember that the abstract value of a complex number `` a + ib ``, i.e.,  `` |a + ib| = \sqrt{a^2 + b^2} `` and `` \sin^2θ + \cos^2θ = 1. ``

Considering a qubit, `` a|0〉 + b|0〉 , `` passing it through the `` R_\frac{𝞹}{2}^Z `` gate is equivalent to `` \begin{bmatrix}1 & 0\\0 & e^{𝜋i/2}\end{bmatrix} `` . And since `` \cos(\frac{𝞹}{2}) = 0`` and `` \sin(\frac{𝞹}{2}) = 1 , `` we can rewrite the above as, ``\begin{bmatrix}1 & 0\\0 & i\end{bmatrix}`` `` \begin{bmatrix}a\\b\end{bmatrix} ``."

# ╔═╡ 8cca417e-1dd4-11eb-2a18-bf9b80da3fd2
md"Lets try the above in Yao! The `` R_𝜑^Z `` gate can be used in Yao with the shift *block*. "

# ╔═╡ 1f7c993e-1dd5-11eb-2f13-7711d2426eb1
begin
	qubit = rand_state(1)
	state(qubit)
end

# ╔═╡ 5960e486-1dd5-11eb-2253-838f59b06d4c
state(qubit |> chain(1, put(1=>shift(π/2))))

# ╔═╡ ac10f6f8-1dd5-11eb-2cf1-29480480572a
md"As expected the output was `` \begin{bmatrix}a\\ib\end{bmatrix} ``. Remember, `` i = \sqrt{-1} `` and `` i^2 = -1 ``. (Also, note that in Julia, imaginary number 𝑖 is represented by im.)"

# ╔═╡ af22f8a8-1ddc-11eb-2958-bb45cdadf44f
md"Also, `` R_𝜋^Z `` gate is equivalent to Z gate."

# ╔═╡ df5a774e-1ddc-11eb-2e86-1dc8dd502e26
round.(Matrix(chain(1, put(1=>shift(π))))) == round.(Matrix(chain(1, put(1=>Z)))) #The round functions "rounds-off" the elements of the matrices

# ╔═╡ c1f87994-1e7c-11eb-3da3-21aa52e80520
md"Its represented in a circuit diagram by,"

# ╔═╡ cda277ce-1e7c-11eb-246c-0b97bb66ab73
plot(chain(1, put(1=>shift(π/3))))

# ╔═╡ a1173196-1dd9-11eb-394d-271a3871e181
md"#### The T Gate"

# ╔═╡ b1cc417a-1dd9-11eb-0156-67fd10743000
md"The T gate is equivalent to `` R_\frac{𝞹}{4}^Z . `` In its matrix form, it can be written as `` \begin{bmatrix}1 & 0\\0 & \frac{1 + i}{\sqrt{2}}\end{bmatrix} ``. Nevertheless, in Yao, it can be used by using the **T** *block* ."

# ╔═╡ 67b2bc30-1dda-11eb-2d81-edd411a207b2
state(qubit |> chain(1, put(1=>T)))

# ╔═╡ 7d546e10-1dda-11eb-2b9a-abe61e9edf6c
md"Also,"

# ╔═╡ 83dfc128-1dda-11eb-37d9-37b959ea5f58
Matrix(chain(1, put(1=>shift(π/4)))) == Matrix(chain(1, put(1=>T)))

# ╔═╡ ea6c8192-1e7c-11eb-30e5-8d65801ce703
md"Its circuit diagram representation looks somewhat like - "

# ╔═╡ feab87f0-1e7c-11eb-075b-a55e9d86184c
plot(chain(1, put(1=>T)))

# ╔═╡ 408ce964-1de4-11eb-2830-f14bdc87f899
md"#### The `` R_𝜑^X `` gate"

# ╔═╡ 09a069f0-1df6-11eb-19de-5da4bb572510
md"Similar to the `` R_𝜑^Z `` gate, the `` R_𝜑^X `` gate can be represented by `` \begin{bmatrix}\cos(\frac{𝜑}{2}) & -\sin(\frac{𝜑}{2})i\\-\sin(\frac{𝜑}{2})i & \cos(\frac{𝜑}{2})\end{bmatrix} ``."

# ╔═╡ 5a6f6786-1df7-11eb-0f06-5b229b0935bb
md"#### The `` R_𝜑^Y `` gate"

# ╔═╡ 6a775062-1df7-11eb-1ea2-7d5517774328
md"Similar to the `` R_𝜑^Z `` gate, the `` R_𝜑^X `` gate can be represented by `` \begin{bmatrix}\cos(\frac{𝜑}{2}) & -\sin(\frac{𝜑}{2})\\\sin(\frac{𝜑}{2}) & \cos(\frac{𝜑}{2})\end{bmatrix} ``."

# ╔═╡ 8d65504c-1df7-11eb-0f3a-93796bf5b7f9
md"They can be represented in Yao using the **Rx** and **Ry** *blocks* respectively"

# ╔═╡ e10cd2c0-1e00-11eb-3c86-2d958e7aa92f
state(qubit |> chain(1, put(1=>Rx(π))))

# ╔═╡ 38071b04-1e00-11eb-1c6c-03f03ce37a92
state(qubit |> chain(1, put(1=>Ry(π))))

# ╔═╡ 451b4c6a-1e1f-11eb-18f0-e301d5d34c93
md"There's also an **Rz** *block* which which represents the alternative form of `` R_𝜑^Z `` matrix, i.e., `` \begin{bmatrix}e^{-𝜑i/2} & 0\\0 & e^{𝜑i/2}\end{bmatrix} ``."

# ╔═╡ ea7cda52-1e1f-11eb-2cdc-e9a04181acfb
state(qubit |> chain(1, put(1=>Rz(π))))

# ╔═╡ 8c6c62c4-1e20-11eb-1017-f51535591214
md"Note that the absolute value of both the shift and Rz blocks are same."

# ╔═╡ a214d034-1e20-11eb-2975-6f440fee2ee5
abs.(Matrix(chain(1, put(1=>shift(π/5))))) == abs.(Matrix(chain(1, put(1=>Rz(π/5)))))

# ╔═╡ 16d650bc-1e7d-11eb-39d0-33ba9cb031ed
md"The circuit diagram representations of Rx, Ry and Rz blocks, respectively"

# ╔═╡ 275a85be-1e7d-11eb-1114-f9872f35915d
plot(chain(1, put(1=>Rx(π/5)), put(1=>Ry(π/3)), put(1=>Rz(π/7))))

# ╔═╡ 3355c7c4-1e21-11eb-1d88-6f5a6b5c51a8
md"### Multi-qubit Gates"

# ╔═╡ 4afe09cc-1e21-11eb-1eec-6b6d78e76aae
md"#### The SWAP Gate"

# ╔═╡ 812b528e-1e21-11eb-3219-a9ce5de87933
md"The SWAP gate swaps the state of two qubits. It can be represented by the matrix,
$\begin{bmatrix}1 & 0 & 0 & 0\\0 & 0 & 1 & 0\\0 & 1 & 0 & 0\\0 & 0 & 0 & 1\end{bmatrix}$. Its represented in Yao via the **swap** *block*."

# ╔═╡ b85aaad6-1e79-11eb-2d2a-2f0a8fdb9336
begin
	q = rand_state(2)
	state(q)
end

# ╔═╡ 63dd442a-1e7c-11eb-3c73-7186108b87b4
state(q |> chain(2, swap(1,2)))

# ╔═╡ 5c1a804c-1e7a-11eb-22ae-df5da105a8d7
md"The SWAP gate has the following circuit diagram representation"

# ╔═╡ cd7061a8-1e7d-11eb-3239-8f8410988851
plot(chain(2, swap(1,2)))

# ╔═╡ 5890257c-1e7e-11eb-2880-1d4257468ffc
md"There's a Toffoli gate, an S gate, a CSWAP gate, a `` \; C R_𝜑^{X,Y,Z} `` gate and probably a lot more. They can all be constructed using the existing blocks in Yao."

# ╔═╡ c1d38e3e-1e7e-11eb-07ac-5345d88b4781
let
	toffoli_Gate = chain(3, control(1:2, 3=>X)) #The toffoli gate
	plot(toffoli_Gate)
end

# ╔═╡ 512086b4-1e7f-11eb-1f4b-cde0db3c9734
let
	S_Gate = chain(1, put(1 => label(shift(π/2), "S"))) #The S gate
	plot(S_Gate)
end

# ╔═╡ 8c44d66e-1e7f-11eb-2821-5532c1fdb89d
let
	CRXYZ_Gates = chain(2, control(1, 2=>Rx(π/5)), control(1, 2=>Ry(π/7)), control(1, 2=>Rz(π/11)))
	plot(CRXYZ_Gates)
end

# ╔═╡ 07b469b8-1e85-11eb-1b2c-f3fb85382b66
let
	CSWAP = chain(3, control(1, 2:3=>SWAP)) #The CSWAP gate
	plot(CSWAP)
end

# ╔═╡ 96d92a88-4459-11eb-29ee-336cbba060a7
md"### The Measure Gate"

# ╔═╡ b23f15ca-4459-11eb-3089-f7ea50aaf5de
md"We already know how to measure the qubits. We can do it in the circuit itself too."

# ╔═╡ c9233378-4459-11eb-0aa9-af1cd8cef6ca
begin
	MeasureGate = chain(2, repeat(H, 1:2), Measure(2, locs=1:2))
	plot(MeasureGate)
end

# ╔═╡ 1615f882-445a-11eb-3a16-c76bb2305ff5
md"Note that now, when we measure them using the measure block, the output remains unchanged, even though we should've a 25% chance of getting `` |00〉, |01〉, |10〉 `` or `` |11〉 ``."

# ╔═╡ 9626495a-445a-11eb-0a2d-2168c66ad1e0
zero_state(2) |> MeasureGate |> r->measure(r, nshots=1024)

# ╔═╡ cd54fd7c-445a-11eb-0015-bd814cd6293f
md"Without the Measure gate,"

# ╔═╡ dbe0b9ee-445a-11eb-3526-29900707092e
zero_state(2) |> repeat(2, H, 1:2) |> r->measure(r, nshots=1024)

# ╔═╡ 06d1ed92-4485-11eb-3485-d1be36b9f68b
md"### The LabeledBlock"

# ╔═╡ 15f64d2c-4485-11eb-3749-a33d80bf5b6b
md"Its used for easily plotting circuits as boxes for simpler visualization."

# ╔═╡ 2c8d9af4-4485-11eb-238e-b36d910466fc
let
	a = chain(3, repeat(H, 1:2), put(3=>X))
	b = chain(2, repeat(Y, 1:2))
	circuit = chain(3, put(1:3 => label(a, "circuit a")), put(2:3 => label(b, "circuit b")))
	plot(circuit)
end

# ╔═╡ f5e2eb8a-4484-11eb-0957-41f14bbfc3d9
md"### Daggered Block"

# ╔═╡ 00b09d5a-4485-11eb-2857-011786a85254
md"We use Daggered block to build circuits which undo the effects of a particular circuit.

Let's take an example. Remember Bell Circuit?"

# ╔═╡ ec2e6fb4-4485-11eb-351d-4781925a47d3
begin
	bellcircuit = chain(2, put(1=>H), control(1, 2=>X))
	plot(bellcircuit)
end

# ╔═╡ 1791bfd0-4486-11eb-0107-e345d88b2ddb
state(zero_state(2) |> bellcircuit)

# ╔═╡ 61a0a0a0-4486-11eb-2f6a-750de59d5988
md"Building the Reverse Bell Circuit is as easy as,"

# ╔═╡ 771faeba-4486-11eb-07ac-4dc4d52d4726
state(zero_state(2) |> bellcircuit |> Daggered(bellcircuit))

# ╔═╡ b7d76bd4-4486-11eb-3565-ab0ef9238acf
md"Plotting the Daggered Block is a bit tricky! YaoPlots doesn't support DaggeredBlock yet. There are two alternatives to this."

# ╔═╡ 39890364-4488-11eb-22d7-33a8db179379
plot(Daggered(bellcircuit))

# ╔═╡ 4eb6a852-4488-11eb-0e42-7be38fca7571
md"One way is to use the label block"

# ╔═╡ 11a084ac-4487-11eb-30ea-539b38e0231f
let
	reversebellcircuit = put(2, 1:2 => label(Daggered(bellcircuit), "Bell\n Circuit†"))
	plot(reversebellcircuit)
end

# ╔═╡ 49ed44ca-4488-11eb-07dd-27f8ca6a1d41
md"Another way, is to use ` ' `."

# ╔═╡ 6325cf34-4488-11eb-0415-1d1f2fa0f1ba
plot(bellcircuit')

# ╔═╡ 6f5420f8-4488-11eb-3cf9-03675dc3c8db
md"Whats the difference between the DaggeredBlock and adjoint ` ' ` ?

One is a function of Yao, while another of Julia `Base`. Both perform the same operation on their input.

**Note:** _Its recommended to always use `'` over DaggeredBlock_."

# ╔═╡ f612e17a-475c-11eb-080a-ab0967411f64
md"### The Kron Block"

# ╔═╡ fc6c126c-475c-11eb-3efa-d5ffc4347358
md"Consider you've to make the below circuit."

# ╔═╡ 0c13552c-475d-11eb-23a8-ad67f6bb93bc
plot(kron(X,Y,Z,H,Rx(15)))

# ╔═╡ 2c87c612-475d-11eb-3af7-9542b157252f
md"Tired of using `put` block after `put` block?

Presenting the kron block, where you can just input the gates on every qubit, one by one."

# ╔═╡ 616d3b96-475d-11eb-1500-5bbf4a263df8
plot(chain(5, put(1:3 => kron(X,Y,Z)), put(3:5 => kron(H,Rx(15),T))))

# ╔═╡ d643aa86-475d-11eb-3af7-0195206e2f64
md"### The Rotation Gate"

# ╔═╡ e39911c6-475d-11eb-2ef5-f743c27d0d07
md"Its the general version of the Rx, Ry and Rz gates you saw above"

# ╔═╡ fb813b74-475d-11eb-1323-6faf1b221dad
plot(chain(1, rot(X, 15), rot(Y, 16)))

# ╔═╡ Cell order:
# ╠═11c4ccdd-3d86-48ce-b471-3a581902c345
# ╟─3c87d83e-1dca-11eb-2068-6fa8ceb4849a
# ╟─5545465e-1dca-11eb-0273-6d32f757b4a2
# ╟─8def8f9e-1dca-11eb-32cf-fdb2038a8feb
# ╟─83f65eb6-1dca-11eb-32cb-c577ad1d0b68
# ╟─9d8d835c-1dca-11eb-0ab7-01a80ca23e30
# ╟─8cca417e-1dd4-11eb-2a18-bf9b80da3fd2
# ╠═f6fa5d36-1dd4-11eb-1c33-3344c7697e82
# ╠═1f7c993e-1dd5-11eb-2f13-7711d2426eb1
# ╠═5960e486-1dd5-11eb-2253-838f59b06d4c
# ╟─ac10f6f8-1dd5-11eb-2cf1-29480480572a
# ╟─af22f8a8-1ddc-11eb-2958-bb45cdadf44f
# ╠═df5a774e-1ddc-11eb-2e86-1dc8dd502e26
# ╟─c1f87994-1e7c-11eb-3da3-21aa52e80520
# ╠═cda277ce-1e7c-11eb-246c-0b97bb66ab73
# ╟─a1173196-1dd9-11eb-394d-271a3871e181
# ╟─b1cc417a-1dd9-11eb-0156-67fd10743000
# ╠═67b2bc30-1dda-11eb-2d81-edd411a207b2
# ╟─7d546e10-1dda-11eb-2b9a-abe61e9edf6c
# ╠═83dfc128-1dda-11eb-37d9-37b959ea5f58
# ╟─ea6c8192-1e7c-11eb-30e5-8d65801ce703
# ╟─feab87f0-1e7c-11eb-075b-a55e9d86184c
# ╟─408ce964-1de4-11eb-2830-f14bdc87f899
# ╟─09a069f0-1df6-11eb-19de-5da4bb572510
# ╟─5a6f6786-1df7-11eb-0f06-5b229b0935bb
# ╟─6a775062-1df7-11eb-1ea2-7d5517774328
# ╟─8d65504c-1df7-11eb-0f3a-93796bf5b7f9
# ╠═e10cd2c0-1e00-11eb-3c86-2d958e7aa92f
# ╠═38071b04-1e00-11eb-1c6c-03f03ce37a92
# ╟─451b4c6a-1e1f-11eb-18f0-e301d5d34c93
# ╠═ea7cda52-1e1f-11eb-2cdc-e9a04181acfb
# ╟─8c6c62c4-1e20-11eb-1017-f51535591214
# ╠═a214d034-1e20-11eb-2975-6f440fee2ee5
# ╟─16d650bc-1e7d-11eb-39d0-33ba9cb031ed
# ╟─275a85be-1e7d-11eb-1114-f9872f35915d
# ╟─3355c7c4-1e21-11eb-1d88-6f5a6b5c51a8
# ╟─4afe09cc-1e21-11eb-1eec-6b6d78e76aae
# ╟─812b528e-1e21-11eb-3219-a9ce5de87933
# ╠═b85aaad6-1e79-11eb-2d2a-2f0a8fdb9336
# ╠═63dd442a-1e7c-11eb-3c73-7186108b87b4
# ╟─5c1a804c-1e7a-11eb-22ae-df5da105a8d7
# ╟─cd7061a8-1e7d-11eb-3239-8f8410988851
# ╟─5890257c-1e7e-11eb-2880-1d4257468ffc
# ╠═c1d38e3e-1e7e-11eb-07ac-5345d88b4781
# ╠═512086b4-1e7f-11eb-1f4b-cde0db3c9734
# ╠═8c44d66e-1e7f-11eb-2821-5532c1fdb89d
# ╠═07b469b8-1e85-11eb-1b2c-f3fb85382b66
# ╟─96d92a88-4459-11eb-29ee-336cbba060a7
# ╟─b23f15ca-4459-11eb-3089-f7ea50aaf5de
# ╠═c9233378-4459-11eb-0aa9-af1cd8cef6ca
# ╟─1615f882-445a-11eb-3a16-c76bb2305ff5
# ╠═9626495a-445a-11eb-0a2d-2168c66ad1e0
# ╟─cd54fd7c-445a-11eb-0015-bd814cd6293f
# ╠═dbe0b9ee-445a-11eb-3526-29900707092e
# ╟─06d1ed92-4485-11eb-3485-d1be36b9f68b
# ╟─15f64d2c-4485-11eb-3749-a33d80bf5b6b
# ╠═2c8d9af4-4485-11eb-238e-b36d910466fc
# ╟─f5e2eb8a-4484-11eb-0957-41f14bbfc3d9
# ╟─00b09d5a-4485-11eb-2857-011786a85254
# ╠═ec2e6fb4-4485-11eb-351d-4781925a47d3
# ╠═1791bfd0-4486-11eb-0107-e345d88b2ddb
# ╟─61a0a0a0-4486-11eb-2f6a-750de59d5988
# ╠═771faeba-4486-11eb-07ac-4dc4d52d4726
# ╟─b7d76bd4-4486-11eb-3565-ab0ef9238acf
# ╠═39890364-4488-11eb-22d7-33a8db179379
# ╟─4eb6a852-4488-11eb-0e42-7be38fca7571
# ╠═11a084ac-4487-11eb-30ea-539b38e0231f
# ╟─49ed44ca-4488-11eb-07dd-27f8ca6a1d41
# ╠═6325cf34-4488-11eb-0415-1d1f2fa0f1ba
# ╟─6f5420f8-4488-11eb-3cf9-03675dc3c8db
# ╟─f612e17a-475c-11eb-080a-ab0967411f64
# ╟─fc6c126c-475c-11eb-3efa-d5ffc4347358
# ╠═0c13552c-475d-11eb-23a8-ad67f6bb93bc
# ╟─2c87c612-475d-11eb-3af7-9542b157252f
# ╠═616d3b96-475d-11eb-1500-5bbf4a263df8
# ╟─d643aa86-475d-11eb-3af7-0195206e2f64
# ╟─e39911c6-475d-11eb-2ef5-f743c27d0d07
# ╠═fb813b74-475d-11eb-1323-6faf1b221dad


struct ScratchThermalGlobal{T}
    ndim::T # model dimension
    nvert::T # number of vertices per element
    nnodel::T # number of nodes per element
    nel::T # number of elements
    nip::T # number of vertices per element
end

struct ShapeFunctionsThermal{T}
    N::Vector{SMatrix{1,6,T,6}}
    ∇N::Vector{SMatrix{2,6,T,12}}
    dN3ds::SMatrix{2,3,T,6}
    w_ip::SVector{7,T}
    N3::Vector{SMatrix{1,3,T,3}}
end

struct ShapeFunctionsStokes{T}
    N::Vector{SMatrix{1,6,T,6}}
    ∇N::Vector{SMatrix{2,6,T,12}}
    NP::Vector{SMatrix{1,3,Float64,3}}
    dN3ds::SMatrix{2,3,T,6}
    w_ip::SVector{7,T}
    N3::Vector{SMatrix{1,3,T,3}}
end

struct ShapeFunctionsStress{A,B,C,D,E,F}
    N::A
    ∇N::B
    NP::C
    dN3ds::D
    w_ip::E
    N3::F
end

@inline function _get_SF(nip, nnodel)
    x_ip, w_ip = ip_triangle(nip)
        # local coordinates and weights of points for integration of
        # velocity/pressure matrices
    N,dNds = shape_functions_triangles(x_ip,nnodel)
        # velocity shape functions and their derivatives
    dN3ds = @SMatrix [-1.0   1.0   0.0 # w.r.t. r
                      -1.0   0.0   1.0]
        # derivatives of linear (3-node) shape functions; used to calculate
        # each element's Jacobian

    return N,dNds,dN3ds,w_ip
end

function ip_triangle(nip)
    if nip === 3
        ipx,ipw = ip_triangle3()
    elseif nip === 6
        ipx,ipw = ip_triangle6()
    elseif nip === 7
        ipx,ipw = ip_triangle7()
    end
    return ipx, ipw
end

ip_triangle(::Val{3}) = ip_triangle3()
ip_triangle(::Val{6}) = ip_triangle6()
ip_triangle(::Val{7}) = ip_triangle7()

function ip_triangle3()
    # -- Integration point coordinates
    ipx = @SMatrix [
        1/6 1/6
        2/3 1/6
        1/6 2/3
    ]

    # -- Weights
    ipw = @SVector [
        1/6,
        1/6,
        1/6,
    ]

    return ipx, ipw

end

function ip_triangle6()

     # -- Integration point coordinates
    g1  = (8.0-sqrt(10.0) + sqrt(38.0-44.0*sqrt(2.0/5.0)))/18.0
    g2  = (8.0-sqrt(10.0) - sqrt(38.0-44.0*sqrt(2.0/5.0)))/18.0
    ipx = @SMatrix [
        1.0-2.0*g1  g1
        g1          1.0-2.0*g1
        g1          g1
        1.0-2.0*g2  g2
        g2          1.0-2.0*g2
        g2          g2
    ]

    # -- Weights
    w1 = (620.0 + sqrt(213125.0-53320.0*sqrt(10.0)))/3720.0;
    w2 = (620.0 - sqrt(213125.0-53320.0*sqrt(10.0)))/3720.0;
    ipw =  @SVector [
        w1,                       #0.223381589678011;
        w1,                       #0.223381589678011;
        w1,                       #0.223381589678011;
        w2,                       #0.109951743655322;
        w2,                       #0.109951743655322;
        w2,                       #0.109951743655322;
    ]

    return ipx, 0.5*ipw

end

function ip_triangle7()

    # -- Integration point coordinates
    g1       = (6.0 - sqrt(15.0))/21.0;
    g2       = (6.0 + sqrt(15.0))/21.0;
    ipx = @SMatrix [
        1.0/3.0     1.0/3.0
        1.0-2.0*g1  g1
        g1          1.0-2.0*g1
        g1          g1
        1.0-2.0*g2  g2
        g2          1.0-2.0*g2
        g2          g2
    ]

    # -- Weights
    w1 = (155.0 - sqrt(15.0))/1200.0;
    w2 = (155.0 + sqrt(15.0))/1200.0;
    ipw = @SVector [
        0.225,
        w1,                       #0.223381589678011;
        w1,                       #0.223381589678011;
        w1,                       #0.223381589678011;
        w2,                       #0.109951743655322;
        w2,                       #0.109951743655322;
        w2,                       #0.109951743655322;
    ]

    return ipx, 0.5*ipw

end

function shape_functions_triangles(lc,nnodel)
    r    = lc[:,1]
    s    = lc[:,2]
    npt  = length(r)
    N, dN = get_N_∇N(r,s,npt,nnodel)
    return N, dN
end

function get_N_∇N(r,s,npt,::Val{3})
    N = [sf_N_tri3(r[ip],s[ip]) for ip in 1:npt]
    dN = [sf_dN_tri3(r[ip],s[ip]) for ip in 1:npt]
    return N, dN
end

function get_N_∇N(r,s,npt,::Val{6})
    N = [sf_N_tri6(r[ip],s[ip]) for ip in 1:npt]
    dN = [sf_dN_tri6(r[ip],s[ip]) for ip in 1:npt]
    return N, dN
end

function get_N_∇N(r,s,npt,::Val{7})
    N = [sf_N_tri7(r[ip],s[ip]) for ip in 1:npt]
    dN = [sf_dN_tri7(r[ip],s[ip]) for ip in 1:npt]
    return N, dN
end

function sf_N_tri3(r,s)
    # Find shape functions and their derivatives at given points on the
    # master element for 3 node triangle
    # 3-node triangle (node numbering is important)
    #
    #        3
    #        | \
    # s-axis |   \
    #        |     \
    #        1 - - - 2
    #          r axis -->
    t = 1.0-r-s;
    N = @SMatrix [t r s]   # N3 at coordinate (r,s)
    return N
end

function sf_dN_tri3(r,s)
    # Find shape functions and their derivatives at given points on the
    # master element for 3 node triangle
    # 3-node triangle (node numbering is important)
    #
    #        3
    #        | \
    # s-axis |   \
    #        |     \
    #        1 - - - 2
    #          r axis -->
    t       = 1.0-r-s;
    dN = @SMatrix [-1.0   1.0   0.0 # w.r.t. r
          -1.0   0.0   1.0]; # w.r.t. s
    return dN
end

function sf_N_tri7(r,s)
    # Find shape functions and their derivatives at given points on the
    # master element for a 7 node triangle
    # 7-node triangle (node numbering is important)
    #
    #        3
    #        | \
    # s-axis 6   5
    #        | 7   \
    #        1 - 4 - 2
    #          r-axis
    t       = 1-r-s
    N = @SMatrix [t*(2*t-1)+3*r*s*t r*(2*r-1)+3*r*s*t s*(2*s-1)+3*r*s*t 4*r*t-12*r*s*t 4*r*s-12*r*s*t 4*s*t-12*r*s*t 27*r*s*t]

    return N
end

function sf_dN_tri7(r,s)
    # Find shape functions and their derivatives at given points on the
    # master element for a 7 node triangle
    # 7-node triangle (node numbering is important)
    #
    #        3
    #        | \
    # s-axis 6   5
    #        | 7   \
    #        1 - 4 - 2
    #          r-axis
    t       = 1-r-s
    dN= @SMatrix [1-4*t+3*s*t-3*r*s  -1+4*r+3*s*t-3*r*s  3*s*t-3*r*s        4*t-4*r+12*r*s-12*s*t 4*s+12*r*s-12*s*t -4*s+12*r*s-12*s*t    -27*r*s+27*s*t
               1-4*t+3*r*t-3*r*s   3*r*t-3*r*s       -1+4*s+3*r*t-3*r*s -4*r-12*r*t+12*r*s     4*r-12*r*t+12*r*s  4*t-4*s-12*r*t+12*r*s 27*r*t-27*r*s]

    return dN
end

function sf_N_tri6(r,s)
    # Find shape functions and their derivatives at given points on the
    # master element for a 6 node triangle (node numbering is important)
    #
    #        3
    #        | \
    # s-axis 6   5
    #        |     \
    #        1 - 4 - 2
    #          r-axis
    #
    t = 1.0 - r - s
    # N1 at coordinate (r,s), N2 at coordinate (r,s), etc
    N = @SMatrix [t*(2.0*t-1.0)  r*(2.0*r-1.0) s*(2.0*s-1.0) 4.0*r*t 4.0*r*s 4.0*s*t]
    #     dN1       dN2    dN3    dN4       dN5    dN6

    return N
end

function sf_dN_tri6(r,s)
    # Find shape functions and their derivatives at given points on the
    # master element for a 6 node triangle (node numbering is important)
    #
    #        3
    #        | \
    # s-axis 6   5
    #        |     \
    #        1 - 4 - 2
    #          r-axis
    #
    t = 1.0 - r - s
    # N1 at coordinate (r,s), N2 at coordinate (r,s), etc
    #     dN1       dN2    dN3    dN4       dN5    dN6
    dN = @SMatrix [-(4.0*t-1.0)  4.0*r-1  0.0         4.0*(t-r)  4.0*s   -4.0*s     # w.r.t. r
               -(4.0*t-1.0)  0.0      4.0*s-1.0  -4.0*r      4.0*r   4.0*(t -s)]; # w.r.t. s

    return dN
end


"""
    SupportedMarksDistributions

Union type for the marks distributions supported by the package for the SEMPP models.
"""
SupportedMarksDistributions = Union{Type{Distributions.GeneralizedPareto}, Type{EGPD.EGPpower}}

"""
SEMPPExpKern(data, μ, ϕ, γ, markdens, ξ, α, β, κ)

Self-exciting marked point process with exponential kernel model type such that for all t, and all m there is an event in t with mark m with probability λ_g(t)f(m|t)) where :
```
λ_g(t) = μ + ϕ * ν(t)
```

with `ν(t) = \\sum_{t_k < t} (1 + δ*m_k) (\\exp(-γ(t-t_k))`

f(m|t) is markdens with scale `σ_t = β + α * ν(t)`

`t_k` are the timestamps of all events and `m_k` their marks.
"""
mutable struct SEMPPExpKern <: SEPP
    data::Union{MarkedTimeSeries, Nothing}
    μ::Real
    ϕ::Real
    γ::Real
    δ::Real
    impact_function::Function
    markdens::SupportedMarksDistributions
    ξ::Real
    α::Real
    β::Real
    κ::Real
    cov_mat::Union{Array{Float64, 2}, Nothing}

    function SEMPPExpKern(data::Union{MarkedTimeSeries, Nothing} = nothing; μ::Real = rand(), ϕ::Real = rand(), γ::Real = rand(), δ::Real = rand(), impact_function::Function = (d -> 1 + tanh(d)), markdens::SupportedMarksDistributions = Distributions.GeneralizedPareto, ξ::Real = rand(), α::Real = rand(), β::Real = rand(), κ::Real = rand(), cov_mat::Union{Array{Float64, 2}, Nothing} = nothing)
        if any((μ, ϕ, γ, α, β, δ, κ) .< 0)
            error("paramaters except for ξ must be positive or zero")
        else
            new(data, μ, ϕ, γ, δ, impact_function, markdens, ξ, α, β, κ, cov_mat)
        end
    end

end

params(sepp::SEMPPExpKern) = Dict(:μ => sepp.μ, :ϕ => sepp.ϕ, :γ => sepp.γ, :impact_function => sepp.impact_function, :δ => sepp.δ, :markdens => sepp.markdens, :ξ => sepp.ξ, :α => sepp.α, :β => sepp.β, :κ => sepp.κ)

"""
    lin_coeff_impact(sepp::SEMPPExpKern)::Real

Get the linear coefficient impact of the marks on the rate.
"""
lin_coeff_impact(sepp::SEMPPExpKern)::Real = sepp.δ

"""
    marks_scale_params(sepp::SEMPPExpKern)::Real

Get the paramaters weighing the volatility function into the scale parameter of the mark density: `\\sigma_t = \\beta + \\alpha \\nu(t)`.
"""
marks_scale_params(sepp::SEMPPExpKern)::Real = (sepp.α, sepp.β)

shape(sepp::SEMPPExpKern) = sepp.κ

decay(sepp::SEMPPExpKern) = sepp.ξ


export projGNexplicit

"""
function projGNexplicit(gc,pMis,pInv,sig,dsig,d2F,d2R,Active)

explicitly builds and solves projected normal equation.

Inputs:

	gc        - gradient
	pMis      - misfit params
	pInv      - inverse param
	sig,dsig  - current model and derivative
	d2F,d2R   - Hessians os misfit and regularizer
	Active    - indicator of active set, default false

Outputs:

	dm        - search direction
	times     - array containing time to build and solve Hessian
"""
function projGNexplicit(gc::Vector,pMis,pInv::InverseParam,sig::Vector,dsig,d2F,d2R,Active=falses(length(gc)))
	if all(Active)
		return 0*gc,[0.0;0.0]
	end
	# get Hessian of misfit
	timeBuild = @elapsed begin
	Hm = getHessian(sig,pMis,d2F)
	end;

	# build overall Hessian
	H  = dsig'*Hm*dsig + d2R

	# remove Active constraints
	timeSolve = @elapsed begin
	Hr = H[.!(Active),.!(Active)]
	gr = gc[.!(Active)]
	dm = 0*gc
	dr = -(Hr\gr)
	if any(isinf.(dr)) || any(isnan.(dr))
		dr = -(pinv(Hr)*gr)
	end
	dm[.!(Active)] = dr
	end

	# solve and return
	return dm,[timeBuild;timeSolve]
end


using GlobalSensitivity, Test, StableRNGs

function ishi_batch(X)
    A= 7
    B= 0.1
    @. sin(X[1,:]) + A*sin(X[2,:])^2+ B*X[3,:]^4 *sin(X[1,:])
end
function ishi(X)
    A= 7
    B= 0.1
    sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1])
end

function linear_batch(X)
    A= 7
    B= 0.1
    @. A*X[1,:]+B*X[2,:]
end
function linear(X)
    A= 7
    B= 0.1
    A*X[1]+B*X[2]
end

lb = -ones(4)*π
ub = ones(4)*π

rng = StableRNG(123)
# res1 = gsa(ishi,GlobalSensitivity.RBDFAST(), num_params = 4, N=5000, rng=rng)
# res2 = gsa(ishi_batch,RBDFAST(),N=1000, batch=true, rng=rng)

# res1efast = gsa(ishi,eFAST(),[[lb[i],ub[i]] for i in 1:4],n=15000)
# res2efast = gsa(ishi_batch,eFAST(),[[lb[i],ub[i]] for i in 1:4],n=15000,batch=true)

# @test res1 ≈ res1efast.S1[1,:] atol = 3e-2
# @test res2.S1 ≈ res2efast.S1[1,:] atol = 3e-2

res1 = gsa(linear,GlobalSensitivity.RBDFAST(),num_params = 4, N=15000)
res2 = gsa(linear_batch,GlobalSensitivity.RBDFAST(),num_params = 4, batch=true, N=15000)

res1efast = gsa(linear,eFAST(),[[lb[i],ub[i]] for i in 1:4])
res2efast = gsa(linear_batch,eFAST(),[[lb[i],ub[i]] for i in 1:4],batch=true)

@test res1 ≈ res1efast.S1[1,:] atol = 3e-2
@test res2 ≈ res2efast.S1[1,:] atol = 3e-2


"""
    DiagonalStandardNormal <: AbstactMvNormal

Multivariate Normal distribution with zero mean vector and identity covariance.
"""
struct DiagonalStandardNormal <: AbstractMvNormal
    D::Int
end

const DiagStdNormal = DiagonalStandardNormal

dim(d::DiagonalStandardNormal) = d.D

rand(rng::AbstractRNG, d::DiagonalStandardNormal) = randn(rng, dim(d))
rand(rng::AbstractRNG, d::DiagonalStandardNormal, N::Int) = randn(rng, dim(d), N)

@unionise function logpdf(d::DiagonalStandardNormal, x::AbstractVector{<:Real})
    @assert dim(d) == length(x)
    return -0.5 * (dim(d) * log(2π) + sum(abs2, x))
end
@unionise function logpdf(d::DiagonalStandardNormal, X::AbstractMatrix{<:Real})
    @assert dim(d) == size(X, 1)
    return -0.5 .* ((dim(d) * log(2π)) .+ reshape(mapreducedim(abs2, +, X, 1), size(X, 2)))
end


@testset "Tensor decomposition" begin
    @testset "k = $k" for k in 1:7
        @testset "p = $p" for p in 1:7
            # Check that inverse of decomposition yields the original sequence
            # of integers.
            decomp = tensordecomp(k, p)
            original_sequence = collect(0:(k^p - 1))
            @test all(decomp * k .^ collect(0:p-1) .== original_sequence)
        end
    end
end


# An Initializer is used to initialize network parameters
export Initializer
export ConstantInitializer
export XavierInitializer
export GaussianInitializer
export OrthogonalInitializer

@compat abstract type Initializer end # The root type of all initializer

immutable NullInitializer <: Initializer end
function init(::NullInitializer, blob::Blob)
  # do nothing
end

immutable ConstantInitializer <: Initializer
  value
end

function init(initializer::ConstantInitializer, blob::Blob)
  fill!(blob, initializer.value)
end

################################################################################
# An initializer  based on the paper [Bengio and Glorot 2010]: Understanding
# the difficulty of training deep feedforward neuralnetworks, but does not
# use the fan_out value.
#
# It fills the incoming matrix by randomly sampling uniform data from
# [-scale, scale] where scale = sqrt(3 / fan_in) where fan_in is the number
# of input nodes.
#
# For a ND-tensor blob parameter, the product of the 1 ~ (N-1) dimensions
# are considered as fan-in, and the last dimension is considered as fan-out.
################################################################################
immutable XavierInitializer <: Initializer
end
function init(initializer::XavierInitializer, blob::Blob)
  fan_in = get_fea_size(blob)
  scale = convert(eltype(blob), sqrt(3.0 / fan_in))
  init_val = rand(eltype(blob), size(blob)) * 2scale - scale
  copy!(blob, init_val)
end

immutable GaussianInitializer <: Initializer
  mean :: AbstractFloat
  std  :: AbstractFloat
end
GaussianInitializer(;mean=0.0, std=1.0) = GaussianInitializer(mean, std)

function init(initializer::GaussianInitializer, blob::Blob)
  init_val = randn(size(blob)) * initializer.std + initializer.mean
  init_val = convert(Array{eltype(blob)}, init_val)
  copy!(blob, init_val)
end

####### Orthogonal Initializer ##############################################
#
# Based on https://github.com/Lasagne/Lasagne/blob/master/lasagne/init.py
#
#############################################################################

immutable OrthogonalInitializer <: Initializer
  gain::AbstractFloat
end
OrthogonalInitializer() = OrthogonalInitializer(1.0) # but use OrthogonalInitializer(sqrt(2)) for ReLU units

function init(initializer::OrthogonalInitializer, blob::Blob)
  dims = size(blob)
  if length(dims) < 2
    error("Must have length(size(blob)) >= 2")
  end
  flatshape = (dims[1], prod(dims[2:end]))
  x = randn(flatshape)
  u, _, v = svd(x)
  x = (size(u) == flatshape) ? u : v
  copy!(blob, convert(Array{eltype(blob)}, x*initializer.gain))
end


"""
    varN(B::LogBinner[, lvl])

Calculates the variance/N of a given level in the Binning Analysis.
"""
function varN(B::LogBinner, lvl::Integer = _reliable_level(B))
    n = B.count[lvl]
    var(B, lvl) / n
end

"""
    var(B::LogBinner[, lvl])

Calculates the variance of a given level in the Binning Analysis.
"""
function var(B::LogBinner) end

function var(
        B::LogBinner{T,N},
        lvl::Integer = _reliable_level(B)
    ) where {N, T <: Real}

    n = B.count[lvl]
    X = B.x_sum[lvl]
    X2 = B.x2_sum[lvl]

    # lvl = 1 <=> original values
    # correct variance:
    # (∑ xᵢ^2) / (N-1) - (∑ xᵢ)(∑ xᵢ) / (N(N-1))
    X2 / (n - 1) - X^2 / (n*(n - 1))
end

function var(
        B::LogBinner{T,N},
        lvl::Integer = _reliable_level(B)
    ) where {N, T <: Complex}

    n = B.count[lvl]
    X = B.x_sum[lvl]
    X2 = B.x2_sum[lvl]

    # lvl = 1 <=> original values
    (real(X2) + imag(X2)) / (n - 1) - (real(X)^2 + imag(X)^2) / (n*(n - 1))
end

function var(
        B::LogBinner{<: AbstractArray{T, D}, N},
        lvl::Integer = _reliable_level(B)
    ) where {N, D, T <: Real}

    n = B.count[lvl]
    X = B.x_sum[lvl]
    X2 = B.x2_sum[lvl]

    @. X2 / (n - 1) - X^2 / (n*(n - 1))
end

function var(
        B::LogBinner{<: AbstractArray{T, D}, N},
        lvl::Integer = _reliable_level(B)
    ) where {N, D, T <: Complex}

    n = B.count[lvl]
    X = B.x_sum[lvl]
    X2 = B.x2_sum[lvl]

    @. (real(X2) + imag(X2)) / (n - 1) - (real(X)^2 + imag(X)^2) / (n*(n - 1))
end

"""
    all_vars(B::LogBinner)

Calculates the variance for each level of the Binning Analysis.
"""
function all_vars(B::LogBinner{T,N}) where {T,N}
    [var(B, lvl) for lvl in 1:N if B.count[lvl] > 1]
end

"""
    all_varNs(B::LogBinner)

Calculates the variance/N for each level of the Binning Analysis.
"""
function all_varNs(B::LogBinner{T,N}) where {T,N}
    [varN(B, lvl) for lvl in 1:N if B.count[lvl] > 1]
end

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

# NOTE works for all
"""
    mean(B::LogBinner[, lvl])

Calculates the mean for a given level in the Binning Analysis.
"""
function mean(B::LogBinner, lvl::Integer = 1)
    B.x_sum[lvl] / B.count[lvl]
end

# NOTE works for all
"""
    all_means(B::LogBinner)

Calculates the mean for each level of the `LogBinner`.
"""
function all_means(B::LogBinner{T,N}) where {T,N}
    [mean(B, lvl) for lvl in 1:N if B.count[lvl] > 1]
end

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

"""
    tau(B::LogBinner[, lvl])

Calculates the autocorrelation time tau.
"""
function tau(B::LogBinner{T,N}, lvl::Integer = _reliable_level(B)) where {N , T <: Number}
    var_0 = varN(B, 1)
    var_l = varN(B, lvl)
    0.5 * (var_l / var_0 - 1)
end
function tau(B::LogBinner{T,N}, lvl::Integer = _reliable_level(B)) where {N , T <: AbstractArray}
    var_0 = varN(B, 1)
    var_l = varN(B, lvl)
    @. 0.5 * (var_l / var_0 - 1)
end

"""
    all_taus(B::LogBinner)

Calculates the autocorrelation time tau for each level of the `LogBinner`.
"""
function all_taus(B::LogBinner{T,N}) where {T,N}
    [tau(B, lvl) for lvl in 1:N if B.count[lvl] > 1]
end

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

# Heuristic for selecting the level with the (presumably) most reliable
# standard error estimate:
# Take the highest lvl with at least 32 bins.
# (Chose 32 based on https://doi.org/10.1119/1.3247985)
function _reliable_level(B::LogBinner{T,N})::Int64 where {T,N}
    isempty(B) && (return 1)                # results in NaN in std_error
    i = findlast(x -> x >= 32, B.count)
    something(i, 1)
end

 """
    std_error(B::LogBinner[, lvl])

Calculates the standard error of the mean.
"""
function std_error(B::LogBinner) end

function std_error(B::LogBinner{T,N}, lvl::Integer=_reliable_level(B)) where {N, T <: Number}
    sqrt(varN(B, lvl))
end
function std_error(B::LogBinner{T,N}, lvl::Integer=_reliable_level(B)) where {N, T <: AbstractArray}
    sqrt.(varN(B, lvl))
end

"""
    all_std_errors(B::LogBinner)

Calculates the standard error for each level of the Binning Analysis.
"""
function all_std_errors(B::LogBinner) end

all_std_errors(B::LogBinner{T,N}) where {N, T <: Number} = sqrt.(all_varNs(B))
all_std_errors(B::LogBinner{T,N}) where {N, T <: AbstractArray} = (x -> sqrt.(x)).(all_varNs(B))

"""
    convergence(B::LogBinner, lvl)

Computes the difference between the variance of this lvl and the last,
normalized to the last lvl. If this value tends to 0, the Binning Analysis has
converged.
"""
function convergence(B::LogBinner) end

function convergence(B::LogBinner{T,N}, lvl::Integer=_reliable_level(B)) where {N, T <: Number}
    abs((varN(B, lvl+1) - varN(B, lvl)) / varN(B, lvl))
end
function convergence(B::LogBinner{T,N}, lvl::Integer=_reliable_level(B)) where {N, T <: AbstractArray}
    mean(abs.((varN(B, lvl+1) .- varN(B, lvl)) ./ varN(B, lvl)))
end

"""
    has_converged(B::LogBinner, lvl[, threshhold = 0.05])

Returns true if the Binning Analysis has converged for a given lvl.
"""
function has_converged(B::LogBinner, lvl::Integer=_reliable_level(B), threshhold::Float64 = 0.05)
    convergence(B, lvl) <= threshhold
end


## Bayesian Linear Regression
## ==========================

using Sigma
using Distributions

Σ_1 = 2 * eye(5,5)
# Zero mean
μ = MvNormal(zeros(5), Σ_1)
Σ_2 = eye(5,5)
Σ_prior = Wishart(5,Σ_2)
w = MvNormal(rand(μ), rand(Σ_prior))
x = [Uniform(-1,1) for j = 1:5, i = 1:500]
τ = Gamma(0.5,2)
ϵ = [Normal(0,1/rand(τ)) for i = 1:500]
y = [rand(w)' * rand(x[i,j]) + ϵ[i] for j=1:5, i = 1:500]

dist(w, y = x[])

# Nubmer of random variables is
num_rvs = 5*5 + 5*500 + 1 * 500