mutable struct Edge
    rev::Int
    from::Int
    to::Int
    cap::Int
end

function push!(G::Vector{Vector{Edge}}, from, to, cap)
    fromrev = length(G[from]) + 1
    torev = length(G[to]) + 1
    Base.push!(G[from], Edge(torev, from, to, cap))
    Base.push!(G[to], Edge(fromrev, to, from, 0))
end

function run_flow!(G::Vector{Vector{Edge}}, e::Edge, f)
    e.cap -= f
    G[e.to][e.rev].cap += f
    return G
end

function fodfs(G, v, t, f, seens)
    v == t && return f

    seens[v] = true
    for e ∈ G[v]
        if seens[e.to] || e.cap == 0
            # Skip
        else
            flow = fodfs(G, e.to, t, min(f, e.cap), seens)
            if flow == 0
                # Skip
            else
                run_flow!(G, e, flow)
                return flow
            end
        end
    end

    return 0
end

function FordFulkerson(G, s, t)
    INF = 10^9
    res = 0

    while true
        flow = fodfs(G, s, t, INF, falses(length(G)))
        if flow == 0
            return res
        else
            res += flow
        end
    end

    return 0
end

function main()
    N, G, E = parse.(Int, split(readline()))
    ps = parse.(Int, split(readline())) .+ 1

    H = [Vector{Edge}() for _ = 1:N+1] # Count the super node
    for _ = 1:E
        a, b = parse.(Int, split(readline())) .+ 1
        push!(H, a, b, 1)
        push!(H, b, a, 1)
    end

    for p in ps
        push!(H, p, N+1, 1)
        push!(H, N+1, p, 1)
    end

    println(FordFulkerson(H, 1, N+1))
end

main()


using Distributed
using StatsPlots

addprocs(4)

@everywhere using CSV
@everywhere using DataFrames
@everywhere using StatsFuns: logistic
@everywhere using Turing

"""
    get_data(path, :y_label)

Returns data for the logistic model split by X, y structure.
Retruns included breast cancer data by default.

# Arguments:
- `path::String="data/breast_cancer.csv"` path to csv file.
- `y_label::Symbol=:diagnosis` y label for target/dependent variable

# Returns:
- `X` Matrix containing X data/exogenous variables.
- `y` Vector containing y label data/endogenous variables.
"""
@everywhere function get_data(
        path::String="data/breast_cancer.csv",
        y_label::Symbol=:diagnosis
    )
    df = DataFrame(CSV.File(path))
    y = df[!, y_label]
    X = convert(Matrix, select!(df, Not(y_label)))
    X = (X .- mean(X, dims=1)) ./ std(X, dims=1)
    return X, y
end

"""
    logistic_regression(x::Matrix{Float64}, y::Vector{Int64})

Creates a logistic model to be used as an input argument in a sampler.

# Arguments:
- `X::Matrix{Float64}`: Exogenous input data X in matrix format.
- `y::Vector{Int64}`: Endogenous/dependent variables vector.
"""
@everywhere @model function logistic_regression(
        X::Matrix{Float64},
        y::Vector{Int64}
    )
    intercept ~ Normal(0, sqrt(25))

    nfeatures = size(X, 2)
    n = size(X, 1)
    slopes ~ MvNormal(nfeatures, sqrt(25))
    p = intercept .+ X * slopes
    for i = 1:n

        y[i] ~ Bernoulli(logistic(p[i]))

    end
end

"""
    profiler(X::Matrix{Float64}, y::Vector{Int64}, max_iters::Int64=10)

# Arguments:
- `X::Matrix{Float64}`: Exogenous variables
- `y::Matrix{Int64}`: Endogenous variables
- `max_iters::Int64`: The specified number of iterations to execute the sampler for logistic regression.

# Returns:
- `times::Vector{Any}`: Execution times for each iteration.
- `trace::Chains{Any}`: The trace generated in the final sample.
"""
@everywhere function profiler(
        X::Matrix{Float64},
        y::Vector{Int64},
        max_iters::Int64=10
    )
    times = []
    for iteration in 1:max_iters

        time = @elapsed global trace = sample(
                logistic_regression(X, y),
                NUTS(1000, 0.90),
                MCMCDistributed(),
                1000,
                4
            )
        append!(times, time)
    end
    return times, trace
end

@everywhere function main()
    X, y = get_data()
    times, traces = profiler(X, y)
    println(times)
    return times, trace
end

times, trace = main()


###################################################
## logarithmic / discrete return transformations ##
###################################################

function disc2log(tn::AbstractTimenum; percent = false)
    ## discrete net return to logarithmic return
    if percent
        rets = log(tn./100 .+ 1).*100
    else
        rets = log(tn .+ 1)
    end
    return rets
end

function log2disc(tn::AbstractTimenum; percent = false)
    ## logarithmic return to discrete net return
    if percent
        rets = (exp(tn./100) .- 1).*100
    else
        rets = exp(tn) .- 1
    end
    return rets
end

###############
## price2ret ##
###############

## price2ret for DataArrays
##-------------------------

function price2ret(prices::DataArray; log = true)
    ## find NAs
    naInds = find(isna, prices)
    nNAs = length(naInds)

    ## create price series with NAs replaced by last value
    pricesNoNA = deepcopy(prices)
    impute!(pricesNoNA, "single last")

    ## calculate returns
    if log
        rets = pricesNoNA[2:end] .- pricesNoNA[1:(end-1)]
    else
        rets = (pricesNoNA[2:end] .- pricesNoNA[1:(end-1)]) ./
        pricesNoNA[1:(end-1)]
    end

    ## fill in NAs again
    for ii=1:nNAs
        if naInds[ii] > 1
            rets[naInds[ii] - 1] = NA
        end
    end
    return rets
end

## price2ret for TimeData types
##-----------------------------

function price2ret(tn::AbstractTimenum; log = true)
    ## dealing with NAs through ret2price for DataArrays

    rets = DataFrame()
    for (nam, col) in eachcol(tn.vals)
        rets[nam] = price2ret(col, log = log)
    end

    return TimeData.Timenum(rets, tn.idx[2:end])
end

