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

"""
    ce = 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)`.
"""

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

"""
    cross_entropy(ŷ, y::Vec{<:Finite})

Given an abstract vector of `UnivariateFinite` distributions `ŷ` (ie,
probabilistic predictions) and an abstract vector of true observations
`y`, return the negative log-probability that each observation would
occur, according to the corresponding probabilistic prediction.

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

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)

# 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

# TODO: support many distributions/samplers D below:
struct BrierScore{D} <: Measure end

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

Given an abstract vector of distributions `ŷ` 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``

For more information, run `info(brier_score)`.
"""
function BrierScore(; distribution=UnivariateFinite)
    distribution == UnivariateFinite ||
        error("Only `UnivariateFinite` Brier scores currently supported. ")
    return BrierScore{distribution}()
end

metadata_measure(BrierScore;
    target_scitype           = Vec{<:Finite},
    prediction_type          = :probabilistic,
    orientation              = :score,
    reports_each_observation = true,
    is_feature_dependent     = false,
    supports_weights         = true,
    distribution_type        = UnivariateFinite)

# adjustments
MMI.name(::Type{<:BrierScore{D}}) where D      = "BrierScore($(string(D)))"
MMI.docstring(::Type{<:BrierScore{D}}) where D =
    "Brier proper scoring rule for distributions of type $D; " *
    "aliases: `BrierScore($D)`."

# 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

# ---------------------------------------------------
## CLASSIFICATION METRICS (DETERMINISTIC PREDICTIONS)
# ==> MisclassificationRate / MCR
# ==> Accuracy
# Binary only
# = order independent
# ==> AUC
# = order dependent
# ==> Recall
# ==> Precision
# ---------------------------------------------------

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."

struct MisclassificationRate <: Measure end

"""
    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)`.
You can also equivalently use `mcr`.
"""
const misclassification_rate = MisclassificationRate()
const mcr = misclassification_rate

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`.")

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)

struct Accuracy <: Measure end

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

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`.")

struct BalancedAccuracy <: Measure end

const BACC = BalancedAccuracy