function price2ret(tm::AbstractTimematr; log = true)
    ## get discrete net returns from historic prices
    if log
        rets = tm[2:end, :] .- tm[1:(end-1), :] # time index of first
                                        # part is automatically taken!
    else
        rets = (tm[2:end, :] .- tm[1:(end-1), :]) ./ tm[1:(end-1), :]
    end
    return rets
end

###############
## ret2price ##
###############

## ret2price for DataArrays
##-------------------------

function ret2price(da::DataArray; log = true)
    ## append 0 or NA and transform returns to prices
    nObs = size(da, 1)
    prices = DataArray(eltype(da), nObs+1)

    if isna(da[1])
        ## find first element
        indFirstRet = find(x -> !isna(x), da)[1]
        prices[1:(indFirstRet-1)] = NA
    else
        indFirstRet = 1
    end

    ## different initial value for log / normal prices
    if log
        prices[indFirstRet] = 0
    else
        prices[indFirstRet] = 1
    end

    for ii=(indFirstRet+1):(nObs+1)
        if isna(da[ii-1])
            prices[ii] = prices[ii-1]
        else
            ## different aggregation formulas
            if log
                prices[ii] = prices[ii-1] + da[ii-1]
            else
                prices[ii] = prices[ii-1] .* (1 .+ da[ii-1])
            end
        end
    end

    ## insert NAs again
    inds = find(isna(da))
    for eachNA in inds
        prices[eachNA + 1] = NA
    end

    ## again insert initial price in case it was overwritten
    if log
        prices[indFirstRet] = 0
    else
        prices[indFirstRet] = 1
    end
    return prices
end

## ret2price for TimeData types
##-----------------------------

## ret2price for Timenum
##----------------------

function ret2price(tn::AbstractTimenum; log = true)
    ## dealing with NAs through ret2price for DataArrays

    prices = DataFrame()
    for (nam, col) in eachcol(tn.vals)
        prices[nam] = ret2price(col, log = log)
    end

    ## get previous day / could be not a business day
    initDate = TimeData.idx(tn)[1] - Dates.Day(1)

    ## get indices
    idxs = [initDate; tn.idx]

    return TimeData.Timenum(prices, idxs)
end

## ret2price for Timematr
##-----------------------

function ret2price(tm::AbstractTimematr; log = true)
    ## get normed prices from returns
    (nObs, nVars) = size(tm)

    if log
        prices = cumsum(tm, 1)

        ## get values of first day
        initPrices = zeros(1, nVars)
    else
        prices = cumprod(tm .+ 1, 1)

        ## get values of first day
        initPrices = ones(1, nVars)
    end

    ## get previous day / could be not a business day
    initDate = TimeData.idx(tm)[1] - Dates.Day(1)

    ## get first day as Timematr
    initPricesDf = TimeData.composeDataFrame(initPrices, names(tm))
    initPricesTm = TimeData.Timematr(initPricesDf,
                                     [initDate])

    return [initPricesTm; prices]

end

#########################################
## ret2price with given price metadata ##
#########################################

## ret2price with prices - Timenum
##--------------------------------

function ret2price(rets::AbstractTimenum,
                   prices::AbstractTimenum; log = true)

    ## find first prices
    (nObs, nVars) = size(rets)

    initPrices = DataFrame()
    for (nam, col) in eachcol(prices.vals)
        if isna(col[1])
            ## find first element
            indFirstPrice = find(x -> !isna(x), col)[1]
            initPrice = col[indFirstPrice]
        else
            initPrice = col[1]
        end
        initPrices[nam] = initPrice*ones(nObs+1)
    end

    ## get normalized prices via DataArrays
    normedPrices = DataFrame()
    for (nam, col) in eachcol(rets.vals)
        normedPrices[nam] = ret2price(col, log = log)
    end
    normedPricesIdxs = [prices.idx[1]; rets.idx]
    normedPricesTn = TimeData.Timenum(normedPrices, normedPricesIdxs)

    ## add / multiply initPrices to each element
    initPricesTn = TimeData.Timenum(initPrices, normedPricesIdxs)

    if log
        newPricesTn = normedPricesTn .+ initPricesTn
    else
        newPricesTn = normedPricesTn .* initPricesTn
    end

    return newPricesTn
end

## ret2price with prices - Timematr
##---------------------------------

function ret2price(tm::AbstractTimematr,
                   prices::AbstractTimenum; log = true)

    normalizedPrices = ret2price(tm, log = log)
    (nObs, nVars) = size(normalizedPrices)

    ## add initial prices to each element
    initPricesTm = convert(Timematr, prices[1, :])
    initPrices = composeDataFrame(repmat(core(initPricesTm), nObs, 1),
                                  names(normalizedPrices))
    initPricesTm = Timematr(initPrices, normalizedPrices.idx)

    if log
        newPrices = normalizedPrices .+ initPricesTm
    else
        newPrices = normalizedPrices .* initPricesTm
    end

    ## fix first date
    newPrices.idx[1] = prices.idx[1]

    return newPrices
end

## #####################################
## ## low-level aggregation functions ##
## #####################################

## function aggrLog(rets::Array{Float64, 2})
##     return sum(rets, 1)
## end

## function aggrLogPercent(rets::Array{Float64, 2})
##     return aggrLog(rets)
## end

## function aggrDiscr(rets::Array{Float64, 2})
##     return prod(rets+1, 1) - 1
## end

## function aggrDiscrPercent(rets::Array{Float64, 2})
##     return (prod(rets/100 + 1, 1) - 1)*100
## end

## #################
## ## aggregation ##
## #################

## function getLastPeriodIndices(n, by)
##     ## get start index
##     startInd = mod(n, by) + 1
##     return [(startInd+by-1):by:n]
## end

## function aggregate(tm::Timematr, by::Integer = 20,
##                    aggrFunc::Function = aggrLog)
##     ## aggregate by using most recent observations

##     (nObs, nAss) = size(tm)
##     vals = core(tm)

##     ## how many aggregations
##     nAggrRets = div(nObs, by)

##     ## get aggregation intervals and last interval dates
##     aggrIndices = getLastPeriodIndices(nObs, by)
##     aggrDates = idx(tm)[aggrIndices]

##     aggrRets = ones(nAggrRets, nAss)
##     for ii=1:nAggrRets
##         startInd = aggrIndices[ii]-by+1
##         aggrRets[ii, :] = aggrFunc(vals[startInd:aggrIndices[ii], :])
##     end

##     return Timematr(aggrRets, aggrDates)
## end

## ######################################
## ## get correlations for frequencies ##
## ######################################

## ## select pair of assets
## ## for each frequency
## ## - get number of possible starting points: remain + by
## ## - aggregate according to frequency in multiple ways
## ## - for each aggregation, get correlation
## ##

## function frequCorrPlot(tm::Timematr, ind1::Integer = 1,
##                        ind2::Integer = 2,
##                        maxFreq::Integer = 50)
##     ## select data
##     data = tm[:, [ind1, ind2]]
##     nObs = size(data, 1)

##     corrs = Float64[]
##     frequencies = Float64[]
##     for frequ=1:maxFreq
##         ## get number of possible phases
##         nPhases = mod(nObs, frequ) + frequ
##         nObsPerPhase = (div(nObs, frequ) - 1)*frequ
##         ## corrsCurrentFrequ = ones(nPhases)

##         for ii=1:nPhases
##             ## get current data
##             corrData = aggregate(data[ii:(ii + nObsPerPhase - 1), :],
##                                  frequ)
##             ## corrsCurrentFrequ[ii] = cor(corrData)[1, 2]
##             corr = cor(corrData)[1, 2]
##             push!(frequencies, frequ)
##             push!(corrs, corr)
##         end

##         ## push!(a, [frequ*ones(nPhases) corrsCurrentFrequ])
##     end

##     return [frequencies corrs]
## end


using StatsBase, Plots ; gr()

function proportionFished(gF,sF,n,N,withReplacement = false)
    function fishing()
        fishInPond = [ones(Int64,gF); zeros(Int64,sF)]
        fishCaught = Int64[]

        for fish in 1:n
            fished = rand(fishInPond)
            push!(fishCaught,fished)
            if withReplacement == false
                deleteat!(fishInPond, findfirst(x->x==fished, fishInPond))
            end
        end
        sum(fishCaught)
    end

    simulations = [fishing() for _ in 1:N]
    proportions = counts(simulations,0:n)/N

    if withReplacement
        plot!(0:n, proportions,
             line=:stem, marker=:circle, c=:blue, ms=6, msw=0,
             label="With replacement",
             xlabel="n",
             ylims=(0, 0.6), ylabel="Probability")
    else
        plot!(0:n, proportions,
            line=:stem, marker=:xcross, c=:red, ms=6, msw=0,
            label="Without replacement")
    end
end

N = 10^6
goldFish, silverFish, n = 3, 4, 3
plot()
proportionFished(goldFish, silverFish, n, N)
proportionFished(goldFish, silverFish, n, N, true)


#=

=#

using DataStructures
using IterTools
function value_iteration(S,A,P,R,Θ)
γ = 0.9
V₀ = DefaultDict{Tuple{Int,Int},Int}(0)
Vk::Array{Any} = [V₀]
k = 1
V₁::Dict{Tuple{Int,Int},Float64} = Dict{Tuple{Int,Int},Float64}()
# update value
while true
    k = k + 1
    push!(Vk,Dict{Tuple{Int,Int},Float64}())
    for s in S
        Vp = []
         for a in A
            #=@show s,a
            println(broadcast(s′ -> R(s,a,s′),[S...]))
            println(broadcast(s′ -> P(s′,s,a),[S...]))=#
            #push!(Vp,sum(broadcast(s′ -> P(s′,s,a)*R(s,a,s′),[S...])))
            push!(Vp, sum(broadcast(s′ -> P(s′,s,a)*(R(s,a,s′)+γ*Vk[k-1][s′]),[S...])))
        end
        #@show Vp
        Vk[k][s] = maximum(Vp)

    end
    abs(sum(abs,values(Vk[k])) - sum(abs,values(Vk[k-1]))) < Θ && break;
end
@show k
Vk[k]
end
ActionToVec = Dict(1 => (0,1),2 => (1,0), 3 => (0,-1), 4 => (-1,0))
function P(s′,s,a)

    Action = [0.1, 0.1, 0.1, 0.1]
    Action[a] = 0.7
    Vec = map(t -> ActionToVec[t],[1,2,3,4])
    Near = map(x -> s.+x,Vec)
    !any([s′] .== Near) && return 0

    Action[findin(Near,[s′])][1]

end
function R(s,a,s′)
    s == (0,0) && return 10
    s == (1,1) && return -10
    0
end


# # Classification of penguin species
#
# ## Packages

using AlgebraOfGraphics
using CairoMakie
using CalibrationErrors
using DataFrames
using Distributions
using MLJ
using MLJNaiveBayesInterface
using PalmerPenguins

using Random

## Plotting settings
set_aog_theme!()
CairoMakie.activate!(; type="svg")

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

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

penguins_mapping =
    data(penguins) * mapping(
        :bill_length_mm => "bill length (mm)", :flipper_length_mm => "flipper length (mm)"
    )
draw(penguins_mapping * mapping(; color=:species) * visual(; alpha=0.7))

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

Random.seed!(1234)
n = nrow(penguins)
k = floor(Int, 0.7 * n)
Random.seed!(100)
penguins.train = shuffle!(vcat(trues(k), falses(n - k)))

## Plot the training and validation data
dataset = :train => renamer(true => "training", false => "validation") => "Dataset"
plt = penguins_mapping * mapping(; color=:species, col=dataset) * visual(; alpha=0.7)
draw(plt; axis=(height=300,))

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

y, X = unpack(
    penguins,
    ==(:species),
    x -> x === :bill_length_mm || x === :flipper_length_mm;
    :species => Multiclass,
    :bill_length_mm => MLJ.Continuous,
    :flipper_length_mm => MLJ.Continuous,
)
model = fit!(machine(GaussianNBClassifier(), X, y); rows=penguins.train);

# We plot the estimated normal distributions.

## plot datasets
fg = draw(plt; axis=(height=300,))