"""
    balanced_accuracy(ŷ, y [, w])
    bacc(ŷ, y [, w])
    bac(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 = BACC()
const bacc = balanced_accuracy
const bac  = bacc

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

metadata_measure(BACC;
    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`.")

## Binary but order independent

struct MatthewsCorrelation <: Measure end

const MCC = MatthewsCorrelation

"""
    matthews_correlation(ŷ, y)
    mcc(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 = MCC()
const mcc = matthews_correlation

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

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`")

struct AUC <: Measure end

"""
    auc(ŷ, y)

Return the Area Under the (ROC) Curve for probabilistic prediction `ŷ` given
true observations `y`.
$INVARIANT_LABEL_BINARY

For more information, run `info(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

metadata_measure(AUC;
    name                     = "auc",
    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: `auc`.")

## Binary and order dependent

const CM2 = ConfusionMatrix{2}

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

# synonyms
const TPR = TruePositiveRate
const TNR = TrueNegativeRate
const FPR = FalsePositiveRate
const FNR = FalseNegativeRate

const FDR = FalseDiscoveryRate
const PPV = Precision

const Recall      = TPR
const Specificity = TNR

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

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

## Names and synonyms
# NOTE: nothing for precision as there is Base.precision

const truepositive  = TruePositive()
const truenegative  = TrueNegative()
const falsepositive = FalsePositive()
const falsenegative = FalseNegative()

const tp = truepositive
const tn = truenegative
const fp = falsepositive
const fn = falsenegative

const truepositive_rate  = TPR()
const truenegative_rate  = TNR()
const falsepositive_rate = FPR()
const falsenegative_rate = FNR()

const tpr = truepositive_rate
const tnr = truenegative_rate
const fpr = falsepositive_rate
const fnr = falsenegative_rate

const falsediscovery_rate = FDR()

const fdr = falsediscovery_rate
const negativepredictive_value = NPV()
const positivepredictive_value = Precision()
const npv = NPV()
const ppv = precision

const recall       = truepositive_rate
const sensitivity  = recall
const hit_rate     = recall
const miss_rate    = falsenegative_rate
const fallout      = falsepositive_rate
const specificity  = truenegative_rate
const selectivity  = specificity
const f1score      = FScore{1}()

const balanced_accuracy = BACC()
const bacc = balanced_accuracy

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.((FPR, FNR, FDR);
    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.((TPR, TNR, Precision, FScore, 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})       = "tp"
MMI.docstring(::Type{<:TruePositive})  = "Number of true positives; " *
                                         "aliases: `truepositive`, `tp`."
MMI.name(::Type{<:TrueNegative})       = "tn"
MMI.docstring(::Type{<:TrueNegative})  = "Number of true negatives; " *
                                         "aliases: `truenegative`, `tn`."
MMI.name(::Type{<:FalsePositive})      = "fp"
MMI.docstring(::Type{<:FalsePositive}) = "Number of false positives; " *
                                         "aliases: `falsepositive`, `fp`."
MMI.name(::Type{<:FalseNegative})      = "fn"
MMI.docstring(::Type{<:FalseNegative}) = "Number of false negatives; " *
                                         "aliases: `falsenegative`, `fn`."

MMI.name(::Type{<:TPR})      = "tpr"
MMI.docstring(::Type{<:TPR}) = "True postive rate; aliases: " *
                               "`truepositive_rate`, `tpr`, `sensitivity`, " *
                               "`recall`, `hit_rate`."
MMI.name(::Type{<:TNR})      = "tnr"
MMI.docstring(::Type{<:TNR}) = "true negative rate; aliases: " *
                               "`truenegative_rate`, `tnr`, `specificity`, " *
                               "`selectivity`."
MMI.name(::Type{<:FPR})      = "fpr"
MMI.docstring(::Type{<:FPR}) = "false positive rate; aliases: " *
                               "`falsepositive_rate`, `fpr`, `fallout`."
MMI.name(::Type{<:FNR})      = "fnr"
MMI.docstring(::Type{<:FNR}) = "false negative rate; aliases: " *
                               "`falsenegative_rate`, `fnr`, `miss_rate`."
MMI.name(::Type{<:FDR})      = "fdr"
MMI.docstring(::Type{<:FDR}) = "false discovery rate; aliases: " *
                               "`falsediscovery_rate`, `fdr`."
MMI.name(::Type{<:NPV})      = "npv"
MMI.docstring(::Type{<:NPV}) = "negative predictive value; aliases: " *
                               "`negativepredictive_value`, `npv`."

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

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


# Unstable Bickley jet in GeophysicalFlows

ENV["GKSwstype"] = "nul"
using Printf
using Random
using Plots

using FFTW: rfft, irfft
using CUDA
using JLD2

using FourierFlows
using GeophysicalFlows: TwoDNavierStokes

include("Bickley.jl")

using .Bickley

function run(; dev = CPU(), nx = 128, dt = 1e-2)

    nsubs = ceil(Int, 2 / dt) # number of steps between each plot
    nsteps = 100 * nsubs # total number of steps

    problem = TwoDNavierStokes.Problem(dev; nx = nx,
                                            Lx = 4π,
                                            ny = nx,
                                            Ly = 4π,
                                            dt = dt,
                                            stepper = "FilteredRK4")

    ϵ = 0.1
    k = 0.5
    ℓ = 0.5

    DeviceArray = dev isa CPU ? Array : CuArray

    grid = problem.grid
    x = repeat(reshape(grid.x, grid.nx, 1), 1, grid.ny)
    y = repeat(reshape(grid.y, 1, grid.ny), grid.nx, 1)

    ψᵢ = DeviceArray(@. Bickley.Ψ(y, grid.Ly) + ϵ * Bickley.ψ̃(x, y, k, ℓ))

    ψᵢh = rfft(ψᵢ)
    CUDA.@allowscalar ψᵢh[1, 1] = 0

    ψᵢ = irfft(ψᵢh, grid.nx)

    @. problem.vars.zetah = - grid.Krsq * ψᵢh

    ζᵢ = irfft(problem.vars.zetah, problem.grid.nx)

    TwoDNavierStokes.set_zeta!(problem, ζᵢ)
    TwoDNavierStokes.set_c!(problem, DeviceArray(sin.(y/2)))

    # ## Output

    # We choose folder for outputing `.jld2` files and snapshots (`.png` files).
    filename = "geophysical_flows_unstable_bickley_jet_Nh$nx.jld2"

    isfile(filename) && rm(filename, force=true)

    output = Output(problem, filename,
        (:ζh, p -> Array(p.vars.zetah)),
        (:ζ,  p -> Array(irfft(p.vars.zetah, p.grid.nx))),
        (:ch, p -> Array(p.vars.ch)),
        (:c,  p -> Array(irfft(p.vars.ch, p.grid.nx))),
    )

    saveproblem(output)

    start_time = time_ns()

    e₀ = TwoDNavierStokes.energy(problem)
    χ₀ = TwoDNavierStokes.tracer_variance(problem)

    for j = 0:Int(nsteps/nsubs)

        stepforward!(problem, nsubs)
        TwoDNavierStokes.updatevars!(problem)
        saveoutput(output)

        e = TwoDNavierStokes.energy(problem)
        χ = TwoDNavierStokes.tracer_variance(problem)

        log = @sprintf("Step: %04d, t: %d, Δe/e₀: %.2e, Δχ/χ₀: %.2e, wall_time: %.2f seconds",
                       problem.clock.step,
                       problem.clock.t,
                       e / e₀ - 1,
                       χ / χ₀ - 1,
                       (time_ns() - start_time) * 1e-9,
                      )
        println(log)

    end

    run_time = (time_ns() - start_time) * 1e-9

    return output.path, run_time, nsteps
end

function visualize(filename)

    file = jldopen(filename)

    iterations = parse.(Int, keys(file["snapshots/t"]))

    x = file["grid/x"]
    y = file["grid/y"]

    animation = @animate for (i, iter) in enumerate(iterations)

        @info "Plotting frame $i of $(length(iterations))"

        t = file["snapshots/t/$iter"]
        ζ = file["snapshots/ζ/$iter"]
        c = file["snapshots/c/$iter"]

        kwargs = Dict(
                      :aspectratio => 1,
                      :linewidth => 0,
                      :ticks => nothing,
                      :colorbar => :none,
                      :clims => (-1, 1),
                      :xlims => (-2π, 2π),
                      :ylims => (-2π, 2π),
                     )

        ζ_plot = heatmap(x, y, clamp.(ζ, -1, 1)'; color=:balance, kwargs...)
        c_plot = heatmap(x, y, clamp.(c, -1, 1)'; color=:thermal, kwargs...)

        ω_title = @sprintf("ζ at t = %.1f", t)
        c_title = @sprintf("c at t = %.1f", t)

        plot(ζ_plot, c_plot, title = [ω_title c_title], size=(4000, 2000))
    end

    close(file)

    name = filename[1:end-5]

    gif(animation, name * ".gif", fps=8)

    return nothing
end

cfl = 0.1
for nx in (32,
           64, 128, 256, 512,
           1024, 2048)
           #32, 64, 128, 256, 512)

    @show dt = cfl * 4π / nx
    filename, run_time, nsteps = run(dev=GPU(), nx=nx, dt=dt)

    DOF = nx^2
    @info "Results for nx = $nx !"
    @show cost = run_time / (nsteps * DOF)

    filename = "geophysical_flows_unstable_bickley_jet_Nh$nx.jld2"
    visualize(filename)
end


"""
find_kbest_assignments(C, k, maximize=false) -> kbest_sols

Find the k lowest [or highest] cost 2D assignments for the two-dimensional
assignment problem with a rectangular cost matrix C.

# Example
```julia
julia> M=rand(1:100,3,4)
3×4 Matrix{Int64}:
 77  51  42  67
 72  53  47   4
 24  50  77  96

 julia> sols = find_kbest_assigments(M, 5)
 5-element Vector{Assignment.AssignmentSolution{Int64, Int64}}:
  AssignmentSolution(CartesianIndex.(1:3, [3, 4, 1]), 70)
  AssignmentSolution(CartesianIndex.(1:3, [2, 4, 1]), 79)
  AssignmentSolution(CartesianIndex.(1:3, [3, 4, 2]), 96)
  AssignmentSolution(CartesianIndex.(1:3, [3, 2, 1]), 119)
  AssignmentSolution(CartesianIndex.(1:3, [2, 3, 1]), 122)

 julia> max_sols = find_kbest_assignments(M, 5, true)
 5-element Vector{Assignment.AssignmentSolution{Int64, Int64}}:
  AssignmentSolution(CartesianIndex.(1:3, [1, 2, 4]), 226)
  AssignmentSolution(CartesianIndex.(1:3, [1, 3, 4]), 220)
  AssignmentSolution(CartesianIndex.(1:3, [2, 1, 4]), 219)
  AssignmentSolution(CartesianIndex.(1:3, [4, 1, 3]), 216)
  AssignmentSolution(CartesianIndex.(1:3, [3, 1, 4]), 210)
```

This is an implementation of Murty's method, which is described in [1].
The algorithm relies on the existence of a 2D assignment algorithm. The 2D
assignment algorithm of [2] is used. Additionally, the dual variable
inheritance methods described in [3] is used to reduce the computational
complexity of the technique.

Murty's algorithm runs 2D assignment algorithms a number of times with an
increasing number of constraints. Instances of MurtyData are stored in an
ordered list implemented using the BinaryMinHeap class.

This code is a port of Matlab code released in the public domain by the
US Naval Research Laboratory [4].

This work is not affliated with or endorsed by the US Naval Research Laboratory.

# References
[1] K. G. Murty, "An algorithm for ranking all the assignments in order of
   increasing cost;" Operations Research; vol. 16; no. 3; pp. 682-687
   May-Jun. 1968.
[2] D. F. Crouse, "On Implementing 2D Rectangular Assignment Algorithms,"
   IEEE Transactions on Aerospace & Electronic Systems; vol. 52; no. 4
   pp. 1679-1696; Aug. 2016.
[3] M. L. Miller, H. S. Stone, & J. Cox, Ingemar, "Optimizing Murty's
   ranked assignment method;" IEEE Transactions on Aerospace &
   Electronic Systems; vol. 33; no. 3; pp. 851-862; Jul. 1997.
[4] D. F. Crouse, "The Tracker Component Library: Free Routines for Rapid
   Prototyping," IEEE Aerospace and Electronic Systems Magazine, vol. 32,
   no. 5, pp. 18-27, May. 2017

October 2013 David F. Crouse; Naval Research Laboratory; Washington D.C.
(UNCLASSIFIED) DISTRIBUTION STATEMENT A. Approved for public release.
"""
function find_kbest_assignments(C::AbstractMatrix{T}, k::Integer, maximize=false) where T
    numRow, numCol = size(C)

	if numRow > numCol
		sols = find_kbest_assignments(C', k, maximize)
		for (i, sol) in enumerate(sols)
			sols[i] = sol'
		end
		return sols
	end

	#The cost matrix must have all non-negative elements for the assignment
	#algorithm to work. This forces all of the elements to be positive. The
	#delta is added back in when computing the gain in the end.
    if maximize
        CDelta = maximum(C)
        C = -C.+ CDelta
	else
        CDelta = minimum(C)
        C =  C .- CDelta
    end

	#Step 1: Find the best assignment. It will become the root from which()
	#everything else splits.
    #MUST PAD IT TO MAKE IT SQUARE; because we are using suboptimal
    #dual variables.
	numPad = numCol-numRow
	numPad > 0 && (C = [C; zeros(T, numPad, numCol)])

    LCHyp = MurtyData(C, numRow)

	#Check for feasibility.
    if isnothing(LCHyp.gainFull)
        return AssignmentSolution{T, _dual_eltype(T)}[]
    end

	k_best = [AssignmentSolution(LCHyp, CDelta, maximize)]
	sizehint!(k_best, k)

	hyp_list = BinaryMinHeap([LCHyp])

	ScannedRows = falses(numCol)
	ScannedCol = falses(numCol)
	pred = zeros(Int, numCol)
	shortestPathCost = fill(typemax(float(T)), numCol)

    for curSweep in 2:k
        #We have to successively split the LC hypothesis either k times |
        #until splitting is no longer possible & get rid of the
        #hypothesis from the list.
		data = pop!(hyp_list)

		#Basically; it is all columns that have not been assigned in
		#one of the subproblems.
		col2Scan = @view data.col4rowLCFull[data.activeRow:end]

		#The loop ends at numVarRow; because we are not going to allow
		#variable hypotheses to be applied to dummy rows.
		for curRow=data.activeRow:data.numVarRow
			#forbiddenColumns contains a 1 for every column that one is
			#not allowed to assign to curRow.
			if curRow==data.activeRow
				forbiddenColumns = data.forbiddenActiveCol
			else
				forbiddenColumns = falses(numCol)
				forbiddenColumns[data.col4rowLCFull[curRow]] = true
			end

			#Get rid of the current assignment.
			row4colInit = data.row4colLCFull[:]
			col4rowInit = data.col4rowLCFull[:]
			row4colInit[col4rowInit[curRow]] = 0
			col4rowInit[curRow] = 0

			splitHyp = ConstrainedShortestPath!(
				data.A, data.numVarRow, curRow,
				forbiddenColumns, col4rowInit, row4colInit, col2Scan[:], data.u[:], data.v[:],
				ScannedRows, ScannedCol, pred, shortestPathCost
			)

			#If a valid assignment was produced.
			!isnothing(splitHyp.gainFull) && push!(hyp_list, splitHyp)

			#Remove the current assigned column from the list of
			#columns that can be scanned.
			col2Scan = @view col2Scan[col2Scan .!= data.col4rowLCFull[curRow]]
		end

        if !isempty(hyp_list)
			#Save the ordered best solutions.
			push!(k_best, AssignmentSolution(first(hyp_list), CDelta, maximize))
        else
            break #All possible hypotheses have been enumerated.
        end
    end

	return k_best
end

struct MurtyData{T, M<:AbstractMatrix{T}, S}
	A::M
	numVarRow::Int
	activeRow::Int
	forbiddenActiveCol::BitVector
	col4rowLCFull::Vector{Int}
	row4colLCFull::Vector{Int}
	gainFull::Union{T, Nothing}
	u::Vector{S}
	v::Vector{S}
end

function MurtyData(A::AbstractMatrix, numVarRow::Integer)
	numRow, numCol = size(A)

	sol = find_best_assignment(A')
	forbiddenActiveCol = falses(numCol)
	if !isempty(sol.col4row)
		forbiddenActiveCol[sol.col4row[1]] = 1
	end

	return MurtyData(
		A, numVarRow, 1, forbiddenActiveCol,
		sol.col4row, sol.row4col, sol.cost, sol.v, sol.u)
end

Base.isless(a::MurtyData, b::MurtyData) = Base.isless(a.gainFull, b.gainFull)

for op in (:<, :>, :(<=), :(>=), :(==), :(!=))
	@eval begin
		$op(a::MurtyData, B::MurtyData) = $op(a.gainFull, b.gainFull)
		$op(a::MurtyData, B) = $op(a.gainFull, b)
		$op(a, B::MurtyData) = $op(a, b.gainFull)
	end
end

function AssignmentSolution(data::MurtyData{T, M, S}, CDelta=T(0), maximize=false) where {T, M, S}
	numCol = size(data.A, 2)
	numRow = data.numVarRow

	col4row = data.col4rowLCFull[1:numRow]
	row4col = data.row4colLCFull[:]
	if numRow < numCol
		for (i, j) in enumerate(row4col)
			if j > numRow
				row4col[i] = 0
			end
		end
	end
	gain =  maximize ? CDelta * numRow  - data.gainFull : CDelta * numRow + data.gainFull

	return AssignmentSolution{T, S}(col4row, row4col, T(gain), data.u[:], data.v[1:numRow])
end

function ConstrainedShortestPath!(
		C::AbstractMatrix{T},
		numVarRow::Integer,
		activeRow::Integer,
		forbiddenActiveCols::AbstractVector{Bool},
		col4row::AbstractVector{Int},
		row4col::AbstractVector{Int},
		col2Scan::AbstractVector{Int},
		u::AbstractVector{S},
		v::AbstractVector{S},
		ScannedRows::AbstractVector{Bool} = falses(size(C,1)),
		ScannedCol::AbstractVector{Bool} = falses(size(C,2)),
    	pred::AbstractVector{Int} = zeros(Int, size(C,2)),
		shortestPathCost::AbstractVector = fill(Inf, size(C,2))
	) where {T, S}

	numRow = size(C,1)
	numCol = size(C,2)
    numCol2Scan=length(col2Scan)
	ScannedRows[:] .= false
	ScannedCol[:] .= false
	pred[:] .= 0
	shortestPathCost[:] .= Inf

	sink = 0
	delta = S(0)
    curRow = activeRow
	closestColScan=0

	while sink==0
		ScannedRows[curRow]=1;

        # Scan all of the columns that have not already been scanned.
        minVal=Inf;
        @inbounds for curColScan=1:numCol2Scan
            curCol=col2Scan[curColScan];

			# Do not scan forbidden columns.
            curRow==activeRow && forbiddenActiveCols[curCol] && continue

            reducedCost=delta+C[curRow,curCol]-u[curRow]-v[curCol]
            if reducedCost < shortestPathCost[curCol]
                pred[curCol] = curRow
                shortestPathCost[curCol] = reducedCost
            end

            # Find the minimum unassigned column that was
            # scanned.
            if(shortestPathCost[curCol]<minVal)
                minVal=shortestPathCost[curCol];
                closestColScan=curColScan;
            end
        end

		# If the minimum cost column is not finite, then the problem is
		# not feasible.
		!isfinite(minVal) && return MurtyData(
			C, numVarRow, activeRow, forbiddenActiveCols,
			Int[], Int[], nothing, u, v)

        closestCol=col2Scan[closestColScan];

        # Add the column to the list of scanned columns and delete it from
        # the list of columns to scan.
        ScannedCol[closestCol] = 1
        numCol2Scan -= 1
		@inbounds for i in closestColScan:numCol2Scan
			col2Scan[i] = col2Scan[i + 1]
		end

        delta = S(shortestPathCost[closestCol])

        # If we have reached an unassigned row.
        if row4col[closestCol]==0
            sink=closestCol
        else
            curRow=row4col[closestCol]
        end
    end

	# Dual Update Step
    # Update the first row in the augmenting path.
    u[activeRow] = u[activeRow] + delta;
    # Update the rest of the rows in the agumenting path.
	@inbounds for (i, col) in enumerate(ScannedRows)
		if col != 0 && i != activeRow
			u[i] += delta - S(shortestPathCost[col4row[i]])
		end
	end
    # Update the scanned columns in the augmenting path.
	@inbounds for (i, row) in enumerate(ScannedCol)
		if row != 0
			v[i] -= delta - S(shortestPathCost[i])
		end
	end

	# Augmentation
    # We have to remove node k from those that must be assigned.
    j = sink
    while true
        i = row4col[j] = pred[j]
		j, col4row[i] = col4row[i], j
		i==activeRow && break
    end

    # Calculate the gain that should be returned. Because the duality gap is
    # zero, the gain can be most easily obtained from the dual cost
    # function.
    gain=T(0)
	@inbounds for (i, j) in enumerate(col4row)
        gain=gain+C[i, j];
    end

	forbiddenActiveCols[col4row[activeRow]] = true

	return MurtyData(
		C, numVarRow, activeRow, forbiddenActiveCols,
		col4row, row4col, gain, u, v)
end


using StudentHousing, StatsBase, Distributions, JuMP, Gurobi
include("utils.jl")

srand(32)

# =============================================================================
# Toy problem data
# =============================================================================
nhouses = 50
nbedrooms_range = collect(1:3)
nbedrooms_frequency = Weights([0.5, 0.25, 0.25])
nbathrooms_range = collect(1:2)
nbathrooms_frequency = Weights([0.5, 0.5])
prices_range_pp = [800.0, 1000.0]
area_ranges = [0.0, 800.0]
market_data = StudentHousing.MarketData(nbedrooms_range, nbedrooms_frequency,
    nbathrooms_range, nbathrooms_frequency, prices_range_pp, area_ranges)
budget = 1e3
d = StudentHousingData(market_data, nhouses = nhouses, budget = budget,
                  demand_distribution = Uniform(0.0, 10.0), easy=false)
nhouses = length(d.houses)

# =============================================================================
# Solve with direct formulation
# =============================================================================
m1 = StudentHousing.ia_model(d)
@assert solve(m1, relaxation=true) == :Optimal
getobjectivevalue(m1)

# =============================================================================
# Solve problem with column generation algorithm
# =============================================================================
m, V_generated, λ_generated, house_choice = solve_ia_generation(d)

# =============================================================================
# Look at solution
# =============================================================================
println(getobjectivevalue(m))

# If not integer...
for i = 1:nhouses
    setcategory(m[:λ][i], :Bin)
end
for λnew in λ_generated
    setcategory(λnew, :Bin)
end
solve(m)

# Recover the solution

# In case one of our initial columns were in the optimal basis
results = zeros(nhouses, length(d.patterns_allow))
for k = 1:nhouses
    if getvalue(m[:λ][k]) ≈ 1.0
        results[k, StudentHousing.house_allowedby(k, d)[1]] = 1.0
    end
end
# Check all the other columns we added
for k = 1:length(V_generated)
    if getvalue(λ_generated[k]) ≈ 1.0
        # Find the assignment that we picked for our house
        patterns = find(V_generated[k] .≈ 1.0)
        # Check the house index
        h = house_choice[k]
        # Mark the assignment
        results[h, patterns] .= 1.0
    end
end

println("House 1 assigned to:", find(results[1,:] .≈ 1.0))
# House 1 assigned to:[6, 7]

# sanity check
StudentHousing.house_allowedby(1, d)
# 5-element Array{Int64,1}:
#   6
#   7
#  11
#  12
#  13

println("House 2 assigned to:", find(results[2,:] .≈ 1.0))
# House 2 assigned to:[3]
# sanity check
StudentHousing.house_allowedby(2, d)
# 11-element Array{Int64,1}:
#   3
#   4
#   5
#   6
#   7
#   8
#   9
#  10
#  11
#  12
#  13

for i = 1:nhouses
    println(println("House $i assigned to:", find(results[i,:] .≈ 1.0)))
end

# =============================================================================
# Compare with original single-stage MIP, 329 pattern example
# =============================================================================
nbedrooms_range = collect(1:2)
nbedrooms_frequency = Weights([0.75, 0.25])
nbathrooms_range = collect(1:2)
nbathrooms_frequency = Weights([0.75; 0.25])
prices_range_pp = [800.0, 1000.0]
area_ranges = [0.0, 700.0, 850.0]
budget = 1000.0
market_data = StudentHousing.MarketData(nbedrooms_range, nbedrooms_frequency,
    nbathrooms_range, nbathrooms_frequency, prices_range_pp, area_ranges)
nhouses = 50
srand(1)
d = StudentHousingData(market_data, nhouses = nhouses, budget = budget,
                  demand_distribution = Uniform(0.0, 100.0))

# =============================================================================
# Solve with direct formulation
# =============================================================================
m1 = StudentHousing.ia_model(d)
@assert solve(m1) == :Optimal
println("Shortage = ", sum(sum(d.demands)) - getobjectivevalue(m1))
# Shortage = 16616.0
@assert solve(m1, relaxation=true) == :Optimal
println("Shortage = ", sum(sum(d.demands)) - getobjectivevalue(m1))
# Shortage = 16615.18328390582
# These are the same objectives we had with the original single stage
# formulation

# =============================================================================
# Solve problem with column generation algorithm, house subproblems
# =============================================================================
tic()
m2, V_generated, λ_generated, house_choice = solve_ia_generation(d)
toc()
# elapsed time: 2.031915379 seconds

# =============================================================================
# Look at solution
# =============================================================================
println(getobjectivevalue(m2))
# 25.63353800440406
println("Shortage = ", sum(sum(d.demands)) - getobjectivevalue(m2))
# Shortage = 16615.366461995596
# a little stronger than the LP relaxation

# Not integer, so get integer solution with existing columns
r = recover_integer_soln(m2, d)
println("Integer shortage = ", sum(sum(d.demands)) - sum(sum(r)))
# Integer shortage = 16616.0
# ... which matches the MIP solution


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

#=
Williams A, Barrus S, Morley R K, et al., 2005.
(https://dl.acm.org/doi/abs/10.1145/1198555.1198748)
=#
function intersection(f, r::Ray{Dim,T}, b::Box{Dim,T}) where {Dim,T}
  invdir = one(T) ./ direction(r)
  lo, up = coordinates.(extrema(b))
  orig = coordinates(origin(r))

  tmin = zero(T)
  tmax = typemax(T)

  # check for intersection with slabs along with each axis
  for i in 1:Dim
    imin = (lo[i] - orig[i]) * invdir[i]
    imax = (up[i] - orig[i]) * invdir[i]

    # swap variables if necessary
    invdir[i] < zero(T) && ((imin, imax) = (imax, imin))

    # the ray is on a face of the box, avoid NaN
    (isnan(imin) || isnan(imax)) && continue

    (tmin > imax || imin > tmax) && return @IT NoIntersection nothing f

    tmin = max(tmin, imin)
    tmax = min(tmax, imax)
  end

  tmin ≈ tmax && return @IT TouchingRayBox r(tmin) f

  return @IT CrossingRayBox Segment(r(tmin), r(tmax)) f
end


if basename(pwd()) == "aoc"
    cd("2021/15")
end
using Pkg
Pkg.activate(".")
using Base.Iterators
using Graphs, SimpleWeightedGraphs

function loadfile(filename::AbstractString)
    transpose(reshape(
        [parse(Int, c) for line in eachline(filename) for c in line],
        length(readline(filename)),
        :
    ))
end

function loadbig(filename::AbstractString)
    A = loadfile(filename)
    B = hcat((A .+ n for n = 0:4)...)
    C = vcat((B .+ n for n = 0:4)...)
    (C .- 1) .% 9 .+ 1
end

function neighbourhood(A::AbstractMatrix, i::CartesianIndex)
    (
        n + i
        for n in CartesianIndex.([(1, 0), (0, 1), (-1, 0), (0, -1)])
        if n + i in CartesianIndices(A)
    )
end

tolinear(A, i::CartesianIndex) = size(A, 2) * (i[1] - 1) + i[2]

function todigraph(A::AbstractMatrix{Int})::SimpleWeightedDiGraph
    g = SimpleWeightedDiGraph{Int,Int}(length(A))
    for i in CartesianIndices(A)
        for n in neighbourhood(A, i)
            @assert add_edge!(g, tolinear(A, i), tolinear(A, n), A[n])
        end
    end
    g
end

minpath(g::AbstractGraph) = last(dijkstra_shortest_paths(g, 1).dists)
minpath(A::AbstractMatrix) = minpath(todigraph(A))

part1(filename) = minpath(loadfile(filename))
part2(filename) = minpath(loadbig(filename))


using Test
import DifferentiableBackwardEuler
import DPFEHM
import NonlinearEquations
import PyPlot
import SparseArrays
import Zygote

#solve u_t + div(v * u) - div(D * grad(u)) - Q = 0 on an unstructured mesh using an upwind discretization of the velocity term (or downwind if use_upwind=false)
@NonlinearEquations.equations exclude=(areasoverlengths, volumes, dirichletnodes, coords) function transport(u, vxs, vys, vzs, Ds, neighbors, areasoverlengths, dirichletnodes, dirichletus, Qs, volumes, coords; use_upwind=true)
	@assert size(coords, 1) == 3#it is assumed that this is a 3d problem
	NonlinearEquations.setnumequations(length(u))
	for i = 1:length(u)
		if i in dirichletnodes
			NonlinearEquations.addterm(i, u[i] - dirichletus[i])
		else
			NonlinearEquations.addterm(i, Qs[i] / (volumes[i]))
		end
	end
	delta = zeros(eltype(coords), size(coords, 1))
	for (i, (node_a, node_b)) in enumerate(neighbors)
		for (node1, node2) in [(node_a, node_b), (node_b, node_a)]
			#add the advection term
			delta .= coords[:, node1] .- coords[:, node2]
			is_upwind = ((vxs[i] * delta[1] + vys[i] * delta[2] + vzs[i] * delta[3]) > 0)
			if (is_upwind && use_upwind) || (!is_upwind && !use_upwind)#!use_upwind basically means use_downwind and !is_upwind basically means is_downwind
				if !(node1 in dirichletnodes) && !(node2 in dirichletnodes)
					NonlinearEquations.addterm(node1, -(vxs[i] * delta[1] + vys[i] * delta[2] + vzs[i] * delta[3]) * (u[node1] - u[node2]) * areasoverlengths[i] / (volumes[node1]))
				elseif !(node1 in dirichletnodes) && node2 in dirichletnodes
					NonlinearEquations.addterm(node1, -(vxs[i] * delta[1] + vys[i] * delta[2] + vzs[i] * delta[3]) * (u[node1] - dirichletus[node2]) * areasoverlengths[i] / (volumes[node1]))
				end
			end
			#add the diffusion term
			if !(node1 in dirichletnodes) && !(node2 in dirichletnodes)
				NonlinearEquations.addterm(node1, Ds[i] * (u[node2] - u[node1]) * areasoverlengths[i] / (volumes[node1]))
			elseif !(node1 in dirichletnodes) && node2 in dirichletnodes
				NonlinearEquations.addterm(node1, Ds[i] * (dirichletus[node2] - u[node1]) * areasoverlengths[i] / (volumes[node1]))
			end
		end
	end
end

mins = [0, 0]; maxs = [3, 2]#size of the domain, in meters
ns = [30, 20]#number of nodes on the grid
coords, neighbors, areasoverlengths, volumes = DPFEHM.regulargrid2d(mins, maxs, ns, 1.0)#build the grid
coords = vcat(coords, zeros(size(coords, 2))')#turn it into a 3d coords

c = 1.0
tspan = [0.0, 1.0 / abs(c)]
u0 = map(i->coords[1, i] < 1 / 2  && coords[1, i] > 0 && coords[2, i] > 0.25 && coords[2, i] < 0.75 ? 1.0 : 0.0, 1:size(coords, 2)) + map(i->exp(-((coords[1, i] - 1.0) ^ 2 + (coords[2, i] - 1.0) ^ 2) * 64), 1:size(coords, 2))
vxs = map(neighbor->c, neighbors)
vys = map(neighbor->c * 0.5, neighbors)
vzs = zeros(length(neighbors))
Ds = 1e-2 * ones(length(neighbors))
Qs = zeros(size(coords, 2))
dirichletnodes = [collect(1:ns[2]); collect(size(coords, 2) - ns[2] + 1:size(coords, 2))]
dirichletus = zeros(size(coords, 2))
f(u, p, t) = transport_residuals(u, p, vys, vzs, Ds, neighbors, areasoverlengths, dirichletnodes, dirichletus, Qs, volumes, coords)
f_u(u, p, t) = transport_u(u, p, vys, vzs, Ds, neighbors, areasoverlengths, dirichletnodes, dirichletus, Qs, volumes, coords)
f_p(u, p, t) = transport_vxs(u, p, vys, vzs, Ds, neighbors, areasoverlengths, dirichletnodes, dirichletus, Qs, volumes, coords)
f_t(u, p, t) = zeros(length(u))
function solveforu(vxs)
	DifferentiableBackwardEuler.steps(u0, f, f_u, f_p, f_t, vxs, tspan[1], tspan[2]; abstol=1e-4, reltol=1e-4)
end
@time u_implicit = solveforu(vxs)

fig, axs = PyPlot.subplots(2, 1; figsize=(8, 8))
img = axs[1].imshow(reshape(u0, reverse(ns)...), extent=[mins[1], maxs[1], mins[2], maxs[2]], origin="lower")
fig.colorbar(img, ax=axs[1])
img = axs[2].imshow(reshape(u_implicit[:, end], reverse(ns)...), extent=[mins[1], maxs[1], mins[2], maxs[2]], origin="lower")
fig.colorbar(img, ax=axs[2])
fig.tight_layout()
display(fig)
PyPlot.println()
PyPlot.close(fig)

_, gradient_node = findmin(map(i->sum((coords[:, i] .- [1.5, 1.0, 0.0]) .^ 2), 1:size(coords, 2)))
g = x->solveforu(x)[gradient_node, end]
@time gradg = Zygote.gradient(g, vxs)[1]
function checkgradientquickly(f, x0, gradf, n; delta::Float64=1e-8, kwargs...)
	indicestocheck = sort(collect(1:length(x0)), by=i->abs(gradf[i]), rev=true)[1:n]
	f0 = f(x0)
	for i in indicestocheck
		x = copy(x0)
		x[i] += delta
		fval = f(x)
		grad_f_i = (fval - f0) / delta
		@test isapprox(gradf[i], grad_f_i; kwargs...)
	end
end
checkgradientquickly(g, vxs, gradg, 3; atol=1e-4, rtol=1e-3)
nothing


using Commas
using HMMs
using PyPlot

spx = Commas.readcomma( joinpath( homedir(), "myenvironment", "dev", "SPXHMM", "data", "SPX" ) )
y = log.(1 .+ spx[:vwretd])

parameters = zeros(3)
HMMs.fit_mle!( HMMs.GenTDist, parameters, y, ones(length(y)), Dict{Symbol,Any}() )

delta = 0.0002
bins = collect(-0.01:delta:0.01)

clf()
hist( y, bins=bins )

midpoints = (bins[1:end-1] + bins[2:end])/2
densities = delta*length(y).*pdf.( HMMs.GenTDist( parameters... ), midpoints )
plot( midpoints, densities )

delta = 0.001
bins = collect(-0.25:delta:0.25)

clf()
hist( y, bins=bins, log=true )

midpoints = (bins[1:end-1] + bins[2:end])/2
densities = delta*length(y).*pdf.( HMMs.GenTDist( parameters... ), midpoints )
plot( midpoints, densities )

mu = mean(y)
sigma = std(y)
abs(mean(((y .- mu)./sigma).^3))^(1/3)

symy = [y; 2*mu .- y]
symparameters = zeros(3)
HMMs.fit_mle!( HMMs.GenTDist, symparameters, symy, ones(length(symy)), Dict{Symbol,Any}() )
parameters

t = HMMs.GenTDist( hmm.stateparameters[:,2]... )
using PyPlot
delta = 0.0005
bins = collect(-0.1:delta:0.1)
midpoints = (bins[1:end-1] + bins[2:end])/2
tdensities = pdf.( t, midpoints )
clf()
semilogy( midpoints, tdensities )
mu = hmm.stateparameters[1,2]
sigma = hmm.stateparameters[2,2]
nu = hmm.stateparameters[3,2]
densities = pdf.( Normal( mu, sigma*sqrt(nu/(nu-2)) ), midpoints )
plot( midpoints, densities )


using FEM

# Load the VTK exporter
FEM.vtkexportmod()
using FEM.VTKExportMod

# Generates a mesh of the shape known as the "Cook membrane" which is a
# quadraterial with corners [0.0, 0.0], [0.0, 44.0], [48.0, 60.0], [48.0, 44.0]
# Possible mesh elements are
# - GeoTrig for 3 node triangles,
# - GeoQTrig for 6 node triangles
# - GeoQuad for 4 node quadraterials
n_ele_x = 100
n_ele_y = 100
geomesh = gencook(n_ele_x, n_ele_y, GeoQTrig)

# We create two node sets, one on the right edge and one on the left.
# This is done by giving an anonymous function that is satisfied by
# the edges.
push!(geomesh, gennodeset(n->n.coords.x>47.999999, "right", geomesh.nodes))
push!(geomesh, gennodeset(n->n.coords.x<0.00001, "left", geomesh.nodes))

# We create an element set containing all the elements
push!(geomesh, ElementSet("all", collect(1:length(geomesh.elements))))

# We create a material section of a linear isotropic material
# and assigns that section to all the elements.
mat_section = MaterialSection(FEM.linearisotropicmod().LinearIsotropic(1, 0.3))
push!(mat_section, geomesh.element_sets["all"])

# We create an element section and assign it to all the elements.
ele_section = ElementSection(FEM.quadtrigmod().QuadTrig)
push!(ele_section, geomesh.element_sets["all"])

# Apply Dirichlet BC to the left side in dofs for x-displacement (Du)
# and dofs in y-displacement (Dv)
bcs = Any[DirichletBC("0.0", [FEM.Du], geomesh.node_sets["left"]),
          DirichletBC("0.0", [FEM.Dv], geomesh.node_sets["left"])]

# Since we currently don't have edge load we give a nodal load
# on the node set on the right edge.
loads = Any[NodeLoad("1/($n_ele_y+1)", [FEM.Dv], geomesh.node_sets["right"])]

# Create the fe problem
fp = FEM.create_feproblem("cook_example_quad", geomesh, [ele_section], [mat_section], bcs, loads)

vtkexp = VTKExporter()
# Output fields are added by pushing them into the exporter.
# We want to export the stress and strain tensor as well as
# the von Mises stress so we push them.
push!(vtkexp, Stress)
push!(vtkexp, Strain)
push!(vtkexp, VonMises)
# We set the output of the exporter to be in ascii
set_binary!(vtkexp, false)

solver = NRSolver(rel_tol = 1e-7, max_iters = 4)

# Solve the fe problem using the solver and using the vtk exporter.
solve(solver, fp, vtkexp)


using DrWatson
@quickactivate "StatReth"

# %%
using CSV
using DataFrames
using Turing
using Plots

include(srcdir("quap.jl"))
include(srcdir("tools.jl"))

# %%
d = DataFrame(CSV.File(datadir("exp_raw/Howell_1.csv")))
sort!(d, :weight)  # to make the plots look nice
d2 = d[d.age .< 18, :]

z_weight, unz_weight = zscore_transform(d2.weight)
d2.weight_s = z_weight(d2.weight)

@model function height_linear(weight_s, height)
    α ~ Normal(100, 20)
    β ~ LogNormal(2, 1)
    σ ~ Exponential(5)
    μ = f_linear(weight_s, α, β)
    height ~ MvNormal(μ, σ)
end

f_linear(weight_s, α, β) = α .+ β .* weight_s

q = quap(height_linear(d2.weight_s, d2.height))
post = DataFrame(rand(q.distr, 1000)', ["α"; "β"; "σ"])

# %%
estimmean, estimlower, estimupper = estimparam(post)
"""
All values heavily rounded.
α = 108 ± 1 cm
β = 23 ± 1 cm/stddev(kg)
σ = 8.5 ± 0.5 cm
For β we need to convert this result into cm/kg
β = 23 / std(d2.weight) = 2.5

Average height of the children is 108 cm. On average they are 2.5 cm taller for every 1 kg
they are heavier. Children with the same weight have an expected standard deviation of their
height of 8.5 cm
"""

# %%
mu = f_linear(d2.weight_s, post.α', post.β')
mu_m, mu_lower, mu_upper = meanlowerupper(mu)

scatter(d2.weight, d2.height, alpha = 0.3, legend = false)
plot!(d2.weight, mu_m, ribbon = (mu_m .- mu_lower, mu_upper .- mu_m))

pred = rand.(Normal.(mu, post.σ'))
_, pred_lower, pred_upper = meanlowerupper(pred)

plot!(d2.weight, pred_lower, fillrange = pred_upper, alpha = 0.3, linealpha = 0.0, c = 2)

# %%
"""
Well, I'm concerned about the fact that I'm using a linear model even though the data
doesn't seem to be particularly linear. I mean, it's predicting there's going to be a 60 cm
tall child that doesn't weigh anything. Negative weights seen to be a-ok, too.
I suspect/hope that doing something like sqrt/cbrt or the likes would be better.
"""


module Contractors
using ValidatedNumerics
using DualNumbers

export N_rig, root, range_estimate, ShootingMethod

function N_rig(f, f′, x::Interval{T}) where {T}
	x_mid = Interval{T}(mid(x))
	return intersect(x, x_mid-f(x_mid)/f′(x))
end

root(f, x::Interval, ϵ; max_iter = 100) = root(f, x->f(Dual(x, 1)).epsilon, x, ϵ; max_iter = max_iter)

function root(f, f′, x::Interval, ϵ; max_iter = 100)
	for i in 1:max_iter
		x = N_rig(f, f′, x)
		if diam(x)<ϵ
			return x
		end
		if isempty(x)
			return x
		end
	end
	@info "Maximum iterates reached" max_iter
	return x
end

# superseded by IntervalOptimisation.jl
function range_estimate(f, domain, recstep = 5)
	if recstep == 1
		return f(domain)
	else
		a, b = bisect(domain)
		Iₐ = range_estimate(f, a, recstep-1)
		Iᵦ = range_estimate(f, b, recstep-1)
		return union(Iₐ, Iᵦ)
	end
end

# a shooting method to find the k-th preimage of a point

#function that gives me the type of the endpoints of the interval

using LinearAlgebra

# this function generates the Jacobian for
# f(x_0)=x_1, f(x_1)=x_2, ..., f(x_{n-1})=y

coeff_interval(x::Array{Interval{T}, 1}) where {T} = T

function Jac(fprime, v::Vector{T}) where {T}
    dv = fprime.(v)
    ev = -ones(T, length(v)-1)
    return Bidiagonal{T}(dv, ev, :U)
end

function ShootingMethod(f, fprime, n, x, y, rigstep = 10)
	F = x->(f.(x)-[x[2:end]; y])

	for i in 1:rigstep
		x_mid = Interval{coeff_interval(x)}.(mid.(x))
		x = intersect.(x, x_mid-Jac(fprime, x)\F(x_mid))
	end
	return x
end

end


# function Phi(u::Float64, v::Float64)
#     BIG = 10000000
#     if v>BIG
#         return (u,1,0)
#     else
#         sq = sqrt(u^2+v^2)
#         p = u+v-sq
#         J_u = 1 - u/sq
#         J_v = 1 - v/sq
#         return (p,J_u,J_v)
# end

# function Phi(u::Point{d},v::Point{d}) where d
#     rr = [Phi(u[i], v[i]) for i=1:d]
#     p = [r[1] for r in rr]
#     J_u = [r[2] for r in rr]
#     J_v = [r[3] for r in rr]
#     return SVector(p...), SDiagonal(J_u...), SDiagonal(J_v...)
# end

function Phi(u::Point{d},v::Point{d}) where d
    sq = sqrt.(u.^2+v.^2)
    p = u+v-sq
    J_u = 1 - u./sq
    J_v = 1 - v./sq
    return p, SDiagonal(J_u), SDiagonal(J_v)
end

function PhiPhi(f::Point{d}, x::Point{d}, a::Point{d}, b::Point{d}) where d
    y, y_f, y_x = Phi(f,x-a)
    z, z_y, z_x = Phi(-y,b-x)
    return z, (-z_y*y_f), (-z_y*y_x - z_x)
end

function PhiPhi0(f::Point{d}, x::Point{d}, a::Point{d}, b::Point{d}) where d
    y = Phi(f,x-a)[1]
    z = Phi(-y,b-x)[1]
    return z
end

function PhiPhi(f::Point{d}, D::SMatrix{d,d,Float64,q}, x::Point{d}, a::Point{d}, b::Point{d}) where d where q
    z, z_f, z_x = PhiPhi(f,x,a,b)
    z, z_f*D + z_x
end

function PhiPhi!(F::Vector{Vector{Point{d}}},
                 X::Vector{Vector{Point{d}}},
                 A::Vector{Vector{Point{d}}},
                 B::Vector{Vector{Point{d}}},
                 D::Vector{v},
                 J::Matrix{w}) where v<:AbstractVector{SMatrix{d,d,Float64,q}} where w<:AbstractVector{SMatrix{d,d,Float64,q}} where d where q

    n_m, n_M = size(J)
    N = length(F[1])
    for i=1:n_m
        for n=1:N
            f = F[i][n]
            x = X[i][n]
            a = A[i][n]
            b = B[i][n]
            z, z_f, z_x = PhiPhi(f,x,a,b)
            F[i][n] = z
            D[i][n] = z_f*D[i][n] + z_x
            for j=1:n_M
                J[i,j][n] = z_f*J[i,j][n]
            end
        end
    end
end

function PhiPhi!(F::Vector{Vector{Point{d}}},X::Vector{Vector{Point{d}}},A::Vector{Vector{Point{d}}},B::Vector{Vector{Point{d}}},D::Vector{v}) where v<:AbstractVector{SMatrix{d,d,Float64,q}} where d where q

    n_m = size(D,1)
    N = length(F[1])
    for i=1:n_m
        for n=1:N
            f = F[i][n]
            x = X[i][n]
            a = A[i][n]
            b = B[i][n]
            z, z_f, z_x = PhiPhi(f,x,a,b)
            F[i][n] = z
            D[i][n] = z_f*D[i][n] + z_x
        end
    end
end

function PhiPhi(F::Vector{Vector{Point{d}}},X::Vector{Vector{Point{d}}},A::Vector{Vector{Point{d}}},B::Vector{Vector{Point{d}}},D::Vector{Vector{SMatrix{d,d,Float64,q}}}, J::Matrix{Vector{SMatrix{d,d,Float64,q}}}) where d where q
    FF = deepcopy(F)
    DD = deepcopy(D)
    JJ = deepcopy(J)
    PhiPhi!(FF,X,A,B,DD,JJ)
    return FF,DD,JJ
end


function solve_polymerflood(Si::T, Sj::T, kr1::RelPerms, kr2::RelPerms, D::T, μw::T, μo::T) where T <: Float64

    f1(s) = fractional_flow(s, kr1, μw, μo)
    df1(s) = fw_derivative(s, k1, μw, μo)
    f2(s) = fractional_flow(s, kr2, μw, μo)
    df2(s) = fw_derivative(s, kr2, μw, μo)

    # Velocity of polymer shock
    vΔc(s) = f2(s) ./ (s + D)

    # * Step 1 - Calculate Ŝ1 and vΔc
    Ŝ = fzero(s -> vΔc.(s) - df2(s), 0.6)
    VΔc = vΔc(Ŝ)

    # * Step2 - Calculate S1B and v2b
    Sb = fzero(s -> (s + D) * vΔc(Ŝ) - f1(s), 0.4)
    Vb = (f1(Sb) - f1(Si)) / (Sb - Si)

    PolymerFlooding(Ŝ, Si, Sj, kr1, kr2, μw, μo, Sb, D, Vb, VΔc)

end

function saturation_profile(pf::PolymerFlooding, time::Real)
    si = pf.Si
    sj = pf.Sj
    sb = pf.Sb
    ŝ = pf.S̃

    df(s) = fw_derivative(s, pf.kr2, pf.μw, pf.μo)

    s = vcat([si, si, sb ,sb, ŝ], collect(ŝ:0.01:sj))
    x = zeros(length(s))
    x[1] = 1.0
    x[2:3] .= pf.Vb * time
    x[4:5] .= pf.VΔc * time
    x[6:end] = df(s[6:end]) * time

    return x, s
end

function polymer_profile(pf::PolymerFlooding, time::Real)
    sj = pf.Sj
    ŝ = pf.S̃

    df(s) = fw_derivative(s, pf.kr2, pf.μw, pf.μo)

    s = collect(ŝ:0.01:sj)
    x = zeros(length(s))
    x = df(s) * time

    return x, s
end


using InteractiveDynamics, GLMakie
using DynamicalSystems

i = 1
p_index = 1

systems = [(Systems.logistic(), 2.5, 4.0, "r", "logistic map"),
           (Systems.henon(), 0.8, 1.4, "a", "Hénon map"),
           (Systems.standardmap(), 0.0, 1.2, "k", "standard map")]

ds, p_min, p_max, parname, t = systems[2]
t = "orbit diagram for the "*t

oddata = interactive_orbitdiagram(ds, p_index, p_min, p_max, i;
                                  parname = parname, title = t)

ps, us = scaleod(oddata)


"""
    dynamicRM(trr::tree, RM::Array{Any,1} )

Implements the nested version of RM for the tree specified by trr.
The user has to specify a vector of risk measures of lenght N, where N is the number of stages of trr.
Furthermore the risk measures RM[i] should be of the form RM[i](states, prob).

"""
function dynamicRM(trr::Tree; RM::Array{Any,1})

    if !(length(RM) == maximum(stage(trr)) )
        println("Error, number of risk measures should be equal to the number of stages")
        return(nothing)
    end
# recursively go through tree and apply RM
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]
        T0.state[i] = RM[i](states,prob)
    end

    return(T0[1],T0)
end

"""
    VaR(trr::Tree,alpha::Float32)
implements the nested version of the Value-at-Risk at level ``\\alpha`` for the Tree ``trr``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version.
"""
function VaR(trr::Tree,alpha::Float32)
    # look for smallest x such that  P(-states <= x) >alpha , i.e.
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.VaR(states,prob,alpha)
    end
    return(T0.state[1],T0)
end

"""
    CTE(trr::Tree,alpha::Float32)
implements the nested version of the Conditional Tail Expectation at level ``\\alpha`` for the Tree ``trr``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version.
"""
function CTE(trr::Tree,alpha::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]
        println(states)
        println(prob)
        T0.state[i] = RiskMeasures.CTE(states,prob,alpha)
        println("risk value:")
        println(T0.state[i])
    end
    return(T0.state[1],T0)
end

"""
    EVaR2(trr::Tree,beta::Float32)
implements the nested version of the Entropic Value-at-Risk at level ``\\beta`` for the Tree ``trr``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version.
"""
function EVaR2(trr::Tree,beta::Float32)
	T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.EVaR2(states,prob,beta)[1]
    end
    return(T0.state[1],T0)
end

"""
    EVaR(trr::Tree,beta::Float32)
implements the nested version of the Entropic Value-at-Risk at level ``\\beta`` for the Tree ``trr``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version. This version uses JuMP and Ipopt.
"""
function EVaR(trr::Tree,beta::Float32)
	T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.EVaR(states,prob,beta)[1]
    end
    return(T0.state[1],T0)
end

"""
    AVaR(trr::Tree,beta::Float32)
implements the nested version of the Average Value-at-Risk at level ``\\alpha`` for the Tree ``trr``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version.
"""
function AVaR(trr::Tree, alpha::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]
        println(states)
        println(prob)

        T0.state[i] = RiskMeasures.AVaR(states,prob,alpha)[1]
    end
    return(T0.state[1],T0)
end

#########################################
# From basicRM.jl
"""
    Expectation(states,prob)

Computes the expectation for the random variable ``Y`` with values given by `states` under
the probability measure ``Q`` given by `prob`:

```math
\\mathbb{E} Y = states \\cdot prob
```
"""
function Expectation(trr::Tree, alpha::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.Expectation(states,prob)
    end
    return(T0.state[1],T0)
end

"""
    entropic(states, prob,theta::Float32)

implements the entropic risk measure defined as
```math
\\rho_\\theta(Y) = \\theta \\cdot \\log \\mathbb{E} e^{\\frac{Y}{\\theta}},
```
where ``\\theta`` is greater 0 and ``Y`` is the random variable defined by `states` and `prob`.
"""
function entropic(trr::Tree, theta::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.entropic(states,prob,theta)
    end
    return(T0.state[1],T0)
end

#########################################################
#Mean Semi-deviation
#########################################################
"""
    mSD(states,prob,beta::Float32,p::Float32)

implements the mean semi-deviation of order `p` which is a coherent risk measure defined by
```math
mSD_\\beta^p (Y) = \\mathbb{E} Y  + \\beta \\lvert \\left( Y - \\mathbb{E}Y \\right)_+ \\rvert_p,
```
for the random variable ``Y`` defined by `states` and `prob`.
"""
function mSD(trr::Tree, beta::Float32,p::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures. mSD(states,prob,beta::Float32,p::Float32)
    end
    return(T0.state[1],T0)
end

"""
    meanVariance(states, prob,c::Float32)

implements the mean Variance risk measure defined by
```math
\\rho_c(Y) := \\mathbb{E}Y + c \\cdot \\mathbb{E} \\left( Y- \\mathbb{E}Y)^2 \\right),
```
where ``c >0`` and ``Y`` is the random variable defined by `states` and `prob`.
"""
function meanVariance(trr::Tree, c::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.meanVariance(states, prob,c::Float32)
    end
    return(T0.state[1],T0)
end

"""
    meanDeviation(states, prob,c::Float32,p::Float32)

implements the mean deviation risk measure of order ``p\\geq 1`` defined by
```math
\\rho_c^p(Y) := \\mathbb{E}Y + c \\cdot \\lVert \\left( Y- \\mathbb{E}Y \\right)^2 \\rVert_p,
```
where ``c >0`` and ``Y`` is the random variable defined by `states` and `prob`.
"""
function meanDeviation(trr::Tree, c::Float32,p::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.meanDeviation(states, prob,c::Float32,p::Float32)
    end
    return(T0.state[1],T0)
end

"""
    meanSemiVariance(states, prob,c::Float32,t::Float32)

implements the mean upper-semi variance risk measure from a target ``t`` defined by
```math
\\rho_{c,t}(Y) = \\mathbb{E}Y + c \\cdot \\mathbb{E} \\left( Y - t \\right)^2_+ ,
```
where ``c >0`` and ``Y`` is the random variable defined by `states` and `prob`.
"""
function meanSemiVariance(trr::Tree, c::Float32,t::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.meanSemiVariance(states, prob,c::Float32,t::Float32)
    end
    return(T0.state[1],T0)
end

"""
    meanSemiDevi(states, prob,c::Float32,target::Float32,p::Float32)

implements the mean upper-semi variance risk measure of order ``p \\geq 1`` from a target ``t`` defined by
```math
\\rho_{c,t}^p(Y) = \\mathbb{E}Y + c \\cdot \\lVert \\left( Y - t \\right)_+ \\rVert_p ,
```
where ``c >0`` and ``Y`` is the random variable defined by `states` and `prob`.
"""
function meanSemiDevi(trr::Tree, c::Float32,target::Float32,p::Float32)
    T0 = deepcopy(trr);

    # For every parent node going backwards
    tmp = reverse(unique(T0.parent)); pop!(tmp)
    for i in tmp
        #Note that the probs are the conditional ones
        #Risk measure w/ adjustment of values
        states = T0.state[T0.children[i+1]]
        prob = T0.probability[T0.children[i+1]]

        T0.state[i] = RiskMeasures.meanSemiDevi(states, prob,c::Float32,target::Float32,p::Float32)
    end
    return(T0.state[1],T0)
end


###############################################################################
#
#   GenericTypes.jl : Generic Types
#
###############################################################################

###############################################################################
#
#   PermGroup / perm
#
###############################################################################

const PermID = ObjectIdDict()

mutable struct PermGroup <: Nemo.Group
   n::Int

   function PermGroup(n::Int, cached=true)
      if haskey(PermID, n)
         return PermID[n]::PermGroup
      else
         z = new(n)
         if cached
            PermID[n] = z
         end
         return z
      end
   end
end

mutable struct perm <: Nemo.GroupElem
   d::Array{Int, 1}
   cycles::Vector{Vector{Int}}
   parent::PermGroup

   function perm(n::Int)
      return new(collect(1:n))
   end

   function perm(a::Array{Int, 1})
      return new(a)
   end
end

doc"""
    AllPerms(n::Int)
> Returns an iterator over arrays representing all permutations of `1:n`.
> Similar to `Combinatorics.permutations(1:n)`
"""
struct AllPerms
   n::Int
   all::Int
   AllPerms(n::Int) = new(n, factorial(n))
end

###############################################################################
#
#   Partition
#
###############################################################################

doc"""
    Partition(part::Vector{Int}, check::Bool=true)
> Partition represents integer partition into numbers in non-increasing order.
> It is a thin wrapper over `Vector{Int}`
"""
mutable struct Partition <: AbstractVector{Int}
   n::Int
   part::Vector{Int}

   function Partition(part::Vector{Int}, check::Bool=true)
      if check
         all(diff(part) .<= 0) || throw("Partition must be decreasing!")
         if length(part) > 0
            part[end] >=1 || throw("Found non-positive entry in partition!")
         end
      end
      return new(sum(part), part)
   end
end

doc"""
   Partitions(n::Int)
> Returns an iterator over all integer `Partition`s of `n`. They come in
> ascending order. See also `Combinatorics.partitions(n)`.
"""
struct Partitions
    n::Int
end

###############################################################################
#
#   SkewDiagram
#
###############################################################################

doc"""
    SkewDiagram(lambda::Partition, mu::Partition)
> Implements a skew diagram, i.e. a difference of two Young diagrams
> represented by partitions `lambda` and `mu`.
"""
struct SkewDiagram
   lam::Partition
   mu::Partition

   function SkewDiagram(lambda, mu)
      lambda.n >= mu.n || throw("Can't create SkewDiagram: $mu is partition of  $(mu.n) > $(lambda.n).")
      length(lambda) >= length(mu) || throw("Can't create SkewDiagram: $mu is longer than $(lambda)!")
      for (l, m) in zip(lambda, mu)
         l >= m || throw("a row of $mu is longer than a row of $lambda")
      end
      return new(lambda, mu)
   end
end

###############################################################################
#
#   Young Tableau
#
###############################################################################

doc"""
    YoungTableau(part::Partition, tab::Array{Int, 2})
> Returns the Young tableaux of partition `part` of `n`, filled linearly
> (row-major) by `fill` vector.
"""
struct YoungTableau <: AbstractArray{Int, 2}
   part::Partition
   tab::Array{Int,2}
end

###############################################################################
#
#   PolyRing / Poly
#
###############################################################################

mutable struct PolyRing{T <: RingElement} <: Nemo.PolyRing{T}
   base_ring::Ring
   S::Symbol

   function PolyRing{T}(R::Ring, s::Symbol, cached::Bool = true) where T <: RingElement
      if haskey(PolyID, (R, s))
         return PolyID[R, s]::PolyRing{T}
      else
         z = new{T}(R, s)
         if cached
           PolyID[R, s] = z
         end
         return z
      end
   end
end

const PolyID = Dict{Tuple{Ring, Symbol}, Ring}()

mutable struct Poly{T <: RingElement} <: Nemo.PolyElem{T}
   coeffs::Array{T, 1}
   length::Int
   parent::PolyRing{T}

   Poly{T}() where T <: RingElement = new{T}(Array{T}(0), 0)

   function Poly{T}(b::Array{T, 1}) where T <: RingElement
      z = new{T}(b)
      z.length = normalise(z, length(b))
      return z
   end

   Poly{T}(a::T) where T <: RingElement = iszero(a) ? new{T}(Array{T}(0), 0) : new{T}([a], 1)
end

###############################################################################
#
#   MPolyRing / MPoly
#
###############################################################################

# S is a Symbol which can take the values:
# :lex
# :revlex
# :deglex
# :degrevlex
#
# T is an Int which is the number of variables
# (plus one if ordered by total degree)

mutable struct MPolyRing{T <: RingElement} <: Nemo.MPolyRing{T}
   base_ring::Ring
   S::Array{Symbol, 1}
   ord::Symbol
   num_vars::Int
   N::Int

   function MPolyRing{T}(R::Ring, s::Array{Symbol, 1}, ord::Symbol, N::Int,
                         cached::Bool = true) where T <: RingElement
      if haskey(MPolyID, (R, s, ord, N))
         return MPolyID[R, s, ord, N]::MPolyRing{T}
      else
         z = new{T}(R, s, ord, length(s), N)
         if cached
           MPolyID[R, s, ord, N] = z
         end
         return z
      end
   end
end

const MPolyID = Dict{Tuple{Ring, Array{Symbol, 1}, Symbol, Int}, Ring}()

mutable struct MPoly{T <: RingElement} <: Nemo.MPolyElem{T}
   coeffs::Array{T, 1}
   exps::Array{UInt, 2}
   length::Int
   parent::MPolyRing{T}

   function MPoly{T}(R::MPolyRing) where T <: RingElement
      N = R.N
      return new{T}(Array{T}(0), Array{UInt}(N, 0), 0, R)
   end

   MPoly{T}(R::MPolyRing, a::Array{T, 1}, b::Array{UInt, 2}) where T <: RingElement = new{T}(a, b, length(a), R)

   function MPoly{T}(R::MPolyRing, a::T) where T <: RingElement
      N = R.N
      return iszero(a) ? new{T}(Array{T}(0), Array{UInt}(N, 0), 0, R) :
                                          new{T}([a], zeros(UInt, N, 1), 1, R)
   end
end

###############################################################################
#
#   SparsePolyRing / SparsePoly
#
###############################################################################

mutable struct SparsePolyRing{T <: RingElement} <: Nemo.Ring
   base_ring::Ring
   S::Symbol
   num_vars::Int

   function SparsePolyRing{T}(R::Ring, s::Symbol, cached::Bool = true) where T <: RingElement
      if haskey(SparsePolyID, (R, s))
         return SparsePolyID[R, s]::SparsePolyRing{T}
      else
         z = new{T}(R, s)
         if cached
           SparsePolyID[R, s] = z
         end
         return z
      end
   end
end

const SparsePolyID = Dict{Tuple{Ring, Symbol}, SparsePolyRing}()

mutable struct SparsePoly{T <: RingElement} <: Nemo.RingElem
   coeffs::Array{T, 1}
   exps::Array{UInt}
   length::Int
   parent::SparsePolyRing{T}

   SparsePoly{T}() where T <: RingElement = new{T}(Array{T}(0), Array{UInt}(0), 0)

   SparsePoly{T}(a::Array{T, 1}, b::Array{UInt, 1}) where T <: RingElement = new{T}(a, b, length(a))

   SparsePoly{T}(a::T) where T <: RingElement = iszero(a) ? new{T}(Array{T}(0), Array{UInt}(0), 0) :
                                               new{T}([a], [UInt(0)], 1)
end

###############################################################################
#
#   ResRing / Res
#
###############################################################################

mutable struct ResRing{T <: RingElement} <: Nemo.ResRing{T}
   base_ring::Ring
   modulus::T

   function ResRing{T}(modulus::T, cached::Bool = true) where T <: RingElement
      c = canonical_unit(modulus)
      if !isone(c)
        modulus = divexact(modulus, c)
      end
      if haskey(ModulusDict, (parent(modulus), modulus))
         return ModulusDict[parent(modulus), modulus]::ResRing{T}
      else
         z = new{T}(parent(modulus), modulus)
         if cached
            ModulusDict[parent(modulus), modulus] = z
         end
         return z
      end
   end
end

const ModulusDict = Dict{Tuple{Ring, RingElement}, Ring}()

mutable struct Res{T <: RingElement} <: Nemo.ResElem{T}
   data::T
   parent::ResRing{T}

   Res{T}(a::T) where T <: RingElement = new{T}(a)
end

###############################################################################
#
#   RelSeriesRing / RelSeries
#
###############################################################################

mutable struct RelSeriesRing{T <: RingElement} <: Nemo.SeriesRing{T}
   base_ring::Ring
   prec_max::Int
   S::Symbol

   function RelSeriesRing{T}(R::Ring, prec::Int, s::Symbol, cached::Bool = true) where T <: RingElement
      if haskey(RelSeriesID, (R, prec, s))
         return RelSeriesID[R, prec, s]::RelSeriesRing{T}
      else
         z = new{T}(R, prec, s)
         if cached
            RelSeriesID[R, prec, s] = z
         end
         return z
      end
   end
end

const RelSeriesID = Dict{Tuple{Ring, Int, Symbol}, Ring}()

mutable struct RelSeries{T <: RingElement} <: Nemo.RelSeriesElem{T}
   coeffs::Array{T, 1}
   length::Int
   prec::Int
   val::Int
   parent::RelSeriesRing{T}

   function RelSeries{T}(a::Array{T, 1}, length::Int, prec::Int, val::Int) where T <: RingElement
      new{T}(a, length, prec, val)
   end

   RelSeries{T}(a::RelSeries{T}) where T <: RingElement = a
end

###############################################################################
#
#   AbsSeriesRing / AbsSeries
#
###############################################################################

mutable struct AbsSeriesRing{T <: RingElement} <: Nemo.SeriesRing{T}
   base_ring::Ring
   prec_max::Int
   S::Symbol

   function AbsSeriesRing{T}(R::Ring, prec::Int, s::Symbol, cached::Bool = true) where T <: RingElement
      if haskey(AbsSeriesID, (R, prec, s))
         return AbsSeriesID[R, prec, s]::AbsSeriesRing{T}
      else
         z = new{T}(R, prec, s)
         if cached
            AbsSeriesID[R, prec, s] = z
         end
         return z
      end
   end
end

const AbsSeriesID = Dict{Tuple{Ring, Int, Symbol}, Ring}()

mutable struct AbsSeries{T <: RingElement} <: Nemo.AbsSeriesElem{T}
   coeffs::Array{T, 1}
   length::Int
   prec::Int
   parent::AbsSeriesRing{T}

   AbsSeries{T}(a::Array{T, 1}, length::Int, prec::Int) where T <: RingElement = new{T}(a, length, prec)
   AbsSeries{T}(a::AbsSeries{T}) where T <: RingElement = a
end

###############################################################################
#
#   LaurentSeriesRing / LarentSeriesRingElem
#
###############################################################################

mutable struct LaurentSeriesRing{T <: RingElement} <: Nemo.Ring
   base_ring::Ring
   prec_max::Int
   S::Symbol

   function LaurentSeriesRing{T}(R::Ring, prec::Int, s::Symbol, cached::Bool = true) where T <: RingElement
      if haskey(LaurentSeriesID, (R, prec, s))
         return LaurentSeriesID[R, prec, s]::LaurentSeriesRing{T}
      else
         z = new{T}(R, prec, s)
         if cached
            LaurentSeriesID[R, prec, s] = z
         end
         return z
      end
   end
end

const LaurentSeriesID = Dict{Tuple{Ring, Int, Symbol}, Ring}()

mutable struct LaurentSeriesRingElem{T <: RingElement} <: Nemo.RingElem
   coeffs::Array{T, 1}
   length::Int
   prec::Int
   val::Int
   parent::LaurentSeriesRing{T}

   function LaurentSeriesRingElem{T}(a::Array{T, 1}, length::Int, prec::Int, val::Int) where T <: RingElement
      new{T}(a, length, prec, val)
   end

   LaurentSeriesRingElem{T}(a::LaurentSeriesRingElem{T}) where T <: RingElement = a
end

###############################################################################
#
#   LaurentSeriesField / LarentSeriesFieldElem
#
###############################################################################

mutable struct LaurentSeriesField{T <: FieldElement} <: Nemo.Field
   base_ring::Field
   prec_max::Int
   S::Symbol

   function LaurentSeriesField{T}(R::Field, prec::Int, s::Symbol, cached::Bool = true) where T <: FieldElement
      if haskey(LaurentSeriesID, (R, prec, s))
         return LaurentSeriesID[R, prec, s]::LaurentSeriesField{T}
      else
         z = new{T}(R, prec, s)
         if cached
            LaurentSeriesID[R, prec, s] = z
         end
         return z
      end
   end
end

mutable struct LaurentSeriesFieldElem{T <: FieldElement} <: Nemo.FieldElem
   coeffs::Array{T, 1}
   length::Int
   prec::Int
   val::Int
   parent::LaurentSeriesField{T}

   function LaurentSeriesFieldElem{T}(a::Array{T, 1}, length::Int, prec::Int, val::Int) where T <: FieldElement
      new{T}(a, length, prec, val)
   end

   LaurentSeriesFieldElem{T}(a::LaurentSeriesFieldElem{T}) where T <: FieldElement = a
end

const LaurentSeriesElem{T} = Union{LaurentSeriesRingElem{T}, LaurentSeriesFieldElem{T}} where T <: RingElement

###############################################################################
#
#   FracField / Frac
#
###############################################################################

mutable struct FracField{T <: RingElem} <: Nemo.FracField{T}
   base_ring::Ring

   function FracField{T}(R::Ring, cached::Bool = true) where T <: RingElem
      if haskey(FracDict, R)
         return FracDict[R]::FracField{T}
      else
         z = new{T}(R)
         if cached
            FracDict[R] = z
         end
         return z
      end
   end
end

const FracDict = Dict{Ring, Ring}()

mutable struct Frac{T <: RingElem} <: Nemo.FracElem{T}
   num::T
   den::T
   parent::FracField{T}

   Frac{T}(num::T, den::T) where T <: RingElem = new{T}(num, den)
end

###############################################################################
#
#   MatSpace / Mat
#
###############################################################################

# not really a mathematical ring
mutable struct MatSpace{T <: RingElement} <: Nemo.MatSpace{T}
   rows::Int
   cols::Int
   base_ring::Ring

   function MatSpace{T}(R::Ring, r::Int, c::Int, cached::Bool = true) where T <: RingElement
      if haskey(MatDict, (R, r, c))
         return MatDict[R, r, c]::MatSpace{T}
      else
         z = new{T}(r, c, R)
         if cached
            MatDict[R, r, c] = z
         end
         return z
      end
   end
end

const MatDict = Dict{Tuple{Ring, Int, Int}, Ring}()

mutable struct Mat{T <: RingElement} <: Nemo.MatElem{T}
   entries::Array{T, 2}
   base_ring::Ring

   function Mat{T}(A::Array{T, 2}) where T <: RingElement
      return new{T}(A)
   end

   function Mat{T}(r::Int, c::Int, A::Array{T, 1}) where T <: RingElement
      t = Array{T}(r, c)
      for i = 1:r
         for j = 1:c
            t[i, j] = A[(i - 1) * c + j]
         end
      end
      return new{T}(t)
   end
end


using DiffEqBase, DiffEqBase.InternalEuler

# Try it
ff(t,u) = u
ff(::Val{:jac}, t, u) = 1
ff(::Type{Val{:analytic}},t,u0) = u0*exp(t)

dt = 0.01
prob = ODEProblem(ff, 1.0, (0.0,1.0) )
sol = solve(prob, InternalEuler.FwdEulerAlg(), tstops=0:dt:1)
sol2 = solve(prob, InternalEuler.BwdEulerAlg(), tstops=0:dt:1)

#using Plots
#plot(sol)
#plot!(sol2)
#plot!(sol2, plot_analytic=true)

#using DiffEqDevTools
#dts = 1./2.^(8:-1:4)
#sim = test_convergence(dts,p2,BwdEulerAlg())
#@show sim.𝒪est[:final]
#plot(sim)


# make your scripts automatically re-activate your project
cd(@__DIR__)
using Pkg; Pkg.activate("../..")

# load packages
using ml_for_pt
using DelimitedFiles
using LaTeXStrings
using Plots
using Random
using Flux
ENV["GKSwstype"]="nul"

# set linear system size and corresponding data folder
L = 10
# L=60
data_folder = "../../data/ising/L="*string(L)*"/"

# set path to save folder
save_folder = "./results/L="*string(L)*"/"

# critical parameter value for the Ising model
p_crit = Float32(2/log(1+sqrt(2)))

# define parameter ranges
p_min = 0.05f0
p_max = 10.0f0
dp = 0.05f0
p_range = collect(p_min:dp:p_max)
p_range_LBC = collect(p_min-dp/2:dp:p_max+dp/2)

# import Monte Carlo data

# unique energies at p=p_1
energy = readdlm(data_folder*"E_T="*string(0)*".txt", eltype(dp))

# number of times each energy is drawn at T=T_1
num = readdlm(data_folder*"num_T="*string(0)*".txt", eltype(dp))

# iterate over entire parameter range
energies = [energy]
numbers = [num]
for index in 1:length(p_range)-1
  energyy = readdlm(data_folder*"E_T="*string(index)*".txt", eltype(dp))
  numm = readdlm(data_folder*"num_T="*string(index)*".txt", eltype(dp))
  push!(energies, energyy)
  push!(numbers, numm)
end
unique_energies = unique(vcat(energies...))
n_samples = sum(numbers[1])

# construct probability distribution from Monte Carlo data
# distr(x,p) gives probability to draw the sample x at parameter value p
# samples correspond to enumeration of unique energies with integers
distr,samples = MLP.distr_approx(energies, unique_energies, numbers, p_range)

# optionally run methods using the analytical expressions to compare with neural network results

# supervised learning using analytical expression
# returns optimal predictions, indicator, and loss
pred_opt_SL, indicator_opt_SL, loss_opt_SL = MLP.get_indicators_SL_analytical(samples, distr, p_range, dp, p_min, p_max)

# prediction-based method using analytical expression
# returns optimal predictions, indicator, and loss
pred_opt_PBM, indicator_opt_PBM, loss_opt_PBM = MLP.get_indicators_PBM_analytical(samples, distr, p_range, dp)

# learning by confusion using analytical expression
# returns optimal indicator and loss
indicator_opt_LBC, loss_opt_LBC = MLP.get_indicators_LBC_analytical(samples, distr, p_range, p_range_LBC)

# construct dataset for training neural networks
dataset = MLP.get_dataset(p_range, distr, samples)
dataset_train_SL = MLP.get_modified_dataset_train_SL(distr, samples, p_range, n_samples)

# standardize inputs
inputs_one_hot = reshape(unique_energies, 1, length(unique_energies))
mean_train,std_train = MLP.get_dataset_stats(dataset, inputs_one_hot, length(p_range))
inputs_one_hot_stand = (inputs_one_hot.-mean_train)./std_train
mean_train_SL,std_train_SL = MLP.get_dataset_stats(dataset_train_SL, inputs_one_hot, length(p_range))
inputs_one_hot_stand_SL = (inputs_one_hot.-mean_train_SL)./std_train_SL

# supervised learning using neural networks

# initialize neural network
n_nodes_SL = [size(inputs_one_hot)[1], 64, 64, 64, 2]
NN_SL = Chain(
  Dense(n_nodes_SL[1], n_nodes_SL[2], relu),
  Dense(n_nodes_SL[2], n_nodes_SL[3], relu),
  Dense(n_nodes_SL[3], n_nodes_SL[4], relu),
  Dense(n_nodes_SL[4], n_nodes_SL[end], identity))
pnn_SL, re_SL = Flux.destructure(NN_SL)

# set hyperparameters
lr_SL = 0.001f0
epochs_SL = 10000
saveat_SL = 100
opt_SL = ADAM(lr_SL)
verbose = true

# train neural network
# returns predictions, indicators, and loss saved at epochs specified by saveat_SL variable
pred_logger_SL, losses_SL, NN_logger_SL = MLP.get_indicators_SL_numerical(pnn_SL,  re_SL, dataset_train_SL, dataset, epochs_SL, p_range, dp, p_min, p_max, opt_SL, inputs_one_hot_stand, verbose=verbose, saveat=saveat_SL)

# extract results at last save point
pred_NN_SL = pred_logger_SL[end][1]
indicator_NN_SL = pred_logger_SL[end][2]
loss_NN_SL = pred_logger_SL[end][3]

# save results of supervised learning in save_folder
open(save_folder*"hyperparams_SL.txt", "w") do io
  writedlm(io, [saveat_SL, lr_SL, epochs_SL])
end

open(save_folder*"NN_logger_SL.txt", "w") do io
  writedlm(io, NN_logger_SL)
end

for i in 1:length(pred_logger_SL)
  open(save_folder*"pred_logger_SL_"*string(i)*".txt", "w") do io
  writedlm(io, pred_logger_SL[i][1])
  end

  open(save_folder*"indicator_logger_SL_"*string(i)*".txt", "w") do io
  writedlm(io, pred_logger_SL[i][2])
  end

  open(save_folder*"loss_logger_SL_"*string(i)*".txt", "w") do io
  writedlm(io, pred_logger_SL[i][3])
  end

end

open(save_folder*"pred_NN_SL.txt", "w") do io
  writedlm(io, pred_NN_SL)
end

open(save_folder*"indicator_NN_SL.txt", "w") do io
  writedlm(io, indicator_NN_SL)
end

open(save_folder*"loss_NN_SL.txt", "w") do io
  writedlm(io, loss_NN_SL)
end

open(save_folder*"NN_structure_SL.txt", "w") do io
  writedlm(io, n_nodes_SL)
end

open(save_folder*"losses_SL.txt", "w") do io
  writedlm(io, losses_SL)
end

# plot and save results of supervised learning in save_folder
plt = plot(p_range, pred_NN_SL, dpi=300, c="black", label="NN", ylims=(0.0,1.0))
plot!(p_range, pred_opt_SL, dpi=300, c="blue", label="analytical", ylims=(0.0,1.0))
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\hat{y}_{\mathrm{SL}}$")
savefig(save_folder*"pred_NN_SL.png")

plt = plot(p_range[2:end-1], indicator_NN_SL, dpi=300, c="black", label="NN")
plot!(p_range[2:end-1], indicator_opt_SL, dpi=300, c="blue", label="analytical")
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I_{\mathrm{SL}}$")
savefig(save_folder*"indicator_NN_SL.png")

plt = plot(collect(1:length(losses_SL)), losses_SL, dpi=300, c="black", label="NN")
hline!([loss_opt_SL], label="analytical", c="blue")
xlabel!("Epochs")
ylabel!(L"$\mathcal{L}_{\mathrm{SL}}$")
savefig(save_folder*"loss_NN_SL.png")

# prediction-based method using neural networks

# initialize neural network
n_nodes_PBM = [size(inputs_one_hot)[1], 64, 64, 64, 1]
NN_PBM = Chain(
  Dense(n_nodes_PBM[1], n_nodes_PBM[2], relu),
  Dense(n_nodes_PBM[2], n_nodes_PBM[3], relu),
  Dense(n_nodes_PBM[3], n_nodes_PBM[4], relu),
  Dense(n_nodes_PBM[4], n_nodes_PBM[end], identity))
pnn_PBM, re_PBM = Flux.destructure(NN_PBM)

# set hyperparameters
lr_PBM = 0.001f0
epochs_PBM = 1000
saveat_PBM = 100
opt_PBM = ADAM(lr_PBM)
verbose = true

# train neural network
# returns predictions, indicators, and loss saved at epochs specified by saveat_PBM variable
pred_logger_PBM, losses_PBM, NN_logger_PBM =MLP.get_indicators_PBM_numerical(pnn_PBM, re_PBM, dataset, epochs_PBM, p_range, dp, opt_PBM, inputs_one_hot_stand, verbose=verbose, saveat=saveat_PBM)

# extract results at last save point
pred_NN_PBM = pred_logger_PBM[end][1]
indicator_NN_PBM = pred_logger_PBM[end][2]
loss_NN_PBM = pred_logger_PBM[end][3]

# save results of prediction-based method in save_folder
open(save_folder*"hyperparams_PBM.txt", "w") do io
  writedlm(io, [saveat_PBM, lr_PBM, epochs_PBM])
end

open(save_folder*"NN_logger_PBM.txt", "w") do io
  writedlm(io, NN_logger_PBM)
end

for i in 1:length(pred_logger_PBM)
  open(save_folder*"pred_logger_PBM_"*string(i)*".txt", "w") do io
  writedlm(io, pred_logger_PBM[i][1])
  end

  open(save_folder*"indicator_logger_PBM_"*string(i)*".txt", "w") do io
  writedlm(io, pred_logger_PBM[i][2])
  end

  open(save_folder*"loss_logger_PBM_"*string(i)*".txt", "w") do io
  writedlm(io, pred_logger_PBM[i][3])
  end

end

open(save_folder*"pred_NN_PBM.txt", "w") do io
  writedlm(io, pred_NN_PBM)
end

open(save_folder*"indicator_NN_PBM.txt", "w") do io
  writedlm(io, indicator_NN_PBM)
end

open(save_folder*"loss_NN_PBM.txt", "w") do io
  writedlm(io, loss_NN_PBM)
end

open(save_folder*"NN_structure_PBM.txt", "w") do io
  writedlm(io, n_nodes_PBM)
end

open(save_folder*"losses_PBM.txt", "w") do io
  writedlm(io, losses_PBM)
end

# plot and save results of prediction-based method in save_folder
plt = plot(p_range, pred_NN_PBM, dpi=300, c="black", label="NN", ylims=(p_range[1],p_range[end]))
plot!(p_range, pred_opt_PBM, dpi=300, c="blue", label="analytical", ylims=(p_range[1],p_range[end]))
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\hat{y}_{\mathrm{PBM}}$")
savefig(save_folder*"pred_NN_PBM.png")

plt = plot(p_range[2:end-1], indicator_NN_PBM, dpi=300, c="black", label="NN")
plot!(p_range[2:end-1], indicator_opt_PBM, dpi=300, c="blue", label="analytical")
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I^{\mathrm{NN}}_{\mathrm{PBM}}$")
savefig(save_folder*"indicator_NN_PBM.png")

plt = plot(collect(1:length(losses_PBM)), losses_PBM, dpi=300, c="black", label="NN")
hline!([loss_opt_PBM], label="analytical", c="blue")
xlabel!("Epochs")
ylabel!(L"$\mathcal{L}_{\mathrm{PBM}}$")
savefig(save_folder*"loss_NN_PBM.png")

# learning by confusion using neural networks

# initialize neural network
n_nodes_LBC = [size(inputs_one_hot)[1], 64, 64, 64, 2]
NN_LBC = Chain(
  Dense(n_nodes_LBC[1], n_nodes_LBC[2], relu),
  Dense(n_nodes_LBC[2], n_nodes_LBC[3], relu),
  Dense(n_nodes_LBC[3], n_nodes_LBC[4], relu),
  Dense(n_nodes_LBC[4], n_nodes_LBC[end], identity))
pnn_LBC, re_LBC = Flux.destructure(NN_LBC)

# set hyperparameters
lr_LBC = 0.001f0
epochs_LBC = 100
saveat_LBC = epochs_LBC
opt_LBC = ADAM(lr_LBC)
verbose = true

indicator_NN_LBC = []
loss_NN_LBC = []
for indx in collect(1:length(p_range_LBC))
  pred_logger_LBC, losses_LBC, NN_logger_LBC = MLP.get_indicators_LBC_numerical_fixed_p(deepcopy(pnn_LBC), re_LBC, dataset, epochs_LBC, p_range, dp, opt_LBC, p_range_LBC, indx, inputs_one_hot_stand, saveat=saveat_LBC, verbose=verbose)

  indicator_NN_LBC_p = pred_logger_LBC[end][1][1]
  push!(indicator_NN_LBC, indicator_NN_LBC_p)
  loss_NN_LBC_p = pred_logger_LBC[end][2][1]
  push!(loss_NN_LBC, loss_NN_LBC_p)
end

# save results of learning by confusion in save_folder
open(save_folder*"hyperparams_LBC.txt", "w") do io
  writedlm(io, [saveat_LBC, lr_LBC, epochs_LBC])
end

open(save_folder*"indicator_NN_LBC.txt", "w") do io
  writedlm(io, indicator_NN_LBC)
end

open(save_folder*"loss_NN_LBC.txt", "w") do io
  writedlm(io, loss_NN_LBC)
end

open(save_folder*"NN_structure_LBC.txt", "w") do io
  writedlm(io, n_nodes_LBC)
end

# plot and save results of learning by confusion in save_folder
plt = plot(p_range_LBC, indicator_NN_LBC, dpi=300, c="black", label="NN")
plot!(p_range_LBC, indicator_opt_LBC, dpi=300, c="blue", label="analytical")
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I_{\mathrm{LBC}}$")
savefig(save_folder*"indicator_NN_LBC.png")

plt = plot(p_range_LBC, loss_NN_LBC, dpi=300, c="black", label="NN")
plot!(p_range_LBC, loss_opt_LBC, dpi=300, c="blue", label="analytical")
vline!([p_crit], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\mathcal{L}_{\mathrm{LBC}}$")
savefig(save_folder*"loss_NN_LBC.png")