## plot Gaussian distributions
xgrid = range(extrema(penguins.bill_length_mm)...; length=100)
ygrid = range(extrema(penguins.flipper_length_mm)...; length=100)
let f = (x, y, dist) -> pdf(dist, [x, y])
    for (class, color) in zip(classes(y), Makie.wong_colors())
        pdfs = f.(xgrid, ygrid', Ref(model.fitresult.gaussians[class]))
        contour!(fg.figure[1, 1], xgrid, ygrid, pdfs; color=color)
        contour!(fg.figure[1, 2], xgrid, ygrid, pdfs; color=color)
    end
end

fg

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

predictions = MLJ.predict(model)
train_predict = predictions[penguins.train]
val_predict = predictions[.!penguins.train]

## Plot datasets
fg = draw(plt; axis=(height=300,))

## Plot predictions
predictions_grid = reshape(
    MLJ.predict(model, reduce(hcat, vcat.(xgrid, ygrid'))'), length(xgrid), length(ygrid)
)
for (class, color) in zip(classes(y), Makie.wong_colors())
    p = pdf.(predictions_grid, class)
    contour!(fg.figure[1, 1], xgrid, ygrid, p; color=color)
    contour!(fg.figure[1, 2], xgrid, ygrid, p; color=color)
end

fg

# ## Evaluation
#
# We evaluate the probabilistic predictions of the naive Bayes classifier that we just
# trained.
#
# ### Log-likelihood
#
# We compute the average log-likelihood of the validation data. It is equivalent to the
# negative cross-entropy.

val_y = y[.!penguins.train]
-mean(cross_entropy(val_predict, val_y))

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

mean(brier_score(val_predict, val_y))

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

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

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

# We have to work with a numerical encoding of the true penguin species and a
# corresponding vector of predictions. We use [`RowVecs`](https://juliagaussianprocesses.github.io/KernelFunctions.jl/stable/api/#KernelFunctions.RowVecs)
# to indicate that the rows in the matrix of probabilities returned by `pdf`
# are the predictions. If we would provide predictions as columns of a matrix, we would have
# to use [`ColVecs`](https://juliagaussianprocesses.github.io/KernelFunctions.jl/stable/api/#KernelFunctions.ColVecs).

val_yint = map(MLJ.levelcode, val_y)
val_probs = RowVecs(pdf(val_predict, MLJ.classes(y)));

# We compute the estimate on the validation data:

ece(val_probs, val_yint)

# For the squared Euclidean distance we obtain:

ece = ECE(UniformBinning(10), SqEuclidean())
ece(val_probs, val_yint)

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

ece = ECE(MedianVarianceBinning(5), (μ, y) -> kl_divergence(y, μ))
ece(val_probs, val_yint)

# For the squared Euclidean distance we obtain:

ece = ECE(MedianVarianceBinning(5), SqEuclidean())
ece(val_probs, val_yint)

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

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

distances = pairwise(SqEuclidean(), RowVecs(pdf(train_predict, MLJ.classes(y))))
ν = sqrt(median(distances[i] for i in CartesianIndices(distances) if i[1] < i[2]))
kernel = with_lengthscale(GaussianKernel(), ν) ⊗ WhiteKernel();

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

skce = BiasedSKCE(kernel)
skce(val_probs, val_yint)

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

skce = UnbiasedSKCE(kernel)
skce(val_probs, val_yint)

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

skce = BlockUnbiasedSKCE(kernel, 2)
skce(val_probs, val_yint)


# Simple example of a homogeneous heat equation:
#
#           y
#             |
#             |     u_y = 0
#          Ly ----------------------------------
#             |                                |
#             |                                |
#             |                                |
#     u_x = 0 |    u_t - a(u_xx + u_yy) = 0    | u_x = 0
#             |                                |
#             |                                |
#             |                                |
#           -----------------------------------|-----  x
#                         u_y = 0              Lx
#
# Here, a>0 and the initial data is
#
#      u_0(x, y) = 1 + cos(pi x / Lx) cos(pi y / Ly).
#
# The exact solution is
#
#      u(x, y, t) = 1 + exp(-lambda t) cos(pi x / Lx) cos(pi y / Ly )
#
# with
#
#      lambda = a pi^2 ( 1 / Lx^2 + 1 / Ly^2 )
#
using PyPlot
using FinElt
using FinElt.PlanarPoisson
using LinearAlgebra: norm
include("ode23s.jl")

const a = 0.1
const Lx = 1.0
const Ly = 1.0
const lambda = a * pi^2 * ( 1/Lx^2 + 1/Ly^2 )
const T = 2.0

function exact_u(x, t)
    return 1 + exp(-lambda * t ) * cos(pi*x[1]/Lx) * cos(pi*x[2]/Ly)
end

function RHS(t, u, S)
    return -S * u
end

maxnorm(x) = norm(x, Inf)

start = time()
mesh = read_msh_file("../rectangle/rect2.msh")
dof = degrees_of_freedom(mesh, String[])
M = assembled_matrix("Omega", func_times_func!, 1.0, mesh, dof)
S = assembled_matrix("Omega", grad_dot_grad!, a, mesh, dof)
u0 = get_nodal_vals(x->exact_u(x,0.0), mesh)
t, u = ode23s((t,u)->RHS(t, u, S), u0, [0,T];
             reltol=1.0e-5, abstol=1.0e-5, mass=M, norm=maxnorm)

N = length(t)
err = zeros(N)
for n = 1:N
    un = get_nodal_vals(x->exact_u(x,t[n]), mesh)
    err[n] = maxnorm(un-u[n])
end

finish = time()
elapsed = finish - start

figure(1)
plot(t, err)
xlabel(L"t")
title("Error solving the heat equation")
grid(true)


# Check of handedness
@test rotationmatrix_normalized(qrotation_x(π/2)) * [0, 1, 0] ≈ [0, 0, 1]
@test rotationmatrix_normalized(qrotation_y(π/2)) * [1, 0, 0] ≈ [0, 0, -1]
@test rotationmatrix_normalized(qrotation_z(π/2)) * [1, 0, 0] ≈ [0, 1, 0]

# Simple composition of rotations
@test qrotation_x(π/3) * qrotation_x(π/3) ≈ qrotation_x(2π/3)
@test qrotation_y(π/3) * qrotation_y(π/3) ≈ qrotation_y(2π/3)
@test qrotation_z(π/3) * qrotation_z(π/3) ≈ qrotation_z(2π/3)

# More complex compositions
comp = qrotation_y(π/2) * qrotation_x(π/2)
@test rotationmatrix_normalized(comp) * [0, 1, 0] ≈ [1, 0, 0]

comp = qrotation_z(π/2) * qrotation_y(π/2)
@test rotationmatrix_normalized(comp) * [0, 0, 1] ≈ [0, 1, 0]


using Printf
import Flux: σ
using ModelingToolkit
using GalacticOptim
using Optim
using DiffEqFlux
using NeuralPDE
using Quadrature, Cubature, Cuba
using Plots

@parameters t,x
@variables c(..)
@derivatives Dt'~t
@derivatives Dxx''~x
@derivatives Dx'~x

# Parameters

v = 1
R = 0
D = 0.1 # diffusion
t_max = 2.0
x_min = -1.0
x_max = 1.0

# Equations, initial and boundary conditions
eqs = [ Dt(c(t, x)) ~ D * Dxx(c(t,x)) - Dx(c(t,x)) ]

bcs = [
        c(0, x) ~ cos(π*x) + 1.0,
        c(t, x_min) ~ c(t, x_max)

]

# Space and time domains
domains = [t ∈ IntervalDomain(0.0,t_max),
        x ∈ IntervalDomain(x_min,x_max)
]

# Discretization
nx = 32
dx = (x_max-x_min) / (nx - 1)
dt = 0.01

# Neural network
dim = length(domains)
output = length(eqs)
hidden = 8

chain = FastChain( FastDense(dim, hidden, σ),
                    FastDense(hidden, hidden, σ),
                    FastDense(hidden, 1))

strategy = GridTraining(dx=[dt,dx])

discretization = PhysicsInformedNN(chain, strategy=strategy)

pde_system = PDESystem(eqs, bcs, domains, [t,x], [c])
prob = discretize(pde_system,discretization)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

res = GalacticOptim.solve(prob,Optim.BFGS();cb=cb,maxiters=1200)

# Plots

phi = discretization.phi

initθ = discretization.initθ

acum =  [0;accumulate(+, length.(initθ))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
minimizers = [res.minimizer[s] for s in sep]
ts,xs = [domain.domain.lower:dx:domain.domain.upper for domain in domains]

anim = @animate for (i, t) in enumerate(0:dt:t_max)
    @info "Animating frame $i..."
    c_predict = reshape([phi([t, x], res.minimizer)[1] for x in xs], length(xs))
    title = @sprintf("Advection-diffusion t = %.3f", t)
    plot(xs, c_predict, label="", title=title , ylims=(0., 2))
end

gif(anim, "advection_diffusion_pinn.gif", fps=15)

c_predict = reshape([ phi([0, x], res.minimizer)[1] for x in xs], length(xs))
plot(xs,c_predict)

# Plot correct solution
using JLD2
file = jldopen("advection_diffusion/simulation/cosine_advection_diffusion.jld2")
iterations = parse.(Int, keys(file["timeseries/t"]))

anim = @animate for (i, iter) in enumerate(iterations)
    @info "Animating frame $i..."
    Hx = file["grid/Hx"]
    x = file["grid/xC"][1+Hx:end-Hx]
    t = file["timeseries/t/$iter"]
    c = file["timeseries/c/$iter"][:]

    title = @sprintf("Advection-diffusion t = %.3f", t)
    p = plot(x, c .+ 1, linewidth=2, title=title, label="Oceananigans",
             xlabel="x", ylabel="Tracer", xlims=(-1, 1), ylims=(0, 2))

    c_predict = reshape([phi([t, x], res.minimizer)[1] for x in xs], length(xs))

    plot!(p, x, c_predict, linewidth=2, label="Neural PDE")
end

gif(anim, "advection_diffusion_comparison.gif", fps=15)


module SimpleDrawing
using Plots, LinearAlgebra

export newdraw, draw_circle, draw_arc, draw_segment, draw_point, draw_vector
export find_center, non_colinear_check
export finish, draw

"""
`newdraw()` clears the screen and sets up a blank drawing canvas with
no axis, grid, or legend, and sets the aspect ratio to 1.
"""
function newdraw()
    plot(aspectratio=1, legend=false, axis=false, grid=false)
end

function draw()
end

"""
`draw_circle(x,y,r)` (or `draw_circle(z,r)`) draws a circle centered
at `(x,y)` (at `z`) with radius `r`.
"""
function draw_circle(x::Real, y::Real, r::Real; opts...)
    f(t) = r*cos(t) + x
    g(t) = r*sin(t) + y
    plot!(f,g,0,2pi;opts...)
end

draw_circle(z::Complex, r::Real; opts...) =
    draw_circle(real(z), imag(z), r; opts...)

"""
`draw_circle(a,b,c)` draws a circle through the points given by
the three complex arguments.
"""
function draw_circle(a::Complex, b::Complex, c::Complex; opts...)
    z = find_center(a,b,c)
    r = abs(a-z)
    x,y = reim(z)
    draw_circle(x,y,r; opts...)
end

"""
`draw_arc(x,y,r,t1,t2)` draws the shorter arc centered at `(x,y)`
with radius `r` between angles `t1` and `t2`.
"""
function draw_arc(x::Real, y::Real, r::Real, t1::Real, t2::Real; opts...)
    f(t) = r*cos(t) + x
    g(t) = r*sin(t) + y

    plot!(f,g,t1,t2;opts...)
end

"""
`draw_arc(a,b,c)` where the arguments are `Complex`: Draw the
arc from `a` through `b` to `c`.
"""
function draw_arc(a::Complex, b::Complex, c::Complex; opts...)
    m2pi(x::Real) = mod(x,2π)

    z = find_center(a,b,c)
    x,y = reim(z)
    r = abs(b-z)
    ta = m2pi(angle(a-z))
    tb = m2pi(angle(b-z))
    tc = m2pi(angle(c-z))

    if ta > tc
        ta,tc = tc,ta
    end

    if ta<tb<tc
        return draw_arc(x,y,r,ta,tc; opts...)
    end

    return draw_arc(x,y,r,tc-2pi,ta; opts...)
end

"""
`draw_segment(x,y,xx,yy)` draws a line segment from `(x,y)` to
`(xx,yy)`. Also `draw_segment(z,zz)` for `Complex` arguments.
"""
function draw_segment(a::Real,b::Real,c::Real,d::Real; opts...)
    plot!([a,c],[b,d];opts...)
end

function draw_segment(a::Complex, b::Complex; opts...)
    x,y = reim(a)
    xx,yy = reim(b)
    draw_segment(x,y,xx,yy; opts...)
end

"""
`draw_point(x,y)` [or `draw_point(z)`] draws a point at coordinates `(x,y)`
[or complex location `z`].

`draw_point(list)` plots all the points in the one-dimensional `list` of
complex values.
"""
function draw_point(x::Real, y::Real; opts...)
    plot!([x],[y];marker=1,opts...)
end

draw_point(z::Complex; opts...) = draw_point(real(z),imag(z); opts...)

function draw_point(pts::Array{Complex{T},1};opts...) where T
    for p in pts
        draw_point(p;opts...)
    end
end

"""
`draw_vector` is used to draw vectors (line segments with an arrow at one end).
The variations are:
+ `draw_vector(x,y)` draws a vector from the origin to `(x,y)`.
+ `draw_vector(x,y,basex,basey)` draws a vector from `(basex,basey)` to
`(basex+x,basey+y)`.
+ `draw_vector(z)` draws a vector from the origin to the complex value `z`.
+ `draw_vector(z,basez)` draws a vector from the complex location `basez` to
`z+basez`.
"""
draw_vector(x::Real,y::Real; opts...) = draw_segment(0,0,x,y;arrow=:arrow, opts...)

draw_vector(z::Complex;opts...) = draw_segment(0+0im, z; arrow=:arrow, opts...)

draw_vector(x::Real,y::Real,basex::Real,basey::Real; opts...) =
    draw_segment(basex,basey,basex+x,basey+y; arrow=:arrow, opts...)

draw_vector(z::Complex, basez::Complex; opts...) =
    draw_segment(basez, basez+z; arrow=:arrow, opts...)

"""
`finish()` is used to clean up a drawing after various calls to `draw`.
It removes the axes and the grid lines, and sets the aspect ratio to one.
"""
function finish()
    plot!(aspectratio=1, legend=false, axis=false, grid=false)
end

export find_center, non_colinear_check

"""
`find_center(a,b,c)`: Given three points in the complex plane,
find the point `z` that is equidistant from all three. If the three
points are collinear then return `Inf + Inf*im`.
"""
function find_center(a::Complex, b::Complex, c::Complex)::Complex
    if !non_colinear_check(a,b,c)
        return Inf + im*Inf
    end
    A = collect(reim(a))
    B = collect(reim(b))
    C = collect(reim(c))
    AB = 0.5*(A+B)
    BC = 0.5*(B+C)
    M = [(A-B)'; (B-C)']
    rhs = [dot(AB,A-B), dot(BC,B-C)]

    Z = M\rhs
    return Z[1] + im*Z[2]
end

"""
`non_colinear_check(a,b,c)`: test if the complex numbers are distinct
and noncollinear.
"""
function non_colinear_check(a::Complex,b::Complex, c::Complex)::Bool
    if a==b || b==c || a==c
        return false
    end
    z = (b-a)/(c-a)
    return imag(z) != 0
end

include("cubic.jl")
include("curve.jl")

end


# Magic Square Generator
# ======================
#
# # Reference
#
# - Moler, Cleve B. Experiments with MATLAB. Society for Industrial and Applied Mathematics, 2011.

# The JuliaLang community is famous for naming their functions as something like
# `verylongfunctionnamethatistotallyunreadable` and abusing unicode characters all over the source.
# A more appropriate name for this is probably `is🎩`, but let's stick with the old style for now.
function is_magic(m::AbstractMatrix{<:Integer})
    n_r, n_c = size(m)
    # Is square?
    if n_r != n_c
        return false
    end

    n = n_r
    magic_const = (1 + n ^ 2) * n ÷ 2

    # Are all rows equal to magic constant?
    if !all(sum(m, dims=2) .== magic_const)
        return false
    end
    # Are all columns equal to magic constant?
    if !all(sum(m, dims=1) .== magic_const)
        return false
    end
    # And the two diagonals?
    if sum([m[i, i] for i in 1:n]) != magic_const
        return false
    end
    if sum([m[p...] for p in zip(1:n, n:-1:1)]) != magic_const
        return false
    end

    true
end

# Odd ordered magic square generator
#
# It is easy to see the correctness of this algorithm. For example, compare the elements in each
# column with the one to the northeast of itself, there are `n - 1` elements that are 1 less while
# exactly 1 that is `n - 1` larger, which makes the two columns add up to the same value.

"Generates a magic square of an odd order `n` using Siamese method (De la Loubère method)."
function odd_magic(n::Integer)
    @assert isodd(n)

    m = zeros(typeof(n), n, n)
    r, c = 1, div(n + 1, 2)

    for i in 1 : n*n
        m[r, c] = i

        # No need to do bounds checking, as only the last one will fall out of the board.
        p = r + 1, c
        r, c = mod1(r - 1, n), mod1(c + 1, n)
        if m[r, c] != 0
            r, c = p
        end
    end

    m
end

# A vectorized version of `odd_magic`.
function odd_magic_vec(n::Integer)
    @assert isodd(n)
    i = 1:n
    offset = div(n - 3, 2)
    A = @. mod(i + i' + offset, n)
    B = @. mod(i + i' * 2 - 2, n)
    @. n * A + B + 1
end

# TODO: Even ordered magic square generator
function singly_even_magic(n::Integer)
    error("unimplemented")
end

function doubly_even_magic(n::Integer)
    error("unimplemented")
end

function magic(n::Integer)
    if isodd(n)
        odd_magic(n)
    else
        half_n = div(n, 2)
        if isodd(half_n)
            singly_even_magic(n)
        else
            doubly_even_magic(n)
        end
    end
end

# Unit tests
using Test

@testset "Odd magic square tests" begin
    orders = 1:2:21
    for n in orders
        m = odd_magic(n)
        @test m == odd_magic_vec(n)
        @test is_magic(m)
    end
end


#provide a pair of points in latitude and longitude, in 'degrees' to get
#the distance back in Km
function getDistanceFromLatLongInKm(lat1,long1,lat2,long2)
  RadiusEarth = 6371; # Earth's radius in km
  dLat = deg2rad(lat2-lat1); # deg2rad, built in fn in Julia
  dLong = deg2rad(long2-long1);
  h = sin(dLat/2)*sin(dLat/2) + cos(deg2rad(lat1))*cos(deg2rad(lat2))*sin(dLong/2)*sin(dLong/2)
  d = RadiusEarth * 2 * atan2(sqrt(h),sqrt(1-h)); # Distance in km
  return d;
end


#
# Denoise an SPD Example with Cyclic Proximal Point applied to the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data as for the corresponding CP algorithm
#
using Manopt
using Images, CSV, DataFrames, LinearAlgebra, JLD2
#
# Settings
ExportResult = true
ExportOrig = true
ExportResultVideo = false
ExportTable = true
#
# Manifold and Data
f = artificialSPDImage2(32)
pixelM = SymmetricPositiveDefinite(3)
resultsFolder = "src/examples/Total_Variation/SPD_TV/"
experimentName = "ImageCPPA"
if !isdir(resultsFolder)
    mkdir(resultsFolder)
end
if ExportOrig
    renderAsymptote(resultsFolder*experimentName*"-orig.asy"; data=f, scaleAxes=(7.5,7.5,7.5))
end
#
# Parameters
α = 6.
maxIterations = 4000
#
# Build Problem for L2-TV
M = Power(pixelM,size(f))
d = length(size(f))
rep(d) = (d>1) ? [ones(Int,d)...,d] : d
fidelity(x) = 1/2*distance(M,x,f)^2
Λ(x) = forwardLogs(M,x) # on T_xN
prior(x) = norm(norm.(Ref(pixelM),getValue(repeat(x,rep(d)...)),getValue(Λ(x))),1)
#
# Setup and Optimize
cost(x) =  fidelity(x) + α*prior(x)
proximalMaps = [(λ,x) -> proxDistance(M,λ,f,x,2), (λ,x) -> proxTV(M,α*λ,x,1)]
x0 = f
@time y, yRec = cyclicProximalPoint(M,cost,proximalMaps,x0;
debug = [:Iteration," | ", DebugProximalParameter()," | ", :Change, " | ", :Cost, "\n",100,:Stop],
record = [:Iteration, :Iterate, :Cost],
stoppingCriterion = stopAfterIteration(maxIterations)
)
#
# Results
if ExportResult
    renderAsymptote(resultsFolder*experimentName*"-result-$(maxIterations)-α$(replace(string(α), "." => "-")).asy"; data=y, render=4, scaleAxes=(7.5,7.5,7.5) )
end
if ExportTable
    A = cat( [ y[1] for y in yRec], [y[3] for y in yRec]; dims=2 )
    CSV.write(string(resultsFolder*experimentName*"ResultCost.csv"),
    DataFrame(A), writeheader=false);
    save(resultsFolder*experimentName*"-CostValue.jld2",
    Dict("compareCostFunctionValue" => last(yRec)[3])
    )
end


## Environment and packages
cd(@__DIR__)
using Pkg; Pkg.activate("."); Pkg.instantiate()

using OrdinaryDiffEq
using ModelingToolkit
using DataDrivenDiffEq
using LinearAlgebra, Optim
using DiffEqFlux, Flux
using Plots
gr()
using JLD2, FileIO
using Statistics
# Set a random seed for reproduceable behaviour
using Random
Random.seed!(1235)

# Create a name for saving ( basically a prefix )
svname = "Scenario_3_"

## Data Generation
# Parameters
D = 0.01f0; #diffusion
r = 1.0f0; #reaction rate

# Domain
X = 1.0f0; T = 5f0;
dx = 0.04f0; dt = T/10;
x = collect(0:dx:X);
t = collect(0:dt:T);
Nx = Int64(X/dx+1);
Nt = Int64(T/dt+1);

#initial conditions
Amp = 1.0f0;
Delta = 0.5f0;
rho_ic(Δ) = Amp*(tanh.((x .- (0.5f0 - Δ/2))/(Δ/10)) - tanh.((x .- (0.5f0 + Δ/2))/(Δ/10)))/2

# Reaction term
reaction(u) = r * u .* (1f0 .- u)
# Difference Op
lap = Float32.(diagm(0 => -2.0f0 * ones(Nx), 1=> ones(Nx-1), -1 => ones(Nx-1)) ./ dx^2)
#Periodic BC
lap[1,end] = 1.0f0/dx^2
lap[end,1] = 1.0f0/dx^2

function rc_ode(rho, p, t)
    #finite difference
    D * lap * rho + reaction.(rho)
end

# Generate random measurements
probpde = ODEProblem(rc_ode, rho_ic(Delta), (0.0f0, T), saveat=dt)
solution = solve(probpde, Tsit5());

# Ideal data
X = Array(solution)
t = solution.t

# Add noise in terms of the mean
x̄ = mean(X, dims = 2)
noise_magnitude = Float32(2.5e-2)
Xₙ = X .+ (noise_magnitude*x̄) .* randn(eltype(X), size(X))

pl_solution_1 = plot(solution, color = :black, ylabel = "rho(x,t)", label = ["True Data" [nothing for i in 1:Nx-1]...])
scatter!(t, transpose(Xₙ), color = :red,label = ["Noisy Data" [nothing for i in 1:Nx-1]...], legend =:bottomright)
# Create dataset plot
pl_contour = contour(t,x,Xₙ, xlabel = "t", ylabel = "x", fill = (true, :thermal))
pl_initial_data = plot(pl_solution_1, pl_contour, layout = (1,2))
savefig(pl_initial_data, joinpath(pwd(), "plots","$(svname)full_data_$(noise_magnitude).pdf"))

## Neural Network setup
## Define the network
# Gaussian RBF as activation
rbf(x) = exp.(-(x.^2))

# Create the network to recover the reaction
rx_nn = FastChain(
    FastDense(1, 5, rbf),
    FastDense(5, 5, rbf),
    FastDense(5, 5, rbf),
    FastDense(5, 1)
)

#initialize D0 close to D/dx^2
D0 = Float32[6.5]
p1s = initial_params(rx_nn)
p2s = zeros(Float32, 4)
p = [p1s;p2s;D0]

function nn_ode(u,p,t)
    # This is the dircetional derivative
    u_cnn_1   = [p[end-4] * u[end] + p[end-3] * u[1] + p[end-2] * u[2]]
    u_cnn     = [p[end-4] * u[i-1] + p[end-3] * u[i] + p[end-2] * u[i+1] for i in 2:size(u, 1)-1]
    u_cnn_end = [p[end-4] * u[end-1] + p[end-3] * u[end] + p[end-2] * u[1]]

    [rx_nn([ui], p[1:length(p1s)])[1] for ui in u] + p[end] * vcat(u_cnn_1, u_cnn, u_cnn_end)
end

prob_nn = ODEProblem(nn_ode, Xₙ[:, 1], (t[1], t[end]), p)

## Necessary functions

function predict_pde(θ, X = Xₙ[:, 1], T = t)
  # No ReverseDiff if using Flux
  Array(solve(prob_nn, Vern7(),
        u0 = X, p = θ, tspan = (T[1], T[end]), saveat=T,
        sensealg=ForwardDiffSensitivity(convert_tspan = false)#, reltol = 1e-4, abstol = 1e-4
        ))
end

#match data and force the weights of the CNN to add up to zero
function objective_pde(p)
    pred = predict_pde(p)
    sum(abs2, pred .- Xₙ) + abs(sum(p[end-4 : end-2]))
end

# Container to track the loss
losses = Float32[]

# Callback to see the training
callback(θ,l) = begin
    push!(losses, l)
    if length(losses)%1==0
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    false
end

## Train the model
# First train with ADAM
res1 = DiffEqFlux.sciml_train(objective_pde, p, ADAM(0.1f0), cb=callback, maxiters = 40)
# Train with BFGS
res2 = DiffEqFlux.sciml_train(objective_pde, res1.minimizer, BFGS(initial_stepnorm=0.1f0), cb=callback, maxiters = 10000)
p_trained = res2.minimizer

## Data evaluation -> we subsample here!
tsample = t[1]:1.0*mean(diff(t)):t[end]
X̂ = predict_pde(p_trained, Xₙ[:,1], tsample)
# Trained on noisy data vs real solution
pl_trajectory = plot(tsample, transpose(X̂), xlabel = "t", ylabel ="rho(x,t)", color = :red,
 label = ["UDE Approximation" [nothing for i in 1:25]...])
scatter!(t, transpose(Xₙ), color = :black, label = ["Measurements"  [nothing for i in 1:25]...], legend = :bottomright)
savefig(pl_trajectory, joinpath(pwd(), "plots", "$(svname)_trajectory_reconstruction.pdf"))

# Create estimates
R̂ = similar(X̂)
R̄ = similar(X̂)
for i in 1:size(X̂, 1), j in 1:size(X̂, 2)
    R̂[i,j] = rx_nn(X̂[i,j], p_trained[1:length(p1s)])[1]
    R̄[i,j] = reaction(X̂[i,j])
end

pl_reconstruction = plot(tsample, R̂', label = ["UDE Approximation" [nothing for i in 1:25]...], color = :red, xlabel = "t", ylabel = "Reaction(x,t)")
plot!(tsample, R̄', label = ["True Reaction" [nothing for i in 1:25]...], color = :black)
pl_reconstruction_error = plot(tsample, norm.(eachcol(R̂ - R̄)), color = :red, xlabel = "t", ylabel = "L2 Error")
pl_missing = plot(pl_reconstruction, pl_reconstruction_error, layout = (2,1))
savefig(pl_missing, joinpath(pwd(), "plots", "$(svname)_missingterm_reconstruction_and_error.pdf"))
pl_overall = plot(pl_trajectory, pl_missing)
savefig(pl_overall, joinpath(pwd(), "plots", "$(svname)_reconstruction.pdf"))
## Symbolic regression via sparse regression / SINDy using DataDrivenDiffEq

# Reshape the reactions
@variables u[1:1]
b = Basis(monomial_basis(u, 10), u)
# We assume (via the modeling) a common mapping of the states
# via the reaction. In other words a global function. Hence, we can
# use all measurements as one variable
Xs = Matrix(vcat(eachrow(X̂[1:end, :])...)')
Rs = Matrix(vcat(eachrow(R̂[1:end, :])...)')

# Technically this is cheating until a general DataDrivenProblem is defined properly.
# we work with the continuous form
nn_prob = ContinuousDataDrivenProblem(Xs, zeros(Float32, size(Xs, 2)), DX = Rs)

λs = Float32.(exp10.(-3:0.01:5))
opt = STLSQ(λs)

nn_res = solve(nn_prob, b, opt, maxiters = 1000, progress = true, denoise = true, normalize = false)
# Print the output
println(nn_res)
println(result(nn_res))
println(parameters(nn_res))
## Save the results

save(joinpath(pwd(), "results" ,"$(svname)recovery_$(noise_magnitude).jld2"),
    "solution", solution, "X", Xₙ, "t" , t, "neural_network" , rx_nn, "initial_parameters", p, "trained_parameters" , p_trained, # Training
    "losses", losses, "result", nn_res,  # Recovery
    ) # Estimation


"""
    GoldenSection <: UnivariateAlgorithm

The golden section search algorithm optimize a univariate function by iteratively shrinking
the initial bracket around the minimum with the golden ratio as an approximation to the
Fibonacci sequence.
"""
struct GoldenSection <: UnivariateAlgorithm end

function _optimize(f, lower::T, upper::T, alg::GoldenSection;
        reltol, abstol, maxiter) where {T}
    p = φ - 1
    x = T(p * lower + (1 - p) * upper)
    yx = f(x)
    converged = false
    iter = 1
    while true
        x_tol = reltol * abs(x) + abstol
        if abs(upper - lower) < 2x_tol
            converged = true
            break
        end
        iter == maxiter && break
        iter += 1
        new_x = T(p * upper + (1 - p) * lower)
        new_yx = f(new_x)
        if yx > new_yx
            lower, x, yx = x, new_x, new_yx
        else
            lower, upper = new_x, lower
        end
    end
    return Solution(converged, iter, x, yx)
end